Knowledge Graph Construction

from neo4j import GraphDatabase

class ResearchGraph:
    def __init__(self, uri, user, password):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def add_paper(self, paper):
        with self.driver.session() as session:
            session.run("""
                MERGE (p:Paper {doi: $doi})
                SET p.title = $title, p.year = $year, p.abstract = $abstract
            """, **paper)

    def add_citation(self, citing_doi, cited_doi):
        with self.driver.session() as session:
            session.run("""
                MATCH (a:Paper {doi: $citing})
                MATCH (b:Paper {doi: $cited})
                MERGE (a)-[:CITES]->(b)
            """, citing=citing_doi, cited=cited_doi)

    def add_method_usage(self, paper_doi, method_name):
        with self.driver.session() as session:
            session.run("""
                MATCH (p:Paper {doi: $doi})
                MERGE (m:Method {name: $method})
                MERGE (p)-[:USES_METHOD]->(m)
            """, doi=paper_doi, method=method_name)

    def find_papers_using_method(self, method_name):
        with self.driver.session() as session:
            result = session.run("""
                MATCH (p:Paper)-[:USES_METHOD]->(m:Method {name: $method})
                RETURN p.title AS title, p.year AS year, p.doi AS doi
                ORDER BY p.year DESC
            """, method=method_name)
            return [dict(record) for record in result]

    def find_common_methods(self, doi1, doi2):
        with self.driver.session() as session:
            result = session.run("""
                MATCH (p1:Paper {doi: $doi1})-[:USES_METHOD]->(m:Method)
                      <-[:USES_METHOD]-(p2:Paper {doi: $doi2})
                RETURN m.name AS method
            """, doi1=doi1, doi2=doi2)
            return [record['method'] for record in result]

import networkx as nx
import json

def build_research_graph(extracted_data_list):
    """Build a NetworkX graph from extracted paper data."""
    G = nx.MultiDiGraph()

    for data in extracted_data_list:
        for entity in data['entities']:
            G.add_node(
                entity['name'],
                type=entity['type'],
                **entity.get('properties', {})
            )

        for rel in data['relations']:
            G.add_edge(
                rel['source'],
                rel['target'],
                relation=rel['relation']
            )

    return G

# Query the graph
def get_method_landscape(G):
    """Find which methods are most used across papers."""
    methods = [n for n, d in G.nodes(data=True) if d.get('type') == 'Method']
    method_usage = {}
    for method in methods:
        papers = [n for n in G.predecessors(method)
                  if G.nodes[n].get('type') == 'Paper']
        method_usage[method] = len(papers)
    return sorted(method_usage.items(), key=lambda x: x[1], reverse=True)

def kg_rag_query(question, graph, embedding_model, llm):
    """Answer a research question using KG-enhanced RAG."""

    # Step 1: Extract entities from the question
    question_entities = extract_entities(question)

    # Step 2: Retrieve relevant subgraph
    subgraph_nodes = set()
    for entity in question_entities:
        if entity in graph:
            # Get 2-hop neighborhood
            neighbors = nx.ego_graph(graph, entity, radius=2)
            subgraph_nodes.update(neighbors.nodes())

    # Step 3: Format context from subgraph
    context_parts = []
    for node in subgraph_nodes:
        node_data = graph.nodes[node]
        edges = list(graph.edges(node, data=True))
        context_parts.append(
            f"{node} ({node_data.get('type', 'Unknown')}): "
            f"{', '.join(f'{e[2].get(\"relation\", \"related_to\")} {e[1]}' for e in edges[:5])}"
        )
    context = '\n'.join(context_parts[:20])

    # Step 4: Generate answer with LLM
    prompt = f"""Based on the following knowledge graph context, answer the question.

Context:
{context}

Question: {question}

Provide a detailed answer citing specific papers, methods, and findings from the context."""

    return llm.generate(prompt)