From 6fdb38e321de3112a46636a77974438aed822cc2 Mon Sep 17 00:00:00 2001 From: cpacker Date: Mon, 30 Dec 2024 20:41:41 +0100 Subject: [PATCH] fix: patch bug in json generator for composio --- letta/functions/schema_generator.py | 53 ++--------------------------- letta/schemas/tool.py | 1 + tests/test_tool_schema_parsing.py | 24 +++++++++++-- 3 files changed, 25 insertions(+), 53 deletions(-) diff --git a/letta/functions/schema_generator.py b/letta/functions/schema_generator.py index 89409cb2..6f5bb52f 100644 --- a/letta/functions/schema_generator.py +++ b/letta/functions/schema_generator.py @@ -396,44 +396,6 @@ def generate_schema(function, name: Optional[str] = None, description: Optional[ return schema -def generate_schema_from_args_schema_v1( - args_schema: Type[V1BaseModel], name: Optional[str] = None, description: Optional[str] = None, append_heartbeat: bool = True -) -> Dict[str, Any]: - properties = {} - required = [] - for field_name, field in args_schema.__fields__.items(): - if field.type_ == str: - field_type = "string" - elif field.type_ == int: - field_type = "integer" - elif field.type_ == bool: - field_type = "boolean" - else: - field_type = field.type_.__name__ - - properties[field_name] = { - "type": field_type, - "description": field.field_info.description, - } - if field.required: - required.append(field_name) - - function_call_json = { - "name": name, - "description": description, - "parameters": {"type": "object", "properties": properties, "required": required}, - } - - if append_heartbeat: - function_call_json["parameters"]["properties"]["request_heartbeat"] = { - "type": "boolean", - "description": "Request an immediate heartbeat after function execution. Set to `True` if you want to send a follow-up message or run a follow-up function.", - } - function_call_json["parameters"]["required"].append("request_heartbeat") - - return function_call_json - - def generate_schema_from_args_schema_v2( args_schema: Type[BaseModel], name: Optional[str] = None, description: Optional[str] = None, append_heartbeat: bool = True ) -> Dict[str, Any]: @@ -441,19 +403,8 @@ def generate_schema_from_args_schema_v2( required = [] for field_name, field in args_schema.model_fields.items(): field_type_annotation = field.annotation - if field_type_annotation == str: - field_type = "string" - elif field_type_annotation == int: - field_type = "integer" - elif field_type_annotation == bool: - field_type = "boolean" - else: - field_type = field_type_annotation.__name__ - - properties[field_name] = { - "type": field_type, - "description": field.description, - } + properties[field_name] = type_to_json_schema_type(field_type_annotation) + properties[field_name]["description"] = field.description if field.is_required(): required.append(field_name) diff --git a/letta/schemas/tool.py b/letta/schemas/tool.py index 8066f9b2..4870356c 100644 --- a/letta/schemas/tool.py +++ b/letta/schemas/tool.py @@ -127,6 +127,7 @@ class ToolCreate(LettaBase): source_type = "python" tags = ["composio"] wrapper_func_name, wrapper_function_str = generate_composio_tool_wrapper(action_name) + print("composio_tool.args_schema", str(composio_tool.args_schema)) json_schema = generate_schema_from_args_schema_v2(composio_tool.args_schema, name=wrapper_func_name, description=description) return cls( diff --git a/tests/test_tool_schema_parsing.py b/tests/test_tool_schema_parsing.py index f6738a06..e6afdc29 100644 --- a/tests/test_tool_schema_parsing.py +++ b/tests/test_tool_schema_parsing.py @@ -5,6 +5,7 @@ import pytest from letta.functions.functions import derive_openai_json_schema from letta.llm_api.helpers import convert_to_structured_output, make_post_request +from letta.schemas.tool import ToolCreate def _clean_diff(d1, d2): @@ -154,8 +155,10 @@ def _load_schema_from_source_filename(filename: str) -> dict: # @pytest.mark.parametrize("openai_model", ["gpt-4o-mini"]) # @pytest.mark.parametrize("structured_output", [True]) -@pytest.mark.parametrize("openai_model", ["gpt-4", "gpt-4o"]) -@pytest.mark.parametrize("structured_output", [True, False]) +# @pytest.mark.parametrize("openai_model", ["gpt-4", "gpt-4o"]) +# @pytest.mark.parametrize("structured_output", [True, False]) +@pytest.mark.parametrize("openai_model", ["gpt-4o-mini"]) +@pytest.mark.parametrize("structured_output", [True]) def test_valid_schemas_via_openai(openai_model: str, structured_output: bool): """Test that we can send the schemas to OpenAI and get a tool call back.""" @@ -176,3 +179,20 @@ def test_valid_schemas_via_openai(openai_model: str, structured_output: bool): _openai_payload(openai_model, schema, structured_output) else: _openai_payload(openai_model, schema, structured_output) + + +@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.""" + + assert os.getenv("COMPOSIO_API_KEY") is not None, "COMPOSIO_API_KEY must be set" + + for action_name in [ + "CAL_GET_AVAILABLE_SLOTS_INFO", # has an array arg, needs to be converted properly + ]: + tool_create = ToolCreate.from_composio(action_name=action_name) + print(tool_create) + + schema = tool_create.json_schema + _openai_payload(openai_model, schema, structured_output)