|
14 | 14 | GET /api/v2/prd/{id}/diff - Diff two versions |
15 | 15 | """ |
16 | 16 |
|
| 17 | +import json |
17 | 18 | import logging |
18 | | -from typing import Optional |
| 19 | +import os |
| 20 | +from typing import AsyncGenerator, Optional |
19 | 21 |
|
20 | 22 | from fastapi import APIRouter, Depends, HTTPException, Query, Request |
| 23 | +from fastapi.responses import StreamingResponse |
21 | 24 | from pydantic import BaseModel, Field |
22 | 25 |
|
23 | 26 | from codeframe.core.workspace import Workspace |
@@ -186,6 +189,122 @@ async def get_latest_prd( |
186 | 189 | return _prd_to_response(record) |
187 | 190 |
|
188 | 191 |
|
| 192 | +def _sse(event: dict) -> str: |
| 193 | + """Format a stress-test event dict as an SSE ``data:`` frame.""" |
| 194 | + return f"data: {json.dumps(event)}\n\n" |
| 195 | + |
| 196 | + |
| 197 | +async def _stress_test_event_stream( |
| 198 | + workspace: Workspace, |
| 199 | + max_depth: int, |
| 200 | + request: Optional[Request] = None, |
| 201 | +) -> AsyncGenerator[str, None]: |
| 202 | + """Yield SSE frames for a PRD stress-test. |
| 203 | +
|
| 204 | + Recoverable problems (missing PRD, missing ``ANTHROPIC_API_KEY``) are |
| 205 | + surfaced as in-stream ``error`` events rather than HTTP errors, so a |
| 206 | + browser ``EventSource`` can display them via its message handler. |
| 207 | +
|
| 208 | + Stops early if the client disconnects, so an abandoned stream does not keep |
| 209 | + issuing LLM calls — mirroring ``event_stream_generator`` in streaming_v2. |
| 210 | + """ |
| 211 | + from codeframe.core.prd_stress_test import stress_test_prd_stream |
| 212 | + |
| 213 | + record = prd.get_latest(workspace) |
| 214 | + if not record: |
| 215 | + yield _sse({ |
| 216 | + "type": "error", |
| 217 | + "message": "No PRD found. Add or generate a PRD first.", |
| 218 | + }) |
| 219 | + return |
| 220 | + |
| 221 | + # Resolve the LLM provider following the documented chain: |
| 222 | + # env var → workspace config (.codeframe/config.yaml) → default "anthropic". |
| 223 | + # (No CLI flag here — this is the web surface.) Mirrors runtime.py. |
| 224 | + from codeframe.adapters.llm import get_provider |
| 225 | + from codeframe.core.config import load_environment_config |
| 226 | + |
| 227 | + env_cfg = load_environment_config(workspace.repo_path) |
| 228 | + llm_cfg = env_cfg.llm if (env_cfg and env_cfg.llm) else None |
| 229 | + provider_type = ( |
| 230 | + os.getenv("CODEFRAME_LLM_PROVIDER") |
| 231 | + or (llm_cfg.provider if llm_cfg else None) |
| 232 | + or "anthropic" |
| 233 | + ) |
| 234 | + |
| 235 | + # Only the Anthropic provider needs an API key up front; local providers |
| 236 | + # (ollama/vllm/compatible) do not. |
| 237 | + if provider_type == "anthropic" and not os.getenv("ANTHROPIC_API_KEY"): |
| 238 | + yield _sse({ |
| 239 | + "type": "error", |
| 240 | + "message": "ANTHROPIC_API_KEY environment variable required.", |
| 241 | + }) |
| 242 | + return |
| 243 | + |
| 244 | + provider_kwargs: dict = {} |
| 245 | + model_override = os.getenv("CODEFRAME_LLM_MODEL") or ( |
| 246 | + llm_cfg.model if llm_cfg else None |
| 247 | + ) |
| 248 | + base_url_override = (llm_cfg.base_url if llm_cfg else None) or os.getenv( |
| 249 | + "OPENAI_BASE_URL" |
| 250 | + ) |
| 251 | + if model_override: |
| 252 | + provider_kwargs["model"] = model_override |
| 253 | + if base_url_override: |
| 254 | + provider_kwargs["base_url"] = base_url_override |
| 255 | + |
| 256 | + try: |
| 257 | + provider = get_provider(provider_type, **provider_kwargs) |
| 258 | + except ValueError as exc: |
| 259 | + yield _sse({"type": "error", "message": str(exc)}) |
| 260 | + return |
| 261 | + |
| 262 | + async for event in stress_test_prd_stream( |
| 263 | + record.content, provider, max_depth=max_depth, |
| 264 | + ): |
| 265 | + # If the browser has gone away, stop iterating the core generator so its |
| 266 | + # next (blocking, billable) LLM call is never made. |
| 267 | + if request is not None and await request.is_disconnected(): |
| 268 | + logger.info("Client disconnected from stress-test stream; aborting") |
| 269 | + break |
| 270 | + yield _sse(event) |
| 271 | + |
| 272 | + |
| 273 | +@router.get("/stress-test") |
| 274 | +@rate_limit_standard() |
| 275 | +async def stress_test_prd_stream_endpoint( |
| 276 | + request: Request, |
| 277 | + max_depth: int = Query(3, ge=1, le=10, description="Maximum recursion depth"), |
| 278 | + workspace: Workspace = Depends(get_v2_workspace), |
| 279 | +) -> StreamingResponse: |
| 280 | + """Stream a PRD stress-test (recursive decomposition) via SSE. |
| 281 | +
|
| 282 | + Runs the headless ``stress_test_prd_stream`` core generator over the |
| 283 | + latest PRD and emits its progress events as Server-Sent Events. This is |
| 284 | + the web equivalent of ``cf prd stress-test``. |
| 285 | +
|
| 286 | + Declared as GET (not POST) so it is reachable from a browser |
| 287 | + ``EventSource``, matching ``GET /api/v2/tasks/{task_id}/stream``. No custom |
| 288 | + auth headers are required (cookie-based auth via ``withCredentials``). |
| 289 | +
|
| 290 | + Event payloads (JSON in the SSE ``data:`` field, ``type`` field): |
| 291 | + - ``goals_extracted``: high-level goals parsed from the PRD |
| 292 | + - ``goal_analyzed``: one per top-level goal (classification + running |
| 293 | + ambiguity count) |
| 294 | + - ``complete``: ambiguity count + rendered tech spec / ambiguity report |
| 295 | + - ``error``: no PRD, missing API key, or decomposition failure |
| 296 | + """ |
| 297 | + return StreamingResponse( |
| 298 | + _stress_test_event_stream(workspace, max_depth, request), |
| 299 | + media_type="text/event-stream", |
| 300 | + headers={ |
| 301 | + "Cache-Control": "no-cache", |
| 302 | + "Connection": "keep-alive", |
| 303 | + "X-Accel-Buffering": "no", |
| 304 | + }, |
| 305 | + ) |
| 306 | + |
| 307 | + |
189 | 308 | @router.get("/{prd_id}", response_model=PrdResponse) |
190 | 309 | @rate_limit_standard() |
191 | 310 | async def get_prd( |
|
0 commit comments