feat(agents-api): Add fixtures for testing workflows

Signed-off-by: Diwank Tomer <[email protected]>

agents-api/agents_api/activities/embed_docs.py

@@ -1,27 +1,39 @@
-from pydantic import UUID4
+from uuid import UUID
+
 from temporalio import activity
 
 from agents_api.clients import embed as embedder
+from agents_api.clients.cozo import get_cozo_client
 from agents_api.models.docs.embed_snippets import embed_snippets as embed_snippets_query
 
 snippet_embed_instruction = "Encode this passage for retrieval: "
 
 
 @activity.defn
-async def embed_docs(doc_id: UUID4, title: str, content: list[str]) -> None:
+async def embed_docs(
+    developer_id: UUID,
+    doc_id: UUID,
+    title: str,
+    content: list[str],
+    include_title: bool = True,
+    cozo_client=None,
+) -> None:
     indices, snippets = list(zip(*enumerate(content)))
+
     embeddings = await embedder.embed(
         [
             {
                 "instruction": snippet_embed_instruction,
-                "text": title + "\n\n" + snippet,
+                "text": (title + "\n\n" + snippet) if include_title else snippet,
             }
             for snippet in snippets
         ]
     )
 
     embed_snippets_query(
+        developer_id=developer_id,
         doc_id=doc_id,
         snippet_indices=indices,
         embeddings=embeddings,
+        client=cozo_client or get_cozo_client(),
     )

agents-api/agents_api/activities/mem_mgmt.py

@@ -4,9 +4,9 @@
 
 from temporalio import activity
 
-from agents_api.clients import litellm
-
-from .types import ChatML, MemoryManagementTaskArgs
+from ..autogen.openapi_model import InputChatMLMessage
+from ..clients import litellm
+from .types import MemoryManagementTaskArgs
 
 example_previous_memory = """
 Speaker 1: Composes and listens to music. Likes to buy basketball shoes but doesn't wear them often.
@@ -118,7 +118,7 @@ def make_prompt(
 
 
 async def run_prompt(
-    dialog: list[ChatML],
+    dialog: list[InputChatMLMessage],
     session_id: UUID,
     previous_memories: list[str] = [],
     model: str = "gpt-4o",
@@ -156,7 +156,9 @@ async def run_prompt(
 
 @activity.defn
 async def mem_mgmt(
-    dialog: list[ChatML], session_id: UUID, previous_memories: list[str] = []
+    dialog: list[InputChatMLMessage],
+    session_id: UUID,
+    previous_memories: list[str] = [],
 ) -> None:
     # session_id = UUID(session_id)
     # entries = [