Semantic Search with OpenAI and PostgreSQL in 10 Minutes

!pip install pandas pyarrow pgvector openai psycopg2

import pandas as pd
import openai
import psycopg2
import os
from ast import literal_eval

client = openai.OpenAI(api_key= OPENAI_API_KEY)

conn = psycopg2.connect(DATABASE_CONNECTION_STRING)

def setup_database():
   pgvector = """CREATE EXTENSION IF NOT EXISTS vector"""
   pgai = """CREATE EXTENSION IF NOT EXISTS ai CASCADE"""
   pgvectorscale = """CREATE EXTENSION IF NOT EXISTS vectorscale CASCADE"""

   with conn.cursor() as cursor:
      cursor.execute(pgvector)
      cursor.execute(pgvectorscale)
      cursor.execute(pgai)

setup_database()

CREATE TABLE IF NOT EXISTS movies (
    id bigint primary key generated by default as identity,
    title TEXT NOT NULL,
    overview TEXT,
    genres TEXT,
    producer TEXT,
    "cast" TEXT
);

# Use ai.load_dataset to load the dataset
with conn.cursor() as cur:
    # Load the 'movies' dataset from Hugging Face
    cur.execute("""
        SELECT ai.load_dataset(
            name => 'Cohere/movies',
            table_name => 'movies',
            if_table_exists => 'drop',
            field_types => '{"title": "TEXT", "overview": "TEXT", "genres": "TEXT", "producer": "TEXT", "cast": "TEXT"}'::jsonb
        );
    """)
    conn.commit()

print("Dataset loaded successfully into the 'movies' table.")

with conn.cursor() as cur:
  cur.execute("""
  SELECT ai.create_vectorizer(
      'movies'::regclass,
      destination => 'movies_embeddings_vectorized',
      embedding => ai.embedding_openai('text-embedding-3-small', 1536),
      chunking => ai.chunking_recursive_character_text_splitter('overview'),
      formatting => ai.formatting_python_template(
          '$title: $chunk'
      )
  );
  """)
print("Vectorizer created successfully.")

def semantic_search(conn, query):
    sql = """
    WITH query_embedding AS (
        SELECT openai_embed(
           'text-embedding-3-small'
           , %s
           , _api_key=>%s
       ) AS embedding
    )
    SELECT title, overview
    FROM movies, query_embedding
    ORDER BY movies.embedding <#> query_embedding.embedding
    LIMIT 5
    """
    with conn.cursor() as cur:
        cur.execute(sql, (query,))
        return cur.fetchall()

query = "A thrilling mystery with unexpected twists and turns."
results = semantic_search(conn, query)

[
    ('Shutter Island',
     'World War II soldier-turned-U.S. Marshal Teddy Daniels investigates the disappearance of a patient from a hospital for the criminally insane, but his efforts are compromised by his troubling visions and also by a mysterious doctor.'),
    ('The Sixth Sense',
     'A psychological thriller about an eight year old boy named Cole Sear who believes he can see into the world of the dead. A child psychologist named Malcolm Crowe comes to Cole to help him deal with his problem, learning that he really can see ghosts of dead people.'),
    ('Angels & Demons',
     'Harvard symbologist Robert Langdon investigates a mysterious symbol seared into the chest of a murdered physicist. He discovers evidence of the unimaginable, the rebirth of an ancient secret brotherhood known as the Illuminati, the most powerful underground organization ever to walk the earth.'),


('Gone Girl',
     'With his wife's disappearance having become the focus of an intense media circus, a man sees the spotlight turned on him when it's suspected that he may not be innocent.'),


    ('The Girl with the Dragon Tattoo',
     'This English-language adaptation of the Swedish novel by Stieg Larsson follows a disgraced journalist, Mikael Blomkvist, as he investigates the disappearance of a weary patriarch's niece from 40 years ago. He is aided by the pierced, tattooed, punk computer hacker named Lisbeth Salander. As they work together in the investigation, Blomkvist and Salander uncover immense corruption beyond anything they have ever imagined.')
]