Merge pull request #948 from julep-ai/f/add-columns-to-indices

feat(agents-api): Add missing columns to indices
julep-ai · Dec 12, 2024 · ce3e3fc · ce3e3fc
2 parents bf48324 + 24e9afb
commit ce3e3fc
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 7 deletions.
diff --git a/agents-api/agents_api/models/docs/search_docs_by_embedding.py b/agents-api/agents_api/models/docs/search_docs_by_embedding.py
@@ -91,7 +91,7 @@ def search_docs_by_embedding(
         snippet_counter[count(item)] :=
             owners[owner_type, owner_id_str],
             owner_id = to_uuid(owner_id_str),
-            *docs {{
+            *docs:owner_id_metadata_doc_id_idx {{
                 owner_type,
                 owner_id,
                 doc_id: item,
@@ -142,13 +142,16 @@ def search_docs_by_embedding(
                 input[owner_type, owner_id, query],
 
                 # Restrict the search to all documents that match the owner
-                *docs {{
+                *docs:owner_id_metadata_doc_id_idx {{
                     owner_type,
                     owner_id,
                     doc_id,
-                    title,
                     metadata,
                 }},
+                *docs {{
+                    doc_id,
+                    title,
+                }},
 
                 # Search for snippets in the embedding space
                 ~snippets:embedding_space {{

diff --git a/agents-api/agents_api/models/docs/search_docs_by_text.py b/agents-api/agents_api/models/docs/search_docs_by_text.py
@@ -90,7 +90,7 @@ def search_docs_by_text(
 
         candidate[doc_id] :=
             input[owner_type, owner_id],
-            *docs {{
+            *docs:owner_id_metadata_doc_id_idx {{
                 owner_type,
                 owner_id,
                 doc_id,
@@ -145,13 +145,16 @@ def search_docs_by_text(
             metadata,
         ] :=
             candidate[doc_id],
-            *docs {{
+            *docs:owner_id_metadata_doc_id_idx {{
                 owner_type,
                 owner_id,
                 doc_id,
-                title,
                 metadata,
             }},
+            *docs {{
+                doc_id,
+                title,
+            }},
             search_result [
                 doc_id,
                 snippet_data,

diff --git a/agents-api/agents_api/models/execution/update_execution.py b/agents-api/agents_api/models/execution/update_execution.py
@@ -79,7 +79,7 @@ def update_execution(
 
     validate_status_query = """
     valid_status[count(status)] :=
-        *executions {
+        *executions:execution_id_status_idx {
             status,
             execution_id: to_uuid($execution_id),
             task_id: to_uuid($task_id),

diff --git a/agents-api/migrations/migrate_1733985509_add_columns_to_indices.py b/agents-api/migrations/migrate_1733985509_add_columns_to_indices.py
@@ -0,0 +1,26 @@
+# /usr/bin/env python3
+
+MIGRATION_ID = "add_columns_to_indices"
+CREATED_AT = 1733985509.241258
+
+
+def up(client):
+    client.run(
+        """
+        { ::index drop executions:execution_id_status_idx }
+        { ::index create executions:execution_id_status_idx { task_id, execution_id, status } }
+        { ::index drop docs:owner_id_metadata_doc_id_idx }
+        { ::index create docs:owner_id_metadata_doc_id_idx { owner_type, owner_id, metadata, doc_id } }
+        """
+    )
+
+
+def down(client):
+    client.run(
+        """
+        { ::index drop executions:execution_id_status_idx }
+        { ::index create executions:execution_id_status_idx { execution_id, status } }
+        { ::index drop docs:owner_id_metadata_doc_id_idx }
+        { ::index create docs:owner_id_metadata_doc_id_idx { owner_id, metadata, doc_id } }
+        """
+    )