""" RUBRA — utils.py Helper functions: queue system, rate limit retry, zip creation, code parsing. """ import re import io import os import time import asyncio import logging import zipfile import hashlib from typing import List, Dict, Optional, AsyncIterator from collections import deque from dataclasses import dataclass, field from datetime import datetime log = logging.getLogger("rubra.utils") # ═══════════════════════════════════════════════════════ # REQUEST QUEUE — heavy requests don't block others # ═══════════════════════════════════════════════════════ @dataclass class QueuedRequest: id: str prompt: str framework: str created_at: float = field(default_factory=time.time) status: str = "pending" # pending | processing | done | error result: dict = field(default_factory=dict) class RequestQueue: """Simple async queue for heavy generation requests.""" def __init__(self, max_concurrent: int = 2): self.queue = deque() self.processing = {} self.results = {} self.max_concurrent = max_concurrent self._lock = asyncio.Lock() async def enqueue(self, prompt: str, framework: str) -> str: req_id = hashlib.md5(f"{prompt}{time.time()}".encode()).hexdigest()[:8] req = QueuedRequest(id=req_id, prompt=prompt, framework=framework) async with self._lock: self.queue.append(req) log.info(f"[QUEUE] Enqueued {req_id} (depth={len(self.queue)})") return req_id async def get_status(self, req_id: str) -> dict: if req_id in self.results: return {"status": "done", "result": self.results[req_id]} if req_id in self.processing: return {"status": "processing"} for r in self.queue: if r.id == req_id: pos = list(self.queue).index(r) return {"status": "pending", "position": pos + 1} return {"status": "not_found"} async def next(self) -> Optional[QueuedRequest]: async with self._lock: if self.queue and len(self.processing) < self.max_concurrent: req = self.queue.popleft() self.processing[req.id] = req return req return None async def complete(self, req_id: str, result: dict): async with self._lock: self.processing.pop(req_id, None) self.results[req_id] = result # Keep last 100 results if len(self.results) > 100: oldest = list(self.results.keys())[0] del self.results[oldest] # Global queue instance request_queue = RequestQueue(max_concurrent=2) # ═══════════════════════════════════════════════════════ # RATE LIMIT HANDLER — smart retry with backoff # ═══════════════════════════════════════════════════════ class RateLimitError(Exception): pass async def with_retry( coro_factory, max_retries: int = 3, base_delay: float = 1.0, max_delay: float = 30.0, on_retry=None, ) -> any: """ Execute async function with exponential backoff retry. Handles rate limits gracefully — user never sees raw errors. """ last_err = None for attempt in range(1, max_retries + 1): try: return await coro_factory() except Exception as e: last_err = e err_str = str(e).lower() # Rate limit / quota → wait and retry is_rate = any(t in err_str for t in [ "429", "rate limit", "quota", "too many requests", "overloaded", "capacity", "retry after", ]) if is_rate and attempt < max_retries: # Parse retry-after if available retry_after = 0 match = re.search(r'retry.?after[:\s]+(\d+)', err_str) if match: retry_after = int(match.group(1)) delay = max(retry_after, min(base_delay * (2 ** (attempt - 1)), max_delay)) log.warning(f"[RETRY] Rate limited, waiting {delay:.1f}s (attempt {attempt}/{max_retries})") if on_retry: await on_retry(attempt, delay, str(e)) await asyncio.sleep(delay) continue # Non-rate-limit error → don't retry if not is_rate: raise e raise last_err # ═══════════════════════════════════════════════════════ # ZIP CREATOR # ═══════════════════════════════════════════════════════ def create_zip(files: List[Dict[str, str]]) -> bytes: """ Create ZIP from list of {path, content} dicts. Returns raw bytes for download. """ buf = io.BytesIO() with zipfile.ZipFile(buf, mode='w', compression=zipfile.ZIP_DEFLATED) as zf: for f in files: path = f.get("path", "file.txt") content = f.get("content", "") if isinstance(content, str): content = content.encode("utf-8") zf.writestr(path, content) buf.seek(0) return buf.read() # ═══════════════════════════════════════════════════════ # CODE PARSER — extract code blocks from LLM response # ═══════════════════════════════════════════════════════ def extract_code_blocks(text: str) -> List[Dict[str, str]]: """Extract all fenced code blocks from markdown text.""" pattern = r'```(\w+)?\n([\s\S]*?)```' blocks = [] for match in re.finditer(pattern, text): lang = match.group(1) or "text" content = match.group(2) blocks.append({"language": lang, "content": content}) return blocks def extract_file_blocks(text: str) -> List[Dict[str, str]]: """ Extract file blocks from LLM output. Looks for patterns like: ### src/App.jsx ```jsx ...code... ``` or: // FILE: src/App.jsx ``` ...code... ``` """ files = [] # Pattern 1: ### filename\n```lang\n...code...\n``` pattern1 = r'###\s+([^\n]+)\n```\w*\n([\s\S]*?)```' for m in re.finditer(pattern1, text): path = m.group(1).strip().strip('`').strip('"').strip("'") content = m.group(2) if path and '.' in path.split('/')[-1]: files.append({"path": path, "content": content}) # Pattern 2: // FILE: path or # FILE: path pattern2 = r'(?://|#)\s*FILE:\s*([^\n]+)\n```\w*\n([\s\S]*?)```' for m in re.finditer(pattern2, text): path = m.group(1).strip() content = m.group(2) files.append({"path": path, "content": content}) # Pattern 3: **`path`** or **path** pattern3 = r'\*\*`?([^*`\n]+\.\w+)`?\*\*\s*\n```\w*\n([\s\S]*?)```' for m in re.finditer(pattern3, text): path = m.group(1).strip() content = m.group(2) if '.' in path.split('/')[-1]: files.append({"path": path, "content": content}) # Deduplicate by path seen = set() dedup = [] for f in files: if f["path"] not in seen: seen.add(f["path"]) dedup.append(f) return dedup def clean_llm_code(code: str) -> str: """Remove markdown fences and extra whitespace from code.""" code = re.sub(r'^```\w*\n?', '', code.strip()) code = re.sub(r'\n?```$', '', code.strip()) return code.strip() # ═══════════════════════════════════════════════════════ # STREAMING HELPERS # ═══════════════════════════════════════════════════════ def sse_event(data: dict) -> str: """Format dict as SSE event.""" import json return f"data: {json.dumps(data, ensure_ascii=False)}\n\n" def sse_done() -> str: return "data: [DONE]\n\n" async def stream_text_chunks(text: str, chunk_size: int = 15, delay: float = 0.02) -> AsyncIterator[str]: """Stream text character by character for demo/fallback.""" for i in range(0, len(text), chunk_size): yield text[i:i+chunk_size] await asyncio.sleep(delay) # ═══════════════════════════════════════════════════════ # PROMPT BUILDERS # ═══════════════════════════════════════════════════════ def build_file_prompt( user_prompt: str, file_path: str, file_hint: str, framework: str, other_files: List[str] = None, ) -> str: """Build prompt for generating a single file in a project.""" other = f"\nOther files in project: {', '.join(other_files)}" if other_files else "" return f"""Generate the file `{file_path}` for a {framework} project. Project description: {user_prompt} File purpose: {file_hint}{other} Rules: - Write ONLY the file content, no explanation - Complete, production-ready code - No placeholders or TODOs - Proper imports and exports Output just the raw file content (no markdown fences):""" def build_project_prompt(user_prompt: str, framework: str, file_list: List[str]) -> str: """Build prompt for generating all project files at once.""" files_str = "\n".join(f"- {f}" for f in file_list) return f"""Generate a complete {framework} project based on this description: {user_prompt} Generate these files: {files_str} For each file, use this exact format: ### path/to/file.ext ```language [file content here] ``` Rules: - Write COMPLETE file contents — no truncation - Production-ready code, no placeholders - Follow {framework} best practices - Include all imports""" # ═══════════════════════════════════════════════════════ # LANGUAGE DETECTION FOR EXPLANATION # ═══════════════════════════════════════════════════════ def detect_code_language(code: str, hint: str = "") -> str: if hint: return hint.lower() if re.search(r'^\s*(def |import |from .+import|class .+:)', code, re.M): return "python" if re.search(r'(const |let |var |=>|async function|require\()', code): return "javascript" if re.search(r'(<[A-Z][a-zA-Z]+|jsx|tsx)', code): return "react" if re.search(r'(. That's a SOFT failure, and it was # silently being counted as "done". This section closes that gap: every # generated file gets checked against what it's actually supposed to look # like, and a soft failure rotates to a genuinely different model — not # just a second roll of the same dice — before giving up. _REFUSAL_MARKERS = ( "i cannot", "i can't", "i'm sorry", "as an ai", "i am unable", "i don't have the ability", "cannot provide", "i won't be able", "i'm not able to", ) def validate_file_content(path: str, content: str) -> "tuple[bool, str]": """ Dependency-free sanity check that a generated file is real and usable — not empty, not a refusal, not cut off mid-block, roughly shaped like its extension expects. NOT a full compiler — it's the floor that catches the actual failure mode reported in practice: empty/truncated/placeholder output getting reported to the user as "done" when it plainly wasn't. Returns (is_valid, reason_if_invalid). """ if not content or not content.strip(): return False, "empty content" stripped = content.strip() low = stripped.lower() if len(stripped) < 3: return False, "suspiciously short (<3 chars) — likely not real content" if any(marker in low[:300] for marker in _REFUSAL_MARKERS): return False, "looks like a refusal/apology, not actual file content" if stripped.startswith("```") or stripped.endswith("```"): return False, "still wrapped in markdown fences — not raw file content" ext = path.rsplit(".", 1)[-1].lower() if "." in path else "" if ext in ("html", "htm"): if "<" not in stripped or ">" not in stripped: return False, "doesn't look like HTML — no tags found" opens = len(re.findall(r'<(html|head|body|div|main|section|nav)\b', low)) closes = len(re.findall(r'', low)) if opens >= 3 and closes < opens - 2: # generous slack — this is not a real parser return False, f"structural tags look unbalanced ({opens} opened vs {closes} closed) — likely truncated" elif ext == "css": if stripped.count("{") == 0: return False, "no CSS rules found" if stripped.count("{") != stripped.count("}"): return False, "unbalanced braces — likely truncated" elif ext in ("js", "jsx", "ts", "tsx"): if stripped.count("{") != stripped.count("}"): return False, "unbalanced braces — likely truncated mid-block" if stripped.count("(") != stripped.count(")"): return False, "unbalanced parentheses — likely truncated mid-call" elif ext == "json": import json as _json try: _json.loads(stripped) except Exception as e: return False, f"invalid JSON: {str(e)[:80]}" elif ext == "py": try: compile(stripped, path or "", "exec") except SyntaxError as e: return False, f"Python syntax error: {str(e)[:80]}" return True, "" async def generate_diff_edit_with_rotation( path: str, current_content: str, instruction: str, framework: str, cascade: list, stream_llm_func, max_rounds: int = 2, max_tokens: int = 8192, on_attempt=None, ) -> dict: """ The 'edit, don't regenerate' engine. Asks the model for a SEARCH/REPLACE diff (or FULL_REWRITE for pervasive changes) instead of the whole file, applies it via merge_engine's 3-tier matcher, and validates the RESULT — same model-rotation safety net as generate_with_model_rotation(), plus one extra grounding step: a merge failure (ambiguous/not-found SEARCH block) re-prompts using merge_engine's auto-generated retry prompt, which carries the file's ACTUAL current content, instead of blindly repeating the same request that just failed. Returns: {"content","ok","reason","model","attempts","tiers"} """ from merge_engine import merge_diff def _initial_prompt() -> str: return ( f"Edit the file `{path}` ({framework} project).\n\n" f"=== CURRENT CONTENT ===\n{current_content}\n=== END CURRENT CONTENT ===\n\n" f"Requested change: \"{instruction}\"\n\n" f"Respond with ONE OR MORE SEARCH/REPLACE blocks in this exact format:\n" f"<<<<<<< SEARCH\n...exact text copied from CURRENT CONTENT above...\n" f"=======\n...replacement...\n>>>>>>> REPLACE\n\n" f"Use multiple blocks if the change touches multiple separate places. " f"Include at least 2-3 lines of context in each SEARCH block (not a " f"single short line) so it's uniquely identifiable in the file. " f"Copy the SEARCH text EXACTLY as it appears above — do not paraphrase, " f"re-indent, or reconstruct it from memory.\n\n" f"If the change is too pervasive for a patch (e.g. a full redesign), " f"reply instead with a single line `FULL_REWRITE` followed by the " f"complete new file content.\n\n" f"Output ONLY the diff/rewrite — no explanation, no markdown fences " f"wrapped around the whole thing." ) system_msg = {"role": "system", "content": f"Expert {framework} developer making a precise, minimal edit to an existing file."} prompt_text = _initial_prompt() attempts = 0 total = max_rounds * len(cascade) last_reason = "no attempts made" for _round in range(max_rounds): for cfg in cascade: attempts += 1 label = cfg.get("label", cfg.get("model", "?")) if on_attempt: try: await on_attempt(label, attempts, total) except Exception: pass messages = [system_msg, {"role": "user", "content": prompt_text}] diff_text = "" try: async for tok in stream_llm_func( messages, cfg["url"], cfg["key"], cfg["model"], cfg.get("temp", 0.1), max_tokens ): diff_text += tok except Exception as e: last_reason = f"{label} request failed: {str(e)[:120]}" continue diff_text = diff_text.strip() if diff_text.startswith("```"): diff_text = re.sub(r'^```\w*\n?', '', diff_text) diff_text = re.sub(r'\n?```$', '', diff_text) merge_result = merge_diff(path, current_content, diff_text) if not merge_result.success: last_reason = f"{label}: merge failed — {merge_result.error}" # Ground the NEXT attempt in the real failure + real file # content instead of repeating the same prompt blind. prompt_text = merge_result.retry_prompt continue ok, reason = validate_file_content(path, merge_result.content) if ok: return {"content": merge_result.content, "ok": True, "reason": "", "model": label, "attempts": attempts, "tiers": merge_result.match_tiers} last_reason = f"{label}: merged but failed content validation — {reason}" prompt_text = _initial_prompt() # reset to a clean ask, last one wasn't the diff's fault return {"content": current_content, "ok": False, "reason": last_reason, "model": "exhausted", "attempts": attempts, "tiers": []} async def generate_with_model_rotation( messages: list, cascade: list, stream_llm_func, path: str = "", max_rounds: int = 2, max_tokens: int = 8192, on_attempt=None, # optional async callback(model_label, attempt_no, total) ) -> dict: """ Try every model in `cascade`, in order, up to `max_rounds` full passes — stopping the instant validate_file_content() passes. A model that "succeeds" with garbage content gets skipped to the NEXT model, not retried on itself; the same broken prompt + same broken model usually just produces the same broken answer twice. `cascade` is a list of dicts: {"url","key","model","temp","label"} — the same shape as brain.CODING_MODEL_CASCADE. Returns: {"content","ok","reason","model","attempts"} """ best_content = "" best_reason = "no attempts made" attempts = 0 total = max_rounds * len(cascade) for _round in range(max_rounds): for cfg in cascade: attempts += 1 label = cfg.get("label", cfg.get("model", "?")) if on_attempt: try: await on_attempt(label, attempts, total) except Exception: pass content = "" try: async for tok in stream_llm_func( messages, cfg["url"], cfg["key"], cfg["model"], cfg.get("temp", 0.1), max_tokens ): content += tok except Exception as e: best_reason = f"{label} request failed: {str(e)[:100]}" continue content = clean_llm_code(content) ok, reason = validate_file_content(path, content) if ok: return {"content": content, "ok": True, "reason": "", "model": label, "attempts": attempts} if len(content) > len(best_content): best_content, best_reason = content, f"{label}: {reason}" # Every model, every round, failed validation — return the best partial # result found so the caller can still keep the old file rather than # nothing, but flag ok=False so it's reported honestly, not as "done". return {"content": best_content, "ok": False, "reason": best_reason, "model": "exhausted", "attempts": attempts}