96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
"""
|
|
Example showing how to use Ollama tool calls with parameter extraction.
|
|
Both synchronous and asynchronous examples are provided.
|
|
|
|
To run this example:
|
|
1. Make sure you have `ollama serve` running
|
|
2. Run: python examples/tool_calls/ollama_tool_calls.py
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Dict, Any
|
|
|
|
from neo4j_graphrag.llm import OllamaLLM
|
|
from neo4j_graphrag.llm.types import ToolCallResponse
|
|
from neo4j_graphrag.tool import (
|
|
Tool,
|
|
ObjectParameter,
|
|
StringParameter,
|
|
IntegerParameter,
|
|
)
|
|
|
|
|
|
# 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_calls(response: ToolCallResponse) -> Dict[str, Any]:
|
|
"""Process all tool calls in the response and return the extracted parameters."""
|
|
if not response.tool_calls:
|
|
raise ValueError("No tool calls found in response")
|
|
|
|
print(f"\nNumber of tool calls: {len(response.tool_calls)}")
|
|
print(f"Additional content: {response.content or 'None'}")
|
|
|
|
results = []
|
|
for i, tool_call in enumerate(response.tool_calls):
|
|
print(f"\nTool call #{i + 1}: {tool_call.name}")
|
|
print(f"Arguments: {tool_call.arguments}")
|
|
results.append(tool_call.arguments)
|
|
|
|
# For backward compatibility, return the first tool call's arguments
|
|
return results[0] if results else {}
|
|
|
|
|
|
async def main() -> None:
|
|
async with OllamaLLM(
|
|
model_name="mistral:latest", model_params={"options": {"temperature": 0}}
|
|
) as llm:
|
|
# 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_calls(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_calls(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())
|