Update api/providers.py

api/providers.py

@@ -4,8 +4,8 @@ from __future__ import annotations
 
 import json
 import uuid
-from aiohttp import ClientSession, ClientTimeout, ClientResponseError
-from typing import AsyncGenerator, List, Dict, Any
+from aiohttp import ClientSession, ClientTimeout
+from typing import AsyncGenerator, List, Dict, Any, Union
 
 from api.logger import setup_logger
 
@@ -77,7 +77,7 @@ class AmigoChat:
         messages: List[Dict[str, Any]],
         stream: bool = False,
         proxy: str = None,
-    ) -> AsyncGenerator[str, None]:
+    ) -> Union[AsyncGenerator[str, None], str]:
         model = cls.get_model(model)
         device_uuid = str(uuid.uuid4())
 
@@ -103,7 +103,6 @@ class AmigoChat:
 
         async with ClientSession(headers=headers) as session:
             if model in cls.chat_models:
-                # Chat completion
                 data = {
                     "messages": messages,
                     "model": model,
@@ -124,50 +123,37 @@ class AmigoChat:
                             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
+                            async def stream_content():
+                                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
+                            return stream_content()
                         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
+                                return content
+                            else:
+                                return ""
                 except Exception as e:
                     logger.error(f"Error during request: {e}")
                     raise
-
             else:
-                # Image generation
-                prompt = messages[-1]['content']
-                data = {
-                    "prompt": prompt,
-                    "model": model,
-                    "personaId": cls.get_persona_id(model)
-                }
-                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
+                # Handle image models or other cases if necessary
+                pass