models.py

from pydantic import BaseModel, Field
from typing import Literal
from uuid import uuid4

ModelID = Literal[
    "openai/gpt-4o-mini",
    "meta-llama/llama-3-70b-instruct",
    "anthropic/claude-3.5-sonnet",
    "deepseek/deepseek-coder",
    "anthropic/claude-3-haiku",
    "openai/gpt-3.5-turbo-instruct",
    "qwen/qwen-72b-chat",
    "google/gemma-2-27b-it"
]

class QueryModel(BaseModel):
    user_query: str = Field(..., description="User's coding query")
    model_id: ModelID = Field(
        default="meta-llama/llama-3-70b-instruct",
        description="ID of the model to use for response generation"
    )
    conversation_id: str = Field(default_factory=lambda: str(uuid4()), description="Unique identifier for the conversation")
    user_id: str = Field(..., description="Unique identifier for the user")

    class Config:
        schema_extra = {
            "example": {
                "user_query": "How do I implement a binary search in Python?",
                "model_id": "meta-llama/llama-3-70b-instruct",
                "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
                "user_id": "user123"
            }
        }

class NewsQueryModel(BaseModel):
    query: str = Field(..., description="News topic to search for")
    model_id: ModelID = Field(
        default="openai/gpt-4o-mini",
        description="ID of the model to use for response generation"
    )
    class Config:
        schema_extra = {
            "example": {
                "query": "Latest developments in AI",
                "model_id": "openai/gpt-4o-mini"
            }
        }

class FollowupQueryModel(BaseModel):
    query: str = Field(..., description="User's query for the followup agent")
    model_id: ModelID = Field(
        default="openai/gpt-4o-mini",
        description="ID of the model to use for response generation"
    )
    conversation_id: str = Field(default_factory=lambda: str(uuid4()), description="Unique identifier for the conversation")
    user_id: str = Field(..., description="Unique identifier for the user")
    tool_call: Literal["web", "news", "auto"] = Field(
        default="auto",
        description="Type of tool to call (web, news, auto)"
    )

    class Config:
        schema_extra = {
            "example": {
                "query": "How can I improve my productivity?",
                "model_id": "openai/gpt-4o-mini",
                "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
                "user_id": "user123",
                "tool_call": "auto"
            }
        }