add:添加tokens限制版本。

This commit is contained in:
taynpg 2025-03-05 09:20:03 +08:00
parent 1648b6c961
commit abe5590ec9
7 changed files with 33 additions and 5 deletions

View File

@ -13,7 +13,8 @@
enum FrameType : int16_t {
TYPE_REQUEST = 0,
TYPE_RESPONSE_SUCCESS,
TYPE_RESPONSE_ERROR
TYPE_RESPONSE_ERROR,
TYPE_OUT_OF_LIMIT,
};
struct FrameData {

View File

@ -4,3 +4,4 @@ ApiEnvKey = TENCENT_DEEPSEEK_KEY
UserName = user
ModelName = deepseek-r1
AssistantName = assistant
MaxTokens = 50000

View File

@ -43,5 +43,11 @@ bool CConfig::parse_config(ConfigInfo& config, const std::string& config_path)
}
config.assistant_name = ini_handle.GetValue("Config", "AssistantName");
if (!ini_handle.KeyExists("Config", "MaxTokens")) {
std::cerr << "Not Key Found Config/MaxTokens in deepseek.ini" << std::endl;
return false;
}
config.max_tokens = ini_handle.GetLongValue("Config", "MaxTokens");
return true;
}

View File

@ -8,6 +8,7 @@ struct ConfigInfo {
std::string user_name;
std::string model_name;
std::string assistant_name;
long max_tokens{};
};
class CConfig

View File

@ -52,6 +52,7 @@ int main(int argc, char* argv[])
std::cout << "assis_name:" << config.assistant_name << std::endl;
std::cout << "base_url:" << config.base_url << std::endl;
std::cout << "model_name:" << config.model_name << std::endl;
std::cout << "max_tokens:" << config.max_tokens << std::endl;
std::cout << "api:" << show_api(key) << std::endl;
if (show_api(key) == "NULL") {
@ -85,6 +86,7 @@ int main(int argc, char* argv[])
// std::cout << "success." << std::endl;
server.set_worker(api, json);
server.set_token(config.max_tokens);
server.start();
io_context.run();

View File

@ -37,6 +37,11 @@ void Server::set_worker(std::shared_ptr<COpenAI> worker, std::shared_ptr<CJsonOp
json_ = json;
}
void Server::set_token(long tokens)
{
tokens_ = tokens;
}
void Server::do_accept()
{
auto socket = std::make_shared<asio::ip::tcp::socket>(io_context_);
@ -91,6 +96,13 @@ void Server::th_client(const std::shared_ptr<asio::ip::tcp::socket>& socket, con
if (frame == nullptr) {
break;
}
if (use_tokens_ > tokens_) {
std::cout << client_key << " tokens not enough" << std::endl;
FrameData req;
req.type = FrameType::TYPE_OUT_OF_LIMIT;
send_frame(socket, req);
continue;
}
std::cout << client_key << " 's data." << std::endl;
if (frame->type == FrameType::TYPE_REQUEST) {
ask_mutex_.lock();
@ -109,6 +121,10 @@ void Server::th_client(const std::shared_ptr<asio::ip::tcp::socket>& socket, con
req.data = new char[req.len];
req.protk = parse.prompt_tokens;
req.coptk = parse.completion_tokens;
use_tokens_ += req.protk;
use_tokens_ += req.coptk;
std::cout << "Already use " << use_tokens_ << " tokens.\r";
std::cout.flush();
memcpy(req.data, parse.message_content.c_str(), parse.message_content.size());
req.data[req.len - 1] = '\0';
send_frame(socket, req);

View File

@ -1,10 +1,10 @@
#ifndef SERVER_H
#define SERVER_H
#include "communicate.hpp"
#include "handle.h"
#include "jsondata.h"
#include "zapi.h"
#include "communicate.hpp"
#include <asio.hpp>
#include <mutex>
#include <unordered_map>
@ -27,11 +27,10 @@ public:
public:
void set_worker(std::shared_ptr<COpenAI> worker, std::shared_ptr<CJsonOper> json);
void set_token(long tokens);
private:
void do_accept();
void th_client(const std::shared_ptr<asio::ip::tcp::socket>& socket,
const std::string& client_key);
void th_client(const std::shared_ptr<asio::ip::tcp::socket>& socket, const std::string& client_key);
std::string post_data(const std::string& data);
bool send_frame(const std::shared_ptr<asio::ip::tcp::socket>& socket, FrameData& data);
@ -48,6 +47,8 @@ private:
std::map<std::string, std::shared_ptr<ClientCache>> client_map_;
CMutBuffer buffer_{};
short port_;
long tokens_{};
long use_tokens_{};
};
#endif