File size: 13,559 Bytes
aca8ab4 |
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 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 |
"""
Trace reader for querying LangFuse observability data.
Provides Python API for programmatic access to traces, spans, and generations.
"""
import logging
from typing import List, Optional, Dict, Any
from datetime import datetime, timedelta
from pydantic import BaseModel, Field
from utils.langfuse_client import get_langfuse_client, is_langfuse_enabled
logger = logging.getLogger(__name__)
class TraceInfo(BaseModel):
"""Pydantic model for trace information."""
id: str
name: str
user_id: Optional[str] = None
session_id: Optional[str] = None
timestamp: datetime
metadata: Dict[str, Any] = Field(default_factory=dict)
input: Optional[Any] = None
output: Optional[Any] = None
duration_ms: Optional[float] = None
total_cost: Optional[float] = None
token_usage: Dict[str, int] = Field(default_factory=dict)
class SpanInfo(BaseModel):
"""Pydantic model for span information."""
id: str
trace_id: str
name: str
start_time: datetime
end_time: Optional[datetime] = None
duration_ms: Optional[float] = None
metadata: Dict[str, Any] = Field(default_factory=dict)
input: Optional[Any] = None
output: Optional[Any] = None
level: str = "DEFAULT"
class GenerationInfo(BaseModel):
"""Pydantic model for LLM generation information."""
id: str
trace_id: str
name: str
model: Optional[str] = None
prompt: Optional[str] = None
completion: Optional[str] = None
usage: Dict[str, int] = Field(default_factory=dict)
cost: Optional[float] = None
start_time: datetime
end_time: Optional[datetime] = None
duration_ms: Optional[float] = None
metadata: Dict[str, Any] = Field(default_factory=dict)
class TraceReader:
"""
Read and query LangFuse traces programmatically.
Usage:
reader = TraceReader()
traces = reader.get_traces(limit=10)
trace = reader.get_trace_by_id("trace-123")
agent_traces = reader.filter_by_agent("retriever_agent")
"""
def __init__(self):
"""Initialize trace reader with LangFuse client."""
if not is_langfuse_enabled():
logger.warning("LangFuse is not enabled. TraceReader will return empty results.")
self.client = None
else:
self.client = get_langfuse_client()
logger.info("TraceReader initialized with LangFuse client")
def get_traces(
self,
limit: int = 50,
user_id: Optional[str] = None,
session_id: Optional[str] = None,
from_timestamp: Optional[datetime] = None,
to_timestamp: Optional[datetime] = None,
) -> List[TraceInfo]:
"""
Get traces with optional filters.
Args:
limit: Maximum number of traces to return
user_id: Filter by user ID
session_id: Filter by session ID
from_timestamp: Filter traces after this timestamp
to_timestamp: Filter traces before this timestamp
Returns:
List of TraceInfo objects
"""
if not self.client:
logger.warning("LangFuse client not available")
return []
try:
# Build filter params
params = {"limit": limit}
if user_id:
params["user_id"] = user_id
if session_id:
params["session_id"] = session_id
if from_timestamp:
params["from_timestamp"] = from_timestamp
if to_timestamp:
params["to_timestamp"] = to_timestamp
# Fetch traces from LangFuse
traces_data = self.client.get_traces(**params)
# Convert to TraceInfo objects
traces = []
for trace in traces_data.data:
trace_info = TraceInfo(
id=trace.id,
name=trace.name,
user_id=trace.user_id,
session_id=trace.session_id,
timestamp=trace.timestamp,
metadata=trace.metadata or {},
input=trace.input,
output=trace.output,
duration_ms=self._calculate_duration(trace),
total_cost=getattr(trace, "total_cost", None),
token_usage=self._extract_token_usage(trace),
)
traces.append(trace_info)
logger.info(f"Retrieved {len(traces)} traces")
return traces
except Exception as e:
logger.error(f"Error fetching traces: {e}")
return []
def get_trace_by_id(self, trace_id: str) -> Optional[TraceInfo]:
"""
Get a specific trace by ID.
Args:
trace_id: Trace identifier
Returns:
TraceInfo object or None if not found
"""
if not self.client:
logger.warning("LangFuse client not available")
return None
try:
trace = self.client.get_trace(trace_id)
if not trace:
logger.warning(f"Trace {trace_id} not found")
return None
trace_info = TraceInfo(
id=trace.id,
name=trace.name,
user_id=trace.user_id,
session_id=trace.session_id,
timestamp=trace.timestamp,
metadata=trace.metadata or {},
input=trace.input,
output=trace.output,
duration_ms=self._calculate_duration(trace),
total_cost=getattr(trace, "total_cost", None),
token_usage=self._extract_token_usage(trace),
)
logger.info(f"Retrieved trace {trace_id}")
return trace_info
except Exception as e:
logger.error(f"Error fetching trace {trace_id}: {e}")
return None
def filter_by_agent(
self,
agent_name: str,
limit: int = 50,
from_timestamp: Optional[datetime] = None,
) -> List[SpanInfo]:
"""
Filter traces by agent name.
Args:
agent_name: Name of the agent (e.g., "retriever_agent", "analyzer_agent")
limit: Maximum number of results
from_timestamp: Filter traces after this timestamp
Returns:
List of SpanInfo objects for the specified agent
"""
if not self.client:
logger.warning("LangFuse client not available")
return []
try:
# Get observations filtered by name
params = {"limit": limit, "name": agent_name, "type": "SPAN"}
if from_timestamp:
params["from_timestamp"] = from_timestamp
observations = self.client.get_observations(**params)
spans = []
for obs in observations.data:
span_info = SpanInfo(
id=obs.id,
trace_id=obs.trace_id,
name=obs.name,
start_time=obs.start_time,
end_time=obs.end_time,
duration_ms=self._calculate_duration(obs),
metadata=obs.metadata or {},
input=obs.input,
output=obs.output,
level=getattr(obs, "level", "DEFAULT"),
)
spans.append(span_info)
logger.info(f"Retrieved {len(spans)} spans for agent '{agent_name}'")
return spans
except Exception as e:
logger.error(f"Error filtering by agent {agent_name}: {e}")
return []
def filter_by_date_range(
self,
from_date: datetime,
to_date: datetime,
limit: int = 100,
) -> List[TraceInfo]:
"""
Filter traces by date range.
Args:
from_date: Start date
to_date: End date
limit: Maximum number of traces
Returns:
List of TraceInfo objects within date range
"""
return self.get_traces(
limit=limit,
from_timestamp=from_date,
to_timestamp=to_date,
)
def get_generations(
self,
trace_id: Optional[str] = None,
limit: int = 50,
) -> List[GenerationInfo]:
"""
Get LLM generations (optionally filtered by trace).
Args:
trace_id: Optional trace ID to filter generations
limit: Maximum number of generations
Returns:
List of GenerationInfo objects
"""
if not self.client:
logger.warning("LangFuse client not available")
return []
try:
params = {"limit": limit, "type": "GENERATION"}
if trace_id:
params["trace_id"] = trace_id
observations = self.client.get_observations(**params)
generations = []
for obs in observations.data:
gen_info = GenerationInfo(
id=obs.id,
trace_id=obs.trace_id,
name=obs.name,
model=getattr(obs, "model", None),
prompt=getattr(obs, "input", None),
completion=getattr(obs, "output", None),
usage=self._extract_token_usage(obs),
cost=getattr(obs, "calculated_total_cost", None),
start_time=obs.start_time,
end_time=obs.end_time,
duration_ms=self._calculate_duration(obs),
metadata=obs.metadata or {},
)
generations.append(gen_info)
logger.info(f"Retrieved {len(generations)} generations")
return generations
except Exception as e:
logger.error(f"Error fetching generations: {e}")
return []
def export_traces_to_json(
self,
traces: List[TraceInfo],
output_file: str,
) -> bool:
"""
Export traces to JSON file.
Args:
traces: List of TraceInfo objects
output_file: Path to output JSON file
Returns:
True if successful, False otherwise
"""
try:
import json
data = [trace.dict() for trace in traces]
with open(output_file, 'w') as f:
json.dump(data, f, indent=2, default=str)
logger.info(f"Exported {len(traces)} traces to {output_file}")
return True
except Exception as e:
logger.error(f"Error exporting traces: {e}")
return False
def export_traces_to_csv(
self,
traces: List[TraceInfo],
output_file: str,
) -> bool:
"""
Export traces to CSV file.
Args:
traces: List of TraceInfo objects
output_file: Path to output CSV file
Returns:
True if successful, False otherwise
"""
try:
import csv
if not traces:
logger.warning("No traces to export")
return False
# Define CSV columns
fieldnames = [
"id", "name", "user_id", "session_id", "timestamp",
"duration_ms", "total_cost", "input_tokens", "output_tokens"
]
with open(output_file, 'w', newline='') as f:
writer = csv.DictWriter(f, fieldnames=fieldnames)
writer.writeheader()
for trace in traces:
row = {
"id": trace.id,
"name": trace.name,
"user_id": trace.user_id or "",
"session_id": trace.session_id or "",
"timestamp": trace.timestamp.isoformat(),
"duration_ms": trace.duration_ms or 0,
"total_cost": trace.total_cost or 0,
"input_tokens": trace.token_usage.get("input", 0),
"output_tokens": trace.token_usage.get("output", 0),
}
writer.writerow(row)
logger.info(f"Exported {len(traces)} traces to {output_file}")
return True
except Exception as e:
logger.error(f"Error exporting traces to CSV: {e}")
return False
# Helper methods
def _calculate_duration(self, obj: Any) -> Optional[float]:
"""Calculate duration in milliseconds from start and end times."""
try:
if hasattr(obj, 'start_time') and hasattr(obj, 'end_time') and obj.end_time:
duration = (obj.end_time - obj.start_time).total_seconds() * 1000
return duration
return None
except Exception:
return None
def _extract_token_usage(self, obj: Any) -> Dict[str, int]:
"""Extract token usage from observation."""
usage = {}
try:
if hasattr(obj, 'usage') and obj.usage:
usage["input"] = getattr(obj.usage, "prompt_tokens", 0) or getattr(obj.usage, "input", 0)
usage["output"] = getattr(obj.usage, "completion_tokens", 0) or getattr(obj.usage, "output", 0)
usage["total"] = getattr(obj.usage, "total_tokens", 0) or getattr(obj.usage, "total", 0)
except Exception:
pass
return usage
|