code: import os
  import argparse
  import lancedb
  from lancedb.context import contextualize
  from lancedb.embeddings import with_embeddings
  from datasets import load_dataset
  import openai
  import pytest
  import subprocess
  from main import embed_func, create_prompt, complete
  
  # DOWNLOAD ==============================================================
  
  subprocess.Popen(
      "wget -c https://eto-public.s3.us-west-2.amazonaws.com/datasets/youtube_transcript/youtube-transcriptions_sample.jsonl",
      shell=True,
  ).wait()
  
  # Testing ===========================================================
  
  
  @pytest.fixture
  def mock_embed_func(monkeypatch):
      def mock_api_call(*args, **kwargs):
          return {"data": [{"embedding": [0.5]} for _ in range(10)]}
  
      monkeypatch.setattr(openai.Embedding, "create", mock_api_call)
  
  
  @pytest.fixture
  def mock_complete(monkeypatch):
      def mock_api_call(*args, **kwargs):
          return {"choices": [{"text": "test"}]}
  
      monkeypatch.setattr(openai.Completion, "create", mock_api_call)
  
  
  def test_main(mock_embed_func, mock_complete):
      args = argparse.Namespace(
          query="test",
          context_length=3,
          window_size=20,
          stride=4,
          openai_key="test",
          model="test",
      )
  
      db = lancedb.connect("~/tmp/lancedb")
      table_name = "youtube-chatbot"
      if table_name not in db.table_names():
          data = load_dataset("jamescalam/youtube-transcriptions", split="train")
          df = (
              contextualize(data.to_pandas())
              .groupby("title")
              .text_col("text")
              .window(args.window_size)
              .stride(args.stride)
              .to_df()
          )
          df = df.iloc[:10].reset_index(drop=True)
          print(df.shape)
          data = with_embeddings(embed_func, df, show_progress=True)
          data.to_pandas().head(1)
          tbl = db.create_table(table_name, data)
          print(f"Created LaneDB table of length: {len(tbl)}")
      else:
          tbl = db.open_table(table_name)
  
      load_dataset("jamescalam/youtube-transcriptions", split="train")
      emb = embed_func(args.query)[0]
      context = tbl.search(emb).limit(args.context_length).to_df()
      prompt = create_prompt(args.query, context)
      complete(prompt)
      top_match = context.iloc[0]
      print(f"Top Match: {top_match['url']}&t={top_match['start']}")
  
apis: lancedb.embeddings.with_embeddings lancedb.connect
extract_api: [((1071, 1184), 'argparse.Namespace', 'argparse.Namespace', ([], {'query': '"""test"""', 'context_length': '(3)', 'window_size': '(20)', 'stride': '(4)', 'openai_key': '"""test"""', 'model': '"""test"""'}), "(query='test', context_length=3, window_size=20, stride=4,\n    openai_key='test', model='test')\n", (1089, 1184), False, 'import argparse\n'), ((1246, 1278), 'lancedb.connect', 'lancedb.connect', (['"""~/tmp/lancedb"""'], {}), "('~/tmp/lancedb')\n", (1261, 1278), False, 'import lancedb\n'), ((1995, 2059), 'datasets.load_dataset', 'load_dataset', (['"""jamescalam/youtube-transcriptions"""'], {'split': '"""train"""'}), "('jamescalam/youtube-transcriptions', split='train')\n", (2007, 2059), False, 'from datasets import load_dataset\n'), ((2174, 2208), 'main.create_prompt', 'create_prompt', (['args.query', 'context'], {}), '(args.query, context)\n', (2187, 2208), False, 'from main import embed_func, create_prompt, complete\n'), ((2213, 2229), 'main.complete', 'complete', (['prompt'], {}), '(prompt)\n', (2221, 2229), False, 'from main import embed_func, create_prompt, complete\n'), ((339, 498), 'subprocess.Popen', 'subprocess.Popen', (['"""wget -c https://eto-public.s3.us-west-2.amazonaws.com/datasets/youtube_transcript/youtube-transcriptions_sample.jsonl"""'], {'shell': '(True)'}), "(\n    'wget -c https://eto-public.s3.us-west-2.amazonaws.com/datasets/youtube_transcript/youtube-transcriptions_sample.jsonl'\n    , shell=True)\n", (355, 498), False, 'import subprocess\n'), ((1372, 1436), 'datasets.load_dataset', 'load_dataset', (['"""jamescalam/youtube-transcriptions"""'], {'split': '"""train"""'}), "('jamescalam/youtube-transcriptions', split='train')\n", (1384, 1436), False, 'from datasets import load_dataset\n'), ((1746, 1797), 'lancedb.embeddings.with_embeddings', 'with_embeddings', (['embed_func', 'df'], {'show_progress': '(True)'}), '(embed_func, df, show_progress=True)\n', (1761, 1797), False, 'from lancedb.embeddings import with_embeddings\n'), ((2070, 2092), 'main.embed_func', 'embed_func', (['args.query'], {}), '(args.query)\n', (2080, 2092), False, 'from main import embed_func, create_prompt, complete\n')]