Add arcee-ai/prime model

models.py

@@ -4,39 +4,32 @@ from enum import Enum
 
 
 class OpenRouterModel(str, Enum):
-    CLAUDE_35_SONNET = "claude-3.5-sonnet"
-    CLAUDE_3_OPUS = "claude-3-opus"
-    GPT_4O = "gpt-4o"
-    GPT_4O_MINI = "gpt-4o-mini"
-    LLAMA_31_70B = "llama-3.1-70b"
-    LLAMA_31_8B = "llama-3.1-8b"
+    ARCEE_PRIME = "arcee-ai/prime"
 
 
 class SamplerType(str, Enum):
     CATEGORY = "CATEGORY"
     UNIFORM = "UNIFORM"
     GAUSSIAN = "GAUSSIAN"
-    UUID = "UUID"
-    DATETIME = "DATETIME"
 
 
 class ColumnConfig(BaseModel):
     name: str = Field(..., description="Column name")
-    type: str = Field(..., description="Column type: sampler, llm_text, llm_code, llm_structured")
-    params: dict = Field(default_factory=dict, description="Column-specific parameters")
+    type: str = Field(..., description="Column type: sampler, llm_text")
+    params: dict = Field(default_factory=dict)
 
 
 class GenerateRequest(BaseModel):
-    num_records: int = Field(default=10, ge=1, le=1000, description="Number of records to generate")
-    model: OpenRouterModel = Field(default=OpenRouterModel.GPT_4O_MINI, description="OpenRouter model to use")
+    num_records: int = Field(default=10, ge=1, le=100)
+    model: OpenRouterModel = Field(default=OpenRouterModel.ARCEE_PRIME)
     columns: list[ColumnConfig] = Field(..., description="Column configurations")
     temperature: float = Field(default=0.7, ge=0.0, le=2.0)
     max_tokens: int = Field(default=512, ge=64, le=4096)
 
 
 class PreviewRequest(BaseModel):
-    model: OpenRouterModel = Field(default=OpenRouterModel.GPT_4O_MINI)
-    columns: list[ColumnConfig] = Field(..., description="Column configurations")
+    model: OpenRouterModel = Field(default=OpenRouterModel.ARCEE_PRIME)
+    columns: list[ColumnConfig] = Field(...)
     temperature: float = Field(default=0.7)
     max_tokens: int = Field(default=512)