Initial rag (#52)

* update signature

* optionally trim context

* handle context limits

* bump tomls

* cleanup debug

* update api name

* change to rag

* update test name

* try dropping

* update ci

* re-create ext w/ each test

* fix format

.github/workflows/extension_ci.yml

@@ -75,7 +75,7 @@ jobs:
         with:
           working-directory: ./
       - name: Cargo format
-        run: cargo +nightly fmt --all --check
+        run: cargo fmt --all --check
       - name: Clippy
         run: cargo clippy
 
@@ -123,13 +123,13 @@ jobs:
           rm -rf ./target/pgrx-test-data-* || true
       - name: unit-test
         run: |
-          cargo pgrx test
+          make test-unit
       - name: integration-test
         run: |
           pgrx15_config=$(/usr/local/bin/stoml ~/.pgrx/config.toml configs.pg15)
           pg_version=$(/usr/local/bin/stoml Cargo.toml features.default)
           echo "\q" | make run
-          cargo test -- --ignored --test-threads=1
+          make test-integration
 
   publish:
     if: github.event_name == 'release'

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "vectorize"
-version = "0.9.0"
+version = "0.10.0"
 edition = "2021"
 publish = false
 

Makefile

@@ -26,3 +26,10 @@ pgxn-zip: $(DISTNAME)-$(DISTVERSION).zip
 
 clean:
 	@rm -rf META.json $(DISTNAME)-$(DISTVERSION).zip
+
+
+test-integration:
+	cargo test -- --ignored --test-threads=1
+
+test-unit:
+	cargo pgrx test

Trunk.toml

@@ -6,7 +6,7 @@ description = "The simplest way to orchestrate vector search on Postgres."
 homepage = "https://github.com/tembo-io/pg_vectorize"
 documentation = "https://github.com/tembo-io/pg_vectorize"
 categories = ["orchestration", "machine_learning"]
-version = "0.9.0"
+version = "0.10.0"
 
 [build]
 postgres_version = "15"

src/api.rs

@@ -89,7 +89,7 @@ fn transform_embeddings(
 
 #[allow(clippy::too_many_arguments)]
 #[pg_extern]
-fn chat_table(
+fn init_rag(
     agent_name: &str,
     table_name: &str,
     unique_record_id: &str,
@@ -121,14 +121,30 @@ fn chat_table(
 
 /// creates an table indexed with embeddings for chat completion workloads
 #[pg_extern]
-fn chat(
+fn rag(
     agent_name: &str,
     query: &str,
+    // chat models: currently only supports gpt 3.5 and 4
+    // https://platform.openai.com/docs/models/gpt-3-5-turbo
+    // https://platform.openai.com/docs/models/gpt-4-and-gpt-4-turbo
     chat_model: default!(&str, "'gpt-3.5-turbo'"),
+    // points to the type of prompt template to use
     task: default!(&str, "'question_answer'"),
     api_key: default!(Option<&str>, "NULL"),
+    // number of records to include in the context
+    num_context: default!(i32, 2),
+    // truncates context to fit the model's context window
+    force_trim: default!(bool, false),
 ) -> Result<TableIterator<'static, (name!(chat_results, pgrx::JsonB),)>> {
-    let resp = call_chat(agent_name, query, chat_model, task, api_key)?;
+    let resp = call_chat(
+        agent_name,
+        query,
+        chat_model,
+        task,
+        api_key,
+        num_context,
+        force_trim,
+    )?;
     let iter = vec![(pgrx::JsonB(serde_json::to_value(resp)?),)];
     Ok(TableIterator::new(iter))
 }