Add tool calling to the LLM base class, implement in OpenAI

Author: oskarhane

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Added
 
+- Added tool calling functionality to the LLM base class with OpenAI implementation, enabling structured parameter extraction and function calling.
 - Added the `run_with_context` method to `Component`. This method includes a `context_` parameter, which provides information about the pipeline from which the component is executed (e.g., the `run_id`). It also enables the component to send events to the pipeline's callback function.
 
 

examples/tool_calls/openai_tool_calls.py

@@ -0,0 +1,95 @@
+"""
+Example showing how to use OpenAI tool calls with parameter extraction.
+Both synchronous and asynchronous examples are provided.
+
+To run this example:
+1. Make sure you have the OpenAI API key in your .env file:
+   OPENAI_API_KEY=your-api-key
+2. Run: python examples/tool_calls/openai_tool_calls.py
+"""
+
+import asyncio
+import json
+import os
+from typing import Dict, Any
+
+from dotenv import load_dotenv
+
+from neo4j_graphrag.llm import OpenAILLM
+from neo4j_graphrag.llm.types import ToolCallResponse
+
+# 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"],
+            },
+        },
+    }
+]
+
+
+def process_tool_call(response: ToolCallResponse) -> Dict[str, Any]:
+    """Process the tool call response and return the extracted parameters."""
+    if not response.tool_calls:
+        raise ValueError("No tool calls found in response")
+
+    tool_call = response.tool_calls[0]
+    print(f"\nTool called: {tool_call.name}")
+    print(f"Arguments: {tool_call.arguments}")
+    print(f"Additional content: {response.content or 'None'}")
+    return tool_call.arguments
+
+
+async def main() -> None:
+    # Initialize the OpenAI LLM
+    llm = OpenAILLM(
+        api_key=os.getenv("OPENAI_API_KEY"),
+        model_name="gpt-4o",
+        model_params={"temperature": 0},
+    )
+
+    # Example text containing information about a person
+    text = "Stella Hane is a 35-year-old software engineer who loves coding."
+
+    print("\n=== Synchronous Tool Call ===")
+    # Make a synchronous tool call
+    sync_response = llm.invoke_with_tools(
+        input=f"Extract information about the person from this text: {text}",
+        tools=TOOLS,
+    )
+    sync_result = process_tool_call(sync_response)
+    print("\n=== Synchronous Tool Call Result ===")
+    print(json.dumps(sync_result, indent=2))
+
+    print("\n=== Asynchronous Tool Call ===")
+    # Make an asynchronous tool call with a different text
+    text2 = "Molly Hane, 32, works as a data scientist and enjoys machine learning."
+    async_response = await llm.ainvoke_with_tools(
+        input=f"Extract information about the person from this text: {text2}",
+        tools=TOOLS,
+    )
+    async_result = process_tool_call(async_response)
+    print("\n=== Asynchronous Tool Call Result ===")
+    print(json.dumps(async_result, indent=2))
+
+
+if __name__ == "__main__":
+    # Run the async main function
+    asyncio.run(main())

src/neo4j_graphrag/llm/base.py

@@ -15,12 +15,12 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from neo4j_graphrag.message_history import MessageHistory
 from neo4j_graphrag.types import LLMMessage
 
-from .types import LLMResponse
+from .types import LLMResponse, ToolCallResponse
 
 
 class LLMInterface(ABC):
@@ -84,3 +84,57 @@ async def ainvoke(
         Raises:
             LLMGenerationError: If anything goes wrong.
         """
+
+    def invoke_with_tools(
+        self,
+        input: str,
+        tools: List[Dict[str, Any]],
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> ToolCallResponse:
+        """Sends a text input to the LLM with tool definitions and retrieves a tool call response.
+
+        This is a default implementation that should be overridden by LLM providers that support tool/function calling.
+
+        Args:
+            input (str): Text sent to the LLM.
+            tools (List[Dict[str, Any]]): List of tool definitions 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.
+
+        Returns:
+            ToolCallResponse: The response from the LLM containing a tool call.
+
+        Raises:
+            LLMGenerationError: If anything goes wrong.
+            NotImplementedError: If the LLM provider does not support tool calling.
+        """
+        raise NotImplementedError("This LLM provider does not support tool calling.")
+
+    async def ainvoke_with_tools(
+        self,
+        input: str,
+        tools: List[Dict[str, Any]],
+        message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
+        system_instruction: Optional[str] = None,
+    ) -> ToolCallResponse:
+        """Asynchronously sends a text input to the LLM with tool definitions and retrieves a tool call response.
+
+        This is a default implementation that should be overridden by LLM providers that support tool/function calling.
+
+        Args:
+            input (str): Text sent to the LLM.
+            tools (List[Dict[str, Any]]): List of tool definitions 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.
+
+        Returns:
+            ToolCallResponse: The response from the LLM containing a tool call.
+
+        Raises:
+            LLMGenerationError: If anything goes wrong.
+            NotImplementedError: If the LLM provider does not support tool calling.
+        """
+        raise NotImplementedError("This LLM provider does not support tool calling.")