mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 16:16:48 +00:00
fix: minor usage tracking improvements
This commit is contained in:
1 parent
0adf88abed
commit
f33bdaa52d
6 files changed
+197
-14
No files matched your search
@@ -194,6 +194,31 @@ def _patch_v1_db(conn):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_response_usage_reports_cumulative_turn_totals():
|
||||
from types import SimpleNamespace
|
||||
|
||||
from application.api.v1.routes import _response_usage
|
||||
|
||||
# Multi-round tool turns must report the accumulator's turn total,
|
||||
# not the final round's provider snapshot.
|
||||
agent = SimpleNamespace(
|
||||
llm=SimpleNamespace(
|
||||
_last_usage={
|
||||
"prompt_tokens": 5,
|
||||
"completion_tokens": 1,
|
||||
"total_tokens": 6,
|
||||
},
|
||||
token_usage={"prompt_tokens": 30, "generated_tokens": 12},
|
||||
)
|
||||
)
|
||||
assert _response_usage(agent) == {
|
||||
"prompt_tokens": 30,
|
||||
"completion_tokens": 12,
|
||||
"total_tokens": 42,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestLookupAgentHappy:
|
||||
def test_returns_agent_for_valid_key(self, pg_conn):
|
||||
|
||||
Reference in new issue
Block a user