Add tool calling support.

Author: qiuosier

ads/llm/autogen/client_v02.py

@@ -9,17 +9,60 @@
 """
 import copy
 import importlib
+import json
 import logging
 from typing import Dict, List, Union
 from types import SimpleNamespace
 
-
 from autogen import ModelClient
+from autogen.oai.client import OpenAIWrapper
 from langchain_core.messages import AIMessage
 
+
 logger = logging.getLogger(__name__)
 
 
+def register_custom_client(client_class):
+    """Registers custom client with AutoGen."""
+    if not hasattr(OpenAIWrapper, "custom_clients"):
+        raise AttributeError(
+            "The AutoGen version you install does not support auto custom client registration."
+        )
+    if client_class not in OpenAIWrapper.custom_clients:
+        OpenAIWrapper.custom_clients.append(client_class)
+
+
+def _convert_to_langchain_tool(tool):
+    if tool["type"] == "function":
+        tool = tool["function"]
+        required = tool["parameters"]["required"]
+        properties = copy.deepcopy(tool["parameters"]["properties"])
+        for key in properties.keys():
+            val = properties[key]
+            val["default"] = key in required
+        return {
+            "title": tool["name"],
+            "description": tool["description"],
+            "properties": properties,
+        }
+    raise NotImplementedError(f"Type {tool['type']} is not supported.")
+
+
+def _convert_to_openai_tool_call(tool_call):
+    return {
+        "id": tool_call.get("id"),
+        "function": {
+            "name": tool_call.get("name"),
+            "arguments": (
+                ""
+                if tool_call.get("args") is None
+                else json.dumps(tool_call.get("args"))
+            ),
+        },
+        "type": "function",
+    }
+
+
 class Message(AIMessage):
     """Represents message returned from the LLM."""
 
@@ -28,6 +71,9 @@ def from_message(cls, message: AIMessage):
         """Converts from LangChain AIMessage."""
         message = copy.deepcopy(message)
         message.__class__ = cls
+        message.tool_calls = [
+            _convert_to_openai_tool_call(tool) for tool in message.tool_calls
+        ]
         return message
 
     @property
@@ -42,23 +88,55 @@ class LangChainModelClient(ModelClient):
     def __init__(self, config: dict, **kwargs) -> None:
         super().__init__()
         logger.info("LangChain model client config: %s", str(config))
+        # Make a copy of the config since we are popping the keys
+        config = copy.deepcopy(config)
         self.client_class = config.pop("model_client_cls")
-        # Parameters for the model
+
+        self.function_call_params = config.pop("function_call_params", {})
+
         self.model_name = config.get("model")
+
         # Import the LangChain class
         if "langchain_cls" not in config:
             raise ValueError("Missing langchain_cls in LangChain Model Client config.")
         module_cls = config.pop("langchain_cls")
         module_name, cls_name = str(module_cls).rsplit(".", 1)
         langchain_module = importlib.import_module(module_name)
         langchain_cls = getattr(langchain_module, cls_name)
+
         # Initialize the LangChain client
         self.model = langchain_cls(**config)
 
     def create(self, params) -> ModelClient.ModelClientResponseProtocol:
+        """Creates a LLM completion for a given config.
+
+        Parameters
+        ----------
+        params : dict
+            OpenAI API compatible parameters, including all the keys from llm_config.
+
+        Returns
+        -------
+        ModelClientResponseProtocol
+            Response from LLM
+
+        """
         streaming = params.get("stream", False)
+        # TODO: num_of_responses
         num_of_responses = params.get("n", 1)
-        messages = params.get("messages", [])
+        messages = params.pop("messages", [])
+
+        invoke_params = {}
+
+        tools = params.get("tools")
+        if tools:
+            model = self.model.bind_tools(
+                [_convert_to_langchain_tool(tool) for tool in tools]
+            )
+            # invoke_params["tools"] = tools
+            invoke_params.update(self.function_call_params)
+        else:
+            model = self.model
 
         response = SimpleNamespace()
         response.choices = []
@@ -69,7 +147,7 @@ def create(self, params) -> ModelClient.ModelClientResponseProtocol:
             raise NotImplementedError()
         else:
             # If streaming is not enabled, send a regular chat completion request
-            ai_message = self.model.invoke(messages)
+            ai_message = model.invoke(messages, **invoke_params)
             choice = SimpleNamespace()
             choice.message = Message.from_message(ai_message)
             response.choices.append(choice)
@@ -84,7 +162,7 @@ def message_retrieval(
         NOTE: if a list of Choice.Message is returned, it currently needs to contain the fields of OpenAI's ChatCompletion Message object,
         since that is expected for function or tool calling in the rest of the codebase at the moment, unless a custom agent is being used.
         """
-        return [choice.message.content for choice in response.choices]
+        return [choice.message for choice in response.choices]
 
     def cost(self, response: ModelClient.ModelClientResponseProtocol) -> float:
         response.cost = 0