diff --git a/.DS_Store b/.DS_Store
deleted file mode 100644
index 7211b8ec9..000000000
Binary files a/.DS_Store and /dev/null differ
diff --git a/.env.example b/.env.example
index fc2bdd7fb..e3cda1a97 100644
--- a/.env.example
+++ b/.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
 
diff --git a/.gitignore b/.gitignore
index b68b9ff1e..7d1d2ca75 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+.DS_Store
 *.swp
 ngrok*
 .env*
@@ -6,4 +7,4 @@ ngrok*
 *.pyc
 */node_modules/
 .aider*
-.vscode/
\ No newline at end of file
+.vscode/
diff --git a/agents-api/agents_api/activities/task_steps/raise_complete_async.py b/agents-api/agents_api/activities/task_steps/raise_complete_async.py
index 0748f08b3..3944c914f 100644
--- a/agents-api/agents_api/activities/task_steps/raise_complete_async.py
+++ b/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",
diff --git a/agents-api/agents_api/activities/task_steps/transition_step.py b/agents-api/agents_api/activities/task_steps/transition_step.py
index cd3c09724..04ce6e8fd 100644
--- a/agents-api/agents_api/activities/task_steps/transition_step.py
+++ b/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
-)
\ No newline at end of file
+)
diff --git a/agents-api/agents_api/clients/litellm.py b/agents-api/agents_api/clients/litellm.py
index 463f6fa67..9b321cc29 100644
--- a/agents-api/agents_api/clients/litellm.py
+++ b/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,
     )
diff --git a/agents-api/agents_api/models/agent/create_or_update_agent.py b/agents-api/agents_api/models/agent/create_or_update_agent.py
index 64b008d44..156c8ae61 100644
--- a/agents-api/agents_api/models/agent/create_or_update_agent.py
+++ b/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,
     ]
 
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
index acebd09cd..d6bed97fc 100644
--- a/agents-api/agents_api/models/docs/search_docs_by_embedding.py
+++ b/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,
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
index 0662aa84d..eeae8362c 100644
--- a/agents-api/agents_api/models/docs/search_docs_by_text.py
+++ b/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
diff --git a/agents-api/agents_api/models/docs/search_docs_hybrid.py b/agents-api/agents_api/models/docs/search_docs_hybrid.py
index 03fb44037..598600511 100644
--- a/agents-api/agents_api/models/docs/search_docs_hybrid.py
+++ b/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
     }
diff --git a/agents-api/agents_api/models/session/create_or_update_session.py b/agents-api/agents_api/models/session/create_or_update_session.py
index 60c0b7f71..6c81b2563 100644
--- a/agents-api/agents_api/models/session/create_or_update_session.py
+++ b/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),
diff --git a/agents-api/agents_api/models/session/create_session.py b/agents-api/agents_api/models/session/create_session.py
index a83837ffd..32a74533a 100644
--- a/agents-api/agents_api/models/session/create_session.py
+++ b/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),
diff --git a/agents-api/agents_api/models/session/prepare_session_data.py b/agents-api/agents_api/models/session/prepare_session_data.py
index 49cdc8cbc..e8cfb7fc7 100644
--- a/agents-api/agents_api/models/session/prepare_session_data.py
+++ b/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,
diff --git a/agents-api/agents_api/routers/agents/create_agent.py b/agents-api/agents_api/routers/agents/create_agent.py
index d1cac0d6b..650b0b531 100644
--- a/agents-api/agents_api/routers/agents/create_agent.py
+++ b/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,
diff --git a/agents-api/agents_api/routers/agents/create_agent_tool.py b/agents-api/agents_api/routers/agents/create_agent_tool.py
index dc934bee7..783dec3c9 100644
--- a/agents-api/agents_api/routers/agents/create_agent_tool.py
+++ b/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(
diff --git a/agents-api/agents_api/routers/agents/create_or_update_agent.py b/agents-api/agents_api/routers/agents/create_or_update_agent.py
index fc2fa5563..ef76ae627 100644
--- a/agents-api/agents_api/routers/agents/create_or_update_agent.py
+++ b/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,
diff --git a/agents-api/agents_api/routers/agents/delete_agent.py b/agents-api/agents_api/routers/agents/delete_agent.py
index 4603c2217..2bacdeb45 100644
--- a/agents-api/agents_api/routers/agents/delete_agent.py
+++ b/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)
diff --git a/agents-api/agents_api/routers/agents/delete_agent_tool.py b/agents-api/agents_api/routers/agents/delete_agent_tool.py
index c220a2cf7..1dd6da83e 100644
--- a/agents-api/agents_api/routers/agents/delete_agent_tool.py
+++ b/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,
diff --git a/agents-api/agents_api/routers/agents/get_agent_details.py b/agents-api/agents_api/routers/agents/get_agent_details.py
index 04511527a..a714c0b2c 100644
--- a/agents-api/agents_api/routers/agents/get_agent_details.py
+++ b/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)
diff --git a/agents-api/agents_api/routers/agents/list_agent_tools.py b/agents-api/agents_api/routers/agents/list_agent_tools.py
index ebc10b0ef..f8af9da9a 100644
--- a/agents-api/agents_api/routers/agents/list_agent_tools.py
+++ b/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,
diff --git a/agents-api/agents_api/routers/agents/list_agents.py b/agents-api/agents_api/routers/agents/list_agents.py
index f64d4bdbf..d3f014063 100644
--- a/agents-api/agents_api/routers/agents/list_agents.py
+++ b/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",
diff --git a/agents-api/agents_api/routers/agents/patch_agent.py b/agents-api/agents_api/routers/agents/patch_agent.py
index 29596447e..fa96a4169 100644
--- a/agents-api/agents_api/routers/agents/patch_agent.py
+++ b/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(
diff --git a/agents-api/agents_api/routers/agents/patch_agent_tool.py b/agents-api/agents_api/routers/agents/patch_agent_tool.py
index 843fa91eb..e979ad30f 100644
--- a/agents-api/agents_api/routers/agents/patch_agent_tool.py
+++ b/agents-api/agents_api/routers/agents/patch_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 (
     PatchToolRequest,
@@ -15,7 +15,7 @@
 
 @router.patch("/agents/{agent_id}/tools/{tool_id}", tags=["agents"])
 async def patch_agent_tool(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     agent_id: UUID,
     tool_id: UUID,
     data: PatchToolRequest,
diff --git a/agents-api/agents_api/routers/agents/update_agent.py b/agents-api/agents_api/routers/agents/update_agent.py
index 865954a6d..e51c459c7 100644
--- a/agents-api/agents_api/routers/agents/update_agent.py
+++ b/agents-api/agents_api/routers/agents/update_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 ResourceUpdatedResponse, UpdateAgentRequest
@@ -17,8 +17,8 @@
     tags=["agents"],
 )
 async def update_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: UpdateAgentRequest,
 ) -> ResourceUpdatedResponse:
     return update_agent_query(
diff --git a/agents-api/agents_api/routers/agents/update_agent_tool.py b/agents-api/agents_api/routers/agents/update_agent_tool.py
index 60e38ad76..b0633b73b 100644
--- a/agents-api/agents_api/routers/agents/update_agent_tool.py
+++ b/agents-api/agents_api/routers/agents/update_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 (
     ResourceUpdatedResponse,
@@ -15,7 +15,7 @@
 
 @router.put("/agents/{agent_id}/tools/{tool_id}", tags=["agents"])
 async def update_agent_tool(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     agent_id: UUID,
     tool_id: UUID,
     data: UpdateToolRequest,
diff --git a/agents-api/agents_api/routers/docs/create_doc.py b/agents-api/agents_api/routers/docs/create_doc.py
index dd41620ae..e71eed59f 100644
--- a/agents-api/agents_api/routers/docs/create_doc.py
+++ b/agents-api/agents_api/routers/docs/create_doc.py
@@ -2,7 +2,7 @@
 from uuid import UUID, uuid4
 
 from fastapi import BackgroundTasks, Depends
-from pydantic import UUID4
+from uuid import UUID
 from starlette.status import HTTP_201_CREATED
 from temporalio.client import Client as TemporalClient
 
@@ -54,9 +54,9 @@ async def run_embed_docs_task(
 
 @router.post("/users/{user_id}/docs", status_code=HTTP_201_CREATED, tags=["docs"])
 async def create_user_doc(
-    user_id: UUID4,
+    user_id: UUID,
     data: CreateDocRequest,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     background_tasks: BackgroundTasks,
 ) -> ResourceCreatedResponse:
     doc = create_doc_query(
@@ -84,9 +84,9 @@ async def create_user_doc(
 
 @router.post("/agents/{agent_id}/docs", status_code=HTTP_201_CREATED, tags=["docs"])
 async def create_agent_doc(
-    agent_id: UUID4,
+    agent_id: UUID,
     data: CreateDocRequest,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     background_tasks: BackgroundTasks,
 ) -> ResourceCreatedResponse:
     doc = create_doc_query(
diff --git a/agents-api/agents_api/routers/docs/delete_doc.py b/agents-api/agents_api/routers/docs/delete_doc.py
index c31bf4051..bfd6a4799 100644
--- a/agents-api/agents_api/routers/docs/delete_doc.py
+++ b/agents-api/agents_api/routers/docs/delete_doc.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
@@ -14,9 +14,9 @@
     "/agents/{agent_id}/docs/{doc_id}", status_code=HTTP_202_ACCEPTED, tags=["docs"]
 )
 async def delete_agent_doc(
-    doc_id: UUID4,
-    agent_id: UUID4,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    doc_id: UUID,
+    agent_id: UUID,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> ResourceDeletedResponse:
     return delete_doc_query(
         developer_id=x_developer_id,
@@ -30,9 +30,9 @@ async def delete_agent_doc(
     "/users/{user_id}/docs/{doc_id}", status_code=HTTP_202_ACCEPTED, tags=["docs"]
 )
 async def delete_user_doc(
-    doc_id: UUID4,
-    user_id: UUID4,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    doc_id: UUID,
+    user_id: UUID,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> ResourceDeletedResponse:
     return delete_doc_query(
         developer_id=x_developer_id,
diff --git a/agents-api/agents_api/routers/docs/embed.py b/agents-api/agents_api/routers/docs/embed.py
index 1de99bfce..2c6b7b641 100644
--- a/agents-api/agents_api/routers/docs/embed.py
+++ b/agents-api/agents_api/routers/docs/embed.py
@@ -1,7 +1,7 @@
 from typing import Annotated
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 
 import agents_api.clients.embed as embedder
 
@@ -15,7 +15,7 @@
 
 @router.post("/embed", tags=["docs"])
 async def embed(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     data: EmbedQueryRequest,
 ) -> EmbedQueryResponse:
     text_to_embed: str | list[str] = data.text
diff --git a/agents-api/agents_api/routers/docs/get_doc.py b/agents-api/agents_api/routers/docs/get_doc.py
index febebf1bd..9a7b23ccf 100644
--- a/agents-api/agents_api/routers/docs/get_doc.py
+++ b/agents-api/agents_api/routers/docs/get_doc.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 Doc
 from ...dependencies.developer_id import get_developer_id
@@ -11,7 +11,7 @@
 
 @router.get("/docs/{doc_id}", tags=["docs"])
 async def get_doc(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    doc_id: UUID4,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
+    doc_id: UUID,
 ) -> Doc:
     return get_doc_query(developer_id=x_developer_id, doc_id=doc_id)
diff --git a/agents-api/agents_api/routers/docs/list_docs.py b/agents-api/agents_api/routers/docs/list_docs.py
index 80a6ba6ae..fa0931f4b 100644
--- a/agents-api/agents_api/routers/docs/list_docs.py
+++ b/agents-api/agents_api/routers/docs/list_docs.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 Doc, ListResponse
 from ...dependencies.developer_id import get_developer_id
@@ -13,8 +13,8 @@
 
 @router.get("/users/{user_id}/docs", tags=["docs"])
 async def list_user_docs(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    user_id: UUID4,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
+    user_id: UUID,
     limit: int = 100,
     offset: int = 0,
     sort_by: Literal["created_at", "updated_at"] = "created_at",
@@ -45,8 +45,8 @@ async def list_user_docs(
 
 @router.get("/agents/{agent_id}/docs", tags=["docs"])
 async def list_agent_docs(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    agent_id: UUID4,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
+    agent_id: UUID,
     limit: int = 100,
     offset: int = 0,
     sort_by: Literal["created_at", "updated_at"] = "created_at",
diff --git a/agents-api/agents_api/routers/docs/search_docs.py b/agents-api/agents_api/routers/docs/search_docs.py
index f2864164e..83acf1ed4 100644
--- a/agents-api/agents_api/routers/docs/search_docs.py
+++ b/agents-api/agents_api/routers/docs/search_docs.py
@@ -2,7 +2,7 @@
 from typing import Annotated, Any, Dict, List, Optional, Tuple, Union
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 
 from ...autogen.openapi_model import (
     DocSearchResponse,
@@ -63,11 +63,11 @@ def get_search_fn_and_params(
 
 @router.post("/users/{user_id}/search", tags=["docs"])
 async def search_user_docs(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     search_params: (
         TextOnlyDocSearchRequest | VectorDocSearchRequest | HybridDocSearchRequest
     ),
-    user_id: UUID4,
+    user_id: UUID,
 ) -> DocSearchResponse:
     search_fn, params = get_search_fn_and_params(search_params)
 
@@ -90,11 +90,11 @@ async def search_user_docs(
 
 @router.post("/agents/{agent_id}/search", tags=["docs"])
 async def search_agent_docs(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     search_params: (
         TextOnlyDocSearchRequest | VectorDocSearchRequest | HybridDocSearchRequest
     ),
-    agent_id: UUID4,
+    agent_id: UUID,
 ) -> DocSearchResponse:
     search_fn, params = get_search_fn_and_params(search_params)
 
diff --git a/agents-api/agents_api/routers/jobs/routers.py b/agents-api/agents_api/routers/jobs/routers.py
index c9783053c..8095af86e 100644
--- a/agents-api/agents_api/routers/jobs/routers.py
+++ b/agents-api/agents_api/routers/jobs/routers.py
@@ -1,7 +1,7 @@
 from typing import Literal
 
 from fastapi import APIRouter
-from pydantic import UUID4
+from uuid import UUID
 from temporalio.client import WorkflowExecutionStatus
 
 from agents_api.autogen.openapi_model import JobStatus
@@ -41,7 +41,7 @@ def map_job_status(
 
 
 @router.get("/jobs/{job_id}", tags=["jobs"])
-async def get_job_status(job_id: UUID4) -> JobStatus:
+async def get_job_status(job_id: UUID) -> JobStatus:
     client = await get_client()
     handle = client.get_workflow_handle(
         workflow_id=str(job_id),
diff --git a/agents-api/agents_api/routers/sessions/chat.py b/agents-api/agents_api/routers/sessions/chat.py
index b7bf96d3a..27199c58d 100644
--- a/agents-api/agents_api/routers/sessions/chat.py
+++ b/agents-api/agents_api/routers/sessions/chat.py
@@ -1,7 +1,7 @@
-from typing import Annotated
+from typing import Annotated, Optional
 from uuid import UUID, uuid4
 
-from fastapi import BackgroundTasks, Depends
+from fastapi import BackgroundTasks, Depends, Header
 from starlette.status import HTTP_201_CREATED
 
 from ...autogen.openapi_model import (
@@ -33,6 +33,7 @@ async def chat(
     session_id: UUID,
     chat_input: ChatInput,
     background_tasks: BackgroundTasks,
+    x_custom_api_key: Optional[str] = Header(None, alias="X-Custom-Api-Key"),
 ) -> ChatResponse:
     if chat_input.stream:
         raise NotImplementedError("Streaming is not yet implemented")
@@ -58,7 +59,13 @@ async def chat(
 
     # Prepare the environment
     env: dict = chat_context.get_chat_environment()
-    env["docs"] = doc_references
+    env["docs"] = [
+        dict(
+            title=ref.title,
+            content=[snippet.content for snippet in ref.snippets],
+        )
+        for ref in doc_references
+    ]
 
     # Render the system message
     if situation := chat_context.session.situation:
@@ -107,6 +114,7 @@ async def chat(
         tools=tools or None,
         user=str(developer.id),  # For tracking usage
         tags=developer.tags,  # For filtering models in litellm
+        custom_api_key=x_custom_api_key,
         **settings,
     )
 
@@ -120,6 +128,7 @@ async def chat(
         ]
 
         # Add the response to the new entries
+        # FIXME: We need to save all the choices
         new_entries.append(
             CreateEntryRequest.from_model_input(
                 model=settings["model"],
diff --git a/agents-api/agents_api/routers/sessions/create_or_update_session.py b/agents-api/agents_api/routers/sessions/create_or_update_session.py
index 8ed9ff7ba..99f98652b 100644
--- a/agents-api/agents_api/routers/sessions/create_or_update_session.py
+++ b/agents-api/agents_api/routers/sessions/create_or_update_session.py
@@ -2,12 +2,12 @@
 from uuid import UUID
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 from starlette.status import HTTP_201_CREATED
 
 from ...autogen.openapi_model import (
     CreateOrUpdateSessionRequest,
-    ResourceCreatedResponse,
+    ResourceUpdatedResponse,
 )
 from ...dependencies.developer_id import get_developer_id
 from ...models.session.create_or_update_session import (
@@ -18,17 +18,14 @@
 
 @router.post("/sessions/{session_id}", status_code=HTTP_201_CREATED, tags=["sessions"])
 async def create_or_update_session(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     session_id: UUID,
     data: CreateOrUpdateSessionRequest,
-) -> ResourceCreatedResponse:
-    session = create_session_query(
+) -> ResourceUpdatedResponse:
+    session_updated = create_session_query(
         developer_id=x_developer_id,
         session_id=session_id,
         data=data,
     )
 
-    return ResourceCreatedResponse(
-        id=session.id,
-        created_at=session.created_at,
-    )
+    return session_updated
diff --git a/agents-api/agents_api/routers/sessions/create_session.py b/agents-api/agents_api/routers/sessions/create_session.py
index 2262d8e7c..f65948991 100644
--- a/agents-api/agents_api/routers/sessions/create_session.py
+++ b/agents-api/agents_api/routers/sessions/create_session.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
 
 from ...autogen.openapi_model import (
@@ -15,7 +15,7 @@
 
 @router.post("/sessions", status_code=HTTP_201_CREATED, tags=["sessions"])
 async def create_session(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     data: CreateSessionRequest,
 ) -> ResourceCreatedResponse:
     session = create_session_query(
diff --git a/agents-api/agents_api/routers/sessions/delete_session.py b/agents-api/agents_api/routers/sessions/delete_session.py
index 9645eb8de..310516af9 100644
--- a/agents-api/agents_api/routers/sessions/delete_session.py
+++ b/agents-api/agents_api/routers/sessions/delete_session.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
@@ -14,6 +14,6 @@
     "/sessions/{session_id}", status_code=HTTP_202_ACCEPTED, tags=["sessions"]
 )
 async def delete_session(
-    session_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)]
+    session_id: UUID, x_developer_id: Annotated[UUID, Depends(get_developer_id)]
 ) -> ResourceDeletedResponse:
     return delete_session_query(developer_id=x_developer_id, session_id=session_id)
diff --git a/agents-api/agents_api/routers/sessions/get_session.py b/agents-api/agents_api/routers/sessions/get_session.py
index c357394fa..141a432e7 100644
--- a/agents-api/agents_api/routers/sessions/get_session.py
+++ b/agents-api/agents_api/routers/sessions/get_session.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 Session
 from ...dependencies.developer_id import get_developer_id
@@ -11,6 +11,6 @@
 
 @router.get("/sessions/{session_id}", tags=["sessions"])
 async def get_session(
-    session_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)]
+    session_id: UUID, x_developer_id: Annotated[UUID, Depends(get_developer_id)]
 ) -> Session:
     return get_session_query(developer_id=x_developer_id, session_id=session_id)
diff --git a/agents-api/agents_api/routers/sessions/get_session_history.py b/agents-api/agents_api/routers/sessions/get_session_history.py
index 64be57fcb..d1a5d1f80 100644
--- a/agents-api/agents_api/routers/sessions/get_session_history.py
+++ b/agents-api/agents_api/routers/sessions/get_session_history.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 History
 from ...dependencies.developer_id import get_developer_id
@@ -11,6 +11,6 @@
 
 @router.get("/sessions/{session_id}/history", tags=["sessions"])
 async def get_session_history(
-    session_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)]
+    session_id: UUID, x_developer_id: Annotated[UUID, Depends(get_developer_id)]
 ) -> History:
     return get_history_query(developer_id=x_developer_id, session_id=session_id)
diff --git a/agents-api/agents_api/routers/sessions/list_sessions.py b/agents-api/agents_api/routers/sessions/list_sessions.py
index bf5458887..8a7fa0884 100644
--- a/agents-api/agents_api/routers/sessions/list_sessions.py
+++ b/agents-api/agents_api/routers/sessions/list_sessions.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 ListResponse, Session
 from ...dependencies.developer_id import get_developer_id
@@ -13,7 +13,7 @@
 
 @router.get("/sessions", tags=["sessions"])
 async def list_sessions(
-    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",
diff --git a/agents-api/agents_api/routers/sessions/patch_session.py b/agents-api/agents_api/routers/sessions/patch_session.py
index 365fa49ca..522ecf57c 100644
--- a/agents-api/agents_api/routers/sessions/patch_session.py
+++ b/agents-api/agents_api/routers/sessions/patch_session.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 (
     PatchSessionRequest,
@@ -14,8 +14,8 @@
 
 @router.patch("/sessions/{session_id}", tags=["sessions"])
 async def patch_session(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    session_id: UUID4,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
+    session_id: UUID,
     data: PatchSessionRequest,
 ) -> ResourceUpdatedResponse:
     return patch_session_query(
diff --git a/agents-api/agents_api/routers/sessions/update_session.py b/agents-api/agents_api/routers/sessions/update_session.py
index 8d7c28b33..14d09ead7 100644
--- a/agents-api/agents_api/routers/sessions/update_session.py
+++ b/agents-api/agents_api/routers/sessions/update_session.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 (
     ResourceUpdatedResponse,
@@ -14,8 +14,8 @@
 
 @router.put("/sessions/{session_id}", tags=["sessions"])
 async def update_session(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    session_id: UUID4,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
+    session_id: UUID,
     data: UpdateSessionRequest,
 ) -> ResourceUpdatedResponse:
     return update_session_query(
diff --git a/agents-api/agents_api/routers/tasks/create_or_update_task.py b/agents-api/agents_api/routers/tasks/create_or_update_task.py
index 5d424eb5b..23cae0f91 100644
--- a/agents-api/agents_api/routers/tasks/create_or_update_task.py
+++ b/agents-api/agents_api/routers/tasks/create_or_update_task.py
@@ -3,7 +3,7 @@
 from fastapi import Depends, HTTPException
 from jsonschema import validate
 from jsonschema.exceptions import SchemaError, ValidationError
-from pydantic import UUID4
+from uuid import UUID
 from starlette.status import HTTP_201_CREATED
 
 from agents_api.autogen.openapi_model import (
@@ -23,9 +23,9 @@
 )
 async def create_or_update_task(
     data: CreateOrUpdateTaskRequest,
-    agent_id: UUID4,
-    task_id: UUID4,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    agent_id: UUID,
+    task_id: UUID,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> ResourceUpdatedResponse:
     # TODO: Do thorough validation of the task spec
     # SCRUM-10
diff --git a/agents-api/agents_api/routers/tasks/create_task.py b/agents-api/agents_api/routers/tasks/create_task.py
index b124c3dee..d6b4bc70b 100644
--- a/agents-api/agents_api/routers/tasks/create_task.py
+++ b/agents-api/agents_api/routers/tasks/create_task.py
@@ -3,7 +3,7 @@
 from fastapi import Depends, HTTPException
 from jsonschema import validate
 from jsonschema.exceptions import SchemaError, ValidationError
-from pydantic import UUID4
+from uuid import UUID
 from starlette.status import HTTP_201_CREATED
 
 from ...autogen.openapi_model import (
@@ -18,8 +18,8 @@
 @router.post("/agents/{agent_id}/tasks", status_code=HTTP_201_CREATED, tags=["tasks"])
 async def create_task(
     data: CreateTaskRequest,
-    agent_id: UUID4,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    agent_id: UUID,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> ResourceCreatedResponse:
     # TODO: Do thorough validation of the task spec
     # SCRUM-10
diff --git a/agents-api/agents_api/routers/tasks/create_task_execution.py b/agents-api/agents_api/routers/tasks/create_task_execution.py
index d3d22750e..3d77f80c6 100644
--- a/agents-api/agents_api/routers/tasks/create_task_execution.py
+++ b/agents-api/agents_api/routers/tasks/create_task_execution.py
@@ -7,7 +7,7 @@
 from jsonschema import validate
 from jsonschema.exceptions import ValidationError
 from pycozo.client import QueryException
-from pydantic import UUID4
+from uuid import UUID
 from starlette.status import HTTP_201_CREATED
 from temporalio.client import WorkflowHandle
 
@@ -92,9 +92,9 @@ async def start_execution(
     tags=["executions"],
 )
 async def create_task_execution(
-    task_id: UUID4,
+    task_id: UUID,
     data: CreateExecutionRequest,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     background_tasks: BackgroundTasks,
 ) -> ResourceCreatedResponse:
     try:
diff --git a/agents-api/agents_api/routers/tasks/get_execution_details.py b/agents-api/agents_api/routers/tasks/get_execution_details.py
index 49f6ae2d2..8da7b98c8 100644
--- a/agents-api/agents_api/routers/tasks/get_execution_details.py
+++ b/agents-api/agents_api/routers/tasks/get_execution_details.py
@@ -1,4 +1,4 @@
-from pydantic import UUID4
+from uuid import UUID
 
 from agents_api.autogen.openapi_model import (
     Execution,
@@ -11,5 +11,5 @@
 
 
 @router.get("/executions/{execution_id}", tags=["executions"])
-async def get_execution_details(execution_id: UUID4) -> Execution:
+async def get_execution_details(execution_id: UUID) -> Execution:
     return get_execution_query(execution_id=execution_id)
diff --git a/agents-api/agents_api/routers/tasks/get_task_details.py b/agents-api/agents_api/routers/tasks/get_task_details.py
index 2121397e7..836a4d640 100644
--- a/agents-api/agents_api/routers/tasks/get_task_details.py
+++ b/agents-api/agents_api/routers/tasks/get_task_details.py
@@ -2,7 +2,7 @@
 
 from fastapi import Depends, HTTPException, status
 from pycozo.client import QueryException
-from pydantic import UUID4
+from uuid import UUID
 
 from ...autogen.openapi_model import (
     Task,
@@ -14,8 +14,8 @@
 
 @router.get("/tasks/{task_id}", tags=["tasks"])
 async def get_task_details(
-    task_id: UUID4,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    task_id: UUID,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> Task:
     not_found = HTTPException(
         status_code=status.HTTP_404_NOT_FOUND, detail="Task not found"
diff --git a/agents-api/agents_api/routers/tasks/list_execution_transitions.py b/agents-api/agents_api/routers/tasks/list_execution_transitions.py
index 5919d152b..9f06fdcf1 100644
--- a/agents-api/agents_api/routers/tasks/list_execution_transitions.py
+++ b/agents-api/agents_api/routers/tasks/list_execution_transitions.py
@@ -1,6 +1,6 @@
 from typing import Literal
 
-from pydantic import UUID4
+from uuid import UUID
 
 from agents_api.autogen.openapi_model import (
     ListResponse,
@@ -15,7 +15,7 @@
 
 @router.get("/executions/{execution_id}/transitions", tags=["executions"])
 async def list_execution_transitions(
-    execution_id: UUID4,
+    execution_id: UUID,
     limit: int = 100,
     offset: int = 0,
     sort_by: Literal["created_at", "updated_at"] = "created_at",
@@ -35,8 +35,8 @@ async def list_execution_transitions(
 # TODO: Do we need this?
 # @router.get("/executions/{execution_id}/transitions/{transition_id}", tags=["tasks"])
 # async def get_execution_transition(
-#     execution_id: UUID4,
-#     transition_id: UUID4,
+#     execution_id: UUID,
+#     transition_id: UUID,
 # ) -> Transition:
 #     try:
 #         res = [
diff --git a/agents-api/agents_api/routers/tasks/list_task_executions.py b/agents-api/agents_api/routers/tasks/list_task_executions.py
index 27ea782c4..6dc2635c9 100644
--- a/agents-api/agents_api/routers/tasks/list_task_executions.py
+++ b/agents-api/agents_api/routers/tasks/list_task_executions.py
@@ -1,7 +1,7 @@
 from typing import Annotated, Literal
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 
 from agents_api.autogen.openapi_model import (
     Execution,
@@ -17,8 +17,8 @@
 
 @router.get("/tasks/{task_id}/executions", tags=["tasks"])
 async def list_task_executions(
-    task_id: UUID4,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    task_id: UUID,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     limit: int = 100,
     offset: int = 0,
     sort_by: Literal["created_at", "updated_at"] = "created_at",
diff --git a/agents-api/agents_api/routers/tasks/list_tasks.py b/agents-api/agents_api/routers/tasks/list_tasks.py
index 5066fcc96..6656738bd 100644
--- a/agents-api/agents_api/routers/tasks/list_tasks.py
+++ b/agents-api/agents_api/routers/tasks/list_tasks.py
@@ -1,7 +1,7 @@
 from typing import Annotated, Literal
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 
 from agents_api.autogen.openapi_model import (
     ListResponse,
@@ -15,8 +15,8 @@
 
 @router.get("/agents/{agent_id}/tasks", tags=["tasks"])
 async def list_tasks(
-    agent_id: UUID4,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    agent_id: UUID,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     limit: int = 100,
     offset: int = 0,
     sort_by: Literal["created_at", "updated_at"] = "created_at",
diff --git a/agents-api/agents_api/routers/tasks/patch_execution.py b/agents-api/agents_api/routers/tasks/patch_execution.py
index 795b1b212..d0a8ee779 100644
--- a/agents-api/agents_api/routers/tasks/patch_execution.py
+++ b/agents-api/agents_api/routers/tasks/patch_execution.py
@@ -1,7 +1,7 @@
 from typing import Annotated
 
 from fastapi import Depends
-from pydantic import UUID4
+from uuid import UUID
 
 from agents_api.autogen.openapi_model import (
     ResourceUpdatedResponse,
@@ -17,9 +17,9 @@
 
 @router.patch("/tasks/{task_id}/executions/{execution_id}", tags=["tasks"])
 async def patch_execution(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    task_id: UUID4,
-    execution_id: UUID4,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
+    task_id: UUID,
+    execution_id: UUID,
     data: UpdateExecutionRequest,
 ) -> ResourceUpdatedResponse:
     return update_execution_query(
diff --git a/agents-api/agents_api/routers/tasks/stream_transitions_events.py b/agents-api/agents_api/routers/tasks/stream_transitions_events.py
index 36e95aa2b..8202881f0 100644
--- a/agents-api/agents_api/routers/tasks/stream_transitions_events.py
+++ b/agents-api/agents_api/routers/tasks/stream_transitions_events.py
@@ -6,7 +6,7 @@
 import anyio
 from anyio.streams.memory import MemoryObjectSendStream
 from fastapi import Depends, Query
-from pydantic import UUID4
+from uuid import UUID
 from sse_starlette.sse import EventSourceResponse
 from starlette.requests import Request
 from temporalio.api.enums.v1 import EventType
@@ -81,8 +81,8 @@ async def event_publisher(
 
 @router.get("/executions/{execution_id}/transitions.stream", tags=["executions"])
 async def stream_transitions_events(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    execution_id: UUID4,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
+    execution_id: UUID,
     req: Request,
     next_page_token: Annotated[str | None, Query()] = None,
 ):
diff --git a/agents-api/agents_api/routers/tasks/update_execution.py b/agents-api/agents_api/routers/tasks/update_execution.py
index 7b60d8bdf..9adfe0b94 100644
--- a/agents-api/agents_api/routers/tasks/update_execution.py
+++ b/agents-api/agents_api/routers/tasks/update_execution.py
@@ -1,7 +1,7 @@
 from typing import Annotated
 
 from fastapi import Depends, HTTPException
-from pydantic import UUID4
+from uuid import UUID
 
 from agents_api.autogen.openapi_model import (
     ResumeExecutionRequest,
@@ -21,8 +21,8 @@
 
 @router.put("/executions/{execution_id}", tags=["executions"])
 async def update_execution(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    execution_id: UUID4,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
+    execution_id: UUID,
     data: ResumeExecutionRequest | StopExecutionRequest,
 ):
     temporal_client = await get_client()
diff --git a/agents-api/agents_api/routers/users/create_or_update_user.py b/agents-api/agents_api/routers/users/create_or_update_user.py
index 331f2c4d7..17102cf42 100644
--- a/agents-api/agents_api/routers/users/create_or_update_user.py
+++ b/agents-api/agents_api/routers/users/create_or_update_user.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
 
 from ...autogen.openapi_model import CreateOrUpdateUserRequest, ResourceCreatedResponse
@@ -14,8 +14,8 @@
 
 @router.post("/users/{user_id}", status_code=HTTP_201_CREATED, tags=["users"])
 async def create_or_update_user(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    user_id: UUID4,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
+    user_id: UUID,
     data: CreateOrUpdateUserRequest,
 ) -> ResourceCreatedResponse:
     user = create_or_update_user_query(
diff --git a/agents-api/agents_api/routers/users/create_user.py b/agents-api/agents_api/routers/users/create_user.py
index fcf8bf89b..53225cd53 100644
--- a/agents-api/agents_api/routers/users/create_user.py
+++ b/agents-api/agents_api/routers/users/create_user.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
 
 from ...autogen.openapi_model import CreateUserRequest, ResourceCreatedResponse
@@ -13,7 +13,7 @@
 @router.post("/users", status_code=HTTP_201_CREATED, tags=["users"])
 async def create_user(
     data: CreateUserRequest,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> ResourceCreatedResponse:
     user = create_user_query(
         developer_id=x_developer_id,
diff --git a/agents-api/agents_api/routers/users/delete_user.py b/agents-api/agents_api/routers/users/delete_user.py
index fd1d02a94..efbc2afdc 100644
--- a/agents-api/agents_api/routers/users/delete_user.py
+++ b/agents-api/agents_api/routers/users/delete_user.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("/users/{user_id}", status_code=HTTP_202_ACCEPTED, tags=["users"])
 async def delete_user(
-    user_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)]
+    user_id: UUID, x_developer_id: Annotated[UUID, Depends(get_developer_id)]
 ) -> ResourceDeletedResponse:
     return delete_user_query(developer_id=x_developer_id, user_id=user_id)
diff --git a/agents-api/agents_api/routers/users/get_user_details.py b/agents-api/agents_api/routers/users/get_user_details.py
index 0bc0460ca..3fcf0c5de 100644
--- a/agents-api/agents_api/routers/users/get_user_details.py
+++ b/agents-api/agents_api/routers/users/get_user_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 User
 from ...dependencies.developer_id import get_developer_id
@@ -11,7 +11,7 @@
 
 @router.get("/users/{user_id}", tags=["users"])
 async def get_user_details(
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    user_id: UUID4,
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
+    user_id: UUID,
 ) -> User:
     return get_user_query(developer_id=x_developer_id, user_id=user_id)
diff --git a/agents-api/agents_api/routers/users/list_users.py b/agents-api/agents_api/routers/users/list_users.py
index bfa5be128..e8ffc1a57 100644
--- a/agents-api/agents_api/routers/users/list_users.py
+++ b/agents-api/agents_api/routers/users/list_users.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 ListResponse, User
 from ...dependencies.developer_id import get_developer_id
@@ -13,7 +13,7 @@
 
 @router.get("/users", tags=["users"])
 async def list_users(
-    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",
diff --git a/agents-api/agents_api/routers/users/patch_user.py b/agents-api/agents_api/routers/users/patch_user.py
index fcd1e9380..54bc254f3 100644
--- a/agents-api/agents_api/routers/users/patch_user.py
+++ b/agents-api/agents_api/routers/users/patch_user.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 PatchUserRequest, ResourceUpdatedResponse
 from ...dependencies.developer_id import get_developer_id
@@ -11,9 +11,9 @@
 
 @router.patch("/users/{user_id}", tags=["users"])
 async def patch_user(
-    user_id: UUID4,
+    user_id: UUID,
     data: PatchUserRequest,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> ResourceUpdatedResponse:
     return patch_user_query(
         developer_id=x_developer_id,
diff --git a/agents-api/agents_api/routers/users/update_user.py b/agents-api/agents_api/routers/users/update_user.py
index 258023173..e13c38d92 100644
--- a/agents-api/agents_api/routers/users/update_user.py
+++ b/agents-api/agents_api/routers/users/update_user.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 ResourceUpdatedResponse, UpdateUserRequest
 from ...dependencies.developer_id import get_developer_id
@@ -11,9 +11,9 @@
 
 @router.put("/users/{user_id}", tags=["users"])
 async def update_user(
-    user_id: UUID4,
+    user_id: UUID,
     data: UpdateUserRequest,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> ResourceUpdatedResponse:
     return update_user_query(
         developer_id=x_developer_id,
diff --git a/agents-api/agents_api/workflows/task_execution/__init__.py b/agents-api/agents_api/workflows/task_execution/__init__.py
index b5f0b0a29..2ff45e55c 100644
--- a/agents-api/agents_api/workflows/task_execution/__init__.py
+++ b/agents-api/agents_api/workflows/task_execution/__init__.py
@@ -416,10 +416,10 @@ async def run(
                         task_steps.prompt_step,
                         context,
                         schedule_to_close_timeout=timedelta(
-                            seconds=30 if debug or testing else 600),
+                            seconds=30 if debug or testing else 600
+                        ),
                     )
-                    state = PartialTransition(
-                        output=new_response.output, type="resume")
+                    state = PartialTransition(output=new_response.output, type="resume")
 
             # case PromptStep(), StepOutcome(
             #     output=response
diff --git a/docs/js-sdk-docs/modules/api.md b/docs/js-sdk-docs/modules/api.md
index 9723c536f..bc4a9e82b 100644
--- a/docs/js-sdk-docs/modules/api.md
+++ b/docs/js-sdk-docs/modules/api.md
@@ -199,7 +199,7 @@ ___
 | :------ | :------ | :------ |
 | `messages` | [`InputChatMLMessage`](api.md#inputchatmlmessage)[] | A list of new input messages comprising the conversation so far. |
 | `tool_choice?` | [`ToolChoiceOption`](api.md#toolchoiceoption) \| [`NamedToolChoice`](api.md#namedtoolchoice) \| ``null`` | Can be one of existing tools given to the agent earlier or the ones included in the request |
-| `tools?` | [`Tool`](api.md#tool)[] \| ``null`` | (Advanced) List of tools that are provided in addition to agent's default set of tools. Functions of same name in agent set are overriden |
+| `tools?` | [`Tool`](api.md#tool)[] \| ``null`` | (Advanced) List of tools that are provided in addition to agent's default set of tools. Functions of same name in agent set are overridden |
 
 #### Defined in
 
diff --git a/llm-proxy/docker-compose.yml b/llm-proxy/docker-compose.yml
index 681ea4596..bbcb184ee 100644
--- a/llm-proxy/docker-compose.yml
+++ b/llm-proxy/docker-compose.yml
@@ -2,7 +2,7 @@ name: julep-llm-proxy
 
 services:
   litellm:
-    image: ghcr.io/berriai/litellm:main-stable
+    image: ghcr.io/berriai/litellm-database:main-v1.46.6
     volumes:
       - ./litellm-config.yaml:/app/config.yaml
       - .keys:/app/.keys
diff --git a/scheduler/docker-compose.yml b/scheduler/docker-compose.yml
index 3a0af3662..6e452eb8d 100644
--- a/scheduler/docker-compose.yml
+++ b/scheduler/docker-compose.yml
@@ -6,9 +6,9 @@ services:
       temporal-db:
         condition: service_started
     environment:
-      - DB=postgres12
-      - DB_PORT=5432
-      - DB_HOST=temporal-db
+      - DB=${TEMPORAL_POSTGRES_DB:-postgres12}
+      - DB_PORT=${TEMPORAL_POSTGRES_PORT:-5432}
+      - DB_HOST=${TEMPORAL_POSTGRES_HOST:-temporal-db}
       - POSTGRES_USER=${TEMPORAL_POSTGRES_USER:-temporal}
       - POSTGRES_PWD=${TEMPORAL_POSTGRES_PASSWORD}
       - POSTGRES_SEEDS=temporal-db
diff --git a/typespec/chat/endpoints.tsp b/typespec/chat/endpoints.tsp
index 64eb6700d..cceaca337 100644
--- a/typespec/chat/endpoints.tsp
+++ b/typespec/chat/endpoints.tsp
@@ -23,6 +23,10 @@ interface Endpoints {
         @doc("The session ID")
         id: uuid;
 
+        @header
+        @doc("Custom API key")
+        "X-Custom-Api-Key"?: string;
+
         @bodyRoot
         @doc("Request to generate a response from the model")
         body: ChatInput;