From 7a34090175335bccfb6f25f0e74166e794eb90a3 Mon Sep 17 00:00:00 2001 From: jayy-77 <1427jay@gmail.com> Date: Sun, 25 Jan 2026 23:30:47 +0530 Subject: [PATCH] add model inference module for self-hosted providers --- litellm/model_inference.py | 263 +++++++++++++++++++++++++++++++++++++ 1 file changed, 263 insertions(+) create mode 100644 litellm/model_inference.py diff --git a/litellm/model_inference.py b/litellm/model_inference.py new file mode 100644 index 00000000000..eedab1cef11 --- /dev/null +++ b/litellm/model_inference.py @@ -0,0 +1,263 @@ +""" +Model inference for self-hosted providers. + +When using self-hosted providers (hosted_vllm, openai_like, ollama, etc.) with custom model names, +this module infers model capabilities from similar models already in the registry. + +Example: + hosted_vllm/my-custom-llama-70b → infers from known llama-70b variants + openai_like/mistral-7b-custom → infers from known mistral-7b variants +""" +from typing import Dict, List, Optional, Any +import re +from litellm._logging import verbose_logger +import litellm + + +# Self-hosted providers that support model inference +SELF_HOSTED_PROVIDERS = [ + "hosted_vllm", + "openai_like", + "ollama", + "ollama_chat", + "lm_studio", + "llamafile", +] + + +def extract_base_model_patterns(model_name: str) -> List[str]: + """ + Extract potential base model patterns from a model name. + + Args: + model_name: Original model name (e.g., "my-custom-llama-3.1-70b-instruct") + + Returns: + List of patterns to search for, from most specific to least specific + + Example: + "my-custom-llama-3.1-70b-instruct" → [ + "llama-3.1-70b-instruct", + "llama-3.1-70b", + "llama-3.1", + "llama-3", + "llama" + ] + """ + # Clean up the model name + name = model_name.lower().strip() + + # Remove common prefixes + for prefix in ["my-", "custom-", "fine-tuned-", "ft-"]: + if name.startswith(prefix): + name = name[len(prefix):] + + patterns = [] + + # Try to extract model family patterns + # Match patterns like: llama-3.1-70b, mistral-7b, qwen-72b, etc. + + # Pattern 1: family-version-size-variant (e.g., llama-3.1-70b-instruct) + match = re.search(r'(llama|mistral|qwen|yi|phi|gemma|mixtral)[-_]?(\d+\.?\d*)?[-_]?(\d+[bkm])?[-_]?(\w+)?', name, re.IGNORECASE) + + if match: + family = match.group(1).lower() + version = match.group(2) if match.group(2) else None + size = match.group(3).lower() if match.group(3) else None + variant = match.group(4).lower() if match.group(4) else None + + # Build patterns from most specific to least specific + if variant and size and version: + patterns.append(f"{family}-{version}-{size}-{variant}") + if size and version: + patterns.append(f"{family}-{version}-{size}") + if version: + patterns.append(f"{family}-{version}") + if size: + patterns.append(f"{family}-{size}") + patterns.append(family) + else: + # Fallback: just use the name with common suffixes removed + clean_name = re.sub(r'[-_](instruct|chat|base|v\d+|fp16|gguf).*$', '', name, flags=re.IGNORECASE) + if clean_name and clean_name != name: + patterns.append(clean_name) + patterns.append(name) + + verbose_logger.debug(f"Extracted patterns from '{model_name}': {patterns}") + return patterns + + +def find_similar_models(patterns: List[str]) -> List[Dict[str, Any]]: + """ + Search model_cost for models matching the given patterns. + + Args: + patterns: List of patterns to search for + + Returns: + List of matching model info dicts from model_cost + """ + matches = [] + seen_keys = set() + + for pattern in patterns: + pattern_lower = pattern.lower() + + for model_key, model_info in litellm.model_cost.items(): + if model_key in seen_keys: + continue + + model_key_lower = model_key.lower() + + # Check if pattern is in the model key + if pattern_lower in model_key_lower: + # Skip if this is an inferred or not_found marker + if model_info.get("_inferred") or model_info.get("_not_found"): + continue + + matches.append({ + "key": model_key, + "info": model_info, + "pattern": pattern + }) + seen_keys.add(model_key) + + # If we found matches with this pattern, stop searching + # (we want the most specific matches) + if matches: + break + + verbose_logger.debug(f"Found {len(matches)} similar models") + return matches + + +def aggregate_capabilities(matches: List[Dict[str, Any]]) -> Dict[str, Any]: + """ + Aggregate capabilities from multiple matching models. + + Uses MAX for numeric values and ANY for boolean capabilities. + + Args: + matches: List of matching model dicts + + Returns: + Aggregated model info dict + """ + if not matches: + return {} + + aggregated: Dict[str, Any] = { + "input_cost_per_token": 0.0, # Zero out costs for self-hosted + "output_cost_per_token": 0.0, + } + + # Numeric fields - use MAX + numeric_fields = [ + "max_tokens", + "max_input_tokens", + "max_output_tokens", + "output_vector_size", + ] + + # Boolean fields - use ANY (True if any model supports it) + boolean_fields = [ + "supports_system_messages", + "supports_response_schema", + "supports_vision", + "supports_function_calling", + "supports_tool_choice", + "supports_assistant_prefill", + "supports_prompt_caching", + "supports_audio_input", + "supports_audio_output", + "supports_pdf_input", + "supports_web_search", + "supports_url_context", + "supports_reasoning", + "supports_computer_use", + ] + + # Aggregate numeric fields + for field in numeric_fields: + values = [m["info"].get(field) for m in matches if m["info"].get(field) is not None] + if values: + aggregated[field] = max(values) + + # Aggregate boolean fields + for field in boolean_fields: + values = [m["info"].get(field) for m in matches if m["info"].get(field) is not None] + if values: + aggregated[field] = any(values) + + # Use mode from first match + if matches[0]["info"].get("mode"): + aggregated["mode"] = matches[0]["info"]["mode"] + + verbose_logger.debug(f"Aggregated capabilities from {len(matches)} models") + return aggregated + + +def infer_model_capabilities( + model: str, + custom_llm_provider: str +) -> Optional[Dict[str, Any]]: + """ + Infer model capabilities from similar models in the registry. + + This is the main entry point for model inference. + + Args: + model: Model name (e.g., "my-custom-llama-70b") + custom_llm_provider: Provider name (e.g., "hosted_vllm") + + Returns: + Inferred model info dict or None if inference fails + """ + # Check if provider supports inference + if custom_llm_provider not in SELF_HOSTED_PROVIDERS: + return None + + # Check cache first + cache_key = f"{custom_llm_provider}/{model}" + if cache_key in litellm.model_cost: + cached = litellm.model_cost[cache_key] + if cached.get("_not_found"): + verbose_logger.debug(f"Model {cache_key} previously marked as not found") + return None + return cached + + verbose_logger.info(f"Attempting to infer capabilities for {cache_key}") + + # Extract patterns to search for + patterns = extract_base_model_patterns(model) + + # Find similar models + matches = find_similar_models(patterns) + + if not matches: + verbose_logger.info(f"No similar models found for {model}, caching not_found marker") + # Cache negative result + litellm.model_cost[cache_key] = {"_not_found": True} + return None + + # Aggregate capabilities + inferred_info = aggregate_capabilities(matches) + + if not inferred_info: + litellm.model_cost[cache_key] = {"_not_found": True} + return None + + # Add metadata + inferred_info["_inferred"] = True + inferred_info["_inferred_from"] = [m["key"] for m in matches[:3]] # Track sources + inferred_info["litellm_provider"] = custom_llm_provider + inferred_info["key"] = cache_key + + # Cache result + litellm.model_cost[cache_key] = inferred_info + + verbose_logger.info( + f"Successfully inferred capabilities for {cache_key} from {len(matches)} similar models" + ) + + return inferred_info