fix(agents-api): Minor fixes

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

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/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/routers/docs/create_doc.py

@@ -1,15 +1,35 @@
 from typing import Annotated
+from uuid import UUID, uuid4
 
 from fastapi import Depends
 from pydantic import UUID4
 from starlette.status import HTTP_201_CREATED
+from temporalio.client import Client as TemporalClient
 
 from ...autogen.openapi_model import CreateDocRequest, ResourceCreatedResponse
+from ...clients import temporal
 from ...dependencies.developer_id import get_developer_id
 from ...models.docs.create_doc import create_doc as create_doc_query
 from .router import router
 
 
+async def run_embed_docs_task(
+    doc_id: UUID,
+    title: str,
+    content: list[str],
+    job_id: UUID,
+    client: TemporalClient | None = None,
+):
+    client = client or (await temporal.get_client())
+
+    await client.execute_workflow(
+        "EmbedDocsWorkflow",
+        args=[str(doc_id), title, content],
+        task_queue="memory-task-queue",
+        id=str(job_id),
+    )
+
+
 @router.post("/users/{user_id}/docs", status_code=HTTP_201_CREATED, tags=["docs"])
 async def create_user_doc(
     user_id: UUID4,
@@ -23,7 +43,18 @@ async def create_user_doc(
         data=data,
     )
 
-    return ResourceCreatedResponse(id=doc.id, created_at=doc.created_at, jobs=[])
+    embed_job_id = uuid4()
+
+    await run_embed_docs_task(
+        doc_id=doc.id,
+        title=doc.title,
+        content=doc.content,
+        job_id=embed_job_id,
+    )
+
+    return ResourceCreatedResponse(
+        id=doc.id, created_at=doc.created_at, jobs=[embed_job_id]
+    )
 
 
 @router.post("/agents/{agent_id}/docs", status_code=HTTP_201_CREATED, tags=["docs"])
@@ -39,4 +70,15 @@ async def create_agent_doc(
         data=data,
     )
 
-    return ResourceCreatedResponse(id=doc.id, created_at=doc.created_at, jobs=[])
+    embed_job_id = uuid4()
+
+    await run_embed_docs_task(
+        doc_id=doc.id,
+        title=doc.title,
+        content=doc.content,
+        job_id=embed_job_id,
+    )
+
+    return ResourceCreatedResponse(
+        id=doc.id, created_at=doc.created_at, jobs=[embed_job_id]
+    )

agents-api/tests/fixtures.py

@@ -74,6 +74,18 @@ async def temporal_worker(wf_env=workflow_environment):
         await c
 
 
+@fixture(scope="test")
+def patch_temporal_get_client(
+    wf_env=workflow_environment,
+    temporal_worker=temporal_worker,
+):
+    mock_client = wf_env.client
+
+    with patch("agents_api.clients.temporal.get_client") as get_client:
+        get_client.return_value = mock_client
+        yield get_client
+
+
 @fixture(scope="global")
 def test_developer_id(cozo_client=cozo_client):
     developer_id = uuid4()

agents-api/tests/test_activities.py

@@ -16,7 +16,18 @@
 # from agents_api.common.protocol.entries import Entry
 
 
-@test("activity: embed_docs")
+@test("activity: check that workflow environment and worker are started correctly")
+async def _(
+    workflow_environment=workflow_environment,
+    worker=temporal_worker,
+):
+    async with workflow_environment as wf_env:
+        assert wf_env is not None
+        assert worker is not None
+        assert worker.is_running
+
+
+@test("activity: call direct embed_docs")
 async def _(
     cozo_client=cozo_client,
     developer_id=test_developer_id,
@@ -41,17 +52,6 @@ async def _(
     embed.assert_called_once()
 
 
-@test("activity: check that workflow environment and worker are started correctly")
-async def _(
-    workflow_environment=workflow_environment,
-    worker=temporal_worker,
-):
-    async with workflow_environment as wf_env:
-        assert wf_env is not None
-        assert worker is not None
-        assert worker.is_running
-
-
 # @test("get extra entries, do not strip system message")
 # def _():
 #     session_ids = [uuid.uuid4()] * 3

agents-api/tests/test_docs_routes.py

@@ -25,6 +25,9 @@ def _(make_request=make_request, user=test_user):
 
     assert response.status_code == 201
 
+    result = response.json()
+    assert len(result["jobs"]) > 0
+
 
 @test("route: create agent doc")
 def _(make_request=make_request, agent=test_agent):
@@ -41,9 +44,8 @@ def _(make_request=make_request, agent=test_agent):
 
     assert response.status_code == 201
 
-    # FIXME: Should create a job to process the document
-    # result = response.json()
-    # assert len(result["jobs"]) > 0
+    result = response.json()
+    assert len(result["jobs"]) > 0
 
 
 @test("route: delete doc")