From d23944bca745571b2bc9b1f2fc164bff27a9fa69 Mon Sep 17 00:00:00 2001
From: Prakash Maheshwaran
 <73785492+Prakashmaheshwaran@users.noreply.github.com>
Date: Tue, 1 Jul 2025 14:02:22 -0400
Subject: [PATCH] fixed the openrouter stuff (#2630)

---
 skyvern/forge/sdk/api/llm/config_registry.py  |  16 +++
 .../unit_tests/test_openrouter_integration.py | 103 ++++++++++++++++++
 2 files changed, 119 insertions(+)
 create mode 100644 tests/unit_tests/test_openrouter_integration.py

diff --git a/skyvern/forge/sdk/api/llm/config_registry.py b/skyvern/forge/sdk/api/llm/config_registry.py
index 51b42ba8..252007b2 100644
--- a/skyvern/forge/sdk/api/llm/config_registry.py
+++ b/skyvern/forge/sdk/api/llm/config_registry.py
@@ -40,6 +40,22 @@ class LLMConfigRegistry:
             # If the key is not found in registered configs, treat it as a general model
             if not llm_key:
                 raise InvalidLLMConfigError(f"LLM_KEY not set for {llm_key}")
+
+            if llm_key.startswith("openrouter/"):
+                return LLMConfig(
+                    llm_key,
+                    ["OPENROUTER_API_KEY"],
+                    supports_vision=settings.LLM_CONFIG_SUPPORT_VISION,
+                    add_assistant_prefix=settings.LLM_CONFIG_ADD_ASSISTANT_PREFIX,
+                    max_completion_tokens=settings.LLM_CONFIG_MAX_TOKENS,
+                    litellm_params=LiteLLMParams(
+                        api_key=settings.OPENROUTER_API_KEY,
+                        api_base=settings.OPENROUTER_API_BASE,
+                        api_version=None,
+                        model_info={"model_name": llm_key},
+                    ),
+                )
+
             return LLMConfig(
                 llm_key,  # Use the LLM_KEY as the model name
                 ["LLM_API_KEY"],
diff --git a/tests/unit_tests/test_openrouter_integration.py b/tests/unit_tests/test_openrouter_integration.py
new file mode 100644
index 00000000..5b28731e
--- /dev/null
+++ b/tests/unit_tests/test_openrouter_integration.py
@@ -0,0 +1,103 @@
+import importlib
+import types
+from unittest.mock import AsyncMock
+
+import pytest
+
+from skyvern.config import Settings
+from skyvern.forge import app
+from skyvern.forge.sdk.api.llm import api_handler_factory, config_registry
+from skyvern.forge.sdk.settings_manager import SettingsManager
+
+
+class DummyResponse(dict):
+    def __init__(self, content: str):
+        super().__init__({"choices": [{"message": {"content": content}}], "usage": {}})
+        self.choices = [types.SimpleNamespace(message=types.SimpleNamespace(content=content))]
+
+    def model_dump_json(self, indent: int = 2):
+        import json
+
+        return json.dumps(self, indent=indent)
+
+
+class DummyArtifactManager:
+    async def create_llm_artifact(self, *args, **kwargs):
+        return None
+
+
+@pytest.mark.asyncio
+async def test_openrouter_basic_completion(monkeypatch):
+    settings = Settings(
+        ENABLE_OPENROUTER=True,
+        OPENROUTER_API_KEY="key",
+        OPENROUTER_MODEL="test-model",
+        LLM_KEY="OPENROUTER",
+    )
+    SettingsManager.set_settings(settings)
+    importlib.reload(config_registry)
+
+    monkeypatch.setattr(app, "ARTIFACT_MANAGER", DummyArtifactManager())
+
+    async_mock = AsyncMock(return_value=DummyResponse('{"result": "ok"}'))
+    monkeypatch.setattr(api_handler_factory.litellm, "acompletion", async_mock)
+
+    handler = api_handler_factory.LLMAPIHandlerFactory.get_llm_api_handler("OPENROUTER")
+    result = await handler("hi", "test")
+    assert result == {"result": "ok"}
+    async_mock.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_openrouter_dynamic_model(monkeypatch):
+    settings = Settings(
+        ENABLE_OPENROUTER=True,
+        OPENROUTER_API_KEY="key",
+        OPENROUTER_MODEL="base-model",
+        LLM_KEY="OPENROUTER",
+    )
+    SettingsManager.set_settings(settings)
+    importlib.reload(config_registry)
+
+    monkeypatch.setattr(app, "ARTIFACT_MANAGER", DummyArtifactManager())
+    async_mock = AsyncMock(return_value=DummyResponse('{"status": "ok"}'))
+    monkeypatch.setattr(api_handler_factory.litellm, "acompletion", async_mock)
+
+    base_handler = api_handler_factory.LLMAPIHandlerFactory.get_llm_api_handler("OPENROUTER")
+    override_handler = api_handler_factory.LLMAPIHandlerFactory.get_override_llm_api_handler(
+        "openrouter/other-model", default=base_handler
+    )
+    result = await override_handler("hi", "test")
+    assert result == {"status": "ok"}
+    called_model = async_mock.call_args.kwargs.get("model")
+    assert called_model == "openrouter/other-model"
+
+
+@pytest.mark.asyncio
+async def test_openrouter_error_propagation(monkeypatch):
+    class DummyAPIError(Exception):
+        pass
+
+    settings = Settings(
+        ENABLE_OPENROUTER=True,
+        OPENROUTER_API_KEY="key",
+        OPENROUTER_MODEL="test-model",
+        LLM_KEY="OPENROUTER",
+    )
+    SettingsManager.set_settings(settings)
+    importlib.reload(config_registry)
+
+    monkeypatch.setattr(app, "ARTIFACT_MANAGER", DummyArtifactManager())
+
+    async def _raise(*args, **kwargs):
+        raise DummyAPIError()
+
+    fake_litellm = types.SimpleNamespace(
+        acompletion=_raise,
+        exceptions=types.SimpleNamespace(APIError=DummyAPIError),
+    )
+    monkeypatch.setattr(api_handler_factory, "litellm", fake_litellm)
+
+    handler = api_handler_factory.LLMAPIHandlerFactory.get_llm_api_handler("OPENROUTER")
+    with pytest.raises(api_handler_factory.LLMProviderErrorRetryableTask):
+        await handler("hi", "test")