fix: abstract graph local model

ScrapeGraphAI · Aug 26, 2024 · 04128e7 · 04128e7
1 parent c1ce9c6
commit 04128e7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py
@@ -170,8 +170,8 @@ def handle_model(model_name, provider, token_key, default_token=8192):
             elif "ollama" in llm_params["model"]:
                 model_name = llm_params["model"].split("ollama/")[-1]
                 token_key = model_name if "model_tokens" not in llm_params else None
-                explicit_model_tokens = 8192 if "model_tokens" not in llm_params else llm_params["model_tokens"]
-                return handle_model(model_name, "ollama", token_key, explicit_model_tokens)
+                model_tokens = 8192 if "model_tokens" not in llm_params else llm_params["model_tokens"]
+                return handle_model(model_name, "ollama", token_key, model_tokens)
 
             elif "claude-3-" in llm_params["model"]:
                 return handle_model(llm_params["model"], "anthropic", "claude3")

diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py
@@ -74,6 +74,7 @@
             "llama3.1:70b": 128000,
             "lama3.1:405b": 128000,
             "scrapegraph": 8192,
+            "mistral": 8192,
             "llava": 4096,
             "mixtral:8x22b-instruct": 65536,
             "mistral-openorca": 32000,