Create graph.py

graph.py

@@ -0,0 +1,495 @@
+from __future__ import annotations
+
+import json
+from typing import Any, Dict, List, TypedDict
+
+from langgraph.graph import END, StateGraph
+
+from agents import (
+    build_code_analysis_agent,
+    build_feedback_agent,
+    build_spec_agent,
+    build_test_generator_agent,
+    build_test_plan_agent,
+)
+from llms import build_llm
+from schemas import (
+    CodeAnalysis,
+    FeedbackSignal,
+    FinalReport,
+    Spec,
+    StudentTestSuite,
+    TestCaseList,
+    TestCase,
+    TestPlan,
+)
+
+
+class GraphState(TypedDict):
+    problem: str
+    description: str
+    constraints: str
+    code: str
+    language: str
+    per_category: int
+    student_count: int
+    iteration: int
+    spec: Spec
+    analysis: CodeAnalysis
+    plan: TestPlan
+    suites: List[StudentTestSuite]
+    feedback: FeedbackSignal
+    issues: List[str]
+
+
+def _category_targets(per_category: int) -> Dict[str, int]:
+    categories = [
+        "Basic cases",
+        "Boundary cases",
+        "Random cases",
+        "Stress cases",
+        "Invalid/robustness cases",
+        "Bug-targeted cases",
+    ]
+    return {category: per_category for category in categories}
+
+
+def _normalize_category(label: str) -> str:
+    lower = label.strip().lower()
+    if "basic" in lower:
+        return "Basic cases"
+    if "boundary" in lower or "edge" in lower:
+        return "Boundary cases"
+    if "random" in lower:
+        return "Random cases"
+    if "stress" in lower:
+        return "Stress cases"
+    if "invalid" in lower or "robust" in lower:
+        return "Invalid/robustness cases"
+    if "bug" in lower:
+        return "Bug-targeted cases"
+    return label
+
+
+def _enforce_targets(
+    cases: List[TestCase], targets: Dict[str, int]
+) -> tuple[List[TestCase], Dict[str, int]]:
+    by_category: Dict[str, List[TestCase]] = {category: [] for category in targets}
+    for case in cases:
+        normalized = _normalize_category(case.category)
+        case.category = normalized
+        if normalized in by_category:
+            by_category[normalized].append(case)
+
+    enforced: List[TestCase] = []
+    missing: Dict[str, int] = {}
+    for category, count in targets.items():
+        selected = by_category.get(category, [])[:count]
+        enforced.extend(selected)
+        remaining = count - len(selected)
+        if remaining > 0:
+            missing[category] = remaining
+    return enforced, missing
+
+
+def _strip_markdown(text: str) -> str:
+    stripped = text.strip()
+    if stripped.startswith("```") and stripped.endswith("```"):
+        lines = stripped.splitlines()
+        if len(lines) >= 2:
+            return "\n".join(lines[1:-1]).strip()
+    return stripped
+
+
+def _extract_json_blob(text: str) -> str:
+    start_obj = text.find("{")
+    end_obj = text.rfind("}")
+    if start_obj != -1 and end_obj != -1 and end_obj > start_obj:
+        return text[start_obj : end_obj + 1]
+    start_list = text.find("[")
+    end_list = text.rfind("]")
+    if start_list != -1 and end_list != -1 and end_list > start_list:
+        return text[start_list : end_list + 1]
+    return text
+
+
+def _scan_string(text: str, start: int) -> int:
+    index = start + 1
+    escaped = False
+    while index < len(text):
+        char = text[index]
+        if escaped:
+            escaped = False
+        elif char == "\\":
+            escaped = True
+        elif char == '"':
+            return index + 1
+        index += 1
+    return len(text)
+
+
+def _find_expr_end(text: str, start: int) -> int:
+    index = start
+    in_string = False
+    escaped = False
+    while index < len(text):
+        char = text[index]
+        if in_string:
+            if escaped:
+                escaped = False
+            elif char == "\\":
+                escaped = True
+            elif char == '"':
+                in_string = False
+        else:
+            if char == '"':
+                in_string = True
+            elif char in {",", "}", "]"}:
+                return index
+        index += 1
+    return len(text)
+
+
+def _tokenize_expr(expr: str) -> List[tuple[str, Any]] | None:
+    tokens: List[tuple[str, Any]] = []
+    index = 0
+    while index < len(expr):
+        char = expr[index]
+        if char.isspace():
+            index += 1
+            continue
+        if char == '"':
+            end = _scan_string(expr, index)
+            literal = expr[index:end]
+            try:
+                value = json.loads(literal)
+            except json.JSONDecodeError:
+                return None
+            tokens.append(("str", value))
+            index = end
+            continue
+        if char.isdigit():
+            end = index
+            while end < len(expr) and expr[end].isdigit():
+                end += 1
+            tokens.append(("int", int(expr[index:end])))
+            index = end
+            continue
+        if char in {"+", "*"}:
+            tokens.append(("op", char))
+            index += 1
+            continue
+        return None
+    return tokens
+
+
+def _eval_string_expression(expr: str) -> str | None:
+    tokens = _tokenize_expr(expr)
+    if not tokens:
+        return None
+
+    has_string = any(token[0] == "str" for token in tokens)
+    if not has_string:
+        return None
+
+    def parse_term(pos: int) -> tuple[str | None, int]:
+        if pos >= len(tokens):
+            return None, pos
+        if tokens[pos][0] == "str":
+            value = tokens[pos][1]
+        elif tokens[pos][0] == "int":
+            value = str(tokens[pos][1])
+        else:
+            return None, pos
+        pos += 1
+        while pos + 1 < len(tokens) and tokens[pos] == ("op", "*"):
+            if tokens[pos + 1][0] != "int":
+                return None, pos
+            repeat = tokens[pos + 1][1]
+            value = value * repeat
+            pos += 2
+        return value, pos
+
+    result, pos = parse_term(0)
+    if result is None:
+        return None
+    while pos < len(tokens):
+        if tokens[pos] != ("op", "+"):
+            return None
+        term, pos = parse_term(pos + 1)
+        if term is None:
+            return None
+        result += term
+    return result
+
+
+def _cap_string(value: str, limit: int = 200) -> str:
+    if len(value) <= limit:
+        return value
+    return value[:limit]
+
+
+def _rewrite_repeat_calls(text: str) -> str:
+    output: List[str] = []
+    index = 0
+    while index < len(text):
+        char = text[index]
+        if char == '"':
+            start = index
+            end = _scan_string(text, index)
+            output.append(text[start:end])
+            probe = end
+            while probe < len(text) and text[probe].isspace():
+                probe += 1
+            if text.startswith(".repeat", probe):
+                cursor = probe + len(".repeat")
+                while cursor < len(text) and text[cursor].isspace():
+                    cursor += 1
+                if cursor < len(text) and text[cursor] == "(":
+                    cursor += 1
+                    while cursor < len(text) and text[cursor].isspace():
+                        cursor += 1
+                    number_start = cursor
+                    while cursor < len(text) and text[cursor].isdigit():
+                        cursor += 1
+                    number = text[number_start:cursor]
+                    while cursor < len(text) and text[cursor].isspace():
+                        cursor += 1
+                    if number and cursor < len(text) and text[cursor] == ")":
+                        output.append(f" * {number}")
+                        index = cursor + 1
+                        continue
+            index = end
+            continue
+        output.append(char)
+        index += 1
+    return "".join(output)
+
+
+def _replace_string_expressions(text: str) -> str:
+    output: List[str] = []
+    index = 0
+    while index < len(text):
+        char = text[index]
+        if char == '"':
+            start = index
+            end = _scan_string(text, index)
+            probe = end
+            while probe < len(text) and text[probe].isspace():
+                probe += 1
+            if probe < len(text) and text[probe] in {"+", "*"}:
+                expr_end = _find_expr_end(text, start)
+                expr_text = text[start:expr_end]
+                evaluated = _eval_string_expression(expr_text)
+                if evaluated is not None:
+                    output.append(json.dumps(_cap_string(evaluated)))
+                    index = expr_end
+                    continue
+            output.append(text[start:end])
+            index = end
+            continue
+        if char.isdigit():
+            start = index
+            end = index
+            while end < len(text) and text[end].isdigit():
+                end += 1
+            probe = end
+            while probe < len(text) and text[probe].isspace():
+                probe += 1
+            if probe < len(text) and text[probe] == "*":
+                expr_end = _find_expr_end(text, start)
+                expr_text = text[start:expr_end]
+                evaluated = _eval_string_expression(expr_text)
+                if evaluated is not None:
+                    output.append(json.dumps(_cap_string(evaluated)))
+                    index = expr_end
+                    continue
+        output.append(char)
+        index += 1
+    return "".join(output)
+
+
+def _parse_case_list(raw_text: str) -> TestCaseList:
+    cleaned = _strip_markdown(raw_text)
+    rewritten = _rewrite_repeat_calls(cleaned)
+    repaired = _replace_string_expressions(rewritten)
+    blob = _extract_json_blob(repaired)
+    try:
+        data = json.loads(blob)
+        if isinstance(data, list):
+            data = {"cases": data}
+        return TestCaseList.model_validate(data)
+    except json.JSONDecodeError:
+        return TestCaseList(cases=[])
+
+
+def node_spec(state: GraphState) -> Dict[str, Any]:
+    llm = build_llm("gemini-3-flash-preview", temperature=0.2)
+    prompt, parser = build_spec_agent(llm)
+    chain = prompt | llm | parser
+    spec = chain.invoke(
+        {
+            "problem": state["problem"],
+            "description": state["description"],
+            "constraints": state["constraints"],
+            "language": state["language"],
+            "format_instructions": parser.get_format_instructions(),
+        }
+    )
+    return {"spec": spec}
+
+
+def node_analysis(state: GraphState) -> Dict[str, Any]:
+    if not state["code"].strip():
+        return {"analysis": CodeAnalysis()}
+    llm = build_llm("gemini-2.5-flash", temperature=0.2)
+    prompt, parser = build_code_analysis_agent(llm)
+    chain = prompt | llm | parser
+    analysis = chain.invoke(
+        {
+            "code": state["code"],
+            "language": state["language"],
+            "format_instructions": parser.get_format_instructions(),
+        }
+    )
+    return {"analysis": analysis}
+
+
+def node_start(state: GraphState) -> Dict[str, Any]:
+    return {"iteration": 0}
+
+
+def node_plan(state: GraphState) -> Dict[str, Any]:
+    llm = build_llm("gemini-3.1-flash-lite-preview", temperature=0.3)
+    prompt, parser = build_test_plan_agent(llm)
+    chain = prompt | llm | parser
+    per_category = max(2, min(3, state["per_category"]))
+    plan = chain.invoke(
+        {
+            "spec": state["spec"].model_dump(),
+            "analysis": state["analysis"].model_dump(),
+            "issues": state.get("issues", []),
+            "per_category": per_category,
+            "format_instructions": parser.get_format_instructions(),
+        }
+    )
+    plan.targets = _category_targets(per_category)
+    plan.categories = list(plan.targets.keys())
+    return {"plan": plan}
+
+
+def node_generate(state: GraphState) -> Dict[str, Any]:
+    llm = build_llm("gemini-2.5-flash-lite", temperature=0.5)
+    prompt, parser = build_test_generator_agent(llm)
+    chain = prompt | llm
+    suites: List[StudentTestSuite] = []
+    issues: List[str] = []
+    for student_id in range(1, state["student_count"] + 1):
+        response = chain.invoke(
+            {
+                "spec": state["spec"].model_dump(),
+                "plan": state["plan"].model_dump(),
+                "student_id": student_id,
+                "format_instructions": parser.get_format_instructions(),
+            }
+        )
+        raw_text = response.content if hasattr(response, "content") else str(response)
+        case_list = _parse_case_list(raw_text)
+        if not case_list.cases:
+            issues.append(f"Student {student_id} output parsing failed")
+            continue
+        enforced, missing = _enforce_targets(case_list.cases, state["plan"].targets)
+        suites.append(StudentTestSuite(student_id=student_id, cases=enforced))
+        if missing:
+            issues.append(
+                f"Student {student_id} missing categories: {sorted(missing.keys())}"
+            )
+    return {"suites": suites, "issues": issues}
+
+
+def node_feedback(state: GraphState) -> Dict[str, Any]:
+    llm = build_llm("gemini-3-flash-preview", temperature=0.2)
+    prompt, parser = build_feedback_agent(llm)
+    chain = prompt | llm | parser
+    issues = state.get("issues", [])
+    feedback = chain.invoke(
+        {
+            "spec": state["spec"].model_dump(),
+            "plan": state["plan"].model_dump(),
+            "issues": issues,
+            "format_instructions": parser.get_format_instructions(),
+        }
+    )
+    needs_refine = feedback.needs_refine or bool(issues)
+    iteration = state.get("iteration", 0) + (1 if needs_refine else 0)
+    return {"feedback": feedback, "iteration": iteration}
+
+
+def should_refine(state: GraphState) -> str:
+    max_refines = 1
+    if state.get("iteration", 0) > max_refines:
+        return "final"
+    if state.get("issues"):
+        return "refine"
+    return "refine" if state["feedback"].needs_refine else "final"
+
+
+def build_graph():
+    graph = StateGraph(GraphState)
+    graph.add_node("start", node_start)
+    graph.add_node("spec", node_spec)
+    graph.add_node("analysis", node_analysis)
+    graph.add_node("plan", node_plan)
+    graph.add_node("generate", node_generate)
+    graph.add_node("feedback", node_feedback)
+
+    graph.set_entry_point("start")
+    graph.add_edge("start", "spec")
+    graph.add_edge("start", "analysis")
+    graph.add_edge("spec", "plan")
+    graph.add_edge("analysis", "plan")
+    graph.add_edge("plan", "generate")
+    graph.add_edge("generate", "feedback")
+    graph.add_conditional_edges(
+        "feedback",
+        should_refine,
+        {
+            "refine": "plan",
+            "final": END,
+        },
+    )
+
+    return graph.compile()
+
+
+def run_pipeline(
+    *,
+    problem: str,
+    description: str,
+    constraints: str,
+    code: str,
+    language: str,
+    student_count: int,
+    per_category: int,
+    issues: List[str] | None = None,
+) -> FinalReport:
+    app = build_graph()
+    state = app.invoke(
+        {
+            "problem": problem,
+            "description": description,
+            "constraints": constraints,
+            "code": code,
+            "language": language,
+            "student_count": student_count,
+            "per_category": per_category,
+            "issues": issues or [],
+        }
+    )
+    return FinalReport(
+        spec=state["spec"],
+        analysis=state["analysis"],
+        plan=state["plan"],
+        suites=state["suites"],
+        feedback=state["feedback"],
+    )