jmlon/DataStructures_and_Algorithms

1import logging2 3import os4from pydantic import BaseModel, Field5from typing import Literal, List, Any, Annotated6 7from typing_extensions import TypedDict8from langchain.schema import Document9from langchain_core.prompts import ChatPromptTemplate10from langchain_core.prompts.prompt import PromptTemplate11from langchain_core.messages import HumanMessage, AIMessage, AnyMessage12from langgraph.graph import END, StateGraph, MessagesState, START13from langgraph.graph.message import add_messages14from huggingface_hub import InferenceClient15 16from dotenv import load_dotenv17load_dotenv(verbose=True)18assert os.getenv("PINECONE_API_KEY") is not None19assert os.getenv("HUGGINGFACEHUB_EMBEDDINGS_MODEL") is not None20assert os.getenv("TAVILY_API_KEY") is not None21 22 23logger = logging.getLogger(__name__)        # Child logger for this module24logger.setLevel(logging.INFO)25logger.info(f"""correctiveRag.py:Config:26GROQ_MODEL = {os.getenv('GROQ_MODEL')}27HUGGINGFACEHUB_EMBEDDINGS_MODEL = {os.getenv('HUGGINGFACEHUB_EMBEDDINGS_MODEL')}28PINECONE_API_KEY = {os.getenv("PINECONE_API_KEY")[:5]}29""")30 31# Prepare the LLM32from langchain_groq import ChatGroq33assert os.getenv('GROQ_MODEL') is not None, "GROQ_MODEL not set"34assert os.getenv('GROQ_API_KEY') is not None, "GROQ_API_KEY not set"35llm = ChatGroq(model_name=os.getenv('GROQ_MODEL'), temperature=0, verbose=True)36 37# For using Grok38# from langchain_openai import ChatOpenAI39# assert os.getenv('XAI_API_KEY') is not None, "XAI_API_KEY not set"40# assert os.getenv('XAI_MODEL') is not None, "XAI_MODEL not set"41# assert os.getenv('XAI_BASE_URL') is not None, "XAI_BASE_URL not set"42# llm = ChatOpenAI(43#     api_key=os.getenv("XAI_API_KEY"), 44#     base_url=os.getenv("XAI_BASE_URL"), 45#     model=os.getenv("XAI_MODEL"), 46#     temperature=0.1)47 48 49# from langchain_openai import ChatOpenAI50# assert os.getenv('OPENAI_MODEL_NAME') is not None, "GROQ_MODEL not set"51# llm = ChatOpenAI(model=os.getenv("OPENAI_MODEL_NAME"), temperature=0.1, verbose=True)52 53# Huggingface - Does not support structured_output54# llm = InferenceClient("HuggingFaceH4/zephyr-7b-beta")55 56 57# Prepare the retriever58from langchain_huggingface import HuggingFaceEmbeddings59from langchain_pinecone import PineconeVectorStore60index_name, namespace = 'courses', 'dsa'61 62# Simple RAG63# embeddings = HuggingFaceEmbeddings(model_name=os.getenv("HUGGINGFACEHUB_EMBEDDINGS_MODEL"))64# docsearch = PineconeVectorStore.from_existing_index(embedding=embeddings, index_name=index_name, namespace=namespace)65# retriever = docsearch.as_retriever(search_type="mmr", search_kwargs={ 'k': 5 })66 67# Large-Small RAG68def larger_from_nearby(vectorstore, doc: Document, range:int) -> Document:69    """70    Given a document, find the "parent" document as a range of chunks around the central chunk71    """72    filter0 = { "document" : doc.metadata['document'] }73    filter1 = { "chunk": { "$gte" : doc.metadata['chunk']-range } }74    filter2 = { "chunk": { "$lte" : doc.metadata['chunk']+range } }75    and_filter = { "$and" : [ filter0, filter1, filter2 ] }76    range_docs = vectorstore.similarity_search(query='', k=2*range+1, filter=and_filter)77    content = ''78    for doc in range_docs:79        content += doc.page_content80    full_document = Document(page_content=content, metadata=doc.metadata)81    return full_document82 83def larger_retriever(vectorstore, query:str, topK:int):84    RANGE=2 # -RANGE...+RANGE85    logger.info(f'larger_retriever: with RANGE={RANGE}')86    docs = vectorstore.similarity_search(query, k=topK)87    larger_documents = list(map(lambda d: larger_from_nearby(vectorstore, d, RANGE), docs))88    logger.info(f'larger_retriever: Found {len(larger_documents)} documents.')89    return larger_documents90 91embeddings = HuggingFaceEmbeddings(model_name=os.getenv("HUGGINGFACEHUB_EMBEDDINGS_MODEL"))92vectorstore = PineconeVectorStore.from_existing_index(embedding=embeddings, index_name=index_name, namespace=namespace)93# docs = larger_retriever(vectorstore, query, 5)94retriever = lambda query: larger_retriever(vectorstore, query, 5)   # TODO topK95 96 97# Classify question98class ClassifyQuestion(BaseModel):99    """Binary score to decide if need to retrieve documents from the vectorstore about data structures and algorithms.100    The binary_score is "yes" to indicate that document retrieval is needed, otherwise is "no"."""101    binary_score: str = Field(description="If the question is about data structures and algorithms answer `yes`, otherwise answer `no`")102    # justification: str = Field(description="Explained reasoning for giving the yes/no score")103# LLM with function call104structured_llm_grader = llm.with_structured_output(ClassifyQuestion)105# Prompt106system = """You are an expert at classifying user questions.107    If the question are specific about data structures and algorithms, then answer `yes` to indicate that document retrieval is needed.108    Otherwise, it is a question as a general question, answer `no`.109"""110grade_prompt = ChatPromptTemplate.from_messages(111    [112        ("system", system),113        ("human", "Question: {question}"),114    ]115)116retriever_grader = grade_prompt | structured_llm_grader117 118 119# Retrieval grader120class GradeDocuments(BaseModel):121    """Binary score for relevance check on retrieved documents."""122    binary_score: str = Field(123        description="Documents are relevant to the question, 'yes' or 'no'"124    )125# LLM with function call126structured_llm_grader = llm.with_structured_output(GradeDocuments)127# Prompt128system = """You are a grader assessing relevance of a retrieved document to a user question.129    If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant.130    It does not need to be a stringent test. The goal is to filter out erroneous retrievals.131    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""132grade_prompt = ChatPromptTemplate.from_messages(133    [134        ("system", system),135        ("human", "Retrieved document: \n\n {document} \n\n User question: {question}"),136    ]137)138retrieval_grader = grade_prompt | structured_llm_grader139 140 141# Create the RAG chain142from langchain import hub143from langchain_core.output_parsers import StrOutputParser144# prompt = hub.pull("rlm/rag-prompt")145# print('----', prompt, '---')146template = """You are an assistant for question-answering tasks. 147Use the following pieces of retrieved context to answer the question. 148If you don't know the answer, just say that you don't know. 149Please keep the answer concise and to the point.150 151Context: {context} 152 153Question: {question}154 155Answer:156"""157prompt_template = PromptTemplate.from_template(template=template)158rag_chain = prompt_template | llm | StrOutputParser()159 160 161# Question rewriter162system = """You a question re-writer that converts an input question to a better version that is optimized163     for vectorstore retrieval. Look at the input and try to reason about the underlying semantic intent / meaning.164     Return only the re-written question. Do not return anything else.165     """166re_write_prompt = ChatPromptTemplate.from_messages(167    [168        ("system", system),169        ("human", "Here is the initial question: \n\n {question} \n Formulate an improved question."),170    ]171)172question_rewriter = re_write_prompt | llm | StrOutputParser()173 174 175# Web search tool176from langchain_community.tools.tavily_search import TavilySearchResults177web_search_tool = TavilySearchResults(k=3)178 179 180# Define the workflow Graph181 182class GraphState(TypedDict):183    """184    Represents the state of our graph.185 186    Attributes:187        messages: conversation history188        generation: LLM generation189        web_search: whether to add search190        documents: list of documents191        question: the last user question192    """193    messages: Annotated[list[AnyMessage], add_messages]194    generation: str195    web_search: str196    documents: List[str]197    question: str198 199def chatbot(state: GraphState):200    logger.info("---GENERATE (no context)---")201    logger.info(state)202    chain = llm | StrOutputParser()203    generation = chain.invoke(state["messages"])204    logger.info(generation)205    return { "messages": [AIMessage(content=generation)], "generation": generation }206 207def retrieve(state):208    """209    Retrieve documents210 211    Args:212        state (dict): The current graph state213 214    Returns:215        state (dict): New key added to state, documents, that contains retrieved documents216    """217    logger.info("---RETRIEVE---")218    question = state['messages'][-1].content     # Last Human message219    logger.info(f'question: {question}')220    # Retrieval221    # documents = retriever.invoke(question)222    documents = retriever(question) # Large-small retriever223    # logger.debug(documents)224    logger.info([ (doc.metadata['id'], doc.page_content[:20])for doc in documents ])225    return {"documents": documents, "question": question}226 227 228def generate_with_context(state):229    """230    Generate answer231 232    Args:233        state (dict): The current graph state234 235    Returns:236        state (dict): New key added to state, generation, that contains LLM generation237    """238    logger.debug("---GENERATE WITH CONTEXT---")239    logger.debug(f'state: {state}')240    question = state["question"]241    documents = state["documents"]242    # RAG generation243    generation = rag_chain.invoke({"context": documents, "question": question})244    logger.debug(generation)245    return {"documents": documents, "question": question, "generation": generation}246 247 248def web_search(state):249    """250    Web search based on the re-phrased question.251 252    Args:253        state (dict): The current graph state254 255    Returns:256        state (dict): Updates documents key with appended web results257    """258    logger.debug("---WEB SEARCH---")259    question = state["question"]260    documents = state["documents"]261    # Web search262    logger.debug(f'question: {question}')263    docs = web_search_tool.invoke({"query": question})  # Returns str if error264    logger.debug(f'type(docs) = {type(docs)}')265    logger.debug(docs)266    web_results = "\n".join([d["content"] for d in docs])267    web_results = Document(page_content=web_results)268    documents.append(web_results)269    return {"documents": web_results, "question": question}270 271 272def grade_documents(state):273    """274    Determines whether the retrieved documents are relevant to the question.275 276    Args:277        state (dict): The current graph state278 279    Returns:280        state (dict): Updates documents key with only filtered relevant documents281    """282    logger.debug("---CHECK DOCUMENT RELEVANCE TO QUESTION---")283    question = state["question"]284    documents = state["documents"]285    # Score each doc286    filtered_docs = []287    web_search = "No"288    for d in documents:289        score = retrieval_grader.invoke({"question": question, "document": d.page_content})290        grade = score.binary_score291        if grade == "yes":292            logger.debug("---GRADE: DOCUMENT RELEVANT---")293            filtered_docs.append(d)294        else:295            logger.debug("---GRADE: DOCUMENT NOT RELEVANT---")296            web_search = "Yes"297            continue298    return {"documents": filtered_docs, "question": question, "web_search": web_search}299 300 301def transform_query(state):302    """303    Transform the query to produce a better question.304 305    Args:306        state (dict): The current graph state307 308    Returns:309        state (dict): Updates question key with a re-phrased question310    """311    logger.debug("---TRANSFORM QUERY---")312    question = state["question"]313    documents = state["documents"]314    # Re-write question315    better_question = question_rewriter.invoke({"question": question})316    return {"documents": documents, "question": better_question}317 318 319 320### Edges ###321 322# For conditional edges323def decide_to_retrieve(state):324    """325    Determines whether to retrieve a context for answering a question.326 327    Args:328        state (dict): The current graph state329 330    Returns:331        str: Binary decision for next node to call332    """333    logger.debug("---ASSESS NEED FOR RETRIEVAL---")334    # logger.debug(state)335    question = state['messages'][-1].content     # Last Human message336    logger.debug(question)337    response = retriever_grader.invoke({ 'question': question })338    logger.debug(response)339    logger.debug(response.binary_score)340 341    if response.binary_score == "yes":342        # All documents have been filtered check_relevance343        # We will re-generate a new query344        logger.debug("---DECISION: RETRIEVE DOCUMENTS---")345        return "retrieve"346    else:347        # We have relevant documents, so generate answer348        logger.debug("---DECISION: GENERAL QUESTION, NO RETRIEVAL---")349        # state['question'] = question350        return "chatbot"351 352 353def decide_to_generate(state):354    """355    Determines whether to generate an answer, or re-generate a question.356 357    Args:358        state (dict): The current graph state359 360    Returns:361        str: Binary decision for next node to call362    """363    logger.debug("---ASSESS GRADED DOCUMENTS---")364    state["question"]365    web_search = state["web_search"]366    state["documents"]367 368    if web_search == "Yes":369        # All documents have been filtered check_relevance370        # We will re-generate a new query371        logger.debug(372            "---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---"373        )374        return "transform_query"375    else:376        # We have relevant documents, so generate answer377        logger.debug("---DECISION: GENERATE---")378        return "generate_with_context"379 380 381# Prepare and compile the Graph382workflow = StateGraph(GraphState)383 384# Define the nodes385workflow.add_node("chatbot", chatbot)                 # retrieve386workflow.add_node("retrieve", retrieve)                 # retrieve387workflow.add_node("grade_documents", grade_documents)   # grade documents388workflow.add_node("generate_with_context", generate_with_context) # generate389workflow.add_node("transform_query", transform_query)   # transform_query390workflow.add_node("web_search_node", web_search)        # web search391 392# Build graph393# workflow.add_edge(START, "retrieve")394workflow.add_conditional_edges(395    START,396    decide_to_retrieve                           397)398workflow.add_edge("retrieve", "grade_documents")399workflow.add_conditional_edges(400    "grade_documents",401    decide_to_generate,402    {403        "transform_query": "transform_query",404        "generate_with_context": "generate_with_context",405    },406)407workflow.add_edge("transform_query", "web_search_node")408workflow.add_edge("web_search_node", "generate_with_context")409workflow.add_edge("generate_with_context", "chatbot")410workflow.add_edge("chatbot", END)411 412# Compile413from langgraph.checkpoint.memory import MemorySaver414memory = MemorySaver()415app = workflow.compile(checkpointer=memory, debug=False)416 417 418if __name__ == "__main__":419    # Use the graph420    from pprint import pprint421 422    # print(retriever.invoke("What is an algorithm?"))423 424    config = {"configurable": {"thread_id": "abc123"}}425 426 427    def query_graph(question:str):428        inputs = { "question": question }429        messages = [HumanMessage(inputs['question'])]430        response = app.invoke({"messages": messages}, config)431        # print('TYPE >>', type(response))    # langgraph.pregel.io.AddableValuesDict432        return response433 434    def print_generation(response:str):435        # pprint(type(response['generation']))    # str436        # pprint(response['messages'])437        pprint(response['generation'])      # AIMessage (no context)438 439    # question = "Cual es el orden de ingreso y egresos de elementos en un Queue?"440    # pprint(query_graph(question))441 442    # print_generation(query_graph("Hi, my name is George and I would like to learn about algorithms"))443    # print_generation(query_graph("Do you remember my name? What algorithm would you use to reverse the letters in my name?"))444    # print_generation(query_graph("Que es un algoritmo?"))445    # print_generation(query_graph("Que es una heuristica?"))446    # print_generation(query_graph("Que se entiende por orden de crecimiento de un algoritmo?"))447    # print_generation(query_graph("Que es la función tilde?"))448    # print_generation(query_graph("Cuál es la diferencia entre función tilde y orden de crecimiento?"))449 450    # In stream mode, returns the full 'chatbot' message451    for x in app.stream({"messages": "What is the answer to the question of everything?"}, config):452        print(x)453