|
|
@@ -0,0 +1,413 @@
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import time
|
|
|
+from copy import deepcopy
|
|
|
+from typing import Any, Callable
|
|
|
+
|
|
|
+from content_agent.business_modules import content_discovery, platform_access, rule_judgment
|
|
|
+from content_agent.business_modules.content_discovery import pattern_recall
|
|
|
+from content_agent.errors import ContentAgentError, ErrorCode
|
|
|
+from content_agent.interfaces import GeminiVideoClient, PlatformSearchClient, RuntimeFileStore
|
|
|
+
|
|
|
+
|
|
|
+INITIAL_BATCH_SIZE = 3
|
|
|
+MAX_PAGES = 3
|
|
|
+SEARCH_MIN_INTERVAL_SECONDS = 15.0
|
|
|
+PASS_ACTION = "ADD_TO_CONTENT_POOL"
|
|
|
+
|
|
|
+
|
|
|
+class SearchCallLimiter:
|
|
|
+ def __init__(
|
|
|
+ self,
|
|
|
+ min_interval_seconds: float = SEARCH_MIN_INTERVAL_SECONDS,
|
|
|
+ *,
|
|
|
+ now_fn: Callable[[], float] = time.monotonic,
|
|
|
+ sleep_fn: Callable[[float], None] = time.sleep,
|
|
|
+ ) -> None:
|
|
|
+ self.min_interval_seconds = min_interval_seconds
|
|
|
+ self.now_fn = now_fn
|
|
|
+ self.sleep_fn = sleep_fn
|
|
|
+ self._last_call_at: float | None = None
|
|
|
+
|
|
|
+ def wait(self) -> None:
|
|
|
+ if self._last_call_at is not None:
|
|
|
+ remaining = self.min_interval_seconds - (self.now_fn() - self._last_call_at)
|
|
|
+ if remaining > 0:
|
|
|
+ self.sleep_fn(remaining)
|
|
|
+ self._last_call_at = self.now_fn()
|
|
|
+
|
|
|
+
|
|
|
+def run(
|
|
|
+ *,
|
|
|
+ run_id: str,
|
|
|
+ policy_run_id: str,
|
|
|
+ search_queries: list[dict[str, Any]],
|
|
|
+ source_context: dict[str, Any],
|
|
|
+ policy_bundle: dict[str, Any],
|
|
|
+ platform_client: PlatformSearchClient,
|
|
|
+ runtime: RuntimeFileStore,
|
|
|
+ gemini_video_client: GeminiVideoClient,
|
|
|
+ limiter: SearchCallLimiter | None = None,
|
|
|
+) -> dict[str, list[dict[str, Any]]]:
|
|
|
+ if limiter is None and getattr(platform_client, "requires_progressive_search_rate_limit", False):
|
|
|
+ limiter = SearchCallLimiter()
|
|
|
+ context = _ProgressiveContext(
|
|
|
+ run_id=run_id,
|
|
|
+ policy_run_id=policy_run_id,
|
|
|
+ source_context=source_context,
|
|
|
+ policy_bundle=policy_bundle,
|
|
|
+ platform_client=platform_client,
|
|
|
+ runtime=runtime,
|
|
|
+ gemini_video_client=gemini_video_client,
|
|
|
+ limiter=limiter,
|
|
|
+ )
|
|
|
+ for query in search_queries:
|
|
|
+ context.process_query(query)
|
|
|
+
|
|
|
+ if search_queries and context.query_failures and not context.platform_results:
|
|
|
+ raise ContentAgentError(
|
|
|
+ ErrorCode.PLATFORM_REQUEST_FAILED,
|
|
|
+ "all platform queries failed",
|
|
|
+ {"query_failures": context.query_failures},
|
|
|
+ )
|
|
|
+
|
|
|
+ context.write_runtime()
|
|
|
+ return {
|
|
|
+ "platform_results": context.platform_results,
|
|
|
+ "query_failures": context.query_failures,
|
|
|
+ "discovered_content_items": context.discovered_content_items,
|
|
|
+ "content_media_records": context.content_media_records,
|
|
|
+ "pattern_recall_evidence": context.pattern_recall_evidence,
|
|
|
+ "evidence_bundles": context.evidence_bundles,
|
|
|
+ "rule_decisions": context.rule_decisions,
|
|
|
+ }
|
|
|
+
|
|
|
+
|
|
|
+class _ProgressiveContext:
|
|
|
+ def __init__(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ run_id: str,
|
|
|
+ policy_run_id: str,
|
|
|
+ source_context: dict[str, Any],
|
|
|
+ policy_bundle: dict[str, Any],
|
|
|
+ platform_client: PlatformSearchClient,
|
|
|
+ runtime: RuntimeFileStore,
|
|
|
+ gemini_video_client: GeminiVideoClient,
|
|
|
+ limiter: SearchCallLimiter | None,
|
|
|
+ ) -> None:
|
|
|
+ self.run_id = run_id
|
|
|
+ self.policy_run_id = policy_run_id
|
|
|
+ self.source_context = source_context
|
|
|
+ self.policy_bundle = policy_bundle
|
|
|
+ self.platform_client = platform_client
|
|
|
+ self.runtime = runtime
|
|
|
+ self.gemini_video_client = gemini_video_client
|
|
|
+ self.limiter = limiter
|
|
|
+ self.platform_results: list[dict[str, Any]] = []
|
|
|
+ self.query_failures: list[dict[str, Any]] = []
|
|
|
+ self.discovered_content_items: list[dict[str, Any]] = []
|
|
|
+ self.content_media_records: list[dict[str, Any]] = []
|
|
|
+ self.pattern_recall_evidence: list[dict[str, Any]] = []
|
|
|
+ self.evidence_bundles: list[dict[str, Any]] = []
|
|
|
+ self.rule_decisions: list[dict[str, Any]] = []
|
|
|
+ self._platform_result_by_key: dict[tuple[str, str], dict[str, Any]] = {}
|
|
|
+ self._record_indexes_by_key: dict[tuple[str, str], dict[str, int]] = {}
|
|
|
+ self._queries_with_consumed_next_pages: set[str] = set()
|
|
|
+
|
|
|
+ def process_query(self, query: dict[str, Any]) -> None:
|
|
|
+ first_page = self._fetch_page(query, page_number=1, cursor=query.get("page_cursor"))
|
|
|
+ if first_page is None:
|
|
|
+ return
|
|
|
+ page_count = len(first_page)
|
|
|
+ if page_count == 0:
|
|
|
+ return
|
|
|
+
|
|
|
+ first_batch = first_page[:INITIAL_BATCH_SIZE]
|
|
|
+ remainder_batch = first_page[INITIAL_BATCH_SIZE:]
|
|
|
+ passed_first = self._process_batch(
|
|
|
+ query,
|
|
|
+ first_batch,
|
|
|
+ page_number=1,
|
|
|
+ page_item_count=page_count,
|
|
|
+ batch_kind="initial_top3",
|
|
|
+ rank_offset=0,
|
|
|
+ )
|
|
|
+ if not passed_first or not remainder_batch:
|
|
|
+ return
|
|
|
+
|
|
|
+ passed_remainder = self._process_batch(
|
|
|
+ query,
|
|
|
+ remainder_batch,
|
|
|
+ page_number=1,
|
|
|
+ page_item_count=page_count,
|
|
|
+ batch_kind="page_remainder",
|
|
|
+ rank_offset=INITIAL_BATCH_SIZE,
|
|
|
+ )
|
|
|
+ if not passed_remainder:
|
|
|
+ return
|
|
|
+
|
|
|
+ cursor = _page_next_cursor(first_page)
|
|
|
+ has_more = _page_has_more(first_page)
|
|
|
+ page_number = 2
|
|
|
+ while page_number <= MAX_PAGES and has_more and cursor:
|
|
|
+ page = self._fetch_page(query, page_number=page_number, cursor=cursor)
|
|
|
+ if page is None or not page:
|
|
|
+ return
|
|
|
+ self._queries_with_consumed_next_pages.add(str(query["search_query_id"]))
|
|
|
+ page_count = len(page)
|
|
|
+ passed_page = self._process_batch(
|
|
|
+ query,
|
|
|
+ page,
|
|
|
+ page_number=page_number,
|
|
|
+ page_item_count=page_count,
|
|
|
+ batch_kind=f"page_{page_number:02d}",
|
|
|
+ rank_offset=0,
|
|
|
+ )
|
|
|
+ if page_number >= MAX_PAGES or not passed_page:
|
|
|
+ return
|
|
|
+ cursor = _page_next_cursor(page)
|
|
|
+ has_more = _page_has_more(page)
|
|
|
+ page_number += 1
|
|
|
+
|
|
|
+ def _fetch_page(
|
|
|
+ self,
|
|
|
+ query: dict[str, Any],
|
|
|
+ *,
|
|
|
+ page_number: int,
|
|
|
+ cursor: Any,
|
|
|
+ ) -> list[dict[str, Any]] | None:
|
|
|
+ page_query = dict(query)
|
|
|
+ if page_number > 1:
|
|
|
+ page_query["page_cursor"] = str(cursor or "")
|
|
|
+ try:
|
|
|
+ if self.limiter is not None:
|
|
|
+ self.limiter.wait()
|
|
|
+ search = getattr(self.platform_client, "search_full_page", None)
|
|
|
+ if not callable(search):
|
|
|
+ search = self.platform_client.search
|
|
|
+ return list(search(page_query))
|
|
|
+ except Exception as exc:
|
|
|
+ failure_query = page_query if page_number == 1 else {
|
|
|
+ **page_query,
|
|
|
+ "search_query_id": f"{query['search_query_id']}_progressive_page_{page_number:03d}",
|
|
|
+ }
|
|
|
+ failure = platform_access._query_failure(failure_query, exc) # noqa: SLF001
|
|
|
+ failure["progressive_page_number"] = page_number
|
|
|
+ self.query_failures.append(failure)
|
|
|
+ return None
|
|
|
+
|
|
|
+ def _process_batch(
|
|
|
+ self,
|
|
|
+ query: dict[str, Any],
|
|
|
+ page_results: list[dict[str, Any]],
|
|
|
+ *,
|
|
|
+ page_number: int,
|
|
|
+ page_item_count: int,
|
|
|
+ batch_kind: str,
|
|
|
+ rank_offset: int,
|
|
|
+ ) -> bool:
|
|
|
+ batch = []
|
|
|
+ for rank, result in enumerate(page_results, start=1):
|
|
|
+ prepared = self._prepare_result(
|
|
|
+ query,
|
|
|
+ result,
|
|
|
+ page_number=page_number,
|
|
|
+ page_item_count=page_item_count,
|
|
|
+ batch_kind=batch_kind,
|
|
|
+ rank=rank_offset + rank,
|
|
|
+ )
|
|
|
+ key = _content_key(prepared)
|
|
|
+ if not key:
|
|
|
+ batch.append(prepared)
|
|
|
+ continue
|
|
|
+ existing = self._platform_result_by_key.get(key)
|
|
|
+ if existing:
|
|
|
+ platform_access._append_query_source(existing, query) # noqa: SLF001
|
|
|
+ self._propagate_query_source_merge(key, existing)
|
|
|
+ continue
|
|
|
+ self._platform_result_by_key[key] = prepared
|
|
|
+ batch.append(prepared)
|
|
|
+ if not batch:
|
|
|
+ return False
|
|
|
+
|
|
|
+ discovered = content_discovery.run(
|
|
|
+ self.run_id,
|
|
|
+ self.policy_run_id,
|
|
|
+ batch,
|
|
|
+ self.source_context,
|
|
|
+ self.runtime,
|
|
|
+ write_runtime=False,
|
|
|
+ )
|
|
|
+ recalled = pattern_recall.run(
|
|
|
+ self.run_id,
|
|
|
+ self.policy_run_id,
|
|
|
+ discovered["discovered_content_items"],
|
|
|
+ discovered["content_media_records"],
|
|
|
+ discovered["evidence_bundles"],
|
|
|
+ self.source_context,
|
|
|
+ self.runtime,
|
|
|
+ self.gemini_video_client,
|
|
|
+ start_index=len(self.discovered_content_items) + 1,
|
|
|
+ write_runtime=False,
|
|
|
+ )
|
|
|
+ decisions = rule_judgment.run(
|
|
|
+ self.run_id,
|
|
|
+ self.policy_run_id,
|
|
|
+ recalled["evidence_bundles"],
|
|
|
+ self.policy_bundle,
|
|
|
+ self.runtime,
|
|
|
+ start_index=len(self.rule_decisions) + 1,
|
|
|
+ write_runtime=False,
|
|
|
+ )
|
|
|
+ self._accumulate(batch, recalled, decisions)
|
|
|
+ return any(decision.get("decision_action") == PASS_ACTION for decision in decisions)
|
|
|
+
|
|
|
+ def _prepare_result(
|
|
|
+ self,
|
|
|
+ query: dict[str, Any],
|
|
|
+ result: dict[str, Any],
|
|
|
+ *,
|
|
|
+ page_number: int,
|
|
|
+ page_item_count: int,
|
|
|
+ batch_kind: str,
|
|
|
+ rank: int,
|
|
|
+ ) -> dict[str, Any]:
|
|
|
+ prepared = platform_access._with_query_source(dict(result), query) # noqa: SLF001
|
|
|
+ original_has_more = bool(prepared.get("has_more"))
|
|
|
+ original_next_cursor = str(prepared.get("next_cursor") or "")
|
|
|
+ prepared.update(
|
|
|
+ {
|
|
|
+ "search_query_id": query["search_query_id"],
|
|
|
+ "content_discovery_id": _content_discovery_id(query["search_query_id"], page_number, rank),
|
|
|
+ "progressive_batch_id": (
|
|
|
+ f"{query['search_query_id']}_p{page_number:02d}_{batch_kind}"
|
|
|
+ ),
|
|
|
+ "progressive_page_number": page_number,
|
|
|
+ "progressive_batch_kind": batch_kind,
|
|
|
+ "progressive_page_item_count": page_item_count,
|
|
|
+ "progressive_item_rank_in_page": rank,
|
|
|
+ "progressive_original_has_more": original_has_more,
|
|
|
+ "progressive_original_next_cursor": original_next_cursor,
|
|
|
+ }
|
|
|
+ )
|
|
|
+ return prepared
|
|
|
+
|
|
|
+ def _accumulate(
|
|
|
+ self,
|
|
|
+ batch: list[dict[str, Any]],
|
|
|
+ recalled: dict[str, Any],
|
|
|
+ decisions: list[dict[str, Any]],
|
|
|
+ ) -> None:
|
|
|
+ for offset, result in enumerate(batch):
|
|
|
+ key = _content_key(result)
|
|
|
+ if key:
|
|
|
+ self._record_indexes_by_key[key] = {
|
|
|
+ "platform_result": len(self.platform_results) + offset,
|
|
|
+ "discovered": len(self.discovered_content_items) + offset,
|
|
|
+ "evidence_bundle": len(self.evidence_bundles) + offset,
|
|
|
+ "decision": len(self.rule_decisions) + offset,
|
|
|
+ }
|
|
|
+ self.platform_results.extend(batch)
|
|
|
+ self.discovered_content_items.extend(recalled["discovered_content_items"])
|
|
|
+ self.content_media_records.extend(recalled["content_media_records"])
|
|
|
+ self.pattern_recall_evidence.extend(recalled["pattern_recall_evidence"])
|
|
|
+ self.evidence_bundles.extend(recalled["evidence_bundles"])
|
|
|
+ self.rule_decisions.extend(decisions)
|
|
|
+
|
|
|
+ def _propagate_query_source_merge(
|
|
|
+ self,
|
|
|
+ key: tuple[str, str],
|
|
|
+ merged_result: dict[str, Any],
|
|
|
+ ) -> None:
|
|
|
+ indexes = self._record_indexes_by_key.get(key)
|
|
|
+ if not indexes:
|
|
|
+ return
|
|
|
+ fields = [
|
|
|
+ "query_sources",
|
|
|
+ "matched_search_query_ids",
|
|
|
+ "matched_search_queries",
|
|
|
+ "matched_search_query_generation_methods",
|
|
|
+ ]
|
|
|
+ for collection_name in ["discovered_content_items"]:
|
|
|
+ row = getattr(self, collection_name)[indexes["discovered"]]
|
|
|
+ _copy_fields(row, merged_result, fields)
|
|
|
+ bundle = self.evidence_bundles[indexes["evidence_bundle"]]
|
|
|
+ _copy_fields(bundle.setdefault("source_evidence", {}), merged_result, fields)
|
|
|
+ decision = self.rule_decisions[indexes["decision"]]
|
|
|
+ _copy_fields(decision.setdefault("source_evidence", {}), merged_result, fields)
|
|
|
+ raw_payload = decision.setdefault("raw_payload", {})
|
|
|
+ if isinstance(raw_payload.get("source_evidence"), dict):
|
|
|
+ _copy_fields(raw_payload["source_evidence"], merged_result, fields)
|
|
|
+
|
|
|
+ def write_runtime(self) -> None:
|
|
|
+ self.runtime.append_jsonl(self.run_id, "content_media_records.jsonl", self.content_media_records)
|
|
|
+ self.runtime.append_jsonl(self.run_id, "pattern_recall_evidence.jsonl", self.pattern_recall_evidence)
|
|
|
+ self._hide_consumed_query_cursors()
|
|
|
+ self.runtime.append_jsonl(self.run_id, "discovered_content_items.jsonl", self.discovered_content_items)
|
|
|
+ self.runtime.append_jsonl(self.run_id, "rule_decisions.jsonl", self.rule_decisions)
|
|
|
+
|
|
|
+ def _hide_consumed_query_cursors(self) -> None:
|
|
|
+ if not self._queries_with_consumed_next_pages:
|
|
|
+ return
|
|
|
+ for row in self.platform_results:
|
|
|
+ self._hide_cursor_if_consumed(row)
|
|
|
+ for row in self.discovered_content_items:
|
|
|
+ self._hide_cursor_if_consumed(row)
|
|
|
+ for bundle in self.evidence_bundles:
|
|
|
+ source_evidence = bundle.get("source_evidence")
|
|
|
+ if isinstance(source_evidence, dict):
|
|
|
+ self._hide_cursor_if_consumed(source_evidence)
|
|
|
+ for decision in self.rule_decisions:
|
|
|
+ source_evidence = decision.get("source_evidence")
|
|
|
+ if isinstance(source_evidence, dict):
|
|
|
+ self._hide_cursor_if_consumed(source_evidence)
|
|
|
+ raw_payload = decision.get("raw_payload")
|
|
|
+ if isinstance(raw_payload, dict) and isinstance(raw_payload.get("source_evidence"), dict):
|
|
|
+ self._hide_cursor_if_consumed(raw_payload["source_evidence"])
|
|
|
+
|
|
|
+ def _hide_cursor_if_consumed(self, row: dict[str, Any]) -> None:
|
|
|
+ if str(row.get("search_query_id") or "") not in self._queries_with_consumed_next_pages:
|
|
|
+ return
|
|
|
+ row["has_more"] = False
|
|
|
+ row["next_cursor"] = ""
|
|
|
+ raw_payload = row.get("raw_payload")
|
|
|
+ if isinstance(raw_payload, dict):
|
|
|
+ raw_payload["has_more"] = False
|
|
|
+ raw_payload["next_cursor"] = ""
|
|
|
+
|
|
|
+
|
|
|
+def _content_key(result: dict[str, Any]) -> tuple[str, str] | None:
|
|
|
+ content_id = str(result.get("platform_content_id") or "")
|
|
|
+ if not content_id:
|
|
|
+ return None
|
|
|
+ return str(result.get("platform") or ""), content_id
|
|
|
+
|
|
|
+
|
|
|
+def _content_discovery_id(search_query_id: str, page_number: int, rank: int) -> str:
|
|
|
+ if page_number == 1:
|
|
|
+ return f"{search_query_id}_content_{rank:03d}"
|
|
|
+ return f"{search_query_id}_page_{page_number:03d}_content_{rank:03d}"
|
|
|
+
|
|
|
+
|
|
|
+def _page_has_more(results: list[dict[str, Any]]) -> bool:
|
|
|
+ return any(bool(result.get("has_more")) for result in results)
|
|
|
+
|
|
|
+
|
|
|
+def _page_next_cursor(results: list[dict[str, Any]]) -> str:
|
|
|
+ for result in results:
|
|
|
+ cursor = str(result.get("next_cursor") or "")
|
|
|
+ if cursor:
|
|
|
+ return cursor
|
|
|
+ return ""
|
|
|
+
|
|
|
+
|
|
|
+def _copy_fields(target: dict[str, Any], source: dict[str, Any], fields: list[str]) -> None:
|
|
|
+ for field in fields:
|
|
|
+ if field in source:
|
|
|
+ target[field] = deepcopy(source[field])
|
|
|
+ raw_payload = target.get("raw_payload")
|
|
|
+ if isinstance(raw_payload, dict):
|
|
|
+ for field in fields:
|
|
|
+ if field in source:
|
|
|
+ raw_payload[field] = deepcopy(source[field])
|