Rework chat rendering.

This commit is contained in:
Bartosz Taudul
2025-06-03 21:39:56 +02:00
parent 33e55c32f7
commit 5ccd05d2a5
5 changed files with 299 additions and 377 deletions

View File

@@ -118,6 +118,7 @@ if(NOT EMSCRIPTEN)
list(APPEND SERVER_FILES
TracyLlm.cpp
TracyLlmApi.cpp
TracyLlmChat.cpp
TracyLlmEmbeddings.cpp
TracyLlmTools.cpp
)

View File

@@ -1,3 +1,4 @@
#include <array>
#include <curl/curl.h>
#include <stdint.h>
#include <stdlib.h>
@@ -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 = "<tool_output>\n...";
constexpr size_t InputBufferSize = 1024;
TracyLlm::TracyLlm()
@@ -45,6 +46,7 @@ TracyLlm::TracyLlm()
ResetChat();
m_api = std::make_unique<TracyLlmApi>();
m_chatUi = std::make_unique<TracyLlmChat>();
m_tools = std::make_unique<TracyLlmTools>();
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<const std::string&>();
if( roleStr == "system" ) continue;
const auto& contentNode = line["content"];
if( !contentNode.is_string() ) continue;
const auto& content = contentNode.get_ref<const std::string&>();
const auto posStart = ImGui::GetCursorPos().x;
const auto& role = line["role"].get_ref<const std::string&>();
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<const std::string&>().starts_with( "<tool_output>\n" );
const auto isForgotten = isToolResponse && line["content"].is_string() && line["content"].get_ref<const std::string&>() == 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( "<tool_output>\n" ) ) role = TracyLlmChat::TurnRole::Assistant;
else if( content.starts_with( "<debug>" ) ) 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( "<debug>" ) ) 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<const std::string&>().c_str() + sizeof( "<tool_output>" ) );
ImGui::PopFont();
ImGui::TreePop();
}
ImGui::PopID();
}
else if( isAssistant )
{
const auto& content = line["content"].get_ref<const std::string&>();
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 == "<think>" )
{
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 != "</think>" )
{
PrintLine( thinkCtx, *it++, num++ );
}
CleanContext( thinkCtx );
if( it != cache.lines.end() ) ++it;
ImGui::TreePop();
}
else
{
while( it != cache.lines.end() && *it != "</think>" ) ++it;
if( it != cache.lines.end() ) ++it;
}
ImGui::PopID();
ImGui::PopStyleColor();
}
else if( line == "<tool>" )
{
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 != "</tool>" )
{
ImGui::TextWrapped( "%s", (*it).c_str() );
++it;
}
if( it != cache.lines.end() ) ++it;
ImGui::PopFont();
ImGui::TreePop();
}
else
{
while( it != cache.lines.end() && *it != "</tool>" ) ++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<const std::string&>();
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();
}
}
}

View File

@@ -12,12 +12,12 @@
#include <vector>
#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<void()> callback;
};
struct ChatCache
{
std::vector<std::string> lines;
size_t parsedLen;
};
struct LineContext
{
bool codeBlock;
};
public:
TracyLlm();
~TracyLlm();
@@ -66,13 +55,8 @@ private:
void SendMessage( std::unique_lock<std::mutex>& 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<TracyLlmApi> m_api;
std::unique_ptr<TracyLlmChat> m_chatUi;
std::unique_ptr<TracyLlmTools> m_tools;
int m_modelIdx;
@@ -96,7 +80,6 @@ private:
char* m_input;
char* m_apiInput;
std::vector<nlohmann::json> m_chat;
unordered_flat_map<size_t, ChatCache> m_chatCache;
std::shared_ptr<EmbedData> m_systemPrompt;
std::shared_ptr<EmbedData> m_systemReminder;

View File

@@ -0,0 +1,218 @@
#include <array>
#include <assert.h>
#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<NumRoles; ++i )
{
m_width[i] = ImGui::CalcTextSize( roles[i].icon ).x;
max = std::max( max, m_width[i] );
}
m_maxWidth = max;
m_role = TurnRole::None;
m_thinkActive = false;
m_thinkOpen = false;
m_thinkIdx = 0;
m_subIdx = 0;
}
void TracyLlmChat::End()
{
if( m_role != TurnRole::None )
{
NormalScope();
ImGui::EndGroup();
}
}
void TracyLlmChat::Turn( TurnRole role, const std::string& content )
{
const auto& roleData = roles[(int)role];
if( role != m_role )
{
if( m_role != TurnRole::None )
{
NormalScope();
ImGui::EndGroup();
}
m_role = role;
m_thinkActive = false;
m_thinkOpen = false;
const auto diff = m_maxWidth - m_width[(int)role];
const auto offset = diff / 2;
ImGui::Dummy( ImVec2( offset, 0 ) );
ImGui::SameLine( 0, 0 );
ImGui::TextColored( roleData.iconColor, "%s", roleData.icon );
ImGui::SameLine( 0, 0 );
ImGui::Dummy( ImVec2( diff - offset, 0 ) );
ImGui::SameLine();
ImGui::BeginGroup();
}
ImGui::PushStyleColor( ImGuiCol_Text, roleData.textColor );
if( role != TurnRole::Assistant )
{
ImGui::TextWrapped( "%s", content.c_str() );
}
else if( content.starts_with( "<tool_output>\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( "<tool_output>\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( "<think>", pos );
auto posTool = content.find( "<tool>", 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( "<think>" ) - 1;
while( content[pos] == '\n' ) pos++;
auto endThink = content.find( "</think>", pos );
ThinkScope();
if( m_thinkOpen ) PrintThink( content.substr( pos, endThink - pos ).c_str() );
if( endThink == std::string::npos ) break;
pos = endThink + sizeof( "</think>" ) - 1;
while( content[pos] == '\n' ) pos++;
}
else
{
assert( minPos == posTool );
pos += sizeof( "<tool>" ) - 1;
while( content[pos] == '\n' ) pos++;
auto endTool = content.find( "</tool>", pos );
ThinkScope();
if( m_thinkOpen ) PrintToolCall( content.substr( pos, endTool - pos ).c_str() );
if( endTool == std::string::npos ) break;
pos = endTool + sizeof( "</tool>" ) - 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();
}
}

View File

@@ -0,0 +1,52 @@
#ifndef __TRACYLLMCHAT_HPP__
#define __TRACYLLMCHAT_HPP__
#include <string>
namespace tracy
{
class TracyLlmChat
{
public:
static constexpr const char* ForgetMsg = "<tool_output>\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