#!/usr/bin/env python
# coding: utf-8

# # Rag From Scratch: Routing
# 
# ![image.png](attachment:c02ab9b5-38f9-451a-b202-62b54ab9c87a.png)
# 
# ## Enviornment
# 
# `(1) Packages`

# In[ ]:


get_ipython().system(' pip install langchain_community tiktoken langchain-openai langchainhub chromadb langchain youtube-transcript-api pytube')


# `(2) LangSmith`
# 
# https://docs.smith.langchain.com/

# In[ ]:


import os
os.environ['LANGCHAIN_TRACING_V2'] = 'true'
os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'
os.environ['LANGCHAIN_API_KEY'] = <your-api-key>


# `(3) API Keys`

# In[ ]:


os.environ['OPENAI_API_KEY'] = <your-api-key>


# ## Part 10: Logical and Semantic routing 
# 
# Use function-calling for classification.
# 
# Flow: 
# 
# ![Screenshot 2024-03-15 at 3.29.30 PM.png](attachment:b6699c4f-6188-4e0e-8ba4-21582dbca9ef.png)
# 
# Docs:
# 
# https://python.langchain.com/docs/use_cases/query_analysis/techniques/routing#routing-to-multiple-indexes

# In[1]:


from typing import Literal

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI

# Data model
class RouteQuery(BaseModel):
    """Route a user query to the most relevant datasource."""

    datasource: Literal["python_docs", "js_docs", "golang_docs"] = Field(
        ...,
        description="Given a user question choose which datasource would be most relevant for answering their question",
    )

# LLM with function call 
llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
structured_llm = llm.with_structured_output(RouteQuery)

# Prompt 
system = """You are an expert at routing a user question to the appropriate data source.

Based on the programming language the question is referring to, route it to the relevant data source."""

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "{question}"),
    ]
)

# Define router 
router = prompt | structured_llm


# Note: we used function calling to produce structured output.
# 
# ![Screenshot 2024-03-16 at 12.38.23 PM.png](attachment:1c7e2e9e-e85f-490f-9591-883a4070bdb2.png)

# In[2]:


question = """Why doesn't the following code work:

from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_messages(["human", "speak in {language}"])
prompt.invoke("french")
"""

result = router.invoke({"question": question})


# In[3]:


result


# In[4]:


result.datasource


# Once we have this, it is trivial to define a branch that uses `result.datasource`
# 
# https://python.langchain.com/docs/expression_language/how_to/routing

# In[39]:


def choose_route(result):
    if "python_docs" in result.datasource.lower():
        ### Logic here 
        return "chain for python_docs"
    elif "js_docs" in result.datasource.lower():
        ### Logic here 
        return "chain for js_docs"
    else:
        ### Logic here 
        return "golang_docs"

from langchain_core.runnables import RunnableLambda

full_chain = router | RunnableLambda(choose_route)


# In[40]:


full_chain.invoke({"question": question})


# Trace:
# 
# https://smith.langchain.com/public/c2ca61b4-3810-45d0-a156-3d6a73e9ee2a/r

# ### Semantic routing
# 
# Flow:
# 
# ![Screenshot 2024-03-15 at 3.30.08 PM.png](attachment:77626ada-cabe-4ecb-a8eb-22992883c5dc.png)
# 
# Docs:
# 
# https://python.langchain.com/docs/expression_language/cookbook/embedding_router

# In[5]:


from langchain.utils.math import cosine_similarity
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

# Two prompts
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:
{query}"""

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:
{query}"""

# Embed prompts
embeddings = OpenAIEmbeddings()
prompt_templates = [physics_template, math_template]
prompt_embeddings = embeddings.embed_documents(prompt_templates)

# Route question to prompt 
def prompt_router(input):
    # Embed question
    query_embedding = embeddings.embed_query(input["query"])
    # Compute similarity
    similarity = cosine_similarity([query_embedding], prompt_embeddings)[0]
    most_similar = prompt_templates[similarity.argmax()]
    # Chosen prompt 
    print("Using MATH" if most_similar == math_template else "Using PHYSICS")
    return PromptTemplate.from_template(most_similar)


chain = (
    {"query": RunnablePassthrough()}
    | RunnableLambda(prompt_router)
    | ChatOpenAI()
    | StrOutputParser()
)

print(chain.invoke("What's a black hole"))


# Trace: 
# 
# https://smith.langchain.com/public/98c25405-2631-4de8-b12a-1891aded3359/r

# # Rag From Scratch: Query Construction
# 
# ![Screenshot 2024-03-25 at 8.20.28 PM.png](attachment:25eea077-2a6a-4787-ad6c-fa6ba8e0f9c1.png)
# 
# For graph and SQL, see helpful resources:
# 
# https://blog.langchain.dev/query-construction/
# 
# https://blog.langchain.dev/enhancing-rag-based-applications-accuracy-by-constructing-and-leveraging-knowledge-graphs/

# ## Part 11: Query structuring for metadata filters
# 
# Flow:
# 
# ![Screenshot 2024-03-16 at 1.12.10 PM.png](attachment:3d933538-e73d-4922-8fe6-dbd2fc2cf2f5.png)
# 
# Many vectorstores contain metadata fields. 
# 
# This makes it possible to filter for specific chunks based on metadata.
# 
# Let's look at some example metadata we might see in a database of YouTube transcripts.
# 
# Docs:
# 
# https://python.langchain.com/docs/use_cases/query_analysis/techniques/structuring

# In[11]:


from langchain_community.document_loaders import YoutubeLoader

docs = YoutubeLoader.from_youtube_url(
    "https://www.youtube.com/watch?v=pbAd8O1Lvm4", add_video_info=True
).load()

docs[0].metadata


# Let’s assume we’ve built an index that:
# 
# 1. Allows us to perform unstructured search over the `contents` and `title` of each document
# 2. And to use range filtering on `view count`, `publication date`, and `length`.
# 
# We want to convert natural langugae into structured search queries.
# 
# We can define a schema for structured search queries.

# In[12]:


import datetime
from typing import Literal, Optional, Tuple
from langchain_core.pydantic_v1 import BaseModel, Field

class TutorialSearch(BaseModel):
    """Search over a database of tutorial videos about a software library."""

    content_search: str = Field(
        ...,
        description="Similarity search query applied to video transcripts.",
    )
    title_search: str = Field(
        ...,
        description=(
            "Alternate version of the content search query to apply to video titles. "
            "Should be succinct and only include key words that could be in a video "
            "title."
        ),
    )
    min_view_count: Optional[int] = Field(
        None,
        description="Minimum view count filter, inclusive. Only use if explicitly specified.",
    )
    max_view_count: Optional[int] = Field(
        None,
        description="Maximum view count filter, exclusive. Only use if explicitly specified.",
    )
    earliest_publish_date: Optional[datetime.date] = Field(
        None,
        description="Earliest publish date filter, inclusive. Only use if explicitly specified.",
    )
    latest_publish_date: Optional[datetime.date] = Field(
        None,
        description="Latest publish date filter, exclusive. Only use if explicitly specified.",
    )
    min_length_sec: Optional[int] = Field(
        None,
        description="Minimum video length in seconds, inclusive. Only use if explicitly specified.",
    )
    max_length_sec: Optional[int] = Field(
        None,
        description="Maximum video length in seconds, exclusive. Only use if explicitly specified.",
    )

    def pretty_print(self) -> None:
        for field in self.__fields__:
            if getattr(self, field) is not None and getattr(self, field) != getattr(
                self.__fields__[field], "default", None
            ):
                print(f"{field}: {getattr(self, field)}")


# Now, we prompt the LLM to produce queries.

# In[13]:


from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

system = """You are an expert at converting user questions into database queries. \
You have access to a database of tutorial videos about a software library for building LLM-powered applications. \
Given a question, return a database query optimized to retrieve the most relevant results.

If there are acronyms or words you are not familiar with, do not try to rephrase them."""
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "{question}"),
    ]
)
llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
structured_llm = llm.with_structured_output(TutorialSearch)
query_analyzer = prompt | structured_llm


# In[14]:


query_analyzer.invoke({"question": "rag from scratch"}).pretty_print()


# In[15]:


query_analyzer.invoke(
    {"question": "videos on chat langchain published in 2023"}
).pretty_print()


# In[16]:


query_analyzer.invoke(
    {"question": "videos that are focused on the topic of chat langchain that are published before 2024"}
).pretty_print()


# In[14]:


query_analyzer.invoke(
    {
        "question": "how to use multi-modal models in an agent, only videos under 5 minutes"
    }
).pretty_print()


# To then connect this to various vectorstores, you can follow [here](https://python.langchain.com/docs/modules/data_connection/retrievers/self_query#constructing-from-scratch-with-lcel).

# In[ ]: