Idea: token-budgeted connected subgraph retrieval for LLM context

[berendmarkhorst]

Hi! First of all, thanks for building Graphify! I came across the project recently and really like the idea of representing a codebase as an explicit knowledge graph rather than relying only on vector-based retrieval.

I maintain SteinerPy, a Python package for Steiner tree problems and related graph optimization problems. Looking at Graphify's graph.json output made me wonder whether graph optimization could be useful for selecting the context that is eventually passed to an LLM.

The idea

Suppose an agent asks a question about a large codebase. Graphify gives us a graph containing functions, classes, files, documentation, and their relations.

Each node could additionally be assigned:

• a relevance score for the current question;
• a cost, for example the number of tokens required to include the corresponding source fragment in the LLM context.

We could then select a connected subgraph that maximizes the relevance of the included information while respecting a fixed token budget.

In a simplified form:

$$ \max \sum_{v \in V} r_v x_v $$

subject to

$$ \sum_{v \in V} c_v x_v \leq B, $$

where $r_v$ is the relevance of node $v$, $c_v$ its token cost, and $B$ the available context budget, together with constraints requiring the selected nodes to form a connected subgraph.

This is closely related to the budgeted maximum-weight connected subgraph problem, for which SteinerPy already contains an implementation.

Why connectivity might help

A simple alternative would of course be to take the $k$ most relevant nodes. The possible advantage of using the graph structure is that some information may not be particularly relevant on its own, but may be important for explaining how two highly relevant pieces of information are related.

For example:

highly relevant function
        |
   helper function
        |
highly relevant test

The helper function may receive a relatively low relevance score and therefore be excluded by top-(k) retrieval. A connected-subgraph formulation can still include it if it is a useful and inexpensive connection between the other two nodes.

Possible experiment

I think a small experiment could already show whether this is useful. For example, compare:

1. Graphify's current query/retrieval approach;
2. the top-(k) most relevant Graphify nodes;
3. a token-budgeted connected subgraph selected using SteinerPy.

All methods could receive exactly the same token budget, e.g. 4k, 8k, or 16k tokens. We could then compare both the number of tokens used and the quality of the resulting answer or coding task.

Importantly, I don't think this necessarily needs to become a dependency of Graphify. A first prototype could simply consume the existing graph.json, solve the selection problem externally, and map the selected nodes back to their source fragments.

Would you consider this useful for Graphify? If so, I would be happy to experiment with a small prototype. I am also curious whether you think this fits better as part of Graphify itself or as an external integration.