sunholo 0.100.2__tar.gz → 0.101.0__tar.gz

{sunholo-0.100.2 → sunholo-0.101.0}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: sunholo
-Version: 0.100.2
+Version: 0.101.0
 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.100.2.tar.gz
+Download-URL: https://github.com/sunholo-data/sunholo-py/archive/refs/tags/v0.101.0.tar.gz
 Author: Holosun ApS
 Author-email: multivac@sunholo.com
 License: Apache License, Version 2.0

{sunholo-0.100.2 → sunholo-0.101.0}/setup.py

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

{sunholo-0.100.2 → sunholo-0.101.0}/sunholo/gcs/add_file.py

@@ -25,6 +25,40 @@ except ImportError:
 from ..custom_logging import log
 from ..utils import ConfigManager
 
+def guess_image_type(file_path: str) -> str:
+    """
+    Guess the image type based on the file extension.
+
+    Args:
+        file_path (str): The path or URL of the image file.
+
+    Returns:
+        str: The guessed image type (e.g., "jpeg", "png", "gif", etc.)
+             or None if the extension is not recognized.
+    """
+    # Extract the file extension
+    _, ext = os.path.splitext(file_path)
+    
+    # Normalize and remove the leading dot
+    ext = ext.lower().strip('.')
+
+    # Mapping of common file extensions to image types
+    extension_to_type = {
+        "jpg": "image/jpeg",
+        "jpeg": "image/jpeg",
+        "png": "image/png",
+        "gif": "image/gif",
+        "bmp": "image/bmp",
+        "tiff": "image/tiff",
+        "tif": "image/tiff",
+        "webp": "image/webp",
+        "ico": "image/ico",
+        "svg": "image/svg",
+        "pdf": "application/pdf",
+    }
+    
+    return extension_to_type.get(ext.lower(), None)
+
 
 def handle_base64_image(base64_data: str, vector_name: str, extension: str):
     """
@@ -58,14 +92,7 @@ def handle_base64_image(base64_data: str, vector_name: str, extension: str):
         os.remove(filename)  # Clean up the saved file
 
         # Determine MIME type based on extension
-        mime_type = {
-            ".jpg": "image/jpeg",
-            ".jpeg": "image/jpeg",
-            ".png": "image/png",
-            ".gif": "image/gif",
-            ".bmp": "image/bmp",
-            ".tiff": "image/tiff"
-        }.get(extension.lower(), "application/octet-stream")  # Default MIME type if unknown
+        mime_type = guess_image_type(extension) or "application/octet-stream"
 
         return image_uri, mime_type
     except Exception as e:

{sunholo-0.100.2 → sunholo-0.101.0}/sunholo/gcs/download_url.py

@@ -176,6 +176,7 @@ def construct_download_link_simple(bucket_name:str, object_name:str) -> Tuple[st
     else:
         public_url = f"https://storage.cloud.google.com/{bucket_name}/{quote(object_name)}"
     filename = os.path.basename(object_name)
+    log.info(f"Created simple download link: {public_url}")
     return public_url, filename, False
 
 def parse_gs_uri(gs_uri: str) -> Tuple[str, str]:

sunholo-0.101.0/sunholo/gcs/extract_and_sign.py

@@ -0,0 +1,41 @@
+import re
+import asyncio
+
+from .download_url import construct_download_link
+from ..utils.mime import guess_mime_type
+from ..custom_logging import log
+
+async def extract_gs_uris_and_sign(content, pattern=r'gs://[^\n]+\.(?:png|jpg|jpeg|pdf|txt|md)'):
+
+    gs_matches = re.findall(pattern, content)
+    unique_gs_matches = set(gs_matches)
+    image_signed_urls = []
+    if unique_gs_matches:
+        log.info(f"Got gs matches: {unique_gs_matches}")
+
+        async def process_link(gs_url):
+            log.info(f"Processing {gs_url}")
+            link, encoded_filename, signed = await asyncio.to_thread(construct_download_link, gs_url)
+            if signed:
+                try:
+                    mime_type = guess_mime_type(gs_url)
+                except Exception as err:
+                    log.error(f"Could not find mime_type for {link} - {str(err)}")
+                    mime_type = "application/octet-stream"
+                    
+                return {
+                    "original": gs_url,
+                    "link": link,
+                    "name": encoded_filename,
+                    "mime": mime_type,
+                    "signed": signed
+                }
+            else:
+                log.info(f"Could not sign this GS_URI: {gs_url} - skipping")
+                return None
+
+        # Gather all tasks and run them concurrently
+        image_signed_urls = await asyncio.gather(*(process_link(gs_url) for gs_url in unique_gs_matches))
+
+    log.info(f"found files to msg: {image_signed_urls}")
+    return image_signed_urls

sunholo-0.101.0/sunholo/genai/images.py

@@ -0,0 +1,38 @@
+import re
+from ..utils.mime import guess_mime_type
+from ..gcs import get_bytes_from_gcs
+from ..custom_logging import log
+import io
+import os
+try:
+    import google.generativeai as genai
+except ImportError:
+    genai = None
+
+def extract_gs_images_and_genai_upload(content:str):
+    # Regular expression to find gs:// URLs 
+    pattern = r'gs://[^ ]+\.(?:png|jpg|jpeg|pdf)'
+
+    gs_matches = re.findall(pattern, content)
+    unique_gs_matches = set(gs_matches)
+    output_gs_images = []
+    for gs_uri in unique_gs_matches:
+        mime_type = guess_mime_type(gs_uri)
+        if mime_type is None:
+            continue
+        
+        log.info(f"Getting bytes from GCS: {gs_uri}")
+        image_bytes = get_bytes_from_gcs(gs_uri)
+        if image_bytes is None:
+            continue
+        image_file = io.BytesIO(image_bytes)
+        image_file.name = os.path.basename(gs_uri)  # Assign a name, as some APIs require it
+        
+        try:
+            uploaded_file = genai.upload_file(image_file)
+            output_gs_images.append(uploaded_file)
+        
+        except Exception as e:
+            log.error(f"Error adding {gs_uri} to base64: {str(e)}")
+    
+    return output_gs_images

{sunholo-0.100.2 → sunholo-0.101.0}/sunholo/genai/process_funcs_cls.py

@@ -17,9 +17,13 @@ try:
 except ImportError:
     genai = None
 
+from .images import extract_gs_images_and_genai_upload
+
 if TYPE_CHECKING:
     from google.generativeai.protos import Part
 
+
+
 class GenAIFunctionProcessor:
     """
     A generic class for processing function calls from google.generativeai function calling models.
@@ -366,20 +370,22 @@ class GenAIFunctionProcessor:
         
         return clean
 
-    def run_agent_loop(self, chat, content, callback, guardrail_max=10, loop_return=3):
+    def run_agent_loop(self, chat, content, callback=None, guardrail_max=10, loop_return=3):
         """
         Runs the agent loop, sending messages to the orchestrator, processing responses, and executing functions.
 
         Args:
             chat: The chat object for interaction with the orchestrator.
             content: The initial content to send to the agent.
-            callback: The callback object for handling intermediate responses.
+            callback: The callback object for handling intermediate responses. If not supplied will use self.IOCallback()
             guardrail_max (int): The maximum number of iterations for the loop.
             loop_return (int): The number of last loop iterations to return. Default 3 will return last 3 iterations. If loop_return > guardrail_max then all iterations are returned.
 
         Returns:
             tuple: (big_text, usage_metadata) from the loop execution.
         """
+        if not callback:
+            callback = self.IOCallback()
         guardrail = 0
         big_result = []
         usage_metadata = {
@@ -535,6 +541,13 @@ class GenAIFunctionProcessor:
 
             if this_text:
                 content.append(f"Agent: {this_text}")    
+                # if text includes gs:// try to download it
+                image_uploads = extract_gs_images_and_genai_upload(this_text)
+                if image_uploads:
+                    for img in image_uploads:
+                        log.info(f"Adding {img=}")
+                        content.append(img)
+                        content.append(f"{img.name} was created by agent and added")
                 log.info(f"[{guardrail}] Updated content:\n{this_text}")
                 big_result.append(this_text)
             else:
@@ -567,6 +580,15 @@ class GenAIFunctionProcessor:
 
         return big_text, usage_metadata
 
+    class IOCallback:
+        """
+        This is a default callback that will print to console any tokens it recieves.
+        """
+        def on_llm_new_token(self, token:str):
+            print(token)
+        def on_llm_end(self, response):
+            print(f"\nFull response: \n{response}")
+
     @staticmethod
     def decide_to_go_on(go_on: bool, chat_summary: str) -> dict:
         """

sunholo-0.101.0/sunholo/invoke/async_class.py

@@ -0,0 +1,178 @@
+import asyncio
+from typing import Callable, Any, AsyncGenerator, Dict
+import time
+import traceback
+from ..custom_logging import setup_logging
+from tenacity import AsyncRetrying, retry_if_exception_type, wait_random_exponential, stop_after_attempt
+
+log = setup_logging("sunholo_AsyncTaskRunner")
+
+class AsyncTaskRunner:
+    def __init__(self, retry_enabled=False, retry_kwargs=None):
+        self.tasks = []
+        self.retry_enabled = retry_enabled
+        self.retry_kwargs = retry_kwargs or {}
+
+    def add_task(self, func: Callable[..., Any], *args: Any):
+        """
+        Adds a task to the list of tasks to be executed.
+        """
+        log.info(f"Adding task: {func.__name__} with args: {args}")
+        self.tasks.append((func.__name__, func, args))
+
+    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.
+        """
+        log.info("Running tasks asynchronously and yielding results as they complete")
+
+        # Create a queue for inter-coroutine communication
+        queue = asyncio.Queue()
+
+        # List to keep track of all running tasks and their heartbeats
+        task_infos = []
+
+        # Start all tasks and their corresponding heartbeats
+        for name, func, args in self.tasks:
+            # Create an event to signal task completion to the heartbeat
+            completion_event = asyncio.Event()
+
+            # Start the main task with retries
+            task_coro = self._run_with_retries(name, func, *args, queue=queue, completion_event=completion_event)
+            task = asyncio.create_task(task_coro)
+
+            # Start the heartbeat coroutine
+            heartbeat_coro = self._send_heartbeat(name, completion_event, queue)
+            heartbeat_task = asyncio.create_task(heartbeat_coro)
+
+            # Store task information for management
+            task_infos.append({
+                'name': name,
+                'task': task,
+                'heartbeat_task': heartbeat_task,
+                'completion_event': completion_event
+            })
+
+            log.info(f"Started task '{name}' and its heartbeat")
+
+        log.info(f"Started async run with {len(self.tasks)} tasks and heartbeats")
+
+        # Create a monitor task to detect when all tasks and heartbeats are done
+        monitor = asyncio.create_task(self._monitor_tasks(task_infos, queue))
+
+        # Continuously yield messages from the queue until sentinel is received
+        while True:
+            message = await queue.get()
+            if message is None:
+                log.info("Received sentinel. Exiting message loop.")
+                break  # Sentinel received, all tasks and heartbeats are done
+            log.info(f"Received message from queue: {message}")
+            yield message
+
+        # Wait for the monitor to finish
+        await monitor
+
+        log.info("All tasks and heartbeats have completed")
+
+    async def _monitor_tasks(self, task_infos, queue):
+        """
+        Monitors the tasks and heartbeats, and sends a sentinel to the queue when done.
+        """
+        # Wait for all main tasks to complete
+        main_tasks = [info['task'] for info in task_infos]
+        log.info("Monitor: Waiting for all main tasks to complete")
+        await asyncio.gather(*main_tasks, return_exceptions=True)
+        log.info("Monitor: All main tasks have completed")
+
+        # Cancel all heartbeat tasks
+        for info in task_infos:
+            info['heartbeat_task'].cancel()
+            try:
+                await info['heartbeat_task']
+            except asyncio.CancelledError:
+                pass
+            log.info(f"Monitor: Heartbeat for task '{info['name']}' has been canceled")
+
+        # Send a sentinel to indicate completion
+        await queue.put(None)
+        log.info("Monitor: Sent sentinel to queue")
+
+    async def _run_with_retries(self, name: str, func: Callable[..., Any], *args: Any, queue: asyncio.Queue, completion_event: asyncio.Event) -> None:
+        """
+        Executes a task with optional retries and sends completion or error messages to the queue.
+        """
+        try:
+            if self.retry_enabled:
+                retry_kwargs = {
+                    'wait': wait_random_exponential(multiplier=1, max=60),
+                    'stop': stop_after_attempt(5),
+                    'retry': retry_if_exception_type(Exception),
+                    **self.retry_kwargs
+                }
+                async for attempt in AsyncRetrying(**retry_kwargs):
+                    with attempt:
+                        log.info(f"Starting task '{name}' with retry")
+                        result = await self._execute_task(func, *args)
+                        await queue.put({'type': 'task_complete', 'func_name': name, 'result': result})
+                        log.info(f"Sent 'task_complete' message for task '{name}'")
+                        return
+            else:
+                log.info(f"Starting task '{name}' with no retry")
+                result = await self._execute_task(func, *args)
+                await queue.put({'type': 'task_complete', 'func_name': name, 'result': result})
+                log.info(f"Sent 'task_complete' message for task '{name}'")
+        except Exception as e:
+            log.error(f"Error in task '{name}': {e}\n{traceback.format_exc()}")
+            await queue.put({'type': 'task_error', 'func_name': name, 'error': f'{e}\n{traceback.format_exc()}'})
+            log.info(f"Sent 'task_error' message for task '{name}'")
+        finally:
+            log.info(f"Task '{name}' completed.")
+            # Set the completion event after sending the message
+            completion_event.set()
+
+    async def _execute_task(self, func: Callable[..., Any], *args: Any) -> Any:
+        """
+        Executes the given task function and returns its result.
+
+        Args:
+            func (Callable): The callable to execute.
+            *args: Arguments to pass to the callable.
+
+        Returns:
+            Any: The result of the task.
+        """
+        if asyncio.iscoroutinefunction(func):
+            return await func(*args)
+        else:
+            return await asyncio.to_thread(func, *args)
+
+    async def _send_heartbeat(self, func_name: str, completion_event: asyncio.Event, queue: asyncio.Queue, interval: int = 2):
+        """
+        Sends periodic heartbeat updates to indicate the task is still in progress.
+
+        Args:
+            func_name (str): The name of the task function.
+            completion_event (asyncio.Event): Event to signal when the task is completed.
+            queue (asyncio.Queue): The queue to send heartbeat messages to.
+            interval (int): How frequently to send heartbeat messages (in seconds).
+        """
+        start_time = time.time()
+        log.info(f"Starting heartbeat for task '{func_name}' with interval {interval} seconds")
+        try:
+            while not completion_event.is_set():
+                await asyncio.sleep(interval)
+                elapsed_time = int(time.time() - start_time)
+                heartbeat_message = {
+                    'type': 'heartbeat',
+                    'name': func_name,
+                    'interval': interval,
+                    'elapsed_time': elapsed_time
+                }
+                log.info(f"Sending heartbeat for task '{func_name}', running for {elapsed_time} seconds")
+                await queue.put(heartbeat_message)
+        except asyncio.CancelledError:
+            log.info(f"Heartbeat for task '{func_name}' has been canceled")
+        finally:
+            log.info(f"Heartbeat for task '{func_name}' stopped")
+

{sunholo-0.100.2 → sunholo-0.101.0}/sunholo/langfuse/prompts.py

@@ -2,13 +2,14 @@ from ..custom_logging import log
 from ..utils import ConfigManager
 
 # Load the YAML file
-def load_prompt_from_yaml(key, prefix="sunholo", load_from_file=False):
+def load_prompt_from_yaml(key, prefix="sunholo", load_from_file=False, f_string=True):
     """
-    Returns a string you can use with Langfuse PromptTemplate.from_template() 
+    Returns a string you can use with prompts.
 
     Will first try to load from the Langfuse prompt library, if unavailable will look in promptConfig type file.
 
-    Langfuse prompts have {{ two braces }}, Langchain prompts have { one brace }.
+    If f_string is True will be in a Langchain style prompt e.g. { one brace }
+    If f_string is False will be Langfuse style prompt e.g. {{ two braces }}
 
     Example:
 
@@ -36,7 +37,10 @@ def load_prompt_from_yaml(key, prefix="sunholo", load_from_file=False):
             
         langfuse_prompt = langfuse.get_prompt(langfuse_template, cache_ttl_seconds=300)
 
-        return langfuse_prompt.get_langchain_prompt()
+        if f_string:
+            return langfuse_prompt.get_langchain_prompt()
+        
+        return langfuse_prompt
     
     except Exception as err:
         log.warning(f"Could not find langfuse template: {langfuse_template} - {str(err)} - attempting to load from promptConfig")

{sunholo-0.100.2 → sunholo-0.101.0}/sunholo/streaming/content_buffer.py

@@ -106,7 +106,6 @@ class ContentBuffer:
         """
         async with self.lock:
             self.content = ""
-            log.debug("Content buffer cleared")
             self.content_available.clear()
 
 

sunholo-0.101.0/sunholo/utils/mime.py

@@ -0,0 +1,63 @@
+import os
+
+def guess_mime_type(file_path: str) -> str:
+    """
+    Guess the mime type based on the file extension.
+
+    Args:
+        file_path (str): The path or URL of the image file.
+
+    Returns:
+        str: The guessed image type (e.g., "jpeg", "png", "gif", etc.)
+             or None if the extension is not recognized.
+    """
+    # Extract the file extension
+    _, ext = os.path.splitext(file_path)
+    
+    # Normalize and remove the leading dot
+    ext = ext.lower().strip('.')
+
+    # Mapping of common file extensions to file types
+    extension_to_type = {
+        "jpg": "image/jpeg",
+        "jpeg": "image/jpeg",
+        "png": "image/png",
+        "gif": "image/gif",
+        "bmp": "image/bmp",
+        "tiff": "image/tiff",
+        "tif": "image/tiff",
+        "webp": "image/webp",
+        "ico": "image/vnd.microsoft.icon",
+        "svg": "image/svg+xml",
+        "pdf": "application/pdf",
+        "txt": "text/plain",
+        "md": "text/markdown",
+        "html": "text/html",
+        "css": "text/css",
+        "js": "application/javascript",
+        "json": "application/json",
+        "xml": "application/xml",
+        "csv": "text/csv",
+        "py": "text/x-python",
+        "java": "text/x-java-source",
+        "c": "text/x-c",
+        "cpp": "text/x-c++",
+        "h": "text/x-c",
+        "hpp": "text/x-c++",
+        "sh": "application/x-sh",
+        "bat": "application/x-msdos-program",
+        "php": "application/x-httpd-php",
+        "rb": "application/x-ruby",
+        "pl": "application/x-perl",
+        "swift": "application/x-swift",
+        "r": "text/x-r",
+        "go": "text/x-go",
+        "sql": "application/sql",
+        "yaml": "text/yaml",
+        "yml": "text/yaml",
+        "ts": "application/typescript",
+        "tsx": "text/tsx",
+        "jsx": "text/jsx",
+    }
+    
+    return extension_to_type.get(ext, "")

{sunholo-0.100.2 → sunholo-0.101.0}/sunholo.egg-info/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: sunholo
-Version: 0.100.2
+Version: 0.101.0
 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.100.2.tar.gz
+Download-URL: https://github.com/sunholo-data/sunholo-py/archive/refs/tags/v0.101.0.tar.gz
 Author: Holosun ApS
 Author-email: multivac@sunholo.com
 License: Apache License, Version 2.0

{sunholo-0.100.2 → sunholo-0.101.0}/sunholo.egg-info/SOURCES.txt

@@ -94,8 +94,10 @@ sunholo/gcs/__init__.py
 sunholo/gcs/add_file.py
 sunholo/gcs/download_folder.py
 sunholo/gcs/download_url.py
+sunholo/gcs/extract_and_sign.py
 sunholo/gcs/metadata.py
 sunholo/genai/__init__.py
+sunholo/genai/images.py
 sunholo/genai/init.py
 sunholo/genai/process_funcs_cls.py
 sunholo/genai/safety.py
@@ -143,6 +145,7 @@ sunholo/utils/config_class.py
 sunholo/utils/config_schema.py
 sunholo/utils/gcp.py
 sunholo/utils/gcp_project.py
+sunholo/utils/mime.py
 sunholo/utils/parsers.py
 sunholo/utils/timedelta.py
 sunholo/utils/user_ids.py
@@ -155,5 +158,6 @@ sunholo/vertex/init.py
 sunholo/vertex/memory_tools.py
 sunholo/vertex/safety.py
 sunholo/vertex/type_dict_to_json.py
+tests/test_async.py
 tests/test_chat_history.py
 tests/test_config.py

sunholo-0.101.0/tests/test_async.py

@@ -0,0 +1,93 @@
+import asyncio
+from sunholo.invoke import AsyncTaskRunner
+# Mock logger for demonstration
+import logging
+logging.basicConfig(level=logging.INFO)
+log = logging.getLogger(__name__)
+
+
+# Dummy task functions for testing
+async def google_search(query, config_manager):
+    await asyncio.sleep(5)  # Simulate a long-running task
+    return f"Results for '{query}'"
+
+async def use_pdfs(query, arg2, arg3):
+    await asyncio.sleep(2)  # Simulate a medium-running task
+    return "No image_uri specified\n"
+
+async def process_urls(url):
+    await asyncio.sleep(1)  # Simulate a short-running task
+    return "No URLs were found"
+
+# Mock callback for demonstration
+class MockCallback:
+    async def async_on_llm_new_token(self, token):
+        log.info(f"Callback received token: {token}")
+
+# Mock functions for demonstration
+def format_token_output(func_name, tokens, result=None, error=None):
+    if error:
+        return f"Error in {func_name}: {error}"
+    return f"{func_name} completed with {tokens} tokens. Result: {result}"
+
+async def count_tokens(result):
+    return len(result.split())
+
+# Example usage
+async def main():
+    runner = AsyncTaskRunner(retry_enabled=True)
+    
+    # Add your tasks
+    runner.add_task(google_search, "<original user question>...", "<config>")
+    runner.add_task(use_pdfs, "<original user question>...", None, None)
+    runner.add_task(process_urls, "please give me a forecast of ppas for spain using wind energy")
+    
+    callback = MockCallback()
+    answers = {}
+    total_context_tokens = 0
+    
+    async for message in runner.run_async_as_completed():
+        log.info(f"Runner message={message}")
+        if message['type'] == 'heartbeat':
+            func_name = message['name']
+            elapsed_time = message['elapsed_time']
+            log.info(f"Runner Heartbeat for {func_name}, elapsed_time={elapsed_time}")
+            # Send heartbeat to callback
+            update_html = (
+                f'<div style="display: none;" data-update-id="{func_name}-spinner">'
+                f'<span class="elapsed-time">{elapsed_time} seconds elapsed</span>'
+                f'</div>'
+            )
+            await callback.async_on_llm_new_token(token=f"[[HEARTBEAT]]{update_html}[[/HEARTBEAT]]")
+        elif message['type'] == 'task_complete':
+            func_name = message['func_name']
+            result = message['result']
+            log.info(f"Runner completed task: {func_name}")
+            # Process result
+            if isinstance(result, Exception):
+                log.info(f"Error Exception for {func_name}: {str(result)}")
+                formatted_output = format_token_output(func_name, 0, error="No results")
+                await callback.async_on_llm_new_token(token=formatted_output)
+            else:
+                # Stream the result to the callback
+                tokens = await count_tokens(result)
+                total_context_tokens += tokens
+                log.info(f"Got task {func_name} to stream result length [{len(result)}] [{tokens} tokens]")
+                formatted_output = format_token_output(func_name, tokens, result=result)
+                await callback.async_on_llm_new_token(token=formatted_output)
+                answers[func_name] = result
+        elif message['type'] == 'task_error':
+            func_name = message['func_name']
+            error = message['error']
+            log.info(f"Error Exception for {func_name}: {str(error)}")
+            formatted_output = format_token_output(func_name, 0, error="No results")
+            await callback.async_on_llm_new_token(token=formatted_output)
+
+    await callback.async_on_llm_new_token(
+        token=f'<div style="margin-top: 20px; color: #333;"><strong>-- Using [{total_context_tokens}] tokens in context for final answer --</strong></div>'
+    )
+    log.info("All tasks have been processed.")
+
+# Run the main coroutine
+if __name__ == "__main__":
+    asyncio.run(main())