broadfield-dev commited on
Commit
1896e9d
·
verified ·
1 Parent(s): 22f2898

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +647 -0
app.py ADDED
@@ -0,0 +1,647 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Overthinker - Local 4B Quantized Edition (Nemotron 3 Nano 4B)
4
+ Uses a local 4B model (NVIDIA Nemotron 3 Nano 4B) loaded in 4-bit quantization if supported,
5
+ otherwise falls back to BF16 (which fits easily on 24GB GPUs).
6
+ """
7
+
8
+ import os
9
+ import re
10
+ import json
11
+ import uuid
12
+ import sqlite3
13
+ import torch
14
+ from pathlib import Path
15
+ from typing import Optional, Dict, List, Any
16
+
17
+ from gradio import Server
18
+ from fastapi import HTTPException
19
+ from starlette.responses import HTMLResponse, PlainTextResponse, JSONResponse
20
+ from datasets import Dataset, concatenate_datasets, load_dataset
21
+ import pandas as pd
22
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
23
+ from bag import (
24
+ BASE_URL,
25
+ LLMS_TXT,
26
+ SITEMAP_XML,
27
+ ROBOTS_TXT,
28
+ OVERSEER_JSON,
29
+ VIDEO_PAGE_HTML,
30
+ README_MD
31
+ )
32
+
33
+ # ---------------------------------------------------------------------------
34
+ # Application Setup
35
+ # ---------------------------------------------------------------------------
36
+ app = Server()
37
+ PORT = 7860
38
+ DATA_DIR = Path("data")
39
+ DATA_DIR.mkdir(exist_ok=True)
40
+
41
+ # ---------- Local Model Configuration ----------
42
+ # Using NVIDIA Nemotron 3 Nano 4B (BF16) - a compact Mamba2-Transformer hybrid SLM
43
+ # 4-bit quantization via BitsAndBytes may not support Mamba layers fully;
44
+ # we attempt it first, then fall back to BF16 (model is ~8GB, fits on A10G/T4)
45
+ MODEL_NAME = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16"
46
+
47
+ print("[Overthinker] Attempting to load Nemotron 3 Nano 4B with 4-bit quantization...")
48
+
49
+ # Try 4-bit first; if incompatibility with Mamba layers, fallback to BF16
50
+ bnb_config = BitsAndBytesConfig(
51
+ load_in_4bit=True,
52
+ bnb_4bit_use_double_quant=True,
53
+ bnb_4bit_quant_type="nf4",
54
+ bnb_4bit_compute_dtype=torch.bfloat16
55
+ )
56
+
57
+ try:
58
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=False)
59
+ model = AutoModelForCausalLM.from_pretrained(
60
+ MODEL_NAME,
61
+ quantization_config=bnb_config,
62
+ device_map="auto",
63
+ trust_remote_code=True,
64
+ torch_dtype=torch.bfloat16
65
+ )
66
+ print(f"[Overthinker] Model loaded in 4-bit quantization on device: {model.device}")
67
+ loaded_quantized = True
68
+ except Exception as e:
69
+ print(f"[Overthinker] 4-bit quantization failed: {e}")
70
+ print("[Overthinker] Falling back to BF16 (no quantization) - model is only ~8GB.")
71
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=False)
72
+ model = AutoModelForCausalLM.from_pretrained(
73
+ MODEL_NAME,
74
+ device_map="auto",
75
+ trust_remote_code=True,
76
+ torch_dtype=torch.bfloat16
77
+ )
78
+ loaded_quantized = False
79
+ print(f"[Overthinker] Model loaded in BF16 on device: {model.device}")
80
+
81
+ pipe = pipeline(
82
+ "text-generation",
83
+ model=model,
84
+ tokenizer=tokenizer,
85
+ max_new_tokens=2048,
86
+ temperature=0.8,
87
+ do_sample=True,
88
+ top_p=0.9
89
+ )
90
+
91
+ HF_TOKEN = os.getenv('HF_TOKEN', '')
92
+ HF_DATASET_REPO = os.getenv('HF_DATASET_REPO', 'build-small-hackathon/Overthinker-traces')
93
+
94
+ # ---------------------------------------------------------------------------
95
+ # Database Helpers (same as OVERTHINKER_FINAL)
96
+ # ---------------------------------------------------------------------------
97
+
98
+ def get_db_path(session_id: str) -> Path:
99
+ return DATA_DIR / f"session_{session_id}.db"
100
+
101
+ def init_session(session_id: str):
102
+ db_path = get_db_path(session_id)
103
+ if db_path.exists():
104
+ return
105
+ conn = sqlite3.connect(str(db_path))
106
+ conn.execute("""
107
+ CREATE TABLE nodes (
108
+ id TEXT PRIMARY KEY,
109
+ parent_id TEXT,
110
+ type TEXT NOT NULL,
111
+ label TEXT NOT NULL,
112
+ description TEXT DEFAULT '',
113
+ emoji TEXT DEFAULT '\U0001f539',
114
+ tips TEXT DEFAULT '[]',
115
+ order_index INTEGER DEFAULT 0,
116
+ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
117
+ )
118
+ """)
119
+ root_id = str(uuid.uuid4())
120
+ conn.execute(
121
+ "INSERT INTO nodes (id, parent_id, type, label, description, emoji) VALUES (?, ?, ?, ?, ?, ?)",
122
+ (root_id, None, "root", "What decision do you want to explore?", "", "\U0001f333")
123
+ )
124
+ conn.commit()
125
+ conn.close()
126
+
127
+ def get_node_db(session_id: str, node_id: str) -> Optional[Dict]:
128
+ db_path = get_db_path(session_id)
129
+ if not db_path.exists():
130
+ return None
131
+ conn = sqlite3.connect(str(db_path))
132
+ conn.row_factory = sqlite3.Row
133
+ row = conn.execute("SELECT * FROM nodes WHERE id=?", (node_id,)).fetchone()
134
+ conn.close()
135
+ if row is None:
136
+ return None
137
+ result = dict(row)
138
+ try:
139
+ result['tips'] = json.loads(result.get('tips', '[]'))
140
+ except:
141
+ result['tips'] = []
142
+ return result
143
+
144
+ def get_children_db(session_id: str, parent_id: str) -> List[Dict]:
145
+ db_path = get_db_path(session_id)
146
+ if not db_path.exists():
147
+ return []
148
+ conn = sqlite3.connect(str(db_path))
149
+ conn.row_factory = sqlite3.Row
150
+ rows = conn.execute(
151
+ "SELECT * FROM nodes WHERE parent_id=? ORDER BY order_index",
152
+ (parent_id,)
153
+ ).fetchall()
154
+ conn.close()
155
+ result = []
156
+ for row in rows:
157
+ d = dict(row)
158
+ try:
159
+ d['tips'] = json.loads(d.get('tips', '[]'))
160
+ except:
161
+ d['tips'] = []
162
+ result.append(d)
163
+ return result
164
+
165
+ def add_node_db(session_id: str, parent_id: str, node_type: str, label: str,
166
+ description: str = "", emoji: str = "\U0001f539",
167
+ tips: list = None, order_index: int = 0) -> Dict:
168
+ node_id = str(uuid.uuid4())
169
+ tips_json = json.dumps(tips or [])
170
+ db_path = get_db_path(session_id)
171
+ conn = sqlite3.connect(str(db_path))
172
+ conn.execute(
173
+ "INSERT INTO nodes (id, parent_id, type, label, description, emoji, tips, order_index) VALUES (?,?,?,?,?,?,?,?)",
174
+ (node_id, parent_id, node_type, label, description, emoji, tips_json, order_index)
175
+ )
176
+ conn.commit()
177
+ conn.close()
178
+ return {
179
+ "id": node_id,
180
+ "parent_id": parent_id,
181
+ "type": node_type,
182
+ "label": label,
183
+ "description": description,
184
+ "emoji": emoji,
185
+ "tips": tips or [],
186
+ "order_index": order_index
187
+ }
188
+
189
+ def update_root_db(session_id: str, label: str, description: str = ""):
190
+ db_path = get_db_path(session_id)
191
+ conn = sqlite3.connect(str(db_path))
192
+ conn.execute(
193
+ "UPDATE nodes SET label=?, description=? WHERE parent_id IS NULL",
194
+ (label, description)
195
+ )
196
+ conn.commit()
197
+ conn.close()
198
+
199
+ def get_path_db(session_id: str, node_id: str) -> List[Dict]:
200
+ path = []
201
+ current_id = node_id
202
+ while current_id:
203
+ node = get_node_db(session_id, current_id)
204
+ if node is None:
205
+ break
206
+ path.append(node)
207
+ current_id = node.get("parent_id")
208
+ path.reverse()
209
+ return path
210
+
211
+ def build_path_string(session_id: str, node_id: str) -> str:
212
+ nodes = get_path_db(session_id, node_id)
213
+ parts = []
214
+ for n in nodes:
215
+ t = n["type"]
216
+ label = n["label"]
217
+ if t == "root":
218
+ parts.append(f"[ROOT] {label}")
219
+ elif t == "input":
220
+ parts.append(f"[INPUT] {label}")
221
+ elif t == "outcome":
222
+ parts.append(f"[OUTCOME] {label}")
223
+ return " \u2192 ".join(parts)
224
+
225
+ def get_root_node(session_id: str) -> Optional[Dict]:
226
+ db_path = get_db_path(session_id)
227
+ if not db_path.exists():
228
+ return None
229
+ conn = sqlite3.connect(str(db_path))
230
+ conn.row_factory = sqlite3.Row
231
+ row = conn.execute("SELECT * FROM nodes WHERE parent_id IS NULL LIMIT 1").fetchone()
232
+ conn.close()
233
+ if row is None:
234
+ return None
235
+ result = dict(row)
236
+ try:
237
+ result['tips'] = json.loads(result.get('tips', '[]'))
238
+ except:
239
+ result['tips'] = []
240
+ return result
241
+
242
+ def get_all_node_ids(session_id: str) -> List[str]:
243
+ db_path = get_db_path(session_id)
244
+ if not db_path.exists():
245
+ return []
246
+ conn = sqlite3.connect(str(db_path))
247
+ rows = conn.execute("SELECT id FROM nodes").fetchall()
248
+ conn.close()
249
+ return [r[0] for r in rows]
250
+
251
+ def build_tree_nested(session_id: str) -> Optional[Dict]:
252
+ root = get_root_node(session_id)
253
+ if not root:
254
+ return None
255
+ def build_tree(node):
256
+ children = get_children_db(session_id, node['id'])
257
+ node_copy = dict(node)
258
+ if isinstance(node_copy.get('tips'), str):
259
+ try:
260
+ node_copy['tips'] = json.loads(node_copy['tips'])
261
+ except:
262
+ node_copy['tips'] = []
263
+ node_copy['children'] = [build_tree(c) for c in children]
264
+ return node_copy
265
+ return build_tree(root)
266
+
267
+ # ---------------------------------------------------------------------------
268
+ # Prompt Builders (same as OVERTHINKER_FINAL)
269
+ # ---------------------------------------------------------------------------
270
+
271
+ def build_root_prompt(decision: str) -> str:
272
+ return f'''You are an AI that helps people explore decisions by generating decision trees.
273
+
274
+ Generate a ROOT decision node for the following decision:
275
+
276
+ "{decision}"
277
+
278
+ Return ONLY valid JSON with exactly this structure (no markdown, no backticks):
279
+ {{
280
+ "label": "A concise label for this decision tree (3-6 words)",
281
+ "description": "A 1-2 sentence description of this decision context",
282
+ "emoji": "An emoji representing this decision",
283
+ "tips": ["One actionable tip for approaching this decision"]
284
+ }}'''
285
+
286
+ def build_options_prompt(decision_label: str, decision_desc: str, count: int, path_context: str, comment: str = "") -> str:
287
+ path_section = f'\nFull path from root to this node: "{path_context}"' if path_context else ''
288
+ comment_section = f'\nUser context: "{comment}"' if comment else ''
289
+ return f'''You are an AI that helps explore decisions by generating decision tree branches.
290
+
291
+ Parent node: "{decision_label}"
292
+ Description: "{decision_desc}"{path_section}{comment_section}
293
+
294
+ Generate EXACTLY {count} child nodes that represent different OPTIONS or CHOICES the person could take.
295
+
296
+ IMPORTANT: Frame each child as an OPTION or CHOICE, not as an outcome.
297
+
298
+ Consider the full decision path above to ensure the options are contextually relevant.
299
+
300
+ Return ONLY valid JSON with exactly this structure (no markdown, no backticks):
301
+ {{
302
+ "children": [
303
+ {{
304
+ "id": "child_1",
305
+ "label": "Short option label (3-6 words)",
306
+ "description": "1-2 sentence description",
307
+ "emoji": "An emoji",
308
+ "tips": ["One practical tip"]
309
+ }},
310
+ ...
311
+ ]
312
+ }}
313
+
314
+ Ensure children have unique IDs like child_1, child_2, etc.'''
315
+
316
+ def build_outcomes_prompt(decision_label: str, decision_desc: str, count: int, path_context: str, comment: str = "") -> str:
317
+ path_section = f'\nFull path from root to this node: "{path_context}"' if path_context else ''
318
+ comment_section = f'\nUser context: "{comment}"' if comment else ''
319
+ return f'''You are an AI that helps explore decisions by generating decision tree branches.
320
+
321
+ Parent node: "{decision_label}"
322
+ Description: "{decision_desc}"{path_section}{comment_section}
323
+
324
+ Generate EXACTLY {count} child nodes that represent a DIVERSE RANGE of possible OUTCOMES. Include a MIX of positive, neutral, and negative outcomes.
325
+
326
+ IMPORTANT: Frame each child as an OUTCOME or CONSEQUENCE, not as a choice someone makes.
327
+
328
+ Consider the full decision path above to ensure the outcomes are contextually relevant.
329
+
330
+ Return ONLY valid JSON with exactly this structure (no markdown, no backticks):
331
+ {{
332
+ "children": [
333
+ {{
334
+ "id": "child_1",
335
+ "label": "Short outcome label (3-6 words)",
336
+ "description": "1-2 sentence description",
337
+ "emoji": "An emoji",
338
+ "tips": ["One practical tip"]
339
+ }},
340
+ ...
341
+ ]
342
+ }}
343
+
344
+ Ensure children have unique IDs. Make sure the first child is POSITIVE, the second is NEUTRAL, and the third is NEGATIVE.'''
345
+
346
+ # ---------------------------------------------------------------------------
347
+ # AI Call (Local Model via pipeline)
348
+ # ---------------------------------------------------------------------------
349
+
350
+ def call_model(prompt: str, system_prompt: str = "You are a helpful assistant that generates decision trees.") -> Optional[str]:
351
+ messages = [
352
+ {"role": "system", "content": system_prompt},
353
+ {"role": "user", "content": prompt}
354
+ ]
355
+ try:
356
+ outputs = pipe(messages, max_new_tokens=2048, temperature=0.8, do_sample=True)
357
+ response_text = outputs[0]["generated_text"][-1]["content"]
358
+ return response_text
359
+ except Exception as e:
360
+ print(f"[Local Model Error] {e}")
361
+ return None
362
+
363
+ def parse_json_response(text: str) -> Optional[dict]:
364
+ if not text:
365
+ return None
366
+ text = text.strip()
367
+ text = re.sub(r'```json\s*', '', text)
368
+ text = re.sub(r'```\s*', '', text)
369
+ text = text.strip()
370
+ start = text.find('{')
371
+ end = text.rfind('}')
372
+ if start >= 0 and end > start:
373
+ text = text[start:end+1]
374
+ try:
375
+ return json.loads(text)
376
+ except json.JSONDecodeError as e:
377
+ print(f"[JSON Parse Error] {e}")
378
+ print(f"[Raw text] {text[:500]}")
379
+ return None
380
+
381
+ # ---------------------------------------------------------------------------
382
+ # Routes (All POST, no GET except for serving index)
383
+ # ---------------------------------------------------------------------------
384
+
385
+ @app.get("/")
386
+ async def index():
387
+ html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates", "index.html")
388
+ if os.path.exists(html_path):
389
+ with open(html_path, "r", encoding="utf-8") as f:
390
+ return HTMLResponse(content=f.read(), status_code=200)
391
+ return HTMLResponse(content="<h1>Overthinker</h1><p>index.html not found</p>", status_code=404)
392
+
393
+ @app.post("/root")
394
+ async def create_root(request: dict):
395
+ session_id = request.get('session_id', str(uuid.uuid4()))
396
+ init_session(session_id)
397
+ root = get_root_node(session_id)
398
+ if root is None:
399
+ raise HTTPException(status_code=500, detail="Could not initialize session.")
400
+ return {"session_id": session_id, "node": root}
401
+
402
+ @app.post("/create_tree")
403
+ async def create_tree(request: dict):
404
+ session_id = request.get('session_id', str(uuid.uuid4()))
405
+ decision = request.get('decision', '')
406
+ if not decision:
407
+ raise HTTPException(status_code=400, detail="Decision text is required.")
408
+ init_session(session_id)
409
+ prompt = build_root_prompt(decision)
410
+ ai_response = call_model(prompt)
411
+ parsed = parse_json_response(ai_response) if ai_response else None
412
+ if not parsed:
413
+ raise HTTPException(status_code=500, detail="Failed to generate root node. Please check model availability.")
414
+ label = parsed.get('label', f'Overthinking: {decision[:40]}')
415
+ description = parsed.get('description', f'You are overthinking: {decision}')
416
+ emoji = parsed.get('emoji', '\U0001f333')
417
+ tips = parsed.get('tips', ['Start by exploring options.'])
418
+ update_root_db(session_id, label, description)
419
+ db_path = get_db_path(session_id)
420
+ conn = sqlite3.connect(str(db_path))
421
+ conn.execute("UPDATE nodes SET emoji=?, tips=? WHERE parent_id IS NULL", (emoji, json.dumps(tips)))
422
+ conn.commit()
423
+ conn.close()
424
+ root = get_root_node(session_id)
425
+ return {'session_id': session_id, 'node': root}
426
+
427
+ @app.post("/get_node")
428
+ async def get_node_endpoint(request: dict):
429
+ session_id = request.get('session_id')
430
+ node_id = request.get('node_id')
431
+ if not session_id or not node_id:
432
+ raise HTTPException(status_code=400, detail="Missing session_id or node_id")
433
+ init_session(session_id)
434
+ node = get_node_db(session_id, node_id)
435
+ if node is None:
436
+ raise HTTPException(status_code=404, detail="Node not found")
437
+ children = get_children_db(session_id, node_id)
438
+ path_context = build_path_string(session_id, node_id)
439
+ return {
440
+ 'node': node,
441
+ 'children': children,
442
+ 'path_context': path_context
443
+ }
444
+
445
+ @app.post("/get_children")
446
+ async def get_children(request: dict):
447
+ session_id = request.get('session_id')
448
+ node_id = request.get('node_id')
449
+ count = request.get('count', 3)
450
+ node_type = request.get('node_type', 'outcome')
451
+ comment = request.get('comment', '')
452
+ if not session_id or not node_id:
453
+ raise HTTPException(status_code=400, detail="Missing session_id or node_id")
454
+ init_session(session_id)
455
+ parent = get_node_db(session_id, node_id)
456
+ if parent is None:
457
+ raise HTTPException(status_code=404, detail="Node not found")
458
+ path_context = build_path_string(session_id, node_id)
459
+ next_type_map = {'root': 'input', 'input': 'outcome', 'outcome': 'input'}
460
+ next_type = next_type_map.get(node_type, 'outcome')
461
+ parent_label = parent.get('label', 'Unknown')
462
+ parent_desc = parent.get('description', '')
463
+ if next_type == 'input':
464
+ prompt = build_options_prompt(parent_label, parent_desc, count, path_context, comment)
465
+ else:
466
+ prompt = build_outcomes_prompt(parent_label, parent_desc, count, path_context, comment)
467
+ ai_response = call_model(prompt)
468
+ parsed = parse_json_response(ai_response) if ai_response else None
469
+ if not parsed or 'children' not in parsed or not isinstance(parsed['children'], list):
470
+ raise HTTPException(status_code=500, detail="Generation failed. Please try again.")
471
+ children_data = parsed['children']
472
+ children = []
473
+ for i, child in enumerate(children_data):
474
+ label = child.get('label', 'Unknown')
475
+ description = child.get('description', '')
476
+ emoji = child.get('emoji', '\U0001f539')
477
+ tips = child.get('tips', [f'Consider this {next_type}.'])
478
+ existing = get_children_db(session_id, node_id)
479
+ existing_labels = [c['label'] for c in existing]
480
+ if label in existing_labels or label in [c['label'] for c in children]:
481
+ label = f"{label} ({i+1})"
482
+ child_node = add_node_db(session_id, node_id, next_type, label, description, emoji, tips, order_index=i)
483
+ child_node['type'] = next_type
484
+ children.append(child_node)
485
+ return {'children': children, 'next_type': next_type}
486
+
487
+ @app.post("/add_options")
488
+ async def add_options(request: dict):
489
+ session_id = request.get('session_id')
490
+ node_id = request.get('node_id')
491
+ count = request.get('count', 3)
492
+ comment = request.get('comment', '')
493
+ if not session_id or not node_id:
494
+ raise HTTPException(status_code=400, detail="Missing session_id or node_id")
495
+ init_session(session_id)
496
+ parent = get_node_db(session_id, node_id)
497
+ if parent is None:
498
+ raise HTTPException(status_code=404, detail="Node not found")
499
+ path_context = build_path_string(session_id, node_id)
500
+ next_type_map = {'root': 'input', 'input': 'outcome', 'outcome': 'input'}
501
+ next_type = next_type_map.get(parent.get('type', 'root'), 'outcome')
502
+ parent_label = parent.get('label', 'Unknown')
503
+ parent_desc = parent.get('description', '')
504
+ if next_type == 'input':
505
+ prompt = build_options_prompt(parent_label, parent_desc, count, path_context, comment)
506
+ else:
507
+ prompt = build_outcomes_prompt(parent_label, parent_desc, count, path_context, comment)
508
+ ai_response = call_model(prompt)
509
+ parsed = parse_json_response(ai_response) if ai_response else None
510
+ if not parsed or 'children' not in parsed or not isinstance(parsed['children'], list):
511
+ raise HTTPException(status_code=500, detail="Failed to add options. Please try again.")
512
+ children_data = parsed['children']
513
+ children = []
514
+ for i, child in enumerate(children_data):
515
+ label = child.get('label', 'Unknown')
516
+ description = child.get('description', '')
517
+ emoji = child.get('emoji', '\U0001f539')
518
+ tips = child.get('tips', [f'Additional {next_type}.'])
519
+ existing = get_children_db(session_id, node_id)
520
+ existing_labels = [c['label'] for c in existing]
521
+ if label in existing_labels or label in [c['label'] for c in children]:
522
+ label = f"{label} ({i+1})"
523
+ child_node = add_node_db(session_id, node_id, next_type, label, description, emoji, tips, order_index=i)
524
+ child_node['type'] = next_type
525
+ children.append(child_node)
526
+ return {'children': children, 'next_type': next_type}
527
+
528
+ @app.post("/upload_trace")
529
+ async def upload_trace(request: dict):
530
+ session_id = request.get('session_id')
531
+ if not session_id:
532
+ raise HTTPException(status_code=400, detail="Missing session_id")
533
+
534
+ if not HF_TOKEN or not HF_DATASET_REPO:
535
+ raise HTTPException(status_code=500, detail="HF_TOKEN and HF_DATASET_REPO must be configured.")
536
+
537
+ tree = build_tree_nested(session_id)
538
+ if tree is None:
539
+ raise HTTPException(status_code=404, detail="No tree found for this session.")
540
+
541
+ try:
542
+ row = {
543
+ 'session_id': session_id,
544
+ 'tree_json': json.dumps(tree),
545
+ 'created_at': str(tree.get('created_at', ''))
546
+ }
547
+ df = pd.DataFrame([row])
548
+ new_dataset = Dataset.from_pandas(df)
549
+
550
+ try:
551
+ existing_dataset = load_dataset(HF_DATASET_REPO, split='train', token=HF_TOKEN)
552
+ combined = concatenate_datasets([existing_dataset, new_dataset])
553
+ except Exception:
554
+ combined = new_dataset
555
+
556
+ combined.push_to_hub(HF_DATASET_REPO, token=HF_TOKEN, private=False)
557
+ return {'status': 'success', 'message': 'Trace uploaded successfully!'}
558
+ except Exception as e:
559
+ print(f"[Upload Trace Error] {e}")
560
+ raise HTTPException(status_code=500, detail=f"Failed to upload trace: {str(e)}")
561
+
562
+ @app.post("/export_json")
563
+ async def export_json(request: dict):
564
+ session_id = request.get('session_id')
565
+ if not session_id:
566
+ raise HTTPException(status_code=400, detail="Missing session_id")
567
+ root = get_root_node(session_id)
568
+ if not root:
569
+ raise HTTPException(status_code=404, detail="No tree found")
570
+ def build_tree(node):
571
+ children = get_children_db(session_id, node['id'])
572
+ node_copy = dict(node)
573
+ node_copy['children'] = [build_tree(c) for c in children]
574
+ return node_copy
575
+ full_tree = build_tree(root)
576
+ return full_tree
577
+
578
+ @app.post("/export_path_json")
579
+ async def export_path_json(request: dict):
580
+ session_id = request.get('session_id')
581
+ node_id = request.get('node_id')
582
+ if not session_id or not node_id:
583
+ raise HTTPException(status_code=400, detail="Missing session_id or node_id")
584
+ path_nodes = get_path_db(session_id, node_id)
585
+ return {'path': path_nodes}
586
+
587
+ @app.post("/export_path_md")
588
+ async def export_path_md(request: dict):
589
+ session_id = request.get('session_id')
590
+ node_id = request.get('node_id')
591
+ if not session_id or not node_id:
592
+ raise HTTPException(status_code=400, detail="Missing session_id or node_id")
593
+ path = get_path_db(session_id, node_id)
594
+ md = '# \U0001f9e0 Overthinker \u2014 Decision Path\n\n'
595
+ for i, node in enumerate(path):
596
+ indent = ' ' * i
597
+ emoji = {'root': '\U0001f333', 'input': '\U0001f9e0', 'outcome': '\U0001f4ca'}.get(node.get('type', ''), '\U0001f4cc')
598
+ md += f'{indent}{emoji} **{node.get("label", "")}**\n'
599
+ if node.get('description'):
600
+ md += f'{indent} > {node.get("description", "")}\n'
601
+ if node.get('tips') and len(node['tips']) > 0:
602
+ md += f'{indent} > \U0001f4a1 {node["tips"][0]}\n'
603
+ md += '\n'
604
+ return PlainTextResponse(content=md, status_code=200)
605
+
606
+ @app.get("/llms.txt", response_class=PlainTextResponse)
607
+ async def get_llms_txt():
608
+ return PlainTextResponse(LLMS_TXT)
609
+
610
+ @app.get("/readme.md", response_class=PlainTextResponse)
611
+ async def get_readme_md():
612
+ return PlainTextResponse(README_MD)
613
+
614
+ @app.get("/sitemap.xml", response_class=HTMLResponse)
615
+ async def get_sitemap():
616
+ return HTMLResponse(content=SITEMAP_XML, media_type="application/xml")
617
+
618
+ @app.get("/robots.txt", response_class=PlainTextResponse)
619
+ async def get_robots():
620
+ return PlainTextResponse(ROBOTS_TXT)
621
+
622
+ @app.get("/overthinker.json", response_class=JSONResponse)
623
+ async def get_overthinker_json():
624
+ return JSONResponse(content=OVERSEER_JSON, media_type="application/json")
625
+
626
+ @app.get("/video", response_class=HTMLResponse)
627
+ async def get_video():
628
+ return HTMLResponse(content=VIDEO_PAGE_HTML)
629
+
630
+ # ---------------------------------------------------------------------------
631
+ # Launch
632
+ # ---------------------------------------------------------------------------
633
+ if __name__ == "__main__":
634
+ print(f"\U0001f9e0 Overthinker \u2014 Local 4B Quantized Edition on port {PORT}")
635
+ print(f"\U0001f916 Model: {MODEL_NAME}")
636
+ if loaded_quantized:
637
+ print("\U0001f4be Quantization: 4-bit NF4 (BitsAndBytes)")
638
+ else:
639
+ print("\U0001f4be Quantization: None (BF16 fallback)")
640
+ print(f"\U0001f310 Open http://localhost:{PORT} in your browser")
641
+ if not HF_TOKEN or not HF_DATASET_REPO:
642
+ print("\u26a0\ufe0f No HF_TOKEN or HF_DATASET_REPO set. Upload will fail.")
643
+ app.launch(
644
+ server_port=PORT,
645
+ show_error=True,
646
+ share=False
647
+ )