nlproxy/server/schemas.py

"""Pydantic schemas for NLProxy server APIs.

Author: IntelliDeep Labs Team
License: BSL 1.1
"""

from __future__ import annotations

import time
import uuid
from typing import Any, Dict, List, Optional, Literal
from pydantic import BaseModel, Field, ConfigDict, field_validator

from nlproxy.llm.client import LLMProvider


class Message(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: str
    name: Optional[str] = None

    @field_validator("content")
    @classmethod
    def validate_content(cls, value: str) -> str:
        if not value or not value.strip():
            raise ValueError("content must not be empty")
        return value.strip()


class ChatCompletionRequest(BaseModel):
    model: str = Field(default="gpt-4", description="Target LLM model identifier")
    provider: Optional[str] = Field(default=None, description="Preferred provider")
    messages: List[Message] = Field(..., min_length=1)
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: int = Field(default=512, ge=1, le=8192)
    top_p: float = Field(default=0.95, ge=0.0, le=1.0)
    top_k: int = Field(default=40, ge=1)
    stop: Optional[List[str]] = Field(default=None)
    stream: bool = Field(default=False)
    manual_restrictions: Optional[List[Dict[str, str]]] = Field(default=None)
    language: Optional[str] = Field(default=None)
    privacy_mode: Optional[bool] = Field(default=None)
    mode: Literal["general", "code", "finance", "legal"] = Field(default="general")
    aggressiveness: float = Field(default=0.2, ge=0.0, le=1.0)
    min_confidence: float = Field(default=0.6, ge=0.0, le=1.0)
    auto_correct: bool = Field(default=True)
    use_perplexity: bool = Field(default=False)

    @field_validator("privacy_mode")
    @classmethod
    def resolve_privacy_mode(cls, value: Optional[bool]) -> bool:
        from .config import settings
        return value if value is not None else settings.privacy_mode_default


class ChatCompletionResponse(BaseModel):
    id: str = Field(default_factory=lambda: f"nlproxy-{uuid.uuid4().hex[:12]}")
    object: str = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: List[Dict[str, Any]]
    usage: Optional[Dict[str, int]] = None
    system_fingerprint: Optional[str] = None
    nlproxy: Optional[Dict[str, Any]] = None


class HealthResponse(BaseModel):
    status: Literal["healthy", "unhealthy"]
    version: str
    timestamp: str
    components: Dict[str, Dict[str, Any]]


class MetricsResponse(BaseModel):
    requests_total: int
    requests_in_progress: int
    cache_hits: int
    cache_misses: int
    avg_latency_ms: float
    error_rate: float
    llm_providers: Dict[str, Dict[str, Any]]