From 5ccd05d2a5faf5348d6e404c34de035ffc99dfba Mon Sep 17 00:00:00 2001 From: Bartosz Taudul Date: Tue, 3 Jun 2025 21:39:56 +0200 Subject: [PATCH] Rework chat rendering. --- profiler/CMakeLists.txt | 1 + profiler/src/profiler/TracyLlm.cpp | 384 ++----------------------- profiler/src/profiler/TracyLlm.hpp | 21 +- profiler/src/profiler/TracyLlmChat.cpp | 218 ++++++++++++++ profiler/src/profiler/TracyLlmChat.hpp | 52 ++++ 5 files changed, 299 insertions(+), 377 deletions(-) create mode 100644 profiler/src/profiler/TracyLlmChat.cpp create mode 100644 profiler/src/profiler/TracyLlmChat.hpp diff --git a/profiler/CMakeLists.txt b/profiler/CMakeLists.txt index 51c82c44..4435bb15 100644 --- a/profiler/CMakeLists.txt +++ b/profiler/CMakeLists.txt @@ -118,6 +118,7 @@ if(NOT EMSCRIPTEN) list(APPEND SERVER_FILES TracyLlm.cpp TracyLlmApi.cpp + TracyLlmChat.cpp TracyLlmEmbeddings.cpp TracyLlmTools.cpp ) diff --git a/profiler/src/profiler/TracyLlm.cpp b/profiler/src/profiler/TracyLlm.cpp index b51da03e..49c9795e 100644 --- a/profiler/src/profiler/TracyLlm.cpp +++ b/profiler/src/profiler/TracyLlm.cpp @@ -1,3 +1,4 @@ +#include #include #include #include @@ -7,6 +8,7 @@ #include "TracyImGui.hpp" #include "TracyLlm.hpp" #include "TracyLlmApi.hpp" +#include "TracyLlmChat.hpp" #include "TracyLlmTools.hpp" #include "TracyPrint.hpp" #include "TracyWeb.hpp" @@ -20,7 +22,6 @@ namespace tracy extern double s_time; -constexpr const char* ForgetMsg = "\n..."; constexpr size_t InputBufferSize = 1024; TracyLlm::TracyLlm() @@ -45,6 +46,7 @@ TracyLlm::TracyLlm() ResetChat(); m_api = std::make_unique(); + m_chatUi = std::make_unique(); m_tools = std::make_unique(); m_busy = true; @@ -403,219 +405,42 @@ void TracyLlm::Draw() else { ImGui::PushID( m_chatId ); - int cacheIdx = 0; - int treeIdx = 0; - int num = 0; + m_chatUi->Begin(); + for( auto& line : m_chat ) { - const auto uw = ImGui::CalcTextSize( ICON_FA_USER ).x; - const auto rw = ImGui::CalcTextSize( ICON_FA_ROBOT ).x; - const auto ew = ImGui::CalcTextSize( ICON_FA_CIRCLE_EXCLAMATION ).x; - const auto yw = ImGui::CalcTextSize( ICON_FA_REPLY ).x; - const auto cw = ImGui::CalcTextSize( ICON_FA_RECYCLE ).x; - const auto mw = std::max( { uw, rw, ew, yw, cw } ); + const auto& roleStr = line["role"].get_ref(); + if( roleStr == "system" ) continue; + const auto& contentNode = line["content"]; + if( !contentNode.is_string() ) continue; + const auto& content = contentNode.get_ref(); - const auto posStart = ImGui::GetCursorPos().x; - const auto& role = line["role"].get_ref(); + TracyLlmChat::TurnRole role = TracyLlmChat::TurnRole::None; + if( roleStr == "user" ) role = TracyLlmChat::TurnRole::User; + else if( roleStr == "error" ) role = TracyLlmChat::TurnRole::Error; + else if( roleStr == "assistant" ) role = TracyLlmChat::TurnRole::Assistant; + else assert( false ); - if( role == "system" ) continue; - - const auto isUser = role == "user"; - const auto isError = role == "error"; - const auto isAssistant = role == "assistant"; - const auto isToolResponse = isUser && line["content"].get_ref().starts_with( "\n" ); - const auto isForgotten = isToolResponse && line["content"].is_string() && line["content"].get_ref() == ForgetMsg; - - float diff, offset; - if( isForgotten ) + if( role == TracyLlmChat::TurnRole::User ) { - diff = mw - cw; - offset = diff / 2; - ImGui::Dummy( ImVec2( offset, 0 ) ); - ImGui::SameLine( 0, 0 ); - ImGui::TextColored( style.Colors[ImGuiCol_TextDisabled], ICON_FA_RECYCLE ); + if( content.starts_with( "\n" ) ) role = TracyLlmChat::TurnRole::Assistant; + else if( content.starts_with( "" ) ) role = TracyLlmChat::TurnRole::UserDebug; } - else if( isToolResponse ) + else if( role == TracyLlmChat::TurnRole::Assistant ) { - diff = mw - yw; - offset = diff / 2; - ImGui::Dummy( ImVec2( offset, 0 ) ); - ImGui::SameLine( 0, 0 ); - ImGui::TextColored( style.Colors[ImGuiCol_TextDisabled], ICON_FA_REPLY ); - } - else if( isUser ) - { - diff = mw - uw; - offset = diff / 2; - ImGui::Dummy( ImVec2( offset, 0 ) ); - ImGui::SameLine( 0, 0 ); - ImGui::TextColored( ImVec4( 0.75f, 1.f, 0.25f, 1.f ), ICON_FA_USER ); - } - else if( isError ) - { - diff = mw - ew; - offset = diff / 2; - ImGui::Dummy( ImVec2( offset, 0 ) ); - ImGui::SameLine( 0, 0 ); - ImGui::TextColored( ImVec4( 1.f, 0.25f, 0.25f, 1.f ), ICON_FA_CIRCLE_EXCLAMATION ); - } - else if( isAssistant ) - { - diff = mw - rw; - offset = diff / 2; - ImGui::Dummy( ImVec2( offset, 0 ) ); - ImGui::SameLine( 0, 0 ); - ImGui::TextColored( ImVec4( 0.4f, 0.5f, 1.f, 1.f ), ICON_FA_ROBOT ); - } - else - { - assert( false ); + if( content.starts_with( "" ) ) role = TracyLlmChat::TurnRole::AssistantDebug; } - ImGui::SameLine( 0, 0 ); - ImGui::Dummy( ImVec2( diff - offset, 0 ) ); - ImGui::SameLine(); - ImGui::BeginGroup(); - - if( isToolResponse ) - { - ImGui::PushStyleColor( ImGuiCol_Text, style.Colors[ImGuiCol_TextDisabled] ); - } - else if( isUser ) - { - ImGui::PushStyleColor( ImGuiCol_Text, ImVec4( 0.64f, 0.76f, 0.41f, 1.f ) ); - } - else if( isError ) - { - ImGui::PushStyleColor( ImGuiCol_Text, ImVec4( 1.f, 0.25f, 0.25f, 1.f ) ); - } - else if( isAssistant ) - { - ImGui::PushStyleColor( ImGuiCol_Text, style.Colors[ImGuiCol_Text] ); - } - else - { - assert( false ); - } - - if( isForgotten ) - { - ImGui::TextUnformatted( "Tool response removed to save context space" ); - treeIdx++; - } - else if( isToolResponse ) - { - ImGui::PushID( treeIdx++ ); - auto expand = ImGui::TreeNode( "Tool response..." ); - if( line.contains( "images" ) ) - { - ImGui::SameLine(); - ImGui::TextUnformatted( ICON_FA_FILE_IMAGE ); - } - if( expand ) - { - ImGui::PushFont( g_fonts.mono ); - ImGui::TextWrapped( "%s", line["content"].get_ref().c_str() + sizeof( "" ) ); - ImGui::PopFont(); - ImGui::TreePop(); - } - ImGui::PopID(); - } - else if( isAssistant ) - { - const auto& content = line["content"].get_ref(); - - auto cit = m_chatCache.find( cacheIdx ); - if( cit == m_chatCache.end() ) cit = m_chatCache.emplace( cacheIdx, ChatCache {} ).first; - auto& cache = cit->second; - - if( cache.parsedLen != content.size() ) - { - UpdateCache( cache, content ); - assert( cache.parsedLen == content.size() ); - } - - if( cache.lines.empty() && m_responding ) - { - tracy::TextDisabledUnformatted( "\xe2\x80\xa6" ); - } - else - { - LineContext ctx = {}; - auto it = cache.lines.begin(); - while( it != cache.lines.end() ) - { - auto& line = *it++; - if( line == "" ) - { - ImGui::PushStyleColor( ImGuiCol_Text, ImVec4( 0.5f, 0.5f, 0.3f, 1.f ) ); - ImGui::PushID( treeIdx++ ); - if( ImGui::TreeNode( ICON_FA_LIGHTBULB " Internal thoughts..." ) ) - { - LineContext thinkCtx = {}; - while( it != cache.lines.end() && *it != "" ) - { - PrintLine( thinkCtx, *it++, num++ ); - } - CleanContext( thinkCtx ); - 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 == "" ) - { - ImGui::PushStyleColor( ImGuiCol_Text, style.Colors[ImGuiCol_TextDisabled] ); - ImGui::PushID( treeIdx++ ); - if( ImGui::TreeNode( "Tool query..." ) ) - { - ImGui::PushFont( g_fonts.mono ); - while( it != cache.lines.end() && *it != "" ) - { - ImGui::TextWrapped( "%s", (*it).c_str() ); - ++it; - } - if( it != cache.lines.end() ) ++it; - ImGui::PopFont(); - ImGui::TreePop(); - } - else - { - while( it != cache.lines.end() && *it != "" ) ++it; - if( it != cache.lines.end() ) ++it; - } - ImGui::PopID(); - ImGui::PopStyleColor(); - } - else - { - PrintLine( ctx, line, num++ ); - } - } - CleanContext( ctx ); - } - } - else if( line["content"].is_string() ) - { - auto& string = line["content"].get_ref(); - PrintMarkdown( string.c_str() ); - } - ImGui::PopStyleColor(); - ImGui::EndGroup(); - cacheIdx++; + m_chatUi->Turn( role, content ); } + m_chatUi->End(); + ImGui::PopID(); + if( ImGui::GetScrollY() >= ImGui::GetScrollMaxY() ) { ImGui::SetScrollHereY( 1.f ); } - ImGui::PopID(); } ImGui::EndChild(); ImGui::Spacing(); @@ -767,7 +592,6 @@ void TracyLlm::ResetChat() m_usedCtx = 0; m_chatId++; m_chat.clear(); - m_chatCache.clear(); AddMessage( std::move( systemPrompt ), "system" ); } @@ -814,8 +638,8 @@ void TracyLlm::ManageContext() for( auto& v : toolOutputs ) { m_usedCtx -= v.first / 4; - m_chat[v.second]["content"] = ForgetMsg; - m_usedCtx += strlen( ForgetMsg ) / 4; + m_chat[v.second]["content"] = TracyLlmChat::ForgetMsg; + m_usedCtx += strlen( TracyLlmChat::ForgetMsg ) / 4; if( m_usedCtx < quota ) break; } } @@ -997,160 +821,4 @@ bool TracyLlm::OnResponse( const nlohmann::json& json ) 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; -} - -static bool IsHeading( const char* str ) -{ - if( *str != '#' ) return false; - while( *str == '#' ) str++; - return *str == ' ' || *str == '\t'; -} - -void TracyLlm::PrintLine( LineContext& ctx, const std::string& str, int num ) -{ - if( str.empty() ) return; - - auto ptr = str.c_str(); - while( *ptr == ' ' || *ptr == '\t' ) ptr++; - if( strncmp( ptr, "```", 3 ) == 0 ) - { - if( ctx.codeBlock ) - { - ImGui::PopFont(); - ImGui::EndChild(); - ctx.codeBlock = false; - } - else - { - char tmp[64]; - snprintf( tmp, sizeof( tmp ), "##chat_code_%d", num ); - ImGui::BeginChild( tmp, ImVec2( 0, 0 ), ImGuiChildFlags_FrameStyle | ImGuiChildFlags_Borders | ImGuiChildFlags_AutoResizeY ); - if( ptr[3] ) - { - ImGui::PushFont( g_fonts.small ); - ImGui::SetCursorPosX( ImGui::GetContentRegionAvail().x - ImGui::CalcTextSize( ptr + 3 ).x ); - ImGui::TextUnformatted( ptr + 3 ); - ImGui::PopFont(); - } - ImGui::PushFont( g_fonts.mono ); - ctx.codeBlock = true; - } - } - else - { - ImGui::PushTextWrapPos( 0 ); - if( ctx.codeBlock ) - { - ImGui::TextUnformatted( str.c_str() ); - } - else if( str == "---" ) - { - ImGui::Spacing(); - ImGui::Separator(); - ImGui::Spacing(); - } - else if( IsHeading( str.c_str() ) ) - { - ImGui::PushFont( g_fonts.big ); - ImGui::TextUnformatted( str.c_str() ); - ImGui::PopFont(); - } - else - { - const auto begin = str.c_str(); - ptr = begin; - while( *ptr == ' ' || *ptr == '\t' ) ptr++; - if( ( ptr[0] == '*' || ptr[0] == '-' ) && ptr[1] == ' ' ) - { - ImGui::TextUnformatted( std::string( ptr - begin, ' ' ).c_str() ); - ImGui::SameLine(); - ImGui::Bullet(); - ImGui::SameLine(); - ptr++; - } - PrintMarkdown( ptr ); - } - ImGui::PopTextWrapPos(); - } -} - -void TracyLlm::PrintMarkdown( const char* str ) -{ - auto& style = ImGui::GetStyle(); - ImGui::PushStyleVar( ImGuiStyleVar_ItemSpacing, ImVec2( style.ItemSpacing.x, 0.0f ) ); - - auto end = str + strlen( str ); - bool first = true; - bool isCode = false; - - while( str != end ) - { - if( first ) - { - first = false; - } - else - { - ImGui::SameLine( 0, 0 ); - } - - auto next = str; - while( next != end && *next != '`' ) next++; - if( *next == '`' ) - { - PrintTextWrapped( str, next ); - str = next + 1; - - isCode = !isCode; - if( isCode ) - { - ImGui::PushFont( g_fonts.mono ); - } - else - { - ImGui::PopFont(); - } - } - else - { - PrintTextWrapped( str, next ); - str = next; - } - } - - if( isCode ) ImGui::PopFont(); - - ImGui::PopStyleVar(); - ImGui::SetCursorPosY( ImGui::GetCursorPosY() + style.ItemSpacing.y ); -} - -void TracyLlm::CleanContext( LineContext& ctx) -{ - if( ctx.codeBlock ) - { - ImGui::PopFont(); - ImGui::EndChild(); - } -} - } diff --git a/profiler/src/profiler/TracyLlm.hpp b/profiler/src/profiler/TracyLlm.hpp index 436296e9..7c08cad3 100644 --- a/profiler/src/profiler/TracyLlm.hpp +++ b/profiler/src/profiler/TracyLlm.hpp @@ -12,12 +12,12 @@ #include #include "TracyEmbed.hpp" -#include "tracy_robin_hood.h" namespace tracy { class TracyLlmApi; +class TracyLlmChat; class TracyLlmTools; class TracyLlm @@ -34,17 +34,6 @@ class TracyLlm std::function callback; }; - struct ChatCache - { - std::vector lines; - size_t parsedLen; - }; - - struct LineContext - { - bool codeBlock; - }; - public: TracyLlm(); ~TracyLlm(); @@ -66,13 +55,8 @@ private: void SendMessage( std::unique_lock& lock ); bool OnResponse( const nlohmann::json& json ); - void UpdateCache( ChatCache& cache, const std::string& str ); - - void PrintLine( LineContext& ctx, const std::string& str, int num ); - void PrintMarkdown( const char* str ); - void CleanContext( LineContext& ctx); - std::unique_ptr m_api; + std::unique_ptr m_chatUi; std::unique_ptr m_tools; int m_modelIdx; @@ -96,7 +80,6 @@ private: char* m_input; char* m_apiInput; std::vector m_chat; - unordered_flat_map m_chatCache; std::shared_ptr m_systemPrompt; std::shared_ptr m_systemReminder; diff --git a/profiler/src/profiler/TracyLlmChat.cpp b/profiler/src/profiler/TracyLlmChat.cpp new file mode 100644 index 00000000..7e8cb72f --- /dev/null +++ b/profiler/src/profiler/TracyLlmChat.cpp @@ -0,0 +1,218 @@ +#include +#include + +#include "TracyImGui.hpp" +#include "TracyLlmChat.hpp" +#include "../Fonts.hpp" + +namespace tracy +{ + +constexpr auto ThinkColor = ImVec4( 0.5f, 0.5f, 0.3f, 1.f ); + +struct RoleData +{ + const char* icon; + ImVec4 iconColor; + ImVec4 textColor; +}; + +constexpr std::array roles = { + RoleData { ICON_FA_USER, ImVec4( 0.75f, 1.f, 0.25f, 1.f ), ImVec4( 0.64f, 0.76f, 0.41f, 1.f ) }, + RoleData { ICON_FA_TERMINAL, ImVec4( 1.f, 0.5f, 0.5f, 1.f ), ImVec4( 1.f, 0.65f, 0.65f, 1.f ) }, + RoleData { ICON_FA_ROBOT, ImVec4( 0.4f, 0.5f, 1.f, 1.f ), ImVec4( 1.f, 1.f, 1.f, 1.f ) }, + RoleData { ICON_FA_CODE, ImVec4( 1.0f, 0.5f, 1.f, 1.f ), ImVec4( 1.f, 0.65f, 1.f, 1.f ) }, + RoleData { ICON_FA_CIRCLE_EXCLAMATION, ImVec4( 1.f, 0.25f, 0.25f, 1.f ), ImVec4( 1.f, 0.25f, 0.25f, 1.f ) } +}; +constexpr size_t NumRoles = roles.size(); + +static_assert( NumRoles == (int)TracyLlmChat::TurnRole::None ); + + +TracyLlmChat::TracyLlmChat() + : m_width( new float[NumRoles] ) +{ +} + +TracyLlmChat::~TracyLlmChat() +{ + delete[] m_width; +} + +void TracyLlmChat::Begin() +{ + float max = 0; + for( size_t i=0; i\n" ) ) + { + ThinkScope(); + if( m_thinkOpen ) + { + ImGui::PushStyleColor( ImGuiCol_Text, ImVec4( 0.5f, 0.5f, 0.5f, 1.f ) ); + if( content == ForgetMsg ) + { + ImGui::TextUnformatted( ICON_FA_RECYCLE " Tool response removed to save context space" ); + m_subIdx++; + } + else + { + ImGui::PushID( m_subIdx++ ); + if( ImGui::TreeNode( ICON_FA_REPLY " Tool response..." ) ) + { + ImGui::PushFont( g_fonts.mono ); + ImGui::TextWrapped( "%s", content.c_str() + sizeof( "\n" ) - 1 ); + ImGui::PopFont(); + ImGui::TreePop(); + } + ImGui::PopID(); + } + ImGui::PopStyleColor(); + } + else + { + m_subIdx++; + } + } + else + { + size_t pos = 0; + size_t end = content.size(); + while( pos < end ) + { + auto posThink = content.find( "", pos ); + auto posTool = content.find( "", pos ); + auto minPos = std::min( posThink, posTool ); + + if( pos != minPos ) + { + NormalScope(); + PrintMarkdown( content.substr( pos, minPos - pos ).c_str() ); + } + + pos = minPos; + if( pos == std::string::npos ) break; + + if( minPos == posThink ) + { + pos += sizeof( "" ) - 1; + while( content[pos] == '\n' ) pos++; + auto endThink = content.find( "", pos ); + ThinkScope(); + if( m_thinkOpen ) PrintThink( content.substr( pos, endThink - pos ).c_str() ); + if( endThink == std::string::npos ) break; + pos = endThink + sizeof( "" ) - 1; + while( content[pos] == '\n' ) pos++; + } + else + { + assert( minPos == posTool ); + pos += sizeof( "" ) - 1; + while( content[pos] == '\n' ) pos++; + auto endTool = content.find( "", pos ); + ThinkScope(); + if( m_thinkOpen ) PrintToolCall( content.substr( pos, endTool - pos ).c_str() ); + if( endTool == std::string::npos ) break; + pos = endTool + sizeof( "" ) - 1; + while( content[pos] == '\n' ) pos++; + } + } + } + ImGui::PopStyleColor(); +} + +void TracyLlmChat::NormalScope() +{ + if( !m_thinkActive ) return; + if( m_thinkOpen ) + { + ImGui::TreePop(); + m_thinkOpen = false; + } + ImGui::PopStyleColor(); + ImGui::PopID(); + m_thinkActive = false; +} + +void TracyLlmChat::ThinkScope() +{ + if( m_thinkActive ) return; + m_thinkActive = true; + ImGui::PushID( m_thinkIdx++ ); + ImGui::PushStyleColor( ImGuiCol_Text, ThinkColor ); + m_thinkOpen = ImGui::TreeNode( ICON_FA_LIGHTBULB " Internal thoughts..." ); +} + +void TracyLlmChat::PrintMarkdown( const char* str ) +{ + ImGui::TextWrapped( "%s", str ); +} + +void TracyLlmChat::PrintThink( const char* str ) +{ + ImGui::PushStyleColor( ImGuiCol_Text, ThinkColor ); + PrintMarkdown( str ); + ImGui::PopStyleColor(); +} + +void TracyLlmChat::PrintToolCall( const char* str ) +{ + ImGui::PushStyleColor( ImGuiCol_Text, ImVec4( 0.5f, 0.5f, 0.5f, 1.f ) ); + ImGui::PushFont( g_fonts.mono ); + ImGui::TextWrapped( "%s", str ); + ImGui::PopFont(); + ImGui::PopStyleColor(); +} + +} diff --git a/profiler/src/profiler/TracyLlmChat.hpp b/profiler/src/profiler/TracyLlmChat.hpp new file mode 100644 index 00000000..2be92b24 --- /dev/null +++ b/profiler/src/profiler/TracyLlmChat.hpp @@ -0,0 +1,52 @@ +#ifndef __TRACYLLMCHAT_HPP__ +#define __TRACYLLMCHAT_HPP__ + +#include + +namespace tracy +{ + +class TracyLlmChat +{ +public: + static constexpr const char* ForgetMsg = "\n..."; + + enum class TurnRole + { + User, + UserDebug, + Assistant, + AssistantDebug, + Error, + None + }; + + TracyLlmChat(); + ~TracyLlmChat(); + + void Begin(); + void End(); + + void Turn( TurnRole role, const std::string& content ); + +private: + void NormalScope(); + void ThinkScope(); + + void PrintMarkdown( const char* str ); + void PrintThink( const char* str ); + void PrintToolCall( const char* str ); + + float* m_width; + float m_maxWidth; + + TurnRole m_role; + bool m_thinkActive; + bool m_thinkOpen; + int m_thinkIdx; + int m_subIdx; +}; + +} + +#endif