Update app.py

app.py

@@ -11,8 +11,260 @@ from dotenv import load_dotenv
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
+load_dotenv()import os
+import gradio as gr
+from openai import OpenAI
+from PyPDF2 import PdfReader
+import requests
+from youtube_transcript_api import YouTubeTranscriptApi
+from urllib.parse import urlparse, parse_qs
+from pinecone import Pinecone
+import uuid
+from dotenv import load_dotenv
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
 load_dotenv()
 
+# Set up OpenAI client
+client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+# Set up Pinecone
+pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
+
+index_name = "main"  # Your index name
+index = pc.Index(index_name)
+
+def get_embedding(text):
+    response = client.embeddings.create(input=text, model="text-embedding-3-large")
+    return response.data[0].embedding
+
+def process_pdf(file):
+    reader = PdfReader(file)
+    text = ""
+    for page in reader.pages:
+        text += page.extract_text() + "\n"
+    return text
+
+def process_web_link(url):
+    response = requests.get(url)
+    return response.text
+
+def process_youtube_link(url):
+    video_id = extract_video_id(url)
+    transcript = YouTubeTranscriptApi.get_transcript(video_id)
+    return " ".join([entry['text'] for entry in transcript])
+
+def extract_video_id(url):
+    parsed_url = urlparse(url)
+    if parsed_url.hostname == 'youtu.be':
+        return parsed_url.path[1:]
+    if parsed_url.hostname in ('www.youtube.com', 'youtube.com'):
+        if parsed_url.path == '/watch':
+            return parse_qs(parsed_url.query)['v'][0]
+        if parsed_url.path[:7] == '/embed/':
+            return parsed_url.path.split('/')[2]
+        if parsed_url.path[:3] == '/v/':
+            return parsed_url.path.split('/')[2]
+    return None
+
+def process_upload(upload_type, file_or_link, file_name=None):
+    print(f"Starting process_upload for {upload_type}")
+    doc_id = str(uuid.uuid4())
+    print(f"Generated doc_id: {doc_id}")
+
+    if upload_type == "PDF":
+        content = process_pdf(file_or_link)
+        doc_name = file_name or "Uploaded PDF"
+    elif upload_type == "Web Link":
+        content = process_web_link(file_or_link)
+        doc_name = file_or_link
+    elif upload_type == "YouTube Link":
+        content = process_youtube_link(file_or_link)
+        doc_name = f"YouTube: {file_or_link}"
+    else:
+        print("Invalid upload type")
+        return "Invalid upload type"
+
+    content_length = len(content)
+    print(f"Content extracted, length: {content_length}")
+
+    # Dynamically adjust chunk size based on content length
+    if content_length < 10000:
+        chunk_size = 1000
+    elif content_length < 100000:
+        chunk_size = 2000
+    else:
+        chunk_size = 4000
+    print(f"Using chunk size: {chunk_size}")
+
+    chunks = [content[i:i+chunk_size] for i in range(0, content_length, chunk_size)]
+    
+    vectors = []
+    with ThreadPoolExecutor() as executor:
+        futures = [executor.submit(process_chunk, chunk, doc_id, i, upload_type, doc_name) for i, chunk in enumerate(chunks)]
+        
+        for future in as_completed(futures):
+            vectors.append(future.result())
+            # Progress can be handled via logging or status messages
+
+    print(f"Generated {len(vectors)} vectors")
+    
+    index.upsert(vectors=vectors)
+    print("Vectors upserted to Pinecone")
+    
+    return f"Processing complete for {upload_type}. Document Name: {doc_name}"
+
+def process_chunk(chunk, doc_id, i, upload_type, doc_name):
+    embedding = get_embedding(chunk)
+    return (f"{doc_id}_{i}", embedding, {
+        "text": chunk, 
+        "type": upload_type,
+        "doc_id": doc_id,
+        "doc_name": doc_name,
+        "chunk_index": i
+    })
+
+def get_relevant_context(query, top_k=5):
+    print(f"Getting relevant context for query: {query}")
+    query_embedding = get_embedding(query)
+    
+    search_results = index.query(vector=query_embedding, top_k=top_k, include_metadata=True)
+    print(f"Found {len(search_results['matches'])} relevant results")
+    
+    # Sort results by doc_id and chunk_index to maintain document structure
+    sorted_results = sorted(search_results['matches'], key=lambda x: (x['metadata']['doc_id'], x['metadata']['chunk_index']))
+    
+    context = "\n".join([result['metadata']['text'] for result in sorted_results])
+    return context, sorted_results
+
+def chat_with_ai(message):
+    print(f"Chatting with AI, message: {message}")
+    context, results = get_relevant_context(message)
+    print(f"Retrieved context, length: {len(context)}")
+    
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant. Use the following information to answer the user's question, but don't mention the context directly in your response. If the information isn't in the context, say you don't know."},
+        {"role": "system", "content": f"Context: {context}"},
+        {"role": "user", "content": message}
+    ]
+    
+    response = client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=messages
+    )
+    print("Received response from OpenAI")
+    
+    ai_response = response.choices[0].message.content
+    
+    # Prepare source information
+    sources = [
+        {
+            "doc_id": result['metadata']['doc_id'],
+            "doc_name": result['metadata']['doc_name'],
+            "chunk_index": result['metadata']['chunk_index'],
+            "text": result['metadata']['text'],
+            "type": result['metadata']['type']
+        }
+        for result in results
+    ]
+    
+    return ai_response, sources
+
+def clear_database():
+    print("Clearing database...")
+    index.delete(delete_all=True)
+    print("Database cleared")
+    return "Database cleared successfully."
+
+# Gradio Interface Components
+
+def handle_uploads(pdf, web_link, youtube_link):
+    results = []
+    if pdf:
+        pdf_result = process_upload("PDF", pdf, pdf.name)
+        results.append(pdf_result)
+    if web_link:
+        web_result = process_upload("Web Link", web_link)
+        results.append(web_result)
+    if youtube_link:
+        youtube_result = process_upload("YouTube Link", youtube_link)
+        results.append(youtube_result)
+    
+    if results:
+        return "\n".join([f"✅ {res}" for res in results])
+    else:
+        return "⚠️ No content uploaded. Please provide at least one input."
+
+def handle_chat(user_input, state):
+    if not user_input:
+        return "⚠️ Please enter a question.", state
+    
+    response, sources = chat_with_ai(user_input)
+    state = sources  # Update state with sources
+    return f"**You:** {user_input}\n\n**AI:** {response}", state
+
+def handle_clear_database():
+    result = clear_database()
+    return result
+
+def display_sources(sources):
+    if not sources:
+        return "ℹ️ Ask a question to see source chunks here."
+    
+    source_texts = []
+    for i, source in enumerate(sources, 1):
+        source_text = f"**Source {i} - {source['type']} ({source['doc_name']})**\n\n" \
+                      f"**Chunk Index:** {source['chunk_index']}\n" \
+                      f"{source['text']}\n\n---\n"
+        source_texts.append(source_text)
+    return "\n".join(source_texts)
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 📄 Upload and Chat with PDFs, Web Links, and YouTube Videos")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("## 📤 Upload")
+            
+            pdf_input = gr.File(label="Choose a PDF file", file_types=[".pdf"])
+            web_link_input = gr.Textbox(label="Enter a Web Link", placeholder="https://example.com")
+            youtube_link_input = gr.Textbox(label="Enter a YouTube Link", placeholder="https://youtube.com/watch?v=...")
+            
+            upload_button = gr.Button("Process All")
+            clear_db_button = gr.Button("Clear Database")
+            upload_output = gr.Markdown()
+        
+        with gr.Column(scale=1):
+            gr.Markdown("## 💬 Chat")
+            user_input = gr.Textbox(label="Ask a question about the uploaded content:", placeholder="Your question here...")
+            chat_button = gr.Button("Send")
+            chat_output = gr.Markdown()
+        
+        with gr.Column(scale=1):
+            gr.Markdown("## 📚 Source Chunks")
+            sources_display = gr.Markdown()
+    
+    # Hidden state to store sources
+    state = gr.State([])
+    
+    # Define interactions
+    upload_button.click(handle_uploads, inputs=[pdf_input, web_link_input, youtube_link_input], outputs=upload_output)
+    clear_db_button.click(handle_clear_database, inputs=None, outputs=upload_output)
+    chat_button.click(handle_chat, inputs=[user_input, state], outputs=[chat_output, state])
+    state.change(display_sources, inputs=state, outputs=sources_display)
+    
+    # Alternatively, use an event to update sources_display when state changes
+    def update_sources(sources):
+        return display_sources(sources)
+    
+    chat_button.click(update_sources, inputs=state, outputs=sources_display)
+
+# Launch the Gradio app
+if __name__ == "__main__":
+    demo.launch()
+
+
 # Set up OpenAI client
 client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))