Slackbot AI follow up + tests

Author: haoruizhou

server/installer/sandbox/code_generator.py

@@ -23,7 +23,7 @@
 from langchain_anthropic import ChatAnthropic
 from langchain_chroma import Chroma
 from langchain_community.embeddings import FastEmbedEmbeddings
-from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langgraph.graph import END, StateGraph
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
@@ -406,6 +406,7 @@ def _format_response(state: AgentState) -> dict:
     retry_history = state.get("retry_history", [])
     return {
         "code": state.get("code", ""),
+        "rag_context": state.get("rag_context", ""),
         "result": {
             "status": "success" if result.get("ok") else "error",
             "output": result.get("std_out", "").strip(),
@@ -542,6 +543,199 @@ def generate_code_endpoint():
         }), 500
 
 
+@app.route("/api/generate-code-followup", methods=["POST"])
+def generate_code_followup_endpoint():
+    """Continue an AI conversation with prior context.
+
+    Body:
+        prompt: str          — the user's follow-up instruction
+        history: list[dict]  — prior turns, each with:
+            role: "user" | "assistant"
+            prompt: str      — the instruction or description
+            code: str        — generated code (assistant turns)
+            output: str      — execution stdout (assistant turns)
+            error: str       — execution error (assistant turns)
+            rag_context: str — RAG context from that turn (optional)
+    """
+    data = request.get_json() or {}
+    user_prompt = data.get("prompt", "").strip()
+    history = data.get("history", [])
+
+    if not user_prompt:
+        return jsonify({"error": "prompt is required"}), 400
+
+    try:
+        # ── Step 1: RAG retrieval on the new prompt ────────────────────────
+        index_data_dir()
+        rag_parts: list[str] = []
+
+        try:
+            sensor_docs = sensors_store.similarity_search(user_prompt, k=RAG_SENSOR_K)
+            if sensor_docs:
+                names = [d.page_content for d in sensor_docs]
+                rag_parts.append("RELEVANT SENSORS:\n" + "\n".join(f"  - {n}" for n in names))
+        except Exception:
+            logger.warning("Followup: sensor RAG failed", exc_info=True)
+
+        try:
+            run_docs = runs_store.similarity_search(user_prompt, k=RAG_RUN_K)
+            if run_docs:
+                descs = [d.page_content for d in run_docs]
+                rag_parts.append("RELEVANT RUNS:\n" + "\n".join(f"  - {d}" for d in descs))
+        except Exception:
+            logger.warning("Followup: runs RAG failed", exc_info=True)
+
+        try:
+            solution_docs = solutions_store.similarity_search(user_prompt, k=RAG_SOLUTION_K)
+            if solution_docs:
+                summaries = [d.page_content for d in solution_docs]
+                rag_parts.append("SUCCESSFUL EXAMPLES:\n" + "\n".join(f"  - {s}" for s in summaries))
+        except Exception:
+            logger.warning("Followup: solutions RAG failed", exc_info=True)
+
+        new_rag_context = "\n\n".join(rag_parts)
+
+        # ── Step 2: Build multi-turn message list ──────────────────────────
+        guide = _load_guide()
+        system_content = guide
+        # Merge RAG context: combine history's RAG with freshly retrieved
+        all_rag_parts: list[str] = []
+        for turn in history:
+            rc = turn.get("rag_context", "")
+            if rc:
+                all_rag_parts.append(rc)
+        if new_rag_context:
+            all_rag_parts.append(new_rag_context)
+        combined_rag = "\n\n".join(dict.fromkeys(all_rag_parts))  # deduplicate, preserve order
+        if combined_rag:
+            system_content += f"\n\n--- RETRIEVED CONTEXT ---\n{combined_rag}"
+
+        messages = [SystemMessage(content=system_content)]
+
+        for turn in history:
+            role = turn.get("role", "user")
+            if role == "user":
+                messages.append(HumanMessage(content=turn.get("prompt", "")))
+            elif role == "assistant":
+                code = turn.get("code", "")
+                output = turn.get("output", "")
+                error = turn.get("error", "")
+                parts = []
+                if code:
+                    parts.append(f"```python\n{code}\n```")
+                if output:
+                    parts.append(f"Output:\n{output}")
+                if error:
+                    parts.append(f"Error:\n{error}")
+                messages.append(AIMessage(content="\n\n".join(parts) or "(no output)"))
+
+        # The new follow-up instruction
+        messages.append(HumanMessage(content=user_prompt))
+
+        logger.info(
+            "Followup: %d history turns, %d total messages, prompt=%s",
+            len(history), len(messages), user_prompt[:80],
+        )
+
+        # ── Step 3: LLM call ───────────────────────────────────────────────
+        response = llm.invoke(messages)
+        code = _extract_python(response.content)
+
+        # ── Step 4: Execute in sandbox ─────────────────────────────────────
+        retry_history: list[dict] = []
+        retries = 0
+        final_result = {}
+        sandbox_ms = 0.0
+
+        while True:
+            try:
+                t0 = _time.perf_counter()
+                resp = requests.post(SANDBOX_URL, json={"code": code}, timeout=120)
+                sandbox_ms = (_time.perf_counter() - t0) * 1000
+                resp.raise_for_status()
+                final_result = resp.json()
+            except Exception as e:
+                err = str(e)
+                retries += 1
+                retry_history.append({"attempt": retries, "error": err})
+                if retries <= MAX_RETRIES:
+                    # Retry: ask LLM to fix
+                    messages.append(AIMessage(content=f"```python\n{code}\n```"))
+                    messages.append(HumanMessage(
+                        content=f"ATTEMPT {retries} FAILED:\n{err}\n\nFix the error above."
+                    ))
+                    response = llm.invoke(messages)
+                    code = _extract_python(response.content)
+                    continue
+                final_result = {}
+                break
+
+            if not final_result.get("ok"):
+                parts_err: list[str] = []
+                rc = final_result.get("return_code")
+                if rc is not None and rc != 0:
+                    parts_err.append(f"Exit code: {rc}")
+                if final_result.get("std_err"):
+                    parts_err.append(f"STDERR:\n{final_result['std_err'].strip()}")
+                if final_result.get("std_out"):
+                    parts_err.append(f"STDOUT:\n{final_result['std_out'].strip()}")
+                err = "\n".join(parts_err)
+                retries += 1
+                retry_history.append({"attempt": retries, "error": err})
+                if retries <= MAX_RETRIES:
+                    messages.append(AIMessage(content=f"```python\n{code}\n```"))
+                    messages.append(HumanMessage(
+                        content=f"ATTEMPT {retries} FAILED:\n{err}\n\nFix the error above."
+                    ))
+                    response = llm.invoke(messages)
+                    code = _extract_python(response.content)
+                    continue
+                break
+            else:
+                break
+
+        # ── Step 5: Format response ────────────────────────────────────────
+        state: AgentState = {
+            "prompt": user_prompt,
+            "rag_context": new_rag_context,
+            "code": code,
+            "sandbox_result": final_result,
+            "error": retry_history[-1]["error"] if retry_history and not final_result.get("ok") else "",
+            "retries": retries,
+            "retry_history": retry_history,
+            "llm_cache_hit": False,
+            "exec_cache_hit": False,
+        }
+        formatted = _format_response(state)
+
+        # Record metrics
+        ok = formatted["result"]["status"] == "success"
+        creator = data.get("creator", "") or ""
+        prompt_key = "followup:" + hashlib.sha256(user_prompt.encode()).hexdigest()
+        try:
+            record_metrics(
+                prompt_hash=prompt_key,
+                llm_cache_hit=False,
+                exec_cache_hit=False,
+                retry_count=retries,
+                sandbox_ms=sandbox_ms,
+                success=ok,
+                creator=creator,
+            )
+        except Exception:
+            logger.exception("Failed to record followup metrics")
+
+        return jsonify(formatted)
+
+    except Exception as e:
+        logger.exception("Followup generation failed for prompt: %s", user_prompt[:80])
+        return jsonify({
+            "error": str(e),
+            "rag_context": "",
+            "result": {"status": "error", "output": "", "error": str(e), "files": []},
+        }), 500
+
+
 @app.route("/api/metrics", methods=["GET"])
 def metrics_endpoint():
     """Return aggregated observability stats and a dashboard PNG."""

server/installer/sandbox/test_code_generator.py

@@ -158,6 +158,155 @@ def test_feedback_endpoint():
     return feedback_resp.status_code == 200
 
 
+def test_rag_context_in_response():
+    """Test that /api/generate-code now returns rag_context."""
+    print("Testing rag_context in generate-code response...")
+
+    prompt = "Print the numbers 1 through 5"
+    response = requests.post(
+        f"{CODE_GENERATOR_URL}/api/generate-code",
+        json={"prompt": prompt},
+        timeout=120,
+    )
+
+    if response.status_code != 200:
+        print(f"⚠️ Generate-code returned {response.status_code}")
+        return False
+
+    result = response.json()
+    has_rag_context = "rag_context" in result
+    print(f"  rag_context present: {has_rag_context}")
+    print(f"  rag_context type: {type(result.get('rag_context'))}")
+    if has_rag_context:
+        ctx = result["rag_context"]
+        print(f"  rag_context length: {len(ctx)} chars")
+        if ctx:
+            print(f"  rag_context preview: {ctx[:200]}...")
+    print()
+    return has_rag_context
+
+
+def test_followup_conversation():
+    """Test multi-turn follow-up conversation via /api/generate-code-followup."""
+    print("Testing follow-up conversation...")
+
+    # Step 1: Generate initial code
+    initial_prompt = "Print the numbers 1 through 5, one per line"
+    print(f"  Step 1 — Initial prompt: {initial_prompt}")
+
+    gen_resp = requests.post(
+        f"{CODE_GENERATOR_URL}/api/generate-code",
+        json={"prompt": initial_prompt},
+        timeout=120,
+    )
+    if gen_resp.status_code != 200:
+        print(f"  ⚠️ Initial generate-code failed ({gen_resp.status_code})")
+        return False
+
+    gen_result = gen_resp.json()
+    exec_result = gen_result.get("result", {})
+    if exec_result.get("status") != "success":
+        print(f"  ⚠️ Initial code execution failed: {exec_result.get('error', 'unknown')}")
+        return False
+
+    initial_code = gen_result.get("code", "")
+    initial_output = exec_result.get("output", "")
+    initial_rag = gen_result.get("rag_context", "")
+
+    print(f"  ✓ Initial code generated ({len(initial_code)} chars)")
+    print(f"  ✓ Initial output: {initial_output[:100]}")
+
+    # Step 2: Send a follow-up
+    history = [
+        {"role": "user", "prompt": initial_prompt, "code": "", "output": "", "error": "", "rag_context": ""},
+        {"role": "assistant", "prompt": "", "code": initial_code, "output": initial_output, "error": "", "rag_context": initial_rag},
+    ]
+
+    followup_prompt = "Now print them in reverse order (5 to 1)"
+    print(f"  Step 2 — Follow-up prompt: {followup_prompt}")
+
+    followup_resp = requests.post(
+        f"{CODE_GENERATOR_URL}/api/generate-code-followup",
+        json={"prompt": followup_prompt, "history": history},
+        timeout=120,
+    )
+
+    if followup_resp.status_code != 200:
+        print(f"  ⚠️ Follow-up failed ({followup_resp.status_code})")
+        return False
+
+    followup_result = followup_resp.json()
+    followup_exec = followup_result.get("result", {})
+
+    print(f"  Follow-up status: {followup_exec.get('status')}")
+    print(f"  Follow-up code:\n{'=' * 40}")
+    print(followup_result.get("code", "No code"))
+    print("=" * 40)
+    print(f"  Follow-up output: {followup_exec.get('output', 'no output')[:200]}")
+    print(f"  Follow-up rag_context present: {'rag_context' in followup_result}")
+
+    if followup_exec.get("error"):
+        print(f"  Follow-up error: {followup_exec['error'][:200]}")
+
+    retries = followup_result.get("retries", [])
+    if retries:
+        print(f"  Retries: {len(retries)}")
+
+    success = followup_exec.get("status") == "success"
+
+    # Step 3: Send a second follow-up (3-turn conversation)
+    if success:
+        history.append({"role": "user", "prompt": followup_prompt, "code": "", "output": "", "error": "", "rag_context": ""})
+        history.append({
+            "role": "assistant", "prompt": "",
+            "code": followup_result.get("code", ""),
+            "output": followup_exec.get("output", ""),
+            "error": "",
+            "rag_context": followup_result.get("rag_context", ""),
+        })
+
+        second_followup = "Now print them with their squares, like '1: 1', '2: 4', '3: 9' etc."
+        print(f"  Step 3 — Second follow-up: {second_followup}")
+
+        resp3 = requests.post(
+            f"{CODE_GENERATOR_URL}/api/generate-code-followup",
+            json={"prompt": second_followup, "history": history},
+            timeout=120,
+        )
+
+        if resp3.status_code == 200:
+            r3 = resp3.json()
+            status3 = r3.get("result", {}).get("status")
+            print(f"  3rd turn status: {status3}")
+            print(f"  3rd turn output: {r3.get('result', {}).get('output', '')[:200]}")
+            if status3 == "success":
+                print("  ✓ 3-turn conversation succeeded!")
+            else:
+                print(f"  ⚠️ 3rd turn failed: {r3.get('result', {}).get('error', '')[:200]}")
+        else:
+            print(f"  ⚠️ 3rd turn HTTP error: {resp3.status_code}")
+
+    print()
+    return success
+
+
+def test_followup_empty_prompt():
+    """Test that followup endpoint rejects empty prompts."""
+    print("Testing followup with empty prompt...")
+
+    resp = requests.post(
+        f"{CODE_GENERATOR_URL}/api/generate-code-followup",
+        json={"prompt": "", "history": []},
+        timeout=10,
+    )
+
+    success = resp.status_code == 400
+    print(f"  Status: {resp.status_code} (expected 400)")
+    print(f"  Response: {resp.json()}")
+    print()
+    return success
+
+
 def main():
     """Run all tests."""
     print("=" * 60)
@@ -169,6 +318,9 @@ def main():
         ("Simple Code Generation", test_simple_code_generation),
         ("Error Handling", test_error_with_retry),
         ("Feedback / Verified Solutions", test_feedback_endpoint),
+        ("RAG Context in Response", test_rag_context_in_response),
+        ("Follow-up Conversation", test_followup_conversation),
+        ("Follow-up Empty Prompt", test_followup_empty_prompt),
     ]
 
     results = []