sunholo 0.99.13__tar.gz → 0.100.1__tar.gz

{sunholo-0.99.13 → sunholo-0.100.1}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: sunholo
-Version: 0.99.13
+Version: 0.100.1
 Summary: Large Language Model DevOps - a package to help deploy LLMs to the Cloud.
 Home-page: https://github.com/sunholo-data/sunholo-py
-Download-URL: https://github.com/sunholo-data/sunholo-py/archive/refs/tags/v0.99.13.tar.gz
+Download-URL: https://github.com/sunholo-data/sunholo-py/archive/refs/tags/v0.100.1.tar.gz
 Author: Holosun ApS
 Author-email: multivac@sunholo.com
 License: Apache License, Version 2.0

{sunholo-0.99.13 → sunholo-0.100.1}/setup.py

@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-version = '0.99.13'
+version = '0.100.1'
 
 setup(
     name='sunholo',

{sunholo-0.99.13 → sunholo-0.100.1}/sunholo/invoke/async_class.py

@@ -15,47 +15,62 @@ class AsyncTaskRunner:
         log.info(f"Adding task: {func.__name__} with args: {args}")
         self.tasks.append((func.__name__, func, args))
 
-    async def run_async_as_completed(self, callback=None) -> AsyncGenerator[Dict[str, Any], None]:
+    async def run_async_as_completed(self) -> AsyncGenerator[Dict[str, Any], None]:
         """
-        Runs all tasks concurrently and yields results as they complete, while periodically sending heartbeat messages.
-
-        Args:
-            callback: The callback object that will receive heartbeat messages.
+        Runs all tasks concurrently and yields results and heartbeats as they are produced.
         """
-        log.info("Running tasks asynchronously and yielding results as they complete")
+        log.info("Running tasks asynchronously and yielding results and heartbeats as they occur")
+        queue = asyncio.Queue()
         tasks = {}
+        completed_tasks = set()
+
         for name, func, args in self.tasks:
-            # Pass the callback down to _task_wrapper
-            coro = self._task_wrapper(name, func, args, callback)
+            coro = self._task_wrapper(name, func, args, queue)
             task = asyncio.create_task(coro)
             tasks[task] = name
 
-        log.info(f"Start async run with {len(self.tasks)} runners")
-        while tasks:
-            done, _ = await asyncio.wait(tasks.keys(), return_when=asyncio.FIRST_COMPLETED)
-            for task in done:
-                name = tasks.pop(task)
-                try:
-                    result = await task
-                    yield {name: result}
-                except Exception as e:
-                    log.error(f"Task {name} resulted in an error: {e}\n{traceback.format_exc()}")
-                    yield {name: e}
+        while tasks or not queue.empty():
+            if not queue.empty():
+                message = await queue.get()
+                log.info(f"Found queue message: {message}")
+                # Ignore heartbeats from completed tasks
+                if message['type'] == 'heartbeat' and message['func_name'] in completed_tasks:
+                    continue
+                yield message
+            else:
+                done, _ = await asyncio.wait(
+                    list(tasks.keys()),
+                    timeout=0.1,
+                    return_when=asyncio.FIRST_COMPLETED
+                )
+                for task in done:
+                    name = tasks.pop(task)
+                    completed_tasks.add(name)
+                    try:
+                        result = await task
+                        await queue.put({'type': 'task_complete', 'func_name': name, 'result': result})
+                    except Exception as e:
+                        log.error(f"Task {name} resulted in an error: {e}\n{traceback.format_exc()}")
+                        await queue.put({'type': 'task_error', 'func_name': name, 'error': e})
 
-    async def _task_wrapper(self, name: str, func: Callable[..., Any], args: Any, callback=None) -> Any:
+        # Process any remaining messages in the queue
+        while not queue.empty():
+            message = await queue.get()
+            log.info(f"Found queue message: {message}")
+            if message['type'] == 'heartbeat' and message['func_name'] in completed_tasks:
+                continue
+            yield message
+
+    async def _task_wrapper(self, name: str, func: Callable[..., Any], args: Any, queue: asyncio.Queue) -> Any:
         """Wraps the task function to process its output and handle retries, while managing heartbeat updates."""
         async def run_func():
             if asyncio.iscoroutinefunction(func):
-                # If the function is async, await it
                 return await func(*args)
             else:
-                # If the function is sync, run it in a thread to prevent blocking
                 return await asyncio.to_thread(func, *args)
 
-        # Start the heartbeat task if a callback is provided
-        heartbeat_task = None
-        if callback:
-            heartbeat_task = asyncio.create_task(self._send_heartbeat(callback, name))
+        # Start the heartbeat task
+        heartbeat_task = asyncio.create_task(self._send_heartbeat(queue, name))
 
         try:
             if self.retry_enabled:
@@ -76,29 +91,16 @@ class AsyncTaskRunner:
                     raise
         finally:
             # Stop the heartbeat task
-            if heartbeat_task:
-                heartbeat_task.cancel()
-                # Let the heartbeat_task finish in the background without awaiting it
-                try:
-                    await asyncio.shield(heartbeat_task)  # Ensure cancellation is handled cleanly
-                except asyncio.CancelledError:
-                    pass
-
-            # Send a message indicating task completion to update the spinner's state
-            if callback:
-                completion_html = (
-                    f'<div style="display: none;" data-complete-id="{name}-spinner"></div>'
-                )
-                await callback.async_on_llm_new_token(token=completion_html)
+            heartbeat_task.cancel()
+            # Wait for the heartbeat task to finish
+            try:
+                await heartbeat_task
+            except asyncio.CancelledError:
+                pass
 
-    async def _send_heartbeat(self, callback, func_name, interval=2):
+    async def _send_heartbeat(self, queue: asyncio.Queue, func_name: str, interval=2):
         """
         Sends a periodic heartbeat to keep the task alive and update the spinner with elapsed time.
-
-        Args:
-            callback: The callback to notify that the task is still working.
-            func_name: The name of the task function.
-            interval: How frequently to send heartbeat messages (in seconds).
         """
         # Send the initial spinner HTML
         spinner_html = (
@@ -107,23 +109,30 @@ class AsyncTaskRunner:
             f'  <span class="elapsed-time">Task {func_name} is still running... 0s elapsed</span>'
             f'</div>'
         )
-        await callback.async_on_llm_new_token(token=spinner_html)
+        log.info(f"Heartbeat started for task {func_name}")
+
+        await queue.put({'type': 'heartbeat', 'func_name': func_name, 'token': spinner_html})
 
         # Keep track of elapsed time
         elapsed_time = 0
 
-        # Keep sending heartbeats until task completes
-        while True:
-            try:
+        try:
+            while True:
                 await asyncio.sleep(interval)  # Sleep for the interval
                 elapsed_time += interval  # Increment elapsed time
-
+                log.info(f"Sending heartbeat for {func_name}: {elapsed_time}s elapsed")
                 # Update spinner with the elapsed time
                 update_html = (
                     f'<div style="display: none;" data-update-id="{func_name}-spinner">'
                     f'<span class="elapsed-time">Task {func_name} is still running... {elapsed_time}s elapsed</span>'
                     f'</div>'
                 )
-                await callback.async_on_llm_new_token(token=update_html)
-            except asyncio.CancelledError:
-                break  # Exit the loop if the heartbeat task is canceled
+                await queue.put({'type': 'heartbeat', 'func_name': func_name, 'token': update_html})
+        except asyncio.CancelledError:
+            log.info(f"Heartbeat task for {func_name} has been cancelled.")
+        finally:
+            # Send a message indicating task completion to update the spinner's state
+            completion_html = (
+                f'<div style="display: none;" data-complete-id="{func_name}-spinner"></div>'
+            )
+            await queue.put({'type': 'heartbeat', 'func_name': func_name, 'token': completion_html})

{sunholo-0.99.13 → sunholo-0.100.1}/sunholo/streaming/content_buffer.py

@@ -16,6 +16,7 @@ from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from langchain.schema import LLMResult
 
 import threading
+import asyncio
 import re
 from ..custom_logging import log
 
@@ -39,6 +40,8 @@ class ContentBuffer:
         that the buffer has been created.
         """
         self.content = ""
+        self.lock = asyncio.Lock()  # Use an async lock to ensure atomic writes
+        self.content_available = asyncio.Event()
         log.debug("Content buffer initialized")
     
     def write(self, text: str):
@@ -61,7 +64,9 @@ class ContentBuffer:
 
         Adds the given text to the existing content of the buffer.
         """
-        self.content += text
+        async with self.lock:
+            self.content += text
+            self.content_available.set()
     
     def read(self) -> str:
         """
@@ -81,7 +86,9 @@ class ContentBuffer:
         Returns:
             str: The content of the buffer.
         """
-        return self.content
+        async with self.lock:
+            content = self.content
+            return content
 
     def clear(self):
         """
@@ -97,7 +104,10 @@ class ContentBuffer:
 
         Empties the buffer content, resetting it to an empty string.
         """
-        self.content = ""
+        async with self.lock:
+            self.content = ""
+            log.debug("Content buffer cleared")
+            self.content_available.clear()
 
 
 class BufferStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
@@ -137,11 +147,20 @@ class BufferStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
         self.tokens = tokens
         self.buffer = ""
         self.stream_finished = threading.Event()
+
         self.in_code_block = False
         self.in_question_block = False
         self.question_buffer = ""
         log.info("Starting to stream LLM")
 
+    def _is_heartbeat_token(self, token: str) -> bool:
+        """Detects if the token is a heartbeat message."""
+        return token.startswith('[[HEARTBEAT]]') and token.endswith('[[/HEARTBEAT]]')
+
+    def _strip_heartbeat_markers(self, token: str) -> str:
+        """Removes the [[HEARTBEAT]] markers from the token."""
+        return token[len('[[HEARTBEAT]]'):-len('[[/HEARTBEAT]]')]
+
     def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
         """
         Processes a new token from the LLM output.
@@ -154,33 +173,22 @@ class BufferStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
         The buffer content is written to the content buffer when appropriate tokens or
         patterns are detected.
         """
-        log.debug(f"on_llm_new_token: {token}")
-        
-        self.buffer += token
-
-        if '```' in token:
-            self.in_code_block = not self.in_code_block
-
-        if not self.in_code_block:
-            self._process_buffer()
-
-    async def async_on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """
-        Asynchronously processes a new token from the LLM output.
-
-        Args:
-            token (str): The new token generated by the LLM.
-            **kwargs: Additional keyword arguments.
-        """
-        log.debug(f"async_on_llm_new_token: {token}")
-        
-        self.buffer += token
+        #log.debug(f"on_llm_new_token: {token}")
+
+        # Check if the token is a heartbeat message
+        if self._is_heartbeat_token(token):
+            # Strip the [[HEARTBEAT]] markers and write immediately
+            heartbeat_content = self._strip_heartbeat_markers(token)
+            #log.info(f"Heartbeat token detected, writing immediately: {heartbeat_content}")
+            self.content_buffer.write(heartbeat_content)
+        else:
+            self.buffer += token
 
-        if '```' in token:
-            self.in_code_block = not self.in_code_block
+            if '```' in token:
+                self.in_code_block = not self.in_code_block
 
-        if not self.in_code_block:
-            await self._async_process_buffer()
+            if not self.in_code_block:
+                self._process_buffer()
 
     def _process_buffer(self):
         """
@@ -200,24 +208,6 @@ class BufferStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
                 self.content_buffer.write(self.buffer)
                 self.buffer = ""
 
-    async def _async_process_buffer(self):
-        """
-        Asynchronously processes the buffer content and writes to the content buffer.
-
-        If the buffer ends with a numbered list pattern or specified tokens, the buffer is flushed
-        to the content buffer. Otherwise, the buffer is left intact for further accumulation.
-        """
-        matches = list(re.finditer(r'\n(\d+\.\s)', self.buffer))
-        if matches:
-            last_match = matches[-1]
-            start_of_last_match = last_match.start() + 1
-            await self.content_buffer.async_write(self.buffer[:start_of_last_match])
-            self.buffer = self.buffer[start_of_last_match:]
-        else:
-            if any(self.buffer.endswith(t) for t in self.tokens):
-                await self.content_buffer.async_write(self.buffer)
-                self.buffer = ""
-
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """
         Handles the end of LLM streaming.
@@ -237,6 +227,90 @@ class BufferStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
         self.stream_finished.set()
         log.info("Streaming LLM response ended successfully")
 
+
+
+class BufferStreamingStdOutCallbackHandlerAsync(StreamingStdOutCallbackHandler):
+    """
+    An async callback handler for streaming LLM output to a content buffer.
+
+    This class handles the streaming of output from a large language model (LLM),
+    processes tokens from the model output, and writes them to a ContentBuffer.
+    It supports handling different types of tokens and keeps track of code blocks
+    and questions.
+
+    Attributes:
+        content_buffer (ContentBuffer): The buffer to which content is streamed.
+        tokens (str): Tokens that indicate the end of a statement, for buffer flushing.
+        buffer (str): Temporary storage for accumulating streamed tokens.
+        stream_finished (asyncio.Event): Signals when the streaming is finished.
+        in_code_block (bool): Indicates whether the current context is a code block.
+    """
+
+    def __init__(self, content_buffer: ContentBuffer, tokens: str = ".?!\n", *args, **kwargs):
+        """
+        Initializes a new BufferStreamingStdOutCallbackHandler instance.
+
+        Args:
+            content_buffer (ContentBuffer): The buffer to which content will be written.
+            tokens (str): Tokens that indicate the end of a statement (default: ".?!\n").
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Sets up the callback handler with the given content buffer and tokens.
+        Initializes tracking variables for code blocks, buffer content, and the finished signal.
+        """
+        super().__init__(*args, **kwargs)
+        self.content_buffer = content_buffer
+        self.tokens = tokens
+        self.buffer = ""
+        self.stream_finished = asyncio.Event()  # Use asyncio.Event for async compatibility
+        self.in_code_block = False
+        log.info("Starting to stream LLM")
+
+    async def async_on_llm_new_token(self, token: str, **kwargs: Any) -> None:
+        #log.debug(f"async_on_llm_new_token: {token}")
+
+        # Check if the token is a heartbeat message
+        if self._is_heartbeat_token(token):
+            # Strip the [[HEARTBEAT]] markers and write immediately
+            heartbeat_content = self._strip_heartbeat_markers(token)
+            #log.info(f"Heartbeat token detected, writing immediately: {heartbeat_content}")
+            await self.content_buffer.async_write(heartbeat_content)
+        else:
+            self.buffer += token
+
+            if '```' in token:
+                self.in_code_block = not self.in_code_block
+
+            if not self.in_code_block:
+                await self._async_process_buffer()
+
+    def _is_heartbeat_token(self, token: str) -> bool:
+        """Detects if the token is a heartbeat message."""
+        return token.startswith('[[HEARTBEAT]]') and token.endswith('[[/HEARTBEAT]]')
+
+    def _strip_heartbeat_markers(self, token: str) -> str:
+        """Removes the [[HEARTBEAT]] markers from the token."""
+        return token[len('[[HEARTBEAT]]'):-len('[[/HEARTBEAT]]')]
+
+    async def _async_process_buffer(self):
+        """
+        Asynchronously processes the buffer content and writes to the content buffer.
+
+        If the buffer ends with a numbered list pattern or specified tokens, the buffer is flushed
+        to the content buffer. Otherwise, the buffer is left intact for further accumulation.
+        """
+        matches = list(re.finditer(r'\n(\d+\.\s)', self.buffer))
+        if matches:
+            last_match = matches[-1]
+            start_of_last_match = last_match.start() + 1
+            await self.content_buffer.async_write(self.buffer[:start_of_last_match])
+            self.buffer = self.buffer[start_of_last_match:]
+        else:
+            if any(self.buffer.endswith(t) for t in self.tokens):
+                await self.content_buffer.async_write(self.buffer)
+                self.buffer = ""
+
     async def async_on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """
         Asynchronously handles the end of LLM streaming.
@@ -251,4 +325,5 @@ class BufferStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
             log.info("Flushing remaining LLM response buffer")
 
         self.stream_finished.set()
-        log.info("Streaming LLM response ended successfully")
+        log.info("Streaming LLM response ended successfully")
+

{sunholo-0.99.13 → sunholo-0.100.1}/sunholo/streaming/streaming.py

@@ -17,7 +17,7 @@ import json
 import time
 import asyncio
 
-from .content_buffer import ContentBuffer, BufferStreamingStdOutCallbackHandler
+from .content_buffer import ContentBuffer, BufferStreamingStdOutCallbackHandler, BufferStreamingStdOutCallbackHandlerAsync
 
 from ..qna.parsers import parse_output
 
@@ -159,7 +159,7 @@ async def start_streaming_chat_async(question, vector_name, qna_func_async, chat
         yield "No **kwargs in qna_func_async - please add it"
         
     content_buffer = ContentBuffer()
-    chat_callback_handler = BufferStreamingStdOutCallbackHandler(content_buffer=content_buffer, tokens=".!?\n")
+    chat_callback_handler = BufferStreamingStdOutCallbackHandlerAsync(content_buffer=content_buffer, tokens=".!?\n")
 
     result_queue = Queue()
     exception_queue = Queue()
@@ -179,32 +179,26 @@ async def start_streaming_chat_async(question, vector_name, qna_func_async, chat
     await asyncio.sleep(0)
 
     # Read and yield any initial content from the content buffer
-    content_to_send = content_buffer.read()
+    content_to_send = await content_buffer.async_read()
     if content_to_send:
         log.info(f"Initial content: {content_to_send}")
         yield content_to_send
-        content_buffer.clear()
-
-    start = time.time()
+        await content_buffer.async_clear()
 
     while not chat_callback_handler.stream_finished.is_set() and not stop_event.is_set():
-        await asyncio.sleep(wait_time)  # Use asyncio.sleep for async compatibility
-        log.info(f"async heartbeat - {round(time.time() - start, 2)} seconds")
-        
-        while not exception_queue.empty():
-            exception = exception_queue.get_nowait()
-            raise exception
+        try:
+            await asyncio.wait_for(content_buffer.content_available.wait(), timeout=timeout)
+        except asyncio.TimeoutError:
+            log.warning(f"Content production has timed out after {timeout} seconds")
+            break
 
-        content_to_send = content_buffer.read()
+        content_to_send = await content_buffer.async_read()
         if content_to_send:
-            log.info(f"==\n{content_to_send}")
+            log.info(f"==Async\n{content_to_send}")
             yield content_to_send
-            content_buffer.clear()
-            start = time.time()
+            await content_buffer.async_clear()
         else:
-            if time.time() - start > timeout:
-                log.warning(f"Content production has timed out after {timeout} seconds")
-                break
+            log.debug("Content available event set but no content found.")
 
     stop_event.set()
     await chat_task  # Ensure the async task is awaited
@@ -217,8 +211,6 @@ async def start_streaming_chat_async(question, vector_name, qna_func_async, chat
             yield json.dumps(parsed_final_result)
 
 
-
-
 def generate_proxy_stream(stream_to_f, user_input, vector_name, chat_history, generate_f_output, **kwargs):
     """
     Generator function for proxying and processing streaming output from an underlying model.

{sunholo-0.99.13 → sunholo-0.100.1}/sunholo.egg-info/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: sunholo
-Version: 0.99.13
+Version: 0.100.1
 Summary: Large Language Model DevOps - a package to help deploy LLMs to the Cloud.
 Home-page: https://github.com/sunholo-data/sunholo-py
-Download-URL: https://github.com/sunholo-data/sunholo-py/archive/refs/tags/v0.99.13.tar.gz
+Download-URL: https://github.com/sunholo-data/sunholo-py/archive/refs/tags/v0.100.1.tar.gz
 Author: Holosun ApS
 Author-email: multivac@sunholo.com
 License: Apache License, Version 2.0