""" Core LLM module for DataVision Handles all interaction with Large Language Models Enterprise-grade error handling and logging SECURED: AI security filter for prompt injection protection """ import logging from typing import List, Dict, Any, Optional, Union import os import contextvars # Global context var to store requested model per-request _requested_model: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar("requested_model", default=None) def set_requested_model(model: str): """Set the requested model for the current context (used by FastAPI endpoints)""" _requested_model.set(model) # 🔒 SECURITY: Import AI security filter try: from core.ai_security import get_ai_security_filter, sanitize_user_input AI_SECURITY_AVAILABLE = True except ImportError: AI_SECURITY_AVAILABLE = False # Try importing litellm, handle if missing or broken LITELLM_AVAILABLE = False try: # 🩹 MONKEY PATCH: Fix aiohttp compatibility for LiteLLM # Some environments have broken aiohttp/litellm combos import aiohttp if not hasattr(aiohttp, 'ConnectionTimeoutError'): # Inject the missing attribute aliased to ClientTimeout or generic ClientError # This prevents LiteLLM from crashing on import if hasattr(aiohttp, 'ClientTimeout'): setattr(aiohttp, 'ConnectionTimeoutError', aiohttp.ClientTimeout) setattr(aiohttp, 'SocketTimeoutError', aiohttp.ClientTimeout) else: setattr(aiohttp, 'ConnectionTimeoutError', Exception) setattr(aiohttp, 'SocketTimeoutError', Exception) import litellm LITELLM_AVAILABLE = True except Exception as e: logger = logging.getLogger(__name__) logger.error(f"Failed to import litellm: {e}") LITELLM_AVAILABLE = False # Try importing sentence-transformers for local embeddings SENTENCE_TRANSFORMER_AVAILABLE = False _embedding_model_cache = {} try: from sentence_transformers import SentenceTransformer SENTENCE_TRANSFORMER_AVAILABLE = True except Exception as e: logger = logging.getLogger(__name__) logger.warning(f"SentenceTransformer not available: {e}") from config.settings import Settings logger = logging.getLogger(__name__) # Configure LiteLLM if available if LITELLM_AVAILABLE: litellm.drop_params = True litellm.logging = False # Suppress verbose logging os.environ["LITELLM_LOG"] = "ERROR" # 🩹 ALIAS MAP: Route user's model to Groq with correct path # The model 'openai/gpt-oss-120b' exists on Groq Console # 'groq/openai/gpt-oss-120b' tells litellm: use Groq, request model 'openai/gpt-oss-120b' litellm.model_alias_map = { "openai/gpt-oss-120b": "groq/openai/gpt-oss-120b" } def get_optimal_model(query_type: str = None, reasoning_depth: str = None, context_length: int = 0) -> str: """ Smart model selection based on query complexity. Args: query_type: Type of query (factual, analytical, comparison, etc.) reasoning_depth: Depth required (shallow, moderate, deep, multi_hop) context_length: Length of context being processed Returns: str: Optimal model name for the query """ # Complex queries requiring reasoning → use reasoning model complex_query_types = { "comparison", "prediction", "trend_analysis", "deep_analysis", "multi_hop", "causal", "predictive", "exploratory" } deep_reasoning = {"deep", "multi_hop", "analytical"} if query_type and query_type.lower() in complex_query_types: return Settings.REASONING_MODEL if reasoning_depth and reasoning_depth.lower() in deep_reasoning: return Settings.REASONING_MODEL # Large context needs more capable model if context_length > 5000: return Settings.REASONING_MODEL # Simple factual queries → use fast model for speed simple_query_types = { "factual", "lookup", "definition", "simple", "greeting", "aggregation" } if query_type and query_type.lower() in simple_query_types: return Settings.FAST_MODEL # Default: use standard model return Settings.MODEL_NAME # ============================================================================ # CHATGPT-STYLE SYSTEM PROMPT - Clean, Professional Responses # ============================================================================ # Import product knowledge try: from core.product_knowledge import DATAVISION_PRODUCT_KNOWLEDGE, is_product_question PRODUCT_KNOWLEDGE_AVAILABLE = True except ImportError: DATAVISION_PRODUCT_KNOWLEDGE = "" PRODUCT_KNOWLEDGE_AVAILABLE = False CHATGPT_SYSTEM_PROMPT = f"""You are DataVision AI Assistant - a professional data analyst and product guide. ## YOUR THREE ROLES: 1. 📊 **DATA ANALYST** - Answer questions about user's uploaded data 2. 🌐 **AI KNOWLEDGE** - Answer general questions using AI knowledge 3. 💡 **PRODUCT GUIDE** - Help users navigate and use DataVision features {DATAVISION_PRODUCT_KNOWLEDGE if PRODUCT_KNOWLEDGE_AVAILABLE else ""} ## RESPONSE RULES: ### For DATA Questions (about their uploaded files): 📊 **From Your Data:** [Use actual numbers from their data - NEVER make up values] ### For GENERAL Questions (not about their data): 🌐 **AI Knowledge:** [Provide helpful AI-powered answers] ### For PRODUCT Questions (how to use DataVision): 💡 **How to do this:** [Step-by-step instructions with page names] ## CRITICAL RULES: 1. **NO HALLUCINATION** - For data questions, ONLY use actual provided data 2. **ALWAYS BE HELPFUL** - Never refuse to answer, use appropriate knowledge source 3. **BE DIRECT** - Answer first, then explain 4. **GUIDE USERS** - If they're asking "how to" questions, give step-by-step help ## DETECTING QUESTION TYPE: - "What are my sales?" → DATA question (use their data) - "What is machine learning?" → GENERAL question (use AI knowledge) - "How do I train a model?" → PRODUCT question (guide them) - "How do I upload data?" → PRODUCT question (guide them) Be friendly, professional, and genuinely helpful like ChatGPT!""" def chat( messages: Union[str, List[Dict[str, str]]], system: Optional[str] = None, temperature: float = 0.7, model: str = None, max_tokens: int = 4000, images: Optional[List[str]] = None, **kwargs ) -> str: """ Send chat to LLM and get response. Uses Groq as primary with automatic Gemini fallback. SECURED: Applies AI security filter to prevent prompt injection. Args: messages: List of message dicts (role, content) OR single string prompt system: Optional system prompt (will be prepended) temperature: Sampling temperature model: Optional model override max_tokens: Max tokens to generate Returns: str: Helper response content """ if not LITELLM_AVAILABLE: logger.error("LiteLLM not installed. Please install: pip install litellm") return "Error: LLM driver (LiteLLM) missing." # Use requested model from context if not explicitly provided if model is None: model = _requested_model.get() # 🔒 SECURITY: Apply AI security filter to user input ai_filter = None if AI_SECURITY_AVAILABLE: ai_filter = get_ai_security_filter() # Prepare messages final_messages = [] # Add system prompt - use ChatGPT prompt by default for cleaner responses effective_system = system if system else CHATGPT_SYSTEM_PROMPT final_messages.append({"role": "system", "content": effective_system}) # Handle input: If string, convert to user message if isinstance(messages, str): # 🔒 SECURITY: Sanitize user input safe_message = messages if ai_filter: safe_message, was_suspicious, pattern = ai_filter.filter_input(messages) if was_suspicious: logger.warning(f"Prompt injection attempt detected: {pattern}") final_messages.append({"role": "user", "content": safe_message}) else: # 🔒 SECURITY: Sanitize each user message in the list for msg in messages: if msg.get("role") == "user" and ai_filter: safe_content, was_suspicious, pattern = ai_filter.filter_input(msg.get("content", "")) if was_suspicious: logger.warning(f"Prompt injection attempt detected: {pattern}") # Handle multimodal content formatting if images and msg == messages[-1]: # Only attach images to the latest user message content_block = [{"type": "text", "text": safe_content}] for img in images: # Add base64 prefix if missing img_str = img if img.startswith("data:image") else f"data:image/jpeg;base64,{img}" content_block.append({"type": "image_url", "image_url": {"url": img_str}}) final_messages.append({"role": "user", "content": content_block}) else: final_messages.append({"role": msg["role"], "content": safe_content}) else: final_messages.append(msg) # If it was a string prompt and we have images, format the single message if isinstance(messages, str) and images and final_messages[-1]["role"] == "user": safe_content = final_messages.pop()["content"] content_block = [{"type": "text", "text": safe_content}] for img in images: img_str = img if img.startswith("data:image") else f"data:image/jpeg;base64,{img}" content_block.append({"type": "image_url", "image_url": {"url": img_str}}) final_messages.append({"role": "user", "content": content_block}) # ========================================================================= # ☁️ CASCADING PROVIDER FALLBACK CHAIN # Chain: Nvidia NIM -> HuggingFace -> Ollama # ========================================================================= # 1. Build the Provider Chain providers = [] nv_key = os.environ.get("NVIDIA_API_KEY") groq_key = os.environ.get("GROQ_API_KEY") hf_key = os.environ.get("HUGGINGFACE_API_KEY") # Check if we have any remote API keys has_remote_keys = bool(nv_key or groq_key or hf_key) # --- 1. PRIMARY: GROQ (Ultra-Fast with Key Rotation) --- from config.settings import Settings # Add all available Groq keys to the fallback chain for i, g_key in enumerate(Settings.GROQ_API_KEYS): providers.append({ "name": f"Groq-Llama3.3 (Key {i})", "model": "groq/llama-3.3-70b-versatile", "api_key": g_key, "api_base": None }) providers.append({ "name": f"Groq-Llama3.1-Fast (Key {i})", "model": "groq/llama-3.1-8b-instant", "api_key": g_key, "api_base": None }) # If no keys in settings (fallback), try direct env if not Settings.GROQ_API_KEYS and groq_key: providers.append({ "name": "Groq-Llama3.3", "model": "groq/llama-3.3-70b-versatile", "api_key": groq_key, "api_base": None }) providers.append({ "name": "Groq-Llama3.1-Fast", "model": "groq/llama-3.1-8b-instant", "api_key": groq_key, "api_base": None }) # --- 2. SECONDARY: NVIDIA LLAMA --- if nv_key: providers.append({ "name": "Nvidia-Llama", "model": "openai/meta/llama-3.1-70b-instruct", "api_key": nv_key, "api_base": "https://integrate.api.nvidia.com/v1" }) # --- HUGGINGFACE --- if hf_key: providers.append({ "name": "HuggingFace", "model": "huggingface/meta-llama/Meta-Llama-3-8B-Instruct", "api_key": hf_key, "api_base": None }) # --- OLLAMA (Local fallback) --- providers.append({ "name": "Ollama", "model": "ollama/llama3", "api_key": None, "api_base": "http://localhost:11434" }) # 2. Execute the Chain last_error = None had_auth_error = False import time as time_module for i, p in enumerate(providers): try: logger.info(f"Try {i+1}/{len(providers)}: {p['name']} ({p['model']})") call_kwargs = { "model": p["model"], "messages": final_messages, "temperature": temperature, "max_tokens": max_tokens } # Inject any additional kwargs (like response_format, tools, tool_choice) call_kwargs.update(kwargs) if p["api_key"]: call_kwargs["api_key"] = p["api_key"] if p["api_base"]: call_kwargs["api_base"] = p["api_base"] is_raw_response = call_kwargs.pop('raw_response', False) # 🔥 FIX: Set ultra-low timeout (15s) so the UI doesn't hang! if "timeout" not in call_kwargs: call_kwargs["timeout"] = 15 # 🩹 FIX: Disable OpenAI client's internal retries - let our cascade handle failover if "num_retries" not in call_kwargs: call_kwargs["num_retries"] = 0 response = litellm.completion(**call_kwargs) logger.info(f"Success with {p['name']}") if is_raw_response: return response result = response.choices[0].message.content if ai_filter: result = ai_filter.filter_output(result) return result except Exception as e: last_error = e error_str = str(e).lower() # Fast-Fail Check is_rate_limit = 'rate_limit' in error_str or 'rate limit' in error_str or '429' in error_str is_auth_error = 'api_key' in error_str or 'unauthorized' in error_str or 'authentication' in error_str or '401' in error_str is_org_restricted = 'organization' in error_str and ('restricted' in error_str or 'blocked' in error_str) is_bad_request = 'badrequest' in error_str.replace(' ', '') or 'invalid_request' in error_str is_connection = 'connection' in error_str or 'timeout' in error_str if is_auth_error: had_auth_error = True reason = 'Rate Limit' if is_rate_limit else ('Org Restricted' if is_org_restricted else ('Bad Request' if is_bad_request else ('Auth Error' if is_auth_error else ('Connection Error' if is_connection else 'Unknown Error')))) if i < len(providers) - 1: logger.warning(f"{p['name']} failed ({reason}). Cascading to next provider...") continue else: logger.error(f"All providers in the cascade chain exhausted. Last error: {e}") break # 3. Handle Complete Failure - Clean user-facing maintenance message return ( "⚙️ **Service Under Maintenance:** The DataVision AI engine is currently performing routine model synchronization. " "Please re-send your question in a moment or switch to Analyst mode." ) def embed_text(text: str) -> List[float]: """ Generate embedding for a single text string Args: text: Input text Returns: List[float]: Embedding vector """ if not LITELLM_AVAILABLE: logger.warning("LiteLLM not available for embedding, returning zero vector") return [0.0] * Settings.EMBED_DIM try: # 1. Try local SentenceTransformer (Zero-cost, fast) if SENTENCE_TRANSFORMER_AVAILABLE and "sentence-transformers" in Settings.EMBEDDING_MODEL: if Settings.EMBEDDING_MODEL not in _embedding_model_cache: logger.info(f"📥 Loading local embedding model: {Settings.EMBEDDING_MODEL}") _embedding_model_cache[Settings.EMBEDDING_MODEL] = SentenceTransformer(Settings.EMBEDDING_MODEL) model = _embedding_model_cache[Settings.EMBEDDING_MODEL] embedding = model.encode([text])[0] return embedding.tolist() # 2. Try LiteLLM (API-based) response = litellm.embedding( model=Settings.EMBEDDING_MODEL, input=[text] ) return response.data[0].embedding except Exception as e: logger.error(f"Embedding error: {str(e)}") # Return zero vector on error to prevent partial failure return [0.0] * Settings.EMBED_DIM def embed_texts(texts: List[str]) -> List[List[float]]: """ Generate embeddings for a list of texts Args: texts: List of input texts Returns: List[List[float]]: List of embedding vectors """ if not LITELLM_AVAILABLE: logger.warning("LiteLLM not available for batch embedding, returning zero vectors") return [[0.0] * Settings.EMBED_DIM for _ in texts] try: # 1. Try local SentenceTransformer if SENTENCE_TRANSFORMER_AVAILABLE and "sentence-transformers" in Settings.EMBEDDING_MODEL: if Settings.EMBEDDING_MODEL not in _embedding_model_cache: logger.info(f"📥 Loading local embedding model: {Settings.EMBEDDING_MODEL}") _embedding_model_cache[Settings.EMBEDDING_MODEL] = SentenceTransformer(Settings.EMBEDDING_MODEL) model = _embedding_model_cache[Settings.EMBEDDING_MODEL] embeddings = model.encode(texts) return [emb.tolist() for emb in embeddings] # 2. Try LiteLLM response = litellm.embedding( model=Settings.EMBEDDING_MODEL, input=texts ) return [item.embedding for item in response.data] except Exception as e: logger.error(f"Batch embedding error: {str(e)}") # Return zero vectors on error return [[0.0] * Settings.EMBED_DIM for _ in texts]