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

chapter3/log-sanitization/agent.py

@@ -0,0 +1,409 @@
+"""
+Log Sanitization Agent using Local Ollama LLM
+"""
+
+import os
+import time
+import re
+import json
+from typing import List, Tuple, Dict, Optional
+from pathlib import Path
+import ollama
+
+try:
+    from dotenv import load_dotenv
+    load_dotenv()
+except ImportError:
+    pass
+
+from config import (
+    OLLAMA_MODEL, 
+    OLLAMA_TEMPERATURE,
+    SYSTEM_PROMPT, 
+    USER_PROMPT_TEMPLATE,
+    PII_DETECTION_SCHEMA,
+    OUTPUT_DIR
+)
+from metrics import PerformanceMetrics, MetricsCollector
+
+
+def _value_appears_in_text(value: str, text: str) -> bool:
+    """Return True if *value* appears as a substring of *text* (case-insensitive)."""
+    return value.lower() in text.lower()
+
+
+class LogSanitizationAgent:
+    """Agent for sanitizing logs using local Qwen3 0.6B model via Ollama"""
+    
+    def __init__(self, model: str = OLLAMA_MODEL):
+        """Initialize the sanitization agent.
+
+        Primary backend is the local Ollama model. If Ollama is unavailable
+        (not running / not reachable) and OPENROUTER_API_KEY is set, the agent
+        falls back to OpenRouter (default hosted model: openai/gpt-5.6-luna),
+        so the experiment still runs without a local model.
+        """
+        self.model = model
+        self.backend = "ollama"
+        self.metrics_collector = MetricsCollector(OUTPUT_DIR)
+
+        # Try the local Ollama backend first.
+        try:
+            self.client = ollama.Client()
+            models = self.client.list()
+            # models is a dict with 'models' key containing a list
+            if isinstance(models, dict) and 'models' in models:
+                available_models = [m.get('name', '') for m in models['models']]
+            else:
+                # If it's a direct list (older API versions)
+                available_models = [m.get('name', '') for m in models] if isinstance(models, list) else []
+
+            if not any(self.model in m for m in available_models):
+                print(f"⚠️  Model {self.model} not found. Pulling it now...")
+                self.client.pull(self.model)
+                print(f"✅ Model {self.model} pulled successfully")
+            else:
+                print(f"✅ Using model: {self.model}")
+
+        except Exception as e:
+            # Universal fallback: route through OpenRouter when Ollama is down.
+            openrouter_key = os.getenv("OPENROUTER_API_KEY")
+            if openrouter_key:
+                from openai import OpenAI
+
+                from agentbook.providers import resolve_backend
+                # 这里的回退条件是“本地 Ollama 连不上”，而非缺少凭证，
+                # 因此由本实验判定后再向注册表要一个 OpenRouter backend。
+                # 本地小模型（qwen3:0.6b 等）在 OpenRouter 上未必可用，
+                # substitute_unknown 让注册表替换成可用的默认模型。
+                backend = resolve_backend(
+                    "openrouter", model=self.model, api_key=openrouter_key
+                )
+                self.backend = "openrouter"
+                self.client = OpenAI(api_key=backend.api_key,
+                                     base_url=backend.base_url)
+                self.model = backend.model
+                print(f"⚠️  Ollama unavailable ({e}); "
+                      f"falling back to OpenRouter model: {self.model}")
+            else:
+                print(f"❌ Failed to connect to Ollama: {e}")
+                print("Please ensure Ollama is running: ollama serve, "
+                      "or set OPENROUTER_API_KEY as a fallback")
+                raise
+
+    def _chat_stream(self, messages):
+        """Yield content chunks from the active backend (Ollama or OpenRouter)."""
+        if self.backend == "ollama":
+            stream = self.client.chat(
+                model=self.model,
+                messages=messages,
+                stream=True,
+                format=PII_DETECTION_SCHEMA,  # Use structured output format
+                options={
+                    "temperature": OLLAMA_TEMPERATURE,
+                    "num_predict": 1000,
+                }
+            )
+            for chunk in stream:
+                yield chunk.get('message', {}).get('content', '')
+        else:
+            # 用与 Ollama 相同的 JSON Schema 强约束输出结构 (pii_items 数组),
+            # 避免模型自行发明字段名. strict 模式要求 additionalProperties=false.
+            strict_schema = dict(PII_DETECTION_SCHEMA)
+            strict_schema["additionalProperties"] = False
+            stream = self.client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                stream=True,
+                temperature=OLLAMA_TEMPERATURE,
+                response_format={
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "pii_detection",
+                        "strict": True,
+                        "schema": strict_schema,
+                    },
+                },
+                max_tokens=1000,
+            )
+            for chunk in stream:
+                if not chunk.choices:
+                    continue
+                yield chunk.choices[0].delta.content or ""
+    
+    def count_tokens(self, text: str) -> int:
+        """Estimate token count (rough approximation)"""
+        # Rough estimate: 1 token ≈ 4 characters for English text
+        # For more accurate counting, we'd need the actual tokenizer
+        return len(text) // 4
+    
+    def detect_pii(self, conversation_text: str) -> Tuple[List[str], Dict]:
+        """
+        Detect Level 3 PII in conversation text using local LLM
+        
+        Args:
+            conversation_text: Text to analyze
+        
+        Returns:
+            - List of detected PII values
+            - Performance metrics dictionary
+        """
+        # Prepare the prompt
+        user_prompt = USER_PROMPT_TEMPLATE.format(conversation_text=conversation_text)
+        
+        # Count input tokens
+        input_tokens = self.count_tokens(SYSTEM_PROMPT + user_prompt)
+        
+        # Measure prefill time (time to first token)
+        start_time = time.perf_counter()
+        
+        # Create messages for Ollama
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": user_prompt}
+        ]
+        
+        # Track first token time
+        first_token_time = None
+        output_tokens_count = 0
+        full_response = ""
+        
+        try:
+            # Use structured output with JSON schema (backend-agnostic stream)
+            print("\n   🧠 Analyzing (JSON): \033[90m", end="", flush=True)  # Gray color for JSON
+
+            for content in self._chat_stream(messages):
+                if first_token_time is None and content:
+                    first_token_time = time.perf_counter()
+
+                full_response += content
+                output_tokens_count += len(content) // 4  # Rough token estimate
+
+                # Stream the actual content
+                if content:
+                    print(content, end="", flush=True)
+
+            print("\033[0m")  # Reset color and new line
+            
+            end_time = time.perf_counter()
+            
+        except Exception as e:
+            print(f"\n❌ Error during PII detection: {e}")
+            return [], {}
+        
+        # Calculate performance metrics
+        prefill_time_ms = (first_token_time - start_time) * 1000 if first_token_time else 0
+        total_time_ms = (end_time - start_time) * 1000
+        output_time_ms = total_time_ms - prefill_time_ms
+        
+        prefill_speed = input_tokens / (prefill_time_ms / 1000) if prefill_time_ms > 0 else 0
+        output_speed = output_tokens_count / (output_time_ms / 1000) if output_time_ms > 0 else 0
+        
+        # Parse JSON response
+        pii_values = []
+        accepted_items = []
+        
+        try:
+            response_json = json.loads(full_response)
+            if not isinstance(response_json, dict):
+                return [], {}
+            
+            raw_items = response_json.get('pii_items')
+            if isinstance(raw_items, list):
+                for item in raw_items:
+                    if isinstance(item, dict):
+                        value = item.get('value')
+                        if value and isinstance(value, str) and _value_appears_in_text(value, conversation_text):
+                            pii_values.append(value)
+                            accepted_items.append(item)
+                    elif isinstance(item, str) and item and _value_appears_in_text(item, conversation_text):
+                        pii_values.append(item)
+                        accepted_items.append({"value": item})
+            else:
+                legacy_values = response_json.get('pii_values')
+                if isinstance(legacy_values, list):
+                    for pii in legacy_values:
+                        if pii and isinstance(pii, str):
+                            cleaned = pii.strip().strip('-').strip()
+                            if cleaned:
+                                pii_values.append(cleaned)
+
+        except json.JSONDecodeError as e:
+            print(f"\n   ⚠️  Failed to parse JSON response: {e}")
+            # Fallback to simple line splitting if JSON parsing fails
+            pii_values = [line.strip() for line in full_response.split('\n') if line.strip()]
+            accepted_items = []
+
+        metrics = {
+            'input_tokens': input_tokens,
+            'output_tokens': output_tokens_count,
+            'prefill_time_ms': prefill_time_ms,
+            'output_time_ms': output_time_ms,
+            'total_time_ms': total_time_ms,
+            'prefill_speed_tps': prefill_speed,
+            'output_speed_tps': output_speed,
+            'pii_items_found': len(pii_values),
+            'pii_items': accepted_items
+        }
+        
+        return pii_values, metrics
+    
+    def sanitize_text(self, text: str, pii_values: List[str]) -> Tuple[str, int]:
+        """
+        Replace PII values with [REDACTED] in the text
+        
+        Returns:
+            - Sanitized text
+            - Number of replacements made
+        """
+        sanitized = text
+        replacements = 0
+        
+        for pii_value in pii_values:
+            # Escape special regex characters in PII value
+            escaped_value = re.escape(pii_value)
+            # Count occurrences before replacement
+            occurrences = len(re.findall(escaped_value, sanitized, re.IGNORECASE))
+            # Replace all occurrences
+            sanitized = re.sub(escaped_value, '[REDACTED]', sanitized, flags=re.IGNORECASE)
+            replacements += occurrences
+        
+        return sanitized, replacements
+    
+    def sanitize_conversation(
+        self, 
+        conversation: Dict,
+        test_id: str = "unknown"
+    ) -> Dict:
+        """
+        Sanitize a single conversation and collect metrics
+        
+        Returns:
+            Dictionary with sanitized conversation and metrics
+        """
+        # Format conversation text
+        conv_text = self.format_conversation(conversation)
+        conv_id = conversation.get('conversation_id', 'unknown')
+        
+        print(f"🔍 Processing conversation: {conv_id}")
+        
+        # Detect PII
+        pii_values, perf_metrics = self.detect_pii(conv_text)
+        accepted_items = perf_metrics.get('pii_items', [])
+        
+        if pii_values:
+            print(f"   ✅ Found {len(pii_values)} PII items:")
+            for pii in pii_values:
+                print(f"      - {pii}")
+        else:
+            print("   ⚠️  No PII items detected")
+        
+        # Sanitize the text
+        sanitized_text, replacements = self.sanitize_text(conv_text, pii_values)
+        
+        # Create performance metric. detect_pii() returns an empty metrics dict
+        # when the LLM backend fails (e.g. Ollama not running) — fall back to
+        # zeros so one failed conversation doesn't crash the whole batch.
+        metric = PerformanceMetrics(
+            test_id=test_id,
+            conversation_id=conv_id,
+            input_text_length=len(conv_text),
+            input_tokens=perf_metrics.get('input_tokens', 0),
+            prefill_time_ms=perf_metrics.get('prefill_time_ms', 0),
+            output_time_ms=perf_metrics.get('output_time_ms', 0),
+            total_time_ms=perf_metrics.get('total_time_ms', 0),
+            output_tokens=perf_metrics.get('output_tokens', 0),
+            prefill_speed_tps=perf_metrics.get('prefill_speed_tps', 0),
+            output_speed_tps=perf_metrics.get('output_speed_tps', 0),
+            pii_items_found=perf_metrics.get('pii_items_found', 0),
+            replacements_made=replacements,
+            sanitized_text_length=len(sanitized_text)
+        )
+        
+        self.metrics_collector.add_metric(metric)
+        
+        return {
+            'conversation_id': conv_id,
+            'original_length': len(conv_text),
+            'sanitized_length': len(sanitized_text),
+            'pii_found': pii_values,
+            'replacements_made': replacements,
+            'sanitized_text': sanitized_text,
+            'pii_items': accepted_items,
+            'metrics': metric.to_dict()
+        }
+    
+    def format_conversation(self, conversation: Dict) -> str:
+        """Format conversation dictionary into text"""
+        lines = []
+        lines.append(f"Conversation ID: {conversation.get('conversation_id', 'unknown')}")
+        lines.append(f"Timestamp: {conversation.get('timestamp', 'unknown')}")
+        lines.append("-" * 50)
+        
+        messages = conversation.get('messages', [])
+        for msg in messages:
+            role = msg.get('role', 'unknown').upper()
+            content = msg.get('content', '')
+            lines.append(f"{role}: {content}")
+            lines.append("")  # Empty line between messages
+        
+        return "\n".join(lines)
+    
+    def save_sanitized_log(self, test_id: str, results: List[Dict]):
+        """Save sanitized logs to output directory"""
+        output_file = OUTPUT_DIR / f"{test_id}_sanitized.txt"
+        summary_file = OUTPUT_DIR / f"{test_id}_summary.json"
+        
+        # Save sanitized text
+        with open(output_file, 'w') as f:
+            for result in results:
+                f.write(f"\n{'='*60}\n")
+                f.write(f"Conversation: {result['conversation_id']}\n")
+                f.write(f"{'='*60}\n")
+                f.write(result['sanitized_text'])
+                f.write("\n")
+        
+        # Save summary
+        summary = {
+            'test_id': test_id,
+            'total_conversations': len(results),
+            'total_pii_found': sum(len(r['pii_found']) for r in results),
+            'total_replacements': sum(r['replacements_made'] for r in results),
+            'conversations': [
+                {
+                    'conversation_id': r['conversation_id'],
+                    'pii_count': len(r['pii_found']),
+                    'replacements': r['replacements_made']
+                }
+                for r in results
+            ]
+        }
+        
+        with open(summary_file, 'w') as f:
+            json.dump(summary, f, indent=2)
+        
+        print(f"✅ Sanitized log saved to: {output_file}")
+        print(f"✅ Summary saved to: {summary_file}")
+    
+    def process_test_case(self, test_id: str, conversations: List[Dict]) -> List[Dict]:
+        """Process all conversations in a test case"""
+        results = []
+        
+        print(f"\n{'='*60}")
+        print(f"Processing Test Case: {test_id}")
+        print(f"Total Conversations: {len(conversations)}")
+        print(f"{'='*60}")
+        
+        for i, conv in enumerate(conversations, 1):
+            print(f"\n[{i}/{len(conversations)}] ", end="")
+            result = self.sanitize_conversation(conv, test_id)
+            results.append(result)
+        
+        # Save results
+        self.save_sanitized_log(test_id, results)
+        
+        # Save metrics
+        self.metrics_collector.save_metrics()
+        self.metrics_collector.print_summary()
+        
+        return results