Back to blog
Search & Retrieval

Optimize LLM Context: The Power of Dependency-Aware Code Graphs

5 min read
Summary

TL;DR Feeding large, undifferentiated code to LLMs like Claude Code inflates token costs and dilutes response quality. Implement dependency-aware code gra...

TL;DR

  • Feeding large, undifferentiated code to LLMs like Claude Code inflates token costs and dilutes response quality.
  • Implement dependency-aware code graph traversal to retrieve only essential, structurally relevant context for engineering tasks.

The Cost of Undifferentiated Context

Engineering teams increasingly leverage large language models (LLMs) for tasks like debugging, refactoring, and code generation. A common, costly pitfall is providing the LLM with an excessive, undifferentiated context window. Engineers often default to sending entire files, modules, or even large sections of a repository, hoping the LLM will discern relevance. This "dump and pray" strategy is inefficient and counterproductive.

Consider a request: "Fix a bug in the calculateTax() function." A naive approach might send the entire TaxService.java file, which could be thousands of lines, including unrelated methods, helper functions, and data models. The LLM must then sift through this noise to find the pertinent information.

This approach fails for several critical reasons:

  • Token Bloat and Cost Escalation: Every token sent to the LLM incurs a cost. Unnecessary context directly translates to higher API bills.
  • Context Dilution: Irrelevant code snippets compete for the LLM's attention, diminishing the signal-to-noise ratio. This can lead to less accurate or less specific outputs.
  • Performance Degradation: Larger context windows require more processing time, increasing latency for responses.
  • Increased Hallucination Risk: When overwhelmed with noise, LLMs are more prone to fabricating details or misinterpreting the actual problem scope.

Limitations of Basic Retrieval-Augmented Generation (RAG)

Many current LLM integrations employ a form of Retrieval-Augmented Generation (RAG), where user queries are embedded, and a vector database is searched for semantically similar code snippets or documents. While effective for broad information retrieval, basic RAG falls short in the context of precise engineering tasks.

The core issue is that semantic similarity does not equate to structural dependency. A function A might be semantically distinct from a type definition T in a different module, but A cannot compile or function correctly without T. A purely semantic search might miss this crucial dependency, leading to an incomplete context.

For instance, debugging a type error in a function might require the definition of a specific interface or class from a different file that shares no direct semantic keywords with the function itself. Standard RAG, operating on embedding similarity, would likely fail to retrieve this structurally vital, yet semantically distant, piece of code. This leads to either over-fetching (retrieving too much unrelated similar code) or under-fetching (missing critical dependencies).

Architectural Alternative: Dependency-Aware Context Graph

A durable alternative involves constructing and leveraging a dependency-aware code graph. This approach moves beyond keyword matching or semantic similarity to a deep structural understanding of the codebase.

The core idea is to parse the entire codebase into an Abstract Syntax Tree (AST) for each file, and then build a comprehensive graph that maps functions, classes, variables, and their relationships across the entire project. This graph represents:

  • Call Graph: Which functions call which other functions.
  • Data Flow Graph: How data moves between different parts of the code.
  • Type Dependencies: How types are defined, inherited, and used across the codebase.
  • Module Imports/Exports: The relationships between different modules and packages.

When a user submits a query targeting a specific code entity (e.g., "refactor OrderService.processOrder"), the system does not perform a broad search. Instead, it identifies the target entity in the code graph and then intelligently traverses the graph to collect only its direct and transitive dependencies.

Implementing Fine-Grained Context Extraction

Implementing this architecture requires static analysis tooling and a robust graph traversal algorithm:

  1. Codebase Parsing: Utilize language-specific parsers (e.g., Tree-sitter for multiple languages, ANTLR, or built-in compiler APIs) to generate ASTs for all source files.
  2. Graph Construction: From the ASTs, extract symbols (functions, classes, variables) and their relationships. This forms the directed dependency graph , where are code entities and are dependencies. For example, if function calls , there is an edge . If class extends , there is an edge .
  3. Query Analysis: Identify the specific target entity (e.g., a function, a class, a variable) from the user's request. This might involve a preliminary semantic search to map natural language to code symbols.
  4. Context Expansion Algorithm:
    • Initialize the context set with the code of .
    • Perform a graph traversal (e.g., Breadth-First Search or Depth-First Search) starting from .
    • At each step, include directly dependent entities:
      • Functions called by (callees).
      • Functions that call (callers, up to a configurable depth).
      • Type definitions of parameters, return types, and local variables used within .
      • Inherited classes or implemented interfaces.
      • Relevant module imports.
    • Limit traversal depth to prevent context bloat while maintaining sufficiency. A typical depth of 2-3 hops often suffices for localized tasks.
    • Prioritize local definitions over global or widely used library code unless the query specifically implicates them.
  5. Context Serialization: Convert the identified minimal set of code snippets back into a textual representation, ensuring it's syntactically valid and complete for the LLM.

This process ensures the LLM receives only the critical code snippets required to understand the problem and generate an accurate solution. For example, UserService.createUser(User user) would include the createUser method's body, the User class definition, the UserRepository.save method signature, and the EmailService.sendWelcomeEmail method signature, but not unrelated methods within UserService or EmailService.

The primary trade-off is the initial investment in tooling and maintaining the code graph. This requires robust static analysis, potentially integrating with CI/CD pipelines to keep the graph up-to-date with every code change. However, the long-term benefits in reduced LLM costs, faster response times, and significantly improved output accuracy far outweigh this initial setup. This architecture transforms context provision from a speculative exercise into a precise, deterministic operation.

Invest in tools that understand your codebase's structure, not just its words. By shifting from broad context provision to targeted, dependency-aware retrieval, engineering teams can unlock the full potential of LLMs while maintaining cost efficiency and architectural stability. The future of LLM-assisted engineering lies in precise context.

Back to all posts