langchain_postgres_pgvector_document_embedding_similarity_search.py

1import asyncio
2from langchain_core.documents import Document
3from langchain_openai import OpenAIEmbeddings
4from langchain_postgres import PGVector
5from langchain_postgres.vectorstores import PGVector
6
7# Connection string for PostgreSQL
8connection = "postgresql+psycopg://langchain:langchain@localhost:6024/langchain"
9collection_name = "my_docs"
10embeddings = OpenAIEmbeddings()
11
12async def main():
13    # Initialize the vector store
14    vector_store = PGVector(
15        embeddings=embeddings,
16        collection_name=collection_name,
17        connection=connection,
18        use_jsonb=True,
19    )
20
21    # Prepare documents to add
22    docs = [
23        Document(
24            page_content="there are cats in the pond",
25            metadata={"id": 1, "location": "pond", "topic": "animals"},
26        ),
27        Document(
28            page_content="ducks are also in the pond",
29            metadata={"id": 2, "location": "pond", "topic": "animals"},
30        ),
31        Document(
32            page_content="fresh apples are available at the market",
33            metadata={"id": 3, "location": "market", "topic": "food"},
34        ),
35    ]
36
37    # Add documents to the vector store
38    await vector_store.aadd_documents(docs, ids=[doc.metadata["id"] for doc in docs])
39
40    # Perform a similarity search
41    results = await vector_store.asimilarity_search(
42        "birds in the water", k=1, filter={"location": "pond"}
43    )
44    
45    for doc in results:
46        print(f"* {doc.page_content} [{doc.metadata}]")
47
48if __name__ == "__main__":
49    asyncio.run(main())