mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-27 23:47:46 +00:00
TPTBusiness
732361bb90
- Path injection (B614): centralized safe_resolve_path in core/utils.py, refactored 6 UI modules to use it with safe_root validation - B701: added explicit autoescape=select_autoescape() to Jinja2 Environment() calls in 3 files - B101: replaced assert statements with proper if/raise patterns in 12+ files (partial) - B112: added logger.warning() to bare except:continue blocks in 5 files
66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
"""
|
|
Adapter tools for pydantic-ai
|
|
"""
|
|
|
|
import os
|
|
|
|
from litellm.utils import get_llm_provider
|
|
from pydantic_ai.models.openai import OpenAIChatModel, OpenAIChatModelSettings
|
|
from pydantic_ai.providers.litellm import LiteLLMProvider
|
|
|
|
from rdagent.oai.backend.litellm import LiteLLMAPIBackend
|
|
from rdagent.oai.llm_conf import LLM_SETTINGS
|
|
from rdagent.oai.llm_utils import APIBackend
|
|
|
|
# NOTE:
|
|
# LiteLLM's code is not well orgnized.
|
|
# we can't reuse any component to map the provider to the env name
|
|
# So we have to hardcode on here.
|
|
PROVIDER_TO_ENV_MAP = {
|
|
"openai": "OPENAI",
|
|
"azure_ai": "AZURE_AI",
|
|
"azure": "AZURE",
|
|
"litellm_proxy": "LITELLM_PROXY",
|
|
}
|
|
|
|
|
|
def get_agent_model() -> OpenAIChatModel:
|
|
"""
|
|
Converting LiteLLM to a pydantic-ai model. So you can use like this
|
|
|
|
.. code-block:: python
|
|
|
|
from rdagent.oai.backend.pydantic_ai import get_agent_model
|
|
model = get_agent_model()
|
|
agent = Agent(model)
|
|
|
|
"""
|
|
backend = APIBackend()
|
|
if not isinstance(backend, LiteLLMAPIBackend):
|
|
raise TypeError("Only LiteLLMAPIBackend is supported")
|
|
|
|
compl_kwargs = backend.get_complete_kwargs()
|
|
|
|
selected_model = compl_kwargs["model"]
|
|
|
|
_, custom_llm_provider, _, _ = get_llm_provider(selected_model)
|
|
if custom_llm_provider not in PROVIDER_TO_ENV_MAP:
|
|
raise ValueError(
|
|
f"Provider {custom_llm_provider} not supported. Please add it into `PROVIDER_TO_ENV_MAP`"
|
|
)
|
|
prefix = PROVIDER_TO_ENV_MAP[custom_llm_provider]
|
|
api_key = os.getenv(f"{prefix}_API_KEY", None)
|
|
api_base = os.getenv(f"{prefix}_API_BASE", None)
|
|
|
|
kwargs = {
|
|
"openai_reasoning_effort": compl_kwargs.get("reasoning_effort"),
|
|
"max_tokens": compl_kwargs.get("max_tokens"),
|
|
"temperature": compl_kwargs.get("temperature"),
|
|
}
|
|
if compl_kwargs.get("max_tokens") is None:
|
|
kwargs["max_tokens"] = LLM_SETTINGS.chat_max_tokens
|
|
settings = OpenAIChatModelSettings(**kwargs)
|
|
return OpenAIChatModel(
|
|
selected_model, provider=LiteLLMProvider(api_base=api_base, api_key=api_key), settings=settings
|
|
)
|