diff --git "a/main.py" "b/main.py" new file mode 100644--- /dev/null +++ "b/main.py" @@ -0,0 +1,3074 @@ +import os +import sys +import uuid +import asyncio +import aiofiles +import concurrent.futures +import logging +import json +import re +import time +import random +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 +from fastapi.responses import JSONResponse, FileResponse +from fastapi.middleware.cors import CORSMiddleware +from fastapi.staticfiles import StaticFiles +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 + +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 +from system_chart_extraction import system_chart_extraction +from system_chart_service import SystemChartService +from hypotheses_service import HypothesesService +from prompts import create_hypotheses_generation_prompt +from conversation_storage import save_conversation, get_conversations, update_conversation + +# 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(): + """Return base directory, preferring persistent cloud storage when provided.""" + persistent_dir = os.environ.get("PERSISTENT_STORAGE_PATH") + if persistent_dir: + return persistent_dir + if hasattr(sys, "_MEIPASS"): # Running as PyInstaller executable + return os.path.dirname(sys.executable) + return os.path.dirname(os.path.abspath(__file__)) + +# Define paths relative to the base directory for persistent storage +UPLOAD_DIR = os.path.join(get_base_directory(), "uploaded files") +os.makedirs(UPLOAD_DIR, exist_ok=True) +logger.info(f"UPLOAD_DIR set to: {os.path.abspath(UPLOAD_DIR)}") + +@asynccontextmanager +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=["*"] +) + +# Use the directory name for the mount path, not the full path +upload_dir_name = os.path.basename(UPLOAD_DIR) +logger.info(f"Mounting uploaded files at: /{upload_dir_name}") + +app.mount(f"/{upload_dir_name}", StaticFiles(directory=UPLOAD_DIR), name="uploaded") + +# 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 = os.path.join(get_base_directory(), "root databases") + + 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 + +@app.get("/") +def get_index(): + return FileResponse(resource_path("index.html")) + + +# -------------------- 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" +DEFAULT_MODEL_QA = DEFAULT_MODEL +DEFAULT_MODEL_SYSTEM_CHART_EXTRACTION = DEFAULT_MODEL + +# -------------- Global placeholders (ensure they're declared) -------------- # +api_key = None +embedding_model = None +llm_QA = None +llm_system_chart_extraction = None + +def get_model_provider(model: str) -> Optional[str]: + if model in OPENAI_MODELS: + return "openai" + if model in GEMINI_MODELS: + return "gemini" + return None + +MODEL_ADVANCED_PARAMETERS: Dict[str, Dict[str, Any]] = { + "gpt-5-nano": { + "reasoning_effort": {"choices": ["minimal", "low", "medium", "high"]}, + "verbosity": {"choices": ["low", "medium", "high"]}, + }, + "gpt-5-mini": { + "reasoning_effort": {"choices": ["minimal", "low", "medium", "high"]}, + "verbosity": {"choices": ["low", "medium", "high"]}, + }, + "gpt-5.2": { + "reasoning_effort": {"choices": ["none", "low", "medium", "high", "xhigh"]}, + "verbosity": {"choices": ["low", "medium", "high"]}, + }, + "gemini-2.5-flash-lite": { + "temperature": {"min": 0.0, "max": 2.0}, + "top_p": {"min": 0.0, "max": 1.0}, + }, + "gemini-3-flash-preview": { + "temperature": {"min": 0.0, "max": 2.0}, + "top_p": {"min": 0.0, "max": 1.0}, + "thinking_level": {"choices": ["minimal", "low", "medium", "high"]}, + }, + "gemini-3-pro-preview": { + "temperature": {"min": 0.0, "max": 2.0}, + "top_p": {"min": 0.0, "max": 1.0}, + "thinking_level": {"choices": ["low", "high"]}, + }, +} + + +def get_model_advanced_parameter_rules(model: str) -> Dict[str, Any]: + return MODEL_ADVANCED_PARAMETERS.get(model, {}) + + +def parse_optional_float(value: Optional[str], field_name: str) -> Tuple[Optional[float], Optional[str]]: + if value is None: + return None, None + value_str = value.strip() + if not value_str: + return None, None + try: + return float(value_str), None + except ValueError: + return None, f"Invalid {field_name}: '{value}'." + + +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) + +@app.post("/set_api_key") +async def set_api_key( + api_key_input: Optional[str] = Form(None), + model: str = Form(DEFAULT_MODEL), + temperature: Optional[str] = Form(None), + top_p: Optional[str] = Form(None), + reasoning_effort: Optional[str] = Form(None), + verbosity: Optional[str] = Form(None), + thinking_level: Optional[str] = Form(None), +): + """ + Sets up global LLM and embeddings for proprietary models. + """ + global api_key + global embedding_model + global llm_QA, llm_system_chart_extraction + global DEFAULT_MODEL_QA, DEFAULT_MODEL_SYSTEM_CHART_EXTRACTION + + 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 = get_model_provider(model) + if not model_provider: + return JSONResponse( + content={"message": f"Invalid proprietary model: {model}"}, + status_code=400 + ) + model_advanced_rules = get_model_advanced_parameter_rules(model) + 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 + ) + + llm_config: Dict[str, Any] = {} + allowed_advanced_params = set(model_advanced_rules.keys()) + + parsed_temperature, temp_error = parse_optional_float(temperature, "temperature") + if temp_error: + return JSONResponse(content={"message": temp_error}, status_code=400) + if parsed_temperature is not None: + temperature_rules = model_advanced_rules.get("temperature", {}) + min_temp = float(temperature_rules.get("min", 0.0)) + max_temp = float(temperature_rules.get("max", 2.0)) + if not min_temp <= parsed_temperature <= max_temp: + return JSONResponse( + content={"message": f"temperature must be between {min_temp} and {max_temp}."}, + status_code=400 + ) + if "temperature" not in allowed_advanced_params: + return JSONResponse( + content={"message": f"Model '{model}' does not support temperature."}, + status_code=400 + ) + llm_config["temperature"] = parsed_temperature + + parsed_top_p, top_p_error = parse_optional_float(top_p, "top_p") + if top_p_error: + return JSONResponse(content={"message": top_p_error}, status_code=400) + if parsed_top_p is not None: + top_p_rules = model_advanced_rules.get("top_p", {}) + min_top_p = float(top_p_rules.get("min", 0.0)) + max_top_p = float(top_p_rules.get("max", 1.0)) + if not min_top_p <= parsed_top_p <= max_top_p: + return JSONResponse( + content={"message": f"top_p must be between {min_top_p} and {max_top_p}."}, + status_code=400 + ) + if "top_p" not in allowed_advanced_params: + return JSONResponse( + content={"message": f"Model '{model}' does not support top_p."}, + status_code=400 + ) + llm_config["top_p"] = parsed_top_p + + if reasoning_effort: + normalized_reasoning_effort = reasoning_effort.strip().lower() + if normalized_reasoning_effort: + allowed_reasoning_efforts = model_advanced_rules.get("reasoning_effort", {}).get("choices", []) + if normalized_reasoning_effort not in allowed_reasoning_efforts: + return JSONResponse( + content={"message": f"Invalid reasoning_effort: '{reasoning_effort}'."}, + status_code=400 + ) + if "reasoning_effort" not in allowed_advanced_params: + return JSONResponse( + content={"message": f"Model '{model}' does not support reasoning_effort."}, + status_code=400 + ) + llm_config["reasoning_effort"] = normalized_reasoning_effort + + if verbosity: + normalized_verbosity = verbosity.strip().lower() + if normalized_verbosity: + allowed_verbosity = model_advanced_rules.get("verbosity", {}).get("choices", []) + if normalized_verbosity not in allowed_verbosity: + return JSONResponse( + content={"message": f"Invalid verbosity: '{verbosity}'."}, + status_code=400 + ) + if "verbosity" not in allowed_advanced_params: + return JSONResponse( + content={"message": f"Model '{model}' does not support verbosity."}, + status_code=400 + ) + llm_config["verbosity"] = normalized_verbosity + + if thinking_level: + normalized_thinking_level = thinking_level.strip().lower() + if normalized_thinking_level: + allowed_thinking_levels = model_advanced_rules.get("thinking_level", {}).get("choices", []) + if normalized_thinking_level not in allowed_thinking_levels: + return JSONResponse( + content={"message": f"Invalid thinking_level: '{thinking_level}'."}, + status_code=400 + ) + if "thinking_level" not in allowed_advanced_params: + return JSONResponse( + content={"message": f"Model '{model}' does not support thinking_level."}, + status_code=400 + ) + llm_config["thinking_level"] = normalized_thinking_level + + if model_provider == "openai": + embedding_model = OpenAIEmbeddings(openai_api_key=api_key) + embedding_model = ensure_async_embeddings(embedding_model) + openai_model_kwargs: Dict[str, Any] = {} + if "reasoning_effort" in llm_config: + openai_model_kwargs["reasoning"] = { + "effort": llm_config["reasoning_effort"] + } + if "verbosity" in llm_config: + openai_model_kwargs["text"] = { + "verbosity": llm_config["verbosity"] + } + openai_common_kwargs: Dict[str, Any] = { + "openai_api_key": api_key, + "model": DEFAULT_MODEL_QA, + "use_responses_api": True + } + openai_chart_kwargs: Dict[str, Any] = { + "openai_api_key": api_key, + "model": DEFAULT_MODEL_SYSTEM_CHART_EXTRACTION, + "use_responses_api": True + } + if openai_model_kwargs: + openai_common_kwargs["model_kwargs"] = openai_model_kwargs + openai_chart_kwargs["model_kwargs"] = openai_model_kwargs + 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 + ) + embedding_model = GoogleGenerativeAIEmbeddings( + google_api_key=api_key, + model="models/embedding-001" + ) + embedding_model = ensure_async_embeddings(embedding_model) + gemini_common_kwargs: Dict[str, Any] = { + "google_api_key": api_key, + "model": DEFAULT_MODEL_QA + } + gemini_chart_kwargs: Dict[str, Any] = { + "google_api_key": api_key, + "model": DEFAULT_MODEL_SYSTEM_CHART_EXTRACTION + } + if "temperature" in llm_config: + gemini_common_kwargs["temperature"] = llm_config["temperature"] + gemini_chart_kwargs["temperature"] = llm_config["temperature"] + if "top_p" in llm_config: + gemini_common_kwargs["top_p"] = llm_config["top_p"] + gemini_chart_kwargs["top_p"] = llm_config["top_p"] + + gemini_model_kwargs: Dict[str, Any] = {} + if "thinking_level" in llm_config: + gemini_model_kwargs["thinking_level"] = llm_config["thinking_level"] + if gemini_model_kwargs: + gemini_common_kwargs["model_kwargs"] = gemini_model_kwargs + gemini_chart_kwargs["model_kwargs"] = gemini_model_kwargs + llm_QA = ChatGoogleGenerativeAI( + **gemini_common_kwargs + ) + llm_system_chart_extraction = ChatGoogleGenerativeAI( + **gemini_chart_kwargs + ) + provider_label = "Gemini" + else: + return JSONResponse( + content={"message": f"Unsupported model provider for {model}."}, + status_code=400 + ) + + 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 + ) + +@app.get("/available_models") +async def get_available_models(): + return JSONResponse(content={ + "models": AVAILABLE_MODELS + }) + +# -------------------- 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 + +# Non-parallel version +async def embed_texts_async(texts: List[str]): + """ + Asynchronously embed a list of texts with the current embedding_model, in batches. + """ + global embedding_model + if embedding_model is None: + raise ValueError("Embedding model not set. Call /set_api_key first.") + + batch_size = 1024 # tune per provider/rate limits + out = [] + for i in range(0, len(texts), batch_size): + j = min(i + batch_size, len(texts)) + out.extend(await embedding_model.aembed_documents(texts[i:j])) + return out + +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 -------------------- # +@app.get("/embedding_progress/{database}") +async def get_embedding_progress(database: str): + """ + Returns the current progress of embedding tasks for a specific database. + """ + + if database 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[database].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) + +@app.get("/download_extraction/{extraction_id}") +async def download_extraction(extraction_id: str): + """ + Download the system chart extraction ZIP file directly from memory. + """ + if extraction_id not in EMBEDDING_PROGRESS: + return JSONResponse(content={"error": "Extraction not found"}, status_code=404) + + progress_data = EMBEDDING_PROGRESS[extraction_id] + + 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): + """ + 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_model, 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(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(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, selected_dbs: list = None): + """ + 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, # Store list instead of single content + 'databases': selected_dbs or [], # Use selected databases instead of hardcoded 'hypo' + 'material_goal_1': material_goal_1, + 'material_goal_2': material_goal_2, + 'iterations_count': iterations_count + } + print(f"šŸ’¾ Stored hypothesis generation context for {extraction_id} with databases: {selected_dbs}") + + # 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"]: + 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 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 = getattr(response, "content", str(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 = getattr(final_response, "content", str(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}" + } + +@app.post("/upload") +async def upload_files( + background_tasks: BackgroundTasks, + category: str = Form(...), + files: List[UploadFile] = File(...) +): + """ + 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 + # 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 + ) + + 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 '{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)) + 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)) + 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." + }) + +@app.post("/upload_zip") +async def upload_zip( + background_tasks: BackgroundTasks, + category: str = Form(...), + zip_file: UploadFile = File(...) +): + """ + 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 + if embedding_model is None: + return JSONResponse( + content={"message": "No embedding model set. Please call /set_api_key first."}, + status_code=400 + ) + + 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 '{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)) + 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 + }) + +@app.get("/category_info/{category}") +async def get_category_info(category: str): + """ + Get information about the files in a category to determine available modes. + """ + if category not in FILE_REGISTRY or not FILE_REGISTRY[category]: + return JSONResponse(content={ + "has_files": False, + "file_types": { + "text_pdf_count": 0, + "image_count": 0, + "pdf_count": 0, + "has_images": False, + "has_text_pdf": False, + "has_pdfs": False + }, + "available_modes": [] + }) + + # 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 FILE_REGISTRY[category].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"]) + + return JSONResponse(content={ + "has_files": True, + "file_types": { + "text_pdf_count": text_pdf_count, + "image_count": image_count, + "pdf_count": pdf_count, + "has_images": image_count > 0, + "has_text_pdf": text_pdf_count > 0, + "has_pdfs": pdf_count > 0 + }, + "available_modes": available_modes, + "total_files": text_pdf_count + image_count + }) + +@app.get("/available_databases") +async def get_available_databases(): + """ + Returns a list of all available databases/categories with their file information. + """ + global FILE_REGISTRY + + databases = [] + for category, files in FILE_REGISTRY.items(): + if files: # Only include 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": category, + "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}) + +@app.get("/metadata/{database}") +async def get_database_metadata(database: str): + """ + Get metadata information for all files in a database. + Useful for debugging and displaying file information including original PDF titles. + """ + if database not in FILE_REGISTRY: + return JSONResponse(content={ + "error": f"Database '{database}' not found" + }, status_code=404) + + metadata = load_metadata(database) + + # Combine file registry info with metadata + files_info = [] + for file_hash, filename in FILE_REGISTRY[database].items(): + file_metadata = metadata.get(file_hash, {}) + file_info = { + "file_hash": file_hash, + "file_hash_short": file_hash[:16] + "...", + "saved_filename": filename, + "original_filename": file_metadata.get("original_filename", filename), + "title": file_metadata.get("title", "Unknown"), + "file_type": file_metadata.get("file_type", "unknown"), + "upload_timestamp": file_metadata.get("upload_timestamp", "Unknown"), + "file_size": file_metadata.get("file_size", 0), + "metadata_found": bool(file_metadata) + } + files_info.append(file_info) + + return JSONResponse(content={ + "database": database, + "file_count": len(files_info), + "files": files_info, + "metadata_keys": [h[:16] + "..." for h in metadata.keys()], + "debug": { + "total_metadata_entries": len(metadata), + "files_with_metadata": len([f for f in files_info if f["metadata_found"]]), + "files_without_metadata": len([f for f in files_info if not f["metadata_found"]]) + } + }) + +# -------------------- Conversation Management Endpoints -------------------- # + +@app.get("/conversations/{extraction_id}") +async def get_extraction_conversations(extraction_id: str): + """ + Get all conversations for a given extraction ID. + Returns a list of conversation metadata for frontend dropdown. + """ + conversations = get_conversations(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) + }) + +@app.post("/continue_conversation/{extraction_id}/{conversation_key}") +async def continue_conversation( + extraction_id: str, + conversation_key: str, + question: str = Form(...), + background_tasks: BackgroundTasks = None +): + """ + Continue a conversation from system chart extraction with a follow-up question. + """ + global llm_system_chart_extraction + + 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 + 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) + ai_response = AIMessage(content=response.content) + + # Update the conversation + update_conversation(extraction_id, conversation_key, new_message, ai_response) + + return JSONResponse(content={ + "question": question, + "answer": response.content, + "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 + ) + +@app.get("/conversation_history/{extraction_id}/{conversation_key}") +async def get_conversation_history(extraction_id: str, conversation_key: str): + """ + Get the full conversation history for a specific conversation. + """ + 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] + + # 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 -------------------- # +@app.post("/ask") +async def ask_question( + background_tasks: BackgroundTasks, + categories: str = Form(...), # Changed from category to categories (comma-separated) + question: str = Form(...), + k: int = Form(10), + goal: str = Form("QA"), + enable_visualization: bool = Form(False), + run_system_chart_normalization: bool = Form(False) +): + """ + 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) + """ + global embedding_model, llm_QA, llm_system_chart_extraction + 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 category in FILE_REGISTRY and FILE_REGISTRY[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[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, + valid_categories, + llm_system_chart_extraction, + 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 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[qa_id] = {"status": "starting", "progress": 0} + + # fire-and-forget + background_tasks.add_task( + process_qa_async, + question, + valid_categories, + k, + embedding_model, + llm_QA, + 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 -------------------- # +@app.post("/generate_hypotheses") +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), + categories: str = Form("") +): + """ + Generate synergistic scientific hypotheses from multiple uploaded system charts files. + Accepts multiple system_charts.txt files and generates hypotheses targeting specific material design goals. + """ + global llm_QA, llm_system_chart_extraction + + # 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 + ) + + selected_dbs = [c.strip() for c in categories.split(",") if c.strip() in FILE_REGISTRY] + + # Initialize progress tracking for hypotheses generation + extraction_id = f"hypotheses_{uuid.uuid4().hex[:8]}" + HYPOTHESIS_GENERATION_CONTEXT[extraction_id] = { + 'system_charts_content_list': system_charts_content_list, # Store list instead of single content + 'databases': selected_dbs, + '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[extraction_id] = { + "status": "starting", + "progress": 0, + "message": "Initializing hypothesis generation..." + } + print(f"šŸ“ Initial hypotheses progress set for {extraction_id}: {EMBEDDING_PROGRESS[extraction_id]}") + + # 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, + extraction_id, + material_goal_1, + material_goal_2, + iterations_count, + selected_dbs # Pass the selected databases + ) + 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 + ) + +# -------------------- Hypothesis Continuous Conversation Endpoint -------------------- # +@app.post("/start_hypothesis_conversation") +async def start_hypothesis_conversation( + extraction_id: str = Form(...), + hypothesis_text: str = Form(...), + original_paper_content: str = Form(...) +): + """ + Start continuous conversation for a hypothesis using user-provided paper content. + Accepts an extraction_id from completed hypothesis generation, the hypothesis text to analyze, + and the original paper content provided by the user. + """ + global HYPOTHESIS_GENERATION_CONTEXT + + print(f"šŸ” Starting hypothesis conversation for extraction_id: {extraction_id}") + print(f"šŸ“ Hypothesis text length: {len(hypothesis_text)} characters") + print(f"šŸ“„ Original paper content length: {len(original_paper_content)} characters") + + if extraction_id not in HYPOTHESIS_GENERATION_CONTEXT: + print(f"āŒ Hypothesis generation context not found for {extraction_id}") + return JSONResponse( + content={"error": "Hypothesis generation context not found. Please ensure the hypothesis generation process has completed."}, + status_code=404 + ) + + try: + context = HYPOTHESIS_GENERATION_CONTEXT[extraction_id] + + print(f"šŸ“‹ Found context for {extraction_id}") + + # Get system charts content - handle both new list format and old single content format + system_charts_content_list = context.get('system_charts_content_list') + if system_charts_content_list: + # New format: combine multiple files into a single string for processing + combined_system_charts_content = '\n\n'.join(system_charts_content_list) + else: + # Old format: fallback to single content + combined_system_charts_content = context.get('system_charts_content', '') + + # Parse system charts to get citation mapping and filter by detected markers + from system_chart_service import SystemChartService + charts = SystemChartService.parse_system_charts_file(combined_system_charts_content) + print(f"ļæ½ Parsed {len(charts)} system charts from content") + + # Create citation number ā†’ chart mapping for all charts + all_charts = {} + for i, chart in enumerate(charts, 1): + all_charts[i] = { + 'title': chart.get('title', ''), + 'citation': chart.get('citation', ''), + 'content': chart.get('content', ''), + 'chart_number': chart.get('chart_number', str(i)) + } + + # Extract citation markers from hypothesis to get only referenced charts + citation_numbers = extract_citation_markers(hypothesis_text) + print(f"šŸ“ Citation numbers found in hypothesis: {citation_numbers}") + + # Filter charts to only include those referenced in the hypothesis + filtered_charts = {} + for citation_num in citation_numbers: + if citation_num in all_charts: + filtered_charts[citation_num] = all_charts[citation_num] + print(f"šŸ“Š Including chart [{citation_num}]: {all_charts[citation_num]['title'][:50]}...") + else: + print(f"āš ļø Chart [{citation_num}] referenced but not found in system charts") + + print(f"šŸ“‹ Filtered to {len(filtered_charts)} charts based on hypothesis citations (from {len(all_charts)} total)") + + # Create conversation context using user-provided paper content and filtered charts + print(f"šŸ“ Creating conversation context with user-provided paper content and filtered charts...") + conversation_context = create_conversation_context_with_user_content( + original_paper_content, + filtered_charts, + hypothesis_text + ) + + print(f"āœ… Successfully created conversation context") + print(f"šŸ“Š Results: User-provided paper content + {len(filtered_charts)} charts") + + # Store the conversation context for follow-up questions + HYPOTHESIS_GENERATION_CONTEXT[extraction_id].update({ + 'conversation_context': conversation_context, + 'user_provided_paper_content': original_paper_content, + 'resolved_charts': filtered_charts, + 'conversation_history': [] # Initialize conversation history + }) + + return JSONResponse(content={ + "success": True, + "conversation_context": conversation_context, + "charts_found": len(filtered_charts), + "chart_titles": [chart['title'] for chart in filtered_charts.values()], + "material_goals": [context['material_goal_1'], context['material_goal_2']] + }) + + except Exception as e: + print(f"āŒ Error starting hypothesis conversation: {str(e)}") + return JSONResponse( + content={"error": f"Error starting hypothesis conversation: {str(e)}"}, + status_code=500 + ) + +# -------------------- Hypothesis Conversation Message Endpoint -------------------- # +@app.post("/hypothesis_conversation_message") +async def hypothesis_conversation_message( + extraction_id: str = Form(...), + message: str = Form(...) +): + """ + Send a message in the hypothesis conversation and get an AI response with full context. + """ + global HYPOTHESIS_GENERATION_CONTEXT, llm_QA + + print(f"šŸ’¬ Processing hypothesis conversation message for {extraction_id}") + print(f"šŸ“ User message: {message[:100]}...") + + if extraction_id not in HYPOTHESIS_GENERATION_CONTEXT: + return JSONResponse( + content={"error": "Hypothesis conversation not found. Please start a conversation first."}, + status_code=404 + ) + + context = HYPOTHESIS_GENERATION_CONTEXT[extraction_id] + + if 'conversation_context' not in context: + return JSONResponse( + content={"error": "Conversation context not loaded. Please start a conversation first."}, + status_code=400 + ) + + if llm_QA is None: + return JSONResponse( + content={"error": "LLM not configured. Please set up the model first."}, + status_code=400 + ) + + try: + # Create the full prompt with context and conversation history + full_context = context['conversation_context'] + conversation_history = context.get('conversation_history', []) + + # Build the prompt with conversation history + prompt = f"""{full_context} + +CONVERSATION HISTORY: +""" + + # Add previous conversation history + for entry in conversation_history[-10:]: # Keep last 10 exchanges + prompt += f"Human: {entry['human']}\nAI: {entry['ai']}\n\n" + + # Add current question + 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 original papers and system charts +2. When referencing information, mention which paper or chart it comes from (e.g., "Paper 1:", "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 prompt to LLM (length: {len(prompt)} characters)") + + # Get response from LLM + from langchain_core.messages import HumanMessage + response = await llm_QA.ainvoke([HumanMessage(content=prompt)]) + ai_response = response.content + + print(f"āœ… Received LLM response (length: {len(ai_response)} characters)") + + # Store in conversation history + conversation_history.append({ + 'human': message, + 'ai': ai_response + }) + context['conversation_history'] = conversation_history + + # Return the response + return JSONResponse(content={ + "success": True, + "response": ai_response, + "has_user_paper_content": 'user_provided_paper_content' in context, + "charts_count": len(context.get('resolved_charts', {})), + "conversation_length": len(conversation_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 + ) + +@app.post("/debug/reload_databases") +async def reload_databases(): + """ + Debug endpoint to reload databases and trigger metadata migration. + """ + global FILE_REGISTRY + + # Clear current registry + FILE_REGISTRY = {} + + # Reload databases (this will trigger migration) + load_existing_databases() + + return JSONResponse(content={ + "message": "Databases reloaded with metadata migration", + "databases": list(FILE_REGISTRY.keys()), + "total_files": sum(len(files) for files in FILE_REGISTRY.values()) + }) + +def extract_citation_markers(hypothesis_text): + """ + DEPRECATED: Extract citation markers like [1], [2], etc. from hypothesis text + This function is no longer used in the new hypothesis conversation flow + where users manually provide paper content. + """ + citation_pattern = r'\[(\d+)\]' + citations = re.findall(citation_pattern, hypothesis_text) + print(f"šŸ” Found citation markers in hypothesis: {citations}") + return [int(c) for c in citations] + +def parse_system_charts_for_citations(system_charts_content): + """ + DEPRECATED: Parse system charts file and create citation mapping + This function is no longer used in the new hypothesis conversation flow + where users manually provide paper content. + """ + from system_chart_service import SystemChartService + + # Parse individual charts + charts = SystemChartService.parse_system_charts_file(system_charts_content) + print(f"šŸ“Š Parsed {len(charts)} system charts from content") + + # Create citation number ā†’ chart mapping + citation_to_chart = {} + for i, chart in enumerate(charts, 1): + citation_to_chart[i] = { + 'title': chart.get('title', ''), + 'citation': chart.get('citation', ''), + 'content': chart.get('content', ''), + 'chart_number': chart.get('chart_number', str(i)) + } + print(f"šŸ“‹ Chart [{i}]: {chart.get('title', 'No title')[:50]}...") + + return citation_to_chart + +def resolve_papers_and_charts(hypothesis_text, system_charts_content, database): + """ + DEPRECATED: Resolve citation markers to original papers and system charts + This function is no longer used in the new hypothesis conversation flow + where users manually provide paper content. + + Args: + hypothesis_text: The hypothesis containing citation markers + system_charts_content: Content of the system charts file (in memory) + database: Database name(s) where original papers are stored (string or list) + + Returns: + dict: { + 'papers': {title: full_text}, + 'charts': {citation_number: chart_info}, + 'citation_numbers': [1, 2, 3, ...] + } + """ + print(f"šŸ” Resolving papers and charts for hypothesis with {len(hypothesis_text)} characters") + + # Extract citation markers from hypothesis + citation_numbers = extract_citation_markers(hypothesis_text) + print(f"šŸ“ Citation numbers found: {citation_numbers}") + + # Parse system charts file + citation_to_chart = parse_system_charts_for_citations(system_charts_content) + print(f"šŸ“Š Available charts: {list(citation_to_chart.keys())}") + + # Get all original papers from database(s) - support both single and multiple databases + all_papers = {} + if isinstance(database, (list, tuple, set)): + for db in database: + for title, text in obtain_text_from_files(db).items(): + all_papers[title] = text + print(f"šŸ“š Available papers from databases {database}: {len(all_papers)}") + else: + all_papers = obtain_text_from_files(database) # title -> full_text mapping + print(f"šŸ“š Available papers in database '{database}': {len(all_papers)}") + + for title in all_papers.keys(): + print(f" - {title[:60]}...") + + # Find relevant papers and charts + relevant_papers = {} + relevant_charts = {} + + for citation_num in citation_numbers: + if citation_num not in citation_to_chart: + continue + + chart_info = citation_to_chart[citation_num] + paper_title = (chart_info.get('title') or '').strip() + matched_title = None + + # exact / substring first + for stored_title in all_papers.keys(): + if (paper_title == stored_title + or (paper_title and paper_title in stored_title) + or (stored_title and stored_title in paper_title)): + matched_title = stored_title + break + + # fuzzy fallback + if not matched_title and paper_title: + best_title = None + best_score = 0 + for stored_title in all_papers.keys(): + score = fuzz.token_set_ratio(paper_title, stored_title) + if score > best_score: + best_score = score + best_title = stored_title + if best_score >= 80: # threshold you can tune + matched_title = best_title + + # ALWAYS keep the chart + relevant_charts[citation_num] = chart_info + + # add paper if we found a match + if matched_title: + relevant_papers[matched_title] = all_papers[matched_title] + else: + print(f"āŒ No matching paper found for chart title: {paper_title}") + + print(f"\nšŸ“Š Final results: {len(relevant_papers)} papers, {len(relevant_charts)} charts") + + return { + 'papers': relevant_papers, + 'charts': relevant_charts, + 'citation_numbers': citation_numbers + } + +def create_conversation_context_with_user_content(user_paper_content, charts, hypothesis_text): + """Create conversation context using user-provided paper content""" + print(f"šŸ“ Creating conversation context with user-provided paper content ({len(user_paper_content)} chars) and {len(charts)} charts") + + context = "=== RESEARCH CONTEXT FOR HYPOTHESIS ANALYSIS ===\n\n" + + # Add user-provided paper content + context += "ORIGINAL PAPER CONTENT:\n" + context += "=" * 50 + "\n" + context += user_paper_content + "\n\n" + print(f"šŸ“„ Added user-provided paper content ({len(user_paper_content)} chars)") + + # Add system charts + context += "\nSYSTEM CHARTS:\n" + context += "=" * 50 + "\n" + for citation_num, chart_info in charts.items(): + context += f"Chart [{citation_num}]: {chart_info['title']}\n" + context += f"Citation: {chart_info['citation']}\n" + context += "-" * 30 + "\n" + context += chart_info['content'] + "\n\n" + print(f"šŸ“Š Added chart [{citation_num}]: {chart_info['title'][:50]}... ({len(chart_info['content'])} chars)") + + # Add the hypothesis + context += "\nHYPOTHESIS TO ANALYZE:\n" + context += "=" * 50 + "\n" + context += hypothesis_text + "\n\n" + + context += "=== END OF CONTEXT ===\n\n" + context += "You can now ask questions about this hypothesis, the provided paper content, or the system charts. How would you like to proceed?" + + print(f"šŸ“ Created context with {len(context)} total characters") + + return context + +def create_conversation_context(papers, charts, hypothesis_text): + """Create pre-populated conversation context""" + print(f"šŸ“ Creating conversation context with {len(papers)} papers and {len(charts)} charts") + + context = "=== RESEARCH CONTEXT FOR HYPOTHESIS ANALYSIS ===\n\n" + + # Add original papers + context += "ORIGINAL PAPERS:\n" + context += "=" * 50 + "\n" + for title, content in papers.items(): + context += f"Paper: {title}\n" + context += "-" * 30 + "\n" + context += content + "\n\n" + print(f"šŸ“„ Added paper: {title[:50]}... ({len(content)} chars)") + + # Add system charts + context += "\nSYSTEM CHARTS:\n" + context += "=" * 50 + "\n" + for citation_num, chart_info in charts.items(): + context += f"Chart [{citation_num}]: {chart_info['title']}\n" + context += f"Citation: {chart_info['citation']}\n" + context += "-" * 30 + "\n" + context += chart_info['content'] + "\n\n" + print(f"šŸ“Š Added chart [{citation_num}]: {chart_info['title'][:50]}... ({len(chart_info['content'])} chars)") + + # Add the hypothesis + context += "\nHYPOTHESIS TO ANALYZE:\n" + context += "=" * 50 + "\n" + context += hypothesis_text + "\n\n" + + context += "=== END OF CONTEXT ===\n\n" + context += "You can now ask questions about this hypothesis, the supporting papers, or the system charts. How would you like to proceed?" + + print(f"šŸ“ Created context with {len(context)} total characters") + + return context + +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)}") + from vectorstore import CHROMA_DB_BASE + 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) \ No newline at end of file