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/capability"
	"git.cin.su/m/Assistant/contracts/cognitive"
	"git.cin.su/m/Assistant/contracts/identity"
	"git.cin.su/m/Assistant/runtime/health"
	"io"
	"log"
	"net/http"
	"os"
	"time"
)

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()
	resolver, err := newHTTPBindingResolver(envOr("CORE_RUNTIME_URL", "http://core-runtime:8080"), &http.Client{Timeout: 12 * time.Second})
	if err != nil {
		log.Fatal(err)
	}
	providers := map[capability.ProviderKind][]providerRegistration{}
	discoveries := map[string]providerRegistration{}
	cpuAdapter, err := models.NewOpenAICompatibleLocalAdapter(envOr("LOCAL_CPU_PROVIDER_URL", "http://docker.cin.su:18082/v1"), nil)
	if err != nil {
		log.Fatal(err)
	}
	cpu := providerRegistration{ProviderInstanceID: "local-cpu-primary", Kind: "openai-compatible-chat-v1", EndpointRef: "local-cpu-primary-v1", Provider: cpuAdapter, Discovery: cpuAdapter}
	providers[cpu.Kind] = append(providers[cpu.Kind], cpu)
	discoveries["cpu"] = cpu
	gpuAdapter, err := models.NewOllamaAdapter(envOr("LOCAL_GPU_PROVIDER_URL", "http://192.168.220.91:11434"), nil)
	if err != nil {
		log.Fatal(err)
	}
	gpu := providerRegistration{ProviderInstanceID: "local-gpu-opportunistic", Kind: "ollama-v1", EndpointRef: "local-gpu-opportunistic-v1", Provider: gpuAdapter, Discovery: gpuAdapter}
	providers[gpu.Kind] = append(providers[gpu.Kind], gpu)
	discoveries["gpu"] = gpu
	if openAIProvider != nil {
		kind := capability.ProviderKind(envOr("OPENAI_PROVIDER_KIND", "openai-responses-v1"))
		providers[kind] = append(providers[kind], providerRegistration{ProviderInstanceID: "openai-bootstrap", Kind: kind, EndpointRef: capability.EndpointBindingRef(envOr("OPENAI_ENDPOINT_REF", "openai-responses-default")), CredentialRef: capability.CredentialRef(envOr("OPENAI_CREDENTIAL_REF", "openai-advisory-v1")), Provider: openAIProvider})
	}
	gateway := cognitiveGateway{resolver: resolver, providers: providers, now: time.Now, allowCanary: os.Getenv("COGNITIVE_ALLOW_CANARY") == "true", canaryMaxOutputBytes: 4096}
	mux.Handle("/healthz", health.Handler{Runtime: "cognitive-worker", SemanticGeneration: 1, ArchitectureGeneration: 1})
	mux.HandleFunc("POST /v1/proposals/deterministic", deterministicProposal)
	mux.HandleFunc("POST /v1/advisory-policy", gateway.policyHandler)
	mux.HandleFunc("POST /v1/advisories", gateway.advisoryHandler)
	discovery := discoveryService{resolver: resolver, providers: discoveries, owner: identity.PrincipalID(envOr("COGNITIVE_OWNER_PRINCIPAL_ID", "domovoy-creator-principal")), now: time.Now}
	mux.HandleFunc("POST /v1/provider-discovery/{provider}", discovery.handler)
	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))
}

func envOr(name, fallback string) string {
	if value := os.Getenv(name); value != "" {
		return value
	}
	return fallback
}