diff --git a/profiler/src/profiler/TracyLlm.cpp b/profiler/src/profiler/TracyLlm.cpp index 19aa8bbb..33b2c039 100644 --- a/profiler/src/profiler/TracyLlm.cpp +++ b/profiler/src/profiler/TracyLlm.cpp @@ -109,6 +109,7 @@ void TracyLlm::Draw() { if( m_responding ) m_stop = true; m_chat = std::make_unique(); + m_chatCache.clear(); *m_input = 0; } ImGui::SameLine(); @@ -221,8 +222,7 @@ void TracyLlm::Draw() ImGui::SameLine( 0, 0 ); ImGui::Dummy( ImVec2( diff - offset, 0 ) ); ImGui::SameLine(); - - const auto indent = ImGui::GetCursorPos().x - posStart; + ImGui::BeginGroup(); auto& style = ImGui::GetStyle(); if( isUser ) @@ -244,36 +244,49 @@ void TracyLlm::Draw() if( isAssistant ) { - auto str = line["content"].get(); - if( strncmp( str.c_str(), "", 7 ) == 0 ) + const auto& content = line["content"]; + const auto contentSize = content.get_ref().size(); + + auto cit = m_chatCache.find( idx ); + if( cit == m_chatCache.end() ) cit = m_chatCache.emplace( idx, ChatCache {} ).first; + auto& cache = cit->second; + + if( cache.parsedLen != contentSize ) { - int strip = 7; - while( str[strip] == '\n' ) strip++; - str = str.substr( strip ); - auto pos = str.find( "\n" ); - ImGui::PushStyleColor( ImGuiCol_Text, ImVec4( 0.5f, 0.5f, 0.3f, 1.f ) ); - ImGui::PushID( idx ); - if( ImGui::TreeNode( ICON_FA_LIGHTBULB " Internal thoughts..." ) ) - { - ImGui::Indent( indent ); - ImGui::TextWrapped( "%s", str.substr( 0, pos ).c_str() ); - ImGui::Unindent( indent ); - ImGui::TreePop(); - } - ImGui::PopID(); - ImGui::PopStyleColor(); - if( pos != std::string::npos ) - { - strip = pos + 9; - while( str[strip] == '\n' ) strip++; - ImGui::Indent( indent ); - ImGui::TextWrapped( "%s", str.substr( strip ).c_str() ); - ImGui::Unindent( indent ); - } + UpdateCache( cache, content.get() ); + assert( cache.parsedLen == contentSize ); } - else + + auto it = cache.lines.begin(); + while( it != cache.lines.end() ) { - ImGui::TextWrapped( "%s", line["content"].get().c_str() ); + auto& line = *it++; + if( line == "" ) + { + ImGui::PushStyleColor( ImGuiCol_Text, ImVec4( 0.5f, 0.5f, 0.3f, 1.f ) ); + ImGui::PushID( idx ); + if( ImGui::TreeNode( ICON_FA_LIGHTBULB " Internal thoughts..." ) ) + { + while( it != cache.lines.end() && *it != "" ) + { + if( !it->empty() ) ImGui::TextWrapped( "%s", it->c_str() ); + it++; + } + if( it != cache.lines.end() ) ++it; + ImGui::TreePop(); + } + else + { + while( it != cache.lines.end() && *it != "" ) ++it; + if( it != cache.lines.end() ) ++it; + } + ImGui::PopID(); + ImGui::PopStyleColor(); + } + else + { + if( !line.empty() ) ImGui::TextWrapped( "%s", line.c_str() ); + } } } else @@ -281,6 +294,7 @@ void TracyLlm::Draw() ImGui::TextWrapped( "%s", line["content"].get().c_str() ); } ImGui::PopStyleColor(); + ImGui::EndGroup(); idx++; } @@ -468,4 +482,26 @@ bool TracyLlm::OnResponse( const ollama::response& response ) return true; } +void TracyLlm::UpdateCache( ChatCache& cache, const std::string& str ) +{ + const auto sz = str.size(); + auto pos = cache.parsedLen; + while( pos < sz ) + { + const auto isNewLine = pos == 0 || str[pos - 1] == '\n'; + auto next = str.find( '\n', pos ); + if( next == std::string::npos ) next = sz; + if( isNewLine ) + { + cache.lines.emplace_back( str.substr( pos, next - pos ) ); + } + else + { + auto& line = cache.lines.back().append( str, pos, next - pos ); + } + pos = next + 1; + } + cache.parsedLen = sz; +} + } diff --git a/profiler/src/profiler/TracyLlm.hpp b/profiler/src/profiler/TracyLlm.hpp index 6ea12fe1..0519666b 100644 --- a/profiler/src/profiler/TracyLlm.hpp +++ b/profiler/src/profiler/TracyLlm.hpp @@ -12,6 +12,8 @@ #include +#include "tracy_robin_hood.h" + class Ollama; namespace ollama @@ -39,6 +41,12 @@ class TracyLlm std::unique_ptr chat; }; + struct ChatCache + { + std::vector lines; + size_t parsedLen; + }; + public: struct LlmModel { @@ -70,6 +78,8 @@ private: void SendMessage( ollama::messages&& messages ); bool OnResponse( const ollama::response& response ); + void UpdateCache( ChatCache& cache, const std::string& str ); + std::unique_ptr m_ollama; mutable std::mutex m_modelsLock; @@ -91,6 +101,7 @@ private: char* m_input; std::unique_ptr m_chat; + unordered_flat_map m_chatCache; ImFont* m_font; ImFont* m_smallFont;