Add google drive connector (#150)

Author: Adityavardhan Agrawal

core/api.py

@@ -3,6 +3,7 @@
 import json
 import logging
 import uuid
+from contextlib import asynccontextmanager
 from datetime import UTC, datetime, timedelta
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
@@ -11,15 +12,17 @@
 import jwt
 import tomli
 from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, UploadFile
-from fastapi.middleware.cors import CORSMiddleware
+from fastapi.middleware.cors import CORSMiddleware  # Import CORSMiddleware
 from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+from starlette.middleware.sessions import SessionMiddleware
 
 from core.agent import MorphikAgent
 from core.auth_utils import verify_token
 from core.cache.llama_cache_factory import LlamaCacheFactory
 from core.completion.litellm_completion import LiteLLMCompletionModel
 from core.config import get_settings
 from core.database.postgres_database import PostgresDatabase
+from core.dependencies import get_redis_pool
 from core.embedding.colpali_api_embedding_model import ColpaliApiEmbeddingModel
 from core.embedding.colpali_embedding_model import ColpaliEmbeddingModel
 from core.embedding.litellm_embedding import LiteLLMEmbeddingModel
@@ -51,31 +54,82 @@
 from core.vector_store.pgvector_store import PGVectorStore
 
 # Initialize FastAPI app
-app = FastAPI(title="Morphik API")
 logger = logging.getLogger(__name__)
 
+# Global variable for redis_pool, primarily for shutdown if app.state access fails.
+_global_redis_pool: Optional[arq.ArqRedis] = None
 
-# Add health check endpoints
-@app.get("/health")
-async def health_check():
-    """Basic health check endpoint."""
-    return {"status": "healthy"}
-
-
-@app.get("/health/ready")
-async def readiness_check():
-    """Readiness check that verifies the application is initialized."""
-    return {
-        "status": "ready",
-        "components": {
-            "database": settings.DATABASE_PROVIDER,
-            "vector_store": settings.VECTOR_STORE_PROVIDER,
-            "embedding": settings.EMBEDDING_PROVIDER,
-            "completion": settings.COMPLETION_PROVIDER,
-            "storage": settings.STORAGE_PROVIDER,
-        },
-    }
 
+@asynccontextmanager
+async def lifespan(app_instance: FastAPI):
+    # --- BEGIN MOVED STARTUP LOGIC ---
+    logger.info("Lifespan: Initializing Database...")
+    try:
+        success = await database.initialize()
+        if success:
+            logger.info("Lifespan: Database initialization successful")
+        else:
+            logger.error("Lifespan: Database initialization failed")
+    except Exception as e:
+        logger.error(f"Lifespan: CRITICAL - Failed to initialize Database: {e}", exc_info=True)
+        raise  # Or handle more gracefully if appropriate
+
+    logger.info("Lifespan: Initializing Vector Store...")
+    try:
+        global vector_store
+        if hasattr(vector_store, "initialize"):
+            await vector_store.initialize()
+        logger.info("Lifespan: Vector Store initialization successful (or not applicable).")
+    except Exception as e:
+        logger.error(f"Lifespan: CRITICAL - Failed to initialize Vector Store: {e}", exc_info=True)
+        # Decide if this is fatal
+        # raise
+
+    logger.info("Lifespan: Attempting to initialize Redis connection pool...")
+    global _global_redis_pool  # Ensure we're using the global for assignment
+    try:
+        redis_settings_obj = arq.connections.RedisSettings(
+            host=settings.REDIS_HOST,
+            port=settings.REDIS_PORT,
+        )
+        logger.info(f"Lifespan: Redis settings for pool: host={settings.REDIS_HOST}, port={settings.REDIS_PORT}")
+        current_redis_pool = await arq.create_pool(redis_settings_obj)
+        if current_redis_pool:
+            app_instance.state.redis_pool = current_redis_pool  # Use app_instance from lifespan
+            _global_redis_pool = current_redis_pool
+            logger.info("Lifespan: Successfully initialized Redis connection pool and stored on app.state.")
+        else:
+            logger.error("Lifespan: arq.create_pool returned None or a falsey value for Redis pool.")
+            raise RuntimeError("Lifespan: Failed to create Redis pool - arq.create_pool returned non-truthy value.")
+    except Exception as e:
+        logger.error(f"Lifespan: CRITICAL - Failed to initialize Redis connection pool: {e}", exc_info=True)
+        raise RuntimeError(f"Lifespan: CRITICAL - Failed to initialize Redis connection pool: {e}") from e
+    # --- END MOVED STARTUP LOGIC ---
+
+    logger.info("Lifespan: Core startup logic executed.")
+    yield
+    # Shutdown logic
+    logger.info("Lifespan: Shutdown initiated.")
+    # Use app_instance.state to get the pool for shutdown too
+    pool_to_close = getattr(app_instance.state, "redis_pool", _global_redis_pool)
+    if pool_to_close:
+        logger.info("Closing Redis connection pool from lifespan...")
+        pool_to_close.close()
+        await pool_to_close.wait_closed()
+        logger.info("Redis connection pool closed from lifespan.")
+    logger.info("Lifespan: Shutdown complete.")
+
+
+app = FastAPI(lifespan=lifespan)
+
+# Add CORSMiddleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],  # Allows all headers
+)
 
 # Initialize telemetry
 telemetry = TelemetryService()
@@ -90,112 +144,17 @@ async def readiness_check():
     tracer_provider=None,  # Use the global tracer provider
 )
 
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
 # Initialize service
 settings = get_settings()
 
+# Add SessionMiddleware
+app.add_middleware(SessionMiddleware, secret_key=settings.SESSION_SECRET_KEY)
+
 # Initialize database
 if not settings.POSTGRES_URI:
     raise ValueError("PostgreSQL URI is required for PostgreSQL database")
 database = PostgresDatabase(uri=settings.POSTGRES_URI)
 
-# Redis settings already imported at top of file
-
-
-@app.on_event("startup")
-async def initialize_database():
-    """Initialize database tables and indexes on application startup."""
-    logger.info("Initializing database...")
-    success = await database.initialize()
-    if success:
-        logger.info("Database initialization successful")
-    else:
-        logger.error("Database initialization failed")
-        # We don't raise an exception here to allow the app to continue starting
-        # even if there are initialization errors
-
-
-@app.on_event("startup")
-async def initialize_vector_store():
-    """Initialize vector store tables and indexes on application startup."""
-    # First initialize the primary vector store (PGVectorStore if using pgvector)
-    logger.info("Initializing primary vector store...")
-    if hasattr(vector_store, "initialize"):
-        success = await vector_store.initialize()
-        if success:
-            logger.info("Primary vector store initialization successful")
-        else:
-            logger.error("Primary vector store initialization failed")
-    else:
-        logger.warning("Primary vector store does not have an initialize method")
-
-    # Then initialize the multivector store if enabled
-    if settings.ENABLE_COLPALI and colpali_vector_store:
-        logger.info("Initializing multivector store...")
-        # Handle both synchronous and asynchronous initialize methods
-        if hasattr(colpali_vector_store.initialize, "__awaitable__"):
-            success = await colpali_vector_store.initialize()
-        else:
-            success = colpali_vector_store.initialize()
-
-        if success:
-            logger.info("Multivector store initialization successful")
-        else:
-            logger.error("Multivector store initialization failed")
-
-
-@app.on_event("startup")
-async def initialize_user_limits_database():
-    """Initialize user service on application startup."""
-    logger.info("Initializing user service...")
-    if settings.MODE == "cloud":
-        from core.database.user_limits_db import UserLimitsDatabase
-
-        user_limits_db = UserLimitsDatabase(uri=settings.POSTGRES_URI)
-        await user_limits_db.initialize()
-
-
-@app.on_event("startup")
-async def initialize_redis_pool():
-    """Initialize the Redis connection pool for background tasks."""
-    global redis_pool
-    logger.info("Initializing Redis connection pool...")
-
-    # Get Redis settings from configuration
-    redis_host = settings.REDIS_HOST
-    redis_port = settings.REDIS_PORT
-
-    # Log the Redis connection details
-    logger.info(f"Connecting to Redis at {redis_host}:{redis_port}")
-
-    redis_settings = arq.connections.RedisSettings(
-        host=redis_host,
-        port=redis_port,
-    )
-
-    redis_pool = await arq.create_pool(redis_settings)
-    logger.info("Redis connection pool initialized successfully")
-
-
-@app.on_event("shutdown")
-async def close_redis_pool():
-    """Close the Redis connection pool on application shutdown."""
-    global redis_pool
-    if redis_pool:
-        logger.info("Closing Redis connection pool...")
-        redis_pool.close()
-        await redis_pool.wait_closed()
-        logger.info("Redis connection pool closed")
-
-
 # Initialize vector store
 if not settings.POSTGRES_URI:
     raise ValueError("PostgreSQL URI is required for pgvector store")
@@ -279,18 +238,22 @@ async def close_redis_pool():
 
 # Initialize document service with configured components
 document_service = DocumentService(
-    storage=storage,
     database=database,
     vector_store=vector_store,
+    storage=storage,
+    parser=parser,
     embedding_model=embedding_model,
     completion_model=completion_model,
-    parser=parser,
-    reranker=reranker,
     cache_factory=cache_factory,
-    enable_colpali=(settings.COLPALI_MODE != "off"),
+    reranker=reranker,
+    enable_colpali=settings.ENABLE_COLPALI,
     colpali_embedding_model=colpali_embedding_model,
     colpali_vector_store=colpali_vector_store,
 )
+# Store document_service on app.state immediately after it's created.
+# This must happen before _init_ee_app(app) is called.
+app.state.document_service = document_service
+logger.info("Document service initialized and stored on app.state")
 
 # Initialize the MorphikAgent once to load tool definitions and avoid repeated I/O
 morphik_agent = MorphikAgent(document_service=document_service)
@@ -304,10 +267,13 @@ async def close_redis_pool():
     from ee.routers import init_app as _init_ee_app  # type: ignore
 
     _init_ee_app(app)  # noqa: SLF001 – runtime extension
-    logger.info("Enterprise routes mounted (ee package detected).")
 except ModuleNotFoundError:
     # Expected in OSS builds – silently ignore.
     logger.debug("Enterprise package not found – running in community mode.")
+except ImportError as e:
+    logger.error(f"Failed to import init_app from ee.routers: {e}", exc_info=True)
+except Exception as e:
+    logger.error(f"An unexpected error occurred during EE app initialization: {e}", exc_info=True)
 
 
 @app.post("/ingest/text", response_model=Document)
@@ -349,15 +315,6 @@ async def ingest_text(
         raise HTTPException(status_code=403, detail=str(e))
 
 
-# Redis pool for background tasks
-redis_pool = None
-
-
-def get_redis_pool():
-    """Get the global Redis connection pool for background tasks."""
-    return redis_pool
-
-
 @app.post("/ingest/file", response_model=Document)
 @telemetry.track(operation_type="queue_ingest_file", metadata_resolver=telemetry.ingest_file_metadata)
 async def ingest_file(

core/config.py

@@ -13,6 +13,7 @@ class Settings(BaseSettings):
 
     # Environment variables
     JWT_SECRET_KEY: str
+    SESSION_SECRET_KEY: str
     POSTGRES_URI: Optional[str] = None
     UNSTRUCTURED_API_KEY: Optional[str] = None
     AWS_ACCESS_KEY: Optional[str] = None
@@ -155,6 +156,7 @@ def get_settings() -> Settings:
     auth_config = {
         "JWT_ALGORITHM": config["auth"]["jwt_algorithm"],
         "JWT_SECRET_KEY": os.environ.get("JWT_SECRET_KEY", "dev-secret-key"),  # Default for dev mode
+        "SESSION_SECRET_KEY": os.environ.get("SESSION_SECRET_KEY", "super-secret-dev-session-key"),
         "dev_mode": config["auth"].get("dev_mode", False),
         "dev_entity_type": config["auth"].get("dev_entity_type", "developer"),
         "dev_entity_id": config["auth"].get("dev_entity_id", "dev_user"),
@@ -164,6 +166,13 @@ def get_settings() -> Settings:
     # Only require JWT_SECRET_KEY in non-dev mode
     if not auth_config["dev_mode"] and "JWT_SECRET_KEY" not in os.environ:
         raise ValueError("JWT_SECRET_KEY is required when dev_mode is disabled")
+    # Also require SESSION_SECRET_KEY in non-dev mode
+    if not auth_config["dev_mode"] and "SESSION_SECRET_KEY" not in os.environ:
+        # Or, if we want to be more strict and always require it via ENV:
+        # if "SESSION_SECRET_KEY" not in os.environ:
+        #     raise ValueError("SESSION_SECRET_KEY environment variable is required.")
+        # For now, align with JWT_SECRET_KEY's dev mode leniency.
+        pass  # Dev mode has a default, production should use ENV.
 
     # Load registered models if available
     registered_models = {}

core/dependencies.py

@@ -0,0 +1,19 @@
+from typing import TYPE_CHECKING
+
+import arq
+from fastapi import Request
+
+if TYPE_CHECKING:
+    from core.services.document_service import DocumentService
+
+
+async def get_redis_pool(request: Request) -> arq.ArqRedis:
+    if not hasattr(request.app.state, "redis_pool") or request.app.state.redis_pool is None:
+        raise RuntimeError("Redis pool not initialized or not available on app.state")
+    return request.app.state.redis_pool
+
+
+async def get_document_service(request: Request) -> "DocumentService":
+    if not hasattr(request.app.state, "document_service") or request.app.state.document_service is None:
+        raise RuntimeError("Document service not initialized or not available on app.state")
+    return request.app.state.document_service

core/models/documents.py

@@ -37,7 +37,7 @@ class Document(BaseModel):
     filename: Optional[str] = None
     metadata: Dict[str, Any] = Field(default_factory=dict)
     """user-defined metadata"""
-    storage_info: Dict[str, str] = Field(default_factory=dict)
+    storage_info: Dict[str, Any] = Field(default_factory=dict)
     """Legacy field for backwards compatibility - for single file storage"""
     storage_files: List[StorageFileInfo] = Field(default_factory=list)
     """List of files associated with this document"""
@@ -57,6 +57,13 @@ class Document(BaseModel):
     access_control: Dict[str, List[str]] = Field(default_factory=lambda: {"readers": [], "writers": [], "admins": []})
     chunk_ids: List[str] = Field(default_factory=list)
 
+    # Ensure storage_info values are strings to maintain backward compatibility
+    @field_validator("storage_info", mode="before")
+    def _coerce_storage_info_values(cls, v):
+        if isinstance(v, dict):
+            return {k: str(val) if val is not None else "" for k, val in v.items()}
+        return v
+
     def __hash__(self):
         return hash(self.external_id)