diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index 21f90460c..f0b596c30 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -78,6 +78,14 @@ def __init__( self.timeout = self.config.get("timeout", 480) self.graph = self._create_graph() + # Report the token window through the execution info as well. The + # warning emitted for an unknown model only reaches stderr, which is + # lost in batch and async contexts, so a caller had no way to tell a + # truncating 8192 fallback from a real limit by looking at the result. + # _create_llm is overridable and does not set model_token on every + # path, so fall back to the defaults BaseGraph already declares. + self.graph.model_token = getattr(self, "model_token", None) + self.graph.model_tokens_defaulted = self.model_tokens_defaulted self.final_state = None self.execution_info = None @@ -319,6 +327,13 @@ def get_execution_info(self): """ Returns the execution information of the graph. + The final "TOTAL RESULT" entry also carries the token window the run + actually used: ``effective_model_tokens`` and ``model_tokens_defaulted``, + the latter being True when no limit was known for the configured model + and the 8192 fallback was applied. A defaulted window chunks long pages + and can change the answer without raising, so batch callers should check + this flag rather than rely on the warning logged to stderr. + Returns: dict: The execution information of the graph. """ diff --git a/scrapegraphai/graphs/base_graph.py b/scrapegraphai/graphs/base_graph.py index d0977ece8..eca32c1d4 100644 --- a/scrapegraphai/graphs/base_graph.py +++ b/scrapegraphai/graphs/base_graph.py @@ -69,6 +69,12 @@ def __init__( self.graph_name = graph_name self.initial_state = {} self.callback_manager = CustomLLMCallbackManager() + # Effective input-token window used to chunk documents, and whether it + # is the 8192 fallback rather than the model's real limit. AbstractGraph + # fills these in after building the graph; they are reported in the + # "TOTAL RESULT" entry of the execution info (see #1121). + self.model_token = None + self.model_tokens_defaulted = False if nodes[0].node_name != entry_point.node_name: warnings.warn( @@ -316,6 +322,8 @@ def _execute_standard(self, initial_state: dict) -> Tuple[dict, list]: "successful_requests": cb_total["successful_requests"], "total_cost_USD": cb_total["total_cost_USD"], "exec_time": total_exec_time, + "effective_model_tokens": self.model_token, + "model_tokens_defaulted": self.model_tokens_defaulted, } ) diff --git a/tests/graphs/abstract_graph_test.py b/tests/graphs/abstract_graph_test.py index 6b57666a4..1f7488fb7 100644 --- a/tests/graphs/abstract_graph_test.py +++ b/tests/graphs/abstract_graph_test.py @@ -369,3 +369,79 @@ def test_get_execution_info(self): graph.execution_info = dummy_info info = graph.get_execution_info() assert info == dummy_info + + +class _StubNode: + """Minimal node used to drive BaseGraph without any model or network call.""" + + def __init__(self, node_name="Stub"): + self.node_name = node_name + self.node_type = "node" + self.node_config = {} + + def execute(self, state): + return state + + +def _run_stub_graph(model_token, model_tokens_defaulted): + """Executes a one-node graph and returns its "TOTAL RESULT" entry.""" + stub = _StubNode() + graph = BaseGraph(nodes=[stub], edges=[], entry_point=stub) + graph.model_token = model_token + graph.model_tokens_defaulted = model_tokens_defaulted + + with patch("scrapegraphai.graphs.base_graph.log_graph_execution"): + _, exec_info = graph.execute({}) + + return next(entry for entry in exec_info if entry["node_name"] == "TOTAL RESULT") + + +def test_execution_info_reports_defaulted_token_window(): + """The 8192 fallback must be visible in the returned execution info. + + Reported in #1121: the warning for an unknown model goes to stderr, so it + is lost in batch, worker and async contexts. A caller that only has the + returned object could not tell a truncating 8192 window from a real limit. + """ + total = _run_stub_graph(8192, True) + + assert total["effective_model_tokens"] == 8192 + assert total["model_tokens_defaulted"] is True + + +def test_execution_info_reports_known_token_window(): + """A model with a known limit is reported without the defaulted flag.""" + total = _run_stub_graph(1000000, False) + + assert total["effective_model_tokens"] == 1000000 + assert total["model_tokens_defaulted"] is False + + +def test_abstract_graph_propagates_defaulted_token_window(monkeypatch): + """AbstractGraph hands the token window to the graph it builds.""" + from scrapegraphai.graphs import abstract_graph + + monkeypatch.setattr( + abstract_graph, "models_tokens", {"openai": {"gpt-3.5-turbo": 4096}} + ) + llm_config = {"model": "openai/not-known-model", "openai_api_key": "test"} + with patch.object(TestGraph, "_create_graph", return_value=Mock(nodes=[])): + graph = TestGraph("Test prompt", {"llm": llm_config}) + + assert graph.graph.model_token == 8192 + assert graph.graph.model_tokens_defaulted is True + + +def test_abstract_graph_propagates_known_token_window(monkeypatch): + """A known model is propagated with the defaulted flag left off.""" + from scrapegraphai.graphs import abstract_graph + + monkeypatch.setattr( + abstract_graph, "models_tokens", {"openai": {"gpt-3.5-turbo": 4096}} + ) + llm_config = {"model": "openai/gpt-3.5-turbo", "openai_api_key": "test"} + with patch.object(TestGraph, "_create_graph", return_value=Mock(nodes=[])): + graph = TestGraph("Test prompt", {"llm": llm_config}) + + assert graph.graph.model_token == 4096 + assert graph.graph.model_tokens_defaulted is False