forked from bojieli/ai-agent-book
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_stream_empty_choices.py
More file actions
64 lines (51 loc) · 1.72 KB
/
Copy pathtest_stream_empty_choices.py
File metadata and controls
64 lines (51 loc) · 1.72 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
"""Regression tests for providers that end streams with ``choices=[]``."""
import os
import sys
from types import SimpleNamespace
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from agent import UserMemoryAgent
from conversational_agent import ConversationConfig, ConversationalAgent
class StubCompletions:
def create(self, **kwargs):
return [
SimpleNamespace(
choices=[
SimpleNamespace(delta=SimpleNamespace(content="Hello"))
]
),
SimpleNamespace(choices=[]),
]
def stub_client():
return SimpleNamespace(
chat=SimpleNamespace(completions=StubCompletions())
)
def test_conversational_agent_ignores_empty_choices_chunk():
agent = ConversationalAgent.__new__(ConversationalAgent)
agent.config = ConversationConfig(
enable_memory_context=False,
enable_conversation_history=False,
)
agent.verbose = False
agent.model = "test-model"
agent.client = stub_client()
agent.conversation = []
agent.conversation_history = None
agent.session_id = "session-test"
assert agent.chat("Hi") == "Hello"
assert agent.conversation[-1] == {
"role": "assistant",
"content": "Hello",
}
def test_user_memory_agent_ignores_empty_choices_chunk():
agent = UserMemoryAgent.__new__(UserMemoryAgent)
agent._get_memory_context = lambda: ""
agent.model = "test-model"
agent.client = stub_client()
agent.conversation = []
agent.conversation_history = None
agent.session_id = "session-test"
assert agent._chat_stream("Hi") == "Hello"
assert agent.conversation[-1] == {
"role": "assistant",
"content": "Hello",
}