fix: bug for abstract graph

VinciGit00 · VinciGit00 · commit cf7388345172 · 2024-08-27T18:05:34.000+02:00
diff --git a/examples/local_models/smart_scraper_ollama.py b/examples/local_models/smart_scraper_ollama.py
@@ -9,7 +9,7 @@
 
 graph_config = {
     "llm": {
-        "model": "ollama/mistral",
+        "model": "ollama/llama3.1",
         "temperature": 0,
         "format": "json",  # Ollama needs the format to be specified explicitly
         # "base_url": "http://localhost:11434", # set ollama URL arbitrarily
diff --git a/requirements-dev.lock b/requirements-dev.lock
@@ -130,7 +130,6 @@ graphviz==0.20.3
     # via burr
 greenlet==3.0.3
     # via playwright
-    # via sqlalchemy
 grpcio==1.65.4
     # via google-api-core
     # via grpcio-status
diff --git a/requirements.lock b/requirements.lock
@@ -83,7 +83,6 @@ googleapis-common-protos==1.63.2
     # via grpcio-status
 greenlet==3.0.3
     # via playwright
-    # via sqlalchemy
 grpcio==1.65.1
     # via google-api-core
     # via grpcio-status
diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py
@@ -139,7 +139,7 @@ def _create_llm(self, llm_config: dict) -> object:
             raise ValueError(f"Provider {llm_params['model_provider']} is not supported. If possible, try to use a model instance instead.")
         
         try:
-            self.model_token = models_tokens[llm_params["model"]][llm_params["model"]]
+            self.model_token = models_tokens[llm_params["model_provider"]].get(llm_params["model"][0])
         except KeyError:
             print("Model not found, using default token size (8192)")
             self.model_token = 8192
diff --git a/tests/graphs/scrape_json_ollama_test.py b/tests/graphs/scrape_json_ollama_test.py
