Add Tool class

To not rely on json schema from openai

Author: oskarhane

examples/tool_calls/openai_tool_calls.py

@@ -17,32 +17,32 @@
 
 from neo4j_graphrag.llm import OpenAILLM
 from neo4j_graphrag.llm.types import ToolCallResponse
+from neo4j_graphrag.tool import Tool, ObjectParameter, StringParameter, IntegerParameter
 
 # Load environment variables from .env file
 load_dotenv()
 
-# Define a tool for extracting information from text
-TOOLS = [
-    {
-        "type": "function",
-        "function": {
-            "name": "extract_person_info",
-            "description": "Extract information about a person from text",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "name": {"type": "string", "description": "The person's full name"},
-                    "age": {"type": "integer", "description": "The person's age"},
-                    "occupation": {
-                        "type": "string",
-                        "description": "The person's occupation",
-                    },
-                },
-                "required": ["name"],
-            },
-        },
-    }
-]
+
+# Create a custom Tool implementation for person info extraction
+parameters = ObjectParameter(
+    description="Parameters for extracting person information",
+    properties={
+        "name": StringParameter(description="The person's full name"),
+        "age": IntegerParameter(description="The person's age"),
+        "occupation": StringParameter(description="The person's occupation"),
+    },
+    required_properties=["name"],
+    additional_properties=False,
+)
+person_info_tool = Tool(
+    name="extract_person_info",
+    description="Extract information about a person from text",
+    parameters=parameters,
+    execute_func=lambda **kwargs: kwargs,
+)
+
+# Create the tool instance
+TOOLS = [person_info_tool]
 
 
 def process_tool_call(response: ToolCallResponse) -> Dict[str, Any]:

src/neo4j_graphrag/llm/base.py

@@ -15,13 +15,15 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, List, Optional, Sequence, Union
 
 from neo4j_graphrag.message_history import MessageHistory
 from neo4j_graphrag.types import LLMMessage
 
 from .types import LLMResponse, ToolCallResponse
 
+from neo4j_graphrag.tool import Tool
+
 
 class LLMInterface(ABC):
     """Interface for large language models.
@@ -88,7 +90,7 @@ async def ainvoke(
     def invoke_with_tools(
         self,
         input: str,
-        tools: List[Dict[str, Any]],
+        tools: Sequence[Tool],
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
         system_instruction: Optional[str] = None,
     ) -> ToolCallResponse:
@@ -98,7 +100,7 @@ def invoke_with_tools(
 
         Args:
             input (str): Text sent to the LLM.
-            tools (List[Dict[str, Any]]): List of tool definitions for the LLM to choose from.
+            tools (Sequence[Tool]): Sequence of Tools for the LLM to choose from. Each LLM implementation should handle the conversion to its specific format.
             message_history (Optional[Union[List[LLMMessage], MessageHistory]]): A collection previous messages,
                 with each message having a specific role assigned.
             system_instruction (Optional[str]): An option to override the llm system message for this invocation.
@@ -115,7 +117,7 @@ def invoke_with_tools(
     async def ainvoke_with_tools(
         self,
         input: str,
-        tools: List[Dict[str, Any]],
+        tools: Sequence[Tool],
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
         system_instruction: Optional[str] = None,
     ) -> ToolCallResponse:
@@ -125,7 +127,7 @@ async def ainvoke_with_tools(
 
         Args:
             input (str): Text sent to the LLM.
-            tools (List[Dict[str, Any]]): List of tool definitions for the LLM to choose from.
+            tools (Sequence[Tool]): Sequence of Tools for the LLM to choose from. Each LLM implementation should handle the conversion to its specific format.
             message_history (Optional[Union[List[LLMMessage], MessageHistory]]): A collection previous messages,
                 with each message having a specific role assigned.
             system_instruction (Optional[str]): An option to override the llm system message for this invocation.

src/neo4j_graphrag/llm/openai_llm.py

@@ -16,7 +16,17 @@
 
 import abc
 import json
-from typing import TYPE_CHECKING, Any, List, Optional, Iterable, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Optional,
+    Iterable,
+    Sequence,
+    Union,
+    cast,
+)
 from openai.types.chat import (
     ChatCompletionMessageParam,
     ChatCompletionToolParam,
@@ -39,6 +49,8 @@
     UserMessage,
 )
 
+from neo4j_graphrag.tool import Tool
+
 if TYPE_CHECKING:
     import openai
 
@@ -91,6 +103,27 @@ def get_messages(
         messages.append(UserMessage(content=input).model_dump())
         return messages  # type: ignore
 
+    def _convert_tool_to_openai_format(self, tool: Tool) -> Dict[str, Any]:
+        """Convert a Tool object to OpenAI's expected format.
+
+        Args:
+            tool: A Tool object to convert to OpenAI's format.
+
+        Returns:
+            A dictionary in OpenAI's tool format.
+        """
+        try:
+            return {
+                "type": "function",
+                "function": {
+                    "name": tool.get_name(),
+                    "description": tool.get_description(),
+                    "parameters": tool.get_parameters(),
+                },
+            }
+        except AttributeError:
+            raise LLMGenerationError(f"Tool {tool} is not a valid Tool object")
+
     def invoke(
         self,
         input: str,
@@ -128,7 +161,7 @@ def invoke(
     def invoke_with_tools(
         self,
         input: str,
-        tools: List[dict[str, Any]],  # Tools definition as a list of dictionaries
+        tools: Sequence[Tool],  # Tools definition as a sequence of Tool objects
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
         system_instruction: Optional[str] = None,
     ) -> ToolCallResponse:
@@ -137,7 +170,7 @@ def invoke_with_tools(
 
         Args:
             input (str): Text sent to the LLM.
-            tools (List[Dict[str, Any]]): List of tool definitions for the LLM to choose from.
+            tools (List[Tool]): List of Tools for the LLM to choose from.
             message_history (Optional[Union[List[LLMMessage], MessageHistory]]): A collection previous messages,
                 with each message having a specific role assigned.
             system_instruction (Optional[str]): An option to override the llm system message for this invocation.
@@ -159,7 +192,8 @@ def invoke_with_tools(
             # Convert tools to OpenAI's expected type
             openai_tools: List[ChatCompletionToolParam] = []
             for tool in tools:
-                openai_tools.append(cast(ChatCompletionToolParam, tool))
+                openai_format_tool = self._convert_tool_to_openai_format(tool)
+                openai_tools.append(cast(ChatCompletionToolParam, openai_format_tool))
 
             response = self.client.chat.completions.create(
                 messages=self.get_messages(input, message_history, system_instruction),
@@ -235,7 +269,7 @@ async def ainvoke(
     async def ainvoke_with_tools(
         self,
         input: str,
-        tools: List[dict[str, Any]],  # Tools definition as a list of dictionaries
+        tools: Sequence[Tool],  # Tools definition as a sequence of Tool objects
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
         system_instruction: Optional[str] = None,
     ) -> ToolCallResponse:
@@ -244,7 +278,7 @@ async def ainvoke_with_tools(
 
         Args:
             input (str): Text sent to the LLM.
-            tools (List[Dict[str, Any]]): List of tool definitions for the LLM to choose from.
+            tools (List[Tool]): List of Tools for the LLM to choose from.
             message_history (Optional[Union[List[LLMMessage], MessageHistory]]): A collection previous messages,
                 with each message having a specific role assigned.
             system_instruction (Optional[str]): An option to override the llm system message for this invocation.
@@ -266,7 +300,8 @@ async def ainvoke_with_tools(
             # Convert tools to OpenAI's expected type
             openai_tools: List[ChatCompletionToolParam] = []
             for tool in tools:
-                openai_tools.append(cast(ChatCompletionToolParam, tool))
+                openai_format_tool = self._convert_tool_to_openai_format(tool)
+                openai_tools.append(cast(ChatCompletionToolParam, openai_format_tool))
 
             response = await self.async_client.chat.completions.create(
                 messages=self.get_messages(input, message_history, system_instruction),