Spaces:
Running
Running
| import json | |
| import pathlib | |
| from enum import StrEnum | |
| from typing import Any, Dict, List, Optional, Set | |
| from pydantic import BaseModel | |
| # ── Authorized Representatives lookup map ────────────────────────────────────── | |
| # Keys: full name (displayed in dropdown) | |
| # Values: dict of data_key → value for the other rep fields | |
| REPRESENTATIVES: Dict[str, Dict[str, str]] = { | |
| "Ahmad Royhan": { | |
| "rep_id": "3174050905990008", | |
| "rep_address": "Jl. Kramat 1 No 15, Kwitang", | |
| "rep_city": "Jakarta Selatan", | |
| }, | |
| # Add more representatives here: | |
| # "Full Name": { | |
| # "rep_id": "...", | |
| # "rep_address": "...", | |
| # "rep_city": "...", | |
| # }, | |
| } | |
| class DocumentType(StrEnum): | |
| COVER_LETTER_SCHENGEN = "cover_letter_schengen" | |
| COVER_LETTER_UK = "cover_letter_uk" | |
| LETTER_OF_AUTHORIZATION = "letter_of_authorization" | |
| LETTER_OF_AUTHORIZATION_MINOR = "letter_of_authorization_minor" | |
| # INVITATION_LETTER = "invitation_letter" | |
| # BANK_REFERENCE = "bank_reference" | |
| class SectionConfig(BaseModel): | |
| """Metadata for one form section, including explicit display order.""" | |
| key: str | |
| title: str | |
| order: int = 0 | |
| # data_keys whose widgets are rendered as read-only in this section | |
| disabled_keys: Set[str] = set() | |
| class FieldSchema(BaseModel): | |
| """Describes one editable field for a document type.""" | |
| widget_key: str # Streamlit session-state key | |
| data_key: str # Key inside the data dict | |
| label: str | |
| placeholder: str = "" | |
| description: Optional[str] = None # Sent to /prefill-document; falls back to placeholder | |
| field_type: str = "text_input" # "text_input" | "text_area" | "select" | |
| options: Optional[List[str]] = None # For field_type="select" | |
| section: str = "general" # Maps to a SectionConfig.key | |
| nested_under: Optional[str] = None # e.g. "personal_details" | |
| order: int = 0 # Explicit render position within the section; lower = earlier | |
| static_value: Optional[str] = None # Pre-seeded value; applied as fallback when DB prefill returns null | |
| class DocumentSchema(BaseModel): | |
| """Full metadata for a document type. | |
| Drives form rendering, session-state key management, and API calls. | |
| Loaded from a JSON file in the src/schemas/ directory. | |
| """ | |
| document_type: DocumentType | |
| title: str | |
| icon: str = "📝" | |
| sections: List[SectionConfig] | |
| fields: List[FieldSchema] | |
| has_trip_type: bool = False | |
| has_group_members: bool = False | |
| def ordered_sections(self) -> List[SectionConfig]: | |
| """Return sections sorted by SectionConfig.order (stable sort).""" | |
| return sorted(self.sections, key=lambda s: s.order) | |
| def section_title(self, key: str) -> str: | |
| """Return the display title for a section key.""" | |
| for s in self.sections: | |
| if s.key == key: | |
| return s.title | |
| return key | |
| def section_disabled_keys(self, key: str) -> Set[str]: | |
| """Return the disabled_keys set for a section key.""" | |
| for s in self.sections: | |
| if s.key == key: | |
| return s.disabled_keys | |
| return set() | |
| def fields_in_section(self, section: str) -> List[FieldSchema]: | |
| """Return fields belonging to section, sorted by FieldSchema.order.""" | |
| return sorted( | |
| [f for f in self.fields if f.section == section], | |
| key=lambda f: f.order, | |
| ) | |
| def top_level_fields(self, section: str) -> List[FieldSchema]: | |
| """Return non-nested fields in section, sorted by order.""" | |
| return sorted( | |
| [f for f in self.fields if f.section == section and f.nested_under is None], | |
| key=lambda f: f.order, | |
| ) | |
| def widget_keys(self) -> List[str]: | |
| """All Streamlit widget keys that must be cleared on undo/redo.""" | |
| keys: List[str] = [] | |
| if self.has_trip_type: | |
| keys.append(f"{self.document_type}_trip_type") | |
| keys.extend(f.widget_key for f in self.fields) | |
| if self.has_group_members: | |
| keys.append(f"{self.document_type}_group_members_editor") | |
| return keys | |
| def nested_fields(self, nested_under: str) -> List[FieldSchema]: | |
| return [f for f in self.fields if f.nested_under == nested_under] | |
| # ── JSON-based schema loader ──────────────────────────────────────────────────── | |
| _SCHEMA_DIR: pathlib.Path = pathlib.Path(__file__).parent / "schemas" | |
| _REPRESENTATIVES_MARKER: str = "$REPRESENTATIVES" | |
| def _resolve_options(options: Any) -> Optional[List[str]]: | |
| """Replace the $REPRESENTATIVES marker with the actual representative names list.""" | |
| if options == _REPRESENTATIVES_MARKER: | |
| return list(REPRESENTATIVES.keys()) | |
| return options # type: ignore[return-value] | |
| def _load_document_schema(path: pathlib.Path) -> DocumentSchema: | |
| """Load and validate a DocumentSchema from a JSON file.""" | |
| raw: Dict[str, Any] = json.loads(path.read_text(encoding="utf-8")) | |
| # Resolve dynamic markers in field options before Pydantic validation | |
| for field_data in raw.get("fields", []): | |
| if "options" in field_data: | |
| field_data["options"] = _resolve_options(field_data["options"]) | |
| return DocumentSchema.model_validate(raw) | |
| def _build_registry() -> Dict[str, DocumentSchema]: | |
| """Scan the schemas/ directory and build DOCUMENT_REGISTRY from JSON files.""" | |
| registry: Dict[str, DocumentSchema] = {} | |
| for json_file in sorted(_SCHEMA_DIR.glob("*.json")): | |
| schema: DocumentSchema = _load_document_schema(json_file) | |
| registry[schema.document_type] = schema | |
| return registry | |
| DOCUMENT_REGISTRY: Dict[str, DocumentSchema] = _build_registry() | |
| # ── API request models ───────────────────────────────────── | |
| class PrefillDocumentRequest(BaseModel): | |
| """ | |
| Sent to /prefill-document. | |
| `structure` describes the fields the LLM should map from DB data. | |
| `sections` (optional) — if provided, the response is nested per section: | |
| { "sections": [{key, title, order, fields: [{key, label, order, value}]}], | |
| "_missing_required": [...] } | |
| If `sections` is omitted, falls back to the old flat dict response. | |
| """ | |
| application_id: int | |
| structure: List[Dict[str, Any]] | |
| sections: Optional[List[Dict[str, Any]]] = None # [{key, title, order}, ...] | |
| class GenerateDocumentRequest(BaseModel): | |
| """ | |
| Sent to /generate-document. | |
| `doc_data` carries all form fields as-is; the backend merges them on top | |
| of the DB data and generates the document. | |
| Adding a new document type requires no change to this model. | |
| """ | |
| application_id: int | |
| document_name: str # value from DocumentType | |
| doc_data: Optional[Dict[str, Any]] = None | |
| class DocChatRequest(BaseModel): | |
| """Sent to /generate-document/chat (new stateless draft-revision flow).""" | |
| query: str | |
| history: List[Dict[str, str]] | |
| session_uuid: Optional[str] = None | |
| current_document: Optional[Dict[str, Any]] = None # structured {title, city_date, main} | |
| current_document_content: str = "" # legacy plain-markdown fallback | |
| structure: Optional[List[Dict[str, Any]]] = None | |
| class GenerateDraftRequest(BaseModel): | |
| """ | |
| Sent to /generate-document/draft. | |
| Generates a Markdown draft without creating a Google Doc. | |
| `doc_type` and `structure` are stored on the frontend for later handover to the core API. | |
| """ | |
| doc_type: str | |
| data: Dict[str, Any] | |
| application_id: Optional[int] = None # used by the API to load persona from DB | |
| structure: Optional[List[Dict[str, Any]]] = None # field definitions for core API handover | |
| session_uuid: Optional[str] = None | |
| class DraftField(BaseModel): | |
| """A single structured field from /generate-document/draft.""" | |
| style: str | |
| content: str | |
| class ExportDocumentRequest(BaseModel): | |
| """Sent to /export-document once the user is satisfied with the draft.""" | |
| document_content: Optional[str] = None # legacy | |
| title: Optional[DraftField] = None | |
| city_date: Optional[DraftField] = None | |
| main: Optional[DraftField] = None | |
| export_format: str = "gdocs" | |
| document_url: Optional[str] = None # existing Google Doc URL; if set, AI edits instead of creating | |
| # ── Shared sub-models (used for type hints / validation internally) ── | |
| class PersonalDetails(BaseModel): | |
| name: Optional[str] = None | |
| dob: Optional[str] = None | |
| nationality: Optional[str] = None | |
| occupation: Optional[str] = None | |
| passport_number: Optional[str] = None | |
| class GroupMember(BaseModel): | |
| relationship: Optional[str] = None | |
| name: Optional[str] = None | |
| dob: Optional[str] = None | |
| occupation: Optional[str] = None | |
| nationality: Optional[str] = None | |
| passport_number: Optional[str] = None | |