af_agent/

model.rs

//! The chat-model abstraction the agent loop runs against.
//!
//! Decoupling the loop from a concrete client means: (a) the loop is unit-
//! testable with a scripted mock — no network; (b) any backend (the real
//! [`LlmClient`], a local model, a replay harness) plugs in by implementing one
//! method.

use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;

use af_llm::{CompletionRequest, CompletionResponse, LlmClient, LlmError};
use async_trait::async_trait;
use tokio::sync::mpsc::UnboundedSender;

/// Registered model and the request options it supports.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ModelDescriptor {
    /// Stable identifier of this record.
    pub id: String,
    /// Name of the provider that produced or serves this record.
    pub provider: String,
    /// Reasoning efforts the provider accepts (`low`, `medium`, `high`).
    #[serde(default)]
    pub reasoning_efforts: BTreeSet<String>,
    /// Input modalities the model accepts (for example `text`, `image`).
    #[serde(default)]
    pub input_modalities: BTreeSet<String>,
}

/// Model registration or resolution failure.
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum ModelRegistryError {
    /// Model descriptor requires id and provider.
    #[error("model descriptor requires id and provider")]
    InvalidDescriptor,
    /// Model is already registered.
    #[error("model is already registered: {0}")]
    Duplicate(String),
    /// Model is not registered.
    #[error("model is not registered: {0}")]
    Unknown(String),
    /// Preferred model is not allowed.
    #[error("preferred model is not allowed: {0}")]
    NotAllowed(String),
}

/// Exact-id registry of models a Profile may select.
#[derive(Default)]
pub struct ModelRegistry {
    models: BTreeMap<String, (ModelDescriptor, Arc<dyn ChatModel>)>,
}

impl ModelRegistry {
    /// Register a model; duplicate ids are rejected.
    pub fn register(
        &mut self,
        descriptor: ModelDescriptor,
        model: Arc<dyn ChatModel>,
    ) -> Result<(), ModelRegistryError> {
        if descriptor.id.trim().is_empty() || descriptor.provider.trim().is_empty() {
            return Err(ModelRegistryError::InvalidDescriptor);
        }
        if self.models.contains_key(&descriptor.id) {
            return Err(ModelRegistryError::Duplicate(descriptor.id));
        }
        self.models
            .insert(descriptor.id.clone(), (descriptor, model));
        Ok(())
    }

    /// Registered descriptors, sorted by id.
    pub fn descriptors(&self) -> impl Iterator<Item = &ModelDescriptor> {
        self.models.values().map(|(descriptor, _)| descriptor)
    }

    /// Resolve an exact model id within `allowed`, validating the requested reasoning effort.
    pub fn resolve(
        &self,
        preferred: &str,
        allowed: &BTreeSet<String>,
    ) -> Result<Arc<dyn ChatModel>, ModelRegistryError> {
        if !allowed.contains(preferred) {
            return Err(ModelRegistryError::NotAllowed(preferred.into()));
        }
        self.models
            .get(preferred)
            .map(|(descriptor, model)| {
                Arc::new(RegisteredModel {
                    model: Arc::clone(model),
                    images: descriptor.input_modalities.contains("image"),
                    provider: descriptor.provider.clone(),
                }) as Arc<dyn ChatModel>
            })
            .ok_or_else(|| ModelRegistryError::Unknown(preferred.into()))
    }
}

struct RegisteredModel {
    model: Arc<dyn ChatModel>,
    images: bool,
    provider: String,
}

#[async_trait]
impl ChatModel for RegisteredModel {
    fn provider_name(&self) -> Option<&str> {
        Some(&self.provider)
    }
    async fn complete_streaming(
        &self,
        request: &CompletionRequest,
        delta_tx: UnboundedSender<(String, bool)>,
    ) -> Result<CompletionResponse, LlmError> {
        if !self.images
            && request
                .messages
                .iter()
                .any(|message| !message.images.is_empty())
        {
            return Err(LlmError::InvalidInput(
                "selected model does not support images".into(),
            ));
        }
        self.model.complete_streaming(request, delta_tx).await
    }
}

/// Anything that can turn a chat-completion request into a response.
#[async_trait]
pub trait ChatModel: Send + Sync {
    /// Stable provider attribution; custom adapters may leave this unknown.
    fn provider_name(&self) -> Option<&str> {
        None
    }
    /// Stream cumulative `(content, has_tool_calls)` updates and return the
    /// canonical terminal response. This is the only model invocation path.
    async fn complete_streaming(
        &self,
        request: &CompletionRequest,
        delta_tx: UnboundedSender<(String, bool)>,
    ) -> Result<CompletionResponse, LlmError>;
}

/// The production model is the real LLM client.
#[async_trait]
impl ChatModel for LlmClient {
    async fn complete_streaming(
        &self,
        request: &CompletionRequest,
        delta_tx: UnboundedSender<(String, bool)>,
    ) -> Result<CompletionResponse, LlmError> {
        self.complete_stream_single_attempt(request, |content, has_tools| {
            let _ = delta_tx.send((content.to_string(), has_tools));
        })
        .await
    }
}

/// A type-erased model, so callers (e.g. an HTTP service holding one `Agent` in
/// shared state) can pick the backend at runtime — real client vs stub vs
/// replay — without the loop being generic over every concrete type.
#[async_trait]
impl ChatModel for Arc<dyn ChatModel> {
    fn provider_name(&self) -> Option<&str> {
        (**self).provider_name()
    }
    async fn complete_streaming(
        &self,
        request: &CompletionRequest,
        delta_tx: UnboundedSender<(String, bool)>,
    ) -> Result<CompletionResponse, LlmError> {
        (**self).complete_streaming(request, delta_tx).await
    }
}

#[cfg(test)]
mod registry_tests {
    use super::*;
    use af_llm::{ChatMessage, Choice, CompletionResponse, LlmError};

    struct Fixed;

    #[async_trait]
    impl ChatModel for Fixed {
        async fn complete_streaming(
            &self,
            _: &CompletionRequest,
            _: UnboundedSender<(String, bool)>,
        ) -> Result<CompletionResponse, LlmError> {
            Ok(CompletionResponse {
                id: "fixed".into(),
                choices: vec![Choice {
                    index: 0,
                    message: ChatMessage::assistant("ok"),
                    output_blocks: Vec::new(),
                    finish_reason: None,
                }],
                usage: None,
            })
        }
    }

    fn descriptor(id: &str) -> ModelDescriptor {
        ModelDescriptor {
            id: id.into(),
            provider: "test".into(),
            reasoning_efforts: BTreeSet::from(["low".into()]),
            input_modalities: BTreeSet::from(["text".into()]),
        }
    }

    #[tokio::test]
    async fn registered_modality_is_enforced_before_invocation() {
        let mut registry = ModelRegistry::default();
        registry
            .register(descriptor("text"), Arc::new(Fixed))
            .unwrap();
        let mut vision = descriptor("vision");
        vision.input_modalities.insert("image".into());
        registry.register(vision, Arc::new(Fixed)).unwrap();
        let mut message = ChatMessage::user("describe");
        message.images.push(af_llm::InputImage {
            asset_id: "asset".parse().unwrap(),
            media_type: "image/png".into(),
        });
        for (id, success) in [("text", false), ("vision", true)] {
            let model = registry.resolve(id, &BTreeSet::from([id.into()])).unwrap();
            let (tx, _) = tokio::sync::mpsc::unbounded_channel();
            let result = model
                .complete_streaming(&CompletionRequest::new(id, vec![message.clone()]), tx)
                .await;
            assert_eq!(result.is_ok(), success);
        }
    }

    #[test]
    fn registry_is_exact_sorted_and_rejects_duplicates() {
        let mut registry = ModelRegistry::default();
        registry.register(descriptor("z"), Arc::new(Fixed)).unwrap();
        registry.register(descriptor("a"), Arc::new(Fixed)).unwrap();
        assert_eq!(
            registry
                .descriptors()
                .map(|value| value.id.as_str())
                .collect::<Vec<_>>(),
            vec!["a", "z"]
        );
        assert_eq!(
            registry
                .register(descriptor("a"), Arc::new(Fixed))
                .unwrap_err(),
            ModelRegistryError::Duplicate("a".into())
        );
        assert!(registry.resolve("a", &BTreeSet::from(["a".into()])).is_ok());
        assert!(matches!(
            registry.resolve("z", &BTreeSet::from(["a".into()])),
            Err(ModelRegistryError::NotAllowed(value)) if value == "z"
        ));
    }
}