From 4e9b31baf885c6889d685b255ceb9bee6db7c06b Mon Sep 17 00:00:00 2001
From: Grail Finder <wohilas@gmail.com>
Date: Sun, 2 Feb 2025 10:05:43 +0300
Subject: Feat: remove thinking and tool msgs from ctx

---
 main_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 main_test.go

(limited to 'main_test.go')

diff --git a/main_test.go b/main_test.go
new file mode 100644
index 0000000..0046ca2
--- /dev/null
+++ b/main_test.go
@@ -0,0 +1,41 @@
+package main
+
+import (
+	"elefant/models"
+	"fmt"
+	"strings"
+	"testing"
+)
+
+func TestRemoveThinking(t *testing.T) {
+	cases := []struct {
+		cb       *models.ChatBody
+		toolMsgs uint8
+	}{
+		{cb: &models.ChatBody{
+			Stream: true,
+			Messages: []models.RoleMsg{
+				{Role: "tool", Content: "should be ommited"},
+				{Role: "system", Content: "should stay"},
+				{Role: "user", Content: "hello, how are you?"},
+				{Role: "assistant", Content: "Oh, hi. <think>I should thank user and continue the conversation</think> I am geat, thank you! How are you?"},
+			},
+		},
+			toolMsgs: uint8(1),
+		},
+	}
+	for i, tc := range cases {
+		t.Run(fmt.Sprintf("run_%d", i), func(t *testing.T) {
+			mNum := len(tc.cb.Messages)
+			removeThinking(tc.cb)
+			if len(tc.cb.Messages) != mNum-int(tc.toolMsgs) {
+				t.Error("failed to delete tools msg", tc.cb.Messages, cfg.ToolRole)
+			}
+			for _, msg := range tc.cb.Messages {
+				if strings.Contains(msg.Content, "<think>") {
+					t.Errorf("msg contains think tag; msg: %s\n", msg.Content)
+				}
+			}
+		})
+	}
+}
-- 
cgit v1.2.3