Build AI Agents with MongoDB

from pymongo import MongoClient
from openai import OpenAI
import voyageai
from dotenv import load_dotenv
import os
# Load environment variables from .env file
load_dotenv()
# Environment variables (private)
MONGODB_URI = os.getenv("MONGODB_URI")
VOYAGE_API_KEY = os.getenv("VOYAGE_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
# MongoDB cluster configuration
mongo_client = MongoClient(MONGODB_URI)
agent_db = mongo_client["ai_agent_db"]
vector_collection = agent_db["embeddings"]
memory_collection = agent_db["chat_history"]
# Model configuration
voyage_client = voyageai.Client(api_key=VOYAGE_API_KEY)
openai_client = OpenAI(api_key=OPENAI_API_KEY)
VOYAGE_MODEL = "voyage-4-large"
OPENAI_MODEL = "gpt-4o"

from config import vector_collection, voyage_client, VOYAGE_MODEL
from pymongo.operations import SearchIndexModel
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters 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 ---
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 = lambda index: index.get("queryable") is True
    max_retries = 24  # ~2 minutes at 5s intervals
    for _ in range(max_retries):
        indices = list(vector_collection.list_search_indexes(index_name))
        if len(indices) and predicate(indices[0]):
            break
        time.sleep(5)
    else:
        raise TimeoutError(f"Index '{index_name}' did not become queryable in time.")
    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) -> list:
    query_embedding = get_embedding(user_input, input_type="query")
    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()

import dotenv from 'dotenv';
import { MongoClient } from 'mongodb';
import OpenAI from 'openai';
// Load environment variables from .env file
dotenv.config();
// MongoDB cluster configuration
export const MONGODB_URI = process.env.MONGODB_URI;
export const mongoClient = new MongoClient(MONGODB_URI);
export const agentDb = mongoClient.db('ai_agent_db');
export const vectorCollection = agentDb.collection('embeddings');
export const memoryCollection = agentDb.collection('chat_history');
// Model Configuration
export const OPENAI_MODEL = 'gpt-4o';
export const VOYAGE_MODEL = 'voyage-4-large';
export const VOYAGE_API_KEY = process.env.VOYAGE_API_KEY;
// Initialize OpenAI Client
export const openAIClient = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });

import { PDFLoader } from '@langchain/community/document_loaders/fs/pdf';
import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
import { vectorCollection } from './config.js';
import { VOYAGE_API_KEY, VOYAGE_MODEL } from './config.js';
import { VoyageAIClient } from 'voyageai';
import * as fs from 'fs';
const EMBEDDING_DIMENSIONS = 1024;
// Use Voyage AI Client SDK to get embeddings
export async function getEmbedding(data, input_type) {
  if (!VOYAGE_API_KEY) {
    throw new Error('VOYAGE_API_KEY is not set in environment variables.');
  }
  try {
    const client = new VoyageAIClient({ apiKey: VOYAGE_API_KEY });
    const response = await client.embed({
      input: [data],
      model: VOYAGE_MODEL,
      input_type: input_type, // "document" or "query"
    });
    if (response.data && response.data.length > 0) {
      return response.data[0].embedding;
    }
    throw new Error('No embedding data found from Voyage AI response.');
  } catch (error) {
    console.error('Error generating Voyage AI embedding:', error);
    return null;
  }
}
// Ingest data from a PDF, generate embeddings, and store in MongoDB
export async function ingestData() {
  try {
    // download PDF
    const rawData = await fetch('https://investors.mongodb.com/node/13176/pdf');
    const pdfBuffer = await rawData.arrayBuffer();
    fs.writeFileSync('investor-report.pdf', Buffer.from(pdfBuffer));
    // load and split PDF
    const loader = new PDFLoader('investor-report.pdf');
    const data = await loader.load();
    const textSplitter = new RecursiveCharacterTextSplitter({
      chunkSize: 400,
      chunkOverlap: 20,
    });
    const docs = await textSplitter.splitDocuments(data);
    console.log(`Chunked PDF into ${docs.length} documents.`);
    // generate embeddings and insert
    const insertDocuments = await Promise.all(
      docs.map(async (doc) => ({
        document: doc,
        embedding: await getEmbedding(doc.pageContent, 'document'),
      }))
    );
    const result = await vectorCollection.insertMany(insertDocuments, {
      ordered: false,
    });
    console.log('Inserted documents:', result.insertedCount);
  } catch (err) {
    console.error('Ingestion error:', err);
  }
}
// Create a vector search index
export async function createVectorIndex() {
  try {
    // check if the index already exists
    const existingIndexes = await vectorCollection
      .listSearchIndexes()
      .toArray();
    if (existingIndexes.some((index) => index.name === 'vector_index')) {
      console.log('Vector index already exists. Skipping creation.');
      return;
    }
    // define your Vector Search index
    const index = {
      name: 'vector_index',
      type: 'vectorSearch',
      definition: {
        fields: [
          {
            type: 'vector',
            path: 'embedding',
            numDimensions: EMBEDDING_DIMENSIONS,
            similarity: 'cosine',
          },
        ],
      },
    };
    // run the helper method to ensure the index is created
    const result = await vectorCollection.createSearchIndex(index);
    console.log(`New index named ${result} is building.`);
    // wait for the index to be ready to query
    console.log(
      'Polling to check if the index is ready. This may take up to a minute.'
    );
    const maxRetries = 24; // ~2 minutes at 5s intervals
    let isQueryable = false;
    for (let attempt = 0; attempt < maxRetries; attempt++) {
      const indexes = await vectorCollection.listSearchIndexes().toArray();
      const target = indexes.find((index) => index.name === result);
      if (target && target.queryable) {
        console.log(`${result} is ready for querying.`);
        isQueryable = true;
        break;
      }
      await new Promise((resolve) => setTimeout(resolve, 5000));
    }
    if (!isQueryable) {
      throw new Error(`Index '${result}' did not become queryable in time.`);
    }
  } catch (err) {
    console.error('Error creating vector index:', err);
    throw err;
  }
}

import { getEmbedding } from './ingest-data.js';
import { vectorCollection } from './config.js';
import { evaluate } from 'mathjs';
// Vector search tool
export async function vectorSearchTool(userInput) {
  const queryEmbedding = await getEmbedding(userInput, 'query');
  const pipeline = [
    {
      $vectorSearch: {
        index: 'vector_index',
        queryVector: queryEmbedding,
        path: 'embedding',
        exact: true,
        limit: 5,
      },
    },
    {
      $project: {
        _id: 0,
        'document.pageContent': 1,
      },
    },
  ];
  const cursor = vectorCollection.aggregate(pipeline);
  const results = await cursor.toArray();
  return results;
}
// Simple calculator tool
export function calculatorTool(userInput) {
  try {
    const result = evaluate(userInput);
    return String(result);
  } catch (e) {
    return `Error: ${e.message}`;
  }
}

import { memoryCollection } from './config.js';
/**
 * Store a chat message in the memory collection.
 * @param {string} sessionId - unique identifier for the chat session
 * @param {string} role - role of the sender (user or system)
 * @param {string} content - content of the message
 */
export async function storeChatMessage(sessionId, role, content) {
  const message = {
    session_id: sessionId,
    role,
    content,
    timestamp: new Date(), // use JS date for timestamp
  };
  await memoryCollection.insertOne(message);
}
/**
 * Retrieve the chat history for a session.
 * @param {string} sessionId - unique identifier for the chat session
 * @returns {Promise<Array<{role: string, content: string}>>}
 */
export async function retrieveSessionHistory(sessionId) {
  const cursor = memoryCollection
    .find({ session_id: sessionId })
    .sort({ timestamp: 1 });
  const messages = [];
  await cursor.forEach((msg) => {
    messages.push({ role: msg.role, content: msg.content });
  });
  return messages;
}

import { vectorSearchTool, calculatorTool } from './tools.js';
import { storeChatMessage, retrieveSessionHistory } from './memory.js';
import { openAIClient, OPENAI_MODEL } from './config.js';
// OpenAI chat completion helper
export async function openAIChatCompletion(messages) {
  try {
    const completion = await openAIClient.chat.completions.create({
      model: OPENAI_MODEL,
      messages,
      max_tokens: 1024,
    });
    return completion.choices[0].message.content;
  } catch (error) {
    console.error('Error in openAIChatCompletion:', error);
    throw error;
  }
}
// Tool selector function to determine which tool to use based on user input and session history
export async function toolSelector(userInput, sessionHistory = []) {
  const systemPrompt = `
Select the appropriate tool from the options below. Consider the full context of the conversation before deciding.
Tools available:
- vector_search_tool: Retrieve specific context about recent MongoDB earnings and announcements
- calculator_tool: For mathematical operations
- none: For general questions without additional context
Process for making your decision:
1. Analyze if the current question relates to or follows up on a previous vector search query
2. For follow-up questions, incorporate context from previous exchanges to create a comprehensive search query
3. Only use calculator_tool for explicit mathematical operations
4. Default to none only when certain the other tools won't help
When continuing a conversation:
- Identify the specific topic being discussed
- Include relevant details from previous exchanges
- Formulate a query that stands alone but preserves conversation context
Return a JSON object only: {"tool": "selected_tool", "input": "your_query"}
    `.trim();
  const messages = [
    { role: 'system', content: systemPrompt },
    ...sessionHistory,
    { role: 'user', content: userInput },
  ];
  try {
    const response = await openAIChatCompletion(messages);
    let toolCall;
    try {
      toolCall = JSON.parse(response);
    } catch {
      try {
        toolCall = eval(`(${response})`);
      } catch {
        return { tool: 'none', input: userInput };
      }
    }
    return {
      tool: toolCall.tool || 'none',
      input: toolCall.input || userInput,
    };
  } catch (err) {
    console.error('Error in toolSelector:', err);
    return { tool: 'none', input: userInput };
  }
}
// Function to get LLM response based on messages and system message content
async function getLlmResponse(messages, systemMessageContent) {
  const systemMessage = { role: 'system', content: systemMessageContent };
  let fullMessages;
  if (messages.some((msg) => msg.role === 'system')) {
    fullMessages = [...messages, systemMessage];
  } else {
    fullMessages = [systemMessage, ...messages];
  }
  const response = await openAIChatCompletion(fullMessages);
  return response;
}
// Function to generate response based on user input
export async function generateResponse(sessionId, userInput) {
  await storeChatMessage(sessionId, 'user', userInput);
  const sessionHistory = await retrieveSessionHistory(sessionId);
  const llmInput = [...sessionHistory, { role: 'user', content: userInput }];
  const { tool, input: toolInput } = await toolSelector(
    userInput,
    sessionHistory
  );
  console.log('Tool selected:', tool);
  let response;
  if (tool === 'vector_search_tool') {
    const contextResults = await vectorSearchTool(toolInput);
    const context = contextResults
      .map((doc) => doc.document?.pageContent || JSON.stringify(doc))
      .join('\n---\n');
    const systemMessageContent = `
Answer the user's question based on the retrieved context and conversation history.
1. First, understand what specific information the user is requesting
2. Then, locate the most relevant details in the context provided
3. Finally, provide a clear, accurate response that directly addresses the question
If the current question builds on previous exchanges, maintain continuity in your answer.
Only state facts clearly supported by the provided context. If information is not available, say 'I DON'T KNOW'.
Context:
${context}
        `.trim();
    response = await getLlmResponse(llmInput, systemMessageContent);
  } else if (tool === 'calculator_tool') {
    response = calculatorTool(toolInput);
  } else {
    const systemMessageContent =
      "You are a helpful assistant. Respond to the user's prompt as best as you can based on the conversation history.";
    response = await getLlmResponse(llmInput, systemMessageContent);
  }
  await storeChatMessage(sessionId, 'system', response);
  return response;
}

import readline from 'readline';
import { mongoClient } from './config.js';
import { ingestData, createVectorIndex } from './ingest-data.js';
import { generateResponse } from './planning.js';
// Prompt for user input
async function prompt(question) {
  const rl = readline.createInterface({
    input: process.stdin,
    output: process.stdout,
  });
  return new Promise((resolve) =>
    rl.question(question, (answer) => {
      rl.close();
      resolve(answer);
    })
  );
}
async function main() {
  try {
    await mongoClient.connect();
    const runIngest = await prompt('Ingest sample data? (y/n): ');
    if (runIngest.trim().toLowerCase() === 'y') {
      await ingestData();
      console.log('\nAttempting to create/verify Vector Search Index...');
      await createVectorIndex();
    } else {
      await createVectorIndex(); // ensure index exists even if not ingesting data
    }
    const sessionId = await prompt('Enter a session ID: ');
    while (true) {
      const userQuery = await prompt(
        "\nEnter your query (or type 'quit' to exit): "
      );
      if (userQuery.trim().toLowerCase() === 'quit') break;
      if (!userQuery.trim()) {
        console.log('Query cannot be empty. Please try again.');
        continue;
      }
      const answer = await generateResponse(sessionId, userQuery);
      console.log('\nAnswer:');
      console.log(answer);
    }
  } finally {
    await mongoClient.close();
  }
}
main();