feat: Add EnsureDownloaded fluent method (#121)

* feat:  Add EnsureDownloaded fluent method

Fixes #115

- add EnsureDownloaded method

* cleanup and use EnsureDownloaded in the example

* error is throw when invalid model type is provided

Author: srebrek

Examples/Examples/Chat/ChatExample.cs

@@ -12,6 +12,7 @@ public async Task Start()
         // Using strongly-typed model
         await AIHub.Chat()
             .WithModel<Gemma2_2b>()
+            .EnsureModelDownloaded()
             .WithMessage("Where do hedgehogs goes at night?")
             .CompleteAsync(interactive: true);
     }

src/MaIN.Core/Hub/Contexts/AgentContext.cs

@@ -1,16 +1,17 @@
 using MaIN.Core.Hub.Contexts.Interfaces.AgentContext;
+using MaIN.Core.Hub.Utils;
 using MaIN.Domain.Configuration;
 using MaIN.Domain.Entities;
 using MaIN.Domain.Entities.Agents;
 using MaIN.Domain.Entities.Agents.AgentSource;
-using MaIN.Domain.Models;
-using MaIN.Services.Services.Abstract;
-using MaIN.Services.Services.Models;
-using MaIN.Core.Hub.Utils;
 using MaIN.Domain.Entities.Agents.Knowledge;
 using MaIN.Domain.Entities.Tools;
 using MaIN.Domain.Exceptions.Agents;
+using MaIN.Domain.Models;
+using MaIN.Domain.Models.Abstract;
 using MaIN.Services.Constants;
+using MaIN.Services.Services.Abstract;
+using MaIN.Services.Services.Models;
 
 namespace MaIN.Core.Hub.Contexts;
 
@@ -20,6 +21,7 @@ public sealed class AgentContext : IAgentBuilderEntryPoint, IAgentConfigurationB
     private InferenceParams? _inferenceParams;
     private MemoryParams? _memoryParams;
     private bool _disableCache;
+    private bool _ensureModelDownloaded;
     private readonly Agent _agent;
     internal Knowledge? _knowledge;
 
@@ -60,8 +62,8 @@ internal AgentContext(IAgentService agentService, Agent existingAgent)
     public async Task<Agent?> GetAgentById(string id) => await _agentService.GetAgentById(id);
     public async Task Delete() => await _agentService.DeleteAgent(_agent.Id);
     public async Task<bool> Exists() => await _agentService.AgentExists(_agent.Id);
-    
-    
+
+
     public IAgentConfigurationBuilder WithModel(string model)
     {
         _agent.Model = model;
@@ -82,18 +84,18 @@ public async Task<IAgentContextExecutor> FromExisting(string agentId)
         {
             throw new AgentNotFoundException(agentId);
         }
-            
+
         var context = new AgentContext(_agentService, existingAgent);
         context.LoadExistingKnowledgeIfExists();
         return context;
     }
-    
+
     public IAgentConfigurationBuilder WithInitialPrompt(string prompt)
     {
         _agent.Context.Instruction = prompt;
         return this;
     }
-    
+
     public IAgentConfigurationBuilder WithId(string id)
     {
         _agent.Id = id;
@@ -112,6 +114,12 @@ public IAgentConfigurationBuilder DisableCache()
         return this;
     }
 
+    public IAgentConfigurationBuilder EnsureModelDownloaded()
+    {
+        _ensureModelDownloaded = true;
+        return this;
+    }
+
     public IAgentConfigurationBuilder WithSource(IAgentSource source, AgentSourceType type)
     {
         _agent.Context.Source = new AgentSource()
@@ -121,7 +129,7 @@ public IAgentConfigurationBuilder WithSource(IAgentSource source, AgentSourceTyp
         };
         return this;
     }
-    
+
     public IAgentConfigurationBuilder WithName(string name)
     {
         _agent.Name = name;
@@ -143,7 +151,7 @@ public IAgentConfigurationBuilder WithMcpConfig(Mcp mcpConfig)
         _agent.Context.McpConfig = mcpConfig;
         return this;
     }
-    
+
     public IAgentConfigurationBuilder WithInferenceParams(InferenceParams inferenceParams)
     {
         _inferenceParams = inferenceParams;
@@ -174,7 +182,7 @@ public IAgentConfigurationBuilder WithKnowledge(KnowledgeBuilder knowledge)
         _knowledge = knowledge.ForAgent(_agent).Build();
         return this;
     }
-    
+
     public IAgentConfigurationBuilder WithKnowledge(Knowledge knowledge)
     {
         _knowledge = knowledge;
@@ -189,7 +197,7 @@ public IAgentConfigurationBuilder WithInMemoryKnowledge(Func<KnowledgeBuilder, K
         _knowledge = knowledgeConfig(builder).Build();
         return this;
     }
-    
+
     public IAgentConfigurationBuilder WithBehaviour(string name, string instruction)
     {
         _agent.Behaviours ??= new Dictionary<string, string>();
@@ -200,10 +208,15 @@ public IAgentConfigurationBuilder WithBehaviour(string name, string instruction)
 
     public async Task<IAgentContextExecutor> CreateAsync(bool flow = false, bool interactiveResponse = false)
     {
+        if (_ensureModelDownloaded && !string.IsNullOrWhiteSpace(_agent.Model))
+        {
+            await AIHub.Model().EnsureDownloadedAsync(_agent.Model);
+        }
+
         await _agentService.CreateAgent(_agent, flow, interactiveResponse, _inferenceParams, _memoryParams, _disableCache);
         return this;
     }
-    
+
     public IAgentContextExecutor Create(bool flow = false, bool interactiveResponse = false)
     {
         _ = _agentService.CreateAgent(_agent, flow, interactiveResponse, _inferenceParams, _memoryParams, _disableCache).Result;
@@ -215,7 +228,7 @@ public IAgentConfigurationBuilder WithTools(ToolsConfiguration toolsConfiguratio
         _agent.ToolsConfiguration = toolsConfiguration;
         return this;
     }
-    
+
     internal void LoadExistingKnowledgeIfExists()
     {
         _knowledge ??= new Knowledge(_agent);
@@ -229,7 +242,7 @@ internal void LoadExistingKnowledgeIfExists()
             Console.WriteLine("Knowledge cannot be loaded - new one will be created");
         }
     }
-    
+
     public async Task<ChatResult> ProcessAsync(Chat chat, bool translate = false)
     {
         if (_knowledge == null)
@@ -247,7 +260,7 @@ public async Task<ChatResult> ProcessAsync(Chat chat, bool translate = false)
             CreatedAt = DateTime.Now
         };
     }
-    
+
     public async Task<ChatResult> ProcessAsync(
         string message,
         bool translate = false,
@@ -276,8 +289,8 @@ public async Task<ChatResult> ProcessAsync(
             CreatedAt = DateTime.Now
         };
     }
-    
-    public async Task<ChatResult> ProcessAsync(Message message, 
+
+    public async Task<ChatResult> ProcessAsync(Message message,
         bool translate = false,
         Func<LLMTokenValue, Task>? tokenCallback = null,
         Func<ToolInvocation, Task>? toolCallback = null)
@@ -288,7 +301,7 @@ public async Task<ChatResult> ProcessAsync(Message message,
         }
         var chat = await _agentService.GetChatByAgent(_agent.Id);
         chat.Messages.Add(message);
-        var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, tokenCallback, toolCallback);;
+        var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, tokenCallback, toolCallback);
         var messageResult = result.Messages.LastOrDefault()!;
         return new ChatResult()
         {
@@ -298,7 +311,7 @@ public async Task<ChatResult> ProcessAsync(Message message,
             CreatedAt = DateTime.Now
         };
     }
-    
+
     public async Task<ChatResult> ProcessAsync(
         IEnumerable<Message> messages,
         bool translate = false,
@@ -317,7 +330,7 @@ public async Task<ChatResult> ProcessAsync(
             chat.Messages.Add(systemMsg);
         }
         chat.Messages.AddRange(messages);
-        var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, tokenCallback, toolCallback);;
+        var result = await _agentService.Process(chat, _agent.Id, _knowledge, translate, tokenCallback, toolCallback);
         var messageResult = result.Messages.LastOrDefault()!;
         return new ChatResult()
         {
@@ -335,7 +348,7 @@ public static async Task<AgentContext> FromExisting(IAgentService agentService,
         {
             throw new AgentNotFoundException(agentId);
         }
-            
+
         var context = new AgentContext(agentService, existingAgent);
         context.LoadExistingKnowledgeIfExists();
         return context;
@@ -345,8 +358,8 @@ public static async Task<AgentContext> FromExisting(IAgentService agentService,
 public static class AgentExtensions
 {
     public static async Task<ChatResult> ProcessAsync(
-        this Task<AgentContext> agentTask, 
-        string message, 
+        this Task<AgentContext> agentTask,
+        string message,
         bool translate = false)
     {
         var agent = await agentTask;

src/MaIN.Core/Hub/Contexts/ChatContext.cs

@@ -18,6 +18,7 @@ public sealed class ChatContext : IChatBuilderEntryPoint, IChatMessageBuilder, I
 {
     private readonly IChatService _chatService;
     private bool _preProcess;
+    private bool _ensureModelDownloaded;
     private readonly Chat _chat;
     private List<FileInfo> _files = [];
 
@@ -88,7 +89,13 @@ public IChatMessageBuilder EnableVisual()
         _chat.Visual = true;
         return this;
     }
-    
+
+    public IChatMessageBuilder EnsureModelDownloaded()
+    {
+        _ensureModelDownloaded = true;
+        return this;
+    }
+
     public IChatConfigurationBuilder WithInferenceParams(InferenceParams inferenceParams)
     {
         _chat.InterferenceParams = inferenceParams;
@@ -194,7 +201,7 @@ public IChatConfigurationBuilder DisableCache()
         _chat.Properties.AddProperty(ServiceConstants.Properties.DisableCacheProperty);
         return this;
     }
-    
+
     public async Task<ChatResult> CompleteAsync(
         bool translate = false, // Move to WithTranslate
         bool interactive = false, // Move to WithInteractive
@@ -208,13 +215,18 @@ public async Task<ChatResult> CompleteAsync(
         {
             throw new EmptyChatException(_chat.Id);
         }
-        
+
+        if (_ensureModelDownloaded)
+        {
+            await AIHub.Model().EnsureDownloadedAsync(_chat.ModelId);
+        }
+
         _chat.Messages.Last().Files = _files;
-        if(_preProcess)
+        if (_preProcess)
         {
             _chat.Messages.Last().Properties.AddProperty(ServiceConstants.Properties.PreProcessProperty);
         }
-        
+
         if (!await ChatExists(_chat.Id))
         {
             await _chatService.Create(_chat);
@@ -227,8 +239,8 @@ public async Task<ChatResult> CompleteAsync(
     public async Task<IChatConfigurationBuilder> FromExisting(string chatId)
     {
         var existing = await _chatService.GetById(chatId);
-        return existing == null 
-            ? throw new ChatNotFoundException(chatId) 
+        return existing == null
+            ? throw new ChatNotFoundException(chatId)
             : new ChatContext(_chatService, existing);
     }
 
@@ -244,12 +256,9 @@ private async Task<bool> ChatExists(string id)
             return false;
         }
     }
-    
-    IChatMessageBuilder IChatMessageBuilder.EnableVisual() => EnableVisual();
 
- 
     public string GetChatId() => _chat.Id;
-    
+
     public async Task<Chat> GetCurrentChat()
     {
         if (_chat.Id == null)
@@ -271,7 +280,7 @@ public async Task DeleteChat()
 
         await _chatService.Delete(_chat.Id);
     }
-    
+
     public List<MessageShort> GetChatHistory()
     {
         return [.. _chat.Messages.Select(x => new MessageShort()

src/MaIN.Core/Hub/Contexts/Interfaces/AgentContext/IAgentConfigurationBuilder.cs

@@ -9,6 +9,14 @@ namespace MaIN.Core.Hub.Contexts.Interfaces.AgentContext;
 
 public interface IAgentConfigurationBuilder : IAgentActions
 {
+    /// <summary>
+    /// Flags the agent to automatically ensure the selected local model is downloaded before creation.
+    /// If the model is already present the download is skipped; cloud models are silently ignored.
+    /// The actual download is deferred until <see cref="CreateAsync"/> is called.
+    /// </summary>
+    /// <returns>The context instance implementing <see cref="IAgentConfigurationBuilder"/> for method chaining.</returns>
+    IAgentConfigurationBuilder EnsureModelDownloaded();
+
     /// <summary>
     /// Sets the initial prompt for the agent. This prompt serves as an instruction or context that guides the agent's behavior during its execution.
     /// </summary>

src/MaIN.Core/Hub/Contexts/Interfaces/ChatContext/IChatMessageBuilder.cs

@@ -10,6 +10,14 @@ public interface IChatMessageBuilder : IChatActions
     /// </summary>
     /// <returns>The context instance implementing <see cref="IChatMessageBuilder"/> for method chaining.</returns>
     IChatMessageBuilder EnableVisual();
+
+    /// <summary>
+    /// Flags the chat to automatically ensure the selected local model is downloaded before completing.
+    /// If the model is already present the download is skipped; cloud models are silently ignored.
+    /// The actual download is deferred until <see cref="IChatConfigurationBuilder.CompleteAsync"/> is called.
+    /// </summary>
+    /// <returns>The context instance implementing <see cref="IChatMessageBuilder"/> for method chaining.</returns>
+    IChatMessageBuilder EnsureModelDownloaded();
 
     /// <summary>
     /// Adds a user message to the chat. This method captures the message content and assigns the "User" role to it.

src/MaIN.Core/Hub/Contexts/Interfaces/ModelContext/IModelContext.cs

@@ -52,33 +52,27 @@ public interface IModelContext
     Task<IModelContext> DownloadAsync(string modelId, CancellationToken cancellationToken = default);
 
     /// <summary>
-    /// Asynchronously downloads a custom model from a specified URL. This method allows downloading models that are not part
-    /// of the known models collection, adding them to the system after download.
+    /// Ensures a known local model is downloaded before use. If the model is already present on disk the call
+    /// returns immediately; if not, the model is downloaded. Cloud models are silently skipped.
+    /// Thread-safe: concurrent calls for the same model will not trigger duplicate downloads.
     /// </summary>
-    /// <param name="model">The name to assign to the downloaded model.</param>
-    /// <param name="url">The URL from which to download the model.</param>
-    /// <returns>A task that represents the asynchronous download operation that completes when the download finishes,
-    /// returning the context instance implementing <see cref="IModelContext"/> for method chaining.</returns>
-    Task<IModelContext> DownloadAsync(string model, string url, CancellationToken cancellationToken = default);
-
-    /// <summary>
-    /// Synchronously downloads a known model from its configured download URL. This is the blocking version of the download operation
-    /// with progress tracking.
-    /// </summary>
-    /// <param name="modelName">The name of the model to download.</param>
-    /// <returns>The context instance implementing <see cref="IModelContext"/> for method chaining.</returns>
-    [Obsolete("Use DownloadAsync instead")]
-    IModelContext Download(string modelName);
+    /// <param name="modelId">The id of the model to ensure is downloaded.</param>
+    /// <param name="cancellationToken">Optional cancellation token to abort the download operation.</param>
+    /// <returns>A task that represents the asynchronous operation, returning the context instance implementing
+    /// <see cref="IModelContext"/> for method chaining.</returns>
+    Task<IModelContext> EnsureDownloadedAsync(string modelId, CancellationToken cancellationToken = default);
 
     /// <summary>
-    /// Synchronously downloads a custom model from a specified URL. This method provides blocking download functionality
-    /// for custom models not in the known models collection.
+    /// Ensures a known local model is downloaded before use using a strongly-typed model reference.
+    /// If the model is already present on disk the call returns immediately; if not, the model is downloaded.
+    /// Cloud models are silently skipped.
+    /// Thread-safe: concurrent calls for the same model will not trigger duplicate downloads.
     /// </summary>
-    /// <param name="model">The name to assign to the downloaded model.</param>
-    /// <param name="url">The URL from which to download the model.</param>
-    /// <returns>The context instance implementing <see cref="IModelContext"/> for method chaining.</returns>
-    [Obsolete("Use DownloadAsync instead")]
-    IModelContext Download(string model, string url);
+    /// <typeparam name="TModel">A <see cref="LocalModel"/> type with a parameterless constructor.</typeparam>
+    /// <param name="cancellationToken">Optional cancellation token to abort the download operation.</param>
+    /// <returns>A task that represents the asynchronous operation, returning the context instance implementing
+    /// <see cref="IModelContext"/> for method chaining.</returns>
+    Task<IModelContext> EnsureDownloadedAsync<TModel>(CancellationToken cancellationToken = default) where TModel : LocalModel, new();
 
     /// <summary>
     /// Loads a model into the memory cache for faster access during inference operations. This method preloads the model to avoid loading