Plugins

type Plugin interface {
    GetName() string
    PreHook(ctx *context.Context, req *BifrostRequest) (*BifrostRequest, *PluginShortCircuit, error)
    PostHook(ctx *context.Context, result *BifrostResponse, err *BifrostError) (*BifrostResponse, *BifrostError, error)
    Cleanup() error
}

type LoggingPlugin struct {
    logger *log.Logger
}

func (p *LoggingPlugin) GetName() string {
    return "logging"
}

func (p *LoggingPlugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.PluginShortCircuit, error) {
    p.logger.Printf("Request: Provider=%s, Model=%s", req.Provider, req.Model)
    return req, nil, nil
}

func (p *LoggingPlugin) PostHook(ctx *context.Context, result *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError, error) {
    if err != nil {
        p.logger.Printf("Error: %s", err.Error.Message)
    } else {
        p.logger.Printf("Success: Provider=%s", result.ExtraFields.Provider)
    }
    return result, err, nil
}

func (p *LoggingPlugin) Cleanup() error {
    return nil
}

type RateLimitPlugin struct {
    requests map[string]int
    mu       sync.Mutex
    limit    int
}

func (p *RateLimitPlugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.PluginShortCircuit, error) {
    userID := p.extractUserID(*ctx)

    p.mu.Lock()
    count := p.requests[userID]
    if count >= p.limit {
        p.mu.Unlock()

        // Rate limit exceeded - short circuit
        return req, &schemas.PluginShortCircuit{
            Error: &schemas.BifrostError{
                StatusCode: &[]int{429}[0],
                Error: schemas.ErrorField{
                    Message: "Rate limit exceeded",
                },
            },
        }, nil
    }

    p.requests[userID] = count + 1
    p.mu.Unlock()

    return req, nil, nil
}

type CachePlugin struct {
    cache map[string]*schemas.BifrostResponse
    mu    sync.RWMutex
}

func (p *CachePlugin) PreHook(ctx *context.Context, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.PluginShortCircuit, error) {
    cacheKey := p.generateCacheKey(req)

    p.mu.RLock()
    cached, exists := p.cache[cacheKey]
    p.mu.RUnlock()

    if exists {
        // Return cached response - short circuit
        return req, &schemas.PluginShortCircuit{
            Response: cached,
        }, nil
    }

    return req, nil, nil
}

func (p *CachePlugin) PostHook(ctx *context.Context, result *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError, error) {
    if result != nil {
        cacheKey := p.generateCacheKeyFromResponse(result)

        p.mu.Lock()
        p.cache[cacheKey] = result
        p.mu.Unlock()
    }

    return result, err, nil
}

// Add plugins to your Bifrost client
client, initErr := bifrost.Init(schemas.BifrostConfig{
    Account: &MyAccount{},
    Plugins: []schemas.Plugin{
        NewLoggingPlugin(),
        NewRateLimitPlugin(100), // 100 requests per user
        NewCachePlugin(time.Hour), // 1 hour cache
    },
})
defer client.Cleanup() // Calls Cleanup() on all plugins