Enhancement: Smarter RAG Retrieval with Top-K and Code Filtering

[deoxyforge]

This update improves the retrieval mechanism in the vector store to provide more relevant and meaningful context to the LLM.

Key Improvements:

• Enabled Top-K semantic retrieval instead of single result
• Normalized and cleaned retrieved documents
• Added filtering logic to prioritize code blocks (functions/classes)
• Improved overall signal-to-noise ratio in retrieved context

Result:

• More accurate and context-aware code reviews
• Better downstream performance for test generation

This builds on Phase 2 (ChromaDB integration) by improving retrieval quality.

[pleasingsunlight]

seems good to me!

[github-actions]

Automated Code Review Comment:

Bugs

• The line if doc and len(doc.strip()) is incomplete and should be if doc and len(doc.strip()) > 0 to remove empty strings.

Improvements

• The function query_embeddings now filters out empty or tiny chunks and prioritizes code with functions or classes, improving the relevance of the retrieved results.

Suggestions

• Consider adding error handling for the case where results does not contain the key "documents" or is empty.
• The magic number 20 in the line if doc and len(doc.strip()) > 20 could be replaced with a named constant for better readability.
• The function query_embeddings now prints the number of retrieved chunks, but it might be more useful to log this information instead of printing it directly.

---

Suggested Tests

import pytest
from agent.indexing.vector_store import query_embeddings

def test_query_embeddings_empty():
    assert query_embeddings("") == []

def test_query_embeddings_k_zero():
    assert query_embeddings("test", k=0) == []

def test_query_embeddings_k_large():
    assert len(query_embeddings("test", k=10)) <= 10

def test_query_embeddings_no_useful_code():
    assert query_embeddings("test") == []

def test_query_embeddings_useful_code():
    assert "def " in query_embeddings("def test")[0] or "class " in query_embeddings("class test")[0]

[haddybhaiya]

Congrats 👏