rag_query.py

import sqlite3
import numpy as np
import openai
import argparse
from typing import List, Dict, Any, Tuple
from sklearn.metrics.pairwise import cosine_similarity
import os
from dotenv import load_dotenv

load_dotenv()

class YouTubeRAG:
    def __init__(self, db_path: str = "youtube_embeddings.db"):
        openai_api_key = os.getenv('OPENAI_API_KEY')
        
        if not openai_api_key:
            raise ValueError("OPENAI_API_KEY not found in environment variables")
        
        self.db_path = db_path
        self.client = openai.OpenAI(api_key=openai_api_key)
    
    def search_similar_chunks(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
        try:
            response = self.client.embeddings.create(
                model="text-embedding-3-small",
                input=query
            )
            query_embedding = np.array(response.data[0].embedding, dtype=np.float32)
        except Exception as e:
            print(f"Error generating query embedding: {e}")
            return []
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT video_id, video_title, chunk_text, chunk_index, 
                   start_time, end_time, embedding
            FROM video_chunks
        ''')
        
        results = []
        for row in cursor.fetchall():
            stored_embedding = np.frombuffer(row[6], dtype=np.float32)
            similarity = cosine_similarity([query_embedding], [stored_embedding])[0][0]
            
            results.append({
                'video_id': row[0],
                'video_title': row[1],
                'chunk_text': row[2],
                'chunk_index': row[3],
                'start_time': row[4],
                'end_time': row[5],
                'similarity': similarity
            })
        
        conn.close()
        
        results.sort(key=lambda x: x['similarity'], reverse=True)
        return results[:top_k]
    
    def format_context(self, chunks: List[Dict[str, Any]]) -> str:
        context = ""
        current_video = None
        
        for chunk in chunks:
            if chunk['video_id'] != current_video:
                context += f"\n\n--- Video: {chunk['video_title']} ---\n"
                current_video = chunk['video_id']
            
            minutes = int(chunk['start_time'] // 60)
            seconds = int(chunk['start_time'] % 60)
            context += f"[{minutes}:{seconds:02d}] {chunk['chunk_text']}\n"
        
        return context
    
    def generate_answer(self, query: str, context: str) -> str:
        prompt = f"""Based solely on the following YouTube video context, answer the user's question.

Context:
{context}

Question: {query}

Instructions:
- Answer only with information present in the context
- If information is not available in the context, say you don't have sufficient information
- Include references to videos and timestamps when relevant
- Respond clearly and concisely in English

Answer:"""

        try:
            response = self.client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": "You are an assistant that answers questions based solely on YouTube video transcripts provided as context."},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=1000,
                temperature=0.3
            )
            return response.choices[0].message.content
        except Exception as e:
            return f"Error generating response: {e}"
    
    def query(self, question: str, top_k: int = 5) -> Dict[str, Any]:
        print(f"Searching for information about: {question}")
        
        similar_chunks = self.search_similar_chunks(question, top_k)
        
        if not similar_chunks:
            return {
                'answer': "No relevant information found in the database.",
                'sources': []
            }
        
        context = self.format_context(similar_chunks)
        answer = self.generate_answer(question, context)
        
        sources = []
        for chunk in similar_chunks:
            minutes = int(chunk['start_time'] // 60)
            seconds = int(chunk['start_time'] % 60)
            sources.append({
                'video_title': chunk['video_title'],
                'video_id': chunk['video_id'],
                'timestamp': f"{minutes}:{seconds:02d}",
                'similarity': round(chunk['similarity'], 3),
                'url': f"https://youtube.com/watch?v={chunk['video_id']}&t={int(chunk['start_time'])}s"
            })
        
        return {
            'answer': answer,
            'sources': sources
        }
    
    def interactive_mode(self):
        print("=== YouTube RAG System ===")
        print("Type 'exit' to quit\n")
        
        while True:
            question = input("Question: ").strip()
            
            if question.lower() in ['exit', 'quit']:
                break
            
            if not question:
                continue
            
            result = self.query(question)
            
            print(f"\nAnswer:\n{result['answer']}\n")
            
            if result['sources']:
                print("Sources:")
                for i, source in enumerate(result['sources'][:3], 1):
                    print(f"{i}. {source['video_title']} [{source['timestamp']}] (similarity: {source['similarity']})")
                    print(f"   {source['url']}")
                print()

def main():
    parser = argparse.ArgumentParser(description='Query YouTube embeddings database')
    parser.add_argument('--db-path', default='youtube_embeddings.db', help='Database path')
    parser.add_argument('--query', help='Single query mode')
    parser.add_argument('--top-k', type=int, default=5, help='Number of results to retrieve')
    
    args = parser.parse_args()
    
    rag = YouTubeRAG(args.db_path)
    
    if args.query:
        result = rag.query(args.query, args.top_k)
        print(f"Answer:\n{result['answer']}\n")
        
        if result['sources']:
            print("Sources:")
            for i, source in enumerate(result['sources'], 1):
                print(f"{i}. {source['video_title']} [{source['timestamp']}]")
                print(f"   {source['url']}")
    else:
        rag.interactive_mode()

if __name__ == "__main__":
    main()