fireworks[patch]: remove custom async and stream implementations (#18363)

libs/partners/fireworks/tests/integration_tests/test_llms.py

@@ -39,3 +39,64 @@ async def test_fireworks_acall() -> None:
     output_text = output.generations[0][0].text
     assert isinstance(output_text, str)
     assert output_text.count("bar") <= 1
+
+
+def test_stream() -> None:
+    """Test streaming tokens from OpenAI."""
+    llm = Fireworks(model="accounts/fireworks/models/mixtral-8x7b-instruct")
+
+    for token in llm.stream("I'm Pickle Rick"):
+        assert isinstance(token, str)
+
+
+async def test_astream() -> None:
+    """Test streaming tokens from OpenAI."""
+    llm = Fireworks(model="accounts/fireworks/models/mixtral-8x7b-instruct")
+
+    async for token in llm.astream("I'm Pickle Rick"):
+        assert isinstance(token, str)
+
+
+async def test_abatch() -> None:
+    """Test streaming tokens from Fireworks."""
+    llm = Fireworks(model="accounts/fireworks/models/mixtral-8x7b-instruct")
+
+    result = await llm.abatch(["I'm Pickle Rick", "I'm not Pickle Rick"])
+    for token in result:
+        assert isinstance(token, str)
+
+
+async def test_abatch_tags() -> None:
+    """Test batch tokens from Fireworks."""
+    llm = Fireworks(model="accounts/fireworks/models/mixtral-8x7b-instruct")
+
+    result = await llm.abatch(
+        ["I'm Pickle Rick", "I'm not Pickle Rick"], config={"tags": ["foo"]}
+    )
+    for token in result:
+        assert isinstance(token, str)
+
+
+def test_batch() -> None:
+    """Test batch tokens from Fireworks."""
+    llm = Fireworks(model="accounts/fireworks/models/mixtral-8x7b-instruct")
+
+    result = llm.batch(["I'm Pickle Rick", "I'm not Pickle Rick"])
+    for token in result:
+        assert isinstance(token, str)
+
+
+async def test_ainvoke() -> None:
+    """Test invoke tokens from Fireworks."""
+    llm = Fireworks(model="accounts/fireworks/models/mixtral-8x7b-instruct")
+
+    result = await llm.ainvoke("I'm Pickle Rick", config={"tags": ["foo"]})
+    assert isinstance(result, str)
+
+
+def test_invoke() -> None:
+    """Test invoke tokens from Fireworks."""
+    llm = Fireworks(model="accounts/fireworks/models/mixtral-8x7b-instruct")
+
+    result = llm.invoke("I'm Pickle Rick", config=dict(tags=["foo"]))
+    assert isinstance(result, str)