From 07c78d90099ab04584b92122f718b60f83023aae Mon Sep 17 00:00:00 2001 From: Mitul Shah Date: Wed, 27 May 2026 15:44:29 +0530 Subject: [PATCH 1/5] feat: ship langchain adapter --- ferrolabsai/client.py | 60 +++- ferrolabsai/types.py | 2 +- .../langchain-ferrolabsai/CHANGELOG.md | 52 ++- integrations/langchain-ferrolabsai/README.md | 80 ++++- .../langchain_ferrolabsai/__init__.py | 39 +-- .../langchain_ferrolabsai/_messages.py | 79 +++++ .../langchain_ferrolabsai/chat_models.py | 325 ++++++++++++++++++ .../langchain_ferrolabsai/embeddings.py | 77 +++++ .../langchain_ferrolabsai/llms.py | 73 ++++ .../langchain-ferrolabsai/pyproject.toml | 4 +- .../langchain-ferrolabsai/tests/__init__.py | 0 .../langchain-ferrolabsai/tests/conftest.py | 78 +++++ .../tests/test_chat_models.py | 306 +++++++++++++++++ .../tests/test_embeddings.py | 93 +++++ .../langchain-ferrolabsai/tests/test_llms.py | 56 +++ .../tests/test_placeholder.py | 31 -- tests/test_sdk.py | 52 +++ 17 files changed, 1332 insertions(+), 75 deletions(-) create mode 100644 integrations/langchain-ferrolabsai/langchain_ferrolabsai/_messages.py create mode 100644 integrations/langchain-ferrolabsai/langchain_ferrolabsai/chat_models.py create mode 100644 integrations/langchain-ferrolabsai/langchain_ferrolabsai/embeddings.py create mode 100644 integrations/langchain-ferrolabsai/langchain_ferrolabsai/llms.py create mode 100644 integrations/langchain-ferrolabsai/tests/__init__.py create mode 100644 integrations/langchain-ferrolabsai/tests/conftest.py create mode 100644 integrations/langchain-ferrolabsai/tests/test_chat_models.py create mode 100644 integrations/langchain-ferrolabsai/tests/test_embeddings.py create mode 100644 integrations/langchain-ferrolabsai/tests/test_llms.py delete mode 100644 integrations/langchain-ferrolabsai/tests/test_placeholder.py diff --git a/ferrolabsai/client.py b/ferrolabsai/client.py index 5155275..017bd06 100644 --- a/ferrolabsai/client.py +++ b/ferrolabsai/client.py @@ -167,7 +167,7 @@ def _request( return response if response.status_code == 204 or not response.content: return {} - return cast("dict[str, Any]", response.json()) + return _with_response_metadata(cast("dict[str, Any]", response.json()), response) except httpx.HTTPStatusError as e: _raise_api_error(e) except httpx.ConnectError as e: @@ -306,7 +306,7 @@ async def _request( response.raise_for_status() if response.status_code == 204 or not response.content: return {} - return cast("dict[str, Any]", response.json()) + return _with_response_metadata(cast("dict[str, Any]", response.json()), response) except httpx.HTTPStatusError as e: _raise_api_error(e) except httpx.ConnectError as e: @@ -345,6 +345,62 @@ def __init__(self, client: AsyncFerroClient) -> None: # ------------------------------------------------------------------ +def _with_response_metadata(data: dict[str, Any], response: httpx.Response) -> dict[str, Any]: + """Copy gateway metadata headers into parsed response bodies. + + Successful SDK calls return dataclasses, so header-only metadata such as + X-Request-ID must be preserved before resource classes construct them. + Body fields stay authoritative when both sources are present. + """ + trace_id = ( + response.headers.get("x-request-id") + or response.headers.get("x-trace-id") + or response.headers.get("x-ferro-request-id") + ) + if trace_id and "trace_id" not in data and "x_ferro_trace_id" not in data: + data["trace_id"] = trace_id + + provider = response.headers.get("x-ferro-provider") + if provider and "provider" not in data and "x_ferro_provider" not in data: + data["provider"] = provider + usage = data.get("usage") + if isinstance(usage, dict) and "provider" not in usage: + usage["provider"] = provider + + latency_ms = _header_int(response.headers.get("x-ferro-latency-ms")) + if latency_ms is not None and "latency_ms" not in data and "x_ferro_latency_ms" not in data: + data["x_ferro_latency_ms"] = latency_ms + + cost_usd = _header_float(response.headers.get("x-ferro-cost-usd")) + if cost_usd is not None: + usage = data.get("usage") + if not isinstance(usage, dict): + usage = {} + data["usage"] = usage + if "cost_usd" not in usage: + usage["cost_usd"] = cost_usd + + return data + + +def _header_int(value: str | None) -> int | None: + if value is None or value == "": + return None + try: + return int(float(value)) + except ValueError: + return None + + +def _header_float(value: str | None) -> float | None: + if value is None or value == "": + return None + try: + return float(value) + except ValueError: + return None + + def _raise_api_error(e: httpx.HTTPStatusError) -> None: from .exceptions import ( FerroAPIError, diff --git a/ferrolabsai/types.py b/ferrolabsai/types.py index 65a5120..bcc7412 100644 --- a/ferrolabsai/types.py +++ b/ferrolabsai/types.py @@ -95,7 +95,7 @@ def from_dict(cls, d: dict[str, Any]) -> ChatCompletion: usage=Usage.from_dict(d["usage"]) if d.get("usage") else None, trace_id=d.get("x_ferro_trace_id") or d.get("trace_id"), provider=d.get("x_ferro_provider") or d.get("provider"), - latency_ms=d.get("x_ferro_latency_ms"), + latency_ms=d.get("x_ferro_latency_ms") or d.get("latency_ms"), ) @property diff --git a/integrations/langchain-ferrolabsai/CHANGELOG.md b/integrations/langchain-ferrolabsai/CHANGELOG.md index 4283883..a536041 100644 --- a/integrations/langchain-ferrolabsai/CHANGELOG.md +++ b/integrations/langchain-ferrolabsai/CHANGELOG.md @@ -9,16 +9,50 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Planned -- `FerroChatModel` — `langchain_core.language_models.chat_models.BaseChatModel` adapter wrapping `ferrolabsai.FerroClient.chat.completions`. -- `FerroEmbeddings` — `langchain_core.embeddings.Embeddings` adapter wrapping `ferrolabsai.FerroClient.embeddings`. -- `FerroLLM` — completion-style adapter for legacy LangChain chains. -- Streaming + async + tool calling. -- `provider`, `cost_usd`, `latency_ms`, `trace_id` exposed via `response_metadata`. -- Native support for Ferro extras (`route_tag`, `template_id`, `template_variables`). -- pytest-httpx based test suite mirroring the parent SDK's mocking pattern. +- Async surfaces (`_agenerate`, `_astream`, `aembed_documents`, `aembed_query`). +- `with_structured_output()` helper for JSON-mode + Pydantic-schema responses. +- Native multi-modal message support (image inputs) once the gateway exposes a + stable contract. --- -## [0.0.1] — TBD +## [0.1.0] — 2026-05-25 + +First functional release. Replaces the `0.0.1` placeholder. + +### Added + +- **`FerroChatModel`** — `langchain_core.language_models.chat_models.BaseChatModel` + adapter wrapping `ferrolabsai.FerroClient.chat.completions`. Supports sync + generation, streaming via `_stream`, and tool binding via `bind_tools` + (compatible with LangGraph agents). +- **`FerroEmbeddings`** — `langchain_core.embeddings.Embeddings` adapter wrapping + `ferrolabsai.FerroClient.embeddings`. `embed_documents` preserves input order + even when the gateway returns embeddings out of order. +- **`FerroLLM`** — completion-style adapter for legacy LangChain chains. Wraps + chat completions with a single user message. +- **`trace_id` surfacing** — every chat response carries `trace_id`, `provider`, + `latency_ms`, `cost_usd`, and `cache_hit` (when present) in + `response_metadata`. `trace_id` is the join key for the v1.2 observability + bridge plugins (LangSmith, Langfuse, Phoenix, …). +- **Native support for Ferro extras** — `route_tag`, `template_id`, + `template_variables`, and `user` are first-class fields on `FerroChatModel` + and forwarded on every request. +- **`pytest-httpx`-based test suite** mirroring the parent SDK's mocking + pattern. No real gateway or network access required to run tests. + +### Notes + +- Async support is intentionally deferred to a follow-up release to keep the + initial diff reviewable. LangChain's default sync-fallback async behaviour + works in the meantime. +- Streaming surfaces incremental content chunks and OpenAI-style streamed + `delta.tool_calls` as LangChain `tool_call_chunks` for tool-using agents. -Placeholder release to reserve the `langchain-ferrolabsai` name on PyPI. No working implementation; importing the package raises `NotImplementedError` with a link to the roadmap. +--- + +## [0.0.1] — 2026-05-13 + +Placeholder release to reserve the `langchain-ferrolabsai` name on PyPI. No +working implementation; importing the package raised `NotImplementedError` +with a link to the roadmap. diff --git a/integrations/langchain-ferrolabsai/README.md b/integrations/langchain-ferrolabsai/README.md index b88ecc2..b1ce926 100644 --- a/integrations/langchain-ferrolabsai/README.md +++ b/integrations/langchain-ferrolabsai/README.md @@ -5,8 +5,6 @@ LangChain integration for [Ferro Labs AI Gateway](https://github.com/ferro-labs/ai-gateway) — route LangChain chat, streaming, tool-calling, and embedding workloads across **30+ LLM providers** through a single OpenAI-compatible endpoint, with automatic fallback, load balancing, cost tracking, and observability. -> **Status: 0.0.1 placeholder.** The full adapter (`FerroChatModel`, `FerroEmbeddings`, `FerroLLM`) is in active development as part of the [OSS Ecosystem Roadmap, Appendix A, Phase C](https://github.com/ferro-labs/ai-gateway-workspace/blob/main/docs/OSS-ECOSYSTEM-ROADMAP.md). The 0.0.1 release exists to reserve the package name on PyPI and signal upcoming work. - --- ## Install @@ -15,10 +13,13 @@ LangChain integration for [Ferro Labs AI Gateway](https://github.com/ferro-labs/ pip install langchain-ferrolabsai ``` -## Planned API +## Quick start + +### Chat ```python -from langchain_ferrolabsai import FerroChatModel, FerroEmbeddings +from langchain_ferrolabsai import FerroChatModel +from langchain_core.messages import HumanMessage llm = FerroChatModel( model="gpt-4o", @@ -26,27 +27,86 @@ llm = FerroChatModel( api_key="sk-ferro-...", ) -response = llm.invoke("Hello, world") +response = llm.invoke([HumanMessage(content="Hello, world")]) print(response.content) print(response.response_metadata["provider"]) # which provider handled it print(response.response_metadata["cost_usd"]) # cost for this request print(response.response_metadata["latency_ms"]) # observed latency -print(response.response_metadata["trace_id"]) # gateway trace ID +print(response.response_metadata["trace_id"]) # gateway trace ID (x-trace-id) +``` + +Swap providers without changing the model class — Ferro auto-routes by model +name: + +```python +claude = FerroChatModel(model="claude-3-5-sonnet-20241022", base_url="...", api_key="...") +gemini = FerroChatModel(model="gemini-1.5-flash", base_url="...", api_key="...") +``` + +### Streaming + +```python +for chunk in llm.stream([HumanMessage(content="Tell me a story")]): + print(chunk.content, end="", flush=True) +``` + +### Tool calling / LangGraph agents + +```python +from langchain_core.tools import tool + +@tool +def add(a: int, b: int) -> int: + """Add two integers.""" + return a + b + +agent_llm = llm.bind_tools([add]) +response = agent_llm.invoke([HumanMessage(content="What is 4 + 7?")]) +print(response.tool_calls) +``` + +### Embeddings + +```python +from langchain_ferrolabsai import FerroEmbeddings + +embed = FerroEmbeddings(model="text-embedding-3-small", base_url="...", api_key="...") +vectors = embed.embed_documents(["hello", "world"]) +query_vec = embed.embed_query("hello") +``` + +### Legacy `LLM` interface + +```python +from langchain_ferrolabsai import FerroLLM + +llm = FerroLLM(model="gpt-4o", base_url="...", api_key="...") +print(llm.invoke("Write a haiku about gateways")) ``` ## Why use this instead of `ChatOpenAI(base_url=...)`? -`ChatOpenAI` pointed at a Ferro Labs gateway already works as a drop-in. This package adds: +`ChatOpenAI` pointed at a Ferro Labs gateway works as a drop-in. This package adds: - First-class `provider`, `cost_usd`, `latency_ms`, `trace_id` exposure on `response_metadata`. - Native support for Ferro extras: `route_tag`, `template_id`, `template_variables`. -- Streaming + async + tool calling that surfaces all 30+ providers transparently. -- Optional LangSmith bridge so non-OpenAI providers (Anthropic, Bedrock, Vertex, etc.) appear in your existing LangSmith dashboards. +- `trace_id` is the **join key** for the v1.2 observability bridge plugins + (LangSmith, Langfuse, Phoenix, Datadog, …) shipping from the + [`ferro-labs/ai-gateway-plugins`](https://github.com/ferro-labs) repo — + any provider's calls become visible in your existing LLMOps backend without + per-provider wiring. + +## Status & roadmap + +`0.1.0` is the **first functional release** of the adapter. See +[`CHANGELOG.md`](CHANGELOG.md) for what shipped and what's planned. Async +surfaces and `with_structured_output()` are the next two items. ## Related - [`ferrolabsai`](https://pypi.org/project/ferrolabsai/) — the core Python SDK this package wraps. -- [Ferro Labs AI Gateway](https://github.com/ferro-labs/ai-gateway) — the open-source gateway server. +- [Ferro Labs AI Gateway](https://github.com/ferro-labs/ai-gateway) — the open-source gateway server (v1.1.0+ OTel-native). +- [`ai-gateway-cookbook`](https://github.com/ferro-labs/ai-gateway-cookbook) — runnable recipes (start with `python/02-langgraph-multi-provider-agent`). - [Documentation](https://docs.ferrolabs.ai) ## License diff --git a/integrations/langchain-ferrolabsai/langchain_ferrolabsai/__init__.py b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/__init__.py index 54bbb1c..fa5bc47 100644 --- a/integrations/langchain-ferrolabsai/langchain_ferrolabsai/__init__.py +++ b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/__init__.py @@ -1,31 +1,30 @@ """LangChain integration for Ferro Labs AI Gateway. -This is a 0.0.1 placeholder release. The full adapter (FerroChatModel, -FerroEmbeddings, FerroLLM) is under active development. Track progress at: +Public API:: - https://github.com/ferro-labs/ai-gateway-workspace/blob/main/docs/OSS-ECOSYSTEM-ROADMAP.md + from langchain_ferrolabsai import FerroChatModel, FerroEmbeddings, FerroLLM -Once shipped, the public API will be: + chat = FerroChatModel(model="gpt-4o", api_key="sk-ferro-...") + embed = FerroEmbeddings(model="text-embedding-3-small", api_key="sk-ferro-...") + legacy = FerroLLM(model="gpt-4o", api_key="sk-ferro-...") - from langchain_ferrolabsai import FerroChatModel, FerroEmbeddings, FerroLLM +All three classes route through a Ferro Labs AI Gateway endpoint and expose +the gateway's ``trace_id`` (frozen contract since ``ai-gateway v1.1.0``) via +``response_metadata`` — the join key for the v1.2 observability bridge plugins +(LangSmith, Langfuse, Phoenix, …). """ from __future__ import annotations -__version__ = "0.0.1" - -__all__ = ["__version__"] - - -def _not_implemented(name: str) -> None: - raise NotImplementedError( - f"langchain_ferrolabsai.{name} is not implemented in 0.0.1. " - "This release reserves the package name. " - "Track progress at https://github.com/ferro-labs/ai-gateway-workspace." - ) +from .chat_models import FerroChatModel +from .embeddings import FerroEmbeddings +from .llms import FerroLLM +__version__ = "0.1.0" -def __getattr__(name: str) -> object: - if name in {"FerroChatModel", "FerroEmbeddings", "FerroLLM"}: - _not_implemented(name) - raise AttributeError(f"module 'langchain_ferrolabsai' has no attribute {name!r}") +__all__ = [ + "__version__", + "FerroChatModel", + "FerroEmbeddings", + "FerroLLM", +] diff --git a/integrations/langchain-ferrolabsai/langchain_ferrolabsai/_messages.py b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/_messages.py new file mode 100644 index 0000000..10b8d00 --- /dev/null +++ b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/_messages.py @@ -0,0 +1,79 @@ +"""Conversion helpers between LangChain ``BaseMessage`` types and the +OpenAI-compatible dicts the Ferro gateway expects. + +Kept tiny and dependency-light so the same module can be reused by the chat +model, the legacy ``FerroLLM`` adapter, and tests. +""" + +from __future__ import annotations + +from typing import Any + +from langchain_core.messages import ( + AIMessage, + BaseMessage, + ChatMessage, + FunctionMessage, + HumanMessage, + SystemMessage, + ToolMessage, +) + + +def message_to_ferro_dict(message: BaseMessage) -> dict[str, Any]: + """Convert a single LangChain ``BaseMessage`` to the OpenAI message shape. + + The Ferro gateway forwards this dict verbatim to the upstream provider, + so the encoding must match the OpenAI chat-completions schema. + """ + if isinstance(message, HumanMessage): + return {"role": "user", "content": message.content} + + if isinstance(message, SystemMessage): + return {"role": "system", "content": message.content} + + if isinstance(message, AIMessage): + d: dict[str, Any] = {"role": "assistant", "content": message.content or ""} + tool_calls = getattr(message, "tool_calls", None) or [] + if tool_calls: + d["tool_calls"] = [ + { + "id": tc.get("id", ""), + "type": "function", + "function": { + "name": tc.get("name", ""), + "arguments": tc.get("args", "") + if isinstance(tc.get("args"), str) + else _json_dumps(tc.get("args") or {}), + }, + } + for tc in tool_calls + ] + return d + + if isinstance(message, ToolMessage): + return { + "role": "tool", + "tool_call_id": message.tool_call_id, + "content": message.content, + } + + if isinstance(message, FunctionMessage): + # Legacy LangChain function-style messages map onto OpenAI ``function`` role. + return {"role": "function", "name": message.name, "content": message.content} + + if isinstance(message, ChatMessage): + return {"role": message.role, "content": message.content} + + raise TypeError(f"Unsupported LangChain message type: {type(message).__name__}") + + +def messages_to_ferro_dicts(messages: list[BaseMessage]) -> list[dict[str, Any]]: + """Convert a list of LangChain messages, preserving order.""" + return [message_to_ferro_dict(m) for m in messages] + + +def _json_dumps(value: Any) -> str: + import json + + return json.dumps(value, separators=(",", ":"), default=str) diff --git a/integrations/langchain-ferrolabsai/langchain_ferrolabsai/chat_models.py b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/chat_models.py new file mode 100644 index 0000000..220e7f4 --- /dev/null +++ b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/chat_models.py @@ -0,0 +1,325 @@ +"""FerroChatModel — LangChain ``BaseChatModel`` backed by Ferro Labs AI Gateway. + +A single ``FerroChatModel`` instance can address any of the gateway's 30+ +providers by name (e.g. ``"gpt-4o"``, ``"claude-3-5-sonnet-20241022"``, +``"gemini-1.5-flash"``) without changing the model class. + +Every response surfaces ``trace_id`` (the Ferro request ID propagated via the +``x-trace-id`` header — frozen contract since ``ai-gateway v1.1.0``) in +``response_metadata``. That value is the join key for any downstream +observability bridge plugin (LangSmith, Langfuse, Phoenix, …) that ships in +``ferro-labs/ai-gateway-plugins``. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterator, Sequence +from typing import Any + +from ferrolabsai import ChatCompletion, FerroClient +from langchain_core.callbacks import CallbackManagerForLLMRun +from langchain_core.language_models import BaseChatModel, LanguageModelInput +from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage +from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult +from langchain_core.runnables import Runnable +from langchain_core.tools import BaseTool +from langchain_core.utils.function_calling import convert_to_openai_tool +from pydantic import ConfigDict, Field, PrivateAttr, SecretStr + +from ._messages import messages_to_ferro_dicts + + +class FerroChatModel(BaseChatModel): + """Chat model that talks to the Ferro Labs AI Gateway. + + Example:: + + from langchain_ferrolabsai import FerroChatModel + from langchain_core.messages import HumanMessage + + chat = FerroChatModel(model="gpt-4o", api_key="sk-ferro-...") + response = chat.invoke([HumanMessage(content="Hello")]) + print(response.content) + print(response.response_metadata["trace_id"]) # Ferro request ID + """ + + model: str = Field(..., description="Model name routed by the gateway.") + base_url: str | None = Field( + default=None, + description="Gateway URL. Defaults to FERRO_BASE_URL env var or http://localhost:8080.", + ) + api_key: SecretStr | None = Field( + default=None, + description="API key. Defaults to FERRO_API_KEY env var.", + ) + timeout: float = 120.0 + max_retries: int = 2 + + temperature: float | None = None + max_tokens: int | None = None + top_p: float | None = None + frequency_penalty: float | None = None + presence_penalty: float | None = None + stop: list[str] | None = None + + # Ferro-specific extras + route_tag: str | None = Field( + default=None, + description="Override the gateway's routing strategy for this caller.", + ) + template_id: str | None = None + template_variables: dict[str, Any] | None = None + user: str | None = None + + default_headers: dict[str, str] | None = None + model_kwargs: dict[str, Any] = Field(default_factory=dict) + + model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True) + + _client_instance: FerroClient | None = PrivateAttr(default=None) + + # ------------------------------------------------------------------ + # LangChain identification + # ------------------------------------------------------------------ + + @property + def _llm_type(self) -> str: + return "ferro-labs-chat" + + @property + def _identifying_params(self) -> dict[str, Any]: + return { + "model": self.model, + "base_url": self.base_url, + "temperature": self.temperature, + "max_tokens": self.max_tokens, + } + + # ------------------------------------------------------------------ + # Client access + # ------------------------------------------------------------------ + + def _get_client(self) -> FerroClient: + if self._client_instance is None: + self._client_instance = FerroClient( + api_key=self.api_key.get_secret_value() if self.api_key else None, + base_url=self.base_url, + timeout=self.timeout, + max_retries=self.max_retries, + default_headers=self.default_headers, + ) + return self._client_instance + + # ------------------------------------------------------------------ + # Request payload assembly + # ------------------------------------------------------------------ + + def _build_request_params( + self, + stop: list[str] | None, + **kwargs: Any, + ) -> dict[str, Any]: + params: dict[str, Any] = {"model": self.model} + if self.temperature is not None: + params["temperature"] = self.temperature + if self.max_tokens is not None: + params["max_tokens"] = self.max_tokens + if self.top_p is not None: + params["top_p"] = self.top_p + if self.frequency_penalty is not None: + params["frequency_penalty"] = self.frequency_penalty + if self.presence_penalty is not None: + params["presence_penalty"] = self.presence_penalty + effective_stop = stop if stop is not None else self.stop + if effective_stop: + params["stop"] = effective_stop + if self.route_tag is not None: + params["route_tag"] = self.route_tag + if self.template_id is not None: + params["template_id"] = self.template_id + if self.template_variables is not None: + params["template_variables"] = self.template_variables + if self.user is not None: + params["user"] = self.user + # model_kwargs first so explicit per-call kwargs win. + params.update(self.model_kwargs) + params.update(kwargs) + return params + + # ------------------------------------------------------------------ + # Generation + # ------------------------------------------------------------------ + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + params = self._build_request_params(stop, **kwargs) + response = self._get_client().chat.completions.create( + messages=messages_to_ferro_dicts(messages), + **params, + ) + return _completion_to_chat_result(response) + + def _stream( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> Iterator[ChatGenerationChunk]: + params = self._build_request_params(stop, **kwargs) + stream = self._get_client().chat.completions.create( + messages=messages_to_ferro_dicts(messages), + stream=True, + **params, + ) + for chunk in stream: + if not chunk.choices: + continue + delta = chunk.choices[0].delta + content = delta.content or "" + ai_chunk = AIMessageChunk( + content=content, + tool_call_chunks=_extract_tool_call_chunks(delta.tool_calls), + ) + generation_chunk = ChatGenerationChunk( + message=ai_chunk, + generation_info={"finish_reason": chunk.choices[0].finish_reason} + if chunk.choices[0].finish_reason + else None, + ) + if run_manager is not None: + run_manager.on_llm_new_token(content, chunk=generation_chunk) + yield generation_chunk + + # ------------------------------------------------------------------ + # Tool binding (LangGraph / agent support) + # ------------------------------------------------------------------ + + def bind_tools( + self, + tools: Sequence[dict[str, Any] | type | BaseTool | Any], + *, + tool_choice: Any | None = None, + **kwargs: Any, + ) -> Runnable[LanguageModelInput, BaseMessage]: + formatted = [convert_to_openai_tool(t) for t in tools] + bind_kwargs: dict[str, Any] = {"tools": formatted} + if tool_choice is not None: + bind_kwargs["tool_choice"] = tool_choice + bind_kwargs.update(kwargs) + return super().bind(**bind_kwargs) + + +# --------------------------------------------------------------------------- +# Response mapping +# --------------------------------------------------------------------------- + + +def _completion_to_chat_result(response: ChatCompletion) -> ChatResult: + """Convert a Ferro ``ChatCompletion`` into a LangChain ``ChatResult``.""" + if not response.choices: + empty = AIMessage( + content="", + response_metadata=_response_metadata(response), + ) + return ChatResult(generations=[ChatGeneration(message=empty)]) + + choice = response.choices[0] + tool_calls = _extract_tool_calls(choice.message.tool_calls) + ai_message = AIMessage( + content=choice.message.content or "", + tool_calls=tool_calls, + response_metadata=_response_metadata(response), + usage_metadata=_usage_metadata(response), + ) + generation = ChatGeneration( + message=ai_message, + generation_info={"finish_reason": choice.finish_reason} if choice.finish_reason else None, + ) + return ChatResult( + generations=[generation], + llm_output={ + "model": response.model, + "trace_id": response.trace_id, + "provider": response.provider, + }, + ) + + +def _response_metadata(response: ChatCompletion) -> dict[str, Any]: + """The Ferro-specific surface every consumer (incl. v1.2 observability bridges) reads.""" + metadata: dict[str, Any] = { + "model": response.model, + "id": response.id, + # ``trace_id`` is the canonical join key. Frozen via x-trace-id since + # ai-gateway v1.1.0; mirrored by every Ferro observability bridge plugin. + "trace_id": response.trace_id, + "provider": response.provider, + "latency_ms": response.latency_ms, + } + if response.usage is not None: + metadata["cost_usd"] = response.usage.cost_usd + metadata["cache_hit"] = response.usage.cache_hit + return {k: v for k, v in metadata.items() if v is not None} + + +def _usage_metadata(response: ChatCompletion) -> dict[str, int] | None: + if response.usage is None: + return None + return { + "input_tokens": response.usage.prompt_tokens, + "output_tokens": response.usage.completion_tokens, + "total_tokens": response.usage.total_tokens, + } + + +def _extract_tool_call_chunks(raw: list[dict[str, Any]] | None) -> list[dict[str, Any]]: + """Map OpenAI streaming tool-call deltas to LangChain chunk shape.""" + if not raw: + return [] + result: list[dict[str, Any]] = [] + for call in raw: + function = call.get("function", {}) or {} + chunk: dict[str, Any] = {"type": "tool_call_chunk"} + if call.get("id") is not None: + chunk["id"] = call.get("id") + if call.get("index") is not None: + chunk["index"] = call.get("index") + if function.get("name") is not None: + chunk["name"] = function.get("name") + if function.get("arguments") is not None: + chunk["args"] = function.get("arguments") + result.append(chunk) + return result + + +def _extract_tool_calls(raw: list[dict[str, Any]] | None) -> list[dict[str, Any]]: + """Map OpenAI-style tool calls to LangChain's expected shape.""" + if not raw: + return [] + result: list[dict[str, Any]] = [] + for call in raw: + function = call.get("function", {}) + args_raw = function.get("arguments", "{}") + if isinstance(args_raw, str): + try: + args: Any = json.loads(args_raw) if args_raw else {} + except json.JSONDecodeError: + args = {"_raw": args_raw} + else: + args = args_raw + result.append( + { + "id": call.get("id", ""), + "name": function.get("name", ""), + "args": args, + "type": "tool_call", + } + ) + return result diff --git a/integrations/langchain-ferrolabsai/langchain_ferrolabsai/embeddings.py b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/embeddings.py new file mode 100644 index 0000000..35afb3c --- /dev/null +++ b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/embeddings.py @@ -0,0 +1,77 @@ +"""FerroEmbeddings — LangChain ``Embeddings`` backed by Ferro Labs AI Gateway. + +A single ``FerroEmbeddings`` instance routes to any embeddings-capable +provider the gateway knows about by model name (e.g. +``"text-embedding-3-small"``, ``"voyage-3"``, ``"cohere.embed-english-v3"``). +""" + +from __future__ import annotations + +from typing import Any + +from ferrolabsai import FerroClient +from langchain_core.embeddings import Embeddings +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, SecretStr + + +class FerroEmbeddings(BaseModel, Embeddings): + """LangChain embeddings adapter for the Ferro gateway. + + Example:: + + from langchain_ferrolabsai import FerroEmbeddings + + embed = FerroEmbeddings(model="text-embedding-3-small", api_key="sk-ferro-...") + vectors = embed.embed_documents(["hello", "world"]) + query_vec = embed.embed_query("hello") + """ + + model: str = Field(..., description="Embeddings model routed by the gateway.") + base_url: str | None = None + api_key: SecretStr | None = None + timeout: float = 120.0 + max_retries: int = 2 + dimensions: int | None = None + encoding_format: str | None = None + user: str | None = None + default_headers: dict[str, str] | None = None + + model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True) + + _client_instance: FerroClient | None = PrivateAttr(default=None) + + def _get_client(self) -> FerroClient: + if self._client_instance is None: + self._client_instance = FerroClient( + api_key=self.api_key.get_secret_value() if self.api_key else None, + base_url=self.base_url, + timeout=self.timeout, + max_retries=self.max_retries, + default_headers=self.default_headers, + ) + return self._client_instance + + def _build_kwargs(self) -> dict[str, Any]: + kwargs: dict[str, Any] = {"model": self.model} + if self.dimensions is not None: + kwargs["dimensions"] = self.dimensions + if self.encoding_format is not None: + kwargs["encoding_format"] = self.encoding_format + if self.user is not None: + kwargs["user"] = self.user + return kwargs + + def embed_documents(self, texts: list[str]) -> list[list[float]]: + if not texts: + return [] + response = self._get_client().embeddings.create(input=texts, **self._build_kwargs()) + # Preserve input order by sorting on `index` — the gateway / provider + # may return data out of order. + ordered = sorted(response.data, key=lambda d: d.index) + return [d.embedding for d in ordered] + + def embed_query(self, text: str) -> list[float]: + response = self._get_client().embeddings.create(input=text, **self._build_kwargs()) + if not response.data: + return [] + return response.data[0].embedding diff --git a/integrations/langchain-ferrolabsai/langchain_ferrolabsai/llms.py b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/llms.py new file mode 100644 index 0000000..2e68c37 --- /dev/null +++ b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/llms.py @@ -0,0 +1,73 @@ +"""FerroLLM — legacy completion-style LangChain ``LLM`` for the Ferro gateway. + +Most new code should use :class:`langchain_ferrolabsai.FerroChatModel` instead. +``FerroLLM`` exists for chains still built on the legacy ``LLM`` interface and +internally wraps the gateway's chat-completions endpoint with a single user +message. +""" + +from __future__ import annotations + +from typing import Any + +from ferrolabsai import FerroClient +from langchain_core.callbacks import CallbackManagerForLLMRun +from langchain_core.language_models import LLM +from pydantic import ConfigDict, Field, PrivateAttr, SecretStr + + +class FerroLLM(LLM): + """Legacy completion-style adapter — wraps chat completions with a single user message.""" + + model: str = Field(...) + base_url: str | None = None + api_key: SecretStr | None = None + timeout: float = 120.0 + max_retries: int = 2 + temperature: float | None = None + max_tokens: int | None = None + route_tag: str | None = None + default_headers: dict[str, str] | None = None + + model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True) + + _client_instance: FerroClient | None = PrivateAttr(default=None) + + @property + def _llm_type(self) -> str: + return "ferro-labs" + + def _get_client(self) -> FerroClient: + if self._client_instance is None: + self._client_instance = FerroClient( + api_key=self.api_key.get_secret_value() if self.api_key else None, + base_url=self.base_url, + timeout=self.timeout, + max_retries=self.max_retries, + default_headers=self.default_headers, + ) + return self._client_instance + + def _call( + self, + prompt: str, + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> str: + params: dict[str, Any] = {"model": self.model} + if self.temperature is not None: + params["temperature"] = self.temperature + if self.max_tokens is not None: + params["max_tokens"] = self.max_tokens + if stop is not None: + params["stop"] = stop + if self.route_tag is not None: + params["route_tag"] = self.route_tag + params.update(kwargs) + + response = self._get_client().chat.completions.create( + messages=[{"role": "user", "content": prompt}], + **params, + ) + return response.content or "" diff --git a/integrations/langchain-ferrolabsai/pyproject.toml b/integrations/langchain-ferrolabsai/pyproject.toml index f8c33bb..68ea54c 100644 --- a/integrations/langchain-ferrolabsai/pyproject.toml +++ b/integrations/langchain-ferrolabsai/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "langchain-ferrolabsai" -version = "0.0.1" +version = "0.1.0" description = "LangChain integration for Ferro Labs AI Gateway — chat, streaming, embeddings, and tool calling across 30+ LLM providers via a single OpenAI-compatible endpoint" readme = "README.md" license = { text = "Apache-2.0" } @@ -18,7 +18,7 @@ keywords = [ "langchain-integration", "generative-ai", ] classifiers = [ - "Development Status :: 3 - Alpha", + "Development Status :: 4 - Beta", "Intended Audience :: Developers", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", diff --git a/integrations/langchain-ferrolabsai/tests/__init__.py b/integrations/langchain-ferrolabsai/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/integrations/langchain-ferrolabsai/tests/conftest.py b/integrations/langchain-ferrolabsai/tests/conftest.py new file mode 100644 index 0000000..a5802a5 --- /dev/null +++ b/integrations/langchain-ferrolabsai/tests/conftest.py @@ -0,0 +1,78 @@ +"""Shared test fixtures for langchain-ferrolabsai. + +Follows the same pytest-httpx mocking pattern as the parent ferrolabsai SDK so +no real gateway is required to run the suite. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +BASE_URL = "http://test-gateway:8080" +API_KEY = "sk-ferro-test" + + +def make_chat_completion( + *, + content: str = "Hello back", + model: str = "gpt-4o", + provider: str = "openai", + trace_id: str = "trace-abc-123", + latency_ms: int = 42, + cost_usd: float = 0.000123, + tool_calls: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + """Build a Ferro chat-completion response payload for use with httpx_mock.""" + message: dict[str, Any] = {"role": "assistant", "content": content} + if tool_calls is not None: + message["tool_calls"] = tool_calls + return { + "id": "cmpl-1", + "object": "chat.completion", + "created": 1_700_000_000, + "model": model, + "choices": [ + { + "index": 0, + "message": message, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 3, + "total_tokens": 8, + "cost_usd": cost_usd, + "provider": provider, + }, + "x_ferro_trace_id": trace_id, + "x_ferro_provider": provider, + "x_ferro_latency_ms": latency_ms, + } + + +def make_embedding_response( + *, + vectors: list[list[float]], + model: str = "text-embedding-3-small", +) -> dict[str, Any]: + return { + "object": "list", + "model": model, + "data": [ + {"index": i, "embedding": v, "object": "embedding"} for i, v in enumerate(vectors) + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 0, "total_tokens": 1}, + } + + +@pytest.fixture +def base_url() -> str: + return BASE_URL + + +@pytest.fixture +def api_key() -> str: + return API_KEY diff --git a/integrations/langchain-ferrolabsai/tests/test_chat_models.py b/integrations/langchain-ferrolabsai/tests/test_chat_models.py new file mode 100644 index 0000000..6df7ede --- /dev/null +++ b/integrations/langchain-ferrolabsai/tests/test_chat_models.py @@ -0,0 +1,306 @@ +"""Tests for FerroChatModel.""" + +from __future__ import annotations + +import json + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage +from langchain_core.tools import tool +from pytest_httpx import HTTPXMock + +from langchain_ferrolabsai import FerroChatModel + +from .conftest import BASE_URL, make_chat_completion + + +def _build_chat(**overrides) -> FerroChatModel: + kwargs = {"model": "gpt-4o", "base_url": BASE_URL, "api_key": "sk-ferro-test"} + kwargs.update(overrides) + return FerroChatModel(**kwargs) # type: ignore[arg-type] + + +class TestBasicGeneration: + def test_invoke_returns_ai_message_with_content(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion(content="hi there"), + ) + chat = _build_chat() + result = chat.invoke([HumanMessage(content="Hello")]) + assert isinstance(result, AIMessage) + assert result.content == "hi there" + + def test_invoke_surfaces_trace_id_in_response_metadata(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion(trace_id="my-trace-xyz"), + ) + chat = _build_chat() + result = chat.invoke([HumanMessage(content="Hello")]) + # trace_id is the join key for v1.2 observability bridges — MUST be present. + assert result.response_metadata["trace_id"] == "my-trace-xyz" + assert result.response_metadata["provider"] == "openai" + assert result.response_metadata["latency_ms"] == 42 + assert result.response_metadata["cost_usd"] == 0.000123 + + def test_invoke_surfaces_header_only_gateway_metadata(self, httpx_mock: HTTPXMock): + body = make_chat_completion(trace_id="body-trace", provider="body-provider") + body.pop("x_ferro_trace_id") + body.pop("x_ferro_provider") + body.pop("x_ferro_latency_ms") + body["usage"].pop("cost_usd") + body["usage"].pop("provider") + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=body, + headers={ + "X-Request-ID": "header-trace", + "x-ferro-provider": "openai", + "x-ferro-latency-ms": "42", + "x-ferro-cost-usd": "0.000123", + }, + ) + chat = _build_chat() + result = chat.invoke([HumanMessage(content="Hello")]) + assert result.response_metadata["trace_id"] == "header-trace" + assert result.response_metadata["provider"] == "openai" + assert result.response_metadata["latency_ms"] == 42 + assert result.response_metadata["cost_usd"] == 0.000123 + + def test_invoke_attaches_usage_metadata(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion(), + ) + chat = _build_chat() + result = chat.invoke([HumanMessage(content="Hello")]) + assert result.usage_metadata == { + "input_tokens": 5, + "output_tokens": 3, + "total_tokens": 8, + } + + +class TestMessageConversion: + def test_system_human_messages_serialized_correctly(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion(), + ) + chat = _build_chat() + chat.invoke([SystemMessage(content="be terse"), HumanMessage(content="hi")]) + + body = json.loads(httpx_mock.get_requests()[0].content) + assert body["messages"] == [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"}, + ] + + def test_tool_messages_carry_tool_call_id(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion(), + ) + chat = _build_chat() + chat.invoke( + [ + HumanMessage(content="what is 1+1"), + AIMessage( + content="", + tool_calls=[{"id": "c1", "name": "add", "args": {"a": 1, "b": 1}}], + ), + ToolMessage(content="2", tool_call_id="c1"), + ] + ) + body = json.loads(httpx_mock.get_requests()[0].content) + tool_msg = body["messages"][-1] + assert tool_msg["role"] == "tool" + assert tool_msg["tool_call_id"] == "c1" + assert tool_msg["content"] == "2" + + +class TestRequestParams: + def test_sends_auth_header(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion(), + ) + chat = _build_chat(api_key="sk-ferro-prod") + chat.invoke([HumanMessage(content="hi")]) + request = httpx_mock.get_requests()[0] + assert request.headers["Authorization"] == "Bearer sk-ferro-prod" + + def test_forwards_temperature_and_max_tokens(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion(), + ) + chat = _build_chat(temperature=0.2, max_tokens=64) + chat.invoke([HumanMessage(content="hi")]) + body = json.loads(httpx_mock.get_requests()[0].content) + assert body["temperature"] == 0.2 + assert body["max_tokens"] == 64 + + def test_forwards_ferro_extras(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion(), + ) + chat = _build_chat( + route_tag="premium", + template_id="customer-support", + template_variables={"tone": "friendly"}, + user="user-123", + ) + chat.invoke([HumanMessage(content="hi")]) + body = json.loads(httpx_mock.get_requests()[0].content) + # Ferro forwards `route_tag` as `x_route_tag` internally; we just + # check the field round-trips through the SDK's request builder. + assert body.get("x_route_tag") == "premium" + assert body["template_id"] == "customer-support" + assert body["template_variables"] == {"tone": "friendly"} + assert body["user"] == "user-123" + + +class TestToolBinding: + def test_bind_tools_forwards_openai_tool_schema(self, httpx_mock: HTTPXMock): + @tool + def add(a: int, b: int) -> int: + """Add two integers.""" + return a + b + + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion( + content="", + tool_calls=[ + { + "id": "call_1", + "type": "function", + "function": {"name": "add", "arguments": '{"a":1,"b":2}'}, + } + ], + ), + ) + chat = _build_chat().bind_tools([add]) + result = chat.invoke([HumanMessage(content="add 1 and 2")]) + + body = json.loads(httpx_mock.get_requests()[0].content) + assert body["tools"][0]["type"] == "function" + assert body["tools"][0]["function"]["name"] == "add" + + assert result.tool_calls == [ + {"id": "call_1", "name": "add", "args": {"a": 1, "b": 2}, "type": "tool_call"} + ] + + +class TestStreaming: + def test_stream_yields_chunks(self, httpx_mock: HTTPXMock): + sse_body = ( + 'data: {"id":"1","object":"chat.completion.chunk","created":1,"model":"gpt-4o",' + '"choices":[{"index":0,"delta":{"role":"assistant","content":"Hel"},"finish_reason":null}]}\n\n' + 'data: {"id":"1","object":"chat.completion.chunk","created":1,"model":"gpt-4o",' + '"choices":[{"index":0,"delta":{"content":"lo"},"finish_reason":null}]}\n\n' + 'data: {"id":"1","object":"chat.completion.chunk","created":1,"model":"gpt-4o",' + '"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n' + "data: [DONE]\n\n" + ) + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + content=sse_body.encode("utf-8"), + headers={"Content-Type": "text/event-stream"}, + ) + chat = _build_chat() + chunks = list(chat.stream([HumanMessage(content="hi")])) + assert "".join(c.content for c in chunks) == "Hello" + + def test_stream_yields_tool_call_chunks(self, httpx_mock: HTTPXMock): + frames = [ + { + "id": "1", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "add", "arguments": '{"a":1'}, + } + ] + }, + "finish_reason": None, + } + ], + }, + { + "id": "1", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + {"index": 0, "function": {"arguments": ',"b":2}'}} + ] + }, + "finish_reason": None, + } + ], + }, + ] + sse_body = "".join(f"data: {json.dumps(frame)}\n\n" for frame in frames) + sse_body += "data: [DONE]\n\n" + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + content=sse_body.encode("utf-8"), + headers={"Content-Type": "text/event-stream"}, + ) + chat = _build_chat() + chunks = list(chat.stream([HumanMessage(content="add 1 and 2")])) + tool_chunks = [tc for chunk in chunks for tc in chunk.tool_call_chunks] + assert tool_chunks[0]["id"] == "call_1" + assert tool_chunks[0]["name"] == "add" + assert tool_chunks[0]["args"] == '{"a":1' + assert tool_chunks[1]["args"] == ',"b":2}' + + + +class TestIdentity: + def test_llm_type(self): + assert _build_chat()._llm_type == "ferro-labs-chat" + + def test_identifying_params_include_model(self): + params = _build_chat(temperature=0.5)._identifying_params + assert params["model"] == "gpt-4o" + assert params["temperature"] == 0.5 + + +class TestApiKeyHandling: + def test_missing_api_key_raises(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("FERRO_API_KEY", raising=False) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + chat = FerroChatModel(model="gpt-4o", base_url=BASE_URL) # no api_key + # FerroClient construction happens lazily on first call. + with pytest.raises(Exception): + chat.invoke([HumanMessage(content="hi")]) diff --git a/integrations/langchain-ferrolabsai/tests/test_embeddings.py b/integrations/langchain-ferrolabsai/tests/test_embeddings.py new file mode 100644 index 0000000..9488e97 --- /dev/null +++ b/integrations/langchain-ferrolabsai/tests/test_embeddings.py @@ -0,0 +1,93 @@ +"""Tests for FerroEmbeddings.""" + +from __future__ import annotations + +import json + +from pytest_httpx import HTTPXMock + +from langchain_ferrolabsai import FerroEmbeddings + +from .conftest import BASE_URL, make_embedding_response + + +def _build(**overrides) -> FerroEmbeddings: + kwargs = { + "model": "text-embedding-3-small", + "base_url": BASE_URL, + "api_key": "sk-ferro-test", + } + kwargs.update(overrides) + return FerroEmbeddings(**kwargs) # type: ignore[arg-type] + + +class TestEmbedDocuments: + def test_returns_vector_per_input(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/embeddings", + json=make_embedding_response(vectors=[[0.1, 0.2], [0.3, 0.4]]), + ) + embed = _build() + result = embed.embed_documents(["hello", "world"]) + assert result == [[0.1, 0.2], [0.3, 0.4]] + + def test_preserves_input_order_when_gateway_returns_out_of_order(self, httpx_mock: HTTPXMock): + # Simulate the gateway returning data with index 1 before index 0. + out_of_order = { + "object": "list", + "model": "text-embedding-3-small", + "data": [ + {"index": 1, "embedding": [0.3, 0.4], "object": "embedding"}, + {"index": 0, "embedding": [0.1, 0.2], "object": "embedding"}, + ], + } + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/embeddings", + json=out_of_order, + ) + embed = _build() + result = embed.embed_documents(["a", "b"]) + assert result == [[0.1, 0.2], [0.3, 0.4]] + + def test_empty_input_returns_empty_list_without_request(self, httpx_mock: HTTPXMock): + # No httpx_mock response registered — would error if a request was made. + embed = _build() + assert embed.embed_documents([]) == [] + assert httpx_mock.get_requests() == [] + + def test_forwards_optional_params(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/embeddings", + json=make_embedding_response(vectors=[[0.0]]), + ) + embed = _build(dimensions=512, encoding_format="float", user="u1") + embed.embed_documents(["x"]) + body = json.loads(httpx_mock.get_requests()[0].content) + assert body["dimensions"] == 512 + assert body["encoding_format"] == "float" + assert body["user"] == "u1" + + +class TestEmbedQuery: + def test_returns_single_vector(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/embeddings", + json=make_embedding_response(vectors=[[0.9, 0.8, 0.7]]), + ) + embed = _build() + assert embed.embed_query("hello") == [0.9, 0.8, 0.7] + + def test_sends_single_string_input(self, httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/embeddings", + json=make_embedding_response(vectors=[[0.0]]), + ) + embed = _build() + embed.embed_query("hello") + body = json.loads(httpx_mock.get_requests()[0].content) + assert body["input"] == "hello" diff --git a/integrations/langchain-ferrolabsai/tests/test_llms.py b/integrations/langchain-ferrolabsai/tests/test_llms.py new file mode 100644 index 0000000..a1fc9c7 --- /dev/null +++ b/integrations/langchain-ferrolabsai/tests/test_llms.py @@ -0,0 +1,56 @@ +"""Tests for the legacy FerroLLM completion-style adapter.""" + +from __future__ import annotations + +import json + +from pytest_httpx import HTTPXMock + +from langchain_ferrolabsai import FerroLLM + +from .conftest import BASE_URL, make_chat_completion + + +def _build(**overrides) -> FerroLLM: + kwargs = {"model": "gpt-4o", "base_url": BASE_URL, "api_key": "sk-ferro-test"} + kwargs.update(overrides) + return FerroLLM(**kwargs) # type: ignore[arg-type] + + +def test_invoke_returns_string(httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion(content="hello world"), + ) + llm = _build() + assert llm.invoke("hi") == "hello world" + + +def test_wraps_prompt_in_single_user_message(httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion(), + ) + llm = _build() + llm.invoke("ping") + body = json.loads(httpx_mock.get_requests()[0].content) + assert body["messages"] == [{"role": "user", "content": "ping"}] + + +def test_llm_type(): + assert _build()._llm_type == "ferro-labs" + + +def test_forwards_temperature_and_max_tokens(httpx_mock: HTTPXMock): + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=make_chat_completion(), + ) + llm = _build(temperature=0.7, max_tokens=128) + llm.invoke("hi") + body = json.loads(httpx_mock.get_requests()[0].content) + assert body["temperature"] == 0.7 + assert body["max_tokens"] == 128 diff --git a/integrations/langchain-ferrolabsai/tests/test_placeholder.py b/integrations/langchain-ferrolabsai/tests/test_placeholder.py deleted file mode 100644 index a7acb80..0000000 --- a/integrations/langchain-ferrolabsai/tests/test_placeholder.py +++ /dev/null @@ -1,31 +0,0 @@ -"""Sanity tests for the langchain-ferrolabsai 0.0.1 placeholder release. - -These tests verify that the placeholder package imports cleanly, exposes the -expected version string, and that any attempt to use the planned public API -fails with a clear NotImplementedError pointing at the roadmap. -""" - -from __future__ import annotations - -import pytest - -import langchain_ferrolabsai - - -def test_version_is_placeholder() -> None: - assert langchain_ferrolabsai.__version__ == "0.0.1" - - -@pytest.mark.parametrize("name", ["FerroChatModel", "FerroEmbeddings", "FerroLLM"]) -def test_planned_api_raises_not_implemented(name: str) -> None: - with pytest.raises(NotImplementedError) as excinfo: - getattr(langchain_ferrolabsai, name) - message = str(excinfo.value) - assert name in message - assert "0.0.1" in message - assert "ferro-labs" in message.lower() - - -def test_unknown_attribute_raises_attribute_error() -> None: - with pytest.raises(AttributeError): - langchain_ferrolabsai.NonExistentClass # type: ignore[attr-defined] diff --git a/tests/test_sdk.py b/tests/test_sdk.py index 41c183b..e0dfe71 100644 --- a/tests/test_sdk.py +++ b/tests/test_sdk.py @@ -140,6 +140,31 @@ def test_basic_create(self, client, httpx_mock: HTTPXMock): assert response.content == "Hello from Ferro!" assert response.usage.cost_usd == 0.000075 + def test_success_metadata_can_come_from_headers(self, client, httpx_mock: HTTPXMock): + body = dict(COMPLETION_RESPONSE) + body["usage"] = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=body, + headers={ + "X-Request-ID": "trace-from-header", + "x-ferro-provider": "openai", + "x-ferro-latency-ms": "42", + "x-ferro-cost-usd": "0.000075", + }, + ) + response = client.chat.completions.create( + model="gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + ) + assert response.trace_id == "trace-from-header" + assert response.provider == "openai" + assert response.latency_ms == 42 + assert response.usage is not None + assert response.usage.cost_usd == 0.000075 + assert response.usage.provider == "openai" + def test_sends_correct_headers(self, client, httpx_mock: HTTPXMock): httpx_mock.add_response( method="POST", @@ -650,6 +675,33 @@ async def test_forwards_route_tag(self, async_client, httpx_mock: HTTPXMock): body = json.loads(httpx_mock.get_requests()[0].content) assert body["x_route_tag"] == "fast" + @pytest.mark.asyncio + async def test_success_metadata_can_come_from_headers( + self, async_client, httpx_mock: HTTPXMock + ): + body = dict(COMPLETION_RESPONSE) + body["usage"] = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + httpx_mock.add_response( + method="POST", + url=f"{BASE_URL}/v1/chat/completions", + json=body, + headers={ + "X-Request-ID": "async-trace-from-header", + "x-ferro-provider": "anthropic", + "x-ferro-latency-ms": "123", + "x-ferro-cost-usd": "0.001", + }, + ) + response = await async_client.chat.completions.create( + model="gpt-4o", + messages=[{"role": "user", "content": "Hi"}], + ) + assert response.trace_id == "async-trace-from-header" + assert response.provider == "anthropic" + assert response.latency_ms == 123 + assert response.usage is not None + assert response.usage.cost_usd == 0.001 + # ------------------------------------------------------------------ # P0-3: BYOC http_client merges auth headers From 0c5cd5aa142d447f4899b4aa173132bbc2505079 Mon Sep 17 00:00:00 2001 From: Mitul Shah Date: Thu, 28 May 2026 11:52:50 +0530 Subject: [PATCH 2/5] fix: stabilize langchain integration CI --- .../workflows/publish-langchain-ferrolabsai.yml | 4 +++- .../langchain_ferrolabsai/chat_models.py | 15 +++++++++------ integrations/langchain-ferrolabsai/pyproject.toml | 5 +++-- 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/.github/workflows/publish-langchain-ferrolabsai.yml b/.github/workflows/publish-langchain-ferrolabsai.yml index 1023710..14ccd32 100644 --- a/.github/workflows/publish-langchain-ferrolabsai.yml +++ b/.github/workflows/publish-langchain-ferrolabsai.yml @@ -32,7 +32,9 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install dependencies - run: pip install -e ".[dev]" + run: | + pip install -e ../.. + pip install -e ".[dev]" - name: Lint run: ruff check langchain_ferrolabsai/ diff --git a/integrations/langchain-ferrolabsai/langchain_ferrolabsai/chat_models.py b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/chat_models.py index 220e7f4..2830da0 100644 --- a/integrations/langchain-ferrolabsai/langchain_ferrolabsai/chat_models.py +++ b/integrations/langchain-ferrolabsai/langchain_ferrolabsai/chat_models.py @@ -15,7 +15,7 @@ import json from collections.abc import Iterator, Sequence -from typing import Any +from typing import TYPE_CHECKING, Any, cast from ferrolabsai import ChatCompletion, FerroClient from langchain_core.callbacks import CallbackManagerForLLMRun @@ -29,6 +29,9 @@ from ._messages import messages_to_ferro_dicts +if TYPE_CHECKING: + from langchain_core.messages import ToolCallChunk + class FerroChatModel(BaseChatModel): """Chat model that talks to the Ferro Labs AI Gateway. @@ -207,13 +210,13 @@ def bind_tools( *, tool_choice: Any | None = None, **kwargs: Any, - ) -> Runnable[LanguageModelInput, BaseMessage]: + ) -> Runnable[LanguageModelInput, AIMessage]: formatted = [convert_to_openai_tool(t) for t in tools] bind_kwargs: dict[str, Any] = {"tools": formatted} if tool_choice is not None: bind_kwargs["tool_choice"] = tool_choice bind_kwargs.update(kwargs) - return super().bind(**bind_kwargs) + return cast("Runnable[LanguageModelInput, AIMessage]", super().bind(**bind_kwargs)) # --------------------------------------------------------------------------- @@ -279,11 +282,11 @@ def _usage_metadata(response: ChatCompletion) -> dict[str, int] | None: } -def _extract_tool_call_chunks(raw: list[dict[str, Any]] | None) -> list[dict[str, Any]]: +def _extract_tool_call_chunks(raw: list[dict[str, Any]] | None) -> list[ToolCallChunk]: """Map OpenAI streaming tool-call deltas to LangChain chunk shape.""" if not raw: return [] - result: list[dict[str, Any]] = [] + result: list[ToolCallChunk] = [] for call in raw: function = call.get("function", {}) or {} chunk: dict[str, Any] = {"type": "tool_call_chunk"} @@ -295,7 +298,7 @@ def _extract_tool_call_chunks(raw: list[dict[str, Any]] | None) -> list[dict[str chunk["name"] = function.get("name") if function.get("arguments") is not None: chunk["args"] = function.get("arguments") - result.append(chunk) + result.append(cast("ToolCallChunk", chunk)) return result diff --git a/integrations/langchain-ferrolabsai/pyproject.toml b/integrations/langchain-ferrolabsai/pyproject.toml index 68ea54c..08529c0 100644 --- a/integrations/langchain-ferrolabsai/pyproject.toml +++ b/integrations/langchain-ferrolabsai/pyproject.toml @@ -40,7 +40,7 @@ dev = [ "pytest>=7.0", "pytest-asyncio>=0.21", "pytest-httpx>=0.22", - "mypy>=1.0", + "mypy>=1.0,<2", "ruff>=0.1.0", ] @@ -62,9 +62,10 @@ target-version = "py39" select = ["E", "F", "I", "UP"] [tool.mypy] -python_version = "3.9" +python_version = "3.10" strict = true ignore_missing_imports = true +exclude = '^(\.venv|build|dist)/' [tool.pytest.ini_options] asyncio_mode = "auto" From 550a85dac5e74568444773b7912c9b7c6133ae84 Mon Sep 17 00:00:00 2001 From: Mitul Shah Date: Thu, 28 May 2026 11:54:59 +0530 Subject: [PATCH 3/5] fix: support langchain integration on python 3.9 --- integrations/langchain-ferrolabsai/pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/integrations/langchain-ferrolabsai/pyproject.toml b/integrations/langchain-ferrolabsai/pyproject.toml index 08529c0..70103f3 100644 --- a/integrations/langchain-ferrolabsai/pyproject.toml +++ b/integrations/langchain-ferrolabsai/pyproject.toml @@ -33,6 +33,7 @@ classifiers = [ dependencies = [ "ferrolabsai>=0.1.0", "langchain-core>=0.3.0", + "eval-type-backport>=0.2.0; python_version < '3.10'", ] [project.optional-dependencies] From 2887373d07f16225f27d0149a2688ff6b777d8e4 Mon Sep 17 00:00:00 2001 From: Mitul Shah Date: Thu, 28 May 2026 13:29:59 +0530 Subject: [PATCH 4/5] Development (#4) (#6) * feat: ship langchain adapter * fix: support langchain integration on python 3.9 From 41849c8ab4d605a420cca2a2fc2a41c4b74e58fe Mon Sep 17 00:00:00 2001 From: Dhyana Date: Mon, 1 Jun 2026 19:07:31 +0800 Subject: [PATCH 5/5] test: cover admin plugins list responses --- CHANGELOG.md | 5 +++++ tests/test_sdk.py | 22 ++++++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index fae7614..e23211d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,11 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm --- +## [Unreleased] + +### Added +- Test coverage for `admin.plugins.list()` response shapes. + ## [0.2.0] — 2026-05-14 ### Added diff --git a/tests/test_sdk.py b/tests/test_sdk.py index e0dfe71..2294ce6 100644 --- a/tests/test_sdk.py +++ b/tests/test_sdk.py @@ -522,6 +522,28 @@ def test_logs_stats(self, client, httpx_mock: HTTPXMock): assert stats["total"] == 42 +class TestAdminPlugins: + def test_list_plugins_bare_array(self, client, httpx_mock: HTTPXMock): + payload = [{"name": "logger", "enabled": True}] + httpx_mock.add_response( + method="GET", + url=f"{BASE_URL}/admin/plugins", + json=payload, + ) + plugins = client.admin.plugins.list() + assert plugins == payload + + def test_list_plugins_data_wrapper(self, client, httpx_mock: HTTPXMock): + payload = [{"name": "cache", "enabled": False}] + httpx_mock.add_response( + method="GET", + url=f"{BASE_URL}/admin/plugins", + json={"data": payload}, + ) + plugins = client.admin.plugins.list() + assert plugins == payload + + # ------------------------------------------------------------------ # Error handling # ------------------------------------------------------------------