# How to Perform AI-Powered Code Optimization with Code-Graph-RAG: A Complete Guide

> Learn AI-powered code optimization with Code-Graph-RAG. Discover how this powerful tool analyzes and optimizes code using LLMs and property graphs.

- Repository: [Vitali Avagyan/code-graph-rag](https://github.com/vitali87/code-graph-rag)
- Tags: how-to-guide
- Published: 2026-09-06

---

**Code-Graph-RAG combines deterministic property graphs with semantic embeddings to find, analyze, and automatically optimize code using LLMs.**

AI-powered code optimization requires more than feeding raw files to a language model. The vitali87/code-graph-rag repository solves this with a **two-layer architecture**: a deterministic graph that maps exact symbols and relationships, plus a vector store that surfaces semantically similar patterns. This guide walks through the complete workflow—from graph construction to automated patch application—using the actual source code implementation as implemented in [vitali87/code-graph-rag](https://github.com/vitali87/code-graph-rag).

## The Two-Layer Architecture for AI Code Optimization

Code-Graph-RAG separates deterministic structure from semantic similarity. This separation ensures reproducible results while enabling AI-driven insights.

| Layer | Responsibility | Key File |
|-------|--------------|----------|
| **Property Graph** | Maps functions, classes, and their relationships with sorted, reproducible queries | [`codebase_rag/graph_query.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/graph_query.py) |
| **Vector Store** | Caches semantic embeddings for similarity search across code fragments | [`codebase_rag/embedder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/embedder.py), [`codebase_rag/vector_store.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/vector_store.py) |

The graph layer uses Cypher-style queries to resolve exact symbols. The embedding layer retrieves analogous patterns for LLM context. Together, they feed a structured optimization prompt that produces reviewable, automatically applicable patches.

## Step 1: Build and Load the Code Graph

Graph construction starts with tree-sitter parsers that handle C, C++, C#, Go, Java, JavaScript, Python, Rust, and more. The [`codebase_rag/graph_updater.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/graph_updater.py) module extracts nodes (functions, classes, modules) and edges (calls, imports, references), then serializes to JSON.

Load an existing graph with `GraphLoader`:

```python
from codebase_rag.graph_loader import load_graph

graph = load_graph("/path/to/my_repo.graph.json")

```

The loader in [`codebase_rag/graph_loader.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/graph_loader.py) indexes nodes by `id`, `label`, and property for fast lookups without reloading the full structure on every query.

## Step 2: Locate Your Optimization Target

Use deterministic queries in [`codebase_rag/graph_query.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/graph_query.py) to pinpoint exact symbols. The `resolve` function finds functions by qualified name, while `callers` and `callees` map dependency relationships:

```python
from codebase_rag.graph_query import resolve

# Find the exact function node to optimize

func_rows = resolve(graph.load, project_name="my_repo", target="utils.cleanup")
target = func_rows[0]          # deterministic, sorted result

source = target["source"]      # raw source code string

```

All query results are **sorted for reproducibility**—critical for CI/CD integration where identical inputs must yield identical outputs.

## Step 3: Retrieve Semantically Similar Patterns

The embedding layer in [`codebase_rag/embedder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/embedder.py) manages caching and model selection. The `_cache_namespace` ensures embeddings are isolated by provider and model, preventing cross-contamination:

```python
from codebase_rag.embedder import get_embedding_cache

cache = get_embedding_cache()
embedding = cache.get(source) or cache.put(source, cache.get(source))

```

Query the vector store (backed by Qdrant) for nearest neighbors:

```python
from codebase_rag.vector_store import vector_store

similar = vector_store.nearest(embedding, k=5)

```

These similar snippets provide the LLM with proven patterns from your own codebase, reducing hallucinated optimizations.

## Step 4: Construct the Optimization Prompt

Combine the target function with retrieved examples into a structured prompt. The embedding layer supports OpenAI and UnixCoder models; configure your provider in [`codebase_rag/config.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/config.py):

```python
prompt = f"""\
You are an expert Python engineer. Optimize the following function for speed
and readability while preserving its behaviour. Also, keep the same
public API.

Original:
{source}

Similar patterns (for inspiration):
{chr(10).join([s['source'] for s in similar])}
"""

```

Send to your LLM through the project's client wrapper:

```python
from codebase_rag.llm import ask_llm

optimised_code = ask_llm(prompt)

```

## Step 5: Apply the Patch Automatically

The editing pipeline in [`codebase_rag/editing/patcher.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/editing/patcher.py) and [`codebase_rag/editing/transaction.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/editing/transaction.py) translates LLM suggestions into concrete source changes. `PatchTransaction` manages atomicity, allowing rollback if verification fails:

```python
from codebase_rag.editing.patcher import apply_patch

apply_patch(
    file_path=target["path"],
    old_source=source,
    new_source=optimised_code,
    description="AI‑driven optimisation of utils.cleanup"
)

```

The patcher uses AST-aware diffs to minimize merge conflicts and preserve formatting outside the optimized region.

## Complete AI-Powered Optimization Workflow

Here's the full integration as implemented in [vitali87/code-graph-rag](https://github.com/vitali87/code-graph-rag):

```python
from codebase_rag.graph_loader import load_graph
from codebase_rag.graph_query import resolve
from codebase_rag.embedder import get_embedding_cache
from codebase_rag.vector_store import vector_store
from codebase_rag.llm import ask_llm
from codebase_rag.editing.patcher import apply_patch

# 1. Load the graph

graph = load_graph("/path/to/my_repo.graph.json")

# 2. Resolve target function

func_rows = resolve(graph.load, project_name="my_repo", target="utils.cleanup")
target = func_rows[0]
source = target["source"]

# 3. Get embedding (cached)

cache = get_embedding_cache()
embedding = cache.get(source) or cache.put(source, cache.get(source))

# 4. Find similar patterns

similar = vector_store.nearest(embedding, k=5)

# 5. Prompt LLM for optimization

prompt = f"""\
You are an expert Python engineer. Optimize the following function for speed
and readability while preserving its behaviour. Also, keep the same
public API.

Original:
{source}

Similar patterns (for inspiration):
{chr(10).join([s['source'] for s in similar])}
"""
optimised_code = ask_llm(prompt)

# 6. Apply patch

apply_patch(
    file_path=target["path"],
    old_source=source,
    new_source=optimised_code,
    description="AI‑driven optimisation of utils.cleanup"
)

```

## Why Deterministic Graphs Matter for CI/CD

**Deterministic outputs** distinguish Code-Graph-RAG from naive RAG approaches. The [`codebase_rag/graph_query.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/graph_query.py) module sorts all query results, and the embedding cache uses namespaced keys. Repeated runs with identical inputs produce identical retrieval sets—essential for:

- Reproducible optimization reports
- Automated regression testing of the optimization pipeline itself
- Safe integration into build systems without nondeterministic failures

## Key Files for AI-Powered Code Optimization

| Component | File | Purpose |
|-----------|------|---------|
| Graph construction | [`codebase_rag/graph_updater.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/graph_updater.py) | Parses source, extracts symbols, writes JSON graph |
| Graph access | [`codebase_rag/graph_loader.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/graph_loader.py) | Lazy loading, indexed lookups |
| Deterministic queries | [`codebase_rag/graph_query.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/graph_query.py) | `resolve`, `callers`, `callees`, `test_reachable` |
| Embeddings | [`codebase_rag/embedder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/embedder.py) | `EmbeddingCache`, provider selection, batching |
| Vector search | [`codebase_rag/vector_store.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/vector_store.py) | Qdrant interface for nearest-neighbor queries |
| Patch application | [`codebase_rag/editing/patcher.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/editing/patcher.py) | `apply_patch`, AST-aware diffs |
| Transaction safety | [`codebase_rag/editing/transaction.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/editing/transaction.py) | `PatchTransaction` for atomic changes |
| Configuration | [`codebase_rag/config.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/config.py) | Model providers, vector DB settings |
| CLI entry | [`main.py`](https://github.com/vitali87/code-graph-rag/blob/main/main.py) | Top-level orchestration |

## Summary

- **Code-Graph-RAG** enables AI-powered code optimization through a deterministic property graph paired with semantic embeddings
- The **graph layer** ([`graph_query.py`](https://github.com/vitali87/code-graph-rag/blob/main/graph_query.py)) provides exact symbol resolution with sorted, reproducible results
- The **embedding layer** ([`embedder.py`](https://github.com/vitali87/code-graph-rag/blob/main/embedder.py), [`vector_store.py`](https://github.com/vitali87/code-graph-rag/blob/main/vector_store.py)) retrieves semantically similar patterns for LLM context
- The **editing pipeline** ([`patcher.py`](https://github.com/vitali87/code-graph-rag/blob/main/patcher.py), [`transaction.py`](https://github.com/vitali87/code-graph-rag/blob/main/transaction.py)) converts LLM suggestions into verified, atomic source patches
- Cache namespacing and sorted queries ensure **deterministic, CI-safe** optimization workflows

## Frequently Asked Questions

### What programming languages does Code-Graph-RAG support?

Code-Graph-RAG supports C, C++, C#, Go, Java, JavaScript, Python, Rust, and additional languages through tree-sitter parsers as implemented in [`codebase_rag/graph_updater.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/graph_updater.py). The modular parser architecture allows extension to new languages by adding the corresponding tree-sitter grammar.

### How does the deterministic query engine ensure reproducible results?

The [`codebase_rag/graph_query.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/graph_query.py) module sorts all Cypher-style query results before returning them. Combined with the `_cache_namespace` isolation in [`codebase_rag/embedder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/embedder.py), this guarantees that identical inputs produce identical retrieval sets across repeated runs—critical for CI/CD integration.

### Can I use a different embedding model or vector database?

Yes. The [`codebase_rag/embedder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/embedder.py) module supports pluggable providers including OpenAI and UnixCoder through configuration in [`codebase_rag/config.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/config.py). The [`codebase_rag/vector_store.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/vector_store.py) interface abstracts the underlying database; Qdrant is the default implementation but compatible alternatives can be substituted.

### What happens if the LLM generates invalid or broken code?

The [`codebase_rag/editing/patcher.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/editing/patcher.py) module parses the LLM response into structured `Patch` objects, and [`codebase_rag/editing/transaction.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/editing/transaction.py) wraps changes in `PatchTransaction` for atomic application. Failed patches can be rejected or rolled back before reaching your codebase, with the `description` field tracking each optimization attempt for audit purposes.