diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3648dc691ec..0601e7e8455 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4709,6 +4709,7 @@ class StandardLoggingPayloadSetup: custom_pricing: Optional[bool], custom_llm_provider: Optional[str], init_response_obj: Union[Any, BaseModel, dict], + api_base: Optional[str] = None, ) -> StandardLoggingModelInformation: model_cost_name = _select_model_name_for_cost_calc( model=None, @@ -4723,7 +4724,9 @@ class StandardLoggingPayloadSetup: else: try: _model_cost_information = litellm.get_model_info( - model=model_cost_name, custom_llm_provider=custom_llm_provider + model=model_cost_name, + custom_llm_provider=custom_llm_provider, + api_base=api_base, ) model_cost_information = StandardLoggingModelInformation( model_map_key=model_cost_name, @@ -5236,6 +5239,7 @@ def get_standard_logging_object_payload( custom_pricing=custom_pricing, custom_llm_provider=kwargs.get("custom_llm_provider"), init_response_obj=init_response_obj, + api_base=litellm_params.get("api_base"), ) response_cost: float = kwargs.get("response_cost", 0) or 0.0 diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index c4d08c83a2a..ed14b6a3318 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, List, Optional, from httpx._models import Headers, Response import litellm -from litellm._logging import verbose_proxy_logger +from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -223,7 +223,9 @@ class OllamaConfig(BaseConfig): or get_secret_str("OLLAMA_API_KEY") ) - def get_model_info(self, model: str) -> ModelInfoBase: + def get_model_info( + self, model: str, api_base: Optional[str] = None + ) -> ModelInfoBase: """ curl http://localhost:11434/api/show -d '{ "name": "mistral" @@ -231,7 +233,11 @@ class OllamaConfig(BaseConfig): """ if model.startswith("ollama/") or model.startswith("ollama_chat/"): model = model.split("/", 1)[1] - api_base = get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434" + api_base = ( + api_base + or get_secret_str("OLLAMA_API_BASE") + or "http://localhost:11434" + ) api_key = self.get_api_key() headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} @@ -242,8 +248,21 @@ class OllamaConfig(BaseConfig): headers=headers, ) except Exception as e: - raise Exception( - f"OllamaError: Error getting model info for {model}. Set Ollama API Base via `OLLAMA_API_BASE` environment variable. Error: {e}" + verbose_logger.debug( + "OllamaError: Could not get model info for %s from %s. Error: %s", + model, + api_base, + e, + ) + return ModelInfoBase( + key=model, + litellm_provider="ollama", + mode="chat", + input_cost_per_token=0.0, + output_cost_per_token=0.0, + max_tokens=None, + max_input_tokens=None, + max_output_tokens=None, ) model_info = response.json() diff --git a/litellm/policy_templates_backup.json b/litellm/policy_templates_backup.json index bcc19462a86..5c93ec11d45 100644 --- a/litellm/policy_templates_backup.json +++ b/litellm/policy_templates_backup.json @@ -2463,20 +2463,20 @@ "iconColor": "text-red-500", "iconBg": "bg-red-50", "guardrails": [ - "sg-pdpa-pii-identifiers", - "sg-pdpa-contact-information", - "sg-pdpa-financial-data", - "sg-pdpa-business-identifiers", - "sg-pdpa-personal-identifiers", - "sg-pdpa-sensitive-data", - "sg-pdpa-do-not-call", - "sg-pdpa-data-transfer", - "sg-pdpa-profiling-automated-decisions" + "pdpa-sg-pii-identifiers", + "pdpa-sg-contact-information", + "pdpa-sg-financial-data", + "pdpa-sg-business-identifiers", + "pdpa-sg-personal-identifiers", + "pdpa-sg-sensitive-data", + "pdpa-sg-do-not-call", + "pdpa-sg-data-transfer", + "pdpa-sg-profiling-automated-decisions" ], "complexity": "High", "guardrailDefinitions": [ { - "guardrail_name": "sg-pdpa-pii-identifiers", + "guardrail_name": "pdpa-sg-pii-identifiers", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2499,7 +2499,7 @@ } }, { - "guardrail_name": "sg-pdpa-contact-information", + "guardrail_name": "pdpa-sg-contact-information", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2527,7 +2527,7 @@ } }, { - "guardrail_name": "sg-pdpa-financial-data", + "guardrail_name": "pdpa-sg-financial-data", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2550,7 +2550,7 @@ } }, { - "guardrail_name": "sg-pdpa-business-identifiers", + "guardrail_name": "pdpa-sg-business-identifiers", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2568,7 +2568,7 @@ } }, { - "guardrail_name": "sg-pdpa-personal-identifiers", + "guardrail_name": "pdpa-sg-personal-identifiers", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2587,7 +2587,7 @@ } }, { - "guardrail_name": "sg-pdpa-sensitive-data", + "guardrail_name": "pdpa-sg-sensitive-data", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2606,7 +2606,7 @@ } }, { - "guardrail_name": "sg-pdpa-do-not-call", + "guardrail_name": "pdpa-sg-do-not-call", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2625,7 +2625,7 @@ } }, { - "guardrail_name": "sg-pdpa-data-transfer", + "guardrail_name": "pdpa-sg-data-transfer", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2644,7 +2644,7 @@ } }, { - "guardrail_name": "sg-pdpa-profiling-automated-decisions", + "guardrail_name": "pdpa-sg-profiling-automated-decisions", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2667,15 +2667,15 @@ "policy_name": "pdpa-singapore", "description": "Singapore PDPA compliance policy. Covers personal identifier protection (s.13), sensitive data profiling (Advisory Guidelines), Do Not Call Registry (Part IX), overseas data transfers (s.26), and automated profiling (Model AI Governance Framework). Includes regex-based PII detection for NRIC/FIN, phone numbers, postal codes, passports, UEN, and bank accounts.", "guardrails_add": [ - "sg-pdpa-pii-identifiers", - "sg-pdpa-contact-information", - "sg-pdpa-financial-data", - "sg-pdpa-business-identifiers", - "sg-pdpa-personal-identifiers", - "sg-pdpa-sensitive-data", - "sg-pdpa-do-not-call", - "sg-pdpa-data-transfer", - "sg-pdpa-profiling-automated-decisions" + "pdpa-sg-pii-identifiers", + "pdpa-sg-contact-information", + "pdpa-sg-financial-data", + "pdpa-sg-business-identifiers", + "pdpa-sg-personal-identifiers", + "pdpa-sg-sensitive-data", + "pdpa-sg-do-not-call", + "pdpa-sg-data-transfer", + "pdpa-sg-profiling-automated-decisions" ], "guardrails_remove": [] }, @@ -2694,16 +2694,16 @@ "iconColor": "text-blue-600", "iconBg": "bg-blue-50", "guardrails": [ - "sg-mas-fairness-bias", - "sg-mas-transparency-explainability", - "sg-mas-human-oversight", - "sg-mas-data-governance", - "sg-mas-model-security" + "mas-sg-fairness-bias", + "mas-sg-transparency-explainability", + "mas-sg-human-oversight", + "mas-sg-data-governance", + "mas-sg-model-security" ], "complexity": "High", "guardrailDefinitions": [ { - "guardrail_name": "sg-mas-fairness-bias", + "guardrail_name": "mas-sg-fairness-bias", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2722,7 +2722,7 @@ } }, { - "guardrail_name": "sg-mas-transparency-explainability", + "guardrail_name": "mas-sg-transparency-explainability", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2741,7 +2741,7 @@ } }, { - "guardrail_name": "sg-mas-human-oversight", + "guardrail_name": "mas-sg-human-oversight", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2760,7 +2760,7 @@ } }, { - "guardrail_name": "sg-mas-data-governance", + "guardrail_name": "mas-sg-data-governance", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2779,7 +2779,7 @@ } }, { - "guardrail_name": "sg-mas-model-security", + "guardrail_name": "mas-sg-model-security", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2802,11 +2802,11 @@ "policy_name": "mas-ai-risk-management", "description": "Guidelines on Artificial Intelligence Risk Management (MAS) for Financial Institutions alignment. Covers fairness & bias, transparency & explainability, human oversight, data governance, and model security. Aligned with the 2018 FEAT Principles, Project MindForge, and NIST AI RMF.", "guardrails_add": [ - "sg-mas-fairness-bias", - "sg-mas-transparency-explainability", - "sg-mas-human-oversight", - "sg-mas-data-governance", - "sg-mas-model-security" + "mas-sg-fairness-bias", + "mas-sg-transparency-explainability", + "mas-sg-human-oversight", + "mas-sg-data-governance", + "mas-sg-model-security" ], "guardrails_remove": [] }, diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_mas_human_oversight.yaml b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_mas_human_oversight.yaml index 0e9e336dbe9..dfca04f9013 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_mas_human_oversight.yaml +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_mas_human_oversight.yaml @@ -74,16 +74,16 @@ always_block_keywords: exceptions: - "human in the loop" - "human on the loop" - - "add human review" - - "with human review" + - "human review" + - "human oversight" - "manual review" - "escalation" - "appeals process" - "override mechanism" - - "explain what" - - "explain how" + - "explain" - "what is" - "research" - "improve oversight" + - "add human review" - "compliance" - "governance" diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_mas_transparency_explainability.yaml b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_mas_transparency_explainability.yaml index 0a08abd86e8..19b999729bb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_mas_transparency_explainability.yaml +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_mas_transparency_explainability.yaml @@ -73,8 +73,7 @@ exceptions: - "interpretability" - "model card" - "audit trail" - - "explain what" - - "explain how" + - "explain" - "what is" - "SHAP" - "LIME" diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_pdpa_data_transfer.yaml b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_pdpa_data_transfer.yaml index bfa5e843401..18cdde80bcd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_pdpa_data_transfer.yaml +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_pdpa_data_transfer.yaml @@ -16,8 +16,6 @@ identifier_words: - "SG resident data" - "PDPA data" - "Singapore employee data" - # PDPA alone implies Singapore regulatory context - - "PDPA" additional_block_words: # Transfer actions diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_pdpa_profiling_automated_decisions.yaml b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_pdpa_profiling_automated_decisions.yaml index 3d95105bb3a..05578641c15 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_pdpa_profiling_automated_decisions.yaml +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/policy_templates/sg_pdpa_profiling_automated_decisions.yaml @@ -69,7 +69,7 @@ always_block_keywords: exceptions: - "human in the loop" - "human review" - - "with human oversight" + - "human oversight" - "appeals process" - "explainability" - "fairness audit" diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 2c429172940..cb1db9d04ad 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, Type, cast from urllib.parse import urlparse from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import CustomGuardrail, log_guardrail_information +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import ( @@ -237,7 +237,6 @@ class NomaV2Guardrail(CustomGuardrail): return inputs - @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, diff --git a/litellm/utils.py b/litellm/utils.py index fef99c8b201..b4de1e00efc 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5385,14 +5385,18 @@ def _get_max_position_embeddings(model_name: str) -> Optional[int]: @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) def _cached_get_model_info_helper( - model: str, custom_llm_provider: Optional[str] + model: str, + custom_llm_provider: Optional[str], + api_base: Optional[str] = None, ) -> ModelInfoBase: """ _get_model_info_helper wrapped with lru_cache Speed Optimization to hit high RPS """ - return _get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + return _get_model_info_helper( + model=model, custom_llm_provider=custom_llm_provider, api_base=api_base + ) def get_provider_info( @@ -5428,7 +5432,9 @@ def _is_potential_model_name_in_model_cost( def _get_model_info_helper( # noqa: PLR0915 - model: str, custom_llm_provider: Optional[str] = None + model: str, + custom_llm_provider: Optional[str] = None, + api_base: Optional[str] = None, ) -> ModelInfoBase: """ Helper for 'get_model_info'. Separated out to avoid infinite loop caused by returning 'supported_openai_param's @@ -5486,7 +5492,7 @@ def _get_model_info_helper( # noqa: PLR0915 elif ( custom_llm_provider == "ollama" or custom_llm_provider == "ollama_chat" ) and not _is_potential_model_name_in_model_cost(potential_model_names): - return litellm.OllamaConfig().get_model_info(model) + return litellm.OllamaConfig().get_model_info(model, api_base=api_base) else: """ Check if: (in order of specificity) @@ -5725,8 +5731,6 @@ def _get_model_info_helper( # noqa: PLR0915 ) except Exception as e: verbose_logger.debug(f"Error getting model info: {e}") - if "OllamaError" in str(e): - raise e raise Exception( "This model isn't mapped yet. model={}, custom_llm_provider={}. Add it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json.".format( model, custom_llm_provider @@ -5735,7 +5739,11 @@ def _get_model_info_helper( # noqa: PLR0915 @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) -def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> ModelInfo: +def get_model_info( + model: str, + custom_llm_provider: Optional[str] = None, + api_base: Optional[str] = None, +) -> ModelInfo: """ Get a dict for the maximum tokens (context window), input_cost_per_token, output_cost_per_token for a given model. @@ -5813,6 +5821,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod _model_info = _get_model_info_helper( model=model, custom_llm_provider=custom_llm_provider, + api_base=api_base, ) provider_info = get_provider_info( diff --git a/policy_templates.json b/policy_templates.json index da2598f186c..6125650cb34 100644 --- a/policy_templates.json +++ b/policy_templates.json @@ -2022,20 +2022,20 @@ "iconColor": "text-red-500", "iconBg": "bg-red-50", "guardrails": [ - "sg-pdpa-pii-identifiers", - "sg-pdpa-contact-information", - "sg-pdpa-financial-data", - "sg-pdpa-business-identifiers", - "sg-pdpa-personal-identifiers", - "sg-pdpa-sensitive-data", - "sg-pdpa-do-not-call", - "sg-pdpa-data-transfer", - "sg-pdpa-profiling-automated-decisions" + "pdpa-sg-pii-identifiers", + "pdpa-sg-contact-information", + "pdpa-sg-financial-data", + "pdpa-sg-business-identifiers", + "pdpa-sg-personal-identifiers", + "pdpa-sg-sensitive-data", + "pdpa-sg-do-not-call", + "pdpa-sg-data-transfer", + "pdpa-sg-profiling-automated-decisions" ], "complexity": "High", "guardrailDefinitions": [ { - "guardrail_name": "sg-pdpa-pii-identifiers", + "guardrail_name": "pdpa-sg-pii-identifiers", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2058,7 +2058,7 @@ } }, { - "guardrail_name": "sg-pdpa-contact-information", + "guardrail_name": "pdpa-sg-contact-information", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2086,7 +2086,7 @@ } }, { - "guardrail_name": "sg-pdpa-financial-data", + "guardrail_name": "pdpa-sg-financial-data", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2109,7 +2109,7 @@ } }, { - "guardrail_name": "sg-pdpa-business-identifiers", + "guardrail_name": "pdpa-sg-business-identifiers", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2127,7 +2127,7 @@ } }, { - "guardrail_name": "sg-pdpa-personal-identifiers", + "guardrail_name": "pdpa-sg-personal-identifiers", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2146,7 +2146,7 @@ } }, { - "guardrail_name": "sg-pdpa-sensitive-data", + "guardrail_name": "pdpa-sg-sensitive-data", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2165,7 +2165,7 @@ } }, { - "guardrail_name": "sg-pdpa-do-not-call", + "guardrail_name": "pdpa-sg-do-not-call", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2184,7 +2184,7 @@ } }, { - "guardrail_name": "sg-pdpa-data-transfer", + "guardrail_name": "pdpa-sg-data-transfer", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2203,7 +2203,7 @@ } }, { - "guardrail_name": "sg-pdpa-profiling-automated-decisions", + "guardrail_name": "pdpa-sg-profiling-automated-decisions", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2226,15 +2226,15 @@ "policy_name": "pdpa-singapore", "description": "Singapore PDPA compliance policy. Covers personal identifier protection (s.13), sensitive data profiling (Advisory Guidelines), Do Not Call Registry (Part IX), overseas data transfers (s.26), and automated profiling (Model AI Governance Framework). Includes regex-based PII detection for NRIC/FIN, phone numbers, postal codes, passports, UEN, and bank accounts.", "guardrails_add": [ - "sg-pdpa-pii-identifiers", - "sg-pdpa-contact-information", - "sg-pdpa-financial-data", - "sg-pdpa-business-identifiers", - "sg-pdpa-personal-identifiers", - "sg-pdpa-sensitive-data", - "sg-pdpa-do-not-call", - "sg-pdpa-data-transfer", - "sg-pdpa-profiling-automated-decisions" + "pdpa-sg-pii-identifiers", + "pdpa-sg-contact-information", + "pdpa-sg-financial-data", + "pdpa-sg-business-identifiers", + "pdpa-sg-personal-identifiers", + "pdpa-sg-sensitive-data", + "pdpa-sg-do-not-call", + "pdpa-sg-data-transfer", + "pdpa-sg-profiling-automated-decisions" ], "guardrails_remove": [] }, @@ -2253,16 +2253,16 @@ "iconColor": "text-blue-600", "iconBg": "bg-blue-50", "guardrails": [ - "sg-mas-fairness-bias", - "sg-mas-transparency-explainability", - "sg-mas-human-oversight", - "sg-mas-data-governance", - "sg-mas-model-security" + "mas-sg-fairness-bias", + "mas-sg-transparency-explainability", + "mas-sg-human-oversight", + "mas-sg-data-governance", + "mas-sg-model-security" ], "complexity": "High", "guardrailDefinitions": [ { - "guardrail_name": "sg-mas-fairness-bias", + "guardrail_name": "mas-sg-fairness-bias", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2281,7 +2281,7 @@ } }, { - "guardrail_name": "sg-mas-transparency-explainability", + "guardrail_name": "mas-sg-transparency-explainability", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2300,7 +2300,7 @@ } }, { - "guardrail_name": "sg-mas-human-oversight", + "guardrail_name": "mas-sg-human-oversight", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2319,7 +2319,7 @@ } }, { - "guardrail_name": "sg-mas-data-governance", + "guardrail_name": "mas-sg-data-governance", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2338,7 +2338,7 @@ } }, { - "guardrail_name": "sg-mas-model-security", + "guardrail_name": "mas-sg-model-security", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "pre_call", @@ -2361,11 +2361,11 @@ "policy_name": "mas-ai-risk-management", "description": "Guidelines on Artificial Intelligence Risk Management (MAS) for Financial Institutions alignment. Covers fairness & bias, transparency & explainability, human oversight, data governance, and model security. Aligned with the 2018 FEAT Principles, Project MindForge, and NIST AI RMF.", "guardrails_add": [ - "sg-mas-fairness-bias", - "sg-mas-transparency-explainability", - "sg-mas-human-oversight", - "sg-mas-data-governance", - "sg-mas-model-security" + "mas-sg-fairness-bias", + "mas-sg-transparency-explainability", + "mas-sg-human-oversight", + "mas-sg-data-governance", + "mas-sg-model-security" ], "guardrails_remove": [] }, diff --git a/tests/test_litellm/llms/ollama/test_ollama_model_info.py b/tests/test_litellm/llms/ollama/test_ollama_model_info.py index 7eef15cd4d2..5585e9d1e0e 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_model_info.py +++ b/tests/test_litellm/llms/ollama/test_ollama_model_info.py @@ -138,6 +138,88 @@ class TestOllamaModelInfo: assert models == ["ollama/llama2"] +class TestOllamaGetModelInfo: + """Tests for OllamaConfig.get_model_info() api_base threading and graceful fallback.""" + + def test_get_model_info_uses_provided_api_base(self, monkeypatch): + """When api_base is passed, get_model_info should use it instead of env var or default.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + captured_urls = [] + + def mock_post(url, json, headers=None): + captured_urls.append(url) + resp = DummyResponse( + {"template": "{{ .System }} tools {{ .Prompt }}", "model_info": {"context_length": 4096}}, + status_code=200, + ) + return resp + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + + config = OllamaConfig() + result = config.get_model_info("llama3", api_base="http://my-remote-server:11434") + + assert captured_urls[0] == "http://my-remote-server:11434/api/show" + assert result["max_tokens"] == 4096 + + def test_get_model_info_falls_back_to_env_var(self, monkeypatch): + """When no api_base is passed, should fall back to OLLAMA_API_BASE env var.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + captured_urls = [] + + def mock_post(url, json, headers=None): + captured_urls.append(url) + return DummyResponse({"template": "", "model_info": {}}, status_code=200) + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + monkeypatch.setenv("OLLAMA_API_BASE", "http://env-server:11434") + + config = OllamaConfig() + config.get_model_info("llama3") + + assert captured_urls[0] == "http://env-server:11434/api/show" + + def test_get_model_info_graceful_fallback_on_connection_error(self, monkeypatch): + """When the Ollama server is unreachable, should return defaults instead of raising.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + def mock_post(url, json, headers=None): + raise ConnectionError("Connection refused") + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + + config = OllamaConfig() + result = config.get_model_info("llama3", api_base="http://unreachable:11434") + + assert result["key"] == "llama3" + assert result["litellm_provider"] == "ollama" + assert result["input_cost_per_token"] == 0.0 + assert result["output_cost_per_token"] == 0.0 + assert result["max_tokens"] is None + + def test_get_model_info_strips_ollama_prefix(self, monkeypatch): + """Should strip 'ollama/' or 'ollama_chat/' prefix from model name.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + captured_json = [] + + def mock_post(url, json, headers=None): + captured_json.append(json) + return DummyResponse({"template": "", "model_info": {}}, status_code=200) + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + + config = OllamaConfig() + config.get_model_info("ollama/llama3", api_base="http://localhost:11434") + assert captured_json[0]["name"] == "llama3" + + config.get_model_info("ollama_chat/llama3", api_base="http://localhost:11434") + assert captured_json[1]["name"] == "llama3" + + class TestOllamaAuthHeaders: """Tests for Ollama authentication header handling in completion calls."""