feat: thinking + streaming

src/agent/inference.py

@@ -32,131 +32,164 @@ class MistralAgent:
             ]
         ).get('tools')
 
-    def make_initial_request(self, input):
-        """Make the initial request to the agent with the given input."""
-        with open("./prompt.md", 'r', encoding='utf-8') as file:
-            self.prompt = file.read()
-        messages = [
-            {"role": "system", "content": self.prompt},
-            {"role": "user", "content": input},
-            {
-                "role": "assistant",
-                "content": "THINKING:\nLet's tackle this problem, ",
-                "prefix": True,
-            },
-        ]
-        payload = {
-            "agent_id": self.agent_id,
-            "messages": messages,
-            "max_tokens": None,
-            "stream": False,
-            "stop": None,
-            "random_seed": None,
-            "response_format": None,
-            "tools": self.tools,
-            "tool_choice": 'auto',
-            "presence_penalty": 0,
-            "frequency_penalty": 0,
-            "n": 1,
-            "prediction": None,
-            "parallel_tool_calls": None
-        }
-        return self.client.agents.complete(**payload), messages
-
-    def run(self, input):
-        """Run the agent with the given input and process the response."""
-        print("\n===== Asking the agent =====\n")
-        response, messages = self.make_initial_request(input)
-        first_iteration = True
-
-        while True:
-            time.sleep(1)
-            if hasattr(response, 'choices') and response.choices:
-                choice = response.choices[0]
-
-                if first_iteration:
-                    messages = [message for message in messages if not message.get("prefix")]
-                    messages.append(
-                        {
-                            "role": "assistant",
-                            "content": choice.message.content,
-                            "prefix": True,
-                        },
-                    )
-                    first_iteration = False
-                else:
-                    if choice.message.tool_calls:
-                        results = []
-
-                        for tool_call in choice.message.tool_calls:
-                            function_name = tool_call.function.name
-                            function_params = json.loads(tool_call.function.arguments)
-
-                            try:
-                                function_result = self.names_to_functions[function_name](**function_params)
-                                results.append((tool_call.id, function_name, function_result))
-
-                            except Exception as e:
-                                results.append((tool_call.id, function_name, None))
+    #def make_initial_request(self, input):
+    #    """Make the initial request to the agent with the given input."""
+    #    with open("./prompt.md", 'r', encoding='utf-8') as file:
+    #        self.prompt = file.read()
+    #    messages = [
+    #        {"role": "system", "content": self.prompt},
+    #        {"role": "user", "content": input},
+    #        {
+    #            "role": "assistant",
+    #            "content": "THINKING:\nLet's tackle this problem, ",
+    #            "prefix": True,
+    #        },
+    #    ]
+    #    payload = {
+    #        "agent_id": self.agent_id,
+    #        "messages": messages,
+    #        "max_tokens": None,
+    #        "stream": True,
+    #        "stop": None,
+    #        "random_seed": None,
+    #        "response_format": None,
+    #        "tools": self.tools,
+    #        "tool_choice": 'auto',
+    #        "presence_penalty": 0,
+    #        "frequency_penalty": 0,
+    #        "n": 1,
+    #        "prediction": None,
+    #        "parallel_tool_calls": None
+    #    }
+    #    stream = self.client.agents.complete(**payload)
+    #    return stream, messages
+#
+    #def run(self, input):
+    #    """Run the agent with the given input and process the response."""
+    #    print("\n===== Asking the agent =====\n")
+    #    stream, messages = self.make_initial_request(input)
+#
+    #    for data in stream:
+    #        # Si `stream` renvoie des chaînes brutes de type `data: {...}`
+    #        if isinstance(data, str) and data.startswith("data: "):
+    #            try:
+    #                json_str = data[len("data: "):].strip()
+    #                if json_str == "[DONE]":
+    #                    break
+    #                chunk = json.loads(json_str)
+    #                delta = chunk.get("choices", [{}])[0].get("delta", {})
+    #                content = delta.get("content")
+    #                if content:
+    #                    yield content
+#
+    #                # Fin de réponse
+    #                if chunk["choices"][0].get("finish_reason") is not None:
+    #                    break
+    #            except json.JSONDecodeError:
+    #                continue
+#
+    #        # Si `stream` donne directement des dicts (selon ton client)
+    #        elif isinstance(data, dict):
+    #            delta = data.get("choices", [{}])[0].get("delta", {})
+    #            content = delta.get("content")
+    #            if content:
+    #                yield content
+#
+    #            if data["choices"][0].get("finish_reason") is not None:
+    #                break
 
-                        for tool_call_id, function_name, function_result in results:
-                            messages.append({
-                                "role": "assistant",
-                                "tool_calls": [
-                                    {
-                                        "id": tool_call_id,
-                                        "type": "function",
-                                        "function": {
-                                            "name": function_name,
-                                            "arguments": json.dumps(function_params),
-                                        }
-                                    }
-                                ]
-                            })
-                            messages.append(
-                                {
-                                    "role": "tool",
-                                    "content": function_result if function_result is not None else f"Error occurred: {function_name} failed to execute",
-                                    "tool_call_id": tool_call_id,
-                                },
-                            )
-                            for message in messages:
-                                if "prefix" in message:
-                                    del message["prefix"]
-                            messages.append(
-                                {
-                                    "role": "assistant",
-                                    "content": f"Based on the results, ",
-                                    "prefix": True,
-                                }
-                            )
-                    else:
-                        for message in messages:
-                            if "prefix" in message:
-                                del message["prefix"]
-                        messages.append(
-                            {
-                                "role": "assistant",
-                                "content": choice.message.content,
-                            }
-                        )
-                        if 'FINAL ANSWER:' in choice.message.content:
-                            print("\n===== END OF REQUEST =====\n", json.dumps(messages, indent=2))
-                            ans = choice.message.content.split('FINAL ANSWER:')[1].strip()
 
-                            timestamp = time.strftime("%Y%m%d-%H%M%S")
-                            output_file = f"chat_{timestamp}.json"
-                            with open(output_file, "w", encoding="utf-8") as f:
-                                json.dump(messages, f, indent=2, ensure_ascii=False)
-                            print(f"Conversation enregistrée dans {output_file}")
 
-                            return ans
+        #first_iteration = True
 
-                print("\n===== MESSAGES BEFORE API CALL =====\n", json.dumps(messages, indent=2))
-                time.sleep(1)
-                response = self.client.agents.complete(
-                    agent_id=self.agent_id,
-                    messages=messages,
-                    tools=self.tools,
-                    tool_choice='auto',
-                )
+        #while True:
+        #    time.sleep(1)
+        #    if hasattr(response, 'choices') and response.choices:
+        #        choice = response.choices[0]
+#
+        #        if first_iteration:
+        #            messages = [message for message in messages if not message.get("prefix")]
+        #            messages.append(
+        #                {
+        #                    "role": "assistant",
+        #                    "content": choice.message.content,
+        #                    "prefix": True,
+        #                },
+        #            )
+        #            first_iteration = False
+        #        else:
+        #            if choice.message.tool_calls:
+        #                results = []
+#
+        #                for tool_call in choice.message.tool_calls:
+        #                    function_name = tool_call.function.name
+        #                    function_params = json.loads(tool_call.function.arguments)
+#
+        #                    try:
+        #                        function_result = self.names_to_functions[function_name](**function_params)
+        #                        results.append((tool_call.id, function_name, function_result))
+#
+        #                    except Exception as e:
+        #                        results.append((tool_call.id, function_name, None))
+#
+        #                for tool_call_id, function_name, function_result in results:
+        #                    messages.append({
+        #                        "role": "assistant",
+        #                        "tool_calls": [
+        #                            {
+        #                                "id": tool_call_id,
+        #                                "type": "function",
+        #                                "function": {
+        #                                    "name": function_name,
+        #                                    "arguments": json.dumps(function_params),
+        #                                }
+        #                            }
+        #                        ]
+        #                    })
+        #                    messages.append(
+        #                        {
+        #                            "role": "tool",
+        #                            "content": function_result if function_result is not None else f"Error occurred: {function_name} failed to execute",
+        #                            "tool_call_id": tool_call_id,
+        #                        },
+        #                    )
+        #                    for message in messages:
+        #                        if "prefix" in message:
+        #                            del message["prefix"]
+        #                    messages.append(
+        #                        {
+        #                            "role": "assistant",
+        #                            "content": f"Based on the results, ",
+        #                            "prefix": True,
+        #                        }
+        #                    )
+        #            else:
+        #                for message in messages:
+        #                    if "prefix" in message:
+        #                        del message["prefix"]
+        #                messages.append(
+        #                    {
+        #                        "role": "assistant",
+        #                        "content": choice.message.content,
+        #                    }
+        #                )
+        #                if 'FINAL ANSWER:' in choice.message.content:
+        #                    print("\n===== END OF REQUEST =====\n", json.dumps(messages, indent=2))
+        #                    ans = choice.message.content.split('FINAL ANSWER:')[1].strip()
+#
+        #                    timestamp = time.strftime("%Y%m%d-%H%M%S")
+        #                    output_file = f"chat_{timestamp}.json"
+        #                    with open(output_file, "w", encoding="utf-8") as f:
+        #                        json.dump(messages, f, indent=2, ensure_ascii=False)
+        #                    print(f"Conversation enregistrée dans {output_file}")
+#
+        #                    return ans
+#
+        #        print("\n===== MESSAGES BEFORE API CALL =====\n", json.dumps(messages, indent=2))
+        #        time.sleep(1)
+        #        response = self.client.agents.complete(
+        #            agent_id=self.agent_id,
+        #            messages=messages,
+        #            tools=self.tools,
+        #            tool_choice='auto',
+        #        )

src/ui/sidebar.py

@@ -1,10 +1,84 @@
 import gradio as gr
+import json
+from gradio import ChatMessage
 
 from src.agent.inference import MistralAgent
 
-def respond(gr_message, history=None):
-    agent = MistralAgent()
-    yield agent.run(gr_message)
+agent = MistralAgent()
+
+async def respond(message, history=None):
+
+    if history is None:
+        history = []
+    history.append(ChatMessage(role="user", content=message))
+
+    thinking_msg = ChatMessage(
+        role="assistant",
+        content="",
+        metadata={"title": "Thinking", "status": "pending"}
+    )
+    history.append(thinking_msg)
+    yield history
+
+    with open("./prompt.md", encoding="utf-8") as f:
+        prompt = f.read()
+
+    messages = [
+        {"role": "system", "content": prompt},
+        {"role": "user", "content": message},
+        #{
+        #    "role": "assistant",
+        #    "content": "THINKING:\nLet's tackle this problem",
+        ##    "prefix": True
+        #},
+    ]
+    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
+    }
+
+    response = await agent.client.agents.stream_async(**payload)
+
+    full = ""
+    thinking = ""
+    final = ""
+
+    async for chunk in response:
+        delta = chunk.data.choices[0].delta
+        content = delta.content or ""
+        full += content
+
+        if "FINAL ANSWER:" in full:
+            parts = full.split("FINAL ANSWER:", 1)
+            thinking = parts[0].replace("THINKING:", "").strip()
+            final = parts[1].strip()
+        else:
+            thinking = full.strip()
+            final = ""
+
+        history[-1] = ChatMessage(
+            role="assistant",
+            content=thinking,
+            metadata={"title": "Thinking", "status": "pending"}
+        )
+        yield history
+
+    history[-1] = ChatMessage(
+        role="assistant",
+        content=thinking,
+        metadata={"title": "Thinking", "status": "done"}
+    )
+
+    history.append(ChatMessage(role="assistant", content=final))
+    yield history
+
 
 
 def sidebar_ui(state, width=700, visible=True):