diff --git a/letta/schemas/tool.py b/letta/schemas/tool.py index e703291c..2d44f41e 100644 --- a/letta/schemas/tool.py +++ b/letta/schemas/tool.py @@ -16,12 +16,10 @@ from letta.constants import ( # MCP Tool metadata constants for schema health status MCP_TOOL_METADATA_SCHEMA_STATUS = f"{MCP_TOOL_TAG_NAME_PREFIX}:SCHEMA_STATUS" MCP_TOOL_METADATA_SCHEMA_WARNINGS = f"{MCP_TOOL_TAG_NAME_PREFIX}:SCHEMA_WARNINGS" -from letta.functions.ast_parsers import get_function_name_and_docstring from letta.functions.composio_helpers import generate_composio_tool_wrapper -from letta.functions.functions import derive_openai_json_schema, get_json_schema_from_module +from letta.functions.functions import get_json_schema_from_module from letta.functions.mcp_client.types import MCPTool from letta.functions.schema_generator import ( - generate_schema_from_args_schema_v2, generate_tool_schema_for_composio, generate_tool_schema_for_mcp, ) @@ -80,46 +78,19 @@ class Tool(BaseTool): def refresh_source_code_and_json_schema(self): """ Refresh name, description, source_code, and json_schema. + + Note: Schema generation for custom tools is now handled at creation/update time in ToolManager. + This method only handles built-in Letta tools. """ - from letta.functions.helpers import generate_model_from_args_json_schema - - if self.tool_type == ToolType.CUSTOM and not self.json_schema: - # attempt various fallbacks to get the JSON schema - if not self.source_code: - logger.error("Custom tool with id=%s is missing source_code field", self.id) - raise ValueError(f"Custom tool with id={self.id} is missing source_code field.") - - if self.source_type == ToolSourceType.typescript: - # TypeScript tools don't support args_json_schema, only direct schema generation - if not self.json_schema: - try: - from letta.functions.typescript_parser import derive_typescript_json_schema - - self.json_schema = derive_typescript_json_schema(source_code=self.source_code) - except Exception as e: - logger.error("Failed to derive TypeScript json schema for tool with id=%s name=%s: %s", self.id, self.name, e) - elif ( - self.source_type == ToolSourceType.python or self.source_type is None - ): # default to python if not provided for backwards compatability - # Python tool handling - # Always derive json_schema for freshest possible json_schema - if self.args_json_schema is not None: - name, description = get_function_name_and_docstring(self.source_code, self.name) - args_schema = generate_model_from_args_json_schema(self.args_json_schema) - self.json_schema = generate_schema_from_args_schema_v2( - args_schema=args_schema, - name=name, - description=description, - append_heartbeat=False, - ) - else: # elif not self.json_schema: # TODO: JSON schema is not being derived correctly the first time? - # If there's not a json_schema provided, then we need to re-derive - try: - self.json_schema = derive_openai_json_schema(source_code=self.source_code) - except Exception as e: - logger.error("Failed to derive json schema for tool with id=%s name=%s: %s", self.id, self.name, e) - else: - raise ValueError(f"Unknown tool source type: {self.source_type}") + if self.tool_type == ToolType.CUSTOM: + # Custom tools should already have their schema set during creation/update + # No schema generation happens here anymore + if not self.json_schema: + logger.warning( + "Custom tool with id=%s name=%s is missing json_schema. Schema should be set during creation/update.", + self.id, + self.name, + ) elif self.tool_type in {ToolType.LETTA_CORE, ToolType.LETTA_MEMORY_CORE, ToolType.LETTA_SLEEPTIME_CORE}: # If it's letta core tool, we generate the json_schema on the fly here self.json_schema = get_json_schema_from_module(module_name=LETTA_CORE_TOOL_MODULE_NAME, function_name=self.name) @@ -139,23 +110,6 @@ class Tool(BaseTool): # Composio schemas handled separately pass - # At this point, we need to validate that at least json_schema is populated - if not self.json_schema: - logger.error("Tool with id=%s name=%s tool_type=%s is missing a json_schema", self.id, self.name, self.tool_type) - raise ValueError(f"Tool with id={self.id} name={self.name} tool_type={self.tool_type} is missing a json_schema.") - - # Derive name from the JSON schema if not provided - if not self.name: - # TODO: This in theory could error, but name should always be on json_schema - # TODO: Make JSON schema a typed pydantic object - self.name = self.json_schema.get("name") - - # Derive description from the JSON schema if not provided - if not self.description: - # TODO: This in theory could error, but description should always be on json_schema - # TODO: Make JSON schema a typed pydantic object - self.description = self.json_schema.get("description") - return self @@ -253,6 +207,7 @@ class ToolUpdate(LettaBase): npm_requirements: list[NpmRequirement] | None = Field(None, description="Optional list of npm packages required by this tool.") metadata_: Optional[Dict[str, Any]] = Field(None, description="A dictionary of additional metadata for the tool.") default_requires_approval: Optional[bool] = Field(None, description="Whether or not to require approval before executing this tool.") + # name: Optional[str] = Field(None, description="The name of the tool (must match the JSON schema name and source code function name).") model_config = ConfigDict(extra="ignore") # Allows extra fields without validation errors # TODO: Remove this, and clean usage of ToolUpdate everywhere else diff --git a/letta/server/rest_api/app.py b/letta/server/rest_api/app.py index 147496f5..6123c193 100644 --- a/letta/server/rest_api/app.py +++ b/letta/server/rest_api/app.py @@ -20,6 +20,8 @@ from letta.constants import ADMIN_PREFIX, API_PREFIX, OPENAI_API_PREFIX from letta.errors import ( BedrockPermissionError, LettaAgentNotFoundError, + LettaToolCreateError, + LettaToolNameConflictError, LettaUserNotFoundError, LLMAuthenticationError, LLMError, @@ -240,6 +242,8 @@ def create_application() -> "FastAPI": app.add_exception_handler(LettaUserNotFoundError, _error_handler_404_user) app.add_exception_handler(ForeignKeyConstraintViolationError, _error_handler_409) app.add_exception_handler(UniqueConstraintViolationError, _error_handler_409) + app.add_exception_handler(LettaToolCreateError, _error_handler_400) + app.add_exception_handler(LettaToolNameConflictError, _error_handler_400) @app.exception_handler(IncompatibleAgentType) async def handle_incompatible_agent_type(request: Request, exc: IncompatibleAgentType): diff --git a/letta/server/rest_api/routers/v1/tools.py b/letta/server/rest_api/routers/v1/tools.py index 4f51cf01..c030592c 100644 --- a/letta/server/rest_api/routers/v1/tools.py +++ b/letta/server/rest_api/routers/v1/tools.py @@ -287,19 +287,9 @@ async def create_tool( """ Create a new tool """ - try: - actor = await server.user_manager.get_actor_or_default_async(actor_id=headers.actor_id) - tool = Tool(**request.model_dump(exclude_unset=True)) - return await server.tool_manager.create_tool_async(pydantic_tool=tool, actor=actor) - except UniqueConstraintViolationError as e: - clean_error_message = "Tool with this name already exists." - raise HTTPException(status_code=409, detail=clean_error_message) - except LettaToolCreateError as e: - # HTTP 400 == Bad Request - raise HTTPException(status_code=400, detail=str(e)) - except Exception as e: - # Catch other unexpected errors and raise an internal server error - raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}") + actor = await server.user_manager.get_actor_or_default_async(actor_id=headers.actor_id) + tool = Tool(**request.model_dump(exclude_unset=True)) + return await server.tool_manager.create_or_update_tool_async(pydantic_tool=tool, actor=actor) @router.put("/", response_model=Tool, operation_id="upsert_tool") @@ -311,21 +301,9 @@ async def upsert_tool( """ Create or update a tool """ - try: - actor = await server.user_manager.get_actor_or_default_async(actor_id=headers.actor_id) - tool = await server.tool_manager.create_or_update_tool_async( - pydantic_tool=Tool(**request.model_dump(exclude_unset=True)), actor=actor - ) - return tool - except UniqueConstraintViolationError as e: - # Log the error and raise a conflict exception - raise HTTPException(status_code=409, detail=str(e)) - except LettaToolCreateError as e: - # HTTP 400 == Bad Request - raise HTTPException(status_code=400, detail=str(e)) - except Exception as e: - # Catch other unexpected errors and raise an internal server error - raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}") + actor = await server.user_manager.get_actor_or_default_async(actor_id=headers.actor_id) + tool = await server.tool_manager.create_or_update_tool_async(pydantic_tool=Tool(**request.model_dump(exclude_unset=True)), actor=actor) + return tool @router.patch("/{tool_id}", response_model=Tool, operation_id="modify_tool") @@ -338,19 +316,9 @@ async def modify_tool( """ Update an existing tool """ - try: - actor = await server.user_manager.get_actor_or_default_async(actor_id=headers.actor_id) - tool = await server.tool_manager.update_tool_by_id_async(tool_id=tool_id, tool_update=request, actor=actor) - return tool - except LettaToolNameConflictError as e: - # HTTP 409 == Conflict - raise HTTPException(status_code=409, detail=str(e)) - except LettaToolCreateError as e: - # HTTP 400 == Bad Request - raise HTTPException(status_code=400, detail=str(e)) - except Exception as e: - # Catch other unexpected errors and raise an internal server error - raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}") + actor = await server.user_manager.get_actor_or_default_async(actor_id=headers.actor_id) + tool = await server.tool_manager.update_tool_by_id_async(tool_id=tool_id, tool_update=request, actor=actor) + return tool @router.post("/add-base-tools", response_model=List[Tool], operation_id="add_base_tools") @@ -376,25 +344,17 @@ async def run_tool_from_source( """ actor = await server.user_manager.get_actor_or_default_async(actor_id=headers.actor_id) - try: - return await server.run_tool_from_source( - tool_source=request.source_code, - tool_source_type=request.source_type, - tool_args=request.args, - tool_env_vars=request.env_vars, - tool_name=request.name, - tool_args_json_schema=request.args_json_schema, - tool_json_schema=request.json_schema, - pip_requirements=request.pip_requirements, - actor=actor, - ) - except LettaToolCreateError as e: - # HTTP 400 == Bad Request - raise HTTPException(status_code=400, detail=str(e)) - - except Exception as e: - # Catch other unexpected errors and raise an internal server error - raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}") + return await server.run_tool_from_source( + tool_source=request.source_code, + tool_source_type=request.source_type, + tool_args=request.args, + tool_env_vars=request.env_vars, + tool_name=request.name, + tool_args_json_schema=request.args_json_schema, + tool_json_schema=request.json_schema, + pip_requirements=request.pip_requirements, + actor=actor, + ) # Specific routes for Composio diff --git a/letta/server/server.py b/letta/server/server.py index 77b85cef..5ff6eb59 100644 --- a/letta/server/server.py +++ b/letta/server/server.py @@ -1181,6 +1181,8 @@ class SyncServer(object): ) -> ToolReturnMessage: """Run a tool from source code""" + from letta.services.tool_schema_generator import generate_schema_for_tool_creation, generate_schema_for_tool_update + if tool_source_type not in (None, ToolSourceType.python, ToolSourceType.typescript): raise ValueError("Tool source type is not supported at this time. Found {tool_source_type}") @@ -1203,6 +1205,11 @@ class SyncServer(object): source_type=tool_source_type, ) + # try to get the schema + if not tool.name: + if not tool.json_schema: + tool.json_schema = generate_schema_for_tool_creation(tool) + tool.name = tool.json_schema.get("name") assert tool.name is not None, "Failed to create tool object" # TODO eventually allow using agent state in tools diff --git a/letta/services/tool_manager.py b/letta/services/tool_manager.py index 7e480332..29943595 100644 --- a/letta/services/tool_manager.py +++ b/letta/services/tool_manager.py @@ -32,6 +32,7 @@ from letta.schemas.user import User as PydanticUser from letta.server.db import db_registry from letta.services.helpers.agent_manager_helper import calculate_multi_agent_tools from letta.services.mcp.types import SSEServerConfig, StdioServerConfig +from letta.services.tool_schema_generator import generate_schema_for_tool_creation, generate_schema_for_tool_update from letta.settings import settings from letta.utils import enforce_types, printd @@ -47,8 +48,29 @@ class ToolManager: self, pydantic_tool: PydanticTool, actor: PydanticUser, bypass_name_check: bool = False ) -> PydanticTool: """Create a new tool based on the ToolCreate schema.""" - tool_id = await self.get_tool_id_by_name_async(tool_name=pydantic_tool.name, actor=actor) - if tool_id: + if pydantic_tool.tool_type == ToolType.CUSTOM and not pydantic_tool.json_schema: + generated_schema = generate_schema_for_tool_creation(pydantic_tool) + if generated_schema: + pydantic_tool.json_schema = generated_schema + else: + raise ValueError("Failed to generate schema for tool", pydantic_tool.source_code) + + print("SCHEMA", pydantic_tool.json_schema) + + # make sure the name matches the json_schema + if not pydantic_tool.name: + pydantic_tool.name = pydantic_tool.json_schema.get("name") + else: + if pydantic_tool.name != pydantic_tool.json_schema.get("name"): + raise LettaToolNameSchemaMismatchError( + tool_name=pydantic_tool.name, + json_schema_name=pydantic_tool.json_schema.get("name"), + source_code=pydantic_tool.source_code, + ) + + # check if the tool name already exists + 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 update_data = pydantic_tool.model_dump(exclude_unset=True, exclude_none=True) update_data["organization_id"] = actor.organization_id @@ -61,17 +83,17 @@ class ToolManager: if "tool_type" in update_data: updated_tool_type = update_data.get("tool_type") tool = await self.update_tool_by_id_async( - tool_id, ToolUpdate(**update_data), actor, updated_tool_type=updated_tool_type, bypass_name_check=bypass_name_check + current_tool.id, ToolUpdate(**update_data), actor, updated_tool_type=updated_tool_type ) else: 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(tool_id, actor=actor) - else: - tool = await self.create_tool_async(pydantic_tool, actor=actor) + tool = await self.get_tool_by_id_async(current_tool.id, actor=actor) + return tool + + return await self.create_tool_async(pydantic_tool, actor=actor) - return tool @enforce_types async def create_mcp_server( @@ -115,9 +137,11 @@ class ToolManager: @trace_method async def create_tool_async(self, pydantic_tool: PydanticTool, actor: PydanticUser) -> PydanticTool: """Create a new tool based on the ToolCreate schema.""" + # Generate schema only if not provided (only for custom tools) + async with db_registry.async_session() as session: # Auto-generate description if not provided - if pydantic_tool.description is None: + if pydantic_tool.description is None and pydantic_tool.json_schema: pydantic_tool.description = pydantic_tool.json_schema.get("description", None) tool_data = pydantic_tool.model_dump(to_orm=True) # Set the organization id at the ORM layer @@ -166,6 +190,11 @@ class ToolManager: if not pydantic_tools: return [] + # get schemas if not provided + for tool in pydantic_tools: + if tool.json_schema is None: + tool.json_schema = generate_schema_for_tool_creation(tool) + # auto-generate descriptions if not provided for tool in pydantic_tools: if tool.description is None: @@ -494,34 +523,56 @@ class ToolManager: bypass_name_check: bool = False, ) -> PydanticTool: """Update a tool by its ID with the given ToolUpdate object.""" - # First, check if source code update would cause a name conflict - update_data = tool_update.model_dump(to_orm=True, exclude_none=True) - new_name = None - new_schema = None - # Fetch current tool early to allow conditional logic based on tool type current_tool = await self.get_tool_by_id_async(tool_id=tool_id, actor=actor) - # Do NOT derive schema from Python source. Trust provided JSON schema. - # Prefer provided json_schema; fall back to current - if "json_schema" in update_data: - new_schema = update_data["json_schema"].copy() + # Handle schema updates for custom tools + new_schema = None + if current_tool.tool_type == ToolType.CUSTOM: + if tool_update.json_schema is not None: + new_schema = tool_update.json_schema + elif tool_update.args_json_schema is not None: + # Generate full schema from args_json_schema + generated_schema = generate_schema_for_tool_update( + current_tool=current_tool, + json_schema=None, + args_json_schema=tool_update.args_json_schema, + source_code=tool_update.source_code, + source_type=tool_update.source_type, + ) + if generated_schema: + tool_update.json_schema = generated_schema + new_schema = generated_schema + + # Now model_dump with the potentially updated schema + update_data = tool_update.model_dump(to_orm=True, exclude_none=True) + + # Determine the final schema and name + if new_schema: + new_name = new_schema.get("name", current_tool.name) + elif "json_schema" in update_data: + new_schema = update_data["json_schema"] new_name = new_schema.get("name", current_tool.name) else: + # Keep existing schema new_schema = current_tool.json_schema new_name = current_tool.name - # original tool may no have a JSON schema at all for legacy reasons - # in this case, fallback to dangerous schema generation - if new_schema is None: - # Get source_type from update_data if present, otherwise use current tool's source_type - source_type = update_data.get("source_type", current_tool.source_type) - if source_type == "typescript": - from letta.functions.typescript_parser import derive_typescript_json_schema - - new_schema = derive_typescript_json_schema(source_code=update_data["source_code"]) - else: - new_schema = derive_openai_json_schema(source_code=update_data["source_code"]) + # Handle explicit name updates + if "name" in update_data and update_data["name"] != current_tool.name: + # Name is being explicitly changed + new_name = update_data["name"] + # Update the json_schema name to match if there's a schema + if new_schema: + new_schema = new_schema.copy() + new_schema["name"] = new_name + update_data["json_schema"] = new_schema + elif new_schema and new_name != current_tool.name: + # Schema provides a different name but name wasn't explicitly changed + update_data["name"] = new_name + #raise ValueError( + # f"JSON schema name '{new_name}' conflicts with current tool name '{current_tool.name}'. Update the name field explicitly if you want to rename the tool." + #) # If name changes, enforce uniqueness if new_name != current_tool.name: @@ -532,7 +583,9 @@ class ToolManager: # NOTE: EXTREMELEY HACKY, we need to stop making assumptions about the source_code if "source_code" in update_data and f"def {new_name}" not in update_data.get("source_code", ""): raise LettaToolNameSchemaMismatchError( - tool_name=new_name, json_schema_name=new_schema.get("name"), source_code=update_data.get("source_code") + tool_name=new_name, + json_schema_name=new_schema.get("name") if new_schema else None, + source_code=update_data.get("source_code"), ) # Now perform the update within the session diff --git a/letta/services/tool_schema_generator.py b/letta/services/tool_schema_generator.py new file mode 100644 index 00000000..4868d7d2 --- /dev/null +++ b/letta/services/tool_schema_generator.py @@ -0,0 +1,123 @@ +"""Schema generation utilities for tool creation and updates.""" + +from typing import Optional + +from letta.functions.ast_parsers import get_function_name_and_docstring +from letta.functions.functions import derive_openai_json_schema +from letta.functions.helpers import generate_model_from_args_json_schema +from letta.functions.schema_generator import generate_schema_from_args_schema_v2 +from letta.log import get_logger +from letta.schemas.enums import ToolSourceType, ToolType +from letta.schemas.tool import Tool as PydanticTool + +logger = get_logger(__name__) + + +def generate_schema_for_tool_creation( + tool: PydanticTool, +) -> Optional[dict]: + """ + Generate JSON schema for tool creation based on the provided parameters. + + Args: + tool: The tool being created + + Returns: + Generated JSON schema or None if not applicable + """ + # Only generate schema for custom tools + if tool.tool_type != ToolType.CUSTOM: + return None + + # If json_schema is already provided, use it + if tool.json_schema: + return tool.json_schema + + # Must have source code for custom tools + if not tool.source_code: + logger.error("Custom tool is missing source_code field") + raise ValueError("Custom tool is missing source_code field.") + + # TypeScript tools + if tool.source_type == ToolSourceType.typescript: + try: + from letta.functions.typescript_parser import derive_typescript_json_schema + + return derive_typescript_json_schema(source_code=tool.source_code) + except Exception as e: + logger.error(f"Failed to derive TypeScript json schema: {e}") + raise ValueError(f"Failed to derive TypeScript json schema: {e}") + + # Python tools (default if not specified for backwards compatibility) + elif tool.source_type == ToolSourceType.python or tool.source_type is None: + # If args_json_schema is provided, use it to generate full schema + if tool.args_json_schema: + name, description = get_function_name_and_docstring(tool.source_code, tool.name) + args_schema = generate_model_from_args_json_schema(tool.args_json_schema) + return generate_schema_from_args_schema_v2( + args_schema=args_schema, + name=name, + description=description, + append_heartbeat=False, + ) + # Otherwise, attempt to parse from docstring with best effort + else: + try: + return derive_openai_json_schema(source_code=tool.source_code) + except Exception as e: + logger.error(f"Failed to derive json schema: {e}") + raise ValueError(f"Failed to derive json schema: {e}") + else: + # TODO: convert to explicit error + raise ValueError(f"Unknown tool source type: {tool.source_type}") + + +def generate_schema_for_tool_update( + current_tool: PydanticTool, + json_schema: Optional[dict] = None, + args_json_schema: Optional[dict] = None, + source_code: Optional[str] = None, + source_type: Optional[ToolSourceType] = None, +) -> Optional[dict]: + """ + Generate JSON schema for tool update based on the provided parameters. + + Args: + current_tool: The current tool being updated + json_schema: Directly provided JSON schema (takes precedence) + args_json_schema: Schema for just the arguments + source_code: Updated source code (only used if explicitly updating source) + source_type: Source type for the tool + + Returns: + Updated JSON schema or None if no update needed + """ + # Only handle custom tools + if current_tool.tool_type != ToolType.CUSTOM: + return None + + # If json_schema is directly provided, use it + if json_schema is not None: + # If args_json_schema is also provided, that's an error + if args_json_schema is not None: + raise ValueError("Cannot provide both json_schema and args_json_schema in update") + return json_schema + + # If args_json_schema is provided, generate full schema from it + if args_json_schema is not None: + # Use updated source_code if provided, otherwise use current + code_to_parse = source_code if source_code is not None else current_tool.source_code + if not code_to_parse: + raise ValueError("Source code required when updating with args_json_schema") + + name, description = get_function_name_and_docstring(code_to_parse, current_tool.name) + args_schema = generate_model_from_args_json_schema(args_json_schema) + return generate_schema_from_args_schema_v2( + args_schema=args_schema, + name=name, + description=description, + append_heartbeat=False, + ) + + # Otherwise, no schema updates (don't parse docstring) + return None diff --git a/tests/managers/test_tool_manager.py b/tests/managers/test_tool_manager.py index bf3e2125..6f041f50 100644 --- a/tests/managers/test_tool_manager.py +++ b/tests/managers/test_tool_manager.py @@ -96,6 +96,7 @@ from letta.server.server import SyncServer from letta.services.block_manager import BlockManager from letta.services.helpers.agent_manager_helper import calculate_base_tools, calculate_multi_agent_tools, validate_agent_exists_async from letta.services.step_manager import FeedbackType +from letta.services.tool_schema_generator import generate_schema_for_tool_creation from letta.settings import settings, tool_settings from letta.utils import calculate_file_defaults_based_on_context_window from tests.helpers.utils import comprehensive_agent_checks, validate_context_window_overview @@ -689,23 +690,25 @@ async def test_list_tools_with_tool_types(server: SyncServer, default_user): # create custom tools custom_tool1 = PydanticTool( - name="calculator", + name="calculator_tool", description="Math tool", source_code=parse_source_code(calculator_tool), source_type="python", tool_type=ToolType.CUSTOM, ) - custom_tool1.json_schema = derive_openai_json_schema(source_code=custom_tool1.source_code, name=custom_tool1.name) + # Use generate_schema_for_tool_creation to generate schema + custom_tool1.json_schema = generate_schema_for_tool_creation(custom_tool1) custom_tool1 = await server.tool_manager.create_or_update_tool_async(custom_tool1, actor=default_user) custom_tool2 = PydanticTool( - name="weather", + # name="weather_tool", description="Weather tool", source_code=parse_source_code(weather_tool), source_type="python", tool_type=ToolType.CUSTOM, ) - custom_tool2.json_schema = derive_openai_json_schema(source_code=custom_tool2.source_code, name=custom_tool2.name) + # Use generate_schema_for_tool_creation to generate schema + custom_tool2.json_schema = generate_schema_for_tool_creation(custom_tool2) custom_tool2 = await server.tool_manager.create_or_update_tool_async(custom_tool2, actor=default_user) # test filtering by single tool type @@ -744,13 +747,13 @@ async def test_list_tools_with_exclude_tool_types(server: SyncServer, default_us return msg special = PydanticTool( - name="special", + name="special_tool", description="Special tool", source_code=parse_source_code(special_tool), source_type="python", tool_type=ToolType.CUSTOM, ) - special.json_schema = derive_openai_json_schema(source_code=special.source_code, name=special.name) + special.json_schema = generate_schema_for_tool_creation(special) special = await server.tool_manager.create_or_update_tool_async(special, actor=default_user) # test excluding EXTERNAL_MCP (should get all tools since none are MCP) @@ -796,15 +799,15 @@ async def test_list_tools_with_names(server: SyncServer, default_user): return "gamma" alpha = PydanticTool(name="alpha_tool", description="Alpha", source_code=parse_source_code(alpha_tool), source_type="python") - alpha.json_schema = derive_openai_json_schema(source_code=alpha.source_code, name=alpha.name) + alpha.json_schema = generate_schema_for_tool_creation(alpha) alpha = await server.tool_manager.create_or_update_tool_async(alpha, actor=default_user) beta = PydanticTool(name="beta_tool", description="Beta", source_code=parse_source_code(beta_tool), source_type="python") - beta.json_schema = derive_openai_json_schema(source_code=beta.source_code, name=beta.name) + beta.json_schema = generate_schema_for_tool_creation(beta) beta = await server.tool_manager.create_or_update_tool_async(beta, actor=default_user) gamma = PydanticTool(name="gamma_tool", description="Gamma", source_code=parse_source_code(gamma_tool), source_type="python") - gamma.json_schema = derive_openai_json_schema(source_code=gamma.source_code, name=gamma.name) + gamma.json_schema = generate_schema_for_tool_creation(gamma) gamma = await server.tool_manager.create_or_update_tool_async(gamma, actor=default_user) # test filtering by single name @@ -852,15 +855,15 @@ async def test_list_tools_with_tool_ids(server: SyncServer, default_user): return "3" t1 = PydanticTool(name="tool1", description="First", source_code=parse_source_code(tool1), source_type="python") - t1.json_schema = derive_openai_json_schema(source_code=t1.source_code, name=t1.name) + t1.json_schema = generate_schema_for_tool_creation(t1) t1 = await server.tool_manager.create_or_update_tool_async(t1, actor=default_user) t2 = PydanticTool(name="tool2", description="Second", source_code=parse_source_code(tool2), source_type="python") - t2.json_schema = derive_openai_json_schema(source_code=t2.source_code, name=t2.name) + t2.json_schema = generate_schema_for_tool_creation(t2) t2 = await server.tool_manager.create_or_update_tool_async(t2, actor=default_user) t3 = PydanticTool(name="tool3", description="Third", source_code=parse_source_code(tool3), source_type="python") - t3.json_schema = derive_openai_json_schema(source_code=t3.source_code, name=t3.name) + t3.json_schema = generate_schema_for_tool_creation(t3) t3 = await server.tool_manager.create_or_update_tool_async(t3, actor=default_user) # test filtering by single id @@ -910,19 +913,19 @@ async def test_list_tools_with_search(server: SyncServer, default_user): calc_add = PydanticTool( name="calculator_add", description="Add numbers", source_code=parse_source_code(calculator_add), source_type="python" ) - calc_add.json_schema = derive_openai_json_schema(source_code=calc_add.source_code, name=calc_add.name) + calc_add.json_schema = generate_schema_for_tool_creation(calc_add) calc_add = await server.tool_manager.create_or_update_tool_async(calc_add, actor=default_user) calc_sub = PydanticTool( name="calculator_subtract", description="Subtract numbers", source_code=parse_source_code(calculator_subtract), source_type="python" ) - calc_sub.json_schema = derive_openai_json_schema(source_code=calc_sub.source_code, name=calc_sub.name) + calc_sub.json_schema = generate_schema_for_tool_creation(calc_sub) calc_sub = await server.tool_manager.create_or_update_tool_async(calc_sub, actor=default_user) weather = PydanticTool( name="weather_forecast", description="Weather", source_code=parse_source_code(weather_forecast), source_type="python" ) - weather.json_schema = derive_openai_json_schema(source_code=weather.source_code, name=weather.name) + weather.json_schema = generate_schema_for_tool_creation(weather) weather = await server.tool_manager.create_or_update_tool_async(weather, actor=default_user) # test searching for "calculator" (should find both calculator tools) @@ -965,7 +968,7 @@ async def test_list_tools_return_only_letta_tools(server: SyncServer, default_us source_type="python", tool_type=ToolType.CUSTOM, ) - custom.json_schema = derive_openai_json_schema(source_code=custom.source_code, name=custom.name) + custom.json_schema = generate_schema_for_tool_creation(custom) custom = await server.tool_manager.create_or_update_tool_async(custom, actor=default_user) # test without filter (should get custom tool + all letta tools) @@ -1013,48 +1016,45 @@ async def test_list_tools_combined_filters(server: SyncServer, default_user): return "weather" calc1 = PydanticTool( - name="calculator_add", description="Add", source_code=parse_source_code(calc_add), source_type="python", tool_type=ToolType.CUSTOM + name="calc_add", description="Add", source_code=parse_source_code(calc_add), source_type="python", tool_type=ToolType.CUSTOM ) - calc1.json_schema = derive_openai_json_schema(source_code=calc1.source_code, name=calc1.name) + calc1.json_schema = generate_schema_for_tool_creation(calc1) calc1 = await server.tool_manager.create_or_update_tool_async(calc1, actor=default_user) calc2 = PydanticTool( - name="calculator_multiply", description="Multiply", source_code=parse_source_code(calc_multiply), source_type="python", tool_type=ToolType.CUSTOM, ) - calc2.json_schema = derive_openai_json_schema(source_code=calc2.source_code, name=calc2.name) + calc2.json_schema = generate_schema_for_tool_creation(calc2) calc2 = await server.tool_manager.create_or_update_tool_async(calc2, actor=default_user) weather = PydanticTool( - name="weather_current", + name="weather_tool", description="Weather", source_code=parse_source_code(weather_tool), source_type="python", tool_type=ToolType.CUSTOM, ) - weather.json_schema = derive_openai_json_schema(source_code=weather.source_code, name=weather.name) + weather.json_schema = generate_schema_for_tool_creation(weather) weather = await server.tool_manager.create_or_update_tool_async(weather, actor=default_user) # combine search with tool_types tools = await server.tool_manager.list_tools_async( - actor=default_user, search="calculator", tool_types=[ToolType.CUSTOM.value], upsert_base_tools=False + actor=default_user, search="calc", tool_types=[ToolType.CUSTOM.value], upsert_base_tools=False ) assert len(tools) == 2 - assert all("calculator" in t.name and t.tool_type == ToolType.CUSTOM for t in tools) + assert all("calc" in t.name and t.tool_type == ToolType.CUSTOM for t in tools) # combine names with tool_ids - tools = await server.tool_manager.list_tools_async( - actor=default_user, names=["calculator_add"], tool_ids=[calc1.id], upsert_base_tools=False - ) + tools = await server.tool_manager.list_tools_async(actor=default_user, names=["calc_add"], tool_ids=[calc1.id], upsert_base_tools=False) assert len(tools) == 1 assert tools[0].id == calc1.id # combine search with exclude_tool_types tools = await server.tool_manager.list_tools_async( - actor=default_user, search="calculator", exclude_tool_types=[ToolType.EXTERNAL_MCP.value], upsert_base_tools=False + actor=default_user, search="cal", exclude_tool_types=[ToolType.EXTERNAL_MCP.value], upsert_base_tools=False ) assert len(tools) == 2 @@ -1091,13 +1091,13 @@ async def test_count_tools_async(server: SyncServer, default_user): ta = PydanticTool( name="tool_a", description="A", source_code=parse_source_code(tool_a), source_type="python", tool_type=ToolType.CUSTOM ) - ta.json_schema = derive_openai_json_schema(source_code=ta.source_code, name=ta.name) + ta.json_schema = generate_schema_for_tool_creation(ta) ta = await server.tool_manager.create_or_update_tool_async(ta, actor=default_user) tb = PydanticTool( name="tool_b", description="B", source_code=parse_source_code(tool_b), source_type="python", tool_type=ToolType.CUSTOM ) - tb.json_schema = derive_openai_json_schema(source_code=tb.source_code, name=tb.name) + tb.json_schema = generate_schema_for_tool_creation(tb) tb = await server.tool_manager.create_or_update_tool_async(tb, actor=default_user) # upsert base tools to ensure we have Letta tools for counting @@ -1167,8 +1167,8 @@ async def test_update_tool_by_id(server: SyncServer, print_tool, default_user): assert updated_tool.tool_type == ToolType.EXTERNAL_MCP -#@pytest.mark.asyncio -#async def test_update_tool_source_code_refreshes_schema_and_name(server: SyncServer, print_tool, default_user): +# @pytest.mark.asyncio +# async def test_update_tool_source_code_refreshes_schema_and_name(server: SyncServer, print_tool, default_user): # def counter_tool(counter: int): # """ # Args: @@ -1205,8 +1205,8 @@ async def test_update_tool_by_id(server: SyncServer, print_tool, default_user): # assert updated_tool.tool_type == ToolType.CUSTOM -#@pytest.mark.asyncio -#async def test_update_tool_source_code_refreshes_schema_only(server: SyncServer, print_tool, default_user): +# @pytest.mark.asyncio +# async def test_update_tool_source_code_refreshes_schema_only(server: SyncServer, print_tool, default_user): # def counter_tool(counter: int): # """ # Args: @@ -1701,7 +1701,7 @@ async def test_create_tool_with_pip_requirements(server: SyncServer, default_use tool = PydanticTool( description=description, tags=tags, source_code=source_code, source_type=source_type, metadata_=metadata, pip_requirements=pip_reqs ) - derived_json_schema = derive_openai_json_schema(source_code=tool.source_code, name=tool.name) + derived_json_schema = generate_schema_for_tool_creation(tool) derived_name = derived_json_schema["name"] tool.json_schema = derived_json_schema tool.name = derived_name @@ -1767,7 +1767,7 @@ async def test_update_tool_clear_pip_requirements(server: SyncServer, default_us tool = PydanticTool( description=description, tags=tags, source_code=source_code, source_type=source_type, metadata_=metadata, pip_requirements=pip_reqs ) - derived_json_schema = derive_openai_json_schema(source_code=tool.source_code, name=tool.name) + derived_json_schema = generate_schema_for_tool_creation(tool) derived_name = derived_json_schema["name"] tool.json_schema = derived_json_schema tool.name = derived_name @@ -1817,7 +1817,7 @@ async def test_pip_requirements_roundtrip(server: SyncServer, default_user, defa tool = PydanticTool( description=description, tags=tags, source_code=source_code, source_type=source_type, metadata_=metadata, pip_requirements=pip_reqs ) - derived_json_schema = derive_openai_json_schema(source_code=tool.source_code, name=tool.name) + derived_json_schema = generate_schema_for_tool_creation(tool) derived_name = derived_json_schema["name"] tool.json_schema = derived_json_schema tool.name = derived_name @@ -1859,3 +1859,360 @@ async def test_update_default_requires_approval(server: SyncServer, bash_tool, d # Assertions assert updated_tool.default_requires_approval == True + + +# ====================================================================================================================== +# ToolManager Schema tests +# ====================================================================================================================== + + +async def test_create_tool_with_json_schema(server: SyncServer, default_user, default_organization): + """Test that json_schema is used when provided at creation.""" + tool_manager = server.tool_manager + + source_code = """ +def test_function(arg1: str) -> str: + return arg1 +""" + + json_schema = { + "name": "test_function", + "description": "A test function", + "parameters": {"type": "object", "properties": {"arg1": {"type": "string"}}, "required": ["arg1"]}, + } + + tool = PydanticTool( + name="test_function", + tool_type=ToolType.CUSTOM, + source_code=source_code, + json_schema=json_schema, + ) + + created_tool = await tool_manager.create_tool_async(tool, default_user) + + assert created_tool.json_schema == json_schema + assert created_tool.name == "test_function" + assert created_tool.description == "A test function" + + +async def test_create_tool_with_args_json_schema(server: SyncServer, default_user, default_organization): + """Test that schema is generated from args_json_schema at creation.""" + tool_manager = server.tool_manager + + source_code = """ +def test_function(arg1: str, arg2: int) -> str: + '''This is a test function''' + return f"{arg1} {arg2}" +""" + + args_json_schema = { + "type": "object", + "properties": { + "arg1": {"type": "string", "description": "First argument"}, + "arg2": {"type": "integer", "description": "Second argument"}, + }, + "required": ["arg1", "arg2"], + } + + tool = PydanticTool( + name="test_function", + tool_type=ToolType.CUSTOM, + source_code=source_code, + args_json_schema=args_json_schema, + ) + + created_tool = await tool_manager.create_or_update_tool_async(tool, default_user) + + assert created_tool.json_schema is not None + assert created_tool.json_schema["name"] == "test_function" + assert created_tool.json_schema["description"] == "This is a test function" + assert "parameters" in created_tool.json_schema + assert created_tool.json_schema["parameters"]["properties"]["arg1"]["type"] == "string" + assert created_tool.json_schema["parameters"]["properties"]["arg2"]["type"] == "integer" + + +async def test_create_tool_with_docstring_no_schema(server: SyncServer, default_user, default_organization): + """Test that schema is generated from docstring when no schema provided.""" + tool_manager = server.tool_manager + + source_code = """ +def test_function(arg1: str, arg2: int = 5) -> str: + ''' + This is a test function + + Args: + arg1: First argument + arg2: Second argument + + Returns: + A string result + ''' + return f"{arg1} {arg2}" +""" + + tool = PydanticTool( + name="test_function", + tool_type=ToolType.CUSTOM, + source_code=source_code, + ) + + created_tool = await tool_manager.create_or_update_tool_async(tool, default_user) + + assert created_tool.json_schema is not None + assert created_tool.json_schema["name"] == "test_function" + assert "This is a test function" in created_tool.json_schema["description"] + assert "parameters" in created_tool.json_schema + + +async def test_create_tool_with_docstring_and_args_schema(server: SyncServer, default_user, default_organization): + """Test that args_json_schema takes precedence over docstring.""" + tool_manager = server.tool_manager + + source_code = """ +def test_function(old_arg: str) -> str: + '''Old docstring that should be overridden''' + return old_arg +""" + + args_json_schema = { + "type": "object", + "properties": {"new_arg": {"type": "string", "description": "New argument from schema"}}, + "required": ["new_arg"], + } + + tool = PydanticTool( + name="test_function", + tool_type=ToolType.CUSTOM, + source_code=source_code, + args_json_schema=args_json_schema, + ) + + created_tool = await tool_manager.create_or_update_tool_async(tool, default_user) + + assert created_tool.json_schema is not None + assert created_tool.json_schema["name"] == "test_function" + # The description should come from the docstring + assert created_tool.json_schema["description"] == "Old docstring that should be overridden" + # But parameters should come from args_json_schema + assert "new_arg" in created_tool.json_schema["parameters"]["properties"] + assert "old_arg" not in created_tool.json_schema["parameters"]["properties"] + + +async def test_error_no_docstring_or_schema(server: SyncServer, default_user, default_organization): + """Test error when no docstring or schema provided (minimal function).""" + tool_manager = server.tool_manager + + # Minimal function with no docstring - should still derive basic schema + source_code = """ +def test_function(): + pass +""" + + tool = PydanticTool( + name="test_function", + tool_type=ToolType.CUSTOM, + source_code=source_code, + ) + + with pytest.raises(ValueError) as exc_info: + created_tool = await tool_manager.create_or_update_tool_async(tool, default_user) + + +async def test_error_on_create_tool_with_name_conflict(server: SyncServer, default_user, default_organization): + """Test error when json_schema name conflicts with function name.""" + tool_manager = server.tool_manager + + source_code = """ +def test_function(arg1: str) -> str: + return arg1 +""" + + # JSON schema with conflicting name + json_schema = { + "name": "different_name", + "description": "A test function", + "parameters": {"type": "object", "properties": {"arg1": {"type": "string"}}, "required": ["arg1"]}, + } + + tool = PydanticTool( + name="test_function", + tool_type=ToolType.CUSTOM, + source_code=source_code, + json_schema=json_schema, + ) + + # This should succeed at creation - the tool name takes precedence + created_tool = await tool_manager.create_tool_async(tool, default_user) + assert created_tool.name == "test_function" + + +async def test_update_tool_with_json_schema(server: SyncServer, default_user, default_organization): + """Test update with a new json_schema.""" + tool_manager = server.tool_manager + + # Create initial tool + source_code = """ +def test_function() -> str: + return "hello" +""" + + tool = PydanticTool( + name="test_update_json_schema", + tool_type=ToolType.CUSTOM, + source_code=source_code, + json_schema={"name": "test_update_json_schema", "description": "Original"}, + ) + + created_tool = await tool_manager.create_tool_async(tool, default_user) + + # Update with new json_schema + new_schema = { + "name": "test_update_json_schema", + "description": "Updated description", + "parameters": {"type": "object", "properties": {"new_arg": {"type": "string"}}, "required": ["new_arg"]}, + } + + update = ToolUpdate(json_schema=new_schema) + updated_tool = await tool_manager.update_tool_by_id_async(created_tool.id, update, default_user) + + assert updated_tool.json_schema == new_schema + assert updated_tool.json_schema["description"] == "Updated description" + + +async def test_update_tool_with_args_json_schema(server: SyncServer, default_user, default_organization): + """Test update with args_json_schema.""" + tool_manager = server.tool_manager + + # Create initial tool + source_code = """ +def test_function() -> str: + '''Original function''' + return "hello" +""" + + tool = PydanticTool( + name="test_function", + tool_type=ToolType.CUSTOM, + source_code=source_code, + ) + + created_tool = await tool_manager.create_or_update_tool_async(tool, default_user) + + # Update with args_json_schema + new_source_code = """ +def test_function(new_arg: str) -> str: + '''Updated function''' + return new_arg +""" + + args_json_schema = { + "type": "object", + "properties": {"new_arg": {"type": "string", "description": "New argument"}}, + "required": ["new_arg"], + } + + update = ToolUpdate(source_code=new_source_code, args_json_schema=args_json_schema) + updated_tool = await tool_manager.update_tool_by_id_async(created_tool.id, update, default_user) + + assert updated_tool.json_schema is not None + assert updated_tool.json_schema["description"] == "Updated function" + assert "new_arg" in updated_tool.json_schema["parameters"]["properties"] + + +async def test_update_tool_with_no_schema(server: SyncServer, default_user, default_organization): + """Test update with no schema changes.""" + tool_manager = server.tool_manager + + # Create initial tool + original_schema = { + "name": "test_no_schema_update", + "description": "Original description", + "parameters": {"type": "object", "properties": {}}, + } + + tool = PydanticTool( + name="test_no_schema_update", + tool_type=ToolType.CUSTOM, + source_code="def test_function(): pass", + json_schema=original_schema, + ) + + created_tool = await tool_manager.create_tool_async(tool, default_user) + + # Update with only description (no schema change) + update = ToolUpdate(description="New description") + updated_tool = await tool_manager.update_tool_by_id_async(created_tool.id, update, default_user) + + # Schema should remain unchanged + assert updated_tool.json_schema == original_schema + assert updated_tool.description == "New description" + + +async def test_update_tool_name(server: SyncServer, default_user, default_organization): + """Test various name update scenarios.""" + tool_manager = server.tool_manager + + # Create initial tool + original_schema = {"name": "original_name", "description": "Test", "parameters": {"type": "object", "properties": {}}} + + tool = PydanticTool( + name="original_name", + tool_type=ToolType.CUSTOM, + source_code="def original_name(): pass", + json_schema=original_schema, + ) + + created_tool = await tool_manager.create_or_update_tool_async(tool, default_user) + assert created_tool.name == "original_name" + assert created_tool.json_schema["name"] == "original_name" + + matching_schema = {"name": "matched_name", "description": "Test", "parameters": {"type": "object", "properties": {}}} + update = ToolUpdate(json_schema=matching_schema) + updated_tool3 = await tool_manager.update_tool_by_id_async(created_tool.id, update, default_user) + assert updated_tool3.name == "matched_name" + assert updated_tool3.json_schema["name"] == "matched_name" + + +@pytest.mark.asyncio +async def test_list_tools_with_corrupted_tool(server: SyncServer, default_user, print_tool): + """Test that list_tools still works even if there's a corrupted tool (missing json_schema) in the database.""" + + # First, verify we have a normal tool + tools = await server.tool_manager.list_tools_async(actor=default_user, upsert_base_tools=False) + initial_tool_count = len(tools) + assert any(t.id == print_tool.id for t in tools) + + # Now insert a corrupted tool directly into the database (bypassing normal validation) + # This simulates a tool that somehow got corrupted in the database + from letta.orm.tool import Tool as ToolModel + + async with db_registry.async_session() as session: + # Create a tool with no json_schema (corrupted state) + corrupted_tool = ToolModel( + id=f"tool-corrupted-{uuid.uuid4()}", + name="corrupted_tool", + description="This tool has no json_schema", + tool_type=ToolType.CUSTOM, + source_code="def corrupted_tool(): pass", + json_schema=None, # Explicitly set to None to simulate corruption + organization_id=default_user.organization_id, + created_by_id=default_user.id, + last_updated_by_id=default_user.id, + tags=["corrupted"], + ) + + session.add(corrupted_tool) + await session.commit() + corrupted_tool_id = corrupted_tool.id + + # Now try to list tools - it should still work and not include the corrupted tool + # The corrupted tool should be automatically excluded from results + tools = await server.tool_manager.list_tools_async(actor=default_user, upsert_base_tools=False) + + # Verify listing still works + assert len(tools) == initial_tool_count # Corrupted tool should not be in the results + assert any(t.id == print_tool.id for t in tools) # Normal tool should still be there + assert not any(t.id == corrupted_tool_id for t in tools) # Corrupted tool should not be there + + # Verify the corrupted tool's name is not in the results + assert not any(t.name == "corrupted_tool" for t in tools) diff --git a/tests/test_sdk_client.py b/tests/test_sdk_client.py index a8629670..1327c367 100644 --- a/tests/test_sdk_client.py +++ b/tests/test_sdk_client.py @@ -50,7 +50,8 @@ def run_server(): start_server(debug=True) -@pytest.fixture(scope="module") +# @pytest.fixture(scope="module") +@pytest.fixture(scope="function") def client() -> LettaSDKClient: # Get URL from environment or start server server_url = os.getenv("LETTA_SERVER_URL", f"http://localhost:{SERVER_PORT}") @@ -65,8 +66,8 @@ def client() -> LettaSDKClient: yield client -@pytest.fixture(scope="module") -def server(): +@pytest.fixture(scope="function") +async def server(): """ Creates a SyncServer instance for testing. @@ -74,7 +75,9 @@ def server(): """ config = LettaConfig.load() config.save() - return SyncServer() + server = SyncServer() + await server.init_async() + return server @pytest.fixture(scope="function") diff --git a/tests/test_tool_schema_parsing.py b/tests/test_tool_schema_parsing.py index 71bdf88b..d4399826 100644 --- a/tests/test_tool_schema_parsing.py +++ b/tests/test_tool_schema_parsing.py @@ -277,42 +277,6 @@ def _run_composio_test(action_name, openai_model, structured_output): return (action_name, False, str(e)) # Failure with error message -@pytest.mark.parametrize("openai_model", ["gpt-4o-mini"]) -@pytest.mark.parametrize("structured_output", [True]) -def test_composio_tool_schema_generation(openai_model: str, structured_output: bool): - """Test that we can generate the schemas for some Composio tools.""" - - if not os.getenv("COMPOSIO_API_KEY"): - pytest.skip("COMPOSIO_API_KEY not set") - - start_time = time.time() - - action_names = [ - "GITHUB_STAR_A_REPOSITORY_FOR_THE_AUTHENTICATED_USER", # Simple - "CAL_GET_AVAILABLE_SLOTS_INFO", # has an array arg, needs to be converted properly - "SALESFORCE_RETRIEVE_LEAD_BY_ID", # has an array arg, needs to be converted properly - "FIRECRAWL_SEARCH", # has an optional array arg, needs to be converted properly - ] - - # Create a pool of processes - pool = mp.Pool(processes=min(mp.cpu_count(), len(action_names))) - - # Map the work to the pool - func = partial(_run_composio_test, openai_model=openai_model, structured_output=structured_output) - results = pool.map(func, action_names) - - # Check results - for action_name, success, error_message in results: - print(f"Test for {action_name}: {'SUCCESS' if success else 'FAILED - ' + error_message}") - assert success, f"Test for {action_name} failed: {error_message}" - - pool.close() - pool.join() - - end_time = time.time() - print(f"Total execution time: {end_time - start_time:.2f} seconds") - - # Helper function for pydantic args schema test def _run_pydantic_args_test(filename, openai_model, structured_output): """Run a single pydantic args schema test case""" @@ -342,6 +306,9 @@ def _run_pydantic_args_test(filename, openai_model, structured_output): source_code=last_function_source, args_json_schema=args_schema, ) + from letta.services.tool_schema_generator import generate_schema_for_tool_creation + + tool.json_schema = generate_schema_for_tool_creation(tool) schema = tool.json_schema # We expect this to fail for all_python_complex with structured_output=True