Chat with a YouTube Video: How to build an AI chatbot using YouTube video transcripts

video_ids = ['dQw4w9WgXcQ']

from dataclasses import dataclass
from youtube_transcript_api import YouTubeTranscriptApi
from youtube_transcript_api._errors import CouldNotRetrieveTranscript


@dataclass
class Transcript:
    video_id: str
    text: str


def __get_transcript_from_video_id(video_id):
    try:
        transcript = YouTubeTranscriptApi.get_transcript(video_id)
        return "\n".join([line["text"] for line in transcript])
    except CouldNotRetrieveTranscript as e:
        # I don't include logging in the article, but you should do it in your code.
        #logger.warning(f"Could not retrieve transcript for video {video_id}: {e}")
        return None

transcripts = []
for video_id in video_ids:
    transcript = __get_transcript_from_video_id(video_id)
    if transcript:
        transcripts.append(Transcript(video_id, transcript))

QA_PROMPT_TEMPLATE = """Given a transcript of a YouTube video. Prepare a list of questions, answers, and quotes. Put every question, quote, and answer inside <qa> </qa> tags.
Include only topics that may be helpful when answering questions about {topic}. Don't include information about the next video, newsletters, gifts, products sold by the author, etc.

Format:
---
<qa>
Question: the question
Quote: A quote from the video that answers the question
Answer: A short answer
</qa>
---
"""

def __prepare_question_and_answer_extraction_prompt(topic_description: str) -> str:
    return QA_PROMPT_TEMPLATE.format(topic=topic_description)

import re
import openai
from tqdm import tqdm
openai.api_key = ... # put your API key here
from langchain.text_splitter import RecursiveCharacterTextSplitter
from llama_index import Document


@dataclass
class QuestionAndAnswer:
    video_id: str
    content: str

    def to_document(self) -> Document:
        return Document(
            text=self.content,
            metadata={"video_id": self.video_id},
        )


def get_questions_and_answers_from_transcripts(
    transcripts: List[Transcript], topic_description: str
) -> List[QuestionAndAnswer]:
    prompt = __prepare_question_and_answer_extraction_prompt(topic_description)
    questions_and_answers = []

    for transcript in tqdm(
        transcripts, desc="Extracting questions and answers from transcripts"
    ):
        questions_and_answers.extend(
            __get_questions_and_answers_from_single_transcript(transcript, prompt)
        )
    return questions_and_answers

def __get_questions_and_answers_from_single_transcript(
    transcript: Transcript, prompt: str
) -> List[QuestionAndAnswer]:
    chunks = __split_transcript(transcript)
    questions_and_answers = []
    for chunk in tqdm(chunks, desc="Extracting questions and answers from chunks"):
        QAs = __extract_questions_and_answers(chunk, prompt)
        for qa in QAs:
            questions_and_answers.append(QuestionAndAnswer(transcript.video_id, qa))
    return questions_and_answers


def __extract_questions_and_answers(chunk_content: str, prompt: str) -> List[str]:
    final_result = []

    messages = [
        {"role": "system", "content": prompt},
        {"role": "user", "content": chunk_content},
    ]
    chat_completion = openai.ChatCompletion.create(
        model="gpt-3.5-turbo", max_tokens=1000, messages=messages, temperature=0.0
    )

    response = chat_completion["choices"][0]["message"]["content"]
    pattern = r"<qa>(.*?)<\/qa>"
    qa_blocks = re.findall(pattern, response, re.DOTALL)
    final_result.extend(qa_blocks)

    return final_result

topic = "ChatGPT, AI, the future of AI"
QAs = get_questions_and_answers_from_transcripts(transcripts, topic)

from typing import List

import chromadb
from llama_index import VectorStoreIndex, ServiceContext
from llama_index.vector_stores import ChromaVectorStore
from llama_index.storage.storage_context import StorageContext
from llama_index.tools import QueryEngineTool, ToolMetadata
from llama_index.embeddings import HuggingFaceEmbedding
from llama_index.storage.storage_context import StorageContext
from .transcripts import QuestionAndAnswer # !!!CHANGE `.transcripts` TO THE MODULE WITH THE `QuestionAndAnswer` CLASS!!!

class Database:
    def __init__(self, database_path: str, topic: str):
        self.topic = topic
        db = chromadb.PersistentClient(path=database_path)

        chroma_collection = db.get_or_create_collection("videos")
        self.vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

        embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-base-en-v1.5")
        self.storage_context = StorageContext.from_defaults(
            vector_store=self.vector_store
        )
        self.service_context = ServiceContext.from_defaults(embed_model=embed_model)

# This is a continuation of the previous code snippet
# put it inside the Database class
def insert(self, questions_and_answers: List[QuestionAndAnswer]):
    QAs = [qa.to_document() for qa in questions_and_answers]
    VectorStoreIndex.from_documents(
        QAs,
        storage_context=self.storage_context,
        service_context=self.service_context,
    )

database_path = "chroma_db"
database = Database(database_path, topic)
database.insert(QAs)

# This is a continuation of the Database class !!!
def to_qa_tool(self, tool_name: str) -> QueryEngineTool:
    index = VectorStoreIndex.from_vector_store(
        self.vector_store,
        service_context=self.service_context,
    )
    video_query_engine_tool = QueryEngineTool(
        query_engine=index.as_query_engine(),
        metadata=ToolMetadata(
            name=tool_name,
            description=f"useful for when you want to answer questions about {self.topic}",
        ),
    )
    return video_query_engine_tool

from .database import Database
from llama_index.agent import OpenAIAgent


class Chatbot:
    def __init__(self, database: Database, verbose: bool = False):
        self.tool_name = "database"
        tool = database.to_qa_tool(self.tool_name)
        self.prompt_template = f"Use the tool: {self.tool_name} to answer the question. Question: "
        self.agent = OpenAIAgent.from_tools([tool], verbose=verbose)

    def ask(self, question: str) -> str:
        prompt = self.prompt_template.format(question=question)
        answer = self.agent.chat(prompt)
        return answer.response

from llama_index.llms.base import ChatMessage, MessageRole


# Put the following methods inside the Chatbot class
def ask_with_history(self, question: str, history: List[Dict]) -> str:
        prompt = self.prompt_template.format(question=question)
        chat_history = self.__convert_history(history)

        answer = self.agent.chat(prompt, chat_history=chat_history)
        return answer.response

def __convert_history(self, history: List[Dict]) -> List[ChatMessage]:
    return [
        ChatMessage(
            content=msg["content"],
            role=MessageRole(msg["role"])
        )
        for msg in history if msg["content"] is not None
    ]

import streamlit as st

from chatbot.database import Database
from chatbot.chat import Chatbot


# use the database_path and topic defined earlier
database = Database(database_path, topic)
chatbot = Chatbot(database)

st.title("Chat with Youtube Videos")
st.subheader(f"Ask a question about {topic}")

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("What is up?"):
    chat_history = list(st.session_state.messages)

    with st.chat_message("user"):
        st.markdown(prompt)

    st.session_state.messages.append({"role": "user", "content": prompt})

    answer = chatbot.ask_with_history(prompt, chat_history)

    with st.chat_message("assistant"):
        st.markdown(answer)

    st.session_state.messages.append({"role": "assistant", "content": answer})