update tool calling APIs

Author: dlqqq

packages/jupyter-ai/jupyter_ai/litellm_utils/__init__.py

@@ -1,2 +1,3 @@
 from .toolcall_list import *
 from .streaming_utils import *
+from .run_tools import *

packages/jupyter-ai/jupyter_ai/litellm_utils/run_tools.py

@@ -0,0 +1,50 @@
+import asyncio
+from pydantic import BaseModel
+from .toolcall_list import ToolCallList
+from ..tools import Toolkit
+
+
+class ToolCallOutput(BaseModel):
+    tool_call_id: str
+    role: str = "tool"
+    name: str
+    content: str
+
+
+async def run_tools(tool_call_list: ToolCallList, toolkit: Toolkit) -> list[dict]:
+    """
+    Runs the tools specified in the list of tool calls returned by
+    `self.stream_message()`. 
+    
+    Returns `list[ToolCallOutput]`. The outputs should be appended directly to
+    the message history on the next request made to the LLM.
+    """
+    tool_calls = tool_call_list.resolve()
+    if not len(tool_calls):
+        return []
+
+    tool_outputs: list[dict] = []
+    for tool_call in tool_calls:
+        # Get tool definition from the correct toolkit
+        # TODO: validation?
+        tool_name = tool_call.function.name
+        tool_defn = toolkit.get_tool_unsafe(tool_name)
+
+        # Run tool and store its output
+        try:
+            output = tool_defn.callable(**tool_call.function.arguments)
+            if asyncio.iscoroutine(output):
+                output = await output
+        except Exception as e:
+            output = str(e)
+
+        # Store the tool output in a dictionary accepted by LiteLLM
+        output_dict = {
+            "tool_call_id": tool_call.id,
+            "role": "tool",
+            "name": tool_call.function.name,
+            "content": output,
+        }
+        tool_outputs.append(output_dict)
+    
+    return tool_outputs

packages/jupyter-ai/jupyter_ai/litellm_utils/streaming_utils.py

@@ -7,7 +7,7 @@ class StreamResult(BaseModel):
     ID of the new message.
     """
 
-    tool_calls: ToolCallList
+    tool_call_list: ToolCallList
     """
     Tool calls requested by the LLM in its streamed response.
     """

packages/jupyter-ai/jupyter_ai/litellm_utils/toolcall_list.py

@@ -170,4 +170,8 @@ def to_json(self) -> list[dict[str, Any]]:
         return [
             model.model_dump() for model in self._aggregate
         ]
+    
+
+    def __len__(self) -> int:
+        return len(self._aggregate)
 

packages/jupyter-ai/jupyter_ai/personas/base_persona.py

@@ -17,7 +17,7 @@
 from traitlets.config import LoggingConfigurable
 
 from .persona_awareness import PersonaAwareness
-from ..litellm_utils import ToolCallList, StreamResult, ResolvedToolCall
+from ..litellm_utils import ToolCallList, StreamResult, run_tools, ToolCallOutput
 
 # Import toolkits
 from ..tools.default_toolkit import DEFAULT_TOOLKIT
@@ -255,7 +255,7 @@ async def stream_message(
         """
         stream_id: Optional[str] = None
         stream_interrupted = False
-        tool_calls = ToolCallList()
+        tool_call_list = ToolCallList()
         try:
             self.awareness.set_local_state_field("isWriting", True)
 
@@ -319,7 +319,7 @@ async def stream_message(
                         append=True,
                     )
                 if toolcalls_delta:
-                    tool_calls += toolcalls_delta
+                    tool_call_list += toolcalls_delta
 
         except Exception as e:
             self.log.error(
@@ -348,15 +348,13 @@ async def stream_message(
                     return None
 
             # TODO: determine where this should live
-            resolved_toolcalls = tool_calls.resolve()
-            if len(resolved_toolcalls):
-                count = len(resolved_toolcalls)
-                names = sorted([tc.function.name for tc in resolved_toolcalls])
-                self.log.info(f"AI response triggered {count} tool calls: {names}")
+            count = len(tool_call_list)
+            if count > 0:
+                self.log.info(f"AI response triggered {count} tool calls.")
 
             return StreamResult(
                 id=stream_id,
-                tool_calls=tool_calls
+                tool_call_list=tool_call_list
             )
 
 
@@ -520,38 +518,13 @@ def get_tools(self, model_id: str) -> list[dict]:
         return tool_descriptions
 
 
-    async def run_tools(self, tools: list[ResolvedToolCall]) -> list[dict]:
+    async def run_tools(self, tool_call_list: ToolCallList) -> list[ToolCallOutput]:
         """
-        Runs the tools specified in the list of tool calls returned by
-        `self.stream_message()`. Returns a list of dictionaries
-        `toolcall_outputs: list[dict]`, which should be appended directly to the
-        message history on the next invocation of the LLM.
+        Runs the tools specified in a given tool call list using the default
+        toolkit.
         """
-        if not len(tools):
-            return []
-
-        tool_outputs: list[dict] = []
-        for tool_call in tools:
-            # Get tool definition from the correct toolkit
-            # TODO: validation?
-            tool_name = tool_call.function.name
-            tool_defn = DEFAULT_TOOLKIT.get_tool_unsafe(tool_name)
-
-            # Run tool and store its output
-            output = tool_defn.callable(**tool_call.function.arguments)
-            if asyncio.iscoroutine(output):
-                output = await output
-
-            # Store the tool output in a dictionary accepted by LiteLLM
-            output_dict = {
-                "tool_call_id": tool_call.id,
-                "role": "tool",
-                "name": tool_call.function.name,
-                "content": output,
-            }
-            tool_outputs.append(output_dict)
-        
-        self.log.info(f"Ran {len(tools)} tool functions.")
+        tool_outputs = await run_tools(tool_call_list, toolkit=DEFAULT_TOOLKIT)
+        self.log.info(f"Ran {len(tool_outputs)} tool functions.")
         return tool_outputs
 
 

packages/jupyter-ai/jupyter_ai/personas/jupyternaut/jupyternaut.py

@@ -62,7 +62,7 @@ async def process_message(self, message: Message) -> None:
 
             # Append new reply to `messages`
             reply = self.ychat.get_message(result.id)
-            tool_calls_json = result.tool_calls.to_json()
+            tool_calls_json = result.tool_call_list.to_json()
             messages.append({
                 "role": "assistant",
                 "content": reply.body,
@@ -80,7 +80,7 @@ async def process_message(self, message: Message) -> None:
                 ), append=True)
 
             # Run tools and append outputs to `messages`
-            tool_call_outputs = await self.run_tools(result.tool_calls.resolve())
+            tool_call_outputs = await self.run_tools(result.tool_call_list)
             messages.extend(tool_call_outputs)
 
             # Add tool call outputs to YChat (not synced with `messages`)