feat: integrate NpuClient for AI screenshot privacy redaction
Add PII redaction to cloud LLM path in hybrid_planner: - NpuClient: async httpx client for NPU Daemon (redact_image, ocr_analyze) - _redact_for_cloud(): intercepts screenshots before remote LLM calls - Mixed mode: text_only (no image sent) vs image (redacted JPEG) - Graceful degradation: NPU unavailable → send original with warning - Privacy metrics: redactions, findings_total, text_only/image mode counts - Config: npu_daemon_url, privacy_redact_enabled, privacy_redact_types 12 new tests for NpuClient + HybridPlanner privacy integration. All 462 tests passing. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -7,6 +7,7 @@ from .hybrid_planner import HybridPlanner, PlanResult
|
||||
from .kvm_client import KVMClient
|
||||
from .llm_planner import LLMPlanner
|
||||
from .memory_client import MemoryClient
|
||||
from .npu_client import NpuClient
|
||||
from .perception import PerceptionEngine, SceneGraph, UIElement
|
||||
from .runner import AutonomousRunner, TaskQueue
|
||||
from .template_recorder import TemplateRecorder
|
||||
@@ -24,6 +25,7 @@ __all__ = [
|
||||
"KVMClient",
|
||||
"LLMPlanner",
|
||||
"MemoryClient",
|
||||
"NpuClient",
|
||||
"OperationTemplate",
|
||||
"PerceptionEngine",
|
||||
"PlanResult",
|
||||
|
||||
@@ -22,6 +22,7 @@ from .config import AgentConfig
|
||||
from .kvm_client import KVMClient
|
||||
from .llm_planner import LLMPlanner
|
||||
from .memory_client import MemoryClient
|
||||
from .npu_client import NpuClient
|
||||
from .perception import PerceptionEngine
|
||||
from .runner import AutonomousRunner, TaskQueue
|
||||
from .template_recorder import TemplateRecorder
|
||||
@@ -46,6 +47,11 @@ def _build_stack(cfg: AgentConfig):
|
||||
if cfg.memory_enabled:
|
||||
memory = MemoryClient(cfg.memory_base_url)
|
||||
|
||||
# NPU client for PII redaction before cloud LLM calls
|
||||
npu = None
|
||||
if cfg.privacy_redact_enabled:
|
||||
npu = NpuClient(cfg.npu_daemon_url)
|
||||
|
||||
agent = KVMAgent(
|
||||
kvm=kvm,
|
||||
planner=planner,
|
||||
@@ -58,7 +64,7 @@ def _build_stack(cfg: AgentConfig):
|
||||
window_manager=window_mgr if cfg.fullscreen_enabled else None,
|
||||
config=cfg,
|
||||
)
|
||||
return kvm, planner, agent, memory, perception
|
||||
return kvm, planner, agent, memory, perception, npu
|
||||
|
||||
|
||||
# ── Subcommand: run ──────────────────────────────────────────────
|
||||
@@ -74,7 +80,7 @@ async def _cmd_run(args, cfg: AgentConfig):
|
||||
if args.step_delay:
|
||||
cfg.step_delay = args.step_delay
|
||||
|
||||
kvm, planner, agent, memory, perception = _build_stack(cfg)
|
||||
kvm, planner, agent, memory, perception, npu = _build_stack(cfg)
|
||||
|
||||
try:
|
||||
result = await agent.run_task(args.task)
|
||||
@@ -103,7 +109,7 @@ async def _cmd_run(args, cfg: AgentConfig):
|
||||
|
||||
async def _cmd_daemon(args, cfg: AgentConfig):
|
||||
"""Start the autonomous daemon with embedded API server."""
|
||||
kvm, planner, agent, memory, perception = _build_stack(cfg)
|
||||
kvm, planner, agent, memory, perception, npu = _build_stack(cfg)
|
||||
|
||||
queue = TaskQueue(
|
||||
host=cfg.db_host,
|
||||
|
||||
@@ -10,6 +10,7 @@ Perceive-decide-act loop enhanced with:
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional
|
||||
|
||||
@@ -18,6 +19,7 @@ from .kvm_client import KVMClient
|
||||
from .llm_planner import LLMPlanner
|
||||
from .screen_state import ScreenStateDetector, PCState
|
||||
from .validator import ValidationResult
|
||||
from .workflow_hooks import EventBus, StepEvent
|
||||
from . import safety
|
||||
from . import mouse_ops
|
||||
|
||||
@@ -80,6 +82,7 @@ class KVMAgent:
|
||||
validator=None,
|
||||
window_manager=None,
|
||||
config=None,
|
||||
hooks: Optional[EventBus] = None,
|
||||
):
|
||||
self.kvm = kvm
|
||||
self.planner = planner
|
||||
@@ -91,6 +94,7 @@ class KVMAgent:
|
||||
self.click_delay = click_delay
|
||||
self.force_control = force_control
|
||||
self._running = False
|
||||
self._hooks = hooks
|
||||
|
||||
# v2 components
|
||||
self.perception = perception
|
||||
@@ -148,18 +152,44 @@ class KVMAgent:
|
||||
)
|
||||
|
||||
# ── Perceive (single OCR call per step) ────
|
||||
_t0 = time.monotonic()
|
||||
before_screenshot = await self.kvm.screenshot()
|
||||
_ss_ms = (time.monotonic() - _t0) * 1000
|
||||
if self._hooks:
|
||||
await self._hooks.emit(StepEvent(
|
||||
"screenshot", step, "kvm_capture", _ss_ms,
|
||||
"capture", f"{len(before_screenshot)}B",
|
||||
))
|
||||
scene = None
|
||||
if self.perception:
|
||||
_t0 = time.monotonic()
|
||||
scene = await self.perception.perceive(before_screenshot)
|
||||
_ocr_ms = (time.monotonic() - _t0) * 1000
|
||||
if self._hooks:
|
||||
await self._hooks.emit(StepEvent(
|
||||
"ocr_complete", step, "ppocrv4_det+rec", _ocr_ms,
|
||||
f"{len(before_screenshot)}B image",
|
||||
f"{len(scene.elements)} elements",
|
||||
{"element_count": len(scene.elements)},
|
||||
))
|
||||
|
||||
# ── State check ────────────────────────────
|
||||
if self.state_detector and scene:
|
||||
_t0 = time.monotonic()
|
||||
detection = self.state_detector.detect(
|
||||
before_screenshot,
|
||||
scene.raw_ocr_text,
|
||||
bool(scene.elements),
|
||||
)
|
||||
if self._hooks:
|
||||
await self._hooks.emit(StepEvent(
|
||||
"state_detected", step, "keyword_matcher",
|
||||
(time.monotonic() - _t0) * 1000,
|
||||
f"{len(scene.raw_ocr_text)} chars",
|
||||
f"{detection.state.value} ({detection.confidence:.0%})",
|
||||
{"state": detection.state.value,
|
||||
"confidence": detection.confidence},
|
||||
))
|
||||
if detection.state == PCState.SLEEP:
|
||||
logger.info(
|
||||
"Target PC sleeping (%s) — attempting wake",
|
||||
@@ -223,6 +253,7 @@ class KVMAgent:
|
||||
)
|
||||
|
||||
# ── Decide ───────────────────────────────────
|
||||
_t0 = time.monotonic()
|
||||
action = await self.planner.plan_action(
|
||||
before_screenshot,
|
||||
task_description,
|
||||
@@ -231,6 +262,16 @@ class KVMAgent:
|
||||
max_steps=self.max_steps,
|
||||
scene_text=scene_text,
|
||||
)
|
||||
if self._hooks:
|
||||
_plan_ms = (time.monotonic() - _t0) * 1000
|
||||
_src = getattr(action, '_source', 'unknown')
|
||||
await self._hooks.emit(StepEvent(
|
||||
"action_planned", step, _src, _plan_ms,
|
||||
task_description[:80],
|
||||
f"{action.type}: {action.reason}",
|
||||
{"action_type": action.type, "x": action.x,
|
||||
"y": action.y, "reason": action.reason},
|
||||
))
|
||||
|
||||
action_desc = f"[{step + 1}] {action.type}: {action.reason}"
|
||||
action_history.append(action_desc)
|
||||
@@ -254,10 +295,19 @@ class KVMAgent:
|
||||
)
|
||||
|
||||
# ── Act + Validate ───────────────────────────
|
||||
_t0 = time.monotonic()
|
||||
validated = await self._execute_with_validation(
|
||||
action, before_screenshot, scene, step,
|
||||
)
|
||||
after_screenshot = await self.kvm.screenshot()
|
||||
if self._hooks:
|
||||
await self._hooks.emit(StepEvent(
|
||||
"action_executed", step, "hid_controller",
|
||||
(time.monotonic() - _t0) * 1000,
|
||||
f"{action.type}: {action.reason}",
|
||||
"validated" if validated else "failed",
|
||||
{"validated": validated},
|
||||
))
|
||||
|
||||
step_records.append(
|
||||
StepRecord(
|
||||
|
||||
@@ -8,6 +8,7 @@ Runs inside the same asyncio event loop as AutonomousRunner.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
@@ -16,6 +17,7 @@ from aiohttp import web
|
||||
|
||||
from .config import AgentConfig
|
||||
from .runner import TaskQueue
|
||||
from .workflow_hooks import EventBus, WorkflowHook, StepEvent
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .runner import AutonomousRunner
|
||||
@@ -75,6 +77,7 @@ class AgentAPIServer:
|
||||
config: AgentConfig,
|
||||
autonomous_runner: AutonomousRunner | None = None,
|
||||
hybrid_planner=None,
|
||||
event_bus: EventBus | None = None,
|
||||
host: str = "0.0.0.0",
|
||||
port: int = 8890,
|
||||
):
|
||||
@@ -82,6 +85,7 @@ class AgentAPIServer:
|
||||
self._config = config
|
||||
self._autonomous_runner = autonomous_runner
|
||||
self._hybrid_planner = hybrid_planner
|
||||
self._event_bus = event_bus
|
||||
self._host = host
|
||||
self._port = port
|
||||
self._runner: Optional[web.AppRunner] = None
|
||||
@@ -97,6 +101,7 @@ class AgentAPIServer:
|
||||
app.router.add_get("/api/v1/agent/config", self._handle_get_config)
|
||||
app.router.add_patch("/api/v1/agent/config", self._handle_patch_config)
|
||||
app.router.add_get("/api/v1/agent/metrics", self._handle_metrics)
|
||||
app.router.add_get("/api/v1/agent/events", self._handle_events)
|
||||
|
||||
self._runner = web.AppRunner(app)
|
||||
await self._runner.setup()
|
||||
@@ -243,3 +248,58 @@ class AgentAPIServer:
|
||||
"remote_llm_calls": 0,
|
||||
"estimated_tokens_saved": 0,
|
||||
})
|
||||
|
||||
async def _handle_events(self, request: web.Request) -> web.StreamResponse:
|
||||
"""SSE stream of workflow step events."""
|
||||
if not self._event_bus:
|
||||
return _error("Event bus not configured", status=503)
|
||||
|
||||
resp = web.StreamResponse(headers={
|
||||
"Content-Type": "text/event-stream",
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
**CORS_HEADERS,
|
||||
})
|
||||
await resp.prepare(request)
|
||||
|
||||
hook = _SSEHook(resp)
|
||||
self._event_bus.register(hook)
|
||||
try:
|
||||
# Keep connection open; send keepalive every 15s
|
||||
while not resp.task.done():
|
||||
await asyncio.sleep(15)
|
||||
try:
|
||||
await resp.write(b": keepalive\n\n")
|
||||
except (ConnectionResetError, ConnectionAbortedError):
|
||||
break
|
||||
finally:
|
||||
self._event_bus.unregister(hook)
|
||||
return resp
|
||||
|
||||
|
||||
class _SSEHook(WorkflowHook):
|
||||
"""Writes StepEvents to an aiohttp SSE stream."""
|
||||
|
||||
def __init__(self, response: web.StreamResponse):
|
||||
self._resp = response
|
||||
|
||||
async def on_event(self, event: StepEvent) -> None:
|
||||
line = f"data: {json.dumps(event.to_dict())}\n\n"
|
||||
try:
|
||||
await self._resp.write(line.encode())
|
||||
except (ConnectionResetError, ConnectionAbortedError):
|
||||
pass
|
||||
|
||||
async def on_task_start(self, task_id: str, description: str) -> None:
|
||||
data = {"type": "task_start", "task_id": task_id, "description": description}
|
||||
try:
|
||||
await self._resp.write(f"data: {json.dumps(data)}\n\n".encode())
|
||||
except (ConnectionResetError, ConnectionAbortedError):
|
||||
pass
|
||||
|
||||
async def on_task_end(self, task_id: str, success: bool, reason: str) -> None:
|
||||
data = {"type": "task_end", "task_id": task_id, "success": success, "reason": reason}
|
||||
try:
|
||||
await self._resp.write(f"data: {json.dumps(data)}\n\n".encode())
|
||||
except (ConnectionResetError, ConnectionAbortedError):
|
||||
pass
|
||||
|
||||
@@ -51,6 +51,11 @@ class AgentConfig:
|
||||
llm_image_detail: str = "auto"
|
||||
adaptive_image: bool = True
|
||||
|
||||
# NPU Daemon (privacy redaction)
|
||||
npu_daemon_url: str = "http://localhost:8004"
|
||||
privacy_redact_enabled: bool = True # redact PII before sending to cloud LLM
|
||||
privacy_redact_types: str = "id_card,phone,bank_card,email"
|
||||
|
||||
# Memory (mem-bridge)
|
||||
memory_enabled: bool = True
|
||||
memory_base_url: str = "http://localhost:8003"
|
||||
@@ -99,6 +104,8 @@ class AgentConfig:
|
||||
"HIAPI_KEY": "llm_api_key",
|
||||
"LLM_MODEL": "llm_model",
|
||||
"MEMORY_BASE_URL": "memory_base_url",
|
||||
"NPU_DAEMON_URL": "npu_daemon_url",
|
||||
"PRIVACY_REDACT_ENABLED": "privacy_redact_enabled",
|
||||
"KVM_AGENT_DB_HOST": "db_host",
|
||||
"KVM_AGENT_DB_NAME": "db_name",
|
||||
"KVM_AGENT_DB_USER": "db_user",
|
||||
|
||||
@@ -12,18 +12,21 @@ Decision flow:
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from .actions import Action
|
||||
from .npu_client import NpuClient, RedactResult
|
||||
from .template_store import (
|
||||
Condition,
|
||||
OperationTemplate,
|
||||
TemplateStore,
|
||||
)
|
||||
from .validator import StepValidator, ValidationResult
|
||||
from .workflow_hooks import EventBus, StepEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -41,6 +44,10 @@ class PlannerMetrics:
|
||||
local_llm_failures: int = 0
|
||||
remote_llm_calls: int = 0
|
||||
estimated_tokens_saved: int = 0
|
||||
privacy_redactions: int = 0
|
||||
privacy_findings_total: int = 0
|
||||
privacy_text_only_mode: int = 0
|
||||
privacy_image_mode: int = 0
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
@@ -50,6 +57,10 @@ class PlannerMetrics:
|
||||
"local_llm_failures": self.local_llm_failures,
|
||||
"remote_llm_calls": self.remote_llm_calls,
|
||||
"estimated_tokens_saved": self.estimated_tokens_saved,
|
||||
"privacy_redactions": self.privacy_redactions,
|
||||
"privacy_findings_total": self.privacy_findings_total,
|
||||
"privacy_text_only_mode": self.privacy_text_only_mode,
|
||||
"privacy_image_mode": self.privacy_image_mode,
|
||||
}
|
||||
|
||||
|
||||
@@ -78,6 +89,10 @@ class HybridPlanner:
|
||||
validator: Optional[StepValidator] = None,
|
||||
memory_client=None,
|
||||
local_llm_url: str = "http://localhost:8891",
|
||||
hooks: Optional[EventBus] = None,
|
||||
npu_client: Optional[NpuClient] = None,
|
||||
privacy_redact_enabled: bool = True,
|
||||
privacy_redact_types: Optional[list[str]] = None,
|
||||
):
|
||||
self._llm = llm_planner
|
||||
self._templates = template_store
|
||||
@@ -86,6 +101,12 @@ class HybridPlanner:
|
||||
self._memory = memory_client
|
||||
self._local_llm_url = local_llm_url.rstrip("/")
|
||||
self._local_llm_available: Optional[bool] = None
|
||||
self._hooks = hooks
|
||||
self._npu = npu_client
|
||||
self._privacy_redact_enabled = privacy_redact_enabled
|
||||
self._privacy_redact_types = privacy_redact_types or [
|
||||
"id_card", "phone", "bank_card", "email",
|
||||
]
|
||||
|
||||
# State for active template replay
|
||||
self._active_template: Optional[OperationTemplate] = None
|
||||
@@ -147,10 +168,19 @@ class HybridPlanner:
|
||||
if self._active_template and self._template_step < len(
|
||||
self._active_template.steps
|
||||
):
|
||||
_t0 = time.monotonic()
|
||||
result = await self._replay_step(screenshot)
|
||||
if result is not None:
|
||||
self.metrics.template_hits += 1
|
||||
self.metrics.estimated_tokens_saved += 3600
|
||||
if self._hooks:
|
||||
await self._hooks.emit(StepEvent(
|
||||
"template_lookup", step, "template_replay",
|
||||
(time.monotonic() - _t0) * 1000,
|
||||
self._active_template.task_pattern,
|
||||
f"step {result.template_step_index}: {result.action.type}",
|
||||
{"source": "template"},
|
||||
))
|
||||
return result
|
||||
# Postcondition failed — abandon template
|
||||
logger.warning("Template postcondition failed at step %d, falling back",
|
||||
@@ -159,12 +189,20 @@ class HybridPlanner:
|
||||
|
||||
# ── Path 2: Local RKLLM for simple tasks ──────────────
|
||||
if await self._should_use_local(task, history):
|
||||
_t0 = time.monotonic()
|
||||
local_result = await self._local_llm_plan(
|
||||
screenshot, task, step, history, max_steps,
|
||||
)
|
||||
_local_ms = (time.monotonic() - _t0) * 1000
|
||||
if local_result is not None:
|
||||
self._local_fail_count = 0 # Reset on success
|
||||
self.metrics.local_llm_calls += 1
|
||||
if self._hooks:
|
||||
await self._hooks.emit(StepEvent(
|
||||
"local_llm", step, "rkllm-qwen2.5-1.5b", _local_ms,
|
||||
task[:80], f"{local_result.action.type}: {local_result.action.reason}",
|
||||
{"source": "local_llm"},
|
||||
))
|
||||
return local_result
|
||||
# Local failed — track fallback
|
||||
self._local_fail_count += 1
|
||||
@@ -172,9 +210,19 @@ class HybridPlanner:
|
||||
logger.warning("RKLLM fallback #%d → remote LLM (cost warning)", self._local_fail_count)
|
||||
|
||||
# ── Path 3: Remote LLM (full capability) ─────────────
|
||||
return await self._llm_plan_with_memory(
|
||||
_t0 = time.monotonic()
|
||||
result = await self._llm_plan_with_memory(
|
||||
screenshot, task, step, history, max_steps,
|
||||
)
|
||||
if self._hooks:
|
||||
_model = getattr(self._llm, 'model', 'remote_llm')
|
||||
await self._hooks.emit(StepEvent(
|
||||
"remote_llm", step, _model,
|
||||
(time.monotonic() - _t0) * 1000,
|
||||
task[:80], f"{result.action.type}: {result.action.reason}",
|
||||
{"source": "remote_llm"},
|
||||
))
|
||||
return result
|
||||
|
||||
async def _replay_step(self, screenshot: bytes) -> Optional[PlanResult]:
|
||||
"""Replay the current template step.
|
||||
@@ -231,14 +279,39 @@ class HybridPlanner:
|
||||
history: list[str],
|
||||
max_steps: int,
|
||||
) -> PlanResult:
|
||||
"""Plan via LLM with injected memory context and OCR scene text."""
|
||||
"""Plan via remote LLM with PII redaction + memory context.
|
||||
|
||||
Data flow:
|
||||
1. screenshot → NPU Daemon /privacy/redact-image (P0)
|
||||
2. redacted content → cloud LLM
|
||||
3. original screenshot + findings → audit log
|
||||
"""
|
||||
self._remote_call_count += 1
|
||||
self.metrics.remote_llm_calls += 1
|
||||
scene_text = ""
|
||||
context_hint = ""
|
||||
screenshot_for_llm = screenshot # default: original
|
||||
|
||||
# Get OCR scene graph
|
||||
if self._perception:
|
||||
# ── PII Redaction (before sending to cloud) ──────────
|
||||
redact_result = await self._redact_for_cloud(screenshot, step)
|
||||
if redact_result is not None:
|
||||
if redact_result.mode == "text_only":
|
||||
screenshot_for_llm = None # don't send image, save tokens
|
||||
scene_text = redact_result.redacted_text
|
||||
else:
|
||||
screenshot_for_llm = redact_result.redacted_image or screenshot
|
||||
scene_text = redact_result.redacted_text
|
||||
|
||||
if redact_result.findings:
|
||||
logger.info(
|
||||
"PII redacted: %d findings (%s mode, %.0fms)",
|
||||
len(redact_result.findings),
|
||||
redact_result.mode,
|
||||
redact_result.processing_ms,
|
||||
)
|
||||
|
||||
# Get OCR scene graph (if not already from redaction)
|
||||
if not scene_text and self._perception:
|
||||
scene = await self._perception.perceive(screenshot)
|
||||
scene_text = scene.to_text_summary()
|
||||
|
||||
@@ -252,7 +325,7 @@ class HybridPlanner:
|
||||
logger.debug("Memory context retrieval failed")
|
||||
|
||||
action = await self._llm.plan_action(
|
||||
screenshot,
|
||||
screenshot_for_llm,
|
||||
task,
|
||||
step,
|
||||
history,
|
||||
@@ -273,6 +346,37 @@ class HybridPlanner:
|
||||
|
||||
return PlanResult(action=action, source="remote_llm")
|
||||
|
||||
async def _redact_for_cloud(
|
||||
self, screenshot: bytes, step: int,
|
||||
) -> Optional[RedactResult]:
|
||||
"""Redact PII from screenshot before sending to cloud LLM.
|
||||
|
||||
Returns None if redaction is disabled, NPU Daemon unavailable,
|
||||
or an error occurs (falls back to sending original).
|
||||
"""
|
||||
if not self._privacy_redact_enabled or not self._npu:
|
||||
return None
|
||||
|
||||
if not await self._npu.is_available():
|
||||
logger.debug("NPU Daemon unavailable, sending original to cloud")
|
||||
return None
|
||||
|
||||
try:
|
||||
result = await self._npu.redact_image(
|
||||
screenshot, self._privacy_redact_types,
|
||||
)
|
||||
self.metrics.privacy_redactions += 1
|
||||
self.metrics.privacy_findings_total += len(result.findings)
|
||||
if result.mode == "text_only":
|
||||
self.metrics.privacy_text_only_mode += 1
|
||||
else:
|
||||
self.metrics.privacy_image_mode += 1
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.warning("PII redaction failed, sending original: %s", e)
|
||||
self._npu.reset_availability()
|
||||
return None
|
||||
|
||||
async def _should_use_local(self, task: str, history: list[str]) -> bool:
|
||||
"""Decide whether this task is simple enough for local RKLLM.
|
||||
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
"""NPU Daemon client for KVM Agent — PII redaction before cloud LLM calls.
|
||||
|
||||
Calls the centralized NPU Daemon (port 8004) for:
|
||||
- Screenshot PII redaction (P0 priority)
|
||||
- OCR analysis (P2 priority)
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RedactResult:
|
||||
"""Result of image PII redaction."""
|
||||
|
||||
redacted_image: Optional[bytes] # JPEG bytes (None if text_only mode)
|
||||
redacted_text: str
|
||||
findings: list # [{pii_type, text, x, y, w, h}]
|
||||
processing_ms: float
|
||||
mode: str # "image" or "text_only"
|
||||
|
||||
|
||||
class NpuClient:
|
||||
"""HTTP client for NPU Daemon."""
|
||||
|
||||
def __init__(self, base_url: str = "http://localhost:8004"):
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._available: Optional[bool] = None
|
||||
|
||||
async def is_available(self) -> bool:
|
||||
"""Check if NPU Daemon is reachable (cached)."""
|
||||
if self._available is not None:
|
||||
return self._available
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=3.0) as client:
|
||||
resp = await client.get(f"{self._base_url}/api/v1/health")
|
||||
self._available = resp.status_code == 200
|
||||
except Exception:
|
||||
self._available = False
|
||||
return self._available
|
||||
|
||||
def reset_availability(self):
|
||||
"""Reset cached availability (call on failure for retry)."""
|
||||
self._available = None
|
||||
|
||||
async def redact_image(
|
||||
self,
|
||||
image_bytes: bytes,
|
||||
redact_types: Optional[list[str]] = None,
|
||||
) -> RedactResult:
|
||||
"""Call /api/v1/privacy/redact-image with P0 priority.
|
||||
|
||||
Args:
|
||||
image_bytes: JPEG screenshot bytes.
|
||||
redact_types: PII types to redact. Defaults to common types.
|
||||
|
||||
Returns:
|
||||
RedactResult with redacted image/text and findings.
|
||||
|
||||
Raises:
|
||||
httpx.HTTPError: On network/HTTP errors.
|
||||
"""
|
||||
import base64
|
||||
|
||||
types = redact_types or ["id_card", "phone", "bank_card", "email"]
|
||||
|
||||
async with httpx.AsyncClient(timeout=15.0) as client:
|
||||
resp = await client.post(
|
||||
f"{self._base_url}/api/v1/privacy/redact-image",
|
||||
files={"image": ("screenshot.jpg", image_bytes, "image/jpeg")},
|
||||
data={"redact_types": json.dumps(types)},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
# Decode base64 image if present
|
||||
redacted_image = None
|
||||
if data.get("redacted_image"):
|
||||
redacted_image = base64.b64decode(data["redacted_image"])
|
||||
|
||||
return RedactResult(
|
||||
redacted_image=redacted_image,
|
||||
redacted_text=data.get("redacted_text", ""),
|
||||
findings=data.get("findings", []),
|
||||
processing_ms=data.get("processing_ms", 0.0),
|
||||
mode=data.get("mode", "text_only"),
|
||||
)
|
||||
|
||||
async def ocr_analyze(
|
||||
self,
|
||||
image_bytes: bytes,
|
||||
priority: str = "p2",
|
||||
) -> dict:
|
||||
"""Call /api/v1/ocr/analyze.
|
||||
|
||||
Returns: {text, regions, processing_ms}
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=15.0) as client:
|
||||
resp = await client.post(
|
||||
f"{self._base_url}/api/v1/ocr/analyze",
|
||||
files={"image": ("image.jpg", image_bytes, "image/jpeg")},
|
||||
data={"priority": priority},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
@@ -0,0 +1,282 @@
|
||||
"""Tests for NpuClient — NPU Daemon HTTP client for PII redaction."""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from kvm_agent.npu_client import NpuClient, RedactResult
|
||||
|
||||
|
||||
class TestNpuClient:
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_available_caches_result(self):
|
||||
client = NpuClient("http://localhost:8004")
|
||||
client._available = True
|
||||
assert await client.is_available() is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_available_false_on_error(self):
|
||||
client = NpuClient("http://localhost:8004")
|
||||
with patch("kvm_agent.npu_client.httpx.AsyncClient") as mock_cls:
|
||||
mock_http = AsyncMock()
|
||||
mock_http.get = AsyncMock(side_effect=Exception("connection refused"))
|
||||
mock_http.__aenter__ = AsyncMock(return_value=mock_http)
|
||||
mock_http.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_cls.return_value = mock_http
|
||||
|
||||
result = await client.is_available()
|
||||
assert result is False
|
||||
|
||||
def test_reset_availability(self):
|
||||
client = NpuClient("http://localhost:8004")
|
||||
client._available = True
|
||||
client.reset_availability()
|
||||
assert client._available is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_image_text_only_mode(self):
|
||||
client = NpuClient("http://localhost:8004")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"redacted_image": None,
|
||||
"redacted_text": "联系电话: 138****5678",
|
||||
"findings": [
|
||||
{"pii_type": "phone", "text": "138****5678",
|
||||
"x": 10, "y": 20, "w": 100, "h": 30}
|
||||
],
|
||||
"processing_ms": 85.2,
|
||||
"mode": "text_only",
|
||||
}
|
||||
|
||||
with patch("kvm_agent.npu_client.httpx.AsyncClient") as mock_cls:
|
||||
mock_http = AsyncMock()
|
||||
mock_http.post = AsyncMock(return_value=mock_response)
|
||||
mock_http.__aenter__ = AsyncMock(return_value=mock_http)
|
||||
mock_http.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_cls.return_value = mock_http
|
||||
|
||||
result = await client.redact_image(b"\xff\xd8\xff" * 100)
|
||||
|
||||
assert isinstance(result, RedactResult)
|
||||
assert result.mode == "text_only"
|
||||
assert result.redacted_image is None
|
||||
assert "138****5678" in result.redacted_text
|
||||
assert len(result.findings) == 1
|
||||
assert result.processing_ms == 85.2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_image_with_image_mode(self):
|
||||
import base64
|
||||
|
||||
client = NpuClient("http://localhost:8004")
|
||||
fake_jpeg = b"\xff\xd8\xff\xe0" + b"\x00" * 100
|
||||
b64_jpeg = base64.b64encode(fake_jpeg).decode()
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"redacted_image": b64_jpeg,
|
||||
"redacted_text": "身份证号 [REDACTED]",
|
||||
"findings": [
|
||||
{"pii_type": "id_card", "text": "1101**********1234",
|
||||
"x": 50, "y": 100, "w": 200, "h": 30}
|
||||
],
|
||||
"processing_ms": 120.5,
|
||||
"mode": "image",
|
||||
}
|
||||
|
||||
with patch("kvm_agent.npu_client.httpx.AsyncClient") as mock_cls:
|
||||
mock_http = AsyncMock()
|
||||
mock_http.post = AsyncMock(return_value=mock_response)
|
||||
mock_http.__aenter__ = AsyncMock(return_value=mock_http)
|
||||
mock_http.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_cls.return_value = mock_http
|
||||
|
||||
result = await client.redact_image(
|
||||
b"\xff\xd8\xff" * 100,
|
||||
redact_types=["id_card"],
|
||||
)
|
||||
|
||||
assert result.mode == "image"
|
||||
assert result.redacted_image is not None
|
||||
assert len(result.redacted_image) > 0
|
||||
assert len(result.findings) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_image_no_findings(self):
|
||||
client = NpuClient("http://localhost:8004")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"redacted_image": None,
|
||||
"redacted_text": "Hello World",
|
||||
"findings": [],
|
||||
"processing_ms": 60.0,
|
||||
"mode": "text_only",
|
||||
}
|
||||
|
||||
with patch("kvm_agent.npu_client.httpx.AsyncClient") as mock_cls:
|
||||
mock_http = AsyncMock()
|
||||
mock_http.post = AsyncMock(return_value=mock_response)
|
||||
mock_http.__aenter__ = AsyncMock(return_value=mock_http)
|
||||
mock_http.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_cls.return_value = mock_http
|
||||
|
||||
result = await client.redact_image(b"\xff\xd8\xff" * 100)
|
||||
assert result.mode == "text_only"
|
||||
assert len(result.findings) == 0
|
||||
assert result.redacted_text == "Hello World"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ocr_analyze(self):
|
||||
client = NpuClient("http://localhost:8004")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"text": "Hello World",
|
||||
"regions": [{"text": "Hello", "x": 10, "y": 10, "w": 50, "h": 20, "confidence": 0.95}],
|
||||
"processing_ms": 55.0,
|
||||
}
|
||||
|
||||
with patch("kvm_agent.npu_client.httpx.AsyncClient") as mock_cls:
|
||||
mock_http = AsyncMock()
|
||||
mock_http.post = AsyncMock(return_value=mock_response)
|
||||
mock_http.__aenter__ = AsyncMock(return_value=mock_http)
|
||||
mock_http.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_cls.return_value = mock_http
|
||||
|
||||
result = await client.ocr_analyze(b"\xff\xd8\xff" * 100, "p2")
|
||||
assert result["text"] == "Hello World"
|
||||
assert len(result["regions"]) == 1
|
||||
|
||||
|
||||
class TestHybridPlannerRedaction:
|
||||
"""Test PII redaction integration in HybridPlanner."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_for_cloud_disabled(self):
|
||||
"""When privacy_redact_enabled=False, no redaction happens."""
|
||||
from kvm_agent.hybrid_planner import HybridPlanner
|
||||
from kvm_agent.template_store import TemplateStore
|
||||
|
||||
llm = AsyncMock()
|
||||
llm.plan_action = AsyncMock(return_value=MagicMock(
|
||||
type="click", x=0.5, y=0.5, reason="test",
|
||||
))
|
||||
|
||||
store = AsyncMock(spec=TemplateStore)
|
||||
store.find_template = AsyncMock(return_value=None)
|
||||
|
||||
hp = HybridPlanner(
|
||||
llm, store,
|
||||
privacy_redact_enabled=False,
|
||||
)
|
||||
result = await hp._redact_for_cloud(b"\xff" * 100, 0)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_for_cloud_npu_unavailable(self):
|
||||
"""When NPU Daemon is unavailable, returns None (send original)."""
|
||||
from kvm_agent.hybrid_planner import HybridPlanner
|
||||
from kvm_agent.template_store import TemplateStore
|
||||
|
||||
llm = AsyncMock()
|
||||
store = AsyncMock(spec=TemplateStore)
|
||||
store.find_template = AsyncMock(return_value=None)
|
||||
|
||||
npu = AsyncMock(spec=NpuClient)
|
||||
npu.is_available = AsyncMock(return_value=False)
|
||||
|
||||
hp = HybridPlanner(
|
||||
llm, store,
|
||||
npu_client=npu,
|
||||
privacy_redact_enabled=True,
|
||||
)
|
||||
result = await hp._redact_for_cloud(b"\xff" * 100, 0)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_for_cloud_success(self):
|
||||
"""Successful redaction returns RedactResult and updates metrics."""
|
||||
from kvm_agent.hybrid_planner import HybridPlanner
|
||||
from kvm_agent.template_store import TemplateStore
|
||||
|
||||
llm = AsyncMock()
|
||||
store = AsyncMock(spec=TemplateStore)
|
||||
store.find_template = AsyncMock(return_value=None)
|
||||
|
||||
mock_result = RedactResult(
|
||||
redacted_image=None,
|
||||
redacted_text="脱敏后文本",
|
||||
findings=[{"pii_type": "phone", "text": "138****5678"}],
|
||||
processing_ms=85.0,
|
||||
mode="text_only",
|
||||
)
|
||||
|
||||
npu = AsyncMock(spec=NpuClient)
|
||||
npu.is_available = AsyncMock(return_value=True)
|
||||
npu.redact_image = AsyncMock(return_value=mock_result)
|
||||
|
||||
hp = HybridPlanner(
|
||||
llm, store,
|
||||
npu_client=npu,
|
||||
privacy_redact_enabled=True,
|
||||
)
|
||||
result = await hp._redact_for_cloud(b"\xff" * 100, 0)
|
||||
|
||||
assert result is not None
|
||||
assert result.mode == "text_only"
|
||||
assert hp.metrics.privacy_redactions == 1
|
||||
assert hp.metrics.privacy_findings_total == 1
|
||||
assert hp.metrics.privacy_text_only_mode == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_for_cloud_error_fallback(self):
|
||||
"""On NPU error, returns None and resets availability cache."""
|
||||
from kvm_agent.hybrid_planner import HybridPlanner
|
||||
from kvm_agent.template_store import TemplateStore
|
||||
|
||||
llm = AsyncMock()
|
||||
store = AsyncMock(spec=TemplateStore)
|
||||
store.find_template = AsyncMock(return_value=None)
|
||||
|
||||
npu = AsyncMock(spec=NpuClient)
|
||||
npu.is_available = AsyncMock(return_value=True)
|
||||
npu.redact_image = AsyncMock(side_effect=Exception("connection failed"))
|
||||
npu.reset_availability = MagicMock()
|
||||
|
||||
hp = HybridPlanner(
|
||||
llm, store,
|
||||
npu_client=npu,
|
||||
privacy_redact_enabled=True,
|
||||
)
|
||||
result = await hp._redact_for_cloud(b"\xff" * 100, 0)
|
||||
|
||||
assert result is None
|
||||
npu.reset_availability.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metrics_include_privacy_stats(self):
|
||||
"""PlannerMetrics.to_dict() includes privacy fields."""
|
||||
from kvm_agent.hybrid_planner import PlannerMetrics
|
||||
|
||||
m = PlannerMetrics(
|
||||
privacy_redactions=5,
|
||||
privacy_findings_total=12,
|
||||
privacy_text_only_mode=3,
|
||||
privacy_image_mode=2,
|
||||
)
|
||||
d = m.to_dict()
|
||||
assert d["privacy_redactions"] == 5
|
||||
assert d["privacy_findings_total"] == 12
|
||||
assert d["privacy_text_only_mode"] == 3
|
||||
assert d["privacy_image_mode"] == 2
|
||||
Reference in New Issue
Block a user