swarmauri_llm_openrouter 0.1.0.dev3__py3-none-any.whl

swarmauri_llm_openrouter/OpenRouterImgGenModel.py

@@ -0,0 +1,169 @@
+"""Provide OpenRouter image-generation support."""
+
+import asyncio
+import base64
+from typing import Any, Literal
+
+from pydantic import Field, PrivateAttr, SecretStr
+from swarmauri_base.ComponentBase import ComponentBase
+from swarmauri_base.image_gens.ImageGenBase import ImageGenBase
+from swarmauri_core.transports.i_http_client import IHttpClient
+from swarmauri_core.transports.i_http_retry_policy import IHttpRetryPolicy
+from swarmauri_provider_openrouter import (
+    OpenRouterModelCatalog,
+    OpenRouterProvider,
+    OpenRouterProviderPreferences,
+)
+
+from ._client import (
+    aclose_client,
+    arequest,
+    build_openrouter_client,
+    close_client,
+    request,
+)
+from ._internal.common import compact
+
+
+@ComponentBase.register_type(ImageGenBase, "OpenRouterImgGenModel")
+class OpenRouterImgGenModel(ImageGenBase):
+    """Generate images through OpenRouter's dedicated image API."""
+
+    api_key: SecretStr | None = Field(default=None, exclude=True)
+    name: str = "openai/gpt-image-1"
+    type: Literal["OpenRouterImgGenModel"] = "OpenRouterImgGenModel"
+    base_url: str = "https://openrouter.ai/api/v1"
+    site_url: str | None = None
+    app_name: str | None = None
+    timeout: float = 600.0
+    max_retries: int = Field(default=3, ge=1)
+    routing: OpenRouterProviderPreferences | None = None
+    allowed_models: list[str] = Field(default_factory=list)
+    _provider: OpenRouterProvider = PrivateAttr()
+    _http_client: IHttpClient = PrivateAttr()
+
+    def __init__(
+        self,
+        *,
+        provider: OpenRouterProvider
+        | OpenRouterProviderPreferences
+        | None = None,
+        routing: OpenRouterProviderPreferences | None = None,
+        http_client: IHttpClient | None = None,
+        retry_policy: IHttpRetryPolicy | None = None,
+        **data: Any,
+    ) -> None:
+        if isinstance(provider, OpenRouterProviderPreferences):
+            routing = provider
+            provider = None
+        if routing is not None:
+            data["routing"] = routing
+        super().__init__(**data)
+        self._provider = provider or OpenRouterProvider(
+            api_key=self.api_key,
+            base_url=self.base_url,
+            site_url=self.site_url,
+            app_name=self.app_name,
+            timeout=self.timeout,
+            max_retries=self.max_retries,
+        )
+        if self._provider.api_key is None:
+            raise ValueError(
+                "api_key or an authenticated provider is required"
+            )
+        self.base_url = self._provider.base_url or self.base_url
+        self._http_client = build_openrouter_client(
+            provider=self._provider,
+            http_client=http_client,
+            retry_policy=retry_policy,
+        )
+
+    @property
+    def catalog(self) -> OpenRouterModelCatalog:
+        return self._provider.catalog
+
+    def _decode_images(self, response: dict[str, Any]) -> list[bytes]:
+        images: list[bytes] = []
+        for item in response.get("data", []):
+            encoded = item.get("b64_json") or item.get("b64")
+            if encoded:
+                images.append(base64.b64decode(encoded))
+            elif item.get("url"):
+                downloaded = self._http_client.request("GET", item["url"])
+                downloaded.raise_for_status()
+                images.append(downloaded.read())
+        return images
+
+    async def _adecode_images(self, response: dict[str, Any]) -> list[bytes]:
+        images: list[bytes] = []
+        for item in response.get("data", []):
+            encoded = item.get("b64_json") or item.get("b64")
+            if encoded:
+                images.append(base64.b64decode(encoded))
+            elif item.get("url"):
+                downloaded = await self._http_client.arequest(
+                    "GET", item["url"]
+                )
+                downloaded.raise_for_status()
+                images.append(await downloaded.aread())
+        return images
+
+    def _payload(self, prompt: str, **kwargs: Any) -> dict[str, Any]:
+        return compact(
+            {
+                "model": self.name,
+                "prompt": prompt,
+                "provider": self.routing.model_dump(exclude_none=True)
+                if self.routing
+                else None,
+                **kwargs,
+            }
+        )
+
+    def generate_image(
+        self, prompt: str, **kwargs: Any
+    ) -> bytes | list[bytes]:
+        response = request(
+            self._provider,
+            self._http_client,
+            "images.generate",
+            model=self.name,
+            json_data=self._payload(prompt, **kwargs),
+        )
+        images = self._decode_images(response.json())
+        return images[0] if len(images) == 1 else images
+
+    async def agenerate_image(
+        self, prompt: str, **kwargs: Any
+    ) -> bytes | list[bytes]:
+        response = await arequest(
+            self._provider,
+            self._http_client,
+            "images.generate",
+            model=self.name,
+            json_data=self._payload(prompt, **kwargs),
+        )
+        images = await self._adecode_images(response.json())
+        return images[0] if len(images) == 1 else images
+
+    def batch_generate(
+        self, prompts: list[str], **kwargs: Any
+    ) -> list[bytes | list[bytes]]:
+        return [self.generate_image(prompt, **kwargs) for prompt in prompts]
+
+    async def abatch_generate(
+        self, prompts: list[str], max_concurrent: int = 5, **kwargs: Any
+    ) -> list[bytes | list[bytes]]:
+        semaphore = asyncio.Semaphore(max_concurrent)
+
+        async def run(prompt: str) -> bytes | list[bytes]:
+            async with semaphore:
+                return await self.agenerate_image(prompt, **kwargs)
+
+        return await asyncio.gather(*(run(prompt) for prompt in prompts))
+
+    def close(self) -> None:
+        close_client(self._http_client)
+
+    async def aclose(self) -> None:
+        await aclose_client(self._http_client)

swarmauri_llm_openrouter/OpenRouterModel.py

@@ -0,0 +1,192 @@
+"""Provide OpenRouter text and reasoning model support."""
+
+import asyncio
+from collections.abc import AsyncIterator, Iterator
+from typing import Any, ClassVar, Literal
+
+from pydantic import Field, PrivateAttr, SecretStr
+from swarmauri_base.ComponentBase import ComponentBase
+from swarmauri_base.llms.LLMBase import LLMBase
+from swarmauri_core.transports.i_http_client import IHttpClient
+from swarmauri_core.transports.i_http_retry_policy import IHttpRetryPolicy
+from swarmauri_provider_openrouter import (
+    OpenRouterModelCatalog,
+    OpenRouterProvider,
+    OpenRouterProviderPreferences,
+)
+from swarmauri_standard.messages.AgentMessage import AgentMessage
+
+from ._client import (
+    aclose_client,
+    arequest,
+    astream_events,
+    build_openrouter_client,
+    close_client,
+    request,
+    stream_events,
+)
+from ._internal.common import (
+    add_assistant_response,
+    compact,
+    delta_text,
+    format_messages,
+)
+
+
+@ComponentBase.register_type(LLMBase, "OpenRouterModel")
+class OpenRouterModel(LLMBase):
+    """Route text and reasoning conversations through OpenRouter."""
+
+    api_key: SecretStr | None = Field(default=None, exclude=True)
+    name: str = "openrouter/auto"
+    type: Literal["OpenRouterModel"] = "OpenRouterModel"
+    base_url: str = "https://openrouter.ai/api/v1"
+    site_url: str | None = None
+    app_name: str | None = None
+    routing: OpenRouterProviderPreferences | None = None
+    allowed_models: list[str] = Field(default_factory=list)
+    capabilities: ClassVar[frozenset[str]] = frozenset(
+        {"predict", "apredict", "stream", "astream", "batch", "abatch"}
+    )
+    _provider: OpenRouterProvider = PrivateAttr()
+    _http_client: IHttpClient = PrivateAttr()
+
+    def __init__(
+        self,
+        *,
+        provider: OpenRouterProvider
+        | OpenRouterProviderPreferences
+        | None = None,
+        routing: OpenRouterProviderPreferences | None = None,
+        http_client: IHttpClient | None = None,
+        retry_policy: IHttpRetryPolicy | None = None,
+        **data: Any,
+    ) -> None:
+        if isinstance(provider, OpenRouterProviderPreferences):
+            routing = provider
+            provider = None
+        if routing is not None:
+            data["routing"] = routing
+        super().__init__(**data)
+        self._provider = provider or OpenRouterProvider(
+            api_key=self.api_key,
+            base_url=self.base_url,
+            site_url=self.site_url,
+            app_name=self.app_name,
+            timeout=self.timeout,
+            max_retries=self.max_retries,
+        )
+        if self._provider.api_key is None:
+            raise ValueError(
+                "api_key or an authenticated provider is required"
+            )
+        self.base_url = self._provider.base_url or self.base_url
+        self._http_client = build_openrouter_client(
+            provider=self._provider,
+            http_client=http_client,
+            retry_policy=retry_policy,
+        )
+
+    @property
+    def catalog(self) -> OpenRouterModelCatalog:
+        return self._provider.catalog
+
+    def refresh_catalog(self) -> list[dict[str, Any]]:
+        response = request(self._provider, self._http_client, "models.list")
+        return self.catalog.load_response(response.json())
+
+    async def arefresh_catalog(self) -> list[dict[str, Any]]:
+        response = await arequest(
+            self._provider, self._http_client, "models.list"
+        )
+        return self.catalog.load_response(response.json())
+
+    def _request_headers(self) -> dict[str, str]:
+        config = self._provider.connection_config(
+            "chat.completions", model=self.name
+        )
+        return self._provider.authorize(config)
+
+    def _payload(self, conversation: Any, **kwargs: Any) -> dict[str, Any]:
+        return compact(
+            {
+                "model": self.name,
+                "messages": format_messages(conversation.history),
+                "provider": self.routing.model_dump(exclude_none=True)
+                if self.routing
+                else None,
+                **kwargs,
+            }
+        )
+
+    def predict(self, conversation: Any, **kwargs: Any) -> Any:
+        response = request(
+            self._provider,
+            self._http_client,
+            "chat.completions",
+            model=self.name,
+            json_data=self._payload(conversation, **kwargs),
+        )
+        return add_assistant_response(conversation, response.json())
+
+    async def apredict(self, conversation: Any, **kwargs: Any) -> Any:
+        response = await arequest(
+            self._provider,
+            self._http_client,
+            "chat.completions",
+            model=self.name,
+            json_data=self._payload(conversation, **kwargs),
+        )
+        return add_assistant_response(conversation, response.json())
+
+    def stream(self, conversation: Any, **kwargs: Any) -> Iterator[str]:
+        content: list[str] = []
+        payload = self._payload(conversation, stream=True, **kwargs)
+        for event in stream_events(
+            self._provider,
+            self._http_client,
+            "chat.completions",
+            payload,
+            model=self.name,
+        ):
+            if delta := delta_text(event):
+                content.append(delta)
+                yield delta
+        conversation.add_message(AgentMessage(content="".join(content)))
+
+    async def astream(
+        self, conversation: Any, **kwargs: Any
+    ) -> AsyncIterator[str]:
+        content: list[str] = []
+        payload = self._payload(conversation, stream=True, **kwargs)
+        async for event in astream_events(
+            self._provider,
+            self._http_client,
+            "chat.completions",
+            payload,
+            model=self.name,
+        ):
+            if delta := delta_text(event):
+                content.append(delta)
+                yield delta
+        conversation.add_message(AgentMessage(content="".join(content)))
+
+    def batch(self, conversations: list[Any], **kwargs: Any) -> list[Any]:
+        return [self.predict(item, **kwargs) for item in conversations]
+
+    async def abatch(
+        self, conversations: list[Any], max_concurrent: int = 5, **kwargs: Any
+    ) -> list[Any]:
+        semaphore = asyncio.Semaphore(max_concurrent)
+
+        async def run(item: Any) -> Any:
+            async with semaphore:
+                return await self.apredict(item, **kwargs)
+
+        return await asyncio.gather(*(run(item) for item in conversations))
+
+    def close(self) -> None:
+        close_client(self._http_client)
+
+    async def aclose(self) -> None:
+        await aclose_client(self._http_client)

swarmauri_llm_openrouter/OpenRouterRetryPolicy.py

@@ -0,0 +1,34 @@
+"""Retry policy for OpenRouter model operations."""
+
+from swarmauri_core.transports.i_http_response import IHttpResponse
+from swarmauri_transport_retry import ExponentialBackoffRetryPolicy
+
+
+class OpenRouterRetryPolicy(ExponentialBackoffRetryPolicy):
+    """Retry safe reads and rate-limited model submissions."""
+
+    def __init__(self, **kwargs) -> None:
+        kwargs.setdefault("retryable_methods", {"GET", "POST"})
+        super().__init__(**kwargs)
+
+    def should_retry(
+        self,
+        *,
+        method: str,
+        attempt: int,
+        response: IHttpResponse | None = None,
+        error: Exception | None = None,
+    ) -> bool:
+        if method.upper() == "POST":
+            return (
+                attempt < self.max_attempts
+                and error is None
+                and response is not None
+                and response.status_code == 429
+            )
+        return super().should_retry(
+            method=method,
+            attempt=attempt,
+            response=response,
+            error=error,
+        )