Преглед изворни кода

Update image_describer: cache singleton

StrayWarrior пре 3 недеља
родитељ
комит
acb43e935d
1 измењених фајлова са 8 додато и 1 уклоњено
  1. 8 1
      pqai_agent/toolkit/image_describer.py

+ 8 - 1
pqai_agent/toolkit/image_describer.py

@@ -1,4 +1,5 @@
 import diskcache
+import threading
 
 from pqai_agent import chat_service
 from pqai_agent.chat_service import VOLCENGINE_MODEL_DOUBAO_1_5_VISION_PRO
@@ -6,6 +7,9 @@ from pqai_agent.logging_service import logger
 from pqai_agent.toolkit.base import BaseToolkit
 from pqai_agent.toolkit.function_tool import FunctionTool
 
+# 不同实例间复用cache,但不是很好的实践
+_image_describer_caches = {}
+_cache_mutex = threading.Lock()
 
 class ImageDescriber(BaseToolkit):
     def __init__(self, cache_dir: str = None):
@@ -13,7 +17,10 @@ class ImageDescriber(BaseToolkit):
         self.llm_client = chat_service.OpenAICompatible.create_client(self.model)
         if not cache_dir:
             cache_dir = 'image_descriptions_cache'
-        self.cache = diskcache.Cache(cache_dir, size_limit=100*1024*1024)
+        if cache_dir not in _image_describer_caches:
+            with _cache_mutex:
+                _image_describer_caches[cache_dir] = diskcache.Cache(cache_dir, size_limit=100*1024*1024)
+        self.cache = _image_describer_caches[cache_dir]
         super().__init__()
 
     def analyse_image(self, image_url: str):