Update api/providers.py

api/providers.py

@@ -5,10 +5,11 @@ from __future__ import annotations
 import json
 import uuid
 from aiohttp import ClientSession, ClientTimeout, ClientResponseError
-
 from typing import AsyncGenerator, List, Dict, Any
-from api.models import Message
-from api.image import ImageResponse
+
+from api.logger import setup_logger
+
+logger = setup_logger(__name__)
 
 class AmigoChat:
     url = "https://amigochat.io"
@@ -116,36 +117,41 @@ class AmigoChat:
                 }
 
                 timeout = ClientTimeout(total=300)
-                async with session.post(cls.chat_api_endpoint, json=data, proxy=proxy, timeout=timeout) as response:
-                    if response.status not in (200, 201):
-                        error_text = await response.text()
-                        raise Exception(f"Error {response.status}: {error_text}")
-
-                    if stream:
-                        async for line in response.content:
-                            line = line.decode('utf-8').strip()
-                            if line.startswith('data: '):
-                                if line == 'data: [DONE]':
-                                    break
-                                try:
-                                    chunk = json.loads(line[6:])
-                                    if 'choices' in chunk and len(chunk['choices']) > 0:
-                                        choice = chunk['choices'][0]
-                                        if 'delta' in choice:
-                                            content = choice['delta'].get('content')
-                                        elif 'text' in choice:
-                                            content = choice['text']
-                                        else:
-                                            content = None
-                                        if content:
-                                            yield content
-                                except json.JSONDecodeError:
-                                    pass
-                    else:
-                        response_data = await response.json()
-                        if 'choices' in response_data and len(response_data['choices']) > 0:
-                            content = response_data['choices'][0]['message']['content']
-                            yield content
+                try:
+                    async with session.post(cls.chat_api_endpoint, json=data, proxy=proxy, timeout=timeout) as response:
+                        if response.status not in (200, 201):
+                            error_text = await response.text()
+                            raise Exception(f"Error {response.status}: {error_text}")
+
+                        if stream:
+                            async for line in response.content:
+                                line = line.decode('utf-8').strip()
+                                if line.startswith('data: '):
+                                    if line == 'data: [DONE]':
+                                        break
+                                    try:
+                                        chunk = json.loads(line[6:])
+                                        if 'choices' in chunk and len(chunk['choices']) > 0:
+                                            choice = chunk['choices'][0]
+                                            if 'delta' in choice:
+                                                content = choice['delta'].get('content')
+                                            elif 'text' in choice:
+                                                content = choice['text']
+                                            else:
+                                                content = None
+                                            if content:
+                                                yield content
+                                    except json.JSONDecodeError:
+                                        pass
+                        else:
+                            response_data = await response.json()
+                            if 'choices' in response_data and len(response_data['choices']) > 0:
+                                content = response_data['choices'][0]['message']['content']
+                                yield content
+                except Exception as e:
+                    logger.error(f"Error during request: {e}")
+                    raise
+
             else:
                 # Image generation
                 prompt = messages[-1]['content']
@@ -154,10 +160,14 @@ class AmigoChat:
                     "model": model,
                     "personaId": cls.get_persona_id(model)
                 }
-                async with session.post(cls.image_api_endpoint, json=data, proxy=proxy) as response:
-                    response.raise_for_status()
-                    response_data = await response.json()
-                    if "data" in response_data:
-                        image_urls = [item["url"] for item in response_data["data"] if "url" in item]
-                        if image_urls:
-                            yield ImageResponse(image_urls, prompt)
+                try:
+                    async with session.post(cls.image_api_endpoint, json=data, proxy=proxy) as response:
+                        response.raise_for_status()
+                        response_data = await response.json()
+                        if "data" in response_data:
+                            image_urls = [item["url"] for item in response_data["data"] if "url" in item]
+                            if image_urls:
+                                yield json.dumps({"images": image_urls, "prompt": prompt})
+                except Exception as e:
+                    logger.error(f"Error during image generation: {e}")
+                    raise