Added subtasks

Author: RijulTP

client.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"strings"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
@@ -16,9 +17,11 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
+
+
 type BulkJob struct {
-	Args     DPromptsJobArgs        `json:"args"`
-	Metadata map[string]interface{} `json:"metadata,omitempty"`
+	SubTasks   []DPromptsSubTask `json:"sub_tasks"`
+	BasePrompt string            `json:"base_prompt,omitempty"`
 }
 
 // RunClient enqueues a job with args and metadata as JSON strings.
@@ -109,32 +112,38 @@ func processJSONArray(ctx context.Context, decoder *json.Decoder, riverClient *r
 
 		params, err := toInsertParams(job)
 		if err != nil {
+			log.Error().
+				Int("job_index", total).
+				Err(err).
+				Msg("Failed to convert job to InsertManyParams")
 			return err
 		}
 
 		batch = append(batch, params)
 		total++
 		count++
 
-		if total%100 == 0 {
-			log.Info().Msgf("Loaded %d jobs...", total)
+		if total%50 == 0 {
+			log.Info().Msgf("Loaded %d jobs into batch...", total)
 		}
 
 		if count == batchSize {
+			log.Info().Msgf("Inserting batch of %d jobs (total so far: %d)", batchSize, total)
 			if err := insertBatch(ctx, riverClient, dbPool, batch); err != nil {
+				log.Error().Err(err).Msg("Failed to insert batch")
 				return err
 			}
-			log.Info().Msgf("Inserted batch of %d jobs (total: %d)", batchSize, total)
 			batch = batch[:0]
 			count = 0
 		}
 	}
 
 	if len(batch) > 0 {
+		log.Info().Msgf("Inserting final batch of %d jobs (total: %d)", len(batch), total)
 		if err := insertBatch(ctx, riverClient, dbPool, batch); err != nil {
+			log.Error().Err(err).Msg("Failed to insert final batch")
 			return err
 		}
-		log.Info().Msgf("Inserted final batch of %d jobs (total: %d)", len(batch), total)
 	}
 
 	log.Info().Msgf("Bulk insert complete. Total jobs inserted: %d", total)
@@ -160,38 +169,43 @@ func processNDJSON(ctx context.Context, decoder *json.Decoder, riverClient *rive
 
 		params, err := toInsertParams(job)
 		if err != nil {
+			log.Error().
+				Int("job_index", total).
+				Err(err).
+				Msg("Failed to convert job to InsertManyParams")
 			return err
 		}
 
 		batch = append(batch, params)
 		total++
 		count++
 
-		if total%100 == 0 {
-			log.Info().Msgf("Loaded %d jobs...", total)
+		if total%50 == 0 {
+			log.Info().Msgf("Loaded %d jobs into batch...", total)
 		}
 
 		if count == batchSize {
+			log.Info().Msgf("Inserting batch of %d jobs (total so far: %d)", batchSize, total)
 			if err := insertBatch(ctx, riverClient, dbPool, batch); err != nil {
+				log.Error().Err(err).Msg("Failed to insert batch")
 				return err
 			}
-			log.Info().Msgf("Inserted batch of %d jobs (total: %d)", batchSize, total)
 			batch = batch[:0]
 			count = 0
 		}
 	}
 
 	if len(batch) > 0 {
+		log.Info().Msgf("Inserting final batch of %d jobs (total: %d)", len(batch), total)
 		if err := insertBatch(ctx, riverClient, dbPool, batch); err != nil {
+			log.Error().Err(err).Msg("Failed to insert final batch")
 			return err
 		}
-		log.Info().Msgf("Inserted final batch of %d jobs (total: %d)", len(batch), total)
 	}
 
 	log.Info().Msgf("Bulk insert complete. Total jobs inserted: %d", total)
 	return nil
 }
-
 // Helper: Read first non-space token
 func nextNonSpaceToken(dec *json.Decoder) (json.Token, error) {
 	for {
@@ -206,35 +220,67 @@ func nextNonSpaceToken(dec *json.Decoder) (json.Token, error) {
 }
 
 func toInsertParams(job BulkJob) (river.InsertManyParams, error) {
-	var insertOpts *river.InsertOpts
-	if job.Metadata != nil {
-		metaBytes, err := json.Marshal(job.Metadata)
-		if err != nil {
-			return river.InsertManyParams{}, err
+	if len(job.SubTasks) == 0 {
+		return river.InsertManyParams{}, fmt.Errorf("job has no sub_tasks")
+	}
+
+	for i, st := range job.SubTasks {
+		if strings.TrimSpace(st.Prompt) == "" {
+			return river.InsertManyParams{}, fmt.Errorf("sub_task[%d] has empty prompt", i)
+		}
+	}
+
+	var opts *river.InsertOpts
+	if job.SubTasks[0].Metadata != nil {
+		metadataBytes, _ := json.Marshal(job.SubTasks[0].Metadata)
+		opts = &river.InsertOpts{
+			Metadata: metadataBytes,
 		}
-		insertOpts = &river.InsertOpts{Metadata: metaBytes}
 	}
+
 	return river.InsertManyParams{
-		Args:       job.Args,
-		InsertOpts: insertOpts,
+		Args: DPromptsJobArgs{
+			BasePrompt: job.BasePrompt,
+			SubTasks:   job.SubTasks,
+		},
+		InsertOpts: opts,
 	}, nil
 }
 
-func insertBatch(ctx context.Context, riverClient *river.Client[pgx.Tx], dbPool *pgxpool.Pool, batch []river.InsertManyParams) error {
+
+
+
+
+
+func insertBatch(
+	ctx context.Context,
+	riverClient *river.Client[pgx.Tx],
+	dbPool *pgxpool.Pool,
+	batch []river.InsertManyParams,
+) error {
+
 	tx, err := dbPool.Begin(ctx)
 	if err != nil {
 		return err
 	}
-	defer tx.Rollback(ctx)
+	defer func() {
+		_ = tx.Rollback(ctx) // safe no-op if already committed
+	}()
 
 	if _, err := riverClient.InsertManyTx(ctx, tx, batch); err != nil {
 		return err
 	}
 
-	return tx.Commit(ctx)
+	if err := tx.Commit(ctx); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 
+
+
 func newRiverClient(driver *riverpgxv5.Driver) (*river.Client[pgx.Tx], error) {
 	return river.NewClient[pgx.Tx](driver, &river.Config{})
 }

main.go

@@ -17,20 +17,19 @@ func main() {
 
 	rootCmd := &cobra.Command{
 		Use:   "dpr",
-		Short: "dpr CLI tool for job management",
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
 			if configPath == "" {
 				home, err := os.UserHomeDir()
 				if err != nil {
 					return err
 				}
-				configPath = home + "/.dprompt.toml"
+				configPath = home + "/.dprompts.toml"
 			}
 			return nil
 		},
 	}
 
-	rootCmd.PersistentFlags().StringVar(&configPath, "config", "", "Path to config file (default: $HOME/.dprompt.toml)")
+	rootCmd.PersistentFlags().StringVar(&configPath, "config", "", "Path to config file (default: $HOME/.dprompts.toml)")
 
 	// ---- Client subcommand ----
 	var argsJSON, metadataJSON, bulkFile string

ollama.go

@@ -3,13 +3,12 @@ package main
 import (
 	"bytes"
 	"encoding/json"
-	"errors"
+	"fmt"
 	"io"
 	"net/http"
 	"time"
 
 	"github.com/BurntSushi/toml"
-	"github.com/rs/zerolog/log"
 )
 
 
@@ -25,27 +24,25 @@ func LoadLLMConfig(configPath string) (*LLMConfig, error) {
 	return &conf.LLM, nil
 }
 
-func CallOllama(prompt string, schema interface{}, configPath string, groupName string, base_prompt string) (string, error) {
-	startTotal := time.Now()
+func CallOllama(
+	prompt string,
+	schema interface{},
+	configPath string,
+	basePrompt string,
+) (string, error) {
 
-	// 1️⃣ Load config
-	start := time.Now()
+	// Load config
 	llmConfig, err := LoadLLMConfig(configPath)
 	if err != nil {
 		return "", err
 	}
-	log.Info().
-		Dur("duration", time.Since(start)).
-		Msg("LoadLLMConfig duration")
 
-	// 2️⃣ Build request body
-	start = time.Now()
-	req := map[string]interface{}{
+	// Build request
+	req := map[string]any{
 		"model":  llmConfig.Model,
 		"stream": false,
-		"session_id": groupName,
 		"messages": []map[string]string{
-			{"role": "system", "content": base_prompt },
+			{"role": "system", "content": basePrompt},
 			{"role": "user", "content": prompt},
 		},
 		"options": map[string]float64{
@@ -64,52 +61,36 @@ func CallOllama(prompt string, schema interface{}, configPath string, groupName
 	if err != nil {
 		return "", err
 	}
-	log.Info().
-		Dur("duration", time.Since(start)).
-		Msg("Marshal request duration")
 
-	// 3️⃣ HTTP POST
-	start = time.Now()
+	// HTTP call
 	client := &http.Client{Timeout: 360 * time.Second}
-	resp, err := client.Post(llmConfig.APIEndpoint, "application/json", bytes.NewReader(reqBody))
+	resp, err := client.Post(
+		llmConfig.APIEndpoint,
+		"application/json",
+		bytes.NewReader(reqBody),
+	)
 	if err != nil {
 		return "", err
 	}
 	defer resp.Body.Close()
 
-	log.Info().
-		Dur("duration", time.Since(start)).
-		Msg("HTTP POST duration")
-
 	if resp.StatusCode != http.StatusOK {
-		return "", errors.New("ollama API returned non-200 status: " + resp.Status)
+		return "", fmt.Errorf("ollama API returned %s", resp.Status)
 	}
 
-	// 4️⃣ Read response
-	start = time.Now()
+	// Read & decode response
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return "", err
 	}
-	log.Info().
-		Dur("duration", time.Since(start)).
-		Msg("Read response duration")
 
-	// 5️⃣ Unmarshal response
-	start = time.Now()
 	var ollamaResp OllamaResponse
 	if err := json.Unmarshal(body, &ollamaResp); err != nil {
 		return "", err
 	}
-	log.Info().
-		Dur("duration", time.Since(start)).
-		Msg("Unmarshal response duration")
-
-	log.Info().
-		Dur("total_duration", time.Since(startTotal)).
-		Msg("Ollama call completed")
 
 	return ollamaResp.Message.Content, nil
 }
 
 
+

types.go

@@ -9,13 +9,18 @@ type DBConfig struct {
 	Port     string
 }
 
-
-type DPromptsJobArgs struct {
+type DPromptsSubTask struct {
 	Prompt       string      `json:"prompt"`
 	Schema       interface{} `json:"schema,omitempty"`
-	BasePrompt string      `json:"base_prompt,omitempty"`
+	Metadata map[string]interface{} `json:"metadata,omitempty"` // <-- new
+
 }
 
+type DPromptsJobArgs struct {
+	SubTasks   []DPromptsSubTask `json:"sub_tasks"`
+	BasePrompt string            `json:"base_prompt,omitempty"`
+
+}
 
 type DPromptsJobResult struct {
 	Response string `json:"response"`