šŸ““ TruLens Quickstart - šŸ¦‘ TruLens

šŸ““ TruLens Quickstart

In this quickstart you will create a RAG from scratch, trace the execution and get feedback on an LLM response.

For evaluation, we will leverage the "hallucination triad" of groundedness, context relevance and answer relevance.

InĀ [Ā ]:

# !pip install trulens trulens-providers-openai chromadb openai

InĀ [Ā ]:

import os

if "OPENAI_API_KEY" not in os.environ:
    os.environ["OPENAI_API_KEY"] = "sk-proj-..."

Get Data

In this case, we'll just initialize some simple text in the notebook.

InĀ [Ā ]:

seattle_info = """
Seattle, a city on Puget Sound in the Pacific Northwest, is surrounded by water, mountains and evergreen forests, and contains thousands of acres of parkland.
It's home to a large tech industry, with Microsoft and Amazon headquartered in its metropolitan area.
The futuristic Space Needle, a legacy of the 1962 World's Fair, is its most iconic landmark.
"""

starbucks_info = """
Starbucks Corporation is an American multinational chain of coffeehouses and roastery reserves headquartered in Seattle, Washington.
As the world's largest coffeehouse chain, Starbucks is seen to be the main representation of the United States' second wave of coffee culture.
"""

coffee_culture_info = """
Coffee culture has evolved through three distinct waves. The first wave focused on convenience and mass production,
exemplified by brands like Folgers and Maxwell House. The second wave, led by Starbucks, introduced espresso-based drinks,
customization, and the cafe as a 'third place' between work and home. The third wave treats coffee as artisanal food,
emphasizing origin, processing methods, and brewing techniques.
"""

seattle_coffee_info = """
Seattle became the epicenter of American coffee culture, birthplace of Starbucks in 1971, and home to numerous independent roasters.
The city's coffee scene has evolved from the second wave dominance of Starbucks to embrace third wave coffee shops
that focus on single-origin beans and precise brewing methods.
"""

ocean_waves_info = """
Ocean waves along the United States coastline provide significant renewable energy opportunities through wave power generation.
The Pacific Northwest, particularly off the coasts of Washington and Oregon, has some of the best wave energy resources in the nation.
Wave energy converters can harness this power to generate electricity.
"""

radio_waves_info = """
Radio waves were first successfully transmitted across the Atlantic Ocean in 1901 by Guglielmo Marconi.
In the United States, radio wave technology revolutionized communication and entertainment throughout the 20th century.
The Federal Communications Commission regulates radio wave frequencies to prevent interference.
"""

starbucks_stock_info = """
Starbucks Corporation (NASDAQ: SBUX) has shown volatile stock performance in recent years, influenced by market trends,
expansion strategies, and consumer spending patterns. The company's stock price reflects broader economic conditions
and competition in the quick-service restaurant sector.
"""

heat_waves_info = """
Heat waves in the United States have become more frequent and intense due to climate change.
The Pacific Northwest, including Washington State, experienced record-breaking temperatures in recent heat waves,
challenging infrastructure designed for milder climates.
"""

Create Vector Store

Create a chromadb vector store in memory.

InĀ [Ā ]:

import chromadb
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction

embedding_function = OpenAIEmbeddingFunction(
    api_key=os.environ.get("OPENAI_API_KEY"),
    model_name="text-embedding-3-small",
)

chroma_client = chromadb.Client()
vector_store = chroma_client.get_or_create_collection(
    name="Washington", embedding_function=embedding_function
)

Populate the vector store.

InĀ [Ā ]:

vector_store.add("seattle_info", documents=seattle_info)
vector_store.add("starbucks_info", documents=starbucks_info)
vector_store.add("coffee_culture_info", documents=coffee_culture_info)
vector_store.add("seattle_coffee_info", documents=seattle_coffee_info)
vector_store.add("ocean_waves_info", documents=ocean_waves_info)
vector_store.add("radio_waves_info", documents=radio_waves_info)
vector_store.add("starbucks_stock_info", documents=starbucks_stock_info)
vector_store.add("heat_waves_info", documents=heat_waves_info)

Build RAG from scratch

Build a custom RAG from scratch, and add TruLens custom instrumentation.

InĀ [Ā ]:

from trulens.core import TruSession

session = TruSession()
session.reset_database()

InĀ [Ā ]:

from openai import OpenAI
from trulens.core.otel.instrument import instrument
from trulens.otel.semconv.trace import SpanAttributes

oai_client = OpenAI()

class RAG:
    def __init__(self, model_name: str = "gpt-5"):
        self.model_name = model_name

@instrument(
        span_type=SpanAttributes.SpanType.RETRIEVAL,
        attributes={
            SpanAttributes.RETRIEVAL.QUERY_TEXT: "query",
            SpanAttributes.RETRIEVAL.RETRIEVED_CONTEXTS: "return",
        },
    )
    def retrieve(self, query: str) -> list:
        results = vector_store.query(
            query_texts=query, n_results=4
        )  # Get more results
        return [doc for sublist in results["documents"] for doc in sublist]

@instrument(span_type=SpanAttributes.SpanType.GENERATION)
    def generate_completion(self, query: str, context_list: list) -> str:
        """Generate answer from context with improved prompting."""
        if len(context_list) == 0:
            return "I don't have enough relevant information to answer this question."

context = (
            "\n---\n".join(context_list)
            if isinstance(context_list, list)
            else context_list
        )

completion = (
            oai_client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {
                        "role": "system",
                        "content": "You are a helpful assistant. ",
                    },
                    {
                        "role": "user",
                        "content": f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer using the context above:",
                    },
                ],
            )
            .choices[0]
            .message.content
        )
        return (
            completion
            if completion
            else "I don't have enough information to answer this question."
        )

@instrument(
        span_type=SpanAttributes.SpanType.RECORD_ROOT,
        attributes={
            SpanAttributes.RECORD_ROOT.INPUT: "query",
            SpanAttributes.RECORD_ROOT.OUTPUT: "return",
        },
    )
    def query(self, query: str) -> str:
        context_list = self.retrieve(query=query)
        completion = self.generate_completion(
            query=query, context_list=context_list
        )
        return completion

Feedback functions

InĀ [Ā ]:

import numpy as np
from trulens.core import Metric
from trulens.core import Selector
from trulens.providers.openai import OpenAI

provider = OpenAI(model_engine="gpt-5-nano")

# Define a groundedness metric
f_groundedness = Metric(
    implementation=provider.groundedness_measure_with_cot_reasons_consider_answerability,
    name="Groundedness",
    selectors={
        "source": Selector.select_context(collect_list=True),
        "statement": Selector.select_record_output(),
        "question": Selector.select_record_input(),
    },
)

# Question/answer relevance between overall question and answer.
f_answer_relevance = Metric(
    implementation=provider.relevance_with_cot_reasons,
    name="Answer Relevance",
    selectors={
        "prompt": Selector.select_record_input(),
        "response": Selector.select_record_output(),
    },
)

# Context relevance between question and each context chunk.
f_context_relevance = Metric(
    implementation=provider.context_relevance_with_cot_reasons,
    name="Context Relevance",
    selectors={
        "question": Selector.select_record_input(),
        "context": Selector.select_context(collect_list=False),
    },
    agg=np.mean,  # choose a different aggregation method if you wish
)

Construct the app

Wrap the custom RAG with TruApp, add list of feedbacks for eval

InĀ [Ā ]:

from trulens.apps.app import TruApp

rag = RAG(model_name="gpt-5-mini")

tru_rag = TruApp(
    rag,
    app_name="RAG",
    app_version="base",
    feedbacks=[f_groundedness, f_answer_relevance, f_context_relevance],
)

Run the app

Use tru_rag as a context manager for the custom RAG-from-scratch app.

InĀ [Ā ]:

test_queries = [
    "What wave of coffee culture does Starbucks represent?",
    "Describe climate challenges for Starbucks.",
]

InĀ [Ā ]:

with tru_rag as recording:
    for query in test_queries:
        rag.query(query)

Check results

We can view results in the leaderboard.

InĀ [Ā ]:

session.get_leaderboard()

InĀ [Ā ]:

from trulens.dashboard import run_dashboard

run_dashboard(session)