Source code for afterimage.types

import uuid
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import (
    Any,
    Dict,
    Generic,
    List,
    Literal,
    Optional,
    TypeAlias,
    TypedDict,
    TypeVar,
)

from pydantic import BaseModel, Field

from .metadata_utils import extract_unique_context_ids

T = TypeVar("T")

ModelProviderName: TypeAlias = Literal[
    "gemini",
    "openai",
    "deepseek",
    "local",
    "openrouter",
]

MODEL_PROVIDER_NAMES: frozenset[str] = frozenset(
    ("gemini", "openai", "deepseek", "local", "openrouter")
)




class GradeSchema(str, Enum):
    PERFECT = "perfect"
    GOOD = "good"
    NEEDS_IMPROVEMENT = "needs_improvement"
    BAD = "bad"
    NOT_ACCEPTABLE = "not_acceptable"





class EvaluationEntrySchema(BaseModel):
    feedback: str
    score: float





class EvaluationSchema(BaseModel):
    coherence: EvaluationEntrySchema
    factuality: EvaluationEntrySchema
    grounding: EvaluationEntrySchema
    helpfulness: EvaluationEntrySchema
    relevance: EvaluationEntrySchema
    overall_grade: GradeSchema





class Role(str, Enum):
    USER = "user"
    ASSISTANT = "assistant"





class ConversationEntry(BaseModel):
    role: Role
    content: str
    reasoning_content: str | None = None





class Conversation(BaseModel):
    conversations: List[ConversationEntry]
    metadata: dict[str, Any] = Field(default_factory=dict)





class ConversationWithContext(Conversation):
    instruction_context: str | None = None
    response_context: str | None = None
    persona: str | None = None





class EvaluatedConversationWithContext(ConversationWithContext):
    evaluation: EvaluationSchema | None = None
    final_score: Optional[float] = 0.0





class PersonaEntry(BaseModel):
    """Represents a set of generated personas for a source document."""

    descriptions: list[str] = Field(..., description="List of persona descriptions")
    metadata: dict = Field(default_factory=dict, description="Any associated metadata")





class StructuredGenerationRow(BaseModel, Generic[T]):
    """Represents a single row of structured generation, including metadata."""

    instruction: str = Field(
        ..., description="The user instructions/prompt used to generate the output."
    )
    context: Optional[str] = Field(
        None, description="The context provided to the model."
    )
    persona: Optional[str] = Field(
        None, description="The persona adopted by the respondent."
    )
    output: T = Field(..., description="The structured output generated by the model.")
    metadata: Dict[str, Any] = Field(
        default_factory=dict, description="Metadata about the generation."
    )





class GeneratedResponsePrompt(BaseModel):
    """Output of RespondentPromptModifier."""

    prompt: str = Field(..., description="Modified respondent prompt")
    context: Optional[str] = Field(
        None, description="Context used in respondent prompt"
    )
    metadata: dict[str, Any] = Field(
        default_factory=dict,
        description=(
            "Additional metadata about respondent prompt generation. "
            "When using :class:`~afterimage.callbacks.WithRAGRespondentPromptModifier`, "
            "retrievers that implement ``get_context_with_metadata`` / "
            "``aget_context_with_metadata`` may attach citation-style payloads under "
            "the key ``retrieval`` (see :data:`afterimage.retrievers.RETRIEVAL_METADATA_KEY`)."
        ),
    )





class Document(BaseModel):
    id: str = Field(
        default_factory=lambda: str(uuid.uuid4()),
        description="Id of the document. Autogenerated if not specified",
    )
    text: Optional[str] = None
    personas: list[PersonaEntry] = Field(
        default_factory=list, description="List of persona generations"
    )
    metadata: dict[str, Any] = Field(
        default_factory=dict, description="Any associated metadata"
    )





class GenerationMetadata(TypedDict, total=False):
    """Structured metadata for generation samples."""

    context_id: str
    context_ids: list[str]
    persona_name: Optional[str]
    persona_generation_depth: Optional[int]
    instruction_index: int
    batch_id: str
    session_id: str





@dataclass
class GenerationState:
    """Current state of the generation process."""

    num_generated: int = 0
    num_requested: int = 0
    start_time: datetime = field(default_factory=datetime.now)
    last_item: Optional[BaseModel] = None
    monitor: Optional[Any] = None  # GenerationMonitor
    metadata: Dict[str, Any] = field(default_factory=dict)
    stop_event: Optional[Any] = None  # asyncio.Event

    # Context coverage tracking
    context_counts: Dict[str, int] = field(default_factory=dict)
    unique_personas: list[str] = field(default_factory=list)



    def update(self, item: Any):
        """Update state with a new generated item."""
        self.num_generated += 1
        self.last_item = item if isinstance(item, BaseModel) else None

        # Extract metadata if available
        meta = None
        if hasattr(item, "metadata"):
            meta = item.metadata
        elif isinstance(item, dict):
            meta = item.get("metadata")

        if meta is not None and isinstance(meta, dict):
            for ctx_id in extract_unique_context_ids(meta):
                self.context_counts[ctx_id] = self.context_counts.get(ctx_id, 0) + 1

        persona = getattr(item, "persona", None) or (
            meta.get("persona_name") if isinstance(meta, dict) else None
        )
        if persona and persona not in self.unique_personas:
            self.unique_personas.append(persona)