[Backend] API 사용량 추정

사용된 토큰 추정

OpenAI

OpenAI의 경우 LangChain 라이브러리에서, 사용된 토큰 수와 그로 인해 발생한 비용을 계산하고 반환해줍니다.
그렇기 때문에 OpenAI 사용량은 Anthropic과 별개로 LangChain을 활용하여 정리합니다.

CRUD 코드

from sqlalchemy import func
from models.project import User
from models.api import ApiKey

def update_usage(db : Session, user_email : str, provider : str, usage : int):
    user = db.query(User).filter(User.email == user_email).first()
    id = user.id
    api = db.query(ApiKey).filter(
        ApiKey.user_id == id,
        ApiKey.status == 'Active',
        func.lower(ApiKey.provider_name) == provider.lower()
    ).first()
    total_usage = api.usage_count + usage
    if api.usage_limit > total_usage :
        api.usage_count = total_usage
        db.commit()
        db.refresh(user)
    return total_usage

데이터베이스 내의 API 키 테이블에서 usage_count 컬럼 값을 업데이트하는 코드입니다.
만약 한 사용자에게 같은 Provider의 API 키가 여러 개 존재할 경우,
가장 첫 번째로 DB에 등록된 API 키를 사용하도록 설정했습니다.

이후에는 가장 사용 가능한 토큰 수가 많이 남은 API 키를 가져오고나,
토큰 한도가 가장 큰 API를 가져오는 방식으로 수정해야 할 것 같습니다.

QA Chain 변경 사항

from langchain_community.callbacks.manager import get_openai_callback
from crud.user import update_usage
from langchain_core.output_parsers import StrOutputParser

if provider == "openai":
    with get_openai_callback() as cb:
        response_text = chain.invoke({"history": formatted_history, "input": conversation})
        update_usage(db = db, user_email = user_email, provider = provider, usage = cb.total_tokens)

함수에 입력된 Provider의 이름이 OpenAI일 경우,
LangChain의 Callback 패키지를 활용하여 해당 질문과 답변으로 발생한 총 토큰 수와 비용 정보를 얻습니다.
그리고 이것을 그대로 Update Usage 함수로 DB에 저장합니다.

적용 후 UI

Anthropic

Anthropic의 경우 OpenAI와 달리, LangChain을 사용해도 정확한 토큰 수와 비용을 추정해주지 않습니다.
그렇기 때문에 토큰 수에 따른 비용을 직접 추정해야 합니다.

토큰 수 추정

def count_tokens(text: str):
    return int(len(text) / 4)

현재는 간단하게 4글자 당 토큰 하나를 사용했다고 추정 중입니다.
이 부분은 이후 Claude의 공식 문서를 확인하고, 정확한 토큰 계산식으로 변경할 필요가 있습니다.

토큰에 따른 비용 추정

def estimate_claude_cost(model: str, prompt_tokens: int, completion_tokens: int):
    prices = {
        "claude-3-haiku-20240307": (0.00025, 0.00125),
        "claude-3-sonnet-20240229": (0.003, 0.015),
        "claude-3-opus-20240229": (0.015, 0.075),
    }

    prompt_price, completion_price = prices.get(model, (0.0, 0.0))
    cost = (prompt_tokens * prompt_price + completion_tokens * completion_price) / 1000

    return {
        "prompt": prompt_tokens,
        "completion": completion_tokens,
        "total": prompt_tokens + completion_tokens,
        "cost": round(cost, 6)
    }

Claude는 어떤 모델을 사용 하는지에 따라 비용의 차이가 크기 때문에
각 모델에 대해 각각의 계산을 적용합니다.

QA Chain에 적용

elif provider == "anthropic":
    prompt_tokens = count_tokens(formatted_history)
    response_text = chain.invoke({"history": formatted_history, "input": conversation})
    completion_tokens = count_tokens(response_text)
    cost_data = estimate_claude_cost(model, prompt_tokens, completion_tokens)
    update_usage(db=db, user_email=user_email, provider=provider, usage=cost_data['total'])