Update rag_engine.py

rag_engine.py

@@ -0,0 +1,111 @@
+# rag_engine.py
+
+import os
+import json
+import time
+import faiss
+import numpy as np
+from sentence_transformers import SentenceTransformer
+import requests
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+class RAGEngine:
+    def __init__(self):
+        # Load model for embedding
+        self.embedder = SentenceTransformer('all-MiniLM-L6-v2')
+        self.embedding_dim = 384
+        self.index = faiss.IndexFlatL2(self.embedding_dim)
+        self.texts = []
+
+        # Load documents
+        self.documents = self.load_documents()
+        self.create_vector_store()
+
+        # Hugging Face API details
+        self.api_token = os.getenv("HF_API_TOKEN")
+        self.model_url = "https://api-inference.huggingface.co/models/deepseek-ai/deepseek-llm-7b-instruct"
+
+    def load_documents(self):
+        docs = []
+        data_folder = "data/"
+        for file_name in os.listdir(data_folder):
+            if file_name.endswith(".json"):
+                with open(os.path.join(data_folder, 'r', encoding='utf-8') as f:
+                    data = json.load(f)
+                    docs.extend(self.flatten_data(data))
+        return docs
+
+    def flatten_data(self, data):
+        flattened = []
+        if isinstance(data, list):
+            for item in data:
+                text = json.dumps(item, ensure_ascii=False)
+                flattened.append({"text": text})
+        elif isinstance(data, dict):
+            text = json.dumps(data, ensure_ascii=False)
+            flattened.append({"text": text})
+        return flattened
+
+    def create_vector_store(self):
+        embeddings = []
+        for doc in self.documents:
+            emb = self.embedder.encode(doc['text'])
+            embeddings.append(emb)
+            self.texts.append(doc['text'])
+        embeddings = np.array(embeddings)
+        self.index.add(embeddings)
+
+    def search_documents(self, query, top_k=5):
+        query_emb = self.embedder.encode(query)
+        query_emb = np.expand_dims(query_emb, axis=0)
+        distances, indices = self.index.search(query_emb, top_k)
+        results = [self.texts[i] for i in indices[0] if i < len(self.texts)]
+        return results
+
+    def ask_deepseek(self, context, query, retries=3, wait_time=5):
+        # 🔥 More detailed prompt
+        prompt = (
+            "You are an expert Honkai Star Rail Build Advisor.\n"
+            "You specialize in optimizing character performance based on Light Cones, Relics, Stats, Eidolons, and Team Synergies.\n"
+            "Provide detailed build advice for the given query using the provided context.\n"
+            "Always prioritize the most effective and meta-relevant recommendations.\n\n"
+            "Format your answer like this:\n"
+            "- Best Light Cones (Top 3)\n"
+            "- Recommended Relic Sets and Main Stats\n"
+            - Important Substats to Prioritize\n"
+            "- Optimal Eidolon Level (if necessary)\n"
+            "- Best Team Compositions (Synergies and Playstyle)\n"
+            "- Any Special Notes\n\n"
+            f"Context:\n{context}\n\n"
+            f"Question:\n{query}\n"
+            "Answer:"
+        )
+        headers = {
+            "Authorization": f"Bearer {self.api_token}",
+            "Content-Type": "application/json"
+        }
+        payload = {
+            "inputs": prompt,
+            "parameters": {"temperature": 0.7, "max_new_tokens": 800}
+        }
+
+        # 🚀 Retry logic
+        for attempt in range(retries):
+            response = requests.post(self.model_url, headers=headers, json=payload)
+            if response.status_code == 200:
+                generated_text = response.json()[0]["generated_text"]
+                return generated_text.split("Answer:")[-1].strip()
+            else:
+                print(f"Request failed (attempt {attempt+1}/{retries}): {response.status_code}")
+                if attempt < retries - 1:
+                    time.sleep(wait_time)  # Wait then retry
+        return f"Error: Could not get a valid response after {retries} attempts."
+
+    def answer_query(self, query):
+        relevant_docs = self.search_documents(query)
+        context = "\n".join(relevant_docs)
+        answer = self.ask_deepseek(context, query)
+        return answer