helper.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. from __future__ import annotations
  2. import random
  3. import string
  4. from ..typing import Messages, Cookies, AsyncIterator, Iterator
  5. from .. import debug
  6. def format_prompt(messages: Messages, add_special_tokens: bool = False, do_continue: bool = False) -> str:
  7. """
  8. Format a series of messages into a single string, optionally adding special tokens.
  9. Args:
  10. messages (Messages): A list of message dictionaries, each containing 'role' and 'content'.
  11. add_special_tokens (bool): Whether to add special formatting tokens.
  12. Returns:
  13. str: A formatted string containing all messages.
  14. """
  15. if not add_special_tokens and len(messages) <= 1:
  16. return messages[0]["content"]
  17. formatted = "\n".join([
  18. f'{message["role"].capitalize()}: {message["content"]}'
  19. for message in messages
  20. ])
  21. if do_continue:
  22. return formatted
  23. return f"{formatted}\nAssistant:"
  24. def format_prompt_max_length(messages: Messages, max_lenght: int) -> str:
  25. prompt = format_prompt(messages)
  26. start = len(prompt)
  27. if start > max_lenght:
  28. if len(messages) > 6:
  29. prompt = format_prompt(messages[:3] + messages[-3:])
  30. if len(prompt) > max_lenght:
  31. if len(messages) > 2:
  32. prompt = format_prompt([m for m in messages if m["role"] == "system"] + messages[-1:])
  33. if len(prompt) > max_lenght:
  34. prompt = messages[-1]["content"]
  35. debug.log(f"Messages trimmed from: {start} to: {len(prompt)}")
  36. return prompt
  37. def get_random_string(length: int = 10) -> str:
  38. """
  39. Generate a random string of specified length, containing lowercase letters and digits.
  40. Args:
  41. length (int, optional): Length of the random string to generate. Defaults to 10.
  42. Returns:
  43. str: A random string of the specified length.
  44. """
  45. return ''.join(
  46. random.choice(string.ascii_lowercase + string.digits)
  47. for _ in range(length)
  48. )
  49. def get_random_hex(length: int = 32) -> str:
  50. """
  51. Generate a random hexadecimal string with n length.
  52. Returns:
  53. str: A random hexadecimal string of n characters.
  54. """
  55. return ''.join(
  56. random.choice("abcdef" + string.digits)
  57. for _ in range(length)
  58. )
  59. def filter_none(**kwargs) -> dict:
  60. return {
  61. key: value
  62. for key, value in kwargs.items()
  63. if value is not None
  64. }
  65. async def async_concat_chunks(chunks: AsyncIterator) -> str:
  66. return concat_chunks([chunk async for chunk in chunks])
  67. def concat_chunks(chunks: Iterator) -> str:
  68. return "".join([
  69. str(chunk) for chunk in chunks
  70. if chunk and not isinstance(chunk, Exception)
  71. ])
  72. def format_cookies(cookies: Cookies) -> str:
  73. return "; ".join([f"{k}={v}" for k, v in cookies.items()])