diff --git a/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/__init__.py b/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/__init__.py
index 4a88963..3a4113a 100644
--- a/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/__init__.py
+++ b/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/__init__.py
@@ -1,6 +1,6 @@
 from typing import Any
 
-from llama_stack.distribution.datatypes import AccessRule, Api
+from llama_stack.apis.datatypes import AccessRule, Api
 
 from .config import LightspeedAgentsImplConfig
 
diff --git a/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/agent_instance.py b/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/agent_instance.py
index d2dea4d..165bf20 100644
--- a/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/agent_instance.py
+++ b/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/agent_instance.py
@@ -3,7 +3,7 @@
 from collections.abc import AsyncGenerator
 
 from llama_stack.apis.agents import AgentConfig, AgentTurnCreateRequest, StepType
-from llama_stack.distribution.datatypes import AccessRule
+from llama_stack.apis.datatypes import AccessRule
 from llama_stack.log import get_logger
 from llama_stack.providers.inline.agents.meta_reference.agent_instance import ChatAgent
 from llama_stack.providers.utils.telemetry import tracing
diff --git a/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/agents.py b/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/agents.py
index 6079663..c2567c5 100644
--- a/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/agents.py
+++ b/lightspeed_stack_providers/providers/inline/agents/lightspeed_inline_agent/agents.py
@@ -9,7 +9,7 @@
 from llama_stack.apis.safety import Safety
 from llama_stack.apis.tools import ToolGroups, ToolRuntime
 from llama_stack.apis.vector_io import VectorIO
-from llama_stack.distribution.datatypes import AccessRule
+from llama_stack.apis.datatypes import AccessRule
 
 from .config import LightspeedAgentsImplConfig
 
diff --git a/lightspeed_stack_providers/providers/inline/safety/lightspeed_question_validity/safety.py b/lightspeed_stack_providers/providers/inline/safety/lightspeed_question_validity/safety.py
index 56a9f4d..1fbb299 100644
--- a/lightspeed_stack_providers/providers/inline/safety/lightspeed_question_validity/safety.py
+++ b/lightspeed_stack_providers/providers/inline/safety/lightspeed_question_validity/safety.py
@@ -7,7 +7,7 @@
 )
 
 from llama_stack.apis.shields import Shield
-from llama_stack.distribution.datatypes import Api
+from llama_stack.apis.datatypes import Api
 from llama_stack.providers.datatypes import ShieldsProtocolPrivate
 from llama_stack.apis.safety import (
     SafetyViolation,
diff --git a/lightspeed_stack_providers/providers/inline/safety/lightspeed_redaction/lightspeed_redaction/redaction.py b/lightspeed_stack_providers/providers/inline/safety/lightspeed_redaction/lightspeed_redaction/redaction.py
index 960dfed..8964781 100644
--- a/lightspeed_stack_providers/providers/inline/safety/lightspeed_redaction/lightspeed_redaction/redaction.py
+++ b/lightspeed_stack_providers/providers/inline/safety/lightspeed_redaction/lightspeed_redaction/redaction.py
@@ -8,7 +8,7 @@
 )
 
 from llama_stack.apis.shields import Shield
-from llama_stack.distribution.datatypes import Api
+from llama_stack.apis.datatypes import Api
 from llama_stack.providers.datatypes import ShieldsProtocolPrivate
 from llama_stack.apis.safety import (
     RunShieldResponse,
diff --git a/lightspeed_stack_providers/providers/remote/tool_runtime/lightspeed/__init__.py b/lightspeed_stack_providers/providers/remote/tool_runtime/lightspeed/__init__.py
index ad66dbc..7b2f157 100644
--- a/lightspeed_stack_providers/providers/remote/tool_runtime/lightspeed/__init__.py
+++ b/lightspeed_stack_providers/providers/remote/tool_runtime/lightspeed/__init__.py
@@ -2,7 +2,7 @@
 
 from pydantic import BaseModel
 
-from llama_stack.distribution.datatypes import Api
+from llama_stack.apis.datatypes import Api
 
 
 from .config import LightspeedToolConfig