Update rag_engine.py

rag_engine.py

@@ -5,9 +5,10 @@ import json
 import time
 import faiss
 import numpy as np
-from sentence_transformers import SentenceTransformer
 import requests
 from dotenv import load_dotenv
+from sentence_transformers import SentenceTransformer
+import streamlit as st
 
 # Load environment variables
 load_dotenv()
@@ -33,7 +34,7 @@ class RAGEngine:
         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:
+                with open(os.path.join(data_folder, file_name), 'r', encoding='utf-8') as f:
                     data = json.load(f)
                     docs.extend(self.flatten_data(data))
         return docs
@@ -42,13 +43,29 @@ class RAGEngine:
         flattened = []
         if isinstance(data, list):
             for item in data:
-                text = json.dumps(item, ensure_ascii=False)
-                flattened.append({"text": text})
+                flattened.extend(self.extract_fields(item))
         elif isinstance(data, dict):
-            text = json.dumps(data, ensure_ascii=False)
-            flattened.append({"text": text})
+            flattened.extend(self.extract_fields(data))
         return flattened
 
+    def extract_fields(self, item):
+        """Smart chunking: extract key fields instead of dumping full JSON."""
+        chunks = []
+        if isinstance(item, dict):
+            for key, value in item.items():
+                if isinstance(value, (str, int, float)):
+                    text = f"{key}: {value}"
+                    chunks.append({"text": text})
+                elif isinstance(value, dict):
+                    for sub_key, sub_value in value.items():
+                        text = f"{key} -> {sub_key}: {sub_value}"
+                        chunks.append({"text": text})
+                elif isinstance(value, list):
+                    for idx, sub_item in enumerate(value):
+                        text = f"{key}[{idx}]: {sub_item}"
+                        chunks.append({"text": text})
+        return chunks
+
     def create_vector_store(self):
         embeddings = []
         for doc in self.documents:
@@ -66,7 +83,6 @@ class RAGEngine:
         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"
@@ -75,7 +91,7 @@ class RAGEngine:
             "Format your answer like this:\n"
             "- Best Light Cones (Top 3)\n"
             "- Recommended Relic Sets and Main Stats\n"
-            - Important Substats to Prioritize\n"
+            "- Important Substats to Prioritize\n"
             "- Optimal Eidolon Level (if necessary)\n"
             "- Best Team Compositions (Synergies and Playstyle)\n"
             "- Any Special Notes\n\n"
@@ -92,7 +108,6 @@ class RAGEngine:
             "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:
@@ -101,11 +116,18 @@ class RAGEngine:
             else:
                 print(f"Request failed (attempt {attempt+1}/{retries}): {response.status_code}")
                 if attempt < retries - 1:
-                    time.sleep(wait_time)  # Wait then retry
+                    time.sleep(wait_time)
         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
+        return answer
+
+    def stream_answer(self, query):
+        """Streamed generation for Streamlit."""
+        answer = self.answer_query(query)
+        for word in answer.split():
+            yield word + " "
+            time.sleep(0.02)  # Feel free to tweak typing speed