From b33110e53e7704bc7736bcf6a029d221f51fae4d Mon Sep 17 00:00:00 2001
From: Bartosz Taudul <wolf@nereid.pl>
Date: Thu, 22 May 2025 21:17:38 +0200
Subject: [PATCH] Retrieve context size for LM Studio models.

---
 profiler/src/profiler/TracyLlm.cpp    |  2 +-
 profiler/src/profiler/TracyLlmApi.cpp | 19 +++++++++++++++++--
 profiler/src/profiler/TracyLlmApi.hpp |  3 ++-
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/profiler/src/profiler/TracyLlm.cpp b/profiler/src/profiler/TracyLlm.cpp
index 713794da..21e5c95c 100644
--- a/profiler/src/profiler/TracyLlm.cpp
+++ b/profiler/src/profiler/TracyLlm.cpp
@@ -602,7 +602,7 @@ void TracyLlm::SendMessage( const std::vector<nlohmann::json>& messages )
         req["stream"] = true;
         if( m_setTemperature ) req["temperature"] = m_temperature;
 
-        res = m_api->ChatCompletion( req, [this]( const nlohmann::json& response ) -> bool { return OnResponse( response ); } );
+        res = m_api->ChatCompletion( req, [this]( const nlohmann::json& response ) -> bool { return OnResponse( response ); }, m_modelIdx );
     }
     catch( std::exception& e )
     {
diff --git a/profiler/src/profiler/TracyLlmApi.cpp b/profiler/src/profiler/TracyLlmApi.cpp
index b569bc28..a197b341 100644
--- a/profiler/src/profiler/TracyLlmApi.cpp
+++ b/profiler/src/profiler/TracyLlmApi.cpp
@@ -72,6 +72,7 @@ bool TracyLlmApi::Connect( const char* url )
                     continue;
                 }
                 m_models.back().quant = json2["quantization"].get_ref<const std::string&>();
+                if( json2.contains( "loaded_context_length" ) ) m_models.back().contextSize = json2["loaded_context_length"].get<int>();
             }
             else if( PostRequest( m_url + "/api/show", "{\"name\":\"" + id + "\"}", buf2 ) == 200 )
             {
@@ -122,7 +123,7 @@ static size_t StreamFn( void* _data, size_t size, size_t num, void* ptr )
     return sz;
 }
 
-bool TracyLlmApi::ChatCompletion( const nlohmann::json& req, const std::function<bool(const nlohmann::json&)>& callback )
+bool TracyLlmApi::ChatCompletion( const nlohmann::json& req, const std::function<bool(const nlohmann::json&)>& callback, int modelIdx )
 {
     assert( m_curl );
     StreamData data = { .callback = callback };
@@ -147,7 +148,21 @@ bool TracyLlmApi::ChatCompletion( const nlohmann::json& req, const std::function
 
     int64_t http_code = 0;
     curl_easy_getinfo( m_curl, CURLINFO_RESPONSE_CODE, &http_code );
-    return http_code == 200;
+    if( http_code == 200 )
+    {
+        if( m_type == Type::LmStudio && m_models[modelIdx].contextSize <= 0 )
+        {
+            curl_easy_reset( m_curl );
+            SetupCurl();
+            std::string buf;
+            if( GetRequest( m_url + "/api/v0/models/" + m_models[modelIdx].name, buf ) == 200 )
+            {
+                m_models[modelIdx].contextSize = nlohmann::json::parse( buf )["loaded_context_length"].get<int>();
+            }
+        }
+        return true;
+    }
+    return false;
 }
 
 int64_t TracyLlmApi::GetRequest( const std::string& url, std::string& response )
diff --git a/profiler/src/profiler/TracyLlmApi.hpp b/profiler/src/profiler/TracyLlmApi.hpp
index d5845c24..67cb2b72 100644
--- a/profiler/src/profiler/TracyLlmApi.hpp
+++ b/profiler/src/profiler/TracyLlmApi.hpp
@@ -14,6 +14,7 @@ struct LlmModel
 {
     std::string name;
     std::string quant;
+    int contextSize = -1;
 };
 
 class TracyLlmApi
@@ -29,7 +30,7 @@ public:
     ~TracyLlmApi();
 
     bool Connect( const char* url );
-    bool ChatCompletion( const nlohmann::json& req, const std::function<bool(const nlohmann::json&)>& callback );;
+    bool ChatCompletion( const nlohmann::json& req, const std::function<bool(const nlohmann::json&)>& callback, int modelIdx );
 
     [[nodiscard]] bool IsConnected() const { return m_curl != nullptr; }
     [[nodiscard]] const std::vector<LlmModel>& GetModels() const { return m_models; }