Merge branch '423-add-vertex-ai-integration' into support

Author: f-aguzzi

pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "langchain==0.1.15",
     "langchain-openai==0.1.6",
     "langchain-google-genai==1.0.3",
+    "langchain-google-vertexai==1.0.6",
     "langchain-groq==0.1.3",
     "langchain-aws==0.1.3",
     "langchain-anthropic==0.1.11",

scrapegraphai/graphs/abstract_graph.py

@@ -10,10 +10,10 @@
 from langchain_aws import BedrockEmbeddings
 from langchain_community.embeddings import HuggingFaceHubEmbeddings, OllamaEmbeddings
 from langchain_google_genai import GoogleGenerativeAIEmbeddings
+from langchain_google_vertexai import VertexAIEmbeddings
 from langchain_google_genai.embeddings import GoogleGenerativeAIEmbeddings
 from langchain_fireworks import FireworksEmbeddings
 from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
-
 from ..helpers import models_tokens
 from ..models import (
     Anthropic,
@@ -25,7 +25,8 @@
     Ollama,
     OpenAI,
     OneApi,
-    Fireworks
+    Fireworks,
+    VertexAI
 )
 from ..models.ernie import Ernie
 from ..utils.logging import set_verbosity_debug, set_verbosity_warning, set_verbosity_info
@@ -73,7 +74,7 @@ def __init__(self, prompt: str, config: dict,
         self.config = config
         self.schema = schema
         self.llm_model = self._create_llm(config["llm"], chat=True)
-        self.embedder_model = self._create_default_embedder(llm_config=config["llm"]                                                            ) if "embeddings" not in config else self._create_embedder(
+        self.embedder_model = self._create_default_embedder(llm_config=config["llm"]) if "embeddings" not in config else self._create_embedder(
             config["embeddings"])
         self.verbose = False if config is None else config.get(
             "verbose", False)
@@ -179,7 +180,6 @@ def _create_llm(self, llm_config: dict, chat=False) -> object:
             except KeyError as exc:
                 raise KeyError("Model not supported") from exc
             return AzureOpenAI(llm_params)
-
         elif "gemini" in llm_params["model"]:
             llm_params["model"] = llm_params["model"].split("/")[-1]
             try:
@@ -194,6 +194,12 @@ def _create_llm(self, llm_config: dict, chat=False) -> object:
             except KeyError as exc:
                 raise KeyError("Model not supported") from exc
             return Anthropic(llm_params)
+        elif llm_params["model"].startswith("vertexai"):
+            try:
+                self.model_token = models_tokens["vertexai"][llm_params["model"]]
+            except KeyError as exc:
+                raise KeyError("Model not supported") from exc
+            return VertexAI(llm_params)
         elif "ollama" in llm_params["model"]:
             llm_params["model"] = llm_params["model"].split("ollama/")[-1]
 
@@ -287,9 +293,12 @@ def _create_default_embedder(self, llm_config=None) -> object:
                 google_api_key=llm_config["api_key"], model="models/embedding-001"
             )
         if isinstance(self.llm_model, OpenAI):
-            return OpenAIEmbeddings(api_key=self.llm_model.openai_api_key, base_url=self.llm_model.openai_api_base)
+            return OpenAIEmbeddings(api_key=self.llm_model.openai_api_key, 
+                                    base_url=self.llm_model.openai_api_base)
         elif isinstance(self.llm_model, DeepSeek):
             return OpenAIEmbeddings(api_key=self.llm_model.openai_api_key)
+        elif isinstance(self.llm_model, VertexAI):
+            return VertexAIEmbeddings()
         elif isinstance(self.llm_model, AzureOpenAIEmbeddings):
             return self.llm_model
         elif isinstance(self.llm_model, AzureOpenAI):

scrapegraphai/helpers/models_tokens.py

@@ -80,6 +80,11 @@
         "claude3": 200000,
         "claude3.5": 200000
     },
+    "vertexai": {
+        "gemini-1.5-flash": 128000,
+        "gemini-1.5-pro": 128000,
+        "gemini-1.0-pro": 128000
+    },
     "bedrock": {
         "anthropic.claude-3-haiku-20240307-v1:0": 200000,
         "anthropic.claude-3-sonnet-20240229-v1:0": 200000,

scrapegraphai/models/__init__.py

@@ -15,3 +15,4 @@
 from .deepseek import DeepSeek
 from .oneapi import OneApi
 from .fireworks import Fireworks
+from .vertex import VertexAI

scrapegraphai/models/vertex.py

@@ -0,0 +1,16 @@
+""" 
+VertexAI Module
+"""
+from langchain_google_vertexai import ChatVertexAI
+
+class VertexAI(ChatVertexAI):
+    """
+    A wrapper for the ChatVertexAI 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)