96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
"""
|
|
Copyright 2024-2026 ReplyPier
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
from typing import List
|
|
|
|
import pytest
|
|
|
|
from app.agents.chat_agent import ChatAgent
|
|
from app.agents.history_memory import BoundedHistoryMemory
|
|
|
|
OVERFLOW = Exception(
|
|
"This model's maximum context length is 128000 tokens. However, your messages resulted in 145155 tokens"
|
|
)
|
|
|
|
|
|
class _FakeAgent:
|
|
"""Stands in for the agno Agent: each arun() consumes the next outcome and
|
|
records the history budget in force at that moment."""
|
|
|
|
def __init__(self, outcomes: List):
|
|
self.outcomes = list(outcomes)
|
|
self.memory = BoundedHistoryMemory()
|
|
self.calls = 0
|
|
self.budgets: List[int] = []
|
|
|
|
async def arun(self, **kwargs):
|
|
self.calls += 1
|
|
self.budgets.append(getattr(self.memory, "history_max_chars", None))
|
|
outcome = self.outcomes.pop(0)
|
|
if isinstance(outcome, BaseException):
|
|
raise outcome
|
|
return outcome
|
|
|
|
|
|
def _chat_agent(outcomes: List) -> ChatAgent:
|
|
shell = ChatAgent.__new__(ChatAgent)
|
|
shell.agent = _FakeAgent(outcomes)
|
|
return shell
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overflow_retries_once_without_history_then_restores_budget():
|
|
shell = _chat_agent([OVERFLOW, "ok"])
|
|
assert await shell._arun("hi", "s1") == "ok"
|
|
assert shell.agent.calls == 2
|
|
assert shell.agent.budgets[1] == 0
|
|
assert shell.agent.memory.history_max_chars > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_other_errors_are_not_retried():
|
|
shell = _chat_agent([Exception("rate limit exceeded")])
|
|
with pytest.raises(Exception, match="rate limit"):
|
|
await shell._arun("hi", "s1")
|
|
assert shell.agent.calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_second_overflow_surfaces():
|
|
shell = _chat_agent([OVERFLOW, OVERFLOW])
|
|
with pytest.raises(Exception, match="maximum context length"):
|
|
await shell._arun("hi", "s1")
|
|
assert shell.agent.calls == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_timeout_is_not_retried():
|
|
shell = _chat_agent([asyncio.TimeoutError()])
|
|
with pytest.raises(asyncio.TimeoutError):
|
|
await shell._arun("hi", "s1")
|
|
assert shell.agent.calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_plain_memory_is_not_retried():
|
|
shell = _chat_agent([OVERFLOW])
|
|
shell.agent.memory = SimpleNamespace()
|
|
with pytest.raises(Exception, match="maximum context length"):
|
|
await shell._arun("hi", "s1")
|
|
assert shell.agent.calls == 1
|