letta-nightly 0.7.20.dev20250521104258__py3-none-any.whl → 0.7.21.dev20250521233415__py3-none-any.whl

letta/services/step_manager.py

@@ -2,6 +2,7 @@ from datetime import datetime
 from typing import List, Literal, Optional
 
 from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import Session
 
 from letta.orm.errors import NoResultFound
@@ -12,6 +13,7 @@ from letta.schemas.openai.chat_completion_response import UsageStatistics
 from letta.schemas.step import Step as PydanticStep
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
+from letta.services.helpers.noop_helper import singleton
 from letta.tracing import get_trace_id
 from letta.utils import enforce_types
 
@@ -57,12 +59,14 @@ class StepManager:
         actor: PydanticUser,
         agent_id: str,
         provider_name: str,
+        provider_category: str,
         model: str,
         model_endpoint: Optional[str],
         context_window_limit: int,
         usage: UsageStatistics,
         provider_id: Optional[str] = None,
         job_id: Optional[str] = None,
+        step_id: Optional[str] = None,
     ) -> PydanticStep:
         step_data = {
             "origin": None,
@@ -70,6 +74,7 @@ class StepManager:
             "agent_id": agent_id,
             "provider_id": provider_id,
             "provider_name": provider_name,
+            "provider_category": provider_category,
             "model": model,
             "model_endpoint": model_endpoint,
             "context_window_limit": context_window_limit,
@@ -81,6 +86,8 @@ class StepManager:
             "tid": None,
             "trace_id": get_trace_id(),  # Get the current trace ID
         }
+        if step_id:
+            step_data["id"] = step_id
         with db_registry.session() as session:
             if job_id:
                 self._verify_job_access(session, job_id, actor, access=["write"])
@@ -88,6 +95,48 @@ class StepManager:
             new_step.create(session)
             return new_step.to_pydantic()
 
+    @enforce_types
+    async def log_step_async(
+        self,
+        actor: PydanticUser,
+        agent_id: str,
+        provider_name: str,
+        provider_category: str,
+        model: str,
+        model_endpoint: Optional[str],
+        context_window_limit: int,
+        usage: UsageStatistics,
+        provider_id: Optional[str] = None,
+        job_id: Optional[str] = None,
+        step_id: Optional[str] = None,
+    ) -> PydanticStep:
+        step_data = {
+            "origin": None,
+            "organization_id": actor.organization_id,
+            "agent_id": agent_id,
+            "provider_id": provider_id,
+            "provider_name": provider_name,
+            "provider_category": provider_category,
+            "model": model,
+            "model_endpoint": model_endpoint,
+            "context_window_limit": context_window_limit,
+            "completion_tokens": usage.completion_tokens,
+            "prompt_tokens": usage.prompt_tokens,
+            "total_tokens": usage.total_tokens,
+            "job_id": job_id,
+            "tags": [],
+            "tid": None,
+            "trace_id": get_trace_id(),  # Get the current trace ID
+        }
+        if step_id:
+            step_data["id"] = step_id
+        async with db_registry.async_session() as session:
+            if job_id:
+                await self._verify_job_access_async(session, job_id, actor, access=["write"])
+            new_step = StepModel(**step_data)
+            await new_step.create_async(session)
+            return new_step.to_pydantic()
+
     @enforce_types
     def get_step(self, step_id: str, actor: PydanticUser) -> PydanticStep:
         with db_registry.session() as session:
@@ -147,3 +196,100 @@ class StepManager:
         if not job:
             raise NoResultFound(f"Job with id {job_id} does not exist or user does not have access")
         return job
+
+    async def _verify_job_access_async(
+        self,
+        session: AsyncSession,
+        job_id: str,
+        actor: PydanticUser,
+        access: List[Literal["read", "write", "delete"]] = ["read"],
+    ) -> JobModel:
+        """
+        Verify that a job exists and the user has the required access asynchronously.
+
+        Args:
+            session: The async database session
+            job_id: The ID of the job to verify
+            actor: The user making the request
+
+        Returns:
+            The job if it exists and the user has access
+
+        Raises:
+            NoResultFound: If the job does not exist or user does not have access
+        """
+        job_query = select(JobModel).where(JobModel.id == job_id)
+        job_query = JobModel.apply_access_predicate(job_query, actor, access, AccessType.USER)
+        result = await session.execute(job_query)
+        job = result.scalar_one_or_none()
+        if not job:
+            raise NoResultFound(f"Job with id {job_id} does not exist or user does not have access")
+        return job
+
+
+@singleton
+class NoopStepManager(StepManager):
+    """
+    Noop implementation of StepManager.
+    Temporarily used for migrations, but allows for different implementations in the future.
+    Will not allow for writes, but will still allow for reads.
+    """
+
+    @enforce_types
+    def log_step(
+        self,
+        actor: PydanticUser,
+        agent_id: str,
+        provider_name: str,
+        provider_category: str,
+        model: str,
+        model_endpoint: Optional[str],
+        context_window_limit: int,
+        usage: UsageStatistics,
+        provider_id: Optional[str] = None,
+        job_id: Optional[str] = None,
+        step_id: Optional[str] = None,
+    ) -> PydanticStep:
+        return
+
+    @enforce_types
+    async def log_step_async(
+        self,
+        actor: PydanticUser,
+        agent_id: str,
+        provider_name: str,
+        provider_category: str,
+        model: str,
+        model_endpoint: Optional[str],
+        context_window_limit: int,
+        usage: UsageStatistics,
+        provider_id: Optional[str] = None,
+        job_id: Optional[str] = None,
+        step_id: Optional[str] = None,
+    ) -> PydanticStep:
+        step_data = {
+            "origin": None,
+            "organization_id": actor.organization_id,
+            "agent_id": agent_id,
+            "provider_id": provider_id,
+            "provider_name": provider_name,
+            "provider_category": provider_category,
+            "model": model,
+            "model_endpoint": model_endpoint,
+            "context_window_limit": context_window_limit,
+            "completion_tokens": usage.completion_tokens,
+            "prompt_tokens": usage.prompt_tokens,
+            "total_tokens": usage.total_tokens,
+            "job_id": job_id,
+            "tags": [],
+            "tid": None,
+            "trace_id": get_trace_id(),  # Get the current trace ID
+        }
+        if step_id:
+            step_data["id"] = step_id
+        async with db_registry.async_session() as session:
+            if job_id:
+                await self._verify_job_access_async(session, job_id, actor, access=["write"])
+            new_step = StepModel(**step_data)
+            await new_step.create_async(session)
+            return new_step.to_pydantic()

letta/services/telemetry_manager.py

@@ -0,0 +1,58 @@
+from letta.helpers.json_helpers import json_dumps, json_loads
+from letta.orm.provider_trace import ProviderTrace as ProviderTraceModel
+from letta.schemas.provider_trace import ProviderTrace as PydanticProviderTrace
+from letta.schemas.provider_trace import ProviderTraceCreate
+from letta.schemas.step import Step as PydanticStep
+from letta.schemas.user import User as PydanticUser
+from letta.server.db import db_registry
+from letta.services.helpers.noop_helper import singleton
+from letta.utils import enforce_types
+
+
+class TelemetryManager:
+    @enforce_types
+    async def get_provider_trace_by_step_id_async(
+        self,
+        step_id: str,
+        actor: PydanticUser,
+    ) -> PydanticProviderTrace:
+        async with db_registry.async_session() as session:
+            provider_trace = await ProviderTraceModel.read_async(db_session=session, step_id=step_id, actor=actor)
+            return provider_trace.to_pydantic()
+
+    @enforce_types
+    async def create_provider_trace_async(self, actor: PydanticUser, provider_trace_create: ProviderTraceCreate) -> PydanticProviderTrace:
+        async with db_registry.async_session() as session:
+            provider_trace = ProviderTraceModel(**provider_trace_create.model_dump())
+            if provider_trace_create.request_json:
+                request_json_str = json_dumps(provider_trace_create.request_json)
+                provider_trace.request_json = json_loads(request_json_str)
+
+            if provider_trace_create.response_json:
+                response_json_str = json_dumps(provider_trace_create.response_json)
+                provider_trace.response_json = json_loads(response_json_str)
+            await provider_trace.create_async(session, actor=actor)
+            return provider_trace.to_pydantic()
+
+    @enforce_types
+    def create_provider_trace(self, actor: PydanticUser, provider_trace_create: ProviderTraceCreate) -> PydanticProviderTrace:
+        with db_registry.session() as session:
+            provider_trace = ProviderTraceModel(**provider_trace_create.model_dump())
+            provider_trace.create(session, actor=actor)
+            return provider_trace.to_pydantic()
+
+
+@singleton
+class NoopTelemetryManager(TelemetryManager):
+    """
+    Noop implementation of TelemetryManager.
+    """
+
+    async def create_provider_trace_async(self, actor: PydanticUser, provider_trace_create: ProviderTraceCreate) -> PydanticProviderTrace:
+        return
+
+    async def get_provider_trace_by_step_id_async(self, step_id: str, actor: PydanticUser) -> PydanticStep:
+        return
+
+    def create_provider_trace(self, actor: PydanticUser, provider_trace_create: ProviderTraceCreate) -> PydanticProviderTrace:
+        return

letta/services/tool_executor/tool_execution_manager.py

@@ -8,9 +8,14 @@ from letta.schemas.sandbox_config import SandboxConfig
 from letta.schemas.tool import Tool
 from letta.schemas.tool_execution_result import ToolExecutionResult
 from letta.schemas.user import User
+from letta.services.agent_manager import AgentManager
+from letta.services.block_manager import BlockManager
+from letta.services.message_manager import MessageManager
+from letta.services.passage_manager import PassageManager
 from letta.services.tool_executor.tool_executor import (
     ExternalComposioToolExecutor,
     ExternalMCPToolExecutor,
+    LettaBuiltinToolExecutor,
     LettaCoreToolExecutor,
     LettaMultiAgentToolExecutor,
     SandboxToolExecutor,
@@ -28,15 +33,30 @@ class ToolExecutorFactory:
         ToolType.LETTA_MEMORY_CORE: LettaCoreToolExecutor,
         ToolType.LETTA_SLEEPTIME_CORE: LettaCoreToolExecutor,
         ToolType.LETTA_MULTI_AGENT_CORE: LettaMultiAgentToolExecutor,
+        ToolType.LETTA_BUILTIN: LettaBuiltinToolExecutor,
         ToolType.EXTERNAL_COMPOSIO: ExternalComposioToolExecutor,
         ToolType.EXTERNAL_MCP: ExternalMCPToolExecutor,
     }
 
     @classmethod
-    def get_executor(cls, tool_type: ToolType) -> ToolExecutor:
+    def get_executor(
+        cls,
+        tool_type: ToolType,
+        message_manager: MessageManager,
+        agent_manager: AgentManager,
+        block_manager: BlockManager,
+        passage_manager: PassageManager,
+        actor: User,
+    ) -> ToolExecutor:
         """Get the appropriate executor for the given tool type."""
         executor_class = cls._executor_map.get(tool_type, SandboxToolExecutor)
-        return executor_class()
+        return executor_class(
+            message_manager=message_manager,
+            agent_manager=agent_manager,
+            block_manager=block_manager,
+            passage_manager=passage_manager,
+            actor=actor,
+        )
 
 
 class ToolExecutionManager:
@@ -44,11 +64,19 @@ class ToolExecutionManager:
 
     def __init__(
         self,
+        message_manager: MessageManager,
+        agent_manager: AgentManager,
+        block_manager: BlockManager,
+        passage_manager: PassageManager,
         agent_state: AgentState,
         actor: User,
         sandbox_config: Optional[SandboxConfig] = None,
         sandbox_env_vars: Optional[Dict[str, Any]] = None,
     ):
+        self.message_manager = message_manager
+        self.agent_manager = agent_manager
+        self.block_manager = block_manager
+        self.passage_manager = passage_manager
         self.agent_state = agent_state
         self.logger = get_logger(__name__)
         self.actor = actor
@@ -68,7 +96,14 @@ class ToolExecutionManager:
             Tuple containing the function response and sandbox run result (if applicable)
         """
         try:
-            executor = ToolExecutorFactory.get_executor(tool.tool_type)
+            executor = ToolExecutorFactory.get_executor(
+                tool.tool_type,
+                message_manager=self.message_manager,
+                agent_manager=self.agent_manager,
+                block_manager=self.block_manager,
+                passage_manager=self.passage_manager,
+                actor=self.actor,
+            )
             return executor.execute(
                 function_name,
                 function_args,
@@ -98,9 +133,18 @@ class ToolExecutionManager:
         Execute a tool asynchronously and persist any state changes.
         """
         try:
-            executor = ToolExecutorFactory.get_executor(tool.tool_type)
+            executor = ToolExecutorFactory.get_executor(
+                tool.tool_type,
+                message_manager=self.message_manager,
+                agent_manager=self.agent_manager,
+                block_manager=self.block_manager,
+                passage_manager=self.passage_manager,
+                actor=self.actor,
+            )
             # TODO: Extend this async model to composio
-            if isinstance(executor, (SandboxToolExecutor, ExternalComposioToolExecutor)):
+            if isinstance(
+                executor, (SandboxToolExecutor, ExternalComposioToolExecutor, LettaBuiltinToolExecutor, LettaMultiAgentToolExecutor)
+            ):
                 result = await executor.execute(function_name, function_args, self.agent_state, tool, self.actor)
             else:
                 result = executor.execute(function_name, function_args, self.agent_state, tool, self.actor)

letta/services/tool_executor/tool_execution_sandbox.py

@@ -73,6 +73,7 @@ class ToolExecutionSandbox:
         self.force_recreate = force_recreate
         self.force_recreate_venv = force_recreate_venv
 
+    @trace_method
     def run(
         self,
         agent_state: Optional[AgentState] = None,
@@ -321,6 +322,7 @@ class ToolExecutionSandbox:
 
     # e2b sandbox specific functions
 
+    @trace_method
     def run_e2b_sandbox(
         self,
         agent_state: Optional[AgentState] = None,
@@ -352,10 +354,22 @@ class ToolExecutionSandbox:
         if additional_env_vars:
             env_vars.update(additional_env_vars)
         code = self.generate_execution_script(agent_state=agent_state)
+        log_event(
+            "e2b_execution_started",
+            {"tool": self.tool_name, "sandbox_id": sbx.sandbox_id, "code": code, "env_vars": env_vars},
+        )
         execution = sbx.run_code(code, envs=env_vars)
 
         if execution.results:
             func_return, agent_state = self.parse_best_effort(execution.results[0].text)
+            log_event(
+                "e2b_execution_succeeded",
+                {
+                    "tool": self.tool_name,
+                    "sandbox_id": sbx.sandbox_id,
+                    "func_return": func_return,
+                },
+            )
         elif execution.error:
             logger.error(f"Executing tool {self.tool_name} raised a {execution.error.name} with message: \n{execution.error.value}")
             logger.error(f"Traceback from e2b sandbox: \n{execution.error.traceback}")
@@ -363,7 +377,25 @@ class ToolExecutionSandbox:
                 function_name=self.tool_name, exception_name=execution.error.name, exception_message=execution.error.value
             )
             execution.logs.stderr.append(execution.error.traceback)
+            log_event(
+                "e2b_execution_failed",
+                {
+                    "tool": self.tool_name,
+                    "sandbox_id": sbx.sandbox_id,
+                    "error_type": execution.error.name,
+                    "error_message": execution.error.value,
+                    "func_return": func_return,
+                },
+            )
         else:
+            log_event(
+                "e2b_execution_empty",
+                {
+                    "tool": self.tool_name,
+                    "sandbox_id": sbx.sandbox_id,
+                    "status": "no_results_no_error",
+                },
+            )
             raise ValueError(f"Tool {self.tool_name} returned execution with None")
 
         return ToolExecutionResult(
@@ -395,16 +427,31 @@ class ToolExecutionSandbox:
 
         return None
 
+    @trace_method
     def create_e2b_sandbox_with_metadata_hash(self, sandbox_config: SandboxConfig) -> "Sandbox":
         from e2b_code_interpreter import Sandbox
 
         state_hash = sandbox_config.fingerprint()
         e2b_config = sandbox_config.get_e2b_config()
+        log_event(
+            "e2b_sandbox_create_started",
+            {
+                "sandbox_fingerprint": state_hash,
+                "e2b_config": e2b_config.model_dump(),
+            },
+        )
         if e2b_config.template:
             sbx = Sandbox(sandbox_config.get_e2b_config().template, metadata={self.METADATA_CONFIG_STATE_KEY: state_hash})
         else:
             # no template
             sbx = Sandbox(metadata={self.METADATA_CONFIG_STATE_KEY: state_hash}, **e2b_config.model_dump(exclude={"pip_requirements"}))
+        log_event(
+            "e2b_sandbox_create_finished",
+            {
+                "sandbox_id": sbx.sandbox_id,
+                "sandbox_fingerprint": state_hash,
+            },
+        )
 
         # install pip requirements
         if e2b_config.pip_requirements: