started integrating fireworks

Author: cosminacho

packages/uipath_langchain_client/pyproject.toml

@@ -29,8 +29,11 @@ azure = [
 vertexai = [
     "langchain-google-vertexai>=3.2.1",
 ]
+fireworks = [
+    "langchain-fireworks>=1.1.0",
+]
 all = [
-    "uipath-langchain-client[openai,aws,google,anthropic,azure,vertexai]"
+    "uipath-langchain-client[openai,aws,google,anthropic,azure,vertexai,fireworks]"
 ]
 
 [build-system]

packages/uipath_langchain_client/src/uipath_langchain_client/base_client.py

@@ -100,8 +100,8 @@ class UiPathBaseLLMClient(BaseModel):
         default=None,
         description="Client-side request timeout in seconds",
     )
-    max_retries: int = Field(
-        default=1,
+    max_retries: int | None = Field(
+        default=None,
         description="Maximum number of retries for failed requests",
     )
     retry_config: RetryConfig | None = Field(

packages/uipath_langchain_client/src/uipath_langchain_client/clients/fireworks/__init__.py

@@ -0,0 +1,4 @@
+from uipath_langchain_client.clients.fireworks.chat_models import UiPathChatFireworks
+from uipath_langchain_client.clients.fireworks.embeddings import UiPathFireworksEmbeddings
+
+__all__ = ["UiPathChatFireworks", "UiPathFireworksEmbeddings"]

packages/uipath_langchain_client/src/uipath_langchain_client/clients/fireworks/chat_models.py

@@ -0,0 +1,40 @@
+from typing import Self
+
+from pydantic import Field, SecretStr, model_validator
+from uipath_langchain_client.base_client import UiPathBaseLLMClient
+from uipath_langchain_client.settings import UiPathAPIConfig
+
+try:
+    from langchain_fireworks.chat_models import ChatFireworks
+
+    from fireworks.client import AsyncFireworks, Fireworks
+    from fireworks.client.api_client import FireworksClient as FireworksClientV1
+except ImportError as e:
+    raise ImportError(
+        "The 'fireworks' extra is required to use UiPathChatFireworks. "
+        "Install it with: uv add uipath-langchain-client[fireworks]"
+    ) from e
+
+
+class UiPathChatFireworks(UiPathBaseLLMClient, ChatFireworks):
+    api_config: UiPathAPIConfig = UiPathAPIConfig(
+        api_type="completions",
+        client_type="passthrough",
+        freeze_base_url=True,
+    )
+
+    # Override fields to avoid errors when instantiating the class
+    fireworks_api_base: str | None = Field(alias="base_url", default="PLACEHOLDER")
+    fireworks_api_key: SecretStr = Field(default=SecretStr("PLACEHOLDER"), alias="api_key")
+
+    @model_validator(mode="after")
+    def setup_uipath_client(self) -> Self:
+        fireworks_client_v1 = FireworksClientV1(
+            api_key=self.fireworks_api_key.get_secret_value(),
+            base_url=self.fireworks_api_base,
+        )
+        fireworks_client_v1._client = self.uipath_sync_client
+        fireworks_client_v1._async_client = self.uipath_async_client
+        self.client = Fireworks(client=fireworks_client_v1)
+
+        return self

packages/uipath_langchain_client/src/uipath_langchain_client/clients/fireworks/embeddings.py

@@ -0,0 +1,60 @@
+from typing import Self
+
+from pydantic import model_validator
+from uipath_langchain_client.base_client import UiPathBaseLLMClient
+from uipath_langchain_client.settings import UiPathAPIConfig
+
+try:
+    from langchain_fireworks.embeddings import FireworksEmbeddings
+    from openai import AsyncOpenAI, OpenAI
+except ImportError as e:
+    raise ImportError(
+        "The 'fireworks' extra is required to use UiPathFireworksEmbeddings. "
+        "Install it with: uv add uipath-langchain-client[fireworks]"
+    ) from e
+
+
+class UiPathFireworksEmbeddings(UiPathBaseLLMClient, FireworksEmbeddings):
+    api_config: UiPathAPIConfig = UiPathAPIConfig(
+        api_type="embeddings",
+        client_type="passthrough",
+        vendor_type="fireworks",
+        freeze_base_url=True,
+    )
+
+    @model_validator(mode="after")
+    def setup_uipath_client(self) -> Self:
+        self.client = OpenAI(
+            api_key="PLACEHOLDER",
+            timeout=None,  # handled by the UiPath client
+            max_retries=1,  # handled by the UiPath client
+            http_client=self.uipath_sync_client,
+        )
+        self.async_client = AsyncOpenAI(
+            api_key="PLACEHOLDER",
+            timeout=None,  # handled by the UiPath client
+            max_retries=1,  # handled by the UiPath client
+            http_client=self.uipath_async_client,
+        )
+        return self
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed search docs."""
+        return [
+            i.embedding for i in self.client.embeddings.create(input=texts, model=self.model).data
+        ]
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embed query text."""
+        return self.embed_documents([text])[0]
+
+    async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed search docs asynchronously."""
+        return [
+            i.embedding
+            for i in (await self.async_client.embeddings.create(input=texts, model=self.model)).data
+        ]
+
+    async def aembed_query(self, text: str) -> list[float]:
+        """Embed query text asynchronously."""
+        return (await self.aembed_documents([text]))[0]

packages/uipath_langchain_client/src/uipath_langchain_client/clients/fireworks/utils.py

@@ -0,0 +1,15 @@
+from fireworks.client import Fireworks, AsyncFireworks
+from fireworks.client.api_client import FireworksClient as FireworksClientV1
+
+from httpx import Client, AsyncClient
+
+class UiPathFireworksClient(FireworksClientV1):
+    def __init__(self, *args, http_client: Client | None = None, async_http_client: AsyncClient | None = None, **kwargs):
+        super().__init__(*args, **kwargs)
+        if http_client is not None:
+            self._client = http_client
+        if async_http_client is not None:
+            self._async_client = async_http_client
+
+class UiPathFireworks(Fireworks):
+