hyperguild/ingestion/internal/mcp/tools_graph.go

package mcp

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/mathiasbq/hyperguild/ingestion/internal/graphstore"
)

// graphReader is the read-side surface of graphstore.PGStore the
// brain_graph handler needs. Splitting it out (vs. depending on the
// concrete *PGStore) lets tests inject a fake without standing up
// postgres, and keeps the write-side graphsync.Store interface free
// of query concerns.
type graphReader interface {
	Neighbors(ctx context.Context, slug, edgeType string, limit int) ([]graphstore.Neighbor, error)
	Subgraph(ctx context.Context, origin string, depth int) ([]graphstore.Neighbor, error)
	Path(ctx context.Context, src, dst string, maxDepth int) ([]graphstore.PathStep, error)
}

// Compile-time check that *graphstore.PGStore satisfies graphReader.
var _ graphReader = (*graphstore.PGStore)(nil)

type brainGraphArgs struct {
	Op       string `json:"op"`
	Slug     string `json:"slug,omitempty"`
	Src      string `json:"src,omitempty"`
	Dst      string `json:"dst,omitempty"`
	EdgeType string `json:"edge_type,omitempty"`
	Limit    int    `json:"limit,omitempty"`
	Depth    int    `json:"depth,omitempty"`
}

func (s *Server) brainGraph(ctx context.Context, args json.RawMessage) (json.RawMessage, error) {
	reader, ok := s.graph.(graphReader)
	if s.graph == nil || !ok {
		return nil, fmt.Errorf("brain graph not configured: set BRAIN_GRAPH_ENABLED=true")
	}
	var a brainGraphArgs
	if err := json.Unmarshal(args, &a); err != nil {
		return nil, fmt.Errorf("parse args: %w", err)
	}

	switch a.Op {
	case "neighbors":
		if a.Slug == "" {
			return nil, fmt.Errorf("slug is required for op=neighbors")
		}
		ns, err := reader.Neighbors(ctx, a.Slug, a.EdgeType, a.Limit)
		if err != nil {
			return nil, fmt.Errorf("neighbors: %w", err)
		}
		return json.Marshal(map[string]any{"results": neighborsView(ns)})

	case "subgraph":
		if a.Slug == "" {
			return nil, fmt.Errorf("slug is required for op=subgraph")
		}
		ns, err := reader.Subgraph(ctx, a.Slug, a.Depth)
		if err != nil {
			return nil, fmt.Errorf("subgraph: %w", err)
		}
		return json.Marshal(map[string]any{"results": neighborsView(ns)})

	case "path":
		if a.Src == "" || a.Dst == "" {
			return nil, fmt.Errorf("src and dst are required for op=path")
		}
		steps, err := reader.Path(ctx, a.Src, a.Dst, a.Depth)
		if err != nil {
			return nil, fmt.Errorf("path: %w", err)
		}
		return json.Marshal(map[string]any{"steps": pathView(steps)})

	default:
		return nil, fmt.Errorf("unknown op %q (want neighbors|subgraph|path)", a.Op)
	}
}

type neighborView struct {
	Slug     string `json:"slug"`
	Type     string `json:"type,omitempty"`
	Wing     string `json:"wing,omitempty"`
	Hall     string `json:"hall,omitempty"`
	DocPath  string `json:"doc_path,omitempty"`
	Title    string `json:"title,omitempty"`
	EdgeType string `json:"edge_type"`
	Distance int    `json:"distance"`
}

func neighborsView(ns []graphstore.Neighbor) []neighborView {
	out := make([]neighborView, 0, len(ns))
	for _, n := range ns {
		out = append(out, neighborView{
			Slug: n.Slug, Type: n.Type, Wing: n.Wing, Hall: n.Hall,
			DocPath: n.DocPath, Title: n.Title,
			EdgeType: n.EdgeType, Distance: n.Distance,
		})
	}
	return out
}

type pathStepView struct {
	From     string `json:"from"`
	To       string `json:"to"`
	EdgeType string `json:"edge_type"`
}

func pathView(steps []graphstore.PathStep) []pathStepView {
	out := make([]pathStepView, 0, len(steps))
	for _, s := range steps {
		out = append(out, pathStepView{From: s.FromSlug, To: s.ToSlug, EdgeType: s.EdgeType})
	}
	return out
}