agents/data_knowledge.py

"""
Data knowledge agent: builds the graph node that retrieves parameter/configuration
information based on user queries (e.g. market mappings, cluster definitions).
"""
from __future__ import annotations

from tools import parameter_reader
from tools.state import GraphState


def _tool_call(tool_or_fn, **kwargs):
    callable_obj = getattr(tool_or_fn, "func", tool_or_fn)
    return callable_obj(**kwargs)


def build_data_knowledge_node():
    """
    Factory that returns a data knowledge node for the graph.
    The node uses the parameter reader to fetch relevant parameter info for the user query.
    """
    def data_knowledge_node(state: GraphState) -> dict:
        user_query = state.get("user_query", "")
        rows = _tool_call(parameter_reader, user_query=user_query)

        error_message = state.get("error_message", "")
        if not rows and not error_message:
            error_message = "I couldn't find relevant parameter information for your request."

        return {
            "parameter_data": rows,
            "error_message": error_message,
        }

    return data_knowledge_node