d0d36a4b07
* patch integration test * create default user and org * rm * patch * add testing * remove validation for schemas from pydantic object * add file * patch tests * fix more tests * fix managers * fix sdk test * patch schema tests * Comment out name in update * patch test * patch * add another test
124 lines
4.7 KiB
Python
124 lines
4.7 KiB
Python
"""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
|