Update api/utils.py

api/utils.py

@@ -1,20 +1,93 @@
+# api/utils.py
+
+from datetime import datetime
+import json
+import uuid
+from typing import Any, Dict, Optional, AsyncGenerator, List
+
 from fastapi import HTTPException
-from api.config import MODEL_PROVIDER_MAPPING
-from api.provider import amigochat
+from api.models import ChatRequest, Message
 from api.logger import setup_logger
+from api.providers import AmigoChat
 
 logger = setup_logger(__name__)
 
-async def process_streaming_response(request_data):
-    provider_name = MODEL_PROVIDER_MAPPING.get(request_data.get('model'))
-    if provider_name == 'amigochat':
-        return await amigochat.process_streaming_response(request_data)
-    else:
-        raise HTTPException(status_code=400, detail=f"Model {request_data.get('model')} is not supported for streaming.")
-
-async def process_non_streaming_response(request_data):
-    provider_name = MODEL_PROVIDER_MAPPING.get(request_data.get('model'))
-    if provider_name == 'amigochat':
-        return await amigochat.process_non_streaming_response(request_data)
-    else:
-        raise HTTPException(status_code=400, detail=f"Model {request_data.get('model')} is not supported.")
+async def process_streaming_response(request: ChatRequest):
+    logger.info("Processing streaming response with AmigoChat")
+    messages = [msg.dict() for msg in request.messages]
+
+    try:
+        async for content in AmigoChat.generate_response(
+            model=request.model,
+            messages=messages,
+            stream=True
+        ):
+            timestamp = int(datetime.now().timestamp())
+            yield json.dumps({
+                "id": f"chatcmpl-{uuid.uuid4()}",
+                "object": "chat.completion.chunk",
+                "created": timestamp,
+                "model": request.model,
+                "choices": [
+                    {
+                        "index": 0,
+                        "delta": {"content": content},
+                        "finish_reason": None,
+                    }
+                ],
+                "usage": None,
+            }) + "\n"
+
+        # Indicate the end of the stream
+        yield json.dumps({
+            "id": f"chatcmpl-{uuid.uuid4()}",
+            "object": "chat.completion.chunk",
+            "created": int(datetime.now().timestamp()),
+            "model": request.model,
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": None,
+        }) + "\n"
+
+    except Exception as e:
+        logger.error(f"Error in streaming response: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+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 {
+            "id": f"chatcmpl-{uuid.uuid4()}",
+            "object": "chat.completion",
+            "created": int(datetime.now().timestamp()),
+            "model": request.model,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": full_response},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": None,
+        }
+
+    except Exception as e:
+        logger.error(f"Error in non-streaming response: {e}")
+        raise HTTPException(status_code=500, detail=str(e))