OpenRAG 0.1.25__py3-none-any.whl

api/auth.py

@@ -0,0 +1,84 @@
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+
+async def auth_init(request: Request, auth_service, session_manager):
+    """Initialize OAuth flow for authentication or data source connection"""
+    try:
+        data = await request.json()
+        connector_type = data.get("connector_type")
+        purpose = data.get("purpose", "data_source")
+        connection_name = data.get("name", f"{connector_type}_{purpose}")
+        redirect_uri = data.get("redirect_uri")
+
+        user = getattr(request.state, "user", None)
+        user_id = user.user_id if user else None
+
+        result = await auth_service.init_oauth(
+            connector_type, purpose, connection_name, redirect_uri, user_id
+        )
+        return JSONResponse(result)
+
+    except Exception as e:
+        import traceback
+
+        traceback.print_exc()
+        return JSONResponse(
+            {"error": f"Failed to initialize OAuth: {str(e)}"}, status_code=500
+        )
+
+
+async def auth_callback(request: Request, auth_service, session_manager):
+    """Handle OAuth callback - exchange authorization code for tokens"""
+    try:
+        data = await request.json()
+        connection_id = data.get("connection_id")
+        authorization_code = data.get("authorization_code")
+        state = data.get("state")
+
+        result = await auth_service.handle_oauth_callback(
+            connection_id, authorization_code, state, request
+        )
+
+        # If this is app auth, set JWT cookie
+        if result.get("purpose") == "app_auth" and result.get("jwt_token"):
+            response = JSONResponse(
+                {k: v for k, v in result.items() if k != "jwt_token"}
+            )
+            response.set_cookie(
+                key="auth_token",
+                value=result["jwt_token"],
+                httponly=True,
+                secure=False,
+                samesite="lax",
+                max_age=7 * 24 * 60 * 60,  # 7 days
+            )
+            return response
+        else:
+            return JSONResponse(result)
+
+    except Exception as e:
+        import traceback
+
+        traceback.print_exc()
+        return JSONResponse({"error": f"Callback failed: {str(e)}"}, status_code=500)
+
+
+async def auth_me(request: Request, auth_service, session_manager):
+    """Get current user information"""
+    result = await auth_service.get_user_info(request)
+    return JSONResponse(result)
+
+
+async def auth_logout(request: Request, auth_service, session_manager):
+    """Logout user by clearing auth cookie"""
+    response = JSONResponse(
+        {"status": "logged_out", "message": "Successfully logged out"}
+    )
+
+    # Clear the auth cookie
+    response.delete_cookie(
+        key="auth_token", httponly=True, secure=False, samesite="lax"
+    )
+
+    return response

api/chat.py

@@ -0,0 +1,181 @@
+from starlette.requests import Request
+from starlette.responses import JSONResponse, StreamingResponse
+from utils.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+async def chat_endpoint(request: Request, chat_service, session_manager):
+    """Handle chat requests"""
+    data = await request.json()
+    prompt = data.get("prompt", "")
+    previous_response_id = data.get("previous_response_id")
+    stream = data.get("stream", False)
+    filters = data.get("filters")
+    limit = data.get("limit", 10)
+    score_threshold = data.get("scoreThreshold", 0)
+
+    user = request.state.user
+    user_id = user.user_id
+
+    jwt_token = session_manager.get_effective_jwt_token(user_id, request.state.jwt_token)
+
+    if not prompt:
+        return JSONResponse({"error": "Prompt is required"}, status_code=400)
+
+    # Set context variables for search tool (similar to search endpoint)
+    if filters:
+        from auth_context import set_search_filters
+
+        set_search_filters(filters)
+
+    from auth_context import set_search_limit, set_score_threshold
+
+    set_search_limit(limit)
+    set_score_threshold(score_threshold)
+
+    if stream:
+        return StreamingResponse(
+            await chat_service.chat(
+                prompt,
+                user_id,
+                jwt_token,
+                previous_response_id=previous_response_id,
+                stream=True,
+            ),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "Access-Control-Allow-Origin": "*",
+                "Access-Control-Allow-Headers": "Cache-Control",
+            },
+        )
+    else:
+        result = await chat_service.chat(
+            prompt,
+            user_id,
+            jwt_token,
+            previous_response_id=previous_response_id,
+            stream=False,
+        )
+        return JSONResponse(result)
+
+
+async def langflow_endpoint(request: Request, chat_service, session_manager):
+    """Handle Langflow chat requests"""
+    data = await request.json()
+    prompt = data.get("prompt", "")
+    previous_response_id = data.get("previous_response_id")
+    stream = data.get("stream", False)
+    filters = data.get("filters")
+    limit = data.get("limit", 10)
+    score_threshold = data.get("scoreThreshold", 0)
+
+    user = request.state.user
+    user_id = user.user_id
+
+    jwt_token = session_manager.get_effective_jwt_token(user_id, request.state.jwt_token)
+
+    if not prompt:
+        return JSONResponse({"error": "Prompt is required"}, status_code=400)
+
+    # Set context variables for search tool (similar to chat endpoint)
+    if filters:
+        from auth_context import set_search_filters
+
+        set_search_filters(filters)
+
+    from auth_context import set_search_limit, set_score_threshold
+
+    set_search_limit(limit)
+    set_score_threshold(score_threshold)
+
+    try:
+        if stream:
+            return StreamingResponse(
+                await chat_service.langflow_chat(
+                    prompt,
+                    user_id,
+                    jwt_token,
+                    previous_response_id=previous_response_id,
+                    stream=True,
+                ),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "Access-Control-Allow-Origin": "*",
+                    "Access-Control-Allow-Headers": "Cache-Control",
+                },
+            )
+        else:
+            result = await chat_service.langflow_chat(
+                prompt,
+                user_id,
+                jwt_token,
+                previous_response_id=previous_response_id,
+                stream=False,
+            )
+            return JSONResponse(result)
+
+    except Exception as e:
+        import traceback
+
+        traceback.print_exc()
+        logger.error("Langflow request failed", error=str(e))
+        return JSONResponse(
+            {"error": f"Langflow request failed: {str(e)}"}, status_code=500
+        )
+
+
+async def chat_history_endpoint(request: Request, chat_service, session_manager):
+    """Get chat history for a user"""
+    user = request.state.user
+    user_id = user.user_id
+
+    try:
+        history = await chat_service.get_chat_history(user_id)
+        return JSONResponse(history)
+    except Exception as e:
+        return JSONResponse(
+            {"error": f"Failed to get chat history: {str(e)}"}, status_code=500
+        )
+
+
+async def langflow_history_endpoint(request: Request, chat_service, session_manager):
+    """Get langflow chat history for a user"""
+    user = request.state.user
+    user_id = user.user_id
+
+    try:
+        history = await chat_service.get_langflow_history(user_id)
+        return JSONResponse(history)
+    except Exception as e:
+        return JSONResponse(
+            {"error": f"Failed to get langflow history: {str(e)}"}, status_code=500
+        )
+
+
+async def delete_session_endpoint(request: Request, chat_service, session_manager):
+    """Delete a chat session"""
+    user = request.state.user
+    user_id = user.user_id
+    session_id = request.path_params["session_id"]
+
+    try:
+        # Delete from both local storage and Langflow
+        result = await chat_service.delete_session(user_id, session_id)
+
+        if result.get("success"):
+            return JSONResponse({"message": "Session deleted successfully"})
+        else:
+            return JSONResponse(
+                {"error": result.get("error", "Failed to delete session")},
+                status_code=500
+            )
+    except Exception as e:
+        logger.error(f"Error deleting session: {e}")
+        return JSONResponse(
+            {"error": f"Failed to delete session: {str(e)}"}, status_code=500
+        )

api/connector_router.py

@@ -0,0 +1,78 @@
+"""Connector router that automatically routes based on configuration settings."""
+
+from starlette.requests import Request
+
+from config.settings import (
+    DISABLE_INGEST_WITH_LANGFLOW,
+    clients,
+    INDEX_NAME,
+    INDEX_BODY,
+)
+from utils.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class ConnectorRouter:
+    """
+    Router that automatically chooses between LangflowConnectorService and ConnectorService
+    based on the DISABLE_INGEST_WITH_LANGFLOW configuration.
+
+    - If DISABLE_INGEST_WITH_LANGFLOW is False (default): uses LangflowConnectorService
+    - If DISABLE_INGEST_WITH_LANGFLOW is True: uses traditional ConnectorService
+    """
+
+    def __init__(self, langflow_connector_service, openrag_connector_service):
+        self.langflow_connector_service = langflow_connector_service
+        self.openrag_connector_service = openrag_connector_service
+        logger.debug(
+            "ConnectorRouter initialized",
+            disable_langflow_ingest=DISABLE_INGEST_WITH_LANGFLOW,
+        )
+
+    def get_active_service(self):
+        """Get the currently active connector service based on configuration."""
+        if DISABLE_INGEST_WITH_LANGFLOW:
+            logger.debug("Using traditional OpenRAG connector service")
+            return self.openrag_connector_service
+        else:
+            logger.debug("Using Langflow connector service")
+            return self.langflow_connector_service
+
+    # Proxy all connector service methods to the active service
+
+    async def initialize(self):
+        """Initialize the active connector service."""
+        # Initialize OpenSearch index if using traditional OpenRAG connector service
+
+        return await self.get_active_service().initialize()
+
+    @property
+    def connection_manager(self):
+        """Get the connection manager from the active service."""
+        return self.get_active_service().connection_manager
+
+    async def get_connector(self, connection_id: str):
+        """Get a connector instance from the active service."""
+        return await self.get_active_service().get_connector(connection_id)
+
+    async def sync_specific_files(
+        self, connection_id: str, user_id: str, file_list: list, jwt_token: str = None
+    ):
+        """Sync specific files using the active service."""
+        return await self.get_active_service().sync_specific_files(
+            connection_id, user_id, file_list, jwt_token
+        )
+
+    def __getattr__(self, name):
+        """
+        Proxy any other method calls to the active service.
+        This ensures compatibility with any methods we might have missed.
+        """
+        active_service = self.get_active_service()
+        if hasattr(active_service, name):
+            return getattr(active_service, name)
+        else:
+            raise AttributeError(
+                f"'{type(active_service).__name__}' object has no attribute '{name}'"
+            )