wehub-resource-sync
446 lines
13 KiB
Go
446 lines
13 KiB
Go
package agents
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"github.com/mudler/cogito"
|
|
openai "github.com/sashabaranov/go-openai"
|
|
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
)
|
|
|
|
// mockLLM implements cogito.LLM and returns a fixed response.
|
|
type mockLLM struct {
|
|
response string
|
|
callCount atomic.Int32
|
|
}
|
|
|
|
func (m *mockLLM) Ask(ctx context.Context, f cogito.Fragment) (cogito.Fragment, error) {
|
|
m.callCount.Add(1)
|
|
return f.AddMessage(cogito.AssistantMessageRole, m.response), nil
|
|
}
|
|
|
|
func (m *mockLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) {
|
|
m.callCount.Add(1)
|
|
msg := openai.ChatCompletionMessage{
|
|
Role: "assistant",
|
|
Content: m.response,
|
|
}
|
|
return cogito.LLMReply{
|
|
ChatCompletionResponse: openai.ChatCompletionResponse{
|
|
Choices: []openai.ChatCompletionChoice{{Message: msg}},
|
|
},
|
|
}, cogito.LLMUsage{}, nil
|
|
}
|
|
|
|
type toolCallingMockLLM struct {
|
|
createResponses []openai.ChatCompletionResponse
|
|
askResponse string
|
|
callCount atomic.Int32
|
|
}
|
|
|
|
func (m *toolCallingMockLLM) Ask(ctx context.Context, f cogito.Fragment) (cogito.Fragment, error) {
|
|
m.callCount.Add(1)
|
|
return f.AddMessage(cogito.AssistantMessageRole, m.askResponse), nil
|
|
}
|
|
|
|
func (m *toolCallingMockLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (cogito.LLMReply, cogito.LLMUsage, error) {
|
|
idx := int(m.callCount.Add(1)) - 1
|
|
if idx >= len(m.createResponses) {
|
|
return cogito.LLMReply{
|
|
ChatCompletionResponse: openai.ChatCompletionResponse{
|
|
Choices: []openai.ChatCompletionChoice{{
|
|
Message: openai.ChatCompletionMessage{
|
|
Role: "assistant",
|
|
Content: "No more tools needed.",
|
|
},
|
|
}},
|
|
},
|
|
}, cogito.LLMUsage{}, nil
|
|
}
|
|
return cogito.LLMReply{ChatCompletionResponse: m.createResponses[idx]}, cogito.LLMUsage{}, nil
|
|
}
|
|
|
|
// statusCollector records status callbacks in a thread-safe way.
|
|
type statusCollector struct {
|
|
mu sync.Mutex
|
|
statuses []string
|
|
}
|
|
|
|
func (sc *statusCollector) onStatus(s string) {
|
|
sc.mu.Lock()
|
|
defer sc.mu.Unlock()
|
|
sc.statuses = append(sc.statuses, s)
|
|
}
|
|
|
|
func (sc *statusCollector) get() []string {
|
|
sc.mu.Lock()
|
|
defer sc.mu.Unlock()
|
|
cp := make([]string, len(sc.statuses))
|
|
copy(cp, sc.statuses)
|
|
return cp
|
|
}
|
|
|
|
var _ = DescribeTable("stripThinkingTags",
|
|
func(input, want string) {
|
|
Expect(stripThinkingTags(input)).To(Equal(want))
|
|
},
|
|
Entry("empty string", "", ""),
|
|
Entry("no tags", "Hello, world!", "Hello, world!"),
|
|
Entry("single tag pair", "before<thinking>secret thoughts</thinking>after", "beforeafter"),
|
|
Entry("multiple tag pairs", "a<thinking>one</thinking>b<thinking>two</thinking>c", "abc"),
|
|
Entry("nested tags", "<thinking>outer<thinking>inner</thinking>still outer</thinking>visible", "still outer</thinking>visible"),
|
|
Entry("unclosed opening tag", "hello<thinking>this is unclosed", "hello<thinking>this is unclosed"),
|
|
Entry("only closing tag", "hello</thinking>world", "hello</thinking>world"),
|
|
Entry("tags with whitespace around content", "before<thinking> spaced out </thinking>after", "beforeafter"),
|
|
Entry("empty thinking block", "before<thinking></thinking>after", "beforeafter"),
|
|
Entry("multiline thinking block", "before<thinking>\nline1\nline2\n</thinking>after", "beforeafter"),
|
|
Entry("adjacent tag pairs", "<thinking>a</thinking><thinking>b</thinking>", ""),
|
|
)
|
|
|
|
var _ = DescribeTable("appendKBCitations",
|
|
func(response, collection, userID string, citations []KBCitation, want string) {
|
|
Expect(AppendKBCitations(response, collection, userID, citations)).To(Equal(want))
|
|
},
|
|
Entry("leaves responses without citations unchanged",
|
|
"answer",
|
|
"agent",
|
|
"",
|
|
nil,
|
|
"answer",
|
|
),
|
|
Entry("leaves blank responses unchanged",
|
|
"",
|
|
"agent",
|
|
"",
|
|
[]KBCitation{{FileName: "source.pdf", EntryKey: "uuid/source.pdf"}},
|
|
"",
|
|
),
|
|
Entry("appends clickable source links",
|
|
"answer",
|
|
"my-agent",
|
|
"",
|
|
[]KBCitation{{FileName: "new feature.pdf", EntryKey: "uuid/new feature.pdf"}},
|
|
"answer\n\nSources:\n[1] [new feature.pdf](/api/agents/collections/my-agent/entries-raw/uuid/new%20feature.pdf)",
|
|
),
|
|
Entry("deduplicates citations by entry key",
|
|
"answer",
|
|
"agent",
|
|
"",
|
|
[]KBCitation{
|
|
{FileName: "first.pdf", EntryKey: "uuid/shared.pdf"},
|
|
{FileName: "second.pdf", EntryKey: "uuid/shared.pdf"},
|
|
},
|
|
"answer\n\nSources:\n[1] [first.pdf](/api/agents/collections/agent/entries-raw/uuid/shared.pdf)",
|
|
),
|
|
Entry("uses plain text when entry key is missing",
|
|
"answer",
|
|
"agent",
|
|
"",
|
|
[]KBCitation{{FileName: "source.pdf"}},
|
|
"answer\n\nSources:\n[1] source.pdf",
|
|
),
|
|
Entry("uses entry basename when filename is missing",
|
|
"answer",
|
|
"agent",
|
|
"",
|
|
[]KBCitation{{EntryKey: "uuid/source.pdf"}},
|
|
"answer\n\nSources:\n[1] [source.pdf](/api/agents/collections/agent/entries-raw/uuid/source.pdf)",
|
|
),
|
|
Entry("adds user id query when present",
|
|
"answer",
|
|
"agent",
|
|
"user 1",
|
|
[]KBCitation{{FileName: "source.pdf", EntryKey: "uuid/source.pdf"}},
|
|
"answer\n\nSources:\n[1] [source.pdf](/api/agents/collections/agent/entries-raw/uuid/source.pdf?user_id=user+1)",
|
|
),
|
|
Entry("escapes collection, path segments, and markdown link text",
|
|
"answer",
|
|
"agent one",
|
|
"",
|
|
[]KBCitation{{FileName: "source [draft].pdf", EntryKey: "uuid/source [draft].pdf"}},
|
|
`answer
|
|
|
|
Sources:
|
|
[1] [source \[draft\].pdf](/api/agents/collections/agent%20one/entries-raw/uuid/source%20%5Bdraft%5D.pdf)`,
|
|
),
|
|
)
|
|
|
|
var _ = Describe("ExecuteChatWithLLM", func() {
|
|
var (
|
|
ctx context.Context
|
|
sc *statusCollector
|
|
cb Callbacks
|
|
)
|
|
|
|
BeforeEach(func() {
|
|
ctx = context.Background()
|
|
sc = &statusCollector{}
|
|
cb = Callbacks{
|
|
OnStatus: sc.onStatus,
|
|
}
|
|
})
|
|
|
|
Context("basic chat completion", func() {
|
|
It("returns the LLM response", func() {
|
|
llm := &mockLLM{response: "Hello from the agent!"}
|
|
cfg := &AgentConfig{
|
|
Name: "test-agent",
|
|
Model: "test-model",
|
|
}
|
|
|
|
result, err := ExecuteChatWithLLM(ctx, llm, cfg, "Hi there", cb)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(result).To(Equal("Hello from the agent!"))
|
|
|
|
statuses := sc.get()
|
|
Expect(statuses).To(ContainElement("processing"))
|
|
Expect(statuses).To(ContainElement("completed"))
|
|
})
|
|
})
|
|
|
|
Context("empty model name", func() {
|
|
It("still succeeds because ExecuteChatWithLLM receives a pre-built LLM", func() {
|
|
// ExecuteChatWithLLM does not check cfg.Model — that's ExecuteChat's job.
|
|
// Verify it does not error with an empty model name.
|
|
llm := &mockLLM{response: "ok"}
|
|
cfg := &AgentConfig{
|
|
Name: "no-model-agent",
|
|
Model: "",
|
|
}
|
|
|
|
result, err := ExecuteChatWithLLM(ctx, llm, cfg, "test", cb)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(result).To(Equal("ok"))
|
|
})
|
|
})
|
|
|
|
Context("ExecuteChat rejects empty model", func() {
|
|
It("returns an error when model is empty", func() {
|
|
cfg := &AgentConfig{
|
|
Name: "empty-model",
|
|
Model: "",
|
|
}
|
|
|
|
_, err := ExecuteChat(ctx, "http://localhost", "key", cfg, "test", cb)
|
|
Expect(err).To(HaveOccurred())
|
|
Expect(err.Error()).To(ContainSubstring("no model configured"))
|
|
})
|
|
})
|
|
|
|
Context("StripThinkingTags flag", func() {
|
|
It("strips thinking tags from the response when enabled", func() {
|
|
llm := &mockLLM{response: "before<thinking>secret</thinking>after"}
|
|
cfg := &AgentConfig{
|
|
Name: "strip-agent",
|
|
Model: "test-model",
|
|
StripThinkingTags: true,
|
|
}
|
|
|
|
result, err := ExecuteChatWithLLM(ctx, llm, cfg, "test", cb)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(result).To(Equal("beforeafter"))
|
|
})
|
|
|
|
It("preserves thinking tags when flag is disabled", func() {
|
|
llm := &mockLLM{response: "before<thinking>secret</thinking>after"}
|
|
cfg := &AgentConfig{
|
|
Name: "no-strip-agent",
|
|
Model: "test-model",
|
|
StripThinkingTags: false,
|
|
}
|
|
|
|
result, err := ExecuteChatWithLLM(ctx, llm, cfg, "test", cb)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(result).To(Equal("before<thinking>secret</thinking>after"))
|
|
})
|
|
})
|
|
|
|
Context("OnMessage callback", func() {
|
|
It("delivers the agent response via OnMessage", func() {
|
|
var msgSender, msgContent string
|
|
cb.OnMessage = func(sender, content, messageID string) {
|
|
msgSender = sender
|
|
msgContent = content
|
|
}
|
|
|
|
llm := &mockLLM{response: "agent reply"}
|
|
cfg := &AgentConfig{
|
|
Name: "msg-agent",
|
|
Model: "test-model",
|
|
}
|
|
|
|
_, err := ExecuteChatWithLLM(ctx, llm, cfg, "hello", cb)
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(msgSender).To(Equal("agent"))
|
|
Expect(msgContent).To(Equal("agent reply"))
|
|
})
|
|
})
|
|
|
|
Context("knowledge base citations", func() {
|
|
It("appends KB sources to the returned response and callback message", func() {
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/api/agents/collections/kb-agent/search", func(w http.ResponseWriter, r *http.Request) {
|
|
Expect(r.URL.Query().Get("user_id")).To(Equal("user-1"))
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{
|
|
"results": [
|
|
{
|
|
"content": "KB content",
|
|
"id": "result-1",
|
|
"similarity": 0.99,
|
|
"metadata": {
|
|
"file_name": "new feature.pdf",
|
|
"source": "uuid/new feature.pdf"
|
|
}
|
|
}
|
|
],
|
|
"count": 1
|
|
}`))
|
|
})
|
|
server := httptest.NewServer(mux)
|
|
defer server.Close()
|
|
|
|
var msgContent string
|
|
cb.OnMessage = func(sender, content, messageID string) {
|
|
msgContent = content
|
|
}
|
|
|
|
llm := &mockLLM{response: "agent reply"}
|
|
cfg := &AgentConfig{
|
|
Name: "kb-agent",
|
|
Model: "test-model",
|
|
EnableKnowledgeBase: true,
|
|
KBMode: KBModeAutoSearch,
|
|
}
|
|
|
|
result, err := ExecuteChatWithLLM(ctx, llm, cfg, "hello", cb, ExecuteChatOpts{
|
|
APIURL: server.URL,
|
|
UserID: "user-1",
|
|
})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(result).To(Equal("agent reply\n\nSources:\n[1] [new feature.pdf](/api/agents/collections/kb-agent/entries-raw/uuid/new%20feature.pdf?user_id=user-1)"))
|
|
Expect(msgContent).To(Equal(result))
|
|
})
|
|
|
|
It("collects citations from the search_memory tool", func() {
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/api/agents/collections/kb-agent/search", func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{
|
|
"results": [
|
|
{
|
|
"content": "Tool KB content",
|
|
"id": "result-1",
|
|
"similarity": 0.99,
|
|
"metadata": {
|
|
"file_name": "tool source.pdf",
|
|
"source": "uuid/tool source.pdf"
|
|
}
|
|
}
|
|
],
|
|
"count": 1
|
|
}`))
|
|
})
|
|
server := httptest.NewServer(mux)
|
|
defer server.Close()
|
|
|
|
collector := &kbCitationList{}
|
|
tool := KBSearchMemoryTool{
|
|
APIURL: server.URL,
|
|
Collection: "kb-agent",
|
|
CitationCollector: collector,
|
|
}
|
|
|
|
result, _, err := tool.Run(KBSearchMemoryArgs{Query: "hello"})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(result).To(ContainSubstring("Tool KB content"))
|
|
Expect(collector.Citations()).To(Equal([]KBCitation{{FileName: "tool source.pdf", EntryKey: "uuid/tool source.pdf"}}))
|
|
})
|
|
|
|
It("appends KB sources found through tools-only search_memory calls", func() {
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/api/agents/collections/kb-agent/search", func(w http.ResponseWriter, r *http.Request) {
|
|
Expect(r.URL.Query().Get("user_id")).To(Equal("user-1"))
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{
|
|
"results": [
|
|
{
|
|
"content": "Tool KB content",
|
|
"id": "result-1",
|
|
"similarity": 0.99,
|
|
"metadata": {
|
|
"file_name": "tool source.pdf",
|
|
"source": "uuid/tool source.pdf"
|
|
}
|
|
}
|
|
],
|
|
"count": 1
|
|
}`))
|
|
})
|
|
server := httptest.NewServer(mux)
|
|
defer server.Close()
|
|
|
|
llm := &toolCallingMockLLM{
|
|
askResponse: "agent reply from tool context",
|
|
createResponses: []openai.ChatCompletionResponse{
|
|
{
|
|
Choices: []openai.ChatCompletionChoice{
|
|
{
|
|
Message: openai.ChatCompletionMessage{
|
|
Role: "assistant",
|
|
ToolCalls: []openai.ToolCall{
|
|
{
|
|
ID: "call-1",
|
|
Type: openai.ToolTypeFunction,
|
|
Function: openai.FunctionCall{
|
|
Name: "search_memory",
|
|
Arguments: `{"query":"hello"}`,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
cfg := &AgentConfig{
|
|
Name: "kb-agent",
|
|
Model: "test-model",
|
|
EnableKnowledgeBase: true,
|
|
KBMode: KBModeTools,
|
|
}
|
|
|
|
result, err := ExecuteChatWithLLM(ctx, llm, cfg, "hello", cb, ExecuteChatOpts{
|
|
APIURL: server.URL,
|
|
UserID: "user-1",
|
|
})
|
|
Expect(err).ToNot(HaveOccurred())
|
|
Expect(result).To(Equal("agent reply from tool context\n\nSources:\n[1] [tool source.pdf](/api/agents/collections/kb-agent/entries-raw/uuid/tool%20source.pdf?user_id=user-1)"))
|
|
})
|
|
})
|
|
|
|
Context("context cancellation", func() {
|
|
It("returns an error when context is already cancelled", func() {
|
|
cancelledCtx, cancel := context.WithCancel(ctx)
|
|
cancel() // immediately cancel
|
|
|
|
llm := &mockLLM{response: "should not reach"}
|
|
cfg := &AgentConfig{
|
|
Name: "cancel-agent",
|
|
Model: "test-model",
|
|
}
|
|
|
|
_, err := ExecuteChatWithLLM(cancelledCtx, llm, cfg, "test", cb)
|
|
Expect(err).To(HaveOccurred())
|
|
Expect(err.Error()).To(ContainSubstring("agent execution failed"))
|
|
})
|
|
})
|
|
})
|