Add thread-safe model manager and combined search engines

- Implemented ModelManager class with thread-safe lazy loading of ML models
- Added Brave Search API integration alongside DuckDuckGo search
- Enhanced CORS configuration with restricted origins for security
- Converted synchronous model inference to async with thread pool execution
- Improved error handling and expanded search result context
- Updated API version to 3.0.0

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

app.py

@@ -3,10 +3,15 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import torch
 from transformers import pipeline
+import requests
+import os
 from duckduckgo_search import DDGS
 from typing import Optional, List, Dict, Any
 import logging
 import re
+import asyncio
+import threading
+from functools import wraps
 
 
 import spacy
@@ -24,15 +29,22 @@ logger = logging.getLogger(__name__)
 app = FastAPI(
     title="Enhanced Chat API with Dynamic Model Selection",
     description="API with query classification and dynamic model loading for QA/Summarization",
-    version="2.0.0"
+    version="3.0.0"
 )
 
-# Configure CORS
+# Configure CORS with restricted origins for security
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    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=["*"],
+    allow_headers=["Content-Type", "Authorization"],
 )
 
 # Request/Response models
@@ -51,10 +63,81 @@ class SearchRequest(BaseModel):
     query: str
     max_results: int = 5
 
-# Global pipelines with lazy loading
-classifier_pipeline = None
-qa_pipeline = None
-summarization_pipeline = None
+# Thread-safe model manager
+class ModelManager:
+    def __init__(self):
+        self._classifier_pipeline = None
+        self._qa_pipeline = None
+        self._summarization_pipeline = None
+        self._classifier_lock = threading.RLock()
+        self._qa_lock = threading.RLock()
+        self._summarization_lock = threading.RLock()
+        
+    def get_classifier(self):
+        if self._classifier_pipeline is None:
+            with self._classifier_lock:
+                if self._classifier_pipeline is None:
+                    self._load_classifier()
+        return self._classifier_pipeline
+    
+    def get_qa_model(self):
+        if self._qa_pipeline is None:
+            with self._qa_lock:
+                if self._qa_pipeline is None:
+                    self._load_qa_model()
+        return self._qa_pipeline
+    
+    def get_summarization_model(self):
+        if self._summarization_pipeline is None:
+            with self._summarization_lock:
+                if self._summarization_pipeline is None:
+                    self._load_summarization_model()
+        return self._summarization_pipeline
+    
+    def _load_classifier(self):
+        """Load zero-shot classification model"""
+        try:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            self._classifier_pipeline = pipeline(
+                "zero-shot-classification",
+                model="valhalla/distilbart-mnli-12-3",
+                device=0 if device == "cuda" else -1
+            )
+            logger.info("Zero-shot classifier loaded")
+        except Exception as e:
+            logger.error(f"Classifier load error: {e}")
+            raise
+    
+    def _load_qa_model(self):
+        """Load question-answering model"""
+        try:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            self._qa_pipeline = pipeline(
+                "question-answering",
+                model="distilbert-base-uncased-distilled-squad",
+                device=0 if device == "cuda" else -1
+            )
+            logger.info("QA model loaded")
+        except Exception as e:
+            logger.error(f"QA model load error: {e}")
+            raise
+    
+    def _load_summarization_model(self):
+        """Load summarization model"""
+        try:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            self._summarization_pipeline = pipeline(
+                "summarization",
+                model="sshleifer/distilbart-cnn-6-6",
+                device=0 if device == "cuda" else -1
+            )
+            logger.info("Summarization model loaded")
+        except Exception as e:
+            logger.error(f"Summarization model load error: {e}")
+            raise
+
+# Global thread-safe model manager and NLP tools
+model_manager = ModelManager()
 nlp = spacy.load("en_core_web_sm")
 rake = Rake()
 
@@ -108,52 +191,20 @@ def clean_terms(terms: List[str]) -> List[str]:
     
     return final_terms
 
-def load_classifier():
-    """Load zero-shot classification model"""
-    global classifier_pipeline
-    try:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        classifier_pipeline = pipeline(
-            "zero-shot-classification",
-            model="valhalla/distilbart-mnli-12-3",
-            device=0 if device == "cuda" else -1
-        )
-        logger.info("Zero-shot classifier loaded")
-    except Exception as e:
-        logger.error(f"Classifier load error: {e}")
-
-def load_qa_model():
-    """Load question-answering model on demand"""
-    global qa_pipeline
-    try:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        qa_pipeline = pipeline(
-            "question-answering",
-            model="distilbert-base-uncased-distilled-squad",
-            device=0 if device == "cuda" else -1
-        )
-        logger.info("QA model loaded")
-    except Exception as e:
-        logger.error(f"QA model load error: {e}")
-
-def load_summarization_model():
-    """Load summarization model on demand"""
-    global summarization_pipeline
-    try:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        summarization_pipeline = pipeline(
-            "summarization",
-            model="sshleifer/distilbart-cnn-6-6",
-            device=0 if device == "cuda" else -1
-        )
-        logger.info("Summarization model loaded")
-    except Exception as e:
-        logger.error(f"Summarization model load error: {e}")
+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
 
+@run_in_threadpool
 def classify_query(prompt: str) -> str:
     """Classify query using zero-shot learning"""
     candidate_labels = ['question answering', 'summarization']
     try:
+        classifier_pipeline = model_manager.get_classifier()
         result = classifier_pipeline(
             prompt,
             candidate_labels,
@@ -164,24 +215,97 @@ def classify_query(prompt: str) -> str:
         logger.error(f"Classification failed: {e}")
         return 'question answering'  # Fallback to QA
 
-def search_web(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
-    """Enhanced web search with error handling"""
+def search_brave(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
+    """Search using Brave Search API"""
+    try:
+        api_key = os.getenv('BRAVE_API_KEY') or 'BSAaiYwrOKAj6njwCZ5IZVaJBfvCVNL'
+        if not api_key:
+            logger.error("BRAVE_API_KEY environment variable not set")
+            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'
+        }
+        
+        response = requests.get(
+            'https://api.search.brave.com/res/v1/web/search',
+            headers=headers,
+            params=params,
+            timeout=10
+        )
+        response.raise_for_status()
+        
+        data = response.json()
+        results = []
+        
+        for result in data.get('web', {}).get('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}")
+        return []
+
+def search_duckduckgo(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
+    """Search using DuckDuckGo"""
     try:
         with DDGS() as ddgs:
-            return [{
-                "title": r.get("title", ""),
-                "body": re.sub(r'\s+', ' ', r.get("body", "")).strip(),
-                "href": r.get("href", "")
-            } for r in ddgs.text(query, safesearch='off', max_results=max_results)]
+            results = []
+            for result in ddgs.text(query, safesearch='moderate', max_results=max_results):
+                results.append({
+                    "title": result.get("title", ""),
+                    "body": re.sub(r'\s+', ' ', result.get("body", "")).strip(),
+                    "href": result.get("href", ""),
+                    "source": "DuckDuckGo"
+                })
+            return results
     except Exception as e:
-        logger.error(f"Search error: {e}")
+        logger.error(f"DuckDuckGo Search error: {e}")
         return []
 
+def search_web_combined(query: str, max_results: int = 10) -> List[Dict[str, Any]]:
+    """Combined web search using both Brave and DuckDuckGo"""
+    # Get 5 results from each search engine
+    brave_results = search_brave(query, 5)
+    duckduckgo_results = search_duckduckgo(query, 5)
+    
+    # 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
+    return unique_results[:max_results]
+
 def format_search_context(results: List[Dict[str, Any]]) -> str:
-    """Create condensed context from search results"""
-    return  "\n".join(
-        f"{i+1}. {res['title']}: {res['body'][:200]}"
-        for i, res in enumerate(results[:5])
+    """Create expanded context from combined search results"""
+    return "\n".join(
+        f"{i+1}. [{res.get('source', 'Unknown')}] {res['title']}: {res['body'][:800]}"
+        for i, res in enumerate(results[:10])
     )
 
 def preprocess_text(text: str) -> str:
@@ -191,18 +315,16 @@ def preprocess_text(text: str) -> str:
     return " ".join([
         token.lemma_ for token in doc
         if not token.is_stop and not token.is_punct
-    ])[:1024]
+    ])[:2048]
     
 @app.on_event("startup")
 async def startup_event():
-    """Initialize core models on startup"""
-    load_classifier()
-    load_qa_model()
-    load_summarization_model()
+    """Initialize model manager on startup"""
+    logger.info("Model manager initialized - models will load on demand")
 
 @app.post("/chat", response_model=ChatResponse)
 async def chat_endpoint(request: ChatRequest):
-    """Enhanced chat endpoint with dynamic model selection"""
+    """Enhanced chat endpoint with dynamic model selection and combined search"""
     logger.info(f"Request: {request.prompt}")  
     try:
         search_results = []
@@ -211,27 +333,23 @@ async def chat_endpoint(request: ChatRequest):
         if request.use_search:
             search_terms = extract_search_terms(request.prompt.lower())
             search_query = " ".join(search_terms) or request.prompt
-            search_results = search_web(search_query)
+            search_results = search_web_combined(search_query, 10)
             search_context = format_search_context(search_results)
 
         logger.info(f"Search Context: {search_context}")
         # Query classification
-        task_type = classify_query(request.prompt)
+        task_type = await classify_query(request.prompt)
         logger.info(f"Classified task: {task_type}")
 
         if task_type == 'question answering':
-            response = qa_pipeline(
-                question=request.prompt,
-                context=search_context or request.prompt,
-                max_answer_len=100
-            )['answer']
+            response = await run_qa_inference(
+                request.prompt,
+                search_context or request.prompt
+            )
         else:
-            response = summarization_pipeline(
-                search_context.lower() or request.prompt.lower(),
-                max_length=150,
-                min_length=30,
-                do_sample=False
-            )[0]['summary_text']
+            response = await run_summarization_inference(
+                search_context.lower() or request.prompt.lower()
+            )
 
         return ChatResponse(
             response=response,
@@ -245,9 +363,9 @@ async def chat_endpoint(request: ChatRequest):
 
 @app.post("/search")
 async def search_endpoint(request: SearchRequest):
-    """Search endpoint with improved error handling"""
+    """Search endpoint with combined search engines"""
     try:
-        return {"results": search_web(request.query, request.max_results)}
+        return {"results": search_web_combined(request.query, request.max_results)}
     except Exception as e:
         logger.error(f"Search endpoint error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
@@ -256,12 +374,13 @@ async def search_endpoint(request: SearchRequest):
 async def root():
     """Enhanced health check with model status"""
     return {
-        "message": "Open Source Chat API is running!",
+        "message": "Open Source Chat API with Combined Search is running!",
         "models": {
-            "classifier": bool(classifier_pipeline),
-            "qa": bool(qa_pipeline),
-            "summarization": bool(summarization_pipeline)
+            "classifier": bool(model_manager._classifier_pipeline),
+            "qa": bool(model_manager._qa_pipeline),
+            "summarization": bool(model_manager._summarization_pipeline)
         },
+        "search_engines": ["Brave", "DuckDuckGo"],
         "endpoints": {
             "chat": "/chat",
             "search": "/search",
@@ -269,6 +388,29 @@ async def root():
         }
     }
 
+@run_in_threadpool
+def run_qa_inference(question: str, context: str) -> str:
+    """Run QA model inference in thread pool"""
+    qa_pipeline = model_manager.get_qa_model()
+    result = qa_pipeline(
+        question=question,
+        context=context[:4000],
+        max_answer_len=200
+    )
+    return result['answer']
+
+@run_in_threadpool
+def run_summarization_inference(text: str) -> str:
+    """Run summarization model inference in thread pool"""
+    summarization_pipeline = model_manager.get_summarization_model()
+    result = summarization_pipeline(
+        text,
+        max_length=150,
+        min_length=30,
+        do_sample=False
+    )
+    return result[0]['summary_text']
+
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)