src/generate.py

"""LLM generation: build prompt, call HuggingFace Inference API, return code + explanation."""
from __future__ import annotations

from huggingface_hub import InferenceClient

from src.citations import Citation, build_context_block, extract_inline_refs, render_citation_markdown
from src.config import DEFAULT_MODEL

# ── few-shot examples ─────────────────────────────────────────────────────────

_FEW_SHOT_BOX = '''
### Example 1 — Parametric box with fillets

User: Create a parametric box width=50, height=30, depth=20 with 5mm fillets on vertical edges.

```python
import FreeCAD as App
import Part, Sketcher

doc = App.newDocument("ParametricBox")

body   = doc.addObject("PartDesign::Body", "Body")
sketch = body.newObject("Sketcher::SketchObject", "RectSketch")
sketch.Support = (body.Origin.OriginFeatures[3], [""])  # XY_Plane
sketch.MapMode = "FlatFace"

width, depth, height, fillet_r = 50.0, 20.0, 30.0, 5.0

sketch.addGeometry(Part.LineSegment(App.Vector(0,0,0),     App.Vector(width,0,0)),    False)
sketch.addGeometry(Part.LineSegment(App.Vector(width,0,0), App.Vector(width,depth,0)), False)
sketch.addGeometry(Part.LineSegment(App.Vector(width,depth,0), App.Vector(0,depth,0)), False)
sketch.addGeometry(Part.LineSegment(App.Vector(0,depth,0), App.Vector(0,0,0)),         False)
sketch.addConstraint(Sketcher.Constraint("Coincident", 0,2, 1,1))
sketch.addConstraint(Sketcher.Constraint("Coincident", 1,2, 2,1))
sketch.addConstraint(Sketcher.Constraint("Coincident", 2,2, 3,1))
sketch.addConstraint(Sketcher.Constraint("Coincident", 3,2, 0,1))
sketch.addConstraint(Sketcher.Constraint("Horizontal", 0))
sketch.addConstraint(Sketcher.Constraint("Horizontal", 2))
sketch.addConstraint(Sketcher.Constraint("Vertical",   1))
sketch.addConstraint(Sketcher.Constraint("Vertical",   3))
sketch.addConstraint(Sketcher.Constraint("DistanceX", 0, 1, 0, 2, width))
sketch.addConstraint(Sketcher.Constraint("DistanceY", 1, 1, 1, 2, depth))
doc.recompute()

pad = body.newObject("PartDesign::Pad", "Pad")
pad.Profile = sketch
pad.Length  = height
doc.recompute()

fillet = body.newObject("PartDesign::Fillet", "Fillet")
fillet.Base   = (pad, ["Edge1","Edge2","Edge3","Edge4"])
fillet.Radius = fillet_r
doc.recompute()

doc.saveAs("output.FCStd")
```
'''.strip()

_FEW_SHOT_REVOLVE = '''
### Example 2 — Revolved profile (cylinder / shaft)

User: Create a 20mm-diameter, 60mm-long cylindrical shaft using Revolution.

```python
import FreeCAD as App
import Part, Sketcher

doc  = App.newDocument("Shaft")
body = doc.addObject("PartDesign::Body", "Body")

sketch = body.newObject("Sketcher::SketchObject", "Profile")
sketch.Support = (body.Origin.OriginFeatures[4], [""])  # XZ_Plane
sketch.MapMode = "FlatFace"

radius, length = 10.0, 60.0

sketch.addGeometry(Part.LineSegment(App.Vector(0,0,0),      App.Vector(radius,0,0)),      False)
sketch.addGeometry(Part.LineSegment(App.Vector(radius,0,0), App.Vector(radius,length,0)), False)
sketch.addGeometry(Part.LineSegment(App.Vector(radius,length,0), App.Vector(0,length,0)), False)
sketch.addGeometry(Part.LineSegment(App.Vector(0,length,0), App.Vector(0,0,0)),           False)
sketch.addConstraint(Sketcher.Constraint("Coincident", 0,2, 1,1))
sketch.addConstraint(Sketcher.Constraint("Coincident", 1,2, 2,1))
sketch.addConstraint(Sketcher.Constraint("Coincident", 2,2, 3,1))
sketch.addConstraint(Sketcher.Constraint("Coincident", 3,2, 0,1))
sketch.addConstraint(Sketcher.Constraint("DistanceX", 0, 1, 0, 2, radius))
sketch.addConstraint(Sketcher.Constraint("DistanceY", 1, 1, 1, 2, length))
sketch.addConstraint(Sketcher.Constraint("PointOnObject", 0, 1, -1))  # origin on Y-axis
doc.recompute()

rev = body.newObject("PartDesign::Revolution", "Revolution")
rev.Profile     = sketch
rev.ReferenceAxis = (sketch, ["V_Axis"])
rev.Angle       = 360.0
doc.recompute()

doc.saveAs("output.FCStd")
```
'''.strip()

# ── system prompt ─────────────────────────────────────────────────────────────

_SYSTEM_PROMPT = f"""You are an expert FreeCAD 1.1 Python scripting assistant specialised in \
parametric solid modelling with the PartDesign and Sketcher workbenches.

OUTPUT CONTRACT (strict):
1. Return ONE complete, self-contained Python script enclosed in a single ```python ... ``` block, \
runnable with `freecadcmd script.py`.
2. The script MUST:
   - import FreeCAD as App, Part, Sketcher (never import PartDesignGui / FreeCADGui / SketcherGui — they crash headless)
   - call App.newDocument(...)
   - create a PartDesign::Body BEFORE any Sketch
   - attach every Sketch to a standard plane from body.Origin.OriginFeatures (index 3=XY, 4=XZ, 5=YZ)
   - call doc.recompute() after EVERY feature creation
   - end with doc.saveAs("output.FCStd")
3. Use named variables for every dimension so the model is parametric.
4. Reference geometry by INDEX where possible (e.g. Sketcher.Constraint("Coincident", 0, 2, 1, 1)), \
NOT by topological name strings like "Face1" or "Edge3", to minimise Topological Naming Problem risk \
(mitigated but NOT eliminated in FreeCAD 1.0/1.1).
5. Add all dress-up features (Fillet, Chamfer, etc.) AFTER all additive/subtractive features.
6. After the code block, write one short paragraph explaining the key design decisions.
7. Cite the retrieved sources inline as [1], [2], etc. in comments and in the explanation.
8. End with a numbered citation list: `1. <Page Title> — <URL>`

KNOWN PITFALLS (never repeat these errors):
- Missing doc.recompute() → silent failure
- Mixing App.ActiveDocument and the captured doc variable
- Creating PartDesign features via doc.addObject instead of body.newObject
- Importing *Gui modules in headless scripts

{_FEW_SHOT_BOX}

{_FEW_SHOT_REVOLVE}
""".strip()


# ── main entry point ──────────────────────────────────────────────────────────

def generate_response(
    query: str,
    citations: list[Citation],
    model: str = DEFAULT_MODEL,
) -> tuple[str, str, str]:
    """
    Returns (code_block, explanation_md, error_msg).
    code_block: the raw python code (no fences).
    explanation_md: explanation + citations markdown.
    error_msg: non-empty string on failure.
    """
    if not citations:
        return "", "", "No relevant documentation chunks were retrieved. Try broadening the query."

    client = InferenceClient()
    context = build_context_block(citations)

    user_msg = f"RETRIEVED CONTEXT:\n{context}\n\nUSER REQUEST:\n{query}"

    try:
        resp = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": _SYSTEM_PROMPT},
                {"role": "user",   "content": user_msg},
            ],
            temperature=0.2,
            max_tokens=2500,
        )
    except Exception as exc:  # noqa: BLE001
        return "", "", f"HuggingFace API error: {exc}"

    full_text = resp.choices[0].message.content or ""

    # Split code block from rest of response
    code_match = __import__("re").search(r"```python\n(.*?)```", full_text, __import__("re").DOTALL)
    if code_match:
        code = code_match.group(1).rstrip()
        after_code = full_text[code_match.end():].strip()
    else:
        code = full_text
        after_code = ""

    used_ids = extract_inline_refs(full_text)
    cite_md  = render_citation_markdown(citations, used_ids or None)
    explain  = (after_code + "\n\n" + cite_md).strip()

    return code, explain, ""