M03B: Context Engineering — Curating What the Model Sees

# WHAT: Decide which layers to include based on query type.
# WHY:  Unused layers waste tokens every call. 4 tools = ~920 overhead tokens.
# GOTCHA: Never crop the system prompt — it contains safety rules and persona.

from openai import OpenAI

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

TOOLS = [
    {"type": "function", "function": {
        "name": "track_order", "description": "Track shipment status by order ID",
        "parameters": {"type": "object", "properties": {"order_id": {"type": "string"}}, "required": ["order_id"]}
    }},
    {"type": "function", "function": {
        "name": "list_orders", "description": "List recent orders for a customer",
        "parameters": {"type": "object", "properties": {"customer_id": {"type": "string"}}, "required": ["customer_id"]}
    }},
]

def needs_tools(user_message: str) -> bool:
    """Heuristic: skip tool definitions if query is clearly conversational."""
    tool_keywords = ["order", "track", "ship", "status", "where", "when", "delivery"]
    msg_lower = user_message.lower()
    return any(kw in msg_lower for kw in tool_keywords)

def call_model(system: str, messages: list, user_message: str) -> str:
    include_tools = needs_tools(user_message)
    full_messages = messages + [{"role": "user", "content": user_message}]

    kwargs = {
        "model": "mistral",
        "messages": [{"role": "system", "content": system}] + full_messages,
        "max_tokens": 512,
    }
    if include_tools:
        kwargs["tools"] = TOOLS  # Only inject when relevant — saves ~920 tokens otherwise

    try:
        response = client.chat.completions.create(**kwargs)
        return response.choices[0].message.content or ""
    except Exception as e:
        return f"[error: {e}]"

# Test it
print(call_model(
    system="You are a friendly order-tracking assistant.",
    messages=[],
    user_message="Hi! How are you today?"  # no tools needed — saves 920 tokens
))

# WHAT: Extract only the fields the model needs from a large tool result.
# WHY:  Raw API responses can be 2,000+ tokens. We need <100.
# GOTCHA: Never compress status codes or IDs — they're cheap and critical.

import json
from typing import Any

def compress_order_result(raw: dict[str, Any]) -> str:
    """Compress a full order API response to ~80 tokens."""
    fields = {
        "order_id":        raw.get("order_id", "unknown"),
        "status":          raw.get("status", "unknown"),
        "carrier":         raw.get("shipment", {}).get("carrier", "N/A"),
        "tracking":        raw.get("shipment", {}).get("tracking_number", "N/A"),
        "estimated_delivery": raw.get("shipment", {}).get("estimated_delivery", "unknown"),
        "last_update":     raw.get("shipment", {}).get("events", [{}])[-1].get("description", "none"),
    }
    return json.dumps(fields, separators=(",", ":"))

# Example: raw = 2,150 tokens; compressed = 68 tokens
raw_order = {
    "order_id": "ORD-88421", "customer_id": "C-4512",
    "created_at": "2024-11-02T08:15:00Z", "updated_at": "2024-11-05T14:30:00Z",
    "line_items": [{"sku": "SKU-991", "qty": 2, "unit_price": 49.99, "description": "Widget Pro"}],
    "subtotal": 99.98, "tax": 8.75, "shipping_cost": 12.00, "total": 120.73,
    "status": "in_transit",
    "shipment": {
        "carrier": "FedEx", "tracking_number": "4782991234500012",
        "estimated_delivery": "2024-11-08",
        "events": [
            {"ts": "2024-11-02", "description": "Label created"},
            {"ts": "2024-11-03", "description": "Picked up by carrier"},
            {"ts": "2024-11-05", "description": "In transit: Memphis hub"},
        ]
    },
    "billing_address": {"street": "123 Main St", "city": "Boston", "state": "MA", "zip": "02101"},
    "notes": []
}

print(compress_order_result(raw_order))
# {"order_id":"ORD-88421","status":"in_transit","carrier":"FedEx",
#  "tracking":"4782991234500012","estimated_delivery":"2024-11-08",
#  "last_update":"In transit: Memphis hub"}

# WHAT: Replace old turns with a model-generated summary. Keep recent k turns verbatim.
# WHY:  History is the biggest layer in long sessions. Old turns have low relevance/token.
# GOTCHA: Use a SEPARATE model call for the summary — never ask the main model to
#         summarize in the same call. It can't summarize what it's currently reading.

from openai import OpenAI

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

SUMMARY_PROMPT = """Summarize this conversation in 3-5 sentences.
CRITICAL: preserve every order ID, customer ID, tracking number, exact dollar amount,
and explicit user decision. Drop only retry errors, duplicate results, and resolved detours.

Conversation:
{transcript}

Summary:"""

def summarize_history(history: list[dict], keep_recent: int = 4) -> list[dict]:
    """Summarize older turns; keep the last `keep_recent` turns verbatim."""
    if len(history) <= keep_recent:
        return history  # nothing to do

    older = history[:-keep_recent]
    recent = history[-keep_recent:]

    transcript = "\n".join(f"{m['role']}: {m['content']}" for m in older)

    try:
        result = client.chat.completions.create(
            model="mistral",  # use a fast/cheap local model for summarization
            messages=[{"role": "user",
                       "content": SUMMARY_PROMPT.format(transcript=transcript)}],
            max_tokens=250,
        )
        summary_text = result.choices[0].message.content or ""
    except Exception as e:
        print(f"Summarization failed: {e}; falling back to truncation.")
        return recent  # safe fallback: just keep recent turns

    summary_msg = {
        "role": "user",
        "content": f"[Summary of {len(older)} earlier turns] {summary_text}",
    }
    return [summary_msg] + recent

import json
import tiktoken  # pip install tiktoken
from dataclasses import dataclass, field
from typing import Any

# cl100k_base is OpenAI's tokenizer — a good approximation for Mistral/Llama
_enc = tiktoken.get_encoding("cl100k_base")

def count_tokens(text: str | list | dict) -> int:
    """Count tokens in a string or JSON-serializable object."""
    if not isinstance(text, str):
        text = json.dumps(text, separators=(",", ":"))
    return max(1, len(_enc.encode(text)))

MODEL_LIMITS = {
    "mistral":  32_768,
    "mixtral":  32_768,
    "llama3":  131_072,
    "gemma2":    8_192,
}

@dataclass
class ContextBudget:
    """Accounts for and curates the six layers of a local model context."""
    system_prompt: str = ""
    tool_definitions: list = field(default_factory=list)
    retrieved_docs: list[str] = field(default_factory=list)   # semi-static, goes near top
    history: list[dict] = field(default_factory=list)          # [{"role","content"}]
    tool_results: list[str] = field(default_factory=list)
    current_user_message: str = ""
    model: str = "mistral"
    reserve_output: int = 2048   # tokens reserved for model response

    @property
    def max_tokens(self) -> int:
        return MODEL_LIMITS.get(self.model, 32_768) - self.reserve_output

    def count(self, text: str | list | dict) -> int:
        return count_tokens(text)

    def account(self) -> dict[str, int]:
        """Per-layer token breakdown."""
        return {
            "system":       count_tokens(self.system_prompt),
            "tools":        count_tokens(self.tool_definitions) if self.tool_definitions else 0,
            "retrieved":    sum(count_tokens(d) for d in self.retrieved_docs),
            "history":      sum(count_tokens(m["content"]) for m in self.history),
            "tool_results": sum(count_tokens(r) for r in self.tool_results),
            "current":      count_tokens(self.current_user_message),
        }

    def total(self) -> int:
        return sum(self.account().values())

    def remaining(self) -> int:
        return self.max_tokens - self.total()

    def fits(self) -> bool:
        return self.total() <= self.max_tokens

    def strategy(self) -> str:
        """Return recommended strategy based on budget utilization.
        Returns: "ok" | "compress" | "summarize" | "critical"
        """
        used = self.total()
        budget = self.max_tokens
        utilization = used / budget

        if utilization < 0.60:
            return "ok"           # plenty of room — no action needed
        elif utilization < 0.75:
            return "compress"     # crop tool definitions, compress tool results
        elif utilization < 0.90:
            return "summarize"    # summarize old history turns
        else:
            return "critical"     # emergency: keep only system + last 2 turns + current

    def crop(self, target: int | None = None) -> "ContextBudget":
        """Apply the recommended strategy in place. Returns self for chaining."""
        strat = self.strategy()
        if strat == "ok":
            return self

        if strat in ("compress", "summarize", "critical"):
            # Crop tool definitions when in trouble (saves ~900 tok for 4 tools)
            self.tool_definitions = []

        if strat in ("summarize", "critical"):
            # Hard-truncate history: keep only last 4 (summarize) or 2 (critical) turns
            keep = 2 if strat == "critical" else 4
            self.history = self.history[-keep:]
            self.tool_results = self.tool_results[-2:]  # keep last 2 tool results only

        if strat == "critical":
            # Nuclear option: drop all retrieved docs too
            self.retrieved_docs = []

        return self

    def checkpoint(self, keep_recent: int = 4) -> "ContextBudget":
        """Summarize old history turns. Call when strategy() == 'summarize'.
        This method is synchronous but calls the Ollama API internally.
        """
        self.history = summarize_history(self.history, keep_recent=keep_recent)
        return self

# Usage pattern — full agent turn with automatic budget management:
def agent_turn(budget: ContextBudget, user_message: str) -> str:
    from openai import OpenAI
    client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")

    budget.current_user_message = user_message

    strat = budget.strategy()
    if strat == "compress":
        budget.crop()
    elif strat == "summarize":
        budget.checkpoint()          # async summarization + crop
        budget.crop()
    elif strat == "critical":
        budget.crop()                # hard truncation only

    system, msgs = budget.build_messages()

    try:
        response = client.chat.completions.create(
            model=budget.model,
            messages=[{"role": "system", "content": system}] + msgs,
            max_tokens=budget.reserve_output,
        )
        reply = response.choices[0].message.content or ""
    except Exception as e:
        reply = f"[error: {e}]"

    # Append this turn to history for next call
    budget.history.append({"role": "user", "content": user_message})
    budget.history.append({"role": "assistant", "content": reply})
    return reply

    def build_messages(self) -> tuple[str, list[dict]]:
        """Assemble (system, messages) for client.chat.completions.create().
        Order: system (start, high recall) → history → retrieved docs → current message (end, high recall).
        Tool results are interleaved within history turns as they occurred.
        """
        # Retrieved docs appended to the CURRENT message, not the system prompt.
        # End position = maximum recall on Mistral 7B.
        retrieved_block = ""
        if self.retrieved_docs:
            retrieved_block = (
                "\n\n<reference_docs>\n"
                + "\n---\n".join(self.retrieved_docs)
                + "\n</reference_docs>\n\n"
            )

        messages = list(self.history)
        messages.append({
            "role": "user",
            "content": retrieved_block + self.current_user_message,
        })
        return self.system_prompt, messages

# Full example
if __name__ == "__main__":
    budget = ContextBudget(
        model="mistral",
        system_prompt="You are an order-tracking assistant. Always cite order IDs.",
        current_user_message="Where is my order ORD-88421?",
        retrieved_docs=["ORD-88421: in transit, estimated delivery Nov 8 via FedEx."],
    )
    print(f"Tokens: {budget.total()} / {budget.max_tokens} ({budget.strategy()})")
    system, messages = budget.build_messages()
    print(f"Messages: {len(messages)}, last content: {messages[-1]['content'][:80]}")

#!/usr/bin/env python3
"""diagnose.py — Demonstrate the poisoned-transcript effect and fix."""

import json, time
from openai import OpenAI
from context_budget import ContextBudget, summarize_history

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

def ask_model(budget: ContextBudget) -> tuple[str, dict]:
    system, messages = budget.build_messages()
    t0 = time.time()
    try:
        response = client.chat.completions.create(
            model="mistral",
            messages=[{"role": "system", "content": system}] + messages,
            max_tokens=300,
        )
        answer = response.choices[0].message.content or ""
        usage = {
            "input_tokens":  response.usage.prompt_tokens,
            "output_tokens": response.usage.completion_tokens,
            "latency_s": round(time.time() - t0, 2),
        }
    except Exception as e:
        answer, usage = f"[error: {e}]", {"input_tokens": 0, "output_tokens": 0, "latency_s": 0}
    return answer, usage

def main():
    fixture = json.load(open("poisoned_transcript.json"))

    base_args = dict(
        model="mistral",
        system_prompt=fixture["system_prompt"],
        tool_definitions=fixture["tool_definitions"],
        history=fixture["history"],
        current_user_message=fixture["current_user_message"],
    )

    # --- Run 1: rotted context ---
    rotted = ContextBudget(**base_args)
    breakdown = rotted.account()
    print("=== Token Breakdown (rotted) ===")
    for layer, tok in breakdown.items():
        print(f"  {layer:15s}: {tok:,} tokens")
    print(f"  {'TOTAL':15s}: {rotted.total():,} / {rotted.max_tokens:,}  strategy={rotted.strategy()!r}")

    print("\n>>> Run 1: ROTTED context (no fix)")
    answer_a, usage_a = ask_model(rotted)
    print(f"Tokens: {usage_a['input_tokens']} in, {usage_a['output_tokens']} out  ({usage_a['latency_s']}s)")
    print(f"Answer: {answer_a}\n")

    # --- Run 2: after checkpoint ---
    fixed = ContextBudget(**base_args)
    fixed.history = summarize_history(fixed.history, keep_recent=4)
    print(">>> Run 2: COMPRESSED context (after checkpoint)")
    fixed_breakdown = fixed.account()
    print(f"  Total after compression: {fixed.total():,} tokens  strategy={fixed.strategy()!r}")
    answer_b, usage_b = ask_model(fixed)
    print(f"Tokens: {usage_b['input_tokens']} in, {usage_b['output_tokens']} out  ({usage_b['latency_s']}s)")
    print(f"Answer: {answer_b}\n")

    print("=" * 60)
    print(f"Token delta:   {usage_a['input_tokens'] - usage_b['input_tokens']:+,} input tokens")
    print(f"Latency delta: {usage_a['latency_s'] - usage_b['latency_s']:+.2f}s")

if __name__ == "__main__":
    main()