Using Not Diamond to Select LLMs For Indexes

In this tutorial, we demonstrate how to use a router query engine with a selector powered by Not Diamond. You can automatically route a query to one of several available LLMs, which will then select the best index for your needs.

Setup

python

%pip install -q llama-index-llms-anthropic llama-index-llms-openai

python

!pip install -q llama-index notdiamond

python

# NOTE: This is ONLY necessary in jupyter notebook.
# Details: Jupyter runs an event-loop behind the scenes.
#          This results in nested event-loops when we start an event-loop to make async queries.
#          This is normally not allowed, we use nest_asyncio to allow it for convenience.
import nest_asyncio

nest_asyncio.apply()

python

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

Routing Queries With Not Diamond

python

import os
from typing import List

os.environ["OPENAI_API_KEY"] = "sk-..."
os.environ["ANTHROPIC_API_KEY"] = "sk-ant-..."
os.environ["NOTDIAMOND_API_KEY"] = "sk-..."

Create Indexes

python

from llama_index.core import (
    SimpleDirectoryReader,
    VectorStoreIndex,
    SummaryIndex,
    Settings,
)
from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.tools import QueryEngineTool
from llama_index.selectors.notdiamond.base import NotDiamondSelector

# load documents
documents = SimpleDirectoryReader("data/paul_graham").load_data()
nodes = Settings.node_parser.get_nodes_from_documents(documents)

# index documents
vector_index = VectorStoreIndex.from_documents(documents)
summary_index = SummaryIndex.from_documents(documents)
query_text = "What was Paul Graham's role at Yahoo?"

Set up Tools for the QueryEngine

python

list_query_engine = summary_index.as_query_engine(
    response_mode="tree_summarize",
    use_async=True,
)
vector_query_engine = vector_index.as_query_engine()

list_tool = QueryEngineTool.from_defaults(
    query_engine=list_query_engine,
    description=(
        "Useful for summarization questions related to Paul Graham eassy on"
        " What I Worked On."
    ),
)

vector_tool = QueryEngineTool.from_defaults(
    query_engine=vector_query_engine,
    description=(
        "Useful for retrieving specific context from Paul Graham essay on What"
        " I Worked On."
    ),
)

Create a NotDiamondSelector and RouterQueryEngine

python

from notdiamond import NotDiamond

client = NotDiamond(
    api_key=os.environ["NOTDIAMOND_API_KEY"],
    llm_configs=["openai/gpt-4o", "anthropic/claude-3-5-sonnet-20240620"],
)
preference_id = client.create_preference_id()
client.preference_id = preference_id

nd_selector = NotDiamondSelector(client=client)

query_engine = RouterQueryEngine(
    selector=nd_selector,
    query_engine_tools=[
        list_tool,
        vector_tool,
    ],
)

Use Not Diamond to Query Indexes

Once we've set up our indexes and query engine, we can submit queries as usual.

python

response = query_engine.query(
    "Please summarize Paul Graham's working experience."
)
print(str(response))

python

response = query_engine.query("What did Paul Graham do after RICS?")
print(str(response))

Using NotDiamondSelector as a standalone selector

As with LlamaIndex's built-in selectors, you can also use the NotDiamondSelector to select an index.

python

from llama_index.core.tools import ToolMetadata
from llama_index.selectors.notdiamond.base import NotDiamondSelector

from notdiamond import NotDiamond, Metric

choices = [
    ToolMetadata(
        name="vector_index",
        description="Great for asking questions about recipes.",
    ),
    ToolMetadata(
        name="list_index", description="Great for summarizing recipes."
    ),
]

llm_configs = ["openai/gpt-4o", "anthropic/claude-3-5-sonnet-20240620"]
nd_client = NotDiamond(
    api_key=os.environ["NOTDIAMOND_API_KEY"],
    llm_configs=llm_configs,
    preference_id=preference_id,
)
preference_id = nd_client.create_preference_id()
nd_client.preference_id = preference_id
nd_selector = NotDiamondSelector(client=nd_client)

nd_result = nd_selector.select(
    choices, query="What is the summary of this recipe for deviled eggs?"
)
print(nd_result)

# Use the result's session_id to customize your routing logic.
metric = Metric("accuracy")
score = metric.feedback(
    session_id=nd_result.session_id,
    llm_config=nd_result.llm,
    value=1,
)