Update main.py

main.py

@@ -99,15 +99,12 @@ class ModelNotWorkingException(Exception):
         self.message = f"The model '{model}' is currently not working. Please try another model or wait for it to be fixed."
         super().__init__(self.message)
 
-# Mock implementations for ImageResponse and to_data_uri
+# ImageResponse class
 class ImageResponse:
-    def __init__(self, images: str, alt: str):
-        self.images = images
+    def __init__(self, url: str, alt: str):
+        self.url = url
         self.alt = alt
 
-def to_data_uri(image: Any) -> str:
-    return "data:image/png;base64,..."  # Replace with actual base64 data
-
 # Placeholder classes for AsyncGeneratorProvider and ProviderModelMixin
 class AsyncGeneratorProvider:
     pass  # Implement as per your actual provider's requirements
@@ -390,31 +387,35 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
 
         async with ClientSession(headers=common_headers) as session:
             try:
-                # Send initial chat request
+                # Send initial chat request with streaming
                 async with session.post(
                     cls.api_endpoint,
                     headers=headers_api_chat_combined,
                     json=payload_api_chat,
-                    proxy=proxy
+                    proxy=proxy,
+                    timeout=ClientTimeout(total=600),  # Adjust timeout as needed
                 ) as response_api_chat:
                     response_api_chat.raise_for_status()
-                    text = await response_api_chat.text()
-                    cleaned_response = cls.clean_response(text)
-
-                    if model in cls.image_models:
-                        match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
-                        if match:
-                            image_url = match.group(1)
-                            image_response = ImageResponse(images=image_url, alt="Generated Image")
+                    
+                    # Stream the response in chunks
+                    async for data in response_api_chat.content.iter_chunked(1024):
+                        decoded_data = data.decode('utf-8', errors='ignore')
+                        cleaned_data = cls.clean_response(decoded_data)
+                        
+                        # Check for image response
+                        image_match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_data)
+                        if image_match:
+                            image_url = image_match.group(1)
+                            image_response = ImageResponse(url=image_url, alt="Generated Image")
                             yield image_response
-                        else:
-                            yield cleaned_response
-                    else:
+                            continue  # Continue to the next chunk
+
+                        # Check for web search sources
                         if websearch:
-                            match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
-                            if match:
-                                source_part = match.group(1).strip()
-                                answer_part = cleaned_response[match.end():].strip()
+                            source_match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_data, re.DOTALL)
+                            if source_match:
+                                source_part = source_match.group(1).strip()
+                                answer_part = cleaned_data[source_match.end():].strip()
                                 try:
                                     sources = json.loads(source_part)
                                     source_formatted = "**Sources:**\n"
@@ -426,15 +427,12 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                                     final_response = f"{answer_part}\n\n{source_formatted}"
                                 except json.JSONDecodeError:
                                     final_response = f"{answer_part}\n\nSource information is unavailable."
-                            else:
-                                final_response = cleaned_response
-                        else:
-                            if '$~~~$' in cleaned_response:
-                                final_response = cleaned_response.split('$~~~$')[0].strip()
-                            else:
-                                final_response = cleaned_response
+                                yield final_response
+                                continue  # Continue to the next chunk
 
-                        yield final_response
+                        # Yield the cleaned data chunk
+                        if cleaned_data.strip():
+                            yield cleaned_data.strip()
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -450,15 +448,21 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
             chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
 
             try:
-                # Send follow-up chat request
+                # Send follow-up chat request (if necessary)
                 async with session.post(
                     chat_url,
                     headers=headers_chat_combined,
                     data=data_chat,
-                    proxy=proxy
+                    proxy=proxy,
+                    timeout=ClientTimeout(total=600),
                 ) as response_chat:
                     response_chat.raise_for_status()
-                    pass
+                    # If there's additional streaming data from the chat URL, handle it here
+                    async for data in response_chat.content.iter_chunked(1024):
+                        decoded_data = data.decode('utf-8', errors='ignore')
+                        cleaned_data = cls.clean_response(decoded_data)
+                        if cleaned_data.strip():
+                            yield cleaned_data.strip()
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -582,15 +586,19 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             async def generate():
                 try:
                     assistant_content = ""
+                    prompt_tokens = 0
+                    completion_tokens = 0
+
                     async for chunk in async_generator:
                         if isinstance(chunk, ImageResponse):
                             # Handle image responses if necessary
-                            image_markdown = f"![image]({chunk.images})\n"
+                            image_markdown = f"![{chunk.alt}]({chunk.url})\n"
                             assistant_content += image_markdown
                             response_chunk = create_response(image_markdown, request.model, finish_reason=None)
+                            yield f"data: {json.dumps(response_chunk)}\n\n"
                         else:
+                            # Assuming 'chunk' is a string of text
                             assistant_content += chunk
-                            # Yield the chunk as a partial choice
                             response_chunk = {
                                 "id": f"chatcmpl-{uuid.uuid4()}",
                                 "object": "chat.completion.chunk",
@@ -605,9 +613,9 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                                 ],
                                 "usage": None,  # Usage can be updated if you track tokens in real-time
                             }
-                        yield f"data: {json.dumps(response_chunk)}\n\n"
-                    
-                    # After all chunks are sent, send the final message with finish_reason
+                            yield f"data: {json.dumps(response_chunk)}\n\n"
+
+                    # After all chunks are sent, calculate tokens and estimated cost
                     prompt_tokens = sum(len(msg['content'].split()) for msg in request.messages)
                     completion_tokens = len(assistant_content.split())
                     total_tokens = prompt_tokens + completion_tokens
@@ -650,7 +658,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             response_content = ""
             async for chunk in async_generator:
                 if isinstance(chunk, ImageResponse):
-                    response_content += f"![image]({chunk.images})\n"
+                    response_content += f"![{chunk.alt}]({chunk.url})\n"
                 else:
                     response_content += chunk