MCP RAG

"""Embedding model management for MCP-RAG.""" import logging from abc import ABC, abstractmethod from typing import List, Optional from .config import settings logger = logging.getLogger(__name__) try: import httpx HTTPX_AVAILABLE = True except ImportError: HTTPX_AVAILABLE = False logger.warning("httpx not available, Doubao embedding will not work") try: from sentence_transformers import SentenceTransformer SENTENCE_TRANSFORMERS_AVAILABLE = True except ImportError: SENTENCE_TRANSFORMERS_AVAILABLE = False SentenceTransformer = None logger.warning("sentence-transformers not available, local embedding models will not work. Install with: pip install mcp-rag[local-embeddings]") DOUBAO_AVAILABLE = HTTPX_AVAILABLE class EmbeddingModel(ABC): """Abstract base class for embedding models.""" @abstractmethod async def encode(self, texts: List[str]) -> List[List[float]]: """Encode texts to embeddings.""" pass @abstractmethod async def encode_single(self, text: str) -> List[float]: """Encode single text to embedding.""" pass class SentenceTransformerModel(EmbeddingModel): """SentenceTransformer-based embedding model.""" def __init__(self, model_name: str = "m3e-small", device: str = "cpu", cache_dir: Optional[str] = None): if not SENTENCE_TRANSFORMERS_AVAILABLE: raise RuntimeError("sentence-transformers not available. Install with: pip install mcp-rag[local-embeddings]") self.model_name = model_name self.device = device self.cache_dir = cache_dir self.model = None async def initialize(self) -> None: """Initialize the embedding model.""" try: # Map model names to actual model identifiers model_mapping = { "m3e-small": "moka-ai/m3e-small", "e5-small": "intfloat/e5-small-v2" } actual_model_name = model_mapping.get(self.model_name, self.model_name) self.model = SentenceTransformer( actual_model_name, device=self.device, cache_folder=self.cache_dir ) logger.info(f"Initialized embedding model: {actual_model_name} on {self.device}") except Exception as e: logger.error(f"Failed to initialize embedding model {self.model_name}: {e}") raise async def encode(self, texts: List[str]) -> List[List[float]]: """Encode multiple texts to embeddings.""" if not self.model: raise RuntimeError("Model not initialized") try: embeddings = self.model.encode(texts, convert_to_numpy=True).tolist() return embeddings except Exception as e: logger.error(f"Failed to encode texts: {e}") raise async def encode_single(self, text: str) -> List[float]: """Encode single text to embedding.""" embeddings = await self.encode([text]) return embeddings[0] class DoubaoEmbeddingModel(EmbeddingModel): """Doubao (豆包) embedding model.""" def __init__(self, api_key: Optional[str] = None, base_url: str = "https://ark.cn-beijing.volces.com/api/v3", model: str = "doubao-embedding-text-240715"): if not DOUBAO_AVAILABLE: raise RuntimeError("httpx not available. Please install it with: pip install httpx") self.api_key = api_key self.base_url = base_url.rstrip('/') self.model = model self.client: Optional[httpx.AsyncClient] = None async def initialize(self) -> None: """Initialize the Doubao client.""" try: if not self.api_key: raise ValueError("API key is required for Doubao embedding") self.client = httpx.AsyncClient( base_url=self.base_url, headers={ "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json" }, timeout=30.0 ) logger.info(f"Initialized Doubao embedding model: {self.model}") except Exception as e: logger.error(f"Failed to initialize Doubao embedding model: {e}") raise async def encode(self, texts: List[str]) -> List[List[float]]: """Encode multiple texts to embeddings.""" if not self.client: raise RuntimeError("Client not initialized") try: response = await self.client.post( "/embeddings", json={ "model": self.model, "input": texts, "encoding_format": "float" } ) if response.status_code != 200: raise RuntimeError(f"Doubao API error: {response.status_code} - {response.text}") data = response.json() embeddings = [item["embedding"] for item in data["data"]] return embeddings except Exception as e: logger.error(f"Failed to encode texts with Doubao: {e}") raise async def encode_single(self, text: str) -> List[float]: """Encode single text to embedding.""" embeddings = await self.encode([text]) return embeddings[0] async def close(self) -> None: """Close the HTTP client.""" if self.client: await self.client.aclose() # Global embedding model instance embedding_model: Optional[EmbeddingModel] = None async def get_embedding_model() -> EmbeddingModel: """Get the global embedding model instance.""" global embedding_model if embedding_model is None: if settings.embedding_provider == "doubao": embedding_model = DoubaoEmbeddingModel( api_key=settings.embedding_api_key, base_url=settings.embedding_base_url, model=settings.embedding_model ) else: # Local SentenceTransformer models embedding_model = SentenceTransformerModel( model_name=settings.embedding_model, device=settings.embedding_device, cache_dir=settings.embedding_cache_dir ) await embedding_model.initialize() return embedding_model