codemaster-cli 2.2.0__py3-none-any.whl

vibe/core/types.py

@@ -0,0 +1,405 @@
+from __future__ import annotations
+
+from abc import ABC
+from collections import OrderedDict
+from collections.abc import Awaitable, Callable
+import copy
+from enum import StrEnum, auto
+from typing import TYPE_CHECKING, Annotated, Any, Literal
+from uuid import uuid4
+
+if TYPE_CHECKING:
+    from vibe.core.tools.base import BaseTool
+else:
+    BaseTool = Any
+
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    PrivateAttr,
+    computed_field,
+    model_validator,
+)
+
+
+class AgentStats(BaseModel):
+    steps: int = 0
+    session_prompt_tokens: int = 0
+    session_completion_tokens: int = 0
+    tool_calls_agreed: int = 0
+    tool_calls_rejected: int = 0
+    tool_calls_failed: int = 0
+    tool_calls_succeeded: int = 0
+
+    context_tokens: int = 0
+
+    last_turn_prompt_tokens: int = 0
+    last_turn_completion_tokens: int = 0
+    last_turn_duration: float = 0.0
+    tokens_per_second: float = 0.0
+
+    input_price_per_million: float = 0.0
+    output_price_per_million: float = 0.0
+
+    _listeners: dict[str, Callable[[AgentStats], None]] = PrivateAttr(
+        default_factory=dict
+    )
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        super().__setattr__(name, value)
+        if name in self._listeners:
+            self._listeners[name](self)
+
+    def trigger_listeners(self) -> None:
+        for listener in self._listeners.values():
+            listener(self)
+
+    def add_listener(
+        self, attr_name: str, listener: Callable[[AgentStats], None]
+    ) -> None:
+        self._listeners[attr_name] = listener
+
+    @computed_field
+    @property
+    def session_total_llm_tokens(self) -> int:
+        return self.session_prompt_tokens + self.session_completion_tokens
+
+    @computed_field
+    @property
+    def last_turn_total_tokens(self) -> int:
+        return self.last_turn_prompt_tokens + self.last_turn_completion_tokens
+
+    @computed_field
+    @property
+    def session_cost(self) -> float:
+        """Calculate the total session cost in dollars based on token usage and pricing.
+
+        NOTE: This is a rough estimate and is worst-case scenario.
+        The actual cost may be lower due to prompt caching.
+        If the model changes mid-session, this uses current pricing for all tokens.
+        """
+        input_cost = (
+            self.session_prompt_tokens / 1_000_000
+        ) * self.input_price_per_million
+        output_cost = (
+            self.session_completion_tokens / 1_000_000
+        ) * self.output_price_per_million
+        return input_cost + output_cost
+
+    def update_pricing(self, input_price: float, output_price: float) -> None:
+        """Update pricing info when model changes.
+
+        NOTE: session_cost will be recalculated using new pricing for all
+        accumulated tokens. This is a known approximation when models change.
+        This should not be a big issue, pricing is only used for max_price which is in
+        programmatic mode, so user should not update models there.
+        """
+        self.input_price_per_million = input_price
+        self.output_price_per_million = output_price
+
+    def reset_context_state(self) -> None:
+        """Reset context-related fields while preserving cumulative session stats.
+
+        Used after config reload or similar operations where the context
+        changes but we want to preserve session totals.
+        """
+        self.context_tokens = 0
+        self.last_turn_prompt_tokens = 0
+        self.last_turn_completion_tokens = 0
+        self.last_turn_duration = 0.0
+        self.tokens_per_second = 0.0
+
+
+class SessionInfo(BaseModel):
+    session_id: str
+    start_time: str
+    message_count: int
+    stats: AgentStats
+    save_dir: str
+
+
+class SessionMetadata(BaseModel):
+    session_id: str
+    start_time: str
+    end_time: str | None
+    git_commit: str | None
+    git_branch: str | None
+    environment: dict[str, str | None]
+    username: str
+
+
+StrToolChoice = Literal["auto", "none", "any", "required"]
+
+
+class AvailableFunction(BaseModel):
+    name: str
+    description: str
+    parameters: dict[str, Any]
+
+
+class AvailableTool(BaseModel):
+    type: Literal["function"] = "function"
+    function: AvailableFunction
+
+
+class FunctionCall(BaseModel):
+    name: str | None = None
+    arguments: str | None = None
+
+
+class ToolCall(BaseModel):
+    id: str | None = None
+    index: int | None = None
+    function: FunctionCall = Field(default_factory=FunctionCall)
+    type: str = "function"
+
+
+def _content_before(v: Any) -> str:
+    if isinstance(v, str):
+        return v
+    if isinstance(v, list):
+        parts: list[str] = []
+        for p in v:
+            if isinstance(p, dict) and isinstance(p.get("text"), str):
+                parts.append(p["text"])
+            else:
+                parts.append(str(p))
+        return "\n".join(parts)
+    return str(v)
+
+
+Content = Annotated[str, BeforeValidator(_content_before)]
+
+
+class Role(StrEnum):
+    system = auto()
+    user = auto()
+    assistant = auto()
+    tool = auto()
+
+
+class ApprovalResponse(StrEnum):
+    YES = "y"
+    NO = "n"
+
+
+class LLMMessage(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    role: Role
+    content: Content | None = None
+    reasoning_content: Content | None = None
+    reasoning_signature: str | None = None
+    tool_calls: list[ToolCall] | None = None
+    name: str | None = None
+    tool_call_id: str | None = None
+    message_id: str | None = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def _from_any(cls, v: Any) -> dict[str, Any] | Any:
+        if isinstance(v, dict):
+            v.setdefault("content", "")
+            v.setdefault("role", "assistant")
+            if "message_id" not in v and v.get("role") != "tool":
+                v["message_id"] = str(uuid4())
+            return v
+        role = str(getattr(v, "role", "assistant"))
+        return {
+            "role": role,
+            "content": getattr(v, "content", ""),
+            "reasoning_content": getattr(v, "reasoning_content", None),
+            "reasoning_signature": getattr(v, "reasoning_signature", None),
+            "tool_calls": getattr(v, "tool_calls", None),
+            "name": getattr(v, "name", None),
+            "tool_call_id": getattr(v, "tool_call_id", None),
+            "message_id": getattr(v, "message_id", None)
+            or (str(uuid4()) if role != "tool" else None),
+        }
+
+    def __add__(self, other: LLMMessage) -> LLMMessage:
+        """Careful: this is not commutative!"""
+        if self.role != other.role:
+            raise ValueError("Can't accumulate messages with different roles")
+
+        if self.name != other.name:
+            raise ValueError("Can't accumulate messages with different names")
+
+        if self.tool_call_id != other.tool_call_id:
+            raise ValueError("Can't accumulate messages with different tool_call_ids")
+
+        content = (self.content or "") + (other.content or "")
+        if not content:
+            content = None
+
+        reasoning_content = (self.reasoning_content or "") + (
+            other.reasoning_content or ""
+        )
+        if not reasoning_content:
+            reasoning_content = None
+
+        reasoning_signature = (self.reasoning_signature or "") + (
+            other.reasoning_signature or ""
+        )
+        if not reasoning_signature:
+            reasoning_signature = None
+
+        tool_calls_map = OrderedDict[int, ToolCall]()
+        for tool_calls in [self.tool_calls or [], other.tool_calls or []]:
+            for tc in tool_calls:
+                if tc.index is None:
+                    raise ValueError("Tool call chunk missing index")
+                if tc.index not in tool_calls_map:
+                    tool_calls_map[tc.index] = copy.deepcopy(tc)
+                else:
+                    existing_name = tool_calls_map[tc.index].function.name
+                    new_name = tc.function.name
+                    if existing_name and new_name and existing_name != new_name:
+                        raise ValueError(
+                            "Can't accumulate messages with different tool call names"
+                        )
+                    if new_name and not existing_name:
+                        tool_calls_map[tc.index].function.name = new_name
+                    new_args = (tool_calls_map[tc.index].function.arguments or "") + (
+                        tc.function.arguments or ""
+                    )
+                    tool_calls_map[tc.index].function.arguments = new_args
+
+        return LLMMessage(
+            role=self.role,
+            content=content,
+            reasoning_content=reasoning_content,
+            reasoning_signature=reasoning_signature,
+            tool_calls=list(tool_calls_map.values()) or None,
+            name=self.name,
+            tool_call_id=self.tool_call_id,
+            message_id=self.message_id,
+        )
+
+
+class LLMUsage(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+
+    def __add__(self, other: LLMUsage) -> LLMUsage:
+        return LLMUsage(
+            prompt_tokens=self.prompt_tokens + other.prompt_tokens,
+            completion_tokens=self.completion_tokens + other.completion_tokens,
+        )
+
+
+class LLMChunk(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    message: LLMMessage
+    usage: LLMUsage | None = None
+
+    def __add__(self, other: LLMChunk) -> LLMChunk:
+        if self.usage is None and other.usage is None:
+            new_usage = None
+        else:
+            new_usage = (self.usage or LLMUsage()) + (other.usage or LLMUsage())
+        return LLMChunk(message=self.message + other.message, usage=new_usage)
+
+
+class BaseEvent(BaseModel, ABC):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+
+class UserMessageEvent(BaseEvent):
+    content: str
+    message_id: str
+
+
+class AssistantEvent(BaseEvent):
+    content: str
+    stopped_by_middleware: bool = False
+    message_id: str | None = None
+
+    def __add__(self, other: AssistantEvent) -> AssistantEvent:
+        return AssistantEvent(
+            content=self.content + other.content,
+            stopped_by_middleware=self.stopped_by_middleware
+            or other.stopped_by_middleware,
+            message_id=self.message_id or other.message_id,
+        )
+
+
+class ReasoningEvent(BaseEvent):
+    content: str
+    message_id: str | None = None
+
+
+class ToolCallEvent(BaseEvent):
+    tool_name: str
+    tool_class: type[BaseTool]
+    args: BaseModel
+    tool_call_id: str
+
+
+class ToolResultEvent(BaseEvent):
+    tool_name: str
+    tool_class: type[BaseTool] | None
+    result: BaseModel | None = None
+    error: str | None = None
+    skipped: bool = False
+    skip_reason: str | None = None
+    duration: float | None = None
+    tool_call_id: str
+
+
+class ToolStreamEvent(BaseEvent):
+    tool_name: str
+    message: str
+    tool_call_id: str
+
+
+class CompactStartEvent(BaseEvent):
+    current_context_tokens: int
+    threshold: int
+    # WORKAROUND: Using tool_call to communicate compact events to the client.
+    # This should be revisited when the ACP protocol defines how compact events
+    # should be represented.
+    # [RFD](https://agentclientprotocol.com/rfds/session-usage)
+    tool_call_id: str
+
+
+class CompactEndEvent(BaseEvent):
+    old_context_tokens: int
+    new_context_tokens: int
+    summary_length: int
+    # WORKAROUND: Using tool_call to communicate compact events to the client.
+    # This should be revisited when the ACP protocol defines how compact events
+    # should be represented.
+    # [RFD](https://agentclientprotocol.com/rfds/session-usage)
+    tool_call_id: str
+
+
+class OutputFormat(StrEnum):
+    TEXT = auto()
+    JSON = auto()
+    STREAMING = auto()
+
+
+type AsyncApprovalCallback = Callable[
+    [str, BaseModel, str], Awaitable[tuple[ApprovalResponse, str | None]]
+]
+
+type SyncApprovalCallback = Callable[
+    [str, BaseModel, str], tuple[ApprovalResponse, str | None]
+]
+
+type ApprovalCallback = AsyncApprovalCallback | SyncApprovalCallback
+
+type UserInputCallback = Callable[[BaseModel], Awaitable[BaseModel]]
+
+
+class RateLimitError(Exception):
+    def __init__(self, provider: str, model: str) -> None:
+        self.provider = provider
+        self.model = model
+        super().__init__(
+            "Rate limits exceeded. Please wait a moment before trying again."
+        )