analysis_endpoints.py

import json
import asyncio
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from typing import List
from uuid import UUID
import logging
from datetime import datetime

from app.api.dependencies import (
    get_analysis_service,
    get_orchestrator_service,
    get_current_user,
    get_claim_service,
    get_together_orchestrator_service,
)
from app.models.domain.user import User
from app.services.analysis_service import AnalysisService
from app.services.claim_service import ClaimService
from app.services.analysis_orchestrator import AnalysisOrchestrator
from app.schemas.analysis_schema import AnalysisRead
from app.core.exceptions import NotFoundException
from fastapi.responses import StreamingResponse
from app.core.scoring import get_percentile

router = APIRouter(prefix="/analysis", tags=["analysis"])
logger = logging.getLogger(__name__)


@router.post("/create", response_model=AnalysisRead)
async def create_analysis_test(data: AnalysisRead) -> AnalysisRead:
    pass


@router.get("/{analysis_id}", response_model=AnalysisRead)
async def get_analysis(
    analysis_id: UUID,
    include_sources: bool = Query(False),
    include_feedback: bool = Query(False),
    current_user: User = Depends(get_current_user),
    analysis_service: AnalysisService = Depends(get_analysis_service),
) -> AnalysisRead:
    try:
        analysis = await analysis_service.get_analysis(
            analysis_id=analysis_id, include_sources=include_sources, include_feedback=include_feedback
        )
        raw_score = analysis.confidence_score
        percentile = (get_percentile(raw_score)) / 100.0
        analysis = AnalysisRead.model_validate(analysis)
        analysis.confidence_percentile = percentile
        return analysis
    except NotFoundException as e:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))


@router.get("/claim/{claim_id}/stream", response_class=StreamingResponse)
async def stream_claim_analysis(
    request: Request,
    claim_id: UUID,
    current_user: User = Depends(get_current_user),
    analysis_orchestrator: AnalysisOrchestrator = Depends(get_orchestrator_service),
    claim_service: ClaimService = Depends(get_claim_service),
) -> StreamingResponse:
    """Stream the analysis process for a claim in real-time."""
    try:
        # current_user = await auth_middleware.authenticate_request(request)

        claim = await claim_service.get_claim(claim_id=claim_id, user_id=current_user.id)

        session = claim_service._claim_repo._session

        async def event_generator():
            try:
                logger.info(f"Starting analysis stream for claim {claim_id}")
                yield f"data: {json.dumps({'type': 'status', 'content': 'Initializing analysis...'})}\n\n"

                async for event in analysis_orchestrator.analyze_claim_stream(claim=claim, user_id=current_user.id):
                    if isinstance(event, dict):
                        yield f"data: {json.dumps(event)}\n\n"

            except Exception as e:
                logger.error(f"Error in analysis stream: {str(e)}", exc_info=True)
                yield f"data: {json.dumps({'type': 'error', 'content': str(e)})}\n\n"
            finally:
                await session.close()
                yield "data: [DONE]\n\n"

        return StreamingResponse(
            event_generator(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
                "Access-Control-Allow-Origin": "*",
                "Access-Control-Allow-Credentials": "true",
            },
        )
    except Exception as e:
        logger.error(f"Stream error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/experiment/claim/{claim_id}/stream", response_class=StreamingResponse)
async def stream_claim_analysis_exp(
    request: Request,
    claim_id: UUID,
    current_user: User = Depends(get_current_user),
    analysis_orchestrator: AnalysisOrchestrator = Depends(get_together_orchestrator_service),
    claim_service: ClaimService = Depends(get_claim_service),
) -> StreamingResponse:
    """Stream the analysis process for a claim in real-time."""
    try:
        claim = await claim_service.get_claim(claim_id=claim_id, user_id=current_user.id)

        session = claim_service._claim_repo._session

        await session.rollback()

        async def event_generator():
            try:
                logger.info(f"Starting analysis stream for claim {claim_id}")
                yield f"data: {json.dumps({'type': 'status', 'content': 'Initializing analysis...'})}\n\n"

                orchestrator_stream = analysis_orchestrator.analyze_claim_stream(
                    claim=claim, user_id=current_user.id, default=False
                )
                # ---------------------------------------------------------
                # THE HEALTH CHECK LOOP
                # ---------------------------------------------------------
                next_event_task = None

                while True:
                    # Only create a new task if we don't already have one waiting
                    if next_event_task is None:
                        next_event_task = asyncio.create_task(anext(orchestrator_stream))

                    # Wait for the task to finish, but only wait 15 seconds
                    done, pending = await asyncio.wait(
                        [next_event_task], timeout=15.0, return_when=asyncio.FIRST_COMPLETED
                    )

                    if next_event_task in done:
                        # The LLM yielded a chunk! Let's process it.
                        try:
                            event = next_event_task.result()
                            if isinstance(event, dict):
                                yield f"data: {json.dumps(event)}\n\n"

                            # Reset the task so we grab the next chunk on the next loop
                            next_event_task = None

                        except StopAsyncIteration:
                            # The stream finished normally!
                            break
                        except Exception as e:
                            # If the orchestrator crashed, catch it here
                            raise e

                    else:
                        # The task is in 'pending'. 15 seconds passed, but the LLM is still thinking.
                        # We yield a heartbeat, but we DO NOT reset next_event_task.
                        # It will keep running safely in the background on the next loop!
                        logger.debug("Stream idle for 15s. Sending health check ping...")
                        yield ": healthcheck\n\n"

                yield "data: [DONE]\n\n"

            except asyncio.CancelledError:
                logger.warning(f"Client disconnected during stream for claim {claim_id}")
                raise

            # async for event in analysis_orchestrator.analyze_claim_stream(
            #     claim=claim, user_id=current_user.id, default=False
            # ):
            #     if isinstance(event, dict):
            #         yield f"data: {json.dumps(event)}\n\n"

            # yield "data: [DONE]\n\n"

            # except asyncio.CancelledError:
            #     # THE FIX: The user closed their browser!
            #     logger.info(f"Client disconnected during stream for claim {claim_id}")
            #     await session.rollback()  # Explicitly release the lock!
            #     raise

            except Exception as e:
                logger.error(f"Error in analysis stream: {str(e)}", exc_info=True)
                yield f"data: {json.dumps({'type': 'error', 'content': str(e)})}\n\n"
            finally:
                # async def force_cleanup():
                #     try:
                #         await session.rollback()
                #     except Exception as e:
                #         logger.error(f"Force rollback failed: {e}")
                #     finally:
                #         await session.close()

                # # Fire and forget. FastAPI cannot cancel this!
                # asyncio.create_task(force_cleanup())
                if next_event_task and not next_event_task.done():
                    logger.debug("Cancelling background orchestrator task...")
                    next_event_task.cancel()

        return StreamingResponse(
            event_generator(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
                # "Access-Control-Allow-Origin": "*",
                # "Access-Control-Allow-Credentials": "true",
            },
        )
    except Exception as e:
        logger.error(f"Stream error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/claim/{claim_id}", response_model=List[AnalysisRead])
async def get_claim_analyses(
    claim_id: UUID,
    include_sources: bool = Query(False),
    include_feedback: bool = Query(False),
    current_user: User = Depends(get_current_user),
    analysis_service: AnalysisService = Depends(get_analysis_service),
) -> List[AnalysisRead]:
    try:
        logger.info("Fetching analysis...")
        analyses = await analysis_service.get_claim_analyses(
            claim_id=claim_id, include_sources=include_sources, include_feedback=include_feedback
        )
        return [AnalysisRead.model_validate(a) for a in analyses]
    except Exception as e:
        logger.error("Could not find the analysis for the claim")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/avg/total", response_model=dict, summary="Get average reliability score for claims by language")
async def get_claim(
    start_date: datetime,
    end_date: datetime,
    language: str = "english",
    analysis_service: AnalysisService = Depends(get_analysis_service),
) -> dict:
    """Get average reliability score for claims by language."""
    try:
        analyses = await analysis_service.get_analysis_list(start_date=start_date, end_date=end_date, language=language)

        if not analyses:
            return {"avg_score": 0.0}  # Return 0 if the list is empty to avoid division by zero

        total_score = sum(analysis.veracity_score for analysis in analyses)
        average_score = total_score / len(analyses)

        return {"avg_score": average_score}
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to get analysis list: {str(e)}"
        )


@router.get("/{analysis_id}/assertiveness", response_model=AnalysisRead, summary="Modify the assertiveness level")
async def vary_assert(
    analysis_id: UUID,
    analysis_orchestrator: AnalysisOrchestrator = Depends(get_orchestrator_service),
) -> AnalysisRead:
    """Get average reliability score for claims by language."""
    try:
        analysis = await analysis_orchestrator.vary_analysis_assertiveness(analysis_id=analysis_id)
        raw_score = analysis.confidence_score
        percentile = (get_percentile(raw_score)) / 100.0
        analysis = AnalysisRead.model_validate(analysis)
        analysis.confidence_percentile = percentile
        return analysis
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to get alter assertivity : {str(e)}"
        )