documentQA.js

import { openai } from './openai.js'
import { Document } from '@langchain/core/documents'
import { MemoryVectorStore } from '@langchain/classic/vectorstores/memory'
import { OpenAIEmbeddings } from '@langchain/openai'
import { CharacterTextSplitter } from '@langchain/textsplitters'
import { PDFLoader } from '@langchain/community/document_loaders/fs/pdf'
import { YoutubeTranscript } from '@danielxceron/youtube-transcript'

const video = 'https://www.youtube.com/watch?v=zR_iuq2evXo'
const question = process.argv[2] || 'HI'

const createStore = (docs) =>
  MemoryVectorStore.fromDocuments(docs, new OpenAIEmbeddings())

const docsFromYoutubeVideo = async (video) => {
  const transcript = await YoutubeTranscript.fetchTranscript(video)
  const text = transcript.map((t) => t.text).join(' ')

  const youtubeDoc = new Document({
    pageContent: text,
    metadata: { source: 'youtube' },
  })

  const splitter = new CharacterTextSplitter({
    separator: ' ',
    chunkSize: 2500,
    chunkOverlap: 100,
  })

  return splitter.splitDocuments([youtubeDoc])
}

const docsFromPDF = async () => {
  const loader = new PDFLoader('xbox.pdf')
  const pdfDoc = await loader.load()

  const splitter = new CharacterTextSplitter({
    separator: '. ',
    chunkSize: 2500,
    chunkOverlap: 200,
  })
  return splitter.splitDocuments(pdfDoc)
}

const loadStore = async () => {
  const videoDocs = await docsFromYoutubeVideo(video)
  const pdfDocs = await docsFromPDF()
  return createStore([...videoDocs, ...pdfDocs])
}

const query = async () => {
  const store = await loadStore()
  const results = await store.similaritySearch(question, 2)
  const response = await openai.chat.completions.create({
    model: 'gpt-4',
    temperature: 0,
    messages: [
      {
        role: 'system',
        content: 'You are an AI assistant. Answer questions to the best of your ability.',
      },
      {
        role: 'user',
        content: `Answer the following question using the provided context only. If you cannot answer the question with the context, don't lie and make up stuff. Just say you need more context.
    Question: ${question}
    Context: ${results.map((r) => r.pageContent).join('\n')}`,
      },
    ],
  })

  console.log(
    `Answer: ${response.choices[0].message.content}\nSources: ${results.map((r) => r.metadata.source).join(', ')}`,
  )
}

query()