diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c9d3c7..d8cf4dc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,10 +10,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - **Automatic provider-icon sync** — new `_sync_model_icons()` method writes provider icons directly into Open WebUI's Models database so they appear in the UI; controlled by the `SYNC_PROVIDER_ICONS` valve (default: enabled). Models with a manually-set icon are never overwritten +- **`_is_owui_managed_icon()` helper** — distinguishes OWUI-default icons (`data:` URLs) and our own provider icons from user-set custom icons, enabling safe icon updates without clobbering user customisations ### Fixed - **Icon sync: correct prefixed model IDs** — `_sync_model_icons()` now discovers the pipe's `function_id` via `type(self).__module__` and writes DB records with the full prefixed ID (e.g. `openrouter_pipe.openai/gpt-4o`) matching what Open WebUI's frontend requests at `/models/model/profile/image` +- **Icon sync: icons now actually appear in the UI** — five bugs prevented provider icons from ever showing after the first pipe load: + - *Wrong skip condition* — `if existing_icon:` skipped any model with *any* icon (including the generic `data:` SVG that OWUI assigns by default), so provider icons were never applied; fixed to skip only user-set custom URLs + - *Race condition* — `_sync_model_icons()` was called before `pipes()` returned, i.e. before OWUI registered the models; OWUI then overwrote the early insert with its own default icon; fixed by also calling `_sync_model_icons()` on cache-hit paths (until all models are confirmed synced) + - *Exception swallowed retry* — DB errors added the model to `_icons_synced` anyway, permanently preventing retry; removed the erroneous add + - *Insert marked as synced prematurely* — after `insert_new_model` the model was marked synced even though OWUI could overwrite it; the insert path no longer updates `_icons_synced` + - *User params clobbered* — `update_model_by_id` used an empty `ModelParams()`, erasing user-configured temperature/system-prompt/etc.; now preserves `existing.params` +- **Icon sync: `function_id` cached at init** — `type(self).__module__` is evaluated once in `__init__` instead of on every `_sync_model_icons()` call - **Streaming status event** — the "done" status event is now correctly emitted at the end of streaming responses (async generator wrapper replaces sync generator that could not `await`) - **Dead provider-icon code removed** — `info.meta.profile_image_url` was included in model dicts returned by `pipes()` but Open WebUI ignores all fields except `id` and `name`; the field has been removed in favour of the new DB-sync approach - **`pipes()` response always closed** — added `finally: response.close()` to guarantee HTTP connections are returned to the session pool in all code paths (auth errors, JSON decode failures, unexpected exceptions) diff --git a/openrouter_pipe.py b/openrouter_pipe.py index b3ba401..c3c7c98 100644 --- a/openrouter_pipe.py +++ b/openrouter_pipe.py @@ -72,6 +72,21 @@ def _is_safe_url(url: str) -> bool: return isinstance(url, str) and url.lower().startswith(("http://", "https://")) +def _is_owui_managed_icon(url: str) -> bool: + """Return True if the icon URL was set by OWUI or our sync logic. + + data: URLs are the pipe's own SVG icon that OWUI assigns as default to all + manifold child models. openrouter.ai/images/models/ URLs are the provider + icons we write. Any other URL is assumed to be a user-set custom icon and + must not be overwritten. + """ + return ( + not url + or url.startswith("data:") + or url.startswith("https://openrouter.ai/images/models/") + ) + + def _insert_citations(text: str, citations: Optional[List[str]]) -> str: """Replace [n] references with markdown links (only safe HTTP URLs).""" if not citations or not text: @@ -240,6 +255,11 @@ def __init__(self) -> None: self._models_cache_key: str = "" # Track which model IDs already have icons synced (avoids repeated DB writes) self._icons_synced: set = set() + # Cache function_id once: OWUI sets __module__ to "function_{id}" at load time + _fm = type(self).__module__ or "" + self._function_id: Optional[str] = ( + _fm[len("function_"):] if _fm.startswith("function_") else None + ) if not self.valves.OPENROUTER_API_KEY: print("[OpenRouter Pipe] Warning: OPENROUTER_API_KEY not set") @@ -274,6 +294,11 @@ def pipes(self) -> List[dict]: # Return cached models if still valid if self._models_cache_valid() and self._models_cache is not None: + # Continue syncing icons on cache hits until all models are confirmed. + # This resolves the race condition where OWUI registers models (and may + # overwrite icons) only after the first pipes() call returns. + if self.valves.SYNC_PROVIDER_ICONS and len(self._icons_synced) < len(self._models_cache): + self._sync_model_icons(self._models_cache) return self._models_cache headers = self._build_headers(include_content_type=False) @@ -456,14 +481,18 @@ def _sync_model_icons(self, models: List[dict]) -> None: """Write provider icons into Open WebUI's Models DB. Open WebUI serves model icons from its database, not from the dicts - returned by ``pipes()``. Crucially, Open WebUI prefixes every pipe - model ID with ``{function_id}.`` (e.g. ``openrouter_pipe.openai/gpt-4o``) - and the frontend requests icons using that prefixed ID. This method - discovers the pipe's own *function_id* from ``type(self).__module__`` - (set to ``function_{id}`` by Open WebUI's module loader) and writes - DB records keyed by the full prefixed ID. - - Models that already have a custom icon in the DB are skipped. + returned by ``pipes()``. OWUI prefixes every pipe model ID with + ``{function_id}.`` (e.g. ``openrouter_pipe.openai/gpt-4o``) and the + frontend requests icons using that prefixed ID. + + Called both on cache miss and on subsequent cache hits (until all + models are confirmed synced). The cache-hit path is needed because + OWUI registers models *after* ``pipes()`` returns, potentially + overwriting any early insert with its own default icon; the second + call finds the models already in DB and updates them correctly. + + User-set custom icons (any URL that is not a ``data:`` URL and does not + start with ``https://openrouter.ai/images/models/``) are preserved. This is a best-effort operation — failures are silently logged. """ try: @@ -477,14 +506,10 @@ def _sync_model_icons(self, models: List[dict]) -> None: # Running outside Open WebUI (e.g. standalone tests) — skip silently return - # Discover the pipe's function_id. Open WebUI loads pipe modules as - # ``function_{function_id}`` so type(self).__module__ exposes it. - func_module = type(self).__module__ or "" - if func_module.startswith("function_"): - function_id = func_module[len("function_"):] - else: - # Cannot determine the pipe's function_id — skip icon sync + # function_id was resolved once in __init__ from type(self).__module__ + if not self._function_id: return + function_id = self._function_id for model in models: model_id = model.get("id", "") @@ -509,17 +534,28 @@ def _sync_model_icons(self, models: List[dict]) -> None: try: existing = Models.get_model_by_id(db_model_id) if existing: - # If model already has a custom icon, don't overwrite it existing_icon = "" if hasattr(existing, "meta") and existing.meta: existing_icon = ( getattr(existing.meta, "profile_image_url", "") or "" ) - if existing_icon: + + # Skip if icon is already the correct provider URL + if existing_icon == icon_url: + self._icons_synced.add(model_id) + continue + + # Skip if icon was set by the user (not by OWUI or our sync). + # data: URLs are OWUI defaults; openrouter.ai URLs are ours. + if existing_icon and not _is_owui_managed_icon(existing_icon): self._icons_synced.add(model_id) continue - # Update existing model with icon + # Proceed: icon is empty, an OWUI default, or one of our URLs + # Update existing model with icon, preserving user-set params + existing_params = ModelParams() + if hasattr(existing, "params") and existing.params: + existing_params = existing.params Models.update_model_by_id( db_model_id, ModelForm( @@ -530,26 +566,34 @@ def _sync_model_icons(self, models: List[dict]) -> None: else model.get("name", model_id) ), meta=ModelMeta(profile_image_url=icon_url), - params=ModelParams(), + params=existing_params, ), ) else: - # Insert new model record with icon - Models.insert_new_model( - ModelForm( - id=db_model_id, - name=model.get("name", model_id), - meta=ModelMeta(profile_image_url=icon_url), - params=ModelParams(), - ), - user_id="pipe:openrouter", - ) + # Model not yet in DB — best-effort early insert. + # OWUI will register models after pipes() returns and may + # overwrite this record, so do NOT mark as synced here. + # The next cache-hit call to _sync_model_icons will find the + # model in DB and update it correctly. + try: + Models.insert_new_model( + ModelForm( + id=db_model_id, + name=model.get("name", model_id), + meta=ModelMeta(profile_image_url=icon_url), + params=ModelParams(), + ), + user_id="pipe:openrouter", + ) + except Exception: + pass + continue # do not add to _icons_synced yet self._icons_synced.add(model_id) except Exception as exc: # Best-effort — don't let icon sync break model listing + # Do NOT add to _icons_synced: allow retry on next call print(f"[OpenRouter Pipe] Icon sync failed for {db_model_id}: {exc}") - self._icons_synced.add(model_id) # Don't retry on every refresh @staticmethod def get_provider_icon(provider: str) -> Optional[str]: diff --git a/test_pipe.py b/test_pipe.py index c546831..2fd99de 100644 --- a/test_pipe.py +++ b/test_pipe.py @@ -43,6 +43,7 @@ _insert_citations = mod._insert_citations _format_citation_list = mod._format_citation_list _OWUI_INTERNAL_KEYS = mod._OWUI_INTERNAL_KEYS +_is_owui_managed_icon = mod._is_owui_managed_icon # ── Helpers ─────────────────────────────────────────────────────────────────── _PASS = 0 @@ -1380,7 +1381,9 @@ async def _test_pipe_no_msgs_key(): # 25d. Valve default is True _assert(Pipe.Valves(OPENROUTER_API_KEY="k").SYNC_PROVIDER_ICONS is True, "SYNC_PROVIDER_ICONS default is True") -# 25e. Non-function module → returns early without DB calls +# 25e. No function_id → returns early without DB calls +# function_id is cached in __init__; pipes created in tests have _function_id=None +# because the test module name doesn't start with "function_". _pipe_nofunc = Pipe() _pipe_nofunc.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SYNC_PROVIDER_ICONS=True) _mock_Models_nf = MagicMock() @@ -1389,25 +1392,25 @@ async def _test_pipe_no_msgs_key(): _fake_owui_nf.ModelForm = MagicMock() _fake_owui_nf.ModelMeta = MagicMock() _fake_owui_nf.ModelParams = MagicMock() -_orig_module_nf = Pipe.__module__ try: sys.modules["open_webui.models.models"] = _fake_owui_nf - # Module name doesn't start with "function_" → should skip - Pipe.__module__ = "openrouter_pipe" + # _pipe_nofunc._function_id is None (no "function_" prefix at init time) → skip + _assert(_pipe_nofunc._function_id is None, "_sync_model_icons: _function_id is None outside OWUI") _pipe_nofunc._sync_model_icons([{"id": "openai/gpt-4o", "name": "GPT-4o"}]) _assert( not _mock_Models_nf.get_model_by_id.called, - "_sync_model_icons: skips DB when module is not function_*", + "_sync_model_icons: skips DB when _function_id is None", ) finally: - Pipe.__module__ = _orig_module_nf sys.modules.pop("open_webui.models.models", None) -# 25f. With function_ module → uses prefixed IDs for DB operations +# 25f. With function_id set → uses prefixed IDs; does NOT add to _icons_synced after insert +# (OWUI may overwrite the record after pipes() returns; confirmed on next call) _pipe_func = Pipe() _pipe_func.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SYNC_PROVIDER_ICONS=True) +_pipe_func._function_id = "openrouter_pipe" # simulate OWUI module naming _mock_Models_f = MagicMock() -_mock_Models_f.get_model_by_id.return_value = None # No existing record +_mock_Models_f.get_model_by_id.return_value = None # No existing record yet _mock_ModelForm_f = MagicMock() _mock_ModelMeta_f = MagicMock() _mock_ModelParams_f = MagicMock() @@ -1416,10 +1419,8 @@ async def _test_pipe_no_msgs_key(): _fake_owui_f.ModelForm = _mock_ModelForm_f _fake_owui_f.ModelMeta = _mock_ModelMeta_f _fake_owui_f.ModelParams = _mock_ModelParams_f -_orig_module_f = Pipe.__module__ try: sys.modules["open_webui.models.models"] = _fake_owui_f - Pipe.__module__ = "function_openrouter_pipe" _pipe_func._sync_model_icons([ {"id": "openai/gpt-4o", "name": "GPT-4o"}, {"id": "anthropic/claude-3.5-sonnet", "name": "Claude 3.5"}, @@ -1440,9 +1441,6 @@ async def _test_pipe_no_msgs_key(): "_sync_model_icons: insert_new_model called for new models", ) # Verify the ModelForm ID is prefixed - _insert_form = _mock_Models_f.insert_new_model.call_args_list[0].args[0] - _insert_id = _mock_ModelForm_f.call_args_list[0].kwargs.get("id", _mock_ModelForm_f.call_args_list[0].args[0] if _mock_ModelForm_f.call_args_list[0].args else "") - # ModelForm was called with id=prefixed_id _form_calls = _mock_ModelForm_f.call_args_list _form_ids = [c.kwargs.get("id", "") for c in _form_calls] _assert( @@ -1453,45 +1451,89 @@ async def _test_pipe_no_msgs_key(): "openrouter_pipe.anthropic/claude-3.5-sonnet" in _form_ids, "_sync_model_icons: ModelForm uses prefixed ID (anthropic)", ) - # Verify _icons_synced tracks the RAW model IDs + # After insert (model not yet registered by OWUI), _icons_synced must NOT be updated. + # The next cache-hit call will confirm the icon is set correctly. _assert( - "openai/gpt-4o" in _pipe_func._icons_synced, - "_sync_model_icons: _icons_synced uses raw model ID", + "openai/gpt-4o" not in _pipe_func._icons_synced, + "_sync_model_icons: _icons_synced NOT updated after insert (allows retry)", ) finally: - Pipe.__module__ = _orig_module_f sys.modules.pop("open_webui.models.models", None) -# 25g. Existing model with icon → skips overwrite +# 25g. Existing model with user custom icon → skips overwrite _pipe_skip = Pipe() _pipe_skip.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SYNC_PROVIDER_ICONS=True) +_pipe_skip._function_id = "openrouter_pipe" # simulate OWUI module naming _mock_Models_s = MagicMock() _existing_model = MagicMock() _existing_model.meta.profile_image_url = "https://custom-icon.example.com/icon.png" _existing_model.name = "Custom GPT" +_existing_model.params = None _mock_Models_s.get_model_by_id.return_value = _existing_model _fake_owui_s = ModuleType("open_webui.models.models") _fake_owui_s.Models = _mock_Models_s _fake_owui_s.ModelForm = MagicMock() _fake_owui_s.ModelMeta = MagicMock() _fake_owui_s.ModelParams = MagicMock() -_orig_module_s = Pipe.__module__ try: sys.modules["open_webui.models.models"] = _fake_owui_s - Pipe.__module__ = "function_openrouter_pipe" _pipe_skip._sync_model_icons([{"id": "openai/gpt-4o", "name": "GPT-4o"}]) _assert( not _mock_Models_s.update_model_by_id.called, - "_sync_model_icons: skips update when model has custom icon", + "_sync_model_icons: skips update when model has user custom icon", ) _assert( not _mock_Models_s.insert_new_model.called, - "_sync_model_icons: skips insert when model has custom icon", + "_sync_model_icons: skips insert when model has user custom icon", + ) + _assert( + "openai/gpt-4o" in _pipe_skip._icons_synced, + "_sync_model_icons: adds to _icons_synced when skipping custom icon", ) finally: - Pipe.__module__ = _orig_module_s sys.modules.pop("open_webui.models.models", None) +# 25h. Existing model with OWUI default (data: URL) icon → updates with provider icon +_pipe_update = Pipe() +_pipe_update.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SYNC_PROVIDER_ICONS=True) +_pipe_update._function_id = "openrouter_pipe" # simulate OWUI module naming +_mock_Models_u = MagicMock() +_existing_default = MagicMock() +_existing_default.name = "GPT-4o" +_existing_default.meta.profile_image_url = "data:image/svg+xml;base64,ABC123==" +_existing_default.params = None +_mock_Models_u.get_model_by_id.return_value = _existing_default +_fake_owui_u = ModuleType("open_webui.models.models") +_fake_owui_u.Models = _mock_Models_u +_fake_owui_u.ModelForm = MagicMock() +_fake_owui_u.ModelMeta = MagicMock() +_fake_owui_u.ModelParams = MagicMock() +try: + sys.modules["open_webui.models.models"] = _fake_owui_u + _pipe_update._sync_model_icons([{"id": "openai/gpt-4o", "name": "GPT-4o"}]) + _assert( + _mock_Models_u.update_model_by_id.called, + "_sync_model_icons: updates model when icon is OWUI default (data: URL)", + ) + _assert( + not _mock_Models_u.insert_new_model.called, + "_sync_model_icons: does not insert when model already exists", + ) + _assert( + "openai/gpt-4o" in _pipe_update._icons_synced, + "_sync_model_icons: adds to _icons_synced after successful update", + ) +finally: + sys.modules.pop("open_webui.models.models", None) + +# 25i. _is_owui_managed_icon helper +_is_owui = mod._is_owui_managed_icon +_assert(_is_owui(""), "_is_owui_managed_icon: empty string → True (no icon)") +_assert(_is_owui("data:image/svg+xml;base64,ABC"), "_is_owui_managed_icon: data: URL → True") +_assert(_is_owui("https://openrouter.ai/images/models/openai.svg"), "_is_owui_managed_icon: openrouter.ai URL → True") +_assert(not _is_owui("https://custom-icon.example.com/icon.png"), "_is_owui_managed_icon: external URL → False") +_assert(not _is_owui("https://cdn.openai.com/logo.png"), "_is_owui_managed_icon: other https URL → False") + # ══════════════════════════════════════════════════════════════════════════════ # Summary # ══════════════════════════════════════════════════════════════════════════════