letta-nightly 0.8.0.dev20250606195656__py3-none-any.whl → 0.8.3.dev20250607000559__py3-none-any.whl

letta/server/rest_api/app.py

@@ -256,13 +256,15 @@ def create_application() -> "FastAPI":
         print(f"▶ Using OTLP tracing with endpoint: {otlp_endpoint}")
         env_name_suffix = os.getenv("ENV_NAME")
         service_name = f"letta-server-{env_name_suffix.lower()}" if env_name_suffix else "letta-server"
-        from letta.tracing import setup_tracing
+        from letta.otel.metrics import setup_metrics
+        from letta.otel.tracing import setup_tracing
 
         setup_tracing(
             endpoint=otlp_endpoint,
             app=app,
             service_name=service_name,
         )
+        setup_metrics(endpoint=otlp_endpoint, app=app, service_name=service_name)
 
     for route in v1_routes:
         app.include_router(route, prefix=API_PREFIX)
@@ -331,7 +333,7 @@ def start_server(
     if (os.getenv("LOCAL_HTTPS") == "true") or "--localhttps" in sys.argv:
         print(f"▶ Server running at: https://{host or 'localhost'}:{port or REST_DEFAULT_PORT}")
         print(f"▶ View using ADE at: https://app.letta.com/development-servers/local/dashboard\n")
-        if importlib.util.find_spec("granian") is not None and settings.use_uvloop:
+        if importlib.util.find_spec("granian") is not None and settings.use_granian:
             from granian import Granian
 
             # Experimental Granian engine
@@ -339,14 +341,14 @@ def start_server(
                 target="letta.server.rest_api.app:app",
                 # factory=True,
                 interface="asgi",
-                address=host or "localhost",
+                address=host or "127.0.0.1",  # Note granian address must be an ip address
                 port=port or REST_DEFAULT_PORT,
                 workers=settings.uvicorn_workers,
                 # threads=
                 reload=reload or settings.uvicorn_reload,
                 reload_ignore_patterns=["openapi_letta.json"],
                 reload_ignore_worker_failure=True,
-                reload_tick=100,
+                reload_tick=4000,  # set to 4s to prevent crashing on weird state
                 # log_level="info"
                 ssl_keyfile="certs/localhost-key.pem",
                 ssl_cert="certs/localhost.pem",
@@ -380,14 +382,14 @@ def start_server(
                 target="letta.server.rest_api.app:app",
                 # factory=True,
                 interface="asgi",
-                address=host or "localhost",
+                address=host or "127.0.0.1",  # Note granian address must be an ip address
                 port=port or REST_DEFAULT_PORT,
                 workers=settings.uvicorn_workers,
                 # threads=
                 reload=reload or settings.uvicorn_reload,
                 reload_ignore_patterns=["openapi_letta.json"],
                 reload_ignore_worker_failure=True,
-                reload_tick=100,
+                reload_tick=4000,  # set to 4s to prevent crashing on weird state
                 # log_level="info"
             ).serve()
         else:

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

@@ -12,16 +12,18 @@ from sqlalchemy.exc import IntegrityError, OperationalError
 from starlette.responses import Response, StreamingResponse
 
 from letta.agents.letta_agent import LettaAgent
-from letta.constants import CORE_MEMORY_SOURCE_CHAR_LIMIT, DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
+from letta.constants import DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
 from letta.groups.sleeptime_multi_agent_v2 import SleeptimeMultiAgentV2
 from letta.helpers.datetime_helpers import get_utc_timestamp_ns
 from letta.log import get_logger
 from letta.orm.errors import NoResultFound
+from letta.otel.context import get_ctx_attributes
+from letta.otel.metric_registry import MetricRegistry
 from letta.schemas.agent import AgentState, AgentType, CreateAgent, UpdateAgent
 from letta.schemas.block import Block, BlockUpdate
 from letta.schemas.group import Group
 from letta.schemas.job import JobStatus, JobUpdate, LettaRequestConfig
-from letta.schemas.letta_message import LettaMessageUnion, LettaMessageUpdateUnion
+from letta.schemas.letta_message import LettaMessageUnion, LettaMessageUpdateUnion, MessageType
 from letta.schemas.letta_request import LettaRequest, LettaStreamingRequest
 from letta.schemas.letta_response import LettaResponse
 from letta.schemas.memory import ContextWindowOverview, CreateArchivalMemory, Memory
@@ -149,7 +151,7 @@ def export_agent_serialized(
 
 
 @router.post("/import", response_model=AgentState, operation_id="import_agent_serialized")
-async def import_agent_serialized(
+def import_agent_serialized(
     file: UploadFile = File(...),
     server: "SyncServer" = Depends(get_letta_server),
     actor_id: Optional[str] = Header(None, alias="user_id"),
@@ -167,10 +169,10 @@ async def import_agent_serialized(
     """
     Import a serialized agent file and recreate the agent in the system.
     """
-    actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
+    actor = server.user_manager.get_user_or_default(user_id=actor_id)
 
     try:
-        serialized_data = await file.read()
+        serialized_data = file.file.read()
         agent_json = json.loads(serialized_data)
 
         # Validate the JSON against AgentSchema before passing it to deserialize
@@ -311,20 +313,21 @@ async def attach_source(
     actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
     agent_state = await server.agent_manager.attach_source_async(agent_id=agent_id, source_id=source_id, actor=actor)
 
-    files = await server.source_manager.list_files(source_id, actor)
+    # Check if the agent is missing any files tools
+    agent_state = await server.agent_manager.attach_missing_files_tools_async(agent_state=agent_state, actor=actor)
+
+    files = await server.source_manager.list_files(source_id, actor, include_content=True)
     texts = []
     file_ids = []
+    file_names = []
     for f in files:
-        passages = await server.passage_manager.list_passages_by_file_id_async(file_id=f.id, actor=actor)
-        passage_text = ""
-        for p in passages:
-            if len(passage_text) <= CORE_MEMORY_SOURCE_CHAR_LIMIT:
-                passage_text += p.text
-
-        texts.append(passage_text)
+        texts.append(f.content if f.content else "")
         file_ids.append(f.id)
+        file_names.append(f.file_name)
 
-    await server.insert_files_into_context_window(agent_state=agent_state, texts=texts, file_ids=file_ids, actor=actor)
+    await server.insert_files_into_context_window(
+        agent_state=agent_state, texts=texts, file_ids=file_ids, file_names=file_names, actor=actor
+    )
 
     if agent_state.enable_sleeptime:
         source = await server.source_manager.get_source_by_id(source_id=source_id)
@@ -347,6 +350,10 @@ async def detach_source(
     """
     actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
     agent_state = await server.agent_manager.detach_source_async(agent_id=agent_id, source_id=source_id, actor=actor)
+
+    if not agent_state.sources:
+        agent_state = await server.agent_manager.detach_all_files_tools_async(agent_state=agent_state, actor=actor)
+
     files = await server.source_manager.list_files(source_id, actor)
     file_ids = [f.id for f in files]
     await server.remove_files_from_context_window(agent_state=agent_state, file_ids=file_ids, actor=actor)
@@ -451,7 +458,7 @@ async def list_blocks(
     """
     Retrieve the core memory blocks of a specific agent.
     """
-    actor = server.user_manager.get_user_or_default(user_id=actor_id)
+    actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
     try:
         agent = await server.agent_manager.get_agent_by_id_async(agent_id=agent_id, include_relationships=["memory"], actor=actor)
         return agent.memory.blocks
@@ -658,19 +665,18 @@ async def send_message(
     Process a user message and return the agent's response.
     This endpoint accepts a message from a user and processes it through the agent.
     """
+    MetricRegistry().user_message_counter.add(1, get_ctx_attributes())
+
     actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
     request_start_timestamp_ns = get_utc_timestamp_ns()
-    user_eligible = True
     # TODO: This is redundant, remove soon
     agent = await server.agent_manager.get_agent_by_id_async(agent_id, actor, include_relationships=["multi_agent_group"])
     agent_eligible = agent.enable_sleeptime or agent.agent_type == AgentType.sleeptime_agent or not agent.multi_agent_group
-    experimental_header = request_obj.headers.get("X-EXPERIMENTAL") or "false"
-    feature_enabled = settings.use_experimental or experimental_header.lower() == "true"
     model_compatible = agent.llm_config.model_endpoint_type in ["anthropic", "openai", "together", "google_ai", "google_vertex"]
 
-    if user_eligible and agent_eligible and feature_enabled and model_compatible:
+    if agent_eligible and model_compatible:
         if agent.enable_sleeptime and agent.agent_type != AgentType.voice_convo_agent:
-            experimental_agent = SleeptimeMultiAgentV2(
+            agent_loop = SleeptimeMultiAgentV2(
                 agent_id=agent_id,
                 message_manager=server.message_manager,
                 agent_manager=server.agent_manager,
@@ -682,7 +688,7 @@ async def send_message(
                 group=agent.multi_agent_group,
             )
         else:
-            experimental_agent = LettaAgent(
+            agent_loop = LettaAgent(
                 agent_id=agent_id,
                 message_manager=server.message_manager,
                 agent_manager=server.agent_manager,
@@ -693,11 +699,12 @@ async def send_message(
                 telemetry_manager=server.telemetry_manager if settings.llm_api_logging else NoopTelemetryManager(),
             )
 
-        result = await experimental_agent.step(
+        result = await agent_loop.step(
             request.messages,
             max_steps=10,
             use_assistant_message=request.use_assistant_message,
             request_start_timestamp_ns=request_start_timestamp_ns,
+            include_return_message_types=request.include_return_message_types,
         )
     else:
         result = await server.send_message_to_agent(
@@ -710,6 +717,7 @@ async def send_message(
             use_assistant_message=request.use_assistant_message,
             assistant_message_tool_name=request.assistant_message_tool_name,
             assistant_message_tool_kwarg=request.assistant_message_tool_kwarg,
+            include_return_message_types=request.include_return_message_types,
         )
     return result
 
@@ -739,22 +747,20 @@ async def send_message_streaming(
     This endpoint accepts a message from a user and processes it through the agent.
     It will stream the steps of the response always, and stream the tokens if 'stream_tokens' is set to True.
     """
-    request_start_timestamp_ns = get_utc_timestamp_ns()
+    MetricRegistry().user_message_counter.add(1, get_ctx_attributes())
+
     actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
-    user_eligible = actor.organization_id not in ["org-4a3af5dd-4c6a-48cb-ac13-3f73ecaaa4bf", "org-4ab3f6e8-9a44-4bee-aeb6-c681cbbc7bf6"]
     # TODO: This is redundant, remove soon
     agent = await server.agent_manager.get_agent_by_id_async(agent_id, actor, include_relationships=["multi_agent_group"])
     agent_eligible = agent.enable_sleeptime or agent.agent_type == AgentType.sleeptime_agent or not agent.multi_agent_group
-    experimental_header = request_obj.headers.get("X-EXPERIMENTAL") or "false"
-    feature_enabled = settings.use_experimental or experimental_header.lower() == "true"
     model_compatible = agent.llm_config.model_endpoint_type in ["anthropic", "openai", "together", "google_ai", "google_vertex"]
     model_compatible_token_streaming = agent.llm_config.model_endpoint_type in ["anthropic", "openai"]
     not_letta_endpoint = not ("inference.letta.com" in agent.llm_config.model_endpoint)
     request_start_timestamp_ns = get_utc_timestamp_ns()
 
-    if user_eligible and agent_eligible and feature_enabled and model_compatible:
+    if agent_eligible and model_compatible:
         if agent.enable_sleeptime and agent.agent_type != AgentType.voice_convo_agent:
-            experimental_agent = SleeptimeMultiAgentV2(
+            agent_loop = SleeptimeMultiAgentV2(
                 agent_id=agent_id,
                 message_manager=server.message_manager,
                 agent_manager=server.agent_manager,
@@ -768,7 +774,7 @@ async def send_message_streaming(
                 group=agent.multi_agent_group,
             )
         else:
-            experimental_agent = LettaAgent(
+            agent_loop = LettaAgent(
                 agent_id=agent_id,
                 message_manager=server.message_manager,
                 agent_manager=server.agent_manager,
@@ -782,21 +788,23 @@ async def send_message_streaming(
 
         if request.stream_tokens and model_compatible_token_streaming and not_letta_endpoint:
             result = StreamingResponseWithStatusCode(
-                experimental_agent.step_stream(
+                agent_loop.step_stream(
                     input_messages=request.messages,
                     max_steps=10,
                     use_assistant_message=request.use_assistant_message,
                     request_start_timestamp_ns=request_start_timestamp_ns,
+                    include_return_message_types=request.include_return_message_types,
                 ),
                 media_type="text/event-stream",
             )
         else:
             result = StreamingResponseWithStatusCode(
-                experimental_agent.step_stream_no_tokens(
+                agent_loop.step_stream_no_tokens(
                     request.messages,
                     max_steps=10,
                     use_assistant_message=request.use_assistant_message,
                     request_start_timestamp_ns=request_start_timestamp_ns,
+                    include_return_message_types=request.include_return_message_types,
                 ),
                 media_type="text/event-stream",
             )
@@ -812,6 +820,7 @@ async def send_message_streaming(
             assistant_message_tool_name=request.assistant_message_tool_name,
             assistant_message_tool_kwarg=request.assistant_message_tool_kwarg,
             request_start_timestamp_ns=request_start_timestamp_ns,
+            include_return_message_types=request.include_return_message_types,
         )
 
     return result
@@ -826,6 +835,7 @@ async def process_message_background(
     use_assistant_message: bool,
     assistant_message_tool_name: str,
     assistant_message_tool_kwarg: str,
+    include_return_message_types: Optional[List[MessageType]] = None,
 ) -> None:
     """Background task to process the message and update job status."""
     try:
@@ -841,6 +851,7 @@ async def process_message_background(
             assistant_message_tool_kwarg=assistant_message_tool_kwarg,
             metadata={"job_id": job_id},  # Pass job_id through metadata
             request_start_timestamp_ns=request_start_timestamp_ns,
+            include_return_message_types=include_return_message_types,
         )
 
         # Update job status to completed
@@ -878,6 +889,7 @@ async def send_message_async(
     Asynchronously process a user message and return a run object.
     The actual processing happens in the background, and the status can be checked using the run ID.
     """
+    MetricRegistry().user_message_counter.add(1, get_ctx_attributes())
     actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
 
     # Create a new job
@@ -907,6 +919,7 @@ async def send_message_async(
         use_assistant_message=request.use_assistant_message,
         assistant_message_tool_name=request.assistant_message_tool_name,
         assistant_message_tool_kwarg=request.assistant_message_tool_kwarg,
+        include_return_message_types=request.include_return_message_types,
     )
 
     return run
@@ -953,17 +966,13 @@ async def summarize_agent_conversation(
     This endpoint summarizes the current message history for a given agent,
     truncating and compressing it down to the specified `max_message_length`.
     """
-    actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
 
-    # user_eligible = actor.organization_id not in ["org-4a3af5dd-4c6a-48cb-ac13-3f73ecaaa4bf", "org-4ab3f6e8-9a44-4bee-aeb6-c681cbbc7bf6"]
-    # TODO: This is redundant, remove soon
+    actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
     agent = await server.agent_manager.get_agent_by_id_async(agent_id, actor, include_relationships=["multi_agent_group"])
     agent_eligible = agent.enable_sleeptime or agent.agent_type == AgentType.sleeptime_agent or not agent.multi_agent_group
-    experimental_header = request_obj.headers.get("X-EXPERIMENTAL") or "false"
-    feature_enabled = settings.use_experimental or experimental_header.lower() == "true"
     model_compatible = agent.llm_config.model_endpoint_type in ["anthropic", "openai", "together", "google_ai", "google_vertex"]
 
-    if agent_eligible and feature_enabled and model_compatible:
+    if agent_eligible and model_compatible:
         agent = LettaAgent(
             agent_id=agent_id,
             message_manager=server.message_manager,

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

@@ -86,7 +86,7 @@ def create_group(
 
 
 @router.patch("/{group_id}", response_model=Group, operation_id="modify_group")
-def modify_group(
+async def modify_group(
     group_id: str,
     group: GroupUpdate = Body(...),
     server: "SyncServer" = Depends(get_letta_server),
@@ -97,8 +97,8 @@ def modify_group(
     Create a new multi-agent group with the specified configuration.
     """
     try:
-        actor = server.user_manager.get_user_or_default(user_id=actor_id)
-        return server.group_manager.modify_group(group_id=group_id, group_update=group, actor=actor)
+        actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
+        return await server.group_manager.modify_group_async(group_id=group_id, group_update=group, actor=actor)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 

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

@@ -27,6 +27,11 @@ from letta.utils import safe_create_task, sanitize_filename
 
 logger = get_logger(__name__)
 
+mimetypes.add_type("text/markdown", ".md")
+mimetypes.add_type("text/markdown", ".markdown")
+mimetypes.add_type("application/jsonl", ".jsonl")
+mimetypes.add_type("application/x-jsonlines", ".jsonl")
+
 
 router = APIRouter(prefix="/sources", tags=["sources"])
 
@@ -174,7 +179,15 @@ async def upload_file_to_source(
     """
     Upload a file to a data source.
     """
-    allowed_media_types = {"application/pdf", "text/plain", "application/json"}
+    allowed_media_types = {
+        "application/pdf",
+        "text/plain",
+        "text/markdown",
+        "text/x-markdown",
+        "application/json",
+        "application/jsonl",
+        "application/x-jsonlines",
+    }
 
     # Normalize incoming Content-Type header (strip charset or any parameters).
     raw_ct = file.content_type or ""
@@ -192,6 +205,9 @@ async def upload_file_to_source(
                 ".pdf": "application/pdf",
                 ".txt": "text/plain",
                 ".json": "application/json",
+                ".md": "text/markdown",
+                ".markdown": "text/markdown",
+                ".jsonl": "application/jsonl",
             }
             media_type = ext_map.get(ext, media_type)
 
@@ -270,14 +286,21 @@ async def list_source_files(
     source_id: str,
     limit: int = Query(1000, description="Number of files to return"),
     after: Optional[str] = Query(None, description="Pagination cursor to fetch the next set of results"),
+    include_content: bool = Query(False, description="Whether to include full file content"),
     server: "SyncServer" = Depends(get_letta_server),
-    actor_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
+    actor_id: Optional[str] = Header(None, alias="user_id"),
 ):
     """
     List paginated files associated with a data source.
     """
     actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
-    return await server.source_manager.list_files(source_id=source_id, limit=limit, after=after, actor=actor)
+    return await server.source_manager.list_files(
+        source_id=source_id,
+        limit=limit,
+        after=after,
+        actor=actor,
+        include_content=include_content,
+    )
 
 
 # it's redundant to include /delete in the URL path. The HTTP verb DELETE already implies that action.

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

@@ -501,7 +501,8 @@ async def add_mcp_server_to_config(
         if isinstance(request, StdioServerConfig):
             mapped_request = MCPServer(server_name=request.server_name, server_type=request.type, stdio_config=request)
             # don't allow stdio servers
-            raise HTTPException(status_code=400, detail="StdioServerConfig is not supported")
+            if tool_settings.mcp_disable_stdio:  # protected server
+                raise HTTPException(status_code=400, detail="StdioServerConfig is not supported")
         elif isinstance(request, SSEServerConfig):
             mapped_request = MCPServer(server_name=request.server_name, server_type=request.type, server_url=request.server_url)
         # TODO: add HTTP streaming
@@ -530,4 +531,8 @@ async def delete_mcp_server_from_config(
         # log to DB
         actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
         mcp_server_id = await server.mcp_manager.get_mcp_server_id_by_name(mcp_server_name, actor)
-        return server.mcp_manager.delete_mcp_server_by_id(mcp_server_id, actor=actor)
+        await server.mcp_manager.delete_mcp_server_by_id(mcp_server_id, actor=actor)
+
+        # TODO: don't do this in the future (just return MCPServer)
+        all_servers = await server.mcp_manager.list_mcp_servers(actor=actor)
+        return [server.to_config() for server in all_servers]

letta/server/rest_api/utils.py

@@ -15,9 +15,12 @@ from pydantic import BaseModel
 
 from letta.constants import DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG, FUNC_FAILED_HEARTBEAT_MESSAGE, REQ_HEARTBEAT_MESSAGE
 from letta.errors import ContextWindowExceededError, RateLimitExceededError
-from letta.helpers.datetime_helpers import get_utc_time, get_utc_timestamp_ns
+from letta.helpers.datetime_helpers import get_utc_time, get_utc_timestamp_ns, ns_to_ms
 from letta.helpers.message_helper import convert_message_creates_to_messages
 from letta.log import get_logger
+from letta.otel.context import get_ctx_attributes
+from letta.otel.metric_registry import MetricRegistry
+from letta.otel.tracing import tracer
 from letta.schemas.enums import MessageRole
 from letta.schemas.letta_message_content import OmittedReasoningContent, ReasoningContent, RedactedReasoningContent, TextContent
 from letta.schemas.llm_config import LLMConfig
@@ -27,7 +30,6 @@ from letta.schemas.usage import LettaUsageStatistics
 from letta.schemas.user import User
 from letta.server.rest_api.interface import StreamingServerInterface
 from letta.system import get_heartbeat, package_function_response
-from letta.tracing import tracer
 
 if TYPE_CHECKING:
     from letta.server.server import SyncServer
@@ -81,8 +83,12 @@ async def sse_async_generator(
             if first_chunk and ttft_span is not None:
                 now = get_utc_timestamp_ns()
                 ttft_ns = now - request_start_timestamp_ns
-                ttft_span.add_event(name="time_to_first_token_ms", attributes={"ttft_ms": ttft_ns // 1_000_000})
+                ttft_span.add_event(name="time_to_first_token_ms", attributes={"ttft_ms": ns_to_ms(ttft_ns)})
                 ttft_span.end()
+                metric_attributes = get_ctx_attributes()
+                if llm_config:
+                    metric_attributes["model.name"] = llm_config.model
+                    MetricRegistry().ttft_ms_histogram.record(ns_to_ms(ttft_ns), metric_attributes)
                 first_chunk = False
 
             # yield f"data: {json.dumps(chunk)}\n\n"
@@ -190,7 +196,6 @@ def create_letta_messages_from_llm_response(
     add_heartbeat_request_system_message: bool = False,
     reasoning_content: Optional[List[Union[TextContent, ReasoningContent, RedactedReasoningContent, OmittedReasoningContent]]] = None,
     pre_computed_assistant_message_id: Optional[str] = None,
-    pre_computed_tool_message_id: Optional[str] = None,
     llm_batch_item_id: Optional[str] = None,
     step_id: str | None = None,
 ) -> List[Message]:
@@ -245,8 +250,6 @@ def create_letta_messages_from_llm_response(
             )
         ],
     )
-    if pre_computed_tool_message_id:
-        tool_message.id = pre_computed_tool_message_id
     messages.append(tool_message)
 
     if add_heartbeat_request_system_message:

letta/server/server.py

@@ -21,7 +21,7 @@ import letta.system as system
 from letta.agent import Agent, save_agent
 from letta.agents.letta_agent import LettaAgent
 from letta.config import LettaConfig
-from letta.constants import CORE_MEMORY_SOURCE_CHAR_LIMIT, LETTA_TOOL_EXECUTION_DIR
+from letta.constants import LETTA_TOOL_EXECUTION_DIR
 from letta.data_sources.connectors import DataConnector, load_data
 from letta.errors import HandleNotFoundError
 from letta.functions.mcp_client.types import MCPServerType, MCPTool, SSEServerConfig, StdioServerConfig
@@ -34,6 +34,7 @@ from letta.interface import AgentInterface  # abstract
 from letta.interface import CLIInterface  # for printing to terminal
 from letta.log import get_logger
 from letta.orm.errors import NoResultFound
+from letta.otel.tracing import log_event, trace_method
 from letta.prompts.gpt_system import get_system_text
 from letta.schemas.agent import AgentState, AgentType, CreateAgent, UpdateAgent
 from letta.schemas.block import Block, BlockUpdate, CreateBlock
@@ -44,7 +45,7 @@ from letta.schemas.enums import JobStatus, MessageStreamStatus, ProviderCategory
 from letta.schemas.environment_variables import SandboxEnvironmentVariableCreate
 from letta.schemas.group import GroupCreate, ManagerType, SleeptimeManager, VoiceSleeptimeManager
 from letta.schemas.job import Job, JobUpdate
-from letta.schemas.letta_message import LegacyLettaMessage, LettaMessage, ToolReturnMessage
+from letta.schemas.letta_message import LegacyLettaMessage, LettaMessage, MessageType, ToolReturnMessage
 from letta.schemas.letta_message_content import TextContent
 from letta.schemas.letta_response import LettaResponse
 from letta.schemas.llm_config import LLMConfig
@@ -101,7 +102,6 @@ from letta.services.tool_executor.tool_execution_manager import ToolExecutionMan
 from letta.services.tool_manager import ToolManager
 from letta.services.user_manager import UserManager
 from letta.settings import model_settings, settings, tool_settings
-from letta.tracing import log_event, trace_method
 from letta.utils import get_friendly_error_msg, get_persona_text, make_key
 
 config = LettaConfig.load()
@@ -1108,13 +1108,11 @@ class SyncServer(Server):
         after: Optional[str] = None,
         before: Optional[str] = None,
         limit: Optional[int] = 100,
-        order_by: Optional[str] = "created_at",
-        reverse: Optional[bool] = False,
         query_text: Optional[str] = None,
         ascending: Optional[bool] = True,
     ) -> List[Passage]:
         # iterate over records
-        records = await self.agent_manager.list_passages_async(
+        records = await self.agent_manager.list_agent_passages_async(
             actor=actor,
             agent_id=agent_id,
             after=after,
@@ -1368,12 +1366,13 @@ class SyncServer(Server):
             )
         await self.agent_manager.delete_agent_async(agent_id=sleeptime_agent_state.id, actor=actor)
 
-    async def _upsert_file_to_agent(self, agent_id: str, text: str, file_id: str, actor: User) -> None:
+    async def _upsert_file_to_agent(self, agent_id: str, text: str, file_id: str, file_name: str, actor: User) -> None:
         """
         Internal method to create or update a file <-> agent association
         """
-        truncated_text = text[:CORE_MEMORY_SOURCE_CHAR_LIMIT]
-        await self.file_agent_manager.attach_file(agent_id=agent_id, file_id=file_id, actor=actor, visible_content=truncated_text)
+        await self.file_agent_manager.attach_file(
+            agent_id=agent_id, file_id=file_id, file_name=file_name, actor=actor, visible_content=text
+        )
 
     async def _remove_file_from_agent(self, agent_id: str, file_id: str, actor: User) -> None:
         """
@@ -1389,7 +1388,7 @@ class SyncServer(Server):
             logger.info(f"File {file_id} already removed from agent {agent_id}, skipping...")
 
     async def insert_file_into_context_windows(
-        self, source_id: str, text: str, file_id: str, actor: User, agent_states: Optional[List[AgentState]] = None
+        self, source_id: str, text: str, file_id: str, file_name: str, actor: User, agent_states: Optional[List[AgentState]] = None
     ) -> List[AgentState]:
         """
         Insert the uploaded document into the context window of all agents
@@ -1404,11 +1403,13 @@ class SyncServer(Server):
         logger.info(f"Inserting document into context window for source: {source_id}")
         logger.info(f"Attached agents: {[a.id for a in agent_states]}")
 
-        await asyncio.gather(*(self._upsert_file_to_agent(agent_state.id, text, file_id, actor) for agent_state in agent_states))
+        await asyncio.gather(*(self._upsert_file_to_agent(agent_state.id, text, file_id, file_name, actor) for agent_state in agent_states))
 
         return agent_states
 
-    async def insert_files_into_context_window(self, agent_state: AgentState, texts: List[str], file_ids: List[str], actor: User) -> None:
+    async def insert_files_into_context_window(
+        self, agent_state: AgentState, texts: List[str], file_ids: List[str], file_names: List[str], actor: User
+    ) -> None:
         """
         Insert the uploaded documents into the context window of an agent
         attached to the given source.
@@ -1418,7 +1419,12 @@ class SyncServer(Server):
         if len(texts) != len(file_ids):
             raise ValueError(f"Mismatch between number of texts ({len(texts)}) and file ids ({len(file_ids)})")
 
-        await asyncio.gather(*(self._upsert_file_to_agent(agent_state.id, text, file_id, actor) for text, file_id in zip(texts, file_ids)))
+        await asyncio.gather(
+            *(
+                self._upsert_file_to_agent(agent_state.id, text, file_id, file_name, actor)
+                for text, file_id, file_name in zip(texts, file_ids, file_names)
+            )
+        )
 
     async def remove_file_from_context_windows(self, source_id: str, file_id: str, actor: User) -> None:
         """
@@ -2231,6 +2237,7 @@ class SyncServer(Server):
         assistant_message_tool_kwarg: str = constants.DEFAULT_MESSAGE_TOOL_KWARG,
         metadata: Optional[dict] = None,
         request_start_timestamp_ns: Optional[int] = None,
+        include_return_message_types: Optional[List[MessageType]] = None,
     ) -> Union[StreamingResponse, LettaResponse]:
         """Split off into a separate function so that it can be imported in the /chat/completion proxy."""
         # TODO: @charles is this the correct way to handle?
@@ -2336,6 +2343,11 @@ class SyncServer(Server):
 
                 # Get rid of the stream status messages
                 filtered_stream = [d for d in generated_stream if not isinstance(d, MessageStreamStatus)]
+
+                # Apply message type filtering if specified
+                if include_return_message_types is not None:
+                    filtered_stream = [msg for msg in filtered_stream if msg.message_type in include_return_message_types]
+
                 usage = await task
 
                 # By default the stream will be messages of type LettaMessage or LettaLegacyMessage