mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-07-23 04:12:13 +00:00
feat: add model provider InfiniAI (#2653)
Co-authored-by: yaozhuyu <yaozhuyu@infini-ai.com>
This commit is contained in:
parent
1b77ed6319
commit
445076b433
40
configs/dbgpt-proxy-infiniai.toml
Normal file
40
configs/dbgpt-proxy-infiniai.toml
Normal file
@ -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}"
|
@ -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",
|
||||
]
|
||||
|
187
packages/dbgpt-core/src/dbgpt/model/proxy/llms/infiniai.py
Normal file
187
packages/dbgpt-core/src/dbgpt/model/proxy/llms/infiniai.py
Normal file
@ -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
|
||||
],
|
||||
)
|
@ -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",
|
||||
]
|
||||
|
@ -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
|
||||
<https://docs.infini-ai.com/gen-studio/api/tutorial-rerank.html>`_ 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
|
||||
)
|
||||
|
Loading…
Reference in New Issue
Block a user