참고소스 수정본
This commit is contained in:
LASTA_DEV01\lasta
@@ -0,0 +1,174 @@
|
||||
# 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 with a Cypher template.
|
||||
|
||||
This example shows:
|
||||
1. How to create a tool from a Cypher template
|
||||
2. How to use the ToolsRetriever to select and execute tools based on a query
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
|
||||
import neo4j
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from neo4j_graphrag.llm.openai_llm import OpenAILLM
|
||||
from neo4j_graphrag.retrievers.base import Retriever
|
||||
from neo4j_graphrag.retrievers.tools_retriever import ToolsRetriever
|
||||
from neo4j_graphrag.types import RawSearchResult, RetrieverResultItem
|
||||
|
||||
# Load environment variables from .env file (OPENAI_API_KEY required for this example)
|
||||
load_dotenv()
|
||||
|
||||
# Define database credentials
|
||||
URI = "neo4j+s://demo.neo4jlabs.com"
|
||||
AUTH = ("recommendations", "recommendations")
|
||||
|
||||
# Cypher template to count actors in a specific movie
|
||||
CYPHER_TEMPLATE = """
|
||||
MATCH (m:Movie {title: $title})
|
||||
OPTIONAL MATCH (m)<-[:ACTED_IN]-(a:Actor)
|
||||
WITH m, collect(a.name) AS actor_names, count(a) AS actor_count
|
||||
RETURN m.title AS movie_title,
|
||||
m.plot AS plot,
|
||||
m.released AS year,
|
||||
actor_count,
|
||||
actor_names
|
||||
"""
|
||||
|
||||
|
||||
class CypherTemplateRetriever(Retriever):
|
||||
"""
|
||||
Custom retriever that executes a parameterized Cypher query template.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
driver: neo4j.Driver,
|
||||
cypher_template: str,
|
||||
neo4j_database: Optional[str] = None,
|
||||
result_formatter: Optional[
|
||||
Callable[[neo4j.Record], RetrieverResultItem]
|
||||
] = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
driver: Neo4j driver instance
|
||||
cypher_template: Cypher query with parameters (e.g., "MATCH (m:Movie {title: $title}) RETURN m")
|
||||
neo4j_database: Optional database name
|
||||
result_formatter: Optional function to format results
|
||||
"""
|
||||
super().__init__(driver, neo4j_database)
|
||||
self.cypher_template = cypher_template
|
||||
self.result_formatter = result_formatter
|
||||
|
||||
def get_search_results(
|
||||
self,
|
||||
query_vector: Optional[list[float]] = None,
|
||||
query_text: Optional[str] = None,
|
||||
top_k: int = 5,
|
||||
query_params: Optional[Dict[str, Any]] = None,
|
||||
**kwargs: Any,
|
||||
) -> RawSearchResult:
|
||||
"""
|
||||
Execute the Cypher template with provided parameters.
|
||||
|
||||
Args:
|
||||
query_text: Can be used as a parameter (e.g., for movie title)
|
||||
query_params: Dictionary of parameters for the Cypher query
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
RawSearchResult containing neo4j.Record objects
|
||||
"""
|
||||
# Prepare parameters for the Cypher query
|
||||
parameters = query_params or {}
|
||||
|
||||
# Optionally use query_text as a parameter
|
||||
if query_text and "query_text" not in parameters:
|
||||
parameters["query_text"] = query_text
|
||||
|
||||
# Add any additional kwargs as parameters
|
||||
parameters.update(kwargs)
|
||||
|
||||
# Execute the query
|
||||
try:
|
||||
records, summary, keys = self.driver.execute_query(
|
||||
self.cypher_template,
|
||||
parameters_=parameters,
|
||||
database_=self.neo4j_database,
|
||||
routing_=neo4j.RoutingControl.READ,
|
||||
)
|
||||
|
||||
return RawSearchResult(
|
||||
records=records,
|
||||
metadata={
|
||||
"cypher_query": self.cypher_template,
|
||||
"parameters": parameters,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to execute Cypher template: {e}") from e
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Run the example."""
|
||||
driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
|
||||
try:
|
||||
neo4j_retriever = CypherTemplateRetriever(
|
||||
driver=driver,
|
||||
cypher_template=CYPHER_TEMPLATE,
|
||||
)
|
||||
|
||||
# Convert retriever to tool
|
||||
neo4j_tool = neo4j_retriever.convert_to_tool(
|
||||
name="movie_info_tool",
|
||||
description=(
|
||||
"Retrieves the total number of actors in a specific movie and returns "
|
||||
"aggregation metrics including actor count, list of actor names, and movie details. "
|
||||
"Use this when the user asks about cast size, number of actors, or movie cast information."
|
||||
),
|
||||
parameter_descriptions={
|
||||
"query_params": (
|
||||
"Dictionary containing 'title' key with the movie title to analyze. "
|
||||
"Example: {'title': 'The Matrix'}"
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
llm = OpenAILLM(
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
model_name="gpt-5",
|
||||
model_params={"temperature": 0.2},
|
||||
)
|
||||
|
||||
tools_retriever = ToolsRetriever(
|
||||
driver=driver,
|
||||
llm=llm,
|
||||
tools=[neo4j_tool],
|
||||
)
|
||||
|
||||
query_text = "How many actors are there in Around the World in 80 Days?"
|
||||
print("Query:", query_text)
|
||||
result = tools_retriever.search(query_text=query_text, return_context=True)
|
||||
print("Result:", result)
|
||||
finally:
|
||||
driver.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,151 @@
|
||||
# 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 create multiple domain-specific tools from retrievers.
|
||||
|
||||
This example shows:
|
||||
1. How to create multiple tools from the same retriever type for different use cases
|
||||
2. How to provide custom parameter descriptions for each tool
|
||||
3. How type inference works automatically while descriptions are explicit
|
||||
"""
|
||||
|
||||
import neo4j
|
||||
from typing import cast, Any, Optional
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from neo4j_graphrag.retrievers.base import Retriever
|
||||
from neo4j_graphrag.types import RawSearchResult
|
||||
|
||||
|
||||
class MockVectorRetriever(Retriever):
|
||||
"""A mock vector retriever for demonstration purposes."""
|
||||
|
||||
VERIFY_NEO4J_VERSION = False
|
||||
|
||||
def __init__(self, driver: neo4j.Driver, index_name: str):
|
||||
super().__init__(driver)
|
||||
self.index_name = index_name
|
||||
|
||||
def get_search_results(
|
||||
self,
|
||||
query_vector: Optional[list[float]] = None,
|
||||
query_text: Optional[str] = None,
|
||||
top_k: int = 5,
|
||||
effective_search_ratio: int = 1,
|
||||
filters: Optional[dict[str, Any]] = None,
|
||||
) -> RawSearchResult:
|
||||
"""Get vector search results (mocked for demonstration)."""
|
||||
# Return empty results for demo
|
||||
return RawSearchResult(records=[], metadata={"index": self.index_name})
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Demonstrate creating multiple domain-specific tools from retrievers."""
|
||||
|
||||
# Create mock driver (in real usage, this would be actual Neo4j driver)
|
||||
driver = cast(Any, MagicMock())
|
||||
|
||||
# Create retrievers for different domains using the same retriever type
|
||||
# In practice, these would point to different vector indexes
|
||||
|
||||
# Movie recommendations retriever
|
||||
movie_retriever = MockVectorRetriever(driver=driver, index_name="movie_embeddings")
|
||||
|
||||
# Product search retriever
|
||||
product_retriever = MockVectorRetriever(
|
||||
driver=driver, index_name="product_embeddings"
|
||||
)
|
||||
|
||||
# Document search retriever
|
||||
document_retriever = MockVectorRetriever(
|
||||
driver=driver, index_name="document_embeddings"
|
||||
)
|
||||
|
||||
# Convert each retriever to a domain-specific tool with custom descriptions
|
||||
|
||||
# 1. Movie recommendation tool
|
||||
movie_tool = movie_retriever.convert_to_tool(
|
||||
name="movie_search",
|
||||
description="Find movie recommendations based on plot, genre, or actor preferences",
|
||||
parameter_descriptions={
|
||||
"query_text": "Movie title, plot description, genre, or actor name",
|
||||
"query_vector": "Pre-computed embedding vector for movie search",
|
||||
"top_k": "Number of movie recommendations to return (1-20)",
|
||||
"filters": "Optional filters for genre, year, rating, etc.",
|
||||
"effective_search_ratio": "Search pool multiplier for better accuracy",
|
||||
},
|
||||
)
|
||||
|
||||
# 2. Product search tool
|
||||
product_tool = product_retriever.convert_to_tool(
|
||||
name="product_search",
|
||||
description="Search for products matching customer needs and preferences",
|
||||
parameter_descriptions={
|
||||
"query_text": "Product name, description, or customer need",
|
||||
"query_vector": "Pre-computed embedding for product matching",
|
||||
"top_k": "Maximum number of product results (1-50)",
|
||||
"filters": "Filters for price range, brand, category, availability",
|
||||
"effective_search_ratio": "Breadth vs precision trade-off for search",
|
||||
},
|
||||
)
|
||||
|
||||
# 3. Document search tool
|
||||
document_tool = document_retriever.convert_to_tool(
|
||||
name="document_search",
|
||||
description="Find relevant documents and knowledge articles",
|
||||
parameter_descriptions={
|
||||
"query_text": "Question, keywords, or topic to search for",
|
||||
"query_vector": "Semantic embedding for document retrieval",
|
||||
"top_k": "Number of relevant documents to retrieve (1-10)",
|
||||
"filters": "Document type, date range, or department filters",
|
||||
},
|
||||
)
|
||||
|
||||
# Demonstrate that each tool has distinct, meaningful descriptions
|
||||
tools = [movie_tool, product_tool, document_tool]
|
||||
|
||||
for tool in tools:
|
||||
print(f"\n=== {tool.get_name().upper()} ===")
|
||||
print(f"Description: {tool.get_description()}")
|
||||
print("Parameters:")
|
||||
|
||||
params = tool.get_parameters()
|
||||
for param_name, param_def in params["properties"].items():
|
||||
required = (
|
||||
"required" if param_name in params.get("required", []) else "optional"
|
||||
)
|
||||
print(
|
||||
f" - {param_name} ({param_def['type']}, {required}): {param_def['description']}"
|
||||
)
|
||||
|
||||
# Show how the same parameter type gets different contextual descriptions
|
||||
print("\n=== PARAMETER COMPARISON ===")
|
||||
print("Same parameter 'query_text' with different contextual descriptions:")
|
||||
|
||||
for tool in tools:
|
||||
params = tool.get_parameters()
|
||||
query_text_desc = params["properties"]["query_text"]["description"]
|
||||
print(f" {tool.get_name()}: {query_text_desc}")
|
||||
|
||||
print("\nSame parameter 'top_k' with different contextual descriptions:")
|
||||
for tool in tools:
|
||||
params = tool.get_parameters()
|
||||
top_k_desc = params["properties"]["top_k"]["description"]
|
||||
print(f" {tool.get_name()}: {top_k_desc}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,114 @@
|
||||
# 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 convert a retriever to a tool.
|
||||
|
||||
This example shows:
|
||||
1. How to convert a custom StaticRetriever to a Tool using the convert_to_tool method
|
||||
2. How to define parameters for the tool in the retriever class
|
||||
3. How to execute the tool
|
||||
"""
|
||||
|
||||
import neo4j
|
||||
from typing import Optional, Any, cast
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from neo4j_graphrag.retrievers.base import Retriever
|
||||
from neo4j_graphrag.types import RawSearchResult
|
||||
|
||||
|
||||
# Create a Retriever that returns static results about Neo4j
|
||||
# This would illustrate the conversion process of any Retriever (Vector, Hybrid, etc.)
|
||||
class StaticRetriever(Retriever):
|
||||
"""A retriever that returns static results 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 static information about Neo4j regardless of the query.
|
||||
|
||||
Args:
|
||||
query_text (Optional[str]): The query about Neo4j (any query will return general Neo4j information)
|
||||
**kwargs (Any): Additional keyword arguments (not used)
|
||||
|
||||
Returns:
|
||||
RawSearchResult: Static Neo4j information with metadata
|
||||
"""
|
||||
# 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})
|
||||
|
||||
|
||||
def main() -> None:
|
||||
# Convert a StaticRetriever to a tool using the new convert_to_tool method
|
||||
static_retriever = StaticRetriever(driver=cast(Any, MagicMock()))
|
||||
|
||||
# Convert the retriever to a tool with custom parameter descriptions
|
||||
static_tool = static_retriever.convert_to_tool(
|
||||
name="Neo4jInfoTool",
|
||||
description="Get general information about Neo4j graph database",
|
||||
parameter_descriptions={
|
||||
"query_text": "Any query about Neo4j (the tool returns general information regardless)"
|
||||
},
|
||||
)
|
||||
|
||||
# Print tool information
|
||||
print("Example: StaticRetriever with specific parameters")
|
||||
print(f"Tool Name: {static_tool.get_name()}")
|
||||
print(f"Tool Description: {static_tool.get_description()}")
|
||||
print(f"Tool Parameters: {static_tool.get_parameters()}")
|
||||
print()
|
||||
|
||||
# Execute the tools (in a real application, this would be done by instructions from an LLM)
|
||||
try:
|
||||
# Execute the static retriever tool
|
||||
print("\nExecuting the static retriever tool...")
|
||||
static_result = static_tool.execute(
|
||||
query_text="What is Neo4j?",
|
||||
)
|
||||
print("Static Search Results:")
|
||||
for i, item in enumerate(static_result):
|
||||
print(f"{i + 1}. {str(item)[:100]}...")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error executing tool: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,335 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user