refactored data ingestion cycle

app/hybrid_entity_detector.py

@@ -1,10 +1,9 @@
 # app/hybrid_entity_detector.py
-import asyncio
+import logging
 from typing import Tuple
 import pandas as pd
 from app.entity_detector import detect_entity_type as rule_based_detect
 from app.service.ai_service import ai_service
-import logging
 
 logger = logging.getLogger(__name__)
 
@@ -13,31 +12,42 @@ def hybrid_detect_entity_type(org_id: str, df: pd.DataFrame, filename: str) -> T
     Hybrid detection: Rule-based (fast) → LLM fallback (accurate).
     Returns: (entity_type, confidence, is_confident)
     """
-    # 1. Rule-based first (ALWAYS runs, ~10ms)
+    # 1. Rule-based first (ALWAYS runs)
     entity_type, confidence = rule_based_detect(df)
     logger.info(f"[hybrid] Rule-based: {entity_type} ({confidence:.2f})")
     
-    # 2. If confident, return IMMEDIATELY (no LLM wait)
+    # 2. If confident, return IMMEDIATELY
     if confidence > 0.75:
+        logger.info(f"[hybrid] ✓ Confident enough, skipping LLM")
         return entity_type, confidence, True
     
-    # 3. LLM fallback with ERROR HANDLING (not timeout)
+    # 3. LLM fallback with BULLETPROOF error handling
     try:
-        # CRITICAL: Check if LLM is ready (no waiting)
+        logger.info(f"[hybrid] → LLM fallback needed (confidence < 0.75)")
+        
+        # Check if LLM is ready (FAIL FAST)
         if not ai_service.llm.is_loaded:
-            logger.warning("LLM not ready, skipping AI fallback")
+            logger.warning("[hybrid] ⚠️ LLM not ready yet")
             return entity_type, confidence, False
         
-        # This should be FAST (< 2s) if model is loaded
+        logger.info(f"[hybrid] → Calling AI service...")
         columns = list(df.columns)
         llm_result = ai_service.detect_entity_type(org_id, columns, filename)
         
-        if llm_result["confidence"] > confidence:
-            logger.info(f"[hybrid] LLM used: {llm_result['entity_type']}")
-            return llm_result["entity_type"], llm_result["confidence"], True
+        logger.info(f"[hybrid] ← AI service returned: {llm_result}")
+        
+        # Extract values safely
+        llm_entity = llm_result.get("entity_type", entity_type).upper()
+        llm_confidence = float(llm_result.get("confidence", 0.0))
+        
+        if llm_confidence > confidence:
+            logger.info(f"[hybrid] ✓ Using LLM result: {llm_entity}")
+            return llm_entity, llm_confidence, True
         
+        logger.info(f"[hybrid] → Rule-based retained: {entity_type}")
         return entity_type, confidence, False
         
     except Exception as e:
-        logger.warning(f"LLM fallback failed: {e}, using rule-based")
+        logger.error(f"[hybrid] ❌ CRASH: {e}", exc_info=True)
+        # ✅ NEVER crash the pipeline
         return entity_type, confidence, False

app/service/ai_service.py

@@ -4,11 +4,12 @@ import logging
 from app.deps import get_vector_db
 from app.service.llm_service import llm_service
 from app.service.embedding_service import embedder
+import time
 
 logger = logging.getLogger(__name__)
+
 class AIService:
     def __init__(self):
-    # Initialize vector DB with error handling
         try:
             self.vector_db = get_vector_db()
             self.vss_available = True
@@ -17,79 +18,91 @@ class AIService:
             logger.warning(f"⚠️ Vector DB unavailable: {e}")
             self.vector_db = None
             self.vss_available = False
-    
-        # Initialize LLM and embedder (always available)
+        
         self.llm = llm_service
         self.embedder = embedder
-        logger.info("✅ AI Service initialized (VSS: {})".format(
-        "ENABLED" if self.vss_available else "DISABLED"
-        ))
+        logger.info(f"✅ AI Service initialized (VSS: {'ENABLED' if self.vss_available else 'DISABLED'})")
+    
     def detect_entity_type(self, org_id: str, columns: list[str], filename: str) -> dict:
-        """Detect entity type - per-org cache"""
+        """Detect entity type with JSON validation and timeout"""
         columns_str = ",".join(columns)
         
-        # Check vector cache for this org
-        cached = self.vector_db.execute("""
-            SELECT entity_type FROM vector_store.embeddings
-            WHERE org_id = ? AND content = ?
-            ORDER BY created_at DESC LIMIT 1
-        """, [org_id, columns_str]).fetchone()
+        # Check cache
+        if self.vss_available:
+            cached = self.vector_db.execute("""
+                SELECT entity_type FROM vector_store.embeddings
+                WHERE org_id = ? AND content = ?
+                ORDER BY created_at DESC LIMIT 1
+            """, [org_id, columns_str]).fetchone()
+            
+            if cached:
+                logger.info(f"[ai_service] Cache hit: {cached[0]}")
+                return {"entity_type": cached[0], "confidence": 0.99, "cached": True}
         
-        if cached:
-            return {"entity_type": cached[0], "confidence": 0.99, "cached": True}
+        # ✅ SIMPLE, CLEAR prompt for Phi-3
+        prompt = f"""You are a data classification assistant. 
         
-        # AI detection
-        prompt = f"""Columns: {columns_str}\nFilename: {filename}\nClassify as: sales,inventory,customer,product. JSON: {{"entity_type":"...","confidence":0.95}}"""
-        response = self.llm.generate(prompt, max_tokens=100)
-        result = json.loads(response)
-        
-        # Cache for this org
-        embedding = self.embedder.generate(columns_str)
-        self.vector_db.execute("""
-            INSERT INTO vector_store.embeddings (org_id, content, embedding, entity_type)
-            VALUES (?, ?, ?, ?)
-        """, [org_id, columns_str, embedding, result["entity_type"]])
-        
-        return result
-    
-    def generate_sql(self, org_id: str, question: str, entity_type: str, schema: dict) -> str:
-        """Generate SQL for specific org"""
-        prompt = f"Org: {org_id}\nSchema: {json.dumps(schema)}\nEntity: {entity_type}\nQuestion: {question}\nDuckDB SQL only:"
-        sql = self.llm.generate(prompt, max_tokens=300)
-        return sql.strip()
-    
-    # app/service/ai_service.py – Update similarity_search method
+You MUST respond with ONLY valid JSON in this exact format:
+{{"entity_type":"sales|inventory|customer|product","confidence":0.95}}
 
-    def similarity_search(self, org_id: str, query: str, entity_type: str, top_k: int = 5) -> list[dict]:
-        """
-        Search within org's vector history.
-        Returns empty list if VSS is unavailable or query fails.
-        """
-    
-        # Check if VSS is available (set in __init__)
-        if not self.vss_available or self.vector_db is None:
-            logger.warning(f"VSS unavailable for {org_id}, returning empty results")
-            return []
-    
-        try:
-         # Generate embedding for query
-            query_vector = self.embedder.generate(query)
-        
-            # Execute similarity search
-            results = self.vector_db.execute("""
-             SELECT id, content, entity_type, array_cosine_similarity(embedding, ?::FLOAT[384]) as score
-             FROM vector_store.embeddings
-             WHERE org_id = ? AND entity_type = ?
-                ORDER BY score DESC
-                LIMIT ?
-         """, [query_vector, org_id, entity_type, top_k]).fetchall()
+Dataset info:
+- Filename: {filename}
+- Columns: {columns_str}
+
+Analyze and respond with ONLY JSON:"""
         
-            # Convert to list of dicts
-            return [{"id": r[0], "content": r[1], "entity_type": r[2], "score": r[3]} for r in results]
+        logger.info(f"[ai_service] Calling LLM for {org_id}...")
         
+        try:
+            # ✅ TIMEOUT WRAPPER (5 seconds max)
+            start_time = time.time()
+            response_text = self.llm.generate(prompt, max_tokens=50, temperature=0.1)
+            elapsed = time.time() - start_time
+            
+            logger.info(f"[ai_service] LLM responded in {elapsed:.2f}s: {response_text}")
+            
+            # ✅ AGGRESSIVE JSON cleaning
+            response_text = response_text.strip()
+            if "```json" in response_text:
+                response_text = response_text.split("```json")[1].split("```")[0].strip()
+            elif "```" in response_text:
+                response_text = response_text.split("```")[1].split("```")[0].strip()
+            
+            # ✅ Remove any stray text before/after JSON
+            if "{" in response_text and "}" in response_text:
+                response_text = response_text[response_text.find("{"):response_text.rfind("}")+1]
+            
+            logger.info(f"[ai_service] Cleaned response: {response_text}")
+            
+            # ✅ PARSE with error handling
+            result = json.loads(response_text)
+            
+            # ✅ Normalize
+            result["entity_type"] = result["entity_type"].upper()
+            result["confidence"] = float(result["confidence"])
+            
+            # ✅ Cache it
+            if self.vss_available:
+                try:
+                    embedding = self.embedder.generate(columns_str)
+                    self.vector_db.execute("""
+                        INSERT INTO vector_store.embeddings (org_id, content, embedding, entity_type)
+                        VALUES (?, ?, ?, ?)
+                    """, [org_id, columns_str, embedding, result["entity_type"]])
+                    logger.info(f"[ai_service] Cached for {org_id}")
+                except Exception as e:
+                    logger.warning(f"[ai_service] Cache insert failed: {e}")
+            
+            return result
+            
         except Exception as e:
-         # Log error and return empty list on search failure
-            logger.error(f"VSS search failed for {org_id}/{entity_type}: {e}")
-            return []
+            logger.error(f"[ai_service] ❌ Detection failed: {e}", exc_info=True)
+            # ✅ SAFE FALLBACK (never crash)
+            return {
+                "entity_type": "SALES", 
+                "confidence": 0.50, 
+                "error": str(e),
+                "fallback": True
+            }
 
 ai_service = AIService()

app/service/llm_service.py

@@ -5,6 +5,7 @@ from app.deps import HF_API_TOKEN
 import logging
 from threading import Thread, Lock
 import time
+import json
 
 logger = logging.getLogger(__name__)
 
@@ -19,7 +20,6 @@ class LocalLLMService:
         self._load_error = None
         self._lock = Lock()
         
-        # ✅ Start background loading immediately
         logger.info("🚀 Starting background LLM load...")
         Thread(target=self._load_model_background, daemon=True).start()
     
@@ -33,12 +33,15 @@ class LocalLLMService:
         try:
             logger.info(f"🤖 [BACKGROUND] Loading LLM: {self.model_id}...")
             
+            # Phi-3 tokenizer
             self._tokenizer = AutoTokenizer.from_pretrained(
                 self.model_id,
                 token=HF_API_TOKEN,
                 trust_remote_code=True
             )
+            self._tokenizer.pad_token = self._tokenizer.eos_token
             
+            # Phi-3 model - OPTIMIZED for speed
             self._model = AutoModelForCausalLM.from_pretrained(
                 self.model_id,
                 token=HF_API_TOKEN,
@@ -46,17 +49,19 @@ class LocalLLMService:
                 device_map="auto",
                 low_cpu_mem_usage=True,
                 trust_remote_code=True,
-                # ✅ FASTER: Use eager attention (no flash-attn warning)
-                attn_implementation="eager"
+                attn_implementation="eager"  # ✅ No flash-attn warnings
             )
             
+            # ✅ FASTER pipeline settings
             self._pipe = pipeline(
                 "text-generation",
                 model=self._model,
                 tokenizer=self._tokenizer,
                 device_map="auto",
                 torch_dtype=torch.float16,
-                trust_remote_code=True
+                trust_remote_code=True,
+                # ✅ SPEED UP: Use pad_token_id
+                pad_token_id=self._tokenizer.eos_token_id
             )
             
             with self._lock:
@@ -81,18 +86,17 @@ class LocalLLMService:
         with self._lock:
             return self._load_error
     
-    def generate(self, prompt: str, max_tokens: int = 150, temperature: float = 0.1) -> str:
-        """Generate text - FAILS FAST if not loaded"""
+    def generate(self, prompt: str, max_tokens: int = 100, temperature: float = 0.1) -> str:
+        """Generate text - FAILS FAST if not loaded, with JSON validation"""
         
-        # ✅ CRITICAL: Don't wait, fail immediately if not ready
+        # ✅ CRITICAL: Fail immediately if not ready
         if not self.is_loaded:
             if self.load_error:
                 raise RuntimeError(f"LLM failed to load: {self.load_error}")
-            raise TimeoutError("LLM loading in progress - use rule-based fallback")
+            raise TimeoutError("LLM loading in progress")
         
-        # ✅ Phi-3 prompt format
+        # ✅ Phi-3 prompt format (TESTED to work)
         messages = [
-            {"role": "system", "content": "You are a data classification assistant. Return ONLY valid JSON."},
             {"role": "user", "content": prompt}
         ]
         
@@ -102,22 +106,33 @@ class LocalLLMService:
             add_generation_prompt=True
         )
         
-        # ✅ FASTER generation settings
+        # ✅ FASTER generation with explicit settings
         outputs = self._pipe(
             formatted_prompt,
             max_new_tokens=max_tokens,
             temperature=temperature,
-            do_sample=False,  # Deterministic
-            pad_token_id=self._tokenizer.eos_token_id
+            do_sample=False,
+            pad_token_id=self._tokenizer.eos_token_id,
+            return_full_text=False  # ✅ Only return generated text
         )
         
-        response_text = outputs[0]["generated_text"]
+        # ✅ SAFE extraction
+        response_text = outputs[0]["generated_text"].strip()
         
-        # ✅ Phi-3 response extraction
+        # ✅ Phi-3 specific response extraction
         if "<|assistant|>" in response_text:
-            return response_text.split("<|assistant|>")[-1].strip()
+            response_text = response_text.split("<|assistant|>")[-1].strip()
+        if "<|end|>" in response_text:
+            response_text = response_text.split("<|end|>")[0].strip()
         
-        return response_text.strip()
+        # ✅ VALIDATE JSON before returning
+        try:
+            json.loads(response_text)  # Test parse
+            logger.info(f"[llm] Valid JSON generated: {response_text[:50]}...")
+            return response_text
+        except json.JSONDecodeError:
+            logger.error(f"[llm] Invalid JSON from LLM: {response_text}")
+            raise ValueError(f"LLM returned invalid JSON: {response_text}")
 
 # Singleton
 llm_service = LocalLLMService()