Spaces:
Running
Running
File size: 2,600 Bytes
d0dd276 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
from pydantic import BaseModel, ConfigDict # Field removed
from typing import List, Dict, Any, Optional, Union, Literal
from app.utils.logging import vertex_log
# 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')
class TokenUsage(BaseModel):
completion_tokens: int = 0
prompt_tokens: int = 0
total_tokens: int = 0
class GeminiMessage(BaseModel):
role: str # 'user' or 'model'
content: str
name: Optional[str] = None
class GeminiChatRequest(BaseModel):
model: str
messages: List[GeminiMessage]
temperature: Optional[float] = 0.7
top_p: Optional[float] = 0.95
top_k: Optional[int] = 40
max_output_tokens: Optional[int] = 2048
stream: Optional[bool] = False
def log_request(self):
vertex_log('info', f"Chat request for model: {self.model}")
vertex_log('debug', f"Request parameters: temp={self.temperature}, top_p={self.top_p}, max_tokens={self.max_output_tokens}")
class GeminiCompletionRequest(BaseModel):
model: str
prompt: str
temperature: Optional[float] = 0.7
top_p: Optional[float] = 0.95
top_k: Optional[int] = 40
max_output_tokens: Optional[int] = 2048
stream: Optional[bool] = False
def log_request(self):
vertex_log('info', f"Completion request for model: {self.model}")
vertex_log('debug', f"Request parameters: temp={self.temperature}, top_p={self.top_p}, max_tokens={self.max_output_tokens}")
prompt_preview = self.prompt[:50] + "..." if len(self.prompt) > 50 else self.prompt
vertex_log('debug', f"Prompt preview: {prompt_preview}") |