Production RAG agents with Contextual AI in 15 minutes

pip install contextual-client matplotlib tqdm requests pandas python-dotenv

import os
import json
import requests
from pathlib import Path
from typing import List, Optional, Dict

import pandas as pd
from contextual import ContextualAI

from dotenv import load_dotenv

try:
    # Google Colab users can use Colab Secrets instead
    from google.colab import userdata
    API_KEY = userdata.get("CONTEXTUAL_API_KEY")
except ImportError:
    load_dotenv()
    API_KEY = os.getenv("CONTEXTUAL_API_KEY")

if not API_KEY:
    raise ValueError(
        "Set CONTEXTUAL_API_KEY in your .env file or as an environment variable."
    )

client = ContextualAI(api_key=API_KEY)

datastore_name = "Financial_Demo_RAG"

# Reuse an existing datastore if one with this name already exists
datastores = client.datastores.list()
existing_datastore = next(
    (ds for ds in datastores if ds.name == datastore_name), None
)

if existing_datastore:
    datastore_id = existing_datastore.id
    print(f"Using existing datastore: {datastore_id}")
else:
    result = client.datastores.create(name=datastore_name)
    datastore_id = result.id
    print(f"Created new datastore: {datastore_id}")

if not os.path.exists("data"):
    os.makedirs("data")

files_to_upload = [
    (
        "A_Rev_by_Mkt_Qtrly_Trend_Q425.pdf",
        "https://raw.githubusercontent.com/ContextualAI/examples/refs/heads/main/08-ai-workshop/data/A_Rev_by_Mkt_Qtrly_Trend_Q425.pdf",
    ),
    (
        "B_Q423-Qtrly-Revenue-by-Market-slide.pdf",
        "https://raw.githubusercontent.com/ContextualAI/examples/refs/heads/main/08-ai-workshop/data/B_Q423-Qtrly-Revenue-by-Market-slide.pdf",
    ),
    (
        "C_Neptune.pdf",
        "https://raw.githubusercontent.com/ContextualAI/examples/refs/heads/main/08-ai-workshop/data/C_Neptune.pdf",
    ),
    (
        "D_Unilever.pdf",
        "https://raw.githubusercontent.com/ContextualAI/examples/refs/heads/main/08-ai-workshop/data/D_Unilever.pdf",
    ),
]

document_ids = []

for filename, url in files_to_upload:
    file_path = f"data/{filename}"

    if not os.path.exists(file_path):
        print(f"Fetching {file_path}")
        try:
            response = requests.get(url)
            response.raise_for_status()
            with open(file_path, "wb") as f:
                f.write(response.content)
        except Exception as e:
            print(f"Error downloading {filename}: {e}")
            continue

    try:
        with open(file_path, "rb") as f:
            ingestion_result = client.datastores.documents.ingest(datastore_id, file=f)
            document_ids.append(ingestion_result.id)
            print(f"Uploaded {filename}")
    except Exception as e:
        print(f"Error uploading {filename}: {e}")

print(f"\nUploaded {len(document_ids)} documents")
print(f"Document IDs: {document_ids}")

metadata = client.datastores.documents.metadata(
    datastore_id=datastore_id,
    document_id=document_ids[0],
)
print("Document metadata:", metadata)

system_prompt = """
You are a helpful AI assistant created by Contextual AI to answer questions
about relevant documentation provided to you. Your responses should be
precise, accurate, and sourced exclusively from the provided information.

Guidelines:
* Only use information from the provided documentation.
* Use the exact terminology found in the documentation.
* Keep answers concise and relevant to the question.
* Apply markdown for lists, tables, or code.
* Directly answer the question, then STOP.
* If the information is not in the documentation, say so and stop.
"""

agent_name = "Demo"

agents = client.agents.list()
existing_agent = next((a for a in agents if a.name == agent_name), None)

if existing_agent:
    agent_id = existing_agent.id
    print(f"Using existing agent: {agent_id}")
else:
    app_response = client.agents.create(
        name=agent_name,
        description="Helpful Grounded AI Assistant",
        datastore_ids=[datastore_id],
        agent_configs={
            "global_config": {
                "enable_multi_turn": False,  # Deterministic for evaluation
            }
        },
        suggested_queries=[
            "What was NVIDIA's annual revenue by fiscal year 2022 to 2025?",
            "When did NVIDIA's data center revenue overtake gaming revenue?",
            "What's the correlation between Neptune's distance from the Sun and US burglary rates?",
            "What's the correlation between Unilever Group's revenue and Google searches for 'lost my wallet'?",
            "Does this imply that Unilever Group's revenue is derived from lost wallets?",
        ],
    )
    agent_id = app_response.id
    print(f"Agent created: {agent_id}")

query_result = client.agents.query.create(
    agent_id=agent_id,
    messages=[{
        "content": "What was NVIDIA's annual revenue by fiscal year 2022 to 2025?",
        "role": "user",
    }],
)
print(query_result.message.content)

For Fiscal Year 2025, the quarterly revenues were $39,331M in Q4, $35,082M
in Q3, $30,040M in Q2, and $26,044M in Q1.[1]()

For Fiscal Year 2024, the quarterly figures were $22,103M in Q4, $18,120M
in Q3, $13,507M in Q2, and $7,192M in Q1.[1]()

import base64, io
from PIL import Image
import matplotlib.pyplot as plt


def display_base64_image(base64_string, title="Document"):
    img_data = base64.b64decode(base64_string)
    img = Image.open(io.BytesIO(img_data))
    plt.figure(figsize=(10, 10))
    plt.imshow(img)
    plt.axis("off")
    plt.title(title)
    plt.show()


for i, retrieval_content in enumerate(query_result.retrieval_contents):
    ret_result = client.agents.query.retrieval_info(
        message_id=query_result.message_id,
        agent_id=agent_id,
        content_ids=[retrieval_content.content_id],
    )
    if ret_result.content_metadatas and ret_result.content_metadatas[0].page_img:
        display_base64_image(
            ret_result.content_metadatas[0].page_img,
            title=f"Source page {i + 1}",
        )

unit_tests = [
    "Does the response accurately extract specific numerical data from the documents?",
    "Does the agent properly distinguish between correlation and causation?",
    "Are multi-document comparisons performed correctly with accurate calculations?",
    "Are potential limitations or uncertainties in the data clearly acknowledged?",
    "Are quantitative claims properly supported with specific evidence from the source documents?",
    "Does the response avoid unnecessary information?",
]

response = client.lmunit.create(
    query="What was NVIDIA's Data Center revenue in Q4 FY25?",
    response="""NVIDIA's Data Center revenue for Q4 FY25 was $35,580 million.[1]()

    This represents an increase from Q3 FY25 ($30,771M), Q2 FY25 ($26,272M),
    and Q1 FY25 ($22,563M).[1]()""",
    unit_test="Does the response avoid unnecessary information?",
)
print(response)
# LMUnitCreateResponse(score=2.338)