adding google gemini

Author: wjohnsto

services/video-search/.env.example

@@ -5,16 +5,26 @@ REDIS_URL=<redis[s]://[[username][:password]@][host][:port][/db-number]>
 SEARCHAPI_API_KEY=<https://www.searchapi.io/>
 YOUTUBE_VIDEOS=<video-id-1,video-id-2,video-id-3>
 
+GOOGLE_API_KEY=<https://console.cloud.google.com/apis/credentials>
+GOOGLE_EMBEDDING_MODEL=<https://ai.google.dev/models/gemini#model_variations>
+GOOGLE_SUMMARY_MODEL=<https://ai.google.dev/models/gemini#model_variations>
+GOOGLE_VIDEO_INDEX_NAME=<redis-video-index-name>
+GOOGLE_VIDEO_PREFIX=<redis-video-prefix>
+GOOGLE_VECTOR_SET=<redis-vector-set>
+GOOGLE_SUMMARY_PREFIX=<redis-summary-prefix>
+
+HF_EMBEDDING_MODEL=<https://huggingface.co/models?pipeline_tag=feature-extraction&library=transformers.js>
+HF_SUMMARY_MODEL=<https://huggingface.co/models?pipeline_tag=summarization&library=transformers.js>
+HF_VIDEO_INDEX_NAME=<redis-video-index-name>
+HF_VIDEO_PREFIX=<redis-video-prefix>
+HF_VECTOR_SET=<redis-vector-set>
+HF_SUMMARY_PREFIX=<redis-summary-prefix>
+
 OPENAI_API_KEY=<https://platform.openai.com/api-keys>
 OPENAI_ORGANIZATION=<https://platform.openai.com/account/organization>
 OPENAI_EMBEDDING_MODEL=<https://platform.openai.com/account/limits>
 OPENAI_SUMMARY_MODEL=<https://platform.openai.com/account/limits>
 OPENAI_VIDEO_INDEX_NAME=<redis-video-index-name>
 OPENAI_VIDEO_PREFIX=<redis-video-prefix>
 OPENAI_VECTOR_SET=<redis-vector-set>
-
-HF_EMBEDDING_MODEL=<https://huggingface.co/models?pipeline_tag=feature-extraction&library=transformers.js>
-HF_SUMMARY_MODEL=<https://huggingface.co/models?pipeline_tag=summarization&library=transformers.js>
-HF_VIDEO_INDEX_NAME=<redis-video-index-name>
-HF_VIDEO_PREFIX=<redis-video-prefix>
-HF_VECTOR_SET=<redis-vector-set>
+OPENAI_SUMMARY_PREFIX=<redis-summary-prefix>

services/video-search/package-lock.json

@@ -18,6 +18,7 @@
     "watch": "nodemon --watch dist dist/index.js"
   },
   "dependencies": {
+    "@langchain/google-genai": "^0.0.5",
     "@tensorflow-models/mobilenet": "^2.1.1",
     "@tensorflow/tfjs": "^4.15.0",
     "@tensorflow/tfjs-node": "^4.15.0",

services/video-search/package.json

@@ -10,20 +10,35 @@ const {
   YOUTUBE_VIDEOS,
   REDIS_URL,
   SEARCHAPI_API_KEY,
+  GOOGLE_VIDEO_INDEX_NAME,
+  GOOGLE_VIDEO_PREFIX,
+  GOOGLE_API_KEY,
+  GOOGLE_EMBEDDING_MODEL,
+  GOOGLE_SUMMARY_MODEL,
+  GOOGLE_VECTOR_SET,
+  GOOGLE_SUMMARY_PREFIX,
   HF_VIDEO_INDEX_NAME,
   HF_VIDEO_PREFIX,
   HF_EMBEDDING_MODEL,
   HF_SUMMARY_MODEL,
   HF_VECTOR_SET,
+  HF_SUMMARY_PREFIX,
   OPENAI_VIDEO_INDEX_NAME,
+  OPENAI_VIDEO_PREFIX,
   OPENAI_API_KEY,
   OPENAI_ORGANIZATION,
   OPENAI_EMBEDDING_MODEL,
   OPENAI_SUMMARY_MODEL,
   OPENAI_VECTOR_SET,
+  OPENAI_SUMMARY_PREFIX,
   USE,
 } = process.env;
 
+const DEFAULT_VIDEO_INDEX_NAME = 'idx-videos';
+const DEFAULT_VIDEO_PREFIX = 'video';
+const DEFAULT_VECTOR_SET = 'video-vectors';
+const DEFAULT_SUMMARY_PREFIX = 'video-summary';
+
 export default {
   app: {
     NAME: npm_package_name ?? 'video-search',
@@ -55,23 +70,39 @@ export default {
     API_KEY: SEARCHAPI_API_KEY ?? '',
   },
   hf: {
-    VIDEO_INDEX_NAME: HF_VIDEO_INDEX_NAME ?? 'idx-videos-hf',
-    VIDEO_PREFIX: HF_VIDEO_PREFIX ?? 'video-hf:',
+    VIDEO_INDEX_NAME: HF_VIDEO_INDEX_NAME ?? `${DEFAULT_VIDEO_INDEX_NAME}-hf`,
+    VIDEO_PREFIX: HF_VIDEO_PREFIX ?? `${DEFAULT_VIDEO_PREFIX}-hf:`,
     EMBEDDING_MODEL: HF_EMBEDDING_MODEL ?? 'Xenova/all-MiniLM-L6-v2',
     SUMMARY_MODEL: HF_SUMMARY_MODEL ?? 'Xenova/paraphrase-albert-small-v2',
-    VECTOR_SET: HF_VECTOR_SET ?? 'video-vectors-hf',
+    VECTOR_SET: HF_VECTOR_SET ?? `${DEFAULT_VECTOR_SET}-hf`,
+    SUMMARY_PREFIX: HF_SUMMARY_PREFIX ?? `${DEFAULT_SUMMARY_PREFIX}-hf:`,
+  },
+  google: {
+    VIDEO_INDEX_NAME:
+      GOOGLE_VIDEO_INDEX_NAME ?? `${DEFAULT_VIDEO_INDEX_NAME}-google`,
+    VIDEO_PREFIX: GOOGLE_VIDEO_PREFIX ?? `${DEFAULT_VIDEO_PREFIX}-google:`,
+    API_KEY: GOOGLE_API_KEY,
+    EMBEDDING_MODEL: GOOGLE_EMBEDDING_MODEL ?? 'embedding-001',
+    SUMMARY_MODEL: GOOGLE_SUMMARY_MODEL ?? 'gemini-pro',
+    VECTOR_SET: GOOGLE_VECTOR_SET ?? `${DEFAULT_VECTOR_SET}-google`,
+    SUMMARY_PREFIX:
+      GOOGLE_SUMMARY_PREFIX ?? `${DEFAULT_SUMMARY_PREFIX}-google:`,
   },
   openai: {
-    VIDEO_INDEX_NAME: OPENAI_VIDEO_INDEX_NAME ?? 'idx-videos',
-    VIDEO_PREFIX: OPENAI_VIDEO_INDEX_NAME ?? 'video:',
+    VIDEO_INDEX_NAME:
+      OPENAI_VIDEO_INDEX_NAME ?? `${DEFAULT_VIDEO_INDEX_NAME}-openai`,
+    VIDEO_PREFIX: OPENAI_VIDEO_PREFIX ?? `${DEFAULT_VIDEO_PREFIX}-openai:`,
     API_KEY: OPENAI_API_KEY,
     ORGANIZATION: OPENAI_ORGANIZATION,
     EMBEDDING_MODEL: OPENAI_EMBEDDING_MODEL ?? 'gpt-4',
     SUMMARY_MODEL: OPENAI_SUMMARY_MODEL ?? 'gpt-4',
-    VECTOR_SET: OPENAI_VECTOR_SET ?? 'video-vectors',
+    VECTOR_SET: OPENAI_VECTOR_SET ?? `${DEFAULT_VECTOR_SET}-openai`,
+    SUMMARY_PREFIX:
+      OPENAI_SUMMARY_PREFIX ?? `${DEFAULT_SUMMARY_PREFIX}-openai:`,
   },
   use: {
-    OPENAI: USE === 'OPENAI',
+    GOOGLE: USE === 'GOOGLE',
     HF: USE === 'HF',
+    OPENAI: USE === 'OPENAI',
   },
 };

services/video-search/src/config.ts

@@ -0,0 +1,23 @@
+import config from '../config.js';
+import { client } from '../db.js';
+import { RedisVectorStore } from 'langchain/vectorstores/redis';
+import { ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings } from '@langchain/google-genai';
+
+export const llm = new ChatGoogleGenerativeAI({
+  apiKey: config.google.API_KEY,
+  modelName: config.google.SUMMARY_MODEL,
+  maxOutputTokens: 10000,
+});
+
+export function getEmbeddings(modelName?: string) {
+  return new GoogleGenerativeAIEmbeddings({
+    apiKey: config.google.API_KEY,
+    modelName: modelName ?? config.google.EMBEDDING_MODEL
+  });
+}
+
+export const vectorStore = new RedisVectorStore(getEmbeddings(), {
+  redisClient: client,
+  indexName: config.google.VIDEO_INDEX_NAME,
+  keyPrefix: config.google.VIDEO_PREFIX,
+});

services/video-search/src/google/config.ts

@@ -0,0 +1,3 @@
+export * as summarize from './summarize.js';
+export { search } from './search.js';
+export { store } from './store.js';

services/video-search/src/google/index.ts

@@ -0,0 +1,42 @@
+import * as summarize from './summarize.js';
+import { VideoDocument } from '../transcripts/index.js';
+import { vectorStore } from './config.js';
+import log from '../log.js';
+
+async function getVideos(question: string) {
+  log.debug(
+    `Performing similarity search for videos that answer: ${question}`,
+    {
+      question,
+      location: 'google.search.search',
+    },
+  );
+
+  const KNN = 3;
+  /* Simple standalone search in the vector DB */
+  return vectorStore.similaritySearch(question, KNN) as Promise<
+    VideoDocument[]
+  >;
+}
+
+export async function search(question: string) {
+  log.debug(`Original question: ${question}`, {
+    location: 'google.search.search',
+  });
+  const semanticQuestion = await summarize.question(question);
+
+  log.debug(`Semantic question: ${semanticQuestion}`, {
+    location: 'google.search.search',
+  });
+  let videos = await getVideos(semanticQuestion);
+
+  if (videos.length === 0) {
+    log.debug('No videos found for semantic question, trying with original question', {
+      location: 'google.search.search',
+    });
+
+    videos = await getVideos(question);
+  }
+
+  return videos;
+}

services/video-search/src/google/search.ts

@@ -0,0 +1,41 @@
+import { client } from '../db.js';
+import { vectorStore } from './config.js';
+import { VideoDocument } from '../transcripts/index.js';
+import config from '../config.js';
+import log from '../log.js';
+
+export async function store(documents: VideoDocument[]) {
+  log.debug('Storing documents...', {
+    location: 'google.store.store',
+  });
+  const newDocuments: VideoDocument[] = [];
+
+  await Promise.all(
+    documents.map(async (doc) => {
+      const exists = await client.sIsMember(
+        config.google.VECTOR_SET,
+        doc.metadata.id,
+      );
+
+      if (!exists) {
+        newDocuments.push(doc);
+      }
+    }),
+  );
+
+  log.debug(`Found ${newDocuments.length} new documents`, {
+    location: 'google.store.store',
+  });
+
+  if (newDocuments.length === 0) {
+    return;
+  }
+
+  await vectorStore.addDocuments(newDocuments);
+
+  await Promise.all(
+    newDocuments.map(async (doc) => {
+      await client.sAdd(config.google.VECTOR_SET, doc.metadata.id);
+    }),
+  );
+}

services/video-search/src/google/store.ts

@@ -0,0 +1,80 @@
+import { Document } from 'langchain/document';
+import { TokenTextSplitter } from 'langchain/text_splitter';
+import { VideoDocument } from '../transcripts/index.js';
+import {
+  QUESTION_PROMPT,
+  SUMMARY_PROMPT,
+  SUMMARY_REFINE_PROMPT,
+} from '../templates/index.js';
+import { loadSummarizationChain } from 'langchain/chains';
+import { llm } from './config.js';
+import { StringOutputParser } from 'langchain/schema/output_parser';
+import { cacheAside } from '../db.js';
+import log from '../log.js';
+import config from '../config.js';
+
+const splitter = new TokenTextSplitter({
+  chunkSize: 10000,
+  chunkOverlap: 250,
+});
+
+const videoSummarizeChain = loadSummarizationChain(llm, {
+  type: 'refine',
+  questionPrompt: SUMMARY_PROMPT,
+  refinePrompt: SUMMARY_REFINE_PROMPT,
+});
+
+const questionSummarizeChain = QUESTION_PROMPT.pipe(llm).pipe(
+  new StringOutputParser(),
+);
+
+const cache = cacheAside(config.google.SUMMARY_PREFIX);
+
+export async function docs(allDocs: VideoDocument[][]) {
+  const summarizedDocs: VideoDocument[] = [];
+
+  for (const docs of allDocs) {
+    log.debug(`Summarizing ${docs[0].metadata.link}`, {
+      ...docs[0].metadata,
+      location: 'google.summarize.docs',
+    });
+    const existingSummary = await cache.get(docs[0].metadata.id);
+
+    if (existingSummary) {
+      summarizedDocs.push(
+        new Document({
+          metadata: docs[0].metadata,
+          pageContent: existingSummary,
+        }),
+      );
+
+      continue;
+    }
+
+    const docsSummary = await splitter.splitDocuments(docs);
+    const summary = await videoSummarizeChain.run(docsSummary);
+
+    log.debug(`Summarized ${docs[0].metadata.link}:\n ${summary}`, {
+        summary,
+        location: 'google.summarize.docs',
+    });
+    await cache.set(docs[0].metadata.id, summary);
+
+    summarizedDocs.push(
+      new Document({
+        metadata: docs[0].metadata,
+        pageContent: summary,
+      }),
+    );
+  }
+
+  return summarizedDocs;
+}
+
+export async function question(question: string) {
+  const summary = await questionSummarizeChain.invoke({
+    question,
+  });
+
+  return summary;
+}