mirror of
https://github.com/hwchase17/langchain.git
synced 2026-06-09 10:17:00 +00:00
test(fireworks): stabilize integration tests with rate limiting and retries (#37590)
Fireworks integration tests have been flaky against the live API with 429s. Adds a shared, xdist-aware rate limiter and a global retry policy so transient rate-limit errors no longer fail the suite. Mirrors the same fix recently applied to `langchain-mistralai`.
This commit is contained in:
@@ -0,0 +1,18 @@
|
||||
"""Shared rate limiter for Fireworks integration tests.
|
||||
|
||||
Scaled by ``PYTEST_XDIST_WORKER_COUNT`` so aggregate QPS across all xdist
|
||||
workers stays bounded near the target rate.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from langchain_core.rate_limiters import InMemoryRateLimiter
|
||||
|
||||
_TARGET_REQUESTS_PER_SECOND = 0.5
|
||||
_WORKER_COUNT = max(1, int(os.environ.get("PYTEST_XDIST_WORKER_COUNT", "1")))
|
||||
|
||||
rate_limiter = InMemoryRateLimiter(
|
||||
requests_per_second=_TARGET_REQUESTS_PER_SECOND / _WORKER_COUNT,
|
||||
)
|
||||
@@ -14,6 +14,7 @@ from pydantic import BaseModel, Field
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langchain_fireworks import ChatFireworks
|
||||
from tests.integration_tests._rate_limiter import rate_limiter
|
||||
|
||||
_MODEL = "accounts/fireworks/models/gpt-oss-120b"
|
||||
|
||||
@@ -21,7 +22,9 @@ _MODEL = "accounts/fireworks/models/gpt-oss-120b"
|
||||
@pytest.mark.parametrize("strict", [None, True, False])
|
||||
def test_tool_choice_bool(strict: bool | None) -> None: # noqa: FBT001
|
||||
"""Test that tool choice is respected with different strict values."""
|
||||
llm = ChatFireworks(model="accounts/fireworks/models/kimi-k2p6")
|
||||
llm = ChatFireworks(
|
||||
model="accounts/fireworks/models/kimi-k2p6", rate_limiter=rate_limiter
|
||||
)
|
||||
|
||||
class MyTool(BaseModel):
|
||||
name: str
|
||||
@@ -59,7 +62,9 @@ def test_tool_choice_bool(strict: bool | None) -> None: # noqa: FBT001
|
||||
|
||||
async def test_astream() -> None:
|
||||
"""Test streaming tokens from ChatFireworks."""
|
||||
llm = ChatFireworks(model="accounts/fireworks/models/kimi-k2p6")
|
||||
llm = ChatFireworks(
|
||||
model="accounts/fireworks/models/kimi-k2p6", rate_limiter=rate_limiter
|
||||
)
|
||||
|
||||
full: BaseMessageChunk | None = None
|
||||
chunks_with_token_counts = 0
|
||||
@@ -96,7 +101,7 @@ async def test_astream() -> None:
|
||||
|
||||
async def test_abatch_tags() -> None:
|
||||
"""Test batch tokens from ChatFireworks."""
|
||||
llm = ChatFireworks(model=_MODEL)
|
||||
llm = ChatFireworks(model=_MODEL, rate_limiter=rate_limiter)
|
||||
|
||||
result = await llm.abatch(
|
||||
["I'm Pickle Rick", "I'm not Pickle Rick"], config={"tags": ["foo"]}
|
||||
@@ -107,7 +112,7 @@ async def test_abatch_tags() -> None:
|
||||
|
||||
async def test_ainvoke() -> None:
|
||||
"""Test invoke tokens from ChatFireworks."""
|
||||
llm = ChatFireworks(model=_MODEL)
|
||||
llm = ChatFireworks(model=_MODEL, rate_limiter=rate_limiter)
|
||||
|
||||
result = await llm.ainvoke("I'm Pickle Rick", config={"tags": ["foo"]})
|
||||
assert isinstance(result.content, str)
|
||||
@@ -115,7 +120,7 @@ async def test_ainvoke() -> None:
|
||||
|
||||
def test_invoke() -> None:
|
||||
"""Test invoke tokens from ChatFireworks."""
|
||||
llm = ChatFireworks(model=_MODEL)
|
||||
llm = ChatFireworks(model=_MODEL, rate_limiter=rate_limiter)
|
||||
|
||||
result = llm.invoke("I'm Pickle Rick", config={"tags": ["foo"]})
|
||||
assert isinstance(result.content, str)
|
||||
@@ -157,7 +162,9 @@ def _get_joke_class(
|
||||
|
||||
@pytest.mark.parametrize("schema_type", ["pydantic", "typeddict", "json_schema"])
|
||||
def test_structured_output_json_schema(schema_type: str) -> None:
|
||||
llm = ChatFireworks(model="accounts/fireworks/models/kimi-k2p6")
|
||||
llm = ChatFireworks(
|
||||
model="accounts/fireworks/models/kimi-k2p6", rate_limiter=rate_limiter
|
||||
)
|
||||
schema, validation_function = _get_joke_class(schema_type) # type: ignore[arg-type]
|
||||
chat = llm.with_structured_output(schema, method="json_schema")
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ from langchain_tests.integration_tests import ( # type: ignore[import-not-found
|
||||
)
|
||||
|
||||
from langchain_fireworks import ChatFireworks
|
||||
from tests.integration_tests._rate_limiter import rate_limiter
|
||||
|
||||
|
||||
class TestFireworksStandard(ChatModelIntegrationTests):
|
||||
@@ -20,6 +21,7 @@ class TestFireworksStandard(ChatModelIntegrationTests):
|
||||
return {
|
||||
"model": "accounts/fireworks/models/kimi-k2p6",
|
||||
"temperature": 0,
|
||||
"rate_limiter": rate_limiter,
|
||||
}
|
||||
|
||||
@pytest.mark.xfail(reason="Not yet implemented.")
|
||||
|
||||
Reference in New Issue
Block a user