chore: create specific tools for LLM vs. SDK

main.go

@@ -30,7 +30,13 @@ env vars: GPTSCRIPT_WORKSPACE_DIR`)
 	case "listElements":
 		tools.ListElements(os.Getenv("DATASETID"))
 	case "getElement":
-		tools.GetElement(os.Getenv("DATASETID"), os.Getenv("ELEMENT"))
+		tools.GetElementLLM(os.Getenv("DATASETID"), os.Getenv("ELEMENT"))
+	case "getAllElements":
+		tools.GetAllElementsLLM(os.Getenv("DATASETID"))
+	case "getElementSDK":
+		tools.GetElementSDK(os.Getenv("DATASETID"), os.Getenv("ELEMENT"))
+	case "getAllElementsSDK":
+		tools.GetAllElementsSDK(os.Getenv("DATASETID"))
 	case "createDataset":
 		tools.CreateDataset(os.Getenv("DATASETNAME"), os.Getenv("DATASETDESCRIPTION"))
 	case "addElement":
@@ -43,8 +49,6 @@ env vars: GPTSCRIPT_WORKSPACE_DIR`)
 		}
 
 		addElements(os.Getenv("DATASETID"), elementInputs)
-	case "getAllElements":
-		tools.GetAllElements(os.Getenv("DATASETID"))
 	default:
 		fmt.Printf("unknown command: %s\n", os.Args[1])
 		os.Exit(1)

pkg/tools/getAllElements.go

@@ -2,14 +2,42 @@ package tools
 
 import (
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/gptscript-ai/datasets/pkg/dataset"
 )
 
-func GetAllElements(datasetID string) {
+func GetAllElementsLLM(datasetID string) {
+	elems := getAllElements(datasetID)
+
+	var elemStrings []string
+	for _, e := range elems {
+		rawContents, err := base64.StdEncoding.DecodeString(e.Contents)
+		if err != nil {
+			rawContents = []byte(e.Contents)
+		}
+		elemStrings = append(elemStrings, fmt.Sprintf(`{"name": %q, "description": %q, "contents": %q}`, e.Name, e.Description, string(rawContents)))
+	}
+
+	fmt.Printf("[%s]", strings.Join(elemStrings, ","))
+}
+
+func GetAllElementsSDK(datasetID string) {
+	elems := getAllElements(datasetID)
+	elemsJSON, err := json.Marshal(elems)
+	if err != nil {
+		fmt.Printf("failed to marshal elements: %v\n", err)
+		os.Exit(1)
+	}
+
+	fmt.Println(string(elemsJSON))
+}
+
+func getAllElements(datasetID string) []elem {
 	m, err := dataset.NewManager()
 	if err != nil {
 		fmt.Printf("failed to create dataset manager: %v\n", err)
@@ -38,11 +66,5 @@ func GetAllElements(datasetID string) {
 		})
 	}
 
-	elemsJSON, err := json.Marshal(elems)
-	if err != nil {
-		fmt.Printf("failed to marshal elements: %v\n", err)
-		os.Exit(1)
-	}
-
-	fmt.Println(string(elemsJSON))
+	return elems
 }

pkg/tools/getElement.go

@@ -2,6 +2,7 @@ package tools
 
 import (
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"os"
@@ -15,7 +16,30 @@ type elem struct {
 	Description string `json:"description,omitempty"`
 }
 
-func GetElement(datasetID, elementName string) {
+func GetElementLLM(datasetID, elementName string) {
+	element := getElement(datasetID, elementName)
+	// Attempt to base64 decode the contents.
+	rawContents, err := base64.StdEncoding.DecodeString(element.Contents)
+	if err != nil {
+		// If it's not base64, just use the contents.
+		rawContents = []byte(element.Contents)
+	}
+
+	fmt.Printf(`{"name": %q, "description": %q, "contents": %q}`, element.Name, element.Description, string(rawContents))
+}
+
+func GetElementSDK(datasetID, elementName string) {
+	element := getElement(datasetID, elementName)
+	elementJSON, err := json.Marshal(element)
+	if err != nil {
+		fmt.Printf("failed to marshal element: %v\n", err)
+		os.Exit(1)
+	}
+
+	fmt.Print(string(elementJSON))
+}
+
+func getElement(datasetID, elementName string) elem {
 	m, err := dataset.NewManager()
 	if err != nil {
 		fmt.Printf("failed to create dataset manager: %v\n", err)
@@ -34,17 +58,9 @@ func GetElement(datasetID, elementName string) {
 		os.Exit(1)
 	}
 
-	element := elem{
+	return elem{
 		Contents:    string(elementContents),
 		Name:        e.Name,
 		Description: e.Description,
 	}
-
-	elementJSON, err := json.Marshal(element)
-	if err != nil {
-		fmt.Printf("failed to marshal element: %v\n", err)
-		os.Exit(1)
-	}
-
-	fmt.Print(string(elementJSON))
 }

tool.gpt

@@ -10,14 +10,36 @@ Param: datasetID: the ID of the dataset
 
 #!${GPTSCRIPT_TOOL_DIR}/bin/gptscript-go-tool listElements
 
+---
+Name: Get Element SDK
+Description: Gets a particular element's metadata and contents. For SDK use only.
+Param: datasetID: the ID of the dataset
+Param: element: the name of the element
+
+#!${GPTSCRIPT_TOOL_DIR}/bin/gptscript-go-tool getElement
+
+---
+Name: Get All Elements SDK
+Description: Gets the contents of all elements in a dataset. For SDK use only.
+Param: datasetID: the ID of the dataset
+
+#!${GPTSCRIPT_TOOL_DIR}/bin/gptscript-go-tool getAllElements
+
 ---
 Name: Get Element
-Description: Gets a particular element's metadata and contents
+Description: Gets a particular element's metadata and contents.
 Param: datasetID: the ID of the dataset
 Param: element: the name of the element
 
 #!${GPTSCRIPT_TOOL_DIR}/bin/gptscript-go-tool getElement
 
+---
+Name: Get All Elements
+Description: Gets the contents of all elements in a dataset.
+Param: datasetID: the ID of the dataset
+
+#!${GPTSCRIPT_TOOL_DIR}/bin/gptscript-go-tool getAllElements
+
 ---
 Name: Create Dataset
 Description: Creates a new, empty dataset
@@ -44,9 +66,4 @@ Param: elements: a JSON array of elements to add, like [{"name":"one","descripti
 
 #!${GPTSCRIPT_TOOL_DIR}/bin/gptscript-go-tool addElements
 
----
-Name: Get All Elements
-Description: Gets the contents of all elements in a dataset
-Param: datasetID: the ID of the dataset
 
-#!${GPTSCRIPT_TOOL_DIR}/bin/gptscript-go-tool getAllElements