complete restructure

vibe_bot/llm/__init__.py

@@ -0,0 +1 @@
+"""Submodules of the OpenAI-compatible client layer (chat, images, registry)."""

vibe_bot/llm/chat.py

@@ -0,0 +1,206 @@
+"""Core async chat completion with iterative tool calling, plus adapters."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING, Any, cast
+
+import openai
+
+if TYPE_CHECKING:
+    from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
+
+logger = logging.getLogger(__name__)
+
+ToolExecutor = Callable[[str, dict[str, str]], str]
+ToolCallNotifier = Callable[[str, dict[str, str]], None | Awaitable[None]]
+
+_chat_client: openai.AsyncOpenAI | None = None
+
+
+def get_chat_client() -> openai.AsyncOpenAI:
+    """Return the shared async chat client, building it once."""
+    global _chat_client
+    if _chat_client is None:
+        from vibe_bot.config import CHAT_ENDPOINT, CHAT_ENDPOINT_KEY
+
+        _chat_client = openai.AsyncOpenAI(
+            base_url=CHAT_ENDPOINT, api_key=CHAT_ENDPOINT_KEY
+        )
+    return _chat_client
+
+
+async def chat_complete(
+    messages: list[ChatCompletionMessageParam],
+    *,
+    model: str,
+    max_tokens: int,
+    seed: int | None = None,
+    tools: list[dict[str, object]] | None = None,
+    tool_executor: ToolExecutor | None = None,
+    tool_call_notifier: ToolCallNotifier | None = None,
+    max_tool_rounds: int = 5,
+    timeout: float = 60.0,
+) -> str:
+    """Send a chat completion, optionally with iterative tool calling.
+
+    Args:
+        messages: The conversation messages (system/user/assistant/tool).
+        model: The model to use for completion.
+        max_tokens: Maximum number of tokens to generate.
+        seed: Optional sampling seed.
+        tools: Optional list of tool definitions in OpenAI format.
+        tool_executor: Sync callable (tool_name, tool_args) -> result string.
+        tool_call_notifier: Optional sync-or-async callback invoked before each
+            tool call with (tool_name, tool_args).
+        max_tool_rounds: Maximum tool call rounds before giving up.
+        timeout: Per-request timeout in seconds.
+
+    Returns:
+        The model's final response text, stripped of whitespace ("" on failure).
+
+    """
+    client = get_chat_client()
+    messages = list(messages)
+
+    for _round in range(max_tool_rounds):
+        kwargs: dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "timeout": timeout,
+        }
+        if seed is not None:
+            kwargs["seed"] = seed
+        if tools:
+            kwargs["tools"] = cast("list[Any]", tools)
+
+        response = cast(
+            "ChatCompletion", await client.chat.completions.create(**kwargs)
+        )
+        if not response.choices:
+            return ""
+
+        message = response.choices[0].message
+        tool_calls = message.tool_calls
+        if tool_calls and tool_executor is not None:
+            assistant_msg: dict[str, object] = {
+                "role": "assistant",
+                "content": message.content or "",
+            }
+            tool_call_dicts: list[dict[str, object]] = []
+            for tool_call in tool_calls:
+                if tool_call.type != "function":
+                    continue
+                tool_call_dicts.append(
+                    {
+                        "id": tool_call.id,
+                        "type": "function",
+                        "function": {
+                            "name": tool_call.function.name,
+                            "arguments": tool_call.function.arguments,
+                        },
+                    },
+                )
+            assistant_msg["tool_calls"] = tool_call_dicts
+            messages.append(cast("ChatCompletionMessageParam", assistant_msg))
+
+            for tool_call in tool_calls:
+                if tool_call.type != "function":
+                    continue
+                tool_name = tool_call.function.name
+                tool_args = json.loads(tool_call.function.arguments)
+
+                if tool_call_notifier is not None:
+                    result = tool_call_notifier(tool_name, tool_args)
+                    if result is not None:
+                        await result
+
+                tool_result = await asyncio.to_thread(
+                    tool_executor, tool_name, tool_args
+                )
+                messages.append(
+                    cast(
+                        "ChatCompletionMessageParam",
+                        {
+                            "role": "tool",
+                            "tool_call_id": tool_call.id,
+                            "content": tool_result,
+                        },
+                    ),
+                )
+            continue
+
+        content = message.content
+        if content:
+            return content.strip()
+        return ""
+
+    return ""
+
+
+async def chat_completion_instruct(
+    system_prompt: str,
+    user_prompt: str,
+    *,
+    model: str,
+    max_tokens: int = 1000,
+) -> str:
+    """Instruction-based completion over :func:`chat_complete`."""
+    messages: list[ChatCompletionMessageParam] = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt},
+    ]
+    return await chat_complete(messages, model=model, max_tokens=max_tokens, seed=-1)
+
+
+async def chat_completion_with_history(
+    system_prompt: str,
+    prompts: list[dict[str, str]],
+    *,
+    model: str,
+    max_tokens: int = 1000,
+) -> str:
+    """Completion with conversation history over :func:`chat_complete`."""
+    messages: list[ChatCompletionMessageParam] = [
+        cast(
+            "ChatCompletionMessageParam",
+            {"role": "system", "content": system_prompt},
+        ),
+    ]
+    messages.extend(cast("list[ChatCompletionMessageParam]", prompts))
+    return await chat_complete(messages, model=model, max_tokens=max_tokens, seed=-1)
+
+
+async def chat_completion_with_tools(
+    system_prompt: str,
+    prompts: list[dict[str, str]],
+    tools: list[dict[str, object]],
+    tool_executor: ToolExecutor,
+    *,
+    model: str,
+    max_tokens: int = 1000,
+    max_tool_rounds: int = 5,
+    tool_call_notifier: ToolCallNotifier | None = None,
+) -> str:
+    """Tool-capable completion over :func:`chat_complete`."""
+    messages: list[ChatCompletionMessageParam] = [
+        cast(
+            "ChatCompletionMessageParam",
+            {"role": "system", "content": system_prompt},
+        ),
+    ]
+    messages.extend(cast("list[ChatCompletionMessageParam]", prompts))
+    return await chat_complete(
+        messages,
+        model=model,
+        max_tokens=max_tokens,
+        seed=-1,
+        tools=tools,
+        tool_executor=tool_executor,
+        tool_call_notifier=tool_call_notifier,
+        max_tool_rounds=max_tool_rounds,
+    )

vibe_bot/llm/images.py

@@ -0,0 +1,92 @@
+"""Async image generation and editing clients."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import openai
+
+if TYPE_CHECKING:
+    from io import BufferedReader, BytesIO
+
+logger = logging.getLogger(__name__)
+
+_image_gen_client: openai.AsyncOpenAI | None = None
+_image_edit_client: openai.AsyncOpenAI | None = None
+
+
+def get_image_gen_client() -> openai.AsyncOpenAI:
+    """Return the shared async image-generation client, building it once."""
+    global _image_gen_client
+    if _image_gen_client is None:
+        from vibe_bot.config import IMAGE_GEN_ENDPOINT, IMAGE_GEN_ENDPOINT_KEY
+
+        _image_gen_client = openai.AsyncOpenAI(
+            base_url=IMAGE_GEN_ENDPOINT,
+            api_key=IMAGE_GEN_ENDPOINT_KEY,
+            max_retries=0,
+        )
+    return _image_gen_client
+
+
+def get_image_edit_client() -> openai.AsyncOpenAI:
+    """Return the shared async image-edit client, building it once."""
+    global _image_edit_client
+    if _image_edit_client is None:
+        from vibe_bot.config import IMAGE_EDIT_ENDPOINT, IMAGE_EDIT_ENDPOINT_KEY
+
+        _image_edit_client = openai.AsyncOpenAI(
+            base_url=IMAGE_EDIT_ENDPOINT, api_key=IMAGE_EDIT_ENDPOINT_KEY
+        )
+    return _image_edit_client
+
+
+async def image_generation(
+    prompt: str,
+    *,
+    model: str = "gen",
+    n: int = 1,
+    size: str = "1024x1024",
+) -> str:
+    """Generate an image; return base64 data ("" on failure)."""
+    client = get_image_gen_client()
+    try:
+        response = await client.images.generate(
+            prompt=prompt,
+            n=n,
+            size=size,
+            model=model,
+            timeout=300.0,
+        )
+    except openai.OpenAIError as e:
+        logger.warning("Image generation failed: %s", e)
+        return ""
+    if response.data:
+        return response.data[0].b64_json or ""
+    return ""
+
+
+async def image_edit(
+    image: BufferedReader | BytesIO | list[BufferedReader] | list[BytesIO],
+    prompt: str,
+    *,
+    model: str = "edit",
+    n: int = 1,
+) -> str:
+    """Edit an image; return base64 data ("" on failure)."""
+    client = get_image_edit_client()
+    try:
+        response = await client.images.edit(
+            image=image,
+            prompt=prompt,
+            n=n,
+            size="768x768",
+            model=model,
+        )
+    except openai.OpenAIError as e:
+        logger.warning("Image edit failed: %s", e)
+        return ""
+    if response.data:
+        return response.data[0].b64_json or ""
+    return ""

vibe_bot/llm/registry.py

@@ -0,0 +1,106 @@
+"""Tool registry: OpenAI schemas plus dispatch to sync implementations."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, cast
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+ToolImpl = Callable[..., str]
+
+
+class _RegisteredTool:
+    """A registered tool: its OpenAI schema plus its synchronous impl."""
+
+    __slots__ = ("args_schema", "description", "impl", "name")
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        args_schema: dict[str, object],
+        impl: ToolImpl,
+    ) -> None:
+        self.name = name
+        self.description = description
+        self.args_schema = args_schema
+        self.impl = impl
+
+
+class ToolRegistry:
+    """Holds tool schemas and dispatches tool calls to their implementations."""
+
+    def __init__(self) -> None:
+        self._tools: dict[str, _RegisteredTool] = {}
+
+    def register(
+        self,
+        name: str,
+        description: str,
+        args_schema: dict[str, object],
+        impl: ToolImpl,
+    ) -> None:
+        """Register a tool under ``name`` with its OpenAI args schema."""
+        self._tools[name] = _RegisteredTool(name, description, args_schema, impl)
+
+    def to_openai_tools(self) -> list[dict[str, object]]:
+        """Return the registered tools in OpenAI function-calling format."""
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "parameters": tool.args_schema,
+                },
+            }
+            for tool in self._tools.values()
+        ]
+
+    def execute(self, name: str, args: dict[str, str], **impl_kwargs: Any) -> str:
+        """Dispatch a tool call; unknown tools yield a friendly message.
+
+        ``impl_kwargs`` (e.g. ``channel``) are forwarded to the impl so tools
+        can access per-invocation context.
+        """
+        tool = self._tools.get(name)
+        if tool is None:
+            return f"Unknown tool: {name}"
+        return tool.impl(name, args, **impl_kwargs)
+
+
+_default_registry: ToolRegistry | None = None
+
+
+def get_tool_registry() -> ToolRegistry:
+    """Return the shared tool registry, seeded with the channel-members tool."""
+    global _default_registry
+    if _default_registry is None:
+        from vibe_bot.tools import get_channel_members
+
+        raw_schema = get_channel_members.args_schema
+        if isinstance(raw_schema, dict):
+            args_schema: dict[str, object] = raw_schema
+        else:
+            # A LangChain @tool exposes args_schema as a pydantic model class.
+            args_schema = cast("type[BaseModel]", raw_schema).model_json_schema()
+
+        registry = ToolRegistry()
+        registry.register(
+            get_channel_members.name,
+            get_channel_members.description or "",
+            args_schema,
+            _channel_members_tool,
+        )
+        _default_registry = registry
+    return _default_registry
+
+
+def _channel_members_tool(name: str, args: dict[str, str], **kwargs: Any) -> str:
+    """Adapt the registry dispatch to ``get_channel_members_impl(channel)``."""
+    from vibe_bot.tools import get_channel_members_impl
+
+    channel = kwargs.get("channel")
+    return get_channel_members_impl(channel)