feat(brain): add brain_answer and brain_classify MCP tools
Adds two new LLM-backed MCP tools to the ingestion service: - brain_answer(query): BM25 retrieval + LLM synthesis → answer + sources - brain_classify(text): classifies doc into type/title/tags via LLM Adds llm.Router for primary→fallback routing (berget.ai → iguana). Wired via BRAIN_LLM_PRIMARY_URL/BRAIN_LLM_FALLBACK_URL env vars; no-op when unset so existing deployments are unaffected. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
29
ingestion/internal/llm/router.go
Normal file
29
ingestion/internal/llm/router.go
Normal file
@@ -0,0 +1,29 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// Router calls Primary first; on any error falls back to Fallback.
|
||||
// Fallback may be nil, in which case primary errors are returned directly.
|
||||
type Router struct {
|
||||
Primary *Client
|
||||
Fallback *Client
|
||||
}
|
||||
|
||||
// Complete implements pipeline.CompleteFunc, routing through Primary then Fallback.
|
||||
func (r *Router) Complete(ctx context.Context, system, user string) (string, error) {
|
||||
out, err := r.Primary.Complete(ctx, system, user)
|
||||
if err == nil {
|
||||
return out, nil
|
||||
}
|
||||
if r.Fallback == nil {
|
||||
return "", fmt.Errorf("primary llm: %w", err)
|
||||
}
|
||||
out, err2 := r.Fallback.Complete(ctx, system, user)
|
||||
if err2 != nil {
|
||||
return "", fmt.Errorf("primary llm: %w; fallback llm: %v", err, err2)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
71
ingestion/internal/llm/router_test.go
Normal file
71
ingestion/internal/llm/router_test.go
Normal file
@@ -0,0 +1,71 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestRouter_PrimarySucceeds(t *testing.T) {
|
||||
primary := mockServer(t, "from-primary")
|
||||
defer primary.Close()
|
||||
fallback := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
t.Error("fallback must not be called when primary succeeds")
|
||||
}))
|
||||
defer fallback.Close()
|
||||
|
||||
r := &Router{
|
||||
Primary: New(primary.URL, "", "m", time.Second),
|
||||
Fallback: New(fallback.URL, "", "m", time.Second),
|
||||
}
|
||||
out, err := r.Complete(context.Background(), "sys", "user")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "from-primary", out)
|
||||
}
|
||||
|
||||
func TestRouter_FallsBackOnPrimaryError(t *testing.T) {
|
||||
primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Error(w, "unavailable", http.StatusServiceUnavailable)
|
||||
}))
|
||||
defer primary.Close()
|
||||
fallback := mockServer(t, "from-fallback")
|
||||
defer fallback.Close()
|
||||
|
||||
r := &Router{
|
||||
Primary: New(primary.URL, "", "m", time.Second),
|
||||
Fallback: New(fallback.URL, "", "m", time.Second),
|
||||
}
|
||||
out, err := r.Complete(context.Background(), "sys", "user")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "from-fallback", out)
|
||||
}
|
||||
|
||||
func TestRouter_BothFail(t *testing.T) {
|
||||
fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Error(w, "err", http.StatusBadGateway)
|
||||
}))
|
||||
defer fail.Close()
|
||||
|
||||
r := &Router{
|
||||
Primary: New(fail.URL, "", "m", time.Second),
|
||||
Fallback: New(fail.URL, "", "m", time.Second),
|
||||
}
|
||||
_, err := r.Complete(context.Background(), "sys", "user")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestRouter_NilFallback(t *testing.T) {
|
||||
fail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Error(w, "err", http.StatusBadGateway)
|
||||
}))
|
||||
defer fail.Close()
|
||||
|
||||
r := &Router{Primary: New(fail.URL, "", "m", time.Second)}
|
||||
_, err := r.Complete(context.Background(), "sys", "user")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
Reference in New Issue
Block a user