Aryan2301/YouTube_RAG_Intelligence

1from fastapi import FastAPI, HTTPException2from fastapi.responses import StreamingResponse3from pydantic import BaseModel4from typing import List, Dict, Any5import json6import asyncio7import logging8from langgraph.checkpoint.memory import MemorySaver9from contextlib import asynccontextmanager10 11from core.graph import build_langgraph12from core.vectorstore import build_retrievers13from services.transcript_service import fetch_transcript14 15logger = logging.getLogger(__name__)16 17@asynccontextmanager18async def lifespan(app: FastAPI):19    # In-memory checkpointer for conversation memory20    app.state.checkpointer = MemorySaver()21    app.state.rag_graph = build_langgraph(checkpointer=app.state.checkpointer)22    yield23 24app = FastAPI(title="YouTube RAG Intelligence API", lifespan=lifespan)25 26# Simple in-memory cache for vector stores and retrievers27# In production, use Redis or a persistent ChromaDB path28app.state.video_stores = {}29 30class InitRequest(BaseModel):31    video_id: str32    transcript_segments: List[Dict[str, Any]] = []33 34class ChatRequest(BaseModel):35    video_id: str36    chat_id: str37    query: str38    chat_history: List[Dict[str, str]] = []39 40@app.post("/init")41async def init_video(req: InitRequest):42    """Initialize a video: fetch transcript and build vector store."""43    try:44        if req.video_id in app.state.video_stores:45            return {"status": "already initialized"}46            47        if not req.transcript_segments:48            transcript_text, transcript_segments = await asyncio.to_thread(fetch_transcript, req.video_id)49        else:50            transcript_segments = req.transcript_segments51            52        vector_store, bm25_retriever = await asyncio.to_thread(build_retrievers, req.video_id, transcript_segments)53        54        app.state.video_stores[req.video_id] = {55            "vector_store": vector_store,56            "bm25_retriever": bm25_retriever57        }58        59        return {"status": "success", "segments": len(transcript_segments)}60    except Exception as e:61        logger.error(f"Error initializing video {req.video_id}: {str(e)}")62        raise HTTPException(status_code=500, detail=str(e))63 64@app.post("/chat/stream")65async def chat_stream(req: ChatRequest):66    """Stream chat response and graph events using Server-Sent Events (SSE)."""67    68    if req.video_id not in app.state.video_stores:69        return {"error": "Video not initialized. Call /init first."}70        71    stores = app.state.video_stores[req.video_id]72    73    inputs = {74        "query": req.query,75        "chat_history": req.chat_history,76        "video_id": req.video_id,77        "retrieval_attempt": 0,78        "context": ""79    }80    81    thread_id = f"{req.video_id}_{req.chat_id}"82    config = {83        "configurable": {84            "thread_id": thread_id,85            "vector_store": stores["vector_store"],86            "bm25_retriever": stores["bm25_retriever"],87            # LLM needs to be initialized. We can import here to avoid global init issues.88            "llm": __import__("core.llm", fromlist=["load_llm"]).load_llm()89        }90    }91 92    async def event_generator():93        # Yield meaningful Server-Sent Events (SSE) by streaming the graph execution94        try:95            async for event in app.state.rag_graph.astream_events(inputs, config=config, version="v1"):96                kind = event["event"]97                name = event["name"]98                99                # Stream graph node transitions (e.g., hybrid_retrieve, evaluate_retrieval)100                if kind == "on_chain_start" and name in ["hybrid_retrieve", "evaluate_retrieval", "correct_query", "web_search", "generate_answer"]:101                    yield f"event: status\ndata: {json.dumps({'message': f'Starting node: {name}'})}\n\n"102                    103                # Stream the actual tokens from the LLM inside generate_answer104                elif kind == "on_chat_model_stream":105                    chunk = event["data"]["chunk"].content106                    if isinstance(chunk, list):107                        chunk = "".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in chunk)108                    if chunk:109                        yield f"event: token\ndata: {json.dumps({'token': chunk})}\n\n"110                        111            yield "event: end\ndata: {}\n\n"112        except Exception as e:113            yield f"event: error\ndata: {json.dumps({'error': str(e)})}\n\n"114 115    return StreamingResponse(event_generator(), media_type="text/event-stream")116