1 files changed, 63 insertions, 31 deletions

diff --git a/g4f/Provider/MetaAI.py b/g4f/Provider/MetaAI.py
index e64a96d5..fb7790f9 100644
--- a/g4f/Provider/MetaAI.py
+++ b/g4f/Provider/MetaAI.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 import uuid
 import random
@@ -8,6 +10,8 @@ from aiohttp import ClientSession, BaseConnector
 
 from ..typing import AsyncResult, Messages, Cookies
 from ..requests import raise_for_status, DEFAULT_HEADERS
+from ..image import ImageResponse, ImagePreview
+from ..errors import ResponseError
 from .base_provider import AsyncGeneratorProvider
 from .helper import format_prompt, get_connector
 
@@ -22,6 +26,7 @@ class AbraGeoBlockedError(Exception):
     pass
 
 class MetaAI(AsyncGeneratorProvider):
+    label = "Meta AI"
     url = "https://www.meta.ai"
     working = True
 
@@ -38,11 +43,10 @@ class MetaAI(AsyncGeneratorProvider):
         proxy: str = None,
         **kwargs
     ) -> AsyncResult:
-        #cookies = get_cookies(".meta.ai", False, True)
         async for chunk in cls(proxy).prompt(format_prompt(messages)):
             yield chunk
 
-    async def get_access_token(self, birthday: str = "1999-01-01") -> str:
+    async def update_access_token(self, birthday: str = "1999-01-01"):
         url = "https://www.meta.ai/api/graphql/"
 
         payload = {
@@ -66,25 +70,37 @@ class MetaAI(AsyncGeneratorProvider):
         async with self.session.post(url, headers=headers, cookies=self.cookies, data=payload) as response:
             await raise_for_status(response, "Fetch access_token failed")
             auth_json = await response.json(content_type=None)
-            access_token = auth_json["data"]["xab_abra_accept_terms_of_service"]["new_temp_user_auth"]["access_token"]
-            return access_token
+            self.access_token = auth_json["data"]["xab_abra_accept_terms_of_service"]["new_temp_user_auth"]["access_token"]
 
     async def prompt(self, message: str, cookies: Cookies = None) -> AsyncResult:
+        if self.cookies is None:
+            await self.update_cookies(cookies)
         if cookies is not None:
-            self.cookies = cookies
             self.access_token = None
-        if self.cookies is None:
-            self.cookies = await self.get_cookies()
-        if self.access_token is None:
-            self.access_token  = await self.get_access_token()
+        if self.access_token is None and cookies is None:
+            await self.update_access_token()
 
-        url = "https://graph.meta.ai/graphql?locale=user"
-        #url = "https://www.meta.ai/api/graphql/"
+        if self.access_token is None:
+            url = "https://www.meta.ai/api/graphql/"
+            payload = {"lsd": self.lsd, 'fb_dtsg': self.dtsg}
+            headers = {'x-fb-lsd': self.lsd}
+        else:
+            url = "https://graph.meta.ai/graphql?locale=user"
+            payload = {"access_token": self.access_token}
+            headers = {}
+        headers = {
+            'content-type': 'application/x-www-form-urlencoded',
+            'cookie': "; ".join([f"{k}={v}" for k, v in cookies.items()]),
+            'origin': 'https://www.meta.ai',
+            'referer': 'https://www.meta.ai/',
+            'x-asbd-id': '129477',
+            'x-fb-friendly-name': 'useAbraSendMessageMutation',
+            **headers
+        }
         payload = {
-            "access_token": self.access_token,
-            #"lsd": cookies["lsd"],
-            "fb_api_caller_class": "RelayModern",
-            "fb_api_req_friendly_name": "useAbraSendMessageMutation",
+            **payload,
+            'fb_api_caller_class': 'RelayModern',
+            'fb_api_req_friendly_name': 'useAbraSendMessageMutation',
             "variables": json.dumps({
                 "message": {"sensitive_string_value": message},
                 "externalConversationId": str(uuid.uuid4()),
@@ -98,19 +114,16 @@ class MetaAI(AsyncGeneratorProvider):
                 "__relay_internal__pv__AbraDebugDevOnlyrelayprovider": False,
                 "__relay_internal__pv__WebPixelRatiorelayprovider": 1,
             }),
-            "server_timestamps": "true",
-            "doc_id": "7783822248314888",
-        }
-        headers = {
-            "x-asbd-id": "129477",
-            "x-fb-friendly-name": "useAbraSendMessageMutation",
-            #"x-fb-lsd": cookies["lsd"],
+            'server_timestamps': 'true',
+            'doc_id': '7783822248314888'
         }
-        async with self.session.post(url, headers=headers, cookies=self.cookies, data=payload) as response:
+        async with self.session.post(url, headers=headers, data=payload) as response:
             await raise_for_status(response, "Fetch response failed")
             last_snippet_len = 0
             fetch_id = None
             async for line in response.content:
+                if b"<h1>Something Went Wrong</h1>" in line:
+                    raise ResponseError("Response: Something Went Wrong")
                 try:
                     json_line = json.loads(line)
                 except json.JSONDecodeError:
@@ -119,7 +132,14 @@ class MetaAI(AsyncGeneratorProvider):
                 streaming_state = bot_response_message.get("streaming_state")
                 fetch_id = bot_response_message.get("fetch_id") or fetch_id
                 if streaming_state in ("STREAMING", "OVERALL_DONE"):
-                    #imagine_card = bot_response_message["imagine_card"]
+                    imagine_card = bot_response_message.get("imagine_card")
+                    if imagine_card is not None:
+                        imagine_session = imagine_card.get("session")
+                        if imagine_session is not None:
+                            imagine_medias = imagine_session.get("media_sets", {}).pop().get("imagine_media")
+                            if imagine_medias is not None:
+                                image_class = ImageResponse if streaming_state == "OVERALL_DONE" else ImagePreview
+                                yield image_class([media["uri"] for media in imagine_medias], imagine_medias[0]["prompt"])
                     snippet =  bot_response_message["snippet"]
                     new_snippet_len = len(snippet)
                     if new_snippet_len > last_snippet_len:
@@ -135,7 +155,7 @@ class MetaAI(AsyncGeneratorProvider):
                 if sources is not None:
                     yield sources 
 
-    async def get_cookies(self, cookies: Cookies = None) -> Cookies:
+    async def update_cookies(self, cookies: Cookies = None):
         async with self.session.get("https://www.meta.ai/", cookies=cookies) as response:
             await raise_for_status(response, "Fetch home failed")
             text = await response.text()
@@ -148,12 +168,20 @@ class MetaAI(AsyncGeneratorProvider):
                     "datr": self.extract_value(text, "datr"),
                 }
             self.lsd = self.extract_value(text, start_str='"LSD",[],{"token":"', end_str='"}')
-            return cookies
+            self.dtsg = self.extract_value(text, start_str='"DTSGInitialData",[],{"token":"', end_str='"}')
+            self.cookies = cookies
 
     async def fetch_sources(self, fetch_id: str) -> Sources:
-        url = "https://graph.meta.ai/graphql?locale=user"
+        if self.access_token is None:
+            url = "https://www.meta.ai/api/graphql/"
+            payload = {"lsd": self.lsd, 'fb_dtsg': self.dtsg}
+            headers = {'x-fb-lsd': self.lsd}
+        else:
+            url = "https://graph.meta.ai/graphql?locale=user"
+            payload = {"access_token": self.access_token}
+            headers = {}
         payload = {
-            "access_token": self.access_token,
+            **payload,
             "fb_api_caller_class": "RelayModern",
             "fb_api_req_friendly_name": "AbraSearchPluginDialogQuery",
             "variables": json.dumps({"abraMessageFetchID": fetch_id}),
@@ -163,18 +191,22 @@ class MetaAI(AsyncGeneratorProvider):
         headers = {
             "authority": "graph.meta.ai",
             "x-fb-friendly-name": "AbraSearchPluginDialogQuery",
+            **headers
         }
         async with self.session.post(url, headers=headers, cookies=self.cookies, data=payload) as response:
             await raise_for_status(response)
-            response_json = await response.json()
+            text = await response.text()
+            if "<h1>Something Went Wrong</h1>" in text:
+                raise ResponseError("Response: Something Went Wrong")
             try:
+                response_json = json.loads(text)
                 message = response_json["data"]["message"]
                 if message is not None:
                     searchResults = message["searchResults"]
                     if searchResults is not None:
                         return Sources(searchResults["references"])
-            except (KeyError, TypeError):
-                raise RuntimeError(f"Response: {response_json}")
+            except (KeyError, TypeError, json.JSONDecodeError):
+                raise RuntimeError(f"Response: {text}")
 
     @staticmethod
     def extract_value(text: str, key: str = None, start_str = None, end_str = '",') -> str: