From b4552938e68de47fb9cc8af987b79403e5a146c6 Mon Sep 17 00:00:00 2001
From: shenchucheng <chuchengshen@fuzhi.ai>
Date: Sat, 23 Dec 2023 22:45:20 +0800
Subject: [PATCH] add llm stream log

---
 metagpt/logs.py                       | 13 +++++++++++++
 metagpt/provider/google_gemini_api.py |  4 ++--
 metagpt/provider/ollama_api.py        |  4 ++--
 metagpt/provider/openai_api.py        |  4 ++--
 metagpt/provider/zhipuai_api.py       |  4 ++--
 5 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/metagpt/logs.py b/metagpt/logs.py
index ab1bc4e94..fb0fdd553 100644
--- a/metagpt/logs.py
+++ b/metagpt/logs.py
@@ -8,6 +8,7 @@
 
 import sys
 from datetime import datetime
+from functools import partial
 
 from loguru import logger as _logger
 
@@ -26,3 +27,15 @@ def define_log_level(print_level="INFO", logfile_level="DEBUG"):
 
 
 logger = define_log_level()
+
+
+def log_llm_stream(msg):
+    _llm_stream_log(msg)
+
+
+def set_llm_stream_logfunc(func):
+    global _llm_stream_log
+    _llm_stream_log = func
+
+
+_llm_stream_log = partial(print, end="")
diff --git a/metagpt/provider/google_gemini_api.py b/metagpt/provider/google_gemini_api.py
index 682f7b507..63a2ff687 100644
--- a/metagpt/provider/google_gemini_api.py
+++ b/metagpt/provider/google_gemini_api.py
@@ -20,7 +20,7 @@ from tenacity import (
 )
 
 from metagpt.config import CONFIG, LLMProviderEnum
-from metagpt.logs import logger
+from metagpt.logs import log_llm_stream, logger
 from metagpt.provider.base_gpt_api import BaseGPTAPI
 from metagpt.provider.llm_provider_registry import register_provider
 from metagpt.provider.openai_api import CostManager, log_and_reraise
@@ -119,7 +119,7 @@ class GeminiGPTAPI(BaseGPTAPI):
         collected_content = []
         async for chunk in resp:
             content = chunk.text
-            print(content, end="")
+            log_llm_stream(content, end="")
             collected_content.append(content)
 
         full_content = "".join(collected_content)
diff --git a/metagpt/provider/ollama_api.py b/metagpt/provider/ollama_api.py
index a15c46458..d668d3af1 100644
--- a/metagpt/provider/ollama_api.py
+++ b/metagpt/provider/ollama_api.py
@@ -15,7 +15,7 @@ from tenacity import (
 
 from metagpt.config import CONFIG, LLMProviderEnum
 from metagpt.const import LLM_API_TIMEOUT
-from metagpt.logs import logger
+from metagpt.logs import log_llm_stream, logger
 from metagpt.provider.base_gpt_api import BaseGPTAPI
 from metagpt.provider.general_api_requestor import GeneralAPIRequestor
 from metagpt.provider.llm_provider_registry import register_provider
@@ -127,7 +127,7 @@ class OllamaGPTAPI(BaseGPTAPI):
             if not chunk.get("done", False):
                 content = self.get_choice_text(chunk)
                 collected_content.append(content)
-                print(content, end="")
+                log_llm_stream(content, end="")
             else:
                 # stream finished
                 usage = self.get_usage(chunk)
diff --git a/metagpt/provider/openai_api.py b/metagpt/provider/openai_api.py
index 8fd8959b2..0b6fdd869 100644
--- a/metagpt/provider/openai_api.py
+++ b/metagpt/provider/openai_api.py
@@ -29,7 +29,7 @@ from tenacity import (
 )
 
 from metagpt.config import CONFIG, Config, LLMProviderEnum
-from metagpt.logs import logger
+from metagpt.logs import log_llm_stream, logger
 from metagpt.provider.base_gpt_api import BaseGPTAPI
 from metagpt.provider.constant import GENERAL_FUNCTION_SCHEMA, GENERAL_TOOL_CHOICE
 from metagpt.provider.llm_provider_registry import register_provider
@@ -222,7 +222,7 @@ class OpenAIGPTAPI(BaseGPTAPI, RateLimiter):
                 chunk_message = chunk.choices[0].delta  # extract the message
                 collected_messages.append(chunk_message)  # save the message
                 if chunk_message.content:
-                    print(chunk_message.content, end="")
+                    log_llm_stream(chunk_message.content)
         print()
 
         full_reply_content = "".join([m.content for m in collected_messages if m.content])
diff --git a/metagpt/provider/zhipuai_api.py b/metagpt/provider/zhipuai_api.py
index 60d9a0777..650720d6f 100644
--- a/metagpt/provider/zhipuai_api.py
+++ b/metagpt/provider/zhipuai_api.py
@@ -17,7 +17,7 @@ from tenacity import (
 )
 
 from metagpt.config import CONFIG, LLMProviderEnum
-from metagpt.logs import logger
+from metagpt.logs import log_llm_stream, logger
 from metagpt.provider.base_gpt_api import BaseGPTAPI
 from metagpt.provider.llm_provider_registry import register_provider
 from metagpt.provider.openai_api import CostManager, log_and_reraise
@@ -94,7 +94,7 @@ class ZhiPuAIGPTAPI(BaseGPTAPI):
             if event.event == ZhiPuEvent.ADD.value:
                 content = event.data
                 collected_content.append(content)
-                print(content, end="")
+                log_llm_stream(content)
             elif event.event == ZhiPuEvent.ERROR.value or event.event == ZhiPuEvent.INTERRUPTED.value:
                 content = event.data
                 logger.error(f"event error: {content}", end="")