Merge pull request #6 from Strvm/feat/auth_and_image_generation

feat: FB auth + image generation

README.md

@@ -10,6 +10,7 @@ Meta AI is running Llama 3 LLM.
 
 ## Features
 - **Prompt AI**: Send a message to the AI and get a response from Llama 3.
+- **Image Generation**: Generate images using the AI. (Only for FB authenticated users)
 - **Get Up To Date Information**: Get the latest information from the AI thanks to its connection to the internet.
 - **Get Sources**: Get the sources of the information provided by the AI.
 - **Streaming**: Stream the AI's response in real-time or get the final response.
@@ -124,3 +125,49 @@ for r in response:
 ...
 {'message': "The Golden State Warriors' last game was against the Sacramento Kings on April 16, 2024, at the Golden 1 Center in Sacramento, California. The Kings won the game with a score of 118-94, with the Warriors scoring 22 points in the first quarter, 28 in the second, 26 in the third and 18 in the fourth quarter ¹.\n", 'sources': [{'link': 'https://sportradar.com/', 'title': 'Game Info of NBA from sportradar.com'}]}
 ```
+
+**Generate Image**:
+
+By default image generation is only available for FB authenticated users. If you go on https://www.meta.ai/ , and ask the AI to generate an image, you will be prompted to authenticate with Facebook.
+So if you want to generate images using this library, you need to authenticate using your FB credentials.
+
+**Note**: There seems to be higher rate limits for authenticated users. So only authenticate to generate images.
+
+```python
+from meta_ai_api import MetaAI
+ai = MetaAI(fb_email="your_fb_email", fb_password="your_fb_password")
+resp = ai.prompt(message="Generate an image of a tech CEO")
+print(resp)
+```
+
+```json
+{
+   "message":"\n",
+   "sources":[
+
+   ],
+   "media":[
+      {
+         "url":"https://scontent-lax3-1.xx.fbcdn.net/o1/v/t0/f1/m247/4282108942387920518_1946149595_21-04-2024-14-17-48.jpeg?_nc_ht=scontent-lax3-1.xx.fbcdn.net&_nc_cat=103&ccb=9-4&oh=00_AfCnbCX7nl_J5kF6mahnams4d99Rs5WZA780HGS_scfc6A&oe=662771EE&_nc_sid=5b3566",
+         "type":"IMAGE",
+         "prompt":"a tech CEO"
+      },
+      {
+         "url":"https://scontent-lax3-1.xx.fbcdn.net/o1/v/t0/f1/m247/3356467762794691754_1025991308_21-04-2024-14-17-48.jpeg?_nc_ht=scontent-lax3-1.xx.fbcdn.net&_nc_cat=108&ccb=9-4&oh=00_AfBLmSbTSqshNAL82KIFk8hGXyL8iK_CZLGcMmmddPrxuA&oe=66276EDD&_nc_sid=5b3566",
+         "type":"IMAGE",
+         "prompt":"a tech CEO"
+      },
+      {
+         "url":"https://scontent-lax3-1.xx.fbcdn.net/o1/v/t0/f1/m247/127487551948523111_2181921077_21-04-2024-14-17-48.jpeg?_nc_ht=scontent-lax3-1.xx.fbcdn.net&_nc_cat=104&ccb=9-4&oh=00_AfAejXKeKPA4vyKXoc6UR0rEirvZwi41P3KiCSQmHRHsEw&oe=66276E45&_nc_sid=5b3566",
+         "type":"IMAGE",
+         "prompt":"a tech CEO"
+      },
+      {
+         "url":"https://scontent-lax3-1.xx.fbcdn.net/o1/v/t0/f1/m247/3497663176351797954_3954783377_21-04-2024-14-17-47.jpeg?_nc_ht=scontent-lax3-1.xx.fbcdn.net&_nc_cat=110&ccb=9-4&oh=00_AfBp3bAfcuofqtI-z9D4bHw-GuGgCNPH_xhMM0PG_95S9Q&oe=66277AE9&_nc_sid=5b3566",
+         "type":"IMAGE",
+         "prompt":"a tech CEO"
+      }
+   ]
+}
+```
+![Tech CEO](https://scontent-lax3-1.xx.fbcdn.net/o1/v/t0/f1/m247/3497663176351797954_3954783377_21-04-2024-14-17-47.jpeg?_nc_ht=scontent-lax3-1.xx.fbcdn.net&_nc_cat=110&ccb=9-4&oh=00_AfBp3bAfcuofqtI-z9D4bHw-GuGgCNPH_xhMM0PG_95S9Q&oe=66277AE9&_nc_sid=5b3566)

pyproject.toml

@@ -13,6 +13,7 @@ python = "^3.7"
 requests = "2.31.0"
 requests-html = "0.10.0"
 lxml_html_cleaner = "0.1.1"
+bs4 = "0.0.2"
 
 [build]
 script = "build.py"

requirements.txt

@@ -1,3 +1,4 @@
 requests==2.31.0
 requests-html==0.10.0
-lxml_html_clean==0.1.1
+lxml_html_clean==0.1.1
+bs4==0.0.2

src/meta_ai_api/__init__.py

@@ -1,2 +1,2 @@
-__version__ = "1.1.2"
+__version__ = "1.1.3"
 from .main import MetaAI  # noqa

src/meta_ai_api/exceptions.py

@@ -0,0 +1,2 @@
+class FacebookInvalidCredentialsException(Exception):
+    pass

src/meta_ai_api/main.py

@@ -14,6 +14,7 @@
     format_response,
 )
 
+from meta_ai_api.utils import get_fb_session
 
 MAX_RETRIES = 3
 
@@ -24,7 +25,7 @@ class MetaAI:
     and receiving messages from the Meta AI Chat API.
     """
 
-    def __init__(self):
+    def __init__(self, fb_email: str = None, fb_password: str = None):
         self.session = requests.Session()
         self.session.headers.update(
             {
@@ -33,7 +34,10 @@ def __init__(self):
             }
         )
         self.access_token = None
-        self.cookies = None
+        self.fb_email = fb_email
+        self.fb_password = fb_password
+        self.is_authed = fb_password is not None and fb_email is not None
+        self.cookies = self.get_cookies()
 
     def get_access_token(self) -> str:
         """
@@ -42,9 +46,7 @@ def get_access_token(self) -> str:
         Returns:
             str: A valid access token.
         """
-        self.cookies = self.get_cookies()
         url = "https://www.meta.ai/api/graphql/"
-
         payload = {
             "lsd": self.cookies["lsd"],
             "fb_api_caller_class": "RelayModern",
@@ -89,15 +91,20 @@ def prompt(
         Raises:
             Exception: If unable to obtain a valid response after several attempts.
         """
-        if not self.access_token:
+        if not self.access_token and not self.is_authed:
             self.access_token = self.get_access_token()
+            auth_payload = {"access_token": self.access_token}
+            url = "https://graph.meta.ai/graphql?locale=user"
+
+        else:
+            auth_payload = {"fb_dtsg": self.cookies["fb_dtsg"]}
+            url = "https://www.meta.ai/api/graphql/"
 
         # Need to sleep for a bit, for some reason the API doesn't like it when we send request too quickly
         # (maybe Meta needs to register Cookies on their side?)
         time.sleep(1)
-        url = "https://graph.meta.ai/graphql?locale=user"
         payload = {
-            "access_token": self.access_token,
+            **auth_payload,
             "fb_api_caller_class": "RelayModern",
             "fb_api_req_friendly_name": "useAbraSendMessageMutation",
             "variables": json.dumps(
@@ -123,6 +130,10 @@ def prompt(
             "content-type": "application/x-www-form-urlencoded",
             "x-fb-friendly-name": "useAbraSendMessageMutation",
         }
+        if self.is_authed:
+            headers["cookie"] = f'abra_sess={self.cookies["abra_sess"]}'
+            # Recreate the session to avoid cookie leakage when user is authenticated
+            self.session = requests.Session()
 
         response = self.session.post(url, headers=headers, data=payload, stream=stream)
         if not stream:
@@ -220,33 +231,78 @@ def extract_data(self, json_line: dict):
         response = format_response(response=json_line)
         fetch_id = bot_response_message.get("fetch_id")
         sources = self.fetch_sources(fetch_id) if fetch_id else []
-        return {"message": response, "sources": sources}
+        medias = self.extract_media(bot_response_message)
+        return {"message": response, "sources": sources, "media": medias}
 
-    @staticmethod
-    def get_cookies() -> dict:
+    def extract_media(self, json_line: dict) -> List[Dict]:
+        """
+        Extract media from a parsed JSON line.
+
+        Args:
+            json_line (dict): Parsed JSON line.
+
+        Returns:
+            list: A list of dictionaries containing the extracted media.
+        """
+        medias = []
+        imagine_card = json_line.get("imagine_card", {})
+        session = imagine_card.get("session", {}) if imagine_card else {}
+        media_sets = (
+            (json_line.get("imagine_card", {}).get("session", {}).get("media_sets", []))
+            if imagine_card and session
+            else []
+        )
+        for media_set in media_sets:
+            imagine_media = media_set.get("imagine_media", [])
+            for media in imagine_media:
+                medias.append(
+                    {
+                        "url": media.get("uri"),
+                        "type": media.get("media_type"),
+                        "prompt": media.get("prompt"),
+                    }
+                )
+        return medias
+
+    def get_cookies(self) -> dict:
         """
         Extracts necessary cookies from the Meta AI main page.
 
         Returns:
             dict: A dictionary containing essential cookies.
         """
         session = HTMLSession()
-        response = session.get("https://www.meta.ai/")
-        return {
+        headers = {}
+        if self.fb_email is not None and self.fb_password is not None:
+            fb_session = get_fb_session(self.fb_email, self.fb_password)
+            headers = {"cookie": f"abra_sess={fb_session['abra_sess']}"}
+        response = session.get(
+            "https://www.meta.ai/",
+            headers=headers,
+        )
+        cookies = {
             "_js_datr": extract_value(
                 response.text, start_str='_js_datr":{"value":"', end_str='",'
             ),
-            "abra_csrf": extract_value(
-                response.text, start_str='abra_csrf":{"value":"', end_str='",'
-            ),
             "datr": extract_value(
                 response.text, start_str='datr":{"value":"', end_str='",'
             ),
             "lsd": extract_value(
                 response.text, start_str='"LSD",[],{"token":"', end_str='"}'
             ),
+            "fb_dtsg": extract_value(
+                response.text, start_str='DTSGInitData",[],{"token":"', end_str='"'
+            ),
         }
 
+        if len(headers) > 0:
+            cookies["abra_sess"] = fb_session["abra_sess"]
+        else:
+            cookies["abra_csrf"] = extract_value(
+                response.text, start_str='abra_csrf":{"value":"', end_str='",'
+            )
+        return cookies
+
     def fetch_sources(self, fetch_id: str) -> List[Dict]:
         """
         Fetches sources from the Meta AI API based on the given query.
@@ -274,13 +330,18 @@ def fetch_sources(self, fetch_id: str) -> List[Dict]:
             "authority": "graph.meta.ai",
             "accept-language": "en-US,en;q=0.9,fr-FR;q=0.8,fr;q=0.7",
             "content-type": "application/x-www-form-urlencoded",
-            "cookie": f'dpr=2; abra_csrf={self.cookies["abra_csrf"]}; datr={self.cookies["datr"]}; ps_n=1; ps_l=1',
+            "cookie": f'dpr=2; abra_csrf={self.cookies.get("abra_csrf")}; datr={self.cookies.get("datr")}; ps_n=1; ps_l=1',
             "x-fb-friendly-name": "AbraSearchPluginDialogQuery",
         }
 
         response = self.session.post(url, headers=headers, data=payload)
         response_json = response.json()
-        search_results = response_json["data"]["message"]["searchResults"]
+        message = response_json.get("data", {}).get("message", {})
+        search_results = (
+            (response_json.get("data", {}).get("message", {}).get("searchResults"))
+            if message
+            else None
+        )
         if search_results is None:
             return []
 
@@ -291,7 +352,4 @@ def fetch_sources(self, fetch_id: str) -> List[Dict]:
 if __name__ == "__main__":
     meta = MetaAI()
     resp = meta.prompt("What was the Warriors score last game?", stream=False)
-    # for r in resp:
-    #     print(r)
-
     print(resp)