x

app/schemas.py

@@ -0,0 +1,26 @@
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class GenerateRequest(BaseModel):
+    prompt: str
+    max_tokens: int = Field(default=256, ge=1, le=4096)
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
+    top_p: float = Field(default=0.95, gt=0.0, le=1.0)
+    repetition_penalty: float = Field(default=1.0, ge=0.5, le=2.0)
+    stop: Optional[List[str]] = None
+
+
+class GenerateResponse(BaseModel):
+    text: str
+    prompt: str
+    model: str
+    usage_prompt_tokens: int
+    usage_completion_tokens: int
+    usage_total_tokens: int
+
+
+class HealthResponse(BaseModel):
+    status: str
+    model: str