File size: 5,312 Bytes
71d239c 7603aa2 71d239c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """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"],
}
|