feat(AR-12 ) - 🚧 WIP 🚧 Integrate a locally running LLM into our CLI

.gitignore

@@ -0,0 +1,6 @@
+
+.DS_Store
+debug.log
+bin/LLMGen_index/
+bin/LLMGen/
+LLMGen_logs.log

.vscode/launch.json

@@ -0,0 +1,24 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Debug LLMGen command - knowledgebase",
+            "type": "go",
+            "request": "launch",
+            "mode": "debug",
+            "program": "${workspaceFolder}/main.go",
+            "args": ["llmgen", "Please give a high-level summary of what the DSP is and how it works", "kb"]
+        },
+        {
+            "name": "Debug LLMGen command - cliAgent",
+            "type": "go",
+            "request": "launch",
+            "mode": "debug",
+            "program": "${workspaceFolder}/main.go",
+            "args": ["llmgen", "Please give a high-level summary of what the DSP is and how it works", "cliAgent"]
+        }
+    ]
+}

cmd/llmgen.go

@@ -0,0 +1,42 @@
+package cmd
+
+import (
+	"fmt"
+
+	"github.com/opentdf/tructl/pkg/llmgen"
+	"github.com/spf13/cobra"
+)
+
+var llmgenCmd = &cobra.Command{
+	Use:   "llmgen",
+	Short: "Just a simple test mode for our llmgen integration",
+	Run: func(cmd *cobra.Command, args []string) {
+		var userQuery string
+		// var route string
+		// var function string
+		if len(args) > 0 {
+			userQuery = args[0]
+		}
+
+		// if len(args) > 1 {
+		// 	route = args[1]
+		// }
+
+		if userQuery == "" {
+			fmt.Println("Error: No query provided. Please provide a query.")
+			return
+		}
+		print("thinking...\n\n")
+		// llmgen.RawInference(function, userQuery)
+		// llmgen.RawInference(userQuery)
+		// if route == "cliAgent" {
+		// 	llmgen.CLIAgent(userQuery)
+		// } else {
+		llmgen.KnowledgeBaseChat(userQuery)
+		// }
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(llmgenCmd)
+}

cmd/root.go

@@ -10,6 +10,7 @@ import (
 )
 
 var cfgFile string
+var host string
 
 // rootCmd represents the base command when called without any subcommands
 var rootCmd = &cobra.Command{

go.sum

@@ -73,10 +73,6 @@ github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N
 github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY=
 github.com/opentdf/opentdf-v2-poc v0.0.0-20240103192012-f7b86a6a5af7 h1:HNjw5eh96WrW2n/3synRy1HLWXq3B74vexlCvm0S4FI=
 github.com/opentdf/opentdf-v2-poc v0.0.0-20240103192012-f7b86a6a5af7/go.mod h1:LLpGcqmnvPAhOuLnaP2Tsej2pKOllsoM1Ar80dqJuqk=
-github.com/opentdf/opentdf-v2-poc/sdk v0.0.0-20240117134358-3d2a619ec033 h1:ayfatFfvfxZMoquVHqWMSG9E8NgoH0LXlPWguthQd4Q=
-github.com/opentdf/opentdf-v2-poc/sdk v0.0.0-20240117134358-3d2a619ec033/go.mod h1:jjvtfVYKaaQTbovM40KRC5aVmgGmkxgacxgtEd8W7zM=
-github.com/opentdf/opentdf-v2-poc/sdk v0.0.0-20240117153907-4c2ad4cfa630 h1:U2/7fQ5OHFcHeJ/FApCOJJCCVD6SbHcbN9qvL4nZXOw=
-github.com/opentdf/opentdf-v2-poc/sdk v0.0.0-20240117153907-4c2ad4cfa630/go.mod h1:jjvtfVYKaaQTbovM40KRC5aVmgGmkxgacxgtEd8W7zM=
 github.com/opentdf/opentdf-v2-poc/sdk v0.0.0-20240117221500-8dfca6a159a8 h1:gzCWVrbeAF585gCWC5WhacHcjoTeYh+e0XaslIgbUFk=
 github.com/opentdf/opentdf-v2-poc/sdk v0.0.0-20240117221500-8dfca6a159a8/go.mod h1:jjvtfVYKaaQTbovM40KRC5aVmgGmkxgacxgtEd8W7zM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=

pkg/llmgen/llmgen.go

@@ -0,0 +1,143 @@
+package llmgen
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"log"
+	"os/exec"
+)
+
+// LLMGenOutputObject struct to hold the output of the command.
+type LLMGenOutputObject struct {
+	Output        string
+	OriginalQuery string
+	ExtraData     string
+}
+
+type LLMGenArgs struct {
+	Query    string
+	Function string
+}
+
+func NewLLMGenArgs(args ...string) LLMGenArgs {
+	var query string
+	var function string
+
+	if len(args) > 0 {
+		query = args[0]
+	}
+
+	if len(args) > 1 {
+		function = args[1]
+	}
+
+	return LLMGenArgs{
+		Query:    query,
+		Function: function,
+	}
+}
+
+func _run_binary(args LLMGenArgs) (*exec.Cmd, *bufio.Reader, error) {
+
+	// Setup the command to run the external binary.
+	cmd := exec.Command("./bin/LLMGen/LLMGen", args.Query, args.Function)
+
+	// Create a pipe to the standard output of the command.
+	stdoutPipe, err := cmd.StdoutPipe()
+	if err != nil {
+		log.Fatalf("Failed to create stdout pipe: %v", err)
+	}
+
+	// Start the command.
+	if err := cmd.Start(); err != nil {
+		log.Fatalf("Failed to start command: %v", err)
+	}
+
+	// Use a buffered reader to read the command's output in real-time.
+	reader := bufio.NewReader(stdoutPipe)
+
+	// Wait for the command to finish.
+	// if err := cmd.Wait(); err != nil {
+	// 	log.Fatalf("Command execution failed: %v", err)
+	// }
+	return cmd, reader, nil
+}
+
+func _collectAndPrintOutput(reader *bufio.Reader) (string, error) {
+	var output []byte
+	for {
+		// Read each line of the output.
+		line, err := reader.ReadBytes('\n')
+		// line = line[:len(line)-1]
+		output = append(output, line...)
+		fmt.Print(string(line))
+
+		// Break the loop if an error occurred.
+		if err != nil {
+			if err != io.EOF {
+				log.Printf("Error reading stdout: %v", err)
+			}
+			break
+		}
+	}
+	return string(output), nil
+}
+
+// RawInference executes a given command and streams the output.
+// It returns LLMGenOutputObject containing the complete output after execution.
+// func RawInference(function string, query string) (LLMGenOutputObject, error) {
+func RawInference(query string) (LLMGenOutputObject, error) {
+	_, reader, _ := _run_binary(NewLLMGenArgs(query, "/api/raw"))
+	var output string
+	output, _ = _collectAndPrintOutput(reader)
+
+	// Return the captured output.
+	return LLMGenOutputObject{
+		Output:        output,
+		OriginalQuery: query,
+		ExtraData:     string(output),
+	}, nil
+}
+
+func KnowledgeBaseChat(query string) (LLMGenOutputObject, error) {
+	_, reader, _ := _run_binary(NewLLMGenArgs(query, "/api/knowledgebase/chat"))
+	var output string
+	output, _ = _collectAndPrintOutput(reader)
+
+	// Return the captured output.
+	return LLMGenOutputObject{
+		Output:        output,
+		OriginalQuery: query,
+		ExtraData:     string(output),
+	}, nil
+}
+
+func CLIAgent(query string) (LLMGenOutputObject, error) {
+	_, reader, _ := _run_binary(NewLLMGenArgs(query, "/api/cli/help"))
+	var output string
+	output, _ = _collectAndPrintOutput(reader)
+
+	// Return the captured output.
+	return LLMGenOutputObject{
+		Output:        output,
+		OriginalQuery: query,
+		ExtraData:     string(output),
+	}, nil
+}
+
+// RawInference executes a given command and streams the output.
+// It returns LLMGenOutputObject containing the complete output after execution.
+// func RawInference(function string, query string) (LLMGenOutputObject, error) {
+func Classify(query string) (LLMGenOutputObject, error) {
+	_, reader, _ := _run_binary(NewLLMGenArgs(query, "/api/dlp/classify"))
+
+	var output string
+	output, _ = _collectAndPrintOutput(reader)
+
+	// Return the captured output.
+	return LLMGenOutputObject{
+		Output:        output,
+		OriginalQuery: query,
+	}, nil
+}