fix(main): move qdrant collection_name to retriever_config

Author: magurh

src/flare_ai_rag/main.py

@@ -32,7 +32,6 @@ def setup_retriever(
     qdrant_client: QdrantClient,
     input_config: dict,
     df_docs: pd.DataFrame,
-    collection_name: str | None = None,
 ) -> QdrantRetriever:
     """Initialize the Qdrant retriever."""
     # Set up Qdrant config
@@ -41,17 +40,16 @@ def setup_retriever(
     # Set up Gemini Embedding client
     embedding_client = GeminiEmbedding(settings.gemini_api_key)
     # (Re)generate qdrant collection
-    if collection_name:
-        generate_collection(
-            df_docs,
-            qdrant_client,
-            retriever_config,
-            collection_name=collection_name,
-            embedding_client=embedding_client,
-        )
-        logger.info(
-            "The Qdrant collection has been generated.", collection_name=collection_name
-        )
+    generate_collection(
+        df_docs,
+        qdrant_client,
+        retriever_config,
+        embedding_client=embedding_client,
+    )
+    logger.info(
+        "The Qdrant collection has been generated.",
+        collection_name=retriever_config.collection_name,
+    )
     # Return retriever
     return QdrantRetriever(
         client=qdrant_client,
@@ -100,9 +98,7 @@ def main() -> None:
     qdrant_client = setup_qdrant(input_config)
 
     # Set up retriever. (Use Gemini Embedding.)
-    retriever = setup_retriever(
-        qdrant_client, input_config, df_docs, collection_name="docs_collection"
-    )
+    retriever = setup_retriever(qdrant_client, input_config, df_docs)
 
     # Set up responder. (Use Gemini Provider.)
     responder = setup_responder(input_config)

src/flare_ai_rag/responder/prompts.py

@@ -6,7 +6,8 @@
 
 Guidelines:
 - Use the provided context to support your answer. If applicable,
-include citations referring to the context (e.g., "[Document <name>]" or "[Source <name>]").
+include citations referring to the context (e.g., "[Document <name>]" or
+"[Source <name>]").
 - Be clear, factual, and concise. Do not introduce any information that isn't
 explicitly supported by the context.
 - Maintain a professional tone and ensure that all technical details are accurate.

src/flare_ai_rag/retriever/qdrant_collection.py

@@ -28,13 +28,16 @@ def generate_collection(
     df_docs: pd.DataFrame,
     qdrant_client: QdrantClient,
     retriever_config: RetrieverConfig,
-    collection_name: str,
     embedding_client: GeminiEmbedding,
 ) -> None:
     """Routine for generating a Qdrant collection for a specific CSV file type."""
     # Create the collection.
-    _create_collection(qdrant_client, collection_name, retriever_config.vector_size)
-    logger.info("Created the collection.", collection_name=collection_name)
+    _create_collection(
+        qdrant_client, retriever_config.collection_name, retriever_config.vector_size
+    )
+    logger.info(
+        "Created the collection.", collection_name=retriever_config.collection_name
+    )
 
     # For each document in the CSV, compute its embedding and prepare a Qdrant point.
     points = []
@@ -74,10 +77,12 @@ def generate_collection(
 
     if points:
         # Upload the points into the Qdrant collection.
-        qdrant_client.upsert(collection_name=collection_name, points=points)
+        qdrant_client.upsert(
+            collection_name=retriever_config.collection_name, points=points
+        )
         logger.info(
             "Collection generated and documents inserted into Qdrant successfully.",
-            collection_name=collection_name,
+            collection_name=retriever_config.collection_name,
             num_points=len(points),
         )
     else:

tests/test_generate_collection.py

@@ -14,7 +14,7 @@
 def main() -> None:
     # Load Qdrant config
     config_json = load_json(settings.input_path / "input_parameters.json")
-    retriever_config = RetrieverConfig.load(config_json["qdrant_config"])
+    retriever_config = RetrieverConfig.load(config_json["retriever_config"])
 
     # Load the CSV file.
     df_docs = pd.read_csv(settings.data_path / "docs.csv", delimiter=",")
@@ -30,7 +30,6 @@ def main() -> None:
         df_docs,
         client,
         retriever_config,
-        collection_name="docs_collection",
         embedding_client=embedding_client,
     )
 

tests/test_qdrant_retriever.py

@@ -12,7 +12,7 @@
 def main() -> None:
     # Load Qdrant config
     config_json = load_json(settings.input_path / "input_parameters.json")
-    retriever_config = RetrieverConfig.load(config_json["qdrant_config"])
+    retriever_config = RetrieverConfig.load(config_json["retriever_config"])
 
     # Initialize Qdrant client
     client = QdrantClient(host=retriever_config.host, port=retriever_config.port)