"""Shared AgentCore invoke logic (monorepo entrypoint + packaged runtime).""" from __future__ import annotations import os import sys from collections.abc import AsyncIterator from pathlib import Path from session_store import ( append_turn, clear_session, get_messages, stringify_message_content, ) def configure_import_paths(app_root: Path | None = None) -> tuple[Path, Path]: """ Ensure imports resolve in the monorepo or a packaged AgentCore app folder. Returns ``(repo_root, agent_redact_root)`` for bootstrap_pi_config. """ root = (app_root or Path(__file__).resolve().parent).resolve() agent_redact = root shared_dir = root / "shared" for path in (root, agent_redact, shared_dir): text = str(path) if text not in sys.path: sys.path.insert(0, text) repo_root = root if (root / "agent-redact").is_dir(): repo_root = root agent_redact = root / "agent-redact" elif ( root.name == "RedactionAgent" and (root.parent.parent / "agent-redact").is_dir() ): repo_root = root.parent.parent.parent agent_redact = repo_root / "agent-redact" return repo_root, agent_redact def bootstrap_runtime_env(app_root: Path) -> None: """ Lightweight env setup for AgentCore (no Pi skills sync or monorepo ``tools/``). Full :func:`bootstrap_pi_config.ensure_pi_config_env` pulls in ``pi_workspace_skills`` and repo ``skills/`` — not vendored in the CodeZip bundle and will fail or stall runtime init on AWS. """ from dotenv import load_dotenv root = app_root.resolve() for env_name in ("agentcore.env", ".env"): env_file = root / env_name if env_file.is_file(): load_dotenv(env_file, override=False) os.environ.setdefault("AGENT_WORKSPACE_DIR", "/tmp/agentcore-workspace") os.environ.setdefault("AGENT_REDACTION_SPLIT_BACKEND", "true") os.environ.setdefault("AGENT_DEFAULT_PROVIDER", "amazon-bedrock") os.environ.setdefault("AGENT_DEFAULT_MODEL", "anthropic.claude-sonnet-4-6") os.environ.setdefault( "AWS_REGION", os.environ.get("AWS_DEFAULT_REGION", "eu-west-2") ) os.environ.setdefault("AWS_DEFAULT_REGION", os.environ["AWS_REGION"]) Path(os.environ["AGENT_WORKSPACE_DIR"]).mkdir(parents=True, exist_ok=True) INVOKE_RUNTIME_CONFIG_KEYS = frozenset( { "DOC_REDACTION_GRADIO_URL", "DOC_REDACTION_GRADIO_AUTH_USER", "DOC_REDACTION_GRADIO_AUTH_PASSWORD", # CloudFront magic-link cookie: the AgentCore runtime runs outside the VPC # and reaches doc_redaction through CloudFront, so the token forwarded by # build_agentcore_invoke_runtime_config must be applied here or every # backend request hits the login wall ("credentials were not provided"). "DOC_REDACTION_AUTH_TOKEN", "DOC_REDACTION_AUTH_COOKIE_NAME", "AGENT_DEFAULT_OCR_METHOD", "AGENT_DEFAULT_PII_METHOD", "AGENT_DEFAULT_PROVIDER", "AGENT_DEFAULT_MODEL", "AWS_REGION", "AWS_DEFAULT_REGION", "HF_TOKEN", "DOC_REDACTION_HF_TOKEN", } ) def apply_invoke_runtime_config(request: dict) -> None: """ Apply per-invoke backend settings from the Gradio UI (overrides agentcore.env). The AgentCore runtime on AWS has its own ``agentcore.env``; without this, a deployed HF Space URL can win over the operator's local ``agent.env``. """ raw = request.get("runtime_config") or request.get("runtime_env") or {} if not isinstance(raw, dict): return for key in INVOKE_RUNTIME_CONFIG_KEYS: value = raw.get(key) if value is None: continue text = str(value).strip() if text: os.environ[key] = text async def invoke_redaction_agent(request: dict) -> AsyncIterator[dict]: """Stream LangGraph agent events for one user prompt (multi-turn per session_hash).""" # Must run before LangChain imports (OpenInference patch order). try: from eval.arize_monitoring import setup_arize_ax_tracing setup_arize_ax_tracing() except ImportError: pass from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from workspace_sync import ( apply_workspace_files, collect_workspace_files_for_sync, ) apply_invoke_runtime_config(request) prompt = str(request.get("prompt") or request.get("message") or "").strip() session_hash = str(request.get("session_hash") or "").strip() or None if request.get("new_session"): clear_session(session_hash) if not prompt: yield {"type": "error", "message": "prompt is required"} return incoming_files = request.get("workspace_files") or [] if isinstance(incoming_files, list) and incoming_files: written = apply_workspace_files(session_hash, incoming_files) if written: yield { "type": "status", "message": f"Synced {len(written)} file(s) into AgentCore workspace.", } backend_url = (os.environ.get("DOC_REDACTION_GRADIO_URL") or "").strip().rstrip("/") if backend_url: yield { "type": "status", "message": f"Redaction backend for this turn: {backend_url}", } from eval.arize_monitoring import ( agent_turn_span, arize_session_context, langgraph_trace_config, ) from redaction_langgraph.graph import build_redaction_agent, graph_recursion_limit from redaction_langgraph.llm_errors import ( is_context_overflow_error, is_tool_call_json_parse_error, ) from redaction_langgraph.message_context import ( get_trim_stats, reset_trim_stats, set_aggressive_trim, ) from redaction_langgraph.workflow_continue import build_tool_call_json_retry_prompt graph, system_message = build_redaction_agent(session_hash) prior = get_messages(session_hash) inputs = {"messages": [system_message, *prior, HumanMessage(content=prompt)]} yield {"type": "agent_start"} assistant_chunks: list[str] = [] stream_config = langgraph_trace_config( session_hash, recursion_limit=graph_recursion_limit() ) def _emit_stream(active_graph, active_inputs): """Yield LangGraph node updates as they complete (do not buffer the run).""" reset_trim_stats() compaction_noted = False for event in active_graph.stream( active_inputs, stream_mode="updates", config=stream_config ): stats = get_trim_stats() if stats is not None and stats.trimmed and not compaction_noted: compaction_noted = True yield { "type": "status", "message": ( f"Context compaction ({stats.tokens_before:,} → " f"{stats.tokens_after:,} tokens)." ), } for node, update in event.items(): messages = update.get("messages") or [] for message in messages: if isinstance(message, AIMessage): text = stringify_message_content(message.content) if text: assistant_chunks.append(text) yield { "type": "message_update", "node": node, "role": "assistant", "content": text, "tool_calls": message.tool_calls or [], } elif isinstance(message, ToolMessage): yield { "type": "message_update", "node": node, "role": "tool", "content": stringify_message_content(message.content), "tool_name": str(message.name or "tool"), } else: content = getattr(message, "content", "") yield { "type": "message_update", "node": node, "role": getattr(message, "type", "unknown"), "content": content, } try: with arize_session_context(session_hash): with agent_turn_span( "agentcore.agent", session_hash=session_hash, message=prompt, agent_name="agentcore", ): try: for item in _emit_stream(graph, inputs): yield item except Exception as exc: if is_context_overflow_error(exc): assistant_chunks.clear() yield { "type": "status", "message": ( "Prompt exceeded model context — retrying once with " "aggressive compaction…" ), } set_aggressive_trim(True) try: graph, _ = build_redaction_agent( session_hash, aggressive_compaction=True ) for item in _emit_stream(graph, inputs): yield item finally: set_aggressive_trim(False) elif is_tool_call_json_parse_error(exc): assistant_chunks.clear() yield { "type": "status", "message": ( "Tool-call JSON was truncated or invalid — retrying once " "with a compact-script nudge…" ), } retry_inputs = { "messages": [ *(inputs.get("messages") or []), HumanMessage( content=build_tool_call_json_retry_prompt() ), ] } for item in _emit_stream(graph, retry_inputs): yield item else: raise except Exception as exc: yield {"type": "error", "message": f"LangGraph agent failed: {exc}"} return append_turn( session_hash, user_text=prompt, assistant_text="\n".join(assistant_chunks), ) if request.get("sync_workspace_files"): for item in collect_workspace_files_for_sync(session_hash): yield {"type": "workspace_file", **item} yield {"type": "agent_end", "message": "Agent finished."}