Запросы Neo4j из инструментов агента

Почему Neo4j подходит для знаний агента

Подключение к Neo4j

from neo4j import GraphDatabase
import os

URI = os.environ.get('NEO4J_URI', 'bolt://localhost:7687')
USER = os.environ.get('NEO4J_USER', 'neo4j')
PASSWORD = os.environ.get('NEO4J_PASSWORD', 'password')

driver = GraphDatabase.driver(URI, auth=(USER, PASSWORD))

def test_connection():
    with driver.session() as session:
        result = session.run('RETURN "Connected to Neo4j" AS message')
        record = result.single()
        print(record['message'])

test_connection()

# Always close driver when application exits
# driver.close()

Основные запросы Cypher

from neo4j import GraphDatabase

driver = GraphDatabase.driver('bolt://localhost:7687', auth=('neo4j', 'password'))

def find_company_for_person(person_name: str) -> list:
    with driver.session() as session:
        result = session.run(
            'MATCH (p:Person {name: $name})-[:WORKS_AT]->(c:Company) '
            'RETURN c.name AS company, c.industry AS industry',
            name=person_name
        )
        return [dict(record) for record in result]

def find_colleagues(person_name: str) -> list:
    with driver.session() as session:
        result = session.run(
            'MATCH (p:Person {name: $name})-[:WORKS_AT]->(c:Company) '
            '<-[:WORKS_AT]-(colleague:Person) '
            'WHERE colleague.name <> $name '
            'RETURN DISTINCT colleague.name AS name',
            name=person_name
        )
        return [r['name'] for r in result]

companies = find_company_for_person('Alice Johnson')
print('Works at:', companies)

Параметризованные запросы

from neo4j import GraphDatabase

driver = GraphDatabase.driver('bolt://localhost:7687', auth=('neo4j', 'password'))

# WRONG: vulnerable to injection
def bad_query(name):
    query = f'MATCH (p:Person {{name: "{name}"}}) RETURN p'
    # Never do this
    pass

# RIGHT: parameterized
def good_query(name: str, company: str) -> list:
    with driver.session() as session:
        result = session.run(
            'MATCH (p:Person {name: $name})-[:WORKS_AT]->(c:Company {name: $company}) '
            'RETURN p.name AS person, p.title AS title, c.name AS company',
            name=name,
            company=company
        )
        return [dict(r) for r in result]

# Multiple parameters via dict
def find_by_params(params: dict) -> list:
    with driver.session() as session:
        result = session.run(
            'MATCH (p:Person) WHERE p.name = $name AND p.department = $dept RETURN p',
            **params
        )
        return [dict(r) for r in result]

print('Good query defined (parameterized)')

Запись данных в граф

from neo4j import GraphDatabase

driver = GraphDatabase.driver('bolt://localhost:7687', auth=('neo4j', 'password'))

def upsert_person_works_at_company(person_name: str, company_name: str, title: str):
    with driver.session() as session:
        session.run(
            'MERGE (p:Person {name: $person}) '
            'MERGE (c:Company {name: $company}) '
            'MERGE (p)-[r:WORKS_AT]->(c) '
            'SET r.title = $title, r.updated_at = datetime()',
            person=person_name,
            company=company_name,
            title=title
        )

def create_entity_with_properties(label: str, properties: dict):
    props_string = ', '.join([f'{k}: ${k}' for k in properties.keys()])
    query = f'MERGE (n:{label} {{{props_string}}}) RETURN n'
    with driver.session() as session:
        result = session.run(query, **properties)
        return result.single()

upsert_person_works_at_company('Alice', 'Acme Corp', 'Senior Engineer')
print('Graph data written')

Генерация Cypher из естественного языка

import openai

client = openai.OpenAI(api_key='sk-...')

GRAPH_SCHEMA = '''
Nodes:
- Person: {name, title, email}
- Company: {name, industry, founded_year}
- Product: {name, category, version}

Relationships:
- (Person)-[:WORKS_AT {title, start_date}]->(Company)
- (Person)-[:FOUNDED]->(Company)
- (Company)-[:MAKES]->(Product)
- (Person)-[:USES]->(Product)
'''

def nl_to_cypher(question: str) -> str:
    prompt = (
        f'Graph schema:\n{GRAPH_SCHEMA}\n\n'
        f'Convert this natural language question to a Cypher query:\n{question}\n\n'
        'Return only the Cypher query, no explanation.'
    )
    response = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=[{'role': 'user', 'content': prompt}]
    )
    return response.choices[0].message.content.strip()

question = 'Who are all the people who work at companies that make AI products?'
cypher = nl_to_cypher(question)
print('Generated Cypher:')
print(cypher)

Безопасное выполнение сгенерированного Cypher

import re
from neo4j import GraphDatabase

driver = GraphDatabase.driver('bolt://localhost:7687', auth=('neo4j', 'password'))

MUTATION_KEYWORDS = ['CREATE', 'DELETE', 'MERGE', 'SET', 'REMOVE', 'DROP']

def is_read_only_cypher(cypher: str) -> bool:
    upper = cypher.upper()
    for keyword in MUTATION_KEYWORDS:
        # Check if mutation keyword appears outside of comments
        if re.search(r'\b' + keyword + r'\b', upper):
            return False
    return True

def execute_agent_query(cypher: str, allow_writes=False) -> list:
    if not allow_writes and not is_read_only_cypher(cypher):
        raise ValueError(f'Mutation query blocked. Query: {cypher[:100]}')
    
    with driver.session() as session:
        result = session.run(cypher)
        return [dict(r) for r in result]

# Read query: allowed
read_cypher = 'MATCH (p:Person)-[:WORKS_AT]->(c:Company) RETURN p.name, c.name LIMIT 10'
if is_read_only_cypher(read_cypher):
    print('Read query: safe to execute')

# Write query: blocked
write_cypher = 'DELETE (p:Person {name: "Alice"})'
if not is_read_only_cypher(write_cypher):
    print('Write query: blocked')

Разбор и форматирование результатов запросов

def format_graph_results(records: list, question: str) -> str:
    if not records:
        return f'No results found for: {question}'
    
    # Format as a simple table
    if not records[0]:
        return 'Query returned no data'
    
    headers = list(records[0].keys())
    rows = []
    for record in records:
        row = [str(record.get(h, '')) for h in headers]
        rows.append(' | '.join(row))
    
    header_line = ' | '.join(headers)
    separator = '-' * len(header_line)
    table = '\n'.join([header_line, separator] + rows[:20])  # Cap at 20 rows
    
    result = f'Results for: {question}\n{table}'
    if len(records) > 20:
        result += f'\n... and {len(records) - 20} more results'
    return result

# Simulate some results
sample = [
    {'person': 'Alice', 'company': 'Acme Corp'},
    {'person': 'Bob', 'company': 'TechCo'},
]
print(format_graph_results(sample, 'Who works where?'))

Инструмент агента: поиск по графу

import openai
import json

client = openai.OpenAI(api_key='sk-...')

def graph_lookup_tool(question: str) -> str:
    try:
        # Step 1: Generate Cypher
        cypher = nl_to_cypher(question)
        print(f'Generated Cypher: {cypher}')
        
        # Step 2: Validate
        if not is_read_only_cypher(cypher):
            return 'Error: Generated query contains write operations'
        
        # Step 3: Execute
        records = execute_agent_query(cypher)
        
        # Step 4: Format
        return format_graph_results(records, question)
    
    except Exception as e:
        return f'Graph lookup failed: {str(e)}'

# Register as OpenAI tool
graph_lookup_schema = {
    'type': 'function',
    'function': {
        'name': 'graph_lookup',
        'description': 'Query the knowledge graph to answer questions about entities and their relationships',
        'parameters': {
            'type': 'object',
            'properties': {
                'question': {
                    'type': 'string',
                    'description': 'Natural language question about entities or relationships'
                }
            },
            'required': ['question']
        }
    }
}

print('Graph lookup tool registered')

Многошаговый обход графа

from neo4j import GraphDatabase

driver = GraphDatabase.driver('bolt://localhost:7687', auth=('neo4j', 'password'))

def find_connected_companies(person_name: str, max_hops: int = 3) -> list:
    with driver.session() as session:
        # Variable-length path: 1 to max_hops relationships
        result = session.run(
            f'MATCH (p:Person {{name: $name}})-[:WORKS_AT|FOUNDED*1..{max_hops}]->(c:Company) '
            'RETURN DISTINCT c.name AS company, c.industry AS industry',
            name=person_name
        )
        return [dict(r) for r in result]

def find_shortest_path(entity1: str, entity2: str) -> dict:
    with driver.session() as session:
        result = session.run(
            'MATCH path = shortestPath((a {name: $name1})-[*..6]-(b {name: $name2})) '
            'RETURN [node in nodes(path) | node.name] AS path_nodes, '
            'length(path) AS hops',
            name1=entity1,
            name2=entity2
        )
        record = result.single()
        if record:
            return {'path': record['path_nodes'], 'hops': record['hops']}
        return {'path': [], 'hops': -1}

print('Multi-hop traversal functions defined')

Агрегация в Cypher

from neo4j import GraphDatabase

driver = GraphDatabase.driver('bolt://localhost:7687', auth=('neo4j', 'password'))

def company_employee_stats() -> list:
    with driver.session() as session:
        result = session.run(
            'MATCH (p:Person)-[:WORKS_AT]->(c:Company) '
            'RETURN c.name AS company, '
            'COUNT(p) AS employee_count, '
            'COLLECT(p.name) AS employees '
            'ORDER BY employee_count DESC '
            'LIMIT 10'
        )
        return [dict(r) for r in result]

def count_connections(person_name: str) -> dict:
    with driver.session() as session:
        result = session.run(
            'MATCH (p:Person {name: $name}) '
            'OPTIONAL MATCH (p)-[:WORKS_AT]->(c:Company) '
            'OPTIONAL MATCH (p)-[:FOUNDED]->(fc:Company) '
            'RETURN COUNT(DISTINCT c) AS employers, COUNT(DISTINCT fc) AS founded_companies',
            name=person_name
        )
        record = result.single()
        return dict(record) if record else {}

stats = company_employee_stats()
print('Company stats:', stats[:3])

Проверка знаний: Neo4j для агентов

Итоги инструментов агента Neo4j