Skip to content

Commit

Permalink
Use EventStream from botocore
Browse files Browse the repository at this point in the history
  • Loading branch information
vblagoje committed Dec 10, 2024
1 parent 405865b commit 3990dfd
Showing 1 changed file with 3 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from typing import Any, Callable, Dict, List, Optional

from botocore.config import Config
from botocore.eventstream import EventStream
from botocore.exceptions import ClientError
from haystack import component, default_from_dict, default_to_dict
from haystack.dataclasses import ChatMessage, ChatRole, StreamingChunk
Expand Down Expand Up @@ -240,7 +241,7 @@ def run(
try:
if callback:
response = self.client.converse_stream(**params)
response_stream = response.get("stream")
response_stream: EventStream = response.get("stream")
if not response_stream:
msg = "No stream found in the response."
raise AmazonBedrockInferenceError(msg)
Expand Down Expand Up @@ -287,7 +288,7 @@ def extract_replies_from_response(self, response_body: Dict[str, Any]) -> List[C
return replies

def process_streaming_response(
self, response_stream, streaming_callback: Callable[[StreamingChunk], None]
self, response_stream: EventStream, streaming_callback: Callable[[StreamingChunk], None]
) -> List[ChatMessage]:
replies = []
current_content = ""
Expand Down

0 comments on commit 3990dfd

Please sign in to comment.