Usar a Atlas Vector Search do MongoDB com LangChain

1
{
2
  "fields": [
3
    {
4
      "type": "vector",
5
      "path": "embedding",  
6
      "numDimensions": 1536, 
7
      "similarity": "cosine"  
8
    }
9
  ]
10
}

1
pip3 install pip --upgrade
2
pip3 install langchain

1
pip3 install pymongo openai pypdf tiktoken

1
export OPENAI_API_KEY="xxxxxxxxxxx"
2
export ATLAS_CONNECTION_STRING="mongodb+srv://user:passwd@vectorsearch.abc.mongodb.net/?retryWrites=true"

1
import os
2
from pymongo import MongoClient
3
from langchain.document_loaders import PyPDFLoader
4
from langchain.text_splitter import RecursiveCharacterTextSplitter
5
from langchain.embeddings import OpenAIEmbeddings
6
from langchain_mongodb import MongoDBAtlasVectorSearch
7
8
# Define the URL of the PDF MongoDB Atlas Best Practices document
9
pdf_url = "https://query.prod.cms.rt.microsoft.com/cms/api/am/binary/RE4HkJP"
10
11
# Retrieve environment variables for sensitive information
12
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
13
if not OPENAI_API_KEY:
14
    raise ValueError("The OPENAI_API_KEY environment variable is not set.")
15
16
ATLAS_CONNECTION_STRING = os.getenv('ATLAS_CONNECTION_STRING')
17
if not ATLAS_CONNECTION_STRING:
18
    raise ValueError("The ATLAS_CONNECTION_STRING environment variable is not set.")
19
20
# Connect to MongoDB Atlas cluster using the connection string
21
cluster = MongoClient(ATLAS_CONNECTION_STRING)
22
23
# Define the MongoDB database and collection names
24
DB_NAME = "langchain"
25
COLLECTION_NAME = "vectorSearch"
26
27
# Connect to the specific collection in the database
28
MONGODB_COLLECTION = cluster[DB_NAME][COLLECTION_NAME]
29
30
# Initialize the PDF loader with the defined URL
31
loader = PyPDFLoader(pdf_url)
32
data = loader.load()
33
34
# Initialize the text splitter
35
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
36
37
# Split the document into manageable segments
38
docs = text_splitter.split_documents(data)
39
40
# Initialize MongoDB Atlas vector search with the document segments
41
vector_search = MongoDBAtlasVectorSearch.from_documents(
42
    documents=docs,
43
    embedding=OpenAIEmbeddings(),
44
    collection=MONGODB_COLLECTION,
45
    index_name="default"  # Use a predefined index name
46
)
47
# At this point, 'docs' are split and indexed in MongoDB Atlas, enabling text search capabilities.

1
def create_vector_search():
2
    """
3
    Creates a MongoDBAtlasVectorSearch object using the connection string, database, and collection names, along with the OpenAI embeddings and index configuration.
4
5
    :return: MongoDBAtlasVectorSearch object
6
    """
7
    vector_search = MongoDBAtlasVectorSearch.from_connection_string(
8
        ATLAS_CONNECTION_STRING,
9
        f"{DB_NAME}.{COLLECTION_NAME}",
10
        OpenAIEmbeddings(),
11
        index_name="default"
12
    )
13
    return vector_search

1
def perform_similarity_search(query, top_k=3):
2
    """
3
    This function performs a similarity search within a MongoDB Atlas collection. It leverages the capabilities of the MongoDB Atlas Search, which under the hood, may use the `$vectorSearch` operator, to find and return the top `k` documents that match the provided query semantically.
4
5
    :param query: The search query string.
6
    :param top_k: Number of top matches to return.
7
    :return: A list of the top `k` matching documents with their similarity scores.
8
    """
9
10
   # Get the MongoDBAtlasVectorSearch object
11
    vector_search = create_vector_search()
12
    
13
    # Execute the similarity search with the given query
14
    results = vector_search.similarity_search_with_score(
15
        query=query,
16
        k=top_k,
17
    )
18
    
19
    return results
20
21
# Example of calling the function directly
22
search_results = perform_similarity_search("MongoDB Atlas auditing")