Split chat responses into lines for easier parsing.

This commit is contained in:
Bartosz Taudul
2025-05-12 23:30:27 +02:00
parent cfd3260ea0
commit 60230c2d7e
2 changed files with 76 additions and 29 deletions

View File

@@ -109,6 +109,7 @@ void TracyLlm::Draw()
{
if( m_responding ) m_stop = true;
m_chat = std::make_unique<ollama::messages>();
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<std::string>();
if( strncmp( str.c_str(), "<think>", 7 ) == 0 )
const auto& content = line["content"];
const auto contentSize = content.get_ref<const nlohmann::json::string_t&>().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( "</think>\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<std::string>() );
assert( cache.parsedLen == contentSize );
}
else
auto it = cache.lines.begin();
while( it != cache.lines.end() )
{
ImGui::TextWrapped( "%s", line["content"].get<std::string>().c_str() );
auto& line = *it++;
if( line == "<think>" )
{
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 != "</think>" )
{
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 != "</think>" ) ++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<std::string>().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;
}
}

View File

@@ -12,6 +12,8 @@
#include <imgui.h>
#include "tracy_robin_hood.h"
class Ollama;
namespace ollama
@@ -39,6 +41,12 @@ class TracyLlm
std::unique_ptr<ollama::messages> chat;
};
struct ChatCache
{
std::vector<std::string> 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<Ollama> m_ollama;
mutable std::mutex m_modelsLock;
@@ -91,6 +101,7 @@ private:
char* m_input;
std::unique_ptr<ollama::messages> m_chat;
unordered_flat_map<size_t, ChatCache> m_chatCache;
ImFont* m_font;
ImFont* m_smallFont;