Update rag/data.py

rag/data.py

@@ -1,60 +1,18 @@
 import json
-from dataclasses import dataclass
-from pathlib import Path
-from typing import List, Optional
-
+from typing import List, Dict, Any
 from rag.logging_utils import get_logger
 
 logger = get_logger(__name__)
 
-@dataclass
-class Doc:
-    doc_name: str
-    image_path: str
-    text: str
-    company: str  # "Apple" | "Microsoft" | "Unknown"
-
-def infer_company(doc_name: str) -> str:
-    dn = (doc_name or "").lower()
-    if "microsoft" in dn or "msft" in dn:
-        return "Microsoft"
-    if "apple" in dn or "aapl" in dn:
-        return "Apple"
-    return "Unknown"
-
-def load_dataset(dataset_path: str) -> List[Doc]:
-    path = Path(dataset_path)
-    if not path.exists():
-        logger.warning("Dataset file not found: %s", dataset_path)
+def load_dataset(path: str) -> List[Dict[str, Any]]:
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        if not isinstance(data, list):
+            logger.warning("Dataset JSON is not a list. Found: %s", type(data))
+            return []
+        logger.info("Loaded dataset: %d docs", len(data))
+        return data
+    except Exception as e:
+        logger.warning("⚠️ Dataset not found/invalid (%s): %s", path, e)
         return []
-
-    with path.open("r", encoding="utf-8") as f:
-        raw = json.load(f)
-
-    docs: List[Doc] = []
-    for item in raw:
-        doc_name = item.get("doc_name", "Unknown Document")
-        image_path = item.get("image_path", "")
-        text = (item.get("text") or "").strip()
-        company = item.get("company") or infer_company(doc_name)
-
-        if not text:
-            continue
-
-        docs.append(
-            Doc(
-                doc_name=doc_name,
-                image_path=image_path,
-                text=text,
-                company=company,
-            )
-        )
-
-    logger.info("Loaded %d docs", len(docs))
-    return docs
-
-def filter_docs_by_companies(docs: List[Doc], companies: Optional[List[str]]) -> List[Doc]:
-    if not companies:
-        return docs
-    allowed = set(companies)
-    return [d for d in docs if d.company in allowed]