langchain_chroma_vector_store_with_openai_embeddings_similarity_search.py

1import langchain_chroma
2from langchain_openai import OpenAIEmbeddings
3from langchain_core.documents import Document
4
5# Initialize embeddings
6embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
7
8# Create a list of documents
9documents = [
10    Document(
11        page_content="Dogs are faithful companions, known for their loyalty and affection.",
12        metadata={"source": "mammal-pets-doc"},
13    ),
14    Document(
15        page_content="Cats are independent pets that often enjoy their own space.",
16        metadata={"source": "mammal-pets-doc"},
17    ),
18]
19
20# Initialize Chroma vector store and add documents
21vector_store = langchain_chroma.Chroma.from_documents(
22    documents=documents,
23    embedding=embeddings,
24)
25
26# Perform a similarity search
27results = vector_store.similarity_search(
28    "Are dogs loyal?",
29    k=1,
30)
31
32# Display results
33for res in results:
34    print(f"* {res.page_content} [{res.metadata}]")