|
|
| import os
|
| from typing import TypedDict, Annotated, List
|
| from langchain_core.messages import BaseMessage, ToolMessage
|
| from langchain_core.prompts import ChatPromptTemplate
|
| from langchain_core.output_parsers import StrOutputParser
|
| from langchain_core.tools import tool
|
| from langchain_groq import ChatGroq
|
| from langgraph.graph import StateGraph, END
|
| from langgraph.prebuilt import ToolNode
|
| from utils import process_documents
|
| from dotenv import load_dotenv
|
|
|
| load_dotenv()
|
|
|
|
|
| retriever_for_tool = None
|
|
|
| class AgentState(TypedDict):
|
| messages: Annotated[List[BaseMessage], lambda x, y: x + y]
|
| question: str
|
| answer: str
|
|
|
|
|
|
|
| def set_retriever_for_tool(retriever):
|
| global retriever_for_tool
|
| retriever_for_tool = retriever
|
|
|
| @tool
|
| def retrieve_clauses(query: str) -> str:
|
| """
|
| Searches the policy document for relevant clauses based on a user's query.
|
| Use this tool to find information within the document to answer a question.
|
| """
|
| global retriever_for_tool
|
| if retriever_for_tool is None:
|
| return "No retriever set. Please process the document first."
|
| retrieved_docs = retriever_for_tool.invoke(query)
|
| context = "\n\n---\n\n".join([doc.page_content for doc in retrieved_docs])
|
| return context
|
|
|
| llm = ChatGroq(model="meta-llama/llama-4-scout-17b-16e-instruct", temperature=0, max_tokens=1024)
|
| tools = [retrieve_clauses]
|
| llm_with_tools = llm.bind_tools(tools)
|
|
|
| final_answer_prompt = ChatPromptTemplate.from_template(
|
| """
|
| You are a precise and helpful assistant. Your task is to answer the user's question based ONLY on the provided context document. You must not use any external knowledge or make assumptions beyond what is explicitly stated in the document.
|
|
|
| ### Context ###
|
| {context}
|
|
|
| ### Question ###
|
| {question}
|
|
|
| ### Instructions ###
|
| 1. Read the context carefully to find the information relevant to the question.
|
| 2. Formulate a concise answer based solely on the information found in the context. The answer should be clear, direct, and based entirely on the context provided.
|
| 3. No answer should exceed 50 words.
|
| 4. If the answer to the question cannot be found in the provided context, you MUST respond with the exact phrase: 'The answer to your question is not available in the provided document.'
|
| 5. Do not, under any circumstances, invent an answer or use information from outside the provided context.
|
| 6. Do not use special characters. The output should be a single line of text without any formatting or additional characters.
|
| ### Answer ###
|
| """
|
| )
|
| string_parser = StrOutputParser()
|
| generation_chain = final_answer_prompt | llm | string_parser
|
|
|
| def call_model(state: AgentState):
|
| messages = state["messages"]
|
| response = llm_with_tools.invoke(messages)
|
| return {"messages": [response]}
|
|
|
| def should_continue(state: AgentState):
|
| if state["messages"][-1].tool_calls:
|
| return "continue"
|
| return "end"
|
|
|
| def generate_final_answer(state: AgentState):
|
| question = state["question"]
|
| context = ""
|
| for msg in reversed(state["messages"]):
|
| if isinstance(msg, ToolMessage):
|
| context = msg.content
|
| break
|
| response = generation_chain.invoke({"context": context, "question": question})
|
| return {"answer": response}
|
|
|
| workflow = StateGraph(AgentState)
|
| workflow.add_node("agent", call_model)
|
| workflow.add_node("action", ToolNode(tools))
|
| workflow.add_node("final_answer", generate_final_answer)
|
|
|
| workflow.set_entry_point("agent")
|
| workflow.add_conditional_edges("agent", should_continue, {"continue": "action", "end": "final_answer"})
|
| workflow.add_edge("action", "agent")
|
| workflow.add_edge("final_answer", END)
|
|
|
| agent_app = workflow.compile() |
