phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-04 20:14:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix non-streaming api in inference server

Browse source
This commit is contained in:
Hardik Shah 2024-07-30 14:25:50 -07:00
parent 404af06e02
commit cc98fbb058
3 changed files with 32 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

23
llama_toolchain/inference/event_logger.py
View file

@ -6,7 +6,10 @@
from termcolor import cprint
from llama_toolchain.inference.api import ChatCompletionResponseEventType
from llama_toolchain.inference.api import (
ChatCompletionResponseEventType,
ChatCompletionResponseStreamChunk
)
class LogEvent:
@ -25,12 +28,16 @@ class LogEvent:
class EventLogger:
async def log(self, event_generator, stream=True):
async def log(self, event_generator):
async for chunk in event_generator:
event = chunk.event
if event.event_type == ChatCompletionResponseEventType.start:
if isinstance(chunk, ChatCompletionResponseStreamChunk):
event = chunk.event
if event.event_type == ChatCompletionResponseEventType.start:
yield LogEvent("Assistant> ", color="cyan", end="")
elif event.event_type == ChatCompletionResponseEventType.progress:
yield LogEvent(event.delta, color="yellow", end="")
elif event.event_type == ChatCompletionResponseEventType.complete:
yield LogEvent("")
else:
yield LogEvent("Assistant> ", color="cyan", end="")
elif event.event_type == ChatCompletionResponseEventType.progress:
yield LogEvent(event.delta, color="yellow", end="")
elif event.event_type == ChatCompletionResponseEventType.complete:
yield LogEvent("")
yield LogEvent(chunk.completion_message.content, color="yellow")