feat(lumo): add a --web-search flag to serve

internal/cli/lumo/handler.go

@@ -13,9 +13,20 @@ import (
 	"github.com/major0/proton-utils/api/lumo"
 )
 
+// webSearchTools returns the tool list to pass to Generate: Lumo's native
+// web_search when enabled, otherwise none. Kept in one place so the streaming
+// and non-streaming paths request the same tools.
+func webSearchTools(enabled bool) []lumo.ToolName {
+	if enabled {
+		return []lumo.ToolName{lumo.ToolWebSearch}
+	}
+	return nil
+}
+
 // chatHandler returns an http.HandlerFunc that proxies OpenAI-format chat
-// completion requests to Lumo via the provided client.
-func chatHandler(client *lumo.Client) http.HandlerFunc {
+// completion requests to Lumo via the provided client. When webSearch is set,
+// Lumo's native web_search tool is enabled for every request.
+func chatHandler(client *lumo.Client, webSearch bool) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		var req lumo.ChatCompletionRequest
 		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
@@ -28,15 +39,15 @@ func chatHandler(client *lumo.Client) http.HandlerFunc {
 		model := "lumo"
 
 		if req.Stream {
-			streamResponse(r.Context(), w, client, turns, id, model)
+			streamResponse(r.Context(), w, client, turns, id, model, webSearch)
 		} else {
-			nonStreamResponse(r.Context(), w, client, turns, id, model)
+			nonStreamResponse(r.Context(), w, client, turns, id, model, webSearch)
 		}
 	}
 }
 
 // streamResponse handles streaming (SSE) chat completions.
-func streamResponse(ctx context.Context, w http.ResponseWriter, client *lumo.Client, turns []lumo.Turn, id, model string) {
+func streamResponse(ctx context.Context, w http.ResponseWriter, client *lumo.Client, turns []lumo.Turn, id, model string, webSearch bool) {
 	flusher, ok := w.(http.Flusher)
 	if !ok {
 		writeError(w, http.StatusInternalServerError, "server_error", "Streaming not supported")
@@ -49,6 +60,7 @@ func streamResponse(ctx context.Context, w http.ResponseWriter, client *lumo.Cli
 	w.WriteHeader(http.StatusOK)
 
 	err := client.Generate(ctx, turns, lumo.GenerateOpts{
+		Tools: webSearchTools(webSearch),
 		ChunkCallback: func(msg lumo.GenerationResponseMessage) {
 			chunk, ok := ChunkToSSEEvent(msg, id, model)
 			if !ok {
@@ -84,10 +96,11 @@ func streamResponse(ctx context.Context, w http.ResponseWriter, client *lumo.Cli
 }
 
 // nonStreamResponse handles non-streaming chat completions.
-func nonStreamResponse(ctx context.Context, w http.ResponseWriter, client *lumo.Client, turns []lumo.Turn, id, model string) {
+func nonStreamResponse(ctx context.Context, w http.ResponseWriter, client *lumo.Client, turns []lumo.Turn, id, model string, webSearch bool) {
 	var content strings.Builder
 
 	err := client.Generate(ctx, turns, lumo.GenerateOpts{
+		Tools: webSearchTools(webSearch),
 		ChunkCallback: func(msg lumo.GenerationResponseMessage) {
 			if msg.Type == "token_data" && msg.Content != "" {
 				content.WriteString(msg.Content)

internal/cli/lumo/serve.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
+	"strings"
 	"syscall"
 	"time"
 
@@ -23,6 +24,7 @@ type serveFlags struct {
 	tlsCert   string
 	tlsKey    string
 	noTLS     bool
+	webSearch bool
 }
 
 var sFlags serveFlags
@@ -46,13 +48,18 @@ func init() {
 	if v := os.Getenv("LUMO_API_KEY"); v != "" {
 		apiKeyDefault = v
 	}
+	webSearchDefault := false
+	if v := os.Getenv("LUMO_WEB_SEARCH"); v == "1" || strings.EqualFold(v, "true") {
+		webSearchDefault = true
+	}
 
 	serveCmd.Flags().StringVar(&sFlags.addr, "addr", addrDefault, "Listen address and port [$LUMO_ENDPOINT]")
 	serveCmd.Flags().StringVar(&sFlags.apiKey, "api-key", apiKeyDefault, "Use this API key (not persisted) [$LUMO_API_KEY]")
 	serveCmd.Flags().BoolVar(&sFlags.newAPIKey, "new-api-key", false, "Generate and persist a new API key")
 	serveCmd.Flags().StringVar(&sFlags.tlsCert, "tls-cert", "", "Custom TLS certificate path")
 	serveCmd.Flags().StringVar(&sFlags.tlsKey, "tls-key", "", "Custom TLS key path")
 	serveCmd.Flags().BoolVar(&sFlags.noTLS, "no-tls", false, "Disable TLS, serve plain HTTP")
+	serveCmd.Flags().BoolVar(&sFlags.webSearch, "web-search", webSearchDefault, "Enable Lumo's native web search tool [$LUMO_WEB_SEARCH]")
 }
 
 func runServe(cmd *cobra.Command, _ []string) error {
@@ -76,7 +83,7 @@ func runServe(cmd *cobra.Command, _ []string) error {
 	}
 
 	mux := http.NewServeMux()
-	mux.HandleFunc("POST /v1/chat/completions", chatHandler(client))
+	mux.HandleFunc("POST /v1/chat/completions", chatHandler(client, sFlags.webSearch))
 	mux.HandleFunc("GET /v1/models", modelsHandler())
 
 	handler := authMiddleware(apiKey, loggingMiddleware(mux))

internal/cli/lumo/websearch_test.go

@@ -0,0 +1,17 @@
+package lumoCmd
+
+import (
+	"testing"
+
+	"github.com/major0/proton-utils/api/lumo"
+)
+
+func TestWebSearchTools(t *testing.T) {
+	if got := webSearchTools(false); got != nil {
+		t.Fatalf("webSearchTools(false) = %v, want nil", got)
+	}
+	got := webSearchTools(true)
+	if len(got) != 1 || got[0] != lumo.ToolWebSearch {
+		t.Fatalf("webSearchTools(true) = %v, want [%q]", got, lumo.ToolWebSearch)
+	}
+}