from typing import List

from pydantic import BaseModel


class QueryRequest(BaseModel):
    """
    Represents the request model for querying similar prompts.
    """

    query: str
    n: int = 5


class SimilarPrompt(BaseModel):
    """
    Represents a single similar prompt with its similarity score.
    """

    score: float
    prompt: str


class QueryResponse(BaseModel):
    """
    Represents the response model containing a list of similar prompts.
    """

    similar_prompts: List[SimilarPrompt]