""" Ollama Embedding Router Provides intelligent routing for embeddings based on model capabilities and dimensions. Integrates with ModelDiscoveryService for real-time dimension detection and supports automatic fallback strategies for optimal performance across distributed Ollama instances. """ from dataclasses import dataclass from typing import Any from ...config.logfire_config import get_logger from .model_discovery_service import model_discovery_service from .routing.fallback_strategy import FallbackStrategy from .routing.vector_normalization import VectorNormalization logger = get_logger(__name__) @dataclass class RoutingDecision: """Represents a routing decision for embedding generation.""" target_column: str model_name: str instance_url: str dimensions: int confidence: float # 0.0 to 1.0 fallback_applied: bool = False routing_strategy: str = "auto-detect" # auto-detect, model-mapping, fallback @dataclass class EmbeddingRoute: """Configuration for embedding routing.""" model_name: str instance_url: str dimensions: int column_name: str performance_score: float = 1.0 # Higher is better class EmbeddingRouter: """ Intelligent router for Ollama embedding operations with dimension-aware routing. Features: - Automatic dimension detection from model capabilities - Intelligent routing to appropriate database columns - Fallback strategies for unknown models - Performance optimization for different vector sizes - Multi-instance load balancing consideration """ def __init__(self): self.routing_cache: dict[str, RoutingDecision] = {} self.cache_ttl = 300 # 5 minutes cache TTL async def route_embedding( self, model_name: str, instance_url: str, text_content: str | None = None ) -> RoutingDecision: """ Determine the optimal routing for an embedding operation. Args: model_name: Name of the embedding model to use instance_url: URL of the Ollama instance text_content: Optional text content for dynamic optimization Returns: RoutingDecision with target column and routing information """ # Check cache first cache_key = f"{model_name}@{instance_url}" if cache_key in self.routing_cache: cached_decision = self.routing_cache[cache_key] logger.debug(f"Using cached routing decision for {model_name}") return cached_decision try: logger.info(f"Determining routing for model {model_name} on {instance_url}") # Step 1: Auto-detect dimensions from model capabilities dimensions = await self._detect_model_dimensions(model_name, instance_url) if dimensions: # Step 2: Route to appropriate column based on detected dimensions decision = await self._route_by_dimensions(model_name, instance_url, dimensions, strategy="auto-detect") logger.info(f"Auto-detected routing: {model_name} -> {decision.target_column} ({dimensions}D)") else: # Step 3: Fallback to model name mapping decision = await self._route_by_model_mapping(model_name, instance_url) logger.warning(f"Fallback routing applied for {model_name} -> {decision.target_column}") # Cache the decision self.routing_cache[cache_key] = decision return decision except Exception as e: logger.error(f"Error routing embedding for {model_name}: {e}") # Emergency fallback to largest supported dimension return RoutingDecision( target_column="embedding_3072", model_name=model_name, instance_url=instance_url, dimensions=3072, confidence=0.1, fallback_applied=True, routing_strategy="emergency-fallback", ) async def _detect_model_dimensions(self, model_name: str, instance_url: str) -> int | None: """ Detect embedding dimensions using the ModelDiscoveryService. Args: model_name: Name of the model instance_url: Ollama instance URL Returns: Detected dimensions or None if detection failed """ try: # Get model info from discovery service model_info = await model_discovery_service.get_model_info(model_name, instance_url) if model_info and model_info.embedding_dimensions: dimensions = model_info.embedding_dimensions logger.debug(f"Detected {dimensions} dimensions for {model_name}") return dimensions # Try capability detection if model info doesn't have dimensions capabilities = await model_discovery_service._detect_model_capabilities(model_name, instance_url) if capabilities.embedding_dimensions: dimensions = capabilities.embedding_dimensions logger.debug(f"Detected {dimensions} dimensions via capabilities for {model_name}") return dimensions logger.warning(f"Could not detect dimensions for {model_name}") return None except Exception as e: logger.error(f"Error detecting dimensions for {model_name}: {e}") return None async def _route_by_dimensions( self, model_name: str, instance_url: str, dimensions: int, strategy: str ) -> RoutingDecision: """ Route embedding based on detected dimensions. Args: model_name: Name of the model instance_url: Ollama instance URL dimensions: Detected embedding dimensions strategy: Routing strategy used Returns: RoutingDecision for the detected dimensions """ # Get target column for dimensions target_column = VectorNormalization.get_target_column(dimensions) # Calculate confidence based on exact dimension match confidence = 1.0 if dimensions in VectorNormalization.DIMENSION_COLUMNS else 0.7 # Check if fallback was applied fallback_applied = dimensions not in VectorNormalization.DIMENSION_COLUMNS if fallback_applied: logger.warning( f"Model {model_name} dimensions {dimensions} not directly supported, " f"using {target_column} with padding/truncation" ) return RoutingDecision( target_column=target_column, model_name=model_name, instance_url=instance_url, dimensions=dimensions, confidence=confidence, fallback_applied=fallback_applied, routing_strategy=strategy, ) async def _route_by_model_mapping(self, model_name: str, instance_url: str) -> RoutingDecision: """ Route embedding based on model name mapping when auto-detection fails. Args: model_name: Name of the model instance_url: Ollama instance URL Returns: RoutingDecision based on model name mapping """ # Use the existing fallback strategy for model mapping dimensions, target_column = FallbackStrategy.get_fallback_dimensions_and_column(model_name) return RoutingDecision( target_column=target_column, model_name=model_name, instance_url=instance_url, dimensions=dimensions, confidence=0.8, # Medium confidence for model mapping fallback_applied=True, routing_strategy="model-mapping", ) def get_optimal_index_type(self, dimensions: int) -> str: """ Get the optimal index type for the given dimensions. Args: dimensions: Embedding dimensions Returns: Recommended index type (ivfflat or hnsw) """ return VectorNormalization.get_optimal_index_type(dimensions) async def get_available_embedding_routes(self, instance_urls: list[str]) -> list[EmbeddingRoute]: """ Get all available embedding routes across multiple instances. Args: instance_urls: List of Ollama instance URLs to check Returns: List of available embedding routes with performance scores """ routes = [] try: # Discover models from all instances discovery_result = await model_discovery_service.discover_models_from_multiple_instances(instance_urls) # Process embedding models for embedding_model in discovery_result["embedding_models"]: model_name = embedding_model["name"] instance_url = embedding_model["instance_url"] dimensions = embedding_model.get("dimensions") if dimensions: target_column = VectorNormalization.get_target_column(dimensions) # Calculate performance score based on dimension efficiency performance_score = VectorNormalization.calculate_performance_score(dimensions) route = EmbeddingRoute( model_name=model_name, instance_url=instance_url, dimensions=dimensions, column_name=target_column, performance_score=performance_score, ) routes.append(route) # Sort by performance score (highest first) routes.sort(key=lambda r: r.performance_score, reverse=True) logger.info(f"Found {len(routes)} embedding routes across {len(instance_urls)} instances") except Exception as e: logger.error(f"Error getting embedding routes: {e}") return routes async def validate_routing_decision(self, decision: RoutingDecision) -> bool: """ Validate that a routing decision is still valid. Args: decision: RoutingDecision to validate Returns: True if decision is valid, False otherwise """ try: # Check if the model still supports embeddings is_valid = await model_discovery_service.validate_model_capabilities( decision.model_name, decision.instance_url, "embedding" ) if not is_valid: logger.warning(f"Routing decision invalid: {decision.model_name} no longer supports embeddings") # Remove from cache if invalid cache_key = f"{decision.model_name}@{decision.instance_url}" if cache_key in self.routing_cache: del self.routing_cache[cache_key] return is_valid except Exception as e: logger.error(f"Error validating routing decision: {e}") return False def clear_routing_cache(self) -> None: """Clear the routing decision cache.""" self.routing_cache.clear() logger.info("Routing cache cleared") def get_routing_statistics(self) -> dict[str, Any]: """ Get statistics about current routing decisions. Returns: Dictionary with routing statistics """ # Use explicit counters with proper types auto_detect_routes = 0 model_mapping_routes = 0 fallback_routes = 0 dimension_distribution: dict[str, int] = {} confidence_high = 0 confidence_medium = 0 confidence_low = 0 for decision in self.routing_cache.values(): # Count routing strategies if decision.routing_strategy == "auto-detect": auto_detect_routes += 1 elif decision.routing_strategy == "model-mapping": model_mapping_routes += 1 else: fallback_routes += 1 # Count dimensions dim_key = f"{decision.dimensions}D" dimension_distribution[dim_key] = dimension_distribution.get(dim_key, 0) + 1 # Count confidence levels if decision.confidence >= 0.9: confidence_high += 1 elif decision.confidence >= 0.7: confidence_medium += 1 else: confidence_low += 1 return { "total_cached_routes": len(self.routing_cache), "auto_detect_routes": auto_detect_routes, "model_mapping_routes": model_mapping_routes, "fallback_routes": fallback_routes, "dimension_distribution": dimension_distribution, "confidence_distribution": {"high": confidence_high, "medium": confidence_medium, "low": confidence_low}, } # Global service instance embedding_router = EmbeddingRouter()