Update app.py

app.py

@@ -1,592 +1,415 @@
-#!/usr/bin/env python3
-"""
-Overthinker — Gradio.Server Backend with SQLite Session Isolation + HF Trace Upload
-
-"""
-
 import os
-import re
-import json
 import uuid
+import json
 import sqlite3
+import httpx
 import requests
-
-from pathlib import Path
-from typing import Optional, Dict, List, Any
-
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import HTMLResponse, PlainTextResponse, Response, JSONResponse
+from fastapi.staticfiles import StaticFiles
 from gradio import Server
-from fastapi import HTTPException
-from starlette.responses import HTMLResponse, PlainTextResponse
-from datasets import Dataset, concatenate_datasets, load_dataset
-import pandas as pd
-# ---------------------------------------------------------------------------
-# Application Setup
-# ---------------------------------------------------------------------------
-app = Server()
-PORT = 7860
-DATA_DIR = Path("data")
-DATA_DIR.mkdir(exist_ok=True)
-
-OPENROUTER_API_KEY = os.getenv('OPENROUTER_API_KEY', '')
-OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
-DEFAULT_MODEL = "nvidia/nemotron-3-nano-30b-a3b"
 
-HF_TOKEN = os.getenv('HF_TOKEN', '')
-HF_DATASET_REPO = os.getenv('HF_DATASET_REPO', 'broadfield-dev/Overthinker-traces')
+# Import static strings from bag.py
+from bag import (
+    BASE_URL,
+    LLMS_TXT,
+    SITEMAP_XML,
+    ROBOTS_TXT,
+    OVERSEER_JSON,
+    VIDEO_PAGE_HTML
+)
+
+app = FastAPI()
 
-# ---------------------------------------------------------------------------
-# Database Helpers
-# ---------------------------------------------------------------------------
+# --- Database helpers ---
+DATA_DIR = "data"
+os.makedirs(DATA_DIR, exist_ok=True)
 
-def get_db_path(session_id: str) -> Path:
-    return DATA_DIR / f"session_{session_id}.db"
+def get_db_path(session_id: str) -> str:
+    return os.path.join(DATA_DIR, f"session_{session_id}.db")
 
-def init_session(session_id: str):
+def init_session_db(session_id: str):
     db_path = get_db_path(session_id)
-    if db_path.exists():
-        return
-    conn = sqlite3.connect(str(db_path))
-    conn.execute("""
-        CREATE TABLE nodes (
-            id TEXT PRIMARY KEY,
-            parent_id TEXT,
-            type TEXT NOT NULL,
-            label TEXT NOT NULL,
-            description TEXT DEFAULT '',
-            emoji TEXT DEFAULT '\U0001f539',
-            tips TEXT DEFAULT '[]',
-            order_index INTEGER DEFAULT 0,
-            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-        )
-    """)
-    root_id = str(uuid.uuid4())
-    conn.execute(
-        "INSERT INTO nodes (id, parent_id, type, label, description, emoji) VALUES (?, ?, ?, ?, ?, ?)",
-        (root_id, None, "root", "What decision do you want to explore?", "", "\U0001f333")
-    )
+    conn = sqlite3.connect(db_path)
+    conn.execute('''CREATE TABLE IF NOT EXISTS nodes (
+        id TEXT PRIMARY KEY,
+        parent_id TEXT,
+        node_type TEXT NOT NULL,
+        label TEXT NOT NULL,
+        description TEXT DEFAULT '',
+        emoji TEXT DEFAULT '',
+        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+    )''')
+    conn.execute('''CREATE TABLE IF NOT EXISTS roots (
+        id TEXT PRIMARY KEY,
+        decision TEXT NOT NULL,
+        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+    )''')
+    # Ensure root node exists
+    root = conn.execute("SELECT id FROM roots LIMIT 1").fetchone()
+    if not root:
+        root_id = str(uuid.uuid4())
+        conn.execute("INSERT INTO roots (id, decision) VALUES (?, 'New Decision')", (root_id,))
+        conn.execute("INSERT INTO nodes (id, parent_id, node_type, label, description) VALUES (?, NULL, 'root', 'What decision do you want to explore?', 'Enter a decision at the top of the page to begin.')", (root_id,))
     conn.commit()
     conn.close()
 
-def get_node_db(session_id: str, node_id: str) -> Optional[Dict]:
+def get_tree_nested(session_id: str) -> dict:
     db_path = get_db_path(session_id)
-    if not db_path.exists():
-        return None
-    conn = sqlite3.connect(str(db_path))
+    conn = sqlite3.connect(db_path)
     conn.row_factory = sqlite3.Row
-    row = conn.execute("SELECT * FROM nodes WHERE id=?", (node_id,)).fetchone()
+    rows = conn.execute("SELECT * FROM nodes ORDER BY created_at").fetchall()
     conn.close()
-    if row is None:
-        return None
-    result = dict(row)
-    try:
-        result['tips'] = json.loads(result.get('tips', '[]'))
-    except:
-        result['tips'] = []
-    return result
+    # Build tree recursively
+    node_map = {}
+    for row in rows:
+        node_map[row['id']] = {
+            'id': row['id'],
+            'parent_id': row['parent_id'],
+            'type': row['node_type'],
+            'label': row['label'],
+            'description': row['description'],
+            'emoji': row['emoji'],
+            'children': []
+        }
+    root = None
+    for nid, node in node_map.items():
+        if node['parent_id'] is None:
+            root = node
+        else:
+            parent = node_map.get(node['parent_id'])
+            if parent:
+                parent['children'].append(node)
+    return root or {'id': 'error', 'label': 'No root found', 'children': []}
 
-def get_children_db(session_id: str, parent_id: str) -> List[Dict]:
+def build_path_string(session_id: str, node_id: str) -> str:
     db_path = get_db_path(session_id)
-    if not db_path.exists():
-        return []
-    conn = sqlite3.connect(str(db_path))
+    conn = sqlite3.connect(db_path)
     conn.row_factory = sqlite3.Row
-    rows = conn.execute(
-        "SELECT * FROM nodes WHERE parent_id=? ORDER BY order_index",
-        (parent_id,)
-    ).fetchall()
+    path_parts = []
+    current_id = node_id
+    while current_id:
+        row = conn.execute("SELECT id, parent_id, node_type, label FROM nodes WHERE id=?", (current_id,)).fetchone()
+        if not row:
+            break
+        path_parts.append(f"[{row['node_type'].upper()}] {row['label']}")
+        current_id = row['parent_id']
     conn.close()
-    result = []
-    for row in rows:
-        d = dict(row)
-        try:
-            d['tips'] = json.loads(d.get('tips', '[]'))
-        except:
-            d['tips'] = []
-        result.append(d)
-    return result
-
-def add_node_db(session_id: str, parent_id: str, node_type: str, label: str,
-                description: str = "", emoji: str = "\U0001f539",
-                tips: list = None, order_index: int = 0) -> Dict:
-    node_id = str(uuid.uuid4())
-    tips_json = json.dumps(tips or [])
+    path_parts.reverse()
+    return " → ".join(path_parts) if path_parts else node_id
+
+def get_node_db(session_id: str, node_id: str) -> dict:
     db_path = get_db_path(session_id)
-    conn = sqlite3.connect(str(db_path))
-    conn.execute(
-        "INSERT INTO nodes (id, parent_id, type, label, description, emoji, tips, order_index) VALUES (?,?,?,?,?,?,?,?)",
-        (node_id, parent_id, node_type, label, description, emoji, tips_json, order_index)
-    )
-    conn.commit()
+    conn = sqlite3.connect(db_path)
+    conn.row_factory = sqlite3.Row
+    row = conn.execute("SELECT * FROM nodes WHERE id=?", (node_id,)).fetchone()
     conn.close()
-    return {
-        "id": node_id,
-        "parent_id": parent_id,
-        "type": node_type,
-        "label": label,
-        "description": description,
-        "emoji": emoji,
-        "tips": tips or [],
-        "order_index": order_index
-    }
+    if row:
+        return dict(row)
+    return None
 
-def update_root_db(session_id: str, label: str, description: str = ""):
+def add_node_db(session_id: str, parent_id: str, node_type: str, label: str, description: str = "", emoji: str = ""):
     db_path = get_db_path(session_id)
-    conn = sqlite3.connect(str(db_path))
+    conn = sqlite3.connect(db_path)
+    node_id = str(uuid.uuid4())
     conn.execute(
-        "UPDATE nodes SET label=?, description=? WHERE parent_id IS NULL",
-        (label, description)
+        "INSERT INTO nodes (id, parent_id, node_type, label, description, emoji) VALUES (?, ?, ?, ?, ?, ?)",
+        (node_id, parent_id, node_type, label, description, emoji)
     )
     conn.commit()
     conn.close()
+    return node_id
 
-def get_path_db(session_id: str, node_id: str) -> List[Dict]:
-    path = []
-    current_id = node_id
-    while current_id:
-        node = get_node_db(session_id, current_id)
-        if node is None:
-            break
-        path.append(node)
-        current_id = node.get("parent_id")
-    path.reverse()
-    return path
+# --- AI Generation ---
+DEFAULT_MODEL = "nvidia/nemotron-3-nano-30b-a3b"
+OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")
 
-def build_path_string(session_id: str, node_id: str) -> str:
-    nodes = get_path_db(session_id, node_id)
-    parts = []
-    for n in nodes:
-        t = n["type"]
-        label = n["label"]
-        if t == "root":
-            parts.append(f"[ROOT] {label}")
-        elif t == "input":
-            parts.append(f"[INPUT] {label}")
-        elif t == "outcome":
-            parts.append(f"[OUTCOME] {label}")
-    return " → ".join(parts)
-
-def get_root_node(session_id: str) -> Optional[Dict]:
-    db_path = get_db_path(session_id)
-    if not db_path.exists():
-        return None
-    conn = sqlite3.connect(str(db_path))
-    conn.row_factory = sqlite3.Row
-    row = conn.execute("SELECT * FROM nodes WHERE parent_id IS NULL LIMIT 1").fetchone()
-    conn.close()
-    if row is None:
-        return None
-    result = dict(row)
+def call_api(prompt: str, max_tokens: int = 1024) -> str:
+    if not OPENROUTER_API_KEY:
+        raise HTTPException(status_code=500, detail="OPENROUTER_API_KEY not set")
+    response = requests.post(
+        url="https://openrouter.ai/api/v1/chat/completions",
+        headers={
+            "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+            "Content-Type": "application/json"
+        },
+        json={
+            "model": DEFAULT_MODEL,
+            "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": max_tokens,
+            "temperature": 0.8
+        },
+        timeout=60
+    )
+    if response.status_code != 200:
+        raise HTTPException(status_code=500, detail=f"API error: {response.status_code} - {response.text}")
+    data = response.json()
+    choices = data.get("choices", [])
+    if not choices:
+        raise HTTPException(status_code=500, detail="No choices in response")
+    return choices[0].get("message", {}).get("content", "")
+
+def parse_children(text: str) -> list:
+    """Parse AI response into list of dicts with label, description, emoji."""
+    children = []
     try:
-        result['tips'] = json.loads(result.get('tips', '[]'))
-    except:
-        result['tips'] = []
-    return result
+        # Try JSON parsing first
+        data = json.loads(text)
+        if isinstance(data, list):
+            children = data
+        elif isinstance(data, dict) and "children" in data:
+            children = data["children"]
+    except json.JSONDecodeError:
+        # Fallback: split by lines
+        lines = text.strip().split('\n')
+        for line in lines:
+            line = line.strip()
+            if line.startswith('-') or line.startswith('*'):
+                label = line[1:].strip()
+                if label:
+                    children.append({"label": label, "description": "", "emoji": ""})
+    return children
 
-def get_all_node_ids(session_id: str) -> List[str]:
-    """Get IDs of all nodes in the tree (for full export)."""
-    db_path = get_db_path(session_id)
-    if not db_path.exists():
-        return []
-    conn = sqlite3.connect(str(db_path))
-    rows = conn.execute("SELECT id FROM nodes").fetchall()
-    conn.close()
-    return [r[0] for r in rows]
+def build_options_prompt(path_context: str, parent_label: str, parent_desc: str, count: int, comment: str) -> str:
+    return f"""You are generating OPTIONS (choices/decisions) for a decision tree.
 
-def build_tree_nested(session_id: str) -> Optional[Dict]:
-    """Build a nested tree structure from the SQLite DB."""
-    root = get_root_node(session_id)
-    if not root:
-        return None
-    def build_tree(node):
-        children = get_children_db(session_id, node['id'])
-        node_copy = dict(node)
-        if isinstance(node_copy.get('tips'), str):
-            try:
-                node_copy['tips'] = json.loads(node_copy['tips'])
-            except:
-                node_copy['tips'] = []
-        node_copy['children'] = [build_tree(c) for c in children]
-        return node_copy
-    return build_tree(root)
-
-# ---------------------------------------------------------------------------
-# Prompt Builders (with path_context)
-# ---------------------------------------------------------------------------
-
-def build_root_prompt(decision: str) -> str:
-    return f'''You are an AI that helps people explore decisions by generating decision trees.
-
-Generate a ROOT decision node for the following decision:
-
-"{decision}"
-
-Return ONLY valid JSON with exactly this structure (no markdown, no backticks):
-{{
-    "label": "A concise label for this decision tree (3-6 words)",
-    "description": "A 1-2 sentence description of this decision context",
-    "emoji": "An emoji representing this decision",
-    "tips": ["One actionable tip for approaching this decision"]
-}}'''
-
-def build_options_prompt(decision_label: str, decision_desc: str, count: int, path_context: str, comment: str = "") -> str:
-    path_section = f'\nFull path from root to this node: "{path_context}"' if path_context else ''
-    comment_section = f'\nUser context: "{comment}"' if comment else ''
-    return f'''You are an AI that helps explore decisions by generating decision tree branches.
-
-Parent node: "{decision_label}"
-Description: "{decision_desc}"{path_section}{comment_section}
-
-Generate EXACTLY {count} child nodes that represent different OPTIONS or CHOICES the person could take.
-
-IMPORTANT: Frame each child as an OPTION or CHOICE, not as an outcome.
-
-Consider the full decision path above to ensure the options are contextually relevant.
-
-Return ONLY valid JSON with exactly this structure (no markdown, no backticks):
-{{
-    "children": [
-        {{
-            "id": "child_1",
-            "label": "Short option label (3-6 words)",
-            "description": "1-2 sentence description",
-            "emoji": "An emoji",
-            "tips": ["One practical tip"]
-        }},
-        ...
-    ]
-}}
-
-Ensure children have unique IDs like child_1, child_2, etc.'''
-
-def build_outcomes_prompt(decision_label: str, decision_desc: str, count: int, path_context: str, comment: str = "") -> str:
-    path_section = f'\nFull path from root to this node: "{path_context}"' if path_context else ''
-    comment_section = f'\nUser context: "{comment}"' if comment else ''
-    return f'''You are an AI that helps explore decisions by generating decision tree branches.
-
-Parent node: "{decision_label}"
-Description: "{decision_desc}"{path_section}{comment_section}
-
-Generate EXACTLY {count} child nodes that represent a DIVERSE RANGE of possible OUTCOMES. Include a MIX of positive, neutral, and negative outcomes.
-
-IMPORTANT: Frame each child as an OUTCOME or CONSEQUENCE, not as a choice someone makes.
-
-Consider the full decision path above to ensure the outcomes are contextually relevant.
-
-Return ONLY valid JSON with exactly this structure (no markdown, no backticks):
-{{
-    "children": [
-        {{
-            "id": "child_1",
-            "label": "Short outcome label (3-6 words)",
-            "description": "1-2 sentence description",
-            "emoji": "An emoji",
-            "tips": ["One practical tip"]
-        }},
-        ...
-    ]
-}}
-
-Ensure children have unique IDs. Make sure the first child is POSITIVE, the second is NEUTRAL, and the third is NEGATIVE.'''
-
-# ---------------------------------------------------------------------------
-# AI Call (using OpenRouter via requests)
-# ---------------------------------------------------------------------------
-
-def call_api(prompt: str, system_prompt: str = "You are a helpful assistant that generates decision trees.") -> Optional[str]:
-    if not OPENROUTER_API_KEY:
-        print("[OpenRouter Error] No API key configured")
-        return None
-    try:
-        headers = {
-            'Authorization': f'Bearer {OPENROUTER_API_KEY}',
-            'Content-Type': 'application/json',
-            'HTTP-Referer': 'http://localhost:7860',
-            'X-Title': 'Overthinker'
-        }
-        data = {
-            'model': DEFAULT_MODEL,
-            'messages': [
-                {'role': 'system', 'content': system_prompt},
-                {'role': 'user', 'content': prompt}
-            ],
-            'temperature': 0.8,
-            'max_tokens': 2048
-        }
-        response = requests.post(
-            OPENROUTER_URL,
-            headers=headers,
-            json=data,
-            timeout=30
-        )
-        if response.status_code == 200:
-            result = response.json()
-            return result['choices'][0]['message']['content']
-        else:
-            print(f"[OpenRouter Error] {response.status_code}: {response.text}")
-    except Exception as e:
-        print(f"[OpenRouter Exception] {e}")
-    return None
+Full path from root to this node:
+{path_context}
 
-def parse_json_response(text: str) -> Optional[dict]:
-    if not text:
-        return None
-    text = text.strip()
-    text = re.sub(r'```json\s*', '', text)
-    text = re.sub(r'```\s*', '', text)
-    text = text.strip()
-    start = text.find('{')
-    end = text.rfind('}')
-    if start >= 0 and end > start:
-        text = text[start:end+1]
-    try:
-        return json.loads(text)
-    except json.JSONDecodeError as e:
-        print(f"[JSON Parse Error] {e}")
-        print(f"[Raw text] {text[:500]}")
-        return None
-
-# ---------------------------------------------------------------------------
-# Routes (All POST, no GET except for serving index)
-# ---------------------------------------------------------------------------
-
-@app.get("/")
-async def index():
-    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates", "index.html")
-    if os.path.exists(html_path):
-        with open(html_path, "r", encoding="utf-8") as f:
-            return HTMLResponse(content=f.read(), status_code=200)
-    return HTMLResponse(content="<h1>Overthinker</h1><p>index.html not found</p>", status_code=404)
+Current node: {parent_label}
+Description: {parent_desc}
+
+Generate {count} distinct, creative options that follow from this node. Each option should be a possible action, choice, or path forward that makes sense given the full context above.
+
+CRITICAL: Respond ONLY with a valid JSON array of objects. Each object must have:
+- "label": A short, punchy title (2-6 words)
+- "description": 1-2 sentence explanation of this option
+- "emoji": A single emoji character representing this option
+
+Example:
+[
+  {{"label": "Start freelancing", "description": "Begin working independently as a freelancer", "emoji": "💼"}},
+  {{"label": "Take a course", "description": "Enroll in a structured learning program", "emoji": "📚"}}
+]
+
+IMPORTANT: Your response must be ONLY the JSON array. No markdown, no explanations, no code blocks."""
+
+def build_outcomes_prompt(path_context: str, parent_label: str, parent_desc: str, count: int, comment: str) -> str:
+    return f"""You are generating OUTCOMES (results/consequences) for a decision tree.
+
+Full path from root to this node:
+{path_context}
+
+Current node: {parent_label}
+Description: {parent_desc}
+
+Generate {count} distinct, realistic outcomes that could result from this choice. Each outcome should feel like a natural consequence given the full decision history above.
+
+CRITICAL: Respond ONLY with a valid JSON array of objects. Each object must have:
+- "label": A short, punchy title (2-6 words)
+- "description": 1-2 sentence explanation of this outcome
+- "emoji": A single emoji character representing this outcome
+
+Example:
+[
+  {{"label": "Financial stability improves", "description": "The freelancer enjoys a steady income over time", "emoji": "💰"}},
+  {{"label": "Loneliness sets in", "description": "Working alone leads to feelings of isolation", "emoji": "😔"}}
+]
+
+IMPORTANT: Your response must be ONLY the JSON array. No markdown, no explanations, no code blocks."""
+
+# --- API Endpoints ---
+
+@app.get("/llms.txt", response_class=PlainTextResponse)
+async def get_llms_txt():
+    return PlainTextResponse(LLMS_TXT)
+
+@app.get("/sitemap.xml", response_class=Response)
+async def get_sitemap():
+    return Response(content=SITEMAP_XML, media_type="application/xml")
+
+@app.get("/robots.txt", response_class=PlainTextResponse)
+async def get_robots():
+    return PlainTextResponse(ROBOTS_TXT)
+
+@app.get("/overthinker.json", response_class=Response)
+async def get_overthinker_json():
+    return Response(content=OVERSEER_JSON, media_type="application/json")
+
+@app.get("/video", response_class=HTMLResponse)
+async def get_video():
+    return HTMLResponse(content=VIDEO_PAGE_HTML)
 
 @app.post("/root")
-async def create_root(request: dict):
-    session_id = request.get('session_id', str(uuid.uuid4()))
-    init_session(session_id)
-    root = get_root_node(session_id)
-    if root is None:
-        raise HTTPException(status_code=500, detail="Could not initialize session.")
-    return {"session_id": session_id, "node": root}
-
-@app.post("/create_tree")
-async def create_tree(request: dict):
-    session_id = request.get('session_id', str(uuid.uuid4()))
-    decision = request.get('decision', '')
-    if not decision:
-        raise HTTPException(status_code=400, detail="Decision text is required.")
-    init_session(session_id)
-    prompt = build_root_prompt(decision)
-    ai_response = call_api(prompt)
-    parsed = parse_json_response(ai_response) if ai_response else None
-    if not parsed:
-        raise HTTPException(status_code=500, detail="Failed to generate root node. Please check your API key and try again.")
-    label = parsed.get('label', f'Overthinking: {decision[:40]}')
-    description = parsed.get('description', f'You are overthinking: {decision}')
-    emoji = parsed.get('emoji', '\U0001f333')
-    tips = parsed.get('tips', ['Start by exploring options.'])
-    update_root_db(session_id, label, description)
+async def create_root(request: Request):
+    body = await request.json()
+    session_id = body.get("session_id", str(uuid.uuid4()))
+    decision = body.get("decision", "")
+    init_session_db(session_id)
     db_path = get_db_path(session_id)
-    conn = sqlite3.connect(str(db_path))
-    conn.execute("UPDATE nodes SET emoji=?, tips=? WHERE parent_id IS NULL", (emoji, json.dumps(tips)))
+    conn = sqlite3.connect(db_path)
+    if decision:
+        conn.execute("UPDATE roots SET decision=? WHERE rowid=1", (decision,))
+        root_row = conn.execute("SELECT id FROM roots LIMIT 1").fetchone()
+        if root_row:
+            conn.execute("UPDATE nodes SET label=? WHERE id=?", (decision, root_row[0]))
     conn.commit()
     conn.close()
-    root = get_root_node(session_id)
-    return {'session_id': session_id, 'node': root}
-
-@app.post("/get_node")
-async def get_node_endpoint(request: dict):
-    session_id = request.get('session_id')
-    node_id = request.get('node_id')
-    if not session_id or not node_id:
-        raise HTTPException(status_code=400, detail="Missing session_id or node_id")
-    init_session(session_id)
-    node = get_node_db(session_id, node_id)
-    if node is None:
-        raise HTTPException(status_code=404, detail="Node not found")
-    children = get_children_db(session_id, node_id)
-    path_context = build_path_string(session_id, node_id)
-    return {
-        'node': node,
-        'children': children,
-        'path_context': path_context
-    }
+    tree = get_tree_nested(session_id)
+    path = build_path_string(session_id, tree['id'])
+    return {"session_id": session_id, "tree": tree, "path": path}
 
 @app.post("/get_children")
-async def get_children(request: dict):
-    session_id = request.get('session_id')
-    node_id = request.get('node_id')
-    count = request.get('count', 3)
-    node_type = request.get('node_type', 'outcome')
-    comment = request.get('comment', '')
-    if not session_id or not node_id:
-        raise HTTPException(status_code=400, detail="Missing session_id or node_id")
-    init_session(session_id)
+async def get_children(request: Request):
+    body = await request.json()
+    session_id = body.get("session_id")
+    node_id = body.get("node_id")
+    count = body.get("count", 3)
+    node_type = body.get("node_type", "outcome")
+    comment = body.get("comment", "")
+    
+    init_session_db(session_id)
     parent = get_node_db(session_id, node_id)
-    if parent is None:
+    if not parent:
         raise HTTPException(status_code=404, detail="Node not found")
-    path_context = build_path_string(session_id, node_id)
-    next_type_map = {'root': 'input', 'input': 'outcome', 'outcome': 'input'}
-    next_type = next_type_map.get(node_type, 'outcome')
+    
     parent_label = parent.get('label', 'Unknown')
     parent_desc = parent.get('description', '')
+    path_context = build_path_string(session_id, node_id)
+    
+    next_type = "input" if node_type == "outcome" else "outcome"
+    
     if next_type == 'input':
-        prompt = build_options_prompt(parent_label, parent_desc, count, path_context, comment)
+        prompt = build_options_prompt(path_context, parent_label, parent_desc, count, comment)
     else:
-        prompt = build_outcomes_prompt(parent_label, parent_desc, count, path_context, comment)
-    ai_response = call_api(prompt)
-    parsed = parse_json_response(ai_response) if ai_response else None
-    if not parsed or 'children' not in parsed or not isinstance(parsed['children'], list):
-        raise HTTPException(status_code=500, detail="Generation failed. Please check your API key and try again.")
-    children_data = parsed['children']
-    children = []
-    for i, child in enumerate(children_data):
-        label = child.get('label', 'Unknown')
-        description = child.get('description', '')
-        emoji = child.get('emoji', '\U0001f539')
-        tips = child.get('tips', [f'Consider this {next_type}.'])
-        existing = get_children_db(session_id, node_id)
-        existing_labels = [c['label'] for c in existing]
-        if label in existing_labels or label in [c['label'] for c in children]:
-            label = f"{label} ({i+1})"
-        child_node = add_node_db(session_id, node_id, next_type, label, description, emoji, tips, order_index=i)
-        child_node['type'] = next_type
-        children.append(child_node)
-    return {'children': children, 'next_type': next_type}
+        prompt = build_outcomes_prompt(path_context, parent_label, parent_desc, count, comment)
+    
+    try:
+        text = call_api(prompt, max_tokens=2048)
+        children = parse_children(text)
+        if not children:
+            raise HTTPException(status_code=500, detail="Generation failed. AI returned empty results.")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
+    
+    # Save children to DB
+    child_ids = []
+    for child in children:
+        cid = add_node_db(session_id, node_id, next_type, child.get('label', ''), child.get('description', ''), child.get('emoji', ''))
+        child_ids.append(cid)
+    
+    # Fetch saved children
+    db_path = get_db_path(session_id)
+    conn = sqlite3.connect(db_path)
+    conn.row_factory = sqlite3.Row
+    saved_children = []
+    for cid in child_ids:
+        row = conn.execute("SELECT * FROM nodes WHERE id=?", (cid,)).fetchone()
+        if row:
+            saved_children.append(dict(row))
+    conn.close()
+    
+    parent_label = parent.get('label', '')
+    parent_desc = parent.get('description', '')
+    path_context = build_path_string(session_id, node_id)
+    next_type = "input" if node_type == "outcome" else "outcome"
+    
+    return {
+        "children": saved_children,
+        "parent_label": parent_label,
+        "parent_desc": parent_desc,
+        "path_context": path_context,
+        "next_type": next_type
+    }
 
 @app.post("/add_options")
-async def add_options(request: dict):
-    session_id = request.get('session_id')
-    node_id = request.get('node_id')
-    count = request.get('count', 3)
-    comment = request.get('comment', '')
-    if not session_id or not node_id:
-        raise HTTPException(status_code=400, detail="Missing session_id or node_id")
-    init_session(session_id)
+async def add_options(request: Request):
+    body = await request.json()
+    session_id = body.get("session_id")
+    node_id = body.get("node_id")
+    count = body.get("count", 3)
+    comment = body.get("comment", "")
+    
+    init_session_db(session_id)
     parent = get_node_db(session_id, node_id)
-    if parent is None:
+    if not parent:
         raise HTTPException(status_code=404, detail="Node not found")
-    path_context = build_path_string(session_id, node_id)
-    next_type_map = {'root': 'input', 'input': 'outcome', 'outcome': 'input'}
-    next_type = next_type_map.get(parent.get('type', 'root'), 'outcome')
-    parent_label = parent.get('label', 'Unknown')
+    
+    parent_label = parent.get('label', '')
     parent_desc = parent.get('description', '')
+    path_context = build_path_string(session_id, node_id)
+    next_type = "input" if parent['node_type'] == "outcome" else "outcome"
+    
     if next_type == 'input':
-        prompt = build_options_prompt(parent_label, parent_desc, count, path_context, comment)
+        prompt = build_options_prompt(path_context, parent_label, parent_desc, count, comment)
     else:
-        prompt = build_outcomes_prompt(parent_label, parent_desc, count, path_context, comment)
-    ai_response = call_api(prompt)
-    parsed = parse_json_response(ai_response) if ai_response else None
-    if not parsed or 'children' not in parsed or not isinstance(parsed['children'], list):
-        raise HTTPException(status_code=500, detail="Failed to add options. Please try again.")
-    children_data = parsed['children']
-    children = []
-    for i, child in enumerate(children_data):
-        label = child.get('label', 'Unknown')
-        description = child.get('description', '')
-        emoji = child.get('emoji', '\U0001f539')
-        tips = child.get('tips', [f'Additional {next_type}.'])
-        existing = get_children_db(session_id, node_id)
-        existing_labels = [c['label'] for c in existing]
-        if label in existing_labels or label in [c['label'] for c in children]:
-            label = f"{label} ({i+1})"
-        child_node = add_node_db(session_id, node_id, next_type, label, description, emoji, tips, order_index=i)
-        child_node['type'] = next_type
-        children.append(child_node)
-    return {'children': children, 'next_type': next_type}
+        prompt = build_outcomes_prompt(path_context, parent_label, parent_desc, count, comment)
+    
+    try:
+        text = call_api(prompt, max_tokens=2048)
+        children = parse_children(text)
+        if not children:
+            raise HTTPException(status_code=500, detail="Generation failed. AI returned empty results.")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
+    
+    child_ids = []
+    for child in children:
+        cid = add_node_db(session_id, node_id, next_type, child.get('label', ''), child.get('description', ''), child.get('emoji', ''))
+        child_ids.append(cid)
+    
+    db_path = get_db_path(session_id)
+    conn = sqlite3.connect(db_path)
+    conn.row_factory = sqlite3.Row
+    saved_children = []
+    for cid in child_ids:
+        row = conn.execute("SELECT * FROM nodes WHERE id=?", (cid,)).fetchone()
+        if row:
+            saved_children.append(dict(row))
+    conn.close()
+    
+    return {
+        "children": saved_children,
+        "parent_label": parent_label,
+        "parent_desc": parent_desc,
+        "path_context": path_context,
+        "next_type": next_type
+    }
 
 @app.post("/upload_trace")
-async def upload_trace(request: dict):
-    """Serialize the full tree from SQLite and push to HuggingFace dataset."""
-    session_id = request.get('session_id')
+async def upload_trace(request: Request):
+    body = await request.json()
+    session_id = body.get("session_id")
     if not session_id:
-        raise HTTPException(status_code=400, detail="Missing session_id")
+        raise HTTPException(status_code=400, detail="session_id required")
     
-    if not HF_TOKEN or not HF_DATASET_REPO:
-        raise HTTPException(status_code=500, detail="HF_TOKEN and HF_DATASET_REPO must be configured in environment.")
+    tree = get_tree_nested(session_id)
+    if not tree:
+        raise HTTPException(status_code=404, detail="No tree found")
     
-    tree = build_tree_nested(session_id)
-    if tree is None:
-        raise HTTPException(status_code=404, detail="No tree found for this session.")
+    # Upload to Hugging Face Dataset via REST API
+    hf_token = os.environ.get("HF_TOKEN", "")
+    dataset_repo = os.environ.get("HF_DATASET_REPO", "build-small-hackathon/Overthinker-trace")
+    if not hf_token or not dataset_repo:
+        raise HTTPException(status_code=500, detail="HF_TOKEN or HF_DATASET_REPO not set")
     
-    try:
+    import json as json_module
+    trace_data = json_module.dumps(tree, indent=2)
+    filename = f"trace_{session_id}.json"
+    
+    url = f"https://huggingface.co/api/datasets/{dataset_repo}/upload"
+    files = {'file': (filename, trace_data, 'application/json')}
+    headers = {'Authorization': f'Bearer {hf_token}'}
+    
+    response = requests.post(url, headers=headers, files=files)
+    if response.status_code not in (200, 201):
+        raise HTTPException(status_code=500, detail=f"Upload failed: {response.status_code} - {response.text}")
+    
+    return {"status": "ok", "filename": filename}
 
-        
-        row = {
-            'session_id': session_id,
-            'tree_json': json.dumps(tree),
-            'created_at': str(tree.get('created_at', ''))
-        }
-        df = pd.DataFrame([row])
-        new_dataset = Dataset.from_pandas(df)
-        
-        try:
-            existing_dataset = load_dataset(HF_DATASET_REPO, split='train', token=HF_TOKEN)
-            combined = concatenate_datasets([existing_dataset, new_dataset])
-        except Exception:
-            combined = new_dataset
-        
-        combined.push_to_hub(HF_DATASET_REPO, token=HF_TOKEN, private=False)
-        
-        return {'status': 'success', 'message': 'Trace uploaded successfully!'}
-    except Exception as e:
-        print(f"[Upload Trace Error] {e}")
-        raise HTTPException(status_code=500, detail=f"Failed to upload trace: {str(e)}")
+# --- Serve static frontend ---
+app.mount("/", StaticFiles(directory="templates", html=True), name="templates")
 
-@app.post("/export_json")
-async def export_json(request: dict):
-    session_id = request.get('session_id')
-    if not session_id:
-        raise HTTPException(status_code=400, detail="Missing session_id")
-    root = get_root_node(session_id)
-    if not root:
-        raise HTTPException(status_code=404, detail="No tree found")
-    def build_tree(node):
-        children = get_children_db(session_id, node['id'])
-        node_copy = dict(node)
-        node_copy['children'] = [build_tree(c) for c in children]
-        return node_copy
-    full_tree = build_tree(root)
-    return full_tree
-
-@app.post("/export_path_json")
-async def export_path_json(request: dict):
-    session_id = request.get('session_id')
-    node_id = request.get('node_id')
-    if not session_id or not node_id:
-        raise HTTPException(status_code=400, detail="Missing session_id or node_id")
-    path_nodes = get_path_db(session_id, node_id)
-    return {'path': path_nodes}
-
-@app.post("/export_path_md")
-async def export_path_md(request: dict):
-    session_id = request.get('session_id')
-    node_id = request.get('node_id')
-    if not session_id or not node_id:
-        raise HTTPException(status_code=400, detail="Missing session_id or node_id")
-    path = get_path_db(session_id, node_id)
-    md = '# \U0001f9e0 Overthinker — Decision Path\n\n'
-    for i, node in enumerate(path):
-        indent = '  ' * i
-        emoji = {'root': '\U0001f333', 'input': '\U0001f9e0', 'outcome': '\U0001f4ca'}.get(node.get('type', ''), '\U0001f4cc')
-        md += f'{indent}{emoji} **{node.get("label", "")}**\n'
-        if node.get('description'):
-            md += f'{indent}  > {node.get("description", "")}\n'
-        if node.get('tips') and len(node['tips']) > 0:
-            md += f'{indent}  > \U0001f4a1 {node["tips"][0]}\n'
-        md += '\n'
-    return PlainTextResponse(content=md, status_code=200)
-
-# ---------------------------------------------------------------------------
-# Launch
-# ---------------------------------------------------------------------------
 if __name__ == "__main__":
-    print(f"\U0001f9e0 Overthinker — SQLite Session Mode + HF Trace Upload on port {PORT}")
-    print(f"\U0001f916 Model: {DEFAULT_MODEL}")
-    print(f"\U0001f310 Open http://localhost:{PORT} in your browser")
-    if not OPENROUTER_API_KEY:
-        print("\u26a0\ufe0f  No OPENROUTER_API_KEY found. Add to .env or environment. Generation will fail.")
-    if not HF_TOKEN or not HF_DATASET_REPO:
-        print("\u26a0\ufe0f  No HF_TOKEN or HF_DATASET_REPO set. Upload will fail.")
-    app.launch(
-        server_port=PORT,
-        show_error=True,
-        share=False
-    )
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)