From 2536942be2423aef355d1ec730f944e274da0b13 Mon Sep 17 00:00:00 2001 From: jnjpng Date: Mon, 8 Dec 2025 17:53:42 -0800 Subject: [PATCH] fix: combined tool manager improvements - tracing and redundant fetches (#6570) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: combined tool manager improvements - tracing and redundant fetches This PR combines improvements from #6530 and #6535: - Add tracer import to enable proper tracing spans - Improve update check logic to verify actual field changes before updating - Return current_tool directly when no update is needed (avoids redundant fetch) - Add structured tracing spans to update_tool_by_id_async for better observability - Fix decorator order for better error handling (raise_on_invalid_id before trace_method) - Remove unnecessary tracing spans in create_or_update_tool_async 🐾 Generated with [Letta Code](https://letta.com) Co-Authored-By: Letta * revert: remove tracing spans from update_tool_by_id_async Remove the tracer span additions from update_tool_by_id_async while keeping all other improvements (decorator order fix, redundant fetch removal, and improved update check logic). 🐾 Generated with [Letta Code](https://letta.com) Co-Authored-By: Letta --------- Co-authored-by: Letta Bot --- letta/services/tool_manager.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/letta/services/tool_manager.py b/letta/services/tool_manager.py index 6b52cd70..246f493b 100644 --- a/letta/services/tool_manager.py +++ b/letta/services/tool_manager.py @@ -31,7 +31,7 @@ from letta.log import get_logger # TODO: Remove this once we translate all of these to the ORM from letta.orm.errors import NoResultFound from letta.orm.tool import Tool as ToolModel -from letta.otel.tracing import trace_method +from letta.otel.tracing import trace_method, tracer from letta.schemas.agent import AgentState from letta.schemas.enums import PrimitiveType, SandboxType, ToolType from letta.schemas.tool import Tool as PydanticTool, ToolCreate, ToolUpdate @@ -205,8 +205,6 @@ class ToolManager: Uses atomic PostgreSQL ON CONFLICT DO UPDATE to prevent race conditions during concurrent upserts. """ - from letta.otel.tracing import tracer - if pydantic_tool.tool_type == ToolType.CUSTOM and not pydantic_tool.json_schema: with tracer.start_as_current_span("generate_schema_for_tool_creation"): generated_schema = generate_schema_for_tool_creation(pydantic_tool) @@ -238,21 +236,23 @@ class ToolManager: # Fallback for SQLite: use non-atomic check-then-act pattern current_tool = await self.get_tool_by_name_async(tool_name=pydantic_tool.name, actor=actor) + if current_tool: # Put to dict and remove fields that should not be reset - with tracer.start_as_current_span("pydantic_tool.model_dump"): - update_data = pydantic_tool.model_dump(exclude_unset=True, exclude_none=True) - update_data["organization_id"] = actor.organization_id + update_data = pydantic_tool.model_dump(exclude_unset=True, exclude_none=True) - # If there's anything to update - if update_data: + # Check if any field in update_data actually differs from the current tool + current_tool_data = current_tool.model_dump() + needs_update = any(current_tool_data.get(key) != value for key, value in update_data.items()) + + if needs_update: # In case we want to update the tool type # Useful if we are shuffling around base tools updated_tool_type = None if "tool_type" in update_data: updated_tool_type = update_data.get("tool_type") - with tracer.start_as_current_span("ToolUpdate_initialization"): - tool_update = ToolUpdate(**update_data) + + tool_update = ToolUpdate(**update_data) tool = await self.update_tool_by_id_async( current_tool.id, tool_update, @@ -264,7 +264,7 @@ class ToolManager: printd( f"`create_or_update_tool` was called with user_id={actor.id}, organization_id={actor.organization_id}, name={pydantic_tool.name}, but found existing tool with nothing to update." ) - tool = await self.get_tool_by_id_async(current_tool.id, actor=actor) + return current_tool return tool return await self.create_tool_async(pydantic_tool, actor=actor, modal_sandbox_enabled=modal_sandbox_enabled) @@ -515,8 +515,8 @@ class ToolManager: return await self._upsert_tools_individually(pydantic_tools, actor, override_existing_tools) @enforce_types - @trace_method @raise_on_invalid_id(param_name="tool_id", expected_prefix=PrimitiveType.TOOL) + @trace_method async def get_tool_by_id_async(self, tool_id: str, actor: PydanticUser) -> PydanticTool: """Fetch a tool by its ID.""" async with db_registry.async_session() as session: @@ -548,8 +548,8 @@ class ToolManager: return None @enforce_types - @trace_method @raise_on_invalid_id(param_name="tool_id", expected_prefix=PrimitiveType.TOOL) + @trace_method async def tool_exists_async(self, tool_id: str, actor: PydanticUser) -> bool: """Check if a tool exists and belongs to the user's organization (lightweight check).""" async with db_registry.async_session() as session: @@ -865,8 +865,8 @@ class ToolManager: return await ToolModel.size_async(db_session=session, actor=actor, name=LETTA_TOOL_SET) @enforce_types - @trace_method @raise_on_invalid_id(param_name="tool_id", expected_prefix=PrimitiveType.TOOL) + @trace_method async def update_tool_by_id_async( self, tool_id: str, @@ -1060,8 +1060,8 @@ class ToolManager: return updated_tool @enforce_types - @trace_method # @raise_on_invalid_id This is commented out bc it's called by _list_tools_async, when it encounters malformed tools (i.e. if id is invalid will fail validation on deletion) + @trace_method async def delete_tool_by_id_async(self, tool_id: str, actor: PydanticUser) -> None: """Delete a tool by its ID.""" async with db_registry.async_session() as session: