diff --git a/profiler/src/profiler/TracyLlmApi.cpp b/profiler/src/profiler/TracyLlmApi.cpp index 7dfece3f..a1832faa 100644 --- a/profiler/src/profiler/TracyLlmApi.cpp +++ b/profiler/src/profiler/TracyLlmApi.cpp @@ -239,6 +239,29 @@ bool TracyLlmApi::Embeddings( const nlohmann::json& req, nlohmann::json& respons return true; } +int TracyLlmApi::Tokenize( const std::string& text, int modelIdx ) +{ + if( m_type == Type::LlamaSwap ) + { + std::string buf; + nlohmann::json req = { { "content", text } }; + auto res = PostRequest( m_url + "/upstream/" + m_models[modelIdx].name + "/tokenize", req.dump( -1, ' ', false, nlohmann::json::error_handler_t::replace ), buf, true ); + if( res != 200 ) return -1; + + try + { + auto json = nlohmann::json::parse( buf ); + return json["tokens"].size(); + } + catch( const std::exception& ) + { + return -1; + } + } + + return -1; +} + int64_t TracyLlmApi::GetRequest( const std::string& url, std::string& response ) { assert( m_curl ); diff --git a/profiler/src/profiler/TracyLlmApi.hpp b/profiler/src/profiler/TracyLlmApi.hpp index 13eeedbd..a1bb9d76 100644 --- a/profiler/src/profiler/TracyLlmApi.hpp +++ b/profiler/src/profiler/TracyLlmApi.hpp @@ -35,6 +35,7 @@ public: bool Connect( const char* url ); bool ChatCompletion( const nlohmann::json& req, const std::function& callback, int modelIdx ); bool Embeddings( const nlohmann::json& req, nlohmann::json& response, bool separateConnection = false ); + [[nodiscard]] int Tokenize( const std::string& text, int modelIdx ); [[nodiscard]] bool IsConnected() const { return m_curl != nullptr; } [[nodiscard]] const std::vector& GetModels() const { return m_models; }