fleet-memory/hindsight-api-slim/tests/test_llm_token_metrics.py
Nicolò Boschi 15ea23d5d6
feat: introduce hindsight-api-slim and hindsight-all-slim packages (#560)
* feat: introduce hindsight-api-slim and hindsight-all-slim packages

Closes #552

- Move all source code from hindsight-api/ to new hindsight-api-slim/
- hindsight-api-slim has heavy ML deps (torch, sentence-transformers,
  transformers, einops, flashrank, mlx, mlx-lm, safetensors) and
  pg0-embedded as optional extras: [local-ml], [embedded-db], [all]
- hindsight-api becomes a zero-code meta-package depending on
  hindsight-api-slim[all] for full backward compatibility
- Add hindsight-all-slim meta-package: hindsight-api-slim + client + embed
- hindsight-all updated to depend on hindsight-api-slim[all]
- pg0.py: lazy-import pg0 with clear ImportError pointing to [embedded-db]
- Dockerfile: replace sed hack with proper uv sync --extra flags
- Update release.yml, test.yml, lint.sh, release.sh, CLAUDE.md and
  all path references throughout the repo

* refactor: rename hindsight/ directory to hindsight-all/

* docs: document hindsight-api-slim and hindsight-all-slim package variants

Add package variants table and extras explanation to installation.md

* docs: remove emojis from installation.md, use professional tone

* docs: link Docker slim variant to pip package variants section

* docs: consolidate Docker image variants into single table

* ci: fix working-directory paths after package restructure

- Replace all hindsight-api → hindsight-api-slim in test.yml
- Replace hindsight → hindsight-all in test.yml
- Add --extra embedded-db to test-embed API install step

* ci: add local-ml and embedded-db extras to API sync steps

These extras were previously implicit in the old hindsight-api package
(which bundled everything). Now that hindsight-api-slim uses optional
extras, we must explicitly request local-ml and embedded-db in CI.

* ci: add API install step with embedded-db to test-embed smoke test

The smoke test starts hindsight-api as a daemon, which requires pg0-embedded.
Add a dedicated install step for hindsight-api-slim with embedded-db extra
so the daemon can start successfully.

* ci: remove --no-install-project when using optional extras

When --no-install-project is combined with --extra, the optional deps
are not installed because extras require the project to be active.
Remove --no-install-project from steps that need local-ml or embedded-db.

* ci: fix ordering of uv sync steps to preserve optional extras

When uv sync runs for a different workspace member, it removes optional
extras installed for other members. Fix by always running extra-requiring
API sync last, after other workspace member syncs.

Also remove --no-install-project from embedded-db sync in test-embed,
as --no-install-project prevents optional extras from being active.

* ci: add local-ml extra to test-embed API install for smoke test

The smoke test starts the full API server which needs sentence-transformers
for local embeddings (default provider). Add local-ml extra to the install.

* ci: simplify extras with --all-extras and add slim pip smoke test

- Replace explicit --extra local-ml --extra embedded-db with --all-extras
  for cleaner, more maintainable sync steps
- Add test-pip-slim job: tests hindsight-api-slim[embedded-db] without
  local ML models, using Cohere for embeddings/reranking (mirrors Docker
  slim smoke test approach)

* ci: simplify slim smoke test to health check only (mirrors Docker test)
2026-03-13 13:50:03 +01:00

247 lines
8.6 KiB
Python

"""
Test that LLM calls record token metrics via the metrics collector.
"""
import os
from unittest.mock import MagicMock, patch
import pytest
from hindsight_api.engine.llm_wrapper import LLMProvider
from hindsight_api.metrics import (
MetricsCollector,
NoOpMetricsCollector,
get_metrics_collector,
)
def get_groq_api_key() -> str | None:
"""Get Groq API key from environment."""
return os.getenv("GROQ_API_KEY")
@pytest.mark.asyncio
async def test_llm_metrics_recorded_for_groq():
"""
Test that LLM metrics are recorded when making LLM calls via Groq.
Uses openai/gpt-oss-20b as recommended by Hindsight.
"""
api_key = get_groq_api_key()
if not api_key:
pytest.skip("Skipping: GROQ_API_KEY not set")
# Create a mock metrics collector to track record_llm_call calls
mock_collector = MagicMock(spec=MetricsCollector)
# Patch the provider module where get_metrics_collector is actually called
with patch("hindsight_api.engine.providers.openai_compatible_llm.get_metrics_collector", return_value=mock_collector):
llm = LLMProvider(
provider="groq",
api_key=api_key,
base_url="",
model="openai/gpt-oss-20b",
)
# Make an LLM call with clear instruction
response = await llm.call(
messages=[
{"role": "system", "content": "You are a helpful assistant. Always respond."},
{"role": "user", "content": "What is 2+2? Reply with just the number."}
],
max_completion_tokens=50,
scope="test_metrics",
)
# Verify record_llm_call was called - this is the main test
assert mock_collector.record_llm_call.called, "record_llm_call should have been called"
# Get the call arguments
call_kwargs = mock_collector.record_llm_call.call_args.kwargs
# Verify the call had correct structure
assert call_kwargs["provider"] == "groq", f"Expected provider='groq', got {call_kwargs}"
assert call_kwargs["model"] == "openai/gpt-oss-20b", f"Expected model='openai/gpt-oss-20b', got {call_kwargs}"
assert call_kwargs["scope"] == "test_metrics", f"Expected scope='test_metrics', got {call_kwargs}"
assert call_kwargs["duration"] > 0, f"Expected duration > 0, got {call_kwargs['duration']}"
assert call_kwargs["input_tokens"] > 0, f"Expected input_tokens > 0, got {call_kwargs['input_tokens']}"
assert call_kwargs["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {call_kwargs['output_tokens']}"
assert call_kwargs["success"] is True, f"Expected success=True, got {call_kwargs['success']}"
print(f"\nLLM metrics recorded:")
print(f" provider: {call_kwargs['provider']}")
print(f" model: {call_kwargs['model']}")
print(f" scope: {call_kwargs['scope']}")
print(f" duration: {call_kwargs['duration']:.3f}s")
print(f" input_tokens: {call_kwargs['input_tokens']}")
print(f" output_tokens: {call_kwargs['output_tokens']}")
print(f" response: {response}")
@pytest.mark.asyncio
async def test_llm_metrics_recorded_for_structured_output():
"""
Test that LLM metrics are recorded for structured output (JSON) calls.
"""
api_key = get_groq_api_key()
if not api_key:
pytest.skip("Skipping: GROQ_API_KEY not set")
from pydantic import BaseModel
class SimpleResponse(BaseModel):
greeting: str
language: str
mock_collector = MagicMock(spec=MetricsCollector)
# Patch the provider module where get_metrics_collector is actually called
with patch("hindsight_api.engine.providers.openai_compatible_llm.get_metrics_collector", return_value=mock_collector):
llm = LLMProvider(
provider="groq",
api_key=api_key,
base_url="",
model="openai/gpt-oss-20b",
)
# Make a structured output call
response = await llm.call(
messages=[{"role": "user", "content": "Say hello in French. Return greeting and language."}],
response_format=SimpleResponse,
max_completion_tokens=100,
scope="structured_output_test",
)
# Verify structured response
assert isinstance(response, SimpleResponse)
assert response.greeting is not None
assert response.language is not None
# Verify record_llm_call was called
assert mock_collector.record_llm_call.called, "record_llm_call should have been called"
call_kwargs = mock_collector.record_llm_call.call_args.kwargs
assert call_kwargs["input_tokens"] > 0
assert call_kwargs["output_tokens"] > 0
print(f"\nStructured output LLM metrics:")
print(f" greeting: {response.greeting}")
print(f" language: {response.language}")
print(f" input_tokens: {call_kwargs['input_tokens']}")
print(f" output_tokens: {call_kwargs['output_tokens']}")
@pytest.mark.asyncio
async def test_noop_collector_when_metrics_disabled():
"""
Test that NoOpMetricsCollector is returned when metrics are not initialized.
This verifies the fallback behavior doesn't break LLM calls.
"""
api_key = get_groq_api_key()
if not api_key:
pytest.skip("Skipping: GROQ_API_KEY not set")
# Without initializing metrics, get_metrics_collector returns NoOpMetricsCollector
collector = get_metrics_collector()
assert isinstance(collector, NoOpMetricsCollector), "Should return NoOpMetricsCollector when not initialized"
# Make an LLM call - should work fine with NoOp collector
llm = LLMProvider(
provider="groq",
api_key=api_key,
base_url="",
model="openai/gpt-oss-20b",
)
response = await llm.call(
messages=[{"role": "user", "content": "Say 'test' in one word."}],
max_completion_tokens=50,
)
assert response is not None
print(f"\nLLM call succeeded with NoOpMetricsCollector: {response}")
@pytest.mark.asyncio
async def test_return_usage_returns_tuple():
"""
Test that return_usage=True returns (result, TokenUsage) tuple.
"""
from hindsight_api.engine.response_models import TokenUsage
api_key = get_groq_api_key()
if not api_key:
pytest.skip("Skipping: GROQ_API_KEY not set")
llm = LLMProvider(
provider="groq",
api_key=api_key,
base_url="",
model="openai/gpt-oss-20b",
)
# Call with return_usage=True
result, usage = await llm.call(
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is 2+2? Reply with just the number."}
],
max_completion_tokens=50,
return_usage=True,
)
# Verify result is the response text
assert result is not None
assert isinstance(result, str)
# Verify usage is TokenUsage model with valid counts
assert isinstance(usage, TokenUsage)
assert usage.input_tokens > 0, f"Expected input_tokens > 0, got {usage.input_tokens}"
assert usage.output_tokens >= 0, f"Expected output_tokens >= 0, got {usage.output_tokens}"
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
print(f"\nreturn_usage=True test:")
print(f" result: {result}")
print(f" usage: {usage}")
@pytest.mark.asyncio
async def test_return_usage_with_structured_output():
"""
Test that return_usage=True works with structured output (JSON).
"""
from pydantic import BaseModel
from hindsight_api.engine.response_models import TokenUsage
api_key = get_groq_api_key()
if not api_key:
pytest.skip("Skipping: GROQ_API_KEY not set")
class MathAnswer(BaseModel):
answer: int
explanation: str
llm = LLMProvider(
provider="groq",
api_key=api_key,
base_url="",
model="openai/gpt-oss-20b",
)
# Call with return_usage=True and structured output
result, usage = await llm.call(
messages=[{"role": "user", "content": "What is 5+3? Return the answer and a brief explanation."}],
response_format=MathAnswer,
max_completion_tokens=100,
return_usage=True,
)
# Verify result is the parsed response
assert isinstance(result, MathAnswer)
assert result.answer == 8
assert result.explanation is not None
# Verify usage is TokenUsage model
assert isinstance(usage, TokenUsage)
assert usage.input_tokens > 0
assert usage.output_tokens > 0
print(f"\nStructured output with return_usage=True:")
print(f" result: {result}")
print(f" usage: {usage}")