src/agent/stream.py

from gradio import ChatMessage
import json
import asyncio
import re

from src.agent.mistral_agent import MistralAgent

from src.agent.utils.call import call_tool


agent = MistralAgent()
api_lock = asyncio.Lock()
tool_lock = asyncio.Lock()

with open("./prompt.md", encoding="utf-8") as f:
    SYSTEM_PROMPT = f.read()

def extract_phases(text):
    """Découpe le contenu en THINK / ACT / OBSERVE / FINAL ANSWER"""
    phases = {'think': '', 'act': '', 'observe': '', 'final': ''}
    matches = list(re.finditer(r'(THINK:|ACT:|OBSERVE:|FINAL ANSWER:)', text))

    for i, match in enumerate(matches):
        phase = match.group(1).lower().replace(":", "").replace("final answer", "final")
        start = match.end()
        end = matches[i+1].start() if i + 1 < len(matches) else len(text)
        phases[phase] = text[start:end].strip()

    return phases

async def respond(message, history=None):
    if history is None:
        history = []

    if not history or history[-1].role != "assistant" or history[-1].metadata.get("status") == "done":
        history.append(ChatMessage(role="assistant", content="", metadata={"title": "Thinking...", "status": "pending"}))
    yield history

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": message},
        {"role": "assistant", "content": "THINK: Let's start thinking, ", "prefix": True},
    ]

    phase_order = ["think", "act", "observe", "final"]
    current_phase_index = 0
    done = False

    final_full = ""
    while not done:
        current_phase = phase_order[current_phase_index]
        if current_phase != "final":
            full = ""
        else:
            full = final_full

        print('\n', '---' * 15)
        print(f">>> messages before payload [phase {current_phase_index}] :", json.dumps([m for m in messages if m.get("role") != "system"], indent=2))
        payload = {
            "agent_id": agent.agent_id,
            "messages": messages,
            "stream": True,
            "max_tokens": None,
            "tools": agent.tools,
            "tool_choice": "auto",
            "presence_penalty": 0,
            "frequency_penalty": 0,
            "n": 1
        }

        async with api_lock:
            response = await agent.client.agents.stream_async(**payload)

            async for chunk in response:
                delta = chunk.data.choices[0].delta
                content = delta.content or ""
                full += content
                if current_phase == "final":
                    final_full = full

                phases = extract_phases(full)
                buffer = phases.get(current_phase, "")
                if current_phase == "think":
                    history[-1] = ChatMessage(role="assistant", content=buffer, metadata={"title": "Thinking...", "status": "pending"})
                elif current_phase == "act":
                    history[-1] = ChatMessage(role="assistant", content=buffer, metadata={"title": "Acting...", "status": "pending"})
                elif current_phase == "observe":
                    history[-1] = ChatMessage(role="assistant", content=buffer, metadata={"title": "Observing...", "status": "pending"})
                yield history

                if current_phase == "final":
                    delta_content = delta.content or ""
                    final_full += delta_content
                    phases = extract_phases(final_full)
                    buffer = phases.get("final", "")
                    yield history
                    if delta_content == "" and buffer:
                        done = True
                        break

        if current_phase_index == 0:
            messages = [msg for msg in messages if not msg.get("prefix")]
            if buffer:
                prefix_label = current_phase.upper() if current_phase != "final" else "FINAL ANSWER"
                messages.append({
                    "role": "assistant",
                    "content": f"{prefix_label}: {buffer}\n\nACT: Let's using some tools to solve the problem.",
                    "prefix": True
                })

        elif current_phase_index == 1:
            for message in messages:
                if "prefix" in message:
                    del message["prefix"]

        if current_phase_index == 2:
            for message in messages:
                if "prefix" in message:
                    del message["prefix"]
            messages.append({
                "role": "assistant",
                "content": "OBSERVE: Based on the results, let's observe the situation and see if we need to adjust our approach.",
                "prefix": True
            })

        yield history

        if current_phase == "act":
            tool_calls = getattr(delta, "tool_calls", None)
            if tool_calls and tool_calls != [] and str(tool_calls) != "Unset()":
                async with tool_lock:
                    messages = call_tool(
                        agent,
                        tool_calls,
                        messages,
                    )
                    last_tool_response = next((m for m in reversed(messages) if m["role"] == "tool"), None)
                    if last_tool_response and last_tool_response.get("content"):
                        buffer += "\n\n" + last_tool_response["content"]
                        history[-1] = ChatMessage(role="assistant", content=buffer,  metadata={"title": "Acting...", "status": "pending"})
                yield history

        if not done:
            current_phase_index += 1
            if current_phase_index < len(phase_order):
                pass
            else:
                done = True

    observe_text = phases.get("observe", "")
    final_text = phases.get("final", "")

    if observe_text:
        history[-1] = ChatMessage(role="assistant", content=observe_text, metadata={"title": "Observing...", "status": "done"})

    if final_text:
        history.append(ChatMessage(role="assistant", content=final_text))

    yield history