fix event logger

meta-llama · Sep 24, 2024 · 61523ce · 61523ce
1 parent f4cd276
commit 61523ce
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/src/llama_stack_client/lib/agents/event_logger.py b/src/llama_stack_client/lib/agents/event_logger.py
@@ -6,8 +6,9 @@
 
 from typing import List, Optional, Union
 
-from llama_stack.types import ToolResponseMessage
-from llama_stack.types.agents import AgentsTurnStreamChunk
+from llama_stack_client.types import ToolResponseMessage
+
+from llama_stack_client.types.agents import AgentsTurnStreamChunk
 from termcolor import cprint
 
 

diff --git a/src/llama_stack_client/lib/inference/event_logger.py b/src/llama_stack_client/lib/inference/event_logger.py
@@ -6,7 +6,10 @@
 
 from typing import List, Optional, Union
 
-from llama_stack.types import ChatCompletionStreamChunk, InferenceChatCompletionResponse
+from llama_stack_client.types import (
+    ChatCompletionStreamChunk,
+    InferenceChatCompletionResponse,
+)
 from termcolor import cprint