validate.py

from __future__ import annotations

import json
import os
import signal
import socket
import subprocess
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path

import httpx

PROJECT_ROOT = Path(__file__).resolve().parent
TASKS = [
    "refund_triage_easy",
    "cross_function_brief_medium",
    "executive_escalation_hard",
]


def print_check(name: str, passed: bool, detail: str = "") -> None:
    status = "PASS" if passed else "FAIL"
    suffix = f" - {detail}" if detail else ""
    print(f"{status}: {name}{suffix}")


def find_free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind(("127.0.0.1", 0))
        sock.listen(1)
        return int(sock.getsockname()[1])


def wait_for_server(base_url: str, timeout: float = 20.0) -> bool:
    deadline = time.time() + timeout
    with httpx.Client(timeout=2.0) as client:
        while time.time() < deadline:
            try:
                response = client.get(f"{base_url}/health")
                if response.status_code == 200:
                    return True
            except Exception:
                time.sleep(0.5)
    return False


def greedy_action(observation: dict) -> dict:
    query_terms = set(observation["query"].lower().split())
    selected = set(observation.get("selected_chunks", []))
    available = [chunk for chunk in observation["available_chunks"] if chunk["chunk_id"] not in selected]
    remaining_budget = observation["token_budget"] - observation["total_tokens_used"]

    def overlap(chunk: dict) -> tuple[float, int, str]:
        keyword_terms = set(" ".join(chunk["keywords"]).lower().split())
        union = query_terms | keyword_terms
        score = (len(query_terms & keyword_terms) / len(union)) if union else 0.0
        return (-score, chunk["tokens"], chunk["chunk_id"])

    if selected and (
        observation["step_number"] >= 3
        or observation["total_tokens_used"] >= int(observation["token_budget"] * 0.7)
    ):
        if not observation.get("plan_draft"):
            return {"action_type": "set_resolution_plan", "plan": "Verify evidence, protect customers, and publish only grounded actions."}
        return {"action_type": "submit_report", "answer": "A concise grounded incident operations brief using the prioritized artifacts."}

    if selected:
        heavy = sorted(
            [chunk for chunk in available + observation["available_chunks"] if chunk["chunk_id"] in selected],
            key=lambda chunk: (-chunk["tokens"], chunk["chunk_id"]),
        )
        if heavy and heavy[0]["tokens"] > max(120, observation["token_budget"] // 3):
            return {
                "action_type": "summarize_artifact",
                "artifact_id": heavy[0]["chunk_id"],
                "compression_ratio": 0.5,
            }

    for chunk in sorted(available, key=overlap):
        return {"action_type": "inspect_artifact", "artifact_id": chunk["chunk_id"]}

    return {"action_type": "submit_report", "answer": "A concise grounded incident operations brief using the prioritized artifacts."}


def planner_action(client: httpx.Client, base_url: str, fallback_observation: dict) -> dict:
    try:
        response = client.post(f"{base_url}/optimize-step")
        if response.status_code == 200:
            return response.json()
    except Exception:
        pass
    return greedy_action(fallback_observation)


def run_task(client: httpx.Client, base_url: str, task_name: str) -> tuple[bool, float]:
    reset = client.post(f"{base_url}/reset", json={"task_name": task_name})
    if reset.status_code != 200:
        print_check(f"reset {task_name}", False, reset.text)
        return False, 0.0

    observation = reset.json()["observation"]
    done = False
    final_score = 0.0
    while not done:
        action = planner_action(client, base_url, observation)
        step = client.post(f"{base_url}/step", json=action)
        if step.status_code != 200:
            print_check(f"step {task_name}", False, step.text)
            return False, 0.0
        body = step.json()
        observation = body["observation"]
        done = body["done"]
        final_score = float(body["reward"])
    in_range = 0.0 <= final_score <= 1.0
    print_check(f"task {task_name} score range", in_range, f"score={final_score:.4f}")
    return in_range, final_score


def run_inference_script(base_url: str) -> bool:
    proxy_port = find_free_port()
    requests_seen: list[dict[str, str | None]] = []

    class ProxyHandler(BaseHTTPRequestHandler):
        def do_POST(self):
            length = int(self.headers.get("Content-Length", "0"))
            body = self.rfile.read(length).decode("utf-8")
            requests_seen.append(
                {
                    "path": self.path,
                    "authorization": self.headers.get("Authorization"),
                    "body": body,
                }
            )
            payload = {
                "id": "chatcmpl-validate",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": "validator-proxy",
                "choices": [
                    {
                        "index": 0,
                        "message": {
                            "role": "assistant",
                            "content": json.dumps(
                                {
                                    "action_type": "submit_report",
                                    "answer": "Validated via proxy [support_003]",
                                }
                            ),
                        },
                        "finish_reason": "stop",
                    }
                ],
            }
            encoded = json.dumps(payload).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(encoded)))
            self.end_headers()
            self.wfile.write(encoded)

        def log_message(self, format: str, *args):
            return

    proxy_server = HTTPServer(("127.0.0.1", proxy_port), ProxyHandler)
    proxy_thread = threading.Thread(target=proxy_server.serve_forever, daemon=True)
    proxy_thread.start()

    try:
        env = os.environ.copy()
        env["RAG_ENV_URL"] = base_url
        env.pop("ALLOW_BASELINE_FALLBACK", None)
        env["API_BASE_URL"] = f"http://127.0.0.1:{proxy_port}/v1"
        env["API_KEY"] = "offline-validation-token"
        env["HF_TOKEN"] = "legacy-should-not-win"
        process = subprocess.run(
            [sys.executable, "inference.py"],
            cwd=PROJECT_ROOT,
            capture_output=True,
            text=True,
            timeout=120,
            env=env,
        )
        stdout = process.stdout or ""
        has_start = "[START]" in stdout
        has_end = "[END]" in stdout
        end_has_score = " score=" in stdout
        proxy_called = any(request["path"] == "/v1/chat/completions" for request in requests_seen)
        auth_ok = any(request["authorization"] == "Bearer offline-validation-token" for request in requests_seen)
        return process.returncode == 0 and has_start and has_end and end_has_score and proxy_called and auth_ok
    finally:
        proxy_server.shutdown()
        proxy_server.server_close()


def main() -> int:
    port = find_free_port()
    base_url = f"http://127.0.0.1:{port}"
    command = [sys.executable, "-m", "uvicorn", "app:app", "--host", "127.0.0.1", "--port", str(port)]
    process = subprocess.Popen(command, cwd=PROJECT_ROOT)

    try:
        if not wait_for_server(base_url):
            print_check("server startup", False, "Timed out waiting for /health")
            return 1
        print_check("server startup", True)

        all_passed = True
        with httpx.Client(timeout=10.0) as client:
            health = client.get(f"{base_url}/health")
            health_ok = health.status_code == 200 and health.json().get("status") == "ok"
            print_check("GET /health", health_ok)
            all_passed &= health_ok

            reset = client.post(f"{base_url}/reset", json={"task_name": "refund_triage_easy"})
            reset_ok = reset.status_code == 200 and "observation" in reset.json()
            print_check("POST /reset", reset_ok)
            all_passed &= reset_ok

            initial_observation = reset.json().get("observation", {})
            first_chunk_id = None
            for chunk in initial_observation.get("available_chunks", []):
                if chunk.get("chunk_id"):
                    first_chunk_id = chunk["chunk_id"]
                    break
            step_payload = {"action_type": "inspect_artifact", "artifact_id": first_chunk_id} if first_chunk_id else {
                "action_type": "submit_report",
                "answer": "No chunk available for validation.",
            }
            step = client.post(f"{base_url}/step", json=step_payload)
            step_ok = step.status_code == 200 and isinstance(step.json().get("reward"), float)
            print_check("POST /step", step_ok)
            all_passed &= step_ok

            state = client.get(f"{base_url}/state")
            state_ok = state.status_code == 200 and "selected_chunks" in state.json()
            print_check("GET /state", state_ok)
            all_passed &= state_ok

            optimize_prompt = client.post(
                f"{base_url}/optimize-prompt",
                json={
                    "prompt": "Draft a customer-safe admin compromise update with rollback safeguards and cite evidence.",
                    "corpus_family": "enterprise_v2",
                    "compression_mode": "grounded",
                },
            )
            optimize_body = optimize_prompt.json() if optimize_prompt.status_code == 200 else {}
            optimize_ok = (
                optimize_prompt.status_code == 200
                and "optimized_prompt" in optimize_body
                and "context_tuning" in optimize_body
                and "grounding" in optimize_body
                and optimize_body.get("optimization_mode") == "grounded"
                and bool(optimize_body.get("grounding", {}).get("citation_ready"))
            )
            print_check("POST /optimize-prompt", optimize_ok)
            all_passed &= optimize_ok

            inference_ok = run_inference_script(base_url)
            print_check("python inference.py", inference_ok)
            all_passed &= inference_ok

            for task_name in TASKS:
                passed, _ = run_task(client, base_url, task_name)
                all_passed &= passed

        return 0 if all_passed else 1
    finally:
        if process.poll() is None:
            process.terminate()
            try:
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                if os.name == "nt":
                    process.kill()
                else:
                    process.send_signal(signal.SIGKILL)


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