Go code example

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"

	openai "github.com/openai/openai-go"
)

func main() {
	// Get your Trend Vision One API key from environment variable
	apiKey := os.Getenv("V1_API_KEY")
	if apiKey == "" {
		panic("Missing V1_API_KEY environment variable")
	}

	// User prompt stored in a variable
	userPrompt := "Explain the concept of machine learning in simple terms."

	// Prepare the request payload
	payload := map[string]string{
		"guard": userPrompt,
	}
	jsonPayload, err := json.Marshal(payload)
	if err != nil {
		panic(err)
	}

	// Prepare the HTTP request
	url := "https://api.<region>.xdr.trendmicro.com/beta/aiSecurity"
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonPayload))
	if err != nil {
		panic(err)
	}
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("detailedResponse", "false") // Optional: Set to true for more detailed responses

	// Send the request
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	// Check the response for the 'action' header
	unsafe := false
	if action := resp.Header.Get("action"); action != "" && action == "block" {
		unsafe = true
	}

	fmt.Printf("Unsafe: %v\n", unsafe)

	if !unsafe {
		openaiApiKey := os.Getenv("OPENAI_API_KEY")
		if openaiApiKey == "" {
			panic("Missing OPENAI_API_KEY environment variable")
		}

		// Use OpenAI Go SDK
		ctx := context.Background()
		openaiClient := openai.NewClient(openaiApiKey)
		resp, err := openaiClient.CreateChatCompletion(ctx, &openai.CreateChatCompletionRequest{
			Model: "gpt-4",
			Messages: []openai.ChatCompletionMessage{{
				Role:    "user",
				Content: userPrompt,
			}},
			MaxTokens:   150,
			Temperature: 0.7,
		})
		if err != nil {
			panic(err)
		}

		// Marshal OpenAI response to send to TrendMicro endpoint
		guardRespPayload, err := json.Marshal(resp)
		if err != nil {
			panic(err)
		}
		guardReq, err := http.NewRequest("POST", url, bytes.NewBuffer(guardRespPayload))
		if err != nil {
			panic(err)
		}
		guardReq.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey))
		guardReq.Header.Set("Content-Type", "application/json")
		guardReq.Header.Set("detailedResponse", "false")

		guardResp, err := client.Do(guardReq)
		if err != nil {
			panic(err)
		}
		defer guardResp.Body.Close()

		guardAction := guardResp.Header.Get("action")
		if guardAction != "" && guardAction == "block" {
			fmt.Println("LLM response is considered unsafe. No response will be shown.")
			os.Exit(0)
		}

		// Print the LLM response (extracting the text)
		if len(resp.Choices) > 0 {
			fmt.Println(resp.Choices[0].Message.Content)
		}
	} else {
		fmt.Println("User prompt is considered unsafe. No response will be generated.")
		os.Exit(0)
	}
}