Building an Advanced RAG System With Self-Querying Retrieval

1
!pip install -qU langgraph openai pymongo sentence-transformers "unstructured-ingest[pdf, s3, mongodb, embed-huggingface]"

1
UNSTRUCTURED_API_KEY = "your-api-key"
2
UNSTRUCTURED_URL = "your-api-url"

1
# Your MongoDB connection string (uri), and collection/database names
2
MONGODB_URI = "your-mongodb-uri"
3
MONGODB_DB_NAME = "your-db-name"
4
MONGODB_COLLECTION = "your-collection-name"
5
# Instantiate the MongoDB client
6
mongodb_client = MongoClient(
7
    MONGODB_URI, appname="devrel.content.selfquery_mongodb_unstructured"
8
)

1
# Your AWS authentication credentials
2
AWS_KEY = "your-aws-access-key-id"
3
AWS_SECRET = "your-aws-secret-access-key"
4
# S3 URI for the Access Point to the bucket with PDF files
5
AWS_S3_NAME = ""

1
# Embedding model to use
2
EMBEDDING_MODEL_NAME = "BAAI/bge-base-en-v1.5"
3
# Completion model to use
4
COMPLETION_MODEL_NAME = "gpt-4o-2024-08-06"

1
from unstructured_ingest.v2.pipeline.pipeline import Pipeline
2
from unstructured_ingest.v2.interfaces import ProcessorConfig
3
from unstructured_ingest.v2.processes.partitioner import PartitionerConfig
4
from unstructured_ingest.v2.processes.chunker import ChunkerConfig
5
from unstructured_ingest.v2.processes.embedder import EmbedderConfig
6
from unstructured_ingest.v2.processes.connectors.fsspec.s3 import (
7
    S3ConnectionConfig,
8
    S3DownloaderConfig,
9
    S3IndexerConfig,
10
    S3AccessConfig,
11
)
12
from unstructured_ingest.v2.processes.connectors.local import LocalUploaderConfig
13
14
WORK_DIR = "/content/temp"
15
16
Pipeline.from_configs(
17
    context=ProcessorConfig(
18
        verbose=True, tqdm=True, num_processes=5, work_dir=WORK_DIR
19
    ),
20
    indexer_config=S3IndexerConfig(remote_url=AWS_S3_NAME),
21
    downloader_config=S3DownloaderConfig(),
22
    source_connection_config=S3ConnectionConfig(
23
        access_config=S3AccessConfig(key=AWS_KEY, secret=AWS_SECRET)
24
    ),
25
    partitioner_config=PartitionerConfig(
26
        partition_by_api=True,
27
        api_key=UNSTRUCTURED_API_KEY,
28
        partition_endpoint=UNSTRUCTURED_URL,
29
        strategy="hi_res",
30
        additional_partition_args={
31
            "split_pdf_page": True,
32
            "split_pdf_allow_failed": True,
33
            "split_pdf_concurrency_level": 15,
34
        },
35
    ),
36
    chunker_config=ChunkerConfig(
37
        chunking_strategy="by_title",
38
        chunk_max_characters=1500,
39
        chunk_overlap=150,
40
    ),
41
    embedder_config=EmbedderConfig(
42
        embedding_provider="langchain-huggingface",
43
        embedding_model_name=EMBEDDING_MODEL_NAME,
44
    ),
45
    uploader_config=LocalUploaderConfig(output_dir="/content/ingest-outputs"),
46
).run()

1
def get_fiscal_year(elements: dict) -> int:
2
    # Regular expression pattern to find the element containing the fiscal year
3
    pattern = r"for the (fiscal\s+)?year ended.*?(\d{4})"
4
    year = 0
5
    for i in range(len(elements)):
6
        match = re.search(pattern, elements[i]["text"], re.IGNORECASE)
7
        if match:
8
            year = match.group(0)[-4:]
9
            try:
10
                year = int(year)
11
            except:
12
                year = 0
13
    return year

1
def get_company_name(elements: dict) -> str:
2
    name = ""
3
    # In most cases the name of the company is right before/above the following line
4
    substring = "(Exact name of registrant as specified"
5
    for i in range(len(elements)):
6
        if substring.lower() in elements[i]["text"].lower():
7
            pattern = (
8
                r"([A-Z][A-Za-z\s&.,]+?)\s*\(Exact name of registrant as specified"
9
            )
10
            match = re.search(pattern, elements[i]["text"], re.IGNORECASE)
11
            if match:
12
                name = match.group(1).strip()
13
                name = name.split("\n\n")[-1]
14
15
    if name == "":
16
        for i in range(len(elements)):
17
            # In some cases, the name of the company is right after/below the following line
18
            match = re.search(
19
                r"Exact name of registrant as specified in its charter:\n\n(.*?)\n\n",
20
                elements[i]["text"],
21
            )
22
            if match:
23
                name = match.group(1)
24
            else:
25
                # In other cases, the name follows the "Commission File Number [Number]" line
26
                match = re.search(
27
                    r"Commission File Number.*\n\n(.*?)\n\n", elements[i]["text"]
28
                )
29
                if match:
30
                    name = match.group(1)
31
    return name

1
directory = f"{WORK_DIR}/embed"
2
3
for filename in os.listdir(directory):
4
    if filename.endswith(".json"):
5
        file_path = os.path.join(directory, filename)
6
        print(f"Processing file {filename}")
7
        try:
8
            with open(file_path, "r") as file:
9
                data = json.load(file)
10
11
            company_name = get_company_name(data)
12
            fiscal_year = get_fiscal_year(data)
13
14
            # Add custom metadata fields to each entry
15
            for entry in data:
16
                entry["metadata"]["custom_metadata"] = {}
17
                entry["metadata"]["custom_metadata"]["company"] = company_name
18
                entry["metadata"]["custom_metadata"]["year"] = fiscal_year
19
20
            with open(file_path, "w") as file:
21
                json.dump(data, file, indent=2)
22
23
            print(f"Successfully updated {file_path} with custom metadata fields.")
24
        except json.JSONDecodeError as e:
25
            print(f"Error parsing JSON in {file_path}: {e}")
26
        except IOError as e:
27
            print(f"Error reading from or writing to {file_path}: {e}")

1
from unstructured_ingest.v2.processes.connectors.mongodb import (
2
    MongoDBConnectionConfig,
3
    MongoDBUploadStagerConfig,
4
    MongoDBUploaderConfig,
5
    MongoDBAccessConfig,
6
)
7
8
Pipeline.from_configs(
9
    context=ProcessorConfig(
10
        verbose=True, tqdm=True, num_processes=5, work_dir=WORK_DIR
11
    ),
12
    indexer_config=S3IndexerConfig(remote_url=AWS_S3_NAME),
13
    downloader_config=S3DownloaderConfig(),
14
    source_connection_config=S3ConnectionConfig(
15
        access_config=S3AccessConfig(key=AWS_KEY, secret=AWS_SECRET)
16
    ),
17
    partitioner_config=PartitionerConfig(
18
        partition_by_api=True,
19
        api_key=UNSTRUCTURED_API_KEY,
20
        partition_endpoint=UNSTRUCTURED_URL,
21
        strategy="hi_res",
22
        additional_partition_args={
23
            "split_pdf_page": True,
24
            "split_pdf_allow_failed": True,
25
            "split_pdf_concurrency_level": 15,
26
        },
27
    ),
28
    chunker_config=ChunkerConfig(
29
        chunking_strategy="by_title",
30
        chunk_max_characters=1500,
31
        chunk_overlap=150,
32
    ),
33
    embedder_config=EmbedderConfig(
34
        embedding_provider="langchain-huggingface",
35
        embedding_model_name=EMBEDDING_MODEL_NAME,
36
    ),
37
    destination_connection_config=MongoDBConnectionConfig(
38
        access_config=MongoDBAccessConfig(uri=MONGODB_URI),
39
        collection=MONGODB_COLLECTION,
40
        database=MONGODB_DB_NAME,
41
    ),
42
    stager_config=MongoDBUploadStagerConfig(),
43
    uploader_config=MongoDBUploaderConfig(batch_size=100),
44
).run()

1
class GraphState(TypedDict):
2
    """
3
    Represents the state of the graph.
4
5
    Attributes:
6
        question: User query
7
        metadata: Extracted metadata
8
        filter: Filter definition
9
        documents: List of retrieved documents from vector search
10
        memory: Conversational history
11
    """
12
13
    question: str
14
    metadata: Dict
15
    filter: Dict
16
    context: List[str]
17
    memory: Annotated[list, add_messages]

1
class Metadata(BaseModel):
2
    """Metadata to use for pre-filtering."""
3
4
    company: List[str] = Field(description="List of company names")
5
    year: List[str] = Field(description="List containing start year and end year")

1
def extract_metadata(state: Dict) -> Dict:
2
    """
3
    Extract metadata from natural language query.
4
5
    Args:
6
        state (Dict): The current graph state
7
8
    Returns:
9
        Dict: New key added to state i.e. metadata containing the metadata extracted from the user query.
10
    """
11
    print("---EXTRACTING METADATA---")
12
    question = state["question"]
13
    system = f"""Extract the specified metadata from the user question:
14
    - company: List of company names, eg: Google, Adobe etc. Match the names to companies on this list: {companies}
15
    - year: List of [start year, end year]. Guidelines for extracting dates:
16
        - If a single date is found, only include that.
17
        - For phrases like 'in the past X years/last year', extract the start year by subtracting X from the current year. The current year is {datetime.now().year}.
18
        - If more than two dates are found, only include the smallest and the largest year."""
19
    completion = openai_client.beta.chat.completions.parse(
20
        model=COMPLETION_MODEL_NAME,
21
        messages=[
22
            {"role": "system", "content": system},
23
            {"role": "user", "content": question},
24
        ],
25
        response_format=Metadata,
26
    )
27
    result = completion.choices[0].message.parsed
28
    # If no metadata is extracted return an empty dictionary
29
    if len(result.company) == 0 and len(result.year) == 0:
30
        return {"metadata": {}}
31
    metadata = {
32
        "metadata.custom_metadata.company": result.company,
33
        "metadata.custom_metadata.year": result.year,
34
    }
35
    return {"metadata": metadata}

1
def generate_filter(state: Dict) -> Dict:
2
    """
3
    Generate MongoDB Query API filter definition.
4
5
    Args:
6
        state (Dict): The current graph state
7
8
    Returns:
9
        Dict: New key added to state i.e. filter.
10
    """
11
    print("---GENERATING FILTER DEFINITION---")
12
    metadata = state["metadata"]
13
    system = """Generate a MongoDB filter definition from the provided fields. Follow the guidelines below:
14
    - Respond in JSON with the filter assigned to a `filter` key.
15
    - The field `metadata.custom_metadata.company` is a list of companies.
16
    - The field `metadata.custom_metadata.year` is a list of one or more years.
17
    - If any of the provided fields are empty lists, DO NOT include them in the filter.
18
    - If both the metadata fields are empty lists, return an empty dictionary {{}}.
19
    - The filter should only contain the fields `metadata.custom_metadata.company` and `metadata.custom_metadata.year`
20
    - The filter can only contain the following MongoDB Query API match expressions:
21
        - $gt: Greater than
22
        - $lt: Lesser than
23
        - $gte: Greater than or equal to
24
        - $lte: Less than or equal to
25
        - $eq: Equal to
26
        - $ne: Not equal to
27
        - $in: Specified field value equals any value in the specified array
28
        - $nin: Specified field value is not present in the specified array
29
        - $nor: Logical NOR operation
30
        - $and: Logical AND operation
31
        - $or: Logical OR operation
32
    - If the `metadata.custom_metadata.year` field has multiple dates, create a date range filter using expressions such as $gt, $lt, $lte and $gte
33
    - If the `metadata.custom_metadata.company` field contains a single company, use the $eq expression
34
    - If the `metadata.custom_metadata.company` field contains multiple companies, use the $in expression
35
    - To combine date range and company filters, use the $and operator
36
    """
37
    completion = openai_client.chat.completions.create(
38
        model=COMPLETION_MODEL_NAME,
39
        temperature=0,
40
        messages=[
41
            {"role": "system", "content": system},
42
            {"role": "user", "content": f"Fields: {metadata}"},
43
        ],
44
        response_format={"type": "json_object"},
45
    )
46
    result = json.loads(completion.choices[0].message.content)
47
    return {"filter": result.get("filter", {})}

1
collection = mongodb_client[MONGODB_DB_NAME][MONGODB_COLLECTION]
2
VECTOR_SEARCH_INDEX_NAME = "vector_index"
3
4
model = {
5
    "name": VECTOR_SEARCH_INDEX_NAME,
6
    "type": "vectorSearch",
7
    "definition": {
8
        "fields": [
9
            {
10
                "type": "vector",
11
                "path": "embeddings",
12
                "numDimensions": 768,
13
                "similarity": "cosine",
14
            },
15
            {"type": "filter", "path": "metadata.custom_metadata.company"},
16
            {"type": "filter", "path": "metadata.custom_metadata.year"},
17
        ]
18
    },
19
}
20
collection.create_search_index(model=model)

1
def vector_search(state: Dict) -> Dict:
2
    """
3
    Get relevant information using MongoDB Atlas Vector Search
4
5
    Args:
6
        state (Dict): The current graph state
7
8
    Returns:
9
        Dict: New key added to state i.e. documents.
10
    """
11
    print("---PERFORMING VECTOR SEARCH---")
12
    question = state["question"]
13
    filter = state["filter"]
14
    # We always want a valid filter object
15
    if not filter:
16
        filter = {}
17
    query_embedding = embedding_model.encode(question).tolist()
18
    pipeline = [
19
        {
20
            "$vectorSearch": {
21
                "index": VECTOR_SEARCH_INDEX_NAME,
22
                "path": "embeddings",
23
                "queryVector": query_embedding,
24
                "numCandidates": 150,
25
                "limit": 5,
26
                "filter": filter,
27
            }
28
        },
29
        {
30
            "$project": {
31
                "_id": 0,
32
                "text": 1,
33
                "score": {"$meta": "vectorSearchScore"},
34
            }
35
        },
36
    ]
37
    # Execute the aggregation pipeline
38
    results = collection.aggregate(pipeline)
39
    # Drop documents with cosine similarity score < 0.8
40
    relevant_results = [doc["text"] for doc in results if doc["score"] >= 0.8]
41
    context = "\n\n".join([doc for doc in relevant_results])
42
    return {"context": context}

1
def generate_answer(state: Dict) -> Dict:
2
    """
3
    Generate the final answer to the user query
4
5
    Args:
6
        state (Dict): The current graph state
7
8
    Returns:
9
        Dict: New key added to state i.e. generation.
10
    """
11
    print("---GENERATING THE ANSWER---")
12
    question = state["question"]
13
    context = state["context"]
14
    memory = state["memory"]
15
    system = f"Answer the question based only on the following context. If the context is empty or if it doesn't provide enough information to answer the question, say I DON'T KNOW"
16
    completion = openai_client.chat.completions.create(
17
        model=COMPLETION_MODEL_NAME,
18
        temperature=0,
19
        messages=[
20
            {"role": "system", "content": system},
21
            {
22
                "role": "user",
23
                "content": f"Context:\n{context}\n\n{memory}\n\nQuestion:{question}",
24
            },
25
        ],
26
    )
27
    answer = completion.choices[0].message.content
28
    return {"memory": [HumanMessage(content=context), AIMessage(content=answer)]}

1
def check_metadata_extracted(state: Dict) -> str:
2
    """
3
    Check if any metadata is extracted.
4
5
    Args:
6
        state (Dict): The current graph state
7
8
    Returns:
9
        str: Binary decision for next node to call
10
    """
11
12
    print("---CHECK FOR METADATA---")
13
    metadata = state["metadata"]
14
    # If no metadata is extracted, skip the generate filter step
15
    if not metadata:
16
        print("---DECISION: SKIP TO VECTOR SEARCH---")
17
        return "vector_search"
18
    # If metadata is extracted, generate filter definition
19
    else:
20
        print("---DECISION: GENERATE FILTER---")
21
        return "generate_filter"

1
from langgraph.graph import END, StateGraph, START
2
from langgraph.checkpoint.memory import MemorySaver
3
from IPython.display import Image, display
4
5
workflow = StateGraph(GraphState)
6
# Adding memory to the graph
7
memory = MemorySaver()
8
9
# Add nodes
10
workflow.add_node("extract_metadata", extract_metadata)
11
workflow.add_node("generate_filter", generate_filter)
12
workflow.add_node("vector_search", vector_search)
13
workflow.add_node("generate_answer", generate_answer)
14
15
# Add edges
16
workflow.add_edge(START, "extract_metadata")
17
workflow.add_conditional_edges(
18
    "extract_metadata",
19
    check_metadata_extracted,
20
    {
21
        "vector_search": "vector_search",
22
        "generate_filter": "generate_filter",
23
    },
24
)
25
workflow.add_edge("generate_filter", "vector_search")
26
workflow.add_edge("vector_search", "generate_answer")
27
workflow.add_edge("generate_answer", END)
28
29
# Compile the graph
30
app = workflow.compile(checkpointer=memory)

1
try:
2
    display(Image(app.get_graph().draw_mermaid_png()))
3
except Exception:
4
    # This requires some extra dependencies and is optional
5
    pass

1
def execute_graph(thread_id: str, question: str) -> None:
2
    """
3
    Execute the graph and stream its output
4
5
    Args:
6
        thread_id (str): Conversation thread ID
7
        question (str): User question
8
    """
9
    # Add question to the question and memory attributes of the graph state
10
    inputs = {"question": question, "memory": [HumanMessage(content=question)]}
11
    config = {"configurable": {"thread_id": thread_id}}
12
    # Stream outputs as they come
13
    for output in app.stream(inputs, config):
14
        for key, value in output.items():
15
            print(f"Node {key}:")
16
            print(value)
17
    print("---FINAL ANSWER---")
18
    print(value["memory"][-1].content)

1
execute_graph("1", "Sales summary for Walmart for 2023.")
2
3
---EXTRACTING METADATA---
4
---CHECK FOR METADATA---
5
---DECISION: GENERATE FILTER---
6
'Node extract_metadata:'
7
{'metadata': {'metadata.custom_metadata.company': ['WALMART INC.'], 'metadata.custom_metadata.year': ['2023']}}
8
---GENERATING FILTER DEFINITION---
9
'Node generate_filter:'
10
{'filter': {'$and': [{'metadata.custom_metadata.company': {'$eq': 'WALMART INC.'}}, {'metadata.custom_metadata.year': {'$eq': 2023}}]}}
11
---PERFORMING VECTOR SEARCH---
12
'Node vector_search:'
13
{'context': 'DOCUMENTS INCORPORATED BY REFERENCE...'}
14
---GENERATING THE ANSWER---
15
'Node generate_answer:'
16
{'memory': [HumanMessage(content='DOCUMENTS INCORPORATED BY REFERENCE...'), AIMessage(content='In fiscal 2023, Walmart generated total revenues of $611.3 billion, primarily comprised of net sales amounting to $605.9 billion. Walmart International contributed $101.0 billion to the fiscal 2023 consolidated net sales, representing 17% of the total.')]}
17
---FINAL ANSWER---
18
In fiscal 2023, Walmart generated total revenues of $611.3 billion, primarily comprised of net sales amounting to $605.9 billion. Walmart International contributed $101.0 billion to the fiscal 2023 consolidated net sales, representing 17% of the total.

1
execute_graph("1", "What did I just ask you?")
2
3
---EXTRACTING METADATA---
4
---CHECK FOR METADATA---
5
---DECISION: SKIP TO VECTOR SEARCH---
6
'Node extract_metadata:'
7
{'metadata': {}}
8
---PERFORMING VECTOR SEARCH---
9
'Node vector_search:'
10
{'context': ''}
11
---GENERATING THE ANSWER---
12
'Node generate_answer:'
13
{'memory': [HumanMessage(content=''), AIMessage(content='You asked for the sales summary for Walmart for 2023.')]}
14
---FINAL ANSWER---
15
You asked for the sales summary for Walmart for 2023.