Adopt new function calling that supports Groq and other models (#569)

Co-authored-by: Hejia Zhang <[email protected]>

pyproject.toml

@@ -28,7 +28,7 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
 ]
 dependencies = [
-    "openai >= 1.6.1, < 2.0.0",
+    "openai >= 1.34.0, < 2.0.0",
     "typer >= 0.7.0, < 1.0.0",
     "click >= 7.1.1, < 9.0.0",
     "rich >= 13.1.0, < 14.0.0",

sgpt/function.py

@@ -2,7 +2,7 @@
 import sys
 from abc import ABCMeta
 from pathlib import Path
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, Union
 
 from .config import cfg
 
@@ -59,4 +59,15 @@ def get_function(name: str) -> Callable[..., Any]:
 
 
 def get_openai_schemas() -> List[Dict[str, Any]]:
-    return [function.openai_schema for function in functions]
+    transformed_schemas = []
+    for function in functions:
+        schema = {
+            "type": "function",
+            "function": {
+                "name": function.openai_schema["name"],
+                "description": function.openai_schema.get("description", ""),
+                "parameters": function.openai_schema.get("parameters", {}),
+            },
+        }
+        transformed_schemas.append(schema)
+    return transformed_schemas

sgpt/handlers/handler.py

@@ -9,6 +9,7 @@
 from ..role import DefaultRoles, SystemRole
 
 completion: Callable[..., Any] = lambda *args, **kwargs: Generator[Any, None, None]
+
 base_url = cfg.get("API_BASE_URL")
 use_litellm = cfg.get("USE_LITELLM") == "true"
 additional_kwargs = {
@@ -89,36 +90,43 @@ def get_completion(
         messages: List[Dict[str, Any]],
         functions: Optional[List[Dict[str, str]]],
     ) -> Generator[str, None, None]:
+
         name = arguments = ""
         is_shell_role = self.role.name == DefaultRoles.SHELL.value
         is_code_role = self.role.name == DefaultRoles.CODE.value
         is_dsc_shell_role = self.role.name == DefaultRoles.DESCRIBE_SHELL.value
         if is_shell_role or is_code_role or is_dsc_shell_role:
             functions = None
 
+        if functions:
+            additional_kwargs["tool_choice"] = "auto"
+            additional_kwargs["tools"] = functions
+            additional_kwargs["parallel_tool_calls"] = False
+
         response = completion(
             model=model,
             temperature=temperature,
             top_p=top_p,
             messages=messages,
-            functions=functions,
             stream=True,
             **additional_kwargs,
         )
 
         try:
             for chunk in response:
                 delta = chunk.choices[0].delta
+
                 # LiteLLM uses dict instead of Pydantic object like OpenAI does.
-                function_call = (
-                    delta.get("function_call") if use_litellm else delta.function_call
+                tool_calls = (
+                    delta.get("tool_calls") if use_litellm else delta.tool_calls
                 )
-                if function_call:
-                    if function_call.name:
-                        name = function_call.name
-                    if function_call.arguments:
-                        arguments += function_call.arguments
-                if chunk.choices[0].finish_reason == "function_call":
+                if tool_calls:
+                    for tool_call in tool_calls:
+                        if tool_call.function.name:
+                            name = tool_call.function.name
+                        if tool_call.function.arguments:
+                            arguments += tool_call.function.arguments
+                if chunk.choices[0].finish_reason == "tool_calls":
                     yield from self.handle_function_call(messages, name, arguments)
                     yield from self.get_completion(
                         model=model,

tests/test_default.py

@@ -209,7 +209,6 @@ def test_llm_options(completion):
         model=args["--model"],
         temperature=args["--temperature"],
         top_p=args["--top-p"],
-        functions=None,
     )
     completion.assert_called_once_with(**expected_args)
     assert result.exit_code == 0

tests/utils.py

@@ -54,7 +54,6 @@ def comp_args(role, prompt, **kwargs):
         "model": cfg.get("DEFAULT_MODEL"),
         "temperature": 0.0,
         "top_p": 1.0,
-        "functions": None,
         "stream": True,
         **kwargs,
     }