Add token and context window tracking
This commit is contained in:
@@ -12,6 +12,7 @@ install_deps() {
|
|||||||
command -v socat &>/dev/null || needed="$needed socat"
|
command -v socat &>/dev/null || needed="$needed socat"
|
||||||
command -v jq &>/dev/null || needed="$needed jq"
|
command -v jq &>/dev/null || needed="$needed jq"
|
||||||
command -v sqlite3 &>/dev/null || needed="$needed sqlite3"
|
command -v sqlite3 &>/dev/null || needed="$needed sqlite3"
|
||||||
|
command -v bc &>/dev/null || needed="$needed bc"
|
||||||
|
|
||||||
if [ -n "$needed" ]; then
|
if [ -n "$needed" ]; then
|
||||||
echo "Installing dependencies:$needed"
|
echo "Installing dependencies:$needed"
|
||||||
@@ -63,6 +64,10 @@ CREATE TABLE IF NOT EXISTS conversations (
|
|||||||
instance_id TEXT,
|
instance_id TEXT,
|
||||||
started_at TEXT DEFAULT CURRENT_TIMESTAMP,
|
started_at TEXT DEFAULT CURRENT_TIMESTAMP,
|
||||||
status TEXT DEFAULT 'active',
|
status TEXT DEFAULT 'active',
|
||||||
|
tokens_in INTEGER DEFAULT 0,
|
||||||
|
tokens_out INTEGER DEFAULT 0,
|
||||||
|
total_tokens INTEGER DEFAULT 0,
|
||||||
|
max_context INTEGER DEFAULT 200000,
|
||||||
FOREIGN KEY (instance_id) REFERENCES instances(id)
|
FOREIGN KEY (instance_id) REFERENCES instances(id)
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -72,6 +77,7 @@ CREATE TABLE IF NOT EXISTS messages (
|
|||||||
role TEXT,
|
role TEXT,
|
||||||
content TEXT,
|
content TEXT,
|
||||||
tool_calls TEXT,
|
tool_calls TEXT,
|
||||||
|
tokens INTEGER DEFAULT 0,
|
||||||
timestamp TEXT DEFAULT CURRENT_TIMESTAMP,
|
timestamp TEXT DEFAULT CURRENT_TIMESTAMP,
|
||||||
FOREIGN KEY (conversation_id) REFERENCES conversations(id)
|
FOREIGN KEY (conversation_id) REFERENCES conversations(id)
|
||||||
);
|
);
|
||||||
|
|||||||
@@ -287,27 +287,90 @@ delete_instance() {
|
|||||||
send_json '{"deleted":true}'
|
send_json '{"deleted":true}'
|
||||||
}
|
}
|
||||||
|
|
||||||
|
# Estimate tokens from text (~4 chars per token)
|
||||||
|
estimate_tokens() {
|
||||||
|
local text="$1"
|
||||||
|
local chars=${#text}
|
||||||
|
echo $(( (chars + 3) / 4 ))
|
||||||
|
}
|
||||||
|
|
||||||
# Store conversation in DB
|
# Store conversation in DB
|
||||||
store_conversation() {
|
store_conversation() {
|
||||||
local conv_id="$1"
|
local conv_id="$1"
|
||||||
local instance_id="$2"
|
local instance_id="$2"
|
||||||
|
local max_context="${3:-200000}"
|
||||||
local db="$AGENTS_DATA_DIR/agents.db"
|
local db="$AGENTS_DATA_DIR/agents.db"
|
||||||
|
|
||||||
sqlite3 "$db" "INSERT OR IGNORE INTO conversations (id, instance_id, status) VALUES ('$conv_id', '$instance_id', 'active');"
|
sqlite3 "$db" "INSERT OR IGNORE INTO conversations (id, instance_id, status, max_context) VALUES ('$conv_id', '$instance_id', 'active', $max_context);"
|
||||||
}
|
}
|
||||||
|
|
||||||
# Store message in DB
|
# Store message in DB and update token counts
|
||||||
store_message() {
|
store_message() {
|
||||||
local conv_id="$1"
|
local conv_id="$1"
|
||||||
local role="$2"
|
local role="$2"
|
||||||
local content="$3"
|
local content="$3"
|
||||||
local db="$AGENTS_DATA_DIR/agents.db"
|
local db="$AGENTS_DATA_DIR/agents.db"
|
||||||
local msg_id=$(gen_id)
|
local msg_id=$(gen_id)
|
||||||
|
local tokens=$(estimate_tokens "$content")
|
||||||
|
|
||||||
# Escape single quotes
|
# Escape single quotes
|
||||||
content=$(echo "$content" | sed "s/'/''/g")
|
local escaped_content=$(echo "$content" | sed "s/'/''/g")
|
||||||
|
|
||||||
sqlite3 "$db" "INSERT INTO messages (id, conversation_id, role, content) VALUES ('$msg_id', '$conv_id', '$role', '$content');"
|
sqlite3 "$db" "INSERT INTO messages (id, conversation_id, role, content, tokens) VALUES ('$msg_id', '$conv_id', '$role', '$escaped_content', $tokens);"
|
||||||
|
|
||||||
|
# Update conversation token counts
|
||||||
|
if [ "$role" = "user" ]; then
|
||||||
|
sqlite3 "$db" "UPDATE conversations SET tokens_in = tokens_in + $tokens, total_tokens = total_tokens + $tokens WHERE id = '$conv_id';"
|
||||||
|
else
|
||||||
|
sqlite3 "$db" "UPDATE conversations SET tokens_out = tokens_out + $tokens, total_tokens = total_tokens + $tokens WHERE id = '$conv_id';"
|
||||||
|
fi
|
||||||
|
}
|
||||||
|
|
||||||
|
# Get conversation usage
|
||||||
|
get_conversation_usage() {
|
||||||
|
local conv_id="$1"
|
||||||
|
local db="$AGENTS_DATA_DIR/agents.db"
|
||||||
|
|
||||||
|
local row=$(sqlite3 -json "$db" "SELECT id, tokens_in, tokens_out, total_tokens, max_context FROM conversations WHERE id='$conv_id';" 2>/dev/null)
|
||||||
|
|
||||||
|
if [ -z "$row" ] || [ "$row" = "[]" ]; then
|
||||||
|
echo "{}"
|
||||||
|
return
|
||||||
|
fi
|
||||||
|
|
||||||
|
# Calculate usage percentage and cost
|
||||||
|
local total=$(echo "$row" | jq -r '.[0].total_tokens // 0')
|
||||||
|
local max=$(echo "$row" | jq -r '.[0].max_context // 200000')
|
||||||
|
local tokens_in=$(echo "$row" | jq -r '.[0].tokens_in // 0')
|
||||||
|
local tokens_out=$(echo "$row" | jq -r '.[0].tokens_out // 0')
|
||||||
|
|
||||||
|
local usage_pct=0
|
||||||
|
if [ "$max" -gt 0 ]; then
|
||||||
|
usage_pct=$(echo "scale=1; $total * 100 / $max" | bc 2>/dev/null || echo "0")
|
||||||
|
fi
|
||||||
|
|
||||||
|
# Estimate cost (Claude Sonnet: $3/1M input, $15/1M output)
|
||||||
|
local cost_in=$(echo "scale=4; $tokens_in * 0.000003" | bc 2>/dev/null || echo "0")
|
||||||
|
local cost_out=$(echo "scale=4; $tokens_out * 0.000015" | bc 2>/dev/null || echo "0")
|
||||||
|
local total_cost=$(echo "scale=4; $cost_in + $cost_out" | bc 2>/dev/null || echo "0")
|
||||||
|
|
||||||
|
local status="ok"
|
||||||
|
if [ "${usage_pct%.*}" -ge 90 ]; then
|
||||||
|
status="critical"
|
||||||
|
elif [ "${usage_pct%.*}" -ge 75 ]; then
|
||||||
|
status="warning"
|
||||||
|
fi
|
||||||
|
|
||||||
|
jq -n \
|
||||||
|
--arg id "$conv_id" \
|
||||||
|
--argjson tokens_in "$tokens_in" \
|
||||||
|
--argjson tokens_out "$tokens_out" \
|
||||||
|
--argjson total "$total" \
|
||||||
|
--argjson max "$max" \
|
||||||
|
--arg usage_pct "$usage_pct" \
|
||||||
|
--arg cost "$total_cost" \
|
||||||
|
--arg status "$status" \
|
||||||
|
'{conversation_id:$id, tokens_in:$tokens_in, tokens_out:$tokens_out, total_tokens:$total, max_context:$max, usage_percent:($usage_pct|tonumber), estimated_cost_usd:($cost|tonumber), status:$status}'
|
||||||
}
|
}
|
||||||
|
|
||||||
# Get conversation messages
|
# Get conversation messages
|
||||||
@@ -426,13 +489,18 @@ case "$method:$path" in
|
|||||||
result=$(list_conversations)
|
result=$(list_conversations)
|
||||||
send_json "$result"
|
send_json "$result"
|
||||||
;;
|
;;
|
||||||
|
GET:/conversations/*/usage)
|
||||||
|
conv_id=$(echo "$path" | sed 's|/conversations/\([^/]*\)/usage|\1|')
|
||||||
|
result=$(get_conversation_usage "$conv_id")
|
||||||
|
send_json "$result"
|
||||||
|
;;
|
||||||
GET:/conversations/*)
|
GET:/conversations/*)
|
||||||
conv_id="${path#/conversations/}"
|
conv_id="${path#/conversations/}"
|
||||||
result=$(get_conversation "$conv_id")
|
result=$(get_conversation "$conv_id")
|
||||||
send_json "$result"
|
send_json "$result"
|
||||||
;;
|
;;
|
||||||
*)
|
*)
|
||||||
send_error "404 Not Found" "Endpoints: GET /personalities, GET /instances, POST /instances, GET /conversations"
|
send_error "404 Not Found" "Endpoints: GET /personalities, GET /instances, POST /instances, GET /conversations, GET /conversations/{id}/usage"
|
||||||
;;
|
;;
|
||||||
esac
|
esac
|
||||||
HANDLER
|
HANDLER
|
||||||
|
|||||||
Reference in New Issue
Block a user