Accelerate Your AI journey: Simplify Gen AI RAG With MongoDB Atlas & Google’s Vertex AI Reasoning Engine

1
def get_text_embeddings(chunks):
2
from vertexai.language_models import TextEmbeddingModel
3
model = TextEmbeddingModel.from_pretrained("text-embedding-004")
4
inputs = chunks[0]
5
embeddings = model.get_embeddings(chunks)
6
return [embedding.values for embedding in embeddings]

1
def write_to_mongoDB(embeddings, chunks, db_name, coll_name):
2
from pymongo import MongoClient
3
client = MongoClient("URI", tlsCAFile=certifi.where())
4
db = client[db_name]
5
collection = db[coll_name]
6
7
for i in range(len(chunks)):
8
    collection.insert_one({
9
        "chunk": chunks[i],
10
        "embedding": embeddings[i]
11
    })

1
model = "gemini-1.5-pro-001"

1
def star_wars_query_tool(
2
query: str) :
3
"""
4
Retrieves vectors from a MongoDB database and uses them to answer a question related to Star wars.
5
6
Args:
7
    query: The question to be answered about star wars.
8
9
Returns:
10
    A dictionary containing the response to the question.
11
"""
12
from langchain.chains import ConversationalRetrievalChain, RetrievalQA
13
from langchain_mongodb import MongoDBAtlasVectorSearch
14
from langchain_google_vertexai import VertexAIEmbeddings, ChatVertexAI
15
from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory
16
from pymongo import MongoClient
17
18
from langchain.prompts import PromptTemplate
19
20
21
prompt_template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Do not return any answers from your own knowledge.
22
23
{context}
24
Question: {question}
25
"""
26
# create prompt for LLM
27
PROMPT = PromptTemplate(
28
    template=prompt_template, input_variables=["context", "question"]
29
)
30
31
# Add your connection string in srv format below in place of URI
32
client = MongoClient("URI")
33
db = client["embeddings"]
34
35
embeddings =VertexAIEmbeddings(model_name="text-embedding-004")
36
37
# initilize the vector store
38
vs = MongoDBAtlasVectorSearch(
39
    collection=db["sample_starwars_embeddings"],
40
    embedding=embeddings,
41
    index_name="vector_index",
42
    embedding_key="embedding",
43
    text_key="chunk",
44
)
45
46
# initilize LLM
47
llm = ChatVertexAI(
48
    model_name="gemini-1.5-pro",
49
    convert_system_message_to_human=True,
50
    max_output_tokens=1000,
51
)
52
53
# initilize retriver for the vector store object created
54
retriever = vs.as_retriever(
55
    search_type="mmr", search_kwargs={"k": 10, "lambda_mult": 0.25}
56
)
57
memory = ConversationBufferWindowMemory(
58
    memory_key="chat_history", k=5, return_messages=True
59
)
60
61
# initilize the conversation chain
62
conversation_chain = ConversationalRetrievalChain.from_llm(
63
    llm=llm,
64
    retriever=retriever,
65
    memory=memory,
66
    combine_docs_chain_kwargs={"prompt": PROMPT},
67
)
68
69
# query and get the response from conversation chain
70
response = conversation_chain({"question": query})
71
72
return response

1
agent = reasoning_engines.LangchainAgent(
2
model=model,
3
tools=[star_wars_query_tool, star_trek_query_tool],
4
agent_executor_kwargs={"return_intermediate_steps": True},
5
) 
6
agent.query(input="tell me about star wars?")

1
remote_agent = reasoning_engines.ReasoningEngine.create(
2
agent,
3
requirements=[
4
    "google-cloud-aiplatform[langchain,reasoningengine]",
5
    "cloudpickle==3.0.0",
6
    "pydantic==2.7.4",
7
    "langchain-mongodb",
8
    "pymongo",
9
    "langchain-google-vertexai",
10
11
],
12
)

1
INFO:vertexai.reasoning_engines._reasoning_engines:reasoning_engine = vertexai.preview.reasoning_engines.ReasoningEngine('projects/project-id/locations/us-central1/reasoningEngines/reasoning-engine-id')
2
3
from vertexai.preview import reasoning_engines
4
REASONING_ENGINE_RESOURCE_NAME = "projects/project-id/locations/us-central1/reasoningEngines/reasoning-engine-id" 
5
remote_agent = reasoning_engines.ReasoningEngine(REASONING_ENGINE_RESOURCE_NAME) 
6
response = remote_agent.query(input="Tell me about episode 1 from wars")