Spaces:
Paused
Paused
File size: 4,155 Bytes
f647629 |
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 |
"""
Data models for the Weave API.
This module defines the data structures used across the Weave API client.
"""
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional, Union
from pydantic import BaseModel, Field
class SortDirection(str, Enum):
"""Sort direction options."""
ASC = "asc"
DESC = "desc"
class FilterOperator(str, Enum):
"""Operators for query filters."""
EQUALS = "$eq"
GREATER_THAN = "$gt"
GREATER_THAN_EQUAL = "$gte"
LESS_THAN = "$lt"
LESS_THAN_EQUAL = "$lte"
CONTAINS = "$contains"
@dataclass
class TimeRange:
"""Time range for filtering traces."""
start: Optional[datetime] = None
end: Optional[datetime] = None
@dataclass
class AttributeFilter:
"""Filter for trace attributes."""
path: str
value: Any
operator: FilterOperator = FilterOperator.EQUALS
@dataclass
class QueryFilter:
"""Query filter parameters."""
display_name: Optional[str] = None
op_name: Optional[str] = None
op_name_contains: Optional[str] = None
trace_id: Optional[str] = None
trace_ids: Optional[List[str]] = None
call_ids: Optional[List[str]] = None
parent_ids: Optional[List[str]] = None
status: Optional[str] = None
time_range: Optional[TimeRange] = None
latency: Optional[Dict[str, float]] = None
attributes: Optional[List[AttributeFilter]] = None
has_exception: Optional[bool] = None
trace_roots_only: Optional[bool] = None
@dataclass
class QueryParams:
"""Parameters for a Weave trace query."""
entity_name: str
project_name: str
filters: Optional[QueryFilter] = None
sort_by: str = "started_at"
sort_direction: SortDirection = SortDirection.DESC
limit: Optional[int] = None
offset: int = 0
include_costs: bool = True
include_feedback: bool = True
columns: Optional[List[str]] = None
expand_columns: Optional[List[str]] = None
truncate_length: Optional[int] = 200
return_full_data: bool = False
metadata_only: bool = False
@dataclass
class TraceCost:
"""Cost information for a model in a trace."""
prompt_tokens: int = 0
completion_tokens: int = 0
prompt_tokens_total_cost: float = 0.0
completion_tokens_total_cost: float = 0.0
total_cost: float = 0.0
class TokenCounts(BaseModel):
"""Token count information for traces."""
total_tokens: int = 0
input_tokens: int = 0
output_tokens: int = 0
average_tokens_per_trace: Union[int, float] = 0
class TraceMetadata(BaseModel):
"""Metadata about a collection of traces."""
total_traces: int = 0
token_counts: Dict[str, Union[int, float]] = Field(default_factory=dict)
time_range: Dict[str, Optional[datetime]] = Field(default_factory=dict)
status_summary: Dict[str, int] = Field(default_factory=dict)
op_distribution: Dict[str, int] = Field(default_factory=dict)
class WeaveTrace(BaseModel):
"""Representation of a Weave trace."""
id: str
project_id: str
op_name: str
trace_id: str
started_at: datetime
display_name: Optional[str] = None
parent_id: Optional[str] = None
ended_at: Optional[datetime] = None
inputs: Optional[Dict[str, Any]] = None
output: Optional[Any] = None
exception: Optional[str] = None
attributes: Dict[str, Any] = Field(default_factory=dict)
summary: Dict[str, Any] = Field(default_factory=dict)
costs: Dict[str, TraceCost] = Field(default_factory=dict)
feedback: Dict[str, Any] = Field(default_factory=dict)
status: Optional[str] = None
latency_ms: Optional[int] = None
wb_run_id: Optional[str] = None
wb_user_id: Optional[str] = None
deleted_at: Optional[datetime] = None
class Config:
"""Pydantic model configuration."""
arbitrary_types_allowed = True
class QueryResult(BaseModel):
"""Result of a Weave trace query."""
metadata: TraceMetadata
traces: Optional[List[Union[WeaveTrace, Dict[str, Any]]]] = None
class Config:
"""Pydantic model configuration."""
arbitrary_types_allowed = True
|