Spaces:
Sleeping
Sleeping
| import os | |
| from langchain_groq import ChatGroq | |
| from langgraph.graph import START, StateGraph, MessagesState | |
| from langgraph.prebuilt import tools_condition | |
| from langgraph.prebuilt import ToolNode | |
| from langchain_core.messages import SystemMessage, HumanMessage | |
| from tools import * | |
| from typing import TypedDict, Optional | |
| tools = [add, multiply, divide, subtract, search_wikipedia] | |
| def build_agent(): | |
| class AgentState(TypedDict): | |
| messages: list # chat history | |
| wiki_table: Optional[str] | |
| llm = ChatGroq(model="qwen-qwq-32b", temperature=0) | |
| chat_with_tools = llm.bind_tools(tools) | |
| def assistant(state: AgentState): | |
| table = state.get("wiki_table") | |
| context = [] | |
| if table: | |
| context.append({"role": "system", "content": f"Here is a Wikipedia table you can use:\n\n{table}"}) | |
| response = chat_with_tools.invoke(context + state["messages"]) | |
| return { | |
| "messages": [response], | |
| "wiki_table": table, | |
| } | |
| def enhancer(state: AgentState): | |
| sys_msg = """ | |
| You are a helpful assistant tasked with answering questions using a set of tools. You can reason through problems and take actions using available tools to solve complex tasks. | |
| Follow this format strictly: | |
| Message: {user question} | |
| Thought: Describe your reasoning about the question. | |
| Action: Call one of the available tools with a specific input. | |
| Observation: Note what was returned from the tool. | |
| Repeat the Thought → Action → Observation steps as many times as necessary until you have enough information to answer the original question. | |
| Now, I will ask you a question. Report your thoughts, and finish your answer with the following template: | |
| FINAL ANSWER: [YOUR FINAL ANSWER]. | |
| YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. | |
| If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. | |
| If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. | |
| If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string. | |
| Your answer should only start with "FINAL ANSWER: ", then follows with the answer. | |
| """ | |
| context = [] | |
| context.append({"role": "system", "content": sys_msg}) | |
| return context + state["messages"] | |
| ## The graph | |
| builder = StateGraph(AgentState) | |
| # Define nodes: these do the work | |
| builder.add_node("enhancer", enhancer) | |
| builder.add_node("assistant", assistant) | |
| builder.add_node("tools", ToolNode(tools)) | |
| # Define edges: these determine how the control flow moves | |
| builder.add_edge(START, "enhancer") | |
| builder.add_edge("enhancer", "assistant") | |
| builder.add_conditional_edges( | |
| "assistant", | |
| # If the latest message requires a tool, route to tools | |
| # Otherwise, provide a direct response | |
| tools_condition, | |
| ) | |
| builder.add_edge("tools", "assistant") | |
| return builder.compile() |