!pip install --quiet --upgrade langchain langchain-community langchain-chroma pandas notdiamond[create]

import os

os.environ["NOTDIAMOND_API_KEY"] = 'NOTDIAMOND_API_KEY'
os.environ["OPENAI_API_KEY"] = "OPENAI_API_KEY"
os.environ["ANTHROPIC_API_KEY"] = 'ANTHROPIC_API_KEY'

import bs4
from langchain_chroma import Chroma
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

# Load, chunk and index the contents of the blog.
agents_blog_loader = WebBaseLoader(
    web_paths=("https://lilianweng.github.io/posts/2023-06-23-agent/",),
    bs_kwargs=dict(
        parse_only=bs4.SoupStrainer(
            class_=("post-content", "post-title", "post-header")
        )
    ),
)
agent_blog_docs = agents_blog_loader.load()

agent_blog_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
agent_blog_splits = agent_blog_splitter.split_documents(agent_blog_docs)
agent_blog_vectorstore = Chroma.from_documents(documents=agent_blog_splits, embedding=OpenAIEmbeddings())

# Retrieve and generate using the relevant snippets of the blog.
agent_blog_retriever = agent_blog_vectorstore.as_retriever()

rorf_loader = WebBaseLoader(
    web_paths=("https://www.notdiamond.ai/blog/rorf",)
)
rorf_docs = rorf_loader.load()

rorf_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=900)
rorf_splits = rorf_splitter.split_documents(rorf_docs)
rorf_vectorstore = Chroma.from_documents(documents=rorf_splits, embedding=OpenAIEmbeddings())

rorf_retriever = rorf_vectorstore.as_retriever()

from notdiamond.llms.config import LLMConfig

agentic_blog_agent = LLMConfig(
    provider="agent",
    model="agentic_blog",
    is_custom=True,
    context_length=128000,
    input_price=2.5, # USD per million tokens
    output_price=10.0, # USD per million tokens
    latency=1 # time to first token (seconds)
)

rorf_agent = LLMConfig(
    provider="agent",
    model="rorf_blog",
    is_custom=True,
    context_length=200000,
    input_price=3., # USD per million tokens
    output_price=15., # USD per million tokens
    latency=1 # time to first token (seconds)
)

agents = [
    agentic_blog_agent,
    rorf_agent
]

import pandas as pd
from pprint import pprint

train_data = {}
for provider in agents:
    df = {"input": [], "response": [], "score": []}

    for doc in agent_blog_splits:
        score = 1 if provider.model == "agentic_blog" else 0
        df["input"].append(doc.page_content)
        df["response"].append("")
        df["score"].append(score)

    for doc in rorf_splits:
        score = 1 if provider.model == "rorf_blog" else 0
        df["input"].append(doc.page_content)
        df["response"].append("")
        df["score"].append(score)

    train_data[provider] = pd.DataFrame(df)

provider = agents[0]
pprint(f"Train samples: {len(train_data[provider])}")

from notdiamond.toolkit import CustomRouter

# Initialize the CustomRouter object for training
trainer = CustomRouter(
    language="english",
    maximize=True,  # Indicate if higher scores are better (setting to False indicates the opposite)
)

# Train the model using your dataset
preference_id = trainer.fit(
    dataset=train_data,  # The dataset containing inputs, responses, and scores
    prompt_column="input",  # Column name for the input prompts
    response_column="response",  # Column name for the model responses
    score_column="score"  # Column name for the scores
)
print(preference_id)

from notdiamond import NotDiamond
from notdiamond.toolkit.langchain import NotDiamondRoutedRunnable, NotDiamondRunnable

from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate

# This client is instantiated with the preference_id, which will call the custom router we trained
custom_router_client = NotDiamond(preference_id=preference_id)

# This client uses Not Diamond's out-of-the-box (OOB) router
oob_client = NotDiamond()

system_prompt = """
You are an assistant for question-answering tasks.
Use the following pieces of retrieved context to answer the question.
If you don't know the answer, just say that you don't know.
Use three sentences maximum and keep the answer concise.
Context: {context}
"""

# Define the prompt template to use for the RAG request
rag_prompt = ChatPromptTemplate(
    [
        ("system", system_prompt),
        ("human", "{question}")
    ]
)

def format_docs(docs):
		"""Concatenate all the relevant documents into a single string"""
    return "\n\n".join(doc.page_content for doc in docs)

# These are the LLMs we want to route between depending on query difficulty
# in the rorf_agent
rorf_llms = [
    "anthropic/claude-3-haiku-20240307", 
    "anthropic/claude-3-5-sonnet-20240620"
]

# Define the runnable that will return the recommended LLM
rorf_router = NotDiamondRunnable(
    nd_client=oob_client,
    nd_kwargs={
        'tradeoff': "cost", # By setting 'tradeoff': 'cost' Not Diamond will route queries to cheaper LLMs when doing so does not affect quality
        "model": rorf_llms
        }
)

# Define the runnable that will return the recommended LLM's response
rorf_routed_runnable = NotDiamondRoutedRunnable(
    nd_client=oob_client, 
    nd_kwargs={
        'tradeoff': "cost", 
        "model": rorf_llms
        }
    )
    
def rorf_prompt_formatter(input):
		"""Print the recommended model and format the prompt template"""
    model = rorf_router.invoke(input["question"])
    print(f"This prompt is routed to the RoRF blog agent using {model}")
    return rag_prompt.invoke(input)

agentic_blog_llms = [
    "openai/gpt-4o-mini", 
    "openai/gpt-4o"
]

agentic_blog_router = NotDiamondRunnable(
    nd_client=oob_client,
    nd_kwargs={
        'tradeoff': "cost", 
        "model": agentic_blog_llms
        }
)

agentic_blog_routed_runnable = NotDiamondRoutedRunnable(
    nd_client=oob_client,
    nd_kwargs={
        'tradeoff': "cost", 
        "model": agentic_blog_llms
        }
    )

def agentic_prompt_formatter(input):
    model = agentic_blog_router.invoke(input["question"])
    print(f"This prompt is routed to the agentic blog agent using {model}")
    return rag_prompt.invoke(input)

agent_blog_agent = (
    {"context": agent_blog_retriever | format_docs, "question": RunnablePassthrough()}
    | RunnableLambda(agentic_prompt_formatter)
    | agentic_blog_routed_runnable
    | StrOutputParser()
)

rorf_agent = (
    {"context": rorf_retriever | format_docs, "question": RunnablePassthrough()}
    | RunnableLambda(rorf_prompt_formatter)
    | rorf_routed_runnable
    | StrOutputParser()
)

def routes(query):
    if query["agent"] == "agent/agentic_blog":
        return agent_blog_agent.invoke(query["question"])
    else:
        return rorf_agent.invoke(query["question"])

def agent_router(query):
    messages = [
        {"role": "user", "content": query}
    ]
    session_id, provider = custom_router_client.chat.completions.model_select(
        messages=messages,
        model=agents
    )
    return str(provider)

full_chain = (
    {"agent": RunnableLambda(agent_router), "question": RunnablePassthrough()} 
    | RunnableLambda(routes)
    )

full_chain.invoke("How many models does RoRF LLM router support?")

This prompt is routed to the RoRF blog agent using anthropic/claude-3-haiku-20240307

According to the context, the RoRF LLM router supports routing between two models. The passage states that "RoRF can only route between two models and it does not have the ability to learn in real-time based on feedback."

full_chain.invoke("What is Task Decomposition?")

This prompt is routed to the agentic blog agent using openai/gpt-4o-mini

Task decomposition is the process of breaking down a complicated task into smaller, more manageable steps. It often involves techniques like Chain of Thought (CoT), where the model is prompted to "think step by step" to enhance its performance on complex tasks. This approach helps clarify the model\'s reasoning process and makes it easier to tackle difficult tasks.

full_chain.invoke("Explain self-reflection for agentic workflows and write a short example using langchain")

This prompt is routed to the agentic blog agent using openai/gpt-4o

Self-reflection in agentic workflows refers to the process by which autonomous agents iteratively refine their decision-making by analyzing and learning from past actions and mistakes. This capability is crucial for improving performance over time, particularly in complex environments where trial and error are common. Through self-reflection, agents can dynamically adjust their strategies to better achieve their objectives.\n\nExample using Langchain:\nImagine an agent using Langchain to manage a customer support chatbot. After each interaction, the agent logs the conversation and evaluates whether the customer issue was resolved satisfactorily. If not, the agent reviews the conversation to identify where the response could be improved, updates its response database, and applies this learning in future interactions to provide better support.