diff --git a/src/python/PythonSDK/foundationallm/langchain/common/foundationallm_tool_base.py b/src/python/PythonSDK/foundationallm/langchain/common/foundationallm_tool_base.py
index 43dd7d010..16033380b 100644
--- a/src/python/PythonSDK/foundationallm/langchain/common/foundationallm_tool_base.py
+++ b/src/python/PythonSDK/foundationallm/langchain/common/foundationallm_tool_base.py
@@ -3,6 +3,8 @@
 Description: FoundationaLLM base class for tools that uses the AgentTool model for its configuration.
 """
 
+from typing import List, Dict
+
 # Platform imports
 from azure.identity import DefaultAzureCredential
 from logging import Logger
@@ -23,7 +25,7 @@ class FoundationaLLMToolBase(BaseTool):
 
     response_format: str = 'content_and_artifact'
 
-    def __init__(self, tool_config: AgentTool, objects:dict, user_identity:UserIdentity, config: Configuration):
+    def __init__(self, tool_config: AgentTool, objects:Dict, user_identity:UserIdentity, config: Configuration):
         """ Initializes the FoundationaLLMToolBase class with the tool configuration. """
         super().__init__(
             name=tool_config.name,
diff --git a/src/python/PythonSDK/foundationallm/langchain/common/foundationallm_workflow_base.py b/src/python/PythonSDK/foundationallm/langchain/common/foundationallm_workflow_base.py
index c16fe40cc..6177bdd58 100644
--- a/src/python/PythonSDK/foundationallm/langchain/common/foundationallm_workflow_base.py
+++ b/src/python/PythonSDK/foundationallm/langchain/common/foundationallm_workflow_base.py
@@ -3,14 +3,19 @@
 Description: FoundationaLLM base class for tools that uses the agent workflow model for its configuration.
 """
 from abc import ABC, abstractmethod
+from typing import List
+
 from azure.identity import DefaultAzureCredential
+
 from langchain_core.messages import BaseMessage
-from typing import List
+
 from foundationallm.config import Configuration, UserIdentity
 from foundationallm.models.agents import AgentTool, ExternalAgentWorkflow
 from foundationallm.models.orchestration import CompletionResponse
 from foundationallm.telemetry import Telemetry
 
+from .foundationallm_tool_base import FoundationaLLMToolBase
+
 class FoundationaLLMWorkflowBase(ABC):
     """
     FoundationaLLM base class for workflows that uses the agent workflow model for its configuration.
@@ -18,7 +23,7 @@ class FoundationaLLMWorkflowBase(ABC):
     def __init__(self,
                  workflow_config: ExternalAgentWorkflow,
                  objects: dict,
-                 tools: List[AgentTool],
+                 tools: List[FoundationaLLMToolBase],
                  user_identity: UserIdentity,
                  config: Configuration):
         """
@@ -30,7 +35,7 @@ def __init__(self,
             The workflow assigned to the agent.
         objects : dict
             The exploded objects assigned from the agent.
-        tools : List[AgentTool]
+        tools : List[FoundationaLLMToolBase]
             The tools assigned to the agent.
         user_identity : UserIdentity
             The user identity of the user initiating the request.