Merge pull request #958 from julep-ai/x/misc-fixes

fix(agents-api): Miscellaneous fixes in `sessions`
julep-ai · Dec 14, 2024 · caa0b27 · caa0b27
2 parents 6e5de98 + 756a547
commit caa0b27
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 7 deletions.
diff --git a/agents-api/agents_api/models/chat/gather_messages.py b/agents-api/agents_api/models/chat/gather_messages.py
@@ -193,6 +193,12 @@ async def gather_messages(
         # MMR is not applied to text search
         and recall_options.mode != "text"
     ):
+        # FIXME: This is a temporary fix to ensure that the MMR algorithm works.
+        # We shouldn't be having references without embeddings.
+        doc_references = [
+            doc for doc in doc_references if doc.snippet.embedding is not None
+        ]
+
         # Apply MMR
         indices = maximal_marginal_relevance(
             np.asarray(query_embedding),

diff --git a/agents-api/agents_api/models/session/create_session.py b/agents-api/agents_api/models/session/create_session.py
@@ -104,14 +104,15 @@ def create_session(
 
     create_query = """
         # Insert the new session data into the 'session' table with the specified columns.
-        ?[session_id, developer_id, situation, metadata, render_templates, token_budget, context_overflow] <- [[
+        ?[session_id, developer_id, situation, metadata, render_templates, token_budget, context_overflow, recall_options] <- [[
             $session_id,
             $developer_id,
             $situation,
             $metadata,
             $render_templates,
             $token_budget,
             $context_overflow,
+            $recall_options,
         ]]
 
         :insert sessions {
@@ -122,6 +123,7 @@ def create_session(
             render_templates,
             token_budget,
             context_overflow,
+            recall_options,
         }
         # Specify the data to return after the query execution, typically the newly created session's ID.
         :returning

diff --git a/agents-api/agents_api/models/utils.py b/agents-api/agents_api/models/utils.py
@@ -94,6 +94,7 @@ def mark_session_updated_query(developer_id: UUID | str, session_id: UUID | str)
         token_budget,
         context_overflow,
         updated_at,
+        recall_options,
     ] :=
         input[developer_id, session_id],
         *sessions {{
@@ -105,6 +106,7 @@ def mark_session_updated_query(developer_id: UUID | str, session_id: UUID | str)
             render_templates,
             token_budget,
             context_overflow,
+            recall_options,
             @ 'END'
         }},
         updated_at = [floor(now()), true]
@@ -119,6 +121,7 @@ def mark_session_updated_query(developer_id: UUID | str, session_id: UUID | str)
         render_templates,
         token_budget,
         context_overflow,
+        recall_options,
         updated_at,
     }}
     """

diff --git a/integrations-service/integrations/models/spider.py b/integrations-service/integrations/models/spider.py
@@ -6,7 +6,7 @@
 
 
 class SpiderResponse(BaseModel):
-    content: Optional[str] = None
+    content: Optional[str | list[str]] = None
     error: Optional[str] = None
     status: Optional[int] = None
     costs: Optional[dict[Any, Any]] = None

diff --git a/integrations-service/integrations/utils/integrations/spider.py b/integrations-service/integrations/utils/integrations/spider.py
@@ -28,11 +28,11 @@ def get_api_key(setup: SpiderSetup) -> str:
 def create_spider_response(pages: list[dict]) -> list[SpiderResponse]:
     return [
         SpiderResponse(
-            url=page.get("url", "Not available"),
-            content=(page.get("content") or "Not available").strip(),
-            error=page.get("error", "Not available"),
-            status=page.get("status", "Not available"),
-            costs=page.get("costs", "Not available"),
+            url=page.get("url"),
+            content=(page.get("content")),
+            error=page.get("error"),
+            status=page.get("status"),
+            costs=page.get("costs"),
         )
         for page in pages
     ]