from typing import Dict, Any, AsyncGenerator

from langgraph.graph.state import CompiledStateGraph, RunnableConfig

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, inputs: dict, context: dict, config: RunnableConfig) -> dict:
        return self._graph.invoke(inputs, config, context=context)
    
    async def stream_events_async(
        self, 
        inputs: dict, 
        context: dict, 
        config: RunnableConfig
    ) -> AsyncGenerator[Dict[str, Any], None]:
        async for event in self._graph.astream_events(
            inputs, 
            config, 
            version="v2", 
            context=context
        ):
            yield event

