Skip to content

Commit 13aa063

Browse files
srebreksrebrek
committed
Separate AIModel and Chat
- remove Backend field from Chat as it is contained in the AIModel
1 parent a32de60 commit 13aa063

18 files changed

Lines changed: 450 additions & 461 deletions

File tree

‎src/MaIN.Core.UnitTests/ChatContextTests.cs‎

Lines changed: 44 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ public void Constructor_ShouldInitializeNewChat()
2727
{
2828
// Act
2929
var chatId = _chatContext.GetChatId();
30-
30+
3131
// Assert
3232
Assert.NotNull(chatId);
3333
Assert.NotEmpty(chatId);
@@ -39,7 +39,7 @@ public void WithMessage_ShouldAddMessage()
3939
// Act
4040
_chatContext.WithMessage("Hello, world!");
4141
var messages = _chatContext.GetChatHistory();
42-
42+
4343
// Assert
4444
Assert.Single(messages);
4545
Assert.Equal("Hello, world!", messages[0].Content);
@@ -53,7 +53,7 @@ public void WithSystemPrompt_ShouldInsertSystemMessageAtBeginning()
5353
_chatContext.WithMessage("User message");
5454
_chatContext.WithSystemPrompt("System prompt");
5555
var messages = _chatContext.GetChatHistory();
56-
56+
5757
// Assert
5858
Assert.Equal(2, messages.Count);
5959
Assert.Equal("System", messages[0].Role);
@@ -65,11 +65,14 @@ public void WithFiles_ShouldAttachFilesToLastMessage()
6565
{
6666
// Arrange
6767
_chatContext.WithMessage("User message");
68-
var files = new List<FileInfo> { new() { Name = "file.txt", Path = "/path/file.txt", Extension = "txt"} };
69-
68+
var files = new List<FileInfo>
69+
{
70+
new() { Name = "file.txt", Path = "/path/file.txt", Extension = "txt" }
71+
};
72+
7073
// Act
7174
_chatContext.WithFiles(files);
72-
75+
7376
// Assert
7477
var lastMessage = _chatContext.GetChatHistory().Last();
7578
Assert.NotNull(lastMessage);
@@ -79,55 +82,74 @@ public void WithFiles_ShouldAttachFilesToLastMessage()
7982
public async Task CompleteAsync_ShouldCallChatService()
8083
{
8184
// Arrange
82-
var chatResult = new ChatResult(){ Model = "test-model", Message = new Message
85+
var chatResult = new ChatResult()
86+
{
87+
Model = "test-model",
88+
Message = new Message
8389
{
8490
Role = "Assistant",
8591
Content = "test-message",
8692
Type = MessageType.LocalLLM
8793
}
8894
};
89-
9095

91-
_mockChatService.Setup(s => s.Completions(It.IsAny<Chat>(), It.IsAny<bool>(), It.IsAny<bool>(), null, It.IsAny<CancellationToken>()))
96+
_mockChatService
97+
.Setup(s => s.Completions(
98+
It.IsAny<Chat>(),
99+
It.IsAny<bool>(),
100+
It.IsAny<bool>(),
101+
null,
102+
It.IsAny<CancellationToken>()))
92103
.ReturnsAsync(chatResult);
93-
104+
94105
_chatContext.WithMessage("User message");
95-
_chatContext.WithModel(new GenericLocalModel("test-model"));
106+
_chatContext.WithModel(_testModelId);
96107

97108
// Act
98109
var result = await _chatContext.CompleteAsync();
99-
110+
100111
// Assert
101-
_mockChatService.Verify(s => s.Completions(It.IsAny<Chat>(), false, false, null, It.IsAny<CancellationToken>()), Times.Once);
112+
_mockChatService.Verify(
113+
s => s.Completions(
114+
It.IsAny<Chat>(),
115+
false,
116+
false,
117+
null,
118+
It.IsAny<CancellationToken>()),
119+
Times.Once);
102120
Assert.Equal(chatResult, result);
103121
}
104122

105123
[Fact]
106124
public async Task GetCurrentChat_ShouldCallChatService()
107125
{
108126
// Arrange
109-
var chat = new Chat { Id = _chatContext.GetChatId(), ModelId = _testModelId , Name = "test"};
127+
var chat = new Chat { Id = _chatContext.GetChatId(), ModelId = _testModelId, Name = "test" };
110128
_mockChatService.Setup(s => s.GetById(chat.Id)).ReturnsAsync(chat);
111-
129+
112130
// Act
113131
var result = await _chatContext.GetCurrentChat();
114-
132+
115133
// Assert
116134
Assert.Equal(chat, result);
117135
}
118136

119137
[Fact]
120138
public async Task WithModel_ShouldSetModelIdAndInstance()
121139
{
122-
// Arrange
123-
var model = new GenericLocalModel(_testModelId);
124-
125140
// Act
126-
await _chatContext.WithModel(model)
141+
await _chatContext.WithModel(_testModelId)
127142
.WithMessage("User message")
128143
.CompleteAsync();
129-
144+
130145
// Assert
131-
_mockChatService.Verify(s => s.Completions(It.Is<Chat>(c => c.ModelId == _testModelId && c.ModelInstance == model), It.IsAny<bool>(), It.IsAny<bool>(), null, It.IsAny<CancellationToken>()), Times.Once);
146+
_mockChatService.Verify(s =>
147+
s.Completions(
148+
It.Is<Chat>(c => c.ModelId == _testModelId),
149+
It.IsAny<bool>(),
150+
It.IsAny<bool>(),
151+
null,
152+
It.IsAny<CancellationToken>()),
153+
Times.Once);
132154
}
133155
}

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

Lines changed: 38 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
using MaIN.Core.Hub.Contexts.Interfaces.ChatContext;
2-
using MaIN.Domain.Configuration;
32
using MaIN.Domain.Entities;
43
using MaIN.Domain.Entities.Tools;
54
using MaIN.Domain.Exceptions.Chats;
@@ -42,50 +41,41 @@ internal ChatContext(IChatService chatService, Chat existingChat)
4241

4342
public IChatMessageBuilder WithModel(AIModel model, bool? imageGen = null)
4443
{
45-
SetModel(model);
46-
_chat.ImageGen = imageGen ?? model is IImageGenerationModel;
44+
ModelRegistry.RegisterOrReplace(model);
45+
_chat.ModelId = model.Id;
46+
_chat.ImageGen = imageGen ?? (model is IImageGenerationModel);
47+
_chat.ImageGen = model.HasImageGeneration;
4748
return this;
4849
}
4950

51+
[Obsolete("Use WithModel(string modelId) or WithModel(AIModel model) instead.")]
5052
public IChatMessageBuilder WithModel<TModel>() where TModel : AIModel, new()
5153
{
5254
var model = new TModel();
53-
return WithModel(model);
55+
ModelRegistry.RegisterOrReplace(model);
56+
_chat.ModelId = model.Id;
57+
return this;
5458
}
5559

56-
[Obsolete("Use WithModel(AIModel model) or WithModel<TModel>() instead.")]
5760
public IChatMessageBuilder WithModel(string modelId)
5861
{
59-
AIModel model;
60-
try
62+
if (!ModelRegistry.Exists(modelId))
6163
{
62-
model = ModelRegistry.GetById(modelId);
64+
throw new ModelNotRegisteredException(modelId);
6365
}
64-
catch (Exception ex)
65-
{
6666

67-
throw new Exception($"Model with ID '{modelId}' not found in registry. Use WithModel(AIModel model) or WithModel<TModel>() instead.", ex);
68-
}
69-
SetModel(model);
67+
_chat.ModelId = modelId;
7068
return this;
7169
}
7270

73-
[Obsolete("Use WithModel<TModel>() instead.")]
71+
[Obsolete("Use WithModel(AIModel model) instead.")]
7472
public IChatMessageBuilder WithCustomModel(string model, string path, string? mmProject = null)
7573
{
7674
KnownModels.AddModel(model, path, mmProject);
7775
_chat.ModelId = model;
7876
return this;
7977
}
8078

81-
private void SetModel(AIModel model)
82-
{
83-
_chat.ModelId = model.Id;
84-
_chat.ModelInstance = model;
85-
_chat.Backend = model.Backend;
86-
_chat.ImageGen = model.HasImageGeneration;
87-
}
88-
8979
public IChatMessageBuilder EnsureModelDownloaded()
9080
{
9181
_ensureModelDownloaded = true;
@@ -117,12 +107,6 @@ public IChatConfigurationBuilder Speak(TextToSpeechParams speechParams)
117107
return this;
118108
}
119109

120-
public IChatConfigurationBuilder WithBackend(BackendType backendType)
121-
{
122-
_chat.Backend = backendType;
123-
return this;
124-
}
125-
126110
public IChatConfigurationBuilder WithSystemPrompt(string systemPrompt)
127111
{
128112
var message = new Message
@@ -139,7 +123,15 @@ public IChatConfigurationBuilder WithSystemPrompt(string systemPrompt)
139123

140124
public IChatConfigurationBuilder WithMessage(string content)
141125
{
142-
_chat.Messages.Add(new Message { Role = "User", Content = content, Type = MessageType.LocalLLM, Time = DateTime.Now });
126+
var message = new Message
127+
{
128+
Role = "User",
129+
Content = content,
130+
Type = MessageType.NotSet,
131+
Time = DateTime.Now
132+
};
133+
134+
_chat.Messages.Add(message);
143135
return this;
144136
}
145137

@@ -204,18 +196,19 @@ public async Task<ChatResult> CompleteAsync(
204196
Func<LLMTokenValue?, Task>? changeOfValue = null,
205197
CancellationToken cancellationToken = default)
206198
{
207-
if (_chat.ModelInstance is null)
199+
if (string.IsNullOrEmpty(_chat.ModelId))
208200
{
209-
throw new MissingModelInstanceException();
201+
throw new MissingModelIdException(nameof(_chat.ModelId));
210202
}
203+
211204
if (_chat.Messages.Count == 0)
212205
{
213206
throw new EmptyChatException(_chat.Id);
214207
}
215208

216209
if (_ensureModelDownloaded)
217210
{
218-
await AIHub.Model().EnsureDownloadedAsync(_chat.ModelId);
211+
await AIHub.Model().EnsureDownloadedAsync(_chat.ModelId, cancellationToken);
219212
}
220213

221214
_chat.Messages.Last().Files = _files;
@@ -228,15 +221,21 @@ public async Task<ChatResult> CompleteAsync(
228221
{
229222
await _chatService.Create(_chat);
230223
}
231-
var result = await _chatService.Completions(_chat, translate, interactive, changeOfValue, cancellationToken);
224+
225+
var result = await _chatService.Completions(
226+
_chat,
227+
translate,
228+
interactive,
229+
changeOfValue,
230+
cancellationToken);
232231
_files = [];
233232
return result;
234233
}
235234

236235
public async Task<IChatConfigurationBuilder> FromExisting(string chatId)
237236
{
238237
var existing = await _chatService.GetById(chatId);
239-
return existing == null
238+
return existing is null
240239
? throw new ChatNotFoundException(chatId)
241240
: new ChatContext(_chatService, existing);
242241
}
@@ -258,19 +257,16 @@ private async Task<bool> ChatExists(string id)
258257

259258
public async Task<Chat> GetCurrentChat()
260259
{
261-
if (_chat.Id == null)
262-
{
263-
throw new ChatNotInitializedException();
264-
}
265-
266-
return await _chatService.GetById(_chat.Id);
260+
return _chat.Id is null
261+
? throw new ChatNotInitializedException()
262+
: await _chatService.GetById(_chat.Id);
267263
}
268264

269265
public async Task<List<Chat>> GetAllChats() => await _chatService.GetAll();
270266

271267
public async Task DeleteChat()
272268
{
273-
if (_chat.Id == null)
269+
if (_chat.Id is null)
274270
{
275271
throw new ChatNotInitializedException();
276272
}
@@ -287,4 +283,4 @@ public List<MessageShort> GetChatHistory()
287283
Time = x.Time
288284
})];
289285
}
290-
}
286+
}

0 commit comments

Comments
 (0)