Create video_rag_routes.py

video_rag_routes.py

@@ -0,0 +1,148 @@
+# video_rag_routes.py
+
+import os
+import uuid
+from fastapi import APIRouter, HTTPException, UploadFile, File
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.chains import ConversationalRetrievalChain
+from langchain_core.prompts import ChatPromptTemplate
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_groq import ChatGroq
+from google import genai
+from google.genai import types
+
+router = APIRouter(prefix="/video_rag", tags=["video_rag"])
+
+# ——— Helpers ——————————————————————————————————————————————
+
+def init_google_client():
+    api_key = os.getenv("GOOGLE_API_KEY", "")
+    if not api_key:
+        raise ValueError("GOOGLE_API_KEY must be set")
+    return genai.Client(api_key=api_key)
+
+def get_llm():
+    api_key = os.getenv("CHATGROQ_API_KEY", "")
+    if not api_key:
+        raise ValueError("CHATGROQ_API_KEY must be set")
+    return ChatGroq(
+        model="llama-3.3-70b-versatile",
+        temperature=0,
+        max_tokens=1024,
+        api_key=api_key,
+    )
+
+def get_embeddings():
+    return HuggingFaceEmbeddings(
+        model_name="BAAI/bge-small-en",
+        model_kwargs={"device": "cpu"},
+        encode_kwargs={"normalize_embeddings": True},
+    )
+
+# Simple prompt template for RAG
+quiz_prompt = """
+You are an assistant specialized in answering questions based on the provided context.
+If the context does not contain the answer, reply “I don't know.”
+Context:
+{context}
+
+Question:
+{question}
+
+Answer:
+"""
+chat_prompt = ChatPromptTemplate.from_messages([
+    ("system", quiz_prompt),
+    ("human", "{question}"),
+])
+
+def create_chain(retriever):
+    return ConversationalRetrievalChain.from_llm(
+        llm=get_llm(),
+        retriever=retriever,
+        return_source_documents=True,
+        chain_type="stuff",
+        combine_docs_chain_kwargs={"prompt": chat_prompt},
+        verbose=False,
+    )
+
+# In-memory session store
+sessions: dict[str, dict] = {}
+
+def process_transcription(text: str) -> str:
+    # split → embed → index → store retriever & empty history
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=20)
+    chunks = splitter.split_text(text)
+    vs = FAISS.from_texts(chunks, get_embeddings())
+    retr = vs.as_retriever(search_kwargs={"k": 3})
+    sid = str(uuid.uuid4())
+    sessions[sid] = {"retriever": retr, "history": []}
+    return sid
+
+# ——— Endpoints ———————————————————————————————————————————
+
+class URLIn(BaseModel):
+    youtube_url: str
+
+@router.post("/transcribe_video")
+async def transcribe_url(body: URLIn):
+    client = init_google_client()
+    try:
+        resp = client.models.generate_content(
+            model="models/gemini-2.0-flash",
+            contents=types.Content(parts=[
+                types.Part(text="Transcribe the video"),
+                types.Part(file_data=types.FileData(file_uri=body.youtube_url))
+            ])
+        )
+        txt = resp.candidates[0].content.parts[0].text
+        sid = process_transcription(txt)
+        return {"session_id": sid}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@router.post("/upload_video")
+async def upload_file(
+    file: UploadFile = File(...),
+    prompt: str = "Transcribe the video",
+):
+    data = await file.read()
+    client = init_google_client()
+    try:
+        resp = client.models.generate_content(
+            model="models/gemini-2.0-flash",
+            contents=types.Content(parts=[
+                types.Part(text=prompt),
+                types.Part(inline_data=types.Blob(data=data, mime_type=file.content_type))
+            ])
+        )
+        txt = resp.candidates[0].content.parts[0].text
+        sid = process_transcription(txt)
+        return {"session_id": sid}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+class QueryIn(BaseModel):
+    session_id: str
+    query: str
+
+@router.post("/vid_query")
+async def query_rag(body: QueryIn):
+    sess = sessions.get(body.session_id)
+    if not sess:
+        raise HTTPException(status_code=404, detail="Session not found")
+    chain = create_chain(sess["retriever"])
+    result = chain({
+        "question": body.query,
+        "chat_history": sess["history"]
+    })
+    answer = result.get("answer", "I don't know.")
+    # update history
+    sess["history"].append((body.query, answer))
+    # collect source snippets
+    docs = result.get("source_documents") or []
+    srcs = [getattr(d, "page_content", str(d)) for d in docs]
+    return {"answer": answer, "source_documents": srcs}