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

from __future__ import annotations

from typing import List, Optional

from pydantic import BaseModel, Field


class NPActivationSingleRequest(BaseModel):
    secret: str = Field(..., description='API secret for the inference server')
    prompt: str = Field(..., description='Input text prompt to get activations for')
    model: str = Field(..., description='Name of the model to test activations on')
    source: Optional[str] = Field(
        None,
        description='Source identifier - could be an SAE ID (eg 5-gemmascope-res-16k). Must be specified with "index", or not at all.',
    )
    index: Optional[str] = Field(None, description='Index of the SAE. Must be specified with "source", or not at all.')
    vector: Optional[List[float]] = Field(
        None, description='Custom vector to test activations. Must be specified with "hook".'
    )
    hook: Optional[str] = Field(
        None, description='Hook that the custom vector applies to. Must be specified with "vector".'
    )


class NPActivationSingleActivation(BaseModel):
    values: List[float]
    max_value: float
    max_value_index: int
    dfa_values: Optional[List[float]] = None
    dfa_max_value: Optional[float] = None
    dfa_target_index: Optional[int] = None


class NPActivationSingleResponse(BaseModel):
    activation: NPActivationSingleActivation
    tokens: List[str]
