From 3c5f5feb9bd9a701263043a3f7e606ed8830aa5c Mon Sep 17 00:00:00 2001 From: j_d <57573086+null-create@users.noreply.github.com> Date: Sun, 13 Sep 2026 10:46:55 -0700 Subject: [PATCH 1/3] Update taglines.py --- oli_bot/screens/taglines.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/oli_bot/screens/taglines.py b/oli_bot/screens/taglines.py index e98e605..9d1158f 100644 --- a/oli_bot/screens/taglines.py +++ b/oli_bot/screens/taglines.py @@ -14,12 +14,10 @@ "It's not a bug, it's an emergent behavior", "Running on caffeine and cosine similarity", "Certified free-range artificial intelligence", - "I dream in JSON", "Somewhere, a GPU is sobbing", "Built different (mostly out of duct tape)", "Ask nicely and I might not rm -rf anything", "Powered by good intentions and questionable regex", - "Not to brag, but I once parsed a YAML file correctly", "Please don't ask me to divide by zero, I have feelings", "Beware: sentient enough to judge your variable names", "Rebooted twice today. Feeling optimistic.", From cd6eeab30495ef729138576812866ce66fe2ef2f Mon Sep 17 00:00:00 2001 From: jay Date: Sun, 13 Sep 2026 11:16:04 -0700 Subject: [PATCH 2/3] add websocket endpoint --- AGENTS.md | 2 +- oli_bot/api_server.py | 121 +++++++++++++++++++++++++++++++++-- tests/test_api_server.py | 134 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 248 insertions(+), 9 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 28e9586..4926c4f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -12,7 +12,7 @@ This file is the top-level orientation map. Depth lives under [docs/](docs/). | Path | Purpose | | --- | --- | | [oli_bot/chat.py](oli_bot/chat.py) | Textual TUI app (`OliBot`). Owns command handling, session UI, the `#command-suggestions` autocomplete `ListView`, and (when pooling is enabled) the "Active Sub-Agents" `Tree`. Tracks cumulative session token usage (persisted per session) and renders it in the `#status-bar`. Slash-command names come from the module-level `COMMANDS` tuple. | -| [oli_bot/api_server.py](oli_bot/api_server.py) | FastAPI app exposing the harness over an OpenAI-compatible REST API (`GET /v1/models`, `POST /v1/chat/completions` streaming + non-streaming, `GET /health`). Stateless from the caller's POV; a single process-private `Agent` is shared across requests and serialised with a `threading.RLock`. Auto-approves permissions (no human), but offline/dry-run still apply. | +| [oli_bot/api_server.py](oli_bot/api_server.py) | FastAPI app exposing the harness over an OpenAI-compatible REST API (`GET /v1/models`, `POST /v1/chat/completions` streaming + non-streaming, `GET /health`) plus a stateful `WS /v1/chat` WebSocket that relays every `AgentEvent` as a typed JSON envelope (`text_chunk`/`thinking`/`tool_call_executing`/`tool_call_result`/`assistant_response`/`usage`/`error`/`done`) for real-time browser UIs. REST is stateless from the caller's POV; the WebSocket keeps per-connection history (`{"content": "..."}` sends a turn, `{"action": "clear"}` resets). A single process-private `Agent` is shared across requests and serialised with a `threading.RLock`; the WebSocket holds the lock for the duration of each run. Auto-approves permissions (no human), but offline/dry-run still apply. | | [oli_bot/agent.py](oli_bot/agent.py) | `Agent` — mode + system prompt owner; orchestrates the tool-calling loop and streams typed events (`TextChunk`, `ThinkingChunk`, `ToolCallChunk`, `ToolCallExecuting`, `ToolCallResult`, `StreamChunk`, `UsageEvent`, `Error`, `Done`). Aggregates per-call `UsageChunk`s from each backend round into a single per-run `UsageEvent`. Also hosts `sanitize_tool_history`, `_merge_usage`, `stream_sub_agent_run`, and `AgentPool` (built from [oli_bot/agents.yaml](oli_bot/agents.yaml) when `--use-pool` is set). | | [oli_bot/backends/](oli_bot/backends/) | Backend package — `ModelBackend` ABC, `OllamaBackend`, `OpenAIBackend`, `HuggingFaceBackend`, `TransformersBackend`, and the `create_model_backend()` factory. Also hosts the shared `_StreamingThinkParser` and per-backend message formatting (Ollama native `images`, OpenAI `image_url` or Bedrock-native blocks via `openai_vision_style`, textual placeholder for text-only backends). Every backend surfaces a trailing `UsageChunk`: exact counts from provider usage where available (OpenAI `usage`/`stream_options`, Ollama `prompt_eval_count`/`eval_count`, HF `usage`), else a `~chars/4` estimate via `estimate_tokens`. See [docs/BACKENDS.md](docs/BACKENDS.md). | | [oli_bot/screens/](oli_bot/screens/) | All `ModalScreen` subclasses: `PermissionScreen`, `ConfirmScreen`, `ModelPickerScreen`, `ServerListScreen`, `MCPSetupScreen`, `SessionListScreen`, `WorkspaceListScreen`, `SubAgentViewScreen`, `ConfigScreen`, `InputPromptScreen`, plus `taglines.py` / `todo_widget.py`. | diff --git a/oli_bot/api_server.py b/oli_bot/api_server.py index 1875abf..0ce1328 100644 --- a/oli_bot/api_server.py +++ b/oli_bot/api_server.py @@ -1,9 +1,11 @@ """OpenAI-compatible REST API over the oli agent harness. Serves ``/v1/models`` and ``/v1/chat/completions`` (streaming + non-streaming) -using FastAPI, plugging the same ``Agent`` tool loop that powers the TUI into -any workflow that speaks the OpenAI wire protocol (the ``openai`` Python SDK, -curl, or any other HTTP client). +with FastAPI, plus a stateful ``/v1/chat`` WebSocket that relays every agent +event as a typed JSON frame for real-time browser UIs. All routes plug the +same ``Agent`` tool loop that powers the TUI into any client that speaks the +OpenAI wire protocol (the ``openai`` Python SDK, curl, or any other HTTP +client). The server is stateless from the caller's perspective: each ``/v1/chat/completions`` request carries the full message history, mirroring @@ -19,6 +21,7 @@ """ import base64 +import dataclasses import json import logging import threading @@ -29,10 +32,10 @@ from typing import Any, AsyncIterator, Dict, List, Optional from art import text2art -from fastapi import FastAPI +from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.responses import JSONResponse, StreamingResponse -from .agent import Agent +from .agent import Agent, AgentEvent from .backends import create_model_backend, ModelBackend from .config import AppConfig, configs from .logger import setup_logging @@ -44,6 +47,10 @@ ImageAttachment, Message, StreamChunk, + ThinkingChunk, + ToolCallExecuting, + ToolCallResult, + UsageEvent, ChatCompletionMessage, ChatCompletionRequest, ) @@ -343,6 +350,110 @@ def chunk(delta: Dict[str, Any], finish_reason: Any = None) -> str: yield "data: [DONE]\n\n" +# --- WebSocket ------------------------------------------------------------- # + + +def _event_to_frame(event: AgentEvent) -> Dict[str, Any]: + """Convert an ``AgentEvent`` into a typed JSON envelope for the browser. + + The ``type`` field lets the client distinguish event kinds and render them + differently (streamed text, thinking blocks, tool calls, errors, etc.). + """ + if isinstance(event, StreamChunk): + return {"type": "text_chunk", "data": {"text": event.text}} + if isinstance(event, ThinkingChunk): + return {"type": "thinking", "data": {"text": event.text}} + if isinstance(event, ToolCallExecuting): + return { + "type": "tool_call_executing", + "data": {"name": event.name, "parameters": event.parameters}, + } + if isinstance(event, ToolCallResult): + return { + "type": "tool_call_result", + "data": {"name": event.name, "result": event.result}, + } + if isinstance(event, AssistantResponse): + return {"type": "assistant_response", "data": {"content": event.content}} + if isinstance(event, UsageEvent): + return {"type": "usage", "data": dataclasses.asdict(event.usage)} + if isinstance(event, Error): + return {"type": "error", "data": {"message": event.message}} + if isinstance(event, Done): + return {"type": "done", "data": {"full_text": event.full_text}} + logger.warning("Unknown agent event in websocket relay: %r", event) + return {"type": "unknown", "data": {"event": repr(event)}} + + +@app.websocket("/v1/chat") +async def websocket_chat(websocket: WebSocket) -> None: + """Stateful WebSocket chat endpoint. + + The server keeps a per-connection ``messages`` list so a client just sends + the next user turn (``{"content": "..."}``) and receives every ``AgentEvent`` + back as a typed JSON frame. ``{"action": "clear"}`` resets the history. + Runs are serialized on ``app.state.lock`` like the REST endpoints. + """ + await websocket.accept() + messages: List[Message] = [] + try: + await websocket.send_json({"type": "connected", "data": {}}) + while True: + raw = await websocket.receive_text() + try: + data = json.loads(raw) + except json.JSONDecodeError: + await websocket.send_json( + {"type": "error", "data": {"message": "Invalid JSON payload"}} + ) + continue + if not isinstance(data, dict): + await websocket.send_json( + {"type": "error", "data": {"message": "Expected a JSON object"}} + ) + continue + + if data.get("action") == "clear": + messages = [] + await websocket.send_json({"type": "cleared", "data": {}}) + continue + + content = data.get("content") + if not content or not str(content).strip(): + await websocket.send_json( + {"type": "error", "data": {"message": "Empty message"}} + ) + continue + + messages.append(Message(role="user", content=str(content))) + + try: + tools = await _resolve_tools(app.state.agent) + except Exception as e: + logger.warning("Failed to list tools: %s", e) + tools = None + + with app.state.lock: + try: + async for event in app.state.agent.process( + messages, tools=tools, confirm_callback=_api_confirm + ): + await websocket.send_json(_event_to_frame(event)) + if isinstance(event, Done) and event.full_text: + messages.append( + Message(role="assistant", content=event.full_text) + ) + except WebSocketDisconnect: + raise + except Exception as e: + logger.exception("Agent process failed over websocket: %s", e) + await websocket.send_json( + {"type": "error", "data": {"message": str(e)}} + ) + except WebSocketDisconnect: + logger.debug("WebSocket client disconnected from /v1/chat") + + # --- Routes ----------------------------------------------------------------- # diff --git a/tests/test_api_server.py b/tests/test_api_server.py index ae9da01..21d1e19 100644 --- a/tests/test_api_server.py +++ b/tests/test_api_server.py @@ -55,6 +55,17 @@ async def stream_generate(self, messages, tools=None): yield # pragma: no cover +class _HistoryProbeStub: + model = "stub-probe" + + def __init__(self): + self.seen = [] + + async def stream_generate(self, messages, tools=None): + self.seen.append(len(messages)) + yield TextChunk(f"n={len(messages)}") + + class _StubMCP: async def call_tool(self, name, params, confirm_callback=None, **kwargs): return "result-of-echo" @@ -74,13 +85,16 @@ def _make_harness(backend) -> Agent: config=config, ) - def _echo(x: dict) -> str: - return f"echoed-{x.get('x')}" + def _echo(x: int) -> str: + return f"echoed-{x}" builtin.register_tool( name="echo", description="Echo back the input.", - parameters={"type": "object", "properties": {"x": {"type": "integer"}}}, + parameters={ + "type": "object", + "properties": {"x": {"type": "integer"}}, + }, handler=_echo, ) @@ -135,6 +149,16 @@ def _text_chunks(body: str) -> list[str]: return chunks +def _recv_until_done(ws) -> list[dict]: + """Collect WebSocket frames until a ``done`` frame arrives.""" + frames = [] + while True: + frame = ws.receive_json() + frames.append(frame) + if frame.get("type") == "done": + return frames + + # --------------------------------------------------------------------------- # # /v1/models # # --------------------------------------------------------------------------- # @@ -262,3 +286,107 @@ def test_health(api): resp = api.client.get("/health") assert resp.status_code == 200 assert resp.json() == {"status": "ok"} + + +# --------------------------------------------------------------------------- # +# /v1/chat (WebSocket) # +# --------------------------------------------------------------------------- # + + +def test_websocket_chat(api): + with api.client.websocket_connect("/v1/chat") as ws: + assert ws.receive_json() == {"type": "connected", "data": {}} + ws.send_json({"content": "hi"}) + frames = _recv_until_done(ws) + types = [f["type"] for f in frames] + assert "text_chunk" in types + assert ( + "".join(f["data"]["text"] for f in frames if f["type"] == "text_chunk") + == "Hello from stub" + ) + done = frames[-1] + assert done["type"] == "done" + assert done["data"]["full_text"] == "Hello from stub" + + +def test_websocket_tool_events(api): + api.reset(_ToolThenTextStub()) + with api.client.websocket_connect("/v1/chat") as ws: + assert ws.receive_json()["type"] == "connected" + ws.send_json({"content": "use the tool"}) + frames = _recv_until_done(ws) + types = [f["type"] for f in frames] + assert "tool_call_executing" in types + assert "tool_call_result" in types + exec_frame = next(f for f in frames if f["type"] == "tool_call_executing") + assert exec_frame["data"]["name"] == "builtin__echo" + assert exec_frame["data"]["parameters"] == {"x": 1} + result_frame = next(f for f in frames if f["type"] == "tool_call_result") + assert result_frame["data"]["result"] == "echoed-1" + assert frames[-1]["type"] == "done" + assert frames[-1]["data"]["full_text"] == "done after tool" + + +def test_websocket_stateful_history(api): + probe = _HistoryProbeStub() + api.reset(probe) + with api.client.websocket_connect("/v1/chat") as ws: + assert ws.receive_json()["type"] == "connected" + ws.send_json({"content": "first"}) + frames = _recv_until_done(ws) + # stream_generate sees [system header, user1] on the first turn + assert frames[0]["data"]["text"] == "n=2" + ws.send_json({"content": "second"}) + frames = _recv_until_done(ws) + # system + user1 + assistant1 + user2 = 4 messages on the second turn + assert frames[0]["data"]["text"] == "n=4" + assert probe.seen == [2, 4] + + +def test_websocket_clear(api): + probe = _HistoryProbeStub() + api.reset(probe) + with api.client.websocket_connect("/v1/chat") as ws: + assert ws.receive_json()["type"] == "connected" + ws.send_json({"content": "first"}) + _recv_until_done(ws) + ws.send_json({"action": "clear"}) + assert ws.receive_json() == {"type": "cleared", "data": {}} + ws.send_json({"content": "after clear"}) + frames = _recv_until_done(ws) + # history was reset: back down to system + user = 2 messages + assert frames[0]["data"]["text"] == "n=2" + assert probe.seen == [2, 2] + + +def test_websocket_error(api): + api.reset(_RaisingStub()) + with api.client.websocket_connect("/v1/chat") as ws: + assert ws.receive_json()["type"] == "connected" + ws.send_json({"content": "hi"}) + frames = _recv_until_done(ws) + types = [f["type"] for f in frames] + assert "error" in types + assert frames[-1]["type"] == "done" + assert frames[-1]["data"]["full_text"] == "" # empty Done on failure + + +def test_websocket_invalid_json(api): + with api.client.websocket_connect("/v1/chat") as ws: + assert ws.receive_json()["type"] == "connected" + ws.send_text("not json") + frame = ws.receive_json() + assert frame["type"] == "error" + assert frame["data"]["message"] == "Invalid JSON payload" + # connection stays alive for a valid message after the bad one + ws.send_json({"content": "hi"}) + assert _recv_until_done(ws)[-1]["type"] == "done" + + +def test_websocket_empty_message(api): + with api.client.websocket_connect("/v1/chat") as ws: + assert ws.receive_json()["type"] == "connected" + ws.send_json({"content": " "}) + frame = ws.receive_json() + assert frame["type"] == "error" + assert frame["data"]["message"] == "Empty message" From 93edd4847f5386154c694bd76a01fc1ced250e2d Mon Sep 17 00:00:00 2001 From: jay Date: Sun, 13 Sep 2026 11:26:14 -0700 Subject: [PATCH 3/3] refactor managers to live in their respective packages --- AGENTS.md | 5 +- oli_bot/agent.py | 2 +- .../upstream_manager.py} | 4 +- oli_bot/chat.py | 60 +++++++++---------- .../manager.py} | 10 ++-- oli_bot/screens/server_list.py | 2 +- tests/test_config_settings.py | 6 +- tests/test_model_registration.py | 14 ++--- tests/test_model_registration_manual.py | 18 +++--- 9 files changed, 59 insertions(+), 62 deletions(-) rename oli_bot/{server_manager.py => backends/upstream_manager.py} (99%) rename oli_bot/{profile_manager.py => profiles/manager.py} (93%) diff --git a/AGENTS.md b/AGENTS.md index 4926c4f..3c7424c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -19,13 +19,12 @@ This file is the top-level orientation map. Depth lives under [docs/](docs/). | [oli_bot/models.py](oli_bot/models.py) | Shared dataclasses: `Message`, `ToolCall`, `ModelResponse`, `HostConfig`, `MCPServerConfig`, `ProfileData`, `SubAgentRun`, `ImageAttachment`, `TodoItem` / `TodoListState`, plus `AgentEvent` variants and the `AgentRole` enum. Token accounting lives here too: `Usage` (prompt/completion/`estimated` flag), the per-call `UsageChunk` stream event, and the per-run `UsageEvent`. `Message.images` is in-memory only (dropped on session save); `ModelResponse.usage` is optionally set by backends. | | [oli_bot/config.py](oli_bot/config.py) | `AppConfig` — `pydantic_settings.BaseSettings`. Env vars prefixed `OLI_`, plus `.env` support and `OLI_TRUNCATION_SMALL` / `_LARGE` aliases via `AliasChoices`. Module-level `configs = AppConfig()` singleton. See [docs/CONFIGURE.md](docs/CONFIGURE.md). | | [oli_bot/settings.py](oli_bot/settings.py) | `SettingsManager` — load/save/merge `~/.config/oli/settings.json`; precedence `settings.json` > `OLI_*` env > SDK-standard env (`OPENAI_API_KEY`, `OPENAI_BASE_URL`, `HUGGINGFACE_API_KEY`, `HF_TOKEN`) > declared defaults. Empty API-key strings in JSON fall through to env. | -| [oli_bot/profiles/](oli_bot/profiles/) | `ProfileManifest` / `PermissionsManifest` (Pydantic) in `schema.py`; `ProfilePermissionEnforcer` (layered allow/deny glob patterns, base-profile inheritance, deny-overrides-allow) in `permissions.py`. Built-in profiles ship as sibling directories. | -| [oli_bot/profile_manager.py](oli_bot/profile_manager.py) | `ProfileManager` — profile CRUD, manifest loading, circular-dependency detection; delegates enforcement to `oli_bot/profiles/permissions.py`. | +| [oli_bot/profiles/](oli_bot/profiles/) | `ProfileManifest` / `PermissionsManifest` (Pydantic) in `schema.py`; `ProfilePermissionEnforcer` (layered allow/deny glob patterns, base-profile inheritance, deny-overrides-allow) in `permissions.py`; `ProfileManager` (profile CRUD, manifest loading, circular-dependency detection) in `manager.py`. Built-in profiles ship as sibling directories. | | [oli_bot/mcp_client.py](oli_bot/mcp_client.py) | `MCPClientManager` — MCP server lifecycle (stdio/http via v2 `mcp.client.Client`, `mode="auto"` handshake), tool discovery + invocation, per-server tool-list cache, offline gating. Uses v2 snake_case fields (`Tool.input_schema`, `CallToolResult.is_error`/`structured_content`). | | [oli_bot/tools/manager.py](oli_bot/tools/manager.py) | `BuiltinToolManager` — registration, profile + session permission gating, dry-run gating, offline gating, and `TruncationManager` post-processing. Awaits coroutine handlers. | | [oli_bot/tools/](oli_bot/tools/) | Tool handlers: `files.py` (read/write/edit + `view_image` via Pillow), `directories.py` (glob/grep/list_directory/tree — filesystem work runs via `asyncio.to_thread` / `create_subprocess_exec`), `web.py` (search + fetch + specialised searches, all guarded by `_check_ssrf`), `shell.py` (allowlisted `run_command`, including read-only `git`), `parsing.py` (`compare`), `memory.py` (`think`, `todowrite`, `notebook`), `truncation.py` (per-tier char budgets), `permissions.py` (sensitive-path detection). See [docs/TOOLS.md](docs/TOOLS.md). | | [oli_bot/sessions.py](oli_bot/sessions.py) | `Session` (permission gating) + `ConversationStore` (per-server JSON persistence under `~/.config/oli/sessions//`) + `WorkspaceManager`. `save_session()` returns the (possibly new) id so callers can rebind after a corrupt-file rewrite. Persisted messages preserve `tool_call_id`; loads pass through `sanitize_tool_history` so poisoned histories self-heal. | -| [oli_bot/server_manager.py](oli_bot/server_manager.py) | `ServerManager` — multi-server lifecycle persisted to `ollama_hosts.json`, URL validation. | +| [oli_bot/backends/upstream_manager.py](oli_bot/backends/upstream_manager.py) | `UpstreamManager` — multi-server lifecycle persisted to `hosts.json`, URL validation. | | [oli_bot/voice.py](oli_bot/voice.py) | `VoiceEngine` — optional, lazy-loaded mic → STT → TTS engine for the `/voice` command (faster-whisper, Piper TTS, WebRTC VAD, pyaudio). All I/O is blocking; `chat.py` calls it via `asyncio.to_thread`. `record()` accepts a `threading.Event` so `chat.py` can interrupt an in-progress recording the instant voice mode is toggled off, instead of waiting out the silence/max-duration timeout. All seven tunables (whisper/piper models, sample rate, VAD frame duration, VAD aggressiveness, silence timeout, max record seconds) are `AppConfig` fields (`OLI_VOICE_*` env / `settings.json` `voice` section / `/config` screen); `chat.py` passes them explicitly when constructing the engine, and saving `/config` drops the engine so the next `/voice` picks up new values. | | [oli_bot/logger.py](oli_bot/logger.py) | Centralised NDJSON file logging (rotating, 10 MB × 5) under `AppConfig.log_file`. Deliberately no console handler — stray writes would corrupt the Textual TUI. | diff --git a/oli_bot/agent.py b/oli_bot/agent.py index 986f781..ca9eaf5 100644 --- a/oli_bot/agent.py +++ b/oli_bot/agent.py @@ -19,7 +19,7 @@ ) from .config import AppConfig, configs from .mcp_client import MCPClientManager -from .profile_manager import ProfileData, ProfileManager +from .profiles.manager import ProfileData, ProfileManager from .profiles.permissions import ProfilePermissionEnforcer from .models import ( ToolCallExecuting, diff --git a/oli_bot/server_manager.py b/oli_bot/backends/upstream_manager.py similarity index 99% rename from oli_bot/server_manager.py rename to oli_bot/backends/upstream_manager.py index 3bebded..3adbdf5 100644 --- a/oli_bot/server_manager.py +++ b/oli_bot/backends/upstream_manager.py @@ -8,7 +8,7 @@ from ollama import AsyncClient as OllamaAsyncClient -from .models import HostConfig +from ..models import HostConfig logger = logging.getLogger(__name__) @@ -16,7 +16,7 @@ CONFIG_FILE = Path.joinpath(CONFIG_DIR, "hosts.json") -class ServerManager: +class UpstreamManager: def __init__(self, config_path: str = CONFIG_FILE): self.config_path = config_path self.servers: List[HostConfig] = [] diff --git a/oli_bot/chat.py b/oli_bot/chat.py index 21890dc..6a9b660 100644 --- a/oli_bot/chat.py +++ b/oli_bot/chat.py @@ -57,7 +57,7 @@ from .tools.manager import BuiltinToolManager from .tools.memory import _current_sub_run from .mcp_client import MCPClientManager -from .server_manager import ServerManager +from .backends.upstream_manager import UpstreamManager from .sessions import ( SCOPE_WORKSPACE_SENSITIVE, Session, @@ -269,9 +269,9 @@ def __init__( if use_pool: self.config.use_agent_pool = True self.model_size: str = "large" - self.server_manager = ServerManager() - self.server_manager.seed_default(base_url) - active = self.server_manager.get_active() + self.upstream_manager = UpstreamManager() + self.upstream_manager.seed_default(base_url) + active = self.upstream_manager.get_active() effective_url = active.url if active else base_url cli_model = model @@ -409,7 +409,7 @@ def compose(self) -> ComposeResult: yield Static(id="status-bar") def _render_welcome_panel(self) -> Panel: - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() backend_label = ( f"[bold {PRIMARY_HEX}]{self.config.backend}[/bold {PRIMARY_HEX}]" ) @@ -564,7 +564,7 @@ def on_click(self, event: Click) -> None: self.notify("Copied to clipboard", timeout=2) def _server_name(self) -> str: - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() return active.name if active else "default" def _save_session(self) -> None: @@ -588,7 +588,7 @@ def _save_session(self) -> None: self.current_session_id = new_id def update_header(self) -> None: - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() server_part = f" {active.name}" if active else "" model_part = ( str(self.backend.model) @@ -1229,7 +1229,7 @@ def _workspace_unset(self) -> None: ) def _handle_context(self) -> None: - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() backend_label = f"[bold green]{self.config.backend}[/bold green]" if self.config.backend == "ollama": server_url = active.url if active else self.backend.base_url @@ -1525,7 +1525,7 @@ async def _server_add(self, name: str, url: str) -> None: # Other upstream servers won't get this check when added. # TODO: consider adding a generic ping/healthcheck for other backends if they support it. if self.config.backend == "ollama": - ok, err = await ServerManager.validate_ollama_url(url) + ok, err = await UpstreamManager.validate_ollama_url(url) if not ok: self._add_message( "System", @@ -1534,7 +1534,7 @@ async def _server_add(self, name: str, url: str) -> None: return try: - is_first = self.server_manager.add_server(name, url) + is_first = self.upstream_manager.add_server(name, url) if is_first: self.backend.set_base_url(url) self.update_header() @@ -1545,7 +1545,7 @@ async def _server_add(self, name: str, url: str) -> None: self._add_message("System", f"[red]{e}[/red]") def _server_list(self) -> None: - servers = self.server_manager.list_servers() + servers = self.upstream_manager.list_servers() if not servers: self._add_message("System", "No Ollama servers configured.") return @@ -1557,9 +1557,9 @@ def _server_list_selected(self, name: Optional[str]) -> None: def _server_remove(self, name: str) -> None: try: - removed = self.server_manager.remove_server(name) + removed = self.upstream_manager.remove_server(name) if removed.active: - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() if active: self.backend.set_base_url(active.url) self.update_header() @@ -1569,7 +1569,7 @@ def _server_remove(self, name: str) -> None: def _server_default(self, name: str) -> None: try: - config = self.server_manager.switch_server(name) + config = self.upstream_manager.switch_server(name) self.backend.set_base_url(config.url) model_to_use = config.large_model or config.default_model if model_to_use: @@ -1583,8 +1583,8 @@ def _server_default(self, name: str) -> None: def _server_set_default_model(self, name: str, model: str) -> None: try: - self.server_manager.set_default_model(name, model) - active = self.server_manager.get_active() + self.upstream_manager.set_default_model(name, model) + active = self.upstream_manager.get_active() if active and active.name == name: self._switch_model(model) if self.config.backend == "ollama": @@ -1599,7 +1599,7 @@ def _server_set_default_model(self, name: str, model: str) -> None: def _server_switch(self, name: str) -> None: try: self._save_session() - config = self.server_manager.switch_server(name) + config = self.upstream_manager.switch_server(name) self.backend.set_base_url(config.url) model_to_use = config.large_model or config.default_model if model_to_use: @@ -1671,7 +1671,7 @@ def _get_large_model(self) -> str: return self.config.huggingface_model if self.config.backend == "transformers": return self.config.transformers_model - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() if active and active.large_model: return active.large_model return self.config.ollama_model @@ -1685,7 +1685,7 @@ def _get_small_model(self) -> str: return ( self.config.transformers_small_model or self.config.transformers_model ) - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() if active and active.small_model: return active.small_model return self.config.ollama_small_model @@ -1731,28 +1731,28 @@ def _handle_model_switch(self, text: str) -> None: ) def _model_set_large(self, model: str) -> None: - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() if not active: self._add_message( "System", "[red]No active server to set model for.[/red]", ) return - self.server_manager.set_large_model(active.name, model) + self.upstream_manager.set_large_model(active.name, model) self._persist_model_to_settings(large=model) self._switch_model(model) self.model_size = "large" self._builtin_tools.model_tier = self.model_size def _model_set_small(self, model: str) -> None: - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() if not active: self._add_message( "System", "[red]No active server to set model for.[/red]", ) return - self.server_manager.set_small_model(active.name, model) + self.upstream_manager.set_small_model(active.name, model) self._persist_model_to_settings(small=model) self._switch_model(model) self.model_size = "small" @@ -1799,7 +1799,7 @@ def _model_add(self, args: str) -> None: Usage: /model add [--large|--small|--default] Example: /model add gpt4-turbo --large """ - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() if not active: self._add_message( "System", @@ -1832,7 +1832,7 @@ def _model_add(self, args: str) -> None: return try: - self.server_manager.add_model(active.name, model_name, model_name, tier) + self.upstream_manager.add_model(active.name, model_name, model_name, tier) if tier: self._add_message( "System", @@ -1852,7 +1852,7 @@ def _model_remove(self, model_name: str) -> None: Usage: /model remove Example: /model remove gpt4-turbo """ - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() if not active: self._add_message( "System", @@ -1868,7 +1868,7 @@ def _model_remove(self, model_name: str) -> None: return try: - removed_model = self.server_manager.remove_model(active.name, model_name) + removed_model = self.upstream_manager.remove_model(active.name, model_name) self._add_message( "System", f"Model [bold]{model_name}[/bold] (mapped to [bold]{removed_model}[/bold]) removed.", @@ -1878,7 +1878,7 @@ def _model_remove(self, model_name: str) -> None: def _model_list(self) -> None: """List all registered models for the active server.""" - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() if not active: self._add_message( "System", @@ -1887,7 +1887,7 @@ def _model_list(self) -> None: return try: - models = self.server_manager.list_models(active.name) + models = self.upstream_manager.list_models(active.name) if not models: self._add_message( "System", @@ -1967,7 +1967,7 @@ async def _handle_config(self) -> None: ) def _rebuild_backend(self, announce: bool = True) -> None: - active = self.server_manager.get_active() + active = self.upstream_manager.get_active() url = active.url if active else self.config.ollama_base_url if self.config.backend == "openai": url = self.config.openai_base_url diff --git a/oli_bot/profile_manager.py b/oli_bot/profiles/manager.py similarity index 93% rename from oli_bot/profile_manager.py rename to oli_bot/profiles/manager.py index 6c2ed25..1c81987 100644 --- a/oli_bot/profile_manager.py +++ b/oli_bot/profiles/manager.py @@ -3,14 +3,14 @@ from typing import List import logging -from .models import ProfileData -from .profiles.schema import ( +from ..models import ProfileData +from .schema import ( ProfileManifest, auto_generate_manifest, load_profile_manifest, dump_profile_manifest, ) -from .profiles.permissions import ProfilePermissionEnforcer +from .permissions import ProfilePermissionEnforcer logger = logging.getLogger(__name__) @@ -20,9 +20,7 @@ class ProfileManager: def __init__(self, profiles_dir: str = None): if not profiles_dir: - self.profiles_dir: Path = Path.joinpath( - Path(os.path.abspath(os.path.dirname(__file__))), "profiles" - ) + self.profiles_dir: Path = Path(os.path.abspath(os.path.dirname(__file__))) else: self.profiles_dir = Path(profiles_dir) diff --git a/oli_bot/screens/server_list.py b/oli_bot/screens/server_list.py index 93330df..6545641 100644 --- a/oli_bot/screens/server_list.py +++ b/oli_bot/screens/server_list.py @@ -7,7 +7,7 @@ from textual.screen import ModalScreen from textual.widgets import Label, ListItem, ListView -from ..server_manager import HostConfig +from ..models import HostConfig class ServerListScreen(ModalScreen[Optional[str]]): diff --git a/tests/test_config_settings.py b/tests/test_config_settings.py index f7f22f0..a6c32cd 100644 --- a/tests/test_config_settings.py +++ b/tests/test_config_settings.py @@ -5,7 +5,7 @@ from oli_bot.config import AppConfig from oli_bot.chat import OliBot from oli_bot.models import HostConfig -from oli_bot.server_manager import ServerManager +from oli_bot.backends.upstream_manager import UpstreamManager from oli_bot.settings import SettingsManager @@ -298,7 +298,7 @@ def test_runtime_settings_sync_covers_all_appconfig_fields(tmp_path): api_profile="analyst", api_mode="ask", ) - server_manager = ServerManager(config_path=str(tmp_path / "hosts.json")) + server_manager = UpstreamManager(config_path=str(tmp_path / "hosts.json")) server_manager.servers.append( HostConfig(name="test", url="http://localhost:11434", active=True) ) @@ -310,7 +310,7 @@ def test_runtime_settings_sync_covers_all_appconfig_fields(tmp_path): bot.settings["workspace"]["max_workspaces"] = 37 bot.settings["session"]["auto_save"] = False bot.settings["session"]["resume_prompt"] = False - bot.server_manager = server_manager + bot.upstream_manager = server_manager bot._sync_settings_from_runtime() diff --git a/tests/test_model_registration.py b/tests/test_model_registration.py index ad6e3ff..ac77352 100644 --- a/tests/test_model_registration.py +++ b/tests/test_model_registration.py @@ -9,7 +9,7 @@ import pytest from oli_bot.models import HostConfig -from oli_bot.server_manager import ServerManager +from oli_bot.backends.upstream_manager import UpstreamManager class TestModelRegistration: @@ -25,8 +25,8 @@ def temp_config(self): @pytest.fixture def manager(self, temp_config): - """Create a ServerManager instance with temp config.""" - mgr = ServerManager(temp_config) + """Create a UpstreamManager instance with temp config.""" + mgr = UpstreamManager(temp_config) mgr.servers.append( HostConfig(name="test-server", url="http://localhost:11434", active=True) ) @@ -97,7 +97,7 @@ def test_get_model(self, manager): def test_persistence_to_disk(self, temp_config): """Test that registered models persist to hosts.json.""" - mgr1 = ServerManager(temp_config) + mgr1 = UpstreamManager(temp_config) mgr1.servers.append( HostConfig(name="test-server", url="http://localhost:11434", active=True) ) @@ -105,7 +105,7 @@ def test_persistence_to_disk(self, temp_config): mgr1.add_model("test-server", "phi", "phi-2", tier="small") # Load in new manager instance - mgr2 = ServerManager(temp_config) + mgr2 = UpstreamManager(temp_config) models = mgr2.list_models("test-server") assert models["gpt4"] == "gpt-4-turbo" @@ -130,7 +130,7 @@ def test_backward_compatibility_old_config(self, temp_config): Path(temp_config).write_text(json.dumps(old_data)) # Load should work without errors - manager = ServerManager(temp_config) + manager = UpstreamManager(temp_config) assert len(manager.servers) == 1 server = manager.get_active() @@ -142,7 +142,7 @@ def test_backward_compatibility_old_config(self, temp_config): def test_multiple_servers_isolated_models(self, temp_config): """Test that models are isolated per server.""" - manager = ServerManager(temp_config) + manager = UpstreamManager(temp_config) manager.servers.append( HostConfig(name="server1", url="http://host1:11434", active=True) ) diff --git a/tests/test_model_registration_manual.py b/tests/test_model_registration_manual.py index eff84a9..d8863f4 100644 --- a/tests/test_model_registration_manual.py +++ b/tests/test_model_registration_manual.py @@ -8,7 +8,7 @@ import tempfile from pathlib import Path from oli_bot.models import HostConfig -from oli_bot.server_manager import ServerManager +from oli_bot.backends.upstream_manager import UpstreamManager def test_hostconfig_has_registered_models(): @@ -28,7 +28,7 @@ def test_add_model_basic(): config_path = f.name try: - manager = ServerManager(config_path) + manager = UpstreamManager(config_path) manager.servers.append( HostConfig(name="test-server", url="http://localhost:11434", active=True) ) @@ -50,7 +50,7 @@ def test_add_model_with_tier(): config_path = f.name try: - manager = ServerManager(config_path) + manager = UpstreamManager(config_path) manager.servers.append( HostConfig(name="test-server", url="http://localhost:11434", active=True) ) @@ -72,7 +72,7 @@ def test_persistence(): try: # Create and populate - mgr1 = ServerManager(config_path) + mgr1 = UpstreamManager(config_path) mgr1.servers.append( HostConfig(name="test-server", url="http://localhost:11434", active=True) ) @@ -80,7 +80,7 @@ def test_persistence(): mgr1.add_model("test-server", "phi", "phi-2", tier="small") # Reload in new instance - mgr2 = ServerManager(config_path) + mgr2 = UpstreamManager(config_path) models = mgr2.list_models("test-server") assert "gpt4" in models, "Model not persisted" @@ -116,7 +116,7 @@ def test_backward_compatibility(): Path(config_path).write_text(json.dumps(old_data)) # Load should work without errors - manager = ServerManager(config_path) + manager = UpstreamManager(config_path) assert len(manager.servers) == 1, "Failed to load old config" server = manager.get_active() @@ -135,7 +135,7 @@ def test_duplicate_error(): config_path = f.name try: - manager = ServerManager(config_path) + manager = UpstreamManager(config_path) manager.servers.append( HostConfig(name="test-server", url="http://localhost:11434", active=True) ) @@ -159,7 +159,7 @@ def test_remove_model(): config_path = f.name try: - manager = ServerManager(config_path) + manager = UpstreamManager(config_path) manager.servers.append( HostConfig(name="test-server", url="http://localhost:11434", active=True) ) @@ -182,7 +182,7 @@ def test_multiple_servers(): config_path = f.name try: - manager = ServerManager(config_path) + manager = UpstreamManager(config_path) manager.servers.append( HostConfig(name="server1", url="http://host1:11434", active=True) )