Spaces:

NITISHRG15102007
/

rag-context-optimizer

Running

App Files Files Community

NITISHRG15102007 commited on Apr 7

Commit

f8d0eb2

verified ·

1 Parent(s): de095eb

Fix custom prompt UI and token optimizer flow

Browse files

Files changed (6) hide show

__pycache__/app.cpython-314.pyc +0 -0
__pycache__/streamlit_app.cpython-314.pyc +0 -0
app.py +618 -220
env/__pycache__/environment.cpython-314.pyc +0 -0
env/environment.py +30 -3
streamlit_app.py +123 -33

__pycache__/app.cpython-314.pyc CHANGED Viewed

Binary files a/__pycache__/app.cpython-314.pyc and b/__pycache__/app.cpython-314.pyc differ

__pycache__/streamlit_app.cpython-314.pyc CHANGED Viewed

Binary files a/__pycache__/streamlit_app.cpython-314.pyc and b/__pycache__/streamlit_app.cpython-314.pyc differ

app.py CHANGED Viewed

@@ -20,6 +20,9 @@ from env.tasks import ALL_TASKS, TASKS_BY_NAME
 class ResetRequest(BaseModel):
     task_name: Literal["single_domain_qa", "cross_domain_synthesis", "adversarial_compression"]
 @asynccontextmanager
@@ -49,251 +52,478 @@ app.add_middleware(
 UI_HTML = """
 <!doctype html>
 <html lang="en">
 <head>
   <meta charset="utf-8" />
   <meta name="viewport" content="width=device-width, initial-scale=1" />
   <title>rag-context-optimizer</title>
   <style>
     :root {
-      --bg: #f4f7fb;
       --panel: #ffffff;
-      --ink: #122033;
-      --muted: #64748b;
-      --line: #dbe4f0;
       --accent: #0f766e;
-      --accent-2: #0ea5e9;
       --warn: #b45309;
     }
     * { box-sizing: border-box; }
     body {
       margin: 0;
       font-family: ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
       color: var(--ink);
-      background:
-        radial-gradient(circle at top left, #d9f99d 0, transparent 22%),
-        radial-gradient(circle at top right, #bfdbfe 0, transparent 25%),
-        linear-gradient(180deg, #f8fbff 0%, var(--bg) 100%);
-    }
-    .shell {
-      max-width: 1240px;
-      margin: 0 auto;
-      padding: 28px 20px 48px;
-    }
-    .hero {
-      display: grid;
-      grid-template-columns: 1.2fr 0.8fr;
-      gap: 18px;
-      margin-bottom: 18px;
     }
     .card {
       background: var(--panel);
       border: 1px solid var(--line);
-      border-radius: 20px;
       padding: 18px;
-      box-shadow: 0 10px 30px rgba(15, 23, 42, 0.05);
     }
     h1, h2, h3, p { margin-top: 0; }
-    h1 { font-size: 2.1rem; margin-bottom: 10px; }
-    .subtitle { color: var(--muted); line-height: 1.5; }
-    .stat-grid {
-      display: grid;
-      grid-template-columns: repeat(3, 1fr);
-      gap: 12px;
-      margin-top: 18px;
-    }
-    .stat {
-      border: 1px solid var(--line);
-      border-radius: 16px;
-      padding: 12px;
-      background: #fcfdff;
-    }
-    .stat .label { font-size: 0.82rem; color: var(--muted); }
-    .stat .value { font-size: 1.4rem; font-weight: 700; margin-top: 6px; }
-    .layout {
-      display: grid;
-      grid-template-columns: 320px 1fr;
-      gap: 18px;
-    }
-    .task-list { display: grid; gap: 10px; }
     .task-btn {
       width: 100%;
       text-align: left;
       border: 1px solid var(--line);
       background: #fff;
-      border-radius: 14px;
-      padding: 12px 14px;
       cursor: pointer;
-      transition: 0.18s ease;
     }
     .task-btn:hover, .task-btn.active {
-      border-color: var(--accent-2);
-      transform: translateY(-1px);
-      box-shadow: 0 8px 18px rgba(14, 165, 233, 0.12);
-    }
-    .task-btn .name { font-weight: 700; }
-    .task-btn .meta { color: var(--muted); font-size: 0.85rem; margin-top: 4px; }
-    .toolbar {
-      display: flex;
-      flex-wrap: wrap;
-      gap: 10px;
-      margin-bottom: 14px;
     }
-    button.action, select, textarea {
-      border-radius: 12px;
-      border: 1px solid var(--line);
-      font: inherit;
-    }
-    button.action {
-      background: linear-gradient(135deg, var(--accent), var(--accent-2));
-      color: white;
-      padding: 10px 14px;
       cursor: pointer;
-      border: none;
       font-weight: 700;
     }
-    button.action.secondary {
-      background: white;
-      color: var(--ink);
-      border: 1px solid var(--line);
-    }
-    .query {
-      font-size: 1.05rem;
-      font-weight: 600;
       padding: 14px;
-      border-radius: 16px;
-      background: #f8fafc;
-      border: 1px solid var(--line);
-      margin-bottom: 14px;
-    }
-    .grid {
-      display: grid;
-      grid-template-columns: repeat(auto-fill, minmax(240px, 1fr));
-      gap: 12px;
-    }
-    .chunk {
       border: 1px solid var(--line);
       border-radius: 16px;
-      padding: 12px;
-      background: #fff;
-    }
-    .chunk.selected {
-      border-color: var(--accent);
-      background: #f0fdfa;
     }
-    .chunk h4 { margin: 0 0 8px; font-size: 1rem; }
     .pill {
       display: inline-block;
       border-radius: 999px;
-      padding: 3px 8px;
-      background: #eef6ff;
       color: #1d4ed8;
       font-size: 0.75rem;
-      margin-right: 6px;
-      margin-bottom: 6px;
-    }
-    .row {
-      display: flex;
-      gap: 8px;
-      align-items: center;
-      flex-wrap: wrap;
-      margin-top: 10px;
-    }
-    .answer-box {
-      width: 100%;
-      min-height: 110px;
-      padding: 12px;
-      resize: vertical;
-    }
-    .panel-grid {
-      display: grid;
-      grid-template-columns: 1fr 1fr;
-      gap: 12px;
-      margin-top: 14px;
     }
     pre {
       margin: 0;
       white-space: pre-wrap;
       word-break: break-word;
-      font-size: 0.88rem;
       background: #0f172a;
       color: #e2e8f0;
       padding: 14px;
-      border-radius: 16px;
-      min-height: 160px;
-    }
-    .feedback { color: var(--warn); font-weight: 600; min-height: 24px; }
-    @media (max-width: 980px) {
-      .hero, .layout, .panel-grid { grid-template-columns: 1fr; }
-      .stat-grid { grid-template-columns: 1fr; }
     }
   </style>
 </head>
 <body>
-  <div class="shell">
-    <section class="hero">
-      <div class="card">
-        <h1>RAG Context Optimizer</h1>
-        <p class="subtitle">
-          Interactively test retrieval, compression, and answer quality tradeoffs.
-          This UI sits on top of the benchmark API, so you can explore the environment visually
-          without losing the validator-friendly `/reset`, `/step`, and `/state` endpoints.
-        </p>
-        <div class="stat-grid">
-          <div class="stat"><div class="label">Token Budget</div><div class="value" id="budgetStat">-</div></div>
-          <div class="stat"><div class="label">Tokens Used</div><div class="value" id="usedStat">-</div></div>
-          <div class="stat"><div class="label">Step</div><div class="value" id="stepStat">-</div></div>
-        </div>
       </div>
-      <div class="card">
-        <h3>Quick Start</h3>
-        <p class="subtitle">
-          Pick a task, press reset, select evidence chunks, optionally compress them, and submit an answer.
-          The right-hand panels mirror the raw observation and server state so debugging stays easy.
-        </p>
-        <div class="row">
-          <a href="/docs" target="_blank"><button class="action secondary" type="button">Open API Docs</button></a>
-          <a href="/health" target="_blank"><button class="action secondary" type="button">Health JSON</button></a>
         </div>
-      </div>
-    </section>
-    <section class="layout">
-      <aside class="card">
-        <h3>Tasks</h3>
-        <div id="taskList" class="task-list"></div>
       </aside>
-      <main class="card">
-        <div class="toolbar">
-          <button id="resetBtn" class="action" type="button">Reset Task</button>
-          <button id="refreshBtn" class="action secondary" type="button">Refresh State</button>
         </div>
-        <div id="queryBox" class="query">Choose a task to begin.</div>
-        <div class="feedback" id="feedbackBox"></div>
-        <h3>Available Chunks</h3>
-        <div id="chunkGrid" class="grid"></div>
-        <h3 style="margin-top:18px;">Answer</h3>
-        <textarea id="answerInput" class="answer-box" placeholder="Write a concise answer here before submitting..."></textarea>
-        <div class="row">
-          <button id="submitBtn" class="action" type="button">Submit Answer</button>
         </div>
-        <div class="panel-grid">
-          <div>
             <h3>Observation</h3>
             <pre id="observationPanel">{}</pre>
           </div>
-          <div>
             <h3>State</h3>
             <pre id="statePanel">{}</pre>
           </div>
         </div>
-      </main>
-    </section>
   </div>
   <script>
     const taskList = document.getElementById("taskList");
     const chunkGrid = document.getElementById("chunkGrid");
@@ -304,121 +534,207 @@ UI_HTML = """
     const budgetStat = document.getElementById("budgetStat");
     const usedStat = document.getElementById("usedStat");
     const stepStat = document.getElementById("stepStat");
     const answerInput = document.getElementById("answerInput");
     let selectedTask = "single_domain_qa";
     let currentObservation = null;
     async function fetchJson(url, options = {}) {
       const response = await fetch(url, {
         headers: { "Content-Type": "application/json" },
-        ...options,
       });
       const body = await response.json();
-      if (!response.ok) {
-        throw new Error(body.detail || JSON.stringify(body));
-      }
       return body;
     }
     function renderTasks(tasks) {
       taskList.innerHTML = "";
-      tasks.forEach(task => {
         const btn = document.createElement("button");
         btn.className = "task-btn" + (task.name === selectedTask ? " active" : "");
-        btn.innerHTML = `<div class="name">${task.name}</div><div class="meta">${task.difficulty} · budget ${task.token_budget}</div><div class="meta">${task.description}</div>`;
-        btn.onclick = () => {
           selectedTask = task.name;
           renderTasks(tasks);
-        };
         taskList.appendChild(btn);
       });
     }
     function renderObservation(observation) {
       currentObservation = observation;
-      queryBox.textContent = observation.query;
       budgetStat.textContent = observation.token_budget;
       usedStat.textContent = observation.total_tokens_used;
       stepStat.textContent = observation.step_number;
-      feedbackBox.textContent = observation.last_action_feedback || "";
       observationPanel.textContent = JSON.stringify(observation, null, 2);
       const selected = new Set(observation.selected_chunks || []);
       chunkGrid.innerHTML = "";
-      observation.available_chunks.forEach(chunk => {
         const card = document.createElement("div");
-        card.className = "chunk" + (selected.has(chunk.chunk_id) ? " selected" : "");
-        const pills = chunk.keywords.map(keyword => `<span class="pill">${keyword}</span>`).join("");
-        const selectAction = selected.has(chunk.chunk_id)
-          ? `<button class="action secondary" data-action="deselect" data-id="${chunk.chunk_id}">Deselect</button>`
-          : `<button class="action secondary" data-action="select" data-id="${chunk.chunk_id}">Select</button>`;
         card.innerHTML = `
           <h4>${chunk.chunk_id}</h4>
-          <div class="row"><span class="pill">${chunk.domain}</span><span class="pill">${chunk.tokens} tokens</span></div>
-          <div style="margin-top:8px;">${pills}</div>
-          <div class="row">
-            ${selectAction}
-            <button class="action secondary" data-action="compress" data-id="${chunk.chunk_id}">Compress 50%</button>
           </div>
         `;
         chunkGrid.appendChild(card);
       });
-      chunkGrid.querySelectorAll("button[data-action]").forEach(btn => {
-        btn.addEventListener("click", async () => {
-          const action = btn.dataset.action;
-          const chunkId = btn.dataset.id;
-          if (action === "select") {
-            await step({ action_type: "select_chunk", chunk_id: chunkId });
-          } else if (action === "deselect") {
-            await step({ action_type: "deselect_chunk", chunk_id: chunkId });
-          } else if (action === "compress") {
-            await step({ action_type: "compress_chunk", chunk_id: chunkId, compression_ratio: 0.5 });
           }
         });
       });
     }
     async function refreshState() {
-      const state = await fetchJson("/state", { method: "GET" });
-      statePanel.textContent = JSON.stringify(state, null, 2);
     }
     async function resetTask() {
-      const body = await fetchJson("/reset", {
-        method: "POST",
-        body: JSON.stringify({ task_name: selectedTask }),
-      });
-      renderObservation(body.observation);
-      await refreshState();
     }
     async function step(payload) {
       try {
         const body = await fetchJson("/step", {
           method: "POST",
-          body: JSON.stringify(payload),
         });
         renderObservation(body.observation);
         if (body.info && body.info.grader_breakdown) {
-          feedbackBox.textContent = "Final score: " + body.reward.toFixed(4) + " | " + JSON.stringify(body.info.grader_breakdown);
         }
         await refreshState();
       } catch (error) {
         feedbackBox.textContent = error.message;
       }
     }
     document.getElementById("resetBtn").addEventListener("click", resetTask);
     document.getElementById("refreshBtn").addEventListener("click", refreshState);
     document.getElementById("submitBtn").addEventListener("click", async () => {
-      await step({ action_type: "submit_answer", answer: answerInput.value || "A concise answer synthesized from the selected evidence." });
     });
     (async function init() {
-      const tasks = await fetchJson("/tasks", { method: "GET" });
-      renderTasks(tasks);
-      await resetTask();
     })();
   </script>
 </body>
@@ -436,7 +752,7 @@ async def log_requests(request: Request, call_next):
 @app.get("/", response_class=HTMLResponse)
 async def home_page():
-    return HTMLResponse(UI_HTML)
 def _serialize_observation(observation: Any) -> dict[str, Any]:
@@ -475,11 +791,83 @@ def _is_bad_action_event(event: str | None) -> bool:
     }
 @app.post("/reset")
 async def reset_endpoint(payload: ResetRequest):
     if payload.task_name not in TASKS_BY_NAME:
         raise HTTPException(status_code=400, detail="Unknown task_name.")
-    env = RagContextOptimizerEnv(task_name=payload.task_name)
     app.state.env = env
     result = await env.reset()
     return _serialize_step_result(result, reset=True)
@@ -519,11 +907,21 @@ async def tasks_endpoint():
             "description": task.description,
             "difficulty": task.difficulty,
             "token_budget": task.token_budget,
         }
         for task in ALL_TASKS
     ]
 if __name__ == "__main__":
     import uvicorn

 class ResetRequest(BaseModel):
     task_name: Literal["single_domain_qa", "cross_domain_synthesis", "adversarial_compression"]
+    custom_query: str | None = None
+    token_budget: int | None = None
+    max_steps: int | None = None
 @asynccontextmanager
 UI_HTML = """
 <!doctype html>
 <html lang="en">
+<head>
+  <meta charset="utf-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1" />
+  <title>rag-context-optimizer</title>
+  <style>
+    :root { --bg:#f5f7fb; --panel:#fff; --ink:#132238; --muted:#64748b; --line:#dbe4f0; --accent:#0f766e; --accent2:#0284c7; }
+    * { box-sizing:border-box; }
+    body { margin:0; font-family:ui-sans-serif,system-ui,Segoe UI,sans-serif; color:var(--ink); background:linear-gradient(180deg,#f8fbff,var(--bg)); }
+    .wrap { max-width:1280px; margin:0 auto; padding:24px; }
+    .hero, .main, .panels { display:grid; gap:16px; }
+    .hero { grid-template-columns:1.2fr 0.8fr; margin-bottom:16px; }
+    .main { grid-template-columns:320px 1fr; }
+    .panels { grid-template-columns:1fr 1fr; margin-top:16px; }
+    .card { background:var(--panel); border:1px solid var(--line); border-radius:20px; padding:18px; box-shadow:0 10px 24px rgba(15,23,42,.05); }
+    h1,h2,h3,p { margin-top:0; }
+    .muted { color:var(--muted); line-height:1.5; }
+    .stats { display:grid; grid-template-columns:repeat(3,1fr); gap:12px; margin-top:14px; }
+    .stat { border:1px solid var(--line); border-radius:16px; padding:12px; }
+    .stat .k { font-size:.82rem; color:var(--muted); }
+    .stat .v { font-size:1.35rem; font-weight:700; margin-top:6px; }
+    .task-list { display:grid; gap:10px; }
+    .task-btn { width:100%; text-align:left; border:1px solid var(--line); border-radius:14px; padding:12px; background:#fff; cursor:pointer; }
+    .task-btn.active, .task-btn:hover { border-color:var(--accent2); box-shadow:0 6px 14px rgba(2,132,199,.12); }
+    .toolbar, .row { display:flex; gap:10px; flex-wrap:wrap; align-items:center; }
+    .toolbar { margin-bottom:12px; }
+    button, input, textarea { font:inherit; border-radius:12px; }
+    button { border:none; padding:10px 14px; cursor:pointer; font-weight:700; }
+    .primary { background:linear-gradient(135deg,var(--accent),var(--accent2)); color:#fff; }
+    .secondary { background:#fff; color:var(--ink); border:1px solid var(--line); }
+    input, textarea { width:100%; border:1px solid var(--line); padding:12px; }
+    textarea { min-height:96px; resize:vertical; }
+    .query { padding:14px; background:#f8fafc; border:1px solid var(--line); border-radius:16px; font-weight:600; margin-bottom:12px; }
+    .feedback { min-height:22px; color:#b45309; font-weight:600; }
+    .chunks { display:grid; grid-template-columns:repeat(auto-fill,minmax(250px,1fr)); gap:12px; }
+    .chunk { border:1px solid var(--line); border-radius:16px; padding:12px; background:#fff; }
+    .chunk.selected { border-color:var(--accent); background:#f0fdfa; }
+    .pill { display:inline-block; margin:0 6px 6px 0; padding:3px 8px; border-radius:999px; background:#eef6ff; color:#1d4ed8; font-size:.75rem; }
+    details { margin-top:16px; }
+    pre { white-space:pre-wrap; word-break:break-word; font-size:.86rem; background:#0f172a; color:#e2e8f0; padding:14px; border-radius:14px; overflow:auto; }
+    .small-input { width:160px; }
+    @media (max-width:980px) { .hero,.main,.panels,.stats { grid-template-columns:1fr; } .small-input { width:100%; } }
+  </style>
+</head>
+<body>
+  <div class="wrap">
+    <div class="hero">
+      <div class="card">
+        <h1>RAG Context Optimizer</h1>
+        <p class="muted">Use any prompt you want. The UI can reset a session with your custom query, token budget, and step limit, then help minimize token usage by selecting only high-value chunks and compressing heavy evidence before answering.</p>
+        <div class="stats">
+          <div class="stat"><div class="k">Token Budget</div><div class="v" id="budgetStat">-</div></div>
+          <div class="stat"><div class="k">Tokens Used</div><div class="v" id="usedStat">-</div></div>
+          <div class="stat"><div class="k">Step</div><div class="v" id="stepStat">-</div></div>
+        </div>
+      </div>
+      <div class="card">
+        <h3>How To Use It</h3>
+        <p class="muted">Pick a task as a corpus shape, then overwrite the prompt with anything you want. Use Auto Optimize to choose chunks by relevance-per-token, or run manually. The API endpoints remain available at <code>/reset</code>, <code>/step</code>, and <code>/state</code>.</p>
+        <div class="row">
+          <a href="/docs" target="_blank"><button class="secondary" type="button">API Docs</button></a>
+          <a href="/health" target="_blank"><button class="secondary" type="button">Health</button></a>
+        </div>
+      </div>
+    </div>
+    <div class="main">
+      <aside class="card">
+        <h3>Task Presets</h3>
+        <div id="taskList" class="task-list"></div>
+      </aside>
+      <section class="card">
+        <div class="toolbar">
+          <button id="resetBtn" class="primary" type="button">Start / Reset</button>
+          <button id="autoStepBtn" class="secondary" type="button">Auto Optimize Step</button>
+          <button id="autoRunBtn" class="secondary" type="button">Auto Run</button>
+          <button id="refreshBtn" class="secondary" type="button">Refresh State</button>
+        </div>
+        <div class="row">
+          <input id="tokenBudgetInput" class="small-input" type="number" min="50" step="10" placeholder="Token budget" />
+          <input id="maxStepsInput" class="small-input" type="number" min="1" step="1" placeholder="Max steps" />
+        </div>
+        <div style="margin-top:12px;">
+          <label for="customQueryInput"><strong>Prompt / Query</strong></label>
+          <textarea id="customQueryInput" placeholder="Enter any prompt you want to optimize for minimal token usage."></textarea>
+        </div>
+        <div id="queryBox" class="query">Set a prompt and press Start / Reset.</div>
+        <div id="feedbackBox" class="feedback"></div>
+        <h3>Available Chunks</h3>
+        <div id="chunkGrid" class="chunks"></div>
+        <h3 style="margin-top:18px;">Answer</h3>
+        <textarea id="answerInput" placeholder="Write your final answer here."></textarea>
+        <div class="row" style="margin-top:10px;">
+          <button id="submitBtn" class="primary" type="button">Submit Answer</button>
+        </div>
+        <div class="panels">
+          <details class="card" open>
+            <summary><strong>Observation</strong></summary>
+            <pre id="observationPanel">{}</pre>
+          </details>
+          <details class="card">
+            <summary><strong>State</strong></summary>
+            <pre id="statePanel">{}</pre>
+          </details>
+        </div>
+      </section>
+    </div>
+  </div>
+  <script>
+    const taskList = document.getElementById("taskList");
+    const chunkGrid = document.getElementById("chunkGrid");
+    const queryBox = document.getElementById("queryBox");
+    const observationPanel = document.getElementById("observationPanel");
+    const statePanel = document.getElementById("statePanel");
+    const feedbackBox = document.getElementById("feedbackBox");
+    const budgetStat = document.getElementById("budgetStat");
+    const usedStat = document.getElementById("usedStat");
+    const stepStat = document.getElementById("stepStat");
+    const answerInput = document.getElementById("answerInput");
+    const customQueryInput = document.getElementById("customQueryInput");
+    const tokenBudgetInput = document.getElementById("tokenBudgetInput");
+    const maxStepsInput = document.getElementById("maxStepsInput");
+    let selectedTask = "single_domain_qa";
+    let currentObservation = null;
+    let allTasks = [];
+    async function fetchJson(url, options = {}) {
+      const response = await fetch(url, { headers: { "Content-Type": "application/json" }, ...options });
+      const body = await response.json();
+      if (!response.ok) throw new Error(body.detail || JSON.stringify(body));
+      return body;
+    }
+    function syncInputsFromTask() {
+      const task = allTasks.find(item => item.name === selectedTask);
+      if (!task) return;
+      if (!customQueryInput.dataset.userEdited || !customQueryInput.value.trim()) customQueryInput.value = task.query;
+      tokenBudgetInput.value = task.token_budget;
+      maxStepsInput.value = task.max_steps;
+    }
+    function renderTasks(tasks) {
+      taskList.innerHTML = "";
+      tasks.forEach(task => {
+        const btn = document.createElement("button");
+        btn.className = "task-btn" + (task.name === selectedTask ? " active" : "");
+        btn.innerHTML = `<div><strong>${task.name}</strong></div><div class="muted">${task.difficulty} · budget ${task.token_budget} · max steps ${task.max_steps}</div><div class="muted">${task.description}</div>`;
+        btn.onclick = () => {
+          selectedTask = task.name;
+          customQueryInput.dataset.userEdited = "";
+          renderTasks(tasks);
+          syncInputsFromTask();
+        };
+        taskList.appendChild(btn);
+      });
+    }
+    function renderObservation(observation) {
+      currentObservation = observation;
+      queryBox.textContent = observation.query;
+      budgetStat.textContent = observation.token_budget;
+      usedStat.textContent = observation.total_tokens_used;
+      stepStat.textContent = observation.step_number;
+      feedbackBox.textContent = observation.last_action_feedback || "";
+      observationPanel.textContent = JSON.stringify(observation, null, 2);
+      const selected = new Set(observation.selected_chunks || []);
+      chunkGrid.innerHTML = "";
+      observation.available_chunks.forEach(chunk => {
+        const card = document.createElement("div");
+        card.className = "chunk" + (selected.has(chunk.chunk_id) ? " selected" : "");
+        card.innerHTML = `
+          <h4>${chunk.chunk_id}</h4>
+          <div class="row"><span class="pill">${chunk.domain}</span><span class="pill">${chunk.tokens} tokens</span></div>
+          <div>${chunk.keywords.map(keyword => `<span class="pill">${keyword}</span>`).join("")}</div>
+          <div class="row" style="margin-top:10px;">
+            ${selected.has(chunk.chunk_id)
+              ? `<button class="secondary" data-action="deselect" data-id="${chunk.chunk_id}" type="button">Deselect</button>`
+              : `<button class="secondary" data-action="select" data-id="${chunk.chunk_id}" type="button">Select</button>`}
+            <button class="secondary" data-action="compress" data-id="${chunk.chunk_id}" type="button">Compress 50%</button>
+          </div>
+        `;
+        chunkGrid.appendChild(card);
+      });
+      chunkGrid.querySelectorAll("button[data-action]").forEach(btn => {
+        btn.addEventListener("click", async () => {
+          const action = btn.dataset.action;
+          const chunkId = btn.dataset.id;
+          if (action === "select") await step({ action_type: "select_chunk", chunk_id: chunkId });
+          if (action === "deselect") await step({ action_type: "deselect_chunk", chunk_id: chunkId });
+          if (action === "compress") await step({ action_type: "compress_chunk", chunk_id: chunkId, compression_ratio: 0.5 });
+        });
+      });
+    }
+    async function refreshState() {
+      const state = await fetchJson("/state", { method: "GET" });
+      statePanel.textContent = JSON.stringify(state, null, 2);
+    }
+    async function resetTask() {
+      try {
+        const body = await fetchJson("/reset", {
+          method: "POST",
+          body: JSON.stringify({
+            task_name: selectedTask,
+            custom_query: customQueryInput.value,
+            token_budget: tokenBudgetInput.value ? Number(tokenBudgetInput.value) : null,
+            max_steps: maxStepsInput.value ? Number(maxStepsInput.value) : null
+          }),
+        });
+        renderObservation(body.observation);
+        await refreshState();
+      } catch (error) {
+        feedbackBox.textContent = error.message;
+      }
+    }
+    async function step(payload) {
+      try {
+        const body = await fetchJson("/step", { method: "POST", body: JSON.stringify(payload) });
+        renderObservation(body.observation);
+        if (body.info && body.info.grader_breakdown) {
+          feedbackBox.textContent = `Final score: ${Number(body.reward).toFixed(4)} | ${JSON.stringify(body.info.grader_breakdown)}`;
+        }
+        await refreshState();
+      } catch (error) {
+        feedbackBox.textContent = error.message;
+      }
+    }
+    async function optimizeStep() {
+      try {
+        const suggestion = await fetchJson("/optimize-step", { method: "POST" });
+        feedbackBox.textContent = "Optimizer suggestion: " + JSON.stringify(suggestion);
+        await step(suggestion);
+      } catch (error) {
+        feedbackBox.textContent = error.message;
+      }
+    }
+    async function autoRun() {
+      for (let i = 0; i < 12; i += 1) {
+        if (!currentObservation) break;
+        const suggestion = await fetchJson("/optimize-step", { method: "POST" });
+        await step(suggestion);
+        if (suggestion.action_type === "submit_answer") break;
+      }
+    }
+    document.getElementById("resetBtn").addEventListener("click", resetTask);
+    document.getElementById("autoStepBtn").addEventListener("click", optimizeStep);
+    document.getElementById("autoRunBtn").addEventListener("click", autoRun);
+    document.getElementById("refreshBtn").addEventListener("click", refreshState);
+    document.getElementById("submitBtn").addEventListener("click", async () => {
+      await step({ action_type: "submit_answer", answer: answerInput.value || "A concise answer synthesized from the selected evidence." });
+    });
+    customQueryInput.addEventListener("input", () => { customQueryInput.dataset.userEdited = "true"; });
+    (async function init() {
+      allTasks = await fetchJson("/tasks", { method: "GET" });
+      renderTasks(allTasks);
+      syncInputsFromTask();
+      await resetTask();
+    })();
+  </script>
+</body>
+</html>
+"""
+UI_HTML_V2 = """
+<!doctype html>
+<html lang="en">
 <head>
   <meta charset="utf-8" />
   <meta name="viewport" content="width=device-width, initial-scale=1" />
   <title>rag-context-optimizer</title>
   <style>
     :root {
+      --bg: #f3f7fb;
       --panel: #ffffff;
+      --ink: #10243c;
+      --muted: #61748a;
+      --line: #d8e2ef;
       --accent: #0f766e;
+      --accent-alt: #0369a1;
       --warn: #b45309;
+      --selected: #ecfeff;
     }
     * { box-sizing: border-box; }
     body {
       margin: 0;
       font-family: ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
+      background: radial-gradient(circle at top left, #ffffff 0%, var(--bg) 65%);
       color: var(--ink);
     }
+    .wrap { max-width: 1360px; margin: 0 auto; padding: 24px; }
+    .grid { display: grid; grid-template-columns: 320px 1fr; gap: 16px; }
+    .stack { display: grid; gap: 16px; }
     .card {
       background: var(--panel);
       border: 1px solid var(--line);
+      border-radius: 22px;
       padding: 18px;
+      box-shadow: 0 10px 24px rgba(15, 23, 42, 0.05);
     }
     h1, h2, h3, p { margin-top: 0; }
+    .muted { color: var(--muted); line-height: 1.5; }
+    .stats { display: grid; grid-template-columns: repeat(4, 1fr); gap: 12px; margin-top: 16px; }
+    .stat { border: 1px solid var(--line); border-radius: 16px; padding: 12px; background: #fbfdff; }
+    .stat .label { font-size: 0.8rem; color: var(--muted); }
+    .stat .value { margin-top: 6px; font-size: 1.35rem; font-weight: 700; }
+    .task-list, .chunk-grid { display: grid; gap: 12px; }
+    .chunk-grid { grid-template-columns: repeat(auto-fill, minmax(255px, 1fr)); }
     .task-btn {
       width: 100%;
       text-align: left;
       border: 1px solid var(--line);
+      border-radius: 16px;
+      padding: 12px;
       background: #fff;
       cursor: pointer;
+      font: inherit;
     }
     .task-btn:hover, .task-btn.active {
+      border-color: var(--accent-alt);
+      box-shadow: 0 8px 18px rgba(3, 105, 161, 0.10);
     }
+    .toolbar, .row { display: flex; gap: 10px; flex-wrap: wrap; align-items: center; }
+    button, input, textarea { border-radius: 12px; font: inherit; }
+    button {
       cursor: pointer;
+      padding: 10px 14px;
+      border: 1px solid transparent;
       font-weight: 700;
     }
+    .primary { background: linear-gradient(135deg, var(--accent), var(--accent-alt)); color: #fff; }
+    .secondary { background: #fff; border-color: var(--line); color: var(--ink); }
+    input, textarea { width: 100%; border: 1px solid var(--line); padding: 12px; background: #fff; }
+    textarea { resize: vertical; min-height: 110px; }
+    .query-box {
       padding: 14px;
       border: 1px solid var(--line);
       border-radius: 16px;
+      background: #f8fafc;
+      font-weight: 600;
+      min-height: 60px;
     }
+    .feedback { min-height: 24px; color: var(--warn); font-weight: 600; }
+    .chunk-card { border: 1px solid var(--line); border-radius: 18px; padding: 12px; background: #fff; }
+    .chunk-card.selected { border-color: var(--accent); background: var(--selected); }
     .pill {
       display: inline-block;
+      margin: 0 6px 6px 0;
+      padding: 4px 8px;
       border-radius: 999px;
+      background: #eff6ff;
       color: #1d4ed8;
       font-size: 0.75rem;
     }
+    .json-panels { display: grid; grid-template-columns: 1fr 1fr; gap: 16px; margin-top: 16px; }
     pre {
       margin: 0;
       white-space: pre-wrap;
       word-break: break-word;
+      font-size: 0.86rem;
+      line-height: 1.45;
       background: #0f172a;
       color: #e2e8f0;
       padding: 14px;
+      border-radius: 14px;
+      overflow: auto;
+      max-height: 420px;
     }
+    .helper-list { padding-left: 18px; margin: 0; color: var(--muted); }
+    .helper-list li + li { margin-top: 8px; }
+    .mono { font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; }
+    @media (max-width: 980px) { .grid, .json-panels, .stats { grid-template-columns: 1fr; } }
   </style>
 </head>
 <body>
+  <div class="wrap">
+    <div class="card" style="margin-bottom:16px;">
+      <h1>RAG Context Optimizer</h1>
+      <p class="muted">
+        This UI works with any prompt you provide. Pick a preset only to choose a corpus shape and default budget, then replace the query with your own prompt and let the optimizer choose the highest-value evidence for the fewest tokens possible.
+      </p>
+      <div class="stats">
+        <div class="stat"><div class="label">Task Preset</div><div class="value" id="taskStat">-</div></div>
+        <div class="stat"><div class="label">Token Budget</div><div class="value" id="budgetStat">-</div></div>
+        <div class="stat"><div class="label">Tokens Used</div><div class="value" id="usedStat">-</div></div>
+        <div class="stat"><div class="label">Step</div><div class="value" id="stepStat">-</div></div>
       </div>
+    </div>
+    <div class="grid">
+      <aside class="stack">
+        <div class="card">
+          <h3>Task Presets</h3>
+          <p class="muted">These only set the default corpus flavor, budget, and step limit. Your prompt can be anything.</p>
+          <div id="taskList" class="task-list"></div>
+        </div>
+        <div class="card">
+          <h3>How It Optimizes</h3>
+          <ul class="helper-list">
+            <li>Ranks chunks by relevance per token, not just raw overlap.</li>
+            <li>Compresses heavier evidence before wasting budget.</li>
+            <li>Stops early once it has enough support to answer.</li>
+            <li>Keeps the OpenEnv endpoints live at <span class="mono">/reset</span>, <span class="mono">/step</span>, and <span class="mono">/state</span>.</li>
+          </ul>
+          <div class="toolbar" style="margin-top:16px;">
+            <a href="/docs" target="_blank"><button class="secondary" type="button">API Docs</button></a>
+            <a href="/health" target="_blank"><button class="secondary" type="button">Health</button></a>
+          </div>
         </div>
       </aside>
+      <section class="stack">
+        <div class="card">
+          <div class="toolbar" style="margin-bottom:12px;">
+            <button id="resetBtn" class="primary" type="button">Start / Reset</button>
+            <button id="autoStepBtn" class="secondary" type="button">Auto Optimize Step</button>
+            <button id="autoRunBtn" class="secondary" type="button">Auto Run</button>
+            <button id="refreshBtn" class="secondary" type="button">Refresh State</button>
+          </div>
+          <div class="row">
+            <input id="tokenBudgetInput" type="number" min="50" step="10" placeholder="Token budget" style="max-width:180px;" />
+            <input id="maxStepsInput" type="number" min="1" step="1" placeholder="Max steps" style="max-width:180px;" />
+          </div>
+          <div style="margin-top:12px;">
+            <label for="customQueryInput"><strong>Custom Prompt</strong></label>
+            <textarea id="customQueryInput" placeholder="Enter any prompt you want to optimize for minimal token usage."></textarea>
+          </div>
+          <div style="margin-top:12px;">
+            <label for="answerInput"><strong>Manual Final Answer</strong></label>
+            <textarea id="answerInput" placeholder="Optional. If left blank, auto-optimize will still submit an answer."></textarea>
+          </div>
+          <div class="toolbar" style="margin-top:12px;">
+            <button id="submitBtn" class="primary" type="button">Submit Manual Answer</button>
+          </div>
         </div>
+        <div class="card">
+          <h3>Active Query</h3>
+          <div id="queryBox" class="query-box">No episode has started yet. Enter your prompt, then press Start / Reset.</div>
+          <div id="feedbackBox" class="feedback" style="margin-top:12px;"></div>
+        </div>
+        <div class="card">
+          <div class="toolbar" style="justify-content:space-between;">
+            <h3 style="margin-bottom:0;">Available Chunks</h3>
+            <div class="muted">Select manually or let Auto Optimize choose the best evidence per token.</div>
+          </div>
+          <div id="chunkGrid" class="chunk-grid" style="margin-top:12px;"></div>
         </div>
+        <div class="json-panels">
+          <div class="card">
             <h3>Observation</h3>
             <pre id="observationPanel">{}</pre>
           </div>
+          <div class="card">
             <h3>State</h3>
             <pre id="statePanel">{}</pre>
           </div>
         </div>
+      </section>
+    </div>
   </div>
   <script>
     const taskList = document.getElementById("taskList");
     const chunkGrid = document.getElementById("chunkGrid");
     const budgetStat = document.getElementById("budgetStat");
     const usedStat = document.getElementById("usedStat");
     const stepStat = document.getElementById("stepStat");
+    const taskStat = document.getElementById("taskStat");
     const answerInput = document.getElementById("answerInput");
+    const customQueryInput = document.getElementById("customQueryInput");
+    const tokenBudgetInput = document.getElementById("tokenBudgetInput");
+    const maxStepsInput = document.getElementById("maxStepsInput");
     let selectedTask = "single_domain_qa";
     let currentObservation = null;
+    let currentDone = false;
+    let allTasks = [];
     async function fetchJson(url, options = {}) {
       const response = await fetch(url, {
         headers: { "Content-Type": "application/json" },
+        ...options
       });
       const body = await response.json();
+      if (!response.ok) throw new Error(body.detail || JSON.stringify(body));
       return body;
     }
+    function getSelectedTaskMeta() {
+      return allTasks.find((task) => task.name === selectedTask);
+    }
+    function syncInputsFromTask() {
+      const task = getSelectedTaskMeta();
+      if (!task) return;
+      tokenBudgetInput.value = task.token_budget;
+      maxStepsInput.value = task.max_steps;
+    }
     function renderTasks(tasks) {
       taskList.innerHTML = "";
+      tasks.forEach((task) => {
         const btn = document.createElement("button");
+        btn.type = "button";
         btn.className = "task-btn" + (task.name === selectedTask ? " active" : "");
+        btn.innerHTML = `
+          <div><strong>${task.name}</strong></div>
+          <div class="muted">${task.difficulty} | budget ${task.token_budget} | max steps ${task.max_steps}</div>
+          <div class="muted">${task.description}</div>
+        `;
+        btn.addEventListener("click", () => {
           selectedTask = task.name;
           renderTasks(tasks);
+          syncInputsFromTask();
+          taskStat.textContent = task.name;
+          feedbackBox.textContent = "Preset changed. Add your prompt and press Start / Reset.";
+        });
         taskList.appendChild(btn);
       });
     }
     function renderObservation(observation) {
       currentObservation = observation;
+      queryBox.textContent = observation.query || "No active query.";
       budgetStat.textContent = observation.token_budget;
       usedStat.textContent = observation.total_tokens_used;
       stepStat.textContent = observation.step_number;
+      taskStat.textContent = observation.task_name;
+      if (observation.last_action_feedback) feedbackBox.textContent = observation.last_action_feedback;
       observationPanel.textContent = JSON.stringify(observation, null, 2);
       const selected = new Set(observation.selected_chunks || []);
       chunkGrid.innerHTML = "";
+      observation.available_chunks.forEach((chunk) => {
         const card = document.createElement("div");
+        card.className = "chunk-card" + (selected.has(chunk.chunk_id) ? " selected" : "");
+        const keywordHtml = (chunk.keywords || []).map((keyword) => `<span class="pill">${keyword}</span>`).join("");
         card.innerHTML = `
           <h4>${chunk.chunk_id}</h4>
+          <div style="margin-bottom:8px;">
+            <span class="pill">${chunk.domain}</span>
+            <span class="pill">${chunk.tokens} tokens</span>
+          </div>
+          <div>${keywordHtml}</div>
+          <div class="toolbar" style="margin-top:12px;">
+            <button class="secondary" type="button" data-action="${selected.has(chunk.chunk_id) ? "deselect_chunk" : "select_chunk"}" data-id="${chunk.chunk_id}">
+              ${selected.has(chunk.chunk_id) ? "Deselect" : "Select"}
+            </button>
+            <button class="secondary" type="button" data-action="compress_chunk" data-id="${chunk.chunk_id}">Compress 50%</button>
           </div>
         `;
         chunkGrid.appendChild(card);
       });
+      chunkGrid.querySelectorAll("button[data-action]").forEach((button) => {
+        button.addEventListener("click", async () => {
+          const actionType = button.dataset.action;
+          const chunkId = button.dataset.id;
+          if (actionType === "compress_chunk") {
+            await step({ action_type: actionType, chunk_id: chunkId, compression_ratio: 0.5 });
+          } else {
+            await step({ action_type: actionType, chunk_id: chunkId });
           }
         });
       });
     }
     async function refreshState() {
+      try {
+        const state = await fetchJson("/state", { method: "GET" });
+        statePanel.textContent = JSON.stringify(state, null, 2);
+      } catch (error) {
+        statePanel.textContent = JSON.stringify({ error: error.message }, null, 2);
+      }
     }
     async function resetTask() {
+      const customQuery = customQueryInput.value.trim();
+      if (!customQuery) {
+        feedbackBox.textContent = "Please enter your own prompt first.";
+        return;
+      }
+      try {
+        const body = await fetchJson("/reset", {
+          method: "POST",
+          body: JSON.stringify({
+            task_name: selectedTask,
+            custom_query: customQuery,
+            token_budget: tokenBudgetInput.value ? Number(tokenBudgetInput.value) : null,
+            max_steps: maxStepsInput.value ? Number(maxStepsInput.value) : null
+          })
+        });
+        currentDone = false;
+        renderObservation(body.observation);
+        feedbackBox.textContent = "Episode started with your custom prompt.";
+        await refreshState();
+      } catch (error) {
+        feedbackBox.textContent = error.message;
+      }
     }
     async function step(payload) {
+      if (!currentObservation) {
+        feedbackBox.textContent = "Start an episode first.";
+        return null;
+      }
       try {
         const body = await fetchJson("/step", {
           method: "POST",
+          body: JSON.stringify(payload)
         });
         renderObservation(body.observation);
+        currentDone = Boolean(body.done);
         if (body.info && body.info.grader_breakdown) {
+          feedbackBox.textContent = `Final score ${Number(body.reward || 0).toFixed(4)} | ${JSON.stringify(body.info.grader_breakdown)}`;
         }
         await refreshState();
+        return body;
+      } catch (error) {
+        feedbackBox.textContent = error.message;
+        return null;
+      }
+    }
+    async function optimizeStep() {
+      if (!currentObservation) {
+        feedbackBox.textContent = "Start an episode first.";
+        return;
+      }
+      try {
+        const suggestion = await fetchJson("/optimize-step", { method: "POST" });
+        feedbackBox.textContent = "Optimizer chose: " + JSON.stringify(suggestion);
+        await step(suggestion);
       } catch (error) {
         feedbackBox.textContent = error.message;
       }
     }
+    async function autoRun() {
+      if (!currentObservation) {
+        feedbackBox.textContent = "Start an episode first.";
+        return;
+      }
+      for (let index = 0; index < 20; index += 1) {
+        if (currentDone) break;
+        const suggestion = await fetchJson("/optimize-step", { method: "POST" });
+        const result = await step(suggestion);
+        if (!result || result.done || suggestion.action_type === "submit_answer") break;
+      }
+    }
     document.getElementById("resetBtn").addEventListener("click", resetTask);
+    document.getElementById("autoStepBtn").addEventListener("click", optimizeStep);
+    document.getElementById("autoRunBtn").addEventListener("click", autoRun);
     document.getElementById("refreshBtn").addEventListener("click", refreshState);
     document.getElementById("submitBtn").addEventListener("click", async () => {
+      const manualAnswer = answerInput.value.trim() || "Concise answer synthesized from the selected evidence.";
+      await step({ action_type: "submit_answer", answer: manualAnswer });
     });
     (async function init() {
+      allTasks = await fetchJson("/tasks", { method: "GET" });
+      renderTasks(allTasks);
+      syncInputsFromTask();
+      taskStat.textContent = selectedTask;
+      feedbackBox.textContent = "Add any prompt you want, then press Start / Reset.";
+      observationPanel.textContent = JSON.stringify({ message: "No active episode yet." }, null, 2);
+      statePanel.textContent = JSON.stringify({ message: "No active episode yet." }, null, 2);
     })();
   </script>
 </body>
 @app.get("/", response_class=HTMLResponse)
 async def home_page():
+    return HTMLResponse(UI_HTML_V2)
 def _serialize_observation(observation: Any) -> dict[str, Any]:
     }
+def _tokenize(text: str) -> set[str]:
+    import re
+    return set(re.findall(r"[a-z0-9]+", text.lower()))
+def _suggest_action(env: RagContextOptimizerEnv) -> dict[str, Any]:
+    observation = env._build_observation()
+    query_terms = _tokenize(observation.query)
+    selected = set(observation.selected_chunks)
+    remaining_budget = observation.token_budget - observation.total_tokens_used
+    def score_chunk(chunk: Any) -> tuple[float, int, str]:
+        keyword_terms = _tokenize(" ".join(chunk.keywords))
+        overlap = len(query_terms & keyword_terms)
+        union = len(query_terms | keyword_terms) or 1
+        relevance = overlap / union
+        ratio = relevance / max(chunk.tokens, 1)
+        return (-ratio, chunk.tokens, chunk.chunk_id)
+    selected_chunks = [chunk for chunk in observation.available_chunks if chunk.chunk_id in selected]
+    if selected_chunks and (
+        observation.total_tokens_used >= int(observation.token_budget * 0.65)
+        or observation.step_number >= 3
+    ):
+        heavy = sorted(selected_chunks, key=lambda chunk: (-chunk.tokens, chunk.chunk_id))
+        if heavy and heavy[0].tokens > max(120, observation.token_budget // 4):
+            return {
+                "action_type": "compress_chunk",
+                "chunk_id": heavy[0].chunk_id,
+                "compression_ratio": 0.5,
+            }
+    if len(selected) >= 2 or observation.step_number >= max(2, env.task.max_steps - 2):
+        chosen_keywords: list[str] = []
+        for chunk in selected_chunks[:3]:
+            chosen_keywords.extend(chunk.keywords[:2])
+        answer = (
+            "Optimized answer based on selected evidence: "
+            + ", ".join(chosen_keywords[:6])
+            if chosen_keywords
+            else "Optimized answer based on the currently selected evidence."
+        )
+        return {"action_type": "submit_answer", "answer": answer}
+    available = [chunk for chunk in observation.available_chunks if chunk.chunk_id not in selected]
+    for chunk in sorted(available, key=score_chunk):
+        if chunk.tokens <= remaining_budget:
+            return {"action_type": "select_chunk", "chunk_id": chunk.chunk_id}
+    if selected_chunks:
+        return {
+            "action_type": "submit_answer",
+            "answer": "Optimized answer based on the currently selected evidence.",
+        }
+    if available:
+        smallest_chunk = min(available, key=lambda chunk: (chunk.tokens, chunk.chunk_id))
+        return {
+            "action_type": "submit_answer",
+            "answer": (
+                "No chunk fits within the current token budget. "
+                f"Increase the budget to at least {smallest_chunk.tokens} tokens or choose a broader budget."
+            ),
+        }
+    return {"action_type": "submit_answer", "answer": "No usable evidence was available."}
 @app.post("/reset")
 async def reset_endpoint(payload: ResetRequest):
     if payload.task_name not in TASKS_BY_NAME:
         raise HTTPException(status_code=400, detail="Unknown task_name.")
+    env = RagContextOptimizerEnv(
+        task_name=payload.task_name,
+        query_override=payload.custom_query,
+        token_budget_override=payload.token_budget,
+        max_steps_override=payload.max_steps,
+    )
     app.state.env = env
     result = await env.reset()
     return _serialize_step_result(result, reset=True)
             "description": task.description,
             "difficulty": task.difficulty,
             "token_budget": task.token_budget,
+            "query": task.query,
+            "max_steps": task.max_steps,
         }
         for task in ALL_TASKS
     ]
+@app.post("/optimize-step")
+async def optimize_step_endpoint():
+    env = getattr(app.state, "env", None)
+    if env is None:
+        raise HTTPException(status_code=400, detail="Environment is not initialized. Call /reset first.")
+    return _suggest_action(env)
 if __name__ == "__main__":
     import uvicorn

env/__pycache__/environment.cpython-314.pyc CHANGED Viewed

Binary files a/env/__pycache__/environment.cpython-314.pyc and b/env/__pycache__/environment.cpython-314.pyc differ

env/environment.py CHANGED Viewed

@@ -4,7 +4,7 @@ Main OpenEnv-style environment for rag-context-optimizer.
 from __future__ import annotations
-from dataclasses import asdict, dataclass, is_dataclass
 from pathlib import Path
 from typing import Any
@@ -24,7 +24,13 @@ class StepResult:
 class RagContextOptimizerEnv:
-    def __init__(self, task_name: str = "single_domain_qa"):
         if task_name not in TASKS_BY_NAME:
             raise ValueError(f"Unknown task_name: {task_name}")
@@ -32,7 +38,12 @@ class RagContextOptimizerEnv:
         self._all_chunks = load_corpus(self._corpus_path)
         self.retriever = HybridRetriever(self._all_chunks)
         self.grader = TaskGrader()
-        self.task: Task = TASKS_BY_NAME[task_name]
         self._available_chunks: list[Chunk] = []
         self._selected_chunks: list[str] = []
@@ -42,6 +53,22 @@ class RagContextOptimizerEnv:
         self._last_action_feedback: str | None = None
         self._last_answer = ""
     async def reset(self) -> StepResult:
         self._available_chunks = self._filter_chunks_for_task(self.task)
         self._selected_chunks = []

 from __future__ import annotations
+from dataclasses import asdict, dataclass, is_dataclass, replace
 from pathlib import Path
 from typing import Any
 class RagContextOptimizerEnv:
+    def __init__(
+        self,
+        task_name: str = "single_domain_qa",
+        query_override: str | None = None,
+        token_budget_override: int | None = None,
+        max_steps_override: int | None = None,
+    ):
         if task_name not in TASKS_BY_NAME:
             raise ValueError(f"Unknown task_name: {task_name}")
         self._all_chunks = load_corpus(self._corpus_path)
         self.retriever = HybridRetriever(self._all_chunks)
         self.grader = TaskGrader()
+        self.task: Task = self._build_task(
+            TASKS_BY_NAME[task_name],
+            query_override=query_override,
+            token_budget_override=token_budget_override,
+            max_steps_override=max_steps_override,
+        )
         self._available_chunks: list[Chunk] = []
         self._selected_chunks: list[str] = []
         self._last_action_feedback: str | None = None
         self._last_answer = ""
+    @staticmethod
+    def _build_task(
+        base_task: Task,
+        query_override: str | None = None,
+        token_budget_override: int | None = None,
+        max_steps_override: int | None = None,
+    ) -> Task:
+        updated_task = base_task
+        if query_override and query_override.strip():
+            updated_task = replace(updated_task, query=query_override.strip())
+        if token_budget_override is not None and token_budget_override > 0:
+            updated_task = replace(updated_task, token_budget=token_budget_override)
+        if max_steps_override is not None and max_steps_override > 0:
+            updated_task = replace(updated_task, max_steps=max_steps_override)
+        return updated_task
     async def reset(self) -> StepResult:
         self._available_chunks = self._filter_chunks_for_task(self.task)
         self._selected_chunks = []

streamlit_app.py CHANGED Viewed

@@ -6,9 +6,9 @@ import streamlit as st
 API_URL = st.secrets.get("API_URL", "http://localhost:7860") if hasattr(st, "secrets") else "http://localhost:7860"
-st.set_page_config(page_title="rag-context-optimizer", page_icon="📚", layout="wide")
 st.title("RAG Context Optimizer")
-st.caption("Streamlit control panel for the benchmark API.")
 def api_get(path: str):
@@ -17,52 +17,142 @@ def api_get(path: str):
     return response.json()
-def api_post(path: str, payload: dict):
-    response = requests.post(f"{API_URL}{path}", json=payload, timeout=20)
     response.raise_for_status()
     return response.json()
-tasks = api_get("/tasks")
-task_names = [task["name"] for task in tasks]
-selected_task = st.sidebar.selectbox("Task", task_names)
-if st.sidebar.button("Reset Task", use_container_width=True):
-    st.session_state["reset_payload"] = api_post("/reset", {"task_name": selected_task})
-if "reset_payload" not in st.session_state:
-    st.session_state["reset_payload"] = api_post("/reset", {"task_name": selected_task})
-payload = st.session_state["reset_payload"]
 observation = payload["observation"]
-col1, col2, col3 = st.columns(3)
-col1.metric("Budget", observation["token_budget"])
-col2.metric("Used", observation["total_tokens_used"])
-col3.metric("Step", observation["step_number"])
-st.subheader("Query")
 st.info(observation["query"])
-st.subheader("Chunks")
-for chunk in observation["available_chunks"]:
-    with st.expander(f"{chunk['chunk_id']} · {chunk['domain']} · {chunk['tokens']} tokens"):
-        st.write(", ".join(chunk["keywords"]))
-        c1, c2 = st.columns(2)
-        if c1.button(f"Select {chunk['chunk_id']}", key=f"select-{chunk['chunk_id']}"):
-            st.session_state["reset_payload"] = api_post("/step", {"action_type": "select_chunk", "chunk_id": chunk["chunk_id"]})
-            st.rerun()
-        if c2.button(f"Compress {chunk['chunk_id']}", key=f"compress-{chunk['chunk_id']}"):
-            st.session_state["reset_payload"] = api_post("/step", {"action_type": "compress_chunk", "chunk_id": chunk["chunk_id"], "compression_ratio": 0.5})
-            st.rerun()
-answer = st.text_area("Answer", "A concise answer synthesized from the selected evidence.")
-if st.button("Submit Answer", type="primary"):
-    st.session_state["reset_payload"] = api_post("/step", {"action_type": "submit_answer", "answer": answer})
     st.rerun()
 st.subheader("Observation")
-st.json(st.session_state["reset_payload"])
 st.subheader("State")
 st.json(api_get("/state"))

 API_URL = st.secrets.get("API_URL", "http://localhost:7860") if hasattr(st, "secrets") else "http://localhost:7860"
+st.set_page_config(page_title="rag-context-optimizer", page_icon="R", layout="wide")
 st.title("RAG Context Optimizer")
+st.caption("Use any prompt, keep the token budget tight, and let the optimizer pick the best evidence per token.")
 def api_get(path: str):
     return response.json()
+def api_post(path: str, payload: dict | None = None):
+    response = requests.post(f"{API_URL}{path}", json=payload or {}, timeout=20)
     response.raise_for_status()
     return response.json()
+def start_episode(task_name: str, query: str, token_budget: int, max_steps: int):
+    st.session_state["payload"] = api_post(
+        "/reset",
+        {
+            "task_name": task_name,
+            "custom_query": query,
+            "token_budget": token_budget,
+            "max_steps": max_steps,
+        },
+    )
+def do_step(payload: dict):
+    st.session_state["payload"] = api_post("/step", payload)
+tasks = api_get("/tasks")
+task_map = {task["name"]: task for task in tasks}
+selected_task = st.sidebar.selectbox("Task preset", list(task_map))
+task_meta = task_map[selected_task]
+default_query = st.session_state.get("custom_query", "")
+custom_query = st.sidebar.text_area(
+    "Custom prompt",
+    value=default_query,
+    height=180,
+    placeholder="Enter any prompt you want to optimize for minimal token usage.",
+)
+token_budget = st.sidebar.number_input(
+    "Token budget",
+    min_value=50,
+    value=int(task_meta["token_budget"]),
+    step=10,
+)
+max_steps = st.sidebar.number_input(
+    "Max steps",
+    min_value=1,
+    value=int(task_meta["max_steps"]),
+    step=1,
+)
+st.session_state["custom_query"] = custom_query
+sidebar_cols = st.sidebar.columns(2)
+if sidebar_cols[0].button("Start / Reset", use_container_width=True):
+    if not custom_query.strip():
+        st.sidebar.error("Enter a custom prompt first.")
+    else:
+        start_episode(selected_task, custom_query.strip(), int(token_budget), int(max_steps))
+        st.rerun()
+if sidebar_cols[1].button("Refresh", use_container_width=True):
+    st.rerun()
+if "payload" not in st.session_state:
+    st.info("Add your prompt in the sidebar and press Start / Reset.")
+    st.stop()
+payload = st.session_state["payload"]
 observation = payload["observation"]
+col1, col2, col3, col4 = st.columns(4)
+col1.metric("Task", observation["task_name"])
+col2.metric("Budget", observation["token_budget"])
+col3.metric("Used", observation["total_tokens_used"])
+col4.metric("Step", observation["step_number"])
+st.subheader("Active Query")
 st.info(observation["query"])
+feedback = observation.get("last_action_feedback")
+if feedback:
+    st.warning(feedback)
+if payload.get("info", {}).get("grader_breakdown"):
+    st.success(f"Final score: {payload.get('reward', 0):.4f}")
+    st.json(payload["info"]["grader_breakdown"])
+action_cols = st.columns(3)
+if action_cols[0].button("Auto Optimize Step", use_container_width=True):
+    suggestion = api_post("/optimize-step")
+    do_step(suggestion)
+    st.rerun()
+if action_cols[1].button("Auto Run", use_container_width=True):
+    for _ in range(20):
+        suggestion = api_post("/optimize-step")
+        do_step(suggestion)
+        if suggestion["action_type"] == "submit_answer" or st.session_state["payload"]["done"]:
+            break
     st.rerun()
+manual_answer = action_cols[2].text_input("Manual answer", value="")
+if st.button("Submit Manual Answer", type="primary", use_container_width=True):
+    do_step(
+        {
+            "action_type": "submit_answer",
+            "answer": manual_answer.strip() or "Concise answer synthesized from the selected evidence.",
+        }
+    )
+    st.rerun()
+st.subheader("Available Chunks")
+chunk_columns = st.columns(2)
+for index, chunk in enumerate(observation["available_chunks"]):
+    selected = chunk["chunk_id"] in set(observation["selected_chunks"])
+    container = chunk_columns[index % 2].container(border=True)
+    container.markdown(f"**{chunk['chunk_id']}**")
+    container.caption(f"{chunk['domain']} | {chunk['tokens']} tokens")
+    container.write(", ".join(chunk["keywords"]))
+    c1, c2 = container.columns(2)
+    if selected:
+      if c1.button("Deselect", key=f"deselect-{chunk['chunk_id']}", use_container_width=True):
+          do_step({"action_type": "deselect_chunk", "chunk_id": chunk["chunk_id"]})
+          st.rerun()
+    else:
+      if c1.button("Select", key=f"select-{chunk['chunk_id']}", use_container_width=True):
+          do_step({"action_type": "select_chunk", "chunk_id": chunk["chunk_id"]})
+          st.rerun()
+    if c2.button("Compress 50%", key=f"compress-{chunk['chunk_id']}", use_container_width=True):
+        do_step(
+            {
+                "action_type": "compress_chunk",
+                "chunk_id": chunk["chunk_id"],
+                "compression_ratio": 0.5,
+            }
+        )
+        st.rerun()
 st.subheader("Observation")
+st.json(payload)
 st.subheader("State")
 st.json(api_get("/state"))