Update main.py

main.py

@@ -98,8 +98,8 @@ class ModelNotWorkingException(Exception):
 
 # Define ImageResponseData class
 class ImageResponseData:
-    def __init__(self, url: str, alt: str):
-        self.url = url
+    def __init__(self, images: str, alt: str):
+        self.images = images
         self.alt = alt
 
 # Mock implementation of to_data_uri
@@ -140,7 +140,6 @@ class Blackbox:
         'ReactAgent',
         'XcodeAgent',
         'AngularJSAgent',
-        'Niansuh',
     ]
 
     # Filter models based on AVAILABLE_MODELS
@@ -266,9 +265,7 @@ class Blackbox:
         model: str,
         messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
-        image: Any = None,
-        image_name: Optional[str] = None,
-        webSearchMode: bool = False,
+        websearch: bool = False,
         **kwargs
     ) -> AsyncGenerator[Union[str, ImageResponseData], None]:
         """
@@ -278,9 +275,7 @@ class Blackbox:
             model (str): Model to use for generating responses.
             messages (List[Dict[str, str]]): Message history.
             proxy (Optional[str]): Proxy URL, if needed.
-            image (Any): Image data, if applicable.
-            image_name (Optional[str]): Image name, if applicable.
-            webSearchMode (bool): Enables or disables web search mode.
+            websearch (bool): Enables or disables web search mode.
             **kwargs: Additional keyword arguments.
 
         Yields:
@@ -350,7 +345,7 @@ class Blackbox:
             "trendingAgentMode": trending_agent_mode,
             "isMicMode": False,
             "userSystemPrompt": None,
-            "maxTokens": 99999999,
+            "maxTokens": 1024,
             "playgroundTopP": 0.9,
             "playgroundTemperature": 0.5,
             "isChromeExt": False,
@@ -360,96 +355,91 @@ class Blackbox:
             "clickedForceWebSearch": False,
             "visitFromDelta": False,
             "mobileClient": False,
-            "webSearchMode": webSearchMode,
-            "userSelectedModel": None,
+            "webSearchMode": websearch,
+            "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
-        if model in cls.agentMode:
-            payload_api_chat["agentMode"] = cls.agentMode[model]
-        elif model in cls.trendingAgentMode:
-            payload_api_chat["trendingAgentMode"] = cls.trendingAgentMode[model]
-        elif model in cls.userSelectedModel:
-            payload_api_chat["userSelectedModel"] = cls.userSelectedModel[model]
-
-        logger.info(f"Sending request to {cls.api_endpoint} with data (excluding messages).")
+        headers_chat = {
+            'Accept': 'text/x-component',
+            'Content-Type': 'text/plain;charset=UTF-8',
+            'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
+            'next-action': next_action,
+            'next-router-state-tree': next_router_state_tree,
+            'next-url': '/'
+        }
+        headers_chat_combined = {**common_headers, **headers_chat}
 
-        timeout = ClientTimeout(total=60)  # Set an appropriate timeout
-        retry_attempts = 10  # Set the number of retry attempts
+        data_chat = '[]'
 
         async with ClientSession(headers=common_headers) as session:
-            for attempt in range(retry_attempts):
-                try:
-                    async with session.post(cls.api_endpoint, json=payload_api_chat, proxy=proxy) as response:
-                        response.raise_for_status()
-                        logger.info(f"Received response with status {response.status}")
-                        response_text = await response.text()
-                        cleaned_response = cls.clean_response(response_text)
-
-                        if model in cls.image_models:
-                            url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', cleaned_response)
-                            if url_match:
-                                image_url = url_match.group(0)
-                                logger.info(f"Image URL found: {image_url}")
-                                yield ImageResponseData(url=image_url, alt=messages[-1]['content'])
+            try:
+                async with session.post(
+                    cls.api_endpoint,
+                    headers=headers_api_chat_combined,
+                    json=payload_api_chat,
+                    proxy=proxy
+                ) 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 = ImageResponseData(images=image_url, alt="Generated Image")
+                            yield image_response
+                        else:
+                            yield cleaned_response
+                    else:
+                        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()
+                                try:
+                                    sources = json.loads(source_part)
+                                    source_formatted = "**Sources:**\n"
+                                    for item in sources[:5]:
+                                        title = item.get('title', 'No Title')
+                                        link = item.get('link', '#')
+                                        position = item.get('position', '')
+                                        source_formatted += f"- [{title}]({link})\n"
+                                    final_response = f"{answer_part}\n\n{source_formatted}"
+                                except json.JSONDecodeError:
+                                    final_response = f"{answer_part}\n\nSource information is unavailable."
                             else:
-                                logger.error("Image URL not found in the response.")
-                                raise Exception("Image URL not found in the response")
+                                final_response = cleaned_response
                         else:
-                            full_response = ""
-                            search_results_json = ""
-                            # Handle streaming-like responses if applicable
-                            # Assuming the response contains '$~~~$' for search results
-                            if webSearchMode:
-                                match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
-                                if match:
-                                    search_results_json = match.group(1)
-                                    answer_part = cleaned_response[:match.start()].strip()
-                                    try:
-                                        search_results = json.loads(search_results_json)
-                                        formatted_results = "\n\n**Sources:**\n"
-                                        for i, result in enumerate(search_results[:5], 1):
-                                            formatted_results += f"{i}. [{result['title']}]({result['link']})\n"
-                                        final_response = f"{answer_part}\n\n{formatted_results}"
-                                        yield final_response
-                                    except json.JSONDecodeError as je:
-                                        logger.error("Failed to parse search results JSON.")
-                                        yield f"{cleaned_response}\n\n**Sources:** Information unavailable."
-                                else:
-                                    yield cleaned_response
+                            if '$~~~$' in cleaned_response:
+                                final_response = cleaned_response.split('$~~~$')[0].strip()
                             else:
-                                yield cleaned_response
-                    break  # Exit the retry loop if successful
-                except ClientError as ce:
-                    logger.error(f"Client error occurred: {ce}. Retrying attempt {attempt + 1}/{retry_attempts}")
-                    if attempt == retry_attempts - 1:
-                        raise HTTPException(status_code=502, detail="Error communicating with the external API.")
-                except asyncio.TimeoutError:
-                    logger.error(f"Request timed out. Retrying attempt {attempt + 1}/{retry_attempts}")
-                    if attempt == retry_attempts - 1:
-                        raise HTTPException(status_code=504, detail="External API request timed out.")
-                except Exception as e:
-                    logger.error(f"Unexpected error: {e}. Retrying attempt {attempt + 1}/{retry_attempts}")
-                    if attempt == retry_attempts - 1:
-                        raise HTTPException(status_code=500, detail=str(e))
+                                final_response = cleaned_response
 
-            # Additional request to /chat/{chat_id} endpoint if necessary
-            chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
+                        yield final_response
+            except ClientResponseError as e:
+                error_text = f"Error {e.status}: {e.message}"
+                try:
+                    error_response = await e.response.text()
+                    cleaned_error = cls.clean_response(error_response)
+                    error_text += f" - {cleaned_error}"
+                except Exception:
+                    pass
+                yield error_text
+            except Exception as e:
+                yield f"Unexpected error during /api/chat request: {str(e)}"
 
-            headers_chat = {
-                'Accept': 'text/x-component',
-                'Content-Type': 'text/plain;charset=UTF-8',
-                'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
-                'next-action': cls.generate_next_action(),
-                'next-router-state-tree': cls.generate_next_router_state_tree(),
-                'next-url': '/'
-            }
-            headers_chat_combined = {**common_headers, **headers_chat}
-            data_chat = '[]'
+            chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
 
             try:
-                async with session.post(chat_url, headers=headers_chat_combined, data=data_chat, proxy=proxy) as response_chat:
+                async with session.post(
+                    chat_url,
+                    headers=headers_chat_combined,
+                    data=data_chat,
+                    proxy=proxy
+                ) as response_chat:
                     response_chat.raise_for_status()
-                    logger.info(f"Successfully posted to {chat_url}")
+                    # No action needed based on the original code
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -458,12 +448,9 @@ class Blackbox:
                     error_text += f" - {cleaned_error}"
                 except Exception:
                     pass
-                logger.error(error_text)
                 yield error_text
             except Exception as e:
-                error_text = f"Unexpected error during /chat/{chat_id} request: {str(e)}"
-                logger.error(error_text)
-                yield error_text
+                yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
 
 # FastAPI app setup
 app = FastAPI()
@@ -515,7 +502,7 @@ class ChatRequest(BaseModel):
     frequency_penalty: Optional[float] = 0.0
     logit_bias: Optional[Dict[str, float]] = None
     user: Optional[str] = None
-    webSearchMode: Optional[bool] = False  # Custom parameter
+    websearch: Optional[bool] = False  # Custom parameter
 
 class TokenizerRequest(BaseModel):
     text: str
@@ -566,7 +553,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         async_generator = Blackbox.create_async_generator(
             model=request.model,
             messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],  # Actual message content used here
-            webSearchMode=request.webSearchMode
+            websearch=request.websearch
         )
 
         if request.stream:
@@ -576,27 +563,46 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                     async for chunk in async_generator:
                         if isinstance(chunk, ImageResponseData):
                             # Handle image responses if necessary
-                            image_markdown = f"![image]({chunk.url})\n"
+                            image_markdown = f"![image]({chunk.images})\n"
                             assistant_content += image_markdown
-                            response_chunk = create_response(image_markdown, request.model, finish_reason=None)
+                            response_chunk = {
+                                "id": f"chatcmpl-{uuid.uuid4()}",
+                                "object": "chat.completion",
+                                "created": int(datetime.now().timestamp()),
+                                "model": request.model,
+                                "choices": [
+                                    {
+                                        "index": 0,
+                                        "message": {
+                                            "role": "assistant",
+                                            "content": image_markdown
+                                        },
+                                        "finish_reason": None
+                                    }
+                                ],
+                                "usage": None,  # Usage can be updated if you track tokens in real-time
+                            }
                         else:
                             assistant_content += chunk
                             # Yield the chunk as a partial choice
                             response_chunk = {
                                 "id": f"chatcmpl-{uuid.uuid4()}",
-                                "object": "chat.completion.chunk",
+                                "object": "chat.completion",
                                 "created": int(datetime.now().timestamp()),
                                 "model": request.model,
                                 "choices": [
                                     {
                                         "index": 0,
-                                        "delta": {"content": chunk, "role": "assistant"},
-                                        "finish_reason": None,
+                                        "message": {
+                                            "role": "assistant",
+                                            "content": chunk
+                                        },
+                                        "finish_reason": None
                                     }
                                 ],
                                 "usage": None,  # Usage can be updated if you track tokens in real-time
                             }
-                        yield f"data: {json.dumps(response_chunk)}\n\n"
+                        yield f"{json.dumps(response_chunk)}\n\n"
                     
                     # After all chunks are sent, send the final message with finish_reason
                     prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
@@ -626,22 +632,22 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                             "estimated_cost": estimated_cost
                         },
                     }
-                    yield f"data: {json.dumps(final_response)}\n\n"
-                    yield "data: [DONE]\n\n"
+                    yield f"{json.dumps(final_response)}\n\n"
+                    yield "DONE\n\n"
                 except HTTPException as he:
                     error_response = {"error": he.detail}
-                    yield f"data: {json.dumps(error_response)}\n\n"
+                    yield f"{json.dumps(error_response)}\n\n"
                 except Exception as e:
                     logger.exception(f"Error during streaming response generation from IP: {client_ip}.")
                     error_response = {"error": str(e)}
-                    yield f"data: {json.dumps(error_response)}\n\n"
+                    yield f"{json.dumps(error_response)}\n\n"
 
-            return StreamingResponse(generate(), media_type="text/event-stream")
+            return StreamingResponse(generate(), media_type="text/plain")
         else:
             response_content = ""
             async for chunk in async_generator:
                 if isinstance(chunk, ImageResponseData):
-                    response_content += f"![image]({chunk.url})\n"
+                    response_content += f"![image]({chunk.images})\n"
                 else:
                     response_content += chunk