use BaseChatModel for type hint

vizro-ai/src/vizro_ai/chains/_llm_chain.py

@@ -8,8 +8,7 @@
 from langchain.prompts import PromptTemplate
 from langchain.schema import ChatGeneration, Generation
 from langchain.schema.messages import AIMessage
-
-from vizro_ai.chains._llm_models import LLM_MODELS
+from langchain_core.language_models.chat_models import BaseChatModel
 
 logger = logging.getLogger(__name__)
 
@@ -38,7 +37,7 @@ class FunctionCallChain(VizroBaseChain, ABC):
 
     def __init__(  # noqa: PLR0913
         self,
-        llm: LLM_MODELS,
+        llm: BaseChatModel,
         raw_prompt: str,
         partial_vars_map: Optional[Dict[Any, Any]] = None,
         llm_kwargs: Optional[Dict[str, Any]] = None,

vizro-ai/src/vizro_ai/chains/_llm_models.py

@@ -1,12 +1,10 @@
 from typing import Dict, Optional, Union
 
+from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_openai import ChatOpenAI
 
-# TODO add new wrappers in if new model support is added
-LLM_MODELS = Union[ChatOpenAI]
-
 # TODO constant of model inventory, can be converted to yaml and link to docs
-PREDEFINED_MODELS: Dict[str, Dict[str, Union[int, LLM_MODELS]]] = {
+PREDEFINED_MODELS: Dict[str, Dict[str, Union[int, BaseChatModel]]] = {
     "gpt-3.5-turbo-0613": {
         "max_tokens": 4096,
         "wrapper": ChatOpenAI,
@@ -41,7 +39,7 @@
 DEFAULT_TEMPERATURE = 0
 
 
-def _get_llm_model(model: Optional[Union[ChatOpenAI, str]] = None) -> LLM_MODELS:
+def _get_llm_model(model: Optional[Union[ChatOpenAI, str]] = None) -> BaseChatModel:
     """Fetches and initializes an instance of the LLM.
 
     Args:

vizro-ai/src/vizro_ai/components/_base.py

@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
 from typing import Dict, Tuple, Union
 
+from langchain_core.language_models.chat_models import BaseChatModel
+
 from vizro_ai.chains import FunctionCallChain
-from vizro_ai.chains._llm_models import LLM_MODELS
 
 
 class VizroAiComponentBase(ABC):
@@ -22,7 +23,7 @@ class VizroAiComponentBase(ABC):
 
     prompt: str = "default prompt place holder"
 
-    def __init__(self, llm: LLM_MODELS):
+    def __init__(self, llm: BaseChatModel):
         """Initialize Vizro-AI base component.
 
         Args:

vizro-ai/src/vizro_ai/components/chart_selection.py

@@ -9,8 +9,9 @@
 except ImportError:  # pragma: no cov
     from pydantic import BaseModel, Field
 
+from langchain_core.language_models.chat_models import BaseChatModel
+
 from vizro_ai.chains._chain_utils import _log_time
-from vizro_ai.chains._llm_models import LLM_MODELS
 from vizro_ai.components import VizroAiComponentBase
 from vizro_ai.schema_manager import SchemaManager
 
@@ -48,7 +49,7 @@ class GetChartSelection(VizroAiComponentBase):
 
     prompt: str = chart_type_prompt
 
-    def __init__(self, llm: LLM_MODELS):
+    def __init__(self, llm: BaseChatModel):
         """Initialization of Chart Selection components.
 
         Args:

vizro-ai/src/vizro_ai/components/code_validation.py

@@ -8,8 +8,9 @@
 except ImportError:  # pragma: no cov
     from pydantic import BaseModel, Field
 
+from langchain_core.language_models.chat_models import BaseChatModel
+
 from vizro_ai.chains._chain_utils import _log_time
-from vizro_ai.chains._llm_models import LLM_MODELS
 from vizro_ai.components import VizroAiComponentBase
 from vizro_ai.schema_manager import SchemaManager
 
@@ -44,7 +45,7 @@ class GetDebugger(VizroAiComponentBase):
 
     prompt: str = debugging_prompt
 
-    def __init__(self, llm: LLM_MODELS):
+    def __init__(self, llm: BaseChatModel):
         """Initialization of Chart Selection components.
 
         Args:

vizro-ai/src/vizro_ai/components/custom_chart_wrap.py

@@ -8,8 +8,9 @@
 except ImportError:  # pragma: no cov
     from pydantic import BaseModel, Field
 
+from langchain_core.language_models.chat_models import BaseChatModel
+
 from vizro_ai.chains._chain_utils import _log_time
-from vizro_ai.chains._llm_models import LLM_MODELS
 from vizro_ai.components import VizroAiComponentBase
 from vizro_ai.schema_manager import SchemaManager
 
@@ -48,7 +49,7 @@ class GetCustomChart(VizroAiComponentBase):
 
     prompt: str = custom_chart_prompt
 
-    def __init__(self, llm: LLM_MODELS):
+    def __init__(self, llm: BaseChatModel):
         """Initialization of custom chart components.
 
         Args:

vizro-ai/src/vizro_ai/components/dataframe_craft.py

@@ -11,8 +11,9 @@
 except ImportError:  # pragma: no cov
     from pydantic import BaseModel, Field
 
+from langchain_core.language_models.chat_models import BaseChatModel
+
 from vizro_ai.chains._chain_utils import _log_time
-from vizro_ai.chains._llm_models import LLM_MODELS
 from vizro_ai.components import VizroAiComponentBase
 from vizro_ai.schema_manager import SchemaManager
 
@@ -51,7 +52,7 @@ class GetDataFrameCraft(VizroAiComponentBase):
 
     prompt: str = dataframe_prompt
 
-    def __init__(self, llm: LLM_MODELS):
+    def __init__(self, llm: BaseChatModel):
         """Initialization of dataframe crafting components.
 
         Args:

vizro-ai/src/vizro_ai/components/explanation.py

@@ -7,8 +7,9 @@
 except ImportError:  # pragma: no cov
     from pydantic import BaseModel, Field
 
+from langchain_core.language_models.chat_models import BaseChatModel
+
 from vizro_ai.chains._chain_utils import _log_time
-from vizro_ai.chains._llm_models import LLM_MODELS
 from vizro_ai.components import VizroAiComponentBase
 from vizro_ai.schema_manager import SchemaManager
 
@@ -44,7 +45,7 @@ class GetCodeExplanation(VizroAiComponentBase):
 
     prompt: str = code_explanation_prompt
 
-    def __init__(self, llm: LLM_MODELS):
+    def __init__(self, llm: BaseChatModel):
         """Initialization of Code Explanation components.
 
         Args:

vizro-ai/src/vizro_ai/components/visual_code.py

@@ -7,8 +7,9 @@
 except ImportError:  # pragma: no cov
     from pydantic import BaseModel, Field
 
+from langchain_core.language_models.chat_models import BaseChatModel
+
 from vizro_ai.chains._chain_utils import _log_time
-from vizro_ai.chains._llm_models import LLM_MODELS
 from vizro_ai.components import VizroAiComponentBase
 from vizro_ai.schema_manager import SchemaManager
 
@@ -43,7 +44,7 @@ class GetVisualCode(VizroAiComponentBase):
 
     prompt: str = visual_code_prompt
 
-    def __init__(self, llm: LLM_MODELS):
+    def __init__(self, llm: BaseChatModel):
         """Initialization of Chart Selection components.
 
         Args:

vizro-ai/src/vizro_ai/task_pipeline/_pipeline.py

@@ -2,13 +2,13 @@
 
 from typing import Any, Dict, List, Optional
 
-from vizro_ai.chains._llm_models import LLM_MODELS
+from langchain_core.language_models.chat_models import BaseChatModel
 
 
 class Pipeline:
     """A pipeline to manage the flow of tasks in a sequence."""
 
-    def __init__(self, llm: LLM_MODELS):
+    def __init__(self, llm: BaseChatModel):
         """Initialize the Pipeline.
 
         Args:

vizro-ai/src/vizro_ai/task_pipeline/_pipeline_manager.py

@@ -1,14 +1,14 @@
 """Pipeline Manager."""
 
-from vizro_ai.chains._llm_models import LLM_MODELS
+from langchain_core.language_models.chat_models import BaseChatModel
 from vizro_ai.components import GetChartSelection, GetCustomChart, GetDataFrameCraft, GetVisualCode
 from vizro_ai.task_pipeline._pipeline import Pipeline
 
 
 class PipelineManager:
     """Task pipeline manager."""
 
-    def __init__(self, llm: LLM_MODELS = None):
+    def __init__(self, llm: BaseChatModel = None):
         """Initialize the Pipeline Manager.
 
         Args: