ir_builder.py

from typing import List, Dict, Any
import json

from utils import (
    call_llm,
    safe_json_loads,
    normalize_ir,
    validate_ir,
    build_fallback_ir
)


# =====================================================
# IR VALIDATION
# =====================================================

def validate_ir_list(data):

    if not isinstance(data, list):
        return False

    for item in data:

        if not isinstance(item, dict):
            return False

        if "nodes" not in item:
            return False

        if "edges" not in item:
            return False

    return True


# =====================================================
# IR BUILDER
# =====================================================

def build_ir_variants(
    equations: List[Dict[str, str]],
    num_variants: int = 3,
    provider: str = "openai"
) -> List[Dict[str, Any]]:

    if not equations:

        return []

    # -------------------------------------------------
    # FORMAT EQUATIONS
    # -------------------------------------------------

    eq_text = "\n".join([

        f"{i+1}. {eq.get('latex', '')}"

        for i, eq in enumerate(equations)

    ])

    # -------------------------------------------------
    # PROMPT
    # -------------------------------------------------

    prompt = f"""
Tu es un expert en :
- compilation mathématique
- graphes de calcul
- IR scientifiques
- optimisation tensorielle

Construis EXACTEMENT {num_variants}
variantes IR différentes.

Chaque IR doit contenir :
- name
- strategy
- nodes
- edges

Chaque node peut représenter :
- add
- multiply
- divide
- tensor
- matrix
- convolution
- reduction
- activation
- equation

IMPORTANT :
- retourne UNIQUEMENT du JSON valide
- aucune balise markdown
- aucun texte hors JSON

FORMAT STRICT :

[
  {{
    "name": "Naive Graph",
    "strategy": "naive",
    "nodes": [
      {{
        "id": "mul_1",
        "type": "multiply",
        "inputs": ["m", "c2"],
        "output": "E"
      }}
    ],
    "edges": [
      {{
        "from": "m",
        "to": "mul_1"
      }}
    ]
  }}
]

EQUATIONS :
{eq_text}
"""

    # -------------------------------------------------
    # LLM CALL
    # -------------------------------------------------

    response = call_llm(
        prompt,
        provider=provider,
        max_tokens=4000
    )

    parsed = safe_json_loads(response)

    # -------------------------------------------------
    # VALIDATION
    # -------------------------------------------------

    valid_variants = []

    if validate_ir_list(parsed):

        for item in parsed:

            normalized = normalize_ir(item)

            if validate_ir(normalized):

                valid_variants.append(normalized)

    # -------------------------------------------------
    # SUCCESS
    # -------------------------------------------------

    if len(valid_variants) > 0:

        return valid_variants[:num_variants]

    # -------------------------------------------------
    # FALLBACK
    # -------------------------------------------------

    return [
        build_fallback_ir(equations)
        for _ in range(num_variants)
    ]