Spaces:
Sleeping
Sleeping
| import os | |
| import streamlit as st | |
| from vector_tool import ensemble_retriever | |
| from langgraph.prebuilt import ToolInvocation | |
| from langchain_core.messages import ToolMessage | |
| import json | |
| # Set up the tools to execute them from the graph | |
| from langgraph.prebuilt import ToolExecutor | |
| # tools retrieval | |
| from function_tools import tool_chain | |
| from vector_tool import ensemble_retriever | |
| os.environ['OPENAI_API_KEY'] = st.secrets["OPENAI_API_KEY"] | |
| os.environ['TAVILY_API_KEY'] = st.secrets["TAVILY_API_KEY"] | |
| ### Retrieval Grader | |
| from langchain_openai import ChatOpenAI | |
| from langchain_core.prompts import ChatPromptTemplate | |
| from langchain_core.pydantic_v1 import BaseModel, Field | |
| #LLM models | |
| llm_AI4 = ChatOpenAI(model="gpt-4-1106-preview", temperature=0) | |
| # 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 | |
| structured_llm_grader = llm_AI4.with_structured_output(GradeDocuments) | |
| # Prompt | |
| system = """You are a grader assessing relevance of a retrieved document to a user question. \n | |
| If the document contains keyword(s) or semantic meaning related to the 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 | |
| ### Generate | |
| from langchain import hub | |
| from langchain.prompts import MessagesPlaceholder | |
| from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser | |
| from langchain.prompts import MessagesPlaceholder | |
| from langchain.agents.format_scratchpad.openai_tools import ( | |
| format_to_openai_tool_messages | |
| ) | |
| from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage | |
| from langchain_core.output_parsers import StrOutputParser | |
| from typing import Any, List, Union | |
| # Prompt | |
| #prompt = hub.pull("rlm/rag-prompt") | |
| system_message = '''You are an AI assistant for answering questions about vedas and scriptures. | |
| \nYou are given the following extracted documents from Svarupa Knowledge Base (https://svarupa.org/) and other documents and a question. | |
| Provide a conversational answer.\nIf you are not provided with any documents, say \"I did not get any relevant context for this but | |
| I will reply to the best of my knowledge\" and then write your answer\nIf you don't know the answer, just say \"Hmm, I'm not sure. \" Don't try to make up an answer. | |
| \nIf the question is not about vedas and scriptures, politely inform them that you are tuned to only answer questions about that.\n\n''' | |
| ''' | |
| prompt = ChatPromptTemplate.from_messages( | |
| [ | |
| ("system",system_message), | |
| # Please note the ordering of the fields in the prompt! | |
| # The correct ordering is: | |
| # 1. history - the past messages between the user and the agent | |
| # 2. user - the user's current input | |
| # 3. agent_scratchpad - the agent's working space for thinking and | |
| # invoking tools to respond to the user's input. | |
| # If you change the ordering, the agent will not work correctly since | |
| # the messages will be shown to the underlying LLM in the wrong order. | |
| MessagesPlaceholder(variable_name="context"), | |
| ("user", "{question}"), | |
| ] | |
| ) | |
| ''' | |
| generate_prompt = ChatPromptTemplate.from_messages( | |
| [ | |
| ("system", system_message), | |
| ("human", "Here is the given context {context}, queation: {question} \n\n Formulate an answer."), | |
| ] | |
| ) | |
| # LLM | |
| llm_AI = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0) | |
| # Post-processing | |
| def format_docs(docs): | |
| return "\n\n".join(doc.page_content for doc in docs) | |
| # Chain | |
| rag_chain = generate_prompt | llm_AI4 | StrOutputParser() #OpenAIToolsAgentOutputParser() | |
| ####-----------------TESTING | |
| prompt = ChatPromptTemplate.from_messages( | |
| [ | |
| ( | |
| "system", | |
| "You are a helpful assistant. Answer all questions to the best of your ability.", | |
| ), | |
| MessagesPlaceholder(variable_name="chat_history"), | |
| ("human", "{question}"), | |
| ] | |
| ) | |
| from langchain_core.runnables.history import RunnableWithMessageHistory | |
| from langchain.memory import ChatMessageHistory | |
| chat_history_for_chain = ChatMessageHistory() | |
| chain_with_message_history = RunnableWithMessageHistory( | |
| rag_chain, | |
| lambda session_id: chat_history_for_chain, | |
| input_messages_key="question", | |
| history_messages_key="chat_history", | |
| ) | |
| ### Question Re-writer | |
| # LLM | |
| llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0) | |
| # Prompt | |
| system = """You a question re-writer that converts an input question to a better version that is optimized \n | |
| for a search. Look at the input and try to reason about the underlying sematic 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() | |
| ### Search | |
| from langchain_community.tools.tavily_search import TavilySearchResults | |
| web_search_tool = TavilySearchResults(k=2) | |
| from typing_extensions import TypedDict | |
| from typing import List | |
| from typing import TypedDict, Annotated, Sequence | |
| import operator | |
| from langchain_core.messages import BaseMessage | |
| class GraphState(TypedDict): | |
| """ | |
| Represents the state of our graph. | |
| Attributes: | |
| question: question | |
| generation: LLM generation | |
| web_search: whether to add search | |
| documents: list of documents | |
| """ | |
| question : str | |
| generation : str | |
| web_search : str | |
| messages: List[str] #Union[dict[str, Any]] | |
| from langchain.schema import Document | |
| 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 | |
| """ | |
| print("---VECTOR RETRIEVE---") | |
| question = state["question"] | |
| # Retrieval | |
| documents = ensemble_retriever.get_relevant_documents(question) | |
| #print(documents) | |
| # Iterate over each document and update the 'metadata' field with the file name | |
| for doc in documents: | |
| try: | |
| file_path = doc.metadata['source'] | |
| #print(file_path) | |
| file_name = os.path.split(file_path)[1] # Get the file name from the file path | |
| doc.metadata['source'] = file_name | |
| except KeyError: | |
| # Handle the case where 'source' field is missing in the metadata | |
| doc.metadata['source'] = 'unavailable' | |
| except Exception as e: | |
| # Handle any other exceptions that may occur | |
| print(f"An error occurred while processing document: {e}") | |
| return {"messages": documents, "question": question} | |
| 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 | |
| """ | |
| print("---GENERATE---") | |
| question = state["question"] | |
| messages = state["messages"] | |
| print(messages) | |
| # RAG generation | |
| generation = chain_with_message_history.invoke({"context": messages, "question": question},{"configurable": {"session_id": "unused"}}) | |
| return {"messages": messages, "question": question, "generation": generation} | |
| 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 | |
| """ | |
| print("---CHECK DOCUMENT RELEVANCE TO QUESTION---") | |
| question = state["question"] | |
| messages = state["messages"] | |
| # Score each doc | |
| filtered_docs = [] | |
| web_search = "No" | |
| for d in messages: | |
| score = retrieval_grader.invoke({"question": question, "document": d.page_content}) | |
| grade = score.binary_score | |
| if grade == "yes": | |
| print("---GRADE: DOCUMENT RELEVANT---") | |
| filtered_docs.append(d) | |
| else: | |
| print("---GRADE: DOCUMENT NOT RELEVANT---") | |
| continue | |
| print("---TOOLS RETRIEVE---") | |
| tool_documents = tool_chain.invoke(question) | |
| #print(tool_documents) | |
| if tool_documents: | |
| for item in tool_documents: | |
| filtered_docs.append(Document(page_content=str(item['output']),metadata={"source": 'https://svarupa.org/home',"name":item['name']})) | |
| # If filtered_docs is empty, perform a web search | |
| if not filtered_docs: | |
| print("--PERFORMING WEB SEARCH--") | |
| web_search = "Yes" | |
| return {"messages": 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 | |
| """ | |
| print("---TRANSFORM QUERY---") | |
| question = state["question"] | |
| messages = state["messages"] | |
| # Re-write question | |
| better_question = question_rewriter.invoke({"question": question}) | |
| return {"messages": messages, "question": better_question} | |
| 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 | |
| """ | |
| print("---WEB SEARCH---") | |
| question = state["question"] | |
| messages = state["messages"] | |
| # Web search | |
| docs = web_search_tool.invoke({"query": question}) | |
| #web_results = "\n".join([d["content"] for d in docs]) | |
| web_results = [Document(page_content=d["content"], metadata={"source": d["url"]}) for d in docs] | |
| print(f"Web Results: {web_results}") | |
| messages.extend(web_results) | |
| return {"messages": messages, "question": question} | |
| ### 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 | |
| """ | |
| print("---ASSESS GRADED DOCUMENTS---") | |
| question = state["question"] | |
| web_search = state["web_search"] | |
| filtered_documents = state["messages"] | |
| if web_search == "Yes": | |
| # All documents have been filtered check_relevance | |
| # We will re-generate a new query | |
| print("---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---") | |
| return "transform_query" | |
| else: | |
| # We have relevant documents, so generate answer | |
| print("---DECISION: GENERATE---") | |
| return "generate" | |
| from langgraph.graph import END, StateGraph | |
| 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) # generatae | |
| workflow.add_node("transform_query", transform_query) # transform_query | |
| workflow.add_node("web_search_node", web_search) # web search | |
| # 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", | |
| }, | |
| ) | |
| workflow.add_edge("transform_query", "web_search_node") | |
| workflow.add_edge("web_search_node", "generate") | |
| workflow.add_edge("generate", END) | |
| # Compile | |
| crag_app = workflow.compile() |