f/switch to pg

agents-api/.gitignore

@@ -1,5 +1,6 @@
 # Local database files
-cozo.db
+cozo*
+.cozo*
 temporal.db
 *.bak
 *.dat

agents-api/agents_api/activities/demo.py

@@ -1,5 +1,3 @@
-from typing import Callable
-
 from temporalio import activity
 
 from ..env import testing
@@ -14,6 +12,6 @@ async def mock_demo_activity(a: int, b: int) -> int:
     return a + b
 
 
-demo_activity: Callable[[int, int], int] = activity.defn(name="demo_activity")(
+demo_activity = activity.defn(name="demo_activity")(
     demo_activity if not testing else mock_demo_activity
 )

agents-api/agents_api/activities/execute_system.py

@@ -6,7 +6,6 @@
 
 from beartype import beartype
 from box import Box, BoxList
-from fastapi import HTTPException
 from fastapi.background import BackgroundTasks
 from temporalio import activity
 

agents-api/agents_api/activities/truncation.py

@@ -14,10 +14,10 @@
 def get_extra_entries(messages: list[Entry], token_count_threshold: int) -> list[UUID]:
     raise NotImplementedError()
 
-    if not len(messages):
-        return messages
+    # if not len(messages):
+    #     return messages
 
-    _token_cnt, _offset = 0, 0
+    # _token_cnt, _offset = 0, 0
     # if messages[0].role == Role.system:
     #     token_cnt, offset = messages[0].token_count, 1
 
@@ -36,7 +36,8 @@ def get_extra_entries(messages: list[Entry], token_count_threshold: int) -> list
 @activity.defn
 @beartype
 async def truncation(session_id: str, token_count_threshold: int) -> None:
-    session_id = UUID(session_id)
+    raise NotImplementedError()
+    # session_id = UUID(session_id)
 
     # delete_entries(
     #     get_extra_entries(

agents-api/agents_api/common/utils/cozo.py

@@ -22,5 +22,5 @@
 
 
 @beartype
-def uuid_int_list_to_uuid4(data: list[int]) -> UUID:
+def uuid_int_list_to_uuid(data: list[int]) -> UUID:
     return UUID(bytes=b"".join([i.to_bytes(1, "big") for i in data]))

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

@@ -4,12 +4,13 @@
 """
 
 from typing import Any, TypeVar
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from beartype import beartype
 from fastapi import HTTPException
 from pycozo.client import QueryException
 from pydantic import ValidationError
+from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import Agent, CreateAgentRequest
 from ...common.utils.cozo import cozo_process_mutate_data
@@ -78,7 +79,7 @@ def create_agent(
         Agent: The newly created agent record.
     """
 
-    agent_id = agent_id or uuid4()
+    agent_id = agent_id or uuid7()
 
     # Extract the agent data from the payload
     data.metadata = data.metadata or {}

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

@@ -1,10 +1,11 @@
 from typing import Any, Literal, TypeVar
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from beartype import beartype
 from fastapi import HTTPException
 from pycozo.client import QueryException
 from pydantic import ValidationError
+from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import CreateDocRequest, Doc
 from ...common.utils.cozo import cozo_process_mutate_data
@@ -58,7 +59,7 @@ def create_doc(
         data (CreateDocRequest): The content of the document.
     """
 
-    doc_id = str(doc_id or uuid4())
+    doc_id = str(doc_id or uuid7())
     owner_id = str(owner_id)
 
     if isinstance(data.content, str):

agents-api/agents_api/models/entry/create_entries.py

@@ -1,10 +1,11 @@
 from typing import Any, TypeVar
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from beartype import beartype
 from fastapi import HTTPException
 from pycozo.client import QueryException
 from pydantic import ValidationError
+from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import CreateEntryRequest, Entry, Relation
 from ...common.utils.cozo import cozo_process_mutate_data
@@ -58,7 +59,7 @@ def create_entries(
     for item in data_dicts:
         item["content"] = content_to_json(item["content"] or [])
         item["session_id"] = session_id
-        item["entry_id"] = item.pop("id", None) or str(uuid4())
+        item["entry_id"] = item.pop("id", None) or str(uuid7())
         item["created_at"] = (item.get("created_at") or utcnow()).timestamp()
 
     cols, rows = cozo_process_mutate_data(data_dicts)

agents-api/agents_api/models/entry/get_history.py

@@ -7,7 +7,7 @@
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import History
-from ...common.utils.cozo import uuid_int_list_to_uuid4 as fix_uuid
+from ...common.utils.cozo import uuid_int_list_to_uuid as fix_uuid
 from ..utils import (
     cozo_query,
     partialclass,

agents-api/agents_api/models/execution/create_execution.py

@@ -1,10 +1,11 @@
 from typing import Annotated, Any, TypeVar
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from beartype import beartype
 from fastapi import HTTPException
 from pycozo.client import QueryException
 from pydantic import ValidationError
+from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import CreateExecutionRequest, Execution
 from ...common.utils.cozo import cozo_process_mutate_data
@@ -47,7 +48,7 @@ def create_execution(
     execution_id: UUID | None = None,
     data: Annotated[CreateExecutionRequest | dict, dict_like(CreateExecutionRequest)],
 ) -> tuple[list[str], dict]:
-    execution_id = execution_id or uuid4()
+    execution_id = execution_id or uuid7()
 
     developer_id = str(developer_id)
     task_id = str(task_id)

agents-api/agents_api/models/execution/create_execution_transition.py

@@ -1,9 +1,10 @@
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from beartype import beartype
 from fastapi import HTTPException
 from pycozo.client import QueryException
 from pydantic import ValidationError
+from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import (
     CreateTransitionRequest,
@@ -38,7 +39,7 @@ def _create_execution_transition(
     update_execution_status: bool = False,
     task_id: UUID | None = None,
 ) -> tuple[list[str | None], dict]:
-    transition_id = transition_id or uuid4()
+    transition_id = transition_id or uuid7()
     data.metadata = data.metadata or {}
     data.execution_id = execution_id
 

agents-api/agents_api/models/files/create_file.py

@@ -6,12 +6,13 @@
 import base64
 import hashlib
 from typing import Any, TypeVar
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from beartype import beartype
 from fastapi import HTTPException
 from pycozo.client import QueryException
 from pydantic import ValidationError
+from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import CreateFileRequest, File
 from ...metrics.counters import increase_counter
@@ -79,7 +80,7 @@ def create_file(
         developer_id (UUID): The unique identifier for the developer creating the file.
     """
 
-    file_id = file_id or uuid4()
+    file_id = file_id or uuid7()
     file_data = data.model_dump(exclude={"content"})
 
     content_bytes = base64.b64decode(data.content)

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

@@ -4,12 +4,13 @@
 """
 
 from typing import Any, TypeVar
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from beartype import beartype
 from fastapi import HTTPException
 from pycozo.client import QueryException
 from pydantic import ValidationError
+from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import CreateSessionRequest, Session
 from ...metrics.counters import increase_counter
@@ -57,7 +58,7 @@ def create_session(
     Constructs and executes a datalog query to create a new session in the database.
     """
 
-    session_id = session_id or uuid4()
+    session_id = session_id or uuid7()
 
     data.metadata = data.metadata or {}
     session_data = data.model_dump(exclude={"auto_run_tools", "disable_cache"})

agents-api/agents_api/models/task/create_task.py

@@ -4,12 +4,13 @@
 """
 
 from typing import Any, TypeVar
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from beartype import beartype
 from fastapi import HTTPException
 from pycozo.client import QueryException
 from pydantic import ValidationError
+from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import (
     CreateTaskRequest,
@@ -74,7 +75,7 @@ def create_task(
     data.metadata = data.metadata or {}
     data.input_schema = data.input_schema or {}
 
-    task_id = task_id or uuid4()
+    task_id = task_id or uuid7()
     task_spec = task_to_spec(data)
 
     # Prepares the update data by filtering out None values and adding user_id and developer_id.

agents-api/agents_api/models/tools/create_tools.py

@@ -1,12 +1,13 @@
 """This module contains functions for creating tools in the CozoDB database."""
 
 from typing import Any, TypeVar
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from beartype import beartype
 from fastapi import HTTPException
 from pycozo.client import QueryException
 from pydantic import ValidationError
+from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import CreateToolRequest, Tool
 from ...metrics.counters import increase_counter
@@ -70,7 +71,7 @@ def create_tools(
     tools_data = [
         [
             str(agent_id),
-            str(uuid4()),
+            str(uuid7()),
             tool.type,
             tool.name,
             getattr(tool, tool.type) and getattr(tool, tool.type).model_dump(),

agents-api/agents_api/models/user/create_user.py

@@ -4,12 +4,13 @@
 """
 
 from typing import Any, TypeVar
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from beartype import beartype
 from fastapi import HTTPException
 from pycozo.client import QueryException
 from pydantic import ValidationError
+from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import CreateUserRequest, User
 from ...metrics.counters import increase_counter
@@ -80,7 +81,7 @@ def create_user(
         pd.DataFrame: A DataFrame containing the result of the query execution.
     """
 
-    user_id = user_id or uuid4()
+    user_id = user_id or uuid7()
     data.metadata = data.metadata or {}
     user_data = data.model_dump()
 

agents-api/agents_api/models/utils.py

@@ -14,7 +14,7 @@
 from pydantic import BaseModel
 from requests.exceptions import ConnectionError, Timeout
 
-from ..common.utils.cozo import uuid_int_list_to_uuid4
+from ..common.utils.cozo import uuid_int_list_to_uuid
 from ..env import do_verify_developer, do_verify_developer_owns_resource
 
 P = ParamSpec("P")
@@ -36,7 +36,7 @@ def fix_uuid(
     fixed = {
         **item,
         **{
-            attr: uuid_int_list_to_uuid4(item[attr])
+            attr: uuid_int_list_to_uuid(item[attr])
             for attr in id_attrs
             if isinstance(item[attr], list)
         },

agents-api/agents_api/routers/docs/create_doc.py

@@ -1,9 +1,10 @@
 from typing import Annotated
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from fastapi import BackgroundTasks, Depends
 from starlette.status import HTTP_201_CREATED
 from temporalio.client import Client as TemporalClient
+from uuid_extensions import uuid7
 
 from ...activities.types import EmbedDocsPayload
 from ...autogen.openapi_model import CreateDocRequest, Doc, ResourceCreatedResponse
@@ -82,7 +83,7 @@ async def create_user_doc(
         data=data,
     )
 
-    embed_job_id = uuid4()
+    embed_job_id = uuid7()
 
     await run_embed_docs_task(
         developer_id=x_developer_id,
@@ -113,7 +114,7 @@ async def create_agent_doc(
         data=data,
     )
 
-    embed_job_id = uuid4()
+    embed_job_id = uuid7()
 
     await run_embed_docs_task(
         developer_id=x_developer_id,

agents-api/agents_api/routers/sessions/chat.py

@@ -1,8 +1,9 @@
 from typing import Annotated, Optional
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from fastapi import BackgroundTasks, Depends, Header, HTTPException, status
 from starlette.status import HTTP_201_CREATED
+from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import (
     ChatInput,
@@ -236,7 +237,7 @@ async def chat(
         ChunkChatResponse if chat_input.stream else MessageChatResponse
     )
     chat_response: ChatResponse = chat_response_class(
-        id=uuid4(),
+        id=uuid7(),
         created_at=utcnow(),
         jobs=jobs,
         docs=doc_references,