diff --git a/agent/indexing/vector_store.py b/agent/indexing/vector_store.py index 41b615a..d180095 100644 --- a/agent/indexing/vector_store.py +++ b/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] \ No newline at end of file + return unique_docs[:k] \ No newline at end of file diff --git a/tests/test_generated.py b/tests/test_generated.py index a96f787..48db62f 100644 --- a/tests/test_generated.py +++ b/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() != "" \ No newline at end of file +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 \ No newline at end of file