move litellm_utils => litellm_lib

Author: dlqqq

packages/jupyter-ai/jupyter_ai/litellm_utils/__init__.py renamed to packages/jupyter-ai/jupyter_ai/litellm_lib/__init__.py

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

packages/jupyter-ai/jupyter_ai/litellm_utils/run_tools.py renamed to packages/jupyter-ai/jupyter_ai/litellm_lib/run_tools.py

@@ -1,16 +1,8 @@
 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

packages/jupyter-ai/jupyter_ai/litellm_utils/test_toolcall_list.py renamed to packages/jupyter-ai/jupyter_ai/litellm_lib/test_toolcall_list.py

@@ -29,10 +29,9 @@ class ResolvedToolCall(BaseModel):
     `litellm.utils.ChatCompletionDeltaToolCall`.
     """
 
-    id: str | None
+    id: str
     """
-    The ID of the tool call. This should always be provided by LiteLLM, this
-    type is left optional as we do not use this attribute.
+    The ID of the tool call.
     """
 
     type: str
@@ -62,7 +61,7 @@ class ToolCallList(BaseModel):
     is used to aggregate the tool call deltas yielded from a LiteLLM response
     stream and produce a list of tool calls.
 
-    After all tool call deltas are added, the `process()` method may be called
+    After all tool call deltas are added, the `resolve()` method may be called
     to return a list of resolved tool calls.
 
     Example usage:
@@ -75,7 +74,7 @@ class ToolCallList(BaseModel):
         tool_call_delta = chunk.choices[0].delta.tool_calls
         tool_call_list += tool_call_delta
     
-    tool_call_list.resolve()
+    tool_calls = tool_call_list.resolve()
     ```
     """
 
@@ -128,7 +127,11 @@ def __add__(self, other: list[ChatCompletionDeltaToolCall] | None) -> 'ToolCallL
 
     def resolve(self) -> list[ResolvedToolCall]:
         """
-        Resolve the aggregated tool call delta lists into a list of tool calls.
+        Returns the aggregated tool calls as `list[ResolvedToolCall]`.
+
+        Raises an exception if any function arguments could not be parsed from
+        JSON into a dictionary. This method should only be called after the
+        stream completed without errors.
         """
         resolved_toolcalls: list[ResolvedToolCall] = []
         for i, raw_toolcall in enumerate(self._aggregate):

packages/jupyter-ai/jupyter_ai/litellm_utils/toolcall_list.py renamed to packages/jupyter-ai/jupyter_ai/litellm_lib/toolcall_list.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from pydantic import BaseModel
 from .toolcall_list import ToolCallList
 
@@ -11,3 +12,10 @@ class StreamResult(BaseModel):
     """
     Tool calls requested by the LLM in its streamed response.
     """
+
+class ToolCallOutput(BaseModel):
+    tool_call_id: str
+    role: str = "tool"
+    name: str
+    content: str
+

packages/jupyter-ai/jupyter_ai/litellm_utils/streaming_utils.py renamed to packages/jupyter-ai/jupyter_ai/litellm_lib/types.py

@@ -17,15 +17,12 @@
 from traitlets.config import LoggingConfigurable
 
 from .persona_awareness import PersonaAwareness
-from ..litellm_utils import ToolCallList, StreamResult, run_tools, ToolCallOutput
-
-# Import toolkits
+from ..litellm_lib import ToolCallList, StreamResult, run_tools
 from ..tools.default_toolkit import DEFAULT_TOOLKIT
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator
     from .persona_manager import PersonaManager
-    from ..tools import Toolkit
 
 class PersonaDefaults(BaseModel):
     """

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

@@ -1,11 +1,10 @@
 from typing import Any, Optional
 import time
-import json
 
 from jupyterlab_chat.models import Message
 from litellm import acompletion
 
-from ...litellm_utils import StreamResult, ToolCallOutput
+from ...litellm_lib import StreamResult, ToolCallOutput
 from ..base_persona import BasePersona, PersonaDefaults
 from ..persona_manager import SYSTEM_USERNAME
 from .prompt_template import (
@@ -110,15 +109,6 @@ def render_tool_calls(self, stream_result: StreamResult):
 
 
     def render_tool_call_outputs(self, message_id: str, tool_call_outputs: list[dict]):
-        # TODO
-        # self.ychat.update_message(Message(
-        #     id=message_id,
-        #     body=f"\n\n```\n{json.dumps(tool_call_outputs, indent=2)}\n```\n",
-        #     sender=self.id,
-        #     time=time.time(),
-        #     raw_time=False
-        # ), append=True)
-
         # Updates the content of the last message directly
         message = self.ychat.get_message(message_id)
         body = message.body
@@ -134,7 +124,6 @@ def render_tool_call_outputs(self, message_id: str, tool_call_outputs: list[dict
                 f'<jai-tool-call id="{tool_id}" output="{tool_output}"',
             )
 
-        self.log.error(body)
         self.ychat.update_message(Message(
             id=message.id,
             time=time.time(),