Support for paint endpoint

cohere/__init__.py

@@ -8,6 +8,7 @@
 EMBED_URL = 'embed'
 CLASSIFY_URL = 'classify'
 EXTRACT_URL = 'extract'
+PAINT_URL = 'paint'
 
 CHECK_API_KEY_URL = 'check-api-key'
 TOKENIZE_URL = 'tokenize'

cohere/client.py

@@ -19,6 +19,7 @@
 from cohere.extract import Extraction, Extractions
 from cohere.generation import Generations
 from cohere.tokenize import Tokens
+from cohere.paint import Images
 
 use_xhr_client = False
 try:
@@ -211,6 +212,10 @@ def detokenize(self, tokens: List[int]) -> Detokenization:
         json_body = {'tokens': tokens}
         return Detokenization(_future=self._executor.submit(self.__request, cohere.DETOKENIZE_URL, json=json_body))
 
+    def paint(self, prompt: str) -> Images:
+        response = self.__request(cohere.PAINT_URL, json={ 'prompt': prompt })
+        return Images(response['image'])
+
     def __print_warning_msg(self, response: Response):
         if 'X-API-Warning' in response.headers:
             print("\033[93mWarning: {}\n\033[0m".format(response.headers['X-API-Warning']), file=sys.stderr)

cohere/paint.py

@@ -0,0 +1,10 @@
+from cohere.response import CohereObject
+
+
+class Images(CohereObject):
+    '''
+    Represents the main response of calling the Image API. An Images is a single base64 encoded string.
+    '''
+
+    def __init__(self, image: str) -> None:
+        self.image = image