Update : clean and optimized retrieval from Chromadb

agent/indexing/vector_store.py

@@ -1,7 +1,9 @@
 import chromadb
 
+# safer collection creation
 client = chromadb.Client()
-collection = client.create_collection(name="codebase")
+collection = client.get_or_create_collection(name="codebase")
+
 
 def store_embeddings(chunks, embeddings):
     for i, chunk in enumerate(chunks):
@@ -11,23 +13,48 @@ def store_embeddings(chunks, embeddings):
             ids=[str(i)]
         )
 
+
+MIN_CHUNK_LENGTH = 30  # Phase 4 threshold
+
+
 def query_embeddings(query_embedding, k=5):
     results = collection.query(
         query_embeddings=[query_embedding],
         n_results=k
     )
 
-    docs = results.get("documents", [[]])[0]
+    docs = results.get("documents", [[]])
+
+    # safety check
+    if not docs or not docs[0]:
+        print("No documents retrieved")
+        return []
+
+    docs = docs[0]
+
+    # 🧹 CLEAN: remove empty / noisy chunks
+    cleaned = [
+        doc for doc in docs
+        if doc and len(doc.strip()) > MIN_CHUNK_LENGTH
+    ]
+
+    # 🎯 PRIORITIZE: functions/classes
+    filtered = [
+        doc for doc in cleaned
+        if "def " in doc or "class " in doc
+    ]
+
+    # fallback if nothing found
+    if not filtered:
+        filtered = cleaned
 
-    # remove empty / tiny chunks
-    cleaned = [doc for doc in docs if doc and len(doc.strip()) > 20]
+    # 🔁 REMOVE DUPLICATES
+    unique_docs = list(dict.fromkeys(filtered))
 
-    # 🔥 prioritize useful code (functions/classes)
-    filtered = []
-    for doc in cleaned:
-        if "def " in doc or "class " in doc:
-            filtered.append(doc)
+    # 🧪 DEBUG LOGS
+    print(f"Retrieved {len(unique_docs)} relevant chunks")
 
-    print(f"Retrieved {len(filtered)} relevant chunks")
+    for i, doc in enumerate(unique_docs[:2]):
+        print(f"Chunk {i+1} preview:", doc[:100])
 
-    return filtered[:k]
+    return unique_docs[:k]

tests/test_generated.py

@@ -1,28 +1,34 @@
-def test_infer_intent_empty_diff():
-    assert infer_intent("") == {"purpose": "", "properties": [], "edge_cases": [], "error": ""}
+import pytest
+from your_module import store_embeddings, query_embeddings, client
 
-def test_infer_intent_short_diff():
-    diff = "diff --git a/file.py b/file.py"
-    intent = infer_intent(diff)
-    assert "purpose" in intent
-    assert "properties" in intent
-    assert "edge_cases" in intent
+def test_store_embeddings():
+    chunks = ["chunk1", "chunk2"]
+    embeddings = [[1, 2], [3, 4]]
+    store_embeddings(chunks, embeddings)
 
-def test_infer_intent_long_diff():
-    diff = "diff --git a/file.py b/file.py" * 10000
-    intent = infer_intent(diff)
-    assert "purpose" in intent
-    assert "properties" in intent
-    assert "edge_cases" in intent
+def test_query_embeddings():
+    query_embedding = [1, 2]
+    results = query_embeddings(query_embedding)
+    assert isinstance(results, list)
 
-def test_generate_tests_no_context_no_intent():
-    diff = "diff --git a/file.py b/file.py"
-    tests = generate_tests(diff)
-    assert tests.strip() != ""
+def test_query_embeddings_empty():
+    collection = client.get_or_create_collection(name="empty")
+    query_embedding = [1, 2]
+    results = collection.query(query_embeddings=[query_embedding], n_results=5)
+    assert results.get("documents", []) == [[]]
 
-def test_generate_tests_with_context_and_intent():
-    diff = "diff --git a/file.py b/file.py"
-    context = ["context1", "context2"]
-    intent = {"purpose": "test", "properties": ["prop1", "prop2"], "edge_cases": ["case1", "case2"]}
-    tests = generate_tests(diff, context, intent)
-    assert tests.strip() != ""
+def test_query_embeddings_min_length():
+    query_embedding = [1, 2]
+    chunks = ["a" * 29, "b" * 31]
+    embeddings = [[1, 2], [3, 4]]
+    store_embeddings(chunks, embeddings)
+    results = query_embeddings(query_embedding)
+    assert len(results) == 1
+
+def test_query_embeddings_k():
+    query_embedding = [1, 2]
+    chunks = ["a" * 31, "b" * 31, "c" * 31]
+    embeddings = [[1, 2], [3, 4], [5, 6]]
+    store_embeddings(chunks, embeddings)
+    results = query_embeddings(query_embedding, k=2)
+    assert len(results) == 2