Update api/utils.py

api/utils.py

@@ -12,86 +12,49 @@ from api.providers import AmigoChat
 
 logger = setup_logger(__name__)
 
-async def process_streaming_response(request: ChatRequest):
+async def process_streaming_response(request: ChatRequest) -> AsyncGenerator[str, None]:
     logger.info("Processing streaming response with AmigoChat")
     messages = [msg.dict() for msg in request.messages]
 
     try:
-        async def event_generator():
-            async for content in AmigoChat.generate_response(
-                model=request.model,
-                messages=messages,
-                stream=True
-            ):
-                timestamp = int(datetime.now().timestamp())
-                chunk = {
-                    "id": f"chatcmpl-{uuid.uuid4()}",
-                    "object": "chat.completion.chunk",
-                    "created": timestamp,
-                    "model": request.model,
-                    "choices": [
-                        {
-                            "index": 0,
-                            "delta": {"content": content},
-                            "finish_reason": None,
-                        }
-                    ],
-                }
-                yield f"data: {json.dumps(chunk)}\n\n"
-
-            # Indicate the end of the stream
-            end_chunk = {
+        async for content in AmigoChat.generate_response(
+            model=request.model,
+            messages=messages,
+            stream=True
+        ):
+            timestamp = int(datetime.now().timestamp())
+            chunk = {
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion.chunk",
-                "created": int(datetime.now().timestamp()),
+                "created": timestamp,
                 "model": request.model,
                 "choices": [
                     {
                         "index": 0,
-                        "delta": {},
-                        "finish_reason": "stop",
+                        "delta": {"content": content},
+                        "finish_reason": None,
                     }
                 ],
             }
-            yield f"data: {json.dumps(end_chunk)}\n\n"
-            yield "data: [DONE]\n\n"
-
-        return event_generator()
-
-    except Exception as e:
-        logger.error(f"Error in streaming response: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+            yield f"data: {json.dumps(chunk)}\n\n"
 
-async def process_non_streaming_response(request: ChatRequest):
-    logger.info("Processing non-streaming response with AmigoChat")
-    messages = [msg.dict() for msg in request.messages]
-
-    try:
-        responses = []
-        async for content in AmigoChat.generate_response(
-            model=request.model,
-            messages=messages,
-            stream=False
-        ):
-            responses.append(content)
-
-        full_response = ''.join(responses)
-
-        return {
+        # Indicate the end of the stream
+        end_chunk = {
             "id": f"chatcmpl-{uuid.uuid4()}",
-            "object": "chat.completion",
+            "object": "chat.completion.chunk",
             "created": int(datetime.now().timestamp()),
             "model": request.model,
             "choices": [
                 {
                     "index": 0,
-                    "message": {"role": "assistant", "content": full_response},
+                    "delta": {},
                     "finish_reason": "stop",
                 }
             ],
-            "usage": None,
         }
+        yield f"data: {json.dumps(end_chunk)}\n\n"
+        yield "data: [DONE]\n\n"
 
     except Exception as e:
-        logger.error(f"Error in non-streaming response: {e}")
+        logger.error(f"Error in streaming response: {e}")
         raise HTTPException(status_code=500, detail=str(e))