ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/log-sanitization/tests/conftest.py

@@ -0,0 +1,22 @@
+"""Shared bootstrap for log-sanitization regression tests."""
+
+import sys
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+try:
+    import ollama  # noqa: F401
+except ImportError:
+    ollama_stub = ModuleType("ollama")
+    ollama_stub.Client = object
+    sys.modules["ollama"] = ollama_stub
+
+try:
+    import dotenv  # noqa: F401
+except ImportError:
+    sys.modules["dotenv"] = SimpleNamespace(load_dotenv=lambda *args, **kwargs: None)

chapter3/log-sanitization/tests/manual/loader_debug.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+"""Debug script to test loading conversations"""
+
+import sys
+from pathlib import Path
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from test_loader import TestCaseLoader
+
+def main():
+    loader = TestCaseLoader()
+    
+    # Get all test cases
+    print("Getting all test cases...")
+    all_cases = loader.get_all_test_cases()
+    print(f"Found {len(all_cases)} test cases")
+    
+    # Get Layer 3 test cases
+    layer3_cases = loader.get_layer3_test_cases()
+    print(f"Found {len(layer3_cases)} Layer 3 test cases")
+    
+    if layer3_cases:
+        # Try to load the first one
+        first_case = layer3_cases[0]
+        print(f"\nTrying to load: {first_case['test_id']}")
+        
+        conversations = loader.get_test_case_conversations(first_case['test_id'])
+        
+        if conversations:
+            print(f"Successfully loaded {len(conversations)} conversations")
+            # Print first conversation snippet
+            if conversations[0]['messages']:
+                print(f"First message: {conversations[0]['messages'][0]['content'][:100]}...")
+        else:
+            print("Failed to load conversations")
+
+if __name__ == "__main__":
+    main()

chapter3/log-sanitization/tests/test_authorization_basic.py

@@ -0,0 +1,40 @@
+"""Authorization: Basic credentials must be redacted like Bearer tokens."""
+from regex_sanitizer import sanitize
+
+
+def test_authorization_basic_redacted():
+    cred = "dXNlcjpwYXNzd29yZA=="
+    text, hits = sanitize(f"Authorization: Basic {cred}")
+    assert cred not in text
+    assert "[REDACTED_BASIC_AUTH]" in text
+    assert any(h["category"] == "basic_auth" for h in hits)
+
+
+def test_authorization_basic_case_insensitive():
+    cred = "YWRtaW46c2VjcmV0"
+    text, hits = sanitize(f"authorization: basic {cred}")
+    assert cred not in text
+    assert "[REDACTED_BASIC_AUTH]" in text
+
+
+def test_bearer_still_redacted():
+    token = "aaaaaaaaaaaaaaaaaaaa"
+    text, hits = sanitize(f"Authorization: Bearer {token}")
+    assert token not in text
+    assert "[REDACTED_BEARER_TOKEN]" in text
+    assert any(h["category"] == "bearer_token" for h in hits)
+
+
+def test_english_basic_prose_not_redacted():
+    prose = "Basic knowledge of Python is required."
+    text, hits = sanitize(prose)
+    assert text == prose
+    assert not any(h["category"] == "basic_auth" for h in hits)
+
+
+def test_www_authenticate_basic_realm_not_treated_as_credential():
+    # Challenge header names the scheme; it does not carry the password blob.
+    line = 'WWW-Authenticate: Basic realm="api"'
+    text, hits = sanitize(line)
+    assert text == line
+    assert not any(h["category"] == "basic_auth" for h in hits)

chapter3/log-sanitization/tests/test_github_pat.py

@@ -0,0 +1,18 @@
+"""Regression: fine-grained github_pat_* tokens must be redacted."""
+from regex_sanitizer import sanitize
+
+
+def test_github_pat_fine_grained_redacted():
+    token = "github_pat_" + "A" * 20 + "_" + "B" * 40
+    text, hits = sanitize(f"Authorization: {token}")
+    assert token not in text
+    assert "[REDACTED_GITHUB_TOKEN]" in text
+    assert any(h["category"] == "github_token" for h in hits)
+
+
+def test_classic_github_token_still_redacted():
+    token = "gh" + "p_" + "x" * 36
+    text, hits = sanitize(token)
+    assert token not in text
+    assert "[REDACTED_GITHUB_TOKEN]" in text
+    assert any(h["category"] == "github_token" for h in hits)

chapter3/log-sanitization/tests/test_null_pii_values.py

@@ -0,0 +1,33 @@
+"""detect_pii must treat JSON null pii_values like an empty list."""
+import json
+from unittest.mock import patch
+
+from agent import LogSanitizationAgent
+
+
+def test_null_pii_values_returns_empty_list():
+    agent = object.__new__(LogSanitizationAgent)
+    agent.count_tokens = lambda text: len(text) // 4
+    agent._chat_stream = lambda messages: iter(
+        [json.dumps({"pii_values": None})]
+    )
+
+    with patch("builtins.print"):
+        pii_values, metrics = agent.detect_pii("Alice phone 555-0100")
+
+    assert pii_values == []
+    assert metrics["pii_items_found"] == 0
+
+
+def test_list_pii_values_still_cleaned():
+    agent = object.__new__(LogSanitizationAgent)
+    agent.count_tokens = lambda text: len(text) // 4
+    agent._chat_stream = lambda messages: iter(
+        [json.dumps({"pii_values": ["  Alice  ", "-", "Bob"]})]
+    )
+
+    with patch("builtins.print"):
+        pii_values, metrics = agent.detect_pii("text")
+
+    assert pii_values == ["Alice", "Bob"]
+    assert metrics["pii_items_found"] == 2

chapter3/log-sanitization/tests/test_quoted_secret_spaces.py

@@ -0,0 +1,23 @@
+"""Regression: quoted secret assignments must redact the full value, including spaces."""
+from regex_sanitizer import sanitize
+
+
+def test_double_quoted_password_with_spaces():
+    text, hits = sanitize('password="hunter 2 with spaces"')
+    assert "hunter" not in text
+    assert "spaces" not in text
+    assert "[REDACTED_SECRET]" in text
+    assert any(h["category"] == "secret_assignment" for h in hits)
+
+
+def test_single_quoted_password_with_spaces():
+    text, hits = sanitize("api_key='test-api-key with spaces'")
+    assert "test-api-key" not in text
+    assert "ghi" not in text
+    assert "[REDACTED_SECRET]" in text
+
+
+def test_unquoted_secret_still_redacted():
+    text, hits = sanitize("password=hunter2xyz")
+    assert "hunter2xyz" not in text
+    assert "[REDACTED_SECRET]" in text

chapter3/log-sanitization/tests/test_sanitize_conversation.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+"""Regression tests for sanitize_conversation() in agent.py.
+
+Bug: detect_pii() catches any backend exception and returns ([], {}), but
+sanitize_conversation then subscripted the empty metrics dict
+(perf_metrics['input_tokens']) -> KeyError that killed the whole batch.
+Fixed with .get(..., 0) defaults so a dead backend degrades gracefully.
+"""
+
+import pytest
+
+from agent import LogSanitizationAgent
+from metrics import MetricsCollector
+
+
+def _make_agent(client, tmp_path):
+    """Build an agent without __init__ (which requires a live Ollama)."""
+    ag = LogSanitizationAgent.__new__(LogSanitizationAgent)
+    ag.model = "qwen3:0.6b"
+    ag.backend = "ollama"
+    ag.metrics_collector = MetricsCollector(tmp_path)
+    ag.client = client
+    return ag
+
+
+CONV = {
+    "conversation_id": "demo_001",
+    "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
+}
+
+
+class _DeadClient:
+    def chat(self, **kwargs):
+        raise ConnectionError("[test] Ollama server is not running")
+
+
+class _FakeClient:
+    """Mimics ollama.Client.chat(stream=True) chunk shape."""
+
+    def __init__(self, payload: str):
+        self.payload = payload
+
+    def chat(self, **kwargs):
+        return [{"message": {"content": self.payload}}]
+
+
+def test_dead_backend_returns_result_with_zero_metrics(tmp_path):
+    ag = _make_agent(_DeadClient(), tmp_path)
+    result = ag.sanitize_conversation(CONV, "t1")  # must not raise
+    assert result["pii_found"] == []
+    assert result["replacements_made"] == 0
+    assert result["metrics"]["input_tokens"] == 0
+    assert result["metrics"]["pii_items_found"] == 0
+
+
+def test_working_backend_still_detects_pii(tmp_path):
+    ag = _make_agent(_FakeClient('{"pii_values": ["123-45-6789"]}'), tmp_path)
+    result = ag.sanitize_conversation(CONV, "t1")
+    assert result["pii_found"] == ["123-45-6789"]
+    assert result["replacements_made"] == 1
+    assert "[REDACTED]" in result["sanitized_text"]
+    assert result["metrics"]["pii_items_found"] == 1
+
+
+def test_working_backend_with_structured_pii_items(tmp_path):
+    conv = {
+        "conversation_id": "demo_002",
+        "messages": [
+            {
+                "role": "user",
+                "content": "My SSN is 000-00-0000 and my card is 0000-0000-0000-0000.",
+            }
+        ],
+    }
+    payload = (
+        '{"pii_items": ['
+        '{"type": "social_security_number", "value": "000-00-0000"}, '
+        '{"type": "credit_card_number", "value": "0000-0000-0000-0000"}'
+        ']}'
+    )
+    ag = _make_agent(_FakeClient(payload), tmp_path)
+    result = ag.sanitize_conversation(conv, "t2")
+    assert result["pii_found"] == ["000-00-0000", "0000-0000-0000-0000"]
+    assert result["replacements_made"] == 2
+    assert result["sanitized_text"].count("[REDACTED]") == 2
+    assert result["metrics"]["pii_items_found"] == 2
+
+
+def test_structured_pii_items_preserve_original_value(tmp_path):
+    conv = {
+        "conversation_id": "demo_003",
+        "messages": [
+            {
+                "role": "user",
+                "content": "The secret is -abc- and the password is   p4ss  .",
+            }
+        ],
+    }
+    payload = (
+        '{"pii_items": ['
+        '{"type": "secret", "value": "-abc-"}, '
+        '{"type": "password", "value": "  p4ss  "}'
+        ']}'
+    )
+    ag = _make_agent(_FakeClient(payload), tmp_path)
+    result = ag.sanitize_conversation(conv, "t3")
+    assert result["pii_found"] == ["-abc-", "  p4ss  "]
+    assert result["replacements_made"] == 2
+    assert result["sanitized_text"].count("[REDACTED]") == 2
+
+
+def test_pii_items_metric_excludes_rejected_and_malformed_items(tmp_path):
+    conv = {
+        "conversation_id": "demo_004",
+        "messages": [
+            {
+                "role": "user",
+                "content": "My email is alice@example.com.",
+            }
+        ],
+    }
+    payload = (
+        '{"pii_items": ['
+        '{"type": "email", "value": "alice@example.com"}, '
+        '{"type": "ssn", "value": "999-99-9999"}, '
+        '{"type": "unknown", "value": ""}, '
+        '{"type": "broken"}, '
+        '"just a string"'
+        ']}'
+    )
+    ag = _make_agent(_FakeClient(payload), tmp_path)
+    result = ag.sanitize_conversation(conv, "t4")
+    assert result["pii_found"] == ["alice@example.com"]
+    assert result["replacements_made"] == 1
+    assert result["metrics"]["pii_items_found"] == 1
+    items = result["pii_items"]
+    assert len(items) == 1
+    assert items[0]["type"] == "email"
+    assert items[0]["value"] == "alice@example.com"

chapter3/log-sanitization/tests/test_truncated_pem.py

@@ -0,0 +1,36 @@
+"""Truncated PEM (BEGIN without END) must be redacted, not leaked."""
+
+from regex_sanitizer import sanitize
+
+PEM_HEADER = "-----BEGIN " + "RSA PRIVATE KEY-----\n"
+PEM_FOOTER = "-----END " + "RSA PRIVATE KEY-----"
+
+
+def test_truncated_rsa_pem_without_end_redacted():
+    blob = (
+        PEM_HEADER +
+        "MIIEowIBAAKCAQEA0Z3VS5JJcds3xfn/ygWyF6PZGFw7\n"
+        "ygWyF6PZGFw7morekeymaterialHERE"
+    )
+    text, hits = sanitize(f"key dump:\n{blob}\n")
+    assert "MIIEowIBAAKCAQEA" not in text
+    assert "[REDACTED_PRIVATE_KEY]" in text
+    assert any(h["category"] == "private_key" for h in hits)
+
+
+def test_complete_pem_still_redacted():
+    blob = (
+        PEM_HEADER +
+        "MIIEowIBAAKCAQEA0Z3VS5JJcds3xfn/ygWyF6PZGFw7\n" +
+        PEM_FOOTER
+    )
+    text, hits = sanitize(blob)
+    assert "MIIEowIBAAKCAQEA" not in text
+    assert text.strip() == "[REDACTED_PRIVATE_KEY]"
+    assert any(h["category"] == "private_key" for h in hits)
+
+
+def test_non_key_text_unchanged():
+    text, hits = sanitize("no secrets here, only BEGIN of a story")
+    assert text == "no secrets here, only BEGIN of a story"
+    assert hits == []

chapter3/log-sanitization/tests/test_url_credential_colon_slash.py

@@ -0,0 +1,23 @@
+"""Regression: URL passwords containing ':' or '/' must be fully redacted."""
+from regex_sanitizer import sanitize
+
+
+def test_password_with_slash_redacted():
+    text, hits = sanitize("DATABASE_URL=postgres://alice:a/b@db.example:5432/app")
+    assert "a/b" not in text
+    assert "[REDACTED_URL_CRED]" in text
+    assert any(h["category"] == "url_credential" for h in hits)
+
+
+def test_password_with_colon_redacted():
+    text, hits = sanitize("redis://default:foo:bar@10.0.0.1:6379/0")
+    assert "foo:bar" not in text
+    assert "[REDACTED_URL_CRED]" in text
+    assert any(h["category"] == "url_credential" for h in hits)
+
+
+def test_simple_password_still_redacted():
+    text, hits = sanitize("postgres://alice:secret@db.example:5432/app")
+    assert "secret" not in text
+    assert "[REDACTED_URL_CRED]" in text
+    assert any(h["category"] == "url_credential" for h in hits)

chapter3/log-sanitization/tests/test_url_credential_empty_user.py

@@ -0,0 +1,22 @@
+"""Regression: URL credentials with an empty username must be redacted."""
+from regex_sanitizer import sanitize
+
+
+def test_redis_empty_user_password_redacted():
+    text, hits = sanitize("redis://:secretpass@10.0.0.1:6379/0")
+    assert "secretpass" not in text
+    assert "[REDACTED_URL_CRED]" in text
+    assert any(h["category"] == "url_credential" for h in hits)
+
+
+def test_postgres_empty_user_password_redacted():
+    text, hits = sanitize("DATABASE_URL=postgres://:hunter2@localhost:5432/db")
+    assert "hunter2" not in text
+    assert "[REDACTED_URL_CRED]" in text
+    assert any(h["category"] == "url_credential" for h in hits)
+
+
+def test_named_user_password_still_redacted():
+    text, hits = sanitize("redis://default:secretpass@10.0.0.1:6379/0")
+    assert "secretpass" not in text
+    assert "[REDACTED_URL_CRED]" in text