Merge pull request #64 from waldronlab/refactor/improve-maintainability

RonaldRonnie · web-flow · commit ea46e4229eed · 2026-01-08T10:44:44.000+03:00
Refactor: improve code maintainability and error handling
diff --git a/README.md b/README.md
diff --git a/app/api/app.py b/app/api/app.py
@@ -1,17 +1,12 @@
 """FastAPI application for BioAnalyzer backend API."""
 
+import logging
+import os
+import traceback
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi.exceptions import RequestValidationError
-import logging
-import os
-import sys
-import traceback
-
-sys.path.append(
-    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-)
 
 from app.api.routers import (
     bugsigdb_analysis,
@@ -119,19 +114,20 @@ async def validation_exception_handler(request: Request, exc: RequestValidationE
 
 @app.exception_handler(Exception)
 async def global_exception_handler(request: Request, exc: Exception):
-    """Handle unexpected exceptions."""
+    """Handle unexpected exceptions with credential masking."""
     from app.utils.credential_masking import mask_exception_message, mask_string
 
-    # Mask any credentials in exception message and traceback
     safe_exc_msg = mask_exception_message(exc)
     safe_traceback = mask_string(traceback.format_exc())
+    request_id = getattr(request.state, "request_id", None)
 
     logger.error(
         f"Unhandled exception: {safe_exc_msg}\n"
         f"Traceback: {safe_traceback}\n"
-        f"Request ID: {getattr(request.state, 'request_id', None)}"
+        f"Request ID: {request_id}"
     )
 
+    # Hide internal details in production
     if ENVIRONMENT == "production":
         detail = "An internal error occurred. Please try again later."
     else:
@@ -142,7 +138,7 @@ async def global_exception_handler(request: Request, exc: Exception):
         content={
             "error": "Internal Server Error",
             "detail": detail,
-            "request_id": getattr(request.state, "request_id", None),
+            "request_id": request_id,
         },
     )
 
diff --git a/app/models/unified_qa.py b/app/models/unified_qa.py
@@ -28,7 +28,11 @@
 
 
 class UnifiedQA:
-    """Unified QA system supporting multiple LLM providers."""
+    """Wrapper around LLM providers for question answering.
+
+    Tries LiteLLM first (supports multiple providers), falls back to Paper-QA,
+    then GeminiQA if those fail. Auto-detects provider from available API keys.
+    """
 
     def __init__(
         self,
@@ -38,7 +42,15 @@ def __init__(
         gemini_api_key: Optional[str] = None,
         use_paperqa: bool = True,
     ):
-        """Initialize QA system with specified provider and model."""
+        """Set up QA system with the specified provider.
+
+        Args:
+            provider: LLM provider (gemini, openai, anthropic, ollama). Auto-detects if None.
+            model: Specific model name. Uses provider default if None.
+            use_gemini: Deprecated. Use provider='gemini' instead.
+            gemini_api_key: Override GEMINI_API_KEY env var.
+            use_paperqa: Try Paper-QA as fallback (default: True).
+        """
         if use_gemini is not None:
             logger.warning(
                 "use_gemini parameter is deprecated. Use provider='gemini' instead."
diff --git a/app/services/advanced_rag.py b/app/services/advanced_rag.py
@@ -1,4 +1,9 @@
-"""Advanced RAG service with contextual summarization and re-ranking."""
+"""RAG pipeline for better field extraction accuracy.
+
+Re-ranks text chunks by relevance to each field query, then generates
+query-aware summaries. More accurate than simple analysis but slower and
+more expensive (multiple LLM calls per field).
+"""
 
 import logging
 from typing import List, Dict, Optional, Tuple, Any
@@ -26,7 +31,11 @@
 
 
 class AdvancedRAGService:
-    """Advanced RAG service combining contextual summarization and chunk re-ranking."""
+    """RAG pipeline: re-rank chunks, then summarize before querying LLM.
+
+    This is what makes v2 API more accurate than v1. The tradeoff is speed
+    and cost - expect 2-3x more LLM API calls per field.
+    """
 
     def __init__(
         self,
@@ -38,7 +47,17 @@ def __init__(
         max_sources: Optional[int] = None,
         use_10_scale: bool = True,
     ):
-        """Initialize advanced RAG service."""
+        """Set up RAG service with configuration.
+
+        Args:
+            summary_provider: LLM provider for summarization (default: auto-detect).
+            summary_model: Model for summarization (default: provider default).
+            rerank_method: keyword|llm|hybrid (default: hybrid).
+            cache_dir: Where to cache summaries (default: cache/).
+            evidence_k: Initial chunks to retrieve before re-ranking.
+            max_sources: Max chunks to use after re-ranking.
+            use_10_scale: Use 0-10 relevance scale instead of 0-1.
+        """
         summary_config = SummarizationConfig(
             summary_length=RAG_SUMMARY_LENGTH,
             quality=RAG_SUMMARY_QUALITY,
diff --git a/app/services/bugsigdb_analyzer.py b/app/services/bugsigdb_analyzer.py
@@ -1,4 +1,8 @@
-"""Analysis service for extracting BugSigDB fields from papers."""
+"""Extracts six BugSigDB fields from papers using LLMs.
+
+This is the main analysis service. It orchestrates paper retrieval, text preparation,
+and field extraction. Results are cached to avoid redundant API calls.
+"""
 
 import logging
 from typing import Dict, Optional, List
@@ -76,7 +80,14 @@ def get_cache_manager():
 
 
 async def analyze_paper_simple(pmid: str) -> Optional[Dict]:
-    """Analyze paper and extract BugSigDB fields."""
+    """Extract BugSigDB fields from a paper.
+
+    Uses v1 API flow: direct LLM queries per field. Fast but less accurate than RAG.
+    Checks cache first, then fetches from PubMed if needed.
+
+    Returns:
+        Dict with field results, or None if paper can't be retrieved.
+    """
     try:
         cache_manager = get_cache_manager()
         pubmed_retriever = get_pubmed_retriever()
diff --git a/app/services/data_retrieval.py b/app/services/data_retrieval.py
@@ -27,7 +27,11 @@ class PubMedRetrieverError(Exception):
 
 
 class PubMedRetriever:
-    """Retrieves paper metadata and full text from PubMed using NCBI E-Utilities API."""
+    """Fetches papers from PubMed/PMC via NCBI E-utilities.
+
+    Handles rate limiting (NCBI requires 0.34s between requests) and retries.
+    Caches results to avoid hitting API limits.
+    """
 
     BASE_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
     DEFAULT_TIMEOUT = 10
@@ -36,22 +40,27 @@ class PubMedRetriever:
     def __init__(
         self, api_key: Optional[str] = None, email: str = "bioanalyzer@example.com"
     ):
-        """Initialize PubMed retriever with API key and email."""
+        """Create retriever instance.
+
+        Args:
+            api_key: NCBI API key (optional but recommended for higher rate limits).
+            email: Contact email for NCBI requests (required by their ToS).
+        """
         self.api_key = api_key
         self.email = email
         self.session = self._create_session()
         self._verify_connectivity()
 
     def _create_session(self) -> requests.Session:
-        """Create a configured requests session."""
+        """Set up HTTP session with proper User-Agent header."""
         session = requests.Session()
         session.headers.update(
             {"User-Agent": f"BioAnalyzer/1.0 (contact: {self.email})"}
         )
         return session
 
     def _verify_connectivity(self, retries: int = 3) -> None:
-        """Test NCBI E-utilities reachability on startup with retries."""
+        """Verify NCBI E-utilities connectivity on startup."""
         test_url = f"{self.BASE_URL}/esearch.fcgi"
         params = {"db": "pubmed", "term": "cancer", "retmax": 1}
         for attempt in range(retries):
@@ -235,12 +244,19 @@ def fetch_paper_metadata(self, pmid: str) -> Dict[str, Any]:
                     fields["authors"] = [a.text for a in item.findall("Item") if a.text]
             fields.setdefault("abstract", "")
             return fields
-        except Exception as e:
+        except (ElementTree.ParseError, AttributeError, ValueError) as e:
             safe_error = mask_exception_message(e)
             logger.error(
                 f"Error parsing fallback esummary for PMID {pmid}: {safe_error}"
             )
             return {"error": "Fallback retrieval failed."}
+        except Exception as e:
+            # Catch-all for unexpected parsing errors
+            safe_error = mask_exception_message(e)
+            logger.error(
+                f"Unexpected error parsing fallback esummary for PMID {pmid}: {safe_error}"
+            )
+            return {"error": "Fallback retrieval failed."}
 
     def search(self, query: str, max_results: int = 10) -> List[str]:
         xml_data = self._make_request(
@@ -271,9 +287,17 @@ def get_pmc_fulltext(self, pmid: str) -> str:
 
             return self._get_pmc_fulltext_by_id(pmc_id)
 
+        except (requests.exceptions.RequestException, PubMedRetrieverError) as e:
+            safe_error = mask_exception_message(e)
+            logger.warning(f"Network error retrieving full text for PMID {pmid}: {safe_error}")
+            return ""
+        except (ElementTree.ParseError, ValueError) as e:
+            safe_error = mask_exception_message(e)
+            logger.warning(f"Parse error retrieving full text for PMID {pmid}: {safe_error}")
+            return ""
         except Exception as e:
             safe_error = mask_exception_message(e)
-            logger.warning(f"Error retrieving full text for PMID {pmid}: {safe_error}")
+            logger.warning(f"Unexpected error retrieving full text for PMID {pmid}: {safe_error}")
             return ""
 
     def _get_pmc_id_from_pmid(self, pmid: str) -> Optional[str]:
@@ -305,9 +329,17 @@ def _get_pmc_id_from_pmid(self, pmid: str) -> Optional[str]:
                                 return pmc_id
             return None
 
+        except (requests.exceptions.RequestException, PubMedRetrieverError) as e:
+            safe_error = mask_exception_message(e)
+            logger.warning(f"Network error getting PMC ID for PMID {pmid}: {safe_error}")
+            return None
+        except (ElementTree.ParseError, AttributeError) as e:
+            safe_error = mask_exception_message(e)
+            logger.warning(f"Parse error getting PMC ID for PMID {pmid}: {safe_error}")
+            return None
         except Exception as e:
             safe_error = mask_exception_message(e)
-            logger.warning(f"Error getting PMC ID for PMID {pmid}: {safe_error}")
+            logger.warning(f"Unexpected error getting PMC ID for PMID {pmid}: {safe_error}")
             return None
 
     def _get_pmc_fulltext_by_id(self, pmc_id: str) -> str:
@@ -351,20 +383,25 @@ def _get_pmc_fulltext_by_id(self, pmc_id: str) -> str:
 
             return "\n\n".join(full_text_parts)
 
+        except (requests.exceptions.RequestException, PubMedRetrieverError) as e:
+            safe_error = mask_exception_message(e)
+            logger.warning(f"Network error retrieving PMC full text for {pmc_id}: {safe_error}")
+            return ""
+        except (ElementTree.ParseError, AttributeError, ValueError) as e:
+            safe_error = mask_exception_message(e)
+            logger.warning(f"Parse error retrieving PMC full text for {pmc_id}: {safe_error}")
+            return ""
         except Exception as e:
             safe_error = mask_exception_message(e)
-            logger.warning(f"Error retrieving PMC full text for {pmc_id}: {safe_error}")
+            logger.warning(f"Unexpected error retrieving PMC full text for {pmc_id}: {safe_error}")
             return ""
 
     async def get_pmc_fulltext_async(self, pmid: str) -> str:
         """Async wrapper for PMC full text retrieval."""
         return await asyncio.to_thread(self.get_pmc_fulltext, pmid)
 
     def get_full_paper_data(self, pmid: str) -> Dict[str, Any]:
-        """
-        Retrieve complete paper data including metadata and full text.
-        This is the main method for comprehensive paper retrieval.
-        """
+        """Retrieve complete paper data including metadata and full text."""
         try:
             logger.info(f"Retrieving full paper data for PMID: {pmid}")
 
@@ -392,14 +429,14 @@ def get_full_paper_data(self, pmid: str) -> Dict[str, Any]:
             logger.info(f"Successfully retrieved paper data for PMID: {pmid}")
             return paper_data
 
-        except Exception as e:
+        except (requests.exceptions.RequestException, PubMedRetrieverError) as e:
             safe_error = mask_exception_message(e)
             logger.error(
-                f"Error retrieving full paper data for PMID {pmid}: {safe_error}"
+                f"Network error retrieving full paper data for PMID {pmid}: {safe_error}"
             )
             return {
                 "pmid": pmid,
-                "error": f"Failed to retrieve paper data: {str(e)}",
+                "error": "Failed to retrieve paper data due to network error",
                 "title": "",
                 "abstract": "",
                 "journal": "",
@@ -420,19 +457,33 @@ async def fetch_metadata():
                 return await asyncio.wait_for(
                     self.get_paper_metadata_async(pmid), timeout=6
                 )
+            except asyncio.TimeoutError:
+                logger.error(f"Timeout fetching metadata for PMID {pmid}")
+                return {}
+            except (requests.exceptions.RequestException, PubMedRetrieverError) as e:
+                safe_error = mask_exception_message(e)
+                logger.error(f"Network error fetching metadata for PMID {pmid}: {safe_error}")
+                return {}
             except Exception as e:
                 safe_error = mask_exception_message(e)
-                logger.error(f"Metadata fetch error for PMID {pmid}: {safe_error}")
+                logger.error(f"Unexpected error fetching metadata for PMID {pmid}: {safe_error}")
                 return {}
 
         async def fetch_fulltext():
             try:
                 return await asyncio.wait_for(
                     self.get_pmc_fulltext_async(pmid), timeout=8
                 )
+            except asyncio.TimeoutError:
+                logger.warning(f"Timeout fetching full text for PMID {pmid}")
+                return ""
+            except (requests.exceptions.RequestException, PubMedRetrieverError) as e:
+                safe_error = mask_exception_message(e)
+                logger.warning(f"Network error fetching full text for PMID {pmid}: {safe_error}")
+                return ""
             except Exception as e:
                 safe_error = mask_exception_message(e)
-                logger.warning(f"Full text fetch error for PMID {pmid}: {safe_error}")
+                logger.warning(f"Unexpected error fetching full text for PMID {pmid}: {safe_error}")
                 return ""
 
         if USE_FULLTEXT:
diff --git a/app/utils/config.py b/app/utils/config.py
@@ -64,11 +64,15 @@ def validate_gemini_key():
     try:
         genai.configure(api_key=GEMINI_API_KEY)
         return True
-    except Exception as e:
-        # Mask any potential credentials in error message
+    except (ValueError, AttributeError) as e:
         safe_error = mask_exception_message(e)
         print(f"Gemini API key validation failed: {safe_error}")
         return False
+    except Exception as e:
+        # Catch-all for unexpected errors during configuration
+        safe_error = mask_exception_message(e)
+        print(f"Unexpected error validating Gemini key: {safe_error}")
+        return False
 
 
 def validate_env_vars():
@@ -82,20 +86,20 @@ def validate_env_vars():
     if not GEMINI_API_KEY:
         missing_vars.append("GEMINI_API_KEY")
 
-    # Check if at least one AI model is available
+    # At least one LLM provider must be available
     if not AVAILABLE_MODELS:
         missing_vars.append("GEMINI_API_KEY")
 
     if missing_vars:
         print(
-            f"Warning: The following environment variables are missing: {', '.join(missing_vars)}"
+            f"Warning: Missing environment variables: {', '.join(missing_vars)}"
         )
-        print("Please set them in your .env file or environment.")
+        print("Set them in your .env file or environment.")
 
     return len(missing_vars) == 0
 
 
-# Call validation when module is imported
+# Validate on import
 validate_env_vars()
 
 
@@ -209,10 +213,9 @@ def check_required_vars():
 
 
 def setup_logging():
-    """Setup comprehensive logging configuration with file rotation.
+    """Configure logging with file rotation.
 
-    Handles permission errors gracefully, falling back to console-only logging
-    if file handlers cannot be created (e.g., during testing).
+    Falls back to console-only logging if file handlers can't be created.
     """
     import logging.handlers
 
diff --git a/docs/README.md b/docs/README.md
diff --git a/main.py b/main.py
