cache¶

async def before_model_callback(
    self,
    callback_context: CallbackContext,
    llm_request: LlmRequest,
) -> Optional[LlmResponse]:
  """Check cache before making LLM call.

  Args:
      callback_context: The callback context with session info.
      llm_request: The LLM request being made.

  Returns:
      LlmResponse if cache hit, None to proceed with LLM call.
  """
  if self._config.first_message_only and not self._is_first_message(
      callback_context
  ):
    logger.debug("Not first message, skipping cache check")
    return None

  prompt = self._extract_prompt(llm_request)
  if not prompt:
    logger.debug("No prompt found in request")
    return None

  cache_key = self._build_cache_key(prompt, callback_context)
  cache_entry = await self._provider.check(cache_key)

  if cache_entry:
    logger.info("Cache hit for prompt: %s", prompt[:50])
    return LlmResponse(
        content=types.Content(
            role="model",
            parts=[types.Part(text=cache_entry.response)],
        )
    )

  logger.debug("Cache miss for prompt: %s", prompt[:50])
  # Store prompt for after_model_callback
  session_key = self._get_session_key(callback_context)
  self._pending_prompts[session_key] = cache_key
  return None

async def after_model_callback(
    self,
    callback_context: CallbackContext,
    llm_response: LlmResponse,
) -> Optional[LlmResponse]:
  """Store response in cache after LLM call.

  Args:
      callback_context: The callback context with session info.
      llm_response: The LLM response to cache.

  Returns:
      None to pass through the original response.
  """
  session_key = self._get_session_key(callback_context)
  cache_key = self._pending_prompts.pop(session_key, None)

  if not cache_key:
    logger.debug("No pending prompt for session, skipping cache store")
    return None

  # Don't cache error responses
  if llm_response.error_message:
    logger.debug("Response has error, skipping cache store")
    return None

  # Don't cache function call responses
  if llm_response.content and llm_response.content.parts:
    for part in llm_response.content.parts:
      if hasattr(part, "function_call") and part.function_call:
        logger.debug("Response is function call, skipping cache store")
        return None

  response_text = self._extract_response_text(llm_response)
  if not response_text:
    logger.debug("No text in response, skipping cache store")
    return None

  await self._provider.store(cache_key, response_text)
  logger.info("Cached response for prompt: %s", cache_key[:50])
  return None

@abstractmethod
async def check(self, prompt: str, **kwargs: Any) -> Optional[CacheEntry]:
  """Check if a semantically similar prompt exists in the cache.

  Args:
    prompt: The prompt to look up.
    **kwargs: Provider-specific options.

  Returns:
    A CacheEntry on a hit (with entry_id populated when the backend
    provides one), None on a miss.
  """
  pass

@abstractmethod
async def store(
    self,
    prompt: str,
    response: str,
    metadata: Optional[dict[str, Any]] = None,
    **kwargs: Any,
) -> Optional[str]:
  """Store a prompt-response pair in the cache.

  Args:
    prompt: The prompt text.
    response: The response text.
    metadata: Optional metadata to store alongside the entry.
    **kwargs: Provider-specific options.

  Returns:
    The identifier of the newly written entry when the backend
    provides one, otherwise None. The identifier can be passed to
    delete_by_id() to retire exactly this entry.
  """
  pass

@abstractmethod
async def delete_by_id(self, entry_id: str, **kwargs: Any) -> None:
  """Delete exactly one entry from the cache by its identifier.

  Unlike clear(), which removes every entry in the cache, this targets
  a single entry, identified by the entry_id obtained from check() or
  store(). Deleting an identifier that no longer exists is a no-op.

  Args:
    entry_id: The identifier of the entry to delete, as returned by
      check() (CacheEntry.entry_id) or store() on this provider.
    **kwargs: Provider-specific options.
  """
  pass

@abstractmethod
async def clear(self, **kwargs: Any) -> None:
  """Clear all entries from the cache.

  For removing a single entry, use delete_by_id() instead.
  """
  pass

@abstractmethod
async def close(self) -> None:
  """Close the cache provider and release resources."""
  pass

async def check(self, prompt: str, **kwargs: Any) -> Optional[CacheEntry]:
  """Check for a semantically similar prompt in LangCache.

  Args:
    prompt: The prompt to check.
    **kwargs: Additional keyword arguments (e.g., distance_threshold).

  Returns:
    A CacheEntry if a cache hit is found, None otherwise. entry_id is
    the LangCache entry ID of the matched entry and can be passed to
    delete_by_id().
  """
  distance_threshold = kwargs.get(
      "distance_threshold", self._config.distance_threshold
  )
  acheck_kwargs: dict[str, Any] = {"prompt": prompt}
  if distance_threshold is not None:
    acheck_kwargs["distance_threshold"] = distance_threshold
  result = await self._cache.acheck(**acheck_kwargs)
  if result:
    logger.debug("LangCache hit for prompt: %s", prompt[:50])
    return CacheEntry(
        prompt=prompt,
        response=result[0]["response"],
        distance=result[0].get("vector_distance"),
        metadata=result[0].get("metadata"),
        entry_id=result[0].get("entry_id"),
    )
  logger.debug("LangCache miss for prompt: %s", prompt[:50])
  return None

async def store(
    self,
    prompt: str,
    response: str,
    metadata: Optional[dict[str, Any]] = None,
    **kwargs: Any,
) -> Optional[str]:
  """Store a prompt-response pair in LangCache.

  Args:
    prompt: The prompt text.
    response: The response text.
    metadata: Optional metadata to store alongside the entry.
    **kwargs: Additional keyword arguments (e.g., ttl).

  Returns:
    The LangCache entry ID of the newly written entry. It can be
    passed to delete_by_id() to retire exactly this entry.
  """
  astore_kwargs: dict[str, Any] = {
      "prompt": prompt,
      "response": response,
      "metadata": metadata,
  }
  # Only override TTL when an explicit, non-None value is provided.
  ttl = kwargs.get("ttl")
  if ttl is not None:
    astore_kwargs["ttl"] = ttl
  entry_id = await self._cache.astore(**astore_kwargs)
  logger.debug("LangCache stored response for prompt: %s", prompt[:50])
  return entry_id or None

async def delete_by_id(self, entry_id: str, **kwargs: Any) -> None:
  """Delete exactly one LangCache entry by its entry ID.

  Unlike clear(), which removes every entry, this deletes only the
  entry identified by the LangCache entry ID obtained from check()
  or store().

  Args:
    entry_id: The LangCache entry ID of the entry to delete.
    **kwargs: Additional keyword arguments (unused).
  """
  await self._cache.adelete_by_id(entry_id)
  logger.debug("LangCache deleted entry: %s", entry_id)

async def clear(self, **kwargs: Any) -> None:
  """Clear all entries from the LangCache.

  For removing a single entry, use delete_by_id() instead.
  """
  await self._cache.aclear()
  logger.info("LangCache cleared")

async def close(self) -> None:
  """Close the LangCache provider and release resources."""
  await self._cache.adisconnect()
  logger.debug("LangCache provider closed")

async def check(self, prompt: str, **kwargs: Any) -> Optional[CacheEntry]:
  """Check for a semantically similar prompt in the cache.

  Args:
    prompt: The prompt to look up.
    **kwargs: Additional keyword arguments (unused).

  Returns:
    A CacheEntry on a hit, None on a miss. entry_id is the RedisVL entry
    ID of the matched entry and can be passed to delete_by_id().
  """
  result = await asyncio.to_thread(self._cache.check, prompt=prompt)
  if result:
    logger.debug("Cache hit for prompt: %s", prompt[:50])
    return CacheEntry(
        prompt=prompt,
        response=result[0]["response"],
        distance=result[0].get("vector_distance"),
        entry_id=result[0].get("entry_id"),
    )
  logger.debug("Cache miss for prompt: %s", prompt[:50])
  return None

async def store(
    self,
    prompt: str,
    response: str,
    metadata: Optional[dict[str, Any]] = None,
    **kwargs: Any,
) -> Optional[str]:
  """Store a prompt-response pair in the cache.

  Args:
    prompt: The prompt text.
    response: The response text.
    metadata: Optional metadata (unused by this provider).
    **kwargs: Additional keyword arguments (unused).

  Returns:
    The RedisVL entry ID of the newly written entry. It can be passed to
    delete_by_id() to retire exactly this entry.
  """
  key: str = await asyncio.to_thread(
      self._cache.store, prompt=prompt, response=response
  )
  key_prefix = f"{self._config.name}:"
  entry_id = key.removeprefix(key_prefix)
  logger.debug("Stored response for prompt: %s", prompt[:50])
  return entry_id

async def delete_by_id(self, entry_id: str, **kwargs: Any) -> None:
  """Delete exactly one cache entry by its RedisVL entry ID.

  Unlike clear(), which removes every entry, this drops only the entry
  identified by the RedisVL entry ID obtained from check() or store().

  Args:
    entry_id: The RedisVL entry ID of the entry to delete.
    **kwargs: Additional keyword arguments (unused).
  """
  await asyncio.to_thread(self._cache.drop, ids=[entry_id])
  logger.debug("Dropped cache entry: %s", entry_id)

async def clear(self, **kwargs: Any) -> None:
  """Clear all entries from the cache.

  For removing a single entry, use delete_by_id() instead.
  """
  await asyncio.to_thread(self._cache.clear)
  logger.info("Cache cleared")

async def close(self) -> None:
  """Close the cache provider and release resources."""
  if hasattr(self._cache, "_index") and hasattr(self._cache._index, "client"):
    await asyncio.to_thread(self._cache._index.client.close)
  logger.debug("RedisVL cache provider closed")

async def before_tool_callback(
    self,
    tool: BaseTool,
    args: Dict[str, Any],
    tool_context: ToolContext,
) -> Optional[Dict[str, Any]]:
  """Check cache before executing tool.

  Args:
      tool: The tool being called.
      args: The arguments for the tool call.
      tool_context: The tool context with session info.

  Returns:
      Dict if cache hit (skips tool execution), None to proceed.
  """
  tool_name = tool.name
  if not self._should_cache_tool(tool_name):
    logger.debug("Tool %s not in cache list, skipping", tool_name)
    return None

  cache_key = self._build_cache_key(tool_name, args, tool_context)
  cache_entry = await self._provider.check(cache_key)

  if cache_entry:
    logger.info("Cache hit for tool: %s", tool_name)
    try:
      return json.loads(cache_entry.response)  # type: ignore[no-any-return]
    except (json.JSONDecodeError, TypeError):
      return {"result": cache_entry.response}

  logger.debug("Cache miss for tool: %s", tool_name)
  # Store cache key for after_tool_callback
  session_key = self._get_session_key(tool_context)
  self._pending_calls[session_key] = cache_key
  return None

async def after_tool_callback(
    self,
    tool: BaseTool,
    args: Dict[str, Any],
    tool_context: ToolContext,
    tool_response: Dict[str, Any],
) -> Optional[Dict[str, Any]]:
  """Store tool result in cache after execution.

  Args:
      tool: The tool that was called.
      args: The arguments used for the tool call.
      tool_context: The tool context with session info.
      tool_response: The result from the tool execution.

  Returns:
      None to pass through the original response.
  """
  session_key = self._get_session_key(tool_context)
  cache_key = self._pending_calls.pop(session_key, None)

  if not cache_key:
    logger.debug("No pending cache key for tool, skipping store")
    return None

  # Serialize the tool response
  try:
    response_str = json.dumps(tool_response, default=str)
  except (TypeError, ValueError):
    response_str = str(tool_response)

  await self._provider.store(cache_key, response_str)
  logger.info("Cached result for tool: %s", tool.name)
  return None