M08: Conversation Management

from openai import OpenAI
# pip install openai
from dataclasses import dataclass, field
from typing import Optional

client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")

@dataclass
class ConversationManager:
    """Manages multi-turn conversations with Mistral via Ollama."""
    system_prompt: str = "You are a helpful assistant."
    model: str = "mistral"
    messages: list = field(default_factory=list)

    def __post_init__(self):
        self.total_input_tokens = 0
        self.total_output_tokens = 0

    def add_user_message(self, content: str) -> None:
        """Add a user message to the conversation history."""
        self.messages.append({"role": "user", "content": content})

    def add_assistant_message(self, content: str) -> None:
        """Add an assistant message to the conversation history."""
        self.messages.append({"role": "assistant", "content": content})

    def get_messages(self) -> list:
        """Return messages with system prompt prepended."""
        return [{"role": "system", "content": self.system_prompt}] + list(self.messages)

    def send(self, user_message: str) -> str:
        """Send a message and get the model's response."""
        self.add_user_message(user_message)

        try:
            response = client.chat.completions.create(
                model=self.model,
                messages=self.get_messages(),
            )

            # Track token usage
            if response.usage:
                self.total_input_tokens += response.usage.prompt_tokens
                self.total_output_tokens += response.usage.completion_tokens

            assistant_text = response.choices[0].message.content
            self.add_assistant_message(assistant_text)
            return assistant_text

        except Exception as e:
            # Remove the user message we just added on failure
            self.messages.pop()
            raise RuntimeError(f"API call failed: {e}") from e

    def get_token_usage(self) -> dict:
        """Return cumulative token usage."""
        return {
            "total_input_tokens": self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
            "message_count": len(self.messages),
        }


# Usage
manager = ConversationManager(
    system_prompt="You are a coding tutor. Be concise."
)

reply1 = manager.send("What is a list comprehension in Python?")
print(reply1)

reply2 = manager.send("Show me an example with filtering.")
print(reply2)

print(manager.get_token_usage())

class SlidingWindowManager(ConversationManager):
    """Extends ConversationManager with a sliding window."""

    def __init__(self, window_size: int = 10, **kwargs):
        super().__init__(**kwargs)
        self.window_size = window_size

    def get_messages(self) -> list:
        """Return system prompt + only the most recent N messages."""
        if len(self.messages) <= self.window_size:
            windowed = list(self.messages)
        else:
            windowed = self.messages[-self.window_size:]

        # Ensure we start with a user message (valid alternation)
        if windowed and windowed[0]["role"] == "assistant":
            windowed = windowed[1:]

        return [{"role": "system", "content": self.system_prompt}] + windowed


# Usage — only last 6 messages sent per call
manager = SlidingWindowManager(
    window_size=6,
    system_prompt="You are a concise coding tutor.",
)
for q in [
    "What is Python?",
    "What are variables?",
    "Explain loops.",
    "What are functions?",
    "Explain classes.",
    "What is inheritance?",
]:
    print(f"Q: {q}")
    print(f"A: {manager.send(q)}\n")

# All 12 messages stored, but only last 6 sent to API
print(f"Stored: {len(manager.messages)} messages")
print(f"Sent (excl. system): {len(manager.get_messages()) - 1} messages")

import json
import time
from pathlib import Path

class SmartConversationManager(ConversationManager):
    """Full-featured manager with summarization and persistence."""

    def __init__(
        self,
        token_threshold: int = 8_000,   # Lower for Mistral's 32K context
        recent_turns_to_keep: int = 4,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.token_threshold = token_threshold
        self.recent_turns_to_keep = recent_turns_to_keep
        self.summary: Optional[str] = None
        self.summary_history: list = []
        self.last_input_tokens = 0

    def _should_summarize(self) -> bool:
        return self.last_input_tokens > self.token_threshold

    def _summarize_old_messages(self) -> None:
        """Use Mistral to summarize older messages."""
        keep_count = self.recent_turns_to_keep * 2  # user+assistant pairs
        if len(self.messages) <= keep_count:
            return

        old_messages = self.messages[:-keep_count]
        recent_messages = self.messages[-keep_count:]

        summary_prompt = (
            "Summarize this conversation concisely. "
            "Preserve: key decisions, user preferences, "
            "important facts. Skip: greetings, filler.\n\n"
        )
        for msg in old_messages:
            summary_prompt += f"{msg['role']}: {msg['content']}\n"

        try:
            response = client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "You are a conversation summarizer. Be concise."},
                    {"role": "user", "content": summary_prompt}
                ],
            )
            new_summary = response.choices[0].message.content

            if self.summary:
                new_summary = f"Previous context: {self.summary}\n\nRecent: {new_summary}"

            self.summary = new_summary
            self.summary_history.append({
                "timestamp": time.time(),
                "messages_summarized": len(old_messages),
            })

            # Replace old messages with a summary message pair
            self.messages = [
                {"role": "user", "content": f"[Conversation summary: {self.summary}]"},
                {"role": "assistant", "content": "Understood. I have the conversation context."},
                *recent_messages,
            ]
        except Exception:
            # If summarization fails, fall back to sliding window
            self.messages = recent_messages

    def send(self, user_message: str) -> str:
        """Send with automatic summarization when needed."""
        self.add_user_message(user_message)

        try:
            msgs_to_send = [{"role": "system", "content": self.system_prompt}] + self.messages
            response = client.chat.completions.create(
                model=self.model,
                messages=msgs_to_send,
            )
            if response.usage:
                self.last_input_tokens = response.usage.prompt_tokens
                self.total_input_tokens += response.usage.prompt_tokens
                self.total_output_tokens += response.usage.completion_tokens

            assistant_text = response.choices[0].message.content
            self.add_assistant_message(assistant_text)

            # Check if we need to summarize
            if self._should_summarize():
                self._summarize_old_messages()

            return assistant_text

        except Exception as e:
            self.messages.pop()
            raise RuntimeError(f"API call failed: {e}") from e

    def save(self, filepath: str) -> None:
        """Persist conversation state to a JSON file."""
        state = {
            "system_prompt": self.system_prompt,
            "model": self.model,
            "messages": self.messages,
            "summary": self.summary,
            "summary_history": self.summary_history,
            "total_input_tokens": self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
            "saved_at": time.time(),
        }
        Path(filepath).write_text(json.dumps(state, indent=2))

    @classmethod
    def load(cls, filepath: str) -> "SmartConversationManager":
        """Load conversation state from a JSON file."""
        data = json.loads(Path(filepath).read_text())
        mgr = cls(
            system_prompt=data["system_prompt"],
            model=data["model"],
        )
        mgr.messages = data["messages"]
        mgr.summary = data.get("summary")
        mgr.summary_history = data.get("summary_history", [])
        mgr.total_input_tokens = data.get("total_input_tokens", 0)
        mgr.total_output_tokens = data.get("total_output_tokens", 0)
        return mgr


# Usage
manager = SmartConversationManager(
    token_threshold=3_000,   # Low threshold for testing
    recent_turns_to_keep=4,
    system_prompt="You are a helpful coding assistant.",
)

reply = manager.send("Help me build a REST API with FastAPI.")
print(reply)

# Save and restore across sessions
manager.save("conversation_state.json")
restored = SmartConversationManager.load("conversation_state.json")
reply2 = restored.send("Where were we?")
print(reply2)

mkdir conversation-lab && cd conversation-lab
python -m venv venv && source venv/bin/activate   # Windows: venv\Scripts\activate
pip install openai
# Make sure Ollama is running: ollama serve && ollama run mistral

from openai import OpenAI
import json
import time
from pathlib import Path
from dataclasses import dataclass, field
from typing import Optional

# pip install openai
client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")

# ── Step 1: Basic ConversationManager ────────────────────────
@dataclass
class ConversationManager:
    """Full-history conversation manager with token tracking."""
    system_prompt: str = "You are a helpful assistant."
    model: str = "mistral"
    messages: list = field(default_factory=list)

    def __post_init__(self):
        self.total_input_tokens = 0
        self.total_output_tokens = 0

    def add_user_message(self, content: str) -> None:
        self.messages.append({"role": "user", "content": content})

    def add_assistant_message(self, content: str) -> None:
        self.messages.append({"role": "assistant", "content": content})

    def get_messages(self) -> list:
        return [{"role": "system", "content": self.system_prompt}] + list(self.messages)

    def send(self, user_message: str) -> str:
        self.add_user_message(user_message)
        try:
            response = client.chat.completions.create(
                model=self.model,
                messages=self.get_messages(),
            )
            if response.usage:
                self.total_input_tokens += response.usage.prompt_tokens
                self.total_output_tokens += response.usage.completion_tokens
            assistant_text = response.choices[0].message.content
            self.add_assistant_message(assistant_text)
            return assistant_text
        except Exception as e:
            self.messages.pop()  # remove failed user message
            raise RuntimeError(f"API call failed: {e}") from e

    def get_token_usage(self) -> dict:
        return {
            "total_input": self.total_input_tokens,
            "total_output": self.total_output_tokens,
            "messages": len(self.messages),
        }

# ── Test Step 1 ──────────────────────────────────────────────
if __name__ == "__main__":
    print("═" * 50)
    print("TEST 1: Basic ConversationManager (Full History)")
    print("═" * 50)

    mgr = ConversationManager(system_prompt="You are a concise coding tutor. Answer in 1-2 sentences.")

    questions = [
        "What is a list in Python?",
        "How do I add an item to it?",
        "What about removing items?",
    ]

    for i, q in enumerate(questions, 1):
        print(f"\n  Turn {i}: {q}")
        reply = mgr.send(q)
        print(f"  Mistral: {reply[:100]}...")
        usage = mgr.get_token_usage()
        print(f"  [Messages: {usage['messages']} | Input tokens so far: {usage['total_input']}]")

    print(f"\n  ✓ Final: {usage['messages']} messages, {usage['total_input']} total input tokens")
    print(f"  Note: Input tokens grow each turn because ALL messages are resent!")

# ── Step 2a: Sliding Window Manager ──────────────────────────
class SlidingWindowManager(ConversationManager):
    """Sends only the last N messages to the API."""

    def __init__(self, window_size: int = 6, **kwargs):
        super().__init__(**kwargs)
        self.window_size = window_size

    def get_messages(self) -> list:
        if len(self.messages) <= self.window_size:
            windowed = list(self.messages)
        else:
            windowed = self.messages[-self.window_size:]
        # Ensure we start with a user message (valid alternation)
        if windowed and windowed[0]["role"] == "assistant":
            windowed = windowed[1:]
        return [{"role": "system", "content": self.system_prompt}] + windowed


# ── Step 2b: Smart Manager with Summarization + Persistence ──
class SmartConversationManager(ConversationManager):
    """Auto-summarizes old turns and saves/loads state."""

    def __init__(self, token_threshold: int = 3_000, recent_turns: int = 4, **kwargs):
        super().__init__(**kwargs)
        self.token_threshold = token_threshold
        self.recent_turns = recent_turns
        self.summary: Optional[str] = None
        self.last_input_tokens = 0

    def send(self, user_message: str) -> str:
        self.add_user_message(user_message)
        try:
            msgs_to_send = [{"role": "system", "content": self.system_prompt}] + self.messages
            response = client.chat.completions.create(
                model=self.model,
                messages=msgs_to_send,
            )
            if response.usage:
                self.last_input_tokens = response.usage.prompt_tokens
                self.total_input_tokens += response.usage.prompt_tokens
                self.total_output_tokens += response.usage.completion_tokens

            assistant_text = response.choices[0].message.content
            self.add_assistant_message(assistant_text)

            # Auto-summarize if threshold exceeded
            if self.last_input_tokens > self.token_threshold:
                self._summarize()

            return assistant_text
        except Exception as e:
            self.messages.pop()
            raise RuntimeError(f"API call failed: {e}") from e

    def _summarize(self) -> None:
        keep_count = self.recent_turns * 2  # user+assistant pairs
        if len(self.messages) <= keep_count:
            return

        old_msgs = self.messages[:-keep_count]
        recent_msgs = self.messages[-keep_count:]

        prompt = "Summarize this conversation concisely. Preserve: key decisions, facts. Skip: greetings, filler.\n\n"
        for m in old_msgs:
            prompt += f"{m['role']}: {m['content']}\n"

        try:
            resp = client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "You are a conversation summarizer."},
                    {"role": "user", "content": prompt}
                ],
            )
            summary_text = resp.choices[0].message.content
            if self.summary:
                summary_text = f"Previous: {self.summary}\n\nRecent: {summary_text}"
            self.summary = summary_text
            self.messages = [
                {"role": "user", "content": f"[Context summary: {self.summary}]"},
                {"role": "assistant", "content": "Understood. I have the context."},
                *recent_msgs,
            ]
            print(f"    ⚡ Summarized {len(old_msgs)} old messages → {len(self.messages)} total")
        except Exception:
            self.messages = recent_msgs  # fallback

    def save(self, filepath: str) -> None:
        state = {
            "system_prompt": self.system_prompt,
            "model": self.model,
            "messages": self.messages,
            "summary": self.summary,
            "total_input_tokens": self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
        }
        Path(filepath).write_text(json.dumps(state, indent=2))
        print(f"    💾 Saved to {filepath}")

    @classmethod
    def load(cls, filepath: str) -> "SmartConversationManager":
        data = json.loads(Path(filepath).read_text())
        mgr = cls(system_prompt=data["system_prompt"], model=data["model"])
        mgr.messages = data["messages"]
        mgr.summary = data.get("summary")
        mgr.total_input_tokens = data.get("total_input_tokens", 0)
        mgr.total_output_tokens = data.get("total_output_tokens", 0)
        return mgr


# ── Test All Three Strategies ────────────────────────────────
if __name__ == "__main__":
    questions = [
        "What is a Python list?",
        "How do I sort a list?",
        "What about list comprehensions?",
        "How do dictionaries differ from lists?",
        "What are sets?",
        "When should I use tuples?",
    ]

    # --- Test 1: Full History ---
    print("═" * 55)
    print("TEST 1: Full History (all messages sent every time)")
    print("═" * 55)
    mgr1 = ConversationManager(system_prompt="Answer in 1 sentence.")
    for i, q in enumerate(questions, 1):
        mgr1.send(q)
        u = mgr1.get_token_usage()
        print(f"  Turn {i}: stored={u['messages']} sent={u['messages']+1} input_tokens={u['total_input']}")

    # --- Test 2: Sliding Window ---
    print(f"\n{'═' * 55}")
    print("TEST 2: Sliding Window (last 4 messages only)")
    print("═" * 55)
    mgr2 = SlidingWindowManager(window_size=4, system_prompt="Answer in 1 sentence.")
    for i, q in enumerate(questions, 1):
        mgr2.send(q)
        u = mgr2.get_token_usage()
        sent = len(mgr2.get_messages()) - 1  # exclude system
        print(f"  Turn {i}: stored={u['messages']} sent={sent} input_tokens={u['total_input']}")

    # --- Test 3: Smart Manager + Persistence ---
    print(f"\n{'═' * 55}")
    print("TEST 3: Smart Manager (auto-summarization + save/load)")
    print("═" * 55)
    mgr3 = SmartConversationManager(
        token_threshold=200,  # very low to force summarization in test
        recent_turns=2,
        system_prompt="Answer in 1 sentence.",
    )
    for i, q in enumerate(questions, 1):
        mgr3.send(q)
        u = mgr3.get_token_usage()
        print(f"  Turn {i}: messages={u['messages']} input_tokens={u['total_input']}")

    mgr3.save("test_session.json")
    restored = SmartConversationManager.load("test_session.json")
    reply = restored.send("What did we discuss?")
    print(f"\n  Restored reply: {reply[:120]}...")

    # --- Comparison ---
    print(f"\n{'═' * 55}")
    print("COMPARISON (total input tokens after 6 turns):")
    print(f"  Full History:    {mgr1.total_input_tokens:,} tokens")
    print(f"  Sliding Window:  {mgr2.total_input_tokens:,} tokens")
    print(f"  Smart Manager:   {mgr3.total_input_tokens:,} tokens")
    print("═" * 55)