Spaces:
Configuration error
Configuration error
| import os | |
| import sys | |
| import uuid | |
| import asyncio | |
| import aiofiles | |
| import concurrent.futures | |
| import logging | |
| import json | |
| import re | |
| import time | |
| import random | |
| import hmac | |
| import shutil | |
| from datetime import datetime | |
| from typing import List, Tuple, Optional, Dict, Any | |
| from contextlib import asynccontextmanager | |
| from rapidfuzz import fuzz | |
| from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks, Depends, HTTPException, Request | |
| from fastapi.responses import JSONResponse, FileResponse | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from starlette.middleware.wsgi import WSGIMiddleware | |
| from zipfile import ZipFile | |
| # Setup logging | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', | |
| handlers=[ | |
| logging.StreamHandler() | |
| ] | |
| ) | |
| logger = logging.getLogger(__name__) | |
| from langchain_core.prompts import PromptTemplate | |
| from langchain_core.documents import Document | |
| # 1) Imports for OpenAI (OpenAIEmbeddings, ChatOpenAI) | |
| from langchain_openai import ChatOpenAI, OpenAIEmbeddings | |
| # 1b) Imports for Gemini (Google Generative AI) | |
| try: | |
| from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings | |
| except ImportError: | |
| ChatGoogleGenerativeAI = None | |
| GoogleGenerativeAIEmbeddings = None | |
| import requests as _requests | |
| try: | |
| from openai import OpenAI as OpenAIClient | |
| except ImportError: | |
| OpenAIClient = None | |
| from pdf_text_processing import ( | |
| extract_text_from_pdf, | |
| extract_and_split, | |
| get_pdf_title, | |
| entity_match | |
| ) | |
| from vectorstore import get_vectorstore, compute_file_hash, CHROMA_DB_BASE | |
| from system_chart_extraction import system_chart_extraction | |
| from system_chart_service import SystemChartService | |
| from hypotheses_service import HypothesesService | |
| from hypotheses_generation import PATTERN as CHART_PATTERN | |
| from prompts import create_hypotheses_generation_prompt, create_design_goals_analysis_prompt | |
| from conversation_storage import save_conversation, get_conversations, get_all_conversations, update_conversation | |
| from plotextract.app import app as plotextract_app | |
| # Metadata management functions | |
| def load_metadata(database: str) -> Dict[str, Dict[str, Any]]: | |
| """ | |
| Load metadata from the METADATA file in the database directory. | |
| Returns a dictionary with file_hash -> metadata mapping. | |
| """ | |
| category_upload_dir = os.path.join(UPLOAD_DIR, database) | |
| metadata_file = os.path.join(category_upload_dir, "METADATA") | |
| if not os.path.exists(metadata_file): | |
| logger.debug(f"📋 No METADATA file found for {database}") | |
| return {} | |
| try: | |
| with open(metadata_file, 'r', encoding='utf-8') as f: | |
| metadata = json.load(f) | |
| logger.debug(f"📋 Loaded metadata for {database}: {len(metadata)} entries") | |
| logger.debug(f"📋 Hash keys: {[h[:16] + '...' for h in metadata.keys()]}") | |
| return metadata | |
| except Exception as e: | |
| logger.error(f"Error loading metadata for {database}: {str(e)}") | |
| return {} | |
| async def embed_stream_into_vs( | |
| texts, metas, vectorstore, embedding_model, | |
| batch_size: int = 512, | |
| max_concurrency: int = 12, | |
| progress_key: str = None | |
| ): | |
| """ | |
| Embeds texts in parallel batches and writes each batch immediately to the vectorstore. | |
| Retries on transient errors with exponential backoff. | |
| Provides real-time progress tracking for large document uploads. | |
| """ | |
| assert len(texts) == len(metas) | |
| n = len(texts) | |
| if n == 0: | |
| return | |
| batches = [(i, min(i + batch_size, n)) for i in range(0, n, batch_size)] | |
| sem = asyncio.Semaphore(max_concurrency) | |
| # Thread-safe progress tracking | |
| progress_lock = asyncio.Lock() | |
| processed_count = 0 | |
| async def run_batch(batch_idx, bi, bj): | |
| nonlocal processed_count | |
| async with sem: | |
| delay = 1.0 | |
| for attempt in range(6): | |
| try: | |
| # Generate embeddings for this batch | |
| embeds = await embedding_model.aembed_documents(texts[bi:bj]) | |
| # Add to vectorstore | |
| vectorstore.add_texts( | |
| texts=texts[bi:bj], | |
| metadatas=metas[bi:bj], | |
| embeddings=embeds | |
| ) | |
| # Update progress atomically | |
| async with progress_lock: | |
| processed_count += (bj - bi) | |
| if progress_key and progress_key in EMBEDDING_PROGRESS: | |
| # Calculate progress (10% to 90% for embedding phase) | |
| progress_percent = min(90, int(10 + (processed_count / n * 80))) | |
| EMBEDDING_PROGRESS[progress_key].update({ | |
| "status": "embedding", | |
| "processed_chunks": processed_count, | |
| "progress": progress_percent, | |
| "message": f"Embedded {processed_count}/{n} chunks ({len(batches)} batches)" | |
| }) | |
| logger.debug(f"Batch {batch_idx + 1}/{len(batches)} completed: {processed_count}/{n} chunks ({progress_percent}%)") | |
| return | |
| except Exception as e: | |
| # Handle retryable errors with exponential backoff | |
| msg = str(e).lower() | |
| if any(keyword in msg for keyword in ["rate", "429", "timeout", "temporarily", "overloaded", "quota"]): | |
| if attempt < 5: # Don't sleep on last attempt | |
| jitter = random.uniform(0, 0.5) | |
| sleep_time = delay + jitter | |
| logger.warning(f"Retryable error in batch {batch_idx + 1}, attempt {attempt + 1}/6: {str(e)}. Sleeping {sleep_time:.2f}s") | |
| await asyncio.sleep(sleep_time) | |
| delay *= 1.5 # More conservative backoff | |
| continue | |
| # Log the error with context | |
| logger.error(f"Failed to embed batch {batch_idx + 1}/{len(batches)} (chunks {bi}-{bj}) after {attempt + 1} attempts: {str(e)}") | |
| raise | |
| # Run all batches in parallel | |
| try: | |
| await asyncio.gather(*[ | |
| run_batch(idx, bi, bj) | |
| for idx, (bi, bj) in enumerate(batches) | |
| ]) | |
| # Final progress update before persisting | |
| if progress_key and progress_key in EMBEDDING_PROGRESS: | |
| EMBEDDING_PROGRESS[progress_key].update({ | |
| "status": "saving", | |
| "progress": 95, | |
| "message": f"Saving {n} embeddings to vector database..." | |
| }) | |
| # Persist the vectorstore | |
| vectorstore.persist() | |
| logger.info(f"✅ Successfully embedded and saved {n} chunks in {len(batches)} parallel batches") | |
| except Exception as e: | |
| # Update progress on failure | |
| if progress_key and progress_key in EMBEDDING_PROGRESS: | |
| EMBEDDING_PROGRESS[progress_key].update({ | |
| "status": "error", | |
| "progress": EMBEDDING_PROGRESS[progress_key].get("progress", 0), | |
| "message": f"Embedding failed: {str(e)}" | |
| }) | |
| logger.error(f"❌ Embedding failed after processing {processed_count}/{n} chunks: {str(e)}") | |
| raise | |
| def upsert_metadata(database: str, new_entries: Dict[str, Dict[str, Any]]) -> None: | |
| category_upload_dir = os.path.join(UPLOAD_DIR, database) | |
| os.makedirs(category_upload_dir, exist_ok=True) | |
| metadata_file = os.path.join(category_upload_dir, "METADATA") | |
| try: | |
| existing = {} | |
| if os.path.exists(metadata_file): | |
| with open(metadata_file, "r", encoding="utf-8") as f: | |
| try: | |
| existing = json.load(f) | |
| except Exception: | |
| logger.warning("⚠️ METADATA unreadable; rebuilding.") | |
| existing = {} | |
| existing.update(new_entries) | |
| with open(metadata_file, "w", encoding="utf-8") as f: | |
| json.dump(existing, f, indent=2, ensure_ascii=False) | |
| logger.info(f"📋 METADATA updated; total entries: {len(existing)}") | |
| except Exception as e: | |
| logger.error(f"Error updating METADATA: {e}") | |
| def save_metadata(database: str, metadata: Dict[str, Dict[str, Any]]): | |
| """ | |
| Save metadata to the METADATA file in the database directory. | |
| PROTECTION: Will not overwrite existing METADATA files to prevent changes. | |
| """ | |
| category_upload_dir = os.path.join(UPLOAD_DIR, database) | |
| os.makedirs(category_upload_dir, exist_ok=True) | |
| metadata_file = os.path.join(category_upload_dir, "METADATA") | |
| # PROTECTION: Check if METADATA file already exists | |
| if os.path.exists(metadata_file): | |
| logger.warning(f"🔒 METADATA file already exists for {database}. Skipping save to prevent modification.") | |
| return | |
| try: | |
| with open(metadata_file, 'w', encoding='utf-8') as f: | |
| json.dump(metadata, f, indent=2, ensure_ascii=False) | |
| logger.info(f"📋 Created new METADATA file for {database}") | |
| except Exception as e: | |
| logger.error(f"Error saving metadata for {database}: {str(e)}") | |
| def extract_file_metadata(file_path: str, original_filename: str, file_type: str, file_hash: str) -> Dict[str, Any]: | |
| """ | |
| Extract metadata from a file including title, size, and other relevant information. | |
| """ | |
| metadata = { | |
| "original_filename": original_filename, | |
| "saved_filename": os.path.basename(file_path), | |
| "file_type": file_type, | |
| "file_hash": file_hash, | |
| "upload_timestamp": datetime.now().isoformat(), | |
| "file_size": os.path.getsize(file_path) if os.path.exists(file_path) else 0 | |
| } | |
| # Extract title based on file type | |
| if file_type == "pdf": | |
| # For PDFs, extract title from the PDF content | |
| try: | |
| from pdf_text_processing import get_pdf_title | |
| title = get_pdf_title(file_path) | |
| metadata["title"] = title if title else original_filename | |
| except Exception as e: | |
| logger.warning(f"Could not extract title from PDF {original_filename}: {str(e)}") | |
| metadata["title"] = original_filename | |
| elif file_type == "text": | |
| # For text files, try to extract title from the file content | |
| try: | |
| from pdf_text_processing import get_file_title | |
| title = get_file_title(file_path) | |
| metadata["title"] = title if title else original_filename | |
| except Exception as e: | |
| logger.warning(f"Could not extract title from text file {original_filename}: {str(e)}") | |
| metadata["title"] = original_filename | |
| else: | |
| # For other files (images), use original filename as title | |
| metadata["title"] = original_filename | |
| return metadata | |
| def get_file_metadata(database: str, file_hash: str) -> Dict[str, Any]: | |
| """ | |
| Get metadata for a specific file by its hash. | |
| """ | |
| metadata = load_metadata(database) | |
| file_metadata = metadata.get(file_hash, {}) | |
| # Debug logging | |
| if file_metadata: | |
| logger.debug(f"✅ Found metadata for hash {file_hash[:16]}... in database {database}") | |
| else: | |
| logger.debug(f"❌ No metadata found for hash {file_hash[:16]}... in database {database}") | |
| logger.debug(f"📋 Available metadata hashes: {[h[:16] + '...' for h in metadata.keys()]}") | |
| return file_metadata | |
| def get_file_title_from_metadata(database: str, file_hash: str) -> str: | |
| """ | |
| Get the title of a file from its metadata. | |
| """ | |
| metadata = get_file_metadata(database, file_hash) | |
| title = metadata.get("title", "Unknown") | |
| # Debug logging | |
| if title == "Unknown": | |
| logger.debug(f"🔍 No title found in metadata for hash {file_hash[:16]}... in database {database}") | |
| else: | |
| logger.debug(f"✅ Found title in metadata for hash {file_hash[:16]}...: {title}") | |
| return title | |
| def resource_path(relative_path: str) -> str: | |
| """ | |
| Get the absolute path to a resource, works for dev and for PyInstaller. | |
| """ | |
| if hasattr(sys, '_MEIPASS'): | |
| return os.path.join(sys._MEIPASS, relative_path) | |
| return os.path.join(os.path.dirname(__file__), relative_path) | |
| def get_base_directory() -> str: | |
| """ | |
| HF Spaces-safe data root: | |
| - Prefer explicit override via APP_DATA_DIR or PERSISTENT_STORAGE_PATH | |
| - Else use /data/<app> only if it exists and is writable | |
| - Else use $HOME/.cache/<app> | |
| - Else use /tmp/<app> | |
| """ | |
| override = (os.environ.get("APP_DATA_DIR") or os.environ.get("PERSISTENT_STORAGE_PATH") or "").strip() | |
| if override: | |
| os.makedirs(override, exist_ok=True) | |
| return override | |
| if os.path.isdir("/data") and os.access("/data", os.W_OK): | |
| base = "/data/pdf-rag" | |
| os.makedirs(base, exist_ok=True) | |
| return base | |
| home = os.environ.get("HOME") | |
| if home: | |
| base = os.path.join(home, ".cache", "pdf-rag") | |
| os.makedirs(base, exist_ok=True) | |
| return base | |
| base = "/tmp/pdf-rag" | |
| os.makedirs(base, exist_ok=True) | |
| return base | |
| # Define paths relative to the base directory for persistent storage | |
| APP_DATA_ROOT = get_base_directory() | |
| UPLOAD_DIR = os.path.join(APP_DATA_ROOT, "uploaded_files") | |
| os.makedirs(UPLOAD_DIR, exist_ok=True) | |
| logger.info(f"UPLOAD_DIR set to: {os.path.abspath(UPLOAD_DIR)}") | |
| async def lifespan(app: FastAPI): | |
| """ | |
| FastAPI lifespan: | |
| 1) Load existing databases from previous sessions | |
| """ | |
| load_existing_databases() | |
| yield | |
| app = FastAPI(lifespan=lifespan) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"] | |
| ) | |
| app.mount("/plotextract", WSGIMiddleware(plotextract_app)) | |
| # A registry of hash -> filename for each "database" category | |
| FILE_REGISTRY = {} | |
| # Progress tracking for embedding tasks | |
| EMBEDDING_PROGRESS = {} | |
| # Hypothesis generation context storage for continuous conversation | |
| HYPOTHESIS_GENERATION_CONTEXT = {} | |
| # Conversation storage for system chart extraction is now handled by conversation_storage module | |
| def load_existing_databases(): | |
| """ | |
| Scan the structured uploaded files directory and existing databases to populate FILE_REGISTRY | |
| with information about previously uploaded files. Now uses category subdirectories. | |
| Note: Only loads text and image files since PDFs are now converted to text during upload. | |
| Also handles METADATA files for preserving file information. | |
| """ | |
| global FILE_REGISTRY | |
| logger.info("🔍 Scanning for existing databases...") | |
| # Get all database directories from root databases | |
| database_dirs = [] | |
| chroma_db_base = CHROMA_DB_BASE | |
| if os.path.exists(chroma_db_base): | |
| for item in os.listdir(chroma_db_base): | |
| db_path = os.path.join(chroma_db_base, item) | |
| if os.path.isdir(db_path) and not item.startswith('.'): | |
| database_dirs.append(item) | |
| logger.info(f"Found database directories: {database_dirs}") | |
| # For each database directory, check if there's a corresponding uploaded files subdirectory | |
| for db_name in database_dirs: | |
| if db_name not in FILE_REGISTRY: | |
| FILE_REGISTRY[db_name] = {} | |
| # Check for corresponding uploaded files subdirectory | |
| category_upload_dir = os.path.join(UPLOAD_DIR, db_name) | |
| if os.path.exists(category_upload_dir) and os.path.isdir(category_upload_dir): | |
| logger.info(f"📁 Found uploaded files directory for {db_name}") | |
| # Load existing metadata | |
| existing_metadata = load_metadata(db_name) | |
| # Scan all files in this category directory | |
| files_found = 0 | |
| metadata_needs_migration = False | |
| migrated_metadata = {} | |
| for file in os.listdir(category_upload_dir): | |
| if not file.startswith('.') and file != "METADATA": | |
| file_path = os.path.join(category_upload_dir, file) | |
| if os.path.isfile(file_path): | |
| try: | |
| # Calculate file hash and add to registry | |
| with open(file_path, 'rb') as f: | |
| content = f.read() | |
| file_hash = compute_file_hash(content) | |
| FILE_REGISTRY[db_name][file_hash] = file | |
| files_found += 1 | |
| # Check if this file exists in metadata with the current hash | |
| if file_hash in existing_metadata: | |
| # Perfect match - use existing metadata | |
| migrated_metadata[file_hash] = existing_metadata[file_hash] | |
| else: | |
| # Hash mismatch - try to find by filename pattern for .txt files (migrated from PDF) | |
| found_metadata = None | |
| if file.endswith('.txt'): | |
| # Try to find metadata entry that matches this file's original filename pattern | |
| for old_hash, meta in existing_metadata.items(): | |
| if meta.get('saved_filename') == file: | |
| logger.info(f"🔄 Migrating metadata for {file}: {old_hash[:16]}... -> {file_hash[:16]}...") | |
| found_metadata = meta.copy() | |
| found_metadata['file_hash'] = file_hash # Update to new hash | |
| migrated_metadata[file_hash] = found_metadata | |
| metadata_needs_migration = True | |
| break | |
| # If still not found, try matching by original filename | |
| if not found_metadata: | |
| for old_hash, meta in existing_metadata.items(): | |
| # Extract the original filename without extension and compare | |
| original_name = os.path.splitext(meta.get('original_filename', ''))[0] | |
| file_name_part = file.split('_', 1)[1] if '_' in file else file # Remove UUID prefix | |
| file_name_part = os.path.splitext(file_name_part)[0] # Remove .txt extension | |
| if original_name and original_name == file_name_part: | |
| logger.info(f"🔄 Migrating metadata by filename match for {file}: {original_name}") | |
| found_metadata = meta.copy() | |
| found_metadata['file_hash'] = file_hash # Update to new hash | |
| found_metadata['saved_filename'] = file # Update saved filename | |
| migrated_metadata[file_hash] = found_metadata | |
| metadata_needs_migration = True | |
| break | |
| # If no metadata found, create new metadata | |
| if not found_metadata: | |
| logger.info(f"📋 Creating metadata for existing file: {file}") | |
| # Determine file type | |
| file_ext = os.path.splitext(file.lower())[1] | |
| if file_ext in ['.txt', '.text']: | |
| file_type = "text" # Could be from PDF or original text | |
| elif file_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']: | |
| file_type = "image" | |
| else: | |
| file_type = "unknown" | |
| # Create metadata for the file | |
| file_metadata = extract_file_metadata(file_path, file, file_type, file_hash) | |
| migrated_metadata[file_hash] = file_metadata | |
| metadata_needs_migration = True | |
| except Exception as e: | |
| logger.warning(f"Error processing file {file}: {e}") | |
| # Update existing_metadata with migrated data | |
| existing_metadata = migrated_metadata | |
| # Log migration results | |
| if metadata_needs_migration: | |
| logger.info(f"🔄 Metadata migration completed for {db_name} database") | |
| # Save updated metadata (only if METADATA file doesn't exist yet) | |
| if existing_metadata: | |
| metadata_file = os.path.join(category_upload_dir, "METADATA") | |
| if not os.path.exists(metadata_file): | |
| save_metadata(db_name, existing_metadata) | |
| else: | |
| logger.info(f"🔒 METADATA file already exists for {db_name}. Skipping metadata updates.") | |
| logger.info(f"📁 Added {files_found} files to {db_name} database") | |
| else: | |
| logger.info(f"⚠️ No uploaded files directory found for {db_name} (expected: {category_upload_dir})") | |
| # Also check for any uploaded files subdirectories that don't have corresponding databases yet | |
| if os.path.exists(UPLOAD_DIR): | |
| for item in os.listdir(UPLOAD_DIR): | |
| item_path = os.path.join(UPLOAD_DIR, item) | |
| if os.path.isdir(item_path) and not item.startswith('.'): | |
| if item not in FILE_REGISTRY: | |
| # Found a files directory without a corresponding database | |
| logger.info(f"📂 Found uploaded files directory '{item}' without corresponding database") | |
| FILE_REGISTRY[item] = {} | |
| # Load existing metadata | |
| existing_metadata = load_metadata(item) | |
| # Add files from this directory | |
| files_found = 0 | |
| metadata_needs_migration = False | |
| migrated_metadata = {} | |
| for file in os.listdir(item_path): | |
| if not file.startswith('.') and file != "METADATA": | |
| file_path = os.path.join(item_path, file) | |
| if os.path.isfile(file_path): | |
| try: | |
| with open(file_path, 'rb') as f: | |
| content = f.read() | |
| file_hash = compute_file_hash(content) | |
| FILE_REGISTRY[item][file_hash] = file | |
| files_found += 1 | |
| # Check if this file exists in metadata with the current hash | |
| if file_hash in existing_metadata: | |
| # Perfect match - use existing metadata | |
| migrated_metadata[file_hash] = existing_metadata[file_hash] | |
| else: | |
| # Hash mismatch - try to find by filename pattern for .txt files (migrated from PDF) | |
| found_metadata = None | |
| if file.endswith('.txt'): | |
| # Try to find metadata entry that matches this file's original filename pattern | |
| for old_hash, meta in existing_metadata.items(): | |
| if meta.get('saved_filename') == file: | |
| logger.info(f"🔄 Migrating metadata for {file}: {old_hash[:16]}... -> {file_hash[:16]}...") | |
| found_metadata = meta.copy() | |
| found_metadata['file_hash'] = file_hash # Update to new hash | |
| migrated_metadata[file_hash] = found_metadata | |
| metadata_needs_migration = True | |
| break | |
| # If still not found, try matching by original filename | |
| if not found_metadata: | |
| for old_hash, meta in existing_metadata.items(): | |
| # Extract the original filename without extension and compare | |
| original_name = os.path.splitext(meta.get('original_filename', ''))[0] | |
| file_name_part = file.split('_', 1)[1] if '_' in file else file # Remove UUID prefix | |
| file_name_part = os.path.splitext(file_name_part)[0] # Remove .txt extension | |
| if original_name and original_name == file_name_part: | |
| logger.info(f"🔄 Migrating metadata by filename match for {file}: {original_name}") | |
| found_metadata = meta.copy() | |
| found_metadata['file_hash'] = file_hash # Update to new hash | |
| found_metadata['saved_filename'] = file # Update saved filename | |
| migrated_metadata[file_hash] = found_metadata | |
| metadata_needs_migration = True | |
| break | |
| # If no metadata found, create new metadata | |
| if not found_metadata: | |
| logger.info(f"📋 Creating metadata for existing file: {file}") | |
| # Determine file type | |
| file_ext = os.path.splitext(file.lower())[1] | |
| if file_ext in ['.txt', '.text']: | |
| file_type = "text" | |
| elif file_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']: | |
| file_type = "image" | |
| else: | |
| file_type = "unknown" | |
| # Create metadata for the file | |
| file_metadata = extract_file_metadata(file_path, file, file_type, file_hash) | |
| migrated_metadata[file_hash] = file_metadata | |
| metadata_needs_migration = True | |
| except Exception as e: | |
| logger.warning(f"Error processing file {file}: {e}") | |
| # Update existing_metadata with migrated data | |
| existing_metadata = migrated_metadata | |
| # Log migration results | |
| if metadata_needs_migration: | |
| logger.info(f"🔄 Metadata migration completed for {item} database") | |
| # Save updated metadata (only if METADATA file doesn't exist yet) | |
| if existing_metadata: | |
| metadata_file = os.path.join(item_path, "METADATA") | |
| if not os.path.exists(metadata_file): | |
| save_metadata(item, existing_metadata) | |
| else: | |
| logger.info(f"🔒 METADATA file already exists for {item}. Skipping metadata updates.") | |
| logger.info(f"📁 Added {files_found} files to {item} (files-only database)") | |
| # Clean up empty databases | |
| empty_dbs = [db for db, files in FILE_REGISTRY.items() if not files] | |
| for db in empty_dbs: | |
| del FILE_REGISTRY[db] | |
| logger.info(f"🗑️ Removed empty database: {db}") | |
| logger.info(f"✅ Loaded {len(FILE_REGISTRY)} databases with files:") | |
| for db_name, files in FILE_REGISTRY.items(): | |
| logger.info(f" 📚 {db_name}: {len(files)} files") | |
| # Utility functions for file type detection | |
| def is_image_file(filename: str) -> bool: | |
| """Check if a file is an image based on its extension.""" | |
| image_extensions = ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp'] | |
| file_ext = os.path.splitext(filename.lower())[1] | |
| return file_ext in image_extensions | |
| def is_text_or_pdf_file(filename: str) -> bool: | |
| """Check if a file is a PDF or text file based on its extension.""" | |
| text_pdf_extensions = ['.pdf', '.txt', '.text'] | |
| file_ext = os.path.splitext(filename.lower())[1] | |
| return file_ext in text_pdf_extensions | |
| def get_file_type(filename: str) -> str: | |
| """Return the file type category: 'image', 'text_pdf', or 'unknown'.""" | |
| if is_image_file(filename): | |
| return 'image' | |
| elif is_text_or_pdf_file(filename): | |
| return 'text_pdf' | |
| else: | |
| return 'unknown' | |
| # Conversation management functions are now handled by conversation_storage module | |
| # -------------------- Per-browser workspaces -------------------- # | |
| # Each browser gets a random workspace id in a cookie. Every server-side key | |
| # (database, progress/extraction id, conversation, hypothesis run) is stored as | |
| # "<workspace>-<name>", so users only see their own data; the client only ever | |
| # sees the unprefixed names. API keys and models are held per workspace too. | |
| WORKSPACE_COOKIE = "madtool_workspace" | |
| WORKSPACE_MODELS: Dict[str, Dict[str, Any]] = {} # workspace -> {"embedding", "qa", "chart"} | |
| def is_workspace_id(value: Optional[str]) -> bool: | |
| return bool(value) and re.fullmatch(r"[0-9a-f]{32}", value) is not None | |
| def workspace(request: Request) -> str: | |
| ws = request.cookies.get(WORKSPACE_COOKIE) | |
| if not is_workspace_id(ws): | |
| raise HTTPException(status_code=401, detail="Your session has no workspace. Please reload the page.") | |
| return ws | |
| def scoped(ws: str, name: str) -> str: | |
| return f"{ws}-{name}" | |
| def unscoped(ws: str, key: str) -> Optional[str]: | |
| """The client-facing name of a key in workspace `ws`, or None if the key isn't theirs.""" | |
| prefix = f"{ws}-" | |
| return key[len(prefix):] if key.startswith(prefix) else None | |
| def get_index(request: Request): | |
| response = FileResponse(resource_path("index.html")) | |
| if not is_workspace_id(request.cookies.get(WORKSPACE_COOKIE)): | |
| # Deployed, the site is always HTTPS (HF's proxy terminates TLS and doesn't say so), | |
| # and SameSite=None lets the cookie work inside the huggingface.co Space iframe. | |
| # Local http dev keeps a plain Lax cookie. | |
| secure = request.url.hostname not in ("localhost", "127.0.0.1") | |
| response.set_cookie( | |
| WORKSPACE_COOKIE, uuid.uuid4().hex, max_age=365 * 24 * 3600, httponly=True, | |
| secure=secure, samesite="none" if secure else "lax", | |
| ) | |
| return response | |
| # -------------------- Model / API related info -------------------- # | |
| OPENAI_MODELS = [ | |
| "gpt-5-nano", | |
| "gpt-5-mini", | |
| "gpt-5.2" | |
| ] | |
| GEMINI_MODELS = [ | |
| "gemini-3-flash-preview", | |
| "gemini-3-pro-preview", | |
| "gemini-2.5-flash-lite" | |
| ] | |
| AVAILABLE_MODELS = OPENAI_MODELS + GEMINI_MODELS | |
| DEFAULT_MODEL = "gpt-5.2" | |
| # Shared output-length cap for QA, system chart extraction, and hypotheses. | |
| # LangChain's max_tokens maps to provider-specific names: | |
| # OpenAI max_output_tokens / max_completion_tokens, Gemini maxOutputTokens. | |
| DEFAULT_MAX_OUTPUT_TOKENS = 8192 | |
| ADMIN_DELETE_CONFIRMATION_TEXT = "DELETE APP DATA" | |
| def get_model_provider(model: str) -> Optional[str]: | |
| if model in OPENAI_MODELS: | |
| return "openai" | |
| if model in GEMINI_MODELS: | |
| return "gemini" | |
| return None | |
| def detect_provider_from_key(api_key: str) -> Optional[str]: | |
| """Infer provider from an API key prefix. | |
| AIza... → Gemini, sk-... → OpenAI. | |
| Anthropic (sk-ant-...) keys are detected but unsupported — they must be | |
| checked *before* the generic sk- prefix so they are not misread as OpenAI; | |
| callers reject the "anthropic" value with a clear message. | |
| """ | |
| key = (api_key or "").strip() | |
| if not key: | |
| return None | |
| if key.startswith("sk-ant-"): | |
| return "anthropic" | |
| if key.startswith("AIza"): | |
| return "gemini" | |
| if key.startswith("sk-"): | |
| return "openai" | |
| return None | |
| UNSUPPORTED_ANTHROPIC_KEY_MESSAGE = ( | |
| "Anthropic (sk-ant-...) keys are not supported: Anthropic does not provide an " | |
| "embedding model, which this tool requires. Please use an OpenAI (sk-...) or " | |
| "Gemini (AIza...) key." | |
| ) | |
| def provider_display_name(provider: str) -> str: | |
| return {"openai": "OpenAI", "gemini": "Gemini"}.get(provider, provider or "Unknown") | |
| # Substrings of OpenAI model ids that are not text-chat models (same list | |
| # PlotExtract uses, plus completions-only "instruct" models). "search" also | |
| # covers the tool-only deep-research models. | |
| OPENAI_NON_CHAT_MARKERS = ( | |
| "audio", | |
| "dall-e", | |
| "embedding", | |
| "gpt-image", | |
| "instruct", | |
| "moderation", | |
| "realtime", | |
| "search", | |
| "transcribe", | |
| "tts", | |
| "whisper", | |
| ) | |
| def is_openai_chat_model(model_id: str) -> bool: | |
| mid = model_id.lower() | |
| is_family = "gpt" in mid or re.match(r"^o\d+(?:-|$)", mid) # GPT and o-series reasoning models | |
| return bool(is_family) and not any(marker in mid for marker in OPENAI_NON_CHAT_MARKERS) | |
| def list_openai_chat_models(api_key: str) -> List[Dict[str, str]]: | |
| if OpenAIClient is None: | |
| raise RuntimeError("openai SDK is not installed.") | |
| client = OpenAIClient(api_key=api_key) | |
| page = client.models.list() | |
| out: List[Dict[str, Any]] = [] | |
| seen = set() | |
| for item in getattr(page, "data", []) or []: | |
| model_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) | |
| if not model_id or model_id in seen: | |
| continue | |
| if not is_openai_chat_model(model_id): | |
| continue | |
| seen.add(model_id) | |
| created = getattr(item, "created", None) or (item.get("created") if isinstance(item, dict) else None) | |
| out.append({"id": model_id, "display_name": model_id, "created": created or 0}) | |
| # Newest first by release timestamp | |
| out.sort(key=lambda m: (m["created"], m["id"]), reverse=True) | |
| for m in out: | |
| m.pop("created") | |
| return out | |
| def _fetch_gemini_model_list(api_key: str) -> List[Dict[str, Any]]: | |
| url = "https://generativelanguage.googleapis.com/v1beta/models" | |
| response = _requests.get(url, params={"key": api_key}, timeout=15) | |
| if response.status_code != 200: | |
| try: | |
| err = response.json().get("error", {}).get("message") or response.text | |
| except Exception: | |
| err = response.text | |
| raise RuntimeError(f"Gemini API returned {response.status_code}: {err}") | |
| payload = response.json() or {} | |
| return payload.get("models", []) or [] | |
| # Substrings of Gemini model ids that support generateContent but are not | |
| # text-chat models (speech, image generation, live audio, tool-only agents). | |
| GEMINI_NON_CHAT_MARKERS = ( | |
| "aqa", | |
| "audio", | |
| "computer-use", | |
| "embedding", | |
| "image", | |
| "live", | |
| "robotics", | |
| "tts", | |
| ) | |
| def gemini_sort_key(model_id: str) -> Tuple: | |
| """Newest-first ordering for Gemini ids, which carry no release date: the | |
| version number tracks release order, then pro > flash > flash-lite, then | |
| stable before preview/experimental. "-latest" aliases float to the top. | |
| """ | |
| mid = model_id.lower() | |
| version = re.search(r"-(\d+(?:\.\d+)?)(?:-|$)", mid) | |
| if "lite" in mid: | |
| tier = 1 | |
| elif "flash" in mid: | |
| tier = 2 | |
| elif "pro" in mid: | |
| tier = 3 | |
| else: | |
| tier = 0 | |
| return ( | |
| "latest" in mid, | |
| mid.startswith("gemini"), | |
| float(version.group(1)) if version else 0.0, | |
| tier, | |
| not any(tag in mid for tag in ("preview", "exp")), | |
| mid, | |
| ) | |
| def list_gemini_models(api_key: str) -> List[Dict[str, str]]: | |
| """List Gemini chat models via the public REST API to avoid pinning conflicts | |
| between google-generativeai and langchain-google-genai. | |
| """ | |
| out: List[Dict[str, str]] = [] | |
| seen = set() | |
| for item in _fetch_gemini_model_list(api_key): | |
| methods = item.get("supportedGenerationMethods") or [] | |
| if "generateContent" not in methods: | |
| continue | |
| raw_id = item.get("name", "") | |
| model_id = raw_id.split("/", 1)[-1] if raw_id else "" | |
| if not model_id or model_id in seen: | |
| continue | |
| if any(marker in model_id.lower() for marker in GEMINI_NON_CHAT_MARKERS): | |
| continue | |
| seen.add(model_id) | |
| display = item.get("displayName") or model_id | |
| out.append({"id": model_id, "display_name": str(display)}) | |
| out.sort(key=lambda m: gemini_sort_key(m["id"]), reverse=True) | |
| return out | |
| def pick_gemini_embedding_model(api_key: str) -> str: | |
| """Find a Gemini embedding model this key actually has access to. | |
| The v1beta endpoint exposes different embedding models depending on the key's | |
| tier and rollout — `embedding-001` and `text-embedding-004` are sometimes 404. | |
| We list models, filter to those supporting `embedContent`, and pick the | |
| newest one we recognise. | |
| """ | |
| try: | |
| embedding_models: List[str] = [] | |
| for item in _fetch_gemini_model_list(api_key): | |
| methods = item.get("supportedGenerationMethods") or [] | |
| if "embedContent" not in methods: | |
| continue | |
| name = item.get("name") or "" | |
| if name: | |
| embedding_models.append(name) | |
| for preferred_prefix in ( | |
| "models/gemini-embedding", | |
| "models/text-embedding", | |
| "models/embedding", | |
| ): | |
| for name in embedding_models: | |
| if name.startswith(preferred_prefix): | |
| return name | |
| if embedding_models: | |
| return embedding_models[0] | |
| except Exception as exc: | |
| logger.warning(f"Failed to pick Gemini embedding model dynamically, falling back: {exc}") | |
| return "models/gemini-embedding-001" | |
| def fetch_models_for_provider(provider: str, api_key: str) -> List[Dict[str, str]]: | |
| if provider == "openai": | |
| return list_openai_chat_models(api_key) | |
| if provider == "gemini": | |
| return list_gemini_models(api_key) | |
| return [] | |
| def choose_default_model(models: List[Dict[str, str]], provider: str) -> str: | |
| # The provider lists are already sorted newest-first, so the default is | |
| # simply the top entry. | |
| return models[0]["id"] if models else "" | |
| class AsyncEmbeddingAdapter: | |
| def __init__(self, embedder): | |
| self._embedder = embedder | |
| async def aembed_documents(self, texts): | |
| if hasattr(self._embedder, "aembed_documents"): | |
| return await self._embedder.aembed_documents(texts) | |
| return await asyncio.to_thread(self._embedder.embed_documents, texts) | |
| def __getattr__(self, name): | |
| return getattr(self._embedder, name) | |
| def ensure_async_embeddings(embedder): | |
| if hasattr(embedder, "aembed_documents"): | |
| return embedder | |
| return AsyncEmbeddingAdapter(embedder) | |
| def llm_content_to_text(content: Any) -> str: | |
| """ | |
| Normalize LLM content payloads (string/list/dict) into plain text. | |
| Some providers (e.g., Responses API) may return structured blocks. | |
| """ | |
| if content is None: | |
| return "" | |
| if isinstance(content, str): | |
| return content | |
| if isinstance(content, list): | |
| parts: List[str] = [] | |
| for item in content: | |
| text = llm_content_to_text(item) | |
| if text: | |
| parts.append(text) | |
| return "\n".join(parts) | |
| if isinstance(content, dict): | |
| # Common structured message shapes from provider SDKs. | |
| for key in ("text", "content", "output_text"): | |
| if key in content: | |
| return llm_content_to_text(content.get(key)) | |
| return json.dumps(content, ensure_ascii=False) | |
| return str(content) | |
| async def set_api_key( | |
| api_key_input: Optional[str] = Form(None), | |
| model: str = Form(DEFAULT_MODEL), | |
| ws: str = Depends(workspace), | |
| ): | |
| """ | |
| Sets up this workspace's LLM and embeddings for proprietary models. | |
| """ | |
| if not api_key_input: | |
| return JSONResponse( | |
| {"message": "Error: API key is required for proprietary models."}, | |
| status_code=400 | |
| ) | |
| api_key = api_key_input.strip() | |
| model_provider = detect_provider_from_key(api_key) | |
| if model_provider == "anthropic": | |
| return JSONResponse( | |
| content={"message": UNSUPPORTED_ANTHROPIC_KEY_MESSAGE}, | |
| status_code=400 | |
| ) | |
| if not model_provider: | |
| return JSONResponse( | |
| content={"message": "Unsupported API key format. Use an OpenAI (sk-...) or Gemini (AIza...) key."}, | |
| status_code=400 | |
| ) | |
| if not model: | |
| return JSONResponse( | |
| content={"message": "Please select a model."}, | |
| status_code=400 | |
| ) | |
| DEFAULT_MODEL_QA = model | |
| DEFAULT_MODEL_SYSTEM_CHART_EXTRACTION = model | |
| try: | |
| if not api_key: | |
| return JSONResponse( | |
| content={"message": "Error: API key is required for proprietary models."}, | |
| status_code=400 | |
| ) | |
| if model_provider == "openai": | |
| embedding_model = OpenAIEmbeddings(openai_api_key=api_key) | |
| embedding_model = ensure_async_embeddings(embedding_model) | |
| openai_common_kwargs: Dict[str, Any] = { | |
| "openai_api_key": api_key, | |
| "model": DEFAULT_MODEL_QA, | |
| "use_responses_api": True, | |
| "max_tokens": DEFAULT_MAX_OUTPUT_TOKENS, | |
| } | |
| openai_chart_kwargs: Dict[str, Any] = { | |
| "openai_api_key": api_key, | |
| "model": DEFAULT_MODEL_SYSTEM_CHART_EXTRACTION, | |
| "use_responses_api": True, | |
| "max_tokens": DEFAULT_MAX_OUTPUT_TOKENS, | |
| } | |
| llm_QA = ChatOpenAI( | |
| **openai_common_kwargs | |
| ) | |
| llm_system_chart_extraction = ChatOpenAI( | |
| **openai_chart_kwargs | |
| ) | |
| provider_label = "OpenAI" | |
| elif model_provider == "gemini": | |
| if ChatGoogleGenerativeAI is None or GoogleGenerativeAIEmbeddings is None: | |
| return JSONResponse( | |
| content={"message": "Gemini models require langchain-google-genai. Please install it to continue."}, | |
| status_code=500 | |
| ) | |
| gemini_embedding_model_name = pick_gemini_embedding_model(api_key) | |
| logger.info(f"Using Gemini embedding model: {gemini_embedding_model_name}") | |
| embedding_model = GoogleGenerativeAIEmbeddings( | |
| google_api_key=api_key, | |
| model=gemini_embedding_model_name | |
| ) | |
| embedding_model = ensure_async_embeddings(embedding_model) | |
| gemini_common_kwargs: Dict[str, Any] = { | |
| "google_api_key": api_key, | |
| "model": DEFAULT_MODEL_QA, | |
| "max_tokens": DEFAULT_MAX_OUTPUT_TOKENS, | |
| } | |
| gemini_chart_kwargs: Dict[str, Any] = { | |
| "google_api_key": api_key, | |
| "model": DEFAULT_MODEL_SYSTEM_CHART_EXTRACTION, | |
| "max_tokens": DEFAULT_MAX_OUTPUT_TOKENS, | |
| } | |
| llm_QA = ChatGoogleGenerativeAI( | |
| **gemini_common_kwargs | |
| ) | |
| llm_system_chart_extraction = ChatGoogleGenerativeAI( | |
| **gemini_chart_kwargs | |
| ) | |
| provider_label = "Gemini" | |
| else: | |
| return JSONResponse( | |
| content={"message": f"Unsupported provider: {model_provider}."}, | |
| status_code=400 | |
| ) | |
| WORKSPACE_MODELS[ws] = {"embedding": embedding_model, "qa": llm_QA, "chart": llm_system_chart_extraction} | |
| return JSONResponse( | |
| content={"message": f"{provider_label} model '{model}' set successfully!"}, | |
| status_code=200 | |
| ) | |
| except Exception as e: | |
| return JSONResponse( | |
| content={"message": f"Error setting up model: {str(e)}"}, | |
| status_code=500 | |
| ) | |
| async def get_available_models(): | |
| return JSONResponse(content={ | |
| "models": AVAILABLE_MODELS | |
| }) | |
| async def list_models_for_key(api_key_input: Optional[str] = Form(None)): | |
| """Return the models available for the provider that the supplied API key belongs to. | |
| Provider is inferred from the key prefix; the dropdown only ever shows models | |
| matching the detected provider. | |
| """ | |
| key = (api_key_input or "").strip() | |
| if not key: | |
| return JSONResponse({"message": "API key is required."}, status_code=400) | |
| provider = detect_provider_from_key(key) | |
| if provider == "anthropic": | |
| return JSONResponse( | |
| {"message": UNSUPPORTED_ANTHROPIC_KEY_MESSAGE}, | |
| status_code=400, | |
| ) | |
| if not provider: | |
| return JSONResponse( | |
| {"message": "Unsupported API key format. Use an OpenAI (sk-...) or Gemini (AIza...) key."}, | |
| status_code=400, | |
| ) | |
| try: | |
| models = fetch_models_for_provider(provider, key) | |
| except Exception as exc: | |
| return JSONResponse( | |
| {"message": f"Failed to fetch models from {provider_display_name(provider)}.", "details": str(exc)}, | |
| status_code=400, | |
| ) | |
| if not models: | |
| return JSONResponse( | |
| {"message": f"No models available for this {provider_display_name(provider)} key."}, | |
| status_code=400, | |
| ) | |
| return JSONResponse({ | |
| "provider": provider, | |
| "provider_label": provider_display_name(provider), | |
| "models": models, | |
| "default_model": choose_default_model(models, provider), | |
| }) | |
| # -------------------- Upload -------------------- # | |
| async def save_file(file: UploadFile, database: str) -> Tuple[str, str]: | |
| """ | |
| Saves text or image files to disk in a subdirectory for the database. | |
| For PDFs: extracts text and saves only the .txt file (not the original PDF). | |
| For text files: saves the text file as-is. | |
| For images: saves the image file as-is. | |
| Also extracts and saves metadata including paper titles to a METADATA file. | |
| Returns (file_path, duplicate_filename). | |
| If duplicate is detected, returns (None, original_filename). | |
| """ | |
| # Check if file type is supported | |
| file_ext = os.path.splitext(file.filename.lower())[1] | |
| supported_extensions = ['.pdf', '.txt', '.text', '.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp'] | |
| if file_ext not in supported_extensions: | |
| logger.warning(f"⚠️ Unsupported file type: {file.filename}") | |
| return None, f"Unsupported file type: {file.filename}" | |
| content = await file.read() | |
| # For non-PDF files, check duplicates immediately | |
| # For PDFs, duplicate check happens after text extraction (due to hash consistency) | |
| if file_ext != '.pdf': | |
| file_hash = compute_file_hash(content) | |
| if file_hash in FILE_REGISTRY[database]: | |
| logger.warning(f"⚠️ Duplicate detected: {file.filename}") | |
| return None, file.filename | |
| # Create category-specific directory | |
| category_upload_dir = os.path.join(UPLOAD_DIR, database) | |
| os.makedirs(category_upload_dir, exist_ok=True) | |
| # PROTECTION: Check if METADATA file already exists | |
| metadata_file = os.path.join(category_upload_dir, "METADATA") | |
| if os.path.exists(metadata_file): | |
| error_msg = f"🔒 Database '{database}' is locked. Cannot add files to databases with existing METADATA files." | |
| logger.warning(error_msg) | |
| return None, error_msg | |
| file_id = str(uuid.uuid4())[:8] | |
| image_extensions = ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp'] | |
| if file_ext == '.pdf': | |
| # For PDFs: Extract text AND preserve original PDF file | |
| logger.info(f"Processing PDF: {file.filename} (will save both text and original PDF)") | |
| # Create temporary file to extract text and metadata from | |
| import tempfile | |
| with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as temp_pdf: | |
| temp_pdf.write(content) | |
| temp_pdf_path = temp_pdf.name | |
| try: | |
| # Extract text from the temporary PDF | |
| docs = extract_text_from_pdf(temp_pdf_path) | |
| if not docs: | |
| logger.warning(f"⚠️ No text could be extracted from PDF: {file.filename}") | |
| os.unlink(temp_pdf_path) # Clean up temp file | |
| return None, f"No text could be extracted from: {file.filename}" | |
| # Save the extracted text file | |
| text_filename = f"{file_id}_{os.path.splitext(file.filename)[0]}.txt" | |
| text_path = os.path.join(category_upload_dir, text_filename) | |
| async with aiofiles.open(text_path, "w", encoding="utf-8") as tfile: | |
| await tfile.write(docs[0].page_content) | |
| # ALSO save the original PDF file | |
| pdf_filename = f"{file_id}_{file.filename}" | |
| pdf_path = os.path.join(category_upload_dir, pdf_filename) | |
| async with aiofiles.open(pdf_path, "wb") as pfile: | |
| await pfile.write(content) | |
| # CRITICAL FIX: Recompute hash from the saved .txt file content, not original PDF | |
| # This ensures hash consistency between upload and loading phases | |
| with open(text_path, "rb") as txt_file: | |
| txt_content = txt_file.read() | |
| file_hash = compute_file_hash(txt_content) # Use TXT hash instead of PDF hash | |
| # Check for duplicates with the new hash | |
| if file_hash in FILE_REGISTRY[database]: | |
| logger.warning(f"⚠️ Duplicate detected after text extraction: {file.filename}") | |
| os.unlink(temp_pdf_path) # Clean up temp file | |
| os.unlink(pdf_path) # Clean up saved PDF | |
| return None, file.filename | |
| # Extract and save metadata (including title from original PDF and PDF path) | |
| file_metadata = extract_file_metadata(temp_pdf_path, file.filename, "pdf", file_hash) | |
| file_metadata["pdf_path"] = pdf_filename # Store PDF filename | |
| upsert_metadata(database, {file_hash: file_metadata}) | |
| # Register the text file in FILE_REGISTRY (for text-based operations) | |
| FILE_REGISTRY[database][file_hash] = text_filename | |
| logger.info(f"📄 Extracted and saved text: {text_filename} in {database}") | |
| logger.info(f"📄 Also saved original PDF: {pdf_filename}") | |
| logger.info(f"📋 Saved metadata with title: {file_metadata.get('title', 'Unknown')}") | |
| logger.info(f"🔑 Using TXT file hash: {file_hash[:16]}... for consistency") | |
| os.unlink(temp_pdf_path) # Clean up temp file | |
| return text_path, None | |
| except Exception as e: | |
| logger.error(f"❌ Error processing PDF {file.filename}: {str(e)}") | |
| os.unlink(temp_pdf_path) # Clean up temp file | |
| return None, f"Error processing PDF: {str(e)}" | |
| elif file_ext in image_extensions: | |
| # For images: Save the image file as-is | |
| filename = f"{file_id}_{file.filename}" | |
| file_path = os.path.join(category_upload_dir, filename) | |
| logger.info(f"Saving image to: {os.path.abspath(file_path)}") | |
| async with aiofiles.open(file_path, "wb") as f: | |
| await f.write(content) | |
| # Extract and save metadata for image (file_hash already computed above) | |
| file_metadata = extract_file_metadata(file_path, file.filename, "image", file_hash) | |
| upsert_metadata(database, {file_hash: file_metadata}) | |
| FILE_REGISTRY[database][file_hash] = filename | |
| logger.info(f"📸 Saved image: {filename} in {database}") | |
| logger.info(f"📋 Saved metadata for image: {file.filename}") | |
| return file_path, None | |
| else: | |
| # For text files: Save as text with UTF-8 encoding | |
| filename = f"{file_id}_{file.filename}" | |
| file_path = os.path.join(category_upload_dir, filename) | |
| logger.info(f"Saving text file to: {os.path.abspath(file_path)}") | |
| async with aiofiles.open(file_path, "w", encoding="utf-8") as f: | |
| await f.write(content.decode('utf-8')) | |
| # Extract and save metadata for text file (file_hash already computed above) | |
| file_metadata = extract_file_metadata(file_path, file.filename, "text", file_hash) | |
| upsert_metadata(database, {file_hash: file_metadata}) | |
| FILE_REGISTRY[database][file_hash] = filename | |
| logger.info(f"📝 Saved text file: {filename} in {database}") | |
| logger.info(f"📋 Saved metadata with title: {file_metadata.get('title', 'Unknown')}") | |
| return file_path, None | |
| def save_file_from_bytes(content: bytes, original_filename: str, database: str) -> Tuple[str, str]: | |
| """ | |
| Save a single file provided as raw bytes with duplicate checks and metadata handling. | |
| Returns (file_path, duplicate_or_error_message). If duplicate or error, file_path is None. | |
| Mirrors logic in save_file() but for in-memory content (e.g., ZIP entries). | |
| """ | |
| # Check if file type is supported | |
| file_ext = os.path.splitext(original_filename.lower())[1] | |
| supported_extensions = ['.pdf', '.txt', '.text', '.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp'] | |
| if file_ext not in supported_extensions: | |
| logger.warning(f"⚠️ Unsupported file type in ZIP: {original_filename}") | |
| return None, f"Unsupported file type: {original_filename}" | |
| # For non-PDF files, check duplicates immediately | |
| if file_ext != '.pdf': | |
| file_hash = compute_file_hash(content) | |
| if file_hash in FILE_REGISTRY[database]: | |
| logger.warning(f"⚠️ Duplicate detected in ZIP: {original_filename}") | |
| return None, original_filename | |
| # Create category-specific directory | |
| category_upload_dir = os.path.join(UPLOAD_DIR, database) | |
| os.makedirs(category_upload_dir, exist_ok=True) | |
| file_id = str(uuid.uuid4())[:8] | |
| image_extensions = ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp'] | |
| if file_ext == '.pdf': | |
| # For PDFs: Extract text AND preserve original PDF file | |
| logger.info(f"Processing PDF from ZIP: {original_filename} (will save both text and original PDF)") | |
| import tempfile | |
| # Create temporary file to extract text and metadata from | |
| with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as temp_pdf: | |
| temp_pdf.write(content) | |
| temp_pdf_path = temp_pdf.name | |
| try: | |
| # Extract text from the temporary PDF | |
| docs = extract_text_from_pdf(temp_pdf_path) | |
| if not docs: | |
| logger.warning(f"⚠️ No text could be extracted from PDF (ZIP): {original_filename}") | |
| os.unlink(temp_pdf_path) | |
| return None, f"No text could be extracted from: {original_filename}" | |
| # Save the extracted text file | |
| text_filename = f"{file_id}_{os.path.splitext(original_filename)[0]}.txt" | |
| text_path = os.path.join(category_upload_dir, text_filename) | |
| with open(text_path, "w", encoding="utf-8") as tfile: | |
| tfile.write(docs[0].page_content) | |
| # ALSO save the original PDF file | |
| pdf_filename = f"{file_id}_{original_filename}" | |
| pdf_path = os.path.join(category_upload_dir, pdf_filename) | |
| with open(pdf_path, "wb") as pfile: | |
| pfile.write(content) | |
| # Recompute hash from the saved .txt file content | |
| with open(text_path, "rb") as txt_file: | |
| txt_content = txt_file.read() | |
| file_hash = compute_file_hash(txt_content) | |
| # Check for duplicates with the new hash | |
| if file_hash in FILE_REGISTRY[database]: | |
| logger.warning(f"⚠️ Duplicate detected after text extraction (ZIP): {original_filename}") | |
| os.unlink(temp_pdf_path) | |
| os.unlink(pdf_path) | |
| try: | |
| os.unlink(text_path) | |
| except Exception: | |
| pass | |
| return None, original_filename | |
| # Extract and save metadata (including title from original PDF and PDF path) | |
| file_metadata = extract_file_metadata(temp_pdf_path, original_filename, "pdf", file_hash) | |
| file_metadata["pdf_path"] = pdf_filename | |
| upsert_metadata(database, {file_hash: file_metadata}) | |
| # Register the text file in FILE_REGISTRY | |
| FILE_REGISTRY[database][file_hash] = text_filename | |
| logger.info(f"📄 Extracted and saved text (ZIP): {text_filename} in {database}") | |
| logger.info(f"📄 Also saved original PDF (ZIP): {pdf_filename}") | |
| os.unlink(temp_pdf_path) | |
| return text_path, None | |
| except Exception as e: | |
| logger.error(f"❌ Error processing PDF from ZIP {original_filename}: {str(e)}") | |
| try: | |
| os.unlink(temp_pdf_path) | |
| except Exception: | |
| pass | |
| return None, f"Error processing PDF: {str(e)}" | |
| elif file_ext in image_extensions: | |
| # For images: Save the image file as-is | |
| filename = f"{file_id}_{original_filename}" | |
| file_path = os.path.join(category_upload_dir, filename) | |
| logger.info(f"Saving image (ZIP) to: {os.path.abspath(file_path)}") | |
| with open(file_path, "wb") as f: | |
| f.write(content) | |
| # Extract and save metadata for image | |
| file_hash = compute_file_hash(content) | |
| file_metadata = extract_file_metadata(file_path, original_filename, "image", file_hash) | |
| upsert_metadata(database, {file_hash: file_metadata}) | |
| FILE_REGISTRY[database][file_hash] = filename | |
| logger.info(f"📸 Saved image (ZIP): {filename} in {database}") | |
| return file_path, None | |
| else: | |
| # For text files: Save as text with UTF-8 encoding | |
| filename = f"{file_id}_{original_filename}" | |
| file_path = os.path.join(category_upload_dir, filename) | |
| logger.info(f"Saving text file (ZIP) to: {os.path.abspath(file_path)}") | |
| with open(file_path, "w", encoding="utf-8") as f: | |
| f.write(content.decode('utf-8')) | |
| # Extract and save metadata for text file | |
| file_hash = compute_file_hash(content) | |
| file_metadata = extract_file_metadata(file_path, original_filename, "text", file_hash) | |
| upsert_metadata(database, {file_hash: file_metadata}) | |
| FILE_REGISTRY[database][file_hash] = filename | |
| logger.info(f"📝 Saved text file (ZIP): {filename} in {database}") | |
| return file_path, None | |
| def obtain_text_from_files(database: str): | |
| """ | |
| Returns a dict of title -> full file text for all text files in the given database. | |
| Now that PDFs are converted to .txt files during upload, we only need to read the text files. | |
| Uses metadata to get proper titles including original PDF titles. | |
| """ | |
| file_dict = {} | |
| if database not in FILE_REGISTRY or not FILE_REGISTRY[database]: | |
| logger.warning(f"Database '{database}' not found or empty; returning no papers.") | |
| return {} | |
| # Create path to category-specific directory | |
| category_upload_dir = os.path.join(UPLOAD_DIR, database) | |
| for file_hash, filename in FILE_REGISTRY[database].items(): | |
| file_path = os.path.join(category_upload_dir, filename) | |
| # Skip image files - we only want text content | |
| if is_image_file(filename): | |
| continue | |
| # Since PDFs are now converted to .txt files during upload, | |
| # we only need to read text files directly | |
| if filename.lower().endswith(('.txt', '.text')): | |
| try: | |
| # Get title from metadata (preserves original PDF titles) | |
| file_title = get_file_title_from_metadata(database, file_hash) | |
| if not file_title or file_title == "Unknown": | |
| # Fallback to extracting title from file content | |
| logger.info(f"🔍 Metadata lookup failed for {filename} (hash: {file_hash[:16]}...), using fallback title extraction") | |
| from pdf_text_processing import get_file_title | |
| temp_title = get_file_title(file_path) | |
| file_title = temp_title if temp_title else filename | |
| logger.info(f"📝 Fallback title for {filename}: {file_title}") | |
| else: | |
| logger.info(f"✅ Using title from metadata for {filename}: {file_title}") | |
| # Read the text file content | |
| with open(file_path, "r", encoding="utf-8") as f: | |
| file_text = f.read() | |
| file_dict[file_title] = file_text | |
| except Exception as e: | |
| logger.error(f"❌ Error reading file {filename}: {str(e)}") | |
| continue | |
| return file_dict | |
| def obtain_text_from_pdf(database: str): | |
| """ | |
| Legacy function name for backward compatibility. | |
| Now handles both PDF and text files. | |
| """ | |
| return obtain_text_from_files(database) | |
| def extract_all_files(file_paths: List[str]) -> List[Document]: | |
| """ | |
| Runs text extraction + chunk-splitting on each file (text files only, since PDFs are now converted to .txt during upload). | |
| """ | |
| all_docs = [] | |
| with concurrent.futures.ThreadPoolExecutor( | |
| max_workers=min(32, os.cpu_count()) | |
| ) as executor: | |
| futures = [executor.submit(extract_and_split, path) for path in file_paths] | |
| for future in concurrent.futures.as_completed(futures): | |
| result = future.result() | |
| if result: | |
| all_docs.extend(result) | |
| return all_docs | |
| def save_to_vectorstore(embeddings, texts, metadatas, vectorstore): | |
| """ | |
| Save a batch of (embeddings, texts, metadatas) to the vectorstore in chunks. | |
| """ | |
| try: | |
| batch_size = 4000 # tune as needed (e.g., 500-2000) | |
| n = len(texts) | |
| for i in range(0, n, batch_size): | |
| j = min(i + batch_size, n) | |
| vectorstore.add_texts( | |
| texts=texts[i:j], | |
| metadatas=metadatas[i:j], | |
| embeddings=embeddings[i:j], | |
| ) | |
| vectorstore.persist() | |
| logger.info(f"✅ Vectorstore saved: {n} chunks (in batches of {batch_size})") | |
| except Exception as e: | |
| logger.error(f"❌ Vectorstore save error: {e}") | |
| def extract_texts_and_metas(docs: List[Document]): | |
| texts = [doc.page_content for doc in docs] | |
| metas = [doc.metadata for doc in docs] | |
| return texts, metas | |
| # -------------------- Progress tracking endpoints -------------------- # | |
| async def get_embedding_progress(database: str, ws: str = Depends(workspace)): | |
| """ | |
| Returns the current progress of embedding tasks for a specific database. | |
| """ | |
| key = scoped(ws, database) | |
| if key not in EMBEDDING_PROGRESS: | |
| print(f"❌ No task found for {database}") | |
| return JSONResponse(content={"status": "no_task", "progress": 0, "message": "No embedding task in progress"}) | |
| progress_data = EMBEDDING_PROGRESS[key].copy() | |
| # Remove zip_data from response to avoid sending large binary data | |
| if "zip_data" in progress_data: | |
| del progress_data["zip_data"] | |
| return JSONResponse(content=progress_data) | |
| async def download_extraction(extraction_id: str, ws: str = Depends(workspace)): | |
| """ | |
| Download the system chart extraction ZIP file directly from memory. | |
| """ | |
| progress_data = EMBEDDING_PROGRESS.get(scoped(ws, extraction_id)) | |
| if progress_data is None: | |
| return JSONResponse(content={"error": "Extraction not found"}, status_code=404) | |
| if progress_data.get("status") != "completed": | |
| return JSONResponse(content={"error": "Extraction not completed yet"}, status_code=400) | |
| zip_data = progress_data.get("zip_data") | |
| if not zip_data: | |
| return JSONResponse(content={"error": "Extraction data not available"}, status_code=404) | |
| filename = progress_data.get("filename", "system_chart_extraction.zip") | |
| # Return the ZIP file as a direct download | |
| from fastapi.responses import Response | |
| return Response( | |
| content=zip_data, | |
| media_type="application/zip", | |
| headers={ | |
| "Content-Disposition": f"attachment; filename={filename}", | |
| "Content-Length": str(len(zip_data)), | |
| "Content-Type": "application/zip", | |
| "Cache-Control": "no-cache, no-store, must-revalidate", | |
| "Pragma": "no-cache", | |
| "Expires": "0" | |
| } | |
| ) | |
| async def process_files_async_all(file_paths: List[str], database: str, embedding_model): | |
| """ | |
| Process all files with parallel embedding for improved performance. | |
| Uses the new embed_stream_into_vs function for concurrent batch processing. | |
| Note: PDFs are now converted to .txt files during upload, so we only process text files. | |
| """ | |
| global EMBEDDING_PROGRESS | |
| if embedding_model is None: | |
| logger.error("❌ Embedding model is None. Did you forget /set_api_key?") | |
| return | |
| total_files = len(file_paths) | |
| logger.info(f"🚀 Starting parallel processing of {total_files} files for database '{database}'") | |
| # Initialize progress tracking | |
| EMBEDDING_PROGRESS[database] = { | |
| "status": "extracting", | |
| "progress": 0, | |
| "total_files": total_files, | |
| "processed_files": 0, | |
| "message": f"Extracting text from {total_files} files..." | |
| } | |
| try: | |
| # Extract text from all files | |
| vectorstore = get_vectorstore(database, embedding_model) | |
| all_docs = extract_all_files(file_paths) | |
| if not all_docs: | |
| EMBEDDING_PROGRESS[database] = { | |
| "status": "completed", | |
| "progress": 100, | |
| "message": "No documents found to embed" | |
| } | |
| logger.warning(f"⚠️ No documents extracted from {total_files} files") | |
| return | |
| # Extract texts and metadata | |
| texts, metas = extract_texts_and_metas(all_docs) | |
| total_chunks = len(texts) | |
| logger.info(f"📄 Extracted {total_chunks} chunks from {total_files} files") | |
| # Update progress after extraction | |
| EMBEDDING_PROGRESS[database] = { | |
| "status": "embedding", | |
| "progress": 10, | |
| "total_chunks": total_chunks, | |
| "processed_chunks": 0, | |
| "message": f"Starting parallel embedding of {total_chunks} chunks..." | |
| } | |
| # Use parallel embedding with progress tracking | |
| await embed_stream_into_vs( | |
| texts, metas, vectorstore, embedding_model, | |
| batch_size=512, | |
| max_concurrency=12, | |
| progress_key=database | |
| ) | |
| # Final completion status | |
| EMBEDDING_PROGRESS[database] = { | |
| "status": "completed", | |
| "progress": 100, | |
| "message": f"Successfully embedded {total_chunks} chunks from {total_files} files using parallel processing" | |
| } | |
| logger.info(f"✅ Completed parallel processing: {total_chunks} chunks from {total_files} files") | |
| except asyncio.CancelledError: | |
| EMBEDDING_PROGRESS[database] = { | |
| "status": "error", | |
| "progress": EMBEDDING_PROGRESS[database].get("progress", 0), | |
| "message": "Embedding process was cancelled" | |
| } | |
| logger.warning(f"⚠️ Embedding cancelled for database '{database}'") | |
| raise | |
| except Exception as e: | |
| current_progress = EMBEDDING_PROGRESS[database].get("progress", 0) | |
| EMBEDDING_PROGRESS[database] = { | |
| "status": "error", | |
| "progress": current_progress, | |
| "message": f"Error during parallel embedding: {str(e)}" | |
| } | |
| logger.error(f"❌ Embedding failed for database '{database}': {str(e)}") | |
| raise | |
| # Removed per-file embedding mode; only all-at-once embedding is supported. | |
| async def process_system_chart_extraction_async(category: str, llm, extraction_id: str, enable_visualization: bool = False, run_system_chart_normalization: bool = False): | |
| """ | |
| Background task for system chart extraction with progress tracking. | |
| Uses the SystemChartService for business logic while handling API-specific concerns. | |
| """ | |
| global EMBEDDING_PROGRESS | |
| print(f"🚀 System Chart Background task started for extraction_id: {extraction_id}") | |
| print(f"📊 Visualization enabled: {enable_visualization}") | |
| # Initialize progress tracking | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "starting", | |
| "progress": 0, | |
| "message": "Initializing system chart extraction..." | |
| } | |
| try: | |
| # Create system chart service instance | |
| system_chart_service = SystemChartService() | |
| # Get PDF text dictionary | |
| pdf_dict = obtain_text_from_pdf(category) | |
| if not pdf_dict: | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "error", | |
| "progress": 0, | |
| "message": "No files found in the specified category." | |
| } | |
| return | |
| # Define progress callback for API-specific progress tracking | |
| def progress_callback(progress_data): | |
| print(f"📊 Progress update: {progress_data}") | |
| EMBEDDING_PROGRESS[extraction_id] = progress_data | |
| print(f"🔄 EMBEDDING_PROGRESS[{extraction_id}] updated to: {EMBEDDING_PROGRESS[extraction_id]}") | |
| print(f"📋 Current keys in EMBEDDING_PROGRESS: {list(EMBEDDING_PROGRESS.keys())}") | |
| # Perform system chart extraction using the service | |
| run_norm_effective = bool(enable_visualization or run_system_chart_normalization) | |
| result = await system_chart_service.extract_system_charts( | |
| category=category, | |
| llm=llm, | |
| pdf_dict=pdf_dict, | |
| enable_visualization=enable_visualization, | |
| run_system_chart_normalization=run_norm_effective, | |
| progress_callback=progress_callback, | |
| extraction_id=extraction_id # Pass extraction_id to save conversations | |
| ) | |
| # Handle the result | |
| if result["success"]: | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "completed", | |
| "progress": 100, | |
| "message": result["message"], | |
| "zip_data": result["zip_data"], | |
| "filename": result["filename"] | |
| } | |
| else: | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "error", | |
| "progress": 0, | |
| "message": result["message"] | |
| } | |
| except Exception as e: | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "error", | |
| "progress": 0, | |
| "message": f"Error during system chart extraction: {str(e)}" | |
| } | |
| logger.error(f"System chart extraction error: {str(e)}") | |
| async def process_multi_database_system_chart_extraction_async(ws: str, categories: List[str], llm, extraction_id: str, enable_visualization: bool = False, run_system_chart_normalization: bool = False): | |
| """ | |
| Background task for system chart extraction across multiple databases with progress tracking. | |
| Combines text from all databases and runs system chart extraction on the combined dataset. | |
| """ | |
| global EMBEDDING_PROGRESS | |
| print(f"🚀 Multi-Database System Chart Background task started for extraction_id: {extraction_id}") | |
| print(f"📊 Processing {len(categories)} databases: {categories}") | |
| print(f"📊 Visualization enabled: {enable_visualization}") | |
| # Initialize progress tracking | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "gathering_data", | |
| "progress": 5, | |
| "message": f"Gathering text from {len(categories)} databases...", | |
| "databases": categories, | |
| "current_database": None, | |
| "databases_processed": 0, | |
| "total_databases": len(categories) | |
| } | |
| try: | |
| # Create system chart service instance | |
| system_chart_service = SystemChartService() | |
| # Gather text from all databases | |
| combined_pdf_dict = {} | |
| databases_processed = 0 | |
| for i, category in enumerate(categories): | |
| EMBEDDING_PROGRESS[extraction_id].update({ | |
| "progress": 5 + (i * 20 // len(categories)), # 5-25% for data gathering | |
| "message": f"Gathering text from database: {category}", | |
| "current_database": category, | |
| "databases_processed": i | |
| }) | |
| print(f"📁 Processing database {i+1}/{len(categories)}: {category}") | |
| # Get PDF text dictionary for this database | |
| pdf_dict = obtain_text_from_pdf(scoped(ws, category)) | |
| if pdf_dict: | |
| # Add database prefix to avoid title conflicts | |
| for title, content in pdf_dict.items(): | |
| prefixed_title = f"[{category}] {title}" | |
| combined_pdf_dict[prefixed_title] = content | |
| print(f"✅ Added {len(pdf_dict)} documents from {category}") | |
| else: | |
| print(f"⚠️ No files found in database: {category}") | |
| databases_processed += 1 | |
| # Check if we have any documents to process | |
| if not combined_pdf_dict: | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "error", | |
| "progress": 0, | |
| "message": f"No files found in any of the selected databases: {', '.join(categories)}" | |
| } | |
| return | |
| print(f"📚 Combined {len(combined_pdf_dict)} documents from {len(categories)} databases") | |
| # Update progress for starting extraction | |
| EMBEDDING_PROGRESS[extraction_id].update({ | |
| "status": "extracting", | |
| "progress": 30, | |
| "message": f"Starting system chart extraction on {len(combined_pdf_dict)} documents...", | |
| "current_database": "combined", | |
| "databases_processed": len(categories), | |
| "total_documents": len(combined_pdf_dict) | |
| }) | |
| # Define progress callback for API-specific progress tracking | |
| def progress_callback(progress_data): | |
| # Adjust progress to account for the 30% already used for data gathering | |
| adjusted_progress = 30 + (progress_data.get("progress", 0) * 0.7) # 30-100% | |
| # Merge with our multi-database tracking data | |
| updated_data = { | |
| "status": progress_data.get("status", "processing"), | |
| "progress": int(adjusted_progress), | |
| "message": progress_data.get("message", "Processing..."), | |
| "databases": categories, | |
| "databases_processed": len(categories), | |
| "total_databases": len(categories), | |
| "total_documents": len(combined_pdf_dict) | |
| } | |
| # Add specific progress details if available | |
| for key in ["current_paper", "total_papers", "current_question", "total_questions", | |
| "completed_questions", "current_paper_type", "questions_in_paper"]: | |
| if key in progress_data: | |
| updated_data[key] = progress_data[key] | |
| print(f"📊 Multi-DB Progress update: {updated_data}") | |
| EMBEDDING_PROGRESS[extraction_id] = updated_data | |
| run_norm_effective = bool(enable_visualization or run_system_chart_normalization) | |
| # Perform system chart extraction using the service on combined data | |
| # Use a combined category name for the service | |
| combined_category_name = "_".join(categories) | |
| result = await system_chart_service.extract_system_charts( | |
| category=combined_category_name, | |
| llm=llm, | |
| pdf_dict=combined_pdf_dict, | |
| enable_visualization=enable_visualization, | |
| run_system_chart_normalization=run_norm_effective, | |
| progress_callback=progress_callback, | |
| extraction_id=extraction_id # Pass extraction_id to save conversations | |
| ) | |
| # Handle the result | |
| if result["success"]: | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "completed", | |
| "progress": 100, | |
| "message": f"System chart extraction completed for {len(categories)} databases ({len(combined_pdf_dict)} documents)", | |
| "zip_data": result["zip_data"], | |
| "filename": result["filename"], | |
| "databases": categories, | |
| "databases_processed": len(categories), | |
| "total_databases": len(categories), | |
| "total_documents": len(combined_pdf_dict) | |
| } | |
| else: | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "error", | |
| "progress": 0, | |
| "message": f"Error during extraction: {result['message']}", | |
| "databases": categories, | |
| "total_documents": len(combined_pdf_dict) | |
| } | |
| except Exception as e: | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "error", | |
| "progress": 0, | |
| "message": f"Error during multi-database system chart extraction: {str(e)}", | |
| "databases": categories | |
| } | |
| logger.error(f"Multi-database system chart extraction error: {str(e)}") | |
| async def process_hypotheses_generation_async(system_charts_content_list: List[str], llm, extraction_id: str, material_goal_1: str, material_goal_2: str, iterations_count: int = 1): | |
| """ | |
| Background task for synergistic hypotheses generation with progress tracking. | |
| Uses the HypothesesService for business logic while handling API-specific concerns. | |
| """ | |
| global EMBEDDING_PROGRESS, HYPOTHESIS_GENERATION_CONTEXT | |
| print(f"🚀 Hypotheses Generation Background task started for extraction_id: {extraction_id}") | |
| # Store context for continuous conversation | |
| HYPOTHESIS_GENERATION_CONTEXT[extraction_id] = { | |
| 'system_charts_content_list': system_charts_content_list, | |
| 'material_goal_1': material_goal_1, | |
| 'material_goal_2': material_goal_2, | |
| 'iterations_count': iterations_count | |
| } | |
| # Initialize progress tracking | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "starting", | |
| "progress": 0, | |
| "message": "Initializing hypothesis generation..." | |
| } | |
| try: | |
| # Create hypotheses service instance | |
| hypotheses_service = HypothesesService() | |
| # Validate the system charts content list first | |
| validation = HypothesesService.validate_system_charts_content(system_charts_content_list) | |
| if not validation["valid"]: | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "error", | |
| "progress": 0, | |
| "message": validation["message"] | |
| } | |
| return | |
| print(f"📊 Validated {validation['valid_charts']} system charts from {len(system_charts_content_list)} files for hypothesis generation") | |
| # Define progress callback for API-specific progress tracking | |
| def progress_callback(progress_data): | |
| EMBEDDING_PROGRESS[extraction_id] = progress_data | |
| # Perform hypotheses generation using the service | |
| result = await hypotheses_service.generate_hypotheses( | |
| system_charts_content_list=system_charts_content_list, # Pass the list | |
| llm=llm, | |
| material_goal_1=material_goal_1, | |
| material_goal_2=material_goal_2, | |
| iterations_count=iterations_count, | |
| progress_callback=progress_callback | |
| ) | |
| # Handle the result | |
| if result["success"]: | |
| HYPOTHESIS_GENERATION_CONTEXT[extraction_id]['hypotheses'] = result["hypotheses"] | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "completed", | |
| "progress": 100, | |
| "message": result["message"], | |
| "zip_data": result["zip_data"], | |
| "filename": result["filename"], | |
| "iterations_count": iterations_count, | |
| "hypotheses": result["hypotheses"] | |
| } | |
| else: | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "error", | |
| "progress": 0, | |
| "message": result["message"] | |
| } | |
| except Exception as e: | |
| EMBEDDING_PROGRESS[extraction_id] = { | |
| "status": "error", | |
| "progress": 0, | |
| "message": f"Error during hypothesis generation: {str(e)}" | |
| } | |
| logger.error(f"Hypothesis generation error: {str(e)}") | |
| # ------------ NEW: background QA with chunk-level progress ------------ # | |
| async def process_qa_async( | |
| question: str, | |
| categories: List[str], | |
| k: int, | |
| embedding_model, | |
| llm_QA, | |
| progress_id: str, | |
| ): | |
| """ | |
| End-to-end QA answer generation that updates EMBEDDING_PROGRESS | |
| after each chunk is summarised. | |
| """ | |
| try: | |
| # ---- gather relevant chunks ------------------------------------------------- | |
| all_docs = [] | |
| for cat in categories: | |
| vs = get_vectorstore(cat, embedding_model) | |
| q_emb = await embedding_model.aembed_documents([question]) | |
| k_cat = k # Use k chunks from each database (not divided across databases) | |
| res = vs._collection.query(query_embeddings=q_emb, n_results=k_cat) | |
| for doc_txt, meta in zip(res["documents"][0], res["metadatas"][0]): | |
| meta["database"] = cat | |
| all_docs.append(Document(page_content=doc_txt, metadata=meta)) | |
| total_chunks = len(all_docs) | |
| if not total_chunks: | |
| EMBEDDING_PROGRESS[progress_id] = { | |
| "status": "error", | |
| "progress": 0, | |
| "message": "No relevant chunks found." | |
| } | |
| return | |
| EMBEDDING_PROGRESS[progress_id] = { | |
| "status": "processing", | |
| "progress": 0, | |
| "total_chunks": total_chunks, | |
| "processed_chunks": 0, | |
| "answer": None | |
| } | |
| # ---- chunk-wise map --------------------------------------------------------- | |
| async def summarise(idx, doc): | |
| from langchain_core.messages import HumanMessage | |
| response = await llm_QA.ainvoke( | |
| [HumanMessage(content=f"Summarise this in relation to '{question}':\n{doc.page_content}")] | |
| ) | |
| summary = llm_content_to_text(getattr(response, "content", response)) | |
| # update progress *after* the chunk is done | |
| done = EMBEDDING_PROGRESS[progress_id]["processed_chunks"] + 1 | |
| EMBEDDING_PROGRESS[progress_id]["processed_chunks"] = done | |
| EMBEDDING_PROGRESS[progress_id]["progress"] = int(done / total_chunks * 90) # keep 10 % for reduce | |
| return summary | |
| summaries = await asyncio.gather( | |
| *[summarise(i, d) for i, d in enumerate(all_docs)] | |
| ) | |
| # ---- reduce ----------------------------------------------------------------- | |
| from langchain_core.messages import HumanMessage | |
| final_response = await llm_QA.ainvoke( | |
| [HumanMessage(content="Combine the following partial answers into one:\n" + "\n".join(summaries))] | |
| ) | |
| final = llm_content_to_text(getattr(final_response, "content", final_response)) | |
| EMBEDDING_PROGRESS[progress_id] = { | |
| "status": "completed", | |
| "progress": 100, | |
| "answer": final | |
| } | |
| except Exception as e: | |
| EMBEDDING_PROGRESS[progress_id] = { | |
| "status": "error", | |
| "progress": 0, | |
| "message": f"QA failed: {e}" | |
| } | |
| async def upload_files( | |
| background_tasks: BackgroundTasks, | |
| category: str = Form(...), | |
| files: List[UploadFile] = File(...), | |
| ws: str = Depends(workspace), | |
| ): | |
| """ | |
| Upload PDF, text, and image files to the specified category ('database'). | |
| PDFs are converted to .txt files (original PDF is not saved to save storage space). | |
| Text files are saved as-is. Image files are saved as-is. | |
| Background embedding tasks are started for text files only. | |
| Supported file types: .pdf, .txt, .text, .png, .jpg, .jpeg, .gif, .bmp, .tiff, .webp | |
| """ | |
| global FILE_REGISTRY | |
| embedding_model = WORKSPACE_MODELS.get(ws, {}).get("embedding") | |
| # Basic check to ensure we have an embedding_model: | |
| if embedding_model is None: | |
| return JSONResponse( | |
| content={ | |
| "message": "No embedding model set. Call /set_api_key first." | |
| }, | |
| status_code=400 | |
| ) | |
| display_category, category = category, scoped(ws, category) | |
| if category not in FILE_REGISTRY: | |
| FILE_REGISTRY[category] = {} | |
| # PROTECTION: Check if database is locked before processing files | |
| category_upload_dir = os.path.join(UPLOAD_DIR, category) | |
| metadata_file = os.path.join(category_upload_dir, "METADATA") | |
| if os.path.exists(metadata_file): | |
| return JSONResponse( | |
| content={ | |
| "message": f"🔒 Database '{display_category}' is locked. Cannot add files to databases with existing METADATA files. Once a database is created and has a METADATA file, it becomes immutable to prevent data corruption." | |
| }, | |
| status_code=400 | |
| ) | |
| # Auto-detect single ZIP upload and route through ZIP logic | |
| if len(files) == 1 and files[0].filename.lower().endswith('.zip'): | |
| # Reuse the ZIP handler logic inline to keep response identical to /upload_zip | |
| import io | |
| import zipfile | |
| zip_bytes = await files[0].read() | |
| zf = zipfile.ZipFile(io.BytesIO(zip_bytes)) | |
| saved_file_paths, duplicates = [], [] | |
| uploaded_files_info = [] | |
| for info in zf.infolist(): | |
| if info.is_dir(): | |
| continue | |
| inner_name = info.filename | |
| base_name = os.path.basename(inner_name) | |
| if not base_name: | |
| continue | |
| try: | |
| data = zf.read(info) | |
| except Exception as e: | |
| logger.warning(f"⚠️ Skipping '{inner_name}' due to read error: {e}") | |
| continue | |
| path, dup = save_file_from_bytes(data, base_name, category) | |
| if dup: | |
| duplicates.append(dup) | |
| continue | |
| if path: | |
| saved_file_paths.append(path) | |
| original_ext = os.path.splitext(base_name.lower())[1] | |
| if original_ext == '.pdf' or original_ext in ['.txt', '.text']: | |
| file_type = "text_pdf" | |
| elif original_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']: | |
| file_type = "image" | |
| else: | |
| file_type = "unknown" | |
| uploaded_files_info.append({ | |
| "path": path, | |
| "filename": os.path.basename(path), | |
| "original_filename": base_name, | |
| "type": file_type | |
| }) | |
| # Separate files by type | |
| text_pdf_files = [info["path"] for info in uploaded_files_info if info["type"] == "text_pdf"] | |
| image_files = [info["path"] for info in uploaded_files_info if info["type"] == "image"] | |
| # Kick off background tasks for embedding (only for PDF/text files) | |
| if text_pdf_files and embedding_model is not None: | |
| start_time = time.perf_counter() | |
| task = asyncio.create_task(process_files_async_all(text_pdf_files, category, embedding_model)) | |
| task.add_done_callback(lambda _, s = start_time: print(f"Embedding time in {(time.perf_counter() - s) * 1000:.2f} ms")) | |
| return JSONResponse(content={ | |
| "message": "✅ ZIP upload complete.", | |
| "uploaded": [info["filename"] for info in uploaded_files_info], | |
| "duplicates": duplicates, | |
| "file_types": { | |
| "text_pdf_count": len(text_pdf_files), | |
| "image_count": len(image_files), | |
| "has_images": len(image_files) > 0, | |
| "has_text_pdf": len(text_pdf_files) > 0 | |
| }, | |
| "embedding_started": len(text_pdf_files) > 0 and embedding_model is not None | |
| }) | |
| tasks = [(save_file(file, category), file.filename) for file in files] | |
| results = await asyncio.gather(*[task[0] for task in tasks]) | |
| saved_file_paths, duplicates = [], [] | |
| uploaded_files_info = [] | |
| for (path, dup), (_, original_filename) in zip(results, tasks): | |
| if dup: | |
| duplicates.append(dup) | |
| elif path: | |
| saved_file_paths.append(path) | |
| # Determine file type based on original filename, not the saved path | |
| # (since PDFs are now converted to .txt files) | |
| original_ext = os.path.splitext(original_filename.lower())[1] | |
| if original_ext == '.pdf' or original_ext in ['.txt', '.text']: | |
| file_type = "text_pdf" | |
| elif original_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']: | |
| file_type = "image" | |
| else: | |
| file_type = "unknown" | |
| uploaded_files_info.append({ | |
| "path": path, | |
| "filename": os.path.basename(path), | |
| "original_filename": original_filename, | |
| "type": file_type | |
| }) | |
| # Separate files by type | |
| text_pdf_files = [info["path"] for info in uploaded_files_info if info["type"] == "text_pdf"] | |
| image_files = [info["path"] for info in uploaded_files_info if info["type"] == "image"] | |
| # Kick off background tasks for embedding (only for PDF/text files) | |
| if text_pdf_files and embedding_model is not None: | |
| start_time = time.perf_counter() | |
| task = asyncio.create_task(process_files_async_all(text_pdf_files, category, embedding_model)) | |
| task.add_done_callback(lambda _, s = start_time: print(f"Embedding time in {(time.perf_counter() - s) * 1000:.2f} ms")) | |
| # Prepare response with file type information | |
| return JSONResponse(content={ | |
| "message": "Upload complete.", | |
| "uploaded": [info["filename"] for info in uploaded_files_info], | |
| "duplicates": duplicates, | |
| "file_types": { | |
| "text_pdf_count": len(text_pdf_files), | |
| "image_count": len(image_files), | |
| "has_images": len(image_files) > 0, | |
| "has_text_pdf": len(text_pdf_files) > 0 | |
| }, | |
| "embedding_started": len(text_pdf_files) > 0 and embedding_model is not None, | |
| "note": "PDFs are converted to text files during upload to save storage space. Only the extracted text is saved." | |
| }) | |
| async def upload_zip( | |
| background_tasks: BackgroundTasks, | |
| category: str = Form(...), | |
| zip_file: UploadFile = File(...), | |
| ws: str = Depends(workspace), | |
| ): | |
| """ | |
| Upload a single ZIP file containing PDFs, text files, and/or images. | |
| Files are extracted server-side and saved using the same rules as /upload. | |
| Then, background embedding is kicked off for text/PDF-derived files. | |
| """ | |
| global FILE_REGISTRY | |
| embedding_model = WORKSPACE_MODELS.get(ws, {}).get("embedding") | |
| if embedding_model is None: | |
| return JSONResponse( | |
| content={"message": "No embedding model set. Please call /set_api_key first."}, | |
| status_code=400 | |
| ) | |
| display_category, category = category, scoped(ws, category) | |
| if category not in FILE_REGISTRY: | |
| FILE_REGISTRY[category] = {} | |
| # PROTECTION: Check if database is locked before processing files | |
| category_upload_dir = os.path.join(UPLOAD_DIR, category) | |
| metadata_file = os.path.join(category_upload_dir, "METADATA") | |
| if os.path.exists(metadata_file): | |
| return JSONResponse( | |
| content={ | |
| "message": f"🔒 Database '{display_category}' is locked. Cannot add files to databases with existing METADATA files. Once a database is created and has a METADATA file, it becomes immutable to prevent data corruption." | |
| }, | |
| status_code=400 | |
| ) | |
| import io | |
| import zipfile | |
| # Read ZIP into memory (could be optimized to temp file for very large zips) | |
| zip_bytes = await zip_file.read() | |
| zf = zipfile.ZipFile(io.BytesIO(zip_bytes)) | |
| saved_file_paths, duplicates = [], [] | |
| uploaded_files_info = [] | |
| for info in zf.infolist(): | |
| if info.is_dir(): | |
| continue | |
| inner_name = info.filename | |
| base_name = os.path.basename(inner_name) | |
| if not base_name: # skip entries without a filename | |
| continue | |
| try: | |
| data = zf.read(info) | |
| except Exception as e: | |
| logger.warning(f"⚠️ Skipping '{inner_name}' due to read error: {e}") | |
| continue | |
| path, dup = save_file_from_bytes(data, base_name, category) | |
| if dup: | |
| duplicates.append(dup) | |
| continue | |
| if path: | |
| saved_file_paths.append(path) | |
| original_ext = os.path.splitext(base_name.lower())[1] | |
| if original_ext == '.pdf' or original_ext in ['.txt', '.text']: | |
| file_type = "text_pdf" | |
| elif original_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp']: | |
| file_type = "image" | |
| else: | |
| file_type = "unknown" | |
| uploaded_files_info.append({ | |
| "path": path, | |
| "filename": os.path.basename(path), | |
| "original_filename": base_name, | |
| "type": file_type | |
| }) | |
| # Separate files by type | |
| text_pdf_files = [info["path"] for info in uploaded_files_info if info["type"] == "text_pdf"] | |
| image_files = [info["path"] for info in uploaded_files_info if info["type"] == "image"] | |
| # Kick off background tasks for embedding (only for PDF/text files) | |
| if text_pdf_files and embedding_model is not None: | |
| start_time = time.perf_counter() | |
| task = asyncio.create_task(process_files_async_all(text_pdf_files, category, embedding_model)) | |
| task.add_done_callback(lambda _, s = start_time: print(f"Embedding time in {(time.perf_counter() - s) * 1000:.2f} ms")) | |
| return JSONResponse(content={ | |
| "message": "✅ ZIP upload complete.", | |
| "uploaded": [info["filename"] for info in uploaded_files_info], | |
| "duplicates": duplicates, | |
| "file_types": { | |
| "text_pdf_count": len(text_pdf_files), | |
| "image_count": len(image_files), | |
| "has_images": len(image_files) > 0, | |
| "has_text_pdf": len(text_pdf_files) > 0 | |
| }, | |
| "embedding_started": len(text_pdf_files) > 0 and embedding_model is not None | |
| }) | |
| async def get_available_databases(ws: str = Depends(workspace)): | |
| """ | |
| Returns this workspace's databases/categories with their file information. | |
| """ | |
| global FILE_REGISTRY | |
| databases = [] | |
| for category, files in FILE_REGISTRY.items(): | |
| name = unscoped(ws, category) | |
| if name is not None and files: # Only this workspace's categories that have files | |
| # Count file types in the category | |
| text_pdf_count = 0 | |
| image_count = 0 | |
| pdf_count = 0 | |
| # Load metadata to check original file types (especially for converted PDFs) | |
| metadata_dict = load_metadata(category) | |
| for file_hash, filename in files.items(): | |
| file_type = get_file_type(filename) | |
| if file_type == "text_pdf": | |
| text_pdf_count += 1 | |
| elif file_type == "image": | |
| image_count += 1 | |
| # Check for PDFs using metadata (for converted PDFs) or filename (for direct PDFs) | |
| if filename.lower().endswith('.pdf'): | |
| pdf_count += 1 | |
| elif file_hash in metadata_dict: | |
| original_filename = metadata_dict[file_hash].get('original_filename', '') | |
| if original_filename.lower().endswith('.pdf'): | |
| pdf_count += 1 | |
| # Determine available modes | |
| available_modes = [] | |
| if text_pdf_count > 0: | |
| available_modes.extend(["QA", "System Chart Extraction"]) | |
| databases.append({ | |
| "name": name, | |
| "file_count": len(files), | |
| "text_pdf_count": text_pdf_count, | |
| "image_count": image_count, | |
| "pdf_count": pdf_count, | |
| "available_modes": available_modes, | |
| "file_types": { | |
| "has_images": image_count > 0, | |
| "has_text_pdf": text_pdf_count > 0, | |
| "has_pdfs": pdf_count > 0 | |
| } | |
| }) | |
| return JSONResponse(content={"databases": databases}) | |
| # -------------------- Conversation Management Endpoints -------------------- # | |
| async def get_every_conversation(ws: str = Depends(workspace)): | |
| """ | |
| Get conversations across all of this workspace's extractions (and therefore all | |
| its databases), so the frontend dropdown is not limited to the most recent run. | |
| """ | |
| all_convs = get_all_conversations() | |
| conversation_list = [] | |
| for key, conversations in all_convs.items(): | |
| ext_id = unscoped(ws, key) | |
| if ext_id is None: | |
| continue | |
| for conv_key, conv_data in conversations.items(): | |
| conversation_list.append({ | |
| "extraction_id": ext_id, | |
| "key": conv_key, | |
| "title": conv_data["paper_title"], | |
| "paper_type": conv_data["paper_type"], | |
| "created_at": conv_data["created_at"], | |
| "last_updated": conv_data["last_updated"], | |
| "message_count": len(conv_data["conversation_history"]) | |
| }) | |
| conversation_list.sort(key=lambda x: x["created_at"], reverse=True) | |
| return JSONResponse(content={ | |
| "conversations": conversation_list, | |
| "total_count": len(conversation_list) | |
| }) | |
| async def get_extraction_conversations(extraction_id: str, ws: str = Depends(workspace)): | |
| """ | |
| Get all conversations for a given extraction ID. | |
| Returns a list of conversation metadata for frontend dropdown. | |
| """ | |
| conversations = get_conversations(scoped(ws, extraction_id)) | |
| if not conversations: | |
| return JSONResponse(content={ | |
| "conversations": [], | |
| "message": "No conversations found for this extraction ID" | |
| }) | |
| # Format conversations for frontend dropdown | |
| conversation_list = [] | |
| for conv_key, conv_data in conversations.items(): | |
| conversation_list.append({ | |
| "key": conv_key, | |
| "title": conv_data["paper_title"], | |
| "paper_type": conv_data["paper_type"], | |
| "created_at": conv_data["created_at"], | |
| "last_updated": conv_data["last_updated"], | |
| "message_count": len(conv_data["conversation_history"]) | |
| }) | |
| # Sort by creation time (newest first) | |
| conversation_list.sort(key=lambda x: x["created_at"], reverse=True) | |
| return JSONResponse(content={ | |
| "conversations": conversation_list, | |
| "total_count": len(conversation_list) | |
| }) | |
| async def continue_conversation( | |
| extraction_id: str, | |
| conversation_key: str, | |
| question: str = Form(...), | |
| background_tasks: BackgroundTasks = None, | |
| ws: str = Depends(workspace), | |
| ): | |
| """ | |
| Continue a conversation from system chart extraction with a follow-up question. | |
| """ | |
| llm_system_chart_extraction = WORKSPACE_MODELS.get(ws, {}).get("chart") | |
| if llm_system_chart_extraction is None: | |
| return JSONResponse( | |
| content={"error": "No model configured. Please set up a model first."}, | |
| status_code=400 | |
| ) | |
| # Get the conversation | |
| extraction_id = scoped(ws, extraction_id) | |
| conversations = get_conversations(extraction_id) | |
| if conversation_key not in conversations: | |
| return JSONResponse( | |
| content={"error": "Conversation not found"}, | |
| status_code=404 | |
| ) | |
| conversation = conversations[conversation_key] | |
| try: | |
| # Import the required message classes | |
| from langchain_core.messages import HumanMessage, AIMessage | |
| # Get the current conversation history | |
| current_history = conversation["conversation_history"].copy() | |
| # Add the new question with scientific formatting instructions | |
| enhanced_question = f""" | |
| {question} | |
| Please provide a detailed, scientific answer. Follow these guidelines: | |
| 1. Provide quantitative details including exact values, percentages, and units when available | |
| 2. When referencing specific information, mention which paper or section it comes from | |
| 3. For compositional data, organize information clearly | |
| 4. Include relevant scientific context and mechanisms | |
| """ | |
| new_message = HumanMessage(content=enhanced_question) | |
| current_history.append(new_message) | |
| # Get response from the model | |
| response = await llm_system_chart_extraction.ainvoke(current_history) | |
| answer_text = llm_content_to_text(response.content) | |
| ai_response = AIMessage(content=answer_text) | |
| # Update the conversation | |
| update_conversation(extraction_id, conversation_key, new_message, ai_response) | |
| return JSONResponse(content={ | |
| "question": question, | |
| "answer": answer_text, | |
| "paper_title": conversation["paper_title"], | |
| "paper_type": conversation["paper_type"], | |
| "conversation_updated": True | |
| }) | |
| except Exception as e: | |
| logger.error(f"Error continuing conversation: {str(e)}") | |
| return JSONResponse( | |
| content={"error": f"Error processing question: {str(e)}"}, | |
| status_code=500 | |
| ) | |
| async def get_conversation_history(extraction_id: str, conversation_key: str, ws: str = Depends(workspace)): | |
| """ | |
| Get the full conversation history for a specific conversation. | |
| """ | |
| conversations = get_conversations(scoped(ws, extraction_id)) | |
| if conversation_key not in conversations: | |
| return JSONResponse( | |
| content={"error": "Conversation not found"}, | |
| status_code=404 | |
| ) | |
| conversation = conversations[conversation_key] | |
| # Format the conversation history for frontend display | |
| formatted_history = [] | |
| for i, message in enumerate(conversation["conversation_history"]): | |
| # Check if it's a human or AI message | |
| if hasattr(message, 'type'): | |
| message_type = message.type | |
| content = message.content | |
| else: | |
| # Fallback for different message formats | |
| message_type = "human" if i % 2 == 0 else "ai" | |
| content = str(message) | |
| formatted_history.append({ | |
| "type": message_type, | |
| "content": content, | |
| "index": i | |
| }) | |
| return JSONResponse(content={ | |
| "paper_title": conversation["paper_title"], | |
| "paper_type": conversation["paper_type"], | |
| "created_at": conversation["created_at"], | |
| "last_updated": conversation["last_updated"], | |
| "history": formatted_history | |
| }) | |
| # -------------------- Ask Endpoint -------------------- # | |
| async def ask_question( | |
| background_tasks: BackgroundTasks, | |
| categories: str = Form(...), # Changed from category to categories (comma-separated) | |
| # Default to "" instead of required: newer FastAPI treats an empty form value as | |
| # missing, and the hidden question field submits "" in System Chart Extraction mode. | |
| question: str = Form(""), | |
| k: int = Form(10), | |
| goal: str = Form("QA"), | |
| enable_visualization: bool = Form(False), | |
| run_system_chart_normalization: bool = Form(False), | |
| ws: str = Depends(workspace), | |
| ): | |
| """ | |
| Ask a question with different 'goal' modes across multiple databases: | |
| - QA (Chunk-based Q&A across multiple databases) | |
| - System Chart Extraction (pulls pure text from PDFs, special method) | |
| """ | |
| models = WORKSPACE_MODELS.get(ws, {}) | |
| embedding_model, llm_QA, llm_system_chart_extraction = models.get("embedding"), models.get("qa"), models.get("chart") | |
| if embedding_model is None or llm_QA is None: | |
| return JSONResponse( | |
| content={"answer": "No model/embedding set. Please call /set_api_key first."}, | |
| status_code=400 | |
| ) | |
| # Parse categories (comma-separated string) | |
| category_list = [cat.strip() for cat in categories.split(",") if cat.strip()] | |
| if not category_list: | |
| return JSONResponse(content={ | |
| "answer": "No databases selected. Please select at least one database." | |
| }) | |
| # Validate that all categories exist and have files | |
| valid_categories = [] | |
| for category in category_list: | |
| if FILE_REGISTRY.get(scoped(ws, category)): | |
| valid_categories.append(category) | |
| if not valid_categories: | |
| return JSONResponse(content={ | |
| "answer": f"No files found in the selected databases: {', '.join(category_list)}" | |
| }) | |
| # System Chart Extraction can now work with multiple databases | |
| if goal == "System Chart Extraction": | |
| # Create a combined extraction ID that represents all databases | |
| database_names = "_".join(valid_categories) | |
| extraction_id = f"system_chart_multi_{uuid.uuid4().hex[:8]}" | |
| print(f"🎯 Created multi-database extraction_id: {extraction_id}") | |
| print(f"📊 Will process databases: {valid_categories}") | |
| EMBEDDING_PROGRESS[scoped(ws, extraction_id)] = { | |
| "status": "starting", | |
| "progress": 0, | |
| "message": f"Initializing system chart extraction for {len(valid_categories)} databases...", | |
| "databases": valid_categories, | |
| "total_databases": len(valid_categories) | |
| } | |
| # Start background task for multi-database system chart extraction | |
| background_tasks.add_task( | |
| process_multi_database_system_chart_extraction_async, | |
| ws, | |
| valid_categories, | |
| llm_system_chart_extraction, | |
| scoped(ws, extraction_id), | |
| enable_visualization, | |
| run_system_chart_normalization | |
| ) | |
| return JSONResponse(content={ | |
| "answer": f"System chart extraction started for {len(valid_categories)} databases: {', '.join(valid_categories)}", | |
| "extraction_id": extraction_id, | |
| "check_progress_url": f"/embedding_progress/{extraction_id}", | |
| "databases": valid_categories, | |
| "total_databases": len(valid_categories) | |
| }) | |
| # For QA goal, we can query multiple databases | |
| if goal == "QA": | |
| if not question.strip(): | |
| return JSONResponse( | |
| content={"answer": "Please enter a question before submitting."}, | |
| status_code=400 | |
| ) | |
| if llm_QA is None: | |
| return JSONResponse( | |
| content={"answer": "QA model not configured."}, | |
| status_code=400 | |
| ) | |
| qa_id = f"qa_{uuid.uuid4().hex[:8]}" | |
| EMBEDDING_PROGRESS[scoped(ws, qa_id)] = {"status": "starting", "progress": 0} | |
| # fire-and-forget | |
| background_tasks.add_task( | |
| process_qa_async, | |
| question, | |
| [scoped(ws, c) for c in valid_categories], | |
| k, | |
| embedding_model, | |
| llm_QA, | |
| scoped(ws, qa_id) | |
| ) | |
| return JSONResponse({ | |
| "answer": "Your question is being processed.", | |
| "extraction_id": qa_id, # same field name as other tasks | |
| "check_progress_url": f"/embedding_progress/{qa_id}", | |
| "databases_searched": valid_categories | |
| }) | |
| # If we get here, it's an unknown goal | |
| return JSONResponse(content={ | |
| "answer": f"Unknown goal: {goal}" | |
| }) | |
| # -------------------- Hypotheses Generation Endpoint -------------------- # | |
| async def generate_hypotheses( | |
| background_tasks: BackgroundTasks, | |
| system_charts_files: List[UploadFile] = File(...), | |
| material_goal_1: str = Form("Property #1 for xxx material"), | |
| material_goal_2: str = Form("Property #2 for xxx material"), | |
| iterations_count: int = Form(1), | |
| ws: str = Depends(workspace), | |
| ): | |
| """ | |
| Generate synergistic scientific hypotheses from multiple uploaded system charts files. | |
| Accepts multiple system_charts.txt files and generates hypotheses targeting specific material design goals. | |
| """ | |
| llm_QA = WORKSPACE_MODELS.get(ws, {}).get("qa") | |
| # Check if models are configured | |
| if llm_QA is None: | |
| return JSONResponse( | |
| content={"answer": "No model/embedding set. Please call /set_api_key first."}, | |
| status_code=400 | |
| ) | |
| # Use the QA LLM for hypothesis generation | |
| selected_llm = llm_QA | |
| # Validate that we have at least one file | |
| if not system_charts_files or len(system_charts_files) == 0: | |
| return JSONResponse( | |
| content={"answer": "Please upload at least one .txt file containing system charts with citations."}, | |
| status_code=400 | |
| ) | |
| # Validate file types | |
| invalid_files = [f.filename for f in system_charts_files if not f.filename.endswith('.txt')] | |
| if invalid_files: | |
| return JSONResponse( | |
| content={"answer": f"All files must be .txt files. Invalid files: {', '.join(invalid_files)}"}, | |
| status_code=400 | |
| ) | |
| # Validate iterations count | |
| if iterations_count < 1 or iterations_count > 100: | |
| return JSONResponse( | |
| content={"answer": "Number of iterations must be between 1 and 100."}, | |
| status_code=400 | |
| ) | |
| try: | |
| # Read all uploaded file contents | |
| system_charts_content_list = [] | |
| for file in system_charts_files: | |
| file_content = await file.read() | |
| content = file_content.decode('utf-8') | |
| system_charts_content_list.append(content) | |
| # Validate that all contents are suitable for synergistic hypothesis generation | |
| validation = HypothesesService.validate_system_charts_content(system_charts_content_list) | |
| if not validation["valid"]: | |
| return JSONResponse( | |
| content={"answer": f"Content validation failed: {validation['message']}"}, | |
| status_code=400 | |
| ) | |
| # Initialize progress tracking for hypotheses generation | |
| extraction_id = f"hypotheses_{uuid.uuid4().hex[:8]}" | |
| run_key = scoped(ws, extraction_id) | |
| HYPOTHESIS_GENERATION_CONTEXT[run_key] = { | |
| 'system_charts_content_list': system_charts_content_list, | |
| 'material_goal_1': material_goal_1, | |
| 'material_goal_2': material_goal_2, | |
| 'iterations_count': iterations_count | |
| } | |
| print(f"🎯 Created hypotheses extraction_id: {extraction_id}") | |
| EMBEDDING_PROGRESS[run_key] = { | |
| "status": "starting", | |
| "progress": 0, | |
| "message": "Initializing hypothesis generation..." | |
| } | |
| print(f"📝 Initial hypotheses progress set for {extraction_id}: {EMBEDDING_PROGRESS[run_key]}") | |
| # Start background task for hypotheses generation | |
| print(f"🚀 Adding hypotheses background task for {extraction_id}") | |
| background_tasks.add_task( | |
| process_hypotheses_generation_async, | |
| system_charts_content_list, # Pass the list of contents | |
| selected_llm, | |
| run_key, | |
| material_goal_1, | |
| material_goal_2, | |
| iterations_count | |
| ) | |
| print(f"✅ Hypotheses background task added for {extraction_id}") | |
| file_names = [f.filename for f in system_charts_files] | |
| return JSONResponse(content={ | |
| "answer": f"Synergistic hypothesis generation started for {len(system_charts_files)} file{'s' if len(system_charts_files) > 1 else ''}: {', '.join(file_names)} targeting '{material_goal_1}' and '{material_goal_2}' with {iterations_count} iteration{'s' if iterations_count > 1 else ''}.", | |
| "extraction_id": extraction_id, | |
| "check_progress_url": f"/embedding_progress/{extraction_id}", | |
| "data_entries": validation.get('data_entries', 0), | |
| "method": "synergistic_analysis", | |
| "material_goals": [material_goal_1, material_goal_2], | |
| "charts_found": validation.get('chart_count', len(system_charts_files)), | |
| "files_processed": len(system_charts_files) | |
| }) | |
| except UnicodeDecodeError: | |
| return JSONResponse( | |
| content={"answer": "Error: Could not read the uploaded file. Please ensure it's a valid UTF-8 text file."}, | |
| status_code=400 | |
| ) | |
| except Exception as e: | |
| return JSONResponse( | |
| content={"answer": f"Error processing uploaded file: {str(e)}"}, | |
| status_code=500 | |
| ) | |
| async def analyze_design_goals(description: str = Form(...), ws: str = Depends(workspace)): | |
| """ | |
| Analyze a free-form materials-design description and suggest the two target | |
| properties for hypothesis generation. Runs a single LLM call and returns the | |
| two properties so the frontend can pre-fill the editable goal fields. | |
| """ | |
| llm_QA = WORKSPACE_MODELS.get(ws, {}).get("qa") | |
| if llm_QA is None: | |
| return JSONResponse( | |
| content={"error": "No model set. Please call /set_api_key first."}, | |
| status_code=400 | |
| ) | |
| description = (description or "").strip() | |
| if not description: | |
| return JSONResponse( | |
| content={"error": "Please describe your materials design goals first."}, | |
| status_code=400 | |
| ) | |
| try: | |
| from langchain_core.messages import HumanMessage | |
| prompt = create_design_goals_analysis_prompt(description) | |
| response = await llm_QA.ainvoke([HumanMessage(content=prompt)]) | |
| raw = llm_content_to_text(response.content).strip() | |
| print(f"🎯 Design-goals analysis raw output: {raw[:300]}") | |
| # Parse the strict "Property 1: ... / Property 2: ..." format robustly. | |
| goal_1, goal_2 = _parse_design_goals(raw) | |
| if not goal_1 or not goal_2: | |
| return JSONResponse( | |
| content={ | |
| "error": "Could not confidently extract two design properties. Please refine your description or fill in the fields manually.", | |
| "raw": raw | |
| }, | |
| status_code=422 | |
| ) | |
| return JSONResponse(content={ | |
| "success": True, | |
| "goal_1": goal_1, | |
| "goal_2": goal_2 | |
| }) | |
| except Exception as e: | |
| print(f"❌ Error analyzing design goals: {str(e)}") | |
| return JSONResponse( | |
| content={"error": f"Error analyzing design goals: {str(e)}"}, | |
| status_code=500 | |
| ) | |
| def _parse_design_goals(raw: str): | |
| """Extract the two properties from the model output. | |
| Handles the expected 'Property 1: ...' / 'Property 2: ...' format and a few | |
| lenient fallbacks (numbered lists, bullet points). | |
| """ | |
| goal_1 = "" | |
| goal_2 = "" | |
| def _clean(value: str) -> str: | |
| return value.strip().strip('"').strip("'").strip("*").strip("-").strip().rstrip(".") | |
| m1 = re.search(r"Property\s*1\s*[:\-]\s*(.+)", raw, re.IGNORECASE) | |
| m2 = re.search(r"Property\s*2\s*[:\-]\s*(.+)", raw, re.IGNORECASE) | |
| if m1: | |
| goal_1 = _clean(m1.group(1).splitlines()[0]) | |
| if m2: | |
| goal_2 = _clean(m2.group(1).splitlines()[0]) | |
| # Fallback: take the first two non-empty lines if labels were not found. | |
| if not (goal_1 and goal_2): | |
| lines = [l.strip() for l in raw.splitlines() if l.strip()] | |
| cleaned = [] | |
| for line in lines: | |
| line = re.sub(r"^\s*(?:property\s*\d+|[0-9]+[.)]|[-*•])\s*[:\-]?\s*", "", line, flags=re.IGNORECASE) | |
| line = _clean(line) | |
| if line: | |
| cleaned.append(line) | |
| if not goal_1 and len(cleaned) >= 1: | |
| goal_1 = cleaned[0] | |
| if not goal_2 and len(cleaned) >= 2: | |
| goal_2 = cleaned[1] | |
| return goal_1, goal_2 | |
| # -------------------- Hypothesis Continuous Conversation Endpoints -------------------- # | |
| async def start_hypothesis_conversation( | |
| extraction_id: str = Form(...), | |
| hypothesis_number: int = Form(...), | |
| pasted_papers: str = Form("{}"), | |
| ws: str = Depends(workspace), | |
| ): | |
| """ | |
| Open a conversation about one generated hypothesis. Resolves the charts it cites | |
| and their source papers from every database in this workspace. `pasted_papers` is a JSON | |
| object {chart_number: paper_text} supplied by the user for papers not found. | |
| Replaces any previous conversation for this run. | |
| """ | |
| run = HYPOTHESIS_GENERATION_CONTEXT.get(scoped(ws, extraction_id)) | |
| if not run or 'hypotheses' not in run: | |
| return JSONResponse(content={"error": "Hypothesis run not found. Please generate hypotheses first."}, status_code=404) | |
| if WORKSPACE_MODELS.get(ws, {}).get("qa") is None: | |
| return JSONResponse(content={"error": "LLM not configured. Please set up the model first."}, status_code=400) | |
| try: | |
| pasted = {int(k): v for k, v in json.loads(pasted_papers).items() if str(v).strip()} | |
| context, sources = build_hypothesis_context(ws, run, hypothesis_number, pasted) | |
| except ValueError as e: | |
| return JSONResponse(content={"error": str(e)}, status_code=400) | |
| except Exception as e: | |
| return JSONResponse(content={"error": f"Error starting hypothesis conversation: {str(e)}"}, status_code=500) | |
| run.update({'conversation_context': context, 'conversation_history': []}) | |
| print(f"💬 Hypothesis #{hypothesis_number} conversation started for {extraction_id}: {len(context)} chars, sources={sources}") | |
| return JSONResponse(content={"success": True, "hypothesis_number": hypothesis_number, "sources": sources}) | |
| async def hypothesis_conversation_message( | |
| extraction_id: str = Form(...), | |
| message: str = Form(...), | |
| ws: str = Depends(workspace), | |
| ): | |
| """ | |
| Send a message in the hypothesis conversation and get an AI response with full context. | |
| """ | |
| run = HYPOTHESIS_GENERATION_CONTEXT.get(scoped(ws, extraction_id)) | |
| llm_QA = WORKSPACE_MODELS.get(ws, {}).get("qa") | |
| if not run or 'conversation_context' not in run: | |
| return JSONResponse(content={"error": "Hypothesis conversation not found. Please click Discuss on a hypothesis first."}, status_code=404) | |
| if llm_QA is None: | |
| return JSONResponse(content={"error": "LLM not configured. Please set up the model first."}, status_code=400) | |
| try: | |
| history = run['conversation_history'] | |
| prompt = run['conversation_context'] + "\nCONVERSATION HISTORY:\n" | |
| for entry in history[-10:]: | |
| prompt += f"Human: {entry['human']}\nAI: {entry['ai']}\n\n" | |
| prompt += f"""CURRENT QUESTION: | |
| Human: {message} | |
| Please provide a detailed, scientific answer based on the research context above. Follow these guidelines: | |
| 1. Include specific information from the source papers and system charts | |
| 2. Cite sources with the same number used in the hypothesis, e.g. "Paper [2]" or "Chart [2]" | |
| 3. For compositional data, organize information clearly | |
| 4. Provide quantitative details including exact values, percentages, and units when available | |
| 5. Explain the scientific significance of the data | |
| AI: """ | |
| print(f"🤖 Sending hypothesis prompt to LLM ({len(prompt)} chars)") | |
| from langchain_core.messages import HumanMessage | |
| response = await llm_QA.ainvoke([HumanMessage(content=prompt)]) | |
| ai_response = llm_content_to_text(response.content) | |
| history.append({'human': message, 'ai': ai_response}) | |
| return JSONResponse(content={"success": True, "response": ai_response, "conversation_length": len(history)}) | |
| except Exception as e: | |
| print(f"❌ Error processing hypothesis conversation message: {str(e)}") | |
| return JSONResponse(content={"error": f"Error processing message: {str(e)}"}, status_code=500) | |
| def get_admin_password() -> Optional[str]: | |
| password = os.environ.get("ADMIN_PASSWORD") | |
| if password: | |
| return password.strip() | |
| return None | |
| def get_admin_wipe_targets() -> List[str]: | |
| return [ | |
| os.path.abspath(UPLOAD_DIR), | |
| os.path.abspath(CHROMA_DB_BASE), | |
| ] | |
| def is_safe_target(path: str, data_root: str) -> bool: | |
| normalized_path = os.path.abspath(path) | |
| normalized_root = os.path.abspath(data_root) | |
| if normalized_path in {"/", ""}: | |
| return False | |
| try: | |
| return os.path.commonpath([normalized_root, normalized_path]) == normalized_root | |
| except ValueError: | |
| return False | |
| def delete_directory_contents(path: str) -> Dict[str, int]: | |
| deleted_files = 0 | |
| deleted_dirs = 0 | |
| if not os.path.isdir(path): | |
| return {"deleted_files": 0, "deleted_dirs": 0} | |
| for name in os.listdir(path): | |
| target = os.path.join(path, name) | |
| if os.path.isfile(target) or os.path.islink(target): | |
| os.unlink(target) | |
| deleted_files += 1 | |
| elif os.path.isdir(target): | |
| shutil.rmtree(target) | |
| deleted_dirs += 1 | |
| return {"deleted_files": deleted_files, "deleted_dirs": deleted_dirs} | |
| async def admin_verify_password(password: str = Form(...)): | |
| admin_password = get_admin_password() | |
| if not admin_password: | |
| return JSONResponse( | |
| content={"ok": False, "message": "ADMIN_PASSWORD is not configured on server."}, | |
| status_code=500 | |
| ) | |
| if not hmac.compare_digest(password or "", admin_password): | |
| return JSONResponse( | |
| content={"ok": False, "message": "Invalid admin password."}, | |
| status_code=403 | |
| ) | |
| return JSONResponse(content={"ok": True, "message": "Password verified."}, status_code=200) | |
| async def admin_wipe_data( | |
| password: str = Form(...), | |
| confirmation_text: str = Form(...) | |
| ): | |
| admin_password = get_admin_password() | |
| if not admin_password: | |
| return JSONResponse( | |
| content={"ok": False, "message": "ADMIN_PASSWORD is not configured on server."}, | |
| status_code=500 | |
| ) | |
| if not hmac.compare_digest(password or "", admin_password): | |
| return JSONResponse( | |
| content={"ok": False, "message": "Invalid admin password."}, | |
| status_code=403 | |
| ) | |
| if (confirmation_text or "").strip() != ADMIN_DELETE_CONFIRMATION_TEXT: | |
| return JSONResponse( | |
| content={ | |
| "ok": False, | |
| "message": f"Invalid confirmation text. Type exactly: {ADMIN_DELETE_CONFIRMATION_TEXT}" | |
| }, | |
| status_code=400 | |
| ) | |
| targets = get_admin_wipe_targets() | |
| data_root = os.path.abspath(APP_DATA_ROOT) | |
| for target in targets: | |
| if not is_safe_target(target, data_root): | |
| return JSONResponse( | |
| content={ | |
| "ok": False, | |
| "message": f"Refusing to wipe unsafe target: {target}", | |
| "data_root": data_root, | |
| "targets": targets | |
| }, | |
| status_code=400 | |
| ) | |
| try: | |
| results = [] | |
| total_files = 0 | |
| total_dirs = 0 | |
| for target in targets: | |
| if not os.path.exists(target): | |
| results.append({"target": target, "status": "missing"}) | |
| continue | |
| result = delete_directory_contents(target) | |
| total_files += result["deleted_files"] | |
| total_dirs += result["deleted_dirs"] | |
| results.append({"target": target, "status": "wiped", **result}) | |
| FILE_REGISTRY.clear() | |
| EMBEDDING_PROGRESS.clear() | |
| HYPOTHESIS_GENERATION_CONTEXT.clear() | |
| logger.warning( | |
| "Admin data wipe executed for targets=%s (files=%s, dirs=%s)", | |
| targets, | |
| total_files, | |
| total_dirs | |
| ) | |
| return JSONResponse( | |
| content={ | |
| "ok": True, | |
| "message": "Wipe completed.", | |
| "data_root": data_root, | |
| "targets": results, | |
| "deleted_files": total_files, | |
| "deleted_dirs": total_dirs | |
| }, | |
| status_code=200 | |
| ) | |
| except Exception as e: | |
| logger.exception("Admin data wipe failed: %s", str(e)) | |
| return JSONResponse( | |
| content={"ok": False, "message": f"Failed to wipe data: {str(e)}"}, | |
| status_code=500 | |
| ) | |
| # -------------------- Hypothesis conversation helpers -------------------- # | |
| HYPOTHESIS_PAPER_CHAR_CAP = 60_000 # per source paper | |
| HYPOTHESIS_CONTEXT_CHAR_CAP = 250_000 # all source papers combined | |
| def load_all_papers(ws: str) -> List[tuple]: | |
| """The workspace's papers as (database, title, text), databases in alphabetical order.""" | |
| return [(unscoped(ws, key), title, text) for key in sorted(FILE_REGISTRY) if unscoped(ws, key) is not None | |
| for title, text in obtain_text_from_files(key).items()] | |
| def match_paper(title: str, papers: List[tuple]): | |
| """Find a paper by title: exact or substring first, then fuzzy (score >= 80). Returns (database, text).""" | |
| title = (title or "").strip() | |
| best = (0, None, None) | |
| for db, stored, text in papers: | |
| if title and (title == stored or title in stored or stored in title): | |
| return db, text | |
| score = fuzz.token_set_ratio(title, stored) if title else 0 | |
| if score > best[0]: | |
| best = (score, db, text) | |
| return (best[1], best[2]) if best[0] >= 80 else (None, None) | |
| def build_hypothesis_context(ws: str, run: dict, hypothesis_number: int, pasted: Dict[int, str]): | |
| """ | |
| Build the fixed conversation context for one hypothesis: the hypothesis text plus, | |
| for every chart it cites as [n], that chart and its source paper (full text, capped). | |
| Returns (context_text, sources) where sources describes how each citation resolved. | |
| """ | |
| hypothesis = next((h for h in run['hypotheses'] if h['number'] == hypothesis_number), None) | |
| if hypothesis is None: | |
| raise ValueError(f"Hypothesis #{hypothesis_number} not found in this run.") | |
| charts = [(m.group('paper').strip(), m.group('chart').strip()) | |
| for blob in run['system_charts_content_list'] for m in CHART_PATTERN.finditer(blob)] | |
| cited = sorted({int(n) for n in re.findall(r'\[(\d+)\]', hypothesis['text']) if 0 < int(n) <= len(charts)}) | |
| papers = load_all_papers(ws) if any(n not in pasted for n in cited) else [] | |
| context = f"=== HYPOTHESIS #{hypothesis_number} UNDER DISCUSSION ===\n{hypothesis['text']}\n\n" | |
| sources, budget = [], HYPOTHESIS_CONTEXT_CHAR_CAP | |
| for n in cited: | |
| title, chart = charts[n - 1] | |
| db, text = ("pasted", pasted[n]) if n in pasted else match_paper(title, papers) | |
| truncated = False | |
| if text: | |
| allowed = min(HYPOTHESIS_PAPER_CHAR_CAP, budget) | |
| truncated = len(text) > allowed | |
| text = text[:allowed] | |
| budget -= len(text) | |
| sources.append({'number': n, 'title': title, 'database': db, 'truncated': truncated}) | |
| context += (f"=== SOURCE [{n}]: {title} ===\n" | |
| f"CHART [{n}]:\n{chart}\n\n" | |
| f"PAPER [{n}]:\n{text or 'NOT AVAILABLE'}\n\n") | |
| context += "=== END OF CONTEXT ===\n\n" | |
| return context, sources | |
| if __name__ == "__main__": | |
| import uvicorn | |
| # Log important startup information for diagnostics | |
| logger.info("=" * 50) | |
| logger.info("Starting PDF-RAG application") | |
| logger.info(f"Running as executable: {hasattr(sys, '_MEIPASS')}") | |
| logger.info(f"Current working directory: {os.getcwd()}") | |
| if hasattr(sys, '_MEIPASS'): | |
| logger.info(f"PyInstaller _MEIPASS: {sys._MEIPASS}") | |
| logger.info(f"Executable path: {sys.executable}") | |
| logger.info(f"Base directory for storage: {get_base_directory()}") | |
| logger.info(f"UPLOAD_DIR: {os.path.abspath(UPLOAD_DIR)}") | |
| logger.info(f"CHROMA_DB_BASE: {os.path.abspath(CHROMA_DB_BASE)}") | |
| logger.info("=" * 50) | |
| print("Starting server...") | |
| uvicorn.run(app, host="127.0.0.1", port=8000) | |