repo_description: A LLM semantic caching system aiming to enhance user experience by reducing response time via cached query-result pairs.
repo_topics: llm semantic-cache
code: # -*- coding: utf-8 -*-
  from contextlib import asynccontextmanager
  import uvicorn
  import json
  import asyncio
  from fastapi import FastAPI, WebSocket
  from starlette.websockets import WebSocketDisconnect
  from modelcache.cache import Cache
  from modelcache.embedding import EmbeddingModel
  
  @asynccontextmanager
  async def lifespan(app: FastAPI):
      global cache
      cache, _ = await Cache.init(
          sql_storage="mysql",
          vector_storage="milvus",
          embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
          embedding_workers_num=8
      )
      yield
  
  app = FastAPI(lifespan=lifespan)
  cache: Cache = None
  
  @app.websocket("/modelcache")
  async def user_backend(websocket: WebSocket):
      await websocket.accept()
      try:
          while True:
              data = await websocket.receive_text()
              asyncio.create_task(handle_message(websocket, data))
      except WebSocketDisconnect as e:
          print(e)
  
  
  async def handle_message(websocket,message):
      try:
          param_dict = json.loads(message)
      except Exception:
          await websocket.send_json({"errorCode": 400, "errorDesc": "bad request", "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''})
          return
  
      request_id = param_dict.get("requestId")
      request_payload = param_dict.get("payload")
      if not request_id or not request_payload:
          await websocket.send_json({"errorCode": 400, "errorDesc": "bad request", "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''})
          return
      try:
          result = await cache.handle_request(request_payload)
          await websocket.send_json({"requestId": request_id,"result": result})
      except Exception as e:
          error_result = {"errorCode": 500, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
          cache.save_query_resp(error_result, model='', query='', delta_time=0)
          await websocket.send_json(error_result)
  
  
  if __name__ == '__main__':
      uvicorn.run(app, host='0.0.0.0', port=5000, loop="asyncio", http="httptools")
  