49 lines
1.1 KiB
Python
49 lines
1.1 KiB
Python
from dataclasses import dataclass
|
|
from typing import Any, List, Literal, Optional, TypedDict
|
|
|
|
|
|
class _FunctionToolCall(TypedDict):
|
|
function_name: str
|
|
arguments: Optional[dict[str, Any]]
|
|
|
|
|
|
class ToolCall(TypedDict):
|
|
"""Represents a tool call in the AI model."""
|
|
|
|
id: str
|
|
function: _FunctionToolCall
|
|
type: Literal["function"]
|
|
|
|
|
|
class CompletionMessage(TypedDict):
|
|
"""Represents a message in the AI model."""
|
|
|
|
content: Any
|
|
role: str = "assistant"
|
|
|
|
|
|
@dataclass
|
|
class PromptEvent:
|
|
"""Represents an input event for the AI model."""
|
|
|
|
content: Any
|
|
role: str = "user"
|
|
tool_calls: Optional[List[ToolCall]] = None
|
|
|
|
|
|
@dataclass
|
|
class CompletionEvent:
|
|
"""Represents a completion event for the AI model."""
|
|
|
|
index: int
|
|
message: CompletionMessage
|
|
finish_reason: str = "unknown"
|
|
tool_calls: Optional[List[ToolCall]] = None
|
|
|
|
@property
|
|
def total_tokens(self) -> Optional[int]:
|
|
"""Returns the total number of tokens used in the event."""
|
|
if self.input_tokens is None or self.output_tokens is None:
|
|
return None
|
|
return self.input_tokens + self.output_tokens
|