diff --git a/docs/design/openviking-usage-reporter-sink-design.md b/docs/design/openviking-usage-reporter-sink-design.md new file mode 100644 index 0000000000..0e32b0186c --- /dev/null +++ b/docs/design/openviking-usage-reporter-sink-design.md @@ -0,0 +1,432 @@ +# OpenViking Usage Reporter / Sink 技术方案 + +## 1. 背景 + +OpenViking 在 session commit 时会收到完整的 session messages。Agent runtime 调用 tool 后,会在 session messages 中留下 tool parts。部分 tool parts 可以表达某个记忆文件被检索、读取、注入等行为。 + +当前需求是:OpenViking 需要从 session 中解析出这些使用行为,并把结构化数据回流到不同下游。 + +不同部署形态的下游不同: + +- 火山商业化版本:回流到指定 Kafka。 +- 私有化部署:回流到客户自己的 Kafka、Webhook、日志系统或数据库。 +- 开源本地版本:可以不回流,或写本地 JSONL 方便调试。 + +因此 OpenViking 开源版不应该硬编码 Kafka,也不应该 import Kafka 依赖,而是提供一套通用的 Usage Reporter / Sink 扩展机制。 + +## 2. 业界做法 + +这种模式是开源基础设施里常见的设计。 + +OpenTelemetry Collector 把数据链路拆成 receivers、processors、exporters、pipelines。exporters 专门负责把数据发送到不同 backend 或 destination,例如 OTLP、Kafka、Prometheus、file 等。配置 exporter 本身不代表启用,需要在 pipeline 中声明。 + +参考:https://opentelemetry.io/docs/collector/configuration/ + +Vector 使用 sinks 概念,把 observability data 投递到不同目的地,例如 File、HTTP、Kafka、S3、ClickHouse、Prometheus remote write 等。 + +参考:https://vector.dev/docs/reference/configuration/sinks/ + +Fluent Bit 使用 Outputs 概念。官方定义里,Outputs 用来定义数据目的地,常见目的地包括远端服务、本地文件系统或标准接口,并且 Outputs 以 plugin 形式实现。 + +参考:https://docs.fluentbit.io/manual/data-pipeline/outputs + +所以 OpenViking 采用“核心定义事件 + 插件式 Sink 输出”是合理的。它的核心价值是把“事件产生”和“事件去哪”解耦。 + +## 3. 设计目标 + +- OpenViking 内核只定义 UsageEvent 标准结构。 +- OpenViking 内核负责从 session commit 中解析 UsageEvent。 +- OpenViking 内核不绑定 Kafka、Webhook、火山内部服务。 +- 开源版本不 import Kafka 依赖。 +- 商业化和私有化部署可以通过自定义 Sink 接入自己的数据系统。 +- Sink 失败默认不影响 session commit。 +- 默认不上报完整 session,只上报解析后的结构化事件,并保留 session evidence 指针。 + +## 4. 总体架构 + +数据流: + +```text +Agent runtime 调用 tool +-> session message 留下 tool part +-> client 上传 session 并 commit +-> OpenViking archive session +-> UsageExtractor 从 session messages 解析 UsageEvent +-> UsageReporter 分发 UsageEvent +-> UsageSink 写入目标系统 +``` + +模块拆分: + +```text +UsageExtractor:负责从 session 里解析事件 +UsageEvent:标准结构化事件 +UsageReporter:负责分发事件 +UsageSink:负责写入不同下游 +``` + +## 5. UsageEvent + +UsageEvent 是 OpenViking 内核和外部 Sink 之间的稳定协议。 + +示例: + +```json +{ + "schema_version": "v1", + "event_type": "memory.injected", + "memory_uri": "viking://user/test/memories/experiences/xxx.md", + "memory_type": "experience", + "account_id": "new", + "user_id": "test", + "session_id": "510bb5f9-4671-498e-adf4-27bb1b3691fe", + "archive_uri": "viking://user/test/sessions/510bb5f9/history/archive_001", + "task_id": "b174eb56-e7d4-4fee-98a6-c53c0ddf62ed", + "occurred_at": "2026-07-09T12:00:00Z", + "source": { + "tool_name": "read_experience", + "tool_status": "completed" + }, + "evidence": { + "message_index": 12, + "part_index": 0, + "tool_call_id": "call_xxx" + }, + "attributes": {} +} +``` + +默认只上报结构化事件,不上报完整 session 内容。 + +`memory_uri` 必须属于当前 `UsageContext.user_id` 对应的规范 experience +目录。客户端传入的其他用户 URI 不生成 UsageEvent。 + +## 6. UsageSink 机制 + +OpenViking 开源包只定义 Sink 抽象: + +```python +class UsageSink: + async def write(self, events: list[UsageEvent], context: UsageContext) -> None: + ... +``` + +可扩展 Sink: + +```text +kafka +webhook +customer_custom_sink +``` + +Kafka 不进入开源主包。商业化部署通过动态加载方式接入。 + +配置示例: + +```yaml +usage_reporter: + enabled: true + sinks: + - type: custom + class_path: volc_ov_usage.kafka_sink.KafkaUsageSink + config: + bootstrap_servers: kafka-1:9092,kafka-2:9092 + topic: ov_memory_usage +``` + +OpenViking 用 `importlib` 动态加载: + +```python +import importlib + +def load_class(class_path: str): + module_name, class_name = class_path.rsplit(".", 1) + module = importlib.import_module(module_name) + return getattr(module, class_name) +``` + +只有配置了该 Sink 时才 import 对应模块。开源版本不会 import Kafka,也不会安装 Kafka 依赖。 + +每个 Sink 的 `write()` 调用最多等待 5 秒。超时或异常只记录日志,不影响其他 Sink。Reporter 在应用生命周期内只创建一次,应用退出时调用 Sink 可选的 `close()` 方法。 + +## 7. 配置设计 + +默认关闭: + +```yaml +usage_reporter: + enabled: false +``` + +商业化 Kafka: + +```yaml +usage_reporter: + enabled: true + sinks: + - type: custom + class_path: volc_ov_usage.kafka_sink.KafkaUsageSink + config: + bootstrap_servers: kafka-1:9092,kafka-2:9092 + topic: ov_memory_usage +``` + +私有化 Webhook: + +```yaml +usage_reporter: + enabled: true + sinks: + - type: custom + class_path: customer_ov_usage.webhook_sink.WebhookUsageSink + config: + url: https://customer.example.com/openviking/usage +``` + +## 8. 对 OpenViking 的侵入 + +侵入点控制在 4 个地方。 + +1. 新增 config + +增加 `usage_reporter` 配置段。 + +2. 新增数据模型 + +增加 `UsageEvent`、`UsageContext`。 + +3. session commit 增加 hook + +在 session archive 成功后触发: + +```text +archive session success +-> usage extractor +-> usage reporter +``` + +4. 新增 reporter/sink 模块 + +新增通用扩展点和 custom sink 动态加载能力。 + +不侵入的地方: + +- 不改 `find/search` 语义。 +- 不改 `read` 语义。 +- 不把 Kafka 写死进 session commit。 +- 不默认上传完整 session。 +- 不强制写 MEMORY_FIELDS。 +- 不强制写 search_tags。 +- 不影响 snapshot。 +- Sink 调用具有 5 秒超时边界,失败不会中断 phase2。 + +整体侵入属于低到中等,核心主链路只增加一个旁路 hook。 + +## 9. 可靠性策略 + +失败处理: + +- Sink 成功:正常返回。 +- Sink 失败:记录日志,不影响 commit。 +- 多个 Sink:某个 Sink 失败不影响其他 Sink。 +- 事件带幂等 key,方便下游去重。 + +幂等 key: + +```text +schema_version ++ event_type ++ account_id ++ user_id ++ session_id ++ archive_uri ++ message_index ++ part_index ++ memory_uri +``` + +后续如果需要更高可靠性,可以增加本地 buffer / dead letter file,但不作为第一期必需能力。 + +## 10. 当前需求落地:记忆文件使用次数回流 + +### 10.1 使用次数定义 + +本需求里的“某个记忆文件被使用”,拆成两个事件。 + +`memory.recalled` 表示记忆文件被检索命中。 + +来源: + +```text +tool_name == search_experience +tool_status == completed +tool_output.results[].uri 包含该 memory uri +``` + +`memory.injected` 表示记忆文件被读取并注入 Agent 上下文。 + +来源: + +```text +tool_name == read_experience +tool_status == completed +tool_input.uri 或 tool_output.uri 指向该 memory uri +``` + +产品上如果只展示“使用次数”,默认使用 `memory.injected` 的计数。 + +### 10.2 当前需求的数据流 + +```text +Agent 调用 search_experience / read_experience +-> tool parts 留在 session messages +-> 用户 commit session +-> OpenViking archive session +-> MemoryUsageExtractor 解析 tool parts +-> 生成 memory.recalled / memory.injected events +-> UsageReporter 分发到配置的 Sink +``` + +### 10.3 MemoryUsageExtractor + +本期内置一个 extractor: + +```text +MemoryUsageExtractor +``` + +解析范围: + +```text +parts[].type == "tool" +tool_status == "completed" +``` + +识别 tool: + +```text +search_experience +read_experience +``` + +解析规则: + +```text +search_experience: + 从 tool_output.results[].uri 生成 memory.recalled + +read_experience: + 从 tool_input.uri 或 tool_output.uri 生成 memory.injected +``` + +事件的 `occurred_at` 优先取 tool part 所在 message 的 `created_at`,并统一 +转换为 UTC ISO 8601;消息时间缺失或非法时才使用 commit 解析时间。 + +不处理: + +- 普通 `find/search` API 调用。 +- 普通 `read` API 调用。 +- 非 completed tool part。 +- 非官方 tool name。 + +### 10.4 商业化版本如何回流 Kafka + +商业化版本不改 OpenViking 开源内核。 + +商业化部署额外提供包: + +```text +volc_ov_usage/ + kafka_sink.py +``` + +实现: + +```python +class KafkaUsageSink(UsageSink): + async def write(self, events, context): + ... +``` + +商业化镜像额外安装 Kafka 依赖,例如: + +```text +confluent-kafka +``` + +OpenViking 配置: + +```yaml +usage_reporter: + enabled: true + sinks: + - type: custom + class_path: volc_ov_usage.kafka_sink.KafkaUsageSink + config: + topic: ov_memory_usage +``` + +这样开源主包不 import Kafka,商业化版本通过 custom sink 注入 Kafka 能力。 + +### 10.5 私有化如何接入 + +私有化客户可以实现自己的 Sink: + +```python +class CustomerUsageSink(UsageSink): + async def write(self, events, context): + ... +``` + +配置: + +```yaml +usage_reporter: + enabled: true + sinks: + - type: custom + class_path: customer_ov_usage.CustomerUsageSink +``` + +客户可以写入: + +- 自己的 Kafka。 +- HTTP endpoint。 +- 数据库。 +- 日志系统。 +- 本地文件。 + +OpenViking 不需要知道目标系统细节。 + +### 10.6 开源版本默认能力 + +开源版默认关闭 usage reporting,不内置具体 Sink。部署方通过 custom sink 接入自己的目标系统。 + +## 11. 当前需求不在内核里做的事 + +本方案只解决“事件如何回流出去”。 + +不在 OpenViking 内核中强行完成: + +- 不直接计算控制台 7 天/30 天聚合。 +- 不直接把统计写进 MEMORY_FIELDS。 +- 不直接把统计写进 search_tags。 +- 不直接修改 experience 文件。 +- 不直接推完整 session。 + +聚合统计可以由 Kafka 消费端、私有化 Sink、控制台 BFF 或后续专门的 Usage Store 来做。 + +## 12. 工作项 + +1. 新增 `usage_reporter` config。 +2. 新增 `UsageEvent` / `UsageContext` 模型。 +3. 新增 `UsageExtractor` 抽象。 +4. 实现 `MemoryUsageExtractor`。 +5. 新增 `UsageSink` 抽象。 +6. 实现 custom sink 动态加载。 +7. 在 session archive 成功后挂接 UsageReporter。 +8. 增加 Sink 超时、best-effort 失败处理和日志。 +9. 增加 Reporter/Sink 生命周期关闭处理。 +10. 增加单测:tool parts -> UsageEvent。 +11. 增加集成测试:commit session -> custom test sink 写出事件。 +12. 增加文档:商业化和私有化 custom sink 接入方式。 diff --git a/examples/skills/ov-experience-memory/SKILL.md b/examples/skills/ov-experience-memory/SKILL.md new file mode 100644 index 0000000000..447406d3ee --- /dev/null +++ b/examples/skills/ov-experience-memory/SKILL.md @@ -0,0 +1,155 @@ +--- +name: ov-experience-memory +description: > + Use OpenViking experience memories during task execution. Search relevant + experiences with search_experience, read selected experiences with + read_experience, and leave standard tool parts in the committed session so + OpenViking can report recall and injection usage. +version: 2026.7.9 +tags: + - openviking + - experience-memory + - agent-memory + - usage-reporting +--- + +# OpenViking Experience Memory + +Use this skill when the current user request starts or continues an executable +task, especially tasks involving tools, files, code changes, data operations, +workflow decisions, or multi-step actions. + +Do not use this skill for casual chat, pure explanation, or one-off factual Q&A +that does not require operational guidance. + +## Runtime Contract + +The agent runtime must expose two tools with these exact names: + +- `search_experience` +- `read_experience` + +OpenViking usage reporting recognizes only completed tool parts with these exact +tool names. Calls to generic `find`, `search`, `read`, `ov_search`, or `ov_read` +do not count as experience recall or injection events. + +## Tool: search_experience + +Purpose: search reusable execution experiences from the OpenViking experience +library before assembling task context. + +Input schema: + +```json +{ + "query": "string", + "limit": 5, + "target_uri": "viking://user/memories/experiences/" +} +``` + +Output schema: + +```json +{ + "results": [ + { + "uri": "viking://user//memories/experiences/example.md", + "title": "example", + "score": 0.82, + "snippet": "Short summary or matched situation" + } + ] +} +``` + +Implementation: + +Call OpenViking `POST /api/v1/search/find` with `target_uri` set to the +current-user shorthand `viking://user/memories/experiences/`. OpenViking resolves +this shorthand against the authenticated request user. Return only canonical +experience memory URIs for that user; never hardcode `default` or another user ID. + +Usage reporting: + +A completed `search_experience` tool part is counted as an experience recall +event for every `results[].uri` value. + +## Tool: read_experience + +Purpose: read the full Markdown body of a selected experience and inject it into +the agent prompt as task execution guidance. + +Input schema: + +```json +{ + "uri": "viking://user//memories/experiences/example.md" +} +``` + +Output schema: + +```json +{ + "uri": "viking://user//memories/experiences/example.md", + "content": "Experience Markdown body" +} +``` + +Implementation: + +Call OpenViking `GET /api/v1/content/read?uri=` for the selected +experience URI. Always pass the canonical URI returned by `search_experience`; +do not construct a URI with a hardcoded user ID. The returned content should be +inserted into the prompt as operational guidance, not as user profile facts. + +Usage reporting: + +A completed `read_experience` tool part is counted as an experience injection +event for `tool_input.uri` or `tool_output.uri`. In this design, reading an +experience through `read_experience` means the experience was injected into the +prompt. + +## Recommended Flow + +1. When a task begins, build a short query from the latest user instruction, + current plan, active skill name, and important tool/environment context. +2. Call `search_experience` before final prompt assembly. +3. Review returned titles/snippets and select only experiences likely to affect + execution. +4. Call `read_experience` for selected experience URIs. +5. Inject the returned Markdown into the prompt under an explicit experience + section. +6. Continue task execution. +7. Commit the session normally. The committed session must include the + `search_experience` and `read_experience` tool parts so OpenViking can report + usage. + +## Prompt Injection Format + +Use a compact and explicit block: + +```text + +The following guidance was retrieved from prior task execution experience. +Use it as operational guidance. Do not treat it as user identity or preference. + + +...experience markdown... + + +``` + +## Commit Requirements + +The session committed to OpenViking must preserve tool parts with: + +- `tool_name` +- `tool_status` +- `tool_input` +- `tool_output` +- `tool_id` + +Only `tool_status == "completed"` is counted. Failed, cancelled, or skipped tool +parts are ignored by usage reporting. diff --git a/openviking/server/app.py b/openviking/server/app.py index 992bc736c9..f25026f847 100644 --- a/openviking/server/app.py +++ b/openviking/server/app.py @@ -217,14 +217,29 @@ def create_app( validate_server_config(config) - def _configure_session_tool_outputs(service_obj) -> None: # noqa: ANN001 + usage_reporter_unset = object() + usage_reporter = usage_reporter_unset + + def _get_usage_reporter(): # noqa: ANN202 + nonlocal usage_reporter + if usage_reporter is usage_reporter_unset: + from openviking.usage_reporter.config import build_usage_reporter + + usage_reporter = build_usage_reporter(config.usage_reporter) + return usage_reporter + + def _configure_session_runtime(service_obj) -> None: # noqa: ANN001 sessions = getattr(service_obj, "sessions", None) - setter = getattr(sessions, "set_tool_output_externalization_config", None) - if callable(setter): - setter(config.tool_output_externalization) + tool_output_setter = getattr(sessions, "set_tool_output_externalization_config", None) + if callable(tool_output_setter): + tool_output_setter(config.tool_output_externalization) + + usage_reporter_setter = getattr(sessions, "set_usage_reporter", None) + if callable(usage_reporter_setter): + usage_reporter_setter(_get_usage_reporter()) if service is not None: - _configure_session_tool_outputs(service) + _configure_session_runtime(service) async def _deferred_init(service, app, config): """Retained for tests that validate deferred-init callback behavior.""" @@ -239,7 +254,7 @@ async def lifespan(app: FastAPI): service = OpenVikingService() assert service is not None - _configure_session_tool_outputs(service) + _configure_session_runtime(service) set_service(service) from openviking.metrics.global_api import ( @@ -320,6 +335,8 @@ async def _oauth_gc_loop(store) -> None: # noqa: ANN001 logger.warning(f"OpenVikingService close cancelled during shutdown: {e}") except Exception as e: logger.warning(f"OpenVikingService close failed during shutdown: {e}") + if usage_reporter is not usage_reporter_unset and usage_reporter is not None: + await usage_reporter.close() app = FastAPI( title="OpenViking API", diff --git a/openviking/server/config.py b/openviking/server/config.py index dcbb95c779..7b3c840989 100644 --- a/openviking/server/config.py +++ b/openviking/server/config.py @@ -170,6 +170,26 @@ class UsageAuditConfig(BaseModel): model_config = {"extra": "forbid"} +class UsageReporterSinkConfig(BaseModel): + """Usage reporter sink configuration.""" + + type: Literal["custom"] = "custom" + class_path: Optional[str] = None + config: Dict[str, object] = Field(default_factory=dict) + + model_config = {"extra": "forbid"} + + +class UsageReporterConfig(BaseModel): + """Usage event reporter configuration.""" + + enabled: bool = False + extractors: List[Literal["memory_usage"]] = Field(default_factory=lambda: ["memory_usage"]) + sinks: List[UsageReporterSinkConfig] = Field(default_factory=list) + + model_config = {"extra": "forbid"} + + class TraceDumpBodyConfig(BaseModel): """HTTP body dump configuration. @@ -234,6 +254,7 @@ class ServerConfig(BaseModel): encryption_enabled: bool = False # Whether file-level AES encryption is enabled api_key_hashing_enabled: bool = False # Whether API key Argon2id hashing is enabled (default: false, rely on file encryption) observability: ObservabilityConfig = Field(default_factory=ObservabilityConfig) + usage_reporter: UsageReporterConfig = Field(default_factory=UsageReporterConfig) # Public-facing base URL emitted in MCP-issued upload instructions. See # ``openviking.server.mcp_endpoint._resolve_public_base_url`` for the full # resolution chain: env var > this field > X-Forwarded-Host/Proto > Host header diff --git a/openviking/service/session_service.py b/openviking/service/session_service.py index 9a7991d68e..81ab0d94f1 100644 --- a/openviking/service/session_service.py +++ b/openviking/service/session_service.py @@ -28,6 +28,7 @@ if TYPE_CHECKING: from openviking.session.compressor_v2 import SessionCompressorV2 + from openviking.usage_reporter import UsageReporter class SessionService: @@ -43,6 +44,7 @@ def __init__( self._viking_fs = viking_fs self._session_compressor = session_compressor self._tool_output_externalization_config = ToolOutputExternalizationConfig() + self._usage_reporter: Optional["UsageReporter"] = None def set_dependencies( self, @@ -61,6 +63,10 @@ def set_tool_output_externalization_config( """Set tool output externalization controls for newly created sessions.""" self._tool_output_externalization_config = config.model_copy(deep=True) + def set_usage_reporter(self, usage_reporter: Optional["UsageReporter"]) -> None: + """Set the usage reporter for newly created sessions.""" + self._usage_reporter = usage_reporter + def _ensure_initialized(self) -> None: """Ensure all dependencies are initialized.""" if not self._viking_fs: @@ -120,6 +126,7 @@ def session( session_id=session_id, session_uri=session_uri, tool_output_externalization_config=self._tool_output_externalization_config, + usage_reporter=self._usage_reporter, ) async def create( diff --git a/openviking/session/session.py b/openviking/session/session.py index aa63f47ddc..6e72cb33e9 100644 --- a/openviking/session/session.py +++ b/openviking/session/session.py @@ -46,6 +46,7 @@ from openviking.session.compressor_v2 import SessionCompressorV2 as SessionCompressor from openviking.storage import VikingDBManager from openviking.storage.viking_fs import VikingFS + from openviking.usage_reporter import UsageReporter logger = get_logger(__name__) @@ -361,6 +362,7 @@ def __init__( session_uri: Optional[str] = None, auto_commit_threshold: int = 8000, tool_output_externalization_config: Optional[ToolOutputExternalizationConfig] = None, + usage_reporter: Optional["UsageReporter"] = None, ): self._viking_fs = viking_fs self._vikingdb_manager = vikingdb_manager @@ -388,6 +390,7 @@ def __init__( if tool_output_externalization_config is not None else ToolOutputExternalizationConfig() ) + self._usage_reporter = usage_reporter logger.info(f"Session created: {self.session_id} for user {self.user}") @@ -1263,6 +1266,28 @@ async def commit_async( "trace_id": trace_id, } + async def _run_usage_reporting( + self, + *, + task_id: str, + archive_uri: str, + messages: List[Message], + ) -> list[Any]: + reporter = getattr(self, "_usage_reporter", None) + if reporter is None: + return [] + + from openviking.usage_reporter import UsageContext + + context = UsageContext( + account_id=self.ctx.account_id, + user_id=self.ctx.user.user_id, + session_id=self.session_id, + archive_uri=archive_uri, + task_id=task_id, + ) + return await reporter.extract_and_report(messages=messages, context=context) + @tracer("session.commit.phase2", ignore_result=True, ignore_args=True) async def _run_memory_extraction( self, @@ -1286,6 +1311,7 @@ async def _run_memory_extraction( request_wait_tracker = get_request_wait_tracker() memories_extracted: Dict[str, int] = {} + usage_events_extracted = 0 extracted_skill_results: list[dict] = [] active_count_updated = 0 memory_diff_uri: Optional[str] = None @@ -1333,6 +1359,13 @@ async def _run_memory_extraction( else "" ) extraction_messages = await self._hydrate_tool_outputs_for_extraction(messages) + usage_events_extracted = len( + await self._run_usage_reporting( + task_id=task_id, + archive_uri=archive_uri, + messages=extraction_messages, + ) + ) async def _run_archive_summary() -> None: if not working_memory_enabled: @@ -1636,6 +1669,7 @@ async def _run_execution_memory_extraction() -> Any: for item in extracted_skill_results if isinstance(item, dict) and (item.get("uri") or item.get("root_uri")) ], + "usage_events_extracted": usage_events_extracted, "active_count_updated": active_count_updated, "token_usage": { "llm": dict(self._meta.llm_token_usage), diff --git a/openviking/usage_reporter/__init__.py b/openviking/usage_reporter/__init__.py new file mode 100644 index 0000000000..7382cbbb45 --- /dev/null +++ b/openviking/usage_reporter/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Usage reporting extension points for OpenViking.""" + +from .extractors import MemoryUsageExtractor, UsageExtractor +from .models import UsageContext, UsageEvent +from .reporter import UsageReporter +from .sinks import UsageSink + +__all__ = [ + "MemoryUsageExtractor", + "UsageContext", + "UsageEvent", + "UsageExtractor", + "UsageReporter", + "UsageSink", +] diff --git a/openviking/usage_reporter/config.py b/openviking/usage_reporter/config.py new file mode 100644 index 0000000000..f1ea78cdc8 --- /dev/null +++ b/openviking/usage_reporter/config.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Build usage reporter instances from server config.""" + +from __future__ import annotations + +import importlib +from typing import Any + +from openviking.server.config import UsageReporterConfig, UsageReporterSinkConfig + +from .extractors import MemoryUsageExtractor +from .reporter import UsageReporter +from .sinks import UsageSink + + +def _load_class(class_path: str) -> type: + module_name, class_name = class_path.rsplit(".", 1) + module = importlib.import_module(module_name) + cls = getattr(module, class_name) + if not isinstance(cls, type): + raise TypeError(f"{class_path} does not resolve to a class") + return cls + + +def _build_sink(config: UsageReporterSinkConfig) -> UsageSink: + if config.type == "custom": + if not config.class_path: + raise ValueError("custom usage sink requires class_path") + cls = _load_class(config.class_path) + kwargs: dict[str, Any] = dict(config.config) + return cls(**kwargs) + + raise ValueError(f"Unsupported usage sink type: {config.type}") + + +def build_usage_reporter(config: UsageReporterConfig) -> UsageReporter | None: + if not config.enabled: + return None + + extractors = [] + for name in config.extractors: + if name == "memory_usage": + extractors.append(MemoryUsageExtractor()) + else: + raise ValueError(f"Unsupported usage extractor: {name}") + + sinks = [_build_sink(sink_config) for sink_config in config.sinks] + return UsageReporter( + extractors=extractors, + sinks=sinks, + ) diff --git a/openviking/usage_reporter/extractors.py b/openviking/usage_reporter/extractors.py new file mode 100644 index 0000000000..0afbdd2457 --- /dev/null +++ b/openviking/usage_reporter/extractors.py @@ -0,0 +1,183 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Usage event extractors.""" + +from __future__ import annotations + +import json +from datetime import datetime +from typing import Any, Iterable, Protocol + +from openviking.message import Message, ToolPart +from openviking.utils.time_utils import format_iso8601, parse_iso_datetime + +from .models import UsageContext, UsageEvent, utc_now_iso + + +class UsageExtractor(Protocol): + name: str + + async def extract( + self, + *, + messages: list[Message], + context: UsageContext, + ) -> list[UsageEvent]: ... + + +def _load_mapping(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return value + if isinstance(value, str) and value.strip(): + try: + parsed = json.loads(value) + except json.JSONDecodeError: + return {} + return parsed if isinstance(parsed, dict) else {} + return {} + + +def _is_experience_uri(uri: str, context: UsageContext) -> bool: + prefix = f"viking://user/{context.user_id}/memories/experiences/" + return uri.startswith(prefix) and bool(uri.removeprefix(prefix).strip("/")) + + +def _event_time(message: Message) -> str: + value = message.created_at + try: + if isinstance(value, datetime): + return format_iso8601(value) + if isinstance(value, str) and value.strip(): + return format_iso8601(parse_iso_datetime(value.strip())) + except (TypeError, ValueError): + pass + return utc_now_iso() + + +class MemoryUsageExtractor: + """Extract memory recalled/injected usage events from official tool parts.""" + + name = "memory_usage" + + async def extract( + self, + *, + messages: list[Message], + context: UsageContext, + ) -> list[UsageEvent]: + events: list[UsageEvent] = [] + for message_index, message in enumerate(messages): + for part_index, part in enumerate(message.parts): + if not isinstance(part, ToolPart): + continue + if part.tool_status != "completed": + continue + if part.tool_name == "search_experience": + events.extend( + self._extract_search_events( + part, + context=context, + message=message, + message_index=message_index, + part_index=part_index, + ) + ) + elif part.tool_name == "read_experience": + event = self._extract_read_event( + part, + context=context, + message=message, + message_index=message_index, + part_index=part_index, + ) + if event is not None: + events.append(event) + return events + + def _extract_search_events( + self, + part: ToolPart, + *, + context: UsageContext, + message: Message, + message_index: int, + part_index: int, + ) -> Iterable[UsageEvent]: + output = _load_mapping(part.tool_output) + results = output.get("results", []) + if not isinstance(results, list): + return [] + + events: list[UsageEvent] = [] + for result in results: + if not isinstance(result, dict): + continue + uri = str(result.get("uri") or "").strip() + if not uri or not _is_experience_uri(uri, context): + continue + events.append( + self._build_event( + event_type="memory.recalled", + memory_uri=uri, + part=part, + context=context, + message=message, + message_index=message_index, + part_index=part_index, + ) + ) + return events + + def _extract_read_event( + self, + part: ToolPart, + *, + context: UsageContext, + message: Message, + message_index: int, + part_index: int, + ) -> UsageEvent | None: + tool_input = part.tool_input if isinstance(part.tool_input, dict) else {} + output = _load_mapping(part.tool_output) + uri = str(tool_input.get("uri") or output.get("uri") or "").strip() + if not uri or not _is_experience_uri(uri, context): + return None + return self._build_event( + event_type="memory.injected", + memory_uri=uri, + part=part, + context=context, + message=message, + message_index=message_index, + part_index=part_index, + ) + + def _build_event( + self, + *, + event_type: str, + memory_uri: str, + part: ToolPart, + context: UsageContext, + message: Message, + message_index: int, + part_index: int, + ) -> UsageEvent: + return UsageEvent( + event_type=event_type, + memory_uri=memory_uri, + memory_type="experience", + account_id=context.account_id, + user_id=context.user_id, + session_id=context.session_id, + archive_uri=context.archive_uri, + task_id=context.task_id, + occurred_at=_event_time(message), + source={"tool_name": part.tool_name, "tool_status": part.tool_status}, + evidence={ + "message_index": message_index, + "message_id": message.id, + "part_index": part_index, + "tool_call_id": part.tool_id, + }, + ) diff --git a/openviking/usage_reporter/models.py b/openviking/usage_reporter/models.py new file mode 100644 index 0000000000..0d2553ae2a --- /dev/null +++ b/openviking/usage_reporter/models.py @@ -0,0 +1,42 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Usage reporter data models.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from typing import Any, Dict, Optional + + +def utc_now_iso() -> str: + return datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") + + +@dataclass(frozen=True) +class UsageContext: + account_id: str + user_id: str + session_id: str + archive_uri: str + task_id: Optional[str] = None + + +@dataclass(frozen=True) +class UsageEvent: + event_type: str + memory_uri: str + memory_type: str + account_id: str + user_id: str + session_id: str + archive_uri: str + occurred_at: str + schema_version: str = "v1" + task_id: Optional[str] = None + source: Dict[str, Any] = field(default_factory=dict) + evidence: Dict[str, Any] = field(default_factory=dict) + attributes: Dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> Dict[str, Any]: + return asdict(self) diff --git a/openviking/usage_reporter/reporter.py b/openviking/usage_reporter/reporter.py new file mode 100644 index 0000000000..ba257b64d6 --- /dev/null +++ b/openviking/usage_reporter/reporter.py @@ -0,0 +1,86 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Usage reporter dispatcher.""" + +from __future__ import annotations + +import asyncio +import inspect +from dataclasses import dataclass, field +from typing import Iterable + +from openviking_cli.utils import get_logger + +from .extractors import UsageExtractor +from .models import UsageContext, UsageEvent +from .sinks import UsageSink + +logger = get_logger(__name__) + + +@dataclass +class UsageReporter: + extractors: list[UsageExtractor] = field(default_factory=list) + sinks: list[UsageSink] = field(default_factory=list) + sink_timeout_seconds: float = 5.0 + + async def extract( + self, + *, + messages, + context: UsageContext, + ) -> list[UsageEvent]: + events: list[UsageEvent] = [] + for extractor in self.extractors: + try: + events.extend(await extractor.extract(messages=messages, context=context)) + except Exception: + logger.exception("Usage extractor failed: %s", getattr(extractor, "name", "")) + return events + + async def report( + self, + *, + events: Iterable[UsageEvent], + context: UsageContext, + ) -> None: + event_list = list(events) + if not event_list or not self.sinks: + return + for sink in self.sinks: + try: + await asyncio.wait_for( + sink.write(events=event_list, context=context), + timeout=self.sink_timeout_seconds, + ) + except TimeoutError: + logger.warning( + "Usage sink timed out after %.1fs: %s", + self.sink_timeout_seconds, + type(sink).__name__, + ) + except Exception: + logger.exception("Usage sink failed: %s", type(sink).__name__) + + async def extract_and_report(self, *, messages, context: UsageContext) -> list[UsageEvent]: + events = await self.extract(messages=messages, context=context) + await self.report(events=events, context=context) + return events + + async def close(self) -> None: + for sink in self.sinks: + close = getattr(sink, "close", None) + if not callable(close): + continue + try: + result = close() + if inspect.isawaitable(result): + await asyncio.wait_for(result, timeout=self.sink_timeout_seconds) + except TimeoutError: + logger.warning( + "Usage sink close timed out after %.1fs: %s", + self.sink_timeout_seconds, + type(sink).__name__, + ) + except Exception: + logger.exception("Usage sink close failed: %s", type(sink).__name__) diff --git a/openviking/usage_reporter/sinks.py b/openviking/usage_reporter/sinks.py new file mode 100644 index 0000000000..04477fc584 --- /dev/null +++ b/openviking/usage_reporter/sinks.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Usage event sinks.""" + +from __future__ import annotations + +from typing import Protocol + +from .models import UsageContext, UsageEvent + + +class UsageSink(Protocol): + async def write(self, *, events: list[UsageEvent], context: UsageContext) -> None: ... diff --git a/tests/session/test_session_usage_reporter.py b/tests/session/test_session_usage_reporter.py new file mode 100644 index 0000000000..149f9800c7 --- /dev/null +++ b/tests/session/test_session_usage_reporter.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +import pytest + +from openviking.message import ToolPart +from openviking.service.task_tracker import get_task_tracker +from openviking.usage_reporter import MemoryUsageExtractor, UsageReporter + + +async def _wait_for_task(task_id: str, timeout: float = 30.0) -> dict: + tracker = get_task_tracker() + for _ in range(int(timeout / 0.1)): + task = await tracker.get(task_id) + if task and task.status.value in ("completed", "failed"): + return task.to_dict() + import asyncio + + await asyncio.sleep(0.1) + raise TimeoutError(f"Task {task_id} did not complete within {timeout}s") + + +@pytest.mark.asyncio +async def test_commit_reports_memory_usage_events_to_sink(session): + reported_events = [] + + class RecordingSink: + async def write(self, *, events, context): + reported_events.extend(events) + + experience_uri = "viking://user/default/memories/experiences/no-order-exchange.md" + session._usage_reporter = UsageReporter( + extractors=[MemoryUsageExtractor()], + sinks=[RecordingSink()], + ) + session.add_message( + "user", + [ + ToolPart( + tool_id="call-read", + tool_name="read_experience", + tool_status="completed", + tool_input={"uri": experience_uri}, + ) + ], + ) + + result = await session.commit_async( + memory_policy={ + "self": {"enabled": False}, + "peer": {"enabled": False}, + "working_memory": {"enabled": False}, + } + ) + task = await _wait_for_task(result["task_id"]) + + assert task["status"] == "completed" + assert task["result"]["usage_events_extracted"] == 1 + assert "usage_events_reported" not in task["result"] + assert len(reported_events) == 1 + assert reported_events[0].event_type == "memory.injected" + assert reported_events[0].memory_uri == experience_uri + assert reported_events[0].session_id == session.session_id + assert reported_events[0].archive_uri == result["archive_uri"] + assert reported_events[0].task_id == result["task_id"] diff --git a/tests/unit/usage_reporter/test_memory_usage_extractor.py b/tests/unit/usage_reporter/test_memory_usage_extractor.py new file mode 100644 index 0000000000..ac3c9a806b --- /dev/null +++ b/tests/unit/usage_reporter/test_memory_usage_extractor.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +import pytest + +from openviking.message import Message, TextPart, ToolPart +from openviking.usage_reporter import ( + MemoryUsageExtractor, + UsageContext, +) + + +def _context() -> UsageContext: + return UsageContext( + account_id="new", + user_id="test", + session_id="session-1", + archive_uri="viking://user/test/sessions/session-1/history/archive_001", + task_id="task-1", + ) + + +@pytest.mark.asyncio +async def test_memory_usage_extractor_emits_recall_and_injection_events(): + experience_uri = "viking://user/test/memories/experiences/no-order-exchange.md" + messages = [ + Message( + id="msg-1", + role="user", + parts=[ + TextPart("我要处理无订单号换货"), + ToolPart( + tool_id="call-search", + tool_name="search_experience", + tool_status="completed", + tool_input={"query": "无订单号换货"}, + tool_output={"results": [{"uri": experience_uri}, {"uri": "viking://other"}]}, + ), + ], + ), + Message( + id="msg-2", + role="user", + parts=[ + ToolPart( + tool_id="call-read", + tool_name="read_experience", + tool_status="completed", + tool_input={"uri": experience_uri}, + tool_output="## Situation\n用户未提供订单号但要求换货。", + ) + ], + ), + Message( + id="msg-3", + role="user", + parts=[ + ToolPart( + tool_id="call-pending", + tool_name="read_experience", + tool_status="running", + tool_input={"uri": experience_uri}, + ) + ], + ), + ] + + events = await MemoryUsageExtractor().extract(messages=messages, context=_context()) + + assert [event.event_type for event in events] == ["memory.recalled", "memory.injected"] + assert [event.memory_uri for event in events] == [experience_uri, experience_uri] + assert events[0].source["tool_name"] == "search_experience" + assert events[0].evidence == { + "message_index": 0, + "message_id": "msg-1", + "part_index": 1, + "tool_call_id": "call-search", + } + assert events[1].source["tool_name"] == "read_experience" + assert events[1].evidence["part_index"] == 0 + + +@pytest.mark.asyncio +async def test_memory_usage_extractor_uses_message_time_for_event_time(): + experience_uri = "viking://user/test/memories/experiences/no-order-exchange.md" + messages = [ + Message( + id="msg-1", + role="user", + created_at="2026-07-10T20:30:40.123456+08:00", + parts=[ + ToolPart( + tool_id="call-read", + tool_name="read_experience", + tool_status="completed", + tool_input={"uri": experience_uri}, + ) + ], + ) + ] + + events = await MemoryUsageExtractor().extract(messages=messages, context=_context()) + + assert events[0].occurred_at == "2026-07-10T12:30:40.123Z" + + +@pytest.mark.asyncio +async def test_memory_usage_extractor_ignores_non_experience_memory_uris(): + messages = [ + Message( + id="msg-1", + role="user", + parts=[ + ToolPart( + tool_id="call-search", + tool_name="search_experience", + tool_status="completed", + tool_output={ + "results": [ + {"uri": "viking://user/default/memories/trajectories/a.md"}, + {"uri": "viking://user/default/memories/preferences/a.md"}, + ] + }, + ), + ToolPart( + tool_id="call-read", + tool_name="read_experience", + tool_status="completed", + tool_input={"uri": "viking://user/default/memories/trajectories/a.md"}, + ), + ], + ) + ] + + events = await MemoryUsageExtractor().extract(messages=messages, context=_context()) + + assert events == [] + + +@pytest.mark.asyncio +async def test_memory_usage_extractor_ignores_other_users_experience_uris(): + own_uri = "viking://user/test/memories/experiences/own.md" + other_uri = "viking://user/other/memories/experiences/other.md" + messages = [ + Message( + id="msg-1", + role="user", + parts=[ + ToolPart( + tool_id="call-search", + tool_name="search_experience", + tool_status="completed", + tool_output={"results": [{"uri": own_uri}, {"uri": other_uri}]}, + ), + ToolPart( + tool_id="call-read", + tool_name="read_experience", + tool_status="completed", + tool_input={"uri": other_uri}, + ), + ], + ) + ] + + events = await MemoryUsageExtractor().extract(messages=messages, context=_context()) + + assert [event.memory_uri for event in events] == [own_uri] diff --git a/tests/unit/usage_reporter/test_reporter.py b/tests/unit/usage_reporter/test_reporter.py new file mode 100644 index 0000000000..ba9f2e0278 --- /dev/null +++ b/tests/unit/usage_reporter/test_reporter.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +import asyncio + +from openviking.usage_reporter import UsageContext, UsageEvent, UsageReporter + + +def _context() -> UsageContext: + return UsageContext( + account_id="new", + user_id="test", + session_id="session-1", + archive_uri="viking://user/test/sessions/session-1/history/archive_001", + task_id="task-1", + ) + + +def _event() -> UsageEvent: + context = _context() + return UsageEvent( + event_type="memory.injected", + memory_uri="viking://user/test/memories/experiences/a.md", + memory_type="experience", + account_id=context.account_id, + user_id=context.user_id, + session_id=context.session_id, + archive_uri=context.archive_uri, + task_id=context.task_id, + occurred_at="2026-07-10T12:00:00Z", + ) + + +async def test_extractor_failure_is_ignored(): + class FailingExtractor: + name = "failing" + + async def extract(self, *, messages, context): + raise RuntimeError("extract failed") + + reporter = UsageReporter(extractors=[FailingExtractor()]) + + assert await reporter.extract(messages=[], context=_context()) == [] + + +async def test_sink_failure_does_not_stop_later_sinks(): + writes = [] + + class FailingSink: + async def write(self, *, events, context): + raise RuntimeError("sink failed") + + class RecordingSink: + async def write(self, *, events, context): + writes.extend(events) + + event = _event() + reporter = UsageReporter(sinks=[FailingSink(), RecordingSink()]) + + await reporter.report(events=[event], context=_context()) + + assert writes == [event] + + +async def test_sink_timeout_does_not_stop_later_sinks(): + writes = [] + + class HangingSink: + async def write(self, *, events, context): + await asyncio.sleep(1) + + class RecordingSink: + async def write(self, *, events, context): + writes.extend(events) + + event = _event() + reporter = UsageReporter( + sinks=[HangingSink(), RecordingSink()], + sink_timeout_seconds=0.01, + ) + + await reporter.report(events=[event], context=_context()) + + assert writes == [event] + + +async def test_close_calls_optional_sink_close(): + closed = [] + + class ClosableSink: + async def write(self, *, events, context): + return None + + async def close(self): + closed.append(True) + + reporter = UsageReporter(sinks=[ClosableSink()]) + + await reporter.close() + + assert closed == [True] diff --git a/tests/unit/usage_reporter/test_usage_reporter_config.py b/tests/unit/usage_reporter/test_usage_reporter_config.py new file mode 100644 index 0000000000..f8a1abcb12 --- /dev/null +++ b/tests/unit/usage_reporter/test_usage_reporter_config.py @@ -0,0 +1,137 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +import json +import sys +from contextlib import asynccontextmanager + +import pytest + +from openviking.server.app import create_app +from openviking.server.config import ServerConfig, UsageReporterConfig +from openviking.usage_reporter import UsageContext, UsageEvent +from openviking.usage_reporter.config import build_usage_reporter + + +@pytest.mark.asyncio +async def test_custom_sink_is_loaded_from_class_path(tmp_path, monkeypatch): + module_dir = tmp_path / "custom_sink_pkg" + module_dir.mkdir() + (module_dir / "__init__.py").write_text("", encoding="utf-8") + output_path = tmp_path / "custom-events.jsonl" + (module_dir / "sink.py").write_text( + """ +import json + +class CustomUsageSink: + def __init__(self, path): + self.path = path + + async def write(self, *, events, context): + with open(self.path, "a", encoding="utf-8") as f: + for event in events: + f.write(json.dumps(event.to_dict(), ensure_ascii=False) + "\\n") +""", + encoding="utf-8", + ) + monkeypatch.syspath_prepend(str(tmp_path)) + sys.modules.pop("custom_sink_pkg.sink", None) + + reporter = build_usage_reporter( + UsageReporterConfig( + enabled=True, + extractors=["memory_usage"], + sinks=[ + { + "type": "custom", + "class_path": "custom_sink_pkg.sink.CustomUsageSink", + "config": {"path": str(output_path)}, + } + ], + ) + ) + context = UsageContext( + account_id="new", + user_id="test", + session_id="session-1", + archive_uri="viking://user/test/sessions/session-1/history/archive_001", + task_id="task-1", + ) + event = UsageEvent( + event_type="memory.injected", + memory_uri="viking://user/default/memories/experiences/a.md", + memory_type="experience", + account_id="new", + user_id="test", + session_id="session-1", + archive_uri=context.archive_uri, + task_id="task-1", + occurred_at="2026-07-09T12:00:00Z", + ) + + await reporter.report(events=[event], context=context) + + payload = json.loads(output_path.read_text(encoding="utf-8").splitlines()[0]) + assert payload["event_type"] == "memory.injected" + assert payload["memory_uri"] == "viking://user/default/memories/experiences/a.md" + + +async def test_app_reuses_and_closes_usage_reporter(monkeypatch): + built_reporters = [] + assigned_reporters = [] + + class Reporter: + closed = False + + async def close(self): + self.closed = True + + reporter = Reporter() + + def build_reporter(config): + del config + built_reporters.append(reporter) + return reporter + + class Sessions: + def set_tool_output_externalization_config(self, config): + del config + + def set_usage_reporter(self, value): + assigned_reporters.append(value) + + class Service: + sessions = Sessions() + + class TaskTracker: + def start_cleanup_loop(self): + return None + + def stop_cleanup_loop(self): + return None + + async def initialize_runtime_state(app, service, config): + del app, service, config + + @asynccontextmanager + async def mcp_lifespan(): + yield + + monkeypatch.setattr( + "openviking.usage_reporter.config.build_usage_reporter", + build_reporter, + ) + monkeypatch.setattr( + "openviking.server.app._initialize_runtime_state", + initialize_runtime_state, + ) + monkeypatch.setattr("openviking.server.app.get_task_tracker", lambda: TaskTracker()) + monkeypatch.setattr("openviking.server.mcp_endpoint.mcp_lifespan", mcp_lifespan) + + app = create_app(config=ServerConfig(), service=Service()) + async with app.router.lifespan_context(app): + pass + + assert built_reporters == [reporter] + assert assigned_reporters == [reporter, reporter] + assert reporter.closed is True