Managing Context Effectively with the Model Context Protocol

In this tutorial, we guide you through a practical implementation of the Model Context Protocol (MCP) by building a ModelContextManager

```python

import torch

import numpy as np

import typing

from dataclasses import dataclass

import time

import gc

from tqdm.notebook import tqdm

from sentence_transformers import SentenceTransformer

from transformers import GPT2Tokenizer, FLAN_T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM

import math

MAX_TOKENS = 8000

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

NUM_CHUNKS = 50

CHUNK_SIZE = 100

RELEVANCE_THRESHOLD = 0.1

IMPORTANCE_FACTOR = 1.0

RECENCY_FACTOR = 0.5

VISUALIZE_CONTEXT = True

BATCH_SIZE = 32

class ContextChunk(typing.NamedTuple):

content: str

embedding: np.array

importance: float = 1.0

timestamp: float = time.time()

metadata: dict = None

def __post_init__(self):

if self.metadata is None:

self.metadata = {}

class ModelContextManager:

def __init__(self, context_chunks:typing.List[ContextChunk]=None, max_tokens:int=MAX_TOKENS, token_limit:int=0, gpt2_tokenizer:GPT2Tokenizer=None):

self.max_tokens = max_tokens

self.token_limit = token_limit

self.context_chunks = context_chunks or []

self.used_tokens = 0

self.last_chunk_index = 0

self.total_chunks = 0

if gpt2_tokenizer is None:

self.gpt2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

else:

self.gpt2_tokenizer = gpt2_tokenizer

self.sentence_transformer = SentenceTransformer('all-mpnet-base-v2')

def add_chunk(self, chunk_text:str, importance:float=1.0):

encoded_input = self.gpt2_tokenizer(chunk_text, return_tensors='pt')

self.used_tokens += int(encoded_input[0].shape[1])

chunk_embedding = self.sentence_transformer.encode(chunk_text, batch_size=BATCH_SIZE)

new_chunk = ContextChunk(content=chunk_text, embedding=chunk_embedding, importance=importance)

self.context_chunks.append(new_chunk)

self.last_chunk_index += 1

self.total_chunks += 1

print(f"Added chunk with {int(encoded_input[0].shape[1])} tokens and importance {importance}. Total used tokens: {self.used_tokens}, total chunks: {self.total_chunks}")

def optimize_context_window(self, query:str, min_chunks:int=3):

if len(self.context_chunks) <= min_chunks:

return []

query_embedding = self.sentence_transformer.encode(query, batch_size=BATCH_SIZE)

chunks_to_keep = []

remaining_tokens = self.max_tokens - self.used_tokens

if remaining_tokens < 0:

print("Warning: token limit exceeded by %s tokens" % -remaining_tokens)

for i in range(min_chunks, len(self.context_chunks) - 1, -1):

chunk = self.context_chunks[i]

if i == len(self.context_chunks) - 1:

chunks_to_keep.append(i)

continue

chunk_importance = chunk.importance * IMPORTANCE_FACTOR

chunk_recency = (time.time() - chunk.timestamp) * RECENCY_FACTOR

relevant_scores = np.array([cosine_similarity(chunk.embedding, x) for x in query_embedding])

max_relevant_score = np.max(relevant_scores)

total_score = chunk_importance + chunk_recency + max_relevant_score

if total_score >= RELEVANCE_THRESHOLD:

encoded_input = self.gpt2_tokenizer(chunk.content, return_tensors='pt')

chunk_token_count = int(encoded_input[0].shape[1])

if remaining_tokens >= chunk_token_count:

chunks_to_keep.append(i)

remaining_