mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
34329840e1
* Add Neo4j GraphRAG samples * Fix sample CI issues * Address sample review feedback * Move Neo4j Python sample to end-to-end * Make Neo4j GraphRAG sample self-contained * Remove unused central package versions
113 lines
3.8 KiB
Python
113 lines
3.8 KiB
Python
# /// script
|
|
# requires-python = ">=3.10"
|
|
# dependencies = [
|
|
# "agent-framework-foundry",
|
|
# "agent-framework-neo4j",
|
|
# ]
|
|
# ///
|
|
|
|
# Copyright (c) Microsoft. All rights reserved.
|
|
|
|
import asyncio
|
|
import os
|
|
|
|
from agent_framework import Agent
|
|
from agent_framework.foundry import FoundryChatClient
|
|
from agent_framework_neo4j import Neo4jContextProvider, Neo4jSettings
|
|
from azure.identity.aio import AzureCliCredential
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
"""
|
|
This sample demonstrates how to use the Neo4j GraphRAG context provider with
|
|
Agent Framework and Azure AI Foundry.
|
|
|
|
Environment variables:
|
|
FOUNDRY_PROJECT_ENDPOINT — Azure AI Foundry project endpoint
|
|
FOUNDRY_MODEL — Model deployment name (e.g. gpt-4o)
|
|
NEO4J_URI — Neo4j connection URI
|
|
NEO4J_USERNAME — Neo4j username
|
|
NEO4J_PASSWORD — Neo4j password
|
|
NEO4J_FULLTEXT_INDEX_NAME — Optional fulltext index name (defaults to search_chunks)
|
|
"""
|
|
|
|
USER_INPUTS = [
|
|
"What products does Microsoft offer?",
|
|
"What risks does Apple face?",
|
|
"Tell me about NVIDIA's AI business and risk factors.",
|
|
]
|
|
|
|
# Optional graph-enrichment query: retrieval works without this, but supplying
|
|
# a query lets the sample attach related company, product, and risk metadata to
|
|
# each retrieved chunk.
|
|
RETRIEVAL_QUERY = """
|
|
MATCH (node)-[:FROM_DOCUMENT]->(doc:Document)<-[:FILED]-(company:Company)
|
|
OPTIONAL MATCH (company)-[:FACES_RISK]->(risk:RiskFactor)
|
|
WITH node, score, company, doc, collect(DISTINCT risk.name)[0..5] AS risks
|
|
OPTIONAL MATCH (company)-[:MENTIONS]->(product:Product)
|
|
WITH node, score, company, doc, risks, collect(DISTINCT product.name)[0..5] AS products
|
|
RETURN
|
|
node.text AS text,
|
|
score,
|
|
company.name AS company,
|
|
company.ticker AS ticker,
|
|
doc.title AS title,
|
|
risks,
|
|
products
|
|
ORDER BY score DESC
|
|
"""
|
|
|
|
|
|
async def main() -> None:
|
|
# 1. Load and validate the Neo4j connection settings.
|
|
settings = Neo4jSettings()
|
|
if not settings.is_configured:
|
|
raise RuntimeError("Set NEO4J_URI, NEO4J_USERNAME, and NEO4J_PASSWORD before running this sample.")
|
|
|
|
# 2. Read the Azure AI Foundry project endpoint and model configuration.
|
|
project_endpoint = os.environ.get("FOUNDRY_PROJECT_ENDPOINT")
|
|
if not project_endpoint:
|
|
raise RuntimeError("Set FOUNDRY_PROJECT_ENDPOINT before running this sample.")
|
|
|
|
model = os.environ.get("FOUNDRY_MODEL") or "gpt-4o"
|
|
|
|
# 3. Create the Neo4j context provider and Foundry-backed agent, then ask sample questions.
|
|
async with (
|
|
AzureCliCredential() as credential,
|
|
Neo4jContextProvider(
|
|
source_id="neo4j_graphrag",
|
|
uri=settings.uri,
|
|
username=settings.username,
|
|
password=settings.get_password(),
|
|
index_name=settings.fulltext_index_name,
|
|
index_type="fulltext",
|
|
retrieval_query=RETRIEVAL_QUERY,
|
|
top_k=5,
|
|
) as provider,
|
|
Agent(
|
|
client=FoundryChatClient(
|
|
project_endpoint=project_endpoint,
|
|
model=model,
|
|
credential=credential,
|
|
),
|
|
name="Neo4jGraphRAGAgent",
|
|
instructions=(
|
|
"You are a helpful assistant. Use the Neo4j context provider results to answer accurately. "
|
|
"If the retrieved context is insufficient, say so plainly."
|
|
),
|
|
context_providers=[provider],
|
|
) as agent,
|
|
):
|
|
session = agent.create_session()
|
|
print("=== Neo4j GraphRAG Context Provider ===\n")
|
|
|
|
for user_input in USER_INPUTS:
|
|
print(f"User: {user_input}")
|
|
result = await agent.run(user_input, session=session)
|
|
print(f"Agent: {getattr(result, 'text', result)}\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|