Calculation of manual embeddings.

profiler/src/profiler/TracyLlmTools.cpp

@@ -1,3 +1,4 @@
+#include <algorithm>
 #include <curl/curl.h>
 #include <nlohmann/json.hpp>
 #include <libbase64.h>
@@ -6,8 +7,12 @@
 #include <tidybuffio.h>
 #include <time.h>
 
+#include "TracyEmbed.hpp"
+#include "TracyLlmApi.hpp"
 #include "TracyLlmTools.hpp"
 
+#include "data/Manual.hpp"
+
 constexpr const char* NoNetworkAccess = "Internet access is disabled by the user. You may inform the user that he can enable it in the settings, so that you can use the tools to gather information.";
 
 #define NetworkCheckString if( !m_netAccess ) return NoNetworkAccess
@@ -42,6 +47,11 @@ static std::string UrlEncode( const std::string& str )
     return out;
 }
 
+TracyLlmTools::~TracyLlmTools()
+{
+    CancelManualEmbeddings();
+}
+
 TracyLlmTools::ToolReply TracyLlmTools::HandleToolCalls( const std::string& name, const std::vector<std::string>& args, int contextSize )
 {
     m_ctxSize = contextSize;
@@ -88,6 +98,81 @@ std::string TracyLlmTools::GetCurrentTime()
     return buffer;
 }
 
+TracyLlmTools::EmbeddingState TracyLlmTools::GetManualEmbeddingsState() const
+{
+    std::lock_guard lock( m_lock );
+    return m_manualEmbeddingState;
+}
+
+void TracyLlmTools::BuildManualEmbeddings( const std::string& model, TracyLlmApi& api )
+{
+    std::unique_lock lock( m_lock );
+    assert( !m_manualEmbeddingState.inProgress );
+    if( m_manualEmbeddingState.done && m_manualEmbeddingState.model == model ) return;
+
+    lock.unlock();
+    if( m_thread.joinable() ) m_thread.join();
+
+    assert( !m_cancel );
+    m_manualEmbeddingState = { .model = model, .inProgress = true };
+    m_thread = std::thread( [this, &api] { ManualEmbeddingsWorker( api ); } );
+}
+
+void TracyLlmTools::ManualEmbeddingsWorker( TracyLlmApi& api )
+{
+    constexpr auto Chunk = 1000;
+    constexpr auto Overlap = 200;
+
+    auto manual = Unembed( Manual );
+
+    const auto sz = (int)manual->size();
+    const auto chunks = ( sz + Chunk - 1 ) / Chunk;
+    for( int i=0; i<chunks; i++ )
+    {
+        std::unique_lock lock( m_lock );
+        if( m_cancel )
+        {
+            m_manualEmbeddingState.inProgress = false;
+            m_manualEmbeddingState.done = false;
+            return;
+        }
+        m_manualEmbeddingState.progress = (float)i / chunks;
+        lock.unlock();
+
+        const auto start = std::max( 0, Chunk * i - Overlap );
+        const auto end = std::min( sz, Chunk * ( i + 1 ) + Overlap );
+
+        nlohmann::json req;
+        req["input"] = std::string( manual->data() + start, end - start );
+        req["model"] = m_manualEmbeddingState.model;
+
+        nlohmann::json response;
+        api.Embeddings( req, response );
+
+        std::vector<float> embeddings;
+        for( auto& item : response["data"][0]["embedding"] )
+        {
+            embeddings.emplace_back( item.get<float>() );
+        }
+    }
+
+    std::lock_guard lock( m_lock );
+    m_manualEmbeddingState.inProgress = false;
+    m_manualEmbeddingState.done = true;
+}
+
+void TracyLlmTools::CancelManualEmbeddings()
+{
+    if( m_thread.joinable() )
+    {
+        m_lock.lock();
+        m_cancel = true;
+        m_lock.unlock();
+        m_thread.join();
+        m_cancel = false;
+    }
+}
+
 int TracyLlmTools::CalcMaxSize() const
 {
     if( m_ctxSize <= 0 ) return 32*1024;

profiler/src/profiler/TracyLlmTools.hpp

@@ -1,7 +1,9 @@
 #ifndef __TRACYLLMTOOLS_HPP__
 #define __TRACYLLMTOOLS_HPP__
 
+#include <mutex>
 #include <string>
+#include <thread>
 #include <vector>
 
 #include "tracy_robin_hood.h"
@@ -9,6 +11,8 @@
 namespace tracy
 {
 
+class TracyLlmApi;
+
 class TracyLlmTools
 {
 public:
@@ -18,9 +22,23 @@ public:
         std::string image;
     };
 
+    struct EmbeddingState
+    {
+        std::string model;
+        bool done = false;
+        bool inProgress = false;
+        float progress = 0;
+    };
+
+    ~TracyLlmTools();
+
     ToolReply HandleToolCalls( const std::string& name, const std::vector<std::string>& args, int contextSize );
     std::string GetCurrentTime();
 
+    [[nodiscard]] EmbeddingState GetManualEmbeddingsState() const;
+    void BuildManualEmbeddings( const std::string& model, TracyLlmApi& api );
+    void CancelManualEmbeddings();
+
     bool m_netAccess = true;
 
 private:
@@ -33,9 +51,16 @@ private:
     std::string GetDictionary( std::string word, const std::string& lang );
     std::string SearchWeb( std::string query );
 
+    void ManualEmbeddingsWorker( TracyLlmApi& api );
+
     unordered_flat_map<std::string, std::string> m_webCache;
 
     int m_ctxSize;
+
+    mutable std::mutex m_lock;
+    std::thread m_thread;
+    bool m_cancel = false;
+    EmbeddingState m_manualEmbeddingState;
 };
 
 }