repograph | khushi keswani

what i built

ast extraction — clones a repo, parses every .py file, and extracts files, functions, classes, methods and imports as typed nodes.
call resolution — builds CONTAINS, IMPORTS and CALLS edges. same-named functions in different files stay distinct.
hybrid retrieval — embeds each function, class and method, finds the closest ones by cosine similarity, then expands along the call graph before asking the llm.
interactive viewer — a pyvis graph where you pick a file and click a node to highlight its connections.

how it fits together

figure 1build once per repo, query many times.(scroll the figure sideways)

ids that survive name collisions

a function's identity is its path plus name, not its name. methods add the class. so three different load_params in three files are three nodes, and lookups are keyed on (file, name).

"id": f"{relpath}::{node.name}"                    # function
"id": f"{relpath}::{node.name}::{child.name}"      # method

function_lookup[(node["file"], node["name"])] = node["id"]

build.py

resolving a call is a small type-inference problem

caller() is easy. tts.speak() needs to know what tts is. for each function body, the builder watches assignments like tts = Tts(), records that tts is an instance of that class, and then resolves tts.speak() to the exact method node. self.method() resolves within the current class, and imported names resolve by matching the module path to a file.

if isinstance(inner, ast.Assign) and isinstance(inner.value, ast.Call):
    class_id = class_lookup.get((relpath, inner.value.func.id))
    if class_id:
        for t in inner.targets:
            object_types[t.id] = class_id      # tts -> Tts

# later: tts.speak()
target = method_lookup.get((class_file, class_name, method_name))

build.py

calls it can't resolve (third-party code, dynamic dispatch) are skipped, not guessed.

why graph expansion beats plain rag

vector search returns code that sounds like the question. ask "how is data cleaned" and you get preprocess_data, but not the function it calls next. a breadth-first walk of up to 2 hops over CALLS edges from each of the top 3 seeds pulls in that connected code, labeled by file and function name. the prompt tells the model to say so if the answer isn't in the context.

def get_2hop_nodes(G, start_node, max_hops=2):
    visited, queue = {start_node}, [(start_node, 0)]
    while queue:
        node, depth = queue.pop(0)
        if depth == max_hops:
            continue
        for nb in G.successors(node):
            if any(e.get("type") == "CALLS"
                   for e in G.get_edge_data(node, nb).values()):
                if nb not in visited:
                    visited.add(nb)
                    queue.append((nb, depth + 1))
    return visited

query.py

on an aapl stock-prediction repo it traced raw data through load_data → preprocess_data → scaling_data → create_sequences → build_model across three files, and separated cleaning from feature engineering.

metrics

stage	setting
graph size	1,000+ nodes and 500+ edges on the reference repo
node types	file · function · class · method
edge types	`CONTAINS` · `IMPORTS` · `CALLS`
embeddings	minilm, one vector per function, class and method
retrieval	top 3 seeds by cosine similarity
expansion	2 hops along `CALLS` edges from each seed

k	precision@k	recall@k
1	0.480	0.480
3	0.307	0.720
5	0.240	0.880
10	0.156	1.000

measured on a labeled set of `(query, expected function)` pairs. precision@k is how much of the top k is relevant; recall@k is how much of the relevant set was recovered.

metric	value
mrr	0.636
relevant-context recall	0.767
cross-file recall	0.840

precision falls as k grows while recall climbs to 1.000, which is the expected shape: at k=10 the answer set is complete but noisier, and at k=1 just under half the time the single best hit is the function you wanted.

because “search the codebase” deserved better.

next project → conversational chatbot all work ↑