Natural Language to Graph Query Translation with Ollama

from rdflib.plugins.sparql import prepareQuery

def validate_sparql(query: str) -> tuple[bool, str]:
    try:
        prepareQuery(query)
        return True, "Valid"
    except Exception as e:
        return False, str(e)

arq --query myquery.sparql --syntax

from neo4j import GraphDatabase
from neo4j.exceptions import CypherSyntaxError

def validate_cypher(driver, query: str) -> tuple[bool, str]:
    try:
        with driver.session() as session:
            session.run(f"EXPLAIN {query}").consume()
        return True, "Valid"
    except CypherSyntaxError as e:
        return False, str(e)
    except Exception as e:
        # Fail open if Neo4j is temporarily unreachable
        return True, f"Skipped: {e}"

import requests
from pydantic import BaseModel
from rdflib.plugins.sparql import prepareQuery
from ollama import chat

# ── TBox Ontology (inject your full OWL/Turtle TBox here) ─────────────────────

TBOX_TURTLE = """
@prefix rdf:  <http://www.w3.org/1999/02/22-rdf-syntax-ns#> .
@prefix rdfs: <http://www.w3.org/2000/01/rdf-schema#> .
@prefix owl:  <http://www.w3.org/2002/07/owl#> .
@prefix xsd:  <http://www.w3.org/2001/XMLSchema#> .
@prefix ex:   <http://example.org/ontology#> .

### Classes
ex:Person        a owl:Class .
ex:Researcher    a owl:Class ; rdfs:subClassOf ex:Person .
ex:Organization  a owl:Class .
ex:Publication   a owl:Class .

### Object Properties
ex:worksAt        a owl:ObjectProperty ;
                  rdfs:domain ex:Person ;
                  rdfs:range  ex:Organization .
ex:authored       a owl:ObjectProperty ;
                  rdfs:domain ex:Researcher ;
                  rdfs:range  ex:Publication .
ex:affiliatedWith a owl:ObjectProperty ;
                  rdfs:domain ex:Publication ;
                  rdfs:range  ex:Organization .

### Datatype Properties
ex:name      a owl:DatatypeProperty ; rdfs:range xsd:string .
ex:birthYear a owl:DatatypeProperty ; rdfs:range xsd:integer .
ex:title     a owl:DatatypeProperty ; rdfs:range xsd:string .
ex:year      a owl:DatatypeProperty ; rdfs:range xsd:integer .
ex:industry  a owl:DatatypeProperty ; rdfs:range xsd:string .
ex:founded   a owl:DatatypeProperty ; rdfs:range xsd:integer .
"""

# ── System Prompt ─────────────────────────────────────────────────────────────

SPARQL_SYSTEM_PROMPT = f"""You are an expert SPARQL 1.1 query generator for RDF knowledge graphs.
Given an OWL/Turtle ontology TBox and a natural language question, generate a valid SPARQL query.

Rules:
- Use ONLY the classes and properties defined in the TBox
- Always declare PREFIX statements at the top
- Use OPTIONAL {{ }} for properties that may not exist on all instances
- Use FILTER for conditions on literals
- Use SELECT DISTINCT unless aggregation is required
- For aggregations, use GROUP BY with the appropriate aggregate function

Ontology TBox (Turtle format):
{TBOX_TURTLE}

Examples:
Q: Who are all the researchers?
SPARQL:
PREFIX ex: <http://example.org/ontology#>
SELECT DISTINCT ?name WHERE {{
  ?r a ex:Researcher ; ex:name ?name .
}}

Q: Which organizations are in the healthcare industry?
SPARQL:
PREFIX ex: <http://example.org/ontology#>
SELECT DISTINCT ?orgName WHERE {{
  ?o a ex:Organization ; ex:industry "healthcare" ; ex:name ?orgName .
}}
"""

# ── Structured Output Schema ──────────────────────────────────────────────────

class SPARQLResult(BaseModel):
    query: str               # Full SPARQL query including PREFIX declarations
    query_type: str          # "SELECT" | "ASK" | "CONSTRUCT" | "DESCRIBE"
    classes_used: list[str]  # TBox classes referenced, e.g. ["ex:Researcher"]
    explanation: str         # One-sentence description of what the query does

# ── Validation ────────────────────────────────────────────────────────────────

def validate_sparql(query: str) -> tuple[bool, str]:
    try:
        prepareQuery(query)
        return True, "Valid"
    except Exception as e:
        return False, str(e)

# ── Generation with Retry Loop ────────────────────────────────────────────────

def nl_to_sparql(
    question: str,
    model: str = "qwen2.5:32b",
    max_retries: int = 3
) -> SPARQLResult:
    error_context = ""

    for attempt in range(1, max_retries + 1):
        user_content = f"Question: {question}"
        if error_context:
            user_content += (
                f"\n\nYour previous attempt produced an invalid SPARQL query.\n"
                f"Error: {error_context}\n"
                f"Please correct the query."
            )

        response = chat(
            model=model,
            messages=[
                {"role": "system", "content": SPARQL_SYSTEM_PROMPT},
                {"role": "user",   "content": user_content}
            ],
            format=SPARQLResult.model_json_schema(),
            options={"temperature": 0.0}
        )

        result = SPARQLResult.model_validate_json(response.message.content)
        is_valid, error = validate_sparql(result.query)

        print(f"  [Attempt {attempt}] {'✓ Valid' if is_valid else f'✗ {error}'}")

        if is_valid:
            return result
        error_context = error

    raise ValueError(f"Failed to generate valid SPARQL after {max_retries} attempts.")

# ── Execute Against Apache Jena Fuseki ────────────────────────────────────────

def run_sparql(
    query: str,
    endpoint: str = "http://localhost:3030/ds/sparql"
) -> list[dict]:
    resp = requests.get(
        endpoint,
        params={"query": query},
        headers={"Accept": "application/sparql-results+json"},
        timeout=30
    )
    resp.raise_for_status()
    bindings = resp.json().get("results", {}).get("bindings", [])
    return [{k: v["value"] for k, v in row.items()} for row in bindings]

# ── Main ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    questions = [
        "List all researchers and the organizations they work at.",
        "How many publications were authored by researchers in the AI industry?",
        "Find all publications affiliated with organizations founded before 2000.",
    ]

    for q in questions:
        print(f"\n{'='*60}")
        print(f"Q: {q}")
        result = nl_to_sparql(q)
        print(f"Type       : {result.query_type}")
        print(f"Classes    : {result.classes_used}")
        print(f"Explanation: {result.explanation}")
        print(f"Query:\n{result.query}")

        # Uncomment to run against a live Fuseki endpoint:
        # rows = run_sparql(result.query)
        # for row in rows:
        #     print(row)

{
  "query": "PREFIX ex: <http://example.org/ontology#>\nSELECT (COUNT(DISTINCT ?pub) AS ?pubCount)\nWHERE {\n  ?r a ex:Researcher ;\n     ex:worksAt ?org ;\n     ex:authored ?pub .\n  ?org ex:industry \"AI\" .\n}",
  "query_type": "SELECT",
  "classes_used": ["ex:Researcher", "ex:Organization", "ex:Publication"],
  "explanation": "Counts distinct publications authored by researchers at organizations in the AI industry."
}

from pydantic import BaseModel
from neo4j import GraphDatabase
from neo4j.exceptions import CypherSyntaxError
from ollama import chat

# ── LPG Schema ────────────────────────────────────────────────────────────────

LPG_SCHEMA = """
Node Labels:
- Person(id: STRING, name: STRING, birthYear: INT, role: STRING)
- Organization(id: STRING, name: STRING, industry: STRING, founded: INT)
- Publication(id: STRING, title: STRING, year: INT, doi: STRING)

Relationship Types:
- (Person)-[:WORKS_AT {since: INT, position: STRING}]->(Organization)
- (Person)-[:AUTHORED {contribution: STRING}]->(Publication)
- (Person)-[:COLLABORATES_WITH]->(Person)
- (Organization)-[:FUNDED]->(Publication)
"""

# ── System Prompt ─────────────────────────────────────────────────────────────

CYPHER_SYSTEM_PROMPT = f"""You are an expert Neo4j Cypher query generator.
Given a graph schema and a natural language question, generate a valid, executable Cypher query.

Rules:
- Use ONLY the node labels, properties, and relationship types defined in the schema
- Use MATCH...WHERE...RETURN structure
- Prefer OPTIONAL MATCH for potentially absent relationships
- Use WITH for intermediate aggregations
- Use ORDER BY and LIMIT where appropriate

Graph Schema:
{LPG_SCHEMA}

Examples:
Q: Find all people who work in the technology industry.
Cypher:
MATCH (p:Person)-[:WORKS_AT]->(o:Organization)
WHERE o.industry = "technology"
RETURN p.name, p.role, o.name AS organization

Q: Which people have authored publications?
Cypher:
MATCH (p:Person)-[:AUTHORED]->(pub:Publication)
RETURN DISTINCT p.name, collect(pub.title) AS publications
"""

# ── Structured Output Schema ──────────────────────────────────────────────────

class CypherResult(BaseModel):
    query: str                    # Full executable Cypher query
    query_intent: str             # "MATCH" | "AGGREGATION" | "PATH" | "EXISTENCE"
    node_labels_used: list[str]   # e.g. ["Person", "Organization"]
    relationships_used: list[str] # e.g. ["WORKS_AT", "AUTHORED"]
    explanation: str              # One-sentence description

# ── Validator ─────────────────────────────────────────────────────────────────

class CypherValidator:
    def __init__(
        self,
        uri: str = "bolt://localhost:7687",
        user: str = "neo4j",
        password: str = "password"
    ):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def validate(self, query: str) -> tuple[bool, str]:
        try:
            with self.driver.session() as session:
                session.run(f"EXPLAIN {query}").consume()
            return True, "Valid"
        except CypherSyntaxError as e:
            return False, str(e)
        except Exception as e:
            # Fail open on connectivity issues — don't block the pipeline
            print(f"  Warning: Neo4j unreachable, skipping validation: {e}")
            return True, "Skipped (no connection)"

    def run(self, query: str) -> list[dict]:
        with self.driver.session() as session:
            result = session.run(query)
            return [dict(record) for record in result]

    def close(self):
        self.driver.close()

# ── Generation with Retry Loop ────────────────────────────────────────────────

def nl_to_cypher(
    question: str,
    validator: CypherValidator,
    model: str = "qwen2.5:32b",
    max_retries: int = 3
) -> CypherResult:
    error_context = ""

    for attempt in range(1, max_retries + 1):
        user_content = f"Question: {question}"
        if error_context:
            user_content += (
                f"\n\nYour previous attempt produced an invalid Cypher query.\n"
                f"Error: {error_context}\n"
                f"Please correct the query."
            )

        response = chat(
            model=model,
            messages=[
                {"role": "system", "content": CYPHER_SYSTEM_PROMPT},
                {"role": "user",   "content": user_content}
            ],
            format=CypherResult.model_json_schema(),
            options={"temperature": 0.0}
        )

        result = CypherResult.model_validate_json(response.message.content)
        is_valid, error = validator.validate(result.query)

        print(f"  [Attempt {attempt}] {'✓ Valid' if is_valid else f'✗ {error}'}")

        if is_valid:
            return result
        error_context = error

    raise ValueError(f"Failed to generate valid Cypher after {max_retries} attempts.")

# ── Main ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    validator = CypherValidator(
        uri="bolt://localhost:7687",
        user="neo4j",
        password="your-password"
    )

    questions = [
        "Who are all the people working at organizations in the AI industry?",
        "Find all publications authored by researchers who collaborate with each other.",
        "Which organizations have funded more than 5 publications since 2020?",
    ]

    try:
        for q in questions:
            print(f"\n{'='*60}")
            print(f"Q: {q}")
            result = nl_to_cypher(q, validator)
            print(f"Intent        : {result.query_intent}")
            print(f"Nodes         : {result.node_labels_used}")
            print(f"Relationships : {result.relationships_used}")
            print(f"Explanation   : {result.explanation}")
            print(f"Query:\n{result.query}")

            # Uncomment to execute against a live Neo4j instance:
            # rows = validator.run(result.query)
            # for row in rows:
            #     print(row)
    finally:
        validator.close()

{
  "query": "MATCH (o:Organization)-[:FUNDED]->(pub:Publication)\nWHERE pub.year >= 2020\nWITH o, COUNT(pub) AS pubCount\nWHERE pubCount > 5\nRETURN o.name, pubCount\nORDER BY pubCount DESC",
  "query_intent": "AGGREGATION",
  "node_labels_used": ["Organization", "Publication"],
  "relationships_used": ["FUNDED"],
  "explanation": "Finds organizations that have funded more than 5 publications from 2020 onward, sorted by count."
}