From 4259c5b9b5fba8523e0268a5dfd6f09ae5c5ea79 Mon Sep 17 00:00:00 2001 From: jimmyzhuu Date: Fri, 26 Jun 2026 10:52:00 +0800 Subject: [PATCH] feat: support OpenAI-compatible channel endpoints --- README.md | 13 ++++ channels/example/tools.toml | 5 ++ tagopen/gateway/router.py | 4 +- tagopen/llm.py | 51 +++++++++++--- tagopen/memory/store.py | 7 +- tests/unit/test_llm.py | 136 ++++++++++++++++++++++++++++++++++++ 6 files changed, 201 insertions(+), 15 deletions(-) create mode 100644 tests/unit/test_llm.py diff --git a/README.md b/README.md index 33775e17..f33b5b1d 100644 --- a/README.md +++ b/README.md @@ -356,6 +356,19 @@ Uses [LiteLLM](https://github.com/BerriAI/litellm) — one interface for every p model = "claude-opus-4-8" ``` +Channels can also point at OpenAI-compatible endpoints, including self-hosted +gateways and cloud providers that expose the OpenAI chat completions API: + +```toml +[llm] +model = "openai/ernie-4.5-turbo-32k" +api_base = "https://qianfan.baidubce.com/v2" +api_key_env = "QIANFAN_API_KEY" +``` + +`api_key_env` names the environment variable TagOpen should read at runtime, so +secrets stay out of channel config files. + --- ## Built-in Tools diff --git a/channels/example/tools.toml b/channels/example/tools.toml index bb1a80e8..4e0a0a44 100644 --- a/channels/example/tools.toml +++ b/channels/example/tools.toml @@ -7,6 +7,11 @@ # model = "claude-sonnet-4-6" # use Anthropic # model = "gemini/gemini-2.0-flash" # use Gemini # model = "groq/llama-3.3-70b-versatile" # use Groq +# +# OpenAI-compatible endpoint override for this channel: +# model = "openai/ernie-4.5-turbo-32k" +# api_base = "https://qianfan.baidubce.com/v2" +# api_key_env = "QIANFAN_API_KEY" # read from process env, not this file # Example: GitHub MCP server # [[mcp_server]] diff --git a/tagopen/gateway/router.py b/tagopen/gateway/router.py index a062bb59..2234a357 100644 --- a/tagopen/gateway/router.py +++ b/tagopen/gateway/router.py @@ -1,5 +1,6 @@ """Channel router — maps (workspace_id, channel_id) to an AgentSession and runs the loop.""" +import asyncio import logging from dataclasses import dataclass, field from typing import TYPE_CHECKING @@ -27,9 +28,6 @@ class AgentSession: _lock: asyncio.Lock = field(default_factory=lambda: __import__("asyncio").Lock()) -import asyncio # noqa: E402 — needed for Lock reference above - - def get_or_create_session(workspace_id: str, channel_id: str) -> AgentSession: key = (workspace_id, channel_id) if key not in _sessions: diff --git a/tagopen/llm.py b/tagopen/llm.py index 1717da08..02e4a144 100644 --- a/tagopen/llm.py +++ b/tagopen/llm.py @@ -14,7 +14,7 @@ import logging import os -from pathlib import Path +from dataclasses import dataclass import litellm import toml @@ -27,6 +27,13 @@ litellm.suppress_debug_info = True +@dataclass(frozen=True) +class ChannelLLMConfig: + model: str | None = None + api_base: str | None = None + api_key_env: str | None = None + + def configure() -> None: """Sync API keys from settings → os.environ so LiteLLM can pick them up.""" _set_if_nonempty("ANTHROPIC_API_KEY", settings.anthropic_api_key) @@ -67,24 +74,50 @@ def resolve_model(channel_id: str | None = None) -> str: 2. LLM_MODEL env var / settings.llm_model """ if channel_id: - override = _channel_model_override(channel_id) - if override: - return override + config = _channel_llm_config(channel_id) + if config.model: + return config.model return settings.llm_model -def _channel_model_override(channel_id: str) -> str | None: +def _channel_llm_config(channel_id: str) -> ChannelLLMConfig: tools_toml = settings.channels_dir / channel_id / "tools.toml" if not tools_toml.exists(): - return None + return ChannelLLMConfig() try: config = toml.loads(tools_toml.read_text()) - return config.get("llm", {}).get("model") or None + llm_config = config.get("llm", {}) + return ChannelLLMConfig( + model=llm_config.get("model") or None, + api_base=llm_config.get("api_base") or None, + api_key_env=llm_config.get("api_key_env") or None, + ) except Exception: - return None + return ChannelLLMConfig() + + +def _apply_channel_llm_config(kwargs: dict, channel_id: str | None) -> None: + if not channel_id: + kwargs.setdefault("model", settings.llm_model) + return + + config = _channel_llm_config(channel_id) + kwargs.setdefault("model", config.model or settings.llm_model) + + if config.api_base: + kwargs.setdefault("api_base", config.api_base) + + if config.api_key_env: + api_key = os.environ.get(config.api_key_env) + if not api_key: + raise ValueError( + f"Channel {channel_id} config references missing environment " + f"variable {config.api_key_env!r} for LLM api_key" + ) + kwargs.setdefault("api_key", api_key) async def acompletion(channel_id: str | None = None, **kwargs): """Thin wrapper around litellm.acompletion that injects the resolved model.""" - kwargs.setdefault("model", resolve_model(channel_id)) + _apply_channel_llm_config(kwargs, channel_id) return await litellm.acompletion(**kwargs) diff --git a/tagopen/memory/store.py b/tagopen/memory/store.py index d930bdc5..16bc589c 100644 --- a/tagopen/memory/store.py +++ b/tagopen/memory/store.py @@ -3,7 +3,6 @@ from __future__ import annotations import logging -from functools import lru_cache from pathlib import Path import aiosqlite @@ -76,7 +75,9 @@ async def add_message( ) -> None: assert self._db await self._db.execute( - """INSERT INTO messages (ts, thread_ts, channel_id, role, user_id, display_name, content, tool_calls) + """INSERT INTO messages ( + ts, thread_ts, channel_id, role, user_id, display_name, content, tool_calls + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", (ts, thread_ts, self._channel_id, role, user_id, display_name, content, tool_calls), ) @@ -88,7 +89,7 @@ async def get_recent_messages(self, limit: int = 50) -> list[aiosqlite.Row]: """SELECT ts, role, user_id, display_name, content FROM messages WHERE channel_id = ? - ORDER BY created_at DESC + ORDER BY id DESC LIMIT ?""", (self._channel_id, limit), ) as cursor: diff --git a/tests/unit/test_llm.py b/tests/unit/test_llm.py new file mode 100644 index 00000000..355d43bb --- /dev/null +++ b/tests/unit/test_llm.py @@ -0,0 +1,136 @@ +"""Tests for LiteLLM channel-level configuration.""" + +from __future__ import annotations + +import pytest + +from tagopen import llm +from tagopen.config import settings + + +@pytest.fixture +def capture_litellm(monkeypatch): + captured: dict = {} + + async def fake_acompletion(**kwargs): + captured.update(kwargs) + return {"ok": True} + + monkeypatch.setattr(llm.litellm, "acompletion", fake_acompletion) + return captured + + +@pytest.fixture +def use_tmp_data_dir(tmp_path, monkeypatch): + monkeypatch.setattr(settings, "data_dir", tmp_path) + monkeypatch.setattr(settings, "llm_model", "global-model") + return tmp_path + + +def write_tools_toml(data_dir, channel_id: str, content: str) -> None: + channel_dir = data_dir / "channels" / channel_id + channel_dir.mkdir(parents=True) + (channel_dir / "tools.toml").write_text(content) + + +async def test_acompletion_uses_global_model_without_channel_config( + use_tmp_data_dir, + capture_litellm, +): + result = await llm.acompletion(channel_id="C123", messages=[]) + + assert result == {"ok": True} + assert capture_litellm["model"] == "global-model" + assert "api_base" not in capture_litellm + assert "api_key" not in capture_litellm + + +async def test_acompletion_injects_channel_endpoint_and_api_key( + use_tmp_data_dir, + capture_litellm, + monkeypatch, +): + write_tools_toml( + use_tmp_data_dir, + "C123", + """ +[llm] +model = "openai/ernie-4.5-turbo-32k" +api_base = "https://qianfan.baidubce.com/v2" +api_key_env = "QIANFAN_API_KEY" +""", + ) + monkeypatch.setenv("QIANFAN_API_KEY", "secret-key") + + await llm.acompletion(channel_id="C123", messages=[]) + + assert capture_litellm["model"] == "openai/ernie-4.5-turbo-32k" + assert capture_litellm["api_base"] == "https://qianfan.baidubce.com/v2" + assert capture_litellm["api_key"] == "secret-key" + + +async def test_acompletion_fails_when_api_key_env_is_missing( + use_tmp_data_dir, + capture_litellm, +): + write_tools_toml( + use_tmp_data_dir, + "C123", + """ +[llm] +model = "openai/ernie-4.5-turbo-32k" +api_base = "https://qianfan.baidubce.com/v2" +api_key_env = "MISSING_QIANFAN_API_KEY" +""", + ) + + with pytest.raises(ValueError) as exc_info: + await llm.acompletion(channel_id="C123", messages=[]) + + message = str(exc_info.value) + assert "MISSING_QIANFAN_API_KEY" in message + assert "C123" in message + assert capture_litellm == {} + + +async def test_acompletion_keeps_explicit_kwargs( + use_tmp_data_dir, + capture_litellm, + monkeypatch, +): + write_tools_toml( + use_tmp_data_dir, + "C123", + """ +[llm] +model = "openai/ernie-4.5-turbo-32k" +api_base = "https://qianfan.baidubce.com/v2" +api_key_env = "QIANFAN_API_KEY" +""", + ) + monkeypatch.setenv("QIANFAN_API_KEY", "channel-key") + + await llm.acompletion( + channel_id="C123", + messages=[], + model="explicit-model", + api_base="https://example.invalid/v1", + api_key="explicit-key", + ) + + assert capture_litellm["model"] == "explicit-model" + assert capture_litellm["api_base"] == "https://example.invalid/v1" + assert capture_litellm["api_key"] == "explicit-key" + + +async def test_acompletion_falls_back_when_tools_toml_is_invalid( + use_tmp_data_dir, + capture_litellm, +): + write_tools_toml(use_tmp_data_dir, "C123", "[llm") + + await llm.acompletion(channel_id="C123", messages=[]) + + assert capture_litellm["model"] == "global-model" + assert "api_base" not in capture_litellm + assert "api_key" not in capture_litellm