토큰 윈도우

# 컨텍스트 윈도우 관리
# pip install tiktoken openai

import tiktoken

class ContextWindowManager:
    def __init__(self, model="gpt-4", max_tokens=128000, reserve_for_response=4096):
        self.encoder = tiktoken.encoding_for_model(model)
        self.max_tokens = max_tokens
        self.reserve_for_response = reserve_for_response
        self.available_tokens = max_tokens - reserve_for_response

    def count_tokens(self, text):
        """텍스트의 토큰 수 계산"""
        return len(self.encoder.encode(text))

    def count_messages_tokens(self, messages):
        """OpenAI 메시지 형식의 토큰 수 계산"""
        total = 0
        for msg in messages:
            total += 4  # 메시지 오버헤드
            for key, value in msg.items():
                total += self.count_tokens(str(value))
        total += 2  # 응답 시작 오버헤드
        return total

    def fit_messages(self, system_prompt, messages, max_history=10):
        """컨텍스트 윈도우에 맞게 메시지 조정"""
        system_tokens = self.count_tokens(system_prompt)
        available = self.available_tokens - system_tokens

        # 최신 메시지부터 역순으로 추가
        fitted_messages = []
        current_tokens = 0

        for msg in reversed(messages[-max_history:]):
            msg_tokens = self.count_tokens(msg.get("content", "")) + 4
            if current_tokens + msg_tokens > available:
                break
            fitted_messages.insert(0, msg)
            current_tokens += msg_tokens

        return fitted_messages, current_tokens + system_tokens

# 사용 예시
manager = ContextWindowManager(model="gpt-4", max_tokens=128000)

system_prompt = "당신은 친절한 AI 어시스턴트입니다."
messages = [
    {"role": "user", "content": "안녕하세요"},
    {"role": "assistant", "content": "안녕하세요! 무엇을 도와드릴까요?"},
    {"role": "user", "content": "오늘 날씨가 어때요?"},
    {"role": "assistant", "content": "죄송합니다, 실시간 날씨 정보는 제공하지 못합니다."},
]

fitted, total_tokens = manager.fit_messages(system_prompt, messages)
print(f"포함된 메시지: {len(fitted)}개")
print(f"총 토큰 수: {total_tokens}")
print(f"남은 토큰: {manager.available_tokens - total_tokens}")

# 슬라이딩 윈도우 대화 관리
# pip install tiktoken

import tiktoken
from collections import deque

class SlidingWindowChat:
    def __init__(self, model="gpt-4", window_size=8000):
        self.encoder = tiktoken.encoding_for_model(model)
        self.window_size = window_size
        self.messages = deque()
        self.current_tokens = 0

    def add_message(self, role, content):
        """새 메시지 추가 (윈도우 초과 시 오래된 메시지 제거)"""
        msg_tokens = len(self.encoder.encode(content)) + 4  # 오버헤드

        # 윈도우 초과 시 오래된 메시지 제거
        while self.current_tokens + msg_tokens > self.window_size and self.messages:
            old_msg = self.messages.popleft()
            old_tokens = len(self.encoder.encode(old_msg["content"])) + 4
            self.current_tokens -= old_tokens

        # 새 메시지 추가
        self.messages.append({"role": role, "content": content})
        self.current_tokens += msg_tokens

    def get_messages(self):
        """현재 윈도우의 메시지 반환"""
        return list(self.messages)

    def get_stats(self):
        """현재 상태 정보"""
        return {
            "message_count": len(self.messages),
            "current_tokens": self.current_tokens,
            "available_tokens": self.window_size - self.current_tokens
        }

# 사용 예시
chat = SlidingWindowChat(window_size=1000)

# 대화 시뮬레이션
conversations = [
    ("user", "인공지능에 대해 설명해주세요."),
    ("assistant", "인공지능(AI)은 인간의 지능을 모방하여 학습, 추론, 문제 해결 등의 작업을 수행하는 컴퓨터 시스템입니다..."),
    ("user", "머신러닝과의 차이점은요?"),
    ("assistant", "머신러닝은 AI의 하위 분야로, 명시적 프로그래밍 없이 데이터로부터 학습하는 방식입니다..."),
]

for role, content in conversations:
    chat.add_message(role, content)
    stats = chat.get_stats()
    print(f"[{role}] 추가 후: {stats['message_count']}개 메시지, {stats['current_tokens']} 토큰")

print("\n최종 메시지:")
for msg in chat.get_messages():
    print(f"  {msg['role']}: {msg['content'][:50]}...")