helper.py 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. from __future__ import annotations
  2. import random
  3. import secrets
  4. import string
  5. from aiohttp import BaseConnector
  6. from ..typing import Messages, Optional
  7. from ..errors import MissingRequirementsError
  8. from ..cookies import get_cookies
  9. def format_prompt(messages: Messages, add_special_tokens=False) -> str:
  10. """
  11. Format a series of messages into a single string, optionally adding special tokens.
  12. Args:
  13. messages (Messages): A list of message dictionaries, each containing 'role' and 'content'.
  14. add_special_tokens (bool): Whether to add special formatting tokens.
  15. Returns:
  16. str: A formatted string containing all messages.
  17. """
  18. if not add_special_tokens and len(messages) <= 1:
  19. return messages[0]["content"]
  20. formatted = "\n".join([
  21. f'{message["role"].capitalize()}: {message["content"]}'
  22. for message in messages
  23. ])
  24. return f"{formatted}\nAssistant:"
  25. def get_random_string(length: int = 10) -> str:
  26. """
  27. Generate a random string of specified length, containing lowercase letters and digits.
  28. Args:
  29. length (int, optional): Length of the random string to generate. Defaults to 10.
  30. Returns:
  31. str: A random string of the specified length.
  32. """
  33. return ''.join(
  34. random.choice(string.ascii_lowercase + string.digits)
  35. for _ in range(length)
  36. )
  37. def get_random_hex() -> str:
  38. """
  39. Generate a random hexadecimal string of a fixed length.
  40. Returns:
  41. str: A random hexadecimal string of 32 characters (16 bytes).
  42. """
  43. return secrets.token_hex(16).zfill(32)
  44. def get_connector(connector: BaseConnector = None, proxy: str = None) -> Optional[BaseConnector]:
  45. if proxy and not connector:
  46. try:
  47. from aiohttp_socks import ProxyConnector
  48. connector = ProxyConnector.from_url(proxy)
  49. except ImportError:
  50. raise MissingRequirementsError('Install "aiohttp_socks" package for proxy support')
  51. return connector