File size: 7,657 Bytes
42da440 | 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 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 | import os
import yaml
import uuid
import json
import asyncio
from typing import Callable, Dict, List, Any, Union
def create_log_observer(
log_dir: str, active: bool = False
) -> (
Callable
): # TODO maybe put these in a class so that I can have type checking and warn when I am putting an input observer in an output observer list?
"""
Creates an output observer that logs all inputs and outputs to files.
Args:
log_dir: Directory path where logs will be saved
Returns:
A callback function that logs inputs and outputs
"""
full_output_path = os.path.join(log_dir, "debug_outputs")
def log_observer(
input_data: Union[str, List[Dict[str, str]]],
output: str,
completion_mode: bool,
*args,
**kwargs,
) -> None:
"""
Logs the input and output to a file with a UUID.
Args:
input_data: Either a prompt string (completion mode) or messages list (chat mode)
output: The model's response
completion_mode: Whether this was a completion (True) or chat (False) request
"""
if not active:
return
# Extract prefix from input_data
prefix = ""
if completion_mode:
# For completion mode, get first 10 chars from the string
prefix = input_data[:25].lower().replace(" ", "_")
else:
# For chat mode, get first 10 chars from the first message content
if input_data and len(input_data) > 0 and "content" in input_data[0]:
prefix = input_data[0]["content"][:25].lower().replace(" ", "_")
# Create log ID with prefix
log_id = f"{prefix}_{str(uuid.uuid4())}"
os.makedirs(full_output_path, exist_ok=True)
if completion_mode:
print(f"Saving output to {log_id}.txt")
# For completion mode, save as .txt
file_path = os.path.join(full_output_path, f"{log_id}.txt")
with open(file_path, "w", encoding="utf-8") as f:
f.write(input_data + output)
else:
print(f"Saving output to {log_id}.yaml")
# For chat mode, save as .yaml
file_path = os.path.join(full_output_path, f"{log_id}.yaml")
# Create a full conversation history including the assistant's response
full_conversation = input_data.copy() # Copy the input messages
full_conversation.append({"role": "assistant", "content": output})
with open(file_path, "w", encoding="utf-8") as f:
yaml.dump(
full_conversation, f, default_flow_style=False, allow_unicode=True
)
return log_observer
def create_input_token_counter(
counter: Dict[str, float],
cost_per_million: float,
count_tokens_fn: Callable,
persistence_path: str = None,
) -> Callable:
"""
Creates an input observer that counts tokens and calculates cost.
Args:
counter: A dictionary to track token counts and costs
cost_per_million: Cost per million tokens
count_tokens_fn: Function to count tokens in text
persistence_path: Optional path to persist counter data to a JSON file
Returns:
A callback function that updates the counter
"""
# Initialize lock for file access
json_path = None
if persistence_path:
os.makedirs(os.path.dirname(persistence_path), exist_ok=True)
json_path = persistence_path
# Load existing counter data if available
if os.path.exists(json_path):
try:
with open(json_path, "r", encoding="utf-8") as f:
saved_data = json.load(f)
# Update the counter with saved values
counter.update(saved_data)
except (json.JSONDecodeError, FileNotFoundError):
# If file is corrupted or doesn't exist, start fresh
pass
def input_token_counter(
input_data: Union[str, List[Dict[str, str]]],
completion_mode: bool,
*args,
**kwargs,
) -> None:
"""
Counts tokens in the input and updates the counter.
Args:
input_data: Either a prompt string (completion mode) or messages list (chat mode)
completion_mode: Whether this was a completion (True) or chat (False) request
"""
if completion_mode:
# For completion mode, input is a string
token_count = count_tokens_fn(input_data)
else:
# For chat mode, input is a list of message dicts
token_count = 0
for message in input_data:
token_count += count_tokens_fn(message["content"])
# Update counter
if "input_tokens" not in counter:
counter["input_tokens"] = 0
if "input_cost" not in counter:
counter["input_cost"] = 0.0
counter["input_tokens"] += token_count
counter["input_cost"] += (token_count / 1_000_000) * cost_per_million
# Persist counter to file if path was provided
if json_path:
with open(json_path, "w", encoding="utf-8") as f:
json.dump(counter, f)
return input_token_counter
def create_output_token_counter(
counter: Dict[str, float],
cost_per_million: float,
count_tokens_fn: Callable,
persistence_path: str = None,
) -> Callable:
"""
Creates an output observer that counts tokens and calculates cost.
Args:
counter: A dictionary to track token counts and costs
cost_per_million: Cost per million tokens
count_tokens_fn: Function to count tokens in text
persistence_path: Optional path to persist counter data to a JSON file
Returns:
A callback function that updates the counter
"""
# Initialize lock for file access
json_path = None
if persistence_path:
os.makedirs(os.path.dirname(persistence_path), exist_ok=True)
json_path = persistence_path
# Load existing counter data if available
if os.path.exists(json_path):
try:
with open(json_path, "r", encoding="utf-8") as f:
saved_data = json.load(f)
# Update the counter with saved values
counter.update(saved_data)
except (json.JSONDecodeError, FileNotFoundError):
# If file is corrupted or doesn't exist, start fresh
pass
def output_token_counter(
input_data: Union[str, List[Dict[str, str]]],
output: str,
completion_mode: bool,
*args,
**kwargs,
) -> None:
"""
Counts tokens in the output and updates the counter.
Args:
input_data: Either a prompt string (completion mode) or messages list (chat mode)
output: The model's response
completion_mode: Whether this was a completion (True) or chat (False) request
"""
token_count = count_tokens_fn(output)
# Update counter
if "output_tokens" not in counter:
counter["output_tokens"] = 0
if "output_cost" not in counter:
counter["output_cost"] = 0.0
counter["output_tokens"] += token_count
counter["output_cost"] += (token_count / 1_000_000) * cost_per_million
# Persist counter to file if path was provided
if json_path:
with open(json_path, "w", encoding="utf-8") as f:
json.dump(counter, f)
return output_token_counter
|