From 000a653d9115d1faa3a519eafc5b6e35d6c77fb3 Mon Sep 17 00:00:00 2001 From: Aarish Alam Date: Tue, 5 May 2026 20:00:45 +0530 Subject: [PATCH] feat: add LiteLLM as an embedded AI gateway provider (#3043) --- configs/dbgpt-proxy-litellm.toml | 48 ++ packages/dbgpt-core/pyproject.toml | 1 + .../src/dbgpt/model/proxy/__init__.py | 3 + .../src/dbgpt/model/proxy/llms/litellm.py | 412 ++++++++++++++++++ .../dbgpt/model/proxy/tests/test_litellm.py | 400 +++++++++++++++++ 5 files changed, 864 insertions(+) create mode 100644 configs/dbgpt-proxy-litellm.toml create mode 100644 packages/dbgpt-core/src/dbgpt/model/proxy/llms/litellm.py create mode 100644 packages/dbgpt-core/src/dbgpt/model/proxy/tests/test_litellm.py diff --git a/configs/dbgpt-proxy-litellm.toml b/configs/dbgpt-proxy-litellm.toml new file mode 100644 index 000000000..1939aba58 --- /dev/null +++ b/configs/dbgpt-proxy-litellm.toml @@ -0,0 +1,48 @@ +[system] +# Load language from environment variable(It is set by the hook) +language = "${env:DBGPT_LANG:-en}" +api_keys = [] +encrypt_key = "your_secret_key" + +# Server Configurations +[service.web] +host = "0.0.0.0" +port = 5670 + +[service.web.database] +type = "sqlite" +path = "pilot/meta_data/dbgpt.db" +[service.model.worker] +host = "127.0.0.1" + +[rag.storage] +[rag.storage.vector] +type = "chroma" +persist_path = "pilot/data" + +# Model Configurations +# +# LiteLLM is used here as an embedded AI gateway (the Python SDK), NOT as a +# separate proxy server — DB-GPT imports litellm directly and routes every +# completion through litellm.acompletion(). Specify the model with a provider +# prefix (e.g. "anthropic/...", "vertex_ai/...", "bedrock/...", "azure/...", +# "groq/...") and set the matching provider environment variable +# (ANTHROPIC_API_KEY, OPENAI_API_KEY, AWS_ACCESS_KEY_ID, AZURE_API_KEY, ...). +# See https://docs.litellm.ai/docs/providers for the full list. +[models] +[[models.llms]] +name = "anthropic/claude-3-5-sonnet-20241022" +provider = "proxy/litellm" +# api_key and api_base are usually unnecessary — LiteLLM resolves them per +# provider from environment variables. Override only when using a custom or +# OpenAI-compatible endpoint. +# api_key = "your_anthropic_api_key" +# api_base = "https://api.anthropic.com" + +[[models.embeddings]] +name = "BAAI/bge-large-zh-v1.5" +provider = "hf" +# If not provided, the model will be downloaded from the Hugging Face model hub +# uncomment the following line to specify the model path in the local file system +# path = "the-model-path-in-the-local-file-system" +path = "models/bge-large-zh-v1.5" diff --git a/packages/dbgpt-core/pyproject.toml b/packages/dbgpt-core/pyproject.toml index f37cf6a45..fd920a1a3 100644 --- a/packages/dbgpt-core/pyproject.toml +++ b/packages/dbgpt-core/pyproject.toml @@ -135,6 +135,7 @@ proxy_openai = [ "httpx[socks]", ] proxy_anthropic = ["anthropic"] +proxy_litellm = ["litellm>=1.60,<1.85"] # For vision-language models model_vl = [ diff --git a/packages/dbgpt-core/src/dbgpt/model/proxy/__init__.py b/packages/dbgpt-core/src/dbgpt/model/proxy/__init__.py index 3d7b9810b..b88ff37e4 100644 --- a/packages/dbgpt-core/src/dbgpt/model/proxy/__init__.py +++ b/packages/dbgpt-core/src/dbgpt/model/proxy/__init__.py @@ -11,6 +11,7 @@ if TYPE_CHECKING: from dbgpt.model.proxy.llms.gemini import GeminiLLMClient from dbgpt.model.proxy.llms.gitee import GiteeLLMClient from dbgpt.model.proxy.llms.infiniai import InfiniAILLMClient + from dbgpt.model.proxy.llms.litellm import LiteLLMClient from dbgpt.model.proxy.llms.minimax import MiniMaxLLMClient from dbgpt.model.proxy.llms.moonshot import MoonshotLLMClient from dbgpt.model.proxy.llms.ollama import OllamaLLMClient @@ -40,6 +41,7 @@ def __lazy_import(name): "DeepseekLLMClient": "dbgpt.model.proxy.llms.deepseek", "GiteeLLMClient": "dbgpt.model.proxy.llms.gitee", "InfiniAILLMClient": "dbgpt.model.proxy.llms.infiniai", + "LiteLLMClient": "dbgpt.model.proxy.llms.litellm", "MiniMaxLLMClient": "dbgpt.model.proxy.llms.minimax", } @@ -71,5 +73,6 @@ __all__ = [ "DeepseekLLMClient", "GiteeLLMClient", "InfiniAILLMClient", + "LiteLLMClient", "MiniMaxLLMClient", ] diff --git a/packages/dbgpt-core/src/dbgpt/model/proxy/llms/litellm.py b/packages/dbgpt-core/src/dbgpt/model/proxy/llms/litellm.py new file mode 100644 index 000000000..84d2c546e --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/model/proxy/llms/litellm.py @@ -0,0 +1,412 @@ +import logging +from concurrent.futures import Executor +from dataclasses import dataclass, field +from typing import ( + TYPE_CHECKING, + Any, + AsyncIterator, + Dict, + List, + Optional, + Type, + Union, + cast, +) + +from dbgpt.core import MessageConverter, ModelMetadata, ModelOutput, ModelRequest +from dbgpt.core.awel.flow import ( + TAGS_ORDER_HIGH, + ResourceCategory, + auto_register_resource, +) +from dbgpt.model.proxy.llms.proxy_model import ProxyModel, parse_model_request +from dbgpt.util.i18n_utils import _ + +from ..base import ( + AsyncGenerateStreamFunction, + GenerateStreamFunction, + ProxyLLMClient, + register_proxy_model_adapter, +) +from .chatgpt import OpenAICompatibleDeployModelParameters + +if TYPE_CHECKING: + from httpx._types import ProxiesTypes + +logger = logging.getLogger(__name__) + +_DEFAULT_MODEL = "openai/gpt-4o-mini" + + +@auto_register_resource( + label=_("LiteLLM AI Gateway"), + category=ResourceCategory.LLM_CLIENT, + tags={"order": TAGS_ORDER_HIGH}, + description=_( + "LiteLLM as an embedded AI gateway — call 100+ LLM providers (OpenAI, " + "Anthropic, Vertex AI, Bedrock, Azure, Cohere, Mistral, Groq, Ollama, ...) " + "through a single client without running a separate proxy server. Specify " + "the model with a provider prefix, for example " + "'anthropic/claude-3-5-sonnet-20241022', 'vertex_ai/gemini-1.5-pro', " + "'bedrock/anthropic.claude-3-haiku-20240307-v1:0', or 'azure/gpt-4o'." + ), + documentation_url="https://docs.litellm.ai/docs/providers", + show_in_ui=False, +) +@dataclass +class LiteLLMDeployModelParameters(OpenAICompatibleDeployModelParameters): + """Deploy model parameters for the embedded LiteLLM gateway.""" + + provider: str = "proxy/litellm" + + api_base: Optional[str] = field( + default=None, + metadata={ + "help": _( + "Optional API base URL. LiteLLM resolves the per-provider endpoint " + "from the model prefix and provider-specific environment variables; " + "set this only when calling an OpenAI-compatible custom endpoint." + ), + }, + ) + + api_key: Optional[str] = field( + default=None, + metadata={ + "help": _( + "Optional API key. LiteLLM resolves provider-specific keys from " + "environment variables (OPENAI_API_KEY, ANTHROPIC_API_KEY, " + "AZURE_API_KEY, GROQ_API_KEY, ...) by default; set this only if " + "your model requires a single shared key." + ), + "tags": "privacy", + }, + ) + + +async def litellm_generate_stream( + model: ProxyModel, tokenizer, params, device, context_len=2048 +) -> AsyncIterator[ModelOutput]: + client: LiteLLMClient = cast(LiteLLMClient, model.proxy_llm_client) + request = parse_model_request(params, client.default_model, stream=True) + async for r in client.generate_stream(request): + yield r + + +class LiteLLMClient(ProxyLLMClient): + """Embedded LiteLLM AI gateway client. + + Routes every request through ``litellm.acompletion`` (the LiteLLM Python + SDK), so a single DB-GPT model entry can talk to OpenAI, Anthropic, Vertex + AI, Bedrock, Azure, Cohere, Mistral, Groq, Ollama, and 90+ other providers + *without running a separate LiteLLM proxy server*. The model is selected + via the standard LiteLLM provider-prefixed name (``anthropic/...``, + ``vertex_ai/...``, etc.) and credentials are resolved from provider-specific + environment variables (``ANTHROPIC_API_KEY``, ``OPENAI_API_KEY``, + ``AWS_ACCESS_KEY_ID``, ...). + + ``drop_params=True`` is enabled by default so kwargs that some providers + reject (``frequency_penalty`` / ``presence_penalty`` on Anthropic, Gemini, + Bedrock; ``response_format`` on Bedrock; etc.) are silently dropped instead + of raising ``UnsupportedParamsError``. Override via + ``litellm_kwargs={"drop_params": False}``. + """ + + def __init__( + self, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + api_type: Optional[str] = None, + api_version: Optional[str] = None, + model: Optional[str] = None, + proxies: Optional["ProxiesTypes"] = None, + timeout: Optional[int] = 240, + model_alias: Optional[str] = None, + context_length: Optional[int] = None, + litellm_kwargs: Optional[Dict[str, Any]] = None, + **kwargs, + ): + try: + import litellm # noqa: F401 + except ImportError as exc: + raise ValueError( + "Could not import python package: litellm. " + 'Please install it via `pip install "litellm>=1.60,<1.85"`.' + ) from exc + + if not model: + model = _DEFAULT_MODEL + if not model_alias: + model_alias = model + if not context_length: + context_length = 1024 * 8 + + # drop_params silently strips kwargs that the destination provider does + # not support (e.g., presence_penalty on Anthropic). Defaulted on so + # DB-GPT's generic per-request payload doesn't crash provider-specific + # backends. Users can opt out via litellm_kwargs={"drop_params": False}. + merged_kwargs: Dict[str, Any] = {"drop_params": True} + if litellm_kwargs: + merged_kwargs.update(litellm_kwargs) + + self._model = model + self._api_key = self._resolve_env_vars(api_key) + self._api_base = self._resolve_env_vars(api_base) + self._api_type = self._resolve_env_vars(api_type) + self._api_version = self._resolve_env_vars(api_version) + self._proxies = proxies + self._timeout = timeout + self._model_alias = model_alias + self._context_length = context_length + self._litellm_kwargs = merged_kwargs + + super().__init__(model_names=[model_alias], context_length=context_length) + + @classmethod + def param_class(cls) -> Type[LiteLLMDeployModelParameters]: + """Get the deploy model parameters class.""" + return LiteLLMDeployModelParameters + + @classmethod + def new_client( + cls, + model_params: LiteLLMDeployModelParameters, + default_executor: Optional[Executor] = None, + ) -> "LiteLLMClient": + """Create a new client with the deploy model parameters.""" + return cls( + api_key=model_params.api_key, + api_base=model_params.api_base, + api_type=model_params.api_type, + api_version=model_params.api_version, + model=model_params.real_provider_model_name, + proxies=model_params.http_proxy, + model_alias=model_params.real_provider_model_name, + context_length=max(model_params.context_length or 8192, 8192), + ) + + @classmethod + def generate_stream_function( + cls, + ) -> Optional[Union[GenerateStreamFunction, AsyncGenerateStreamFunction]]: + return litellm_generate_stream + + @property + def default_model(self) -> str: + return self._model or _DEFAULT_MODEL + + def _build_request( + self, request: ModelRequest, stream: Optional[bool] = False + ) -> Dict[str, Any]: + payload: Dict[str, Any] = {"stream": stream} + payload["model"] = request.model or self.default_model + # Provider-specific overrides come last so users can shadow defaults. + for k, v in self._litellm_kwargs.items(): + payload[k] = v + if self._api_key: + payload["api_key"] = self._api_key + if self._api_base: + payload["api_base"] = self._api_base + if self._api_version: + payload["api_version"] = self._api_version + if request.temperature is not None: + payload["temperature"] = request.temperature + if request.max_new_tokens: + payload["max_tokens"] = request.max_new_tokens + if request.stop: + payload["stop"] = request.stop + if request.top_p is not None: + payload["top_p"] = request.top_p + if self._timeout: + payload.setdefault("timeout", self._timeout) + if stream: + # Ask LiteLLM/OpenAI for a final usage chunk so we can report tokens. + payload.setdefault("stream_options", {"include_usage": True}) + return payload + + async def generate( + self, + request: ModelRequest, + message_converter: Optional[MessageConverter] = None, + ) -> ModelOutput: + import litellm + + request = self.local_covert_message(request, message_converter) + messages = request.to_common_messages() + payload = self._build_request(request) + logger.info( + f"Send request to litellm, payload: {payload}\n\nmessages:\n{messages}" + ) + try: + response = await litellm.acompletion(messages=messages, **payload) + message_obj = response.choices[0].message + text = message_obj.content + reasoning_content = getattr(message_obj, "reasoning_content", "") or "" + usage = response.usage.model_dump() if response.usage else None + return ModelOutput.build(text, reasoning_content, usage=usage) + except Exception as e: + return ModelOutput( + text=f"**LLMServer Generate Error, Please CheckErrorInfo.**: {e}", + error_code=1, + ) + + async def generate_stream( + self, + request: ModelRequest, + message_converter: Optional[MessageConverter] = None, + ) -> AsyncIterator[ModelOutput]: + import litellm + + request = self.local_covert_message(request, message_converter) + messages = request.to_common_messages() + payload = self._build_request(request, stream=True) + logger.info( + f"Send request to litellm (stream), payload: {payload}\n\n" + f"messages:\n{messages}" + ) + try: + response = await litellm.acompletion(messages=messages, **payload) + text = "" + reasoning_content = "" + usage: Optional[Dict[str, Any]] = None + async for chunk in response: + # Some providers (Anthropic, Bedrock via LiteLLM) attach usage on + # the final content chunk; OpenAI / Azure with + # stream_options=include_usage emit a trailing usage-only chunk + # with empty choices. Capture either form. + chunk_usage = getattr(chunk, "usage", None) + if chunk_usage is not None: + usage = chunk_usage.model_dump() + choices = getattr(chunk, "choices", None) or [] + if not choices: + continue + choice = choices[0] + if choice is None or choice.delta is None: + continue + delta_obj = choice.delta + new_reasoning = "" + if ( + hasattr(delta_obj, "reasoning_content") + and delta_obj.reasoning_content + ): + new_reasoning = delta_obj.reasoning_content + reasoning_content += new_reasoning + new_content = delta_obj.content if delta_obj.content is not None else "" + if new_content: + text += new_content + # Only yield when this chunk carried new content/reasoning so we + # don't emit duplicate frames on finish-only chunks. + if new_content or new_reasoning: + yield ModelOutput.build(text, reasoning_content, usage=usage) + except Exception as e: + yield ModelOutput( + text=( + f"**LLMServer Generate Stream Error, Please CheckErrorInfo.**: {e}" + ), + error_code=1, + ) + + async def models(self) -> List[ModelMetadata]: + return [ + ModelMetadata( + model=self._model_alias, + context_length=await self.get_context_length(), + ) + ] + + async def get_context_length(self) -> int: + return self._context_length + + +register_proxy_model_adapter( + LiteLLMClient, + supported_models=[ + ModelMetadata( + model=[ + "openai/gpt-4o", + "openai/gpt-4o-mini", + "openai/gpt-4-turbo", + "openai/o1", + "openai/o1-mini", + "openai/o3-mini", + ], + context_length=128000, + max_output_length=16384, + description="OpenAI models routed via LiteLLM", + link="https://docs.litellm.ai/docs/providers/openai", + function_calling=True, + ), + ModelMetadata( + model=[ + "anthropic/claude-3-5-sonnet-20241022", + "anthropic/claude-3-5-sonnet-latest", + "anthropic/claude-3-5-haiku-latest", + "anthropic/claude-3-opus-latest", + "anthropic/claude-3-haiku-20240307", + ], + context_length=200000, + max_output_length=8192, + description="Anthropic Claude routed via LiteLLM", + link="https://docs.litellm.ai/docs/providers/anthropic", + function_calling=True, + ), + ModelMetadata( + model=[ + "vertex_ai/gemini-1.5-pro", + "vertex_ai/gemini-1.5-flash", + "vertex_ai/gemini-2.0-flash", + ], + context_length=1048576, + max_output_length=8192, + description="Google Vertex AI Gemini routed via LiteLLM", + link="https://docs.litellm.ai/docs/providers/vertex", + function_calling=True, + ), + ModelMetadata( + model=[ + "bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0", + "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + "bedrock/anthropic.claude-3-haiku-20240307-v1:0", + ], + context_length=200000, + max_output_length=8192, + description="AWS Bedrock Anthropic Claude routed via LiteLLM", + link="https://docs.litellm.ai/docs/providers/bedrock", + function_calling=True, + ), + ModelMetadata( + model=[ + "azure/gpt-4o", + "azure/gpt-4o-mini", + ], + context_length=128000, + max_output_length=16384, + description="Azure OpenAI routed via LiteLLM", + link="https://docs.litellm.ai/docs/providers/azure", + function_calling=True, + ), + ModelMetadata( + model=[ + "groq/llama-3.3-70b-versatile", + "groq/llama-3.1-70b-versatile", + "groq/mixtral-8x7b-32768", + ], + context_length=131072, + max_output_length=8192, + description="Groq routed via LiteLLM", + link="https://docs.litellm.ai/docs/providers/groq", + function_calling=True, + ), + ModelMetadata( + model=[ + "mistral/mistral-large-latest", + "mistral/mistral-small-latest", + ], + context_length=131072, + max_output_length=8192, + description="Mistral routed via LiteLLM", + link="https://docs.litellm.ai/docs/providers/mistral", + function_calling=True, + ), + ], +) diff --git a/packages/dbgpt-core/src/dbgpt/model/proxy/tests/test_litellm.py b/packages/dbgpt-core/src/dbgpt/model/proxy/tests/test_litellm.py new file mode 100644 index 000000000..59b6d758b --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/model/proxy/tests/test_litellm.py @@ -0,0 +1,400 @@ +"""Tests for LiteLLM proxy LLM client.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from dbgpt.model.proxy.llms.litellm import ( + _DEFAULT_MODEL, + LiteLLMClient, + LiteLLMDeployModelParameters, +) + + +class TestLiteLLMDefaults: + """LiteLLM default model + parameter configuration.""" + + def test_default_model_constant(self): + assert _DEFAULT_MODEL == "openai/gpt-4o-mini" + + def test_client_default_model(self): + client = LiteLLMClient() + assert client.default_model == "openai/gpt-4o-mini" + + def test_client_custom_model(self): + client = LiteLLMClient(model="anthropic/claude-3-5-sonnet-20241022") + assert client.default_model == "anthropic/claude-3-5-sonnet-20241022" + + def test_model_alias_falls_back_to_model(self): + client = LiteLLMClient(model="vertex_ai/gemini-1.5-pro") + assert client._model_alias == "vertex_ai/gemini-1.5-pro" + + def test_deploy_params_provider(self): + assert LiteLLMDeployModelParameters.provider == "proxy/litellm" + + +class TestLiteLLMModelRegistration: + """LiteLLM provider/model adapter resolution.""" + + def test_openai_routed_via_litellm_resolves(self): + from dbgpt.model.adapter.base import get_model_adapter + + adapter = get_model_adapter("proxy/litellm", "openai/gpt-4o") + assert adapter is not None + + def test_anthropic_routed_via_litellm_resolves(self): + from dbgpt.model.adapter.base import get_model_adapter + + adapter = get_model_adapter( + "proxy/litellm", "anthropic/claude-3-5-sonnet-20241022" + ) + assert adapter is not None + + def test_bedrock_routed_via_litellm_resolves(self): + from dbgpt.model.adapter.base import get_model_adapter + + adapter = get_model_adapter( + "proxy/litellm", "bedrock/anthropic.claude-3-haiku-20240307-v1:0" + ) + assert adapter is not None + + def test_azure_routed_via_litellm_resolves(self): + from dbgpt.model.adapter.base import get_model_adapter + + adapter = get_model_adapter("proxy/litellm", "azure/gpt-4o") + assert adapter is not None + + +class TestLiteLLMBuildRequest: + """_build_request payload assembly.""" + + def _request(self, **overrides): + from dbgpt.core import ModelMessage, ModelRequest + + defaults = dict( + model="anthropic/claude-3-5-sonnet-20241022", + messages=[ModelMessage(role="user", content="hi")], + ) + defaults.update(overrides) + return ModelRequest(**defaults) + + def test_drop_params_defaulted_on(self): + client = LiteLLMClient() + payload = client._build_request(self._request()) + assert payload["drop_params"] is True + + def test_drop_params_can_be_disabled(self): + client = LiteLLMClient(litellm_kwargs={"drop_params": False}) + payload = client._build_request(self._request()) + assert payload["drop_params"] is False + + def test_user_litellm_kwargs_merged(self): + client = LiteLLMClient(litellm_kwargs={"num_retries": 3}) + payload = client._build_request(self._request()) + assert payload["num_retries"] == 3 + # And drop_params default is preserved. + assert payload["drop_params"] is True + + def test_temperature_and_max_tokens_forwarded(self): + client = LiteLLMClient() + payload = client._build_request( + self._request(temperature=0.3, max_new_tokens=512) + ) + assert payload["temperature"] == 0.3 + assert payload["max_tokens"] == 512 + + def test_stream_options_set_on_streaming(self): + client = LiteLLMClient() + payload = client._build_request(self._request(), stream=True) + assert payload["stream"] is True + assert payload["stream_options"] == {"include_usage": True} + + def test_stream_options_absent_on_non_stream(self): + client = LiteLLMClient() + payload = client._build_request(self._request()) + assert payload["stream"] is False + assert "stream_options" not in payload + + def test_model_resolution_uses_request_model_first(self): + client = LiteLLMClient(model="openai/gpt-4o-mini") + payload = client._build_request( + self._request(model="groq/llama-3.3-70b-versatile") + ) + assert payload["model"] == "groq/llama-3.3-70b-versatile" + + def test_api_credentials_passed_when_provided(self): + client = LiteLLMClient(api_key="sk-test", api_base="https://example.invalid/v1") + payload = client._build_request(self._request()) + assert payload["api_key"] == "sk-test" + assert payload["api_base"] == "https://example.invalid/v1" + + def test_api_credentials_omitted_when_unset(self): + client = LiteLLMClient() + payload = client._build_request(self._request()) + assert "api_key" not in payload + assert "api_base" not in payload + + def test_tools_passthrough(self): + tools = [ + { + "type": "function", + "function": {"name": "get_weather", "parameters": {}}, + } + ] + client = LiteLLMClient(litellm_kwargs={"tools": tools, "tool_choice": "auto"}) + payload = client._build_request(self._request()) + assert payload["tools"] == tools + assert payload["tool_choice"] == "auto" + + +# --------------------------------------------------------------------------- +# Mocked generate() + generate_stream() — exercises the LiteLLM adapter logic +# across the response shapes LiteLLM emits for different upstream providers. +# We patch litellm.acompletion so the tests don't need network access. +# --------------------------------------------------------------------------- + + +def _make_response(content, usage=None, reasoning_content=None): + """Mock the non-streaming litellm.acompletion return value (OpenAI shape).""" + message = SimpleNamespace(content=content) + if reasoning_content is not None: + message.reasoning_content = reasoning_content + choice = SimpleNamespace(message=message) + usage_obj = None + if usage is not None: + usage_obj = SimpleNamespace(model_dump=lambda u=usage: u) + return SimpleNamespace(choices=[choice], usage=usage_obj) + + +def _make_chunk(content=None, reasoning=None, usage=None, finish=False): + """Mock a single streaming chunk.""" + delta = SimpleNamespace(content=content) + if reasoning is not None: + delta.reasoning_content = reasoning + choices = [] + if not (usage is not None and content is None and reasoning is None and not finish): + # Non-usage-only chunks always include a choice. + choices = [SimpleNamespace(delta=delta)] + chunk_usage = None + if usage is not None: + chunk_usage = SimpleNamespace(model_dump=lambda u=usage: u) + return SimpleNamespace(choices=choices, usage=chunk_usage) + + +def _usage_only_chunk(usage): + """Mock the OpenAI stream_options=include_usage trailing chunk: empty + choices, usage attached.""" + chunk_usage = SimpleNamespace(model_dump=lambda u=usage: u) + return SimpleNamespace(choices=[], usage=chunk_usage) + + +async def _async_iter(chunks): + for c in chunks: + yield c + + +class TestLiteLLMGenerate: + """generate() — non-streaming response handling and error path.""" + + def _request(self, **overrides): + from dbgpt.core import ModelMessage, ModelRequest + from dbgpt.core.interface.message import ModelMessageRoleType + + defaults = dict( + model="anthropic/claude-3-5-sonnet-20241022", + messages=[ + ModelMessage(role=ModelMessageRoleType.HUMAN, content="hi"), + ], + ) + defaults.update(overrides) + return ModelRequest(**defaults) + + def test_generate_returns_text_and_usage(self): + client = LiteLLMClient() + usage = {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12} + + async def run(): + with patch( + "litellm.acompletion", + return_value=_make_response("4", usage=usage), + ): + return await client.generate(self._request()) + + out = asyncio.run(run()) + assert out.error_code == 0 + assert out.text == "4" + assert out.usage == usage + + def test_generate_propagates_reasoning_content(self): + client = LiteLLMClient() + + async def run(): + with patch( + "litellm.acompletion", + return_value=_make_response( + "answer", usage=None, reasoning_content="thought" + ), + ): + return await client.generate(self._request()) + + out = asyncio.run(run()) + # ModelOutput.build packs reasoning into a structured content list with + # MediaContent entries of type="thinking" and type="text". + types_to_text = {mc.type: mc.object.data for mc in (out.content or [])} + assert types_to_text.get("thinking") == "thought" + assert types_to_text.get("text") == "answer" + + def test_generate_error_returns_error_code_one(self): + client = LiteLLMClient() + + async def run(): + with patch( + "litellm.acompletion", + side_effect=RuntimeError("boom"), + ): + return await client.generate(self._request()) + + out = asyncio.run(run()) + assert out.error_code == 1 + assert "boom" in out.text + + +class TestLiteLLMGenerateStream: + """generate_stream() — provider-shape coverage without live keys.""" + + def _request(self, **overrides): + from dbgpt.core import ModelMessage, ModelRequest + from dbgpt.core.interface.message import ModelMessageRoleType + + defaults = dict( + model="anthropic/claude-3-5-sonnet-20241022", + messages=[ + ModelMessage(role=ModelMessageRoleType.HUMAN, content="hi"), + ], + ) + defaults.update(overrides) + return ModelRequest(**defaults) + + def _collect(self, client, chunks): + async def run(): + with patch( + "litellm.acompletion", + return_value=_async_iter(chunks), + ): + outputs = [] + async for o in client.generate_stream(self._request()): + outputs.append(o) + return outputs + + return asyncio.run(run()) + + def test_openai_style_separate_usage_chunk(self): + """Azure / OpenAI with stream_options=include_usage: usage arrives in a + trailing choices=[] chunk; finish-only intermediate chunk should not + cause duplicate frames.""" + client = LiteLLMClient() + usage = {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8} + outputs = self._collect( + client, + [ + _make_chunk(content="1"), + _make_chunk(content=", 2"), + _make_chunk(content=", 3"), + _make_chunk(content="", finish=True), # finish-only chunk + _usage_only_chunk(usage), + ], + ) + texts = [o.text for o in outputs] + # No duplicate text frames at the tail. + assert texts == ["1", "1, 2", "1, 2, 3"] + assert outputs[-1].error_code == 0 + + def test_anthropic_style_inline_usage(self): + """Anthropic via LiteLLM attaches usage onto the last content chunk.""" + client = LiteLLMClient() + usage = {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8} + outputs = self._collect( + client, + [ + _make_chunk(content="1"), + _make_chunk(content=", 2"), + _make_chunk(content=", 3", usage=usage), + ], + ) + texts = [o.text for o in outputs] + # Final yield contains both the full text and usage. No duplicate. + assert texts == ["1", "1, 2", "1, 2, 3"] + assert outputs[-1].usage == usage + + def test_no_duplicate_on_finish_only_chunk(self): + """A finish chunk with no new content/reasoning must not emit a frame.""" + client = LiteLLMClient() + outputs = self._collect( + client, + [ + _make_chunk(content="hi"), + # finish chunk: choices present but delta.content is "" / None + SimpleNamespace( + choices=[SimpleNamespace(delta=SimpleNamespace(content=None))], + usage=None, + ), + ], + ) + assert [o.text for o in outputs] == ["hi"] + + def test_error_yields_error_code_one(self): + client = LiteLLMClient() + + async def run(): + with patch("litellm.acompletion", side_effect=RuntimeError("net dead")): + outs = [] + async for o in client.generate_stream(self._request()): + outs.append(o) + return outs + + outputs = asyncio.run(run()) + assert len(outputs) == 1 + assert outputs[0].error_code == 1 + assert "net dead" in outputs[0].text + + +class TestLiteLLMNewClient: + """new_client() — DB-GPT's actual entry point when loading a deploy config.""" + + def test_new_client_constructs_with_params(self): + params = LiteLLMDeployModelParameters( + name="anthropic/claude-3-5-sonnet-20241022", + provider="proxy/litellm", + ) + client = LiteLLMClient.new_client(params) + assert isinstance(client, LiteLLMClient) + assert client.default_model == "anthropic/claude-3-5-sonnet-20241022" + assert client._model_alias == "anthropic/claude-3-5-sonnet-20241022" + + def test_new_client_preserves_drop_params_default(self): + params = LiteLLMDeployModelParameters( + name="groq/llama-3.3-70b-versatile", + provider="proxy/litellm", + ) + client = LiteLLMClient.new_client(params) + assert client._litellm_kwargs.get("drop_params") is True + + +class TestLiteLLMImportGuard: + def test_import_error_message_mentions_pin(self): + # Force-fail the litellm import inside __init__ to verify the user- + # facing error names the supported version range. + import sys + + saved = sys.modules.pop("litellm", None) + sys.modules["litellm"] = None # type: ignore[assignment] + try: + with pytest.raises(ValueError, match=r"litellm>=1\.60"): + LiteLLMClient() + finally: + if saved is not None: + sys.modules["litellm"] = saved + else: + sys.modules.pop("litellm", None)