Add stream text output

src/api/endpoints.py

@@ -1,4 +1,5 @@
 from fastapi import APIRouter, HTTPException, Depends
+from fastapi.responses import StreamingResponse
 from ..schemas.models import ProcessRequest_chat, ProcessResponse_chat, ProcessRequest_get, ProcessResponse_get, ProcessRequest_set, ProcessResponse_set, ProcessResponse_delete_session, ProcessRequest_delete_session
 from fastgpt_client import AsyncChatClient
 from fastgpt_client.exceptions import (
@@ -98,15 +99,104 @@ async def delete_last_two_chat_records(
 #     logging.error(f"Validation Error: {exc.errors()}")
 #     raise HTTPException(status_code=422, detail=exc.errors())
 
+def create_sse_event(event: str, data: dict) -> str:
+    """Format data as an SSE event."""
+    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
+
 @router.post("/chat", response_model=ProcessResponse_chat)
 async def chat(
     request: ProcessRequest_chat,
+    stream: bool = False,
     client: AsyncChatClient = Depends(get_fastgpt_client)
 ):
     """Handle chat completion request."""
     json_data = request.model_dump()
-    logger.info(f"用户请求信息ProcessRequest_chat: {json_data}")
+    logger.info(f"用户请求信息ProcessRequest_chat: {json_data}, stream={stream}")
     
+    if stream:
+        async def event_generator():
+            try:
+                # Use SDK's create_chat_completion with stream=True
+                response = await client.create_chat_completion(
+                    messages=[{"role": "user", "content": json_data['text']}],
+                    chatId=json_data['sessionId'],
+                    stream=True,
+                    detail=True
+                )
+                
+                buffer = ""
+                state_code_found = False
+                
+                async for chunk in response.aiter_lines():
+                    if chunk.startswith('data: '):
+                        data_str = chunk[6:].strip()
+                        if data_str == '[DONE]':
+                            break
+                        try:
+                            data = json.loads(data_str)
+                            try:
+                                delta_content = data['choices'][0]['delta'].get('content', '')
+                            except (KeyError, IndexError):
+                                delta_content = ''
+                            if delta_content:
+                                buffer += delta_content
+                                
+                                if not state_code_found:
+                                    # Check for <state>XXXX</state> pattern
+                                    match = re.search(r"<state>(.*?)</state>", buffer)
+                                    if match:
+                                        state_code = match.group(1)
+                                        
+                                        # Apply logic to map/adjust state code
+                                        nextStageCode = state_code
+                                        if nextStageCode in ['3001', '3002', '1002']:
+                                            nextStageCode = '1002'
+                                        elif nextStageCode == '2006':
+                                            nextStageCode = '2004'
+                                        elif nextStageCode == '2017':
+                                            nextStageCode = '2016'
+                                        elif nextStageCode == '2020':
+                                            nextStageCode = '0002'
+                                        nextStage = STATUS_CODE_MAP.get(nextStageCode, '')
+                                        
+                                        # Send stage code event
+                                        yield create_sse_event("stage_code", {
+                                            "nextStageCode": nextStageCode,
+                                            "nextStage": nextStage
+                                        })
+                                        
+                                        state_code_found = True
+                                        
+                                        # Send remaining content as text_delta
+                                        remaining_content = buffer[match.end():]
+                                        if remaining_content:
+                                            yield create_sse_event("text_delta", {"text": remaining_content})
+                                        buffer = "" # Clear buffer after extracting state
+                                else:
+                                    # State code already found, just stream text
+                                    yield create_sse_event("text_delta", {"text": delta_content})
+                                    buffer = "" # Do not buffer text after state found
+                                    
+                        except json.JSONDecodeError:
+                            continue
+                        except Exception as e:
+                            print(data)
+                            logger.error(f"Error processing chunk: {e}")
+                            continue
+
+                # If stream ends and no state code found (unlikely if format is strict), 
+                # we might want to send what we have
+                if not state_code_found and buffer:
+                     yield create_sse_event("text_delta", {"text": buffer})
+
+                yield create_sse_event("done", {"status": "completed"})
+                
+            except Exception as e:
+                logger.error(f"Streaming error: {e}")
+                yield create_sse_event("error", {"msg": str(e), "code": "500"})
+
+        return StreamingResponse(event_generator(), media_type="text/event-stream")
+
     try:
         # Use SDK's create_chat_completion
         response = await client.create_chat_completion(