diff --git a/llama_stack/providers/tests/agents/test_persistence.py b/llama_stack/providers/tests/agents/test_persistence.py
index e6b1470ef..bb6e280cb 100644
--- a/llama_stack/providers/tests/agents/test_persistence.py
+++ b/llama_stack/providers/tests/agents/test_persistence.py
@@ -9,8 +9,7 @@ import pytest
 from llama_stack.apis.agents import AgentConfig, Turn
 from llama_stack.apis.inference import SamplingParams, UserMessage
 from llama_stack.providers.datatypes import Api
-from llama_stack.providers.utils.kvstore import kvstore_impl
-from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
+from llama_stack.providers.utils.kvstore import kvstore_impl, SqliteKVStoreConfig
 
 from .fixtures import pick_inference_model
 
diff --git a/llama_stack/providers/tests/post_training/test_post_training.py b/llama_stack/providers/tests/post_training/test_post_training.py
index 0c58c1fa0..0645cd555 100644
--- a/llama_stack/providers/tests/post_training/test_post_training.py
+++ b/llama_stack/providers/tests/post_training/test_post_training.py
@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 import pytest
 
-from llama_stack.apis.common.job_types import JobStatus
+from llama_stack.apis.common.type_system import JobStatus
 from llama_stack.apis.post_training import (
     Checkpoint,
     DataConfig,