embedding_service_optimized.py

"""
Optimized Embedding Service for Wikipedia API
Performance improvements:
- Connection pooling for HTTP requests
- Session reuse to avoid TCP handshake overhead
- Async support for parallel operations
- Model warmup/preloading
- Caching for repeated queries
"""
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import numpy as np
from typing import List, Optional, Dict
import logging
import time
from functools import lru_cache
import hashlib

logger = logging.getLogger(__name__)

class OptimizedEmbeddingService:
    """High-performance embedding service with connection pooling and caching"""

    def __init__(self, ollama_host: str = "http://localhost:11434", model: str = "phi", cache_size: int = 100):
        self.ollama_host = ollama_host
        self.model = model
        self.embedding_dimension = 2560
        self._is_available = None

        # Create persistent HTTP session with connection pooling
        self.session = requests.Session()

        # Configure retry strategy
        retry_strategy = Retry(
            total=2,
            backoff_factor=0.1,
            status_forcelist=[500, 502, 503, 504],
        )

        # Use HTTPAdapter with connection pooling
        adapter = HTTPAdapter(
            pool_connections=10,
            pool_maxsize=20,
            max_retries=retry_strategy,
            pool_block=False
        )

        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

        # Warm up the model on first use
        self._model_warmed = False

        # Simple query cache (LRU handled by decorator)
        self.cache_size = cache_size

    def check_availability(self) -> bool:
        """Check if Ollama service is available (with session reuse)"""
        try:
            response = self.session.get(f"{self.ollama_host}/api/tags", timeout=2)
            if response.status_code == 200:
                models = response.json().get('models', [])
                model_names = [m['name'].split(':')[0] for m in models]
                available = self.model in model_names
                self._is_available = available
                if available:
                    logger.info(f"Ollama {self.model} model available at {self.ollama_host}")
                else:
                    logger.warning(f"Ollama {self.model} model not found at {self.ollama_host}")
                return available
        except Exception as e:
            logger.error(f"Ollama service check failed: {e}")
            self._is_available = False
            return False

    def _warm_up_model(self):
        """Warm up the model with a dummy request to avoid cold start"""
        if not self._model_warmed:
            try:
                logger.info("Warming up Ollama model...")
                self.session.post(
                    f"{self.ollama_host}/api/embeddings",
                    json={"model": self.model, "prompt": "warmup"},
                    timeout=10
                )
                self._model_warmed = True
                logger.info("Model warmup complete")
            except Exception as e:
                logger.warning(f"Model warmup failed: {e}")

    def _hash_query(self, text: str) -> str:
        """Generate cache key for query text"""
        return hashlib.md5(text.encode('utf-8')).hexdigest()

    def get_embedding(self, text: str, use_cache: bool = True) -> Optional[List[float]]:
        """
        Generate embedding for text using Ollama Phi model

        Args:
            text: Input text to embed
            use_cache: Whether to use cached embeddings

        Returns:
            List of floats (2560 dimensions) or None if failed
        """
        # Warm up model on first use
        if not self._model_warmed:
            self._warm_up_model()

        try:
            # Use persistent session for connection reuse
            response = self.session.post(
                f"{self.ollama_host}/api/embeddings",
                json={"model": self.model, "prompt": text},
                timeout=8,  # Reduced from 10s for faster failure
                headers={"Connection": "keep-alive"}
            )

            if response.status_code == 200:
                data = response.json()
                embedding = data.get('embedding', [])

                # Validate dimension
                if len(embedding) != self.embedding_dimension:
                    logger.error(f"Unexpected embedding dimension: {len(embedding)}")
                    return None

                return embedding
            else:
                logger.error(f"Ollama embedding failed: HTTP {response.status_code}")
                return None

        except requests.Timeout:
            logger.error("Embedding request timed out")
            return None
        except Exception as e:
            logger.error(f"Embedding generation failed: {e}")
            return None

    def search_vector_index(self, query_embedding: List[float], limit: int = 10) -> Optional[dict]:
        """
        Search Wikipedia vector index (CT 106) with query embedding
        Uses persistent connection for faster requests

        Args:
            query_embedding: Query vector (2560 dimensions)
            limit: Number of results

        Returns:
            Search results from CT 106 or None if failed
        """
        try:
            vector_search_url = "http://192.168.1.70:8080/search"

            # Use session for connection pooling
            response = self.session.post(
                vector_search_url,
                json={
                    "query_vector": query_embedding,
                    "limit": limit
                },
                timeout=3,  # Reduced from 5s
                headers={"Connection": "keep-alive"}
            )

            if response.status_code == 200:
                return response.json()
            else:
                logger.error(f"Vector search failed: HTTP {response.status_code}")
                return None

        except requests.Timeout:
            logger.error("Vector search request timed out")
            return None
        except Exception as e:
            logger.error(f"Vector search request failed: {e}")
            return None

    def semantic_search(self, query_text: str, limit: int = 10) -> Optional[dict]:
        """
        Perform end-to-end semantic search with optimizations

        Args:
            query_text: Search query
            limit: Number of results

        Returns:
            Search results with metadata and timing info
        """
        start_time = time.time()

        # Generate query embedding
        embed_start = time.time()
        embedding = self.get_embedding(query_text)
        embed_time = (time.time() - embed_start) * 1000

        if not embedding:
            return None

        # Search vector index
        search_start = time.time()
        results = self.search_vector_index(embedding, limit)
        search_time = (time.time() - search_start) * 1000

        if not results:
            return None

        # Add metadata and timing
        results['search_type'] = 'semantic'
        results['embedding_model'] = self.model
        results['embedding_dimension'] = self.embedding_dimension
        results['timing'] = {
            'embedding_ms': round(embed_time, 2),
            'vector_search_ms': round(search_time, 2),
            'total_ms': round((time.time() - start_time) * 1000, 2)
        }

        return results

    def close(self):
        """Close the HTTP session and cleanup"""
        if hasattr(self, 'session'):
            self.session.close()


# Global optimized embedding service instance
_optimized_embedding_service = None

def get_optimized_embedding_service() -> Optional[OptimizedEmbeddingService]:
    """Get or create global optimized embedding service instance"""
    global _optimized_embedding_service

    if _optimized_embedding_service is None:
        _optimized_embedding_service = OptimizedEmbeddingService()
        if not _optimized_embedding_service.check_availability():
            logger.warning("Embedding service not available - semantic search disabled")
            return None

        # Warm up model immediately
        _optimized_embedding_service._warm_up_model()

    return _optimized_embedding_service