add colored code

app.py

@@ -32,7 +32,9 @@ import torch
 from langchain.chains import LLMMathChain
 from interpreter.code_interpreter import CodeInterpreter
 
-
+from pygments import highlight
+from pygments.lexers import PythonLexer
+from pygments.formatters import HtmlFormatter
 
 langchain.debug = True
 
@@ -43,6 +45,9 @@ CurrentAgent = 'Structured Zero Short Agent'
 global Audio_output
 Audio_output = []
 
+def colored_code_block(code: str):
+    return highlight(code, PythonLexer(), HtmlFormatter(style="monokai"))
+
 class DB_Search2(BaseTool):
     name = "Vector Database Search"
     description = "This is the internal vector database to search information firstly (i.e. engineering data, acronym.)"
@@ -418,5 +423,12 @@ if prompt := st.chat_input(placeholder="Input Your Request"):
         st_cb = StreamlitCallbackHandler(st.container(), expand_new_thoughts=False)
         response = agent.run(prompt, callbacks=[st_cb])
         # st.write(response)
-        st.markdown(response)
+        # st.markdown(response)
+
+        for content in response:
+            if isinstance(content, str):
+                st.write(content)
+            elif isinstance(content, dict) and content.get('type') == 'code':
+                st.markdown(colored_code_block(content.get('code')), unsafe_allow_html=True)
+
         print(msgs.messages)

requirements.txt

@@ -31,3 +31,4 @@ regex
 graphviz
 streamlit
 streamlit_chat
+pygments