Add image generation functionality to Grokit

- Implement `image()` method to generate and return JPEG binary data
- Add `image_url()` method to get URL of generated image
- Update README with examples for new image generation features
- Modify _stream_response to handle image attachment data
- Enhance error handling for image generation process

Author: brunn3is

README.md

@@ -28,20 +28,36 @@ grok = Grokit(
 )
 ```
 
-## Generate
+## Generate Text
 
 ```python
 response = grok.generate('Who are you?', model_id='grok-2-mini')
 print(response)
 ```
 
-## Stream
+## Stream Text
 
 ```python
 for chunk in grok.stream('Who are you?', model_id='grok-2'):
     print(chunk, end='', flush=True)
 ```
 
+## Generate Image
+
+### JPEG binary
+```python
+image_data = grok.image('An astronaut riding a horse.')
+
+with open('sunset.jpg', 'wb') as f:
+    f.write(image_data)
+```
+
+### image URL
+```python
+image_url = grok.image_url('An astronaut riding a horse.')
+print(image_url)
+```
+
 ## Credentials
 
 To obtain the necessary credentials for using Grokit, follow these steps:

grokit/grokit.py

@@ -94,6 +94,38 @@ def stream(
             model_id,
         )
 
+    def image(self, prompt: str) -> bytes:
+        image_url = self._get_image_url(prompt)
+        image_response = requests.get(image_url)
+        if image_response.status_code == 200:
+            return image_response.content
+        raise ValueError('Failed to download the image')
+
+    def image_url(self, prompt: str) -> str:
+        return self._get_image_url(prompt)
+
+    def _get_image_url(self, prompt: str) -> str:
+        conversation_id = self.create_conversation()
+        if not conversation_id:
+            raise ValueError('Failed to create conversation')
+        image_prompt = 'Generate an image of "{}"'.format(prompt)
+        response = self._stream_response(
+            conversation_id,
+            image_prompt,
+            '',
+            GrokModels.GROK_2_MINI,
+        )
+
+        for chunk in response:
+            try:
+                data = json.loads(chunk)
+                if 'result' in data and 'imageAttachment' in data['result']:
+                    return data['result']['imageAttachment']['imageUrl']
+            except json.JSONDecodeError:
+                continue
+
+        raise ValueError('Failed to generate the image')
+
     def _ensure_conversation_id(self, conversation_id: Optional[str]) -> str:
         if not conversation_id:
             conversation_id = self.create_conversation()
@@ -169,5 +201,8 @@ def _process_response_stream(
         for line in response.iter_lines():
             if line:
                 chunk = json.loads(line)
-                if 'result' in chunk and 'message' in chunk['result']:
-                    yield chunk['result']['message']
+                if 'result' in chunk:
+                    if 'message' in chunk['result']:
+                        yield chunk['result']['message']
+                    elif 'imageAttachment' in chunk['result']:
+                        yield json.dumps(chunk)

grokit/version.py

@@ -1,4 +1,4 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-__version__ = '1.2410.0'
+__version__ = '1.2410.1'