add agent code final

app.py

@@ -6,7 +6,6 @@ from PIL import Image
 import os
 import sys
 sys.path.append(os.path.dirname(__file__))
-import torch
 from download_models import download_model
 
 
@@ -41,8 +40,15 @@ def main(cfg):
     config_dict = OmegaConf.to_container(cfg, resolve=True)
 
     ## download model from modelscope
-    if not os.path.exists(config_dict["llm_model"]):
+    if not config_dict["use_remote"] and not os.path.exists(config_dict["llm_model"]):
         download_model(llm_model_path =config_dict["llm_model"])
+    
+    ## agent mode, used llama-index, rturn off lmdeloy and chroma rag
+    if cfg.agent_mode:
+        ## load wulewule agent 
+        wulewule_assistant = load_wulewule_agent(config_dict) 
+        cfg.use_rag = False
+        cfg.use_lmdepoly = False
 
     if cfg.use_rag:
         ## load rag model
@@ -72,50 +78,85 @@ def main(cfg):
     # 遍历session_state中的所有消息，并显示在聊天界面上
     for msg in st.session_state.messages:
         st.chat_message("user").write(msg["user"])
-        st.chat_message("assistant").write(msg["assistant"])
+        assistant_res = msg["assistant"]
+        if isinstance(assistant_res, str):
+            st.chat_message("assistant").write(assistant_res)
+        elif cfg.agent_mode and isinstance(assistant_res, dict):
+            image_url = assistant_res["image_url"]
+            audio_text = assistant_res["audio_text"]
+            st.chat_message("assistant").write(assistant_res["response"])
+            if image_url:
+                # 使用st.image展示URL图像，并设置使用列宽
+                st.image( image_url, width=256 )
+            if audio_text:
+                # 使用st.audio函数播放音频
+                st.audio("audio.mp3")
+                st.write(f"语音内容为: \n\n{audio_text}")
 
     # Get user input
     if prompt := st.chat_input("请输入你的问题，换行使用Shfit+Enter。"):
         # Display user input
         st.chat_message("user").write(prompt)
-        # 流式显示, used streaming result
-        if cfg.stream_response:
-            # rag
-            ## 初始化完整的回答字符串
-            full_answer = ""
+        ## 初始化完整的回答字符串
+        full_answer = ""
+        if cfg.agent_mode:
             with st.chat_message('robot'):
                 message_placeholder = st.empty()
-                if cfg.use_rag:
-                    for cur_response in wulewule_model.query_stream(prompt):
-                        full_answer += cur_response
-                        # Display robot response in chat message container
-                        message_placeholder.markdown(full_answer + '▌')
-                elif cfg.use_lmdepoly:
-                    # gen_config = GenerationConfig(top_p=0.8,
-                    #             top_k=40,
-                    #             temperature=0.8,
-                    #             max_new_tokens=2048,
-                    #             repetition_penalty=1.05)
-                    messages = [{'role': 'user', 'content': f'{prompt}'}]
-                    for response in wulewule_model.stream_infer(messages):
-                        full_answer += response.text
-                        # Display robot response in chat message container
-                        message_placeholder.markdown(full_answer + '▌')
-
+                response_dict = wulewule_assistant.chat(prompt)
+                image_url = response_dict["image_url"]
+                audio_text = response_dict["audio_text"]
+                for cur_response in response_dict["response"]:
+                    full_answer += cur_response
+                    # Display robot response in chat message container
+                    message_placeholder.markdown(full_answer + '▌')
                 message_placeholder.markdown(full_answer)
-        # 一次性显示结果
+            # 将问答结果添加到 session_state 的消息历史中
+            st.session_state.messages.append({"user": prompt, "assistant": response_dict})
+            if image_url:
+                # 使用st.image展示URL图像，并设置使用列宽
+                st.image( image_url, width=256 )
+
+            if audio_text:
+                # 使用st.audio函数播放音频
+                st.audio("audio.mp3")
+                st.write(f"语音内容为: \n\n{audio_text}")
+
+        # 流式显示, used streaming result
         else:
-            if cfg.use_lmdepoly:
-                    messages = [{'role': 'user', 'content': f'{prompt}'}]
-                    full_answer = wulewule_model(messages).text
-            elif cfg.use_rag:         
-                full_answer = wulewule_model.query(prompt)
-            # 显示回答
-            st.chat_message("assistant").write(full_answer)
-
-        # 将问答结果添加到 session_state 的消息历史中
-        st.session_state.messages.append({"user": prompt, "assistant": full_answer})
-        torch.cuda.empty_cache()
+            if cfg.stream_response:
+                # rag
+                with st.chat_message('robot'):
+                    message_placeholder = st.empty()
+                    if cfg.use_rag:
+                        for cur_response in wulewule_model.query_stream(prompt):
+                            full_answer += cur_response
+                            # Display robot response in chat message container
+                            message_placeholder.markdown(full_answer + '▌')
+                    elif cfg.use_lmdepoly:
+                        # gen_config = GenerationConfig(top_p=0.8,
+                        #             top_k=40,
+                        #             temperature=0.8,
+                        #             max_new_tokens=2048,
+                        #             repetition_penalty=1.05)
+                        messages = [{'role': 'user', 'content': f'{prompt}'}]
+                        for response in wulewule_model.stream_infer(messages):
+                            full_answer += response.text
+                            # Display robot response in chat message container
+                            message_placeholder.markdown(full_answer + '▌')
+
+                    message_placeholder.markdown(full_answer)
+            # 一次性显示结果
+            else:
+                if cfg.use_lmdepoly:
+                        messages = [{'role': 'user', 'content': f'{prompt}'}]
+                        full_answer = wulewule_model(messages).text
+                elif cfg.use_rag:         
+                    full_answer = wulewule_model.query(prompt)
+                # 显示回答
+                st.chat_message("assistant").write(full_answer)
+
+            # 将问答结果添加到 session_state 的消息历史中
+            st.session_state.messages.append({"user": prompt, "assistant": full_answer})
 
 
 if __name__ == "__main__":