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

Author: oskarhane

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())