fix example and type hints

Author: Tialo

dialog2graph/pipelines/d2g_llm/three_stages_llm.py

@@ -14,7 +14,7 @@
 from langchain.output_parsers import PydanticOutputParser, OutputFixingParser
 from langchain.schema import HumanMessage
 from langchain_openai import ChatOpenAI
-from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_huggingface.embeddings import HuggingFaceEmbeddings
 
 
 from dialog2graph import metrics

dialog2graph/pipelines/model_storage.py

@@ -5,14 +5,13 @@
 import yaml
 import re
 import dotenv
-from pydantic._internal._model_construction import ModelMetaclass
-from typing import Union, Dict
+from typing import Union, Dict, Type
 from pathlib import Path
 from pydantic import BaseModel, Field, model_validator
 
-from langchain_community.chat_models import ChatOpenAI
+from langchain_openai.chat_models import ChatOpenAI
 from langchain_core.language_models import BaseChatModel
-from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_huggingface.embeddings import HuggingFaceEmbeddings
 
 from dialog2graph.utils.logger import Logger
 
@@ -62,7 +61,7 @@ class StoredData(BaseModel):
     Attributes:
         key (str): Key for the stored model.
         config (dict): Configuration for the stored model.
-        model_type (Union[Literal["llm"], Literal["emb"]]): Type of the stored model, either "llm" (language model) or "emb" (embedding model).
+        model_type (Type[BaseChatModel | HuggingFaceEmbeddings]): Type of the stored model, for example ChatOpenAI, HuggingFaceEmbeddings.
         model (Union[HuggingFaceEmbeddings, BaseChatModel]): The actual model object, which can either be a HuggingFaceEmbeddings instance or a BaseChatModel instance.
 
     Methods:
@@ -76,7 +75,7 @@ class StoredData(BaseModel):
 
     key: str = Field(description="Key for the stored model")
     config: dict = Field(description="Configuration for the stored model")
-    model_type: ModelMetaclass = Field(description="Type of the stored model")
+    model_type: Type[BaseChatModel | HuggingFaceEmbeddings] = Field(description="Type of the stored model")
     model: Union[
         HuggingFaceEmbeddings,
         BaseChatModel,
@@ -140,7 +139,7 @@ def add(
         self,
         key: str,
         config: dict,
-        model_type: ModelMetaclass,
+        model_type: Type[BaseChatModel | HuggingFaceEmbeddings],
         overwright: bool = False,
     ):
         """
@@ -155,7 +154,7 @@ def add(
                 Args:
                     key (str): The unique identifier for the model configuration.
                     config (dict): The configuration dictionary for initializing the model.
-                    model_type (ModelMetaclass): The type name of the model to be added.
+                    model_type (Type[BaseChatModel | HuggingFaceEmbeddings]): The type of the model to be added.
                     overwright (bool): Whether to overwright model existing under same key
         .
                 Raises:

examples/pipeline_usage/models_config.yml

@@ -1,20 +1,20 @@
-d2g_llm_filling_llm:v1:
+my_filling_model:
   config:
-    name: o3-mini
-    temperature: 1
-  model_type: llm
-d2g_llm_grouping_llm:v1:
+    model_name: gpt-3.5-turbo
+  model_type: ChatOpenAI
+three_stages_llm_formatting_llm:v1:
   config:
-    name: gpt-4o-latest
+    model_name: gpt-4o-mini
     temperature: 0
-  model_type: llm
-d2g_llm_sim_model:v1:
+  model_type: ChatOpenAI
+three_stages_llm_grouping_llm:v1:
+  config:
+    model_name: chatgpt-4o-latest
+    temperature: 0
+  model_type: ChatOpenAI
+three_stages_llm_sim_model:v1:
   config:
     model_kwargs:
       device: cpu
-    model_name: cointegrated/LaBSE-en-ru
-  model_type: emb
-my_model:
-  config:
-    name: gpt-3.5-turbo
-  model_type: llm
+    model_name: BAAI/bge-m3
+  model_type: HuggingFaceEmbeddings