add 'llm' crate as provider

[Tails]

There is this great LLM provider crate: https://github.com/graniet/llm

It would be great to add it as provider so support LLM's that do not neccessarily have an OAI-compatible endpoint.

[kyegomez]

@Tails great idea, does it support llms like anthropic? and deepseek? and others?

we use litellm in our python version is it similiar?

and you should also join our discord community for real-time discussions, since i have so many githubs i don't see notifications easily

https://discord.gg/jM3Z6M9uMq

[Tails]

Ja, I've tried quite a few and 'llm' stuck with me: https://crates.io/crates/llm

I had taken a stab at a draft integration. It has a lot of unimplemented!() parts, but the basic completion and struct definition conversion works. Did not clean it up from Gemini yet.

use core::fmt;
use std::collections::HashMap;
use std::fmt::Formatter;
use std::sync::Arc;

use ::llm::FunctionCall;
use ::llm::LLMProvider; /* LLMProvider seems to be at the root based on one of the compiler messages */
use ::llm::builder::{LLMBackend, LLMBuilder};
use ::llm::chat::ParameterProperty;
use ::llm::chat::ParameterSchemaOrProperty;
use ::llm::chat::ParameterStringProperty;
// Use fully qualified paths for llm crate types as suggested by compiler
use ::llm::chat::{ChatMessage, ChatResponse, Tool as LlmSdkChatTool}; /* Renamed Tool to LlmSdkChatTool to avoid conflict */
use ::llm::chat::{ChatRole, MessageType};
use ::llm::chat::{FunctionTool, ParametersSchema};
use ::llm::error::LLMError;
use anyhow::anyhow;
// Comment out async_trait for now
// use async_trait::async_trait;
use futures::future::BoxFuture;
// use ::llm::tool::{Function as LlmSdkFunction, ToolCall as LlmSdkToolCall, ToolType as LlmSdkToolType};
use serde_json;
use serde_json::{Value, json};
use tracing::debug;

use crate::llm::completion::{
    AssistantContent, Message, Text, ToolCall, ToolFunction, ToolResult, ToolResultContent,
    UserContent,
};
use crate::llm::request::ToolDefinition;
// For logit_bias in Step
use crate::{
    agent::SwarmsAgentBuilder,
    llm::{
        self, CompletionError, Model,
        request::{CompletionRequest, CompletionResponse},
    },
};

#[derive(Clone)]
pub struct GranietLlm {
    backend: LLMBackend,
    api_key: String,
    model_name: String,
    system_prompt: Option<String>,
    base_url: Option<String>, // Keep for potential future use by LLMBuilder
}

impl GranietLlm {
    pub fn try_new<S: Into<String>>(
        backend: LLMBackend,
        api_key: String,
        model: S,
        base_url: Option<String>,
    ) -> Result<Self, LLMError> {
        let model_name_str: String = model.into();

        Ok(Self {
            backend,
            api_key,
            model_name: model_name_str,
            system_prompt: None,
            base_url,
        })
    }

    pub fn from_env_with_backend_and_model<S: Into<String>>(
        backend: LLMBackend,
        model: S,
    ) -> anyhow::Result<Self> {
        let api_key_var = match backend {
            LLMBackend::OpenAI => "OPENAI_API_KEY",
            LLMBackend::Anthropic => "ANTHROPIC_API_KEY",
            LLMBackend::Ollama => {
                unimplemented!()
            },
            LLMBackend::DeepSeek => {
                unimplemented!()
            },
            LLMBackend::XAI => {
                unimplemented!()
            },
            LLMBackend::Phind => {
                unimplemented!()
            },
            LLMBackend::Google => "GEMINI_API_KEY",
            LLMBackend::Groq => {
                unimplemented!()
            },
            LLMBackend::AzureOpenAI => {
                unimplemented!()
            },
            LLMBackend::ElevenLabs => {
                unimplemented!()
            },
        };

        let api_key = std::env::var(api_key_var)?;

        Ok(Self::try_new(backend, api_key, model, None)?)
    }

    pub fn set_model<S: Into<String>>(mut self, model: S) -> Self {
        self.model_name = model.into();
        self
    }

    pub fn set_system_prompt<S: Into<String>>(&mut self, prompt: S) {
        self.system_prompt = Some(prompt.into());
    }

    pub fn agent_builder(&self) -> SwarmsAgentBuilder<Self> {
        SwarmsAgentBuilder::new_with_model(self.clone())
    }
}

impl Into<::llm::chat::Tool> for ToolDefinition {
    fn into(self) -> ::llm::chat::Tool {
        let args: HashMap<String, serde_json::Value> =
            serde_json::from_value(self.parameters.clone()).unwrap();

        // todo: this is not passed yet from swarm-rs Tool!
        let mut required = args["properties"]
            .as_object()
            .unwrap()
            .clone()
            .keys()
            .cloned()
            .collect::<Vec<_>>();

        let mut properties: HashMap<String, ParameterSchemaOrProperty> = HashMap::new();

        fn into_llm_tool_property(
            propname: Option<&str>,
            value: &Value,
        ) -> Result<Option<ParameterSchemaOrProperty>, anyhow::Error> {
            (|| {
                if propname == Some("$defs") {
                    return Ok(None);
                }

                match value {
                    Value::Null => Ok(None),
                    Value::Bool(_) => Ok(Some(
                        ParameterProperty {
                            property_type: "boolean".to_string().into(),
                            description: None,
                            items: None,
                            enum_list: None,
                            one_of: None,
                        }
                        .into(),
                    )),
                    Value::Number(_) => Ok(Some(
                        ParameterProperty {
                            property_type: "number".to_string().into(),
                            description: None,
                            items: None,
                            enum_list: None,
                            one_of: None,
                        }
                        .into(),
                    )),
                    Value::String(_) => Ok(Some(
                        ParameterProperty {
                            property_type: Some("string".to_string()),
                            description: None,
                            items: None,
                            enum_list: None,
                            one_of: None,
                        }
                        .into(),
                    )),
                    Value::Array(v) => {
                        let item_schema = if v.is_empty() {
                            // Handle empty array case: maybe return a generic array or error
                            // For now, let's assume items can be of any type (less strict)
                            // Or, if schema requires a specific type, this might be an error
                            // eprintln!("Warning: Array items schema is undefined for empty array.");
                            None // Or some default schema, or an error
                        } else {
                            into_llm_tool_property(None, &v[0])?.map(Box::new)
                        };
                        Ok(Some(
                            ParameterProperty {
                                property_type: "array".to_string().into(),
                                description: None,
                                items: item_schema,
                                enum_list: None,
                                one_of: None,
                            }
                            .into(),
                        ))
                    },
                    Value::Object(map) => {
                        let object_type = map
                            .get("type")
                            .and_then(Value::as_str)
                            .or_else(|| map.get("oneOf").is_some().then_some("oneOf"))
                            .ok_or_else(|| anyhow!("Object missing 'type' field: {:?}", map))?;

                        match object_type {
                            /// example:
                            ///      {
                            //     "description": String("Tracks task status and provides feedback for incomplete tasks.\n When `Incomplete`, the `context` becomes the system's next prompt."),
                            //     "oneOf": Array [
                            //         Object {
                            //             "const": String("Complete"),
                            //             "description": String("Indicates that the task is complete."),
                            //             "type": String("string"),
                            //         },
                            //         Object {
                            //             "additionalProperties": Bool(false),
                            //             "description": String("Indicates that the task is incomplete."),
                            //             "properties": Object {
                            //                 "Incomplete": Object {
                            //                     "properties": Object {
                            //                         "context": Object {
                            //                             "description": String("Required guidance when task is incomplete:\n - Clear description of missing elements\n - Specific questions needing clarification  \n - Remaining steps to completion"),
                            //                             "type": String("string"),
                            //                         },
                            //                     },
                            //                     "required": Array [
                            //                         String("context"),
                            //                     ],
                            //                     "type": String("object"),
                            //                 },
                            //             },
                            //             "required": Array [
                            //                 String("Incomplete"),
                            //             ],
                            //             "type": String("object"),
                            //         },
                            //     ],
                            // }.
                            "oneOf" => {
                                let one_of_array = map.get("oneOf")
                                    .and_then(Value::as_array)
                                    .ok_or_else(|| anyhow!("'oneOf' key (used as type for propname: {:?}) does not contain an array", propname))?;

                                let mut collected_one_of_schemas = Vec::new();
                                for sub_schema_value in one_of_array {
                                    // Recursively call into_llm_tool_property for each schema in the oneOf array.
                                    // propname is None because these are anonymous schemas within the oneOf list.
                                    match into_llm_tool_property(None, sub_schema_value)? {
                                        Some(schema) => collected_one_of_schemas.push(schema),
                                        None => { /* A schema in oneOf resolved to None (e.g. internal $defs), skip it */ }
                                    }
                                }

                                Ok(Some(
                                    ParametersSchema {
                                        schema_type: "oneOf".to_string(), // No explicit 'type' in the original JSON, 'oneOf' was inferred.
                                        description: map.get("description").and_then(Value::as_str).map(String::from),
                                        one_of: Some(collected_one_of_schemas),
                                        properties: None, // Typically not applicable when schema is primarily 'oneOf'
                                        required: None,   // Typically not applicable
                                        items: None,      // Not an array type itself
                                    }
                                    .into(), // Converts to ParameterSchemaOrProperty::Schema
                                ))
                            }
                            "object" => {
                                let properties_map: Result<HashMap<String, ParameterSchemaOrProperty>, anyhow::Error> =
                                    map.get("properties")
                                        .and_then(Value::as_object)
                                        .ok_or_else(|| anyhow!("Object schema missing 'properties' field"))?
                                        .into_iter()
                                        .filter_map(|(key, val)| {
                                            match into_llm_tool_property(Some(key), val) {
                                                Ok(Some(prop_schema)) => Some(Ok((key.clone(), prop_schema))),
                                                Ok(None) => None, // Skip $defs or other None cases
                                                Err(e) => Some(Err(e)),
                                            }
                                        })
                                        .collect();

                                Ok(Some(
                                    ParametersSchema {
                                        properties: properties_map?.into(),
                                        schema_type: "object".to_string().into(),
                                        one_of: None, // simplified, was complex before
                                        required: None, // todo: populate required from schema if present
                                        description: map.get("description").and_then(Value::as_str).map(String::from),
                                        items: None,
                                    }
                                    .into(),
                                ))
                            }
                            "array" => {
                                let items_schema = map.get("items").ok_or_else(|| anyhow!("Array schema missing 'items' field"))?;
                                Ok(Some(
                                    ParameterProperty {
                                        property_type: "array".to_string().into(),
                                        description: map.get("description").and_then(Value::as_str).map(String::from),
                                        items: into_llm_tool_property(None, items_schema)?.map(Box::new),
                                        enum_list: None,
                                        one_of: None,
                                    }
                                    .into(),
                                ))
                            }
                            other_type => Ok(Some(
                                ParameterProperty {
                                    property_type: other_type.to_string().into(),
                                    description: map.get("description").and_then(Value::as_str).map(String::from),
                                    items: None,
                                    enum_list: None,
                                    one_of: None,
                                }
                                .into(),
                            )),
                        }
                    }
                }
            })()
            .map_err(move |e: anyhow::Error| {
                eprintln!(
                    "Error in into_llm_tool_property for propname: {:?}, value: {:#?}. Error: {}",
                    propname,
                    value,
                    e
                );
                e
            })
        }

        properties = args["properties"]
            .as_object()
            .unwrap()
            .iter()
            .filter_map(|(prop, val)| {
                match into_llm_tool_property(Some(prop), val) {
                    Ok(Some(p)) => Some(Ok((prop.to_string(), p))),
                    Ok(None) => None, // Property was intentionally skipped (e.g. $defs)
                    Err(e) => {
                        eprintln!("Error processing property '{}': {:?}", prop, e);
                        Some(Err(e))
                    },
                }
            })
            .collect::<Result<HashMap<String, ParameterSchemaOrProperty>, anyhow::Error>>()
            .unwrap_or_else(|e| {
                eprintln!("Failed to collect properties, using default: {:?}", e);
                Default::default()
            });

        ::llm::chat::Tool {
            tool_type: "function".to_string(),
            function: FunctionTool {
                name: self.name.clone(),
                description: self.description.clone(),
                parameters: ParametersSchema {
                    schema_type: "object".to_string(),
                    required: Some(required),
                    items: None,
                    properties: properties.into(),
                    description: None,
                    one_of: None,
                },
            },
        }
    }
}

impl Into<::llm::chat::ChatMessage> for Message {
    fn into(self) -> ChatMessage {
        match self {
            Message::User { content } => ChatMessage {
                role: ChatRole::User,
                message_type: content
                    .first()
                    .map(|f| match f {
                        UserContent::Text(_) => MessageType::Text,
                        UserContent::ToolResult(_) => MessageType::ToolResult(todo!()),
                        UserContent::Image(_) => MessageType::Image(todo!()),
                        UserContent::Audio(_) => {
                            unimplemented!()
                        },
                        UserContent::Document(_) => {
                            unimplemented!()
                        },
                    })
                    .unwrap(),
                content: content
                    .into_iter()
                    .map(|c| match c {
                        UserContent::Text(Text { text }) => text,
                        UserContent::ToolResult(ToolResult {
                            id,
                            content: content2,
                        }) => content2
                            .into_iter()
                            .map(|tc| match tc {
                                ToolResultContent::Text(Text { text }) => text,
                                ToolResultContent::Image(_) => {
                                    unimplemented!()
                                },
                            })
                            .collect::<Vec<_>>()
                            .join("\n\n"),
                        UserContent::Image(_) => {
                            unimplemented!()
                        },
                        UserContent::Audio(_) => {
                            unimplemented!()
                        },
                        UserContent::Document(_) => {
                            unimplemented!()
                        },
                    })
                    .collect::<Vec<_>>()
                    .join("\n\n"),
            },
            Message::Assistant { content } => {
                ChatMessage {
                    role: ChatRole::Assistant,
                    message_type: {
                        if matches!(&content[0], AssistantContent::Text(_)) {
                            MessageType::Text
                        } else {
                            MessageType::ToolUse(
                                content
                                    .iter()
                                    .map(|c| match c {
                                        AssistantContent::ToolCall(ToolCall {
                                            id,
                                            function: ToolFunction { name, arguments },
                                        }) => ::llm::ToolCall {
                                            id: id.clone(),
                                            call_type: "function".to_string(),
                                            function: FunctionCall {
                                                name: name.clone(),
                                                arguments: serde_json::to_string(&arguments)
                                                    .unwrap(),
                                            },
                                        },
                                        _ => unimplemented!(),
                                    })
                                    .collect(),
                            )
                        }
                    },
                    content: content
                        .iter()
                        .map(|c| match c {
                            AssistantContent::Text(Text { text }) => text.clone(),
                            // todo: shoujld this be filtered out?
                            AssistantContent::ToolCall(ToolCall { id, function }) => {
                                serde_json::to_string(&function).unwrap()
                            },
                        })
                        .collect::<Vec<_>>()
                        .join("\n\n"),
                }
            },
        }
    }
}

#[derive(Debug, Clone)]
pub struct SimpleChatResponse {
    pub text: Option<String>,
    pub tool_calls: Option<Vec<::llm::ToolCall>>,
    pub thinking: Option<String>,
}

impl ChatResponse for SimpleChatResponse {
    fn text(&self) -> Option<String> {
        self.text.clone()
    }

    fn tool_calls(&self) -> Option<Vec<::llm::ToolCall>> {
        self.tool_calls.clone()
    }

    fn thinking(&self) -> Option<String> {
        self.thinking.clone()
    }
}

impl fmt::Display for SimpleChatResponse {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        todo!()
    }
}

// todo: make better, since not every type maps well
impl From<::llm::error::LLMError> for CompletionError {
    fn from(err: ::llm::error::LLMError) -> Self {
        match err {
            LLMError::HttpError(http) => CompletionError::Other(http),
            LLMError::AuthError(_) => {
                todo!()
            },
            LLMError::InvalidRequest(_) => {
                todo!()
            },
            LLMError::ProviderError(_) => {
                todo!()
            },
            LLMError::ResponseFormatError {
                message,
                raw_response,
            } => CompletionError::Request(
                anyhow!("error in request format: {message}.\n\nResponse: {raw_response}").into(),
            ),
            LLMError::JsonError(err) => CompletionError::Other(err),
            LLMError::ToolConfigError(err) => CompletionError::Other(err),
        }
    }
}

impl Model for GranietLlm {
    type RawCompletionResponse = SimpleChatResponse;

    fn completion(
        &self,
        request: CompletionRequest,
    ) -> BoxFuture<Result<CompletionResponse<Self::RawCompletionResponse>, CompletionError>> {
        Box::pin(async move {
            let input_messages = request
                .chat_history
                .into_iter()
                .map(|m| m.into())
                .collect::<Vec<::llm::chat::ChatMessage>>();

            let input_tools = request
                .tools
                .into_iter()
                .map(|t| t.into())
                .collect::<Vec<_>>();

            // dbg!(&input_tools);

            // Build the provider on demand
            let mut builder = LLMBuilder::new()
                .backend(self.backend.clone()) // LLMBackend should be Clone
                .model(&self.model_name)
                .api_key(self.api_key.clone());
            // base_url would be applied here if LLMBuilder supports a generic way

            if let Some(ref sp) = request.system_prompt {
                builder = builder.system(sp.clone());
            }

            let provider = builder.build()?;

            let raw_response = provider
                .chat_with_tools(&input_messages, Some(&input_tools))
                .await?;

            let simple_response = SimpleChatResponse {
                text: raw_response.text(),
                tool_calls: raw_response.tool_calls(),
                thinking: raw_response.thinking(),
            };

            let mut output_messages = vec![];

            if let Some(txt) = &simple_response.text {
                output_messages.push(AssistantContent::Text(txt.clone().into()));
            }

            if let Some(tools) = simple_response.tool_calls.clone() {
                for ::llm::ToolCall {
                    id,
                    function: ::llm::FunctionCall { name, arguments },
                    ..
                } in tools
                {
                    output_messages.push(AssistantContent::ToolCall(ToolCall {
                        id,
                        function: ToolFunction {
                            name,
                            arguments: serde_json::from_str(&arguments).unwrap(),
                        },
                    }));
                }
            }

            Ok(CompletionResponse {
                // list of well-typed responses
                choice: output_messages,

                // the raw thing in the original type as received form the provider
                raw_response: simple_response,
            })
        })
    }
}

// We might need From/TryFrom implementations for errors if the llm crate has specific error types
// that need to be converted to CompletionError.
// For example:
// impl From<llm::Error> for CompletionError {
//     fn from(error: llm::Error) -> Self {
//         // Conversion logic here
//         CompletionError::Provider(error.to_string())
//     }
// }

[kyegomez]

@Tails this would be extremely helpful if you could submit a PR for this update. We're still building our rust team : )

We're hiring quite extensively: swarms.ai/hiring