Implement collection management in Milvus: drop collection during cleanup and on startup if no documents exist. Update embedding generation to clarify data types and improve error handling. Refactor vector store collection creation logic to avoid unnecessary drops.

app.py

@@ -38,6 +38,13 @@ def cleanup_documents():
     for f in files:
         if os.path.isfile(f):
             os.remove(f)
+    # Also drop the Milvus collection to avoid stale state between restarts
+    try:
+        if milvus_client and milvus_client.has_collection(COLLECTION_NAME):
+            milvus_client.drop_collection(COLLECTION_NAME)
+            print(f"Dropped collection {COLLECTION_NAME} during cleanup.")
+    except Exception as e:
+        print(f"Error dropping collection during cleanup: {e}")
     print("Cleanup complete.")
 
 
@@ -45,6 +52,23 @@ def cleanup_documents():
 atexit.register(cleanup_documents)
 
 
+def reset_collection_if_no_docs():
+    """Drop existing collection on startup if there are no documents on disk."""
+    try:
+        os.makedirs(DOCS_DIR, exist_ok=True)
+        files = glob.glob(os.path.join(DOCS_DIR, "*"))
+        has_docs = any(os.path.isfile(f) for f in files)
+        if (
+            not has_docs
+            and milvus_client
+            and milvus_client.has_collection(COLLECTION_NAME)
+        ):
+            milvus_client.drop_collection(COLLECTION_NAME)
+            print(f"No documents found. Dropped existing collection {COLLECTION_NAME}.")
+    except Exception as e:
+        print(f"Error resetting collection on startup: {e}")
+
+
 def index_documents(file_list):
     """Index documents from a list of files."""
     if not file_list:
@@ -113,4 +137,6 @@ with gr.Blocks() as demo:
 if __name__ == "__main__":
     # Ensure the documents directory exists from the start
     os.makedirs(DOCS_DIR, exist_ok=True)
+    # Reset collection state if there are no documents at startup
+    reset_collection_if_no_docs()
     demo.launch()

src/embedding_generator.py

@@ -32,7 +32,7 @@ def generate_document_embeddings(
     Returns:
         A list of document embeddings
     """
-    binary_embeddings = []
+    binary_embeddings: list[bytes] = []
 
     try:
         for context in batch_iterate(documents, batch_size=512):
@@ -43,10 +43,10 @@ def generate_document_embeddings(
             embeds_array = np.array(batch_embeddings)
             binary_embeds = np.where(embeds_array > 0, 1, 0).astype(np.uint8)
 
-            # convert to bytes array
+            # convert to bytes per vector
             packed_embeds = np.packbits(binary_embeds, axis=1)
-
-            binary_embeddings.extend(packed_embeds)
+            for row in packed_embeds:
+                binary_embeddings.append(row.tobytes())
         return binary_embeddings
     except Exception as e:
         print(f"Error generating document embeddings: {e}")

src/rag_pipeline.py

@@ -1,6 +1,9 @@
+from dotenv import load_dotenv
 from langchain.chat_models import init_chat_model
 from langchain_core.messages import HumanMessage
 
+load_dotenv(override=True)
+
 from .config import MODEL_NAME, MODEL_PROVIDER, PROMPT, TEMPERATURE
 
 llm = init_chat_model(

src/vector_store.py

@@ -32,50 +32,47 @@ def create_collection_if_not_exists(
         dim: The dimension of the binary vector
     """
     try:
+        # Create collection only if it does not exist
+        if not client.has_collection(collection_name):
+            print(f"Collection {collection_name} not found. Creating it...")
 
-        # Drop collection if it exists
-        if client.has_collection(collection_name):
-            print(f"Collection {collection_name} exists, dropping it...")
-            client.drop_collection(collection_name)
-
-            # Initialize client
             schema = client.create_schema(
                 auto_id=True,
                 enable_dynamic_fields=True,
             )
-            # Add primary key field
             schema.add_field(
                 field_name="id",
                 datatype=DataType.INT64,
                 is_primary=True,
                 auto_id=True,
             )
-            # Add fields to schema
             schema.add_field(
                 field_name="context",
                 datatype=DataType.VARCHAR,
-                max_length=65535,  # max length for VARCHAR
+                max_length=65535,
             )
             schema.add_field(
                 field_name="binary_vector",
                 datatype=DataType.BINARY_VECTOR,
                 dim=dim,
             )
-            # Create index params for binary vector
+
             index_params = client.prepare_index_params()
             index_params.add_index(
                 field_name="binary_vector",
                 index_name="binary_vector_index",
-                index_type="BIN_FLAT",  # Exact search for binary vectors
-                metric_type="HAMMING",  # Hamming distance for binary vectors
+                index_type="BIN_FLAT",
+                metric_type="HAMMING",
             )
-            # Create collection with schema and index
+
             client.create_collection(
                 collection_name=collection_name,
                 schema=schema,
                 index_params=index_params,
             )
             print(f"Collection {collection_name} created successfully.")
+        else:
+            print(f"Collection {collection_name} already exists. Skipping creation.")
     except Exception as e:
         print(f"Error creating collection: {e}")
         return None