chore: replace SentenceTransformer embeddings by Gemini embeddings

Author: magurh

pyproject.toml

@@ -6,6 +6,9 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "cryptography>=44.0.1",
+    "google>=3.0.0",
+    "google-genai>=1.2.0",
+    "google-generativeai>=0.8.4",
     "httpx>=0.28.1",
     "openrouter>=1.0",
     "pandas>=2.2.3",

src/data/rag_answer.json

@@ -1,4 +1,4 @@
 {
     "query": "What is the block time for the Flare blockchain?",
-    "answer": "The block time for the Flare blockchain is approximately **1.8 seconds**. This is evidenced by the context provided in **Document 2-feeds.mdx**, which states: \"FTSOv2's block-latency feeds update incrementally with each new block on Flare, approximately every 1.8 seconds\" (Document 2-feeds.mdx). This indicates that new blocks are produced on the Flare blockchain roughly every 1.8 seconds."
+    "answer": "The block time for the Flare blockchain is approximately **1.8 seconds**. This is evidenced by the context provided in Document 0-overview.mdx, which states that the Flare Time Series Oracle (FTSOv2) features block-latency feeds that update \"with each new block on Flare, every \u22481.8 seconds\" [Document 0-overview.mdx]. This indicates that the time interval between consecutive blocks on the Flare blockchain is roughly 1.8 seconds."
 }

src/flare_ai_rag/ai/__init__.py

@@ -0,0 +1,6 @@
+from .base_client import AsyncBaseClient, BaseClient
+from .gemini import GeminiClient
+from .model import Model
+from .openrouter import OpenRouterClient
+
+__all__ = ["AsyncBaseClient", "BaseClient", "GeminiClient", "Model", "OpenRouterClient"]

…c/flare_ai_rag/openrouter/base_client.py src/flare_ai_rag/ai/base_client.pysrc/flare_ai_rag/openrouter/base_client.py renamed to src/flare_ai_rag/ai/base_client.py src/flare_ai_rag/openrouter/base_client.py renamed to src/flare_ai_rag/ai/base_client.py

@@ -0,0 +1,15 @@
+from google import genai
+
+
+class GeminiClient:
+    def __init__(self, api_key: str) -> None:
+        self.client = genai.Client(api_key=api_key)
+
+    def embed_content(self, model: str, contents: str) -> list[float]:
+        result = self.client.models.embed_content(model=model, contents=contents)
+        embedding = result.embeddings
+
+        if not embedding or embedding[0].values is None:
+            msg = "No embedding was returned from the API."
+            raise ValueError(msg)
+        return embedding[0].values

src/flare_ai_rag/ai/gemini.py

@@ -1,4 +1,4 @@
-from flare_ai_rag.openrouter.base_client import AsyncBaseClient, BaseClient
+from flare_ai_rag.ai import AsyncBaseClient, BaseClient
 
 
 class OpenRouterClient(BaseClient):

src/flare_ai_rag/openrouter/model.py src/flare_ai_rag/ai/model.pysrc/flare_ai_rag/openrouter/model.py renamed to src/flare_ai_rag/ai/model.py

@@ -7,7 +7,7 @@
     "qdrant_config": {
         "embedding_model": "all-MiniLM-L6-v2",
         "collection_name": "docs_collection",
-        "vector_size": 384,
+        "vector_size": 768,
         "host": "localhost",
         "port": 6333
     },

src/flare_ai_rag/openrouter/client.py src/flare_ai_rag/ai/openrouter.pysrc/flare_ai_rag/openrouter/client.py renamed to src/flare_ai_rag/ai/openrouter.py

@@ -2,7 +2,7 @@
 import structlog
 from qdrant_client import QdrantClient
 
-from flare_ai_rag.openrouter import OpenRouterClient
+from flare_ai_rag.ai import GeminiClient, OpenRouterClient
 from flare_ai_rag.responder import OpenRouterResponder, ResponderConfig
 from flare_ai_rag.retriever import QdrantConfig, QdrantRetriever, generate_collection
 from flare_ai_rag.router import QueryRouter, RouterConfig
@@ -54,6 +54,7 @@ def setup_retriever(
     qdrant_client: QdrantClient,
     input_config: dict,
     df_docs: pd.DataFrame,
+    gemini_client: GeminiClient,
     collection_name: str | None = None,
 ) -> QdrantRetriever:
     """Initialize the Qdrant retriever."""
@@ -62,13 +63,19 @@ def setup_retriever(
     # (Re)generate qdrant collection
     if collection_name:
         generate_collection(
-            df_docs, qdrant_client, qdrant_config, collection_name=collection_name
+            df_docs,
+            qdrant_client,
+            qdrant_config,
+            collection_name=collection_name,
+            gemini_client=gemini_client,
         )
         logger.info(
             "The Qdrant collection has been generated.", collection_name=collection_name
         )
     # Return retriever
-    return QdrantRetriever(client=qdrant_client, qdrant_config=qdrant_config)
+    return QdrantRetriever(
+        client=qdrant_client, qdrant_config=qdrant_config, gemini_client=gemini_client
+    )
 
 
 def main() -> None:
@@ -93,11 +100,13 @@ def main() -> None:
         logger.info("Loaded CSV Data.", num_rows=len(df_docs))
 
         # Retrieve docs
+        gemini_client = GeminiClient(api_key=settings.gemini_api_key)
         retriever = setup_retriever(
             qdrant_client,
             input_config,
             df_docs,
             collection_name="docs_collection",
+            gemini_client=gemini_client,
         )
         retrieved_docs = retriever.semantic_search(query, top_k=5)
         logger.info("Docs have been retrieved.")