github/autogen
1
1
Fork 0
mirror of https://github.com/microsoft/autogen.git synced 2026-04-20 03:02:16 -04:00
Code Issues Packages Projects Releases Wiki Activity

Add token usage termination (#4035)

Browse Source 
* Add token usage termination

* fix test
This commit is contained in:
Eric Zhu
2024-11-01 15:01:43 -07:00
committed by GitHub
parent ca7caa779d
commit 27ea99a485
3 changed files with 110 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
python/packages/autogen-agentchat/tests/test_termination_condition.py
View File

@@ -1,6 +1,12 @@
import pytest
from autogen_agentchat.messages import StopMessage, TextMessage
from autogen_agentchat.task import MaxMessageTermination, StopMessageTermination, TextMentionTermination
from autogen_agentchat.task import (
MaxMessageTermination,
StopMessageTermination,
TextMentionTermination,
TokenUsageTermination,
)
from autogen_core.components.models import RequestUsage
@pytest.mark.asyncio
@@ -51,6 +57,51 @@ async def test_mention_termination() -> None:
)
@pytest.mark.asyncio
async def test_token_usage_termination() -> None:
termination = TokenUsageTermination(max_total_token=10)
assert await termination([]) is None
await termination.reset()
assert (
await termination(
[
TextMessage(
content="Hello", source="user", model_usage=RequestUsage(prompt_tokens=10, completion_tokens=10)
)
]
)
is not None
)
await termination.reset()
assert (
await termination(
[
TextMessage(
content="Hello", source="user", model_usage=RequestUsage(prompt_tokens=1, completion_tokens=1)
),
TextMessage(
content="World", source="agent", model_usage=RequestUsage(prompt_tokens=1, completion_tokens=1)
),
]
)
is None
)
await termination.reset()
assert (
await termination(
[
TextMessage(
content="Hello", source="user", model_usage=RequestUsage(prompt_tokens=5, completion_tokens=0)
),
TextMessage(
content="stop", source="user", model_usage=RequestUsage(prompt_tokens=0, completion_tokens=5)
),
]
)
is not None
)
@pytest.mark.asyncio
async def test_and_termination() -> None:
termination = MaxMessageTermination(2) & TextMentionTermination("stop")