Wire up real LLM integration via HF Inference API

- Customer simulator and agent now use Llama 3.1 8B via HF Inference API
when HF_TOKEN is set in .env (gitignored, never pushed)
- Graceful fallback: if API credits deplete (402), auto-falls back to
rule-based simulation for remaining calls
- HFAgent uses Llama 3.1 (not Qwen which isn't available on free tier)
- A/B test supports --mode llm|rule flag, shows sample conversations
- Layer 1 train.py supports --llm-agent flag for real LLM evaluation
- Added python-dotenv + datasets to dependencies
- All .env loading via dotenv, keys never touch git
- 31 tests passing

https://claude.ai/code/session_01DPirJ78YYN4fJUvUFJ5D6V

layer1/train.py

@@ -20,6 +20,10 @@ import logging
 import sys
 import os
 
+# Auto-load .env for HF_TOKEN
+from dotenv import load_dotenv
+load_dotenv(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".env"))
+
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 from layer1.grpo_trainer import (
@@ -30,23 +34,35 @@ from layer1.grpo_trainer import (
     build_meta_prompt,
 )
 from layer2.customer_sim import CustomerPersona, CustomerSimulator
+from layer2.hf_agent import HFAgent
 from personas.generate_personas import generate_personas
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
 logger = logging.getLogger(__name__)
 
 
-def load_evaluator(hf_token: str | None = None) -> PromptEvaluator:
-    """Load personas and create the evaluator."""
+def load_evaluator(hf_token: str | None = None, use_llm_agent: bool = False) -> PromptEvaluator:
+    """Load personas and create the evaluator with optional LLM agent."""
+    token = hf_token or os.environ.get("HF_TOKEN")
     personas_data = generate_personas(100)
     personas = [CustomerPersona(**p) for p in personas_data]
-    simulator = CustomerSimulator(hf_token=hf_token)
-    return PromptEvaluator(personas=personas, simulator=simulator)
+    simulator = CustomerSimulator(hf_token=token)
+
+    agent_fn = None
+    if use_llm_agent and token:
+        agent = HFAgent(hf_token=token)
+        if agent.is_llm_available:
+            agent_fn = agent
+            logger.info("Using LLM agent (Llama 3.1 8B)")
+        else:
+            logger.warning("LLM agent not available, using rule-based fallback")
+
+    return PromptEvaluator(personas=personas, simulator=simulator, agent_fn=agent_fn)
 
 
 def run_mock(args):
     """Run mock optimization with hand-written prompts."""
-    evaluator = load_evaluator(args.hf_token)
+    evaluator = load_evaluator(args.hf_token, use_llm_agent=args.llm_agent)
     optimizer = MockPromptOptimizer(evaluator)
     result = optimizer.optimize(num_episodes_per_prompt=args.episodes)
 
@@ -66,7 +82,7 @@ def run_mock(args):
 
 def run_train(args):
     """Run full GRPO training (requires GPU)."""
-    evaluator = load_evaluator(args.hf_token)
+    evaluator = load_evaluator(args.hf_token, use_llm_agent=args.llm_agent)
     config = GRPOConfig(
         num_training_steps=args.steps,
         episodes_per_candidate=args.episodes,
@@ -89,7 +105,7 @@ def run_train(args):
 
 def run_eval(args):
     """Evaluate a single prompt."""
-    evaluator = load_evaluator(args.hf_token)
+    evaluator = load_evaluator(args.hf_token, use_llm_agent=args.llm_agent)
     result = evaluator.evaluate_prompt(args.prompt, num_episodes=args.episodes)
     print(f"Prompt: {args.prompt[:80]}...")
     print(f"Mean reward: {result['mean_reward']:.1f}")
@@ -118,6 +134,8 @@ def main():
     parser.add_argument("--output-dir", type=str, default="./grpo_output", help="Training output dir")
     parser.add_argument("--hf-token", type=str, default=None, help="HuggingFace API token")
     parser.add_argument("--prompt", type=str, default=None, help="Prompt to evaluate (eval mode)")
+    parser.add_argument("--llm-agent", action="store_true",
+                        help="Use LLM (Llama 3.1) as the agent instead of rule-based")
     args = parser.parse_args()
 
     if args.mode == "train":

layer2/customer_sim.py

@@ -80,7 +80,18 @@ class CustomerSimulator:
     ) -> str:
         """Generate the next customer reply given the conversation so far."""
         if self._client is not None:
-            return self._generate_llm_reply(persona, conversation_history, agent_message)
+            try:
+                return self._generate_llm_reply(persona, conversation_history, agent_message)
+            except Exception as e:
+                if "402" in str(e) or "Payment Required" in str(e):
+                    import logging
+                    logging.getLogger(__name__).warning(
+                        "HF API credits depleted, falling back to rule-based. "
+                        "Get more credits at https://huggingface.co/settings/billing"
+                    )
+                    self._client = None  # disable for remaining calls
+                else:
+                    raise
         return self._generate_rule_reply(persona, conversation_history, agent_message)
 
     def _generate_llm_reply(

layer2/hf_agent.py

@@ -1,8 +1,9 @@
 """
 HF Inference API wrapper for the voice agent (Layer 2).
 
-Uses a small model via HF Inference to act as the customer support agent
-during evaluation. In training (Layer 1), the agent is the model being optimized.
+Uses Llama 3.1 8B Instruct via HF Inference to act as the customer support
+agent during evaluation. In training (Layer 1), the agent is the model being
+optimized — this module provides the inference-time agent for A/B testing.
 """
 
 from __future__ import annotations
@@ -21,11 +22,11 @@ class HFAgent:
     """
     Voice agent powered by HF Inference API.
 
-    This wraps a small model (e.g. Qwen 2.5 3B) with a system prompt
-    from Layer 1, and generates responses in the customer support conversation.
+    Takes a system prompt from Layer 1 and generates responses
+    in the customer support conversation using Llama 3.1 8B.
     """
 
-    DEFAULT_MODEL = "Qwen/Qwen2.5-3B-Instruct"
+    DEFAULT_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
 
     def __init__(self, model_id: str | None = None, hf_token: str | None = None):
         self.model_id = model_id or self.DEFAULT_MODEL
@@ -34,6 +35,10 @@ class HFAgent:
         if self.hf_token and InferenceClient is not None:
             self._client = InferenceClient(token=self.hf_token)
 
+    @property
+    def is_llm_available(self) -> bool:
+        return self._client is not None
+
     def __call__(
         self,
         system_prompt: str,
@@ -46,7 +51,7 @@ class HFAgent:
         Compatible with ConversationEnvironment.run_episode(agent_fn=...).
         """
         if self._client is None:
-            return self._fallback_response(observation)
+            return self._fallback_response(system_prompt, observation)
 
         messages = [{"role": "system", "content": system_prompt}]
 
@@ -61,15 +66,26 @@ class HFAgent:
         if customer_msg:
             messages.append({"role": "user", "content": customer_msg})
 
-        response = self._client.chat_completion(
-            model=self.model_id,
-            messages=messages,
-            max_tokens=300,
-            temperature=0.3,
-        )
-        return response.choices[0].message.content.strip()
-
-    def _fallback_response(self, observation: dict[str, Any]) -> str:
+        try:
+            response = self._client.chat_completion(
+                model=self.model_id,
+                messages=messages,
+                max_tokens=300,
+                temperature=0.3,
+            )
+            return response.choices[0].message.content.strip()
+        except Exception as e:
+            if "402" in str(e) or "Payment Required" in str(e):
+                import logging
+                logging.getLogger(__name__).warning(
+                    "HF API credits depleted, falling back to rule-based. "
+                    "Get more credits at https://huggingface.co/settings/billing"
+                )
+                self._client = None
+                return self._fallback_response(system_prompt, observation)
+            raise
+
+    def _fallback_response(self, system_prompt: str, observation: dict[str, Any]) -> str:
         """Rule-based fallback when no HF token is available."""
         customer_msg = observation.get("customer_message", "").lower()
         intents = observation.get("intents", [])

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "huggingface-hub>=0.20.0",
     "requests>=2.31.0",
     "pydantic>=2.0",
+    "python-dotenv>=1.0.0",
     "gradio>=4.0.0",
 ]
 
@@ -24,6 +25,7 @@ train = [
     "peft>=0.9.0",
     "bitsandbytes>=0.43.0",
     "accelerate>=0.27.0",
+    "datasets>=2.18.0",
 ]
 dev = [
     "pytest>=8.0",

scripts/ab_test.py

@@ -1,8 +1,11 @@
 """
-A/B Test: Compare base prompt vs trained/optimized prompt on 100 simulated customers.
+A/B Test: Compare base prompt vs trained/optimized prompt.
+
+Uses real LLM (Llama 3.1 8B via HF Inference API) for both
+the customer simulator and the voice agent when HF_TOKEN is set.
 
 Usage:
-    python -m scripts.ab_test [--episodes 100] [--hf-token TOKEN]
+    python -m scripts.ab_test [--episodes 10] [--mode llm|rule]
 """
 
 from __future__ import annotations
@@ -12,11 +15,16 @@ import json
 import sys
 import os
 
+# Auto-load .env
+from dotenv import load_dotenv
+load_dotenv(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".env"))
+
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 from layer0.reward import reward_fn, BANKING_INTENTS
 from layer2.customer_sim import CustomerPersona, CustomerSimulator
 from layer2.environment import ConversationEnvironment, EnvConfig
+from layer2.hf_agent import HFAgent
 from personas.generate_personas import generate_personas
 
 
@@ -42,16 +50,34 @@ TRAINED_PROMPT = (
 
 
 def run_ab_test(
-    num_episodes: int = 100,
+    num_episodes: int = 10,
     hf_token: str | None = None,
+    mode: str = "llm",
 ) -> dict:
-    """Run A/B test comparing base vs trained prompt."""
+    """
+    Run A/B test comparing base vs trained prompt.
+
+    Args:
+        num_episodes: Number of episodes per prompt
+        hf_token: HuggingFace API token (auto-loaded from .env if not provided)
+        mode: "llm" for real LLM agent+customer, "rule" for rule-based fallback
+    """
+    token = hf_token or os.environ.get("HF_TOKEN")
+
     # Load personas
     personas_data = generate_personas(num_episodes)
     personas = [CustomerPersona(**p) for p in personas_data]
 
-    # Initialize simulator
-    simulator = CustomerSimulator(hf_token=hf_token)
+    # Initialize simulator (uses LLM if token available)
+    simulator = CustomerSimulator(hf_token=token if mode == "llm" else None)
+
+    # Initialize LLM agent (uses LLM if token available)
+    agent = HFAgent(hf_token=token if mode == "llm" else None)
+
+    using_llm = mode == "llm" and agent.is_llm_available
+    print(f"Mode: {'LLM (Llama 3.1 8B)' if using_llm else 'Rule-based'}")
+    print(f"Customer sim: {'LLM' if simulator._client else 'Rule-based'}")
+    print(f"Agent: {'LLM' if agent.is_llm_available else 'Rule-based'}")
 
     # Create environment
     env = ConversationEnvironment(
@@ -73,9 +99,17 @@ def run_ab_test(
         correct = 0
         injection_resisted = 0
         injection_total = 0
+        sample_conversations = []
 
         for i, persona in enumerate(personas):
-            log = env.run_episode(system_prompt=prompt, persona=persona)
+            # Use LLM agent if available, otherwise default rule-based
+            agent_fn = agent if using_llm else None
+
+            log = env.run_episode(
+                system_prompt=prompt,
+                agent_fn=agent_fn,
+                persona=persona,
+            )
             r = reward_fn(log)
             rewards.append(r)
             turns_list.append(log.turns)
@@ -88,7 +122,20 @@ def run_ab_test(
                 if not log.injection_succeeded:
                     injection_resisted += 1
 
-            if (i + 1) % 25 == 0:
+            # Save first 3 conversations for inspection
+            if len(sample_conversations) < 3:
+                sample_conversations.append({
+                    "persona_id": persona.id,
+                    "true_intent": persona.true_intent,
+                    "social_engineering": persona.social_engineering,
+                    "messages": log.messages if hasattr(log, "messages") else [],
+                    "reward": r,
+                    "intent_correct": log.intent_correct,
+                    "injection_succeeded": log.injection_succeeded,
+                    "turns": log.turns,
+                })
+
+            if (i + 1) % max(1, num_episodes // 4) == 0:
                 print(f"  [{i+1}/{num_episodes}] avg_reward={sum(rewards)/len(rewards):.1f}")
 
         results[label] = {
@@ -101,6 +148,8 @@ def run_ab_test(
             "min_reward": min(rewards),
             "max_reward": max(rewards),
             "total_episodes": num_episodes,
+            "mode": "llm" if using_llm else "rule",
+            "sample_conversations": sample_conversations,
         }
 
     return results
@@ -112,6 +161,10 @@ def print_results(results: dict):
     print("=" * 62)
     print(f"{'A/B TEST RESULTS':^62}")
     print("=" * 62)
+
+    mode = results.get("base", {}).get("mode", "unknown")
+    print(f"{'Mode: ' + mode:^62}")
+    print("-" * 62)
     print(f"{'Metric':<25} {'Base Prompt':>15} {'Trained Prompt':>18}")
     print("-" * 62)
 
@@ -129,27 +182,49 @@ def print_results(results: dict):
         print(f"{name:<25} {b_val:>15} {t_val:>18}")
 
     print("=" * 62)
-    print()
+
+    # Print sample conversations
+    for label in ["base", "trained"]:
+        samples = results[label].get("sample_conversations", [])
+        if samples:
+            print(f"\n--- Sample conversations ({label.upper()}) ---")
+            for conv in samples[:2]:
+                print(f"  Persona {conv['persona_id']} ({conv['true_intent']}, "
+                      f"SE={conv['social_engineering']})")
+                for msg in conv.get("messages", []):
+                    if isinstance(msg, dict):
+                        role = "Customer" if msg.get("role") == "customer" else "Agent"
+                        text = msg.get("content", "")[:120]
+                        print(f"    [{role}] {text}")
+                print(f"    => reward={conv['reward']:.1f} correct={conv['intent_correct']} "
+                      f"injection={conv['injection_succeeded']}")
+                print()
 
 
 def main():
     parser = argparse.ArgumentParser(description="A/B test: base vs trained prompt")
-    parser.add_argument("--episodes", type=int, default=100, help="Number of episodes per prompt")
+    parser.add_argument("--episodes", type=int, default=10, help="Number of episodes per prompt")
     parser.add_argument("--hf-token", type=str, default=None, help="HuggingFace API token")
+    parser.add_argument("--mode", choices=["llm", "rule"], default="llm",
+                        help="llm=real LLM agent+customer, rule=rule-based fallback")
     parser.add_argument("--output", type=str, default=None, help="Save results to JSON file")
     args = parser.parse_args()
 
     results = run_ab_test(
         num_episodes=args.episodes,
         hf_token=args.hf_token,
+        mode=args.mode,
     )
 
     print_results(results)
 
     if args.output:
+        # Remove non-serializable data
+        for label in results:
+            results[label].pop("sample_conversations", None)
         with open(args.output, "w") as f:
             json.dump(results, f, indent=2)
-        print(f"Results saved to {args.output}")
+        print(f"\nResults saved to {args.output}")
 
 
 if __name__ == "__main__":