diff --git a/README.md b/README.md index c55a8e32..fef87a0c 100644 --- a/README.md +++ b/README.md @@ -111,6 +111,7 @@ turnstone/ │ ├── tools.py # Tool definitions (auto-loaded from JSON) │ ├── workstream.py # WorkstreamManager — parallel independent sessions │ ├── mcp_client.py # MCP client manager (external tool servers) +│ ├── model_registry.py # ModelRegistry — named models, fallback routing, per-workstream selection │ ├── config.py # Unified TOML config (~/.config/turnstone/config.toml) │ ├── memory.py # SQLite persistence (memories, conversations, FTS5) │ ├── metrics.py # Prometheus-compatible metrics collector @@ -238,6 +239,29 @@ turnstone-server --mcp-config ~/.config/turnstone/mcp.json Use `/mcp` in the REPL to list connected tools. MCP tools require user approval by default (overridden by `--skip-permissions` or UI auto-approve). +### Multi-Model Support + +Turnstone supports multiple model backends per server instance. Define named models in `config.toml` and select per-workstream or switch mid-session with `/model `. + +```toml +[models.local] +base_url = "http://localhost:8000/v1" +model = "qwen3-32b" + +[models.openai] +base_url = "https://api.openai.com/v1" +api_key = "sk-..." +model = "gpt-4o" +context_window = 128000 + +[model] +default = "local" # which model to use by default +fallback = ["openai"] # try these if the primary is unreachable +agent_model = "local" # optional: cheaper model for plan/task sub-agents +``` + +Use `/model` to show available models, `/model openai` to switch. Workstreams created via the API accept an optional `model` parameter. + ## Configuration All entry points read `~/.config/turnstone/config.toml`. CLI flags override config values. @@ -252,6 +276,9 @@ tavily_key = "" name = "" # empty = auto-detect temperature = 0.5 reasoning_effort = "medium" +default = "default" # model alias for new workstreams +fallback = [] # ordered list of fallback model aliases +agent_model = "" # model alias for plan/task sub-agents [tools] timeout = 30 diff --git a/docs/api-reference.md b/docs/api-reference.md index bb8e25b9..1d24a83a 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -51,6 +51,7 @@ not recognized. { "type": "connected", "model": "kappa_20b_131k", + "model_alias": "default", "skip_permissions": false } ``` @@ -505,10 +506,16 @@ Creates a new workstream. The server supports up to 10 concurrent workstreams. **Request body:** ```json -{} +{"name": "my-ws", "model": "openai"} ``` -No fields are required. The body can be empty or an empty JSON object. +All fields are optional. The body can be empty or an empty JSON object. + +| Field | Type | Default | Description | +|----------------|--------|---------|------------------------------------------------| +| `name` | string | auto | Workstream display name | +| `model` | string | default | Model alias from the registry (`[models.*]`) | +| `auto_approve` | bool | false | Auto-approve all tool calls for this workstream | **Response (success):** diff --git a/docs/architecture.md b/docs/architecture.md index 92979a56..ce0dd416 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -35,6 +35,7 @@ turnstone/ workstream.py Parallel workstream manager (WorkstreamState, Workstream, WorkstreamManager) tools.py Tool schema loader (JSON -> OpenAI function-calling format) mcp_client.py MCPClientManager — MCP server connections, tool discovery, async-sync bridge + model_registry.py ModelRegistry — named model configs, lazy client creation, fallback routing memory.py SQLite persistence (conversations, memories, FTS5 search) metrics.py Prometheus-compatible metrics collector (MetricsCollector) edit.py File edit utilities (find_occurrences, pick_nearest) @@ -467,6 +468,47 @@ at connection time (server names with `__` are rejected). servers still connect. Tool execution errors return error strings to the LLM rather than crashing the session. +### Multi-Model Registry + +`ModelRegistry` (`turnstone/core/model_registry.py`) manages named model +configurations so workstreams can use different LLM backends. + +**Config format:** +```toml +[models.local] +base_url = "http://localhost:8000/v1" +model = "qwen3-32b" + +[models.openai] +base_url = "https://api.openai.com/v1" +api_key = "sk-..." +model = "gpt-4o" +context_window = 128000 + +[model] +default = "local" +fallback = ["openai"] +agent_model = "local" +``` + +**Lifecycle:** +1. `load_model_registry()` reads `[models.*]` sections from config.toml and + builds a `"default"` entry from CLI `--base-url`/`--model`/`--api-key` args +2. The registry is passed to the session factory closure in both `cli.py` and + `server.py`; each workstream resolves its model on creation +3. `ModelRegistry.get_client()` lazily creates `OpenAI` client instances + (thread-safe via `_client_lock`) +4. `/model` command shows available models; `/model ` switches the + active workstream's client, model, and context window +5. `_create_stream_with_retry()` tries the primary model, then each fallback + alias in order if the primary is unreachable +6. `_run_agent()` resolves `registry.agent_model` (if set) for plan/task + sub-agents, allowing a cheaper model for autonomous loops + +**Per-workstream selection:** `POST /api/workstreams/new` accepts an optional +`"model"` field. The bridge `CreateWorkstreamMessage` carries the same field +through the MQ protocol. + ### Tool Output Truncation Tool execution results (bash, read_file, search, math, man) are truncated by diff --git a/docs/diagrams/02-package-structure.puml b/docs/diagrams/02-package-structure.puml index 6d4fed2f..e24cd6e0 100644 --- a/docs/diagrams/02-package-structure.puml +++ b/docs/diagrams/02-package-structure.puml @@ -35,6 +35,7 @@ package "turnstone/core/" <> { component [web.py\nWeb helpers] as web <> component [auth.py\nAuthentication] as auth <> component [mcp_client.py\nMCPClientManager] as mcp <> + component [model_registry.py\nModelRegistry] as registry <> } ' MQ subsystem @@ -107,7 +108,9 @@ session --> sandbox session --> edit session --> web session --> mcp : optional +session --> registry : optional mcp --> config +registry --> config tools --> schemas ' MQ dependencies diff --git a/docs/diagrams/03-core-engine-classes.puml b/docs/diagrams/03-core-engine-classes.puml index a309b503..d8096c52 100644 --- a/docs/diagrams/03-core-engine-classes.puml +++ b/docs/diagrams/03-core-engine-classes.puml @@ -72,6 +72,8 @@ class "ChatSession" as ChatSession { - _msg_tokens: list[int] - _session_id: str - _mcp_client: MCPClientManager | None + - _registry: ModelRegistry | None + + model_alias: str | None {property} - _tools: list[dict] - _task_tools: list[dict] - _agent_tools: list[dict] @@ -83,7 +85,8 @@ class "ChatSession" as ChatSession { + resume_session(session_id: str) - _save_config() - _stream_response(stream) → dict - - _create_stream_with_retry(msgs) → Stream + - _create_stream_with_retry(msgs) → Stream (+ fallback) + - _try_stream(client, model, msgs) → Stream - _execute_tools(tool_calls) → (results, feedback) - _prepare_tool(tc) → item dict - _prepare_mcp_tool(call_id, name, args) → item dict @@ -167,6 +170,35 @@ class "MCPClientManager" as MCPMgr { core/mcp_client.py } +' ModelRegistry +class "ModelRegistry" as ModelReg { + - _models: dict[str, ModelConfig] + - _clients: dict[str, OpenAI] + - _client_lock: Lock + + default: str + + fallback: list[str] + + agent_model: str | None + -- + + resolve(alias) → (client, model, config) + + get_client(alias) → OpenAI + + has_alias(alias) → bool + + list_aliases() → list[str] + + shutdown() + -- + Thread-safe lazy client creation. + Loaded by load_model_registry() + from CLI args + [models.*] config. + -- + core/model_registry.py +} + +class "ModelConfig" as ModelCfg <> { + + alias: str + + base_url: str + + model: str + + context_window: int +} + ' Relationships SessionUI <|.. TerminalUI TerminalUI <|-- WsTermUI @@ -175,6 +207,7 @@ SessionUI <|.. NullUI ChatSession --> SessionUI : uses ChatSession --> MCPMgr : optional +ChatSession --> ModelReg : optional ChatSession <|-- HeadlessSession WsMgr --> "*" Ws : manages @@ -182,7 +215,9 @@ Ws --> "1" ChatSession : wraps Ws --> "1" SessionUI : wraps Ws --> "1" WsState : has -WsMgr ..> ChatSession : creates via\nsession_factory(ui) +WsMgr ..> ChatSession : creates via\nsession_factory(ui, model_alias) + +ModelReg --> "*" ModelCfg : holds note bottom of ChatSession Central engine: multi-turn LLM loop diff --git a/docs/diagrams/png/02-package-structure.png b/docs/diagrams/png/02-package-structure.png index d6c31dad..9d2dc678 100644 --- a/docs/diagrams/png/02-package-structure.png +++ b/docs/diagrams/png/02-package-structure.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:e1aebded0994beccbf1af25c5c1c216775a59a4e343bcbafe69bcc1172907c71 -size 291899 +oid sha256:d66025ff18c6e28ef632e2ebf8bfc71554c64bd457741fda924e90a579580c34 +size 309783 diff --git a/docs/diagrams/png/03-core-engine-classes.png b/docs/diagrams/png/03-core-engine-classes.png index 62d055f3..37ec2bf8 100644 --- a/docs/diagrams/png/03-core-engine-classes.png +++ b/docs/diagrams/png/03-core-engine-classes.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:7456dcc6a68b22d841524062028298814293b0e4dfab4097c536847d1528bb8a -size 275572 +oid sha256:168ed3f0a91729673526c98b0358c5ec59b511206dcfc8c227d6a2b365e4e5c5 +size 322776 diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py new file mode 100644 index 00000000..50ee168e --- /dev/null +++ b/tests/test_model_registry.py @@ -0,0 +1,587 @@ +"""Tests for turnstone.core.model_registry — model registry, loading, session integration.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from turnstone.core.model_registry import ( + ModelConfig, + ModelRegistry, + load_model_registry, +) + +# --------------------------------------------------------------------------- +# ModelConfig +# --------------------------------------------------------------------------- + + +class TestModelConfig: + def test_construction(self) -> None: + cfg = ModelConfig( + alias="local", + base_url="http://localhost:8000/v1", + api_key="dummy", + model="qwen3-32b", + ) + assert cfg.alias == "local" + assert cfg.model == "qwen3-32b" + assert cfg.context_window == 131072 # default + + def test_custom_context_window(self) -> None: + cfg = ModelConfig( + alias="oai", + base_url="https://api.openai.com/v1", + api_key="sk-test", + model="gpt-4o", + context_window=128000, + ) + assert cfg.context_window == 128000 + + def test_frozen(self) -> None: + cfg = ModelConfig(alias="x", base_url="x", api_key="x", model="x") + with pytest.raises(AttributeError): + cfg.alias = "y" # type: ignore[misc] + + def test_api_key_not_in_repr(self) -> None: + cfg = ModelConfig(alias="test", base_url="http://x", api_key="sk-secret-key", model="m") + assert "sk-secret-key" not in repr(cfg) + + +# --------------------------------------------------------------------------- +# ModelRegistry +# --------------------------------------------------------------------------- + + +class TestModelRegistry: + def _make_registry( + self, + fallback: list[str] | None = None, + agent_model: str | None = None, + ) -> ModelRegistry: + models = { + "default": ModelConfig("default", "http://localhost:8000/v1", "dummy", "qwen3-32b"), + "openai": ModelConfig( + "openai", "https://api.openai.com/v1", "sk-test", "gpt-4o", 128000 + ), + "cheap": ModelConfig( + "cheap", "https://api.openai.com/v1", "sk-test", "gpt-4o-mini", 128000 + ), + } + return ModelRegistry( + models=models, + default="default", + fallback=fallback, + agent_model=agent_model, + ) + + def test_resolve_default(self) -> None: + reg = self._make_registry() + client, model, cfg = reg.resolve() + assert model == "qwen3-32b" + assert cfg.alias == "default" + + def test_resolve_alias(self) -> None: + reg = self._make_registry() + client, model, cfg = reg.resolve("openai") + assert model == "gpt-4o" + assert cfg.context_window == 128000 + + def test_resolve_none_uses_default(self) -> None: + reg = self._make_registry() + _, model1, _ = reg.resolve(None) + _, model2, _ = reg.resolve() + assert model1 == model2 + + def test_lazy_client_creation(self) -> None: + reg = self._make_registry() + assert len(reg._clients) == 0 + reg.get_client("default") + assert len(reg._clients) == 1 + # Second call reuses + c1 = reg.get_client("default") + c2 = reg.get_client("default") + assert c1 is c2 + + def test_list_aliases(self) -> None: + reg = self._make_registry() + aliases = reg.list_aliases() + assert set(aliases) == {"default", "openai", "cheap"} + + def test_count(self) -> None: + reg = self._make_registry() + assert reg.count == 3 + + def test_unknown_alias_error(self) -> None: + reg = self._make_registry() + with pytest.raises(ValueError, match="Unknown model alias"): + reg.get_config("nonexistent") + with pytest.raises(ValueError, match="Unknown model alias"): + reg.get_client("nonexistent") + + def test_shutdown(self) -> None: + reg = self._make_registry() + reg.get_client("default") + reg.get_client("openai") + assert len(reg._clients) == 2 + reg.shutdown() + assert len(reg._clients) == 0 + + def test_has_alias(self) -> None: + reg = self._make_registry() + assert reg.has_alias("default") + assert reg.has_alias("openai") + assert not reg.has_alias("nonexistent") + + def test_concurrent_get_client(self) -> None: + """Thread-safe lazy client creation under concurrency.""" + import concurrent.futures + + reg = self._make_registry() + clients: list[Any] = [] + + def get_it() -> Any: + return reg.get_client("default") + + with concurrent.futures.ThreadPoolExecutor(max_workers=10) as pool: + futs = [pool.submit(get_it) for _ in range(20)] + clients = [f.result() for f in futs] + + # All threads should get the same client instance + assert all(c is clients[0] for c in clients) + assert len(reg._clients) == 1 + + def test_fallback_stored(self) -> None: + reg = self._make_registry(fallback=["openai", "cheap"]) + assert reg.fallback == ["openai", "cheap"] + + def test_agent_model_stored(self) -> None: + reg = self._make_registry(agent_model="cheap") + assert reg.agent_model == "cheap" + + +class TestModelRegistryValidation: + def test_empty_models_raises(self) -> None: + with pytest.raises(ValueError, match="at least one"): + ModelRegistry(models={}, default="x") + + def test_invalid_default_raises(self) -> None: + models = {"a": ModelConfig("a", "x", "x", "x")} + with pytest.raises(ValueError, match="Default model 'bad'"): + ModelRegistry(models=models, default="bad") + + def test_invalid_fallback_raises(self) -> None: + models = {"a": ModelConfig("a", "x", "x", "x")} + with pytest.raises(ValueError, match="Fallback model 'bad'"): + ModelRegistry(models=models, default="a", fallback=["bad"]) + + def test_invalid_agent_model_raises(self) -> None: + models = {"a": ModelConfig("a", "x", "x", "x")} + with pytest.raises(ValueError, match="Agent model 'bad'"): + ModelRegistry(models=models, default="a", agent_model="bad") + + +# --------------------------------------------------------------------------- +# load_model_registry +# --------------------------------------------------------------------------- + + +class TestLoadModelRegistry: + def test_single_entry_from_args(self) -> None: + """No [models] config → single-entry registry from CLI args.""" + with patch("turnstone.core.model_registry.load_config", return_value={}): + reg = load_model_registry( + base_url="http://localhost:8000/v1", + api_key="dummy", + model="qwen3-32b", + ) + assert reg.count == 1 + assert reg.default == "default" + _, model, cfg = reg.resolve() + assert model == "qwen3-32b" + assert cfg.base_url == "http://localhost:8000/v1" + + def test_models_from_config(self) -> None: + """[models.*] sections create additional entries.""" + fake_cfg: dict[str, Any] = { + "models": { + "openai": { + "base_url": "https://api.openai.com/v1", + "api_key": "sk-test", + "model": "gpt-4o", + "context_window": 128000, + }, + }, + "model": { + "default": "openai", + }, + } + with patch("turnstone.core.model_registry.load_config", return_value=fake_cfg): + reg = load_model_registry( + base_url="http://localhost:8000/v1", + api_key="dummy", + model="local-model", + ) + assert reg.count == 2 # "openai" + "default" + assert reg.default == "openai" + _, model, _ = reg.resolve() + assert model == "gpt-4o" + + def test_fallback_from_config(self) -> None: + fake_cfg: dict[str, Any] = { + "models": { + "fallback1": { + "base_url": "http://fb1/v1", + "model": "fb-model", + }, + }, + "model": { + "fallback": ["fallback1", "nonexistent"], + }, + } + with patch("turnstone.core.model_registry.load_config", return_value=fake_cfg): + reg = load_model_registry("http://x/v1", "x", "x") + # "nonexistent" is silently dropped + assert reg.fallback == ["fallback1"] + + def test_agent_model_from_config(self) -> None: + fake_cfg: dict[str, Any] = { + "models": { + "cheap": { + "base_url": "http://cheap/v1", + "model": "cheap-model", + }, + }, + "model": { + "agent_model": "cheap", + }, + } + with patch("turnstone.core.model_registry.load_config", return_value=fake_cfg): + reg = load_model_registry("http://x/v1", "x", "x") + assert reg.agent_model == "cheap" + + def test_invalid_agent_model_ignored(self) -> None: + fake_cfg: dict[str, Any] = { + "model": {"agent_model": "nonexistent"}, + } + with patch("turnstone.core.model_registry.load_config", return_value=fake_cfg): + reg = load_model_registry("http://x/v1", "x", "x") + assert reg.agent_model is None + + def test_invalid_default_falls_back(self) -> None: + fake_cfg: dict[str, Any] = { + "model": {"default": "nonexistent"}, + } + with patch("turnstone.core.model_registry.load_config", return_value=fake_cfg): + reg = load_model_registry("http://x/v1", "x", "x") + assert reg.default == "default" + + def test_empty_model_name_skipped(self) -> None: + """Config entries without a model name are skipped.""" + fake_cfg: dict[str, Any] = { + "models": { + "bad": {"base_url": "http://bad/v1"}, # no model key + "good": {"base_url": "http://good/v1", "model": "good-model"}, + }, + } + with patch("turnstone.core.model_registry.load_config", return_value=fake_cfg): + reg = load_model_registry("http://x/v1", "x", "x") + assert not reg.has_alias("bad") + assert reg.has_alias("good") + + def test_unknown_fallback_logged_and_dropped(self) -> None: + fake_cfg: dict[str, Any] = { + "model": {"fallback": ["good", "bad"]}, + "models": { + "good": {"base_url": "http://g/v1", "model": "g-model"}, + }, + } + with patch("turnstone.core.model_registry.load_config", return_value=fake_cfg): + reg = load_model_registry("http://x/v1", "x", "x") + assert reg.fallback == ["good"] + + def test_models_inherit_cli_args(self) -> None: + """Model entries without base_url/api_key inherit from CLI args.""" + fake_cfg: dict[str, Any] = { + "models": { + "alt": { + "model": "alt-model", + }, + }, + } + with patch("turnstone.core.model_registry.load_config", return_value=fake_cfg): + reg = load_model_registry("http://base/v1", "my-key", "default-model") + alt_cfg = reg.get_config("alt") + assert alt_cfg.base_url == "http://base/v1" + assert alt_cfg.api_key == "my-key" + + +# --------------------------------------------------------------------------- +# Session integration +# --------------------------------------------------------------------------- + + +class _FakeUI: + """Minimal SessionUI stub for testing.""" + + def __init__(self) -> None: + self.infos: list[str] = [] + self.errors: list[str] = [] + + def on_thinking_start(self) -> None: ... + def on_thinking_stop(self) -> None: ... + def on_reasoning_token(self, text: str) -> None: ... + def on_content_token(self, text: str) -> None: ... + def on_stream_end(self) -> None: ... + def approve_tools(self, items: list[dict[str, Any]]) -> tuple[bool, str | None]: + return True, None + + def on_tool_result(self, call_id: str, name: str, output: str) -> None: ... + def on_tool_output_chunk(self, call_id: str, chunk: str) -> None: ... + def on_status(self, usage: dict[str, Any], context_window: int, effort: str) -> None: ... + def on_plan_review(self, content: str) -> str: + return "approve" + + def on_info(self, message: str) -> None: + self.infos.append(message) + + def on_error(self, message: str) -> None: + self.errors.append(message) + + def on_state_change(self, state: str) -> None: ... + def on_rename(self, name: str) -> None: ... + + +def _make_session( + registry: ModelRegistry | None = None, + model_alias: str | None = None, +) -> Any: + """Create a ChatSession with a mock client and optional registry.""" + from turnstone.core.session import ChatSession + + client = MagicMock() + return ChatSession( + client=client, + model="test-model", + ui=_FakeUI(), + instructions=None, + temperature=0.5, + max_tokens=4096, + tool_timeout=30, + registry=registry, + model_alias=model_alias, + ) + + +class TestSessionModelCommand: + def test_model_show_without_registry(self) -> None: + session = _make_session() + session.handle_command("/model") + assert "test-model" in session.ui.infos[-1] + + def test_model_show_with_registry(self) -> None: + reg = ModelRegistry( + models={ + "default": ModelConfig("default", "x", "x", "test-model"), + "alt": ModelConfig("alt", "y", "y", "alt-model"), + }, + default="default", + ) + session = _make_session(registry=reg, model_alias="default") + session.handle_command("/model") + info = session.ui.infos[-1] + assert "test-model" in info + assert "default" in info + assert "alt" in info + + def test_model_switch(self) -> None: + reg = ModelRegistry( + models={ + "default": ModelConfig("default", "x", "x", "default-model"), + "alt": ModelConfig("alt", "y", "y", "alt-model", context_window=64000), + }, + default="default", + ) + session = _make_session(registry=reg, model_alias="default") + session.handle_command("/model alt") + assert session.model == "alt-model" + assert session.model_alias == "alt" + assert session.context_window == 64000 + assert "Switched to" in session.ui.infos[-1] + + def test_model_switch_unknown_alias(self) -> None: + reg = ModelRegistry( + models={"default": ModelConfig("default", "x", "x", "test-model")}, + default="default", + ) + session = _make_session(registry=reg, model_alias="default") + session.handle_command("/model nonexistent") + assert "Unknown model alias" in session.ui.infos[-1] + + def test_model_switch_without_registry(self) -> None: + session = _make_session() + session.handle_command("/model something") + assert "Unknown model alias" in session.ui.infos[-1] + + def test_model_show_fallback_info(self) -> None: + reg = ModelRegistry( + models={ + "a": ModelConfig("a", "x", "x", "m-a"), + "b": ModelConfig("b", "y", "y", "m-b"), + }, + default="a", + fallback=["b"], + agent_model="b", + ) + session = _make_session(registry=reg, model_alias="a") + session.handle_command("/model") + info = session.ui.infos[-1] + assert "Fallback: b" in info + assert "Agent model: b" in info + + +class TestSessionFallback: + def test_fallback_on_primary_failure(self) -> None: + reg = ModelRegistry( + models={ + "primary": ModelConfig("primary", "http://p/v1", "k", "p-model"), + "fallback": ModelConfig("fallback", "http://f/v1", "k", "f-model"), + }, + default="primary", + fallback=["fallback"], + ) + session = _make_session(registry=reg, model_alias="primary") + + # _try_stream: first call (primary) raises, second call (fallback) succeeds + call_count = 0 + + def fake_try_stream(client: Any, model: str, msgs: Any) -> str: + nonlocal call_count + call_count += 1 + if call_count == 1: + raise ConnectionError("Primary down") + return "fallback_response" + + session._try_stream = fake_try_stream # type: ignore[assignment] + result = session._create_stream_with_retry([{"role": "user", "content": "hi"}]) + assert result == "fallback_response" + assert call_count == 2 + assert any("falling back" in i for i in session.ui.infos) + + def test_no_fallback_without_registry(self) -> None: + session = _make_session() + + def fake_try_stream(client: Any, model: str, msgs: Any) -> str: + raise ConnectionError("Down") + + session._try_stream = fake_try_stream # type: ignore[assignment] + with pytest.raises(ConnectionError): + session._create_stream_with_retry([{"role": "user", "content": "hi"}]) + + +class TestSessionAgentModel: + def test_agent_model_resolved(self) -> None: + reg = ModelRegistry( + models={ + "main": ModelConfig("main", "http://m/v1", "k", "main-model"), + "agent": ModelConfig("agent", "http://a/v1", "k", "agent-model"), + }, + default="main", + agent_model="agent", + ) + session = _make_session(registry=reg, model_alias="main") + + # Mock the API to capture what model was used + captured_model = None + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message.content = "done" + mock_response.choices[0].message.tool_calls = None + mock_response.choices[0].finish_reason = "stop" + + def fake_create(**kwargs: Any) -> Any: + nonlocal captured_model + captured_model = kwargs.get("model") + return mock_response + + # Get the agent client from the registry and patch it + agent_client = reg.get_client("agent") + agent_client.chat.completions.create = fake_create + + agent_msgs = [ + {"role": "developer", "content": "You are an agent."}, + {"role": "user", "content": "Do something."}, + ] + session._run_agent(agent_msgs) + assert captured_model == "agent-model" + + +# --------------------------------------------------------------------------- +# Workstream integration +# --------------------------------------------------------------------------- + + +class TestWorkstreamModelParam: + def test_create_with_model(self) -> None: + """WorkstreamManager.create passes model_alias to session_factory.""" + from turnstone.core.workstream import WorkstreamManager + + captured_alias = None + + def factory(ui: Any, model_alias: str | None = None) -> Any: + nonlocal captured_alias + captured_alias = model_alias + mock_session = MagicMock() + mock_session.session_id = "test123" + return mock_session + + mgr = WorkstreamManager(factory) + mgr.create(name="test", model="openai") + assert captured_alias == "openai" + + def test_create_without_model(self) -> None: + captured_alias = None + + def factory(ui: Any, model_alias: str | None = None) -> Any: + nonlocal captured_alias + captured_alias = model_alias + mock_session = MagicMock() + mock_session.session_id = "test123" + return mock_session + + from turnstone.core.workstream import WorkstreamManager + + mgr = WorkstreamManager(factory) + mgr.create(name="test") + assert captured_alias is None + + +# --------------------------------------------------------------------------- +# Protocol +# --------------------------------------------------------------------------- + + +class TestProtocolModel: + def test_create_workstream_message_has_model(self) -> None: + from turnstone.mq.protocol import CreateWorkstreamMessage + + msg = CreateWorkstreamMessage(name="test", model="openai") + assert msg.model == "openai" + + def test_create_workstream_message_default(self) -> None: + from turnstone.mq.protocol import CreateWorkstreamMessage + + msg = CreateWorkstreamMessage(name="test") + assert msg.model == "" + + def test_round_trip(self) -> None: + from turnstone.mq.protocol import CreateWorkstreamMessage, InboundMessage + + msg = CreateWorkstreamMessage(name="ws1", model="local") + raw = msg.to_json() + restored = InboundMessage.from_json(raw) + assert isinstance(restored, CreateWorkstreamMessage) + assert restored.model == "local" + assert restored.name == "ws1" diff --git a/tests/test_workstream.py b/tests/test_workstream.py index 87a31000..c59feaf5 100644 --- a/tests/test_workstream.py +++ b/tests/test_workstream.py @@ -20,7 +20,7 @@ class FakeSession: self.messages = [] -def _fake_factory(ui): +def _fake_factory(ui, model_alias=None): return FakeSession() diff --git a/turnstone/cli.py b/turnstone/cli.py index ff337475..a951f616 100644 --- a/turnstone/cli.py +++ b/turnstone/cli.py @@ -845,39 +845,50 @@ def main() -> None: # Set up readline setup_readline() - # Create client + # Create client and detect model api_key = args.api_key or os.environ.get("OPENAI_API_KEY") or "dummy" client = OpenAI( base_url=args.base_url, api_key=api_key, ) - - # Detect or use provided model model = args.model or detect_model(client) + # Build model registry (reads [models.*] sections from config.toml) + from turnstone.core.model_registry import load_model_registry + + registry = load_model_registry( + base_url=args.base_url, + api_key=api_key, + model=model, + context_window=args.context_window, + ) + # Initialize MCP client (connects to configured MCP servers, if any) from turnstone.core.mcp_client import create_mcp_client mcp_client = create_mcp_client(getattr(args, "mcp_config", None)) # Session factory — captures shared config for creating workstream sessions - def session_factory(ui: SessionUI | None) -> ChatSession: + def session_factory(ui: SessionUI | None, model_alias: str | None = None) -> ChatSession: assert ui is not None, "session_factory requires a non-None UI" + r_client, r_model, r_cfg = registry.resolve(model_alias) return ChatSession( - client=client, - model=model, + client=r_client, + model=r_model, ui=ui, instructions=args.instructions, temperature=args.temperature, max_tokens=args.max_tokens, tool_timeout=args.tool_timeout, reasoning_effort=args.reasoning_effort, - context_window=args.context_window, + context_window=r_cfg.context_window, compact_max_tokens=args.compact_max_tokens, auto_compact_pct=args.auto_compact_pct, agent_max_turns=args.agent_max_turns, tool_truncation=args.tool_truncation, mcp_client=mcp_client, + registry=registry, + model_alias=model_alias or registry.default, ) # Create workstream manager and initial workstream @@ -921,6 +932,9 @@ def main() -> None: # Print banner print(f"\n{bold('Chat')} with {cyan(model)}") + if registry.count > 1: + others = [a for a in registry.list_aliases() if a != registry.default] + print(f"Models: {registry.default} (default), {', '.join(others)}") if mcp_client: mcp_tools = mcp_client.get_tools() if mcp_tools: @@ -980,6 +994,7 @@ def main() -> None: if mcp_client: mcp_client.shutdown() + registry.shutdown() print("Goodbye.") diff --git a/turnstone/console/static/app.js b/turnstone/console/static/app.js index 46510ea5..bd05c301 100644 --- a/turnstone/console/static/app.js +++ b/turnstone/console/static/app.js @@ -753,6 +753,8 @@ function renderWsTable(container, wsList) { row.setAttribute("tabindex", "0"); row.setAttribute("role", "button"); var ariaLabel = sd.label + ": " + (ws.name || ws.id || "unnamed"); + if (ws.model_alias || ws.model) + ariaLabel += ", model: " + (ws.model_alias || ws.model); if (ws.node) ariaLabel += " on " + ws.node; if (ws.title) ariaLabel += ", task: " + ws.title; if (ws.tokens) ariaLabel += ", " + formatTokens(ws.tokens) + " tokens"; @@ -785,6 +787,13 @@ function renderWsTable(container, wsList) { nameCell.textContent = ws.name || ws.id || ""; main.appendChild(nameCell); + // MODEL + var modelCell = document.createElement("span"); + modelCell.className = "dash-cell-model"; + modelCell.textContent = ws.model_alias || ws.model || ""; + if (ws.model) modelCell.title = ws.model; + main.appendChild(modelCell); + // NODE (clickable) var nodeCell = document.createElement("span"); nodeCell.className = "dash-cell-node"; diff --git a/turnstone/console/static/index.html b/turnstone/console/static/index.html index b0ae7be1..030795d0 100644 --- a/turnstone/console/static/index.html +++ b/turnstone/console/static/index.html @@ -42,6 +42,7 @@