Merge pull request #65 from waldronlab/feature/add-type-hints

Add comprehensive type hints to core functions

Author: RonaldRonnie

app/api/app.py

@@ -3,10 +3,12 @@
 import logging
 import os
 import traceback
+from typing import Dict, Any
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi.exceptions import RequestValidationError
+from app.api.models.api_models import HealthResponse
 
 from app.api.routers import (
     bugsigdb_analysis,
@@ -78,7 +80,7 @@
 
 
 @app.get("/")
-async def root():
+async def root() -> Dict[str, str]:
     """API root endpoint."""
     return {
         "message": "BioAnalyzer Backend API",
@@ -91,15 +93,17 @@ async def root():
 
 
 @app.get("/health")
-async def health_check():
+async def health_check() -> HealthResponse:
     """Health check endpoint."""
     from app.api.routers.system import health_check as system_health_check
 
     return await system_health_check()
 
 
 @app.exception_handler(RequestValidationError)
-async def validation_exception_handler(request: Request, exc: RequestValidationError):
+async def validation_exception_handler(
+    request: Request, exc: RequestValidationError
+) -> JSONResponse:
     """Handle request validation errors."""
     logger.warning(f"Validation error: {exc.errors()}")
     return JSONResponse(
@@ -113,7 +117,9 @@ async def validation_exception_handler(request: Request, exc: RequestValidationE
 
 
 @app.exception_handler(Exception)
-async def global_exception_handler(request: Request, exc: Exception):
+async def global_exception_handler(
+    request: Request, exc: Exception
+) -> JSONResponse:
     """Handle unexpected exceptions with credential masking."""
     from app.utils.credential_masking import mask_exception_message, mask_string
 

app/api/routers/bugsigdb_analysis.py

@@ -2,6 +2,7 @@
 
 from fastapi import APIRouter, HTTPException
 import logging
+from typing import Dict, Any
 from app.utils.credential_masking import mask_exception_message
 from app.api.utils.constants import ESSENTIAL_FIELDS_INFO, STATUS_VALUES
 from app.services.bugsigdb_analyzer import analyze_paper_simple
@@ -10,7 +11,7 @@
 router = APIRouter(prefix="/api/v1", tags=["BugSigDB Analysis"])
 
 
-async def _run_analysis(pmid: str):
+async def _run_analysis(pmid: str) -> Dict[str, Any]:
     """Run analysis with validated PMID."""
     from app.api.utils.api_utils import validate_pmid
 
@@ -22,15 +23,15 @@ async def _run_analysis(pmid: str):
     return result
 
 
-def _handle_analysis_error(pmid: str, e: Exception):
+def _handle_analysis_error(pmid: str, e: Exception) -> None:
     """Handle analysis errors consistently."""
     safe_error = mask_exception_message(e)
     logger.error(f"Error in analysis for PMID {pmid}: {safe_error}")
     raise HTTPException(status_code=500, detail=f"Analysis error: {str(e)}")
 
 
 @router.get("/analyze/{pmid}")
-async def analyze_paper(pmid: str):
+async def analyze_paper(pmid: str) -> Dict[str, Any]:
     """Analyze paper for BugSigDB fields."""
     try:
         return await _run_analysis(pmid)
@@ -41,7 +42,7 @@ async def analyze_paper(pmid: str):
 
 
 @router.post("/analyze/{pmid}")
-async def analyze_paper_post(pmid: str):
+async def analyze_paper_post(pmid: str) -> Dict[str, Any]:
     """Analyze paper for BugSigDB fields (POST method)."""
     try:
         return await _run_analysis(pmid)

app/api/routers/system.py

@@ -34,10 +34,10 @@
 logger = logging.getLogger(__name__)
 router = APIRouter(prefix="/api/v1", tags=["System"])
 
-_unified_qa = None
-_pubmed_retriever = None
+_unified_qa: Optional[UnifiedQA] = None
+_pubmed_retriever: Optional[PubMedRetriever] = None
 
-def get_unified_qa():
+def get_unified_qa() -> Optional[UnifiedQA]:
     """Get or initialize UnifiedQA instance."""
     global _unified_qa
     if _unified_qa is None:
@@ -55,7 +55,7 @@ def get_unified_qa():
 unified_qa = None  # Will be initialized on first use via get_unified_qa()
 
 
-def get_pubmed_retriever():
+def get_pubmed_retriever() -> Optional[PubMedRetriever]:
     """Get or initialize PubMedRetriever instance."""
     global _pubmed_retriever
     if _pubmed_retriever is None:
@@ -69,15 +69,15 @@ def get_pubmed_retriever():
 
 
 @router.get("/")
-async def root():
+async def root() -> Any:
     """Redirect to the frontend application."""
     from fastapi.responses import RedirectResponse
 
     return RedirectResponse(url="/static/index.html")
 
 
 @router.get("/health")
-async def health_check():
+async def health_check() -> HealthResponse:
     """Health check endpoint to verify service is running."""
     try:
         current_time = get_current_timestamp()
@@ -92,7 +92,7 @@ async def health_check():
 
 
 @router.get("/config")
-async def get_config():
+async def get_config() -> ConfigResponse:
     """
     **Get configuration settings for the frontend.**
 

app/api/utils/api_utils.py

@@ -2,7 +2,7 @@
 
 import re
 import logging
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Any
 from datetime import datetime
 import pytz
 
@@ -24,7 +24,7 @@ def extract_taxa(text: str) -> List[str]:
     return list(set(taxa))
 
 
-def create_default_field_structure(field_name: str) -> Dict:
+def create_default_field_structure(field_name: str) -> Dict[str, Any]:
     """Create a default structure for a missing field."""
     field_structures = {
         "host_species": {
@@ -83,7 +83,7 @@ def create_default_field_structure(field_name: str) -> Dict:
     )
 
 
-def validate_field_structure(field_data: Dict, field_name: str) -> bool:
+def validate_field_structure(field_data: Dict[str, Any], field_name: str) -> bool:
     """Validate that a field has the correct structure."""
     required_keys = {
         "status",
@@ -99,7 +99,7 @@ def validate_field_structure(field_data: Dict, field_name: str) -> bool:
     return required_keys.issubset(field_data.keys())
 
 
-def create_comprehensive_fallback_analysis() -> Dict:
+def create_comprehensive_fallback_analysis() -> Dict[str, Any]:
     """Create a comprehensive fallback analysis when parsing completely fails."""
     return {
         "host_species": create_default_field_structure("host_species"),
@@ -131,7 +131,7 @@ def generate_curation_summary(parsed_analysis: Dict, missing_fields: List[str])
 
 def get_paper_metadata_from_csv(
     pmid: str, csv_path: str = "data/full_dump.csv"
-) -> Optional[Dict]:
+) -> Optional[Dict[str, str]]:
     """Get paper metadata from CSV file."""
     try:
         import csv

app/services/advanced_rag.py

@@ -211,7 +211,7 @@ def get_rerank_metrics(self) -> Dict[str, Any]:
         """Get performance metrics from the re-ranker."""
         return self.reranker.get_metrics()
 
-    def get_summary_stats(self, summaries: List[ChunkSummary]) -> Dict:
+    def get_summary_stats(self, summaries: List[ChunkSummary]) -> Dict[str, Any]:
         """Get statistics about summaries."""
         if not summaries:
             return {

app/services/bugsigdb_analyzer.py

@@ -5,7 +5,7 @@
 """
 
 import logging
-from typing import Dict, Optional, List
+from typing import Dict, Optional, List, Any
 import asyncio
 import json
 
@@ -23,12 +23,12 @@
 
 logger = logging.getLogger(__name__)
 
-_unified_qa = None
-_pubmed_retriever = None
-_cache_manager = None
+_unified_qa: Optional[UnifiedQA] = None
+_pubmed_retriever: Optional[PubMedRetriever] = None
+_cache_manager: Optional[CacheManager] = None
 
 
-def get_unified_qa():
+def get_unified_qa() -> Optional[UnifiedQA]:
     """Get or initialize UnifiedQA instance."""
     global _unified_qa
     if _unified_qa is None:
@@ -49,7 +49,7 @@ def get_unified_qa():
     return _unified_qa
 
 
-def get_pubmed_retriever():
+def get_pubmed_retriever() -> Optional[PubMedRetriever]:
     """Get or initialize PubMedRetriever instance."""
     global _pubmed_retriever
     if _pubmed_retriever is None:
@@ -61,15 +61,15 @@ def get_pubmed_retriever():
     return _pubmed_retriever
 
 
-def get_cache_manager():
+def get_cache_manager() -> CacheManager:
     """Get or initialize CacheManager instance."""
     global _cache_manager
     if _cache_manager is None:
         _cache_manager = CacheManager()
     return _cache_manager
 
 
-ESSENTIAL_FIELDS = {
+ESSENTIAL_FIELDS: Dict[str, str] = {
     "host_species": "What host species is being studied in this research?",
     "body_site": "What body site or anatomical location was sampled for microbiome analysis?",
     "condition": "What disease, treatment, or condition is being studied?",
@@ -79,7 +79,7 @@ def get_cache_manager():
 }
 
 
-async def analyze_paper_simple(pmid: str) -> Optional[Dict]:
+async def analyze_paper_simple(pmid: str) -> Optional[Dict[str, Any]]:
     """Extract BugSigDB fields from a paper.
 
     Uses v1 API flow: direct LLM queries per field. Fast but less accurate than RAG.

app/utils/config.py

@@ -1,13 +1,14 @@
 import os
 from pathlib import Path
 import logging
+from typing import List, Optional
 from .credential_masking import mask_exception_message, mask_string
 
 try:
     from dotenv import load_dotenv  # type: ignore
 except ImportError:
 
-    def load_dotenv(*args, **kwargs):  # type: ignore[no-redef]
+    def load_dotenv(*args: object, **kwargs: object) -> None:  # type: ignore[no-redef]
         """Fallback when python-dotenv is not installed."""
         logger = logging.getLogger(__name__)
         logger.warning(
@@ -45,7 +46,7 @@ def load_dotenv(*args, **kwargs):  # type: ignore[no-redef]
 LLM_MODEL = os.getenv("LLM_MODEL", "") or None
 
 DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "gemini")
-AVAILABLE_MODELS = []
+AVAILABLE_MODELS: List[str] = []
 
 if GEMINI_API_KEY:
     AVAILABLE_MODELS.append("gemini")
@@ -57,7 +58,7 @@ def load_dotenv(*args, **kwargs):  # type: ignore[no-redef]
     AVAILABLE_MODELS.append("ollama")
 
 
-def validate_gemini_key():
+def validate_gemini_key() -> bool:
     """Validate Gemini API key by configuring the client."""
     if not GEMINI_API_KEY:
         return False
@@ -75,9 +76,9 @@ def validate_gemini_key():
         return False
 
 
-def validate_env_vars():
+def validate_env_vars() -> bool:
     """Validate that required environment variables are set."""
-    missing_vars = []
+    missing_vars: List[str] = []
 
     if not NCBI_API_KEY:
         missing_vars.append("NCBI_API_KEY")
@@ -103,9 +104,9 @@ def validate_env_vars():
 validate_env_vars()
 
 
-def check_required_vars():
+def check_required_vars() -> bool:
     """Check if all required environment variables are set."""
-    missing_vars = []
+    missing_vars: List[str] = []
 
     if not NCBI_API_KEY:
         missing_vars.append("NCBI_API_KEY")
@@ -212,7 +213,7 @@ def check_required_vars():
 MAX_LOG_FILES = 5  # Keep 5 rotated log files
 
 
-def setup_logging():
+def setup_logging() -> logging.Logger:
     """Configure logging with file rotation.
 
     Falls back to console-only logging if file handlers can't be created.
@@ -237,7 +238,7 @@ def setup_logging():
         root_logger.removeHandler(handler)
 
     # Try to create file handlers, but handle permission errors gracefully
-    file_handlers_created = []
+    file_handlers_created: List[str] = []
 
     try:
         # Main application log handler with rotation

app/utils/credential_masking.py

@@ -6,7 +6,7 @@
 """
 
 import re
-from typing import Optional, Iterable
+from typing import Optional, Iterable, Any, Dict
 
 # Environment / config keys that should always be masked
 API_KEY_ENV_VARS: set[str] = {
@@ -99,7 +99,7 @@ def mask_string(text: str, show_last: int = 4) -> str:
     return masked
 
 
-def mask_dict(data: dict, keys_to_mask: Optional[Iterable[str]] = None) -> dict:
+def mask_dict(data: Dict[str, Any], keys_to_mask: Optional[Iterable[str]] = None) -> Dict[str, Any]:
     """
     Recursively mask sensitive values in dictionaries.
     """
@@ -122,7 +122,7 @@ def mask_dict(data: dict, keys_to_mask: Optional[Iterable[str]] = None) -> dict:
     return masked
 
 
-def safe_log_message(message: str, *args, **kwargs) -> str:
+def safe_log_message(message: str, *args: Any, **kwargs: Any) -> str:
     """
     Format a log message and safely mask any credentials.
     """

cli.py

@@ -95,7 +95,7 @@ def _get_env_file_values(self) -> Dict[str, str]:
         except Exception:
             return {}
 
-    def _collect_env_flags(self) -> list:
+    def _collect_env_flags(self) -> List[str]:
         """Collect docker -e flags for known env vars if present in the host environment."""
         flags = []
         env_file = self._get_env_file_path()
@@ -113,7 +113,7 @@ def _build_api_url(self, path: str) -> str:
         suffix = path.lstrip("/")
         return f"{base}/{suffix}" if suffix else base
 
-    def _validate_environment(self):
+    def _validate_environment(self) -> None:
         """Warn about missing critical env vars before starting containers."""
         env_file = self._get_env_file_path()
         env_file_values = self._get_env_file_values()