embedding_service.py

"""
Embedding Service for Wikipedia API
Uses Ollama with Phi model (2560 dimensions) for semantic search
"""
import requests
import numpy as np
from typing import List, Optional
import logging

logger = logging.getLogger(__name__)

class EmbeddingService:
    """Service for generating embeddings using Ollama Phi model"""

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

    def check_availability(self) -> bool:
        """Check if Ollama service is available"""
        try:
            response = requests.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 get_embedding(self, text: str) -> Optional[List[float]]:
        """
        Generate embedding for text using Ollama Phi model

        Args:
            text: Input text to embed

        Returns:
            List of floats (2560 dimensions) or None if failed
        """
        try:
            response = requests.post(
                f"{self.ollama_host}/api/embeddings",
                json={"model": self.model, "prompt": text},
                timeout=10
            )

            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)} (expected {self.embedding_dimension})")
                    return None

                return embedding
            else:
                logger.error(f"Ollama embedding failed: HTTP {response.status_code}")
                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

        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"

            response = requests.post(
                vector_search_url,
                json={
                    "query_vector": query_embedding,
                    "limit": limit
                },
                timeout=5
            )

            if response.status_code == 200:
                return response.json()
            else:
                logger.error(f"Vector search failed: HTTP {response.status_code}")
                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

        Args:
            query_text: Search query
            limit: Number of results

        Returns:
            Search results with metadata
        """
        # Generate query embedding
        embedding = self.get_embedding(query_text)
        if not embedding:
            return None

        # Search vector index
        results = self.search_vector_index(embedding, limit)
        if not results:
            return None

        # Add metadata
        results['search_type'] = 'semantic'
        results['embedding_model'] = self.model
        results['embedding_dimension'] = self.embedding_dimension

        return results


# Global embedding service instance
_embedding_service = None

def get_embedding_service() -> Optional[EmbeddingService]:
    """Get or create global embedding service instance"""
    global _embedding_service

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

    return _embedding_service