Knowledge Graph Builder

+------------------+     +-------------------+     +------------------+
| Data Sources     | --> | Extraction        | --> | Graph Database   |
| - Documents      |     | - NER             |     | (Person)-[:WORKS_AT]->(Company)
| - APIs           |     | - Relation extract|     | (Paper)-[:CITES]->(Paper)
| - Databases      |     | - Entity linking  |     | (Product)-[:HAS]->(Feature)
+------------------+     +-------------------+     +------------------+
                                                          |
+------------------+     +-------------------+     +------v-----------+
| Applications     | <-- | Query & Reasoning | <-- | Indexing &       |
| - GraphRAG       |     | - Cypher / SPARQL |     | Enrichment       |
| - Recommendations|     | - Path finding    |     | - Embeddings     |
| - Analytics      |     | - Subgraph match  |     | - Community det. |
+------------------+     +-------------------+     +------------------+

Selection guide:
  AWS native? -> Amazon Neptune
  Need RDF/SPARQL? -> Neptune or Apache Jena
  Graph + document hybrid? -> ArangoDB (multi-model)
  SQL-like preferred? -> Neo4j (Cypher is most approachable)
  GraphQL preferred? -> Dgraph
  Enterprise trillions of edges? -> TigerGraph
  Default -> Neo4j (largest community, vector search, mature)

1. Start with competency questions
   "Who authored papers about topic X?"
   "What products share features?"

2. Identify entity types (nouns in questions)

3. Define relationship types (verbs connecting nouns)
   Always directional: (Subject)-[:VERB]->(Object)

4. Add properties (attributes of nodes and edges)
   Include temporal properties (created_at, valid_from)

5. Define constraints (uniqueness, required props, cardinality)

CREATE CONSTRAINT person_id IF NOT EXISTS
FOR (p:Person) REQUIRE p.id IS UNIQUE;

CREATE CONSTRAINT company_id IF NOT EXISTS
FOR (c:Company) REQUIRE c.id IS UNIQUE;

CREATE FULLTEXT INDEX entity_search IF NOT EXISTS
FOR (n:Person|Company|Paper) ON EACH [n.name, n.title, n.description];

CREATE VECTOR INDEX entity_embeddings IF NOT EXISTS
FOR (n:Entity) ON (n.embedding)
OPTIONS {indexConfig: {
  `vector.dimensions`: 1536,
  `vector.similarity_function`: 'cosine'
}};

Temporal:     (Person)-[:WORKS_AT {from: date, to: date}]->(Company)
Hierarchical: (Category)-[:SUBCATEGORY_OF]->(Category)
Provenance:   (Fact)-[:EXTRACTED_FROM {confidence: 0.95}]->(Document)
Events:       (Person)-[:PARTICIPATED_IN]->(Event)-[:OCCURRED_AT]->(Location)
Multi-hop:    (Drug)-[:TREATS]->(Disease)-[:HAS_SYMPTOM]->(Symptom)

from openai import OpenAI
import json

class KnowledgeExtractor:
    def __init__(self, model: str = "gpt-4o-mini"):
        self.client = OpenAI()
        self.model = model

    def extract(self, text: str, schema: dict = None) -> dict:
        schema_inst = ""
        if schema:
            schema_inst = (f"\nEntity types: {schema.get('entity_types', [])}\n"
                          f"Relation types: {schema.get('relation_types', [])}\n")

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "Extract precise entities and relationships. "
                 "Only extract what is explicitly stated or strongly implied."},
                {"role": "user", "content": f"Extract entities and relationships:{schema_inst}\n"
                 f"Return JSON: {{\"entities\": [{{\"id\", \"type\", \"name\", \"properties\"}}], "
                 f"\"relationships\": [{{\"source\", \"target\", \"type\", \"properties\"}}]}}\n\n"
                 f"Text: {text}"},
            ],
            response_format={"type": "json_object"}, temperature=0,
        )
        return json.loads(response.choices[0].message.content)

    def extract_batch(self, documents: list[str], schema: dict = None) -> dict:
        all_entities = {}
        all_rels = []
        for doc in documents:
            result = self.extract(doc, schema)
            for e in result.get("entities", []):
                key = (e["type"], e["name"].lower())
                if key not in all_entities:
                    all_entities[key] = e
                else:
                    all_entities[key]["properties"].update(e.get("properties", {}))
            all_rels.extend(result.get("relationships", []))
        return {"entities": list(all_entities.values()), "relationships": all_rels}

from neo4j import GraphDatabase

class GraphRAGPipeline:
    def __init__(self, neo4j_uri: str, neo4j_auth: tuple, llm_model: str = "gpt-4o"):
        self.driver = GraphDatabase.driver(neo4j_uri, auth=neo4j_auth)
        self.client = OpenAI()
        self.llm_model = llm_model

    def query(self, question: str) -> dict:
        entities = self._extract_query_entities(question)
        subgraph = self._retrieve_subgraph(entities)
        answer = self._generate_answer(question, subgraph)
        return {"answer": answer, "entities_found": entities,
                "subgraph_nodes": len(subgraph["nodes"])}

    def _extract_query_entities(self, question: str) -> list[str]:
        resp = self.client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content":
                f"Extract entity names. JSON: {{\"entities\": [...]}}\nQuestion: {question}"}],
            response_format={"type": "json_object"}, temperature=0,
        )
        return json.loads(resp.choices[0].message.content).get("entities", [])

    def _retrieve_subgraph(self, entities: list[str], hops: int = 2) -> dict:
        nodes, edges = [], []
        with self.driver.session() as session:
            for entity in entities:
                result = session.run("""
                    CALL db.index.fulltext.queryNodes('entity_search', $name)
                    YIELD node, score WHERE score > 0.5
                    WITH node LIMIT 5
                    CALL apoc.path.subgraphAll(node, {maxLevel: $hops, limit: 50})
                    YIELD nodes, relationships
                    RETURN nodes, relationships
                """, name=entity, hops=hops)
                for record in result:
                    nodes.extend([{"labels": list(n.labels), "props": dict(n)} for n in record["nodes"]])
                    edges.extend([{"type": r.type} for r in record["relationships"]])
        return {"nodes": nodes, "edges": edges}

    def _generate_answer(self, question: str, subgraph: dict) -> str:
        context = "\n".join(
            f"[{'/'.join(n['labels'])}] {n['props']}" for n in subgraph["nodes"][:30]
        )
        resp = self.client.chat.completions.create(
            model=self.llm_model,
            messages=[
                {"role": "system", "content": "Answer using the knowledge graph context. Cite entities."},
                {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"},
            ], temperature=0,
        )
        return resp.choices[0].message.content

-- Bound traversal depth (never use unbounded [*])
MATCH (a:Person {name: "Alice"})-[*1..3]->(b:Company)
RETURN b.name;

-- Use full-text index instead of CONTAINS
CALL db.index.fulltext.queryNodes('entity_search', 'Alice') YIELD node;

-- Profile to find bottlenecks
PROFILE MATCH (a:Person)-[:WORKS_AT]->(c:Company)-[:LOCATED_IN]->(city:City)
WHERE city.name = "San Francisco"
RETURN a.name, c.name;

-- Use parameters (enables query plan caching)
MATCH (p:Person {name: $name}) RETURN p;