llama_index_postgres_pgvector_store_document_embedding_similarity_search.py

1import psycopg2
2from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
3from llama_index.vector_stores.postgres import PostgresVectorStore
4
5# Connection parameters
6connection_string = "postgresql://postgres:password@localhost:5432/postgres"
7db_name = "postgres"
8conn = psycopg2.connect(connection_string)
9conn.autocommit = True
10
11# 1. Initialize the PostgresVectorStore
12vector_store = PostgresVectorStore.from_params(
13    host="localhost",
14    port="5432",
15    user="postgres",
16    password="password",
17    database="postgres",
18    table_name="paul_graham_essay",
19    embed_dim=1536,  # openai embedding dimension
20)
21
22# 2. Load documents and build the index
23# Note: This assumes you have a 'data' directory with text files
24documents = SimpleDirectoryReader("./data").load_data()
25
26# 3. Create the storage context and index
27storage_context = StorageContext.from_defaults(vector_store=vector_store)
28index = VectorStoreIndex.from_documents(
29    documents, storage_context=storage_context, show_progress=True
30)
31
32# 4. Query the index
33query_engine = index.as_query_engine()
34response = query_engine.query("What did the author do growing up?")
35print(response)