GraphRAG and Knowledge Graphs

# script_id: day_027_graphrag/triples_example
# A triple is simply (subject, predicate, object)
triples = [
    ("Globex Corp", "ACQUIRED", "Acme Inc"),
    ("Globex Corp", "HAS_CEO", "Jane Smith"),
    ("Acme Inc", "PRODUCES", "Widgets"),
    ("Acme Inc", "LOCATED_IN", "Austin, TX"),
    ("Jane Smith", "GRADUATED_FROM", "MIT"),
]

# Now you can answer multi-hop questions by traversal:
# Q: "Who is the CEO of the company that acquired Acme?"
# Step 1: Find who acquired Acme -> Globex Corp
# Step 2: Find CEO of Globex Corp -> Jane Smith

# script_id: day_027_graphrag/entity_extraction_pipeline
from openai import OpenAI
from pydantic import BaseModel
import json

client = OpenAI()

class Entity(BaseModel):
    name: str
    type: str  # Person, Company, Product, etc.

class Relationship(BaseModel):
    subject: str
    predicate: str
    object: str

class ExtractionResult(BaseModel):
    entities: list[Entity]
    relationships: list[Relationship]

def extract_entities(text: str) -> ExtractionResult:
    """Extract entities and relationships from text using an LLM."""

    response = client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": """Extract all entities and relationships from the text.
Entities should have a name and type (Person, Company, Product, Location, etc.).
Relationships should be triples: (subject, PREDICATE, object).
Use UPPERCASE for relationship predicates.
Be thorough -- extract every factual relationship you can find."""
            },
            {"role": "user", "content": text}
        ],
        response_format=ExtractionResult,
    )

    return response.choices[0].message.parsed

# Example usage
text = """
Globex Corporation announced the acquisition of Acme Inc for $2.3 billion.
Globex CEO Jane Smith said the deal would strengthen their position in the
widget market. Acme, founded by Bob Johnson in 2015 and headquartered in
Austin, Texas, is known for its premium widget product line. After the
acquisition, Bob Johnson will serve as VP of Product at Globex.
"""

result = extract_entities(text)

print("Entities:")
for entity in result.entities:
    print(f"  ({entity.name}: {entity.type})")

print("\nRelationships:")
for rel in result.relationships:
    print(f"  ({rel.subject}) -[{rel.predicate}]-> ({rel.object})")

Entities:
  (Globex Corporation: Company)
  (Acme Inc: Company)
  (Jane Smith: Person)
  (Bob Johnson: Person)
  (Austin, Texas: Location)

Relationships:
  (Globex Corporation) -[ACQUIRED]-> (Acme Inc)
  (Globex Corporation) -[HAS_CEO]-> (Jane Smith)
  (Acme Inc) -[FOUNDED_BY]-> (Bob Johnson)
  (Acme Inc) -[HEADQUARTERED_IN]-> (Austin, Texas)
  (Bob Johnson) -[WILL_SERVE_AS]-> (VP of Product at Globex)

# script_id: day_027_graphrag/neo4j_crud
# pip install neo4j
from neo4j import GraphDatabase

driver = GraphDatabase.driver(
    "bolt://localhost:7687",
    auth=("neo4j", "your-password")
)

def run_query(query: str, parameters: dict = None) -> list:
    """Execute a Cypher query and return results."""
    with driver.session() as session:
        result = session.run(query, parameters or {})
        return [record.data() for record in result]

# script_id: day_027_graphrag/neo4j_crud
# Create entities as nodes
def create_entity(name: str, entity_type: str, properties: dict = None):
    """Create a node in Neo4j."""
    props = properties or {}
    props["name"] = name

    query = f"""
    MERGE (e:{entity_type} {{name: $name}})
    SET e += $props
    RETURN e
    """
    return run_query(query, {"name": name, "props": props})

# Create relationships
def create_relationship(subject: str, predicate: str, obj: str):
    """Create a relationship between two nodes."""
    query = f"""
    MATCH (a {{name: $subject}})
    MATCH (b {{name: $object}})
    MERGE (a)-[r:{predicate}]->(b)
    RETURN a.name, type(r), b.name
    """
    return run_query(query, {"subject": subject, "object": obj})

# Build the graph from extracted entities
create_entity("Globex Corporation", "Company", {"industry": "Manufacturing"})
create_entity("Acme Inc", "Company", {"founded": 2015})
create_entity("Jane Smith", "Person", {"title": "CEO"})
create_entity("Bob Johnson", "Person", {"title": "Founder"})

create_relationship("Globex Corporation", "ACQUIRED", "Acme Inc")
create_relationship("Globex Corporation", "HAS_CEO", "Jane Smith")
create_relationship("Acme Inc", "FOUNDED_BY", "Bob Johnson")

# script_id: day_027_graphrag/neo4j_crud
# 1. Simple match: Find a company's CEO
results = run_query("""
    MATCH (c:Company {name: "Globex Corporation"})-[:HAS_CEO]->(p:Person)
    RETURN p.name AS ceo
""")
# -> [{"ceo": "Jane Smith"}]

# 2. Multi-hop: CEO of the company that acquired Acme
results = run_query("""
    MATCH (target:Company {name: "Acme Inc"})
          <-[:ACQUIRED]-(acquirer:Company)
          -[:HAS_CEO]->(ceo:Person)
    RETURN ceo.name AS ceo, acquirer.name AS company
""")
# -> [{"ceo": "Jane Smith", "company": "Globex Corporation"}]

# 3. Path finding: How are two entities connected?
results = run_query("""
    MATCH path = shortestPath(
        (a {name: "Jane Smith"})-[*]-(b {name: "Bob Johnson"})
    )
    RETURN [node in nodes(path) | node.name] AS path_nodes,
           [rel in relationships(path) | type(rel)] AS path_rels
""")

# 4. Aggregation: Count subsidiaries
results = run_query("""
    MATCH (parent:Company {name: "Globex Corporation"})-[:ACQUIRED]->(sub:Company)
    RETURN count(sub) AS subsidiary_count
""")

# script_id: day_027_graphrag/entity_extraction_pipeline
from openai import OpenAI
from neo4j import GraphDatabase
import chromadb

client = OpenAI()
neo4j_driver = GraphDatabase.driver("bolt://localhost:7687", auth=("neo4j", "password"))
chroma = chromadb.PersistentClient(path="./vectordb")
collection = chroma.get_or_create_collection("documents")

class GraphRAG:
    """Combines knowledge graph traversal with vector retrieval."""

    def __init__(self):
        self.client = client
        self.driver = neo4j_driver
        self.collection = collection

    def index_document(self, doc_id: str, text: str):
        """Index a document into both the graph and vector store."""

        # 1. Extract entities and relationships
        extraction = extract_entities(text)  # From our earlier function

        # 2. Add to knowledge graph
        with self.driver.session() as session:
            for entity in extraction.entities:
                session.run(
                    f"MERGE (e:{entity.type} {{name: $name}})",
                    {"name": entity.name}
                )
            for rel in extraction.relationships:
                session.run(
                    f"""
                    MATCH (a {{name: $subject}})
                    MATCH (b {{name: $object}})
                    MERGE (a)-[:{rel.predicate}]->(b)
                    """,
                    {"subject": rel.subject, "object": rel.object}
                )

        # 3. Add to vector store
        emb = self.client.embeddings.create(
            model="text-embedding-3-small", input=text
        )
        self.collection.add(
            ids=[doc_id],
            embeddings=[emb.data[0].embedding],
            documents=[text],
            metadatas=[{"entities": ", ".join(e.name for e in extraction.entities)}]
        )

    def query(self, question: str, max_hops: int = 2) -> str:
        """Answer a question using both graph and vector retrieval."""

        # 1. Extract entities from the question
        q_extraction = extract_entities(question)
        entity_names = [e.name for e in q_extraction.entities]

        # 2. Graph traversal: find related entities within max_hops
        graph_context = []
        with self.driver.session() as session:
            for name in entity_names:
                result = session.run(f"""
                    MATCH path = (start {{name: $name}})-[*1..{max_hops}]-(connected)
                    RETURN start.name AS source,
                           [rel in relationships(path) | type(rel)] AS rels,
                           connected.name AS target,
                           labels(connected) AS target_type
                    LIMIT 20
                """, {"name": name})

                for record in result:
                    rel_chain = " -> ".join(record["rels"])
                    graph_context.append(
                        f"{record['source']} -[{rel_chain}]-> "
                        f"{record['target']} ({record['target_type'][0]})"
                    )

        # 3. Vector search for relevant text chunks
        emb = self.client.embeddings.create(
            model="text-embedding-3-small", input=question
        )
        vector_results = self.collection.query(
            query_embeddings=[emb.data[0].embedding],
            n_results=5,
            include=["documents"]
        )
        text_context = "\n".join(vector_results["documents"][0])

        # 4. Combine graph + vector context and generate answer
        combined_context = f"""Graph relationships:
{chr(10).join(graph_context) if graph_context else "No graph relationships found."}

Relevant text:
{text_context}"""

        response = self.client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    "Answer the question using the graph relationships and text context. "
                    "The graph relationships show how entities are connected. "
                    "Use both sources to form a complete answer."
                )},
                {"role": "user", "content": f"Context:\n{combined_context}\n\nQuestion: {question}"}
            ],
            temperature=0
        )

        return response.choices[0].message.content


# Usage
rag = GraphRAG()

# Index documents
rag.index_document("doc1", """
Globex Corporation announced the acquisition of Acme Inc for $2.3 billion.
Globex CEO Jane Smith said the deal strengthens their widget market position.
""")

rag.index_document("doc2", """
Acme Inc, founded by Bob Johnson in 2015 in Austin, Texas, is the leading
producer of premium widgets. Their flagship product, the Widget Pro, has
captured 40% market share.
""")

# Now we can answer multi-hop questions
answer = rag.query("Who is the CEO of the company that acquired the Widget Pro maker?")
print(answer)
# -> "Jane Smith is the CEO of Globex Corporation, which acquired Acme Inc,
#     the maker of Widget Pro."

# script_id: day_027_graphrag/entity_extraction_pipeline
def normalize_entity(name: str) -> str:
    """Basic entity normalization."""
    # Strip common suffixes
    for suffix in [" Inc", " Corp", " Corporation", " LLC", " Ltd"]:
        name = name.replace(suffix, "")
    return name.strip()

# Better: Use the LLM to resolve entities
def resolve_entities(entities: list[str]) -> dict[str, str]:
    """Use LLM to map entity variants to canonical names."""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": f"""
Given these entity names, group them by the real-world entity they refer to.
Return a JSON mapping from each name to its canonical form.

Entities: {json.dumps(entities)}
"""}],
        response_format={"type": "json_object"}
    )
    return json.loads(response.choices[0].message.content)

# Example
mapping = resolve_entities(["Globex Corp", "Globex Corporation", "Globex"])
# -> {"Globex Corp": "Globex Corporation", "Globex": "Globex Corporation", ...}

# script_id: day_027_graphrag/allowed_relations
ALLOWED_RELATIONS = [
    "ACQUIRED", "HAS_CEO", "WORKS_AT", "REPORTS_TO",
    "FOUNDED_BY", "LOCATED_IN", "PRODUCES", "PARTNERS_WITH"
]

# Add to your extraction prompt:
# "Only use these relationship types: {ALLOWED_RELATIONS}"

# script_id: day_027_graphrag/quick_reference
# Extract entities with LLM
result = extract_entities("Globex acquired Acme. Jane Smith is Globex CEO.")

# Create Neo4j nodes
session.run("MERGE (c:Company {name: $name})", {"name": "Globex"})

# Create relationships
session.run("""
    MATCH (a {name: $subj}), (b {name: $obj})
    MERGE (a)-[:ACQUIRED]->(b)
""", {"subj": "Globex", "obj": "Acme"})

# Multi-hop Cypher query
session.run("""
    MATCH (target)<-[:ACQUIRED]-(acquirer)-[:HAS_CEO]->(ceo)
    WHERE target.name = $name
    RETURN ceo.name
""", {"name": "Acme"})

# Combine graph + vector context for the LLM
context = f"Graph: {graph_facts}\nText: {vector_chunks}"