Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,34 +1,39 @@
+# -*- coding: utf-8 -*-
 """
-ChatYT Streamlit App (API-Only Version)
+ChatYT Streamlit App (LCEL Chain Version)
 
 This Streamlit app enables you to:
 * Summarise YouTube videos
 * Ask questions about the topics discussed in the video
 
-It uses Google's Gemini APIs for all AI tasks.
+It uses LangChain Expression Language (LCEL) with Google's Gemini APIs.
 """
 
 import streamlit as st
 import yt_dlp
 import os
-import textwrap
+# Corrected import: Document is now in langchain_core.documents
 from langchain_core.documents import Document
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+# Corrected import: RecursiveCharacterTextSplitter is in its own package
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_chroma import Chroma
-from langchain_google_genai import GoogleGenerativeAIEmbeddings
+from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
+# Corrected import: ChatPromptTemplate is now in langchain_core.prompts
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
 import google.generativeai as genai
-from langchain.prompts import ChatPromptTemplate
-import time  # To simulate progress
+import time
 
 # --- App Configuration ---
 st.set_page_config(
-    page_title="ChatYT",
+    page_title="ChatYT (LangChain)",
     page_icon="📺",
     layout="wide",
 )
 
 st.title("📺 ChatYT: Chat with any YouTube Video")
-st.caption("Summarize and ask questions about any YouTube video using Google's Gemini APIs.")
+st.caption("Summarize and ask questions about any YouTube video using LangChain and Google Gemini.")
 
 # --- API Key Handling ---
 GEMINI_API_KEY = st.secrets.get("GEMINI_API_KEY")
@@ -42,6 +47,7 @@ if not GEMINI_API_KEY:
     st.error("Please provide your Gemini API Key in the sidebar to continue.")
     st.stop()
 
+# Configure the genai library (still needed for file upload)
 try:
     genai.configure(api_key=GEMINI_API_KEY)
 except Exception as e:
@@ -76,6 +82,7 @@ def compress_audio(input_file, output_file="compressed.mp3"):
 def speech_to_text(audio_file):
     """
     Transcribes audio using the Gemini API.
+    (This function uses the base genai library for file upload)
     """
     try:
         model = genai.GenerativeModel("gemini-2.5-flash")
@@ -102,110 +109,65 @@ def speech_to_text(audio_file):
 @st.cache_data(show_spinner="Summarizing text...")
 def summarize_text_api(text):
     """
-    Summarizes the text using the Gemini API.
+    Summarizes the text using a LangChain chain.
     """
-    model = genai.GenerativeModel("gemini-2.5-flash")
-    prompt = f"""Please provide a concise, high-level summary of the following text:
+    # 1. Define the LLM
+    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash",
+                               temperature=0.3,
+                               google_api_key=GEMINI_API_KEY)
+    
+    # 2. Define the Prompt
+    prompt_template = """Please provide a concise, high-level summary of the following text:
     ---
     {text}
     ---
     Provide only the summary."""
+    summarize_prompt = ChatPromptTemplate.from_template(prompt_template)
+    
+    # 3. Define the Chain
+    summarize_chain = summarize_prompt | llm | StrOutputParser()
     
     try:
-        response = model.generate_content(prompt)
-        if response.candidates and response.candidates[0].content.parts:
-            return response.candidates[0].content.parts[0].text
-        else:
-            return "Error: Could not summarize text."
+        # 4. Invoke the Chain
+        response = summarize_chain.invoke({"text": text})
+        return response
     except Exception as e:
         st.error(f"An error occurred during summarization: {e}")
         return f"Error: {e}"
 
 @st.cache_data(show_spinner="Generating embeddings...")
-def generate_embeddings(text):
+def generate_embeddings_db(text):
     """
     Splits text, generates embeddings via API, and stores in ChromaDB.
+    Returns the Chroma database object.
     """
     doc = Document(page_content=text, metadata={"source": "youtube"})
+    # This now uses the imported RecursiveCharacterTextSplitter
     splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     chunks = splitter.split_documents([doc])
     
     try:
-        embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
-        # Using a unique persist_directory for each session
-        # Note: In a real-world deployed Streamlit app, this directory is temporary.
-        # For persistence, a proper vector DB server would be needed.
+        embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001",
+                                                  google_api_key=GEMINI_API_KEY)
         db = Chroma.from_documents(chunks, embeddings)
         return db
     except Exception as e:
         st.error(f"An error occurred during embedding generation: {e}")
         return None
 
-# --- Q&A Functions ---
-
-def closest(query, db):
-    """
-    Finds the most relevant text chunks from the vector database.
-    """
-    if db is None:
-        st.warning("Database not initialized.")
-        return None
-    try:
-        results = db.similarity_search(query, k=3)
-        if len(results) == 0:
-            return None
-        return results
-    except Exception as e:
-        st.error(f"Error during similarity search: {e}")
-        return None
-
-
-def create_prompt(results, question):
-    """
-    Creates a prompt for the Q&A model based on retrieved chunks.
-    """
-    PROMPT = """Answer the following questions based only on the following context:
-    {context}
-    ---
-    Answer the question based on the above context:
-    {que}
-    """
-    if not results:
-      return "Sorry, I couldn’t find anything relevant in the video transcript."
-
-    context_text = "\n\n---\n\n".join(
-        doc.page_content for doc in results
-    )
-    prompt_template = ChatPromptTemplate.from_template(PROMPT)
-    return prompt_template.format(context=context_text, que=question)
-
-def answer_llm(question, closest_chunks):
-    """
-    Answers the question using the Gemini API and context.
-    """
-    model = genai.GenerativeModel("gemini-2.5-flash")
-    prompt = create_prompt(closest_chunks, question)
-
-    if prompt == "Sorry, I couldn’t find anything relevant in the video transcript.":
-        return prompt
-
-    try:
-        response = model.generate_content(prompt)
-        if response.candidates and response.candidates[0].content.parts:
-            return response.candidates[0].content.parts[0].text
-        else:
-            return "No answer generated."
-    except Exception as e:
-        st.error(f"An error occurred during Q&A: {e}")
-        return f"Error: {e}"
+def format_docs(docs):
+    """Helper function to format retrieved documents into a string."""
+    if not docs:
+        return "No relevant context found."
+    return "\n\n---\n\n".join(doc.page_content for doc in docs)
 
 # --- Streamlit UI Components ---
 
 # Initialize session state variables
 if "summary" not in st.session_state:
     st.session_state.summary = ""
-if "db" not in st.session_state:
-    st.session_state.db = None
+if "rag_chain" not in st.session_state:
+    st.session_state.rag_chain = None
 if "video_title" not in st.session_state:
     st.session_state.video_title = ""
 if "chat_history" not in st.session_state:
@@ -217,6 +179,12 @@ if st.button("Process Video", key="process_video"):
     if url:
         with st.spinner("Processing video... This may take a few minutes."):
             try:
+                # Reset state
+                st.session_state.summary = ""
+                st.session_state.rag_chain = None
+                st.session_state.video_title = ""
+                st.session_state.chat_history = []
+
                 # 1. Download
                 audio_file, video_title = download_audio(url)
                 st.session_state.video_title = video_title
@@ -230,20 +198,41 @@ if st.button("Process Video", key="process_video"):
                     st.error(f"Failed to transcribe: {text}")
                     st.stop()
                 
-                # 4. Summarize
+                # 4. Summarize (using the new chain function)
                 summary = summarize_text_api(text)
-                if "Error:" in summary:
-                    st.error(f"Failed to summarize: {summary}")
-                    st.session_state.summary = "Could not generate summary."
-                else:
-                    st.session_state.summary = summary
+                st.session_state.summary = summary
                 
-                # 5. Embed
-                db = generate_embeddings(text)
-                if db is None:
-                    st.error("Failed to create vector database.")
+                # 5. Embed and create DB
+                db = generate_embeddings_db(text)
+                
+                if db:
+                    # 6. Create RAG Chain and store it in session state
+                    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash",
+                                               temperature=0.3,
+                                               google_api_key=GEMINI_API_KEY)
+                    
+                    retriever = db.as_retriever(search_kwargs={"k": 3})
+                    
+                    PROMPT_TEMPLATE = """Answer the following questions based only on the following context:
+                    {context}
+                    ---
+                    Answer the question based on the above context:
+                    {question}
+                    """
+                    prompt = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
+                    
+                    # This is the RAG chain
+                    rag_chain = (
+                        {"context": retriever | format_docs, "question": RunnablePassthrough()}
+                        | prompt
+                        | llm
+                        | StrOutputParser()
+                    )
+                    
+                    st.session_state.rag_chain = rag_chain
+                    st.success("Video processed and Q&A chat is ready!")
                 else:
-                    st.session_state.db = db
+                    st.error("Failed to create vector database.")
                 
                 # Clean up local files
                 try:
@@ -252,8 +241,6 @@ if st.button("Process Video", key="process_video"):
                 except OSError as e:
                     st.warning(f"Could not clean up audio files: {e}")
                     
-                st.success("Video processed successfully!")
-                
             except Exception as e:
                 st.error(f"An error occurred during video processing: {e}")
     else:
@@ -274,20 +261,21 @@ if st.session_state.summary:
             
     # Chat input
     if prompt := st.chat_input("Ask a question about the video..."):
-        if st.session_state.db:
+        if st.session_state.rag_chain:
             # Add user message to history
             st.session_state.chat_history.append(("user", prompt))
             with st.chat_message("user"):
                 st.markdown(prompt)
                 
-            # Generate and display bot response
+            # Generate and display bot response by invoking the chain
             with st.chat_message("assistant"):
                 with st.spinner("Thinking..."):
-                    chunks = closest(prompt, st.session_state.db)
-                    answer = answer_llm(prompt, chunks)
+                    # Here we just invoke the chain with the prompt!
+                    answer = st.session_state.rag_chain.invoke(prompt)
                     st.markdown(answer)
                     
             # Add bot message to history
             st.session_state.chat_history.append(("assistant", answer))
         else:
-            st.error("The vector database is not loaded. Please process a video first.")
+            st.error("The Q&A chain is not ready. Please process a video first.")
+