Assistant/cmd/cognitive-worker/main.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"git.cin.su/m/Assistant/adapters/models"
	"git.cin.su/m/Assistant/contracts/cognitive"
	"git.cin.su/m/Assistant/runtime/health"
	"io"
	"log"
	"net/http"
	"os"
)

const maxOpenAISecretBytes = 8 << 10

type advisoryProvider interface {
	Advise(context.Context, cognitive.AdvisoryRequest) (cognitive.AdvisoryResponse, error)
}

func main() {
	addr := os.Getenv("LISTEN_ADDR")
	if addr == "" {
		addr = ":8083"
	}
	openAIProvider, err := loadOpenAIAdvisoryProvider(os.Getenv("OPENAI_API_KEY_FILE"))
	if err != nil {
		log.Fatal(err)
	}
	mux := http.NewServeMux()
	mux.Handle("/healthz", health.Handler{Runtime: "cognitive-worker", SemanticGeneration: 1, ArchitectureGeneration: 1})
	mux.HandleFunc("POST /v1/proposals/deterministic", deterministicProposal)
	mux.HandleFunc("POST /v1/advisories/deterministic", deterministicAdvisory)
	mux.HandleFunc("POST /v1/advisories/openai", openAIAdvisory(openAIProvider))
	log.Printf("cognitive-worker listening on %s", addr)
	if err := http.ListenAndServe(addr, mux); err != nil {
		log.Fatal(err)
	}
}

func deterministicAdvisory(w http.ResponseWriter, r *http.Request) {
	request, ok := decodeAdvisoryRequest(w, r)
	if !ok {
		return
	}
	response, err := models.DeterministicAdvisoryUnavailable(request)
	if err != nil {
		http.Error(w, "invalid advisory request", http.StatusBadRequest)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(response)
}

func openAIAdvisory(provider advisoryProvider) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if provider == nil {
			http.Error(w, "provider unavailable", http.StatusServiceUnavailable)
			return
		}
		request, ok := decodeAdvisoryRequest(w, r)
		if !ok {
			return
		}
		response, err := provider.Advise(r.Context(), request)
		if err != nil {
			http.Error(w, "invalid advisory request", http.StatusBadRequest)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(response)
	}
}

func decodeAdvisoryRequest(w http.ResponseWriter, r *http.Request) (cognitive.AdvisoryRequest, bool) {
	r.Body = http.MaxBytesReader(w, r.Body, 80*1024)
	decoder := json.NewDecoder(r.Body)
	decoder.DisallowUnknownFields()
	var request cognitive.AdvisoryRequest
	if err := decoder.Decode(&request); err != nil {
		http.Error(w, "invalid advisory request", http.StatusBadRequest)
		return cognitive.AdvisoryRequest{}, false
	}
	if err := decoder.Decode(&struct{}{}); err != io.EOF {
		http.Error(w, "invalid advisory request", http.StatusBadRequest)
		return cognitive.AdvisoryRequest{}, false
	}
	return request, true
}

func loadOpenAIAdvisoryProvider(secretPath string) (advisoryProvider, error) {
	if os.Getenv("OPENAI_API_KEY") != "" {
		return nil, errors.New("direct OpenAI credential environment is forbidden; use OPENAI_API_KEY_FILE")
	}
	if secretPath == "" {
		return nil, nil
	}
	file, err := os.Open(secretPath)
	if err != nil {
		return nil, fmt.Errorf("open OpenAI credential file: %w", err)
	}
	defer file.Close()
	raw, err := io.ReadAll(io.LimitReader(file, maxOpenAISecretBytes+1))
	if err != nil {
		return nil, fmt.Errorf("read OpenAI credential file: %w", err)
	}
	defer clear(raw)
	if len(raw) > maxOpenAISecretBytes {
		return nil, errors.New("OpenAI credential file is oversized")
	}
	key := bytes.TrimSpace(raw)
	if len(key) == 0 {
		return nil, errors.New("OpenAI credential file is empty")
	}
	for _, character := range key {
		if character < 0x21 || character > 0x7e {
			return nil, errors.New("OpenAI credential file contains invalid characters")
		}
	}
	provider, err := models.NewOpenAIAdvisoryAdapter(key, nil)
	if err != nil {
		return nil, errors.New("OpenAI credential file is invalid")
	}
	return provider, nil
}

func deterministicProposal(w http.ResponseWriter, r *http.Request) {
	var request models.DeterministicRequest
	if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
		http.Error(w, "invalid deterministic proposal request", http.StatusBadRequest)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(models.Propose(request))
}