diff --git a/requirements.txt b/requirements.txt index b859ad08a..2f806b315 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,6 +5,7 @@ langchain==0.3.27 langchain_community==0.3.29 msal==1.31.0 openai==1.66.2 +litellm>=1.80,<1.87 Pillow==11.3.0 PyPDF2==3.0.1; sys_platform == 'win32' protobuf==4.25.8 diff --git a/ufo/llm/base.py b/ufo/llm/base.py index be81960d5..501dd23c4 100644 --- a/ufo/llm/base.py +++ b/ufo/llm/base.py @@ -42,6 +42,7 @@ def get_service( "custom": "CustomService", "operator": "OperatorServicePreview", "placeholder": "PlaceHolderService", + "litellm": "LiteLLMService", } custom_service_map = { "llava": "LlavaService", @@ -145,6 +146,8 @@ def get_cost_estimator( name = str("gemini/" + model) elif api_type.lower() == "claude": name = str("claude/" + model) + elif api_type.lower() == "litellm": + name = model else: name = model diff --git a/ufo/llm/litellm.py b/ufo/llm/litellm.py new file mode 100644 index 000000000..9a98321c5 --- /dev/null +++ b/ufo/llm/litellm.py @@ -0,0 +1,133 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import logging +import time +from typing import Any, Dict, List, Optional, Tuple + +from ufo.llm.base import BaseService + +logger = logging.getLogger(__name__) + + +class LiteLLMService(BaseService): + """ + A service class for LiteLLM — access 100+ LLM providers + (OpenAI, Anthropic, Google, Azure, Bedrock, Cohere, Mistral, etc.) + through a single unified interface. + + Models are specified with provider-prefixed names, e.g.: + anthropic/claude-sonnet-4-6 + openai/gpt-4o + gemini/gemini-2.5-pro + bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0 + """ + + def __init__(self, config: Dict[str, Any], agent_type: str): + """ + Initialize the LiteLLM service. + :param config: The configuration dictionary. + :param agent_type: The agent type. + """ + self.config_llm = config[agent_type] + self.config = config + self.model = self.config_llm["API_MODEL"] + self.prices = self.config.get("PRICES", {}) + self.max_retry = self.config.get("MAX_RETRY", 3) + self.api_type = self.config_llm["API_TYPE"].lower() + + self.api_key = self.config_llm.get("API_KEY", "") + self.api_base = self.config_llm.get("API_BASE", "") + + try: + import litellm # noqa: F401 + except ImportError: + raise ImportError( + "litellm is not installed. " + "Install with: pip install litellm>=1.80,<1.87" + ) + + def chat_completion( + self, + messages: List[Dict[str, str]], + n: int = 1, + temperature: Optional[float] = None, + max_tokens: Optional[int] = None, + top_p: Optional[float] = None, + **kwargs: Any, + ) -> Tuple[List[str], Optional[float]]: + """ + Generates completions using LiteLLM (100+ providers). + :param messages: The list of messages in the conversation. + :param n: The number of completions to generate. + :param temperature: The temperature parameter for randomness. + :param max_tokens: The maximum number of tokens in the response. + :param top_p: The top-p parameter for nucleus sampling. + :param kwargs: Additional keyword arguments. + :return: A tuple of (list of response strings, estimated cost). + """ + import litellm + + temperature = ( + temperature if temperature is not None else self.config.get("TEMPERATURE", 0.0) + ) + max_tokens = ( + max_tokens if max_tokens is not None else self.config.get("MAX_TOKENS", 2000) + ) + + responses = [] + cost = 0.0 + + call_kwargs = { + "model": self.model, + "messages": messages, + "temperature": temperature, + "max_tokens": max_tokens, + "drop_params": True, + **kwargs, + } + if self.api_key: + call_kwargs["api_key"] = self.api_key + if self.api_base: + call_kwargs["api_base"] = self.api_base + + for _ in range(n): + for attempt in range(self.max_retry): + try: + response = litellm.completion(**call_kwargs) + if not response.choices: + raise ValueError("LiteLLM returned empty choices list") + content = response.choices[0].message.content or "" + responses.append(content) + + usage = getattr(response, "usage", None) + prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0 + completion_tokens = getattr(usage, "completion_tokens", 0) or 0 + + cost += self.get_cost_estimator( + self.api_type, + self.model, + self.prices, + prompt_tokens, + completion_tokens, + ) + break + except Exception as e: + qualname = f"{type(e).__module__}.{type(e).__name__}" + is_transient = qualname in { + "litellm.exceptions.RateLimitError", + "litellm.exceptions.APIConnectionError", + "litellm.exceptions.Timeout", + "litellm.exceptions.InternalServerError", + "litellm.exceptions.ServiceUnavailableError", + "litellm.exceptions.BadGatewayError", + } + if not is_transient or attempt >= self.max_retry - 1: + logger.error(f"LiteLLM API error: {e}") + raise + logger.warning( + f"LiteLLM transient error (attempt {attempt + 1}/{self.max_retry}): {e}" + ) + time.sleep(2**attempt) + + return responses, cost