fix(nodes): expose model identity in node event inputs (#80)

Author: laipz8200

src/graphon/nodes/llm/llm_utils.py

@@ -67,6 +67,17 @@ def fetch_model_schema(*, model_instance: PreparedLLMProtocol) -> AIModelEntity:
     return model_schema
 
 
+def build_model_identity_inputs(
+    *,
+    model_instance: PreparedLLMProtocol,
+) -> dict[str, Any]:
+    """Expose the prepared model identity in node inputs."""
+    return {
+        "model_provider": model_instance.provider,
+        "model_name": model_instance.model_name,
+    }
+
+
 def fetch_files(variable_pool: VariablePool, selector: Sequence[str]) -> Sequence[File]:
     variable = variable_pool.get(selector)
     if variable is None:

src/graphon/nodes/llm/node.py

@@ -253,6 +253,9 @@ def _prepare_run_prompt(
             collected_context=collected_context,
         )
         model_instance = self._prepare_model_instance()
+        node_inputs.update(
+            llm_utils.build_model_identity_inputs(model_instance=model_instance),
+        )
         prompt_messages, stop = LLMNode.fetch_prompt_messages(
             sys_query=self._resolve_memory_query(),
             sys_files=files,
@@ -377,6 +380,9 @@ def _yield_run_completion(
                 )
             break
 
+        node_inputs.update(
+            llm_utils.build_model_identity_inputs(model_instance=self._model_instance),
+        )
         process_data.update(
             self._build_process_data(
                 prompt_messages=prompt_messages,

src/graphon/nodes/parameter_extractor/parameter_extractor_node.py

@@ -384,6 +384,7 @@ def _prepare_run_context(self) -> _ParameterExtractorRunContext:
             "files": [f.to_dict() for f in files],
             "parameters": jsonable_encoder(node_data.parameters),
             "instruction": jsonable_encoder(node_data.instruction),
+            **llm_utils.build_model_identity_inputs(model_instance=model_instance),
         }
         process_data = {
             "model_mode": node_data.model.mode,

src/graphon/nodes/question_classifier/question_classifier_node.py

@@ -275,14 +275,18 @@ def _prepare_run_context(self) -> _QuestionClassifierRunContext:
             model_instance=model_instance,
             files=files,
         )
+        inputs = {
+            "query": query,
+            **llm_utils.build_model_identity_inputs(model_instance=model_instance),
+        }
         rendered_classes = [
             class_.model_copy(
                 update={"name": variable_pool.convert_template(class_.name).text},
             )
             for class_ in node_data.classes
         ]
         return _QuestionClassifierRunContext(
-            inputs={"query": query},
+            inputs=inputs,
             model_instance=model_instance,
             prompt_messages=prompt_messages,
             stop=stop,

tests/nodes/llm/test_node.py

@@ -0,0 +1,78 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from graphon.model_runtime.entities.llm_entities import LLMUsage
+from graphon.node_events.node import ModelInvokeCompletedEvent, StreamCompletedEvent
+from graphon.nodes.llm import LLMNode, LLMNodeData
+from graphon.runtime.graph_runtime_state import GraphRuntimeState
+
+from ...helpers import build_graph_init_params, build_variable_pool
+
+
+def _build_llm_node() -> LLMNode:
+    return LLMNode(
+        node_id="llm",
+        config=LLMNodeData.model_validate({
+            "title": "LLM",
+            "model": {
+                "provider": "openai",
+                "name": "gpt-4o",
+                "mode": "chat",
+                "completion_params": {},
+            },
+            "prompt_template": [
+                {
+                    "role": "user",
+                    "text": "Hello",
+                }
+            ],
+            "context": {"enabled": False},
+        }),
+        graph_init_params=build_graph_init_params(
+            graph_config={"nodes": [], "edges": []}
+        ),
+        graph_runtime_state=GraphRuntimeState(
+            variable_pool=build_variable_pool(),
+            start_at=0.0,
+        ),
+        model_instance=MagicMock(
+            provider="openai",
+            model_name="gpt-4o",
+            parameters={},
+            stop=(),
+        ),
+        llm_file_saver=MagicMock(),
+        prompt_message_serializer=MagicMock(),
+    )
+
+
+def test_run_emits_model_identity_in_node_result_inputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    node = _build_llm_node()
+
+    monkeypatch.setattr(node, "_fetch_inputs", lambda **_: {})
+    monkeypatch.setattr(node, "_fetch_jinja_inputs", lambda **_: {})
+    monkeypatch.setattr(node, "_collect_run_context", lambda **_: iter(()))
+    monkeypatch.setattr(
+        LLMNode, "fetch_prompt_messages", staticmethod(lambda **_: ([], None))
+    )
+    monkeypatch.setattr(
+        "graphon.nodes.llm.node.LLMNode.invoke_llm",
+        lambda **_: iter([
+            ModelInvokeCompletedEvent(
+                text="Hello back",
+                usage=LLMUsage.empty_usage(),
+                finish_reason="stop",
+            ),
+        ]),
+    )
+
+    events = list(node._run())  # noqa: SLF001
+    completed_event = next(
+        event for event in events if isinstance(event, StreamCompletedEvent)
+    )
+
+    assert completed_event.node_run_result.inputs["model_provider"] == "openai"
+    assert completed_event.node_run_result.inputs["model_name"] == "gpt-4o"

tests/nodes/parameter_extractor/test_prompts.py

@@ -1,11 +1,13 @@
 import time
+from types import SimpleNamespace
 from typing import Any, cast
 from unittest.mock import Mock
 
 import pytest
 
-from graphon.model_runtime.entities.llm_entities import LLMMode
+from graphon.model_runtime.entities.llm_entities import LLMMode, LLMUsage
 from graphon.model_runtime.entities.message_entities import PromptMessageRole
+from graphon.nodes.llm import llm_utils
 from graphon.nodes.llm.entities import ModelConfig
 from graphon.nodes.parameter_extractor import parameter_extractor_node
 from graphon.nodes.parameter_extractor.entities import (
@@ -47,7 +49,12 @@ def _build_parameter_extractor_node() -> tuple[ParameterExtractorNode, VariableP
         start_at=time.perf_counter(),
     )
     init_params = build_graph_init_params(graph_config={"nodes": [], "edges": []})
-    model_instance = Mock()
+    model_instance = Mock(
+        provider="test",
+        model_name="test-model",
+        parameters={},
+        stop=(),
+    )
     prompt_message_serializer = Mock()
     node = cast(
         ParameterExtractorNode,
@@ -151,6 +158,65 @@ def test_function_calling_prompt_template_renders_system_message() -> None:
     assert prompt_messages[1].text == "Extract the location from this request."
 
 
+def test_prepare_run_context_exposes_model_identity_in_inputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    node, variable_pool = _build_parameter_extractor_node()
+    variable_pool.add(("start", "query"), "weather in sf")
+
+    monkeypatch.setattr(
+        llm_utils,
+        "resolve_completion_params_variables",
+        lambda parameters, _: parameters,
+    )
+    monkeypatch.setattr(
+        node,
+        "_fetch_llm_model_schema",
+        lambda **_: SimpleNamespace(features=[]),
+    )
+    monkeypatch.setattr(node, "_build_run_prompt", lambda **_: ([], []))
+
+    run_context = node._prepare_run_context()  # noqa: SLF001
+
+    assert run_context.inputs["query"] == "weather in sf"
+    assert run_context.inputs["model_provider"] == "test"
+    assert run_context.inputs["model_name"] == "test-model"
+
+
+def test_parameter_extractor_run_emits_model_identity_in_inputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    node, variable_pool = _build_parameter_extractor_node()
+    variable_pool.add(("start", "query"), "weather in sf")
+
+    monkeypatch.setattr(
+        llm_utils,
+        "resolve_completion_params_variables",
+        lambda parameters, _: parameters,
+    )
+    monkeypatch.setattr(
+        node,
+        "_fetch_llm_model_schema",
+        lambda **_: SimpleNamespace(features=[]),
+    )
+    monkeypatch.setattr(node, "_build_run_prompt", lambda **_: ([], []))
+
+    invoke_result = SimpleNamespace(
+        usage=LLMUsage.empty_usage(),
+        message=SimpleNamespace(
+            get_text_content=lambda: "{}",
+            tool_calls=[],
+        ),
+    )
+    monkeypatch.setattr(node.model_instance, "invoke_llm", lambda **_: invoke_result)
+
+    result = node._run()  # noqa: SLF001
+
+    assert result.inputs["query"] == "weather in sf"
+    assert result.inputs["model_provider"] == "test"
+    assert result.inputs["model_name"] == "test-model"
+
+
 def test_parameter_extractor_accepts_dependency_bundle() -> None:
     variable_pool = build_variable_pool(variables=[])
     runtime_state = GraphRuntimeState(

tests/nodes/question_classifier/test_question_classifier_node.py

@@ -300,6 +300,8 @@ def test_question_classifier_run_returns_custom_class_label(
     assert result.outputs["class_name"] == "Questions about refunds"
     assert result.outputs["class_label"] == "Refund desk"
     assert result.outputs["class_id"] == "refund"
+    assert result.inputs["model_provider"] == "openai"
+    assert result.inputs["model_name"] == "gpt-4o"
 
 
 def test_question_classifier_run_falls_back_to_canonical_class_label(