Skip to content

Commit 911eb8c

Browse files
author
matdev83
committed
feat: Enhance usage tracking and model capabilities
This commit introduces several improvements to the usage tracking and model capabilities systems. - **Refactor Usage Tracking:** The usage tracking statistics endpoint now returns a UsageStatsResponse Pydantic model instead of a raw dictionary. This improves type safety, ensures a consistent data structure, and enables better validation. The UsageController, IUsageTrackingService, and UsageTrackingService have been updated accordingly. - **Improve Model Capabilities:** The default reasoning and execution phase parameters now include easoning_effort. This ensures that even unknown backends have a default reasoning effort setting, improving the robustness of the reasoning system. - **Update Tests:** All relevant unit tests for usage tracking have been updated to align with the new UsageStatsResponse model and to ensure comprehensive coverage of the changes.
1 parent 28017a9 commit 911eb8c

File tree

8 files changed

+228
-117
lines changed

8 files changed

+228
-117
lines changed

data/test_suite_state.json

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
{
2-
"test_count": 5119,
3-
"last_updated": "1762168167.0802596"
2+
"test_count": 5046,
3+
"last_updated": "1762217835.0"
44
}

src/connectors/utils/model_capabilities.py

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -27,18 +27,18 @@
2727
REASONING_TAG_FORMAT: dict[str, tuple[str, str]] = {
2828
"openai": ("<thinking>", "</thinking>"),
2929
"anthropic": ("<thinking>", "</thinking>"),
30-
"qwen": ("<thinking>", "</thinking>"),
31-
"qwen-oauth": ("<thinking>", "</thinking>"),
32-
"deepseek": ("<think>", "</think>"), # DeepSeek uses <think> natively
33-
"minimax": ("<think>", "</think>"),
34-
"gemini": ("<thinking>", "</thinking>"),
35-
"gemini-oauth-plan": ("<thinking>", "</thinking>"),
36-
"gemini-oauth-free": ("<thinking>", "</thinking>"),
37-
"gemini-cli-acp": ("<thinking>", "</thinking>"),
38-
"gemini-cli-cloud-project": ("<thinking>", "</thinking>"),
39-
# Default for others
40-
"_default": ("<reasoning>", "</reasoning>"),
41-
}
30+
"qwen": ("<thinking>", "</thinking>"),
31+
"qwen-oauth": ("<thinking>", "</thinking>"),
32+
"deepseek": ("<think>", "</think>"), # DeepSeek uses <think> natively
33+
"minimax": ("<think>", "</think>"),
34+
"gemini": ("<thinking>", "</thinking>"),
35+
"gemini-oauth-plan": ("<thinking>", "</thinking>"),
36+
"gemini-oauth-free": ("<thinking>", "</thinking>"),
37+
"gemini-cli-acp": ("<thinking>", "</thinking>"),
38+
"gemini-cli-cloud-project": ("<thinking>", "</thinking>"),
39+
# Default for others
40+
"_default": ("<reasoning>", "</reasoning>"),
41+
}
4242

4343
# Reasoning phase parameters - maximize reasoning quality
4444
REASONING_PHASE_PARAMS: dict[str, dict[str, Any]] = {
@@ -63,6 +63,7 @@
6363
"_default": {
6464
# Generic reasoning parameters
6565
"temperature": 0.7, # Balanced for reasoning
66+
"reasoning_effort": "high", # Ensure strong reasoning on unknown backends
6667
},
6768
}
6869

@@ -87,6 +88,7 @@
8788
"_default": {
8889
# Generic execution parameters
8990
"temperature": 0.5, # Lower for consistency
91+
"reasoning_effort": "low", # Minimize reasoning for unknown backends
9092
},
9193
}
9294

src/core/app/controllers/usage_controller.py

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,13 @@
55
from __future__ import annotations
66

77
import logging
8-
from typing import Any
8+
from typing import Any, cast
99

1010
from fastapi import APIRouter, Depends, Query
1111

1212
from src.core.di.services import get_or_build_service_provider
1313
from src.core.domain.usage_data import UsageData
14+
from src.core.domain.usage_stats import UsageStatsResponse
1415
from src.core.interfaces.usage_tracking_interface import IUsageTrackingService
1516

1617
logger = logging.getLogger(__name__)
@@ -31,21 +32,21 @@ def __init__(self, usage_service: IUsageTrackingService | None = None) -> None:
3132

3233
async def get_usage_stats(
3334
self, project: str | None = None, days: int = 30
34-
) -> dict[str, Any]:
35+
) -> UsageStatsResponse:
3536
"""Get usage statistics.
3637
3738
Args:
3839
project: Optional project filter
3940
days: Number of days to include in stats
4041
4142
Returns:
42-
Usage statistics dictionary
43+
Usage statistics response model
4344
"""
4445
if not self.usage_service:
45-
return {"error": "Usage tracking service not available"}
46+
raise RuntimeError("Usage tracking service not available")
4647

4748
result = await self.usage_service.get_usage_stats(project=project, days=days)
48-
return result # type: ignore[no-any-return]
49+
return result
4950

5051
async def get_recent_usage(
5152
self, session_id: str | None = None, limit: int = 100
@@ -68,12 +69,12 @@ async def get_recent_usage(
6869
return result # type: ignore[no-any-return]
6970

7071

71-
@router.get("/stats", response_model=dict[str, Any])
72+
@router.get("/stats", response_model=UsageStatsResponse)
7273
async def get_usage_stats(
7374
project: str | None = Query(None, description="Filter by project name"),
7475
days: int = Query(30, description="Number of days to include in stats"),
7576
service_provider: Any = Depends(get_or_build_service_provider),
76-
) -> dict[str, Any]:
77+
) -> UsageStatsResponse:
7778
"""Get usage statistics.
7879
7980
Args:
@@ -84,9 +85,12 @@ async def get_usage_stats(
8485
Returns:
8586
Usage statistics dictionary
8687
"""
87-
usage_service = service_provider.get_required_service(IUsageTrackingService)
88+
usage_service = cast(
89+
IUsageTrackingService,
90+
service_provider.get_required_service(IUsageTrackingService),
91+
)
8892
result = await usage_service.get_usage_stats(project=project, days=days)
89-
return result # type: ignore[no-any-return]
93+
return result
9094

9195

9296
@router.get("/recent", response_model=list[UsageData])

src/core/interfaces/usage_tracking_interface.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from typing import Any
99

1010
from src.core.domain.usage_data import UsageData
11+
from src.core.domain.usage_stats import UsageStatsResponse
1112

1213

1314
class IUsageTrackingService(abc.ABC):
@@ -44,7 +45,7 @@ async def track_request(
4445
@abc.abstractmethod
4546
async def get_usage_stats(
4647
self, project: str | None = None, days: int = 30
47-
) -> dict[str, Any]:
48+
) -> UsageStatsResponse:
4849
pass
4950

5051
@abc.abstractmethod

src/core/services/usage_tracking_service.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -274,23 +274,27 @@ def _parse_billing_cost(value: Any) -> float | None:
274274

275275
async def get_usage_stats(
276276
self, project: str | None = None, days: int = 30
277-
) -> dict[str, Any]:
278-
"""Get usage statistics.
277+
) -> UsageStatsResponse:
278+
"""Get usage statistics aggregated by model.
279279
280280
Args:
281281
project: Optional project filter
282282
days: Number of days to include in stats
283283
284284
Returns:
285-
Usage statistics dictionary
285+
A :class:`UsageStatsResponse` with aggregated usage metrics.
286286
"""
287287
if days <= 0:
288288
logger.warning(
289289
"Received non-positive days=%s when requesting usage stats; "
290290
"falling back to complete history.",
291291
days,
292292
)
293-
return await self._repository.get_stats(project)
293+
raw_stats = await self._repository.get_stats(project)
294+
stats_response = UsageStatsResponse()
295+
for model_name, payload in raw_stats.items():
296+
stats_response[model_name] = payload
297+
return stats_response
294298

295299
usage_records = await self._repository.get_all()
296300
cutoff = datetime.now(timezone.utc) - timedelta(days=days)
@@ -329,7 +333,7 @@ async def get_usage_stats(
329333
requests=current_stats.requests + 1,
330334
)
331335

332-
return stats.model_dump()
336+
return stats
333337

334338
async def get_recent_usage(
335339
self, session_id: str | None = None, limit: int = 100

0 commit comments

Comments
 (0)