Merge pull request #455 from julep-ai/f/chat-tests

fix(agents-api): Fix chat endpoint behavior

agents-api/agents_api/activities/__init__.py

@@ -2,7 +2,6 @@
 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.
-- `demo.py`: Provides a simple demonstration of defining an activity with Temporal.
 - `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.

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.embed import embed
+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 embed(
+
+    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,8 +4,9 @@
 
 from temporalio import activity
 
-from ..clients.model import julep_client
-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.
@@ -117,10 +118,10 @@ def make_prompt(
 
 
 async def run_prompt(
-    dialog: list[ChatML],
+    dialog: list[InputChatMLMessage],
     session_id: UUID,
     previous_memories: list[str] = [],
-    model: str = "julep-ai/samantha-1-turbo",
+    model: str = "gpt-4o",
     max_tokens: int = 400,
     temperature: float = 0.4,
     parser: Callable[[str], str] = lambda x: x,
@@ -134,7 +135,7 @@ async def run_prompt(
         )
     )
 
-    response = await julep_client.chat.completions.create(
+    response = await litellm.acompletion(
         model=model,
         messages=[
             {
@@ -155,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 = [

agents-api/agents_api/activities/mem_rating.py

@@ -3,7 +3,7 @@
 
 from temporalio import activity
 
-from ..clients.model import julep_client
+from ..clients import litellm
 from .types import MemoryRatingTaskArgs
 
 
@@ -40,14 +40,14 @@ def make_prompt(args: MemoryRatingTaskArgs):
 
 async def run_prompt(
     memory: str,
-    model: str = "julep-ai/samantha-1-turbo",
+    model: str = "gpt-4o",
     max_tokens: int = 400,
     temperature: float = 0.1,
     parser: Callable[[str], str] = lambda x: x,
 ) -> str:
     prompt = make_prompt(MemoryRatingTaskArgs(memory=memory))
 
-    response = await julep_client.chat.completions.create(
+    response = await litellm.acompletion(
         model=model,
         messages=[
             {