crag.py

import os
import streamlit as st
from vector_tool import ensemble_retriever
from langgraph.prebuilt import ToolInvocation
from langchain_core.messages import ToolMessage
import json
# Set up the tools to execute them from the graph
from langgraph.prebuilt import ToolExecutor
# tools retrieval
from function_tools import tool_chain
from vector_tool import ensemble_retriever

os.environ['OPENAI_API_KEY'] = st.secrets["OPENAI_API_KEY"]
os.environ['TAVILY_API_KEY'] = st.secrets["TAVILY_API_KEY"]

### Retrieval Grader 

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

#LLM models
llm_AI4 = ChatOpenAI(model="gpt-4-1106-preview", temperature=0)

# Data model
class GradeDocuments(BaseModel):
    """Binary score for relevance check on retrieved documents."""

    binary_score: str = Field(description="Documents are relevant to the question, 'yes' or 'no'")

# LLM with function call 
structured_llm_grader = llm_AI4.with_structured_output(GradeDocuments)

# Prompt 
system = """You are a grader assessing relevance of a retrieved document to a user question. \n 
    If the document contains keyword(s) or semantic meaning related to the question, grade it as relevant. \n
    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""
grade_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "Retrieved document: \n\n {document} \n\n User question: {question}"),
    ]
)

retrieval_grader = grade_prompt | structured_llm_grader

### Generate
from langchain import hub
from langchain.prompts import MessagesPlaceholder
from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
from langchain.prompts import MessagesPlaceholder
from langchain.agents.format_scratchpad.openai_tools import (
    format_to_openai_tool_messages
    )
from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage
from langchain_core.output_parsers import StrOutputParser
from typing import Any, List, Union
# Prompt
#prompt = hub.pull("rlm/rag-prompt")
system_message = '''You are an AI assistant for answering questions about vedas and scriptures.
                    \nYou are given the following extracted documents from Svarupa Knowledge Base (https://svarupa.org/) and other documents and a question. 
                    Provide a conversational answer.\nIf you are not provided with any documents, say \"I did not get any relevant context for this but 
                    I will reply to the best of my knowledge\" and then write your answer\nIf you don't know the answer, just say \"Hmm, I'm not sure. \" Don't try to make up an answer.
                    \nIf the question is not about vedas and scriptures, politely inform them that you are tuned to only answer questions about that.\n\n'''

generate_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_message),
        ("human", "Here is the given context {context}, queation: {question} \n\n Formulate an answer."),
    ]
)
# LLM
llm_AI = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)

# Post-processing
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

# Chain
rag_chain =  generate_prompt | llm_AI4 | StrOutputParser() #OpenAIToolsAgentOutputParser()

####-----------------TESTING
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful assistant. Answer all questions to the best of your ability.",
        ),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{question}"),
    ]
)
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain.memory import ChatMessageHistory

chat_history_for_chain = ChatMessageHistory()

chain_with_message_history = RunnableWithMessageHistory(
    rag_chain,
    lambda session_id: chat_history_for_chain,
    input_messages_key="question",
    history_messages_key="chat_history",
)

### Question Re-writer

# LLM 
llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)

# Prompt 
system = """You a question re-writer that converts an input question to a better version that is optimized \n 
     for a search. Look at the input and try to reason about the underlying sematic intent / meaning."""
re_write_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "Here is the initial question: \n\n {question} \n Formulate an improved question."),
    ]
)

question_rewriter = re_write_prompt | llm | StrOutputParser()



### Search

from langchain_community.tools.tavily_search import TavilySearchResults
web_search_tool = TavilySearchResults(k=2)

from typing_extensions import TypedDict
from typing import List
from typing import TypedDict, Annotated, Sequence
import operator
from langchain_core.messages import BaseMessage

class GraphState(TypedDict):
    """
    Represents the state of our graph.

    Attributes:
        question: question
        generation: LLM generation
        web_search: whether to add search
        documents: list of documents 
    """
    question : str
    generation : str
    web_search : str
    messages: List[str] #Union[dict[str, Any]] 

from langchain.schema import Document



def retrieve(state):
    """
    Retrieve documents

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, documents, that contains retrieved documents
    """
    print("---VECTOR RETRIEVE---")
    question = state["question"]
    # Retrieval
    documents = ensemble_retriever.get_relevant_documents(question)
    #print(documents)
    # Iterate over each document and update the 'metadata' field with the file name
    for doc in documents:
        try:
            file_path = doc.metadata['source']
            #print(file_path)
            file_name = os.path.split(file_path)[1]  # Get the file name from the file path
            doc.metadata['source'] = file_name
        except KeyError:
            # Handle the case where 'source' field is missing in the metadata
            doc.metadata['source'] = 'unavailable'
        except Exception as e:
            # Handle any other exceptions that may occur
            print(f"An error occurred while processing document: {e}")
    return {"messages": documents, "question": question}


def generate(state):
    """
    Generate answer

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, generation, that contains LLM generation
    """
    print("---GENERATE---")
    question = state["question"]
    messages = state["messages"]
    print(messages)
    # RAG generation
    generation = chain_with_message_history.invoke({"context": messages, "question": question},{"configurable": {"session_id": "unused"}})
    return {"messages": messages, "question": question, "generation": generation}

def grade_documents(state):
    """
    Determines whether the retrieved documents are relevant to the question.

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): Updates documents key with only filtered relevant documents
    """

    print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")
    question = state["question"]
    messages = state["messages"]
    
    # Score each doc
    filtered_docs = []
    web_search = "No"
    for d in messages:
        score = retrieval_grader.invoke({"question": question, "document": d.page_content})
        grade = score.binary_score
        if grade == "yes":
            print("---GRADE: DOCUMENT RELEVANT---")
            filtered_docs.append(d)
        else:
            print("---GRADE: DOCUMENT NOT RELEVANT---")
            continue
    print("---TOOLS RETRIEVE---")
    tool_documents = tool_chain.invoke(question)
    #print(tool_documents)
    if tool_documents:
        for item in tool_documents:
            filtered_docs.append(Document(page_content=str(item['output']),metadata={"source": 'https://svarupa.org/home',"name":item['name']}))        
    # If filtered_docs is empty, perform a web search
    if not filtered_docs:
        print("--PERFORMING WEB SEARCH--")
        web_search = "Yes"

    return {"messages": filtered_docs, "question": question, "web_search": web_search}



def transform_query(state):
    """
    Transform the query to produce a better question.

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): Updates question key with a re-phrased question
    """

    print("---TRANSFORM QUERY---")
    question = state["question"]
    messages = state["messages"]

    # Re-write question
    better_question = question_rewriter.invoke({"question": question})
    return {"messages": messages, "question": better_question}
    
def web_search(state):
    """
    Web search based on the re-phrased question.

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): Updates documents key with appended web results
    """

    print("---WEB SEARCH---")
    question = state["question"]
    messages = state["messages"]

    # Web search
    docs = web_search_tool.invoke({"query": question})
    print(docs)
    #web_results = "\n".join([d["content"] for d in docs])
    web_results = [Document(page_content=d["content"], metadata={"source": d["url"]}) for d in docs]
    print(f"Web Results: {web_results}")
    messages.extend(web_results)
    return {"messages": messages, "question": question}

### Edges

def decide_to_generate(state):
    """
    Determines whether to generate an answer, or re-generate a question.

    Args:
        state (dict): The current graph state

    Returns:
        str: Binary decision for next node to call
    """

    print("---ASSESS GRADED DOCUMENTS---")
    question = state["question"]
    web_search = state["web_search"]
    filtered_documents = state["messages"]

    if web_search == "Yes":
        # All documents have been filtered check_relevance
        # We will re-generate a new query
        print("---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---")
        return "transform_query"
    else:
        # We have relevant documents, so generate answer
        print("---DECISION: GENERATE---")
        return "generate"
from langgraph.graph import END, StateGraph

workflow = StateGraph(GraphState)

# Define the nodes
workflow.add_node("retrieve", retrieve)  # retrieve
workflow.add_node("grade_documents", grade_documents)  # grade documents
workflow.add_node("generate", generate)  # generatae
workflow.add_node("transform_query", transform_query)  # transform_query
workflow.add_node("web_search_node", web_search)  # web search

# Build graph
workflow.set_entry_point("retrieve")
workflow.add_edge("retrieve", "grade_documents")
workflow.add_conditional_edges(
    "grade_documents",
    decide_to_generate,
    {
        "transform_query": "transform_query",
        "generate": "generate",
    },
)
workflow.add_edge("transform_query", "web_search_node")
workflow.add_edge("web_search_node", "generate")
workflow.add_edge("generate", END)

# Compile
crag_app = workflow.compile()