From 35b994a8cd201bcde86038dcde6933946a4fe84f Mon Sep 17 00:00:00 2001 From: Jamie Beck Date: Mon, 26 Aug 2024 11:30:09 -0400 Subject: [PATCH] fix model_tokens not being used for ollama I am passing in the explicit model_tokens from user config as the default_token so it will correctly fallback to the users setting if the model is not found --- scrapegraphai/graphs/abstract_graph.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index ae1e90b2..e76a2ed7 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -174,8 +174,9 @@ class AbstractGraph(ABC): 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 llm_params["model_tokens"] - return handle_model(model_name, "ollama", token_key) + 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) elif "claude-3-" in llm_params["model"]: return handle_model(llm_params["model"], "anthropic", "claude3") @@ -271,4 +272,4 @@ class AbstractGraph(ABC): def run(self) -> str: """ Abstract method to execute the graph and return the result. - """ \ No newline at end of file + """