fix: Make embedding_endpoint_type an enum (#2251)

This commit is contained in:
Matthew Zhou
2024-12-16 17:16:43 -08:00
committed by GitHub
parent c7131ad0da
commit 80f32b41f5

View File

@@ -1,4 +1,4 @@
from typing import Optional
from typing import Literal, Optional
from pydantic import BaseModel, Field
@@ -20,7 +20,25 @@ class EmbeddingConfig(BaseModel):
"""
embedding_endpoint_type: str = Field(..., description="The endpoint type for the model.")
embedding_endpoint_type: Literal[
"openai",
"anthropic",
"cohere",
"google_ai",
"azure",
"groq",
"ollama",
"webui",
"webui-legacy",
"lmstudio",
"lmstudio-legacy",
"llamacpp",
"koboldcpp",
"vllm",
"hugging-face",
"mistral",
"together", # completions endpoint
] = Field(..., description="The endpoint type for the model.")
embedding_endpoint: Optional[str] = Field(None, description="The endpoint for the model (`None` if local).")
embedding_model: str = Field(..., description="The model for the embedding.")
embedding_dim: int = Field(..., description="The dimension of the embedding.")