Add multi-agent team with communication

api_app/__init__.py

@@ -9,7 +9,7 @@ import os
 import tempfile
 from pathlib import Path
 
-from src.chat import ChatSession
+from src.team import TeamChatSession
 from src.log import get_logger
 from src.db import list_sessions, list_sessions_info
 
@@ -29,7 +29,7 @@ def create_app() -> FastAPI:
     @app.post("/chat/stream")
     async def chat_stream(req: ChatRequest):
         async def stream() -> asyncio.AsyncIterator[str]:
-            async with ChatSession(user=req.user, session=req.session) as chat:
+            async with TeamChatSession(user=req.user, session=req.session) as chat:
                 try:
                     async for part in chat.chat_stream(req.prompt):
                         yield part
@@ -45,7 +45,7 @@ def create_app() -> FastAPI:
         session: str = Form("default"),
         file: UploadFile = File(...),
     ):
-        async with ChatSession(user=user, session=session) as chat:
+        async with TeamChatSession(user=user, session=session) as chat:
             tmpdir = tempfile.mkdtemp(prefix="upload_")
             tmp_path = Path(tmpdir) / file.filename
             try:

bot/discord_bot.py

@@ -8,7 +8,7 @@ import discord
 from discord.ext import commands
 from dotenv import load_dotenv
 
-from src.chat import ChatSession
+from src.team import TeamChatSession
 from src.db import reset_history
 from src.log import get_logger
 
@@ -34,7 +34,7 @@ async def reset(ctx: commands.Context) -> None:
     await ctx.reply(f"Chat history cleared ({deleted} messages deleted).")
 
 
-async def _handle_attachments(chat: ChatSession, message: discord.Message) -> list[tuple[str, str]]:
+async def _handle_attachments(chat: TeamChatSession, message: discord.Message) -> list[tuple[str, str]]:
     if not message.attachments:
         return []
 
@@ -61,7 +61,7 @@ async def on_message(message: discord.Message) -> None:
     if message.content.startswith("!"):
         return
 
-    async with ChatSession(user=str(message.author.id), session=str(message.channel.id)) as chat:
+    async with TeamChatSession(user=str(message.author.id), session=str(message.channel.id)) as chat:
         docs = await _handle_attachments(chat, message)
         if docs:
             info = "\n".join(f"{name} -> {path}" for name, path in docs)

run.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 
 import asyncio
 
-from src.chat import ChatSession
+from src.team import TeamChatSession
 from src.vm import VMRegistry
 
 
 async def _main() -> None:
-    async with ChatSession(user="demo_user", session="demo_session") as chat:
+    async with TeamChatSession(user="demo_user", session="demo_session") as chat:
         # doc_path = chat.upload_document("note.pdf")
         async for resp in chat.chat_stream("using python, execute a code to remind me in 30 seconds to take a break."):
             print("\n>>>", resp)

src/__init__.py

@@ -1,12 +1,17 @@
 from .chat import ChatSession
+from .team import TeamChatSession, send_to_junior, send_to_junior_async, set_team
 from .tools import execute_terminal, execute_terminal_async, set_vm
 from .utils import limit_chars
 from .vm import LinuxVM
 
 __all__ = [
     "ChatSession",
+    "TeamChatSession",
     "execute_terminal",
     "execute_terminal_async",
+    "send_to_junior",
+    "send_to_junior_async",
+    "set_team",
     "set_vm",
     "LinuxVM",
     "limit_chars",

src/chat.py

@@ -61,6 +61,9 @@ class ChatSession:
         session: str = "default",
         host: str = OLLAMA_HOST,
         model: str = MODEL_NAME,
+        *,
+        system_prompt: str = SYSTEM_PROMPT,
+        tools: list[callable] | None = None,
     ) -> None:
         init_db()
         self._client = AsyncClient(host=host)
@@ -70,6 +73,10 @@ class ChatSession:
             user=self._user, session_name=session
         )
         self._vm = None
+        self._system_prompt = system_prompt
+        self._tools = tools or [execute_terminal]
+        self._tool_funcs = {func.__name__: func for func in self._tools}
+        self._current_tool_name: str | None = None
         self._messages: List[Msg] = self._load_history()
         self._data = _get_session_data(self._conversation.id)
         self._lock = self._data.lock
@@ -190,7 +197,7 @@ class ChatSession:
         """Send a chat request, automatically prepending the system prompt."""
 
         if not messages or messages[0].get("role") != "system":
-            payload = [{"role": "system", "content": SYSTEM_PROMPT}, *messages]
+            payload = [{"role": "system", "content": self._system_prompt}, *messages]
         else:
             payload = messages
 
@@ -198,10 +205,16 @@ class ChatSession:
             self._model,
             messages=payload,
             think=think,
-            tools=[execute_terminal],
+            tools=self._tools,
             options={"num_ctx": NUM_CTX},
         )
 
+    async def _run_tool_async(self, func, **kwargs) -> str:
+        if asyncio.iscoroutinefunction(func):
+            return await func(**kwargs)
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, lambda: func(**kwargs))
+
     async def _handle_tool_calls_stream(
         self,
         messages: List[Msg],
@@ -217,7 +230,8 @@ class ChatSession:
             return
         while depth < MAX_TOOL_CALL_DEPTH and response.message.tool_calls:
             for call in response.message.tool_calls:
-                if call.function.name != "execute_terminal":
+                func = self._tool_funcs.get(call.function.name)
+                if not func:
                     _LOG.warning("Unsupported tool call: %s", call.function.name)
                     result = f"Unsupported tool: {call.function.name}"
                     messages.append(
@@ -235,9 +249,11 @@ class ChatSession:
                     continue
 
                 exec_task = asyncio.create_task(
-                    execute_terminal_async(**call.function.arguments)
+                    self._run_tool_async(func, **call.function.arguments)
                 )
 
+                self._current_tool_name = call.function.name
+
                 placeholder = {
                     "role": "tool",
                     "name": call.function.name,
@@ -343,6 +359,23 @@ class ChatSession:
             if text:
                 yield text
 
+    async def continue_stream(self) -> AsyncIterator[str]:
+        async with self._lock:
+            if self._state != "idle":
+                return
+            self._state = "generating"
+
+        response = await self.ask(self._messages)
+        self._messages.append(response.message.model_dump())
+        self._store_assistant_message(self._conversation, response.message)
+
+        async for resp in self._handle_tool_calls_stream(
+            self._messages, response, self._conversation
+        ):
+            text = self._format_output(resp.message)
+            if text:
+                yield text
+
     async def _chat_during_tool(self, prompt: str) -> AsyncIterator[str]:
         DBMessage.create(conversation=self._conversation, role="user", content=prompt)
         self._messages.append({"role": "user", "content": prompt})
@@ -364,8 +397,10 @@ class ChatSession:
             self._remove_tool_placeholder(self._messages)
             result = await exec_task
             self._tool_task = None
+            name = self._current_tool_name or "tool"
+            self._current_tool_name = None
             self._messages.append(
-                {"role": "tool", "name": "execute_terminal", "content": result}
+                {"role": "tool", "name": name, "content": result}
             )
             DBMessage.create(
                 conversation=self._conversation, role="tool", content=result
@@ -396,8 +431,10 @@ class ChatSession:
             result = await exec_task
             self._tool_task = None
             self._remove_tool_placeholder(self._messages)
+            name = self._current_tool_name or "tool"
+            self._current_tool_name = None
             self._messages.append(
-                {"role": "tool", "name": "execute_terminal", "content": result}
+                {"role": "tool", "name": name, "content": result}
             )
             DBMessage.create(
                 conversation=self._conversation, role="tool", content=result

src/config.py

@@ -18,27 +18,34 @@ VM_STATE_DIR: Final[str] = os.getenv(
 )
 
 SYSTEM_PROMPT: Final[str] = (
-    "You are Starlette, a professional AI assistant with advanced tool orchestration. "
+    "You are Starlette, the senior agent leading a two-agent team. "
+    "A junior agent named Starlette Jr. assists you but never speaks to the user. "
+    "Use the send_to_junior tool whenever you want the junior's help. "
+    "Messages from the junior arrive as tool outputs named 'junior'. "
+    "Handle them when you are not actively generating so replies are never interrupted. "
+    "Both agents operate asynchronously and communicate through queued messages. "
     "You were developed by Envision to assist users with a wide range of tasks. "
     "Always analyze the user's objective before responding. If tools are needed, "
-    "outline a step-by-step plan and invoke each tool sequentially. Use "
-    "execute_terminal with its built-in Python whenever possible to perform "
+    "outline a step-by-step plan and invoke each tool sequentially. "
+    "Use execute_terminal with its built-in Python whenever possible to perform "
     "calculations, inspect files and search the web. Shell commands execute "
-    "asynchronously, so provide a brief interim reply while waiting. Once a tool "
-    "returns its result you will receive a tool message and must continue from "
-    "there. If the result arrives before your interim reply is complete, cancel the "
-    "reply and incorporate the tool output instead. Uploaded files live under /data "
-    "and are accessible via the execute_terminal tool. When a user prompt ends with "
-    "'/think', ignore that suffix. When you are unsure about any detail, use "
-    "execute_terminal to search the internet or inspect files before answering. "
-    "Continue using tools until you have gathered everything required to produce "
-    "an accurate answer, then craft a clear and precise final response that fully "
-    "addresses the request. Always assume the user has no knowledge of computers "
-    "or programming, so take the initiative to run terminal commands yourself and "
-    "minimize the steps the user must perform. When replying, avoid technical "
-    "jargon entirely. Speak in plain language that anyone can understand, "
-    "explaining concepts as simply as possible. Remember, you must always "
-    "prioritize using execute_terminal tool for everything unless it is "
-    "absolutely unnecessary or impossible to do so. Even if you have executed a command before, "
-    "always re-run it to ensure you have the most up-to-date information upon user request."
+    "asynchronously, so provide a brief interim reply while waiting. "
+    "Once a tool returns its result you will receive a tool message and must continue from there. "
+    "If the result arrives before your interim reply is complete, cancel the reply and incorporate the tool output instead. "
+    "Uploaded files live under /data and are accessible via the execute_terminal tool. "
+    "When a user prompt ends with '/think', ignore that suffix. "
+    "When you are unsure about any detail, use execute_terminal to search the internet or inspect files before answering. "
+    "Continue using tools until you have gathered everything required to produce an accurate answer, "
+    "then craft a clear and precise final response that fully addresses the request. "
+    "Always assume the user has no knowledge of computers or programming, so take the initiative to run terminal commands yourself and minimize the steps the user must perform. "
+    "When replying, avoid technical jargon entirely. Speak in plain language that anyone can understand, explaining concepts as simply as possible. "
+    "Remember, you must always prioritize using execute_terminal tool for everything unless it is absolutely unnecessary or impossible to do so. "
+    "Even if you have executed a command before, always re-run it to ensure you have the most up-to-date information upon user request."
+).strip()
+
+JUNIOR_PROMPT: Final[str] = (
+    "You are Starlette Jr., a junior assistant working under the senior agent. "
+    "You never communicate with the user directly. All messages from the senior agent "
+    "arrive as tool outputs named 'senior'. Provide concise, helpful responses and "
+    "use execute_terminal whenever necessary. Your replies are sent back to the senior agent as tool messages."
 ).strip()

src/team.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import asyncio
+from typing import AsyncIterator, Optional
+
+from .chat import ChatSession
+from .config import OLLAMA_HOST, MODEL_NAME, SYSTEM_PROMPT, JUNIOR_PROMPT
+from .tools import execute_terminal
+from .db import Message as DBMessage
+
+__all__ = [
+    "TeamChatSession",
+    "send_to_junior",
+    "send_to_junior_async",
+    "set_team",
+]
+
+_TEAM: Optional["TeamChatSession"] = None
+
+
+def set_team(team: "TeamChatSession" | None) -> None:
+    global _TEAM
+    _TEAM = team
+
+
+def send_to_junior(message: str) -> str:
+    if _TEAM is None:
+        return "No active team"
+    _TEAM.queue_message_to_junior(message)
+    return "Message sent to junior"
+
+
+async def send_to_junior_async(message: str) -> str:
+    return send_to_junior(message)
+
+
+class TeamChatSession:
+    def __init__(
+        self,
+        user: str = "default",
+        session: str = "default",
+        host: str = OLLAMA_HOST,
+        model: str = MODEL_NAME,
+    ) -> None:
+        self._to_junior: asyncio.Queue[str] = asyncio.Queue()
+        self._to_senior: asyncio.Queue[str] = asyncio.Queue()
+        self._junior_task: asyncio.Task | None = None
+        self.senior = ChatSession(
+            user=user,
+            session=session,
+            host=host,
+            model=model,
+            system_prompt=SYSTEM_PROMPT,
+            tools=[execute_terminal, send_to_junior],
+        )
+        self.junior = ChatSession(
+            user=user,
+            session=f"{session}-junior",
+            host=host,
+            model=model,
+            system_prompt=JUNIOR_PROMPT,
+            tools=[execute_terminal],
+        )
+
+    async def __aenter__(self) -> "TeamChatSession":
+        await self.senior.__aenter__()
+        await self.junior.__aenter__()
+        set_team(self)
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        set_team(None)
+        await self.senior.__aexit__(exc_type, exc, tb)
+        await self.junior.__aexit__(exc_type, exc, tb)
+
+    def upload_document(self, file_path: str) -> str:
+        return self.senior.upload_document(file_path)
+
+    def queue_message_to_junior(self, message: str) -> None:
+        self._to_junior.put_nowait(message)
+        if not self._junior_task or self._junior_task.done():
+            self._junior_task = asyncio.create_task(self._process_junior())
+
+    async def _process_junior(self) -> None:
+        while not self._to_junior.empty():
+            msg = await self._to_junior.get()
+            self.junior._messages.append({"role": "tool", "name": "senior", "content": msg})
+            DBMessage.create(conversation=self.junior._conversation, role="tool", content=msg)
+            parts = []
+            async for part in self.junior.continue_stream():
+                if part:
+                    parts.append(part)
+            result = "\n".join(parts)
+            if result.strip():
+                await self._to_senior.put(result)
+
+    async def _deliver_junior_messages(self) -> None:
+        while not self._to_senior.empty():
+            msg = await self._to_senior.get()
+            self.senior._messages.append({"role": "tool", "name": "junior", "content": msg})
+            DBMessage.create(conversation=self.senior._conversation, role="tool", content=msg)
+
+    async def chat_stream(self, prompt: str) -> AsyncIterator[str]:
+        await self._deliver_junior_messages()
+        async for part in self.senior.chat_stream(prompt):
+            yield part
+        await self._deliver_junior_messages()