refactor(OpenAI): integrate new LangChain chat init
Some checks are pending
/ build (push) Waiting to run

This commit is contained in:
Federico Aguzzi 2024-07-29 11:26:20 +02:00
parent d177afb68b
commit 5e3eb6e43d
2 changed files with 3 additions and 20 deletions

View File

@ -8,6 +8,7 @@ import uuid
from pydantic import BaseModel
from langchain_community.chat_models import ChatOllama
from langchain_openai import ChatOpenAI
from langchain_aws import BedrockEmbeddings
from langchain_community.embeddings import HuggingFaceHubEmbeddings, OllamaEmbeddings
@ -21,7 +22,6 @@ from ..helpers import models_tokens
from ..models import (
Anthropic,
AzureOpenAI,
OpenAI,
Bedrock,
Gemini,
Groq,
@ -37,7 +37,7 @@ from langchain.chat_models import init_chat_model
from ..utils.logging import set_verbosity_debug, set_verbosity_warning, set_verbosity_info
from ..helpers import models_tokens
from ..models import AzureOpenAI, OpenAI, Bedrock, Gemini, Groq, HuggingFace, Anthropic, DeepSeek
from ..models import AzureOpenAI, Bedrock, Gemini, Groq, HuggingFace, Anthropic, DeepSeek
class AbstractGraph(ABC):
@ -311,7 +311,7 @@ class AbstractGraph(ABC):
return GoogleGenerativeAIEmbeddings(
google_api_key=llm_config["api_key"], model="models/embedding-001"
)
if isinstance(self.llm_model, OpenAI):
if isinstance(self.llm_model, ChatOpenAI):
return OpenAIEmbeddings(api_key=self.llm_model.openai_api_key,
base_url=self.llm_model.openai_api_base)
elif isinstance(self.llm_model, DeepSeek):

View File

@ -1,17 +0,0 @@
"""
OpenAI Module
"""
from langchain_openai import ChatOpenAI
class OpenAI(ChatOpenAI):
"""
A wrapper for the ChatOpenAI class that provides default configuration
and could be extended with additional methods if needed.
Args:
llm_config (dict): Configuration parameters for the language model.
"""
def __init__(self, llm_config: dict):
super().__init__(**llm_config)