Coverage for src/appl/core/types/content.py: 84%

1import base64 

2from io import BytesIO 

3from os import PathLike 

4from typing import Optional 

5 

6import PIL.Image 

7from PIL.ImageFile import ImageFile 

8from pydantic import BaseModel, Field 

9 

10from .basic import * 

11from .futures import String, is_string 

12 

13 

14class Image(BaseModel): 

15 """Represent an image in the message.""" 

16 

17 url: str = Field( 

18 ..., description="Either a URL of the image or the base64 encoded image data." 

19 ) 

20 detail: Optional[str] = Field( 

21 None, description="Specifies the detail level of the image." 

22 ) 

23 

24 def __init__(self, url: str, detail: Optional[str] = None) -> None: 

25 """Initialize the image with the URL and detail level. 

26 

27 See [the guide](https://platform.openai.com/docs/guides/vision/low-or-high-fidelity-image-understanding) 

28 for more information about the detail level. 

29 """ 

30 super().__init__(url=url, detail=detail) 

31 

32 @classmethod 

33 def from_image(cls, image: ImageFile, detail: Optional[str] = None) -> "Image": 

34 """Construct an image prompt from a PIL ImageFile.""" 

35 buffered = BytesIO() 

36 # Save the image to the buffer in PNG format 

37 image.save(buffered, format="PNG") 

38 # Get the byte data from the buffer 

39 img_byte = buffered.getvalue() 

40 img_base64 = base64.b64encode(img_byte).decode("utf-8") 

41 return cls(url=f"data:image/png;base64,{img_base64}", detail=detail) 

42 

43 @classmethod 

44 def from_file(cls, file: PathLike, detail: Optional[str] = None) -> "Image": 

45 """Construct an image prompt from an image file.""" 

46 image = PIL.Image.open(file) 

47 return cls.from_image(image, detail) 

48 

49 def __repr__(self) -> str: 

50 return f"Image(url={self.url})" 

51 

52 def __str__(self) -> str: 

53 return f"[Image]{self.url}" 

54 

55 

56StrOrImg = Union[String, Image] 

57"""A type that can be either a string or an image.""" 

58 

59 

60class ContentList(BaseModel): 

61 """Represent a list of contents containing text and images.""" 

62 

63 contents: List[StrOrImg] = Field(..., description="The content of the message") 

64 

65 def __iadd__( 

66 self, other: Union[StrOrImg, List[StrOrImg], "ContentList"] 

67 ) -> "ContentList": 

68 if isinstance(other, ContentList): 

69 self.extend(other.contents) 

70 elif isinstance(other, list): 

71 self.extend(other) 

72 else: 

73 self.append(other) 

74 return self 

75 

76 def append(self, content: StrOrImg) -> None: 

77 """Append a content to the list. 

78 

79 If the last content is a string, it will be concatenated with the new content. 

80 """ 

81 if is_string(content) and len(self.contents) and is_string(self.contents[-1]): 

82 self.contents[-1] += content # type: ignore 

83 else: 

84 self.contents.append(content) 

85 

86 def extend(self, contents: list[StrOrImg]) -> None: 

87 """Extend the list with multiple contents.""" 

88 for content in contents: 

89 self.append(content) 

90 

91 def get_contents(self) -> List[Dict[str, Any]]: 

92 """Return the contents as a list of dictionaries.""" 

93 

94 def get_dict(content): 

95 if isinstance(content, Image): 

96 image_args = {"url": content.url} 

97 if content.detail: 

98 image_args["detail"] = content.detail 

99 return {"type": "image_url", "image_url": image_args} 

100 return {"type": "text", "text": str(content)} 

101 

102 return [get_dict(c) for c in self.contents] 

103 

104 def __str__(self) -> str: 

105 return "\n".join([str(c) for c in self.contents])