File size: 19,426 Bytes
e619a67
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
import os
import streamlit as st
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from pprint import pprint
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain import hub
from langchain_core.output_parsers import StrOutputParser
from typing import List
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, END

# Streamlit setup with new theme and typography
st.set_page_config(page_title="SELF-RAG Workflow Application", page_icon="🤖", layout="centered")
st.markdown(
    """

    <style>

    .main {

        background-color: #272727;

        font-family: 'Helvetica Neue', sans-serif;

    }

    .sidebar .sidebar-content {

        background-color: #2E3944;

        color: #ffffff;

    }

    h1 {

        color: #14A76C;

    }

    .stTextInput {

        border: 1px solid #272727;

        border-radius: 5px;

    }

    </style>

    """,
    unsafe_allow_html=True
)

# Sidebar with instructions and API key input
st.sidebar.title("Instructions")
st.sidebar.write("""

1. Enter your OpenAI API Key.

2. Enter your question in the text box.

3. Provide URLs for the documents you want to use.

4. Click on the 'Run Workflow' button.

5. View the results below.

""")
api_key = st.sidebar.text_input("Enter your OpenAI API Key:", type="password")

st.title("SELF-RAG Workflow Application")
input_text = st.text_input("Enter your question : ")
urls_input = st.text_area("Enter URLs (one per line) :")
urls = [url.strip() for url in urls_input.split('\n') if url.strip()]
inputs = {"question": input_text, "transform_attempts": 0}

if st.button("Run Workflow"):
    if not api_key:
        st.error("Please enter your OpenAI API Key.")
    elif not urls:
        st.error("Please provide at least one URL.")
    elif not input_text:
        st.error("Please enter a question.")
    else:
        # Document loading and processing
        try:
            texts = []
            docs = []
            for url in urls:
                try:
                    docs.extend(WebBaseLoader(url).load())
                except Exception as e:
                    st.error(f"Error loading document from {url}: {e}")
            if not docs:
                st.error("No documents loaded. Please check the URLs.")
            else:
                text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
                    chunk_size=250, chunk_overlap=0
                )
                doc_splits = text_splitter.split_documents(docs)

                # Add to vectorDB
                vectorstore = FAISS.from_documents(
                    documents=doc_splits,
                    embedding=OpenAIEmbeddings(openai_api_key=api_key),
                )
                retriever = vectorstore.as_retriever()

                ### Retrieval Grader
                # 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
                llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0, openai_api_key=api_key)
                structured_llm_grader = llm.with_structured_output(GradeDocuments)

                # Prompt
                system = """You are a grader assessing relevance of a retrieved document to a user question. \n 

                It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \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."""
                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
                question = input_text
                docs = retriever.get_relevant_documents(question)
                if not docs:
                    st.error("No relevant documents found for the question.")
                else:
                    doc_txt = docs[1].page_content

                    ### Generate
                    # Prompt
                    prompt = hub.pull("rlm/rag-prompt")

                    # LLM
                    llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0, openai_api_key=api_key)

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

                    # Chain
                    rag_chain = prompt | llm | StrOutputParser()

                    # Run
                    generation = rag_chain.invoke({"context": docs, "question": question})

                    ### Hallucination Grader
                    # Data model
                    class GradeHallucinations(BaseModel):
                        """Binary score for hallucination present in generation answer."""
                        binary_score: str = Field(description="Answer is grounded in the facts, 'yes' or 'no'")

                    # LLM with function call
                    llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0, openai_api_key=api_key)
                    structured_llm_grader = llm.with_structured_output(GradeHallucinations)

                    # Prompt
                    system = """You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. \n 

                    Give a binary score 'yes' or 'no'. 'Yes' means that the answer is grounded in / supported by the set of facts."""
                    hallucination_prompt = ChatPromptTemplate.from_messages(
                        [
                            ("system", system),
                            ("human", "Set of facts: \n\n {documents} \n\n LLM generation: {generation}"),
                        ]
                    )

                    hallucination_grader = hallucination_prompt | structured_llm_grader

                    ### Answer Grader
                    # Data model
                    class GradeAnswer(BaseModel):
                        """Binary score to assess answer addresses question."""
                        binary_score: str = Field(description="Answer addresses the question, 'yes' or 'no'")

                    # LLM with function call
                    llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0, openai_api_key=api_key)
                    structured_llm_grader = llm.with_structured_output(GradeAnswer)

                    # Prompt
                    system = """You are a grader assessing whether an answer addresses / resolves a question \n 

                    Give a binary score 'yes' or 'no'. Yes' means that the answer resolves the question."""
                    answer_prompt = ChatPromptTemplate.from_messages(
                        [
                            ("system", system),
                            ("human", "User question: \n\n {question} \n\n LLM generation: {generation}"),
                        ]
                    )

                    answer_grader = answer_prompt | structured_llm_grader

                    ### Question Re-writer
                    # LLM
                    llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0, openai_api_key=api_key)

                    # Prompt
                    system = """You a question re-writer that converts an input question to a better version that is optimized \n 

                    for vectorstore retrieval. Look at the input and try to reason about the underlying semantic 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()

                    class GraphState(TypedDict):
                        """

                        Represents the state of our graph.



                        Attributes:

                            question: question

                            generation: LLM generation

                            documents: list of documents

                            transform_attempts: int

                        """
                        question: str
                        generation: str
                        documents: List[str]
                        transform_attempts: int

                    ### Nodes
                    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

                        """
                        texts.append("---RETRIEVE---")
                        question = state["question"]

                        # Retrieval
                        documents = retriever.get_relevant_documents(question)
                        return {"documents": documents, "question": question, "transform_attempts": state.get("transform_attempts", 0)}

                    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

                        """
                        texts.append("---GENERATE---")
                        question = state["question"]
                        documents = state["documents"]

                        # RAG generation
                        generation = rag_chain.invoke({"context": documents, "question": question})
                        return {"documents": documents, "question": question, "generation": generation, "transform_attempts": state.get("transform_attempts", 0)}

                    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

                        """
                        texts.append("---CHECK DOCUMENT RELEVANCE TO QUESTION---")
                        question = state["question"]
                        documents = state["documents"]

                        # Score each doc
                        filtered_docs = []
                        for d in documents:
                            score = retrieval_grader.invoke(
                                {"question": question, "document": d.page_content}
                            )
                            grade = score.binary_score
                            if grade == "yes":
                                texts.append("---GRADE: DOCUMENT RELEVANT---")
                                filtered_docs.append(d)
                            else:
                                texts.append("---GRADE: DOCUMENT NOT RELEVANT---")
                                continue
                        return {"documents": filtered_docs, "question": question, "transform_attempts": state.get("transform_attempts", 0)}

                    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

                        """
                        texts.append("---TRANSFORM QUERY---")
                        question = state["question"]
                        documents = state["documents"]

                        # Re-write question
                        better_question = question_rewriter.invoke({"question": question})
                        return {"documents": documents, "question": better_question, "transform_attempts": state.get("transform_attempts", 0) + 1}

                    ### 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

                        """
                        texts.append("---ASSESS GRADED DOCUMENTS---")
                        filtered_documents = state["documents"]

                        if not filtered_documents:
                            if state.get("transform_attempts", 0) >= 3:
                                return "conclude_no_answer"
                            else:
                                # All documents have been filtered check_relevance
                                # We will re-generate a new query
                                texts.append(
                                    "---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---"
                                )
                                return "transform_query"
                        else:
                            # We have relevant documents, so generate answer
                            texts.append("---DECISION: GENERATE---")
                            return "generate"

                    def grade_generation_v_documents_and_question(state):
                        """

                        Determines whether the generation is grounded in the document and answers question.



                        Args:

                            state (dict): The current graph state



                        Returns:

                            str: Decision for next node to call

                        """
                        texts.append("---CHECK HALLUCINATIONS---")
                        question = state["question"]
                        documents = state["documents"]
                        generation = state["generation"]

                        score = hallucination_grader.invoke(
                            {"documents": documents, "generation": generation}
                        )
                        grade = score.binary_score

                        # Check hallucination
                        if grade == "yes":
                            texts.append("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")
                            # Check question-answering
                            texts.append("---GRADE GENERATION vs QUESTION---")
                            score = answer_grader.invoke({"question": question, "generation": generation})
                            grade = score.binary_score
                            if grade == "yes":
                                texts.append("---DECISION: GENERATION ADDRESSES QUESTION---")
                                return "useful"
                            else:
                                texts.append("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")
                                return "not useful"
                        else:
                            texts.append("---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---")
                            return "not supported"

                    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)  # generate
                    workflow.add_node("transform_query", transform_query)  # transform_query
                    workflow.add_node("conclude_no_answer", lambda state: {"question": state["question"], "generation": "I don't know the answer since none of the given documents are relevant to the question.", "documents": [], "transform_attempts": state.get("transform_attempts", 0)})

                    # 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",
                            "conclude_no_answer": "conclude_no_answer"
                        },
                    )
                    workflow.add_edge("transform_query", "retrieve")
                    workflow.add_conditional_edges(
                        "generate",
                        grade_generation_v_documents_and_question,
                        {
                            "not supported": "generate",
                            "useful": END,
                            "not useful": "transform_query",
                        },
                    )

                    # Compile
                    app = workflow.compile()

                    try:
                        for output in app.stream(inputs):
                            for key, value in output.items():
                                for i in texts:
                                    st.write(i)
                                texts = []
                        # Final generation
                        st.write('## Final Answer')
                        st.write(value["generation"])
                    except Exception as e:
                        st.error(f"Error in workflow execution: {e}")
        except Exception as e:
            st.error(f"Error in document processing: {e}")