Create app.py

app.py

@@ -0,0 +1,161 @@
+import os
+import streamlit as st
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain.schema import Document
+from langchain_openai import ChatOpenAI
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langgraph.graph import StateGraph, END
+from graphviz import Digraph  # For workflow visualization
+from typing_extensions import TypedDict
+from typing import List
+from utils.build_rag import RAG
+
+
+# Fetch API Keys
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
+
+# Check for Missing API Keys
+if not OPENAI_API_KEY or not TAVILY_API_KEY:
+    st.error("❌ API keys missing! Please set `OPENAI_API_KEY` and `TAVILY_API_KEY` in your `.env` file.")
+    st.stop()  # Stop the app execution
+
+# Set up LLM and Tools
+llm = ChatOpenAI(model="gpt-4-1106-preview", openai_api_key=OPENAI_API_KEY)
+web_search_tool = TavilySearchResults(api_key=TAVILY_API_KEY, k=2)
+
+# Prompt templates
+def get_prompt():
+    template = """Answer the question based only on the following context:
+    {context}
+
+    Question: {question}
+    """
+    return ChatPromptTemplate.from_template(template)
+
+# Define Graph State
+class GraphState(TypedDict):
+    question: str
+    generation: str
+    web_search: str
+    documents: List[Document]
+
+# RAG Setup
+rag = RAG()
+retriever = rag.get_retriever()
+prompt = get_prompt()
+output_parser = StrOutputParser()
+
+# Nodes
+def retrieve(state):
+    question = state["question"]
+    documents = retriever.get_relevant_documents(question)
+    st.sidebar.write(f"Retrieved Documents: {len(documents)}")
+    return {"documents": documents, "question": question}
+
+def grade_documents(state):
+    question = state["question"]
+    documents = state["documents"]
+    filtered_docs = []
+    web_search = "No"
+
+    for doc in documents:
+        score = {"binary_score": "yes"}  # Dummy grader; integrate as needed
+        if score["binary_score"] == "yes":
+            filtered_docs.append(doc)
+        else:
+            web_search = "Yes"
+    st.sidebar.write(f"Document Grading Results: {len(filtered_docs)} relevant")
+    return {"documents": filtered_docs, "web_search": web_search, "question": question}
+
+def generate(state):
+    context = "\n".join([doc.page_content for doc in state["documents"]])
+    response = output_parser.parse(llm.invoke({"context": context, "question": state["question"]}).content)
+    return {"generation": response}
+
+def transform_query(state):
+    question = state["question"]
+    new_question = llm.invoke(f"Rewrite: {question}").content
+    st.sidebar.write(f"Rewritten Question: {new_question}")
+    return {"question": new_question}
+
+def web_search(state):
+    question = state["question"]
+    results = web_search_tool.invoke({"query": question})
+    docs = "\n".join([result["content"] for result in results])
+    state["documents"].append(Document(page_content=docs))
+    st.sidebar.write("Web Search Completed")
+    return {"documents": state["documents"], "question": question}
+
+def decide_to_generate(state):
+    return "generate" if state["web_search"] == "No" else "transform_query"
+
+# Build Graph
+workflow = StateGraph(GraphState)
+workflow.add_node("retrieve", retrieve)
+workflow.add_node("grade_documents", grade_documents)
+workflow.add_node("generate", generate)
+workflow.add_node("transform_query", transform_query)
+workflow.add_node("web_search_node", web_search)
+
+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)
+
+app = workflow.compile()
+
+# Visualize Workflows
+def plot_workflow():
+    graph = Digraph()
+    graph.attr(size='6,6')
+
+    # Add nodes
+    graph.node("retrieve", "Retrieve Documents")
+    graph.node("grade_documents", "Grade Documents")
+    graph.node("generate", "Generate Answer")
+    graph.node("transform_query", "Transform Query")
+    graph.node("web_search_node", "Web Search")
+    graph.node("END", "End")
+
+    # Add edges
+    graph.edge("retrieve", "grade_documents")
+    graph.edge("grade_documents", "generate", label="Relevant Docs")
+    graph.edge("grade_documents", "transform_query", label="No Relevant Docs")
+    graph.edge("transform_query", "web_search_node")
+    graph.edge("web_search_node", "generate")
+    graph.edge("generate", "END")
+
+    return graph
+
+# Streamlit App
+st.title("Self-Corrective RAG")
+st.write("### Compare RAG Pipeline Outputs (With and Without Self-Correction)")
+
+# Plot Workflow
+st.subheader("Workflow Visualization")
+st.graphviz_chart(plot_workflow().source)
+
+# User Input
+question = st.text_input("Enter your question:", "What is Llama2?")
+
+if st.button("Run Comparison"):
+    # Run Basic RAG
+    st.subheader("Without Self-Correction:")
+    docs = retriever.invoke(question)
+    basic_context = "\n".join([doc.page_content for doc in docs])
+    basic_response = output_parser.parse(llm.invoke({"context": basic_context, "question": question}).content)
+    st.write(basic_response)
+
+    # Run Self-Corrective RAG
+    st.subheader("With Self-Correction:")
+    inputs = {"question": question}
+    final_generation = ""
+    for output in app.stream(inputs):
+        for key, value in output.items():
+            if key == "generation":
+                final_generation = value
+    st.write(final_generation)