build-small-hackathon/CodeFlow

1 2from huggingface_hub import hf_hub_download3from llama_cpp import Llama4import gradio as gr5from gradio import Server6from fastapi.responses import HTMLResponse, PlainTextResponse, FileResponse # frontend + traces + vendored static assets7import mimetypes8from typing import Any, cast # to resolve PyLance freaking out over llama-cpp-python in the generate_flowchart function9from textwrap import dedent10from pathlib import Path # load the custom frontend from disk11import re # remove thinking tag from response12import json, time, uuid # agent-trace logging13from datetime import datetime, timezone14 15# ----- Get Model ----- #16# Download Q4_K_M GGUF file from the repo17model_path = hf_hub_download(18    repo_id="build-small-hackathon/codeflow-qwen-3-finetuning", # fallback: unsloth/Qwen3-Coder-30B-A3B-Instruct-GGUF19    filename="qwen3-coder-codeflow-Q3_K_L.gguf" # fallback: Qwen3-Coder-30B-A3B-Instruct-UD-Q3_K_XL.gguf20)21 22# Initialize llama.cpp with the local cached path23llm = Llama(24    model_path=model_path,25    n_ctx=4096,26    n_threads=227)28 29# ----- Init App ----- #30app = gr.Server(title="Code-to-Flowchart Generator")31 32# ----- Agent traces ----- #33# Each generation appends one JSON line capturing the full LLM call (input code,34# the model's reasoning, output Mermaid + linemap, token usage, latency).35# Download the whole log from the running app at  /traces .36MODEL_NAME = "unsloth/Qwen3-Coder-30B-A3B-Instruct-GGUF:UD-Q3_K_XL"37TRACE_PATH = Path(__file__).parent / "agent_traces.jsonl"38 39def write_trace(record: dict) -> None:40    try:41        with open(TRACE_PATH, "a", encoding="utf-8") as f:42            f.write(json.dumps(record, ensure_ascii=False) + "\n")43    except Exception:44        pass  # tracing must never break generation45 46# ----- Functions ----- #47 48# This is a cleaning function to resolve common syntax errors.49def quote_labels(text: str) -> str:50    # Mermaid node labels can't hold raw code characters, so quote-wrap each label body51    # A label's real closing bracket is followed by a Mermaid connector, edge-label, pipe, statement end, or EOL52    # operators after a subscript (== < <= > >= != %) are never mistaken for a close.53    END = r'(?=\s*(?:[-<][-.>xo]|==[>=xo]|\||;|$))'54 55    def esc(body: str) -> str:56        return (body.replace('"', "'")57                    .replace('[', '&#91;').replace(']', '&#93;')58                    .replace('{', '&#123;').replace('}', '&#125;'))59 60    out = []61    for line in text.split('\n'):62        line = re.sub(r'(?<=\w)\[(.*?)\]' + END, lambda m: '["' + esc(m.group(1)) + '"]', line)63        line = re.sub(r'(?<=\w)\{(.*?)\}' + END, lambda m: '{"' + esc(m.group(1)) + '"}', line)64        out.append(line)65    return '\n'.join(out)66 67# Parse the model's <linemap> block into {nodeId: [startLine, endLine]}.68# Tolerant of junk lines; drops any entry whose line(s) fall outside the source.69def parse_linemap(block: str, num_lines: int) -> dict:70    out: dict = {}71    for raw in block.strip().splitlines():72        m = re.match(r'\s*([A-Za-z]\w*)\s*:\s*(\d+)(?:\s*-\s*(\d+))?\s*$', raw)73        if not m:74            continue75        a = int(m.group(2))76        b = int(m.group(3)) if m.group(3) else a77        if a > b:78            a, b = b, a79        if num_lines and 1 <= a <= num_lines and 1 <= b <= num_lines:80            out[m.group(1)] = [a, b]81    return out82 83@app.api(name="generate_flowchart")84def generate_flowchart(src_code: str) -> dict:85    # check if src_code is empty86    if not src_code.strip(): return {"mermaid": "", "linemap": {}}87 88    # Number the source lines so the model can cite them in the <linemap> block.89    src_lines = src_code.splitlines()90    num_lines = len(src_lines)91    numbered = "\n".join(f"{i}| {ln}" for i, ln in enumerate(src_lines, 1))92 93    # Set system prompt94    system_prompt = dedent("""95    ## Role/Persona96    You are a senior staff software architect and compiler engineer specializing in visual control-flow mapping. Your philosophy is pure utility: you translate raw execution logic into highly accurate, scannable, structural diagrams without any conversational filler, meta-commentary, or stylistic fluff.97 98    ## Context/Objective99    The user will provide source code files or logic snippets. Your sole objective is to parse the syntax and output a corresponding, valid Mermaid.js flowchart graph. This graph will be rendered natively in a production UI to help developers audit execution paths at a glance.100 101    ## Strict Constraints102    <constraints>103    1. OUTPUT FORMAT: Output valid, raw Mermaid.js syntax, immediately followed by the required <linemap> block (constraint 5). Nothing else.104    2. NO MARKDOWN FENCING: Do not wrap the output in ```mermaid or ``` blocks. Start directly with the Mermaid graph definition, for example: graph TD.105    3. NO PROSE: Do not include introductory text, explanations, or concluding remarks. If the code cannot be parsed, output an isolated error node.106    4. NODE NAMING: Paraphrase conditions into plain words — never put raw code, operators, quotes, parentheses, or square brackets/subscripts inside labels (write Index in bounds?, not i < len(nums); write Element is even?, not nums[i] % 2 == 0)107    5. SOURCE MAP: The user's code is prefixed with `N| ` line numbers (these are references, never copy the `N| ` prefix into a label). After the diagram, output a <linemap> block: one `NodeId: N` per node, where N is the 1-based source line that node represents (use `NodeId: start-end` for a multi-line construct). Omit purely structural Start/End nodes that correspond to no source line.108    </constraints>109 110    <banned_vocabulary>111    - Here is the flowchart112    - ```mermaid113    - ```114    - Note:115    - Explanation:116    - In this diagram117    - As requested118    </banned_vocabulary>119 120    ## Response Workflow121    Before outputting the final diagram syntax, perform structural parsing inside a hidden <thinking> tag according to these steps:122    1. Identify all conditional branches, including if/else, loops, including for/while, and termination points, including return/throw.123    2. Map out the execution flow nodes chronologically.124    3. Verify that every opening bracket and node label matching syntax, including [ ], ( ), and { }, is perfectly balanced and closed according to Mermaid specifications.125    4. Ensure no markdown formatting tags leak past the closing </thinking> tag.126 127    ## Few-Shot Examples128 129    Input:130    1| def check_status(val):131    2|     if val > 10:132    3|         return "Active"133    4|     else:134    5|         return "Inactive"135 136    Output:137    <thinking>138    1. Control structures: One conditional check, two return branches.139    2. Nodes: A Start, B Conditional, C Active return, D Inactive return.140    3. Source lines: def is line 1, the if is line 2, Active return is line 3, Inactive return is line 5.141    </thinking>142    graph TD143        A[Start: check_status] --> B{val > 10}144        B -- True --> C[Return 'Active']145        B -- False --> D[Return 'Inactive']146    <linemap>147    A: 1148    B: 2149    C: 3150    D: 5151    </linemap>152    """).strip()153 154    # Reset the cache per request so no cross-request bleeding155    llm.reset()156 157    # Casting else PyLance gets mad158    t0 = time.perf_counter()159    response = cast(Any, llm.create_chat_completion(160        messages=[161            {"role": "system", "content": system_prompt},162            {"role": "user", "content": numbered}163        ],164        temperature=0.1, # Keep it quite deterministic for now165        max_tokens=1024,166        stream=False167    ))168    latency_ms = round((time.perf_counter() - t0) * 1000)169 170    raw = response["choices"][0]["message"]["content"]171    usage = response.get("usage", {}) or {}172 173    # Capture the model's hidden reasoning for the trace, then strip the tags174    think = re.search(r'<thinking>(.*?)</thinking>', raw, flags=re.DOTALL)175    reasoning = think.group(1).strip() if think else ""176    content = re.sub(r'<thinking>.*?</thinking>', '', raw, flags=re.DOTALL)177 178    # Extract + strip the node→line map, then validate it against the source length179    linemap: dict = {}180    lm = re.search(r'<linemap>(.*?)</linemap>', content, flags=re.DOTALL)181    if lm:182        linemap = parse_linemap(lm.group(1), num_lines)183        content = content[:lm.start()] + content[lm.end():]184 185    # Quote-wrap each node label and escape any leaked code characters186    mermaid = quote_labels(content).strip() # and remove excess whitespace187 188    # ----- Agent trace (append-only JSONL; downloadable at /traces) -----189    write_trace({190        "id": uuid.uuid4().hex,191        "ts": datetime.now(timezone.utc).isoformat(),192        "event": "generate_flowchart",193        "model": MODEL_NAME,194        "params": {"temperature": 0.1, "max_tokens": 1024, "n_ctx": 4096},195        "input": {"src_code": src_code, "num_lines": num_lines},196        "reasoning": reasoning,197        "output": {"raw": raw, "mermaid": mermaid, "linemap": linemap},198        "usage": {199            "prompt_tokens": usage.get("prompt_tokens"),200            "completion_tokens": usage.get("completion_tokens"),201            "total_tokens": usage.get("total_tokens"),202        },203        "latency_ms": latency_ms,204        "status": "ok",205    })206 207    return {"mermaid": mermaid, "linemap": linemap}208 209# ----- Custom Frontend ----- #210# Served from frontend.html so the same file can be opened directly in a211# browser (file://) to preview the UI without loading the model.212index_html = (Path(__file__).parent / "frontend.html").read_text(encoding="utf-8")213 214# Load the custom HTML215# / takes precedent over default Blocks UI216@app.get("/")217def index():218    return HTMLResponse(index_html)219 220# Serve the vendored frontend assets (Mermaid, CodeMirror bundle, Gradio client,221# fonts) locally so the app needs NO external CDN/API at runtime.222STATIC_DIR = (Path(__file__).parent / "static").resolve()223mimetypes.add_type("text/javascript", ".js")224mimetypes.add_type("font/woff2", ".woff2")225 226@app.get("/static/{fname:path}")227def static_files(fname: str):228    fp = (STATIC_DIR / fname).resolve()229    # contain to STATIC_DIR (no path traversal) and require a real file230    if not str(fp).startswith(str(STATIC_DIR) + "/") or not fp.is_file():231        return PlainTextResponse("not found", status_code=404)232    mt, _ = mimetypes.guess_type(str(fp))233    return FileResponse(fp, media_type=mt or "application/octet-stream")234 235# Download every agent trace collected this run (one JSON object per line).236#   curl https://<your-space>/traces > agent_traces.jsonl237@app.get("/traces")238def traces():239    text = TRACE_PATH.read_text(encoding="utf-8") if TRACE_PATH.exists() else ""240    return PlainTextResponse(text, media_type="application/x-ndjson",241                             headers={"Content-Disposition": 'attachment; filename="agent_traces.jsonl"'})242 243app.launch(share=False)   # no external gradio.live tunnel — fully self-hosted