wehub-resource-sync
80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Shared models and utilities for API responses."""
|
|
|
|
import time
|
|
import uuid
|
|
from enum import Enum
|
|
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class IDPrefix(str, Enum):
|
|
"""Prefixes for generated IDs."""
|
|
|
|
CHAT_COMPLETION = "chatcmpl"
|
|
COMPLETION = "cmpl"
|
|
MESSAGE = "msg"
|
|
EMBEDDING = "emb"
|
|
RERANK = "rerank"
|
|
RESPONSE = "resp"
|
|
FUNCTION_CALL = "fc"
|
|
REASONING = "rs"
|
|
|
|
|
|
def generate_id(prefix: IDPrefix, length: int = 8) -> str:
|
|
"""Generate a unique ID with the given prefix.
|
|
|
|
Args:
|
|
prefix: The ID prefix to use
|
|
length: Length of the random suffix (default 8)
|
|
|
|
Returns:
|
|
Generated ID string (e.g., "chatcmpl-abc12345")
|
|
"""
|
|
if prefix == IDPrefix.MESSAGE:
|
|
# Anthropic style: msg_<24-char-hex>
|
|
return f"msg_{uuid.uuid4().hex[:24]}"
|
|
if prefix == IDPrefix.RESPONSE:
|
|
return f"resp_{uuid.uuid4().hex[:24]}"
|
|
if prefix == IDPrefix.FUNCTION_CALL:
|
|
return f"fc_{uuid.uuid4().hex[:8]}"
|
|
if prefix == IDPrefix.REASONING:
|
|
return f"rs_{uuid.uuid4().hex[:24]}"
|
|
return f"{prefix.value}-{uuid.uuid4().hex[:length]}"
|
|
|
|
|
|
def get_unix_timestamp() -> int:
|
|
"""Get current Unix timestamp.
|
|
|
|
Returns:
|
|
Current time as Unix timestamp (integer seconds since epoch)
|
|
"""
|
|
return int(time.time())
|
|
|
|
|
|
class BaseUsage(BaseModel):
|
|
"""Base class for token usage statistics.
|
|
|
|
This provides a foundation for both OpenAI-style (prompt_tokens/completion_tokens)
|
|
and Anthropic-style (input_tokens/output_tokens) usage tracking.
|
|
"""
|
|
|
|
prompt_tokens: int = 0
|
|
completion_tokens: int = 0
|
|
total_tokens: int = 0
|
|
input_tokens: int = 0
|
|
output_tokens: int = 0
|
|
|
|
def model_post_init(self, __context) -> None:
|
|
"""Calculate total_tokens and sync Anthropic-style aliases."""
|
|
if self.total_tokens == 0 and (self.prompt_tokens > 0 or self.completion_tokens > 0):
|
|
object.__setattr__(
|
|
self,
|
|
"total_tokens",
|
|
self.prompt_tokens + self.completion_tokens,
|
|
)
|
|
if self.input_tokens == 0 and self.prompt_tokens > 0:
|
|
object.__setattr__(self, "input_tokens", self.prompt_tokens)
|
|
if self.output_tokens == 0 and self.completion_tokens > 0:
|
|
object.__setattr__(self, "output_tokens", self.completion_tokens)
|