8
8
from fastapi import FastAPI , Request
9
9
from fastapi .responses import JSONResponse , StreamingResponse
10
10
from pydantic import BaseModel
11
+ from asyncio import CancelledError , Task
11
12
12
13
13
14
class ChatCompletionRequest (BaseModel ):
@@ -35,14 +36,17 @@ def __init__(self, interpreter):
35
36
# Setup routes
36
37
self .app .post ("/chat/completions" )(self .chat_completion )
37
38
39
+
38
40
async def chat_completion (self , request : Request ):
39
41
"""Main chat completion endpoint"""
40
42
body = await request .json ()
43
+ if self .interpreter .debug :
44
+ print ("Request body:" , body )
41
45
try :
42
46
req = ChatCompletionRequest (** body )
43
47
except Exception as e :
44
- print ("Validation error:" , str (e )) # Debug print
45
- print ("Request body:" , body ) # Print the request body
48
+ print ("Validation error:" , str (e ))
49
+ print ("Request body:" , body )
46
50
raise
47
51
48
52
# Filter out system message
@@ -75,18 +79,6 @@ async def _stream_response(self):
75
79
delta ["function_call" ] = choice .delta .function_call
76
80
if choice .delta .tool_calls is not None :
77
81
pass
78
- # Convert tool_calls to dict representation
79
- # delta["tool_calls"] = [
80
- # {
81
- # "index": tool_call.index,
82
- # "id": tool_call.id,
83
- # "type": tool_call.type,
84
- # "function": {
85
- # "name": tool_call.function.name,
86
- # "arguments": tool_call.function.arguments
87
- # }
88
- # } for tool_call in choice.delta.tool_calls
89
- # ]
90
82
91
83
choices .append (
92
84
{
@@ -108,11 +100,16 @@ async def _stream_response(self):
108
100
data ["system_fingerprint" ] = chunk .system_fingerprint
109
101
110
102
yield f"data: { json .dumps (data )} \n \n "
111
- except asyncio .CancelledError :
112
- # Set stop flag when stream is cancelled
113
- self .interpreter ._stop_flag = True
103
+
104
+ except CancelledError :
105
+ # Handle cancellation gracefully
106
+ print ("Request cancelled - cleaning up..." )
107
+
114
108
raise
109
+ except Exception as e :
110
+ print (f"Error in stream: { str (e )} " )
115
111
finally :
112
+ # Always send DONE message and cleanup
116
113
yield "data: [DONE]\n \n "
117
114
118
115
def run (self ):
0 commit comments