File size: 3,958 Bytes
d9cd5c8 f970994 d9cd5c8 f970994 d9cd5c8 f970994 8b0215d f970994 d9cd5c8 | 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 |
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()
# The retriever will be set at runtime from the FastAPI endpoint
retriever_for_tool = None
class AgentState(TypedDict):
messages: Annotated[List[BaseMessage], lambda x, y: x + y]
question: str
answer: str
# This setter will be called from FastAPI to set the retriever for each request
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() |