bifrost-prompt-rewriter/prompt-rewriter.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/maximhq/bifrost/core/schemas"
)

// Config that can be set in Bifrost's config.json
type PluginConfig struct {
	SmallModelEndpoint string `json:"small_model_endpoint"` // e.g. "http://localhost:8081/v1/chat/completions"
	SmallModelName     string `json:"small_model_name"`     // e.g. "qwen3-4b-instruct"
	RewriteSystemPrompt string `json:"rewrite_system_prompt"`
	EnableLogging      bool   `json:"enable_logging"`
	TimeoutSeconds     int    `json:"timeout_seconds"`
}

var config PluginConfig

func Init(cfg any) error {
	// Default values
	config = PluginConfig{
		SmallModelEndpoint:  "https://bifrost.home.timandjenni.com/openai/v1/chat/completions", // change to your llama.cpp / ollama / lemonade endpoint
		SmallModelName:      "qwen3-4b-instruct",
		RewriteSystemPrompt: `You are an expert prompt engineer. Rewrite the user's raw request into a clear, well-structured, high-quality prompt that will produce better results from a larger coding/reasoning model. 
Keep the original intent. Make it more precise, add useful structure if helpful, and remove ambiguity. 
Return ONLY the improved prompt — no explanations.`,
		EnableLogging:  true,
		TimeoutSeconds: 15,
	}

	// Override with config from Bifrost if provided
	if cfgMap, ok := cfg.(map[string]any); ok {
		if v, ok := cfgMap["small_model_endpoint"].(string); ok {
			config.SmallModelEndpoint = v
		}
		if v, ok := cfgMap["small_model_name"].(string); ok {
			config.SmallModelName = v
		}
		if v, ok := cfgMap["rewrite_system_prompt"].(string); ok {
			config.RewriteSystemPrompt = v
		}
		if v, ok := cfgMap["enable_logging"].(bool); ok {
			config.EnableLogging = v
		}
	}
	return nil
}

func GetName() string {
	return "prompt-rewriter"
}

func PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) {
	// Only act on chat requests that have messages
	if req.ChatRequest == nil || len(req.ChatRequest.Input) == 0 {
		return req, nil, nil
	}

	// Find the last user message
	var lastUserMsg *schemas.ChatMessage
	var lastUserIdx int
	for i := len(req.ChatRequest.Input) - 1; i >= 0; i-- {
		if req.ChatRequest.Input[i].Role == "user" {
			lastUserMsg = &req.ChatRequest.Input[i]
			lastUserIdx = i
			break
		}
	}
	if lastUserMsg == nil || lastUserMsg.Content == nil || lastUserMsg.Content.ContentStr == nil {
		return req, nil, nil
	}

	originalPrompt := *lastUserMsg.Content.ContentStr
	if strings.TrimSpace(originalPrompt) == "" {
		return req, nil, nil
	}

	if config.EnableLogging {
		ctx.Log(schemas.LogLevelInfo, fmt.Sprintf("[prompt-rewriter] Original prompt: %s", truncate(originalPrompt, 120)))
	}

	// Call the small model to rewrite the prompt
	improved, err := rewritePrompt(originalPrompt)
	if err != nil {
		ctx.Log(schemas.LogLevelWarn, fmt.Sprintf("[prompt-rewriter] Rewrite failed, using original: %v", err))
		return req, nil, nil // fail open — keep original prompt
	}

	if config.EnableLogging {
		ctx.Log(schemas.LogLevelInfo, fmt.Sprintf("[prompt-rewriter] Improved prompt: %s", truncate(improved, 120)))
	}

	// Replace the user message content with the improved version
	req.ChatRequest.Input[lastUserIdx].Content = &schemas.ChatMessageContent{
		ContentStr: &improved,
	}

	return req, nil, nil
}

func rewritePrompt(original string) (string, error) {
	payload := map[string]any{
		"model": config.SmallModelName,
		"messages": []map[string]string{
			{"role": "system", "content": config.RewriteSystemPrompt},
			{"role": "user", "content": original},
		},
		"temperature": 0.3,
		"max_tokens":  1024,
	}

	body, _ := json.Marshal(payload)
	client := &http.Client{Timeout: time.Duration(config.TimeoutSeconds) * time.Second}

	resp, err := client.Post(config.SmallModelEndpoint, "application/json", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		b, _ := io.ReadAll(resp.Body)
		return "", fmt.Errorf("small model returned %d: %s", resp.StatusCode, string(b))
	}

	var result struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", err
	}
	if len(result.Choices) == 0 {
		return "", fmt.Errorf("no choices returned")
	}

	return strings.TrimSpace(result.Choices[0].Message.Content), nil
}

func truncate(s string, max int) string {
	if len(s) <= max {
		return s
	}
	return s[:max] + "..."
}

// Required stubs
func PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { return nil }
func PostLLMHook(_ *schemas.BifrostContext, resp *schemas.BifrostResponse, err *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError, error) {
	return resp, err, nil
}
func Cleanup() error { return nil }