Build AI Agents with Vertex AI Agent Engine and Atlas

def star_wars_query_tool(
    query: str
):
    """
    Retrieves vectors from a MongoDB database and uses them to answer a question related to Star wars.
    Args:
        query: The question to be answered about star wars.
    Returns:
        A dictionary containing the response to the question.
    """
    from langchain.chains import ConversationalRetrievalChain, RetrievalQA
    from langchain_mongodb import MongoDBAtlasVectorSearch
    from langchain_google_vertexai import VertexAIEmbeddings, ChatVertexAI
    from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory
    from langchain.prompts import PromptTemplate
    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. Respond only in 2 or 3 sentences.
    {context}
    Question: {question}
    """
    PROMPT = PromptTemplate(
        template=prompt_template, input_variables=["context", "question"]
    )
    # Replace with your connection string to your Atlas cluster
    connection_string = "<connection-string>"
    embeddings = VertexAIEmbeddings(model_name="text-embedding-005")
    vs = MongoDBAtlasVectorSearch.from_connection_string(
        connection_string=connection_string,
        namespace="AGENT-ENGINE.sample_starwars_embeddings",
        embedding=embeddings,
        index_name="vector_index",
        embedding_key="embedding",
        text_key="chunk",
    )
    llm = ChatVertexAI(
        model_name="gemini-pro",
        convert_system_message_to_human=True,
        max_output_tokens=1000,
    )
    retriever = vs.as_retriever(
        search_type="mmr", search_kwargs={"k": 10, "lambda_mult": 0.25}
    )
    memory = ConversationBufferWindowMemory(
        memory_key="chat_history", k=5, return_messages=True
    )
    conversation_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        memory=memory,
        combine_docs_chain_kwargs={"prompt": PROMPT},
    )
    response = conversation_chain({"question": query})
    return response

def star_trek_query_tool(
    query: str
):
    """
    Retrieves vectors from a MongoDB database and uses them to answer a question related to star trek.
    Args:
        query: The question to be answered about star trek.
    Returns:
        A dictionary containing the response to the question.
    """
    from langchain.chains import ConversationalRetrievalChain, RetrievalQA
    from langchain_mongodb import MongoDBAtlasVectorSearch
    from langchain_google_vertexai import VertexAIEmbeddings, ChatVertexAI
    from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory
    from langchain.prompts import PromptTemplate
    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. Respond only in 2 or 3 sentences.
    {context}
    Question: {question}
    """
    PROMPT = PromptTemplate(
        template=prompt_template, input_variables=["context", "question"]
    )
    # Replace with your connection string to your Atlas cluster
    connection_string = "<connection-string>"
    embeddings = VertexAIEmbeddings(model_name="text-embedding-005")
    vs = MongoDBAtlasVectorSearch.from_connection_string(
        connection_string=connection_string,
        namespace="AGENT-ENGINE.sample_startrek_embeddings",
        embedding=embeddings,
        index_name="vector_index",
        embedding_key="embedding",
        text_key="chunk",
    )
    llm = ChatVertexAI(
        model_name="gemini-pro",
        convert_system_message_to_human=True,
        max_output_tokens=1000,
    )
    retriever = vs.as_retriever(
        search_type="mmr", search_kwargs={"k": 10, "lambda_mult": 0.25}
    )
    memory = ConversationBufferWindowMemory(
        memory_key="chat_history", k=5, return_messages=True
    )
    conversation_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        memory=memory,
        combine_docs_chain_kwargs={"prompt": PROMPT},
    )
    response = conversation_chain({"question": query})
    return response