feat: add tests for the session, ref #17

Author: pminev

ac-local-plugin/code/LocalLlama.cpp

@@ -28,7 +28,7 @@ namespace ac::local {
 
 namespace {
 class ChatSession {
-    llama::Session m_session;
+    llama::Session& m_session;
     const llama::Vocab& m_vocab;
     std::string m_userPrefix;
     std::string m_assistantPrefix;
@@ -149,7 +149,7 @@ class LlamaInstance final : public Instance {
         auto& prompt = params.prompt.value();
         const auto maxTokens = params.maxTokens.value();
 
-        auto s = m_instance.startSession({});
+        auto& s = m_instance.startSession({});
 
         auto promptTokens = m_instance.model().vocab().tokenize(prompt, true, true);
         s.setInitialPrompt(promptTokens);

code/ac/llama/Instance.cpp

@@ -117,10 +117,11 @@ void Instance::warmup() {
 }
 
 Session& Instance::startSession(const Session::InitParams params) {
-    if (!m_session) {
-        m_session.reset(new Session(*this, m_lctx.get(), params));
+    if (m_session) {
+        throw_ex{} << "Session is already started. Stop it to start a new one.";
     }
 
+    m_session.reset(new Session(*this, m_lctx.get(), params));
     return *m_session;
 }
 

code/ac/llama/Session.cpp

@@ -79,13 +79,16 @@ void Session::setInitialPrompt(std::span<const Token> initialPrompt) {
     }
 
     doDecode(initialPrompt, Source::InitialPrompt);
+    m_state.m_phase = State::Phase::Generating;
 }
 
 void Session::pushPrompt(std::span<const Token> prompt) {
     if (m_state.m_phase != State::Phase::Generating) {
         throw_ex{} << "Session hasn't started yet";
     }
 
+    flushPendingState();
+
     if (!prompt.empty()) {
         auto& sampler = m_instance.sampler();
         auto& model = m_instance.model();
@@ -108,10 +111,7 @@ Token Session::getToken() {
         throw_ex{} << "Session hasn't started yet";
     }
 
-    if (m_state.m_currToken != Token_Invalid) {
-        // first yield, then decode, thus we don't decode if the session is aborted
-        doDecode({&m_state.m_currToken, 1}, Source::Generated);
-    }
+    flushPendingState();
 
     auto& sampler = m_instance.sampler();
     auto& vocab = m_instance.model().vocab();
@@ -131,6 +131,8 @@ std::vector<uint8_t> Session::getState() {
         throw_ex{} << "Session hasn't started yet";
     }
 
+    flushPendingState();
+
     const auto size = llama_state_get_size(m_ctx);
     std::vector<uint8_t> state(size);
     if (llama_state_get_data(m_ctx, state.data(), size) != size) {
@@ -147,6 +149,8 @@ bool Session::setState(std::span<uint8_t> state) {
     if (llama_state_set_data(m_ctx, state.data(), state.size()) != state.size()) {
         throw_ex{} << "Failed to set state";
     }
+
+    m_state.m_phase = State::Phase::Generating;
     return true;
 }
 
@@ -235,4 +239,11 @@ void Session::doDecode(std::span<const Token> tokens, Source src) {
 
 }
 
+void Session::flushPendingState() {
+    if (m_state.m_currToken != Token_Invalid) {
+        // first yield, then decode, thus we don't decode if the session is aborted
+        doDecode({&m_state.m_currToken, 1}, Source::Generated);
+        m_state.m_currToken = Token_Invalid;
+    }
+}
 } // namespace ac::llama

code/ac/llama/Session.hpp

@@ -26,13 +26,18 @@ class Session {
         bool infiniteContext = true;
     };
     Session(Instance& instance, llama_context* ctx, InitParams params);
+    Session(const Session&) = delete;
+    Session& operator=(const Session&) = delete;
+    ~Session() = default;
 
+    // initial functions to prepare the session
     void setInitialPrompt(std::span<const Token> prompt);
+    bool setState(std::span<uint8_t> state);
 
+    // main functions to interact with the model
     void pushPrompt(std::span<const Token> prompt);
     Token getToken();
     std::vector<uint8_t> getState();
-    bool setState(std::span<uint8_t> state);
 private:
     enum class Source {
         InitialPrompt,
@@ -41,6 +46,7 @@ class Session {
     };
 
     void doDecode(std::span<const Token> tokens, Source src);
+    void flushPendingState();
 
     struct State {
         enum class Phase {

example/e-basic.cpp

@@ -60,7 +60,7 @@ int main() try {
     std::cout << "Prompt: " << prompt << "\n";
 
     // start session
-    auto session = instance.startSession({});
+    auto& session = instance.startSession({});
     session.setInitialPrompt(model.vocab().tokenize(prompt, true, true));
 
     // generate and print 100 tokens

example/e-gui.cpp

@@ -133,7 +133,7 @@ class UModel {
                 ac::llama::Session::InitParams m_params;
                 std::vector<ac::llama::Token> m_promptTokens;
                 std::string m_text;
-                ac::llama::Session m_session;
+                ac::llama::Session& m_session;
                 ac::llama::AntipromptManager m_antiprompt;
                 uint32_t m_numTokens = 0;
             };

test/CMakeLists.txt

@@ -11,7 +11,6 @@ endmacro()
 
 llama_test(Antiprompt)
 llama_test(ChatFormat)
-llama_test(Session)
 
 add_doctest_lib_test(integration ac-llama
     SOURCES