Spaces:
Runtime error
Runtime error
Merge pull request #63 from EnvisionMindCa/codex/fix-code-to-await-junior-agent-response
Browse files- src/team.py +13 -7
src/team.py
CHANGED
@@ -24,13 +24,12 @@ def set_team(team: "TeamChatSession" | None) -> None:
|
|
24 |
|
25 |
|
26 |
async def send_to_junior(message: str) -> str:
|
27 |
-
"""Forward ``message`` to the junior agent and
|
28 |
|
29 |
if _TEAM is None:
|
30 |
return "No active team"
|
31 |
|
32 |
-
_TEAM.queue_message_to_junior(message)
|
33 |
-
return "Message sent to junior"
|
34 |
|
35 |
|
36 |
# Backwards compatibility ---------------------------------------------------
|
@@ -46,7 +45,7 @@ class TeamChatSession:
|
|
46 |
host: str = OLLAMA_HOST,
|
47 |
model: str = MODEL_NAME,
|
48 |
) -> None:
|
49 |
-
self._to_junior: asyncio.Queue[str] = asyncio.Queue()
|
50 |
self._to_senior: asyncio.Queue[str] = asyncio.Queue()
|
51 |
self._junior_task: asyncio.Task | None = None
|
52 |
self.senior = ChatSession(
|
@@ -80,15 +79,20 @@ class TeamChatSession:
|
|
80 |
def upload_document(self, file_path: str) -> str:
|
81 |
return self.senior.upload_document(file_path)
|
82 |
|
83 |
-
def queue_message_to_junior(self, message: str) ->
|
84 |
-
|
|
|
|
|
|
|
|
|
85 |
if not self._junior_task or self._junior_task.done():
|
86 |
self._junior_task = asyncio.create_task(self._process_junior())
|
|
|
87 |
|
88 |
async def _process_junior(self) -> None:
|
89 |
try:
|
90 |
while not self._to_junior.empty():
|
91 |
-
msg = await self._to_junior.get()
|
92 |
self.junior._messages.append({"role": "tool", "name": "senior", "content": msg})
|
93 |
DBMessage.create(conversation=self.junior._conversation, role="tool", content=msg)
|
94 |
parts: list[str] = []
|
@@ -98,6 +102,8 @@ class TeamChatSession:
|
|
98 |
result = "\n".join(parts)
|
99 |
if result.strip():
|
100 |
await self._to_senior.put(result)
|
|
|
|
|
101 |
|
102 |
if self.senior._state == "idle":
|
103 |
await self._deliver_junior_messages()
|
|
|
24 |
|
25 |
|
26 |
async def send_to_junior(message: str) -> str:
|
27 |
+
"""Forward ``message`` to the junior agent and await the response."""
|
28 |
|
29 |
if _TEAM is None:
|
30 |
return "No active team"
|
31 |
|
32 |
+
return await _TEAM.queue_message_to_junior(message)
|
|
|
33 |
|
34 |
|
35 |
# Backwards compatibility ---------------------------------------------------
|
|
|
45 |
host: str = OLLAMA_HOST,
|
46 |
model: str = MODEL_NAME,
|
47 |
) -> None:
|
48 |
+
self._to_junior: asyncio.Queue[tuple[str, asyncio.Future[str]]] = asyncio.Queue()
|
49 |
self._to_senior: asyncio.Queue[str] = asyncio.Queue()
|
50 |
self._junior_task: asyncio.Task | None = None
|
51 |
self.senior = ChatSession(
|
|
|
79 |
def upload_document(self, file_path: str) -> str:
|
80 |
return self.senior.upload_document(file_path)
|
81 |
|
82 |
+
async def queue_message_to_junior(self, message: str) -> str:
|
83 |
+
"""Send ``message`` to the junior agent and wait for the reply."""
|
84 |
+
|
85 |
+
loop = asyncio.get_running_loop()
|
86 |
+
fut: asyncio.Future[str] = loop.create_future()
|
87 |
+
await self._to_junior.put((message, fut))
|
88 |
if not self._junior_task or self._junior_task.done():
|
89 |
self._junior_task = asyncio.create_task(self._process_junior())
|
90 |
+
return await fut
|
91 |
|
92 |
async def _process_junior(self) -> None:
|
93 |
try:
|
94 |
while not self._to_junior.empty():
|
95 |
+
msg, fut = await self._to_junior.get()
|
96 |
self.junior._messages.append({"role": "tool", "name": "senior", "content": msg})
|
97 |
DBMessage.create(conversation=self.junior._conversation, role="tool", content=msg)
|
98 |
parts: list[str] = []
|
|
|
102 |
result = "\n".join(parts)
|
103 |
if result.strip():
|
104 |
await self._to_senior.put(result)
|
105 |
+
if not fut.done():
|
106 |
+
fut.set_result(result)
|
107 |
|
108 |
if self.senior._state == "idle":
|
109 |
await self._deliver_junior_messages()
|