Add the saved token report, auto flush data (#401)

Signed-off-by: SimFG <[email protected]>

Author: SimFG

gptcache/__init__.py

@@ -1,7 +1,6 @@
 """gptcache version"""
 __version__ = "0.1.28"
 
-from gptcache.client import Client
 from gptcache.config import Config
 from gptcache.core import Cache
 from gptcache.core import cache

gptcache/adapter/adapter.py

@@ -36,7 +36,11 @@ def adapt(llm_handler, cache_data_convert, update_cache_callback, *args, **kwarg
     else:  # temperature <= 0
         cache_skip = kwargs.pop("cache_skip", False)
     cache_factor = kwargs.pop("cache_factor", 1.0)
-    pre_embedding_res = chat_cache.pre_embedding_func(
+    pre_embedding_res = time_cal(
+        chat_cache.pre_embedding_func,
+        func_name="pre_process",
+        report_func=chat_cache.report.pre,
+    )(
         kwargs,
         extra_param=context.get("pre_embedding_func", None),
         prompts=chat_cache.config.prompts,
@@ -81,7 +85,11 @@ def adapt(llm_handler, cache_data_convert, update_cache_callback, *args, **kwarg
             else rank_threshold
         )
         for cache_data in cache_data_list:
-            ret = chat_cache.data_manager.get_scalar_data(
+            ret = time_cal(
+                chat_cache.data_manager.get_scalar_data,
+                func_name="get_data",
+                report_func=chat_cache.report.data,
+            )(
                 cache_data,
                 extra_param=context.get("get_scalar_data", None),
                 session=session,
@@ -112,7 +120,11 @@ def adapt(llm_handler, cache_data_convert, update_cache_callback, *args, **kwarg
                     "search_result": cache_data,
                     "embedding": ret.embedding_data,
                 }
-            rank = chat_cache.similarity_evaluation.evaluation(
+            rank = time_cal(
+                chat_cache.similarity_evaluation.evaluation,
+                func_name="evaluation",
+                report_func=chat_cache.report.evaluation,
+            )(
                 eval_query_data,
                 eval_cache_data,
                 extra_param=context.get("evaluation_func", None),
@@ -129,16 +141,25 @@ def adapt(llm_handler, cache_data_convert, update_cache_callback, *args, **kwarg
         cache_answers = sorted(cache_answers, key=lambda x: x[0], reverse=True)
         answers_dict = dict((d[1], d[2]) for d in cache_answers)
         if len(cache_answers) != 0:
-            if chat_cache.post_process_messages_func is temperature_softmax:
-                return_message = chat_cache.post_process_messages_func(
-                    messages=[t[1] for t in cache_answers],
-                    scores=[t[0] for t in cache_answers],
-                    temperature=temperature,
-                )
-            else:
-                return_message = chat_cache.post_process_messages_func(
-                    [t[1] for t in cache_answers]
-                )
+
+            def post_process():
+                if chat_cache.post_process_messages_func is temperature_softmax:
+                    return_message = chat_cache.post_process_messages_func(
+                        messages=[t[1] for t in cache_answers],
+                        scores=[t[0] for t in cache_answers],
+                        temperature=temperature,
+                    )
+                else:
+                    return_message = chat_cache.post_process_messages_func(
+                        [t[1] for t in cache_answers]
+                    )
+                return return_message
+
+            return_message = time_cal(
+                post_process,
+                func_name="post_process",
+                report_func=chat_cache.report.post,
+            )()
             chat_cache.report.hint_cache()
             if session:
                 chat_cache.data_manager.add_session(
@@ -156,7 +177,9 @@ def adapt(llm_handler, cache_data_convert, update_cache_callback, *args, **kwarg
             llm_handler, cache_data_convert, update_cache_callback, *args, **kwargs
         )
     else:
-        llm_data = llm_handler(*args, **kwargs)
+        llm_data = time_cal(
+            llm_handler, func_name="llm_request", report_func=chat_cache.report.llm
+        )(*args, **kwargs)
 
     if cache_enable:
         try:
@@ -166,13 +189,23 @@ def update_cache_func(handled_llm_data, question=None):
                     question = pre_store_data
                 else:
                     question.content = pre_store_data
-                chat_cache.data_manager.save(
+                time_cal(
+                    chat_cache.data_manager.save,
+                    func_name="save",
+                    report_func=chat_cache.report.save,
+                )(
                     question,
                     handled_llm_data,
                     embedding_data,
                     extra_param=context.get("save_func", None),
                     session=session,
                 )
+                if (
+                    chat_cache.report.op_save.count > 0
+                    and chat_cache.report.op_save.count % chat_cache.config.auto_flush
+                    == 0
+                ):
+                    chat_cache.flush()
 
             llm_data = update_cache_callback(
                 llm_data, update_cache_func, *args, **kwargs

gptcache/adapter/api.py

@@ -17,6 +17,7 @@
     Cohere,
     Rwkv,
     PaddleNLP,
+    UForm,
 )
 from gptcache.embedding.base import BaseEmbedding
 from gptcache.manager import manager_factory
@@ -276,6 +277,8 @@ def _get_model(model_src, model_config=None):
         return Rwkv(**model_config)
     if model_src == "paddlenlp":
         return PaddleNLP(**model_config)
+    if model_src == "uform":
+        return UForm(**model_config)
 
 
 def _get_eval(strategy, kws=None):

gptcache/adapter/openai.py

@@ -5,6 +5,7 @@
 from io import BytesIO
 from typing import Iterator, Any, List
 
+from gptcache import cache
 from gptcache.adapter.adapter import adapt
 from gptcache.adapter.base import BaseCacheLLM
 from gptcache.manager.scalar_data.base import Answer, DataType
@@ -18,6 +19,7 @@
     get_image_from_openai_url,
     get_audio_text_from_openai_answer,
 )
+from gptcache.utils.token import token_counter
 
 import_openai()
 
@@ -80,10 +82,19 @@ def hook_openai_data(it):
 
     @classmethod
     def create(cls, *args, **kwargs):
+        chat_cache = kwargs.get("cache_obj", cache)
+        enable_token_counter = chat_cache.config.enable_token_counter
+
         def cache_data_convert(cache_data):
+            if enable_token_counter:
+                input_token = _num_tokens_from_messages(kwargs.get("messages"))
+                output_token = token_counter(cache_data)
+                saved_token = [input_token, output_token]
+            else:
+                saved_token = [0, 0]
             if kwargs.get("stream", False):
-                return _construct_stream_resp_from_cache(cache_data)
-            return _construct_resp_from_cache(cache_data)
+                return _construct_stream_resp_from_cache(cache_data, saved_token)
+            return _construct_resp_from_cache(cache_data, saved_token)
 
         kwargs = cls.fill_base_args(**kwargs)
         return adapt(
@@ -346,9 +357,10 @@ def create(cls, *args, **kwargs):
         return res
 
 
-def _construct_resp_from_cache(return_message):
+def _construct_resp_from_cache(return_message, saved_token):
     return {
         "gptcache": True,
+        "saved_token": saved_token,
         "choices": [
             {
                 "message": {"role": "assistant", "content": return_message},
@@ -362,7 +374,7 @@ def _construct_resp_from_cache(return_message):
     }
 
 
-def _construct_stream_resp_from_cache(return_message):
+def _construct_stream_resp_from_cache(return_message, saved_token):
     created = int(time.time())
     return [
         {
@@ -388,6 +400,7 @@ def _construct_stream_resp_from_cache(return_message):
             "choices": [{"delta": {}, "finish_reason": "stop", "index": 0}],
             "created": created,
             "object": "chat.completion.chunk",
+            "saved_token": saved_token,
         },
     ]
 
@@ -447,3 +460,19 @@ def _construct_audio_text_from_cache(return_text):
         "gptcache": True,
         "text": return_text,
     }
+
+
+def _num_tokens_from_messages(messages):
+    """Returns the number of tokens used by a list of messages."""
+    tokens_per_message = 3
+    tokens_per_name = 1
+
+    num_tokens = 0
+    for message in messages:
+        num_tokens += tokens_per_message
+        for key, value in message.items():
+            num_tokens += token_counter(value)
+            if key == "name":
+                num_tokens += tokens_per_name
+    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
+    return num_tokens

gptcache/config.py

@@ -15,6 +15,10 @@ class Config:
     :type prompts: Optional[List[str]]
     :param template: optional, if the request content will remove the template string and only keep the parameter value in the template
     :type template: Optional[str]
+    :param auto_flush: it will be automatically flushed every time xx pieces of data are added, default to 20
+    :type auto_flush: int
+    :param enable_token_counter: enable token counter, default to False
+    :type enable_token_counter: bool
 
     Example:
         .. code-block:: python
@@ -30,6 +34,8 @@ def __init__(
             similarity_threshold: float = 0.8,
             prompts: Optional[List[str]] = None,
             template: Optional[str] = None,
+            auto_flush: int = 20,
+            enable_token_counter: bool = True,
     ):
         if similarity_threshold < 0 or similarity_threshold > 1:
             raise CacheError(
@@ -39,3 +45,5 @@ def __init__(
         self.similarity_threshold = similarity_threshold
         self.prompts = prompts
         self.template = template
+        self.auto_flush = auto_flush
+        self.enable_token_counter = enable_token_counter

gptcache/report.py

@@ -2,46 +2,126 @@ class Report:
     """Get GPTCache report including time and counts for different operations."""
 
     def __init__(self):
-        self.embedding_all_time = 0
-        self.embedding_count = 0
-        self.search_all_time = 0
-        self.search_count = 0
+        self.op_pre = OpCounter()
+        self.op_embedding = OpCounter()
+        self.op_search = OpCounter()
+        self.op_data = OpCounter()
+        self.op_evaluation = OpCounter()
+        self.op_post = OpCounter()
+        self.op_llm = OpCounter()
+        self.op_save = OpCounter()
         self.hint_cache_count = 0
 
+    def pre(self, delta_time):
+        """Pre-process counts and time.
+
+        :param delta_time: additional runtime.
+        """
+        self.op_pre.total_time += delta_time
+        self.op_pre.count += 1
+
     def embedding(self, delta_time):
         """Embedding counts and time.
 
         :param delta_time: additional runtime.
         """
-        self.embedding_all_time += delta_time
-        self.embedding_count += 1
+        self.op_embedding.total_time += delta_time
+        self.op_embedding.count += 1
 
     def search(self, delta_time):
         """Search counts and time.
 
         :param delta_time: additional runtime.
         """
-        self.search_all_time += delta_time
-        self.search_count += 1
+        self.op_search.total_time += delta_time
+        self.op_search.count += 1
+
+    def data(self, delta_time):
+        """Get data counts and time.
+
+        :param delta_time: additional runtime.
+        """
+
+        self.op_data.total_time += delta_time
+        self.op_data.count += 1
+
+    def evaluation(self, delta_time):
+        """Evaluation counts and time.
+
+        :param delta_time: additional runtime.
+        """
+        self.op_evaluation.total_time += delta_time
+        self.op_evaluation.count += 1
+
+    def post(self, delta_time):
+        """Post-process counts and time.
+
+        :param delta_time: additional runtime.
+        """
+        self.op_post.total_time += delta_time
+        self.op_post.count += 1
+
+    def llm(self, delta_time):
+        """LLM counts and time.
+
+        :param delta_time: additional runtime.
+        """
+        self.op_llm.total_time += delta_time
+        self.op_llm.count += 1
+
+    def save(self, delta_time):
+        """Save counts and time.
+
+        :param delta_time: additional runtime.
+        """
+        self.op_save.total_time += delta_time
+        self.op_save.count += 1
+
+    def average_pre_time(self):
+        """Average pre-process time."""
+        return self.op_pre.average()
 
     def average_embedding_time(self):
         """Average embedding time."""
-        return round(
-            self.embedding_all_time / self.embedding_count
-            if self.embedding_count != 0
-            else 0,
-            4,
-        )
+        return self.op_embedding.average()
 
     def average_search_time(self):
         """Average search time."""
-        return round(
-            self.search_all_time / self.search_count
-            if self.embedding_count != 0
-            else 0,
-            4,
-        )
+        return self.op_search.average()
+
+    def average_data_time(self):
+        """Average data time."""
+        return self.op_data.average()
+
+    def average_evaluation_time(self):
+        """Average evaluation time."""
+        return self.op_evaluation.average()
+
+    def average_post_time(self):
+        """Average post-process time."""
+        return self.op_post.average()
+
+    def average_llm_time(self):
+        """Average LLM time."""
+        return self.op_llm.average()
+
+    def average_save_time(self):
+        """Average save time."""
+        return self.op_save.average()
 
     def hint_cache(self):
         """hint cache count."""
         self.hint_cache_count += 1
+
+
+class OpCounter:
+    """Operation counter."""
+
+    count = 0
+    """Operation count."""
+    total_time = 0
+    """Total time."""
+
+    def average(self):
+        """Average time."""
+        return round(self.total_time / self.count, 4) if self.count != 0 else 0