Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
104 changes: 74 additions & 30 deletions openrouter_pipe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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")

Expand Down Expand Up @@ -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):
Comment thread
sena-labs marked this conversation as resolved.
self._sync_model_icons(self._models_cache)
return self._models_cache

headers = self._build_headers(include_content_type=False)
Expand Down Expand Up @@ -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:
Expand All @@ -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", "")
Expand All @@ -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(
Expand All @@ -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
Comment thread
sena-labs marked this conversation as resolved.
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]:
Expand Down
88 changes: 65 additions & 23 deletions test_pipe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand All @@ -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()
Expand All @@ -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"},
Expand All @@ -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(
Expand All @@ -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
# ══════════════════════════════════════════════════════════════════════════════
Expand Down
Loading