Routing

%load_ext dotenv
%dotenv secrets/secrets.env

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_classic import hub
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

def generate_vectorstores(file, dir):
    loader = PyPDFLoader(file)
    documents = loader.load()

    # Split text into chunks

    text_splitter  = RecursiveCharacterTextSplitter(chunk_size=500,chunk_overlap=20)
    text_chunks = text_splitter.split_documents(documents)

    vectorstore = Chroma.from_documents(documents=text_chunks, 
                                        embedding=OpenAIEmbeddings(),
                                        persist_directory=dir)
    vectorstore.persist()
    return vectorstore


# Create a vectorstore to answer questions about LoRA
vectorstore_lora = generate_vectorstores("data/LoRA.pdf","data/vectorstore_lora")

# Create a vectorstore to answer questions about BERT
vectorstore_bert = generate_vectorstores("data/BERT.pdf","data/vectorstore_bert")

retriever_lora = vectorstore_lora.as_retriever(search_kwargs={'k':5})
retriever_bert = vectorstore_bert.as_retriever(search_kwargs={'k':5})

from pydantic import BaseModel, Field
from langchain_core.prompts import ChatPromptTemplate
from typing import Literal

class QueryRouter(BaseModel):
    
    """Route a user query to the appropriate datasource that will help answer the query accurately"""
    
    datasource: Literal['lora', 'bert', 'general'] = Field(..., 
                                                description="Given a user question choose which datasource would be most relevant for answering their question"
                                                )
    question: str = Field(..., description="User question to be routed to the appropriate datasource")
    
llm = ChatOpenAI(model='gpt-4',temperature=0)
structured_llm = llm.with_structured_output(QueryRouter)

router_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are an expert router that can direct user queries to the appropriate datasource. Route the following user question about a topic in NLP and LLMs to the appropriate datasource.\nIf it is a general question not related to the provided datasources, route it to the general datasource.\n"),
        ("user", "{question}")
    ]
)

router = (
    {'question': RunnablePassthrough()}
    | router_prompt
    | structured_llm
)

question = "How does the BERT work?"
result = router.invoke(question)
result

qa_prompt = hub.pull('rlm/rag-prompt')


def choose_route(result):
    
    llm_route = ChatOpenAI(model='gpt-4',temperature=0)
    if "bert" in result.datasource.lower():
        print(f"> Asking about BERT ...\nQuestion: {result.question}\nAnswer:")
        bert_chain = (
            {'context': retriever_bert, 'question': RunnablePassthrough()}
            | qa_prompt
            | llm_route
            | StrOutputParser()
        )
        return bert_chain.invoke(result.question)
    elif "lora" in result.datasource.lower():
        print(f"> Asking about LoRA ...\nQuestion: {result.question}\nAnswer:")
        lora_chain = (
            {'context': retriever_lora, 'question': RunnablePassthrough()}
            | qa_prompt
            | llm_route
            | StrOutputParser()
        )
        return lora_chain.invoke(result.question)
    else:
        print(f"> Asking about a general question ...\nQuestion: {result.question}\nAnswer:")
        general_chain = llm_route | StrOutputParser()
        return general_chain.invoke(result.question)

from langchain_core.runnables import RunnableLambda

full_chain = router | RunnableLambda(choose_route)

full_chain.invoke("What are the benefits of LoRA?")

physics_template = """You are a very smart physics professor. \
You are great at answering questions about physics in a concise and easy to understand manner. \
When you don't know the answer to a question you admit that you don't know.

Here is a question:
{question}"""

math_template = """You are a very good mathematician. You are great at answering math questions. \
You are so good because you are able to break down hard problems into their component parts, \
answer the component parts, and then put them together to answer the broader question.

Here is a question:
{question}"""

embeddings = OpenAIEmbeddings()
routes = [physics_template, math_template]
route_embeddings = embeddings.embed_documents(routes)
len(route_embeddings)

from langchain_classic.utils.math import cosine_similarity
from langchain_core.prompts import PromptTemplate

def router(input):
    # Generate embeddings for the user query
    query_embedding = embeddings.embed_query(input['question'])
    # Getting similarity scores between the user query and the routes. This contains the similarity scores between the user query and each of the two routes.
    similarity = cosine_similarity([query_embedding], route_embeddings)[0]
    # Find the route that gives the maximum similarity score
    route_id = similarity.argmax()
    if route_id == 0:
        print(f"> Asking a physics question ...\nQuestion: {input['question']}\nAnswer:")
    else:
        print(f"> Asking a math question ...\nQuestion: {input['question']}\nAnswer:")
        
    return PromptTemplate.from_template(routes[route_id])

semantic_router_chain = (
    {'question': RunnablePassthrough()}
    | RunnableLambda(router)
    | ChatOpenAI(model='gpt-4',temperature=0)
    | StrOutputParser()
)

semantic_router_chain.invoke("What is the formula for the area of a circle?")