"""Router model module.

Provides two routing architectures:

1. RouterMLP: Simple feedforward network for fast inference
2. CompatibilityScorer: Dual-encoder model for query-LLM compatibility scoring

Both models learn to predict which LLM will perform best for a given query
based on learned representations.
"""

from partner_inference_harness.model.hetero_graph import (
    CompatibilityConfig,
    CompatibilityScorer,
    HeteroGraphConfig,
    HeteroGraphRouter,
    load_model as load_graph_model,
    save_model as save_graph_model,
)
from partner_inference_harness.model.router_mlp import (
    RouterConfig,
    RouterMLP,
    load_model,
    save_model,
)

__all__ = [
    # MLP Router
    "RouterConfig",
    "RouterMLP",
    "load_model",
    "save_model",
    # Compatibility Scorer (primary)
    "CompatibilityConfig",
    "CompatibilityScorer",
    "load_graph_model",
    "save_graph_model",
    # Backward compatibility aliases
    "HeteroGraphConfig",
    "HeteroGraphRouter",
]
