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 the database
8connection = "postgresql+psycopg://langchain:langchain@localhost:6024/langchain"
9collection_name = "my_docs"
10embeddings = OpenAIEmbeddings()
11
12# Initialize the vector store
13vector_store = PGVector(
14    embeddings=embeddings,
15    collection_name=collection_name,
16    connection=connection,
17    use_jsonb=True,
18)
19
20# Add documents to the vector store
21docs = [
22    Document(
23        page_content="there are cats in the pond",
24        metadata={"id": 1, "location": "pond", "topic": "animals"},
25    ),
26    Document(
27        page_content="ducks are also in the pond",
28        metadata={"id": 2, "location": "pond", "topic": "animals"},
29    ),
30    Document(
31        page_content="fresh apples are available at the market",
32        metadata={"id": 3, "location": "market", "topic": "food"},
33    ),
34    Document(
35        page_content="the market also sells fresh oranges",
36        metadata={"id": 4, "location": "market", "topic": "food"},
37    ),
38    Document(
39        page_content="the new movie is about a cat",
40        metadata={"id": 5, "location": "cinema", "topic": "entertainment"},
41    ),
42]
43
44vector_store.add_documents(docs, ids=[doc.metadata["id"] for doc in docs])
45
46# Perform a similarity search
47results = vector_store.similarity_search("kitty", k=10)
48
49for doc in results:
50    print(f"* {doc.page_content} [{doc.metadata}]")