Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ This Go-based MCP server acts as a bridge between AI applications and Collibra,
- [`get_lineage_entity`](pkg/tools/get_lineage_entity/) - Get metadata about a specific entity in the technical lineage graph
- [`get_lineage_transformation`](pkg/tools/get_lineage_transformation/) - Get details and logic of a specific data transformation
- [`get_lineage_upstream`](pkg/tools/get_lineage_upstream/) - Get upstream technical lineage (sources) for a data entity
- [`get_maestro_agent`](pkg/tools/get_maestro_agent/) - Read the configuration of an existing AI Maestro agent: name, handle, description, instructions, color, welcome message, sample questions, tools, knowledge base views, sharing, ownership, plus its status and whether its definition is complete. The connecting user must be able to see the agent
- [`get_measure_data`](pkg/tools/get_measure_data/) - Trace a measure back to its underlying physical columns and tables
- [`get_table_semantics`](pkg/tools/get_table_semantics/) - Retrieve the semantic layer for a table: columns, data attributes, and connected measures
- [`list_asset_types`](pkg/tools/list_asset_types/) - List available asset types
Expand All @@ -44,6 +45,8 @@ This Go-based MCP server acts as a bridge between AI applications and Collibra,
- `set_owner` - set the owner by user UUID
- `set_assignees` - replace the assignee list with the given users/groups
- `set_visibility` - set whether the assessment is visible to everyone
- [`create_maestro_agent`](pkg/tools/create_maestro_agent/) - Create an AI Maestro agent. Only `name` and `handle` are required; description, instructions, color, welcome message, sample questions, tools, sharing and ownership are optional. The agent is created in DRAFT status with the connecting user as its creator and owner, and has to be submitted and approved in AI Maestro before end users can use it. Its knowledge base and status are not settable here
- [`edit_maestro_agent`](pkg/tools/edit_maestro_agent/) - Edit an existing AI Maestro agent as a partial update: pass only the fields to change and the rest keep their stored values. Lists (`sampleQuestions`, `tools`, `sharing`, `ownership`) are replaced wholesale, and an empty list clears them. Any edit returns the agent to DRAFT, so a published agent leaves Collibra Copilot until it is approved again. The connecting user must be an owner of the agent, or hold the permission to manage all AI agents
- [`edit_asset`](pkg/tools/edit_asset/) - Edit an existing asset via a list of typed operations:
- `set_attribute`, `add_attribute`, `remove_attribute` - set an attribute value (creates if empty, updates if present), append an extra value to a multi-valued attribute, or clear one (e.g. `Definition`, `Note`)
- `update_property` - rename the asset (`name`), change its `displayName`, or change its `statusId` (status name or UUID accepted)
Expand Down
238 changes: 238 additions & 0 deletions pkg/clients/ai_maestro_client.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,238 @@
package clients

import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"log/slog"
"net/http"
"net/url"
)

// AI Maestro API v1 (/rest/aiMaestro/v1) β€” the AI Maestro application's own REST
// API, where Maestro Agent definitions live. Agents are NOT catalog assets, so
// they are reached here rather than through /rest/2.0. Served on the same host as
// the DGC API, so the shared RoundTripper injects host + auth transparently.
const aiMaestroAPIBasePath = "/rest/aiMaestro/v1"

// MaestroAgent is a Maestro Agent as the API returns it, runtime fields included.
type MaestroAgent struct {
ID string `json:"id"`
Name string `json:"name"`
Handle string `json:"handle"`
Color string `json:"color,omitempty"`
Description string `json:"description,omitempty"`
IsValid bool `json:"isValid"`
Instructions string `json:"instructions,omitempty"`
KnowledgeBase *MaestroKnowledgeBase `json:"knowledgeBase,omitempty"`
WelcomeMessage string `json:"welcomeMessage,omitempty"`
SampleQuestions []string `json:"sampleQuestions,omitempty"`
Tools []string `json:"tools,omitempty"`
Sharing *MaestroSharing `json:"sharing,omitempty"`
Ownership *MaestroOwnership `json:"ownership,omitempty"`
Status string `json:"status,omitempty"`
CreatedBy string `json:"createdBy,omitempty"`
LastModifiedOn string `json:"lastModifiedOn,omitempty"`
LastModifiedBy string `json:"lastModifiedBy,omitempty"`
}

// MaestroKnowledgeBase is the set of Collibra views the agent may draw on. Read
// only: knowledge base views are configured in AI Maestro, so neither agent
// request type carries them.
type MaestroKnowledgeBase struct {
Views []MaestroKnowledgeBaseView `json:"views,omitempty"`
}

type MaestroKnowledgeBaseView struct {
ViewID string `json:"viewId"`
Location string `json:"location"`
SelectedCommunityID string `json:"selectedCommunityId,omitempty"`
}

// MaestroSharing lists who, beside the creator, sees the agent in Collibra Copilot.
// The API requires all three lists once the object is present, so none of them
// carry omitempty.
type MaestroSharing struct {
Roles []string `json:"roles"`
Groups []string `json:"groups"`
Users []string `json:"users"`
}

// MaestroOwnership lists who, beside the creator, may edit the agent.
type MaestroOwnership struct {
Users []string `json:"users"`
}

// CreateAgentRequest is the POST /agents body. status is deliberately absent: the
// endpoint only accepts DRAFT, which is also its default. knowledgeBase is absent
// too β€” views are configured in AI Maestro, not through this client.
type CreateAgentRequest struct {
Name string `json:"name"`
Handle string `json:"handle"`
Color string `json:"color,omitempty"`
Description string `json:"description,omitempty"`
Instructions string `json:"instructions,omitempty"`
WelcomeMessage string `json:"welcomeMessage,omitempty"`
SampleQuestions []string `json:"sampleQuestions,omitempty"`
Tools []string `json:"tools,omitempty"`
Sharing *MaestroSharing `json:"sharing,omitempty"`
Ownership *MaestroOwnership `json:"ownership,omitempty"`
}

// PatchAgentRequest is the PATCH /agents/{agentId} body. Every field is optional
// and nil means "leave unchanged", so pointers are needed throughout: an omitted
// list preserves the stored one, an empty list clears it.
//
// status is deliberately absent. The server reads its absence as "demote to
// DRAFT", which is what an edit should do β€” edited content has to go back through
// review before end users see it again. knowledgeBase is absent so an edit never
// touches the agent's views.
type PatchAgentRequest struct {
Name *string `json:"name,omitempty"`
Handle *string `json:"handle,omitempty"`
Color *string `json:"color,omitempty"`
Description *string `json:"description,omitempty"`
Instructions *string `json:"instructions,omitempty"`
WelcomeMessage *string `json:"welcomeMessage,omitempty"`
SampleQuestions *[]string `json:"sampleQuestions,omitempty"`
Tools *[]string `json:"tools,omitempty"`
Sharing *MaestroSharing `json:"sharing,omitempty"`
Ownership *MaestroOwnership `json:"ownership,omitempty"`
}

// MaestroAgentError is AI Maestro's StandardErrorResponse carried as an error, so
// the machine-readable error code survives to the caller instead of being
// flattened into a message. Codes worth acting on include duplicatedHandle,
// BAD_REQUEST_UNSUPPORTED_TOOL, BAD_REQUEST_CREATOR_IN_OWNERSHIP_USERS and
// agentInvalid.
type MaestroAgentError struct {
StatusCode int `json:"statusCode"`
ErrorCode string `json:"errorCode"`
TitleMessage string `json:"titleMessage,omitempty"`
}

func (e *MaestroAgentError) Error() string {
message := fmt.Sprintf("HTTP %d [%s]", e.StatusCode, e.ErrorCode)
if e.TitleMessage != "" {
message += ": " + e.TitleMessage
}
return message
}

// GetAgent reads a Maestro Agent. The endpoint answers 404 both for an unknown
// agent and for one the caller cannot see; seeing an agent is enough to read it,
// being allowed to edit it is not required.
func GetAgent(ctx context.Context, collibraHttpClient *http.Client, agentID string) (*MaestroAgent, error) {
slog.InfoContext(ctx, fmt.Sprintf("Retrieving Maestro agent ID: %s", agentID))

endpoint := fmt.Sprintf("%s/agents/%s", aiMaestroAPIBasePath, url.PathEscape(agentID))

req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
if err != nil {
return nil, fmt.Errorf("get agent: building request: %w", err)
}
req.Header.Set("Accept", "application/json")

return sendMaestroAgentRequest(collibraHttpClient, req, "get agent")
}

// CreateAgent creates a Maestro Agent. The agent is always created in DRAFT with
// the connecting user as its creator and owner, and has to be submitted and
// approved in AI Maestro before end users can use it.
func CreateAgent(ctx context.Context, collibraHttpClient *http.Client, request CreateAgentRequest) (*MaestroAgent, error) {
slog.InfoContext(ctx, fmt.Sprintf("Creating Maestro agent with handle: %s", request.Handle))

body, err := json.Marshal(request)
if err != nil {
return nil, fmt.Errorf("create agent: marshaling request: %w", err)
}

req, err := http.NewRequestWithContext(ctx, http.MethodPost, aiMaestroAPIBasePath+"/agents", bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("create agent: building request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")

return sendMaestroAgentRequest(collibraHttpClient, req, "create agent")
}

// UpdateAgent applies a partial update to a Maestro Agent (PATCH) and returns the
// agent as stored afterwards. Fields left out of the request keep their stored
// value; because the request carries no status, the server demotes the agent to
// DRAFT, so an edited agent has to be submitted and approved again.
func UpdateAgent(ctx context.Context, collibraHttpClient *http.Client, agentID string, request PatchAgentRequest) (*MaestroAgent, error) {
slog.InfoContext(ctx, fmt.Sprintf("Updating Maestro agent ID: %s", agentID))

body, err := json.Marshal(request)
if err != nil {
return nil, fmt.Errorf("update agent: marshaling request: %w", err)
}

endpoint := fmt.Sprintf("%s/agents/%s", aiMaestroAPIBasePath, url.PathEscape(agentID))

req, err := http.NewRequestWithContext(ctx, http.MethodPatch, endpoint, bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("update agent: building request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")

return sendMaestroAgentRequest(collibraHttpClient, req, "update agent")
}

// sendMaestroAgentRequest runs a request that all three agent endpoints answer the
// same way β€” with the agent as stored β€” and decodes it. operation names the call
// in the errors it returns.
func sendMaestroAgentRequest(client *http.Client, req *http.Request, operation string) (*MaestroAgent, error) {
responseBody, err := executeMaestroAgentRequest(client, req)
if err != nil {
return nil, err
}

var agent MaestroAgent
if err := json.Unmarshal(responseBody, &agent); err != nil {
return nil, fmt.Errorf("%s: decoding response: %w", operation, err)
}
return &agent, nil
}

// executeMaestroAgentRequest returns the success body. Unlike executeRequest it
// does not flatten a failure into a message: a parseable error body becomes a
// *MaestroAgentError so the error code reaches the caller, which is what tells it
// whether the failure is worth recovering from.
func executeMaestroAgentRequest(client *http.Client, req *http.Request) ([]byte, error) {
response, err := client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to make request: %w", err)
}
defer func(body io.ReadCloser) {
_ = body.Close()
}(response.Body)

responseBody, err := io.ReadAll(response.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response body: %w", err)
}

if response.StatusCode < 200 || response.StatusCode >= 300 {
return nil, maestroAgentFailure(response.StatusCode, responseBody)
}

return responseBody, nil
}

// maestroAgentFailure reads the error envelope out of a non-2xx body. Bodies that
// are not the documented envelope β€” a gateway error page, say β€” fall back to a
// plain error carrying the raw body.
func maestroAgentFailure(statusCode int, responseBody []byte) error {
var agentErr MaestroAgentError
if err := json.Unmarshal(responseBody, &agentErr); err == nil && agentErr.ErrorCode != "" {
// The body reports its own statusCode, but the transport's is authoritative.
agentErr.StatusCode = statusCode
return &agentErr
}
return fmt.Errorf("HTTP %d: %s", statusCode, string(responseBody))
}
Loading
Loading