From 4b29b3d790efc21273e7174067fc3125a912c26e Mon Sep 17 00:00:00 2001 From: Christian Date: Sat, 13 Dec 2025 03:30:48 +0100 Subject: [PATCH 01/21] Add PostgreSQL support with SQLAlchemy adapter - Create DatabaseManager using SQLAlchemy Core API for unified database interface - Add database schema definitions using SQLAlchemy Table/Column - Refactor CacheManager to use SQLAlchemy instead of direct sqlite3 - Add pgvector support in enterprise module for PostgreSQL vector search - Update tests to support both SQLite and PostgreSQL adapters - Maintain backward compatibility with existing SQLite deployments --- report_analyst/core/cache_manager.py | 876 ++++++++++-------- report_analyst/core/database_manager.py | 104 +++ report_analyst/core/database_schema.py | 118 +++ report_analyst_enterprise/LICENSE | 17 + report_analyst_enterprise/__init__.py | 9 + .../database/__init__.py | 18 + .../database/pgvector_support.py | 145 +++ report_analyst_enterprise/requirements.txt | 14 + tests/integration/test_streamlit_app.py | 23 +- tests/test_analyzer.py | 8 +- tests/test_cache_manager.py | 172 ++-- tests/test_llm_evidence_separation.py | 33 +- 12 files changed, 1036 insertions(+), 501 deletions(-) create mode 100644 report_analyst/core/database_manager.py create mode 100644 report_analyst/core/database_schema.py create mode 100644 report_analyst_enterprise/LICENSE create mode 100644 report_analyst_enterprise/__init__.py create mode 100644 report_analyst_enterprise/database/__init__.py create mode 100644 report_analyst_enterprise/database/pgvector_support.py create mode 100644 report_analyst_enterprise/requirements.txt diff --git a/report_analyst/core/cache_manager.py b/report_analyst/core/cache_manager.py index 4c655a782..c54d4e76f 100644 --- a/report_analyst/core/cache_manager.py +++ b/report_analyst/core/cache_manager.py @@ -1,7 +1,6 @@ import json import logging import os -import sqlite3 from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional @@ -9,21 +8,45 @@ import numpy as np from llama_index.core import Document, QueryBundle from llama_index.core.indices import VectorStoreIndex +from sqlalchemy import text + +from .database_manager import DatabaseManager +from .database_schema import indexes, metadata logger = logging.getLogger(__name__) class CacheManager: - def __init__(self, db_path: str = None): - if db_path is None: - # Use the project's storage path + def __init__(self, db_path: str = None, database_url: str = None): + """ + Initialize CacheManager. + + Args: + db_path: Path to SQLite database file (for backward compatibility). + If None and database_url is None, uses default SQLite path. + database_url: SQLAlchemy database URL (e.g., 'sqlite:///path' or 'postgresql://...'). + Takes precedence over db_path. + """ + # Determine database URL + if database_url: + db_url = database_url + elif db_path: + # Convert file path to SQLite URL + db_path_obj = Path(db_path) + db_path_obj.parent.mkdir(parents=True, exist_ok=True) + db_url = f"sqlite:///{db_path}" + self.db_path = db_path_obj # Keep for backward compatibility + else: + # Default to SQLite storage_path = os.getenv("STORAGE_PATH", "./storage") db_path = str(Path(storage_path) / "cache" / "analysis.db") + Path(db_path).parent.mkdir(parents=True, exist_ok=True) + db_url = f"sqlite:///{db_path}" + self.db_path = Path(db_path) # Keep for backward compatibility - self.db_path = Path(db_path) - # Create parent directories if they don't exist - self.db_path.parent.mkdir(parents=True, exist_ok=True) - logger.info(f"Initializing CacheManager with db: {self.db_path}") + # Initialize database manager + self.db_manager = DatabaseManager(db_url) + logger.info(f"Initializing CacheManager with database: {self.db_manager._mask_url(db_url)}") self.init_db() # In-memory vector store for current document @@ -31,107 +54,28 @@ def __init__(self, db_path: str = None): self.current_file_path = None def init_db(self): - """Initialize the database schema""" - conn = sqlite3.connect(self.db_path) + """Initialize the database schema using SQLAlchemy""" try: - # Create optimized table for document chunks - conn.execute( - """ - CREATE TABLE IF NOT EXISTS document_chunks ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - file_path TEXT, - chunk_text TEXT, - chunk_size INTEGER, - chunk_overlap INTEGER, - embedding BLOB, -- Store embedding for caching - metadata TEXT, - created_at TIMESTAMP, - UNIQUE(file_path, chunk_text, chunk_size, chunk_overlap) - ) - """ - ) - - # Create indices for better performance - conn.execute( - "CREATE INDEX IF NOT EXISTS idx_file_path ON document_chunks(file_path)" - ) - conn.execute( - "CREATE INDEX IF NOT EXISTS idx_chunk_params ON document_chunks(chunk_size, chunk_overlap)" - ) - - # Store questions separately - conn.execute( - """ - CREATE TABLE IF NOT EXISTS questions ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - question_id TEXT, - question_set TEXT, - question_text TEXT, - guidelines TEXT, - UNIQUE(question_id, question_set) - ) - """ - ) - - # Create analysis cache table - conn.execute( - """ - CREATE TABLE IF NOT EXISTS analysis_cache ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - file_path TEXT, - question_id TEXT, - chunk_size INTEGER, - chunk_overlap INTEGER, - top_k INTEGER, - model TEXT, - question_set TEXT, - result TEXT, - created_at TIMESTAMP, - UNIQUE(file_path, question_id, chunk_size, chunk_overlap, top_k, model, question_set) - ) - """ - ) - - # Store analysis configurations and results - conn.execute( - """ - CREATE TABLE IF NOT EXISTS question_analysis ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - file_path TEXT, - question_id INTEGER, - model TEXT, - top_k INTEGER, - analysis_result TEXT, - version INTEGER DEFAULT 1, - created_at TIMESTAMP, - FOREIGN KEY(question_id) REFERENCES questions(id), - UNIQUE(file_path, question_id, model, top_k, version) - ) - """ - ) - - # Store chunk-question relationships with all scores and ordering - conn.execute( - """ - CREATE TABLE IF NOT EXISTS chunk_relevance ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - question_analysis_id INTEGER, - document_chunk_id INTEGER, - chunk_order INTEGER, - similarity_score REAL, - llm_score REAL, - is_evidence BOOLEAN, - evidence_order INTEGER, - metadata TEXT, - FOREIGN KEY(question_analysis_id) REFERENCES question_analysis(id), - FOREIGN KEY(document_chunk_id) REFERENCES document_chunks(id), - UNIQUE(question_analysis_id, document_chunk_id) - ) - """ - ) - - finally: - conn.close() + engine = self.db_manager.get_engine() + # Create all tables + metadata.create_all(engine) + + # Create indexes (using raw SQL for IF NOT EXISTS support) + # Note: Some databases may not support IF NOT EXISTS in CREATE INDEX + # We'll try to create them and ignore errors if they already exist + with self.db_manager.get_connection() as conn: + for index_sql in indexes: + try: + conn.execute(text(index_sql)) + except Exception as e: + # Index might already exist, which is fine + logger.debug(f"Index creation (may already exist): {e}") + conn.commit() + + logger.info("Database schema initialized successfully") + except Exception as e: + logger.error(f"Error initializing database schema: {str(e)}", exc_info=True) + raise def _load_vector_store(self, file_path: str, chunks: List[Dict]) -> None: """Load chunks into an in-memory vector store.""" @@ -243,65 +187,124 @@ def save_analysis( logger.info(f"Saving analysis for {file_path} - {question_id}") logger.info(f"Configuration: {json.dumps(config, indent=2)}") - # First, ensure question exists in questions table - with sqlite3.connect(self.db_path) as conn: - # Extract question set and number from question_id (format: set_number) - question_set = question_id.split("_")[0] - question_number = int(question_id.split("_")[1]) + # Extract question set and number from question_id (format: set_number) + question_set = question_id.split("_")[0] + with self.db_manager.get_connection() as conn: + # Ensure question exists in questions table logger.info( f"Ensuring question {question_id} exists in questions table" ) - cursor = conn.execute( - """ - SELECT id FROM questions - WHERE question_id = ? AND question_set = ? - """, - (question_id, question_set), + result_obj = conn.execute( + text(""" + SELECT id FROM questions + WHERE question_id = :question_id AND question_set = :question_set + """), + {"question_id": question_id, "question_set": question_set}, ) - row = cursor.fetchone() + row = result_obj.fetchone() if row: question_db_id = row[0] logger.info(f"Found existing question with DB ID: {question_db_id}") else: - # Insert new question - cursor = conn.execute( - """ - INSERT INTO questions (question_id, question_set, question_text, guidelines) - VALUES (?, ?, ?, ?) - RETURNING id - """, - ( - question_id, - question_set, - result.get("question_text", ""), - result.get("guidelines", ""), - ), - ) - question_db_id = cursor.fetchone()[0] + # Insert new question - use dialect-specific upsert + if self.db_manager.is_postgres(): + # PostgreSQL: ON CONFLICT + result_obj = conn.execute( + text(""" + INSERT INTO questions (question_id, question_set, question_text, guidelines) + VALUES (:question_id, :question_set, :question_text, :guidelines) + ON CONFLICT (question_id, question_set) DO UPDATE + SET question_text = EXCLUDED.question_text, + guidelines = EXCLUDED.guidelines + RETURNING id + """), + { + "question_id": question_id, + "question_set": question_set, + "question_text": result.get("question_text", ""), + "guidelines": result.get("guidelines", ""), + }, + ) + else: + # SQLite: INSERT OR REPLACE + result_obj = conn.execute( + text(""" + INSERT OR REPLACE INTO questions (question_id, question_set, question_text, guidelines) + VALUES (:question_id, :question_set, :question_text, :guidelines) + """), + { + "question_id": question_id, + "question_set": question_set, + "question_text": result.get("question_text", ""), + "guidelines": result.get("guidelines", ""), + }, + ) + # Get the ID separately for SQLite + result_obj = conn.execute( + text("SELECT id FROM questions WHERE question_id = :question_id AND question_set = :question_set"), + {"question_id": question_id, "question_set": question_set}, + ) + question_db_id = result_obj.fetchone()[0] logger.info(f"Created new question with DB ID: {question_db_id}") # Save main analysis result logger.info("Saving main analysis result") - cursor = conn.execute( - """ - INSERT OR REPLACE INTO question_analysis - (file_path, question_id, model, top_k, analysis_result, version, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?) - RETURNING id - """, - ( - str(file_path), - question_db_id, - config["model"], - config["top_k"], - json.dumps(result), - 1, # version - datetime.now().isoformat(), - ), - ) - analysis_id = cursor.fetchone()[0] + if self.db_manager.is_postgres(): + result_obj = conn.execute( + text(""" + INSERT INTO question_analysis + (file_path, question_id, model, top_k, analysis_result, version, created_at) + VALUES (:file_path, :question_id, :model, :top_k, :analysis_result, :version, :created_at) + ON CONFLICT (file_path, question_id, model, top_k, version) DO UPDATE + SET analysis_result = EXCLUDED.analysis_result, + created_at = EXCLUDED.created_at + RETURNING id + """), + { + "file_path": str(file_path), + "question_id": question_db_id, + "model": config["model"], + "top_k": config["top_k"], + "analysis_result": json.dumps(result), + "version": 1, + "created_at": datetime.now().isoformat(), + }, + ) + else: + result_obj = conn.execute( + text(""" + INSERT OR REPLACE INTO question_analysis + (file_path, question_id, model, top_k, analysis_result, version, created_at) + VALUES (:file_path, :question_id, :model, :top_k, :analysis_result, :version, :created_at) + """), + { + "file_path": str(file_path), + "question_id": question_db_id, + "model": config["model"], + "top_k": config["top_k"], + "analysis_result": json.dumps(result), + "version": 1, + "created_at": datetime.now().isoformat(), + }, + ) + # Get ID separately for SQLite + result_obj = conn.execute( + text(""" + SELECT id FROM question_analysis + WHERE file_path = :file_path AND question_id = :question_id + AND model = :model AND top_k = :top_k AND version = :version + """), + { + "file_path": str(file_path), + "question_id": question_db_id, + "model": config["model"], + "top_k": config["top_k"], + "version": 1, + }, + ) + analysis_id = result_obj.fetchone()[0] logger.info(f"Analysis ID: {analysis_id}") # Save chunk relevance information @@ -313,37 +316,66 @@ def save_analysis( logger.debug(f"Processing chunk: {json.dumps(chunk, indent=2)}") # Get chunk ID from document_chunks table - cursor = conn.execute( - """ - SELECT id FROM document_chunks - WHERE file_path = ? AND chunk_text = ? - """, - (str(file_path), chunk["text"]), + result_obj = conn.execute( + text(""" + SELECT id FROM document_chunks + WHERE file_path = :file_path AND chunk_text = :chunk_text + """), + {"file_path": str(file_path), "chunk_text": chunk["text"]}, ) - row = cursor.fetchone() + row = result_obj.fetchone() if row: chunk_id = row[0] logger.debug(f"Found chunk ID: {chunk_id}") # Save chunk relevance with all available information - conn.execute( - """ - INSERT OR REPLACE INTO chunk_relevance - (question_analysis_id, document_chunk_id, chunk_order, - similarity_score, llm_score, is_evidence, evidence_order, metadata) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - analysis_id, - chunk_id, - chunk.get("chunk_order", 0), - chunk.get("similarity_score", 0.0), - chunk.get("llm_score", None), - chunk.get("is_evidence", False), - chunk.get("evidence_order"), - json.dumps(chunk.get("metadata", {})), - ), - ) + if self.db_manager.is_postgres(): + conn.execute( + text(""" + INSERT INTO chunk_relevance + (question_analysis_id, document_chunk_id, chunk_order, + similarity_score, llm_score, is_evidence, evidence_order, metadata) + VALUES (:question_analysis_id, :document_chunk_id, :chunk_order, + :similarity_score, :llm_score, :is_evidence, :evidence_order, :metadata) + ON CONFLICT (question_analysis_id, document_chunk_id) DO UPDATE + SET chunk_order = EXCLUDED.chunk_order, + similarity_score = EXCLUDED.similarity_score, + llm_score = EXCLUDED.llm_score, + is_evidence = EXCLUDED.is_evidence, + evidence_order = EXCLUDED.evidence_order, + metadata = EXCLUDED.metadata + """), + { + "question_analysis_id": analysis_id, + "document_chunk_id": chunk_id, + "chunk_order": chunk.get("chunk_order", 0), + "similarity_score": chunk.get("similarity_score", 0.0), + "llm_score": chunk.get("llm_score"), + "is_evidence": chunk.get("is_evidence", False), + "evidence_order": chunk.get("evidence_order"), + "metadata": json.dumps(chunk.get("metadata", {})), + }, + ) + else: + conn.execute( + text(""" + INSERT OR REPLACE INTO chunk_relevance + (question_analysis_id, document_chunk_id, chunk_order, + similarity_score, llm_score, is_evidence, evidence_order, metadata) + VALUES (:question_analysis_id, :document_chunk_id, :chunk_order, + :similarity_score, :llm_score, :is_evidence, :evidence_order, :metadata) + """), + { + "question_analysis_id": analysis_id, + "document_chunk_id": chunk_id, + "chunk_order": chunk.get("chunk_order", 0), + "similarity_score": chunk.get("similarity_score", 0.0), + "llm_score": chunk.get("llm_score"), + "is_evidence": chunk.get("is_evidence", False), + "evidence_order": chunk.get("evidence_order"), + "metadata": json.dumps(chunk.get("metadata", {})), + }, + ) logger.info( f"Saving raw values to DB - similarity_score: {chunk.get('similarity_score')}, llm_score: {chunk.get('llm_score')}, is_evidence: {chunk.get('is_evidence')}" ) @@ -354,25 +386,51 @@ def save_analysis( # Save to analysis cache logger.info("Saving to analysis cache") - conn.execute( - """ - INSERT OR REPLACE INTO analysis_cache - (file_path, question_id, chunk_size, chunk_overlap, top_k, - model, question_set, result, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - str(file_path), - question_id, # Use original question_id here - config["chunk_size"], - config["chunk_overlap"], - config["top_k"], - config["model"], - config["question_set"], - json.dumps(result), - datetime.now().isoformat(), - ), - ) + if self.db_manager.is_postgres(): + conn.execute( + text(""" + INSERT INTO analysis_cache + (file_path, question_id, chunk_size, chunk_overlap, top_k, + model, question_set, result, created_at) + VALUES (:file_path, :question_id, :chunk_size, :chunk_overlap, :top_k, + :model, :question_set, :result, :created_at) + ON CONFLICT (file_path, question_id, chunk_size, chunk_overlap, top_k, model, question_set) DO UPDATE + SET result = EXCLUDED.result, + created_at = EXCLUDED.created_at + """), + { + "file_path": str(file_path), + "question_id": question_id, + "chunk_size": config["chunk_size"], + "chunk_overlap": config["chunk_overlap"], + "top_k": config["top_k"], + "model": config["model"], + "question_set": config["question_set"], + "result": json.dumps(result), + "created_at": datetime.now().isoformat(), + }, + ) + else: + conn.execute( + text(""" + INSERT OR REPLACE INTO analysis_cache + (file_path, question_id, chunk_size, chunk_overlap, top_k, + model, question_set, result, created_at) + VALUES (:file_path, :question_id, :chunk_size, :chunk_overlap, :top_k, + :model, :question_set, :result, :created_at) + """), + { + "file_path": str(file_path), + "question_id": question_id, + "chunk_size": config["chunk_size"], + "chunk_overlap": config["chunk_overlap"], + "top_k": config["top_k"], + "model": config["model"], + "question_set": config["question_set"], + "result": json.dumps(result), + "created_at": datetime.now().isoformat(), + }, + ) logger.info("Successfully saved complete analysis") @@ -400,18 +458,7 @@ def get_analysis( Dict mapping question_ids to their analysis results with chunks """ try: - with sqlite3.connect(self.db_path) as conn: - # First get the analysis results from the cache table - query = """ - SELECT question_id, result - FROM analysis_cache - WHERE file_path = ? - AND chunk_size = ? - AND chunk_overlap = ? - AND top_k = ? - AND model = ? - AND question_set = ? - """ + with self.db_manager.get_connection() as conn: # Map question set to database identifier (same mapping as in save_analysis) question_set_mapping = { "everest": "ev", @@ -423,22 +470,35 @@ def get_analysis( config["question_set"], config["question_set"] ) - params = [ - str(file_path), - config["chunk_size"], - config["chunk_overlap"], - config["top_k"], - config["model"], - db_question_set, # Use mapped question set - ] + # First get the analysis results from the cache table + query = """ + SELECT question_id, result + FROM analysis_cache + WHERE file_path = :file_path + AND chunk_size = :chunk_size + AND chunk_overlap = :chunk_overlap + AND top_k = :top_k + AND model = :model + AND question_set = :question_set + """ + params = { + "file_path": str(file_path), + "chunk_size": config["chunk_size"], + "chunk_overlap": config["chunk_overlap"], + "top_k": config["top_k"], + "model": config["model"], + "question_set": db_question_set, + } if question_ids: - placeholders = ",".join("?" * len(question_ids)) + # Use SQLAlchemy's IN clause with bind parameters + placeholders = ",".join(f":qid_{i}" for i in range(len(question_ids))) query += f" AND question_id IN ({placeholders})" - params.extend(question_ids) + for i, qid in enumerate(question_ids): + params[f"qid_{i}"] = qid - cursor = conn.execute(query, params) - rows = cursor.fetchall() + result_obj = conn.execute(text(query), params) + rows = result_obj.fetchall() # Process results results = {} @@ -452,7 +512,9 @@ def get_analysis( # Now get the chunk information for each question if results: - chunk_query = """ + # Build IN clause for question IDs + qid_placeholders = ",".join(f":qid_{i}" for i in range(len(results))) + chunk_query = f""" SELECT ac.question_id, dc.chunk_text, @@ -468,30 +530,30 @@ def get_analysis( JOIN question_analysis qa ON qa.question_id = q.id AND qa.file_path = ac.file_path JOIN chunk_relevance cr ON cr.question_analysis_id = qa.id JOIN document_chunks dc ON cr.document_chunk_id = dc.id - WHERE ac.file_path = ? - AND ac.chunk_size = ? - AND ac.chunk_overlap = ? - AND ac.top_k = ? - AND ac.model = ? - AND ac.question_set = ? - AND ac.question_id IN ({}) + WHERE ac.file_path = :file_path + AND ac.chunk_size = :chunk_size + AND ac.chunk_overlap = :chunk_overlap + AND ac.top_k = :top_k + AND ac.model = :model + AND ac.question_set = :question_set + AND ac.question_id IN ({qid_placeholders}) ORDER BY ac.question_id, cr.chunk_order - """.format( - ",".join("?" * len(results)) - ) + """ + + chunk_params = { + "file_path": str(file_path), + "chunk_size": config["chunk_size"], + "chunk_overlap": config["chunk_overlap"], + "top_k": config["top_k"], + "model": config["model"], + "question_set": db_question_set, + } + for i, qid in enumerate(results.keys()): + chunk_params[f"qid_{i}"] = qid - chunk_params = [ - str(file_path), - config["chunk_size"], - config["chunk_overlap"], - config["top_k"], - config["model"], - db_question_set, # Use mapped question set - ] + list(results.keys()) - - logger.info(f"Executing chunk query with params: {chunk_params}") - chunk_cursor = conn.execute(chunk_query, chunk_params) - chunk_rows = chunk_cursor.fetchall() + logger.info(f"Executing chunk query with params: {list(chunk_params.keys())}") + chunk_result = conn.execute(text(chunk_query), chunk_params) + chunk_rows = chunk_result.fetchall() logger.info(f"Retrieved {len(chunk_rows)} chunk rows") # Add chunks to their respective questions @@ -543,10 +605,7 @@ def save_vectors(self, file_path: str, chunks: List[Dict[str, Any]]) -> None: ) logger.info(f"Chunk parameters: size={chunk_size}, overlap={chunk_overlap}") - # Begin transaction - with sqlite3.connect(self.db_path) as conn: - cursor = conn.cursor() - + with self.db_manager.get_connection() as conn: # Prepare chunks for insertion chunk_data = [] for i, chunk in enumerate(chunks): @@ -564,17 +623,15 @@ def save_vectors(self, file_path: str, chunks: List[Dict[str, Any]]) -> None: metadata_with_shape["embedding_dtype"] = str(embedding.dtype) # Prepare chunk data - chunk_data.append( - ( - file_path, - chunk["text"], - chunk_size, - chunk_overlap, - embedding_bytes, - json.dumps(metadata_with_shape), - datetime.now().isoformat(), - ) - ) + chunk_data.append({ + "file_path": file_path, + "chunk_text": chunk["text"], + "chunk_size": chunk_size, + "chunk_overlap": chunk_overlap, + "embedding": embedding_bytes, + "metadata": json.dumps(metadata_with_shape), + "created_at": datetime.now().isoformat(), + }) except Exception as e: logger.warning( f"Error preparing chunk {i} for storage: {str(e)}" @@ -582,27 +639,48 @@ def save_vectors(self, file_path: str, chunks: List[Dict[str, Any]]) -> None: continue if chunk_data: - # Insert all chunks in a single transaction - cursor.executemany( - """ - INSERT INTO document_chunks ( - file_path, chunk_text, chunk_size, chunk_overlap, - embedding, metadata, created_at - ) VALUES (?, ?, ?, ?, ?, ?, ?) - """, - chunk_data, - ) + # Insert all chunks - use dialect-specific upsert + if self.db_manager.is_postgres(): + # PostgreSQL: ON CONFLICT + for chunk_row in chunk_data: + conn.execute( + text(""" + INSERT INTO document_chunks + (file_path, chunk_text, chunk_size, chunk_overlap, + embedding, metadata, created_at) + VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, + :embedding, :metadata, :created_at) + ON CONFLICT (file_path, chunk_text, chunk_size, chunk_overlap) DO UPDATE + SET embedding = EXCLUDED.embedding, + metadata = EXCLUDED.metadata, + created_at = EXCLUDED.created_at + """), + chunk_row, + ) + else: + # SQLite: INSERT OR REPLACE + for chunk_row in chunk_data: + conn.execute( + text(""" + INSERT OR REPLACE INTO document_chunks + (file_path, chunk_text, chunk_size, chunk_overlap, + embedding, metadata, created_at) + VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, + :embedding, :metadata, :created_at) + """), + chunk_row, + ) logger.info( f"Successfully saved all {len(chunk_data)} chunks to database" ) # Verify the insertion - cursor.execute( - "SELECT COUNT(*) FROM document_chunks WHERE file_path = ?", - (file_path,), + result_obj = conn.execute( + text("SELECT COUNT(*) FROM document_chunks WHERE file_path = :file_path"), + {"file_path": file_path}, ) - count = cursor.fetchone()[0] + count = result_obj.fetchone()[0] logger.info( f"Verification: Found {count} chunks in database for {file_path}" ) @@ -616,25 +694,26 @@ def save_vectors(self, file_path: str, chunks: List[Dict[str, Any]]) -> None: def get_vectors(self, file_path: str) -> List[Dict[str, Any]]: """Get vector embeddings for a document""" try: - with sqlite3.connect(self.db_path) as conn: + with self.db_manager.get_connection() as conn: + result_obj = conn.execute( + text(""" + SELECT chunk_text, embedding, metadata + FROM document_chunks + WHERE file_path = :file_path + """), + {"file_path": str(file_path)}, + ) chunks = [] - for row in conn.execute( - """ - SELECT chunk_text, embedding, metadata - FROM document_chunks - WHERE file_path = ? - """, - (str(file_path),), - ): - metadata = json.loads(row[2]) if row[2] else {} + for row in result_obj: + metadata_dict = json.loads(row[2]) if row[2] else {} # Reconstruct embedding with proper shape embedding = None if row[1]: try: # Get shape and dtype from metadata - shape = tuple(metadata.get("embedding_shape", [])) - dtype = metadata.get("embedding_dtype", "float32") + shape = tuple(metadata_dict.get("embedding_shape", [])) + dtype = metadata_dict.get("embedding_dtype", "float32") if shape: embedding = np.frombuffer(row[1], dtype=dtype).reshape( @@ -650,7 +729,7 @@ def get_vectors(self, file_path: str) -> List[Dict[str, Any]]: # Remove embedding metadata from the returned metadata clean_metadata = { k: v - for k, v in metadata.items() + for k, v in metadata_dict.items() if k not in ["embedding_shape", "embedding_dtype"] } @@ -670,20 +749,20 @@ def get_vectors(self, file_path: str) -> List[Dict[str, Any]]: def clear_cache(self, file_path: Optional[str] = None): """Clear cache entries""" try: - with sqlite3.connect(self.db_path) as conn: + with self.db_manager.get_connection() as conn: if file_path: conn.execute( - "DELETE FROM analysis_cache WHERE file_path = ?", - (str(file_path),), + text("DELETE FROM analysis_cache WHERE file_path = :file_path"), + {"file_path": str(file_path)}, ) conn.execute( - "DELETE FROM document_chunks WHERE file_path = ?", - (str(file_path),), + text("DELETE FROM document_chunks WHERE file_path = :file_path"), + {"file_path": str(file_path)}, ) logger.info(f"Cleared cache for {file_path}") else: - conn.execute("DELETE FROM analysis_cache") - conn.execute("DELETE FROM document_chunks") + conn.execute(text("DELETE FROM analysis_cache")) + conn.execute(text("DELETE FROM document_chunks")) logger.info("Cleared all cache") except Exception as e: logger.error(f"Error clearing cache: {str(e)}", exc_info=True) @@ -691,27 +770,27 @@ def clear_cache(self, file_path: Optional[str] = None): def check_cache_status(self, file_path: str = None): """Debug method to check cache contents""" try: - with sqlite3.connect(self.db_path) as conn: + with self.db_manager.get_connection() as conn: if file_path: logger.info(f"Checking cache for file: {file_path}") - cursor = conn.execute( - """ - SELECT DISTINCT chunk_size, chunk_overlap, top_k, model, question_set - FROM analysis_cache - WHERE file_path = ? - """, - (str(file_path),), + result_obj = conn.execute( + text(""" + SELECT DISTINCT chunk_size, chunk_overlap, top_k, model, question_set + FROM analysis_cache + WHERE file_path = :file_path + """), + {"file_path": str(file_path)}, ) else: logger.info("Checking all cache entries") - cursor = conn.execute( - """ - SELECT DISTINCT file_path, chunk_size, chunk_overlap, top_k, model, question_set - FROM analysis_cache - """ + result_obj = conn.execute( + text(""" + SELECT DISTINCT file_path, chunk_size, chunk_overlap, top_k, model, question_set + FROM analysis_cache + """) ) - rows = cursor.fetchall() + rows = result_obj.fetchall() logger.info(f"Found {len(rows)} distinct configurations:") for row in rows: logger.info(f"Config: {row}") @@ -725,28 +804,28 @@ def check_cache_status(self, file_path: str = None): def get_all_answers_by_question_set(self, question_set: str) -> Dict[str, Any]: """Get all cached answers for a specific question set""" try: - with sqlite3.connect(self.db_path) as conn: + with self.db_manager.get_connection() as conn: # First get all analysis results - cursor = conn.execute( - """ - SELECT ac.question_id, ac.result, - dc.chunk_text, dc.metadata as chunk_metadata, - cr.chunk_order, cr.similarity_score, - cr.llm_score, cr.is_evidence, cr.evidence_order, - cr.metadata as relevance_metadata - FROM analysis_cache ac - LEFT JOIN questions q ON q.question_id = ac.question_id - LEFT JOIN question_analysis qa ON qa.question_id = q.id - LEFT JOIN chunk_relevance cr ON cr.question_analysis_id = qa.id - LEFT JOIN document_chunks dc ON cr.document_chunk_id = dc.id - WHERE ac.question_set = ? - ORDER BY ac.question_id, cr.chunk_order - """, - (question_set,), + result_obj = conn.execute( + text(""" + SELECT ac.question_id, ac.result, + dc.chunk_text, dc.metadata as chunk_metadata, + cr.chunk_order, cr.similarity_score, + cr.llm_score, cr.is_evidence, cr.evidence_order, + cr.metadata as relevance_metadata + FROM analysis_cache ac + LEFT JOIN questions q ON q.question_id = ac.question_id + LEFT JOIN question_analysis qa ON qa.question_id = q.id + LEFT JOIN chunk_relevance cr ON cr.question_analysis_id = qa.id + LEFT JOIN document_chunks dc ON cr.document_chunk_id = dc.id + WHERE ac.question_set = :question_set + ORDER BY ac.question_id, cr.chunk_order + """), + {"question_set": question_set}, ) results = {} - for row in cursor.fetchall(): + for row in result_obj: question_id = row[0] result_json = row[1] chunk_text = row[2] @@ -797,11 +876,9 @@ def save_document_chunks( logger.info(f"Starting to save {len(chunks)} chunks for {file_path}") logger.info(f"Chunk parameters: size={chunk_size}, overlap={chunk_overlap}") - conn = sqlite3.connect(self.db_path) - try: - conn.execute("BEGIN TRANSACTION") - timestamp = datetime.now().isoformat() + timestamp = datetime.now().isoformat() + with self.db_manager.get_connection() as conn: for i, chunk in enumerate(chunks): logger.debug(f"Processing chunk {i+1}/{len(chunks)}") @@ -815,41 +892,64 @@ def save_document_chunks( metadata_json = json.dumps(chunk.get("metadata", {})) - cursor = conn.execute( - """ - INSERT OR REPLACE INTO document_chunks - (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?) - """, - ( - str(file_path), - chunk["text"], - chunk_size, - chunk_overlap, - embedding_bytes, - metadata_json, - timestamp, - ), - ) - - logger.debug(f"Inserted chunk with ID: {cursor.lastrowid}") + if self.db_manager.is_postgres(): + conn.execute( + text(""" + INSERT INTO document_chunks + (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) + VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) + ON CONFLICT (file_path, chunk_text, chunk_size, chunk_overlap) DO UPDATE + SET embedding = EXCLUDED.embedding, + metadata = EXCLUDED.metadata, + created_at = EXCLUDED.created_at + """), + { + "file_path": str(file_path), + "chunk_text": chunk["text"], + "chunk_size": chunk_size, + "chunk_overlap": chunk_overlap, + "embedding": embedding_bytes, + "metadata": metadata_json, + "created_at": timestamp, + }, + ) + else: + conn.execute( + text(""" + INSERT OR REPLACE INTO document_chunks + (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) + VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) + """), + { + "file_path": str(file_path), + "chunk_text": chunk["text"], + "chunk_size": chunk_size, + "chunk_overlap": chunk_overlap, + "embedding": embedding_bytes, + "metadata": metadata_json, + "created_at": timestamp, + }, + ) - conn.execute("COMMIT") logger.info(f"Successfully saved all {len(chunks)} chunks to database") # Verify chunks were saved - cursor = conn.execute( - "SELECT COUNT(*) FROM document_chunks WHERE file_path = ? AND chunk_size = ? AND chunk_overlap = ?", - (str(file_path), chunk_size, chunk_overlap), + result_obj = conn.execute( + text(""" + SELECT COUNT(*) FROM document_chunks + WHERE file_path = :file_path AND chunk_size = :chunk_size AND chunk_overlap = :chunk_overlap + """), + { + "file_path": str(file_path), + "chunk_size": chunk_size, + "chunk_overlap": chunk_overlap, + }, ) - count = cursor.fetchone()[0] + count = result_obj.fetchone()[0] logger.info( f"Verification: Found {count} chunks in database for {file_path}" ) - finally: - conn.close() - except Exception as e: logger.error(f"Error saving document chunks: {str(e)}", exc_info=True) raise @@ -866,27 +966,27 @@ def get_document_chunks( f"Filters: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}" ) - with sqlite3.connect(self.db_path) as conn: + with self.db_manager.get_connection() as conn: query = """ SELECT id, chunk_text, embedding, metadata, chunk_size, chunk_overlap FROM document_chunks - WHERE file_path = ? + WHERE file_path = :file_path """ - params = [str(file_path)] + params = {"file_path": str(file_path)} if chunk_size is not None: - query += " AND chunk_size = ?" - params.append(chunk_size) + query += " AND chunk_size = :chunk_size" + params["chunk_size"] = chunk_size if chunk_overlap is not None: - query += " AND chunk_overlap = ?" - params.append(chunk_overlap) + query += " AND chunk_overlap = :chunk_overlap" + params["chunk_overlap"] = chunk_overlap logger.debug(f"Executing query: {query}") logger.debug(f"Query parameters: {params}") - cursor = conn.execute(query, params) - rows = cursor.fetchall() + result_obj = conn.execute(text(query), params) + rows = result_obj.fetchall() chunks = [] for row in rows: @@ -951,27 +1051,27 @@ def get_chunks_without_embeddings( f"Filters: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}" ) - with sqlite3.connect(self.db_path) as conn: + with self.db_manager.get_connection() as conn: query = """ SELECT id, chunk_text, metadata, chunk_size, chunk_overlap FROM document_chunks - WHERE file_path = ? AND embedding IS NULL + WHERE file_path = :file_path AND embedding IS NULL """ - params = [str(file_path)] + params = {"file_path": str(file_path)} if chunk_size is not None: - query += " AND chunk_size = ?" - params.append(chunk_size) + query += " AND chunk_size = :chunk_size" + params["chunk_size"] = chunk_size if chunk_overlap is not None: - query += " AND chunk_overlap = ?" - params.append(chunk_overlap) + query += " AND chunk_overlap = :chunk_overlap" + params["chunk_overlap"] = chunk_overlap logger.debug(f"Executing query: {query}") logger.debug(f"Query parameters: {params}") - cursor = conn.execute(query, params) - rows = cursor.fetchall() + result_obj = conn.execute(text(query), params) + rows = result_obj.fetchall() chunks = [] for row in rows: @@ -1005,23 +1105,23 @@ def get_chunks_without_embeddings( def has_chunk_scoring(self, file_path: str, config: Dict) -> bool: """Check if any questions have been scored for this file/config""" try: - with sqlite3.connect(self.db_path) as conn: - cursor = conn.execute( - """ - SELECT COUNT(DISTINCT q.question_id) - FROM questions q - JOIN question_analysis qa ON qa.question_id = q.id - JOIN chunk_relevance cr ON cr.question_analysis_id = qa.id - WHERE qa.file_path = ? AND qa.model = ? AND qa.top_k = ? - """, - ( - str(file_path), - config["model"], - config["top_k"], - ), + with self.db_manager.get_connection() as conn: + result_obj = conn.execute( + text(""" + SELECT COUNT(DISTINCT q.question_id) + FROM questions q + JOIN question_analysis qa ON qa.question_id = q.id + JOIN chunk_relevance cr ON cr.question_analysis_id = qa.id + WHERE qa.file_path = :file_path AND qa.model = :model AND qa.top_k = :top_k + """), + { + "file_path": str(file_path), + "model": config["model"], + "top_k": config["top_k"], + }, ) - count = cursor.fetchone()[0] + count = result_obj.fetchone()[0] return count > 0 except Exception as e: diff --git a/report_analyst/core/database_manager.py b/report_analyst/core/database_manager.py new file mode 100644 index 000000000..efa737247 --- /dev/null +++ b/report_analyst/core/database_manager.py @@ -0,0 +1,104 @@ +""" +Database Manager using SQLAlchemy + +Provides unified database interface for both SQLite and PostgreSQL. +""" + +import logging +import os +from contextlib import contextmanager +from pathlib import Path +from typing import Optional + +from sqlalchemy import create_engine, text +from sqlalchemy.engine import Engine +from sqlalchemy.exc import SQLAlchemyError + +logger = logging.getLogger(__name__) + + +class DatabaseManager: + """Manages database connections using SQLAlchemy.""" + + def __init__(self, database_url: Optional[str] = None): + """ + Initialize database manager. + + Args: + database_url: Database connection string. If None, uses SQLite default. + - SQLite: sqlite:///path/to/db + - PostgreSQL: postgresql://user:pass@host:port/db + """ + if database_url is None: + # Check DATABASE_URL environment variable first + database_url = os.getenv("DATABASE_URL") + if database_url is None: + # Default to SQLite + storage_path = os.getenv("STORAGE_PATH", "./storage") + db_path = str(Path(storage_path) / "cache" / "analysis.db") + # Ensure parent directory exists + Path(db_path).parent.mkdir(parents=True, exist_ok=True) + database_url = f"sqlite:///{db_path}" + + self.database_url = database_url + self._engine: Optional[Engine] = None + self._is_postgres = database_url.startswith(("postgresql://", "postgres://")) + + logger.info(f"Initializing DatabaseManager with URL: {self._mask_url(database_url)}") + logger.info(f"Database type: {'PostgreSQL' if self._is_postgres else 'SQLite'}") + + def _mask_url(self, url: str) -> str: + """Mask password in database URL for logging.""" + if "@" in url: + parts = url.split("@") + if len(parts) == 2: + user_pass = parts[0].split("://")[-1] + if ":" in user_pass: + user = user_pass.split(":")[0] + return url.replace(user_pass, f"{user}:***") + return url + + def get_engine(self) -> Engine: + """Get or create SQLAlchemy engine.""" + if self._engine is None: + # For SQLite, use check_same_thread=False for compatibility + connect_args = {} + if not self._is_postgres: + connect_args["check_same_thread"] = False + + self._engine = create_engine( + self.database_url, + connect_args=connect_args, + echo=False, # Set to True for SQL debugging + ) + logger.info("SQLAlchemy engine created") + return self._engine + + @contextmanager + def get_connection(self): + """Get database connection (context manager).""" + engine = self.get_engine() + conn = engine.connect() + try: + yield conn + conn.commit() + except Exception: + conn.rollback() + raise + finally: + conn.close() + + def execute(self, query: str, params: Optional[dict] = None): + """Execute a query and return result.""" + with self.get_connection() as conn: + result = conn.execute(text(query), params or {}) + return result + + def is_postgres(self) -> bool: + """Check if using PostgreSQL.""" + return self._is_postgres + + def is_sqlite(self) -> bool: + """Check if using SQLite.""" + return not self._is_postgres + diff --git a/report_analyst/core/database_schema.py b/report_analyst/core/database_schema.py new file mode 100644 index 000000000..5ef987e11 --- /dev/null +++ b/report_analyst/core/database_schema.py @@ -0,0 +1,118 @@ +""" +Database Schema Definitions using SQLAlchemy + +Defines all tables for the analysis cache system. +""" + +from datetime import datetime + +from sqlalchemy import ( + Boolean, + Column, + DateTime, + Float, + ForeignKey, + Integer, + LargeBinary, + MetaData, + Table, + Text, + UniqueConstraint, +) + +# Create metadata object +metadata = MetaData() + +# Document chunks table +document_chunks = Table( + "document_chunks", + metadata, + Column("id", Integer, primary_key=True, autoincrement=True), + Column("file_path", Text, nullable=False), + Column("chunk_text", Text, nullable=False), + Column("chunk_size", Integer, nullable=False), + Column("chunk_overlap", Integer, nullable=False), + Column("embedding", LargeBinary, nullable=True), # BLOB/BYTEA + Column("metadata", Text, nullable=True), # JSON stored as text + Column("created_at", DateTime, default=datetime.now), + UniqueConstraint("file_path", "chunk_text", "chunk_size", "chunk_overlap"), +) + +# Questions table +questions = Table( + "questions", + metadata, + Column("id", Integer, primary_key=True, autoincrement=True), + Column("question_id", Text, nullable=False), + Column("question_set", Text, nullable=False), + Column("question_text", Text, nullable=True), + Column("guidelines", Text, nullable=True), + UniqueConstraint("question_id", "question_set"), +) + +# Analysis cache table +analysis_cache = Table( + "analysis_cache", + metadata, + Column("id", Integer, primary_key=True, autoincrement=True), + Column("file_path", Text, nullable=False), + Column("question_id", Text, nullable=False), + Column("chunk_size", Integer, nullable=False), + Column("chunk_overlap", Integer, nullable=False), + Column("top_k", Integer, nullable=False), + Column("model", Text, nullable=False), + Column("question_set", Text, nullable=False), + Column("result", Text, nullable=False), # JSON stored as text + Column("created_at", DateTime, default=datetime.now), + UniqueConstraint( + "file_path", + "question_id", + "chunk_size", + "chunk_overlap", + "top_k", + "model", + "question_set", + ), +) + +# Question analysis table +question_analysis = Table( + "question_analysis", + metadata, + Column("id", Integer, primary_key=True, autoincrement=True), + Column("file_path", Text, nullable=False), + Column("question_id", Integer, ForeignKey("questions.id"), nullable=False), + Column("model", Text, nullable=False), + Column("top_k", Integer, nullable=False), + Column("analysis_result", Text, nullable=False), # JSON stored as text + Column("version", Integer, default=1), + Column("created_at", DateTime, default=datetime.now), + UniqueConstraint("file_path", "question_id", "model", "top_k", "version"), +) + +# Chunk relevance table +chunk_relevance = Table( + "chunk_relevance", + metadata, + Column("id", Integer, primary_key=True, autoincrement=True), + Column("question_analysis_id", Integer, ForeignKey("question_analysis.id"), nullable=False), + Column("document_chunk_id", Integer, ForeignKey("document_chunks.id"), nullable=False), + Column("chunk_order", Integer, nullable=False), + Column("similarity_score", Float, nullable=True), + Column("llm_score", Float, nullable=True), + Column("is_evidence", Boolean, nullable=False, default=False), + Column("evidence_order", Integer, nullable=True), + Column("metadata", Text, nullable=True), # JSON stored as text + UniqueConstraint("question_analysis_id", "document_chunk_id"), +) + +# Indexes (defined separately for clarity) +# Note: SQLAlchemy doesn't support "IF NOT EXISTS" in CREATE INDEX directly, +# so we'll handle these in init_db() using raw SQL +indexes = [ + # Index on file_path for document_chunks + "CREATE INDEX IF NOT EXISTS idx_file_path ON document_chunks(file_path)", + # Index on chunk parameters + "CREATE INDEX IF NOT EXISTS idx_chunk_params ON document_chunks(chunk_size, chunk_overlap)", +] + diff --git a/report_analyst_enterprise/LICENSE b/report_analyst_enterprise/LICENSE new file mode 100644 index 000000000..9964d4b3f --- /dev/null +++ b/report_analyst_enterprise/LICENSE @@ -0,0 +1,17 @@ +Climate+Tech Open License for Good + +This module (report_analyst_enterprise/) is licensed under the Climate+Tech Open License for Good. + +This license allows use for research, educational, and non-commercial purposes. +Commercial use and dual licensing options are available upon request. + +For the full text of the Climate+Tech Open License for Good, licensing inquiries, +or commercial/dual licensing options, please contact Climate+Tech: + +- https://climateandtech.com/en/climate-ai-solutions/opensustainability-analysis-framework +- https://climateandtech.com/en/research-projects/sustainability-ai-benchmark-and-dataset + +Copyright (c) 2025 Climate+Tech + +This software is part of the Open Sustainability Analysis project. + diff --git a/report_analyst_enterprise/__init__.py b/report_analyst_enterprise/__init__.py new file mode 100644 index 000000000..fd5649480 --- /dev/null +++ b/report_analyst_enterprise/__init__.py @@ -0,0 +1,9 @@ +""" +Report Analyst Enterprise Module + +Enterprise features including PostgreSQL support with pgvector. +Licensed under Climate+Tech Open License for Good. +""" + +__version__ = "0.1.0" + diff --git a/report_analyst_enterprise/database/__init__.py b/report_analyst_enterprise/database/__init__.py new file mode 100644 index 000000000..5c1b5fbce --- /dev/null +++ b/report_analyst_enterprise/database/__init__.py @@ -0,0 +1,18 @@ +""" +Database utilities for enterprise module. +""" + +from .pgvector_support import ( + check_pgvector_available, + create_vector_type, + get_vector_distance_func, + setup_pgvector_extension, +) + +__all__ = [ + "check_pgvector_available", + "create_vector_type", + "get_vector_distance_func", + "setup_pgvector_extension", +] + diff --git a/report_analyst_enterprise/database/pgvector_support.py b/report_analyst_enterprise/database/pgvector_support.py new file mode 100644 index 000000000..9a6521f48 --- /dev/null +++ b/report_analyst_enterprise/database/pgvector_support.py @@ -0,0 +1,145 @@ +""" +pgvector Support for PostgreSQL + +Provides vector type and similarity search functions for PostgreSQL with pgvector extension. +""" + +import logging +from typing import Optional + +logger = logging.getLogger(__name__) + + +def check_pgvector_available(connection) -> bool: + """ + Check if pgvector extension is available in the database. + + Args: + connection: SQLAlchemy connection object + + Returns: + True if pgvector is available, False otherwise + """ + try: + from sqlalchemy import text + result = connection.execute( + text(""" + SELECT EXISTS( + SELECT 1 FROM pg_extension WHERE extname = 'vector' + ) + """) + ) + available = result.fetchone()[0] + if available: + logger.info("pgvector extension is available") + else: + logger.info("pgvector extension is not available") + return available + except Exception as e: + logger.warning(f"Error checking pgvector availability: {e}") + return False + + +def setup_pgvector_extension(connection) -> bool: + """ + Attempt to create pgvector extension (requires superuser privileges). + + Args: + connection: SQLAlchemy connection object + + Returns: + True if extension was created or already exists, False otherwise + """ + try: + from sqlalchemy import text + # Check if already exists + if check_pgvector_available(connection): + return True + + # Try to create extension + connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + connection.commit() + logger.info("pgvector extension created successfully") + return True + except Exception as e: + logger.warning(f"Could not create pgvector extension (may require superuser): {e}") + return False + + +def create_vector_type(dimension: Optional[int] = None): + """ + Create a SQLAlchemy type for pgvector. + + Args: + dimension: Optional dimension for the vector type + + Returns: + SQLAlchemy TypeEngine for vector type + """ + from sqlalchemy import TypeDecorator, Text + + class VectorType(TypeDecorator): + """Custom type for pgvector vector column""" + impl = Text + cache_ok = True + + def load_dialect_impl(self, dialect): + if dialect.name == "postgresql": + # Use raw SQL to create vector type + return dialect.type_descriptor(Text()) + return dialect.type_descriptor(Text()) + + def process_bind_param(self, value, dialect): + if value is None: + return None + # Convert numpy array or list to string format for pgvector + if hasattr(value, "tolist"): + value = value.tolist() + # pgvector expects format: [1,2,3] + return str(value).replace(" ", "") + + def process_result_value(self, value, dialect): + if value is None: + return None + # Parse string format back to list + import ast + return ast.literal_eval(value) + + return VectorType() + + +def get_vector_distance_func(embedding_column_name: str, query_vector, distance_type: str = "cosine"): + """ + Get SQL expression for vector distance calculation. + + Args: + embedding_column_name: Name of the embedding column + query_vector: Query vector (numpy array or list) + distance_type: Type of distance ('cosine', 'l2', 'inner_product') + + Returns: + Tuple of (SQL expression string, parameters dict) for use in ORDER BY clause + """ + # Convert query vector to string format for pgvector + if hasattr(query_vector, "tolist"): + query_vector = query_vector.tolist() + query_str = str(query_vector).replace(" ", "") + + # Map distance type to pgvector operator + operators = { + "cosine": "<=>", # Cosine distance + "l2": "<->", # L2 distance + "inner_product": "<#>", # Inner product (negative) + } + + if distance_type not in operators: + logger.warning(f"Unknown distance type {distance_type}, using cosine") + distance_type = "cosine" + + operator = operators[distance_type] + + # Return SQL expression for ORDER BY + # Format: embedding <=> '[1,2,3]'::vector + sql_expr = f"{embedding_column_name} {operator} :query_vector::vector" + return sql_expr, {"query_vector": query_str} + diff --git a/report_analyst_enterprise/requirements.txt b/report_analyst_enterprise/requirements.txt new file mode 100644 index 000000000..a97656bb0 --- /dev/null +++ b/report_analyst_enterprise/requirements.txt @@ -0,0 +1,14 @@ +# Report Analyst Enterprise Module Dependencies +# This module provides PostgreSQL support with pgvector + +# Core dependency - SQLAlchemy (should match core version) +sqlalchemy>=2.0.0 + +# PostgreSQL driver +psycopg2-binary>=2.9.0 + +# pgvector support (Python client) +# Note: pgvector extension must be installed on PostgreSQL server +# For Heroku: heroku-postgresql may include pgvector +# For Vercel: May need to enable pgvector separately + diff --git a/tests/integration/test_streamlit_app.py b/tests/integration/test_streamlit_app.py index 3d42aaf9d..1ddff2658 100644 --- a/tests/integration/test_streamlit_app.py +++ b/tests/integration/test_streamlit_app.py @@ -62,26 +62,11 @@ def test_env(): cache_path = storage_path / "cache" cache_path.mkdir(parents=True) - # Create test database + # Create test database using CacheManager (which will create all tables) db_path = cache_path / "analysis.db" - conn = sqlite3.connect(db_path) - conn.execute( - """ - CREATE TABLE IF NOT EXISTS analysis_cache ( - file_path TEXT, - question_id TEXT, - chunk_size INTEGER, - chunk_overlap INTEGER, - top_k INTEGER, - model TEXT, - question_set TEXT, - result TEXT, - created_at TEXT, - PRIMARY KEY (file_path, question_id, chunk_size, chunk_overlap, top_k, model, question_set) - ) - """ - ) - conn.close() + from report_analyst.core.cache_manager import CacheManager + cache_manager = CacheManager(db_path=str(db_path)) + # Tables are created automatically by CacheManager.init_db() # Create test question set questions_dir = Path(temp_dir) / "questionsets" diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py index dea1a79c0..c88e9fb87 100644 --- a/tests/test_analyzer.py +++ b/tests/test_analyzer.py @@ -47,10 +47,10 @@ def test_db(): def clean_db(test_db): """Provide a clean database for each test function""" print(f"\nCleaning database at: {test_db}") # Debug print - conn = sqlite3.connect(str(test_db)) - conn.execute("DELETE FROM analysis_cache") - conn.commit() - conn.close() + # Use CacheManager to clean the database + from report_analyst.core.cache_manager import CacheManager + cache_manager = CacheManager(db_path=str(test_db)) + cache_manager.clear_cache() # Clear all cache return test_db diff --git a/tests/test_cache_manager.py b/tests/test_cache_manager.py index 230d19d97..b87be7d79 100644 --- a/tests/test_cache_manager.py +++ b/tests/test_cache_manager.py @@ -1,12 +1,12 @@ import json import os import shutil -import sqlite3 import tempfile from datetime import datetime from pathlib import Path import pytest +from sqlalchemy import inspect, text from report_analyst.core.cache_manager import CacheManager @@ -14,39 +14,65 @@ @pytest.fixture(autouse=True) def setup_test_env(): """Setup test environment variables""" + original_storage = os.environ.get("STORAGE_PATH") os.environ["STORAGE_PATH"] = str(Path(__file__).parent / "test_storage") yield # Cleanup after tests if Path(os.environ["STORAGE_PATH"]).exists(): import shutil - shutil.rmtree(os.environ["STORAGE_PATH"]) + # Restore original + if original_storage: + os.environ["STORAGE_PATH"] = original_storage + elif "STORAGE_PATH" in os.environ: + del os.environ["STORAGE_PATH"] @pytest.fixture def temp_db(): - """Create a temporary database for testing""" + """Create a temporary database for testing (SQLite)""" temp_dir = tempfile.mkdtemp() db_path = Path(temp_dir) / "test_cache.db" - cache_manager = CacheManager(str(db_path)) + cache_manager = CacheManager(db_path=str(db_path)) yield cache_manager shutil.rmtree(temp_dir) +@pytest.fixture(params=["sqlite", "postgres"]) +def temp_db_both(request): + """ + Create a temporary database for testing both SQLite and PostgreSQL. + For PostgreSQL, requires DATABASE_URL environment variable or skips test. + """ + temp_dir = tempfile.mkdtemp() + + if request.param == "sqlite": + db_path = Path(temp_dir) / "test_cache.db" + cache_manager = CacheManager(db_path=str(db_path)) + yield cache_manager + else: + # PostgreSQL - check if DATABASE_URL is set + database_url = os.getenv("TEST_POSTGRES_URL") + if not database_url: + pytest.skip("TEST_POSTGRES_URL not set, skipping PostgreSQL test") + cache_manager = CacheManager(database_url=database_url) + yield cache_manager + + shutil.rmtree(temp_dir) + + def test_init_db(temp_db): """Test database initialization""" - # Check if tables exist - with sqlite3.connect(temp_db.db_path) as conn: - cursor = conn.execute( - """ - SELECT name FROM sqlite_master - WHERE type='table' AND (name='analysis_cache' OR name='document_chunks') - """ - ) - tables = [row[0] for row in cursor.fetchall()] + # Check if tables exist using SQLAlchemy inspector + engine = temp_db.db_manager.get_engine() + inspector = inspect(engine) + tables = inspector.get_table_names() assert "analysis_cache" in tables assert "document_chunks" in tables + assert "questions" in tables + assert "question_analysis" in tables + assert "chunk_relevance" in tables def test_save_and_get_analysis(temp_db): @@ -211,68 +237,66 @@ def test_cache_status(temp_db): def test_get_chunks_without_embeddings(temp_db): """Test get_chunks_without_embeddings method""" import numpy as np - import sqlite3 - from datetime import datetime file_path = "test_file.pdf" chunk_size = 500 chunk_overlap = 20 # Insert chunks directly into database (some without embeddings, some with) - with sqlite3.connect(temp_db.db_path) as conn: + with temp_db.db_manager.get_connection() as conn: timestamp = datetime.now().isoformat() # Insert chunks without embeddings conn.execute( - """ - INSERT INTO document_chunks - (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?) - """, - ( - file_path, - "Chunk 1 without embedding", - chunk_size, - chunk_overlap, - None, # No embedding - json.dumps({}), - timestamp, - ), + text(""" + INSERT INTO document_chunks + (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) + VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) + """), + { + "file_path": file_path, + "chunk_text": "Chunk 1 without embedding", + "chunk_size": chunk_size, + "chunk_overlap": chunk_overlap, + "embedding": None, + "metadata": json.dumps({}), + "created_at": timestamp, + }, ) conn.execute( - """ - INSERT INTO document_chunks - (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?) - """, - ( - file_path, - "Chunk 2 without embedding", - chunk_size, - chunk_overlap, - None, # No embedding - json.dumps({}), - timestamp, - ), + text(""" + INSERT INTO document_chunks + (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) + VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) + """), + { + "file_path": file_path, + "chunk_text": "Chunk 2 without embedding", + "chunk_size": chunk_size, + "chunk_overlap": chunk_overlap, + "embedding": None, + "metadata": json.dumps({}), + "created_at": timestamp, + }, ) # Insert chunk with embedding embedding_bytes = np.array([0.1, 0.2, 0.3], dtype=np.float32).tobytes() conn.execute( - """ - INSERT INTO document_chunks - (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?) - """, - ( - file_path, - "Chunk 3 with embedding", - chunk_size, - chunk_overlap, - embedding_bytes, - json.dumps({}), - timestamp, - ), + text(""" + INSERT INTO document_chunks + (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) + VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) + """), + { + "file_path": file_path, + "chunk_text": "Chunk 3 with embedding", + "chunk_size": chunk_size, + "chunk_overlap": chunk_overlap, + "embedding": embedding_bytes, + "metadata": json.dumps({}), + "created_at": timestamp, + }, ) # Get chunks without embeddings @@ -289,8 +313,6 @@ def test_get_chunks_without_embeddings(temp_db): def test_has_chunk_scoring(temp_db): """Test has_chunk_scoring method""" import numpy as np - import sqlite3 - from datetime import datetime file_path = "test_file.pdf" config = { @@ -307,23 +329,23 @@ def test_has_chunk_scoring(temp_db): # First, save the chunk to document_chunks table so it can be referenced chunk_text = "Chunk 1" embedding_bytes = np.array([0.1, 0.2, 0.3], dtype=np.float32).tobytes() - with sqlite3.connect(temp_db.db_path) as conn: + with temp_db.db_manager.get_connection() as conn: timestamp = datetime.now().isoformat() conn.execute( - """ - INSERT INTO document_chunks - (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?) - """, - ( - file_path, - chunk_text, - config["chunk_size"], - config["chunk_overlap"], - embedding_bytes, - json.dumps({}), - timestamp, - ), + text(""" + INSERT INTO document_chunks + (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) + VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) + """), + { + "file_path": file_path, + "chunk_text": chunk_text, + "chunk_size": config["chunk_size"], + "chunk_overlap": config["chunk_overlap"], + "embedding": embedding_bytes, + "metadata": json.dumps({}), + "created_at": timestamp, + }, ) # Save analysis with chunk relevance (scoring) diff --git a/tests/test_llm_evidence_separation.py b/tests/test_llm_evidence_separation.py index f9b86df81..fc09f3cc9 100644 --- a/tests/test_llm_evidence_separation.py +++ b/tests/test_llm_evidence_separation.py @@ -15,7 +15,7 @@ import json import os -import sqlite3 +from sqlalchemy import inspect, text import sys import tempfile from pathlib import Path @@ -55,22 +55,25 @@ def cache_manager(self, temp_db): def test_database_schema_separation(self, cache_manager): """Test that the database schema has separate fields for similarity_score, llm_score, and is_evidence.""" - with sqlite3.connect(cache_manager.db_path) as conn: - cursor = conn.execute("PRAGMA table_info(chunk_relevance)") - columns = {row[1]: row[2] for row in cursor.fetchall()} + from sqlalchemy import inspect + inspector = inspect(cache_manager.db_manager.get_engine()) + columns = {col["name"]: str(col["type"]) for col in inspector.get_columns("chunk_relevance")} - # Verify required columns exist with correct types - required_columns = { - "similarity_score": "REAL", - "llm_score": "REAL", - "is_evidence": "BOOLEAN", - } + # Verify required columns exist with correct types + # SQLAlchemy type names may differ (Float vs REAL, Boolean vs BOOLEAN) + required_columns = { + "similarity_score": ["REAL", "FLOAT", "DOUBLE PRECISION"], + "llm_score": ["REAL", "FLOAT", "DOUBLE PRECISION"], + "is_evidence": ["BOOLEAN", "BOOL"], + } - for col, expected_type in required_columns.items(): - assert col in columns, f"Missing column: {col}" - assert ( - columns[col] == expected_type - ), f"Wrong type for {col}: got {columns[col]}, expected {expected_type}" + for col, expected_types in required_columns.items(): + assert col in columns, f"Missing column: {col}" + col_type = columns[col].upper() + assert any( + expected_type.upper() in col_type + for expected_type in expected_types + ), f"Column {col} has type {columns[col]}, expected one of {expected_types}" def test_independent_value_storage(self, cache_manager): """Test that LLM score and evidence can be stored independently.""" From 350307003cb589501ddc3faaf03675a8821dfc26 Mon Sep 17 00:00:00 2001 From: Christian Date: Sun, 14 Dec 2025 01:07:54 +0100 Subject: [PATCH 02/21] Fix double info icons in Settings page - Remove duplicate CSS rules that added icons to stAlert elements - Keep single icon by adding to markdown container only, not paragraphs - Prevents double icons in Database Configuration and other info messages --- report_analyst/streamlit_app.py | 426 +++++++++++++++++++++++++++++--- 1 file changed, 394 insertions(+), 32 deletions(-) diff --git a/report_analyst/streamlit_app.py b/report_analyst/streamlit_app.py index 4d3a02a86..cc73c9ec1 100644 --- a/report_analyst/streamlit_app.py +++ b/report_analyst/streamlit_app.py @@ -71,6 +71,7 @@ def log_analysis_step(message: str, level: str = "info"): # Keep relative imports from report_analyst.core.analyzer import DocumentAnalyzer +from report_analyst.core.api_key_manager import APIKeyManager from report_analyst.core.dataframe_manager import ( create_analysis_dataframes, create_combined_dataframe, @@ -274,7 +275,7 @@ def process_document( def save_uploaded_file(uploaded_file) -> Optional[str]: - """Save uploaded file to temp directory""" + """Save uploaded file to temp directory or PostgreSQL""" try: if uploaded_file is None: logger.warning("No file was uploaded") @@ -289,10 +290,50 @@ def save_uploaded_file(uploaded_file) -> Optional[str]: if file_key in st.session_state: return st.session_state[file_key] - # Otherwise, handle it as an UploadedFile + # Get file bytes + file_bytes = uploaded_file.getbuffer() + + # Check if PostgreSQL file storage is enabled + use_postgres_storage = st.session_state.get("use_postgres_file_storage", False) + + if use_postgres_storage: + try: + from report_analyst.core.file_storage import get_file_storage + + # Get database URL from session state or environment + database_url = st.session_state.get("database_url") + file_storage = get_file_storage(database_url) + + if file_storage: + # Store in PostgreSQL + file_id = file_storage.store_file( + file_bytes, + uploaded_file.name, + uploaded_file.type + ) + + # Save to temp for processing (retrieve from DB) + temp_path = file_storage.save_to_temp(file_id) + + if temp_path: + # Store both file_id and path in session state + st.session_state[file_key] = temp_path + st.session_state[f"{file_key}_id"] = file_id + logger.info(f"Stored file {uploaded_file.name} in PostgreSQL (ID: {file_id})") + st.session_state.file_processed = False + return temp_path + else: + logger.warning("Failed to save file from PostgreSQL to temp, falling back to local") + else: + logger.warning("PostgreSQL file storage not available, falling back to local") + except Exception as e: + logger.warning(f"PostgreSQL file storage failed: {str(e)}, falling back to local") + + # Fallback to local file storage file_path = Path("temp") / uploaded_file.name + file_path.parent.mkdir(parents=True, exist_ok=True) with open(file_path, "wb") as f: - f.write(uploaded_file.getbuffer()) + f.write(file_bytes) logger.info(f"Successfully saved file: {file_path}") # Store the path in session state @@ -1556,6 +1597,9 @@ def main(): if "use_s3_upload" not in st.session_state: st.session_state.use_s3_upload = os.getenv("USE_S3_UPLOAD", "false").lower() == "true" + # Sync API keys from session state to environment at startup + APIKeyManager.sync_api_keys_to_env(st.session_state) + st.set_page_config(page_title="Report Analyst", layout="wide") # Inject Material Icons link tag at the top @@ -1595,13 +1639,9 @@ def main(): margin-right: 8px; } - /* Add Material Icon to all notifications (unified - using info icon for all) */ - .stAlert [data-testid="stMarkdownContainer"]::before, - [data-testid="stNotification"] [data-testid="stMarkdownContainer"]::before, - .stInfo [data-testid="stMarkdownContainer"]::before, - .stSuccess [data-testid="stMarkdownContainer"]::before, - .stError [data-testid="stMarkdownContainer"]::before, - .stWarning [data-testid="stMarkdownContainer"]::before { + /* Add Material Icon to stAlert elements - only ONE icon per alert */ + /* Add icon only to the markdown container, NOT to paragraphs to avoid duplicates */ + [data-testid="stAlert"] [data-testid="stMarkdownContainer"]::before { content: 'info'; font-family: 'Material Icons'; font-size: 20px; @@ -1610,13 +1650,14 @@ def main(): display: inline-block; } - /* Alternative selectors for notifications without markdown container */ - .stAlert p::before, - [data-testid="stNotification"] p::before, - .stInfo p::before, - .stSuccess p::before, - .stError p::before, - .stWarning p::before { + /* Remove icons from paragraphs inside stAlert to prevent double icons */ + [data-testid="stAlert"] p::before { + content: none !important; + display: none !important; + } + + /* Add icons to custom notifications */ + [data-testid="stNotification"] [data-testid="stMarkdownContainer"]::before { content: 'info'; font-family: 'Material Icons'; font-size: 20px; @@ -1625,6 +1666,12 @@ def main(): display: inline-block; } + /* Remove icons from paragraphs in custom notifications too */ + [data-testid="stNotification"] p::before { + content: none !important; + display: none !important; + } + /* Settings expander icon in sidebar */ [data-testid="stSidebar"] [data-testid="stExpander"] summary::before { content: 'settings'; @@ -2756,8 +2803,8 @@ def main(): with st.sidebar: nav_page = option_menu( menu_title=None, - options=["Upload Report", "Report Analyst", "All Results"], - icons=["house", "file-text", "bar-chart"], + options=["Upload Report", "Report Analyst", "All Results", "Settings"], + icons=["house", "file-text", "bar-chart", "gear"], menu_icon=None, default_index=0, orientation="vertical", @@ -2784,7 +2831,7 @@ def main(): ) except ImportError: # Fallback to regular radio if package not installed - nav_options = ["Upload Report", "Report Analyst", "All Results"] + nav_options = ["Upload Report", "Report Analyst", "All Results", "Settings"] nav_page = st.sidebar.radio( "", nav_options, @@ -2792,26 +2839,241 @@ def main(): label_visibility="collapsed" ) - # Settings section in sidebar (consolidates all integration settings) - st.sidebar.markdown("---") - with st.sidebar.expander("Settings", expanded=False): - # Show Enterprise Mode status at the top - use_s3_upload = st.session_state.get("use_s3_upload", False) - if use_s3_upload and BACKEND_INTEGRATION_AVAILABLE: - st.caption("Enterprise mode") + # Show page-specific content based on navigation + if nav_page == "Settings": + st.title("Settings") + st.caption("Configure application settings and integrations") + + # Open Source Modules Section + st.header("Open Source Modules") + st.caption("Core features available in the open source edition") + + # API Keys Configuration + st.subheader("API Keys") + st.caption("Enter your API keys to enable LLM features. Keys are stored in session state only and not persisted.") + + # Check if keys exist in environment but not in session state + env_openai_key = os.getenv("OPENAI_API_KEY") + env_google_key = os.getenv("GOOGLE_API_KEY") + session_openai_key = st.session_state.get("api_key_openai_api_key") + session_google_key = st.session_state.get("api_key_google_api_key") + + has_env_openai = env_openai_key and not session_openai_key + has_env_google = env_google_key and not session_google_key + + # Get current values (from session state or environment) + current_openai_key = APIKeyManager.get_api_key("OPENAI_API_KEY", st.session_state) + current_google_key = APIKeyManager.get_api_key("GOOGLE_API_KEY", st.session_state) + + # OpenAI API Key section + with st.expander("OpenAI API Key", expanded=not has_env_openai or st.session_state.get("override_openai_key", False)): + # Show status for OpenAI key + if has_env_openai and not st.session_state.get("override_openai_key", False): + st.info("API key is set from environment variable") + if st.button("Override with new key", key="btn_override_openai"): + st.session_state.override_openai_key = True + st.rerun() + elif current_openai_key: + masked_openai = f"{current_openai_key[:8]}...{current_openai_key[-4:]}" if len(current_openai_key) > 12 else "***" + st.caption(f"Current key: `{masked_openai}`") + + # Track override state + override_openai = st.session_state.get("override_openai_key", False) + + if not has_env_openai or override_openai: + # Track previous values to detect changes + prev_openai_key = st.session_state.get("prev_openai_key", current_openai_key) + + # OpenAI API Key input + openai_key_input = st.text_input( + "OpenAI API Key", + value="", # Never show the actual key in the input + type="password", + key="openai_api_key_input", + help="Enter your OpenAI API key to use GPT models. Leave empty to use existing key from environment.", + placeholder="sk-..." if not current_openai_key else "Enter new key to update" + ) + + # Update API key if user entered a new value (different from current) + if openai_key_input and openai_key_input != current_openai_key: + APIKeyManager.set_api_key("OPENAI_API_KEY", openai_key_input, st.session_state) + st.session_state.prev_openai_key = openai_key_input + st.session_state.override_openai_key = False # Reset override state + st.success("OpenAI API key updated") + elif openai_key_input == "" and current_openai_key and not has_env_openai: + # User cleared the input - keep existing key (only if not from env) + st.session_state.prev_openai_key = current_openai_key + else: + st.session_state.prev_openai_key = current_openai_key + + # Cancel override button + if override_openai: + if st.button("Cancel Override", key="cancel_override_openai"): + st.session_state.override_openai_key = False + st.rerun() - # Enterprise Integration (S3+NATS) - always shown first, outside of backend config + # Google/Gemini API Key section + with st.expander("Google/Gemini API Key", expanded=not has_env_google or st.session_state.get("override_google_key", False)): + # Show status for Google key + if has_env_google and not st.session_state.get("override_google_key", False): + st.info("API key is set from environment variable") + if st.button("Override with new key", key="btn_override_google"): + st.session_state.override_google_key = True + st.rerun() + elif current_google_key: + masked_google = f"{current_google_key[:8]}...{current_google_key[-4:]}" if len(current_google_key) > 12 else "***" + st.caption(f"Current key: `{masked_google}`") + + # Track override state + override_google = st.session_state.get("override_google_key", False) + + if not has_env_google or override_google: + # Track previous values to detect changes + prev_google_key = st.session_state.get("prev_google_key", current_google_key) + + # Google/Gemini API Key input + google_key_input = st.text_input( + "Google/Gemini API Key", + value="", # Never show the actual key in the input + type="password", + key="google_api_key_input", + help="Enter your Google API key to use Gemini models. Leave empty to use existing key from environment.", + placeholder="Enter your Google API key" if not current_google_key else "Enter new key to update" + ) + + # Update API key if user entered a new value (different from current) + if google_key_input and google_key_input != current_google_key: + APIKeyManager.set_api_key("GOOGLE_API_KEY", google_key_input, st.session_state) + st.session_state.prev_google_key = google_key_input + st.session_state.override_google_key = False # Reset override state + st.success("Google API key updated") + elif google_key_input == "" and current_google_key and not has_env_google: + # User cleared the input - keep existing key (only if not from env) + st.session_state.prev_google_key = current_google_key + else: + st.session_state.prev_google_key = current_google_key + + # Cancel override button + if override_google: + if st.button("Cancel Override", key="cancel_override_google"): + st.session_state.override_google_key = False + st.rerun() + + # Show clear button if key exists (only for session state keys, not env) + if current_google_key and not has_env_google: + if st.button("Clear Google Key", key="clear_google_key"): + APIKeyManager.set_api_key("GOOGLE_API_KEY", None, st.session_state) + st.rerun() + + # Show clear button for OpenAI if key exists (only for session state keys, not env) + if current_openai_key and not has_env_openai: + if st.button("Clear OpenAI Key", key="clear_openai_key"): + APIKeyManager.set_api_key("OPENAI_API_KEY", None, st.session_state) + st.rerun() + + st.divider() + + # Database Configuration (read-only, from environment variables) + st.subheader("Database Configuration") + + # Get database URL from environment or default + database_url = os.getenv("DATABASE_URL") + if database_url is None: + # Default to SQLite + storage_path = os.getenv("STORAGE_PATH", "./storage") + db_path = str(Path(storage_path) / "cache" / "analysis.db") + database_url = f"sqlite:///{db_path}" + database_type = "SQLite" + st.info(f"**Type:** {database_type}\n\n**Path:** `{db_path}`\n\n*Configure via `STORAGE_PATH` environment variable*") + else: + # Parse PostgreSQL URL to show connection details (masked) + database_type = "PostgreSQL" + try: + # Mask password in display + if "@" in database_url: + parts = database_url.split("@") + if len(parts) == 2: + user_pass = parts[0].split("://")[-1] + if ":" in user_pass: + user = user_pass.split(":")[0] + masked_url = database_url.replace(user_pass, f"{user}:***") + else: + masked_url = database_url + else: + masked_url = database_url + else: + masked_url = database_url + + # Extract connection details for display + if "postgresql://" in database_url or "postgres://" in database_url: + url_part = database_url.split("://")[-1] + if "@" in url_part: + auth, host_db = url_part.split("@") + user = auth.split(":")[0] if ":" in auth else auth + if ":" in host_db: + host, port_db = host_db.split(":") + if "/" in port_db: + port, db = port_db.split("/", 1) + else: + port = port_db + db = "?" + else: + if "/" in host_db: + host, db = host_db.split("/", 1) + port = "5432" + else: + host = host_db + port = "5432" + db = "?" + + st.info(f"**Type:** {database_type}\n\n**Host:** `{host}`\n**Port:** `{port}`\n**Database:** `{db}`\n**User:** `{user}`\n\n*Configure via `DATABASE_URL` environment variable*") + else: + st.info(f"**Type:** {database_type}\n\n**Connection:** `{masked_url}`\n\n*Configure via `DATABASE_URL` environment variable*") + else: + st.info(f"**Type:** {database_type}\n\n**Connection:** `{masked_url}`\n\n*Configure via `DATABASE_URL` environment variable*") + except Exception: + st.info(f"**Type:** {database_type}\n\n**Connection:** `{masked_url}`\n\n*Configure via `DATABASE_URL` environment variable*") + + # Store in session state for use by DocumentAnalyzer + st.session_state.database_url = database_url + + st.divider() + st.divider() + + # Enterprise Modules Section + st.header("Enterprise Modules") + st.caption("Features available in the enterprise edition") + + # Enterprise Integration (S3+NATS) st.subheader("Enterprise Integration") + # In Streamlit, when a widget has a 'key', it automatically syncs with session state + # The widget's return value is the current value from session state (or default if not set) + # IMPORTANT: Don't provide 'value' parameter when using 'key' - let Streamlit manage it + # The widget return value is the source of truth for the current render + st.markdown(""" + + """, unsafe_allow_html=True) use_s3_upload = st.checkbox( "Enable S3+NATS Upload", - value=st.session_state.use_s3_upload, key="use_s3_upload", help="Upload documents via S3 and process via NATS for enterprise integration", ) + # Show Enterprise Mode status only if checkbox is checked AND backend is available + # Check AFTER widget render - use widget return value which reflects current state + # The widget return value is the authoritative source for the current render cycle + if use_s3_upload and BACKEND_INTEGRATION_AVAILABLE: + st.info("Enterprise mode enabled") + st.divider() - # Backend Integration + # Backend Integration (Enterprise feature) if BACKEND_INTEGRATION_AVAILABLE: config = configure_backend_integration() # Store config in session state for access across pages @@ -2821,6 +3083,36 @@ def main(): st.warning("Backend integration modules not available") config = None st.session_state.backend_config = None + + st.divider() + + # File Storage Configuration (Enterprise feature) + st.subheader("File Storage") + st.caption("Configure where uploaded files are stored (Enterprise feature)") + + # Get database URL from session state (set above in Database Configuration) + database_url_enterprise = st.session_state.get("database_url") + is_postgres_enterprise = database_url_enterprise and database_url_enterprise.startswith(("postgresql://", "postgres://")) + + # Initialize use_postgres_file_storage from session state or env + if "use_postgres_file_storage" not in st.session_state: + st.session_state.use_postgres_file_storage = os.getenv("USE_POSTGRES_FILE_STORAGE", "false").lower() == "true" + + if is_postgres_enterprise: + use_postgres_storage = st.checkbox( + "Store files in PostgreSQL", + value=st.session_state.get("use_postgres_file_storage", False), + key="use_postgres_file_storage", + help="Store uploaded files in PostgreSQL database (useful for Heroku deployments). Files are stored as BYTEA/BLOB. This is an enterprise feature.", + ) + + if use_postgres_storage: + st.info("📦 Files will be stored in PostgreSQL database") + else: + st.caption("Files will be stored in local temp directory") + else: + st.info("PostgreSQL file storage requires a PostgreSQL database. Currently using SQLite.") + st.caption("Files are stored in local temp directory") # Show page-specific content based on navigation if nav_page == "Report Analyst": @@ -3683,6 +3975,34 @@ def main(): """, unsafe_allow_html=True) + # Try to import JSON Schema form component (enterprise feature) + try: + # Use the proper Streamlit custom component + from report_analyst_enterprise.components.streamlit_component.backend import json_schema_form + import json + # Path is already imported at the top of the file + + JSON_SCHEMA_FORM_AVAILABLE = True + + # Load PDF upload schema + schema_path = Path(__file__).parent.parent / "report_analyst_enterprise" / "components" / "schemas" / "pdf_upload_schema.json" + ui_schema_path = Path(__file__).parent.parent / "report_analyst_enterprise" / "components" / "schemas" / "pdf_upload_ui_schema.json" + + if schema_path.exists() and ui_schema_path.exists(): + with open(schema_path) as f: + pdf_upload_schema = json.load(f) + with open(ui_schema_path) as f: + pdf_upload_ui_schema = json.load(f) + else: + JSON_SCHEMA_FORM_AVAILABLE = False + pdf_upload_schema = None + pdf_upload_ui_schema = None + except ImportError: + JSON_SCHEMA_FORM_AVAILABLE = False + pdf_upload_schema = None + pdf_upload_ui_schema = None + + # File upload with optional metadata form uploaded_file = st.file_uploader( "Choose a PDF file", type="pdf", @@ -3690,6 +4010,48 @@ def main(): help="Limit 200MB per file • PDF" ) + # Show metadata form if JSON Schema form is available + pdf_metadata = None + company_metadata = None + + if JSON_SCHEMA_FORM_AVAILABLE: + # ESRS Company Information Form + esrs_schema_path = Path(__file__).parent.parent / "report_analyst_enterprise" / "components" / "schemas" / "esrs_company_schema.json" + esrs_ui_schema_path = Path(__file__).parent.parent / "report_analyst_enterprise" / "components" / "schemas" / "esrs_company_ui_schema.json" + + if esrs_schema_path.exists() and esrs_ui_schema_path.exists(): + with open(esrs_schema_path) as f: + esrs_company_schema = json.load(f) + with open(esrs_ui_schema_path) as f: + esrs_company_ui_schema = json.load(f) + + with st.expander("ESRS Company Information", expanded=True): + st.caption("Enter company data aligned with ESRS XBRL taxonomy requirements") + company_metadata = json_schema_form( + schema=esrs_company_schema, + ui_schema=esrs_company_ui_schema, + key="esrs_company_form", + height=700 + ) + if company_metadata and company_metadata.get("type") == "submit": + st.success("Company information saved!") + st.session_state.esrs_company_metadata = company_metadata.get("formData", company_metadata) + + # Basic PDF metadata form + if pdf_upload_schema: + with st.expander("Add Document Metadata (Optional)", expanded=False): + st.caption("Add metadata like category, tags, and description to help organize your documents.") + pdf_metadata = json_schema_form( + schema=pdf_upload_schema, + ui_schema=pdf_upload_ui_schema, + key="pdf_metadata_form", + height=500 + ) + if pdf_metadata: + st.success("Metadata saved!") + # Store in session state for use after upload + st.session_state.pdf_metadata = pdf_metadata + if uploaded_file: # Handle upload based on mode if use_s3_upload and BACKEND_INTEGRATION_AVAILABLE: @@ -3791,9 +4153,9 @@ def main(): ) if selected_set: - # Show question set description + # Show question set description (without info icon to avoid double icons) if selected_set in question_sets: - st.info(question_sets[selected_set]["description"]) + st.caption(question_sets[selected_set]["description"]) # Only show consolidated results display_consolidated_results(analyzer, selected_set) From 0ede76f651cea9ba20c217c8edf53778e8d8a9da Mon Sep 17 00:00:00 2001 From: Christian Date: Sun, 14 Dec 2025 02:02:00 +0100 Subject: [PATCH 03/21] Add PostgreSQL migration system with Alembic - Set up Alembic for database migrations - Create initial migration with all tables and indexes - Add hybrid approach: auto-creation for dev, Alembic for production - Update CacheManager and FileStorage to check USE_ALEMBIC_MIGRATIONS - Add migration utilities for status checking and migration execution - Add Heroku support with Procfile release phase for auto-migrations - Create migration script for manual execution - Add comprehensive documentation in DATABASE_MIGRATIONS.md --- Procfile | 5 + alembic.ini | 120 ++++++++++ alembic/README | 1 + alembic/env.py | 110 ++++++++++ alembic/script.py.mako | 28 +++ alembic/versions/001_initial_schema.py | 132 +++++++++++ docs/DATABASE_MIGRATIONS.md | 289 +++++++++++++++++++++++++ migrations/run_migrations.sh | 30 +++ report_analyst/core/cache_manager.py | 8 +- report_analyst/core/file_storage.py | 252 +++++++++++++++++++++ report_analyst/core/migration_utils.py | 168 ++++++++++++++ 11 files changed, 1142 insertions(+), 1 deletion(-) create mode 100644 Procfile create mode 100644 alembic.ini create mode 100644 alembic/README create mode 100644 alembic/env.py create mode 100644 alembic/script.py.mako create mode 100644 alembic/versions/001_initial_schema.py create mode 100644 docs/DATABASE_MIGRATIONS.md create mode 100755 migrations/run_migrations.sh create mode 100644 report_analyst/core/file_storage.py create mode 100644 report_analyst/core/migration_utils.py diff --git a/Procfile b/Procfile new file mode 100644 index 000000000..9e44064de --- /dev/null +++ b/Procfile @@ -0,0 +1,5 @@ +# Heroku Procfile +# Release phase runs migrations before the web dyno starts +release: python -m alembic upgrade head || echo "Migrations skipped (USE_ALEMBIC_MIGRATIONS not enabled or not PostgreSQL)" +web: streamlit run report_analyst/streamlit_app.py --server.port=$PORT --server.address=0.0.0.0 + diff --git a/alembic.ini b/alembic.ini new file mode 100644 index 000000000..14215507d --- /dev/null +++ b/alembic.ini @@ -0,0 +1,120 @@ +# A generic, single database configuration. + +[alembic] +# path to migration scripts +# Use forward slashes (/) also on windows to provide an os agnostic path +script_location = alembic + +# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s +# Uncomment the line below if you want the files to be prepended with date and time +# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file +# for all available tokens +# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s + +# sys.path path, will be prepended to sys.path if present. +# defaults to the current working directory. +prepend_sys_path = . + +# timezone to use when rendering the date within the migration file +# as well as the filename. +# If specified, requires the python>=3.9 or backports.zoneinfo library and tzdata library. +# Any required deps can installed by adding `alembic[tz]` to the pip requirements +# string value is passed to ZoneInfo() +# leave blank for localtime +# timezone = + +# max length of characters to apply to the "slug" field +# truncate_slug_length = 40 + +# set to 'true' to run the environment during +# the 'revision' command, regardless of autogenerate +# revision_environment = false + +# set to 'true' to allow .pyc and .pyo files without +# a source .py file to be detected as revisions in the +# versions/ directory +# sourceless = false + +# version location specification; This defaults +# to alembic/versions. When using multiple version +# directories, initial revisions must be specified with --version-path. +# The path separator used here should be the separator specified by "version_path_separator" below. +# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions + +# version path separator; As mentioned above, this is the character used to split +# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep. +# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas. +# Valid values for version_path_separator are: +# +# version_path_separator = : +# version_path_separator = ; +# version_path_separator = space +# version_path_separator = newline +# +# Use os.pathsep. Default configuration used for new projects. +version_path_separator = os + +# set to 'true' to search source files recursively +# in each "version_locations" directory +# new in Alembic version 1.10 +# recursive_version_locations = false + +# the output encoding used when revision files +# are written from script.py.mako +# output_encoding = utf-8 + +# sqlalchemy.url is set dynamically in alembic/env.py from DATABASE_URL environment variable +# sqlalchemy.url = driver://user:pass@localhost/dbname + + +[post_write_hooks] +# post_write_hooks defines scripts or Python functions that are run +# on newly generated revision scripts. See the documentation for further +# detail and examples + +# format using "black" - use the console_scripts runner, against the "black" entrypoint +# hooks = black +# black.type = console_scripts +# black.entrypoint = black +# black.options = -l 79 REVISION_SCRIPT_FILENAME + +# lint with attempts to fix using "ruff" - use the exec runner, execute a binary +# hooks = ruff +# ruff.type = exec +# ruff.executable = %(here)s/.venv/bin/ruff +# ruff.options = check --fix REVISION_SCRIPT_FILENAME + +# Logging configuration +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARNING +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARNING +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/alembic/README b/alembic/README new file mode 100644 index 000000000..98e4f9c44 --- /dev/null +++ b/alembic/README @@ -0,0 +1 @@ +Generic single-database configuration. \ No newline at end of file diff --git a/alembic/env.py b/alembic/env.py new file mode 100644 index 000000000..24757d161 --- /dev/null +++ b/alembic/env.py @@ -0,0 +1,110 @@ +import os +from logging.config import fileConfig + +from sqlalchemy import engine_from_config, pool + +from alembic import context + +# Import our database schema +from report_analyst.core.database_schema import metadata as db_metadata +from report_analyst.core.database_manager import DatabaseManager + +# this is the Alembic Config object, which provides +# access to the values within the .ini file in use. +config = context.config + +# Interpret the config file for Python logging. +# This line sets up loggers basically. +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +# Get database URL from environment or use DatabaseManager +database_url = os.getenv("DATABASE_URL") +if database_url is None: + # Use DatabaseManager to get default URL (SQLite) + db_manager = DatabaseManager() + database_url = db_manager.database_url +else: + # Validate it's a valid URL format + db_manager = DatabaseManager(database_url) + +# Set the database URL in config +config.set_main_option("sqlalchemy.url", database_url) + +# Combine metadata from database_schema and file_storage +from sqlalchemy import MetaData, Table, Column, String, Text, DateTime, LargeBinary +from datetime import datetime + +# Use the database_schema metadata as base +target_metadata = db_metadata + +# Add stored_files table to the metadata (from file_storage.py) +# This table is created dynamically in file_storage, so we define it here for migrations +if "stored_files" not in target_metadata.tables: + stored_files = Table( + "stored_files", + target_metadata, + Column("id", String(36), primary_key=True), # UUID as string + Column("filename", Text, nullable=False), + Column("file_data", LargeBinary, nullable=False), # BYTEA in PostgreSQL + Column("content_type", Text, nullable=True), + Column("file_size", Text, nullable=False), # Store as string for large files + Column("created_at", DateTime, default=datetime.now), + ) + +# other values from the config, defined by the needs of env.py, +# can be acquired: +# my_important_option = config.get_main_option("my_important_option") +# ... etc. + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode. + + This configures the context with just a URL + and not an Engine, though an Engine is acceptable + here as well. By skipping the Engine creation + we don't even need a DBAPI to be available. + + Calls to context.execute() here emit the given string to the + script output. + + """ + url = config.get_main_option("sqlalchemy.url") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations in 'online' mode. + + In this scenario we need to create an Engine + and associate a connection with the context. + + """ + connectable = engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure( + connection=connection, target_metadata=target_metadata + ) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/alembic/script.py.mako b/alembic/script.py.mako new file mode 100644 index 000000000..480b130d6 --- /dev/null +++ b/alembic/script.py.mako @@ -0,0 +1,28 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: Union[str, None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + """Upgrade schema.""" + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + """Downgrade schema.""" + ${downgrades if downgrades else "pass"} diff --git a/alembic/versions/001_initial_schema.py b/alembic/versions/001_initial_schema.py new file mode 100644 index 000000000..679e7935a --- /dev/null +++ b/alembic/versions/001_initial_schema.py @@ -0,0 +1,132 @@ +"""initial_schema + +Revision ID: 001_initial_schema +Revises: +Create Date: 2025-12-14 01:00:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from datetime import datetime + + +# revision identifiers, used by Alembic. +revision: str = '001_initial_schema' +down_revision: Union[str, None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Create initial database schema.""" + # Document chunks table + op.create_table( + 'document_chunks', + sa.Column('id', sa.Integer(), nullable=False, autoincrement=True), + sa.Column('file_path', sa.Text(), nullable=False), + sa.Column('chunk_text', sa.Text(), nullable=False), + sa.Column('chunk_size', sa.Integer(), nullable=False), + sa.Column('chunk_overlap', sa.Integer(), nullable=False), + sa.Column('embedding', sa.LargeBinary(), nullable=True), + sa.Column('metadata', sa.Text(), nullable=True), + sa.Column('created_at', sa.DateTime(), default=datetime.now), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('file_path', 'chunk_text', 'chunk_size', 'chunk_overlap') + ) + + # Questions table + op.create_table( + 'questions', + sa.Column('id', sa.Integer(), nullable=False, autoincrement=True), + sa.Column('question_id', sa.Text(), nullable=False), + sa.Column('question_set', sa.Text(), nullable=False), + sa.Column('question_text', sa.Text(), nullable=True), + sa.Column('guidelines', sa.Text(), nullable=True), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('question_id', 'question_set') + ) + + # Analysis cache table + op.create_table( + 'analysis_cache', + sa.Column('id', sa.Integer(), nullable=False, autoincrement=True), + sa.Column('file_path', sa.Text(), nullable=False), + sa.Column('question_id', sa.Text(), nullable=False), + sa.Column('chunk_size', sa.Integer(), nullable=False), + sa.Column('chunk_overlap', sa.Integer(), nullable=False), + sa.Column('top_k', sa.Integer(), nullable=False), + sa.Column('model', sa.Text(), nullable=False), + sa.Column('question_set', sa.Text(), nullable=False), + sa.Column('result', sa.Text(), nullable=False), + sa.Column('created_at', sa.DateTime(), default=datetime.now), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint( + 'file_path', 'question_id', 'chunk_size', 'chunk_overlap', + 'top_k', 'model', 'question_set' + ) + ) + + # Question analysis table + op.create_table( + 'question_analysis', + sa.Column('id', sa.Integer(), nullable=False, autoincrement=True), + sa.Column('file_path', sa.Text(), nullable=False), + sa.Column('question_id', sa.Integer(), nullable=False), + sa.Column('model', sa.Text(), nullable=False), + sa.Column('top_k', sa.Integer(), nullable=False), + sa.Column('analysis_result', sa.Text(), nullable=False), + sa.Column('version', sa.Integer(), default=1), + sa.Column('created_at', sa.DateTime(), default=datetime.now), + sa.ForeignKeyConstraint(['question_id'], ['questions.id']), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('file_path', 'question_id', 'model', 'top_k', 'version') + ) + + # Chunk relevance table + op.create_table( + 'chunk_relevance', + sa.Column('id', sa.Integer(), nullable=False, autoincrement=True), + sa.Column('question_analysis_id', sa.Integer(), nullable=False), + sa.Column('document_chunk_id', sa.Integer(), nullable=False), + sa.Column('chunk_order', sa.Integer(), nullable=False), + sa.Column('similarity_score', sa.Float(), nullable=True), + sa.Column('llm_score', sa.Float(), nullable=True), + sa.Column('is_evidence', sa.Boolean(), nullable=False, server_default='0'), + sa.Column('evidence_order', sa.Integer(), nullable=True), + sa.Column('metadata', sa.Text(), nullable=True), + sa.ForeignKeyConstraint(['question_analysis_id'], ['question_analysis.id']), + sa.ForeignKeyConstraint(['document_chunk_id'], ['document_chunks.id']), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('question_analysis_id', 'document_chunk_id') + ) + + # Stored files table (for PostgreSQL file storage) + op.create_table( + 'stored_files', + sa.Column('id', sa.String(length=36), nullable=False), + sa.Column('filename', sa.Text(), nullable=False), + sa.Column('file_data', sa.LargeBinary(), nullable=False), + sa.Column('content_type', sa.Text(), nullable=True), + sa.Column('file_size', sa.Text(), nullable=False), + sa.Column('created_at', sa.DateTime(), default=datetime.now), + sa.PrimaryKeyConstraint('id') + ) + + # Create indexes + op.create_index('idx_file_path', 'document_chunks', ['file_path'], unique=False) + op.create_index('idx_chunk_params', 'document_chunks', ['chunk_size', 'chunk_overlap'], unique=False) + + +def downgrade() -> None: + """Drop all tables.""" + op.drop_index('idx_chunk_params', table_name='document_chunks') + op.drop_index('idx_file_path', table_name='document_chunks') + op.drop_table('chunk_relevance') + op.drop_table('question_analysis') + op.drop_table('analysis_cache') + op.drop_table('questions') + op.drop_table('document_chunks') + op.drop_table('stored_files') + diff --git a/docs/DATABASE_MIGRATIONS.md b/docs/DATABASE_MIGRATIONS.md new file mode 100644 index 000000000..671e20876 --- /dev/null +++ b/docs/DATABASE_MIGRATIONS.md @@ -0,0 +1,289 @@ +# Database Migrations Guide + +This guide explains how database migrations work in the Report Analyst application, covering both development and production scenarios. + +## Overview + +The application uses a **hybrid migration approach**: +- **Development (SQLite)**: Automatic table creation on first connection (simple, fast) +- **Production (PostgreSQL)**: Alembic migrations for version control and rollback capability + +## Migration System + +### Alembic + +We use [Alembic](https://alembic.sqlalchemy.org/) for database migrations in production. Alembic provides: +- Version control for database schema changes +- Rollback capability (`alembic downgrade`) +- Team collaboration (shared migration history) +- Safe production deployments + +### Environment Variables + +- `DATABASE_URL`: Database connection string (required) + - SQLite: `sqlite:///path/to/db` + - PostgreSQL: `postgresql://user:pass@host:port/db` +- `USE_ALEMBIC_MIGRATIONS`: Enable Alembic migrations (default: `false`) + - Set to `true` for production PostgreSQL deployments + - Set to `false` for development (uses auto-creation) +- `ALEMBIC_CONFIG`: Path to `alembic.ini` (optional, default: `./alembic.ini`) + +## Development Workflow + +### SQLite (Default) + +When using SQLite (default for development), tables are created automatically: + +```python +from report_analyst.core.cache_manager import CacheManager + +# Tables are created automatically on first connection +cache_manager = CacheManager() +# No migration needed! +``` + +**No action required** - just start using the application. + +### Local PostgreSQL (Optional) + +If you want to test with PostgreSQL locally: + +1. **Set up PostgreSQL database:** + ```bash + createdb report_analyst + ``` + +2. **Set environment variable:** + ```bash + export DATABASE_URL="postgresql://user:pass@localhost/report_analyst" + ``` + +3. **Choose migration approach:** + + **Option A: Auto-creation (development)** + ```bash + # Don't set USE_ALEMBIC_MIGRATIONS (or set to false) + # Tables created automatically + ``` + + **Option B: Alembic migrations (production-like)** + ```bash + export USE_ALEMBIC_MIGRATIONS=true + alembic upgrade head + ``` + +## Production Workflow + +### Heroku Deployment + +#### Automatic Migrations (Recommended) + +Migrations run automatically on deploy via the `release` phase in `Procfile`: + +```procfile +release: python -m alembic upgrade head || echo "Migrations skipped..." +web: streamlit run report_analyst/streamlit_app.py --server.port=$PORT --server.address=0.0.0.0 +``` + +**Setup:** +1. Set environment variables on Heroku: + ```bash + heroku config:set DATABASE_URL="postgresql://..." + heroku config:set USE_ALEMBIC_MIGRATIONS=true + ``` + +2. Deploy: + ```bash + git push heroku main + # Migrations run automatically before web dyno starts + ``` + +#### Manual Migrations + +If you prefer to run migrations manually: + +```bash +# Run migrations manually +heroku run alembic upgrade head + +# Or use the migration script +heroku run bash migrations/run_migrations.sh +``` + +### Other Deployments + +For other platforms (Render, Railway, etc.): + +1. **Set environment variables:** + ```bash + DATABASE_URL=postgresql://... + USE_ALEMBIC_MIGRATIONS=true + ``` + +2. **Run migrations before starting the app:** + ```bash + # In your deployment script or startup command + python -m alembic upgrade head + python -m streamlit run report_analyst/streamlit_app.py + ``` + +## Migration Commands + +### Check Migration Status + +```bash +# Check current revision +alembic current + +# Check if migrations are needed +python -c "from report_analyst.core.migration_utils import check_migration_status; print(check_migration_status())" +``` + +### Create New Migration + +```bash +# Auto-generate migration from schema changes +alembic revision --autogenerate -m "description_of_changes" + +# Review the generated migration file in alembic/versions/ +# Edit if needed, then apply: +alembic upgrade head +``` + +### Apply Migrations + +```bash +# Upgrade to latest (head) +alembic upgrade head + +# Upgrade to specific revision +alembic upgrade + +# Upgrade one step +alembic upgrade +1 +``` + +### Rollback Migrations + +```bash +# Downgrade one step +alembic downgrade -1 + +# Downgrade to specific revision +alembic downgrade + +# Downgrade to base (removes all tables - use with caution!) +alembic downgrade base +``` + +### View Migration History + +```bash +# List all migrations +alembic history + +# Show current revision +alembic current + +# Show detailed history +alembic history --verbose +``` + +## Migration Files + +Migrations are stored in `alembic/versions/`: +- Format: `_.py` +- Example: `001_initial_schema.py` + +### Initial Migration + +The initial migration (`001_initial_schema.py`) creates: +- `document_chunks` - Document text chunks with embeddings +- `questions` - Question definitions +- `analysis_cache` - Cached analysis results +- `question_analysis` - Analysis results per question +- `chunk_relevance` - Chunk relevance scores +- `stored_files` - File storage (PostgreSQL only) + +## Troubleshooting + +### Migration Fails on Heroku + +**Problem:** `alembic: command not found` + +**Solution:** Ensure `alembic` is in `requirements.txt`: +```bash +pip freeze | grep alembic >> requirements.txt +``` + +### Tables Already Exist + +**Problem:** Migration fails because tables already exist + +**Solution:** +- If using auto-creation, don't set `USE_ALEMBIC_MIGRATIONS=true` +- If using migrations, mark current state as migrated: + ```bash + alembic stamp head + ``` + +### Migration Out of Sync + +**Problem:** Database schema doesn't match migrations + +**Solution:** +1. Check current state: `alembic current` +2. Check head: `alembic history` +3. If needed, create a new migration to sync: + ```bash + alembic revision --autogenerate -m "sync_schema" + ``` + +### Development vs Production Mismatch + +**Problem:** Different behavior in dev vs production + +**Solution:** +- Development: `USE_ALEMBIC_MIGRATIONS=false` (or unset) +- Production: `USE_ALEMBIC_MIGRATIONS=true` + +## Best Practices + +1. **Always review auto-generated migrations** before applying +2. **Test migrations locally** before deploying to production +3. **Backup database** before running migrations in production +4. **Use descriptive migration names**: `add_user_table`, `add_index_to_chunks` +5. **Keep migrations small and focused** - one logical change per migration +6. **Never edit applied migrations** - create new ones instead + +## Migration Utilities + +The `report_analyst.core.migration_utils` module provides helper functions: + +```python +from report_analyst.core.migration_utils import ( + check_migration_status, + get_current_revision, + get_head_revision, + needs_migration, + run_migrations, +) + +# Check if migration is needed +if needs_migration(): + print("Database needs migration") + run_migrations() + +# Get detailed status +status = check_migration_status() +print(f"Current: {status['current_revision']}") +print(f"Head: {status['head_revision']}") +print(f"Up to date: {status['is_up_to_date']}") +``` + +## Additional Resources + +- [Alembic Documentation](https://alembic.sqlalchemy.org/) +- [SQLAlchemy Migrations Guide](https://docs.sqlalchemy.org/en/20/core/metadata.html) +- [Heroku Release Phase](https://devcenter.heroku.com/articles/release-phase) + diff --git a/migrations/run_migrations.sh b/migrations/run_migrations.sh new file mode 100755 index 000000000..f0e232a46 --- /dev/null +++ b/migrations/run_migrations.sh @@ -0,0 +1,30 @@ +#!/bin/bash +# Script to run database migrations on Heroku or other deployments + +set -e + +echo "Running database migrations..." + +# Check if DATABASE_URL is set +if [ -z "$DATABASE_URL" ]; then + echo "Error: DATABASE_URL environment variable is not set" + exit 1 +fi + +# Check if USE_ALEMBIC_MIGRATIONS is enabled +if [ "$USE_ALEMBIC_MIGRATIONS" != "true" ]; then + echo "Warning: USE_ALEMBIC_MIGRATIONS is not set to 'true'. Skipping migrations." + echo "Set USE_ALEMBIC_MIGRATIONS=true to enable Alembic migrations." + exit 0 +fi + +# Run migrations +python -m alembic upgrade head + +if [ $? -eq 0 ]; then + echo "Migrations completed successfully" +else + echo "Error: Migrations failed" + exit 1 +fi + diff --git a/report_analyst/core/cache_manager.py b/report_analyst/core/cache_manager.py index c54d4e76f..dafd987c2 100644 --- a/report_analyst/core/cache_manager.py +++ b/report_analyst/core/cache_manager.py @@ -47,7 +47,13 @@ def __init__(self, db_path: str = None, database_url: str = None): # Initialize database manager self.db_manager = DatabaseManager(db_url) logger.info(f"Initializing CacheManager with database: {self.db_manager._mask_url(db_url)}") - self.init_db() + + # Check if we should use Alembic migrations instead of auto-creation + use_alembic = os.getenv("USE_ALEMBIC_MIGRATIONS", "false").lower() == "true" + if use_alembic and self.db_manager.is_postgres(): + logger.info("Using Alembic migrations - skipping auto table creation") + else: + self.init_db() # In-memory vector store for current document self.vector_store = None diff --git a/report_analyst/core/file_storage.py b/report_analyst/core/file_storage.py new file mode 100644 index 000000000..478fad5a2 --- /dev/null +++ b/report_analyst/core/file_storage.py @@ -0,0 +1,252 @@ +""" +File Storage Service + +Provides file storage abstraction with support for: +- Local filesystem (default) +- PostgreSQL (for Heroku deployments) +- S3 (enterprise feature) +""" + +import logging +import os +import uuid +from pathlib import Path +from typing import Optional + +from sqlalchemy import LargeBinary, MetaData, Table, Column, Text, DateTime, String +from sqlalchemy import text +from datetime import datetime + +from .database_manager import DatabaseManager + +logger = logging.getLogger(__name__) + + +class FileStorageError(Exception): + """Exception raised for file storage errors""" + pass + + +class PostgreSQLFileStorage: + """Store files in PostgreSQL database as BYTEA/BLOB""" + + def __init__(self, database_url: Optional[str] = None): + """ + Initialize PostgreSQL file storage. + + Args: + database_url: Database connection string. If None, uses environment variable. + """ + if database_url is None: + database_url = os.getenv("DATABASE_URL") + if database_url is None: + raise FileStorageError("DATABASE_URL not set for PostgreSQL file storage") + + self.db_manager = DatabaseManager(database_url) + if not self.db_manager.is_postgres(): + raise FileStorageError("PostgreSQL file storage requires PostgreSQL database") + + # Check if we should use Alembic migrations instead of auto-creation + use_alembic = os.getenv("USE_ALEMBIC_MIGRATIONS", "false").lower() == "true" + if use_alembic: + logger.info("Using Alembic migrations - skipping auto table creation for stored_files") + else: + self._init_table() + + def _init_table(self): + """Initialize the stored_files table""" + try: + metadata = MetaData() + stored_files = Table( + "stored_files", + metadata, + Column("id", String(36), primary_key=True), # UUID as string + Column("filename", Text, nullable=False), + Column("file_data", LargeBinary, nullable=False), # BYTEA in PostgreSQL + Column("content_type", Text, nullable=True), + Column("file_size", Text, nullable=False), # Store as string for large files + Column("created_at", DateTime, default=datetime.now), + ) + + engine = self.db_manager.get_engine() + metadata.create_all(engine, checkfirst=True) + logger.info("stored_files table initialized") + except Exception as e: + logger.error(f"Error initializing stored_files table: {str(e)}") + raise FileStorageError(f"Failed to initialize file storage table: {str(e)}") + + def store_file(self, file_bytes: bytes, filename: str, content_type: Optional[str] = None) -> str: + """ + Store file in PostgreSQL and return file ID. + + Args: + file_bytes: File content as bytes + filename: Original filename + content_type: MIME type (optional) + + Returns: + file_id: Unique identifier for the stored file + """ + try: + file_id = str(uuid.uuid4()) + file_size = len(file_bytes) + + with self.db_manager.get_connection() as conn: + # Use parameterized query for safety + query = text(""" + INSERT INTO stored_files (id, filename, file_data, content_type, file_size, created_at) + VALUES (:id, :filename, :file_data, :content_type, :file_size, :created_at) + """) + conn.execute(query, { + "id": file_id, + "filename": filename, + "file_data": file_bytes, + "content_type": content_type or "application/pdf", + "file_size": str(file_size), + "created_at": datetime.now() + }) + conn.commit() + + logger.info(f"Stored file {filename} (ID: {file_id}, size: {file_size} bytes) in PostgreSQL") + return file_id + except Exception as e: + logger.error(f"Error storing file in PostgreSQL: {str(e)}") + raise FileStorageError(f"Failed to store file: {str(e)}") + + def retrieve_file(self, file_id: str) -> Optional[bytes]: + """ + Retrieve file from PostgreSQL. + + Args: + file_id: Unique identifier for the stored file + + Returns: + File content as bytes, or None if not found + """ + try: + with self.db_manager.get_connection() as conn: + query = text("SELECT file_data FROM stored_files WHERE id = :file_id") + result = conn.execute(query, {"file_id": file_id}) + row = result.fetchone() + + if row: + return bytes(row[0]) + return None + except Exception as e: + logger.error(f"Error retrieving file {file_id} from PostgreSQL: {str(e)}") + raise FileStorageError(f"Failed to retrieve file: {str(e)}") + + def get_file_info(self, file_id: str) -> Optional[dict]: + """ + Get file metadata without retrieving the file data. + + Args: + file_id: Unique identifier for the stored file + + Returns: + Dictionary with file info, or None if not found + """ + try: + with self.db_manager.get_connection() as conn: + query = text(""" + SELECT filename, content_type, file_size, created_at + FROM stored_files WHERE id = :file_id + """) + result = conn.execute(query, {"file_id": file_id}) + row = result.fetchone() + + if row: + return { + "filename": row[0], + "content_type": row[1], + "file_size": int(row[2]) if row[2] else 0, + "created_at": row[3] + } + return None + except Exception as e: + logger.error(f"Error getting file info for {file_id}: {str(e)}") + return None + + def delete_file(self, file_id: str) -> bool: + """ + Delete file from PostgreSQL. + + Args: + file_id: Unique identifier for the stored file + + Returns: + True if deleted, False if not found + """ + try: + with self.db_manager.get_connection() as conn: + query = text("DELETE FROM stored_files WHERE id = :file_id") + result = conn.execute(query, {"file_id": file_id}) + conn.commit() + return result.rowcount > 0 + except Exception as e: + logger.error(f"Error deleting file {file_id}: {str(e)}") + return False + + def save_to_temp(self, file_id: str, temp_dir: Path = Path("temp")) -> Optional[str]: + """ + Retrieve file from PostgreSQL and save to temporary directory. + + Args: + file_id: Unique identifier for the stored file + temp_dir: Directory to save the file to + + Returns: + Path to the temporary file, or None if not found + """ + try: + file_info = self.get_file_info(file_id) + if not file_info: + return None + + file_bytes = self.retrieve_file(file_id) + if not file_bytes: + return None + + # Create temp directory if it doesn't exist + temp_dir.mkdir(parents=True, exist_ok=True) + + # Save to temp file + temp_path = temp_dir / file_info["filename"] + with open(temp_path, "wb") as f: + f.write(file_bytes) + + logger.info(f"Retrieved file {file_id} to {temp_path}") + return str(temp_path) + except Exception as e: + logger.error(f"Error saving file {file_id} to temp: {str(e)}") + return None + + +def get_file_storage(database_url: Optional[str] = None) -> Optional[PostgreSQLFileStorage]: + """ + Get file storage instance if PostgreSQL is configured. + + Args: + database_url: Database connection string (optional) + + Returns: + PostgreSQLFileStorage instance if PostgreSQL is configured, None otherwise + """ + try: + # Check if PostgreSQL file storage is enabled + use_postgres_storage = os.getenv("USE_POSTGRES_FILE_STORAGE", "false").lower() == "true" + if not use_postgres_storage: + return None + + # Check if we have a PostgreSQL database + if database_url is None: + database_url = os.getenv("DATABASE_URL") + + if database_url and database_url.startswith(("postgresql://", "postgres://")): + return PostgreSQLFileStorage(database_url) + + return None + except Exception as e: + logger.warning(f"PostgreSQL file storage not available: {str(e)}") + return None + diff --git a/report_analyst/core/migration_utils.py b/report_analyst/core/migration_utils.py new file mode 100644 index 000000000..2985de3ee --- /dev/null +++ b/report_analyst/core/migration_utils.py @@ -0,0 +1,168 @@ +""" +Migration Utilities + +Helper functions for checking migration status and managing database migrations. +""" + +import logging +import os +from typing import Optional + +from alembic import command +from alembic.config import Config +from alembic.runtime.migration import MigrationContext +from alembic.script import ScriptDirectory + +from .database_manager import DatabaseManager + +logger = logging.getLogger(__name__) + + +def get_alembic_config(database_url: Optional[str] = None) -> Config: + """ + Get Alembic configuration. + + Args: + database_url: Database connection string. If None, uses environment variable. + + Returns: + Alembic Config object + """ + alembic_ini_path = os.getenv("ALEMBIC_CONFIG", "alembic.ini") + config = Config(alembic_ini_path) + + # Set database URL if provided + if database_url: + config.set_main_option("sqlalchemy.url", database_url) + elif os.getenv("DATABASE_URL"): + config.set_main_option("sqlalchemy.url", os.getenv("DATABASE_URL")) + + return config + + +def get_current_revision(database_url: Optional[str] = None) -> Optional[str]: + """ + Get the current database revision. + + Args: + database_url: Database connection string. If None, uses environment variable. + + Returns: + Current revision string, or None if no migrations have been applied + """ + try: + config = get_alembic_config(database_url) + + # Get database manager to create engine + if database_url: + db_manager = DatabaseManager(database_url) + else: + db_manager = DatabaseManager() + + engine = db_manager.get_engine() + + with engine.connect() as connection: + context = MigrationContext.configure(connection) + current_rev = context.get_current_revision() + return current_rev + except Exception as e: + logger.error(f"Error getting current revision: {str(e)}") + return None + + +def get_head_revision() -> Optional[str]: + """ + Get the head (latest) migration revision. + + Returns: + Head revision string, or None if no migrations exist + """ + try: + config = get_alembic_config() + script = ScriptDirectory.from_config(config) + head = script.get_current_head() + return head + except Exception as e: + logger.error(f"Error getting head revision: {str(e)}") + return None + + +def needs_migration(database_url: Optional[str] = None) -> bool: + """ + Check if database needs migration. + + Args: + database_url: Database connection string. If None, uses environment variable. + + Returns: + True if migration is needed, False otherwise + """ + try: + current_rev = get_current_revision(database_url) + head_rev = get_head_revision() + + if head_rev is None: + # No migrations exist + return False + + if current_rev is None: + # Database has no migrations applied + return True + + return current_rev != head_rev + except Exception as e: + logger.error(f"Error checking migration status: {str(e)}") + return False + + +def check_migration_status(database_url: Optional[str] = None) -> dict: + """ + Check migration status and return detailed information. + + Args: + database_url: Database connection string. If None, uses environment variable. + + Returns: + Dictionary with migration status information + """ + try: + current_rev = get_current_revision(database_url) + head_rev = get_head_revision() + + return { + "current_revision": current_rev, + "head_revision": head_rev, + "is_up_to_date": current_rev == head_rev if (current_rev and head_rev) else False, + "needs_migration": needs_migration(database_url), + } + except Exception as e: + logger.error(f"Error checking migration status: {str(e)}") + return { + "current_revision": None, + "head_revision": None, + "is_up_to_date": False, + "needs_migration": False, + "error": str(e), + } + + +def run_migrations(database_url: Optional[str] = None, revision: str = "head") -> bool: + """ + Run database migrations. + + Args: + database_url: Database connection string. If None, uses environment variable. + revision: Target revision (default: "head") + + Returns: + True if successful, False otherwise + """ + try: + config = get_alembic_config(database_url) + command.upgrade(config, revision) + logger.info(f"Successfully upgraded database to revision: {revision}") + return True + except Exception as e: + logger.error(f"Error running migrations: {str(e)}") + return False + From 49b1bb9bd8b7925ed38f138ccf4c58dc9ea8eb9a Mon Sep 17 00:00:00 2001 From: Christian Date: Sun, 14 Dec 2025 02:06:48 +0100 Subject: [PATCH 04/21] Add database URL support, API key management, and PostgreSQL file storage - Update DocumentAnalyzer to accept database_url from environment or Streamlit session state - Add APIKeyManager service for managing API keys in session state without persistence - Add PostgreSQLFileStorage for storing files in PostgreSQL (Heroku deployments) - Add checkbox styling fixes in backend config - Add comprehensive tests for API key manager, file storage, and settings --- report_analyst/core/analyzer.py | 13 +- report_analyst/core/api_key_manager.py | 81 ++++++++ report_analyst/core/file_storage.py | 252 ++++++++++++++++++++++++ report_analyst_search_backend/config.py | 9 + tests/test_api_key_manager.py | 192 ++++++++++++++++++ tests/test_file_storage.py | 97 +++++++++ tests/test_settings_enterprise_mode.py | 95 +++++++++ 7 files changed, 738 insertions(+), 1 deletion(-) create mode 100644 report_analyst/core/api_key_manager.py create mode 100644 report_analyst/core/file_storage.py create mode 100644 tests/test_api_key_manager.py create mode 100644 tests/test_file_storage.py create mode 100644 tests/test_settings_enterprise_mode.py diff --git a/report_analyst/core/analyzer.py b/report_analyst/core/analyzer.py index b6822a1b1..9cca04c47 100644 --- a/report_analyst/core/analyzer.py +++ b/report_analyst/core/analyzer.py @@ -232,7 +232,18 @@ def __init__(self): # Add a cache for loaded answers self._answers_cache = {} - self.cache_manager = CacheManager() + # Get database URL from environment or session state (for Streamlit) + database_url = os.getenv("DATABASE_URL") + if database_url is None: + # Try to get from session state if available (Streamlit context) + try: + import streamlit as st + database_url = st.session_state.get("database_url") + except (ImportError, RuntimeError): + # Not in Streamlit context or streamlit not available + pass + + self.cache_manager = CacheManager(database_url=database_url) logger.info("Initialized DocumentAnalyzer with cache manager") self._initialized = True diff --git a/report_analyst/core/api_key_manager.py b/report_analyst/core/api_key_manager.py new file mode 100644 index 000000000..41e48fe0a --- /dev/null +++ b/report_analyst/core/api_key_manager.py @@ -0,0 +1,81 @@ +""" +API Key Management Service + +Manages API keys in session state and environment variables without storing them persistently. +Keys are only kept in memory for the current session. +""" + +import os +from typing import Optional + + +class APIKeyManager: + """Service to manage API keys in session state and environment without persistence""" + + @staticmethod + def set_api_key(key_name: str, value: Optional[str], session_state: dict) -> None: + """ + Set an API key in session state and environment variables. + + Args: + key_name: The environment variable name (e.g., 'OPENAI_API_KEY') + value: The API key value (None to clear) + session_state: Streamlit session state dictionary + """ + # Store in session state (temporary, per session) + session_key = f"api_key_{key_name.lower()}" + if value: + session_state[session_key] = value + # Set in environment for current process + os.environ[key_name] = value + else: + # Clear from session state + if session_key in session_state: + del session_state[session_key] + # Remove from environment if it was set by us + if key_name in os.environ: + # Only remove if it was set in this session (not from .env file) + # We can't easily track this, so we'll leave env vars that might + # have been set externally. The session state value takes precedence. + pass + + @staticmethod + def get_api_key(key_name: str, session_state: dict) -> Optional[str]: + """ + Get an API key from session state or environment. + + Priority: + 1. Session state (user-entered value) + 2. Environment variable (from .env or system) + + Args: + key_name: The environment variable name (e.g., 'OPENAI_API_KEY') + session_state: Streamlit session state dictionary + + Returns: + The API key value or None + """ + session_key = f"api_key_{key_name.lower()}" + # Check session state first (user-entered value takes precedence) + if session_key in session_state: + return session_state[session_key] + # Fall back to environment variable + return os.getenv(key_name) + + @staticmethod + def sync_api_keys_to_env(session_state: dict) -> None: + """ + Sync all API keys from session state to environment variables. + Called at startup to ensure environment has current values. + + Args: + session_state: Streamlit session state dictionary + """ + api_key_names = ["OPENAI_API_KEY", "GOOGLE_API_KEY"] + for key_name in api_key_names: + session_key = f"api_key_{key_name.lower()}" + if session_key in session_state: + value = session_state[session_key] + if value: + os.environ[key_name] = value + diff --git a/report_analyst/core/file_storage.py b/report_analyst/core/file_storage.py new file mode 100644 index 000000000..478fad5a2 --- /dev/null +++ b/report_analyst/core/file_storage.py @@ -0,0 +1,252 @@ +""" +File Storage Service + +Provides file storage abstraction with support for: +- Local filesystem (default) +- PostgreSQL (for Heroku deployments) +- S3 (enterprise feature) +""" + +import logging +import os +import uuid +from pathlib import Path +from typing import Optional + +from sqlalchemy import LargeBinary, MetaData, Table, Column, Text, DateTime, String +from sqlalchemy import text +from datetime import datetime + +from .database_manager import DatabaseManager + +logger = logging.getLogger(__name__) + + +class FileStorageError(Exception): + """Exception raised for file storage errors""" + pass + + +class PostgreSQLFileStorage: + """Store files in PostgreSQL database as BYTEA/BLOB""" + + def __init__(self, database_url: Optional[str] = None): + """ + Initialize PostgreSQL file storage. + + Args: + database_url: Database connection string. If None, uses environment variable. + """ + if database_url is None: + database_url = os.getenv("DATABASE_URL") + if database_url is None: + raise FileStorageError("DATABASE_URL not set for PostgreSQL file storage") + + self.db_manager = DatabaseManager(database_url) + if not self.db_manager.is_postgres(): + raise FileStorageError("PostgreSQL file storage requires PostgreSQL database") + + # Check if we should use Alembic migrations instead of auto-creation + use_alembic = os.getenv("USE_ALEMBIC_MIGRATIONS", "false").lower() == "true" + if use_alembic: + logger.info("Using Alembic migrations - skipping auto table creation for stored_files") + else: + self._init_table() + + def _init_table(self): + """Initialize the stored_files table""" + try: + metadata = MetaData() + stored_files = Table( + "stored_files", + metadata, + Column("id", String(36), primary_key=True), # UUID as string + Column("filename", Text, nullable=False), + Column("file_data", LargeBinary, nullable=False), # BYTEA in PostgreSQL + Column("content_type", Text, nullable=True), + Column("file_size", Text, nullable=False), # Store as string for large files + Column("created_at", DateTime, default=datetime.now), + ) + + engine = self.db_manager.get_engine() + metadata.create_all(engine, checkfirst=True) + logger.info("stored_files table initialized") + except Exception as e: + logger.error(f"Error initializing stored_files table: {str(e)}") + raise FileStorageError(f"Failed to initialize file storage table: {str(e)}") + + def store_file(self, file_bytes: bytes, filename: str, content_type: Optional[str] = None) -> str: + """ + Store file in PostgreSQL and return file ID. + + Args: + file_bytes: File content as bytes + filename: Original filename + content_type: MIME type (optional) + + Returns: + file_id: Unique identifier for the stored file + """ + try: + file_id = str(uuid.uuid4()) + file_size = len(file_bytes) + + with self.db_manager.get_connection() as conn: + # Use parameterized query for safety + query = text(""" + INSERT INTO stored_files (id, filename, file_data, content_type, file_size, created_at) + VALUES (:id, :filename, :file_data, :content_type, :file_size, :created_at) + """) + conn.execute(query, { + "id": file_id, + "filename": filename, + "file_data": file_bytes, + "content_type": content_type or "application/pdf", + "file_size": str(file_size), + "created_at": datetime.now() + }) + conn.commit() + + logger.info(f"Stored file {filename} (ID: {file_id}, size: {file_size} bytes) in PostgreSQL") + return file_id + except Exception as e: + logger.error(f"Error storing file in PostgreSQL: {str(e)}") + raise FileStorageError(f"Failed to store file: {str(e)}") + + def retrieve_file(self, file_id: str) -> Optional[bytes]: + """ + Retrieve file from PostgreSQL. + + Args: + file_id: Unique identifier for the stored file + + Returns: + File content as bytes, or None if not found + """ + try: + with self.db_manager.get_connection() as conn: + query = text("SELECT file_data FROM stored_files WHERE id = :file_id") + result = conn.execute(query, {"file_id": file_id}) + row = result.fetchone() + + if row: + return bytes(row[0]) + return None + except Exception as e: + logger.error(f"Error retrieving file {file_id} from PostgreSQL: {str(e)}") + raise FileStorageError(f"Failed to retrieve file: {str(e)}") + + def get_file_info(self, file_id: str) -> Optional[dict]: + """ + Get file metadata without retrieving the file data. + + Args: + file_id: Unique identifier for the stored file + + Returns: + Dictionary with file info, or None if not found + """ + try: + with self.db_manager.get_connection() as conn: + query = text(""" + SELECT filename, content_type, file_size, created_at + FROM stored_files WHERE id = :file_id + """) + result = conn.execute(query, {"file_id": file_id}) + row = result.fetchone() + + if row: + return { + "filename": row[0], + "content_type": row[1], + "file_size": int(row[2]) if row[2] else 0, + "created_at": row[3] + } + return None + except Exception as e: + logger.error(f"Error getting file info for {file_id}: {str(e)}") + return None + + def delete_file(self, file_id: str) -> bool: + """ + Delete file from PostgreSQL. + + Args: + file_id: Unique identifier for the stored file + + Returns: + True if deleted, False if not found + """ + try: + with self.db_manager.get_connection() as conn: + query = text("DELETE FROM stored_files WHERE id = :file_id") + result = conn.execute(query, {"file_id": file_id}) + conn.commit() + return result.rowcount > 0 + except Exception as e: + logger.error(f"Error deleting file {file_id}: {str(e)}") + return False + + def save_to_temp(self, file_id: str, temp_dir: Path = Path("temp")) -> Optional[str]: + """ + Retrieve file from PostgreSQL and save to temporary directory. + + Args: + file_id: Unique identifier for the stored file + temp_dir: Directory to save the file to + + Returns: + Path to the temporary file, or None if not found + """ + try: + file_info = self.get_file_info(file_id) + if not file_info: + return None + + file_bytes = self.retrieve_file(file_id) + if not file_bytes: + return None + + # Create temp directory if it doesn't exist + temp_dir.mkdir(parents=True, exist_ok=True) + + # Save to temp file + temp_path = temp_dir / file_info["filename"] + with open(temp_path, "wb") as f: + f.write(file_bytes) + + logger.info(f"Retrieved file {file_id} to {temp_path}") + return str(temp_path) + except Exception as e: + logger.error(f"Error saving file {file_id} to temp: {str(e)}") + return None + + +def get_file_storage(database_url: Optional[str] = None) -> Optional[PostgreSQLFileStorage]: + """ + Get file storage instance if PostgreSQL is configured. + + Args: + database_url: Database connection string (optional) + + Returns: + PostgreSQLFileStorage instance if PostgreSQL is configured, None otherwise + """ + try: + # Check if PostgreSQL file storage is enabled + use_postgres_storage = os.getenv("USE_POSTGRES_FILE_STORAGE", "false").lower() == "true" + if not use_postgres_storage: + return None + + # Check if we have a PostgreSQL database + if database_url is None: + database_url = os.getenv("DATABASE_URL") + + if database_url and database_url.startswith(("postgresql://", "postgres://")): + return PostgreSQLFileStorage(database_url) + + return None + except Exception as e: + logger.warning(f"PostgreSQL file storage not available: {str(e)}") + return None + diff --git a/report_analyst_search_backend/config.py b/report_analyst_search_backend/config.py index 2edde0fee..e9fb73965 100644 --- a/report_analyst_search_backend/config.py +++ b/report_analyst_search_backend/config.py @@ -78,6 +78,15 @@ def configure_backend_integration() -> BackendConfig: # Note: Enterprise Integration (S3+NATS) is now shown in the main Settings section above # Basic backend toggle + st.markdown(""" + + """, unsafe_allow_html=True) use_backend = st.checkbox( "Use Search Backend", value=False, diff --git a/tests/test_api_key_manager.py b/tests/test_api_key_manager.py new file mode 100644 index 000000000..8808293d9 --- /dev/null +++ b/tests/test_api_key_manager.py @@ -0,0 +1,192 @@ +""" +Tests for API Key Manager service. + +Tests that API keys are managed correctly in session state and environment +without being persisted to disk. +""" + +import os +from report_analyst.core.api_key_manager import APIKeyManager + + +def test_set_and_get_api_key(): + """Test setting and getting API keys""" + session_state = {} + + # Set an API key + APIKeyManager.set_api_key("OPENAI_API_KEY", "test-key-123", session_state) + + # Verify it's in session state + assert session_state["api_key_openai_api_key"] == "test-key-123" + + # Verify it's in environment + assert os.getenv("OPENAI_API_KEY") == "test-key-123" + + # Get the key back + retrieved_key = APIKeyManager.get_api_key("OPENAI_API_KEY", session_state) + assert retrieved_key == "test-key-123" + + +def test_get_api_key_from_env(): + """Test getting API key from environment when not in session state""" + session_state = {} + + # Set in environment directly + os.environ["OPENAI_API_KEY"] = "env-key-456" + + # Get should retrieve from environment + retrieved_key = APIKeyManager.get_api_key("OPENAI_API_KEY", session_state) + assert retrieved_key == "env-key-456" + + # Clean up + del os.environ["OPENAI_API_KEY"] + + +def test_session_state_takes_precedence(): + """Test that session state value takes precedence over environment""" + session_state = {} + + # Set in environment + os.environ["OPENAI_API_KEY"] = "env-key-789" + + # Set in session state + APIKeyManager.set_api_key("OPENAI_API_KEY", "session-key-789", session_state) + + # Get should return session state value + retrieved_key = APIKeyManager.get_api_key("OPENAI_API_KEY", session_state) + assert retrieved_key == "session-key-789" + + # Clean up + del os.environ["OPENAI_API_KEY"] + + +def test_clear_api_key(): + """Test clearing an API key""" + session_state = {} + + # Set a key + APIKeyManager.set_api_key("OPENAI_API_KEY", "test-key-clear", session_state) + assert session_state["api_key_openai_api_key"] == "test-key-clear" + + # Clear it + APIKeyManager.set_api_key("OPENAI_API_KEY", None, session_state) + + # Should be removed from session state + assert "api_key_openai_api_key" not in session_state + + +def test_sync_api_keys_to_env(): + """Test syncing API keys from session state to environment""" + session_state = { + "api_key_openai_api_key": "synced-openai-key", + "api_key_google_api_key": "synced-google-key" + } + + # Sync to environment + APIKeyManager.sync_api_keys_to_env(session_state) + + # Verify they're in environment + assert os.getenv("OPENAI_API_KEY") == "synced-openai-key" + assert os.getenv("GOOGLE_API_KEY") == "synced-google-key" + + # Clean up + del os.environ["OPENAI_API_KEY"] + del os.environ["GOOGLE_API_KEY"] + + +def test_sync_only_existing_keys(): + """Test that sync only affects keys in session state""" + session_state = { + "api_key_openai_api_key": "only-openai-key" + } + + # Set a Google key in environment + os.environ["GOOGLE_API_KEY"] = "existing-google-key" + + # Sync + APIKeyManager.sync_api_keys_to_env(session_state) + + # OpenAI should be synced + assert os.getenv("OPENAI_API_KEY") == "only-openai-key" + + # Google should remain unchanged (not in session state) + assert os.getenv("GOOGLE_API_KEY") == "existing-google-key" + + # Clean up + del os.environ["OPENAI_API_KEY"] + del os.environ["GOOGLE_API_KEY"] + + +def test_api_key_is_used_by_llm_provider(): + """Test that API key set via APIKeyManager is actually used by LLM providers""" + import report_analyst.core.llm_providers as llm_module + + session_state = {} + test_openai_key = "test-openai-key-12345" + test_google_key = "test-google-key-67890" + + # Set keys via APIKeyManager + APIKeyManager.set_api_key("OPENAI_API_KEY", test_openai_key, session_state) + APIKeyManager.set_api_key("GOOGLE_API_KEY", test_google_key, session_state) + + # Verify keys are in environment + assert os.getenv("OPENAI_API_KEY") == test_openai_key + assert os.getenv("GOOGLE_API_KEY") == test_google_key + + # Verify that get_llm reads from os.getenv (which now has our key) + # We can't actually initialize the LLM without a real key, but we can verify + # that the function will use the key we set + original_getenv = llm_module.os.getenv + + # Track what keys are accessed + accessed_keys = [] + def tracking_getenv(key, default=None): + if key in ["OPENAI_API_KEY", "GOOGLE_API_KEY"]: + accessed_keys.append(key) + return original_getenv(key, default) + + # Temporarily replace os.getenv in the module + llm_module.os.getenv = tracking_getenv + + try: + # Try to get an OpenAI LLM (will fail without real key, but we can check it reads the env) + try: + from report_analyst.core.llm_providers import get_llm + get_llm("gpt-4o-mini") + except (ValueError, Exception): + # Expected to fail, but we check that it tried to read OPENAI_API_KEY + pass + + # Verify it tried to read the API key from environment + assert "OPENAI_API_KEY" in accessed_keys, "get_llm should read OPENAI_API_KEY from environment" + finally: + # Restore original + llm_module.os.getenv = original_getenv + + # Clean up + if "OPENAI_API_KEY" in os.environ: + del os.environ["OPENAI_API_KEY"] + if "GOOGLE_API_KEY" in os.environ: + del os.environ["GOOGLE_API_KEY"] + + +def test_api_key_from_session_state_overrides_env(): + """Test that API key in session state overrides environment variable for LLM""" + session_state = {} + + # Set key in environment + os.environ["OPENAI_API_KEY"] = "env-key-111" + + # Set different key in session state + APIKeyManager.set_api_key("OPENAI_API_KEY", "session-key-222", session_state) + + # Environment should now have session state value + assert os.getenv("OPENAI_API_KEY") == "session-key-222" + + # Get should return session state value + retrieved = APIKeyManager.get_api_key("OPENAI_API_KEY", session_state) + assert retrieved == "session-key-222" + + # Clean up + del os.environ["OPENAI_API_KEY"] + diff --git a/tests/test_file_storage.py b/tests/test_file_storage.py new file mode 100644 index 000000000..e4e99ebf6 --- /dev/null +++ b/tests/test_file_storage.py @@ -0,0 +1,97 @@ +""" +Tests for PostgreSQL file storage service. + +Tests that files can be stored and retrieved from PostgreSQL. +""" + +import os +import pytest +from report_analyst.core.file_storage import PostgreSQLFileStorage, FileStorageError, get_file_storage + + +@pytest.mark.skipif( + not os.getenv("DATABASE_URL") or not os.getenv("DATABASE_URL").startswith(("postgresql://", "postgres://")), + reason="PostgreSQL not configured" +) +def test_postgres_file_storage_store_and_retrieve(): + """Test storing and retrieving a file from PostgreSQL""" + database_url = os.getenv("DATABASE_URL") + storage = PostgreSQLFileStorage(database_url) + + # Test file content + test_content = b"Test file content for PostgreSQL storage" + filename = "test_file.pdf" + + # Store file + file_id = storage.store_file(test_content, filename, "application/pdf") + assert file_id is not None + assert len(file_id) == 36 # UUID length + + # Retrieve file + retrieved_content = storage.retrieve_file(file_id) + assert retrieved_content == test_content + + # Get file info + file_info = storage.get_file_info(file_id) + assert file_info is not None + assert file_info["filename"] == filename + assert file_info["content_type"] == "application/pdf" + assert file_info["file_size"] == len(test_content) + + # Clean up + storage.delete_file(file_id) + + +@pytest.mark.skipif( + not os.getenv("DATABASE_URL") or not os.getenv("DATABASE_URL").startswith(("postgresql://", "postgres://")), + reason="PostgreSQL not configured" +) +def test_postgres_file_storage_delete(): + """Test deleting a file from PostgreSQL""" + database_url = os.getenv("DATABASE_URL") + storage = PostgreSQLFileStorage(database_url) + + # Store file + test_content = b"Test file to delete" + file_id = storage.store_file(test_content, "delete_test.pdf") + + # Delete file + deleted = storage.delete_file(file_id) + assert deleted is True + + # Verify file is gone + retrieved = storage.retrieve_file(file_id) + assert retrieved is None + + +def test_get_file_storage_without_postgres(): + """Test that get_file_storage returns None when PostgreSQL is not configured""" + # Temporarily unset DATABASE_URL + original_url = os.environ.get("DATABASE_URL") + if "DATABASE_URL" in os.environ: + del os.environ["DATABASE_URL"] + + # Unset USE_POSTGRES_FILE_STORAGE + original_setting = os.environ.get("USE_POSTGRES_FILE_STORAGE") + if "USE_POSTGRES_FILE_STORAGE" in os.environ: + del os.environ["USE_POSTGRES_FILE_STORAGE"] + + try: + storage = get_file_storage() + assert storage is None + finally: + # Restore original values + if original_url: + os.environ["DATABASE_URL"] = original_url + if original_setting: + os.environ["USE_POSTGRES_FILE_STORAGE"] = original_setting + + +def test_postgres_file_storage_requires_postgres(): + """Test that PostgreSQLFileStorage raises error for SQLite""" + # Use SQLite URL + sqlite_url = "sqlite:///test.db" + + with pytest.raises(FileStorageError, match="PostgreSQL"): + PostgreSQLFileStorage(sqlite_url) + diff --git a/tests/test_settings_enterprise_mode.py b/tests/test_settings_enterprise_mode.py new file mode 100644 index 000000000..46404e68c --- /dev/null +++ b/tests/test_settings_enterprise_mode.py @@ -0,0 +1,95 @@ +""" +Test for Settings page enterprise mode checkbox behavior. + +Tests that the "Enterprise mode enabled" message only appears when: +1. The checkbox is checked +2. Backend integration is available +""" +from streamlit.testing.v1 import AppTest + + +def test_enterprise_mode_message_only_when_checked(): + """Test that enterprise mode message only shows when checkbox is checked""" + at = AppTest.from_file("report_analyst/streamlit_app.py") + at.run(timeout=10) + + # Navigate to Settings page + at.session_state["nav_page"] = "Settings" + at.run(timeout=10) + + # Check that Settings page loaded + assert "Settings" in str(at), "Settings page should be visible" + + # Initially, checkbox should be unchecked (default False) + # Find the checkbox + checkboxes = [w for w in at.checkbox if "S3+NATS" in str(w)] + assert len(checkboxes) > 0, "S3+NATS checkbox should exist" + + checkbox = checkboxes[0] + initial_value = checkbox.value + + # If checkbox is checked initially (from env var), uncheck it + if initial_value: + checkbox.set_value(False) + at.run(timeout=10) + + # After unchecking, the message should NOT appear + page_text = str(at) + if "Enterprise mode enabled" in page_text: + # This is the bug - message appears even when unchecked + assert False, "Enterprise mode message should not appear when checkbox is unchecked" + + # Now check the checkbox + checkbox.set_value(True) + at.run(timeout=10) + + # After checking, if backend is available, message should appear + page_text_after = str(at) + # Note: We can't easily test backend availability in AppTest, so we just check + # that the checkbox state is correctly reflected + assert checkbox.value == True, "Checkbox should be checked" + + +def test_enterprise_mode_checkbox_state_persistence(): + """Test that checkbox state persists correctly across reruns""" + at = AppTest.from_file("report_analyst/streamlit_app.py") + at.run(timeout=10) + + # Navigate to Settings + at.session_state["nav_page"] = "Settings" + at.run(timeout=10) + + # Find checkbox + checkboxes = [w for w in at.checkbox if "S3+NATS" in str(w)] + assert len(checkboxes) > 0, "S3+NATS checkbox should exist" + checkbox = checkboxes[0] + + # Set to True + checkbox.set_value(True) + at.run(timeout=10) + assert checkbox.value == True, "Checkbox should be True after setting" + # Note: AppTest session_state doesn't support .get(), access directly + try: + assert at.session_state["use_s3_upload"] == True, "Session state should be True" + except (KeyError, AttributeError): + pass # Session state might use widget ID instead of key + + # Set to False + checkbox.set_value(False) + at.run(timeout=10) + assert checkbox.value == False, "Checkbox should be False after unchecking" + + # Check that enterprise mode message is NOT shown when unchecked + page_text = str(at) + if "Enterprise mode enabled" in page_text: + assert False, "Enterprise mode message should NOT appear when checkbox is unchecked" + + # Rerun - state should persist + at.run(timeout=10) + assert checkbox.value == False, "Checkbox should remain False after rerun" + + # Verify message still doesn't appear after rerun + page_text_rerun = str(at) + if "Enterprise mode enabled" in page_text_rerun: + assert False, "Enterprise mode message should NOT appear after rerun when checkbox is unchecked" + From 121edcbc2d4153d0e219dfc05345bb9327ca98ed Mon Sep 17 00:00:00 2001 From: Christian Date: Mon, 15 Dec 2025 22:58:04 +0100 Subject: [PATCH 05/21] Fix Material Icons rendering for Streamlit components - Add CSS fix for stIconMaterial components to use Material Icons font - Add @font-face fallback for Material Icons - Fixes keyboard_arrow_right and other icon names rendering as text --- report_analyst/streamlit_app.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/report_analyst/streamlit_app.py b/report_analyst/streamlit_app.py index cc73c9ec1..aa87191f4 100644 --- a/report_analyst/streamlit_app.py +++ b/report_analyst/streamlit_app.py @@ -1639,6 +1639,25 @@ def main(): margin-right: 8px; } + /* Fix Material Icons rendering issues for Streamlit's stIconMaterial component */ + [data-testid="stIconMaterial"] { + font-family: 'Material Icons' !important; + font-feature-settings: 'liga' !important; + -webkit-font-feature-settings: 'liga' !important; + font-weight: normal !important; + font-style: normal !important; + text-transform: none !important; + letter-spacing: normal !important; + } + + /* @font-face fallback for Material Icons */ + @font-face { + font-family: 'Material Icons'; + font-style: normal; + font-weight: 400; + src: url(https://fonts.gstatic.com/s/materialicons/v142/flUhRq6tzZclQEJ-Vdg-IuiaDsNc.woff2) format('woff2'); + } + /* Add Material Icon to stAlert elements - only ONE icon per alert */ /* Add icon only to the markdown container, NOT to paragraphs to avoid duplicates */ [data-testid="stAlert"] [data-testid="stMarkdownContainer"]::before { From ee1a15f5e7abea5744df0cec9db394185ac06f4c Mon Sep 17 00:00:00 2001 From: Christian Date: Mon, 15 Dec 2025 23:57:22 +0100 Subject: [PATCH 06/21] Fix linting issues for CI - Add trailing newlines to all files (POSIX standard) - Fix black formatting across 52 files - Fix isort import sorting across 18 files - Update pyproject.toml line-length to 127 to match CI - All critical linting checks now pass --- Procfile | 1 + TEST_FAILURES_ANALYSIS.md | 1 + alembic/env.py | 11 +- alembic/versions/001_initial_schema.py | 186 ++-- docs/DATABASE_MIGRATIONS.md | 1 + migrations/run_migrations.sh | 1 + pyproject.toml | 4 +- report_analyst/LICENSE | 1 + report_analyst/core/analyzer.py | 342 ++---- report_analyst/core/api_key_manager.py | 11 +- report_analyst/core/cache_manager.py | 268 +++-- report_analyst/core/config.py | 4 +- report_analyst/core/database_manager.py | 1 - report_analyst/core/database_schema.py | 8 +- report_analyst/core/dataframe_manager.py | 36 +- report_analyst/core/document_processor.py | 4 +- report_analyst/core/document_sources.py | 8 +- report_analyst/core/file_storage.py | 121 ++- report_analyst/core/llm_providers.py | 16 +- report_analyst/core/migration_utils.py | 43 +- report_analyst/core/question_loader.py | 17 +- report_analyst/core/report_data_client.py | 37 +- report_analyst/core/storage/llama_store.py | 9 +- report_analyst/core/workflow_engine.py | 46 +- report_analyst/gradio_app.py | 36 +- report_analyst/main.py | 8 +- report_analyst/models/requests.py | 12 +- report_analyst/models/responses.py | 8 +- report_analyst/streamlit_app.py | 970 ++++++++---------- report_analyst/streamlit_app_backend.py | 28 +- report_analyst_api/LICENSE | 1 + report_analyst_api/main.py | 4 +- report_analyst_api/schemas.py | 68 +- report_analyst_enterprise/LICENSE | 1 + report_analyst_enterprise/__init__.py | 1 - .../database/__init__.py | 1 - .../database/pgvector_support.py | 15 +- report_analyst_jobs/LICENSE | 1 + report_analyst_jobs/analysis_handler.py | 35 +- report_analyst_jobs/coordinator.py | 12 +- report_analyst_jobs/data_lake_integration.py | 84 +- report_analyst_jobs/event_handlers.py | 21 +- report_analyst_jobs/event_router.py | 79 +- report_analyst_jobs/event_router_example.py | 9 +- report_analyst_jobs/event_routing.yaml | 1 + report_analyst_jobs/integration_examples.py | 4 +- report_analyst_jobs/interfaces.py | 4 +- report_analyst_jobs/llm_integration.py | 34 +- report_analyst_jobs/local_executor.py | 4 +- report_analyst_jobs/nats_integration.py | 171 +-- .../search_backend_integration.py | 20 +- report_analyst_search_backend/LICENSE | 1 + .../backend_service.py | 49 +- report_analyst_search_backend/config.py | 27 +- .../external_service_client.py | 36 +- .../external_service_delivery.py | 8 +- .../external_service_handler.py | 46 +- .../flow_orchestrator.py | 91 +- .../s3_upload_service.py | 10 +- .../service_discovery.py | 83 +- .../streamlit_integration.py | 56 +- run_step_by_step_tests.py | 12 +- tests/conftest.py | 9 +- tests/integration/test_streamlit_app.py | 23 +- tests/test_analyzer.py | 76 +- tests/test_api_key_manager.py | 82 +- tests/test_backend_resource_roundtrip.py | 27 +- tests/test_backend_service_list_reports.py | 29 +- tests/test_cache_manager.py | 37 +- tests/test_dataframe_manager.py | 14 +- tests/test_document_ready_auto_index.py | 29 +- tests/test_document_ready_e2e_router.py | 71 +- tests/test_event_router.py | 8 +- tests/test_event_router_e2e.py | 3 +- tests/test_external_service_integration.py | 40 +- tests/test_file_storage.py | 35 +- tests/test_llm_evidence_separation.py | 30 +- tests/test_question_loader.py | 32 +- tests/test_report_data_client.py | 29 +- tests/test_s3_upload_service.py | 13 +- tests/test_service_discovery.py | 37 +- tests/test_settings_enterprise_mode.py | 32 +- tests/test_similarity_search.py | 21 +- .../test_streamlit_app_backend_integration.py | 77 +- tests/test_streamlit_app_data_display.py | 33 +- tests/test_streamlit_app_file_selection.py | 65 +- tests/test_streamlit_app_processing_steps.py | 24 +- tests/test_streamlit_app_questions.py | 24 +- tests/test_streamlit_app_tabs.py | 20 +- 89 files changed, 1635 insertions(+), 2513 deletions(-) diff --git a/Procfile b/Procfile index 9e44064de..23a3a5aeb 100644 --- a/Procfile +++ b/Procfile @@ -3,3 +3,4 @@ release: python -m alembic upgrade head || echo "Migrations skipped (USE_ALEMBIC_MIGRATIONS not enabled or not PostgreSQL)" web: streamlit run report_analyst/streamlit_app.py --server.port=$PORT --server.address=0.0.0.0 + diff --git a/TEST_FAILURES_ANALYSIS.md b/TEST_FAILURES_ANALYSIS.md index b80f23783..a4ac40886 100644 --- a/TEST_FAILURES_ANALYSIS.md +++ b/TEST_FAILURES_ANALYSIS.md @@ -162,3 +162,4 @@ But the tests expect formatted output with bullet points: 3. **Consider**: Some tests might need to navigate to specific pages first before checking for widgets + diff --git a/alembic/env.py b/alembic/env.py index 24757d161..01ed8ac8a 100644 --- a/alembic/env.py +++ b/alembic/env.py @@ -4,10 +4,10 @@ from sqlalchemy import engine_from_config, pool from alembic import context +from report_analyst.core.database_manager import DatabaseManager # Import our database schema from report_analyst.core.database_schema import metadata as db_metadata -from report_analyst.core.database_manager import DatabaseManager # this is the Alembic Config object, which provides # access to the values within the .ini file in use. @@ -31,10 +31,11 @@ # Set the database URL in config config.set_main_option("sqlalchemy.url", database_url) -# Combine metadata from database_schema and file_storage -from sqlalchemy import MetaData, Table, Column, String, Text, DateTime, LargeBinary from datetime import datetime +# Combine metadata from database_schema and file_storage +from sqlalchemy import Column, DateTime, LargeBinary, MetaData, String, Table, Text + # Use the database_schema metadata as base target_metadata = db_metadata @@ -96,9 +97,7 @@ def run_migrations_online() -> None: ) with connectable.connect() as connection: - context.configure( - connection=connection, target_metadata=target_metadata - ) + context.configure(connection=connection, target_metadata=target_metadata) with context.begin_transaction(): context.run_migrations() diff --git a/alembic/versions/001_initial_schema.py b/alembic/versions/001_initial_schema.py index 679e7935a..b9901d230 100644 --- a/alembic/versions/001_initial_schema.py +++ b/alembic/versions/001_initial_schema.py @@ -5,15 +5,16 @@ Create Date: 2025-12-14 01:00:00.000000 """ + +from datetime import datetime from typing import Sequence, Union -from alembic import op import sqlalchemy as sa -from datetime import datetime +from alembic import op # revision identifiers, used by Alembic. -revision: str = '001_initial_schema' +revision: str = "001_initial_schema" down_revision: Union[str, None] = None branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None @@ -23,110 +24,119 @@ def upgrade() -> None: """Create initial database schema.""" # Document chunks table op.create_table( - 'document_chunks', - sa.Column('id', sa.Integer(), nullable=False, autoincrement=True), - sa.Column('file_path', sa.Text(), nullable=False), - sa.Column('chunk_text', sa.Text(), nullable=False), - sa.Column('chunk_size', sa.Integer(), nullable=False), - sa.Column('chunk_overlap', sa.Integer(), nullable=False), - sa.Column('embedding', sa.LargeBinary(), nullable=True), - sa.Column('metadata', sa.Text(), nullable=True), - sa.Column('created_at', sa.DateTime(), default=datetime.now), - sa.PrimaryKeyConstraint('id'), - sa.UniqueConstraint('file_path', 'chunk_text', 'chunk_size', 'chunk_overlap') + "document_chunks", + sa.Column("id", sa.Integer(), nullable=False, autoincrement=True), + sa.Column("file_path", sa.Text(), nullable=False), + sa.Column("chunk_text", sa.Text(), nullable=False), + sa.Column("chunk_size", sa.Integer(), nullable=False), + sa.Column("chunk_overlap", sa.Integer(), nullable=False), + sa.Column("embedding", sa.LargeBinary(), nullable=True), + sa.Column("metadata", sa.Text(), nullable=True), + sa.Column("created_at", sa.DateTime(), default=datetime.now), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("file_path", "chunk_text", "chunk_size", "chunk_overlap"), ) - + # Questions table op.create_table( - 'questions', - sa.Column('id', sa.Integer(), nullable=False, autoincrement=True), - sa.Column('question_id', sa.Text(), nullable=False), - sa.Column('question_set', sa.Text(), nullable=False), - sa.Column('question_text', sa.Text(), nullable=True), - sa.Column('guidelines', sa.Text(), nullable=True), - sa.PrimaryKeyConstraint('id'), - sa.UniqueConstraint('question_id', 'question_set') + "questions", + sa.Column("id", sa.Integer(), nullable=False, autoincrement=True), + sa.Column("question_id", sa.Text(), nullable=False), + sa.Column("question_set", sa.Text(), nullable=False), + sa.Column("question_text", sa.Text(), nullable=True), + sa.Column("guidelines", sa.Text(), nullable=True), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("question_id", "question_set"), ) - + # Analysis cache table op.create_table( - 'analysis_cache', - sa.Column('id', sa.Integer(), nullable=False, autoincrement=True), - sa.Column('file_path', sa.Text(), nullable=False), - sa.Column('question_id', sa.Text(), nullable=False), - sa.Column('chunk_size', sa.Integer(), nullable=False), - sa.Column('chunk_overlap', sa.Integer(), nullable=False), - sa.Column('top_k', sa.Integer(), nullable=False), - sa.Column('model', sa.Text(), nullable=False), - sa.Column('question_set', sa.Text(), nullable=False), - sa.Column('result', sa.Text(), nullable=False), - sa.Column('created_at', sa.DateTime(), default=datetime.now), - sa.PrimaryKeyConstraint('id'), + "analysis_cache", + sa.Column("id", sa.Integer(), nullable=False, autoincrement=True), + sa.Column("file_path", sa.Text(), nullable=False), + sa.Column("question_id", sa.Text(), nullable=False), + sa.Column("chunk_size", sa.Integer(), nullable=False), + sa.Column("chunk_overlap", sa.Integer(), nullable=False), + sa.Column("top_k", sa.Integer(), nullable=False), + sa.Column("model", sa.Text(), nullable=False), + sa.Column("question_set", sa.Text(), nullable=False), + sa.Column("result", sa.Text(), nullable=False), + sa.Column("created_at", sa.DateTime(), default=datetime.now), + sa.PrimaryKeyConstraint("id"), sa.UniqueConstraint( - 'file_path', 'question_id', 'chunk_size', 'chunk_overlap', - 'top_k', 'model', 'question_set' - ) + "file_path", + "question_id", + "chunk_size", + "chunk_overlap", + "top_k", + "model", + "question_set", + ), ) - + # Question analysis table op.create_table( - 'question_analysis', - sa.Column('id', sa.Integer(), nullable=False, autoincrement=True), - sa.Column('file_path', sa.Text(), nullable=False), - sa.Column('question_id', sa.Integer(), nullable=False), - sa.Column('model', sa.Text(), nullable=False), - sa.Column('top_k', sa.Integer(), nullable=False), - sa.Column('analysis_result', sa.Text(), nullable=False), - sa.Column('version', sa.Integer(), default=1), - sa.Column('created_at', sa.DateTime(), default=datetime.now), - sa.ForeignKeyConstraint(['question_id'], ['questions.id']), - sa.PrimaryKeyConstraint('id'), - sa.UniqueConstraint('file_path', 'question_id', 'model', 'top_k', 'version') + "question_analysis", + sa.Column("id", sa.Integer(), nullable=False, autoincrement=True), + sa.Column("file_path", sa.Text(), nullable=False), + sa.Column("question_id", sa.Integer(), nullable=False), + sa.Column("model", sa.Text(), nullable=False), + sa.Column("top_k", sa.Integer(), nullable=False), + sa.Column("analysis_result", sa.Text(), nullable=False), + sa.Column("version", sa.Integer(), default=1), + sa.Column("created_at", sa.DateTime(), default=datetime.now), + sa.ForeignKeyConstraint(["question_id"], ["questions.id"]), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("file_path", "question_id", "model", "top_k", "version"), ) - + # Chunk relevance table op.create_table( - 'chunk_relevance', - sa.Column('id', sa.Integer(), nullable=False, autoincrement=True), - sa.Column('question_analysis_id', sa.Integer(), nullable=False), - sa.Column('document_chunk_id', sa.Integer(), nullable=False), - sa.Column('chunk_order', sa.Integer(), nullable=False), - sa.Column('similarity_score', sa.Float(), nullable=True), - sa.Column('llm_score', sa.Float(), nullable=True), - sa.Column('is_evidence', sa.Boolean(), nullable=False, server_default='0'), - sa.Column('evidence_order', sa.Integer(), nullable=True), - sa.Column('metadata', sa.Text(), nullable=True), - sa.ForeignKeyConstraint(['question_analysis_id'], ['question_analysis.id']), - sa.ForeignKeyConstraint(['document_chunk_id'], ['document_chunks.id']), - sa.PrimaryKeyConstraint('id'), - sa.UniqueConstraint('question_analysis_id', 'document_chunk_id') + "chunk_relevance", + sa.Column("id", sa.Integer(), nullable=False, autoincrement=True), + sa.Column("question_analysis_id", sa.Integer(), nullable=False), + sa.Column("document_chunk_id", sa.Integer(), nullable=False), + sa.Column("chunk_order", sa.Integer(), nullable=False), + sa.Column("similarity_score", sa.Float(), nullable=True), + sa.Column("llm_score", sa.Float(), nullable=True), + sa.Column("is_evidence", sa.Boolean(), nullable=False, server_default="0"), + sa.Column("evidence_order", sa.Integer(), nullable=True), + sa.Column("metadata", sa.Text(), nullable=True), + sa.ForeignKeyConstraint(["question_analysis_id"], ["question_analysis.id"]), + sa.ForeignKeyConstraint(["document_chunk_id"], ["document_chunks.id"]), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("question_analysis_id", "document_chunk_id"), ) - + # Stored files table (for PostgreSQL file storage) op.create_table( - 'stored_files', - sa.Column('id', sa.String(length=36), nullable=False), - sa.Column('filename', sa.Text(), nullable=False), - sa.Column('file_data', sa.LargeBinary(), nullable=False), - sa.Column('content_type', sa.Text(), nullable=True), - sa.Column('file_size', sa.Text(), nullable=False), - sa.Column('created_at', sa.DateTime(), default=datetime.now), - sa.PrimaryKeyConstraint('id') + "stored_files", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("filename", sa.Text(), nullable=False), + sa.Column("file_data", sa.LargeBinary(), nullable=False), + sa.Column("content_type", sa.Text(), nullable=True), + sa.Column("file_size", sa.Text(), nullable=False), + sa.Column("created_at", sa.DateTime(), default=datetime.now), + sa.PrimaryKeyConstraint("id"), ) - + # Create indexes - op.create_index('idx_file_path', 'document_chunks', ['file_path'], unique=False) - op.create_index('idx_chunk_params', 'document_chunks', ['chunk_size', 'chunk_overlap'], unique=False) + op.create_index("idx_file_path", "document_chunks", ["file_path"], unique=False) + op.create_index( + "idx_chunk_params", + "document_chunks", + ["chunk_size", "chunk_overlap"], + unique=False, + ) def downgrade() -> None: """Drop all tables.""" - op.drop_index('idx_chunk_params', table_name='document_chunks') - op.drop_index('idx_file_path', table_name='document_chunks') - op.drop_table('chunk_relevance') - op.drop_table('question_analysis') - op.drop_table('analysis_cache') - op.drop_table('questions') - op.drop_table('document_chunks') - op.drop_table('stored_files') - + op.drop_index("idx_chunk_params", table_name="document_chunks") + op.drop_index("idx_file_path", table_name="document_chunks") + op.drop_table("chunk_relevance") + op.drop_table("question_analysis") + op.drop_table("analysis_cache") + op.drop_table("questions") + op.drop_table("document_chunks") + op.drop_table("stored_files") diff --git a/docs/DATABASE_MIGRATIONS.md b/docs/DATABASE_MIGRATIONS.md index 671e20876..accce6bd5 100644 --- a/docs/DATABASE_MIGRATIONS.md +++ b/docs/DATABASE_MIGRATIONS.md @@ -287,3 +287,4 @@ print(f"Up to date: {status['is_up_to_date']}") - [SQLAlchemy Migrations Guide](https://docs.sqlalchemy.org/en/20/core/metadata.html) - [Heroku Release Phase](https://devcenter.heroku.com/articles/release-phase) + diff --git a/migrations/run_migrations.sh b/migrations/run_migrations.sh index f0e232a46..7c50caae6 100755 --- a/migrations/run_migrations.sh +++ b/migrations/run_migrations.sh @@ -28,3 +28,4 @@ else exit 1 fi + diff --git a/pyproject.toml b/pyproject.toml index 11c14c343..3efbe5964 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [tool.isort] profile = "black" -line_length = 88 +line_length = 127 [tool.black] -line-length = 88 +line-length = 127 target-version = ['py38', 'py39', 'py310', 'py311', 'py312'] diff --git a/report_analyst/LICENSE b/report_analyst/LICENSE index 34f710549..0f6e744e2 100644 --- a/report_analyst/LICENSE +++ b/report_analyst/LICENSE @@ -18,3 +18,4 @@ For licensing inquiries or commercial use cases, please contact: + diff --git a/report_analyst/core/analyzer.py b/report_analyst/core/analyzer.py index 9cca04c47..f673dd8ea 100644 --- a/report_analyst/core/analyzer.py +++ b/report_analyst/core/analyzer.py @@ -39,9 +39,7 @@ # Check for backend configuration first use_backend = os.getenv("USE_BACKEND", "false").lower() == "true" use_centralized_llm = os.getenv("USE_CENTRALIZED_LLM", "false").lower() == "true" -use_full_backend_analysis = ( - os.getenv("USE_FULL_BACKEND_ANALYSIS", "false").lower() == "true" -) +use_full_backend_analysis = os.getenv("USE_FULL_BACKEND_ANALYSIS", "false").lower() == "true" # Check for required environment variables openai_key = os.getenv("OPENAI_API_KEY") @@ -49,9 +47,7 @@ default_model = os.getenv("OPENAI_API_MODEL", "gpt-3.5-turbo-1106") # Log available model keys -logger.info( - f"API Keys available - OpenAI: {bool(openai_key)}, Gemini: {bool(gemini_key)}" -) +logger.info(f"API Keys available - OpenAI: {bool(openai_key)}, Gemini: {bool(gemini_key)}") logger.info( f"Backend mode - USE_BACKEND: {use_backend}, USE_CENTRALIZED_LLM: {use_centralized_llm}, USE_FULL_BACKEND_ANALYSIS: {use_full_backend_analysis}" ) @@ -68,36 +64,26 @@ # Only check for API keys if not using backend LLM # Check if we need to force the default model based on available keys if default_model.startswith("gemini-") and not gemini_key: - logger.warning( - f"Default model is {default_model} but no GOOGLE_API_KEY is available" - ) + logger.warning(f"Default model is {default_model} but no GOOGLE_API_KEY is available") if openai_key: default_model = "gpt-3.5-turbo-1106" logger.info(f"Switching default model to {default_model}") else: logger.error("No valid API keys available for any models") - raise ValueError( - "No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY" - ) + raise ValueError("No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY") elif default_model.startswith("gpt-") and not openai_key: - logger.warning( - f"Default model is {default_model} but no OPENAI_API_KEY is available" - ) + logger.warning(f"Default model is {default_model} but no OPENAI_API_KEY is available") if gemini_key: default_model = "gemini-pro" logger.info(f"Switching default model to {default_model}") else: logger.error("No valid API keys available for any models") - raise ValueError( - "No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY" - ) + raise ValueError("No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY") # Ensure we have at least one API key for the selected model type if not openai_key and not gemini_key: logger.error("No API keys found - set either OPENAI_API_KEY or GOOGLE_API_KEY") - raise ValueError( - "Set either OPENAI_API_KEY or GOOGLE_API_KEY environment variable" - ) + raise ValueError("Set either OPENAI_API_KEY or GOOGLE_API_KEY environment variable") if not os.getenv("OPENAI_ORGANIZATION"): logger.warning("OPENAI_ORGANIZATION environment variable is not set") @@ -160,9 +146,7 @@ def __init__(self): log_analysis_step(f"Using default model from env: {self.default_model}") # Check if we should use backend for all LLM functionality - self.use_backend_llm = use_backend and ( - use_centralized_llm or use_full_backend_analysis - ) + self.use_backend_llm = use_backend and (use_centralized_llm or use_full_backend_analysis) if self.use_backend_llm: log_analysis_step( @@ -185,39 +169,29 @@ def __init__(self): self.embeddings = OpenAIEmbedding( api_key=openai_key, api_base=os.getenv("OPENAI_API_BASE"), - model_name=os.getenv( - "OPENAI_EMBEDDING_MODEL", "text-embedding-ada-002" - ), + model_name=os.getenv("OPENAI_EMBEDDING_MODEL", "text-embedding-ada-002"), embed_batch_size=100, ) # Configure embeddings globally for LlamaIndex Settings.embed_model = self.embeddings else: - logger.warning( - "No OpenAI API key - embedding functionality will be limited" - ) + logger.warning("No OpenAI API key - embedding functionality will be limited") self.embeddings = None except Exception as e: - log_analysis_step( - f"Error initializing local LLM clients: {str(e)}", "error" - ) + log_analysis_step(f"Error initializing local LLM clients: {str(e)}", "error") if not self.use_backend_llm: raise else: # In backend mode, local LLM failures are not critical - logger.warning( - "Local LLM initialization failed, but using backend mode" - ) + logger.warning("Local LLM initialization failed, but using backend mode") self.llm = None self.embeddings = None # Initialize caching and text processing (these are always needed) self.use_cache = True # Default to True, can be overridden - Settings.ingestion_cache = IngestionCache( - cache_dir=str(self.llm_cache_path), cache_type="local" - ) + Settings.ingestion_cache = IngestionCache(cache_dir=str(self.llm_cache_path), cache_type="local") self.text_splitter = SentenceSplitter(chunk_size=500, chunk_overlap=20) @@ -238,11 +212,12 @@ def __init__(self): # Try to get from session state if available (Streamlit context) try: import streamlit as st + database_url = st.session_state.get("database_url") except (ImportError, RuntimeError): # Not in Streamlit context or streamlit not available pass - + self.cache_manager = CacheManager(database_url=database_url) logger.info("Initialized DocumentAnalyzer with cache manager") @@ -250,7 +225,7 @@ def __init__(self): def _get_cache_key(self, file_path: str) -> str: """Generate a unique cache key based on file and all analysis parameters. - + Handles both local file paths and URNs (backend resources). Maintains backwards compatibility with existing local file cache keys. """ @@ -264,7 +239,7 @@ def _get_cache_key(self, file_path: str) -> str: f"m{model_name}_" # Include LLM model f"qs{self.question_set}" ) # Include question set - + # Handle URNs (backend resources) vs local file paths if file_path.startswith("urn:report-analyst:backend:"): # Extract resource ID from URN for cache key @@ -314,16 +289,12 @@ def _load_chunks_cache(self, cache_key: str) -> Optional[List]: # Convert to LlamaIndex Document objects chunks = [ Document( - text=chunk[ - "page_content" - ], # LlamaIndex uses text instead of page_content + text=chunk["page_content"], # LlamaIndex uses text instead of page_content metadata=chunk["metadata"], ) for chunk in chunk_data ] - logger.info( - f"[ANALYSIS] ✓ Cache HIT: Loaded {len(chunks)} chunks from cache" - ) + logger.info(f"[ANALYSIS] ✓ Cache HIT: Loaded {len(chunks)} chunks from cache") return chunks logger.info("[ANALYSIS] Cache MISS: No cached chunks found") return None @@ -349,9 +320,7 @@ def _save_chunks_cache(self, cache_key: str, chunks: List) -> None: except Exception as e: logger.warning(f"[ANALYSIS] Cache ERROR: Failed to save chunks cache: {e}") - def _load_vector_store( - self, cache_key: str, chunks: List - ) -> Optional[LlamaVectorStore]: + def _load_vector_store(self, cache_key: str, chunks: List) -> Optional[LlamaVectorStore]: """Load vector store from cache if available.""" try: store_dir = self.cache_path / f"{cache_key}_vectors" @@ -363,9 +332,7 @@ def _load_vector_store( vector_store = LlamaVectorStore(store_dir) # Try to load the store - this will verify if it's valid if vector_store.load(): - logger.info( - f"[ANALYSIS] ✓ Cache HIT: Loaded vector store from cache" - ) + logger.info(f"[ANALYSIS] ✓ Cache HIT: Loaded vector store from cache") return vector_store except Exception as inner_e: logger.error( @@ -376,9 +343,7 @@ def _load_vector_store( logger.info("[ANALYSIS] Cache MISS: No cached vector store found") return None except Exception as e: - logger.warning( - f"[ANALYSIS] Cache ERROR: Failed to load vector store cache: {e}" - ) + logger.warning(f"[ANALYSIS] Cache ERROR: Failed to load vector store cache: {e}") logger.debug(f"Full vector store cache error: {str(e)}", exc_info=True) return None @@ -442,9 +407,7 @@ async def score_chunk_relevance(self, question: str, chunk_text: str) -> float: log_analysis_step(f"Error scoring chunk relevance: {str(e)}", "error") return 0.0 - async def score_chunk_relevance_batch( - self, question: str, chunks: List[Dict], single_call: bool = True - ) -> List[float]: + async def score_chunk_relevance_batch(self, question: str, chunks: List[Dict], single_call: bool = True) -> List[float]: """Score a batch of chunks using LLM. Args: @@ -455,12 +418,7 @@ async def score_chunk_relevance_batch( try: if single_call: # Batch scoring - all chunks in one call - chunks_text = "\n\n".join( - [ - f"[CHUNK {i+1}]\n{chunk['text']}" - for i, chunk in enumerate(chunks) - ] - ) + chunks_text = "\n\n".join([f"[CHUNK {i+1}]\n{chunk['text']}" for i, chunk in enumerate(chunks)]) response = await self.llm.achat( prompt=f"""As a senior equity analyst with expertise in climate science evaluating a company's sustainability report, you are tasked with evaluating text fragments for their usefulness in answering specific TCFD questions. @@ -507,14 +465,9 @@ async def score_chunk_relevance_batch( # Parse scores from response try: - scores = [ - float(score.strip()) - for score in response.message.content.strip().split("\n") - ] + scores = [float(score.strip()) for score in response.message.content.strip().split("\n")] if len(scores) != len(chunks): - raise ValueError( - f"Got {len(scores)} scores for {len(chunks)} chunks" - ) + raise ValueError(f"Got {len(scores)} scores for {len(chunks)} chunks") return scores except Exception as e: log_analysis_step(f"Error parsing batch scores: {str(e)}", "error") @@ -567,9 +520,7 @@ def _load_cached_answers(self, file_path: str) -> Dict: with open(cache_file, "r") as f: cached_data = json.load(f) logger.info(f"Loaded cache data with keys: {list(cached_data.keys())}") - logger.info( - f"Cache data structure: {json.dumps(cached_data, indent=2)[:500]}..." - ) # Show first 500 chars + logger.info(f"Cache data structure: {json.dumps(cached_data, indent=2)[:500]}...") # Show first 500 chars return cached_data except Exception as e: @@ -608,7 +559,7 @@ async def process_document( pre_retrieved_chunks: Optional[List[Dict[str, Any]]] = None, ) -> AsyncGenerator[Dict, None]: """Process a document for selected questions - + Args: file_path: Path to document file or URN for backend resources selected_questions: List of question numbers to process @@ -666,9 +617,7 @@ async def process_document( logger.info(f"[ANALYSIS] Retrieved {len(chunks)} chunks from cache") if not chunks: - logger.info( - f"[ANALYSIS] No chunks found in cache with current parameters, creating new chunks" - ) + logger.info(f"[ANALYSIS] No chunks found in cache with current parameters, creating new chunks") # If no chunks in cache with current parameters, create them # Check if file_path is a URN (backend resource) if file_path.startswith("urn:report-analyst:backend:"): @@ -696,43 +645,29 @@ async def process_document( # 2. Process each question for question_number in selected_questions: try: - logger.info( - f"[ANALYSIS] Processing question number {question_number}" - ) + logger.info(f"[ANALYSIS] Processing question number {question_number}") question_data = self.get_question_by_number(question_number) if not question_data: - logger.warning( - f"[ANALYSIS] Question {question_number} not found" - ) + logger.warning(f"[ANALYSIS] Question {question_number} not found") yield {"error": f"Question {question_number} not found"} continue question_id = f"{self.question_set}_{question_number}" logger.info(f"[ANALYSIS] Question ID: {question_id}") - yield { - "status": f"Processing question {question_number}: {question_data['text'][:50]}..." - } + yield {"status": f"Processing question {question_number}: {question_data['text'][:50]}..."} # 3. Get similar chunks using embeddings with current parameters - logger.info( - f"[ANALYSIS] Getting similar chunks for question {question_id}" - ) - similar_chunks = await self._get_similar_chunks( - question_data["text"], chunks, self.chunk_params["top_k"] - ) - logger.info( - f"[ANALYSIS] Found {len(similar_chunks)} similar chunks" - ) + logger.info(f"[ANALYSIS] Getting similar chunks for question {question_id}") + similar_chunks = await self._get_similar_chunks(question_data["text"], chunks, self.chunk_params["top_k"]) + logger.info(f"[ANALYSIS] Found {len(similar_chunks)} similar chunks") # 3.5. Apply LLM scoring to chunks if enabled (INDEPENDENT of evidence determination) if use_llm_scoring: logger.info( f"[ANALYSIS] Applying LLM scoring to {len(similar_chunks)} chunks for question {question_id}" ) - yield { - "status": f"Scoring chunks with LLM for question {question_number}..." - } + yield {"status": f"Scoring chunks with LLM for question {question_number}..."} try: llm_scores = await self.score_chunk_relevance_batch( @@ -745,18 +680,12 @@ async def process_document( for i, chunk in enumerate(similar_chunks): if i < len(llm_scores): chunk["llm_score"] = llm_scores[i] - logger.debug( - f"Applied LLM score {llm_scores[i]:.3f} to chunk {i+1}" - ) + logger.debug(f"Applied LLM score {llm_scores[i]:.3f} to chunk {i+1}") else: chunk["llm_score"] = 0.0 - logger.warning( - f"No LLM score available for chunk {i+1}" - ) + logger.warning(f"No LLM score available for chunk {i+1}") - logger.info( - f"[ANALYSIS] Applied LLM scores to {len(similar_chunks)} chunks" - ) + logger.info(f"[ANALYSIS] Applied LLM scores to {len(similar_chunks)} chunks") except Exception as e: logger.error( @@ -767,38 +696,26 @@ async def process_document( for chunk in similar_chunks: chunk["llm_score"] = 0.0 else: - logger.info( - f"[ANALYSIS] LLM scoring disabled for question {question_id}" - ) + logger.info(f"[ANALYSIS] LLM scoring disabled for question {question_id}") # Ensure llm_score is set to None when not using LLM scoring for chunk in similar_chunks: chunk["llm_score"] = None # 4. Run LLM analysis (evidence determination happens here) - logger.info( - f"[ANALYSIS] Running LLM analysis for question {question_id}" - ) - result = await self._analyze_chunks( - question_data, similar_chunks, use_llm_scoring - ) - logger.info( - f"[ANALYSIS] LLM analysis complete for question {question_id}" - ) + logger.info(f"[ANALYSIS] Running LLM analysis for question {question_id}") + result = await self._analyze_chunks(question_data, similar_chunks, use_llm_scoring) + logger.info(f"[ANALYSIS] LLM analysis complete for question {question_id}") # Process evidence and update chunks if "EVIDENCE" in result: - logger.info( - f"Processing {len(result['EVIDENCE'])} evidence items" - ) + logger.info(f"Processing {len(result['EVIDENCE'])} evidence items") evidence_items = [] for evidence_idx, evidence in enumerate(result["EVIDENCE"]): # Extract chunk number from evidence if isinstance(evidence, dict): chunk_num = evidence.get("chunk") if chunk_num is not None: - chunk_idx = ( - chunk_num - 1 - ) # Convert to 0-based index + chunk_idx = chunk_num - 1 # Convert to 0-based index if 0 <= chunk_idx < len(similar_chunks): # Update chunk information - ONLY set evidence flags, NOT llm_score similar_chunks[chunk_idx].update( @@ -811,17 +728,13 @@ async def process_document( evidence_items.append( { "chunk": chunk_num, - "text": evidence.get( - "text", "" - ), # Keep LLM's evidence text + "text": evidence.get("text", ""), # Keep LLM's evidence text "chunk_text": similar_chunks[chunk_idx][ "text" ], # Store full chunk text separately "score": evidence.get("score", 1.0), "order": evidence_idx + 1, - "metadata": similar_chunks[chunk_idx][ - "metadata" - ], + "metadata": similar_chunks[chunk_idx]["metadata"], } ) logger.info( @@ -835,9 +748,7 @@ async def process_document( logger.warning("No EVIDENCE field found in result") # 5. Save complete analysis - logger.info( - f"[ANALYSIS] Saving analysis result for question {question_id}" - ) + logger.info(f"[ANALYSIS] Saving analysis result for question {question_id}") # Create config dict for cache manager # Safely get model name, fallback to default_model if llm is None @@ -860,9 +771,7 @@ async def process_document( logger.info(f"[ANALYSIS] Analysis saved for question {question_id}") # 8. Yield the result - logger.info( - f"[ANALYSIS] Yielding result for question {question_id}" - ) + logger.info(f"[ANALYSIS] Yielding result for question {question_id}") yield { "question_number": question_number, "question_id": question_id, @@ -874,14 +783,10 @@ async def process_document( f"[ANALYSIS] Error processing question {question_number}: {str(e)}", exc_info=True, ) - yield { - "error": f"Error processing question {question_number}: {str(e)}" - } + yield {"error": f"Error processing question {question_number}: {str(e)}"} except Exception as e: - logger.error( - f"[ANALYSIS] Error processing document: {str(e)}", exc_info=True - ) + logger.error(f"[ANALYSIS] Error processing document: {str(e)}", exc_info=True) yield {"error": f"Error processing document: {str(e)}"} def _create_chunks(self, file_path: str) -> List[Dict[str, Any]]: @@ -918,9 +823,7 @@ def _create_chunks(self, file_path: str) -> List[Dict[str, Any]]: for i in range(0, len(text_chunks), BATCH_SIZE): batch = text_chunks[i : i + BATCH_SIZE] - logger.info( - f"Computing embeddings for batch {i//BATCH_SIZE + 1}/{(len(text_chunks)-1)//BATCH_SIZE + 1}" - ) + logger.info(f"Computing embeddings for batch {i//BATCH_SIZE + 1}/{(len(text_chunks)-1)//BATCH_SIZE + 1}") # Get text from batch and clean it batch_texts = [] @@ -938,21 +841,13 @@ def _create_chunks(self, file_path: str) -> List[Dict[str, Any]]: try: # Only compute embeddings if we have valid texts if batch_texts: - logger.info( - f"Computing embeddings for {len(batch_texts)} texts in batch" - ) - batch_embeddings = self.embeddings.get_text_embedding_batch( - batch_texts - ) - logger.info( - f"Successfully computed {len(batch_embeddings)} embeddings" - ) + logger.info(f"Computing embeddings for {len(batch_texts)} texts in batch") + batch_embeddings = self.embeddings.get_text_embedding_batch(batch_texts) + logger.info(f"Successfully computed {len(batch_embeddings)} embeddings") # Create chunk dictionaries with embeddings for chunk, embedding in zip(batch, batch_embeddings): - if ( - embedding is not None - ): # Only add chunks with valid embeddings + if embedding is not None: # Only add chunks with valid embeddings chunk_dict = { "text": chunk.text, "metadata": chunk.metadata, @@ -961,16 +856,12 @@ def _create_chunks(self, file_path: str) -> List[Dict[str, Any]]: "computed_score": 0.0, # Will be populated during analysis } chunks_data.append(chunk_dict) - logger.debug( - f"Added chunk with text length {len(chunk.text)}" - ) + logger.debug(f"Added chunk with text length {len(chunk.text)}") else: logger.warning(f"Skipping chunk - embedding is None") except Exception as e: - logger.error( - f"Error computing embeddings for batch: {str(e)}", exc_info=True - ) + logger.error(f"Error computing embeddings for batch: {str(e)}", exc_info=True) # Continue with next batch, storing chunks without embeddings for chunk in batch: chunk_dict = { @@ -984,12 +875,8 @@ def _create_chunks(self, file_path: str) -> List[Dict[str, Any]]: logger.warning(f"Added chunk without embedding due to error") # Log embedding statistics - chunks_with_embeddings = sum( - 1 for c in chunks_data if c["embedding"] is not None - ) - logger.info( - f"Created {len(chunks_data)} chunks, {chunks_with_embeddings} with embeddings" - ) + chunks_with_embeddings = sum(1 for c in chunks_data if c["embedding"] is not None) + logger.info(f"Created {len(chunks_data)} chunks, {chunks_with_embeddings} with embeddings") # Only save chunks that have valid embeddings valid_chunks = [c for c in chunks_data if c["embedding"] is not None] @@ -999,9 +886,7 @@ def _create_chunks(self, file_path: str) -> List[Dict[str, Any]]: self.cache_manager.save_vectors(file_path, valid_chunks) logger.info(f"Successfully saved chunks and vectors to cache") except Exception as e: - logger.error( - f"Failed to save vectors to cache: {str(e)}", exc_info=True - ) + logger.error(f"Failed to save vectors to cache: {str(e)}", exc_info=True) else: logger.warning("No valid chunks to save to cache") @@ -1019,18 +904,14 @@ async def _analyze_chunks( ) -> Dict[str, Any]: """Analyze chunks using LLM to extract evidence and generate answer.""" try: - logger.info( - f"Analyzing {len(chunks)} chunks for question: {question_data['text'][:100]}..." - ) + logger.info(f"Analyzing {len(chunks)} chunks for question: {question_data['text'][:100]}...") # Process chunks first processed_chunks = [] for i, chunk in enumerate(chunks): logger.debug(f"Processing chunk {i+1}/{len(chunks)}") # Get similarity score from either 'score' (from vector store) or 'similarity_score' (from cache) - similarity_score = chunk.get( - "score", chunk.get("similarity_score", 0.0) - ) + similarity_score = chunk.get("score", chunk.get("similarity_score", 0.0)) # Get LLM score if it exists (independent of evidence) llm_score = chunk.get("llm_score", None) @@ -1047,9 +928,7 @@ async def _analyze_chunks( "relevance_metadata": {}, } processed_chunks.append(chunk_data) - logger.debug( - f"Processed chunk with similarity score: {similarity_score:.4f}, llm_score: {llm_score}" - ) + logger.debug(f"Processed chunk with similarity score: {similarity_score:.4f}, llm_score: {llm_score}") # Create analysis prompt with indexed chunks messages = self.prompt_manager.get_analysis_messages( @@ -1079,9 +958,7 @@ async def _analyze_chunks( # Get LLM response try: response = await self.llm.achat(messages) - response_text = ( - response.message.content - ) # Changed from response.content to response.message.content + response_text = response.message.content # Changed from response.content to response.message.content logger.info("=== LLM Response ===") logger.info(response_text) logger.info("=== End LLM Response ===") @@ -1150,17 +1027,11 @@ async def _analyze_chunks( evidence_items.append( { "chunk": chunk_num, - "text": evidence.get( - "text", "" - ), # Keep LLM's evidence text - "chunk_text": processed_chunks[chunk_idx][ - "text" - ], # Store full chunk text separately + "text": evidence.get("text", ""), # Keep LLM's evidence text + "chunk_text": processed_chunks[chunk_idx]["text"], # Store full chunk text separately "score": evidence.get("score", 1.0), "order": evidence_idx + 1, - "metadata": processed_chunks[chunk_idx][ - "metadata" - ], + "metadata": processed_chunks[chunk_idx]["metadata"], } ) logger.info( @@ -1175,9 +1046,7 @@ async def _analyze_chunks( # Add processed chunks to result result["chunks"] = processed_chunks - logger.info( - f"Analysis complete. Found {sum(1 for c in processed_chunks if c['is_evidence'])} evidence chunks" - ) + logger.info(f"Analysis complete. Found {sum(1 for c in processed_chunks if c['is_evidence'])} evidence chunks") return result except Exception as e: @@ -1197,15 +1066,9 @@ def _load_questions(self) -> dict: """Load questions from YAML files""" # Look for question set file in multiple possible locations possible_paths = [ - Path(__file__).parent.parent - / "questionsets" - / f"{self.question_set}_questions.yaml", # app/questionsets - Path(__file__).parent.parent.parent - / "questionsets" - / f"{self.question_set}_questions.yaml", # project root - Path.cwd() - / "questionsets" - / f"{self.question_set}_questions.yaml", # current working directory + Path(__file__).parent.parent / "questionsets" / f"{self.question_set}_questions.yaml", # app/questionsets + Path(__file__).parent.parent.parent / "questionsets" / f"{self.question_set}_questions.yaml", # project root + Path.cwd() / "questionsets" / f"{self.question_set}_questions.yaml", # current working directory ] log_analysis_step(f"Looking for {self.question_set}_questions.yaml in:") @@ -1229,9 +1092,7 @@ def _load_questions(self) -> dict: try: with open(yaml_file, "r") as f: config = yaml.safe_load(f) - log_analysis_step( - f"Loaded YAML content: {str(config)[:200]}..." - ) # Show first 200 chars + log_analysis_step(f"Loaded YAML content: {str(config)[:200]}...") # Show first 200 chars questions = {} # Convert the questions list into a structured format @@ -1242,13 +1103,9 @@ def _load_questions(self) -> dict: "text": q.get("text", ""), "guidelines": q.get("guidelines", ""), } - log_analysis_step( - f"Added question {q_id}: {questions[q_id]['text'][:50]}..." - ) + log_analysis_step(f"Added question {q_id}: {questions[q_id]['text'][:50]}...") - log_analysis_step( - f"✓ Loaded {len(questions)} questions for {self.question_set}" - ) + log_analysis_step(f"✓ Loaded {len(questions)} questions for {self.question_set}") log_analysis_step(f"Available question IDs: {list(questions.keys())}") return questions except Exception as e: @@ -1268,9 +1125,7 @@ def get_question_by_number(self, number: int) -> Optional[Dict]: } # Get the correct prefix for the question set - question_prefix = question_set_mapping.get( - self.question_set, self.question_set - ) + question_prefix = question_set_mapping.get(self.question_set, self.question_set) question_key = f"{question_prefix}_{number}" logger.debug(f"Looking for question {number} with key: {question_key}") @@ -1284,9 +1139,7 @@ def get_question_by_number(self, number: int) -> Optional[Dict]: def update_parameters(self, chunk_size: int, chunk_overlap: int, top_k: int): """Update analysis parameters and recreate text splitter.""" - logger.info( - f"Updating parameters: size={chunk_size}, overlap={chunk_overlap}, top_k={top_k}" - ) + logger.info(f"Updating parameters: size={chunk_size}, overlap={chunk_overlap}, top_k={top_k}") self.chunk_params = { "chunk_size": chunk_size, @@ -1295,9 +1148,7 @@ def update_parameters(self, chunk_size: int, chunk_overlap: int, top_k: int): } # Recreate text splitter with new parameters - self.text_splitter = SentenceSplitter( - chunk_size=chunk_size, chunk_overlap=chunk_overlap - ) + self.text_splitter = SentenceSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) logger.info(f"Updated parameters and recreated text splitter") @@ -1327,7 +1178,7 @@ def check_step_completion(self, file_path: str) -> Dict[str, bool]: "chunk_size": self.chunk_params["chunk_size"], "chunk_overlap": self.chunk_params["chunk_overlap"], "top_k": self.chunk_params["top_k"], - "model": self.llm.model if self.llm and hasattr(self.llm, "model") else self.default_model, + "model": (self.llm.model if self.llm and hasattr(self.llm, "model") else self.default_model), "question_set": self.question_set, } @@ -1345,9 +1196,8 @@ def check_step_completion(self, file_path: str) -> Dict[str, bool]: chunk_size=config["chunk_size"], chunk_overlap=config["chunk_overlap"], ) - step2_complete = ( - len(chunks_with_embeddings) > 0 - and any(c.get("embedding") is not None for c in chunks_with_embeddings) + step2_complete = len(chunks_with_embeddings) > 0 and any( + c.get("embedding") is not None for c in chunks_with_embeddings ) # Check step 3: Chunk scoring (check if any questions have been scored) @@ -1411,9 +1261,7 @@ def _parse_config_from_filename(self, filename: str) -> Dict[str, Any]: logger.warning(f"Error parsing config from filename {filename}: {e}") return config - async def _get_similar_chunks( - self, query_text: str, chunks: List[Dict], top_k: int - ) -> List[Dict]: + async def _get_similar_chunks(self, query_text: str, chunks: List[Dict], top_k: int) -> List[Dict]: """ Get chunks most similar to the query text using vector similarity. @@ -1487,9 +1335,7 @@ def _parse_analysis_response(self, response_text: str) -> Dict[str, Any]: if isinstance(parsed_json["SCORE"], str): import re - score_match = re.search( - r"\d+(\.\d+)?", parsed_json["SCORE"] - ) + score_match = re.search(r"\d+(\.\d+)?", parsed_json["SCORE"]) if score_match: result["SCORE"] = float(score_match.group(0)) @@ -1501,8 +1347,7 @@ def _parse_analysis_response(self, response_text: str) -> Dict[str, Any]: chunk = evidence.get("chunk") if chunk is not None: evidence_item = { - "chunk_index": int(chunk) - - 1, # Convert to 0-based index + "chunk_index": int(chunk) - 1, # Convert to 0-based index "order": len(evidence_list) + 1, "score": 1.0, # Default score "text": evidence.get("text", ""), @@ -1575,8 +1420,7 @@ def _parse_analysis_response(self, response_text: str) -> Dict[str, Any]: chunk_match = re.search(r"\[CHUNK (\d+)\]", line) if chunk_match: evidence_item = { - "chunk_index": int(chunk_match.group(1)) - - 1, # Convert to 0-based index + "chunk_index": int(chunk_match.group(1)) - 1, # Convert to 0-based index "order": len(evidence_items) + 1, "score": 1.0, # Default score "text": line, @@ -1585,17 +1429,9 @@ def _parse_analysis_response(self, response_text: str) -> Dict[str, Any]: result["EVIDENCE"] = evidence_items elif section_name == "GAPS": - result["GAPS"] = [ - line.strip() - for line in section_content.split("\n") - if line.strip() - ] + result["GAPS"] = [line.strip() for line in section_content.split("\n") if line.strip()] elif section_name == "SOURCES": - result["SOURCES"] = [ - int(s.strip()) - for s in section_content.split(",") - if s.strip().isdigit() - ] + result["SOURCES"] = [int(s.strip()) for s in section_content.split(",") if s.strip().isdigit()] return result @@ -1624,9 +1460,7 @@ def create_analysis_dataframes(results: Dict) -> pd.DataFrame: continue # Get question text from analyzer's questions data - question_text = questions.get(question_id, {}).get( - "text", f"Question {question_id}" - ) + question_text = questions.get(question_id, {}).get("text", f"Question {question_id}") # Convert lists to strings and ensure proper types row = { diff --git a/report_analyst/core/api_key_manager.py b/report_analyst/core/api_key_manager.py index 41e48fe0a..4a81f3a9f 100644 --- a/report_analyst/core/api_key_manager.py +++ b/report_analyst/core/api_key_manager.py @@ -16,7 +16,7 @@ class APIKeyManager: def set_api_key(key_name: str, value: Optional[str], session_state: dict) -> None: """ Set an API key in session state and environment variables. - + Args: key_name: The environment variable name (e.g., 'OPENAI_API_KEY') value: The API key value (None to clear) @@ -43,15 +43,15 @@ def set_api_key(key_name: str, value: Optional[str], session_state: dict) -> Non def get_api_key(key_name: str, session_state: dict) -> Optional[str]: """ Get an API key from session state or environment. - + Priority: 1. Session state (user-entered value) 2. Environment variable (from .env or system) - + Args: key_name: The environment variable name (e.g., 'OPENAI_API_KEY') session_state: Streamlit session state dictionary - + Returns: The API key value or None """ @@ -67,7 +67,7 @@ def sync_api_keys_to_env(session_state: dict) -> None: """ Sync all API keys from session state to environment variables. Called at startup to ensure environment has current values. - + Args: session_state: Streamlit session state dictionary """ @@ -78,4 +78,3 @@ def sync_api_keys_to_env(session_state: dict) -> None: value = session_state[session_key] if value: os.environ[key_name] = value - diff --git a/report_analyst/core/cache_manager.py b/report_analyst/core/cache_manager.py index dafd987c2..501718993 100644 --- a/report_analyst/core/cache_manager.py +++ b/report_analyst/core/cache_manager.py @@ -47,7 +47,7 @@ def __init__(self, db_path: str = None, database_url: str = None): # Initialize database manager self.db_manager = DatabaseManager(db_url) logger.info(f"Initializing CacheManager with database: {self.db_manager._mask_url(db_url)}") - + # Check if we should use Alembic migrations instead of auto-creation use_alembic = os.getenv("USE_ALEMBIC_MIGRATIONS", "false").lower() == "true" if use_alembic and self.db_manager.is_postgres(): @@ -65,7 +65,7 @@ def init_db(self): engine = self.db_manager.get_engine() # Create all tables metadata.create_all(engine) - + # Create indexes (using raw SQL for IF NOT EXISTS support) # Note: Some databases may not support IF NOT EXISTS in CREATE INDEX # We'll try to create them and ignore errors if they already exist @@ -77,7 +77,7 @@ def init_db(self): # Index might already exist, which is fine logger.debug(f"Index creation (may already exist): {e}") conn.commit() - + logger.info("Database schema initialized successfully") except Exception as e: logger.error(f"Error initializing database schema: {str(e)}", exc_info=True) @@ -113,9 +113,7 @@ def _load_vector_store(self, file_path: str, chunks: List[Dict]) -> None: ) self.current_file_path = file_path - logger.info( - f"Loaded {len(documents)} chunks into vector store for {file_path}" - ) + logger.info(f"Loaded {len(documents)} chunks into vector store for {file_path}") except Exception as e: logger.error(f"Error loading vector store: {str(e)}", exc_info=True) @@ -155,9 +153,7 @@ async def get_similar_chunks( for node in nodes: # Ensure we have a valid similarity score similarity_score = ( - node.score - if hasattr(node, "score") - else node.get_score() if hasattr(node, "get_score") else 0.0 + node.score if hasattr(node, "score") else node.get_score() if hasattr(node, "get_score") else 0.0 ) chunk = { @@ -169,9 +165,7 @@ async def get_similar_chunks( "similarity_score": similarity_score, # Also store as 'similarity_score' for backward compatibility } chunks.append(chunk) - logger.debug( - f"Found chunk with similarity score: {similarity_score:.4f}" - ) + logger.debug(f"Found chunk with similarity score: {similarity_score:.4f}") logger.info(f"Retrieved {len(chunks)} similar chunks for {file_path}") if chunks: @@ -185,9 +179,7 @@ async def get_similar_chunks( logger.error(f"Error getting similar chunks: {str(e)}", exc_info=True) return [] - def save_analysis( - self, file_path: str, question_id: str, result: Dict, config: Dict - ): + def save_analysis(self, file_path: str, question_id: str, result: Dict, config: Dict): """Save analysis result to cache with improved logging""" try: logger.info(f"Saving analysis for {file_path} - {question_id}") @@ -198,14 +190,14 @@ def save_analysis( with self.db_manager.get_connection() as conn: # Ensure question exists in questions table - logger.info( - f"Ensuring question {question_id} exists in questions table" - ) + logger.info(f"Ensuring question {question_id} exists in questions table") result_obj = conn.execute( - text(""" + text( + """ SELECT id FROM questions WHERE question_id = :question_id AND question_set = :question_set - """), + """ + ), {"question_id": question_id, "question_set": question_set}, ) row = result_obj.fetchone() @@ -218,14 +210,16 @@ def save_analysis( if self.db_manager.is_postgres(): # PostgreSQL: ON CONFLICT result_obj = conn.execute( - text(""" + text( + """ INSERT INTO questions (question_id, question_set, question_text, guidelines) VALUES (:question_id, :question_set, :question_text, :guidelines) ON CONFLICT (question_id, question_set) DO UPDATE SET question_text = EXCLUDED.question_text, guidelines = EXCLUDED.guidelines RETURNING id - """), + """ + ), { "question_id": question_id, "question_set": question_set, @@ -236,10 +230,12 @@ def save_analysis( else: # SQLite: INSERT OR REPLACE result_obj = conn.execute( - text(""" + text( + """ INSERT OR REPLACE INTO questions (question_id, question_set, question_text, guidelines) VALUES (:question_id, :question_set, :question_text, :guidelines) - """), + """ + ), { "question_id": question_id, "question_set": question_set, @@ -259,7 +255,8 @@ def save_analysis( logger.info("Saving main analysis result") if self.db_manager.is_postgres(): result_obj = conn.execute( - text(""" + text( + """ INSERT INTO question_analysis (file_path, question_id, model, top_k, analysis_result, version, created_at) VALUES (:file_path, :question_id, :model, :top_k, :analysis_result, :version, :created_at) @@ -267,7 +264,8 @@ def save_analysis( SET analysis_result = EXCLUDED.analysis_result, created_at = EXCLUDED.created_at RETURNING id - """), + """ + ), { "file_path": str(file_path), "question_id": question_db_id, @@ -280,11 +278,13 @@ def save_analysis( ) else: result_obj = conn.execute( - text(""" + text( + """ INSERT OR REPLACE INTO question_analysis (file_path, question_id, model, top_k, analysis_result, version, created_at) VALUES (:file_path, :question_id, :model, :top_k, :analysis_result, :version, :created_at) - """), + """ + ), { "file_path": str(file_path), "question_id": question_db_id, @@ -297,11 +297,13 @@ def save_analysis( ) # Get ID separately for SQLite result_obj = conn.execute( - text(""" + text( + """ SELECT id FROM question_analysis WHERE file_path = :file_path AND question_id = :question_id AND model = :model AND top_k = :top_k AND version = :version - """), + """ + ), { "file_path": str(file_path), "question_id": question_db_id, @@ -315,18 +317,18 @@ def save_analysis( # Save chunk relevance information if "chunks" in result: - logger.info( - f"Processing {len(result['chunks'])} chunks for relevance" - ) + logger.info(f"Processing {len(result['chunks'])} chunks for relevance") for chunk in result["chunks"]: logger.debug(f"Processing chunk: {json.dumps(chunk, indent=2)}") # Get chunk ID from document_chunks table result_obj = conn.execute( - text(""" + text( + """ SELECT id FROM document_chunks WHERE file_path = :file_path AND chunk_text = :chunk_text - """), + """ + ), {"file_path": str(file_path), "chunk_text": chunk["text"]}, ) row = result_obj.fetchone() @@ -337,7 +339,8 @@ def save_analysis( # Save chunk relevance with all available information if self.db_manager.is_postgres(): conn.execute( - text(""" + text( + """ INSERT INTO chunk_relevance (question_analysis_id, document_chunk_id, chunk_order, similarity_score, llm_score, is_evidence, evidence_order, metadata) @@ -350,7 +353,8 @@ def save_analysis( is_evidence = EXCLUDED.is_evidence, evidence_order = EXCLUDED.evidence_order, metadata = EXCLUDED.metadata - """), + """ + ), { "question_analysis_id": analysis_id, "document_chunk_id": chunk_id, @@ -364,13 +368,15 @@ def save_analysis( ) else: conn.execute( - text(""" + text( + """ INSERT OR REPLACE INTO chunk_relevance (question_analysis_id, document_chunk_id, chunk_order, similarity_score, llm_score, is_evidence, evidence_order, metadata) VALUES (:question_analysis_id, :document_chunk_id, :chunk_order, :similarity_score, :llm_score, :is_evidence, :evidence_order, :metadata) - """), + """ + ), { "question_analysis_id": analysis_id, "document_chunk_id": chunk_id, @@ -386,15 +392,14 @@ def save_analysis( f"Saving raw values to DB - similarity_score: {chunk.get('similarity_score')}, llm_score: {chunk.get('llm_score')}, is_evidence: {chunk.get('is_evidence')}" ) else: - logger.warning( - f"Could not find chunk in document_chunks table" - ) + logger.warning(f"Could not find chunk in document_chunks table") # Save to analysis cache logger.info("Saving to analysis cache") if self.db_manager.is_postgres(): conn.execute( - text(""" + text( + """ INSERT INTO analysis_cache (file_path, question_id, chunk_size, chunk_overlap, top_k, model, question_set, result, created_at) @@ -403,7 +408,8 @@ def save_analysis( ON CONFLICT (file_path, question_id, chunk_size, chunk_overlap, top_k, model, question_set) DO UPDATE SET result = EXCLUDED.result, created_at = EXCLUDED.created_at - """), + """ + ), { "file_path": str(file_path), "question_id": question_id, @@ -418,13 +424,15 @@ def save_analysis( ) else: conn.execute( - text(""" + text( + """ INSERT OR REPLACE INTO analysis_cache (file_path, question_id, chunk_size, chunk_overlap, top_k, model, question_set, result, created_at) VALUES (:file_path, :question_id, :chunk_size, :chunk_overlap, :top_k, :model, :question_set, :result, :created_at) - """), + """ + ), { "file_path": str(file_path), "question_id": question_id, @@ -444,9 +452,7 @@ def save_analysis( logger.error(f"Error saving analysis: {str(e)}", exc_info=True) raise - def get_analysis( - self, file_path: str, config: Dict, question_ids: Optional[List[str]] = None - ) -> Dict[str, Any]: + def get_analysis(self, file_path: str, config: Dict, question_ids: Optional[List[str]] = None) -> Dict[str, Any]: """ Get analysis results matching the exact configuration. @@ -472,9 +478,7 @@ def get_analysis( "s4m": "s4m", "lucia": "lucia", } - db_question_set = question_set_mapping.get( - config["question_set"], config["question_set"] - ) + db_question_set = question_set_mapping.get(config["question_set"], config["question_set"]) # First get the analysis results from the cache table query = """ @@ -582,12 +586,8 @@ def get_analysis( # Sort chunks by their order for question_id in results: - results[question_id]["chunks"].sort( - key=lambda x: x["chunk_order"] - ) - logger.info( - f"Question {question_id}: {len(results[question_id]['chunks'])} chunks" - ) + results[question_id]["chunks"].sort(key=lambda x: x["chunk_order"]) + logger.info(f"Question {question_id}: {len(results[question_id]['chunks'])} chunks") if results[question_id]["chunks"]: logger.info( f" Similarity range: {min(c['similarity_score'] for c in results[question_id]['chunks']):.4f} - {max(c['similarity_score'] for c in results[question_id]['chunks']):.4f}" @@ -606,9 +606,7 @@ def save_vectors(self, file_path: str, chunks: List[Dict[str, Any]]) -> None: # Get chunk parameters from first chunk's metadata chunk_size = chunks[0]["metadata"].get("chunk_size", 0) if chunks else 0 - chunk_overlap = ( - chunks[0]["metadata"].get("chunk_overlap", 0) if chunks else 0 - ) + chunk_overlap = chunks[0]["metadata"].get("chunk_overlap", 0) if chunks else 0 logger.info(f"Chunk parameters: size={chunk_size}, overlap={chunk_overlap}") with self.db_manager.get_connection() as conn: @@ -629,19 +627,19 @@ def save_vectors(self, file_path: str, chunks: List[Dict[str, Any]]) -> None: metadata_with_shape["embedding_dtype"] = str(embedding.dtype) # Prepare chunk data - chunk_data.append({ - "file_path": file_path, - "chunk_text": chunk["text"], - "chunk_size": chunk_size, - "chunk_overlap": chunk_overlap, - "embedding": embedding_bytes, - "metadata": json.dumps(metadata_with_shape), - "created_at": datetime.now().isoformat(), - }) - except Exception as e: - logger.warning( - f"Error preparing chunk {i} for storage: {str(e)}" + chunk_data.append( + { + "file_path": file_path, + "chunk_text": chunk["text"], + "chunk_size": chunk_size, + "chunk_overlap": chunk_overlap, + "embedding": embedding_bytes, + "metadata": json.dumps(metadata_with_shape), + "created_at": datetime.now().isoformat(), + } ) + except Exception as e: + logger.warning(f"Error preparing chunk {i} for storage: {str(e)}") continue if chunk_data: @@ -650,7 +648,8 @@ def save_vectors(self, file_path: str, chunks: List[Dict[str, Any]]) -> None: # PostgreSQL: ON CONFLICT for chunk_row in chunk_data: conn.execute( - text(""" + text( + """ INSERT INTO document_chunks (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) @@ -660,26 +659,27 @@ def save_vectors(self, file_path: str, chunks: List[Dict[str, Any]]) -> None: SET embedding = EXCLUDED.embedding, metadata = EXCLUDED.metadata, created_at = EXCLUDED.created_at - """), + """ + ), chunk_row, ) else: # SQLite: INSERT OR REPLACE for chunk_row in chunk_data: conn.execute( - text(""" + text( + """ INSERT OR REPLACE INTO document_chunks (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) - """), + """ + ), chunk_row, ) - logger.info( - f"Successfully saved all {len(chunk_data)} chunks to database" - ) + logger.info(f"Successfully saved all {len(chunk_data)} chunks to database") # Verify the insertion result_obj = conn.execute( @@ -687,9 +687,7 @@ def save_vectors(self, file_path: str, chunks: List[Dict[str, Any]]) -> None: {"file_path": file_path}, ) count = result_obj.fetchone()[0] - logger.info( - f"Verification: Found {count} chunks in database for {file_path}" - ) + logger.info(f"Verification: Found {count} chunks in database for {file_path}") else: logger.warning("No valid chunks to save") @@ -702,11 +700,13 @@ def get_vectors(self, file_path: str) -> List[Dict[str, Any]]: try: with self.db_manager.get_connection() as conn: result_obj = conn.execute( - text(""" + text( + """ SELECT chunk_text, embedding, metadata FROM document_chunks WHERE file_path = :file_path - """), + """ + ), {"file_path": str(file_path)}, ) chunks = [] @@ -722,9 +722,7 @@ def get_vectors(self, file_path: str) -> List[Dict[str, Any]]: dtype = metadata_dict.get("embedding_dtype", "float32") if shape: - embedding = np.frombuffer(row[1], dtype=dtype).reshape( - shape - ) + embedding = np.frombuffer(row[1], dtype=dtype).reshape(shape) else: # Fallback to default shape if not stored embedding = np.frombuffer(row[1], dtype=np.float32) @@ -734,9 +732,7 @@ def get_vectors(self, file_path: str) -> List[Dict[str, Any]]: # Remove embedding metadata from the returned metadata clean_metadata = { - k: v - for k, v in metadata_dict.items() - if k not in ["embedding_shape", "embedding_dtype"] + k: v for k, v in metadata_dict.items() if k not in ["embedding_shape", "embedding_dtype"] } chunks.append( @@ -780,20 +776,24 @@ def check_cache_status(self, file_path: str = None): if file_path: logger.info(f"Checking cache for file: {file_path}") result_obj = conn.execute( - text(""" + text( + """ SELECT DISTINCT chunk_size, chunk_overlap, top_k, model, question_set FROM analysis_cache WHERE file_path = :file_path - """), + """ + ), {"file_path": str(file_path)}, ) else: logger.info("Checking all cache entries") result_obj = conn.execute( - text(""" + text( + """ SELECT DISTINCT file_path, chunk_size, chunk_overlap, top_k, model, question_set FROM analysis_cache - """) + """ + ) ) rows = result_obj.fetchall() @@ -813,7 +813,8 @@ def get_all_answers_by_question_set(self, question_set: str) -> Dict[str, Any]: with self.db_manager.get_connection() as conn: # First get all analysis results result_obj = conn.execute( - text(""" + text( + """ SELECT ac.question_id, ac.result, dc.chunk_text, dc.metadata as chunk_metadata, cr.chunk_order, cr.similarity_score, @@ -826,7 +827,8 @@ def get_all_answers_by_question_set(self, question_set: str) -> Dict[str, Any]: LEFT JOIN document_chunks dc ON cr.document_chunk_id = dc.id WHERE ac.question_set = :question_set ORDER BY ac.question_id, cr.chunk_order - """), + """ + ), {"question_set": question_set}, ) @@ -862,21 +864,15 @@ def get_all_answers_by_question_set(self, question_set: str) -> Dict[str, Any]: # Sort chunks by their order for each result for question_id in results: - results[question_id]["chunks"].sort( - key=lambda x: x.get("chunk_order", 0) - ) + results[question_id]["chunks"].sort(key=lambda x: x.get("chunk_order", 0)) return results except Exception as e: - logger.error( - f"Error retrieving answers for question set {question_set}: {e}" - ) + logger.error(f"Error retrieving answers for question set {question_set}: {e}") raise - def save_document_chunks( - self, file_path: str, chunks: List[Dict], chunk_size: int, chunk_overlap: int - ) -> None: + def save_document_chunks(self, file_path: str, chunks: List[Dict], chunk_size: int, chunk_overlap: int) -> None: """Save document chunks to cache with their embeddings.""" try: logger.info(f"Starting to save {len(chunks)} chunks for {file_path}") @@ -900,7 +896,8 @@ def save_document_chunks( if self.db_manager.is_postgres(): conn.execute( - text(""" + text( + """ INSERT INTO document_chunks (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) @@ -908,7 +905,8 @@ def save_document_chunks( SET embedding = EXCLUDED.embedding, metadata = EXCLUDED.metadata, created_at = EXCLUDED.created_at - """), + """ + ), { "file_path": str(file_path), "chunk_text": chunk["text"], @@ -921,11 +919,13 @@ def save_document_chunks( ) else: conn.execute( - text(""" + text( + """ INSERT OR REPLACE INTO document_chunks (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) - """), + """ + ), { "file_path": str(file_path), "chunk_text": chunk["text"], @@ -941,10 +941,12 @@ def save_document_chunks( # Verify chunks were saved result_obj = conn.execute( - text(""" + text( + """ SELECT COUNT(*) FROM document_chunks WHERE file_path = :file_path AND chunk_size = :chunk_size AND chunk_overlap = :chunk_overlap - """), + """ + ), { "file_path": str(file_path), "chunk_size": chunk_size, @@ -952,25 +954,19 @@ def save_document_chunks( }, ) count = result_obj.fetchone()[0] - logger.info( - f"Verification: Found {count} chunks in database for {file_path}" - ) + logger.info(f"Verification: Found {count} chunks in database for {file_path}") except Exception as e: logger.error(f"Error saving document chunks: {str(e)}", exc_info=True) raise - def get_document_chunks( - self, file_path: str, chunk_size: int = None, chunk_overlap: int = None - ) -> List[Dict]: + def get_document_chunks(self, file_path: str, chunk_size: int = None, chunk_overlap: int = None) -> List[Dict]: """ Get document chunks from cache with improved logging. """ try: logger.info(f"Retrieving chunks for {file_path}") - logger.info( - f"Filters: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}" - ) + logger.info(f"Filters: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}") with self.db_manager.get_connection() as conn: query = """ @@ -1012,9 +1008,7 @@ def get_document_chunks( embedding = None if embedding_bytes: embedding = np.frombuffer(embedding_bytes, dtype=np.float32) - logger.debug( - f"Converted embedding bytes to numpy array, shape: {embedding.shape}" - ) + logger.debug(f"Converted embedding bytes to numpy array, shape: {embedding.shape}") else: logger.debug("No embedding found for chunk") @@ -1022,9 +1016,7 @@ def get_document_chunks( metadata = {} if metadata_json: metadata = json.loads(metadata_json) - logger.debug( - f"Parsed metadata: {json.dumps(metadata, indent=2)}" - ) + logger.debug(f"Parsed metadata: {json.dumps(metadata, indent=2)}") chunks.append( { @@ -1038,24 +1030,18 @@ def get_document_chunks( ) logger.info(f"Retrieved {len(chunks)} chunks") - logger.debug( - f"Chunks have embeddings: {sum(1 for c in chunks if c['embedding'] is not None)}/{len(chunks)}" - ) + logger.debug(f"Chunks have embeddings: {sum(1 for c in chunks if c['embedding'] is not None)}/{len(chunks)}") return chunks except Exception as e: logger.error(f"Error getting document chunks: {str(e)}", exc_info=True) return [] - def get_chunks_without_embeddings( - self, file_path: str, chunk_size: int = None, chunk_overlap: int = None - ) -> List[Dict]: + def get_chunks_without_embeddings(self, file_path: str, chunk_size: int = None, chunk_overlap: int = None) -> List[Dict]: """Get chunks without embeddings (where embedding IS NULL)""" try: logger.info(f"Retrieving chunks without embeddings for {file_path}") - logger.info( - f"Filters: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}" - ) + logger.info(f"Filters: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}") with self.db_manager.get_connection() as conn: query = """ @@ -1103,9 +1089,7 @@ def get_chunks_without_embeddings( return chunks except Exception as e: - logger.error( - f"Error getting chunks without embeddings: {str(e)}", exc_info=True - ) + logger.error(f"Error getting chunks without embeddings: {str(e)}", exc_info=True) return [] def has_chunk_scoring(self, file_path: str, config: Dict) -> bool: @@ -1113,13 +1097,15 @@ def has_chunk_scoring(self, file_path: str, config: Dict) -> bool: try: with self.db_manager.get_connection() as conn: result_obj = conn.execute( - text(""" + text( + """ SELECT COUNT(DISTINCT q.question_id) FROM questions q JOIN question_analysis qa ON qa.question_id = q.id JOIN chunk_relevance cr ON cr.question_analysis_id = qa.id WHERE qa.file_path = :file_path AND qa.model = :model AND qa.top_k = :top_k - """), + """ + ), { "file_path": str(file_path), "model": config["model"], diff --git a/report_analyst/core/config.py b/report_analyst/core/config.py index 6d11462fb..496136282 100644 --- a/report_analyst/core/config.py +++ b/report_analyst/core/config.py @@ -22,9 +22,7 @@ class Config: GOOGLE_API_KEY: Optional[str] = os.getenv("GOOGLE_API_KEY") # Search Backend Configuration - SEARCH_BACKEND_URL: Optional[str] = os.getenv( - "SEARCH_BACKEND_URL", "http://localhost:8001" - ) + SEARCH_BACKEND_URL: Optional[str] = os.getenv("SEARCH_BACKEND_URL", "http://localhost:8001") SEARCH_BACKEND_API_KEY: Optional[str] = os.getenv("SEARCH_BACKEND_API_KEY") # Document Processing diff --git a/report_analyst/core/database_manager.py b/report_analyst/core/database_manager.py index efa737247..e35c1beb2 100644 --- a/report_analyst/core/database_manager.py +++ b/report_analyst/core/database_manager.py @@ -101,4 +101,3 @@ def is_postgres(self) -> bool: def is_sqlite(self) -> bool: """Check if using SQLite.""" return not self._is_postgres - diff --git a/report_analyst/core/database_schema.py b/report_analyst/core/database_schema.py index 5ef987e11..b0c9d0708 100644 --- a/report_analyst/core/database_schema.py +++ b/report_analyst/core/database_schema.py @@ -95,7 +95,12 @@ "chunk_relevance", metadata, Column("id", Integer, primary_key=True, autoincrement=True), - Column("question_analysis_id", Integer, ForeignKey("question_analysis.id"), nullable=False), + Column( + "question_analysis_id", + Integer, + ForeignKey("question_analysis.id"), + nullable=False, + ), Column("document_chunk_id", Integer, ForeignKey("document_chunks.id"), nullable=False), Column("chunk_order", Integer, nullable=False), Column("similarity_score", Float, nullable=True), @@ -115,4 +120,3 @@ # Index on chunk parameters "CREATE INDEX IF NOT EXISTS idx_chunk_params ON document_chunks(chunk_size, chunk_overlap)", ] - diff --git a/report_analyst/core/dataframe_manager.py b/report_analyst/core/dataframe_manager.py index dc39516e4..c9fb524bd 100644 --- a/report_analyst/core/dataframe_manager.py +++ b/report_analyst/core/dataframe_manager.py @@ -40,17 +40,13 @@ def extract_evidence_text(evidence: Any) -> str: return str(evidence) -def create_analysis_dataframes( - cached_results: Dict, file_key: str = None -) -> Tuple[pd.DataFrame, pd.DataFrame]: +def create_analysis_dataframes(cached_results: Dict, file_key: str = None) -> Tuple[pd.DataFrame, pd.DataFrame]: """Create analysis and chunks dataframes from database results.""" try: analysis_rows = [] chunks_rows = [] - logger.info( - f"Processing {len(cached_results)} results for file_key: {file_key}" - ) + logger.info(f"Processing {len(cached_results)} results for file_key: {file_key}") logger.info(f"Input cached_results keys: {list(cached_results.keys())}") # Handle each question's results @@ -58,9 +54,7 @@ def create_analysis_dataframes( try: # Get the result data - it might be nested under 'result' key result = data.get("result", data) - logger.info( - f"Processing question {question_id} with keys: {list(result.keys())}" - ) + logger.info(f"Processing question {question_id} with keys: {list(result.keys())}") # Create analysis row analysis_row = { @@ -76,18 +70,14 @@ def create_analysis_dataframes( # Process chunks - use exactly what's in the database chunks = data.get("chunks", []) - logger.info( - f"Processing {len(chunks)} chunks for question {question_id}" - ) + logger.info(f"Processing {len(chunks)} chunks for question {question_id}") for chunk in chunks: # Create chunk row with exactly what's in the database chunk_row = { "Question ID": question_id, "Chunk Text": chunk["text"], - "Vector Similarity": chunk[ - "similarity_score" - ], # Raw value from DB + "Vector Similarity": chunk["similarity_score"], # Raw value from DB "LLM Score": chunk.get("llm_score"), # Raw value from DB "Is Evidence": chunk.get("is_evidence"), # Raw value from DB "Position": chunk.get("chunk_order"), # Raw value from DB @@ -100,9 +90,7 @@ def create_analysis_dataframes( ) except Exception as e: - logger.error( - f"Error processing result for question {question_id}: {str(e)}" - ) + logger.error(f"Error processing result for question {question_id}: {str(e)}") logger.error(f"Result data: {data}") continue @@ -111,18 +99,14 @@ def create_analysis_dataframes( chunks_df = pd.DataFrame(chunks_rows) if chunks_rows else pd.DataFrame() # Log DataFrame information - logger.info( - f"Created dataframes - Analysis: {len(analysis_df)} rows, Chunks: {len(chunks_df)} rows" - ) + logger.info(f"Created dataframes - Analysis: {len(analysis_df)} rows, Chunks: {len(chunks_df)} rows") if not chunks_df.empty: logger.info(f"Chunks columns: {chunks_df.columns.tolist()}") logger.info( f"Vector similarity range: {chunks_df['Vector Similarity'].min():.4f} - {chunks_df['Vector Similarity'].max():.4f}" ) if "LLM Score" in chunks_df.columns: - logger.info( - f"LLM score range: {chunks_df['LLM Score'].min():.4f} - {chunks_df['LLM Score'].max():.4f}" - ) + logger.info(f"LLM score range: {chunks_df['LLM Score'].min():.4f} - {chunks_df['LLM Score'].max():.4f}") return analysis_df, chunks_df @@ -139,9 +123,7 @@ def is_chunk_referenced(position: int, evidence_list: List[Dict]) -> bool: return False -def create_combined_dataframe( - analysis_df: pd.DataFrame, chunks_df: pd.DataFrame -) -> pd.DataFrame: +def create_combined_dataframe(analysis_df: pd.DataFrame, chunks_df: pd.DataFrame) -> pd.DataFrame: """ Create a combined dataframe with analysis results and their chunks. diff --git a/report_analyst/core/document_processor.py b/report_analyst/core/document_processor.py index 3267da319..c33b8ff48 100644 --- a/report_analyst/core/document_processor.py +++ b/report_analyst/core/document_processor.py @@ -65,9 +65,7 @@ async def process_upload(self, file_path: Union[str, Path]) -> dict: dest_path.unlink() # Clean up on error raise - async def _extract_metadata( - self, file_path: Path, original_filename: str - ) -> DocumentMetadata: + async def _extract_metadata(self, file_path: Path, original_filename: str) -> DocumentMetadata: """Extract metadata from the document using PyMuPDF""" try: file_size = file_path.stat().st_size diff --git a/report_analyst/core/document_sources.py b/report_analyst/core/document_sources.py index 94fecc264..1b505adba 100644 --- a/report_analyst/core/document_sources.py +++ b/report_analyst/core/document_sources.py @@ -44,9 +44,7 @@ async def upload_document(self, file_path: Union[str, Path]) -> str: pass @abstractmethod - async def get_chunks( - self, document_id: str, configuration: Optional[Dict[str, Any]] = None - ) -> List[DocumentChunk]: + async def get_chunks(self, document_id: str, configuration: Optional[Dict[str, Any]] = None) -> List[DocumentChunk]: """ Get chunks for a document. @@ -113,9 +111,7 @@ async def upload_document(self, file_path: Union[str, Path]) -> str: return document_id - async def get_chunks( - self, document_id: str, configuration: Optional[Dict[str, Any]] = None - ) -> List[DocumentChunk]: + async def get_chunks(self, document_id: str, configuration: Optional[Dict[str, Any]] = None) -> List[DocumentChunk]: """Get chunks using existing analyzer logic""" if document_id not in self._document_cache: raise ValueError(f"Document {document_id} not found") diff --git a/report_analyst/core/file_storage.py b/report_analyst/core/file_storage.py index 478fad5a2..480a23a5c 100644 --- a/report_analyst/core/file_storage.py +++ b/report_analyst/core/file_storage.py @@ -10,12 +10,20 @@ import logging import os import uuid +from datetime import datetime from pathlib import Path from typing import Optional -from sqlalchemy import LargeBinary, MetaData, Table, Column, Text, DateTime, String -from sqlalchemy import text -from datetime import datetime +from sqlalchemy import ( + Column, + DateTime, + LargeBinary, + MetaData, + String, + Table, + Text, + text, +) from .database_manager import DatabaseManager @@ -24,16 +32,17 @@ class FileStorageError(Exception): """Exception raised for file storage errors""" + pass class PostgreSQLFileStorage: """Store files in PostgreSQL database as BYTEA/BLOB""" - + def __init__(self, database_url: Optional[str] = None): """ Initialize PostgreSQL file storage. - + Args: database_url: Database connection string. If None, uses environment variable. """ @@ -41,18 +50,18 @@ def __init__(self, database_url: Optional[str] = None): database_url = os.getenv("DATABASE_URL") if database_url is None: raise FileStorageError("DATABASE_URL not set for PostgreSQL file storage") - + self.db_manager = DatabaseManager(database_url) if not self.db_manager.is_postgres(): raise FileStorageError("PostgreSQL file storage requires PostgreSQL database") - + # Check if we should use Alembic migrations instead of auto-creation use_alembic = os.getenv("USE_ALEMBIC_MIGRATIONS", "false").lower() == "true" if use_alembic: logger.info("Using Alembic migrations - skipping auto table creation for stored_files") else: self._init_table() - + def _init_table(self): """Initialize the stored_files table""" try: @@ -67,59 +76,64 @@ def _init_table(self): Column("file_size", Text, nullable=False), # Store as string for large files Column("created_at", DateTime, default=datetime.now), ) - + engine = self.db_manager.get_engine() metadata.create_all(engine, checkfirst=True) logger.info("stored_files table initialized") except Exception as e: logger.error(f"Error initializing stored_files table: {str(e)}") raise FileStorageError(f"Failed to initialize file storage table: {str(e)}") - + def store_file(self, file_bytes: bytes, filename: str, content_type: Optional[str] = None) -> str: """ Store file in PostgreSQL and return file ID. - + Args: file_bytes: File content as bytes filename: Original filename content_type: MIME type (optional) - + Returns: file_id: Unique identifier for the stored file """ try: file_id = str(uuid.uuid4()) file_size = len(file_bytes) - + with self.db_manager.get_connection() as conn: # Use parameterized query for safety - query = text(""" + query = text( + """ INSERT INTO stored_files (id, filename, file_data, content_type, file_size, created_at) VALUES (:id, :filename, :file_data, :content_type, :file_size, :created_at) - """) - conn.execute(query, { - "id": file_id, - "filename": filename, - "file_data": file_bytes, - "content_type": content_type or "application/pdf", - "file_size": str(file_size), - "created_at": datetime.now() - }) + """ + ) + conn.execute( + query, + { + "id": file_id, + "filename": filename, + "file_data": file_bytes, + "content_type": content_type or "application/pdf", + "file_size": str(file_size), + "created_at": datetime.now(), + }, + ) conn.commit() - + logger.info(f"Stored file {filename} (ID: {file_id}, size: {file_size} bytes) in PostgreSQL") return file_id except Exception as e: logger.error(f"Error storing file in PostgreSQL: {str(e)}") raise FileStorageError(f"Failed to store file: {str(e)}") - + def retrieve_file(self, file_id: str) -> Optional[bytes]: """ Retrieve file from PostgreSQL. - + Args: file_id: Unique identifier for the stored file - + Returns: File content as bytes, or None if not found """ @@ -128,52 +142,54 @@ def retrieve_file(self, file_id: str) -> Optional[bytes]: query = text("SELECT file_data FROM stored_files WHERE id = :file_id") result = conn.execute(query, {"file_id": file_id}) row = result.fetchone() - + if row: return bytes(row[0]) return None except Exception as e: logger.error(f"Error retrieving file {file_id} from PostgreSQL: {str(e)}") raise FileStorageError(f"Failed to retrieve file: {str(e)}") - + def get_file_info(self, file_id: str) -> Optional[dict]: """ Get file metadata without retrieving the file data. - + Args: file_id: Unique identifier for the stored file - + Returns: Dictionary with file info, or None if not found """ try: with self.db_manager.get_connection() as conn: - query = text(""" + query = text( + """ SELECT filename, content_type, file_size, created_at FROM stored_files WHERE id = :file_id - """) + """ + ) result = conn.execute(query, {"file_id": file_id}) row = result.fetchone() - + if row: return { "filename": row[0], "content_type": row[1], "file_size": int(row[2]) if row[2] else 0, - "created_at": row[3] + "created_at": row[3], } return None except Exception as e: logger.error(f"Error getting file info for {file_id}: {str(e)}") return None - + def delete_file(self, file_id: str) -> bool: """ Delete file from PostgreSQL. - + Args: file_id: Unique identifier for the stored file - + Returns: True if deleted, False if not found """ @@ -186,15 +202,15 @@ def delete_file(self, file_id: str) -> bool: except Exception as e: logger.error(f"Error deleting file {file_id}: {str(e)}") return False - + def save_to_temp(self, file_id: str, temp_dir: Path = Path("temp")) -> Optional[str]: """ Retrieve file from PostgreSQL and save to temporary directory. - + Args: file_id: Unique identifier for the stored file temp_dir: Directory to save the file to - + Returns: Path to the temporary file, or None if not found """ @@ -202,19 +218,19 @@ def save_to_temp(self, file_id: str, temp_dir: Path = Path("temp")) -> Optional[ file_info = self.get_file_info(file_id) if not file_info: return None - + file_bytes = self.retrieve_file(file_id) if not file_bytes: return None - + # Create temp directory if it doesn't exist temp_dir.mkdir(parents=True, exist_ok=True) - + # Save to temp file temp_path = temp_dir / file_info["filename"] with open(temp_path, "wb") as f: f.write(file_bytes) - + logger.info(f"Retrieved file {file_id} to {temp_path}") return str(temp_path) except Exception as e: @@ -222,13 +238,15 @@ def save_to_temp(self, file_id: str, temp_dir: Path = Path("temp")) -> Optional[ return None -def get_file_storage(database_url: Optional[str] = None) -> Optional[PostgreSQLFileStorage]: +def get_file_storage( + database_url: Optional[str] = None, +) -> Optional[PostgreSQLFileStorage]: """ Get file storage instance if PostgreSQL is configured. - + Args: database_url: Database connection string (optional) - + Returns: PostgreSQLFileStorage instance if PostgreSQL is configured, None otherwise """ @@ -237,16 +255,15 @@ def get_file_storage(database_url: Optional[str] = None) -> Optional[PostgreSQLF use_postgres_storage = os.getenv("USE_POSTGRES_FILE_STORAGE", "false").lower() == "true" if not use_postgres_storage: return None - + # Check if we have a PostgreSQL database if database_url is None: database_url = os.getenv("DATABASE_URL") - + if database_url and database_url.startswith(("postgresql://", "postgres://")): return PostgreSQLFileStorage(database_url) - + return None except Exception as e: logger.warning(f"PostgreSQL file storage not available: {str(e)}") return None - diff --git a/report_analyst/core/llm_providers.py b/report_analyst/core/llm_providers.py index c6a0df4f0..0fd9808c3 100644 --- a/report_analyst/core/llm_providers.py +++ b/report_analyst/core/llm_providers.py @@ -34,12 +34,8 @@ def get_llm(model_name: str, cache_dir: Optional[str] = None, **kwargs) -> Any: if model_name.startswith("gpt-"): api_key = os.getenv("OPENAI_API_KEY") if not api_key: - logger.error( - f"Cannot initialize OpenAI model '{model_name}' - OPENAI_API_KEY environment variable is not set" - ) - raise ValueError( - "OPENAI_API_KEY environment variable is required for OpenAI models" - ) + logger.error(f"Cannot initialize OpenAI model '{model_name}' - OPENAI_API_KEY environment variable is not set") + raise ValueError("OPENAI_API_KEY environment variable is required for OpenAI models") return OpenAI( model=model_name, @@ -53,12 +49,8 @@ def get_llm(model_name: str, cache_dir: Optional[str] = None, **kwargs) -> Any: elif model_name.startswith("gemini-") or model_name.startswith("models/gemini-"): api_key = os.getenv("GOOGLE_API_KEY") if not api_key: - logger.error( - f"Cannot initialize Gemini model '{model_name}' - GOOGLE_API_KEY environment variable is not set" - ) - raise ValueError( - "GOOGLE_API_KEY environment variable is required for Gemini models" - ) + logger.error(f"Cannot initialize Gemini model '{model_name}' - GOOGLE_API_KEY environment variable is not set") + raise ValueError("GOOGLE_API_KEY environment variable is required for Gemini models") # Use the full model path if provided, otherwise prefix with "models/" full_model_name = model_name diff --git a/report_analyst/core/migration_utils.py b/report_analyst/core/migration_utils.py index 2985de3ee..0800f72bc 100644 --- a/report_analyst/core/migration_utils.py +++ b/report_analyst/core/migration_utils.py @@ -21,46 +21,46 @@ def get_alembic_config(database_url: Optional[str] = None) -> Config: """ Get Alembic configuration. - + Args: database_url: Database connection string. If None, uses environment variable. - + Returns: Alembic Config object """ alembic_ini_path = os.getenv("ALEMBIC_CONFIG", "alembic.ini") config = Config(alembic_ini_path) - + # Set database URL if provided if database_url: config.set_main_option("sqlalchemy.url", database_url) elif os.getenv("DATABASE_URL"): config.set_main_option("sqlalchemy.url", os.getenv("DATABASE_URL")) - + return config def get_current_revision(database_url: Optional[str] = None) -> Optional[str]: """ Get the current database revision. - + Args: database_url: Database connection string. If None, uses environment variable. - + Returns: Current revision string, or None if no migrations have been applied """ try: config = get_alembic_config(database_url) - + # Get database manager to create engine if database_url: db_manager = DatabaseManager(database_url) else: db_manager = DatabaseManager() - + engine = db_manager.get_engine() - + with engine.connect() as connection: context = MigrationContext.configure(connection) current_rev = context.get_current_revision() @@ -73,7 +73,7 @@ def get_current_revision(database_url: Optional[str] = None) -> Optional[str]: def get_head_revision() -> Optional[str]: """ Get the head (latest) migration revision. - + Returns: Head revision string, or None if no migrations exist """ @@ -90,25 +90,25 @@ def get_head_revision() -> Optional[str]: def needs_migration(database_url: Optional[str] = None) -> bool: """ Check if database needs migration. - + Args: database_url: Database connection string. If None, uses environment variable. - + Returns: True if migration is needed, False otherwise """ try: current_rev = get_current_revision(database_url) head_rev = get_head_revision() - + if head_rev is None: # No migrations exist return False - + if current_rev is None: # Database has no migrations applied return True - + return current_rev != head_rev except Exception as e: logger.error(f"Error checking migration status: {str(e)}") @@ -118,21 +118,21 @@ def needs_migration(database_url: Optional[str] = None) -> bool: def check_migration_status(database_url: Optional[str] = None) -> dict: """ Check migration status and return detailed information. - + Args: database_url: Database connection string. If None, uses environment variable. - + Returns: Dictionary with migration status information """ try: current_rev = get_current_revision(database_url) head_rev = get_head_revision() - + return { "current_revision": current_rev, "head_revision": head_rev, - "is_up_to_date": current_rev == head_rev if (current_rev and head_rev) else False, + "is_up_to_date": (current_rev == head_rev if (current_rev and head_rev) else False), "needs_migration": needs_migration(database_url), } except Exception as e: @@ -149,11 +149,11 @@ def check_migration_status(database_url: Optional[str] = None) -> dict: def run_migrations(database_url: Optional[str] = None, revision: str = "head") -> bool: """ Run database migrations. - + Args: database_url: Database connection string. If None, uses environment variable. revision: Target revision (default: "head") - + Returns: True if successful, False otherwise """ @@ -165,4 +165,3 @@ def run_migrations(database_url: Optional[str] = None, revision: str = "head") - except Exception as e: logger.error(f"Error running migrations: {str(e)}") return False - diff --git a/report_analyst/core/question_loader.py b/report_analyst/core/question_loader.py index 3c2349add..48778c91c 100644 --- a/report_analyst/core/question_loader.py +++ b/report_analyst/core/question_loader.py @@ -95,14 +95,10 @@ def _load_question_sets(self) -> Dict[str, QuestionSet]: question_sets[question_set_id] = question_set - logger.info( - f"[QUESTION_LOADER] ✓ Loaded {len(questions)} questions for {question_set_id}" - ) + logger.info(f"[QUESTION_LOADER] ✓ Loaded {len(questions)} questions for {question_set_id}") except Exception as e: - logger.error( - f"[QUESTION_LOADER] Error loading question set from {yaml_file}: {str(e)}" - ) + logger.error(f"[QUESTION_LOADER] Error loading question set from {yaml_file}: {str(e)}") continue if not question_sets: @@ -129,10 +125,7 @@ def get_question_set_names(self) -> Dict[str, str]: def get_question_set_info(self) -> Dict[str, Dict[str, str]]: """Get question set info (name and description) for UI display""" question_sets = self.get_question_sets() - return { - qset.id: {"name": qset.name, "description": qset.description} - for qset in question_sets.values() - } + return {qset.id: {"name": qset.name, "description": qset.description} for qset in question_sets.values()} def get_questions(self, question_set_id: str) -> Dict[str, Dict[str, str]]: """Get questions for a specific question set""" @@ -146,9 +139,7 @@ def get_question_set_options(self) -> List[str]: def reload(self): """Force reload of question sets""" self._question_sets = None - logger.info( - "[QUESTION_LOADER] Question sets cache cleared, will reload on next access" - ) + logger.info("[QUESTION_LOADER] Question sets cache cleared, will reload on next access") # Global singleton instance diff --git a/report_analyst/core/report_data_client.py b/report_analyst/core/report_data_client.py index 463322460..529672d27 100644 --- a/report_analyst/core/report_data_client.py +++ b/report_analyst/core/report_data_client.py @@ -35,7 +35,7 @@ def is_local_resource(self) -> bool: def parse_backend_urn(self) -> Optional[Dict[str, str]]: """Parse URN: urn:report-analyst:backend:host:resource_id - + Handles both with and without ports: - urn:report-analyst:backend:localhost:8000:abc-123 (host:port:resource_id) - urn:report-analyst:backend:api.example.com:abc-123 (host:resource_id) @@ -64,11 +64,7 @@ def resolve_to_http_url(self) -> Optional[str]: if not parsed: return None # Determine protocol (http vs https) - could be config-based - protocol = ( - "http" - if "localhost" in parsed["host"] or "127.0.0.1" in parsed["host"] - else "https" - ) + protocol = "http" if "localhost" in parsed["host"] or "127.0.0.1" in parsed["host"] else "https" return f"{protocol}://{parsed['host']}/resources/{parsed['resource_id']}" @@ -79,9 +75,7 @@ def __init__(self, temp_dir: Path = Path("temp")): self.temp_dir = temp_dir self._backend_clients: Dict[str, Any] = {} # Cache backend clients by host - def list_reports( - self, backend_configs: Optional[List[Any]] = None - ) -> List[ReportResource]: + def list_reports(self, backend_configs: Optional[List[Any]] = None) -> List[ReportResource]: """ List all available sustainability reports from all sources. @@ -117,9 +111,7 @@ def _list_local_reports(self) -> List[ReportResource]: file_size = file.stat().st_size if file_size < 100: # Minimum size for a valid PDF - logger.warning( - f"Skipping {file.name}: file too small ({file_size} bytes), likely invalid" - ) + logger.warning(f"Skipping {file.name}: file too small ({file_size} bytes), likely invalid") continue # Try to validate it's a real PDF by attempting to open it @@ -130,14 +122,10 @@ def _list_local_reports(self) -> List[ReportResource]: page_count = doc.page_count doc.close() if page_count == 0: - logger.warning( - f"Skipping {file.name}: PDF has 0 pages, likely invalid" - ) + logger.warning(f"Skipping {file.name}: PDF has 0 pages, likely invalid") continue except Exception as e: - logger.warning( - f"Skipping {file.name}: cannot open as PDF ({str(e)})" - ) + logger.warning(f"Skipping {file.name}: cannot open as PDF ({str(e)})") continue # Create file:// URI @@ -151,9 +139,7 @@ def _list_local_reports(self) -> List[ReportResource]: metadata={"path": str(file.resolve()), "pages": page_count}, ) ) - logger.info( - f"Found valid PDF: {file.name}, size: {file_size} bytes, pages: {page_count}" - ) + logger.info(f"Found valid PDF: {file.name}, size: {file_size} bytes, pages: {page_count}") return files @@ -170,9 +156,7 @@ def _list_backend_reports(self, config: Any) -> List[ReportResource]: return [] -def get_backend_service_for_urn( - urn: str, backend_configs: List[Any] -) -> Optional[Any]: +def get_backend_service_for_urn(urn: str, backend_configs: List[Any]) -> Optional[Any]: """Get BackendService instance for a given backend URN""" if not urn.startswith("urn:report-analyst:backend:"): return None @@ -195,9 +179,7 @@ def get_backend_service_for_urn( return None -def get_chunks_for_backend_resource( - urn: str, backend_configs: List[Any] -) -> Optional[List[Dict[str, Any]]]: +def get_chunks_for_backend_resource(urn: str, backend_configs: List[Any]) -> Optional[List[Dict[str, Any]]]: """ Get chunks for a backend resource identified by URN. @@ -219,4 +201,3 @@ def get_chunks_for_backend_resource( # Use BackendService to get chunks return backend_service.get_chunks(parsed["resource_id"]) - diff --git a/report_analyst/core/storage/llama_store.py b/report_analyst/core/storage/llama_store.py index c53bb357a..ac36e2086 100644 --- a/report_analyst/core/storage/llama_store.py +++ b/report_analyst/core/storage/llama_store.py @@ -49,9 +49,7 @@ def load(self) -> bool: # Load the existing index logger.info(f"Loading vector store from {self.storage_path}") - self.storage_context = StorageContext.from_defaults( - persist_dir=str(self.storage_path) - ) + self.storage_context = StorageContext.from_defaults(persist_dir=str(self.storage_path)) self.store = load_index_from_storage(storage_context=self.storage_context) logger.info("Successfully loaded existing vector store") return True @@ -124,10 +122,7 @@ def similarity_search(self, query: str, k: int = 4) -> List[Tuple[Document, floa # Get nodes and convert to Document format with scores nodes = retriever.retrieve(query) - return [ - (Document(text=node.node.text, metadata=node.node.metadata), node.score) - for node in nodes - ] + return [(Document(text=node.node.text, metadata=node.node.metadata), node.score) for node in nodes] except Exception as e: logger.error(f"Error during similarity search: {e}") diff --git a/report_analyst/core/workflow_engine.py b/report_analyst/core/workflow_engine.py index 3a35684a0..74b52b477 100644 --- a/report_analyst/core/workflow_engine.py +++ b/report_analyst/core/workflow_engine.py @@ -91,9 +91,7 @@ async def execute(self, context: WorkflowContext) -> WorkflowContext: "id": chunk.get("id"), "text": chunk["text"], "chunk_order": i, - "similarity_score": chunk.get( - "similarity_score", chunk.get("score", 0.0) - ), + "similarity_score": chunk.get("similarity_score", chunk.get("score", 0.0)), "llm_score": None, # Will be set in LLM scoring step if enabled "is_evidence": False, # Will be set in evidence assignment step "evidence_order": None, @@ -103,9 +101,7 @@ async def execute(self, context: WorkflowContext) -> WorkflowContext: clean_chunks.append(clean_chunk) context.chunks = clean_chunks - logger.info( - f"[WORKFLOW] Retrieved {len(clean_chunks)} chunks with vector similarity" - ) + logger.info(f"[WORKFLOW] Retrieved {len(clean_chunks)} chunks with vector similarity") return context @@ -138,9 +134,7 @@ async def execute(self, context: WorkflowContext) -> WorkflowContext: try: # Score chunks with LLM - llm_scores = await self.llm_manager.score_chunks( - question=context.question_text, chunks=context.chunks - ) + llm_scores = await self.llm_manager.score_chunks(question=context.question_text, chunks=context.chunks) # Update chunks with LLM scores for chunk in context.chunks: @@ -171,9 +165,7 @@ def validate_input(self, context: WorkflowContext) -> bool: return context.chunks is not None and len(context.chunks) > 0 def validate_output(self, context: WorkflowContext) -> bool: - return ( - context.analysis_result is not None and context.evidence_chunks is not None - ) + return context.analysis_result is not None and context.evidence_chunks is not None async def execute(self, context: WorkflowContext) -> WorkflowContext: """Analyze question and extract evidence""" @@ -182,9 +174,7 @@ async def execute(self, context: WorkflowContext) -> WorkflowContext: use_llm_scoring = context.config.get("use_llm_scoring", False) if use_llm_scoring and context.llm_scores: # Sort by LLM score if available - sorted_chunks = sorted( - context.chunks, key=lambda x: x.get("llm_score", 0.0), reverse=True - ) + sorted_chunks = sorted(context.chunks, key=lambda x: x.get("llm_score", 0.0), reverse=True) else: # Sort by vector similarity sorted_chunks = sorted( @@ -194,18 +184,14 @@ async def execute(self, context: WorkflowContext) -> WorkflowContext: ) # Analyze question with LLM - analysis_result = await self.llm_manager.analyze_question( - question=context.question_text, chunks=sorted_chunks - ) + analysis_result = await self.llm_manager.analyze_question(question=context.question_text, chunks=sorted_chunks) # Extract evidence chunk IDs evidence_chunks = analysis_result.get("evidence_chunks", []) context.analysis_result = analysis_result context.evidence_chunks = evidence_chunks - logger.info( - f"[WORKFLOW] Question analysis complete, found {len(evidence_chunks)} evidence chunks" - ) + logger.info(f"[WORKFLOW] Question analysis complete, found {len(evidence_chunks)} evidence chunks") return context @@ -221,11 +207,7 @@ def __init__(self): super().__init__("evidence_assignment") def validate_input(self, context: WorkflowContext) -> bool: - return ( - context.chunks is not None - and context.evidence_chunks is not None - and context.analysis_result is not None - ) + return context.chunks is not None and context.evidence_chunks is not None and context.analysis_result is not None def validate_output(self, context: WorkflowContext) -> bool: return context.chunks is not None @@ -248,9 +230,7 @@ async def execute(self, context: WorkflowContext) -> WorkflowContext: chunk["is_evidence"] = False chunk["evidence_order"] = None - logger.info( - f"[WORKFLOW] Assigned evidence flags to {len(evidence_mapping)} chunks" - ) + logger.info(f"[WORKFLOW] Assigned evidence flags to {len(evidence_mapping)} chunks") return context @@ -284,17 +264,13 @@ async def execute_workflow( logger.info(f"[WORKFLOW] Starting workflow for question: {question_id}") # For now, just return cached results if available - cached_result = self.cache_manager.get_analysis( - file_path=file_path, config=config, question_ids=[question_id] - ) + cached_result = self.cache_manager.get_analysis(file_path=file_path, config=config, question_ids=[question_id]) if cached_result and question_id in cached_result: logger.info(f"[WORKFLOW] Found cached result for {question_id}") return cached_result[question_id] else: - logger.warning( - f"[WORKFLOW] No cached result for {question_id}, full workflow not implemented" - ) + logger.warning(f"[WORKFLOW] No cached result for {question_id}, full workflow not implemented") return { "error": "Full workflow not implemented - only cached results available", "question_id": question_id, diff --git a/report_analyst/gradio_app.py b/report_analyst/gradio_app.py index b78c3d5ad..705f1aa9e 100644 --- a/report_analyst/gradio_app.py +++ b/report_analyst/gradio_app.py @@ -30,12 +30,8 @@ class DocumentService: def __init__(self): self.analyzer = DocumentAnalyzer() # Get valid question IDs from the loaded questions - self.valid_question_ids = list( - range(1, len(self.analyzer.questions["TCFD Analysis"]["questions"]) + 1) - ) - logger.info( - f"Initialized with {len(self.valid_question_ids)} valid question IDs" - ) + self.valid_question_ids = list(range(1, len(self.analyzer.questions["TCFD Analysis"]["questions"]) + 1)) + logger.info(f"Initialized with {len(self.valid_question_ids)} valid question IDs") def validate_question_ids(self, question_ids: List[int]) -> List[int]: """Validate and filter question IDs""" @@ -49,9 +45,7 @@ def validate_question_ids(self, question_ids: List[int]) -> List[int]: logger.info(f"Validated question IDs: {valid_ids}") return valid_ids - async def process_document( - self, file_path: str, question_ids: List[int] = None - ) -> AsyncGenerator[Dict, None]: + async def process_document(self, file_path: str, question_ids: List[int] = None) -> AsyncGenerator[Dict, None]: """Process uploaded document and stream analysis results""" if not file_path: yield {"error": "No file uploaded"} @@ -70,12 +64,8 @@ async def process_document( temp_file = Path(tempfile.gettempdir()) / f"temp_{uuid.uuid4()}.pdf" try: shutil.copy2(file_path, temp_file) - async for result in self.analyzer.process_document( - str(temp_file), question_ids - ): - logger.info( - f"Processing section: {result.get('section', 'unknown')}" - ) + async for result in self.analyzer.process_document(str(temp_file), question_ids): + logger.info(f"Processing section: {result.get('section', 'unknown')}") yield result finally: if temp_file.exists(): @@ -112,9 +102,7 @@ def create_app(): gr.Markdown("Upload a sustainability report for detailed TCFD analysis") with gr.Row(): - file_input = gr.File( - label="Upload PDF Report", file_types=[".pdf"], type="filepath" - ) + file_input = gr.File(label="Upload PDF Report", file_types=[".pdf"], type="filepath") analyze_btn = gr.Button("Start Analysis", variant="primary") with gr.Row(): @@ -164,9 +152,7 @@ async def process_analysis(file, *selected_questions): return try: - selected_ids = [ - i + 1 for i, selected in enumerate(selected_questions) if selected - ] + selected_ids = [i + 1 for i, selected in enumerate(selected_questions) if selected] if not selected_ids: yield "Please select at least one question", [], {} return @@ -198,16 +184,12 @@ async def process_analysis(file, *selected_questions): if analysis.get("evidence"): question_html += "

Evidence:

    " - question_html += "".join( - [f"
  • {e}
    • " for e in analysis["evidence"]] - ) + question_html += "".join([f"
  • {e}
    • " for e in analysis["evidence"]]) question_html += "
" if analysis.get("gaps"): question_html += "

Gaps:

    " - question_html += "".join( - [f"
  • {g}
    • " for g in analysis["gaps"]] - ) + question_html += "".join([f"
  • {g}
    • " for g in analysis["gaps"]]) question_html += "
" question_html += "" diff --git a/report_analyst/main.py b/report_analyst/main.py index 2727f39b4..d5952192d 100644 --- a/report_analyst/main.py +++ b/report_analyst/main.py @@ -61,9 +61,7 @@ async def upload_document(file: UploadFile = File(...)): async def analyze_document(request: AnalysisRequest): """Analyze a document with specified parameters""" try: - result = await document_analyzer.analyze( - request.document_id, request.analysis_type - ) + result = await document_analyzer.analyze(request.document_id, request.analysis_type) return result except Exception as e: raise HTTPException(status_code=400, detail=str(e)) @@ -73,9 +71,7 @@ async def analyze_document(request: AnalysisRequest): async def ask_question(request: QuestionRequest): """Ask a question about a document""" try: - result = await document_analyzer.ask_question( - request.document_id, request.question, request.context - ) + result = await document_analyzer.ask_question(request.document_id, request.question, request.context) return result except Exception as e: raise HTTPException(status_code=400, detail=str(e)) diff --git a/report_analyst/models/requests.py b/report_analyst/models/requests.py index b6e1eff99..a56885de2 100644 --- a/report_analyst/models/requests.py +++ b/report_analyst/models/requests.py @@ -14,20 +14,14 @@ class AnalysisType(str, Enum): class AnalysisRequest(BaseModel): document_id: str = Field(..., description="The ID of the uploaded document") - analysis_type: AnalysisType = Field( - default=AnalysisType.GENERAL, description="The type of analysis to perform" - ) - custom_instructions: Optional[str] = Field( - None, description="Custom instructions for the analysis" - ) + analysis_type: AnalysisType = Field(default=AnalysisType.GENERAL, description="The type of analysis to perform") + custom_instructions: Optional[str] = Field(None, description="Custom instructions for the analysis") class QuestionRequest(BaseModel): document_id: str = Field(..., description="The ID of the uploaded document") question: str = Field(..., description="The question to ask about the document") - context: Optional[str] = Field( - None, description="Additional context for the question" - ) + context: Optional[str] = Field(None, description="Additional context for the question") class DocumentMetadata(BaseModel): diff --git a/report_analyst/models/responses.py b/report_analyst/models/responses.py index cf1d29102..d75e762af 100644 --- a/report_analyst/models/responses.py +++ b/report_analyst/models/responses.py @@ -12,9 +12,7 @@ class AnalysisResponse(BaseModel): key_points: List[str] topics: List[Dict[str, float]] metadata: DocumentMetadata - confidence_score: float = Field( - ..., ge=0.0, le=1.0, description="Confidence score of the analysis" - ) + confidence_score: float = Field(..., ge=0.0, le=1.0, description="Confidence score of the analysis") class QuestionResponse(BaseModel): @@ -22,9 +20,7 @@ class QuestionResponse(BaseModel): question: str answer: str context_used: Optional[str] = None - confidence_score: float = Field( - ..., ge=0.0, le=1.0, description="Confidence score of the answer" - ) + confidence_score: float = Field(..., ge=0.0, le=1.0, description="Confidence score of the answer") relevant_quotes: List[str] = Field( default_factory=list, description="Relevant quotes from the document supporting the answer", diff --git a/report_analyst/streamlit_app.py b/report_analyst/streamlit_app.py index aa87191f4..e2d3f7fec 100644 --- a/report_analyst/streamlit_app.py +++ b/report_analyst/streamlit_app.py @@ -144,9 +144,7 @@ def __init__(self): # Initialize the real document analyzer self.analyzer = DocumentAnalyzer() self.prompt_manager = PromptManager() - self.cache_manager = ( - self.analyzer.cache_manager - ) # Access the cache manager from the analyzer + self.cache_manager = self.analyzer.cache_manager # Access the cache manager from the analyzer def load_question_set(self, question_set: str) -> Dict: """Load questions from the specified question set using centralized loader""" @@ -187,7 +185,7 @@ async def analyze_document( pre_retrieved_chunks: Optional[List[Dict[str, Any]]] = None, ) -> AsyncGenerator[Dict, None]: """Analyze a document using the provided questions - + Args: file_path: Path to document file or URN for backend resources questions: Dictionary of questions @@ -206,14 +204,10 @@ async def analyze_document( self.analyzer.questions = questions # Convert selected question IDs to numbers for the analyzer - selected_numbers = [ - questions[q_id]["number"] for q_id in selected_questions - ] + selected_numbers = [questions[q_id]["number"] for q_id in selected_questions] # Get the question set prefix from the first selected question - question_set = ( - selected_questions[0].split("_")[0] if selected_questions else "tcfd" - ) + question_set = selected_questions[0].split("_")[0] if selected_questions else "tcfd" self.analyzer.question_set = question_set # Pass use_llm_scoring to process_document @@ -269,9 +263,7 @@ def process_document( force_recompute: bool = False, ): """Delegate to the analyzer's process_document method""" - return self.analyzer.process_document( - file_path, selected_questions, use_llm_scoring, single_call, force_recompute - ) + return self.analyzer.process_document(file_path, selected_questions, use_llm_scoring, single_call, force_recompute) def save_uploaded_file(uploaded_file) -> Optional[str]: @@ -292,29 +284,25 @@ def save_uploaded_file(uploaded_file) -> Optional[str]: # Get file bytes file_bytes = uploaded_file.getbuffer() - + # Check if PostgreSQL file storage is enabled use_postgres_storage = st.session_state.get("use_postgres_file_storage", False) - + if use_postgres_storage: try: from report_analyst.core.file_storage import get_file_storage - + # Get database URL from session state or environment database_url = st.session_state.get("database_url") file_storage = get_file_storage(database_url) - + if file_storage: # Store in PostgreSQL - file_id = file_storage.store_file( - file_bytes, - uploaded_file.name, - uploaded_file.type - ) - + file_id = file_storage.store_file(file_bytes, uploaded_file.name, uploaded_file.type) + # Save to temp for processing (retrieve from DB) temp_path = file_storage.save_to_temp(file_id) - + if temp_path: # Store both file_id and path in session state st.session_state[file_key] = temp_path @@ -328,7 +316,7 @@ def save_uploaded_file(uploaded_file) -> Optional[str]: logger.warning("PostgreSQL file storage not available, falling back to local") except Exception as e: logger.warning(f"PostgreSQL file storage failed: {str(e)}, falling back to local") - + # Fallback to local file storage file_path = Path("temp") / uploaded_file.name file_path.parent.mkdir(parents=True, exist_ok=True) @@ -373,9 +361,7 @@ def display_dataframes(analysis_df: pd.DataFrame, chunks_df: pd.DataFrame): chunks_df, use_container_width=True, column_config={ - "Vector Similarity": st.column_config.NumberColumn( - "Vector Similarity", format="%.3f" - ), + "Vector Similarity": st.column_config.NumberColumn("Vector Similarity", format="%.3f"), "LLM Score": st.column_config.NumberColumn("LLM Score", format="%.3f"), "Chunk Text": st.column_config.TextColumn("Chunk Text", width="large"), }, @@ -387,9 +373,7 @@ def convert_df(df: pd.DataFrame) -> bytes: return df.to_csv(index=False).encode("utf-8") -def display_download_buttons( - analysis_df: pd.DataFrame, chunks_df: pd.DataFrame, file_key: str -): +def display_download_buttons(analysis_df: pd.DataFrame, chunks_df: pd.DataFrame, file_key: str): """Display download buttons for analysis results""" # Generate unique timestamp for this render timestamp = int(time.time() * 1000) @@ -440,11 +424,7 @@ async def analyze_document_and_display( """Analyze document and display results as they come in""" try: selected_questions_list = list(selected_questions) if selected_questions else [] - question_set = ( - selected_questions_list[0].split("_")[0] - if selected_questions_list - else "tcfd" - ) + question_set = selected_questions_list[0].split("_")[0] if selected_questions_list else "tcfd" # Use the helper function to generate file key file_key = generate_file_key(file_path, st) @@ -483,9 +463,7 @@ async def analyze_document_and_display( ) if cached_answers: - log_analysis_step( - f"Found {len(cached_answers)} cached answers for {file_key}" - ) + log_analysis_step(f"Found {len(cached_answers)} cached answers for {file_key}") # Show cache info st.info(f"📁 Loading results from stored: {file_key}") @@ -494,29 +472,19 @@ async def analyze_document_and_display( st.session_state.results["answers"][q_id] = answer # Update display with cached results - logger.info( - f"Creating dataframes with cached results for file_key: {file_key}" - ) + logger.info(f"Creating dataframes with cached results for file_key: {file_key}") logger.info( f"Current session state settings: chunk_size={st.session_state.get('new_chunk_size')}, overlap={st.session_state.get('new_overlap')}, top_k={st.session_state.get('new_top_k')}, llm_model={st.session_state.get('new_llm_model')}, use_llm_scoring={st.session_state.get('new_llm_scoring')}" ) - analysis_df, chunks_df = create_analysis_dataframes( - st.session_state.results["answers"], file_key - ) + analysis_df, chunks_df = create_analysis_dataframes(st.session_state.results["answers"], file_key) st.session_state.analysis_df = analysis_df st.session_state.chunks_df = chunks_df # Determine which questions need processing - questions_to_process = [ - q_id - for q_id in selected_questions_list - if force_recompute or q_id not in cached_answers - ] + questions_to_process = [q_id for q_id in selected_questions_list if force_recompute or q_id not in cached_answers] if questions_to_process: - log_analysis_step( - f"Processing {len(questions_to_process)} uncached questions..." - ) + log_analysis_step(f"Processing {len(questions_to_process)} uncached questions...") # Update analyzer with question set report_analyzer.analyzer.question_set = question_set @@ -524,7 +492,7 @@ async def analyze_document_and_display( # Process only uncached questions # Check if we have pre-retrieved chunks (for backend resources) pre_retrieved_chunks = st.session_state.get("backend_chunks") - + async for result in report_analyzer.analyze_document( file_path, questions, @@ -538,9 +506,7 @@ async def analyze_document_and_display( log_analysis_step(f"Received result: {str(result)[:200]}...") if "error" in result: - log_analysis_step( - f"Error received from analyzer: {result['error']}", "error" - ) + log_analysis_step(f"Error received from analyzer: {result['error']}", "error") st.error(f"Analysis error: {result['error']}") continue @@ -550,24 +516,18 @@ async def analyze_document_and_display( question_id = result.get("question_id") if question_id is None: - log_analysis_step( - f"No question_id in result: {str(result)[:200]}...", "warning" - ) + log_analysis_step(f"No question_id in result: {str(result)[:200]}...", "warning") continue # Store results st.session_state.results["answers"][question_id] = result # Update display - logger.info( - f"Creating dataframes with updated results for file_key: {file_key}" - ) + logger.info(f"Creating dataframes with updated results for file_key: {file_key}") logger.info( f"Current session state settings: chunk_size={st.session_state.get('new_chunk_size')}, overlap={st.session_state.get('new_overlap')}, top_k={st.session_state.get('new_top_k')}, llm_model={st.session_state.get('new_llm_model')}, use_llm_scoring={st.session_state.get('new_llm_scoring')}" ) - analysis_df, chunks_df = create_analysis_dataframes( - st.session_state.results["answers"], file_key - ) + analysis_df, chunks_df = create_analysis_dataframes(st.session_state.results["answers"], file_key) st.session_state.analysis_df = analysis_df st.session_state.chunks_df = chunks_df @@ -577,9 +537,7 @@ async def analyze_document_and_display( else: log_analysis_step("All selected questions have cached answers") # Show success message for cached results - st.success( - f"✓ All {len(selected_questions_list)} selected questions loaded from stored" - ) + st.success(f"✓ All {len(selected_questions_list)} selected questions loaded from stored") # Mark this file as analyzed with current configuration if "analyzed_files" not in st.session_state: @@ -719,9 +677,7 @@ def display_final_results(analysis_df: pd.DataFrame, chunks_df: pd.DataFrame): max_value=1, format="%.3f", ), - "Chunk Text": st.column_config.TextColumn( - "Chunk Text", help="Text content of the chunk", width="large" - ), + "Chunk Text": st.column_config.TextColumn("Chunk Text", help="Text content of the chunk", width="large"), "Evidence Reference": st.column_config.CheckboxColumn( "Used as Evidence", help="Whether this chunk was referenced in the analysis", @@ -753,7 +709,9 @@ def get_uploaded_files_history(backend_config=None) -> List[Dict]: client = ReportDataClient() # Collect backend configs (could be multiple backends in future) - backend_configs = [backend_config] if backend_config and hasattr(backend_config, "use_backend") and backend_config.use_backend else [] + backend_configs = ( + [backend_config] if backend_config and hasattr(backend_config, "use_backend") and backend_config.use_backend else [] + ) resources = client.list_reports(backend_configs=backend_configs) @@ -770,20 +728,20 @@ def get_uploaded_files_history(backend_config=None) -> List[Dict]: else: # Backend resource - keep URI as path for compatibility actual_path = r.uri - - result.append({ - "name": r.name, - "uri": r.uri, # Primary identifier (URN or file://) - "path": actual_path, # Actual file path for local files, URI for backend - "date": r.date, - "size": r.size, - }) + + result.append( + { + "name": r.name, + "uri": r.uri, # Primary identifier (URN or file://) + "path": actual_path, # Actual file path for local files, URI for backend + "date": r.date, + "size": r.size, + } + ) return result -def display_analysis_results( - analysis_df: pd.DataFrame, chunks_df: pd.DataFrame, file_key: str = None -) -> None: +def display_analysis_results(analysis_df: pd.DataFrame, chunks_df: pd.DataFrame, file_key: str = None) -> None: """Display analysis results in a consistent format for both individual and consolidated views""" try: if analysis_df.empty: @@ -897,9 +855,7 @@ def display_consolidated_results(analyzer, question_set): # Get the database identifier for the selected question set db_question_set = question_set_mapping.get(question_set, question_set) - logger.info( - f"Mapping question set '{question_set}' to database identifier '{db_question_set}'" - ) + logger.info(f"Mapping question set '{question_set}' to database identifier '{db_question_set}'") # Get all available cache configurations cache_configs = analyzer.analyzer.cache_manager.check_cache_status() @@ -914,9 +870,7 @@ def display_consolidated_results(analyzer, question_set): for config in cache_configs: if len(config) == 6: # Full config row from cache_status file_path, chunk_size, chunk_overlap, top_k, model, qs = config - if ( - qs == db_question_set - ): # Only show configs for selected question set using database identifier + if qs == db_question_set: # Only show configs for selected question set using database identifier if file_path not in file_configs: file_configs[file_path] = [] file_configs[file_path].append( @@ -956,9 +910,7 @@ def display_consolidated_results(analyzer, question_set): ) if selected_config: - logger.info( - f"Getting results for {Path(file_path).name} with config: {selected_config['config']}" - ) + logger.info(f"Getting results for {Path(file_path).name} with config: {selected_config['config']}") # Add similarity search section for document chunks try: @@ -981,9 +933,7 @@ def display_consolidated_results(analyzer, question_set): col1, col2 = st.columns([1, 1]) with col1: # Question dropdown with shorter, cleaner options - question_options = ["None"] + [ - f"{q_id}" for q_id in questions.keys() - ] + question_options = ["None"] + [f"{q_id}" for q_id in questions.keys()] selected_question_id = st.selectbox( "Select a question to sort by similarity:", options=question_options, @@ -992,13 +942,8 @@ def display_consolidated_results(analyzer, question_set): ) # Show selected question text below dropdown - if ( - selected_question_id != "None" - and selected_question_id in questions - ): - st.caption( - f"**{selected_question_id}:** {questions[selected_question_id]['text'][:100]}..." - ) + if selected_question_id != "None" and selected_question_id in questions: + st.caption(f"**{selected_question_id}:** {questions[selected_question_id]['text'][:100]}...") selected_question = selected_question_id with col2: @@ -1017,108 +962,71 @@ def display_consolidated_results(analyzer, question_set): elif selected_question != "None": if selected_question in questions: query_text = questions[selected_question]["text"] - st.info( - f"Using question {selected_question}: {query_text[:100]}..." - ) + st.info(f"Using question {selected_question}: {query_text[:100]}...") # Process chunks chunks_rows = [] - chunks_with_embeddings = [ - c for c in raw_chunks if c.get("embedding") is not None - ] + chunks_with_embeddings = [c for c in raw_chunks if c.get("embedding") is not None] if query_text and chunks_with_embeddings: # Compute similarity scores try: # Check if embeddings are available - if ( - not analyzer.analyzer.embeddings - or analyzer.analyzer.use_backend_llm - ): + if not analyzer.analyzer.embeddings or analyzer.analyzer.use_backend_llm: st.warning( "Embeddings not available for similarity search. Using backend mode or embeddings not initialized." ) query_text = None else: # Get query embedding - query_embedding = ( - analyzer.analyzer.embeddings.get_text_embedding( - query_text - ) - ) - query_embedding = np.array( - query_embedding, dtype=np.float32 - ) + query_embedding = analyzer.analyzer.embeddings.get_text_embedding(query_text) + query_embedding = np.array(query_embedding, dtype=np.float32) # Compute similarity for each chunk similarities = [] for chunk in raw_chunks: if chunk.get("embedding") is not None: - chunk_embedding = np.frombuffer( - chunk["embedding"], dtype=np.float32 - ) + chunk_embedding = np.frombuffer(chunk["embedding"], dtype=np.float32) # Compute cosine similarity - similarity = np.dot( - query_embedding, chunk_embedding - ) / ( - np.linalg.norm(query_embedding) - * np.linalg.norm(chunk_embedding) + similarity = np.dot(query_embedding, chunk_embedding) / ( + np.linalg.norm(query_embedding) * np.linalg.norm(chunk_embedding) ) similarities.append(similarity) else: similarities.append(0.0) # Sort chunks by similarity - chunk_similarity_pairs = list( - zip(raw_chunks, similarities) - ) - chunk_similarity_pairs.sort( - key=lambda x: x[1], reverse=True - ) + chunk_similarity_pairs = list(zip(raw_chunks, similarities)) + chunk_similarity_pairs.sort(key=lambda x: x[1], reverse=True) # Create rows with similarity scores - for i, (chunk, similarity) in enumerate( - chunk_similarity_pairs - ): + for i, (chunk, similarity) in enumerate(chunk_similarity_pairs): chunk_row = { "Rank": i + 1, "Similarity": similarity, - "Text": chunk.get( - "text", chunk.get("chunk_text", "") - ), - "Has Embedding": chunk.get("embedding") - is not None, - "Chunk Size": chunk.get( - "chunk_size", "N/A" - ), - "Chunk Overlap": chunk.get( - "chunk_overlap", "N/A" - ), + "Text": chunk.get("text", chunk.get("chunk_text", "")), + "Has Embedding": chunk.get("embedding") is not None, + "Chunk Size": chunk.get("chunk_size", "N/A"), + "Chunk Overlap": chunk.get("chunk_overlap", "N/A"), } chunks_rows.append(chunk_row) - st.success( - f"✓ Sorted {len(chunks_rows)} chunks by similarity to query" - ) + st.success(f"✓ Sorted {len(chunks_rows)} chunks by similarity to query") except Exception as e: st.error(f"Error computing similarity: {str(e)}") logger.error( - f"Error computing similarity: {str(e)}", exc_info=True + f"Error computing similarity: {str(e)}", + exc_info=True, ) # Fall back to original display for i, chunk in enumerate(raw_chunks): chunk_row = { "Chunk #": i + 1, - "Text": chunk.get( - "text", chunk.get("chunk_text", "") - ), - "Has Embedding": chunk.get("embedding") - is not None, + "Text": chunk.get("text", chunk.get("chunk_text", "")), + "Has Embedding": chunk.get("embedding") is not None, "Chunk Size": chunk.get("chunk_size", "N/A"), - "Chunk Overlap": chunk.get( - "chunk_overlap", "N/A" - ), + "Chunk Overlap": chunk.get("chunk_overlap", "N/A"), } chunks_rows.append(chunk_row) @@ -1201,18 +1109,12 @@ def display_consolidated_results(analyzer, question_set): column_config=column_config, ) - st.info( - f"✓ Found {len(chunks_rows)} total document chunks in this configuration." - ) + st.info(f"✓ Found {len(chunks_rows)} total document chunks in this configuration.") else: - st.warning( - "No chunks found. Run Step 1 to generate document chunks first." - ) + st.warning("No chunks found. Run Step 1 to generate document chunks first.") except Exception as e: - logger.warning( - f"Error displaying document chunks with similarity search: {str(e)}" - ) + logger.warning(f"Error displaying document chunks with similarity search: {str(e)}") # Continue to show analysis results even if chunk display fails # Get cached results @@ -1234,28 +1136,15 @@ def display_consolidated_results(analyzer, question_set): result = data.get("result", {}) analysis_row = { "Question ID": question_id, - "Question Text": ( - questions[question_id]["text"] - if question_id in questions - else question_id - ), + "Question Text": (questions[question_id]["text"] if question_id in questions else question_id), "Analysis": result.get("ANSWER", ""), "Score": float(result.get("SCORE", 0)), - "Key Evidence": "\n".join( - [ - e.get("text", "") - for e in result.get("EVIDENCE", []) - ] - ), + "Key Evidence": "\n".join([e.get("text", "") for e in result.get("EVIDENCE", [])]), "Gaps": "\n".join(result.get("GAPS", [])), - "Sources": ", ".join( - map(str, result.get("SOURCES", [])) - ), + "Sources": ", ".join(map(str, result.get("SOURCES", []))), } analysis_rows.append(analysis_row) - logger.debug( - f"Added analysis row for {question_id}: {json.dumps(analysis_row, indent=2)}" - ) + logger.debug(f"Added analysis row for {question_id}: {json.dumps(analysis_row, indent=2)}") # Process chunks if available if "chunks" in data: @@ -1263,9 +1152,7 @@ def display_consolidated_results(analyzer, question_set): chunk_row = { "Question ID": question_id, "Text": chunk.get("text", ""), - "Vector Similarity": chunk.get( - "similarity_score", 0.0 - ), + "Vector Similarity": chunk.get("similarity_score", 0.0), "LLM Score": chunk.get("llm_score", 0.0), "Is Evidence": chunk.get("is_evidence", False), "Position": chunk.get("chunk_order", 0), @@ -1282,9 +1169,7 @@ def display_consolidated_results(analyzer, question_set): # Create DataFrames if analysis_rows: analysis_df = pd.DataFrame(analysis_rows) - chunks_df = ( - pd.DataFrame(chunks_rows) if chunks_rows else pd.DataFrame() - ) + chunks_df = pd.DataFrame(chunks_rows) if chunks_rows else pd.DataFrame() # Display results using the existing display function file_key = f"{Path(file_path).stem}_cs{selected_config['config']['chunk_size']}" @@ -1315,11 +1200,7 @@ def display_cache_selector(file_path: str): if "analyzer" in st.session_state: # Show cache status using cache manager try: - cache_entries = ( - st.session_state.analyzer.analyzer.cache_manager.check_cache_status( - file_path - ) - ) + cache_entries = st.session_state.analyzer.analyzer.cache_manager.check_cache_status(file_path) if cache_entries: st.text(f"Found {len(cache_entries)} cached configurations:") for entry in cache_entries: @@ -1332,9 +1213,7 @@ def display_cache_selector(file_path: str): with col2: if st.button("Clear Stored Data for File"): try: - st.session_state.analyzer.analyzer.cache_manager.clear_cache( - file_path - ) + st.session_state.analyzer.analyzer.cache_manager.clear_cache(file_path) st.success(f"Stored data cleared for file.") # Clear results from session state if "results" in st.session_state: @@ -1390,17 +1269,13 @@ def update_analyzer_parameters(): if llm_model.startswith("gemini-") and not os.getenv("GOOGLE_API_KEY"): # If somehow a Gemini model was selected but no API key exists logger.error(f"Attempt to use Gemini model '{llm_model}' without API key") - st.error( - f"Cannot use {llm_model} - No Google API key is set. Defaulting to {OPENAI_MODELS[0]}." - ) + st.error(f"Cannot use {llm_model} - No Google API key is set. Defaulting to {OPENAI_MODELS[0]}.") # Reset to default OpenAI model llm_model = OPENAI_MODELS[0] st.session_state.new_llm_model = llm_model elif llm_model.startswith("gpt-") and not os.getenv("OPENAI_API_KEY"): logger.error(f"Attempt to use OpenAI model '{llm_model}' without API key") - st.error( - f"OPENAI_API_KEY environment variable is not set. OpenAI models will not work correctly." - ) + st.error(f"OPENAI_API_KEY environment variable is not set. OpenAI models will not work correctly.") # Update the analyzer with the new parameters try: @@ -1419,9 +1294,7 @@ def update_analyzer_parameters(): # Sync LLM scoring checkbox with session state if "new_llm_scoring" in st.session_state: st.session_state.use_llm_scoring = st.session_state.new_llm_scoring - logger.info( - f"Updated use_llm_scoring to: {st.session_state.use_llm_scoring}" - ) + logger.info(f"Updated use_llm_scoring to: {st.session_state.use_llm_scoring}") except Exception as e: st.error(f"Error updating parameters: {str(e)}") @@ -1457,12 +1330,8 @@ async def run_analysis(analyzer, file_path, selected_questions, progress_text): logger.info(f"[CACHE] Cache HIT for config: {config}") progress_text.success("Found stored results!") st.session_state.results = cached_results - logger.info( - f"[ANALYSIS] Writing results to session state for file: {file_path}" - ) - logger.info( - f"[ANALYSIS] Attempting to display results for file: {file_path}" - ) + logger.info(f"[ANALYSIS] Writing results to session state for file: {file_path}") + logger.info(f"[ANALYSIS] Attempting to display results for file: {file_path}") return else: logger.info(f"[CACHE] Cache MISS for config: {config}") @@ -1471,9 +1340,7 @@ async def run_analysis(analyzer, file_path, selected_questions, progress_text): # Log the LLM scoring setting llm_scoring_enabled = st.session_state.get("new_llm_scoring", False) - progress_text.info( - f"LLM scoring: {'Enabled' if llm_scoring_enabled else 'Disabled'}" - ) + progress_text.info(f"LLM scoring: {'Enabled' if llm_scoring_enabled else 'Disabled'}") logger.info(f"Starting analysis with LLM scoring: {llm_scoring_enabled}") # Track results @@ -1494,14 +1361,12 @@ async def run_analysis(analyzer, file_path, selected_questions, progress_text): # Check if we have pre-retrieved chunks (for backend resources) pre_retrieved_chunks = st.session_state.get("backend_chunks") - + # First update the analyzer's process_document method to use progress_text instead of yielding status async for result in analyzer.process_document( file_path=file_path, selected_questions=question_numbers, # Pass just the numbers - use_llm_scoring=st.session_state.get( - "new_llm_scoring", False - ), # Use the checkbox value directly + use_llm_scoring=st.session_state.get("new_llm_scoring", False), # Use the checkbox value directly force_recompute=st.session_state.get("force_recompute", False), pre_retrieved_chunks=pre_retrieved_chunks, # Pass backend chunks if available ): @@ -1542,9 +1407,7 @@ async def run_analysis(analyzer, file_path, selected_questions, progress_text): final_results = all_results # When writing results to session state - logger.info( - f"[ANALYSIS] Writing results to session state for file: {file_path}" - ) + logger.info(f"[ANALYSIS] Writing results to session state for file: {file_path}") st.session_state.results = final_results logger.info(f"[ANALYSIS] Attempting to display results for file: {file_path}") progress_text.success("Analysis complete!") @@ -1589,9 +1452,7 @@ def main(): st.session_state.file_processed = False # Initialize file processed flag if "analysis_complete" not in st.session_state: - st.session_state.analysis_complete = ( - False # Initialize analysis complete flag - ) + st.session_state.analysis_complete = False # Initialize analysis complete flag # Initialize use_s3_upload in session state if not already set if "use_s3_upload" not in st.session_state: @@ -1605,9 +1466,9 @@ def main(): # Inject Material Icons link tag at the top st.markdown( '', - unsafe_allow_html=True + unsafe_allow_html=True, ) - + # Inject MINIMAL custom CSS for specific customizations only # Let Streamlit handle most theming automatically try: @@ -2775,7 +2636,7 @@ def main(): """ - + st.markdown(custom_css, unsafe_allow_html=True) except Exception as e: # Fallback if theme detection fails @@ -2807,29 +2668,37 @@ def main(): style="width: 90%; max-width: 200px; height: auto;" /> """, - unsafe_allow_html=True + unsafe_allow_html=True, ) except Exception as e: logger.warning(f"Could not load sidebar logo: {str(e)}") - + # Create sidebar navigation using streamlit-option-menu st.sidebar.markdown("---") - + try: from streamlit_option_menu import option_menu - + # Ensure it's in sidebar context with st.sidebar: nav_page = option_menu( menu_title=None, - options=["Upload Report", "Report Analyst", "All Results", "Settings"], + options=[ + "Upload Report", + "Report Analyst", + "All Results", + "Settings", + ], icons=["house", "file-text", "bar-chart", "gear"], menu_icon=None, default_index=0, orientation="vertical", key="nav_page", styles={ - "container": {"padding": "0", "background-color": "transparent"}, + "container": { + "padding": "0", + "background-color": "transparent", + }, "icon": {"color": "#7872A7", "font-size": "20px"}, "nav-link": { "font-family": "'Afacad', sans-serif", @@ -2846,46 +2715,44 @@ def main(): "color": "#4313C8", "font-weight": "700", }, - } + }, ) except ImportError: # Fallback to regular radio if package not installed nav_options = ["Upload Report", "Report Analyst", "All Results", "Settings"] - nav_page = st.sidebar.radio( - "", - nav_options, - key="nav_page", - label_visibility="collapsed" - ) - + nav_page = st.sidebar.radio("", nav_options, key="nav_page", label_visibility="collapsed") + # Show page-specific content based on navigation if nav_page == "Settings": st.title("Settings") st.caption("Configure application settings and integrations") - + # Open Source Modules Section st.header("Open Source Modules") st.caption("Core features available in the open source edition") - + # API Keys Configuration st.subheader("API Keys") st.caption("Enter your API keys to enable LLM features. Keys are stored in session state only and not persisted.") - + # Check if keys exist in environment but not in session state env_openai_key = os.getenv("OPENAI_API_KEY") env_google_key = os.getenv("GOOGLE_API_KEY") session_openai_key = st.session_state.get("api_key_openai_api_key") session_google_key = st.session_state.get("api_key_google_api_key") - + has_env_openai = env_openai_key and not session_openai_key has_env_google = env_google_key and not session_google_key - + # Get current values (from session state or environment) current_openai_key = APIKeyManager.get_api_key("OPENAI_API_KEY", st.session_state) current_google_key = APIKeyManager.get_api_key("GOOGLE_API_KEY", st.session_state) - + # OpenAI API Key section - with st.expander("OpenAI API Key", expanded=not has_env_openai or st.session_state.get("override_openai_key", False)): + with st.expander( + "OpenAI API Key", + expanded=not has_env_openai or st.session_state.get("override_openai_key", False), + ): # Show status for OpenAI key if has_env_openai and not st.session_state.get("override_openai_key", False): st.info("API key is set from environment variable") @@ -2893,16 +2760,18 @@ def main(): st.session_state.override_openai_key = True st.rerun() elif current_openai_key: - masked_openai = f"{current_openai_key[:8]}...{current_openai_key[-4:]}" if len(current_openai_key) > 12 else "***" + masked_openai = ( + f"{current_openai_key[:8]}...{current_openai_key[-4:]}" if len(current_openai_key) > 12 else "***" + ) st.caption(f"Current key: `{masked_openai}`") - + # Track override state override_openai = st.session_state.get("override_openai_key", False) - + if not has_env_openai or override_openai: # Track previous values to detect changes prev_openai_key = st.session_state.get("prev_openai_key", current_openai_key) - + # OpenAI API Key input openai_key_input = st.text_input( "OpenAI API Key", @@ -2910,9 +2779,9 @@ def main(): type="password", key="openai_api_key_input", help="Enter your OpenAI API key to use GPT models. Leave empty to use existing key from environment.", - placeholder="sk-..." if not current_openai_key else "Enter new key to update" + placeholder=("sk-..." if not current_openai_key else "Enter new key to update"), ) - + # Update API key if user entered a new value (different from current) if openai_key_input and openai_key_input != current_openai_key: APIKeyManager.set_api_key("OPENAI_API_KEY", openai_key_input, st.session_state) @@ -2924,15 +2793,18 @@ def main(): st.session_state.prev_openai_key = current_openai_key else: st.session_state.prev_openai_key = current_openai_key - + # Cancel override button if override_openai: if st.button("Cancel Override", key="cancel_override_openai"): st.session_state.override_openai_key = False st.rerun() - + # Google/Gemini API Key section - with st.expander("Google/Gemini API Key", expanded=not has_env_google or st.session_state.get("override_google_key", False)): + with st.expander( + "Google/Gemini API Key", + expanded=not has_env_google or st.session_state.get("override_google_key", False), + ): # Show status for Google key if has_env_google and not st.session_state.get("override_google_key", False): st.info("API key is set from environment variable") @@ -2940,16 +2812,18 @@ def main(): st.session_state.override_google_key = True st.rerun() elif current_google_key: - masked_google = f"{current_google_key[:8]}...{current_google_key[-4:]}" if len(current_google_key) > 12 else "***" + masked_google = ( + f"{current_google_key[:8]}...{current_google_key[-4:]}" if len(current_google_key) > 12 else "***" + ) st.caption(f"Current key: `{masked_google}`") - + # Track override state override_google = st.session_state.get("override_google_key", False) - + if not has_env_google or override_google: # Track previous values to detect changes prev_google_key = st.session_state.get("prev_google_key", current_google_key) - + # Google/Gemini API Key input google_key_input = st.text_input( "Google/Gemini API Key", @@ -2957,9 +2831,9 @@ def main(): type="password", key="google_api_key_input", help="Enter your Google API key to use Gemini models. Leave empty to use existing key from environment.", - placeholder="Enter your Google API key" if not current_google_key else "Enter new key to update" + placeholder=("Enter your Google API key" if not current_google_key else "Enter new key to update"), ) - + # Update API key if user entered a new value (different from current) if google_key_input and google_key_input != current_google_key: APIKeyManager.set_api_key("GOOGLE_API_KEY", google_key_input, st.session_state) @@ -2971,30 +2845,30 @@ def main(): st.session_state.prev_google_key = current_google_key else: st.session_state.prev_google_key = current_google_key - + # Cancel override button if override_google: if st.button("Cancel Override", key="cancel_override_google"): st.session_state.override_google_key = False st.rerun() - + # Show clear button if key exists (only for session state keys, not env) if current_google_key and not has_env_google: if st.button("Clear Google Key", key="clear_google_key"): APIKeyManager.set_api_key("GOOGLE_API_KEY", None, st.session_state) st.rerun() - + # Show clear button for OpenAI if key exists (only for session state keys, not env) if current_openai_key and not has_env_openai: if st.button("Clear OpenAI Key", key="clear_openai_key"): APIKeyManager.set_api_key("OPENAI_API_KEY", None, st.session_state) st.rerun() - + st.divider() - + # Database Configuration (read-only, from environment variables) st.subheader("Database Configuration") - + # Get database URL from environment or default database_url = os.getenv("DATABASE_URL") if database_url is None: @@ -3003,7 +2877,9 @@ def main(): db_path = str(Path(storage_path) / "cache" / "analysis.db") database_url = f"sqlite:///{db_path}" database_type = "SQLite" - st.info(f"**Type:** {database_type}\n\n**Path:** `{db_path}`\n\n*Configure via `STORAGE_PATH` environment variable*") + st.info( + f"**Type:** {database_type}\n\n**Path:** `{db_path}`\n\n*Configure via `STORAGE_PATH` environment variable*" + ) else: # Parse PostgreSQL URL to show connection details (masked) database_type = "PostgreSQL" @@ -3022,7 +2898,7 @@ def main(): masked_url = database_url else: masked_url = database_url - + # Extract connection details for display if "postgresql://" in database_url or "postgres://" in database_url: url_part = database_url.split("://")[-1] @@ -3044,32 +2920,41 @@ def main(): host = host_db port = "5432" db = "?" - - st.info(f"**Type:** {database_type}\n\n**Host:** `{host}`\n**Port:** `{port}`\n**Database:** `{db}`\n**User:** `{user}`\n\n*Configure via `DATABASE_URL` environment variable*") + + st.info( + f"**Type:** {database_type}\n\n**Host:** `{host}`\n**Port:** `{port}`\n**Database:** `{db}`\n**User:** `{user}`\n\n*Configure via `DATABASE_URL` environment variable*" + ) else: - st.info(f"**Type:** {database_type}\n\n**Connection:** `{masked_url}`\n\n*Configure via `DATABASE_URL` environment variable*") + st.info( + f"**Type:** {database_type}\n\n**Connection:** `{masked_url}`\n\n*Configure via `DATABASE_URL` environment variable*" + ) else: - st.info(f"**Type:** {database_type}\n\n**Connection:** `{masked_url}`\n\n*Configure via `DATABASE_URL` environment variable*") + st.info( + f"**Type:** {database_type}\n\n**Connection:** `{masked_url}`\n\n*Configure via `DATABASE_URL` environment variable*" + ) except Exception: - st.info(f"**Type:** {database_type}\n\n**Connection:** `{masked_url}`\n\n*Configure via `DATABASE_URL` environment variable*") - + st.info( + f"**Type:** {database_type}\n\n**Connection:** `{masked_url}`\n\n*Configure via `DATABASE_URL` environment variable*" + ) + # Store in session state for use by DocumentAnalyzer st.session_state.database_url = database_url - + st.divider() st.divider() - + # Enterprise Modules Section st.header("Enterprise Modules") st.caption("Features available in the enterprise edition") - + # Enterprise Integration (S3+NATS) st.subheader("Enterprise Integration") # In Streamlit, when a widget has a 'key', it automatically syncs with session state # The widget's return value is the current value from session state (or default if not set) # IMPORTANT: Don't provide 'value' parameter when using 'key' - let Streamlit manage it # The widget return value is the source of truth for the current render - st.markdown(""" + st.markdown( + """ - """, unsafe_allow_html=True) + """, + unsafe_allow_html=True, + ) use_s3_upload = st.checkbox( "Enable S3+NATS Upload", key="use_s3_upload", help="Upload documents via S3 and process via NATS for enterprise integration", ) - + # Show Enterprise Mode status only if checkbox is checked AND backend is available # Check AFTER widget render - use widget return value which reflects current state # The widget return value is the authoritative source for the current render cycle if use_s3_upload and BACKEND_INTEGRATION_AVAILABLE: st.info("Enterprise mode enabled") - + st.divider() - + # Backend Integration (Enterprise feature) if BACKEND_INTEGRATION_AVAILABLE: config = configure_backend_integration() @@ -3102,21 +2989,23 @@ def main(): st.warning("Backend integration modules not available") config = None st.session_state.backend_config = None - + st.divider() - + # File Storage Configuration (Enterprise feature) st.subheader("File Storage") st.caption("Configure where uploaded files are stored (Enterprise feature)") - + # Get database URL from session state (set above in Database Configuration) database_url_enterprise = st.session_state.get("database_url") - is_postgres_enterprise = database_url_enterprise and database_url_enterprise.startswith(("postgresql://", "postgres://")) - + is_postgres_enterprise = database_url_enterprise and database_url_enterprise.startswith( + ("postgresql://", "postgres://") + ) + # Initialize use_postgres_file_storage from session state or env if "use_postgres_file_storage" not in st.session_state: st.session_state.use_postgres_file_storage = os.getenv("USE_POSTGRES_FILE_STORAGE", "false").lower() == "true" - + if is_postgres_enterprise: use_postgres_storage = st.checkbox( "Store files in PostgreSQL", @@ -3124,7 +3013,7 @@ def main(): key="use_postgres_file_storage", help="Store uploaded files in PostgreSQL database (useful for Heroku deployments). Files are stored as BYTEA/BLOB. This is an enterprise feature.", ) - + if use_postgres_storage: st.info("📦 Files will be stored in PostgreSQL database") else: @@ -3141,7 +3030,7 @@ def main(): # Get file list for dropdown (including backend resources if enabled) backend_config = st.session_state.get("backend_config") previous_files = get_uploaded_files_history(backend_config=backend_config) - + # Determine selected file for display - check session state first selected_file_for_display = None if previous_files: @@ -3156,11 +3045,11 @@ def main(): if f["name"] == prev_file or f.get("path") == prev_file: selected_file_for_display = f break - + # If no file selected yet, use first one if not selected_file_for_display and previous_files: selected_file_for_display = previous_files[0] - + # Green PDF Display Container with integrated file selector if previous_files: # Use container with unique key for green panel styling @@ -3168,14 +3057,17 @@ def main(): with st.container(key="file-display-panel"): # Use columns for layout icon_col, content_col = st.columns([0.1, 0.9]) - + with icon_col: - st.markdown(""" + st.markdown( + """
description
- """, unsafe_allow_html=True) - + """, + unsafe_allow_html=True, + ) + with content_col: # File selector inside the green container selected_file_dropdown = st.selectbox( @@ -3183,39 +3075,52 @@ def main(): options=previous_files, format_func=lambda x: x["name"], key="previous_file", - label_visibility="collapsed" + label_visibility="collapsed", ) - + # Upload date below selector if selected_file_dropdown: selected_uri = selected_file_dropdown.get("uri", selected_file_dropdown.get("path", "")) is_backend = selected_uri.startswith("urn:report-analyst:backend:") - + if is_backend: # For backend resources, show backend info - from report_analyst.core.report_data_client import ReportResource - resource = ReportResource(name=selected_file_dropdown["name"], uri=selected_uri) + from report_analyst.core.report_data_client import ( + ReportResource, + ) + + resource = ReportResource( + name=selected_file_dropdown["name"], + uri=selected_uri, + ) parsed = resource.parse_backend_urn() if parsed: - st.markdown(f'Backend: {parsed["host"]}', unsafe_allow_html=True) + st.markdown( + f'Backend: {parsed["host"]}', + unsafe_allow_html=True, + ) else: # For local files, show upload date file_path_str = selected_file_dropdown.get("path", "") # Handle file:// URI format if file_path_str.startswith("file://"): file_path_str = file_path_str.replace("file://", "") - + file_path_display = Path(file_path_str) if file_path_display.exists(): import datetime + mod_time = file_path_display.stat().st_mtime upload_date = datetime.datetime.fromtimestamp(mod_time).strftime("%d.%m.%Y") - st.markdown(f'Uploaded, {upload_date}', unsafe_allow_html=True) + st.markdown( + f'Uploaded, {upload_date}', + unsafe_allow_html=True, + ) # Analysis Configuration section - 3 column layout with st.expander("Analysis Configuration", expanded=True): col1, col2, col3 = st.columns([1, 1, 1]) - + # Left column: Question Set with col1: selected_set = st.selectbox( @@ -3226,7 +3131,7 @@ def main(): index=0, # Ensure a default is selected on_change=update_analyzer_parameters, ) - + # Show question set description below if selected_set in question_sets: st.caption(question_sets[selected_set]["description"]) @@ -3234,7 +3139,8 @@ def main(): # Middle column: Processing Steps with col2: # Processing Steps heading with help icon and tooltip - st.markdown(""" + st.markdown( + """ - """, unsafe_allow_html=True) - + """, + unsafe_allow_html=True, + ) + # Create select slider for step selection selected_step = st.select_slider( "Select processing steps", @@ -3402,13 +3307,15 @@ def main(): value=st.session_state.processing_steps_slider, key="processing_steps_slider", help="Select how many processing steps to execute", - label_visibility="collapsed" + label_visibility="collapsed", ) - + # Display all step labels below the slider step_cols = st.columns(4) - selected_index = step_options.index(selected_step) if selected_step in step_options else len(step_options) - 1 - + selected_index = ( + step_options.index(selected_step) if selected_step in step_options else len(step_options) - 1 + ) + # Map step names to status keys step_status_map = { "Chunk": step_status.get("chunks", False), @@ -3416,7 +3323,7 @@ def main(): "Map": step_status.get("scoring", False), "Answer": step_status.get("analysis", False), } - + # Check if there's any stored data for this file # Reuse the file_path_for_status from the step_status check above has_stored_data = False @@ -3428,12 +3335,11 @@ def main(): # Try to find the selected file selected_file_obj = None prev_file = st.session_state.previous_file - + # Handle both dict and string formats if isinstance(prev_file, dict): for f in previous_files: - if (f["name"] == prev_file.get("name") or - f.get("path") == prev_file.get("path")): + if f["name"] == prev_file.get("name") or f.get("path") == prev_file.get("path"): selected_file_obj = f break else: @@ -3441,7 +3347,7 @@ def main(): if f["name"] == prev_file or f.get("path") == prev_file: selected_file_obj = f break - + if selected_file_obj: selected_uri = selected_file_obj.get("uri", selected_file_obj.get("path", "")) is_backend = selected_uri.startswith("urn:report-analyst:backend:") @@ -3449,7 +3355,7 @@ def main(): file_to_check = None # Backend resources don't have local file paths else: file_to_check = Path(selected_file_obj["path"]) - + # If we have a file path, check for stored data if file_to_check and file_to_check.exists(): cache_entries = st.session_state.analyzer.analyzer.cache_manager.check_cache_status( @@ -3460,42 +3366,47 @@ def main(): except Exception as e: logger.debug(f"Error checking stored data: {str(e)}") has_stored_data = False - + for idx, step_short in enumerate(step_options): with step_cols[idx]: step_full = step_full_names[step_short] is_selected = idx <= selected_index is_complete = step_status_map.get(step_short, False) - + # Show checkmark if complete, circle if incomplete indicator = "✓" if is_complete else "○" - + # Visual styling for selected steps if is_selected: - highlight_style = "background-color: rgba(192, 196, 250, 0.1); border: 1px solid #4313C8; color: #4313C8;" + highlight_style = ( + "background-color: rgba(192, 196, 250, 0.1); border: 1px solid #4313C8; color: #4313C8;" + ) else: highlight_style = "background-color: rgba(192, 196, 250, 0.05); border: 1px solid rgba(67, 19, 200, 0.3); color: #718096;" - + # Add status badge next to Chunking step - always show status_badge = "" if step_short == "Chunk": badge_text = "Stored" if has_stored_data else "New" badge_bg = "rgba(192, 196, 250, 0.3)" if has_stored_data else "rgba(192, 196, 250, 0.15)" - status_badge = f'{badge_text}' - - st.markdown(f""" + status_badge = f"{badge_text}" + + st.markdown( + f"""
{indicator} {step_full}{status_badge}
- """, unsafe_allow_html=True) + """, + unsafe_allow_html=True, + ) # Right column: Advanced Parameters with col3: st.markdown("**Advanced Parameters**") - + # Use 2 columns for Advanced Parameters to make it more compact adv_col1, adv_col2 = st.columns(2) - + with adv_col1: new_top_k = st.number_input( "Top K", @@ -3505,7 +3416,7 @@ def main(): key="new_top_k", on_change=update_analyzer_parameters, ) - + new_chunk_size = st.number_input( "Chunk Size", min_value=100, @@ -3514,7 +3425,7 @@ def main(): key="new_chunk_size", on_change=update_analyzer_parameters, ) - + new_overlap = st.number_input( "Overlap", min_value=0, @@ -3523,7 +3434,7 @@ def main(): key="new_overlap", on_change=update_analyzer_parameters, ) - + with adv_col2: new_llm_model = st.selectbox( "LLM Model", @@ -3532,14 +3443,14 @@ def main(): key="new_llm_model", on_change=update_analyzer_parameters, ) - + new_llm_scoring = st.checkbox( "LLM Scoring", value=False, key="new_llm_scoring", on_change=update_analyzer_parameters, ) - + new_batch_scoring = st.checkbox( "Batch Scoring", value=True, @@ -3552,10 +3463,7 @@ def main(): analyzer.analyzer.update_question_set(selected_set) # Clear results if question set changed - if ( - "last_question_set" not in st.session_state - or st.session_state.last_question_set != selected_set - ): + if "last_question_set" not in st.session_state or st.session_state.last_question_set != selected_set: if "results" in st.session_state: del st.session_state.results st.session_state.last_question_set = selected_set @@ -3565,14 +3473,20 @@ def main(): # Check if this is a backend resource (URN) or local file selected_uri = selected_file_dropdown.get("uri", selected_file_dropdown.get("path", "")) is_backend_resource = selected_uri.startswith("urn:report-analyst:backend:") - + if is_backend_resource: # Handle backend resource - from report_analyst.core.report_data_client import get_chunks_for_backend_resource - + from report_analyst.core.report_data_client import ( + get_chunks_for_backend_resource, + ) + backend_config = st.session_state.get("backend_config") - backend_configs = [backend_config] if backend_config and hasattr(backend_config, "use_backend") and backend_config.use_backend else [] - + backend_configs = ( + [backend_config] + if backend_config and hasattr(backend_config, "use_backend") and backend_config.use_backend + else [] + ) + # Get chunks from backend chunks = get_chunks_for_backend_resource(selected_uri, backend_configs) if chunks: @@ -3589,11 +3503,11 @@ def main(): # Handle local file - maintain backwards compatibility # Use absolute path string as before (existing behavior for SQLite cache) file_path_str = selected_file_dropdown.get("path", "") - + # Handle file:// URI format - extract actual path if file_path_str.startswith("file://"): file_path_str = file_path_str.replace("file://", "") - + if file_path_str: file_path_obj = Path(file_path_str) if file_path_obj.exists(): @@ -3610,30 +3524,24 @@ def main(): # No path found in selected file file_path = None logger.warning(f"No path found in selected file: {selected_file_dropdown}") - + # Continue with analysis if we have a valid file path or chunks if (not is_backend_resource and file_path and Path(file_path).exists()) or (is_backend_resource and chunks): # Load questions and handle selection - question_set_data = analyzer.load_question_set( - st.session_state.new_question_set - ) + question_set_data = analyzer.load_question_set(st.session_state.new_question_set) questions = question_set_data["questions"] st.markdown("
", unsafe_allow_html=True) # Add question selection UI - styled table format st.subheader("Select Questions") - + table_key = f"questions_table_{st.session_state.new_question_set}" select_all_key = f"select_all_{st.session_state.new_question_set}" - + # Select All button - styled purple, smaller, below heading - select_all_clicked = st.button( - "Select All", - key=select_all_key, - type="primary" - ) - + select_all_clicked = st.button("Select All", key=select_all_key, type="primary") + # Handle select all button click - toggle state if select_all_clicked: # Toggle the select all state @@ -3641,21 +3549,23 @@ def main(): if toggle_key not in st.session_state: st.session_state[toggle_key] = False st.session_state[toggle_key] = not st.session_state[toggle_key] - + # Get current select all state toggle_key = f"select_all_state_{st.session_state.new_question_set}" select_all = st.session_state.get(toggle_key, False) - + # If select all is active, update all checkboxes if select_all: # Set all questions to selected questions_data = [] for q_id, q_data in questions.items(): - questions_data.append({ - "Select": True, - "QID": q_id, - "QUESTION": q_data['text'] - }) + questions_data.append( + { + "Select": True, + "QID": q_id, + "QUESTION": q_data["text"], + } + ) questions_df = pd.DataFrame(questions_data) else: # Build dataframe - let widget manage its own state, don't sync until analyze is clicked @@ -3663,7 +3573,11 @@ def main(): if table_key in st.session_state: widget_df = st.session_state[table_key] # Check if widget_df is a DataFrame with the expected columns - if isinstance(widget_df, pd.DataFrame) and "QID" in widget_df.columns and "Select" in widget_df.columns: + if ( + isinstance(widget_df, pd.DataFrame) + and "QID" in widget_df.columns + and "Select" in widget_df.columns + ): # Widget has valid state - use it directly questions_data = [] for q_id, q_data in questions.items(): @@ -3672,76 +3586,76 @@ def main(): else: # Question not in widget state yet, default to False is_selected = False - - questions_data.append({ - "Select": is_selected, - "QID": q_id, - "QUESTION": q_data['text'] - }) + + questions_data.append( + { + "Select": is_selected, + "QID": q_id, + "QUESTION": q_data["text"], + } + ) questions_df = pd.DataFrame(questions_data) else: # Widget state exists but has wrong structure - rebuild from scratch questions_data = [] for q_id, q_data in questions.items(): - questions_data.append({ - "Select": False, - "QID": q_id, - "QUESTION": q_data['text'] - }) + questions_data.append( + { + "Select": False, + "QID": q_id, + "QUESTION": q_data["text"], + } + ) questions_df = pd.DataFrame(questions_data) else: # First time - build from scratch (don't use session state to avoid sync issues) questions_data = [] for q_id, q_data in questions.items(): - questions_data.append({ - "Select": False, - "QID": q_id, - "QUESTION": q_data['text'] - }) + questions_data.append( + { + "Select": False, + "QID": q_id, + "QUESTION": q_data["text"], + } + ) questions_df = pd.DataFrame(questions_data) - + # Display as editable table - widget manages its own state edited_df = st.data_editor( questions_df, column_config={ "Select": st.column_config.CheckboxColumn("Select", width=70), "QID": st.column_config.TextColumn("QID", disabled=True, width=120), - "QUESTION": st.column_config.TextColumn("Question", disabled=True) + "QUESTION": st.column_config.TextColumn("Question", disabled=True), }, hide_index=True, use_container_width=True, key=table_key, num_rows="fixed", - column_order=["Select", "QID", "QUESTION"] + column_order=["Select", "QID", "QUESTION"], ) - + # Don't sync to session state here - only sync when analyze button is clicked # Analysis button and results col1, col2 = st.columns([2, 1]) with col1: - analyze_clicked = st.button( - "Analyze Selected Questions", key="analyze_button" - ) + analyze_clicked = st.button("Analyze Selected Questions", key="analyze_button") with col2: - reanalyze_clicked = st.button( - "Reanalyze", key="reanalyze_button" - ) + reanalyze_clicked = st.button("Reanalyze", key="reanalyze_button") if analyze_clicked or reanalyze_clicked: # NOW sync the selection state from the widget # Get selected questions from the edited dataframe selected_questions = edited_df[edited_df["Select"] == True]["QID"].tolist() - + # Update session state for individual question checkboxes (for backward compatibility) for q_id in questions.keys(): is_selected = q_id in selected_questions st.session_state[f"individual_question_{q_id}"] = is_selected - + if not selected_questions: - st.warning( - "Please select at least one question to analyze." - ) + st.warning("Please select at least one question to analyze.") else: try: # Set force_recompute based on which button was clicked @@ -3762,7 +3676,7 @@ def main(): # Check if this is a backend resource is_backend = st.session_state.get("backend_chunks") is not None backend_uri = st.session_state.get("backend_resource_uri") - + # Use URN as file_path for backend resources, absolute path string for local files # This maintains backwards compatibility with SQLite cache if is_backend and backend_uri: @@ -3770,12 +3684,10 @@ def main(): else: # Local file - ensure it's absolute path string (backwards compatible) analysis_file_path = str(Path(file_path).resolve()) if file_path else file_path - + if reanalyze_clicked: # For reanalysis, skip cache check and analyze all selected questions - progress_text.info( - f"Reanalyzing {len(selected_questions)} questions..." - ) + progress_text.info(f"Reanalyzing {len(selected_questions)} questions...") asyncio.run( run_analysis( analyzer, @@ -3798,30 +3710,22 @@ def main(): question_id, result, ) in cached_results.items(): - st.session_state.results["answers"][ - question_id - ] = result + st.session_state.results["answers"][question_id] = result # Generate file key for display - file_key = generate_file_key( - analysis_file_path, st - ) + file_key = generate_file_key(analysis_file_path, st) # Update display - analysis_df, chunks_df = ( - create_analysis_dataframes( - st.session_state.results["answers"], - file_key, - ) + analysis_df, chunks_df = create_analysis_dataframes( + st.session_state.results["answers"], + file_key, ) st.session_state.analysis_df = analysis_df st.session_state.chunks_df = chunks_df st.session_state.analysis_complete = True else: # Run analysis for uncached questions - progress_text.info( - f"Processing {len(selected_questions)} questions..." - ) + progress_text.info(f"Processing {len(selected_questions)} questions...") try: # Run analysis for uncached questions @@ -3836,14 +3740,10 @@ def main(): ) ) - progress_text.success( - "Analysis complete!" - ) + progress_text.success("Analysis complete!") except Exception as e: - st.error( - f"Error during analysis: {str(e)}" - ) + st.error(f"Error during analysis: {str(e)}") st.exception(e) # Get final results @@ -3855,20 +3755,12 @@ def main(): # Process all results into dataframes if all_results: - analysis_df, chunks_df = ( - create_analysis_dataframes(all_results) - ) + analysis_df, chunks_df = create_analysis_dataframes(all_results) file_key = Path(file_path).stem - display_analysis_results( - analysis_df, chunks_df, file_key - ) - progress_text.success( - f"✓ Analysis complete for {len(selected_questions)} questions" - ) + display_analysis_results(analysis_df, chunks_df, file_key) + progress_text.success(f"✓ Analysis complete for {len(selected_questions)} questions") else: - progress_text.error( - "No results found after analysis" - ) + progress_text.error("No results found after analysis") except Exception as e: logger.error( @@ -3909,12 +3801,11 @@ def main(): deployment_type="enterprise", experiment_name="S3+NATS Upload", ) - st.session_state.flow_orchestrator = create_flow_orchestrator( - st.session_state.backend_config - ) + st.session_state.flow_orchestrator = create_flow_orchestrator(st.session_state.backend_config) # Unified upload page styling (shows for both enterprise and regular mode) - st.markdown(""" + st.markdown( + """ - """, unsafe_allow_html=True) - + """, + unsafe_allow_html=True, + ) + # Try to import JSON Schema form component (enterprise feature) try: # Use the proper Streamlit custom component - from report_analyst_enterprise.components.streamlit_component.backend import json_schema_form import json + + from report_analyst_enterprise.components.streamlit_component.backend import ( + json_schema_form, + ) + # Path is already imported at the top of the file - + JSON_SCHEMA_FORM_AVAILABLE = True - + # Load PDF upload schema - schema_path = Path(__file__).parent.parent / "report_analyst_enterprise" / "components" / "schemas" / "pdf_upload_schema.json" - ui_schema_path = Path(__file__).parent.parent / "report_analyst_enterprise" / "components" / "schemas" / "pdf_upload_ui_schema.json" - + schema_path = ( + Path(__file__).parent.parent + / "report_analyst_enterprise" + / "components" + / "schemas" + / "pdf_upload_schema.json" + ) + ui_schema_path = ( + Path(__file__).parent.parent + / "report_analyst_enterprise" + / "components" + / "schemas" + / "pdf_upload_ui_schema.json" + ) + if schema_path.exists() and ui_schema_path.exists(): with open(schema_path) as f: pdf_upload_schema = json.load(f) @@ -4020,42 +3932,54 @@ def main(): JSON_SCHEMA_FORM_AVAILABLE = False pdf_upload_schema = None pdf_upload_ui_schema = None - + # File upload with optional metadata form uploaded_file = st.file_uploader( - "Choose a PDF file", - type="pdf", + "Choose a PDF file", + type="pdf", key="file_uploader", - help="Limit 200MB per file • PDF" + help="Limit 200MB per file • PDF", ) - + # Show metadata form if JSON Schema form is available pdf_metadata = None company_metadata = None - + if JSON_SCHEMA_FORM_AVAILABLE: # ESRS Company Information Form - esrs_schema_path = Path(__file__).parent.parent / "report_analyst_enterprise" / "components" / "schemas" / "esrs_company_schema.json" - esrs_ui_schema_path = Path(__file__).parent.parent / "report_analyst_enterprise" / "components" / "schemas" / "esrs_company_ui_schema.json" - + esrs_schema_path = ( + Path(__file__).parent.parent + / "report_analyst_enterprise" + / "components" + / "schemas" + / "esrs_company_schema.json" + ) + esrs_ui_schema_path = ( + Path(__file__).parent.parent + / "report_analyst_enterprise" + / "components" + / "schemas" + / "esrs_company_ui_schema.json" + ) + if esrs_schema_path.exists() and esrs_ui_schema_path.exists(): with open(esrs_schema_path) as f: esrs_company_schema = json.load(f) with open(esrs_ui_schema_path) as f: esrs_company_ui_schema = json.load(f) - + with st.expander("ESRS Company Information", expanded=True): st.caption("Enter company data aligned with ESRS XBRL taxonomy requirements") company_metadata = json_schema_form( schema=esrs_company_schema, ui_schema=esrs_company_ui_schema, key="esrs_company_form", - height=700 + height=700, ) if company_metadata and company_metadata.get("type") == "submit": st.success("Company information saved!") st.session_state.esrs_company_metadata = company_metadata.get("formData", company_metadata) - + # Basic PDF metadata form if pdf_upload_schema: with st.expander("Add Document Metadata (Optional)", expanded=False): @@ -4064,55 +3988,37 @@ def main(): schema=pdf_upload_schema, ui_schema=pdf_upload_ui_schema, key="pdf_metadata_form", - height=500 + height=500, ) if pdf_metadata: st.success("Metadata saved!") # Store in session state for use after upload st.session_state.pdf_metadata = pdf_metadata - + if uploaded_file: # Handle upload based on mode if use_s3_upload and BACKEND_INTEGRATION_AVAILABLE: # Use S3+NATS enterprise flow - with st.spinner( - "Uploading to S3 and triggering backend processing..." - ): + with st.spinner("Uploading to S3 and triggering backend processing..."): try: # Debug: Check if backend_service exists orchestrator = st.session_state.get("flow_orchestrator") - backend_service = ( - getattr(orchestrator, "backend_service", None) - if orchestrator - else None - ) + backend_service = getattr(orchestrator, "backend_service", None) if orchestrator else None - logger.info( - f"[ENTERPRISE] Debug - orchestrator: {orchestrator is not None}" - ) - logger.info( - f"[ENTERPRISE] Debug - backend_service: {backend_service is not None}" - ) + logger.info(f"[ENTERPRISE] Debug - orchestrator: {orchestrator is not None}") + logger.info(f"[ENTERPRISE] Debug - backend_service: {backend_service is not None}") if not backend_service: - raise Exception( - "Backend service not initialized properly" - ) + raise Exception("Backend service not initialized properly") # Get file bytes file_bytes = uploaded_file.getbuffer() # Process via S3+NATS flow - result = asyncio.run( - backend_service.upload_pdf( - file_bytes, uploaded_file.name - ) - ) + result = asyncio.run(backend_service.upload_pdf(file_bytes, uploaded_file.name)) if result: - st.success( - f"File uploaded via S3+NATS: {uploaded_file.name}" - ) + st.success(f"File uploaded via S3+NATS: {uploaded_file.name}") st.info(f"Document ID: {result}") st.session_state.current_file = result st.session_state.uploaded_file = uploaded_file @@ -4132,13 +4038,11 @@ def main(): st.info("Falling back to local processing...") # Fall through to local processing use_s3_upload = False - + # Local processing (when enterprise mode is off or failed) if not use_s3_upload or not BACKEND_INTEGRATION_AVAILABLE: file_path = save_uploaded_file(uploaded_file) - logger.info( - f"[UPLOAD] Saved uploaded file: {uploaded_file.name} at {file_path}" - ) + logger.info(f"[UPLOAD] Saved uploaded file: {uploaded_file.name} at {file_path}") if file_path and file_path != st.session_state.get("current_file"): st.session_state.current_file = file_path st.session_state.uploaded_file = uploaded_file @@ -4146,13 +4050,9 @@ def main(): st.session_state.analysis_triggered = False if "results" in st.session_state: del st.session_state.results - logger.info( - f"[UPLOAD] Added file to session state: {uploaded_file.name}" - ) + logger.info(f"[UPLOAD] Added file to session state: {uploaded_file.name}") st.success(f"File uploaded successfully: {uploaded_file.name}") - logger.info( - f"[UPLOAD] Displaying cache selector for file: {file_path}" - ) + logger.info(f"[UPLOAD] Displaying cache selector for file: {file_path}") # Removed display_cache_selector - stored data status now shown as pill next to Chunking step if not st.session_state.get("file_processed"): st.session_state.file_processed = True @@ -4182,11 +4082,13 @@ def main(): # Add Climate+Tech footer at the bottom of sidebar # Get current theme for logo selection and encode image as base64 try: - theme = st.context.theme if hasattr(st.context, 'theme') else {} - is_dark = theme.get('base', 'light') == 'dark' if theme else False - logo_filename = "assets/climate-and-tech-logo-dark-mode.png" if is_dark else "assets/climateandtech-logo-new-light-mode.png" + theme = st.context.theme if hasattr(st.context, "theme") else {} + is_dark = theme.get("base", "light") == "dark" if theme else False + logo_filename = ( + "assets/climate-and-tech-logo-dark-mode.png" if is_dark else "assets/climateandtech-logo-new-light-mode.png" + ) logo_path = Path(__file__).parent / logo_filename - + # Read and encode image as base64 if logo_path.exists(): with open(logo_path, "rb") as img_file: @@ -4198,7 +4100,7 @@ def main(): except Exception as e: logger.warning(f"Could not load logo: {str(e)}") logo_src = "" - + # Add footer to sidebar st.sidebar.markdown("---") footer = f""" diff --git a/report_analyst/streamlit_app_backend.py b/report_analyst/streamlit_app_backend.py index e828c468f..7bc3a3708 100644 --- a/report_analyst/streamlit_app_backend.py +++ b/report_analyst/streamlit_app_backend.py @@ -74,9 +74,7 @@ def main(): # Check if backend integration is available if not BACKEND_INTEGRATION_AVAILABLE: - st.warning( - "Backend integration modules not available. Running in fallback mode." - ) + st.warning("Backend integration modules not available. Running in fallback mode.") run_fallback_mode() return @@ -122,9 +120,7 @@ def handle_complete_backend_flow(orchestrator, uploaded_file, config: BackendCon question_set_options = question_loader.get_question_set_options() else: # Fallback: use a generic approach without hardcoded names - question_set_options = ( - [] - ) # No predefined options when core functionality unavailable + question_set_options = [] # No predefined options when core functionality unavailable question_set = st.selectbox( "Select Question Set", options=question_set_options, @@ -146,9 +142,7 @@ def handle_complete_backend_flow(orchestrator, uploaded_file, config: BackendCon display_backend_analysis_results(st.session_state.backend_analysis_result) -def handle_processing_and_analysis_flow( - orchestrator, uploaded_file, config: BackendConfig -): +def handle_processing_and_analysis_flow(orchestrator, uploaded_file, config: BackendConfig): """Handle flows that require processing then analysis""" # Step 1: Process document @@ -199,9 +193,7 @@ def configure_questions() -> List[str]: question_set_options = question_loader.get_question_set_options() + ["custom"] else: # Fallback: use a generic approach without hardcoded names - question_set_options = [ - "custom" - ] # Only custom when core functionality unavailable + question_set_options = ["custom"] # Only custom when core functionality unavailable question_set_name = st.selectbox( "Select Question Set", options=question_set_options, @@ -259,9 +251,7 @@ def display_backend_analysis_results(result: AnalysisResult): with col1: st.write(f"**Analysis Job ID:** {result.analysis_job_id}") - st.write( - f"**Stored in Backend:** {'Yes' if result.stored_in_backend else 'No'}" - ) + st.write(f"**Stored in Backend:** {'Yes' if result.stored_in_backend else 'No'}") with col2: st.write(f"**Analysis Method:** Complete Backend") @@ -369,18 +359,14 @@ def run_fallback_mode(): questions = get_question_set(question_set_name) if st.button("Run Analysis"): - st.success( - f"Loaded {len(questions)} questions from {question_set_name.upper()}" - ) + st.success(f"Loaded {len(questions)} questions from {question_set_name.upper()}") # Display questions (fallback analysis) st.subheader("Questions to Analyze") for i, question in enumerate(questions, 1): st.write(f"**{i}.** {question}") - st.info( - "This is a demo mode. Backend integration modules are not available." - ) + st.info("This is a demo mode. Backend integration modules are not available.") else: st.error("Core functionality not available. Please check your installation.") diff --git a/report_analyst_api/LICENSE b/report_analyst_api/LICENSE index ea3f9d727..559ec71d3 100644 --- a/report_analyst_api/LICENSE +++ b/report_analyst_api/LICENSE @@ -17,3 +17,4 @@ This software is part of the Open Sustainability Analysis project. + diff --git a/report_analyst_api/main.py b/report_analyst_api/main.py index 6b795f03f..7b30993e6 100644 --- a/report_analyst_api/main.py +++ b/report_analyst_api/main.py @@ -84,9 +84,7 @@ async def get_question_sets(): # Convert to API response format question_sets = [] for qset in question_sets_data.values(): - question_sets.append( - QuestionSet(id=qset.id, name=qset.name, description=qset.description) - ) + question_sets.append(QuestionSet(id=qset.id, name=qset.name, description=qset.description)) return question_sets except Exception as e: diff --git a/report_analyst_api/schemas.py b/report_analyst_api/schemas.py index e3c9a299f..c33905d64 100644 --- a/report_analyst_api/schemas.py +++ b/report_analyst_api/schemas.py @@ -24,14 +24,10 @@ class DocumentUpload(BaseModel): """Request model for document upload""" file_path: str = Field(..., description="Path to the document file") - source_type: Optional[str] = Field( - "local", description="Document source type (local or search_backend)" - ) + source_type: Optional[str] = Field("local", description="Document source type (local or search_backend)") class Config: - schema_extra = { - "example": {"file_path": "/path/to/document.pdf", "source_type": "local"} - } + schema_extra = {"example": {"file_path": "/path/to/document.pdf", "source_type": "local"}} class AnalysisConfiguration(BaseModel): @@ -41,13 +37,9 @@ class AnalysisConfiguration(BaseModel): chunk_size: int = Field(500, description="Text chunk size") chunk_overlap: int = Field(20, description="Chunk overlap") top_k: int = Field(5, description="Top K chunks for retrieval") - use_llm_scoring: bool = Field( - False, description="Use LLM for chunk relevance scoring" - ) + use_llm_scoring: bool = Field(False, description="Use LLM for chunk relevance scoring") single_call: bool = Field(True, description="Use single LLM call for analysis") - force_recompute: bool = Field( - False, description="Force recomputation of cached results" - ) + force_recompute: bool = Field(False, description="Force recomputation of cached results") class AnalysisRequest(BaseModel): @@ -55,12 +47,8 @@ class AnalysisRequest(BaseModel): document_id: str = Field(..., description="ID of the uploaded document") question_set_id: str = Field(..., description="ID of the question set to use") - selected_questions: List[str] = Field( - ..., description="List of question IDs to analyze" - ) - configuration: Optional[AnalysisConfiguration] = Field( - default_factory=AnalysisConfiguration - ) + selected_questions: List[str] = Field(..., description="List of question IDs to analyze") + configuration: Optional[AnalysisConfiguration] = Field(default_factory=AnalysisConfiguration) class Config: schema_extra = { @@ -96,9 +84,7 @@ class ChunkRelevance(BaseModel): chunk_id: str = Field(..., description="Chunk identifier") question_id: str = Field(..., description="Question identifier") relevance_score: float = Field(..., ge=0.0, le=1.0, description="Relevance score") - llm_explanation: Optional[str] = Field( - None, description="LLM explanation of relevance" - ) + llm_explanation: Optional[str] = Field(None, description="LLM explanation of relevance") class AnalysisResult(BaseModel): @@ -108,12 +94,8 @@ class AnalysisResult(BaseModel): question_id: str = Field(..., description="Question identifier") question_text: str = Field(..., description="Full question text") answer: str = Field(..., description="Analysis answer") - evidence_chunks: List[Dict[str, Any]] = Field( - ..., description="Supporting evidence chunks" - ) - confidence_score: Optional[float] = Field( - None, ge=0.0, le=1.0, description="Confidence in the answer" - ) + evidence_chunks: List[Dict[str, Any]] = Field(..., description="Supporting evidence chunks") + confidence_score: Optional[float] = Field(None, ge=0.0, le=1.0, description="Confidence in the answer") model_used: str = Field(..., description="LLM model used for analysis") processing_time: float = Field(..., description="Time taken in seconds") @@ -138,9 +120,7 @@ class QuestionSet(BaseModel): class QuestionSetResponse(BaseModel): """Response model for question sets""" - question_sets: Dict[str, Dict[str, Any]] = Field( - ..., description="Available question sets" - ) + question_sets: Dict[str, Dict[str, Any]] = Field(..., description="Available question sets") class DocumentChunkResponse(BaseModel): @@ -149,17 +129,13 @@ class DocumentChunkResponse(BaseModel): chunk_id: str = Field(..., description="Chunk identifier") chunk_text: str = Field(..., description="Chunk text content") chunk_metadata: Dict[str, Any] = Field(..., description="Chunk metadata") - relevance_scores: List[ChunkRelevance] = Field( - default_factory=list, description="Relevance scores" - ) + relevance_scores: List[ChunkRelevance] = Field(default_factory=list, description="Relevance scores") class IntegrationsResponse(BaseModel): """Response model for available integrations""" - available_integrations: Dict[str, bool] = Field( - ..., description="Available integration modules" - ) + available_integrations: Dict[str, bool] = Field(..., description="Available integration modules") document_sources: List[str] = Field(..., description="Available document sources") @@ -177,9 +153,7 @@ class HealthResponse(BaseModel): core_package: bool = Field(..., description="Core package availability") question_loader: bool = Field(..., description="Question loader status") analyzer: bool = Field(..., description="Analyzer status") - available_integrations: Dict[str, bool] = Field( - ..., description="Integration status" - ) + available_integrations: Dict[str, bool] = Field(..., description="Integration status") document_sources: List[str] = Field(..., description="Available document sources") @@ -201,9 +175,7 @@ class AnalysisResponse(BaseModel): question_set: str = Field(..., description="Question set used") results: List[Dict[str, Any]] = Field(..., description="Analysis results") configuration: Dict[str, Any] = Field(..., description="Analysis configuration") - timestamp: datetime = Field( - default_factory=datetime.now, description="Analysis timestamp" - ) + timestamp: datetime = Field(default_factory=datetime.now, description="Analysis timestamp") class AsyncJobResponse(BaseModel): @@ -211,9 +183,7 @@ class AsyncJobResponse(BaseModel): task_id: str = Field(..., description="Unique task identifier") status: str = Field(..., description="Job status") - timestamp: datetime = Field( - default_factory=datetime.now, description="Job creation timestamp" - ) + timestamp: datetime = Field(default_factory=datetime.now, description="Job creation timestamp") class JobStatus(BaseModel): @@ -223,9 +193,5 @@ class JobStatus(BaseModel): status: str = Field(..., description="Current job status") progress: float = Field(0.0, description="Job progress percentage") error: Optional[str] = Field(None, description="Error message if failed") - results: Optional[List[Dict[str, Any]]] = Field( - None, description="Analysis results if completed" - ) - timestamp: datetime = Field( - default_factory=datetime.now, description="Status timestamp" - ) + results: Optional[List[Dict[str, Any]]] = Field(None, description="Analysis results if completed") + timestamp: datetime = Field(default_factory=datetime.now, description="Status timestamp") diff --git a/report_analyst_enterprise/LICENSE b/report_analyst_enterprise/LICENSE index 9964d4b3f..296049dfa 100644 --- a/report_analyst_enterprise/LICENSE +++ b/report_analyst_enterprise/LICENSE @@ -15,3 +15,4 @@ Copyright (c) 2025 Climate+Tech This software is part of the Open Sustainability Analysis project. + diff --git a/report_analyst_enterprise/__init__.py b/report_analyst_enterprise/__init__.py index fd5649480..a58327c69 100644 --- a/report_analyst_enterprise/__init__.py +++ b/report_analyst_enterprise/__init__.py @@ -6,4 +6,3 @@ """ __version__ = "0.1.0" - diff --git a/report_analyst_enterprise/database/__init__.py b/report_analyst_enterprise/database/__init__.py index 5c1b5fbce..8d5408784 100644 --- a/report_analyst_enterprise/database/__init__.py +++ b/report_analyst_enterprise/database/__init__.py @@ -15,4 +15,3 @@ "get_vector_distance_func", "setup_pgvector_extension", ] - diff --git a/report_analyst_enterprise/database/pgvector_support.py b/report_analyst_enterprise/database/pgvector_support.py index 9a6521f48..2cbd3fea3 100644 --- a/report_analyst_enterprise/database/pgvector_support.py +++ b/report_analyst_enterprise/database/pgvector_support.py @@ -22,12 +22,15 @@ def check_pgvector_available(connection) -> bool: """ try: from sqlalchemy import text + result = connection.execute( - text(""" + text( + """ SELECT EXISTS( SELECT 1 FROM pg_extension WHERE extname = 'vector' ) - """) + """ + ) ) available = result.fetchone()[0] if available: @@ -52,6 +55,7 @@ def setup_pgvector_extension(connection) -> bool: """ try: from sqlalchemy import text + # Check if already exists if check_pgvector_available(connection): return True @@ -76,10 +80,11 @@ def create_vector_type(dimension: Optional[int] = None): Returns: SQLAlchemy TypeEngine for vector type """ - from sqlalchemy import TypeDecorator, Text + from sqlalchemy import Text, TypeDecorator class VectorType(TypeDecorator): """Custom type for pgvector vector column""" + impl = Text cache_ok = True @@ -103,6 +108,7 @@ def process_result_value(self, value, dialect): return None # Parse string format back to list import ast + return ast.literal_eval(value) return VectorType() @@ -128,7 +134,7 @@ def get_vector_distance_func(embedding_column_name: str, query_vector, distance_ # Map distance type to pgvector operator operators = { "cosine": "<=>", # Cosine distance - "l2": "<->", # L2 distance + "l2": "<->", # L2 distance "inner_product": "<#>", # Inner product (negative) } @@ -142,4 +148,3 @@ def get_vector_distance_func(embedding_column_name: str, query_vector, distance_ # Format: embedding <=> '[1,2,3]'::vector sql_expr = f"{embedding_column_name} {operator} :query_vector::vector" return sql_expr, {"query_vector": query_str} - diff --git a/report_analyst_jobs/LICENSE b/report_analyst_jobs/LICENSE index fd952fdb5..5eba841db 100644 --- a/report_analyst_jobs/LICENSE +++ b/report_analyst_jobs/LICENSE @@ -17,3 +17,4 @@ This software is part of the Open Sustainability Analysis project. + diff --git a/report_analyst_jobs/analysis_handler.py b/report_analyst_jobs/analysis_handler.py index 0cb8537bf..898506783 100644 --- a/report_analyst_jobs/analysis_handler.py +++ b/report_analyst_jobs/analysis_handler.py @@ -94,9 +94,7 @@ async def _get_document_chunks(self, job: AnalysisJobDefinition): source = source_class() # Get chunks - chunks = await source.get_chunks( - job.document_id, job.parameters.get("configuration", {}) - ) + chunks = await source.get_chunks(job.document_id, job.parameters.get("configuration", {})) logger.info(f"Retrieved {len(chunks)} chunks for document {job.document_id}") return chunks @@ -113,9 +111,7 @@ async def _get_question_set(self, job: AnalysisJobDefinition): if question_id in question_set.questions: filtered_questions[question_id] = question_set.questions[question_id] - logger.info( - f"Using {len(filtered_questions)} questions from set {job.question_set_id}" - ) + logger.info(f"Using {len(filtered_questions)} questions from set {job.question_set_id}") return filtered_questions async def _analyze_document(self, job: AnalysisJobDefinition, chunks, questions): @@ -130,14 +126,10 @@ async def _analyze_document(self, job: AnalysisJobDefinition, chunks, questions) logger.info(f"Analyzing question {question_id}") # Find relevant chunks for this question - relevant_chunks = await self._find_relevant_chunks( - chunks, question_data["text"] - ) + relevant_chunks = await self._find_relevant_chunks(chunks, question_data["text"]) # Perform LLM analysis - analysis_result = await self._analyze_question( - question_data, relevant_chunks, job - ) + analysis_result = await self._analyze_question(question_data, relevant_chunks, job) results.append( { @@ -195,12 +187,7 @@ async def _analyze_question(self, question_data, chunks, job): # Combine chunks into context context = "\n\n".join( - [ - f"Chunk {i+1}: {chunk['text']}" - for i, chunk in enumerate( - chunks[:10] - ) # Limit to avoid token limits - ] + [f"Chunk {i+1}: {chunk['text']}" for i, chunk in enumerate(chunks[:10])] # Limit to avoid token limits ) # Create prompt @@ -299,9 +286,7 @@ async def execute(self, job: JobDefinition) -> JobResult: except Exception as e: if self.progress_callback: - await self.progress_callback( - job.job_id, 0.0, f"Analysis failed: {str(e)}" - ) + await self.progress_callback(job.job_id, 0.0, f"Analysis failed: {str(e)}") return JobResult(job_id=job.job_id, status=JobStatus.FAILED, error=str(e)) @@ -320,12 +305,8 @@ async def _analyze_with_progress(self, job, chunks, questions): ) # Analyze question - relevant_chunks = await self._find_relevant_chunks( - chunks, question_data["text"] - ) - analysis_result = await self._analyze_question( - question_data, relevant_chunks, job - ) + relevant_chunks = await self._find_relevant_chunks(chunks, question_data["text"]) + analysis_result = await self._analyze_question(question_data, relevant_chunks, job) results.append( { diff --git a/report_analyst_jobs/coordinator.py b/report_analyst_jobs/coordinator.py index b61b67d04..253fec1b0 100644 --- a/report_analyst_jobs/coordinator.py +++ b/report_analyst_jobs/coordinator.py @@ -48,9 +48,7 @@ def _initialize_executors(self): nats_executor = NATSJobExecutor( nats_url=nats_config.get("url", "nats://localhost:4222"), stream_name=nats_config.get("stream", "JOBS"), - consumer_name=nats_config.get( - "consumer", "report-analyst-consumer" - ), + consumer_name=nats_config.get("consumer", "report-analyst-consumer"), ) self.executors[ExecutionBackend.NATS] = nats_executor self.default_backend = ExecutionBackend.NATS @@ -82,9 +80,7 @@ def register_handler(self, job_type: str, handler: JobHandler): if hasattr(executor, "register_handler"): executor.register_handler(job_type, handler) - async def submit_job( - self, job: JobDefinition, backend: Optional[ExecutionBackend] = None - ) -> str: + async def submit_job(self, job: JobDefinition, backend: Optional[ExecutionBackend] = None) -> str: """Submit a job for execution""" # Determine backend @@ -114,9 +110,7 @@ async def get_job_status(self, job_id: str) -> JobResult: logger.debug(f"Failed to get job status from {backend}: {e}") # Job not found in any backend - return JobResult( - job_id=job_id, status=JobStatus.FAILED, error="Job not found in any backend" - ) + return JobResult(job_id=job_id, status=JobStatus.FAILED, error="Job not found in any backend") async def cancel_job(self, job_id: str) -> bool: """Cancel a job in any backend""" diff --git a/report_analyst_jobs/data_lake_integration.py b/report_analyst_jobs/data_lake_integration.py index 21c2fe342..9be5740d1 100644 --- a/report_analyst_jobs/data_lake_integration.py +++ b/report_analyst_jobs/data_lake_integration.py @@ -113,9 +113,7 @@ def __post_init__(self): class DataLakeClient: """Client for interacting with the search backend data lake""" - def __init__( - self, backend_url: str = "http://localhost:8000", owner: str = "default" - ): + def __init__(self, backend_url: str = "http://localhost:8000", owner: str = "default"): self.backend_url = backend_url self.owner = owner @@ -128,17 +126,13 @@ async def deploy_configuration(self, config: DeploymentConfig) -> str: "deployed_at": datetime.utcnow().isoformat(), } - async with session.post( - f"{self.backend_url}/deployments/", json=deployment_data - ) as response: + async with session.post(f"{self.backend_url}/deployments/", json=deployment_data) as response: if response.status == 200: result = await response.json() logger.info(f"Deployed configuration {config.id} to data lake") return result.get("deployment_id", config.id) else: - raise Exception( - f"Failed to deploy configuration: {response.status}" - ) + raise Exception(f"Failed to deploy configuration: {response.status}") async def upload_document_with_metadata( self, @@ -160,9 +154,7 @@ async def upload_document_with_metadata( }, } - async with session.post( - f"{self.backend_url}/resources/", json=resource_data - ) as response: + async with session.post(f"{self.backend_url}/resources/", json=resource_data) as response: if response.status == 200: resource = await response.json() logger.info(f"Uploaded document to data lake with metadata") @@ -185,21 +177,15 @@ async def store_analysis_result(self, result: AnalysisResult) -> str: }, } - async with session.post( - f"{self.backend_url}/resources/", json=analysis_data - ) as response: + async with session.post(f"{self.backend_url}/resources/", json=analysis_data) as response: if response.status == 200: resource = await response.json() logger.info(f"Stored analysis result {result.id} in data lake") return resource["id"] else: - raise Exception( - f"Failed to store analysis result: {response.status}" - ) + raise Exception(f"Failed to store analysis result: {response.status}") - async def get_deployment_data( - self, deployment_id: str, data_types: List[str] = None - ) -> Dict[str, Any]: + async def get_deployment_data(self, deployment_id: str, data_types: List[str] = None) -> Dict[str, Any]: """Get all data for a specific deployment""" if data_types is None: data_types = ["documents", "chunks", "analysis_results"] @@ -213,16 +199,12 @@ async def get_deployment_data( async with aiohttp.ClientSession() as session: # Get documents for this deployment if "documents" in data_types: - documents = await self._get_documents_by_deployment( - session, deployment_id - ) + documents = await self._get_documents_by_deployment(session, deployment_id) deployment_data["data"]["documents"] = documents # Get analysis results for this deployment if "analysis_results" in data_types: - results = await self._get_analysis_results_by_deployment( - session, deployment_id - ) + results = await self._get_analysis_results_by_deployment(session, deployment_id) deployment_data["data"]["analysis_results"] = results # Get chunks for documents in this deployment @@ -232,9 +214,7 @@ async def get_deployment_data( return deployment_data - async def list_deployments( - self, deployment_type: Optional[DeploymentType] = None - ) -> List[Dict[str, Any]]: + async def list_deployments(self, deployment_type: Optional[DeploymentType] = None) -> List[Dict[str, Any]]: """List all deployments for this owner""" async with aiohttp.ClientSession() as session: # Filter resources by deployment metadata @@ -242,9 +222,7 @@ async def list_deployments( if deployment_type: params["deployment_type"] = deployment_type.value - async with session.get( - f"{self.backend_url}/deployments/", params=params - ) as response: + async with session.get(f"{self.backend_url}/deployments/", params=params) as response: if response.status == 200: return await response.json() else: @@ -280,9 +258,7 @@ async def get_data_by_type(self, deployment_type: DeploymentType) -> Dict[str, A } for resource in resources: - resource_type = resource.get("resource_metadata", {}).get( - "type", "document" - ) + resource_type = resource.get("resource_metadata", {}).get("type", "document") if resource_type == "analysis_result": grouped_data["analysis_results"].append(resource) else: @@ -292,9 +268,7 @@ async def get_data_by_type(self, deployment_type: DeploymentType) -> Dict[str, A else: return {"error": f"Failed to get data: {response.status}"} - async def _get_documents_by_deployment( - self, session: aiohttp.ClientSession, deployment_id: str - ) -> List[Dict[str, Any]]: + async def _get_documents_by_deployment(self, session: aiohttp.ClientSession, deployment_id: str) -> List[Dict[str, Any]]: """Get documents for a specific deployment""" async with session.get( f"{self.backend_url}/resources/", @@ -316,9 +290,7 @@ async def _get_analysis_results_by_deployment( return await response.json() return [] - async def _get_chunks_by_deployment( - self, session: aiohttp.ClientSession, deployment_id: str - ) -> List[Dict[str, Any]]: + async def _get_chunks_by_deployment(self, session: aiohttp.ClientSession, deployment_id: str) -> List[Dict[str, Any]]: """Get chunks for documents in a specific deployment""" # This would use the search endpoint with deployment filtering async with session.post( @@ -347,9 +319,7 @@ def __init__(self, owner: str, backend_url: str = "http://localhost:8000"): self.client = DataLakeClient(backend_url, owner) self.owner = owner - async def create_experiment( - self, name: str, description: str, question_set: str, config: Dict[str, Any] - ) -> str: + async def create_experiment(self, name: str, description: str, question_set: str, config: Dict[str, Any]) -> str: """Create a new experiment deployment""" deployment = DeploymentConfig( id=str(uuid.uuid4()), @@ -365,9 +335,7 @@ async def create_experiment( return await self.client.deploy_configuration(deployment) - async def promote_to_production( - self, experiment_id: str, production_name: str - ) -> str: + async def promote_to_production(self, experiment_id: str, production_name: str) -> str: """Promote an experiment to production""" # Get experiment data experiment_data = await self.client.get_deployment_data(experiment_id) @@ -387,9 +355,7 @@ async def promote_to_production( return await self.client.deploy_configuration(production_deployment) - async def upload_experimental_document( - self, document_url: str, experiment_id: str, filename: str = None - ) -> str: + async def upload_experimental_document(self, document_url: str, experiment_id: str, filename: str = None) -> str: """Upload document for experimentation""" metadata = DataMetadata( source=DataSource.REPORT_ANALYST, @@ -401,13 +367,9 @@ async def upload_experimental_document( tags=["experiment", "report_analyst"], ) - return await self.client.upload_document_with_metadata( - document_url, metadata, experiment_id - ) + return await self.client.upload_document_with_metadata(document_url, metadata, experiment_id) - async def store_experiment_results( - self, experiment_id: str, resource_id: str, analysis_results: Dict[str, Any] - ) -> str: + async def store_experiment_results(self, experiment_id: str, resource_id: str, analysis_results: Dict[str, Any]) -> str: """Store analysis results for an experiment""" metadata = DataMetadata( source=DataSource.REPORT_ANALYST, @@ -467,14 +429,10 @@ async def example_data_lake_usage(): analysis_results = { "question_set": "tcfd_v2", "model_used": "gpt-4o-mini", - "results": [ - {"question": "Climate risks?", "answer": "Significant risks identified..."} - ], + "results": [{"question": "Climate risks?", "answer": "Significant risks identified..."}], } - result_id = await integration.store_experiment_results( - experiment_id, resource_id, analysis_results - ) + result_id = await integration.store_experiment_results(experiment_id, resource_id, analysis_results) print(f"Stored results: {result_id}") diff --git a/report_analyst_jobs/event_handlers.py b/report_analyst_jobs/event_handlers.py index aa41ad772..976fae414 100644 --- a/report_analyst_jobs/event_handlers.py +++ b/report_analyst_jobs/event_handlers.py @@ -18,10 +18,10 @@ async def handle_document_ready(ctx: EventContext): try: event = DocumentReadyEvent(**ctx.data) logger.info(f"Processing document.ready for resource {event.resource_id}") - + # TODO: Integrate with existing document.ready processing logic # This could call the DocumentReadyProcessingConfig flow - + await ctx.message.ack() except Exception as e: logger.error(f"Error handling document.ready: {e}", exc_info=True) @@ -33,9 +33,9 @@ async def handle_analysis_job(ctx: EventContext): try: job_data = ctx.data logger.info(f"Processing analysis job: {job_data.get('id')}") - + # TODO: Integrate with existing analysis job processing logic - + await ctx.message.ack() except Exception as e: logger.error(f"Error handling analysis job: {e}", exc_info=True) @@ -47,9 +47,9 @@ async def handle_llm_request(ctx: EventContext): try: request_data = ctx.data logger.info(f"Processing LLM request: {request_data.get('request_id')}") - + # TODO: Integrate with existing LLM request processing logic - + await ctx.message.ack() except Exception as e: logger.error(f"Error handling LLM request: {e}", exc_info=True) @@ -61,9 +61,9 @@ async def handle_external_service_ready(ctx: EventContext): try: service_data = ctx.data logger.info(f"Processing external service ready: {service_data.get('service_id')}") - + # TODO: Integrate with existing external service handler - + await ctx.message.ack() except Exception as e: logger.error(f"Error handling external service ready: {e}", exc_info=True) @@ -75,11 +75,10 @@ async def handle_external_service_analysis(ctx: EventContext): try: request_data = ctx.data logger.info(f"Processing external service analysis request: {request_data.get('request_id')}") - + # TODO: Integrate with existing external service analysis logic - + await ctx.message.ack() except Exception as e: logger.error(f"Error handling external service analysis: {e}", exc_info=True) await ctx.message.ack() - diff --git a/report_analyst_jobs/event_router.py b/report_analyst_jobs/event_router.py index 0244435fb..462c4562d 100644 --- a/report_analyst_jobs/event_router.py +++ b/report_analyst_jobs/event_router.py @@ -13,8 +13,8 @@ import importlib import json import logging -from pathlib import Path from dataclasses import dataclass, field +from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Union import nats @@ -33,16 +33,16 @@ class EventActionRule: event_pattern: str """NATS subject pattern (supports wildcards like 'document.*' or 'analysis.job.>')""" - + action: Union[str, Callable] """Action to take: callable handler function, or 'ignore' to skip""" - + description: Optional[str] = None """Optional description of what this rule does""" - + enabled: bool = True """Whether this rule is enabled""" - + priority: int = 0 """Priority for matching (higher = checked first)""" @@ -53,16 +53,16 @@ class EventContext: subject: str """NATS subject/channel the event came from""" - + data: Dict[str, Any] """Parsed event data""" - + raw_data: bytes """Raw message data""" - + message: Any """Original NATS message object""" - + metadata: Dict[str, Any] = field(default_factory=dict) """Additional metadata""" @@ -70,23 +70,23 @@ class EventContext: class EventRouter: """ Simple event-action router for NATS events. - + Maps event patterns to actions using a table-based configuration. Actions can be handler functions or "ignore" to skip processing. - + Example: router = EventRouter(nats_url="nats://localhost:4222") - + # Define actions async def handle_document_ready(ctx: EventContext): print(f"Document ready: {ctx.data}") await ctx.message.ack() - + # Configure routing table router.add_rule("document.ready", handle_document_ready) router.add_rule("document.*", "ignore") # Ignore other document events router.add_rule("analysis.job.submit", handle_analysis_job) - + # Start processing await router.connect() await router.start() @@ -125,7 +125,7 @@ def add_rule( ): """ Add an event-action rule to the routing table. - + Args: event_pattern: NATS subject pattern (supports wildcards) action: Handler function or "ignore" @@ -143,9 +143,7 @@ def add_rule( self.rules.append(rule) # Sort by priority (higher first) self.rules.sort(key=lambda r: r.priority, reverse=True) - logger.info( - f"Added rule: {event_pattern} -> {action if action == IGNORE_ACTION else 'handler'}" - ) + logger.info(f"Added rule: {event_pattern} -> {action if action == IGNORE_ACTION else 'handler'}") def remove_rule(self, event_pattern: str): """Remove a rule by pattern""" @@ -171,16 +169,16 @@ def _match_subject(self, pattern: str, subject: str) -> bool: """ if pattern == subject: return True - + # Handle wildcards if pattern.endswith(".*"): prefix = pattern[:-2] return subject.startswith(prefix + ".") and "." not in subject[len(prefix) + 1 :] - + if pattern.endswith(".>"): prefix = pattern[:-2] return subject.startswith(prefix + ".") - + return False def _find_rule(self, subject: str) -> Optional[EventActionRule]: @@ -256,7 +254,7 @@ async def _handle_message(self, msg): async def start(self, subjects: Optional[List[str]] = None): """ Start processing events. - + Args: subjects: Optional list of subjects to subscribe to. If None, automatically subscribes to all subjects mentioned in rules. @@ -291,7 +289,7 @@ async def start(self, subjects: Optional[List[str]] = None): for subject in subjects: if subject in self._subscribed_subjects: continue - + try: await self.js.subscribe(subject, cb=self._handle_message) self._subscribed_subjects.add(subject) @@ -329,28 +327,28 @@ def from_yaml( ) -> "EventRouter": """ Create EventRouter from YAML configuration file. - + Args: yaml_path: Path to YAML config file. If None, looks for event_routing.yaml in the same directory as this module. handler_registry: Optional dict mapping handler names to callable functions. If None, tries to import handlers from paths specified in YAML. - + Returns: Configured EventRouter instance """ router = cls() - + # Default YAML path if yaml_path is None: yaml_path = Path(__file__).parent / "event_routing.yaml" else: yaml_path = Path(yaml_path) - + if not yaml_path.exists(): logger.warning(f"YAML config file not found: {yaml_path}") return router - + # Load YAML try: with open(yaml_path, "r", encoding="utf-8") as f: @@ -358,7 +356,7 @@ def from_yaml( except Exception as e: logger.error(f"Failed to load YAML config: {e}") return router - + # Load handler registry if handler_registry is None: handler_registry = {} @@ -370,7 +368,7 @@ def from_yaml( handler_registry[handler_name] = handler except Exception as e: logger.warning(f"Failed to load handler {handler_name} from {handler_path}: {e}") - + # Load routing rules routing_rules = config.get("routing", []) for rule_config in routing_rules: @@ -379,11 +377,11 @@ def from_yaml( description = rule_config.get("description") enabled = rule_config.get("enabled", True) priority = rule_config.get("priority", 0) - + if not pattern or not action_name: logger.warning(f"Invalid rule config: {rule_config}") continue - + # Resolve action if action_name == IGNORE_ACTION: action = IGNORE_ACTION @@ -392,7 +390,7 @@ def from_yaml( else: logger.warning(f"Handler not found for action: {action_name}, ignoring rule") continue - + router.add_rule( event_pattern=pattern, action=action, @@ -400,30 +398,29 @@ def from_yaml( enabled=enabled, priority=priority, ) - + logger.info(f"Loaded {len(router.rules)} routing rules from {yaml_path}") return router - + @staticmethod def _load_handler(handler_path: str) -> Callable: """ Load handler function from module path string. - + Args: handler_path: String like "module.path.to.function" - + Returns: Callable handler function """ parts = handler_path.split(".") module_path = ".".join(parts[:-1]) function_name = parts[-1] - + module = importlib.import_module(module_path) handler = getattr(module, function_name) - + if not callable(handler): raise ValueError(f"{handler_path} is not callable") - - return handler + return handler diff --git a/report_analyst_jobs/event_router_example.py b/report_analyst_jobs/event_router_example.py index e844862f2..e2da8588f 100644 --- a/report_analyst_jobs/event_router_example.py +++ b/report_analyst_jobs/event_router_example.py @@ -16,18 +16,18 @@ async def main(): """Example: Start event router from YAML configuration""" # Load router from YAML file (event_routing.yaml) router = EventRouter.from_yaml() - + # Or specify custom path # router = EventRouter.from_yaml("path/to/custom_routing.yaml") - + # Connect to NATS await router.connect() - + # Print routing table for inspection print("\nRouting Table:") for rule in router.get_routing_table(): print(f" {rule['pattern']:30} -> {rule['action']:20} (priority: {rule['priority']})") - + # Start processing events print("\nStarting event router...") await router.start() @@ -36,4 +36,3 @@ async def main(): if __name__ == "__main__": logging.basicConfig(level=logging.INFO) asyncio.run(main()) - diff --git a/report_analyst_jobs/event_routing.yaml b/report_analyst_jobs/event_routing.yaml index 69f891b33..20c5b4de4 100644 --- a/report_analyst_jobs/event_routing.yaml +++ b/report_analyst_jobs/event_routing.yaml @@ -97,3 +97,4 @@ handlers: handle_external_service_ready: "report_analyst_jobs.event_handlers.handle_external_service_ready" handle_external_service_analysis: "report_analyst_jobs.event_handlers.handle_external_service_analysis" + diff --git a/report_analyst_jobs/integration_examples.py b/report_analyst_jobs/integration_examples.py index b7ad1fdf2..0f3c6f6a4 100644 --- a/report_analyst_jobs/integration_examples.py +++ b/report_analyst_jobs/integration_examples.py @@ -215,9 +215,7 @@ def create_fastapi_endpoints(app): from .analysis_toolkit import analyze_document_sync @app.post("/analyze-document") - async def analyze_document_endpoint( - analysis_request: dict, background_tasks: BackgroundTasks - ): + async def analyze_document_endpoint(analysis_request: dict, background_tasks: BackgroundTasks): """FastAPI endpoint with background task""" def run_analysis(): diff --git a/report_analyst_jobs/interfaces.py b/report_analyst_jobs/interfaces.py index 997e705e6..46ac88f6c 100644 --- a/report_analyst_jobs/interfaces.py +++ b/report_analyst_jobs/interfaces.py @@ -85,9 +85,7 @@ def to_dict(self) -> Dict[str, Any]: "error": self.error, "progress": self.progress, "started_at": self.started_at.isoformat() if self.started_at else None, - "completed_at": ( - self.completed_at.isoformat() if self.completed_at else None - ), + "completed_at": (self.completed_at.isoformat() if self.completed_at else None), "metadata": self.metadata or {}, } diff --git a/report_analyst_jobs/llm_integration.py b/report_analyst_jobs/llm_integration.py index 96503dc0e..6b53f7bcf 100644 --- a/report_analyst_jobs/llm_integration.py +++ b/report_analyst_jobs/llm_integration.py @@ -104,18 +104,14 @@ async def disconnect(self): if self.nc: await self.nc.close() - async def analyze_question( - self, question: str, context_chunks: List[str], model: str = "gpt-4o-mini" - ) -> str: + async def analyze_question(self, question: str, context_chunks: List[str], model: str = "gpt-4o-mini") -> str: """ Analyze a question against context chunks using search backend LLM. This replaces direct LLM calls in report-analyst. """ # Build prompt for question analysis - context = "\n\n".join( - [f"Chunk {i+1}: {chunk}" for i, chunk in enumerate(context_chunks)] - ) + context = "\n\n".join([f"Chunk {i+1}: {chunk}" for i, chunk in enumerate(context_chunks)]) prompt = f"""Please analyze the following question based on the provided context: @@ -153,9 +149,7 @@ async def summarize_chunks( ) -> str: """Summarize document chunks using search backend LLM""" - content = "\n\n".join( - [f"Section {i+1}: {chunk}" for i, chunk in enumerate(chunks)] - ) + content = "\n\n".join([f"Section {i+1}: {chunk}" for i, chunk in enumerate(chunks)]) prompt = f"""Please provide a {summary_type} summary of the following document sections: @@ -186,26 +180,18 @@ async def _send_request(self, request: LLMRequest) -> str: self.pending_requests[request.id] = asyncio.Event() # Send request - await self.js.publish( - "llm.request", json.dumps(asdict(request), default=str).encode() - ) + await self.js.publish("llm.request", json.dumps(asdict(request), default=str).encode()) # Wait for response (with timeout) try: - await asyncio.wait_for( - self.pending_requests[request.id].wait(), timeout=60.0 - ) + await asyncio.wait_for(self.pending_requests[request.id].wait(), timeout=60.0) # Get response response_data = self.pending_requests.get(f"{request.id}_response") if response_data and not response_data.get("error"): return response_data["response"] else: - error = ( - response_data.get("error", "Unknown error") - if response_data - else "No response received" - ) + error = response_data.get("error", "Unknown error") if response_data else "No response received" raise Exception(f"LLM request failed: {error}") except asyncio.TimeoutError: @@ -276,9 +262,7 @@ async def _process_request(self, msg): processing_time=1.5, # Could track actual time ) - await self.js.publish( - "llm.response", json.dumps(asdict(response), default=str).encode() - ) + await self.js.publish("llm.response", json.dumps(asdict(response), default=str).encode()) await msg.ack() logger.info(f"LLM request {request.id} completed") @@ -294,9 +278,7 @@ async def _process_request(self, msg): error=str(e), ) - await self.js.publish( - "llm.response", json.dumps(asdict(error_response), default=str).encode() - ) + await self.js.publish("llm.response", json.dumps(asdict(error_response), default=str).encode()) await msg.ack() async def _call_search_backend_llm(self, request: LLMRequest) -> str: diff --git a/report_analyst_jobs/local_executor.py b/report_analyst_jobs/local_executor.py index 5e15f89b2..bdfc093f1 100644 --- a/report_analyst_jobs/local_executor.py +++ b/report_analyst_jobs/local_executor.py @@ -39,9 +39,7 @@ async def submit_job(self, job: JobDefinition) -> str: """Submit a job for local execution""" # Create initial job result - job_result = JobResult( - job_id=job.job_id, status=JobStatus.PENDING, started_at=datetime.utcnow() - ) + job_result = JobResult(job_id=job.job_id, status=JobStatus.PENDING, started_at=datetime.utcnow()) self.job_storage[job.job_id] = job_result # Start job execution task diff --git a/report_analyst_jobs/nats_integration.py b/report_analyst_jobs/nats_integration.py index 877b3ed39..cce37ac95 100644 --- a/report_analyst_jobs/nats_integration.py +++ b/report_analyst_jobs/nats_integration.py @@ -76,7 +76,7 @@ def __post_init__(self): class DocumentReadyProcessingConfig: """ Configuration for automatic document.ready event processing. - + Controls how document.ready events are handled: - Whether to pull chunks from backend or use provided chunks - Whether to run analysis @@ -87,32 +87,34 @@ class DocumentReadyProcessingConfig: # Chunk retrieval strategy pull_chunks: bool = True """Whether to pull chunks from backend. If False, chunks must be provided in event.""" - + # Analysis configuration question_set: str = "tcfd" """Question set to use for analysis""" - - analysis_config: Dict[str, Any] = field(default_factory=lambda: { - "model": "gpt-4o-mini", - "temperature": 0.1, - }) + + analysis_config: Dict[str, Any] = field( + default_factory=lambda: { + "model": "gpt-4o-mini", + "temperature": 0.1, + } + ) """Analysis configuration (model, temperature, etc.)""" - + # Result storage store_to_backend: bool = True """Whether to store analysis results back to backend""" - + # Error handling ack_on_error: bool = True """Whether to acknowledge message even on error (prevents redelivery loops)""" - + # Chunk retrieval options (when pull_chunks=True) chunk_retrieval_method: str = "search" # "search" or "direct" """Method to retrieve chunks: 'search' uses /search/ endpoint, 'direct' uses /resources/{id}/chunks""" - + max_chunks: Optional[int] = None """Maximum number of chunks to retrieve (None = no limit)""" - + def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" return { @@ -151,14 +153,10 @@ async def get_resource_chunks(self, resource_id: str) -> List[Dict[str, Any]]: { "id": chunk_data["chunk"]["id"], "text": chunk_data["chunk"]["chunk_text"], - "metadata": chunk_data["chunk"][ - "chunk_metadata" - ], + "metadata": chunk_data["chunk"]["chunk_metadata"], } ) - logger.info( - f"Retrieved {len(chunks)} chunks for resource {resource_id}" - ) + logger.info(f"Retrieved {len(chunks)} chunks for resource {resource_id}") return chunks else: raise Exception(f"Failed to get chunks: {response.status}") @@ -199,9 +197,7 @@ async def disconnect(self): logger.info("Disconnected from NATS") # For Search Backend to publish when PDF processing is complete - async def publish_document_ready( - self, resource_id: str, document_url: str, chunks_count: int - ): + async def publish_document_ready(self, resource_id: str, document_url: str, chunks_count: int): """Publish document ready event (called from search backend)""" event = DocumentReadyEvent( resource_id=resource_id, @@ -210,16 +206,12 @@ async def publish_document_ready( status="ready", ) - await self.js.publish( - "document.ready", json.dumps(asdict(event), default=str).encode() - ) + await self.js.publish("document.ready", json.dumps(asdict(event), default=str).encode()) logger.info(f"Published document ready event for resource {resource_id}") # For Clients to request analysis - async def submit_analysis_job( - self, resource_id: str, question_set: str, analysis_config: Dict[str, Any] - ) -> str: + async def submit_analysis_job(self, resource_id: str, question_set: str, analysis_config: Dict[str, Any]) -> str: """Submit analysis job for a resource that's already processed""" job_id = str(uuid.uuid4()) job = AnalysisJob( @@ -230,9 +222,7 @@ async def submit_analysis_job( ) # Send job to NATS - await self.js.publish( - "analysis.job.submit", json.dumps(asdict(job), default=str).encode() - ) + await self.js.publish("analysis.job.submit", json.dumps(asdict(job), default=str).encode()) logger.info(f"Submitted analysis job {job_id} for resource {resource_id}") return job_id @@ -265,12 +255,8 @@ async def process_document_ready_events( if config is None: config = DocumentReadyProcessingConfig() - await self.js.subscribe( - "document.ready", cb=lambda msg: self._handle_document_ready(msg, config) - ) - logger.info( - f"Started processing document.ready events from NATS with config: {config.to_dict()}" - ) + await self.js.subscribe("document.ready", cb=lambda msg: self._handle_document_ready(msg, config)) + logger.info(f"Started processing document.ready events from NATS with config: {config.to_dict()}") # Keep processing try: @@ -314,32 +300,22 @@ async def _handle_document_ready( chunks = event.chunks logger.info(f"Using {len(chunks)} chunks provided in event") else: - logger.warning( - f"No chunks provided and pull_chunks=False for resource {event.resource_id}" - ) + logger.warning(f"No chunks provided and pull_chunks=False for resource {event.resource_id}") if config.ack_on_error: await msg.ack() return # Step 2: Run analysis logger.info(f"Running analysis for resource {event.resource_id}") - analysis_result = await self._run_analysis( - chunks, config.question_set, config.analysis_config - ) + analysis_result = await self._run_analysis(chunks, config.question_set, config.analysis_config) # Step 3: Store results back to backend (if configured) if config.store_to_backend: - logger.info( - f"Storing analysis results back to backend for resource {event.resource_id}" - ) - await self._store_analysis_to_backend( - event.resource_id, analysis_result, config.question_set - ) + logger.info(f"Storing analysis results back to backend for resource {event.resource_id}") + await self._store_analysis_to_backend(event.resource_id, analysis_result, config.question_set) await msg.ack() - logger.info( - f"Successfully processed document.ready for resource {event.resource_id}" - ) + logger.info(f"Successfully processed document.ready for resource {event.resource_id}") except Exception as e: logger.error(f"Error processing document.ready event: {e}", exc_info=True) @@ -354,12 +330,12 @@ async def _get_chunks_for_resource( ) -> List[Dict[str, Any]]: """ Get chunks for a resource using specified method. - + Args: resource_id: Resource ID method: "search" (use /search/ endpoint) or "direct" (use /resources/{id}/chunks) max_chunks: Maximum number of chunks to return (None = no limit) - + Returns: List of chunk dictionaries """ @@ -367,9 +343,7 @@ async def _get_chunks_for_resource( # Try direct endpoint first try: async with aiohttp.ClientSession() as session: - async with session.get( - f"{self.search_backend.base_url}/resources/{resource_id}/chunks" - ) as response: + async with session.get(f"{self.search_backend.base_url}/resources/{resource_id}/chunks") as response: if response.status == 200: data = await response.json() chunks = data.get("chunks", []) @@ -377,9 +351,7 @@ async def _get_chunks_for_resource( chunks = chunks[:max_chunks] return chunks except Exception as e: - logger.warning( - f"Direct chunk retrieval failed for {resource_id}, falling back to search: {e}" - ) + logger.warning(f"Direct chunk retrieval failed for {resource_id}, falling back to search: {e}") # Default: use search endpoint (existing method) chunks = await self.search_backend.get_resource_chunks(resource_id) @@ -397,13 +369,12 @@ async def _store_analysis_to_backend( try: # Use BackendService to store results (synchronous, so run in executor) import asyncio + from report_analyst_search_backend.backend_service import BackendService from report_analyst_search_backend.config import BackendConfig # Create BackendService from the base_url - config = BackendConfig( - use_backend=True, backend_url=self.search_backend.base_url - ) + config = BackendConfig(use_backend=True, backend_url=self.search_backend.base_url) backend_service = BackendService(config) # Run synchronous store_analysis_results in executor @@ -418,13 +389,9 @@ async def _store_analysis_to_backend( ) if result_id: - logger.info( - f"Stored analysis results for resource {resource_id}: {result_id}" - ) + logger.info(f"Stored analysis results for resource {resource_id}: {result_id}") else: - logger.warning( - f"Failed to store analysis results for resource {resource_id}" - ) + logger.warning(f"Failed to store analysis results for resource {resource_id}") except Exception as e: logger.error(f"Error storing analysis to backend: {e}") @@ -435,26 +402,20 @@ async def _process_analysis_job(self, msg): job_data = json.loads(msg.data.decode()) job = AnalysisJob(**job_data) - logger.info( - f"Processing analysis job {job.id} for resource {job.resource_id}" - ) + logger.info(f"Processing analysis job {job.id} for resource {job.resource_id}") # Update job status job.status = JobStatus.PROCESSING job.updated_at = datetime.utcnow() # Send status update - await self.js.publish( - "analysis.job.status", json.dumps(asdict(job), default=str).encode() - ) + await self.js.publish("analysis.job.status", json.dumps(asdict(job), default=str).encode()) # Get chunks from search backend (already processed) chunks = await self.search_backend.get_resource_chunks(job.resource_id) # Run analysis using report analyst toolkit - analysis_result = await self._run_analysis( - chunks, job.question_set, job.analysis_config - ) + analysis_result = await self._run_analysis(chunks, job.question_set, job.analysis_config) # Complete job job.status = JobStatus.COMPLETED @@ -467,9 +428,7 @@ async def _process_analysis_job(self, msg): job.updated_at = datetime.utcnow() # Send completion notification - await self.js.publish( - "analysis.job.completed", json.dumps(asdict(job), default=str).encode() - ) + await self.js.publish("analysis.job.completed", json.dumps(asdict(job), default=str).encode()) await msg.ack() logger.info(f"Analysis job {job.id} completed successfully") @@ -482,15 +441,11 @@ async def _process_analysis_job(self, msg): job.error = str(e) job.updated_at = datetime.utcnow() - await self.js.publish( - "analysis.job.failed", json.dumps(asdict(job), default=str).encode() - ) + await self.js.publish("analysis.job.failed", json.dumps(asdict(job), default=str).encode()) await msg.ack() - async def _run_analysis( - self, chunks: List[Dict[str, Any]], question_set: str, config: Dict[str, Any] - ) -> Dict[str, Any]: + async def _run_analysis(self, chunks: List[Dict[str, Any]], question_set: str, config: Dict[str, Any]) -> Dict[str, Any]: """Run the actual analysis using report analyst toolkit""" # Import here to avoid circular imports from .analysis_toolkit import analyze_document_with_chunks @@ -498,14 +453,10 @@ async def _run_analysis( # Convert chunks to the format expected by analysis toolkit formatted_chunks = [] for chunk in chunks: - formatted_chunks.append( - {"text": chunk["text"], "metadata": chunk["metadata"]} - ) + formatted_chunks.append({"text": chunk["text"], "metadata": chunk["metadata"]}) # Run analysis - result = await analyze_document_with_chunks( - chunks=formatted_chunks, question_set=question_set, config=config - ) + result = await analyze_document_with_chunks(chunks=formatted_chunks, question_set=question_set, config=config) return result @@ -523,13 +474,9 @@ async def __aenter__(self): async def __aexit__(self, exc_type, exc_val, exc_tb): await self.coordinator.disconnect() - async def notify_document_ready( - self, resource_id: str, document_url: str, chunks_count: int - ): + async def notify_document_ready(self, resource_id: str, document_url: str, chunks_count: int): """Notify that document processing is complete""" - await self.coordinator.publish_document_ready( - resource_id, document_url, chunks_count - ) + await self.coordinator.publish_document_ready(resource_id, document_url, chunks_count) class NATSAnalysisWorker: @@ -559,32 +506,20 @@ def __init__(self, nats_url: str = "nats://localhost:4222"): async def __aenter__(self): await self.coordinator.connect() # Subscribe to job results - await self.coordinator.js.subscribe( - "analysis.job.completed", cb=self._handle_completed - ) - await self.coordinator.js.subscribe( - "analysis.job.failed", cb=self._handle_failed - ) - await self.coordinator.js.subscribe( - "analysis.job.status", cb=self._handle_status - ) + await self.coordinator.js.subscribe("analysis.job.completed", cb=self._handle_completed) + await self.coordinator.js.subscribe("analysis.job.failed", cb=self._handle_failed) + await self.coordinator.js.subscribe("analysis.job.status", cb=self._handle_status) return self async def __aexit__(self, exc_type, exc_val, exc_tb): await self.coordinator.disconnect() - async def analyze_resource( - self, resource_id: str, question_set: str, analysis_config: Dict[str, Any] - ) -> str: + async def analyze_resource(self, resource_id: str, question_set: str, analysis_config: Dict[str, Any]) -> str: """Submit analysis job for a resource that's already processed""" - job_id = await self.coordinator.submit_analysis_job( - resource_id, question_set, analysis_config - ) + job_id = await self.coordinator.submit_analysis_job(resource_id, question_set, analysis_config) return job_id - async def wait_for_completion( - self, job_id: str, timeout: int = 300 - ) -> Dict[str, Any]: + async def wait_for_completion(self, job_id: str, timeout: int = 300) -> Dict[str, Any]: """Wait for job completion""" start_time = datetime.utcnow() @@ -595,9 +530,7 @@ async def wait_for_completion( return result if (datetime.utcnow() - start_time).seconds > timeout: - raise TimeoutError( - f"Job {job_id} did not complete within {timeout} seconds" - ) + raise TimeoutError(f"Job {job_id} did not complete within {timeout} seconds") await asyncio.sleep(2) diff --git a/report_analyst_jobs/search_backend_integration.py b/report_analyst_jobs/search_backend_integration.py index 7a443dc74..0823bf255 100644 --- a/report_analyst_jobs/search_backend_integration.py +++ b/report_analyst_jobs/search_backend_integration.py @@ -40,12 +40,8 @@ async def notify_document_ready( """ try: async with NATSSearchBackendPublisher(nats_url) as publisher: - await publisher.notify_document_ready( - resource_id, document_url, chunks_count - ) - logger.info( - f"Successfully notified NATS that document {resource_id} is ready" - ) + await publisher.notify_document_ready(resource_id, document_url, chunks_count) + logger.info(f"Successfully notified NATS that document {resource_id} is ready") except Exception as e: logger.error(f"Failed to notify NATS about document {resource_id}: {e}") # Don't raise - this is a nice-to-have feature @@ -63,9 +59,7 @@ def notify_document_ready_sync( Use this in synchronous contexts like Celery tasks. """ try: - asyncio.run( - notify_document_ready(resource_id, document_url, chunks_count, nats_url) - ) + asyncio.run(notify_document_ready(resource_id, document_url, chunks_count, nats_url)) except Exception as e: logger.error(f"Failed to notify NATS about document {resource_id}: {e}") @@ -98,16 +92,12 @@ async def disconnect(self): self._connected = False logger.info("Disconnected from NATS") - async def notify_document_ready( - self, resource_id: str, document_url: str, chunks_count: int - ): + async def notify_document_ready(self, resource_id: str, document_url: str, chunks_count: int): """Notify that document is ready""" if not self._connected: await self.connect() - await self.publisher.notify_document_ready( - resource_id, document_url, chunks_count - ) + await self.publisher.notify_document_ready(resource_id, document_url, chunks_count) logger.info(f"Notified NATS that document {resource_id} is ready") diff --git a/report_analyst_search_backend/LICENSE b/report_analyst_search_backend/LICENSE index efd9a365e..67635e2d8 100644 --- a/report_analyst_search_backend/LICENSE +++ b/report_analyst_search_backend/LICENSE @@ -17,3 +17,4 @@ This software is part of the Open Sustainability Analysis project. + diff --git a/report_analyst_search_backend/backend_service.py b/report_analyst_search_backend/backend_service.py index 1061e9103..b271b3399 100644 --- a/report_analyst_search_backend/backend_service.py +++ b/report_analyst_search_backend/backend_service.py @@ -188,9 +188,7 @@ def get_chunks(self, resource_id: str) -> List[Dict[str, Any]]: return chunks else: - raise BackendServiceError( - f"Failed to get chunks: {response.status_code}" - ) + raise BackendServiceError(f"Failed to get chunks: {response.status_code}") except requests.RequestException as e: raise BackendServiceError(f"Error getting chunks: {str(e)}") @@ -234,16 +232,12 @@ def submit_analysis_job(self, resource_id: str, question_set: str) -> str: job_result = response.json() return job_result.get("job_id") else: - raise BackendServiceError( - f"Failed to submit analysis job: {response.status_code}" - ) + raise BackendServiceError(f"Failed to submit analysis job: {response.status_code}") except requests.RequestException as e: raise BackendServiceError(f"Error submitting analysis job: {str(e)}") - def wait_for_analysis( - self, analysis_job_id: str, timeout: int = 300 - ) -> Dict[str, Any]: + def wait_for_analysis(self, analysis_job_id: str, timeout: int = 300) -> Dict[str, Any]: """ Wait for analysis job to complete. @@ -285,18 +279,14 @@ def wait_for_analysis( # Check timeout elapsed = time.time() - start_time if elapsed > timeout: - raise BackendServiceError( - f"Analysis timed out after {timeout} seconds" - ) + raise BackendServiceError(f"Analysis timed out after {timeout} seconds") time.sleep(5) except requests.RequestException as e: raise BackendServiceError(f"Error checking analysis status: {str(e)}") - def get_analysis_results( - self, analysis_job_id: str = None, resource_id: str = None - ) -> Optional[Dict[str, Any]]: + def get_analysis_results(self, analysis_job_id: str = None, resource_id: str = None) -> Optional[Dict[str, Any]]: """ Get stored analysis results from backend database. @@ -384,25 +374,16 @@ def store_analysis_results( if response.status_code == 200: result = response.json() result_id = result.get("id") or result.get("result_id") - logger.info( - f"Stored analysis results for resource {resource_id}: {result_id}" - ) + logger.info(f"Stored analysis results for resource {resource_id}: {result_id}") return result_id elif response.status_code == 404: # Endpoint doesn't exist, try alternative: use submit_analysis_job pattern - logger.warning( - "/analysis/results/ endpoint not found, " - "trying alternative storage method" - ) + logger.warning("/analysis/results/ endpoint not found, " "trying alternative storage method") # Alternative: Store as a new resource with analysis results - return self._store_analysis_as_resource( - resource_id, analysis_results, question_set, metadata - ) + return self._store_analysis_as_resource(resource_id, analysis_results, question_set, metadata) else: error_text = response.text - logger.error( - f"Failed to store analysis results: {response.status_code} - {error_text}" - ) + logger.error(f"Failed to store analysis results: {response.status_code} - {error_text}") return None except requests.RequestException as e: @@ -442,14 +423,10 @@ def _store_analysis_as_resource( if response.status_code == 200: resource = response.json() result_id = resource.get("id") - logger.info( - f"Stored analysis results as resource for {resource_id}: {result_id}" - ) + logger.info(f"Stored analysis results as resource for {resource_id}: {result_id}") return result_id else: - logger.error( - f"Failed to store analysis as resource: {response.status_code}" - ) + logger.error(f"Failed to store analysis as resource: {response.status_code}") return None except requests.RequestException as e: @@ -459,9 +436,7 @@ def _store_analysis_as_resource( def _get_resources(self) -> List[Dict[str, Any]]: """Get all resources from backend""" try: - response = requests.get( - f"{self.config.backend_url}/resources/", timeout=self.timeout - ) + response = requests.get(f"{self.config.backend_url}/resources/", timeout=self.timeout) return response.json() if response.status_code == 200 else [] except requests.RequestException: return [] diff --git a/report_analyst_search_backend/config.py b/report_analyst_search_backend/config.py index e9fb73965..6d962a405 100644 --- a/report_analyst_search_backend/config.py +++ b/report_analyst_search_backend/config.py @@ -29,11 +29,7 @@ class BackendConfig: @property def has_advanced_features(self) -> bool: """Check if any advanced features are enabled""" - return ( - self.use_centralized_llm - or self.use_data_lake - or self.use_full_backend_analysis - ) + return self.use_centralized_llm or self.use_data_lake or self.use_full_backend_analysis @property def flow_type(self) -> str: @@ -78,7 +74,8 @@ def configure_backend_integration() -> BackendConfig: # Note: Enterprise Integration (S3+NATS) is now shown in the main Settings section above # Basic backend toggle - st.markdown(""" + st.markdown( + """ - """, unsafe_allow_html=True) + """, + unsafe_allow_html=True, + ) use_backend = st.checkbox( "Use Search Backend", value=False, @@ -95,10 +94,10 @@ def configure_backend_integration() -> BackendConfig: # Backend URL (shown even if backend is disabled, for future use) backend_url = st.text_input( - "Backend URL", - value=os.getenv("BACKEND_URL", "http://localhost:8000"), + "Backend URL", + value=os.getenv("BACKEND_URL", "http://localhost:8000"), help="Search backend API URL", - disabled=not use_backend + disabled=not use_backend, ) if not use_backend: @@ -134,7 +133,9 @@ def configure_backend_integration() -> BackendConfig: ) nats_url = st.text_input( - "NATS URL", value=os.getenv("NATS_URL", "nats://localhost:4222"), help="URL of your NATS server" + "NATS URL", + value=os.getenv("NATS_URL", "nats://localhost:4222"), + help="URL of your NATS server", ) # Data lake configuration @@ -210,9 +211,7 @@ def display_config_status(config: BackendConfig): st.info(f"Using centralized LLM via NATS: {config.nats_url}") if config.use_data_lake: - st.info( - f"Data lake enabled for owner: {config.owner} ({config.deployment_type})" - ) + st.info(f"Data lake enabled for owner: {config.owner} ({config.deployment_type})") if config.use_full_backend_analysis: st.info("Complete backend analysis enabled - backend does all the work!") diff --git a/report_analyst_search_backend/external_service_client.py b/report_analyst_search_backend/external_service_client.py index a26c43cc1..6bf8e19c9 100644 --- a/report_analyst_search_backend/external_service_client.py +++ b/report_analyst_search_backend/external_service_client.py @@ -34,9 +34,7 @@ def __init__( base_url: Base URL for HTTP API (defaults to env var) nats_url: NATS server URL (defaults to env var) """ - self.base_url = base_url or os.getenv( - "REPORT_ANALYST_API_URL", "http://localhost:8000" - ) + self.base_url = base_url or os.getenv("REPORT_ANALYST_API_URL", "http://localhost:8000") self.nats_url = nats_url or os.getenv("NATS_URL", "nats://localhost:4222") self.nc = None self.js = None @@ -114,9 +112,7 @@ async def notify_ready( else: return await self._notify_via_http(notification) - async def _notify_via_nats( - self, notification: ExternalServiceReadyEvent - ) -> bool: + async def _notify_via_nats(self, notification: ExternalServiceReadyEvent) -> bool: """Send notification via NATS""" try: await self.connect_nats() @@ -139,8 +135,7 @@ async def _notify_via_nats( message_data = json.dumps(notification_dict).encode() await self.js.publish("external.service.ready", message_data) logger.info( - f"Published external service notification via NATS: " - f"{notification.service_id}/{notification.request_id}" + f"Published external service notification via NATS: " f"{notification.service_id}/{notification.request_id}" ) return True @@ -148,9 +143,7 @@ async def _notify_via_nats( logger.error(f"Failed to notify via NATS: {e}") return False - async def _notify_via_http( - self, notification: ExternalServiceReadyEvent - ) -> bool: + async def _notify_via_http(self, notification: ExternalServiceReadyEvent) -> bool: """Send notification via HTTP""" try: url = f"{self.base_url}/external/services/{notification.service_id}/notify" @@ -179,9 +172,7 @@ async def _notify_via_http( return True else: error_text = await response.text() - logger.error( - f"HTTP notification failed: {response.status} - {error_text}" - ) + logger.error(f"HTTP notification failed: {response.status} - {error_text}") return False except Exception as e: @@ -259,9 +250,7 @@ async def _request_analysis_via_nats(self, request_data: Dict) -> Optional[str]: logger.error(f"Failed to request analysis via NATS: {e}") return None - async def _request_analysis_via_http( - self, service_id: str, request_data: Dict - ) -> Optional[str]: + async def _request_analysis_via_http(self, service_id: str, request_data: Dict) -> Optional[str]: """Request analysis via HTTP""" try: url = f"{self.base_url}/external/services/{service_id}/analyze" @@ -275,18 +264,14 @@ async def _request_analysis_via_http( return request_id else: error_text = await response.text() - logger.error( - f"HTTP analysis request failed: {response.status} - {error_text}" - ) + logger.error(f"HTTP analysis request failed: {response.status} - {error_text}") return None except Exception as e: logger.error(f"Failed to request analysis via HTTP: {e}") return None - async def get_results( - self, service_id: str, request_id: str - ) -> Optional[Dict[str, Any]]: + async def get_results(self, service_id: str, request_id: str) -> Optional[Dict[str, Any]]: """ Poll for analysis results. @@ -311,12 +296,9 @@ async def get_results( return None else: error_text = await response.text() - logger.error( - f"Failed to get results: {response.status} - {error_text}" - ) + logger.error(f"Failed to get results: {response.status} - {error_text}") return None except Exception as e: logger.error(f"Error polling for results: {e}") return None - diff --git a/report_analyst_search_backend/external_service_delivery.py b/report_analyst_search_backend/external_service_delivery.py index 640d165b9..e75ba61c9 100644 --- a/report_analyst_search_backend/external_service_delivery.py +++ b/report_analyst_search_backend/external_service_delivery.py @@ -105,8 +105,7 @@ async def _deliver_via_nats(self, response_data: Dict[str, Any]) -> bool: message_data = json.dumps(response_data).encode() await self.js.publish("external.service.analysis.response", message_data) logger.info( - f"Published analysis results via NATS: " - f"{response_data['service_id']}/{response_data['request_id']}" + f"Published analysis results via NATS: " f"{response_data['service_id']}/{response_data['request_id']}" ) return True @@ -114,9 +113,7 @@ async def _deliver_via_nats(self, response_data: Dict[str, Any]) -> bool: logger.error(f"Failed to deliver results via NATS: {e}") return False - async def _deliver_via_poll( - self, request_id: str, response_data: Dict[str, Any] - ) -> bool: + async def _deliver_via_poll(self, request_id: str, response_data: Dict[str, Any]) -> bool: """Store results for HTTP polling""" try: # Store results in memory (in production, use database) @@ -143,4 +140,3 @@ def get_results(self, request_id: str) -> Optional[Dict[str, Any]]: def clear_results(self, request_id: str): """Clear stored results (cleanup)""" self._result_storage.pop(request_id, None) - diff --git a/report_analyst_search_backend/external_service_handler.py b/report_analyst_search_backend/external_service_handler.py index d50b83525..dc411eb64 100644 --- a/report_analyst_search_backend/external_service_handler.py +++ b/report_analyst_search_backend/external_service_handler.py @@ -58,10 +58,7 @@ def __init__(self): def _init_s3_client(self): """Initialize S3 client if credentials are available""" try: - if all( - os.getenv(var) - for var in ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"] - ): + if all(os.getenv(var) for var in ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"]): self.s3_client = boto3.client( "s3", endpoint_url=os.getenv("S3_ENDPOINT_URL"), @@ -107,30 +104,20 @@ async def handle_external_notification( if notification.content_type == "s3_url": if not notification.s3_url: - return ProcessingResult( - success=False, error="S3 URL not provided" - ) + return ProcessingResult(success=False, error="S3 URL not provided") chunks = await self._process_s3_url(notification.s3_url) return ProcessingResult(success=True, chunks=chunks) elif notification.content_type == "chunks": if not notification.chunks: - return ProcessingResult( - success=False, error="Chunks not provided" - ) - chunks = await self._process_provided_chunks( - notification.chunks, rechunk_mode - ) + return ProcessingResult(success=False, error="Chunks not provided") + chunks = await self._process_provided_chunks(notification.chunks, rechunk_mode) return ProcessingResult(success=True, chunks=chunks) elif notification.content_type == "pages": if not notification.pages: - return ProcessingResult( - success=False, error="Pages not provided" - ) - chunks = await self._process_provided_pages( - notification.pages, rechunk_mode - ) + return ProcessingResult(success=False, error="Pages not provided") + chunks = await self._process_provided_pages(notification.pages, rechunk_mode) return ProcessingResult(success=True, chunks=chunks) else: @@ -177,9 +164,7 @@ async def _process_s3_url(self, s3_url: str) -> List[Dict[str, Any]]: logger.info(f"Downloaded {len(file_bytes)} bytes from S3") # Save to temporary file for processing - with tempfile.NamedTemporaryFile( - delete=False, suffix=".pdf", mode="wb" - ) as tmp_file: + with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf", mode="wb") as tmp_file: tmp_file.write(file_bytes) tmp_file_path = tmp_file.name @@ -232,9 +217,7 @@ async def _process_s3_url(self, s3_url: str) -> List[Dict[str, Any]]: logger.error(f"Error processing S3 URL: {e}") raise - async def _process_provided_chunks( - self, chunks: List[Dict[str, Any]], rechunk_mode: str - ) -> List[Dict[str, Any]]: + async def _process_provided_chunks(self, chunks: List[Dict[str, Any]], rechunk_mode: str) -> List[Dict[str, Any]]: """ Process provided chunks (re-chunk if needed). @@ -252,9 +235,7 @@ async def _process_provided_chunks( elif rechunk_mode == "always": # Always re-chunk - combine all text and re-chunk # This would use the analyzer's chunking logic - combined_text = "\n\n".join( - chunk.get("text", chunk.get("chunk_text", "")) for chunk in chunks - ) + combined_text = "\n\n".join(chunk.get("text", chunk.get("chunk_text", "")) for chunk in chunks) # For now, return normalized chunks - actual re-chunking would be done by analyzer logger.info(f"Re-chunking {len(chunks)} chunks (always mode)") return self._normalize_chunks(chunks) @@ -265,14 +246,10 @@ async def _process_provided_chunks( return self._normalize_chunks(chunks) else: # Re-chunk needed - logger.info( - f"Chunks don't match format, re-chunking {len(chunks)} chunks" - ) + logger.info(f"Chunks don't match format, re-chunking {len(chunks)} chunks") return self._normalize_chunks(chunks) - async def _process_provided_pages( - self, pages: List[Dict[str, Any]], rechunk_mode: str - ) -> List[Dict[str, Any]]: + async def _process_provided_pages(self, pages: List[Dict[str, Any]], rechunk_mode: str) -> List[Dict[str, Any]]: """ Process provided pages (convert to chunks). @@ -334,4 +311,3 @@ def _chunks_match_format(self, chunks: List[Dict[str, Any]]) -> bool: has_id = "id" in first_chunk or "chunk_id" in first_chunk return has_text and has_id - diff --git a/report_analyst_search_backend/flow_orchestrator.py b/report_analyst_search_backend/flow_orchestrator.py index b018feeeb..389301fdb 100644 --- a/report_analyst_search_backend/flow_orchestrator.py +++ b/report_analyst_search_backend/flow_orchestrator.py @@ -87,16 +87,12 @@ def process_document(self, uploaded_file) -> ProcessingResult: elif flow_type == "complete_backend": return self._process_complete_backend(uploaded_file) else: - return ProcessingResult( - success=False, error=f"Unknown flow type: {flow_type}" - ) + return ProcessingResult(success=False, error=f"Unknown flow type: {flow_type}") except Exception as e: logger.error(f"Document processing failed: {e}") return ProcessingResult(success=False, error=str(e)) - def analyze_document( - self, chunks: List[Dict[str, Any]], questions: List[str] - ) -> AnalysisResult: + def analyze_document(self, chunks: List[Dict[str, Any]], questions: List[str]) -> AnalysisResult: """ Analyze document based on configuration. @@ -117,16 +113,12 @@ def analyze_document( elif flow_type == "enhanced_integration": return self._analyze_enhanced(chunks, questions) else: - return AnalysisResult( - success=False, error=f"Analysis not supported for flow: {flow_type}" - ) + return AnalysisResult(success=False, error=f"Analysis not supported for flow: {flow_type}") except Exception as e: logger.error(f"Document analysis failed: {e}") return AnalysisResult(success=False, error=str(e)) - def complete_backend_analysis( - self, uploaded_file, question_set: str - ) -> AnalysisResult: + def complete_backend_analysis(self, uploaded_file, question_set: str) -> AnalysisResult: """ Complete backend analysis flow (Flow 4). @@ -145,16 +137,12 @@ def complete_backend_analysis( file_bytes = uploaded_file.read() import asyncio - resource_id = asyncio.run( - self.backend_service.upload_pdf(file_bytes, uploaded_file.name) - ) + resource_id = asyncio.run(self.backend_service.upload_pdf(file_bytes, uploaded_file.name)) st.success(f"✅ Document uploaded! Resource ID: {resource_id}") # Step 2: Wait for processing with st.spinner("Waiting for PDF processing..."): - processing_success = self.backend_service.wait_for_processing( - resource_id - ) + processing_success = self.backend_service.wait_for_processing(resource_id) if not processing_success: return AnalysisResult(success=False, error="PDF processing failed") @@ -163,16 +151,12 @@ def complete_backend_analysis( # Step 4: Submit analysis job with st.spinner("Submitting analysis job to backend..."): - analysis_job_id = self.backend_service.submit_analysis_job( - resource_id, question_set - ) + analysis_job_id = self.backend_service.submit_analysis_job(resource_id, question_set) st.success(f"Analysis job submitted! Job ID: {analysis_job_id}") # Step 5: Wait for analysis with st.spinner("Backend is running analysis..."): - analysis_results = self.backend_service.wait_for_analysis( - analysis_job_id - ) + analysis_results = self.backend_service.wait_for_analysis(analysis_job_id) st.success("Analysis completed and stored in backend database!") st.info("These results are now available to all authorized users") @@ -217,18 +201,14 @@ def _process_with_backend(self, uploaded_file) -> ProcessingResult: file_bytes = uploaded_file.read() import asyncio - resource_id = asyncio.run( - self.backend_service.upload_pdf(file_bytes, uploaded_file.name) - ) + resource_id = asyncio.run(self.backend_service.upload_pdf(file_bytes, uploaded_file.name)) st.success(f"✅ Document uploaded! Resource ID: {resource_id}") # Wait for processing with st.spinner("Waiting for backend processing..."): success = self.backend_service.wait_for_processing(resource_id) if not success: - return ProcessingResult( - success=False, error="Backend processing failed" - ) + return ProcessingResult(success=False, error="Backend processing failed") # Get chunks with st.spinner("Retrieving chunks..."): @@ -236,9 +216,7 @@ def _process_with_backend(self, uploaded_file) -> ProcessingResult: if not chunks: return ProcessingResult(success=False, error="No chunks retrieved") - return ProcessingResult( - success=True, chunks=chunks, resource_id=resource_id - ) + return ProcessingResult(success=True, chunks=chunks, resource_id=resource_id) except BackendServiceError as e: handle_backend_error(e, "Backend processing") @@ -251,9 +229,7 @@ def _process_complete_backend(self, uploaded_file) -> ProcessingResult: error="Complete backend analysis should use complete_backend_analysis() method", ) - def _analyze_local( - self, chunks: List[Dict[str, Any]], questions: List[str] - ) -> AnalysisResult: + def _analyze_local(self, chunks: List[Dict[str, Any]], questions: List[str]) -> AnalysisResult: """Analyze locally""" st.info("Using local analysis") @@ -275,17 +251,13 @@ def _analyze_local( }, ) - def _analyze_local_with_features( - self, chunks: List[Dict[str, Any]], questions: List[str] - ) -> AnalysisResult: + def _analyze_local_with_features(self, chunks: List[Dict[str, Any]], questions: List[str]) -> AnalysisResult: """Analyze locally with optional features""" # For now, same as local analysis # Could be enhanced with centralized LLM or data lake storage return self._analyze_local(chunks, questions) - def _analyze_enhanced( - self, chunks: List[Dict[str, Any]], questions: List[str] - ) -> AnalysisResult: + def _analyze_enhanced(self, chunks: List[Dict[str, Any]], questions: List[str]) -> AnalysisResult: """Enhanced analysis with centralized LLM and data lake""" # This would use NATS LLM and store in data lake # For now, fallback to local analysis @@ -298,23 +270,15 @@ def _configure_question_set(self, default_question_set: str) -> str: # Get dynamic question set options if QUESTION_LOADER_AVAILABLE: - question_set_options = question_loader.get_question_set_options() + [ - "custom" - ] + question_set_options = question_loader.get_question_set_options() + ["custom"] # Calculate index for default question set try: - index = ( - question_set_options.index(default_question_set) - if default_question_set in question_set_options - else 0 - ) + index = question_set_options.index(default_question_set) if default_question_set in question_set_options else 0 except ValueError: index = 0 else: # Fallback: use a generic approach without hardcoded names - question_set_options = [ - "custom" - ] # Only custom when question loader unavailable + question_set_options = ["custom"] # Only custom when question loader unavailable index = 0 question_set = st.selectbox( @@ -360,10 +324,11 @@ async def external_service_analysis( AnalysisResult with answers and top chunks """ try: - from .external_service_handler import ExternalServiceHandler - from .external_service_delivery import ExternalServiceDelivery from report_analyst.core.analyzer import DocumentAnalyzer + from .external_service_delivery import ExternalServiceDelivery + from .external_service_handler import ExternalServiceHandler + handler = ExternalServiceHandler() delivery = ExternalServiceDelivery() @@ -378,9 +343,7 @@ async def external_service_analysis( content_type="s3_url", s3_url=content, ) - processing_result = await handler.handle_external_notification( - service_id, notification - ) + processing_result = await handler.handle_external_notification(service_id, notification) else: # Pre-processed chunks from .external_service_handler import ExternalServiceReadyEvent @@ -391,20 +354,14 @@ async def external_service_analysis( content_type="chunks", chunks=content, ) - processing_result = await handler.handle_external_notification( - service_id, notification - ) + processing_result = await handler.handle_external_notification(service_id, notification) if not processing_result.success: - return AnalysisResult( - success=False, error=processing_result.error - ) + return AnalysisResult(success=False, error=processing_result.error) chunks = processing_result.chunks if not chunks: - return AnalysisResult( - success=False, error="No chunks available for analysis" - ) + return AnalysisResult(success=False, error="No chunks available for analysis") # Convert chunks to analyzer format analyzer_chunks = [ diff --git a/report_analyst_search_backend/s3_upload_service.py b/report_analyst_search_backend/s3_upload_service.py index 5aa104b25..b9034d91e 100644 --- a/report_analyst_search_backend/s3_upload_service.py +++ b/report_analyst_search_backend/s3_upload_service.py @@ -135,9 +135,7 @@ async def upload_pdf_via_s3_nats(self, file_bytes: bytes, filename: str) -> str: # Simple NATS publish - reliable delivery message_data = json.dumps(control_message).encode() - logger.info( - f"🔍 Publishing {len(message_data)} bytes to subject: {subject}" - ) + logger.info(f"🔍 Publishing {len(message_data)} bytes to subject: {subject}") try: # Use JetStream publish (CLI --jetstream works) @@ -245,7 +243,7 @@ def is_available() -> bool: # Since they're imported at module level, check the module reference # The test patches the module to None, so we check for that from report_analyst_search_backend import s3_upload_service - + if s3_upload_service.boto3 is None or s3_upload_service.nats is None: return False @@ -259,9 +257,7 @@ def is_available() -> bool: # Convenience function for backward compatibility -async def upload_pdf_via_s3_nats( - config: BackendConfig, file_bytes: bytes, filename: str -) -> str: +async def upload_pdf_via_s3_nats(config: BackendConfig, file_bytes: bytes, filename: str) -> str: """ Convenience function to upload PDF via S3+NATS. diff --git a/report_analyst_search_backend/service_discovery.py b/report_analyst_search_backend/service_discovery.py index 3c9daec3e..7a077a053 100644 --- a/report_analyst_search_backend/service_discovery.py +++ b/report_analyst_search_backend/service_discovery.py @@ -49,7 +49,7 @@ class ServiceValidator: def __init__(self, schema_dir: Optional[Path] = None): """ Initialize validator with schema directory. - + Args: schema_dir: Directory containing service contract schemas. Defaults to schemas/service-discovery/ in project root. @@ -57,7 +57,7 @@ def __init__(self, schema_dir: Optional[Path] = None): if schema_dir is None: # Find schema directory relative to this module (in enterprise module) schema_dir = Path(__file__).parent / "schemas" / "service-discovery" - + self.schema_dir = Path(schema_dir) self._contract_schema = None self._asyncapi_schema = None @@ -72,19 +72,19 @@ def _load_schemas(self): if contract_path.exists(): with open(contract_path) as f: self._contract_schema = json.load(f) - + # Load AsyncAPI schema (YAML) asyncapi_path = self.schema_dir / "asyncapi.yaml" if asyncapi_path.exists(): with open(asyncapi_path) as f: self._asyncapi_schema = yaml.safe_load(f) - + # Load OpenAPI schema (YAML) openapi_path = self.schema_dir / "openapi.yaml" if openapi_path.exists(): with open(openapi_path) as f: self._openapi_schema = yaml.safe_load(f) - + logger.info(f"Loaded service contract schemas from {self.schema_dir}") except Exception as e: logger.error(f"Failed to load schemas: {e}") @@ -93,10 +93,10 @@ def _load_schemas(self): def validate_service(self, service_manifest: Dict[str, Any]) -> ValidationResult: """ Validate a service manifest against the service contract schema. - + Args: service_manifest: Service manifest dictionary (must match service-contract.json schema) - + Returns: ValidationResult with validation status and any errors/warnings """ @@ -138,8 +138,7 @@ def _validate_semantics(self, manifest: Dict[str, Any]) -> ValidationResult: contract_version = manifest.get("contract_version", "1.0.0") if contract_version != "1.0.0": warnings.append( - f"Service uses contract version {contract_version}, " - f"validator expects 1.0.0. Compatibility not guaranteed." + f"Service uses contract version {contract_version}, " f"validator expects 1.0.0. Compatibility not guaranteed." ) # Validate NATS channels exist in AsyncAPI schema @@ -147,41 +146,34 @@ def _validate_semantics(self, manifest: Dict[str, Any]) -> ValidationResult: nats_channels = manifest.get("nats_channels", {}) published = [ch["channel"] for ch in nats_channels.get("publishes", [])] subscribed = [ch["channel"] for ch in nats_channels.get("subscribes", [])] - + all_channels = published + subscribed asyncapi_channels = self._asyncapi_schema.get("channels", {}) - + for channel in all_channels: if channel not in asyncapi_channels: - warnings.append( - f"NATS channel '{channel}' not defined in AsyncAPI schema. " - f"May be a custom extension." - ) + warnings.append(f"NATS channel '{channel}' not defined in AsyncAPI schema. " f"May be a custom extension.") # Validate HTTP endpoints exist in OpenAPI schema if self._openapi_schema and manifest.get("protocols", {}).get("http", {}).get("enabled"): http_endpoints = manifest.get("http_endpoints", {}) required_endpoints = http_endpoints.get("required", []) - + openapi_paths = self._openapi_schema.get("paths", {}) - + for endpoint in required_endpoints: path = endpoint.get("path") method = endpoint.get("method", "GET").lower() operation_id = endpoint.get("operation_id") - + # Check if path exists in OpenAPI if path not in openapi_paths: - errors.append( - f"Required endpoint {method.upper()} {path} not defined in OpenAPI schema" - ) + errors.append(f"Required endpoint {method.upper()} {path} not defined in OpenAPI schema") else: # Check if method exists for this path path_item = openapi_paths[path] if method not in path_item: - errors.append( - f"Method {method.upper()} not defined for endpoint {path} in OpenAPI schema" - ) + errors.append(f"Method {method.upper()} not defined for endpoint {path} in OpenAPI schema") else: # Check if operation_id matches operation = path_item[method] @@ -196,13 +188,13 @@ def _validate_semantics(self, manifest: Dict[str, Any]) -> ValidationResult: def get_required_channels(self) -> Dict[str, List[str]]: """ Get list of required NATS channels from AsyncAPI schema. - + Returns: Dictionary with 'publish' and 'subscribe' channel lists """ if not self._asyncapi_schema: return {"publish": [], "subscribe": []} - + channels = self._asyncapi_schema.get("channels", {}) return { "publish": list(channels.keys()), @@ -212,33 +204,35 @@ def get_required_channels(self) -> Dict[str, List[str]]: def get_required_endpoints(self) -> List[Dict[str, str]]: """ Get list of required HTTP endpoints from OpenAPI schema. - + Returns: List of endpoint dictionaries with 'method' and 'path' """ if not self._openapi_schema: return [] - + endpoints = [] paths = self._openapi_schema.get("paths", {}) - + for path, path_item in paths.items(): for method in ["get", "post", "put", "delete", "patch"]: if method in path_item: operation = path_item[method] - endpoints.append({ - "method": method.upper(), - "path": path, - "operation_id": operation.get("operationId", ""), - "summary": operation.get("summary", ""), - }) - + endpoints.append( + { + "method": method.upper(), + "path": path, + "operation_id": operation.get("operationId", ""), + "summary": operation.get("summary", ""), + } + ) + return endpoints def generate_service_template(self) -> Dict[str, Any]: """ Generate a template service manifest based on the contract schema. - + Returns: Template dictionary that can be filled in by service implementers """ @@ -291,18 +285,18 @@ def generate_service_template(self) -> Dict[str, Any]: def validate_service_from_file(manifest_path: Path) -> ValidationResult: """ Convenience function to validate a service manifest from a file. - + Args: manifest_path: Path to service manifest JSON file - + Returns: ValidationResult """ validator = ServiceValidator() - + with open(manifest_path) as f: manifest = json.load(f) - + return validator.validate_service(manifest) @@ -311,15 +305,14 @@ def validate_service_from_file(manifest_path: Path) -> ValidationResult: # Generate template validator = ServiceValidator() template = validator.generate_service_template() - + print("Service Contract Template:") print(json.dumps(template, indent=2)) - + print("\nRequired NATS Channels:") channels = validator.get_required_channels() print(json.dumps(channels, indent=2)) - + print("\nRequired HTTP Endpoints:") endpoints = validator.get_required_endpoints() print(json.dumps(endpoints, indent=2)) - diff --git a/report_analyst_search_backend/streamlit_integration.py b/report_analyst_search_backend/streamlit_integration.py index b8814e7f2..b11aa7afc 100644 --- a/report_analyst_search_backend/streamlit_integration.py +++ b/report_analyst_search_backend/streamlit_integration.py @@ -19,71 +19,51 @@ # Deprecated function - use new architecture def streamlit_enhanced_flow(uploaded_file, config): """Deprecated: Use FlowOrchestrator instead""" - logger.warning( - "streamlit_enhanced_flow is deprecated. Use FlowOrchestrator.process_document() instead." - ) + logger.warning("streamlit_enhanced_flow is deprecated. Use FlowOrchestrator.process_document() instead.") return None # Deprecated function - use new architecture def streamlit_full_backend_flow(uploaded_file, config): """Deprecated: Use FlowOrchestrator instead""" - logger.warning( - "streamlit_full_backend_flow is deprecated. Use FlowOrchestrator.complete_backend_analysis() instead." - ) + logger.warning("streamlit_full_backend_flow is deprecated. Use FlowOrchestrator.complete_backend_analysis() instead.") return None # Keep a few convenience functions for backward compatibility -def upload_pdf_to_backend( - file_bytes: bytes, filename: str, backend_url: str = "http://localhost:8000" -): +def upload_pdf_to_backend(file_bytes: bytes, filename: str, backend_url: str = "http://localhost:8000"): """Deprecated: Use BackendService.upload_pdf() instead""" - logger.warning( - "upload_pdf_to_backend is deprecated. Use BackendService.upload_pdf() instead." - ) + logger.warning("upload_pdf_to_backend is deprecated. Use BackendService.upload_pdf() instead.") return None -def wait_for_processing_polling( - resource_id: str, backend_url: str = "http://localhost:8000", timeout: int = 120 -): +def wait_for_processing_polling(resource_id: str, backend_url: str = "http://localhost:8000", timeout: int = 120): """Deprecated: Use BackendService.wait_for_processing() instead""" - logger.warning( - "wait_for_processing_polling is deprecated. Use BackendService.wait_for_processing() instead." - ) + logger.warning("wait_for_processing_polling is deprecated. Use BackendService.wait_for_processing() instead.") return False def get_backend_chunks(resource_id: str, backend_url: str = "http://localhost:8000"): """Deprecated: Use BackendService.get_chunks() instead""" - logger.warning( - "get_backend_chunks is deprecated. Use BackendService.get_chunks() instead." - ) + logger.warning("get_backend_chunks is deprecated. Use BackendService.get_chunks() instead.") return [] def streamlit_backend_flow(uploaded_file, backend_url: str = "http://localhost:8000"): """Deprecated: Use FlowOrchestrator.process_document() instead""" - logger.warning( - "streamlit_backend_flow is deprecated. Use FlowOrchestrator.process_document() instead." - ) + logger.warning("streamlit_backend_flow is deprecated. Use FlowOrchestrator.process_document() instead.") return None def use_centralized_llm_for_analysis(question: str, context_chunks, config): """Deprecated: Use FlowOrchestrator.analyze_document() instead""" - logger.warning( - "use_centralized_llm_for_analysis is deprecated. Use FlowOrchestrator.analyze_document() instead." - ) + logger.warning("use_centralized_llm_for_analysis is deprecated. Use FlowOrchestrator.analyze_document() instead.") return None def store_analysis_in_data_lake(analysis_results, config, experiment_id=None): """Deprecated: Use FlowOrchestrator.analyze_document() instead""" - logger.warning( - "store_analysis_in_data_lake is deprecated. Use FlowOrchestrator.analyze_document() instead." - ) + logger.warning("store_analysis_in_data_lake is deprecated. Use FlowOrchestrator.analyze_document() instead.") return False @@ -94,19 +74,13 @@ def submit_analysis_job_to_backend( backend_url: str = "http://localhost:8000", ): """Deprecated: Use BackendService.submit_analysis_job() instead""" - logger.warning( - "submit_analysis_job_to_backend is deprecated. Use BackendService.submit_analysis_job() instead." - ) + logger.warning("submit_analysis_job_to_backend is deprecated. Use BackendService.submit_analysis_job() instead.") return None -def wait_for_analysis_completion( - analysis_job_id: str, backend_url: str = "http://localhost:8000", timeout: int = 300 -): +def wait_for_analysis_completion(analysis_job_id: str, backend_url: str = "http://localhost:8000", timeout: int = 300): """Deprecated: Use BackendService.wait_for_analysis() instead""" - logger.warning( - "wait_for_analysis_completion is deprecated. Use BackendService.wait_for_analysis() instead." - ) + logger.warning("wait_for_analysis_completion is deprecated. Use BackendService.wait_for_analysis() instead.") return None @@ -116,7 +90,5 @@ def get_stored_analysis_results( backend_url: str = "http://localhost:8000", ): """Deprecated: Use BackendService.get_analysis_results() instead""" - logger.warning( - "get_stored_analysis_results is deprecated. Use BackendService.get_analysis_results() instead." - ) + logger.warning("get_stored_analysis_results is deprecated. Use BackendService.get_analysis_results() instead.") return None diff --git a/run_step_by_step_tests.py b/run_step_by_step_tests.py index 85492a9c2..c3fdc2198 100644 --- a/run_step_by_step_tests.py +++ b/run_step_by_step_tests.py @@ -58,15 +58,9 @@ def run_tests(verbose=False, specific_test=None): def main(): parser = argparse.ArgumentParser(description="Run step-by-step processing tests") - parser.add_argument( - "--verbose", "-v", action="store_true", help="Run tests with verbose output" - ) - parser.add_argument( - "--specific-test", "-k", help="Run only tests matching this pattern" - ) - parser.add_argument( - "--list-tests", "-l", action="store_true", help="List all available tests" - ) + parser.add_argument("--verbose", "-v", action="store_true", help="Run tests with verbose output") + parser.add_argument("--specific-test", "-k", help="Run only tests matching this pattern") + parser.add_argument("--list-tests", "-l", action="store_true", help="List all available tests") args = parser.parse_args() diff --git a/tests/conftest.py b/tests/conftest.py index 9a26f13dc..807ae4ee8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,7 +10,7 @@ import pytest import yaml -from report_analyst_jobs.event_router import EventRouter, IGNORE_ACTION +from report_analyst_jobs.event_router import IGNORE_ACTION, EventRouter @pytest.fixture @@ -66,22 +66,23 @@ def mock_nats_connection(): def event_router_with_mocks(event_router_yaml_file, mock_nats_connection): """Event router with mocked NATS connection""" mock_nc, mock_js = mock_nats_connection - + router = EventRouter.from_yaml(yaml_path=event_router_yaml_file) router.nc = mock_nc router.js = mock_js - + return router @pytest.fixture def mock_nats_message(): """Factory for creating mock NATS messages""" + def _create_message(subject: str, data: dict): mock_msg = AsyncMock() mock_msg.subject = subject mock_msg.data = json.dumps(data).encode() mock_msg.ack = AsyncMock() return mock_msg - return _create_message + return _create_message diff --git a/tests/integration/test_streamlit_app.py b/tests/integration/test_streamlit_app.py index 1ddff2658..7c0ffd5d1 100644 --- a/tests/integration/test_streamlit_app.py +++ b/tests/integration/test_streamlit_app.py @@ -27,13 +27,9 @@ @pytest.fixture def mock_streamlit(): """Mock main Streamlit functions""" - with patch("streamlit.set_page_config") as mock_config, patch( - "streamlit.title" - ) as mock_title, patch("streamlit.session_state", {}) as mock_state, patch( - "streamlit.selectbox" - ) as mock_select, patch( - "streamlit.expander" - ) as mock_expander, patch( + with patch("streamlit.set_page_config") as mock_config, patch("streamlit.title") as mock_title, patch( + "streamlit.session_state", {} + ) as mock_state, patch("streamlit.selectbox") as mock_select, patch("streamlit.expander") as mock_expander, patch( "streamlit.columns" ) as mock_columns: @@ -65,6 +61,7 @@ def test_env(): # Create test database using CacheManager (which will create all tables) db_path = cache_path / "analysis.db" from report_analyst.core.cache_manager import CacheManager + cache_manager = CacheManager(db_path=str(db_path)) # Tables are created automatically by CacheManager.init_db() @@ -184,13 +181,9 @@ async def mock_process_document(*args, **kwargs): for result in mock_results: yield result - with patch.object( - report_analyzer.analyzer, "process_document", new=mock_process_document - ): + with patch.object(report_analyzer.analyzer, "process_document", new=mock_process_document): results = [] - async for result in report_analyzer.analyze_document( - str(file_path), questions, selected_questions - ): + async for result in report_analyzer.analyze_document(str(file_path), questions, selected_questions): results.append(result) assert len(results) == 3 @@ -271,9 +264,7 @@ def test_check_step_completion(report_analyzer, test_env): assert status["chunks"] is True # Test embeddings complete - mock_get_chunks.return_value = [ - {"id": 1, "text": "chunk1", "embedding": [0.1, 0.2]} - ] + mock_get_chunks.return_value = [{"id": 1, "text": "chunk1", "embedding": [0.1, 0.2]}] status = report_analyzer.analyzer.check_step_completion(str(file_path)) assert status["embeddings"] is True diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py index c88e9fb87..0957834a1 100644 --- a/tests/test_analyzer.py +++ b/tests/test_analyzer.py @@ -49,6 +49,7 @@ def clean_db(test_db): print(f"\nCleaning database at: {test_db}") # Debug print # Use CacheManager to clean the database from report_analyst.core.cache_manager import CacheManager + cache_manager = CacheManager(db_path=str(test_db)) cache_manager.clear_cache() # Clear all cache return test_db @@ -230,10 +231,7 @@ def test_load_questions(analyzer): questions = analyzer._load_questions() assert len(questions) == 11 # TCFD has 11 questions assert "tcfd_1" in questions - assert ( - questions["tcfd_1"]["text"] - == "How does the company's board oversee climate-related risks and opportunities?" - ) + assert questions["tcfd_1"]["text"] == "How does the company's board oversee climate-related risks and opportunities?" assert "guidelines" in questions["tcfd_1"] @@ -241,10 +239,7 @@ def test_get_question_by_number(analyzer): """Test getting question by number""" question = analyzer.get_question_by_number(1) assert question is not None - assert ( - question["text"] - == "How does the company's board oversee climate-related risks and opportunities?" - ) + assert question["text"] == "How does the company's board oversee climate-related risks and opportunities?" assert "guidelines" in question @@ -305,9 +300,7 @@ async def test_process_document_with_cache(analyzer): } # Add to cache - analyzer.cache_manager.save_analysis( - file_path="test.pdf", question_id="tcfd_1", result=test_answer, config=config - ) + analyzer.cache_manager.save_analysis(file_path="test.pdf", question_id="tcfd_1", result=test_answer, config=config) # Process document results = [] @@ -317,9 +310,7 @@ async def test_process_document_with_cache(analyzer): assert result["status"] in ["processing", "complete", "cached"] # Verify we got the cached result - cached_result = analyzer.cache_manager.get_analysis( - file_path="test.pdf", config=config, question_ids=["tcfd_1"] - ) + cached_result = analyzer.cache_manager.get_analysis(file_path="test.pdf", config=config, question_ids=["tcfd_1"]) assert cached_result is not None assert cached_result["tcfd_1"]["result"]["ANSWER"] == test_answer["ANSWER"] @@ -338,7 +329,7 @@ def test_update_llm_model(analyzer): def test_cache_key_generation_with_none_llm(analyzer): """Test cache key generation when llm is None (no API keys available) - + This test would have failed before the fix that handles None llm gracefully. Before the fix, accessing self.llm.model would raise: AttributeError: 'NoneType' object has no attribute 'model' @@ -346,14 +337,14 @@ def test_cache_key_generation_with_none_llm(analyzer): # Simulate the case where no API keys are available (llm is None) analyzer.llm = None analyzer.default_model = "gpt-3.5-turbo-1106" # Set a default model - + # This should not crash - it should fall back to default_model analyzer.update_parameters(500, 20, 5) - + # Before the fix, this would raise: AttributeError: 'NoneType' object has no attribute 'model' # After the fix, it should work and use default_model key = analyzer._get_cache_key("test.pdf") - + # Verify the cache key contains expected components assert "cs500" in key assert "ov20" in key @@ -365,7 +356,7 @@ def test_cache_key_generation_with_none_llm(analyzer): def test_process_document_config_with_none_llm(analyzer, test_env): """Test that process_document creates config dict correctly when llm is None - + This test would have failed before the fix that handles None llm in config creation. """ # Simulate the case where no API keys are available (llm is None) @@ -373,36 +364,36 @@ def test_process_document_config_with_none_llm(analyzer, test_env): analyzer.default_model = "gpt-3.5-turbo-1106" analyzer.update_parameters(500, 20, 5) analyzer.question_set = "tcfd" - + # Create a minimal test PDF file test_file = test_env["storage_path"] / "test_report.pdf" test_file.write_bytes(b"%PDF-1.4\n%Test PDF") - + # Mock the document processing to avoid actual LLM calls # We just want to verify the config dict creation doesn't crash - with patch.object(analyzer, '_create_chunks', return_value=[]): + with patch.object(analyzer, "_create_chunks", return_value=[]): # This should not crash when creating the config dict # The actual processing will fail, but config creation should work try: # We'll catch the error after config is created results = [] + async def collect_results(): - async for result in analyzer.process_document( - str(test_file), [1], force_recompute=True - ): + async for result in analyzer.process_document(str(test_file), [1], force_recompute=True): results.append(result) # Stop after we see the first error or status if "error" in result or "status" in result: break - + import asyncio + asyncio.run(collect_results()) - + # The important thing is that we didn't crash with AttributeError # about 'NoneType' object has no attribute 'model' # If we got here, the config creation worked assert True # Test passes if no AttributeError was raised - + except AttributeError as e: if "'NoneType' object has no attribute 'model'" in str(e): pytest.fail("Config creation failed with None llm - this should be fixed!") @@ -414,7 +405,7 @@ async def test_process_document_pre_retrieved_chunks(analyzer, test_env): """Test processing document with pre-retrieved chunks""" test_file = test_env["storage_path"] / "test_report.pdf" test_file.write_bytes(b"%PDF-1.4\n%Test PDF") - + # Pre-retrieved chunks in backend format pre_chunks = [ { @@ -426,13 +417,13 @@ async def test_process_document_pre_retrieved_chunks(analyzer, test_env): "chunk_metadata": {"page": 2, "source": "backend"}, }, ] - + # Mock LLM to avoid actual API calls with patch.object(analyzer, "llm") as mock_llm: mock_response = Mock() mock_response.message.content = "Test answer" mock_llm.achat = AsyncMock(return_value=mock_response) - + results = [] async for result in analyzer.process_document( str(test_file), @@ -443,7 +434,7 @@ async def test_process_document_pre_retrieved_chunks(analyzer, test_env): results.append(result) if "error" in result or len(results) > 5: # Limit iterations break - + # Should use pre-retrieved chunks instead of creating new ones assert len(results) > 0 # Check that chunks were used (status message should indicate chunks loaded) @@ -460,18 +451,21 @@ async def test_process_document_s3_url_support(analyzer, test_env): s3_chunks = [ { "chunk_text": "Content downloaded from S3 bucket.", - "chunk_metadata": {"source": "s3", "url": "http://s3.example.com/bucket/file.pdf"}, + "chunk_metadata": { + "source": "s3", + "url": "http://s3.example.com/bucket/file.pdf", + }, }, ] - + # Use a temporary file path as identifier s3_file_path = "s3://bucket/file.pdf" - + with patch.object(analyzer, "llm") as mock_llm: mock_response = Mock() mock_response.message.content = "Test answer from S3 content" mock_llm.achat = AsyncMock(return_value=mock_response) - + results = [] async for result in analyzer.process_document( s3_file_path, @@ -482,7 +476,7 @@ async def test_process_document_s3_url_support(analyzer, test_env): results.append(result) if "error" in result or len(results) > 5: break - + # Should process S3 chunks successfully assert len(results) > 0 @@ -517,9 +511,7 @@ def test_get_all_cached_answers(analyzer): # Save answers to database for qid, answer in test_answers.items(): - analyzer.cache_manager.save_analysis( - file_path=f"test_{qid}.pdf", question_id=qid, result=answer, config=config - ) + analyzer.cache_manager.save_analysis(file_path=f"test_{qid}.pdf", question_id=qid, result=answer, config=config) # Get all cached answers and verify answers = analyzer.get_all_cached_answers("tcfd") @@ -559,9 +551,7 @@ async def test_document_analysis_workflow(test_env): # 2. Process document results = [] - async for result in analyzer.process_document( - str(test_env["test_file"]), ["tcfd_1", "tcfd_2"] - ): + async for result in analyzer.process_document(str(test_env["test_file"]), ["tcfd_1", "tcfd_2"]): results.append(result) # Handle both status and error results if "status" in result: diff --git a/tests/test_api_key_manager.py b/tests/test_api_key_manager.py index 8808293d9..e28da87ff 100644 --- a/tests/test_api_key_manager.py +++ b/tests/test_api_key_manager.py @@ -6,22 +6,23 @@ """ import os + from report_analyst.core.api_key_manager import APIKeyManager def test_set_and_get_api_key(): """Test setting and getting API keys""" session_state = {} - + # Set an API key APIKeyManager.set_api_key("OPENAI_API_KEY", "test-key-123", session_state) - + # Verify it's in session state assert session_state["api_key_openai_api_key"] == "test-key-123" - + # Verify it's in environment assert os.getenv("OPENAI_API_KEY") == "test-key-123" - + # Get the key back retrieved_key = APIKeyManager.get_api_key("OPENAI_API_KEY", session_state) assert retrieved_key == "test-key-123" @@ -30,14 +31,14 @@ def test_set_and_get_api_key(): def test_get_api_key_from_env(): """Test getting API key from environment when not in session state""" session_state = {} - + # Set in environment directly os.environ["OPENAI_API_KEY"] = "env-key-456" - + # Get should retrieve from environment retrieved_key = APIKeyManager.get_api_key("OPENAI_API_KEY", session_state) assert retrieved_key == "env-key-456" - + # Clean up del os.environ["OPENAI_API_KEY"] @@ -45,17 +46,17 @@ def test_get_api_key_from_env(): def test_session_state_takes_precedence(): """Test that session state value takes precedence over environment""" session_state = {} - + # Set in environment os.environ["OPENAI_API_KEY"] = "env-key-789" - + # Set in session state APIKeyManager.set_api_key("OPENAI_API_KEY", "session-key-789", session_state) - + # Get should return session state value retrieved_key = APIKeyManager.get_api_key("OPENAI_API_KEY", session_state) assert retrieved_key == "session-key-789" - + # Clean up del os.environ["OPENAI_API_KEY"] @@ -63,14 +64,14 @@ def test_session_state_takes_precedence(): def test_clear_api_key(): """Test clearing an API key""" session_state = {} - + # Set a key APIKeyManager.set_api_key("OPENAI_API_KEY", "test-key-clear", session_state) assert session_state["api_key_openai_api_key"] == "test-key-clear" - + # Clear it APIKeyManager.set_api_key("OPENAI_API_KEY", None, session_state) - + # Should be removed from session state assert "api_key_openai_api_key" not in session_state @@ -79,16 +80,16 @@ def test_sync_api_keys_to_env(): """Test syncing API keys from session state to environment""" session_state = { "api_key_openai_api_key": "synced-openai-key", - "api_key_google_api_key": "synced-google-key" + "api_key_google_api_key": "synced-google-key", } - + # Sync to environment APIKeyManager.sync_api_keys_to_env(session_state) - + # Verify they're in environment assert os.getenv("OPENAI_API_KEY") == "synced-openai-key" assert os.getenv("GOOGLE_API_KEY") == "synced-google-key" - + # Clean up del os.environ["OPENAI_API_KEY"] del os.environ["GOOGLE_API_KEY"] @@ -96,22 +97,20 @@ def test_sync_api_keys_to_env(): def test_sync_only_existing_keys(): """Test that sync only affects keys in session state""" - session_state = { - "api_key_openai_api_key": "only-openai-key" - } - + session_state = {"api_key_openai_api_key": "only-openai-key"} + # Set a Google key in environment os.environ["GOOGLE_API_KEY"] = "existing-google-key" - + # Sync APIKeyManager.sync_api_keys_to_env(session_state) - + # OpenAI should be synced assert os.getenv("OPENAI_API_KEY") == "only-openai-key" - + # Google should remain unchanged (not in session state) assert os.getenv("GOOGLE_API_KEY") == "existing-google-key" - + # Clean up del os.environ["OPENAI_API_KEY"] del os.environ["GOOGLE_API_KEY"] @@ -120,49 +119,51 @@ def test_sync_only_existing_keys(): def test_api_key_is_used_by_llm_provider(): """Test that API key set via APIKeyManager is actually used by LLM providers""" import report_analyst.core.llm_providers as llm_module - + session_state = {} test_openai_key = "test-openai-key-12345" test_google_key = "test-google-key-67890" - + # Set keys via APIKeyManager APIKeyManager.set_api_key("OPENAI_API_KEY", test_openai_key, session_state) APIKeyManager.set_api_key("GOOGLE_API_KEY", test_google_key, session_state) - + # Verify keys are in environment assert os.getenv("OPENAI_API_KEY") == test_openai_key assert os.getenv("GOOGLE_API_KEY") == test_google_key - + # Verify that get_llm reads from os.getenv (which now has our key) # We can't actually initialize the LLM without a real key, but we can verify # that the function will use the key we set original_getenv = llm_module.os.getenv - + # Track what keys are accessed accessed_keys = [] + def tracking_getenv(key, default=None): if key in ["OPENAI_API_KEY", "GOOGLE_API_KEY"]: accessed_keys.append(key) return original_getenv(key, default) - + # Temporarily replace os.getenv in the module llm_module.os.getenv = tracking_getenv - + try: # Try to get an OpenAI LLM (will fail without real key, but we can check it reads the env) try: from report_analyst.core.llm_providers import get_llm + get_llm("gpt-4o-mini") except (ValueError, Exception): # Expected to fail, but we check that it tried to read OPENAI_API_KEY pass - + # Verify it tried to read the API key from environment assert "OPENAI_API_KEY" in accessed_keys, "get_llm should read OPENAI_API_KEY from environment" finally: # Restore original llm_module.os.getenv = original_getenv - + # Clean up if "OPENAI_API_KEY" in os.environ: del os.environ["OPENAI_API_KEY"] @@ -173,20 +174,19 @@ def tracking_getenv(key, default=None): def test_api_key_from_session_state_overrides_env(): """Test that API key in session state overrides environment variable for LLM""" session_state = {} - + # Set key in environment os.environ["OPENAI_API_KEY"] = "env-key-111" - + # Set different key in session state APIKeyManager.set_api_key("OPENAI_API_KEY", "session-key-222", session_state) - + # Environment should now have session state value assert os.getenv("OPENAI_API_KEY") == "session-key-222" - + # Get should return session state value retrieved = APIKeyManager.get_api_key("OPENAI_API_KEY", session_state) assert retrieved == "session-key-222" - + # Clean up del os.environ["OPENAI_API_KEY"] - diff --git a/tests/test_backend_resource_roundtrip.py b/tests/test_backend_resource_roundtrip.py index fe4c8ea38..e5e10fe55 100644 --- a/tests/test_backend_resource_roundtrip.py +++ b/tests/test_backend_resource_roundtrip.py @@ -38,9 +38,7 @@ def temp_dir(): @pytest.fixture def backend_config(): """Create test backend configuration""" - return BackendConfig( - use_backend=True, backend_url="http://localhost:8000" - ) + return BackendConfig(use_backend=True, backend_url="http://localhost:8000") @pytest.fixture @@ -78,9 +76,7 @@ def mock_backend_chunks(): ] -def test_full_roundtrip_list_and_select_backend_resource( - backend_config, mock_backend_resources -): +def test_full_roundtrip_list_and_select_backend_resource(backend_config, mock_backend_resources): """Test Step 1 & 2: List backend resources and verify URN format""" with patch("requests.get") as mock_get: mock_response = Mock() @@ -94,9 +90,7 @@ def test_full_roundtrip_list_and_select_backend_resource( # Verify backend resource is listed assert len(resources) >= 1 - backend_resource = next( - (r for r in resources if r.is_backend_resource), None - ) + backend_resource = next((r for r in resources if r.is_backend_resource), None) assert backend_resource is not None # Step 2: Verify URN format @@ -111,9 +105,7 @@ def test_full_roundtrip_list_and_select_backend_resource( assert parsed["resource_id"] == "test-resource-1" -def test_full_roundtrip_retrieve_chunks( - backend_config, mock_backend_chunks -): +def test_full_roundtrip_retrieve_chunks(backend_config, mock_backend_chunks): """Test Step 3: Retrieve chunks from backend resource""" urn = "urn:report-analyst:backend:localhost:8000:test-resource-1" @@ -157,9 +149,7 @@ def test_full_roundtrip_retrieve_chunks( assert chunks[1]["chunk_text"] == mock_backend_chunks[1]["chunk_text"] -def test_full_roundtrip_analyzer_with_backend_chunks( - temp_dir, backend_config, mock_backend_chunks -): +def test_full_roundtrip_analyzer_with_backend_chunks(temp_dir, backend_config, mock_backend_chunks): """Test Step 4: Use backend chunks in analyzer""" # Create cache manager cache_path = temp_dir / "cache" @@ -176,6 +166,7 @@ def test_full_roundtrip_analyzer_with_backend_chunks( # Convert backend chunks to analyzer format with mock embeddings import numpy as np + backend_chunks = [ { "chunk_id": "chunk-1", @@ -238,6 +229,7 @@ def test_full_roundtrip_cache_compatibility(temp_dir, backend_config): # Create test chunks with mock embeddings (required by cache manager) import numpy as np + test_chunks = [ { "text": "Test chunk 1", @@ -288,9 +280,7 @@ def test_full_roundtrip_cache_compatibility(temp_dir, backend_config): assert urn_cached[0]["text"] == test_chunks[0]["text"] -def test_full_roundtrip_combined_local_and_backend( - temp_dir, backend_config, mock_backend_resources -): +def test_full_roundtrip_combined_local_and_backend(temp_dir, backend_config, mock_backend_resources): """Test combined listing of local and backend resources""" # Create local PDF test_pdf = temp_dir / "local_report.pdf" @@ -322,4 +312,3 @@ def test_full_roundtrip_combined_local_and_backend( # Verify sorting (most recent first) dates = [r.date for r in resources if r.date is not None] assert dates == sorted(dates, reverse=True) - diff --git a/tests/test_backend_service_list_reports.py b/tests/test_backend_service_list_reports.py index 20cf80769..826b09f0d 100644 --- a/tests/test_backend_service_list_reports.py +++ b/tests/test_backend_service_list_reports.py @@ -15,9 +15,7 @@ @pytest.fixture def backend_config(): """Create test backend configuration""" - return BackendConfig( - use_backend=True, backend_url="http://localhost:8000" - ) + return BackendConfig(use_backend=True, backend_url="http://localhost:8000") @pytest.fixture @@ -62,7 +60,7 @@ def test_backend_service_list_reports(backend_config, mock_backend_resources): def test_backend_service_list_reports_with_https(backend_config, mock_backend_resources): """Test URN generation with HTTPS backend URL""" backend_config.backend_url = "https://api.example.com" - + with patch("requests.get") as mock_get: mock_response = Mock() mock_response.json.return_value = mock_backend_resources @@ -81,7 +79,7 @@ def test_backend_service_list_reports_with_https(backend_config, mock_backend_re def test_backend_service_list_reports_with_port(backend_config, mock_backend_resources): """Test URN generation with port in backend URL""" backend_config.backend_url = "http://localhost:8080" - + with patch("requests.get") as mock_get: mock_response = Mock() mock_response.json.return_value = mock_backend_resources @@ -125,15 +123,13 @@ def test_backend_service_list_reports_empty_response(backend_config): def test_backend_service_normalize_backend_url(): """Test URL normalization for URN""" - backend_config = BackendConfig( - use_backend=True, backend_url="https://api.example.com:443" - ) + backend_config = BackendConfig(use_backend=True, backend_url="https://api.example.com:443") service = BackendService(backend_config) - + normalized = service._normalize_backend_url("https://api.example.com:443") assert normalized == "api.example.com:443" assert "https://" not in normalized - + normalized = service._normalize_backend_url("http://localhost:8000") assert normalized == "localhost:8000" assert "http://" not in normalized @@ -141,24 +137,22 @@ def test_backend_service_normalize_backend_url(): def test_backend_service_parse_date(): """Test date parsing for timestamps""" - backend_config = BackendConfig( - use_backend=True, backend_url="http://localhost:8000" - ) + backend_config = BackendConfig(use_backend=True, backend_url="http://localhost:8000") service = BackendService(backend_config) - + # Test ISO format with Z timestamp = service._parse_date("2024-01-01T00:00:00Z") assert timestamp is not None assert isinstance(timestamp, float) - + # Test ISO format without Z timestamp = service._parse_date("2024-01-01T00:00:00") assert timestamp is not None - + # Test None timestamp = service._parse_date(None) assert timestamp is None - + # Test invalid format timestamp = service._parse_date("invalid-date") assert timestamp is None @@ -177,4 +171,3 @@ def test_backend_service_get_resources_public(backend_config, mock_backend_resou assert len(resources) == 2 assert resources[0]["id"] == "test-resource-1" - diff --git a/tests/test_cache_manager.py b/tests/test_cache_manager.py index b87be7d79..1a0baa97a 100644 --- a/tests/test_cache_manager.py +++ b/tests/test_cache_manager.py @@ -20,6 +20,7 @@ def setup_test_env(): # Cleanup after tests if Path(os.environ["STORAGE_PATH"]).exists(): import shutil + shutil.rmtree(os.environ["STORAGE_PATH"]) # Restore original if original_storage: @@ -45,7 +46,7 @@ def temp_db_both(request): For PostgreSQL, requires DATABASE_URL environment variable or skips test. """ temp_dir = tempfile.mkdtemp() - + if request.param == "sqlite": db_path = Path(temp_dir) / "test_cache.db" cache_manager = CacheManager(db_path=str(db_path)) @@ -57,7 +58,7 @@ def temp_db_both(request): pytest.skip("TEST_POSTGRES_URL not set, skipping PostgreSQL test") cache_manager = CacheManager(database_url=database_url) yield cache_manager - + shutil.rmtree(temp_dir) @@ -245,14 +246,16 @@ def test_get_chunks_without_embeddings(temp_db): # Insert chunks directly into database (some without embeddings, some with) with temp_db.db_manager.get_connection() as conn: timestamp = datetime.now().isoformat() - + # Insert chunks without embeddings conn.execute( - text(""" + text( + """ INSERT INTO document_chunks (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) - """), + """ + ), { "file_path": file_path, "chunk_text": "Chunk 1 without embedding", @@ -264,11 +267,13 @@ def test_get_chunks_without_embeddings(temp_db): }, ) conn.execute( - text(""" + text( + """ INSERT INTO document_chunks (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) - """), + """ + ), { "file_path": file_path, "chunk_text": "Chunk 2 without embedding", @@ -279,15 +284,17 @@ def test_get_chunks_without_embeddings(temp_db): "created_at": timestamp, }, ) - + # Insert chunk with embedding embedding_bytes = np.array([0.1, 0.2, 0.3], dtype=np.float32).tobytes() conn.execute( - text(""" + text( + """ INSERT INTO document_chunks (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) - """), + """ + ), { "file_path": file_path, "chunk_text": "Chunk 3 with embedding", @@ -300,9 +307,7 @@ def test_get_chunks_without_embeddings(temp_db): ) # Get chunks without embeddings - result = temp_db.get_chunks_without_embeddings( - file_path, chunk_size, chunk_overlap - ) + result = temp_db.get_chunks_without_embeddings(file_path, chunk_size, chunk_overlap) # Should return only chunks without embeddings assert len(result) == 2 @@ -332,11 +337,13 @@ def test_has_chunk_scoring(temp_db): with temp_db.db_manager.get_connection() as conn: timestamp = datetime.now().isoformat() conn.execute( - text(""" + text( + """ INSERT INTO document_chunks (file_path, chunk_text, chunk_size, chunk_overlap, embedding, metadata, created_at) VALUES (:file_path, :chunk_text, :chunk_size, :chunk_overlap, :embedding, :metadata, :created_at) - """), + """ + ), { "file_path": file_path, "chunk_text": chunk_text, diff --git a/tests/test_dataframe_manager.py b/tests/test_dataframe_manager.py index 8a0344181..9656c75c3 100644 --- a/tests/test_dataframe_manager.py +++ b/tests/test_dataframe_manager.py @@ -1,6 +1,10 @@ -import pytest import pandas as pd -from report_analyst.core.dataframe_manager import create_analysis_dataframes, format_list_field +import pytest + +from report_analyst.core.dataframe_manager import ( + create_analysis_dataframes, + format_list_field, +) def test_format_list_field_with_evidence(): @@ -90,10 +94,7 @@ def test_create_analysis_dataframes_with_evidence(): # Check that analysis dataframe was created correctly assert len(analysis_df) == 1 assert analysis_df.iloc[0]["Question ID"] == "ev_24" - assert ( - analysis_df.iloc[0]["Analysis"] - == "No data reported for Scope 1 CO₂ emissions for the year 2022." - ) + assert analysis_df.iloc[0]["Analysis"] == "No data reported for Scope 1 CO₂ emissions for the year 2022." assert analysis_df.iloc[0]["Score"] == 2 # Most importantly, check that Key Evidence is properly formatted @@ -171,4 +172,3 @@ def test_create_analysis_dataframes_with_chunks(): assert chunks_df.iloc[1]["Is Evidence"] == False # Use == instead of is for pandas bool assert pd.isna(chunks_df.iloc[0]["LLM Score"]) # None should become NaN assert chunks_df.iloc[1]["LLM Score"] == 0.8 - diff --git a/tests/test_document_ready_auto_index.py b/tests/test_document_ready_auto_index.py index 86804591d..2eda86366 100644 --- a/tests/test_document_ready_auto_index.py +++ b/tests/test_document_ready_auto_index.py @@ -70,9 +70,7 @@ async def test_process_document_ready_event(mock_chunks, mock_analysis_result): # Mock search backend client coordinator.search_backend = Mock() coordinator.search_backend.base_url = "http://localhost:8000" - coordinator.search_backend.get_resource_chunks = AsyncMock( - return_value=mock_chunks - ) + coordinator.search_backend.get_resource_chunks = AsyncMock(return_value=mock_chunks) # Mock _get_chunks_for_resource (which calls get_resource_chunks) coordinator._get_chunks_for_resource = AsyncMock(return_value=mock_chunks) @@ -113,19 +111,13 @@ async def test_process_document_ready_event(mock_chunks, mock_analysis_result): await coordinator._handle_document_ready(mock_msg, config) # Verify chunks were retrieved - coordinator._get_chunks_for_resource.assert_called_once_with( - "test-resource-123", "search", None - ) + coordinator._get_chunks_for_resource.assert_called_once_with("test-resource-123", "search", None) # Verify analysis was run - coordinator._run_analysis.assert_called_once_with( - mock_chunks, "tcfd", {"model": "gpt-4o-mini"} - ) + coordinator._run_analysis.assert_called_once_with(mock_chunks, "tcfd", {"model": "gpt-4o-mini"}) # Verify results were stored - coordinator._store_analysis_to_backend.assert_called_once_with( - "test-resource-123", mock_analysis_result, "tcfd" - ) + coordinator._store_analysis_to_backend.assert_called_once_with("test-resource-123", mock_analysis_result, "tcfd") # Verify message was acked mock_msg.ack.assert_called_once() @@ -190,14 +182,10 @@ async def test_process_document_ready_with_provided_chunks(mock_analysis_result) coordinator._get_chunks_for_resource.assert_not_called() # Verify analysis was run with provided chunks - coordinator._run_analysis.assert_called_once_with( - provided_chunks, "tcfd", {"model": "gpt-4o-mini"} - ) + coordinator._run_analysis.assert_called_once_with(provided_chunks, "tcfd", {"model": "gpt-4o-mini"}) # Verify results were stored - coordinator._store_analysis_to_backend.assert_called_once_with( - "test-resource-123", mock_analysis_result, "tcfd" - ) + coordinator._store_analysis_to_backend.assert_called_once_with("test-resource-123", mock_analysis_result, "tcfd") # Verify message was acked mock_msg.ack.assert_called_once() @@ -258,9 +246,7 @@ async def test_store_analysis_to_backend(mock_analysis_result): coordinator.search_backend = Mock() coordinator.search_backend.base_url = "http://localhost:8000" - with patch( - "report_analyst_search_backend.backend_service.BackendService.store_analysis_results" - ) as mock_store: + with patch("report_analyst_search_backend.backend_service.BackendService.store_analysis_results") as mock_store: mock_store.return_value = "stored-result-id-123" await coordinator._store_analysis_to_backend( @@ -306,4 +292,3 @@ async def test_config_options(): assert config_dict["pull_chunks"] is False assert config_dict["chunk_retrieval_method"] == "direct" assert config_dict["max_chunks"] == 100 - diff --git a/tests/test_document_ready_e2e_router.py b/tests/test_document_ready_e2e_router.py index ae2b5a640..7efabfb9c 100644 --- a/tests/test_document_ready_e2e_router.py +++ b/tests/test_document_ready_e2e_router.py @@ -58,14 +58,14 @@ def mock_analysis_result(): def document_ready_handler(mock_chunks, mock_analysis_result): """Create a handler for document.ready events that uses NATSJobCoordinator""" coordinator = NATSJobCoordinator() - + # Mock search backend client coordinator.search_backend = Mock() coordinator.search_backend.base_url = "http://localhost:8000" coordinator._get_chunks_for_resource = AsyncMock(return_value=mock_chunks) coordinator._run_analysis = AsyncMock(return_value=mock_analysis_result) coordinator._store_analysis_to_backend = AsyncMock() - + # Create config config = DocumentReadyProcessingConfig( pull_chunks=True, @@ -73,12 +73,12 @@ def document_ready_handler(mock_chunks, mock_analysis_result): analysis_config={"model": "gpt-4o-mini"}, store_to_backend=True, ) - + async def handler(ctx: EventContext): """Handler that processes document.ready events""" event = DocumentReadyEvent(**ctx.data) await coordinator._handle_document_ready(ctx.message, config) - + return handler, coordinator @@ -92,18 +92,18 @@ async def test_e2e_document_ready_flow_with_router( ): """End-to-end test: document.ready event → router → handler → full processing""" handler, coordinator = document_ready_handler - + # Create router with handler router = EventRouter.from_yaml( yaml_path=event_router_yaml_file, handler_registry={"handle_document_ready": handler}, ) - + # Mock NATS connection mock_nc, mock_js = mock_nats_connection router.nc = mock_nc router.js = mock_js - + # Create document.ready event message event_data = { "resource_id": "test-resource-123", @@ -115,26 +115,20 @@ async def test_e2e_document_ready_flow_with_router( mock_msg.subject = "document.ready" mock_msg.data = json.dumps(event_data).encode() mock_msg.ack = AsyncMock() - + # Process through router await router._handle_message(mock_msg) - + # Verify full flow executed: # 1. Chunks were retrieved - coordinator._get_chunks_for_resource.assert_called_once_with( - "test-resource-123", "search", None - ) - + coordinator._get_chunks_for_resource.assert_called_once_with("test-resource-123", "search", None) + # 2. Analysis was run - coordinator._run_analysis.assert_called_once_with( - mock_chunks, "tcfd", {"model": "gpt-4o-mini"} - ) - + coordinator._run_analysis.assert_called_once_with(mock_chunks, "tcfd", {"model": "gpt-4o-mini"}) + # 3. Results were stored - coordinator._store_analysis_to_backend.assert_called_once_with( - "test-resource-123", mock_analysis_result, "tcfd" - ) - + coordinator._store_analysis_to_backend.assert_called_once_with("test-resource-123", mock_analysis_result, "tcfd") + # 4. Message was acked mock_msg.ack.assert_called_once() @@ -150,24 +144,24 @@ async def test_e2e_document_ready_ignore_with_router( yaml_path=event_router_yaml_file, handler_registry={}, ) - + # Add ignore rule for document.upload router.add_rule("document.upload", "ignore", priority=5) - + # Mock NATS mock_nc, mock_js = mock_nats_connection router.nc = mock_nc router.js = mock_js - + # Create document.upload event (should be ignored) mock_msg = AsyncMock() mock_msg.subject = "document.upload" mock_msg.data = json.dumps({"resource_id": "test-456"}).encode() mock_msg.ack = AsyncMock() - + # Process through router await router._handle_message(mock_msg) - + # Should be acked but not processed mock_msg.ack.assert_called_once() # No handlers should have been called (we didn't register any for upload) @@ -184,27 +178,27 @@ async def test_e2e_document_ready_with_provided_chunks( coordinator._get_chunks_for_resource = AsyncMock() # Should not be called coordinator._run_analysis = AsyncMock(return_value=mock_analysis_result) coordinator._store_analysis_to_backend = AsyncMock() - + config = DocumentReadyProcessingConfig( pull_chunks=False, # Don't pull, use provided chunks question_set="tcfd", analysis_config={"model": "gpt-4o-mini"}, store_to_backend=True, ) - + async def handler(ctx: EventContext): event = DocumentReadyEvent(**ctx.data) await coordinator._handle_document_ready(ctx.message, config) - + router = EventRouter.from_yaml( yaml_path=event_router_yaml_file, handler_registry={"handle_document_ready": handler}, ) - + mock_nc, mock_js = mock_nats_connection router.nc = mock_nc router.js = mock_js - + # Event with chunks included provided_chunks = [ {"id": "chunk-1", "text": "Test chunk 1", "metadata": {}}, @@ -217,23 +211,20 @@ async def handler(ctx: EventContext): "status": "ready", "chunks": provided_chunks, } - + mock_msg = AsyncMock() mock_msg.subject = "document.ready" mock_msg.data = json.dumps(event_data).encode() mock_msg.ack = AsyncMock() - + await router._handle_message(mock_msg) - + # Verify chunks were NOT pulled coordinator._get_chunks_for_resource.assert_not_called() - + # Verify analysis was run with provided chunks - coordinator._run_analysis.assert_called_once_with( - provided_chunks, "tcfd", {"model": "gpt-4o-mini"} - ) - + coordinator._run_analysis.assert_called_once_with(provided_chunks, "tcfd", {"model": "gpt-4o-mini"}) + # Verify results stored coordinator._store_analysis_to_backend.assert_called_once() mock_msg.ack.assert_called_once() - diff --git a/tests/test_event_router.py b/tests/test_event_router.py index 12a08ad5d..c08cdadea 100644 --- a/tests/test_event_router.py +++ b/tests/test_event_router.py @@ -10,10 +10,10 @@ import pytest from report_analyst_jobs.event_router import ( + IGNORE_ACTION, EventActionRule, EventContext, EventRouter, - IGNORE_ACTION, ) @@ -35,6 +35,7 @@ def router(): @pytest.mark.asyncio async def test_add_rule(router): """Test adding rules""" + async def handler(ctx: EventContext): pass @@ -68,6 +69,7 @@ async def test_match_subject(router): @pytest.mark.asyncio async def test_find_rule(router): """Test finding matching rule""" + async def handler(ctx: EventContext): pass @@ -155,6 +157,7 @@ async def test_handle_message_no_rule(router, mock_nats): @pytest.mark.asyncio async def test_priority_ordering(router): """Test that rules are checked in priority order""" + async def handler1(ctx: EventContext): pass @@ -182,6 +185,7 @@ async def test_disabled_rule(router): @pytest.mark.asyncio async def test_routing_table(router): """Test getting routing table""" + async def handler(ctx: EventContext): pass @@ -199,6 +203,7 @@ async def handler(ctx: EventContext): @pytest.mark.asyncio async def test_set_rules(router): """Test setting all rules at once""" + async def handler(ctx: EventContext): pass @@ -212,4 +217,3 @@ async def handler(ctx: EventContext): assert len(router.get_rules()) == 2 # Should be sorted by priority assert router.get_rules()[0].priority == 10 - diff --git a/tests/test_event_router_e2e.py b/tests/test_event_router_e2e.py index b841c3bbe..2daed093b 100644 --- a/tests/test_event_router_e2e.py +++ b/tests/test_event_router_e2e.py @@ -18,7 +18,7 @@ import pytest import yaml -from report_analyst_jobs.event_router import EventContext, EventRouter, IGNORE_ACTION +from report_analyst_jobs.event_router import IGNORE_ACTION, EventContext, EventRouter @pytest.fixture @@ -333,4 +333,3 @@ async def test_e2e_routing_table_inspection(temp_yaml_file): document_upload_rule = next(r for r in table if r["pattern"] == "document.upload") assert document_upload_rule["action"] == IGNORE_ACTION assert document_upload_rule["enabled"] is True - diff --git a/tests/test_external_service_integration.py b/tests/test_external_service_integration.py index 2712775dd..155c1c3c6 100644 --- a/tests/test_external_service_integration.py +++ b/tests/test_external_service_integration.py @@ -14,8 +14,10 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest + try: from aioresponses import aioresponses + HAS_AIORESPONSES = True except ImportError: HAS_AIORESPONSES = False @@ -42,9 +44,7 @@ def external_handler(): @pytest.fixture def external_client(): """Create external service client""" - return ExternalServiceClient( - base_url="http://localhost:8000", nats_url="nats://localhost:4222" - ) + return ExternalServiceClient(base_url="http://localhost:8000", nats_url="nats://localhost:4222") @pytest.fixture @@ -111,9 +111,7 @@ async def test_handle_notification_s3_url(self, external_handler): external_handler.s3_client = mock_s3_client # Mock PyMuPDFReader - with patch( - "llama_index.readers.file.PyMuPDFReader" - ) as mock_reader: + with patch("llama_index.readers.file.PyMuPDFReader") as mock_reader: mock_doc = Mock() mock_doc.text = "Test PDF content" mock_doc.metadata = {} @@ -121,9 +119,7 @@ async def test_handle_notification_s3_url(self, external_handler): mock_reader_instance.load.return_value = [mock_doc] mock_reader.return_value = mock_reader_instance - result = await external_handler.handle_external_notification( - "service-x", notification - ) + result = await external_handler.handle_external_notification("service-x", notification) assert isinstance(result, ProcessingResult) # Should process S3 URL (may fail if S3 client not initialized, but structure is correct) @@ -138,9 +134,7 @@ async def test_handle_notification_chunks(self, external_handler, sample_chunks) chunks=sample_chunks, ) - result = await external_handler.handle_external_notification( - "service-x", notification, rechunk_mode="never" - ) + result = await external_handler.handle_external_notification("service-x", notification, rechunk_mode="never") assert result.success assert result.chunks is not None @@ -157,9 +151,7 @@ async def test_handle_notification_pages(self, external_handler, sample_pages): pages=sample_pages, ) - result = await external_handler.handle_external_notification( - "service-x", notification - ) + result = await external_handler.handle_external_notification("service-x", notification) assert result.success assert result.chunks is not None @@ -176,9 +168,7 @@ async def test_rechunk_mode_never(self, external_handler, sample_chunks): chunks=sample_chunks, ) - result = await external_handler.handle_external_notification( - "service-x", notification, rechunk_mode="never" - ) + result = await external_handler.handle_external_notification("service-x", notification, rechunk_mode="never") assert result.success assert len(result.chunks) == 2 @@ -194,9 +184,7 @@ async def test_rechunk_mode_auto(self, external_handler, sample_chunks): chunks=sample_chunks, ) - result = await external_handler.handle_external_notification( - "service-x", notification, rechunk_mode="auto" - ) + result = await external_handler.handle_external_notification("service-x", notification, rechunk_mode="auto") assert result.success assert result.chunks is not None @@ -338,7 +326,10 @@ async def test_deliver_results_poll(self, external_delivery): service_id="service-x", request_id="req-123", external_request_id="ext-req-123", - results={"answers": [{"question_id": "q1", "answer": "test"}], "top_chunks": []}, + results={ + "answers": [{"question_id": "q1", "answer": "test"}], + "top_chunks": [], + }, response_method="poll", ) @@ -384,9 +375,7 @@ async def test_full_flow_chunks_http(self, external_handler, sample_chunks): chunks=sample_chunks, ) - result = await external_handler.handle_external_notification( - "service-x", notification, rechunk_mode="never" - ) + result = await external_handler.handle_external_notification("service-x", notification, rechunk_mode="never") assert result.success assert result.chunks is not None @@ -419,4 +408,3 @@ async def test_external_service_error_handling(self, external_handler): assert not result2.success assert "S3 URL not provided" in result2.error - diff --git a/tests/test_file_storage.py b/tests/test_file_storage.py index e4e99ebf6..5cf62ece3 100644 --- a/tests/test_file_storage.py +++ b/tests/test_file_storage.py @@ -5,60 +5,66 @@ """ import os + import pytest -from report_analyst.core.file_storage import PostgreSQLFileStorage, FileStorageError, get_file_storage + +from report_analyst.core.file_storage import ( + FileStorageError, + PostgreSQLFileStorage, + get_file_storage, +) @pytest.mark.skipif( not os.getenv("DATABASE_URL") or not os.getenv("DATABASE_URL").startswith(("postgresql://", "postgres://")), - reason="PostgreSQL not configured" + reason="PostgreSQL not configured", ) def test_postgres_file_storage_store_and_retrieve(): """Test storing and retrieving a file from PostgreSQL""" database_url = os.getenv("DATABASE_URL") storage = PostgreSQLFileStorage(database_url) - + # Test file content test_content = b"Test file content for PostgreSQL storage" filename = "test_file.pdf" - + # Store file file_id = storage.store_file(test_content, filename, "application/pdf") assert file_id is not None assert len(file_id) == 36 # UUID length - + # Retrieve file retrieved_content = storage.retrieve_file(file_id) assert retrieved_content == test_content - + # Get file info file_info = storage.get_file_info(file_id) assert file_info is not None assert file_info["filename"] == filename assert file_info["content_type"] == "application/pdf" assert file_info["file_size"] == len(test_content) - + # Clean up storage.delete_file(file_id) @pytest.mark.skipif( not os.getenv("DATABASE_URL") or not os.getenv("DATABASE_URL").startswith(("postgresql://", "postgres://")), - reason="PostgreSQL not configured" + reason="PostgreSQL not configured", ) def test_postgres_file_storage_delete(): """Test deleting a file from PostgreSQL""" database_url = os.getenv("DATABASE_URL") storage = PostgreSQLFileStorage(database_url) - + # Store file test_content = b"Test file to delete" file_id = storage.store_file(test_content, "delete_test.pdf") - + # Delete file deleted = storage.delete_file(file_id) assert deleted is True - + # Verify file is gone retrieved = storage.retrieve_file(file_id) assert retrieved is None @@ -70,12 +76,12 @@ def test_get_file_storage_without_postgres(): original_url = os.environ.get("DATABASE_URL") if "DATABASE_URL" in os.environ: del os.environ["DATABASE_URL"] - + # Unset USE_POSTGRES_FILE_STORAGE original_setting = os.environ.get("USE_POSTGRES_FILE_STORAGE") if "USE_POSTGRES_FILE_STORAGE" in os.environ: del os.environ["USE_POSTGRES_FILE_STORAGE"] - + try: storage = get_file_storage() assert storage is None @@ -91,7 +97,6 @@ def test_postgres_file_storage_requires_postgres(): """Test that PostgreSQLFileStorage raises error for SQLite""" # Use SQLite URL sqlite_url = "sqlite:///test.db" - + with pytest.raises(FileStorageError, match="PostgreSQL"): PostgreSQLFileStorage(sqlite_url) - diff --git a/tests/test_llm_evidence_separation.py b/tests/test_llm_evidence_separation.py index fc09f3cc9..16ebf97ae 100644 --- a/tests/test_llm_evidence_separation.py +++ b/tests/test_llm_evidence_separation.py @@ -15,12 +15,12 @@ import json import os -from sqlalchemy import inspect, text import sys import tempfile from pathlib import Path import pytest +from sqlalchemy import inspect, text # Add the app directory to the path sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) @@ -56,6 +56,7 @@ def test_database_schema_separation(self, cache_manager): """Test that the database schema has separate fields for similarity_score, llm_score, and is_evidence.""" from sqlalchemy import inspect + inspector = inspect(cache_manager.db_manager.get_engine()) columns = {col["name"]: str(col["type"]) for col in inspector.get_columns("chunk_relevance")} @@ -71,8 +72,7 @@ def test_database_schema_separation(self, cache_manager): assert col in columns, f"Missing column: {col}" col_type = columns[col].upper() assert any( - expected_type.upper() in col_type - for expected_type in expected_types + expected_type.upper() in col_type for expected_type in expected_types ), f"Column {col} has type {columns[col]}, expected one of {expected_types}" def test_independent_value_storage(self, cache_manager): @@ -142,24 +142,16 @@ def test_independent_value_storage(self, cache_manager): # Retrieve and verify retrieved = cache_manager.get_analysis(file_path, config, [question_id]) - assert ( - question_id in retrieved - ), f"Failed to retrieve analysis for {question_id}" + assert question_id in retrieved, f"Failed to retrieve analysis for {question_id}" retrieved_chunks = retrieved[question_id]["chunks"] # Verify each chunk maintains independent values for i, chunk in enumerate(retrieved_chunks): original = test_chunks[i] - assert ( - chunk["similarity_score"] == original["similarity_score"] - ), f"Similarity score mismatch for chunk {i+1}" - assert ( - chunk["llm_score"] == original["llm_score"] - ), f"LLM score mismatch for chunk {i+1}" - assert ( - chunk["is_evidence"] == original["is_evidence"] - ), f"Evidence flag mismatch for chunk {i+1}" + assert chunk["similarity_score"] == original["similarity_score"], f"Similarity score mismatch for chunk {i+1}" + assert chunk["llm_score"] == original["llm_score"], f"LLM score mismatch for chunk {i+1}" + assert chunk["is_evidence"] == original["is_evidence"], f"Evidence flag mismatch for chunk {i+1}" def test_independence_principle_examples(self): """Test that demonstrates the independence principle with clear examples.""" @@ -258,9 +250,7 @@ def test_workflow_separation(self): assert chunks_after_llm_scoring[0]["llm_score"] == 0.7 # Now set assert chunks_after_llm_scoring[0]["is_evidence"] is False # Still unchanged - assert ( - chunks_after_evidence_determination[0]["similarity_score"] == 0.8 - ) # Unchanged + assert chunks_after_evidence_determination[0]["similarity_score"] == 0.8 # Unchanged assert chunks_after_evidence_determination[0]["llm_score"] == 0.7 # Unchanged assert chunks_after_evidence_determination[0]["is_evidence"] is True # Now set @@ -296,9 +286,7 @@ def test_workflow_separation(self): test_suite.test_workflow_separation() print("✅ Workflow properly separates concerns") - print( - "\n🎉 All tests passed! LLM score and evidence determination are properly separated." - ) + print("\n🎉 All tests passed! LLM score and evidence determination are properly separated.") finally: try: diff --git a/tests/test_question_loader.py b/tests/test_question_loader.py index 908a11d46..cb63b0645 100644 --- a/tests/test_question_loader.py +++ b/tests/test_question_loader.py @@ -186,9 +186,7 @@ def test_load_real_question_sets(self): # Should have at least the main question sets expected_sets = ["everest", "tcfd", "denali", "kilimanjaro"] for expected_set in expected_sets: - assert ( - expected_set in question_sets - ), f"Expected {expected_set} in question sets" + assert expected_set in question_sets, f"Expected {expected_set} in question sets" qset = question_sets[expected_set] assert qset.name is not None @@ -303,14 +301,10 @@ def test_fallback_behavior_without_core_functionality(self): from report_analyst.core.question_loader import get_question_loader question_loader = get_question_loader() - question_set_options = question_loader.get_question_set_options() + [ - "custom" - ] + question_set_options = question_loader.get_question_set_options() + ["custom"] else: # Fallback: use a generic approach without hardcoded names - question_set_options = [ - "custom" - ] # Only custom when core functionality unavailable + question_set_options = ["custom"] # Only custom when core functionality unavailable assert question_set_options == ["custom"] assert len(question_set_options) == 1 @@ -325,14 +319,10 @@ def test_normal_behavior_with_core_functionality(self): from report_analyst.core.question_loader import get_question_loader question_loader = get_question_loader() - question_set_options = question_loader.get_question_set_options() + [ - "custom" - ] + question_set_options = question_loader.get_question_set_options() + ["custom"] else: # Fallback: use a generic approach without hardcoded names - question_set_options = [ - "custom" - ] # Only custom when core functionality unavailable + question_set_options = ["custom"] # Only custom when core functionality unavailable # Should have all question sets plus custom expected_sets = {"everest", "tcfd", "denali", "kilimanjaro", "lucia", "custom"} @@ -359,16 +349,12 @@ def test_no_hardcoded_question_set_names(self): question_set_options = question_loader.get_question_set_options() else: # Fallback: use a generic approach without hardcoded names - question_set_options = ( - [] - ) # No predefined options when core functionality unavailable + question_set_options = [] # No predefined options when core functionality unavailable # Should not contain any hardcoded question set names hardcoded_names = ["tcfd", "kilimanjaro", "denali", "everest"] for name in hardcoded_names: - assert ( - name not in question_set_options - ), f"Hardcoded name '{name}' found in fallback options" + assert name not in question_set_options, f"Hardcoded name '{name}' found in fallback options" assert question_set_options == [] @@ -387,8 +373,6 @@ def test_question_set_options_consistency(self): # Should contain expected question sets expected_sets = ["everest", "tcfd", "denali", "kilimanjaro", "lucia"] for expected_set in expected_sets: - assert ( - expected_set in options1 - ), f"Expected question set '{expected_set}' not found in options" + assert expected_set in options1, f"Expected question set '{expected_set}' not found in options" assert len(options1) == 5 diff --git a/tests/test_report_data_client.py b/tests/test_report_data_client.py index d7df9ea4c..b80cb9b67 100644 --- a/tests/test_report_data_client.py +++ b/tests/test_report_data_client.py @@ -46,16 +46,12 @@ def backend_config(): """Create test backend configuration""" from report_analyst_search_backend.config import BackendConfig - return BackendConfig( - use_backend=True, backend_url="http://localhost:8000" - ) + return BackendConfig(use_backend=True, backend_url="http://localhost:8000") def test_report_resource_urn_parsing(): """Test parsing backend URNs""" - resource = ReportResource( - name="test.pdf", uri="urn:report-analyst:backend:localhost:8000:abc-123" - ) + resource = ReportResource(name="test.pdf", uri="urn:report-analyst:backend:localhost:8000:abc-123") parsed = resource.parse_backend_urn() assert parsed is not None assert parsed["host"] == "localhost:8000" @@ -64,18 +60,14 @@ def test_report_resource_urn_parsing(): def test_report_resource_resolve_to_http_url(): """Test resolving URN to HTTP URL""" - resource = ReportResource( - name="test.pdf", uri="urn:report-analyst:backend:localhost:8000:abc-123" - ) + resource = ReportResource(name="test.pdf", uri="urn:report-analyst:backend:localhost:8000:abc-123") url = resource.resolve_to_http_url() assert url == "http://localhost:8000/resources/abc-123" def test_report_resource_is_backend_resource(): """Test backend resource detection""" - backend_resource = ReportResource( - name="test.pdf", uri="urn:report-analyst:backend:localhost:8000:abc-123" - ) + backend_resource = ReportResource(name="test.pdf", uri="urn:report-analyst:backend:localhost:8000:abc-123") local_resource = ReportResource(name="test.pdf", uri="file:///path/to/file.pdf") assert backend_resource.is_backend_resource is True @@ -154,9 +146,7 @@ def test_get_backend_service_for_urn_invalid(): """Test getting BackendService with invalid URN""" from report_analyst_search_backend.config import BackendConfig - backend_config = BackendConfig( - use_backend=True, backend_url="http://localhost:8000" - ) + backend_config = BackendConfig(use_backend=True, backend_url="http://localhost:8000") invalid_urn = "file:///path/to/file.pdf" service = get_backend_service_for_urn(invalid_urn, [backend_config]) assert service is None @@ -175,9 +165,7 @@ def test_get_chunks_for_backend_resource(backend_config): } ] - with patch( - "report_analyst.core.report_data_client.get_backend_service_for_urn" - ) as mock_get_service: + with patch("report_analyst.core.report_data_client.get_backend_service_for_urn") as mock_get_service: mock_service = Mock() mock_service.get_chunks.return_value = mock_chunks mock_get_service.return_value = mock_service @@ -215,10 +203,7 @@ def test_report_resource_urn_with_colons_in_resource_id(): def test_report_resource_local_file_uri(): """Test local file URI handling""" - resource = ReportResource( - name="test.pdf", uri="file:///absolute/path/to/file.pdf" - ) + resource = ReportResource(name="test.pdf", uri="file:///absolute/path/to/file.pdf") assert resource.is_local_resource is True assert resource.is_backend_resource is False assert resource.parse_backend_urn() is None - diff --git a/tests/test_s3_upload_service.py b/tests/test_s3_upload_service.py index 9180331cc..9f7f53a49 100644 --- a/tests/test_s3_upload_service.py +++ b/tests/test_s3_upload_service.py @@ -78,9 +78,7 @@ async def test_upload_to_s3(self, s3_service, sample_pdf_bytes): s3_service.s3_client = mock_s3_client s3_service.s3_bucket = "test-bucket" - s3_url = await s3_service._upload_to_s3( - sample_pdf_bytes, "test-key", "test.pdf" - ) + s3_url = await s3_service._upload_to_s3(sample_pdf_bytes, "test-key", "test.pdf") assert s3_url is not None assert "s3.example.com" in s3_url @@ -112,9 +110,7 @@ async def test_upload_pdf_via_s3_nats(self, s3_service, sample_pdf_bytes): mock_nc.jetstream.return_value = mock_js mock_nats_connect.return_value = mock_nc - result = await s3_service.upload_pdf_via_s3_nats( - sample_pdf_bytes, "test.pdf" - ) + result = await s3_service.upload_pdf_via_s3_nats(sample_pdf_bytes, "test.pdf") assert result is not None assert isinstance(result, str) # request_id @@ -148,9 +144,7 @@ async def test_s3_cleanup_on_error(self, s3_service): await s3_service._cleanup_s3_object("test-key") - mock_s3_client.delete_object.assert_called_once_with( - Bucket="test-bucket", Key="test-key" - ) + mock_s3_client.delete_object.assert_called_once_with(Bucket="test-bucket", Key="test-key") def test_is_available(self): """Test availability check""" @@ -172,4 +166,3 @@ def test_is_available(self): # Test when boto3 not available with patch("report_analyst_search_backend.s3_upload_service.boto3", None): assert S3UploadService.is_available() is False - diff --git a/tests/test_service_discovery.py b/tests/test_service_discovery.py index 928849bf4..de0b5fb6a 100644 --- a/tests/test_service_discovery.py +++ b/tests/test_service_discovery.py @@ -49,7 +49,7 @@ def test_validator_initialization(validator): def test_generate_service_template(validator): """Test template generation""" template = validator.generate_service_template() - + assert "service_name" in template assert "version" in template assert "contract_version" in template @@ -63,9 +63,9 @@ def test_validate_valid_service(validator, example_manifest): """Test validation of a valid service manifest""" if example_manifest is None: pytest.skip("Example manifest not found") - + result = validator.validate_service(example_manifest) - + assert isinstance(result, ValidationResult) # Should pass basic validation (may have warnings for extensions) assert result.is_valid or len(result.errors) == 0 @@ -77,9 +77,9 @@ def test_validate_invalid_service(validator): "service_name": "test", # Missing required fields } - + result = validator.validate_service(invalid_manifest) - + assert isinstance(result, ValidationResult) assert not result.is_valid assert len(result.errors) > 0 @@ -92,9 +92,9 @@ def test_validate_missing_required_fields(validator): "version": "1.0.0", # Missing contract_version, protocols, etc. } - + result = validator.validate_service(incomplete_manifest) - + assert not result.is_valid assert any("required" in error.lower() for error in result.errors) @@ -102,7 +102,7 @@ def test_validate_missing_required_fields(validator): def test_get_required_channels(validator): """Test getting required NATS channels""" channels = validator.get_required_channels() - + assert isinstance(channels, dict) assert "publish" in channels assert "subscribe" in channels @@ -113,7 +113,7 @@ def test_get_required_channels(validator): def test_get_required_endpoints(validator): """Test getting required HTTP endpoints""" endpoints = validator.get_required_endpoints() - + assert isinstance(endpoints, list) if len(endpoints) > 0: endpoint = endpoints[0] @@ -125,12 +125,12 @@ def test_get_required_endpoints(validator): def test_validate_from_file(schema_dir): """Test validation from file""" manifest_path = schema_dir / "example-service-manifest.json" - + if not manifest_path.exists(): pytest.skip("Example manifest file not found") - + result = validate_service_from_file(manifest_path) - + assert isinstance(result, ValidationResult) # Should pass basic validation assert result.is_valid or len(result.errors) == 0 @@ -140,7 +140,7 @@ def test_service_manifest_structure(example_manifest): """Test that example manifest has correct structure""" if example_manifest is None: pytest.skip("Example manifest not found") - + # Check required top-level fields assert "service_name" in example_manifest assert "version" in example_manifest @@ -148,17 +148,17 @@ def test_service_manifest_structure(example_manifest): assert "protocols" in example_manifest assert "nats_channels" in example_manifest assert "http_endpoints" in example_manifest - + # Check protocols structure protocols = example_manifest["protocols"] assert "nats" in protocols assert "http" in protocols - + # Check NATS channels structure nats_channels = example_manifest["nats_channels"] assert "publishes" in nats_channels assert "subscribes" in nats_channels - + # Check HTTP endpoints structure http_endpoints = example_manifest["http_endpoints"] assert "required" in http_endpoints @@ -168,9 +168,8 @@ def test_version_compatibility_warning(validator): """Test that version mismatch generates warning""" manifest = validator.generate_service_template() manifest["contract_version"] = "2.0.0" # Different version - + result = validator.validate_service(manifest) - + # Should have warnings about version mismatch assert len(result.warnings) > 0 or result.is_valid # May still be valid but with warnings - diff --git a/tests/test_settings_enterprise_mode.py b/tests/test_settings_enterprise_mode.py index 46404e68c..d2a03a131 100644 --- a/tests/test_settings_enterprise_mode.py +++ b/tests/test_settings_enterprise_mode.py @@ -5,6 +5,7 @@ 1. The checkbox is checked 2. Backend integration is available """ + from streamlit.testing.v1 import AppTest @@ -12,37 +13,37 @@ def test_enterprise_mode_message_only_when_checked(): """Test that enterprise mode message only shows when checkbox is checked""" at = AppTest.from_file("report_analyst/streamlit_app.py") at.run(timeout=10) - + # Navigate to Settings page at.session_state["nav_page"] = "Settings" at.run(timeout=10) - + # Check that Settings page loaded assert "Settings" in str(at), "Settings page should be visible" - + # Initially, checkbox should be unchecked (default False) # Find the checkbox checkboxes = [w for w in at.checkbox if "S3+NATS" in str(w)] assert len(checkboxes) > 0, "S3+NATS checkbox should exist" - + checkbox = checkboxes[0] initial_value = checkbox.value - + # If checkbox is checked initially (from env var), uncheck it if initial_value: checkbox.set_value(False) at.run(timeout=10) - + # After unchecking, the message should NOT appear page_text = str(at) if "Enterprise mode enabled" in page_text: # This is the bug - message appears even when unchecked assert False, "Enterprise mode message should not appear when checkbox is unchecked" - + # Now check the checkbox checkbox.set_value(True) at.run(timeout=10) - + # After checking, if backend is available, message should appear page_text_after = str(at) # Note: We can't easily test backend availability in AppTest, so we just check @@ -54,16 +55,16 @@ def test_enterprise_mode_checkbox_state_persistence(): """Test that checkbox state persists correctly across reruns""" at = AppTest.from_file("report_analyst/streamlit_app.py") at.run(timeout=10) - + # Navigate to Settings at.session_state["nav_page"] = "Settings" at.run(timeout=10) - + # Find checkbox checkboxes = [w for w in at.checkbox if "S3+NATS" in str(w)] assert len(checkboxes) > 0, "S3+NATS checkbox should exist" checkbox = checkboxes[0] - + # Set to True checkbox.set_value(True) at.run(timeout=10) @@ -73,23 +74,22 @@ def test_enterprise_mode_checkbox_state_persistence(): assert at.session_state["use_s3_upload"] == True, "Session state should be True" except (KeyError, AttributeError): pass # Session state might use widget ID instead of key - + # Set to False checkbox.set_value(False) at.run(timeout=10) assert checkbox.value == False, "Checkbox should be False after unchecking" - + # Check that enterprise mode message is NOT shown when unchecked page_text = str(at) if "Enterprise mode enabled" in page_text: assert False, "Enterprise mode message should NOT appear when checkbox is unchecked" - + # Rerun - state should persist at.run(timeout=10) assert checkbox.value == False, "Checkbox should remain False after rerun" - + # Verify message still doesn't appear after rerun page_text_rerun = str(at) if "Enterprise mode enabled" in page_text_rerun: assert False, "Enterprise mode message should NOT appear after rerun when checkbox is unchecked" - diff --git a/tests/test_similarity_search.py b/tests/test_similarity_search.py index 29beff7b0..e33fd459e 100644 --- a/tests/test_similarity_search.py +++ b/tests/test_similarity_search.py @@ -1,13 +1,14 @@ -import pytest -from unittest.mock import Mock, patch, AsyncMock -import numpy as np -from pathlib import Path -import tempfile -import shutil import json +import shutil +import tempfile +from pathlib import Path +from unittest.mock import AsyncMock, Mock, patch + +import numpy as np +import pytest -from report_analyst.core.cache_manager import CacheManager from report_analyst.core.analyzer import DocumentAnalyzer +from report_analyst.core.cache_manager import CacheManager @pytest.fixture @@ -126,9 +127,7 @@ def test_similarity_search_ordering(temp_db, sample_chunks): mock_node3.embedding = sample_chunks[2]["embedding"] # Nodes should be returned in descending order by score - mock_retriever.aretrieve = AsyncMock( - return_value=[mock_node1, mock_node2, mock_node3] - ) + mock_retriever.aretrieve = AsyncMock(return_value=[mock_node1, mock_node2, mock_node3]) temp_db.vector_store = Mock() temp_db.vector_store.as_retriever.return_value = mock_retriever @@ -199,5 +198,3 @@ async def test_chunk_ordering_in_analysis(): assert chunks[0]["chunk_order"] == 0 # Highest similarity = first position assert chunks[1]["chunk_order"] == 1 # Second highest = second position assert chunks[2]["chunk_order"] == 2 # Lowest = last position - - diff --git a/tests/test_streamlit_app_backend_integration.py b/tests/test_streamlit_app_backend_integration.py index f7bc15f85..1ebbfdde9 100644 --- a/tests/test_streamlit_app_backend_integration.py +++ b/tests/test_streamlit_app_backend_integration.py @@ -154,7 +154,7 @@ def test_backend_integration_compatibility(): # Check that navigation page is set in session state assert "nav_page" in at.session_state, "Navigation page not found in session state" - + # Navigate to Report Analyst page to check for title at.session_state["nav_page"] = "Report Analyst" at.run(timeout=10) @@ -167,14 +167,14 @@ def test_backend_integration_compatibility(): def test_backend_resource_full_roundtrip(): """Test full roundtrip: list backend resources, select, retrieve chunks, analyze""" from unittest.mock import Mock, patch - + at = AppTest.from_file("report_analyst/streamlit_app.py") - + # Mock backend configuration mock_backend_config = Mock() mock_backend_config.use_backend = True mock_backend_config.backend_url = "http://localhost:8000" - + # Mock backend resources response mock_resources = [ { @@ -185,7 +185,7 @@ def test_backend_resource_full_roundtrip(): "status": "processed", } ] - + # Mock backend chunks response mock_chunks = [ { @@ -203,74 +203,75 @@ def test_backend_resource_full_roundtrip(): "resource_id": "test-resource-1", }, ] - - with patch("requests.get") as mock_get, \ - patch("requests.post") as mock_post: - + + with patch("requests.get") as mock_get, patch("requests.post") as mock_post: + # Mock resources endpoint mock_resources_response = Mock() mock_resources_response.json.return_value = mock_resources mock_resources_response.status_code = 200 - + # Mock chunks endpoint (search endpoint) mock_chunks_response = Mock() mock_chunks_response.json.return_value = { - "results": [{ - "resource": {"id": "test-resource-1"}, - "chunks": [ - { - "chunk": { - "id": "chunk-1", - "chunk_text": mock_chunks[0]["chunk_text"], - "chunk_metadata": mock_chunks[0]["chunk_metadata"], + "results": [ + { + "resource": {"id": "test-resource-1"}, + "chunks": [ + { + "chunk": { + "id": "chunk-1", + "chunk_text": mock_chunks[0]["chunk_text"], + "chunk_metadata": mock_chunks[0]["chunk_metadata"], + }, + "similarity": mock_chunks[0]["similarity_score"], }, - "similarity": mock_chunks[0]["similarity_score"], - }, - { - "chunk": { - "id": "chunk-2", - "chunk_text": mock_chunks[1]["chunk_text"], - "chunk_metadata": mock_chunks[1]["chunk_metadata"], + { + "chunk": { + "id": "chunk-2", + "chunk_text": mock_chunks[1]["chunk_text"], + "chunk_metadata": mock_chunks[1]["chunk_metadata"], + }, + "similarity": mock_chunks[1]["similarity_score"], }, - "similarity": mock_chunks[1]["similarity_score"], - }, - ], - }], + ], + } + ], } mock_chunks_response.status_code = 200 - + # Setup mock responses def mock_get_side_effect(url, **kwargs): if "/resources/" in url: return mock_resources_response return mock_resources_response - + def mock_post_side_effect(url, **kwargs): if "/search/" in url: return mock_chunks_response return mock_chunks_response - + mock_get.side_effect = mock_get_side_effect mock_post.side_effect = mock_post_side_effect - + # Run app at.run(timeout=10) assert not at.exception, "App failed to load" - + # Set backend config in session state at.session_state["backend_config"] = mock_backend_config - + # Navigate to Report Analyst page at.session_state["nav_page"] = "Report Analyst" at.run(timeout=10) assert not at.exception, "Failed to navigate to Report Analyst page" - + # Verify backend resources are listed # The get_uploaded_files_history should now include backend resources # Check that the dropdown has options (this would include backend resources) - + # Verify URN format is used # This is tested indirectly through the app loading and not crashing # when backend resources are available - + assert not at.exception, "Backend resource roundtrip failed" diff --git a/tests/test_streamlit_app_data_display.py b/tests/test_streamlit_app_data_display.py index 3a5e32b87..7cc1cbbb8 100644 --- a/tests/test_streamlit_app_data_display.py +++ b/tests/test_streamlit_app_data_display.py @@ -41,9 +41,7 @@ def test_file_history_functionality(): # Check for file history selectbox has_file_history = False for sb in at.selectbox: - if "previously analyzed" in str(sb.label).lower() or "previous_file" in str( - sb.key - ): + if "previously analyzed" in str(sb.label).lower() or "previous_file" in str(sb.key): has_file_history = True break @@ -64,10 +62,7 @@ def test_question_display_functionality(): # Look for question-related elements in selectboxes for sb in at.selectbox: - if any( - keyword in str(sb.label).lower() - for keyword in ["question", "tcfd", "everest", "denali", "kilimanjaro"] - ): + if any(keyword in str(sb.label).lower() for keyword in ["question", "tcfd", "everest", "denali", "kilimanjaro"]): has_question_elements = True break @@ -90,9 +85,7 @@ def test_model_selection_display(): has_model_selection = True # Check that it has model options options = [str(opt).lower() for opt in sb.options] - assert any( - "gpt" in opt for opt in options - ), "No GPT models found in options" + assert any("gpt" in opt for opt in options), "No GPT models found in options" break assert has_model_selection, "Model selection not found" @@ -116,15 +109,9 @@ def test_configuration_display(): config_labels.append(str(ni.label).lower()) expected_configs = ["chunk", "overlap", "top", "k"] - found_configs = [ - config - for config in expected_configs - if any(config in label for label in config_labels) - ] - - assert ( - len(found_configs) >= 2 - ), f"Expected at least 2 configuration parameters, found: {found_configs}" + found_configs = [config for config in expected_configs if any(config in label for label in config_labels)] + + assert len(found_configs) >= 2, f"Expected at least 2 configuration parameters, found: {found_configs}" assert not at.exception @@ -140,15 +127,15 @@ def test_analysis_controls_display(): # 1. A file is selected # 2. Questions are loaded # 3. User is on the Report Analyst page - + # Since these are UI elements that depend on user interaction, # we just verify the app loads without errors and can handle the page assert not at.exception, "App should load without errors on Report Analyst page" - + # Check that we're on the right page assert "nav_page" in at.session_state, "Navigation page should be set" assert at.session_state["nav_page"] == "Report Analyst", "Should be on Report Analyst page" - + # Note: UI elements like buttons and checkboxes are conditionally rendered # and may not appear until a file is selected. This is expected behavior. @@ -195,7 +182,7 @@ def test_app_layout_and_structure(): # Check that navigation page is set in session state assert "nav_page" in at.session_state, "Navigation page not found in session state" - + # Navigate to Report Analyst page to check for title at.session_state["nav_page"] = "Report Analyst" at.run(timeout=10) diff --git a/tests/test_streamlit_app_file_selection.py b/tests/test_streamlit_app_file_selection.py index 54010a2a5..4eb10ed13 100644 --- a/tests/test_streamlit_app_file_selection.py +++ b/tests/test_streamlit_app_file_selection.py @@ -5,6 +5,7 @@ import tempfile from pathlib import Path + from streamlit.testing.v1 import AppTest @@ -13,21 +14,24 @@ def test_file_selection_with_file_uri(): # Create a temporary PDF file with tempfile.TemporaryDirectory() as temp_dir: test_pdf = Path(temp_dir) / "test_report.pdf" - test_pdf.write_bytes(b"%PDF-1.4\n%Test PDF\n1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>\nendobj\nxref\n0 4\ntrailer\n<< /Size 4 /Root 1 0 R >>\nstartxref\n100\n%%EOF") - + test_pdf.write_bytes( + b"%PDF-1.4\n%Test PDF\n1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>\nendobj\nxref\n0 4\ntrailer\n<< /Size 4 /Root 1 0 R >>\nstartxref\n100\n%%EOF" + ) + # Set temp directory in environment import os + original_temp = os.environ.get("TEMP_DIR", None) os.environ["TEMP_DIR"] = str(temp_dir) - + try: at = AppTest.from_file("report_analyst/streamlit_app.py") at.session_state["nav_page"] = "Report Analyst" at.run(timeout=10) - + # Verify app loads without errors assert not at.exception, "App should load without errors" - + # Check that file dropdown is available if len(at.selectbox) > 0: file_selectbox = None @@ -35,15 +39,16 @@ def test_file_selection_with_file_uri(): if "previous_file" in str(selectbox.key): file_selectbox = selectbox break - + if file_selectbox and len(file_selectbox.options) > 0: # Verify files are listed assert len(file_selectbox.options) > 0, "Files should be listed in dropdown" - + # The file path should be correctly resolved # This is verified by the app not showing "File not found" error - assert not any("File not found" in str(err) for err in at.error if hasattr(at, 'error')), \ - "Should not show 'File not found' error for valid files" + assert not any( + "File not found" in str(err) for err in at.error if hasattr(at, "error") + ), "Should not show 'File not found' error for valid files" finally: if original_temp: os.environ["TEMP_DIR"] = original_temp @@ -54,39 +59,42 @@ def test_file_selection_with_file_uri(): def test_file_path_resolution_from_uri(): """Test that file:// URIs are correctly converted to file paths""" from report_analyst.streamlit_app import get_uploaded_files_history - + # Create a temporary PDF file with tempfile.TemporaryDirectory() as temp_dir: test_pdf = Path(temp_dir) / "test_report.pdf" - test_pdf.write_bytes(b"%PDF-1.4\n%Test PDF\n1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>\nendobj\nxref\n0 4\ntrailer\n<< /Size 4 /Root 1 0 R >>\nstartxref\n100\n%%EOF") - + test_pdf.write_bytes( + b"%PDF-1.4\n%Test PDF\n1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>\nendobj\nxref\n0 4\ntrailer\n<< /Size 4 /Root 1 0 R >>\nstartxref\n100\n%%EOF" + ) + # Set temp directory in environment import os + original_temp = os.environ.get("TEMP_DIR", None) os.environ["TEMP_DIR"] = str(temp_dir) - + try: # Get file list files = get_uploaded_files_history() - + # Find our test file test_file = None for f in files: if f["name"] == "test_report.pdf": test_file = f break - + if test_file: # Verify path is correctly extracted from file:// URI path = test_file.get("path", "") uri = test_file.get("uri", "") - + # Path should not start with file:// assert not path.startswith("file://"), "Path should not contain file:// prefix" - + # Path should exist assert Path(path).exists(), f"File path should exist: {path}" - + # URI should start with file:// assert uri.startswith("file://"), "URI should start with file://" finally: @@ -101,21 +109,24 @@ def test_file_not_found_error_not_shown(): # Create a temporary PDF file with tempfile.TemporaryDirectory() as temp_dir: test_pdf = Path(temp_dir) / "test_report.pdf" - test_pdf.write_bytes(b"%PDF-1.4\n%Test PDF\n1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>\nendobj\nxref\n0 4\ntrailer\n<< /Size 4 /Root 1 0 R >>\nstartxref\n100\n%%EOF") - + test_pdf.write_bytes( + b"%PDF-1.4\n%Test PDF\n1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>\nendobj\nxref\n0 4\ntrailer\n<< /Size 4 /Root 1 0 R >>\nstartxref\n100\n%%EOF" + ) + # Set temp directory in environment import os + original_temp = os.environ.get("TEMP_DIR", None) os.environ["TEMP_DIR"] = str(temp_dir) - + try: at = AppTest.from_file("report_analyst/streamlit_app.py") at.session_state["nav_page"] = "Report Analyst" at.run(timeout=10) - + # Verify app loads assert not at.exception, "App should load without errors" - + # Check for file dropdown and select a file if len(at.selectbox) > 0: file_selectbox = None @@ -123,21 +134,19 @@ def test_file_not_found_error_not_shown(): if "previous_file" in str(selectbox.key): file_selectbox = selectbox break - + if file_selectbox and len(file_selectbox.options) > 0: # Select first file # Note: We can't easily set the file in session state due to format_func issues # But we can verify the app doesn't show errors at.run(timeout=10) - + # Check that no "File not found" error is shown # This is verified by checking the app doesn't have that error message page_text = str(at) - assert "File not found: None" not in page_text, \ - "Should not show 'File not found: None' error" + assert "File not found: None" not in page_text, "Should not show 'File not found: None' error" finally: if original_temp: os.environ["TEMP_DIR"] = original_temp elif "TEMP_DIR" in os.environ: del os.environ["TEMP_DIR"] - diff --git a/tests/test_streamlit_app_processing_steps.py b/tests/test_streamlit_app_processing_steps.py index 936b52261..1bdf13e4b 100644 --- a/tests/test_streamlit_app_processing_steps.py +++ b/tests/test_streamlit_app_processing_steps.py @@ -11,18 +11,18 @@ def test_processing_steps_slider_exists(): at = AppTest.from_file("report_analyst/streamlit_app.py") at.session_state["nav_page"] = "Report Analyst" at.run(timeout=10) - + # The Processing Steps slider is conditionally rendered # It only appears when a file is selected on the Report Analyst page # Since setting file selection in session state causes issues with format_func, # we just verify the app loads correctly and the page structure is there - + assert not at.exception, "App should load without errors" - + # Verify we're on the Report Analyst page assert "nav_page" in at.session_state, "Navigation page should be set" assert at.session_state["nav_page"] == "Report Analyst", "Should be on Report Analyst page" - + # Note: The processing steps slider uses st.select_slider and is conditionally rendered # It will appear when a file is selected, but we can't easily test that in AppTest # without causing format_func issues. The important thing is the app loads correctly. @@ -33,18 +33,18 @@ def test_processing_steps_slider_interactive(): at = AppTest.from_file("report_analyst/streamlit_app.py") at.session_state["nav_page"] = "Report Analyst" at.run(timeout=10) - + # The Processing Steps slider is conditionally rendered # It only appears when a file is selected on the Report Analyst page # Since setting file selection in session state causes issues with format_func, # we just verify the app loads correctly - + assert not at.exception, "App should load without errors" - + # Verify we're on the Report Analyst page assert "nav_page" in at.session_state, "Navigation page should be set" assert at.session_state["nav_page"] == "Report Analyst", "Should be on Report Analyst page" - + # Note: The processing steps slider uses st.select_slider and is conditionally rendered # Testing interactivity requires a file to be selected, which causes format_func issues # in AppTest. The important thing is the app loads correctly and the page structure is there. @@ -55,20 +55,18 @@ def test_processing_steps_displayed(): at = AppTest.from_file("report_analyst/streamlit_app.py") at.session_state["nav_page"] = "Report Analyst" at.run(timeout=10) - + # Check that Processing Steps section exists # This is verified by the app loading without errors assert not at.exception, "App should load without errors" - + # Check for Processing Steps heading has_processing_steps = False for header in at.header: if "Processing Steps" in str(header.value): has_processing_steps = True break - + # Processing Steps might only show when a file is selected # So we just verify the app loads correctly assert not at.exception - - diff --git a/tests/test_streamlit_app_questions.py b/tests/test_streamlit_app_questions.py index bb2e4f675..dde571ece 100644 --- a/tests/test_streamlit_app_questions.py +++ b/tests/test_streamlit_app_questions.py @@ -36,15 +36,11 @@ def test_question_sets_loaded_dynamically(): if "Question Set" in str(sb.label) or "new_question_set" in str(sb.key): options = [str(opt) for opt in sb.options] # Should have multiple question sets - assert ( - len(options) >= 4 - ), f"Expected at least 4 question sets, got {len(options)}" + assert len(options) >= 4, f"Expected at least 4 question sets, got {len(options)}" # Verify key question sets present options_lower = [opt.lower() for opt in options] assert any("tcfd" in opt for opt in options_lower), "TCFD not in options" - assert any( - "everest" in opt for opt in options_lower - ), "Everest not in options" + assert any("everest" in opt for opt in options_lower), "Everest not in options" break assert not at.exception @@ -74,17 +70,9 @@ def test_question_set_selectbox_has_options(): options_lower = [opt.lower() for opt in options] # Should have the key question sets - assert any( - "tcfd" in opt for opt in options_lower - ), f"TCFD not found in options: {options}" - assert any( - "everest" in opt for opt in options_lower - ), f"Everest not found in options: {options}" - assert any( - "denali" in opt for opt in options_lower - ), f"Denali not found in options: {options}" - assert any( - "kilimanjaro" in opt for opt in options_lower - ), f"Kilimanjaro not found in options: {options}" + assert any("tcfd" in opt for opt in options_lower), f"TCFD not found in options: {options}" + assert any("everest" in opt for opt in options_lower), f"Everest not found in options: {options}" + assert any("denali" in opt for opt in options_lower), f"Denali not found in options: {options}" + assert any("kilimanjaro" in opt for opt in options_lower), f"Kilimanjaro not found in options: {options}" assert not at.exception diff --git a/tests/test_streamlit_app_tabs.py b/tests/test_streamlit_app_tabs.py index 8d3ac9cdb..cfbd0bb0e 100644 --- a/tests/test_streamlit_app_tabs.py +++ b/tests/test_streamlit_app_tabs.py @@ -84,9 +84,7 @@ def test_consolidated_results_tab(): break # The selectbox should exist for question set selection - assert ( - has_consolidated_selectbox - ), "Question set selectbox not found in All Results page" + assert has_consolidated_selectbox, "Question set selectbox not found in All Results page" assert not at.exception @@ -147,18 +145,14 @@ def test_question_set_selection(): assert question_selectbox is not None, "Question set selectbox not found" # Check that it has multiple options - assert ( - len(question_selectbox.options) >= 4 - ), f"Expected at least 4 question sets, got {len(question_selectbox.options)}" + assert len(question_selectbox.options) >= 4, f"Expected at least 4 question sets, got {len(question_selectbox.options)}" # Verify key question sets are present options_lower = [str(opt).lower() for opt in question_selectbox.options] expected_sets = ["tcfd", "everest", "denali", "kilimanjaro"] for expected_set in expected_sets: - assert any( - expected_set in opt for opt in options_lower - ), f"Question set '{expected_set}' not found in options" + assert any(expected_set in opt for opt in options_lower), f"Question set '{expected_set}' not found in options" assert not at.exception @@ -174,13 +168,13 @@ def test_analysis_controls(): # They only appear when a file is selected and questions are loaded # Since setting file selection in session state causes format_func issues in AppTest, # we verify the app loads correctly and the page structure is there - + assert not at.exception, "App should load without errors" - + # Verify we're on the Report Analyst page assert "nav_page" in at.session_state, "Navigation page should be set" assert at.session_state["nav_page"] == "Report Analyst", "Should be on Report Analyst page" - + # Note: UI elements like buttons and checkboxes are conditionally rendered # and may not appear until a file is selected. This is expected behavior. # The important thing is that the app loads correctly and handles the page navigation. @@ -208,7 +202,7 @@ def test_session_state_initialization(): # Navigate to Report Analyst page to check for title at.session_state["nav_page"] = "Report Analyst" at.run(timeout=10) - + # Check that the app has the expected structure assert len(at.title) > 0, "App title not found" assert len(at.expander) > 0, "No expanders found" From 92e846d2cf2c223e07c34c2f5c692cf2642191cc Mon Sep 17 00:00:00 2001 From: Christian Date: Thu, 18 Dec 2025 13:08:09 +0100 Subject: [PATCH 07/21] Fix test dependencies and PostgreSQL file storage persistence - Add missing test dependencies (aioresponses, psycopg2-binary, pytest-postgresql) - Remove fallback logic in tests - fail if dependencies missing - Fix PostgreSQL test database setup with transaction rollback pattern - Fix S3 override persistence across tab navigation - Add API key source visibility (session vs environment) - Fix PostgreSQL file storage checkbox persistence across pages - Add find_by_filename method to retrieve existing files from PostgreSQL - Add tests for find_by_filename and retrieve existing file flows --- report_analyst/core/file_storage.py | 22 +++ report_analyst/streamlit_app.py | 128 ++++++++++++----- requirements.txt | 5 + tests/conftest.py | 154 ++++++++++++++++++++- tests/test_external_service_integration.py | 104 ++++++-------- tests/test_file_storage.py | 82 +++++++++-- tests/test_settings_enterprise_mode.py | 149 +++++++++++--------- 7 files changed, 468 insertions(+), 176 deletions(-) diff --git a/report_analyst/core/file_storage.py b/report_analyst/core/file_storage.py index 480a23a5c..d8e1eb8c9 100644 --- a/report_analyst/core/file_storage.py +++ b/report_analyst/core/file_storage.py @@ -203,6 +203,28 @@ def delete_file(self, file_id: str) -> bool: logger.error(f"Error deleting file {file_id}: {str(e)}") return False + def find_by_filename(self, filename: str) -> Optional[str]: + """ + Find a file by filename and return its ID. + + Args: + filename: Original filename to search for + + Returns: + file_id if found, None otherwise + """ + try: + with self.db_manager.get_connection() as conn: + query = text("SELECT id FROM stored_files WHERE filename = :filename ORDER BY created_at DESC LIMIT 1") + result = conn.execute(query, {"filename": filename}) + row = result.fetchone() + if row: + return row[0] + return None + except Exception as e: + logger.error(f"Error finding file by filename {filename}: {str(e)}") + return None + def save_to_temp(self, file_id: str, temp_dir: Path = Path("temp")) -> Optional[str]: """ Retrieve file from PostgreSQL and save to temporary directory. diff --git a/report_analyst/streamlit_app.py b/report_analyst/streamlit_app.py index 57f67f115..b059e9bdf 100644 --- a/report_analyst/streamlit_app.py +++ b/report_analyst/streamlit_app.py @@ -278,17 +278,23 @@ def save_uploaded_file(uploaded_file) -> Optional[str]: if isinstance(uploaded_file, (str, Path)): return str(uploaded_file) - # Check if file was already saved in this session + # Check if PostgreSQL file storage is enabled (check both keys for persistence) + use_postgres_storage = st.session_state.get("postgres_file_storage_enabled", False) or st.session_state.get( + "use_postgres_file_storage", False + ) + + # Check if file was already saved in this session with the SAME storage mode file_key = f"saved_file_{uploaded_file.name}" - if file_key in st.session_state: + storage_mode_key = f"saved_file_mode_{uploaded_file.name}" + cached_mode = st.session_state.get(storage_mode_key) + + # Only use cache if storage mode matches + if file_key in st.session_state and cached_mode == ("postgres" if use_postgres_storage else "local"): return st.session_state[file_key] # Get file bytes file_bytes = uploaded_file.getbuffer() - # Check if PostgreSQL file storage is enabled - use_postgres_storage = st.session_state.get("use_postgres_file_storage", False) - if use_postgres_storage: try: from report_analyst.core.file_storage import get_file_storage @@ -298,7 +304,22 @@ def save_uploaded_file(uploaded_file) -> Optional[str]: file_storage = get_file_storage(database_url) if file_storage: - # Store in PostgreSQL + # Check if file already exists in PostgreSQL + existing_file_id = file_storage.find_by_filename(uploaded_file.name) + if existing_file_id: + # Retrieve from PostgreSQL instead of re-uploading + temp_path = file_storage.save_to_temp(existing_file_id) + if temp_path: + st.session_state[file_key] = temp_path + st.session_state[f"{file_key}_id"] = existing_file_id + st.session_state[storage_mode_key] = "postgres" + logger.info( + f"Retrieved existing file {uploaded_file.name} from PostgreSQL (ID: {existing_file_id})" + ) + st.session_state.file_processed = False + return temp_path + + # Store new file in PostgreSQL file_id = file_storage.store_file(file_bytes, uploaded_file.name, uploaded_file.type) # Save to temp for processing (retrieve from DB) @@ -308,6 +329,7 @@ def save_uploaded_file(uploaded_file) -> Optional[str]: # Store both file_id and path in session state st.session_state[file_key] = temp_path st.session_state[f"{file_key}_id"] = file_id + st.session_state[storage_mode_key] = "postgres" logger.info(f"Stored file {uploaded_file.name} in PostgreSQL (ID: {file_id})") st.session_state.file_processed = False return temp_path @@ -1465,8 +1487,11 @@ def main(): st.session_state.analysis_complete = False # Initialize analysis complete flag # Initialize use_s3_upload in session state if not already set + # Respect override_s3_upload if user has temporarily disabled it if "use_s3_upload" not in st.session_state: - st.session_state.use_s3_upload = os.getenv("USE_S3_UPLOAD", "false").lower() == "true" + env_s3_upload = os.getenv("USE_S3_UPLOAD", "false").lower() == "true" + override = st.session_state.get("override_s3_upload", False) + st.session_state.use_s3_upload = env_s3_upload and not override # Sync API keys from session state to environment at startup APIKeyManager.sync_api_keys_to_env(st.session_state) @@ -2823,7 +2848,13 @@ def main(): masked_openai = ( f"{current_openai_key[:8]}...{current_openai_key[-4:]}" if len(current_openai_key) > 12 else "***" ) - st.caption(f"Current key: `{masked_openai}`") + # Show source of key + if session_openai_key: + st.success(f"✓ API key set in session: `{masked_openai}`") + elif env_openai_key: + st.info(f"API key from environment: `{masked_openai}`") + else: + st.caption(f"Current key: `{masked_openai}`") # Track override state override_openai = st.session_state.get("override_openai_key", False) @@ -3009,31 +3040,42 @@ def main(): # Enterprise Integration (S3+NATS) st.subheader("Enterprise Integration") - # In Streamlit, when a widget has a 'key', it automatically syncs with session state - # The widget's return value is the current value from session state (or default if not set) - # IMPORTANT: Don't provide 'value' parameter when using 'key' - let Streamlit manage it - # The widget return value is the source of truth for the current render - st.markdown( - """ - - """, - unsafe_allow_html=True, - ) - use_s3_upload = st.checkbox( - "Enable S3+NATS Upload", - key="use_s3_upload", - help="Upload documents via S3 and process via NATS for enterprise integration", - ) - # Show Enterprise Mode status only if checkbox is checked AND backend is available - # Check AFTER widget render - use widget return value which reflects current state - # The widget return value is the authoritative source for the current render cycle + # Check if USE_S3_UPLOAD is set from environment + env_s3_upload = os.getenv("USE_S3_UPLOAD", "").lower() == "true" + + # Show env var status like API keys + if env_s3_upload and not st.session_state.get("override_s3_upload", False): + st.info("S3+NATS upload is enabled via `USE_S3_UPLOAD` environment variable") + col1, col2 = st.columns([1, 3]) + with col1: + if st.button("Disable temporarily", key="btn_override_s3"): + st.session_state.override_s3_upload = True + st.session_state.use_s3_upload = False + st.rerun() + use_s3_upload = True + else: + st.markdown( + """ + + """, + unsafe_allow_html=True, + ) + use_s3_upload = st.checkbox( + "Enable S3+NATS Upload", + key="use_s3_upload", + help="Upload documents via S3 and process via NATS for enterprise integration", + ) + if not env_s3_upload: + st.caption("*Or set `USE_S3_UPLOAD=true` in environment*") + + # Show Enterprise Mode status only if enabled AND backend is available if use_s3_upload and BACKEND_INTEGRATION_AVAILABLE: st.info("Enterprise mode enabled") @@ -3062,17 +3104,21 @@ def main(): ("postgresql://", "postgres://") ) - # Initialize use_postgres_file_storage from session state or env - if "use_postgres_file_storage" not in st.session_state: - st.session_state.use_postgres_file_storage = os.getenv("USE_POSTGRES_FILE_STORAGE", "false").lower() == "true" + # Initialize postgres_file_storage_enabled from session state or env + if "postgres_file_storage_enabled" not in st.session_state: + st.session_state.postgres_file_storage_enabled = ( + os.getenv("USE_POSTGRES_FILE_STORAGE", "false").lower() == "true" + ) if is_postgres_enterprise: use_postgres_storage = st.checkbox( "Store files in PostgreSQL", - value=st.session_state.get("use_postgres_file_storage", False), + value=st.session_state.get("postgres_file_storage_enabled", False), key="use_postgres_file_storage", help="Store uploaded files in PostgreSQL database (useful for Heroku deployments). Files are stored as BYTEA/BLOB. This is an enterprise feature.", ) + # Store in a separate key that persists across page navigation + st.session_state.postgres_file_storage_enabled = use_postgres_storage if use_postgres_storage: st.info("📦 Files will be stored in PostgreSQL database") @@ -3839,8 +3885,14 @@ def main(): # Upload Report page elif nav_page == "Upload Report": - # Check if S3+NATS enterprise integration is enabled (from UI checkbox) - use_s3_upload = st.session_state.get("use_s3_upload", False) + # Check if S3+NATS enterprise integration is enabled + # Respect override_s3_upload if user has temporarily disabled it + if st.session_state.get("override_s3_upload", False): + use_s3_upload = False + else: + use_s3_upload = ( + st.session_state.get("use_s3_upload", False) or os.getenv("USE_S3_UPLOAD", "false").lower() == "true" + ) # Initialize backend integration with S3+NATS enabled if needed if use_s3_upload and BACKEND_INTEGRATION_AVAILABLE: diff --git a/requirements.txt b/requirements.txt index ea369c9b8..03ae41e8b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ aiofiles==23.2.1 +aioresponses==0.7.7 aiohappyeyeballs==2.4.4 aiohttp==3.11.11 aioitertools==0.12.0 @@ -160,8 +161,12 @@ pyparsing==3.2.1 pypdf==5.2.0 PyPika==0.48.9 pyproject_hooks==1.2.0 +psycopg2-binary==2.9.10 pytest==8.3.5 pytest-asyncio==0.25.3 +pytest-cov==7.0.0 +pytest-env==1.1.1 +pytest-postgresql==7.0.2 python-dateutil==2.9.0.post0 python-dotenv==1.0.1 python-multipart==0.0.20 diff --git a/tests/conftest.py b/tests/conftest.py index 807ae4ee8..63b92ffe6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,17 +1,169 @@ """ -Shared pytest fixtures for event router tests +Shared pytest fixtures for all tests. + +Provides test environment configuration and common fixtures. """ import json +import os import tempfile from pathlib import Path from unittest.mock import AsyncMock, Mock import pytest import yaml +from dotenv import load_dotenv + +# Load environment variables from .env file before any other imports +load_dotenv() from report_analyst_jobs.event_router import IGNORE_ACTION, EventRouter +# ============================================================================= +# Test Environment Configuration +# ============================================================================= +# Override environment variables for test isolation +# These are set before any test imports happen + + +def _setup_test_environment(): + """Configure test environment variables. + + This sets up isolated test configuration to avoid affecting production data. + Tests can override these values using pytest fixtures or monkeypatch. + """ + test_env = { + # Use separate test database + "TEST_DATABASE_URL": os.getenv("TEST_DATABASE_URL", "postgresql://analyst:analyst@localhost:5432/reports_test"), + # Disable enterprise features by default in tests + "USE_S3_UPLOAD": os.getenv("TEST_USE_S3_UPLOAD", "false"), + # Use test storage path + "TEST_STORAGE_PATH": os.getenv("TEST_STORAGE_PATH", "storage_test"), + } + return test_env + + +TEST_ENV = _setup_test_environment() + + +def pytest_configure(config): + """Pytest hook to configure test environment.""" + # Register custom markers + config.addinivalue_line("markers", "postgres: mark test as requiring PostgreSQL") + config.addinivalue_line("markers", "integration: mark test as integration test") + + +# ============================================================================= +# PostgreSQL Test Database Configuration +# ============================================================================= +# Uses a dedicated test database (test_reports) with transaction rollback +# for test isolation. Each test runs in a transaction that is rolled back +# after the test, keeping the database clean. + + +def _get_test_database_url(): + """Get the test database URL from environment. + + Uses TEST_DATABASE_URL if set, otherwise constructs from PGHOST/PGPORT/etc. + """ + # First check for explicit TEST_DATABASE_URL + url = os.getenv("TEST_DATABASE_URL") + if url: + return url + + # Construct from PG* environment variables + host = os.getenv("PGHOST", "localhost") + port = os.getenv("PGPORT", "5432") + user = os.getenv("PGUSER", "analyst") + password = os.getenv("PGPASSWORD", "analyst") + dbname = os.getenv("TEST_PGDATABASE", "test_reports") + + return f"postgresql://{user}:{password}@{host}:{port}/{dbname}" + + +def _test_database_connection(url): + """Test if we can connect to the database.""" + try: + from sqlalchemy import create_engine, text + + engine = create_engine(url) + with engine.connect() as conn: + conn.execute(text("SELECT 1")) + engine.dispose() + return True + except Exception: + return False + + +@pytest.fixture(scope="session") +def test_db_engine(): + """Create a SQLAlchemy engine for the test database. + + Session-scoped to reuse the connection pool across all tests. + """ + from sqlalchemy import create_engine, text + + url = _get_test_database_url() + + if not _test_database_connection(url): + pytest.skip(f"Cannot connect to test database: {url}") + return None + + engine = create_engine(url) + + # Set up schema once at the start of the test session + with engine.connect() as conn: + conn.execute( + text( + """ + CREATE TABLE IF NOT EXISTS stored_files ( + id VARCHAR(36) PRIMARY KEY, + filename VARCHAR(255) NOT NULL, + file_data BYTEA NOT NULL, + content_type VARCHAR(100), + file_size VARCHAR(20), + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ) + """ + ) + ) + conn.commit() + + yield engine + + engine.dispose() + + +@pytest.fixture(scope="function") +def test_db_connection(test_db_engine): + """Provide a database connection with transaction rollback. + + Each test runs in a transaction that is rolled back after the test, + ensuring complete test isolation without affecting other tests. + """ + connection = test_db_engine.connect() + transaction = connection.begin() + + yield connection + + # Rollback the transaction after each test + transaction.rollback() + connection.close() + + +@pytest.fixture(scope="function") +def test_database_url(test_db_engine): + """Get the test database URL.""" + return _get_test_database_url() + + +@pytest.fixture +def postgres_available(test_database_url): + """Fixture that skips test if PostgreSQL is not available.""" + if not test_database_url: + pytest.skip("PostgreSQL test database not available") + return test_database_url + @pytest.fixture def event_router_config(): diff --git a/tests/test_external_service_integration.py b/tests/test_external_service_integration.py index 155c1c3c6..cc997b5e4 100644 --- a/tests/test_external_service_integration.py +++ b/tests/test_external_service_integration.py @@ -14,13 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest - -try: - from aioresponses import aioresponses - - HAS_AIORESPONSES = True -except ImportError: - HAS_AIORESPONSES = False +from aioresponses import aioresponses from report_analyst_search_backend.external_service_client import ( ExternalServiceClient, @@ -233,67 +227,57 @@ async def test_notify_ready_nats(self, external_client): @pytest.mark.asyncio async def test_notify_ready_http(self, external_client): """Test notifying via HTTP""" - if HAS_AIORESPONSES: - with aioresponses() as m: - m.post( - "http://localhost:8000/external/services/service-x/notify", - status=200, - ) - result = await external_client.notify_ready( - service_id="service-x", - request_id="req-123", - content_type="chunks", - chunks=[{"id": "1", "text": "test"}], - method="http", - ) - assert result is True - else: - # Fallback: skip if aioresponses not available - pytest.skip("aioresponses not available for HTTP mocking") + with aioresponses() as m: + m.post( + "http://localhost:8000/external/services/service-x/notify", + status=200, + ) + result = await external_client.notify_ready( + service_id="service-x", + request_id="req-123", + content_type="chunks", + chunks=[{"id": "1", "text": "test"}], + method="http", + ) + assert result is True @pytest.mark.asyncio async def test_request_analysis_http(self, external_client): """Test requesting analysis via HTTP""" - if HAS_AIORESPONSES: - with aioresponses() as m: - m.post( - "http://localhost:8000/external/services/service-x/analyze", - status=200, - payload={"request_id": "analysis-123"}, - ) - request_id = await external_client.request_analysis( - service_id="service-x", - external_request_id="req-123", - content=[{"id": "1", "text": "test"}], - question_set="tcfd", - analysis_config={"model": "gpt-4o-mini"}, - method="http", - ) - assert request_id == "analysis-123" - else: - pytest.skip("aioresponses not available for HTTP mocking") + with aioresponses() as m: + m.post( + "http://localhost:8000/external/services/service-x/analyze", + status=200, + payload={"request_id": "analysis-123"}, + ) + request_id = await external_client.request_analysis( + service_id="service-x", + external_request_id="req-123", + content=[{"id": "1", "text": "test"}], + question_set="tcfd", + analysis_config={"model": "gpt-4o-mini"}, + method="http", + ) + assert request_id == "analysis-123" @pytest.mark.asyncio async def test_get_results(self, external_client): """Test polling for results""" - if HAS_AIORESPONSES: - with aioresponses() as m: - m.get( - "http://localhost:8000/external/services/service-x/results/analysis-123", - status=200, - payload={ - "request_id": "analysis-123", - "status": "completed", - "answers": [], - "top_chunks": [], - }, - ) - results = await external_client.get_results("service-x", "analysis-123") - - assert results is not None - assert results["status"] == "completed" - else: - pytest.skip("aioresponses not available for HTTP mocking") + with aioresponses() as m: + m.get( + "http://localhost:8000/external/services/service-x/results/analysis-123", + status=200, + payload={ + "request_id": "analysis-123", + "status": "completed", + "answers": [], + "top_chunks": [], + }, + ) + results = await external_client.get_results("service-x", "analysis-123") + + assert results is not None + assert results["status"] == "completed" class TestExternalServiceDelivery: diff --git a/tests/test_file_storage.py b/tests/test_file_storage.py index 5cf62ece3..f0bc2fb91 100644 --- a/tests/test_file_storage.py +++ b/tests/test_file_storage.py @@ -2,6 +2,10 @@ Tests for PostgreSQL file storage service. Tests that files can be stored and retrieved from PostgreSQL. + +These tests require a running PostgreSQL database. They will be skipped +if the test database is not available. Configure TEST_DATABASE_URL +environment variable to point to your test database. """ import os @@ -15,14 +19,10 @@ ) -@pytest.mark.skipif( - not os.getenv("DATABASE_URL") or not os.getenv("DATABASE_URL").startswith(("postgresql://", "postgres://")), - reason="PostgreSQL not configured", -) -def test_postgres_file_storage_store_and_retrieve(): +@pytest.mark.postgres +def test_postgres_file_storage_store_and_retrieve(postgres_available): """Test storing and retrieving a file from PostgreSQL""" - database_url = os.getenv("DATABASE_URL") - storage = PostgreSQLFileStorage(database_url) + storage = PostgreSQLFileStorage(postgres_available) # Test file content test_content = b"Test file content for PostgreSQL storage" @@ -48,14 +48,10 @@ def test_postgres_file_storage_store_and_retrieve(): storage.delete_file(file_id) -@pytest.mark.skipif( - not os.getenv("DATABASE_URL") or not os.getenv("DATABASE_URL").startswith(("postgresql://", "postgres://")), - reason="PostgreSQL not configured", -) -def test_postgres_file_storage_delete(): +@pytest.mark.postgres +def test_postgres_file_storage_delete(postgres_available): """Test deleting a file from PostgreSQL""" - database_url = os.getenv("DATABASE_URL") - storage = PostgreSQLFileStorage(database_url) + storage = PostgreSQLFileStorage(postgres_available) # Store file test_content = b"Test file to delete" @@ -70,6 +66,64 @@ def test_postgres_file_storage_delete(): assert retrieved is None +@pytest.mark.postgres +def test_postgres_file_storage_find_by_filename(postgres_available): + """Test finding a file by filename in PostgreSQL""" + storage = PostgreSQLFileStorage(postgres_available) + + # Store file + test_content = b"Test file for find by filename" + filename = "find_test_unique.pdf" + file_id = storage.store_file(test_content, filename, "application/pdf") + + try: + # Find by filename + found_id = storage.find_by_filename(filename) + assert found_id == file_id + + # Find non-existent file + not_found = storage.find_by_filename("nonexistent_file.pdf") + assert not_found is None + finally: + # Clean up + storage.delete_file(file_id) + + +@pytest.mark.postgres +def test_postgres_file_storage_retrieve_existing(postgres_available): + """Test the flow where a file already exists in PostgreSQL and is retrieved instead of re-uploaded""" + storage = PostgreSQLFileStorage(postgres_available) + + # Store file first time + test_content = b"Test content for existing file retrieval" + filename = "existing_file_test.pdf" + original_file_id = storage.store_file(test_content, filename, "application/pdf") + + try: + # Simulate new session - find existing file by filename + found_id = storage.find_by_filename(filename) + assert found_id == original_file_id + + # Retrieve the file using the found ID + temp_path = storage.save_to_temp(found_id) + assert temp_path is not None + assert filename in temp_path + + # Verify content matches + with open(temp_path, "rb") as f: + retrieved_content = f.read() + assert retrieved_content == test_content + + # Clean up temp file + import os + + if os.path.exists(temp_path): + os.remove(temp_path) + finally: + # Clean up database + storage.delete_file(original_file_id) + + def test_get_file_storage_without_postgres(): """Test that get_file_storage returns None when PostgreSQL is not configured""" # Temporarily unset DATABASE_URL diff --git a/tests/test_settings_enterprise_mode.py b/tests/test_settings_enterprise_mode.py index d2a03a131..b3bc74a1e 100644 --- a/tests/test_settings_enterprise_mode.py +++ b/tests/test_settings_enterprise_mode.py @@ -2,94 +2,117 @@ Test for Settings page enterprise mode checkbox behavior. Tests that the "Enterprise mode enabled" message only appears when: -1. The checkbox is checked +1. The checkbox is checked (or env var is set and not overridden) 2. Backend integration is available """ +import os +from unittest.mock import patch + from streamlit.testing.v1 import AppTest def test_enterprise_mode_message_only_when_checked(): """Test that enterprise mode message only shows when checkbox is checked""" - at = AppTest.from_file("report_analyst/streamlit_app.py") - at.run(timeout=10) + # Run without USE_S3_UPLOAD env var to test checkbox behavior + with patch.dict(os.environ, {"USE_S3_UPLOAD": "false"}, clear=False): + at = AppTest.from_file("report_analyst/streamlit_app.py") + at.run(timeout=10) - # Navigate to Settings page - at.session_state["nav_page"] = "Settings" - at.run(timeout=10) + # Navigate to Settings page + at.session_state["nav_page"] = "Settings" + at.run(timeout=10) - # Check that Settings page loaded - assert "Settings" in str(at), "Settings page should be visible" + # Check that Settings page loaded + assert "Settings" in str(at), "Settings page should be visible" - # Initially, checkbox should be unchecked (default False) - # Find the checkbox - checkboxes = [w for w in at.checkbox if "S3+NATS" in str(w)] - assert len(checkboxes) > 0, "S3+NATS checkbox should exist" + # Find the checkbox (should exist when env var is not set) + checkboxes = [w for w in at.checkbox if "S3+NATS" in str(w)] + assert len(checkboxes) > 0, "S3+NATS checkbox should exist when env var not set" - checkbox = checkboxes[0] - initial_value = checkbox.value + checkbox = checkboxes[0] + initial_value = checkbox.value - # If checkbox is checked initially (from env var), uncheck it - if initial_value: - checkbox.set_value(False) - at.run(timeout=10) + # If checkbox is checked initially, uncheck it + if initial_value: + checkbox.set_value(False) + at.run(timeout=10) - # After unchecking, the message should NOT appear - page_text = str(at) - if "Enterprise mode enabled" in page_text: - # This is the bug - message appears even when unchecked - assert False, "Enterprise mode message should not appear when checkbox is unchecked" + # After unchecking, the message should NOT appear + page_text = str(at) + if "Enterprise mode enabled" in page_text: + assert False, "Enterprise mode message should not appear when checkbox is unchecked" - # Now check the checkbox - checkbox.set_value(True) - at.run(timeout=10) + # Now check the checkbox + checkbox.set_value(True) + at.run(timeout=10) - # After checking, if backend is available, message should appear - page_text_after = str(at) - # Note: We can't easily test backend availability in AppTest, so we just check - # that the checkbox state is correctly reflected - assert checkbox.value == True, "Checkbox should be checked" + # After checking, checkbox state should be reflected + assert checkbox.value is True, "Checkbox should be checked" def test_enterprise_mode_checkbox_state_persistence(): """Test that checkbox state persists correctly across reruns""" - at = AppTest.from_file("report_analyst/streamlit_app.py") - at.run(timeout=10) + # Run without USE_S3_UPLOAD env var to test checkbox behavior + with patch.dict(os.environ, {"USE_S3_UPLOAD": "false"}, clear=False): + at = AppTest.from_file("report_analyst/streamlit_app.py") + at.run(timeout=10) - # Navigate to Settings - at.session_state["nav_page"] = "Settings" - at.run(timeout=10) + # Navigate to Settings + at.session_state["nav_page"] = "Settings" + at.run(timeout=10) - # Find checkbox - checkboxes = [w for w in at.checkbox if "S3+NATS" in str(w)] - assert len(checkboxes) > 0, "S3+NATS checkbox should exist" - checkbox = checkboxes[0] + # Find checkbox + checkboxes = [w for w in at.checkbox if "S3+NATS" in str(w)] + assert len(checkboxes) > 0, "S3+NATS checkbox should exist" + checkbox = checkboxes[0] - # Set to True - checkbox.set_value(True) - at.run(timeout=10) - assert checkbox.value == True, "Checkbox should be True after setting" - # Note: AppTest session_state doesn't support .get(), access directly - try: - assert at.session_state["use_s3_upload"] == True, "Session state should be True" - except (KeyError, AttributeError): - pass # Session state might use widget ID instead of key + # Set to True + checkbox.set_value(True) + at.run(timeout=10) + assert checkbox.value is True, "Checkbox should be True after setting" - # Set to False - checkbox.set_value(False) - at.run(timeout=10) - assert checkbox.value == False, "Checkbox should be False after unchecking" + # Set to False + checkbox.set_value(False) + at.run(timeout=10) + assert checkbox.value is False, "Checkbox should be False after unchecking" - # Check that enterprise mode message is NOT shown when unchecked - page_text = str(at) - if "Enterprise mode enabled" in page_text: - assert False, "Enterprise mode message should NOT appear when checkbox is unchecked" + # Check that enterprise mode message is NOT shown when unchecked + page_text = str(at) + if "Enterprise mode enabled" in page_text: + assert False, "Enterprise mode message should NOT appear when checkbox is unchecked" + + # Rerun - state should persist + at.run(timeout=10) + assert checkbox.value is False, "Checkbox should remain False after rerun" - # Rerun - state should persist + +def test_enterprise_mode_env_var_detection(): + """Test that the app correctly detects USE_S3_UPLOAD env var""" + # This test verifies the env var detection logic works + # When USE_S3_UPLOAD=true is in env, session state should be True + env_value = os.getenv("USE_S3_UPLOAD", "false").lower() == "true" + + at = AppTest.from_file("report_analyst/streamlit_app.py") + at.run(timeout=10) + + # Navigate to Settings + at.session_state["nav_page"] = "Settings" at.run(timeout=10) - assert checkbox.value == False, "Checkbox should remain False after rerun" - # Verify message still doesn't appear after rerun - page_text_rerun = str(at) - if "Enterprise mode enabled" in page_text_rerun: - assert False, "Enterprise mode message should NOT appear after rerun when checkbox is unchecked" + # Session state should match env var (unless overridden) + # AppTest session_state doesn't support .get(), access with try/except + try: + session_value = at.session_state["use_s3_upload"] + except KeyError: + session_value = False + + # If env var is true, session should be true (unless override) + if env_value: + # When env is true, session should be true unless overridden + try: + override = at.session_state["override_s3_upload"] + except KeyError: + override = False + if not override: + assert session_value is True, "use_s3_upload should be True when USE_S3_UPLOAD=true" From 32f9f7b5ccfaed7f0dbb086947138f59de2210e1 Mon Sep 17 00:00:00 2001 From: Christian Date: Thu, 19 Feb 2026 21:01:55 +0100 Subject: [PATCH 08/21] Fix black formatting in alembic migration --- alembic/versions/001_initial_schema.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/alembic/versions/001_initial_schema.py b/alembic/versions/001_initial_schema.py index b9901d230..53292db82 100644 --- a/alembic/versions/001_initial_schema.py +++ b/alembic/versions/001_initial_schema.py @@ -1,7 +1,7 @@ """initial_schema Revision ID: 001_initial_schema -Revises: +Revises: Create Date: 2025-12-14 01:00:00.000000 """ From 878def230b57675795797400bc119c280df26da9 Mon Sep 17 00:00:00 2001 From: Christian Date: Thu, 19 Feb 2026 21:07:33 +0100 Subject: [PATCH 09/21] Add Render and Streamlit deploy buttons - Add render.yaml for one-click Render deploy - Add deploy badges to README - Add .githooks to strip Co-authored-by from commits - Document git hooks setup in INSTALL.md Closes #25 --- .githooks/README.md | 15 +++++++++++++++ .githooks/prepare-commit-msg | 8 ++++++++ INSTALL.md | 12 ++++++++++++ README.md | 2 ++ render.yaml | 16 ++++++++++++++++ 5 files changed, 53 insertions(+) create mode 100644 .githooks/README.md create mode 100755 .githooks/prepare-commit-msg create mode 100644 render.yaml diff --git a/.githooks/README.md b/.githooks/README.md new file mode 100644 index 000000000..39bc09d22 --- /dev/null +++ b/.githooks/README.md @@ -0,0 +1,15 @@ +# Git Hooks + +These hooks run automatically when using this repository. + +## Setup + +After cloning, run once: + +```bash +git config core.hooksPath .githooks +``` + +## Hooks + +- **prepare-commit-msg**: Strips `Co-authored-by` trailers from commit messages. diff --git a/.githooks/prepare-commit-msg b/.githooks/prepare-commit-msg new file mode 100755 index 000000000..1d0229ee8 --- /dev/null +++ b/.githooks/prepare-commit-msg @@ -0,0 +1,8 @@ +#!/bin/sh +# Strip Co-authored-by trailers - we don't want them in our commits. +COMMIT_MSG_FILE=$1 +if [ -f "$COMMIT_MSG_FILE" ]; then + grep -v '^Co-authored-by:' "$COMMIT_MSG_FILE" > "${COMMIT_MSG_FILE}.tmp" + mv "${COMMIT_MSG_FILE}.tmp" "$COMMIT_MSG_FILE" +fi +exit 0 diff --git a/INSTALL.md b/INSTALL.md index fc9759e9f..e7b8c4bbf 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -130,6 +130,18 @@ For more advanced deployment patterns (NATS workers, search backend, etc.), see: --- +## For Contributors + +After cloning, enable the project's git hooks (strips Co-authored-by from commits): + +```bash +git config core.hooksPath .githooks +``` + +See `.githooks/README.md` for details. + +--- + ## Licensing Notes - **`report_analyst/`** is open-core under the **Reciprocal Public License (RPL)** diff --git a/README.md b/README.md index 482f20ccf..23a6d959a 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,8 @@ **Version 0.8.0-rc (Release Candidate)** +[![Deploy to Render](https://render.com/images/deploy-to-render-button.svg)](https://render.com/deploy?repo=https://github.com/climateandtech/report-analyst) [![Open in Streamlit](https://static.streamlit.io/badges/streamlit_badge_black_white.svg)](https://share.streamlit.io) + Open Sustainability Analyst is the analyst-facing application of the **Open Sustainability Analysis** project by **Climate+Tech**. It helps sustainability and ESG professionals analyze complex sustainability reports with modern AI, while keeping methods transparent and research-based. diff --git a/render.yaml b/render.yaml new file mode 100644 index 000000000..5800878e1 --- /dev/null +++ b/render.yaml @@ -0,0 +1,16 @@ +services: + - type: web + name: report-analyst + runtime: python + plan: free + + buildCommand: pip install -r requirements.txt + startCommand: streamlit run report_analyst/streamlit_app.py --server.port=$PORT --server.address=0.0.0.0 + + envVars: + - key: OPENAI_API_KEY + sync: false + - key: OPENAI_API_MODEL + value: gpt-4o-mini + - key: GOOGLE_API_KEY + sync: false From 6b3f34a2e81dc522ebfc4f76226ac701b7007d1d Mon Sep 17 00:00:00 2001 From: Christian Date: Thu, 19 Feb 2026 21:11:43 +0100 Subject: [PATCH 10/21] Add node_modules to .gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index de53c385f..1b287b70a 100644 --- a/.gitignore +++ b/.gitignore @@ -50,3 +50,4 @@ THEME_SWITCHING.md VERCEL_CHANGES_SUMMARY.md VERCEL_DEPLOYMENT_ASSESSMENT.md VERCEL_MIGRATION_GUIDE.md +node_modules From c8146f970c256b16d448c17c24fbb2132a8c28db Mon Sep 17 00:00:00 2001 From: Christian Date: Fri, 20 Feb 2026 00:27:40 +0100 Subject: [PATCH 11/21] Add black/isort excludes and flake8 config - Exclude venv, node_modules, etc. from black and isort - Add .flake8 to skip venv recursion - Fix pgvector_support formatting --- .flake8 | 4 ++++ pyproject.toml | 11 +++++++++++ .../database/pgvector_support.py | 2 +- 3 files changed, 16 insertions(+), 1 deletion(-) create mode 100644 .flake8 diff --git a/.flake8 b/.flake8 new file mode 100644 index 000000000..22072d256 --- /dev/null +++ b/.flake8 @@ -0,0 +1,4 @@ +[flake8] +max-line-length = 127 +max-complexity = 10 +exclude = .git,__pycache__,venv,venv2,node_modules diff --git a/pyproject.toml b/pyproject.toml index 3efbe5964..050a65857 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,18 @@ [tool.isort] profile = "black" line_length = 127 +skip = [".git", ".hg", "venv", "venv2", "node_modules", "__pycache__"] [tool.black] line-length = 127 target-version = ['py38', 'py39', 'py310', 'py311', 'py312'] +exclude = ''' +/( + \.git + | \.hg + | venv + | venv2 + | node_modules + | __pycache__ +)/ +''' diff --git a/report_analyst_enterprise/database/pgvector_support.py b/report_analyst_enterprise/database/pgvector_support.py index 2cbd3fea3..04563c9b5 100644 --- a/report_analyst_enterprise/database/pgvector_support.py +++ b/report_analyst_enterprise/database/pgvector_support.py @@ -29,7 +29,7 @@ def check_pgvector_available(connection) -> bool: SELECT EXISTS( SELECT 1 FROM pg_extension WHERE extname = 'vector' ) - """ + """ ) ) available = result.fetchone()[0] From 2691483c302fd6630e0c49c76ceae5477e7dff5d Mon Sep 17 00:00:00 2001 From: Christian Date: Fri, 20 Feb 2026 00:40:01 +0100 Subject: [PATCH 12/21] Add Vercel and Heroku deploy buttons and app.json - Add app.json for Heroku one-click deploy - Add Heroku and Vercel badges to README - Add Vercel caveat (Streamlit not supported) --- README.md | 4 +++- app.json | 20 ++++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) create mode 100644 app.json diff --git a/README.md b/README.md index 23a6d959a..98ec1364c 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,9 @@ **Version 0.8.0-rc (Release Candidate)** -[![Deploy to Render](https://render.com/images/deploy-to-render-button.svg)](https://render.com/deploy?repo=https://github.com/climateandtech/report-analyst) [![Open in Streamlit](https://static.streamlit.io/badges/streamlit_badge_black_white.svg)](https://share.streamlit.io) +[![Deploy to Render](https://render.com/images/deploy-to-render-button.svg)](https://render.com/deploy?repo=https://github.com/climateandtech/report-analyst) [![Deploy to Heroku](https://www.herokucdn.com/deploy/button.svg)](https://www.heroku.com/deploy?template=https://github.com/climateandtech/report-analyst) [![Deploy with Vercel](https://vercel.com/button)](https://vercel.com/new/clone?repository-url=https://github.com/climateandtech/report-analyst) [![Open in Streamlit](https://static.streamlit.io/badges/streamlit_badge_black_white.svg)](https://share.streamlit.io) + +Vercel: FastAPI may be deployable; Streamlit is not (see [VERCEL_DEPLOYMENT_ASSESSMENT.md](VERCEL_DEPLOYMENT_ASSESSMENT.md)). Open Sustainability Analyst is the analyst-facing application of the **Open Sustainability Analysis** project by **Climate+Tech**. It helps sustainability and ESG professionals analyze complex sustainability reports with modern AI, while keeping methods transparent and research-based. diff --git a/app.json b/app.json new file mode 100644 index 000000000..8c671a4d9 --- /dev/null +++ b/app.json @@ -0,0 +1,20 @@ +{ + "name": "Open Sustainability Analyst", + "description": "Analyze sustainability reports with AI. Streamlit app for ESG/TCFD analysis.", + "repository": "https://github.com/climateandtech/report-analyst", + "success_url": "/", + "env": { + "OPENAI_API_KEY": { + "description": "Your OpenAI API key for LLM analysis.", + "required": true + }, + "OPENAI_API_MODEL": { + "description": "OpenAI model to use.", + "value": "gpt-4o-mini" + }, + "GOOGLE_API_KEY": { + "description": "Optional Google API key for Gemini.", + "required": false + } + } +} From 230b4c8d4dc1c8da876584dfa8891b4bf345fe0d Mon Sep 17 00:00:00 2001 From: Christian Date: Fri, 20 Feb 2026 00:47:10 +0100 Subject: [PATCH 13/21] Sort deploy badges by license and fix Heroku build - Group deploy by core (RPL) vs API/other modules in README - Add Aptfile and buildpacks to app.json for Heroku - Use module wording only in app.json description --- Aptfile | 3 +++ README.md | 5 ++--- app.json | 7 ++++++- 3 files changed, 11 insertions(+), 4 deletions(-) create mode 100644 Aptfile diff --git a/Aptfile b/Aptfile new file mode 100644 index 000000000..889ce2ea7 --- /dev/null +++ b/Aptfile @@ -0,0 +1,3 @@ +# System dependencies for PyMuPDF (PDF processing) on Heroku +libpoppler-cpp-dev +pkg-config diff --git a/README.md b/README.md index 98ec1364c..ebde83564 100644 --- a/README.md +++ b/README.md @@ -2,9 +2,8 @@ **Version 0.8.0-rc (Release Candidate)** -[![Deploy to Render](https://render.com/images/deploy-to-render-button.svg)](https://render.com/deploy?repo=https://github.com/climateandtech/report-analyst) [![Deploy to Heroku](https://www.herokucdn.com/deploy/button.svg)](https://www.heroku.com/deploy?template=https://github.com/climateandtech/report-analyst) [![Deploy with Vercel](https://vercel.com/button)](https://vercel.com/new/clone?repository-url=https://github.com/climateandtech/report-analyst) [![Open in Streamlit](https://static.streamlit.io/badges/streamlit_badge_black_white.svg)](https://share.streamlit.io) - -Vercel: FastAPI may be deployable; Streamlit is not (see [VERCEL_DEPLOYMENT_ASSESSMENT.md](VERCEL_DEPLOYMENT_ASSESSMENT.md)). +**Deploy (core `report_analyst/`, RPL):** [![Open in Streamlit](https://static.streamlit.io/badges/streamlit_badge_black_white.svg)](https://share.streamlit.io) [![Deploy to Render](https://render.com/images/deploy-to-render-button.svg)](https://render.com/deploy?repo=https://github.com/climateandtech/report-analyst) [![Deploy to Heroku](https://www.herokucdn.com/deploy/button.svg)](https://www.heroku.com/deploy?template=https://github.com/climateandtech/report-analyst) +**Deploy (API and other modules):** [![Deploy with Vercel](https://vercel.com/button)](https://vercel.com/new/clone?repository-url=https://github.com/climateandtech/report-analyst) — `report_analyst_api/` may be deployable; core app is not (see [VERCEL_DEPLOYMENT_ASSESSMENT.md](VERCEL_DEPLOYMENT_ASSESSMENT.md)). Open Sustainability Analyst is the analyst-facing application of the **Open Sustainability Analysis** project by **Climate+Tech**. It helps sustainability and ESG professionals analyze complex sustainability reports with modern AI, while keeping methods transparent and research-based. diff --git a/app.json b/app.json index 8c671a4d9..7a8174954 100644 --- a/app.json +++ b/app.json @@ -1,7 +1,12 @@ { "name": "Open Sustainability Analyst", - "description": "Analyze sustainability reports with AI. Streamlit app for ESG/TCFD analysis.", + "description": "Analyze sustainability reports with AI. Core report_analyst app for ESG/TCFD analysis.", "repository": "https://github.com/climateandtech/report-analyst", + "stack": "heroku-24", + "buildpacks": [ + { "url": "heroku-community/apt" }, + { "url": "heroku/python" } + ], "success_url": "/", "env": { "OPENAI_API_KEY": { From 029f3335cd327ec60060962f650548d5440aea89 Mon Sep 17 00:00:00 2001 From: Christian Date: Fri, 20 Feb 2026 01:06:45 +0100 Subject: [PATCH 14/21] Improve Deploy section and add one-click intro - Move deploy into table above Quick Start with one row per platform - Add intro: one-click install, no card on Streamlit/Render, secrets in Quick Start --- README.md | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index ebde83564..93f353b2b 100644 --- a/README.md +++ b/README.md @@ -2,9 +2,6 @@ **Version 0.8.0-rc (Release Candidate)** -**Deploy (core `report_analyst/`, RPL):** [![Open in Streamlit](https://static.streamlit.io/badges/streamlit_badge_black_white.svg)](https://share.streamlit.io) [![Deploy to Render](https://render.com/images/deploy-to-render-button.svg)](https://render.com/deploy?repo=https://github.com/climateandtech/report-analyst) [![Deploy to Heroku](https://www.herokucdn.com/deploy/button.svg)](https://www.heroku.com/deploy?template=https://github.com/climateandtech/report-analyst) -**Deploy (API and other modules):** [![Deploy with Vercel](https://vercel.com/button)](https://vercel.com/new/clone?repository-url=https://github.com/climateandtech/report-analyst) — `report_analyst_api/` may be deployable; core app is not (see [VERCEL_DEPLOYMENT_ASSESSMENT.md](VERCEL_DEPLOYMENT_ASSESSMENT.md)). - Open Sustainability Analyst is the analyst-facing application of the **Open Sustainability Analysis** project by **Climate+Tech**. It helps sustainability and ESG professionals analyze complex sustainability reports with modern AI, while keeping methods transparent and research-based. @@ -30,6 +27,19 @@ You stay in control of: --- +## Deploy + +You can install and run the app with one click on the platforms below; on **Streamlit Cloud** and **Render** no credit card is required. Follow each platform’s dialogs and enter the secrets (API keys) explained in **Quick Start** below. + +| Platform | Deploys | | +|----------|---------|---| +| **Streamlit Cloud** | Core `report_analyst/` (RPL) | [![Open in Streamlit](https://static.streamlit.io/badges/streamlit_badge_black_white.svg)](https://share.streamlit.io) | +| **Render** | Core `report_analyst/` (RPL) | [![Deploy to Render](https://render.com/images/deploy-to-render-button.svg)](https://render.com/deploy?repo=https://github.com/climateandtech/report-analyst) | +| **Heroku** | Core `report_analyst/` (RPL) | [![Deploy to Heroku](https://www.herokucdn.com/deploy/button.svg)](https://www.heroku.com/deploy?template=https://github.com/climateandtech/report-analyst) | +| **Vercel** | `report_analyst_api/` only; core app not supported | [![Deploy with Vercel](https://vercel.com/button)](https://vercel.com/new/clone?repository-url=https://github.com/climateandtech/report-analyst) | + +--- + ## Quick Start (for Analysts) You need basic command line access, but no deep Python knowledge. From 5901e0afb6923cb18bf6db6e91e4c4f8ec9ca35c Mon Sep 17 00:00:00 2001 From: Christian Date: Fri, 20 Feb 2026 01:11:49 +0100 Subject: [PATCH 15/21] Add Postgres and migrations to Render blueprint - Add report-analyst-db (free) and DATABASE_URL fromDatabase - Set USE_ALEMBIC_MIGRATIONS and USE_POSTGRES_FILE_STORAGE - Add preDeployCommand for alembic upgrade head --- render.yaml | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/render.yaml b/render.yaml index 5800878e1..7c9ee21cf 100644 --- a/render.yaml +++ b/render.yaml @@ -5,12 +5,25 @@ services: plan: free buildCommand: pip install -r requirements.txt + preDeployCommand: python -m alembic upgrade head || echo "Migrations skipped (no Postgres or not enabled)" startCommand: streamlit run report_analyst/streamlit_app.py --server.port=$PORT --server.address=0.0.0.0 envVars: + - key: DATABASE_URL + fromDatabase: + name: report-analyst-db + property: connectionString + - key: USE_ALEMBIC_MIGRATIONS + value: "true" + - key: USE_POSTGRES_FILE_STORAGE + value: "true" - key: OPENAI_API_KEY sync: false - key: OPENAI_API_MODEL value: gpt-4o-mini - key: GOOGLE_API_KEY sync: false + +databases: + - name: report-analyst-db + plan: free From 822bd71318b8569e933028fd45dffe8d62083cb2 Mon Sep 17 00:00:00 2001 From: Christian Date: Fri, 20 Feb 2026 20:17:14 +0100 Subject: [PATCH 16/21] Fix deploy buttons - Render: remove preDeployCommand, run migrations in startCommand for free tier - README: Streamlit badge links to deploy, add Deploy on Streamlit Cloud steps --- README.md | 10 +++++++++- render.yaml | 5 +++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 93f353b2b..c917c5e80 100644 --- a/README.md +++ b/README.md @@ -33,11 +33,19 @@ You can install and run the app with one click on the platforms below; on **Stre | Platform | Deploys | | |----------|---------|---| -| **Streamlit Cloud** | Core `report_analyst/` (RPL) | [![Open in Streamlit](https://static.streamlit.io/badges/streamlit_badge_black_white.svg)](https://share.streamlit.io) | +| **Streamlit Cloud** | Core `report_analyst/` (RPL) | [![Open in Streamlit](https://static.streamlit.io/badges/streamlit_badge_black_white.svg)](https://share.streamlit.io/deploy) | | **Render** | Core `report_analyst/` (RPL) | [![Deploy to Render](https://render.com/images/deploy-to-render-button.svg)](https://render.com/deploy?repo=https://github.com/climateandtech/report-analyst) | | **Heroku** | Core `report_analyst/` (RPL) | [![Deploy to Heroku](https://www.herokucdn.com/deploy/button.svg)](https://www.heroku.com/deploy?template=https://github.com/climateandtech/report-analyst) | | **Vercel** | `report_analyst_api/` only; core app not supported | [![Deploy with Vercel](https://vercel.com/button)](https://vercel.com/new/clone?repository-url=https://github.com/climateandtech/report-analyst) | +### Deploy on Streamlit Cloud + +1. Click the **Open in Streamlit** badge above (or go to [share.streamlit.io](https://share.streamlit.io)). +2. Click **Create app**. +3. Select **Yup, I have an app**, then paste this URL: + `https://github.com/climateandtech/report-analyst/blob/main/report_analyst/streamlit_app.py` +4. In **Advanced settings**, add your secrets (e.g. `OPENAI_API_KEY`, `GOOGLE_API_KEY`) as in **Quick Start** below. + --- ## Quick Start (for Analysts) diff --git a/render.yaml b/render.yaml index 7c9ee21cf..3d27646f4 100644 --- a/render.yaml +++ b/render.yaml @@ -5,8 +5,9 @@ services: plan: free buildCommand: pip install -r requirements.txt - preDeployCommand: python -m alembic upgrade head || echo "Migrations skipped (no Postgres or not enabled)" - startCommand: streamlit run report_analyst/streamlit_app.py --server.port=$PORT --server.address=0.0.0.0 + startCommand: | + python -m alembic upgrade head || true + exec streamlit run report_analyst/streamlit_app.py --server.port=$PORT --server.address=0.0.0.0 envVars: - key: DATABASE_URL From fb204fa309a85073a4609970f7b300fae87607dc Mon Sep 17 00:00:00 2001 From: Christian Date: Fri, 20 Feb 2026 20:29:28 +0100 Subject: [PATCH 17/21] Bump faiss-cpu to 1.12.0 for Render build - 1.10.0 no longer available on PyPI for Render's Python --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 03ae41e8b..3412f5a7d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -36,7 +36,7 @@ dirtyjson==1.0.8 distro==1.9.0 duckduckgo_search==7.5.1 durationpy==0.9 -faiss-cpu==1.10.0 +faiss-cpu==1.12.0 fastapi==0.115.11 ffmpy==0.5.0 filelock==3.17.0 From 897a01f44d907a177794dc12115478685632761f Mon Sep 17 00:00:00 2001 From: Christian Date: Fri, 20 Feb 2026 20:34:12 +0100 Subject: [PATCH 18/21] Pin Python 3.12 for Render build - Avoid Python 3.14 so jiter and other deps use pre-built wheels --- .python-version | 1 + 1 file changed, 1 insertion(+) create mode 100644 .python-version diff --git a/.python-version b/.python-version new file mode 100644 index 000000000..e4fba2183 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.12 From 384d28e769dc279ad85c0a72a080e8a3cbeb0fbe Mon Sep 17 00:00:00 2001 From: Christian Date: Fri, 20 Feb 2026 22:33:15 +0100 Subject: [PATCH 19/21] Add OpenRouter OAuth for one-click AI access without credit card - Allow app startup without API keys when USE_OPENROUTER_OAUTH=true - Add OpenRouter OAuth PKCE flow (connect, callback, key exchange) - Add Connect OpenRouter button in Settings - Support OpenRouter models (openai/gpt-4o-mini, anthropic/claude-3-haiku) - Update render.yaml for deploy with optional keys - Update README with one-click deploy instructions --- README.md | 10 +- render.yaml | 11 +- report_analyst/core/analyzer.py | 75 +++++++++---- report_analyst/core/llm_providers.py | 41 ++++++- report_analyst/core/openrouter_oauth.py | 142 ++++++++++++++++++++++++ report_analyst/streamlit_app.py | 122 ++++++++++++++++++-- 6 files changed, 365 insertions(+), 36 deletions(-) create mode 100644 report_analyst/core/openrouter_oauth.py diff --git a/README.md b/README.md index c917c5e80..944cf8818 100644 --- a/README.md +++ b/README.md @@ -65,7 +65,7 @@ source venv/bin/activate # Windows: venv\Scripts\activate pip install -r requirements.txt ``` -3. **Set your API keys (for LLMs)** +3. **Set your API keys (for LLMs)** – or use OpenRouter one-click on Render Copy the example environment file and add your API keys: @@ -101,6 +101,14 @@ In the web UI you can: For more detailed setup options (API, search backend, jobs), see `INSTALL.md`. +### One-click AI access (Render + OpenRouter) + +When deploying to Render with `USE_OPENROUTER_OAUTH=true`, users get AI access without API keys or credit card: + +1. Deploy to Render (see `render.yaml`). +2. Set `OPENROUTER_CALLBACK_URL` to your app URL. +3. Users open Settings → **Connect OpenRouter** → authorize → done (50 free requests/day). + --- ## Use Cases diff --git a/render.yaml b/render.yaml index 3d27646f4..ec6c68179 100644 --- a/render.yaml +++ b/render.yaml @@ -18,10 +18,17 @@ services: value: "true" - key: USE_POSTGRES_FILE_STORAGE value: "true" + # OpenRouter OAuth: one-click AI access, no credit card + - key: USE_OPENROUTER_OAUTH + value: "true" + - key: OPENAI_API_BASE + value: "https://openrouter.ai/api/v1" + - key: OPENAI_API_MODEL + value: openai/gpt-4o-mini + # RENDER_EXTERNAL_URL is set automatically by Render for OAuth callback + # Optional: override with your own keys for non-OAuth usage - key: OPENAI_API_KEY sync: false - - key: OPENAI_API_MODEL - value: gpt-4o-mini - key: GOOGLE_API_KEY sync: false diff --git a/report_analyst/core/analyzer.py b/report_analyst/core/analyzer.py index 075e3c9c4..2cdaefc9b 100644 --- a/report_analyst/core/analyzer.py +++ b/report_analyst/core/analyzer.py @@ -61,29 +61,53 @@ if not gemini_key: gemini_key = "backend-handles-llm" else: - # Only check for API keys if not using backend LLM - # Check if we need to force the default model based on available keys - if default_model.startswith("gemini-") and not gemini_key: - logger.warning(f"Default model is {default_model} but no GOOGLE_API_KEY is available") - if openai_key: - default_model = "gpt-3.5-turbo-1106" - logger.info(f"Switching default model to {default_model}") - else: - logger.error("No valid API keys available for any models") - raise ValueError("No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY") - elif default_model.startswith("gpt-") and not openai_key: - logger.warning(f"Default model is {default_model} but no OPENAI_API_KEY is available") - if gemini_key: - default_model = "gemini-pro" - logger.info(f"Switching default model to {default_model}") - else: - logger.error("No valid API keys available for any models") - raise ValueError("No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY") + # OpenRouter OAuth mode: allow startup without keys (user connects via OAuth) + use_openrouter_oauth = os.getenv("USE_OPENROUTER_OAUTH", "false").lower() == "true" + if use_openrouter_oauth and not openai_key and not gemini_key: + logger.info( + "OpenRouter OAuth mode - app will start without keys, user connects via OAuth" + ) + default_model = os.getenv("OPENAI_API_MODEL", "openai/gpt-4o-mini") + openai_key = None # Will be set at runtime after OAuth + gemini_key = None + else: + # Only check for API keys if not using OpenRouter OAuth + if default_model.startswith("gemini-") and not gemini_key: + logger.warning( + f"Default model is {default_model} but no GOOGLE_API_KEY is available" + ) + if openai_key: + default_model = "gpt-3.5-turbo-1106" + logger.info(f"Switching default model to {default_model}") + else: + logger.error("No valid API keys available for any models") + raise ValueError( + "No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY" + ) + elif default_model.startswith("gpt-") and not openai_key: + logger.warning( + f"Default model is {default_model} but no OPENAI_API_KEY is available" + ) + if gemini_key: + default_model = "gemini-pro" + logger.info(f"Switching default model to {default_model}") + else: + logger.error("No valid API keys available for any models") + raise ValueError( + "No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY" + ) - # Ensure we have at least one API key for the selected model type - if not openai_key and not gemini_key: - logger.error("No API keys found - set either OPENAI_API_KEY or GOOGLE_API_KEY") - raise ValueError("Set either OPENAI_API_KEY or GOOGLE_API_KEY environment variable") + # Ensure we have at least one API key for the selected model type + if not openai_key and not gemini_key: + logger.error( + "No API keys found - set either OPENAI_API_KEY or GOOGLE_API_KEY" + ) + raise ValueError( + "Set either OPENAI_API_KEY or GOOGLE_API_KEY environment variable" + ) + +# Track OpenRouter OAuth mode for deferred LLM init +use_openrouter_oauth = os.getenv("USE_OPENROUTER_OAUTH", "false").lower() == "true" if not os.getenv("OPENAI_ORGANIZATION"): logger.warning("OPENAI_ORGANIZATION environment variable is not set") @@ -156,6 +180,13 @@ def __init__(self): # Set minimal placeholders for compatibility self.llm = None self.embeddings = None + elif use_openrouter_oauth and not openai_key and not gemini_key: + log_analysis_step( + "OpenRouter OAuth mode - deferring LLM init until user connects", + "info", + ) + self.llm = None + self.embeddings = None else: try: # Initialize LLM with caching using the provider factory diff --git a/report_analyst/core/llm_providers.py b/report_analyst/core/llm_providers.py index 0fd9808c3..918300163 100644 --- a/report_analyst/core/llm_providers.py +++ b/report_analyst/core/llm_providers.py @@ -13,13 +13,23 @@ # Setup logging logger = logging.getLogger(__name__) +# OpenRouter API base - used when OPENAI_API_BASE points to OpenRouter +OPENROUTER_API_BASE = "https://openrouter.ai/api/v1" + + +def _is_openrouter() -> bool: + """Check if we're using OpenRouter (via api_base or explicit flag).""" + api_base = os.getenv("OPENAI_API_BASE") + use_openrouter = os.getenv("USE_OPENROUTER_OAUTH", "false").lower() == "true" + return use_openrouter or (api_base and "openrouter.ai" in api_base) + def get_llm(model_name: str, cache_dir: Optional[str] = None, **kwargs) -> Any: """ Factory function to get LLM implementations based on model name. Args: - model_name: Name of the model to use (e.g., "gpt-4o", "gemini-flash-2.0") + model_name: Name of the model to use (e.g., "gpt-4o", "openai/gpt-4o-mini") cache_dir: Optional directory for LLM response caching **kwargs: Additional keyword arguments to pass to the LLM constructor @@ -30,7 +40,34 @@ def get_llm(model_name: str, cache_dir: Optional[str] = None, **kwargs) -> Any: ValueError: If the API key for the selected model is not available ValueError: If the model type is not supported """ - # OpenAI models + # OpenRouter: uses OpenAI-compatible API with openrouter.ai base + if _is_openrouter(): + api_key = os.getenv("OPENAI_API_KEY") + if not api_key: + logger.error( + f"Cannot initialize OpenRouter model '{model_name}' - " + "OPENAI_API_KEY not set (connect via OpenRouter OAuth)" + ) + raise ValueError( + "Connect OpenRouter to get AI access. No API key available." + ) + api_base = os.getenv("OPENAI_API_BASE", OPENROUTER_API_BASE) + # OpenRouter model IDs: openai/gpt-4o-mini, anthropic/claude-3-haiku, etc. + if not model_name.startswith(("openai/", "anthropic/", "google/", "meta/")): + openrouter_model = ( + f"openai/{model_name}" if "gpt" in model_name else model_name + ) + else: + openrouter_model = model_name + return OpenAI( + model=openrouter_model, + api_key=api_key, + api_base=api_base, + cache_dir=cache_dir, + **kwargs, + ) + + # OpenAI models (direct) if model_name.startswith("gpt-"): api_key = os.getenv("OPENAI_API_KEY") if not api_key: diff --git a/report_analyst/core/openrouter_oauth.py b/report_analyst/core/openrouter_oauth.py new file mode 100644 index 000000000..e12c25c41 --- /dev/null +++ b/report_analyst/core/openrouter_oauth.py @@ -0,0 +1,142 @@ +""" +OpenRouter OAuth PKCE flow for one-click AI access. + +Allows users to connect their OpenRouter account without manually copying API keys. +Uses PKCE (Proof Key for Code Exchange) for secure authorization. +""" + +import base64 +import hashlib +import json +import logging +import os +import secrets +from pathlib import Path +from typing import Optional, Tuple + +import httpx + +logger = logging.getLogger(__name__) + +OPENROUTER_AUTH_URL = "https://openrouter.ai/auth" +OPENROUTER_EXCHANGE_URL = "https://openrouter.ai/api/v1/auth/keys" + + +def _get_state_file() -> Path: + """Get path to OAuth state storage file.""" + storage = Path(__file__).parent.parent.parent / "storage" + storage.mkdir(parents=True, exist_ok=True) + return storage / "openrouter_oauth_state.json" + + +def _load_state() -> dict: + """Load OAuth state from file.""" + path = _get_state_file() + if not path.exists(): + return {} + try: + with open(path) as f: + return json.load(f) + except (json.JSONDecodeError, IOError) as e: + logger.warning(f"Could not load OAuth state: {e}") + return {} + + +def _save_state(data: dict) -> None: + """Save OAuth state to file.""" + path = _get_state_file() + with open(path, "w") as f: + json.dump(data, f, indent=2) + + +def generate_pkce() -> Tuple[str, str]: + """ + Generate PKCE code_verifier and code_challenge (S256). + + Returns: + Tuple of (code_verifier, code_challenge) + """ + code_verifier = secrets.token_urlsafe(32) + digest = hashlib.sha256(code_verifier.encode()).digest() + code_challenge = base64.urlsafe_b64encode(digest).decode().rstrip("=") + return code_verifier, code_challenge + + +def store_oauth_state(state: str, code_verifier: str) -> None: + """Store OAuth state and code_verifier for callback.""" + data = _load_state() + data[state] = { + "code_verifier": code_verifier, + "code_challenge_method": "S256", + } + _save_state(data) + + +def get_and_clear_oauth_state(state: str) -> Optional[dict]: + """Retrieve and remove OAuth state. Returns None if not found.""" + data = _load_state() + entry = data.pop(state, None) + if entry: + _save_state(data) + return entry + + +def get_auth_url(callback_url: str) -> Tuple[str, str]: + """ + Build OpenRouter auth URL and return (url, state). + + Caller must store state with code_verifier before redirecting user. + + Returns: + Tuple of (auth_url, state) + """ + state = secrets.token_urlsafe(16) + code_verifier, code_challenge = generate_pkce() + store_oauth_state(state, code_verifier) + params = { + "callback_url": callback_url, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + } + url = f"{OPENROUTER_AUTH_URL}?callback_url={callback_url}&code_challenge={code_challenge}&code_challenge_method=S256" + return url, state + + +def exchange_code_for_key( + code: str, + code_verifier: str, + code_challenge_method: str = "S256", +) -> str: + """ + Exchange authorization code for OpenRouter API key. + + Args: + code: Authorization code from OpenRouter callback + code_verifier: PKCE code verifier + code_challenge_method: Method used (S256 or plain) + + Returns: + OpenRouter API key + + Raises: + httpx.HTTPStatusError: On exchange failure + ValueError: If response has no key + """ + payload = { + "code": code, + "code_verifier": code_verifier, + "code_challenge_method": code_challenge_method, + } + with httpx.Client() as client: + resp = client.post( + OPENROUTER_EXCHANGE_URL, + json=payload, + headers={"Content-Type": "application/json"}, + timeout=30.0, + ) + resp.raise_for_status() + data = resp.json() + key = data.get("key") + if not key: + raise ValueError("OpenRouter exchange response missing 'key'") + return key diff --git a/report_analyst/streamlit_app.py b/report_analyst/streamlit_app.py index b059e9bdf..ee59f0469 100644 --- a/report_analyst/streamlit_app.py +++ b/report_analyst/streamlit_app.py @@ -77,6 +77,11 @@ def log_analysis_step(message: str, level: str = "info"): create_analysis_dataframes, create_combined_dataframe, ) +from report_analyst.core.openrouter_oauth import ( + exchange_code_for_key, + get_and_clear_oauth_state, + get_auth_url, +) from report_analyst.core.prompt_manager import PromptManager from report_analyst.core.question_loader import get_question_loader @@ -89,14 +94,24 @@ def log_analysis_step(message: str, level: str = "info"): # Define model lists based on available API keys OPENAI_MODELS = ["gpt-4o-mini", "gpt-4o", "gpt-4-turbo", "gpt-3.5-turbo"] +# OpenRouter model IDs (provider/model format) +OPENROUTER_MODELS = [ + "openai/gpt-4o-mini", + "openai/gpt-4o", + "anthropic/claude-3-haiku", + "anthropic/claude-3-sonnet", +] GEMINI_MODELS = ["gemini-1.5-flash", "gemini-1.5-pro"] # Only include models with available API keys LLM_MODELS = OPENAI_MODELS.copy() -# Check for Google API key and add Gemini models if available -if os.getenv("GOOGLE_API_KEY"): +# OpenRouter OAuth mode: use OpenRouter models when enabled +if os.getenv("USE_OPENROUTER_OAUTH", "false").lower() == "true": + LLM_MODELS = OPENROUTER_MODELS.copy() + logger.info("OpenRouter OAuth mode - using OpenRouter model list") +elif os.getenv("GOOGLE_API_KEY"): logger.info("Google API key found - adding Gemini models to available options") LLM_MODELS.extend(GEMINI_MODELS) else: @@ -1298,6 +1313,7 @@ def update_analyzer_parameters(): llm_model = st.session_state.new_llm_model # Validate selected model availability + use_openrouter = os.getenv("USE_OPENROUTER_OAUTH", "false").lower() == "true" if llm_model.startswith("gemini-") and not os.getenv("GOOGLE_API_KEY"): # If somehow a Gemini model was selected but no API key exists logger.error(f"Attempt to use Gemini model '{llm_model}' without API key") @@ -1305,9 +1321,11 @@ def update_analyzer_parameters(): # Reset to default OpenAI model llm_model = OPENAI_MODELS[0] st.session_state.new_llm_model = llm_model - elif llm_model.startswith("gpt-") and not os.getenv("OPENAI_API_KEY"): - logger.error(f"Attempt to use OpenAI model '{llm_model}' without API key") - st.error(f"OPENAI_API_KEY environment variable is not set. OpenAI models will not work correctly.") + elif (llm_model.startswith("gpt-") or llm_model.startswith("openai/") or use_openrouter) and not os.getenv("OPENAI_API_KEY"): + logger.error(f"Attempt to use model '{llm_model}' without API key") + st.error( + "Connect OpenRouter in Settings to get AI access, or set OPENAI_API_KEY." + ) # Update the analyzer with the new parameters try: @@ -1498,6 +1516,39 @@ def main(): st.set_page_config(page_title="Report Analyst", layout="wide") + # OpenRouter OAuth: handle callback and sync key to env + use_openrouter = os.getenv("USE_OPENROUTER_OAUTH", "false").lower() == "true" + if use_openrouter: + os.environ.setdefault("OPENAI_API_BASE", "https://openrouter.ai/api/v1") + # Handle OAuth callback (code + state in URL) + query_params = st.query_params + code = query_params.get("code") + state = query_params.get("state") + if code and state: + try: + oauth_state = get_and_clear_oauth_state(state) + if oauth_state: + api_key = exchange_code_for_key( + code=code, + code_verifier=oauth_state["code_verifier"], + code_challenge_method=oauth_state.get( + "code_challenge_method", "S256" + ), + ) + st.session_state["openrouter_api_key"] = api_key + os.environ["OPENAI_API_KEY"] = api_key + # Clear URL params + st.query_params.clear() + st.success("Connected to OpenRouter! AI access is ready.") + st.rerun() + else: + st.error("Invalid or expired OAuth state. Please try again.") + except Exception as e: + logger.exception("OpenRouter OAuth exchange failed") + st.error(f"Failed to connect: {e}") + elif "openrouter_api_key" in st.session_state: + os.environ["OPENAI_API_KEY"] = st.session_state["openrouter_api_key"] + # Inject Material Icons link tag at the top st.markdown( '', @@ -2734,6 +2785,25 @@ def main(): st.session_state.analyzer = ReportAnalyzer() analyzer = st.session_state.analyzer # Use the stored analyzer + # OpenRouter: refresh LLM when key became available after OAuth + if use_openrouter and os.getenv("OPENAI_API_KEY"): + if analyzer.analyzer.llm is None: + default_model = os.getenv( + "OPENAI_API_MODEL", "openai/gpt-4o-mini" + ) + analyzer.analyzer.update_llm_model(default_model) + # Also init embeddings for OpenRouter + from llama_index.core import Settings + from llama_index.embeddings.openai import OpenAIEmbedding + + analyzer.analyzer.embeddings = OpenAIEmbedding( + api_key=os.getenv("OPENAI_API_KEY"), + api_base="https://openrouter.ai/api/v1", + model_name="text-embedding-ada-002", + embed_batch_size=100, + ) + Settings.embed_model = analyzer.analyzer.embeddings + except Exception as e: st.error(f"Error initializing analyzer: {str(e)}") st.exception(e) @@ -2807,10 +2877,37 @@ def main(): nav_options = ["Upload Report", "Report Analyst", "All Results", "Settings"] nav_page = st.sidebar.radio("", nav_options, key="nav_page", label_visibility="collapsed") - # Show page-specific content based on navigation - if nav_page == "Settings": - st.title("Settings") - st.caption("Configure application settings and integrations") + # Settings section in sidebar (consolidates all integration settings) + st.sidebar.markdown("---") + with st.sidebar.expander("Settings", expanded=False): + # OpenRouter Connect (when USE_OPENROUTER_OAUTH and no key) + if use_openrouter and not os.getenv("OPENAI_API_KEY"): + st.subheader("AI Access") + callback_url = ( + os.getenv("OPENROUTER_CALLBACK_URL") + or os.getenv("RENDER_EXTERNAL_URL", "") + ).rstrip("/") + if callback_url: + auth_url, _state = get_auth_url(callback_url) + st.link_button( + "Connect OpenRouter", + auth_url, + type="primary", + help="One-click AI access, no credit card. Free: 50 req/day.", + ) + st.caption( + "One-click AI access, no credit card. Free tier: 50 requests/day." + ) + else: + st.warning( + "Set OPENROUTER_CALLBACK_URL to your app URL to enable " + "Connect OpenRouter. (On Render, RENDER_EXTERNAL_URL is set automatically.)" + ) + st.divider() + # Show Enterprise Mode status at the top + use_s3_upload = st.session_state.get("use_s3_upload", False) + if use_s3_upload and BACKEND_INTEGRATION_AVAILABLE: + st.caption("Enterprise mode") # Open Source Modules Section st.header("Open Source Modules") @@ -3128,6 +3225,13 @@ def main(): st.info("PostgreSQL file storage requires a PostgreSQL database. Currently using SQLite.") st.caption("Files are stored in local temp directory") + # OpenRouter: show connect banner when no key + if use_openrouter and not os.getenv("OPENAI_API_KEY"): + st.info( + "**Get AI access** – Open Settings in the sidebar and click " + '"Connect OpenRouter" for one-click access (no credit card, 50 free requests/day).' + ) + # Show page-specific content based on navigation if nav_page == "Report Analyst": st.title("Report Analyst") From d524dbac9255de6d0b046d639afab697e1def50b Mon Sep 17 00:00:00 2001 From: Christian Date: Fri, 20 Feb 2026 22:35:05 +0100 Subject: [PATCH 20/21] Use RENDER_EXTERNAL_URL for OpenRouter callback (auto on Render) --- README.md | 5 +++-- render.yaml | 2 +- report_analyst/streamlit_app.py | 1 + 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 944cf8818..a88ed0a9f 100644 --- a/README.md +++ b/README.md @@ -106,8 +106,9 @@ For more detailed setup options (API, search backend, jobs), see `INSTALL.md`. When deploying to Render with `USE_OPENROUTER_OAUTH=true`, users get AI access without API keys or credit card: 1. Deploy to Render (see `render.yaml`). -2. Set `OPENROUTER_CALLBACK_URL` to your app URL. -3. Users open Settings → **Connect OpenRouter** → authorize → done (50 free requests/day). +2. Users open Settings → **Connect OpenRouter** → authorize → done (50 free requests/day). + +(Render sets `RENDER_EXTERNAL_URL` automatically for the OAuth callback.) --- diff --git a/render.yaml b/render.yaml index ec6c68179..c215e9921 100644 --- a/render.yaml +++ b/render.yaml @@ -25,7 +25,7 @@ services: value: "https://openrouter.ai/api/v1" - key: OPENAI_API_MODEL value: openai/gpt-4o-mini - # RENDER_EXTERNAL_URL is set automatically by Render for OAuth callback + # RENDER_EXTERNAL_URL is set automatically by Render for the callback # Optional: override with your own keys for non-OAuth usage - key: OPENAI_API_KEY sync: false diff --git a/report_analyst/streamlit_app.py b/report_analyst/streamlit_app.py index ee59f0469..9e8cc766f 100644 --- a/report_analyst/streamlit_app.py +++ b/report_analyst/streamlit_app.py @@ -2883,6 +2883,7 @@ def main(): # OpenRouter Connect (when USE_OPENROUTER_OAUTH and no key) if use_openrouter and not os.getenv("OPENAI_API_KEY"): st.subheader("AI Access") + # Use OPENROUTER_CALLBACK_URL or Render's automatic RENDER_EXTERNAL_URL callback_url = ( os.getenv("OPENROUTER_CALLBACK_URL") or os.getenv("RENDER_EXTERNAL_URL", "") From a7970d576fc4f6f8fa670cf598f3f54f2d2d9b51 Mon Sep 17 00:00:00 2001 From: Christian Date: Sat, 14 Mar 2026 00:56:17 +0100 Subject: [PATCH 21/21] Fix black formatting and remove unused code in OpenRouter module --- report_analyst/core/analyzer.py | 28 +++++++------------------ report_analyst/core/llm_providers.py | 11 +++------- report_analyst/core/openrouter_oauth.py | 6 ------ report_analyst/streamlit_app.py | 25 +++++++--------------- 4 files changed, 18 insertions(+), 52 deletions(-) diff --git a/report_analyst/core/analyzer.py b/report_analyst/core/analyzer.py index 2cdaefc9b..c59c07519 100644 --- a/report_analyst/core/analyzer.py +++ b/report_analyst/core/analyzer.py @@ -64,47 +64,33 @@ # OpenRouter OAuth mode: allow startup without keys (user connects via OAuth) use_openrouter_oauth = os.getenv("USE_OPENROUTER_OAUTH", "false").lower() == "true" if use_openrouter_oauth and not openai_key and not gemini_key: - logger.info( - "OpenRouter OAuth mode - app will start without keys, user connects via OAuth" - ) + logger.info("OpenRouter OAuth mode - app will start without keys, user connects via OAuth") default_model = os.getenv("OPENAI_API_MODEL", "openai/gpt-4o-mini") openai_key = None # Will be set at runtime after OAuth gemini_key = None else: # Only check for API keys if not using OpenRouter OAuth if default_model.startswith("gemini-") and not gemini_key: - logger.warning( - f"Default model is {default_model} but no GOOGLE_API_KEY is available" - ) + logger.warning(f"Default model is {default_model} but no GOOGLE_API_KEY is available") if openai_key: default_model = "gpt-3.5-turbo-1106" logger.info(f"Switching default model to {default_model}") else: logger.error("No valid API keys available for any models") - raise ValueError( - "No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY" - ) + raise ValueError("No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY") elif default_model.startswith("gpt-") and not openai_key: - logger.warning( - f"Default model is {default_model} but no OPENAI_API_KEY is available" - ) + logger.warning(f"Default model is {default_model} but no OPENAI_API_KEY is available") if gemini_key: default_model = "gemini-pro" logger.info(f"Switching default model to {default_model}") else: logger.error("No valid API keys available for any models") - raise ValueError( - "No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY" - ) + raise ValueError("No valid API keys found. Set either OPENAI_API_KEY or GOOGLE_API_KEY") # Ensure we have at least one API key for the selected model type if not openai_key and not gemini_key: - logger.error( - "No API keys found - set either OPENAI_API_KEY or GOOGLE_API_KEY" - ) - raise ValueError( - "Set either OPENAI_API_KEY or GOOGLE_API_KEY environment variable" - ) + logger.error("No API keys found - set either OPENAI_API_KEY or GOOGLE_API_KEY") + raise ValueError("Set either OPENAI_API_KEY or GOOGLE_API_KEY environment variable") # Track OpenRouter OAuth mode for deferred LLM init use_openrouter_oauth = os.getenv("USE_OPENROUTER_OAUTH", "false").lower() == "true" diff --git a/report_analyst/core/llm_providers.py b/report_analyst/core/llm_providers.py index 918300163..f9961bf79 100644 --- a/report_analyst/core/llm_providers.py +++ b/report_analyst/core/llm_providers.py @@ -45,18 +45,13 @@ def get_llm(model_name: str, cache_dir: Optional[str] = None, **kwargs) -> Any: api_key = os.getenv("OPENAI_API_KEY") if not api_key: logger.error( - f"Cannot initialize OpenRouter model '{model_name}' - " - "OPENAI_API_KEY not set (connect via OpenRouter OAuth)" - ) - raise ValueError( - "Connect OpenRouter to get AI access. No API key available." + f"Cannot initialize OpenRouter model '{model_name}' - " "OPENAI_API_KEY not set (connect via OpenRouter OAuth)" ) + raise ValueError("Connect OpenRouter to get AI access. No API key available.") api_base = os.getenv("OPENAI_API_BASE", OPENROUTER_API_BASE) # OpenRouter model IDs: openai/gpt-4o-mini, anthropic/claude-3-haiku, etc. if not model_name.startswith(("openai/", "anthropic/", "google/", "meta/")): - openrouter_model = ( - f"openai/{model_name}" if "gpt" in model_name else model_name - ) + openrouter_model = f"openai/{model_name}" if "gpt" in model_name else model_name else: openrouter_model = model_name return OpenAI( diff --git a/report_analyst/core/openrouter_oauth.py b/report_analyst/core/openrouter_oauth.py index e12c25c41..deffb55e2 100644 --- a/report_analyst/core/openrouter_oauth.py +++ b/report_analyst/core/openrouter_oauth.py @@ -9,7 +9,6 @@ import hashlib import json import logging -import os import secrets from pathlib import Path from typing import Optional, Tuple @@ -93,11 +92,6 @@ def get_auth_url(callback_url: str) -> Tuple[str, str]: state = secrets.token_urlsafe(16) code_verifier, code_challenge = generate_pkce() store_oauth_state(state, code_verifier) - params = { - "callback_url": callback_url, - "code_challenge": code_challenge, - "code_challenge_method": "S256", - } url = f"{OPENROUTER_AUTH_URL}?callback_url={callback_url}&code_challenge={code_challenge}&code_challenge_method=S256" return url, state diff --git a/report_analyst/streamlit_app.py b/report_analyst/streamlit_app.py index 9e8cc766f..ece80181d 100644 --- a/report_analyst/streamlit_app.py +++ b/report_analyst/streamlit_app.py @@ -1321,11 +1321,11 @@ def update_analyzer_parameters(): # Reset to default OpenAI model llm_model = OPENAI_MODELS[0] st.session_state.new_llm_model = llm_model - elif (llm_model.startswith("gpt-") or llm_model.startswith("openai/") or use_openrouter) and not os.getenv("OPENAI_API_KEY"): + elif (llm_model.startswith("gpt-") or llm_model.startswith("openai/") or use_openrouter) and not os.getenv( + "OPENAI_API_KEY" + ): logger.error(f"Attempt to use model '{llm_model}' without API key") - st.error( - "Connect OpenRouter in Settings to get AI access, or set OPENAI_API_KEY." - ) + st.error("Connect OpenRouter in Settings to get AI access, or set OPENAI_API_KEY.") # Update the analyzer with the new parameters try: @@ -1531,9 +1531,7 @@ def main(): api_key = exchange_code_for_key( code=code, code_verifier=oauth_state["code_verifier"], - code_challenge_method=oauth_state.get( - "code_challenge_method", "S256" - ), + code_challenge_method=oauth_state.get("code_challenge_method", "S256"), ) st.session_state["openrouter_api_key"] = api_key os.environ["OPENAI_API_KEY"] = api_key @@ -2788,9 +2786,7 @@ def main(): # OpenRouter: refresh LLM when key became available after OAuth if use_openrouter and os.getenv("OPENAI_API_KEY"): if analyzer.analyzer.llm is None: - default_model = os.getenv( - "OPENAI_API_MODEL", "openai/gpt-4o-mini" - ) + default_model = os.getenv("OPENAI_API_MODEL", "openai/gpt-4o-mini") analyzer.analyzer.update_llm_model(default_model) # Also init embeddings for OpenRouter from llama_index.core import Settings @@ -2884,10 +2880,7 @@ def main(): if use_openrouter and not os.getenv("OPENAI_API_KEY"): st.subheader("AI Access") # Use OPENROUTER_CALLBACK_URL or Render's automatic RENDER_EXTERNAL_URL - callback_url = ( - os.getenv("OPENROUTER_CALLBACK_URL") - or os.getenv("RENDER_EXTERNAL_URL", "") - ).rstrip("/") + callback_url = (os.getenv("OPENROUTER_CALLBACK_URL") or os.getenv("RENDER_EXTERNAL_URL", "")).rstrip("/") if callback_url: auth_url, _state = get_auth_url(callback_url) st.link_button( @@ -2896,9 +2889,7 @@ def main(): type="primary", help="One-click AI access, no credit card. Free: 50 req/day.", ) - st.caption( - "One-click AI access, no credit card. Free tier: 50 requests/day." - ) + st.caption("One-click AI access, no credit card. Free tier: 50 requests/day.") else: st.warning( "Set OPENROUTER_CALLBACK_URL to your app URL to enable "