Skip to content

Commit 38f138e

Browse files
Royna2544Royna2544
committed
Enhance LLMCore error handling and improve query function signature
1 parent b1d0951 commit 38f138e

3 files changed

Lines changed: 125 additions & 47 deletions

File tree

‎src/command_modules/ask.cpp‎

Lines changed: 59 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,16 @@
66
#include <tgbot/TgException.h>
77

88
#include <algorithm>
9+
#include <cstdint>
910
#include <api/CommandModule.hpp>
1011
#include <api/MessageExt.hpp>
1112
#include <api/TgBotApi.hpp>
1213
#include <filesystem>
14+
#include <memory>
15+
#include <mutex>
1316
#include <nlohmann/json.hpp>
1417
#include <string_view>
18+
#include <unordered_map>
1519

1620
#include "CurlUtils.hpp"
1721
#include "llm/LMStudioApi.hpp"
@@ -21,27 +25,30 @@
2125
#ifdef ASK_ENABLE_LOCAL_LLM
2226
#include "llm/LLMCore.hpp"
2327

24-
constexpr const char* kLLMStorageKey = "LLMCoreInstance";
25-
2628
struct LLMCoreInstance {
2729
LLMCore llm;
2830
LLMCore::Model model;
31+
std::filesystem::path modelPath;
2932
};
3033

3134
// Handler for local model LLM <i.e. With supported GPU>
3235
static void localmodelhandler(TgBotApi::Ptr api, MessageExt* message,
33-
const StringResLoader::PerLocaleMap* res,
34-
const Providers* provider,
3536
const std::filesystem::path& modelPath) {
36-
LLMCoreInstance* instance = nullptr;
37+
static std::mutex llm_mutex;
38+
static std::unique_ptr<LLMCoreInstance> instance;
3739

3840
auto sent = api->sendReplyMessage(message->message(),
3941
"Processing your query, please wait...");
4042

41-
if (provider->globalStorage->hasStorage(kLLMStorageKey)) {
42-
instance = provider->globalStorage->getStorage(kLLMStorageKey)
43-
.getAs<LLMCoreInstance>();
43+
std::unique_lock<std::mutex> lock(llm_mutex, std::defer_lock);
44+
if (!lock.try_lock()) {
45+
api->editMessage(sent,
46+
"Another LLM request is running. Waiting for it to "
47+
"finish...");
48+
lock.lock();
49+
}
4450

51+
if (instance != nullptr && instance->modelPath == modelPath) {
4552
if (!message->has<MessageAttrs::ExtraText>()) {
4653
auto info = instance->model.info();
4754
api->editMessage(
@@ -57,19 +64,23 @@ File Size: {} bytes)",
5764
return;
5865
}
5966
} else {
60-
provider->globalStorage->getStorage(kLLMStorageKey) =
61-
SharedMalloc::fromType<LLMCoreInstance>();
62-
63-
instance = provider->globalStorage->getStorage(kLLMStorageKey)
64-
.getAs<LLMCoreInstance>();
67+
if (!message->has<MessageAttrs::ExtraText>()) {
68+
api->editMessage(
69+
sent,
70+
"Please provide a query after the command to ask the LLM.");
71+
return;
72+
}
6573

6674
api->editMessage(sent, "Initializing LLM core, please wait...");
67-
if (!instance->model.load(modelPath)) {
75+
instance.reset();
76+
auto new_instance = std::make_unique<LLMCoreInstance>();
77+
if (!new_instance->model.load(modelPath)) {
6878
LOG(ERROR) << "Failed to initialize LLM core.";
6979
api->editMessage(sent, "Error: Unable to initialize LLM core.");
70-
provider->globalStorage->removeStorage(kLLMStorageKey);
7180
return;
7281
}
82+
new_instance->modelPath = modelPath;
83+
instance = std::move(new_instance);
7384
}
7485

7586
if (!message->has<MessageAttrs::ExtraText>()) {
@@ -79,9 +90,17 @@ File Size: {} bytes)",
7990
}
8091

8192
std::string query = message->get<MessageAttrs::ExtraText>();
82-
auto response_opt = instance->llm.query(&instance->model, query, 512);
93+
LLMCore::QueryError query_error = LLMCore::QueryError::None;
94+
auto response_opt =
95+
instance->llm.query(&instance->model, query, 512, &query_error);
8396
if (!response_opt) {
8497
LOG(ERROR) << "LLM query failed.";
98+
if (query_error == LLMCore::QueryError::PromptTooLong) {
99+
api->editMessage(sent,
100+
"Error: Prompt is too long for the local LLM "
101+
"context. Please shorten it and try again.");
102+
return;
103+
}
85104
api->editMessage(sent, "Error: LLM query failed.");
86105
return;
87106
}
@@ -107,8 +126,6 @@ File Size: {} bytes)",
107126
}
108127
#endif // ASK_ENABLE_LOCAL_LLM
109128

110-
constexpr float kTemperature = 0.7F;
111-
112129
static void localnetmodelhandler(TgBotApi::Ptr api, MessageExt* message,
113130
const StringResLoader::PerLocaleMap* res,
114131
const Providers* provider,
@@ -165,7 +182,6 @@ static void localnetmodelhandler(TgBotApi::Ptr api, MessageExt* message,
165182
return m.type == LMStudioApi::LLMType::llm;
166183
});
167184

168-
static std::string response_key;
169185
if (that == modelResponse.models.end()) {
170186
LOG(ERROR) << "No LLM type models available from LLM server.";
171187
api->editMessage(
@@ -174,19 +190,18 @@ static void localnetmodelhandler(TgBotApi::Ptr api, MessageExt* message,
174190
}
175191
chatRequest.model = that->key;
176192
LOG(INFO) << "Using model " << chatRequest.model << " for query.";
177-
chatRequest.temperature = 0.2f;
178193
chatRequest.system_prompt = SYSTEM_PROMPT;
179194
chatRequest.input = query;
180-
std::vector<LMStudioApi::ChatRequest::Plugin> plugins{};
181-
plugins.emplace_back(LMStudioApi::ChatRequest::Plugin{
182-
.type = "plugin",
183-
.id = "mcp/playwright",
184-
});
185-
chatRequest.integrations = plugins;
186-
chatRequest.context_length = 64000;
187-
chatRequest.reasoning = LMStudioApi::Reasoning::medium;
188-
if (!response_key.empty()) {
189-
chatRequest.previous_response_id = response_key;
195+
196+
const std::int64_t chat_id = message->get<MessageAttrs::Chat>()->id;
197+
static std::mutex response_keys_mutex;
198+
static std::unordered_map<std::int64_t, std::string> response_keys;
199+
{
200+
std::lock_guard<std::mutex> lock(response_keys_mutex);
201+
if (auto it = response_keys.find(chat_id);
202+
it != response_keys.end() && !it->second.empty()) {
203+
chatRequest.previous_response_id = it->second;
204+
}
190205
}
191206

192207
nlohmann::json payload = chatRequest;
@@ -220,7 +235,14 @@ static void localnetmodelhandler(TgBotApi::Ptr api, MessageExt* message,
220235
"Error: LLM server response has no 'message' type output.");
221236
return;
222237
}
223-
response_key = chatResponse.response_id.value_or("");
238+
{
239+
std::lock_guard<std::mutex> lock(response_keys_mutex);
240+
if (chatResponse.response_id) {
241+
response_keys[chat_id] = *chatResponse.response_id;
242+
} else {
243+
response_keys.erase(chat_id);
244+
}
245+
}
224246
std::initializer_list<
225247
std::pair<absl::string_view, absl::string_view>>
226248
replacements = {
@@ -266,9 +288,12 @@ DECLARE_COMMAND_HANDLER(ask) {
266288

267289
if (type == "local") {
268290
#ifdef ASK_ENABLE_LOCAL_LLM
269-
localmodelhandler(api, message, res, provider, location);
291+
localmodelhandler(api, message, location);
270292
#else
271293
LOG(ERROR) << "Local LLM support is not enabled in this build.";
294+
api->sendMessage(
295+
message->get<MessageAttrs::Chat>(),
296+
"Error: Local LLM support is not enabled in this build.");
272297
#endif
273298
} else if (type == "localnet") {
274299
if (authkey) {
@@ -280,6 +305,8 @@ DECLARE_COMMAND_HANDLER(ask) {
280305
}
281306
} else {
282307
LOG(ERROR) << "Unsupported LLM configuration type: " << type;
308+
api->sendMessage(message->get<MessageAttrs::Chat>(),
309+
"Error: Unsupported LLM configuration type.");
283310
}
284311
}
285312

‎src/command_modules/llm/LLMCore.cpp‎

Lines changed: 53 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
1-
#include "LLMCore.hpp"
1+
#include "LLMCore.hpp"
22

33
#include <absl/log/log.h>
44

55
#include <algorithm>
6+
#include <array>
7+
#include <iomanip>
8+
#include <sstream>
69
#include <string>
710
#include <string_view>
811
#include <vector>
@@ -15,6 +18,7 @@ struct LLMCore::Model::Impl {
1518
llama_context* ctx = nullptr;
1619
llama_sampler* sampler = nullptr;
1720
llama_batch batch{};
21+
int n_batch = 0;
1822
int n_past = 0;
1923
int seq_id = 0;
2024
};
@@ -24,6 +28,8 @@ LLMCore::Model::~Model() { cleanup(); }
2428

2529
bool LLMCore::Model::load(const std::filesystem::path& modelPath,
2630
int n_gpu_layers, int n_ctx, int n_batch) {
31+
cleanup();
32+
2733
llama_model_params model_params = llama_model_default_params();
2834
llama_context_params ctx_params = llama_context_default_params();
2935

@@ -72,6 +78,7 @@ bool LLMCore::Model::load(const std::filesystem::path& modelPath,
7278
llama_sampler_chain_add(impl_->sampler, llama_sampler_init_dist(1234));
7379
// Prepare Batch
7480
impl_->batch = llama_batch_init(n_batch, 0, 1);
81+
impl_->n_batch = n_batch;
7582

7683
info();
7784

@@ -104,16 +111,21 @@ LLMCore::Model::Info LLMCore::Model::info() const {
104111
}
105112

106113
void LLMCore::Model::cleanup() {
114+
if (impl_->sampler != nullptr) {
115+
llama_sampler_free(impl_->sampler);
116+
impl_->sampler = nullptr;
117+
}
107118
if (impl_->ctx != nullptr) {
108119
llama_batch_free(impl_->batch);
109-
llama_sampler_free(impl_->sampler);
110120
llama_free(impl_->ctx);
111121
impl_->ctx = nullptr;
112122
}
113123
if (impl_->model != nullptr) {
114124
llama_model_free(impl_->model);
115125
impl_->model = nullptr;
116126
}
127+
impl_->n_batch = 0;
128+
impl_->n_past = 0;
117129
}
118130

119131
namespace {
@@ -153,6 +165,12 @@ std::vector<llama_token> tokenize(const llama_model* model,
153165
return tokens;
154166
}
155167

168+
void set_error(LLMCore::QueryError* error, LLMCore::QueryError value) {
169+
if (error != nullptr) {
170+
*error = value;
171+
}
172+
}
173+
156174
} // namespace
157175

158176
LLMCore::LLMCore() {
@@ -166,21 +184,21 @@ LLMCore::~LLMCore() {
166184
}
167185

168186
std::optional<LLMCore::Model::Response> LLMCore::query(
169-
const Model* model, const std::string_view prompt, int max_tokens) {
187+
const Model* model, const std::string_view prompt, int max_tokens,
188+
QueryError* error) {
189+
set_error(error, QueryError::None);
170190
auto start_time = std::chrono::steady_clock::now();
171191
LOG(INFO) << "LLMCore::query called with prompt: " << prompt;
172192

173-
// -----------------------
174-
// 1. Context length check
175-
// -----------------------
176-
const auto n_ctx = llama_n_ctx(model->impl_->ctx);
177-
if (model->impl_->n_past + max_tokens > n_ctx) {
178-
LOG(WARNING) << "Context overflow likely. Resetting KV cache.";
179-
model->impl_->n_past = 0;
193+
if (model == nullptr || model->impl_->model == nullptr ||
194+
model->impl_->ctx == nullptr || model->impl_->sampler == nullptr) {
195+
LOG(ERROR) << "LLM query requested before the model was loaded.";
196+
set_error(error, QueryError::ModelNotLoaded);
197+
return std::nullopt;
180198
}
181199

182200
// -----------------------
183-
// 2. Construct Chat Prompt (Mistral format)
201+
// 1. Construct Chat Prompt
184202
// -----------------------
185203
std::string system = SYSTEM_PROMPT;
186204

@@ -196,18 +214,39 @@ std::optional<LLMCore::Model::Response> LLMCore::query(
196214
LOG(INFO) << "Formatted prompt: " << std::quoted(formatted_prompt);
197215

198216
// -----------------------
199-
// 3. Tokenize (do NOT add BOS manually)
217+
// 2. Tokenize (do NOT add BOS manually)
200218
// -----------------------
201219
std::vector<llama_token> tokens =
202220
tokenize(model->impl_->model, formatted_prompt);
203221

204222
if (tokens.empty()) {
205223
LOG(ERROR) << "Tokenization failed for prompt: " << prompt;
224+
set_error(error, QueryError::TokenizationFailed);
206225
return std::nullopt;
207226
}
208227

209228
LOG(INFO) << "Tokenized prompt into " << tokens.size() << " tokens.";
210229

230+
// -----------------------
231+
// 3. Prompt size checks
232+
// -----------------------
233+
const auto n_ctx = llama_n_ctx(model->impl_->ctx);
234+
if (tokens.size() > static_cast<size_t>(model->impl_->n_batch) ||
235+
tokens.size() + static_cast<size_t>(max_tokens) >
236+
static_cast<size_t>(n_ctx)) {
237+
LOG(WARNING) << "Prompt is too long for the local LLM context. tokens="
238+
<< tokens.size() << ", max_tokens=" << max_tokens
239+
<< ", n_batch=" << model->impl_->n_batch
240+
<< ", n_ctx=" << n_ctx;
241+
set_error(error, QueryError::PromptTooLong);
242+
return std::nullopt;
243+
}
244+
245+
// Stateless per request: clear KV memory and sampler history.
246+
llama_memory_clear(llama_get_memory(model->impl_->ctx), true);
247+
llama_sampler_reset(model->impl_->sampler);
248+
model->impl_->n_past = 0;
249+
211250
// -----------------------
212251
// 4. Load prompt tokens
213252
// -----------------------
@@ -221,6 +260,7 @@ std::optional<LLMCore::Model::Response> LLMCore::query(
221260
// Decode
222261
if (llama_decode(model->impl_->ctx, model->impl_->batch) != 0) {
223262
LOG(ERROR) << "Decode failed for prompt";
263+
set_error(error, QueryError::DecodeFailed);
224264
return std::nullopt;
225265
}
226266

@@ -319,4 +359,4 @@ std::optional<LLMCore::Model::Response> LLMCore::query(
319359
r.duration = end_time - start_time;
320360

321361
return r;
322-
}
362+
}

‎src/command_modules/llm/LLMCore.hpp‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,14 @@ class LLMCore {
1414
LLMCore();
1515
~LLMCore();
1616

17+
enum class QueryError {
18+
None,
19+
ModelNotLoaded,
20+
TokenizationFailed,
21+
PromptTooLong,
22+
DecodeFailed,
23+
};
24+
1725
class Model {
1826
public:
1927
friend class LLMCore;
@@ -82,5 +90,8 @@ class LLMCore {
8290
*
8391
* Returns an optional Response struct containing the model's response
8492
*/
85-
std::optional<Model::Response> query(const Model* model, const std::string_view prompt, int max_tokens = 4096);
86-
};
93+
std::optional<Model::Response> query(const Model* model,
94+
std::string_view prompt,
95+
int max_tokens = 4096,
96+
QueryError* error = nullptr);
97+
};

0 commit comments

Comments
 (0)