content-find-agent-new/content_agent/graph.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from langgraph.graph import END, START, StateGraph

from content_agent.business_modules import (
    content_discovery,
    learning_review,
    platform_access,
    policy_version,
    result_source_lookup,
    rule_judgment,
    run_record,
    search_intent,
    walk_engine,
    source_seed,
    walk_strategy,
)
from content_agent.business_modules.content_discovery import pattern_recall
from content_agent.interfaces import (
    CategoryMatchClient,
    DecodeClient,
    PlatformSearchClient,
    PolicyBundleStore,
    QueryVariantClient,
    RuntimeFileStore,
)
from content_agent.models import RunState


@dataclass(frozen=True)
class RunDependencies:
    runtime: RuntimeFileStore
    platform_client: PlatformSearchClient
    policy_store: PolicyBundleStore
    query_variant_client: QueryVariantClient
    decode_client: DecodeClient
    category_match_client: CategoryMatchClient
    pattern_recall_max_wait_seconds: float = 1200.0
    pattern_recall_poll_interval_seconds: float = 5.0


def build_run_graph(deps: RunDependencies):
    graph = StateGraph(RunState)

    def load_source(state: RunState) -> dict[str, Any]:
        result = source_seed.run(
            state["run_id"], state["policy_run_id"], state.get("source"), deps.runtime
        )
        return {**result, "current_step": "load_source"}

    def plan_queries(state: RunState) -> dict[str, Any]:
        search_queries = search_intent.run(
            state["run_id"],
            state["policy_run_id"],
            state["pattern_seed_pack"],
            deps.runtime,
            deps.query_variant_client,
        )
        return {"search_queries": search_queries, "current_step": "plan_queries"}

    def search_platform(state: RunState) -> dict[str, Any]:
        result = platform_access.run(state["search_queries"], deps.platform_client)
        return {**result, "current_step": "search_platform"}

    def build_discovered_content(state: RunState) -> dict[str, Any]:
        result = content_discovery.run(
            state["run_id"],
            state["policy_run_id"],
            state["platform_results"],
            state["source_context"],
            deps.runtime,
        )
        return {**result, "current_step": "build_discovered_content"}

    def recall_pattern(state: RunState) -> dict[str, Any]:
        result = pattern_recall.run(
            state["run_id"],
            state["policy_run_id"],
            state["discovered_content_items"],
            state["content_media_records"],
            state["evidence_bundles"],
            state["source_context"],
            state["pattern_seed_pack"],
            deps.runtime,
            deps.decode_client,
            deps.category_match_client,
            max_wait_seconds=deps.pattern_recall_max_wait_seconds,
            poll_interval_seconds=deps.pattern_recall_poll_interval_seconds,
        )
        return {**result, "current_step": "recall_pattern"}

    def load_policy(state: RunState) -> dict[str, Any]:
        bundle = policy_version.run(state["strategy_version"], deps.policy_store)
        return {
            "policy_bundle": bundle,
            "policy_bundle_id": bundle["policy_bundle_id"],
            "strategy_version": bundle["strategy_version"],
            "strategy_source_ref": bundle["strategy_source_ref"],
            "current_step": "load_policy",
        }

    def evaluate_rules(state: RunState) -> dict[str, Any]:
        decisions = rule_judgment.run(
            state["run_id"],
            state["policy_run_id"],
            state["evidence_bundles"],
            state["policy_bundle"],
            deps.runtime,
        )
        return {"rule_decisions": decisions, "current_step": "evaluate_rules"}

    def execute_walk(state: RunState) -> dict[str, Any]:
        result = walk_engine.run_bounded_walk(
            run_id=state["run_id"],
            policy_run_id=state["policy_run_id"],
            pattern_seed_pack=state["pattern_seed_pack"],
            source_context=state["source_context"],
            search_queries=state["search_queries"],
            discovered_content_items=state["discovered_content_items"],
            content_media_records=state["content_media_records"],
            evidence_bundles=state["evidence_bundles"],
            rule_decisions=state["rule_decisions"],
            policy_bundle=state["policy_bundle"],
            platform_client=deps.platform_client,
            runtime=deps.runtime,
            decode_client=deps.decode_client,
            category_match_client=deps.category_match_client,
            max_wait_seconds=deps.pattern_recall_max_wait_seconds,
            poll_interval_seconds=deps.pattern_recall_poll_interval_seconds,
        )
        return {**result, "current_step": "execute_walk"}

    def plan_walk(state: RunState) -> dict[str, Any]:
        result = walk_strategy.run(
            state["pattern_seed_pack"],
            state["search_queries"],
            state["discovered_content_items"],
            state["rule_decisions"],
        )
        result["walk_actions"] = [*state.get("walk_actions", []), *result["walk_actions"]]
        return {**result, "current_step": "plan_walk"}

    def record_run(state: RunState) -> dict[str, Any]:
        result = run_record.run(
            state["run_id"],
            state["policy_run_id"],
            state["search_queries"],
            state["discovered_content_items"],
            state["rule_decisions"],
            state["source_path_record_basis"],
            state["policy_bundle"],
            deps.runtime,
            walk_actions=state["walk_actions"],
            query_failures=state.get("query_failures", []),
        )
        return {**result, "current_step": "record_run"}

    def commit_results(state: RunState) -> dict[str, Any]:
        final_output = result_source_lookup.run(
            state["run_id"],
            state["policy_run_id"],
            state["policy_bundle"],
            state["discovered_content_items"],
            state["content_media_records"],
            state["rule_decisions"],
            state["source_path_records"],
            state["search_clues"],
            deps.runtime,
        )
        return {"final_output": final_output, "current_step": "commit_results"}

    def review_strategy(state: RunState) -> dict[str, Any]:
        review = learning_review.run(state["run_id"], state["policy_run_id"], deps.runtime)
        return {"strategy_review": review, "current_step": "review_strategy", "status": "success"}

    graph.add_node("load_source", load_source)
    graph.add_node("plan_queries", plan_queries)
    graph.add_node("search_platform", search_platform)
    graph.add_node("build_discovered_content", build_discovered_content)
    graph.add_node("recall_pattern", recall_pattern)
    graph.add_node("load_policy", load_policy)
    graph.add_node("evaluate_rules", evaluate_rules)
    graph.add_node("execute_walk", execute_walk)
    graph.add_node("plan_walk", plan_walk)
    graph.add_node("record_run", record_run)
    graph.add_node("commit_results", commit_results)
    graph.add_node("review_strategy", review_strategy)

    graph.add_edge(START, "load_source")
    graph.add_edge("load_source", "plan_queries")
    graph.add_edge("plan_queries", "search_platform")
    graph.add_edge("search_platform", "build_discovered_content")
    graph.add_edge("build_discovered_content", "recall_pattern")
    graph.add_edge("recall_pattern", "load_policy")
    graph.add_edge("load_policy", "evaluate_rules")
    graph.add_edge("evaluate_rules", "execute_walk")
    graph.add_edge("execute_walk", "plan_walk")
    graph.add_edge("plan_walk", "record_run")
    graph.add_edge("record_run", "commit_results")
    graph.add_edge("commit_results", "review_strategy")
    graph.add_edge("review_strategy", END)
    return graph.compile()