mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-17 11:14:19 +00:00
feat(telemetry): payload v2 with feature usage counters and instance heartbeat
Per-query events now carry event_type='query' and a features JSON object: - tool_calls by source (native/plugin/mcp/skill) via ToolManager - tool_call_rounds, kb usage (count/engine plugins/retrieved entries) via local-agent - sandbox execs/errors via BoxService - activated_skills and bound mcp_servers snapshots New instance_heartbeat event (startup + daily) reports anonymous instance profile: deploy platform, database/vdb kind, box backend/availability, adapter type names, and resource counts. Respects space.disable_telemetry. All collection helpers are defensive and never break the pipeline. Verified: ruff, 37 telemetry unit tests (13 new), 504 box/provider/pipeline tests.
This commit is contained in:
@@ -0,0 +1,92 @@
|
||||
"""Unit tests for telemetry feature counters (pkg/telemetry/features.py)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from importlib import import_module
|
||||
|
||||
|
||||
def get_features_module():
|
||||
return import_module('langbot.pkg.telemetry.features')
|
||||
|
||||
|
||||
class FakeQuery:
|
||||
def __init__(self):
|
||||
self.variables = {}
|
||||
|
||||
|
||||
class TestIncrement:
|
||||
def test_increment_nested_counter(self):
|
||||
features = get_features_module()
|
||||
q = FakeQuery()
|
||||
features.increment(q, 'tool_calls', 'native')
|
||||
features.increment(q, 'tool_calls', 'native')
|
||||
features.increment(q, 'tool_calls', 'mcp')
|
||||
assert q.variables[features.FEATURES_KEY]['tool_calls'] == {'native': 2, 'mcp': 1}
|
||||
|
||||
def test_increment_flat_counter(self):
|
||||
features = get_features_module()
|
||||
q = FakeQuery()
|
||||
features.increment(q, 'something')
|
||||
features.increment(q, 'something', amount=2)
|
||||
assert q.variables[features.FEATURES_KEY]['something'] == 3
|
||||
|
||||
def test_increment_never_raises_on_broken_query(self):
|
||||
features = get_features_module()
|
||||
|
||||
class Broken:
|
||||
@property
|
||||
def variables(self):
|
||||
raise RuntimeError('boom')
|
||||
|
||||
# Must not raise
|
||||
features.increment(Broken(), 'tool_calls', 'native')
|
||||
|
||||
def test_set_value(self):
|
||||
features = get_features_module()
|
||||
q = FakeQuery()
|
||||
features.set_value(q, 'tool_call_rounds', 5)
|
||||
assert q.variables[features.FEATURES_KEY]['tool_call_rounds'] == 5
|
||||
|
||||
|
||||
class TestCollectFeatures:
|
||||
def test_collect_empty(self):
|
||||
features = get_features_module()
|
||||
q = FakeQuery()
|
||||
assert features.collect_features(q) == {}
|
||||
|
||||
def test_collect_combines_counters_and_snapshots(self):
|
||||
features = get_features_module()
|
||||
q = FakeQuery()
|
||||
features.increment(q, 'sandbox', 'execs')
|
||||
features.set_value(q, 'kb', {'kb_count': 2, 'engine_plugins': ['builtin'], 'retrieved_entries': 7})
|
||||
q.variables['_activated_skills'] = {'pdf-tools': {}, 'a-skill': {}}
|
||||
q.variables['_pipeline_bound_mcp_servers'] = ['srv1', 'srv2']
|
||||
|
||||
result = features.collect_features(q)
|
||||
assert result['sandbox'] == {'execs': 1}
|
||||
assert result['kb']['kb_count'] == 2
|
||||
assert result['activated_skills'] == ['a-skill', 'pdf-tools'] # sorted
|
||||
assert result['mcp_servers'] == ['srv1', 'srv2']
|
||||
|
||||
def test_collect_omits_mcp_when_all_enabled(self):
|
||||
"""None means 'all enabled' and is not reported."""
|
||||
features = get_features_module()
|
||||
q = FakeQuery()
|
||||
q.variables['_pipeline_bound_mcp_servers'] = None
|
||||
assert 'mcp_servers' not in features.collect_features(q)
|
||||
|
||||
def test_collect_drops_non_json_serializable(self):
|
||||
features = get_features_module()
|
||||
q = FakeQuery()
|
||||
features.set_value(q, 'good', 1)
|
||||
features.set_value(q, 'bad', object())
|
||||
result = features.collect_features(q)
|
||||
assert result == {'good': 1}
|
||||
|
||||
def test_collect_is_json_serializable(self):
|
||||
import json
|
||||
|
||||
features = get_features_module()
|
||||
q = FakeQuery()
|
||||
features.increment(q, 'tool_calls', 'skill')
|
||||
json.dumps(features.collect_features(q))
|
||||
@@ -0,0 +1,104 @@
|
||||
"""Unit tests for telemetry heartbeat payload (pkg/telemetry/heartbeat.py)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
from importlib import import_module
|
||||
|
||||
|
||||
def get_heartbeat_module():
|
||||
return import_module('langbot.pkg.telemetry.heartbeat')
|
||||
|
||||
|
||||
def make_app():
|
||||
ap = Mock()
|
||||
ap.instance_config = Mock()
|
||||
ap.instance_config.data = {
|
||||
'database': {'use': 'postgresql'},
|
||||
'vdb': {'use': 'chroma'},
|
||||
'box': {'enabled': True, 'backend': 'nsjail'},
|
||||
}
|
||||
|
||||
# persistence counts
|
||||
result = Mock()
|
||||
result.scalar.return_value = 3
|
||||
ap.persistence_mgr = Mock()
|
||||
ap.persistence_mgr.execute_async = AsyncMock(return_value=result)
|
||||
|
||||
# box service
|
||||
ap.box_service = Mock()
|
||||
ap.box_service.enabled = True
|
||||
ap.box_service.available = False
|
||||
ap.box_service.shares_filesystem_with_box = False
|
||||
|
||||
# platform manager with one enabled bot
|
||||
bot = Mock()
|
||||
bot.enable = True
|
||||
bot.adapter = Mock()
|
||||
bot.adapter.__class__.__name__ = 'TelegramAdapter'
|
||||
ap.platform_mgr = Mock()
|
||||
ap.platform_mgr.bots = [bot]
|
||||
|
||||
# plugin connector
|
||||
ap.plugin_connector = Mock()
|
||||
ap.plugin_connector.list_plugins = AsyncMock(return_value=[{}, {}])
|
||||
|
||||
# skills
|
||||
ap.skill_mgr = Mock()
|
||||
ap.skill_mgr.skills = {'a': {}, 'b': {}, 'c': {}}
|
||||
|
||||
return ap
|
||||
|
||||
|
||||
class TestBuildHeartbeatPayload:
|
||||
@pytest.mark.asyncio
|
||||
async def test_payload_shape(self):
|
||||
heartbeat = get_heartbeat_module()
|
||||
ap = make_app()
|
||||
payload = await heartbeat.build_heartbeat_payload(ap)
|
||||
|
||||
assert payload['event_type'] == 'instance_heartbeat'
|
||||
assert payload['query_id'] == ''
|
||||
assert 'timestamp' in payload
|
||||
f = payload['features']
|
||||
assert f['database'] == 'postgresql'
|
||||
assert f['vdb'] == 'chroma'
|
||||
assert f['box'] == {
|
||||
'enabled': True,
|
||||
'available': False,
|
||||
'backend': 'nsjail',
|
||||
'shares_fs': False,
|
||||
}
|
||||
assert f['adapters'] == ['TelegramAdapter']
|
||||
assert f['bot_count'] == 1
|
||||
assert f['plugin_count'] == 2
|
||||
assert f['skill_count'] == 3
|
||||
assert f['pipeline_count'] == 3
|
||||
assert f['mcp_server_count'] == 3
|
||||
assert f['knowledge_base_count'] == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_payload_is_json_serializable(self):
|
||||
heartbeat = get_heartbeat_module()
|
||||
payload = await heartbeat.build_heartbeat_payload(make_app())
|
||||
json.dumps(payload)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_failure_yields_minus_one(self):
|
||||
heartbeat = get_heartbeat_module()
|
||||
ap = make_app()
|
||||
ap.persistence_mgr.execute_async = AsyncMock(side_effect=RuntimeError('db down'))
|
||||
payload = await heartbeat.build_heartbeat_payload(ap)
|
||||
assert payload['features']['pipeline_count'] == -1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_user_content_fields(self):
|
||||
"""The heartbeat must never carry message content / credentials keys."""
|
||||
heartbeat = get_heartbeat_module()
|
||||
payload = await heartbeat.build_heartbeat_payload(make_app())
|
||||
flat = json.dumps(payload).lower()
|
||||
for forbidden in ('api_key', 'password', 'token', 'message_content'):
|
||||
assert forbidden not in flat
|
||||
Reference in New Issue
Block a user