supabase/examples/ai/llamaindex/llamaindex.ipynb

Supabase + LlamaIndex

In this example we'll use PostgreSQL + pgvectors with LlamaIndex's Supabase Vector Store.

Install Dependencies

In [1]:

!pip install -qU vecs datasets llama_index html2text

In [ ]:

import logging
import sys

# Uncomment to see debug logs
# logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

from llama_index.core.storage import StorageContext
from llama_index.readers.web import SimpleWebPageReader
from llama_index.indices.vector_store import VectorStoreIndex
from llama_index.vector_stores.supabase import SupabaseVectorStore
import textwrap
import html2text

Set up OpenAI

OpenAI requires an API key to run their models. Let's store that in an environment variable:

In [ ]:

import os
os.environ['OPENAI_API_KEY'] = "[your_openai_api_key]"

Load the Dataset

Let's load a small data set of Paul Graham's essays:

In [ ]:

essays = [
    'paul_graham_essay.txt'
]
documents = SimpleWebPageReader().load_data([f'https://raw.githubusercontent.com/supabase/supabase/master/examples/ai/llamaindex/data/{essay}' for essay in essays])
print('Document ID:', documents[0].doc_id, 'Document Hash:', documents[0].hash)

Create an index in Supabase

Let's store Paul Graham's essays in Supabase. You can find the Postgres connection string in the Database Settings of your Supabase project.

Note: SQLAlchemy requires the connection string to start with postgresql:// (instead of postgres://). Don't forget to rename this after copying the string from the dashboard.

Note: You must use the "connection pooling" string (domain ending in *.pooler.supabase.com) with Google Colab since Colab does not support IPv6.

This will also work with any other Postgres provider that supports pgvector.

In [5]:

# Substitute your connection string here
DB_CONNECTION = "postgresql://postgres:password@localhost:5431/db"

vector_store = SupabaseVectorStore(
    postgres_connection_string=DB_CONNECTION, 
    collection_name='base_demo'
)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)

Query the index

We can now ask questions using our index.

In [7]:

query_engine = index.as_query_engine()

# Ask a question
response = query_engine.query("What did the author do growing up?")

# Print the response
print(response)