fix: Retrieve correct doc IDs

julep-ai · May 16, 2024 · c93ef10 · c93ef10
1 parent 24e61ad
commit c93ef10
Showing 1 changed file with 30 additions and 8 deletions.
diff --git a/agents-api/agents_api/models/entry/proc_mem_context.py b/agents-api/agents_api/models/entry/proc_mem_context.py
@@ -12,7 +12,7 @@ def proc_mem_context_query(
     tools_confidence: float = 0,
     docs_confidence: float = 0.4,
     k_tools: int = 3,
-    k_docs: int = 2,
+    k_docs: int = 3,
 ) -> tuple[str, dict]:
     """Executes a complex query to retrieve memory context based on session ID, tool and document embeddings.
 
@@ -172,7 +172,7 @@ def proc_mem_context_query(
         }}
     }} {{
         # Collect document information based on agent ID and document query embedding.
-        # Collect docs
+        # Collect agent docs
 
         # Search for agent docs
         ?[role, name, content, token_count, created_at, index, agent_doc_id, user_doc_id] :=
@@ -201,6 +201,21 @@ def proc_mem_context_query(
             index = 5 + (snippet_idx * 0.01),
             user_doc_id = null,
 
+        # Save in temp table
+        :create _agent_docs {{
+            role: String,
+            content: String,
+            token_count: Int,
+            created_at: Float,
+            index: Float,
+            name: String? default null,
+            agent_doc_id: Uuid? default null,
+            user_doc_id: Uuid? default null,
+        }}
+    }} {{
+        # Collect document information based on agent ID and document query embedding.
+        # Collect user docs
+
         # Search for user docs
         ?[role, name, content, token_count, created_at, index, user_doc_id, agent_doc_id] :=
             *_input{{user_id, doc_query}},
@@ -229,15 +244,15 @@ def proc_mem_context_query(
             agent_doc_id = null,
 
         # Save in temp table
-        :create _docs {{
+        :create _user_docs {{
             role: String,
-            name: String?,
             content: String,
             token_count: Int,
             created_at: Float,
             index: Float,
-            agent_doc_id: Uuid default null,
-            user_doc_id: Uuid default null,
+            name: String? default null,
+            agent_doc_id: Uuid? default null,
+            user_doc_id: Uuid? default null,
         }}
     }} {{
         # Collect all entries related to the session.
@@ -302,9 +317,16 @@ def proc_mem_context_query(
             agent_doc_id = null, user_doc_id = null,
 
         ?[role, name, content, token_count, created_at, index, agent_doc_id, user_doc_id] :=
-            *_docs {{
-                role, name, content, token_count, created_at, index, agent_doc_id, user_doc_id
+            *_agent_docs {{
+                role, name, content, token_count, created_at, index, agent_doc_id
             }},
+            user_doc_id = null,
+        
+        ?[role, name, content, token_count, created_at, index, agent_doc_id, user_doc_id] :=
+            *_user_docs {{
+                role, name, content, token_count, created_at, index, user_doc_id
+            }},
+            agent_doc_id = null,
 
         ?[role, name, content, token_count, created_at, index, agent_doc_id, user_doc_id] :=
             *_entries{{