feat: add responder and create RAG routine

Author: magurh

.env.dummy

@@ -0,0 +1,3 @@
+# OpenRouter base url and API key
+OPENROUTER_BASE_URL="https://openrouter.ai/api/v1"
+OPENROUTER_API_KEY=""

src/flare_ai_rag/input_parameters.json

@@ -1,7 +1,7 @@
 {
     "router_model": {
         "id": "qwen/qwen-vl-plus:free",
-        "max_tokens": 20,
+        "max_tokens": 5,
         "temperature": 0
     },
     "qdrant_config": {
@@ -10,5 +10,10 @@
         "vector_size": 384,
         "host": "localhost",
         "port": 6333
+    },
+    "responder_model": {
+        "id": "deepseek/deepseek-chat:free",
+        "max_tokens": 200,
+        "temperature": 0
     }
 }

src/flare_ai_rag/main.py

@@ -0,0 +1,109 @@
+import pandas as pd
+import structlog
+from qdrant_client import QdrantClient
+
+from flare_ai_rag.config import config
+from flare_ai_rag.openrouter.client import OpenRouterClient
+from flare_ai_rag.responder.config import ResponderConfig
+from flare_ai_rag.responder.responder import OpenRouterResponder
+from flare_ai_rag.retriever.config import QdrantConfig
+from flare_ai_rag.retriever.qdrant_collection import generate_collection
+from flare_ai_rag.retriever.qdrant_retriever import QdrantRetriever
+from flare_ai_rag.router.config import RouterConfig
+from flare_ai_rag.router.router import QueryRouter
+from flare_ai_rag.utils import loader
+
+logger = structlog.get_logger(__name__)
+
+
+def setup_clients(input_config: dict) -> tuple[OpenRouterClient, QdrantClient]:
+    """Initialize OpenRouter and Qdrant clients."""
+    # Setup OpenRouter client.
+    openrouter_client = OpenRouterClient(
+        api_key=config.open_router_api_key, base_url=config.open_router_base_url
+    )
+
+    # Setup Qdrant client.
+    qdrant_config = QdrantConfig.load(input_config["qdrant_config"])
+    qdrant_client = QdrantClient(host=qdrant_config.host, port=qdrant_config.port)
+
+    return openrouter_client, qdrant_client
+
+
+def setup_router(
+    openrouter_client: OpenRouterClient, input_config: dict
+) -> QueryRouter:
+    """Initialize the query router."""
+    router_model_config = input_config["router_model"]
+    router_config = RouterConfig.load(router_model_config)
+    return QueryRouter(client=openrouter_client, config=router_config)
+
+
+def setup_responder(
+    openrouter_client: OpenRouterClient, input_config: dict
+) -> OpenRouterResponder:
+    """Initialize the responder."""
+    responder_config = input_config["responder_model"]
+    responder_config = ResponderConfig.load(responder_config)
+    return OpenRouterResponder(
+        client=openrouter_client, responder_config=responder_config
+    )
+
+
+def setup_retriever(
+    qdrant_client: QdrantClient,
+    input_config: dict,
+    df_docs: pd.DataFrame,
+    collection: str | None = None,
+) -> QdrantRetriever:
+    """Initialize the Qdrant retriever."""
+    qdrant_config = QdrantConfig.load(input_config["qdrant_config"])
+
+    # (Re)generate qdrant collection
+    if collection:
+        generate_collection(
+            df_docs, qdrant_client, qdrant_config, collection_name=collection
+        )
+    # Return retriever
+    return QdrantRetriever(client=qdrant_client, qdrant_config=qdrant_config)
+
+
+def main() -> None:
+    # Load input configuration.
+    input_config = loader.load_json(config.input_path / "input_parameters.json")
+
+    # Setup clients.
+    openrouter_client, qdrant_client = setup_clients(input_config)
+
+    # Setup the router.
+    router = setup_router(openrouter_client, input_config)
+
+    # Process user query.
+    query = loader.load_txt(config.input_path / "query.txt")
+    classification = router.route_query(query)
+    logger.info("Queried classified.", classification=classification)
+
+    if classification == "ANSWER":
+        df_docs = pd.read_csv(config.data_path / "docs.csv", delimiter=",")
+        logger.info("Loaded CSV Data.", num_rows=len(df_docs))
+
+        # Retrieve docs
+        retriever = setup_retriever(
+            qdrant_client, input_config, df_docs, collection="docs_collection"
+        )
+        retrieved_docs = retriever.semantic_search(query, top_k=5)
+
+        # Prepare answer
+        responder = setup_responder(openrouter_client, input_config)
+        answer = responder.generate_response(query, retrieved_docs)
+        logger.info("Answer retrieved.", answer=answer)
+    elif classification == "CLARIFY":
+        logger.info("Your query needs clarification. Please provide more details.")
+    elif classification == "REJECT":
+        logger.info("Your query has been rejected as it is out of scope.")
+    else:
+        logger.info("Unexpected classification.", classification=classification)
+
+
+if __name__ == "__main__":
+    main()

src/flare_ai_rag/query.txt

@@ -0,0 +1,2 @@
+What is the block time for the Flare blockchain?
+

src/flare_ai_rag/responder/base_responder.py

@@ -0,0 +1,9 @@
+from abc import ABC, abstractmethod
+
+
+class BaseResponder(ABC):
+    @abstractmethod
+    def generate_response(self, query: str, retrieved_documents: list[dict]) -> str:
+        """
+        Generate a final answer given the query and a list of retrieved documents.
+        """

src/flare_ai_rag/responder/config.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+
+from flare_ai_rag.config import config
+from flare_ai_rag.openrouter.model import Model
+from flare_ai_rag.utils import loader
+
+# Load base prompt
+BASE_PROMPT = loader.load_txt(config.input_path / "responder" / "prompts.txt")
+
+
+@dataclass(frozen=True)
+class ResponderConfig:
+    model: Model
+    base_prompt: str
+
+    @staticmethod
+    def load(model_config: dict) -> "ResponderConfig":
+        """Loads the Responder config."""
+        model = Model(
+            model_id=model_config["id"],
+            max_tokens=model_config["max_tokens"],
+            temperature=model_config["temperature"],
+        )
+
+        return ResponderConfig(model=model, base_prompt=BASE_PROMPT)

src/flare_ai_rag/responder/prompts.txt

@@ -0,0 +1,8 @@
+Answer the following question using the provided context.
+Include citations for supporting evidence in your answer.
+
+Question: {query}
+
+Context: {context}
+
+Answer:

src/flare_ai_rag/responder/responder.py

@@ -0,0 +1,51 @@
+from typing import override
+
+from flare_ai_rag.openrouter.client import OpenRouterClient
+from flare_ai_rag.responder.base_responder import BaseResponder
+from flare_ai_rag.responder.config import ResponderConfig
+from flare_ai_rag.utils import parser
+
+
+class OpenRouterResponder(BaseResponder):
+    def __init__(
+        self, client: OpenRouterClient, responder_config: ResponderConfig
+    ) -> None:
+        """
+        Initialize the responder with an OpenRouter client and the model to use.
+
+        :param client: An instance of OpenRouterClient.
+        :param model: The model identifier to be used by the API.
+        """
+        self.client = client
+        self.responder_config = responder_config
+
+    @override
+    def generate_response(self, query: str, retrieved_documents: list[dict]) -> str:
+        """
+        Generate a final answer using the query and the retrieved context,
+        and include citations.
+
+        :param query: The input query.
+        :param retrieved_documents: A list of dictionaries containing retrieved docs.
+        :return: The generated answer as a string.
+        """
+        context = ""
+
+        # Build context from the retrieved documents.
+        for idx, doc in enumerate(retrieved_documents, start=1):
+            identifier = doc.get("metadata", {}).get("filename", f"Doc{idx}")
+            context += f"Document {identifier}:\n{doc.get('text', '')}\n\n"
+
+        # Compose the prompt
+        prompt = self.responder_config.base_prompt.format(query=query, context=context)
+        # Prepare the payload for the completion endpoint.
+        payload = {
+            "model": self.responder_config.model.model_id,
+            "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": self.responder_config.model.max_tokens,
+            "temperature": self.responder_config.model.temperature,
+        }
+        # Send the prompt to the OpenRouter API.
+        response = self.client.send_chat_completion(payload)
+
+        return parser.parse_openrouter_response(response)

src/flare_ai_rag/retriever/qdrant_collection.py

@@ -1,5 +1,12 @@
+import pandas as pd
+import structlog
 from qdrant_client import QdrantClient
-from qdrant_client.http.models import Distance, VectorParams
+from qdrant_client.http.models import Distance, PointStruct, VectorParams
+from sentence_transformers import SentenceTransformer
+
+from flare_ai_rag.retriever.config import QdrantConfig
+
+logger = structlog.get_logger(__name__)
 
 
 def create_collection(
@@ -15,3 +22,63 @@ def create_collection(
         collection_name=collection_name,
         vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE),
     )
+
+
+def generate_collection(
+    df_docs: pd.DataFrame,
+    client: QdrantClient,
+    qdrant_config: QdrantConfig,
+    collection_name: str,
+) -> None:
+    """Routine for generating a Qdrant collection for a specific CSV file type."""
+    # Create the collection.
+    create_collection(client, collection_name, qdrant_config.vector_size)
+    logger.info("Created the collection.", collection_name=collection_name)
+
+    # Load the embedding model.
+    embedding_model = SentenceTransformer(qdrant_config.embedding_model)
+
+    # For each document in the CSV, compute its embedding and prepare a Qdrant point.
+    points = []
+    for i, row in df_docs.iterrows():
+        doc_id = str(i)
+        content = row["Contents"]
+
+        # Check if content is missing or not a string.
+        if not isinstance(content, str):
+            logger.warning(
+                "Skipping document due to missing or invalid content.",
+                filename=row["Filename"],
+            )
+            continue
+
+        try:
+            # Compute the embedding for the document content.
+            embedding = embedding_model.encode(content).tolist()
+        except Exception as e:
+            logger.exception(
+                "Error encoding document.", filename=row["Filename"], error=str(e)
+            )
+            continue
+
+        # Prepare the payload.
+        payload = {
+            "filename": row["Filename"],
+            "metadata": row["Metadata"],
+            "text": content,
+        }
+
+        # Create a Qdrant point.
+        point = PointStruct(id=doc_id, vector=embedding, payload=payload)
+        points.append(point)
+
+    if points:
+        # Upload the points into the Qdrant collection.
+        client.upsert(collection_name=collection_name, points=points)
+        logger.info(
+            "Collection generated and documents inserted into Qdrant successfully.",
+            collection_name=collection_name,
+            num_points=len(points),
+        )
+    else:
+        logger.warning("No valid documents found to insert.")

src/flare_ai_rag/retriever/qdrant_retriever.py

@@ -31,7 +31,7 @@ def semantic_search(self, query: str, top_k: int = 5) -> list[dict]:
         """
         # Convert the query into a vector embedding using the
         # SentenceTransformer instance.
-        query_vector = self.embedding_model.encode(query)
+        query_vector = self.embedding_model.encode(query).tolist()
 
         # Search Qdrant for similar vectors.
         results = self.client.search(