helper.py 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. from __future__ import annotations
  2. import re
  3. import queue
  4. import threading
  5. import logging
  6. import asyncio
  7. from typing import AsyncIterator, Iterator, AsyncGenerator, Optional
  8. def filter_json(text: str) -> str:
  9. """
  10. Parses JSON code block from a string.
  11. Args:
  12. text (str): A string containing a JSON code block.
  13. Returns:
  14. dict: A dictionary parsed from the JSON code block.
  15. """
  16. match = re.search(r"```(json|)\n(?P<code>[\S\s]+?)\n```", text)
  17. if match:
  18. return match.group("code")
  19. return text
  20. def find_stop(stop: Optional[list[str]], content: str, chunk: str = None):
  21. first = -1
  22. word = None
  23. if stop is not None:
  24. for word in list(stop):
  25. first = content.find(word)
  26. if first != -1:
  27. content = content[:first]
  28. break
  29. if chunk is not None and first != -1:
  30. first = chunk.find(word)
  31. if first != -1:
  32. chunk = chunk[:first]
  33. else:
  34. first = 0
  35. return first, content, chunk
  36. def filter_none(**kwargs) -> dict:
  37. return {
  38. key: value
  39. for key, value in kwargs.items()
  40. if value is not None
  41. }
  42. async def safe_aclose(generator: AsyncGenerator) -> None:
  43. try:
  44. await generator.aclose()
  45. except Exception as e:
  46. logging.warning(f"Error while closing generator: {e}")
  47. # Helper function to convert an async generator to a synchronous iterator
  48. def to_sync_iter(async_gen: AsyncIterator) -> Iterator:
  49. q = queue.Queue()
  50. loop = asyncio.new_event_loop()
  51. done = object()
  52. def _run():
  53. asyncio.set_event_loop(loop)
  54. async def iterate():
  55. try:
  56. async for item in async_gen:
  57. q.put(item)
  58. finally:
  59. q.put(done)
  60. loop.run_until_complete(iterate())
  61. loop.close()
  62. threading.Thread(target=_run).start()
  63. while True:
  64. item = q.get()
  65. if item is done:
  66. break
  67. yield item
  68. # Helper function to convert a synchronous iterator to an async iterator
  69. async def to_async_iterator(iterator: Iterator) -> AsyncIterator:
  70. for item in iterator:
  71. yield item