yujunjun/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2025-11-24 22:16:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Small fix and update tests (#9370)

Browse Source
This commit is contained in:
Sebastian Husch Lee 2025-05-12 22:02:26 +02:00 committed by GitHub
parent da5fc0f16c
commit 9f2c0679d4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 26 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
haystack/components/generators/chat/hugging_face_api.py
View File

@ -401,6 +401,7 @@ class HuggingFaceAPIChatGenerator:
generated_text = ""
first_chunk_time = None
meta: Dict[str, Any] = {}
for chunk in api_output:
# n is unused, so the API always returns only one choice
@ -412,8 +413,6 @@ class HuggingFaceAPIChatGenerator:
generated_text += text
finish_reason = choice.finish_reason
meta: Dict[str, Any] = {}
if finish_reason:
meta["finish_reason"] = finish_reason
@ -426,7 +425,6 @@ class HuggingFaceAPIChatGenerator:
meta.update(
{
"model": self._client.model,
"finish_reason": finish_reason,
"index": 0,
"usage": {"prompt_tokens": 0, "completion_tokens": 0}, # not available in streaming
"completion_start_time": first_chunk_time,
@ -434,7 +432,6 @@ class HuggingFaceAPIChatGenerator:
)
message = ChatMessage.from_assistant(text=generated_text, meta=meta)
return {"replies": [message]}
def _run_non_streaming(
@ -485,6 +482,7 @@ class HuggingFaceAPIChatGenerator:
generated_text = ""
first_chunk_time = None
meta: Dict[str, Any] = {}
async for chunk in api_output:
choice = chunk.choices[0]
@ -493,8 +491,6 @@ class HuggingFaceAPIChatGenerator:
generated_text += text
finish_reason = choice.finish_reason
meta: Dict[str, Any] = {}
if finish_reason:
meta["finish_reason"] = finish_reason
@ -507,7 +503,6 @@ class HuggingFaceAPIChatGenerator:
meta.update(
{
"model": self._async_client.model,
"finish_reason": finish_reason,
"index": 0,
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
"completion_start_time": first_chunk_time,

30
test/components/generators/chat/test_hugging_face_api.py
View File

@ -671,9 +671,15 @@ class TestHuggingFaceAPIChatGenerator:
assert isinstance(response["replies"], list)
assert len(response["replies"]) > 0
assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
assert "usage" in response["replies"][0].meta
assert "prompt_tokens" in response["replies"][0].meta["usage"]
assert "completion_tokens" in response["replies"][0].meta["usage"]
assert response["replies"][0].text is not None
meta = response["replies"][0].meta
assert "usage" in meta
assert "prompt_tokens" in meta["usage"]
assert meta["usage"]["prompt_tokens"] > 0
assert "completion_tokens" in meta["usage"]
assert meta["usage"]["completion_tokens"] > 0
assert meta["model"] == "microsoft/Phi-3.5-mini-instruct"
assert meta["finish_reason"] is not None
@pytest.mark.integration
@pytest.mark.slow
@ -701,13 +707,18 @@ class TestHuggingFaceAPIChatGenerator:
assert isinstance(response["replies"], list)
assert len(response["replies"]) > 0
assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
assert response["replies"][0].text is not None
response_meta = response["replies"][0].meta
assert "completion_start_time" in response_meta
assert datetime.fromisoformat(response_meta["completion_start_time"]) <= datetime.now()
assert "usage" in response_meta
assert "prompt_tokens" in response_meta["usage"]
assert response_meta["usage"]["prompt_tokens"] == 0
assert "completion_tokens" in response_meta["usage"]
assert response_meta["usage"]["completion_tokens"] == 0
assert response_meta["model"] == "microsoft/Phi-3.5-mini-instruct"
assert response_meta["finish_reason"] is not None
@pytest.mark.integration
@pytest.mark.slow
@ -926,9 +937,16 @@ class TestHuggingFaceAPIChatGenerator:
assert isinstance(response["replies"], list)
assert len(response["replies"]) > 0
assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
assert "usage" in response["replies"][0].meta
assert "prompt_tokens" in response["replies"][0].meta["usage"]
assert "completion_tokens" in response["replies"][0].meta["usage"]
assert response["replies"][0].text is not None
meta = response["replies"][0].meta
assert "usage" in meta
assert "prompt_tokens" in meta["usage"]
assert meta["usage"]["prompt_tokens"] > 0
assert "completion_tokens" in meta["usage"]
assert meta["usage"]["completion_tokens"] > 0
assert meta["model"] == "microsoft/Phi-3.5-mini-instruct"
assert meta["finish_reason"] is not None
finally:
await generator._async_client.close()