75 lines
2.8 KiB
Python
75 lines
2.8 KiB
Python
from typing import Literal
|
|
|
|
from letta.log import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
from pydantic import Field
|
|
|
|
from letta.schemas.enums import ProviderCategory, ProviderType
|
|
from letta.schemas.llm_config import LLMConfig
|
|
from letta.schemas.providers.openai import OpenAIProvider
|
|
|
|
MODEL_CONTEXT_WINDOWS = {
|
|
"grok-3-fast": 131_072,
|
|
"grok-3": 131_072,
|
|
"grok-3-mini": 131_072,
|
|
"grok-3-mini-fast": 131_072,
|
|
"grok-4-0709": 256_000,
|
|
"grok-4-fast-reasoning": 2_000_000,
|
|
"grok-4-fast-non-reasoning": 2_000_000,
|
|
"grok-code-fast-1": 256_000
|
|
}
|
|
|
|
|
|
class XAIProvider(OpenAIProvider):
|
|
"""https://docs.x.ai/docs/api-reference"""
|
|
|
|
provider_type: Literal[ProviderType.xai] = Field(ProviderType.xai, description="The type of the provider.")
|
|
provider_category: ProviderCategory = Field(ProviderCategory.base, description="The category of the provider (base or byok)")
|
|
api_key: str | None = Field(None, description="API key for the xAI/Grok API.", deprecated=True)
|
|
base_url: str = Field("https://api.x.ai/v1", description="Base URL for the xAI/Grok API.")
|
|
|
|
def get_model_context_window_size(self, model_name: str) -> int | None:
|
|
# xAI doesn't return context window in the model listing,
|
|
# this is hardcoded from https://docs.x.ai/docs/models
|
|
return MODEL_CONTEXT_WINDOWS.get(model_name)
|
|
|
|
async def list_llm_models_async(self) -> list[LLMConfig]:
|
|
from letta.llm_api.openai import openai_get_model_list_async
|
|
|
|
api_key = await self.api_key_enc.get_plaintext_async() if self.api_key_enc else None
|
|
response = await openai_get_model_list_async(self.base_url, api_key=api_key)
|
|
|
|
data = response.get("data", response)
|
|
|
|
configs = []
|
|
for model in data:
|
|
assert "id" in model, f"xAI/Grok model missing 'id' field: {model}"
|
|
model_name = model["id"]
|
|
|
|
# In case xAI starts supporting it in the future:
|
|
if "context_length" in model:
|
|
context_window_size = model["context_length"]
|
|
else:
|
|
context_window_size = self.get_model_context_window_size(model_name)
|
|
|
|
if not context_window_size:
|
|
logger.warning(f"Couldn't find context window size for model {model_name}")
|
|
continue
|
|
|
|
configs.append(
|
|
LLMConfig(
|
|
model=model_name,
|
|
model_endpoint_type="xai",
|
|
model_endpoint=self.base_url,
|
|
context_window=context_window_size,
|
|
handle=self.get_handle(model_name),
|
|
max_tokens=self.get_default_max_output_tokens(model_name),
|
|
provider_name=self.name,
|
|
provider_category=self.provider_category,
|
|
)
|
|
)
|
|
|
|
return configs
|