Build AI Agents with MongoDB Atlas

from config import vector_collection, voyage_client, VOYAGE_MODEL
from pymongo.operations import SearchIndexModel
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
import time
# Define a function to generate embeddings
def get_embedding(data, input_type = "document"):
  embeddings = voyage_client.embed(
      data, model = VOYAGE_MODEL, input_type = input_type
  ).embeddings
  return embeddings[0]
# --- Ingest embeddings into MongoDB Atlas ---
def ingest_data():
    # Chunk PDF data
    loader = PyPDFLoader("https://investors.mongodb.com/node/13176/pdf")
    data = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=20)
    documents = text_splitter.split_documents(data)
    print(f"Successfully split PDF into {len(documents)} chunks.")
    # Ingest chunked documents into collection
    print("Generating embeddings and ingesting documents...")
    docs_to_insert = []
    for i, doc in enumerate(documents):
        embedding = get_embedding(doc.page_content)
        if embedding:
            docs_to_insert.append({
                "text": doc.page_content,
                "embedding": embedding
            })
    if docs_to_insert:
        result = vector_collection.insert_many(docs_to_insert)
        print(f"Inserted {len(result.inserted_ids)} documents into the collection.")
    else:
        print("No documents were inserted. Check embedding generation process.")
    # --- Create the vector search index ---
    index_name = "vector_index"
    
    search_index_model = SearchIndexModel(
        definition = {
            "fields": [
                {
                    "type": "vector",
                    "numDimensions": 1024, 
                    "path": "embedding",
                    "similarity": "cosine"
                }
            ]
        },
        name=index_name,
        type="vectorSearch"
    )
    try:
        vector_collection.create_search_index(model=search_index_model)
        print(f"Search index '{index_name}' creation initiated.")
    except Exception as e:
        print(f"Error creating search index: {e}")
        return
    # Wait for initial sync to complete
    print("Polling to check if the index is ready. This may take up to a minute.")
    predicate=None
    if predicate is None:
       predicate = lambda index: index.get("queryable") is True
    while True:
       indices = list(vector_collection.list_search_indexes(index_name))
       if len(indices) and predicate(indices[0]):
          break
       time.sleep(5)
    print(index_name + " is ready for querying.")

from config import vector_collection
from ingest_data import get_embedding
# Define a vector search tool
def vector_search_tool(user_input: str) -> str:
    query_embedding = get_embedding(user_input)
    pipeline = [
        {
            "$vectorSearch": {
                "index": "vector_index",
                "queryVector": query_embedding,
                "path": "embedding",
                "exact": True,
                "limit": 5
            }
        }, {
            "$project": {
                "_id": 0,
                "text": 1
            }
        }
    ]
    results = vector_collection.aggregate(pipeline)
    array_of_results = []
    for doc in results:
        array_of_results.append(doc)
    return array_of_results
# Define a simple calculator tool
def calculator_tool(user_input: str) -> str:
    try:
        result = eval(user_input)
        return str(result)
    except Exception as e:
        return f"Error: {str(e)}" 

from config import memory_collection
from datetime import datetime
from typing import List
def store_chat_message(session_id: str, role: str, content: str) -> None:
    message = {
        "session_id": session_id,     # Unique identifier for the chat session
        "role": role,                 # Role of the sender (user or system) 
        "content": content,           # Content of the message
        "timestamp": datetime.now(),  # Timestamp of when the message was sent
    }
    memory_collection.insert_one(message)
def retrieve_session_history(session_id: str) -> List:
    # Query the collection for messages with a specific "session_id" in ascending order
    cursor =  memory_collection.find({"session_id": session_id}).sort("timestamp", 1)
    # Iterate through the cursor and return a JSON object with the message role and content
    if cursor:
        messages = [{"role": msg["role"], "content": msg["content"]} for msg in cursor]
    else:
        messages = []
    return messages

from config import openai_client, OPENAI_MODEL
from tools import vector_search_tool, calculator_tool
from memory import store_chat_message, retrieve_session_history
# Define a tool selector function that decides which tool to use based on user input and message history
def tool_selector(user_input, session_history=None):
    messages = [
        {
            "role": "system",
            "content": (
                "Select the appropriate tool from the options below. Consider the full context of the conversation before deciding.\n\n"
                "Tools available:\n"
                "- vector_search_tool: Retrieve specific context about recent MongoDB earnings and announcements\n"
                "- calculator_tool: For mathematical operations\n"
                "- none: For general questions without additional context\n"
                "Process for making your decision:\n"
                "1. Analyze if the current question relates to or follows up on a previous vector search query\n"
                "2. For follow-up questions, incorporate context from previous exchanges to create a comprehensive search query\n"
                "3. Only use calculator_tool for explicit mathematical operations\n"
                "4. Default to none only when certain the other tools won't help\n\n"
                
                "When continuing a conversation:\n"
                "- Identify the specific topic being discussed\n"
                "- Include relevant details from previous exchanges\n"
                "- Formulate a query that stands alone but preserves conversation context\n\n"
                
                "Return a JSON object only: {\"tool\": \"selected_tool\", \"input\": \"your_query\"}"
           )
        }
    ]
    if session_history:
        messages.extend(session_history)
    messages.append({"role": "user", "content": user_input})
    
    response = openai_client.chat.completions.create(
      model=OPENAI_MODEL,
      messages=messages
    ).choices[0].message.content
    try:
        tool_call = eval(response)
        return tool_call.get("tool"), tool_call.get("input")
    except:
        return "none", user_input
# Define the agent workflow
def generate_response(session_id: str, user_input: str) -> str:
    
    # Store the user input in the chat history collection
    store_chat_message(session_id, "user", user_input)
    # Initialize a list of inputs to pass to the LLM
    llm_input = []
    # Retrieve the session history for the current session and add it to the LLM input
    session_history = retrieve_session_history(session_id)
    llm_input.extend(session_history)
    # Append the user message in the correct format
    user_message = {
        "role": "user",
        "content": user_input
    }
    llm_input.append(user_message)
    # Call the tool_selector function to determine which tool to use
    tool, tool_input = tool_selector(user_input, session_history)
    print("Tool selected: ", tool)
    
    # Process based on selected tool
    if tool == "vector_search_tool":
        context = vector_search_tool(tool_input)
        # Construct the system prompt using the retrieved context and append it to the LLM input
        system_message_content = (
            f"Answer the user's question based on the retrieved context and conversation history.\n"
            f"1. First, understand what specific information the user is requesting\n" 
            f"2. Then, locate the most relevant details in the context provided\n"
            f"3. Finally, provide a clear, accurate response that directly addresses the question\n\n"
            f"If the current question builds on previous exchanges, maintain continuity in your answer.\n"
            f"Only state facts clearly supported by the provided context. If information is not available, say 'I DON'T KNOW'.\n\n"
            f"Context:\n{context}"
        )
        response = get_llm_response(llm_input, system_message_content)
    elif tool == "calculator_tool":
        # Perform the calculation using the calculator tool
        response = calculator_tool(tool_input)
    else:
        system_message_content = "You are a helpful assistant. Respond to the user's prompt as best as you can based on the conversation history."
        response = get_llm_response(llm_input, system_message_content)
    
    # Store the system response in the chat history collection
    store_chat_message(session_id, "system", response)
    return response
# Helper function to get the LLM response
def get_llm_response(messages, system_message_content):
    # Add the system message to the messages list
    system_message = {
        "role": "system",
        "content": system_message_content,
    }
    
    # If the system message should go at the end (for context-based queries)
    if any(msg.get("role") == "system" for msg in messages):
        messages.append(system_message)
    else:
        # For general queries, put system message at beginning
        messages = [system_message] + messages
    
    # Get response from LLM
    response = openai_client.chat.completions.create(
        model=OPENAI_MODEL,
        messages=messages
    ).choices[0].message.content
    
    return response

from config import mongo_client
from ingest_data import ingest_data
from planning import generate_response
    
if __name__ == "__main__":
    try:
        run_ingest = input("Ingest sample data? (y/n): ")
        if run_ingest.lower() == 'y':
            ingest_data()
        session_id = input("Enter a session ID: ")
        while True:
            user_query = input("\nEnter your query (or type 'quit' to exit): ")
            if user_query.lower() == 'quit':
                break
            
            if not user_query.strip():
                print("Query cannot be empty. Please try again.")
                continue
            answer = generate_response(session_id, user_query)
            print("\nAnswer:")
            print(answer)
    finally:
        mongo_client.close()