# generated by datamodel-codegen:
#   filename:  bench.yaml
#   timestamp: 2023-12-15T18:09:28+00:00

from __future__ import annotations

from datetime import datetime
from typing import List, Optional
from uuid import UUID

from arthur.client.rest.common.models import Page, PageSize, TotalCount, TotalPages
from pydantic import BaseModel, Field


class TestCaseRequest(BaseModel):
    """
    An input, reference output pair.
    """

    input: str
    """
    Input to the test case. Does not include the prompt template.
    """
    reference_output: Optional[str] = None
    """
    Reference or "Golden" output for the given input.
    """


class TestCaseResponseItem(BaseModel):
    id: Optional[UUID] = None
    input: Optional[str] = None
    """
    Input to the test case. Does not include the prompt template.
    """
    reference_output: Optional[str] = None
    """
    Reference or "Golden" output for the given input.
    """


class TestCaseResponse(BaseModel):
    """
    A test case (input, reference output pair).
    """

    __root__: List[TestCaseResponseItem]
    """
    A test case (input, reference output pair).
    """


class HistogramItem(BaseModel):
    count: int
    low: float
    high: float


class SummaryItem(BaseModel):
    name: str
    id: UUID
    avg_score: float
    histogram: List[HistogramItem]


class TestSuiteSummaryResponse(BaseModel):
    num_test_cases: float
    summary: List[SummaryItem]
    page: Page
    page_size: PageSize
    total_pages: TotalPages
    total_count: TotalCount


class TestCaseOutput(BaseModel):
    id: Optional[UUID] = None
    output: Optional[str] = None
    score: Optional[float] = None


class CreateRunRequest(BaseModel):
    test_case_outputs: List[TestCaseOutput]
    description: Optional[str] = None
    model_name: Optional[str] = None
    name: str
    foundation_model: Optional[str] = None
    prompt_template: Optional[str] = None
    model_version: Optional[str] = None


class CreateRunResponse(BaseModel):
    id: Optional[UUID] = None


class RunResult(BaseModel):
    id: str
    input: str
    reference_output: Optional[str] = None
    output: str
    score: float


class TestRun(BaseModel):
    id: Optional[UUID] = None
    name: Optional[str] = None
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None
    model_version: Optional[str] = None
    prompt_template: Optional[str] = None
    avg_score: Optional[float] = None


class PaginatedGetRunsForTestSuiteResponse(BaseModel):
    """
    Paginated list of runs for a test suite.
    """

    test_runs: List[TestRun]
    page: Page
    page_size: PageSize
    total_pages: TotalPages
    total_count: TotalCount


class ScoringMethod(BaseModel):
    name: str
    type: str


class PaginatedGetRunResponse(BaseModel):
    """
    Paginated list of prompts, reference outputs, and model outputs for a particular run.
    """

    id: UUID
    created_at: datetime
    updated_at: datetime
    description: Optional[str] = None
    name: Optional[str] = None
    metadata: Optional[str] = None
    test_case_runs: List[RunResult]
    test_suite_id: UUID
    page: Page
    page_size: PageSize
    total_pages: Optional[TotalPages] = None
    total_count: Optional[TotalCount] = None


class HallucinationCheckRequest(BaseModel):
    """
    Model generated response and context used to determine whether the response is a hallucination
    """

    response: str
    """
    model generated response string
    """
    context: str
    """
    string providing context in which model generated response should be consistent with
    """


class HallucinationCheckResponse(BaseModel):
    """
    Binary hallucination classification and reasoning for classification
    """

    hallucination: bool
    """
    if the response contained a hallucination
    """
    reason: str
    """
    justification of the hallucination classification
    """


class TestSuite(BaseModel):
    id: UUID
    name: str
    scoring_method: ScoringMethod
    last_run_time: Optional[datetime] = None
    description: Optional[str] = None
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None


class PaginatedGetTestSuitesResponse(BaseModel):
    test_suites: List[TestSuite]
    page: Page
    page_size: PageSize
    total_pages: TotalPages
    total_count: TotalCount


class TestSuiteRequest(BaseModel):
    name: str
    description: Optional[str] = None
    scoring_method: ScoringMethod
    test_cases: List[TestCaseRequest] = Field(..., min_items=1)
    """
    JSON object containing test case data for the test suite.
    """


class TestSuiteResponse(BaseModel):
    id: UUID
    name: str
    description: Optional[str] = None
    scoring_method: ScoringMethod
    organization_id: Optional[UUID] = None
    created_at: datetime
    updated_at: datetime
    test_cases: List[TestCaseResponse] = Field(..., min_items=1)
    """
    JSON object containing test case data for the test suite.
    """


class PaginatedGetTestSuiteResponse(BaseModel):
    id: UUID
    name: str
    scoring_method: ScoringMethod
    test_cases: List[TestCaseResponse]
