schema.py

"""Pydantic data model for StoryCode.

Two kinds of data live here:

1. The **ProjectModel** — the deterministic, factual structure of the uploaded
   project, produced entirely by static analysis (`analyzer/`). This is the
   single source of truth. The language model never edits it.

2. The model's **narration** (`FileSummary`, `ProjectStory`) — plain-language
   text the model writes *about* the ProjectModel. The JSON schemas at the
   bottom are handed to vLLM as `guided_json` so the model is forced to emit
   conforming JSON (no fence-stripping / repair downstream).
"""
from __future__ import annotations

from pydantic import BaseModel, Field


# --- Factual structure (from static analysis) -------------------------------
class Symbol(BaseModel):
    """A function or class defined in a file."""
    kind: str = Field(description="'function' or 'class'")
    name: str
    lineno: int = 0
    doc: str | None = Field(default=None, description="First line of its docstring/comment")


class FileInfo(BaseModel):
    path: str = Field(description="Project-relative path, e.g. 'analyzer/graph.py'")
    lang: str = Field(description="language key from config.LANG_BY_EXT, or 'other'")
    role: str = Field(default="other", description="role key from config.ROLES")
    loc: int = Field(default=0, description="non-blank lines of code")
    symbols: list[Symbol] = Field(default_factory=list)
    imports: list[str] = Field(default_factory=list, description="raw import targets")
    # Resolved intra-project edges (filled by graph.py): paths this file depends on.
    depends_on: list[str] = Field(default_factory=list)
    # Graph metrics (filled by graph.py).
    fan_in: int = Field(default=0, description="how many project files depend on this one")
    fan_out: int = Field(default=0, description="how many project files this one depends on")
    is_entry: bool = False
    # Safety verdict (config.SAFE / CAREFUL / DANGER) + plain reason.
    safety: str = Field(default="careful")
    safety_reason: str = ""


class Dependency(BaseModel):
    name: str
    manifest: str = Field(description="which file declared it, e.g. 'requirements.txt'")
    plain: str = Field(default="", description="one-line plain-English description")
    risky: bool = Field(default=False, description="flagged old/insecure (best-effort)")


class ProjectModel(BaseModel):
    """The whole analysed project — the factual backbone of every story."""
    name: str = "your project"
    files: list[FileInfo] = Field(default_factory=list)
    deps: list[Dependency] = Field(default_factory=list)
    entry_points: list[str] = Field(default_factory=list)
    languages: list[str] = Field(default_factory=list, description="distinct langs present")
    skipped: list[str] = Field(default_factory=list, description="paths listed but not parsed")
    secrets_found: list[str] = Field(default_factory=list, description="redacted secret hits")
    note: str = Field(default="", description="user-facing ingest caveat, if any")

    def by_path(self, path: str) -> FileInfo | None:
        return next((f for f in self.files if f.path == path), None)


# --- Model narration (text the model writes about the structure) ------------
class FileSummary(BaseModel):
    """The MAP step output: a 2-3 sentence plain summary of one file."""
    path: str
    one_liner: str = Field(description="<=8 words, the file's job, plain English")
    summary: str = Field(description="2-3 sentences, no jargon")


class StorySection(BaseModel):
    heading: str
    body: str


class ProjectStory(BaseModel):
    """The REDUCE step output: the narrated story of the whole project."""
    title: str
    overview: str = Field(description="2-4 sentences: what this project is, plainly")
    steps: list[StorySection] = Field(
        default_factory=list,
        description="the step-by-step flow: what happens when the app runs / a user acts")
    plain_overview: str = Field(
        default="",
        description="the SAME content in flat Simple-Walkthrough English (always shown "
                    "beside a creative style so the metaphor is never the only version)")


# --- guided_json schemas (vLLM XGrammar-friendly subset) ---------------------
def file_summary_schema() -> dict:
    return {
        "type": "object",
        "additionalProperties": False,
        "properties": {
            "one_liner": {"type": "string"},
            "summary": {"type": "string"},
        },
        "required": ["one_liner", "summary"],
    }


def project_story_schema() -> dict:
    return {
        "type": "object",
        "additionalProperties": False,
        "properties": {
            "title": {"type": "string"},
            "overview": {"type": "string"},
            "steps": {
                "type": "array",
                "items": {
                    "type": "object",
                    "additionalProperties": False,
                    "properties": {
                        "heading": {"type": "string"},
                        "body": {"type": "string"},
                    },
                    "required": ["heading", "body"],
                },
            },
            "plain_overview": {"type": "string"},
        },
        "required": ["title", "overview", "steps", "plain_overview"],
    }