Spaces:
Sleeping
Sleeping
File size: 7,316 Bytes
374dd02 |
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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 |
import json
import os
import re
from dotenv import load_dotenv
from langchain_core.messages import (AIMessage, HumanMessage, SystemMessage,
ToolMessage)
from langchain_huggingface import (ChatHuggingFace, HuggingFaceEmbeddings,
HuggingFaceEndpoint)
from langgraph.graph import START, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from tools import (absolute, add, analyze_csv_file, analyze_excel_file,
arvix_search, audio_transcription, compound_interest,
convert_temperature, divide, exponential,
extract_text_from_image, factorial, floor_divide,
get_current_time_in_timezone, greatest_common_divisor,
is_prime, least_common_multiple, logarithm, modulus,
multiply, percentage_calculator, power, python_code_parser,
reverse_sentence, roman_calculator_converter, square_root,
subtract, web_content_extract, web_search, wiki_search)
# Load Constants
load_dotenv()
HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
tools = [
multiply, add, subtract, power, divide, modulus,
square_root, floor_divide, absolute, logarithm,
exponential, web_search, roman_calculator_converter,
get_current_time_in_timezone, compound_interest,
convert_temperature, factorial, greatest_common_divisor,
is_prime, least_common_multiple, percentage_calculator,
wiki_search, analyze_excel_file, arvix_search,
audio_transcription, python_code_parser, analyze_csv_file,
extract_text_from_image, reverse_sentence, web_content_extract,
]
# Load system prompt
system_prompt = """
You are a general AI assistant. I will ask you a question.
Report your thoughts, and finish your answer with only the answer, no extra text, no prefix, and no explanation.
Your 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 a comma to write your number, nor 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 on whether the element to be put in the list is a number or a string.
Format your output as: [{"task_id": ..., "submitted_answer": ...}]
Do NOT include the format string or any JSON inside the submitted_answer field. Only output a single flat list as: [{"task_id": ..., "submitted_answer": ...}]
"""
# System message
sys_msg = SystemMessage(content=system_prompt)
def build_graph():
"""Build the graph"""
# First create the HuggingFaceEndpoint
llm_endpoint = HuggingFaceEndpoint(
repo_id="mistralai/Mistral-7B-Instruct-v0.2",
huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN,
#api_key=GEMINI_API_KEY,
temperature=0.1,
max_new_tokens=1024,
timeout=60,
)
# Then wrap it with ChatHuggingFace to get chat model functionality
llm = ChatHuggingFace(llm=llm_endpoint)
# Bind tools to LLM
llm_with_tools = llm.bind_tools(tools)
# --- Nodes ---
def extract_answer(llm_output):
# Try to parse as JSON if possible
try:
# If the LLM output is a JSON list, extract the answer
parsed = json.loads(llm_output.strip().split('\n')[0])
if isinstance(parsed, list) and isinstance(parsed[0], dict) and "submitted_answer" in parsed[0]:
return parsed[0]["submitted_answer"]
except Exception:
pass
# Otherwise, just return the first line (before any explanation)
return llm_output.strip().split('\n')[0]
def assistant(state: MessagesState):
messages_with_system_prompt = [sys_msg] + state["messages"]
llm_response = llm_with_tools.invoke(messages_with_system_prompt)
answer_text = extract_answer(llm_response.content)
task_id = str(state.get("task_id", "1")) # Ensure task_id is a string
formatted = [{"task_id": task_id, "submitted_answer": answer_text}]
return {"messages": [AIMessage(content=json.dumps(formatted, ensure_ascii=False))]}
# --- Graph Definition ---
builder = StateGraph(MessagesState)
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
builder.add_edge(START, "assistant")
builder.add_conditional_edges("assistant", tools_condition)
builder.add_edge("tools", "assistant")
# Compile graph
return builder.compile()
def is_valid_agent_output(output):
"""
Checks if the output matches the required format:
Answers (answers): [{"task_id": ..., "submitted_answer": ...}]
"""
# Basic regex to check the format
pattern = r'^Answers \(answers\): \[(\{.*\})\]$'
match = re.match(pattern, output.strip())
if not match:
return False
# Try to parse the JSON part
try:
answers_list = json.loads(f'[{match.group(1)}]')
# Check required keys
for ans in answers_list:
if not isinstance(ans, dict):
return False
if "task_id" not in ans or "submitted_answer" not in ans:
return False
return True
except Exception:
return False
def extract_flat_answer(output):
# Try to find the innermost Answers (answers): [{...}]
pattern = r'Answers \(answers\): \[(\{.*?\})\]'
matches = re.findall(pattern, output)
if matches:
# Use the last match (innermost)
try:
answers_list = json.loads(f'[{matches[-1]}]')
if isinstance(answers_list, list) and "task_id" in answers_list[0] and "submitted_answer" in answers_list[0]:
return f'Answers (answers): [{matches[-1]}]'
except Exception:
pass
return output # fallback
# test
if __name__ == "__main__":
question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
# Build the graph
graph = build_graph()
# Run the graph
messages = [HumanMessage(content=question)]
# The initial state for the graph
initial_state = {"messages": messages, "task_id": "test123"}
# Invoke the graph stream to see the steps
for s in graph.stream(initial_state, stream_mode="values"):
message = s["messages"][-1]
if isinstance(message, ToolMessage):
print("---RETRIEVED CONTEXT---")
print(message.content)
print("-----------------------")
else:
output = message.content # This is a string
try:
parsed = json.loads(output)
if isinstance(parsed, list) and "task_id" in parsed[0] and "submitted_answer" in parsed[0]:
print("✅ Output is in the correct format!")
else:
print("❌ Output is NOT in the correct format!")
except Exception as e:
print("❌ Output is NOT in the correct format!", e) |