Caching, Latency & Cost

import hashlib
import json
import uuid
from collections.abc import Callable
from typing import Any

import redis
from langchain_core.documents import Document
from langchain_core.vectorstores import VectorStore


def acl_key(readable_acls: list[str], config_version: str) -> str:
    return hashlib.sha256(("|".join(sorted(readable_acls)) + "#" + config_version).encode()).hexdigest()[:16]


def pgvector_filter(acl: str) -> dict:
    return {"acl_key": {"$eq": acl}}


class SemanticAnswerCache:
    def __init__(self, store: VectorStore, index: redis.Redis, threshold: float = 0.93,
                 acl_filter: Callable[[str], Any] = pgvector_filter):
        self.store, self.index, self.threshold, self.acl_filter = store, index, threshold, acl_filter

    def lookup(self, question: str, acl: str) -> dict | None:
        hits = self.store.similarity_search_with_relevance_scores(question, k=1, filter=self.acl_filter(acl))
        if hits and hits[0][1] >= self.threshold:
            return json.loads(hits[0][0].metadata["payload"])
        return None

    def put(self, question: str, acl: str, answer: dict, cited_chunk_ids: list[str]) -> str:
        entry_id = str(uuid.uuid4())
        self.store.add_documents([Document(page_content=question, metadata={
            "acl_key": acl, "payload": json.dumps(answer)})], ids=[entry_id])
        for chunk_id in cited_chunk_ids:
            self.index.sadd(f"cites:{chunk_id}", entry_id)             # reverse index for invalidation
        return entry_id

    def invalidate_chunk(self, chunk_id: str) -> int:
        entries = list(self.index.smembers(f"cites:{chunk_id}"))
        if entries:
            self.store.delete(ids=entries)
        self.index.delete(f"cites:{chunk_id}")
        return len(entries)

from dataclasses import dataclass

from langchain_core.messages import AIMessage


@dataclass(frozen=True)
class Price:
    input_per_mtok: float          # USD per million input tokens
    output_per_mtok: float
    cached_input_per_mtok: float   # USD per million cache-read input tokens


def request_cost(messages: list[AIMessage], price: Price) -> dict:
    """Sum usage over every model call of a request (an agent makes several) and price it."""
    inp = out = cached = 0
    for m in messages:
        usage = m.usage_metadata or {}
        inp += usage.get("input_tokens", 0)
        out += usage.get("output_tokens", 0)
        cached += (usage.get("input_token_details") or {}).get("cache_read", 0)
    cost = ((inp - cached) * price.input_per_mtok + cached * price.cached_input_per_mtok
            + out * price.output_per_mtok) / 1_000_000
    return {"inputTokens": inp, "outputTokens": out, "cachedInputTokens": cached, "costUsd": round(cost, 6)}

package com.example.docsassistant.rag.cost;

import java.util.Map;

import io.micrometer.core.instrument.MeterRegistry;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.core.Ordered;

public class CostAccountingAdvisor implements CallAdvisor {

    public record Price(double inputPerMtok, double outputPerMtok, double cachedInputPerMtok) {
    }

    private final MeterRegistry meters;
    private final Map<String, Price> prices;       // by model name, from configuration

    public CostAccountingAdvisor(MeterRegistry meters, Map<String, Price> prices) {
        this.meters = meters;
        this.prices = prices;
    }

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
        ChatClientResponse response = chain.nextCall(request);
        if (response.chatResponse() == null) {
            return response;
        }
        var metadata = response.chatResponse().getMetadata();
        Usage usage = metadata.getUsage();
        Price price = prices.get(metadata.getModel());
        if (usage != null && price != null) {
            long input = usage.getPromptTokens() == null ? 0 : usage.getPromptTokens();
            long output = usage.getCompletionTokens() == null ? 0 : usage.getCompletionTokens();
            long cached = usage.getCacheReadInputTokens() == null ? 0 : usage.getCacheReadInputTokens();
            double usd = ((input - cached) * price.inputPerMtok() + cached * price.cachedInputPerMtok()
                    + output * price.outputPerMtok()) / 1_000_000.0;
            meters.counter("docsassistant.llm.cost.usd", "model", metadata.getModel()).increment(usd);
            response.context().put("costUsd", usd);
        }
        return response;
    }

    @Override
    public String getName() {
        return "costAccounting";
    }

    @Override
    public int getOrder() {
        return Ordered.LOWEST_PRECEDENCE - 100;
    }
}

package com.example.docsassistant.l4j.cost;

import java.util.concurrent.atomic.DoubleAdder;

import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.model.chat.listener.ChatModelResponseContext;
import dev.langchain4j.model.output.TokenUsage;

public class CostListener implements ChatModelListener {

    private final double inputPerMtok;
    private final double outputPerMtok;
    private final DoubleAdder totalUsd = new DoubleAdder();

    public CostListener(double inputPerMtok, double outputPerMtok) {
        this.inputPerMtok = inputPerMtok;
        this.outputPerMtok = outputPerMtok;
    }

    @Override
    public void onResponse(ChatModelResponseContext context) {
        TokenUsage usage = context.chatResponse().tokenUsage();
        if (usage == null) {
            return;
        }
        int in = usage.inputTokenCount() == null ? 0 : usage.inputTokenCount();
        int out = usage.outputTokenCount() == null ? 0 : usage.outputTokenCount();
        double usd = (in * inputPerMtok + out * outputPerMtok) / 1_000_000.0;
        totalUsd.add(usd);
        context.attributes().put("costUsd", usd);          // visible to later listeners of the same call
    }

    public double totalUsd() {
        return totalUsd.sum();
    }
}