cache.py 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142
  1. import logging
  2. import os # noqa: F401
  3. from typing import Any
  4. from gptcache import cache # noqa: F401
  5. from gptcache.adapter.adapter import adapt # noqa: F401
  6. from gptcache.config import Config # noqa: F401
  7. from gptcache.manager import get_data_manager
  8. from gptcache.manager.scalar_data.base import Answer
  9. from gptcache.manager.scalar_data.base import DataType as CacheDataType
  10. from gptcache.session import Session
  11. from gptcache.similarity_evaluation.distance import \
  12. SearchDistanceEvaluation # noqa: F401
  13. from gptcache.similarity_evaluation.exact_match import \
  14. ExactMatchEvaluation # noqa: F401
  15. def gptcache_pre_function(data: dict[str, Any], **params: dict[str, Any]):
  16. return data["input_query"]
  17. def gptcache_data_manager(vector_dimension):
  18. return get_data_manager(cache_base="sqlite", vector_base="chromadb", max_size=1000, eviction="LRU")
  19. def gptcache_data_convert(cache_data):
  20. logging.info("[Cache] Cache hit, returning cache data...")
  21. return cache_data
  22. def gptcache_update_cache_callback(llm_data, update_cache_func, *args, **kwargs):
  23. logging.info("[Cache] Cache missed, updating cache...")
  24. update_cache_func(Answer(llm_data, CacheDataType.STR))
  25. return llm_data
  26. def _gptcache_session_hit_func(cur_session_id: str, cache_session_ids: list, cache_questions: list, cache_answer: str):
  27. return cur_session_id in cache_session_ids
  28. def get_gptcache_session(session_id: str):
  29. return Session(name=session_id, check_hit_func=_gptcache_session_hit_func)