src/ai_agent/generator/schema.py

# generator/schema.py
from __future__ import annotations

from enum import Enum
from typing import List, Optional
from pydantic import BaseModel, Field, ConfigDict, field_validator, model_validator


class SupportingData(BaseModel):
    model_config = ConfigDict(populate_by_name=True, extra="ignore")

    content_url: Optional[str] = Field(default=None, alias="contentUrl")
    description: Optional[str] = None
    name: Optional[str] = None
    dataset_format: Optional[str] = Field(default=None, alias="datasetFormat")
    has_dimensionality: Optional[int] = Field(default=None, alias="hasDimensionality")
    body_site: Optional[str] = Field(default=None, alias="bodySite")
    imaging_modality: Optional[str] = Field(default=None, alias="imagingModality")


class RunnableExample(BaseModel):
    model_config = ConfigDict(populate_by_name=True, extra="ignore")

    description: Optional[str] = None
    name: Optional[str] = None
    url: Optional[str] = None
    host_type: Optional[str] = Field(default=None, alias="hostType")


class ExecutableNotebook(BaseModel):
    model_config = ConfigDict(populate_by_name=True, extra="ignore")

    description: Optional[str] = None
    name: Optional[str] = None
    url: Optional[str] = None


class CandidateDoc(BaseModel):
    """
    Minimal view of a software tool passed to the generator (aligned with SoftwareDoc).
    All fields optional/lenient to tolerate catalog variation.
    """

    model_config = ConfigDict(populate_by_name=True, extra="ignore")

    # Core identity
    id: Optional[str] = None
    name: Optional[str] = None
    description: Optional[str] = None
    url: Optional[str] = None

    # Categories & semantics
    category: List[str] = Field(default_factory=list, alias="applicationCategory")
    tasks: List[str] = Field(default_factory=list, alias="featureList")
    keywords: List[str] = Field(default_factory=list)

    # Modality / anatomy / dimensionality
    modality: List[str] = Field(default_factory=list, alias="imagingModality")
    dims: List[int] = Field(default_factory=list)  # normalized to ints (2,3,4)
    anatomy: List[str] = Field(default_factory=list)  # from supportingData[*].bodySite

    # Tech details
    programming_language: Optional[str] = Field(
        default=None, alias="programmingLanguage"
    )
    software_requirements: List[str] = Field(
        default_factory=list, alias="softwareRequirements"
    )
    gpu_required: Optional[bool] = Field(default=None, alias="requiresGPU")
    is_free: Optional[bool] = Field(default=None, alias="isAccessibleForFree")
    is_based_on: List[str] = Field(default_factory=list, alias="isBasedOn")
    plugin_of: List[str] = Field(
        default_factory=list, alias="isPluginModuleOf"
    )  # <-- CHANGED to list
    related_organizations: List[str] = Field(
        default_factory=list, alias="relatedToOrganization"
    )

    # Rich, nested metadata
    supporting_data: List[SupportingData] = Field(
        default_factory=list, alias="supportingData"
    )
    runnable_examples: List[RunnableExample] = Field(
        default_factory=list, alias="runnableExample"
    )
    executable_notebooks: List[ExecutableNotebook] = Field(
        default_factory=list, alias="hasExecutableNotebook"
    )

    # -------- Normalizers --------
    @field_validator("plugin_of", mode="before")
    @classmethod
    def _coerce_plugin_of(cls, v):
        if v is None:
            return []
        if isinstance(v, list):
            return [
                str(x).strip()
                for x in v
                if isinstance(x, (str, int, float)) and str(x).strip()
            ]
        # single value -> list
        s = str(v).strip()
        return [s] if s else []

    @field_validator("dims", mode="before")
    @classmethod
    def _coerce_dims(cls, v):
        """
        Accepts: int/float, '3D'/'2-D', '3', 'volumetric','stack','timeseries', lists, or comma strings.
        Returns de-duped list[int].
        """
        if v is None:
            return []
        items = []
        if isinstance(v, list):
            items = v
        elif isinstance(v, str):
            parts = [p.strip() for p in v.split(",") if p.strip()]
            items = parts or [v]
        else:
            items = [v]

        out: list[int] = []

        def push(x):
            try:
                xi = int(x)
                if xi not in out:
                    out.append(xi)
            except Exception:
                pass

        for it in items:
            if isinstance(it, (int, float)):
                push(it)
                continue
            if not isinstance(it, str):
                continue
            s = it.strip().lower().replace(" ", "")
            if s in {"2", "2d", "2-d"}:
                push(2)
                continue
            if s in {"3", "3d", "3-d", "volume", "volumetric", "stack"}:
                push(3)
                continue
            if s in {"4", "4d", "4-d", "timeseries", "time-series", "temporal"}:
                push(4)
                continue
            digits = "".join(ch for ch in s if ch.isdigit())
            if digits:
                push(digits)
        return out


class PlanAndCode(BaseModel):
    """
    Back-compat schema for the older 'plan + code' generator.
    Your pipeline can ignore steps/code, but we keep it so imports don't fail.
    """

    choice: str
    alternates: List[str] = Field(default_factory=list)
    why: str
    steps: List[str] = Field(default_factory=list)
    code: str = ""


class NoToolReason(str, Enum):
    NO_SUITABLE_TOOL = "no_suitable_tool"
    NO_MODALITY_MATCH = "no_modality_match"
    NO_TASK_MATCH = "no_task_match"
    NO_DIMENSION_MATCH = "no_dimension_match"
    INVALID_FILES = "invalid_files"


class ConversationStatus(str, Enum):
    NEEDS_CLARIFICATION = "needs_clarification"
    COMPLETE = "complete"


class Conversation(BaseModel):
    status: ConversationStatus
    question: Optional[str] = None
    context: Optional[str] = None
    options: Optional[List[str]] = None

    @field_validator("status", mode="before")
    @classmethod
    def _coerce_status(cls, v):
        if isinstance(v, ConversationStatus):
            return v
        s = str(v or "").strip().lower().replace("-", "_").replace(" ", "_")
        if s in {"complete", "done", "finished"}:
            return ConversationStatus.COMPLETE
        if s in {"needs_clarification", "need_clarification", "clarification", "question"}:
            return ConversationStatus.NEEDS_CLARIFICATION
        return v

    @model_validator(mode="after")
    def validate_fields(self) -> "Conversation":
        if self.status == ConversationStatus.NEEDS_CLARIFICATION:
            if not self.question:
                raise ValueError("Question required when status is needs_clarification")
            if not self.context:
                raise ValueError("Context required when status is needs_clarification")
        return self


class ToolChoice(BaseModel):
    name: str
    rank: int
    accuracy: float = Field(ge=0, le=100)  # accuracy score between 0-100
    why: str
    demo_link: Optional[str] = None

    @field_validator("rank", mode="before")
    @classmethod
    def _coerce_rank(cls, v):
        if isinstance(v, int):
            return max(1, v)
        try:
            return max(1, int(float(str(v).strip())))
        except Exception:
            return 1

    @field_validator("accuracy", mode="before")
    @classmethod
    def _coerce_accuracy(cls, v):
        try:
            x = float(v)
        except Exception:
            return 50.0
        # Accept 0..1 confidence and scale to 0..100 when likely intended.
        if 0.0 <= x <= 1.0:
            x *= 100.0
        if x < 0:
            return 0.0
        if x > 100:
            return 100.0
        return x


class ToolSelection(BaseModel):
    conversation: Conversation
    choices: List[ToolChoice] = Field(default_factory=list)
    explanation: Optional[str] = None
    reason: Optional[NoToolReason] = None

    @model_validator(mode="after")
    def normalize(self) -> "ToolSelection":
        has_choices = len(self.choices) > 0
        asked_q = bool(
            self.conversation.question and str(self.conversation.question).strip()
        )

        if asked_q:
            # Clarification turn
            self.conversation.status = ConversationStatus.NEEDS_CLARIFICATION
            self.reason = None
            self.explanation = None
            return self

        if has_choices:
            # Normal success: per-tool 'why' explains rationale; no global reason
            self.conversation.status = ConversationStatus.COMPLETE
            self.reason = None
            # explanation optional but usually unnecessary; keep if you want, or null it:
            # self.explanation = None
            return self

        # No choices
        if self.reason or (self.explanation and self.explanation.strip()):
            # Terminal no-tool
            self.conversation.status = ConversationStatus.COMPLETE
            return self

        # Ambiguous empty -> ask a question
        self.conversation.status = ConversationStatus.NEEDS_CLARIFICATION
        return self

    @field_validator("reason", mode="before")
    def _reason_empty_to_none(cls, v):
        if v is None:
            return None
        s = str(v).strip().lower().replace("-", "_").replace(" ", "_")
        if not s:
            return None
        if s in {"none", "null", "n/a", "na"}:
            return None
        allowed = {
            "no_suitable_tool",
            "no_modality_match",
            "no_task_match",
            "no_dimension_match",
            "invalid_files",
        }
        return s if s in allowed else None

    @field_validator("explanation", mode="before")
    def _expl_empty_to_none(cls, v):
        return None if v is None or str(v).strip() == "" else v


__all__ = [
    "CandidateDoc",
    "PlanAndCode",
    "ToolSelection",
    "Conversation",
    "ConversationStatus",
    "ToolChoice",
    "NoToolReason",
]