File size: 3,032 Bytes
2df8973 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import os
from youtube_transcript_api import YouTubeTranscriptApi, TranscriptsDisabled
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import PromptTemplate
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_groq import ChatGroq
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableParallel, RunnablePassthrough, RunnableLambda
VECTORSTORE_DIR = "vectorstores"
EMBEDDING_MODEL_NAME = os.environ["EMBEDDING_MODEL"]
_embedding_model = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
def load_transcript(video_id: str) -> str:
"""Fetch and flatten a YouTube video's transcript."""
try:
yt_api = YouTubeTranscriptApi()
transcript_list = yt_api.fetch(video_id, languages=["en"])
transcript = " ".join(chunk.text for chunk in transcript_list)
return transcript
except TranscriptsDisabled:
raise ValueError("No captions available for this video.")
def split_transcript(transcript: str):
splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200,
)
return splitter.create_documents([transcript])
def get_or_build_vectorstore(video_id: str):
"""Load a cached FAISS index for this video, or build + save one if missing."""
path = os.path.join(VECTORSTORE_DIR, video_id)
if os.path.exists(path):
return FAISS.load_local(
path,
_embedding_model,
allow_dangerous_deserialization=True,
)
transcript = load_transcript(video_id)
chunks = split_transcript(transcript)
vectorstore = FAISS.from_documents(chunks, _embedding_model)
os.makedirs(VECTORSTORE_DIR, exist_ok=True)
vectorstore.save_local(path)
return vectorstore
def format_docs(retrieved_docs):
return "\n\n".join(doc.page_content for doc in retrieved_docs)
def get_rag_chain(video_id: str):
"""Builds and returns the full RAG chain for a given YouTube video."""
vectorstore = get_or_build_vectorstore(video_id)
retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
llm = ChatGroq(
model="llama-3.1-8b-instant",
temperature=0,
api_key=os.environ["GROQ_API_KEY"],
)
prompt = PromptTemplate(
template="""
You are a helpful assistant.
Answer ONLY from the provided transcript context.
If the context is insufficient, just say you don't know.
{context}
Question: {question}
""",
input_variables=["context", "question"],
)
parser = StrOutputParser()
parallel_chain = RunnableParallel({
"context": retriever | RunnableLambda(format_docs),
"question": RunnablePassthrough(),
})
main_chain = parallel_chain | prompt | llm | parser
return main_chain |