GraphRAG: Build a Knowledge Graph from a Wiki with an LLM and Answer Multi-Hop Questions

"""A knowledge graph of Larkspur Pay's systems, built from the engineering wiki. You write the functions marked
TODO, one step at a time."""
import difflib
import json
import re
from concurrent.futures import ThreadPoolExecutor

import harness as H

RELATIONS = {
    "depends_on": "X needs Y to work: X calls Y, reads from it or writes to it (X is a service or datastore)",
    "owned_by": "X (a service or datastore) is owned and run by team Y",
    "led_by": "team X is led by person Y",
}


# ---------- Step 1: extract triples ----------

def extraction_prompt(page):
    """The prompt that asks the model for the page's (subject, relation, object) triples as JSON."""
    # TODO (Step 1): list the RELATIONS, the rules for what counts as a fact, the JSON format
    # with an example, and the page.
    raise NotImplementedError("Step 1: write extraction_prompt()")


def parse_triples(reply):
    """[(subject, relation, object), ...] from the model's reply: the JSON list inside it, keeping only well-formed
    items whose relation is in RELATIONS. Anything unreadable gives []."""
    # TODO (Step 1): find the JSON list in the reply (re.search(r"\[.*\]", reply, re.S)), json.loads it,
    # keep dicts with non-empty string subject/relation/object whose normalised relation is in RELATIONS.
    raise NotImplementedError("Step 1: write parse_triples()")


def extract_all(pages, workers=8):
    """{page id: parse_triples(the model's reply to extraction_prompt(text))}, asking about up to `workers` pages at
    a time."""
    # TODO (Step 1): with ThreadPoolExecutor(workers), H.ask_model(extraction_prompt(text)) per page, then
    # parse_triples() each reply.
    raise NotImplementedError("Step 1: write extract_all()")


# ---------- Step 2: one node per thing ----------

TYPES = {"depends_on": ({"service", "datastore"}, {"service", "datastore"}),
         "owned_by": ({"service", "datastore"}, {"team"}),
         "led_by": ({"team"}, {"person"})}


class Graph:
    def __init__(self):
        self.edges = set()      # {(subject id, relation, object id)}
        self.sources = {}       # {edge: [page ids that state it]}
        self.dropped = []       # [(page id, (subject, relation, object) as extracted, reason)]

    def add(self, edge, page):
        self.edges.add(edge)
        self.sources.setdefault(edge, [])
        if page not in self.sources[edge]:
            self.sources[edge].append(page)


def normalise(name):
    """Lower case; "-", "_" and "." become spaces; no leading "the "; no trailing " service" or " team"; single spaces."""
    raise NotImplementedError("normalise() arrives in Step 2")


def build_lookup(catalogue):
    """{normalised id, name or alias: catalogue id} for every entity."""
    raise NotImplementedError("build_lookup() arrives in Step 2")


def resolve(name, lookup):
    """The catalogue id for a name: exact match after normalise(), else the closest key by difflib with cutoff 0.85,
    else None."""
    raise NotImplementedError("resolve() arrives in Step 2")


def build_graph(extracted, catalogue):
    """A Graph of resolved triples. A triple is dropped (with a reason) if a name does not resolve, or if the
    types do not fit TYPES (a team cannot depend on anything, for example)."""
    raise NotImplementedError("build_graph() arrives in Step 2")


def precision_recall(edges, labelled):
    """(precision, recall) of a set of edges against the hand-labelled triples."""
    raise NotImplementedError("precision_recall() arrives in Step 2")


# ---------- Step 3: walk the graph ----------

def hop(graph, nodes, step):
    """The nodes one step away: "owned_by" follows edges forwards, "~owned_by" backwards, and a trailing "*" repeats
    the step until nothing new is reached (the result never includes the nodes you started from)."""
    raise NotImplementedError("hop() arrives in Step 3")


def run_plan(graph, catalogue, plan):
    """Sorted ids: start at plan["start"], apply each hop in plan["hops"], then keep only nodes matching every
    plan["where"] condition ("type", "public", or "owned_by": a team id)."""
    raise NotImplementedError("run_plan() arrives in Step 3")


# ---------- Step 4: questions become graph queries ----------

HOP_RE = re.compile(r"^~?(depends_on|owned_by|led_by)\*?$")


def plan_prompt(question, catalogue):
    """The prompt that turns a question into a plan: {"start": id, "hops": [...], "where": {...}} as JSON."""
    raise NotImplementedError("plan_prompt() arrives in Step 4")


def parse_plan(reply, catalogue):
    """The plan in the reply as a dict, or None unless: "start" is a catalogue id, "hops" is a non-empty list of
    hops matching HOP_RE, and "where" (if present) is a dict whose keys are "type", "public" or "owned_by"."""
    raise NotImplementedError("parse_plan() arrives in Step 4")


def graph_answer(graph, catalogue, question):
    """(ids, plan): run_plan() on the plan the model writes for the question, or (None, None) if the plan is
    invalid."""
    raise NotImplementedError("graph_answer() arrives in Step 4")


# ---------- Step 5: which questions each method answers ----------

def answer_f1(pred, gold):
    """F1 between two sets of ids: 1.0 if both are empty, 0.0 if only one is (None counts as empty)."""
    raise NotImplementedError("answer_f1() arrives in Step 5")


def evaluate(answers, questions):
    """{question type: mean answer_f1, ..., "all": mean over every question}; answers[i] answers questions[i]."""
    raise NotImplementedError("evaluate() arrives in Step 5")


def hybrid(graph_ids, vector_ids):
    """The graph's answer when it has one (not None and not empty), otherwise the vector answer."""
    raise NotImplementedError("hybrid() arrives in Step 5")