import fact0
from langgraph.graph import StateGraph

client = fact0.Client(api_key="f0_live_...")

def with_telemetry(name: str, fn):
    def wrapped(state):
        with client.telemetry.execution(
            agent_id=state.get("agent_id", "graph"),
            execution_id=state["execution_id"],
        ) as ex:
            with ex.span(name, span_type="CUSTOM") as span:
                result = fn(state)
                span.complete(output=result)
                return result
    return wrapped

graph = StateGraph(dict)
graph.add_node("plan", with_telemetry("plan", plan_fn))
graph.add_node("act", with_telemetry("act", act_fn))
graph.add_node("reflect", with_telemetry("reflect", reflect_fn))