CodeCow/Scrapegraph-ai
1
0
Fork 0
mirror of https://github.com/VinciGit00/Scrapegraph-ai.git synced 2026-07-04 21:00:36 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

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

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
scrapegraphai/graphs/abstract_graph.py
View File

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

17
scrapegraphai/models/openai.py
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)