Refactored

app.py

@@ -169,24 +169,78 @@ class KnowledgeBase:
         with self.lock:
             if country in self.vstores:
                 return
-        texts = []
-        if self.domain == "banking": texts = self.collector.get_banking_data(country)
-        elif self.domain == "credit": texts = self.collector.get_credit_data(country)
+        # Retrieve raw texts
+        if self.domain == "banking":
+            texts = self.collector.get_banking_data(country)
+        elif self.domain == "credit":
+            texts = self.collector.get_credit_data(country)
         elif self.domain == "tax":
             ti = self.collector.tax_db.get_tax_information(country)
-            texts = ti["regulations"] + ti["treaty"]
-        # else: other domains...
-        splits = self.splitter.split_text("\n\n".join(texts))
+            texts = ti.get("regulations", []) + ti.get("treaty", [])
+        else:
+            texts = []
+
+        if not texts:
+            log_workflow(f"No texts available for domain '{self.domain}' and country '{country}'")
+            with self.lock:
+                self.vstores[country] = None
+                self.retrievers[country] = None
+            return
+
+        # Split texts into chunks
+                # Split texts into chunks
+        splits = self.splitter.split_text("".join(texts))
+        if not splits:
+            log_workflow(f"No splits generated for domain '{self.domain}' and country '{country}'")
+            with self.lock:
+                self.vstores[country] = None
+                self.retrievers[country] = None
+            return
+
+        # Build vector store
         store = Chroma.from_texts(splits, self.embeddings, collection_name=f"{self.domain}_{country}")
+        retr = store.as_retriever(search_kwargs={"k": 3})
         with self.lock:
             self.vstores[country] = store
-            self.retrievers[country] = store.as_retriever(search_kwargs={"k":3})
+            self.retrievers[country] = retr
+        log_workflow(f"Vector store ready for domain '{self.domain}' and country '{country}'")
 
     def retrieve(self, query: str, country: str) -> List[str]:
-        log_workflow(f"Retrieving {self.domain} for {country}")
+        log_workflow(f"Retrieving domain '{self.domain}' for {country}")
         self._init_country(country)
-        docs = self.retrievers[country].get_relevant_documents(query)
-        return [d.page_content for d in docs]
+        retr = self.retrievers.get(country)
+        if not retr:
+            # Fallback to direct collector methods
+            log_workflow(f"Falling back to direct retrieval for domain '{self.domain}' and country '{country}'")
+            if self.domain == "banking":
+                return self.collector.get_banking_data(country)
+            if self.domain == "credit":
+                return self.collector.get_credit_data(country)
+            if self.domain == "tax":
+                info = self.collector.tax_db.get_tax_information(country)
+                return info.get("regulations", []) + info.get("treaty", [])
+            return []
+
+        # Perform similarity search
+        docs = retr.get_relevant_documents(query)
+        results = [d.page_content for d in docs]
+        log_workflow(f"Retrieved {len(results)} docs for domain '{self.domain}' and country '{country}'")
+        return results
+
+# Pre-initialize KnowledgeBase for common domains and countries
+COMMON_COUNTRIES = ["India", "China", "South Korea", "Brazil", "Saudi Arabia", "Canada", "Mexico", "Taiwan", "Japan", "Vietnam"]
+DOMAINS = ["banking", "credit", "tax"]
+
+# Initialize and preload vector stores at startup
+def preload_kbs():
+    for domain in DOMAINS:
+        kb = KnowledgeBase(domain)
+        for country in COMMON_COUNTRIES:
+            # Launch in background to avoid blocking
+            threading.Thread(target=kb._init_country, args=(country,), daemon=True).start()
+
+# Trigger preload
+preload_kbs()
 
 # =======================================
 # Specialist Agents