"""LLM formalizer: natural language geometry → GeometryIR.""" from __future__ import annotations import json import os from typing import Any import httpx from statement_to_tikz.ir import GeometryIR SYSTEM_PROMPT = """You convert olympiad-style plane geometry statements into a JSON GeometryIR. Rules: - Output ONLY valid JSON matching the schema. No markdown. - List every named point in "points". - Put metric/incidence relations in "constraints" (not coordinates). - Do NOT invent numeric coordinates; optional "hint" only if the statement clearly implies a position. - Use segments for edges to draw; set draw=true on segments that should appear. - Use circles with center+through or through (3 points) for circumcircles. - Constraint types (field "type"): equal_length, length, equal_angle, angle_measure, perpendicular, parallel, on_line, on_circle, midpoint, collinear, intersection. - For equal_length: segments a1-a2 and b1-b2 have equal length. - For midpoint: point is midpoint of a-b. - For perpendicular: segment a1-a2 ⊥ b1-b2. - For intersection: point lies on both lines ab and cd. - Always include labels for all points (or leave labels empty to auto-fill). - Keep the original statement in "statement". """ SCHEMA_HINT = { "statement": "string", "points": [{"name": "A", "hint": None}], "segments": [{"a": "A", "b": "B", "draw": True}], "lines": [{"a": "A", "b": "B", "draw": False}], "circles": [ { "id": "omega", "center": "O", "radius": None, "through": ["A", "B", "C"], "draw": True, } ], "angle_marks": [ {"vertex": "A", "a": "B", "b": "C", "right_angle": False, "label": None} ], "labels": [{"point": "A", "text": "A", "position": "above right"}], "constraints": [ {"type": "equal_length", "a1": "A", "a2": "B", "b1": "A", "b2": "C"}, {"type": "midpoint", "point": "D", "a": "B", "b": "C"}, {"type": "perpendicular", "a1": "A", "a2": "D", "b1": "B", "b2": "C"}, ], } def _extract_json(text: str) -> dict[str, Any]: text = text.strip() if text.startswith("```"): lines = text.split("\n") lines = lines[1:] if lines and lines[-1].strip().startswith("```"): lines = lines[:-1] text = "\n".join(lines).strip() return json.loads(text) def formalize_offline(statement: str) -> GeometryIR | None: """Tiny pattern-based formalizer for demos/tests without an API key. Recognizes a few common phrasings; returns None if not matched. """ s = " ".join(statement.lower().split()) # Isosceles triangle with midpoint median if ( "triangle abc" in s and ("ab = ac" in s or "ab=ac" in s) and "midpoint" in s and ("d" in s) ): return GeometryIR.model_validate( { "statement": statement, "points": [ {"name": "A", "hint": [0, 1.2]}, {"name": "B", "hint": [-1, 0]}, {"name": "C", "hint": [1, 0]}, {"name": "D"}, ], "segments": [ {"a": "A", "b": "B"}, {"a": "B", "b": "C"}, {"a": "C", "b": "A"}, {"a": "A", "b": "D"}, ], "constraints": [ {"type": "equal_length", "a1": "A", "a2": "B", "b1": "A", "b2": "C"}, {"type": "midpoint", "point": "D", "a": "B", "b": "C"}, {"type": "length", "a": "B", "b": "C", "value": 2.0}, {"type": "length", "a": "A", "b": "B", "value": 1.5}, ], "labels": [], } ) if "equilateral" in s and "triangle abc" in s: return GeometryIR.model_validate( { "statement": statement, "points": [ {"name": "A", "hint": [0, 0]}, {"name": "B", "hint": [1, 0]}, {"name": "C", "hint": [0.5, 0.8]}, ], "segments": [ {"a": "A", "b": "B"}, {"a": "B", "b": "C"}, {"a": "C", "b": "A"}, ], "constraints": [ {"type": "equal_length", "a1": "A", "a2": "B", "b1": "B", "b2": "C"}, {"type": "equal_length", "a1": "B", "a2": "C", "b1": "C", "b2": "A"}, {"type": "length", "a": "A", "b": "B", "value": 1.0}, ], "labels": [], } ) if "draw triangle abc" in s or s.strip() == "triangle abc": return GeometryIR.model_validate( { "statement": statement, "points": [ {"name": "A", "hint": [0, 1]}, {"name": "B", "hint": [-1, 0]}, {"name": "C", "hint": [1, 0]}, ], "segments": [ {"a": "A", "b": "B"}, {"a": "B", "b": "C"}, {"a": "C", "b": "A"}, ], "constraints": [], "labels": [], } ) return None def formalize_with_llm( statement: str, *, api_key: str | None = None, base_url: str | None = None, model: str | None = None, timeout: float = 60.0, ) -> GeometryIR: api_key = api_key or os.environ.get("OPENAI_API_KEY") if not api_key: raise RuntimeError( "OPENAI_API_KEY not set; use --ir, offline patterns, or set the key" ) base_url = (base_url or os.environ.get("OPENAI_BASE_URL") or "https://api.openai.com/v1").rstrip( "/" ) model = model or os.environ.get("STT_MODEL") or "gpt-4o-mini" user = ( f"Statement:\n{statement}\n\n" f"JSON shape example (illustrative):\n{json.dumps(SCHEMA_HINT, indent=2)}" ) payload = { "model": model, "temperature": 0, "messages": [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user}, ], "response_format": {"type": "json_object"}, } url = f"{base_url}/chat/completions" headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} last_err: Exception | None = None for _ in range(2): try: with httpx.Client(timeout=timeout) as client: resp = client.post(url, headers=headers, json=payload) resp.raise_for_status() data = resp.json() content = data["choices"][0]["message"]["content"] obj = _extract_json(content) if "statement" not in obj: obj["statement"] = statement ir = GeometryIR.model_validate(obj) ir.ensure_labels() return ir except Exception as exc: # noqa: BLE001 last_err = exc continue raise RuntimeError(f"formalizer failed after retry: {last_err}") def formalize(statement: str, *, allow_offline: bool = True, **kwargs: Any) -> GeometryIR: """Formalize statement via LLM, with optional offline pattern fallback.""" if allow_offline: offline = formalize_offline(statement) if offline is not None: offline.ensure_labels() return offline return formalize_with_llm(statement, **kwargs)