tech-envision commited on
Commit
e6ecb98
·
1 Parent(s): 7a9bc48

fix duplicate junior messages

Browse files
Files changed (2) hide show
  1. src/chat.py +10 -19
  2. src/team.py +8 -6
src/chat.py CHANGED
@@ -240,13 +240,10 @@ class ChatSession:
240
  if not func:
241
  _LOG.warning("Unsupported tool call: %s", call.function.name)
242
  result = f"Unsupported tool: {call.function.name}"
243
- messages.append(
244
- {
245
- "role": "tool",
246
- "name": call.function.name,
247
- "content": result,
248
- }
249
  )
 
250
  DBMessage.create(
251
  conversation=conversation,
252
  role="tool",
@@ -262,7 +259,7 @@ class ChatSession:
262
 
263
  placeholder = {
264
  "role": "tool",
265
- "name": call.function.name,
266
  "content": "Awaiting tool response...",
267
  }
268
  messages.append(placeholder)
@@ -286,13 +283,10 @@ class ChatSession:
286
  pass
287
  self._remove_tool_placeholder(messages)
288
  result = await exec_task
289
- messages.append(
290
- {
291
- "role": "tool",
292
- "name": call.function.name,
293
- "content": result,
294
- }
295
  )
 
296
  DBMessage.create(
297
  conversation=conversation,
298
  role="tool",
@@ -313,13 +307,10 @@ class ChatSession:
313
  yield followup
314
  result = await exec_task
315
  self._remove_tool_placeholder(messages)
316
- messages.append(
317
- {
318
- "role": "tool",
319
- "name": call.function.name,
320
- "content": result,
321
- }
322
  )
 
323
  DBMessage.create(
324
  conversation=conversation,
325
  role="tool",
 
240
  if not func:
241
  _LOG.warning("Unsupported tool call: %s", call.function.name)
242
  result = f"Unsupported tool: {call.function.name}"
243
+ name = (
244
+ "junior" if call.function.name == "send_to_junior" else call.function.name
 
 
 
 
245
  )
246
+ messages.append({"role": "tool", "name": name, "content": result})
247
  DBMessage.create(
248
  conversation=conversation,
249
  role="tool",
 
259
 
260
  placeholder = {
261
  "role": "tool",
262
+ "name": "junior" if call.function.name == "send_to_junior" else call.function.name,
263
  "content": "Awaiting tool response...",
264
  }
265
  messages.append(placeholder)
 
283
  pass
284
  self._remove_tool_placeholder(messages)
285
  result = await exec_task
286
+ name = (
287
+ "junior" if call.function.name == "send_to_junior" else call.function.name
 
 
 
 
288
  )
289
+ messages.append({"role": "tool", "name": name, "content": result})
290
  DBMessage.create(
291
  conversation=conversation,
292
  role="tool",
 
307
  yield followup
308
  result = await exec_task
309
  self._remove_tool_placeholder(messages)
310
+ name = (
311
+ "junior" if call.function.name == "send_to_junior" else call.function.name
 
 
 
 
312
  )
313
+ messages.append({"role": "tool", "name": name, "content": result})
314
  DBMessage.create(
315
  conversation=conversation,
316
  role="tool",
src/team.py CHANGED
@@ -29,7 +29,7 @@ async def send_to_junior(message: str) -> str:
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,7 +45,7 @@ class TeamChatSession:
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,12 +79,14 @@ class TeamChatSession:
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
@@ -92,7 +94,7 @@ class TeamChatSession:
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] = []
@@ -100,7 +102,7 @@ class TeamChatSession:
100
  if part:
101
  parts.append(part)
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)
 
29
  if _TEAM is None:
30
  return "No active team"
31
 
32
+ return await _TEAM.queue_message_to_junior(message, enqueue=False)
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], bool]] = 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(
83
+ self, message: str, *, enqueue: bool = True
84
+ ) -> str:
85
  """Send ``message`` to the junior agent and wait for the reply."""
86
 
87
  loop = asyncio.get_running_loop()
88
  fut: asyncio.Future[str] = loop.create_future()
89
+ await self._to_junior.put((message, fut, enqueue))
90
  if not self._junior_task or self._junior_task.done():
91
  self._junior_task = asyncio.create_task(self._process_junior())
92
  return await fut
 
94
  async def _process_junior(self) -> None:
95
  try:
96
  while not self._to_junior.empty():
97
+ msg, fut, enqueue = await self._to_junior.get()
98
  self.junior._messages.append({"role": "tool", "name": "senior", "content": msg})
99
  DBMessage.create(conversation=self.junior._conversation, role="tool", content=msg)
100
  parts: list[str] = []
 
102
  if part:
103
  parts.append(part)
104
  result = "\n".join(parts)
105
+ if enqueue and result.strip():
106
  await self._to_senior.put(result)
107
  if not fut.done():
108
  fut.set_result(result)