72 lines
2.5 KiB
Go
72 lines
2.5 KiB
Go
package ollama
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"github.com/example/glpi-ai-agent/internal/model"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestAnalyseStructured(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var body map[string]any
|
|
_ = json.NewDecoder(r.Body).Decode(&body)
|
|
format, _ := body["format"].(map[string]any)
|
|
if format == nil {
|
|
t.Error("missing schema")
|
|
} else if props, _ := format["properties"].(map[string]any); props != nil {
|
|
if category, _ := props["category"].(map[string]any); category != nil {
|
|
if categoryProps, _ := category["properties"].(map[string]any); categoryProps != nil {
|
|
if _, exists := categoryProps["change"]; exists {
|
|
t.Error("category schema must not let the model decide change=true/false")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
options, _ := body["options"].(map[string]any)
|
|
if options["num_predict"] != float64(256) {
|
|
t.Errorf("unexpected num_predict: %v", options["num_predict"])
|
|
}
|
|
if body["keep_alive"] != "10m0s" {
|
|
t.Errorf("unexpected keep_alive: %v", body["keep_alive"])
|
|
}
|
|
if body["think"] != false {
|
|
t.Errorf("unexpected think: %v", body["think"])
|
|
}
|
|
json.NewEncoder(w).Encode(map[string]any{"message": map[string]any{"content": `{"category":{"id":1,"confidence":0.9},"reply":{"allowed":false,"confidence":0.1,"knowledge_id":""},"reason":"ok"}`}})
|
|
}))
|
|
defer srv.Close()
|
|
c := New(srv.URL, "m", "e", "de-DE", "formal", time.Second, 256, 10*time.Minute, false, 1, 1)
|
|
d, err := c.Analyse(context.Background(), model.Ticket{ID: 1}, []model.Category{{ID: 1}}, nil, model.ContextSnapshot{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if d.Reason != "ok" {
|
|
t.Fatalf("unexpected %+v", d)
|
|
}
|
|
}
|
|
|
|
func TestAnalyseRetriesInvalidJSON(t *testing.T) {
|
|
calls := 0
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
calls++
|
|
content := `{"category":`
|
|
if calls > 1 {
|
|
content = `{"category":{"id":2,"confidence":0.95},"reply":{"allowed":false,"confidence":0.1,"knowledge_id":""},"reason":"ok"}`
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"message": map[string]any{"content": content}})
|
|
}))
|
|
defer srv.Close()
|
|
c := New(srv.URL, "m", "e", "de-DE", "formal", time.Second, 768, time.Minute, false, 1, 1)
|
|
d, err := c.Analyse(context.Background(), model.Ticket{ID: 1}, []model.Category{{ID: 2, Name: "Active Directory"}}, nil, model.ContextSnapshot{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if calls != 2 || d.Category.ID != 2 {
|
|
t.Fatalf("calls=%d decision=%+v", calls, d)
|
|
}
|
|
}
|