feat: stop chat

.gitignore

@@ -1 +1,2 @@
 venv
+local-dev.sh

README.md

@@ -1,5 +1,5 @@
 ---
-title: ODR Demo - Async SSE Chat Interface
+title: ODR Demo
 emoji: 🤖
 colorFrom: blue
 colorTo: purple
@@ -8,7 +8,7 @@ sdk_version: 5.43.1
 app_file: app.py
 pinned: false
 license: mit
-short_description: 一个演示项目，包含异步 SSE（服务器发送事件）接口请求功能，支持实时流式聊天体验
+short_description: a demo
 tags:
 - chatbot
 - sse
@@ -19,151 +19,4 @@ tags:
 suggested_hardware: cpu-basic
 hf_oauth: false
 disable_embedding: false
----
-
-# ODR Demo
-
-一个演示项目，包含异步 SSE（服务器发送事件）接口请求功能，支持实时流式聊天体验。
-
-## 🌟 特性
-
-- 异步 SSE 流式请求
-- 支持自定义参数（深度思考模式、调试模式等）
-- 返回异步迭代器进行实时数据处理
-- 支持原始数据和解析 JSON 数据模式
-- 结构化事件解析，包含 `event` 和 `data` 字段
-- 正确的 SSE 事件类型检测和处理
-- 实时 Markdown 渲染和可折叠工具调用显示
-
-## 🚀 快速开始
-
-这个应用部署在 Hugging Face Spaces 上，你可以直接访问 Web 界面进行交互：
-
-1. 在输入框中输入你的英文问题
-2. 点击 "Run" 按钮开始处理
-3. 观察实时的 Agent 和工具调用过程
-4. 可以随时点击 "Stop" 按钮停止处理
-
-## 📦 本地安装
-
-如果你想在本地运行此项目：
-
-```bash
-# 克隆项目
-git clone <repository-url>
-cd odr-demo
-
-# 创建虚拟环境
-python -m venv venv
-source venv/bin/activate  # 在 Windows 上使用: venv\Scripts\activate
-
-# 安装依赖
-pip install gradio aiohttp
-
-# 设置环境变量
-export API_ENDPOINT="your-api-endpoint-here"
-
-# 运行应用
-python app.py
-```
-
-## 💻 基本用法
-
-### 原始 SSE 事件
-
-```python
-import asyncio
-from api import request_sse_stream
-
-async def main():
-    query = "Hello, please introduce Python"
-    
-    async for event_data in request_sse_stream(query):
-        event_type = event_data.get('event', 'unknown')
-        data_content = event_data.get('data', '')
-        print(f"Event: {event_type}")
-        print(f"Data: {data_content}")
-
-asyncio.run(main())
-```
-
-### 解析后的数据
-
-```python
-import asyncio
-from api import request_sse_stream_parsed
-
-async def main():
-    query = "What is machine learning?"
-    
-    async for event_data in request_sse_stream_parsed(query):
-        event_type = event_data.get('event', 'unknown')
-        parsed_data = event_data.get('data', {})
-        print(f"Event: {event_type}")
-        print(f"Parsed Data: {parsed_data}")
-
-asyncio.run(main())
-```
-
-### 使用类方法（更多控制）
-
-```python
-import asyncio
-from api import SSEClient
-
-async def main():
-    client = SSEClient()
-    
-    async for event_data in client.stream_chat(
-        query="Explain deep learning",
-        deep_thinking_mode=True,  # 启用深度思考
-        debug=True,              # 启用调试模式
-        chat_id="my_custom_id"   # 自定义聊天 ID
-    ):
-        event_type = event_data.get('event', 'unknown')
-        data_content = event_data.get('data', '')
-        print(f"Event: {event_type}")
-        print(f"Data: {data_content}")
-
-asyncio.run(main())
-```
-
-## ⚙️ API 参数
-
-- `query`: 必需，用户查询内容
-- `deep_thinking_mode`: 可选，是否启用深度思考模式，默认 False
-- `search_before_planning`: 可选，是否在规划前搜索，默认 False
-- `debug`: 可选，是否启用调试模式，默认 False
-- `chat_id`: 可选，聊天 ID，如未提供将自动生成
-
-## 📊 数据结构
-
-所有函数返回异步迭代器，产生以下结构的字典：
-
-```python
-{
-    "event": "message",  # SSE 事件类型（如 "message", "error", "data" 等）
-    "data": "..."        # 事件数据内容
-}
-```
-
-- `request_sse_stream()`: 返回原始数据，`data` 字段包含原始字符串
-- `request_sse_stream_parsed()`: 返回解析后的数据，`data` 字段包含 JSON 对象（如果可能）
-
-## 📁 文件说明
-
-- `api.py`: 主要的 SSE 客户端实现
-- `app.py`: Gradio Web 应用程序
-- `utils.py`: 工具函数
-- `README.md`: 项目文档
-
-## 🎯 注意事项
-
-- 目前仅支持英文输入
-- 需要设置有效的 `API_ENDPOINT` 环境变量
-- 实时显示包括代理处理过程和工具调用详情
-- 支持错误处理和中断操作
-
-## 📄 许可证
-
-MIT License
+---

api.py

@@ -4,6 +4,9 @@ import os
 import uuid
 from typing import AsyncIterator, Dict, Any
 import aiohttp
+import logging
+
+logger = logging.getLogger(__name__)
 
 
 class SSEClient:
@@ -164,6 +167,15 @@ async def request_sse_stream_parsed(query: str, **kwargs) -> AsyncIterator[Dict[
         yield event_data
 
 
+async def stop_chat(chat_id: str):
+    url = f"{os.getenv('STOP_CHAT_API_ENDPOINT')}"
+    async with aiohttp.ClientSession() as session:
+        async with session.post(url, json={"chatId": chat_id}) as response:
+            if response.status != 200:
+                logger.error(f"Request failed with status code: {response.status}")
+                raise Exception(f"Request failed with status code: {response.status}")
+            return await response.json()
+
 # Example usage
 async def main():
     """Example usage method"""

app.py

@@ -1,15 +1,18 @@
 import json
 import logging
 import uuid
+import asyncio
 from typing import  Optional
 import gradio as gr
 
-from api import request_sse_stream_parsed
+from api import request_sse_stream_parsed, stop_chat
 
 from utils import contains_chinese, replace_chinese_punctuation
 
 logger = logging.getLogger(__name__)
 
+running_tasks = {}
+
 
 from typing import Optional
 
@@ -89,10 +92,14 @@ def _render_markdown(state: dict) -> str:
         lines.append("\n---\n")
     return "\n".join(lines) if lines else "Waiting..."
 
-def _update_state_with_event(state: dict, message: dict):
+def _update_state_with_event(ui_state: dict, state: dict, message: dict):
     event = message.get("event")
     data = message.get("data", {})
-    if event == "start_of_agent":
+    if event == "job_started":
+        chat_id = data.get("chat_id")
+        state["chat_id"] = chat_id
+        ui_state["chat_id"] = chat_id
+    elif event == "start_of_agent":
         agent_id = data.get("agent_id")
         agent_name = data.get("agent_name", "unknown")
         if agent_id and agent_id not in state["agents"]:
@@ -189,15 +196,17 @@ async def gradio_run(query: str, ui_state: Optional[dict]):
             "we only support English input for the time being.",
             gr.update(interactive=True),
             gr.update(interactive=False),
-            ui_state or {"task_id": None}
+            ui_state or {"chat_id": None}
         )
         return
-    task_id = str(uuid.uuid4())
-    if not ui_state:
-        ui_state = {"task_id": task_id}
-    else:
-        ui_state = {**ui_state, "task_id": task_id}
+
     state = _init_render_state()
+    
+    task_id = str(uuid.uuid4())
+    if ui_state is None:
+        ui_state = {"chat_id": None}
+    ui_state["task_id"] = task_id
+    
     # Initial: disable Run, enable Stop, and show spinner at bottom of text
     yield (
         _render_markdown(state) + _spinner_markup(True),
@@ -205,15 +214,36 @@ async def gradio_run(query: str, ui_state: Optional[dict]):
         gr.update(interactive=True),
         ui_state
     )
-    async for message in request_sse_stream_parsed(query):
-        state = _update_state_with_event(state, message)
-        md = _render_markdown(state)
+    
+    try:
+        current_task = asyncio.current_task()
+        running_tasks[task_id] = current_task
+        
+        async for message in request_sse_stream_parsed(query):
+            if current_task.cancelled():
+                break
+                
+            state = _update_state_with_event(ui_state, state, message)
+            md = _render_markdown(state)
+            yield (
+                md + _spinner_markup(True),
+                gr.update(interactive=False),
+                gr.update(interactive=True),
+                ui_state
+            )
+    except asyncio.CancelledError:
+        state.setdefault("errors", []).append("Task has been cancelled")
         yield (
-            md + _spinner_markup(True),
-            gr.update(interactive=False),
+            _render_markdown(state),
             gr.update(interactive=True),
+            gr.update(interactive=False),
             ui_state
         )
+        return
+    finally:
+        if task_id in running_tasks:
+            del running_tasks[task_id]
+    
     # End: enable Run, disable Stop, remove spinner
     yield (
         _render_markdown(state),
@@ -222,8 +252,25 @@ async def gradio_run(query: str, ui_state: Optional[dict]):
         ui_state
     )
 
-def stop_current(ui_state: Optional[dict]):
-    # Immediately switch button availability: enable Run, disable Stop
+async def stop_current(ui_state: Optional[dict]):
+    if ui_state is None:
+        ui_state = {}
+    
+    task_id = ui_state.get("task_id")
+    if task_id and task_id in running_tasks:
+        task = running_tasks[task_id]
+        if task and not task.done():
+            task.cancel()
+            logger.info(f"Task has been cancelled: {task_id}")
+    
+    chat_id = ui_state.get("chat_id")
+    if chat_id:
+        try:
+            res = await stop_chat(chat_id)
+            logger.info(f"Chat has been stopped: {chat_id}, res: {res}")
+        except Exception as e:
+            logger.error(f"Stop chat API call failed: {e}")
+    
     return (
         gr.update(interactive=True),
         gr.update(interactive=False),
@@ -244,7 +291,7 @@ def build_demo():
             run_btn = gr.Button("Run")
             stop_btn = gr.Button("Stop", variant="stop", interactive=False)
         out_md = gr.Markdown("", elem_id="log-view")
-        ui_state = gr.State({"task_id": None})
+        ui_state = gr.State({"chat_id": None})
         # run: outputs -> markdown, run_btn(update), stop_btn(update), ui_state
         run_btn.click(fn=gradio_run, inputs=[inp, ui_state], outputs=[out_md, run_btn, stop_btn, ui_state])
         # stop: outputs -> run_btn(update), stop_btn(update)