# generated by datamodel-codegen:
#   filename:  steer.yaml

from __future__ import annotations

from enum import Enum
from typing import List, Optional

from pydantic import BaseModel, Field, confloat, conint


class NPFeature(BaseModel):
    model: str = Field(..., example='gpt2-small')
    source: str = Field(..., example='0-res-jb')
    index: int = Field(..., example=14057)


class NPSteerVector(BaseModel):
    steering_vector: List[float]
    strength: float
    hook: str = Field(..., example='blocks.0.hook_resid_pre')


class NPSteerType(Enum):
    STEERED = 'STEERED'
    DEFAULT = 'DEFAULT'


class NPSteeringMethod(Enum):
    simple_additive = 'simple_additive'
    orthogonal_decomp = 'orthogonal_decomp'


class NPSteerCompletionResponse(BaseModel):
    STEERED: Optional[str] = Field(
        None, description='The full text of the steered prompt and response, concatenated as one string.'
    )
    DEFAULT: Optional[str] = Field(
        None, description='The full text of the default (non-steered) prompt and response, concatenated as one string.'
    )


class NPSteerChatTemplate(BaseModel):
    content: str = Field(..., description='The chat message')
    role: str = Field(..., description='The role of the message (eg "model", "user", etc)')


class NPSteerChatResult(BaseModel):
    chat_template: List[NPSteerChatTemplate]
    raw: str


class NPSteerCompletionChatResponse(BaseModel):
    STEERED: Optional[NPSteerChatResult] = None
    DEFAULT: Optional[NPSteerChatResult] = None
    input: Optional[NPSteerChatResult] = None


class NPSteerFeature(NPFeature):
    strength: float = Field(..., example=5)
    steering_vector: Optional[List[float]] = None


class NPSteerCompletionRequest(BaseModel):
    secret: str = Field(..., description='API secret for the inference server')
    prompt: str = Field(..., description='Text to pass the model for completion')
    model: str = Field(..., description='Name of the model')
    steering_method: NPSteeringMethod
    normalize_steering: bool
    types: List[NPSteerType] = Field(
        ...,
        description='Array that specifies whether or not to generate STEERED output, DEFAULT (non-steered) output, or both.',
        min_items=1,
    )
    features: Optional[List[NPSteerFeature]] = Field(None, description='Features to steer towards or away from')
    vectors: Optional[List[NPSteerVector]] = None
    n_completion_tokens: conint(ge=1) = Field(..., description='Number of completion tokens to generate')
    temperature: confloat(ge=0.0)
    strength_multiplier: float = Field(..., description='The steering strength will be multiplied by this number')
    freq_penalty: float
    seed: float


class NPSteerCompletionChatRequest(NPSteerCompletionRequest):
    steer_special_tokens: bool
    prompt: List[NPSteerChatTemplate] = Field(..., description='Array of chat messages to pass to the model')
