feat: support native model ability to invoke tools

aidial_assistant/app.py

@@ -8,7 +8,6 @@
 from aidial_assistant.utils.log_config import get_log_config
 
 log_level = os.getenv("LOG_LEVEL", "INFO")
-config_dir = Path(os.getenv("CONFIG_DIR", "aidial_assistant/configs"))
 
 logging.config.dictConfig(get_log_config(log_level))
 
@@ -22,7 +21,13 @@
     AssistantApplication,
 )
 
+config_dir = Path(os.getenv("CONFIG_DIR", "aidial_assistant/configs"))
+tools_supporting_deployments: set[str] = set(
+    os.getenv(
+        "TOOLS_SUPPORTING_DEPLOYMENTS", "gpt-4-turbo-1106,anthropic.claude-v2-1"
+    ).split(",")
+)
 app.add_chat_completion(
     "assistant",
-    AssistantApplication(config_dir),
+    AssistantApplication(config_dir, tools_supporting_deployments),
 )

aidial_assistant/application/addons_dialogue_limiter.py

@@ -4,7 +4,10 @@
     LimitExceededException,
     ModelRequestLimiter,
 )
-from aidial_assistant.model.model_client import Message, ModelClient
+from aidial_assistant.model.model_client import (
+    ChatCompletionMessageParam,
+    ModelClient,
+)
 
 
 class AddonsDialogueLimiter(ModelRequestLimiter):
@@ -16,7 +19,7 @@ def __init__(self, max_dialogue_tokens: int, model_client: ModelClient):
         self._initial_tokens: int | None = None
 
     @override
-    async def verify_limit(self, messages: list[Message]):
+    async def verify_limit(self, messages: list[ChatCompletionMessageParam]):
         if self._initial_tokens is None:
             self._initial_tokens = await self.model_client.count_tokens(
                 messages

aidial_assistant/application/assistant_application.py

@@ -1,14 +1,13 @@
-import json
 import logging
 from pathlib import Path
+from typing import Tuple
 
 from aidial_sdk.chat_completion import FinishReason
 from aidial_sdk.chat_completion.base import ChatCompletion
-from aidial_sdk.chat_completion.request import Addon
-from aidial_sdk.chat_completion.request import Message as SdkMessage
-from aidial_sdk.chat_completion.request import Request, Role
+from aidial_sdk.chat_completion.request import Addon, Message, Request, Role
 from aidial_sdk.chat_completion.response import Response
 from openai.lib.azure import AsyncAzureOpenAI
+from openai.types.chat import ChatCompletionToolParam
 from pydantic import BaseModel
 
 from aidial_assistant.application.addons_dialogue_limiter import (
@@ -22,28 +21,29 @@
     MAIN_BEST_EFFORT_TEMPLATE,
     MAIN_SYSTEM_DIALOG_MESSAGE,
 )
-from aidial_assistant.chain.command_chain import CommandChain, CommandDict
-from aidial_assistant.chain.command_result import Commands, Responses
-from aidial_assistant.chain.history import History, MessageScope, ScopedMessage
+from aidial_assistant.chain.command_chain import (
+    CommandChain,
+    CommandConstructor,
+    CommandDict,
+)
+from aidial_assistant.chain.history import History
 from aidial_assistant.commands.reply import Reply
 from aidial_assistant.commands.run_plugin import PluginInfo, RunPlugin
 from aidial_assistant.commands.run_tool import RunTool
 from aidial_assistant.model.model_client import (
-    Message,
     ModelClient,
     ReasonLengthException,
-    ToolCall,
 )
-from aidial_assistant.tools_chain.tools_chain import ToolsChain
+from aidial_assistant.tools_chain.tools_chain import (
+    CommandToolDict,
+    ToolsChain,
+    convert_commands_to_tools,
+)
 from aidial_assistant.utils.exceptions import (
     RequestParameterValidationError,
     unhandled_exception_handler,
 )
-from aidial_assistant.utils.open_ai import (
-    FunctionCall,
-    Tool,
-    construct_function,
-)
+from aidial_assistant.utils.open_ai import construct_tool
 from aidial_assistant.utils.open_ai_plugin import (
     AddonTokenSource,
     get_open_ai_plugin_info,
@@ -83,7 +83,7 @@ def _validate_addons(addons: list[Addon] | None) -> list[AddonReference]:
     return addon_references
 
 
-def _validate_messages(messages: list[SdkMessage]) -> None:
+def _validate_messages(messages: list[Message]) -> None:
     if not messages:
         raise RequestParameterValidationError(
             "Message list cannot be empty.", param="messages"
@@ -95,13 +95,8 @@ def _validate_messages(messages: list[SdkMessage]) -> None:
         )
 
 
-def _validate_request(request: Request) -> None:
-    _validate_messages(request.messages)
-    _validate_addons(request.addons)
-
-
-def _construct_function(name: str, description: str) -> Tool:
-    return construct_function(
+def _construct_tool(name: str, description: str) -> ChatCompletionToolParam:
+    return construct_tool(
         name,
         description,
         {
@@ -114,71 +109,12 @@ def _construct_function(name: str, description: str) -> Tool:
     )
 
 
-def _convert_commands_to_tools(
-    scoped_messages: list[ScopedMessage],
-) -> list[Message]:
-    messages: list[Message] = []
-    next_tool_id: int = 0
-    last_call_count: int = 0
-    for scoped_message in scoped_messages:
-        message = scoped_message.message
-        if scoped_message.scope == MessageScope.INTERNAL:
-            if not message.content:
-                raise RequestParameterValidationError(
-                    "State is broken. Content cannot be empty.",
-                    param="messages",
-                )
-
-            if message.role == Role.ASSISTANT:
-                commands: Commands = json.loads(message.content)
-                messages.append(
-                    Message(
-                        role=Role.ASSISTANT,
-                        tool_calls=[
-                            ToolCall(
-                                index=index,
-                                id=str(next_tool_id + index),
-                                function=FunctionCall(
-                                    name=command["command"],
-                                    arguments=json.dumps(command["arguments"]),
-                                ),
-                                type="function",
-                            )
-                            for index, command in enumerate(
-                                commands["commands"]
-                            )
-                        ],
-                    )
-                )
-                last_call_count = len(commands["commands"])
-                next_tool_id += last_call_count
-            elif message.role == Role.USER:
-                responses: Responses = json.loads(message.content)
-                response_count = len(responses["responses"])
-                if response_count != last_call_count:
-                    raise RequestParameterValidationError(
-                        f"Expected {last_call_count} responses, but got {response_count}.",
-                        param="messages",
-                    )
-                first_tool_id = next_tool_id - last_call_count
-                messages.extend(
-                    [
-                        Message(
-                            role=Role.TOOL,
-                            tool_call_id=str(first_tool_id + index),
-                            content=response["response"],
-                        )
-                        for index, response in enumerate(responses["responses"])
-                    ]
-                )
-        else:
-            messages.append(scoped_message.message)
-    return messages
-
-
 class AssistantApplication(ChatCompletion):
-    def __init__(self, config_dir: Path):
+    def __init__(
+        self, config_dir: Path, tools_supporting_deployments: set[str]
+    ):
         self.args = parse_args(config_dir)
+        self.tools_supporting_deployments = tools_supporting_deployments
 
     @unhandled_exception_handler
     async def chat_completion(
@@ -203,12 +139,12 @@ async def chat_completion(
             (addon_reference.url for addon_reference in addon_references),
         )
 
-        addons: list[PluginInfo] = []
+        plugins: list[PluginInfo] = []
         # DIAL Core has own names for addons, so in stages we need to map them to the names used by the user
         addon_name_mapping: dict[str, str] = {}
         for addon_reference in addon_references:
             info = await get_open_ai_plugin_info(addon_reference.url)
-            addons.append(
+            plugins.append(
                 PluginInfo(
                     info=info,
                     auth=get_plugin_auth(
@@ -225,13 +161,13 @@ async def chat_completion(
                     info.ai_plugin.name_for_model
                 ] = addon_reference.name
 
-        if request.model in {"gpt-4-turbo-1106", "anthropic.claude-v2-1"}:
+        if request.model in self.tools_supporting_deployments:
             await AssistantApplication._run_native_tools_chat(
-                model, addons, addon_name_mapping, request, response
+                model, plugins, addon_name_mapping, request, response
             )
         else:
             await AssistantApplication._run_emulated_tools_chat(
-                model, addons, addon_name_mapping, request, response
+                model, plugins, addon_name_mapping, request, response
             )
 
     @staticmethod
@@ -313,34 +249,31 @@ def create_command(addon: PluginInfo):
     @staticmethod
     async def _run_native_tools_chat(
         model: ModelClient,
-        addons: list[PluginInfo],
+        plugins: list[PluginInfo],
         addon_name_mapping: dict[str, str],
         request: Request,
         response: Response,
     ):
-        tools: list[Tool] = [
-            _construct_function(
-                addon.info.ai_plugin.name_for_model,
-                addon.info.ai_plugin.description_for_human,
+        def create_command_tool(
+            plugin: PluginInfo,
+        ) -> Tuple[CommandConstructor, ChatCompletionToolParam]:
+            return lambda: RunTool(model, plugin), _construct_tool(
+                plugin.info.ai_plugin.name_for_model,
+                plugin.info.ai_plugin.description_for_human,
             )
-            for addon in addons
-        ]
 
-        def create_command(addon: PluginInfo):
-            return lambda: RunTool(model, addon)
-
-        command_dict: CommandDict = {
-            addon.info.ai_plugin.name_for_model: create_command(addon)
-            for addon in addons
+        command_tool_dict: CommandToolDict = {
+            plugin.info.ai_plugin.name_for_model: create_command_tool(plugin)
+            for plugin in plugins
         }
-        chain = ToolsChain(model, tools, command_dict)
+        chain = ToolsChain(model, command_tool_dict)
 
         choice = response.create_single_choice()
         choice.open()
 
         callback = AssistantChainCallback(choice, addon_name_mapping)
         finish_reason = FinishReason.STOP
-        messages = _convert_commands_to_tools(parse_history(request.messages))
+        messages = convert_commands_to_tools(parse_history(request.messages))
         try:
             await chain.run_chat(messages, callback)
         except ReasonLengthException: