from fastapi import FastAPI, HTTPException, Header, Response from fastapi.responses import HTMLResponse, StreamingResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel import google.generativeai as genai import httpx import os from dotenv import load_dotenv from duckduckgo_search import DDGS from typing import Optional, List, Dict, Any import logging import re import asyncio import threading import time import hashlib import json from functools import wraps from collections import OrderedDict from dataclasses import dataclass, field # Load environment variables from .env file load_dotenv() import spacy from rake_nltk import Rake import nltk from cache.chromadb_cache import ChromaDBSearchCache from search_optimizer import ( has_meaningful_conversation_history, hybrid_search_decision, extract_search_terms, format_search_context ) nltk.download('stopwords') nltk.download('punkt_tab') # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Initialize FastAPI app app = FastAPI( title="Enhanced Chat API with Dynamic Model Selection", description="API with query classification and dynamic model loading for QA/Summarization", version="3.0.0" ) # Configure CORS with restricted origins for security app.add_middleware( CORSMiddleware, allow_origins=[ "https://huggingface.co", "https://*.hf.space", "http://localhost:3000", "http://localhost:8000", "http://127.0.0.1:3000", "http://127.0.0.1:8000" ], allow_methods=["POST", "GET"], allow_headers=["Content-Type", "Authorization"], ) # Request/Response models class ChatRequest(BaseModel): prompt: str max_new_tokens: int = 500 use_search: bool = True temperature: float = 0.7 user_id: Optional[str] = None history: Optional[List[Dict[str, str]]] = None force_search: Optional[bool] = None search_decision_mode: str = "balanced" # conservative, balanced, aggressive class ChatResponse(BaseModel): response: str search_results: Optional[List[Dict[str, Any]]] = None search_decision: Optional[Dict[str, Any]] = None cache_info: Optional[Dict[str, Any]] = None class SearchRequest(BaseModel): query: str max_results: int = 5 # Configure Google AI try: genai.configure(api_key=os.environ["GOOGLE_API_KEY"]) except KeyError: logger.error("CRITICAL: GOOGLE_API_KEY environment variable not set.") raise # Exit if API key is not set # Initialize the Generative Model model = genai.GenerativeModel('gemini-1.5-flash') # Global NLP tools nlp = spacy.load("en_core_web_sm") rake = Rake() # ===== Phase 3c: ChromaDB Vector Database Cache System ===== # Universal ChromaDB search cache - single global instance for all users # Vector database provides better semantic matching and persistent storage universal_search_cache = None def get_universal_cache() -> ChromaDBSearchCache: """Get the universal ChromaDB search cache instance""" global universal_search_cache if universal_search_cache is None: # Initialize ChromaDB cache with environment configuration cache_db_path = os.getenv('CHROMADB_PATH', 'cache_db') cache_results_path = os.getenv('CACHE_RESULTS_PATH', 'cache_results') embedding_model = os.getenv('CACHE_EMBEDDING_MODEL', 'all-MiniLM-L6-v2') universal_search_cache = ChromaDBSearchCache( max_size=1000, # Large cache size default_ttl=3600, # 1 hour TTL cache_db_path=cache_db_path, cache_results_path=cache_results_path, embedding_model=embedding_model, similarity_threshold=0.7 ) logger.info(f"Initialized ChromaDB universal cache: {cache_db_path}") return universal_search_cache # ===== End Phase 3c ChromaDB Cache System ===== def normalize_user_id(user_id: Optional[str]) -> Optional[str]: """Normalize user_id to None for anonymous requests""" if user_id is None or user_id.strip() == "": return None return user_id.strip() def validate_user_id(user_id: Optional[str]) -> Optional[str]: """Validate user_id format and return normalized value""" # Normalize to None for anonymous users user_id = normalize_user_id(user_id) if user_id is None: return None # Anonymous user # Validate format: non-empty string, max 255 chars, alphanumeric + hyphens + underscores if len(user_id) > 255: raise HTTPException(status_code=400, detail="user_id must be 255 characters or less") # Check allowed characters if not re.match(r'^[a-zA-Z0-9_-]+$', user_id): raise HTTPException(status_code=400, detail="user_id can only contain alphanumeric characters, hyphens, and underscores") return user_id def run_in_threadpool(func): """Decorator to run synchronous model inference in thread pool""" @wraps(func) async def wrapper(*args, **kwargs): loop = asyncio.get_event_loop() return await loop.run_in_executor(None, func, *args, **kwargs) return wrapper async def search_brave(query: str, max_results: int = 5) -> List[Dict[str, Any]]: """Search using Brave Search API with async HTTP client""" try: # Check for API key in environment variable first, then fallback to hardcoded api_key = os.getenv('BRAVE_API_KEY') if not api_key: logger.error("No Brave API key available") return [] headers = { 'Accept': 'application/json', 'Accept-Encoding': 'gzip', 'X-Subscription-Token': api_key } params = { 'q': query, 'count': max_results, 'safesearch': 'moderate', 'search_lang': 'en', 'country': 'US' } async with httpx.AsyncClient(timeout=10.0) as client: response = await client.get( 'https://api.search.brave.com/res/v1/web/search', headers=headers, params=params ) response.raise_for_status() data = response.json() web_results = data.get('web', {}) raw_results = web_results.get('results', []) results = [] for result in raw_results: results.append({ "title": result.get("title", ""), "body": re.sub(r'\s+', ' ', result.get("description", "")).strip(), "href": result.get("url", ""), "source": "Brave" }) return results[:max_results] except Exception as e: logger.error(f"Brave Search error: {e}") logger.error(f"Brave Search error type: {type(e).__name__}") import traceback logger.error(f"Brave Search traceback: {traceback.format_exc()}") return [] async def search_duckduckgo(query: str, max_results: int = 5) -> List[Dict[str, Any]]: """Search using DuckDuckGo with async execution and timeout handling""" try: # Run DuckDuckGo search in thread pool with timeout loop = asyncio.get_event_loop() results = await asyncio.wait_for( loop.run_in_executor(None, _sync_duckduckgo_search, query, max_results), timeout=8.0 # 8 second timeout for Hugging Face compatibility ) return results except asyncio.TimeoutError: logger.warning(f"DuckDuckGo search timed out for query: {query}") return [] except Exception as e: logger.error(f"DuckDuckGo Search error: {e}") return [] def _sync_duckduckgo_search(query: str, max_results: int) -> List[Dict[str, Any]]: """Synchronous DuckDuckGo search helper with retry logic""" max_retries = 2 for attempt in range(max_retries): try: # Configure DDGS with more conservative settings for hosted environments with DDGS(timeout=5) as ddgs: results = [] search_results = ddgs.text( query, safesearch='moderate', max_results=max_results, region='us-en' # Specify region to potentially avoid some blocks ) for result in search_results: results.append({ "title": result.get("title", ""), "body": re.sub(r'\s+', ' ', result.get("body", "")).strip(), "href": result.get("href", ""), "source": "DuckDuckGo" }) logger.info(f"DuckDuckGo search successful on attempt {attempt + 1}") return results except Exception as e: logger.warning(f"DuckDuckGo attempt {attempt + 1} failed: {e}") if attempt == max_retries - 1: # Last attempt logger.error(f"DuckDuckGo search failed after {max_retries} attempts") raise # Wait briefly before retry import time time.sleep(1) return [] async def search_web_combined(query: str, max_results: int = 10) -> List[Dict[str, Any]]: """Combined web search with resilient fallback strategy""" try: logger.info(f"Combined Search: Starting search for '{query}'") # Run both searches concurrently with timeout protection brave_task = asyncio.create_task(search_brave(query, 6)) # Get more from Brave as primary duckduckgo_task = asyncio.create_task(search_duckduckgo(query, 4)) # Fewer from DDG as backup # Wait for both searches to complete with overall timeout try: brave_results, duckduckgo_results = await asyncio.wait_for( asyncio.gather(brave_task, duckduckgo_task, return_exceptions=True), timeout=12.0 # Overall timeout for both searches ) except asyncio.TimeoutError: logger.warning("Combined search timed out, cancelling remaining tasks") brave_task.cancel() duckduckgo_task.cancel() brave_results, duckduckgo_results = [], [] # Handle exceptions and log results if isinstance(brave_results, Exception): logger.error(f"Brave search failed: {brave_results}") brave_results = [] else: logger.info(f"Brave search returned {len(brave_results)} results") if isinstance(duckduckgo_results, Exception): logger.error(f"DuckDuckGo search failed: {duckduckgo_results}") duckduckgo_results = [] else: logger.info(f"DuckDuckGo search returned {len(duckduckgo_results)} results") # If both searches fail, return empty with warning if not brave_results and not duckduckgo_results: logger.warning(f"All search engines failed for query: {query}") return [] # Combine results combined_results = brave_results + duckduckgo_results # Remove duplicates based on URL seen_urls = set() unique_results = [] for result in combined_results: url = result.get("href", "") if url and url not in seen_urls: seen_urls.add(url) unique_results.append(result) # Return top results up to max_results final_results = unique_results[:max_results] logger.info(f"Combined Search: Returning {len(final_results)} total unique results") # Log search engine performance brave_count = len([r for r in final_results if r.get('source') == 'Brave']) ddg_count = len([r for r in final_results if r.get('source') == 'DuckDuckGo']) logger.info(f"Search distribution: Brave={brave_count}, DuckDuckGo={ddg_count}") return final_results except Exception as e: logger.error(f"Combined search error: {e}") return [] def preprocess_text(text: str) -> str: """Use spaCy for fast text cleaning/normalization""" doc = nlp(text) # Lemmatize and remove stopwords return " ".join([ token.lemma_ for token in doc if not token.is_stop and not token.is_punct ])[:2048] def format_conversation_history(history: Optional[List[Dict[str, str]]], max_entries: int = 10) -> str: """Format conversation history for inclusion in AI prompt""" if not history: return "" try: formatted_entries = [] # Take the most recent entries up to max_entries recent_history = history[-max_entries:] if len(history) > max_entries else history for entry in recent_history: # Handle both formats: {"role": "user/assistant", "content": "..."} # and {"user": "...", "assistant": "..."} if "role" in entry and "content" in entry: role = entry["role"].title() # User or Assistant content = entry["content"].strip() if content: formatted_entries.append(f"{role}: {content}") elif "user" in entry and "assistant" in entry: user_msg = entry["user"].strip() assistant_msg = entry["assistant"].strip() if user_msg and assistant_msg: formatted_entries.append(f"User: {user_msg}") formatted_entries.append(f"Assistant: {assistant_msg}") if formatted_entries: return "\n".join(formatted_entries) return "" except Exception as e: logger.warning(f"Error formatting conversation history: {e}") return "" @app.on_event("startup") async def startup_event(): """Perform startup tasks like NLP model loading""" logger.info("Starting NLP model loading...") # spaCy and NLTK data are loaded implicitly on first use or by spacy.load # Test analytics database connection try: from analytics.database import test_connection analytics_connected = await test_connection() if analytics_connected: logger.info("Analytics database connection successful") else: logger.warning("Analytics database connection failed - analytics disabled") except Exception as e: logger.warning(f"Analytics initialization failed: {e}") logger.info("Startup complete - NLP models ready") @app.post("/chat", response_model=ChatResponse) async def chat_endpoint(request: ChatRequest, response: Response, user_agent: str = Header(None), x_session_id: str = Header(None)): """Enhanced chat endpoint with dynamic model selection and combined search""" logger.info(f"Request: {request.prompt}") # Validate and extract user_id user_id = validate_user_id(request.user_id) # Analytics setup session_id = x_session_id message_id = None start_time = None try: # Import analytics (with fallback if not available) try: from analytics.collectors import create_session, get_session, track_message, PerformanceTimer analytics_available = True except ImportError: logger.warning("Analytics not available") analytics_available = False # Create fallback PerformanceTimer when analytics unavailable class PerformanceTimer: def __init__(self): self.start_time = None self.end_time = None def __enter__(self): import time self.start_time = time.time() return self def __exit__(self, exc_type, exc_val, exc_tb): import time self.end_time = time.time() @property def duration_ms(self) -> int: if self.start_time and self.end_time: return int((self.end_time - self.start_time) * 1000) return 0 # Start performance timing with PerformanceTimer() as timer: # Handle session management if analytics_available: if not session_id: # Create new session session = await create_session(user_agent=user_agent, user_id=user_id) session_id = session.session_id else: # Get existing session or create new one if not found session = await get_session(session_id) if not session: session = await create_session(user_agent=user_agent, user_id=user_id) session_id = session.session_id search_results = [] search_context = "" search_decision = None cache_info = None # Phase 3: Get universal cache for result caching search_cache = get_universal_cache() # Phase 3b: Context-Aware Request Flow Optimization logger.info(f"Use Search : {request.use_search}") has_history = has_meaningful_conversation_history(request.history) logger.info(f"Has meaningful conversation history: {has_history}") if request.use_search: # Check if search should be forced (overrides all optimizations) if request.force_search: search_decision = { "should_search": True, "reason": "Search forced by user", "confidence": 1.0, "flow_type": "forced" } perform_search = True logger.info(f"Using forced search flow") elif not has_history: # First Message (No History): Cache-first approach logger.info(f"Using cache-first flow (no conversation history)") search_terms = extract_search_terms(request.prompt.lower(), nlp, rake) logger.info(f"Extract search terms successful: {search_terms}") # Try to get from cache first (skip search decision for performance) cached_entry = search_cache.get(search_terms, use_semantic_matching=True, similarity_threshold=0.7) if cached_entry: # Cache hit! Use cached results, no search needed search_results = cached_entry.results search_context = format_search_context(search_results) cache_info = { "cache_hit": True, "cached_query": cached_entry.search_query, "cache_age_seconds": int(time.time() - cached_entry.timestamp), "hit_count": cached_entry.hit_count, "flow_type": "cache_first_hit", "cache_type": "chromadb_vector" } search_decision = { "should_search": False, "reason": "Cache hit in cache-first flow", "confidence": 1.0, "flow_type": "cache_first_hit", "cache_type": "chromadb_vector" } perform_search = False logger.info(f"ChromaDB Cache-first HIT: Using cached results (age: {cache_info['cache_age_seconds']}s, hits: {cached_entry.hit_count})") else: # Cache miss - perform web search without search decision overhead logger.info(f"ChromaDB Cache-first MISS: Performing web search") search_query = " ".join(search_terms) or request.prompt search_results = await search_web_combined(search_query, 10) search_context = format_search_context(search_results) # Store results in cache if search was successful if search_results: search_cache.put(search_terms, search_query, search_results) logger.info(f"Cached search results in universal cache for future use") cache_info = { "cache_hit": False, "stored_in_cache": len(search_results) > 0, "flow_type": "cache_first_miss", "cache_type": "chromadb_vector" } search_decision = { "should_search": True, "reason": "Cache miss in cache-first flow", "confidence": 1.0, "flow_type": "cache_first_miss", "cache_type": "chromadb_vector" } perform_search = True else: # Follow-up Messages (Has History): Search decision first logger.info(f"Using search-decision-first flow (has conversation history)") # Use hybrid intelligent search decision (Phase 2: AI-enhanced) search_decision = await hybrid_search_decision( request.prompt, request.history, request.search_decision_mode, nlp, model ) search_decision["flow_type"] = "search_decision_first" perform_search = search_decision["should_search"] if perform_search: # Search needed - check cache before web search search_terms = extract_search_terms(request.prompt.lower(), nlp, rake) logger.info(f"Extract search terms successful: {search_terms}") # Try to get from cache first cached_entry = search_cache.get(search_terms, use_semantic_matching=True, similarity_threshold=0.7) if cached_entry: # Cache hit! Use cached results search_results = cached_entry.results search_context = format_search_context(search_results) cache_info = { "cache_hit": True, "cached_query": cached_entry.search_query, "cache_age_seconds": int(time.time() - cached_entry.timestamp), "hit_count": cached_entry.hit_count, "flow_type": "search_decision_cache_hit", "cache_type": "chromadb_vector" } logger.info(f"ChromaDB Search-decision flow Cache HIT: Using cached results (age: {cache_info['cache_age_seconds']}s, hits: {cached_entry.hit_count})") else: # Cache miss - perform web search logger.info(f"ChromaDB Search-decision flow Cache MISS: Performing web search") search_query = " ".join(search_terms) or request.prompt search_results = await search_web_combined(search_query, 10) search_context = format_search_context(search_results) # Store results in cache if search was successful if search_results: search_cache.put(search_terms, search_query, search_results) logger.info(f"Cached search results in universal cache for future use") cache_info = { "cache_hit": False, "stored_in_cache": len(search_results) > 0, "flow_type": "search_decision_cache_miss", "cache_type": "chromadb_vector" } else: # Search not needed based on conversation context logger.info(f"Search skipped: {search_decision['reason']}") cache_info = { "search_skipped": True, "flow_type": "search_decision_skip" } logger.info(f"Search Decision: {search_decision}") if perform_search: logger.info(f"Search processing complete - Results: {len(search_results)}") else: logger.info(f"Search disabled by request") cache_info = {"search_disabled": True} logger.info(f"Search Context: {search_context}") # Format conversation history conversation_history = format_conversation_history(request.history, max_entries=10) logger.info(f"Conversation History: {len(request.history or [])} entries") # Create a prompt for the AI model with conversation history if conversation_history: prompt_template = f""" You are having a conversation with a user. Here is the conversation history: Conversation History: --- {conversation_history} --- Based on the following context from web pages I have read and the conversation history above, please answer the user's question. If the context does not contain the answer and you cannot answer based on the conversation history, say that you don't have enough information. Context: --- {search_context} --- Question: {request.prompt} Answer: """ else: prompt_template = f""" Based on the following context from web pages I have read, please answer the user's question. If the context does not contain the answer, say that you don't have enough information. Context: --- {search_context} --- Question: {request.prompt} Answer: """ response_text = await run_gemini_inference(prompt_template) # Track message analytics if analytics_available and session_id: message = await track_message( session_id=session_id, prompt_length=len(request.prompt), response_length=len(response_text), response_time_ms=timer.duration_ms, used_search=request.use_search, max_tokens=request.max_new_tokens, temperature=request.temperature, success=True, user_id=user_id ) if message: message_id = message.message_id # Add session ID to response headers if session_id: response.headers["X-Session-ID"] = session_id # Prepare response chat_response = ChatResponse( response=response_text, search_results=search_results, search_decision=search_decision, cache_info=cache_info ) return chat_response except Exception as e: # Track failed message if analytics_available and session_id: await track_message( session_id=session_id, prompt_length=len(request.prompt), response_length=0, response_time_ms=timer.duration_ms if 'timer' in locals() else 0, used_search=request.use_search, max_tokens=request.max_new_tokens, temperature=request.temperature, success=False, error_message=str(e), user_id=user_id ) logger.error(f"Chat error: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.post("/search") async def search_endpoint(request: SearchRequest): """Search endpoint with combined search engines""" try: results = await search_web_combined(request.query, request.max_results) return {"results": results} except Exception as e: logger.error(f"Search endpoint error: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.get("/analytics/stats") async def analytics_stats(): """Get basic analytics statistics""" try: from analytics.dashboard import get_basic_stats stats = await get_basic_stats() return stats except Exception as e: logger.error(f"Analytics stats error: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.get("/analytics/dashboard") async def analytics_dashboard(): """Get HTML analytics dashboard""" try: from analytics.dashboard import get_dashboard_data from fastapi.responses import HTMLResponse # Get dashboard data including user metrics data = await get_dashboard_data() # Get user statistics for the dashboard from analytics.dashboard import get_user_statistics, get_authenticated_vs_anonymous_metrics user_stats = await get_user_statistics() comparison_stats = await get_authenticated_vs_anonymous_metrics() # Create HTML dashboard html_content = f"""
Real-time insights into your chat application