chore: import upstream snapshot with attribution

packages/graphrag-llm/graphrag_llm/utils/function_tool_manager.py

@@ -0,0 +1,138 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Function tool manager."""
+
+import json
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+from openai import pydantic_function_tool
+from pydantic import BaseModel
+from typing_extensions import TypedDict
+
+if TYPE_CHECKING:
+    from graphrag_llm.types import LLMCompletionFunctionToolParam, LLMCompletionResponse
+
+FunctionArgumentModel = TypeVar(
+    "FunctionArgumentModel", bound=BaseModel, covariant=True
+)
+
+
+class FunctionDefinition(TypedDict, Generic[FunctionArgumentModel]):
+    """Function definition."""
+
+    name: str
+    description: str
+    input_model: type[FunctionArgumentModel]
+    function: Callable[[FunctionArgumentModel], str]
+
+
+class ToolMessage(TypedDict):
+    """Function tool response message to be added to message history."""
+
+    content: str
+    tool_call_id: str
+
+
+class FunctionToolManager:
+    """Function tool manager."""
+
+    _tools: dict[str, FunctionDefinition[Any]]
+
+    def __init__(self) -> None:
+        """Initialize FunctionToolManager."""
+        self._tools = {}
+
+    def register_function_tool(
+        self,
+        *,
+        name: str,
+        description: str,
+        input_model: type[FunctionArgumentModel],
+        function: Callable[[FunctionArgumentModel], str],
+    ) -> None:
+        """Register function tool.
+
+        Args
+        ----
+            name: str
+                The name of the function tool.
+            description: str
+                The description of the function tool.
+            input_model: type[T]
+                The pydantic model type for the function tool input.
+            function: Callable[[T], str]
+                The function to call for the function tool.
+        """
+        self._tools[name] = {
+            "name": name,
+            "description": description,
+            "input_model": input_model,
+            "function": function,
+        }
+
+    def definitions(self) -> list["LLMCompletionFunctionToolParam"]:
+        """Get function tool definitions.
+
+        Returns
+        -------
+            list[LLMCompletionFunctionToolParam]
+                List of function tool definitions.
+        """
+        return [
+            pydantic_function_tool(
+                tool_def["input_model"],
+                name=tool_def["name"],
+                description=tool_def["description"],
+            )
+            for tool_def in self._tools.values()
+        ]
+
+    def call_functions(self, response: "LLMCompletionResponse") -> list[ToolMessage]:
+        """Call functions based on the response.
+
+        Args
+        ----
+            response: LLMCompletionResponse
+                The LLM completion response.
+
+        Returns
+        -------
+            list[ToolMessage]
+                The list of tool response messages to be added to the message history.
+        """
+        if not response.choices[0].message.tool_calls:
+            return []
+
+        tool_messages: list[ToolMessage] = []
+
+        for tool_call in response.choices[0].message.tool_calls:
+            if tool_call.type != "function":
+                continue
+            tool_id = tool_call.id
+            function_name = tool_call.function.name
+            function_args = tool_call.function.arguments
+
+            if function_name not in self._tools:
+                msg = f"Function '{function_name}' not registered."
+                raise ValueError(msg)
+
+            tool_def = self._tools[function_name]
+            input_model = tool_def["input_model"]
+            function = tool_def["function"]
+
+            try:
+                parsed_args_dict = json.loads(function_args)
+                input_model_instance = input_model(**parsed_args_dict)
+            except Exception as e:
+                msg = f"Failed to parse arguments for function '{function_name}': {e}"
+                raise ValueError(msg) from e
+
+            result = function(input_model_instance)
+            tool_messages.append({
+                "content": result,
+                "tool_call_id": tool_id,
+            })
+
+        return tool_messages