130. Capstone — RAM Manager With LLM

🚀 Aryan’s RAM manager is almost complete. The integration layer that wires psutil → LLM prompt → parse response → save report has several bugs: wrong async pattern, missing error handling, and a report-building function that loses data.

💡 Fun fact: The async with pattern for HTTP clients is not just style — it ensures connection pool cleanup. The httpx.AsyncClient manages a pool of TCP connections under the hood. Without async with, those sockets stay open until the garbage collector notices, which can cause “too many open files” errors on busy servers. Python’s context manager protocol was designed exactly for this.

⚠️ Watch out: Forgetting await on an async function is one of the most common Python bugs. Without await, you get a coroutine object stored in a variable — it looks like it worked, but the function body never ran. Python 3.11+ will warn you about unawaited coroutines, but older versions are silent. Always check: if you call an async def function, you must await it.

🤔 Think about it: Why does dataclasses.asdict() exist when you could just write {'timestamp': snap.timestamp, 'percent': snap.percent, ...} manually? What happens when you add a new field to RamSnapshot six months from now?

Learning objectives

• Use async with to manage HTTP client lifecycle properly
• Always await async function calls
• Use dataclasses.asdict() to serialize nested dataclass trees to JSON-safe dicts
• Use snapshot.timestamp consistently instead of regenerating with datetime.now()
• Write complete session pipelines combining psutil, async HTTP, and file I/O

Key concepts

• async with client — context-managed HTTP client lifecycle
• await — execute coroutine and wait for result
• dataclasses.asdict() — recursive dataclass serialization
• Path.replace() — atomic file write
• asyncio.run() — sync entry point for async code

Try it

Concept detail

Capstone: Full RAM Manager Architecture

This exercise combines all story_3 concepts:

psutil.virtual_memory()        → RamSnapshot (dataclass)
     ↓
async get_llm_advice()         → LlmAdvice (TypedDict)
AsyncClient.post()
     ↓
build_report()                 → dict (JSON-serializable)
dataclasses.asdict()
     ↓
atomic file write              → session_YYYY-MM-DD.json
Path.replace()

Key patterns

Async HTTP with proper cleanup:

async with httpx.AsyncClient() as client:
    resp = await client.post(url, ...)

Always await coroutines:

# Wrong (stores coroutine object)
result = some_async_fn()

# Right
result = await some_async_fn()

Serialize nested dataclasses:

from dataclasses import asdict
# Recursively converts RamSnapshot + nested ProcessInfo
d = asdict(snapshot)
d['advice'] = advice
json.dumps(d)  # fully serializable

Atomic write to prevent corruption:

tmp = out.with_suffix('.tmp')
tmp.write_text(json.dumps(data, indent=2))
tmp.replace(out)   # atomic on POSIX

Solution

import asyncio
import json
import psutil
from dataclasses import dataclass, asdict, field
from datetime import datetime
from pathlib import Path
from typing import TypedDict
from unittest.mock import AsyncMock, MagicMock

# ── Models ──────────────────────────────────────────────────────────────────

@dataclass
class ProcessInfo:
    name: str
    pid: int
    rss_mb: float

@dataclass
class RamSnapshot:
    timestamp: str
    total_gb: float
    used_gb: float
    percent: float
    top_processes: list = field(default_factory=list)

class LlmAdvice(TypedDict):
    summary: str
    recommendations: list

# ── Async HTTP client stand-in (replaced by mock in tests) ───────────────────

class AsyncHttpClient:
    """Stand-in for httpx.AsyncClient — replaced by mock in tests."""
    async def __aenter__(self):
        return self
    async def __aexit__(self, *args):
        pass
    async def post(self, url, headers=None, json=None, timeout=None):
        raise RuntimeError("Real network call — replace with mock in tests")

_http_client_factory = AsyncHttpClient

# ── Core logic ───────────────────────────────────────────────────────────────

def take_snapshot(top_n: int = 5) -> RamSnapshot:
    """Capture current RAM state."""
    mem = psutil.virtual_memory()
    procs = []
    for proc in psutil.process_iter(['pid', 'name', 'memory_info']):
        try:
            procs.append(ProcessInfo(
                name=proc.info['name'],
                pid=proc.info['pid'],
                rss_mb=round(proc.info['memory_info'].rss / 1e6, 1),
            ))
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
    top = sorted(procs, key=lambda p: p.rss_mb, reverse=True)[:top_n]
    return RamSnapshot(
        timestamp=datetime.now().isoformat(),
        total_gb=round(mem.total / 1e9, 1),
        used_gb=round(mem.used / 1e9, 1),
        percent=mem.percent,
        top_processes=top,
    )

async def get_llm_advice(snapshot: RamSnapshot, api_key: str) -> LlmAdvice:
    """Ask the LLM for RAM management advice."""
    lines = [f'  {p.rss_mb} MB  {p.name} (PID {p.pid})' for p in snapshot.top_processes]
    prompt = (
        f'RAM usage: {snapshot.percent}% ({snapshot.used_gb}/{snapshot.total_gb} GB)\n'
        f'Top processes:\n' + '\n'.join(lines) +
        '\n\nGive a 1-sentence summary and 3 bullet recommendations as JSON: '
        '{"summary": "...", "recommendations": ["...", "...", "..."]}'
    )
    # FIX 1: use async with for proper context-managed client lifecycle
    async with _http_client_factory() as client:
        resp = await client.post(
            'https://api.anthropic.com/v1/messages',
            headers={
                'x-api-key': api_key,
                'anthropic-version': '2023-06-01',
                'content-type': 'application/json',
            },
            json={
                'model': 'claude-3-haiku-20240307',
                'max_tokens': 256,
                'messages': [{'role': 'user', 'content': prompt}],
            },
            timeout=30,
        )
    resp.raise_for_status()
    raw = resp.json()['content'][0]['text']
    import re
    match = re.search(r'\{.*\}', raw, re.DOTALL)
    return json.loads(match.group()) if match else {'summary': raw, 'recommendations': []}

def build_report(snapshot: RamSnapshot, advice: LlmAdvice) -> dict:
    """Build a JSON-serializable report dict."""
    # FIX 2+3: asdict() recurses into nested dataclasses and uses snapshot.timestamp
    d = asdict(snapshot)
    d['advice'] = advice
    return d

async def run_session(api_key: str, report_dir: Path):
    """One full monitoring session: snapshot → LLM → save."""
    report_dir.mkdir(parents=True, exist_ok=True)

    snapshot = take_snapshot()
    advice = await get_llm_advice(snapshot, api_key)  # FIX 4: await the coroutine

    report = build_report(snapshot, advice)
    out = report_dir / f'session_{snapshot.timestamp[:10]}.json'

    existing = json.loads(out.read_text()) if out.exists() else []
    existing.append(report)
    tmp = out.with_suffix('.tmp')
    tmp.write_text(json.dumps(existing, indent=2))
    tmp.replace(out)
    return report

Tests

import asyncio
import json
import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

def _mock_mem(total, used, available, percent):
    m = MagicMock()
    m.total = total
    m.used = used
    m.available = available
    m.percent = percent
    return m

def _mock_proc(name, pid, rss):
    p = MagicMock()
    p.info = {'name': name, 'pid': pid, 'memory_info': MagicMock(rss=rss)}
    return p

def _make_client_factory(response_text):
    """Return a factory that injects a mock async HTTP client."""
    mock_resp = MagicMock()
    mock_resp.status_code = 200
    mock_resp.raise_for_status.return_value = None
    mock_resp.json.return_value = {
        'content': [{'text': response_text}]
    }
    mock_client = AsyncMock()
    mock_client.post = AsyncMock(return_value=mock_resp)
    cm = MagicMock()
    cm.__aenter__ = AsyncMock(return_value=mock_client)
    cm.__aexit__ = AsyncMock(return_value=None)
    def factory():
        return cm
    return factory

def test_build_report_includes_top_processes():
    snap = RamSnapshot(
        timestamp='2024-01-15T10:00:00',
        total_gb=8.0, used_gb=5.0, percent=62.5,
        top_processes=[ProcessInfo('chrome', 1, 500.0), ProcessInfo('python', 2, 100.0)],
    )
    advice = {'summary': 'High usage', 'recommendations': ['Kill Chrome']}
    report = build_report(snap, advice)
    assert 'top_processes' in report, 'Report must include top_processes'
    assert len(report['top_processes']) == 2
    assert report['top_processes'][0]['name'] == 'chrome'

def test_build_report_uses_snapshot_timestamp():
    snap = RamSnapshot(
        timestamp='2024-01-15T10:00:00',
        total_gb=8.0, used_gb=5.0, percent=62.5,
    )
    advice = {'summary': 'ok', 'recommendations': []}
    report = build_report(snap, advice)
    assert report['timestamp'] == '2024-01-15T10:00:00', (
        'Use snapshot.timestamp, not datetime.now()'
    )

def test_build_report_includes_advice():
    snap = RamSnapshot('2024-01-15T10:00:00', 8.0, 5.0, 62.5)
    advice = {'summary': 'High', 'recommendations': ['tip1', 'tip2']}
    report = build_report(snap, advice)
    assert report['advice']['summary'] == 'High'
    assert report['advice']['recommendations'] == ['tip1', 'tip2']

def test_run_session_saves_report():
    mem = _mock_mem(8e9, 5e9, 3e9, 62.5)
    procs = [_mock_proc('chrome', 1, 500_000_000)]
    llm_text = '{"summary": "High RAM", "recommendations": ["Kill Chrome"]}'
    factory = _make_client_factory(llm_text)

    old_factory = globals()['_http_client_factory']
    globals()['_http_client_factory'] = factory
    try:
        with patch('psutil.virtual_memory', return_value=mem), \
             patch('psutil.process_iter', return_value=procs):
            with tempfile.TemporaryDirectory() as tmp:
                report = asyncio.run(run_session('fake-key', Path(tmp)))
                assert report is not None
                assert 'advice' in report
                files = list(Path(tmp).glob('session_*.json'))
                assert len(files) == 1
    finally:
        globals()['_http_client_factory'] = old_factory

def test_run_session_awaits_llm():
    """Ensure get_llm_advice is awaited (not stored as coroutine object)."""
    mem = _mock_mem(8e9, 5e9, 3e9, 62.5)
    procs = [_mock_proc('vim', 1, 50_000_000)]
    llm_text = '{"summary": "ok", "recommendations": []}'
    factory = _make_client_factory(llm_text)

    old_factory = globals()['_http_client_factory']
    globals()['_http_client_factory'] = factory
    try:
        with patch('psutil.virtual_memory', return_value=mem), \
             patch('psutil.process_iter', return_value=procs):
            with tempfile.TemporaryDirectory() as tmp:
                report = asyncio.run(run_session('k', Path(tmp)))
        # If get_llm_advice wasn't awaited, advice would be a coroutine object
        assert isinstance(report.get('advice'), dict), (
            'advice must be a dict — did you forget to await get_llm_advice()?'
        )
    finally:
        globals()['_http_client_factory'] = old_factory

def test_get_llm_advice_uses_async_with():
    """get_llm_advice must use async with for the client context manager."""
    entry_count = [0]
    mock_resp = MagicMock()
    mock_resp.status_code = 200
    mock_resp.raise_for_status.return_value = None
    mock_resp.json.return_value = {
        'content': [{'text': '{"summary": "test", "recommendations": []}'}]
    }
    mock_client = AsyncMock()
    mock_client.post = AsyncMock(return_value=mock_resp)
    cm = MagicMock()
    async def fake_aenter():
        entry_count[0] += 1
        return mock_client
    cm.__aenter__ = fake_aenter
    cm.__aexit__ = AsyncMock(return_value=None)

    old_factory = globals()['_http_client_factory']
    globals()['_http_client_factory'] = lambda: cm
    try:
        snap = RamSnapshot('2024-01-01T00:00:00', 8.0, 4.0, 50.0, [])
        asyncio.run(get_llm_advice(snap, 'key'))
        assert entry_count[0] == 1, (
            'get_llm_advice must use async with client — __aenter__ not called'
        )
    finally:
        globals()['_http_client_factory'] = old_factory

Resources

• httpx AsyncClient — Official Docs
• asyncio — Python Docs
• dataclasses.asdict — Python Docs
• Python Async IO in Practice — Real Python
• Python asyncio explained — ArjanCodes (YouTube)