from langchain.chains import create_extraction_chain
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o", temperature=0)

schema = {
    "properties": {
        "entities": {"type": "array", "items": {"type": "string"}},
        "relationships": {"type": "array", "items": {
            "type": "object",
            "properties": {
                "source": {"type": "string"},
                "relation": {"type": "string"},
                "target": {"type": "string"}
            }
        }}
    }
}

extractor = create_extraction_chain(llm, schema)
result = extractor.run(document_text)

from neo4j import GraphDatabase

class KnowledgeGraph:
    def __init__(self, uri, user, password):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))
    
    def add_entity(self, entity_name, entity_type, source):
        with self.driver.session() as session:
            session.run(
                "MERGE (e:Entity {name: $name, type: $type, source: $source})",
                name=entity_name, type=entity_type, source=source
            )
    
    def add_relationship(self, source, relation, target):
        with self.driver.session() as session:
            session.run(
                """MATCH (s:Entity {name: $source})
                   MATCH (t:Entity {name: $target})
                   MERGE (s)-[r:RELATES {type: $relation}]->(t)""",
                source=source, relation=relation, target=target
            )

def retrieve_with_context(query_entity, max_hops=2):
    with driver.session() as session:
        result = session.run(
            """MATCH path = (start:Entity {name: $entity})-[*1..$hops]-(end:Entity)
               RETURN path""",
            entity=query_entity, hops=max_hops
        )
        return [record["path"] for record in result]