336 lines
12 KiB
Python
336 lines
12 KiB
Python
|
# Copyright (c) "Neo4j"
|
||
|
|
# Neo4j Sweden AB [https://neo4j.com]
|
||
|
|
# #
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
# #
|
||
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
# #
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
"""
|
||
|
|
Example demonstrating how to use the ToolsRetriever.
|
||
|
|
|
||
|
|
This example shows:
|
||
|
|
1. How to create tools from different retrievers
|
||
|
|
2. How to use the ToolsRetriever to select and execute tools based on a query
|
||
|
|
"""
|
||
|
|
|
||
|
|
import os
|
||
|
|
from typing import Any, Optional, cast
|
||
|
|
from unittest.mock import MagicMock
|
||
|
|
from dotenv import load_dotenv
|
||
|
|
import requests
|
||
|
|
from datetime import datetime, date
|
||
|
|
|
||
|
|
import neo4j
|
||
|
|
|
||
|
|
from neo4j_graphrag.generation import GraphRAG
|
||
|
|
from neo4j_graphrag.retrievers.base import Retriever
|
||
|
|
from neo4j_graphrag.retrievers.tools_retriever import ToolsRetriever
|
||
|
|
from neo4j_graphrag.types import RawSearchResult
|
||
|
|
from neo4j_graphrag.tool import (
|
||
|
|
ObjectParameter,
|
||
|
|
StringParameter,
|
||
|
|
Tool,
|
||
|
|
)
|
||
|
|
from neo4j_graphrag.llm.openai_llm import OpenAILLM
|
||
|
|
|
||
|
|
# Load environment variables from .env file (OPENAI_API_KEY required for this example)
|
||
|
|
load_dotenv()
|
||
|
|
|
||
|
|
|
||
|
|
# Create a Retriever that returns static results about Neo4j
|
||
|
|
class Neo4jInfoRetriever(Retriever):
|
||
|
|
"""A retriever that returns general information about Neo4j."""
|
||
|
|
|
||
|
|
# Disable Neo4j version verification
|
||
|
|
VERIFY_NEO4J_VERSION = False
|
||
|
|
|
||
|
|
def __init__(self, driver: neo4j.Driver):
|
||
|
|
# Call the parent class constructor with the driver
|
||
|
|
super().__init__(driver)
|
||
|
|
|
||
|
|
def get_search_results(
|
||
|
|
self, query_text: Optional[str] = None, **kwargs: Any
|
||
|
|
) -> RawSearchResult:
|
||
|
|
"""Return general information about Neo4j."""
|
||
|
|
# Create formatted Neo4j information
|
||
|
|
neo4j_info = (
|
||
|
|
"# Neo4j Graph Database\n\n"
|
||
|
|
"Neo4j is a graph database management system developed by Neo4j, Inc. "
|
||
|
|
"It is an ACID-compliant transactional database with native graph storage and processing.\n\n"
|
||
|
|
"## Key Features:\n\n"
|
||
|
|
"- **Cypher Query Language**: Neo4j's intuitive query language designed specifically for working with graph data\n"
|
||
|
|
"- **Property Graphs**: Both nodes and relationships can have properties (key-value pairs)\n"
|
||
|
|
"- **ACID Compliance**: Ensures data integrity with full transaction support\n"
|
||
|
|
"- **Native Graph Storage**: Optimized storage for graph data structures\n"
|
||
|
|
"- **High Availability**: Clustering for enterprise deployments\n"
|
||
|
|
"- **Scalability**: Handles billions of nodes and relationships"
|
||
|
|
)
|
||
|
|
|
||
|
|
# Create a Neo4j record with the information
|
||
|
|
records = [neo4j.Record({"result": neo4j_info})]
|
||
|
|
|
||
|
|
# Return a RawSearchResult with the records and metadata
|
||
|
|
return RawSearchResult(records=records, metadata={"query": query_text})
|
||
|
|
|
||
|
|
|
||
|
|
class CalendarTool(Tool):
|
||
|
|
"""A simple tool to get calendar information."""
|
||
|
|
|
||
|
|
def __init__(self) -> None:
|
||
|
|
"""Initialize the calendar tool."""
|
||
|
|
# Define parameters for the tool
|
||
|
|
parameters = ObjectParameter(
|
||
|
|
description="Parameters for calendar information retrieval",
|
||
|
|
properties={
|
||
|
|
"date": StringParameter(
|
||
|
|
description="The date to check events for in YYYY-MM-DD format (e.g., 2025-04-14)",
|
||
|
|
),
|
||
|
|
},
|
||
|
|
required_properties=["date"],
|
||
|
|
)
|
||
|
|
|
||
|
|
# Sample calendar data with fixed dates
|
||
|
|
self.calendar_data = {
|
||
|
|
"2025-04-15": [
|
||
|
|
{"time": "10:00", "title": "Dentist Appointment"},
|
||
|
|
{"time": "14:00", "title": "Conference Call"},
|
||
|
|
],
|
||
|
|
"2025-04-16": [],
|
||
|
|
}
|
||
|
|
|
||
|
|
# Define a wrapper function that handles parameters correctly
|
||
|
|
def execute_func(**kwargs: Any) -> str:
|
||
|
|
return self.execute_calendar(**kwargs)
|
||
|
|
|
||
|
|
super().__init__(
|
||
|
|
name="calendar_tool",
|
||
|
|
description="Check calendar events for a specific date in YYYY-MM-DD format",
|
||
|
|
parameters=parameters,
|
||
|
|
execute_func=execute_func,
|
||
|
|
)
|
||
|
|
|
||
|
|
def execute_calendar(self, **kwargs: Any) -> str:
|
||
|
|
"""Execute the calendar tool.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
**kwargs: Dictionary of parameters, including 'date'.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
str: The events for the specified date.
|
||
|
|
"""
|
||
|
|
date = kwargs.get("date")
|
||
|
|
if not date:
|
||
|
|
return "Error: No date provided"
|
||
|
|
|
||
|
|
# Check for events on the date
|
||
|
|
if date in self.calendar_data:
|
||
|
|
events_list = self.calendar_data[date]
|
||
|
|
if not events_list:
|
||
|
|
return f"No events scheduled for {date}"
|
||
|
|
|
||
|
|
events_str = "\n".join(
|
||
|
|
f"- {event.get('time', 'All day')}: {event.get('title', 'Untitled event')}"
|
||
|
|
for event in events_list
|
||
|
|
)
|
||
|
|
return f"Events for {date}:\n{events_str}"
|
||
|
|
else:
|
||
|
|
return f"No events found for {date}"
|
||
|
|
|
||
|
|
|
||
|
|
class WeatherTool(Tool):
|
||
|
|
"""A tool to fetch weather in Malmö, Sweden for a specific date."""
|
||
|
|
|
||
|
|
def __init__(self) -> None:
|
||
|
|
"""Initialize the weather tool."""
|
||
|
|
parameters = ObjectParameter(
|
||
|
|
description="Parameters for fetching weather information about a date.",
|
||
|
|
properties={
|
||
|
|
"date": StringParameter(
|
||
|
|
description='The date, in YYYY-MM-DD format. Example: "2025-04-25"'
|
||
|
|
)
|
||
|
|
},
|
||
|
|
required_properties=["date"],
|
||
|
|
)
|
||
|
|
super().__init__(
|
||
|
|
name="weather_tool",
|
||
|
|
description="Check for weather for a specific date in YYYY-MM-DD format",
|
||
|
|
parameters=parameters,
|
||
|
|
execute_func=self.execute_weather_retrieval,
|
||
|
|
)
|
||
|
|
|
||
|
|
def execute_weather_retrieval(self, **kwargs: Any) -> str:
|
||
|
|
"""Fetch historical weather data for a given date in Malmö, Sweden."""
|
||
|
|
date_str = kwargs.get("date")
|
||
|
|
if not date_str:
|
||
|
|
return "Error: Date not provided for weather lookup."
|
||
|
|
|
||
|
|
try:
|
||
|
|
input_date = datetime.strptime(date_str, "%Y-%m-%d").date()
|
||
|
|
except ValueError:
|
||
|
|
return f"Error: Invalid date format '{date_str}'. Please use YYYY-MM-DD."
|
||
|
|
|
||
|
|
today_date = date.today()
|
||
|
|
|
||
|
|
if input_date < today_date:
|
||
|
|
api_url = "https://archive-api.open-meteo.com/v1/archive"
|
||
|
|
else:
|
||
|
|
# For today or future dates, use the forecast API
|
||
|
|
# Note: Forecast API typically has a limit (e.g., 16 days into the future)
|
||
|
|
api_url = "https://api.open-meteo.com/v1/forecast"
|
||
|
|
|
||
|
|
params = {
|
||
|
|
"latitude": 55.6059, # Malmö, Sweden
|
||
|
|
"longitude": 13.0007, # Malmö, Sweden
|
||
|
|
"start_date": date_str,
|
||
|
|
"end_date": date_str,
|
||
|
|
"daily": "temperature_2m_max,sunshine_duration",
|
||
|
|
"timezone": "Europe/Stockholm",
|
||
|
|
}
|
||
|
|
headers = {"Accept": "application/json"}
|
||
|
|
|
||
|
|
try:
|
||
|
|
response = requests.get(api_url, headers=headers, params=params)
|
||
|
|
response.raise_for_status()
|
||
|
|
data = response.json()
|
||
|
|
|
||
|
|
# Try to access keys directly, relying on the existing broader except block for errors
|
||
|
|
daily = data["daily"]
|
||
|
|
temp_max = daily["temperature_2m_max"][0]
|
||
|
|
sunshine_seconds = daily["sunshine_duration"][0]
|
||
|
|
|
||
|
|
sunshine_hours = 0
|
||
|
|
if (
|
||
|
|
sunshine_seconds is not None
|
||
|
|
): # API might return null for sunshine_duration
|
||
|
|
sunshine_hours = round(sunshine_seconds / 3600, 1)
|
||
|
|
|
||
|
|
return (
|
||
|
|
f"Weather for Malmö, Sweden on this day:\n"
|
||
|
|
f"- Max Temperature: {temp_max}°C\n"
|
||
|
|
f"- Sunshine Duration: {sunshine_hours} hours"
|
||
|
|
)
|
||
|
|
except requests.exceptions.RequestException as e:
|
||
|
|
return f"API request failed for weather data: {e}"
|
||
|
|
except (
|
||
|
|
ValueError,
|
||
|
|
KeyError,
|
||
|
|
) as e:
|
||
|
|
return f"Error parsing weather data for Malmö on {date_str}: {e}"
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
"""Run the example."""
|
||
|
|
# Create a mock Neo4j driver
|
||
|
|
driver = cast(neo4j.Driver, MagicMock())
|
||
|
|
|
||
|
|
# Create retrievers
|
||
|
|
neo4j_retriever = Neo4jInfoRetriever(driver=driver)
|
||
|
|
|
||
|
|
# Convert retrievers to tools
|
||
|
|
neo4j_tool = neo4j_retriever.convert_to_tool(
|
||
|
|
name="neo4j_info_tool",
|
||
|
|
description="Get information about Neo4j graph database",
|
||
|
|
parameter_descriptions={
|
||
|
|
"query_text": "The query about Neo4j",
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
# Create a calendar tool
|
||
|
|
calendar_tool = CalendarTool()
|
||
|
|
|
||
|
|
# Create a weather tool
|
||
|
|
weather_tool = WeatherTool()
|
||
|
|
|
||
|
|
# Create an OpenAI LLM
|
||
|
|
llm = OpenAILLM(
|
||
|
|
api_key=os.getenv("OPENAI_API_KEY"),
|
||
|
|
model_name="gpt-5",
|
||
|
|
model_params={
|
||
|
|
"temperature": 0.2,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
# Print tool information for debugging
|
||
|
|
print("\nTool Information:")
|
||
|
|
print(f"Neo4j Tool: {neo4j_tool.get_name()}, {neo4j_tool.get_description()}")
|
||
|
|
print(
|
||
|
|
f"Calendar Tool: {calendar_tool.get_name()}, {calendar_tool.get_description()}"
|
||
|
|
)
|
||
|
|
parameters_description = (
|
||
|
|
weather_tool._parameters.description
|
||
|
|
if weather_tool._parameters
|
||
|
|
else "No parameters description"
|
||
|
|
)
|
||
|
|
print(
|
||
|
|
f"Weather Tool: {weather_tool.get_name()}, {weather_tool.get_description()}: {parameters_description}"
|
||
|
|
)
|
||
|
|
|
||
|
|
# Create a ToolsRetriever with the LLM and tools
|
||
|
|
tools_retriever = ToolsRetriever(
|
||
|
|
driver=driver,
|
||
|
|
llm=llm,
|
||
|
|
tools=[neo4j_tool, calendar_tool, weather_tool],
|
||
|
|
)
|
||
|
|
|
||
|
|
# Test queries
|
||
|
|
test_queries = [
|
||
|
|
"What is Neo4j?",
|
||
|
|
"Do I have any meetings the 15th of April 2025?",
|
||
|
|
"Any information about 2025-04-16?",
|
||
|
|
]
|
||
|
|
|
||
|
|
# Run just the tools retriever directly to show metadata etc.
|
||
|
|
print(f"\n\n{'=' * 80}")
|
||
|
|
print("Retriever call examples, to show metadata etc.")
|
||
|
|
print(f"{'=' * 80}")
|
||
|
|
for query in test_queries:
|
||
|
|
print(f"Query: {query}")
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Get search results through the ToolsRetriever
|
||
|
|
result = tools_retriever.get_search_results(query_text=query)
|
||
|
|
|
||
|
|
# Print metadata
|
||
|
|
if result.metadata is not None:
|
||
|
|
print(f"\nTools selected: {result.metadata.get('tools_selected', [])}")
|
||
|
|
if result.metadata.get("error", ""):
|
||
|
|
print(f"Error: {result.metadata.get('error', '')}")
|
||
|
|
|
||
|
|
# Print results
|
||
|
|
print("\nRESULTS:")
|
||
|
|
for i, record in enumerate(result.records):
|
||
|
|
print(f"\n--- Result {i + 1} ---")
|
||
|
|
print(f"Content: {record.get('content', 'N/A')}")
|
||
|
|
print(f"Tool: {record.get('tool_name', 'Unknown')}")
|
||
|
|
if record.get("metadata"):
|
||
|
|
print(f"Metadata: {record.get('metadata')}")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error: {str(e)}")
|
||
|
|
print(f"{'=' * 80}")
|
||
|
|
|
||
|
|
# For demo purposes, run the queries through GraphRAG to get text input -> text output
|
||
|
|
print(f"\n\n{'=' * 80}")
|
||
|
|
print("Full GraphRAG examples")
|
||
|
|
print(f"{'=' * 80}")
|
||
|
|
for query in test_queries:
|
||
|
|
print(f"Query: {query}")
|
||
|
|
# Full GraphRAG example
|
||
|
|
graphrag = GraphRAG(
|
||
|
|
llm=llm,
|
||
|
|
retriever=tools_retriever,
|
||
|
|
)
|
||
|
|
rag_result = graphrag.search(query_text=query, return_context=False)
|
||
|
|
print(f"Answer: {rag_result.answer}")
|
||
|
|
print(f"{'=' * 80}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|