Upload tests/test_optimizer.py

tests/test_optimizer.py

@@ -0,0 +1,100 @@
+"""Tests for Agent Cost Optimizer."""
+
+import unittest
+from aco import AgentCostOptimizer
+from aco.config import ACOConfig, ModelConfig, ToolConfig, VerifierConfig, RoutingPolicy
+from aco.trace_schema import TaskType, Outcome, FailureTag, ModelCall
+
+
+class TestAgentCostOptimizer(unittest.TestCase):
+    def setUp(self):
+        self.config = ACOConfig(
+            models={
+                "tiny": ModelConfig("tiny", "local", 0.0001, 0.0002, strength_tier=1),
+                "cheap": ModelConfig("cheap", "cloud", 0.0005, 0.001, strength_tier=2),
+                "medium": ModelConfig("medium", "cloud", 0.003, 0.006, strength_tier=3),
+                "frontier": ModelConfig("frontier", "cloud", 0.01, 0.03, strength_tier=4),
+            },
+            tools={
+                "search": ToolConfig("search", 0.002, 500),
+            },
+            verifiers={
+                "v1": VerifierConfig("v1", 0.005, 800, 0.8),
+            },
+            routing_policy=RoutingPolicy("test"),
+        )
+        self.optimizer = AgentCostOptimizer(self.config)
+
+    def test_optimize_quick_answer(self):
+        result = self.optimizer.optimize("What is the capital of France?")
+        self.assertEqual(result.routing_decision.tier, 1)
+        self.assertLess(result.estimated_cost, 0.01)
+
+    def test_optimize_coding_task(self):
+        result = self.optimizer.optimize("Write a Python function to reverse a linked list")
+        self.assertIn(result.routing_decision.tier, [2, 3])
+        self.assertGreater(result.confidence, 0.0)
+
+    def test_optimize_legal_task(self):
+        result = self.optimizer.optimize("Review this contract for liability clauses")
+        self.assertEqual(result.routing_decision.tier, 4)  # frontier for legal
+        self.assertTrue(result.verifier_decision is not None)
+
+    def test_trace_lifecycle(self):
+        from aco.classifier import TaskPrediction
+        
+        trace_id = "test-trace-001"
+        prediction = TaskPrediction(
+            task_type=TaskType.CODING,
+            expected_cost=0.05,
+            expected_model_tier=3,
+            expected_tools_needed=["code_execution"],
+            risk_of_failure=0.5,
+            retrieval_required=False,
+            verifier_required=False,
+            expected_latency_ms=1000,
+            confidence=0.8,
+        )
+        
+        trace = self.optimizer.start_trace(trace_id, "Test request", prediction)
+        self.assertEqual(trace.trace_id, trace_id)
+        
+        model_call = ModelCall(
+            model_id="medium",
+            provider="cloud",
+            input_tokens=1024,
+            output_tokens=512,
+            cost_per_1k_input=0.003,
+            cost_per_1k_output=0.006,
+        )
+        
+        self.optimizer.record_step(
+            trace_id=trace_id,
+            model_call=model_call,
+            context_size_tokens=1024,
+            step_outcome=Outcome.SUCCESS,
+        )
+        
+        finalized = self.optimizer.finalize_trace(
+            trace_id=trace_id,
+            outcome=Outcome.SUCCESS,
+            user_satisfaction=0.95,
+        )
+        
+        self.assertIsNotNone(finalized)
+        self.assertEqual(finalized.final_outcome, Outcome.SUCCESS)
+        self.assertGreater(finalized.total_cost_computed, 0)
+
+    def test_cost_adjusted_score(self):
+        from aco.datasets.synthetic_traces import SyntheticTraceGenerator
+        
+        gen = SyntheticTraceGenerator(seed=42)
+        traces = gen.generate(10)
+        
+        for trace in traces:
+            score = self.optimizer.compute_cost_adjusted_score(trace)
+            self.assertIsInstance(score, float)
+
+
+if __name__ == "__main__":
+    unittest.main()