| """SHARED CONTRACT — the prompt + output schema the model is trained and served on. |
| |
| This file is the single source of truth for: |
| - FIELD_ORDER : canonical key order of the extracted JSON |
| - build_prompt() : the EXACT prompt string used at train AND inference time |
| - format_completion() : training target serialization (also used by eval) |
| - extract_json() : brace-matched JSON extraction from raw model text |
| - HotelExtraction : Pydantic schema for validation |
| - validate() : parse+validate -> clean dict (or None) |
| |
| It is intentionally DEPENDENCY-LIGHT (only json + pydantic) so both the heavy |
| training env and the lean serving env can import it. A COPY of this file lives in |
| BOTH training/ and deployment/; tests/test_contract_parity.py asserts they are |
| byte-identical, so the prompt can never silently drift between train and serve. |
| |
| DO NOT EDIT one copy without the other. Edit the source, re-sync, re-run the test. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from typing import Any, Literal, Optional |
|
|
| from pydantic import BaseModel, ConfigDict, Field, ValidationError |
|
|
| DEFAULT_BASE_MODEL = "Qwen/Qwen3-1.7B" |
|
|
| |
| |
| |
| FIELD_ORDER = [ |
| "destination", "locality", "hotelName", "checkinDate", "checkoutDate", |
| "adultCount", "roomCount", "childCount", "childAges", "infantCount", |
| "sortCriteria", "minPrice", "maxPrice", "filters", "deepSearch", "isNearMe", "resetAction", |
| ] |
|
|
|
|
| def build_prompt(query: str, today: str) -> str: |
| """The exact user-message content used for BOTH training and inference.""" |
| return ( |
| "Extract hotel-search entities.\n" |
| "Return strict JSON only.\n" |
| "Schema:\n" |
| "destination, locality, hotelName, checkinDate, checkoutDate, " |
| "adultCount, roomCount, childCount, childAges, infantCount, " |
| "sortCriteria, minPrice, maxPrice, filters, deepSearch, isNearMe, resetAction.\n\n" |
| "Rules:\n" |
| "- Dates are DDMMYYYY.\n" |
| '- deepSearch and isNearMe must be "true" or "false".\n' |
| "- filters is a comma-separated list of amenity/type phrases (e.g. " |
| '"swimming pool, pet friendly"); prefix removals with "no ".\n' |
| "- Omit unknown optional fields.\n\n" |
| f"today={today}\n" |
| f"query={query}" |
| ) |
|
|
|
|
| def format_completion(expected: dict[str, Any]) -> str: |
| """Stable-key-order minified JSON — the training target / gold serialization.""" |
| ordered = {f: expected[f] for f in FIELD_ORDER if f in expected} |
| for k in expected: |
| if k not in ordered: |
| ordered[k] = expected[k] |
| return json.dumps(ordered, ensure_ascii=False, separators=(",", ":")) |
|
|
|
|
| def extract_json(text: str) -> dict | None: |
| """Pull the first balanced {...} object out of raw model text.""" |
| start = text.find("{") |
| if start < 0: |
| return None |
| depth = 0 |
| in_str = esc = False |
| for i in range(start, len(text)): |
| c = text[i] |
| if esc: |
| esc = False |
| continue |
| if c == "\\": |
| esc = True |
| continue |
| if c == '"': |
| in_str = not in_str |
| continue |
| if in_str: |
| continue |
| if c == "{": |
| depth += 1 |
| elif c == "}": |
| depth -= 1 |
| if depth == 0: |
| try: |
| return json.loads(text[start:i + 1]) |
| except json.JSONDecodeError: |
| return None |
| return None |
|
|
|
|
| class HotelExtraction(BaseModel): |
| """Output contract. `filters` is a comma-separated PHRASE string (e.g. |
| "swimming pool, pet friendly"); a downstream matcher resolves it to codes.""" |
| model_config = ConfigDict(extra="forbid") |
| deepSearch: Literal["true", "false"] |
| isNearMe: Literal["true", "false"] |
| destination: Optional[str] = Field(default=None, min_length=1) |
| locality: Optional[str] = Field(default=None, min_length=1) |
| hotelName: Optional[str] = Field(default=None, min_length=1) |
| checkinDate: Optional[str] = Field(default=None, pattern=r"^\d{8}$") |
| checkoutDate: Optional[str] = Field(default=None, pattern=r"^\d{8}$") |
| adultCount: Optional[int] = Field(default=None, ge=0, le=20) |
| roomCount: Optional[int] = Field(default=None, ge=0, le=20) |
| childCount: Optional[int] = Field(default=None, ge=0, le=20) |
| childAges: Optional[list[int]] = None |
| infantCount: Optional[int] = Field(default=None, ge=0, le=20) |
| sortCriteria: Optional[Literal["SC_P_LH", "SC_P_HL", "SC_UR", "SC_P", "SC_DIST"]] = None |
| minPrice: Optional[int] = Field(default=None, ge=0, le=10_000_000) |
| maxPrice: Optional[int] = Field(default=None, ge=0, le=10_000_000) |
| filters: Optional[str] = Field(default=None, min_length=1) |
| resetAction: Optional[Literal["filters", "guests", "dates", "all"]] = None |
|
|
|
|
| def validate(d: dict) -> dict | None: |
| """Validate a parsed dict against the schema; return clean dict or None.""" |
| try: |
| return HotelExtraction.model_validate(d).model_dump(exclude_none=True) |
| except ValidationError: |
| return None |
|
|
