Create main.py

main.py

@@ -0,0 +1,310 @@
+from fastapi import FastAPI, HTTPException, Body, Query, File, UploadFile, Form
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import List, Optional, Dict, Any, Union
+import uuid
+import os
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Import necessary libraries
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.chains import ConversationalRetrievalChain
+from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
+from langchain_groq import ChatGroq
+from google import genai
+from google.genai import types
+
+# Initialize FastAPI app
+app = FastAPI(title="RAG System API", description="An API for question answering based on YouTube video content or uploaded video files")
+
+# Configure CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Define models
+class TranscriptionRequest(BaseModel):
+    youtube_url: str
+
+class QueryRequest(BaseModel):
+    query: str
+    session_id: Optional[str] = None
+
+class QueryResponse(BaseModel):
+    answer: str
+    session_id: str
+    source_documents: Optional[List[str]] = None
+
+# Global variables
+sessions = {}
+
+# Initialize Google API client
+def init_google_client():
+    api_key = os.getenv("GOOGLE_API_KEY", "")
+    if not api_key:
+        raise ValueError("GOOGLE_API_KEY environment variable not set")
+    return genai.Client(api_key=api_key)
+
+# Get LLM
+def get_llm():
+    """
+    Returns the language model instance (LLM) using ChatGroq API.
+    The LLM used is Llama 3.1 with a versatile 70 billion parameters model.
+    """
+    api_key = os.getenv("GROQ_API_KEY", "")
+    if not api_key:
+        raise ValueError("GROQ_API_KEY environment variable not set")
+    
+    llm = ChatGroq(
+        model="llama-3.3-70b-versatile",
+        temperature=0,
+        max_tokens=1024,
+        api_key=api_key
+    )
+    return llm
+
+# Get embeddings
+def get_embeddings():
+    model_name = "BAAI/bge-small-en"
+    model_kwargs = {"device": "cpu"}
+    encode_kwargs = {"normalize_embeddings": True}
+    embeddings = HuggingFaceBgeEmbeddings(
+        model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs
+    )
+    return embeddings
+
+# Create prompt template
+quiz_solving_prompt = '''
+You are an assistant specialized in solving quizzes. Your goal is to provide accurate, concise, and contextually relevant answers.
+Use the following retrieved context to answer the user's question.
+If the context lacks sufficient information, respond with "I don't know." Do not make up answers or provide unverified information.
+
+Guidelines:
+1. Extract key information from the context to form a coherent response.
+2. Maintain a clear and professional tone.
+3. If the question requires clarification, specify it politely.
+
+Retrieved context:
+{context}
+
+User's question:
+{question}
+
+Your response:
+'''
+
+# Create a prompt template to pass the context and user input to the chain
+user_prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", quiz_solving_prompt),
+        ("human", "{question}"),
+    ]
+)
+
+# Create a chain
+def create_chain(retriever):
+    llm = get_llm()
+    chain = ConversationalRetrievalChain.from_llm(
+        llm=llm,
+        retriever=retriever,
+        return_source_documents=True,
+        chain_type='stuff',
+        combine_docs_chain_kwargs={"prompt": user_prompt},
+        verbose=False,
+    )
+    return chain
+
+# Process transcription and prepare RAG system
+def process_transcription(transcription):
+    # Process the transcription
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=20)
+    all_splits = text_splitter.split_text(transcription)
+    
+    # Create vector store
+    embeddings = get_embeddings()
+    vectorstore = FAISS.from_texts(all_splits, embeddings)
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+    
+    # Create a session ID
+    session_id = str(uuid.uuid4())
+    
+    # Store session data
+    sessions[session_id] = {
+        "retriever": retriever,
+        "chat_history": [],
+        "transcription": transcription
+    }
+    
+    return session_id
+
+@app.post("/transcribe", response_model=Dict[str, str])
+async def transcribe_video(request: TranscriptionRequest):
+    """
+    Transcribe a YouTube video and prepare the RAG system
+    """
+    try:
+        # Initialize Google API client
+        client = init_google_client()
+        
+        # Transcribe the video
+        response = client.models.generate_content(
+            model='models/gemini-2.0-flash',
+            contents=types.Content(
+                parts=[
+                    types.Part(text='Transcribe the Video. Write all the things described in the video'),
+                    types.Part(
+                        file_data=types.FileData(file_uri=request.youtube_url)
+                    )
+                ]
+            )
+        )
+        
+        # Get transcription text
+        transcription = response.candidates[0].content.parts[0].text
+        
+        # Process transcription and get session ID
+        session_id = process_transcription(transcription)
+        
+        return {"session_id": session_id, "message": "YouTube video transcribed and RAG system prepared"}
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error transcribing video: {str(e)}")
+
+@app.post("/upload", response_model=Dict[str, str])
+async def upload_video(file: UploadFile = File(...), prompt: str = Form("Transcribe the Video. Write all the things described in the video")):
+    """
+    Upload a video file (max 20MB), transcribe it and prepare the RAG system
+    """
+    try:
+        # Check file size (20MB limit)
+        contents = await file.read()
+        if len(contents) > 20 * 1024 * 1024:  # 20MB in bytes
+            raise HTTPException(status_code=400, detail="File size exceeds 20MB limit")
+        
+        # Check file type
+        if not file.content_type.startswith('video/'):
+            raise HTTPException(status_code=400, detail="File must be a video")
+        
+        # Initialize Google API client
+        client = init_google_client()
+        
+        # Transcribe the video
+        response = 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=contents, mime_type=file.content_type)
+                    )
+                ]
+            )
+        )
+        
+        # Get transcription text
+        transcription = response.candidates[0].content.parts[0].text
+        
+        # Process transcription and get session ID
+        session_id = process_transcription(transcription)
+        
+        return {"session_id": session_id, "message": "Uploaded video transcribed and RAG system prepared"}
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing uploaded video: {str(e)}")
+    finally:
+        # Reset file pointer
+        await file.seek(0)
+
+@app.post("/query", response_model=QueryResponse)
+async def query_system(request: QueryRequest):
+    """
+    Query the RAG system with a question
+    """
+    try:
+        session_id = request.session_id
+        
+        # Create a new session if none provided
+        if not session_id or session_id not in sessions:
+            raise HTTPException(status_code=404, detail="Session not found. Please transcribe a video first.")
+        
+        # Get session data
+        session = sessions[session_id]
+        retriever = session["retriever"]
+        chat_history = session["chat_history"]
+        
+        # Create chain
+        chain = create_chain(retriever)
+        
+        # Query the chain
+        result = chain({"question": request.query, "chat_history": chat_history})
+        
+        # Update chat history
+        chat_history.append((request.query, result["answer"]))
+        
+        # Prepare source documents
+        source_docs = [doc.page_content[:100] + "..." for doc in result.get("source_documents", [])]
+        
+        return {
+            "answer": result["answer"],
+            "session_id": session_id,
+            "source_documents": source_docs
+        }
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error querying system: {str(e)}")
+
+@app.get("/sessions/{session_id}", response_model=Dict[str, Any])
+async def get_session_info(session_id: str):
+    """
+    Get information about a specific session
+    """
+    if session_id not in sessions:
+        raise HTTPException(status_code=404, detail="Session not found")
+    
+    session = sessions[session_id]
+    
+    return {
+        "session_id": session_id,
+        "chat_history_length": len(session["chat_history"]),
+        "transcription_preview": session["transcription"][:200] + "..."
+    }
+
+@app.delete("/sessions/{session_id}")
+async def delete_session(session_id: str):
+    """
+    Delete a session
+    """
+    if session_id not in sessions:
+        raise HTTPException(status_code=404, detail="Session not found")
+    
+    del sessions[session_id]
+    return {"message": f"Session {session_id} deleted successfully"}
+
+@app.get("/")
+async def root():
+    """
+    API root endpoint
+    """
+    return {
+        "message": "Video Transcription and QA API",
+        "endpoints": {
+            "/transcribe": "Transcribe YouTube videos",
+            "/upload": "Upload and transcribe video files (max 20MB)",
+            "/query": "Query the RAG system",
+            "/sessions/{session_id}": "Get session information",
+        }
+    }
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)