tech-envision commited on
Commit
97c8016
·
unverified ·
2 Parent(s): b9049da 9728a79

Merge pull request #63 from EnvisionMindCa/codex/fix-code-to-await-junior-agent-response

Browse files
Files changed (1) hide show
  1. 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 return a status string."""
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) -> None:
84
- self._to_junior.put_nowait(message)
 
 
 
 
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()