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])}" f"Unexpected embedding response structure: {type(response.data[0])}"
) )
else: else:
logger.error(f"Unexpected response structure: {response}") logger.error(
f"Unexpected response structure: {response}"
)
raise ValueError( raise ValueError(
f"Unexpected embedding response structure: {type(response)}" f"Unexpected embedding response structure: {type(response)}"
) )

@ -168,7 +168,15 @@ class MajorityVoting:
title="Majority Voting", 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. 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 # 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: # 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 consensus_agent_name = self.consensus_agent.agent_name
if streaming_callback is not None: if streaming_callback is not None:
def consensus_streaming_callback(chunk: str): def consensus_streaming_callback(chunk: str):
"""Wrapper for consensus agent streaming callback.""" """Wrapper for consensus agent streaming callback."""
try: try:
if chunk is not None and chunk.strip(): 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: except Exception as callback_error:
if self.verbose: if self.verbose:
logger.warning( logger.warning(
f"[STREAMING] Callback failed for {consensus_agent_name}: {str(callback_error)}" f"[STREAMING] Callback failed for {consensus_agent_name}: {str(callback_error)}"
) )
else: else:
consensus_streaming_callback = None consensus_streaming_callback = None

@ -16,7 +16,9 @@ def pdf_to_text(pdf_path: str) -> str:
try: try:
import pypdf import pypdf
except ImportError: 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: try:
# Open the PDF file # Open the PDF file

@ -1,4 +1,4 @@
import pytest
# from unittest.mock import Mock, patch # from unittest.mock import Mock, patch
from swarms.structs.agent_router import AgentRouter from swarms.structs.agent_router import AgentRouter
@ -133,9 +133,13 @@ def test_add_agent_success():
def streaming_callback(chunk: str): def streaming_callback(chunk: str):
streamed_chunks.append(chunk) 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 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(): def test_add_agents_multiple():
@ -181,9 +185,13 @@ def test_add_agents_multiple():
def streaming_callback(chunk: str): def streaming_callback(chunk: str):
streamed_chunks.append(chunk) 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 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(): def test_find_best_agent_success():
@ -224,9 +232,13 @@ def test_find_best_agent_success():
def streaming_callback(chunk: str): def streaming_callback(chunk: str):
streamed_chunks.append(chunk) 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 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(): def test_find_best_agent_no_agents():
@ -260,7 +272,9 @@ def test_update_agent_history_success():
def streaming_callback(chunk: str): def streaming_callback(chunk: str):
streamed_chunks.append(chunk) 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 # Update agent history
router.update_agent_history("test_agent") router.update_agent_history("test_agent")
@ -356,9 +370,14 @@ def test_router_with_agent_streaming():
if chunk: if chunk:
streamed_chunks.append(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 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(): def test_router_find_and_run_with_streaming():
@ -400,9 +419,13 @@ def test_router_find_and_run_with_streaming():
if chunk: if chunk:
streamed_chunks.append(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 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__": if __name__ == "__main__":
@ -444,6 +467,7 @@ if __name__ == "__main__":
print(f"✗ FAILED: {test_func.__name__}") print(f"✗ FAILED: {test_func.__name__}")
print(f" Error: {str(e)}") print(f" Error: {str(e)}")
import traceback import traceback
traceback.print_exc() traceback.print_exc()
failed += 1 failed += 1

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

Loading…
Cancel
Save