Upload folder using huggingface_hub

README.md

@@ -116,7 +116,9 @@ Interactive docs: `http://localhost:7860/docs`
 ### Prerequisites
 - Python 3.10+
 - Docker
-- `OPENAI_API_KEY` (for baseline only)
+- `API_BASE_URL` (OpenAI-compatible endpoint for inference)
+- `MODEL_NAME` (model identifier for inference)
+- `HF_TOKEN` (API key / bearer token for inference)
 
 ### Local (Python)
 
@@ -135,8 +137,10 @@ docker run -p 7860:7860 -e OPENAI_API_KEY=sk-... sql-optimizer-env
 ### Baseline Inference
 
 ```bash
-export OPENAI_API_KEY=sk-...
-python baseline.py
+$env:API_BASE_URL="https://api.openai.com/v1"
+$env:MODEL_NAME="gpt-4o-mini"
+$env:HF_TOKEN="hf_or_openai_api_key_here"
+python inference.py
 ```
 
 ### OpenEnv Validation
@@ -154,6 +158,16 @@ huggingface-cli login
 openenv push --repo-id your-username/sql-query-optimizer
 ```
 
+### Environment Configuration
+
+Define these variables before running inference or `/baseline`:
+
+```powershell
+$env:API_BASE_URL = "https://api.openai.com/v1"
+$env:MODEL_NAME = "gpt-4o-mini"
+$env:HF_TOKEN = "your_api_key"
+```
+
 ---
 
 ## Baseline Scores

baseline.py

@@ -1,144 +1,8 @@
-"""
-Baseline inference script for the SQL Query Optimizer OpenEnv environment.
-
-Usage:
-    python baseline.py              # human-readable output
-    python baseline.py --json       # JSON output (used by /baseline endpoint)
-
-Requires:
-    OPENAI_API_KEY environment variable
-
-The script runs gpt-4o-mini against all 3 tasks and reports grader scores.
-"""
-from __future__ import annotations
-
-import argparse
-import json
-import os
-import sys
-
-from openai import OpenAI
-
-# ── import env from local package ──────────────────────────────────────────
-sys.path.insert(0, os.path.dirname(__file__))
-from env.environment import SQLOptimizerEnv
-from env.models import Action
-
-# ──────────────────────────────────────────────────────────────────────────────
-MODEL = "gpt-4o-mini"
-MAX_STEPS = 5
-TASKS = [1, 2, 3]
-
-SYSTEM_PROMPT = """You are a database performance engineer.
-You will receive a broken or unoptimised SQL query along with table schema context.
-Your job is to rewrite the query so it is correct and performant.
-
-Respond ONLY with a JSON object with these exact keys:
-{
-  "rewritten_query": "<your improved SQL>",
-  "explanation": "<brief explanation of changes>",
-  "is_done": true
-}
-Do not wrap in markdown. Output raw JSON only."""
-
-
-def _build_user_message(obs_dict: dict) -> str:
-    return (
-        f"Task: {obs_dict['task_name']} ({obs_dict['task_id']} — difficulty: "
-        f"{obs_dict.get('difficulty', 'unknown')})\n\n"
-        f"Description:\n{obs_dict['task_description']}\n\n"
-        f"Schema:\n{obs_dict['schema_context']}\n\n"
-        f"Query to fix:\n{obs_dict['query']}"
-        + (f"\n\nHint: {obs_dict['hint']}" if obs_dict.get("hint") else "")
-    )
-
-
-def run_baseline(verbose: bool = True) -> dict[str, float]:
-    api_key = os.getenv("OPENAI_API_KEY")
-    if not api_key:
-        print("ERROR: OPENAI_API_KEY is not set.", file=sys.stderr)
-        sys.exit(1)
-
-    client = OpenAI(api_key=api_key)
-    env = SQLOptimizerEnv()
-    results: dict[str, float] = {}
-
-    for task_id in TASKS:
-        obs = env.reset(task_id=task_id)
-        obs_dict = obs.model_dump()
-        final_score = 0.0
-
-        if verbose:
-            print(f"\n{'='*60}")
-            print(f"Task {task_id}: {obs_dict['task_name']} [{obs_dict['task_id']}]")
-            print(f"{'='*60}")
-
-        for step_num in range(MAX_STEPS):
-            messages = [
-                {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": _build_user_message(obs_dict)},
-            ]
-
-            try:
-                response = client.chat.completions.create(
-                    model=MODEL,
-                    messages=messages,
-                    temperature=0.0,
-                    max_tokens=1024,
-                )
-                content = response.choices[0].message.content.strip()
-                parsed = json.loads(content)
-                action = Action(
-                    rewritten_query=parsed.get("rewritten_query", ""),
-                    explanation=parsed.get("explanation", ""),
-                    is_done=bool(parsed.get("is_done", False)),
-                )
-            except Exception as exc:
-                if verbose:
-                    print(f"  Step {step_num + 1}: LLM error — {exc}")
-                action = Action(
-                    rewritten_query="",
-                    explanation="error",
-                    is_done=True,
-                )
-
-            obs, reward, done, info = env.step(action)
-            obs_dict = obs.model_dump()
-            final_score = info["grader_score"]
-
-            if verbose:
-                print(
-                    f"  Step {step_num + 1}: grader_score={info['grader_score']:.3f}  "
-                    f"step_reward={reward.score:.4f}  feedback={reward.feedback[:80]}"
-                )
-
-            if done:
-                break
-
-        results[f"task_{task_id}_{env._task.name}"] = round(final_score, 4)
-
-        if verbose:
-            print(f"  → Final grader score: {final_score:.4f}")
-
-    if verbose:
-        print(f"\n{'='*60}")
-        print("BASELINE RESULTS")
-        print(f"{'='*60}")
-        for k, v in results.items():
-            print(f"  {k}: {v:.4f}")
-        avg = sum(results.values()) / len(results)
-        print(f"  Average: {avg:.4f}")
-
-    return results
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="OpenEnv SQL Optimizer — Baseline Inference")
-    parser.add_argument(
-        "--json", action="store_true", help="Output results as JSON (used by /baseline endpoint)"
-    )
-    args = parser.parse_args()
-
-    scores = run_baseline(verbose=not args.json)
-    if args.json:
-        print(json.dumps(scores))
+"""Compatibility wrapper for the required inference.py entrypoint."""
+from __future__ import annotations
+
+from inference import run_inference
+
+
+if __name__ == "__main__":
+    run_inference()

inference.py

@@ -0,0 +1,197 @@
+"""OpenAI-based inference runner for the SQL Query Optimizer OpenEnv environment.
+
+Environment variables:
+    API_BASE_URL: OpenAI-compatible API endpoint
+    MODEL_NAME: model identifier to use for inference
+    HF_TOKEN: API key / bearer token for the LLM provider
+
+The script emits structured stdout logs in three sections only:
+    [START] ...
+    [STEP] ...
+    [END] ...
+"""
+from __future__ import annotations
+
+import json
+import os
+import sys
+from collections import OrderedDict
+from typing import Any, Dict
+
+from openai import OpenAI
+
+sys.path.insert(0, os.path.dirname(__file__))
+
+from env.environment import SQLOptimizerEnv
+from env.models import Action
+
+DEFAULT_MAX_STEPS = 5
+TASK_IDS = (1, 2, 3)
+
+SYSTEM_PROMPT = """You are a database performance engineer.
+You will receive a broken or unoptimised SQL query along with table schema context.
+Your job is to rewrite the query so it is correct and performant.
+
+Respond ONLY with a JSON object with these exact keys:
+{
+  "rewritten_query": "<your improved SQL>",
+  "explanation": "<brief explanation of changes>",
+  "is_done": true
+}
+Do not wrap in markdown. Output raw JSON only."""
+
+
+def _load_runtime_config() -> Dict[str, str]:
+    api_base_url = os.getenv("API_BASE_URL", "").strip()
+    model_name = os.getenv("MODEL_NAME", "").strip()
+    hf_token = os.getenv("HF_TOKEN", "").strip()
+
+    missing = [
+        name
+        for name, value in (
+            ("API_BASE_URL", api_base_url),
+            ("MODEL_NAME", model_name),
+            ("HF_TOKEN", hf_token),
+        )
+        if not value
+    ]
+    if missing:
+        raise RuntimeError(f"Missing required environment variables: {', '.join(missing)}")
+
+    return {
+        "API_BASE_URL": api_base_url,
+        "MODEL_NAME": model_name,
+        "HF_TOKEN": hf_token,
+    }
+
+
+def _build_user_message(obs_dict: dict) -> str:
+    message = (
+        f"Task: {obs_dict['task_name']} ({obs_dict['task_id']} — difficulty: "
+        f"{obs_dict.get('difficulty', 'unknown')})\n\n"
+        f"Description:\n{obs_dict['task_description']}\n\n"
+        f"Schema:\n{obs_dict['schema_context']}\n\n"
+        f"Query to fix:\n{obs_dict['query']}"
+    )
+    if obs_dict.get("hint"):
+        message += f"\n\nHint: {obs_dict['hint']}"
+    return message
+
+
+def _log(prefix: str, payload: Dict[str, Any]) -> None:
+    print(f"{prefix} {json.dumps(payload, ensure_ascii=True, separators=(',', ':'))}")
+
+
+def _parse_json_action(text: str) -> Action:
+    parsed = json.loads(text)
+    return Action(
+        rewritten_query=parsed.get("rewritten_query", ""),
+        explanation=parsed.get("explanation", ""),
+        is_done=bool(parsed.get("is_done", False)),
+    )
+
+
+def run_inference() -> Dict[str, float]:
+    config = _load_runtime_config()
+    client = OpenAI(api_key=config["HF_TOKEN"], base_url=config["API_BASE_URL"])
+    env = SQLOptimizerEnv()
+
+    _log(
+        "[START]",
+        OrderedDict(
+            [
+                ("script", "inference.py"),
+                ("api_base_url", config["API_BASE_URL"]),
+                ("model_name", config["MODEL_NAME"]),
+                ("tasks", list(TASK_IDS)),
+            ]
+        ),
+    )
+
+    results: Dict[str, float] = {}
+    total_score = 0.0
+
+    for task_id in TASK_IDS:
+        observation = env.reset(task_id=task_id)
+        obs_dict = observation.model_dump()
+        final_grader_score = 0.0
+        step_count = 0
+
+        for step_number in range(DEFAULT_MAX_STEPS):
+            messages = [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": _build_user_message(obs_dict)},
+            ]
+
+            try:
+                response = client.chat.completions.create(
+                    model=config["MODEL_NAME"],
+                    messages=messages,
+                    temperature=0.0,
+                    max_tokens=1024,
+                )
+                content = (response.choices[0].message.content or "").strip()
+                action = _parse_json_action(content)
+                llm_status = "ok"
+            except Exception as exc:
+                action = Action(rewritten_query="", explanation=f"error: {exc}", is_done=True)
+                llm_status = "error"
+
+            observation, reward, done, info = env.step(action)
+            obs_dict = observation.model_dump()
+            final_grader_score = float(info.get("grader_score", 0.0))
+            step_count = step_number + 1
+
+            _log(
+                "[STEP]",
+                OrderedDict(
+                    [
+                        ("task_id", task_id),
+                        ("task_name", obs_dict["task_name"]),
+                        ("step", step_count),
+                        ("grader_score", round(final_grader_score, 4)),
+                        ("reward_score", round(float(reward.score), 4)),
+                        ("done", bool(done)),
+                        ("llm_status", llm_status),
+                    ]
+                ),
+            )
+
+            if done:
+                break
+
+        task_key = f"task_{task_id}_{env._task.name}"
+        results[task_key] = round(final_grader_score, 4)
+        total_score += final_grader_score
+
+    average_score = round(total_score / len(TASK_IDS), 4)
+
+    _log(
+        "[END]",
+        OrderedDict(
+            [
+                ("task_results", results),
+                ("average_score", average_score),
+                ("status", "success"),
+            ]
+        ),
+    )
+    return results
+
+
+if __name__ == "__main__":
+    try:
+        run_inference()
+    except Exception as exc:
+        _log(
+            "[END]",
+            OrderedDict(
+                [
+                    ("task_results", {}),
+                    ("average_score", 0.0),
+                    ("status", "error"),
+                    ("error", str(exc)),
+                ]
+            ),
+        )
+        sys.exit(1)

pyproject.toml

@@ -30,9 +30,13 @@ dependencies = [
     "uvicorn[standard]>=0.29.0",
     "pydantic>=2.7.0",
     "openai>=1.30.0",
+    "openenv-core>=0.2.0",
     "pyyaml>=6.0",
 ]
 
+[project.scripts]
+server = "server.app:main"
+
 [project.optional-dependencies]
 dev = [
     "pytest>=7.0",

server/app.py

@@ -19,6 +19,7 @@ from typing import Any, Dict, Optional
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
+import uvicorn
 
 from env.environment import SQLOptimizerEnv
 from env.models import Action, Observation, Reward
@@ -79,6 +80,17 @@ class BaselineResponse(BaseModel):
     message: str
 
 
+def _parse_end_payload(stdout: str) -> Dict[str, Any]:
+    for line in reversed(stdout.splitlines()):
+        if not line.startswith("[END] "):
+            continue
+        payload_text = line[len("[END] ") :].strip()
+        import json
+
+        return json.loads(payload_text)
+    raise ValueError("Could not find [END] payload in inference output")
+
+
 # ──────────────────────────────────────────────────────────────────────────────
 # Endpoints
 # ──────────────────────────────────────────────────────────────────────────────
@@ -157,30 +169,42 @@ def grader() -> GraderResponse:
 @app.post("/baseline", response_model=BaselineResponse, summary="Run baseline inference on all tasks")
 def baseline() -> BaselineResponse:
     """
-    Trigger the baseline inference script (baseline.py) and return scores.
-    Requires OPENAI_API_KEY to be set in the environment.
+    Trigger the baseline inference script (inference.py) and return scores.
+    Requires API_BASE_URL, MODEL_NAME, and HF_TOKEN to be set in the environment.
     """
-    if not os.getenv("OPENAI_API_KEY"):
+    required_vars = ["API_BASE_URL", "MODEL_NAME", "HF_TOKEN"]
+    missing = [name for name in required_vars if not os.getenv(name)]
+    if missing:
         raise HTTPException(
             status_code=400,
-            detail="OPENAI_API_KEY environment variable not set. Cannot run baseline.",
+            detail=f"Missing required environment variables: {', '.join(missing)}",
         )
     try:
         result = subprocess.run(
-            [sys.executable, "baseline.py", "--json"],
+            [sys.executable, "inference.py"],
             capture_output=True,
             text=True,
-            timeout=300,
+            timeout=1200,
         )
         if result.returncode != 0:
             raise HTTPException(
                 status_code=500,
-                detail=f"Baseline script failed:\n{result.stderr}",
+                detail=f"Inference script failed:\n{result.stderr}",
             )
-        import json
-        scores = json.loads(result.stdout)
-        return BaselineResponse(task_results=scores, message="Baseline completed successfully.")
+        payload = _parse_end_payload(result.stdout)
+        return BaselineResponse(
+            task_results=payload.get("task_results", {}),
+            message="Baseline completed successfully.",
+        )
     except subprocess.TimeoutExpired:
-        raise HTTPException(status_code=500, detail="Baseline script timed out after 300s.")
+        raise HTTPException(status_code=500, detail="Inference script timed out after 1200s.")
     except Exception as exc:
         raise HTTPException(status_code=500, detail=str(exc))
+
+
+def main() -> None:
+    uvicorn.run("server.app:app", host="0.0.0.0", port=7860)
+
+
+if __name__ == "__main__":
+    main()