Como criar um sistema RAG usando o Claude 3 Opus e o MongoDB

1
import os
2
3
os.environ["ANTHROPIC_API_KEY"] = ""
4
ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY")
5
6
os.environ["VOYAGE_API_KEY"] = ""
7
VOYAGE_API_KEY = os.environ.get("VOYAGE_API_KEY")
8
9
os.environ["HF_TOKEN"] = ""

1
from datasets import load_dataset
2
import pandas as pd
3
4
# Make sure you have an Hugging Face token(HF_TOKEN) in your development environment before running the code below
5
# How to get a token: https://huggingface.co/docs/hub/en/security-tokens
6
7
# https://huggingface.co/datasets/MongoDB/tech-news-embeddings
8
dataset = load_dataset("MongoDB/tech-news-embeddings", split="train", streaming=True)
9
combined_df = dataset.take(500)
10
11
# Convert the dataset to a pandas dataframe
12
combined_df = pd.DataFrame(combined_df)

1
# Remove the _id coloum from the intital dataset
2
combined_df = combined_df.drop(columns=['_id'])
3
4
# Convert each numpy array in the 'embedding' column to a normal Python list
5
combined_df['embedding'] = combined_df['embedding'].apply(lambda x: x.tolist())

1
import voyageai
2
3
vo = voyageai.Client(api_key=VOYAGE_API_KEY)
4
5
def get_embedding(text: str) -> list[float]:
6
    if not text.strip():
7
      print("Attempted to get embedding for empty text.")
8
      return []
9
10
    embedding = vo.embed(text, model="voyage-large-2", input_type="document")
11
12
    return embedding.embeddings[0]
13
14
combined_df["embedding"] = combined_df["description"].apply(get_embedding)
15
16
combined_df.head()

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

1
import pymongo
2
3
def get_mongo_client(mongo_uri):
4
  """Establish and validate connection to the MongoDB."""
5
6
  client = pymongo.MongoClient(mongo_uri, appname="devrel.showcase.anthropic_rag.python")
7
8
  # Validate the connection
9
  ping_result = client.admin.command('ping')
10
  if ping_result.get('ok') == 1.0:
11
    # Connection successful
12
    print("Connection to MongoDB successful")
13
    return client
14
  else:
15
    print("Connection to MongoDB failed")
16
  return None
17
18
mongo_uri = os.environ["MONGO_URI"]
19
20
if not mongo_uri:
21
  print("MONGO_URI not set in environment variables")
22
23
mongo_client = get_mongo_client(mongo_uri)
24
25
DB_NAME = "knowledge"
26
COLLECTION_NAME = "research_papers"
27
28
db = mongo_client.get_database(DB_NAME)
29
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
# Data Ingestion
2
combined_df_json = combined_df.to_dict(orient='records')
3
collection.insert_many(combined_df_json)

1
def vector_search(user_query, collection):
2
    """
3
    Perform a vector search in the MongoDB collection based on the user query.
4
5
    Args:
6
    user_query (str): The user's query string.
7
    collection (MongoCollection): The MongoDB collection to search.
8
9
    Returns:
10
    list: A list of matching documents.
11
    """
12
13
    # Generate embedding for the user query
14
    query_embedding = get_embedding(user_query)
15
16
    if query_embedding is None:
17
        return "Invalid query or embedding generation failed."
18
19
    # Define the vector search pipeline
20
    pipeline = [
21
        {
22
            "$vectorSearch": {
23
                "index": "vector_index",
24
                "queryVector": query_embedding,
25
                "path": "embedding",
26
                "numCandidates": 150,  # Number of candidate matches to consider
27
                "limit": 5  # Return top 5 matches
28
            }
29
        },
30
        {
31
            "$project": {
32
                "_id": 0,  # Exclude the _id field
33
                "embedding": 0,  # Exclude the embedding field
34
                "score": {
35
                    "$meta": "vectorSearchScore"  # Include the search score
36
                }
37
            }
38
        }
39
    ]
40
41
    # Execute the search
42
    results = collection.aggregate(pipeline)
43
    return list(results)

1
import anthropic
2
client = anthropic.Client(api_key=userdata.get("ANTHROPIC_API_KEY"))

1
def handle_user_query(query, collection):
2
3
  get_knowledge = vector_search(query, collection)
4
5
  search_result = ''
6
  for result in get_knowledge:
7
    search_result += (
8
        f"Title: {result.get('title', 'N/A')}, "
9
        f"Company Name: {result.get('companyName', 'N/A')}, "
10
        f"Company URL: {result.get('companyUrl', 'N/A')}, "
11
        f"Date Published: {result.get('published_at', 'N/A')}, "
12
        f"Article URL: {result.get('url', 'N/A')}, "
13
        f"Description: {result.get('description', 'N/A')}, \n"
14
    )
15
16
  response = client.messages.create(
17
    model="claude-3-opus-20240229",
18
    max_tokens=1024,
19
    system="You are Venture Capital Tech Analyst with access to some tech company articles and information. Use the information you are given to provide advice.",
20
    messages=[
21
        {"role": "user", "content": "Answer this user query: " + query + " with the following context: " + search_result}
22
    ]
23
  )
24
25
  return (response.content[0].text), search_result

1
# Conduct query with retrieval of sources
2
query = "Give me the best tech stock to invest in and tell me why"
3
response, source_information = handle_user_query(query, collection)
4
5
print(f"Response: {response}")
6
print(f"Source Information: \\n{source_information}")