Building an AI Agent With Memory Using MongoDB, Fireworks AI, and LangChain

1
!pip install langchain langchain_openai langchain-fireworks langchain-mongodb arxiv pymupdf datasets `pymongo`

1
    import os
2
3
    # Be sure to have all the API keys in your local environment as shown below
4
    # Do not publish environment keys in production
5
    # os.environ["OPENAI_API_KEY"] = "sk"
6
    # os.environ["FIREWORKS_API_KEY"] = ""
7
    # os.environ["MONGO_URI"] = ""
8
9
    FIREWORKS_API_KEY = os.environ.get("FIREWORKS_API_KEY")
10
    OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
11
    MONGO_URI = os.environ.get("MONGO_URI")

1
    import pandas as pd
2
    from datasets import load_dataset
3
4
    data = load_dataset("MongoDB/subset_arxiv_papers_with_embeddings")
5
    dataset_df = pd.DataFrame(data["train"])

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

1
    from pymongo import MongoClient
2
3
    # Initialize MongoDB python client
4
    client = MongoClient(MONGO_URI)
5
6
    DB_NAME = "agent_demo"
7
    COLLECTION_NAME = "knowledge"
8
    ATLAS_VECTOR_SEARCH_INDEX_NAME = "vector_index"
9
    collection = client.get_database(DB_NAME).get_collection(COLLECTION_NAME)

1
    # Delete any existing records in the collection
2
    collection.delete_many({})
3
4
    # Data Ingestion
5
    records = dataset_df.to_dict('records')
6
    collection.insert_many(records)
7
8
    print("Data ingestion into MongoDB completed")

1
    from langchain_openai import OpenAIEmbeddings
2
    from langchain_mongodb import MongoDBAtlasVectorSearch
3
4
    embedding_model = OpenAIEmbeddings(model="text-embedding-3-small", dimensions=256)
5
6
    # Vector Store Creation
7
    vector_store = MongoDBAtlasVectorSearch.from_connection_string(
8
        connection_string=MONGO_URI,
9
        namespace=DB_NAME + "." + COLLECTION_NAME,
10
        embedding= embedding_model,
11
        index_name=ATLAS_VECTOR_SEARCH_INDEX_NAME,
12
        text_key="abstract"
13
    )
14
15
`retriever` = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})

1
    import os
2
    from langchain_fireworks import Fireworks
3
4
    llm = Fireworks(
5
        model="accounts/fireworks/models/firefunction-v1",
6
        max_tokens=256)

1
    from langchain.agents import tool
2
    from langchain.tools.retriever import create_retriever_tool
3
    from langchain_community.document_loaders import ArxivLoader

1
    @tool
2
    def get_metadata_information_from_arxiv(word: str) -> list:
3
      """
4
      Fetches and returns metadata for a maximum of ten documents from arXiv matching the given query word.
5
6
      Args:
7
        word (str): The search query to find relevant documents on arXiv.
8
9
      Returns:
10
        list: Metadata about the documents matching the query.
11
      """
12
      docs = ArxivLoader(query=word, load_max_docs=10).load()
13
      # Extract just the metadata from each document
14
      metadata_list = [doc.metadata for doc in docs]
15
      return metadata_list

1
    @tool
2
    def get_information_from_arxiv(word: str) -> list:
3
      """
4
      Fetches and returns metadata for a single research paper from arXiv matching the given query word, which is the ID of the paper, for example: 704.0001.
5
6
      Args:
7
        word (str): The search query to find the relevant paper on arXiv using the ID.
8
9
      Returns:
10
        list: Data about the paper matching the query.
11
      """
12
      doc = ArxivLoader(query=word, load_max_docs=1).load()
13
      return doc

1
    retriever_tool = create_retriever_tool(
2
        retriever=retriever,
3
        name="knowledge_base",
4
        description="This serves as the base knowledge source of the agent and contains some records of research papers from Arxiv. This tool is used as the first step for exploration and research efforts." 
5
    )

1
tools = [get_metadata_information_from_arxiv, get_information_from_arxiv, retriever_tool]

1
    from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
2
    agent_purpose = "You are a helpful research assistant"
3
    prompt = ChatPromptTemplate.from_messages(
4
        [
5
            ("system", agent_purpose),
6
            ("human", "{input}"),
7
            MessagesPlaceholder("agent_scratchpad")
8
        ]
9
    )

1
    from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory
2
    from langchain.memory import ConversationBufferMemory
3
4
    def get_session_history(session_id: str) -> MongoDBChatMessageHistory:
5
        return MongoDBChatMessageHistory(MONGO_URI, session_id, database_name=DB_NAME, collection_name="history")
6
7
    memory = ConversationBufferMemory(
8
        memory_key="chat_history", 
9
        chat_memory=get_session_history("my-session")
10
    )

1
    from langchain.agents import AgentExecutor, create_tool_calling_agent
2
    agent = create_tool_calling_agent(llm, tools, prompt)
3
4
    agent_executor = AgentExecutor(
5
        agent=agent,
6
        tools=tools,
7
        verbose=True,
8
        handle_parsing_errors=True,
9
        memory=memory,
10
    )

1
agent_chain.invoke({"input": "Get me a list of research papers on the topic Prompt Compression"})

1
    agent_executor.invoke({"input":"Get me the abstract of the first paper on the list"})