Context Window

import tiktoken
from anthropic import Anthropic

# 토큰 수 계산 (OpenAI 모델용 tiktoken 예시)
def count_tokens_openai(text: str, model: str = "gpt-4") -> int:
    encoding = tiktoken.encoding_for_model(model)
    return len(encoding.encode(text))

# Claude 토큰 수 계산 (Anthropic API 사용)
def count_tokens_claude(text: str) -> int:
    client = Anthropic()
    # count_tokens API는 Claude 모델의 정확한 토큰 수 반환
    result = client.count_tokens(text)
    return result.tokens

# 컨텍스트 윈도우 관리 예시
class ContextManager:
    def __init__(self, max_tokens: int = 128000, reserve_output: int = 4096):
        self.max_tokens = max_tokens
        self.reserve_output = reserve_output  # 출력용 토큰 예약
        self.available = max_tokens - reserve_output

    def fit_messages(self, messages: list[dict], system_prompt: str = "") -> list[dict]:
        """컨텍스트 윈도우에 맞게 메시지 자르기 (최근 메시지 우선)"""
        system_tokens = count_tokens_openai(system_prompt) if system_prompt else 0
        available = self.available - system_tokens

        fitted = []
        current_tokens = 0

        # 최신 메시지부터 역순으로 추가
        for msg in reversed(messages):
            msg_tokens = count_tokens_openai(msg["content"])
            if current_tokens + msg_tokens <= available:
                fitted.insert(0, msg)
                current_tokens += msg_tokens
            else:
                break

        print(f"사용 토큰: {current_tokens + system_tokens}/{self.max_tokens}")
        print(f"포함된 메시지: {len(fitted)}/{len(messages)}")
        return fitted

# 사용 예시
manager = ContextManager(max_tokens=128000)
chat_history = [{"role": "user", "content": "..." }] * 100  # 긴 대화

fitted_messages = manager.fit_messages(chat_history, system_prompt="You are helpful.")
# 결과: 컨텍스트에 맞게 오래된 메시지가 제거됨