fix(agents-api): Minor fix to tests

agents-api/agents_api/activities/__init__.py

@@ -1,13 +1,5 @@
 """
-The `activities` module within the agents-api package is designed to facilitate various activities related to agent interactions. This includes handling memory management, generating insights from dialogues, summarizing relationships, and more. Each file within the module offers specific functionality:
-
-- `co_density.py`: Conducts cognitive density analysis to generate concise, entity-dense summaries.
-- `dialog_insights.py`: Extracts insights from dialogues, identifying details that participants might find interesting.
-- `mem_mgmt.py`: Manages memory by updating and incorporating new personality information from dialogues.
-- `mem_rating.py`: Rates memories based on their poignancy and importance.
-- `relationship_summary.py`: Summarizes the relationship between individuals based on provided statements.
-- `salient_questions.py`: Identifies salient questions from a set of statements.
-- `summarization.py`: Summarizes dialogues and updates memory based on the conversation context.
+The `activities` module within the agents-api package is designed to facilitate various activities related to agent interactions. This includes handling memory management, generating insights from dialogues, summarizing relationships, and more.
 
 This module plays a crucial role in enhancing the capabilities of agents by providing them with the tools to understand and process information more effectively.
 """

agents-api/agents_api/activities/summarization.py

@@ -1,23 +1,13 @@
 #!/usr/bin/env python3
 
-import asyncio
-from textwrap import dedent
-from typing import Callable
-from uuid import UUID
 
 import pandas as pd
 from temporalio import activity
 
-# from agents_api.common.protocol.entries import Entry
 # from agents_api.models.entry.entries_summarization import (
 #     entries_summarization_query,
 #     get_toplevel_entries_query,
 # )
-from agents_api.rec_sum.entities import get_entities
-from agents_api.rec_sum.summarize import summarize_messages
-from agents_api.rec_sum.trim import trim_messages
-
-from ..env import summarization_model_name
 
 
 # TODO: remove stubs

agents-api/agents_api/activities/truncation.py

@@ -2,9 +2,7 @@
 
 from temporalio import activity
 
-# from agents_api.autogen.openapi_model import Role
-from agents_api.common.protocol.entries import Entry
-from agents_api.models.entry.delete_entries import delete_entries
+from agents_api.autogen.openapi_model import Entry
 
 # from agents_api.models.entry.entries_summarization import get_toplevel_entries_query
 
@@ -13,8 +11,7 @@ def get_extra_entries(messages: list[Entry], token_count_threshold: int) -> list
     if not len(messages):
         return messages
 
-    result: list[UUID] = []
-    token_cnt, offset = 0, 0
+    _token_cnt, _offset = 0, 0
     # if messages[0].role == Role.system:
     #     token_cnt, offset = messages[0].token_count, 1
 

agents-api/agents_api/autogen/Entries.py

@@ -45,8 +45,8 @@ class BaseEntry(BaseModel):
     source: Literal[
         "api_request", "api_response", "tool_response", "internal", "summarizer", "meta"
     ]
-    tokenizer: str | None = None
-    token_count: int | None = None
+    tokenizer: str
+    token_count: int
     timestamp: Annotated[float, Field(ge=0.0)]
     """
     This is the time that this event refers to.

agents-api/agents_api/autogen/openapi_model.py

@@ -1,7 +1,9 @@
 # ruff: noqa: F401, F403, F405
-from typing import Annotated, Generic, TypeVar
+from typing import Annotated, Generic, Self, Type, TypeVar
 from uuid import UUID
 
+from litellm.utils import _select_tokenizer as select_tokenizer
+from litellm.utils import token_counter
 from pydantic import AwareDatetime, Field
 from pydantic_partial import create_partial_model
 
@@ -34,14 +36,67 @@
     "metadata",
 )
 
-ChatMLRole = BaseEntry.model_fields["role"].annotation
+ChatMLRole = Literal[
+    "user",
+    "assistant",
+    "system",
+    "function",
+    "function_response",
+    "function_call",
+    "auto",
+]
+
+ChatMLContent = (
+    list[ChatMLTextContentPart | ChatMLImageContentPart]
+    | Tool
+    | ChosenToolCall
+    | str
+    | ToolResponse
+    | list[
+        list[ChatMLTextContentPart | ChatMLImageContentPart]
+        | Tool
+        | ChosenToolCall
+        | str
+        | ToolResponse
+    ]
+)
+
+ChatMLSource = Literal[
+    "api_request", "api_response", "tool_response", "internal", "summarizer", "meta"
+]
 
 
 class CreateEntryRequest(BaseEntry):
     timestamp: Annotated[
         float, Field(ge=0.0, default_factory=lambda: utcnow().timestamp())
     ]
 
+    @classmethod
+    def from_model_input(
+        cls: Type[Self],
+        model: str,
+        *,
+        role: ChatMLRole,
+        content: ChatMLContent,
+        name: str | None = None,
+        source: ChatMLSource,
+        **kwargs: dict,
+    ) -> Self:
+        tokenizer: dict = select_tokenizer(model=model)
+        token_count = token_counter(
+            model=model, messages=[{"role": role, "content": content, "name": name}]
+        )
+
+        return cls(
+            role=role,
+            content=content,
+            name=name,
+            source=source,
+            tokenizer=tokenizer["type"],
+            token_count=token_count,
+            **kwargs,
+        )
+
 
 def make_session(
     *,

agents-api/agents_api/clients/temporal.py

@@ -34,52 +34,6 @@ async def get_client(
     )
 
 
-async def run_summarization_task(
-    session_id: UUID, job_id: UUID, client: Client | None = None
-):
-    client = client or (await get_client())
-
-    await client.execute_workflow(
-        "SummarizationWorkflow",
-        args=[str(session_id)],
-        task_queue="memory-task-queue",
-        id=str(job_id),
-    )
-
-
-async def run_embed_docs_task(
-    doc_id: UUID,
-    title: str,
-    content: list[str],
-    job_id: UUID,
-    client: Client | None = None,
-):
-    client = client or (await get_client())
-
-    await client.execute_workflow(
-        "EmbedDocsWorkflow",
-        args=[str(doc_id), title, content],
-        task_queue="memory-task-queue",
-        id=str(job_id),
-    )
-
-
-async def run_truncation_task(
-    token_count_threshold: int,
-    session_id: UUID,
-    job_id: UUID,
-    client: Client | None = None,
-):
-    client = client or (await get_client())
-
-    await client.execute_workflow(
-        "TruncationWorkflow",
-        args=[str(session_id), token_count_threshold],
-        task_queue="memory-task-queue",
-        id=str(job_id),
-    )
-
-
 async def run_task_execution_workflow(
     execution_input: ExecutionInput,
     job_id: UUID,

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

@@ -34,7 +34,10 @@
     }
 )
 @wrap_in_class(
-    Agent, one=True, transform=lambda d: {"id": UUID(d.pop("agent_id")), **d}
+    Agent,
+    one=True,
+    transform=lambda d: {"id": UUID(d.pop("agent_id")), **d},
+    _kind="inserted",
 )
 @cozo_query
 @beartype

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

@@ -41,6 +41,7 @@
         "deleted_at": utcnow(),
         "jobs": [],
     },
+    _kind="deleted",
 )
 @cozo_query
 @beartype

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

@@ -29,6 +29,7 @@
     ResourceUpdatedResponse,
     one=True,
     transform=lambda d: {"id": d["agent_id"], "jobs": [], **d},
+    _kind="replaced",
 )
 @cozo_query
 @beartype

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

@@ -28,6 +28,7 @@
     ResourceUpdatedResponse,
     one=True,
     transform=lambda d: {"id": d["agent_id"], "jobs": [], **d},
+    _kind="replaced",
 )
 @cozo_query
 @beartype

agents-api/agents_api/models/chat/__init__.py

@@ -0,0 +1,22 @@
+"""
+Module: agents_api/models/docs
+
+This module is responsible for managing document-related operations within the application, particularly for agents and possibly other entities. It serves as a core component of the document management system, enabling features such as document creation, listing, deletion, and embedding of snippets for enhanced search and retrieval capabilities.
+
+Main functionalities include:
+- Creating new documents and associating them with agents or users.
+- Listing documents based on various criteria, including ownership and metadata filters.
+- Deleting documents by their unique identifiers.
+- Embedding document snippets for retrieval purposes.
+
+The module interacts with other parts of the application, such as the agents and users modules, to provide a comprehensive document management system. Its role is crucial in enabling document search, retrieval, and management features within the context of agents and users.
+
+This documentation aims to provide clear, concise, and sufficient context for new developers or contributors to understand the module's role without needing to dive deep into the code immediately.
+"""
+
+# ruff: noqa: F401, F403, F405
+
+from .gather_messages import gather_messages
+from .get_cached_response import get_cached_response
+from .prepare_chat_context import prepare_chat_context
+from .set_cached_response import set_cached_response