Use Mistral embeddings to match original index

app.py

@@ -1,5 +1,6 @@
 import os
 import gradio as gr
+import requests
 from huggingface_hub import InferenceClient
 from langchain_core.embeddings import Embeddings
 from langchain_pinecone import PineconeVectorStore
@@ -13,36 +14,33 @@ from pinecone import Pinecone
 DEFAULT_MODEL = "Qwen/Qwen2.5-72B-Instruct"
 
 
-class HFEmbeddings(Embeddings):
-    """Custom embeddings class using HuggingFace Inference API"""
+class MistralEmbeddings(Embeddings):
+    """Mistral embeddings to match the original index"""
     
-    # Using BGE-large which produces 1024-dimensional embeddings to match Pinecone index
-    def __init__(self, model_name: str = "BAAI/bge-large-en-v1.5"):
-        self.model_name = model_name
-        self.client = InferenceClient(token=os.getenv("HF_TOKEN"))
+    def __init__(self):
+        self.api_key = os.getenv("MISTRAL_API_KEY")
+        self.model = "mistral-embed"
+        self.url = "https://api.mistral.ai/v1/embeddings"
     
-    def _to_float_list(self, result) -> list[float]:
-        """Convert numpy arrays or nested lists to plain Python floats"""
-        import numpy as np
-        if hasattr(result, 'tolist'):
-            return result.tolist()
-        if isinstance(result, (list, tuple)):
-            # Handle nested structure - flatten if needed
-            if len(result) > 0 and isinstance(result[0], (list, tuple, np.ndarray)):
-                result = result[0]
-            return [float(x) for x in result]
-        return [float(x) for x in result]
+    def _get_embeddings(self, texts: list[str]) -> list[list[float]]:
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json"
+        }
+        data = {
+            "model": self.model,
+            "input": texts
+        }
+        response = requests.post(self.url, headers=headers, json=data)
+        response.raise_for_status()
+        result = response.json()
+        return [item["embedding"] for item in result["data"]]
     
     def embed_documents(self, texts: list[str]) -> list[list[float]]:
-        embeddings = []
-        for text in texts:
-            result = self.client.feature_extraction(text, model=self.model_name)
-            embeddings.append(self._to_float_list(result))
-        return embeddings
+        return self._get_embeddings(texts)
     
     def embed_query(self, text: str) -> list[float]:
-        result = self.client.feature_extraction(text, model=self.model_name)
-        return self._to_float_list(result)
+        return self._get_embeddings([text])[0]
 
 
 class ResearchParrot:
@@ -54,7 +52,7 @@ class ResearchParrot:
 
     def embeddings(self):
         if self._embeddings is None:
-            self._embeddings = HFEmbeddings()
+            self._embeddings = MistralEmbeddings()
         return self._embeddings
 
     def vectorstore(self):

requirements.txt

@@ -1,6 +1,7 @@
 gradio>=4.0.0
 huggingface_hub>=0.20.0
 langchain>=0.1.0
-langchain-community>=0.2.0
+langchain-core>=0.1.0
 langchain-pinecone>=0.1.0
 pinecone-client>=3.0.0
+requests>=2.28.0