Merge pull request #1 from HexmosTech/rijul/schema-feature

Added schema support, grouping of tasks, queue management commands and data viewing

Author: RijulTP

README.md

@@ -75,6 +75,45 @@ The JSON file should be an array of job objects:
 ]
 ```
 
+### View Last Results
+
+To view the last `n` results processed by the worker:
+
+```
+dpr --mode=view --n=20
+```
+
+
+### View Groups
+
+List all groups created in the system:
+
+```
+dpr --mode=view --total-groups
+```
+
+### View results for a specific group by its ID
+
+```
+dpr --mode=view --group=1
+```
+
+### View Queued Jobs
+
+```
+dpr --mode=queue --action=view --queue-n=20
+```
+
+### Clear Queued Jobs
+
+```
+dpr --mode=queue --action=clear
+```
+
+
+
+
+
 ## Useful Ollama Commands
 
 - **Run Ollama server:**
@@ -115,4 +154,9 @@ The JSON file should be an array of job objects:
 
 - The `.dprompts.toml` file **must** be placed in your home directory.
 - You can customize job arguments and metadata using the `--args` and `--metadata` flags (as JSON).
-- The worker will process jobs and store results in the configured PostgreSQL database.
+- The worker will process jobs and store results in the configured PostgreSQL database.
+- **PostgreSQL Storage Details:**
+  - `dprompt_results` — stores the results of processed jobs.
+  - `dprompt_groups` — stores job groups with unique `group_name` and `id`.
+    - Groups are a way to organize related jobs that work toward the same goal.
+    - This makes it easy to view or analyze all jobs related to a single goal together.

bulk_jobs_sample_with_grouping.json

@@ -0,0 +1,47 @@
+[
+    {
+      "args": {
+        "prompt": "Explain the difference between TCP and UDP.",
+        "group_name": "networking_basics"
+      },
+      "metadata": {
+        "type": "test",
+        "category": "networking",
+        "filename": "tcp_udp.md"
+      }
+    },
+    {
+      "args": {
+        "prompt": "What is a subnet mask and why is it important?",
+        "group_name": "networking_basics"
+      },
+      "metadata": {
+        "type": "test",
+        "category": "networking",
+        "filename": "subnet_mask.md"
+      }
+    },
+    {
+      "args": {
+        "prompt": "How do you reverse a linked list in Python?",
+        "group_name": "python_algorithms"
+      },
+      "metadata": {
+        "type": "test",
+        "category": "python",
+        "filename": "reverse_linked_list.md"
+      }
+    },
+    {
+      "args": {
+        "prompt": "Explain Python's list comprehension with examples.",
+        "group_name": "python_algorithms"
+      },
+      "metadata": {
+        "type": "test",
+        "category": "python",
+        "filename": "list_comprehension.md"
+      }
+    }
+  ]
+  

client.go

@@ -40,6 +40,7 @@ func RunClient(ctx context.Context, driver *riverpgxv5.Driver, argsJSON string,
 		log.Fatal().Err(err).Msg("Failed to parse args JSON")
 	}
 
+
 	var insertOpts *river.InsertOpts
 	if metadataJSON != "" {
 		var metadata map[string]interface{}
@@ -58,7 +59,11 @@ func RunClient(ctx context.Context, driver *riverpgxv5.Driver, argsJSON string,
 	if _, err := riverClient.Insert(ctx, &args, insertOpts); err != nil {
 		log.Fatal().Err(err).Msg("Failed to enqueue job")
 	}
-	log.Info().Interface("args", args).Interface("metadata", insertOpts).Msg("Enqueued job")
+
+	log.Info().
+		Interface("args", args).
+		Interface("metadata", insertOpts).
+		Msg("Enqueued job")
 }
 
 func enqueueBulkJobsFromFile(ctx context.Context, riverClient *river.Client[pgx.Tx], dbPool *pgxpool.Pool, filename string) error {
@@ -80,17 +85,19 @@ func enqueueBulkJobsFromFile(ctx context.Context, riverClient *river.Client[pgx.
 	defer tx.Rollback(ctx)
 
 	var jobsToInsert []river.InsertManyParams
-	for _, job := range jobs {
+
+	for i := range jobs {
+
 		var insertOpts *river.InsertOpts
-		if job.Metadata != nil {
-			metadataBytes, err := json.Marshal(job.Metadata)
+		if jobs[i].Metadata != nil {
+			metadataBytes, err := json.Marshal(jobs[i].Metadata)
 			if err != nil {
 				return err
 			}
 			insertOpts = &river.InsertOpts{Metadata: metadataBytes}
 		}
 		jobsToInsert = append(jobsToInsert, river.InsertManyParams{
-			Args:       job.Args,
+			Args:       jobs[i].Args,
 			InsertOpts: insertOpts,
 		})
 	}

config.go

@@ -65,4 +65,4 @@ func NewDBPool(ctx context.Context, configPath string) (*pgxpool.Pool, error) {
 	}
 
 	return dbPool, nil
-}
+}

groupmanagement.go

@@ -0,0 +1,26 @@
+package main
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+// DeleteGroupAndResults deletes a group by its ID and all associated results
+func DeleteGroupAndResults(ctx context.Context, db *pgxpool.Pool, groupID int) error {
+	// Delete associated results first
+	res1, err := db.Exec(ctx, `DELETE FROM dprompts_results WHERE group_id = $1`, groupID)
+	if err != nil {
+		return fmt.Errorf("failed to delete results: %w", err)
+	}
+
+	// Delete the group itself
+	res2, err := db.Exec(ctx, `DELETE FROM dprompt_groups WHERE id = $1`, groupID)
+	if err != nil {
+		return fmt.Errorf("failed to delete group: %w", err)
+	}
+
+	fmt.Printf("Deleted %d results and %d group(s) for group ID %d\n", res1.RowsAffected(), res2.RowsAffected(), groupID)
+	return nil
+}

main.go

@@ -19,6 +19,15 @@ func main() {
 	metadataJSON := flag.String("metadata", "", "Job metadata as JSON (for client mode)")
 	bulkFile := flag.String("bulk-from-file", "", "Bulk insert jobs from JSON file")
 	configPath := flag.String("config", "", "Path to config file (default: $HOME/.dprompt.toml)")
+	
+	totalGroups := flag.Bool("total-groups", false, "Display total number of groups (view mode)")
+	groupID := flag.Int("group", 0, "Display results for a specific group ID (view mode)")
+	deleteGroupID := flag.Int("delete-group-id", 0, "Delete a specific group ID and all its associated results")
+
+	n := flag.Int("n", 10, "Number of results to display (view mode)")
+	
+	queueN := flag.Int("queue-n", 10, "Number of queued jobs to display (for view action)")
+	queueAction := flag.String("action", "", "Queue action: 'view' to list queued jobs, 'clear' to delete all queued jobs")
 
 	flag.Parse()
 
@@ -39,6 +48,41 @@ func main() {
 		RunClient(ctx, driver, *argsJSON, *metadataJSON, *bulkFile, dbPool)
 	case "worker":
 		RunWorker(ctx, driver, cancel, dbPool)
+	case "delete-group":
+		if *deleteGroupID == 0 {
+			log.Fatal().Msg("Please provide --delete-group-id")
+		}
+		if err := DeleteGroupAndResults(ctx, dbPool, *deleteGroupID); err != nil {
+			log.Fatal().Err(err).Msg("Failed to delete group and results")
+		}
+		log.Info().Int("group_id", *deleteGroupID).Msg("Deleted group and associated results")
+	case "view":
+		if *totalGroups {
+			if err := viewTotalGroups(ctx, dbPool); err != nil {
+				log.Fatal().Err(err).Msg("Failed to get total groups")
+			}
+		} else if *groupID != 0 {
+			if err := viewResultsByGroup(ctx, dbPool, *groupID); err != nil {
+				log.Fatal().Err(err).Msg("Failed to get results by group")
+			}
+		} else {
+			if err := viewLastResults(ctx, dbPool, *n); err != nil {
+				log.Fatal().Err(err).Msg("Failed to get last results")
+			}
+		}
+	case "queue":
+		switch *queueAction {
+		case "view":
+			if err := ViewQueuedJobs(ctx, dbPool, *queueN); err != nil {
+				log.Fatal().Err(err).Msg("Failed to view queued jobs")
+			}
+		case "clear":
+			if err := ClearQueuedJobs(ctx, dbPool); err != nil {
+				log.Fatal().Err(err).Msg("Failed to clear queued jobs")
+			}
+		default:
+			log.Fatal().Str("action", *queueAction).Msg("Unknown queue action")
+		}
 	default:
 		log.Fatal().Str("mode", *mode).Msg("Unknown mode")
 	}

managequeuedjobs.go

@@ -0,0 +1,48 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+func ViewQueuedJobs(ctx context.Context, db *pgxpool.Pool, n int) error {
+	rows, err := db.Query(ctx, `
+		SELECT id, state, created_at, scheduled_at
+		FROM river_job
+		WHERE state IN ('available', 'scheduled')
+		ORDER BY created_at DESC
+		LIMIT $1
+	`, n)
+	if err != nil {
+		return err
+	}
+	defer rows.Close()
+
+	fmt.Printf("Last %d queued jobs:\n", n)
+	for rows.Next() {
+		var id int64
+		var state string
+		var createdAt, scheduledAt time.Time
+		if err := rows.Scan(&id, &state, &createdAt, &scheduledAt); err != nil {
+			return err
+		}
+		fmt.Printf("ID: %d | State: %s | CreatedAt: %s | ScheduledAt: %s\n",
+			id, state, createdAt.Format(time.RFC3339), scheduledAt.Format(time.RFC3339))
+	}
+	return rows.Err()
+}
+
+func ClearQueuedJobs(ctx context.Context, db *pgxpool.Pool) error {
+	res, err := db.Exec(ctx, `
+		DELETE FROM river_job
+		WHERE state IN ('available', 'scheduled')
+	`)
+	if err != nil {
+		return err
+	}
+	fmt.Printf("Deleted %d queued jobs\n", res.RowsAffected())
+	return nil
+}

ollama.go

@@ -25,23 +25,33 @@ func LoadLLMConfig(configPath string) (*LLMConfig, error) {
 	return &conf.LLM, nil
 }
 
-func CallOllama(prompt string, configPath string) (string, error) {
+func CallOllama(prompt string, schema interface{}, configPath string) (string, error) {
 	llmConfig, err := LoadLLMConfig(configPath)
 	if err != nil {
 		return "", err
 	}
 
-	reqBody, err := json.Marshal(map[string]interface{}{
-		"model":    llmConfig.Model,
-		"messages": []map[string]string{{"role": "user", "content": prompt}},
-		"stream":   false,
-		"format":   "json",
+	// Build request body
+	req := map[string]interface{}{
+		"model":  llmConfig.Model,
+		"stream": false,
+		"messages": []map[string]string{
+			{"role": "user", "content": prompt},
+		},
 		"options": map[string]float64{
 			"temperature": llmConfig.Temperature,
 			"top_p":       llmConfig.TopP,
-			"num_predict": 100,
 		},
-	})
+	}
+
+	// ⬇️ Only include schema if provided
+	if schema != nil {
+		req["format"] = schema
+	} else {
+		req["format"] = "json" // default JSON output
+	}
+
+	reqBody, err := json.Marshal(req)
 	if err != nil {
 		return "", err
 	}

sql-queries/dprompts-groups.sql

@@ -0,0 +1,5 @@
+CREATE TABLE dprompt_groups (
+    id INT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
+    group_name TEXT NOT NULL UNIQUE,
+    created_at TIMESTAMPTZ DEFAULT NOW()
+);

sql-queries/dprompts-results.sql

@@ -0,0 +1,10 @@
+CREATE TABLE dprompts_results (
+    id INT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
+    job_id BIGINT UNIQUE,
+    response JSONB,
+    created_at TIMESTAMPTZ DEFAULT NOW(),
+    group_id INT,
+    CONSTRAINT fk_group
+        FOREIGN KEY (group_id)
+        REFERENCES dprompt_groups(id)
+);