Merge branch 'pre/beta' of https://github.com/VinciGit00/Scrapegraph-ai into pre/beta

Author: VinciGit00

CHANGELOG.md

@@ -1,3 +1,17 @@
+## [0.9.0-beta.7](https://github.com/VinciGit00/Scrapegraph-ai/compare/v0.9.0-beta.6...v0.9.0-beta.7) (2024-05-06)
+
+
+### Bug Fixes
+
+* **llm:** fixed gemini api_key ([fd01b73](https://github.com/VinciGit00/Scrapegraph-ai/commit/fd01b73b71b515206cfdf51c1d52136293494389))
+
+## [0.9.0-beta.6](https://github.com/VinciGit00/Scrapegraph-ai/compare/v0.9.0-beta.5...v0.9.0-beta.6) (2024-05-06)
+
+
+### Features
+
+* Fix bug for gemini case when embeddings config not passed ([726de28](https://github.com/VinciGit00/Scrapegraph-ai/commit/726de288982700dab8ab9f22af8e26f01c6198a7))
+
 ## [0.9.0-beta.5](https://github.com/VinciGit00/Scrapegraph-ai/compare/v0.9.0-beta.4...v0.9.0-beta.5) (2024-05-06)
 
 

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "scrapegraphai"
 
-version = "0.9.0b5"
+version = "0.9.0b7"
 
 description = "A web scraping library based on LangChain which uses LLM and direct graph logic to create scraping pipelines."
 authors = [

scrapegraphai/graphs/abstract_graph.py

@@ -46,7 +46,7 @@ def __init__(self, prompt: str, config: dict, source: Optional[str] = None):
         self.source = source
         self.config = config
         self.llm_model = self._create_llm(config["llm"], chat=True)
-        self.embedder_model = self._create_default_embedder(
+        self.embedder_model = self._create_default_embedder(llm_config=config["llm"]
         ) if "embeddings" not in config else self._create_embedder(
             config["embeddings"])
 
@@ -91,6 +91,13 @@ def _set_model_token(self, llm):
                     self.model_token = models_tokens['mistral'][llm.repo_id]
                 except KeyError:
                     raise KeyError("Model not supported")
+                
+        elif 'Google' in str(type(llm)):
+            try:
+                if 'gemini' in llm.model:
+                    self.model_token = models_tokens['gemini'][llm.model]
+            except KeyError:
+                raise KeyError("Model not supported")
 
     def _create_llm(self, llm_config: dict, chat=False) -> object:
         """
@@ -197,7 +204,7 @@ def _create_llm(self, llm_config: dict, chat=False) -> object:
             raise ValueError(
                 "Model provided by the configuration not supported")
 
-    def _create_default_embedder(self) -> object:
+    def _create_default_embedder(self, llm_config=None) -> object:
         """
         Create an embedding model instance based on the chosen llm model.
 
@@ -207,6 +214,8 @@ def _create_default_embedder(self) -> object:
         Raises:
             ValueError: If the model is not supported.
         """
+        if isinstance(self.llm_model, Gemini):
+            return GoogleGenerativeAIEmbeddings(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)
         elif isinstance(self.llm_model, AzureOpenAIEmbeddings):
@@ -241,7 +250,6 @@ def _create_embedder(self, embedder_config: dict) -> object:
         Raises:
             KeyError: If the model is not supported.
         """
-
         if 'model_instance' in embedder_config:
             return embedder_config['model_instance']
         # Instantiate the embedding model based on the model name

scrapegraphai/helpers/models_tokens.py

@@ -26,6 +26,7 @@
     },
     "gemini": {
         "gemini-pro": 128000,
+        "models/embedding-001": 2048
     },
 
     "ollama": {

scrapegraphai/models/gemini.py

@@ -15,4 +15,6 @@ class Gemini(ChatGoogleGenerativeAI):
     """
 
     def __init__(self, llm_config: dict):
+        # replace "api_key" to "google_api_key"
+        llm_config["google_api_key"] = llm_config.pop("api_key", None)
         super().__init__(**llm_config)