response.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350
  1. from __future__ import annotations
  2. import re
  3. import base64
  4. from typing import Union, Dict, List, Optional
  5. from abc import abstractmethod
  6. from urllib.parse import quote_plus, unquote_plus
  7. def quote_url(url: str) -> str:
  8. """
  9. Quote parts of a URL while preserving the domain structure.
  10. Args:
  11. url: The URL to quote
  12. Returns:
  13. str: The properly quoted URL
  14. """
  15. # Only unquote if needed to avoid double-unquoting
  16. if '%' in url:
  17. url = unquote_plus(url)
  18. url_parts = url.split("//", maxsplit=1)
  19. # If there is no "//" in the URL, then it is a relative URL
  20. if len(url_parts) == 1:
  21. return quote_plus(url_parts[0], '/?&=#')
  22. protocol, rest = url_parts
  23. domain_parts = rest.split("/", maxsplit=1)
  24. # If there is no "/" after the domain, then it is a domain URL
  25. if len(domain_parts) == 1:
  26. return f"{protocol}//{domain_parts[0]}"
  27. domain, path = domain_parts
  28. return f"{protocol}//{domain}/{quote_plus(path, '/?&=#')}"
  29. def quote_title(title: str) -> str:
  30. """
  31. Normalize whitespace in a title.
  32. Args:
  33. title: The title to normalize
  34. Returns:
  35. str: The title with normalized whitespace
  36. """
  37. return " ".join(title.split()) if title else ""
  38. def format_link(url: str, title: Optional[str] = None) -> str:
  39. """
  40. Format a URL and title as a markdown link.
  41. Args:
  42. url: The URL to link to
  43. title: The title to display. If None, extracts from URL
  44. Returns:
  45. str: The formatted markdown link
  46. """
  47. if title is None:
  48. try:
  49. title = unquote_plus(url.split("//", maxsplit=1)[1].split("?")[0].replace("www.", ""))
  50. except IndexError:
  51. title = url
  52. return f"[{quote_title(title)}]({quote_url(url)})"
  53. def format_image(image: str, alt: str, preview: Optional[str] = None) -> str:
  54. """
  55. Formats the given image as a markdown string.
  56. Args:
  57. image: The image to format.
  58. alt: The alt text for the image.
  59. preview: The preview URL format. Defaults to the original image.
  60. Returns:
  61. str: The formatted markdown string.
  62. """
  63. preview_url = preview.replace('{image}', image) if preview else image
  64. return f"[![{quote_title(alt)}]({quote_url(preview_url)})]({quote_url(image)})"
  65. def format_images_markdown(images: Union[str, List[str]], alt: str,
  66. preview: Union[str, List[str]] = None) -> str:
  67. """
  68. Formats the given images as a markdown string.
  69. Args:
  70. images: The image or list of images to format.
  71. alt: The alt text for the images.
  72. preview: The preview URL format or list of preview URLs.
  73. If not provided, original images are used.
  74. Returns:
  75. str: The formatted markdown string.
  76. """
  77. if isinstance(images, list) and len(images) == 1:
  78. images = images[0]
  79. if isinstance(images, str):
  80. result = format_image(images, alt, preview)
  81. else:
  82. result = "\n".join(
  83. format_image(
  84. image,
  85. f"#{idx+1} {alt}",
  86. preview[idx] if isinstance(preview, list) and idx < len(preview) else preview
  87. )
  88. for idx, image in enumerate(images)
  89. )
  90. start_flag = "<!-- generated images start -->\n"
  91. end_flag = "<!-- generated images end -->\n"
  92. return f"\n{start_flag}{result}\n{end_flag}\n"
  93. class ResponseType:
  94. @abstractmethod
  95. def __str__(self) -> str:
  96. """Convert the response to a string representation."""
  97. raise NotImplementedError
  98. class JsonMixin:
  99. def __init__(self, **kwargs) -> None:
  100. """Initialize with keyword arguments as attributes."""
  101. for key, value in kwargs.items():
  102. setattr(self, key, value)
  103. def get_dict(self) -> Dict:
  104. """Return a dictionary of non-private attributes."""
  105. return {
  106. key: value
  107. for key, value in self.__dict__.items()
  108. if not key.startswith("__")
  109. }
  110. def reset(self) -> None:
  111. """Reset all attributes."""
  112. self.__dict__ = {}
  113. class RawResponse(ResponseType, JsonMixin):
  114. pass
  115. class HiddenResponse(ResponseType):
  116. def __str__(self) -> str:
  117. """Hidden responses return an empty string."""
  118. return ""
  119. class FinishReason(JsonMixin, HiddenResponse):
  120. def __init__(self, reason: str) -> None:
  121. """Initialize with a reason."""
  122. self.reason = reason
  123. class ToolCalls(HiddenResponse):
  124. def __init__(self, list: List) -> None:
  125. """Initialize with a list of tool calls."""
  126. self.list = list
  127. def get_list(self) -> List:
  128. """Return the list of tool calls."""
  129. return self.list
  130. class Usage(JsonMixin, HiddenResponse):
  131. pass
  132. class AuthResult(JsonMixin, HiddenResponse):
  133. pass
  134. class TitleGeneration(HiddenResponse):
  135. def __init__(self, title: str) -> None:
  136. """Initialize with a title."""
  137. self.title = title
  138. class DebugResponse(HiddenResponse):
  139. def __init__(self, log: str) -> None:
  140. """Initialize with a log message."""
  141. self.log = log
  142. class Reasoning(ResponseType):
  143. def __init__(
  144. self,
  145. token: Optional[str] = None,
  146. label: Optional[str] = None,
  147. status: Optional[str] = None,
  148. is_thinking: Optional[str] = None
  149. ) -> None:
  150. """Initialize with token, status, and thinking state."""
  151. self.token = token
  152. self.label = label
  153. self.status = status
  154. self.is_thinking = is_thinking
  155. def __str__(self) -> str:
  156. """Return string representation based on available attributes."""
  157. if self.is_thinking is not None:
  158. return self.is_thinking
  159. if self.token is not None:
  160. return self.token
  161. if self.status is not None:
  162. return f"{self.status}\n"
  163. return ""
  164. def __eq__(self, other: Reasoning):
  165. return (self.token == other.token and
  166. self.status == other.status and
  167. self.is_thinking == other.is_thinking)
  168. def get_dict(self) -> Dict:
  169. """Return a dictionary representation of the reasoning."""
  170. if self.label is not None:
  171. return {"label": self.label, "status": self.status}
  172. if self.is_thinking is None:
  173. if self.status is None:
  174. return {"token": self.token}
  175. return {"token": self.token, "status": self.status}
  176. return {"token": self.token, "status": self.status, "is_thinking": self.is_thinking}
  177. class Sources(ResponseType):
  178. def __init__(self, sources: List[Dict[str, str]]) -> None:
  179. """Initialize with a list of source dictionaries."""
  180. self.list = []
  181. for source in sources:
  182. self.add_source(source)
  183. def add_source(self, source: Union[Dict[str, str], str]) -> None:
  184. """Add a source to the list, cleaning the URL if necessary."""
  185. source = source if isinstance(source, dict) else {"url": source}
  186. url = source.get("url", source.get("link", None))
  187. if url is not None:
  188. url = re.sub(r"[&?]utm_source=.+", "", url)
  189. source["url"] = url
  190. self.list.append(source)
  191. def __str__(self) -> str:
  192. """Return formatted sources as a string."""
  193. if not self.list:
  194. return ""
  195. return "\n\n\n\n" + ("\n>\n".join([
  196. f"> [{idx}] {format_link(link['url'], link.get('title', None))}"
  197. for idx, link in enumerate(self.list)
  198. ]))
  199. class YouTube(HiddenResponse):
  200. def __init__(self, ids: List[str]) -> None:
  201. """Initialize with a list of YouTube IDs."""
  202. self.ids = ids
  203. def to_string(self) -> str:
  204. """Return YouTube embeds as a string."""
  205. if not self.ids:
  206. return ""
  207. return "\n\n" + ("\n".join([
  208. f'<iframe type="text/html" src="https://www.youtube.com/embed/{id}"></iframe>'
  209. for id in self.ids
  210. ]))
  211. class AudioResponse(ResponseType):
  212. def __init__(self, data: Union[bytes, str]) -> None:
  213. """Initialize with audio data bytes."""
  214. self.data = data
  215. def to_uri(self) -> str:
  216. if isinstance(self.data, str):
  217. return self.data
  218. """Return audio data as a base64-encoded data URI."""
  219. data_base64 = base64.b64encode(self.data).decode()
  220. return f"data:audio/mpeg;base64,{data_base64}"
  221. def __str__(self) -> str:
  222. """Return audio as html element."""
  223. return f'<audio controls src="{self.to_uri()}"></audio>'
  224. class BaseConversation(ResponseType):
  225. def __str__(self) -> str:
  226. """Return an empty string by default."""
  227. return ""
  228. class JsonConversation(BaseConversation, JsonMixin):
  229. pass
  230. class SynthesizeData(HiddenResponse, JsonMixin):
  231. def __init__(self, provider: str, data: Dict) -> None:
  232. """Initialize with provider and data."""
  233. self.provider = provider
  234. self.data = data
  235. class RequestLogin(HiddenResponse):
  236. def __init__(self, label: str, login_url: str) -> None:
  237. """Initialize with label and login URL."""
  238. self.label = label
  239. self.login_url = login_url
  240. def to_string(self) -> str:
  241. """Return formatted login link as a string."""
  242. return format_link(self.login_url, f"[Login to {self.label}]") + "\n\n"
  243. class MediaResponse(ResponseType):
  244. def __init__(
  245. self,
  246. images: Union[str, List[str]],
  247. alt: str,
  248. options: Dict = {}
  249. ) -> None:
  250. """Initialize with images, alt text, and options."""
  251. self.images = images
  252. self.alt = alt
  253. self.options = options
  254. def get(self, key: str) -> any:
  255. """Get an option value by key."""
  256. return self.options.get(key)
  257. def get_list(self) -> List[str]:
  258. """Return images as a list."""
  259. return [self.images] if isinstance(self.images, str) else self.images
  260. class ImageResponse(MediaResponse):
  261. def __str__(self) -> str:
  262. """Return images as markdown."""
  263. return format_images_markdown(self.images, self.alt, self.get("preview"))
  264. class VideoResponse(MediaResponse):
  265. def __str__(self) -> str:
  266. """Return videos as html elements."""
  267. return "\n".join([f'<video controls src="{video}"></video>' for video in self.get_list()])
  268. class ImagePreview(ImageResponse):
  269. def __str__(self) -> str:
  270. """Return an empty string for preview."""
  271. return ""
  272. def to_string(self) -> str:
  273. """Return images as markdown."""
  274. return super().__str__()
  275. class PreviewResponse(HiddenResponse):
  276. def __init__(self, data: str) -> None:
  277. """Initialize with data."""
  278. self.data = data
  279. def to_string(self) -> str:
  280. """Return data as a string."""
  281. return self.data
  282. class Parameters(ResponseType, JsonMixin):
  283. def __str__(self) -> str:
  284. """Return an empty string."""
  285. return ""
  286. class ProviderInfo(JsonMixin, HiddenResponse):
  287. pass