diff --git a/configs/dbgpt-proxy-infiniai.toml b/configs/dbgpt-proxy-infiniai.toml new file mode 100644 index 000000000..8c0109e59 --- /dev/null +++ b/configs/dbgpt-proxy-infiniai.toml @@ -0,0 +1,40 @@ +[system] +# Load language from environment variable(It is set by the hook) +language = "${env:DBGPT_LANG:-zh}" +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 +[models] +[[models.llms]] +name = "deepseek-v3" +provider = "proxy/infiniai" +api_key = "${env:INFINIAI_API_KEY}" + +[[models.embeddings]] +name = "bge-m3" +provider = "proxy/openai" +api_url = "https://cloud.infini-ai.com/maas/v1/embeddings" +api_key = "${env:INFINIAI_API_KEY}" + +[[models.rerankers]] +type = "reranker" +name = "bge-reranker-v2-m3" +provider = "proxy/infiniai" +api_key = "${env:INFINIAI_API_KEY}" \ No newline at end of file diff --git a/packages/dbgpt-core/src/dbgpt/model/proxy/__init__.py b/packages/dbgpt-core/src/dbgpt/model/proxy/__init__.py index 88973a727..8db369c84 100644 --- a/packages/dbgpt-core/src/dbgpt/model/proxy/__init__.py +++ b/packages/dbgpt-core/src/dbgpt/model/proxy/__init__.py @@ -8,6 +8,7 @@ if TYPE_CHECKING: from dbgpt.model.proxy.llms.deepseek import DeepseekLLMClient 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.moonshot import MoonshotLLMClient from dbgpt.model.proxy.llms.ollama import OllamaLLMClient from dbgpt.model.proxy.llms.siliconflow import SiliconFlowLLMClient @@ -33,6 +34,7 @@ def __lazy_import(name): "OllamaLLMClient": "dbgpt.model.proxy.llms.ollama", "DeepseekLLMClient": "dbgpt.model.proxy.llms.deepseek", "GiteeLLMClient": "dbgpt.model.proxy.llms.gitee", + "InfiniAILLMClient": "dbgpt.model.proxy.llms.infiniai", } if name in module_path: @@ -60,4 +62,5 @@ __all__ = [ "OllamaLLMClient", "DeepseekLLMClient", "GiteeLLMClient", + "InfiniAILLMClient", ] diff --git a/packages/dbgpt-core/src/dbgpt/model/proxy/llms/infiniai.py b/packages/dbgpt-core/src/dbgpt/model/proxy/llms/infiniai.py new file mode 100644 index 000000000..07990f40b --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/model/proxy/llms/infiniai.py @@ -0,0 +1,187 @@ +import os +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union + +from dbgpt.core import ModelMetadata +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, + register_proxy_model_adapter, +) +from .chatgpt import OpenAICompatibleDeployModelParameters, OpenAILLMClient + +if TYPE_CHECKING: + from httpx._types import ProxiesTypes + from openai import AsyncAzureOpenAI, AsyncOpenAI + + ClientType = Union[AsyncAzureOpenAI, AsyncOpenAI] + + +_INFINIAI_DEFAULT_MODEL = "deepseek-v3" + + +@auto_register_resource( + label=_("InfiniAI Proxy LLM"), + category=ResourceCategory.LLM_CLIENT, + tags={"order": TAGS_ORDER_HIGH}, + description=_("InfiniAI proxy LLM configuration."), + documentation_url="https://docs.infini-ai.com/gen-studio/api/tutorial.html", # noqa + show_in_ui=False, +) +@dataclass +class InfiniAIDeployModelParameters(OpenAICompatibleDeployModelParameters): + """Deploy model parameters for InfiniAI.""" + + provider: str = "proxy/infiniai" + + api_base: Optional[str] = field( + default="${env:INFINIAI_API_BASE:-https://cloud.infini-ai.com/maas/v1}", + metadata={ + "help": _("The base url of the InfiniAI API."), + }, + ) + + api_key: Optional[str] = field( + default="${env:INFINIAI_API_KEY}", + metadata={ + "help": _("The API key of the InfiniAI API."), + "tags": "privacy", + }, + ) + + +async def infiniai_generate_stream( + model: ProxyModel, tokenizer, params, device, context_len=2048 +): + client: InfiniAILLMClient = 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 InfiniAILLMClient(OpenAILLMClient): + """InfiniAI LLM Client. + + InfiniAI's API is compatible with OpenAI's API, so we inherit from + OpenAILLMClient. + """ + + 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] = _INFINIAI_DEFAULT_MODEL, + proxies: Optional["ProxiesTypes"] = None, + timeout: Optional[int] = 240, + model_alias: Optional[str] = _INFINIAI_DEFAULT_MODEL, + context_length: Optional[int] = None, + openai_client: Optional["ClientType"] = None, + openai_kwargs: Optional[Dict[str, Any]] = None, + **kwargs, + ): + api_base = ( + api_base + or os.getenv("INFINIAI_API_BASE") + or "https://cloud.infini-ai.com/maas/v1" + ) + api_key = api_key or os.getenv("INFINIAI_API_KEY") + model = model or _INFINIAI_DEFAULT_MODEL + if not context_length: + if "200k" in model: + context_length = 200 * 1024 + else: + context_length = 4096 + + if not api_key: + raise ValueError( + "InfiniAI API key is required, please set 'INFINIAI_API_KEY' in " + "environment or pass it as an argument." + ) + + super().__init__( + api_key=api_key, + api_base=api_base, + api_type=api_type, + api_version=api_version, + model=model, + proxies=proxies, + timeout=timeout, + model_alias=model_alias, + context_length=context_length, + openai_client=openai_client, + openai_kwargs=openai_kwargs, + **kwargs, + ) + + @property + def default_model(self) -> str: + model = self._model + if not model: + model = _INFINIAI_DEFAULT_MODEL + return model + + @classmethod + def param_class(cls) -> Type[InfiniAIDeployModelParameters]: + return InfiniAIDeployModelParameters + + @classmethod + def generate_stream_function( + cls, + ) -> Optional[Union[GenerateStreamFunction, AsyncGenerateStreamFunction]]: + return infiniai_generate_stream + + +register_proxy_model_adapter( + InfiniAILLMClient, + supported_models=[ + ModelMetadata( + model=["deepseek-v3"], + context_length=64 * 1024, + max_output_length=8 * 1024, + description="DeepSeek-V3 by DeepSeek", + link="https://cloud.infini-ai.com/genstudio/model", + function_calling=True, + ), + ModelMetadata( + model=["deepseek-r1"], + context_length=64 * 1024, + max_output_length=8 * 1024, + description="DeepSeek-V3 by DeepSeek", + link="https://cloud.infini-ai.com/genstudio/model", + function_calling=False, + ), + ModelMetadata( + model=["qwq-32b"], + context_length=64 * 1024, + max_output_length=8 * 1024, + description="qwq By Qwen", + link="https://cloud.infini-ai.com/genstudio/model", + function_calling=True, + ), + ModelMetadata( + model=[ + "qwen2.5-72b-instruct", + "qwen2.5-32b-instruct", + "qwen2.5-14b-instruct", + "qwen2.5-7b-instruct", + "qwen2.5-coder-32b-instruct", + ], + context_length=32 * 1024, + max_output_length=4 * 1024, + description="Qwen 2.5 By Qwen", + link="https://cloud.infini-ai.com/genstudio/model", + function_calling=True, + ), + # More models see: https://cloud.infiniai.cn/models + ], +) diff --git a/packages/dbgpt-core/src/dbgpt/rag/embedding/__init__.py b/packages/dbgpt-core/src/dbgpt/rag/embedding/__init__.py index 1baef38b8..affbba870 100644 --- a/packages/dbgpt-core/src/dbgpt/rag/embedding/__init__.py +++ b/packages/dbgpt-core/src/dbgpt/rag/embedding/__init__.py @@ -15,6 +15,7 @@ from .embeddings import ( # noqa: F401 ) from .rerank import ( # noqa: F401 CrossEncoderRerankEmbeddings, + InfiniAIRerankEmbeddings, OpenAPIRerankEmbeddings, SiliconFlowRerankEmbeddings, ) @@ -31,5 +32,6 @@ __ALL__ = [ "OpenAPIEmbeddings", "OpenAPIRerankEmbeddings", "SiliconFlowRerankEmbeddings", + "InfiniAIRerankEmbeddings", "WrappedEmbeddingFactory", ] diff --git a/packages/dbgpt-core/src/dbgpt/rag/embedding/rerank.py b/packages/dbgpt-core/src/dbgpt/rag/embedding/rerank.py index ee57402b8..caf986b7f 100644 --- a/packages/dbgpt-core/src/dbgpt/rag/embedding/rerank.py +++ b/packages/dbgpt-core/src/dbgpt/rag/embedding/rerank.py @@ -493,6 +493,77 @@ class TeiRerankEmbeddings(OpenAPIRerankEmbeddings): return self._parse_results(response_data) +@dataclass +class InfiniAIRerankEmbeddingsParameters(OpenAPIRerankerDeployModelParameters): + """InfiniAI Rerank Embeddings Parameters.""" + + provider: str = "proxy/infiniai" + + api_url: str = field( + default="https://cloud.infini-ai.com/maas/v1/rerank", + metadata={ + "help": _("The URL of the rerank API."), + }, + ) + api_key: Optional[str] = field( + default="${env:INFINIAI_API_KEY}", + metadata={ + "help": _("The API key for the rerank API."), + }, + ) + + +class InfiniAIRerankEmbeddings(OpenAPIRerankEmbeddings): + """InfiniAI Rerank Model. + + See `InfiniAI API + `_ for more details. + """ + + def __init__(self, **kwargs: Any): + """Initialize the InfiniAIRerankEmbeddings.""" + # If the API key is not provided, try to get it from the environment + if "api_key" not in kwargs: + kwargs["api_key"] = os.getenv("InfiniAI_API_KEY") + + if "api_url" not in kwargs: + env_api_url = os.getenv("InfiniAI_API_BASE") + if env_api_url: + env_api_url = env_api_url.rstrip("/") + kwargs["api_url"] = env_api_url + "/rerank" + else: + kwargs["api_url"] = "https://cloud.infini-ai.com/maas/v1/rerank" + + if "model_name" not in kwargs: + kwargs["model_name"] = "bge-reranker-v2-m3" + + super().__init__(**kwargs) + + @classmethod + def param_class(cls) -> Type[InfiniAIRerankEmbeddingsParameters]: + """Get the parameter class.""" + return InfiniAIRerankEmbeddingsParameters + + def _parse_results(self, response: Dict[str, Any]) -> List[float]: + """Parse the response from the API. + + Args: + response: The response from the API. + + Returns: + List[float]: The rank scores of the candidates. + """ + results = response.get("results") + if not results: + raise RuntimeError("Cannot find results in the response") + if not isinstance(results, list): + raise RuntimeError("Results should be a list") + # Sort by index, 0 in the first element + results = sorted(results, key=lambda x: x.get("index", 0)) + scores = [float(result.get("relevance_score")) for result in results] + return scores + + register_embedding_adapter( CrossEncoderRerankEmbeddings, supported_models=RERANKER_COMMON_HF_MODELS ) @@ -505,3 +576,6 @@ register_embedding_adapter( register_embedding_adapter( TeiRerankEmbeddings, supported_models=RERANKER_COMMON_HF_MODELS ) +register_embedding_adapter( + InfiniAIRerankEmbeddings, supported_models=RERANKER_COMMON_HF_MODELS +)