mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-24 20:51:02 +02:00
Fixing tests
This commit is contained in:
parent
6c0147feda
commit
c14116f7a3
1 changed files with 35 additions and 26 deletions
|
|
@ -28,34 +28,35 @@ class TestAgentStreaming:
|
||||||
client = AsyncMock()
|
client = AsyncMock()
|
||||||
|
|
||||||
async def agent_react_streaming(variables, timeout=600, streaming=False, chunk_callback=None):
|
async def agent_react_streaming(variables, timeout=600, streaming=False, chunk_callback=None):
|
||||||
if streaming and chunk_callback:
|
# Both modes return the same text for equivalence
|
||||||
# Simulate streaming response with chunks
|
full_text = """Thought: I need to search for information about machine learning.
|
||||||
chunks = [
|
|
||||||
"Thought: I need to",
|
|
||||||
" search for",
|
|
||||||
" information about",
|
|
||||||
" machine learning.",
|
|
||||||
"\n",
|
|
||||||
"Action: knowledge_query\n",
|
|
||||||
"Args: {\n",
|
|
||||||
' "question": "What is machine learning?"\n',
|
|
||||||
"}"
|
|
||||||
]
|
|
||||||
|
|
||||||
full_text = ""
|
|
||||||
for chunk in chunks:
|
|
||||||
full_text += chunk
|
|
||||||
await chunk_callback(chunk)
|
|
||||||
|
|
||||||
return full_text
|
|
||||||
else:
|
|
||||||
# Non-streaming response
|
|
||||||
return """Thought: I need to search for information about machine learning.
|
|
||||||
Action: knowledge_query
|
Action: knowledge_query
|
||||||
Args: {
|
Args: {
|
||||||
"question": "What is machine learning?"
|
"question": "What is machine learning?"
|
||||||
}"""
|
}"""
|
||||||
|
|
||||||
|
if streaming and chunk_callback:
|
||||||
|
# Simulate streaming response with complete lines as chunks
|
||||||
|
# This mimics how real LLMs stream - complete tokens/words
|
||||||
|
chunks = [
|
||||||
|
"Thought: I need to search",
|
||||||
|
" for information about",
|
||||||
|
" machine learning.\n",
|
||||||
|
"Action: knowledge_query\n",
|
||||||
|
"Args: {\n",
|
||||||
|
' "question": ',
|
||||||
|
'"What is machine learning?"',
|
||||||
|
'\n}'
|
||||||
|
]
|
||||||
|
|
||||||
|
for chunk in chunks:
|
||||||
|
await chunk_callback(chunk)
|
||||||
|
|
||||||
|
return full_text
|
||||||
|
else:
|
||||||
|
# Non-streaming response - same text
|
||||||
|
return full_text
|
||||||
|
|
||||||
client.agent_react.side_effect = agent_react_streaming
|
client.agent_react.side_effect = agent_react_streaming
|
||||||
return client
|
return client
|
||||||
|
|
||||||
|
|
@ -240,11 +241,19 @@ Args: {
|
||||||
assert result is not None
|
assert result is not None
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_agent_streaming_with_conversation_history(self, agent_manager, mock_flow_context,
|
async def test_agent_streaming_with_conversation_history(self, agent_manager, mock_flow_context):
|
||||||
sample_agent_responses):
|
|
||||||
"""Test streaming with existing conversation history"""
|
"""Test streaming with existing conversation history"""
|
||||||
# Arrange
|
# Arrange
|
||||||
history = sample_agent_responses[:1] # Use first response as history
|
# History should be a list of Action objects
|
||||||
|
from trustgraph.agent.react.types import Action
|
||||||
|
history = [
|
||||||
|
Action(
|
||||||
|
thought="I need to search for information about machine learning",
|
||||||
|
name="knowledge_query",
|
||||||
|
arguments={"question": "What is machine learning?"},
|
||||||
|
observation="Machine learning is a subset of AI that enables computers to learn from data."
|
||||||
|
)
|
||||||
|
]
|
||||||
think = AsyncMock()
|
think = AsyncMock()
|
||||||
|
|
||||||
# Act
|
# Act
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue