From d92fa26ccb3060d05d00a36e6970172e27715e4e Mon Sep 17 00:00:00 2001 From: shmukit Date: Thu, 7 May 2026 22:55:06 +0600 Subject: [PATCH 1/8] feat(database): refactor PostgreSQL connection handling and introduce job repository - Replaced direct psycopg2 connection calls with a centralized connect_real_dict function for improved connection management. - Added a new JobRepository interface with in-memory and SQLite implementations for better job tracking in async extract jobs. - Updated various services to utilize the new job repository and handle job states more effectively. - Introduced media encoding utilities for PNG images to streamline data URI generation. - Enhanced question mapping and schema definitions for better integration with the extraction pipeline. --- content_gen/adapters/postgres_adapter.py | 6 +- content_gen/core/media_encoding.py | 15 ++ content_gen/core/model_router.py | 32 +++- content_gen/core/question_mapping.py | 81 +++++++++ content_gen/core/schema.py | 27 ++- content_gen/core/schemas.py | 8 + content_gen/db/__init__.py | 5 + content_gen/db/session.py | 22 +++ .../scripts/pipeline/pipeline_orchestrator.py | 9 +- content_gen/tests/test_adapters.py | 4 +- content_gen/tests/test_question_mapping.py | 21 +++ content_gen/tests/test_router.py | 6 +- qc_viewer/router_v1.py | 60 ++++--- qc_viewer/routers/automation.py | 18 +- qc_viewer/routers/questions.py | 82 ++------- qc_viewer/services/automation_pipeline.py | 169 +++--------------- qc_viewer/services/db_service.py | 13 +- qc_viewer/services/draft_store.py | 14 +- qc_viewer/services/job_repository.py | 119 ++++++++++++ qc_viewer/services/legacy_question_payload.py | 151 ++++++++++++++++ qc_viewer/services/question_repository.py | 99 ++++++++++ .../services/tests/test_job_repository.py | 50 ++++++ .../tests/test_legacy_question_payload.py | 79 ++++++++ qc_viewer/services/tests/test_pipeline.py | 28 +-- 24 files changed, 814 insertions(+), 304 deletions(-) create mode 100644 content_gen/core/media_encoding.py create mode 100644 content_gen/core/question_mapping.py create mode 100644 content_gen/db/__init__.py create mode 100644 content_gen/db/session.py create mode 100644 content_gen/tests/test_question_mapping.py create mode 100644 qc_viewer/services/job_repository.py create mode 100644 qc_viewer/services/legacy_question_payload.py create mode 100644 qc_viewer/services/question_repository.py create mode 100644 qc_viewer/services/tests/test_job_repository.py create mode 100644 qc_viewer/services/tests/test_legacy_question_payload.py diff --git a/content_gen/adapters/postgres_adapter.py b/content_gen/adapters/postgres_adapter.py index 630d477..13de126 100644 --- a/content_gen/adapters/postgres_adapter.py +++ b/content_gen/adapters/postgres_adapter.py @@ -1,10 +1,10 @@ import re -import psycopg2 import uuid from datetime import datetime, timezone from typing import List, Dict, Optional, Any from .base import BaseStorageAdapter +from ..db.session import connect_real_dict from ..core.schemas import ProcessedQuestion, Flashcard from ..core.config_loader import ConfigLoader from ..core.config_schema import EdmateConfig @@ -20,7 +20,7 @@ class PostgresStorageAdapter(BaseStorageAdapter): def __init__(self, connection_string: str, edmate_config: Optional[EdmateConfig] = None): self.conn_str = connection_string - self.conn = psycopg2.connect(connection_string) + self.conn = connect_real_dict(connection_string, connect_timeout=30) self.cur = self.conn.cursor() self._workspace = (edmate_config or ConfigLoader.load_config()).workspace @@ -31,7 +31,7 @@ def initialize_schema(connection_string: str, edmate_config: Optional[EdmateConf Table names come from edmate_config workspace.target_tables when set; otherwise a legacy Cambridge-style multi-table set is created for backward compatibility. """ - conn = psycopg2.connect(connection_string) + conn = connect_real_dict(connection_string, connect_timeout=30) cur = conn.cursor() try: # Shared tables diff --git a/content_gen/core/media_encoding.py b/content_gen/core/media_encoding.py new file mode 100644 index 0000000..4ca73e5 --- /dev/null +++ b/content_gen/core/media_encoding.py @@ -0,0 +1,15 @@ +"""Shared binary → data URI helpers (PNG).""" + +from __future__ import annotations + +import base64 +from pathlib import Path + + +def png_bytes_to_data_uri(png_bytes: bytes) -> str: + return "data:image/png;base64," + base64.b64encode(png_bytes).decode("utf-8") + + +def png_file_to_data_uri(path: Path) -> str: + with open(path, "rb") as f: + return png_bytes_to_data_uri(f.read()) diff --git a/content_gen/core/model_router.py b/content_gen/core/model_router.py index 09addf1..e6d30da 100644 --- a/content_gen/core/model_router.py +++ b/content_gen/core/model_router.py @@ -23,10 +23,18 @@ class ModelRoutingEngine: Includes an automatic 'Economic Kill-Switch' based on configured budget. """ - def __init__(self, config: Optional[EdmateConfig] = None): + def __init__( + self, + config: Optional[EdmateConfig] = None, + *, + api_key: Optional[str] = None, + ): # Load config from YAML/JSON if not provided self.config = config or CoreConfig.load_from_yaml() self.tracker = MetricsTracker() + # Per-request / BYOK key passed to litellm (never mutate os.environ). + _k = (api_key or "").strip() + self._api_key: Optional[str] = _k or None def generate_content( self, @@ -34,7 +42,8 @@ def generate_content( task_type: str = "generation", system_prompt: Optional[str] = None, images: Optional[List[str]] = None, - json_mode: bool = False + json_mode: bool = False, + api_key: Optional[str] = None, ) -> str: """ Routes the task to the appropriate model based on task_type. @@ -70,13 +79,18 @@ def generate_content( else: messages.append({"role": "user", "content": prompt}) - # 3. Execute call - response = litellm.completion( - model=model, - messages=messages, - response_format={"type": "json_object"} if json_mode else None, - timeout=120 if task_type == "extraction" else 180 - ) + # 3. Execute call (optional api_key: per-call override, else instance BYOK) + _key = (api_key or "").strip() or self._api_key + _kwargs: Dict[str, Any] = { + "model": model, + "messages": messages, + "timeout": 120 if task_type == "extraction" else 180, + } + if json_mode: + _kwargs["response_format"] = {"type": "json_object"} + if _key: + _kwargs["api_key"] = _key + response = litellm.completion(**_kwargs) # 4. Log usage (The Analytics Layer) self.tracker.log_usage(response) diff --git a/content_gen/core/question_mapping.py b/content_gen/core/question_mapping.py new file mode 100644 index 0000000..bc5e8f7 --- /dev/null +++ b/content_gen/core/question_mapping.py @@ -0,0 +1,81 @@ +""" +Boundary mappers between pipeline models and public / export DTOs. + +**Canonical pipeline model:** :class:`content_gen.core.schemas.ProcessedQuestion` + — produced by PDF extraction + :class:`content_gen.scripts.processing.content_generator.ContentGenerator`. + Field names are oriented to processing (``question_text``, ``options`` as ``dict[str, str]``, + ``Flashcard`` with ``front_text`` / ``back_text``). + +**Export / API sketch:** :class:`content_gen.core.schema.EdmateQuestion` + — structured Lab_QA JSON with nested ``Metadata``, ``Option`` list, ``Explanations``, etc. + Use when emitting versioned JSON for external consumers. + +Do not merge the two modules into one model in a single step; map at the edges only. +""" + +from __future__ import annotations + +import uuid +from typing import List, Optional + +from content_gen.core.schema import ( + EdmateQuestion, + Explanations, + Flashcard as ExportFlashcard, + Metadata, + Option, +) +from content_gen.core.schemas import Flashcard as PipelineFlashcard, ProcessedQuestion + + +def processed_question_to_edmate_question( + q: ProcessedQuestion, + *, + curriculum: str = "General", + topic: str = "General", +) -> EdmateQuestion: + """ + Best-effort map from pipeline output to :class:`EdmateQuestion`. + + Unknown fields use safe defaults so the object validates. + """ + opts_in = q.options or {} + letters = ("A", "B", "C", "D") + correct = {str(c).upper() for c in (q.correct_options or []) if c} + options_list: List[Option] = [] + for label in letters: + text = str(opts_in.get(label, "") or "").strip() + options_list.append( + Option( + id=label, + text=text or "(empty)", + is_correct=label in correct, + explanation="", + ) + ) + + core = (q.metadata or {}).get("core_concept_generated") or "" + body = (q.explanation_body or "").strip() + final_letter = next(iter(correct), "") if correct else "" + + fc_out: List[ExportFlashcard] = [] + for fc in q.flashcards or []: + if isinstance(fc, PipelineFlashcard): + fc_out.append(ExportFlashcard(front=fc.front_text, back=fc.back_text)) + else: + fc_out.append(ExportFlashcard(front=str(fc), back="")) + + return EdmateQuestion( + id=uuid.uuid4(), + metadata=Metadata(curriculum=curriculum, subject=q.subject, topic=topic), + content={"type": "mcq"}, + question_text=q.question_text, + options=options_list, + explanations=Explanations( + core_concept=str(core) or topic, + detailed_logic=q.option_wise_explanation or body, + final_answer_display=final_letter or "", + ), + flashcards=fc_out, + media=[], + ) diff --git a/content_gen/core/schema.py b/content_gen/core/schema.py index cf39398..fc42b65 100644 --- a/content_gen/core/schema.py +++ b/content_gen/core/schema.py @@ -1,5 +1,12 @@ +""" +Public / export-oriented Lab_QA question shape (nested metadata, option list, etc.). + +The live pipeline uses :mod:`content_gen.core.schemas` (:class:`ProcessedQuestion`). +Map between the two with :mod:`content_gen.core.question_mapping`. +""" + from typing import List, Optional, Literal -from pydantic import BaseModel, Field, UUID4 +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, UUID4 from enum import Enum class DifficultyLevel(str, Enum): @@ -60,19 +67,23 @@ class MediaItem(BaseModel): class EdmateQuestion(BaseModel): """The standard v1.0.0 Edmate Lab_QA Question Schema""" - schema_version: str = Field("1.0.0", alias="$schema_version") + + model_config = ConfigDict(use_enum_values=True, populate_by_name=True) + + # Python attribute is schema_version; JSON may use "$schema_version" (Lab_QA export). + schema_version: str = Field( + default="1.0.0", + validation_alias=AliasChoices("schema_version", "$schema_version"), + serialization_alias="$schema_version", + ) id: Optional[UUID4] = None metadata: Metadata content: dict = Field(..., description="Type-specific content based on QuestionType") # Note: Using dict for content to allow flexible types, can be further specialized - + # Common fields that can be promoted to direct child for convenience - question_text: str + question_text: str options: Optional[List[Option]] = None explanations: Explanations flashcards: List[Flashcard] = [] media: List[MediaItem] = [] - - class Config: - use_enum_values = True - populate_by_name = True diff --git a/content_gen/core/schemas.py b/content_gen/core/schemas.py index ff54470..40e97ee 100644 --- a/content_gen/core/schemas.py +++ b/content_gen/core/schemas.py @@ -1,3 +1,11 @@ +""" +Pipeline DTOs (extraction + generation). + +Canonical in-process model for PDF → LLM → DB flows is :class:`ProcessedQuestion`. +For versioned export JSON see :mod:`content_gen.core.schema` and +:func:`content_gen.core.question_mapping.processed_question_to_edmate_question`. +""" + from typing import List, Dict, Optional, Any from pydantic import BaseModel, Field from datetime import datetime diff --git a/content_gen/db/__init__.py b/content_gen/db/__init__.py new file mode 100644 index 0000000..b4736a4 --- /dev/null +++ b/content_gen/db/__init__.py @@ -0,0 +1,5 @@ +"""Shared database helpers for Edmate Python packages.""" + +from content_gen.db.session import connect_real_dict + +__all__ = ["connect_real_dict"] diff --git a/content_gen/db/session.py b/content_gen/db/session.py new file mode 100644 index 0000000..90eff1b --- /dev/null +++ b/content_gen/db/session.py @@ -0,0 +1,22 @@ +"""Centralized psycopg2 connection factory (RealDictCursor, timeouts).""" + +from __future__ import annotations + +from typing import Any + +import psycopg2 +from psycopg2.extras import RealDictCursor + + +def connect_real_dict(dsn: str, *, connect_timeout: int = 30, **kwargs: Any): + """ + Open a PostgreSQL connection with dict-like rows. + + Extra kwargs are forwarded to psycopg2.connect (e.g. application_name). + """ + return psycopg2.connect( + dsn, + cursor_factory=RealDictCursor, + connect_timeout=connect_timeout, + **kwargs, + ) diff --git a/content_gen/scripts/pipeline/pipeline_orchestrator.py b/content_gen/scripts/pipeline/pipeline_orchestrator.py index 7ba6aaa..9da0fbc 100644 --- a/content_gen/scripts/pipeline/pipeline_orchestrator.py +++ b/content_gen/scripts/pipeline/pipeline_orchestrator.py @@ -4,9 +4,9 @@ from pathlib import Path from typing import List, Dict, Optional, cast import argparse -import base64 # Import modular core +from content_gen.core.media_encoding import png_file_to_data_uri from content_gen.core.model_router import ModelRoutingEngine from content_gen.core.config_schema import ExtractionEngine from content_gen.adapters.postgres_adapter import PostgresStorageAdapter @@ -32,9 +32,6 @@ def __init__( self.storage_bucket = storage_bucket self.db_connection = db_connection self.router = router or ModelRoutingEngine() - - # Initialize components - self.router = router or ModelRoutingEngine() self.generator = ContentGenerator(router=self.router) self.storage = ( PostgresStorageAdapter(db_connection, edmate_config=self.router.config) @@ -73,9 +70,7 @@ def __init__( def _convert_to_base64(self, image_path: Path) -> str: """Converts an image file to a base64 Data URI.""" try: - with open(image_path, "rb") as f: - encoded = base64.b64encode(f.read()).decode("utf-8") - return f"data:image/png;base64,{encoded}" + return png_file_to_data_uri(image_path) except Exception as e: print(f"⚠️ Failed to base64 encode {image_path}: {e}") return "" diff --git a/content_gen/tests/test_adapters.py b/content_gen/tests/test_adapters.py index 3104ecb..e607d49 100644 --- a/content_gen/tests/test_adapters.py +++ b/content_gen/tests/test_adapters.py @@ -15,7 +15,7 @@ def test_storage_adapter_save_question(): # Initialize adapter with mocked connection empty_ws = EdmateConfig(workspace=WorkspaceConfig()) - with patch("psycopg2.connect", return_value=mock_conn): + with patch("content_gen.adapters.postgres_adapter.connect_real_dict", return_value=mock_conn): adapter = PostgresStorageAdapter( "postgres://user:pass@host:5432/db", edmate_config=empty_ws ) @@ -54,7 +54,7 @@ def test_storage_adapter_save_flashcards(): mock_cur = MagicMock() empty_ws = EdmateConfig(workspace=WorkspaceConfig()) - with patch("psycopg2.connect", return_value=mock_conn): + with patch("content_gen.adapters.postgres_adapter.connect_real_dict", return_value=mock_conn): adapter = PostgresStorageAdapter( "postgres://user:pass@host:5432/db", edmate_config=empty_ws ) diff --git a/content_gen/tests/test_question_mapping.py b/content_gen/tests/test_question_mapping.py new file mode 100644 index 0000000..6b24836 --- /dev/null +++ b/content_gen/tests/test_question_mapping.py @@ -0,0 +1,21 @@ +from content_gen.core.question_mapping import processed_question_to_edmate_question +from content_gen.core.schemas import ProcessedQuestion + + +def test_processed_to_edmate_basic(): + q = ProcessedQuestion( + question_number=1, + question_text="What is 2+2?", + options={"A": "3", "B": "4", "C": "5", "D": "6"}, + correct_options=["B"], + subject="Math", + explanation_body="It is four.", + option_wise_explanation="Option B is correct.", + ) + out = processed_question_to_edmate_question(q, curriculum="IGCSE", topic="Arithmetic") + assert out.metadata.subject == "Math" + assert out.metadata.curriculum == "IGCSE" + assert out.question_text == "What is 2+2?" + labels = {o.id: o.text for o in (out.options or [])} + assert labels["B"] == "4" + assert any(o.is_correct for o in (out.options or []) if o.id == "B") diff --git a/content_gen/tests/test_router.py b/content_gen/tests/test_router.py index 92c67cf..1e6b896 100644 --- a/content_gen/tests/test_router.py +++ b/content_gen/tests/test_router.py @@ -47,8 +47,7 @@ def test_router_task_routing(mock_completion): mock_completion.assert_called_with( model="test/extraction", messages=[{"role": "user", "content": "extract this"}], - response_format=None, - timeout=120 + timeout=120, ) # Test default generation task @@ -56,7 +55,6 @@ def test_router_task_routing(mock_completion): mock_completion.assert_called_with( model="test/generation", messages=[{"role": "user", "content": "generate this"}], - response_format=None, - timeout=180 + timeout=180, ) diff --git a/qc_viewer/router_v1.py b/qc_viewer/router_v1.py index e00abf7..ffb14e6 100644 --- a/qc_viewer/router_v1.py +++ b/qc_viewer/router_v1.py @@ -9,12 +9,10 @@ from content_gen.scripts.pipeline.pipeline_orchestrator import PipelineOrchestrator from content_gen.core.model_router import ModelRoutingEngine from content_gen.core.schemas import ProcessedQuestion +from qc_viewer.services.job_repository import get_job_repository router = APIRouter(prefix="/api/v1", tags=["Service API v1"]) -# In-memory job store for the experimental run (use DB/Cache for production) -JOBS = {} - @router.post("/extract") async def extract_content( background_tasks: BackgroundTasks, @@ -41,13 +39,16 @@ async def extract_content( with open(file_path, "wb") as f: shutil.copyfileobj(file.file, f) - JOBS[job_id] = { - "status": "PROCESSING", - "id": job_id, - "subject": subject, - "curriculum": curriculum, - "created_at": datetime.now().isoformat() - } + get_job_repository().put( + job_id, + { + "status": "PROCESSING", + "id": job_id, + "subject": subject, + "curriculum": curriculum, + "created_at": datetime.now().isoformat(), + }, + ) # Run the processing in background background_tasks.add_task( @@ -65,9 +66,10 @@ async def extract_content( @router.get("/jobs/{job_id}") async def get_job_status(job_id: str): - if job_id not in JOBS: + row = get_job_repository().get(job_id) + if not row: raise HTTPException(status_code=404, detail="Job not found") - return JOBS[job_id] + return row async def process_service_job( job_id: str, @@ -79,17 +81,11 @@ async def process_service_job( gemini_key: Optional[str] ): try: - import os - # Determine BYOK with neutral header first, then legacy provider headers. if not api_key: api_key = gemini_key or openai_key - # Initialize engine with BYOK - if api_key: - os.environ["LITELLM_API_KEY"] = api_key - - router = ModelRoutingEngine() + router = ModelRoutingEngine(api_key=api_key) orchestrator = PipelineOrchestrator(router=router) # Execute processing @@ -114,15 +110,21 @@ async def process_service_job( } validated_questions.append(legacy_q) - JOBS[job_id].update({ - "status": "COMPLETED", - "questions": validated_questions, - "completed_at": datetime.now().isoformat() - }) - + get_job_repository().merge( + job_id, + { + "status": "COMPLETED", + "questions": validated_questions, + "completed_at": datetime.now().isoformat(), + }, + ) + except Exception as e: print(f"Job {job_id} failed: {e}") - JOBS[job_id].update({ - "status": "FAILED", - "error": str(e) - }) + get_job_repository().merge( + job_id, + { + "status": "FAILED", + "error": str(e), + }, + ) diff --git a/qc_viewer/routers/automation.py b/qc_viewer/routers/automation.py index 2ebb9a4..1816b78 100644 --- a/qc_viewer/routers/automation.py +++ b/qc_viewer/routers/automation.py @@ -17,6 +17,7 @@ from pydantic import BaseModel from qc_viewer.services import draft_export from qc_viewer.services.draft_store import ( + DraftNotFound, delete_draft_data, ensure_drafts_root, get_draft_dir, @@ -28,6 +29,13 @@ ) +def _require_draft_path(draft_id: str) -> Path: + try: + return resolve_metadata_path(draft_id) + except DraftNotFound: + raise HTTPException(status_code=404, detail="Draft not found") + + class PublishRequest(BaseModel): draft_id: str table_name: str @@ -142,7 +150,7 @@ async def list_drafts(): @router.get("/api/automate/draft/{draft_id}") async def get_draft_results(draft_id: str): - return read_json(resolve_metadata_path(draft_id)) + return read_json(_require_draft_path(draft_id)) @router.get("/api/automate/draft/{draft_id}/export") @@ -150,7 +158,7 @@ async def export_draft(draft_id: str, format: str = "json"): fmt = format.lower() if fmt not in draft_export.SUPPORTED_FORMATS: raise HTTPException(status_code=400, detail=f"Unsupported format: {format}") - meta = read_json(resolve_metadata_path(draft_id)) + meta = read_json(_require_draft_path(draft_id)) body = draft_export.render(meta, fmt) filename = draft_export.safe_filename(meta, fmt) return Response( @@ -164,7 +172,7 @@ async def export_draft(draft_id: str, format: str = "json"): async def stream_draft_progress(draft_id: str, request: Request): try: meta_path = resolve_metadata_path(draft_id) - except HTTPException: + except DraftNotFound: meta_path = DRAFTS_ROOT / draft_id / "metadata.json" if not meta_path.exists(): meta_path = DRAFTS_ROOT / f"{draft_id}.json" @@ -211,7 +219,7 @@ async def stop_draft_processing(draft_id: str): if meta and meta.get("status") == "PROCESSING": meta["status"] = "FAILED" meta["status_message"] = "Stopped by user" - write_json(resolve_metadata_path(draft_id), meta) + write_json(_require_draft_path(draft_id), meta) return {"status": "stopped"} return {"status": "not_running"} @@ -219,7 +227,7 @@ async def stop_draft_processing(draft_id: str): @router.patch("/api/automate/draft/{draft_id}") async def update_draft(draft_id: str, updates: dict): - meta_path = resolve_metadata_path(draft_id) + meta_path = _require_draft_path(draft_id) try: data = read_json(meta_path) data.update(updates) diff --git a/qc_viewer/routers/questions.py b/qc_viewer/routers/questions.py index 87b832a..e841703 100644 --- a/qc_viewer/routers/questions.py +++ b/qc_viewer/routers/questions.py @@ -1,12 +1,12 @@ -from typing import Any, Optional, cast +from typing import Optional from fastapi import APIRouter, HTTPException -from qc_viewer.config import get_allowed_table_ids from qc_viewer.services.db_service import get_db - +from qc_viewer.services.question_repository import QuestionRepository router = APIRouter() +_question_repo = QuestionRepository() @router.get("/api/papers") @@ -14,73 +14,41 @@ async def get_papers(): conn = get_db() if not conn: raise HTTPException(status_code=500, detail="Database connection failed") - papers = [] try: with conn.cursor() as cur: - for table in get_allowed_table_ids(): - cur.execute( - f"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = '{table}')" - ) - table_exists_row = cast(Optional[dict[str, Any]], cur.fetchone()) - if not table_exists_row or not table_exists_row["exists"]: - continue - cur.execute( - r"SELECT DISTINCT regexp_replace(question_identifier, '/Q?\d+$', '') AS code" - f" FROM {table} WHERE question_identifier IS NOT NULL" - ) - for row in cast(list[dict[str, Any]], cur.fetchall()): - code = row["code"] - if code: - papers.append({"code": code, "table": table}) + return _question_repo.list_papers(cur) finally: conn.close() - return sorted(papers, key=lambda x: x["code"]) @router.get("/api/questions") async def get_questions(table: str, paper_code: str): - if table not in get_allowed_table_ids(): + try: + _question_repo.assert_table_allowed(table) + except ValueError: raise HTTPException(status_code=400, detail="Invalid table name") conn = get_db() if not conn: raise HTTPException(status_code=500, detail="Database connection failed") try: with conn.cursor() as cur: - cur.execute( - r""" - SELECT id, question_identifier, - regexp_replace(question_identifier, '^.+/Q?', '') AS question_number, - title, options, - correct_options, option_explanations, - summary_explanation, detailed_explanation, - is_verified, other_contents AS diagrams, - topic_id, subtopic_id - FROM """ - + table - + r""" - WHERE question_identifier LIKE %s - ORDER BY question_identifier; - """, - (paper_code + "%",), - ) - return cur.fetchall() + return _question_repo.fetch_questions_for_paper(cur, table, paper_code) finally: conn.close() @router.post("/api/verify") async def verify_question(table: str, question_id: str, is_verified: bool): - if table not in get_allowed_table_ids(): + try: + _question_repo.assert_table_allowed(table) + except ValueError: raise HTTPException(status_code=400, detail="Invalid table name") conn = get_db() if not conn: raise HTTPException(status_code=500, detail="Database connection failed") try: with conn.cursor() as cur: - cur.execute( - f"UPDATE {table} SET is_verified = %s, updated_at = NOW() WHERE id = %s", - (is_verified, question_id), - ) + _question_repo.set_question_verified(cur, table, question_id, is_verified) conn.commit() return {"status": "success"} finally: @@ -98,38 +66,22 @@ async def get_flashcards( raise HTTPException(status_code=500, detail="Database connection failed") try: with conn.cursor() as cur: - if question_id and question_id != "undefined": - cur.execute('SELECT * FROM flashcards WHERE "questionId" = %s', (question_id,)) - res = cur.fetchall() - if res: - return res - - if subtopic_id and subtopic_id != "undefined" and subtopic_id != "null": - cur.execute('SELECT * FROM flashcards WHERE "subtopicId" = %s', (subtopic_id,)) - res = cur.fetchall() - if res: - return res - - if topic_id and topic_id != "undefined" and topic_id != "null": - cur.execute('SELECT * FROM flashcards WHERE "topicId" = %s', (topic_id,)) - return cur.fetchall() - - return [] + return _question_repo.fetch_flashcards(cur, topic_id, subtopic_id, question_id) finally: conn.close() @router.get("/api/question/details") async def get_question_details(table: str, question_identifier: str): - if table not in get_allowed_table_ids(): + try: + _question_repo.assert_table_allowed(table) + except ValueError: raise HTTPException(status_code=400, detail="Invalid table name") conn = get_db() if not conn: raise HTTPException(status_code=500, detail="Database connection failed") try: with conn.cursor() as cur: - query = "SELECT id, title, options, is_verified FROM " + table + " WHERE question_identifier = %s" - cur.execute(query, (question_identifier,)) - return cur.fetchone() + return _question_repo.fetch_question_details_row(cur, table, question_identifier) finally: conn.close() diff --git a/qc_viewer/services/automation_pipeline.py b/qc_viewer/services/automation_pipeline.py index 748d14c..25d07a9 100644 --- a/qc_viewer/services/automation_pipeline.py +++ b/qc_viewer/services/automation_pipeline.py @@ -1,18 +1,24 @@ import asyncio -import base64 -import os -import re import threading import time from datetime import datetime from pathlib import Path -from typing import Optional +from typing import Optional, Protocol from content_gen.core.model_router import ModelRoutingEngine + + +class _MutableModelRouting(Protocol): + """Assignable routing slots (real :class:`ModelRouting` or test doubles).""" + + extraction: str + generation: str + validation: str from content_gen.core.config_schema import DetectionMode from content_gen.core.pedagogy_engine import PedagogyEngine from content_gen.scripts.pipeline.pipeline_orchestrator import PipelineOrchestrator from qc_viewer.services.draft_store import read_modify_write_json +from qc_viewer.services.legacy_question_payload import build_legacy_question_dict CANCELLATION_EVENTS: dict[str, threading.Event] = {} @@ -55,7 +61,7 @@ def _provider_default_models(provider: str) -> dict[str, str]: def _apply_runtime_model_overrides( - router: ModelRoutingEngine, + model_routing: _MutableModelRouting, llm_provider: Optional[str], model_id: Optional[str], has_api_key: bool = False, @@ -71,9 +77,9 @@ def _apply_runtime_model_overrides( if requested_model_id: normalized = _normalize_model_id(requested_model_id, requested_provider) - router.config.model_routing.extraction = normalized - router.config.model_routing.generation = normalized - router.config.model_routing.validation = normalized + model_routing.extraction = normalized + model_routing.generation = normalized + model_routing.validation = normalized return { "provider": requested_provider, "requested_model_id": requested_model_id, @@ -83,9 +89,9 @@ def _apply_runtime_model_overrides( if requested_provider and has_api_key: defaults = _provider_default_models(requested_provider) if defaults: - router.config.model_routing.extraction = defaults["extraction"] - router.config.model_routing.generation = defaults["generation"] - router.config.model_routing.validation = defaults["validation"] + model_routing.extraction = defaults["extraction"] + model_routing.generation = defaults["generation"] + model_routing.validation = defaults["validation"] return { "provider": requested_provider, "requested_model_id": None, @@ -99,61 +105,6 @@ def _apply_runtime_model_overrides( } -def extract_correct_answer(explanation: str) -> str: - match = re.search(r"Final Correct Answer\s*:\s*([A-D])", explanation or "", re.IGNORECASE) - if match: - return match.group(1).upper() - # Fallback to looking for single letters at the end - match_fallback = re.search(r"(?:Correct Answer|Answer is)\s*[:\s]*([A-D])", explanation or "", re.IGNORECASE) - if match_fallback: - return match_fallback.group(1).upper() - return "N/A" - - -def extract_option_analysis(option_text: str, option_keys: list[str]) -> dict[str, str]: - result = {k: "" for k in option_keys} - if not option_text: - return result - pattern = re.compile( - r"Option\s*([A-D])\s*:\s*(.*?)(?=Option\s*[A-D]\s*:|$)", - re.IGNORECASE | re.DOTALL, - ) - for match in pattern.finditer(option_text): - label = match.group(1).upper() - if label in result: - result[label] = re.sub(r"\s+", " ", match.group(2)).strip() - return result - - -def extract_core_concept(explanation: str) -> str: - if not explanation: - return "" - core_match = re.search( - r"^\s*Core Concept\s*:?\s*(.*?)(?=\n\s*(?:Step\s*1|Analyze Step 1|Final Correct Answer|Option\s+[A-D]:)|$)", - explanation, - re.IGNORECASE | re.DOTALL | re.MULTILINE, - ) - if core_match: - content = re.sub(r"\s+", " ", core_match.group(1)).strip(" -*\n\t") - if content: - return content - # Smarter fallback: first few sentences or first 300 chars (approx 5-6 lines) - lines = [ln.strip() for ln in explanation.splitlines() if ln.strip()] - if not lines: - return "" - - # Try to get the first 2-3 sentences if possible - full_text = " ".join(lines) - sentences = re.split(r'(?<=[.!?])\s+', full_text) - if len(sentences) > 0: - concept = " ".join(sentences[:3]) - if len(concept) > 350: - return concept[:347] + "..." - return concept - - return full_text[:297] + "..." - - def run_automation_pipeline( draft_id: str, subject: str, @@ -194,7 +145,7 @@ def _mut(meta: dict) -> None: print(f"Error updating progress: {e}") try: - router = ModelRoutingEngine() + router = ModelRoutingEngine(api_key=api_key) if curriculum is None or not str(curriculum).strip(): curriculum = ( (router.config.workspace.default_curriculum or "").strip() or "General" @@ -215,12 +166,12 @@ def _mut(meta: dict) -> None: router.config.extraction_settings.max_question_number = max_question_number resolved_model_override = _apply_runtime_model_overrides( - router, llm_provider, model_id, has_api_key=(api_key is not None) + router.config.model_routing, + llm_provider, + model_id, + has_api_key=(api_key is not None), ) - if api_key: - os.environ["LITELLM_API_KEY"] = api_key - orchestrator = PipelineOrchestrator(router=router) draft_dir = str(file_path.parent) @@ -257,81 +208,7 @@ def _mut(meta: dict) -> None: processed_count = i + 1 progress_val = 90 + int((processed_count / total_questions) * 9) _update_progress(progress_val, f"Finalizing question {processed_count} of {total_questions}...") - - diagram_b64 = None - stem_images = q.metadata.get("stem_images", []) - stem_images_b64 = q.metadata.get("stem_images_b64", []) - if stem_images_b64: - diagram_b64 = stem_images_b64[0] - if stem_images and len(stem_images) > 0: - first_img = stem_images[0] - if diagram_b64: - pass - elif str(first_img).startswith("data:image"): - diagram_b64 = first_img - else: - try: - img_path = Path(first_img) - if img_path.exists(): - with open(img_path, "rb") as img_f: - diagram_b64 = ( - f"data:image/png;base64,{base64.b64encode(img_f.read()).decode('utf-8')}" - ) - except Exception as img_e: - print(f"Failed to encode diagram: {img_e}") - - explanation_text = q.explanation_body or "" - option_text = q.option_wise_explanation or "" - option_analysis = extract_option_analysis(option_text, list((q.options or {}).keys())) - - # Prioritize dedicated Core Concept from LLM - core_concept = (q.metadata or {}).get("core_concept_generated") - if not core_concept: - core_concept = extract_core_concept(explanation_text) - - quality_report = (q.metadata or {}).get("generation_quality", {}) - - # Normalize Correct Answer - norm_correct_answer = (q.correct_options[0] if q.correct_options else extract_correct_answer(explanation_text)).strip().upper() - if len(norm_correct_answer) > 1: - norm_correct_answer = norm_correct_answer[0] - if norm_correct_answer not in "ABCD": - norm_correct_answer = "N/A" - - # Normalize Option Analysis keys - final_option_analysis = {k: option_analysis.get(k, "Analysis missing for this option.") for k in ["A", "B", "C", "D"]} - - # Ensure Core Concept is not a duplicate of Detailed Explanation - clean_core_concept = core_concept - if clean_core_concept.lower() in explanation_text.lower() and len(clean_core_concept) > len(explanation_text) * 0.8: - clean_core_concept = "Concept extracted from explanation." - - opts_map = q.options if isinstance(q.options, dict) else {} - opt_vals = [str(opts_map.get(k, "") or "").strip() for k in ("A", "B", "C", "D")] - non_empty_opts = sum(1 for v in opt_vals if v) - extraction_warnings = ( - ["mcq_options_missing"] if non_empty_opts < 2 else [] - ) - - legacy_q = { - "question_number": q.question_number, - "text": q.question_text, - "options": q.options, - "correct_answer": norm_correct_answer, - "status": "Draft", - "diagram_base64": diagram_b64, - "generated_content": { - "core_concept": clean_core_concept, - "detailed_explanation": explanation_text, - "option_analysis": final_option_analysis, - "flashcards": [{"question": f.front_text, "answer": f.back_text} for f in q.flashcards], - }, - "quality_report": quality_report, - "contract_warnings": [k for k, v in quality_report.items() if v is False], - } - if extraction_warnings: - legacy_q["extraction_warnings"] = extraction_warnings - questions_payload.append(legacy_q) + questions_payload.append(build_legacy_question_dict(q)) t_normalization_end = time.time() t_pipeline_end = time.time() diff --git a/qc_viewer/services/db_service.py b/qc_viewer/services/db_service.py index 994c908..000e7f4 100644 --- a/qc_viewer/services/db_service.py +++ b/qc_viewer/services/db_service.py @@ -1,18 +1,15 @@ from typing import Optional -import psycopg2 -from psycopg2.extras import RealDictCursor - +from content_gen.db.session import connect_real_dict from qc_viewer.config import DATABASE_URL def get_db(): + if DATABASE_URL is None: + print("DB connection error: DATABASE_URL is not set") + return None try: - return psycopg2.connect( - DATABASE_URL, - cursor_factory=RealDictCursor, - connect_timeout=3, - ) + return connect_real_dict(DATABASE_URL, connect_timeout=3) except Exception as e: print(f"DB connection error: {e}") return None diff --git a/qc_viewer/services/draft_store.py b/qc_viewer/services/draft_store.py index c18a3e4..d3756d8 100644 --- a/qc_viewer/services/draft_store.py +++ b/qc_viewer/services/draft_store.py @@ -6,13 +6,19 @@ from typing import Any, Callable, Optional import threading -from fastapi import HTTPException - from qc_viewer.config import DRAFTS_ROOT METADATA_LOCK = threading.Lock() +class DraftNotFound(Exception): + """No metadata file exists for this draft_id (neither dir/metadata.json nor legacy flat .json).""" + + def __init__(self, draft_id: str): + self.draft_id = draft_id + super().__init__(f"Draft not found: {draft_id}") + + def ensure_drafts_root() -> None: DRAFTS_ROOT.mkdir(parents=True, exist_ok=True) @@ -36,7 +42,7 @@ def resolve_metadata_path(draft_id: str) -> Path: legacy_path = get_legacy_metadata_path(draft_id) if legacy_path.exists(): return legacy_path - raise HTTPException(status_code=404, detail="Draft not found") + raise DraftNotFound(draft_id) def read_json(path: Path) -> dict[str, Any]: @@ -122,5 +128,5 @@ def delete_draft_data(draft_id: str) -> bool: def load_metadata_if_exists(draft_id: str) -> Optional[dict[str, Any]]: try: return read_json(resolve_metadata_path(draft_id)) - except HTTPException: + except DraftNotFound: return None diff --git a/qc_viewer/services/job_repository.py b/qc_viewer/services/job_repository.py new file mode 100644 index 0000000..91f4c0c --- /dev/null +++ b/qc_viewer/services/job_repository.py @@ -0,0 +1,119 @@ +""" +Pluggable job storage for /api/v1 async extract jobs. + +Default: in-memory (dev). Set EDMATE_JOB_STORE_DB to a filesystem path for a +SQLite-backed store (survives process restarts; use Redis/Postgres for multi-worker). +""" + +from __future__ import annotations + +import json +import os +import sqlite3 +import threading +from abc import ABC, abstractmethod +from typing import Any, Dict, Optional + + +class JobRepository(ABC): + @abstractmethod + def get(self, job_id: str) -> Optional[Dict[str, Any]]: + ... + + @abstractmethod + def put(self, job_id: str, payload: Dict[str, Any]) -> None: + ... + + @abstractmethod + def merge(self, job_id: str, updates: Dict[str, Any]) -> None: + """Shallow merge updates into the existing job record.""" + ... + + +class InMemoryJobRepository(JobRepository): + def __init__(self) -> None: + self._jobs: Dict[str, Dict[str, Any]] = {} + self._lock = threading.Lock() + + def get(self, job_id: str) -> Optional[Dict[str, Any]]: + with self._lock: + row = self._jobs.get(job_id) + return None if row is None else dict(row) + + def put(self, job_id: str, payload: Dict[str, Any]) -> None: + with self._lock: + self._jobs[job_id] = dict(payload) + + def merge(self, job_id: str, updates: Dict[str, Any]) -> None: + with self._lock: + if job_id not in self._jobs: + self._jobs[job_id] = {} + self._jobs[job_id].update(updates) + + +class SqliteJobRepository(JobRepository): + def __init__(self, db_path: str) -> None: + self._path = db_path + self._lock = threading.Lock() + parent = os.path.dirname(os.path.abspath(db_path)) + if parent: + os.makedirs(parent, exist_ok=True) + with self._connect() as conn: + conn.execute( + "CREATE TABLE IF NOT EXISTS jobs (id TEXT PRIMARY KEY, payload TEXT NOT NULL)" + ) + conn.commit() + + def _connect(self) -> sqlite3.Connection: + return sqlite3.connect(self._path, check_same_thread=False) + + def get(self, job_id: str) -> Optional[Dict[str, Any]]: + with self._lock: + with self._connect() as conn: + cur = conn.execute("SELECT payload FROM jobs WHERE id = ?", (job_id,)) + row = cur.fetchone() + if not row: + return None + return json.loads(row[0]) + + def put(self, job_id: str, payload: Dict[str, Any]) -> None: + body = json.dumps(payload, default=str) + with self._lock: + with self._connect() as conn: + conn.execute( + "INSERT INTO jobs(id, payload) VALUES(?, ?) ON CONFLICT(id) DO UPDATE SET payload=excluded.payload", + (job_id, body), + ) + conn.commit() + + def merge(self, job_id: str, updates: Dict[str, Any]) -> None: + current = self.get(job_id) or {} + current.update(updates) + self.put(job_id, current) + + +def _default_repo_from_env() -> JobRepository: + path = (os.environ.get("EDMATE_JOB_STORE_DB") or "").strip() + if path: + return SqliteJobRepository(path) + return InMemoryJobRepository() + + +_default_repo: Optional[JobRepository] = None +_repo_lock = threading.Lock() + + +def get_job_repository() -> JobRepository: + """Singleton repository (swap via set_job_repository in tests).""" + global _default_repo + with _repo_lock: + if _default_repo is None: + _default_repo = _default_repo_from_env() + return _default_repo + + +def set_job_repository(repo: Optional[JobRepository]) -> None: + """Test hook: pass None to reset to env-based default on next get.""" + global _default_repo + with _repo_lock: + _default_repo = repo diff --git a/qc_viewer/services/legacy_question_payload.py b/qc_viewer/services/legacy_question_payload.py new file mode 100644 index 0000000..4a46a44 --- /dev/null +++ b/qc_viewer/services/legacy_question_payload.py @@ -0,0 +1,151 @@ +""" +Pure helpers: map pipeline ProcessedQuestion → legacy QC viewer JSON. + +Used by automation_pipeline and tests; no I/O. +""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import Any, Dict, List, Optional + +from content_gen.core.media_encoding import png_file_to_data_uri +from content_gen.core.schemas import ProcessedQuestion + + +def extract_correct_answer(explanation: str) -> str: + match = re.search(r"Final Correct Answer\s*:\s*([A-D])", explanation or "", re.IGNORECASE) + if match: + return match.group(1).upper() + match_fallback = re.search( + r"(?:Correct Answer|Answer is)\s*[:\s]*([A-D])", explanation or "", re.IGNORECASE + ) + if match_fallback: + return match_fallback.group(1).upper() + return "N/A" + + +def extract_option_analysis(option_text: str, option_keys: list[str]) -> dict[str, str]: + result = {k: "" for k in option_keys} + if not option_text: + return result + pattern = re.compile( + r"Option\s*([A-D])\s*:\s*(.*?)(?=Option\s*[A-D]\s*:|$)", + re.IGNORECASE | re.DOTALL, + ) + for match in pattern.finditer(option_text): + label = match.group(1).upper() + if label in result: + result[label] = re.sub(r"\s+", " ", match.group(2)).strip() + return result + + +def extract_core_concept(explanation: str) -> str: + if not explanation: + return "" + core_match = re.search( + r"^\s*Core Concept\s*:?\s*(.*?)(?=\n\s*(?:Step\s*1|Analyze Step 1|Final Correct Answer|Option\s+[A-D]:)|$)", + explanation, + re.IGNORECASE | re.DOTALL | re.MULTILINE, + ) + if core_match: + content = re.sub(r"\s+", " ", core_match.group(1)).strip(" -*\n\t") + if content: + return content + lines = [ln.strip() for ln in explanation.splitlines() if ln.strip()] + if not lines: + return "" + + full_text = " ".join(lines) + sentences = re.split(r"(?<=[.!?])\s+", full_text) + if len(sentences) > 0: + concept = " ".join(sentences[:3]) + if len(concept) > 350: + return concept[:347] + "..." + return concept + + return full_text[:297] + "..." + + +def resolve_diagram_data_uri(q: ProcessedQuestion) -> Optional[str]: + """Return a data:image PNG URI for the first stem diagram, if any.""" + md = q.metadata or {} + stem_images_b64 = md.get("stem_images_b64", []) + if stem_images_b64: + return stem_images_b64[0] + stem_images = md.get("stem_images", []) + if not stem_images: + return None + first_img = stem_images[0] + if str(first_img).startswith("data:image"): + return str(first_img) + try: + img_path = Path(first_img) + if img_path.exists(): + return png_file_to_data_uri(img_path) + except OSError: + pass + return None + + +def build_legacy_question_dict(q: ProcessedQuestion) -> Dict[str, Any]: + """Single QC-viewer question record after extraction + generation.""" + diagram_b64 = resolve_diagram_data_uri(q) + + explanation_text = q.explanation_body or "" + option_text = q.option_wise_explanation or "" + option_keys = list((q.options or {}).keys()) or ["A", "B", "C", "D"] + option_analysis = extract_option_analysis(option_text, option_keys) + + core_concept = (q.metadata or {}).get("core_concept_generated") + if not core_concept: + core_concept = extract_core_concept(explanation_text) + + quality_report = (q.metadata or {}).get("generation_quality", {}) + + norm_correct_answer = ( + q.correct_options[0] if q.correct_options else extract_correct_answer(explanation_text) + ).strip().upper() + if len(norm_correct_answer) > 1: + norm_correct_answer = norm_correct_answer[0] + if norm_correct_answer not in "ABCD": + norm_correct_answer = "N/A" + + final_option_analysis = { + k: option_analysis.get(k, "Analysis missing for this option.") for k in ["A", "B", "C", "D"] + } + + clean_core_concept = core_concept + if ( + clean_core_concept + and explanation_text + and clean_core_concept.lower() in explanation_text.lower() + and len(clean_core_concept) > len(explanation_text) * 0.8 + ): + clean_core_concept = "Concept extracted from explanation." + + opts_map = q.options if isinstance(q.options, dict) else {} + opt_vals = [str(opts_map.get(k, "") or "").strip() for k in ("A", "B", "C", "D")] + non_empty_opts = sum(1 for v in opt_vals if v) + extraction_warnings: List[str] = ["mcq_options_missing"] if non_empty_opts < 2 else [] + + legacy_q: Dict[str, Any] = { + "question_number": q.question_number, + "text": q.question_text, + "options": q.options, + "correct_answer": norm_correct_answer, + "status": "Draft", + "diagram_base64": diagram_b64, + "generated_content": { + "core_concept": clean_core_concept, + "detailed_explanation": explanation_text, + "option_analysis": final_option_analysis, + "flashcards": [{"question": f.front_text, "answer": f.back_text} for f in q.flashcards], + }, + "quality_report": quality_report, + "contract_warnings": [k for k, v in quality_report.items() if v is False], + } + if extraction_warnings: + legacy_q["extraction_warnings"] = extraction_warnings + return legacy_q diff --git a/qc_viewer/services/question_repository.py b/qc_viewer/services/question_repository.py new file mode 100644 index 0000000..734edbb --- /dev/null +++ b/qc_viewer/services/question_repository.py @@ -0,0 +1,99 @@ +""" +PostgreSQL access for QC question/flashcard routes. + +Routers validate HTTP; this module performs allowlisted table SQL only. +""" + +from __future__ import annotations + +from typing import Any, List, Optional, cast + +from qc_viewer.config import get_allowed_table_ids + + +class QuestionRepository: + """Read/write questions and related flashcards using allowlisted table names.""" + + def list_papers(self, cur) -> List[dict[str, Any]]: + papers: list[dict[str, Any]] = [] + for table in get_allowed_table_ids(): + cur.execute( + f"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = '{table}')" + ) + table_exists_row = cast(Optional[dict[str, Any]], cur.fetchone()) + if not table_exists_row or not table_exists_row["exists"]: + continue + cur.execute( + r"SELECT DISTINCT regexp_replace(question_identifier, '/Q?\d+$', '') AS code" + f" FROM {table} WHERE question_identifier IS NOT NULL" + ) + for row in cast(list[dict[str, Any]], cur.fetchall()): + code = row["code"] + if code: + papers.append({"code": code, "table": table}) + return sorted(papers, key=lambda x: x["code"]) + + def fetch_questions_for_paper(self, cur, table: str, paper_code: str) -> list: + self.assert_table_allowed(table) + cur.execute( + r""" + SELECT id, question_identifier, + regexp_replace(question_identifier, '^.+/Q?', '') AS question_number, + title, options, + correct_options, option_explanations, + summary_explanation, detailed_explanation, + is_verified, other_contents AS diagrams, + topic_id, subtopic_id + FROM """ + + table + + r""" + WHERE question_identifier LIKE %s + ORDER BY question_identifier; + """, + (paper_code + "%",), + ) + return cur.fetchall() + + def set_question_verified(self, cur, table: str, question_id: str, is_verified: bool) -> None: + self.assert_table_allowed(table) + cur.execute( + f"UPDATE {table} SET is_verified = %s, updated_at = NOW() WHERE id = %s", + (is_verified, question_id), + ) + + def fetch_flashcards( + self, + cur, + topic_id: Optional[str] = None, + subtopic_id: Optional[str] = None, + question_id: Optional[str] = None, + ) -> list: + if question_id and question_id != "undefined": + cur.execute('SELECT * FROM flashcards WHERE "questionId" = %s', (question_id,)) + res = cur.fetchall() + if res: + return res + + if subtopic_id and subtopic_id not in ("undefined", "null"): + cur.execute('SELECT * FROM flashcards WHERE "subtopicId" = %s', (subtopic_id,)) + res = cur.fetchall() + if res: + return res + + if topic_id and topic_id not in ("undefined", "null"): + cur.execute('SELECT * FROM flashcards WHERE "topicId" = %s', (topic_id,)) + return cur.fetchall() + + return [] + + def fetch_question_details_row(self, cur, table: str, question_identifier: str): + self.assert_table_allowed(table) + cur.execute( + "SELECT id, title, options, is_verified FROM " + table + " WHERE question_identifier = %s", + (question_identifier,), + ) + return cur.fetchone() + + def assert_table_allowed(self, table: str) -> None: + if table not in get_allowed_table_ids(): + raise ValueError("Invalid table name") diff --git a/qc_viewer/services/tests/test_job_repository.py b/qc_viewer/services/tests/test_job_repository.py new file mode 100644 index 0000000..32c279c --- /dev/null +++ b/qc_viewer/services/tests/test_job_repository.py @@ -0,0 +1,50 @@ +import unittest + +from qc_viewer.services.job_repository import ( + InMemoryJobRepository, + SqliteJobRepository, + get_job_repository, + set_job_repository, +) + + +class TestJobRepository(unittest.TestCase): + def tearDown(self) -> None: + set_job_repository(None) + + def test_in_memory_put_get_merge(self): + r = InMemoryJobRepository() + r.put("j1", {"status": "PROCESSING", "id": "j1"}) + r.merge("j1", {"progress": 50}) + row = r.get("j1") + if row is None: + self.fail("Expected row to exist for j1") + self.assertEqual(row["status"], "PROCESSING") + self.assertEqual(row["progress"], 50) + + def test_sqlite_roundtrip(self): + import os + import tempfile + + fd, path = tempfile.mkstemp(suffix="_jobs.db") + os.close(fd) + try: + r = SqliteJobRepository(path) + r.put("x", {"a": 1}) + self.assertEqual(r.get("x"), {"a": 1}) + r.merge("x", {"b": 2}) + self.assertEqual(r.get("x"), {"a": 1, "b": 2}) + finally: + os.unlink(path) + + def test_singleton_reset(self): + set_job_repository(InMemoryJobRepository()) + get_job_repository().put("singleton", {"ok": True}) + row = get_job_repository().get("singleton") + if row is None: + self.fail("Expected singleton row to exist") + self.assertTrue(row["ok"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/qc_viewer/services/tests/test_legacy_question_payload.py b/qc_viewer/services/tests/test_legacy_question_payload.py new file mode 100644 index 0000000..14c3c9a --- /dev/null +++ b/qc_viewer/services/tests/test_legacy_question_payload.py @@ -0,0 +1,79 @@ +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory + +from content_gen.core.schemas import Flashcard, ProcessedQuestion +from qc_viewer.services.legacy_question_payload import ( + build_legacy_question_dict, + resolve_diagram_data_uri, +) + + +class TestLegacyQuestionPayload(unittest.TestCase): + def test_resolve_diagram_from_b64_metadata(self): + q = ProcessedQuestion( + question_number=1, + question_text="Q?", + options={"A": "1", "B": "2", "C": "3", "D": "4"}, + correct_options=["A"], + subject="Physics", + metadata={"stem_images_b64": ["data:image/png;base64,xxx"]}, + ) + self.assertEqual(resolve_diagram_data_uri(q), "data:image/png;base64,xxx") + + def test_resolve_diagram_from_file(self): + with TemporaryDirectory() as tmp: + p = Path(tmp) / "stem.png" + p.write_bytes(b"\x89PNG\r\n\x1a\n") # minimal PNG header bytes; reader only base64 encodes + q = ProcessedQuestion( + question_number=1, + question_text="Q?", + options={"A": "1"}, + correct_options=[], + subject="Physics", + metadata={"stem_images": [str(p)]}, + ) + uri = resolve_diagram_data_uri(q) + if uri is None: + self.fail("Expected data URI from stem image file") + self.assertTrue(uri.startswith("data:image/png;base64,")) + + def test_build_legacy_question_dict_contract_warnings(self): + q = ProcessedQuestion( + question_number=2, + question_text="Stem", + options={"A": "a", "B": "b", "C": "c", "D": "d"}, + correct_options=["B"], + subject="Chemistry", + explanation_body="Final Correct Answer: B\n", + option_wise_explanation="Option A: x\nOption B: y\nOption C: z\nOption D: w\n", + metadata={ + "generation_quality": { + "has_explanation": True, + "has_final_answer": False, + } + }, + ) + legacy = build_legacy_question_dict(q) + self.assertEqual(legacy["correct_answer"], "B") + self.assertIn("has_final_answer", legacy["contract_warnings"]) + self.assertEqual(len(legacy["generated_content"]["flashcards"]), 0) + + def test_build_legacy_with_flashcards(self): + q = ProcessedQuestion( + question_number=1, + question_text="Q", + options={"A": "1", "B": "2", "C": "3", "D": "4"}, + correct_options=["A"], + subject="Bio", + flashcards=[Flashcard(front_text="F", back_text="B")], + ) + legacy = build_legacy_question_dict(q) + self.assertEqual( + legacy["generated_content"]["flashcards"], + [{"question": "F", "answer": "B"}], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/qc_viewer/services/tests/test_pipeline.py b/qc_viewer/services/tests/test_pipeline.py index b620731..71ef3db 100644 --- a/qc_viewer/services/tests/test_pipeline.py +++ b/qc_viewer/services/tests/test_pipeline.py @@ -1,14 +1,7 @@ import unittest -from typing import Optional -import qc_viewer.services.automation_pipeline as ap -print(f"DEBUG: Module path: {ap.__file__}") -from qc_viewer.services.automation_pipeline import ( - _apply_runtime_model_overrides, - extract_correct_answer, - extract_core_concept, -) -from content_gen.core.model_router import ModelRoutingEngine +from qc_viewer.services.automation_pipeline import _apply_runtime_model_overrides +from qc_viewer.services.legacy_question_payload import extract_correct_answer, extract_core_concept class MockModelRouting: def __init__(self): @@ -37,19 +30,28 @@ def setUp(self): def test_apply_runtime_model_overrides_no_byok(self): # Even if provider is sent, if has_api_key is False, it should NOT override - res = _apply_runtime_model_overrides(self.router, "openai", None, has_api_key=False) + res = _apply_runtime_model_overrides( + self.router.config.model_routing, "openai", None, has_api_key=False + ) self.assertIsNone(res["resolved_model"]) self.assertEqual(self.router.config.model_routing.generation, "default/generation") def test_apply_runtime_model_overrides_with_byok(self): # If has_api_key is True, it SHOULD override with provider defaults - res = _apply_runtime_model_overrides(self.router, "openai", None, has_api_key=True) + res = _apply_runtime_model_overrides( + self.router.config.model_routing, "openai", None, has_api_key=True + ) self.assertEqual(res["resolved_model"], "openai/gpt-4o-mini") self.assertEqual(self.router.config.model_routing.generation, "openai/gpt-4o-mini") def test_apply_runtime_model_overrides_explicit_model(self): # Explicit model ID should always override - res = _apply_runtime_model_overrides(self.router, "gemini", "gemini-1.5-pro", has_api_key=True) + res = _apply_runtime_model_overrides( + self.router.config.model_routing, + "gemini", + "gemini-1.5-pro", + has_api_key=True, + ) self.assertEqual(res["resolved_model"], "gemini/gemini-1.5-pro") def test_extract_correct_answer(self): @@ -68,8 +70,6 @@ def test_extract_core_concept(self): # Test truncation long_line = "This is a very very long first line that should definitely be truncated because it's way too long to be a core concept title and we want to keep it clean " * 5 result = extract_core_concept(long_line) - print(f"DEBUG: Result length: {len(result)}") - print(f"DEBUG: Result: {result}") self.assertTrue(result.endswith("...")) if __name__ == "__main__": From 48d65b67d58cbb4f00f7ce44757713252424b8a1 Mon Sep 17 00:00:00 2001 From: shmukit Date: Thu, 7 May 2026 23:10:23 +0600 Subject: [PATCH 2/8] refactor(model_router): modularize content generation logic and enforce budget checks - Split the content generation process into smaller, private methods for better readability and maintainability. - Introduced `_enforce_budget`, `_select_model`, `_build_messages`, and `_completion_kwargs` methods to encapsulate specific functionalities. - Updated the `generate_content` method to utilize the new modular structure while maintaining existing functionality. - Enhanced error handling and logging for budget enforcement and response processing in the content generation workflow. --- content_gen/core/model_router.py | 83 ++-- .../scripts/processing/content_generator.py | 92 ++-- content_gen/tests/test_regression_guards.py | 2 +- qc_viewer/static/index.html | 439 +---------------- qc_viewer/static/js/api.js | 6 + qc_viewer/static/js/controllers/viewer.js | 463 ++++++++++++++++++ 6 files changed, 584 insertions(+), 501 deletions(-) create mode 100644 qc_viewer/static/js/api.js create mode 100644 qc_viewer/static/js/controllers/viewer.js diff --git a/content_gen/core/model_router.py b/content_gen/core/model_router.py index e6d30da..ea7e969 100644 --- a/content_gen/core/model_router.py +++ b/content_gen/core/model_router.py @@ -36,20 +36,7 @@ def __init__( _k = (api_key or "").strip() self._api_key: Optional[str] = _k or None - def generate_content( - self, - prompt: str, - task_type: str = "generation", - system_prompt: Optional[str] = None, - images: Optional[List[str]] = None, - json_mode: bool = False, - api_key: Optional[str] = None, - ) -> str: - """ - Routes the task to the appropriate model based on task_type. - Checks budget before execution (Economic Kill-Switch). - """ - # 1. Check Budget (The Safety Layer) + def _enforce_budget(self) -> None: current_cost = self.tracker.get_current_cost() if current_cost >= self.config.budget.max_daily_usd: raise BudgetExceededError( @@ -57,51 +44,79 @@ def generate_content( f"has reached the limit (${self.config.budget.max_daily_usd:.2f})." ) - # 2. Determine model + def _select_model(self, task_type: str) -> str: if task_type == "extraction": - model = self.config.model_routing.extraction or "gemini/gemini-1.5-pro" - elif task_type == "validation": - model = self.config.model_routing.validation or "openai/gpt-4o" - else: - model = self.config.model_routing.generation or "anthropic/claude-3-haiku" + return self.config.model_routing.extraction or "gemini/gemini-1.5-pro" + if task_type == "validation": + return self.config.model_routing.validation or "openai/gpt-4o" + return self.config.model_routing.generation or "anthropic/claude-3-haiku" - messages = [] + def _build_messages( + self, + prompt: str, + system_prompt: Optional[str], + images: Optional[List[str]], + ) -> List[Dict[str, Any]]: + messages: List[Dict[str, Any]] = [] if system_prompt: messages.append({"role": "system", "content": system_prompt}) - - # Handle Multimodal if images: content: List[Dict[str, Any]] = [{"type": "text", "text": prompt}] for img in images: - content.append( - {"type": "image_url", "image_url": {"url": img}}) + content.append({"type": "image_url", "image_url": {"url": img}}) messages.append({"role": "user", "content": content}) else: messages.append({"role": "user", "content": prompt}) + return messages - # 3. Execute call (optional api_key: per-call override, else instance BYOK) + def _completion_kwargs( + self, + model: str, + messages: List[Dict[str, Any]], + task_type: str, + json_mode: bool, + api_key: Optional[str], + ) -> Dict[str, Any]: _key = (api_key or "").strip() or self._api_key - _kwargs: Dict[str, Any] = { + kwargs: Dict[str, Any] = { "model": model, "messages": messages, "timeout": 120 if task_type == "extraction" else 180, } if json_mode: - _kwargs["response_format"] = {"type": "json_object"} + kwargs["response_format"] = {"type": "json_object"} if _key: - _kwargs["api_key"] = _key - response = litellm.completion(**_kwargs) - - # 4. Log usage (The Analytics Layer) - self.tracker.log_usage(response) + kwargs["api_key"] = _key + return kwargs + def _extract_message_text(self, response: Any) -> str: response_obj = cast(Any, response) if not hasattr(response_obj, "choices"): raise RuntimeError("Unexpected streaming response received from litellm.completion") - content = response_obj.choices[0].message.content return content if isinstance(content, str) else str(content) + def generate_content( + self, + prompt: str, + task_type: str = "generation", + system_prompt: Optional[str] = None, + images: Optional[List[str]] = None, + json_mode: bool = False, + api_key: Optional[str] = None, + ) -> str: + """ + Routes the task to the appropriate model based on task_type. + Checks budget before execution (Economic Kill-Switch). + """ + self._enforce_budget() + model = self._select_model(task_type) + messages = self._build_messages(prompt, system_prompt, images) + kwargs = self._completion_kwargs(model, messages, task_type, json_mode, api_key) + response = litellm.completion(**kwargs) + self.tracker.log_usage(response) + return self._extract_message_text(response) + def get_summary(self) -> Dict[str, Any]: """Returns the current metrics and routing config.""" return { diff --git a/content_gen/scripts/processing/content_generator.py b/content_gen/scripts/processing/content_generator.py index 42cec6f..bd926c8 100644 --- a/content_gen/scripts/processing/content_generator.py +++ b/content_gen/scripts/processing/content_generator.py @@ -1,10 +1,13 @@ +import logging import os import re -import json -import time from typing import List, Dict, Optional, Callable from pathlib import Path +logger = logging.getLogger(__name__) + +_DEFAULT_LLM_RESPONSE_LOG = Path("content_gen/logs/all_llm_responses.log") + # Try to import LLM libraries try: import google.generativeai as genai @@ -42,11 +45,22 @@ class ContentGenerator: using the modular ModelRoutingEngine. """ - def __init__(self, router: Optional[ModelRoutingEngine] = None): + def __init__( + self, + router: Optional[ModelRoutingEngine] = None, + *, + log_raw_llm_responses: bool = True, + raw_response_log_path: Optional[Path] = None, + raw_response_trace_sink: Optional[Callable[[str], None]] = None, + ): """ - Initialize the generator with a modular router. + log_raw_llm_responses: append full raw model output to ``raw_response_log_path`` (default). + raw_response_trace_sink: if set, receives the same trace block and file logging is skipped. """ self.router = router or ModelRoutingEngine() + self._log_raw_llm_responses = log_raw_llm_responses + self._raw_response_log_path = raw_response_log_path or _DEFAULT_LLM_RESPONSE_LOG + self._raw_response_trace_sink = raw_response_trace_sink def _default_curriculum(self) -> str: raw = getattr(self.router.config.workspace, "default_curriculum", None) @@ -236,49 +250,69 @@ def _parse_flashcards(self, gap_body: str) -> List[Flashcard]: flashcards.append(Flashcard(front_text=front, back_text=back)) return flashcards - def _parse_response(self, response: str, batch_indices: List[int]) -> Dict[int, Dict]: - """ - Parses the LLM output back into a structured dictionary using strict markers. - """ - results = {} - - # Log response - log_dir = Path("content_gen/logs") - log_dir.mkdir(parents=True, exist_ok=True) - with open(log_dir / "all_llm_responses.log", "a", encoding="utf-8") as f: - f.write(f"\n\n{'='*50}\nBatch: {batch_indices}\n{'='*50}\n") - f.write(response) + def _trace_raw_llm_response(self, batch_indices: List[int], response: str) -> None: + """Record full LLM output for debugging (file, custom sink, or skipped).""" + block = f"\n\n{'='*50}\nBatch: {batch_indices}\n{'='*50}\n{response}" + if self._raw_response_trace_sink is not None: + self._raw_response_trace_sink(block) + logger.debug( + "LLM raw response traced via sink for batch %s (%d chars)", + batch_indices, + len(response), + ) + return + if not self._log_raw_llm_responses: + return + try: + path = self._raw_response_log_path + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "a", encoding="utf-8") as f: + f.write(block) + logger.debug( + "LLM raw response appended to %s for batch %s (%d chars)", + path, + batch_indices, + len(response), + ) + except OSError as e: + logger.warning("Could not append LLM trace log: %s", e) - # More robust splitting - # Look for "Question X", "Q1", markdown headers, or just digits followed by period/brace + def _split_response_into_sections(self, response: str) -> List[str]: + """Split flat LLM text on question headers; returns re.split list (captures + bodies).""" header_pattern = r'(?mi)^\s*(?:#+\s*)?(?:Question|Q|Problem)?\s*[:\s]*(\d+)\s*[:\.\-\)]*\s*$' sections = re.split(header_pattern, response) if len(sections) < 3: alt_pattern = r'(?m)^\s*(?:#+\s*)?(\d{1,3})[\)\.\:\-]\s+' sections = re.split(alt_pattern, response) - - # Fallback for single question batches or if headers were omitted if len(sections) < 3: - # If we expected multiple questions but got none, try a simpler digit-only header alt_pattern = r'(?mi)^\s*(\d{1,3})\s*$' sections = re.split(alt_pattern, response) + return sections - if len(sections) < 3: - if len(batch_indices) == 1: - results[batch_indices[0]] = self._parse_single_content(response) - return results - return results - + def _build_results_from_sections(self, sections: List[str]) -> Dict[int, Dict]: + """Requires len(sections) >= 3 (header split produced captures + bodies).""" + results: Dict[int, Dict] = {} for i in range(1, len(sections), 2): try: q_num = int(sections[i]) - content = sections[i+1] + content = sections[i + 1] results[q_num] = self._parse_single_content(content) except (ValueError, IndexError): continue - return results + def _parse_response(self, response: str, batch_indices: List[int]) -> Dict[int, Dict]: + """ + Parses the LLM output back into a structured dictionary using strict markers. + """ + self._trace_raw_llm_response(batch_indices, response) + sections = self._split_response_into_sections(response) + if len(sections) < 3: + if len(batch_indices) == 1: + return {batch_indices[0]: self._parse_single_content(response)} + return {} + return self._build_results_from_sections(sections) + def _parse_single_content(self, content: str) -> Dict: """Helper to parse markers from a single question block.""" # 0. Core Concept diff --git a/content_gen/tests/test_regression_guards.py b/content_gen/tests/test_regression_guards.py index 3c16242..d87d261 100644 --- a/content_gen/tests/test_regression_guards.py +++ b/content_gen/tests/test_regression_guards.py @@ -101,7 +101,7 @@ def test_wrapper_question_number_range_is_configurable(): def test_parse_response_returns_empty_for_multi_without_headers(): - generator = ContentGenerator(router=MagicMock()) + generator = ContentGenerator(router=MagicMock(), log_raw_llm_responses=False) parsed = generator._parse_response( "This is a combined response with no explicit question headers.", [1, 2, 3] diff --git a/qc_viewer/static/index.html b/qc_viewer/static/index.html index f3be1b0..1f86728 100644 --- a/qc_viewer/static/index.html +++ b/qc_viewer/static/index.html @@ -131,446 +131,11 @@

Ready to Verify

diff --git a/qc_viewer/static/js/api.js b/qc_viewer/static/js/api.js new file mode 100644 index 0000000..ab24377 --- /dev/null +++ b/qc_viewer/static/js/api.js @@ -0,0 +1,6 @@ +/** + * Shared path prefixes for QC viewer pages (standard viewer, automate, etc.). + */ +export const API_BASE = '/api'; +export const ROOT_PATH = '/'; +export const AUTOMATE_PATH = '/automate'; diff --git a/qc_viewer/static/js/controllers/viewer.js b/qc_viewer/static/js/controllers/viewer.js new file mode 100644 index 0000000..8e9c998 --- /dev/null +++ b/qc_viewer/static/js/controllers/viewer.js @@ -0,0 +1,463 @@ +/** + * Standard QC viewer: paper search, question list, detail + flashcards. + */ +import { API_BASE } from '/js/api.js'; + +const API = API_BASE; + +let allPapers = []; +let currentTable = null; + +window.fcCards = []; +window.fcIndex = 0; +window.fcFlipped = false; + +async function init() { + try { + const res = await fetch(`${API}/papers`); + if (!res.ok) throw new Error('failed'); + allPapers = await res.json(); + renderDropdown(allPapers); + } catch { + document.getElementById('questionList').innerHTML = + '
⚠ Could not connect to backend.
'; + } +} + +function renderDropdown(papers) { + const list = document.getElementById('dropdownList'); + list.innerHTML = ''; + if (!papers.length) { + list.innerHTML = '
No papers found.
'; + return; + } + papers.forEach((p) => { + const item = document.createElement('div'); + item.className = 'dropdown-item'; + const subject = p.table.replace('_questions', '').replace('igcse_', 'IGCSE '); + item.innerHTML = `${p.code}${subject}`; + item.onclick = () => window.selectPaper(p); + list.appendChild(item); + }); +} + +function bindPaperSelector() { + document.getElementById('searchInput').addEventListener('input', (e) => { + const q = e.target.value.toLowerCase(); + const filtered = allPapers.filter( + (p) => p.code.toLowerCase().includes(q) || p.table.toLowerCase().includes(q), + ); + renderDropdown(filtered); + document.getElementById('dropdownList').classList.add('open'); + }); + + document.getElementById('searchInput').addEventListener('focus', () => { + renderDropdown(allPapers); + document.getElementById('dropdownList').classList.add('open'); + }); + + document.addEventListener('click', (e) => { + if (!document.getElementById('paperSelector').contains(e.target)) { + document.getElementById('dropdownList').classList.remove('open'); + } + }); +} + +window.selectPaper = function (p) { + document.getElementById('selectedBadge').classList.add('visible'); + document.getElementById('badgeText').textContent = `${p.code} · ${p.table.replace('_questions', '').replace('igcse_', 'IGCSE ')}`; + document.getElementById('searchBox').classList.add('hidden'); + document.getElementById('searchInput').value = ''; + document.getElementById('dropdownList').classList.remove('open'); + currentTable = p.table; + loadQuestions(p.table, p.code); +}; + +window.clearSelection = function () { + document.getElementById('selectedBadge').classList.remove('visible'); + document.getElementById('searchBox').classList.remove('hidden'); + document.getElementById('searchInput').value = ''; + document.getElementById('searchInput').focus(); + document.getElementById('questionList').innerHTML = + '
Select a paper to load questions.
'; + document.getElementById('contentView').innerHTML = ` +
+
+ +
+
+
+ 🔬 +

Ready to Verify

+

Select a paper and question to begin QC review.

+
`; + currentTable = null; +}; + +async function loadQuestions(table, code) { + const list = document.getElementById('questionList'); + list.innerHTML = '
Loading questions…
'; + try { + const res = await fetch(`${API}/questions?table=${table}&paper_code=${encodeURIComponent(code)}`); + if (!res.ok) throw new Error('failed'); + const qs = await res.json(); + list.innerHTML = ''; + if (!qs.length) { + list.innerHTML = '
No questions found.
'; + return; + } + list.innerHTML = `
${qs.length} questions
`; + qs.forEach((q, i) => { + const hasContent = !!q.summary_explanation; + const isVerified = q.is_verified; + const div = document.createElement('div'); + div.className = 'q-item'; + div.innerHTML = ` + Q${q.question_number || i + 1} + + `; + div.onclick = () => { + renderQuestion(q, table); + setActive(div); + }; + list.appendChild(div); + }); + } catch { + list.innerHTML = '
Failed to load questions.
'; + } +} + +function setActive(el) { + document.querySelectorAll('.q-item').forEach((e) => e.classList.remove('active')); + el.classList.add('active'); +} + +function markdownToHtml(text) { + if (!text) return ''; + + const replaceSubSup = (match) => + match + .replace(/]*>(.*?)<\/sup>/g, '^{$1}') + .replace(/]*>(.*?)<\/sub>/g, '_{$1}') + .replace(/<\/?(p|strong|em|br)[^>]*>/g, ''); + + text = text.replace(/\\\([\s\S]*?\\\)/g, replaceSubSup); + text = text.replace(/\\\[[\s\S]*?\\\]/g, replaceSubSup); + text = text.replace(/\$\$[\s\S]*?\$\$/g, replaceSubSup); + + const hasHtmlBlocks = /<[a-z][\s\S]*>/i.test(text); + + if (hasHtmlBlocks) { + return text + .replace(/\*\*\*(.+?)\*\*\*/gs, '$1') + .replace(/\*\*(.+?)\*\*/gs, '$1') + .replace(/(?$1') + .replace(/__(.+?)__/g, '$1') + .replace(/(?)\n(?!<)/g, '
'); + } + + let html = text + .replace(/^### (.+)$/gm, '

$1

') + .replace(/^## (.+)$/gm, '

$1

') + .replace(/^# (.+)$/gm, '

$1

') + .replace(/\*\*\*(.+?)\*\*\*/g, '$1') + .replace(/\*\*(.+?)\*\*/g, '$1') + .replace(/\*(.+?)\*/g, '$1') + .replace(/__(.+?)__/g, '$1') + .replace(/_(.+?)_/g, '$1') + .replace(/`([^`]+)`/g, '$1') + .replace(/^---+$/gm, '
') + .replace(/^\s*[-*]\s+(.+)$/gm, '
  • $1
    • ') + .replace(/(
  • [\s\S]*?<\/li>)/g, '
      $1
    ') + .replace(/^\s*\d+\.\s+(.+)$/gm, '
  • $1
    • ') + .replace(/^> (.+)$/gm, '
    $1
    '); + + html = html + .split(/\n{2,}/) + .map((block) => { + const t = block.trim(); + if (!t) return ''; + if (/^<(h[1-3]|ul|ol|li|hr|blockquote)/.test(t)) return block; + return `

    ${block.replace(/\n/g, '
    ')}

    `; + }) + .join('\n'); + + return html; +} + +function typesetMath() { + if (window.MathJax && MathJax.typesetPromise) { + if (MathJax.startup && MathJax.startup.document) { + MathJax.startup.document.clear(); + MathJax.startup.document.updateDocument(); + } + MathJax.typesetPromise().catch(() => {}); + } +} + +function renderQuestion(q, table) { + const subject = (table || '').replace('_questions', '').replace('igcse_', 'IGCSE ').toUpperCase(); + const view = document.getElementById('contentView'); + + let oeHtml = ''; + if (q.option_explanations && q.option_explanations.length > 0) { + oeHtml = q.option_explanations + .map( + (exp, i) => ` +
    +
    Option ${String.fromCharCode(65 + i)} + ${q.correct_options && q.correct_options.includes(i) ? '✓ Correct' : ''} +
    +
    ${markdownToHtml(exp) || 'No detail.'}
    +
    `, + ) + .join(''); + } else { + oeHtml = '

    No option analysis generated yet.

    '; + } + + let diagramHtml = ''; + if (q.diagrams) { + const images = q.diagrams.split('\n').filter((u) => u.trim()); + if (images.length > 0) { + diagramHtml = `
    + ${images.map((url) => `Question diagram`).join('')} +
    `; + } + } + + const rawExp = q.detailed_explanation || q.summary_explanation; + const explanationHtml = rawExp + ? markdownToHtml(rawExp) + : '

    No explanation generated yet. Run the AI pipeline for this paper.

    '; + + view.innerHTML = ` +
    +
    + +
    +
    +
    +
    + ${subject} + ${q.question_identifier || ''} + ${q.is_verified ? '✓ Verified' : ''} +
    +

    Question ${q.question_number || ''}

    +
    ${q.title || '[No question text]'}
    + ${diagramHtml} +
    + ${(q.options || []) + .map( + (opt, i) => ` +
    +
    ${String.fromCharCode(65 + i)}
    +
    ${opt || '[empty]'}
    +
    `, + ) + .join('')} +
    +
    + +
    +
    +
    General Explanation
    +
    Option Analysis
    +
    Flashcards
    +
    + +
    +
    ${explanationHtml}
    +
    + +
    ${oeHtml}
    + +
    +
    +
    Click the Flashcards tab to load…
    +
    +
    +
    + `; + + view.querySelectorAll('.tab').forEach((tab) => { + tab.onclick = () => { + view.querySelectorAll('.tab').forEach((t) => t.classList.remove('active')); + view.querySelectorAll('.tab-content').forEach((c) => c.classList.remove('active')); + tab.classList.add('active'); + view.querySelector(`#${tab.dataset.tab}`).classList.add('active'); + if (tab.dataset.tab === 'tab-flashcards') { + loadFlashcards(q.topic_id, q.subtopic_id, q.id); + } + typesetMath(); + }; + }); + + typesetMath(); +} + +async function loadFlashcards(topicId, subtopicId, questionId) { + const container = document.getElementById('flashcardContainer'); + if (!container) return; + + let isFallback = false; + if (!topicId || topicId === 'unknown-topic') { + topicId = 'unknown-topic'; + isFallback = true; + } + + container.innerHTML = '
    Loading flashcards…
    '; + try { + let url = `${API}/flashcards?topic_id=${encodeURIComponent(topicId)}`; + if (subtopicId && subtopicId !== 'undefined') url += `&subtopic_id=${encodeURIComponent(subtopicId)}`; + if (questionId) url += `&question_id=${encodeURIComponent(questionId)}`; + + const res = await fetch(url); + if (!res.ok) throw new Error(`HTTP ${res.status}`); + const cards = await res.json(); + + if (!cards.length) { + container.innerHTML = ` +
    +
    🃏
    +

    No flashcards found for this question yet.

    +
    `; + return; + } + + window.fcCards = cards; + window.fcIndex = 0; + window.fcFlipped = false; + const isSpecific = cards.length > 0 && cards[0].is_specific; + renderFlashcardViewer(container, isFallback, isSpecific); + typesetMath(); + } catch (e) { + console.error('Flashcard Load Error:', e); + container.innerHTML = '

    Could not load flashcards.

    '; + } +} + +function renderFlashcardViewer(container, isFallback, isSpecific) { + const total = window.fcCards.length; + + const statusIndicator = isSpecific + ? `
    + ✨ Question-Specific AI Flashcards (${total}) +
    ` + : isFallback + ? `
    + ⚠️ Topic Fallback (${total} generic cards) +
    ` + : ''; + + container.innerHTML = ` + ${statusIndicator} +
    +
    +
    +
    +
    +
    +
    Question
    +
    +
    Tap to reveal answer
    +
    +
    +
    Answer
    +
    +
    Tap to flip back
    +
    +
    +
    +
    + +
    + +
    + +
    +
    `; + + document.getElementById('fcScene').onclick = window.flipCard; + document.getElementById('fcFlip').onclick = window.flipCard; + document.getElementById('fcPrev').onclick = window.prevCard; + document.getElementById('fcNext').onclick = window.nextCard; + + updateFlashcardDisplay(); +} + +function updateFlashcardDisplay() { + const total = window.fcCards.length; + const card = window.fcCards[window.fcIndex]; + const pipsEl = document.getElementById('fcPips'); + if (pipsEl) { + pipsEl.innerHTML = window.fcCards + .map((_, i) => { + const cls = i < window.fcIndex ? 'done' : i === window.fcIndex ? 'current' : ''; + return `
    `; + }) + .join(''); + } + const counter = document.getElementById('fcCounter'); + if (counter) counter.textContent = `${window.fcIndex + 1} / ${total}`; + const frontEl = document.getElementById('fcFrontText'); + const backEl = document.getElementById('fcBackText'); + if (frontEl) frontEl.innerHTML = markdownToHtml(card.frontText || 'No front text'); + if (backEl) backEl.innerHTML = markdownToHtml(card.backText || 'No answer'); + window.fcFlipped = false; + const cardEl = document.getElementById('fcCard'); + if (cardEl) cardEl.classList.remove('flipped'); + const prevBtn = document.getElementById('fcPrev'); + const nextBtn = document.getElementById('fcNext'); + if (prevBtn) prevBtn.disabled = window.fcIndex === 0; + if (nextBtn) nextBtn.disabled = window.fcIndex === total - 1; + document.onkeydown = handleFcKey; + typesetMath(); +} + +window.flipCard = function () { + window.fcFlipped = !window.fcFlipped; + const cardEl = document.getElementById('fcCard'); + if (cardEl) cardEl.classList.toggle('flipped', window.fcFlipped); +}; + +window.prevCard = function () { + if (window.fcIndex > 0) { + window.fcIndex--; + updateFlashcardDisplay(); + } +}; + +window.nextCard = function () { + if (window.fcIndex < window.fcCards.length - 1) { + window.fcIndex++; + updateFlashcardDisplay(); + } +}; + +function handleFcKey(e) { + if (e.key === 'ArrowRight' || e.key === 'ArrowDown') { + e.preventDefault(); + window.nextCard(); + } else if (e.key === 'ArrowLeft' || e.key === 'ArrowUp') { + e.preventDefault(); + window.prevCard(); + } else if (e.key === ' ' || e.key === 'Enter') { + e.preventDefault(); + window.flipCard(); + } +} + +/** + * Wire standard viewer after ThemeManager / LayoutManager init. + */ +export function bootStandardViewer() { + bindPaperSelector(); + init(); +} From f8b62e39424b83f8134cdf16310d06c212a05cd4 Mon Sep 17 00:00:00 2001 From: shmukit Date: Thu, 7 May 2026 23:53:00 +0600 Subject: [PATCH 3/8] feat(metrics): enhance metrics tracking with Redis support and improve file handling - Introduced a new `RedisMetricsTracker` class for optional Redis-backed metrics aggregation. - Refactored existing `FileMetricsTracker` and `MemoryMetricsTracker` to improve thread safety and data handling. - Updated `create_metrics_tracker` factory function to support different backend options (file, memory, redis). - Added tests for metrics tracking functionality, ensuring correct behavior across different backends. - Enhanced documentation to clarify architecture and usage of metrics tracking. --- content_gen/core/metrics.py | 183 +++++++++++++++--- content_gen/core/model_router.py | 4 +- content_gen/requirements.txt | 3 + content_gen/scripts/prompts.py | 3 + content_gen/tests/test_metrics.py | 33 ++++ docs/architecture.md | 59 ++++++ qc_viewer/app_factory.py | 3 + qc_viewer/config.py | 14 ++ qc_viewer/middleware/__init__.py | 1 + qc_viewer/middleware/security.py | 90 +++++++++ qc_viewer/routers/automation.py | 20 +- qc_viewer/services/automation_pipeline.py | 45 ++++- .../services/tests/test_automate_publish.py | 73 +++++++ .../tests/test_security_middleware.py | 64 ++++++ 14 files changed, 554 insertions(+), 41 deletions(-) create mode 100644 content_gen/tests/test_metrics.py create mode 100644 docs/architecture.md create mode 100644 qc_viewer/middleware/__init__.py create mode 100644 qc_viewer/middleware/security.py create mode 100644 qc_viewer/services/tests/test_automate_publish.py create mode 100644 qc_viewer/services/tests/test_security_middleware.py diff --git a/content_gen/core/metrics.py b/content_gen/core/metrics.py index 9595a09..4a03e91 100644 --- a/content_gen/core/metrics.py +++ b/content_gen/core/metrics.py @@ -1,59 +1,184 @@ +import importlib import json +import os +import threading +from abc import ABC, abstractmethod from pathlib import Path -from typing import Dict, Any +from typing import Any, Dict, Optional from datetime import datetime -class MetricsTracker: - """ - Tracks LLM usage, tokens, and costs across Edmate sessions. - Persists data to a local JSON file for UI consumption. - """ +class MetricsTrackerBase(ABC): + """LLM usage / cost tracking (file, memory, or Redis).""" + + @property + @abstractmethod + def metrics(self) -> Dict[str, Any]: + """Aggregate counters (mutable dict for file/memory; snapshot for redis).""" + ... + + @abstractmethod + def log_usage(self, response_obj: Any) -> None: + ... + + @abstractmethod + def get_current_cost(self) -> float: + ... + + def reset_if_new_day(self) -> None: + """Placeholder for daily reset logic if needed.""" + pass + + +class FileMetricsTracker(MetricsTrackerBase): + """Persists aggregate metrics to a JSON file (default; single-process friendly).""" def __init__(self, storage_path: str = "content_gen/data/session_metrics.json"): self.storage_path = Path(storage_path) self.storage_path.parent.mkdir(parents=True, exist_ok=True) - self.metrics = self._load_metrics() + self._lock = threading.Lock() + self._metrics = self._load_metrics() + + @property + def metrics(self) -> Dict[str, Any]: + return self._metrics def _load_metrics(self) -> Dict[str, Any]: if not self.storage_path.exists(): return {"total_cost": 0.0, "total_tokens": 0, "last_updated": None, "calls": 0} - try: - with open(self.storage_path, 'r') as f: + with open(self.storage_path, "r", encoding="utf-8") as f: return json.load(f) except Exception: return {"total_cost": 0.0, "total_tokens": 0, "last_updated": None, "calls": 0} - def log_usage(self, response_obj: Any): - """ - Extracts usage data from a LiteLLM response object and updates local metrics. - """ - # LiteLLM provides usage and cost in the response - usage = getattr(response_obj, 'usage', {}) + def log_usage(self, response_obj: Any) -> None: + usage = getattr(response_obj, "usage", {}) + try: + from litellm import completion_cost + + actual_cost = completion_cost(completion_response=response_obj) + except Exception: + actual_cost = 0.0 + + with self._lock: + self._metrics["total_cost"] += actual_cost + self._metrics["total_tokens"] += usage.get("total_tokens", 0) + self._metrics["calls"] = int(self._metrics.get("calls", 0)) + 1 + self._metrics["last_updated"] = datetime.now().isoformat() + self._save_metrics_unlocked() + + def _save_metrics_unlocked(self) -> None: + with open(self.storage_path, "w", encoding="utf-8") as f: + json.dump(self._metrics, f, indent=2) + + def get_current_cost(self) -> float: + with self._lock: + return float(self._metrics.get("total_cost", 0.0)) + + +class MemoryMetricsTracker(MetricsTrackerBase): + """In-process metrics only (safe default under multi-worker: no shared file corruption).""" - # Real cost from LiteLLM is usually in response._hidden_params or similar - # but LiteLLM also provides it in a more direct way if litellm.ModelResponse is used + def __init__(self) -> None: + self._lock = threading.Lock() + self._metrics: Dict[str, Any] = { + "total_cost": 0.0, + "total_tokens": 0, + "last_updated": None, + "calls": 0, + } + + @property + def metrics(self) -> Dict[str, Any]: + return self._metrics + + def log_usage(self, response_obj: Any) -> None: + usage = getattr(response_obj, "usage", {}) try: from litellm import completion_cost + actual_cost = completion_cost(completion_response=response_obj) except Exception: actual_cost = 0.0 - self.metrics["total_cost"] += actual_cost - self.metrics["total_tokens"] += usage.get("total_tokens", 0) - self.metrics["calls"] += 1 - self.metrics["last_updated"] = datetime.now().isoformat() + with self._lock: + self._metrics["total_cost"] += actual_cost + self._metrics["total_tokens"] += usage.get("total_tokens", 0) + self._metrics["calls"] = int(self._metrics.get("calls", 0)) + 1 + self._metrics["last_updated"] = datetime.now().isoformat() - self._save_metrics() + def get_current_cost(self) -> float: + with self._lock: + return float(self._metrics.get("total_cost", 0.0)) - def _save_metrics(self): - with open(self.storage_path, 'w') as f: - json.dump(self.metrics, f, indent=2) + +class RedisMetricsTracker(MetricsTrackerBase): + """Optional Redis-backed aggregates (hash: atomic HINCRBY*).""" + + def __init__(self, url: str, key: str = "edmate:metrics:session"): + try: + redis_mod = importlib.import_module("redis") + except ImportError as e: + raise RuntimeError("redis package required for Redis metrics backend") from e + self._r = redis_mod.from_url(url, decode_responses=True) + self._key = key + + @property + def metrics(self) -> Dict[str, Any]: + raw = self._r.hgetall(self._key) + if not raw: + return {"total_cost": 0.0, "total_tokens": 0, "last_updated": None, "calls": 0} + return { + "total_cost": float(raw.get("total_cost") or 0.0), + "total_tokens": int(raw.get("total_tokens") or 0), + "last_updated": raw.get("last_updated"), + "calls": int(raw.get("calls") or 0), + } + + def log_usage(self, response_obj: Any) -> None: + usage = getattr(response_obj, "usage", {}) + try: + from litellm import completion_cost + + actual_cost = completion_cost(completion_response=response_obj) + except Exception: + actual_cost = 0.0 + + tokens = int(usage.get("total_tokens", 0) or 0) + pipe = self._r.pipeline() + pipe.hincrbyfloat(self._key, "total_cost", float(actual_cost)) + pipe.hincrby(self._key, "total_tokens", tokens) + pipe.hincrby(self._key, "calls", 1) + pipe.hset(self._key, "last_updated", datetime.now().isoformat()) + pipe.execute() def get_current_cost(self) -> float: - return self.metrics.get("total_cost", 0.0) + v = self._r.hget(self._key, "total_cost") + return float(v or 0.0) - def reset_if_new_day(self): - """Placeholder for daily reset logic if needed.""" - pass + +def create_metrics_tracker( + storage_path: Optional[str] = None, +) -> MetricsTrackerBase: + """ + Factory: EDMATE_METRICS_BACKEND=file|memory|redis (default: file). + + - ``memory``: no disk; per-process (recommended for multi-worker until Redis is used). + - ``file``: legacy JSON file (single-worker or dev). + - ``redis``: requires EDMATE_METRICS_REDIS_URL and ``redis`` package. + """ + mode = (os.environ.get("EDMATE_METRICS_BACKEND") or "file").strip().lower() + if mode == "memory": + return MemoryMetricsTracker() + if mode == "redis": + url = (os.environ.get("EDMATE_METRICS_REDIS_URL") or "").strip() + if not url: + return MemoryMetricsTracker() + return RedisMetricsTracker(url) + path = storage_path or os.environ.get("EDMATE_METRICS_FILE") or "content_gen/data/session_metrics.json" + return FileMetricsTracker(path) + + +# Backward-compatible name used by ModelRoutingEngine and tests. +MetricsTracker = FileMetricsTracker diff --git a/content_gen/core/model_router.py b/content_gen/core/model_router.py index ea7e969..2c070ec 100644 --- a/content_gen/core/model_router.py +++ b/content_gen/core/model_router.py @@ -3,7 +3,7 @@ from typing import Optional, List, Dict, Any, cast from content_gen.core.config_schema import EdmateConfig from content_gen.core.config import CoreConfig -from content_gen.core.metrics import MetricsTracker +from content_gen.core.metrics import create_metrics_tracker # Set up callbacks for observability if requested # Users can set LITELLM_CALLBACKS=["opik"] in their .env @@ -31,7 +31,7 @@ def __init__( ): # Load config from YAML/JSON if not provided self.config = config or CoreConfig.load_from_yaml() - self.tracker = MetricsTracker() + self.tracker = create_metrics_tracker() # Per-request / BYOK key passed to litellm (never mutate os.environ). _k = (api_key or "").strip() self._api_key: Optional[str] = _k or None diff --git a/content_gen/requirements.txt b/content_gen/requirements.txt index d85256e..762fc12 100644 --- a/content_gen/requirements.txt +++ b/content_gen/requirements.txt @@ -18,6 +18,9 @@ azure-storage-blob>=12.19.0 # Azure Blob Storage # Database psycopg2-binary>=2.9.0 # PostgreSQL adapter +# Optional: shared LLM metrics (EDMATE_METRICS_BACKEND=redis) +redis>=4.6.0 + # Optional: API Integrations google-generativeai>=0.3.0 # Gemini API openai>=1.0.0 # OpenAI/ChatGPT API diff --git a/content_gen/scripts/prompts.py b/content_gen/scripts/prompts.py index c57bd84..b354f89 100644 --- a/content_gen/scripts/prompts.py +++ b/content_gen/scripts/prompts.py @@ -2,6 +2,9 @@ Central repository for system prompts used in the content generation pipeline. """ +# Bump when CONTENT_GENERATION_PROMPT meaningfully changes (QC / reproducibility). +CONTENT_GENERATION_PROMPT_VERSION = "2026-05-01" + # The primary system prompt for Gemini to generate educational content CONTENT_GENERATION_PROMPT = """ You are an expert [Subject] teacher specializing in [Curriculum] curriculum. diff --git a/content_gen/tests/test_metrics.py b/content_gen/tests/test_metrics.py new file mode 100644 index 0000000..fc3f217 --- /dev/null +++ b/content_gen/tests/test_metrics.py @@ -0,0 +1,33 @@ +import os + +import pytest + +from content_gen.core.metrics import ( + FileMetricsTracker, + MemoryMetricsTracker, + create_metrics_tracker, +) + + +def test_create_metrics_memory(monkeypatch): + monkeypatch.setenv("EDMATE_METRICS_BACKEND", "memory") + t = create_metrics_tracker() + assert isinstance(t, MemoryMetricsTracker) + + +def test_create_metrics_file_default(monkeypatch, tmp_path): + monkeypatch.delenv("EDMATE_METRICS_BACKEND", raising=False) + f = tmp_path / "m.json" + monkeypatch.setenv("EDMATE_METRICS_FILE", str(f)) + t = create_metrics_tracker() + assert isinstance(t, FileMetricsTracker) + + +@pytest.mark.skipif(not os.environ.get("RUN_REDIS_METRICS_TEST"), reason="set RUN_REDIS_METRICS_TEST=1 to run") +def test_create_metrics_redis_integration(monkeypatch): + url = os.environ["EDMATE_METRICS_REDIS_URL"] + monkeypatch.setenv("EDMATE_METRICS_BACKEND", "redis") + monkeypatch.setenv("EDMATE_METRICS_REDIS_URL", url) + pytest.importorskip("redis") + t = create_metrics_tracker() + assert t.get_current_cost() >= 0.0 diff --git a/docs/architecture.md b/docs/architecture.md new file mode 100644 index 0000000..993e429 --- /dev/null +++ b/docs/architecture.md @@ -0,0 +1,59 @@ +# Edmate architecture (qc_viewer + content_gen) + +This document is for operators and developers wiring the QC web app, automation API, and PDF/LLM pipeline. + +## High-level data flow + +1. **Browser** loads static assets from `qc_viewer/static/` and calls FastAPI routes under `/api/automate`, `/api/v1`, and question endpoints. +2. **Draft upload** (`POST /api/automate/draft`) writes `source.pdf` and `metadata.json` under the configured drafts directory, then schedules **`run_automation_pipeline`** as a FastAPI `BackgroundTasks` job (same process as the web worker). +3. **Pipeline** (`qc_viewer/services/automation_pipeline.py`) builds a **`ModelRoutingEngine`** (`content_gen/core/model_router.py`), runs **`PipelineOrchestrator`** (`content_gen/scripts/pipeline/pipeline_orchestrator.py`) for extraction + generation, normalizes rows via **`build_legacy_question_dict`**, and writes results back into **`metadata.json`**. +4. **LLM calls** go through **LiteLLM** (`litellm.completion`). Usage/cost aggregates are recorded by **`create_metrics_tracker()`** (`content_gen/core/metrics.py`). +5. **Publish** (`POST /api/automate/publish`) validates `table_name` against the workspace allowlist (`qc_viewer/config.py`), then uses **`DatabaseService`** (`content_gen/scripts/processing/database_service.py`) to insert into Postgres. +6. **Prompts** for generation live in **`content_gen/scripts/prompts.py`**. The constant **`CONTENT_GENERATION_PROMPT_VERSION`** is copied into draft **`metadata.json`** and pipeline telemetry for reproducibility. + +```mermaid +flowchart LR + UI[Static_UI] --> API[FastAPI] + API --> Disk[Drafts_dir] + API --> PG[(Postgres)] + API --> Pipe[automation_pipeline] + Pipe --> Orch[PipelineOrchestrator] + Orch --> LiteLLM[litellm] + LiteLLM --> Metrics[metrics_tracker] +``` + +## Single-process and scaling caveats + +- **Background tasks** run **in-process**. Multiple Uvicorn/Gunicorn workers each have their own background queue; a draft started on worker A is not visible to worker B’s memory. Prefer **one worker** for predictable automation, or front the API with sticky sessions, or move long jobs to an **external queue** (RQ/Celery/SQS — not bundled here). +- **Metrics**: `EDMATE_METRICS_BACKEND=file` writes JSON under `content_gen/data/` (or `EDMATE_METRICS_FILE`). Multiple workers **corrupt or skew** that file. Use **`memory`** per process (no cross-worker budget) or **`redis`** with `EDMATE_METRICS_REDIS_URL` for shared aggregates. +- **Rate limit** (`EDMATE_RATE_LIMIT_PER_MINUTE`) is **in-process per server IP**; it is not a global cluster limit. +- **Draft storage** and optional **SQLite job store** (if used) are **local disk** unless you configure shared storage / DB. + +## Where to change things + +| Concern | Primary location | +|--------|-------------------| +| HTTP automation routes | `qc_viewer/routers/automation.py` | +| Optional API key + rate limit | `qc_viewer/middleware/security.py` (env below) | +| Pipeline steps / orchestration | `content_gen/scripts/pipeline/pipeline_orchestrator.py` | +| Model routing, budget kill-switch | `content_gen/core/model_router.py` | +| System prompts + `prompt_version` | `content_gen/scripts/prompts.py` | +| Postgres insert shapes (publish/import) | `content_gen/scripts/processing/database_service.py`, adapters under `content_gen/adapters/` | +| Allowed tables for QC / publish | `qc_viewer/config.py` (`get_allowed_table_ids`, `is_publish_table_allowed`) | + +## Deploy environment variables (security & ops) + +| Variable | Purpose | +|----------|---------| +| `EDMATE_REQUIRE_API_KEY` | If truthy (`1`, `true`, …), require `EDMATE_API_KEY` on `/api/automate/*` and `/api/v1/*`. | +| `EDMATE_API_KEY` | Shared secret; client sends `X-API-Key` or `Authorization: Bearer `. | +| `EDMATE_RATE_LIMIT_PER_MINUTE` | Integer cap per client IP per rolling minute (`0` = off). Returns **429** when exceeded. | +| `EDMATE_METRICS_BACKEND` | `file` (default), `memory`, or `redis`. | +| `EDMATE_METRICS_FILE` | Override path when backend is `file`. | +| `EDMATE_METRICS_REDIS_URL` | Redis URL when backend is `redis` (requires `redis` package). | + +**Internal lab**: leave API key off, single worker, `file` metrics. **Shared or public network**: set `EDMATE_REQUIRE_API_KEY`, tighten CORS in `qc_viewer/app_factory.py`, use `redis` metrics if multiple workers, and plan an external worker queue for LLM-heavy load. + +## Related doc + +- Deeper technical notes: `docs/technical/ARCHITECTURE.md` (if present in your tree). diff --git a/qc_viewer/app_factory.py b/qc_viewer/app_factory.py index d3dc976..a0d0950 100644 --- a/qc_viewer/app_factory.py +++ b/qc_viewer/app_factory.py @@ -5,6 +5,7 @@ from fastapi.staticfiles import StaticFiles from qc_viewer.config import DOCS_ROOT, STATIC_ROOT +from qc_viewer.middleware.security import EdmateApiSecurityMiddleware from qc_viewer.router_v1 import router as api_v1_router from qc_viewer.routers.automation import router as automation_router from qc_viewer.routers.questions import router as questions_router @@ -40,6 +41,8 @@ def create_app() -> FastAPI: ], ) + app.add_middleware(EdmateApiSecurityMiddleware) + app.include_router(api_v1_router) app.include_router(static_pages_router) app.include_router(questions_router) diff --git a/qc_viewer/config.py b/qc_viewer/config.py index cae778d..fea02ad 100644 --- a/qc_viewer/config.py +++ b/qc_viewer/config.py @@ -62,6 +62,20 @@ def get_allowed_table_ids() -> Tuple[str, ...]: return tuple(_LEGACY_TABLES) +def is_publish_table_allowed(table_name: str) -> bool: + """ + Tables allowed for /api/automate/publish (defense against SQL injection in dynamic INSERT). + + Includes Prisma-style ``questions`` plus workspace ``target_tables`` / legacy list. + """ + t = (table_name or "").strip() + if not t or not _TABLE_ID_RE.match(t): + return False + if t == "questions": + return True + return t in get_allowed_table_ids() + + def get_workspace_defaults() -> tuple[str, str]: """(default_curriculum, default_subject) from edmate_config workspace.""" try: diff --git a/qc_viewer/middleware/__init__.py b/qc_viewer/middleware/__init__.py new file mode 100644 index 0000000..50ca4f8 --- /dev/null +++ b/qc_viewer/middleware/__init__.py @@ -0,0 +1 @@ +"""Starlette / FastAPI middleware for qc_viewer.""" diff --git a/qc_viewer/middleware/security.py b/qc_viewer/middleware/security.py new file mode 100644 index 0000000..0520f2d --- /dev/null +++ b/qc_viewer/middleware/security.py @@ -0,0 +1,90 @@ +""" +Optional API key gate and simple in-process rate limits for sensitive routes. + +Environment (see docs/architecture.md): + EDMATE_REQUIRE_API_KEY — set to 1/true to require a shared secret on protected paths. + EDMATE_API_KEY — shared secret; client sends X-API-Key or Authorization: Bearer . + EDMATE_RATE_LIMIT_PER_MINUTE — max requests per client IP per rolling minute (0 = disabled). +""" + +from __future__ import annotations + +import os +import threading +import time +from collections import defaultdict, deque +from typing import Deque, Dict + +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import JSONResponse + +_PROTECTED_PREFIXES: tuple[str, ...] = ("/api/automate", "/api/v1") + +_rate_lock = threading.Lock() +_rate_buckets: Dict[str, Deque[float]] = defaultdict(deque) + + +def _truthy(val: str | None) -> bool: + return (val or "").strip().lower() in ("1", "true", "yes", "on") + + +def _client_id(request: Request) -> str: + if request.client and request.client.host: + return request.client.host + return "unknown" + + +def _rate_limit_allows(client_id: str, limit_per_minute: int) -> bool: + if limit_per_minute <= 0: + return True + now = time.monotonic() + window = 60.0 + with _rate_lock: + dq = _rate_buckets[client_id] + while dq and dq[0] < now - window: + dq.popleft() + if len(dq) >= limit_per_minute: + return False + dq.append(now) + return True + + +class EdmateApiSecurityMiddleware(BaseHTTPMiddleware): + """Require API key (optional) and apply per-IP rate limits on automate + v1 API.""" + + async def dispatch(self, request: Request, call_next): + if request.method == "OPTIONS": + return await call_next(request) + + path = request.url.path + if not any(path.startswith(prefix) for prefix in _PROTECTED_PREFIXES): + return await call_next(request) + + if _truthy(os.environ.get("EDMATE_REQUIRE_API_KEY")): + secret = (os.environ.get("EDMATE_API_KEY") or "").strip() + if not secret: + return JSONResponse( + status_code=503, + content={"detail": "EDMATE_REQUIRE_API_KEY is set but EDMATE_API_KEY is empty"}, + ) + token = (request.headers.get("x-api-key") or "").strip() + auth = request.headers.get("authorization") or "" + if auth.lower().startswith("bearer "): + token = token or auth[7:].strip() + if token != secret: + return JSONResponse(status_code=401, content={"detail": "Unauthorized"}) + + try: + rlim = int((os.environ.get("EDMATE_RATE_LIMIT_PER_MINUTE") or "0").strip()) + except ValueError: + rlim = 0 + if rlim > 0: + cid = _client_id(request) + if not _rate_limit_allows(cid, rlim): + return JSONResponse( + status_code=429, + content={"detail": "Rate limit exceeded; try again shortly"}, + ) + + return await call_next(request) diff --git a/qc_viewer/routers/automation.py b/qc_viewer/routers/automation.py index 1816b78..215266c 100644 --- a/qc_viewer/routers/automation.py +++ b/qc_viewer/routers/automation.py @@ -11,8 +11,9 @@ from fastapi.responses import Response, StreamingResponse from content_gen.core.model_router import ModelRoutingEngine +from content_gen.scripts.prompts import CONTENT_GENERATION_PROMPT_VERSION from content_gen.scripts.processing.database_service import DatabaseService -from qc_viewer.config import DRAFTS_ROOT +from qc_viewer.config import DRAFTS_ROOT, is_publish_table_allowed from qc_viewer.services.automation_pipeline import CANCELLATION_EVENTS, run_automation_pipeline from pydantic import BaseModel from qc_viewer.services import draft_export @@ -85,6 +86,8 @@ async def receive_draft( static_dir / "metadata.json", { "id": draft_id, + "pipeline_job_id": draft_id, + "prompt_version": CONTENT_GENERATION_PROMPT_VERSION, "subject": subject, "paper_code": paper_code, "filename": file.filename, @@ -249,6 +252,8 @@ async def delete_draft(draft_id: str): async def publish_draft(request: PublishRequest): if not request.table_name or not request.question_data: raise HTTPException(status_code=400, detail="Missing table_name or question_data") + if not is_publish_table_allowed(request.table_name): + raise HTTPException(status_code=400, detail="Invalid table_name") db = DatabaseService() try: @@ -274,10 +279,21 @@ async def get_metrics(): except Exception: continue + active_drafts = 0 + if drafts_root.exists(): + for d in drafts_root.iterdir(): + if d.is_dir() and (d / "metadata.json").exists(): + try: + st = read_json(d / "metadata.json") + if st.get("status") == "PROCESSING": + active_drafts += 1 + except Exception: + continue + return { "total_cost": total_cost, "total_tokens": total_tokens, - "active_drafts": 5, + "active_drafts": active_drafts, "timestamp": datetime.now().isoformat(), } diff --git a/qc_viewer/services/automation_pipeline.py b/qc_viewer/services/automation_pipeline.py index 25d07a9..75a2220 100644 --- a/qc_viewer/services/automation_pipeline.py +++ b/qc_viewer/services/automation_pipeline.py @@ -1,4 +1,5 @@ import asyncio +import logging import threading import time from datetime import datetime @@ -14,15 +15,20 @@ class _MutableModelRouting(Protocol): extraction: str generation: str validation: str + + from content_gen.core.config_schema import DetectionMode from content_gen.core.pedagogy_engine import PedagogyEngine from content_gen.scripts.pipeline.pipeline_orchestrator import PipelineOrchestrator +from content_gen.scripts.prompts import CONTENT_GENERATION_PROMPT_VERSION from qc_viewer.services.draft_store import read_modify_write_json from qc_viewer.services.legacy_question_payload import build_legacy_question_dict CANCELLATION_EVENTS: dict[str, threading.Event] = {} +_logger = logging.getLogger(__name__) + def _normalize_model_id(model_id: str, provider: Optional[str]) -> str: """Normalize model identifiers into litellm's provider/model format when possible.""" @@ -142,9 +148,14 @@ def _mut(meta: dict) -> None: read_modify_write_json(meta_path, _mut) except Exception as e: - print(f"Error updating progress: {e}") + _logger.warning("progress_update_failed draft_id=%s err=%s", draft_id, e) try: + _logger.info( + "automation_pipeline_start draft_id=%s prompt_version=%s", + draft_id, + CONTENT_GENERATION_PROMPT_VERSION, + ) router = ModelRoutingEngine(api_key=api_key) if curriculum is None or not str(curriculum).strip(): curriculum = ( @@ -187,7 +198,12 @@ def _mut(meta: dict) -> None: progress_callback=_update_progress, ) t_extraction_end = time.time() - print(f"DEBUG: Extracted {len(extracted)} questions from PDF in {t_extraction_end - t_extraction_start:.2f}s.") + _logger.info( + "extraction_done draft_id=%s questions=%d duration_sec=%.2f", + draft_id, + len(extracted), + t_extraction_end - t_extraction_start, + ) _update_progress(60, "Applying Learning Science & Pedagogy Analysis...", processed_count=0, total_count=len(extracted)) @@ -200,7 +216,12 @@ def _mut(meta: dict) -> None: pedagogy_system_prompt=pedagogy_system_prompt, ) t_generation_end = time.time() - print(f"DEBUG: Generated content for {len(generated)} questions in {t_generation_end - t_generation_start:.2f}s.") + _logger.info( + "generation_done draft_id=%s questions=%d duration_sec=%.2f", + draft_id, + len(generated), + t_generation_end - t_generation_start, + ) t_normalization_start = time.time() total_questions = len(generated) @@ -215,11 +236,13 @@ def _mut(meta: dict) -> None: telemetry = { "total_time_sec": round(t_pipeline_end - t_pipeline_start, 2), + "prompt_version": CONTENT_GENERATION_PROMPT_VERSION, + "pipeline_job_id": draft_id, "nodes": { "extraction_sec": round(t_extraction_end - t_extraction_start, 2), "generation_sec": round(t_generation_end - t_generation_start, 2), "normalization_sec": round(t_normalization_end - t_normalization_start, 2), - } + }, } def _finalize(meta: dict) -> None: @@ -232,6 +255,8 @@ def _finalize(meta: dict) -> None: "total_count": total_questions, "status_message": "Generation complete!", "id": draft_id, + "pipeline_job_id": draft_id, + "prompt_version": CONTENT_GENERATION_PROMPT_VERSION, "subject": subject, "paper_code": paper_code, "pedagogy_profile": pedagogy.get_profile_summary(), @@ -244,7 +269,7 @@ def _finalize(meta: dict) -> None: read_modify_write_json(meta_path, _finalize) except InterruptedError: - print(f"Task {draft_id} cancelled.") + _logger.info("automation_pipeline_cancelled draft_id=%s", draft_id) try: read_modify_write_json( meta_path, @@ -257,10 +282,14 @@ def _finalize(meta: dict) -> None: ), ) except Exception as inner_e: - print(f"Failed to update metadata after cancel: {inner_e}") + _logger.warning("metadata_update_failed_after_cancel draft_id=%s err=%s", draft_id, inner_e) except Exception as e: error_str = str(e) - print(f"Background Processing Error: {error_str}") + _logger.exception( + "automation_pipeline_failed draft_id=%s prompt_version=%s", + draft_id, + CONTENT_GENERATION_PROMPT_VERSION, + ) try: def _fail(meta: dict) -> None: @@ -275,4 +304,4 @@ def _fail(meta: dict) -> None: read_modify_write_json(meta_path, _fail) except Exception as inner_e: - print(f"Failed to update metadata with error: {inner_e}") + _logger.warning("metadata_update_failed_after_error draft_id=%s err=%s", draft_id, inner_e) diff --git a/qc_viewer/services/tests/test_automate_publish.py b/qc_viewer/services/tests/test_automate_publish.py new file mode 100644 index 0000000..5318a9a --- /dev/null +++ b/qc_viewer/services/tests/test_automate_publish.py @@ -0,0 +1,73 @@ +"""Regression: publish endpoint must reject non-allowlisted table names.""" + +from unittest.mock import MagicMock, patch + +import pytest +from starlette.testclient import TestClient + + +@pytest.fixture +def client(): + from qc_viewer.app_factory import create_app + + return TestClient(create_app(), raise_server_exceptions=False) + + +def test_publish_rejects_injection_like_table_name(client: TestClient): + with patch("qc_viewer.routers.automation.DatabaseService") as MockSvc: + inst = MagicMock() + MockSvc.return_value = inst + res = client.post( + "/api/automate/publish", + json={ + "draft_id": "draft_x", + "table_name": "chemistry_questions; DROP TABLE users;--", + "question_data": {"question_identifier": "x", "title": "t"}, + }, + ) + assert res.status_code == 400 + assert "Invalid" in (res.json().get("detail") or "") + inst.inject_question.assert_not_called() + + +def test_publish_accepts_allowlisted_table(client: TestClient): + from qc_viewer.config import get_allowed_table_ids + + allowed = list(get_allowed_table_ids()) + assert allowed, "test needs at least one workspace table from config or legacy defaults" + table_name = allowed[0] + with patch("qc_viewer.routers.automation.DatabaseService") as MockSvc: + inst = MagicMock() + inst.inject_question.return_value = True + MockSvc.return_value = inst + res = client.post( + "/api/automate/publish", + json={ + "draft_id": "draft_x", + "table_name": table_name, + "question_data": {"question_identifier": "x", "title": "t"}, + }, + ) + assert res.status_code == 200, res.text + inst.inject_question.assert_called_once() + args, _ = inst.inject_question.call_args + assert args[0] == table_name + + +def test_publish_accepts_questions_table(client: TestClient): + with patch("qc_viewer.routers.automation.DatabaseService") as MockSvc: + inst = MagicMock() + inst.inject_question.return_value = True + MockSvc.return_value = inst + res = client.post( + "/api/automate/publish", + json={ + "draft_id": "draft_x", + "table_name": "questions", + "question_data": {"question_identifier": "x", "title": "t"}, + }, + ) + assert res.status_code == 200 + inst.inject_question.assert_called_once() + args, _ = inst.inject_question.call_args + assert args[0] == "questions" diff --git a/qc_viewer/services/tests/test_security_middleware.py b/qc_viewer/services/tests/test_security_middleware.py new file mode 100644 index 0000000..39f2c14 --- /dev/null +++ b/qc_viewer/services/tests/test_security_middleware.py @@ -0,0 +1,64 @@ +"""Regression tests for optional API key and rate limiting middleware.""" + +import os + +import pytest +from fastapi.testclient import TestClient + +from qc_viewer.app_factory import create_app + + +@pytest.fixture +def client(monkeypatch): + monkeypatch.delenv("EDMATE_REQUIRE_API_KEY", raising=False) + monkeypatch.delenv("EDMATE_API_KEY", raising=False) + monkeypatch.setenv("EDMATE_RATE_LIMIT_PER_MINUTE", "0") + with TestClient(create_app()) as c: + yield c + + +def test_automate_route_ok_without_key_when_not_required(client): + r = client.get("/api/automate/config") + assert r.status_code == 200 + + +def test_api_key_required_401(monkeypatch): + monkeypatch.setenv("EDMATE_REQUIRE_API_KEY", "1") + monkeypatch.setenv("EDMATE_API_KEY", "secret-test-key") + monkeypatch.setenv("EDMATE_RATE_LIMIT_PER_MINUTE", "0") + with TestClient(create_app()) as c: + r = c.get("/api/automate/config") + assert r.status_code == 401 + r2 = c.get("/api/automate/config", headers={"X-API-Key": "secret-test-key"}) + assert r2.status_code == 200 + + +def test_api_key_bearer_accepted(monkeypatch): + monkeypatch.setenv("EDMATE_REQUIRE_API_KEY", "1") + monkeypatch.setenv("EDMATE_API_KEY", "tok") + monkeypatch.setenv("EDMATE_RATE_LIMIT_PER_MINUTE", "0") + with TestClient(create_app()) as c: + r = c.get("/api/automate/config", headers={"Authorization": "Bearer tok"}) + assert r.status_code == 200 + + +def test_rate_limit_429(monkeypatch): + from qc_viewer.middleware import security as sec + + sec._rate_buckets.clear() + monkeypatch.delenv("EDMATE_REQUIRE_API_KEY", raising=False) + monkeypatch.setenv("EDMATE_RATE_LIMIT_PER_MINUTE", "2") + with TestClient(create_app()) as c: + assert c.get("/api/automate/config").status_code == 200 + assert c.get("/api/automate/config").status_code == 200 + r = c.get("/api/automate/config") + assert r.status_code == 429 + + +def test_require_key_without_secret_returns_503(monkeypatch): + monkeypatch.setenv("EDMATE_REQUIRE_API_KEY", "1") + monkeypatch.delenv("EDMATE_API_KEY", raising=False) + monkeypatch.setenv("EDMATE_RATE_LIMIT_PER_MINUTE", "0") + with TestClient(create_app()) as c: + r = c.get("/api/automate/config") + assert r.status_code == 503 From 97f9a1f520068931f5bb86e5a5808b90b5c1b8a9 Mon Sep 17 00:00:00 2001 From: shmukit Date: Thu, 14 May 2026 23:06:43 +0600 Subject: [PATCH 4/8] feat: set vision AI as default extraction engine --- content_gen/core/config_schema.py | 2 +- content_gen/core/schemas.py | 2 +- edmate_config.json.example | 2 +- edmate_config.yaml | 2 +- edmate_config.yaml.example | 2 +- qc_viewer/routers/automation.py | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/content_gen/core/config_schema.py b/content_gen/core/config_schema.py index fa1dc31..7d91a3f 100644 --- a/content_gen/core/config_schema.py +++ b/content_gen/core/config_schema.py @@ -66,7 +66,7 @@ class StorageSettings(BaseModel): class ExtractionSettings(BaseModel): - engine: ExtractionEngine = ExtractionEngine.PDF_EXTRACT_KIT + engine: ExtractionEngine = ExtractionEngine.VISION min_question_number: int = 1 max_question_number: Optional[int] = None question_detection_mode: DetectionMode = DetectionMode.BALANCED diff --git a/content_gen/core/schemas.py b/content_gen/core/schemas.py index 40e97ee..9cf76e7 100644 --- a/content_gen/core/schemas.py +++ b/content_gen/core/schemas.py @@ -66,7 +66,7 @@ class ModelConfig(BaseModel): validation_model: str = "openai/gpt-4o" max_budget: float = 10.0 # USD daily cap placeholder image_mode: str = "cdn" # "cdn" or "base64" - extraction_engine: str = "pdf_extract_kit" # "pdf_extract_kit" | "pymupdf" | "vision" | "multimodal" + extraction_engine: str = "vision" # "pdf_extract_kit" | "pymupdf" | "vision" | "multimodal" min_question_number: int = 1 max_question_number: Optional[int] = None question_detection_mode: str = "balanced" # "strict" | "balanced" | "open" diff --git a/edmate_config.json.example b/edmate_config.json.example index 609ca6f..9977139 100644 --- a/edmate_config.json.example +++ b/edmate_config.json.example @@ -11,7 +11,7 @@ "image_mode": "cdn" }, "extraction_settings": { - "engine": "pdf_extract_kit", + "engine": "vision", "min_question_number": 1, "max_question_number": 40, "question_detection_mode": "balanced" diff --git a/edmate_config.yaml b/edmate_config.yaml index c37fd86..8744d72 100644 --- a/edmate_config.yaml +++ b/edmate_config.yaml @@ -10,7 +10,7 @@ storage_settings: image_mode: "base64" extraction_settings: - engine: "pdf_extract_kit" + engine: "vision" # Optional: override boilerplate regexes for PDF text cleanup (empty list = skip these only) # extraction_noise_patterns: [] segmentation_preset: "bangladeshi" diff --git a/edmate_config.yaml.example b/edmate_config.yaml.example index 1a906f2..1d46372 100644 --- a/edmate_config.yaml.example +++ b/edmate_config.yaml.example @@ -19,7 +19,7 @@ storage_settings: image_mode: "base64" extraction_settings: - engine: "pdf_extract_kit" + engine: "vision" min_question_number: 1 max_question_number: null question_detection_mode: "balanced" diff --git a/qc_viewer/routers/automation.py b/qc_viewer/routers/automation.py index 215266c..6df98a2 100644 --- a/qc_viewer/routers/automation.py +++ b/qc_viewer/routers/automation.py @@ -328,7 +328,7 @@ async def get_config(): kit_path = Path(PROJECT_ROOT) / "content_gen" / "tools" / "PDF-Extract-Kit" kit_present = kit_path.is_dir() and (kit_path / "pdf_extract_kit").is_dir() - engine = (extraction_settings.get("engine") or "pdf_extract_kit").lower() + engine = (extraction_settings.get("engine") or "vision").lower() extraction_hints: dict[str, str] = {} if engine == "pdf_extract_kit": extraction_hints["summary"] = ( From 71cd7a50b0c72573d61e6a9c77e09b4d3eb4c281 Mon Sep 17 00:00:00 2001 From: shmukit Date: Wed, 20 May 2026 13:19:41 +0600 Subject: [PATCH 5/8] feat(ui): implement marketing page, global nav/footer, BYOK modal, and export paywall --- qc_viewer/routers/static_pages.py | 5 + qc_viewer/static/analytics.html | 38 +- qc_viewer/static/automate.html | 95 +++- qc_viewer/static/css/automate.css | 1 + qc_viewer/static/css/modules/marketing.css | 606 +++++++++++++++++++++ qc_viewer/static/how_it_works.html | 32 +- qc_viewer/static/index.html | 367 ++++++++----- qc_viewer/static/js/automate_ui.js | 96 +++- qc_viewer/static/js/controllers/review.js | 15 + qc_viewer/static/viewer.html | 172 ++++++ 10 files changed, 1292 insertions(+), 135 deletions(-) create mode 100644 qc_viewer/static/css/modules/marketing.css create mode 100644 qc_viewer/static/viewer.html diff --git a/qc_viewer/routers/static_pages.py b/qc_viewer/routers/static_pages.py index f5f7991..64bc30a 100644 --- a/qc_viewer/routers/static_pages.py +++ b/qc_viewer/routers/static_pages.py @@ -21,6 +21,11 @@ async def serve_root(): return serve_static_html("index.html", "index.html not found") +@router.get("/viewer") +async def serve_viewer(): + return serve_static_html("viewer.html", "viewer.html not found") + + @router.get("/automate") async def serve_hub(): return serve_static_html("automate.html", "automate.html not found") diff --git a/qc_viewer/static/analytics.html b/qc_viewer/static/analytics.html index 588bee6..adeee08 100644 --- a/qc_viewer/static/analytics.html +++ b/qc_viewer/static/analytics.html @@ -18,12 +18,12 @@ @@ -133,6 +154,17 @@

    Pipeline Analytics

    Recent Draft Activity

    Loading...
    + + diff --git a/qc_viewer/static/automate.html b/qc_viewer/static/automate.html index 3879ea4..41b780a 100644 --- a/qc_viewer/static/automate.html +++ b/qc_viewer/static/automate.html @@ -86,12 +86,12 @@

    Live Preview

    @@ -192,6 +206,17 @@

    Recent Processing Drafts

    Loading drafts...
    + + @@ -460,6 +485,70 @@

    Image Editor

    + + + + + + - + + + Edmate AI | Open-Source Assessment Automation & Pedagogy Engine + + + + + + - + + + + - + - - - - -
    -
    - + +
    + + + + + Edmate AI + + +
    + + + + + GitHub + + Automation Hub +
    +
    -
    -
    -

    Dashboard

    -

    Quality control and verification workspace.

    + +
    +
    + ✨ Open-Source & Enterprise Ready +
    +

    Automate Assessment Workflows with Learning Science

    +

    + Ingest PDFs, extract complex questions with layout-aware vision engines, enrich them with spaced repetition explanations, and PERSIST directly into your databases. +

    +
    + Launch Automation Hub + Star on GitHub
    -
    - -
    - +
    +

    Visual AI Engine

    +

    + Experience Edmate's high-speed Vision Extraction pipeline. Paste your API keys directly to operate fully locally and perform layout-aware extractions for FREE. +

    + Try the Pipeline +
    -
    + + + +
    +

    Engineered for Technical Integrity

    +

    We decouple raw model outputs from rigid schemas to serve curriculum-aligned educational pipelines.

    + +
    + +
    +
    🔑
    +

    Bring Your Own Key (BYOK)

    +

    Plug in your personal LLM API credentials (Gemini, OpenAI, Anthropic). We never persist your keys on our servers — everything operates directly in your browser session.

    +
    + +
    +
    🛡️
    +

    High-Integrity Assessments

    +

    Generate authentic, generative-AI resilient Q&As: AI Critique exercises, Viva Defense probing, and Scaffolded Sequences that combat direct student cheating.

    +
    + +
    +
    💾
    +

    Database Agnostic

    +

    Designed using the Adapter Pattern. Persist approved questions seamlessly into Postgres/Supabase, SQL/NoSQL databases, or export them instantly to CSV, JSON, Word, or ZIP.

    +
    + +
    +
    ⚡
    +

    Open-Core Architecture

    +

    Keep your workflows entirely local. Technicians can clone the repository, run the CLI orchestrator, or deploy the FastAPI server natively on their local devices.

    +
    +
    +
    + + +
    +
    +

    Simple, Dynamic Pricing Plans

    +

    Host the entire engine locally for free, bring your own API keys, or use our high-speed hosted routers.

    + +
    + +
    +
    Free & Open Source
    +

    For developers, hackers, and self-hosted environments. Run the platform entirely on your PC.

    +
    + $0 + / forever +
    +
      +
    • + + Self-Host CLI or local Server +
      • +
    • + + Bring Your Own Key (BYOK) on Web +
      • +
    • + + Full Layout-Aware Vision Extraction +
      • +
    • + + Unlimited Document Exports (with BYOK) +
      • +
    + Star & Clone Repo +
    -
    - 🔬 -

    Ready to Verify

    -

    Select a paper from the sidebar to begin your QC review process.

    -
    -
    -
    + +
    +
    Basic Solutions Plan
    +

    For independent educators and publishers who want a simple, ready-to-run hosted web editor.

    +
    + $10 + / month +
    +
      +
    • + + Access Default Hosted Gemini LLM +
      • +
    • + + Unlock Web Export (CSV/JSON/MD) +
      • +
    • + + Standard API Response Latency +
      • +
    • + + Save Up to 50 active drafts +
      • +
    + Activate Basic Plan +
    + + +
    +
    Pro Solutions Plan
    +

    For institutions and scale operations needing high throughput, embedded graphics, and premium Word formats.

    +
    + $50 + / month +
    +
      +
    • + + All Hosted Default Models Included +
      • +
    • + + Advanced exports: Word & Image ZIPs +
      • +
    • + + High-concurrency parallel ingestion +
      • +
    • + + Dedicated database schema slots +
      • +
    + Get Started Pro +
    + + + - diff --git a/qc_viewer/static/js/automate_ui.js b/qc_viewer/static/js/automate_ui.js index 0258036..7db4086 100644 --- a/qc_viewer/static/js/automate_ui.js +++ b/qc_viewer/static/js/automate_ui.js @@ -19,6 +19,7 @@ export const AutomationUI = { async init() { ThemeManager.init(); this.setupEventListeners(); + this.initByokConfig(); await this.fetchPipelineConfig(); this.fetchDrafts(); window.addEventListener('hashchange', () => this.checkHash()); @@ -26,6 +27,27 @@ export const AutomationUI = { setTimeout(() => this.checkHash(), 400); }, + initByokConfig() { + const savedProvider = localStorage.getItem('edmate-byok-provider'); + const savedKey = localStorage.getItem('edmate-byok-key'); + const savedModel = localStorage.getItem('edmate-byok-model'); + const apiPref = localStorage.getItem('edmate-api-preference'); + + const sidebarProvider = document.getElementById('byokProviderSelect'); + const sidebarKey = document.getElementById('byokApiKey'); + const sidebarModel = document.getElementById('byokModelId'); + + if (sidebarProvider && savedProvider) sidebarProvider.value = savedProvider; + if (sidebarKey && savedKey) sidebarKey.value = savedKey; + if (sidebarModel && savedModel) sidebarModel.value = savedModel; + + // If no preference is set, trigger the onboarding modal + if (!apiPref) { + const modal = document.getElementById('byokOnboardingModal'); + if (modal) modal.style.display = 'flex'; + } + }, + async fetchPipelineConfig() { try { const resp = await fetch('/api/automate/config'); @@ -190,7 +212,79 @@ export const AutomationUI = { // File Upload for Images document.getElementById('imageUploadInput')?.addEventListener('change', (e) => this.handleImageUpload(e)); - + + // BYOK Onboarding Modal Listeners + const onboardingModal = document.getElementById('byokOnboardingModal'); + const btnOnboardByok = document.getElementById('btnOnboardByok'); + const btnOnboardDefault = document.getElementById('btnOnboardDefault'); + const onboardOptions = document.getElementById('onboardOptions'); + const onboardByokForm = document.getElementById('onboardByokForm'); + const btnBackToOptions = document.getElementById('btnBackToOptions'); + const btnSaveOnboardKeys = document.getElementById('btnSaveOnboardKeys'); + + btnOnboardByok?.addEventListener('click', () => { + if (onboardOptions) onboardOptions.style.display = 'none'; + if (onboardByokForm) onboardByokForm.style.display = 'block'; + }); + + btnOnboardDefault?.addEventListener('click', () => { + localStorage.setItem('edmate-api-preference', 'default'); + if (onboardingModal) onboardingModal.style.display = 'none'; + const sidebarKey = document.getElementById('byokApiKey'); + if (sidebarKey) sidebarKey.value = ''; + }); + + btnBackToOptions?.addEventListener('click', () => { + if (onboardByokForm) onboardByokForm.style.display = 'none'; + if (onboardOptions) onboardOptions.style.display = 'flex'; + }); + + btnSaveOnboardKeys?.addEventListener('click', () => { + const provider = document.getElementById('onboardProviderSelect')?.value || 'gemini'; + const key = document.getElementById('onboardApiKey')?.value || ''; + + if (!key) { + alert('Please enter a valid API key.'); + return; + } + + localStorage.setItem('edmate-byok-provider', provider); + localStorage.setItem('edmate-byok-key', key); + localStorage.setItem('edmate-api-preference', 'byok'); + + const sidebarProvider = document.getElementById('byokProviderSelect'); + const sidebarKey = document.getElementById('byokApiKey'); + if (sidebarProvider) sidebarProvider.value = provider; + if (sidebarKey) sidebarKey.value = key; + + if (onboardingModal) onboardingModal.style.display = 'none'; + }); + + // Paywall Modal listeners + const paywallModal = document.getElementById('downloadPaywallModal'); + const btnPaywallByok = document.getElementById('btnPaywallByok'); + const btnClosePaywall = document.getElementById('btnClosePaywall'); + + btnPaywallByok?.addEventListener('click', () => { + if (paywallModal) paywallModal.style.display = 'none'; + if (onboardingModal) { + const savedKey = localStorage.getItem('edmate-byok-key') || ''; + const savedProvider = localStorage.getItem('edmate-byok-provider') || 'gemini'; + const onboardApiKey = document.getElementById('onboardApiKey'); + const onboardProviderSelect = document.getElementById('onboardProviderSelect'); + if (onboardApiKey) onboardApiKey.value = savedKey; + if (onboardProviderSelect) onboardProviderSelect.value = savedProvider; + + if (onboardOptions) onboardOptions.style.display = 'none'; + if (onboardByokForm) onboardByokForm.style.display = 'block'; + onboardingModal.style.display = 'flex'; + } + }); + + btnClosePaywall?.addEventListener('click', () => { + if (paywallModal) paywallModal.style.display = 'none'; + }); + // Analytics Nav document.getElementById('navAnalytics')?.addEventListener('click', (e) => { e.preventDefault(); diff --git a/qc_viewer/static/js/controllers/review.js b/qc_viewer/static/js/controllers/review.js index 555373d..d89747e 100644 --- a/qc_viewer/static/js/controllers/review.js +++ b/qc_viewer/static/js/controllers/review.js @@ -408,6 +408,21 @@ export const ReviewController = { this.showToast('Open a draft in Review to export', 'danger'); return; } + + // Paid solutions / BYOK Paywall Gate (only active on hosted web domains) + const isLocalhost = window.location.hostname === 'localhost' || window.location.hostname === '127.0.0.1'; + if (!isLocalhost) { + const byokKey = document.getElementById('byokApiKey')?.value; + if (!byokKey) { + const paywallModal = document.getElementById('downloadPaywallModal'); + if (paywallModal) { + paywallModal.style.display = 'flex'; + this.closeExportMenu(); + return; + } + } + } + try { await AutomationAPI.exportDraft(this.currentDraftData.id, format); this.showToast(`Downloaded ${String(format).toUpperCase()}`, 'success'); diff --git a/qc_viewer/static/viewer.html b/qc_viewer/static/viewer.html new file mode 100644 index 0000000..0523765 --- /dev/null +++ b/qc_viewer/static/viewer.html @@ -0,0 +1,172 @@ + + + + + +Edmate QC Content Viewer + + + + + + + + + + + + + + + + + +
    +
    + + +
    +
    +

    Dashboard

    +

    Quality control and verification workspace.

    +
    +
    + +
    + +
    +
    +
    + +
    + 🔬 +

    Ready to Verify

    +

    Select a paper from the sidebar to begin your QC review process.

    +
    + + +
    +
    + + + + From c8aa178048203cd8c7a03fea298f256089bcccb2 Mon Sep 17 00:00:00 2001 From: shmukit Date: Wed, 20 May 2026 13:51:36 +0600 Subject: [PATCH 6/8] fix(ui): update logos to beta, refine hero pipeline diagram, fix MathJax race condition --- qc_viewer/static/analytics.html | 2 +- qc_viewer/static/automate.html | 2 +- qc_viewer/static/how_it_works.html | 2 +- qc_viewer/static/index.html | 46 ++++++++++++++++++------- qc_viewer/static/js/controllers/docs.js | 4 +-- qc_viewer/static/viewer.html | 2 +- 6 files changed, 40 insertions(+), 18 deletions(-) diff --git a/qc_viewer/static/analytics.html b/qc_viewer/static/analytics.html index adeee08..53424cd 100644 --- a/qc_viewer/static/analytics.html +++ b/qc_viewer/static/analytics.html @@ -22,7 +22,7 @@ - Edmate AI + Edmate AI (Beta) + Automation Hub @@ -167,7 +173,11 @@

    Simple, Dynamic Pricing Plans

  • - Unlimited Document Exports (with BYOK) + Unlimited Exports (with BYOK) +
    • +
  • + + Optimal: up to 50-page PDFs (max 10 MB)
    • Star & Clone Repo @@ -188,15 +198,19 @@

    Simple, Dynamic Pricing Plans

  • - Unlock Web Export (CSV/JSON/MD) + Unlock Web Export (CSV / JSON / MD / Docx)
  • - Standard API Response Latency + Up to 30 uploads/month · max 10 MB each
  • - Save Up to 50 active drafts + Optimal: docs up to 50 pages +
    • +
  • + + 50 active drafts in workspace
    • Activate Basic Plan @@ -219,6 +233,10 @@

    Simple, Dynamic Pricing Plans

    Advanced exports: Word & Image ZIPs +
  • + + Unlimited uploads · max 25 MB / 200 pages +
  • High-concurrency parallel ingestion @@ -227,6 +245,10 @@

    Simple, Dynamic Pricing Plans

    Dedicated database schema slots
    • +
  • + + Priority support & SLA +
    • Get Started Pro @@ -234,6 +256,67 @@

    Simple, Dynamic Pricing Plans

    + +
    +
    +

    Frequently Asked Questions

    +

    Everything you need to know about using Edmate AI.

    +
    + +
    + +
    +

    Yes. For optimal performance, documents should be under 10 MB. The Vision AI engine must render each page as an image for layout-aware extraction — very large files increase processing time and API costs significantly. If you have a large textbook or report, we recommend splitting it into chapter-level PDFs for best results.

    +
    +
    + +
    + +
    +

    The pipeline is optimised for documents under 50 pages. Longer documents (e.g., 200-page textbooks) will work on the Pro plan, but processing time and token costs scale linearly with page count. We strongly recommend splitting large documents into focused chapters for the highest quality output and lowest cost.

    +
    +
    + +
    + +
    +

    You can still use the platform for free using our default hosted Gemini model. You can view and interact with the generated output in the browser. However, downloading the output in any format (CSV, JSON, Docx, etc.) requires a paid Basic or Pro plan. BYOK users always retain full download access for free.

    +
    +
    + +
    + +
    +

    Edmate AI's pedagogy engine applies evidence-backed techniques including Spaced Repetition, Cognitive Scaffolding, Concept Gap detection, and High-Integrity Assessment (HIA) generation — producing Q&As, Flashcards, Viva Defense prompts, and AI-resilient exercises that are structurally resistant to direct AI-assisted cheating.

    +
    +
    + +
    + +
    +

    Your API keys are never stored on our servers. They exist only in your current browser session and are transmitted directly to the LLM provider. Uploaded documents are processed in-memory and are not permanently retained. For full auditability, you can run the entire platform locally via the open-source repository.

    +
    +
    + +
    +
    +
    +
    @@ -278,5 +361,26 @@

    Resources

    + From 3fe51812878bd1a525433fd5380638e41f4df38d Mon Sep 17 00:00:00 2001 From: shmukit Date: Mon, 8 Jun 2026 16:25:34 +0600 Subject: [PATCH 8/8] feat: implement user authentication, storage backend interface, and export paywalls --- README.md | 12 +- content_gen/.env.example | 15 ++ qc_viewer/app_factory.py | 7 +- qc_viewer/middleware/auth.py | 205 ++++++++++++++++++ qc_viewer/middleware/quota.py | 169 +++++++++++++++ qc_viewer/routers/auth.py | 31 +++ qc_viewer/routers/automation.py | 80 ++++--- qc_viewer/services/cloud_storage.py | 174 +++++++++++++++ qc_viewer/services/draft_store.py | 115 ++++++++-- qc_viewer/static/automate.html | 79 ++++++- .../static/css/modules/marketing/theme.css | 1 + qc_viewer/static/js/auth.js | 194 +++++++++++++++++ qc_viewer/static/js/automate_api.js | 44 ++-- qc_viewer/static/js/controllers/drafts.js | 14 ++ qc_viewer/static/js/controllers/review.js | 24 +- 15 files changed, 1084 insertions(+), 80 deletions(-) create mode 100644 qc_viewer/middleware/auth.py create mode 100644 qc_viewer/middleware/quota.py create mode 100644 qc_viewer/routers/auth.py create mode 100644 qc_viewer/services/cloud_storage.py create mode 100644 qc_viewer/static/js/auth.js diff --git a/README.md b/README.md index 32b8cae..1a17211 100644 --- a/README.md +++ b/README.md @@ -100,7 +100,17 @@ workspace: > [!CAUTION] > **Database Schema Consistency**: Edmate's `database_service.py` currently expects tables to have specific columns (e.g., `title`, `options`, `correct_options`). If your database uses different column names, you must update the SQL queries in `content_gen/scripts/processing/database_service.py` to match your schema. -### 3. Understanding Pipeline Settings +### 3. Auth & Storage Configuration +Edmate supports both local self-hosted deployment and cloud SaaS operation. Configure these settings in `content_gen/.env`: + +- **Authentication (`EDMATE_AUTH_REQUIRED`)**: + - `false` (Default): Operates in "Self-Hosted Mode". All requests bypass auth checks and are treated as unrestricted "pro" users. Perfect for local development. + - `true`: Operates in "Cloud Mode". Requires a valid JWT token via `Authorization: Bearer `. Anonymous users have restricted access (e.g., cannot export without hitting a paywall). +- **Storage Backend (`EDMATE_STORAGE_BACKEND`)**: + - `local` (Default): Saves all drafts and extracted assets locally to the `.drafts` directory. + - `s3`: Placeholder for cloud deployments to store files in AWS S3 or compatible blob storage. + +### 4. Understanding Pipeline Settings The **Automation Hub** provides several "Admin" settings to handle diverse document formats: - **Extraction Guardrails**: Adjust "Detection Mode" to **Strict** for standard papers or **Open** for noisy documents. - **Model Routing**: Strategies to balance cost and quality. Edmate can use cheaper models (like Gemini Flash) for extraction and switch to high-precision models (like GPT-4o) for final content generation. diff --git a/content_gen/.env.example b/content_gen/.env.example index 6695a34..c85a157 100644 --- a/content_gen/.env.example +++ b/content_gen/.env.example @@ -31,3 +31,18 @@ DATABASE_URL=postgresql://user:password@host:5432/edmate # Path to Google Service Account JSON (optional, for Vertex AI/Google Cloud mode) # GOOGLE_APPLICATION_CREDENTIALS=credentials/gemini_creds.json + +# ===== Auth (Optional — omit for self-hosted mode) ===== +EDMATE_AUTH_REQUIRED=false # Set true for SaaS deployment +# AUTH_PROVIDER=postgres # "postgres" | "supabase" | "custom" +# AUTH_DATABASE_URL=... # If different from main DATABASE_URL + +# ===== Storage ===== +EDMATE_STORAGE_BACKEND=local # "local" | "s3" +# S3_ENDPOINT=... +# S3_BUCKET=edmate-uploads +# S3_ACCESS_KEY=... +# S3_SECRET_KEY=... + +# ===== Encryption (for server-side BYOK storage) ===== +# EDMATE_ENCRYPTION_SECRET=... diff --git a/qc_viewer/app_factory.py b/qc_viewer/app_factory.py index a0d0950..c470ace 100644 --- a/qc_viewer/app_factory.py +++ b/qc_viewer/app_factory.py @@ -42,7 +42,12 @@ def create_app() -> FastAPI: ) app.add_middleware(EdmateApiSecurityMiddleware) - + + from qc_viewer.middleware.auth import EdmateAuthMiddleware + from qc_viewer.middleware.quota import EdmateQuotaMiddleware + + app.add_middleware(EdmateQuotaMiddleware) + app.add_middleware(EdmateAuthMiddleware) app.include_router(api_v1_router) app.include_router(static_pages_router) app.include_router(questions_router) diff --git a/qc_viewer/middleware/auth.py b/qc_viewer/middleware/auth.py new file mode 100644 index 0000000..1262b90 --- /dev/null +++ b/qc_viewer/middleware/auth.py @@ -0,0 +1,205 @@ +""" +Provider-agnostic authentication middleware for qc_viewer. + +Environment: + EDMATE_AUTH_REQUIRED — set to 1/true to enforce token-based auth. + When false (default / self-hosted), every request is treated as a 'pro' + user with id 'anonymous'. When true and no valid token is presented the + request proceeds as 'anonymous' on the 'anonymous' (limited) plan. + +The actual identity backend is injected via :func:`set_auth_provider`. +""" + +from __future__ import annotations + +import logging +import os +import threading +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Optional + +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import Response + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Public data structures +# --------------------------------------------------------------------------- + + +@dataclass +class UserInfo: + """Minimal identity payload returned by an :class:`AuthProvider`.""" + + user_id: str + email: Optional[str] = None + display_name: Optional[str] = None + + +# --------------------------------------------------------------------------- +# Auth provider interface + dummy implementation +# --------------------------------------------------------------------------- + + +class AuthProvider(ABC): + """Abstract authentication / user-lookup backend.""" + + @abstractmethod + async def verify_token(self, token: str) -> Optional[UserInfo]: + """Return a :class:`UserInfo` if *token* is valid, else ``None``.""" + ... + + @abstractmethod + async def get_user_plan(self, user_id: str) -> str: + """Return the billing plan name (e.g. ``'free'``, ``'pro'``) for *user_id*.""" + ... + + + @abstractmethod + def save_user_byok(self, user_id: str, provider: str, api_key: str, model_id: Optional[str] = None) -> bool: + """Save a BYOK configuration for a user. Returns True if successful.""" + ... + + +class DummyAuthProvider(AuthProvider): + """Placeholder that rejects every token. + + Wire in a real provider (Firebase, Supabase, …) before enabling + ``EDMATE_AUTH_REQUIRED``. + """ + + async def verify_token(self, token: str) -> Optional[UserInfo]: + return None + + async def get_user_plan(self, user_id: str) -> str: + return "anonymous" + + def save_user_byok(self, user_id: str, provider: str, api_key: str, model_id: Optional[str] = None) -> bool: + return True + + +# --------------------------------------------------------------------------- +# Module-level singleton (same pattern as job_repository.py) +# --------------------------------------------------------------------------- + +_auth_provider: Optional[AuthProvider] = None +_provider_lock = threading.Lock() + + +def get_auth_provider() -> AuthProvider: + """Singleton auth provider (swap via :func:`set_auth_provider` in tests).""" + global _auth_provider + with _provider_lock: + if _auth_provider is None: + _auth_provider = DummyAuthProvider() + return _auth_provider + + +def set_auth_provider(provider: Optional[AuthProvider]) -> None: + """Test hook: pass ``None`` to reset to :class:`DummyAuthProvider` on next get.""" + global _auth_provider + with _provider_lock: + _auth_provider = provider + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_PUBLIC_EXACT: frozenset[str] = frozenset(( + "/", + "/index.html", + "/how_it_works.html", + "/api/health", +)) + +_PUBLIC_PREFIXES: tuple[str, ...] = ( + "/css/", + "/js/", + "/favicon", + "/docs/", +) + + +def _truthy(val: str | None) -> bool: + return (val or "").strip().lower() in ("1", "true", "yes", "on") + + +def _is_public(path: str) -> bool: + """Return ``True`` if *path* should bypass authentication entirely.""" + if path in _PUBLIC_EXACT: + return True + return any(path.startswith(prefix) for prefix in _PUBLIC_PREFIXES) + + +def _extract_bearer_token(request: Request) -> Optional[str]: + """Pull the token from ``Authorization: Bearer ``, or ``None``.""" + auth = request.headers.get("authorization") or "" + if auth.lower().startswith("bearer "): + token = auth[7:].strip() + return token if token else None + return None + + +# --------------------------------------------------------------------------- +# Middleware +# --------------------------------------------------------------------------- + + +class EdmateAuthMiddleware(BaseHTTPMiddleware): + """Populate ``request.state.user_id`` and ``request.state.plan``. + + Three operating modes: + + 1. **Self-hosted** (``EDMATE_AUTH_REQUIRED=false``, the default): + Every request is treated as ``user_id='anonymous'``, + ``plan='pro'`` — full, unrestricted access. + + 2. **Cloud, unauthenticated** (``EDMATE_AUTH_REQUIRED=true``, no token): + ``user_id='anonymous'``, ``plan='anonymous'`` — limited access. + + 3. **Cloud, authenticated** (``EDMATE_AUTH_REQUIRED=true``, valid token): + ``user_id`` and ``plan`` are resolved via the configured + :class:`AuthProvider`. + """ + + async def dispatch(self, request: Request, call_next) -> Response: + # Let CORS preflight and public assets through untouched. + if request.method == "OPTIONS": + return await call_next(request) + + path = request.url.path + if _is_public(path): + return await call_next(request) + + # --- Mode 1: self-hosted (auth not required) --- + if not _truthy(os.environ.get("EDMATE_AUTH_REQUIRED")): + request.state.user_id = "anonymous" + request.state.plan = "pro" + return await call_next(request) + + # --- Auth is required (cloud mode) --- + token = _extract_bearer_token(request) + + if token is None: + # Mode 2: no credentials supplied → limited anonymous access. + request.state.user_id = "anonymous" + request.state.plan = "anonymous" + return await call_next(request) + + # Mode 3: token present → verify with the configured provider. + provider = get_auth_provider() + user = await provider.verify_token(token) + + if user is None: + # Invalid / expired token → treat as anonymous. + request.state.user_id = "anonymous" + request.state.plan = "anonymous" + return await call_next(request) + + request.state.user_id = user.user_id + request.state.plan = await provider.get_user_plan(user.user_id) + return await call_next(request) diff --git a/qc_viewer/middleware/quota.py b/qc_viewer/middleware/quota.py new file mode 100644 index 0000000..2076b1f --- /dev/null +++ b/qc_viewer/middleware/quota.py @@ -0,0 +1,169 @@ +""" +Plan-based quota enforcement for upload and export routes. + +Reads ``request.state.plan`` (set by the auth middleware) and enforces +per-plan limits on uploads and exports. + +Environment: + EDMATE_AUTH_REQUIRED — when falsy (or unset), all quota checks are skipped + so that self-hosted deployments are unrestricted. + +v1 note: upload counters are kept in-process memory and reset each calendar +month. A DB-backed implementation will replace this later. +""" + +from __future__ import annotations + +import os +import threading +import time +from typing import Dict, Tuple + +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import JSONResponse + +# ── plan definitions ──────────────────────────────────────────────────────── + +PLAN_LIMITS: Dict[str, Dict] = { + "anonymous": { + "max_uploads_per_month": 999, + "max_file_size_mb": 50, + "can_export": False, + "history_ttl_days": 1, + }, + "free": { + "max_uploads_per_month": 999, + "max_file_size_mb": 50, + "can_export": False, + "history_ttl_days": 1, + }, + "basic": { + "max_uploads_per_month": 30, + "max_file_size_mb": 10, + "can_export": True, + "history_ttl_days": 30, + }, + "pro": { + "max_uploads_per_month": 999999, + "max_file_size_mb": 25, + "can_export": True, + "history_ttl_days": -1, + }, +} + +# ── in-memory upload counters (v1) ────────────────────────────────────────── + +_counter_lock = threading.Lock() +# key: (user_id, "YYYY-MM") → count +_upload_counts: Dict[Tuple[str, str], int] = {} + + +def _truthy(val: str | None) -> bool: + return (val or "").strip().lower() in ("1", "true", "yes", "on") + + +def _current_month_key() -> str: + """Return the current year-month string, e.g. ``'2026-06'``.""" + return time.strftime("%Y-%m", time.gmtime()) + + +def _get_and_increment_uploads(user_id: str) -> int: + """Atomically read the current month's upload count, increment, and return + the *previous* value (i.e. the count before this request).""" + month = _current_month_key() + key = (user_id, month) + with _counter_lock: + current = _upload_counts.get(key, 0) + _upload_counts[key] = current + 1 + return current + + +def _rollback_upload(user_id: str) -> None: + """Undo the increment when the request is going to be rejected.""" + month = _current_month_key() + key = (user_id, month) + with _counter_lock: + val = _upload_counts.get(key, 0) + if val > 0: + _upload_counts[key] = val - 1 + + +def _quota_error(detail: str) -> JSONResponse: + return JSONResponse( + status_code=402, + content={"detail": detail, "upgrade_required": True}, + ) + + +# ── middleware ────────────────────────────────────────────────────────────── + + +class EdmateQuotaMiddleware(BaseHTTPMiddleware): + """Enforce plan-based quotas on upload and export routes.""" + + async def dispatch(self, request: Request, call_next): + # Skip everything for preflight requests. + if request.method == "OPTIONS": + return await call_next(request) + + # Self-hosted mode: no quotas. + if not _truthy(os.environ.get("EDMATE_AUTH_REQUIRED")): + return await call_next(request) + + path = request.url.path + + # ── upload route ──────────────────────────────────────────────── + if request.method == "POST" and path == "/api/automate/draft": + return await self._check_upload(request, call_next) + + # ── export route ──────────────────────────────────────────────── + if request.method == "GET" and path.startswith("/api/automate/draft/") and path.endswith("/export"): + return await self._check_export(request, call_next) + + return await call_next(request) + + # ── private helpers ───────────────────────────────────────────────── + + async def _check_upload(self, request: Request, call_next): + plan_name = getattr(request.state, "plan", "anonymous") + limits = PLAN_LIMITS.get(plan_name, PLAN_LIMITS["anonymous"]) + user_id = getattr(request.state, "user_id", None) or "anon" + + # File size check via content-length header. + content_length = request.headers.get("content-length") + if content_length is not None: + try: + size_mb = int(content_length) / (1024 * 1024) + except (ValueError, TypeError): + size_mb = 0.0 + max_mb = limits["max_file_size_mb"] + if size_mb > max_mb: + return _quota_error( + f"File size exceeds the {max_mb} MB limit for the " + f"'{plan_name}' plan. Please upgrade your plan.", + ) + + # Monthly upload count. + used = _get_and_increment_uploads(user_id) + max_uploads = limits["max_uploads_per_month"] + if used >= max_uploads: + _rollback_upload(user_id) + return _quota_error( + f"Monthly upload limit ({max_uploads}) reached for the " + f"'{plan_name}' plan. Please upgrade your plan.", + ) + + return await call_next(request) + + async def _check_export(self, request: Request, call_next): + plan_name = getattr(request.state, "plan", "anonymous") + limits = PLAN_LIMITS.get(plan_name, PLAN_LIMITS["anonymous"]) + + if not limits["can_export"]: + return _quota_error( + f"Export is not available on the '{plan_name}' plan. " + f"Please upgrade to a plan that supports exports.", + ) + + return await call_next(request) diff --git a/qc_viewer/routers/auth.py b/qc_viewer/routers/auth.py new file mode 100644 index 0000000..67bc587 --- /dev/null +++ b/qc_viewer/routers/auth.py @@ -0,0 +1,31 @@ +from fastapi import APIRouter, Depends, HTTPException, Request +from pydantic import BaseModel +from typing import Optional + +from qc_viewer.middleware.auth import get_auth_provider + +router = APIRouter() + +class BYOKSaveRequest(BaseModel): + provider: str + api_key: str + model_id: Optional[str] = None + +@router.post("/api/auth/byok") +async def save_byok(request: Request, data: BYOKSaveRequest): + """ + Save BYOK key on server (encrypted) for authenticated users. + Anonymous users cannot use this endpoint. + """ + user_id = getattr(request.state, "user_id", None) + if not user_id or user_id == "anonymous": + raise HTTPException(status_code=401, detail="Must be logged in to save BYOK keys") + + auth_provider = get_auth_provider() + # In a real implementation, you'd encrypt the key and store it in Postgres. + # We simulate it for now. + success = auth_provider.save_user_byok(user_id, data.provider, data.api_key, data.model_id) + if not success: + raise HTTPException(status_code=500, detail="Failed to save BYOK keys") + + return {"status": "success", "message": "BYOK keys saved securely"} diff --git a/qc_viewer/routers/automation.py b/qc_viewer/routers/automation.py index 6df98a2..68ba1a1 100644 --- a/qc_viewer/routers/automation.py +++ b/qc_viewer/routers/automation.py @@ -18,10 +18,13 @@ from pydantic import BaseModel from qc_viewer.services import draft_export from qc_viewer.services.draft_store import ( + ANONYMOUS_USER, DraftNotFound, delete_draft_data, ensure_drafts_root, get_draft_dir, + is_draft_owned_by, + list_draft_metadata, load_metadata_if_exists, read_json, resolve_metadata_path, @@ -30,13 +33,24 @@ ) -def _require_draft_path(draft_id: str) -> Path: +def _get_user_id(request: Request) -> str: + """Extract user_id from request state (set by auth middleware).""" + return getattr(request.state, "user_id", ANONYMOUS_USER) + + +def _require_draft_path(draft_id: str, user_id: Optional[str] = None) -> Path: try: - return resolve_metadata_path(draft_id) + return resolve_metadata_path(draft_id, user_id) except DraftNotFound: raise HTTPException(status_code=404, detail="Draft not found") +def _require_ownership(draft_id: str, user_id: str) -> None: + """Raise 404 if the user does not own this draft.""" + if not is_draft_owned_by(draft_id, user_id): + raise HTTPException(status_code=404, detail="Draft not found") + + class PublishRequest(BaseModel): draft_id: str table_name: str @@ -53,6 +67,7 @@ class RefineRequest(BaseModel): @router.post("/api/automate/draft") async def receive_draft( + request: Request, background_tasks: BackgroundTasks, subject: str = Form(...), paper_code: str = Form(...), @@ -69,13 +84,14 @@ async def receive_draft( ): from qc_viewer.config import get_workspace_defaults + user_id = _get_user_id(request) default_curriculum, _default_subject = get_workspace_defaults() curriculum_resolved = (curriculum or "").strip() or default_curriculum draft_id = f"draft_{uuid.uuid4().hex[:8]}" ensure_drafts_root() - static_dir = get_draft_dir(draft_id) + static_dir = get_draft_dir(draft_id, user_id) static_dir.mkdir(parents=True, exist_ok=True) file_path = static_dir / "source.pdf" @@ -86,6 +102,7 @@ async def receive_draft( static_dir / "metadata.json", { "id": draft_id, + "owner_id": user_id, "pipeline_job_id": draft_id, "prompt_version": CONTENT_GENERATION_PROMPT_VERSION, "subject": subject, @@ -129,39 +146,27 @@ async def receive_draft( @router.get("/api/automate/drafts") -async def list_drafts(): - drafts = [] - if not DRAFTS_ROOT.exists(): - return [] - - for d in DRAFTS_ROOT.iterdir(): - if d.is_dir() and (d / "metadata.json").exists(): - try: - drafts.append(read_json(d / "metadata.json")) - except Exception: - continue - elif d.suffix == ".json" and d.name != "metadata.json": - try: - data = read_json(d) - if "id" in data: - drafts.append(data) - except Exception: - continue - +async def list_drafts(request: Request): + user_id = _get_user_id(request) + drafts = list_draft_metadata(user_id) return sorted(drafts, key=sort_key_from_timestamp, reverse=True) @router.get("/api/automate/draft/{draft_id}") -async def get_draft_results(draft_id: str): - return read_json(_require_draft_path(draft_id)) +async def get_draft_results(request: Request, draft_id: str): + user_id = _get_user_id(request) + _require_ownership(draft_id, user_id) + return read_json(_require_draft_path(draft_id, user_id)) @router.get("/api/automate/draft/{draft_id}/export") -async def export_draft(draft_id: str, format: str = "json"): +async def export_draft(request: Request, draft_id: str, format: str = "json"): + user_id = _get_user_id(request) + _require_ownership(draft_id, user_id) fmt = format.lower() if fmt not in draft_export.SUPPORTED_FORMATS: raise HTTPException(status_code=400, detail=f"Unsupported format: {format}") - meta = read_json(_require_draft_path(draft_id)) + meta = read_json(_require_draft_path(draft_id, user_id)) body = draft_export.render(meta, fmt) filename = draft_export.safe_filename(meta, fmt) return Response( @@ -173,8 +178,9 @@ async def export_draft(draft_id: str, format: str = "json"): @router.get("/api/automate/draft/{draft_id}/stream") async def stream_draft_progress(draft_id: str, request: Request): + user_id = _get_user_id(request) try: - meta_path = resolve_metadata_path(draft_id) + meta_path = resolve_metadata_path(draft_id, user_id) except DraftNotFound: meta_path = DRAFTS_ROOT / draft_id / "metadata.json" if not meta_path.exists(): @@ -213,24 +219,28 @@ async def event_generator(): @router.post("/api/automate/draft/{draft_id}/stop") -async def stop_draft_processing(draft_id: str): +async def stop_draft_processing(request: Request, draft_id: str): + user_id = _get_user_id(request) + _require_ownership(draft_id, user_id) if draft_id in CANCELLATION_EVENTS: CANCELLATION_EVENTS[draft_id].set() return {"status": "stopping"} - meta = load_metadata_if_exists(draft_id) + meta = load_metadata_if_exists(draft_id, user_id) if meta and meta.get("status") == "PROCESSING": meta["status"] = "FAILED" meta["status_message"] = "Stopped by user" - write_json(_require_draft_path(draft_id), meta) + write_json(_require_draft_path(draft_id, user_id), meta) return {"status": "stopped"} return {"status": "not_running"} @router.patch("/api/automate/draft/{draft_id}") -async def update_draft(draft_id: str, updates: dict): - meta_path = _require_draft_path(draft_id) +async def update_draft(request: Request, draft_id: str, updates: dict): + user_id = _get_user_id(request) + _require_ownership(draft_id, user_id) + meta_path = _require_draft_path(draft_id, user_id) try: data = read_json(meta_path) data.update(updates) @@ -242,8 +252,10 @@ async def update_draft(draft_id: str, updates: dict): @router.delete("/api/automate/draft/{draft_id}") -async def delete_draft(draft_id: str): - if delete_draft_data(draft_id): +async def delete_draft(request: Request, draft_id: str): + user_id = _get_user_id(request) + _require_ownership(draft_id, user_id) + if delete_draft_data(draft_id, user_id): return {"status": "deleted"} raise HTTPException(status_code=404, detail="Draft not found") diff --git a/qc_viewer/services/cloud_storage.py b/qc_viewer/services/cloud_storage.py new file mode 100644 index 0000000..17b1893 --- /dev/null +++ b/qc_viewer/services/cloud_storage.py @@ -0,0 +1,174 @@ +""" +Pluggable file-storage adapter for draft assets (PDFs, images, etc.). + +Default: local filesystem. Set EDMATE_STORAGE_BACKEND=s3 to switch to the +(not-yet-implemented) S3 backend. +""" + +from __future__ import annotations + +import asyncio +import os +import shutil +import threading +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Optional + +from qc_viewer.config import DRAFTS_ROOT + + +# --------------------------------------------------------------------------- +# Abstract interface +# --------------------------------------------------------------------------- + +class StorageBackend(ABC): + """Async file-storage interface for per-draft assets.""" + + @abstractmethod + async def upload(self, user_id: str, draft_id: str, data: bytes, filename: str) -> str: + """Persist *data* and return a reference key (e.g. path or URL).""" + ... + + @abstractmethod + async def download(self, user_id: str, draft_id: str, filename: str) -> bytes: + """Return the raw bytes for the given file.""" + ... + + @abstractmethod + async def delete_draft(self, user_id: str, draft_id: str) -> None: + """Remove **all** stored files for a draft.""" + ... + + @abstractmethod + async def list_files(self, user_id: str, draft_id: str) -> list[str]: + """Return the filenames stored under the given draft.""" + ... + + +# --------------------------------------------------------------------------- +# Local-filesystem implementation +# --------------------------------------------------------------------------- + +class LocalStorageBackend(StorageBackend): + """Store draft assets on the local filesystem under ``DRAFTS_ROOT``. + + Directory layout:: + + DRAFTS_ROOT / / / + + For backward compatibility, ``user_id == "anonymous"`` omits the user + directory:: + + DRAFTS_ROOT / / + + All blocking I/O is offloaded to the default executor via + ``asyncio.to_thread`` so the async interface never blocks the event loop. + """ + + def _draft_dir(self, user_id: str, draft_id: str) -> Path: + """Return the directory that holds a draft's files.""" + if user_id == "anonymous": + return DRAFTS_ROOT / draft_id + return DRAFTS_ROOT / user_id / draft_id + + # -- synchronous helpers (run inside threads) --------------------------- + + def _upload_sync(self, user_id: str, draft_id: str, data: bytes, filename: str) -> str: + dest_dir = self._draft_dir(user_id, draft_id) + dest_dir.mkdir(parents=True, exist_ok=True) + dest = dest_dir / filename + dest.write_bytes(data) + return str(dest) + + def _download_sync(self, user_id: str, draft_id: str, filename: str) -> bytes: + path = self._draft_dir(user_id, draft_id) / filename + if not path.exists(): + raise FileNotFoundError(f"No such file: {path}") + return path.read_bytes() + + def _delete_draft_sync(self, user_id: str, draft_id: str) -> None: + draft_dir = self._draft_dir(user_id, draft_id) + if draft_dir.is_dir(): + shutil.rmtree(draft_dir) + + def _list_files_sync(self, user_id: str, draft_id: str) -> list[str]: + draft_dir = self._draft_dir(user_id, draft_id) + if not draft_dir.is_dir(): + return [] + return [p.name for p in draft_dir.iterdir() if p.is_file()] + + # -- async interface ---------------------------------------------------- + + async def upload(self, user_id: str, draft_id: str, data: bytes, filename: str) -> str: + return await asyncio.to_thread(self._upload_sync, user_id, draft_id, data, filename) + + async def download(self, user_id: str, draft_id: str, filename: str) -> bytes: + return await asyncio.to_thread(self._download_sync, user_id, draft_id, filename) + + async def delete_draft(self, user_id: str, draft_id: str) -> None: + await asyncio.to_thread(self._delete_draft_sync, user_id, draft_id) + + async def list_files(self, user_id: str, draft_id: str) -> list[str]: + return await asyncio.to_thread(self._list_files_sync, user_id, draft_id) + + +# --------------------------------------------------------------------------- +# S3 stub (future implementation) +# --------------------------------------------------------------------------- + +class S3StorageBackend(StorageBackend): + """Placeholder for an AWS-S3-backed storage backend.""" + + async def upload(self, user_id: str, draft_id: str, data: bytes, filename: str) -> str: + raise NotImplementedError("S3 storage not yet configured") + + async def download(self, user_id: str, draft_id: str, filename: str) -> bytes: + raise NotImplementedError("S3 storage not yet configured") + + async def delete_draft(self, user_id: str, draft_id: str) -> None: + raise NotImplementedError("S3 storage not yet configured") + + async def list_files(self, user_id: str, draft_id: str) -> list[str]: + raise NotImplementedError("S3 storage not yet configured") + + +# --------------------------------------------------------------------------- +# Singleton factory +# --------------------------------------------------------------------------- + +_BACKENDS = { + "local": LocalStorageBackend, + "s3": S3StorageBackend, +} + +_default_backend: Optional[StorageBackend] = None +_backend_lock = threading.Lock() + + +def _default_backend_from_env() -> StorageBackend: + """Instantiate the backend selected by the ``EDMATE_STORAGE_BACKEND`` env var.""" + name = (os.environ.get("EDMATE_STORAGE_BACKEND") or "local").strip().lower() + cls = _BACKENDS.get(name) + if cls is None: + raise ValueError( + f"Unknown storage backend {name!r}. " + f"Supported: {', '.join(sorted(_BACKENDS))}" + ) + return cls() + + +def get_storage_backend() -> StorageBackend: + """Singleton storage backend (swap via ``set_storage_backend`` in tests).""" + global _default_backend + with _backend_lock: + if _default_backend is None: + _default_backend = _default_backend_from_env() + return _default_backend + + +def set_storage_backend(backend: Optional[StorageBackend]) -> None: + """Test hook: pass ``None`` to reset to env-based default on next get.""" + global _default_backend + with _backend_lock: + _default_backend = backend diff --git a/qc_viewer/services/draft_store.py b/qc_viewer/services/draft_store.py index d3756d8..244933f 100644 --- a/qc_viewer/services/draft_store.py +++ b/qc_viewer/services/draft_store.py @@ -10,6 +10,9 @@ METADATA_LOCK = threading.Lock() +# Sentinel for anonymous / self-hosted users (no user-scoping) +ANONYMOUS_USER = "anonymous" + class DraftNotFound(Exception): """No metadata file exists for this draft_id (neither dir/metadata.json nor legacy flat .json).""" @@ -23,25 +26,57 @@ def ensure_drafts_root() -> None: DRAFTS_ROOT.mkdir(parents=True, exist_ok=True) -def get_draft_dir(draft_id: str) -> Path: - return DRAFTS_ROOT / draft_id +def _user_drafts_root(user_id: Optional[str] = None) -> Path: + """ + Return the base directory for a user's drafts. + + Anonymous / None users use the flat DRAFTS_ROOT (backward compatible). + Authenticated users get a scoped sub-directory: DRAFTS_ROOT / user_id. + """ + if not user_id or user_id == ANONYMOUS_USER: + return DRAFTS_ROOT + scoped = DRAFTS_ROOT / user_id + scoped.mkdir(parents=True, exist_ok=True) + return scoped + + +def get_draft_dir(draft_id: str, user_id: Optional[str] = None) -> Path: + return _user_drafts_root(user_id) / draft_id -def get_draft_metadata_path(draft_id: str) -> Path: - return get_draft_dir(draft_id) / "metadata.json" +def get_draft_metadata_path(draft_id: str, user_id: Optional[str] = None) -> Path: + return get_draft_dir(draft_id, user_id) / "metadata.json" def get_legacy_metadata_path(draft_id: str) -> Path: return DRAFTS_ROOT / f"{draft_id}.json" -def resolve_metadata_path(draft_id: str) -> Path: - meta_path = get_draft_metadata_path(draft_id) - if meta_path.exists(): - return meta_path +def resolve_metadata_path(draft_id: str, user_id: Optional[str] = None) -> Path: + """ + Resolve the metadata.json path for a draft. + + Search order: + 1. User-scoped directory: DRAFTS_ROOT / user_id / draft_id / metadata.json + 2. Flat directory (anonymous/legacy): DRAFTS_ROOT / draft_id / metadata.json + 3. Legacy flat JSON: DRAFTS_ROOT / draft_id.json + """ + # Try user-scoped path first (if a real user_id is provided) + if user_id and user_id != ANONYMOUS_USER: + scoped_path = _user_drafts_root(user_id) / draft_id / "metadata.json" + if scoped_path.exists(): + return scoped_path + + # Fall back to flat directory (anonymous / legacy) + flat_path = DRAFTS_ROOT / draft_id / "metadata.json" + if flat_path.exists(): + return flat_path + + # Fall back to legacy flat JSON legacy_path = get_legacy_metadata_path(draft_id) if legacy_path.exists(): return legacy_path + raise DraftNotFound(draft_id) @@ -75,12 +110,13 @@ def read_modify_write_json(path: Path, mutator: Callable[[dict[str, Any]], None] os.replace(tmp, path) -def list_draft_metadata() -> list[dict[str, Any]]: +def _scan_directory_for_drafts(root: Path) -> list[dict[str, Any]]: + """Scan a single directory for draft metadata (directories with metadata.json + legacy flat .json).""" drafts: list[dict[str, Any]] = [] - if not DRAFTS_ROOT.exists(): + if not root.exists(): return drafts - for d in DRAFTS_ROOT.iterdir(): + for d in root.iterdir(): if d.is_dir() and (d / "metadata.json").exists(): try: drafts.append(read_json(d / "metadata.json")) @@ -96,6 +132,17 @@ def list_draft_metadata() -> list[dict[str, Any]]: return drafts +def list_draft_metadata(user_id: Optional[str] = None) -> list[dict[str, Any]]: + """ + List all draft metadata for a given user. + + - Anonymous / None: scans DRAFTS_ROOT (flat layout, backward compatible). + - Authenticated user: scans DRAFTS_ROOT / user_id only. + """ + root = _user_drafts_root(user_id) + return _scan_directory_for_drafts(root) + + def sort_key_from_timestamp(payload: dict[str, Any]) -> float: ts = payload.get("timestamp") or payload.get("created_at") or "" if not ts: @@ -109,12 +156,20 @@ def sort_key_from_timestamp(payload: dict[str, Any]) -> float: return 0.0 -def delete_draft_data(draft_id: str) -> bool: - draft_dir = get_draft_dir(draft_id) +def delete_draft_data(draft_id: str, user_id: Optional[str] = None) -> bool: + """Delete all data for a draft. Checks user-scoped dir first, then flat/legacy layout.""" + draft_dir = get_draft_dir(draft_id, user_id) if draft_dir.is_dir(): shutil.rmtree(draft_dir) return True + # Also check the flat layout (anonymous / legacy) + if user_id and user_id != ANONYMOUS_USER: + flat_dir = DRAFTS_ROOT / draft_id + if flat_dir.is_dir(): + shutil.rmtree(flat_dir) + return True + legacy_json = get_legacy_metadata_path(draft_id) legacy_pdf = DRAFTS_ROOT / f"{draft_id}.pdf" if legacy_json.exists(): @@ -125,8 +180,38 @@ def delete_draft_data(draft_id: str) -> bool: return False -def load_metadata_if_exists(draft_id: str) -> Optional[dict[str, Any]]: +def load_metadata_if_exists(draft_id: str, user_id: Optional[str] = None) -> Optional[dict[str, Any]]: try: - return read_json(resolve_metadata_path(draft_id)) + return read_json(resolve_metadata_path(draft_id, user_id)) except DraftNotFound: return None + + +def is_draft_owned_by(draft_id: str, user_id: str) -> bool: + """ + Check if a draft belongs to a user. + + Returns True if: + - Draft exists in the user's scoped directory, OR + - Draft metadata has matching owner_id, OR + - user_id is anonymous (anonymous can access flat-layout drafts) + """ + if not user_id or user_id == ANONYMOUS_USER: + # Anonymous users can access any draft in the flat layout + return True + + # Check user-scoped directory + scoped_path = _user_drafts_root(user_id) / draft_id / "metadata.json" + if scoped_path.exists(): + return True + + # Check owner_id in flat-layout metadata + flat_path = DRAFTS_ROOT / draft_id / "metadata.json" + if flat_path.exists(): + try: + meta = read_json(flat_path) + return meta.get("owner_id") == user_id + except Exception: + return False + + return False diff --git a/qc_viewer/static/automate.html b/qc_viewer/static/automate.html index fab6744..5c97f00 100644 --- a/qc_viewer/static/automate.html +++ b/qc_viewer/static/automate.html @@ -40,6 +40,46 @@
    + + + + + +
    @@ -135,13 +175,45 @@

    Live Preview

    Pricing - + GitHub Repo + +
    + + + +
    @@ -551,11 +623,16 @@

    + import { AuthUI } from '/js/auth.js'; import { AutomationUI } from '/js/automate_ui.js'; import { LayoutManager } from '/js/ui.js'; import { ThemeManager } from '/js/theme.js'; + // Expose to window for inline onclick handlers (e.g. in the warning banner) + window.AuthUI = AuthUI; + ThemeManager.init(); + AuthUI.init(); AutomationUI.init(); LayoutManager.init(); diff --git a/qc_viewer/static/css/modules/marketing/theme.css b/qc_viewer/static/css/modules/marketing/theme.css index 89357f6..7d89266 100644 --- a/qc_viewer/static/css/modules/marketing/theme.css +++ b/qc_viewer/static/css/modules/marketing/theme.css @@ -19,6 +19,7 @@ color: var(--primary); } [data-theme="light"] .navbar-brand { + color: var(--text-main); /* Rely on var(--text-main) */ } [data-theme="light"] .navbar-brand svg { diff --git a/qc_viewer/static/js/auth.js b/qc_viewer/static/js/auth.js new file mode 100644 index 0000000..c12e650 --- /dev/null +++ b/qc_viewer/static/js/auth.js @@ -0,0 +1,194 @@ +/** + * Edmate Authentication & User Profile UI + * Handles login modal, session persistence (JWT in memory), and UI state updates. + */ + +export const AuthUI = { + token: null, + userProfile: null, + + init() { + this.checkSession(); + this.setupListeners(); + }, + + setupListeners() { + const btnLogin = document.getElementById('btnLoginNav'); + const btnLogout = document.getElementById('btnLogout'); + const modal = document.getElementById('authModal'); + const closeBtn = document.getElementById('closeAuthModal'); + const authForm = document.getElementById('authForm'); + + if (btnLogin) btnLogin.addEventListener('click', () => this.openModal()); + if (closeBtn) closeBtn.addEventListener('click', () => this.closeModal()); + if (btnLogout) btnLogout.addEventListener('click', () => this.logout()); + + if (authForm) { + authForm.addEventListener('submit', (e) => { + e.preventDefault(); + this.handleAuthSubmit(); + }); + } + + // Tab switching + const tabs = document.querySelectorAll('.auth-tab'); + tabs.forEach(tab => { + tab.addEventListener('click', (e) => { + tabs.forEach(t => t.classList.remove('active')); + e.target.classList.add('active'); + + const mode = e.target.dataset.mode; // 'login' or 'signup' + const btnSubmit = document.getElementById('btnAuthSubmit'); + if (btnSubmit) { + btnSubmit.textContent = mode === 'login' ? 'Sign In' : 'Create Account'; + } + }); + }); + }, + + openModal() { + const modal = document.getElementById('authModal'); + if (modal) modal.style.display = 'flex'; + }, + + closeModal() { + const modal = document.getElementById('authModal'); + if (modal) modal.style.display = 'none'; + + const errorEl = document.getElementById('authError'); + if (errorEl) errorEl.style.display = 'none'; + }, + + async handleAuthSubmit() { + const email = document.getElementById('authEmail')?.value; + const password = document.getElementById('authPassword')?.value; + const errorEl = document.getElementById('authError'); + const btnSubmit = document.getElementById('btnAuthSubmit'); + + if (!email || !password) return; + + // This is a placeholder for the actual auth provider call. + // For now, we simulate a successful login to unblock the UI. + + const activeTab = document.querySelector('.auth-tab.active'); + const mode = activeTab ? activeTab.dataset.mode : 'login'; + + try { + btnSubmit.disabled = true; + btnSubmit.textContent = 'Please wait...'; + errorEl.style.display = 'none'; + + // SIMULATED AUTH DELAY + await new Promise(resolve => setTimeout(resolve, 800)); + + // Simulate receiving a token and profile + this.token = "simulated_jwt_token_for_" + email; + this.userProfile = { + user_id: "user_" + Date.now().toString(36), + email: email, + display_name: email.split('@')[0], + plan: "free" // Default to free plan + }; + + // In a real implementation, you would store a refresh token in httpOnly cookie + // and the short-lived access token in memory. + // For this demo, we use sessionStorage + sessionStorage.setItem('edmate_token', this.token); + sessionStorage.setItem('edmate_profile', JSON.stringify(this.userProfile)); + + this.updateUIState(); + this.closeModal(); + + if (window.AutomationUI && window.AutomationUI.showToast) { + window.AutomationUI.showToast(`Successfully signed in as ${this.userProfile.display_name}`); + } + + // Refresh drafts to show user-scoped drafts + if (window.AutomationUI) { + window.AutomationUI.fetchDrafts(); + } + + } catch (error) { + errorEl.textContent = error.message || 'Authentication failed. Please try again.'; + errorEl.style.display = 'block'; + } finally { + btnSubmit.disabled = false; + btnSubmit.textContent = mode === 'login' ? 'Sign In' : 'Create Account'; + } + }, + + logout() { + this.token = null; + this.userProfile = null; + sessionStorage.removeItem('edmate_token'); + sessionStorage.removeItem('edmate_profile'); + + this.updateUIState(); + + if (window.AutomationUI && window.AutomationUI.showToast) { + window.AutomationUI.showToast('You have been signed out.'); + } + + // Refresh drafts to show anonymous drafts + if (window.AutomationUI) { + window.AutomationUI.fetchDrafts(); + } + }, + + checkSession() { + const token = sessionStorage.getItem('edmate_token'); + const profileStr = sessionStorage.getItem('edmate_profile'); + + if (token && profileStr) { + try { + this.token = token; + this.userProfile = JSON.parse(profileStr); + } catch (e) { + this.logout(); + } + } + + this.updateUIState(); + }, + + updateUIState() { + const navLogin = document.getElementById('navLoginBtn'); + const navProfile = document.getElementById('navUserProfile'); + const profileName = document.getElementById('profileName'); + const profilePlan = document.getElementById('profilePlan'); + const anonBanner = document.getElementById('anonymousWarningBanner'); + + if (this.token && this.userProfile) { + // User is logged in + if (navLogin) navLogin.style.display = 'none'; + if (navProfile) navProfile.style.display = 'flex'; + if (anonBanner) anonBanner.style.display = 'none'; + + if (profileName) profileName.textContent = this.userProfile.display_name || this.userProfile.email; + + if (profilePlan) { + profilePlan.textContent = String(this.userProfile.plan).toUpperCase(); + // Set badge color based on plan + profilePlan.className = 'plan-badge ' + (this.userProfile.plan === 'pro' ? 'plan-pro' : (this.userProfile.plan === 'basic' ? 'plan-basic' : 'plan-free')); + } + } else { + // User is anonymous + if (navLogin) navLogin.style.display = 'block'; + if (navProfile) navProfile.style.display = 'none'; + + // Only show banner if we are not in self-hosted mode + // We can infer this if the user hasn't actively dismissed it + if (anonBanner && !sessionStorage.getItem('edmate_dismiss_anon_banner')) { + anonBanner.style.display = 'flex'; + } + } + }, + + getAuthHeaders() { + const headers = {}; + if (this.token) { + headers['Authorization'] = `Bearer ${this.token}`; + } + return headers; + } +}; diff --git a/qc_viewer/static/js/automate_api.js b/qc_viewer/static/js/automate_api.js index f78476d..d67a304 100644 --- a/qc_viewer/static/js/automate_api.js +++ b/qc_viewer/static/js/automate_api.js @@ -35,13 +35,21 @@ export const AutomationAPI = { const xhr = new XMLHttpRequest(); xhr.open('POST', '/api/automate/draft', true); - // Only set BYOK headers if the user provided a key - if (byokKey) { - const resolvedProvider = byokProvider || providerHint; - if (resolvedProvider) xhr.setRequestHeader('X-LLM-Provider', resolvedProvider); - xhr.setRequestHeader('X-API-Key', byokKey); - if (byokModel) xhr.setRequestHeader('X-Model-ID', byokModel); - } + // Inject Auth headers + import('/js/auth.js').then(({ AuthUI }) => { + const headers = AuthUI.getAuthHeaders(); + for (const [key, value] of Object.entries(headers)) { + xhr.setRequestHeader(key, value); + } + + // Only set BYOK headers if the user provided a key + if (byokKey) { + const resolvedProvider = byokProvider || providerHint; + if (resolvedProvider) xhr.setRequestHeader('X-LLM-Provider', resolvedProvider); + xhr.setRequestHeader('X-API-Key', byokKey); + if (byokModel) xhr.setRequestHeader('X-Model-ID', byokModel); + } + }); xhr.upload.onprogress = (e) => { if (e.lengthComputable && onProgress) { @@ -63,20 +71,26 @@ export const AutomationAPI = { }); }, + async fetchWithAuth(url, options = {}) { + const { AuthUI } = await import('/js/auth.js'); + const headers = { ...options.headers, ...AuthUI.getAuthHeaders() }; + return fetch(url, { ...options, headers }); + }, + async fetchDrafts() { - const response = await fetch('/api/automate/drafts'); + const response = await this.fetchWithAuth('/api/automate/drafts'); if (!response.ok) throw new Error('Failed to fetch drafts'); return await response.json(); }, async getDraft(id) { - const response = await fetch(`/api/automate/draft/${id}`); + const response = await this.fetchWithAuth(`/api/automate/draft/${id}`); if (!response.ok) throw new Error('Failed to fetch draft detail'); return await response.json(); }, async updateDraft(id, updates) { - const response = await fetch(`/api/automate/draft/${id}`, { + const response = await this.fetchWithAuth(`/api/automate/draft/${id}`, { method: 'PATCH', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(updates) @@ -86,20 +100,20 @@ export const AutomationAPI = { }, async deleteDraft(id) { - const response = await fetch(`/api/automate/draft/${id}`, { method: 'DELETE' }); + const response = await this.fetchWithAuth(`/api/automate/draft/${id}`, { method: 'DELETE' }); if (!response.ok) throw new Error('Failed to delete draft'); return await response.json(); }, async stopDraft(id) { - const response = await fetch(`/api/automate/draft/${id}/stop`, { method: 'POST' }); + const response = await this.fetchWithAuth(`/api/automate/draft/${id}/stop`, { method: 'POST' }); if (!response.ok) throw new Error('Failed to stop processing'); return await response.json(); }, async publishQuestion(payload) { // payload: { draft_id, table_name, question_data } - const response = await fetch('/api/automate/publish', { + const response = await this.fetchWithAuth('/api/automate/publish', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(payload) @@ -111,7 +125,7 @@ export const AutomationAPI = { async refineQuestion(payload) { // payload: { feedback, original_q } - const response = await fetch('/api/automate/refine', { + const response = await this.fetchWithAuth('/api/automate/refine', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(payload) @@ -126,7 +140,7 @@ export const AutomationAPI = { */ async exportDraft(id, format) { const q = format === 'markdown' ? 'markdown' : format; - const resp = await fetch(`/api/automate/draft/${encodeURIComponent(id)}/export?format=${encodeURIComponent(q)}`); + const resp = await this.fetchWithAuth(`/api/automate/draft/${encodeURIComponent(id)}/export?format=${encodeURIComponent(q)}`); if (!resp.ok) { let detail = `Export failed (${resp.status})`; try { diff --git a/qc_viewer/static/js/controllers/drafts.js b/qc_viewer/static/js/controllers/drafts.js index f2c4e75..c71bd75 100644 --- a/qc_viewer/static/js/controllers/drafts.js +++ b/qc_viewer/static/js/controllers/drafts.js @@ -378,6 +378,20 @@ export const DraftController = { * Floating export format picker for a draft card (PROCESSED / REVIEW_READY). */ openDraftExportPopover(anchorEl, draftId) { + import('/js/auth.js').then(({ AuthUI }) => { + // Check if user is logged in + if (!AuthUI.token) { + // Anonymous users hit the paywall for exports + const paywallModal = document.getElementById('downloadPaywallModal'); + if (paywallModal) paywallModal.style.display = 'flex'; + return; + } + + this._showExportPopover(anchorEl, draftId); + }); + }, + + _showExportPopover(anchorEl, draftId) { document.querySelector('.draft-export-popover')?.remove(); const pop = document.createElement('div'); pop.className = 'draft-export-popover'; diff --git a/qc_viewer/static/js/controllers/review.js b/qc_viewer/static/js/controllers/review.js index d89747e..2423cf1 100644 --- a/qc_viewer/static/js/controllers/review.js +++ b/qc_viewer/static/js/controllers/review.js @@ -409,11 +409,9 @@ export const ReviewController = { return; } - // Paid solutions / BYOK Paywall Gate (only active on hosted web domains) - const isLocalhost = window.location.hostname === 'localhost' || window.location.hostname === '127.0.0.1'; - if (!isLocalhost) { - const byokKey = document.getElementById('byokApiKey')?.value; - if (!byokKey) { + // Paid solutions / Auth Paywall Gate + import('/js/auth.js').then(async ({ AuthUI }) => { + if (!AuthUI.token) { const paywallModal = document.getElementById('downloadPaywallModal'); if (paywallModal) { paywallModal.style.display = 'flex'; @@ -421,14 +419,14 @@ export const ReviewController = { return; } } - } - try { - await AutomationAPI.exportDraft(this.currentDraftData.id, format); - this.showToast(`Downloaded ${String(format).toUpperCase()}`, 'success'); - this.closeExportMenu(); - } catch (err) { - this.showToast(err.message || 'Export failed', 'danger'); - } + try { + await AutomationAPI.exportDraft(this.currentDraftData.id, format); + this.showToast(`Downloaded ${String(format).toUpperCase()}`, 'success'); + this.closeExportMenu(); + } catch (err) { + this.showToast(err.message || 'Export failed', 'danger'); + } + }); }, };