Skip to content

feat: Add EnsureDownloaded fluent method #121

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
srebrek wants to merge 2 commits into
base: tests/fix-integration-tests
Choose a base branch
from
Open

feat: Add EnsureDownloaded fluent method #121

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cbd2297
feat: Add EnsureDownloaded fluent method
srebrek Feb 12, 2026
47c7f53
cleanup and use EnsureDownloaded in the example
srebrek Feb 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Examples/Examples/Chat/ChatExample.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
69 changes: 43 additions & 26 deletions src/MaIN.Core/Hub/Contexts/AgentContext.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;

Expand All @@ -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;

Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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()
Expand All @@ -121,7 +129,7 @@ public IAgentConfigurationBuilder WithSource(IAgentSource source, AgentSourceTyp
};
return this;
}

public IAgentConfigurationBuilder WithName(string name)
{
_agent.Name = name;
Expand All @@ -143,7 +151,7 @@ public IAgentConfigurationBuilder WithMcpConfig(Mcp mcpConfig)
_agent.Context.McpConfig = mcpConfig;
return this;
}

public IAgentConfigurationBuilder WithInferenceParams(InferenceParams inferenceParams)
{
_inferenceParams = inferenceParams;
Expand Down Expand Up @@ -174,7 +182,7 @@ public IAgentConfigurationBuilder WithKnowledge(KnowledgeBuilder knowledge)
_knowledge = knowledge.ForAgent(_agent).Build();
return this;
}

public IAgentConfigurationBuilder WithKnowledge(Knowledge knowledge)
{
_knowledge = knowledge;
Expand All @@ -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>();
Expand All @@ -200,11 +208,20 @@ public IAgentConfigurationBuilder WithBehaviour(string name, string instruction)

public async Task<IAgentContextExecutor> CreateAsync(bool flow = false, bool interactiveResponse = false)
{
if (_ensureModelDownloaded && !string.IsNullOrWhiteSpace(_agent.Model))
{
var model = ModelRegistry.GetById(_agent.Model);
if (model is LocalModel)
{
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)

public IAgentContextExecutor Create(bool flow = false, bool interactiveResponse = false) // I think it should be removed as there is a deadlock risk.
{
_ = _agentService.CreateAgent(_agent, flow, interactiveResponse, _inferenceParams, _memoryParams, _disableCache).Result;
return this;
Expand All @@ -215,7 +232,7 @@ public IAgentConfigurationBuilder WithTools(ToolsConfiguration toolsConfiguratio
_agent.ToolsConfiguration = toolsConfiguration;
return this;
}

internal void LoadExistingKnowledgeIfExists()
{
_knowledge ??= new Knowledge(_agent);
Expand All @@ -229,7 +246,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)
Expand All @@ -247,7 +264,7 @@ public async Task<ChatResult> ProcessAsync(Chat chat, bool translate = false)
CreatedAt = DateTime.Now
};
}

public async Task<ChatResult> ProcessAsync(
string message,
bool translate = false,
Expand Down Expand Up @@ -276,8 +293,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)
Expand All @@ -288,7 +305,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()
{
Expand All @@ -298,7 +315,7 @@ public async Task<ChatResult> ProcessAsync(Message message,
CreatedAt = DateTime.Now
};
}

public async Task<ChatResult> ProcessAsync(
IEnumerable<Message> messages,
bool translate = false,
Expand All @@ -317,7 +334,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()
{
Expand All @@ -335,7 +352,7 @@ public static async Task<AgentContext> FromExisting(IAgentService agentService,
{
throw new AgentNotFoundException(agentId);
}

var context = new AgentContext(agentService, existingAgent);
context.LoadExistingKnowledgeIfExists();
return context;
Expand All @@ -345,8 +362,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;
Expand Down
33 changes: 21 additions & 12 deletions src/MaIN.Core/Hub/Contexts/ChatContext.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [];

Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -208,13 +215,18 @@ public async Task<ChatResult> CompleteAsync(
{
throw new EmptyChatException(_chat.Id);
}


if (_ensureModelDownloaded && _chat.ModelInstance is LocalModel)
{
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);
Expand All @@ -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);
}

Expand All @@ -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)
Expand All @@ -271,7 +280,7 @@ public async Task DeleteChat()

await _chatService.Delete(_chat.Id);
}

public List<MessageShort> GetChatHistory()
{
return [.. _chat.Messages.Select(x => new MessageShort()
Expand Down
8 changes: 8 additions & 0 deletions src/MaIN.Core/Hub/Contexts/Interfaces/AgentContext/IAgentConfigurationBuilder.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
8 changes: 8 additions & 0 deletions src/MaIN.Core/Hub/Contexts/Interfaces/ChatContext/IChatMessageBuilder.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading