Como melhorar os aplicativos LLM com a recuperação de documentos principais usando MongoDB e LangChain

1
! pip install -qU datasets pymongo langchain langgraph langchain-mongodb langchain-openai

1
os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter your OpenAI API Key:")

1
MONGODB_URI = getpass.getpass("Enter your MongoDB connection string:")
2
mongodb_client = MongoClient(
3
    MONGODB_URI, appname="devrel.showcase.parent_doc_retrieval"
4
)
5
mongodb_client.admin.command("ping")

1
os.environ["HF_TOKEN"] = getpass.getpass("Enter your HF Access Token:")

1
from datasets import load_dataset
2
import pandas as pd
3
4
data = load_dataset("mongodb-eai/docs", streaming=True, split="train")
5
data_head = data.take(1000)
6
df = pd.DataFrame(data_head)

1
from langchain_core.documents import Document
2
3
docs = []
4
metadata_fields = ["updated", "url", "title"]
5
for _, row in df.iterrows():
6
    content = row["body"]
7
    metadata = row["metadata"]
8
    for field in metadata_fields:
9
        metadata[field] = row[field]
10
    docs.append(Document(page_content=content, metadata=metadata))

1
Document(page_content='# View Database Access History\n\n- This feature is not available for `M0` free clusters, `M2`, and `M5` clusters. To learn more, see Atlas M0 (Free Cluster), M2, and M5 Limits', metadata={'contentType': None, 'pageDescription': None, 'productName': 'MongoDB Atlas', 'tags': ['atlas', 'docs'], 'version': None, 'updated': {'$date': '2024-05-20T17:30:49.148Z'}, 'url': 'https://mongodb.com/pt-br/docs/atlas/access-tracking/', 'title': 'View Database Access History'})

1
from langchain_mongodb.retrievers import (
2
    MongoDBAtlasParentDocumentRetriever,
3
)
4
from langchain_text_splitters import RecursiveCharacterTextSplitter
5
from langchain_openai import OpenAIEmbeddings
6
7
embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
8
9
DB_NAME = "langchain"
10
COLLECTION_NAME = "parent_doc"
11
12
def get_splitter(chunk_size: int) -> RecursiveCharacterTextSplitter:
13
    """
14
    Returns a token-based text splitter with overlap
15
16
    Args:
17
        chunk_size (_type_): Chunk size in number of tokens
18
19
    Returns:
20
        RecursiveCharacterTextSplitter: Recursive text splitter object
21
    """
22
    return RecursiveCharacterTextSplitter.from_tiktoken_encoder(
23
        encoding_name="cl100k_base",
24
        chunk_size=chunk_size,
25
        chunk_overlap=0.15 * chunk_size,
26
    )

1
parent_doc_retriever = MongoDBAtlasParentDocumentRetriever.from_connection_string(
2
    connection_string=MONGODB_URI,
3
    embedding_model=embedding_model,
4
    child_splitter=get_splitter(200),
5
    database_name=DB_NAME,
6
    collection_name=COLLECTION_NAME,
7
    text_key="page_content",
8
    search_kwargs={"k": 10},
9
)

1
parent_chunk_retriever = MongoDBAtlasParentDocumentRetriever.from_connection_string(
2
    connection_string=MONGODB_URI,
3
    embedding_model=embedding_model,
4
    child_splitter=get_splitter(200),
5
    parent_splitter=get_splitter(800),
6
    database_name=DB_NAME,
7
    collection_name=COLLECTION_NAME,
8
    text_key="page_content",
9
    search_kwargs={"k": 10},
10
)

1
import asyncio
2
from typing import Generator, List
3
4
BATCH_SIZE = 256
5
MAX_CONCURRENCY = 4
6
7
async def process_batch(batch: Generator, semaphore: asyncio.Semaphore) -> None:
8
    """
9
    Ingest batches of documents into MongoDB
10
11
    Args:
12
        batch (Generator): Chunk of documents to ingest
13
        semaphore (as): Asyncio semaphore
14
    """
15
    async with semaphore:
16
        await parent_doc_retriever.aadd_documents(batch)
17
        print(f"Processed {len(batch)} documents")

1
def get_batches(docs: List[Document], batch_size: int) -> Generator:
2
    """
3
    Return batches of documents to ingest into MongoDB
4
5
    Args:
6
        docs (List[Document]): List of LangChain documents
7
        batch_size (int): Batch size
8
9
    Yields:
10
        Generator: Batch of documents
11
    """
12
    for i in range(0, len(docs), batch_size):
13
        yield docs[i : i + batch_size]

1
async def process_docs(docs: List[Document]) -> List[None]:
2
    """
3
    Asynchronously ingest LangChain documents into MongoDB
4
5
    Args:
6
        docs (List[Document]): List of LangChain documents
7
8
    Returns:
9
        List[None]: Results of the task executions
10
    """
11
    semaphore = asyncio.Semaphore(MAX_CONCURRENCY)
12
    batches = get_batches(docs, BATCH_SIZE)
13
14
    tasks = []
15
    for batch in batches:
16
        tasks.append(process_batch(batch, semaphore))
17
    # Gather results from all tasks
18
    results = await asyncio.gather(*tasks)
19
    return results

1
collection = mongodb_client[DB_NAME][COLLECTION_NAME]
2
# Delete any existing documents from the collection
3
collection.delete_many({})
4
print(f"Deletion complete.")
5
# Ingest LangChain documents into MongoDB
6
results = await process_docs(docs)

1
from pymongo.operations import SearchIndexModel
2
from pymongo.errors import OperationFailure
3
4
VS_INDEX_NAME = "vector_index"
5
6
# Vector search index definition
7
model = SearchIndexModel(
8
    definition={
9
        "fields": [
10
            {
11
                "type": "vector",
12
                "path": "embedding",
13
                "numDimensions": 1536,
14
                "similarity": "cosine",
15
            }
16
        ]
17
    },
18
    name=VS_INDEX_NAME,
19
    type="vectorSearch",
20
)
21
22
# Check if the index already exists, if not create it
23
try:
24
    collection.create_search_index(model=model)
25
    print(
26
        f"Successfully created index {VS_INDEX_NAME} for collection {COLLECTION_NAME}"
27
    )
28
except OperationFailure:
29
    print(
30
        f"Duplicate index {VS_INDEX_NAME} found for collection {COLLECTION_NAME}. Skipping index creation."
31
    )

1
from langchain_openai import ChatOpenAI
2
from langchain_core.prompts import ChatPromptTemplate
3
from langchain_core.runnables import RunnablePassthrough
4
from langchain_core.output_parsers import StrOutputParser
5
6
# Retrieve and parse documents
7
retrieve = {
8
    "context": parent_doc_retriever
9
    | (lambda docs: "\n\n".join([d.page_content for d in docs])),
10
    "question": RunnablePassthrough(),
11
}
12
template = """Answer the question based only on the following context. If no context is provided, respond with I DON'T KNOW: \
13
{context}
14
15
Question: {question}
16
"""
17
# Define the chat prompt
18
prompt = ChatPromptTemplate.from_template(template)
19
# Define the model to be used for chat completion
20
llm = ChatOpenAI(temperature=0, model="gpt-4o-2024-11-20")
21
# Parse output as a string
22
parse_output = StrOutputParser()
23
# Naive RAG chain
24
rag_chain = retrieve | prompt | llm | parse_output

1
print(rag_chain.invoke("How do I improve slow queries in MongoDB?"))
2
3
To improve slow queries in MongoDB, you can follow these steps:
4
5
1. **Use the Performance Advisor**:
6
   - The Performance Advisor monitors slow queries and suggests new indexes to improve query performance.
7
   - Review the suggested indexes, especially those with high Impact scores and low Average Query Targeting scores, and create them if they align with your indexing strategies.
8
9
2. **Analyze Query Performance**:
10
   - Use the **Query Profiler** to explore slow-running operations and their key performance statistics for up to the last 24 hours.
11
   - Use the **Real-Time Performance Panel (RTPP)** to evaluate query execution times and the ratio of documents scanned to documents returned.
12
13
3. **Monitor Query Latency**:
14
   - Use **Namespace Insights** to monitor collection-level query latency and view query latency metrics and statistics.
15
16
4. **Fix Inefficient Queries**:
17
   - Address `Query Targeting` alerts by adding indexes to support inefficient queries.
18
   - Use the `cursor.explain()` command to analyze query plans and identify inefficiencies.
19
20
5. **Follow Best Practices**:
21
   - Create queries that are supported by existing indexes.
22
   - Avoid large array fields in documents that are costly to search and index.
23
   - Optimize and remove unused or inefficient indexes to balance read and write performance.
24
   - Perform rolling index builds to minimize performance impact on replica sets and sharded clusters.
25
26
6. **Configure Slow Query Threshold**:
27
   - Adjust the slow query threshold to identify slow queries more effectively. By default, Atlas dynamically adjusts this threshold, but you can set a fixed threshold of 100 milliseconds if needed.
28
...
29
   - Ensure queries are supported by indexes.
30
   - Optimize queries involving `$lookup` or large array fields.
31
32
By implementing these steps, you can identify and resolve slow queries, improving overall query performance in MongoDB.

1
from langchain.agents import tool
2
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
3
from typing import Annotated, Dict
4
from langgraph.graph.message import add_messages
5
from typing_extensions import TypedDict
6
from langgraph.prebuilt import ToolNode, tools_condition
7
from langgraph.graph import StateGraph, START, END
8
9
# Converting the retriever into an agent tool
10
@tool
11
def get_info_about_mongodb(user_query: str) -> str:
12
    """
13
    Retrieve information about MongoDB.
14
15
    Args:
16
    user_query (str): The user's query string.
17
18
    Returns:
19
    str: The retrieved information formatted as a string.
20
    """
21
    docs = parent_doc_retriever.invoke(user_query)
22
    context = "\n\n".join([d.page_content for d in docs])
23
    return context
24
25
tools = [get_info_about_mongodb]

1
# Define the LLM to use as the brain of the agent
2
llm = ChatOpenAI(temperature=0, model="gpt-4o-2024-11-20")
3
# Agent prompt
4
prompt = ChatPromptTemplate.from_messages(
5
    [
6
        (
7
            "You are a helpful AI assistant."
8
            " You are provided with tools to answer questions about MongoDB."
9
            " Think step-by-step and use these tools to get the information required to answer the user query."
10
            " Do not re-run tools unless absolutely necessary."
11
            " If you are not able to get enough information using the tools, reply with I DON'T KNOW."
12
            " You have access to the following tools: {tool_names}."
13
        ),
14
        MessagesPlaceholder(variable_name="messages"),
15
    ]
16
)
17
# Partial the prompt with tool names
18
prompt = prompt.partial(tool_names=", ".join([tool.name for tool in tools]))
19
# Bind tools to LLM
20
llm_with_tools = prompt | llm.bind_tools(tools)

1
# Define graph state
2
class GraphState(TypedDict):
3
    messages: Annotated[list, add_messages]
4
5
def agent(state: GraphState) -> Dict[str, List]:
6
    """
7
    Agent node
8
9
    Args:
10
        state (GraphState): Graph state
11
12
    Returns:
13
        Dict[str, List]: Updates to the graph state
14
    """
15
    messages = state["messages"]
16
    response = llm_with_tools.invoke(messages)
17
    # We return a list, because this will get added to the existing list
18
    return {"messages": [response]}
19
20
# Convert tools into a graph node
21
tool_node = ToolNode(tools)
22
23
# Parameterize the graph with the state
24
graph = StateGraph(GraphState)
25
# Add graph nodes
26
graph.add_node("agent", agent)
27
graph.add_node("tools", tool_node)
28
# Add graph edges
29
graph.add_edge(START, "agent")
30
graph.add_edge("tools", "agent")
31
graph.add_conditional_edges(
32
    "agent",
33
    tools_condition,
34
    {"tools": "tools", END: END},
35
)
36
# Compile the graph
37
app = graph.compile()
38
39
# Execute the agent and view outputs
40
inputs = {
41
    "messages": [
42
        ("user", "How do I improve slow queries in MongoDB?"),
43
    ]
44
}
45
46
for output in app.stream(inputs):
47
    for key, value in output.items():
48
        print(f"Node {key}:")
49
        print(value)
50
print("---FINAL ANSWER---")
51
print(value["messages"][-1].content)