aibridge

package module
v0.3.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Dec 5, 2025 License: AGPL-3.0 Imports: 45 Imported by: 2

README

aibridge

aibridge is an HTTP gateway that sits between AI clients and upstream AI providers (Anthropic, OpenAI). It intercepts requests to record token usage, prompts, and tool invocations per user. Optionally supports centralized MCP tool injection with allowlist/denylist filtering.

Architecture

┌─────────────────┐     ┌───────────────────────────────────────────┐
│    AI Client    │     │                    aibridge               │
│  (Claude Code,  │────▶│  ┌─────────────────┐    ┌─────────────┐   │
│   Cursor, etc.) │     │  │  RequestBridge  │───▶│  Providers  │   │
└─────────────────┘     │  │  (http.Handler) │    │  (Anthropic │   │
                        │  └─────────────────┘    │   OpenAI)   │   │
                        │                         └──────┬──────┘   │
                        │                                │          │
                        │                                ▼          │    ┌─────────────┐
                        │  ┌─────────────────┐    ┌─────────────┐   │    │  Upstream   │
                        │  │    Recorder     │◀───│ Interceptor │─── ───▶│    API      │
                        │  │ (tokens, tools, │    │ (streaming/ │   │    │ (Anthropic  │
                        │  │  prompts)       │    │  blocking)  │   │    │   OpenAI)   │
                        │  └────────┬────────┘    └──────┬──────┘   │    └─────────────┘
                        │           │                    │          │
                        │           ▼             ┌──────▼──────┐   │
                        │  ┌ ─ ─ ─ ─ ─ ─ ─ ┐      │  MCP Proxy  │   │
                        │  │    Database   │      │   (tools)   │   │
                        │  └ ─ ─ ─ ─ ─ ─ ─ ┘      └─────────────┘   │
                        └───────────────────────────────────────────┘
Components
  • RequestBridge: The main http.Handler that routes requests to providers
  • Provider: Defines bridged routes (intercepted) and passthrough routes (proxied)
  • Interceptor: Handles request/response processing and streaming
  • Recorder: Interface for capturing usage data (tokens, prompts, tools)
  • MCP Proxy (optional): Connects to MCP servers to list tool, inject them into requests, and invoke them in an inner agentic loop

Request Flow

  1. Client sends request to /anthropic/v1/messages or /openai/v1/chat/completions
  2. Actor extraction: Request must have an actor in context (via AsActor()).
  3. Upstream call: Request forwarded to the AI provider
  4. Response relay: Response streamed/sent to client
  5. Recording: Token usage, prompts, and tool invocations recorded

With MCP enabled: Tools from configured MCP servers are centrally defined and injected into requests (prefixed bmcp_). Allowlist/denylist regex patterns control which tools are available. When the model selects an injected tool, the gateway invokes it in an inner agentic loop, and continues the conversation loop until complete.

Passthrough routes (/v1/models, /v1/messages/count_tokens) are reverse-proxied directly.

Observability

Prometheus Metrics

Create metrics with NewMetrics(prometheus.Registerer):

Metric Type Description
interceptions_total Counter Intercepted request count
interceptions_inflight Gauge Currently processing requests
interceptions_duration_seconds Histogram Request duration
tokens_total Counter Token usage (input/output)
prompts_total Counter User prompt count
injected_tool_invocations_total Counter MCP tool invocations
passthrough_total Counter Non-intercepted requests
Recorder Interface

Implement Recorder to persist usage data to your database. The example uses SQLite (example/recorder.go):

  • aibridge_interceptions - request metadata (provider, model, initiator, timestamps)
  • aibridge_token_usages - input/output token counts per response
  • aibridge_user_prompts - user prompts
  • aibridge_tool_usages - tool invocations (injected and client-defined)
type Recorder interface {
    RecordInterception(ctx context.Context, req *InterceptionRecord) error
    RecordInterceptionEnded(ctx context.Context, req *InterceptionRecordEnded) error
    RecordTokenUsage(ctx context.Context, req *TokenUsageRecord) error
    RecordPromptUsage(ctx context.Context, req *PromptUsageRecord) error
    RecordToolUsage(ctx context.Context, req *ToolUsageRecord) error
}

Example

See example/ for a complete runnable example with SQLite persistence and DeepWiki MCP integration.

Setup

  1. Get API keys from the provider consoles:

  2. Set environment variables:

    export ANTHROPIC_API_KEY="sk-ant-..."
export OPENAI_API_KEY="sk-..."

  3. Run the example:

    cd example && go run .

  4. Test with curl:

    curl -X POST http://localhost:8080/anthropic/v1/messages \
  -H "Content-Type: application/json" \
  -d '{
    "model": "claude-sonnet-4-20250514",
    "max_tokens": 1024,
    "messages": [{"role": "user", "content": "Hello!"}],
    "stream": true
  }'

  5. Test with Claude Code: Claude Code allows a base URL override via ANTHROPIC_BASE_URL.

    image with cloude code example

Supported Routes

Provider Route Type
Anthropic /anthropic/v1/messages Bridged (intercepted)
Anthropic /anthropic/v1/models Passthrough
Anthropic /anthropic/v1/messages/count_tokens Passthrough
OpenAI /openai/v1/chat/completions Bridged (intercepted)
OpenAI /openai/v1/models Passthrough

Documentation

Index

Constants

View Source 
const (
	InterceptionCountStatusFailed    = "failed"
	InterceptionCountStatusCompleted = "completed"
)
View Source 
const (
	SSEEventTypeMessage = "message"
	SSEEventTypeError   = "error"
	SSEEventTypePing    = "ping"
)
View Source 
const (
	ProviderAnthropic = "anthropic"
)
View Source 
const (
	ProviderOpenAI = "openai"
)

Variables

View Source 
var ErrEventStreamClosed = errors.New("event stream closed")
View Source 
var UnknownRoute = errors.New("unknown route")

Functions

func AsActor 

func AsActor(ctx context.Context, actorID string, metadata Metadata) context.Context

Types

type AWSBedrockConfig added in v0.1.6 

type AWSBedrockConfig struct {
	Region                     string
	AccessKey, AccessKeySecret string
	Model, SmallFastModel      string
	// EndpointOverride allows overriding the Bedrock endpoint URL for testing.
	// If set, requests will be sent to this URL instead of the default AWS Bedrock endpoint.
	EndpointOverride string
}

type AnthropicConfig added in v0.1.6 

type AnthropicConfig ProviderConfig

type AnthropicErrorResponse 

type AnthropicErrorResponse struct {
	*anthropic.ErrorResponse

	StatusCode int `json:"-"`
}

func (*AnthropicErrorResponse) Error 

func (a *AnthropicErrorResponse) Error() string

type AnthropicMessagesBlockingInterception 

type AnthropicMessagesBlockingInterception struct {
	AnthropicMessagesInterceptionBase
}

func NewAnthropicMessagesBlockingInterception 

func NewAnthropicMessagesBlockingInterception(id uuid.UUID, req *MessageNewParamsWrapper, cfg AnthropicConfig, bedrockCfg *AWSBedrockConfig, tracer trace.Tracer) *AnthropicMessagesBlockingInterception

func (*AnthropicMessagesBlockingInterception) ProcessRequest 

func (i *AnthropicMessagesBlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Request) (outErr error)

func (*AnthropicMessagesBlockingInterception) Setup 

func (i *AnthropicMessagesBlockingInterception) Setup(logger slog.Logger, recorder Recorder, mcpProxy mcp.ServerProxier)

func (*AnthropicMessagesBlockingInterception) Streaming added in v0.2.0 

func (s *AnthropicMessagesBlockingInterception) Streaming() bool

func (*AnthropicMessagesBlockingInterception) TraceAttributes added in v0.3.0 

func (i *AnthropicMessagesBlockingInterception) TraceAttributes(r *http.Request) []attribute.KeyValue

type AnthropicMessagesInterceptionBase 

type AnthropicMessagesInterceptionBase struct {
	// contains filtered or unexported fields
}

func (*AnthropicMessagesInterceptionBase) ID 

func (i *AnthropicMessagesInterceptionBase) ID() uuid.UUID

func (*AnthropicMessagesInterceptionBase) Model 

func (i *AnthropicMessagesInterceptionBase) Model() string

func (*AnthropicMessagesInterceptionBase) Setup 

func (i *AnthropicMessagesInterceptionBase) Setup(logger slog.Logger, recorder Recorder, mcpProxy mcp.ServerProxier)

type AnthropicMessagesStreamingInterception 

type AnthropicMessagesStreamingInterception struct {
	AnthropicMessagesInterceptionBase
}

func NewAnthropicMessagesStreamingInterception 

func NewAnthropicMessagesStreamingInterception(id uuid.UUID, req *MessageNewParamsWrapper, cfg AnthropicConfig, bedrockCfg *AWSBedrockConfig, tracer trace.Tracer) *AnthropicMessagesStreamingInterception

func (*AnthropicMessagesStreamingInterception) ProcessRequest 

func (i *AnthropicMessagesStreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Request) (outErr error)

ProcessRequest handles a request to /v1/messages. This API has a state-machine behind it, which is described in https://docs.claude.com/en/docs/build-with-claude/streaming#event-types.

Each stream uses the following event flow: - `message_start`: contains a Message object with empty content. - A series of content blocks, each of which have a `content_block_start`, one or more `content_block_delta` events, and a `content_block_stop` event. - Each content block will have an index that corresponds to its index in the final Message content array. - One or more `message_delta` events, indicating top-level changes to the final Message object. - A final `message_stop` event.

It will inject any tools which have been provided by the mcp.ServerProxier.

When a response from the server includes an event indicating that a tool must be invoked, a conditional flow takes place:

a) if the tool is not injected (i.e. defined by the client), relay the event unmodified b) if the tool is injected, it will be invoked by the mcp.ServerProxier in the remote MCP server, and its results relayed to the SERVER. The response from the server will be handled synchronously, and this loop can continue until all injected tool invocations are completed and the response is relayed to the client.

func (*AnthropicMessagesStreamingInterception) Setup 

func (s *AnthropicMessagesStreamingInterception) Setup(logger slog.Logger, recorder Recorder, mcpProxy mcp.ServerProxier)

func (*AnthropicMessagesStreamingInterception) Streaming added in v0.2.0 

func (s *AnthropicMessagesStreamingInterception) Streaming() bool

func (*AnthropicMessagesStreamingInterception) TraceAttributes added in v0.3.0 

func (s *AnthropicMessagesStreamingInterception) TraceAttributes(r *http.Request) []attribute.KeyValue

type AnthropicProvider 

type AnthropicProvider struct {
	// contains filtered or unexported fields
}

AnthropicProvider allows for interactions with the Anthropic API.

func NewAnthropicProvider 

func NewAnthropicProvider(cfg AnthropicConfig, bedrockCfg *AWSBedrockConfig) *AnthropicProvider

func (*AnthropicProvider) AuthHeader 

func (p *AnthropicProvider) AuthHeader() string

func (*AnthropicProvider) BaseURL 

func (p *AnthropicProvider) BaseURL() string

func (*AnthropicProvider) BridgedRoutes 

func (p *AnthropicProvider) BridgedRoutes() []string

func (*AnthropicProvider) CreateInterceptor 

func (p *AnthropicProvider) CreateInterceptor(w http.ResponseWriter, r *http.Request, tracer trace.Tracer) (_ Interceptor, outErr error)

func (*AnthropicProvider) InjectAuthHeader 

func (p *AnthropicProvider) InjectAuthHeader(headers *http.Header)

func (*AnthropicProvider) Name 

func (p *AnthropicProvider) Name() string

func (*AnthropicProvider) PassthroughRoutes 

func (p *AnthropicProvider) PassthroughRoutes() []string

type AsyncRecorder 

type AsyncRecorder struct {
	// contains filtered or unexported fields
}

AsyncRecorder calls Recorder methods asynchronously and logs any errors which may occur.

func NewAsyncRecorder 

func NewAsyncRecorder(logger slog.Logger, wrapped Recorder, timeout time.Duration) *AsyncRecorder

func (*AsyncRecorder) RecordInterception 

func (a *AsyncRecorder) RecordInterception(ctx context.Context, req *InterceptionRecord) error

RecordInterception must NOT be called asynchronously. If an interception cannot be recorded, the whole request should fail.

func (*AsyncRecorder) RecordInterceptionEnded added in v0.1.5 

func (a *AsyncRecorder) RecordInterceptionEnded(ctx context.Context, req *InterceptionRecordEnded) error

func (*AsyncRecorder) RecordPromptUsage 

func (a *AsyncRecorder) RecordPromptUsage(ctx context.Context, req *PromptUsageRecord) error

func (*AsyncRecorder) RecordTokenUsage 

func (a *AsyncRecorder) RecordTokenUsage(ctx context.Context, req *TokenUsageRecord) error

func (*AsyncRecorder) RecordToolUsage 

func (a *AsyncRecorder) RecordToolUsage(ctx context.Context, req *ToolUsageRecord) error

func (*AsyncRecorder) Wait 

func (a *AsyncRecorder) Wait()

func (*AsyncRecorder) WithInitiatorID added in v0.2.0 

func (a *AsyncRecorder) WithInitiatorID(initiatorID string)

func (*AsyncRecorder) WithMetrics added in v0.2.0 

func (a *AsyncRecorder) WithMetrics(metrics *Metrics)

func (*AsyncRecorder) WithModel added in v0.2.0 

func (a *AsyncRecorder) WithModel(model string)

func (*AsyncRecorder) WithProvider added in v0.2.0 

func (a *AsyncRecorder) WithProvider(provider string)

type ChatCompletionNewParamsWrapper 

type ChatCompletionNewParamsWrapper struct {
	openai.ChatCompletionNewParams `json:""`
	Stream                         bool `json:"stream,omitempty"`
}

ChatCompletionNewParamsWrapper exists because the "stream" param is not included in openai.ChatCompletionNewParams.

func (*ChatCompletionNewParamsWrapper) LastUserPrompt 

func (c *ChatCompletionNewParamsWrapper) LastUserPrompt() (*string, error)

func (ChatCompletionNewParamsWrapper) MarshalJSON 

func (c ChatCompletionNewParamsWrapper) MarshalJSON() ([]byte, error)

func (*ChatCompletionNewParamsWrapper) UnmarshalJSON 

func (c *ChatCompletionNewParamsWrapper) UnmarshalJSON(raw []byte) error

type Config 

type Config struct {
	OpenAI    ProviderConfig
	Anthropic ProviderConfig
	Bedrock   AWSBedrockConfig
}

type InterceptionRecord 

type InterceptionRecord struct {
	ID                           string
	InitiatorID, Provider, Model string
	Metadata                     Metadata
	StartedAt                    time.Time
}

type InterceptionRecordEnded added in v0.1.5 

type InterceptionRecordEnded struct {
	ID      string
	EndedAt time.Time
}

type Interceptor 

type Interceptor interface {
	// ID returns the unique identifier for this interception.
	ID() uuid.UUID
	// Setup injects some required dependencies. This MUST be called before using the interceptor
	// to process requests.
	Setup(logger slog.Logger, recorder Recorder, mcpProxy mcp.ServerProxier)
	// Model returns the model in use for this [Interceptor].
	Model() string
	// ProcessRequest handles the HTTP request.
	ProcessRequest(w http.ResponseWriter, r *http.Request) error
	// Specifies whether an interceptor handles streaming or not.
	Streaming() bool
	// TraceAttributes returns tracing attributes for this [Interceptor]
	TraceAttributes(*http.Request) []attribute.KeyValue
}

Interceptor describes a (potentially) stateful interaction with an AI provider.

type MessageNewParamsWrapper 

type MessageNewParamsWrapper struct {
	anthropic.MessageNewParams `json:""`
	Stream                     bool `json:"stream,omitempty"`
}

MessageNewParamsWrapper exists because the "stream" param is not included in anthropic.MessageNewParams.

func (*MessageNewParamsWrapper) LastUserPrompt 

func (b *MessageNewParamsWrapper) LastUserPrompt() (*string, error)

func (MessageNewParamsWrapper) MarshalJSON 

func (b MessageNewParamsWrapper) MarshalJSON() ([]byte, error)

func (*MessageNewParamsWrapper) UnmarshalJSON 

func (b *MessageNewParamsWrapper) UnmarshalJSON(raw []byte) error

func (*MessageNewParamsWrapper) UseStreaming 

func (b *MessageNewParamsWrapper) UseStreaming() bool

type Metadata 

type Metadata map[string]any

type Metrics added in v0.2.0 

type Metrics struct {
	// Interception-related metrics.
	InterceptionDuration  *prometheus.HistogramVec
	InterceptionCount     *prometheus.CounterVec
	InterceptionsInflight *prometheus.GaugeVec
	PassthroughCount      *prometheus.CounterVec

	// Prompt-related metrics.
	PromptCount *prometheus.CounterVec

	// Token-related metrics.
	TokenUseCount *prometheus.CounterVec

	// Tool-related metrics.
	InjectedToolUseCount    *prometheus.CounterVec
	NonInjectedToolUseCount *prometheus.CounterVec
}

func NewMetrics added in v0.2.0 

func NewMetrics(reg prometheus.Registerer) *Metrics

NewMetrics creates AND registers metrics. It will panic if a collector has already been registered. Note: we are not specifying namespace in the metrics; the provided registerer may specify a "namespace" using prometheus.WrapRegistererWithPrefix.

type OpenAIBlockingChatInterception 

type OpenAIBlockingChatInterception struct {
	OpenAIChatInterceptionBase
}

func NewOpenAIBlockingChatInterception 

func NewOpenAIBlockingChatInterception(id uuid.UUID, req *ChatCompletionNewParamsWrapper, baseURL, key string, tracer trace.Tracer) *OpenAIBlockingChatInterception

func (*OpenAIBlockingChatInterception) ProcessRequest 

func (i *OpenAIBlockingChatInterception) ProcessRequest(w http.ResponseWriter, r *http.Request) (outErr error)

func (*OpenAIBlockingChatInterception) Setup 

func (s *OpenAIBlockingChatInterception) Setup(logger slog.Logger, recorder Recorder, mcpProxy mcp.ServerProxier)

func (*OpenAIBlockingChatInterception) Streaming added in v0.2.0 

func (s *OpenAIBlockingChatInterception) Streaming() bool

func (*OpenAIBlockingChatInterception) TraceAttributes added in v0.3.0 

func (s *OpenAIBlockingChatInterception) TraceAttributes(r *http.Request) []attribute.KeyValue

type OpenAIChatInterceptionBase 

type OpenAIChatInterceptionBase struct {
	// contains filtered or unexported fields
}

func (*OpenAIChatInterceptionBase) ID 

func (i *OpenAIChatInterceptionBase) ID() uuid.UUID

func (*OpenAIChatInterceptionBase) Model 

func (i *OpenAIChatInterceptionBase) Model() string

func (*OpenAIChatInterceptionBase) Setup 

func (i *OpenAIChatInterceptionBase) Setup(logger slog.Logger, recorder Recorder, mcpProxy mcp.ServerProxier)

type OpenAIConfig added in v0.1.6 

type OpenAIConfig ProviderConfig

type OpenAIErrorResponse 

type OpenAIErrorResponse struct {
	ErrorObject *shared.ErrorObject `json:"error"`
	StatusCode  int                 `json:"-"`
}

func (*OpenAIErrorResponse) Error 

func (a *OpenAIErrorResponse) Error() string

type OpenAIProvider 

type OpenAIProvider struct {
	// contains filtered or unexported fields
}

OpenAIProvider allows for interactions with the OpenAI API.

func NewOpenAIProvider 

func NewOpenAIProvider(cfg OpenAIConfig) *OpenAIProvider

func (*OpenAIProvider) AuthHeader 

func (p *OpenAIProvider) AuthHeader() string

func (*OpenAIProvider) BaseURL 

func (p *OpenAIProvider) BaseURL() string

func (*OpenAIProvider) BridgedRoutes 

func (p *OpenAIProvider) BridgedRoutes() []string

func (*OpenAIProvider) CreateInterceptor 

func (p *OpenAIProvider) CreateInterceptor(w http.ResponseWriter, r *http.Request, tracer trace.Tracer) (_ Interceptor, outErr error)

func (*OpenAIProvider) InjectAuthHeader 

func (p *OpenAIProvider) InjectAuthHeader(headers *http.Header)

func (*OpenAIProvider) Name 

func (p *OpenAIProvider) Name() string

func (*OpenAIProvider) PassthroughRoutes 

func (p *OpenAIProvider) PassthroughRoutes() []string

PassthroughRoutes define the routes which are not currently intercepted but must be passed through to the upstream. The /v1/completions legacy API is deprecated and will not be passed through. See https://platform.openai.com/docs/api-reference/completions.

type OpenAIStreamingChatInterception 

type OpenAIStreamingChatInterception struct {
	OpenAIChatInterceptionBase
}

func NewOpenAIStreamingChatInterception 

func NewOpenAIStreamingChatInterception(id uuid.UUID, req *ChatCompletionNewParamsWrapper, baseURL, key string, tracer trace.Tracer) *OpenAIStreamingChatInterception

func (*OpenAIStreamingChatInterception) ProcessRequest 

func (i *OpenAIStreamingChatInterception) ProcessRequest(w http.ResponseWriter, r *http.Request) (outErr error)

ProcessRequest handles a request to /v1/chat/completions. See https://platform.openai.com/docs/api-reference/chat-streaming/streaming.

It will inject any tools which have been provided by the mcp.ServerProxier.

When a response from the server includes an event indicating that a tool must be invoked, a conditional flow takes place:

a) if the tool is not injected (i.e. defined by the client), relay the event unmodified b) if the tool is injected, it will be invoked by the mcp.ServerProxier in the remote MCP server, and its results relayed to the SERVER. The response from the server will be handled synchronously, and this loop can continue until all injected tool invocations are completed and the response is relayed to the client.

func (*OpenAIStreamingChatInterception) Setup 

func (i *OpenAIStreamingChatInterception) Setup(logger slog.Logger, recorder Recorder, mcpProxy mcp.ServerProxier)

func (*OpenAIStreamingChatInterception) Streaming added in v0.2.0 

func (i *OpenAIStreamingChatInterception) Streaming() bool

func (*OpenAIStreamingChatInterception) TraceAttributes added in v0.3.0 

func (s *OpenAIStreamingChatInterception) TraceAttributes(r *http.Request) []attribute.KeyValue

type PromptUsageRecord 

type PromptUsageRecord struct {
	InterceptionID string
	MsgID, Prompt  string
	Metadata       Metadata
	CreatedAt      time.Time
}

type Provider 

type Provider interface {
	// Name returns the provider's name.
	Name() string
	// BaseURL defines the base URL endpoint for this provider's API.
	BaseURL() string

	// CreateInterceptor starts a new [Interceptor] which is responsible for intercepting requests,
	// communicating with the upstream provider and formulating a response to be sent to the requesting client.
	CreateInterceptor(http.ResponseWriter, *http.Request, trace.Tracer) (Interceptor, error)

	// BridgedRoutes returns a slice of [http.ServeMux]-compatible routes which will have special handling.
	// See https://pkg.go.dev/net/http#hdr-Patterns-ServeMux.
	BridgedRoutes() []string
	// PassthroughRoutes returns a slice of whitelisted [http.ServeMux]-compatible* routes which are
	// not currently intercepted and must be handled by the upstream directly.
	//
	// * only path routes can be specified, not ones containing HTTP methods. (i.e. GET /route).
	// By default, these passthrough routes will accept any HTTP method.
	PassthroughRoutes() []string

	// AuthHeader returns the name of the header which the provider expects to find its authentication
	// token in.
	AuthHeader() string
	// InjectAuthHeader allows [Provider]s to set its authentication header.
	InjectAuthHeader(*http.Header)
}

Provider describes an AI provider client's behaviour. Provider clients are responsible for interacting with upstream AI providers.

type ProviderConfig 

type ProviderConfig struct {
	BaseURL, Key string
}

type Recorder 

type Recorder interface {
	// RecordInterception records metadata about an interception with an upstream AI provider.
	RecordInterception(ctx context.Context, req *InterceptionRecord) error
	// RecordInterceptionEnded records that given interception has completed.
	RecordInterceptionEnded(ctx context.Context, req *InterceptionRecordEnded) error
	// RecordTokenUsage records the tokens used in an interception with an upstream AI provider.
	RecordTokenUsage(ctx context.Context, req *TokenUsageRecord) error
	// RecordPromptUsage records the prompts used in an interception with an upstream AI provider.
	RecordPromptUsage(ctx context.Context, req *PromptUsageRecord) error
	// RecordToolUsage records the tools used in an interception with an upstream AI provider.
	RecordToolUsage(ctx context.Context, req *ToolUsageRecord) error
}

Recorder describes all the possible usage information we need to capture during interactions with AI providers. Additionally, it introduces the concept of an "Interception", which includes information about which provider/model was used and by whom. All usage records should reference this Interception by ID.

type RecorderWrapper 

type RecorderWrapper struct {
	// contains filtered or unexported fields
}

RecorderWrapper is a convenience struct which implements RecorderClient and resolves a client before calling each method. It also sets the start/creation time of each record.

func NewRecorder 

func NewRecorder(logger slog.Logger, tracer trace.Tracer, clientFn func() (Recorder, error)) *RecorderWrapper

func (*RecorderWrapper) RecordInterception 

func (r *RecorderWrapper) RecordInterception(ctx context.Context, req *InterceptionRecord) (outErr error)

func (*RecorderWrapper) RecordInterceptionEnded added in v0.1.5 

func (r *RecorderWrapper) RecordInterceptionEnded(ctx context.Context, req *InterceptionRecordEnded) (outErr error)

func (*RecorderWrapper) RecordPromptUsage 

func (r *RecorderWrapper) RecordPromptUsage(ctx context.Context, req *PromptUsageRecord) (outErr error)

func (*RecorderWrapper) RecordTokenUsage 

func (r *RecorderWrapper) RecordTokenUsage(ctx context.Context, req *TokenUsageRecord) (outErr error)

func (*RecorderWrapper) RecordToolUsage 

func (r *RecorderWrapper) RecordToolUsage(ctx context.Context, req *ToolUsageRecord) (outErr error)

type RequestBridge 

type RequestBridge struct {
	// contains filtered or unexported fields
}

RequestBridge is an http.Handler which is capable of masquerading as AI providers' APIs; specifically, OpenAI's & Anthropic's at present. RequestBridge intercepts requests to - and responses from - these upstream services to provide a centralized governance layer.

RequestBridge has no concept of authentication or authorization. It does have a concept of identity, in the narrow sense that it expects an [actor] to be defined in the context, to record the initiator of each interception.

RequestBridge is safe for concurrent use.

func NewRequestBridge 

func NewRequestBridge(ctx context.Context, providers []Provider, recorder Recorder, mcpProxy mcp.ServerProxier, logger slog.Logger, metrics *Metrics, tracer trace.Tracer) (*RequestBridge, error)

NewRequestBridge creates a new *RequestBridge and registers the HTTP routes defined by the given providers. Any routes which are requested but not registered will be reverse-proxied to the upstream service.

A Recorder is also required to record prompt, tool, and token use.

mcpProxy will be closed when the RequestBridge is closed.

func (*RequestBridge) InflightRequests 

func (b *RequestBridge) InflightRequests() int32

func (*RequestBridge) ServeHTTP 

func (b *RequestBridge) ServeHTTP(rw http.ResponseWriter, r *http.Request)

ServeHTTP exposes the internal http.Handler, which has all [Provider]s' routes registered. It also tracks inflight requests.

func (*RequestBridge) Shutdown 

func (b *RequestBridge) Shutdown(ctx context.Context) error

Shutdown will attempt to gracefully shutdown. This entails waiting for all requests to complete, and shutting down the MCP server proxier. TODO: add tests.

type SSEEvent 

type SSEEvent struct {
	Type  string
	Data  string
	ID    string
	Retry int
}

type SSEParser 

type SSEParser struct {
	// contains filtered or unexported fields
}

func NewSSEParser 

func NewSSEParser() *SSEParser

func (*SSEParser) AllEvents 

func (p *SSEParser) AllEvents() map[string][]SSEEvent

func (*SSEParser) EventsByType 

func (p *SSEParser) EventsByType(eventType string) []SSEEvent

func (*SSEParser) MessageEvents 

func (p *SSEParser) MessageEvents() []SSEEvent

func (*SSEParser) Parse 

func (p *SSEParser) Parse(reader io.Reader) error

type TokenUsageRecord 

type TokenUsageRecord struct {
	InterceptionID string
	MsgID          string
	Input, Output  int64
	// ExtraTokenTypes holds token types which *may* exist over and above input/output.
	// These should ultimately get merged into [Metadata], but it's useful to keep these
	// with their actual type (int64) since [Metadata] is a map[string]any.
	ExtraTokenTypes map[string]int64
	Metadata        Metadata
	CreatedAt       time.Time
}

type ToolArgs 

type ToolArgs any

type ToolUsageRecord 

type ToolUsageRecord struct {
	InterceptionID  string
	MsgID, Tool     string
	ServerURL       *string
	Args            ToolArgs
	Injected        bool
	InvocationError error
	Metadata        Metadata
	CreatedAt       time.Time
}

Source Files

View all Source files

Directories

Path Synopsis
Package mcpmock is a generated GoMock package.
 Package mcpmock is a generated GoMock package.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.