llama_index_postgres_pgvector_document_ingestion_and_similarity_search.py

1import logging
2import sys
3import psycopg2
4from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
5from llama_index.vector_stores.postgres import PostgresVectorStore
6
7# Optional: Set up logging
8logging.basicConfig(stream=sys.stdout, level=logging.INFO)
9logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
10
11# Connection parameters
12connection_string = "postgresql://postgres:password@localhost:5432"
13db_name = "postgres"
14conn = psycopg2.connect(connection_string)
15conn.autocommit = True
16
17# Create the database if it doesn't exist
18with conn.cursor() as c:
19    c.execute(f"DROP DATABASE IF EXISTS {db_name}")
20    c.execute(f"CREATE DATABASE {db_name}")
21
22# Initialize the vector store
23vector_store = PostgresVectorStore.from_params(
24    host="localhost",
25    port="5432",
26    user="postgres",
27    password="password",
28    database="postgres",
29    table_name="paul_graham_essay",
30    embed_dim=1536,  # openai embedding dimension
31)
32
33# Load documents and build index
34# Note: Ensure you have data in the './data' directory
35documents = SimpleDirectoryReader("./data").load_data()
36storage_context = StorageContext.from_defaults(vector_store=vector_store)
37index = VectorStoreIndex.from_documents(
38    documents, storage_context=storage_context, show_progress=True
39)
40
41# Query the index
42query_engine = index.as_query_engine()
43response = query_engine.query("What did the author do growing up?")
44print(response)