Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """ | |
| IC Generate FastAPI — Single-file app combining: | |
| - Featherless API proxy with key rotation | |
| - Concurrent UI question & solution generator | |
| - SQLite database with crash recovery | |
| - REST API for management (Admin-protected) | |
| - API Key management routes | |
| - [ENGINE] Question generation engine with randomized prompt templates | |
| - [ENGINE] Runtime prompt management endpoints | |
| Persistent storage: /data | |
| - /data/ic_data.db — SQLite database | |
| - /data/keys.txt — API keys (one per line) | |
| - /data/exported_code/ — File exports | |
| - /data/prompt_engine_state.json — [ENGINE] Prompt engine state | |
| """ | |
| # === Imports === | |
| import os | |
| import re | |
| import json | |
| import time | |
| import random | |
| import sqlite3 | |
| import logging | |
| import signal | |
| import threading | |
| import itertools | |
| from typing import Optional, List, Dict, Any | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| from contextlib import asynccontextmanager | |
| import httpx | |
| from fastapi import FastAPI, Request, HTTPException, Query, Depends, Security, status | |
| from fastapi.security import APIKeyHeader | |
| from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse | |
| from pydantic import BaseModel | |
| from dotenv import load_dotenv | |
| from openai import OpenAI, APIError, APIConnectionError, RateLimitError | |
| import sys | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| # NOW your imports will work: | |
| from prompt_engine import QuestionGenerationEngine, PromptConfig | |
| load_dotenv() | |
| # === Configuration === | |
| DATA_DIR = os.environ.get("DATA_DIR", "/data") | |
| os.makedirs(DATA_DIR, exist_ok=True) | |
| DB_PATH = os.path.join(DATA_DIR, "ic_data.db") | |
| KEYS_FILE = os.path.join(DATA_DIR, "keys.txt") | |
| EXPORT_DIR = os.path.join(DATA_DIR, "exported_code") | |
| # [ENGINE] State file for prompt engine persistence | |
| PROMPT_STATE_FILE = os.path.join(DATA_DIR, "prompt_engine_state.json") | |
| FEATHERLESS_API_BASE = os.environ.get( | |
| "FEATHERLESS_API_BASE", "https://api.featherless.ai/v1" | |
| ) | |
| PORT = int(os.environ.get("PORT", "7860")) | |
| STACK = "HTML/CSS/JS" | |
| # [ENGINE] Initialize the global prompt engine | |
| prompt_engine = QuestionGenerationEngine( | |
| state_file=PROMPT_STATE_FILE, | |
| min_temperature=float(os.environ.get("PROMPT_MIN_TEMP", "0.85")), | |
| max_temperature=float(os.environ.get("PROMPT_MAX_TEMP", "1.15")), | |
| max_tokens=4096, | |
| ) | |
| # === Admin Authentication === | |
| ADMIN_API_KEY = os.environ.get("ADMIN_API_KEY", "change-me-default-admin-key") | |
| api_key_header_auth = APIKeyHeader(name="X-Admin-Key", auto_error=False) | |
| async def verify_admin_key(api_key: str = Security(api_key_header_auth)): | |
| if api_key == ADMIN_API_KEY: | |
| return api_key | |
| raise HTTPException( | |
| status_code=status.HTTP_401_UNAUTHORIZED, | |
| detail="Invalid or missing Admin API Key. Provide it in the 'X-Admin-Key' header.", | |
| ) | |
| # === Logging === | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s %(levelname)s [%(threadName)s]: %(message)s", | |
| datefmt="%H:%M:%S", | |
| ) | |
| logger = logging.getLogger(__name__) | |
| # === Shutdown Event === | |
| _shutdown = threading.Event() | |
| def _handle_signal(sig, frame): | |
| logger.warning("⚠️ Shutdown requested. Finishing in-flight requests...") | |
| _shutdown.set() | |
| signal.signal(signal.SIGINT, _handle_signal) | |
| signal.signal(signal.SIGTERM, _handle_signal) | |
| # === Database Layer === | |
| DB_LOCK = threading.Lock() | |
| def init_db(db_path: str = DB_PATH): | |
| """Create all tables including pending_jobs for crash recovery.""" | |
| os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True) | |
| conn = sqlite3.connect(db_path, timeout=60) | |
| conn.execute("PRAGMA foreign_keys = ON") | |
| conn.execute("PRAGMA journal_mode = DELETE") | |
| conn.execute("PRAGMA busy_timeout = 30000") | |
| c = conn.cursor() | |
| c.execute( | |
| """CREATE TABLE IF NOT EXISTS questions ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| question_text TEXT NOT NULL, | |
| thinking_trace TEXT, | |
| full_response TEXT, | |
| model TEXT, | |
| finish_reason TEXT, | |
| usage_json TEXT, | |
| raw_chunks_json TEXT, | |
| chunk_count INTEGER, | |
| generation_time_s REAL, | |
| prompt_metadata TEXT, | |
| created_at TEXT NOT NULL DEFAULT (datetime('now')) | |
| )""" | |
| ) | |
| c.execute( | |
| """CREATE TABLE IF NOT EXISTS solutions ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| question_id INTEGER NOT NULL, | |
| stack TEXT NOT NULL DEFAULT 'HTML/CSS/JS', | |
| solution_code TEXT NOT NULL, | |
| thinking_trace TEXT, | |
| full_response TEXT, | |
| model TEXT, | |
| finish_reason TEXT, | |
| usage_json TEXT, | |
| raw_chunks_json TEXT, | |
| chunk_count INTEGER, | |
| generation_time_s REAL, | |
| created_at TEXT NOT NULL DEFAULT (datetime('now')), | |
| FOREIGN KEY(question_id) REFERENCES questions(id) ON DELETE CASCADE | |
| )""" | |
| ) | |
| c.execute( | |
| """CREATE TABLE IF NOT EXISTS pending_jobs ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| question_id INTEGER NOT NULL, | |
| status TEXT NOT NULL DEFAULT 'pending', | |
| attempts INTEGER NOT NULL DEFAULT 0, | |
| last_error TEXT, | |
| started_at TEXT, | |
| completed_at TEXT, | |
| created_at TEXT NOT NULL DEFAULT (datetime('now')), | |
| FOREIGN KEY(question_id) REFERENCES questions(id) ON DELETE CASCADE | |
| )""" | |
| ) | |
| c.execute( | |
| """CREATE TABLE IF NOT EXISTS run_log ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| started_at TEXT NOT NULL DEFAULT (datetime('now')), | |
| finished_at TEXT, | |
| q_model TEXT, s_model TEXT, | |
| requested INTEGER, completed INTEGER, failed INTEGER, pending INTEGER, | |
| config_json TEXT | |
| )""" | |
| ) | |
| # Migrations | |
| _migrate_add_columns( | |
| c, | |
| "questions", | |
| { | |
| "thinking_trace": "TEXT", | |
| "full_response": "TEXT", | |
| "generation_time_s": "REAL", | |
| "prompt_metadata": "TEXT", # [ENGINE] New column | |
| }, | |
| ) | |
| _migrate_add_columns( | |
| c, | |
| "solutions", | |
| {"thinking_trace": "TEXT", "full_response": "TEXT", "generation_time_s": "REAL"}, | |
| ) | |
| _migrate_add_columns( | |
| c, | |
| "pending_jobs", | |
| { | |
| "attempts": "INTEGER DEFAULT 0", | |
| "last_error": "TEXT", | |
| "started_at": "TEXT", | |
| "completed_at": "TEXT", | |
| }, | |
| ) | |
| conn.commit() | |
| conn.close() | |
| logger.info(f"Database ready: {db_path}") | |
| def _migrate_add_columns(cursor, table, columns): | |
| for col, col_type in columns.items(): | |
| try: | |
| cursor.execute(f"ALTER TABLE {table} ADD COLUMN {col} {col_type}") | |
| except sqlite3.OperationalError: | |
| pass | |
| def _safe_db_execute(db_path, operations, max_retries=5): | |
| for attempt in range(max_retries): | |
| try: | |
| with DB_LOCK: | |
| conn = sqlite3.connect(db_path, timeout=60) | |
| conn.execute("PRAGMA journal_mode = DELETE") | |
| conn.execute("PRAGMA busy_timeout = 30000") | |
| cur = conn.cursor() | |
| try: | |
| result = operations(cur) | |
| conn.commit() | |
| return result | |
| except Exception: | |
| conn.rollback() | |
| raise | |
| finally: | |
| conn.close() | |
| except sqlite3.OperationalError as e: | |
| if "locked" in str(e).lower() and attempt < max_retries - 1: | |
| wait = 0.5 * (2**attempt) | |
| logger.warning( | |
| f"DB locked (attempt {attempt+1}/{max_retries}), retrying in {wait:.1f}s..." | |
| ) | |
| time.sleep(wait) | |
| else: | |
| raise | |
| def save_question( | |
| db_path, q_text, q_thinking, q_full, q_model, q_finish, q_usage, q_chunks, gen_time, | |
| prompt_metadata=None, # [ENGINE] New parameter | |
| ): | |
| def ops(cur): | |
| cur.execute( | |
| """INSERT INTO questions | |
| (question_text, thinking_trace, full_response, model, finish_reason, | |
| usage_json, raw_chunks_json, chunk_count, generation_time_s, prompt_metadata) | |
| VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", | |
| ( | |
| q_text, | |
| q_thinking, | |
| q_full, | |
| q_model, | |
| q_finish, | |
| json.dumps(q_usage, ensure_ascii=False) if q_usage else None, | |
| json.dumps(q_chunks, ensure_ascii=False, default=str), | |
| len(q_chunks), | |
| gen_time, | |
| json.dumps(prompt_metadata, ensure_ascii=False) if prompt_metadata else None, # [ENGINE] | |
| ), | |
| ) | |
| question_id = cur.lastrowid | |
| cur.execute( | |
| "INSERT INTO pending_jobs (question_id, status) VALUES (?, 'pending')", | |
| (question_id,), | |
| ) | |
| job_id = cur.lastrowid | |
| return job_id, question_id | |
| return _safe_db_execute(db_path, ops) | |
| def save_solution( | |
| db_path, | |
| job_id, | |
| question_id, | |
| s_code, | |
| s_thinking, | |
| s_full, | |
| s_model, | |
| s_finish, | |
| s_usage, | |
| s_chunks, | |
| gen_time, | |
| ): | |
| def ops(cur): | |
| cur.execute( | |
| """INSERT INTO solutions | |
| (question_id, stack, solution_code, thinking_trace, full_response, | |
| model, finish_reason, usage_json, raw_chunks_json, chunk_count, generation_time_s) | |
| VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", | |
| ( | |
| question_id, | |
| STACK, | |
| s_code, | |
| s_thinking, | |
| s_full, | |
| s_model, | |
| s_finish, | |
| json.dumps(s_usage, ensure_ascii=False) if s_usage else None, | |
| json.dumps(s_chunks, ensure_ascii=False, default=str), | |
| len(s_chunks), | |
| gen_time, | |
| ), | |
| ) | |
| cur.execute( | |
| "UPDATE pending_jobs SET status='done', completed_at=datetime('now') WHERE id=?", | |
| (job_id,), | |
| ) | |
| _safe_db_execute(db_path, ops) | |
| def mark_job_started(db_path, job_id): | |
| def ops(cur): | |
| cur.execute( | |
| "UPDATE pending_jobs SET started_at=datetime('now'), attempts=attempts+1 WHERE id=?", | |
| (job_id,), | |
| ) | |
| _safe_db_execute(db_path, ops) | |
| def mark_job_failed(db_path, job_id, error_msg): | |
| def ops(cur): | |
| cur.execute( | |
| "UPDATE pending_jobs SET last_error=?, status='pending' WHERE id=?", | |
| (str(error_msg)[:500], job_id), | |
| ) | |
| _safe_db_execute(db_path, ops) | |
| def get_pending_jobs(db_path): | |
| conn = sqlite3.connect(db_path, timeout=30) | |
| cur = conn.cursor() | |
| cur.execute( | |
| """SELECT pj.id, pj.question_id, q.question_text, pj.attempts | |
| FROM pending_jobs pj | |
| JOIN questions q ON pj.question_id = q.id | |
| WHERE pj.status = 'pending' | |
| ORDER BY pj.id""" | |
| ) | |
| rows = cur.fetchall() | |
| conn.close() | |
| return [(r[0], r[1], r[2], r[3]) for r in rows] | |
| def log_run( | |
| db_path, q_model, s_model, requested, completed, failed, pending, config | |
| ): | |
| def ops(cur): | |
| cur.execute( | |
| """INSERT INTO run_log | |
| (q_model, s_model, requested, completed, failed, pending, | |
| finished_at, config_json) | |
| VALUES (?, ?, ?, ?, ?, ?, datetime('now'), ?)""", | |
| ( | |
| q_model, | |
| s_model, | |
| requested, | |
| completed, | |
| failed, | |
| pending, | |
| json.dumps(config, ensure_ascii=False), | |
| ), | |
| ) | |
| try: | |
| _safe_db_execute(db_path, ops) | |
| except Exception as e: | |
| logger.warning(f"Failed to log run: {e}") | |
| def delete_question_db(db_path, question_id): | |
| def ops(cur): | |
| cur.execute("DELETE FROM questions WHERE id=?", (question_id,)) | |
| return cur.rowcount > 0 | |
| return _safe_db_execute(db_path, ops) | |
| def get_stats_dict(db_path): | |
| if not os.path.exists(db_path): | |
| return {"error": f"Database {db_path} does not exist."} | |
| with sqlite3.connect(db_path) as conn: | |
| cursor = conn.cursor() | |
| cursor.execute("SELECT COUNT(*) FROM questions") | |
| q_count = cursor.fetchone()[0] | |
| cursor.execute("SELECT COUNT(*) FROM solutions") | |
| s_count = cursor.fetchone()[0] | |
| cursor.execute("SELECT COALESCE(SUM(chunk_count), 0) FROM questions") | |
| q_chunks = cursor.fetchone()[0] | |
| cursor.execute("SELECT COALESCE(SUM(chunk_count), 0) FROM solutions") | |
| s_chunks = cursor.fetchone()[0] | |
| cursor.execute( | |
| "SELECT COUNT(*) FROM questions WHERE thinking_trace IS NOT NULL" | |
| ) | |
| q_with_thinking = cursor.fetchone()[0] | |
| cursor.execute( | |
| "SELECT COUNT(*) FROM solutions WHERE thinking_trace IS NOT NULL" | |
| ) | |
| s_with_thinking = cursor.fetchone()[0] | |
| cursor.execute("SELECT COUNT(*) FROM pending_jobs WHERE status='pending'") | |
| pending = cursor.fetchone()[0] | |
| cursor.execute( | |
| "SELECT COALESCE(AVG(generation_time_s), 0) FROM questions WHERE generation_time_s IS NOT NULL" | |
| ) | |
| avg_q_time = cursor.fetchone()[0] | |
| cursor.execute( | |
| "SELECT COALESCE(AVG(generation_time_s), 0) FROM solutions WHERE generation_time_s IS NOT NULL" | |
| ) | |
| avg_s_time = cursor.fetchone()[0] | |
| cursor.execute( | |
| "SELECT id, question_text, model, created_at FROM questions ORDER BY id DESC LIMIT 5" | |
| ) | |
| recent = cursor.fetchall() | |
| runs = [] | |
| try: | |
| cursor.execute( | |
| "SELECT started_at, completed, failed, pending, q_model, s_model FROM run_log ORDER BY id DESC LIMIT 3" | |
| ) | |
| runs = cursor.fetchall() | |
| except sqlite3.OperationalError: | |
| pass | |
| return { | |
| "questions": q_count, | |
| "questions_with_thinking": q_with_thinking, | |
| "solutions": s_count, | |
| "solutions_with_thinking": s_with_thinking, | |
| "pending_jobs": pending, | |
| "total_chunks": q_chunks + s_chunks, | |
| "question_chunks": q_chunks, | |
| "solution_chunks": s_chunks, | |
| "avg_question_time_s": round(avg_q_time, 1), | |
| "avg_solution_time_s": round(avg_s_time, 1), | |
| "recent_questions": [ | |
| {"id": qid, "model": qm, "created_at": qc, "preview": qt[:80]} | |
| for qid, qt, qm, qc in reversed(recent) | |
| ], | |
| "recent_runs": [ | |
| { | |
| "started_at": r[0], | |
| "completed": r[1], | |
| "failed": r[2], | |
| "pending": r[3], | |
| "q_model": r[4], | |
| "s_model": r[5], | |
| } | |
| for r in reversed(runs) | |
| ], | |
| # [ENGINE] Include prompt engine stats | |
| "prompt_engine_stats": prompt_engine.get_stats(), | |
| } | |
| def export_data_json(db_path): | |
| if not os.path.exists(db_path): | |
| return [] | |
| with sqlite3.connect(db_path) as conn: | |
| cursor = conn.cursor() | |
| cursor.execute( | |
| """ | |
| SELECT q.id, q.question_text, q.thinking_trace, q.full_response, | |
| q.model AS q_model, q.created_at, q.generation_time_s, | |
| q.prompt_metadata, | |
| s.id AS s_id, s.stack, s.solution_code, s.thinking_trace AS s_thinking, | |
| s.full_response AS s_full, | |
| s.model AS s_model, s.finish_reason, s.chunk_count, | |
| s.created_at AS s_created, s.generation_time_s AS s_gen_time | |
| FROM questions q | |
| LEFT JOIN solutions s ON q.id = s.question_id | |
| ORDER BY q.id ASC | |
| """ | |
| ) | |
| rows = cursor.fetchall() | |
| questions_map = {} | |
| for ( | |
| q_id, | |
| q_text, | |
| q_thinking, | |
| q_full, | |
| q_model, | |
| q_created, | |
| q_gen_time, | |
| q_prompt_meta, # [ENGINE] | |
| s_id, | |
| stack, | |
| code, | |
| s_thinking, | |
| s_full, | |
| s_model, | |
| s_finish, | |
| s_chunks, | |
| s_created, | |
| s_gen_time, | |
| ) in rows: | |
| if q_id not in questions_map: | |
| questions_map[q_id] = { | |
| "id": q_id, | |
| "question": q_text, | |
| "thinking_trace": q_thinking, | |
| "full_response": q_full, | |
| "model": q_model, | |
| "created_at": q_created, | |
| "generation_time_s": q_gen_time, | |
| "prompt_metadata": json.loads(q_prompt_meta) if q_prompt_meta else None, # [ENGINE] | |
| "solutions": [], | |
| } | |
| if s_id and code: | |
| questions_map[q_id]["solutions"].append( | |
| { | |
| "id": s_id, | |
| "stack": stack, | |
| "code": code, | |
| "thinking_trace": s_thinking, | |
| "full_response": s_full, | |
| "model": s_model, | |
| "finish_reason": s_finish, | |
| "chunk_count": s_chunks, | |
| "created_at": s_created, | |
| "generation_time_s": s_gen_time, | |
| } | |
| ) | |
| return list(questions_map.values()) | |
| def export_files(db_path, output_dir): | |
| if not os.path.exists(db_path): | |
| return {"error": "Database does not exist."} | |
| os.makedirs(output_dir, exist_ok=True) | |
| with sqlite3.connect(db_path) as conn: | |
| cursor = conn.cursor() | |
| cursor.execute( | |
| """ | |
| SELECT q.id, q.question_text, q.thinking_trace, q.full_response, | |
| q.model AS q_model, q.created_at, q.generation_time_s, | |
| q.prompt_metadata, | |
| s.id AS s_id, s.stack, s.solution_code, s.thinking_trace AS s_thinking, | |
| s.full_response AS s_full, | |
| s.model AS s_model, s.finish_reason, s.chunk_count, | |
| s.created_at AS s_created, s.generation_time_s AS s_gen_time | |
| FROM questions q | |
| LEFT JOIN solutions s ON q.id = s.question_id | |
| ORDER BY q.id ASC | |
| """ | |
| ) | |
| rows = cursor.fetchall() | |
| exported = 0 | |
| for row in rows: | |
| ( | |
| q_id, | |
| q_text, | |
| q_thinking, | |
| q_full, | |
| q_model, | |
| q_created, | |
| q_gen_time, | |
| q_prompt_meta, # [ENGINE] | |
| s_id, | |
| stack, | |
| code, | |
| s_thinking, | |
| s_full, | |
| s_model, | |
| s_finish, | |
| s_chunks, | |
| s_created, | |
| s_gen_time, | |
| ) = row | |
| if not s_id or not code: | |
| continue | |
| safe_q = re.sub(r"[^\w\-_\. ]", "_", q_text)[:50].strip() | |
| q_folder = os.path.join(output_dir, f"Q{q_id}_{safe_q}") | |
| os.makedirs(q_folder, exist_ok=True) | |
| with open(os.path.join(q_folder, f"solution_{s_id}.html"), "w") as f: | |
| f.write(code) | |
| if q_thinking: | |
| with open(os.path.join(q_folder, f"q_thinking_{q_id}.txt"), "w") as f: | |
| f.write(q_thinking) | |
| if s_thinking: | |
| with open(os.path.join(q_folder, f"s_thinking_{s_id}.txt"), "w") as f: | |
| f.write(s_thinking) | |
| if q_full: | |
| with open(os.path.join(q_folder, f"q_full_response_{q_id}.txt"), "w") as f: | |
| f.write(q_full) | |
| if s_full: | |
| with open(os.path.join(q_folder, f"s_full_response_{s_id}.txt"), "w") as f: | |
| f.write(s_full) | |
| meta = { | |
| "question_id": q_id, | |
| "question": q_text, | |
| "q_model": q_model, | |
| "has_q_thinking": q_thinking is not None, | |
| "q_generation_time_s": q_gen_time, | |
| "prompt_metadata": json.loads(q_prompt_meta) if q_prompt_meta else None, # [ENGINE] | |
| "solution_id": s_id, | |
| "stack": stack, | |
| "s_model": s_model, | |
| "has_s_thinking": s_thinking is not None, | |
| "finish_reason": s_finish, | |
| "chunk_count": s_chunks, | |
| "s_generation_time_s": s_gen_time, | |
| "q_created": q_created, | |
| "s_created": s_created, | |
| } | |
| with open(os.path.join(q_folder, f"meta_{s_id}.json"), "w") as f: | |
| json.dump(meta, f, indent=2, ensure_ascii=False) | |
| exported += 1 | |
| return {"exported": exported, "output_dir": output_dir} | |
| # === Key Rotation === | |
| def load_keys(): | |
| """Load API keys from /data/keys.txt or environment variable.""" | |
| keys = [] | |
| if os.path.exists(KEYS_FILE): | |
| with open(KEYS_FILE, "r") as f: | |
| keys = [line.strip() for line in f if line.strip()] | |
| env_keys = os.environ.get("FEATHERLESS_API_KEYS", "") | |
| if env_keys: | |
| keys.extend([k.strip() for k in env_keys.split(",") if k.strip()]) | |
| seen = set() | |
| unique_keys = [] | |
| for k in keys: | |
| if k not in seen: | |
| seen.add(k) | |
| unique_keys.append(k) | |
| if not unique_keys: | |
| logger.warning(f"No keys found in {KEYS_FILE} or env. Using dummy key.") | |
| unique_keys = ["dummy_key"] | |
| return unique_keys | |
| _keys = load_keys() | |
| _key_cycle = itertools.cycle(_keys) | |
| _key_lock = threading.Lock() | |
| def get_next_key(): | |
| with _key_lock: | |
| return next(_key_cycle) | |
| def reload_keys(): | |
| global _keys, _key_cycle | |
| _keys = load_keys() | |
| _key_cycle = itertools.cycle(_keys) | |
| return len(_keys) | |
| # === Retry Logic === | |
| def retry_api_call(func, max_retries=5, delay=2): | |
| last_err = None | |
| for attempt in range(max_retries): | |
| if _shutdown.is_set(): | |
| raise InterruptedError("Shutdown requested") | |
| try: | |
| return func() | |
| except RateLimitError as e: | |
| last_err = e | |
| backoff = 120.0 + random.uniform(0, 10) | |
| logger.warning( | |
| f"Rate limit (429). Waiting {backoff:.0f}s... (attempt {attempt+1}/{max_retries})" | |
| ) | |
| time.sleep(backoff) | |
| except APIConnectionError as e: | |
| last_err = e | |
| backoff = min(60, delay * (2**attempt)) + random.uniform(0, 2) | |
| logger.warning( | |
| f"Connection error: {e}. Retry {attempt+1}/{max_retries} in {backoff:.1f}s..." | |
| ) | |
| time.sleep(backoff) | |
| except APIError as e: | |
| last_err = e | |
| if hasattr(e, "status_code") and e.status_code == 429: | |
| backoff = 120.0 + random.uniform(0, 10) | |
| logger.warning(f"Rate limit (via APIError 429). Waiting {backoff:.0f}s...") | |
| elif hasattr(e, "status_code") and e.status_code in (500, 502, 503, 504): | |
| backoff = min(90, delay * (2**attempt)) + random.uniform(0, 5) | |
| logger.warning( | |
| f"Server error {e.status_code}: {e}. Retry in {backoff:.1f}s..." | |
| ) | |
| else: | |
| raise | |
| time.sleep(backoff) | |
| except Exception: | |
| raise | |
| raise last_err | |
| # === Streaming Completion === | |
| def stream_completion(client, model, messages, temperature=0.2, max_tokens=16384): | |
| raw_chunks = [] | |
| content_parts = [] | |
| thinking_parts = [] | |
| tool_call_parts = [] | |
| function_call_parts = [] | |
| finish_reason = None | |
| usage = None | |
| actual_model = model | |
| t0 = time.monotonic() | |
| def make_call(): | |
| return client.chat.completions.create( | |
| model=model, | |
| messages=messages, | |
| temperature=temperature, | |
| max_tokens=max_tokens, | |
| stream=True, | |
| stream_options={"include_usage": True}, | |
| ) | |
| stream = retry_api_call(make_call) | |
| for chunk in stream: | |
| if _shutdown.is_set(): | |
| break | |
| try: | |
| chunk_dict = chunk.model_dump() | |
| except AttributeError: | |
| try: | |
| chunk_dict = json.loads(chunk.json()) | |
| except Exception: | |
| chunk_dict = {"_raw": str(chunk)} | |
| raw_chunks.append(chunk_dict) | |
| if chunk.choices: | |
| choice = chunk.choices[0] | |
| delta = choice.delta | |
| if delta: | |
| if delta.content: | |
| content_parts.append(delta.content) | |
| for attr in ( | |
| "reasoning_content", | |
| "reasoning", | |
| "thinking", | |
| "thought", | |
| "internal_thoughts", | |
| ): | |
| val = getattr(delta, attr, None) | |
| if val: | |
| thinking_parts.append(val) | |
| break | |
| if hasattr(delta, "tool_calls") and delta.tool_calls: | |
| for tc in delta.tool_calls: | |
| try: | |
| tool_call_parts.append(tc.model_dump()) | |
| except Exception: | |
| tool_call_parts.append(str(tc)) | |
| if hasattr(delta, "function_call") and delta.function_call: | |
| try: | |
| function_call_parts.append(delta.function_call.model_dump()) | |
| except Exception: | |
| function_call_parts.append(str(delta.function_call)) | |
| if choice.finish_reason: | |
| finish_reason = choice.finish_reason | |
| if hasattr(chunk, "model") and chunk.model: | |
| actual_model = chunk.model | |
| if hasattr(chunk, "usage") and chunk.usage: | |
| try: | |
| usage = chunk.usage.model_dump() | |
| except AttributeError: | |
| usage = { | |
| "prompt_tokens": getattr(chunk.usage, "prompt_tokens", None), | |
| "completion_tokens": getattr(chunk.usage, "completion_tokens", None), | |
| "total_tokens": getattr(chunk.usage, "total_tokens", None), | |
| } | |
| gen_time = time.monotonic() - t0 | |
| content = "".join(content_parts) | |
| thinking = "".join(thinking_parts) or None | |
| full_parts = [] | |
| if thinking: | |
| full_parts.append(f"<thinking>\n{thinking}\n</thinking>\n\n") | |
| if content: | |
| full_parts.append(content) | |
| if tool_call_parts: | |
| full_parts.append( | |
| f"\n\n<tool_calls>\n{json.dumps(tool_call_parts, indent=2)}\n</tool_calls>" | |
| ) | |
| if function_call_parts: | |
| full_parts.append( | |
| f"\n\n<function_calls>\n{json.dumps(function_call_parts, indent=2)}\n</function_calls>" | |
| ) | |
| full_response = "".join(full_parts) or None | |
| return ( | |
| content, | |
| thinking, | |
| full_response, | |
| raw_chunks, | |
| usage, | |
| finish_reason, | |
| actual_model, | |
| gen_time, | |
| ) | |
| # === Generation Functions === | |
| # [ENGINE] Modified to use the prompt engine with randomized templates and high temperature | |
| def generate_question(client, model): | |
| """ | |
| Generate a question using the prompt engine with randomized templates. | |
| Returns 9 values: (content, thinking, full, chunks, usage, finish, model, gen_time, config) | |
| """ | |
| try: | |
| config = prompt_engine.generate() | |
| except Exception as e: | |
| logger.warning(f"Prompt engine error, using fallback: {e}") | |
| config = PromptConfig( | |
| system_message=( | |
| "You are an expert frontend developer and technical interviewer. " | |
| "Generate a unique, creative UI coding problem in English. " | |
| "Mid-difficulty, under 1000 words." | |
| ), | |
| user_prompt=( | |
| "Generate 1 unique, creative, and detailed question about building " | |
| "a User Interface (UI) using HTML, CSS, and JavaScript. The question " | |
| "must be in English. The problem should be specific enough that a " | |
| "developer can write a complete, self-contained solution.\n\n" | |
| "Output ONLY the problem description. Do not include greetings, " | |
| "numbering, or any other formatting." | |
| ), | |
| temperature=1.0, | |
| max_tokens=4096, | |
| metadata={"fallback": True, "error": str(e)}, | |
| ) | |
| messages = config.to_messages() | |
| topic_preview = config.metadata.get("topic", "unknown")[:60] | |
| logger.info( | |
| f"📝 Streaming question generation " | |
| f"(temp={config.temperature}, tpl={config.metadata.get('template_id', '?')}, " | |
| f"topic={topic_preview}...)..." | |
| ) | |
| result = stream_completion( | |
| client, model, messages, | |
| temperature=config.temperature, | |
| max_tokens=config.max_tokens, | |
| ) | |
| # Return 9 values (original 8 + config) | |
| return (*result, config) | |
| def generate_solution(client, model, question): | |
| prompt = f"""You are an expert frontend developer. I will give you a UI problem. | |
| You must solve it using ONLY: {STACK} | |
| Return ONLY the raw code. Do not include any explanations, greetings, or commentary. | |
| If multiple files are required, use markdown code blocks and CLEARLY indicate the filename before each block (e.g., `**index.html**`). | |
| CRITICAL INSTRUCTIONS FOR REASONING/THINKING: | |
| 1. Keep your thinking/reasoning process brief (under 2000 words). | |
| 2. DO NOT write long draft, or preview any long code inside your thinking block. | |
| Problem: {question} | |
| """ | |
| messages = [{"role": "user", "content": prompt}] | |
| logger.info("🔧 Streaming solution generation...") | |
| return stream_completion(client, model, messages, temperature=0.2, max_tokens=16384) | |
| def health_check(client, model): | |
| try: | |
| resp = client.chat.completions.create( | |
| model=model, | |
| messages=[{"role": "user", "content": "Say 'ok'"}], | |
| max_tokens=5, | |
| temperature=0, | |
| ) | |
| text = resp.choices[0].message.content if resp.choices else "" | |
| logger.info( | |
| f"✅ Health check passed (model={resp.model}, response='{text.strip()}')" | |
| ) | |
| return True | |
| except Exception as e: | |
| logger.error(f"❌ Health check failed: {e}") | |
| return False | |
| def process_single_job( | |
| job_id, question_id, q_text, s_client, s_model, db_path, max_attempts=3 | |
| ): | |
| if _shutdown.is_set(): | |
| return None | |
| mark_job_started(db_path, job_id) | |
| last_error = None | |
| for attempt in range(max_attempts): | |
| if _shutdown.is_set(): | |
| return None | |
| try: | |
| ( | |
| s_code, | |
| s_thinking, | |
| s_full, | |
| s_chunks, | |
| s_usage, | |
| s_finish, | |
| s_model_actual, | |
| gen_time, | |
| ) = generate_solution(s_client, s_model, q_text) | |
| except InterruptedError: | |
| logger.info(f" Job {job_id} interrupted (shutdown).") | |
| return None | |
| except Exception as e: | |
| last_error = e | |
| if attempt < max_attempts - 1: | |
| wait = 5 * (2**attempt) + random.uniform(0, 3) | |
| logger.warning( | |
| f" Job {job_id} (Q{question_id}) attempt {attempt+1} failed: {e}. " | |
| f"Retrying in {wait:.0f}s..." | |
| ) | |
| time.sleep(wait) | |
| continue | |
| else: | |
| logger.error( | |
| f" Job {job_id} (Q{question_id}) failed after {max_attempts} attempts: {e}" | |
| ) | |
| mark_job_failed(db_path, job_id, str(e)) | |
| return None | |
| if not s_code and not s_full: | |
| last_error = "Empty response (no content or full_response)" | |
| if attempt < max_attempts - 1: | |
| logger.warning( | |
| f" Job {job_id} (Q{question_id}): empty solution, retrying..." | |
| ) | |
| time.sleep(3) | |
| continue | |
| else: | |
| logger.warning( | |
| f" Job {job_id} (Q{question_id}): empty solution after {max_attempts} attempts." | |
| ) | |
| mark_job_failed(db_path, job_id, last_error) | |
| return None | |
| save_solution( | |
| db_path, | |
| job_id, | |
| question_id, | |
| s_code or "(empty content — see full_response)", | |
| s_thinking, | |
| s_full, | |
| s_model_actual, | |
| s_finish, | |
| s_usage, | |
| s_chunks, | |
| gen_time, | |
| ) | |
| thinking_info = f", thinking={len(s_thinking)}ch" if s_thinking else "" | |
| logger.info( | |
| f" ✓ Job {job_id} (Q{question_id}) done " | |
| f"({len(s_chunks)} chunks, {gen_time:.1f}s, finish={s_finish}{thinking_info})" | |
| ) | |
| return question_id | |
| mark_job_failed(db_path, job_id, str(last_error)) | |
| return None | |
| # [ENGINE] Modified to unpack 9 values and pass prompt_metadata to save_question | |
| def generate_questions_batch(q_client, q_model, db_path, count): | |
| jobs = [] | |
| for i in range(count): | |
| if _shutdown.is_set(): | |
| break | |
| logger.info(f"📝 Generating question {i+1}/{count}...") | |
| try: | |
| ( | |
| q_text, | |
| q_thinking, | |
| q_full, | |
| q_chunks, | |
| q_usage, | |
| q_finish, | |
| q_model_actual, | |
| gen_time, | |
| config, # [ENGINE] 9th return value | |
| ) = generate_question(q_client, q_model) | |
| except InterruptedError: | |
| logger.info("Shutdown during question generation.") | |
| break | |
| except Exception as e: | |
| logger.error(f"Question {i+1} failed: {e}") | |
| continue | |
| if not q_text and q_full: | |
| q_text = q_full | |
| if not q_text: | |
| logger.warning(f"Question {i+1}: completely empty. Skipping.") | |
| continue | |
| q_text = q_text.strip() | |
| thinking_info = f" (thinking={len(q_thinking)}ch)" if q_thinking else "" | |
| logger.info( | |
| f" Q{i+1} [{gen_time:.1f}s{thinking_info}]: " | |
| f"{q_text[:100]}{'...' if len(q_text) > 100 else ''}" | |
| ) | |
| job_id, question_id = save_question( | |
| db_path, | |
| q_text, | |
| q_thinking, | |
| q_full, | |
| q_model_actual, | |
| q_finish, | |
| q_usage, | |
| q_chunks, | |
| gen_time, | |
| prompt_metadata=config.metadata, # [ENGINE] Pass metadata | |
| ) | |
| jobs.append((job_id, question_id, q_text, 0)) | |
| if i < count - 1 and not _shutdown.is_set(): | |
| time.sleep(random.uniform(0.3, 1.0)) | |
| return jobs | |
| # === Background Generation Runner === | |
| generation_lock = threading.Lock() | |
| generation_running = False | |
| generation_thread: Optional[threading.Thread] = None | |
| def run_generation( | |
| number, workers, q_model, s_model, max_attempts, skip_health_check | |
| ): | |
| """Background generation task.""" | |
| global generation_running | |
| try: | |
| proxy_url = f"http://127.0.0.1:{PORT}/v1" | |
| api_key = "proxy-key" | |
| http_client = httpx.Client( | |
| timeout=httpx.Timeout(connect=30.0, read=600.0, write=60.0, pool=30.0) | |
| ) | |
| q_client = OpenAI(api_key=api_key, base_url=proxy_url, http_client=http_client) | |
| s_client = OpenAI(api_key=api_key, base_url=proxy_url, http_client=http_client) | |
| logger.info(f"Q client → {proxy_url} (model: {q_model})") | |
| logger.info(f"S client → {proxy_url} (model: {s_model})") | |
| if not skip_health_check: | |
| logger.info("Running health check...") | |
| if not health_check(s_client, s_model): | |
| logger.error("Health check failed. Use skip_health_check=true to skip.") | |
| http_client.close() | |
| return | |
| init_db(DB_PATH) | |
| pending = get_pending_jobs(DB_PATH) | |
| if pending: | |
| logger.info( | |
| f"🔄 Found {len(pending)} pending jobs from previous crash — resuming those first!" | |
| ) | |
| new_jobs = generate_questions_batch(q_client, q_model, DB_PATH, number) | |
| logger.info(f"Generated {len(new_jobs)} new questions.") | |
| all_jobs = pending + new_jobs | |
| if not all_jobs: | |
| logger.info("No jobs to process.") | |
| http_client.close() | |
| return | |
| logger.info( | |
| f"📦 Processing {len(all_jobs)} jobs with {workers} concurrent workers..." | |
| ) | |
| completed = 0 | |
| failed = 0 | |
| t_start = time.monotonic() | |
| with ThreadPoolExecutor(max_workers=workers) as pool: | |
| futures = {} | |
| for job_id, q_id, q_text, _attempts in all_jobs: | |
| if _shutdown.is_set(): | |
| break | |
| f = pool.submit( | |
| process_single_job, | |
| job_id, | |
| q_id, | |
| q_text, | |
| s_client, | |
| s_model, | |
| DB_PATH, | |
| max_attempts, | |
| ) | |
| futures[f] = (job_id, q_id) | |
| for future in as_completed(futures): | |
| job_id, q_id = futures[future] | |
| try: | |
| result = future.result() | |
| if result is not None: | |
| completed += 1 | |
| else: | |
| failed += 1 | |
| except Exception as e: | |
| logger.error(f"Job {job_id} (Q{q_id}) crashed: {e}") | |
| mark_job_failed(DB_PATH, job_id, str(e)) | |
| failed += 1 | |
| if _shutdown.is_set(): | |
| logger.warning( | |
| "Shutdown flag set — remaining jobs stay pending for next run." | |
| ) | |
| break | |
| elapsed = time.monotonic() - t_start | |
| remaining = get_pending_jobs(DB_PATH) | |
| log_run( | |
| DB_PATH, | |
| q_model, | |
| s_model, | |
| len(all_jobs), | |
| completed, | |
| failed, | |
| len(remaining), | |
| { | |
| "workers": workers, | |
| "max_attempts": max_attempts, | |
| "proxy": proxy_url, | |
| "elapsed_s": round(elapsed, 1), | |
| }, | |
| ) | |
| logger.info( | |
| f"✅ Completed: {completed} | ❌ Failed: {failed} | ⏳ Pending: {len(remaining)}" | |
| ) | |
| logger.info(f"⏱️ Total time: {elapsed:.1f}s") | |
| http_client.close() | |
| except Exception as e: | |
| logger.error(f"Generation run failed: {e}", exc_info=True) | |
| def run_resume(workers, s_model, max_attempts): | |
| """Resume pending jobs.""" | |
| global generation_running | |
| try: | |
| proxy_url = f"http://127.0.0.1:{PORT}/v1" | |
| http_client = httpx.Client( | |
| timeout=httpx.Timeout(connect=30.0, read=600.0, write=60.0, pool=30.0) | |
| ) | |
| s_client = OpenAI(api_key="proxy-key", base_url=proxy_url, http_client=http_client) | |
| pending = get_pending_jobs(DB_PATH) | |
| if not pending: | |
| logger.info("No pending jobs to resume.") | |
| http_client.close() | |
| return | |
| logger.info(f"📦 Resuming {len(pending)} pending jobs with {workers} workers...") | |
| completed = 0 | |
| failed = 0 | |
| t_start = time.monotonic() | |
| with ThreadPoolExecutor(max_workers=workers) as pool: | |
| futures = {} | |
| for job_id, q_id, q_text, _attempts in pending: | |
| if _shutdown.is_set(): | |
| break | |
| f = pool.submit( | |
| process_single_job, | |
| job_id, | |
| q_id, | |
| q_text, | |
| s_client, | |
| s_model, | |
| DB_PATH, | |
| max_attempts, | |
| ) | |
| futures[f] = (job_id, q_id) | |
| for future in as_completed(futures): | |
| job_id, q_id = futures[future] | |
| try: | |
| result = future.result() | |
| if result is not None: | |
| completed += 1 | |
| else: | |
| failed += 1 | |
| except Exception as e: | |
| logger.error(f"Job {job_id} (Q{q_id}) crashed: {e}") | |
| mark_job_failed(DB_PATH, job_id, str(e)) | |
| failed += 1 | |
| elapsed = time.monotonic() - t_start | |
| remaining = get_pending_jobs(DB_PATH) | |
| log_run( | |
| DB_PATH, | |
| "resume", | |
| s_model, | |
| len(pending), | |
| completed, | |
| failed, | |
| len(remaining), | |
| {"workers": workers, "max_attempts": max_attempts, "elapsed_s": round(elapsed, 1)}, | |
| ) | |
| logger.info(f"Resume done: ✅{completed} ❌{failed} ⏳{len(remaining)}") | |
| http_client.close() | |
| except Exception as e: | |
| logger.error(f"Resume failed: {e}", exc_info=True) | |
| # === Pydantic Models === | |
| class GenerateRequest(BaseModel): | |
| number: int = 20 | |
| workers: int = 20 | |
| q_model: str = "moonshotai/Kimi-K2.6" | |
| s_model: str = "zai-org/GLM-5.2" | |
| max_attempts: int = 3 | |
| skip_health_check: bool = False | |
| class ResumeRequest(BaseModel): | |
| workers: int = 20 | |
| s_model: str = "zai-org/GLM-5.2" | |
| max_attempts: int = 3 | |
| class AddKeyRequest(BaseModel): | |
| key: str | |
| # [ENGINE] New pydantic models for prompt management | |
| class AddTemplateRequest(BaseModel): | |
| template_id: str | |
| template_text: str | |
| class AddTopicRequest(BaseModel): | |
| category: str | |
| topic: str | |
| class AddSystemMessageRequest(BaseModel): | |
| message: str | |
| class TemperatureRangeRequest(BaseModel): | |
| min_temp: float = 0.85 | |
| max_temp: float = 1.15 | |
| class MaxTokensRequest(BaseModel): | |
| max_tokens: int = 4096 | |
| # === FastAPI App === | |
| timeout_config = httpx.Timeout(connect=10.0, read=300.0, write=300.0, pool=10.0) | |
| proxy_http_client: Optional[httpx.AsyncClient] = None | |
| async def lifespan(app: FastAPI): | |
| global proxy_http_client | |
| init_db(DB_PATH) | |
| proxy_http_client = httpx.AsyncClient(timeout=timeout_config) | |
| logger.info(f"Server starting on port {PORT}") | |
| logger.info(f"Data directory: {DATA_DIR}") | |
| logger.info(f"Database: {DB_PATH}") | |
| logger.info(f"Keys file: {KEYS_FILE}") | |
| logger.info(f"Loaded {len(_keys)} API key(s)") | |
| logger.info(f"Featherless base: {FEATHERLESS_API_BASE}") | |
| # [ENGINE] Log engine info | |
| engine_stats = prompt_engine.get_stats() | |
| logger.info( | |
| f"Prompt engine: {engine_stats['available_system_messages']} system msgs, " | |
| f"{engine_stats['available_templates']} templates, " | |
| f"{engine_stats['available_topics']} topics, " | |
| f"temp range {engine_stats['temperature_range']}" | |
| ) | |
| yield | |
| _shutdown.set() | |
| if proxy_http_client: | |
| await proxy_http_client.aclose() | |
| logger.info("Server shutting down...") | |
| app = FastAPI(title="IC Generate API", lifespan=lifespan) | |
| # === Homepage === | |
| async def homepage(): | |
| return """<!DOCTYPE html> | |
| <html> | |
| <head><title>IC Generate API</title> | |
| <style> | |
| body { font-family: monospace; max-width: 900px; margin: 50px auto; padding: 20px; background: #1a1a2e; color: #e0e0e0; } | |
| h1 { color: #00d4ff; } | |
| h2 { color: #00ffaa; margin-top: 30px; } | |
| h3 { color: #ffaa00; margin-top: 20px; } | |
| code { background: #16213e; padding: 2px 6px; border-radius: 3px; color: #ff6b6b; } | |
| ul { line-height: 1.8; } | |
| a { color: #00d4ff; } | |
| </style> | |
| </head> | |
| <body> | |
| <h1>🚀 IC Generate API</h1> | |
| <p>Combined Featherless proxy + UI question/solution generator with full thinking capture.</p> | |
| <p><b>✨ Now with randomized prompt engine (high temperature, 200+ topics, 8 templates)!</b></p> | |
| <h2>Admin Authentication:</h2> | |
| <p>All management routes require the <code>X-Admin-Key</code> header.<br> | |
| Set the <code>ADMIN_API_KEY</code> environment variable in Hugging Face Spaces Secrets.</p> | |
| <h2>Endpoints (Public):</h2> | |
| <ul> | |
| <li><b>POST /v1/{path}</b> — Proxy to Featherless API (key rotation)</li> | |
| </ul> | |
| <h2>Endpoints (Admin-Protected):</h2> | |
| <h3>Generation</h3> | |
| <ul> | |
| <li><b>POST /generate</b> — Start generation run (background)</li> | |
| <li><b>GET /generate/status</b> — Check if generation is running</li> | |
| <li><b>POST /resume</b> — Resume pending jobs</li> | |
| </ul> | |
| <h3>Data & Stats</h3> | |
| <ul> | |
| <li><b>GET /stats</b> — Database statistics (includes engine stats)</li> | |
| <li><b>GET /questions</b> — List questions (paginated)</li> | |
| <li><b>GET /question/{id}</b> — Get question with solutions</li> | |
| <li><b>GET /export/json</b> — Export all data as JSON</li> | |
| <li><b>GET /export/files</b> — Export to /data/exported_code/</li> | |
| <li><b>GET /pending</b> — List pending jobs</li> | |
| <li><b>DELETE /question/{id}</b> — Delete a question</li> | |
| <li><b>GET /health</b> — Health check</li> | |
| </ul> | |
| <h3>API Keys</h3> | |
| <ul> | |
| <li><b>POST /keys/add</b> — Add a new Featherless API key</li> | |
| <li><b>GET /keys/list</b> — List active Featherless API keys</li> | |
| <li><b>DELETE /keys/{key}</b> — Delete a Featherless API key</li> | |
| <li><b>POST /keys/reload</b> — Reload keys from file</li> | |
| </ul> | |
| <h3>✨ Prompt Engine Management</h3> | |
| <ul> | |
| <li><b>GET /prompts/stats</b> — Engine statistics (generations, unique combos, template usage)</li> | |
| <li><b>GET /prompts/preview</b> — Preview a randomly generated prompt (no API call)</li> | |
| <li><b>GET /prompts/templates</b> — List all prompt templates</li> | |
| <li><b>POST /prompts/templates</b> — Add a custom template</li> | |
| <li><b>DELETE /prompts/templates/{id}</b> — Remove a custom template</li> | |
| <li><b>GET /prompts/topics</b> — List all topics</li> | |
| <li><b>POST /prompts/topics</b> — Add a custom topic</li> | |
| <li><b>DELETE /prompts/topics/{topic}</b> — Remove a custom topic</li> | |
| <li><b>GET /prompts/system-messages</b> — List all system messages</li> | |
| <li><b>POST /prompts/system-messages</b> — Add a custom system message</li> | |
| <li><b>DELETE /prompts/system-messages/{index}</b> — Remove a custom system message</li> | |
| <li><b>POST /prompts/temperature</b> — Set temperature range</li> | |
| <li><b>POST /prompts/max-tokens</b> — Set max tokens</li> | |
| <li><b>POST /prompts/reset</b> — Reset engine deduplication state</li> | |
| </ul> | |
| </body> | |
| </html>""" | |
| # === Proxy Endpoint (Public) === | |
| async def proxy(request: Request, path: str): | |
| """Proxy requests to Featherless API with key rotation.""" | |
| url = f"{FEATHERLESS_API_BASE}/{path}" | |
| api_key = get_next_key() | |
| headers = dict(request.headers) | |
| headers.pop("host", None) | |
| headers.pop("content-length", None) | |
| headers.pop("content-encoding", None) | |
| headers.pop("transfer-encoding", None) | |
| headers["authorization"] = f"Bearer {api_key}" | |
| body = await request.body() | |
| req = proxy_http_client.build_request( | |
| method=request.method, | |
| url=url, | |
| headers=headers, | |
| content=body if body else None, | |
| params=request.query_params, | |
| ) | |
| try: | |
| response = await proxy_http_client.send(req, stream=True) | |
| except httpx.RequestError as e: | |
| raise HTTPException(status_code=502, detail=f"Proxy error: {str(e)}") | |
| async def generate(): | |
| try: | |
| async for chunk in response.aiter_bytes(): | |
| yield chunk | |
| except httpx.ReadTimeout: | |
| pass | |
| finally: | |
| await response.aclose() | |
| response_headers = dict(response.headers) | |
| response_headers.pop("content-encoding", None) | |
| response_headers.pop("transfer-encoding", None) | |
| response_headers.pop("content-length", None) | |
| return StreamingResponse( | |
| generate(), | |
| status_code=response.status_code, | |
| headers=response_headers, | |
| media_type=response.headers.get("content-type"), | |
| ) | |
| # === Generate Endpoint (Admin) === | |
| async def generate(req: GenerateRequest): | |
| """Start a generation run in the background.""" | |
| global generation_running, generation_thread | |
| with generation_lock: | |
| if generation_running: | |
| raise HTTPException( | |
| status_code=409, detail="A generation run is already in progress." | |
| ) | |
| generation_running = True | |
| def task_wrapper(): | |
| global generation_running | |
| try: | |
| run_generation( | |
| req.number, | |
| req.workers, | |
| req.q_model, | |
| req.s_model, | |
| req.max_attempts, | |
| req.skip_health_check, | |
| ) | |
| finally: | |
| with generation_lock: | |
| generation_running = False | |
| generation_thread = threading.Thread(target=task_wrapper, daemon=True) | |
| generation_thread.start() | |
| return { | |
| "status": "started", | |
| "message": f"Generating {req.number} Q&A pairs with {req.workers} workers.", | |
| "config": { | |
| "number": req.number, | |
| "workers": req.workers, | |
| "q_model": req.q_model, | |
| "s_model": req.s_model, | |
| "max_attempts": req.max_attempts, | |
| "skip_health_check": req.skip_health_check, | |
| }, | |
| } | |
| async def generation_status(): | |
| """Check if a generation run is in progress.""" | |
| return {"running": generation_running} | |
| # === Resume Endpoint (Admin) === | |
| async def resume(req: ResumeRequest): | |
| """Resume pending jobs without generating new questions.""" | |
| global generation_running, generation_thread | |
| with generation_lock: | |
| if generation_running: | |
| raise HTTPException( | |
| status_code=409, detail="A generation run is already in progress." | |
| ) | |
| generation_running = True | |
| pending = get_pending_jobs(DB_PATH) | |
| if not pending: | |
| with generation_lock: | |
| generation_running = False | |
| return {"status": "no_pending", "message": "No pending jobs to resume."} | |
| def task_wrapper(): | |
| global generation_running | |
| try: | |
| run_resume(req.workers, req.s_model, req.max_attempts) | |
| finally: | |
| with generation_lock: | |
| generation_running = False | |
| generation_thread = threading.Thread(target=task_wrapper, daemon=True) | |
| generation_thread.start() | |
| return { | |
| "status": "started", | |
| "message": f"Resuming {len(pending)} pending jobs with {req.workers} workers.", | |
| } | |
| async def fix_missing_solutions(): | |
| """Mark questions without solutions as pending so they can be generated.""" | |
| fixed_count = 0 | |
| with sqlite3.connect(DB_PATH) as conn: | |
| c = conn.cursor() | |
| # Find questions that don't have solutions | |
| c.execute(''' | |
| SELECT q.id | |
| FROM questions q | |
| LEFT JOIN solutions s ON q.id = s.question_id | |
| WHERE s.id IS NULL | |
| ''') | |
| missing_qs = c.fetchall() | |
| for q in missing_qs: | |
| q_id = q[0] | |
| # Check if it exists in pending_jobs | |
| c.execute("SELECT id, status FROM pending_jobs WHERE question_id=?", (q_id,)) | |
| job = c.fetchone() | |
| if job: | |
| if job[1] != 'pending': | |
| c.execute("UPDATE pending_jobs SET status='pending', attempts=0, last_error=NULL, completed_at=NULL WHERE id=?", (job[0],)) | |
| fixed_count += 1 | |
| else: | |
| c.execute("INSERT INTO pending_jobs (question_id, status) VALUES (?, 'pending')", (q_id,)) | |
| fixed_count += 1 | |
| conn.commit() | |
| return { | |
| "status": "success", | |
| "message": f"Fixed {fixed_count} pending jobs for questions without solutions." | |
| } | |
| # === Stats Endpoint (Admin) === | |
| async def stats(): | |
| return get_stats_dict(DB_PATH) | |
| # === Questions Endpoints (Admin) === | |
| async def list_questions(limit: int = Query(20, ge=1, le=100), offset: int = Query(0, ge=0)): | |
| """List questions with pagination.""" | |
| with sqlite3.connect(DB_PATH) as conn: | |
| cursor = conn.cursor() | |
| cursor.execute("SELECT COUNT(*) FROM questions") | |
| total = cursor.fetchone()[0] | |
| cursor.execute( | |
| """ | |
| SELECT q.id, q.question_text, q.model, q.created_at, q.generation_time_s, | |
| (SELECT COUNT(*) FROM solutions WHERE question_id = q.id) AS solution_count | |
| FROM questions q | |
| ORDER BY q.id DESC | |
| LIMIT ? OFFSET ? | |
| """, | |
| (limit, offset), | |
| ) | |
| rows = cursor.fetchall() | |
| return { | |
| "total": total, | |
| "limit": limit, | |
| "offset": offset, | |
| "questions": [ | |
| { | |
| "id": r[0], | |
| "question": r[1], | |
| "model": r[2], | |
| "created_at": r[3], | |
| "generation_time_s": r[4], | |
| "solution_count": r[5], | |
| } | |
| for r in rows | |
| ], | |
| } | |
| async def get_question(question_id: int): | |
| """Get a specific question with its solutions.""" | |
| with sqlite3.connect(DB_PATH) as conn: | |
| conn.row_factory = sqlite3.Row | |
| cursor = conn.cursor() | |
| cursor.execute("SELECT * FROM questions WHERE id=?", (question_id,)) | |
| q = cursor.fetchone() | |
| if not q: | |
| raise HTTPException( | |
| status_code=404, detail=f"Question {question_id} not found." | |
| ) | |
| cursor.execute("SELECT * FROM solutions WHERE question_id=?", (question_id,)) | |
| solutions = cursor.fetchall() | |
| question = dict(q) | |
| question.pop("raw_chunks_json", None) | |
| question["usage"] = json.loads(question.pop("usage_json", "null") or "null") | |
| # [ENGINE] Parse prompt_metadata | |
| question["prompt_metadata"] = json.loads( | |
| question.pop("prompt_metadata", "null") or "null" | |
| ) | |
| question["solutions"] = [] | |
| for s in solutions: | |
| sol = dict(s) | |
| sol.pop("raw_chunks_json", None) | |
| sol["usage"] = json.loads(sol.pop("usage_json", "null") or "null") | |
| question["solutions"].append(sol) | |
| return question | |
| # === Export Endpoints (Admin) === | |
| async def export_json(): | |
| data = export_data_json(DB_PATH) | |
| return JSONResponse(content=data) | |
| async def export_files_endpoint(): | |
| result = export_files(DB_PATH, EXPORT_DIR) | |
| return result | |
| # === Pending Jobs Endpoint (Admin) === | |
| async def pending_jobs(): | |
| jobs = get_pending_jobs(DB_PATH) | |
| return { | |
| "count": len(jobs), | |
| "jobs": [ | |
| { | |
| "job_id": j[0], | |
| "question_id": j[1], | |
| "attempts": j[3], | |
| "question_preview": j[2][:100] if j[2] else "", | |
| } | |
| for j in jobs | |
| ], | |
| } | |
| # === Delete Question Endpoint (Admin) === | |
| async def delete_question(question_id: int): | |
| success = delete_question_db(DB_PATH, question_id) | |
| if success: | |
| return {"status": "deleted", "question_id": question_id} | |
| else: | |
| raise HTTPException( | |
| status_code=404, detail=f"Question {question_id} not found." | |
| ) | |
| # === Health Endpoint (Admin) === | |
| async def health(): | |
| return { | |
| "status": "ok", | |
| "data_dir": DATA_DIR, | |
| "db_path": DB_PATH, | |
| "db_exists": os.path.exists(DB_PATH), | |
| "keys_loaded": len(_keys), | |
| "keys_file": KEYS_FILE, | |
| "keys_file_exists": os.path.exists(KEYS_FILE), | |
| "featherless_base": FEATHERLESS_API_BASE, | |
| "port": PORT, | |
| "generation_running": generation_running, | |
| } | |
| # === Keys Management Endpoints (Admin) === | |
| async def keys_info(): | |
| return { | |
| "count": len(_keys), | |
| "keys_file": KEYS_FILE, | |
| "keys_file_exists": os.path.exists(KEYS_FILE), | |
| } | |
| async def list_keys(): | |
| """List all active Featherless API keys.""" | |
| return {"count": len(_keys), "keys": _keys} | |
| async def add_key(req: AddKeyRequest): | |
| """Add a new Featherless API key to the rotation.""" | |
| global _keys, _key_cycle | |
| key = req.key.strip() | |
| if not key: | |
| raise HTTPException(status_code=400, detail="Key cannot be empty") | |
| with open(KEYS_FILE, "a") as f: | |
| f.write(f"{key}\n") | |
| if key not in _keys: | |
| _keys.append(key) | |
| _key_cycle = itertools.cycle(_keys) | |
| return {"status": "success", "message": "Key added successfully", "total_keys": len(_keys)} | |
| async def delete_key(key: str): | |
| """Remove a Featherless API key from the rotation.""" | |
| global _keys, _key_cycle | |
| if key in _keys: | |
| _keys.remove(key) | |
| _key_cycle = itertools.cycle(_keys) | |
| with open(KEYS_FILE, "w") as f: | |
| for k in _keys: | |
| f.write(f"{k}\n") | |
| return {"status": "success", "message": "Key deleted", "total_keys": len(_keys)} | |
| raise HTTPException(status_code=404, detail="Key not found") | |
| async def reload_keys_endpoint(): | |
| count = reload_keys() | |
| return {"status": "reloaded", "count": count} | |
| # ============================================================================ | |
| # [ENGINE] Prompt Engine Management Endpoints (Admin) | |
| # ============================================================================ | |
| async def prompt_engine_stats(): | |
| """Get prompt engine statistics.""" | |
| return prompt_engine.get_stats() | |
| async def prompt_preview(): | |
| """Preview a randomly generated prompt without consuming dedup state or calling the API.""" | |
| config = prompt_engine.peek() | |
| return config.to_dict() | |
| async def list_templates(): | |
| """List all available prompt templates (built-in + custom).""" | |
| templates = prompt_engine.get_all_templates() | |
| return {"count": len(templates), "templates": templates} | |
| async def add_template(req: AddTemplateRequest): | |
| """Add a custom prompt template.""" | |
| success = prompt_engine.add_template(req.template_id, req.template_text) | |
| if success: | |
| return {"status": "success", "message": f"Template '{req.template_id}' added."} | |
| raise HTTPException( | |
| status_code=400, | |
| detail=f"Failed to add template. ID may be duplicate or fields empty." | |
| ) | |
| async def remove_template(template_id: str): | |
| """Remove a custom prompt template by ID.""" | |
| success = prompt_engine.remove_template(template_id) | |
| if success: | |
| return {"status": "success", "message": f"Template '{template_id}' removed."} | |
| raise HTTPException(status_code=404, detail=f"Template '{template_id}' not found in custom templates.") | |
| async def list_topics(): | |
| """List all available topics (built-in + custom).""" | |
| topics = prompt_engine.get_all_topics() | |
| return {"count": len(topics), "topics": topics} | |
| async def add_topic(req: AddTopicRequest): | |
| """Add a custom topic.""" | |
| success = prompt_engine.add_topic(req.category, req.topic) | |
| if success: | |
| return {"status": "success", "message": f"Topic added to category '{req.category}'."} | |
| raise HTTPException( | |
| status_code=400, | |
| detail="Failed to add topic. It may be a duplicate or fields are empty." | |
| ) | |
| async def remove_topic(topic: str): | |
| """Remove a custom topic by topic text.""" | |
| success = prompt_engine.remove_topic(topic) | |
| if success: | |
| return {"status": "success", "message": f"Topic '{topic[:50]}' removed."} | |
| raise HTTPException(status_code=404, detail=f"Topic not found in custom topics.") | |
| async def list_system_messages(): | |
| """List all system messages (built-in + custom).""" | |
| messages = prompt_engine.get_all_system_messages() | |
| return {"count": len(messages), "system_messages": messages} | |
| async def add_system_message(req: AddSystemMessageRequest): | |
| """Add a custom system message.""" | |
| success = prompt_engine.add_system_message(req.message) | |
| if success: | |
| return {"status": "success", "message": "System message added."} | |
| raise HTTPException( | |
| status_code=400, | |
| detail="Failed to add system message. It may be a duplicate or empty." | |
| ) | |
| async def remove_system_message(index: int): | |
| """Remove a custom system message by index (custom index only).""" | |
| success = prompt_engine.remove_system_message(index) | |
| if success: | |
| return {"status": "success", "message": f"Custom system message at index {index} removed."} | |
| raise HTTPException(status_code=404, detail=f"No custom system message at index {index}.") | |
| async def set_temperature_range(req: TemperatureRangeRequest): | |
| """Set the temperature range for question generation.""" | |
| prompt_engine.set_temperature_range(req.min_temp, req.max_temp) | |
| stats = prompt_engine.get_stats() | |
| return { | |
| "status": "success", | |
| "message": f"Temperature range set to [{req.min_temp}, {req.max_temp}].", | |
| "temperature_range": stats["temperature_range"], | |
| } | |
| async def set_max_tokens(req: MaxTokensRequest): | |
| """Set max tokens for question generation.""" | |
| prompt_engine.set_max_tokens(req.max_tokens) | |
| return { | |
| "status": "success", | |
| "message": f"Max tokens set to {req.max_tokens}.", | |
| "max_tokens": req.max_tokens, | |
| } | |
| async def reset_prompt_engine(): | |
| """Reset the prompt engine's deduplication state.""" | |
| prompt_engine.reset_state() | |
| return {"status": "success", "message": "Prompt engine state reset."} | |
| # === Main === | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=PORT) | |