M22: Cost Optimization

"""
model_selector_benchmark.py
Run: python model_selector_benchmark.py
Requires: ollama running locally
"""
import time, statistics, json
from ollama import Client

client = Client(host="http://localhost:11434")

TEST_PROMPTS = [
    "Extract the company name and filing date from: 'UCC-1 filed by Acme LLC on 2024-03-15 in Delaware'",
    "Is this address a PO Box? '123 Main St, Springfield IL 62701'",
    "Summarize in one sentence: the debtor is Riverside Corp, secured party is First National Bank, collateral is all business assets.",
    "Generate a JSON object with keys: debtor, secured_party, collateral. Values: 'Smith Inc', 'Chase Bank', 'equipment'",
    "What is the difference between a UCC-1 and a UCC-3 amendment?",
] * 4  # 20 total prompts

MODELS = ["phi3:mini", "mistral:7b", "llama3.1:8b"]
JUDGE_MODEL = "mistral:7b"

def score_response(prompt: str, response: str) -> float:
    """Ask a local judge model to score quality 1-10."""
    judge_prompt = f"""Score this LLM response 1-10 for correctness and completeness.
Prompt: {prompt[:200]}
Response: {response[:300]}
Reply with ONLY a number 1-10."""
    try:
        result = client.generate(model=JUDGE_MODEL, prompt=judge_prompt)
        score = float(result["response"].strip().split()[0])
        return min(max(score, 1.0), 10.0)
    except Exception:
        return 5.0  # neutral on parse failure

def benchmark_model(model: str) -> dict:
    latencies, scores = [], []
    for prompt in TEST_PROMPTS:
        t0 = time.perf_counter()
        try:
            result = client.generate(model=model, prompt=prompt)
            elapsed = time.perf_counter() - t0
            latencies.append(elapsed)
            scores.append(score_response(prompt, result["response"]))
        except Exception as e:
            print(f"  [ERROR] {model}: {e}")
            latencies.append(999.0)
            scores.append(0.0)

    latencies_sorted = sorted(latencies)
    return {
        "model": model,
        "p50_s": round(statistics.median(latencies), 2),
        "p95_s": round(latencies_sorted[int(len(latencies_sorted) * 0.95)], 2),
        "avg_quality": round(statistics.mean(scores), 1),
        "n": len(latencies),
    }

def main():
    print("Running model selection benchmark (20 prompts per model)...\n")
    results = []
    for model in MODELS:
        print(f"  Testing {model}...")
        r = benchmark_model(model)
        results.append(r)
        print(f"  p50={r['p50_s']}s  p95={r['p95_s']}s  quality={r['avg_quality']}/10")

    results.sort(key=lambda x: (-(x["avg_quality"]), x["p50_s"]))
    print("\n=== Ranked Results ===")
    print(f"{'Model':<25} {'p50 (s)':<10} {'p95 (s)':<10} {'Quality':<10}")
    print("-" * 55)
    for r in results:
        print(f"{r['model']:<25} {r['p50_s']:<10} {r['p95_s']:<10} {r['avg_quality']}/10")

    with open("model_benchmark.json", "w") as f:
        json.dump(results, f, indent=2)
    print("\nSaved: model_benchmark.json")

if __name__ == "__main__":
    main()

"""
prompt_efficiency.py — Five token-reduction techniques for Ollama agents
"""
import json
from ollama import Client

client = Client(host="http://localhost:11434")
MODEL = "mistral:7b-instruct-q4_K_M"

# -----------------------------------------------------------------
# TECHNIQUE 2: max_tokens budget per step
# Without a budget, each step can consume up to context_window tokens
# (4096 for Mistral). Budget to 256 for extraction tasks.
# -----------------------------------------------------------------
def step_with_budget(prompt: str, max_tokens: int = 256) -> str:
    result = client.generate(
        model=MODEL,
        prompt=prompt,
        options={"num_predict": max_tokens},  # Ollama option = max_tokens
    )
    return result["response"]

# -----------------------------------------------------------------
# TECHNIQUE 3: Tool result truncation before appending to messages
# A search tool may return 10 KB of JSON. Truncate to essential fields
# before appending to the conversation context.
# -----------------------------------------------------------------
def truncate_tool_result(result: dict, max_chars: int = 500) -> str:
    """Serialize and truncate tool output to stay within token budget."""
    serialized = json.dumps(result, separators=(",", ":"))
    if len(serialized) <= max_chars:
        return serialized
    # Keep first max_chars and add truncation marker
    return serialized[:max_chars] + "...[truncated]"

# -----------------------------------------------------------------
# TECHNIQUE 4: Sliding window conversation pruning
# Keep system prompt + last N turns. Drop oldest turns first.
# (Adapted from M11: Multi-Layer Memory)
# -----------------------------------------------------------------
class SlidingWindowConversation:
    def __init__(self, system_prompt: str, max_turns: int = 6):
        self.system = system_prompt
        self.max_turns = max_turns
        self.turns: list[dict] = []  # each turn = {"role": ..., "content": ...}

    def add_turn(self, role: str, content: str) -> None:
        self.turns.append({"role": role, "content": content})
        # Prune oldest turns when over budget (keep in pairs to avoid orphans)
        while len(self.turns) > self.max_turns * 2:
            self.turns.pop(0)
            if self.turns and self.turns[0]["role"] == "assistant":
                self.turns.pop(0)  # remove dangling assistant turn

    def build_prompt(self) -> str:
        """Build a single string prompt from system + sliding window turns."""
        parts = [f"System: {self.system}"]
        for turn in self.turns:
            prefix = "User" if turn["role"] == "user" else "Assistant"
            parts.append(f"{prefix}: {turn['content']}")
        parts.append("Assistant:")
        return "\n\n".join(parts)

    def token_estimate(self) -> int:
        """Rough estimate: ~4 chars per token."""
        return len(self.build_prompt()) // 4

# -----------------------------------------------------------------
# TECHNIQUE 5: Structured JSON responses are denser than prose
# Asking for JSON instead of prose can cut response tokens by 40-60%.
# -----------------------------------------------------------------
EXTRACTION_PROMPT = """Extract filing metadata from the text below.
Respond ONLY with a JSON object — no prose, no explanation, no markdown.
Required keys: debtor_name, secured_party, filing_date, collateral_summary.
Text: Acme Corp has granted First Capital Bank a security interest in all current
and future inventory and receivables. Filing recorded 2024-06-01 in Delaware."""

def extract_structured(prompt: str) -> dict:
    """Get structured JSON output — typically 40-60% fewer tokens than prose."""
    response = client.generate(
        model=MODEL,
        prompt=prompt,
        options={"num_predict": 200},  # JSON is compact, 200 is plenty
        format="json",  # Ollama JSON mode enforces valid JSON output
    )
    try:
        return json.loads(response["response"])
    except json.JSONDecodeError as e:
        raise ValueError(f"Model returned invalid JSON: {e}") from e

# -----------------------------------------------------------------
# DEMO: Before/after token count comparison
# -----------------------------------------------------------------
def count_tokens_approx(text: str) -> int:
    """Approximation: ~4 chars per token (good enough for comparison)."""
    return len(text) // 4

def demo_token_savings():
    prose_prompt = "Describe the UCC filing: Acme Corp granted First Capital Bank a security interest in inventory. Filing 2024-06-01 Delaware."
    json_prompt = EXTRACTION_PROMPT

    prose_response = step_with_budget(prose_prompt, max_tokens=300)
    json_response = step_with_budget(json_prompt + "\nRespond in JSON only.", max_tokens=150)

    prose_tokens = count_tokens_approx(prose_response)
    json_tokens = count_tokens_approx(json_response)

    print(f"Prose response: ~{prose_tokens} tokens")
    print(f"JSON response:  ~{json_tokens} tokens")
    print(f"Savings: {prose_tokens - json_tokens} tokens ({(1 - json_tokens/max(prose_tokens,1))*100:.0f}%)")

if __name__ == "__main__":
    demo_token_savings()

"""
batch_inference.py — Async batch processing with Ollama
"""
import asyncio, time, httpx, json
from dataclasses import dataclass

MODEL = "mistral:7b-instruct-q4_K_M"
OLLAMA_URL = "http://localhost:11434/api/generate"
CONCURRENCY_LIMIT = 8  # Safe cap for Mistral-7B on 16 GB VRAM

@dataclass
class BatchResult:
    index: int
    prompt: str
    response: str
    elapsed_s: float
    error: str | None = None

async def single_inference(
    client: httpx.AsyncClient,
    semaphore: asyncio.Semaphore,
    index: int,
    prompt: str,
) -> BatchResult:
    """Run one inference, respecting the concurrency semaphore."""
    async with semaphore:
        t0 = time.perf_counter()
        try:
            resp = await client.post(
                OLLAMA_URL,
                json={"model": MODEL, "prompt": prompt, "stream": False,
                      "options": {"num_predict": 256}},
                timeout=120.0,
            )
            resp.raise_for_status()
            data = resp.json()
            return BatchResult(
                index=index,
                prompt=prompt[:60],
                response=data["response"],
                elapsed_s=round(time.perf_counter() - t0, 2),
            )
        except Exception as e:
            return BatchResult(
                index=index, prompt=prompt[:60], response="",
                elapsed_s=round(time.perf_counter() - t0, 2), error=str(e),
            )

async def batch_infer(prompts: list[str], concurrency: int = CONCURRENCY_LIMIT) -> list[BatchResult]:
    """Process a list of prompts concurrently, capped at `concurrency`."""
    semaphore = asyncio.Semaphore(concurrency)
    async with httpx.AsyncClient() as client:
        tasks = [
            single_inference(client, semaphore, i, p)
            for i, p in enumerate(prompts)
        ]
        return await asyncio.gather(*tasks)

async def throughput_comparison(prompts: list[str]) -> None:
    # Serial baseline
    t0 = time.perf_counter()
    await batch_infer(prompts, concurrency=1)
    serial_s = time.perf_counter() - t0

    # Batch=4
    t0 = time.perf_counter()
    await batch_infer(prompts, concurrency=4)
    batch4_s = time.perf_counter() - t0

    # Batch=8
    t0 = time.perf_counter()
    await batch_infer(prompts, concurrency=8)
    batch8_s = time.perf_counter() - t0

    print(f"Serial (concurrency=1): {serial_s:.1f}s  1.0x")
    print(f"Batch=4:                {batch4_s:.1f}s  {serial_s/batch4_s:.1f}x speedup")
    print(f"Batch=8:                {batch8_s:.1f}s  {serial_s/batch8_s:.1f}x speedup")

if __name__ == "__main__":
    test_prompts = [f"Classify this text as business or personal: sample text #{i}" for i in range(20)]
    asyncio.run(throughput_comparison(test_prompts))

"""
agent_cache.py — Two-layer response cache with TTL
pip install chromadb sentence-transformers ollama
"""
import hashlib, time
from dataclasses import dataclass, field
from typing import Optional
import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
from ollama import Client

@dataclass
class CacheEntry:
    response: str
    created_at: float
    ttl_s: float

    def is_expired(self) -> bool:
        return time.time() > self.created_at + self.ttl_s

class AgentCache:
    """
    Layer 1: Exact-match dict cache (hash of query+model).
    Layer 2: Semantic vector cache via ChromaDB (cosine sim > threshold).
    Both layers support TTL-based expiration.
    """
    def __init__(
        self,
        model: str = "mistral:7b-instruct-q4_K_M",
        embed_model: str = "all-MiniLM-L6-v2",
        similarity_threshold: float = 0.95,
        default_ttl_s: float = 3600.0,  # 1 hour
    ):
        self.model = model
        self.threshold = similarity_threshold
        self.default_ttl = default_ttl_s
        self._exact: dict[str, CacheEntry] = {}  # Layer 1

        # Layer 2: ChromaDB with local sentence-transformer embeddings
        self._chroma = chromadb.Client()
        embed_fn = SentenceTransformerEmbeddingFunction(model_name=embed_model)
        self._collection = self._chroma.get_or_create_collection(
            name="agent_cache",
            embedding_function=embed_fn,
            metadata={"hnsw:space": "cosine"},
        )
        self._ollama = Client(host="http://localhost:11434")
        self._stats = {"hits": 0, "misses": 0}

    def _hash_key(self, query: str) -> str:
        return hashlib.sha256(f"{self.model}:{query}".encode()).hexdigest()[:16]

    def _purge_expired(self) -> None:
        expired = [k for k, v in self._exact.items() if v.is_expired()]
        for k in expired:
            del self._exact[k]

    def get(self, query: str) -> Optional[str]:
        self._purge_expired()

        # Layer 1: exact match
        key = self._hash_key(query)
        if key in self._exact:
            entry = self._exact[key]
            if not entry.is_expired():
                self._stats["hits"] += 1
                return entry.response

        # Layer 2: semantic similarity
        try:
            results = self._collection.query(
                query_texts=[query],
                n_results=1,
                include=["documents", "distances", "metadatas"],
            )
            if results["documents"] and results["documents"][0]:
                distance = results["distances"][0][0]
                similarity = 1 - distance  # ChromaDB returns distance, not similarity
                metadata = results["metadatas"][0][0]
                expires_at = metadata.get("expires_at", 0)
                if similarity >= self.threshold and time.time() < expires_at:
                    self._stats["hits"] += 1
                    return results["documents"][0][0]
        except Exception:
            pass  # Degrade gracefully to live inference on cache errors

        return None

    def set(self, query: str, response: str, ttl_s: Optional[float] = None) -> None:
        ttl = ttl_s or self.default_ttl
        key = self._hash_key(query)

        # Layer 1
        self._exact[key] = CacheEntry(response=response, created_at=time.time(), ttl_s=ttl)

        # Layer 2
        try:
            self._collection.upsert(
                ids=[key],
                documents=[response],
                metadatas=[{"query": query[:200], "expires_at": time.time() + ttl}],
            )
        except Exception:
            pass  # Layer 1 is sufficient fallback

    def infer(self, query: str, ttl_s: Optional[float] = None) -> str:
        """Full cache-aware inference: check cache, fall back to Ollama."""
        cached = self.get(query)
        if cached is not None:
            return cached

        self._stats["misses"] += 1
        result = self._ollama.generate(
            model=self.model,
            prompt=query,
            options={"num_predict": 256},
        )
        response = result["response"]
        self.set(query, response, ttl_s)
        return response

    @property
    def stats(self) -> dict:
        total = self._stats["hits"] + self._stats["misses"]
        return {
            **self._stats,
            "hit_rate": f"{self._stats['hits'] / max(total, 1) * 100:.1f}%",
        }

# Usage example
if __name__ == "__main__":
    cache = AgentCache()
    q1 = "What is the debtor name in UCC filing #12345?"
    q2 = "Who is the debtor listed on UCC-1 number 12345?"  # semantic near-duplicate

    r1 = cache.infer(q1)       # Miss — runs Ollama
    r2 = cache.infer(q1)       # Hit  — exact match
    r3 = cache.infer(q2)       # Likely hit — semantic similarity > 0.95
    print(f"Stats: {cache.stats}")

"""
agent_benchmark.py — Full benchmark harness for Ollama models
Run: python agent_benchmark.py
Output: benchmark_results.csv  +  benchmark_results.md
"""
import asyncio, csv, json, statistics, time
from pathlib import Path
import httpx

OLLAMA_URL = "http://localhost:11434/api/generate"
JUDGE_MODEL = "mistral:7b"  # Judge runs separately

MODELS = ["phi3:mini", "mistral:7b", "llama3.1:8b"]

BENCHMARK_PROMPTS = [
    # Entity extraction
    "Extract: debtor=?, secured_party=? from: 'TechCorp Inc granted Wells Fargo a lien on all software assets'",
    "Parse this UCC filing date: 'filed on the fifteenth day of March two thousand twenty-four'",
    "Is 'ACME LLC' and 'Acme Limited Liability Company' the same entity? Answer yes or no with brief reason.",
    # Summarization
    "Summarize in 20 words: Debtor Riverside Corp grants secured party First Bank interest in all machinery and equipment now owned or hereafter acquired.",
    "One-sentence summary: The amendment extends the maturity date of the original UCC-1 financing statement by 5 years from the original lapse date.",
    # Classification
    "Classify as high-risk or low-risk: secured party is a non-bank fintech, collateral is future receivables, filed in a state with no central registry.",
    "Is this a continuation, amendment, or termination? 'The secured party hereby terminates all security interests in the collateral described in the original filing.'",
    # JSON generation
    "Return JSON with keys: entity_type, risk_level, jurisdiction for: 'Delaware LLC with a federal tax lien from IRS'",
    "Generate a JSON summary of: 'UCC-1 #2024-003-456, debtor: Smith Farms LLC, secured party: AgriCredit Corp, collateral: livestock and crops'",
    # Reasoning
    "If a UCC-1 was filed in 2019 and not renewed, is it still valid in 2025? Explain in 2 sentences.",
    "Which takes priority: a UCC-1 filed today or a mortgage filed 3 years ago? Brief answer.",
    # Code assistance
    "Write a one-liner Python to extract all 10-digit UCC filing numbers from a string using regex.",
    "What HTTP status code does the Illinois SOS lien search API return for a not-found entity?",
    # Math/logic
    "A UCC-1 lapses after 5 years. Filed 2020-04-01. Lapse date?",
    "If processing 10,000 filings at 3 seconds each with batch=8, estimate total runtime in hours.",
    # Long context (tests KV pressure)
    "Given this partial data: entity_id=E001, name=Acme Corp, address=123 Main St, state=IL, filing_count=3, risk_score=0.72 — generate a brief risk summary.",
    "Compare: debtor_a='First National Leasing LLC' vs debtor_b='1st Natl. Leasing, LLC'. Likely same entity? Confidence?",
    # Edge cases
    "What should an agent do if the Ollama API returns a 503 error mid-response?",
    "An entity resolution agent gets conflicting addresses for the same debtor. What disambiguation strategy minimizes false merges?",
    "List 3 signs that a local LLM agent is suffering from context window overflow.",
]

async def run_single(client: httpx.AsyncClient, model: str, prompt: str) -> dict:
    t0 = time.perf_counter()
    try:
        resp = await client.post(
            OLLAMA_URL,
            json={"model": model, "prompt": prompt, "stream": False, "options": {"num_predict": 300}},
            timeout=180.0,
        )
        resp.raise_for_status()
        data = resp.json()
        elapsed = time.perf_counter() - t0
        eval_count = data.get("eval_count", 1)
        eval_duration_ns = data.get("eval_duration", elapsed * 1e9)
        tps = eval_count / (eval_duration_ns / 1e9) if eval_duration_ns > 0 else 0
        return {"prompt": prompt[:80], "response": data["response"], "elapsed_s": elapsed, "tps": tps, "error": None}
    except Exception as e:
        return {"prompt": prompt[:80], "response": "", "elapsed_s": time.perf_counter() - t0, "tps": 0, "error": str(e)}

async def score_batch(client: httpx.AsyncClient, prompts_responses: list[tuple[str, str]]) -> list[float]:
    scores = []
    for prompt, response in prompts_responses:
        judge_prompt = f"Score 1-10 for correctness. Prompt: {prompt[:150]} Response: {response[:250]}\nReply: number only."
        try:
            r = await client.post(OLLAMA_URL, json={"model": JUDGE_MODEL, "prompt": judge_prompt, "stream": False, "options": {"num_predict": 5}}, timeout=60.0)
            score = float(r.json()["response"].strip().split()[0])
            scores.append(min(max(score, 1.0), 10.0))
        except Exception:
            scores.append(5.0)
    return scores

async def benchmark_model(model: str) -> dict:
    print(f"  Benchmarking {model} ({len(BENCHMARK_PROMPTS)} prompts)...")
    async with httpx.AsyncClient() as client:
        results = [await run_single(client, model, p) for p in BENCHMARK_PROMPTS]
        scores = await score_batch(client, [(r["prompt"], r["response"]) for r in results])

    latencies = [r["elapsed_s"] for r in results if r["error"] is None]
    tps_vals = [r["tps"] for r in results if r["tps"] > 0]
    lat_sorted = sorted(latencies)
    return {
        "model": model,
        "n": len(BENCHMARK_PROMPTS),
        "errors": sum(1 for r in results if r["error"]),
        "p50_s": round(statistics.median(latencies), 2) if latencies else 0,
        "p95_s": round(lat_sorted[int(len(lat_sorted) * 0.95)], 2) if latencies else 0,
        "avg_tps": round(statistics.mean(tps_vals), 1) if tps_vals else 0,
        "avg_quality": round(statistics.mean(scores), 1),
    }

def write_csv(results: list[dict], path: str = "benchmark_results.csv") -> None:
    with open(path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=results[0].keys())
        writer.writeheader()
        writer.writerows(results)

def write_markdown(results: list[dict], path: str = "benchmark_results.md") -> None:
    lines = ["# Agent Benchmark Results\n",
             "| Model | p50 (s) | p95 (s) | Avg tok/s | Quality/10 | Errors |",
             "|---|---|---|---|---|---|"]
    for r in results:
        lines.append(f"| {r['model']} | {r['p50_s']} | {r['p95_s']} | {r['avg_tps']} | {r['avg_quality']} | {r['errors']} |")
    Path(path).write_text("\n".join(lines))

async def main():
    print(f"Running benchmark across {len(MODELS)} models × {len(BENCHMARK_PROMPTS)} prompts\n")
    results = [await benchmark_model(m) for m in MODELS]
    results.sort(key=lambda x: (-(x["avg_quality"]), x["p50_s"]))

    print("\n=== Benchmark Results ===")
    print(f"{'Model':<28} {'p50':>6} {'p95':>6} {'tok/s':>7} {'Quality':>9}")
    print("-" * 60)
    for r in results:
        print(f"{r['model']:<28} {r['p50_s']:>6.1f} {r['p95_s']:>6.1f} {r['avg_tps']:>7.0f} {r['avg_quality']:>8.1f}/10")

    write_csv(results)
    write_markdown(results)
    print("\nSaved: benchmark_results.csv, benchmark_results.md")

if __name__ == "__main__":
    asyncio.run(main())