Память сводки и усечение с учётом токенов

from langchain.memory import ConversationSummaryMemory
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model='gpt-4o-mini')

memory = ConversationSummaryMemory(
    llm=llm,
    return_messages=True
)

# Add messages manually
memory.save_context(
    {'input': 'My name is Alice and I am building a RAG system.'},
    {'output': 'Great, I can help you build a RAG system, Alice.'}
)

print(memory.load_memory_variables({}))

from langchain.memory import ConversationSummaryBufferMemory
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model='gpt-4o-mini')

memory = ConversationSummaryBufferMemory(
    llm=llm,
    max_token_limit=500,  # Summarize when older turns exceed 500 tokens
    return_messages=True
)

# As the conversation grows, old turns get summarized automatically
print('Moving summary:', memory.moving_summary_buffer)

import tiktoken

def count_tokens(messages: list, model: str = 'gpt-4o') -> int:
    encoding = tiktoken.encoding_for_model(model)
    total = 0
    for msg in messages:
        # Each message has overhead tokens for role framing
        total += 4
        total += len(encoding.encode(msg.get('content', '')))
    total += 2  # Reply primer
    return total

messages = [
    {'role': 'user', 'content': 'Explain RAG to me.'},
    {'role': 'assistant', 'content': 'RAG stands for Retrieval-Augmented Generation...'}
]
print('Token count:', count_tokens(messages))

def truncate_to_budget(messages: list, budget: int, model: str = 'gpt-4o') -> list:
    '''Remove oldest non-system messages until under budget.'''
    import tiktoken
    encoding = tiktoken.encoding_for_model(model)

    def token_count(msgs):
        total = 2
        for m in msgs:
            total += 4 + len(encoding.encode(m.get('content', '')))
        return total

    result = list(messages)
    while token_count(result) > budget and len(result) > 1:
        # Never remove the system prompt at index 0
        if result[0]['role'] == 'system':
            del result[1]
        else:
            del result[0]
    return result

from openai import OpenAI

client = OpenAI()

def summarize_history(old_summary: str, new_turns: list) -> str:
    turns_text = '\n'.join(
        f'{m["role"].capitalize()}: {m["content"]}' for m in new_turns
    )
    prompt = f'''Current summary:\n{old_summary}\n\nNew conversation turns:\n{turns_text}\n\nWrite an updated summary that captures key facts, user preferences, and open questions.'''

    resp = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=[{'role': 'user', 'content': prompt}]
    )
    return resp.choices[0].message.content

def build_messages_with_summary(summary: str, recent_turns: list, user_input: str) -> list:
    system_content = (
        'You are a helpful AI assistant.'
        + ('\n\n[Conversation summary]\n' + summary if summary else '')
    )

    messages = [{'role': 'system', 'content': system_content}]
    messages.extend(recent_turns)
    messages.append({'role': 'user', 'content': user_input})
    return messages

class SummaryBufferChat:
    def __init__(self, max_tokens=2000):
        self.summary = ''
        self.recent_turns = []
        self.max_tokens = max_tokens

    def chat(self, user_message: str) -> str:
        from openai import OpenAI
        client = OpenAI()

        messages = build_messages_with_summary(self.summary, self.recent_turns, user_message)
        resp = client.chat.completions.create(model='gpt-4o', messages=messages)
        reply = resp.choices[0].message.content

        self.recent_turns.append({'role': 'user', 'content': user_message})
        self.recent_turns.append({'role': 'assistant', 'content': reply})

        if count_tokens(self.recent_turns) > self.max_tokens:
            to_summarize = self.recent_turns[:-2]
            self.recent_turns = self.recent_turns[-2:]
            self.summary = summarize_history(self.summary, to_summarize)

        return reply