Spaces:
Build error
Build error
File size: 2,912 Bytes
1d89d60 | 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 | import os
from typing import Annotated, TypedDict, Sequence
from langgraph.graph import StateGraph, START, END
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage
from langgraph.graph.message import add_messages
from langchain_google_genai import ChatGoogleGenerativeAI as Gemini
from langgraph.prebuilt import ToolNode
from tools import tools
from vector_db import vector_db, add_recent_question
from dotenv import load_dotenv
load_dotenv()
# initiate model
api_key = os.getenv("gemini_api_key")
model = Gemini(model="gemini-2.0-flash", temperature=0.8, api_key=api_key).bind_tools(
tools=tools
)
class MessageState(TypedDict):
messages: Annotated[Sequence[BaseMessage], add_messages]
# creating nodes
def retrieve(state: MessageState) -> MessageState:
user_input = state["messages"][-1].content
similar_questions = vector_db.similarity_search(user_input, k=1)
if similar_questions:
reference_msg = HumanMessage(
content=f"Here is a similar question and answer from history:\n\n{similar_questions[0].page_content}"
)
else:
reference_msg = HumanMessage(
content="No similar question found in the database."
)
add_recent_question(user_input)
return {"messages": [reference_msg]}
def call_model(state: MessageState) -> MessageState:
response = model.invoke(state["messages"])
return {"messages": [response]}
def should_continue(state: MessageState) -> str:
last_message = state["messages"][-1]
if isinstance(last_message, AIMessage) and last_message.tool_calls:
return "tools_node"
return "end"
tools_node = ToolNode(tools=tools)
graph = StateGraph(MessageState)
graph.add_node("retriever_node", retrieve)
graph.add_node("model_node", call_model)
graph.add_node("tools_node", tools_node)
graph.add_edge(START, "retriever_node")
graph.add_edge("retriever_node", "model_node")
graph.add_conditional_edges(
"model_node",
should_continue,
{
"tools_node": "tools_node",
"end": END,
},
)
graph.add_edge("tools_node", "model_node")
agent = graph.compile()
#testing
if __name__ == "__main__":
with open("./system_prompt.txt", "r", encoding="utf-8") as f:
SYSTEM_PROMPT_CONTENT = f.read()
SYSTEM_MESSAGE = SystemMessage(content=SYSTEM_PROMPT_CONTENT)
user_input = ""
while user_input != "exit":
user_input = input("question : ")
user_query = HumanMessage(content=user_input)
response = agent.invoke({"messages": [SYSTEM_MESSAGE] + [user_query]})
last_message = response["messages"][-1]
if isinstance(last_message, AIMessage):
print(f"Agent: {last_message.content[14:]}")
print(f"AI-debug: {last_message}")
else:
print(f"AI-debug: {last_message}")
|