modified app.py to contain graph

AgenticRAG/agentic_rag.py

@@ -0,0 +1,225 @@
+from langchain.tools.retriever import create_retriever_tool
+from typing import Annotated, Literal, Sequence, TypedDict
+from typing import Annotated, Sequence, TypedDict
+from langchain_core.messages import BaseMessage
+from langgraph.graph.message import add_messages
+from langchain import hub
+from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import PromptTemplate
+from langchain_openai import ChatOpenAI
+from pydantic import BaseModel, Field
+from langgraph.prebuilt import tools_condition
+
+from aimakerspace.vectordatabase import VectorDatabase
+
+
+
+retriever_tool = create_retriever_tool(
+    retriever,
+    "retrieve_blog_posts",
+    "Search and return information about the responsible and ethical use of AI along with the development of policies and practices to protect civil rights and promote democratic values in the building, deployment, and government of automated systems.",
+)
+
+tools = [retriever_tool]
+
+
+
+
+class AgentState(TypedDict):
+    # The add_messages function defines how an update should be processed
+    # Default is to replace. add_messages says "append"
+    messages: Annotated[Sequence[BaseMessage], add_messages]
+
+
+
+
+### Edges
+
+
+def grade_documents(state) -> Literal["generate", "rewrite"]:
+    """
+    Determines whether the retrieved documents are relevant to the question.
+
+    Args:
+        state (messages): The current state
+
+    Returns:
+        str: A decision for whether the documents are relevant or not
+    """
+
+    # Data model
+    class grade(BaseModel):
+        """Binary score for relevance check."""
+
+        binary_score: str = Field(description="Relevance score 'yes' or 'no'")
+
+    # LLM
+    model = ChatOpenAI(temperature=0, model="gpt-4o-mini", streaming=True)
+
+    # LLM with tool and validation
+    llm_with_tool = model.with_structured_output(grade)
+
+    # Prompt
+    prompt = PromptTemplate(
+        template="""You are a grader assessing relevance of a retrieved document to a user question. \n 
+        Here is the retrieved document: \n\n {context} \n\n
+        Here is the user question: {question} \n
+        If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n
+        Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.""",
+        input_variables=["context", "question"],
+    )
+
+    # Chain
+    chain = prompt | llm_with_tool
+
+    messages = state["messages"]
+    last_message = messages[-1]
+
+    question = messages[0].content
+    docs = last_message.content
+
+    scored_result = chain.invoke({"question": question, "context": docs})
+
+    score = scored_result.binary_score
+
+    if score == "yes":
+        print("---DECISION: DOCS RELEVANT---")
+        return "generate"
+
+    else:
+        print("---DECISION: DOCS NOT RELEVANT---")
+        print(score)
+        return "rewrite"
+
+### Nodes
+
+
+def agent(state):
+    """
+    Invokes the agent model to generate a response based on the current state. Given
+    the question, it will decide to retrieve using the retriever tool, or simply end.
+
+    Args:
+        state (messages): The current state
+
+    Returns:
+        dict: The updated state with the agent response appended to messages
+    """
+    print("---CALL AGENT---")
+    messages = state["messages"]
+    model = ChatOpenAI(temperature=0, streaming=True, model="gpt-4o-mini")
+    model = model.bind_tools(tools)
+    response = model.invoke(messages)
+    # We return a list, because this will get added to the existing list
+    return {"messages": [response]}
+
+
+def rewrite(state):
+    """
+    Transform the query to produce a better question.
+
+    Args:
+        state (messages): The current state
+
+    Returns:
+        dict: The updated state with re-phrased question
+    """
+
+    print("---TRANSFORM QUERY---")
+    messages = state["messages"]
+    question = messages[0].content
+
+    msg = [
+        HumanMessage(
+            content=f""" \n 
+    Look at the input and try to reason about the underlying semantic intent / meaning. \n 
+    Here is the initial question:
+    \n ------- \n
+    {question} 
+    \n ------- \n
+    Formulate an improved question: """,
+        )
+    ]
+
+    # Grader
+    model = ChatOpenAI(temperature=0, model="gpt-4o-mini", streaming=True)
+    response = model.invoke(msg)
+    return {"messages": [response]}
+
+
+def generate(state):
+    """
+    Generate answer
+
+    Args:
+        state (messages): The current state
+
+    Returns:
+         dict: The updated state with re-phrased question
+    """
+    print("---GENERATE---")
+    messages = state["messages"]
+    question = messages[0].content
+    last_message = messages[-1]
+
+    docs = last_message.content
+
+    # Prompt
+    prompt = hub.pull("rlm/rag-prompt")
+
+    # LLM
+    llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0, streaming=True)
+
+    # Post-processing
+    def format_docs(docs):
+        return "\n\n".join(doc.page_content for doc in docs)
+
+    # Chain
+    rag_chain = prompt | llm | StrOutputParser()
+
+    # Run
+    response = rag_chain.invoke({"context": docs, "question": question})
+    return {"messages": [response]}
+
+from langgraph.graph import END, StateGraph, START
+from langgraph.prebuilt import ToolNode
+
+# Define a new graph
+workflow = StateGraph(AgentState)
+
+# Define the nodes we will cycle between
+workflow.add_node("agent", agent)  # agent
+retrieve = ToolNode([retriever_tool])
+workflow.add_node("retrieve", retrieve)  # retrieval
+workflow.add_node("rewrite", rewrite)  # Re-writing the question
+workflow.add_node(
+    "generate", generate
+)  # Generating a response after we know the documents are relevant
+# Call agent node to decide to retrieve or not
+workflow.add_edge(START, "agent")
+
+# Decide whether to retrieve
+workflow.add_conditional_edges(
+    "agent",
+    # Assess agent decision
+    tools_condition,
+    {
+        # Translate the condition outputs to nodes in our graph
+        "tools": "retrieve",
+        END: END,
+    },
+)
+
+# Edges taken after the `action` node is called.
+workflow.add_conditional_edges(
+    "retrieve",
+    # Assess agent decision
+    grade_documents,
+)
+workflow.add_edge("generate", END)
+workflow.add_edge("rewrite", "agent")
+
+# Compile
+graph = workflow.compile()
+

AgenticRAG/aimakerspace/openai_utils/embedding.py

@@ -4,8 +4,10 @@ import openai
 from typing import List
 import os
 import asyncio
+from dotenv import load_dotenv
 
-
+load_dotenv()
+print(os.getenv('OPENAI_API_KEY'))
 class EmbeddingModel:
     def __init__(self, embeddings_model_name: str = "text-embedding-3-small"):
         load_dotenv()

AgenticRAG/app.py

@@ -14,6 +14,10 @@ import chainlit as cl
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 # from langchain_experimental.text_splitter import SemanticChunker
 # from langchain_openai.embeddings import OpenAIEmbeddings
+import importlib
+
+
+
 
 system_template = """\
 Use the following context to answer a users question. If you cannot find the answer in the context, say you don't know the answer."""
@@ -27,20 +31,18 @@ Question:
 """
 user_role_prompt = UserRolePrompt(user_prompt_template)
 
-class RetrievalAugmentedQAPipeline:
-    def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase) -> None:
-        self.llm = llm
+class AgenticRAGPipeline:
+    def __init__(self, graph: StateGraph, vector_db_retriever: VectorDatabase) -> None:
+        self.graph = graph
         self.vector_db_retriever = vector_db_retriever
 
-    async def arun_pipeline(self, user_query: str):
-        context_list = self.vector_db_retriever.search_by_text(user_query, k=4)
+    async def run_pipeline(self, user_query: str):
+        state = self.graph.execute({"text": user_query, "chunk_size": 100})
+        context_list = state["retriever"]
 
-        context_prompt = ""
-        for context in context_list:
-            context_prompt += context[0] + "\n"
+        context_prompt = "\n".join(context_list)
 
         formatted_system_prompt = system_role_prompt.create_message()
-
         formatted_user_prompt = user_role_prompt.create_message(question=user_query, context=context_prompt)
 
         async def generate_response():
@@ -49,12 +51,12 @@ class RetrievalAugmentedQAPipeline:
 
         return {"response": generate_response(), "context": context_list}
 
+
+
+
 text_splitter = RecursiveCharacterTextSplitter()
-# try:
-#     api_key = os.environ["OPENAI_API_KEY"]
-# except KeyError:
-#     print("Environment variable OPENAI_API_KEY not found")
-# text_splitter = SemanticChunker(OpenAIEmbeddings(api_key=api_key), breakpoint_threshold_type="standard_deviation")
+
+
 
 def process_text_file(file: AskFileResponse):
     import tempfile
@@ -75,11 +77,9 @@ def process_text_file(file: AskFileResponse):
     else:
         raise ValueError("Provide a .txt or .pdf file")
     texts = [x.page_content for x in text_splitter.transform_documents(documents)]
-    # texts = [x.page_content for x in text_splitter.split_documents(documents)]
     return texts
 
 
-
 @cl.on_chat_start
 async def on_chat_start():
     files = None
@@ -111,11 +111,234 @@ async def on_chat_start():
     
     chat_openai = ChatOpenAI()
 
-    # Create a chain
-    retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
-        vector_db_retriever=vector_db,
-        llm=chat_openai
+    retriever = vector_db
+    """Graph code here"""
+
+    from langchain.tools.retriever import create_retriever_tool
+    from typing import Annotated, Literal, Sequence, TypedDict
+    from typing import Annotated, Sequence, TypedDict
+    from langchain_core.messages import BaseMessage
+    from langgraph.graph.message import add_messages
+    from langchain import hub
+    from langchain_core.messages import BaseMessage, HumanMessage
+    from langchain_core.output_parsers import StrOutputParser
+    from langchain_core.prompts import PromptTemplate
+    from langchain_openai import ChatOpenAI
+    from pydantic import BaseModel, Field
+    from langgraph.prebuilt import tools_condition
+
+    from aimakerspace.vectordatabase import VectorDatabase
+
+    retriever_tool = create_retriever_tool(
+        retriever,
+        "retrieve_blog_posts",
+        "Search and return information about the responsible and ethical use of AI along with the development of policies and practices to protect civil rights and promote democratic values in the building, deployment, and government of automated systems.",
+    )
+
+    tools = [retriever_tool]
+
+
+
+    class AgentState(TypedDict):
+        # The add_messages function defines how an update should be processed
+        # Default is to replace. add_messages says "append"
+        messages: Annotated[Sequence[BaseMessage], add_messages]
+
+
+
+
+    ### Edges
+
+
+    def grade_documents(state) -> Literal["generate", "rewrite"]:
+        """
+        Determines whether the retrieved documents are relevant to the question.
+
+        Args:
+            state (messages): The current state
+
+        Returns:
+            str: A decision for whether the documents are relevant or not
+        """
+
+        # Data model
+        class grade(BaseModel):
+            """Binary score for relevance check."""
+
+            binary_score: str = Field(description="Relevance score 'yes' or 'no'")
+
+        # LLM
+        model = ChatOpenAI(temperature=0, model="gpt-4o-mini", streaming=True)
+
+        # LLM with tool and validation
+        llm_with_tool = model.with_structured_output(grade)
+
+        # Prompt
+        prompt = PromptTemplate(
+            template="""You are a grader assessing relevance of a retrieved document to a user question. \n 
+            Here is the retrieved document: \n\n {context} \n\n
+            Here is the user question: {question} \n
+            If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n
+            Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.""",
+            input_variables=["context", "question"],
+        )
+
+        # Chain
+        chain = prompt | llm_with_tool
+
+        messages = state["messages"]
+        last_message = messages[-1]
+
+        question = messages[0].content
+        docs = last_message.content
+
+        scored_result = chain.invoke({"question": question, "context": docs})
+
+        score = scored_result.binary_score
+
+        if score == "yes":
+            print("---DECISION: DOCS RELEVANT---")
+            return "generate"
+
+        else:
+            print("---DECISION: DOCS NOT RELEVANT---")
+            print(score)
+            return "rewrite"
+
+    ### Nodes
+
+
+    def agent(state):
+        """
+        Invokes the agent model to generate a response based on the current state. Given
+        the question, it will decide to retrieve using the retriever tool, or simply end.
+
+        Args:
+            state (messages): The current state
+
+        Returns:
+            dict: The updated state with the agent response appended to messages
+        """
+        print("---CALL AGENT---")
+        messages = state["messages"]
+        model = ChatOpenAI(temperature=0, streaming=True, model="gpt-4o-mini")
+        model = model.bind_tools(tools)
+        response = model.invoke(messages)
+        # We return a list, because this will get added to the existing list
+        return {"messages": [response]}
+
+
+    def rewrite(state):
+        """
+        Transform the query to produce a better question.
+
+        Args:
+            state (messages): The current state
+
+        Returns:
+            dict: The updated state with re-phrased question
+        """
+
+        print("---TRANSFORM QUERY---")
+        messages = state["messages"]
+        question = messages[0].content
+
+        msg = [
+            HumanMessage(
+                content=f""" \n 
+        Look at the input and try to reason about the underlying semantic intent / meaning. \n 
+        Here is the initial question:
+        \n ------- \n
+        {question} 
+        \n ------- \n
+        Formulate an improved question: """,
+            )
+        ]
+
+        # Grader
+        model = ChatOpenAI(temperature=0, model="gpt-4o-mini", streaming=True)
+        response = model.invoke(msg)
+        return {"messages": [response]}
+
+
+    def generate(state):
+        """
+        Generate answer
+
+        Args:
+            state (messages): The current state
+
+        Returns:
+            dict: The updated state with re-phrased question
+        """
+        print("---GENERATE---")
+        messages = state["messages"]
+        question = messages[0].content
+        last_message = messages[-1]
+
+        docs = last_message.content
+
+        # Prompt
+        prompt = hub.pull("rlm/rag-prompt")
+
+        # LLM
+        llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0, streaming=True)
+
+        # Post-processing
+        def format_docs(docs):
+            return "\n\n".join(doc.page_content for doc in docs)
+
+        # Chain
+        rag_chain = prompt | llm | StrOutputParser()
+
+        # Run
+        response = rag_chain.invoke({"context": docs, "question": question})
+        return {"messages": [response]}
+
+    from langgraph.graph import END, StateGraph, START
+    from langgraph.prebuilt import ToolNode
+
+    # Define a new graph
+    workflow = StateGraph(AgentState)
+
+    # Define the nodes we will cycle between
+    workflow.add_node("agent", agent)  # agent
+    retrieve = ToolNode([retriever_tool])
+    workflow.add_node("retrieve", retrieve)  # retrieval
+    workflow.add_node("rewrite", rewrite)  # Re-writing the question
+    workflow.add_node(
+        "generate", generate
+    )  # Generating a response after we know the documents are relevant
+    # Call agent node to decide to retrieve or not
+    workflow.add_edge(START, "agent")
+
+    # Decide whether to retrieve
+    workflow.add_conditional_edges(
+        "agent",
+        # Assess agent decision
+        tools_condition,
+        {
+            # Translate the condition outputs to nodes in our graph
+            "tools": "retrieve",
+            END: END,
+        },
+    )
+
+    # Edges taken after the `action` node is called.
+    workflow.add_conditional_edges(
+        "retrieve",
+        # Assess agent decision
+        grade_documents,
     )
+    workflow.add_edge("generate", END)
+    workflow.add_edge("rewrite", "agent")
+
+    # Compile
+    graph = workflow.compile()
+
+    """END GRAPH CODE"""
+    # Create a chain
+    retrieval_augmented_qa_pipeline = AgenticRAGPipeline(graph=graph, vector_db_retriever=vector_db)
     
     # Let the user know that the system is ready
     msg.content = f"Processing `{file.name}` done. You can now ask questions!"

AgenticRAG/previous_app.py

@@ -0,0 +1,137 @@
+import os
+from typing import List
+from chainlit.types import AskFileResponse
+from aimakerspace.text_utils import CharacterTextSplitter, TextFileLoader
+from aimakerspace.openai_utils.prompts import (
+    UserRolePrompt,
+    SystemRolePrompt,
+    AssistantRolePrompt,
+)
+from aimakerspace.openai_utils.embedding import EmbeddingModel
+from aimakerspace.vectordatabase import VectorDatabase
+from aimakerspace.openai_utils.chatmodel import ChatOpenAI
+import chainlit as cl
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+# from langchain_experimental.text_splitter import SemanticChunker
+# from langchain_openai.embeddings import OpenAIEmbeddings
+
+system_template = """\
+Use the following context to answer a users question. If you cannot find the answer in the context, say you don't know the answer."""
+system_role_prompt = SystemRolePrompt(system_template)
+
+user_prompt_template = """\
+Context:
+{context}
+Question:
+{question}
+"""
+user_role_prompt = UserRolePrompt(user_prompt_template)
+
+class RetrievalAugmentedQAPipeline:
+    def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase) -> None:
+        self.llm = llm
+        self.vector_db_retriever = vector_db_retriever
+
+    async def arun_pipeline(self, user_query: str):
+        context_list = self.vector_db_retriever.search_by_text(user_query, k=4)
+
+        context_prompt = ""
+        for context in context_list:
+            context_prompt += context[0] + "\n"
+
+        formatted_system_prompt = system_role_prompt.create_message()
+
+        formatted_user_prompt = user_role_prompt.create_message(question=user_query, context=context_prompt)
+
+        async def generate_response():
+            async for chunk in self.llm.astream([formatted_system_prompt, formatted_user_prompt]):
+                yield chunk
+
+        return {"response": generate_response(), "context": context_list}
+
+text_splitter = RecursiveCharacterTextSplitter()
+# try:
+#     api_key = os.environ["OPENAI_API_KEY"]
+# except KeyError:
+#     print("Environment variable OPENAI_API_KEY not found")
+# text_splitter = SemanticChunker(OpenAIEmbeddings(api_key=api_key), breakpoint_threshold_type="standard_deviation")
+
+def process_text_file(file: AskFileResponse):
+    import tempfile
+    from langchain_community.document_loaders.pdf import PyPDFLoader
+
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=file.name) as temp_file:
+        temp_file_path = temp_file.name
+
+    with open(temp_file_path, "wb") as f:
+        f.write(file.content)
+
+    if file.type == 'text/plain':
+        text_loader = TextFileLoader(temp_file_path)
+        documents = text_loader.load_documents()
+    elif file.type == 'application/pdf':
+        pdf_loader = PyPDFLoader(temp_file_path)
+        documents = pdf_loader.load()
+    else:
+        raise ValueError("Provide a .txt or .pdf file")
+    texts = [x.page_content for x in text_splitter.transform_documents(documents)]
+    # texts = [x.page_content for x in text_splitter.split_documents(documents)]
+    return texts
+
+
+
+@cl.on_chat_start
+async def on_chat_start():
+    files = None
+
+    # Wait for the user to upload a file
+    while files == None:
+        files = await cl.AskFileMessage(
+            content="Please upload a Text file or a PDF to begin!",
+            accept=["text/plain", "application/pdf"],
+            max_size_mb=12,
+            timeout=180,
+        ).send()
+
+    file = files[0]
+
+    msg = cl.Message(
+        content=f"Processing `{file.name}`...", disable_human_feedback=True
+    )
+    await msg.send()
+
+    # load the file
+    texts = process_text_file(file)
+
+    print(f"Processing {len(texts)} text chunks")
+
+    # Create a dict vector store
+    vector_db = VectorDatabase()
+    vector_db = await vector_db.abuild_from_list(texts)
+    
+    chat_openai = ChatOpenAI()
+
+    # Create a chain
+    retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
+        vector_db_retriever=vector_db,
+        llm=chat_openai
+    )
+    
+    # Let the user know that the system is ready
+    msg.content = f"Processing `{file.name}` done. You can now ask questions!"
+    await msg.update()
+
+    cl.user_session.set("chain", retrieval_augmented_qa_pipeline)
+
+
+@cl.on_message
+async def main(message):
+    chain = cl.user_session.get("chain")
+
+    msg = cl.Message(content="")
+    result = await chain.arun_pipeline(message.content)
+
+    async for stream_resp in result["response"]:
+        await msg.stream_token(stream_resp)
+
+    await msg.send()

AgenticRAG/requirements.txt

@@ -4,4 +4,9 @@ openai
 langchain_community
 langchain_experimental
 langchain_openai
-pypdf
+pypdf
+tiktoken 
+langchainhub 
+langchain 
+langgraph 
+langchain-text-splitters