From dd979e07a82217ca4ca5ff06806d9a593e0fc942 Mon Sep 17 00:00:00 2001 From: Yukinobu Mine Date: Fri, 13 Dec 2024 00:07:00 +0900 Subject: [PATCH] Remove title key from input JSON schema of AgentTool. #629 --- backend/app/agents/tools/agent_tool.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/backend/app/agents/tools/agent_tool.py b/backend/app/agents/tools/agent_tool.py index 7af079e4..42d9b7bd 100644 --- a/backend/app/agents/tools/agent_tool.py +++ b/backend/app/agents/tools/agent_tool.py @@ -11,6 +11,7 @@ from app.repositories.models.custom_bot import BotModel from app.routes.schemas.conversation import type_model_name from pydantic import BaseModel, JsonValue +from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue from mypy_boto3_bedrock_runtime.type_defs import ( ToolSpecificationTypeDef, ) @@ -49,6 +50,16 @@ class InvalidToolError(Exception): pass +class RemoveTitle(GenerateJsonSchema): + def field_title_should_be_set(self, schema) -> bool: + return False + + def generate(self, schema, mode="validation") -> JsonSchemaValue: + value = super().generate(schema, mode) + del value["title"] + return value + + class AgentTool(Generic[T]): def __init__( self, @@ -71,7 +82,7 @@ def __init__( def _generate_input_schema(self) -> dict[str, Any]: """Converts the Pydantic model to a JSON schema.""" - return self.args_schema.model_json_schema() + return self.args_schema.model_json_schema(schema_generator=RemoveTitle) def to_converse_spec(self) -> ToolSpecificationTypeDef: return ToolSpecificationTypeDef(