code cleanup op

pull/1175/head
Kye Gomez 4 days ago
parent 156bc595df
commit 638e9e2ba2

@ -83,7 +83,9 @@ class AgentRouter:
f"Unexpected embedding response structure: {type(response.data[0])}"
)
else:
logger.error(f"Unexpected response structure: {response}")
logger.error(
f"Unexpected response structure: {response}"
)
raise ValueError(
f"Unexpected embedding response structure: {type(response)}"
)

@ -168,7 +168,15 @@ class MajorityVoting:
title="Majority Voting",
)
def run(self, task: str, streaming_callback: Optional[Callable[[str, str, bool], None]] = None, *args, **kwargs) -> List[Any]:
def run(
self,
task: str,
streaming_callback: Optional[
Callable[[str, str, bool], None]
] = None,
*args,
**kwargs,
) -> List[Any]:
"""
Runs the majority voting system with multi-loop functionality and returns the majority vote.
@ -202,22 +210,28 @@ class MajorityVoting:
)
# Set streaming_on for the consensus agent based on the provided streaming_callback
self.consensus_agent.streaming_on = streaming_callback is not None
self.consensus_agent.streaming_on = (
streaming_callback is not None
)
# Instead of a simple passthrough wrapper, match the callback invocation pattern from the provided reference for the consensus agent:
consensus_agent_name = self.consensus_agent.agent_name
if streaming_callback is not None:
def consensus_streaming_callback(chunk: str):
"""Wrapper for consensus agent streaming callback."""
try:
if chunk is not None and chunk.strip():
streaming_callback(consensus_agent_name, chunk, False)
streaming_callback(
consensus_agent_name, chunk, False
)
except Exception as callback_error:
if self.verbose:
logger.warning(
f"[STREAMING] Callback failed for {consensus_agent_name}: {str(callback_error)}"
)
else:
consensus_streaming_callback = None

@ -16,7 +16,9 @@ def pdf_to_text(pdf_path: str) -> str:
try:
import pypdf
except ImportError:
raise ImportError("pypdf is not installed. Please install it using `pip install pypdf`.")
raise ImportError(
"pypdf is not installed. Please install it using `pip install pypdf`."
)
try:
# Open the PDF file

@ -1,4 +1,4 @@
import pytest
# from unittest.mock import Mock, patch
from swarms.structs.agent_router import AgentRouter
@ -133,9 +133,13 @@ def test_add_agent_success():
def streaming_callback(chunk: str):
streamed_chunks.append(chunk)
response = agent.run("Say hello", streaming_callback=streaming_callback)
response = agent.run(
"Say hello", streaming_callback=streaming_callback
)
assert response is not None
assert len(streamed_chunks) > 0 or response != "", "Agent should stream or return response"
assert (
len(streamed_chunks) > 0 or response != ""
), "Agent should stream or return response"
def test_add_agents_multiple():
@ -181,9 +185,13 @@ def test_add_agents_multiple():
def streaming_callback(chunk: str):
streamed_chunks.append(chunk)
response = agent.run("Say hi", streaming_callback=streaming_callback)
response = agent.run(
"Say hi", streaming_callback=streaming_callback
)
assert response is not None
assert len(streamed_chunks) > 0 or response != "", f"Agent {agent.agent_name} should stream or return response"
assert (
len(streamed_chunks) > 0 or response != ""
), f"Agent {agent.agent_name} should stream or return response"
def test_find_best_agent_success():
@ -224,9 +232,13 @@ def test_find_best_agent_success():
def streaming_callback(chunk: str):
streamed_chunks.append(chunk)
response = result.run("Test task", streaming_callback=streaming_callback)
response = result.run(
"Test task", streaming_callback=streaming_callback
)
assert response is not None
assert len(streamed_chunks) > 0 or response != "", "Found agent should stream or return response"
assert (
len(streamed_chunks) > 0 or response != ""
), "Found agent should stream or return response"
def test_find_best_agent_no_agents():
@ -260,7 +272,9 @@ def test_update_agent_history_success():
def streaming_callback(chunk: str):
streamed_chunks.append(chunk)
agent.run("Hello, how are you?", streaming_callback=streaming_callback)
agent.run(
"Hello, how are you?", streaming_callback=streaming_callback
)
# Update agent history
router.update_agent_history("test_agent")
@ -356,9 +370,14 @@ def test_router_with_agent_streaming():
if chunk:
streamed_chunks.append(chunk)
response = agent.run("Tell me a short joke", streaming_callback=streaming_callback)
response = agent.run(
"Tell me a short joke",
streaming_callback=streaming_callback,
)
assert response is not None
assert len(streamed_chunks) > 0 or response != "", f"Agent {agent.agent_name} should stream"
assert (
len(streamed_chunks) > 0 or response != ""
), f"Agent {agent.agent_name} should stream"
def test_router_find_and_run_with_streaming():
@ -400,9 +419,13 @@ def test_router_find_and_run_with_streaming():
if chunk:
streamed_chunks.append(chunk)
response = best_agent.run("What is 2 + 2?", streaming_callback=streaming_callback)
response = best_agent.run(
"What is 2 + 2?", streaming_callback=streaming_callback
)
assert response is not None
assert len(streamed_chunks) > 0 or response != "", "Best agent should stream when run"
assert (
len(streamed_chunks) > 0 or response != ""
), "Best agent should stream when run"
if __name__ == "__main__":
@ -444,6 +467,7 @@ if __name__ == "__main__":
print(f"✗ FAILED: {test_func.__name__}")
print(f" Error: {str(e)}")
import traceback
traceback.print_exc()
failed += 1

@ -202,12 +202,16 @@ def test_majority_voting_different_output_types():
# Assert majority vote is correct
assert majority_vote is not None
def test_streaming_majority_voting():
"""
Test the streaming_majority_voting with logging/try-except and assertion.
"""
logs = []
def streaming_callback(agent_name: str, chunk: str, is_final: bool):
def streaming_callback(
agent_name: str, chunk: str, is_final: bool
):
# Chunk buffer static per call (reset each session)
if not hasattr(streaming_callback, "_buffer"):
streaming_callback._buffer = ""
@ -218,7 +222,10 @@ def test_streaming_majority_voting():
if chunk:
streaming_callback._buffer += chunk
streaming_callback._buffer_size += len(chunk)
if streaming_callback._buffer_size >= min_chunk_size or is_final:
if (
streaming_callback._buffer_size >= min_chunk_size
or is_final
):
if streaming_callback._buffer:
print(streaming_callback._buffer, end="", flush=True)
logs.append(streaming_callback._buffer)

Loading…
Cancel
Save