feat(core): llm integration for cli

.gitignore

@@ -3,6 +3,8 @@ bin/.DS_Store
 target/
 .vscode/launch.json
 otdfctl.yaml
+pkg/chat/log/*
+pkg/chat/log/temp/*
 
 # Ignore the tructl binary
 otdfctl

cmd/chat.go

@@ -0,0 +1,20 @@
+package cmd
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/opentdf/otdfctl/pkg/chat"
+)
+
+func init() {
+	// Load in configs from YAML file - change to `otdfctl.yaml` for production, `otdfctl-example.yaml` for development
+	err := chat.LoadConfig("otdfctl-example.yaml")
+	// err := chat.LoadConfig("otdfctl.yaml")
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error loading chat config: %v\n", err)
+		os.Exit(1)
+	}
+	chat.ConfigureChatCommand()
+	RootCmd.AddCommand(chat.GetChatCommand())
+}

otdfctl-example.yaml

@@ -1,3 +1,11 @@
 output:
   # acceptable formats: json or styled
   format: styled
+# Find full list of attributes for ollama in API documentation: https://github.com/ollama/ollama/blob/main/docs/api.md
+chat:
+  model: "llama3.1:8b-instruct-q5_K_S"
+  apiUrl: "http://localhost:11434/api/generate"
+  logLength: 20
+  verbose: true
+  tokenLimit: 1000
+  numGPU: 1

pkg/chat/config.go

@@ -0,0 +1,94 @@
+package chat
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/spf13/cobra"
+	"gopkg.in/yaml.v2"
+)
+
+// Configs are loaded from a otdfctl.yaml file in home directory where defaults are provided
+type ChatConfig struct {
+	Model      string `yaml:"model" default:"llama3"`
+	ApiURL     string `yaml:"apiUrl" default:"http://localhost:11434/api/generate"`
+	LogLength  int    `yaml:"logLength" default:"100"`
+	Verbose    bool   `yaml:"verbose" default:"true"`
+	TokenLimit int    `yaml:"tokenLimit" default:"1000"`
+	NumGPU     int    `yaml:"numGPU" default:"1"`
+}
+
+type Output struct {
+	Format string `yaml:"format"`
+}
+
+type Config struct {
+	Output Output     `yaml:"output"`
+	Chat   ChatConfig `yaml:"chat"`
+}
+
+var chatConfig Config
+var ask string
+
+func LoadConfig(filename string) error {
+	file, err := os.Open(filename)
+	if err != nil {
+		return fmt.Errorf("could not open config file: %v", err)
+	}
+	defer file.Close()
+
+	decoder := yaml.NewDecoder(file)
+	err = decoder.Decode(&chatConfig)
+	if err != nil {
+		return fmt.Errorf("could not decode config YAML: %v", err)
+	}
+
+	return nil
+}
+
+func ConfigureChatCommand() {
+	chatCmd.PersistentFlags().StringVar(&chatConfig.Chat.Model, "model", chatConfig.Chat.Model, "Model name for Ollama")
+	chatCmd.PersistentFlags().StringVar(&ask, "ask", "", "Ask a one-off question without entering the chat session") // --ask invocation: ./otdfctl chat --ask "[$question_here]"
+}
+
+func RunChatSession(cmd *cobra.Command, args []string) {
+	// Call the Setup function before starting the chat session
+	// Load in configs from YAML file - change to `otdfctl.yaml` for production, `otdfctl-example.yaml` for development
+	configFile := "otdfctl-example.yaml"
+	// configFile = "otdfctl.yaml"
+	if _, err := os.Stat(configFile); os.IsNotExist(err) {
+		configFile = "otdfctl-example.yaml"
+	}
+
+	err := Setup(configFile)
+
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Setup failed: %v\n", err)
+		return
+	}
+	logger, err := NewLogger()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error initializing logger: %v\n", err)
+		return
+	}
+	defer logger.Close()
+
+	if ask != "" {
+		HandleUserInput(ask, logger)
+		return
+	}
+
+	fmt.Println("Starting chat session. Type 'exit' or 'quit' to end.")
+	UserInputLoop(logger)
+}
+
+var chatCmd = &cobra.Command{
+	Use:   "chat",
+	Short: "Start a chat session with a LLM helper aid",
+	Long:  `This command starts an interactive chat session with a local LLM to help with setup, debugging, or generic troubleshooting`,
+	Run:   RunChatSession,
+}
+
+func GetChatCommand() *cobra.Command {
+	return chatCmd
+}

pkg/chat/intake.go

@@ -0,0 +1,211 @@
+package chat
+
+import (
+	"bufio"
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"os"
+	"os/signal"
+	"strings"
+	"syscall"
+	"time"
+)
+
+func UserInputLoop(logger *Logger) {
+	scanner := bufio.NewScanner(os.Stdin)
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+
+	for {
+		fmt.Print("> ")
+		select {
+		case <-sigChan:
+			fmt.Println("\nReceived interrupt signal. Ending chat session.")
+			return
+		default:
+			if !scanner.Scan() {
+				return
+			}
+
+			line := scanner.Text()
+			if strings.TrimSpace(line) == "exit" || strings.TrimSpace(line) == "quit" {
+				fmt.Println("Ending chat session.")
+				return
+			}
+
+			HandleUserInput(line, logger)
+		}
+	}
+}
+
+func HandleUserInput(input string, logger *Logger) {
+	sanitizedInput := SanitizeInput(input)
+
+	// if verbosity is enabled, print the sanitized input, otherwise just log it
+	if chatConfig.Chat.Verbose {
+		fmt.Printf("\n%s\n\n", sanitizedInput)
+	}
+	logger.Log(fmt.Sprintf("User: %s", input))
+	logger.Log(fmt.Sprintf("Sanitized: %s", sanitizedInput))
+
+	// Channels to receive results
+	keywordChan := make(chan []string)
+	apiResponseChan := make(chan *http.Response)
+	errorChan := make(chan error)
+	// Start keyword extraction in a goroutine
+	go func() {
+		keywords, err := ExtractKeywordsFromLLM(sanitizedInput)
+		if err != nil {
+			errorChan <- err
+			return
+		}
+		keywordChan <- keywords
+	}()
+
+	// Start main API call in a goroutine
+	go func() {
+		requestBody, err := CreateRequestBody(sanitizedInput)
+		if err != nil {
+			errorChan <- err
+			return
+		}
+
+		resp, err := SendRequest(requestBody)
+		if err != nil {
+			errorChan <- err
+			return
+		}
+		apiResponseChan <- resp
+	}()
+
+	done := make(chan bool)
+	go LoadingAnimation(done)
+
+	startTime := time.Now() // Start timing before sending the request
+
+	// Wait for both results
+	var keywords []string
+	var resp *http.Response
+	for i := 0; i < 2; i++ {
+		select {
+		case kw := <-keywordChan:
+			keywords = kw
+			// done <- true // Stop the loading animation
+			done <- true                       // Stop the loading animation
+			time.Sleep(100 * time.Millisecond) // Small delay to ensure the line is cleared
+			fmt.Println()
+			fmt.Printf("\rKeywords: [%s]\n", strings.Join(keywords, ", "))
+			fmt.Println()
+			// done <- true // Stop the loading animation
+			logger.Log(fmt.Sprintf("Keywords: [%s]", strings.Join(keywords, ", ")))
+		case r := <-apiResponseChan:
+			resp = r
+		case err := <-errorChan:
+			ReportError("during chat", err)
+			done <- true
+			return
+		}
+	}
+
+	defer resp.Body.Close()
+	ProcessResponse(resp, logger, startTime)
+}
+
+func CreateRequestBody(userInput string) ([]byte, error) {
+	return json.Marshal(map[string]interface{}{
+		"model":      chatConfig.Chat.Model,
+		"prompt":     userInput,
+		"stream":     true,
+		"tokenLimit": chatConfig.Chat.TokenLimit,
+		"options": map[string]interface{}{
+			"num_gpu": chatConfig.Chat.NumGPU,
+		},
+	})
+}
+
+func SendRequest(body []byte) (*http.Response, error) {
+	return http.Post(chatConfig.Chat.ApiURL, "application/json", bytes.NewBuffer(body))
+}
+
+func ProcessResponse(resp *http.Response, logger *Logger, startTime time.Time) {
+	responseScanner := bufio.NewScanner(resp.Body)
+	var responseBuffer bytes.Buffer
+	var tokenBuffer []string
+	tokenCount := 0
+	firstTokenReceived := false
+
+	for responseScanner.Scan() {
+		if !firstTokenReceived {
+			timeBeforeFirstToken = time.Since(startTime)
+			firstTokenReceived = true
+		}
+
+		result, err := DecodeResponse(responseScanner.Bytes())
+		if err != nil {
+			ReportError("decoding response", err)
+			continue
+		}
+
+		if response, ok := result["response"]; ok {
+			fmt.Print(response)
+			responseBuffer.WriteString(fmt.Sprintf("%s", response))
+			tokenBuffer = append(tokenBuffer, fmt.Sprintf("%s", response))
+			tokenCount++
+		}
+		if done, ok := result["done"].(bool); ok && done {
+			fmt.Println()
+			break
+		}
+		TrackStats(responseScanner.Bytes())
+
+		// Log every logLength tokens
+		if tokenCount >= chatConfig.Chat.LogLength {
+			LogWithTimestamp(logger, strings.Join(tokenBuffer, ""))
+			tokenBuffer = tokenBuffer[:0] // Reset the buffer
+			tokenCount = 0
+		}
+	}
+
+	// Log any remaining tokens
+	if tokenCount > 0 {
+		LogWithTimestamp(logger, strings.Join(tokenBuffer, ""))
+	}
+
+	PrintAndResetStats(startTime)
+}
+
+func LogWithTimestamp(logger *Logger, message string) {
+	// Remove newline characters from the message
+	cleanedMessage := strings.ReplaceAll(message, "\n", "")
+	timestamp := time.Now().Format(time.RFC3339)
+	logger.Log(fmt.Sprintf("%s: %s", timestamp, cleanedMessage))
+}
+
+func DecodeResponse(data []byte) (map[string]interface{}, error) {
+	var result map[string]interface{}
+	err := json.Unmarshal(data, &result)
+	return result, err
+}
+
+func ReportError(action string, err error) {
+	fmt.Fprintf(os.Stderr, "Error %s: %v\n", action, err)
+}
+
+func LoadingAnimation(done chan bool) {
+	chars := []rune{'|', '/', '-', '\\'}
+	for {
+		select {
+		case <-done:
+			fmt.Print("\r \r")                 // Clear the loading animation
+			time.Sleep(100 * time.Millisecond) // Small delay to ensure the line is cleared
+			return
+		default:
+			for _, char := range chars {
+				fmt.Printf("\r%c", char)
+				time.Sleep(100 * time.Millisecond)
+			}
+		}
+	}
+}

pkg/chat/logger.go

@@ -0,0 +1,33 @@
+package chat
+
+import (
+	"fmt"
+	"os"
+	"time"
+)
+
+type Logger struct {
+	file *os.File
+}
+
+// New log file with a timestamp in its name and returns a Logger instance.
+func NewLogger() (*Logger, error) {
+	timestamp := time.Now().Format("20060102_150405")
+	filename := fmt.Sprintf("pkg/chat/log/session_%s.txt", timestamp)
+	file, err := os.Create(filename)
+	if err != nil {
+		return nil, fmt.Errorf("could not create log file: %v", err)
+	}
+	return &Logger{file: file}, nil
+}
+
+// Log writes a message to the log file with a timestamp.
+func (l *Logger) Log(message string) error {
+	timestamp := time.Now().Format(time.RFC3339)
+	_, err := l.file.WriteString(fmt.Sprintf("%s: %s\n", timestamp, message))
+	return err
+}
+
+func (l *Logger) Close() {
+	l.file.Close()
+}