File size: 1,094 Bytes
1037adc 67ab62d 1037adc ef7b2fa 1037adc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | from abc import ABC, abstractmethod
from typing import Optional
from langchain_core.embeddings import Embeddings
from langchain_community.chat_models.tongyi import BaseChatModel
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.chat_models.tongyi import ChatTongyi
from utils.config_handler import rag_conf
import os
class BaseModelFactory(ABC):
@abstractmethod
def generator(self) -> Optional[Embeddings | BaseChatModel]:
pass
class ChatModelFactory(BaseModelFactory):
def generator(self) -> Optional[Embeddings | BaseChatModel]:
return ChatTongyi(model=rag_conf["chat_model_name"],
api_key=rag_conf["api_key"])
class EmbeddingsFactory(BaseModelFactory):
def generator(self) -> Optional[Embeddings | BaseChatModel]:
return DashScopeEmbeddings(model=rag_conf["embedding_model_name"],
dashscope_api_key=os.environ.get("API_KEY"))
chat_model = ChatModelFactory().generator()
embed_model = EmbeddingsFactory().generator()
|