First_agent_template / tools /parse_km_reading_tool.py
TomMoeras's picture
Change output_type from "dict" to "any" in multiple tools for improved flexibility
a8853dd
from typing import Any, Dict, Union
from smolagents.tools import Tool
import re
class ParseKmReadingTool(Tool):
name = "parse_km_reading"
description = "Parses kilometer readings from user input."
inputs = {
'input_text': {'type': 'string', 'description': 'Text input from the user that may contain kilometer readings'}
}
output_type = "any"
def __init__(self):
self.is_initialized = True
# Store the last valid kilometer reading to use as start_km for future trips
self.last_reading = None
def forward(self, input_text: str) -> Dict[str, Any]:
"""
Parse kilometer readings from user input.
Args:
input_text: Text input that may contain kilometer readings
Returns:
A dictionary with parsed kilometer readings
"""
try:
# Clean and normalize input
clean_text = input_text.strip().replace(',', '')
# Find all numbers in the input
numbers = re.findall(r'\b\d+\b', clean_text)
numbers = [int(num) for num in numbers]
# Determine what kind of reading this is
result = {
"success": True,
"numbers_found": numbers,
"reading_type": None
}
if len(numbers) == 0:
result["success"] = False
result["error"] = "No kilometer readings found in input"
return result
if len(numbers) == 1:
# Single reading - could be start or end
km_reading = numbers[0]
if self.last_reading is None:
# First reading ever, must be start
result["reading_type"] = "start"
result["start_km"] = km_reading
self.last_reading = km_reading
else:
# We have a previous reading
if km_reading > self.last_reading:
# This is likely an end reading
result["reading_type"] = "end"
result["end_km"] = km_reading
result["start_km"] = self.last_reading
result["km_traveled"] = km_reading - self.last_reading
self.last_reading = km_reading
else:
# This is likely a new start reading
result["reading_type"] = "start"
result["start_km"] = km_reading
self.last_reading = km_reading
elif len(numbers) == 2:
# Two readings - likely start and end
start_km = min(numbers)
end_km = max(numbers)
result["reading_type"] = "complete"
result["start_km"] = start_km
result["end_km"] = end_km
result["km_traveled"] = end_km - start_km
self.last_reading = end_km
else:
# More than 2 numbers - take first and last assuming chronological order
result["reading_type"] = "multiple"
result["start_km"] = numbers[0]
result["end_km"] = numbers[-1]
result["km_traveled"] = numbers[-1] - numbers[0]
result["all_readings"] = numbers
self.last_reading = numbers[-1]
return result
except Exception as e:
return {
"success": False,
"error": str(e)
}