vertex / app /models.py
bibibi12345's picture
complete refactor
3cc1b9e
from pydantic import BaseModel, ConfigDict # Field removed
from typing import List, Dict, Any, Optional, Union, Literal
# Define data models
class ImageUrl(BaseModel):
url: str
class ContentPartImage(BaseModel):
type: Literal["image_url"]
image_url: ImageUrl
class ContentPartText(BaseModel):
type: Literal["text"]
text: str
class OpenAIMessage(BaseModel):
role: str
content: Union[str, List[Union[ContentPartText, ContentPartImage, Dict[str, Any]]]]
class OpenAIRequest(BaseModel):
model: str
messages: List[OpenAIMessage]
temperature: Optional[float] = 1.0
max_tokens: Optional[int] = None
top_p: Optional[float] = 1.0
top_k: Optional[int] = None
stream: Optional[bool] = False
stop: Optional[List[str]] = None
presence_penalty: Optional[float] = None
frequency_penalty: Optional[float] = None
seed: Optional[int] = None
logprobs: Optional[int] = None
response_logprobs: Optional[bool] = None
n: Optional[int] = None # Maps to candidate_count in Vertex AI
# Allow extra fields to pass through without causing validation errors
model_config = ConfigDict(extra='allow')