letta-nightly 0.7.29.dev20250602104315__py3-none-any.whl → 0.8.0.dev20250604104349__py3-none-any.whl

letta/functions/mcp_client/stdio_client.py

@@ -1,108 +1,109 @@
-import asyncio
-import sys
-from contextlib import asynccontextmanager
+# import asyncio
+# import sys
+# from contextlib import asynccontextmanager
+#
+# import anyio
+# import anyio.lowlevel
+# import mcp.types as types
+# from anyio.streams.text import TextReceiveStream
+# from mcp import ClientSession, StdioServerParameters
+# from mcp.client.stdio import get_default_environment
+#
+# from letta.functions.mcp_client.base_client import BaseMCPClient
+# from letta.functions.mcp_client.types import StdioServerConfig
+# from letta.log import get_logger
+#
+# logger = get_logger(__name__)
 
-import anyio
-import anyio.lowlevel
-import mcp.types as types
-from anyio.streams.text import TextReceiveStream
-from mcp import ClientSession, StdioServerParameters
-from mcp.client.stdio import get_default_environment
 
-from letta.functions.mcp_client.base_client import BaseMCPClient
-from letta.functions.mcp_client.types import StdioServerConfig
-from letta.log import get_logger
-
-logger = get_logger(__name__)
-
-
-class StdioMCPClient(BaseMCPClient):
-    def _initialize_connection(self, server_config: StdioServerConfig, timeout: float) -> bool:
-        try:
-            server_params = StdioServerParameters(command=server_config.command, args=server_config.args, env=server_config.env)
-            stdio_cm = forked_stdio_client(server_params)
-            stdio_transport = self.loop.run_until_complete(asyncio.wait_for(stdio_cm.__aenter__(), timeout=timeout))
-            self.stdio, self.write = stdio_transport
-            self.cleanup_funcs.append(lambda: self.loop.run_until_complete(stdio_cm.__aexit__(None, None, None)))
-
-            session_cm = ClientSession(self.stdio, self.write)
-            self.session = self.loop.run_until_complete(asyncio.wait_for(session_cm.__aenter__(), timeout=timeout))
-            self.cleanup_funcs.append(lambda: self.loop.run_until_complete(session_cm.__aexit__(None, None, None)))
-            return True
-        except asyncio.TimeoutError:
-            logger.error(f"Timed out while establishing stdio connection (timeout={timeout}s).")
-            return False
-        except Exception:
-            logger.exception("Exception occurred while initializing stdio client session.")
-            return False
-
-
-@asynccontextmanager
-async def forked_stdio_client(server: StdioServerParameters):
-    """
-    Client transport for stdio: this will connect to a server by spawning a
-    process and communicating with it over stdin/stdout.
-    """
-    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
-    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
-
-    try:
-        process = await anyio.open_process(
-            [server.command, *server.args],
-            env=server.env or get_default_environment(),
-            stderr=sys.stderr,  # Consider logging stderr somewhere instead of silencing it
-        )
-    except OSError as exc:
-        raise RuntimeError(f"Failed to spawn process: {server.command} {server.args}") from exc
-
-    async def stdout_reader():
-        assert process.stdout, "Opened process is missing stdout"
-        buffer = ""
-        try:
-            async with read_stream_writer:
-                async for chunk in TextReceiveStream(
-                    process.stdout,
-                    encoding=server.encoding,
-                    errors=server.encoding_error_handler,
-                ):
-                    lines = (buffer + chunk).split("\n")
-                    buffer = lines.pop()
-                    for line in lines:
-                        try:
-                            message = types.JSONRPCMessage.model_validate_json(line)
-                        except Exception as exc:
-                            await read_stream_writer.send(exc)
-                            continue
-                        await read_stream_writer.send(message)
-        except anyio.ClosedResourceError:
-            await anyio.lowlevel.checkpoint()
-
-    async def stdin_writer():
-        assert process.stdin, "Opened process is missing stdin"
-        try:
-            async with write_stream_reader:
-                async for message in write_stream_reader:
-                    json = message.model_dump_json(by_alias=True, exclude_none=True)
-                    await process.stdin.send(
-                        (json + "\n").encode(
-                            encoding=server.encoding,
-                            errors=server.encoding_error_handler,
-                        )
-                    )
-        except anyio.ClosedResourceError:
-            await anyio.lowlevel.checkpoint()
-
-    async def watch_process_exit():
-        returncode = await process.wait()
-        if returncode != 0:
-            raise RuntimeError(f"Subprocess exited with code {returncode}. Command: {server.command} {server.args}")
-
-    async with anyio.create_task_group() as tg, process:
-        tg.start_soon(stdout_reader)
-        tg.start_soon(stdin_writer)
-        tg.start_soon(watch_process_exit)
-
-        with anyio.move_on_after(0.2):
-            await anyio.sleep_forever()
-
-        yield read_stream, write_stream
+# class StdioMCPClient(BaseMCPClient):
+#    def _initialize_connection(self, server_config: StdioServerConfig, timeout: float) -> bool:
+#        try:
+#            server_params = StdioServerParameters(command=server_config.command, args=server_config.args, env=server_config.env)
+#            stdio_cm = forked_stdio_client(server_params)
+#            stdio_transport = self.loop.run_until_complete(asyncio.wait_for(stdio_cm.__aenter__(), timeout=timeout))
+#            self.stdio, self.write = stdio_transport
+#            self.cleanup_funcs.append(lambda: self.loop.run_until_complete(stdio_cm.__aexit__(None, None, None)))
+#
+#            session_cm = ClientSession(self.stdio, self.write)
+#            self.session = self.loop.run_until_complete(asyncio.wait_for(session_cm.__aenter__(), timeout=timeout))
+#            self.cleanup_funcs.append(lambda: self.loop.run_until_complete(session_cm.__aexit__(None, None, None)))
+#            return True
+#        except asyncio.TimeoutError:
+#            logger.error(f"Timed out while establishing stdio connection (timeout={timeout}s).")
+#            return False
+#        except Exception:
+#            logger.exception("Exception occurred while initializing stdio client session.")
+#            return False
+#
+#
+# @asynccontextmanager
+# async def forked_stdio_client(server: StdioServerParameters):
+#    """
+#    Client transport for stdio: this will connect to a server by spawning a
+#    process and communicating with it over stdin/stdout.
+#    """
+#    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+#    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+#
+#    try:
+#        process = await anyio.open_process(
+#            [server.command, *server.args],
+#            env=server.env or get_default_environment(),
+#            stderr=sys.stderr,  # Consider logging stderr somewhere instead of silencing it
+#        )
+#    except OSError as exc:
+#        raise RuntimeError(f"Failed to spawn process: {server.command} {server.args}") from exc
+#
+#    async def stdout_reader():
+#        assert process.stdout, "Opened process is missing stdout"
+#        buffer = ""
+#        try:
+#            async with read_stream_writer:
+#                async for chunk in TextReceiveStream(
+#                    process.stdout,
+#                    encoding=server.encoding,
+#                    errors=server.encoding_error_handler,
+#                ):
+#                    lines = (buffer + chunk).split("\n")
+#                    buffer = lines.pop()
+#                    for line in lines:
+#                        try:
+#                            message = types.JSONRPCMessage.model_validate_json(line)
+#                        except Exception as exc:
+#                            await read_stream_writer.send(exc)
+#                            continue
+#                        await read_stream_writer.send(message)
+#        except anyio.ClosedResourceError:
+#            await anyio.lowlevel.checkpoint()
+#
+#    async def stdin_writer():
+#        assert process.stdin, "Opened process is missing stdin"
+#        try:
+#            async with write_stream_reader:
+#                async for message in write_stream_reader:
+#                    json = message.model_dump_json(by_alias=True, exclude_none=True)
+#                    await process.stdin.send(
+#                        (json + "\n").encode(
+#                            encoding=server.encoding,
+#                            errors=server.encoding_error_handler,
+#                        )
+#                    )
+#        except anyio.ClosedResourceError:
+#            await anyio.lowlevel.checkpoint()
+#
+#    async def watch_process_exit():
+#        returncode = await process.wait()
+#        if returncode != 0:
+#            raise RuntimeError(f"Subprocess exited with code {returncode}. Command: {server.command} {server.args}")
+#
+#    async with anyio.create_task_group() as tg, process:
+#        tg.start_soon(stdout_reader)
+#        tg.start_soon(stdin_writer)
+#        tg.start_soon(watch_process_exit)
+#
+#        with anyio.move_on_after(0.2):
+#            await anyio.sleep_forever()
+#
+#        yield read_stream, write_stream
+#

letta/functions/schema_generator.py

@@ -7,6 +7,7 @@ from docstring_parser import parse
 from pydantic import BaseModel
 from typing_extensions import Literal
 
+from letta.constants import REQUEST_HEARTBEAT_DESCRIPTION, REQUEST_HEARTBEAT_PARAM
 from letta.functions.mcp_client.types import MCPTool
 
 
@@ -143,7 +144,10 @@ def pydantic_model_to_open_ai(model: Type[BaseModel]) -> dict:
     parameters["required"] = sorted(k for k, v in parameters["properties"].items() if "default" not in v)
 
     if "description" not in schema:
-        if docstring.short_description:
+        # Support multiline docstrings for complex functions, TODO (cliandy): consider having this as a setting
+        if docstring.long_description:
+            schema["description"] = docstring.long_description
+        elif docstring.short_description:
             schema["description"] = docstring.short_description
         else:
             raise ValueError(f"No description found in docstring or description field (model: {model}, docstring: {docstring})")
@@ -330,10 +334,17 @@ def generate_schema(function, name: Optional[str] = None, description: Optional[
     # Parse the docstring
     docstring = parse(function.__doc__)
 
+    if not description:
+        # Support multiline docstrings for complex functions, TODO (cliandy): consider having this as a setting
+        if docstring.long_description:
+            description = docstring.long_description
+        else:
+            description = docstring.short_description
+
     # Prepare the schema dictionary
     schema = {
         "name": function.__name__ if name is None else name,
-        "description": docstring.short_description if description is None else description,
+        "description": description,
         "parameters": {"type": "object", "properties": {}, "required": []},
     }
 
@@ -412,17 +423,6 @@ def generate_schema(function, name: Optional[str] = None, description: Optional[
         # TODO is this not duplicating the other append directly above?
         if param.annotation == inspect.Parameter.empty:
             schema["parameters"]["required"].append(param.name)
-
-    # append the heartbeat
-    # TODO: don't hard-code
-    # TODO: if terminal, don't include this
-    # if function.__name__ not in ["send_message"]:
-    schema["parameters"]["properties"]["request_heartbeat"] = {
-        "type": "boolean",
-        "description": "Request an immediate heartbeat after function execution. Set to `True` if you want to send a follow-up message or run a follow-up function.",
-    }
-    schema["parameters"]["required"].append("request_heartbeat")
-
     return schema
 
 
@@ -445,11 +445,11 @@ def generate_schema_from_args_schema_v2(
     }
 
     if append_heartbeat:
-        function_call_json["parameters"]["properties"]["request_heartbeat"] = {
+        function_call_json["parameters"]["properties"][REQUEST_HEARTBEAT_PARAM] = {
             "type": "boolean",
-            "description": "Request an immediate heartbeat after function execution. Set to `True` if you want to send a follow-up message or run a follow-up function.",
+            "description": REQUEST_HEARTBEAT_DESCRIPTION,
         }
-        function_call_json["parameters"]["required"].append("request_heartbeat")
+        function_call_json["parameters"]["required"].append(REQUEST_HEARTBEAT_PARAM)
 
     return function_call_json
 
@@ -476,11 +476,11 @@ def generate_tool_schema_for_mcp(
 
     # Add the optional heartbeat parameter
     if append_heartbeat:
-        parameters_schema["properties"]["request_heartbeat"] = {
+        parameters_schema["properties"][REQUEST_HEARTBEAT_PARAM] = {
             "type": "boolean",
-            "description": "Request an immediate heartbeat after function execution. Set to `True` if you want to send a follow-up message or run a follow-up function.",
+            "description": REQUEST_HEARTBEAT_DESCRIPTION,
         }
-        parameters_schema["required"].append("request_heartbeat")
+        parameters_schema["required"].append(REQUEST_HEARTBEAT_PARAM)
 
     # Return the final schema
     if strict:
@@ -538,11 +538,11 @@ def generate_tool_schema_for_composio(
 
     # Add the optional heartbeat parameter
     if append_heartbeat:
-        properties_json["request_heartbeat"] = {
+        properties_json[REQUEST_HEARTBEAT_PARAM] = {
             "type": "boolean",
-            "description": "Request an immediate heartbeat after function execution. Set to `True` if you want to send a follow-up message or run a follow-up function.",
+            "description": REQUEST_HEARTBEAT_DESCRIPTION,
         }
-        required_fields.append("request_heartbeat")
+        required_fields.append(REQUEST_HEARTBEAT_PARAM)
 
     # Return the final schema
     if strict:

letta/groups/helpers.py

@@ -2,13 +2,13 @@ import json
 from typing import Dict, Optional, Union
 
 from letta.agent import Agent
-from letta.functions.mcp_client.base_client import BaseMCPClient
 from letta.interface import AgentInterface
 from letta.orm.group import Group
 from letta.orm.user import User
 from letta.schemas.agent import AgentState
 from letta.schemas.group import ManagerType
 from letta.schemas.message import Message
+from letta.services.mcp.base_client import AsyncBaseMCPClient
 
 
 def load_multi_agent(
@@ -16,7 +16,7 @@ def load_multi_agent(
     agent_state: Optional[AgentState],
     actor: User,
     interface: Union[AgentInterface, None] = None,
-    mcp_clients: Optional[Dict[str, BaseMCPClient]] = None,
+    mcp_clients: Optional[Dict[str, AsyncBaseMCPClient]] = None,
 ) -> Agent:
     if len(group.agent_ids) == 0:
         raise ValueError("Empty group: group must have at least one agent")
@@ -76,7 +76,6 @@ def load_multi_agent(
                 agent_state=agent_state,
                 interface=interface,
                 user=actor,
-                mcp_clients=mcp_clients,
                 group_id=group.id,
                 agent_ids=group.agent_ids,
                 description=group.description,
@@ -108,7 +107,7 @@ def stringify_message(message: Message, use_assistant_name: bool = False) -> str
     elif message.role == "tool":
         if message.content:
             content = json.loads(message.content[0].text)
-            if content["message"] != "None" and content["message"] != None:
+            if str(content["message"]) != "None":
                 return f"{assistant_name}: Tool call returned {content['message']}"
         return None
     elif message.role == "system":

letta/groups/sleeptime_multi_agent.py

@@ -1,10 +1,9 @@
 import asyncio
 import threading
 from datetime import datetime, timezone
-from typing import Dict, List, Optional
+from typing import List, Optional
 
 from letta.agent import Agent, AgentState
-from letta.functions.mcp_client.base_client import BaseMCPClient
 from letta.groups.helpers import stringify_message
 from letta.interface import AgentInterface
 from letta.orm import User
@@ -27,7 +26,7 @@ class SleeptimeMultiAgent(Agent):
         interface: AgentInterface,
         agent_state: AgentState,
         user: User,
-        mcp_clients: Optional[Dict[str, BaseMCPClient]] = None,
+        # mcp_clients: Optional[Dict[str, BaseMCPClient]] = None,
         # custom
         group_id: str = "",
         agent_ids: List[str] = [],
@@ -42,7 +41,8 @@ class SleeptimeMultiAgent(Agent):
         self.group_manager = GroupManager()
         self.message_manager = MessageManager()
         self.job_manager = JobManager()
-        self.mcp_clients = mcp_clients
+        # TODO: add back MCP support with new agent loop
+        self.mcp_clients = {}
 
     def _run_async_in_new_thread(self, coro):
         """Run an async coroutine in a new thread with its own event loop"""

letta/groups/sleeptime_multi_agent_v2.py

@@ -21,6 +21,7 @@ from letta.services.message_manager import MessageManager
 from letta.services.passage_manager import PassageManager
 from letta.services.step_manager import NoopStepManager, StepManager
 from letta.services.telemetry_manager import NoopTelemetryManager, TelemetryManager
+from letta.tracing import trace_method
 
 
 class SleeptimeMultiAgentV2(BaseAgent):
@@ -55,11 +56,13 @@ class SleeptimeMultiAgentV2(BaseAgent):
         assert group.manager_type == ManagerType.sleeptime, f"Expected group manager type to be 'sleeptime', got {group.manager_type}"
         self.group = group
 
+    @trace_method
     async def step(
         self,
         input_messages: List[MessageCreate],
         max_steps: int = 10,
         use_assistant_message: bool = True,
+        request_start_timestamp_ns: Optional[int] = None,
     ) -> LettaResponse:
         run_ids = []
 
@@ -119,6 +122,22 @@ class SleeptimeMultiAgentV2(BaseAgent):
         response.usage.run_ids = run_ids
         return response
 
+    @trace_method
+    async def step_stream_no_tokens(
+        self,
+        input_messages: List[MessageCreate],
+        max_steps: int = 10,
+        use_assistant_message: bool = True,
+        request_start_timestamp_ns: Optional[int] = None,
+    ):
+        response = await self.step(input_messages, max_steps, use_assistant_message)
+
+        for message in response.messages:
+            yield f"data: {message.model_dump_json()}\n\n"
+
+        yield f"data: {response.usage.model_dump_json()}\n\n"
+
+    @trace_method
     async def step_stream(
         self,
         input_messages: List[MessageCreate],
@@ -256,6 +275,9 @@ class SleeptimeMultiAgentV2(BaseAgent):
                 actor=self.actor,
                 step_manager=self.step_manager,
                 telemetry_manager=self.telemetry_manager,
+                message_buffer_limit=20,  # TODO: Make this configurable
+                message_buffer_min=8,  # TODO: Make this configurable
+                enable_summarization=False,  # TODO: Make this configurable
             )
 
             # Perform sleeptime agent step

letta/helpers/composio_helpers.py

@@ -20,3 +20,19 @@ def get_composio_api_key(actor: User, logger: Optional[Logger] = None) -> Option
         # Ideally, not tied to a specific sandbox, but for now we just get the first one
         # Theoretically possible for someone to have different composio api keys per sandbox
         return api_keys[0].value
+
+
+async def get_composio_api_key_async(actor: User, logger: Optional[Logger] = None) -> Optional[str]:
+    api_keys = await SandboxConfigManager().list_sandbox_env_vars_by_key_async(key="COMPOSIO_API_KEY", actor=actor)
+    if not api_keys:
+        if logger:
+            logger.debug(f"No API keys found for Composio. Defaulting to the environment variable...")
+        if tool_settings.composio_api_key:
+            return tool_settings.composio_api_key
+        else:
+            return None
+    else:
+        # TODO: Add more protections around this
+        # Ideally, not tied to a specific sandbox, but for now we just get the first one
+        # Theoretically possible for someone to have different composio api keys per sandbox
+        return api_keys[0].value

letta/helpers/converters.py

@@ -7,6 +7,7 @@ from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMe
 from openai.types.chat.chat_completion_message_tool_call import Function as OpenAIFunction
 from sqlalchemy import Dialect
 
+from letta.functions.mcp_client.types import StdioServerConfig
 from letta.schemas.agent import AgentStepState
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import ProviderType, ToolRuleType
@@ -400,3 +401,22 @@ def deserialize_response_format(data: Optional[Dict]) -> Optional[ResponseFormat
         return JsonSchemaResponseFormat(**data)
     if data["type"] == ResponseFormatType.json_object:
         return JsonObjectResponseFormat(**data)
+
+
+# --------------------------
+# MCP Stdio Server Config Serialization
+# --------------------------
+
+
+def serialize_mcp_stdio_config(config: Union[Optional[StdioServerConfig], Dict]) -> Optional[Dict]:
+    """Convert an StdioServerConfig object into a JSON-serializable dictionary."""
+    if config and isinstance(config, StdioServerConfig):
+        return config.to_dict()
+    return config
+
+
+def deserialize_mcp_stdio_config(data: Optional[Dict]) -> Optional[StdioServerConfig]:
+    """Convert a dictionary back into an StdioServerConfig object."""
+    if not data:
+        return None
+    return StdioServerConfig(**data)

letta/helpers/datetime_helpers.py

@@ -16,11 +16,6 @@ def datetime_to_timestamp(dt):
     return int(dt.timestamp())
 
 
-def timestamp_to_datetime(ts):
-    # convert integer timestamp to datetime object
-    return datetime.fromtimestamp(ts)
-
-
 def get_local_time_military():
     # Get the current time in UTC
     current_time_utc = datetime.now(pytz.utc)
@@ -36,7 +31,7 @@ def get_local_time_military():
 
 
 def get_local_time_fast():
-    formatted_time = strftime("%Y-%m-%d %H:%M:%S")
+    formatted_time = strftime("%Y-%m-%d %I:%M:%S %p %Z%z")
 
     return formatted_time
 

letta/helpers/tool_rule_solver.py

@@ -141,7 +141,8 @@ class ToolRulesSolver(BaseModel):
         """Check if the tool is defined as a continue tool in the tool rules."""
         return any(rule.tool_name == tool_name for rule in self.continue_tool_rules)
 
-    def validate_conditional_tool(self, rule: ConditionalToolRule):
+    @staticmethod
+    def validate_conditional_tool(rule: ConditionalToolRule):
         """
         Validate a conditional tool rule
 

letta/interfaces/anthropic_streaming_interface.py

@@ -1,7 +1,7 @@
 import json
 from datetime import datetime, timezone
 from enum import Enum
-from typing import AsyncGenerator, List, Union
+from typing import AsyncGenerator, List, Optional, Union
 
 from anthropic import AsyncStream
 from anthropic.types.beta import (
@@ -23,6 +23,7 @@ from anthropic.types.beta import (
 )
 
 from letta.constants import DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
+from letta.helpers.datetime_helpers import get_utc_timestamp_ns
 from letta.local_llm.constants import INNER_THOUGHTS_KWARG
 from letta.log import get_logger
 from letta.schemas.letta_message import (
@@ -115,12 +116,26 @@ class AnthropicStreamingInterface:
             logger.error("Error checking inner thoughts: %s", e)
             raise
 
-    async def process(self, stream: AsyncStream[BetaRawMessageStreamEvent]) -> AsyncGenerator[LettaMessage, None]:
+    async def process(
+        self,
+        stream: AsyncStream[BetaRawMessageStreamEvent],
+        ttft_span: Optional["Span"] = None,
+        provider_request_start_timestamp_ns: Optional[int] = None,
+    ) -> AsyncGenerator[LettaMessage, None]:
         prev_message_type = None
         message_index = 0
+        first_chunk = True
         try:
             async with stream:
                 async for event in stream:
+                    if first_chunk and ttft_span is not None and provider_request_start_timestamp_ns is not None:
+                        now = get_utc_timestamp_ns()
+                        ttft_ns = now - provider_request_start_timestamp_ns
+                        ttft_span.add_event(
+                            name="anthropic_time_to_first_token_ms", attributes={"anthropic_time_to_first_token_ms": ttft_ns // 1_000_000}
+                        )
+                        first_chunk = False
+
                     # TODO: Support BetaThinkingBlock, BetaRedactedThinkingBlock
                     if isinstance(event, BetaRawContentBlockStartEvent):
                         content = event.content_block

letta/interfaces/openai_chat_completions_streaming_interface.py

@@ -16,6 +16,7 @@ class OpenAIChatCompletionsStreamingInterface:
     """
 
     def __init__(self, stream_pre_execution_message: bool = True):
+        print("CHAT COMPLETITION INTERFACE")
         self.optimistic_json_parser: OptimisticJSONParser = OptimisticJSONParser()
         self.stream_pre_execution_message: bool = stream_pre_execution_message
 

letta/interfaces/openai_streaming_interface.py

@@ -5,6 +5,7 @@ from openai import AsyncStream
 from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
 
 from letta.constants import DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
+from letta.helpers.datetime_helpers import get_utc_timestamp_ns
 from letta.schemas.letta_message import AssistantMessage, LettaMessage, ReasoningMessage, ToolCallDelta, ToolCallMessage
 from letta.schemas.letta_message_content import TextContent
 from letta.schemas.message import Message
@@ -26,7 +27,7 @@ class OpenAIStreamingInterface:
         self.assistant_message_tool_kwarg = DEFAULT_MESSAGE_TOOL_KWARG
 
         self.optimistic_json_parser: OptimisticJSONParser = OptimisticJSONParser()
-        self.function_args_reader = JSONInnerThoughtsExtractor(wait_for_first_key=True)  # TODO: pass in kward
+        self.function_args_reader = JSONInnerThoughtsExtractor(wait_for_first_key=True)  # TODO: pass in kwarg
         self.function_name_buffer = None
         self.function_args_buffer = None
         self.function_id_buffer = None
@@ -64,15 +65,30 @@ class OpenAIStreamingInterface:
             function=FunctionCall(arguments=self.current_function_arguments, name=function_name),
         )
 
-    async def process(self, stream: AsyncStream[ChatCompletionChunk]) -> AsyncGenerator[LettaMessage, None]:
+    async def process(
+        self,
+        stream: AsyncStream[ChatCompletionChunk],
+        ttft_span: Optional["Span"] = None,
+        provider_request_start_timestamp_ns: Optional[int] = None,
+    ) -> AsyncGenerator[LettaMessage, None]:
         """
         Iterates over the OpenAI stream, yielding SSE events.
         It also collects tokens and detects if a tool call is triggered.
         """
+        first_chunk = True
+
         async with stream:
             prev_message_type = None
             message_index = 0
             async for chunk in stream:
+                if first_chunk and ttft_span is not None and provider_request_start_timestamp_ns is not None:
+                    now = get_utc_timestamp_ns()
+                    ttft_ns = now - provider_request_start_timestamp_ns
+                    ttft_span.add_event(
+                        name="openai_time_to_first_token_ms", attributes={"openai_time_to_first_token_ms": ttft_ns // 1_000_000}
+                    )
+                    first_chunk = False
+
                 if not self.model or not self.message_id:
                     self.model = chunk.model
                     self.message_id = chunk.id

letta/jobs/llm_batch_job_polling.py

@@ -238,4 +238,4 @@ async def poll_running_llm_batches(server: "SyncServer") -> List[LettaBatchRespo
         logger.exception("[Poll BatchJob] Unhandled error in poll_running_llm_batches", exc_info=e)
     finally:
         # 7. Log metrics summary
-        metrics.log_summary()
+        metrics.log_summary()

letta/jobs/scheduler.py

@@ -252,4 +252,4 @@ async def shutdown_scheduler_and_release_lock():
         try:
             scheduler.shutdown(wait=False)
         except:
-            pass
+            pass