Como criar um sistema RAG com LlamaIndex, OpenAI e MongoDB Vector Database

1
!pip install llama-index
2
3
!pip install llama-index-vector-stores-mongodb
4
5
!pip install llama-index-embeddings-openai
6
7
!pip install pymongo
8
9
!pip install datasets
10
11
!pip install pandas

1
%env OPENAI\_API\_KEY=openai\_key\_here 

1
from datasets import load_dataset
2
3
importpandasaspd
4
5
# https://huggingface.co/datasets/AIatMongoDB/embedded\_movies
6
7
dataset=load_dataset("AIatMongoDB/embedded\_movies")
8
9
# Convert the dataset to a pandas dataframe
10
11
dataset_df=pd.DataFrame(dataset['train'])
12
13
dataset_df.head(5)

1
# Remove data point where plot column is missing
2
3
dataset_df=dataset_df.dropna(subset=['plot'])
4
5
print("\nNumber of missing values in each column after removal:")
6
7
print(dataset_df.isnull().sum())
8
9
# Remove the plot_embedding from each data point in the dataset as we are going to create new embeddings with the new OpenAI embedding Model "text-embedding-3-small"
10
11
dataset_df=dataset_df.drop(columns=['plot_embedding'])
12
13
dataset_df.head(5)

1
from llama_index.core.settings import Settings
2
3
from llama_index.llms.openai import OpenAI
4
5
from llama_index.embeddings.openai import OpenAIEmbedding
6
7
embed_model=OpenAIEmbedding(model="text-embedding-3-small",dimensions=256)
8
9
llm=OpenAI()
10
11
Settings.llm=llm
12
13
Settings.embed_model=embed_model

1
import json
2
from llama_index.core import Document
3
from llama_index.core.schema import MetadataMode
4
5
# Convert the DataFrame to a JSON string representation
6
documents_json = dataset_df.to_json(orient='records')
7
# Load the JSON string into a Python list of dictionaries
8
documents_list = json.loads(documents_json)
9
10
llama_documents = []
11
12
for document in documents_list:
13
14
  # Value for metadata must be one of (str, int, float, None)
15
  document["writers"] = json.dumps(document["writers"])
16
  document["languages"] = json.dumps(document["languages"])
17
  document["genres"] = json.dumps(document["genres"])
18
  document["cast"] = json.dumps(document["cast"])
19
  document["directors"] = json.dumps(document["directors"])
20
  document["countries"] = json.dumps(document["countries"])
21
  document["imdb"] = json.dumps(document["imdb"])
22
  document["awards"] = json.dumps(document["awards"])
23
24
25
  # Create a Document object with the text and excluded metadata for llm and embedding models
26
  llama_document = Document(
27
      text=document["fullplot"],
28
      metadata=document,
29
      excluded_llm_metadata_keys=["fullplot", "metacritic"],
30
      excluded_embed_metadata_keys=["fullplot", "metacritic", "poster", "num_mflix_comments", "runtime", "rated"],
31
      metadata_template="{key}=>{value}",
32
      text_template="Metadata: {metadata_str}\n-----\nContent: {content}",
33
      )
34
35
  llama_documents.append(llama_document)
36
37
# Observing an example of what the LLM and Embedding model receive as input
38
print(
39
    "\nThe LLM sees this: \n",
40
    llama_documents[0].get_content(metadata_mode=MetadataMode.LLM),
41
)
42
print(
43
    "\nThe Embedding model sees this: \n",
44
    llama_documents[0].get_content(metadata_mode=MetadataMode.EMBED),
45
)

1
from llama_index.core.node_parser import SentenceSplitter
2
3
parser = SentenceSplitter()
4
nodes = parser.get_nodes_from_documents(llama_documents)
5
6
for node in nodes:
7
    node_embedding = embed_model.get_text_embedding(
8
        node.get_content(metadata_mode="all")
9
    )
10
    node.embedding = node_embedding
11
     

1
import pymongo
2
from google.colab import userdata
3
4
def get_mongo_client(mongo_uri):
5
  """Establish connection to the MongoDB."""
6
  try:
7
    client = pymongo.MongoClient(mongo_uri)
8
    print("Connection to MongoDB successful")
9
    return client
10
  except pymongo.errors.ConnectionFailure as e:
11
    print(f"Connection failed: {e}")
12
    return None
13
14
mongo_uri = userdata.get('MONGO_URI_2')
15
if not mongo_uri:
16
  print("MONGO_URI not set in environment variables")
17
18
mongo_client = get_mongo_client(mongo_uri)
19
20
DB_NAME="movies"
21
COLLECTION_NAME="movies_records"
22
23
db = mongo_client[DB_NAME]
24
collection = db[COLLECTION_NAME]

1
# To ensure we are working with a fresh collection 
2
# delete any existing records in the collection
3
4
collection.delete_many({})

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

1
from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch
2
3
vector_store = MongoDBAtlasVectorSearch(mongo_client, db_name=DB_NAME, collection_name=COLLECTION_NAME, index_name="vector_index")
4
vector_store.add(nodes)

1
from llama_index.core import VectorStoreIndex, StorageContext
2
index = VectorStoreIndex.from_vector_store(vector_store)

1
import pprint
2
from llama_index.core.response.notebook_utils import display_response
3
4
query_engine = index.as_query_engine(similarity_top_k=3)
5
query = "Recommend a romantic movie suitable for the christmas season and justify your selecton"
6
response = query_engine.query(query)
7
display_response(response)
8
pprint.pprint(response.source_nodes)