Dev tasks

agents-api/agents_api/activities/task_steps/__init__.py

@@ -6,12 +6,14 @@
 from temporalio import activity
 
 from ...autogen.openapi_model import (
+    CreateTransitionRequest,
     EvaluateStep,
     # ErrorWorkflowStep,
     IfElseWorkflowStep,
     InputChatMLMessage,
     PromptStep,
     ToolCallStep,
+    UpdateExecutionRequest,
     YieldStep,
 )
 from ...clients.worker.types import ChatML
@@ -23,6 +25,9 @@
 from ...models.execution.create_execution_transition import (
     create_execution_transition as create_execution_transition_query,
 )
+from ...models.execution.update_execution import (
+    update_execution as update_execution_query,
+)
 from ...routers.sessions.protocol import Settings
 from ...routers.sessions.session import llm_generate
 
@@ -142,6 +147,15 @@ async def transition_step(
         **transition_data,
     )
 
+    update_execution_query(
+        developer_id=context.developer_id,
+        task_id=context.task.id,
+        execution_id=context.execution.id,
+        data=UpdateExecutionRequest(
+            status="awaiting_input",
+        ),
+    )
+
     # Raise if it's a waiting step
     if transition_info.type == "awaiting_input":
         activity.raise_complete_async()

agents-api/agents_api/autogen/Docs.py

@@ -13,21 +13,34 @@ class BaseDocSearchRequest(BaseModel):
     model_config = ConfigDict(
         populate_by_name=True,
     )
-    confidence: Annotated[float, Field(0.5, ge=0.0, le=1.0)]
+    limit: Annotated[int, Field(10, ge=1, le=100)]
+    lang: Literal["en-US"] = "en-US"
     """
-    The confidence cutoff level
+    The language to be used for text-only search. Support for other languages coming soon.
     """
-    alpha: Annotated[float, Field(0.75, ge=0.0, le=1.0)]
+
+
+class CreateDocRequest(BaseModel):
     """
-    The weight to apply to BM25 vs Vector search results. 0 => pure BM25; 1 => pure vector;
+    Payload for creating a doc
     """
-    mmr: bool = False
+
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    metadata: dict[str, Any] | None = None
+    title: Annotated[
+        str,
+        Field(
+            pattern="^[\\p{L}\\p{Nl}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]+[\\p{ID_Start}\\p{Mn}\\p{Mc}\\p{Nd}\\p{Pc}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]*$"
+        ),
+    ]
     """
-    Whether to include the MMR algorithm in the search. Optimizes for diversity in search results.
+    Title describing what this document contains
     """
-    lang: Literal["en-US"] = "en-US"
+    content: str | list[str]
     """
-    The language to be used for text-only search. Support for other languages coming soon.
+    Contents of the document
     """
 
 
@@ -76,12 +89,23 @@ class DocReference(BaseModel):
     """
     ID of the document
     """
-    snippet_index: list[int]
+    title: str | None = None
+    snippets: Annotated[list[Snippet], Field(min_length=1)]
+    distance: float | None = None
+
+
+class DocSearchResponse(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    docs: list[DocReference]
+    """
+    The documents that were found
     """
-    Snippets referred to of the document
+    time: Annotated[float, Field(gt=0.0)]
+    """
+    The time taken to search in seconds
     """
-    title: str | None = None
-    snippet: str | None = None
 
 
 class EmbedQueryRequest(BaseModel):
@@ -108,31 +132,51 @@ class HybridDocSearchRequest(BaseDocSearchRequest):
     model_config = ConfigDict(
         populate_by_name=True,
     )
-    text: str | list[str]
+    confidence: Annotated[float, Field(0.5, ge=0.0, le=1.0)]
+    """
+    The confidence cutoff level
+    """
+    alpha: Annotated[float, Field(0.75, ge=0.0, le=1.0)]
     """
-    Text or texts to use in the search. In `hybrid` search mode, either `text` or both `text` and `vector` fields are required.
+    The weight to apply to BM25 vs Vector search results. 0 => pure BM25; 1 => pure vector;
+    """
+    text: str
     """
-    vector: list[float] | list[list[float]]
+    Text to use in the search. In `hybrid` search mode, either `text` or both `text` and `vector` fields are required.
     """
-    Vector or vectors to use in the search. Must be the same dimensions as the embedding model or else an error will be thrown.
+    vector: list[float]
     """
+    Vector to use in the search. Must be the same dimensions as the embedding model or else an error will be thrown.
+    """
+
+
+class Snippet(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    index: int
+    content: str
 
 
 class TextOnlyDocSearchRequest(BaseDocSearchRequest):
     model_config = ConfigDict(
         populate_by_name=True,
     )
-    text: str | list[str]
+    text: str
     """
-    Text or texts to use in the search.
+    Text to use in the search.
     """
 
 
 class VectorDocSearchRequest(BaseDocSearchRequest):
     model_config = ConfigDict(
         populate_by_name=True,
     )
-    vector: list[float] | list[list[float]]
+    confidence: Annotated[float, Field(0.5, ge=0.0, le=1.0)]
+    """
+    The confidence cutoff level
+    """
+    vector: list[float]
     """
-    Vector or vectors to use in the search. Must be the same dimensions as the embedding model or else an error will be thrown.
+    Vector to use in the search. Must be the same dimensions as the embedding model or else an error will be thrown.
     """

agents-api/agents_api/autogen/Executions.py

@@ -68,10 +68,6 @@ class TaskTokenResumeExecutionRequest(BaseModel):
         populate_by_name=True,
     )
     status: Literal["running"] = "running"
-    task_token: str
-    """
-    A Task Token is a unique identifier for a specific Task Execution.
-    """
     input: dict[str, Any] | None = None
     """
     The input to resume the execution with

agents-api/agents_api/autogen/openapi_model.py

@@ -1,5 +1,5 @@
 # ruff: noqa: F401, F403, F405
-from typing import Annotated
+from typing import Annotated, Generic, TypeVar
 from uuid import UUID
 
 from pydantic import AwareDatetime, Field
@@ -158,3 +158,10 @@ class UpdateTaskRequest(_UpdateTaskRequest):
             "extra": "allow",
         }
     )
+
+
+DataT = TypeVar("DataT", bound=BaseModel)
+
+
+class ListResponse(BaseModel, Generic[DataT]):
+    items: list[DataT]

agents-api/agents_api/clients/temporal.py

@@ -75,7 +75,7 @@ async def run_task_execution_workflow(
 ):
     client = await get_client()
 
-    await client.execute_workflow(
+    return await client.start_workflow(
         "TaskExecutionWorkflow",
         args=[execution_input, start, previous_inputs],
         task_queue="memory-task-queue",

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

@@ -0,0 +1,22 @@
+from typing import Type
+
+from beartype.vale import Is
+from beartype.vale._core._valecore import BeartypeValidator
+from pydantic import BaseModel
+
+
+def dict_like(pydantic_model_class: Type[BaseModel]) -> BeartypeValidator:
+    required_fields_set: set[str] = set(
+        [
+            field
+            for field, info in pydantic_model_class.model_fields.items()
+            if info.is_required()
+        ]
+    )
+
+    validator = Is[
+        lambda x: isinstance(x, pydantic_model_class)
+        or required_fields_set.issubset(set(x.keys()))
+    ]
+
+    return validator

agents-api/agents_api/models/__init__.py

@@ -5,3 +5,14 @@
 
 This module also integrates with the `common` module for exception handling and utility functions, ensuring robust error management and providing reusable components for data processing and query construction.
 """
+
+# ruff: noqa: F401, F403, F405
+
+import agents_api.models.agent as agent
+import agents_api.models.docs as docs
+import agents_api.models.entry as entry
+import agents_api.models.execution as execution
+import agents_api.models.session as session
+import agents_api.models.task as task
+import agents_api.models.tools as tools
+import agents_api.models.user as user

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

@@ -10,3 +10,13 @@
 
 This module serves as the backbone for agent management within the CozoDB ecosystem, facilitating a wide range of operations necessary for the effective handling of agent data.
 """
+
+# ruff: noqa: F401, F403, F405
+
+from .create_agent import create_agent
+from .create_or_update_agent import create_or_update_agent
+from .delete_agent import delete_agent
+from .get_agent import get_agent
+from .list_agents import list_agents
+from .patch_agent import patch_agent
+from .update_agent import update_agent

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

@@ -43,18 +43,12 @@ def create_agent(
     Constructs and executes a datalog query to create a new agent in the database.
 
     Parameters:
-    - agent_id (UUID): The unique identifier for the agent.
+    - agent_id (UUID | None): The unique identifier for the agent.
     - developer_id (UUID): The unique identifier for the developer creating the agent.
-    - name (str): The name of the agent.
-    - about (str): A description of the agent.
-    - instructions (list[str], optional): A list of instructions for using the agent. Defaults to an empty list.
-    - model (str, optional): The model identifier for the agent. Defaults to "julep-ai/samantha-1-turbo".
-    - metadata (dict, optional): A dictionary of metadata for the agent. Defaults to an empty dict.
-    - default_settings (dict, optional): A dictionary of default settings for the agent. Defaults to an empty dict.
-    - client (CozoClient, optional): The CozoDB client instance to use for the query. Defaults to a preconfigured client instance.
+    - data (CreateAgentRequest): The data for the new agent.
 
     Returns:
-    Agent: The newly created agent record.
+    - Agent: The newly created agent record.
     """
 
     agent_id = agent_id or uuid4()