Searching all of Wikipedia

%%sql
drop database if exists video_game_wikipedia;
create database video_game_wikipedia;
use video_game_wikipedia;

-- create our table (note the usage of the vector data type)
create table vecs(
    id bigint(20),
    url text default null,
    paragraph text default null,
    v vector(1536) not null,
    shard key(id),
    key(id) using hash,
    fulltext (paragraph)
);

%%sql
-- generates random floats in a specified range
create or replace function randbetween(a float, b float) returns float
as
begin
  return (rand()*(b - a) + a);
end ;

%%sql
-- set the sql_mode so that the following function returns the expected result
set sql_mode = pipes_as_concat;

-- creates a vector of given length filled with random values
create or replace function gen_vector(length int) returns text as
declare s text = "[";
begin
  if length < 2 then
    raise user_exception("length too short: " || length);
  end if;

  for i in 1..length-1 loop
    s = s || randbetween(-1,1) || "," ;
  end loop;
  s = s || randbetween(-1,1) || "]";
  return s;
end;

%%sql
-- adjusts a vector to unit length
create or replace function normalize(v blob) returns blob as
declare
  squares blob = vector_mul(v,v);
  length float = sqrt(vector_elements_sum(squares));
begin
  return scalar_vector_mul(1/length, v);
end;

%%sql
-- normalizes a vector of dimension 1536
create or replace function norm1536(v vector(1536)) returns vector(1536) as
begin
  return normalize(v) :> vector(1536);
end;

%%sql
-- generates a normalized, random vector of dimension 1536
create or replace function nrandv1536() returns vector(1536) as
begin
  return norm1536(gen_vector(1536));
end;

%%sql
-- procedure to populate `vecs` with `num_rows` vectors
create or replace procedure insert_vectors(num_rows bigint) as
declare c int;
begin
  select count(*) into c from vecs;
  loop
    insert into vecs (id, v)
    select id + (select max(id) from vecs), nrandv1536()
    from vecs
    where id <= 128 * 1024; /* chunk size 128K so we can see progress */
    select count(*) into c from vecs;
    if c >= num_rows then
      exit;
    end if;
  end loop;
end;

%%sql
-- run the procedure to populate `vecs` with 10,000,000 vectors
-- this will take around 20 min
insert into vecs (id, v) values (1, nrandv1536());
call insert_vectors(10000000);

%%sql
-- since the bucket is open, you can leave the credentials clause as it is
create or replace pipeline `wiki_pipeline` as
load data S3 's3://wikipedia-video-game-data/video-game-embeddings(1).csv'
config '{"region":"us-west-1"}'
credentials '{"aws_access_key_id": "",
            "aws_secret_access_key": ""}'
skip duplicate key errors
into table `vecs`
format csv
fields terminated by ','
enclosed by '"'
lines terminated by '\r\n';

%%sql
-- start the pipeline!
start pipeline `wiki_pipeline`;

%%sql
-- monitor the pipeline!
select DATABASE_NAME, PIPELINE_NAME, BATCH_ID, BATCH_STATE, START_TIME, ROWS_STREAMED, ROWS_PER_SEC
from information_schema.PIPELINES_BATCHES_SUMMARY
order by BATCH_ID;

%%sql
alter table vecs add vector index auto (v) INDEX_OPTIONS '{"index_type":"AUTO"}';

%%sql
alter table vecs add vector index ivf_flat (v) INDEX_OPTIONS '{"index_type":"IVF_FLAT"}';

%%sql
alter table vecs add vector index ivf_pq (v) INDEX_OPTIONS '{"index_type":"IVF_PQ"}';

%%sql
set @qv = (select v from vecs where id = 1125899906845489);

-- NO INDEX: exact kNN search
select paragraph, v <*> @qv as sim
from vecs
order by sim use index () desc
limit 5;

%%sql
set @qv = (select v from vecs where id = 1125899906845489);

-- AUTO index
select paragraph, v <*> @qv as sim
from vecs
order by sim use index (auto) desc
limit 5;

%%sql
set @qv = (select v from vecs where id = 1125899906845489);

-- IVF_FLAT
select paragraph, v <*> @qv as sim
from vecs
order by sim use index (ivf_flat) desc
limit 5;

%%sql
set @qv = (select v from vecs where id = 1125899906845489);

-- IVF_PQ
select paragraph, v <*> @qv as sim
from vecs
order by sim use index (ivf_pq) desc
limit 5;

%%sql

-- set the query vector
set @v_mario = (select v
                from vecs where url = "https://en.wikipedia.org/wiki/Super_Mario_Kart"
                order by id
                limit 1);

-- building the hybrid search
with fts as(
    select id, paragraph, match (paragraph) against ('Mario Kart') as score
    from vecs
    where match (paragraph) against ('Mario Kart')
    order by score desc
    limit 200
),
vs as (
    select id, paragraph, v <*> @v_mario as score
    from vecs
    order by score use index (auto) desc
    limit 200
)
select vs.id,
    vs.paragraph,
    .3 * ifnull(fts.score, 0) + .7 * vs.score as hybrid_score,
    vs.score as vec_score,
    ifnull(fts.score, 0) as ft_score
from fts full outer join vs
    on fts.id = vs.id
order by hybrid_score desc
limit 5;

!pip3 install openai --quiet

import sqlalchemy as sa
from openai import OpenAI
import getpass
import os
import time
import json

# OpenAI connection
OPENAI_API_KEY = os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key: ')
client = OpenAI()
EMBEDDING_MODEL = 'text-embedding-ada-002'
GPT_MODEL = 'gpt-3.5-turbo'

# SingleStore connection
engine = sa.create_engine(connection_url)
connection = engine.connect()

def get_embedding(text, model=EMBEDDING_MODEL):
    '''Generates the OpenAI embedding from an input `text`.'''
    if isinstance(text, str):
        response = client.embeddings.create(input=[text], model=model)
        return json.dumps(response.data[0].embedding)

def search_wiki_page(query, limit=5):
    '''Returns a df of the top k matches to the query ordered by similarity.'''
    query_embedding_vec = get_embedding(query)
    statement = sa.text(
        f'''select paragraph, v <*> :query_embedding :> vector(1536) AS similarity
        from vecs
        order by similarity use index (auto) desc
        limit :limit;'''
    )
    print("Searching for matches...")
    start_time = time.time()
    results = connection.execute(statement, {"query_embedding": query_embedding_vec, "limit": limit})
    end_time = time.time()
    execution_time = end_time - start_time
    print(f"Search complete in {execution_time} seconds.")
    results_as_dict = results.fetchall()
    return results_as_dict

def ask_wiki_page(query, limit=5, temp=0.0):
    '''Uses RAG to answer a question from the wiki page'''
    results = search_wiki_page(query, limit)
    print("Asking Chatbot...")
    prompt = f'''Excerpt from the conversation history:
        {results}
        Question: {query}

        Based on the conversation history, try to provide the most accurate answer to the question.
        Consider the details mentioned in the conversation history to formulate a response that is as
        helpful and precise as possible. please provide links to WIKIPEDIA ARTICLES TO LOOK AT FOR MORE INFORMATION.

        Most importantly, IF THE INFORMATION IS NOT PRESENT IN THE CONVERSATION HISTORY, DO NOT MAKE UP AN ANSWER.'''
    response = client.chat.completions.create(
        model=GPT_MODEL,
        messages=[
            {"role": "system", "content": "You are a helpful assistant who is answering questions about an article."},
            {"role": "user", "content": prompt}
        ],
        temperature=temp
    )
    response_message = response.choices[0].message.content
    return response_message

query = input('Ask me a question about video games!')
ask_wiki_page(query)