Surface LangSmith Gateway Metadata On Successful Fireworks Completions
ChatFireworks does not capture the X-LangSmith-Gateway-Metadata header from successful gateway-routed responses, causing missing or incorrect model/provider information in LangSmith traces. This fix adds gateway metadata capture conditioned on the lsv2_ API key prefix, mirroring OpenAI and Anthropic integrations.
ChatFireworks calls client.create() directly, which discards HTTP response headers. Unlike ChatOpenAI and ChatAnthropic, it never uses with_raw_response to access the X-LangSmith-Gateway-Metadata header sent by the LangSmith gateway on success.
1. Set LANGSMITH_GATEWAY=true and use a Fireworks model through the gateway.\n2. Call ChatFireworks with streaming and non-streaming.\n3. Inspect the LangSmith trace; observe that ls_gateway_info is absent and ls_model_name/ls_provider reflect the requested model rather than the gateway-resolved one.
Fixing Code Block
```python
# Add these methods and modify existing ones in ChatFireworks class
from langchain_core.utils._gateway import _parse_gateway_metadata
class ChatFireworks(BaseChatModel):
# ... existing code ...
@property
def _uses_gateway(self) -> bool:
"""Check if the API key indicates LangSmith gateway usage."""
return self.fireworks_api_key is not None and self.fireworks_api_key.startswith("lsv2_")
def _add_gateway_metadata(self, generation_info: dict, headers: dict) -> None:
"""Parse and attach gateway metadata to generation info."""
gateway_metadata = _parse_gateway_metadata(headers)
if gateway_metadata:
generation_info["ls_gateway_info"] = gateway_metadata
def _completion_with_retry(self, **kwargs):
"""Wrap completion call to optionally capture response headers."""
if not self._uses_gateway:
return self.client.create(**kwargs)
response = self.client.with_raw_response.create(**kwargs)
headers = getattr(response, "headers", {})
parsed_response = response.parse()
return parsed_response, headers
async def _acompletion_with_retry(self, **kwargs):
"""Async version of _completion_with_retry."""
if not self._uses_gateway:
return await self.client.create(**kwargs)
response = await self.client.with_raw_response.create(**kwargs)
headers = getattr(response, "headers", {})
parsed_response = response.parse()
return parsed_response, headers
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
"""Modify to attach gateway metadata to each generation."""
params = self._get_invocation_params(stop=stop, **kwargs)
if self._uses_gateway:
response, headers = self._completion_with_retry(messages=messages, **params)
chat_result = self._create_chat_result(response)
for generation in chat_result.generations:
self._add_gateway_metadata(generation.generation_info, headers)
else:
response = self._completion_with_retry(messages=messages, **params)
chat_result = self._create_chat_result(response)
return chat_result
async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs):
"""Async version of _generate."""
params = self._get_invocation_params(stop=stop, **kwargs)
if self._uses_gateway:
response, headers = await self._acompletion_with_retry(messages=messages, **params)
chat_result = self._create_chat_result(response)
for generation in chat_result.generations:
self._add_gateway_metadata(generation.generation_info, headers)
else:
response = await self._acompletion_with_retry(messages=messages, **params)
chat_result = self._create_chat_result(response)
return chat_result
def _stream(self, messages, stop=None, run_manager=None, **kwargs):
"""Modify to attach gateway metadata to the first streamed chunk."""
params = self._get_invocation_params(stop=stop, **kwargs)
if self._uses_gateway:
response, headers = self._completion_with_retry(messages=messages, stream=True, **params)
# response is a generator/iterator of chunks
stream_iter = iter(response)
try:
first_chunk = next(stream_iter)
except StopIteration:
return iter(())
if first_chunk.generation_info is None:
first_chunk.generation_info = {}
self._add_gateway_metadata(first_chunk.generation_info, headers)
def gen():
yield first_chunk
yield from stream_iter
return gen()
else:
return self._completion_with_retry(messages=messages, stream=True, **params)
async def _astream(self, messages, stop=None, run_manager=None, **kwargs):
"""Async version of _stream."""
params = self._get_invocation_params(stop=stop, **kwargs)
if self._uses_gateway:
response, headers = await self._acompletion_with_retry(messages=messages, stream=True, **params)
stream_iter = response.__aiter__()
try:
first_chunk = await stream_iter.__anext__()
except StopAsyncIteration:
return _empty_async_iter()
if first_chunk.generation_info is None:
first_chunk.generation_info = {}
self._add_gateway_metadata(first_chunk.generation_info, headers)
async def agen():
yield first_chunk
async for chunk in stream_iter:
yield chunk
return agen()
else:
return await self._acompletion_with_retry(messages=messages, stream=True, **params)
```
The fix adds a `_uses_gateway` property to detect the lsv2_ API key prefix. When gateway is active, it uses `client.with_raw_response.create()` to preserve headers, parses them with the shared `_parse_gateway_metadata`, and returns the parsed body plus headers. Non-streaming paths attach metadata to each generation's `generation_info`; streaming paths attach it to the first chunk. Non-gateway code paths remain unchanged.
Edge Case Audit
This change introduces a new tuple return type from `_completion_with_retry` when gateway is active, which could break subclasses or external code that overrides these methods expecting a single return value. The `with_raw_response` API availability depends on the installed fireworks-ai SDK version; older versions may not support it, leading to AttributeError. The async stream path assumes the returned object is an async iterator; if the SDK returns a sync iterator in async mode, it will fail. Custom clients passed by the user that do not provide `with_raw_response` will cause issues when the gateway key is used, but the `_uses_gateway` check should prevent this unless a user erroneously uses a gateway key with a custom client. To rollback, revert the changed methods and remove the `_uses_gateway` property; ensure all tests are rerun.
