Upload 4 files

app.py

@@ -1230,6 +1230,22 @@ def _resolve_example_corpus(corpus_dir: Path, question: str) -> Path:
     return candidate if candidate.exists() else corpus_dir
 
 
+def _cleanup_runtime_state(runtime_state: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+    state = dict(runtime_state or {})
+    corpus_root = state.get("corpus_root")
+    session_dir = state.get("session_dir")
+    for path_value in (corpus_root, session_dir):
+        if not path_value:
+            continue
+        try:
+            path = Path(path_value)
+            if path.exists():
+                shutil.rmtree(path, ignore_errors=True)
+        except Exception:
+            pass
+    return {}
+
+
 def _ensure_corpus_ready(
     use_default: bool,
     uploaded_zip: Optional[str],
@@ -1836,32 +1852,58 @@ def run_search(
     mode: str,
     corpus_source: str,
     uploaded_zip: Optional[str],
-) -> Generator[tuple[str, str], None, None]:
-    """Yields (terminal_html, status_str)."""
+    runtime_state: Optional[Dict[str, Any]],
+) -> Generator[tuple[str, str, Dict[str, Any]], None, None]:
+    """Yields (terminal_html, status_str, runtime_state)."""
 
     if not api_key or not api_key.strip():
-        yield _TERM_IDLE, "⚠  OpenAI API Key is required."
+        yield _TERM_IDLE, "⚠  OpenAI API Key is required.", dict(runtime_state or {})
         return
     if not question or not question.strip():
-        yield _TERM_IDLE, "⚠  Please enter a question."
+        yield _TERM_IDLE, "⚠  Please enter a question.", dict(runtime_state or {})
         return
 
+    runtime = dict(runtime_state or {})
     use_default_corpus = corpus_source == DEFAULT_CORPUS_LABEL
     selected_domain = _selected_example_corpus(question) if use_default_corpus else None
-    try:
-        corpus_root = _ensure_corpus_ready(use_default_corpus, uploaded_zip, selected_domain)
-    except ValueError as exc:
-        yield _TERM_IDLE, f"⚠  {exc}"
-        return
-    except Exception as exc:
-        yield _TERM_IDLE, f"⚠  Corpus error: {exc}"
-        return
+    corpus_key = json.dumps(
+        {
+            "source": corpus_source,
+            "uploaded_zip": uploaded_zip or "",
+            "selected_domain": selected_domain or "",
+        },
+        ensure_ascii=False,
+        sort_keys=True,
+    )
+    corpus_root_value = runtime.get("corpus_root")
+    corpus_root = Path(corpus_root_value) if corpus_root_value else None
+    if runtime.get("corpus_key") != corpus_key or corpus_root is None or not corpus_root.exists():
+        runtime = _cleanup_runtime_state(runtime)
+        try:
+            corpus_root = _ensure_corpus_ready(use_default_corpus, uploaded_zip, selected_domain)
+        except ValueError as exc:
+            yield _TERM_IDLE, f"⚠  {exc}", runtime
+            return
+        except Exception as exc:
+            yield _TERM_IDLE, f"⚠  Corpus error: {exc}", runtime
+            return
+        session_dir = Path(tempfile.mkdtemp(prefix="pi_session_"))
+        runtime.update(
+            {
+                "corpus_key": corpus_key,
+                "corpus_root": str(corpus_root),
+                "session_dir": str(session_dir),
+                "session_file": "",
+                "session_id": "",
+            }
+        )
 
+    assert corpus_root is not None
     cwd = _resolve_example_corpus(corpus_root, question) if use_default_corpus else corpus_root
 
     prompt = build_ir_prompt(question, cwd) if mode == "IR" else build_benchmark_prompt(question, cwd)
     state = _initial_terminal_state(question, cwd, mode, model, max_turns, selected_domain)
-    yield _render_terminal_state(state), "⚙  Running…"
+    yield _render_terminal_state(state), "⚙  Running…", runtime
 
     try:
         for event in run_pi_stream(
@@ -1871,26 +1913,27 @@ def run_search(
             provider="openai",
             model=model,
             max_turns=max_turns,
+            session_dir=Path(runtime["session_dir"]) if runtime.get("session_dir") else None,
+            session_path=Path(runtime["session_file"]) if runtime.get("session_file") else None,
         ):
+            if event.get("type") == "__session__":
+                if event.get("sessionFile"):
+                    runtime["session_file"] = str(event["sessionFile"])
+                if event.get("sessionId"):
+                    runtime["session_id"] = str(event["sessionId"])
+                continue
             if _apply_event(state, event):
                 status = "✓  Completed" if event.get("type") == "__final__" else "⚙  Running…"
-                yield _render_terminal_state(state), status
+                yield _render_terminal_state(state), status, runtime
 
     except Exception as exc:
         _apply_event(state, {"type": "error", "error": str(exc)})
-        yield _render_terminal_state(state), f"⚠  {exc}"
+        yield _render_terminal_state(state), f"⚠  {exc}", runtime
         return
-    finally:
-        try:
-            p = cwd.parent
-            if p.name.startswith("pi_corpus_"):
-                shutil.rmtree(p, ignore_errors=True)
-        except Exception:
-            pass
 
     if not state.get("finalized"):
         _apply_event(state, {"type": "__final__", "text": ""})
-    yield _render_terminal_state(state), "✓  Completed"
+    yield _render_terminal_state(state), "✓  Completed", runtime
 
 
 # ──────────────────────────────────────────────────────────────────────────
@@ -1899,6 +1942,7 @@ def run_search(
 def build_ui() -> gr.Blocks:
     hero_logo_src = f"/gradio_api/file={HERO_LOGO_PATH}"
     with gr.Blocks(title="DCI-Agent Search") as demo:
+        runtime_state = gr.State({})
         with gr.Row(equal_height=False):
 
             # ── Left sidebar ───────────────────────────────────────────────
@@ -2053,28 +2097,30 @@ def build_ui() -> gr.Blocks:
             outputs=[uploaded_zip, corpus_hint],
         )
 
-        def clear_all():
+        def clear_all(runtime_state):
+            cleared_runtime = _cleanup_runtime_state(runtime_state)
             return (
-                _TERM_IDLE, "", "Ready ✔",
+                _TERM_IDLE, "", "Ready ✔", cleared_runtime,
             )
 
         clear_btn.click(
             fn=clear_all,
-            outputs=[terminal, terminal_question, status_text],
+            inputs=[runtime_state],
+            outputs=[terminal, terminal_question, status_text, runtime_state],
         )
 
-        def _run_with_qa_mode(api_key, question, model, max_turns, corpus_source, uploaded_zip):
-            yield from run_search(api_key, question, model, max_turns, "QA", corpus_source, uploaded_zip)
+        def _run_with_qa_mode(api_key, question, model, max_turns, corpus_source, uploaded_zip, runtime_state):
+            yield from run_search(api_key, question, model, max_turns, "QA", corpus_source, uploaded_zip, runtime_state)
 
         run_event = run_btn.click(
             fn=_run_with_qa_mode,
-            inputs=[api_key, terminal_question, model, max_turns, corpus_source, uploaded_zip],
-            outputs=[terminal, status_text],
+            inputs=[api_key, terminal_question, model, max_turns, corpus_source, uploaded_zip, runtime_state],
+            outputs=[terminal, status_text, runtime_state],
         )
         terminal_question.submit(
             fn=_run_with_qa_mode,
-            inputs=[api_key, terminal_question, model, max_turns, corpus_source, uploaded_zip],
-            outputs=[terminal, status_text],
+            inputs=[api_key, terminal_question, model, max_turns, corpus_source, uploaded_zip, runtime_state],
+            outputs=[terminal, status_text, runtime_state],
         )
         stop_btn.click(fn=None, cancels=[run_event])
 

pi_wrapper.py

@@ -93,6 +93,8 @@ class NativePiClient:
         provider: str,
         model: str,
         tools: str,
+        session_dir: Optional[Path] = None,
+        session_path: Optional[Path] = None,
         api_key: str = "",
     ) -> None:
         self.package_dir = package_dir
@@ -101,6 +103,8 @@ class NativePiClient:
         self.provider = provider
         self.model = model
         self.tools = tools
+        self.session_dir = session_dir
+        self.session_path = session_path
         self.api_key = api_key
         self.proc: Optional[subprocess.Popen[bytes]] = None
         self.stderr_chunks: List[str] = []
@@ -132,7 +136,10 @@ class NativePiClient:
             cmd.extend(["--model", self.model])
         if self.tools:
             cmd.extend(["--tools", self.tools])
-        cmd.append("--no-session")
+        if self.session_dir:
+            cmd.extend(["--session-dir", str(self.session_dir)])
+        if self.session_path:
+            cmd.extend(["--session", str(self.session_path)])
         return cmd
 
     def start(self) -> None:
@@ -185,6 +192,17 @@ class NativePiClient:
         self.proc.stdin.write(line.encode("utf-8"))
         self.proc.stdin.flush()
 
+    def call(self, command_type: str, **payload: Any) -> Dict[str, Any]:
+        request_id = self._next_id()
+        message = {"id": request_id, "type": command_type, **payload}
+        self._send(message)
+        while True:
+            event = self._read_json_line()
+            if event.get("type") == "response" and event.get("id") == request_id:
+                if not event.get("success", False):
+                    raise RuntimeError(f"RPC {command_type} failed: {event.get('error', 'unknown error')}")
+                return event
+
     def _read_json_line(self) -> Dict[str, Any]:
         if self.proc is None or self.proc.stdout is None:
             raise RuntimeError("RPC client is not running")
@@ -624,6 +642,8 @@ def run_pi(
     provider: str = "openai",
     model: str = "gpt-4o",
     max_turns: int = 6,
+    session_dir: Optional[Path] = None,
+    session_path: Optional[Path] = None,
 ) -> Generator[Dict[str, Any], None, str]:
     """
     Run PI (native if available, else fallback) and yield events.
@@ -637,6 +657,8 @@ def run_pi(
             provider=provider,
             model=model,
             tools="read,bash",
+            session_dir=session_dir,
+            session_path=session_path,
             api_key=api_key,
         )
         try:
@@ -666,6 +688,8 @@ def run_pi_stream(
     provider: str = "openai",
     model: str = "gpt-4o",
     max_turns: int = 6,
+    session_dir: Optional[Path] = None,
+    session_path: Optional[Path] = None,
 ) -> Iterator[Dict[str, Any]]:
     """
     Iterator that yields PI events.  The *last* item is a sentinel dict
@@ -679,11 +703,20 @@ def run_pi_stream(
             provider=provider,
             model=model,
             tools="read,bash",
+            session_dir=session_dir,
+            session_path=session_path,
             api_key=api_key,
         )
         text_parts: List[str] = []
         try:
             client.start()
+            state_response = client.call("get_state")
+            state_data = state_response.get("data", {}) or {}
+            yield {
+                "type": "__session__",
+                "sessionFile": state_data.get("sessionFile"),
+                "sessionId": state_data.get("sessionId"),
+            }
             request_id = client._next_id()
             client._send({"id": request_id, "type": "prompt", "message": question})