letta-nightly 0.7.0.dev20250423003112__py3-none-any.whl → 0.7.2.dev20250423222439__py3-none-any.whl

letta/schemas/openai/chat_completion_response.py

@@ -39,9 +39,10 @@ class Message(BaseModel):
     tool_calls: Optional[List[ToolCall]] = None
     role: str
     function_call: Optional[FunctionCall] = None  # Deprecated
-    reasoning_content: Optional[str] = None  # Used in newer reasoning APIs
+    reasoning_content: Optional[str] = None  # Used in newer reasoning APIs, e.g. DeepSeek
     reasoning_content_signature: Optional[str] = None  # NOTE: for Anthropic
     redacted_reasoning_content: Optional[str] = None  # NOTE: for Anthropic
+    ommitted_reasoning_content: bool = False  # NOTE: for OpenAI o1/o3
 
 
 class Choice(BaseModel):
@@ -52,16 +53,64 @@ class Choice(BaseModel):
     seed: Optional[int] = None  # found in TogetherAI
 
 
+class UsageStatisticsPromptTokenDetails(BaseModel):
+    cached_tokens: int = 0
+    # NOTE: OAI specific
+    # audio_tokens: int = 0
+
+    def __add__(self, other: "UsageStatisticsPromptTokenDetails") -> "UsageStatisticsPromptTokenDetails":
+        return UsageStatisticsPromptTokenDetails(
+            cached_tokens=self.cached_tokens + other.cached_tokens,
+        )
+
+
+class UsageStatisticsCompletionTokenDetails(BaseModel):
+    reasoning_tokens: int = 0
+    # NOTE: OAI specific
+    # audio_tokens: int = 0
+    # accepted_prediction_tokens: int = 0
+    # rejected_prediction_tokens: int = 0
+
+    def __add__(self, other: "UsageStatisticsCompletionTokenDetails") -> "UsageStatisticsCompletionTokenDetails":
+        return UsageStatisticsCompletionTokenDetails(
+            reasoning_tokens=self.reasoning_tokens + other.reasoning_tokens,
+        )
+
+
 class UsageStatistics(BaseModel):
     completion_tokens: int = 0
     prompt_tokens: int = 0
     total_tokens: int = 0
 
+    prompt_tokens_details: Optional[UsageStatisticsPromptTokenDetails] = None
+    completion_tokens_details: Optional[UsageStatisticsCompletionTokenDetails] = None
+
     def __add__(self, other: "UsageStatistics") -> "UsageStatistics":
+
+        if self.prompt_tokens_details is None and other.prompt_tokens_details is None:
+            total_prompt_tokens_details = None
+        elif self.prompt_tokens_details is None:
+            total_prompt_tokens_details = other.prompt_tokens_details
+        elif other.prompt_tokens_details is None:
+            total_prompt_tokens_details = self.prompt_tokens_details
+        else:
+            total_prompt_tokens_details = self.prompt_tokens_details + other.prompt_tokens_details
+
+        if self.completion_tokens_details is None and other.completion_tokens_details is None:
+            total_completion_tokens_details = None
+        elif self.completion_tokens_details is None:
+            total_completion_tokens_details = other.completion_tokens_details
+        elif other.completion_tokens_details is None:
+            total_completion_tokens_details = self.completion_tokens_details
+        else:
+            total_completion_tokens_details = self.completion_tokens_details + other.completion_tokens_details
+
         return UsageStatistics(
             completion_tokens=self.completion_tokens + other.completion_tokens,
             prompt_tokens=self.prompt_tokens + other.prompt_tokens,
             total_tokens=self.total_tokens + other.total_tokens,
+            prompt_tokens_details=total_prompt_tokens_details,
+            completion_tokens_details=total_completion_tokens_details,
         )
 
 
@@ -70,7 +119,7 @@ class ChatCompletionResponse(BaseModel):
 
     id: str
     choices: List[Choice]
-    created: datetime.datetime
+    created: Union[datetime.datetime, int]
     model: Optional[str] = None  # NOTE: this is not consistent with OpenAI API standard, however is necessary to support local LLMs
     # system_fingerprint: str  # docs say this is mandatory, but in reality API returns None
     system_fingerprint: Optional[str] = None
@@ -138,7 +187,7 @@ class ChatCompletionChunkResponse(BaseModel):
 
     id: str
     choices: List[ChunkChoice]
-    created: Union[datetime.datetime, str]
+    created: Union[datetime.datetime, int]
     model: str
     # system_fingerprint: str  # docs say this is mandatory, but in reality API returns None
     system_fingerprint: Optional[str] = None

letta/schemas/response_format.py

@@ -0,0 +1,78 @@
+from enum import Enum
+from typing import Annotated, Any, Dict, Literal, Union
+
+from pydantic import BaseModel, Field, validator
+
+
+class ResponseFormatType(str, Enum):
+    """Enum defining the possible response format types."""
+
+    text = "text"
+    json_schema = "json_schema"
+    json_object = "json_object"
+
+
+class ResponseFormat(BaseModel):
+    """Base class for all response formats."""
+
+    type: ResponseFormatType = Field(
+        ...,
+        description="The type of the response format.",
+        # why use this?
+        example=ResponseFormatType.text,
+    )
+
+
+# ---------------------
+# Response Format Types
+# ---------------------
+
+# SQLAlchemy type for database mapping
+ResponseFormatDict = Dict[str, Any]
+
+
+class TextResponseFormat(ResponseFormat):
+    """Response format for plain text responses."""
+
+    type: Literal[ResponseFormatType.text] = Field(
+        ResponseFormatType.text,
+        description="The type of the response format.",
+    )
+
+
+class JsonSchemaResponseFormat(ResponseFormat):
+    """Response format for JSON schema-based responses."""
+
+    type: Literal[ResponseFormatType.json_schema] = Field(
+        ResponseFormatType.json_schema,
+        description="The type of the response format.",
+    )
+    json_schema: Dict[str, Any] = Field(
+        ...,
+        description="The JSON schema of the response.",
+    )
+
+    @validator("json_schema")
+    def validate_json_schema(cls, v: Dict[str, Any]) -> Dict[str, Any]:
+        """Validate that the provided schema is a valid JSON schema."""
+        if not isinstance(v, dict):
+            raise ValueError("JSON schema must be a dictionary")
+        if "schema" not in v:
+            raise ValueError("JSON schema should include a $schema property")
+        return v
+
+
+class JsonObjectResponseFormat(ResponseFormat):
+    """Response format for JSON object responses."""
+
+    type: Literal[ResponseFormatType.json_object] = Field(
+        ResponseFormatType.json_object,
+        description="The type of the response format.",
+    )
+
+
+# Pydantic type for validation
+ResponseFormatUnion = Annotated[
+    Union[TextResponseFormat | JsonSchemaResponseFormat | JsonObjectResponseFormat],
+    Field(discriminator="type"),
+]

letta/schemas/tool_execution_result.py

@@ -0,0 +1,14 @@
+from typing import Any, List, Literal, Optional
+
+from pydantic import BaseModel, Field
+
+from letta.schemas.agent import AgentState
+
+
+class ToolExecutionResult(BaseModel):
+    status: Literal["success", "error"] = Field(..., description="The status of the tool execution and return object")
+    func_return: Optional[Any] = Field(None, description="The function return object")
+    agent_state: Optional[AgentState] = Field(None, description="The agent state")
+    stdout: Optional[List[str]] = Field(None, description="Captured stdout (prints, logs) from function invocation")
+    stderr: Optional[List[str]] = Field(None, description="Captured stderr from the function invocation")
+    sandbox_config_fingerprint: Optional[str] = Field(None, description="The fingerprint of the config for the sandbox")

letta/server/rest_api/chat_completions_interface.py

@@ -238,7 +238,7 @@ class ChatCompletionsStreamingInterface(AgentChunkStreamingInterface):
                     return ChatCompletionChunk(
                         id=chunk.id,
                         object=chunk.object,
-                        created=chunk.created.timestamp(),
+                        created=chunk.created,
                         model=chunk.model,
                         choices=[
                             Choice(
@@ -256,7 +256,7 @@ class ChatCompletionsStreamingInterface(AgentChunkStreamingInterface):
                 return ChatCompletionChunk(
                     id=chunk.id,
                     object=chunk.object,
-                    created=chunk.created.timestamp(),
+                    created=chunk.created,
                     model=chunk.model,
                     choices=[
                         Choice(

letta/server/rest_api/interface.py

@@ -1001,7 +1001,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
                 # Example case that would trigger here:
                 # id='chatcmpl-AKtUvREgRRvgTW6n8ZafiKuV0mxhQ'
                 # choices=[ChunkChoice(finish_reason=None, index=0, delta=MessageDelta(content=None, tool_calls=None, function_call=None), logprobs=None)]
-                # created=datetime.datetime(2024, 10, 21, 20, 40, 57, tzinfo=TzInfo(UTC))
+                # created=1713216662
                 # model='gpt-4o-mini-2024-07-18'
                 # object='chat.completion.chunk'
                 warnings.warn(f"Couldn't find delta in chunk: {chunk}")
@@ -1240,10 +1240,11 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
                         and function_call.function.name == self.assistant_message_tool_name
                         and self.assistant_message_tool_kwarg in func_args
                     ):
+                        # Coerce content to `str` in cases where it's a JSON due to `response_format` being a JSON
                         processed_chunk = AssistantMessage(
                             id=msg_obj.id,
                             date=msg_obj.created_at,
-                            content=func_args[self.assistant_message_tool_kwarg],
+                            content=str(func_args[self.assistant_message_tool_kwarg]),
                             name=msg_obj.name,
                             otid=Message.generate_otid_from_id(msg_obj.id, chunk_index) if chunk_index is not None else None,
                         )

letta/server/rest_api/routers/openai/chat_completions/chat_completions.py

@@ -111,7 +111,7 @@ async def send_message_to_agent_chat_completions(
                 server.send_messages,
                 actor=actor,
                 agent_id=letta_agent.agent_state.id,
-                messages=messages,
+                input_messages=messages,
                 interface=streaming_interface,
                 put_inner_thoughts_first=False,
             )

letta/server/rest_api/routers/v1/agents.py

@@ -412,7 +412,7 @@ def list_blocks(
     """
     actor = server.user_manager.get_user_or_default(user_id=actor_id)
     try:
-        agent = server.agent_manager.get_agent_by_id(agent_id, actor=actor)
+        agent = server.agent_manager.get_agent_by_id(agent_id, actor)
         return agent.memory.blocks
     except NoResultFound as e:
         raise HTTPException(status_code=404, detail=str(e))
@@ -640,7 +640,7 @@ async def send_message(
         result = await server.send_message_to_agent(
             agent_id=agent_id,
             actor=actor,
-            messages=request.messages,
+            input_messages=request.messages,
             stream_steps=False,
             stream_tokens=False,
             # Support for AssistantMessage
@@ -703,7 +703,7 @@ async def send_message_streaming(
         result = await server.send_message_to_agent(
             agent_id=agent_id,
             actor=actor,
-            messages=request.messages,
+            input_messages=request.messages,
             stream_steps=True,
             stream_tokens=request.stream_tokens,
             # Support for AssistantMessage
@@ -730,7 +730,7 @@ async def process_message_background(
         result = await server.send_message_to_agent(
             agent_id=agent_id,
             actor=actor,
-            messages=messages,
+            input_messages=messages,
             stream_steps=False,  # NOTE(matt)
             stream_tokens=False,
             use_assistant_message=use_assistant_message,

letta/server/rest_api/routers/v1/groups.py

@@ -128,7 +128,7 @@ async def send_group_message(
     result = await server.send_group_message_to_agent(
         group_id=group_id,
         actor=actor,
-        messages=request.messages,
+        input_messages=request.messages,
         stream_steps=False,
         stream_tokens=False,
         # Support for AssistantMessage
@@ -167,7 +167,7 @@ async def send_group_message_streaming(
     result = await server.send_group_message_to_agent(
         group_id=group_id,
         actor=actor,
-        messages=request.messages,
+        input_messages=request.messages,
         stream_steps=True,
         stream_tokens=request.stream_tokens,
         # Support for AssistantMessage

letta/server/rest_api/routers/v1/messages.py

@@ -1,16 +1,17 @@
 from typing import List, Optional
 
-from fastapi import APIRouter, Body, Depends, Header
+from fastapi import APIRouter, Body, Depends, Header, status
 from fastapi.exceptions import HTTPException
 from starlette.requests import Request
 
 from letta.agents.letta_agent_batch import LettaAgentBatch
 from letta.log import get_logger
 from letta.orm.errors import NoResultFound
-from letta.schemas.job import BatchJob, JobStatus, JobType
+from letta.schemas.job import BatchJob, JobStatus, JobType, JobUpdate
 from letta.schemas.letta_request import CreateBatch
 from letta.server.rest_api.utils import get_letta_server
 from letta.server.server import SyncServer
+from letta.settings import settings
 
 router = APIRouter(prefix="/messages", tags=["messages"])
 
@@ -43,19 +44,26 @@ async def create_messages_batch(
         if length > max_bytes:
             raise HTTPException(status_code=413, detail=f"Request too large ({length} bytes). Max is {max_bytes} bytes.")
 
-    try:
-        actor = server.user_manager.get_user_or_default(user_id=actor_id)
-
-        # Create a new job
-        batch_job = BatchJob(
-            user_id=actor.id,
-            status=JobStatus.created,
-            metadata={
-                "job_type": "batch_messages",
-            },
-            callback_url=str(payload.callback_url),
+    # Reject request if env var is not set
+    if not settings.enable_batch_job_polling:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Server misconfiguration: LETTA_ENABLE_BATCH_JOB_POLLING is set to False.",
         )
 
+    actor = server.user_manager.get_user_or_default(user_id=actor_id)
+    batch_job = BatchJob(
+        user_id=actor.id,
+        status=JobStatus.running,
+        metadata={
+            "job_type": "batch_messages",
+        },
+        callback_url=str(payload.callback_url),
+    )
+
+    try:
+        batch_job = server.job_manager.create_job(pydantic_job=batch_job, actor=actor)
+
         # create the batch runner
         batch_runner = LettaAgentBatch(
             message_manager=server.message_manager,
@@ -67,14 +75,17 @@ async def create_messages_batch(
             job_manager=server.job_manager,
             actor=actor,
         )
-        llm_batch_job = await batch_runner.step_until_request(batch_requests=payload.requests, letta_batch_job_id=batch_job.id)
+        await batch_runner.step_until_request(batch_requests=payload.requests, letta_batch_job_id=batch_job.id)
 
         # TODO: update run metadata
-        batch_job = server.job_manager.create_job(pydantic_job=batch_job, actor=actor)
-    except Exception:
+    except Exception as e:
         import traceback
 
+        print("Error creating batch job", e)
         traceback.print_exc()
+
+        # mark job as failed
+        server.job_manager.update_job_by_id(job_id=batch_job.id, job=BatchJob(status=JobStatus.failed), actor=actor)
         raise
     return batch_job
 
@@ -125,8 +136,19 @@ async def cancel_batch_run(
 
     try:
         job = server.job_manager.get_job_by_id(job_id=batch_id, actor=actor)
-        job.status = JobStatus.cancelled
-        server.job_manager.update_job_by_id(job_id=job, job=job)
-        # TODO: actually cancel it
+        job = server.job_manager.update_job_by_id(job_id=job.id, job_update=JobUpdate(status=JobStatus.cancelled), actor=actor)
+
+        # Get related llm batch jobs
+        llm_batch_jobs = server.batch_manager.list_llm_batch_jobs(letta_batch_id=job.id, actor=actor)
+        for llm_batch_job in llm_batch_jobs:
+            if llm_batch_job.status in {JobStatus.running, JobStatus.created}:
+                # TODO: Extend to providers beyond anthropic
+                # TODO: For now, we only support anthropic
+                # Cancel the job
+                anthropic_batch_id = llm_batch_job.create_batch_response.id
+                await server.anthropic_async_client.messages.batches.cancel(anthropic_batch_id)
+
+                # Update all the batch_job statuses
+                server.batch_manager.update_llm_batch_status(llm_batch_id=llm_batch_job.id, status=JobStatus.cancelled, actor=actor)
     except NoResultFound:
         raise HTTPException(status_code=404, detail="Run not found")

letta/server/server.py

@@ -28,7 +28,6 @@ from letta.functions.mcp_client.types import MCPServerType, MCPTool, SSEServerCo
 from letta.groups.helpers import load_multi_agent
 from letta.helpers.datetime_helpers import get_utc_time
 from letta.helpers.json_helpers import json_dumps, json_loads
-from letta.helpers.message_helper import prepare_input_message_create
 
 # TODO use custom interface
 from letta.interface import AgentInterface  # abstract
@@ -148,7 +147,7 @@ class Server(object):
         raise NotImplementedError
 
     @abstractmethod
-    def send_messages(self, user_id: str, agent_id: str, messages: Union[MessageCreate, List[Message]]) -> None:
+    def send_messages(self, user_id: str, agent_id: str, input_messages: List[MessageCreate]) -> None:
         """Send a list of messages to the agent"""
         raise NotImplementedError
 
@@ -372,19 +371,13 @@ class SyncServer(Server):
         self,
         actor: User,
         agent_id: str,
-        input_messages: Union[Message, List[Message]],
+        input_messages: List[MessageCreate],
         interface: Union[AgentInterface, None] = None,  # needed to getting responses
         put_inner_thoughts_first: bool = True,
         # timestamp: Optional[datetime],
     ) -> LettaUsageStatistics:
         """Send the input message through the agent"""
         # TODO: Thread actor directly through this function, since the top level caller most likely already retrieved the user
-        # Input validation
-        if isinstance(input_messages, Message):
-            input_messages = [input_messages]
-        if not all(isinstance(m, Message) for m in input_messages):
-            raise ValueError(f"messages should be a Message or a list of Message, got {type(input_messages)}")
-
         logger.debug(f"Got input messages: {input_messages}")
         letta_agent = None
         try:
@@ -400,8 +393,9 @@ class SyncServer(Server):
                 metadata = interface.metadata if hasattr(interface, "metadata") else None
             else:
                 metadata = None
+
             usage_stats = letta_agent.step(
-                messages=input_messages,
+                input_messages=input_messages,
                 chaining=self.chaining,
                 max_chaining_steps=self.max_chaining_steps,
                 stream=token_streaming,
@@ -572,23 +566,14 @@ class SyncServer(Server):
             )
 
             # NOTE: eventually deprecate and only allow passing Message types
-            # Convert to a Message object
-            if timestamp:
-                message = Message(
-                    agent_id=agent_id,
-                    role="user",
-                    content=[TextContent(text=packaged_user_message)],
-                    created_at=timestamp,
-                )
-            else:
-                message = Message(
-                    agent_id=agent_id,
-                    role="user",
-                    content=[TextContent(text=packaged_user_message)],
-                )
+            message = MessageCreate(
+                agent_id=agent_id,
+                role="user",
+                content=[TextContent(text=packaged_user_message)],
+            )
 
         # Run the agent state forward
-        usage = self._step(actor=actor, agent_id=agent_id, input_messages=message)
+        usage = self._step(actor=actor, agent_id=agent_id, input_messages=[message])
         return usage
 
     def system_message(
@@ -660,23 +645,14 @@ class SyncServer(Server):
         self,
         actor: User,
         agent_id: str,
-        messages: Union[List[MessageCreate], List[Message]],
+        input_messages: List[MessageCreate],
         wrap_user_message: bool = True,
         wrap_system_message: bool = True,
         interface: Union[AgentInterface, ChatCompletionsStreamingInterface, None] = None,  # needed for responses
         metadata: Optional[dict] = None,  # Pass through metadata to interface
         put_inner_thoughts_first: bool = True,
     ) -> LettaUsageStatistics:
-        """Send a list of messages to the agent.
-
-        If messages are of type MessageCreate, convert them to Message objects before sending.
-        """
-        if all(isinstance(m, MessageCreate) for m in messages):
-            message_objects = [prepare_input_message_create(m, agent_id, wrap_user_message, wrap_system_message) for m in messages]
-        elif all(isinstance(m, Message) for m in messages):
-            message_objects = messages
-        else:
-            raise ValueError(f"All messages must be of type Message or MessageCreate, got {[type(m) for m in messages]}")
+        """Send a list of messages to the agent."""
 
         # Store metadata in interface if provided
         if metadata and hasattr(interface, "metadata"):
@@ -686,7 +662,7 @@ class SyncServer(Server):
         return self._step(
             actor=actor,
             agent_id=agent_id,
-            input_messages=message_objects,
+            input_messages=input_messages,
             interface=interface,
             put_inner_thoughts_first=put_inner_thoughts_first,
         )
@@ -703,8 +679,6 @@ class SyncServer(Server):
     @trace_method
     def get_cached_llm_config(self, **kwargs):
         key = make_key(**kwargs)
-        print(self._llm_config_cache)
-        print("KEY", key)
         if key not in self._llm_config_cache:
             self._llm_config_cache[key] = self.get_llm_config_from_handle(**kwargs)
         return self._llm_config_cache[key]
@@ -1019,12 +993,8 @@ class SyncServer(Server):
         agent = self.load_agent(agent_id=sleeptime_agent.id, actor=actor)
         for passage in self.list_data_source_passages(source_id=source.id, user_id=actor.id):
             agent.step(
-                messages=[
-                    Message(
-                        role="user",
-                        content=[TextContent(text=passage.text)],
-                        agent_id=sleeptime_agent.id,
-                    ),
+                input_messages=[
+                    MessageCreate(role="user", content=passage.text),
                 ]
             )
         self.agent_manager.delete_agent(agent_id=sleeptime_agent.id, actor=actor)
@@ -1182,7 +1152,6 @@ class SyncServer(Server):
             provider = self.get_provider_from_name(provider_name)
 
             llm_configs = [config for config in provider.list_llm_models() if config.handle == handle]
-            print("LLM CONFIGS", llm_configs)
             if not llm_configs:
                 llm_configs = [config for config in provider.list_llm_models() if config.model == model_name]
             if not llm_configs:
@@ -1195,8 +1164,6 @@ class SyncServer(Server):
             if not llm_configs:
                 raise e
 
-        print("CONFIGS", llm_configs)
-
         if len(llm_configs) == 1:
             llm_config = llm_configs[0]
         elif len(llm_configs) > 1:
@@ -1343,17 +1310,17 @@ class SyncServer(Server):
 
         # Next, attempt to run the tool with the sandbox
         try:
-            sandbox_run_result = ToolExecutionSandbox(tool.name, tool_args, actor, tool_object=tool).run(
+            tool_execution_result = ToolExecutionSandbox(tool.name, tool_args, actor, tool_object=tool).run(
                 agent_state=agent_state, additional_env_vars=tool_env_vars
             )
             return ToolReturnMessage(
                 id="null",
                 tool_call_id="null",
                 date=get_utc_time(),
-                status=sandbox_run_result.status,
-                tool_return=str(sandbox_run_result.func_return),
-                stdout=sandbox_run_result.stdout,
-                stderr=sandbox_run_result.stderr,
+                status=tool_execution_result.status,
+                tool_return=str(tool_execution_result.func_return),
+                stdout=tool_execution_result.stdout,
+                stderr=tool_execution_result.stderr,
             )
 
         except Exception as e:
@@ -1567,7 +1534,7 @@ class SyncServer(Server):
         agent_id: str,
         actor: User,
         # role: MessageRole,
-        messages: Union[List[Message], List[MessageCreate]],
+        input_messages: List[MessageCreate],
         stream_steps: bool,
         stream_tokens: bool,
         # related to whether or not we return `LettaMessage`s or `Message`s
@@ -1647,7 +1614,7 @@ class SyncServer(Server):
                     self.send_messages,
                     actor=actor,
                     agent_id=agent_id,
-                    messages=messages,
+                    input_messages=input_messages,
                     interface=streaming_interface,
                     metadata=metadata,
                 )
@@ -1701,7 +1668,7 @@ class SyncServer(Server):
         self,
         group_id: str,
         actor: User,
-        messages: Union[List[Message], List[MessageCreate]],
+        input_messages: Union[List[Message], List[MessageCreate]],
         stream_steps: bool,
         stream_tokens: bool,
         chat_completion_mode: bool = False,
@@ -1751,7 +1718,7 @@ class SyncServer(Server):
         task = asyncio.create_task(
             asyncio.to_thread(
                 letta_multi_agent.step,
-                messages=messages,
+                input_messages=input_messages,
                 chaining=self.chaining,
                 max_chaining_steps=self.max_chaining_steps,
             )

letta/services/agent_manager.py

@@ -161,7 +161,7 @@ class AgentManager:
     # Basic CRUD operations
     # ======================================================================================================================
     @trace_method
-    def create_agent(self, agent_create: CreateAgent, actor: PydanticUser) -> PydanticAgentState:
+    def create_agent(self, agent_create: CreateAgent, actor: PydanticUser, _test_only_force_id: Optional[str] = None) -> PydanticAgentState:
         # validate required configs
         if not agent_create.llm_config or not agent_create.embedding_config:
             raise ValueError("llm_config and embedding_config are required")
@@ -239,6 +239,10 @@ class AgentManager:
                     created_by_id=actor.id,
                     last_updated_by_id=actor.id,
                 )
+
+                if _test_only_force_id:
+                    new_agent.id = _test_only_force_id
+
                 session.add(new_agent)
                 session.flush()
                 aid = new_agent.id
@@ -364,6 +368,7 @@ class AgentManager:
                 "base_template_id": agent_update.base_template_id,
                 "message_buffer_autoclear": agent_update.message_buffer_autoclear,
                 "enable_sleeptime": agent_update.enable_sleeptime,
+                "response_format": agent_update.response_format,
             }
             for col, val in scalar_updates.items():
                 if val is not None: