import re

import streamlit as st

class StreamToStreamlit:

def __init__(self, expander):

self.expander = expander

self.buffer = []

self.colors = ['red', 'green', 'blue', 'orange'] # Define a list of colors

self.color_index = 0 # Initialize color index

def write(self, data):

# Filter out ANSI escape codes using a regular expression

cleaned_data = re.sub(r'\x1B\[[0-9;]*[mK]', '', data)

# Check if the data contains 'task' information

task_match_object = re.search(r'\"task\"\s*:\s*\"(.*?)\"', cleaned_data, re.IGNORECASE)

task_match_input = re.search(r'task\s*:\s*([^\n]*)', cleaned_data, re.IGNORECASE)

task_value = None

if task_match_object:

task_value = task_match_object.group(1)

elif task_match_input:

task_value = task_match_input.group(1).strip()

if task_value:

st.info(":robot_face: " + task_value)

# Check if the text contains the specified phrase and apply color

if "Entering new CrewAgentExecutor chain" in cleaned_data:

# Apply different color and switch color index

self.color_index = (self.color_index + 1) % len(

self.colors) # Increment color index and wrap around if necessary

cleaned_data = cleaned_data.replace("Entering new CrewAgentExecutor chain",

f":{self.colors[self.color_index]}[Entering new CrewAgentExecutor chain]")

if "Market Research Analyst" in cleaned_data:

# Apply different color

cleaned_data = cleaned_data.replace("Market Research Analyst",

f":{self.colors[self.color_index]}[Market Research Analyst]")

if "Business Development Consultant" in cleaned_data:

cleaned_data = cleaned_data.replace("Business Development Consultant",

f":{self.colors[self.color_index]}[Business Development Consultant]")

if "Technology Expert" in cleaned_data:

cleaned_data = cleaned_data.replace("Technology Expert",

f":{self.colors[self.color_index]}[Technology Expert]")

if "Finished chain." in cleaned_data:

cleaned_data = cleaned_data.replace("Finished chain.", f":{self.colors[self.color_index]}[Finished chain.]")

self.buffer.append(cleaned_data)

if "\n" in data:

self.expander.markdown(''.join(self.buffer), unsafe_allow_html=True)

self.buffer = []