Add math_calculate and split_bill agent tools
- math_calculate: safe AST-based math expression evaluator - split_bill: receipt image analysis with even bill splitting - Expose image_key to agent context for tool referencing
This commit is contained in:
@@ -66,6 +66,7 @@ def _build_messages_payload(
|
||||
query_text: str,
|
||||
image_description: str | None = None,
|
||||
system_prompt: str | None = None,
|
||||
image_key: str | None = None,
|
||||
) -> list:
|
||||
recent_messages = (
|
||||
conversation.messages[-10:]
|
||||
@@ -77,12 +78,14 @@ def _build_messages_payload(
|
||||
role = "user" if msg.speaker == "user" else "assistant"
|
||||
text = msg.text
|
||||
if msg.image_key and role == "user":
|
||||
text = f"[User sent an image]\n{text}"
|
||||
text = f"[User sent an image (image_key: {msg.image_key})]\n{text}"
|
||||
messages_payload.append({"role": role, "content": text})
|
||||
|
||||
# Build the current user message with optional image description
|
||||
if image_description:
|
||||
content = f"[Image analysis: {image_description}]"
|
||||
if image_key:
|
||||
content = f"[Image key: {image_key}]\n{content}"
|
||||
if query_text:
|
||||
content = f"{query_text}\n\n{content}"
|
||||
else:
|
||||
@@ -203,7 +206,11 @@ async def stream_query():
|
||||
|
||||
system_prompt = await _build_system_prompt_with_memories(str(user.id))
|
||||
messages_payload = _build_messages_payload(
|
||||
conversation, query_text or "", image_description, system_prompt=system_prompt
|
||||
conversation,
|
||||
query_text or "",
|
||||
image_description,
|
||||
system_prompt=system_prompt,
|
||||
image_key=image_key,
|
||||
)
|
||||
payload = {"messages": messages_payload}
|
||||
agent_config = {
|
||||
|
||||
@@ -1044,8 +1044,235 @@ async def simbakit_drinking_activity(pet_name: str = "", days: int = 7) -> str:
|
||||
return f"Error fetching drinking activity: {str(e)}"
|
||||
|
||||
|
||||
@tool
|
||||
def math_calculate(expression: str) -> str:
|
||||
"""Evaluate a math expression and return the result.
|
||||
|
||||
Use this tool when the user asks you to:
|
||||
- Do arithmetic or math calculations
|
||||
- Calculate percentages, tips, or splits
|
||||
- Solve a math expression
|
||||
- Compare numbers or compute differences
|
||||
|
||||
Supports: +, -, *, /, //, **, %, parentheses, and basic functions
|
||||
(sqrt, abs, round, min, max). Use 'math.' prefix for functions.
|
||||
|
||||
Args:
|
||||
expression: A math expression to evaluate (e.g., "(123.45 * 0.18) / 4")
|
||||
|
||||
Returns:
|
||||
The numeric result of the expression.
|
||||
"""
|
||||
import ast
|
||||
import operator
|
||||
import math as _math
|
||||
|
||||
# Allowed operators and functions
|
||||
allowed_ops = {
|
||||
ast.Add: operator.add,
|
||||
ast.Sub: operator.sub,
|
||||
ast.Mult: operator.mul,
|
||||
ast.Div: operator.truediv,
|
||||
ast.FloorDiv: operator.floordiv,
|
||||
ast.Mod: operator.mod,
|
||||
ast.Pow: operator.pow,
|
||||
ast.USub: operator.neg,
|
||||
ast.UAdd: operator.pos,
|
||||
}
|
||||
allowed_funcs = {
|
||||
"sqrt": _math.sqrt,
|
||||
"abs": abs,
|
||||
"round": round,
|
||||
"min": min,
|
||||
"max": max,
|
||||
"pi": _math.pi,
|
||||
"e": _math.e,
|
||||
}
|
||||
|
||||
def _eval_node(node):
|
||||
if isinstance(node, ast.Expression):
|
||||
return _eval_node(node.body)
|
||||
if isinstance(node, ast.Constant):
|
||||
if isinstance(node.value, (int, float)):
|
||||
return node.value
|
||||
raise ValueError(f"Unsupported constant: {node.value!r}")
|
||||
if isinstance(node, ast.BinOp):
|
||||
op_type = type(node.op)
|
||||
if op_type not in allowed_ops:
|
||||
raise ValueError(f"Operator not allowed: {op_type.__name__}")
|
||||
return allowed_ops[op_type](_eval_node(node.left), _eval_node(node.right))
|
||||
if isinstance(node, ast.UnaryOp):
|
||||
op_type = type(node.op)
|
||||
if op_type not in allowed_ops:
|
||||
raise ValueError(f"Unary operator not allowed: {op_type.__name__}")
|
||||
return allowed_ops[op_type](_eval_node(node.operand))
|
||||
if isinstance(node, ast.Call):
|
||||
if not isinstance(node.func, ast.Name):
|
||||
raise ValueError("Only simple function calls allowed")
|
||||
func_name = node.func.id
|
||||
if func_name not in allowed_funcs:
|
||||
raise ValueError(f"Function not allowed: {func_name}")
|
||||
args = [_eval_node(a) for a in node.args]
|
||||
return allowed_funcs[func_name](*args)
|
||||
if isinstance(node, ast.Name):
|
||||
if node.id in allowed_funcs:
|
||||
return allowed_funcs[node.id]
|
||||
raise ValueError(f"Unknown name: {node.id}")
|
||||
raise ValueError(f"Unsupported expression element: {ast.dump(node)}")
|
||||
|
||||
try:
|
||||
tree = ast.parse(expression, mode="eval")
|
||||
result = _eval_node(tree)
|
||||
# Format: show full precision for decimals, trim trailing zeros
|
||||
if isinstance(result, float):
|
||||
if result == int(result) and abs(result) < 1e15:
|
||||
return str(int(result))
|
||||
return f"{result:.10g}"
|
||||
return str(result)
|
||||
except ZeroDivisionError:
|
||||
return "Error: Division by zero."
|
||||
except Exception as e:
|
||||
return f"Error evaluating expression: {str(e)}"
|
||||
|
||||
|
||||
@tool
|
||||
async def split_bill(
|
||||
image_key: str, num_people: int = 2, include_tip: bool = True
|
||||
) -> str:
|
||||
"""Analyze a receipt image and split the bill evenly among people.
|
||||
|
||||
Use this tool when the user:
|
||||
- Sends a photo of a receipt or bill and asks to split it
|
||||
- Asks "how much does each person owe"
|
||||
- Wants to calculate a fair split of a bill from an image
|
||||
|
||||
Args:
|
||||
image_key: The S3 key of the uploaded receipt image (from the conversation context)
|
||||
num_people: Number of people to split the bill among (default 2)
|
||||
include_tip: Whether to include the tip in the split (default True)
|
||||
|
||||
Returns:
|
||||
The bill breakdown with line items, total, tip, and each person's share.
|
||||
"""
|
||||
import base64
|
||||
import json
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
from utils.s3_client import get_image as s3_get_image
|
||||
|
||||
try:
|
||||
image_bytes, _ = await s3_get_image(image_key)
|
||||
except Exception as e:
|
||||
return f"Error: Could not fetch image '{image_key}': {str(e)}"
|
||||
|
||||
# Normalize to JPEG
|
||||
from utils.image_process import _to_jpeg_bytes
|
||||
|
||||
image_bytes = _to_jpeg_bytes(image_bytes)
|
||||
|
||||
llama_url = os.getenv("LLAMA_SERVER_URL")
|
||||
if llama_url:
|
||||
aclient = AsyncOpenAI(base_url=llama_url, api_key="not-needed")
|
||||
model = os.getenv("LLAMA_MODEL_NAME", "llama-3.1-8b-instruct")
|
||||
else:
|
||||
aclient = AsyncOpenAI()
|
||||
model = "gpt-4o-mini"
|
||||
|
||||
b64 = base64.b64encode(image_bytes).decode("utf-8")
|
||||
|
||||
system_prompt = (
|
||||
"You are a receipt parser. Extract all line items, subtotals, tax, tip, and total "
|
||||
"from the receipt. Respond ONLY with valid JSON matching this schema:\n"
|
||||
"{\n"
|
||||
' "items": [{"name": "string", "price": number}],\n'
|
||||
' "subtotal": number,\n'
|
||||
' "tax": number,\n'
|
||||
' "tip": number,\n'
|
||||
' "total": number\n'
|
||||
"}\n"
|
||||
"If a field is not present on the receipt, use 0. Do not include any text outside the JSON."
|
||||
)
|
||||
|
||||
try:
|
||||
response = await aclient.chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Parse this receipt into structured JSON.",
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:image/jpeg;base64,{b64}"},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
)
|
||||
except Exception as e:
|
||||
return f"Error: Vision model failed to analyze receipt: {str(e)}"
|
||||
|
||||
raw = response.choices[0].message.content.strip()
|
||||
# Strip markdown code fences if present
|
||||
if raw.startswith("```"):
|
||||
raw = raw.split("\n", 1)[-1]
|
||||
if raw.endswith("```"):
|
||||
raw = raw[:-3]
|
||||
|
||||
try:
|
||||
receipt = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
return f"Error: Could not parse receipt data from model response.\n\nRaw: {raw}"
|
||||
|
||||
# Compute the split
|
||||
items = receipt.get("items", [])
|
||||
subtotal = receipt.get("subtotal", sum(i.get("price", 0) for i in items))
|
||||
tax = receipt.get("tax", 0)
|
||||
tip = receipt.get("tip", 0)
|
||||
total = receipt.get("total", subtotal + tax + tip)
|
||||
|
||||
if not include_tip:
|
||||
split_amount = subtotal + tax
|
||||
else:
|
||||
split_amount = total
|
||||
|
||||
per_person = split_amount / num_people
|
||||
|
||||
# Format result
|
||||
lines = ["**Bill Breakdown:**\n"]
|
||||
if items:
|
||||
for item in items:
|
||||
lines.append(f"- {item.get('name', 'Item')}: ${item.get('price', 0):.2f}")
|
||||
lines.append("")
|
||||
lines.append(f"Subtotal: ${subtotal:.2f}")
|
||||
if tax:
|
||||
lines.append(f"Tax: ${tax:.2f}")
|
||||
if tip:
|
||||
lines.append(f"Tip: ${tip:.2f}")
|
||||
lines.append(f"Total: ${total:.2f}")
|
||||
lines.append("")
|
||||
|
||||
tip_note = " (includes tip)" if include_tip else " (excluding tip)"
|
||||
lines.append(
|
||||
f"Split among {num_people} people{tip_note}: **${per_person:.2f} each**"
|
||||
)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
# Create tools list based on what's available
|
||||
tools = [get_current_date, simba_search, web_search, save_user_memory]
|
||||
tools = [
|
||||
get_current_date,
|
||||
simba_search,
|
||||
web_search,
|
||||
save_user_memory,
|
||||
math_calculate,
|
||||
split_bill,
|
||||
]
|
||||
if simbakit_enabled:
|
||||
tools.extend(
|
||||
[
|
||||
|
||||
Reference in New Issue
Block a user