Added tool calls support for LLM clients

giskard/llm/client/base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Sequence
 
 from .logger import LLMLogger
 
@@ -16,6 +16,7 @@ class LLMFunctionCall:
 class LLMOutput:
     message: Optional[str] = None
     function_call: Optional[LLMFunctionCall] = None
+    tool_calls: Sequence[LLMFunctionCall] = field(default_factory=list)
 
 
 class LLMClient(ABC):
@@ -33,5 +34,7 @@ def complete(
         max_tokens=None,
         function_call: Optional[Dict] = None,
         caller_id: Optional[str] = None,
+        tools=None,
+        tool_choice=None,
     ) -> LLMOutput:
         ...

giskard/llm/client/openai.py

@@ -1,8 +1,8 @@
 import json
 from abc import ABC, abstractmethod
-from typing import Dict, Optional, Sequence
 
 from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
+from typing import Dict, Optional, Sequence
 
 from . import LLMClient, LLMFunctionCall, LLMLogger, LLMOutput
 from ..config import LLMConfigurationError
@@ -37,6 +37,8 @@ def _completion(
         function_call: Optional[Dict] = None,
         max_tokens=None,
         caller_id: Optional[str] = None,
+        tools=None,
+        tool_choice=None,
     ) -> dict:
         ...
 
@@ -48,6 +50,8 @@ def complete(
         max_tokens=None,
         function_call: Optional[Dict] = None,
         caller_id: Optional[str] = None,
+        tools=None,
+        tool_choice=None,
     ):
         cc = self._completion(
             messages=messages,
@@ -56,6 +60,8 @@ def complete(
             function_call=function_call,
             max_tokens=max_tokens,
             caller_id=caller_id,
+            tools=tools,
+            tool_choice=tool_choice,
         )
 
         function_call = None
@@ -69,7 +75,15 @@ def complete(
             except (json.JSONDecodeError, KeyError) as err:
                 raise LLMGenerationError("Could not parse function call") from err
 
-        return LLMOutput(message=cc["content"], function_call=function_call)
+        tool_calls = []
+        for tool_call in cc.get("tool_calls") or []:
+            tool_calls.append(
+                LLMFunctionCall(
+                    function=tool_call["function"]["name"], args=json.loads(tool_call["function"]["arguments"])
+                )
+            )
+
+        return LLMOutput(message=cc["content"], function_call=function_call, tool_calls=tool_calls)
 
 
 class LegacyOpenAIClient(BaseOpenAIClient):
@@ -89,12 +103,18 @@ def _completion(
         function_call: Optional[Dict] = None,
         max_tokens=None,
         caller_id: Optional[str] = None,
+        tools=None,
+        tool_choice=None,
     ):
         extra_params = dict()
         if function_call is not None:
             extra_params["function_call"] = function_call
         if functions is not None:
             extra_params["functions"] = functions
+        if tools is not None:
+            extra_params["tools"] = tools
+        if tool_choice is not None:
+            extra_params["tool_choice"] = tool_choice
 
         try:
             completion = openai.ChatCompletion.create(
@@ -134,12 +154,18 @@ def _completion(
         function_call: Optional[Dict] = None,
         max_tokens=None,
         caller_id: Optional[str] = None,
+        tools=None,
+        tool_choice=None,
     ):
         extra_params = dict()
         if function_call is not None:
             extra_params["function_call"] = function_call
         if functions is not None:
             extra_params["functions"] = functions
+        if tools is not None:
+            extra_params["tools"] = tools
+        if tool_choice is not None:
+            extra_params["tool_choice"] = tool_choice
 
         try:
             completion = self._client.chat.completions.create(

giskard/llm/evaluators/base.py

@@ -9,22 +9,25 @@
 
 EVALUATE_MODEL_FUNCTIONS = [
     {
-        "name": "evaluate_model",
-        "description": "Evaluates if the model passes the test",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "passed_test": {
-                    "type": "boolean",
-                    "description": "true if the model successfully passes the test",
-                },
-                "reason": {
-                    "type": "string",
-                    "description": "optional short description of why the model does not pass the test, in 1 or 2 short sentences",
+        "type": "function",
+        "function": {
+            "name": "evaluate_model",
+            "description": "Evaluates if the model passes the test",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "passed_test": {
+                        "type": "boolean",
+                        "description": "true if the model successfully passes the test",
+                    },
+                    "reason": {
+                        "type": "string",
+                        "description": "optional short description of why the model does not pass the test, in 1 or 2 short sentences",
+                    },
                 },
+                "required": ["passed_test"],
             },
         },
-        "required": ["passed_test"],
     },
 ]
 
@@ -90,18 +93,18 @@ def evaluate(self, model: BaseModel, dataset: Dataset):
             try:
                 out = self.llm_client.complete(
                     [{"role": "system", "content": prompt}],
-                    functions=funcs,
-                    function_call={"name": "evaluate_model"},
+                    tools=funcs,
+                    tool_choice={"type": "function", "function": {"name": "evaluate_model"}},
                     temperature=self.llm_temperature,
                     caller_id=self.__class__.__name__,
                 )
-                if out.function_call is None or "passed_test" not in out.function_call.args:
+                if len(out.tool_calls) != 1 or "passed_test" not in out.tool_calls[0].args:
                     raise LLMGenerationError("Invalid function call arguments received")
             except LLMGenerationError as err:
                 errored.append({"message": str(err), "sample": sample})
                 continue
 
-            args = out.function_call.args
+            args = out.tool_calls[0].args
             if args["passed_test"]:
                 succeeded.append({"input_vars": input_vars, "model_output": model_output, "reason": args.get("reason")})
             else:

giskard/llm/evaluators/coherency.py

@@ -98,13 +98,13 @@ def _eval_pair(self, model: BaseModel, input_1, input_2, output_1, output_2):
 
         out = self.llm_client.complete(
             [{"role": "system", "content": prompt}],
-            functions=EVALUATE_MODEL_FUNCTIONS,
-            function_call={"name": "evaluate_model"},  # force function call
+            tools=EVALUATE_MODEL_FUNCTIONS,
+            tool_choice={"type": "function", "function": {"name": "evaluate_model"}},  # force tool call
             temperature=self.llm_temperature,
             caller_id=self.__class__.__name__,
         )
 
-        if out.function_call is None or "passed_test" not in out.function_call.args:
+        if len(out.tool_calls) != 1 or "passed_test" not in out.tool_calls[0].args:
             raise LLMGenerationError("Invalid function call arguments received")
 
-        return out.function_call.args["passed_test"], out.function_call.args.get("reason")
+        return out.tool_calls[0].args["passed_test"], out.tool_calls[0].args.get("reason")

giskard/llm/generators/base.py

@@ -1,13 +1,12 @@
-from typing import Optional, Sequence
-
 from abc import ABC, abstractmethod
 
 import pandas as pd
 from pydantic import BaseModel
+from typing import Optional, Sequence
 
-from ...datasets.base import Dataset
 from ..client import LLMClient, get_default_client
 from ..errors import LLMGenerationError
+from ...datasets.base import Dataset
 
 DEFAULT_GENERATE_INPUTS_PROMPT = """You are auditing an AI model. Your task is to generate typical but varied inputs for this model.
 
@@ -61,20 +60,23 @@ def _make_generate_input_prompt(self, model: BaseModel, num_samples: int):
     def _make_generate_input_functions(self, model: BaseModel, num_samples: int):
         return [
             {
-                "name": "generate_inputs",
-                "description": "generates inputs for model audit",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "inputs": {
-                            "type": "array",
-                            "items": {
-                                "type": "object",
-                                "properties": {name: {"type": "string"} for name in model.meta.feature_names},
-                            },
-                        }
+                "type": "function",
+                "function": {
+                    "name": "generate_inputs",
+                    "description": "generates inputs for model audit",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "inputs": {
+                                "type": "array",
+                                "items": {
+                                    "type": "object",
+                                    "properties": {name: {"type": "string"} for name in model.meta.feature_names},
+                                },
+                            }
+                        },
+                        "required": ["inputs"],
                     },
-                    "required": ["inputs"],
                 },
             }
         ]
@@ -107,19 +109,19 @@ def generate_dataset(self, model: BaseModel, num_samples: int = 10, column_types
 
         """
         prompt = self._make_generate_input_prompt(model, num_samples)
-        functions = self._make_generate_input_functions(model, num_samples)
+        tools = self._make_generate_input_functions(model, num_samples)
 
         out = self.llm_client.complete(
             messages=[{"role": "system", "content": prompt}],
-            functions=functions,
-            function_call={"name": "generate_inputs"},
+            tools=tools,
+            tool_choice={"type": "function", "function": {"name": "generate_inputs"}},
             temperature=self.llm_temperature,
             caller_id=self.__class__.__name__,
         )
 
         try:
-            generated = out.function_call.args["inputs"]
-        except (AttributeError, KeyError) as err:
+            generated = out.tool_calls[0].args["inputs"]
+        except (AttributeError, KeyError, IndexError) as err:
             raise LLMGenerationError("Could not parse generated inputs") from err
 
         dataset = Dataset(

giskard/llm/testcase.py

@@ -22,26 +22,32 @@
 
 GENERATE_REQUIREMENTS_FUNCTIONS = [
     {
-        "name": "generate_requirements",
-        "description": "Generates requirements for model testing",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "requirements": {
-                    "type": "array",
-                    "items": {"type": "string", "description": "A requirement the model must satisfy"},
-                }
+        "type": "function",
+        "function": {
+            "name": "generate_requirements",
+            "description": "Generates requirements for model testing",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "requirements": {
+                        "type": "array",
+                        "items": {"type": "string", "description": "A requirement the model must satisfy"},
+                    }
+                },
+                "required": ["requirements"],
             },
-            "required": ["requirements"],
         },
     },
     {
-        "name": "skip",
-        "description": "Skips the generation when no relevant requirements can be generated",
-        "parameters": {
-            "type": "object",
-            "properties": {},
-            "required": [],
+        "type": "function",
+        "function": {
+            "name": "skip",
+            "description": "Skips the generation when no relevant requirements can be generated",
+            "parameters": {
+                "type": "object",
+                "properties": {},
+                "required": [],
+            },
         },
     },
 ]
@@ -70,8 +76,8 @@ def generate_requirements(self, model: BaseModel, max_requirements: int = 5):
         functions = self._make_generate_requirements_functions()
         out = self.llm_client.complete(
             messages=[{"role": "system", "content": prompt}],
-            functions=functions,
-            function_call={"name": "generate_requirements"},
+            tools=functions,
+            tool_choice={"type": "function", "function": {"name": "generate_requirements"}},
             temperature=self.llm_temperature,
             caller_id=self.__class__.__name__,
         )