fix(settings): update configs for smooth transition between Gemini and OpenRouter

Author: magurh

src/data/rag_answer.json

@@ -1,4 +1,4 @@
 {
     "query": "What is the block time for the Flare blockchain?",
-    "answer": "Based on the provided text, a new block is generated on the Flare blockchain approximately every 1.8 seconds [Document 0].\n"
+    "answer": "The Flare blockchain produces a block approximately every 1.8 seconds [Document 0].\n"
 }

src/flare_ai_rag/ai/model.py

@@ -4,5 +4,5 @@
 @dataclass(frozen=True)
 class Model:
     model_id: str
-    max_tokens: int
-    temperature: float
+    max_tokens: int | None
+    temperature: float | None

src/flare_ai_rag/input_parameters.json

@@ -1,19 +1,15 @@
 {
     "router_model": {
-        "id": "qwen/qwen-vl-plus:free",
-        "max_tokens": 50,
-        "temperature": 0
+        "id": "gemini-1.5-flash"
     },
-    "qdrant_config": {
-        "embedding_model": "all-MiniLM-L6-v2",
-        "collection_name": "docs_collection",
+    "retriever_config": {
+        "embedding_model": "text-embedding-004",
         "vector_size": 768,
+        "collection_name": "docs_collection",
         "host": "localhost",
         "port": 6333
     },
     "responder_model": {
-        "id": "deepseek/deepseek-chat:free",
-        "max_tokens": 200,
-        "temperature": 0
+        "id": "gemini-1.5-flash"
     }
 }

src/flare_ai_rag/main.py

@@ -4,24 +4,14 @@
 
 from flare_ai_rag.ai import GeminiEmbedding, GeminiProvider
 from flare_ai_rag.responder import GeminiResponder, ResponderConfig
-from flare_ai_rag.retriever import QdrantConfig, QdrantRetriever, generate_collection
+from flare_ai_rag.retriever import QdrantRetriever, RetrieverConfig, generate_collection
 from flare_ai_rag.router import GeminiRouter, RouterConfig
 from flare_ai_rag.settings import settings
 from flare_ai_rag.utils import load_json, load_txt, save_json
 
 logger = structlog.get_logger(__name__)
 
 
-def setup_qdrant(input_config: dict) -> QdrantClient:
-    """Initialize Qdrant client."""
-    logger.info("Setting up Qdrant client...")
-    qdrant_config = QdrantConfig.load(input_config["qdrant_config"])
-    qdrant_client = QdrantClient(host=qdrant_config.host, port=qdrant_config.port)
-    logger.info("Qdrant client has been set up.")
-
-    return qdrant_client
-
-
 def setup_router(input_config: dict) -> GeminiRouter:
     """Initialize the Gemini Provider and the Gemini Router."""
     # Setup router config
@@ -31,7 +21,7 @@ def setup_router(input_config: dict) -> GeminiRouter:
     # Setup Gemini client based on Router config
     gemini_provider = GeminiProvider(
         api_key=settings.gemini_api_key,
-        model=settings.gemini_model,
+        model=router_config.model.model_id,
         system_instruction=router_config.system_prompt,
     )
 
@@ -46,7 +36,7 @@ def setup_retriever(
 ) -> QdrantRetriever:
     """Initialize the Qdrant retriever."""
     # Set up Qdrant config
-    qdrant_config = QdrantConfig.load(input_config["qdrant_config"])
+    retriever_config = RetrieverConfig.load(input_config["retriever_config"])
 
     # Set up Gemini Embedding client
     embedding_client = GeminiEmbedding(settings.gemini_api_key)
@@ -55,7 +45,7 @@ def setup_retriever(
         generate_collection(
             df_docs,
             qdrant_client,
-            qdrant_config,
+            retriever_config,
             collection_name=collection_name,
             embedding_client=embedding_client,
         )
@@ -65,11 +55,21 @@ def setup_retriever(
     # Return retriever
     return QdrantRetriever(
         client=qdrant_client,
-        qdrant_config=qdrant_config,
+        retriever_config=retriever_config,
         embedding_client=embedding_client,
     )
 
 
+def setup_qdrant(input_config: dict) -> QdrantClient:
+    """Initialize Qdrant client."""
+    logger.info("Setting up Qdrant client...")
+    retriever_config = RetrieverConfig.load(input_config["retriever_config"])
+    qdrant_client = QdrantClient(host=retriever_config.host, port=retriever_config.port)
+    logger.info("Qdrant client has been set up.")
+
+    return qdrant_client
+
+
 def setup_responder(input_config: dict) -> GeminiResponder:
     """Initialize the responder."""
     # Set up Responder Config.
@@ -79,7 +79,7 @@ def setup_responder(input_config: dict) -> GeminiResponder:
     # Set up a new Gemini Provider based on Responder Config.
     gemini_provider = GeminiProvider(
         api_key=settings.gemini_api_key,
-        model=settings.gemini_model,
+        model=responder_config.model.model_id,
         system_instruction=responder_config.system_prompt,
     )
     return GeminiResponder(client=gemini_provider, responder_config=responder_config)

src/flare_ai_rag/responder/config.py

@@ -1,28 +1,24 @@
 from dataclasses import dataclass
+from typing import Any
 
 from flare_ai_rag.ai import Model
 from flare_ai_rag.responder.prompts import RESPONDER_INSTRUCTION, RESPONDER_PROMPT
 
 
 @dataclass(frozen=True)
 class ResponderConfig:
-    model: Model | None
+    model: Model
     system_prompt: str
     query_prompt: str
 
     @staticmethod
-    def load(model_config: dict | None = None) -> "ResponderConfig":
+    def load(model_config: dict[str, Any]) -> "ResponderConfig":
         """Loads the Responder config."""
-        if not model_config:
-            # When using Gemini
-            model = None
-        else:
-            # When using OpenRouter
-            model = Model(
-                model_id=model_config["id"],
-                max_tokens=model_config["max_tokens"],
-                temperature=model_config["temperature"],
-            )
+        model = Model(
+            model_id=model_config["id"],
+            max_tokens=model_config.get("max_tokens"),
+            temperature=model_config.get("temperature"),
+        )
 
         return ResponderConfig(
             model=model,

src/flare_ai_rag/responder/responder.py

@@ -81,15 +81,16 @@ def generate_response(self, query: str, retrieved_documents: list[dict]) -> str:
         prompt = context + f"User query: {query}\n" + self.responder_config.query_prompt
         # Prepare the payload for the completion endpoint.
         payload: dict[str, Any] = {
+            "model": self.responder_config.model.model_id,
             "messages": [
                 {"role": "system", "content": self.responder_config.system_prompt},
                 {"role": "user", "content": prompt},
-            ]
+            ],
         }
 
-        if self.responder_config.model is not None:
-            payload["model"] = self.responder_config.model.model_id
+        if self.responder_config.model.max_tokens is not None:
             payload["max_tokens"] = self.responder_config.model.max_tokens
+        if self.responder_config.model.temperature is not None:
             payload["temperature"] = self.responder_config.model.temperature
 
         # Send the prompt to the OpenRouter API.

src/flare_ai_rag/retriever/__init__.py

@@ -1,6 +1,6 @@
 from .base import BaseRetriever
-from .config import QdrantConfig
+from .config import RetrieverConfig
 from .qdrant_collection import generate_collection
 from .qdrant_retriever import QdrantRetriever
 
-__all__ = ["BaseRetriever", "QdrantConfig", "QdrantRetriever", "generate_collection"]
+__all__ = ["BaseRetriever", "QdrantRetriever", "RetrieverConfig", "generate_collection"]

src/flare_ai_rag/retriever/config.py

@@ -1,8 +1,9 @@
 from dataclasses import dataclass
+from typing import Any
 
 
 @dataclass(frozen=True)
-class QdrantConfig:
+class RetrieverConfig:
     """Configuration for the embedding model used in the retriever."""
 
     embedding_model: str
@@ -12,8 +13,8 @@ class QdrantConfig:
     port: int
 
     @staticmethod
-    def load(retriever_config: dict) -> "QdrantConfig":
-        return QdrantConfig(
+    def load(retriever_config: dict[str, Any]) -> "RetrieverConfig":
+        return RetrieverConfig(
             embedding_model=retriever_config["embedding_model"],
             collection_name=retriever_config["collection_name"],
             vector_size=retriever_config["vector_size"],

src/flare_ai_rag/retriever/qdrant_collection.py

@@ -4,8 +4,7 @@
 from qdrant_client.http.models import Distance, PointStruct, VectorParams
 
 from flare_ai_rag.ai import GeminiEmbedding
-from flare_ai_rag.retriever.config import QdrantConfig
-from flare_ai_rag.settings import settings
+from flare_ai_rag.retriever.config import RetrieverConfig
 
 logger = structlog.get_logger(__name__)
 
@@ -28,13 +27,13 @@ def _create_collection(
 def generate_collection(
     df_docs: pd.DataFrame,
     qdrant_client: QdrantClient,
-    qdrant_config: QdrantConfig,
+    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, qdrant_config.vector_size)
+    _create_collection(qdrant_client, collection_name, retriever_config.vector_size)
     logger.info("Created the collection.", collection_name=collection_name)
 
     # For each document in the CSV, compute its embedding and prepare a Qdrant point.
@@ -54,7 +53,7 @@ def generate_collection(
         try:
             # Compute the embedding for the document content.
             embedding = embedding_client.embed_content(
-                embedding_model=settings.gemini_embedding_model, contents=content
+                embedding_model=retriever_config.embedding_model, contents=content
             )
         except Exception as e:
             logger.exception(

src/flare_ai_rag/retriever/qdrant_retriever.py

@@ -4,19 +4,19 @@
 
 from flare_ai_rag.ai import GeminiEmbedding
 from flare_ai_rag.retriever.base import BaseRetriever
-from flare_ai_rag.retriever.config import QdrantConfig
+from flare_ai_rag.retriever.config import RetrieverConfig
 
 
 class QdrantRetriever(BaseRetriever):
     def __init__(
         self,
         client: QdrantClient,
-        qdrant_config: QdrantConfig,
+        retriever_config: RetrieverConfig,
         embedding_client: GeminiEmbedding,
     ) -> None:
         """Initialize the QdrantRetriever."""
         self.client = client
-        self.qdrant_config = qdrant_config
+        self.retriever_config = retriever_config
         self.embedding_client = embedding_client
 
     @override
@@ -36,7 +36,7 @@ def semantic_search(self, query: str, top_k: int = 5) -> list[dict]:
 
         # Search Qdrant for similar vectors.
         results = self.client.search(
-            collection_name=self.qdrant_config.collection_name,
+            collection_name=self.retriever_config.collection_name,
             query_vector=query_vector,
             limit=top_k,
         )