import torchimport numpy as npfrom typing import List, Dict, Any, Optional, Union, Tuplefrom dataclasses import dataclassimport timeimport gcfrom tqdm.notebook import tqdm

@dataclassclass ContextChunk:    """A chunk of text with metadata for the Model Context Protocol."""    text: str    embedding: Optional[torch.Tensor] = None    importance: float = 1.0    timestamp: float = 0.0    metadata: Dict[str, Any] = None       def __post_init__(self):        if self.metadata is None:            self.metadata = {}        if self.timestamp == 0.0:            self.timestamp = time.time()

class ModelContextManager:    """    Manager for implementing Model Context Protocol in LLMs on Google Colab.    Handles context window optimization, token management, and relevance scoring.    """       def __init__(        self,        max_context_length: int = 8192,        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",        relevance_threshold: float = 0.7,        recency_weight: float = 0.3,        importance_weight: float = 0.3,        semantic_weight: float = 0.4,        device: str = "cuda" if torch.cuda.is_available() else "cpu"    ):        """        Initialize the Model Context Manager.               Args:            max_context_length: Maximum number of tokens in context window            embedding_model: Model to use for text embeddings            relevance_threshold: Threshold for chunk relevance to be included            recency_weight: Weight for recency in relevance calculation            importance_weight: Weight for importance in relevance calculation            semantic_weight: Weight for semantic similarity in relevance calculation            device: Device to run computations on        """        self.max_context_length = max_context_length        self.device = device        self.chunks = []        self.current_token_count = 0        self.relevance_threshold = relevance_threshold               self.recency_weight = recency_weight        self.importance_weight = importance_weight        self.semantic_weight = semantic_weight               try:            from sentence_transformers import SentenceTransformer            print(f"Loading embedding model {embedding_model}...")            self.embedding_model = SentenceTransformer(embedding_model).to(self.device)            print(f"Embedding model loaded successfully on {self.device}")        except ImportError:            print("Installing sentence-transformers...")            import subprocess            subprocess.run(["pip", "install", "sentence-transformers"])            from sentence_transformers import SentenceTransformer            self.embedding_model = SentenceTransformer(embedding_model).to(self.device)            print(f"Embedding model loaded successfully on {self.device}")                   try:            from transformers import GPT2Tokenizer            self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")        except ImportError:            print("Installing transformers...")            import subprocess            subprocess.run(["pip", "install", "transformers"])            from transformers import GPT2Tokenizer            self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")       def add_chunk(self, text: str, importance: float = 1.0, metadata: Dict[str, Any] = None) -> None:        """        Add a new chunk of text to the context manager.               Args:            text: The text content to add            importance: Importance score (0-1)            metadata: Additional metadata for the chunk        """        with torch.no_grad():            embedding = self.embedding_model.encode(text, convert_to_tensor=True)               chunk = ContextChunk(            text=text,            embedding=embedding,            importance=importance,            timestamp=time.time(),            metadata=metadata or {}        )               self.chunks.append(chunk)        self.current_token_count += len(self.tokenizer.encode(text))               if self.current_token_count > self.max_context_length:            self.optimize_context()       def optimize_context(self) -> None:        """Optimize context by removing less relevant chunks to fit within token limit."""        if not self.chunks:            return                   print("Optimizing context window...")               scores = self.score_chunks()               sorted_indices = np.argsort(scores)[::-1]               new_chunks = []        new_token_count = 0               for idx in sorted_indices:            chunk = self.chunks[idx]            chunk_tokens = len(self.tokenizer.encode(chunk.text))                       if new_token_count + chunk_tokens <= self.max_context_length:                new_chunks.append(chunk)                new_token_count += chunk_tokens            else:                if scores[idx] > self.relevance_threshold * 1.5:                    for i, included_chunk in enumerate(new_chunks):                        included_idx = sorted_indices[i]                        if scores[included_idx] < self.relevance_threshold:                            included_tokens = len(self.tokenizer.encode(included_chunk.text))                            if new_token_count - included_tokens + chunk_tokens <= self.max_context_length:                                new_chunks.remove(included_chunk)                                new_token_count -= included_tokens                                new_chunks.append(chunk)                                new_token_count += chunk_tokens                                break               removed_count = len(self.chunks) - len(new_chunks)        self.chunks = new_chunks        self.current_token_count = new_token_count               print(f"Context optimized: Removed {removed_count} chunks, {len(new_chunks)} remaining, using {new_token_count}/{self.max_context_length} tokens")               gc.collect()        if torch.cuda.is_available():            torch.cuda.empty_cache()       def score_chunks(self, query: str = None) -> np.ndarray:        """        Score chunks based on recency, importance, and semantic relevance.               Args:            query: Optional query to calculate semantic relevance against                   Returns:            Array of scores for each chunk        """        if not self.chunks:            return np.array([])                   current_time = time.time()        max_age = max(current_time - chunk.timestamp for chunk in self.chunks) or 1.0        recency_scores = np.array([            1.0 - ((current_time - chunk.timestamp) / max_age)            for chunk in self.chunks        ])               importance_scores = np.array([chunk.importance for chunk in self.chunks])               if query is not None:            query_embedding = self.embedding_model.encode(query, convert_to_tensor=True)            similarity_scores = np.array([                torch.cosine_similarity(chunk.embedding, query_embedding, dim=0).item()                for chunk in self.chunks            ])                       similarity_scores = (similarity_scores - similarity_scores.min()) / (similarity_scores.max() - similarity_scores.min() + 1e-8)        else:            similarity_scores = np.ones(len(self.chunks))               final_scores = (            self.recency_weight * recency_scores +            self.importance_weight * importance_scores +            self.semantic_weight * similarity_scores        )               return final_scores       def retrieve_context(self, query: str = None, k: int = None) -> str:        """        Retrieve the most relevant context for a given query.               Args:            query: The query to retrieve context for            k: The maximum number of chunks to return (None = all relevant chunks)                   Returns:            String containing the combined relevant context        """        if not self.chunks:            return ""                   scores = self.score_chunks(query)               relevant_indices = np.where(scores >= self.relevance_threshold)[0]               relevant_indices = relevant_indices[np.argsort(scores[relevant_indices])[::-1]]               if k is not None:            relevant_indices = relevant_indices[:k]                   relevant_texts = [self.chunks[i].text for i in relevant_indices]        return "\n\n".join(relevant_texts)       def get_stats(self) -> Dict[str, Any]:        """Get statistics about the current context state."""        return {            "chunk_count": len(self.chunks),            "token_count": self.current_token_count,            "max_tokens": self.max_context_length,            "usage_percentage": self.current_token_count / self.max_context_length * 100 if self.max_context_length else 0,            "avg_chunk_size": self.current_token_count / len(self.chunks) if self.chunks else 0,            "oldest_chunk_age": time.time() - min(chunk.timestamp for chunk in self.chunks) if self.chunks else 0,        }    def visualize_context(self):        """Visualize the current context window distribution."""        try:            import matplotlib.pyplot as plt            import pandas as pd                       if not self.chunks:                print("No chunks to visualize")                return                       scores = self.score_chunks()            chunk_sizes = [len(self.tokenizer.encode(chunk.text)) for chunk in self.chunks]            timestamps = [chunk.timestamp for chunk in self.chunks]            relative_times = [time.time() - ts for ts in timestamps]            importance = [chunk.importance for chunk in self.chunks]                       df = pd.DataFrame({                'Size (tokens)': chunk_sizes,                'Age (seconds)': relative_times,                'Importance': importance,                'Score': scores            })                       fig, axs = plt.subplots(2, 2, figsize=(14, 10))                       axs[0, 0].bar(range(len(chunk_sizes)), chunk_sizes)            axs[0, 0].set_title('Token Distribution by Chunk')            axs[0, 0].set_ylabel('Tokens')            axs[0, 0].set_xlabel('Chunk Index')                       axs[0, 1].scatter(chunk_sizes, scores)            axs[0, 1].set_title('Score vs Chunk Size')            axs[0, 1].set_xlabel('Tokens')            axs[0, 1].set_ylabel('Score')                       axs[1, 0].scatter(relative_times, scores)            axs[1, 0].set_title('Score vs Chunk Age')            axs[1, 0].set_xlabel('Age (seconds)')            axs[1, 0].set_ylabel('Score')                       axs[1, 1].scatter(importance, scores)            axs[1, 1].set_title('Score vs Importance')            axs[1, 1].set_xlabel('Importance')            axs[1, 1].set_ylabel('Score')                       plt.tight_layout()            plt.show()                   except ImportError:            print("Please install matplotlib and pandas for visualization")            print('!pip install matplotlib pandas')

class MCPColabDemo:    """Demonstration of Model Context Protocol in Google Colab with a Language Model."""       def __init__(        self,        model_name: str = "google/flan-t5-base",        max_context_length: int = 2048,        device: str = "cuda" if torch.cuda.is_available() else "cpu"    ):        """        Initialize the MCP Colab demo with a specified model.               Args:            model_name: Hugging Face model name            max_context_length: Maximum context length for the MCP manager            device: Device to run the model on        """        self.device = device        self.context_manager = ModelContextManager(            max_context_length=max_context_length,            device=device        )               try:            from transformers import AutoModelForSeq2SeqLM, AutoTokenizer            print(f"Loading model {model_name}...")            self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)            self.tokenizer = AutoTokenizer.from_pretrained(model_name)            print(f"Model loaded successfully on {device}")        except ImportError:            print("Installing transformers...")            import subprocess            subprocess.run(["pip", "install", "transformers"])            from transformers import AutoModelForSeq2SeqLM, AutoTokenizer            self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)            self.tokenizer = AutoTokenizer.from_pretrained(model_name)            print(f"Model loaded successfully on {device}")       def add_document(self, text: str, chunk_size: int = 512, overlap: int = 50) -> None:        """        Add a document to the context by chunking it appropriately.               Args:            text: Document text            chunk_size: Size of each chunk in characters            overlap: Overlap between chunks in characters        """        chunks = []        for i in range(0, len(text), chunk_size - overlap):            chunk = text[i:i + chunk_size]            if len(chunk) > 20:                  chunks.append(chunk)               print(f"Adding {len(chunks)} chunks to context...")        for i, chunk in enumerate(tqdm(chunks)):            pos = i / len(chunks)            importance = 1.0 - 0.5 * min(pos, 1 - pos)                       self.context_manager.add_chunk(                text=chunk,                importance=importance,                metadata={"source": "document", "position": i, "total_chunks": len(chunks)}            )       def process_query(self, query: str, max_new_tokens: int = 256) -> str:        """        Process a query using the context manager and model.               Args:            query: The query to process            max_new_tokens: Maximum number of tokens in response                   Returns:            Model response        """        self.context_manager.add_chunk(query, importance=1.0, metadata={"type": "query"})               relevant_context = self.context_manager.retrieve_context(query=query)               prompt = f"Context: {relevant_context}\n\nQuestion: {query}\n\nAnswer:"               inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)               print("Generating response...")        with torch.no_grad():            outputs = self.model.generate(                **inputs,                max_new_tokens=max_new_tokens,                do_sample=True,                temperature=0.7,                top_p=0.9,            )               response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)               self.context_manager.add_chunk(            response,            importance=0.9,            metadata={"type": "response", "query": query}        )               return response       def interactive_session(self):        """Run an interactive session in the notebook."""        from IPython.display import clear_output               print("Starting interactive MCP session. Type 'exit' to end.")        conversation_history = []               while True:            query = input("\nYour query: ")                       if query.lower() == 'exit':                break                           if query.lower() == 'stats':                print("\nContext Statistics:")                stats = self.context_manager.get_stats()                for key, value in stats.items():                    print(f"{key}: {value}")                self.context_manager.visualize_context()                continue                           if query.lower() == 'clear':                self.context_manager.chunks = []                self.context_manager.current_token_count = 0                conversation_history = []                clear_output(wait=True)                print("Context cleared!")                continue                       response = self.process_query(query)            conversation_history.append((query, response))                       print("\nResponse:")            print(response)            print("\n" + "-"*50)                       stats = self.context_manager.get_stats()            print(f"Context usage: {stats['token_count']}/{stats['max_tokens']} tokens ({stats['usage_percentage']:.1f}%)")

def run_mcp_demo():    """Run a simple demo of the Model Context Protocol."""    print("Running Model Context Protocol Demo...")       context_manager = ModelContextManager(max_context_length=4096)       print("Adding sample chunks...")       context_manager.add_chunk(        "The Model Context Protocol (MCP) is a framework for managing context "        "windows in large language models. It helps optimize token usage and improve relevance.",        importance=1.0    )       context_manager.add_chunk(        "Context management involves techniques like sliding windows, chunking, "        "and relevance filtering to handle large documents efficiently.",        importance=0.8    )       for i in range(10):        context_manager.add_chunk(            f"This is test chunk {i} with some filler content to simulate a larger context "            f"window that needs optimization. This helps demonstrate the MCP functionality "            f"for context window management in language models on Google Colab.",            importance=0.5 - (i * 0.02)          )       stats = context_manager.get_stats()    print("\nInitial Statistics:")    for key, value in stats.items():        print(f"{key}: {value}")           query = "How does the Model Context Protocol work?"    print(f"\nRetrieving context for: '{query}'")    context = context_manager.retrieve_context(query)    print(f"\nRelevant context:\n{context}")       print("\nVisualizing context:")    context_manager.visualize_context()       print("\nDemo complete!")

if __name__ == "__main__":    run_mcp_demo()