Skip to content

English · Español

Lab 00 — Un planificador bajo JSONSchemaMask

Lee theory/01-react-and-planning.md. No consultes solutions/.

Objetivo

Construye un Planner que, dado un estado (frase + traza), genere el siguiente paso como JSON válido conforme al esquema del planificador — usando el JSONSchemaMask de la Fase 30 para restringir el decoding. El planificador debería ser incapaz de emitir salida inválida por construcción.

Setup

Un archivo nuevo: src/miniagent/planner.py. Importa:

  • MiniGPT (o una variante con fine-tuning) de src/minimodel/mini_gpt.py.
  • JSONSchemaMask de src/ministruct/mask.py (Fase 30).
  • ToolCall, FinalAnswer, Step de src/miniagent/types.py.

Tareas

Tarea 1 — define el esquema del planificador

En src/miniagent/schemas.py, define el JSON Schema para la salida del planificador:

{
  "oneOf": [
    {
      "type": "object",
      "properties": {
        "next": {"const": "tool_call"},
        "tool": {"enum": ["conjugate", "lookup_irregular_verb", "check_subject_verb_agreement", "lookup_spanish"]},
        "args": {"type": "object"}
      },
      "required": ["next", "tool", "args"],
      "additionalProperties": false
    },
    {
      "type": "object",
      "properties": {
        "next": {"const": "final_answer"},
        "answer": {
          "type": "object",
          "properties": {
            "corrected": {"type": ["string", "null"]},
            "rationale": {"type": "array", "items": {"type": "string"}},
            "spanish_gloss": {"type": ["string", "null"]},
            "in_scope": {"type": "boolean"}
          },
          "required": ["corrected", "rationale", "in_scope"],
          "additionalProperties": false
        }
      },
      "required": ["next", "answer"],
      "additionalProperties": false
    }
  ]
}

Valida este esquema en papel antes de codificar — escribe unas cuantas salidas de ejemplo y confirma que parsean (un tool_call con tool="conjugate" y args; un final_answer con payload completo).

Tarea 2 — implementa la clase Planner

class Planner:
    def __init__(self, model: MiniGPT, mask: JSONSchemaMask, tokenizer):
        self.model = model
        self.mask = mask
        self.tokenizer = tokenizer

    def next_step(self, state: PlannerState) -> Step:
        """Generate the next step. Output is guaranteed to be a valid Step by construction."""
        prompt = self._format_prompt(state)
        tokens = self.tokenizer.encode(prompt)
        # Run masked generation until the JSON object closes
        json_str = self._masked_generate(tokens)
        parsed = json.loads(json_str)
        return self._parse_step(parsed)

    def _masked_generate(self, prompt_tokens) -> str:
        """Generate tokens one at a time, applying the mask at each step."""
        # Loop:
        #   logits = self.model(tokens + generated)
        #   masked_logits = self.mask.apply(logits[-1], partial_json=...)
        #   next_tok = sample_or_argmax(masked_logits)
        #   if json_complete(generated): break
        ...

Restricciones:

  • La máscara debe aplicarse en cada token, no solo en puntos estructurales. El esquema restringe cada carácter de la salida.
  • El bucle de decoding debe soportar tanto greedy (argmax) como decoding con sampling por temperatura. Default: greedy.
  • Un presupuesto máximo de generación (p. ej., 256 tokens) evita bucles infinitos si la lógica de la máscara tiene un bug.

Tarea 3 — valida la imposición del esquema

Añade a tests/test_planner.py:

  1. Test de cumplimiento del esquema. Genera 100 salidas del planificador (con prompts variados). Valida cada una contra el esquema con jsonschema.validate. Se esperan cero fallos.
  2. Test del enum de herramientas. Confirma que el valor tool generado está siempre en el conjunto permitido.
  3. Test de sin-basura-al-final. Confirma que la salida generada termina exactamente tras el } de cierre — sin tokens al final.
  4. Comparación con máscara desactivada. Ejecuta los mismos prompts sin la máscara. La mayoría de las salidas deberían ser JSON inválido o fuera de esquema. Esto demuestra que la máscara hace trabajo real.

Tarea 4 — maneja con elegancia el estado no entrenado del modelo

El MiniGPT de la Fase 17 está no entrenado. Sus salidas son aleatorias también bajo decoding enmascarado — pero con la máscara, al menos parsean. El siguiente paso (campo tool, args) será aleatorio — eso está bien. La Fase 32 espera que un modelo entrenado o con fine-tuning de la Fase 28 se enchufe aquí.

En el lab, escribe un MockPlanner que produzca pasos correctos para un conjunto de frases canónicas de prueba. Esto te permite ejercitar el bucle del agente en el Lab 01 incluso antes de cablear un modelo entrenado:

class MockPlanner:
    """For testing only. Returns scripted steps for known sentences."""
    def __init__(self, scripts: dict[str, list[Step]]): ...
    def next_step(self, state: PlannerState) -> Step:
        return self.scripts[state.original][state.step_index]

El lab 01 de la Fase 32 usará MockPlanner para el conjunto canónico de prueba; en producción (Fase 33), el Planner real se enchufa.

Tarea 5 — mide: tokens para emitir un paso

Para una frase de prueba fija y una traza, mide:

  1. El forward del modelo (generación no restringida): ¿cuántos tokens hasta algo que parezca un objeto JSON?
  2. El planificador enmascarado: ¿cuántos tokens hasta un objeto válido completo?
  3. El overhead por token de la máscara (timing).

Esperado: el decoding enmascarado produce un objeto completo en menos tokens (sin preámbulo de "pensar en voz alta") y el overhead por token es < 10× el forward desnudo (la máscara hace un cómputo de seguimiento de esquema por token).

Guarda en experiments/<date>-phase-32-planner/timing.csv.

Medidas a capturar

  • Cumplimiento del esquema: 100/100 salidas generadas validan.
  • Máscara desactivada: conteo de salidas inválidas (debería ser alto).
  • Tokens por paso: enmascarado vs sin máscara.
  • Overhead por token de la máscara.

Aceptación

  • src/miniagent/planner.py y src/miniagent/schemas.py existen.
  • Todas las salidas generadas pasan jsonschema.validate.
  • MockPlanner disponible para uso por el Lab 01.
  • El test tests/test_planner.py está en verde.
  • La comparación máscara-vs-sin-máscara está documentada.

Trampas a esperar

  • El decoding no sabe cuándo parar. La máscara define los siguientes tokens válidos, pero también necesitas detectar cuándo el objeto está completo. Estrategia: cuando la profundidad de corchetes vuelve a cero tras el { de apertura, deja de generar.
  • additionalProperties: false. Fácil de omitir; sin él, el modelo puede emitir campos extra. Testea esto explícitamente.
  • oneOf sobre dos esquemas. La máscara debe seguir qué rama se está comprometiendo en cuanto el modelo se compromete (p. ej., en el momento en que "next": "tool_call" se emite, el esquema colapsa a la rama tool_call). Si tu máscara no maneja oneOf adecuadamente, las salidas pueden estar a caballo entre esquemas.
  • Alineamiento del tokenizer. La lógica de máscara JSON suele operar a nivel de carácter, pero la generación es a nivel de token. Necesitas un JSONSchemaMask que sepa qué tokens corresponden a qué caracteres (o decodifica/re-codifica por paso). Este es el detalle de implementación del decoding con máscara; la Fase 30 debería haberte dado un patrón funcional.

Siguiente: 01-tutor-end-to-end.md