diff --git a/aidial_assistant/commands/open_api.py b/aidial_assistant/commands/open_api.py index 8f2a679..135e51a 100644 --- a/aidial_assistant/commands/open_api.py +++ b/aidial_assistant/commands/open_api.py @@ -8,7 +8,10 @@ ExecutionCallback, ResultObject, ) -from aidial_assistant.open_api.requester import OpenAPIEndpointRequester +from aidial_assistant.open_api.requester import ( + OpenAPIEndpointRequester, + ParamMapping, +) class OpenAPIChatCommand(Command): @@ -16,14 +19,24 @@ class OpenAPIChatCommand(Command): def token() -> str: return "open-api-chat-command" - def __init__(self, op: APIOperation, plugin_auth: str | None): - self.op = op - self.plugin_auth = plugin_auth + def __init__(self, requester: OpenAPIEndpointRequester): + self.requester = requester @override async def execute( self, args: dict[str, Any], execution_callback: ExecutionCallback ) -> ResultObject: - return await OpenAPIEndpointRequester( - self.op, self.plugin_auth - ).execute(args) + return await self.requester.execute(args) + + @classmethod + def create( + cls, base_url: str, operation: APIOperation, auth: str | None + ) -> "OpenAPIChatCommand": + path = base_url.rstrip("/") + operation.path + method = operation.method + param_mapping = ParamMapping( + query_params=operation.query_params, + body_params=operation.body_params, + path_params=operation.path_params, + ) + return cls(OpenAPIEndpointRequester(path, method, param_mapping, auth)) diff --git a/aidial_assistant/commands/plugin_callback.py b/aidial_assistant/commands/plugin_callback.py index 84770be..bf706ac 100644 --- a/aidial_assistant/commands/plugin_callback.py +++ b/aidial_assistant/commands/plugin_callback.py @@ -79,7 +79,7 @@ def on_state(self, request: str, response: str): @override def on_error(self, title: str, error: str): - pass + self.callback(f"```\n{title}: {error}\n```\n") @property def result(self) -> str: diff --git a/aidial_assistant/commands/run_plugin.py b/aidial_assistant/commands/run_plugin.py index 96f2913..8fbd138 100644 --- a/aidial_assistant/commands/run_plugin.py +++ b/aidial_assistant/commands/run_plugin.py @@ -25,7 +25,6 @@ ModelClient, ReasonLengthException, ) -from aidial_assistant.open_api.operation_selector import collect_operations from aidial_assistant.utils.open_ai import user_message from aidial_assistant.utils.open_ai_plugin import OpenAIPluginInfo @@ -62,17 +61,27 @@ async def _run_plugin( self, query: str, execution_callback: ExecutionCallback ) -> ResultObject: info = self.plugin.info - ops = collect_operations(info.open_api, info.ai_plugin.api.url) - api_schema = "\n\n".join([op.to_typescript() for op in ops.values()]) # type: ignore + spec = info.open_api + spec_url = info.get_full_spec_url() + operations = [ + APIOperation.from_openapi_spec(spec, path, method) + for path in spec.paths + for method in spec.get_methods_for_path(path) + ] + api_schema = "\n\n".join([op.to_typescript() for op in operations]) # type: ignore def create_command(op: APIOperation): - return lambda: OpenAPIChatCommand(op, self.plugin.auth) - - command_dict: dict[str, CommandConstructor] = {} - for name, op in ops.items(): # The function is necessary to capture the current value of op. # Otherwise, only first op will be used for all commands - command_dict[name] = create_command(op) + return lambda: OpenAPIChatCommand.create( + spec_url, op, self.plugin.auth + ) + + command_dict: dict[str, CommandConstructor] = { + operation.operation_id: create_command(operation) + for operation in operations + } + command_names = command_dict.keys() if Reply.token() in command_dict: Exception(f"Operation with name '{Reply.token()}' is not allowed.") @@ -80,7 +89,7 @@ def create_command(op: APIOperation): history = History( assistant_system_message_template=ADDON_SYSTEM_DIALOG_MESSAGE.build( - command_names=ops.keys(), + command_names=command_names, api_description=info.ai_plugin.description_for_model, api_schema=api_schema, ), diff --git a/aidial_assistant/commands/run_tool.py b/aidial_assistant/commands/run_tool.py index 5c6b4a2..f82ff51 100644 --- a/aidial_assistant/commands/run_tool.py +++ b/aidial_assistant/commands/run_tool.py @@ -1,14 +1,11 @@ from typing import Any -from langchain_community.tools.openapi.utils.api_models import ( - APIOperation, - APIPropertyBase, -) -from openai.types.chat import ChatCompletionToolParam +from langchain_community.tools.openapi.utils.api_models import APIOperation from typing_extensions import override from aidial_assistant.commands.base import ( Command, + CommandConstructor, ExecutionCallback, ResultObject, TextResult, @@ -21,46 +18,9 @@ ModelClient, ReasonLengthException, ) -from aidial_assistant.open_api.operation_selector import collect_operations -from aidial_assistant.tools_chain.tools_chain import ( - CommandTool, - CommandToolDict, - ToolsChain, -) -from aidial_assistant.utils.open_ai import ( - construct_tool, - system_message, - user_message, -) - - -def _construct_property(p: APIPropertyBase) -> dict[str, Any]: - parameter = { - "type": p.type, - "description": p.description, - } - return {k: v for k, v in parameter.items() if v is not None} - - -def _construct_tool(op: APIOperation) -> ChatCompletionToolParam: - properties = {} - required = [] - for p in op.properties: - properties[p.name] = _construct_property(p) - - if p.required: - required.append(p.name) - - if op.request_body is not None: - for p in op.request_body.properties: - properties[p.name] = _construct_property(p) - - if p.required: - required.append(p.name) - - return construct_tool( - op.operation_id, op.description or "", properties, required - ) +from aidial_assistant.tools_chain.tools_chain import CommandToolDict, ToolsChain +from aidial_assistant.utils.open_ai import system_message, user_message +from aidial_assistant.utils.open_api import construct_tool_from_spec class RunTool(Command): @@ -81,17 +41,26 @@ async def execute( ) -> ResultObject: query = get_required_field(args, "query") - ops = collect_operations( - self.plugin.info.open_api, self.plugin.info.ai_plugin.api.url - ) + spec = self.plugin.info.open_api + spec_url = self.plugin.info.get_full_spec_url() - def create_command_tool(op: APIOperation) -> CommandTool: - return lambda: OpenAPIChatCommand( - op, self.plugin.auth - ), _construct_tool(op) + def create_command(operation: APIOperation) -> CommandConstructor: + # The function is necessary to capture the current value of op. + # Otherwise, only first op will be used for all commands + return lambda: OpenAPIChatCommand.create( + spec_url, operation, self.plugin.auth + ) commands: CommandToolDict = { - name: create_command_tool(op) for name, op in ops.items() + operation.operation_id: (create_command(operation), tool) + for path in spec.paths + for operation, tool in ( + ( + APIOperation.from_openapi_spec(spec, path, method), + construct_tool_from_spec(spec, path, method), + ) + for method in spec.get_methods_for_path(path) + ) } chain = ToolsChain(self.model, commands, self.max_completion_tokens) diff --git a/aidial_assistant/open_api/operation_selector.py b/aidial_assistant/open_api/operation_selector.py deleted file mode 100644 index 9f8a005..0000000 --- a/aidial_assistant/open_api/operation_selector.py +++ /dev/null @@ -1,51 +0,0 @@ -import json -from typing import Union -from urllib.parse import urljoin - -from langchain.tools import APIOperation, OpenAPISpec -from pydantic import BaseModel - - -class OpenAPICommand(BaseModel): - command: str - args: dict - - -class OpenAPIClarification(BaseModel): - user_question: str - - -OpenAPIResponse = Union[OpenAPICommand, OpenAPIClarification] - - -class OpenAPIResponseWrapper(BaseModel): - """Just a wrapper class for the union to ease parsing""" - - resp: OpenAPIResponse - - @staticmethod - def parse_str(s) -> OpenAPIResponse: - return OpenAPIResponseWrapper.parse_obj({"resp": json.loads(s)}).resp - - -OpenAPIOperations = dict[str, APIOperation] - - -def collect_operations(spec: OpenAPISpec, spec_url: str) -> OpenAPIOperations: - operations: dict[str, APIOperation] = {} - - def add_operation(spec, path, method): - operation = APIOperation.from_openapi_spec(spec, path, method) # type: ignore - operation.base_url = urljoin(spec_url, operation.base_url) - operations[operation.operation_id] = operation - - if spec.paths is None: # type: ignore - return operations - - for path, path_item in spec.paths.items(): # type: ignore - if path_item.get is not None: - add_operation(spec, path, "get") - if path_item.post is not None: - add_operation(spec, path, "post") - - return operations diff --git a/aidial_assistant/open_api/requester.py b/aidial_assistant/open_api/requester.py index b46665e..5da3e41 100644 --- a/aidial_assistant/open_api/requester.py +++ b/aidial_assistant/open_api/requester.py @@ -4,7 +4,6 @@ import aiohttp.client_exceptions from aiohttp import hdrs -from langchain.tools.openapi.utils.api_models import APIOperation from aidial_assistant.commands.base import JsonResult, ResultObject, TextResult from aidial_assistant.utils.requests import arequest @@ -12,7 +11,7 @@ logger = logging.getLogger(__name__) -class _ParamMapping(NamedTuple): +class ParamMapping(NamedTuple): """Mapping from parameter name to parameter value.""" query_params: List[str] @@ -25,18 +24,21 @@ class OpenAPIEndpointRequester: Based on OpenAPIEndpointChain from LangChain. """ - def __init__(self, operation: APIOperation, plugin_auth: str | None): - self.operation = operation - self.param_mapping = _ParamMapping( - query_params=operation.query_params, # type: ignore - body_params=operation.body_params, # type: ignore - path_params=operation.path_params, # type: ignore - ) + def __init__( + self, + url: str, + method: str, + param_mapping: ParamMapping, + plugin_auth: str | None, + ): + self.url = url + self.method = method + self.param_mapping = param_mapping self.plugin_auth = plugin_auth def _construct_path(self, args: Dict[str, str]) -> str: """Construct the path from the deserialized input.""" - path = self.operation.base_url.rstrip("/") + self.operation.path # type: ignore + path = self.url for param in self.param_mapping.path_params: path = path.replace(f"{{{param}}}", str(args.pop(param, ""))) return path @@ -87,17 +89,16 @@ async def execute( ) logger.debug(f"Request args: {request_args}") async with arequest( - self.operation.method.value, headers=headers, **request_args # type: ignore + self.method, headers=headers, **request_args # type: ignore ) as response: if response.status != 200: try: return JsonResult(json.dumps(await response.json())) except aiohttp.ContentTypeError: - method_str = str(self.operation.method.value) # type: ignore error_object = { "reason": response.reason, "status_code": response.status, - "method:": method_str.upper(), + "method:": self.method.upper(), "url": request_args["url"], "params": request_args["params"], } diff --git a/aidial_assistant/tools_chain/tools_chain.py b/aidial_assistant/tools_chain/tools_chain.py index ca10121..1e9f0cc 100644 --- a/aidial_assistant/tools_chain/tools_chain.py +++ b/aidial_assistant/tools_chain/tools_chain.py @@ -1,4 +1,5 @@ import json +import logging from typing import Any, Tuple, cast from openai import BadRequestError @@ -37,6 +38,8 @@ from aidial_assistant.utils.exceptions import RequestParameterValidationError from aidial_assistant.utils.open_ai import tool_calls_message, tool_message +logger = logging.getLogger(__name__) + def convert_commands_to_tools( scoped_messages: list[ScopedMessage], @@ -99,12 +102,12 @@ def convert_commands_to_tools( def _publish_command( - command_callback: CommandCallback, name: str, arguments: str + command_callback: CommandCallback, name: str, arguments: dict[str, Any] ): command_callback.on_command(name) args_callback = command_callback.args_callback() args_callback.on_args_start() - args_callback.on_args_chunk(arguments) + args_callback.on_args_chunk(json.dumps(arguments)) args_callback.on_args_end() @@ -209,37 +212,53 @@ async def _run_tools( tool_calls: list[ChatCompletionMessageToolCallParam], callback: ChainCallback, ): - commands: list[CommandInvocation] = [] - command_results: list[CommandResult] = [] + state: list[Tuple[CommandInvocation, CommandResult]] = [] result_messages: list[ChatCompletionMessageParam] = [] for tool_call in tool_calls: function = tool_call["function"] name = function["name"] - arguments: dict[str, Any] = json.loads(function["arguments"]) - with callback.command_callback() as command_callback: - _publish_command(command_callback, name, json.dumps(arguments)) + try: command = self._create_command(name) - result = await self._execute_command( - command, - arguments, - command_callback, + arguments: dict[str, Any] = ToolsChain._parse_arguments( + function["arguments"] ) - result_messages.append( - tool_message( - content=result["response"], - tool_call_id=tool_call["id"], + with callback.command_callback() as command_callback: + _publish_command(command_callback, name, arguments) + result = await self._execute_command( + command, + arguments, + command_callback, ) + result_messages.append( + tool_message( + content=result["response"], + tool_call_id=tool_call["id"], + ) + ) + state.append( + ( + CommandInvocation( + command=name, arguments=arguments + ), + result, + ) + ) + + except AssistantProtocolException as e: + logger.exception("Failed to process model response") + callback.on_error( + "Error", "The model failed to establish the protocol." + ) + result_messages.append( + tool_message(content=str(e), tool_call_id=tool_call["id"]) ) - command_results.append(result) - commands.append( - CommandInvocation(command=name, arguments=arguments) + if state: + commands, results = zip(*state) + callback.on_state( + commands_to_text(commands), responses_to_text(results) ) - callback.on_state( - commands_to_text(commands), responses_to_text(command_results) - ) - return result_messages @staticmethod @@ -257,3 +276,12 @@ async def _execute_command( except Exception as e: command_callback.on_error(e) return CommandResult(status=Status.ERROR, response=str(e)) + + @staticmethod + def _parse_arguments(arguments: str) -> dict[str, Any]: + try: + return json.loads(arguments) + except json.JSONDecodeError as e: + raise AssistantProtocolException( + f"Failed to parse arguments: {e}" + ) from e diff --git a/aidial_assistant/utils/open_ai_plugin.py b/aidial_assistant/utils/open_ai_plugin.py index 822d197..12241b5 100644 --- a/aidial_assistant/utils/open_ai_plugin.py +++ b/aidial_assistant/utils/open_ai_plugin.py @@ -43,6 +43,9 @@ class OpenAIPluginInfo(BaseModel): ai_plugin: AIPluginConf open_api: OpenAPISpec + def get_full_spec_url(self) -> str: + return urljoin(self.ai_plugin.api.url, self.open_api.base_url) + class AddonTokenSource: def __init__(self, headers: Mapping[str, str], urls: Iterable[str]): diff --git a/aidial_assistant/utils/open_api.py b/aidial_assistant/utils/open_api.py new file mode 100644 index 0000000..22a6325 --- /dev/null +++ b/aidial_assistant/utils/open_api.py @@ -0,0 +1,75 @@ +from typing import Any, Iterable, Tuple + +from langchain_community.utilities.openapi import OpenAPISpec +from openai.types.chat import ChatCompletionToolParam +from openapi_pydantic import DataType, Reference, Schema + +from aidial_assistant.utils.open_ai import construct_tool + + +def _resolve_schema(spec: OpenAPISpec, schema: Schema | Reference) -> Schema: + if isinstance(schema, Reference): + return spec.get_referenced_schema(schema) + + return schema + + +def _construct_property( + spec: OpenAPISpec, schema: Schema | Reference +) -> dict[str, Any]: + return _resolve_schema(spec, schema).dict(exclude_none=True) + + +def _extract_body_parameters( + spec: OpenAPISpec, schema: Schema | Reference +) -> Iterable[Tuple[str, dict[str, Any], bool]]: + schema = _resolve_schema(spec, schema) + if schema.type != DataType.OBJECT: + raise ValueError("Body schema must be an object") + + if schema.properties: + required = schema.required or [] + for prop_name, prop_schema in schema.properties.items(): + prop_schema = _resolve_schema(spec, prop_schema) + + yield ( + prop_name, + _construct_property(spec, prop_schema), + prop_name in required, + ) + + +def construct_tool_from_spec( + spec: OpenAPISpec, path: str, method: str +) -> ChatCompletionToolParam: + operation = spec.get_operation(path, method) + properties: dict[str, Any] = {} + required = [] + for p in spec.get_parameters_for_operation(operation): + if p.param_schema is None: + raise ValueError(f"Parameter {p.name} has no schema") + + properties[p.name] = _construct_property(spec, p.param_schema) + + if p.required: + required.append(p.name) + + request_body = spec.get_request_body_for_operation(operation) + if request_body is not None: + for key, media_type in request_body.content.items(): + if key == "application/json": + if media_type.media_type_schema is None: + raise ValueError("Body has no schema") + + for name, prop, is_required in _extract_body_parameters( + spec, media_type.media_type_schema + ): + properties[name] = prop + if is_required: + required.append(name) + break + + operation_id = OpenAPISpec.get_cleaned_operation_id(operation, path, method) + return construct_tool( + operation_id, operation.description or "", properties, required + ) diff --git a/tests/unit_tests/utils/test_open_api.py b/tests/unit_tests/utils/test_open_api.py new file mode 100644 index 0000000..70826ca --- /dev/null +++ b/tests/unit_tests/utils/test_open_api.py @@ -0,0 +1,136 @@ +import pytest +from langchain_community.utilities.openapi import OpenAPISpec + +from aidial_assistant.utils.open_api import construct_tool_from_spec + +OPEN_API_SPEC = { + "openapi": "3.0.2", + "info": { + "title": "Test API title", + "description": "Test API description", + "version": "0.0.1", + }, + "servers": [{"url": ".."}], + "paths": { + "/path_with_parameters": { + "get": { + "description": "Tool with parameters", + "operationId": "id_with_parameters", + "parameters": [ + { + "name": "param1", + "in": "query", + "required": True, + "schema": { + "type": "string", + "description": "First parameter", + }, + }, + { + "name": "param2", + "in": "query", + "schema": { + "type": "string", + "description": "Second parameter", + }, + }, + ], + } + }, + "/path_with_body_parameters": { + "post": { + "description": "Tool with body parameters", + "operationId": "id_with_body_parameters", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SchemaWithBodyParameters" + } + } + }, + }, + } + }, + }, + "components": { + "schemas": { + "SchemaWithBodyParameters": { + "required": ["param1"], + "type": "object", + "properties": { + "param1": { + "type": "string", + "description": "First parameter", + }, + "param2": { + "type": "string", + "description": "Second parameter", + }, + }, + } + } + }, +} + +EXPECTED_TOOLS = [ + ( + "/path_with_parameters", + "get", + { + "type": "function", + "function": { + "name": "id_with_parameters", + "description": "Tool with parameters", + "parameters": { + "type": "object", + "properties": { + "param1": { + "type": "string", + "description": "First parameter", + }, + "param2": { + "type": "string", + "description": "Second parameter", + }, + }, + "required": ["param1"], + }, + }, + }, + ), + ( + "/path_with_body_parameters", + "post", + { + "type": "function", + "function": { + "name": "id_with_body_parameters", + "description": "Tool with body parameters", + "parameters": { + "type": "object", + "properties": { + "param1": { + "type": "string", + "description": "First parameter", + }, + "param2": { + "type": "string", + "description": "Second parameter", + }, + }, + "required": ["param1"], + }, + }, + }, + ), +] + + +@pytest.mark.parametrize("path,method,expected", EXPECTED_TOOLS) +def test_construct_tool_from_spec(path, method, expected): + actual = construct_tool_from_spec( + OpenAPISpec.from_spec_dict(OPEN_API_SPEC), path, method + ) + + assert actual == expected