Initial commit: CarrotAssistant backend (Go + Gin + SQLite, OpenAI-compatible agent runtime)
This commit is contained in:
307
internal/server/chat_test.go
Normal file
307
internal/server/chat_test.go
Normal file
@@ -0,0 +1,307 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"code.littlelan.cn/CarrotAssistant/backend/internal/config"
|
||||
"code.littlelan.cn/CarrotAssistant/backend/internal/security"
|
||||
"code.littlelan.cn/CarrotAssistant/backend/internal/skill"
|
||||
"code.littlelan.cn/CarrotAssistant/backend/internal/store"
|
||||
)
|
||||
|
||||
// newTestDB opens an in-memory-ish SQLite (temp file) and returns the gorm DB
|
||||
// plus a cleanup function.
|
||||
func newTestDB(t *testing.T) (*gorm.DB, func()) {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
db, err := store.Open(dir + "/test.db")
|
||||
if err != nil {
|
||||
t.Fatalf("open db: %v", err)
|
||||
}
|
||||
return db, func() {}
|
||||
}
|
||||
|
||||
// newTestDeps builds a Deps wired to a mock LLM server, with crypto + skills
|
||||
// dir under the test temp directory.
|
||||
func newTestDeps(t *testing.T, mockLLMURL string) (Deps, *store.App, *skill.Skill, string, func()) {
|
||||
t.Helper()
|
||||
db, dbCleanup := newTestDB(t)
|
||||
dir := t.TempDir()
|
||||
crypto, err := security.NewCrypto("0123456789abcdef0123456789abcdef")
|
||||
if err != nil {
|
||||
t.Fatalf("crypto: %v", err)
|
||||
}
|
||||
|
||||
// A model config pointing at the mock server.
|
||||
mc := store.ModelConfig{
|
||||
Name: "mock", BaseURL: mockLLMURL, APIKeyEncrypted: mustEnc(crypto, "sk-mock"),
|
||||
APIKeyPreview: "mock", DefaultModel: "mock-model", SupportsTools: true,
|
||||
}
|
||||
if err := db.Create(&mc).Error; err != nil {
|
||||
t.Fatalf("create mc: %v", err)
|
||||
}
|
||||
|
||||
// An app that uses it, with a known token so the chat client can auth.
|
||||
const appToken = "ca_live_testtoken"
|
||||
app := &store.App{
|
||||
Slug: "demo", Name: "Demo", TokenHash: security.HashToken(appToken),
|
||||
ModelConfigID: &mc.ID, Status: "active",
|
||||
}
|
||||
if err := db.Create(app).Error; err != nil {
|
||||
t.Fatalf("create app: %v", err)
|
||||
}
|
||||
|
||||
// A published skill with one read tool and one mutating tool that
|
||||
// requires confirmation.
|
||||
sk := &skill.Skill{
|
||||
Name: "demo-assistant",
|
||||
Description: "demo",
|
||||
AppSlug: "demo",
|
||||
Permissions: skill.Permissions{
|
||||
DefaultMode: "readonly",
|
||||
RequireConfirmation: []string{"delete_post"},
|
||||
},
|
||||
Tools: []skill.Tool{
|
||||
{
|
||||
Name: "search", Description: "search",
|
||||
Parameters: skill.MakeParameters([]byte(`{"type":"object","properties":{"q":{"type":"string"}}}`)),
|
||||
Endpoint: skill.Endpoint{Method: "GET", URL: "https://example.invalid/search", Query: map[string]string{"q": "{{q}}"}},
|
||||
},
|
||||
{
|
||||
Name: "delete_post", Description: "delete",
|
||||
Parameters: skill.MakeParameters([]byte(`{"type":"object","properties":{"id":{"type":"integer"}}}`)),
|
||||
Endpoint: skill.Endpoint{Method: "DELETE", URL: "https://example.invalid/posts/{{id}}"},
|
||||
},
|
||||
},
|
||||
SystemPrompt: "you are a demo assistant",
|
||||
}
|
||||
mgr := skill.NewManager(dir)
|
||||
if _, err := mgr.Write("demo", "demo-assistant", sk); err != nil {
|
||||
t.Fatalf("write skill: %v", err)
|
||||
}
|
||||
row := store.Skill{
|
||||
AppID: app.ID, Slug: "demo-assistant", Name: sk.Name, FilePath: mgr.RelPath("demo", "demo-assistant"),
|
||||
Status: "published", Source: "manual",
|
||||
}
|
||||
if err := db.Create(&row).Error; err != nil {
|
||||
t.Fatalf("create skill: %v", err)
|
||||
}
|
||||
|
||||
cfg := &config.Config{JWTSecret: "test-jwt"}
|
||||
deps := Deps{DB: db, Config: cfg, Crypto: crypto, SkillsDir: dir}
|
||||
return deps, app, sk, appToken, dbCleanup
|
||||
}
|
||||
|
||||
func mustEnc(c *security.Crypto, s string) string {
|
||||
out, err := c.Encrypt(s)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// mockLLM is a tiny SSE chat-completion server whose response sequence is
|
||||
// driven by a script: a list of chunk-producers invoked in order, one per
|
||||
// LLM call. This lets a test script a tool-call turn followed by a final
|
||||
// answer turn.
|
||||
type mockLLM struct {
|
||||
turns [][]string // each turn is a list of SSE chunks (raw JSON delta lines)
|
||||
idx int
|
||||
}
|
||||
|
||||
func (m *mockLLM) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
// Read the request body to find how many turns we've served (we just
|
||||
// advance the script on each call).
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
_ = body
|
||||
if m.idx >= len(m.turns) {
|
||||
// No more scripted turns: emit a bare stop so the loop ends cleanly.
|
||||
m.idx++
|
||||
writeSSEChunk(w, map[string]any{"choices": []map[string]any{{"finish_reason": "stop", "delta": map[string]any{}}}})
|
||||
fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
return
|
||||
}
|
||||
chunks := m.turns[m.idx]
|
||||
m.idx++
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
for _, c := range chunks {
|
||||
fmt.Fprintf(w, "data: %s\n\n", c)
|
||||
}
|
||||
fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
// writeSSEChunk marshals one chunk object and writes it as a data line.
|
||||
func writeSSEChunk(w http.ResponseWriter, chunk map[string]any) {
|
||||
b, _ := json.Marshal(chunk)
|
||||
fmt.Fprintf(w, "data: %s\n\n", b)
|
||||
}
|
||||
|
||||
// TestChatToolCallFlow drives a chat request against a mock LLM that first
|
||||
// requests the search tool (a read), then produces a final answer. Asserts
|
||||
// the SSE stream emits the expected event sequence.
|
||||
func TestChatToolCallFlow(t *testing.T) {
|
||||
mock := &mockLLM{
|
||||
turns: [][]string{
|
||||
// Turn 1: model asks to call search(q="hi").
|
||||
{
|
||||
chunkJSON("search", "call_1", "", `{"q":"hi"}`, ""),
|
||||
`{"choices":[{"finish_reason":"tool_calls","delta":{}}]}`,
|
||||
},
|
||||
// Turn 2: model produces the final text answer.
|
||||
{
|
||||
`{"choices":[{"finish_reason":"","delta":{"content":"found "}}]}`,
|
||||
`{"choices":[{"finish_reason":"","delta":{"content":"2 posts"}}]}`,
|
||||
`{"choices":[{"finish_reason":"stop","delta":{}}]}`,
|
||||
},
|
||||
},
|
||||
}
|
||||
mockSrv := httptest.NewServer(mock)
|
||||
defer mockSrv.Close()
|
||||
|
||||
deps, _, _, appToken, cleanup := newTestDeps(t, mockSrv.URL)
|
||||
defer cleanup()
|
||||
|
||||
// Also stand up a fake "target site" so the tool call resolves.
|
||||
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte(`["post1","post2"]`))
|
||||
}))
|
||||
defer target.Close()
|
||||
// Patch the search tool URL to point at the fake target. (We do this by
|
||||
// editing the published skill file in place.)
|
||||
patchSkillURL(t, deps, "demo", "demo-assistant", "search", target.URL+"/search")
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
deps.mountChatRoutes(r)
|
||||
|
||||
// POST /api/v1/chat and collect the SSE stream.
|
||||
body := bytes.NewBufferString(`{"message":"hi"}`)
|
||||
req := httptest.NewRequest("POST", "/api/v1/chat", body)
|
||||
req.Header.Set("X-App-Token", appToken)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
r.ServeHTTP(rec, req)
|
||||
|
||||
events := parseSSE(rec.Body.String())
|
||||
kinds := eventKinds(events)
|
||||
wantSeq := []string{"session", "tool_call", "tool_result", "token", "token", "done"}
|
||||
if !containsSeq(kinds, wantSeq) {
|
||||
t.Fatalf("event kinds = %v, want sequence containing %v", kinds, wantSeq)
|
||||
}
|
||||
}
|
||||
|
||||
// patchSkillURL rewrites a tool's endpoint URL in the published skill file
|
||||
// for the test, so tool calls hit the in-process fake target.
|
||||
func patchSkillURL(t *testing.T, deps Deps, appSlug, skillSlug, toolName, newURL string) {
|
||||
t.Helper()
|
||||
mgr := skill.NewManager(deps.SkillsDir)
|
||||
var row store.Skill
|
||||
deps.DB.Where("slug = ? AND app_id IN (SELECT id FROM apps WHERE slug = ?)", skillSlug, appSlug).First(&row)
|
||||
sk, err := mgr.Load(&row)
|
||||
if err != nil {
|
||||
t.Fatalf("load skill: %v", err)
|
||||
}
|
||||
for i := range sk.Tools {
|
||||
if sk.Tools[i].Name == toolName {
|
||||
sk.Tools[i].Endpoint.URL = newURL
|
||||
}
|
||||
}
|
||||
if _, err := mgr.Write(appSlug, skillSlug, sk); err != nil {
|
||||
t.Fatalf("rewrite skill: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// chunkJSON builds a delta chunk for a tool call (function arguments arrive
|
||||
// as a JSON string). When content is non-empty it's a content delta instead.
|
||||
func chunkJSON(name, id, content, argsJSON, _ string) string {
|
||||
if content != "" {
|
||||
b, _ := json.Marshal(map[string]any{"choices": []map[string]any{{
|
||||
"delta": map[string]any{"content": content},
|
||||
}}})
|
||||
return string(b)
|
||||
}
|
||||
// tool-call delta
|
||||
b, _ := json.Marshal(map[string]any{"choices": []map[string]any{{
|
||||
"delta": map[string]any{
|
||||
"tool_calls": []map[string]any{{
|
||||
"index": 0, "id": id,
|
||||
"function": map[string]any{"name": name, "arguments": argsJSON},
|
||||
}},
|
||||
},
|
||||
}}})
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// parseSSE turns a raw SSE body into a list of (event, data) pairs.
|
||||
func parseSSE(raw string) []sseEvent {
|
||||
var out []sseEvent
|
||||
var cur sseEvent
|
||||
sc := bufio.NewScanner(strings.NewReader(raw))
|
||||
sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
|
||||
for sc.Scan() {
|
||||
ln := sc.Text()
|
||||
if ln == "" {
|
||||
if cur.event != "" || cur.data != "" {
|
||||
out = append(out, cur)
|
||||
}
|
||||
cur = sseEvent{}
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(ln, "event: ") {
|
||||
cur.event = strings.TrimPrefix(ln, "event: ")
|
||||
} else if strings.HasPrefix(ln, "data: ") {
|
||||
cur.data += strings.TrimPrefix(ln, "data: ")
|
||||
}
|
||||
}
|
||||
if cur.event != "" || cur.data != "" {
|
||||
out = append(out, cur)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
type sseEvent struct {
|
||||
event string
|
||||
data string
|
||||
}
|
||||
|
||||
func eventKinds(es []sseEvent) []string {
|
||||
out := make([]string, 0, len(es))
|
||||
for _, e := range es {
|
||||
out = append(out, e.event)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// containsSeq reports whether want appears as a contiguous subsequence of got.
|
||||
func containsSeq(got, want []string) bool {
|
||||
if len(want) > len(got) {
|
||||
return false
|
||||
}
|
||||
for i := 0; i+len(want) <= len(got); i++ {
|
||||
match := true
|
||||
for j := range want {
|
||||
if got[i+j] != want[j] {
|
||||
match = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if match {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
Reference in New Issue
Block a user