Merge pull request #4 from HexmosTech/rijul/schema-validation

Add schema validation for requests with defined schemas

Author: RijulTP

go.mod

@@ -26,6 +26,7 @@ require (
 	github.com/riverqueue/river/riverdriver v0.26.0 // indirect
 	github.com/riverqueue/river/rivershared v0.26.0 // indirect
 	github.com/riverqueue/river/rivertype v0.26.0 // indirect
+	github.com/santhosh-tekuri/jsonschema/v5 v5.3.1 // indirect
 	github.com/spf13/cobra v1.10.2 // indirect
 	github.com/spf13/pflag v1.0.9 // indirect
 	github.com/stretchr/testify v1.11.1 // indirect

go.sum

@@ -61,6 +61,8 @@ github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0=
 github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY=
 github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/santhosh-tekuri/jsonschema/v5 v5.3.1 h1:lZUw3E0/J3roVtGQ+SCrUrg3ON6NgVqpn3+iol9aGu4=
+github.com/santhosh-tekuri/jsonschema/v5 v5.3.1/go.mod h1:uToXkOrWAZ6/Oc07xWQrPOhJotwFIyu2bBVN41fcDUY=
 github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU=
 github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4=
 github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY=

main.go

@@ -122,8 +122,8 @@ func main() {
 			}
 		},
 	}
-	viewCmd.Flags().BoolVar(&totalGroups, "total-groups", false, "Display total number of groups")
-	viewCmd.Flags().IntVar(&groupID, "group", 0, "Display results for a specific group ID")
+	// viewCmd.Flags().BoolVar(&totalGroups, "total-groups", false, "Display total number of groups")
+	// viewCmd.Flags().IntVar(&groupID, "group", 0, "Display results for a specific group ID")
 	viewCmd.Flags().IntVarP(&n, "number", "n", 10, "Number of results to display")
 
 	// ---- Delete-group subcommand ----
@@ -348,7 +348,7 @@ func main() {
 	)
 
 	// Add subcommands
-	rootCmd.AddCommand(clientCmd, workerCmd, viewCmd, deleteGroupCmd, queueCmd, exportCmd)
+	rootCmd.AddCommand(clientCmd, workerCmd, viewCmd, queueCmd, exportCmd)
 
 	if err := rootCmd.Execute(); err != nil {
 		log.Fatal().Err(err).Msg("Command execution failed")

ollama.go

@@ -4,14 +4,15 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"github.com/AlecAivazis/survey/v2"
+	"github.com/BurntSushi/toml"
+	"github.com/santhosh-tekuri/jsonschema/v5"
 	"io"
 	"net/http"
 	"os/exec"
 	"runtime"
+	"strings"
 	"time"
-
-	"github.com/AlecAivazis/survey/v2"
-	"github.com/BurntSushi/toml"
 )
 
 func LoadLLMConfig(configPath string) (*LLMConfig, error) {
@@ -41,7 +42,7 @@ func CallOllama(
 	// Build request
 	req := map[string]any{
 		"model":  llmConfig.Model,
-		"stream": false,
+		"stream": true,
 		"messages": []map[string]string{
 			{"role": "system", "content": basePrompt},
 			{"role": "user", "content": prompt},
@@ -79,18 +80,32 @@ func CallOllama(
 		return "", fmt.Errorf("ollama API returned %s", resp.Status)
 	}
 
-	// Read & decode response
-	body, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return "", err
+	// Decode streamed JSON objects one by one
+	decoder := json.NewDecoder(resp.Body)
+
+	var fullContent strings.Builder
+	for {
+		var chunk OllamaResponse
+		if err := decoder.Decode(&chunk); err != nil {
+			if err == io.EOF {
+				break
+			}
+			return "", err
+		}
+
+		fullContent.WriteString(chunk.Message.Content)
 	}
 
-	var ollamaResp OllamaResponse
-	if err := json.Unmarshal(body, &ollamaResp); err != nil {
-		return "", err
+	output := fullContent.String()
+
+	// schema validation (fail job if invalid)
+	if schema != nil {
+		if err := validateJSONAgainstSchema(output, schema); err != nil {
+			return "", err
+		}
 	}
 
-	return ollamaResp.Message.Content, nil
+	return output, nil
 }
 
 func isOllamaRunning() bool {
@@ -221,3 +236,35 @@ func waitForOllama(timeout time.Duration) error {
 
 	return fmt.Errorf("ollama did not start within %s", timeout)
 }
+
+func validateJSONAgainstSchema(output string, schema any) error {
+	if schema == nil {
+		return nil
+	}
+
+	schemaBytes, err := json.Marshal(schema)
+	if err != nil {
+		return fmt.Errorf("failed to marshal schema: %w", err)
+	}
+
+	compiler := jsonschema.NewCompiler()
+	if err := compiler.AddResource("schema.json", bytes.NewReader(schemaBytes)); err != nil {
+		return fmt.Errorf("failed to add schema: %w", err)
+	}
+
+	compiledSchema, err := compiler.Compile("schema.json")
+	if err != nil {
+		return fmt.Errorf("failed to compile schema: %w", err)
+	}
+
+	var instance any
+	if err := json.Unmarshal([]byte(output), &instance); err != nil {
+		return fmt.Errorf("output is not valid JSON: %w", err)
+	}
+
+	if err := compiledSchema.Validate(instance); err != nil {
+		return fmt.Errorf("schema validation failed: %w", err)
+	}
+
+	return nil
+}