main <- dev

.env.example

@@ -6,6 +6,7 @@ COZO_AUTH_TOKEN=<your_cozo_auth_token>
 TEMPORAL_POSTGRES_PASSWORD=<your_temporal_postgres_password>
 LITELLM_POSTGRES_PASSWORD=<your_litellm_postgres_password>
 LITELLM_MASTER_KEY=<your_litellm_master_key>
+LITELLM_SALT_KEY=<your_litellm_salt_key>
 LITELLM_REDIS_PASSWORD=<your_litellm_redis_password>
 EMBEDDING_SERVICE_BASE=http://text-embeddings-inference-<gpu|cpu>  # Use the 'gpu' profile to run on GPU
 

.gitignore

@@ -1,3 +1,4 @@
+.DS_Store
 *.swp
 ngrok*
 .env*
@@ -6,4 +7,4 @@ ngrok*
 *.pyc
 */node_modules/
 .aider*
-.vscode/
+.vscode/

agents-api/agents_api/activities/task_steps/raise_complete_async.py

@@ -13,7 +13,7 @@ async def raise_complete_async(context: StepContext, output: StepOutcome) -> Non
     # TODO: Create a transtition to "wait" and save the captured_token to the transition
 
     captured_token = activity.info().task_token
-    captured_token = captured_token.decode('latin-1')
+    captured_token = captured_token.decode("latin-1")
     transition_info = CreateTransitionRequest(
         current=context.cursor,
         type="wait",

agents-api/agents_api/activities/task_steps/transition_step.py

@@ -29,4 +29,4 @@ async def transition_step(
 
 transition_step = activity.defn(name="transition_step")(
     transition_step if not testing else mock_transition_step
-)
+)

agents-api/agents_api/clients/litellm.py

@@ -14,9 +14,8 @@
 @wraps(_acompletion)
 @beartype
 async def acompletion(
-    *, model: str, messages: list[dict], **kwargs
+    *, model: str, messages: list[dict], custom_api_key: None | str = None, **kwargs
 ) -> ModelResponse | CustomStreamWrapper:
-    model = f"openai/{model}"  # This is here because litellm proxy expects this format
 
     supported_params = get_supported_openai_params(model)
     settings = {k: v for k, v in kwargs.items() if k in supported_params}
@@ -25,6 +24,6 @@ async def acompletion(
         model=model,
         messages=messages,
         **settings,
-        base_url=litellm_url,
-        api_key=litellm_master_key,
+        base_url=None if custom_api_key else litellm_url,
+        api_key=custom_api_key or litellm_master_key,
     )

agents-api/agents_api/models/agent/create_or_update_agent.py

@@ -80,6 +80,29 @@ def create_or_update_agent(
         }
     )
 
+    # TODO: remove this
+    ### # Create default agent settings
+    ### # Construct a query to insert default settings for the new agent
+    ### default_settings_query = f"""
+    ### %if {{
+    ###     len[count(agent_id)] :=
+    ###         *agent_default_settings{{agent_id}},
+    ###         agent_id = to_uuid($agent_id)
+
+    ###     ?[should_create] := len[count], count > 0
+    ### }}
+    ### %then {{
+    ###     ?[{settings_cols}] <- $settings_vals
+
+    ###     :put agent_default_settings {{
+    ###         {settings_cols}
+    ###     }}
+    ### }}
+    ### """
+
+    # FIXME: This create or update query will overwrite the settings
+    #        Need to find a way to only run the insert query if the agent_default_settings
+
     # Create default agent settings
     # Construct a query to insert default settings for the new agent
     default_settings_query = f"""
@@ -89,6 +112,7 @@ def create_or_update_agent(
             {settings_cols}
         }}
     """
+
     # create the agent
     # Construct a query to insert the new agent record into the agents table
     agent_query = """
@@ -127,7 +151,7 @@ def create_or_update_agent(
 
     queries = [
         verify_developer_id_query(developer_id),
-        default_settings_query if default_settings else None,
+        default_settings_query,
         agent_query,
     ]
 

agents-api/agents_api/models/docs/search_docs_by_embedding.py

@@ -194,11 +194,11 @@ def search_docs_by_embedding(
     """
 
     collect_query = """
-        m[
+        n[
             doc_id,
             owner_type,
             owner_id,
-            collect(snippet),
+            unique(snippet_data),
             distance,
             title,
         ] := 
@@ -209,11 +209,30 @@ def search_docs_by_embedding(
                 snippet_data,
                 distance,
                 title,
-            }, snippet = {
-                "index": snippet_data->0,
-                "content": snippet_data->1,
             }
 
+        m[
+            doc_id,
+            owner_type,
+            owner_id,
+            collect(snippet),
+            distance,
+            title,
+        ] := 
+            n[
+                doc_id,
+                owner_type,
+                owner_id,
+                snippet_data,
+                distance,
+                title,
+            ],
+            snippet = {
+                "index": snippet_datum->0,
+                "content": snippet_datum->1
+            },
+            snippet_datum in snippet_data
+
         ?[
             id,
             owner_type,

agents-api/agents_api/models/docs/search_docs_by_text.py

@@ -1,5 +1,6 @@
 """This module contains functions for searching documents in the CozoDB based on embedding queries."""
 
+import json
 from typing import Any, Literal, TypeVar
 from uuid import UUID
 
@@ -61,6 +62,10 @@ def search_docs_by_text(
         [owner_type, str(owner_id)] for owner_type, owner_id in owners
     ]
 
+    # Need to use NEAR/3($query) to search for arbitrary text within 3 words of each other
+    # See: https://docs.cozodb.org/en/latest/vector.html#full-text-search-fts
+    query = f"NEAR/3({json.dumps(query)})"
+
     # Construct the datalog query for searching document snippets
     search_query = f"""
         owners[owner_type, owner_id] <- $owners

agents-api/agents_api/models/docs/search_docs_hybrid.py

@@ -42,8 +42,6 @@ def dbsf_fuse(
     """
     all_docs = {doc.id: doc for doc in text_results + embedding_results}
 
-    assert all(doc.distance is not None in all_docs for doc in text_results)
-
     text_scores: dict[UUID, float] = {
         doc.id: -(doc.distance or 0.0) for doc in text_results
     }

agents-api/agents_api/models/session/create_or_update_session.py

@@ -26,6 +26,7 @@
 
 @rewrap_exceptions(
     {
+        AssertionError: partialclass(HTTPException, status_code=400),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),

agents-api/agents_api/models/session/create_session.py

@@ -27,6 +27,7 @@
 
 @rewrap_exceptions(
     {
+        AssertionError: partialclass(HTTPException, status_code=400),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),

agents-api/agents_api/models/session/prepare_session_data.py

@@ -127,6 +127,33 @@ def prepare_session_data(
             "instructions":      instructions,
         }
 
+    # Version where we don't have default settings
+    agent_data[collect(record)] :=
+        agents[agent_ids],
+        agent_id in agent_ids,
+        *agents{
+            agent_id,
+            model,
+            name,
+            about,
+            created_at,
+            updated_at,
+            metadata,
+            instructions,
+        },
+        not settings_data[agent_id, _],
+        record = {
+            "id":         agent_id,
+            "name":       name,
+            "model":      model,
+            "about":      about,
+            "created_at": created_at,
+            "updated_at": updated_at,
+            "metadata":   metadata,
+            "default_settings":  {},
+            "instructions":      instructions,
+        }
+
     user_data[collect(record)] :=
         users[user_ids],
         user_id in user_ids,

agents-api/agents_api/routers/agents/create_agent.py

@@ -1,7 +1,7 @@
 from typing import Annotated
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 from starlette.status import HTTP_201_CREATED
 
 import agents_api.models as models
@@ -16,9 +16,10 @@
 
 @router.post("/agents", status_code=HTTP_201_CREATED, tags=["agents"])
 async def create_agent(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     data: CreateAgentRequest,
 ) -> ResourceCreatedResponse:
+    # TODO: Validate model name
     agent = models.agent.create_agent(
         developer_id=x_developer_id,
         data=data,

agents-api/agents_api/routers/agents/create_agent_tool.py

@@ -2,7 +2,7 @@
 from uuid import UUID
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 from starlette.status import HTTP_201_CREATED
 
 import agents_api.models as models
@@ -18,7 +18,7 @@
 @router.post("/agents/{agent_id}/tools", status_code=HTTP_201_CREATED, tags=["agents"])
 async def create_agent_tool(
     agent_id: UUID,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     data: CreateToolRequest,
 ) -> ResourceCreatedResponse:
     tool = models.tools.create_tools(

agents-api/agents_api/routers/agents/create_or_update_agent.py

@@ -2,7 +2,7 @@
 from uuid import UUID
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 from starlette.status import HTTP_201_CREATED
 
 import agents_api.models as models
@@ -19,8 +19,9 @@
 async def create_or_update_agent(
     agent_id: UUID,
     data: CreateOrUpdateAgentRequest,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> ResourceCreatedResponse:
+    # TODO: Validate model name
     agent = models.agent.create_or_update_agent(
         developer_id=x_developer_id,
         agent_id=agent_id,

agents-api/agents_api/routers/agents/delete_agent.py

@@ -1,7 +1,7 @@
 from typing import Annotated
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 from starlette.status import HTTP_202_ACCEPTED
 
 from ...autogen.openapi_model import ResourceDeletedResponse
@@ -12,6 +12,6 @@
 
 @router.delete("/agents/{agent_id}", status_code=HTTP_202_ACCEPTED, tags=["agents"])
 async def delete_agent(
-    agent_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)]
+    agent_id: UUID, x_developer_id: Annotated[UUID, Depends(get_developer_id)]
 ) -> ResourceDeletedResponse:
     return delete_agent_query(developer_id=x_developer_id, agent_id=agent_id)

agents-api/agents_api/routers/agents/delete_agent_tool.py

@@ -2,7 +2,7 @@
 from uuid import UUID
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 
 from ...autogen.openapi_model import ResourceDeletedResponse
 from ...dependencies.developer_id import get_developer_id
@@ -14,7 +14,7 @@
 async def delete_agent_tool(
     agent_id: UUID,
     tool_id: UUID,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> ResourceDeletedResponse:
     return delete_tool(
         developer_id=x_developer_id,

agents-api/agents_api/routers/agents/get_agent_details.py

@@ -1,7 +1,7 @@
 from typing import Annotated
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 
 from ...autogen.openapi_model import Agent
 from ...dependencies.developer_id import get_developer_id
@@ -11,7 +11,7 @@
 
 @router.get("/agents/{agent_id}", tags=["agents"])
 async def get_agent_details(
-    agent_id: UUID4,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    agent_id: UUID,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> Agent:
     return get_agent_query(developer_id=x_developer_id, agent_id=agent_id)

agents-api/agents_api/routers/agents/list_agent_tools.py

@@ -2,7 +2,7 @@
 from uuid import UUID
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 
 from ...autogen.openapi_model import ListResponse, Tool
 from ...dependencies.developer_id import get_developer_id
@@ -12,7 +12,7 @@
 
 @router.get("/agents/{agent_id}/tools", tags=["agents"])
 async def list_agent_tools(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     agent_id: UUID,
     limit: int = 100,
     offset: int = 0,

agents-api/agents_api/routers/agents/list_agents.py

@@ -3,7 +3,7 @@
 from typing import Annotated, Literal
 
 from fastapi import Depends, HTTPException, status
-from pydantic import UUID4
+from uuid import UUID
 
 from ...autogen.openapi_model import Agent, ListResponse
 from ...dependencies.developer_id import get_developer_id
@@ -13,7 +13,7 @@
 
 @router.get("/agents", tags=["agents"])
 async def list_agents(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     limit: int = 100,
     offset: int = 0,
     sort_by: Literal["created_at", "updated_at"] = "created_at",

agents-api/agents_api/routers/agents/patch_agent.py

@@ -1,7 +1,7 @@
 from typing import Annotated
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 from starlette.status import HTTP_200_OK
 
 from ...autogen.openapi_model import PatchAgentRequest, ResourceUpdatedResponse
@@ -17,8 +17,8 @@
     tags=["agents"],
 )
 async def patch_agent(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    agent_id: UUID4,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
+    agent_id: UUID,
     data: PatchAgentRequest,
 ) -> ResourceUpdatedResponse:
     return patch_agent_query(