kosong 0.13.0__tar.gz

kosong-0.13.0/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.3
+Name: kosong
+Version: 0.13.0
+Summary: The building blocks of AI agent.
+Requires-Dist: dotenv>=0.9.9
+Requires-Dist: jsonschema>=4.25.1
+Requires-Dist: openai>=1.109.1
+Requires-Dist: pydantic>=2.11.9
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+
+# Kosong
+
+> Kosong is the emptiness.

kosong-0.13.0/README.md

@@ -0,0 +1,3 @@
+# Kosong
+
+> Kosong is the emptiness.

kosong-0.13.0/pyproject.toml

@@ -0,0 +1,40 @@
+[project]
+name = "kosong"
+version = "0.13.0"
+description = "The building blocks of AI agent."
+readme = "README.md"
+requires-python = ">=3.13"
+dependencies = [
+    "dotenv>=0.9.9",
+    "jsonschema>=4.25.1",
+    "openai>=1.109.1",
+    "pydantic>=2.11.9",
+]
+
+[dependency-groups]
+dev = [
+    "pyright>=1.1.405",
+    "pytest>=8.4.2",
+    "pytest-asyncio>=1.2.0",
+    "ruff>=0.13.1",
+]
+
+[build-system]
+requires = ["uv_build>=0.8.5,<0.9.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = ["kosong"]
+
+[tool.ruff]
+line-length = 100
+
+[tool.ruff.lint]
+select = [
+    "E",   # pycodestyle
+    "F",   # Pyflakes
+    "UP",  # pyupgrade
+    "B",   # flake8-bugbear
+    "SIM", # flake8-simplify
+    "I",   # isort
+]

kosong-0.13.0/src/kosong/__init__.py

@@ -0,0 +1,100 @@
+import asyncio
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+
+from kosong.base import generate
+from kosong.base.chat_provider import ChatProvider, StreamedMessagePart, TokenUsage
+from kosong.base.message import Message, ToolCall
+from kosong.chat_provider import ChatProviderError
+from kosong.tooling import ToolResult, ToolResultFuture, Toolset
+from kosong.utils.aio import Callback
+
+
+async def step(
+    chat_provider: ChatProvider,
+    system_prompt: str,
+    toolset: Toolset,
+    history: Sequence[Message],
+    *,
+    on_message_part: Callback[[StreamedMessagePart], None] | None = None,
+    on_tool_result: Callable[[ToolResult], None] | None = None,
+) -> "StepResult":
+    """
+    Run one "step". In one step, the function generates LLM response based on the given context for
+    exactly one time. All new message parts will be streamed to `on_message_part` in real-time if
+    provided. Tool calls will be handled by `context.toolset`. The combined message will be returned
+    in a `StepResult`. Depending on the toolset implementation, the tool calls may be handled
+    asynchronously and the results need to be fetched by `await step_result.tool_results()`.
+
+    The context will NOT be modified in this function.
+
+    The token usage will be returned in the `StepResult` if available.
+    """
+
+    tool_calls: list[ToolCall] = []
+    tool_result_futures: dict[str, ToolResultFuture] = {}
+
+    def future_done_callback(future: ToolResultFuture):
+        if on_tool_result:
+            try:
+                result = future.result()
+                on_tool_result(result)
+            except asyncio.CancelledError:
+                return
+
+    async def on_tool_call(tool_call: ToolCall):
+        tool_calls.append(tool_call)
+        result = toolset.handle(tool_call)
+
+        if isinstance(result, ToolResult):
+            future = ToolResultFuture()
+            future.add_done_callback(future_done_callback)
+            future.set_result(result)
+            tool_result_futures[tool_call.id] = future
+        else:
+            result.add_done_callback(future_done_callback)
+            tool_result_futures[tool_call.id] = result
+
+    try:
+        message, usage = await generate(
+            chat_provider,
+            system_prompt,
+            toolset.tools,
+            history,
+            on_message_part=on_message_part,
+            on_tool_call=on_tool_call,
+        )
+    except (ChatProviderError, asyncio.CancelledError):
+        # cancel all the futures to avoid hanging tasks
+        for future in tool_result_futures.values():
+            future.remove_done_callback(future_done_callback)
+            future.cancel()
+        raise
+
+    return StepResult(message, usage, tool_calls, tool_result_futures)
+
+
+@dataclass(frozen=True)
+class StepResult:
+    message: Message
+    """The combined message generated in this step."""
+
+    usage: TokenUsage | None
+    """The token usage of the generated message."""
+
+    tool_calls: list[ToolCall]
+    """All the tool calls generated in this step."""
+
+    tool_result_futures: dict[str, ToolResultFuture]
+    """The futures of the results of the spawned tool calls."""
+
+    async def tool_results(self) -> list[ToolResult]:
+        """All the tool results returned by corresponding tool calls."""
+        if not self.tool_result_futures:
+            return []
+
+        results: list[ToolResult] = []
+        for tool_call in self.tool_calls:
+            result = await self.tool_result_futures[tool_call.id]
+            results.append(result)
+        return results

kosong-0.13.0/src/kosong/base/__init__.py

@@ -0,0 +1,66 @@
+from collections.abc import Sequence
+
+from kosong.base.chat_provider import ChatProvider, StreamedMessagePart, TokenUsage
+from kosong.base.message import ContentPart, Message, TextPart, ToolCall
+from kosong.base.tool import Tool
+from kosong.utils.aio import Callback, callback
+
+
+async def generate(
+    chat_provider: ChatProvider,
+    system_prompt: str,
+    tools: Sequence[Tool],
+    history: Sequence[Message],
+    *,
+    on_message_part: Callback[[StreamedMessagePart], None] | None = None,
+    on_tool_call: Callback[[ToolCall], None] | None = None,
+) -> tuple[Message, TokenUsage | None]:
+    """
+    Generate one message based on the given context. The given context will remain untouched.
+
+    Parts of the message will be streamed to the given handlers:
+    - `on_message_part` will be called for each raw part which may be incomplete.
+    - `on_tool_call` will be called for each complete tool call.
+
+    The generated message and the token usage will be returned. All parts in the message are
+    guaranteed to be complete and merged as much as possible.
+    """
+    message = Message(role="assistant", content=[])
+    pending_part: StreamedMessagePart | None = None  # message part that is currently incomplete
+
+    stream = await chat_provider.generate(system_prompt, tools, history)
+    async for part in stream:
+        if on_message_part:
+            await callback(on_message_part, part.model_copy(deep=True))
+
+        if pending_part is None:
+            pending_part = part
+        elif not pending_part.merge_in_place(part):  # try merge into the pending part
+            # unmergeable part must push the pending part to the buffer
+            _message_append(message, pending_part)
+            if isinstance(pending_part, ToolCall) and on_tool_call:
+                await callback(on_tool_call, pending_part)
+            pending_part = part
+
+    # end of message
+    if pending_part is not None:
+        _message_append(message, pending_part)
+        if isinstance(pending_part, ToolCall) and on_tool_call:
+            await callback(on_tool_call, pending_part)
+
+    return message, stream.usage
+
+
+def _message_append(message: Message, part: StreamedMessagePart) -> None:
+    match part:
+        case ContentPart():
+            if isinstance(message.content, str):
+                message.content = [TextPart(text=message.content)]
+            message.content.append(part)
+        case ToolCall():
+            if message.tool_calls is None:
+                message.tool_calls = []
+            message.tool_calls.append(part)
+        case _:
+            # may be an orphaned `ToolCallPart`
+            return

kosong-0.13.0/src/kosong/base/chat_provider.py

@@ -0,0 +1,56 @@
+from collections.abc import AsyncIterator, Sequence
+from typing import NamedTuple, Protocol, runtime_checkable
+
+from kosong.base.message import ContentPart, Message, ToolCall, ToolCallPart
+from kosong.base.tool import Tool
+
+
+@runtime_checkable
+class ChatProvider(Protocol):
+    name: str
+    """
+    The name of the chat provider.
+    """
+
+    @property
+    def model_name(self) -> str:
+        """
+        The model name to use for the chat provider.
+        """
+        ...
+
+    async def generate(
+        self,
+        system_prompt: str,
+        tools: Sequence[Tool],
+        history: Sequence[Message],
+    ) -> "StreamedMessage":
+        """
+        Generate a new message based on the given system prompt, tools, and history.
+        """
+        ...
+
+
+type StreamedMessagePart = ContentPart | ToolCall | ToolCallPart
+
+
+@runtime_checkable
+class StreamedMessage(Protocol):
+    def __aiter__(self) -> AsyncIterator[StreamedMessagePart]:
+        """Create an async iterator from the stream."""
+        ...
+
+    @property
+    def usage(self) -> "TokenUsage | None":
+        """The usage of the streamed message."""
+        ...
+
+
+class TokenUsage(NamedTuple):
+    input: int
+    output: int
+    # TODO: support `cached`
+
+    @property
+    def total(self) -> int:
+        return self.input + self.output

kosong-0.13.0/src/kosong/base/message.py

@@ -0,0 +1,212 @@
+from abc import ABC
+from typing import Any, ClassVar, Literal, override
+
+from pydantic import BaseModel, GetCoreSchemaHandler, field_serializer
+from pydantic_core import core_schema
+
+
+class MergableMixin:
+    def merge_in_place(self, other: Any) -> bool:
+        """Merge the other part into the current part. Return True if the merge is successful."""
+        return False
+
+
+class ContentPart(BaseModel, ABC, MergableMixin):
+    """A part of a message content."""
+
+    __content_part_registry: ClassVar[dict[str, type["ContentPart"]]] = {}
+
+    type: str
+    ...  # to be added by subclasses
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+
+        invalid_subclass_error_msg = (
+            f"ContentPart subclass {cls.__name__} must have a `type` field of type `str`"
+        )
+
+        if not hasattr(cls, "type"):
+            raise ValueError(invalid_subclass_error_msg)
+
+        type_value = cls.type
+        if not isinstance(type_value, str):
+            raise ValueError(invalid_subclass_error_msg)
+
+        cls.__content_part_registry[type_value] = cls
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        # If we're dealing with the base ContentPart class, use custom validation
+        if cls.__name__ == "ContentPart":
+
+            def validate_content_part(value: Any) -> Any:
+                # if it's already an instance of a ContentPart subclass, return it
+                if hasattr(value, "__class__") and issubclass(value.__class__, cls):
+                    return value
+
+                # if it's a dict with a type field, dispatch to the appropriate subclass
+                if isinstance(value, dict) and "type" in value:
+                    type_value = value["type"]
+                    if not isinstance(type_value, str):
+                        raise ValueError(f"Cannot validate {value} as ContentPart")
+                    target_class = cls.__content_part_registry[type_value]
+                    return target_class.model_validate(value)
+
+                raise ValueError(f"Cannot validate {value} as ContentPart")
+
+            return core_schema.no_info_plain_validator_function(validate_content_part)
+
+        # for subclasses, use the default schema
+        return handler(source_type)
+
+
+class TextPart(ContentPart):
+    """
+    >>> TextPart(text="Hello, world!").model_dump()
+    {'type': 'text', 'text': 'Hello, world!'}
+    """
+
+    type: str = "text"
+    text: str
+
+    @override
+    def merge_in_place(self, other) -> bool:
+        if not isinstance(other, TextPart):
+            return False
+        self.text += other.text
+        return True
+
+
+class ThinkPart(ContentPart):
+    """
+    >>> ThinkPart(think="I think I need to think about this.").model_dump()
+    {'type': 'think', 'think': 'I think I need to think about this.'}
+    """
+
+    type: str = "think"
+    think: str
+
+    @override
+    def merge_in_place(self, other) -> bool:
+        if not isinstance(other, ThinkPart):
+            return False
+        self.think += other.think
+        return True
+
+
+class ImageURLPart(ContentPart):
+    """
+    >>> ImageURLPart(image_url=ImageURLPart.ImageURL(url="https://example.com/image.png")).model_dump()
+    {'type': 'image_url', 'image_url': {'url': 'https://example.com/image.png', 'id': None}}
+    """
+
+    class ImageURL(BaseModel):
+        url: str
+        """The URL of the image, can be data URI scheme like `data:image/png;base64,...`."""
+        id: str | None = None
+        """The ID of the image, to allow LLMs to distinguish different images."""
+
+    type: str = "image_url"
+    image_url: ImageURL
+
+
+class AudioURLPart(ContentPart):
+    """
+    >>> AudioURLPart(audio_url=AudioURLPart.AudioURL(url="https://example.com/audio.mp3")).model_dump()
+    {'type': 'audio_url', 'audio_url': {'url': 'https://example.com/audio.mp3', 'id': None}}
+    """
+
+    class AudioURL(BaseModel):
+        url: str
+        """The URL of the audio, can be data URI scheme like `data:audio/aac;base64,...`."""
+        id: str | None = None
+        """The ID of the audio, to allow LLMs to distinguish different audios."""
+
+    type: str = "audio_url"
+    audio_url: AudioURL
+
+
+class ToolCall(BaseModel, MergableMixin):
+    """
+    A tool call requested by the assistant.
+
+    >>> ToolCall(
+    ...     id="123",
+    ...     function=ToolCall.FunctionBody(
+    ...         name="function",
+    ...         arguments="{}"
+    ...     ),
+    ... ).model_dump()
+    {'type': 'function', 'id': '123', 'function': {'name': 'function', 'arguments': '{}'}}
+    """
+
+    class FunctionBody(BaseModel):
+        name: str
+        arguments: str | None
+
+    type: Literal["function"] = "function"
+
+    id: str
+    """The ID of the tool call."""
+    function: FunctionBody
+    """The function body of the tool call."""
+
+    @override
+    def merge_in_place(self, other) -> bool:
+        if not isinstance(other, ToolCallPart):
+            return False
+        if self.function.arguments is None:
+            self.function.arguments = other.arguments_part
+        else:
+            self.function.arguments += other.arguments_part or ""
+        return True
+
+
+class ToolCallPart(BaseModel, MergableMixin):
+    """A part of the tool call."""
+
+    arguments_part: str | None = None
+    """A part of the arguments of the tool call."""
+
+    @override
+    def merge_in_place(self, other) -> bool:
+        if not isinstance(other, ToolCallPart):
+            return False
+        if self.arguments_part is None:
+            self.arguments_part = other.arguments_part
+        else:
+            self.arguments_part += other.arguments_part or ""
+        return True
+
+
+class Message(BaseModel):
+    """A message in a conversation."""
+
+    role: Literal[
+        "system",
+        "developer",
+        "user",
+        "assistant",
+        "tool",
+    ]
+    name: str | None = None
+
+    content: str | list[ContentPart]
+    """The content of the message."""
+
+    tool_calls: list[ToolCall] | None = None
+    """In assistant messages, there can be tool calls."""
+
+    tool_call_id: str | None = None
+    """In tool messages, there can be a tool call ID."""
+
+    partial: bool | None = None
+
+    @field_serializer("content")
+    def serialize_content(self, content: str | list[ContentPart]) -> str | list[dict]:
+        if isinstance(content, str):
+            return content
+        return [part.model_dump() for part in content]

kosong-0.13.0/src/kosong/base/tool.py

@@ -0,0 +1,22 @@
+from typing import Any, Self
+
+import jsonschema
+from pydantic import BaseModel, model_validator
+
+type ParametersType = dict[str, Any]
+
+
+class Tool(BaseModel):
+    name: str
+    """The name of the tool."""
+
+    description: str
+    """The description of the tool."""
+
+    parameters: ParametersType
+    """The parameters of the tool, in JSON Schema format."""
+
+    @model_validator(mode="after")
+    def validate_parameters(self) -> Self:
+        jsonschema.validate(self.parameters, jsonschema.Draft202012Validator.META_SCHEMA)
+        return self

kosong-0.13.0/src/kosong/chat_provider/__init__.py

@@ -0,0 +1,53 @@
+from kosong.base.chat_provider import ChatProvider
+
+__all__ = [
+    "OpenAILegacy",
+    "Kimi",
+    # for testing
+    "MockChatProvider",
+    "ChaosChatProvider",
+]
+
+
+def __static_check_types(
+    openai: "OpenAILegacy",
+    kimi: "Kimi",
+    mock: "MockChatProvider",
+    chaos: "ChaosChatProvider",
+):
+    """Use type checking to ensure the types are correct implemented."""
+    _: ChatProvider = openai
+    _: ChatProvider = mock
+    _: ChatProvider = kimi
+    _: ChatProvider = chaos
+
+
+class ChatProviderError(Exception):
+    """The error raised by a chat provider."""
+
+    def __init__(self, message: str):
+        super().__init__(message)
+
+
+class APIConnectionError(ChatProviderError):
+    """The error raised when the API connection fails."""
+
+
+class APITimeoutError(ChatProviderError):
+    """The error raised when the API request times out."""
+
+
+class APIStatusError(ChatProviderError):
+    """The error raised when the API returns a status code of 4xx or 5xx."""
+
+    status_code: int
+
+    def __init__(self, status_code: int, message: str):
+        super().__init__(message)
+        self.status_code = status_code
+
+
+from .chaos import ChaosChatProvider  # noqa: E402
+from .kimi import Kimi  # noqa: E402
+from .mock import MockChatProvider  # noqa: E402
+from .openai_legacy import OpenAILegacy  # noqa: E402

kosong-0.13.0/src/kosong/chat_provider/chaos.py

@@ -0,0 +1,106 @@
+import json
+import os
+import random
+
+import httpx
+from pydantic import BaseModel
+
+from kosong.chat_provider.openai_legacy import OpenAILegacy
+
+
+class ChaosConfig(BaseModel):
+    """Configuration for chaos provider."""
+
+    error_probability: float = 0.3
+    error_types: list[int] = [429, 500, 502, 503]
+    retry_after: int = 2
+    seed: int | None = None
+
+    @classmethod
+    def from_env(cls) -> "ChaosConfig":
+        """Create config from environment variables."""
+        seed_str = os.getenv("CHAOS_SEED")
+        return cls(
+            error_probability=float(os.getenv("CHAOS_ERROR_PROBABILITY", "0.3")),
+            error_types=[
+                int(x.strip()) for x in os.getenv("CHAOS_ERROR_TYPES", "429,500,502,503").split(",")
+            ],
+            retry_after=int(os.getenv("CHAOS_RETRY_AFTER", "2")),
+            seed=int(seed_str) if seed_str else None,
+        )
+
+
+class ChaosTransport(httpx.AsyncBaseTransport):
+    """HTTP transport that randomly injects errors."""
+
+    def __init__(self, wrapped_transport: httpx.AsyncBaseTransport, config: ChaosConfig):
+        self._wrapped = wrapped_transport
+        self._config = config
+        if config.seed is not None:
+            random.seed(config.seed)
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        if self._should_inject_error():
+            error_code = random.choice(self._config.error_types)
+            return self._create_error_response(request, error_code)
+
+        return await self._wrapped.handle_async_request(request)
+
+    def _should_inject_error(self) -> bool:
+        return random.random() < self._config.error_probability
+
+    def _create_error_response(self, request: httpx.Request, status_code: int) -> httpx.Response:
+        error_messages = {
+            429: {"error": {"code": "rate_limit_exceeded", "message": "Rate limit exceeded"}},
+            500: {"error": {"code": "internal_error", "message": "Internal server error"}},
+            502: {"error": {"code": "bad_gateway", "message": "Bad gateway"}},
+            503: {
+                "error": {
+                    "code": "service_unavailable",
+                    "message": "Service temporarily unavailable",
+                }
+            },
+        }
+
+        content = json.dumps(
+            error_messages.get(status_code, {"error": {"message": "Unknown error"}})
+        )
+        headers = {"content-type": "application/json"}
+
+        if status_code == 429:
+            headers["retry-after"] = str(self._config.retry_after)
+
+        return httpx.Response(
+            status_code=status_code,
+            headers=headers,
+            content=content.encode(),
+            request=request,
+        )
+
+
+class ChaosChatProvider(OpenAILegacy):
+    """OpenAI Legacy provider with chaos error injection."""
+
+    def __init__(
+        self,
+        model: str,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        chaos_config: ChaosConfig | None = None,
+        **client_kwargs,
+    ):
+        super().__init__(model=model, api_key=api_key, base_url=base_url, **client_kwargs)
+        self._chaos_config = chaos_config or ChaosConfig.from_env()
+        self._monkey_patch_client()
+
+    def _monkey_patch_client(self):
+        """Inject chaos transport into the client."""
+        original_transport = self._client._client._transport
+        chaos_transport = ChaosTransport(original_transport, self._chaos_config)
+        self._client._client._transport = chaos_transport
+
+    @property
+    def model_name(self) -> str:
+        if self._chaos_config.error_probability > 0:
+            return f"chaos({super().model_name})"
+        return super().model_name