parent
0fd5d9a634
commit
a9b30998e4
@ -1,57 +1,24 @@
|
|||||||
from typing import List, Callable
|
from typing import List
|
||||||
from swarms.workers.worker import Worker
|
from swarms.workers.worker import Worker
|
||||||
|
|
||||||
|
|
||||||
class DialogueSimulator:
|
class DialogueSimulator:
|
||||||
def __init__(
|
def __init__(self, agents: List[Worker]):
|
||||||
self,
|
|
||||||
agents: List[Worker],
|
|
||||||
selection_func: Callable[[int, List[Worker]], int],
|
|
||||||
):
|
|
||||||
self.agents = agents
|
self.agents = agents
|
||||||
self._step = 0
|
|
||||||
self.select_next_speaker = selection_func
|
|
||||||
|
|
||||||
def reset(self):
|
|
||||||
for agent in self.agents:
|
|
||||||
agent.reset()
|
|
||||||
|
|
||||||
def start(self, name: str, message: str):
|
|
||||||
#init conv with a message from name
|
|
||||||
prompt = f"Name {name} and message: {message}"
|
|
||||||
|
|
||||||
for agent in self.agents:
|
|
||||||
agent.run(prompt)
|
|
||||||
|
|
||||||
#increment time
|
|
||||||
self._step += 1
|
|
||||||
|
|
||||||
def inject(self, name: str, message: str):
|
|
||||||
for agent in self.agents:
|
|
||||||
agent.receieve(name, message)
|
|
||||||
|
|
||||||
self._step += 1
|
|
||||||
|
|
||||||
def step(self) -> tuple[str, str]:
|
|
||||||
#choose next speaker
|
|
||||||
speaker_idx = self.select_next_speaker(
|
|
||||||
self._step,
|
|
||||||
self.agents
|
|
||||||
)
|
|
||||||
speaker = self.agents[speaker_idx]
|
|
||||||
|
|
||||||
#2. next speaker ends message
|
|
||||||
message = speaker.run()
|
|
||||||
|
|
||||||
#everyone receives messages
|
def run(self, max_iters: int, name: str = None, message: str = None):
|
||||||
for receiver in self.agents:
|
step = 0
|
||||||
receiver.receive(speaker.name, message)
|
if name and message:
|
||||||
|
prompt = f"Name {name} and message: {message}"
|
||||||
#increment time
|
for agent in self.agents:
|
||||||
self._step += 1
|
agent.run(prompt)
|
||||||
|
step += 1
|
||||||
|
|
||||||
return speaker.name, message
|
while step < max_iters:
|
||||||
|
speaker_idx = step % len(self.agents)
|
||||||
def select_next_speaker(step: int, agents) -> int:
|
speaker = self.agents[speaker_idx]
|
||||||
idx = (step) % len(agents)
|
speaker_message = speaker.run()
|
||||||
return idx
|
for receiver in self.agents:
|
||||||
|
receiver.receive(speaker.name, speaker_message)
|
||||||
|
print(f"({speaker.name}): {speaker_message}")
|
||||||
|
print("\n")
|
||||||
|
step += 1
|
Loading…
Reference in new issue