IPSExtractor/docs/plan_pruebas/scripts/llm_mock_server.py

#!/usr/bin/env python3
"""Mock OpenAI-compatible para ST-EXT-01 (cola ARQ sin coste LLM real).

Endpoints:
  GET  /v1/models
  POST /v1/chat/completions

Responde JSON fijo con delay configurable (default 50 ms).

Ejemplo:
  python docs/plan_pruebas/scripts/llm_mock_server.py --port 8099 --delay-ms 50

Apuntar extraction/workers (Docker):
  EXTRACTOR_BASE_URL=http://host.docker.internal:8099/v1
  NORMALIZER_BASE_URL=http://host.docker.internal:8099/v1
  EXTRACTOR_API_KEY=mock-key
  ASYNC_QUEUE=true
"""
from __future__ import annotations

import argparse
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse


MOCK_SCHEMA = {
    "numero_factura": "str",
    "fecha": "str",
    "nit_prestador": "str",
    "razon_social": "str",
    "valor_total": "number",
}

MOCK_EXTRACTION = {
    "numero_factura": "MOCK-0001",
    "fecha": "2026-01-15",
    "nit_prestador": "900000000",
    "razon_social": "IPS MOCK ST-EXT-01",
    "valor_total": 1000,
}


class Stats:
    def __init__(self) -> None:
        self.lock = threading.Lock()
        self.completions = 0
        self.models = 0

    def inc(self, kind: str) -> None:
        with self.lock:
            if kind == "completions":
                self.completions += 1
            else:
                self.models += 1


STATS = Stats()


def make_handler(delay_ms: float):
    class Handler(BaseHTTPRequestHandler):
        def log_message(self, fmt: str, *args) -> None:  # quieter
            pass

        def _json(self, code: int, payload: dict) -> None:
            body = json.dumps(payload).encode("utf-8")
            self.send_response(code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self) -> None:  # noqa: N802
            path = urlparse(self.path).path.rstrip("/")
            if path in ("/v1/models", "/models"):
                STATS.inc("models")
                self._json(
                    200,
                    {
                        "object": "list",
                        "data": [
                            {"id": "mock-llm", "object": "model", "owned_by": "st-ext-01"},
                            {"id": "gpt-5.4-mini", "object": "model", "owned_by": "st-ext-01"},
                        ],
                    },
                )
                return
            if path in ("/health", "/v1/health"):
                self._json(200, {"status": "ok", "completions": STATS.completions})
                return
            self._json(404, {"error": {"message": f"not found: {path}", "type": "not_found"}})

        def do_POST(self) -> None:  # noqa: N802
            path = urlparse(self.path).path.rstrip("/")
            length = int(self.headers.get("Content-Length") or 0)
            raw = self.rfile.read(length) if length else b"{}"
            try:
                req = json.loads(raw.decode("utf-8") or "{}")
            except Exception:
                req = {}

            if path not in ("/v1/chat/completions", "/chat/completions"):
                self._json(404, {"error": {"message": f"not found: {path}", "type": "not_found"}})
                return

            if delay_ms > 0:
                time.sleep(delay_ms / 1000.0)

            # Heurística: schema discovery suele pedir "campos" / schema; extracción pide factura.
            messages = req.get("messages") or []
            blob = json.dumps(messages, ensure_ascii=False).lower()
            if "schema" in blob or "campos" in blob or "field" in blob or "descub" in blob:
                content = json.dumps(MOCK_SCHEMA, ensure_ascii=False)
            else:
                content = json.dumps(MOCK_EXTRACTION, ensure_ascii=False)

            model = req.get("model") or "mock-llm"
            STATS.inc("completions")
            self._json(
                200,
                {
                    "id": f"chatcmpl-mock-{STATS.completions}",
                    "object": "chat.completion",
                    "created": int(time.time()),
                    "model": model,
                    "choices": [
                        {
                            "index": 0,
                            "message": {"role": "assistant", "content": content},
                            "finish_reason": "stop",
                        }
                    ],
                    "usage": {
                        "prompt_tokens": 100,
                        "completion_tokens": 50,
                        "total_tokens": 150,
                    },
                },
            )

    return Handler


def main() -> int:
    p = argparse.ArgumentParser(description="Mock OpenAI API for ST-EXT-01")
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=int(__import__("os").environ.get("LLM_MOCK_PORT", "8099")))
    p.add_argument("--delay-ms", type=float, default=float(__import__("os").environ.get("LLM_MOCK_DELAY_MS", "50")))
    args = p.parse_args()

    httpd = ThreadingHTTPServer((args.host, args.port), make_handler(args.delay_ms))
    print(f"[llm-mock] http://{args.host}:{args.port}/v1  delay={args.delay_ms}ms")
    print(f"[llm-mock] health: http://127.0.0.1:{args.port}/health")
    print("[llm-mock] Ctrl+C para detener")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\n[llm-mock] stopped")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())