Retrieval-Augmented Generation with LangChain

Exploring Chroma with OpenAI and LangChain

In this demo, you’ll learn how to use Chroma with OpenAI and LangChain. Thanks to LangChain, the interface for working with different vector databases is remarkably consistent. In this section, you’ll focus on Chroma, but remember that you can readily substitute it with another supported database if you prefer.

Getting Started with Chroma

Chroma is an open-source vector database designed with developer productivity in mind. To install the necessary LangChain integration, return to your terminal and execute:

pip install langchain-chroma

Now, create a notebook and set up Chroma:

from langchain_chroma import Chroma

db = Chroma(
  embedding_function=embeddings_model,
)

You’ve initialized Chroma by providing an embedding model. Note that you can leave out the api_key attribute when creating an OpenAI embedding model; it’ll automatically fetch it from your environment, looking for it in an OPENAI_API_KEY variable by default.

By default, Chroma stores data in memory. However, this means your data will be lost when the app restarts. You’ll configure Chroma to store your data on disk instead.

Also, you need to organize your data effectively. Just as you’d use tables in SQL databases or collections in NoSQL databases, you specify a collection name in Chroma to group related data. Update your Chroma initialization code to include these enhancements:

db = Chroma(
  collection_name="speech_collection",
  embedding_function=OpenAIEmbeddings(),
  persist_directory="./chroma_db",
)

With these changes, your data will be saved to disk and organized within the “speech_collection.”

Populating Chroma With Data

Next, insert data into your Chroma database. LangChain abstracts away the low-level details, so you’ll work with LangChain document objects to represent your data.

In a new cell, add the following code:

from uuid import uuid4
from langchain_core.documents import Document

document_1 = Document(
  page_content="20 tons of cocoa have been deposited at Warehouse AX749",
  collection_name="speech_collection",
  embedding_function=OpenAIEmbeddings(),
  persist_directory="./chroma_db",
  metadata={"source": "messaging_api"},
  id=1,
)

document_2 = Document(
  page_content="The National Geographic Society has discovered a new species
    of aquatic animal, off the coast of Miami. They have been exploring at 
    8000 miles deep in the Pacific Ocean. They believe there's a lot 
    more to learn from the oceans.",
  metadata={"source": "news"},
  id=2,
)

document_3 = Document(
  page_content="Martin Luther King's speech, I Have a Dream, remains 
    one of the world's greatest ever. Here's everything he said 
    in 5 minutes.",
  metadata={"source": "website"},
  id=3,
)

document_4 = Document(
  page_content="For the first time in 1200 years, the Kalahari 
    desert receives 200ml of rain.",
  metadata={"source": "tweet"},
  id=4,
)

document_5 = Document(
  page_content="New multi-modal learning content about AI is ready
    from Kodeco.",
  metadata={"source": "kodeco_rss_feed"},
  id=5,
)

documents = [
  document_1,
  document_2,
  document_3,
  document_4,
  document_5,
]
uuids = [str(uuid4()) for _ in range(len(documents))]

db.add_documents(ids=uuids, documents=documents)

In this code, you’ve prepared your data and included metadata for each document, which can be helpful for filtering and identification later. Then, you added these documents to Chroma using add_documents(), along with unique IDs generated using uuid4(). Although this example focuses on text, Chroma supports other data types as well, with methods like add_images() and add_texts(). The metadata attribute gives extra information about a document. This helps with easy filtering and identification during queries.

Unleashing the Power of Semantic Search

So far, so good. Now, here comes some of the beauty of working with vector data stores: the search capability. Traditional SQL or NoSQL databases demand you adhere to specific query syntax, but with vector databases, you interact using natural language — just like talking to a person!

Remember, vector stores arrange data based on semantic meaning. This means search results come with a score indicating how closely they match your query.

Watch it in action. Execute this query in a new cell:

results = db.similarity_search(
  "What's the latest on the warehouse?",
)
for res in results:
  print(f"* {res.page_content}")

You used the similarity_search function to query your database. It returned:

* 20 tons of cocoa have been deposited at Warehouse AX749
* New multi-modal learning content about AI is ready from Kodeco.
* The National Geographic Society has discovered a new species of 
  aquatic animal, off the coast of Miami. They have been exploring 
  at 8000 miles deep in the Pacific Ocean. They believe there's 
  a lot more to learn from the oceans.
* For the first time in 1200 years, the Kalahari desert receives 200ml of rain.

You have stored five documents. When you ran a query, it returned three. However, only the first document directly relates to your query. Do you need that many documents? Additionally, you might notice that the best matching results appear first, with the relevance decreasing for subsequent documents. To address this, you should limit the results to a maximum of two in the next update and use its metadata to improve filtering and enhance the search results.

results = db.similarity_search(
  "What's the latest on the warehouse?",
  k=2,
  filter={"source": "messaging_api"},
)
for res in results:
  print(f"* {res.page_content}")

This time, it returned only one document, which turned out to be the most relevant to the query:

* 20 tons of cocoa have been deposited at Warehouse AX749

Ranking Results With Similarity Scores

Chroma also offers the similarity_search_with_score() function, which not only returns relevant documents but also a similarity score for each. This score quantifies how closely a document’s embedding aligns with your query’s. You can use these scores to filter out less-relevant results or even incorporate them into your application’s logic.

results = db.similarity_search_with_score(
  "Where can I find tutorials on AI?",
  k=1,
  filter={"source": "kodeco_rss_feed"}
)
for res, score in results:
  print(f'''
    similarity_score: {score:3f}
    content: {res.page_content}
    source: {res.metadata['source']}
    ''')

This query fetches the most relevant document from your “kodeco_rss_feed”, along with its similarity score. You’ll get results like the following:

similarity_score: 0.386230
content: New multi-modal learning content about AI is ready from Kodeco.
source: kodeco_rss_feed

There you have it. You’ve successfully stored and retrieved data from a vector database, demonstrating the power behind RAG applications. Continue to this lesson’s concluding segment.