Support Tool Calling with Language Agent (#110)

Author: nqyy

mbodied/agents/backends/anthropic_backend.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from typing import Any, List
 
 import anthropic
@@ -80,7 +81,7 @@ def __init__(self, api_key: str | None, client: anthropic.Anthropic | None = Non
             client: An optional client for the Anthropic service.
             kwargs: Additional keyword arguments.
         """
-        self.api_key = api_key
+        self.api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
         self.client = client
 
         self.model = kwargs.pop("model", self.DEFAULT_MODEL)
@@ -117,13 +118,6 @@ def predict(
         )
         return completion.content[0].text
 
-    async def async_predict(
-        self, message: Message, context: List[Message] | None = None, model: Any | None = None
-    ) -> str:
-        """Asynchronously predict the next message in the conversation."""
-        # For now, we'll use the synchronous method since Anthropic doesn't provide an async API
-        return self.predict(message, context, model)
-
     def stream(
         self, message: Message, context: List[Message] = None, model: str = "claude-3-5-sonnet-20240620", **kwargs
     ):

mbodied/agents/backends/gemini_backend.py

@@ -105,7 +105,7 @@ def __init__(
             client: An optional client for the Gemini service.
             **kwargs: Additional keyword arguments.
         """
-        self.api_key = api_key or os.getenv("GOOGLE_API_KEY") or os.getenv("MBODI_API_KEY")
+        self.api_key = api_key or os.getenv("GOOGLE_API_KEY")
         self.client = client
 
         self.model = kwargs.pop("model", self.DEFAULT_MODEL)

mbodied/agents/backends/openai_backend.py

@@ -25,6 +25,7 @@
 from mbodied.agents.backends.serializer import Serializer
 from mbodied.types.message import Message
 from mbodied.types.sense.vision import Image
+from mbodied.types.tool import Tool, ToolCall
 
 ERRORS = (
     OpenAIRateLimitError,
@@ -99,7 +100,7 @@ def __init__(
             aclient: Whether to use the asynchronous client.
             **kwargs: Additional keyword arguments.
         """
-        self.api_key = api_key or os.getenv("OPENAI_API_KEY") or os.getenv("MBODI_API_KEY")
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
         self.client = client
         if self.client is None:
             from openai import AsyncOpenAI, OpenAI
@@ -119,18 +120,26 @@ def __init__(
         on_backoff=lambda details: print(f"Backing off {details['wait']:.1f} seconds after {details['tries']} tries."),  # noqa
     )
     def predict(
-        self, message: Message, context: List[Message] | None = None, model: Any | None = None, **kwargs
-    ) -> str:
+        self,
+        message: Message,
+        context: List[Message] | None = None,
+        model: Any | None = None,
+        tools: List[Tool] | None = None,
+        **kwargs,
+    ) -> str | tuple[str, List[ToolCall]]:
         """Create a completion based on the given message and context.
 
         Args:
             message (Message): The message to process.
             context (Optional[List[Message]]): The context of messages.
             model (Optional[Any]): The model used for processing the messages.
+            tools (Optional[List[Tool]]): The tools to make available for function calling.
             **kwargs: Additional keyword arguments.
 
         Returns:
-            str: The result of the completion.
+            str | tuple[str, List[ToolCall]]:
+                When tools are not provided: Just the text response
+                When tools are provided: A tuple of (text_response, tool_calls)
         """
         context = context or self.INITIAL_CONTEXT
         model = model or self.DEFAULT_MODEL
@@ -141,18 +150,35 @@ def predict(
             messages=serialized_messages,
             temperature=0,
             max_tokens=1000,
+            tools=tools,
             **kwargs,
         )
+        if tools:
+            tool_calls = []
+            if completion.choices[0].message.tool_calls:
+                for tool_call in completion.choices[0].message.tool_calls:
+                    tool_calls.append(ToolCall.model_validate(tool_call))
+            return completion.choices[0].message.content, tool_calls
+
         return completion.choices[0].message.content
 
-    def stream(self, message: Message, context: List[Message] = None, model: str = "gpt-4o", **kwargs):
+    def stream(
+        self, message: Message, context: List[Message] = None, model: str = "gpt-4o", tools: List[Tool] = None, **kwargs
+    ):
         """Streams a completion for the given messages using the OpenAI API standard.
 
         Args:
             message: Message to be sent to the completion API.
             context: The context of the messages.
             model: The model to be used for the completion.
+            tools: Optional list of tools (function calls) available to the model.
             **kwargs: Additional keyword arguments.
+
+        Yields:
+            When tools is None:
+                str: Content delta chunks
+            When tools is provided:
+                tuple[str, Any]: Tuples of (content_delta, tool_call_delta) where either may be None
         """
         model = model or self.DEFAULT_MODEL
         context = context or self.INITIAL_CONTEXT
@@ -162,19 +188,37 @@ def stream(self, message: Message, context: List[Message] = None, model: str = "
             model=model,
             temperature=0,
             stream=True,
+            tools=tools,
             **kwargs,
         )
-        for chunk in stream:
-            yield chunk.choices[0].delta.content or ""
 
-    async def astream(self, message: Message, context: List[Message] = None, model: str = "gpt-4o", **kwargs):
+        if not tools:
+            for chunk in stream:
+                yield chunk.choices[0].delta.content or ""
+        else:
+            for chunk in stream:
+                delta = chunk.choices[0].delta
+                content = delta.content or ""
+                tool_calls = delta.tool_calls
+                yield content, tool_calls
+
+    async def astream(
+        self, message: Message, context: List[Message] = None, model: str = "gpt-4o", tools: List[Tool] = None, **kwargs
+    ):
         """Streams a completion asynchronously for the given messages using the OpenAI API standard.
 
         Args:
             message: Message to be sent to the completion API.
             context: The context of the messages.
             model: The model to be used for the completion.
+            tools: Optional list of tools (function calls) available to the model.
             **kwargs: Additional keyword arguments.
+
+        Yields:
+            When tools is None:
+                str: Content delta chunks
+            When tools is provided:
+                tuple[str, Any]: Tuples of (content_delta, tool_call_delta) where either may be None
         """
         if not hasattr(self, "aclient"):
             raise AttributeError("AsyncOpenAI client not initialized. Pass in aclient=True to the constructor.")
@@ -186,7 +230,16 @@ async def astream(self, message: Message, context: List[Message] = None, model:
             model=model,
             temperature=0,
             stream=True,
+            tools=tools,
             **kwargs,
         )
-        async for chunk in stream:
-            yield chunk.choices[0].delta.content or ""
+
+        if not tools:
+            async for chunk in stream:
+                yield chunk.choices[0].delta.content or ""
+        else:
+            async for chunk in stream:
+                delta = chunk.choices[0].delta
+                content = delta.content or ""
+                tool_calls = delta.tool_calls
+                yield content, tool_calls