| import time |
| from typing import Any |
|
|
| from src.server.config.logfire_config import get_logger |
|
|
| from .capability_tester import get_model_details_logic |
| from .models import OllamaModel |
|
|
| logger = get_logger(__name__) |
|
|
| |
| EMBEDDING_PATTERNS = ( |
| "embed", |
| "embedding", |
| "bge-", |
| "e5-", |
| "sentence-", |
| "arctic-embed", |
| "nomic-embed", |
| "mxbai-embed", |
| "snowflake-arctic-embed", |
| "gte-", |
| "stella-", |
| ) |
|
|
| CHAT_PATTERNS = ( |
| "phi", |
| "qwen", |
| "llama", |
| "mistral", |
| "gemma", |
| "deepseek", |
| "codellama", |
| "orca", |
| "vicuna", |
| "wizardlm", |
| "solar", |
| "mixtral", |
| "chatglm", |
| "baichuan", |
| "yi", |
| "zephyr", |
| "openchat", |
| "starling", |
| "nous-hermes", |
| ) |
|
|
| FUNCTION_CALLING_PATTERNS = ("qwen", "llama3", "phi3", "mistral") |
| STRUCTURED_OUTPUT_PATTERNS = ("llama", "phi", "gemma") |
| UNKNOWN_EMBEDDING_PATTERNS = ("embed", "embedding", "vector") |
|
|
|
|
| async def enrich_model_capabilities_logic( |
| service_instance, |
| models: list[OllamaModel], |
| instance_url: str, |
| fetch_details: bool = False, |
| ) -> list[OllamaModel]: |
| """Pattern-match and enrich model capabilities from Ollama API.""" |
| start_time = time.time() |
| enriched_models = [] |
| unknown_models = [] |
|
|
| for model in models: |
| model_name_lower = model.name.lower() |
|
|
| is_embedding_model = any(pattern in model_name_lower for pattern in EMBEDDING_PATTERNS) |
|
|
| if is_embedding_model: |
| model.capabilities = ["embedding"] |
| if "nomic" in model_name_lower: |
| model.embedding_dimensions = 768 |
| elif "bge" in model_name_lower: |
| model.embedding_dimensions = 1024 if "large" in model_name_lower else 768 |
| elif "e5" in model_name_lower: |
| model.embedding_dimensions = 1024 if "large" in model_name_lower else 768 |
| elif "arctic" in model_name_lower: |
| model.embedding_dimensions = 1024 |
| else: |
| model.embedding_dimensions = 768 |
| logger.debug(f"Pattern-matched embedding model {model.name}") |
| enriched_models.append(model) |
| else: |
| if any(pattern in model_name_lower for pattern in CHAT_PATTERNS): |
| model.capabilities = ["chat"] |
| if any(p in model_name_lower for p in FUNCTION_CALLING_PATTERNS): |
| model.capabilities.extend(["function_calling", "structured_output"]) |
| elif any(p in model_name_lower for p in STRUCTURED_OUTPUT_PATTERNS): |
| model.capabilities.append("structured_output") |
|
|
| if fetch_details: |
| try: |
| detailed_info = await get_model_details_logic(model.name, instance_url) |
| if detailed_info: |
| _map_details_to_model(model, detailed_info) |
| except Exception as e: |
| logger.debug(f"Could not get details for {model.name}: {e}") |
| enriched_models.append(model) |
| else: |
| unknown_models.append(model) |
|
|
| if unknown_models: |
| for model in unknown_models: |
| model.capabilities = ["chat"] |
| model_name_lower = model.name.lower() |
| if any(h in model_name_lower for h in UNKNOWN_EMBEDDING_PATTERNS): |
| model.capabilities = ["embedding"] |
| model.embedding_dimensions = 768 |
| enriched_models.append(model) |
|
|
| logger.info(f"Model enrichment complete for {instance_url} in {time.time() - start_time:.2f}s") |
| return enriched_models |
|
|
|
|
| def _map_details_to_model(model: OllamaModel, info: dict[str, Any]): |
| """Internal helper to map /api/show dict to OllamaModel object.""" |
| model.context_window = info.get("context_window") |
| model.max_context_length = info.get("max_context_length") |
| model.base_context_length = info.get("base_context_length") |
| model.custom_context_length = info.get("custom_context_length") |
| model.architecture = info.get("architecture") |
| model.block_count = info.get("block_count") |
| model.attention_heads = info.get("attention_heads") |
| model.format = info.get("format") |
| model.parent_model = info.get("parent_model") |
| model.family = info.get("family") |
| model.parameter_size = info.get("parameter_size") |
| model.quantization = info.get("quantization") |
| model.parameter_count = info.get("parameter_count") |
| model.file_type = info.get("file_type") |
| model.quantization_version = info.get("quantization_version") |
| model.basename = info.get("basename") |
| model.size_label = info.get("size_label") |
| model.license = info.get("license") |
| model.finetune = info.get("finetune") |
| model.embedding_dimension = info.get("embedding_dimension") |
|
|
| api_caps = info.get("capabilities", []) |
| if api_caps: |
| model.capabilities = list(set(model.capabilities + api_caps)) |
|
|
| if info.get("parameters"): |
| if model.parameters: |
| model.parameters.update(info["parameters"]) |
| else: |
| model.parameters = info["parameters"] |
|
|
