Simplifique a pesquisa semântica do Atlas com LangChain e MongoDB

1
git clone https://github.com/wbleonard/atlas-langchain.git

1
python3 -m venv env

1
source env/bin/activate

1
pip3 install -r requirements.txt

1
python3 vectorize.py

1
python3 query.py -q "Who started AT&T?"

1
from langchain_community.document_loaders import WebBaseLoader
2
3
# Step 1: Load
4
loaders = [
5
 WebBaseLoader("https://en.wikipedia.org/wiki/AT%26T"),
6
 WebBaseLoader("https://en.wikipedia.org/wiki/Bank_of_America")
7
]
8
9
docs = []
10
11
for loader in loaders:
12
    for doc in loader.lazy_load():
13
        docs.append(doc)

1
# Step 2: Transform (Split)
2
from langchain.text_splitter import RecursiveCharacterTextSplitter
3
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0, separators=[
4
                                              "\n\n", "\n", r"(?<=\. )", " "], length_function=len)
5
docs = text_splitter.split_documents(docs)

1
# Step 3: Embed
2
from langchain_openai import OpenAIEmbeddings
3
embeddings = OpenAIEmbeddings(openai_api_key=params.OPENAI_API_KEY)

1
# Step 4: Store
2
from pymongo import MongoClient
3
from langchain_mongodb.vectorstores import MongoDBAtlasVectorSearch
4
5
client = MongoClient(params.MONGODB_CONN_STRING)
6
collection = client[params.DB_NAME][params.COLL_NAME]
7
8
# Insert the documents in MongoDB Atlas with their embedding
9
docsearch = MongoDBAtlasVectorSearch.from_documents(
10
    docs, embeddings, collection=collection, index_name=index_name
11
)

1
# Step 5: Create Vector Search Index
2
# THIS ONLY WORKS ON DEDICATED CLUSTERS (M10+)
3
docsearch.create_vector_search_index(dimensions=1536, update=True)

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

1
python3 vectorize.py

1
from pymongo import MongoClient
2
from langchain_mongodb.vectorstores import MongoDBAtlasVectorSearch
3
from langchain.embeddings.openai import OpenAIEmbeddings
4
from langchain.llms import OpenAI
5
from langchain.retrievers import ContextualCompressionRetriever
6
from langchain.retrievers.document_compressors import LLMChainExtractor
7
8
# Initialize MongoDB python client
9
client = MongoClient(params.MONGODB_CONN_STRING)
10
collection = client[params.DB_NAME][params.COLL_NAME]
11
12
# initialize vector store
13
vectorStore = MongoDBAtlasVectorSearch(
14
    collection, OpenAIEmbeddings(openai_api_key=params.OPENAI_API_KEY), index_name=params.INDEX_NAME
15
)
16
# perform a search between the embedding of the query and the embeddings of the documents
17
print("\nQuery Response:")
18
print("---------------")
19
docs = vectorStore.max_marginal_relevance_search(query, K=1)
20
#docs = vectorStore.similarity_search(query, K=1)
21
22
print(docs[0].metadata['title'])
23
print(docs[0].page_content)
24
25
# Contextual Compression
26
llm = OpenAI(openai_api_key=params.OPENAI_API_KEY, temperature=0)
27
compressor = LLMChainExtractor.from_llm(llm)
28
29
compression_retriever = ContextualCompressionRetriever(
30
    base_compressor=compressor,
31
    base_retriever=vectorStore.as_retriever()
32
)

1
python3 query.py -q "Who started AT&T?"
2
3
Your question:
4
-------------
5
Who started AT&T?
6
7
AI Response:
8
-----------
9
AT&T - Wikipedia
10
"AT&T was founded as Bell Telephone Company by Alexander Graham Bell, Thomas Watson and Gardiner Greene Hubbard after Bell's patenting of the telephone in 1875."[25] "On December 30, 1899, AT&T acquired the assets of its parent American Bell Telephone, becoming the new parent company."[28]