mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 14:07:55 +00:00
Prathap PandGitHub
c2756964a7
* fix(channels): wire channel→agent handler and fix Telegram send pipeline * format code * add supported tests
113 lines
3.5 KiB
Python
113 lines
3.5 KiB
Python
"""Tests for the instrumented inference wrappers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncIterator, Sequence
|
|
from typing import Any, Dict, List
|
|
from unittest import mock
|
|
|
|
from openjarvis.core.events import EventBus, EventType
|
|
from openjarvis.core.types import Message, Role, TelemetryRecord
|
|
from openjarvis.engine._base import InferenceEngine
|
|
from openjarvis.telemetry.wrapper import instrumented_generate
|
|
|
|
|
|
class _StubEngine(InferenceEngine):
|
|
engine_id = "stub"
|
|
|
|
def __init__(self, response: Dict[str, Any] | None = None) -> None:
|
|
self._response = response or {
|
|
"content": "Hello",
|
|
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
|
|
}
|
|
|
|
def generate(
|
|
self, messages: Sequence[Message], *, model: str, **kwargs: Any
|
|
) -> Dict[str, Any]:
|
|
return self._response
|
|
|
|
async def stream(
|
|
self, messages: Sequence[Message], *, model: str, **kwargs: Any
|
|
) -> AsyncIterator[str]:
|
|
yield "Hello"
|
|
|
|
def list_models(self) -> List[str]:
|
|
return ["stub-model"]
|
|
|
|
def health(self) -> bool:
|
|
return True
|
|
|
|
|
|
class TestInstrumentedGenerate:
|
|
def test_calls_engine(self) -> None:
|
|
engine = _StubEngine()
|
|
bus = EventBus()
|
|
result = instrumented_generate(
|
|
engine,
|
|
[Message(role=Role.USER, content="Hi")],
|
|
model="m",
|
|
bus=bus,
|
|
)
|
|
assert result["content"] == "Hello"
|
|
|
|
def test_publishes_inference_events(self) -> None:
|
|
engine = _StubEngine()
|
|
bus = EventBus(record_history=True)
|
|
instrumented_generate(
|
|
engine,
|
|
[Message(role=Role.USER, content="Hi")],
|
|
model="m",
|
|
bus=bus,
|
|
)
|
|
event_types = [e.event_type for e in bus.history]
|
|
assert EventType.INFERENCE_START in event_types
|
|
assert EventType.INFERENCE_END in event_types
|
|
|
|
def test_publishes_telemetry_record(self) -> None:
|
|
engine = _StubEngine()
|
|
bus = EventBus(record_history=True)
|
|
instrumented_generate(
|
|
engine,
|
|
[Message(role=Role.USER, content="Hi")],
|
|
model="m",
|
|
bus=bus,
|
|
)
|
|
telem_events = [
|
|
e for e in bus.history if e.event_type == EventType.TELEMETRY_RECORD
|
|
]
|
|
assert len(telem_events) == 1
|
|
rec = telem_events[0].data["record"]
|
|
assert isinstance(rec, TelemetryRecord)
|
|
assert rec.model_id == "m"
|
|
|
|
def test_measures_latency(self) -> None:
|
|
engine = _StubEngine()
|
|
bus = EventBus(record_history=True)
|
|
instrumented_generate(
|
|
engine,
|
|
[Message(role=Role.USER, content="Hi")],
|
|
model="m",
|
|
bus=bus,
|
|
)
|
|
telem_events = [
|
|
e for e in bus.history if e.event_type == EventType.TELEMETRY_RECORD
|
|
]
|
|
rec = telem_events[0].data["record"]
|
|
assert rec.latency_seconds >= 0
|
|
|
|
def test_passes_kwargs(self) -> None:
|
|
engine = _StubEngine()
|
|
engine.generate = mock.MagicMock(return_value={"content": "ok", "usage": {}})
|
|
bus = EventBus()
|
|
instrumented_generate(
|
|
engine,
|
|
[Message(role=Role.USER, content="Hi")],
|
|
model="m",
|
|
bus=bus,
|
|
temperature=0.1,
|
|
max_tokens=512,
|
|
)
|
|
_, kwargs = engine.generate.call_args
|
|
assert kwargs["temperature"] == 0.1
|
|
assert kwargs["max_tokens"] == 512
|