from functools import cached_property
from typing import Dict, Generator

from langgraph.graph.state import CompiledStateGraph

from alayaflow.workflow.runnable.base_runnable_workflow import BaseRunnableWorkflow
from alayaflow.workflow.workflow_info import WorkflowInfo

class StateGraphRunnableWorkflow(BaseRunnableWorkflow):
    def __init__(self, info: WorkflowInfo, graph: CompiledStateGraph) -> None:
        super().__init__(info)
        self._graph = graph

    def invoke(self, input_data: dict, user_config: dict) -> dict:
        return self._graph.invoke(input_data, {
            "configurable": user_config
        })
    
    def astream_events(self, input_data: dict, user_config: dict) -> Generator[Dict, None, None]:
        return self._graph.astream_events(input_data, {
            "configurable": user_config
        })

