import logging import os from pydantic import BaseModel, Field from typing import Literal, List, Any, Annotated from typing_extensions import TypedDict from langchain.schema import Document from langchain_core.prompts import ChatPromptTemplate from langchain_core.prompts.prompt import PromptTemplate from langchain_core.messages import HumanMessage, AIMessage, AnyMessage from langgraph.graph import END, StateGraph, MessagesState, START from langgraph.graph.message import add_messages from huggingface_hub import InferenceClient from dotenv import load_dotenv load_dotenv(verbose=True) assert os.getenv("PINECONE_API_KEY") is not None assert os.getenv("HUGGINGFACEHUB_EMBEDDINGS_MODEL") is not None assert os.getenv("TAVILY_API_KEY") is not None logger = logging.getLogger(__name__) # Child logger for this module logger.setLevel(logging.INFO) logger.info(f"""correctiveRag.py:Config: GROQ_MODEL = {os.getenv('GROQ_MODEL')} HUGGINGFACEHUB_EMBEDDINGS_MODEL = {os.getenv('HUGGINGFACEHUB_EMBEDDINGS_MODEL')} PINECONE_API_KEY = {os.getenv("PINECONE_API_KEY")[:5]} """) # Prepare the LLM from langchain_groq import ChatGroq assert os.getenv('GROQ_MODEL') is not None, "GROQ_MODEL not set" assert os.getenv('GROQ_API_KEY') is not None, "GROQ_API_KEY not set" llm = ChatGroq(model_name=os.getenv('GROQ_MODEL'), temperature=0, verbose=True) # For using Grok # from langchain_openai import ChatOpenAI # assert os.getenv('XAI_API_KEY') is not None, "XAI_API_KEY not set" # assert os.getenv('XAI_MODEL') is not None, "XAI_MODEL not set" # assert os.getenv('XAI_BASE_URL') is not None, "XAI_BASE_URL not set" # llm = ChatOpenAI( # api_key=os.getenv("XAI_API_KEY"), # base_url=os.getenv("XAI_BASE_URL"), # model=os.getenv("XAI_MODEL"), # temperature=0.1) # from langchain_openai import ChatOpenAI # assert os.getenv('OPENAI_MODEL_NAME') is not None, "GROQ_MODEL not set" # llm = ChatOpenAI(model=os.getenv("OPENAI_MODEL_NAME"), temperature=0.1, verbose=True) # Huggingface - Does not support structured_output # llm = InferenceClient("HuggingFaceH4/zephyr-7b-beta") # Prepare the retriever from langchain_huggingface import HuggingFaceEmbeddings from langchain_pinecone import PineconeVectorStore index_name, namespace = 'courses', 'dsa' # Simple RAG # embeddings = HuggingFaceEmbeddings(model_name=os.getenv("HUGGINGFACEHUB_EMBEDDINGS_MODEL")) # docsearch = PineconeVectorStore.from_existing_index(embedding=embeddings, index_name=index_name, namespace=namespace) # retriever = docsearch.as_retriever(search_type="mmr", search_kwargs={ 'k': 5 }) # Large-Small RAG def larger_from_nearby(vectorstore, doc: Document, range:int) -> Document: """ Given a document, find the "parent" document as a range of chunks around the central chunk """ filter0 = { "document" : doc.metadata['document'] } filter1 = { "chunk": { "$gte" : doc.metadata['chunk']-range } } filter2 = { "chunk": { "$lte" : doc.metadata['chunk']+range } } and_filter = { "$and" : [ filter0, filter1, filter2 ] } range_docs = vectorstore.similarity_search(query='', k=2*range+1, filter=and_filter) content = '' for doc in range_docs: content += doc.page_content full_document = Document(page_content=content, metadata=doc.metadata) return full_document def larger_retriever(vectorstore, query:str, topK:int): RANGE=2 # -RANGE...+RANGE logger.info(f'larger_retriever: with RANGE={RANGE}') docs = vectorstore.similarity_search(query, k=topK) larger_documents = list(map(lambda d: larger_from_nearby(vectorstore, d, RANGE), docs)) logger.info(f'larger_retriever: Found {len(larger_documents)} documents.') return larger_documents embeddings = HuggingFaceEmbeddings(model_name=os.getenv("HUGGINGFACEHUB_EMBEDDINGS_MODEL")) vectorstore = PineconeVectorStore.from_existing_index(embedding=embeddings, index_name=index_name, namespace=namespace) # docs = larger_retriever(vectorstore, query, 5) retriever = lambda query: larger_retriever(vectorstore, query, 5) # TODO topK # Classify question class ClassifyQuestion(BaseModel): """Binary score to decide if need to retrieve documents from the vectorstore about data structures and algorithms. The binary_score is "yes" to indicate that document retrieval is needed, otherwise is "no".""" binary_score: str = Field(description="If the question is about data structures and algorithms answer `yes`, otherwise answer `no`") # justification: str = Field(description="Explained reasoning for giving the yes/no score") # LLM with function call structured_llm_grader = llm.with_structured_output(ClassifyQuestion) # Prompt system = """You are an expert at classifying user questions. If the question are specific about data structures and algorithms, then answer `yes` to indicate that document retrieval is needed. Otherwise, it is a question as a general question, answer `no`. """ grade_prompt = ChatPromptTemplate.from_messages( [ ("system", system), ("human", "Question: {question}"), ] ) retriever_grader = grade_prompt | structured_llm_grader # Retrieval grader 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.with_structured_output(GradeDocuments) # Prompt system = """You are a grader assessing relevance of a retrieved document to a user question. If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. It does not need to be a stringent test. The goal is to filter out erroneous retrievals. 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 # Create the RAG chain from langchain import hub from langchain_core.output_parsers import StrOutputParser # prompt = hub.pull("rlm/rag-prompt") # print('----', prompt, '---') template = """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. Please keep the answer concise and to the point. Context: {context} Question: {question} Answer: """ prompt_template = PromptTemplate.from_template(template=template) rag_chain = prompt_template | llm | StrOutputParser() # Question rewriter system = """You a question re-writer that converts an input question to a better version that is optimized for vectorstore retrieval. Look at the input and try to reason about the underlying semantic intent / meaning. Return only the re-written question. Do not return anything else. """ 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() # Web search tool from langchain_community.tools.tavily_search import TavilySearchResults web_search_tool = TavilySearchResults(k=3) # Define the workflow Graph class GraphState(TypedDict): """ Represents the state of our graph. Attributes: messages: conversation history generation: LLM generation web_search: whether to add search documents: list of documents question: the last user question """ messages: Annotated[list[AnyMessage], add_messages] generation: str web_search: str documents: List[str] question: str def chatbot(state: GraphState): logger.info("---GENERATE (no context)---") logger.info(state) chain = llm | StrOutputParser() generation = chain.invoke(state["messages"]) logger.info(generation) return { "messages": [AIMessage(content=generation)], "generation": generation } 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 """ logger.info("---RETRIEVE---") question = state['messages'][-1].content # Last Human message logger.info(f'question: {question}') # Retrieval # documents = retriever.invoke(question) documents = retriever(question) # Large-small retriever # logger.debug(documents) logger.info([ (doc.metadata['id'], doc.page_content[:20])for doc in documents ]) return {"documents": documents, "question": question} def generate_with_context(state): """ Generate answer Args: state (dict): The current graph state Returns: state (dict): New key added to state, generation, that contains LLM generation """ logger.debug("---GENERATE WITH CONTEXT---") logger.debug(f'state: {state}') question = state["question"] documents = state["documents"] # RAG generation generation = rag_chain.invoke({"context": documents, "question": question}) logger.debug(generation) return {"documents": documents, "question": question, "generation": generation} 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 """ logger.debug("---WEB SEARCH---") question = state["question"] documents = state["documents"] # Web search logger.debug(f'question: {question}') docs = web_search_tool.invoke({"query": question}) # Returns str if error logger.debug(f'type(docs) = {type(docs)}') logger.debug(docs) web_results = "\n".join([d["content"] for d in docs]) web_results = Document(page_content=web_results) documents.append(web_results) return {"documents": web_results, "question": question} 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 """ logger.debug("---CHECK DOCUMENT RELEVANCE TO QUESTION---") question = state["question"] documents = state["documents"] # Score each doc filtered_docs = [] web_search = "No" for d in documents: score = retrieval_grader.invoke({"question": question, "document": d.page_content}) grade = score.binary_score if grade == "yes": logger.debug("---GRADE: DOCUMENT RELEVANT---") filtered_docs.append(d) else: logger.debug("---GRADE: DOCUMENT NOT RELEVANT---") web_search = "Yes" continue return {"documents": 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 """ logger.debug("---TRANSFORM QUERY---") question = state["question"] documents = state["documents"] # Re-write question better_question = question_rewriter.invoke({"question": question}) return {"documents": documents, "question": better_question} ### Edges ### # For conditional edges def decide_to_retrieve(state): """ Determines whether to retrieve a context for answering a question. Args: state (dict): The current graph state Returns: str: Binary decision for next node to call """ logger.debug("---ASSESS NEED FOR RETRIEVAL---") # logger.debug(state) question = state['messages'][-1].content # Last Human message logger.debug(question) response = retriever_grader.invoke({ 'question': question }) logger.debug(response) logger.debug(response.binary_score) if response.binary_score == "yes": # All documents have been filtered check_relevance # We will re-generate a new query logger.debug("---DECISION: RETRIEVE DOCUMENTS---") return "retrieve" else: # We have relevant documents, so generate answer logger.debug("---DECISION: GENERAL QUESTION, NO RETRIEVAL---") # state['question'] = question return "chatbot" 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 """ logger.debug("---ASSESS GRADED DOCUMENTS---") state["question"] web_search = state["web_search"] state["documents"] if web_search == "Yes": # All documents have been filtered check_relevance # We will re-generate a new query logger.debug( "---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---" ) return "transform_query" else: # We have relevant documents, so generate answer logger.debug("---DECISION: GENERATE---") return "generate_with_context" # Prepare and compile the Graph workflow = StateGraph(GraphState) # Define the nodes workflow.add_node("chatbot", chatbot) # retrieve workflow.add_node("retrieve", retrieve) # retrieve workflow.add_node("grade_documents", grade_documents) # grade documents workflow.add_node("generate_with_context", generate_with_context) # generate workflow.add_node("transform_query", transform_query) # transform_query workflow.add_node("web_search_node", web_search) # web search # Build graph # workflow.add_edge(START, "retrieve") workflow.add_conditional_edges( START, decide_to_retrieve ) workflow.add_edge("retrieve", "grade_documents") workflow.add_conditional_edges( "grade_documents", decide_to_generate, { "transform_query": "transform_query", "generate_with_context": "generate_with_context", }, ) workflow.add_edge("transform_query", "web_search_node") workflow.add_edge("web_search_node", "generate_with_context") workflow.add_edge("generate_with_context", "chatbot") workflow.add_edge("chatbot", END) # Compile from langgraph.checkpoint.memory import MemorySaver memory = MemorySaver() app = workflow.compile(checkpointer=memory, debug=False) if __name__ == "__main__": # Use the graph from pprint import pprint # print(retriever.invoke("What is an algorithm?")) config = {"configurable": {"thread_id": "abc123"}} def query_graph(question:str): inputs = { "question": question } messages = [HumanMessage(inputs['question'])] response = app.invoke({"messages": messages}, config) # print('TYPE >>', type(response)) # langgraph.pregel.io.AddableValuesDict return response def print_generation(response:str): # pprint(type(response['generation'])) # str # pprint(response['messages']) pprint(response['generation']) # AIMessage (no context) # question = "Cual es el orden de ingreso y egresos de elementos en un Queue?" # pprint(query_graph(question)) # print_generation(query_graph("Hi, my name is George and I would like to learn about algorithms")) # print_generation(query_graph("Do you remember my name? What algorithm would you use to reverse the letters in my name?")) # print_generation(query_graph("Que es un algoritmo?")) # print_generation(query_graph("Que es una heuristica?")) # print_generation(query_graph("Que se entiende por orden de crecimiento de un algoritmo?")) # print_generation(query_graph("Que es la función tilde?")) # print_generation(query_graph("Cuál es la diferencia entre función tilde y orden de crecimiento?")) # In stream mode, returns the full 'chatbot' message for x in app.stream({"messages": "What is the answer to the question of everything?"}, config): print(x)