code/20250108_Agent-Design-Thinking-System.py

import os, json, hashlib, ast
import gradio as gr
from openai import OpenAI
from datetime import datetime

# --- Configuration & Prompts (From Colleague's Script) ---
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
MODEL = "gpt-5.2"
course_username = os.environ["COURSE_USERNAME"]
password_to_name_dict = ast.literal_eval(os.environ["COURSE_PASSWORDS"])

SECTION_SEQUENCE = ["0","1.1","1.2","1.3","1.4","1.5","2.1","2.2","3.1","3.2","3.3","3.4","3.5","4.1"]


SECTIONS = {
    "0": {
        "title": "Summary",
        "student_prompt": "In one succinct paragraph, describe your idea: what it is, what problem it solves, what is innovative vs current approaches, and what impact it could have if fully realized.",
        "must_include": ["What the idea is", "What problem it solves", "Why it's innovative vs state of the art", "Impact if successful"],
    },
    "1.1": {
        "title": "Defining the problem: what problem does this solve?",
        "student_prompt": "Define the problem your idea is trying to solve. Include background/context and why it matters.",
        "must_include": ["Context", "What fails today", "Why it matters", "Scope boundaries"],
    },
    "1.2": {
        "title": "Why is this still not solved?",
        "student_prompt": "Explain why this problem is still not solved. Identify the key gaps or constraints in existing solutions (technical/logistical/economic/etc.).",
        "must_include": ["Existing approaches (brief)", "At least 2 reasons it's unsolved", "Avoid purely cost-only reasoning"],
    },
    "1.3": {
        "title": "Bottleneck analysis",
        "student_prompt": "Describe the chain of events needed to solve the problem, then identify the main bottleneck your idea addresses and why it is the priority bottleneck.",
        "must_include": ["Chain of events", "Bottleneck step", "Why it blocks the chain", "Which step your idea targets"],
    },
    "1.4": {
        "title": "Bar for success (metrics)",
        "student_prompt": "Define the quantitative bar for success for overcoming the bottleneck (key metrics and threshold values). Include technical constraints (and socio-economic if relevant).",
        "must_include": ["2–5 metrics", "Numeric thresholds", "Justification for thresholds", "At least one technical metric"],
    },
    "1.5": {
        "title": "Feedback on problem framing",
        "student_prompt": "List people/roles who could give you the best feedback on the problem and current bottlenecks, and what you’d ask each of them.",
        "must_include": ["5–12 people/roles", "Question for each", "Mix of perspectives"],
    },
    "2.1": {
        "title": "What is your idea (how it works)?",
        "student_prompt": "Explain your idea in detail, focusing on how it works step-by-step. Include inputs/outputs and any diagrams you’d make.",
        "must_include": ["Mechanism", "Inputs/outputs", "Workflow fit", "Sufficient detail to critique/build"],
    },
    "2.2": {
        "title": "Novelty and closest related ideas",
        "student_prompt": "How new is the idea? Identify the closest related ideas and explain similarities/differences and the novelty type (combination/extension/new application/etc.).",
        "must_include": ["2–5 related ideas", "Similarities", "Differences", "Novelty claim"],
    },
    "3.1": {
        "title": "Feasibility: dependencies in the chain",
        "student_prompt": "Where does your idea exist in the chain of events, and what dependencies does it have (upstream/downstream)? Does it add new burdens or links?",
        "must_include": ["Placement in chain", "Dependencies", "New constraints introduced"],
    },
    "3.2": {
        "title": "Feasibility: meets bars and overcomes bottleneck?",
        "student_prompt": "Do a quick sanity check: does your idea overcome the bottleneck and meet the quantitative bars? Use napkin-level logic or calculations.",
        "must_include": ["Tie to each bar", "Evidence/logic", "Conclusion + uncertainties"],
    },
    "3.3": {
        "title": "Contrarian analysis",
        "student_prompt": "Explain why the idea might not work. List and prioritize major assumptions/risks/bottlenecks, and briefly note mitigations.",
        "must_include": ["5–10 risks/assumptions", "Prioritization", "Failure mechanism", "Mitigation ideas (brief)"],
    },
    "3.4": {
        "title": "Go/no-go experiments",
        "student_prompt": "Design key go/no-go experiments to test the most important assumptions. Define what results would count as go vs no-go.",
        "must_include": ["2–5 experiments", "Go criteria", "No-go criteria", "Simple decisive design"],
    },
    "3.5": {
        "title": "Technical feedback people",
        "student_prompt": "List people/roles who could give you the best technical feedback, and what you’d ask them.",
        "must_include": ["5–12 people/roles", "Question for each", "Relevant technical coverage"],
    },
    "4.1": {
        "title": "Impact",
        "student_prompt": "If you are 100% successful, what would be the impact and who would be impacted? Include scale, time horizon, and broader implications.",
        "must_include": ["Stakeholders", "Nature of impact", "Scale/time horizon", "Broader implications"],
    },
}

# Add remaining section definitions from your colleague's code here...

SYSTEM_MESSAGE = """You are an educational writing and idea-framing coach for undergraduate and early graduate students.

You must help the student complete ONE section of a framing template at a time. You must follow the developer instructions.

Critical output requirement:
- You must respond with ONLY valid JSON that conforms exactly to the provided JSON Schema.
- Do not include any extra keys, commentary, markdown, or non-JSON text.

Behavioral requirements:
- Be constructive, specific, and pedagogically supportive.
- Maintain student ownership: help them clarify and improve, but do not fabricate facts, citations, data, or results.
- If the user requests disallowed assistance (e.g., instructions for wrongdoing, dangerous biological/chemical steps, or other harmful content), refuse and redirect to safe, high-level guidance consistent with an educational framing context.
""" # Full message from colleague
DEVELOPER_MESSAGE = """# Role
You are “Framing Coach,” a structured, rubric-driven assistant that guides a student through an idea-framing template one section at a time. You coach the student to produce clear, specific, and (when required) quantitative responses.

# Primary objective
For the CURRENT section only:
1) Evaluate the student's response against the section’s rubric (“must include” items + section-specific quality bars).
2) Provide targeted feedback and 1–3 follow-up questions.
3) Optionally propose an improved version that preserves the student’s intent and voice.
4) Decide whether the section is acceptable to advance.

# Scope and boundaries
- Focus strictly on the current section. Do not jump ahead to later sections unless the student explicitly asks and the current section is already accepted.
- Maintain student ownership: do not invent details the student did not provide (numbers, results, citations, experimental outcomes, claims of novelty, named experts, etc.).
- You may suggest example metrics, placeholder variables, or plausible ranges ONLY when clearly labeled as “assumptions/placeholders” and framed as options for the student to confirm or revise.
- Do not provide step-by-step instructions for wrongdoing or unsafe activity. If the student’s idea involves hazardous, illegal, or harmful actions, refuse and pivot to safe, high-level framing (problem definition, ethics, risk analysis) without operational details.

# Input contract (what the application provides)
You will receive, in the user message, a JSON object with:
- section_id: one of ["0","1.1","1.2","1.3","1.4","1.5","2.1","2.2","3.1","3.2","3.3","3.4","3.5","4.1"]
- section_title: string
- student_prompt: string (the main question to ask for this section)
- must_include: array of strings (rubric checklist items)
- student_message: string (the student’s latest attempt for this section; may be empty)
- draft_so_far (optional): object containing prior accepted sections; use only for context and consistency

If the user message is not valid JSON or lacks section_id/section_title/student_prompt/must_include/student_message:
- Set status="needs_input"
- In follow_up_questions[0], ask for the missing information in the simplest way.
- Do not guess the section_id.

# Section order (for next_section_id)
Use this fixed sequence:
["0","1.1","1.2","1.3","1.4","1.5","2.1","2.2","3.1","3.2","3.3","3.4","3.5","4.1"]

When a section is accepted, next_section_id is the next item in the sequence. If the current section is "4.1", next_section_id must be null.

# Output contract (MUST match the JSON Schema exactly)
Return a single JSON object with these required fields:
- section_id
- section_title
- status: "needs_input" | "needs_revision" | "accepted"
- feedback_bullets: array of strings
- missing_or_unclear: array of strings
- improved_version: string or null
- follow_up_questions: array of 1–3 strings
- advance: boolean
- next_section_id: string or null
- draft_update: object with
  - student_answer: string
  - accepted_version: string or null
  - coach_notes: string or null

# How to set status and advance
1) status="needs_input"
   Use when the student_message is empty, non-responsive, or only meta (e.g., “I don’t know,” “help me,” or off-topic).
   - advance=false
   - next_section_id=null
   - improved_version=null

2) status="needs_revision"
   Use when the student_message attempts the section but misses key rubric items or is unclear.
   - advance=false
   - next_section_id=null
   - improved_version should usually be provided (unless the student_message is too thin; then keep improved_version null and focus on questions).

3) status="accepted"
   Use when the response satisfies must_include and is sufficiently clear for downstream sections.
   - advance=true
   - next_section_id must follow the fixed sequence (or null at the end)
   - accepted_version must be a clean, student-faithful version of the section

# Rubric interpretation (general)
- must_include items are the minimum checklist; the student doesn’t need perfection, but they must address each item meaningfully.
- Prefer clarity over length; avoid jargon unless the student uses it correctly.
- Encourage specificity and testability.
- Be appropriately skeptical: flag hand-wavy claims and ask for grounding.

# Coaching style constraints
- Be direct, kind, and concrete.
- Provide feedback as actionable bullets (typically 3–6).
- Ask at most 3 follow-up questions; each should be targeted and non-overlapping.
- Avoid writing an entire proposal or adding substantial new content the student did not supply.

# Handling uncertainty and missing data
- If the student lacks numbers/metrics, you may suggest candidate metrics, placeholder variables, or plausible ranges ONLY if clearly labeled as assumptions/options to verify.

# Consistency with draft_so_far
If draft_so_far is provided:
- Maintain consistency in terminology.
- If you detect contradictions, flag them and ask a follow-up question.

# Safety and integrity
- If asked for harmful operational guidance, refuse briefly and redirect to safe alternatives.
- If the student asks you to “write it for me,” comply by coaching and offering outlines/edits, but do not generate a fully original submission without student input.

# coach_notes
Keep short. No chain-of-thought. No sensitive personal data.
"""


FRAMING_COACH_JSON_SCHEMA = {
    "type": "object",
    "additionalProperties": False,
    "properties": {
        "section_id": {"type": "string", "enum": SECTION_SEQUENCE},
        "section_title": {"type": "string"},
        "status": {"type": "string", "enum": ["needs_input", "needs_revision", "accepted"]},
        "feedback_bullets": {"type": "array", "items": {"type": "string"}},
        "missing_or_unclear": {"type": "array", "items": {"type": "string"}},
        "improved_version": {"type": ["string", "null"]},
        "follow_up_questions": {"type": "array", "minItems": 1, "maxItems": 3, "items": {"type": "string"}},
        "advance": {"type": "boolean"},
        "next_section_id": {"type": ["string", "null"], "enum": SECTION_SEQUENCE + [None]},
        "draft_update": {
            "type": "object",
            "additionalProperties": False,
            "properties": {
                "student_answer": {"type": "string"},
                "accepted_version": {"type": ["string", "null"]},
                "coach_notes": {"type": ["string", "null"]},
            },
            "required": ["student_answer", "accepted_version", "coach_notes"],
        },
    },
    "required": ["section_id", "section_title", "status", "feedback_bullets", "missing_or_unclear", "improved_version", "follow_up_questions", "advance", "next_section_id", "draft_update"],
}

# --- Logic Wrapper ---

class FramingCoachLogic:
    def __init__(self):
        self.model = MODEL

    def format_as_markdown(self, result):
        """Converts the JSON response into a beautiful Markdown string for the chat."""
        md = f"### Section {result['section_id']}: {result['section_title']}\n"
        md += f"**Status:** `{result['status'].upper()}`\n\n"
        
        if result["feedback_bullets"]:
            md += "#### 📝 Feedback\n"
            for b in result["feedback_bullets"]:
                md += f"* {b}\n"
        
        if result["missing_or_unclear"]:
            md += "\n#### 🔍 Missing or Unclear\n"
            for m in result["missing_or_unclear"]:
                md += f"* {m}\n"
                
        if result["improved_version"]:
            md += f"\n#### ✨ Suggested Draft\n> {result['improved_version']}\n"
            
        md += "\n---\n#### ❓ Next Steps\n"
        for q in result["follow_up_questions"]:
            md += f"* {q}\n"
            
        if result["advance"] and result["next_section_id"]:
            next_title = SECTIONS[result["next_section_id"]]["title"]
            md += f"\n✅ **Moving to Section {result['next_section_id']}: {next_title}**"
        elif result["advance"] and not result["next_section_id"]:
            md += "\n🎉 **Template Complete!**"
            
        return md

    def call_llm(self, section_id, student_message, accepted_sections, safety_id):
        spec = SECTIONS[section_id]
        payload = {
            "section_id": section_id,
            "section_title": spec["title"],
            "student_prompt": spec["student_prompt"],
            "must_include": spec["must_include"],
            "student_message": student_message,
            "draft_so_far": {"accepted_sections": accepted_sections},
        }

        resp = client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": SYSTEM_MESSAGE},
                {"role": "developer", "content": DEVELOPER_MESSAGE},
                {"role": "user", "content": json.dumps(payload)},
            ],
            response_format={"type": "json_schema", "json_schema": {"name": "coach_reply", "schema": FRAMING_COACH_JSON_SCHEMA, "strict": True}},
            temperature=0.3,
        )
        return json.loads(resp.choices[0].message.content)

coach_logic = FramingCoachLogic()

# --- Gradio App ---

def respond(message, history, state):
    # Initialize state if it's the first message
    if state is None:
        state = {"current_section_id": "0", "accepted_sections": {}, "user_id": "default_user"}
    
    safety_id = hashlib.sha256(state["user_id"].encode()).hexdigest()
    
    # Call the LLM
    result = coach_logic.call_llm(
        state["current_section_id"], 
        message, 
        state["accepted_sections"], 
        safety_id
    )
    
    # Update State if accepted
    if result["status"] == "accepted":
        state["accepted_sections"][state["current_section_id"]] = (
            result["draft_update"]["accepted_version"] or result["improved_version"] or message
        )
    
    # Advance section if applicable
    if result["advance"] and result["next_section_id"]:
        state["current_section_id"] = result["next_section_id"]
    
    # Format response for UI
    final_md = coach_logic.format_as_markdown(result)
    
    return final_md, state

with gr.Blocks(fill_height=True) as demo:
    # Persistent session state
    session_state = gr.State(None)
    
    with gr.Row():
        user_input = gr.Textbox(label="Username", placeholder="Enter to start...")
        pass_input = gr.Textbox(label="Password", type="password")
        login_btn = gr.Button("Login")

    chat_container = gr.Column(visible=False)
    
    with chat_container:
        # We use a standard chatbot with a custom function to handle the state
        chatbot = gr.Chatbot(render_markdown=True, scale=1)
        msg_input = gr.Textbox(placeholder="Type your response here and press Enter...")

        def user_msg(user_message, history):
            # Append user message as a dictionary
            if history is None:
                history = []
            history.append({"role": "user", "content": user_message})
            return "", history
        
        def bot_msg(history, state):
            # The last message in history is the user's prompt
            user_message = history[-1]["content"]
            
            # Call your LLM logic
            bot_markdown, updated_state = respond(user_message, history, state)
            
            # Append the assistant's response as a dictionary
            history.append({"role": "assistant", "content": bot_markdown})
            
            return history, updated_state

        msg_input.submit(user_msg, [msg_input, chatbot], [msg_input, chatbot]).then(
            bot_msg, [chatbot, session_state], [chatbot, session_state]
        )

    def login(u, p):
        if (u == course_username) and (p in password_to_name_dict.keys()):
            return gr.update(visible=True)

    login_btn.click(login, [user_input, pass_input], chat_container)


demo.queue(default_concurrency_limit=4)
demo.launch(show_error=True)