Como implementar o Agentic RAG usando o Claude 3.5 Sonnet, LlamaIndex e MongoDB

1
!pip install --quiet llama-index  # main llamaindex library
2
3
!pip install --quiet llama-index-vector-stores-MongoDB # mongodb vector database
4
5
!pip install --quiet llama-index-llms-anthropic # anthropic LLM provider
6
7
!pip install --quiet llama-index-embeddings-openai # openai embedding provider
8
9
!pip install --quiet pymongo pandas datasets # others

1
import os
2
os.environ["ANTHROPIC_API_KEY"] = ""
3
os.environ["HF_TOKEN"] = ""
4
os.environ["OPENAI_API_KEY"] = ""
5
# WARNING: Never commit API keys or sensitive information to public repositories

1
from llama_index.embeddings.openai import OpenAIEmbedding
2
from llama_index.llms.anthropic import Anthropic
3
from llama_index.core import Settings
4
llm = Anthropic(model="claude-3-5-sonnet-20240620")
5
6
embed_model = OpenAIEmbedding(
7
    model="text-embedding-3-small", 
8
    dimensions=256,
9
    embed_batch_size=10, 
10
    openai_api_key=os.environ["OPENAI_API_KEY"]
11
)
12
13
Settings.embed_model = embed_model
14
Settings.llm = llm

1
from datasets import load_dataset
2
import pandas as pd
3
4
# https://huggingface.co/datasets/MongoDB/airbnb_embeddings
5
dataset = load_dataset("MongoDB/airbnb_embeddings", split="train", streaming=True)
6
dataset = dataset.take(4000)
7
# Convert the dataset to a pandas dataframe
8
dataset_df = pd.DataFrame(dataset)
9
dataset_df.head(5)

1
# Dataset comes with embeddings created with OpenAI, but we will recreate new ones
2
dataset_df = dataset_df.drop(columns=['text_embeddings'])

1
import json
2
from llama_index.core import Document
3
from llama_index.core.schema import MetadataMode
4
5
documents_json = dataset_df.to_json(orient='records')
6
documents_list = json.loads(documents_json)
7
8
llama_documents = []
9
10
for document in documents_list:
11
    # Convert complex objects to JSON strings
12
    for field in ["amenities", "images", "host", "address", "availability", "review_scores", "reviews", "image_embeddings"]:
13
        document[field] = json.dumps(document[field])
14
15
    # Create a Document object
16
    llama_document = Document(
17
        text=document["description"],
18
        metadata=document,
19
        excluded_llm_metadata_keys=["_id", "transit", "minimum_nights", "maximum_nights", "cancellation_policy", "last_scraped", "calendar_last_scraped", "first_review", "last_review", "security_deposit", "cleaning_fee", "guests_included", "host", "availability", "reviews", "image_embeddings"],
20
        excluded_embed_metadata_keys=["_id", "transit", "minimum_nights", "maximum_nights", "cancellation_policy", "last_scraped", "calendar_last_scraped", "first_review", "last_review", "security_deposit", "cleaning_fee", "guests_included", "host", "availability", "reviews", "image_embeddings"],
21
        metadata_template="{key}=>{value}",
22
        text_template="Metadata: {metadata_str}\n-----\nContent: {content}",
23
    )
24
    llama_documents.append(llama_document)
25
26
# Observing input examples
27
print("\nThe LLM sees this: \n", llama_documents[0].get_content(metadata_mode=MetadataMode.LLM))
28
print("\nThe Embedding model sees this: \n", llama_documents[0].get_content(metadata_mode=MetadataMode.EMBED))

1
from llama_index.core.node_parser import SentenceSplitter, SemanticSplitterNodeParser
2
from llama_index.core.schema import MetadataMode
3
from tqdm import tqdm
4
5
# semantic_splitter = SemanticSplitterNodeParser(
6
#     buffer_size=10, breakpoint_percentile_threshold=95, embed_model=embed_model
7
# )
8
9
base_splitter = SentenceSplitter(chunk_size=5000, chunk_overlap=200)
10
11
nodes = base_splitter.get_nodes_from_documents(llama_documents)
12
13
# Progress bar
14
pbar = tqdm(total=len(nodes), desc="Embedding Progress", unit="node")
15
16
for node in nodes:
17
    node_embedding = embed_model.get_text_embedding(
18
        node.get_content(metadata_mode=MetadataMode.EMBED)
19
    )
20
    node.embedding = node_embedding
21
    
22
    # Update the progress bar
23
    pbar.update(1)
24
25
# Close the progress bar
26
pbar.close()
27
28
print("Embedding process completed!")

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

1
import pymongo
2
3
os.environ["MONGO_URI"] = ""
4
5
def get_mongo_client(mongo_uri):
6
    """Establish and validate connection to the MongoDB."""
7
    
8
    client = pymongo.MongoClient(mongo_uri, appname="devrel.showcase.python")
9
10
    # Validate the connection
11
    ping_result = client.admin.command('ping')
12
    if ping_result.get('ok') == 1.0:
13
        # Connection successful
14
        print("Connection to MongoDB successful")
15
        return client
16
    else:
17
        print("Connection to MongoDB failed")
18
    return None
19
20
21
mongo_client = get_mongo_client(mongo_uri)
22
23
DB_NAME = "airbnb"
24
COLLECTION_NAME = "listings_reviews"
25
26
db = mongo_client.get_database(DB_NAME)
27
collection = db.get_collection(COLLECTION_NAME)

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

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

1
vector_store.add(nodes)

1
from llama_index.core import VectorStoreIndex
2
from llama_index.core.tools import QueryEngineTool, ToolMetadata
3
4
index = VectorStoreIndex.from_vector_store(vector_store)
5
query_engine = index.as_query_engine(similarity_top_k=5, llm=llm)
6
7
query_engine_tool = QueryEngineTool(
8
    query_engine=query_engine,
9
    metadata=ToolMetadata(
10
        name="knowledge_base",
11
        description=(
12
            "Provides information about Airbnb listings and reviews."
13
            "Use a detailed plain text question as input to the tool."
14
        ),
15
    ),
16
)

1
from llama_index.core.agent import FunctionCallingAgentWorker
2
3
agent_worker = FunctionCallingAgentWorker.from_tools(
4
    [query_engine_tool], llm=llm, verbose=True
5
)
6
agent = agent_worker.as_agent()

1
response = agent.chat("Tell me the best listing for a place in New York")
2
print(str(response))