langchain_postgres_pgvector_document_store_similarity_search.py

1from langchain_core.documents import Document
2from langchain_openai import OpenAIEmbeddings
3from langchain_postgres import PGVector
4from langchain_postgres.vectorstores import PGVector
5
6# Connection string for PostgreSQL
7connection = "postgresql+psycopg://langchain:langchain@localhost:6024/langchain"
8collection_name = "my_docs"
9embeddings = OpenAIEmbeddings()
10
11# Initialize the vector store
12vector_store = PGVector(
13    embeddings=embeddings,
14    collection_name=collection_name,
15    connection=connection,
16    use_jsonb=True,
17)
18
19# Add documents to the vector store
20docs = [
21    Document(
22        page_content="there are cats in the pond",
23        metadata={"id": 1, "location": "pond", "topic": "animals"},
24    ),
25    Document(
26        page_content="ducks are also in the pond",
27        metadata={"id": 2, "location": "pond", "topic": "animals"},
28    ),
29    Document(
30        page_content="fresh apple juice is good for health",
31        metadata={"id": 3, "location": "market", "topic": "food"},
32    ),
33    Document(
34        page_content="the market sells fresh apples",
35        metadata={"id": 4, "location": "market", "topic": "food"},
36    ),
37    Document(
38        page_content="the park is full of dogs",
39        metadata={"id": 5, "location": "park", "topic": "animals"},
40    ),
41]
42
43vector_store.add_documents(docs, ids=[doc.metadata["id"] for doc in docs])
44
45# Perform a similarity search
46results = vector_store.similarity_search(
47    "query about animals in the pond", k=2
48)
49
50for res in results:
51    print(f"* {res.page_content} [{res.metadata}]")