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/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,38 @@
+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.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):  # type: ignore[override]
+    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._client_v1 = fireworks_client_v1
+        return self

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

@@ -0,0 +1,62 @@
+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="openai",
+        api_flavor="chat-completions",
+        api_version="2025-03-01-preview",
+        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=0,  # 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=0,  # 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]