ruvector/crates/ruvector-postgres/tests/routing_tests.rs

// Integration tests for Tiny Dancer Routing module
//
// These tests validate the complete routing functionality including
// agent registration, FastGRNN neural network, and routing decisions.

#[cfg(test)]
mod routing_tests {
    use ruvector_postgres::routing::{
        agents::{Agent, AgentRegistry, AgentType},
        fastgrnn::FastGRNN,
        router::{OptimizationTarget, Router, RoutingConstraints},
    };

    #[test]
    fn test_complete_routing_workflow() {
        // Create registry and router
        let registry = AgentRegistry::new();
        let router = Router::with_registry(std::sync::Arc::new(registry));

        // Register diverse agents
        let agents = vec![
            create_agent("gpt-4", 0.03, 500.0, 0.95, vec!["coding", "reasoning"]),
            create_agent("claude-3", 0.025, 400.0, 0.93, vec!["coding", "writing"]),
            create_agent("gpt-3.5", 0.002, 200.0, 0.75, vec!["general", "fast"]),
            create_agent("llama-2", 0.0, 800.0, 0.70, vec!["local", "private"]),
        ];

        for agent in agents {
            router.registry().register(agent).unwrap();
        }

        // Test cost-optimized routing
        let request_emb = vec![0.1; 384];
        let decision = router
            .route(
                &request_emb,
                &RoutingConstraints::new(),
                OptimizationTarget::Cost,
            )
            .unwrap();

        assert_eq!(decision.agent_name, "llama-2"); // Free option
        assert!(decision.confidence > 0.0);

        // Test quality-optimized routing
        let decision = router
            .route(
                &request_emb,
                &RoutingConstraints::new(),
                OptimizationTarget::Quality,
            )
            .unwrap();

        assert_eq!(decision.agent_name, "gpt-4"); // Highest quality

        // Test latency-optimized routing
        let decision = router
            .route(
                &request_emb,
                &RoutingConstraints::new(),
                OptimizationTarget::Latency,
            )
            .unwrap();

        assert_eq!(decision.agent_name, "gpt-3.5"); // Fastest
    }

    #[test]
    fn test_routing_with_constraints() {
        let registry = AgentRegistry::new();
        let router = Router::with_registry(std::sync::Arc::new(registry));

        router
            .registry()
            .register(create_agent(
                "expensive-high-quality",
                1.0,
                200.0,
                0.99,
                vec!["coding"],
            ))
            .unwrap();

        router
            .registry()
            .register(create_agent(
                "cheap-medium-quality",
                0.01,
                200.0,
                0.75,
                vec!["coding"],
            ))
            .unwrap();

        let request_emb = vec![0.1; 384];

        // Constrain by max cost
        let constraints = RoutingConstraints::new()
            .with_max_cost(0.5)
            .with_min_quality(0.7);

        let decision = router
            .route(&request_emb, &constraints, OptimizationTarget::Quality)
            .unwrap();

        // Should pick cheap option due to cost constraint
        assert_eq!(decision.agent_name, "cheap-medium-quality");
    }

    #[test]
    fn test_fastgrnn_routing() {
        let mut router = Router::new();
        router.init_grnn(64);

        router
            .registry()
            .register(create_agent("agent1", 0.05, 200.0, 0.85, vec!["coding"]))
            .unwrap();

        let request_emb = vec![0.1; 384];

        let decision = router
            .route(
                &request_emb,
                &RoutingConstraints::new(),
                OptimizationTarget::Balanced,
            )
            .unwrap();

        // Verify neural network enhanced confidence
        assert!(decision.confidence > 0.0);
        assert!(decision.confidence <= 1.0);
    }

    #[test]
    fn test_capability_based_routing() {
        let registry = AgentRegistry::new();
        let router = Router::with_registry(std::sync::Arc::new(registry));

        router
            .registry()
            .register(create_agent(
                "coder",
                0.05,
                200.0,
                0.90,
                vec!["coding", "debugging"],
            ))
            .unwrap();

        router
            .registry()
            .register(create_agent(
                "writer",
                0.03,
                150.0,
                0.85,
                vec!["writing", "translation"],
            ))
            .unwrap();

        router
            .registry()
            .register(create_agent(
                "generalist",
                0.02,
                300.0,
                0.70,
                vec!["coding", "writing", "general"],
            ))
            .unwrap();

        let request_emb = vec![0.1; 384];

        // Require coding capability
        let constraints = RoutingConstraints::new().with_capability("coding".to_string());

        let decision = router
            .route(&request_emb, &constraints, OptimizationTarget::Quality)
            .unwrap();

        // Should pick specialized coder (highest quality with coding)
        assert!(decision.agent_name == "coder" || decision.agent_name == "generalist");

        // Verify writer was not selected
        assert_ne!(decision.agent_name, "writer");
    }

    #[test]
    fn test_agent_metrics_update() {
        let registry = AgentRegistry::new();
        let mut agent = create_agent("test-agent", 0.05, 200.0, 0.80, vec!["test"]);

        // Initial state
        assert_eq!(agent.performance.total_requests, 0);
        assert_eq!(agent.performance.avg_latency_ms, 200.0);

        // Update with better latency
        agent.update_metrics(150.0, true, Some(0.85));
        assert_eq!(agent.performance.total_requests, 1);
        assert_eq!(agent.performance.avg_latency_ms, 150.0);
        assert_eq!(agent.performance.success_rate, 1.0);

        // Update with worse latency
        agent.update_metrics(250.0, true, Some(0.75));
        assert_eq!(agent.performance.total_requests, 2);
        assert_eq!(agent.performance.avg_latency_ms, 200.0); // Average of 150 and 250
        assert_eq!(agent.performance.success_rate, 1.0);

        // Failed request
        agent.update_metrics(300.0, false, None);
        assert_eq!(agent.performance.total_requests, 3);
        assert!(agent.performance.success_rate < 1.0);
    }

    #[test]
    fn test_fastgrnn_sequence_processing() {
        let grnn = FastGRNN::new(10, 5);

        let sequence = vec![
            vec![1.0, 0.0, 0.0, 0.5, -0.5, 0.2, -0.2, 0.8, -0.8, 0.0],
            vec![0.0, 1.0, 0.0, -0.5, 0.5, -0.2, 0.2, -0.8, 0.8, 0.0],
            vec![0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
        ];

        let outputs = grnn.forward_sequence(&sequence);

        assert_eq!(outputs.len(), 3);
        assert_eq!(outputs[0].len(), 5);

        // Verify state evolution (later states should be different from first)
        let first_state = &outputs[0];
        let last_state = &outputs[2];

        let diff: f32 = first_state
            .iter()
            .zip(last_state.iter())
            .map(|(a, b)| (a - b).abs())
            .sum();

        assert!(diff > 0.0, "Hidden state should evolve across sequence");
    }

    #[test]
    fn test_routing_alternatives() {
        let registry = AgentRegistry::new();
        let router = Router::with_registry(std::sync::Arc::new(registry));

        // Register multiple similar agents
        for i in 0..5 {
            let quality = 0.7 + (i as f32 * 0.05);
            let cost = 0.01 + (i as f32 * 0.01);
            router
                .registry()
                .register(create_agent(
                    &format!("agent-{}", i),
                    cost,
                    200.0,
                    quality,
                    vec!["test"],
                ))
                .unwrap();
        }

        let request_emb = vec![0.1; 384];

        let decision = router
            .route(
                &request_emb,
                &RoutingConstraints::new(),
                OptimizationTarget::Quality,
            )
            .unwrap();

        // Should have alternatives listed
        assert!(!decision.alternatives.is_empty());
        assert!(decision.alternatives.len() <= 3); // Max 3 alternatives

        // Alternatives should have lower scores
        for alt in &decision.alternatives {
            assert!(alt.score < 1.0);
            assert!(!alt.reason.is_empty());
        }
    }

    #[test]
    fn test_excluded_agents() {
        let registry = AgentRegistry::new();
        let router = Router::with_registry(std::sync::Arc::new(registry));

        router
            .registry()
            .register(create_agent("agent-a", 0.05, 200.0, 0.90, vec!["test"]))
            .unwrap();

        router
            .registry()
            .register(create_agent("agent-b", 0.05, 200.0, 0.85, vec!["test"]))
            .unwrap();

        let request_emb = vec![0.1; 384];

        // Exclude the best agent
        let constraints = RoutingConstraints::new().with_excluded_agent("agent-a".to_string());

        let decision = router
            .route(&request_emb, &constraints, OptimizationTarget::Quality)
            .unwrap();

        assert_eq!(decision.agent_name, "agent-b");
    }

    // Helper function to create test agents
    fn create_agent(
        name: &str,
        cost: f32,
        latency: f32,
        quality: f32,
        capabilities: Vec<&str>,
    ) -> Agent {
        let mut agent = Agent::new(
            name.to_string(),
            AgentType::LLM,
            capabilities.iter().map(|s| s.to_string()).collect(),
        );
        agent.cost_model.per_request = cost;
        agent.performance.avg_latency_ms = latency;
        agent.performance.quality_score = quality;
        agent.embedding = Some(vec![0.1; 384]); // Default embedding
        agent
    }
}