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

agent/indexing/vector_store.py

@@ -16,4 +16,18 @@ def query_embeddings(query_embedding, k=5):
         query_embeddings=[query_embedding],
         n_results=k
     )
-    return results["documents"][0]
+
+    docs = results.get("documents", [[]])[0]
+
+    # remove empty / tiny chunks
+    cleaned = [doc for doc in docs if doc and len(doc.strip()) > 20]
+
+    # 🔥 prioritize useful code (functions/classes)
+    filtered = []
+    for doc in cleaned:
+        if "def " in doc or "class " in doc:
+            filtered.append(doc)
+
+    print(f"Retrieved {len(filtered)} relevant chunks")
+
+    return filtered[:k]

agent/llm/test_generator.py

@@ -0,0 +1,40 @@
+from agent.llm.groq_client import client, model
+def generate_tests(diff:str,context:list =None)->str:
+    context_text = ""
+    if context:
+        context_text = "\n\n".join(context[:5])  # Include only the first 5 chunks for context
+    prompt = f"""
+    You are a senior software engineer.
+    Generate concise pytest test cases for the given code diff.
+
+    STRICT RULES:
+    - Output only code
+    - Max  5 test functions
+    - keep tests short
+    - Focus on edge cases and core logic
+    - NO explanations, NO extra text
+
+
+    ---CONTEXT---
+    {context_text}
+
+    ---DIFF---
+    {diff}
+    """
+
+
+    try:
+        response = client.chat.completions.create(
+            model = model,
+            messages = [
+
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ],
+          )
+        return response.choices[0].message.content.strip()
+    except Exception as e:
+        return f"Error generating tests:{str(e)}"
+

agent/main.py

@@ -2,6 +2,7 @@
 from agent.indexing.chunker import chunk_code
 from agent.indexing.embedder import get_embeddings
 from agent.indexing.vector_store import store_embeddings, query_embeddings
+from agent.llm.test_generator import generate_tests
 
 import subprocess
 from agent.llm.groq_client import generate_review
@@ -68,9 +69,19 @@ def main():
     print("Generating review...")
     review = generate_review(diff, context=relevant_chunks)
 
-    # 8. Post comment 
+    # 8. Generate tests
+    print("Generating tests...")
+    tests = generate_tests(diff, context=relevant_chunks)
+
+    # 9. Combine output
+    final_output = f"{review}\n\n---\n\n### Suggested Tests\n{tests}"
+
+    print("\n FINAL OUTPUT:\n")
+    print(final_output)
+
+    # 10. Post comment
     print("Posting comment...")
-    post_comment(review)
+    post_comment(final_output)
 
     print("Done")