File size: 3,166 Bytes
5ef3875 | 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 | """Provider-neutral LLM interface.
The agent only speaks in these types. A provider translates them to and from
its own wire format. Assistant messages carry an opaque ``provider_data``
blob so a provider can replay its native content (e.g. thinking blocks)
exactly as it produced it; other providers ignore it and fall back to the
neutral text/tool_calls fields.
"""
from __future__ import annotations
import abc
import enum
from dataclasses import dataclass, field
from typing import Any, Callable, Union
@dataclass(frozen=True)
class ToolSpec:
name: str
description: str
parameters: dict[str, Any]
@dataclass
class ToolCall:
id: str
name: str
arguments: Any # normally a dict; validated before execution
@dataclass
class UserMessage:
text: str
@dataclass
class AssistantMessage:
text: str = ""
tool_calls: list[ToolCall] = field(default_factory=list)
provider: str | None = None
provider_data: Any = None
@dataclass
class ToolResultBlock:
tool_call_id: str
content: str
is_error: bool = False
@dataclass
class ToolResultsMessage:
results: list[ToolResultBlock]
Message = Union[UserMessage, AssistantMessage, ToolResultsMessage]
class StopReason(enum.Enum):
END_TURN = "end_turn"
TOOL_USE = "tool_use"
MAX_TOKENS = "max_tokens"
REFUSAL = "refusal"
PAUSE = "pause"
OTHER = "other"
@dataclass
class Usage:
input_tokens: int = 0
output_tokens: int = 0
cache_read_tokens: int = 0
cache_write_tokens: int = 0
def add(self, other: "Usage") -> None:
self.input_tokens += other.input_tokens
self.output_tokens += other.output_tokens
self.cache_read_tokens += other.cache_read_tokens
self.cache_write_tokens += other.cache_write_tokens
@dataclass
class LLMResponse:
message: AssistantMessage
stop_reason: StopReason
usage: Usage = field(default_factory=Usage)
stop_detail: str | None = None # e.g. refusal explanation
@dataclass
class StreamCallbacks:
"""Optional hooks a UI can use to show progress while a response streams."""
on_text: Callable[[str], None] | None = None
on_tool_call_start: Callable[[str], None] | None = None
def text(self, chunk: str) -> None:
if self.on_text:
self.on_text(chunk)
def tool_call_start(self, name: str) -> None:
if self.on_tool_call_start:
self.on_tool_call_start(name)
class LLMProvider(abc.ABC):
"""A chat model that supports tool calling."""
name: str = "provider"
# False for models that can only answer in text (no tool calling).
supports_tools: bool = True
def __init__(self, model: str) -> None:
self.model = model
@property
def display_name(self) -> str:
return self.model
@abc.abstractmethod
def complete(
self,
*,
system: str,
messages: list[Message],
tools: list[ToolSpec],
callbacks: StreamCallbacks | None = None,
) -> LLMResponse:
"""Generate the next assistant turn.
Raises :class:`ycode.errors.LLMError` subclasses for provider failures.
"""
|