Tech Stack Advisor - Code Viewer

← Back to File Tree

models.py

Language: python | Path: backend/src/api/models.py | Lines: 239
"""Pydantic models for API requests and responses."""
from typing import Any
from pydantic import BaseModel, Field


class RecommendationRequest(BaseModel):
    """Request model for tech stack recommendations."""

    query: str = Field(
        ...,
        description="Natural language query describing the tech stack requirements",
        min_length=10,
        max_length=1000,
        examples=["I'm building a real-time chat application expecting 100K daily active users"],
    )

    # Optional overrides for parsed context
    dau: int | None = Field(
        None,
        description="Daily active users (optional override)",
        ge=0,
        examples=[100_000],
    )
    budget_target: float | None = Field(
        None,
        description="Monthly budget target in USD (optional)",
        ge=0,
        examples=[500.0],
    )

    # Optional user-provided API key
    api_key: str | None = Field(
        None,
        description="User's Anthropic API key (optional). If provided, bypasses shared budget limits.",
        min_length=20,
        max_length=200,
        examples=["sk-ant-api03-..."],
    )

    class Config:
        json_schema_extra = {
            "example": {
                "query": "I need a database and infrastructure for an e-commerce site with 50k DAU and payment processing",
                "dau": 50_000,
                "budget_target": 1000.0,
            }
        }


class ParsedContext(BaseModel):
    """Parsed context from user query."""

    dau: int = Field(..., description="Daily active users")
    qps: int = Field(..., description="Queries per second")
    data_type: str = Field(..., description="Data type (structured/unstructured/time-series)")
    workload_type: str = Field(..., description="Workload type (web/api/realtime/background)")
    data_sensitivity: str = Field(..., description="Data sensitivity level")
    compliance: list[str] = Field(..., description="Required compliance frameworks")


class AgentResult(BaseModel):
    """Generic agent result."""

    agent: str = Field(..., description="Agent name")
    recommendations: str = Field(..., description="Agent recommendations")
    scale_info: dict[str, Any] | None = Field(None, description="Scale information")
    cost_comparisons: list[dict[str, Any]] | None = Field(None, description="Cost comparisons")
    threat_assessment: dict[str, Any] | None = Field(None, description="Threat assessment")


class RecommendationResponse(BaseModel):
    """Response model for tech stack recommendations."""

    status: str = Field(..., description="Status (success/error)")
    query: str = Field(..., description="Original user query")
    correlation_id: str = Field(..., description="Request correlation ID for tracing")
    parsed_context: ParsedContext | None = Field(None, description="Parsed context from query")
    recommendations: dict[str, Any] | None = Field(None, description="All agent recommendations")
    error: str | None = Field(None, description="Error message if status is error")

    class Config:
        json_schema_extra = {
            "example": {
                "status": "success",
                "query": "Building a chat app for 100K users",
                "correlation_id": "uuid-here",
                "parsed_context": {
                    "dau": 100_000,
                    "qps": 25,
                    "data_type": "structured",
                    "workload_type": "realtime",
                    "data_sensitivity": "medium",
                    "compliance": [],
                },
                "recommendations": {
                    "database": "...",
                    "infrastructure": "...",
                    "cost": "...",
                    "security": "...",
                },
            }
        }


class HealthResponse(BaseModel):
    """Health check response."""

    status: str = Field(..., description="Service status")
    version: str = Field(..., description="API version")
    agents_loaded: int = Field(..., description="Number of agents loaded")
    uptime_seconds: float = Field(..., description="Service uptime in seconds")


class MetricsResponse(BaseModel):
    """Metrics response."""

    total_requests: int = Field(..., description="Total requests processed")
    total_tokens: int = Field(..., description="Total tokens consumed")
    total_cost_usd: float = Field(..., description="Total cost in USD")
    daily_queries: int = Field(..., description="Queries today")
    daily_cost_usd: float = Field(..., description="Cost today in USD")
    budget_remaining_usd: float = Field(..., description="Daily budget remaining")


# Conversation API models
class ConversationStartRequest(BaseModel):
    """Request to start a new conversation."""

    initial_message: str = Field(
        ...,
        description="Initial user message to start the conversation",
        min_length=5,
        max_length=500,
        examples=["I want to build a social media platform"],
    )
    api_key: str | None = Field(
        None,
        description="Optional user-provided Anthropic API key",
    )


class ConversationStartResponse(BaseModel):
    """Response when starting a conversation."""

    session_id: str = Field(..., description="Session ID for this conversation")
    question: str = Field(..., description="Next question from the assistant")
    question_type: str | None = Field(None, description="Question type: 'choice' or 'text'")
    options: list[str] | None = Field(None, description="Choice options if question_type is 'choice'")
    extracted_context: dict[str, Any] = Field(..., description="Context extracted so far")
    completion_percentage: int = Field(..., description="Completion percentage (0-100)")


class ConversationMessageRequest(BaseModel):
    """Request to send a message in an existing conversation."""

    session_id: str = Field(..., description="Session ID from conversation start")
    message: str = Field(
        ...,
        description="User's message/response",
        min_length=1,
        max_length=500,
    )
    api_key: str | None = Field(
        None,
        description="Optional user-provided Anthropic API key",
    )


class ConversationMessageResponse(BaseModel):
    """Response to a conversation message."""

    session_id: str = Field(..., description="Session ID")
    question: str | None = Field(None, description="Next question (null if ready for recommendation)")
    question_type: str | None = Field(None, description="Question type: 'choice' or 'text'")
    options: list[str] | None = Field(None, description="Choice options if question_type is 'choice'")
    extracted_context: dict[str, Any] = Field(..., description="Updated extracted context")
    completion_percentage: int = Field(..., description="Completion percentage (0-100)")
    ready_for_recommendation: bool = Field(..., description="True if ready to generate recommendations")
    conversation_history: list[dict[str, Any]] = Field(..., description="Full conversation history")


class ConversationStatusResponse(BaseModel):
    """Response with conversation status."""

    session_id: str = Field(..., description="Session ID")
    exists: bool = Field(..., description="Whether session exists")
    conversation_history: list[dict[str, Any]] | None = Field(None, description="Conversation history")
    extracted_context: dict[str, Any] | None = Field(None, description="Extracted context")
    completion_percentage: int | None = Field(None, description="Completion percentage")
    ready_for_recommendation: bool | None = Field(None, description="Ready for recommendation")


# Authentication API models
class RegisterRequest(BaseModel):
    """Request to register a new user."""

    email: str = Field(..., description="User email", examples=["user@example.com"])
    password: str = Field(..., description="User password", min_length=8)
    full_name: str | None = Field(None, description="User's full name")


class LoginRequest(BaseModel):
    """Request to login."""

    email: str = Field(..., description="User email")
    password: str = Field(..., description="User password")


class AuthResponse(BaseModel):
    """Response with authentication token."""

    token: str = Field(..., description="JWT access token")
    user: dict[str, Any] = Field(..., description="User information")


class FeedbackRequest(BaseModel):
    """Request to submit feedback."""

    query_id: str | None = Field(None, description="Related query ID (optional)")
    rating: int | None = Field(None, description="Rating 1-5 stars", ge=1, le=5)
    comment: str | None = Field(None, description="Feedback comment")
    feedback_type: str | None = Field(None, description="Type: helpful, not_helpful, inaccurate, etc.")


class SimilarQueriesResponse(BaseModel):
    """Response with similar past queries."""

    queries: list[dict[str, Any]] = Field(..., description="List of similar queries")


class AdminStatsResponse(BaseModel):
    """Response with admin statistics."""

    total_users: int = Field(..., description="Total number of users")
    total_queries: int = Field(..., description="Total queries across all users")
    total_cost_usd: float = Field(..., description="Total cost across all users")
    recent_queries: list[dict[str, Any]] = Field(..., description="Recent queries")
    recent_feedback: list[dict[str, Any]] = Field(..., description="Recent feedback")
    user_stats: list[dict[str, Any]] = Field(..., description="Per-user statistics")