From d50f05c00b0097adeb5e8a0cbf860ba9b8b8df90 Mon Sep 17 00:00:00 2001
From: Bartosz Taudul <wolf@nereid.pl>
Date: Sun, 11 May 2025 19:23:59 +0200
Subject: [PATCH] Rework loading LLM models.

---
 profiler/src/main.cpp              | 16 ++++-----
 profiler/src/profiler/TracyLlm.cpp | 56 ++++++++++++++----------------
 profiler/src/profiler/TracyLlm.hpp | 16 ++++++---
 3 files changed, 45 insertions(+), 43 deletions(-)

diff --git a/profiler/src/main.cpp b/profiler/src/main.cpp
index 8a7478aa..054c2912 100644
--- a/profiler/src/main.cpp
+++ b/profiler/src/main.cpp
@@ -864,8 +864,7 @@ static void DrawContents()
                 {
                     static int llmstatus = 0;
                     static std::string llmVersion;
-                    static std::vector<std::string> llmModels;
-                    static std::vector<size_t> llmCtxSizes;
+                    static std::vector<tracy::TracyLlm::LlmModel> llmModels;
 
                     ImGui::Indent();
                     ImGui::TextUnformatted( "Ollama URL" );
@@ -897,7 +896,6 @@ static void DrawContents()
                         {
                             llmVersion = llm.GetVersion();
                             llmModels = llm.GetModels();
-                            for( auto& v : llmModels ) llmCtxSizes.push_back( llm.GetCtxSize( v ) );
                         }
                     }
 
@@ -930,29 +928,29 @@ static void DrawContents()
                         int sel;
                         for( sel=0; sel<llmModels.size(); sel++ )
                         {
-                            if( llmModels[sel] == s_config.llmModel ) break;
+                            if( llmModels[sel].name == s_config.llmModel ) break;
                         }
                         if( sel == llmModels.size() )
                         {
                             sel = 0;
-                            s_config.llmModel = llmModels[0];
+                            s_config.llmModel = llmModels[0].name;
                             SaveConfig();
                         }
                         ImGui::SameLine();
                         ImGui::SetNextItemWidth( 225 * dpiScale );
-                        if( ImGui::BeginCombo( "##llmmodel", llmModels[sel].c_str() ) )
+                        if( ImGui::BeginCombo( "##llmmodel", llmModels[sel].name.c_str() ) )
                         {
                             for( int i=0; i<llmModels.size(); i++ )
                             {
                                 bool isSelected = ( i == sel );
-                                if( ImGui::Selectable( llmModels[i].c_str(), isSelected ) )
+                                if( ImGui::Selectable( llmModels[i].name.c_str(), isSelected ) )
                                 {
-                                    s_config.llmModel = llmModels[i];
+                                    s_config.llmModel = llmModels[i].name;
                                     SaveConfig();
                                 }
                                 if( isSelected ) ImGui::SetItemDefaultFocus();
                                 ImGui::SameLine();
-                                ImGui::TextDisabled( "(ctx: %s)", tracy::RealToString( llmCtxSizes[i] ) );
+                                ImGui::TextDisabled( "(ctx: %s)", tracy::RealToString( llmModels[i].ctxSize ) );
                             }
                             ImGui::EndCombo();
                         }
diff --git a/profiler/src/profiler/TracyLlm.cpp b/profiler/src/profiler/TracyLlm.cpp
index 704df499..583960dc 100644
--- a/profiler/src/profiler/TracyLlm.cpp
+++ b/profiler/src/profiler/TracyLlm.cpp
@@ -21,21 +21,24 @@ TracyLlm::TracyLlm()
         return;
     }
 
-    const auto models = m_ollama->list_models();
-    if( models.empty() )
+    try
+    {
+        LoadModels();
+    }
+    catch( const std::exception& e )
     {
         m_ollama.reset();
         return;
     }
 
-    m_model = s_config.llmModel;
-    if( std::ranges::find( models, m_model ) == models.end() ) m_model = models[0];
-
-    m_ctxSize = GetCtxSize( m_model );
-    if( m_ctxSize == 0 )
+    auto it = std::ranges::find_if( m_models, []( const auto& model ) { return model.name == s_config.llmModel; } );
+    if( it == m_models.end() )
     {
-        m_ollama.reset();
-        return;
+        m_modelIdx = 0;
+    }
+    else
+    {
+        m_modelIdx = std::distance( m_models.begin(), it );
     }
 }
 
@@ -48,27 +51,6 @@ std::string TracyLlm::GetVersion() const
     return m_ollama->get_version();
 }
 
-std::vector<std::string> TracyLlm::GetModels() const
-{
-    return m_ollama->list_models();
-}
-
-size_t TracyLlm::GetCtxSize( const std::string& model ) const
-{
-    try
-    {
-        const auto info = m_ollama->show_model_info( model );
-        const auto& modelInfo = info["model_info"];
-        const auto architecture = modelInfo["general.architecture"].get<std::string>();
-        const auto& ctx = modelInfo[architecture + ".context_length"];
-        return ctx.get<size_t>();
-    }
-    catch( const std::exception& e )
-    {
-        return 0;
-    }
-}
-
 void TracyLlm::Draw()
 {
     const auto scale = GetScale();
@@ -78,4 +60,18 @@ void TracyLlm::Draw()
     ImGui::End();
 }
 
+void TracyLlm::LoadModels()
+{
+    m_models.clear();
+    const auto models = m_ollama->list_models();
+    for( const auto& model : models )
+    {
+        const auto info = m_ollama->show_model_info( model );
+        const auto& modelInfo = info["model_info"];
+        const auto architecture = modelInfo["general.architecture"].get<std::string>();
+        const auto& ctx = modelInfo[architecture + ".context_length"];
+        m_models.emplace_back( LlmModel { .name = model, .ctxSize = ctx.get<size_t>() } );
+    }
+}
+
 }
diff --git a/profiler/src/profiler/TracyLlm.hpp b/profiler/src/profiler/TracyLlm.hpp
index c07043a1..a5acb8bd 100644
--- a/profiler/src/profiler/TracyLlm.hpp
+++ b/profiler/src/profiler/TracyLlm.hpp
@@ -13,22 +13,30 @@ namespace tracy
 class TracyLlm
 {
 public:
+    struct LlmModel
+    {
+        std::string name;
+        size_t ctxSize;
+    };
+
     TracyLlm();
     ~TracyLlm();
 
     [[nodiscard]] bool IsValid() const { return (bool)m_ollama; }
     [[nodiscard]] std::string GetVersion() const;
-    [[nodiscard]] std::vector<std::string> GetModels() const;
-    [[nodiscard]] size_t GetCtxSize( const std::string& model ) const;
+    [[nodiscard]] std::vector<LlmModel> GetModels() const { return m_models; }
 
     void Draw();
 
     bool m_show = false;
 
 private:
-    std::unique_ptr<Ollama> m_ollama;
+    void LoadModels();
 
-    std::string m_model;
+    std::unique_ptr<Ollama> m_ollama;
+    std::vector<LlmModel> m_models;
+
+    size_t m_modelIdx;
     size_t m_ctxSize;
 };