feat(agents-api): New chat context query and model (#437)

* feat(typespec): Add initial typespec definitions Signed-off-by: Diwank Tomer <diwank@julep.ai> * feat(typespec): Add Tasks typespec Signed-off-by: Diwank Tomer <diwank@julep.ai> * feat(typespec): Add Executions typespec Signed-off-by: Diwank Tomer <diwank@julep.ai> * feat: Add session token budget and context overflow * fix(typespec): Minor spec fixes Signed-off-by: Diwank Tomer <diwank@julep.ai> * feat(typespec): Add name field to Tool Signed-off-by: Diwank Tomer <diwank@julep.ai> * refactor(typespec): Minor refactors to common/ utils Signed-off-by: Diwank Tomer <diwank@julep.ai> * feat(typespec): Add search endpoints Signed-off-by: Diwank Tomer <diwank@julep.ai> * feat(typespec): Add multi-agent multi-user sessions, add template for situation Signed-off-by: Diwank Tomer <diwank@julep.ai> * refactor: Remove mock_openapi.yaml and link openapi.yaml to the typespec generated file Signed-off-by: Diwank Tomer <diwank@julep.ai> * feat: Codegen for all new typespec stuff Signed-off-by: Diwank Tomer <diwank@julep.ai> * feat(fern): Upgrade fern and fern-python Signed-off-by: Diwank Tomer <diwank@julep.ai> * feat(typespec): Various changes to the spec Signed-off-by: Diwank Tomer <diwank@julep.ai> * fix(typespec): Fix docstrings Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Julep Developers <developers@julep.ai> * fix(agents-api): Fix datamodel-codegen options Signed-off-by: Diwank Tomer <diwank@julep.ai> * fix(typespec): Misc fixes Signed-off-by: Diwank Tomer <diwank@julep.ai> * feat(agents-api): Add migrations for updated typespec models Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Julep Developers <developers@julep.ai> * feat(agents-api): Add migrations for developers relations and renaming information_snippets Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Julep Developers <developers@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * fix(agents-api): Update agent/tool models Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * reformat evthing Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * wip Signed-off-by: Diwank Tomer <diwank@julep.ai> * wup Signed-off-by: Diwank Tomer <diwank@julep.ai> * history Signed-off-by: Diwank Tomer <diwank@julep.ai> * feat(agents-api): New chat context query and model Signed-off-by: Diwank Tomer <diwank@julep.ai> --------- Signed-off-by: Diwank Tomer <diwank@julep.ai> Signed-off-by: Julep Developers <developers@julep.ai> Co-authored-by: Diwank Tomer <diwank@julep.ai> Co-authored-by: Dmitry Paramonov <asmatic075@gmail.com> Co-authored-by: Julep Developers <developers@julep.ai>
julep-ai · Jul 30, 2024 · 4808a00 · 4808a00
1 parent 687f487
commit 4808a00
Show file tree

Hide file tree

Showing 65 changed files with 1,020 additions and 1,036 deletions.
diff --git a/agents-api/agents_api/autogen/Entries.py b/agents-api/agents_api/autogen/Entries.py
@@ -11,6 +11,48 @@
 from .Tools import ChosenToolCall, Tool, ToolResponse
 
 
+class BaseEntry(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    role: Literal[
+        "user",
+        "agent",
+        "system",
+        "function",
+        "function_response",
+        "function_call",
+        "auto",
+    ]
+    """
+    ChatML role (system|assistant|user|function_call|function|function_response|auto)
+    """
+    name: str | None = None
+    content: (
+        list[ChatMLTextContentPart | ChatMLImageContentPart]
+        | Tool
+        | ChosenToolCall
+        | str
+        | ToolResponse
+        | list[
+            list[ChatMLTextContentPart | ChatMLImageContentPart]
+            | Tool
+            | ChosenToolCall
+            | str
+            | ToolResponse
+        ]
+    )
+    source: Literal[
+        "api_request", "api_response", "tool_response", "internal", "summarizer", "meta"
+    ]
+    tokenizer: str | None = None
+    token_count: int | None = None
+    timestamp: Annotated[float, Field(ge=0.0)]
+    """
+    This is the time that this event refers to.
+    """
+
+
 class ChatMLImageContentPart(BaseModel):
     model_config = ConfigDict(
         populate_by_name=True,
@@ -73,44 +115,10 @@ class ChatMLTextContentPart(BaseModel):
     """
 
 
-class Entry(BaseModel):
+class Entry(BaseEntry):
     model_config = ConfigDict(
         populate_by_name=True,
     )
-    role: Literal[
-        "user",
-        "agent",
-        "system",
-        "function",
-        "function_response",
-        "function_call",
-        "auto",
-    ]
-    """
-    ChatML role (system|assistant|user|function_call|function|function_response|auto)
-    """
-    name: str | None = None
-    content: (
-        list[ChatMLTextContentPart | ChatMLImageContentPart]
-        | Tool
-        | ChosenToolCall
-        | str
-        | ToolResponse
-        | list[
-            list[ChatMLTextContentPart | ChatMLImageContentPart]
-            | Tool
-            | ChosenToolCall
-            | str
-            | ToolResponse
-        ]
-    )
-    source: Literal[
-        "api_request", "api_response", "tool_response", "internal", "summarizer", "meta"
-    ]
-    timestamp: Annotated[float, Field(ge=0.0)]
-    """
-    This is the time that this event refers to.
-    """
     created_at: Annotated[AwareDatetime, Field(json_schema_extra={"readOnly": True})]
     """
     When this resource was created as UTC date-time
@@ -122,7 +130,7 @@ class History(BaseModel):
     model_config = ConfigDict(
         populate_by_name=True,
     )
-    entries: list[Entry]
+    entries: list[BaseEntry]
     relations: list[Relation]
     session_id: Annotated[UUID, Field(json_schema_extra={"readOnly": True})]
     created_at: Annotated[AwareDatetime, Field(json_schema_extra={"readOnly": True})]

diff --git a/agents-api/agents_api/autogen/openapi_model.py b/agents-api/agents_api/autogen/openapi_model.py
@@ -34,17 +34,13 @@
     "metadata",
 )
 
-ChatMLRole = Entry.model_fields["role"].annotation
+ChatMLRole = BaseEntry.model_fields["role"].annotation
 
 
-class CreateEntryRequest(Entry):
-    id: Annotated[UUID | None, Field(default=None)]
-    created_at: Annotated[AwareDatetime | None, Field(default=None)]
+class CreateEntryRequest(BaseEntry):
     timestamp: Annotated[
         float, Field(ge=0.0, default_factory=lambda: utcnow().timestamp())
     ]
-    tokenizer: str
-    token_count: int
 
 
 def make_session(
@@ -72,3 +68,5 @@ def make_session(
         case _:
             cls = MultiAgentMultiUserSession
             participants = {"agents": agents, "users": users}
+
+    return cls(**{**data, **participants})
diff --git a/agents-api/agents_api/common/protocol/entries.py b/agents-api/agents_api/common/protocol/entries.py
@@ -1,8 +1,8 @@
 import json
 from typing import Literal
-from uuid import UUID, uuid4
+from uuid import UUID
 
-from pydantic import BaseModel, Field, computed_field
+from pydantic import Field, computed_field
 
 from ...autogen.openapi_model import (
     ChatMLImageContentPart,
@@ -12,7 +12,6 @@
 from ...autogen.openapi_model import (
     Entry as BaseEntry,
 )
-from ...common.utils.datetime import utcnow
 
 EntrySource = Literal["api_request", "api_response", "internal", "summarizer"]
 Tokenizer = Literal["character_count"]

diff --git a/agents-api/agents_api/common/protocol/sessions.py b/agents-api/agents_api/common/protocol/sessions.py
@@ -3,11 +3,18 @@
 It includes definitions for session settings and session data models.
 """
 
-from typing import Dict, Optional
-from uuid import UUID
+from typing import Optional
 
 from pydantic import BaseModel
 
+from ...autogen.openapi_model import (
+    Agent,
+    Entry,
+    Session,
+    Settings,
+    Tool,
+    User,
+)
 from .agents import AgentDefaultSettings
 
 
@@ -22,26 +29,19 @@ class SessionSettings(AgentDefaultSettings):
 
 class SessionData(BaseModel):
     """
-    Represents the data associated with a session, including identifiers for the agent, user, and session itself,
-    along with session-specific information such as situation, summary, and timestamps.
+    Represents the data associated with a session, including for agents, and users.
     """
 
-    agent_id: UUID
-    user_id: Optional[UUID]
-    session_id: UUID
-    situation: str
-    summary: Optional[str]
-    user_name: Optional[str]
-    user_about: Optional[str]
-    agent_name: Optional[str]
-    agent_about: str
-    updated_at: float
-    created_at: float
-    model: str
-    default_settings: SessionSettings
-    render_templates: bool = False
-    metadata: Dict = {}
-    user_metadata: Optional[Dict] = None
-    agent_metadata: Dict = {}
-    token_budget: int | None = None
-    context_overflow: str | None = None
+    session: Session
+    agents: list[Agent]
+    users: list[User] = []
+    settings: Optional[Settings] = None
+
+
+class ChatContext(SessionData):
+    """
+    Represents the data associated with a context, including for agents, and users.
+    """
+
+    entries: list[Entry]
+    tools: list[Tool]
diff --git a/agents-api/agents_api/common/protocol/tasks.py b/agents-api/agents_api/common/protocol/tasks.py
@@ -18,7 +18,7 @@
     User,
     YieldStep,
 )
-from ...models.execution.prepare_execution_data import get_execution_input_query
+from ...models.execution.prepare_execution_data import prepare_execution_data
 from ..utils.cozo import uuid_int_list_to_uuid4
 
 WorkflowStep = (
@@ -117,7 +117,7 @@ class ExecutionInput(BaseModel):
     def fetch(
         cls, *, developer_id: UUID4, task_id: UUID4, execution_id: UUID4, client: Any
     ) -> "ExecutionInput":
-        [data] = get_execution_input_query(
+        [data] = prepare_execution_data(
             task_id=task_id,
             execution_id=execution_id,
             client=client,

diff --git a/agents-api/agents_api/models/agent/list_agents.py b/agents-api/agents_api/models/agent/list_agents.py
@@ -32,7 +32,7 @@ def list_agents(
     developer_id: UUID,
     limit: int = 100,
     offset: int = 0,
-    sort_by: Literal["created_at", "updated_at", "deleted_at"] = "created_at",
+    sort_by: Literal["created_at", "updated_at"] = "created_at",
     direction: Literal["asc", "desc"] = "desc",
     metadata_filter: dict[str, Any] = {},
 ) -> tuple[str, dict]:

diff --git a/agents-api/agents_api/models/entry/create_entries.py b/agents-api/agents_api/models/entry/create_entries.py
@@ -5,7 +5,7 @@
 from pycozo.client import QueryException
 from pydantic import ValidationError
 
-from ...autogen.openapi_model import CreateEntryRequest, Entry
+from ...autogen.openapi_model import CreateEntryRequest, Entry, Relation
 from ...common.utils.cozo import cozo_process_mutate_data
 from ...common.utils.datetime import utcnow
 from ...common.utils.messages import content_to_json
@@ -39,7 +39,7 @@ def create_entries(
     *,
     developer_id: UUID,
     session_id: UUID,
-    data: list[Entry],
+    data: list[CreateEntryRequest],
 ) -> tuple[str, dict]:
     developer_id = str(developer_id)
     session_id = str(session_id)
@@ -78,3 +78,44 @@ def create_entries(
     query = f"{{ {query} }}"
 
     return (query, {"rows": rows})
+
+
+@rewrap_exceptions(
+    {
+        QueryException: partialclass(HTTPException, status_code=400),
+        ValidationError: partialclass(HTTPException, status_code=400),
+        TypeError: partialclass(HTTPException, status_code=400),
+    }
+)
+@wrap_in_class(Relation)
+@cozo_query
+@beartype
+def add_entry_relations(
+    *,
+    developer_id: UUID,
+    data: list[Relation],
+) -> tuple[str, dict]:
+    developer_id = str(developer_id)
+
+    data_dicts = [item.model_dump(mode="json") for item in data]
+    cols, rows = cozo_process_mutate_data(data_dicts)
+
+    create_query = f"""
+        ?[{cols}] <- $rows
+
+        :insert relations {{
+            {cols}
+        }}
+
+        :returning
+    """
+
+    queries = [
+        verify_developer_id_query(developer_id),
+        create_query,
+    ]
+
+    query = "}\n\n{\n".join(queries)
+    query = f"{{ {query} }}"
+
+    return (query, {"rows": rows})