Handle context window

Cargo.toml

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

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))
 }