142 lines
3.7 KiB
Go
142 lines
3.7 KiB
Go
// internal/skills/brain/handlers.go
|
|
package brain
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
)
|
|
|
|
// Handle dispatches brain tool calls.
|
|
func (s *Skill) Handle(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) {
|
|
switch tool {
|
|
case "brain_query":
|
|
return s.query(ctx, args)
|
|
case "brain_write":
|
|
return s.write(ctx, args)
|
|
case "brain_ingest":
|
|
return s.ingest(ctx, args)
|
|
case "brain_search":
|
|
return s.search(ctx, args)
|
|
default:
|
|
return nil, fmt.Errorf("unknown brain tool: %s", tool)
|
|
}
|
|
}
|
|
|
|
type queryArgs struct {
|
|
Query string `json:"query"`
|
|
Limit int `json:"limit,omitempty"`
|
|
}
|
|
|
|
func (s *Skill) query(ctx context.Context, args json.RawMessage) (json.RawMessage, error) {
|
|
var a queryArgs
|
|
if err := json.Unmarshal(args, &a); err != nil {
|
|
return nil, fmt.Errorf("parse args: %w", err)
|
|
}
|
|
if a.Query == "" {
|
|
return nil, fmt.Errorf("query is required")
|
|
}
|
|
if a.Limit == 0 {
|
|
a.Limit = 5
|
|
}
|
|
return s.post(ctx, "/query", a)
|
|
}
|
|
|
|
type writeArgs struct {
|
|
Content string `json:"content"`
|
|
Type string `json:"type,omitempty"`
|
|
Domain string `json:"domain,omitempty"`
|
|
Filename string `json:"filename,omitempty"`
|
|
}
|
|
|
|
func (s *Skill) write(ctx context.Context, args json.RawMessage) (json.RawMessage, error) {
|
|
var a writeArgs
|
|
if err := json.Unmarshal(args, &a); err != nil {
|
|
return nil, fmt.Errorf("parse args: %w", err)
|
|
}
|
|
if a.Content == "" {
|
|
return nil, fmt.Errorf("content is required")
|
|
}
|
|
return s.post(ctx, "/write", a)
|
|
}
|
|
|
|
type ingestArgs struct {
|
|
Content string `json:"content"`
|
|
Source string `json:"source"`
|
|
DryRun bool `json:"dry_run,omitempty"`
|
|
}
|
|
|
|
func (s *Skill) ingest(ctx context.Context, args json.RawMessage) (json.RawMessage, error) {
|
|
var a ingestArgs
|
|
if err := json.Unmarshal(args, &a); err != nil {
|
|
return nil, fmt.Errorf("parse args: %w", err)
|
|
}
|
|
if a.Content == "" {
|
|
return nil, fmt.Errorf("content is required")
|
|
}
|
|
if a.Source == "" {
|
|
return nil, fmt.Errorf("source is required")
|
|
}
|
|
if s.cfg.IngestSvcURL == "" {
|
|
return nil, fmt.Errorf("brain_ingest: INGEST_SVC_URL not configured")
|
|
}
|
|
return s.postTo(ctx, s.cfg.IngestSvcURL+"/ingest", a)
|
|
}
|
|
|
|
type searchArgs struct {
|
|
Query string `json:"query"`
|
|
Collection string `json:"collection,omitempty"`
|
|
Limit int `json:"limit,omitempty"`
|
|
}
|
|
|
|
func (s *Skill) search(ctx context.Context, args json.RawMessage) (json.RawMessage, error) {
|
|
var a searchArgs
|
|
if err := json.Unmarshal(args, &a); err != nil {
|
|
return nil, fmt.Errorf("parse args: %w", err)
|
|
}
|
|
if a.Query == "" {
|
|
return nil, fmt.Errorf("query is required")
|
|
}
|
|
if a.Limit == 0 {
|
|
a.Limit = 5
|
|
}
|
|
if s.cfg.KBRetrievalURL == "" {
|
|
return nil, fmt.Errorf("brain_search: KB_RETRIEVAL_URL not configured")
|
|
}
|
|
return s.postTo(ctx, s.cfg.KBRetrievalURL+"/api/v1/search", a)
|
|
}
|
|
|
|
func (s *Skill) post(ctx context.Context, path string, body any) (json.RawMessage, error) {
|
|
return s.postTo(ctx, s.cfg.IngestBaseURL+path, body)
|
|
}
|
|
|
|
func (s *Skill) postTo(ctx context.Context, url string, body any) (json.RawMessage, error) {
|
|
b, err := json.Marshal(body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("marshal request: %w", err)
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("build request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("call ingestion server: %w", err)
|
|
}
|
|
defer func() { _ = resp.Body.Close() }()
|
|
|
|
out, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read response: %w", err)
|
|
}
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("ingestion server returned %d: %s", resp.StatusCode, out)
|
|
}
|
|
return json.RawMessage(out), nil
|
|
}
|