From ce6be37fbc1095afe4df6a2fc206923e477190e5 Mon Sep 17 00:00:00 2001 From: Marco Vinciguerra Date: Sun, 21 Jul 2024 18:37:15 +0200 Subject: [PATCH] fix: abstract_graph moel token bug --- scrapegraphai/graphs/abstract_graph.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index 7f8ec4ea..ca2b2e81 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -148,6 +148,10 @@ class AbstractGraph(ABC): # If model instance is passed directly instead of the model details if "model_instance" in llm_params: + try: + self.model_token = llm_params["model_tokens"] + except KeyError as exc: + raise KeyError("model_tokens not specified") from exc return llm_params["model_instance"] # Instantiate the language model based on the model name