phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-29 15:23:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix return type and remove debug print

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-03 07:31:25 -08:00
parent 0e7e4bfb35
commit 09d5986ed8
1 changed files with 3 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
llama_stack/providers/adapters/inference/bedrock/bedrock.py
View file

@ -295,7 +295,9 @@ class BedrockInferenceAdapter(ModelRegistryHelper, Inference):
tool_prompt_format: Optional[ToolPromptFormat] = ToolPromptFormat.json,
stream: Optional[bool] = False,
logprobs: Optional[LogProbConfig] = None,
) -> AsyncGenerator:
) -> Union[
ChatCompletionResponse, AsyncIterator[ChatCompletionResponseStreamChunk]
]:
request = ChatCompletionRequest(
model=model,
messages=messages,
@ -316,7 +318,6 @@ class BedrockInferenceAdapter(ModelRegistryHelper, Inference):
async def _nonstream_chat_completion(
self, request: ChatCompletionRequest
) -> ChatCompletionResponse:
print("non-streaming chat completion")
params = self._get_params_for_chat_completion(request)
converse_api_res = self.client.converse(**params)
@ -332,7 +333,6 @@ class BedrockInferenceAdapter(ModelRegistryHelper, Inference):
async def _stream_chat_completion(
self, request: ChatCompletionRequest
) -> AsyncGenerator:
print("streaming chat completion")
params = self._get_params_for_chat_completion(request)
converse_stream_api_res = self.client.converse_stream(**params)
event_stream = converse_stream_api_res["stream"]