Refactor job pkg

This commit is contained in:
dwrz
2026-10-30 00:45:54 +00:00
parent 0eca987ba1
commit 7d8f8f1b2a
2 changed files with 105 additions and 962 deletions
+35 -458
View File
@@ -1,45 +1,24 @@
// Package job provides cron-based scheduling for autonomous LLM tasks.
//
// Jobs are defined in the YAML configuration with a cron expression,
// a system message, and a task description. Each job owns its inference
// and tool-execution loop.
// a system message, and a task description. The service creates a
// conversation for each firing; [Job.Run] runs the job's task in that
// conversation through [conversation.Conversation.Reserve].
package job
import (
"context"
"database/sql"
"encoding/json/jsontext"
"encoding/json/v2"
"errors"
"fmt"
"log/slog"
"runtime"
"time"
"github.com/robfig/cron/v3"
"code.chimeric.al/chimerical/odidere/internal/config"
"code.chimeric.al/chimerical/odidere/internal/db"
"code.chimeric.al/chimerical/odidere/internal/conversation"
"code.chimeric.al/chimerical/odidere/internal/llm"
"code.chimeric.al/chimerical/odidere/internal/tool"
)
const (
// interruptedWarning indicates a tool call started but did not finish
// normally, so its effects are unknown.
interruptedWarning = `▸ WARNING: the tool started, but its outcome is unknown.
It may have had effects. Do not assume it is safe to repeat.`
// resultSaveTimeout is the time limit for saving one tool result.
resultSaveTimeout = time.Second
)
// errMaxIterations is the error Run reports when a job's
// MaxIterations limit is reached before the job completes.
// A job with MaxIterations set to zero has no iteration limit and never
// produces this error.
var errMaxIterations = errors.New("max iterations exceeded")
// Job defines a scheduled autonomous task.
//
// Job is immutable after construction via New. All fields are private and
@@ -48,8 +27,6 @@ var errMaxIterations = errors.New("max iterations exceeded")
type Job struct {
// cron is the cron.Schedule derived from schedule during construction.
cron cron.Schedule
// db is the database connection pool used to persist job conversations.
db *sql.DB
// llm is the LLM client used to execute the job.
llm *llm.Client
// log is the structured logger used for job execution.
@@ -70,55 +47,30 @@ type Job struct {
// fires. This is the core of what the job does.
task string
// timeout is the parsed execution time limit.
// Zero means no per-job timeout (falls back to the global LLM timeout).
// Zero means no per-job timeout.
timeout time.Duration
// tools contains the tools available to the job.
tools *tool.Registry
// userId identifies the owner of conversations created by the job.
userId string
}
// Result holds the outcome of a job execution.
type Result struct {
// ConversationID identifies the conversation for this invocation.
ConversationID string
// Duration is how long the job took to run.
Duration time.Duration
// Error is set if the job failed.
Error error
// Messages is the full conversation returned by the LLM agent loop,
// including tool calls and tool results. Empty if the job failed.
Messages []llm.Message
// Name is the job name.
Name string
}
// New constructs a Job from cfg and its runtime dependencies. New parses
// values required for execution and rejects dependencies that Run cannot
// operate without. log must be non-nil; pass slog.New(slog.DiscardHandler)
// to run without logs.
func New(
cfg config.JobConfig,
db *sql.DB,
llmc *llm.Client,
log *slog.Logger,
tools *tool.Registry,
) (*Job, error) {
if err := cfg.Validate(); err != nil {
return nil, fmt.Errorf("invalid config: %w", err)
}
if db == nil {
return nil, fmt.Errorf("missing database")
}
if llmc == nil {
return nil, fmt.Errorf("missing llm client")
}
if log == nil {
return nil, fmt.Errorf("missing logger")
}
if tools == nil {
return nil, fmt.Errorf("missing tool registry")
}
sched, err := cron.ParseStandard(cfg.Schedule)
if err != nil {
@@ -129,7 +81,6 @@ func New(
return &Job{
cron: sched,
db: db,
llm: llmc,
log: log.With(slog.String("job", cfg.Name)),
maxIterations: cfg.MaxIterations,
@@ -139,439 +90,62 @@ func New(
systemMessage: cfg.SystemMessage,
task: cfg.Task,
timeout: cfg.Timeout,
tools: tools,
userId: cfg.UserId,
}, nil
}
// callTool records one tool call as running, runs the tool, and saves
// the tool message that reports the outcome to the model.
//
// The record is a row in the executions table, saved before the tool
// starts. After an interruption, the row shows whether the tool may have
// run:
// - No row: the tool never started. It is safe to run again.
// - Row with status 'running': the tool may have run and had effects.
// - Row with status 'completed': the tool finished or never started,
// and the tool message is saved.
//
// Results:
// - A message and a nil error: the message is saved. The message
// reports the tool's outcome: exit code and output, the error that
// kept the tool from starting, or the error that interrupted it.
// - A nil message and an error: the outcome was not recorded. The tool
// may have run and had effects.
//
// The message contains:
// - If the tool finished: its output.
// - If it never started: the error.
// - If it started but was stopped, or its output was cut off: the
// error, a warning that its effects are unknown, and any output.
//
// Limitations:
// - callTool relies on [tool.Registry.Execute] returning a nil result
// only when the process never started.
func (j *Job) callTool(
ctx context.Context,
log *slog.Logger,
cid, mid string,
call llm.ToolCall,
) (*llm.Message, error) {
log = log.With(
slog.String("tool_call_id", call.ID),
slog.String("name", call.Function.Name),
)
execution, err := db.New(j.db).InsertExecution(ctx, db.InsertExecutionParams{
MessageId: mid,
Status: db.ExecutionStatusRunning,
ToolCallId: call.ID,
})
if err != nil {
return nil, fmt.Errorf(
"insert execution for tool call %q: %w", call.ID, err,
)
}
log.InfoContext(
ctx,
"calling tool",
slog.String("args", call.Function.Arguments),
)
result, toolErr := j.tools.Execute(
ctx, call.Function.Name, call.Function.Arguments,
)
var content string
switch {
case toolErr == nil:
log.InfoContext(
ctx,
"called tool",
slog.Duration(
"duration",
result.FinishedAt.Sub(result.StartedAt),
),
)
content = result.FormatResult()
case result == nil:
log.ErrorContext(
ctx,
"tool did not start",
slog.Any("error", toolErr),
)
content = toolErr.Error()
default:
log.ErrorContext(
ctx,
"tool interrupted",
slog.Any("error", toolErr),
slog.Duration(
"duration",
result.FinishedAt.Sub(result.StartedAt),
),
)
content = toolErr.Error() + "\n" +
interruptedWarning + "\n" +
result.FormatResult()
}
message := llm.Message{
Role: llm.RoleTool,
ContentParts: []llm.ContentPart{{
Type: llm.ContentTypeText,
Text: content,
}},
Name: call.Function.Name,
ToolCallID: call.ID,
}
var update *db.UpdateExecutionStatusParams
if toolErr == nil || result == nil {
update = &db.UpdateExecutionStatusParams{
Id: execution.Id,
Status: db.ExecutionStatusCompleted,
}
}
// Try to persist the outcome even if the job was canceled after the
// tool ran. If saving fails, the running row records that it may have run.
saveCtx, saveCancel := context.WithTimeout(
context.WithoutCancel(ctx), resultSaveTimeout,
)
defer saveCancel()
if _, err := j.insertMessage(saveCtx, cid, message, nil, update); err != nil {
log.ErrorContext(
ctx,
"failed to persist tool result",
slog.Any("error", err),
)
return nil, fmt.Errorf("persist tool call %q: %w", call.ID, err)
}
return &message, nil
}
// callTools runs every tool call in an assistant message and collects its
// persisted result. cid is the conversation's ID; mid is the assistant
// message's ID.
//
// - Calls run concurrently, at most twice the number of CPUs at once.
// - By design, each call saves its own result, so the saved order is
// not fixed. The returned order may differ from the saved order.
// - If a call's row or result cannot be saved, the other calls still run and
// save.
// - The returned error joins every failure. If it is non-nil, at least
// one call has no result, so the next inference request cannot be
// sent.
func (j *Job) callTools(
ctx context.Context,
log *slog.Logger,
cid, mid string,
message llm.Message,
) ([]llm.Message, error) {
type outcome struct {
err error
msg *llm.Message
}
var (
completed = make(chan outcome, len(message.ToolCalls))
sem = make(chan struct{}, runtime.NumCPU()*2)
)
for _, call := range message.ToolCalls {
go func() {
sem <- struct{}{}
defer func() { <-sem }()
msg, err := j.callTool(ctx, log, cid, mid, call)
completed <- outcome{err: err, msg: msg}
}()
}
var (
errs []error
results []llm.Message
)
for range message.ToolCalls {
result := <-completed
if result.err != nil {
log.ErrorContext(
ctx,
"tool call failed",
slog.Any("error", result.err),
)
errs = append(errs, result.err)
continue
}
results = append(results, *result.msg)
}
return results, errors.Join(errs...)
}
// createConversation creates the conversation for one run of the job.
//
// messages is the complete initial history, in the order it should appear.
// The conversation and its initial messages persist together.
func (j *Job) createConversation(
ctx context.Context,
messages []llm.Message,
) (*db.Conversation, error) {
tx, err := j.db.BeginTx(ctx, nil)
if err != nil {
return nil, fmt.Errorf("begin conversation transaction: %w", err)
}
defer tx.Rollback()
q := db.New(tx)
c, err := q.InsertConversation(
ctx,
db.InsertConversationParams{
Title: &j.name,
UserId: j.userId,
},
)
if err != nil {
return nil, fmt.Errorf("insert conversation: %w", err)
}
for sequence, message := range messages {
content, err := json.Marshal(message)
if err != nil {
return nil, fmt.Errorf("marshal initial message: %w", err)
}
if _, err := q.InsertMessage(ctx, db.InsertMessageParams{
Content: content,
ConversationId: c.Id,
Role: db.MessageRole(message.Role),
Sequence: sequence,
}); err != nil {
return nil, fmt.Errorf("insert initial message: %w", err)
}
}
if err := tx.Commit(); err != nil {
return nil, fmt.Errorf(
"commit conversation transaction: %w", err,
)
}
return c, nil
}
// Cron returns the parsed cron schedule.
func (j *Job) Cron() cron.Schedule { return j.cron }
// insertMessage stores a message and its optional inference metadata in the
// same row at the next sequence in a conversation. Nil metadata stores NULL.
// If executionUpdate is non-nil, its status change and the message commit in
// the same transaction. A nil update leaves execution statuses unchanged.
// The caller is responsible for identifying the execution and choosing a
// valid status transition.
func (j *Job) insertMessage(
ctx context.Context,
cid string,
message llm.Message,
metadata *llm.Metadata,
executionUpdate *db.UpdateExecutionStatusParams,
) (*db.Message, error) {
content, err := json.Marshal(message)
if err != nil {
return nil, fmt.Errorf("marshal message: %w", err)
}
var metadataJSON jsontext.Value
if metadata != nil {
encoded, err := json.Marshal(
metadata, json.FormatNilSliceAsNull(true),
)
if err != nil {
return nil, fmt.Errorf("marshal metadata: %w", err)
}
metadataJSON = jsontext.Value(encoded)
}
tx, err := j.db.BeginTx(ctx, nil)
if err != nil {
return nil, fmt.Errorf("begin message transaction: %w", err)
}
defer tx.Rollback()
q := db.New(tx)
sequence, err := q.NextSequence(ctx, cid)
if err != nil {
return nil, fmt.Errorf("get next message sequence: %w", err)
}
m, err := q.InsertMessage(ctx, db.InsertMessageParams{
Content: content,
ConversationId: cid,
Metadata: metadataJSON,
Role: db.MessageRole(message.Role),
Sequence: sequence,
})
if err != nil {
return nil, fmt.Errorf("insert message: %w", err)
}
if executionUpdate != nil {
if err := q.UpdateExecutionStatus(ctx, *executionUpdate); err != nil {
return nil, fmt.Errorf("update execution status: %w", err)
}
}
if err := tx.Commit(); err != nil {
return nil, fmt.Errorf("commit message transaction: %w", err)
}
return m, nil
}
// MaxIterations returns the configured max iterations.
// MaxIterations returns the job's iteration limit. Zero means unlimited.
func (j *Job) MaxIterations() int { return j.maxIterations }
// Model returns the model ID used by the job.
// Model returns the model ID.
func (j *Job) Model() string { return j.model }
// Name returns the job name.
func (j *Job) Name() string { return j.name }
// Run executes the job's inference and tool-execution loop.
func (j *Job) Run(ctx context.Context) Result {
var start = time.Now()
// Run sends the job's task as a user message to c and runs inference
// until it returns. c must already hold the job's system message, if
// any; Run does not add it.
//
// The caller registers c before Run, so clients can subscribe to the
// job and stop it like any other conversation. Run returns
// [conversation.ErrBusy] if c is already reserved.
func (j *Job) Run(ctx context.Context, c *conversation.Conversation) error {
start := time.Now()
j.log.InfoContext(ctx, "starting job")
// Build messages from job config.
var messages []llm.Message
if j.SystemMessage() != "" {
messages = append(messages, llm.Message{
Role: llm.RoleSystem,
ContentParts: []llm.ContentPart{
{
Text: j.SystemMessage(),
Type: llm.ContentTypeText,
},
},
})
}
messages = append(messages, llm.Message{
Role: llm.RoleUser,
ContentParts: []llm.ContentPart{
{
Text: j.Task(),
inf, err := c.Reserve(ctx, j.llm, conversation.InferenceRequest{
Message: &llm.Message{
Role: llm.RoleUser,
ContentParts: []llm.ContentPart{{
Type: llm.ContentTypeText,
},
Text: j.task,
}},
},
Model: j.model,
MaxIterations: j.maxIterations,
Timeout: j.timeout,
})
// Track the end of the initial messages.
// Everything appended after is generated by the loop.
base := len(messages)
// Apply per-job timeout if set.
if j.Timeout() > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, j.Timeout())
defer cancel()
if err == nil {
err = inf.Run()
}
// Create a conversation for this run.
c, err := j.createConversation(ctx, messages)
duration := time.Since(start)
if err != nil {
duration := time.Since(start)
j.log.ErrorContext(
ctx,
"job failed",
slog.Any("error", err),
slog.Duration("duration", duration),
)
return Result{
Name: j.Name(),
Duration: duration,
Error: fmt.Errorf("job %q: %w", j.Name(), err),
}
return fmt.Errorf("job %q: %w", j.name, err)
}
log := j.log.With(slog.String("conversation_id", c.Id))
log.InfoContext(ctx, "created job conversation")
// Run the inference and tool-execution loop.
var completed bool
for i := 0; j.MaxIterations() == 0 || i < j.MaxIterations(); i++ {
var (
result *llm.CompletionResult
tm *db.Message
)
// Request the next completion from the model.
result, err = j.llm.Completions(
ctx,
llm.CompletionsParams{
Model: j.model,
Messages: messages,
Tools: j.tools.ToAPI(),
},
)
if err != nil {
break
}
// Persist the assistant message before acting on it,
// so an interruption leaves a committed record.
message := result.Message
tm, err = j.insertMessage(ctx, c.Id, message, &result.Metadata, nil)
if err != nil {
break
}
messages = append(messages, message)
// No tool calls means the model is done.
if len(message.ToolCalls) == 0 {
completed = true
break
}
// Execute every tool call in the batch, persisting each result.
toolResults, toolErr := j.callTools(
ctx, log, c.Id, tm.Id, message,
)
messages = append(messages, toolResults...)
if toolErr != nil {
err = toolErr
break
}
}
duration := time.Since(start)
if err == nil && !completed {
err = errMaxIterations
}
if err != nil {
log.ErrorContext(
ctx,
"job failed",
slog.Any("error", err),
slog.Duration("duration", duration),
)
return Result{
ConversationID: c.Id,
Name: j.Name(),
Duration: duration,
Error: fmt.Errorf("job %q: %w", j.Name(), err),
}
}
log.InfoContext(
ctx,
"job completed",
slog.Duration("duration", duration),
slog.Int("messages", len(messages)-base),
j.log.InfoContext(
ctx, "job completed", slog.Duration("duration", duration),
)
return Result{
ConversationID: c.Id,
Name: j.Name(),
Duration: duration,
Messages: messages[base:],
}
return nil
}
// Schedule returns the job schedule string.
@@ -585,3 +159,6 @@ func (j *Job) Task() string { return j.task }
// Timeout returns the job's timeout duration.
func (j *Job) Timeout() time.Duration { return j.timeout }
// UserId returns the owner of conversations created by the job.
func (j *Job) UserId() string { return j.userId }
+70 -504
View File
@@ -1,50 +1,31 @@
package job
import (
"database/sql"
"encoding/json/v2"
"fmt"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"sync/atomic"
"testing"
"time"
"code.chimeric.al/chimerical/odidere/internal/config"
"code.chimeric.al/chimerical/odidere/internal/conversation"
"code.chimeric.al/chimerical/odidere/internal/db"
"code.chimeric.al/chimerical/odidere/internal/llm"
"code.chimeric.al/chimerical/odidere/internal/tool"
)
func testDependencies(t *testing.T) (*db.SQLite, *llm.Client, *tool.Registry) {
func testDependencies(t *testing.T) *llm.Client {
t.Helper()
database, err := db.Connect(t.Context(), db.Config{
Path: filepath.Join(t.TempDir(), "job.db"),
})
if err != nil {
t.Fatalf("db.Connect: %v", err)
}
t.Cleanup(func() { database.Close() })
if err := database.ApplySchema(t.Context(), slog.Default()); err != nil {
t.Fatalf("ApplySchema: %v", err)
}
client, err := llm.NewClient(
llm.Config{URL: "http://localhost"}, slog.Default(),
)
if err != nil {
t.Fatalf("llm.NewClient: %v", err)
}
registry, err := tool.NewRegistry(nil)
if err != nil {
t.Fatalf("tool.NewRegistry: %v", err)
}
return database, client, registry
return client
}
func TestConfigValidate(t *testing.T) {
@@ -169,7 +150,7 @@ func TestConfigTimeout(t *testing.T) {
}
func TestNew(t *testing.T) {
database, client, registry := testDependencies(t)
client := testDependencies(t)
cfg := config.JobConfig{
Name: "test",
Task: "do something",
@@ -181,7 +162,7 @@ func TestNew(t *testing.T) {
cfg.Model.Model = "test-model"
cfg.UserId = "default"
j, err := New(cfg, database.Writer, client, slog.Default(), registry)
j, err := New(cfg, client, slog.Default())
if err != nil {
t.Fatalf("New() error = %v", err)
}
@@ -213,7 +194,7 @@ func TestNew(t *testing.T) {
}
func TestNew_InvalidSchedule(t *testing.T) {
database, client, registry := testDependencies(t)
client := testDependencies(t)
cfg := config.JobConfig{
Name: "test",
Task: "do something",
@@ -222,14 +203,14 @@ func TestNew_InvalidSchedule(t *testing.T) {
cfg.Model.Model = "test-model"
cfg.UserId = "default"
_, err := New(cfg, database.Writer, client, slog.Default(), registry)
_, err := New(cfg, client, slog.Default())
if err == nil {
t.Fatal("New() expected error for invalid schedule")
}
}
func TestNew_InvalidTimeout(t *testing.T) {
database, client, registry := testDependencies(t)
client := testDependencies(t)
cfg := config.JobConfig{
Name: "test",
Task: "do something",
@@ -239,69 +220,44 @@ func TestNew_InvalidTimeout(t *testing.T) {
cfg.Model.Model = "test-model"
cfg.UserId = "default"
_, err := New(cfg, database.Writer, client, slog.Default(), registry)
_, err := New(cfg, client, slog.Default())
if err == nil {
t.Fatal("New() expected error for invalid timeout")
}
}
func TestNewRequiresDependencies(t *testing.T) {
database, client, registry := testDependencies(t)
client := testDependencies(t)
cfg := config.JobConfig{
Name: "test",
Task: "do something",
Schedule: "0 9 * * *",
}
for _, test := range []struct {
name string
registry *tool.Registry
database *sql.DB
client *llm.Client
logger *slog.Logger
userID string
name string
client *llm.Client
logger *slog.Logger
userID string
}{
{
name: "nil registry",
database: database.Writer,
client: client,
logger: slog.Default(),
userID: "default",
name: "empty user ID",
client: client,
logger: slog.Default(),
},
{
name: "nil database",
registry: registry,
client: client,
logger: slog.Default(),
userID: "default",
name: "nil llm client",
logger: slog.Default(),
userID: "default",
},
{
name: "empty user ID",
registry: registry,
database: database.Writer,
client: client,
logger: slog.Default(),
},
{
name: "nil llm client",
registry: registry,
database: database.Writer,
logger: slog.Default(),
userID: "default",
},
{
name: "nil logger",
registry: registry,
database: database.Writer,
client: client,
userID: "default",
name: "nil logger",
client: client,
userID: "default",
},
} {
t.Run(test.name, func(t *testing.T) {
cfg.UserId = test.userID
_, err := New(
cfg, test.database, test.client, test.logger,
test.registry,
)
_, err := New(cfg, test.client, test.logger)
if err == nil {
t.Fatal("New() expected error")
}
@@ -309,471 +265,81 @@ func TestNewRequiresDependencies(t *testing.T) {
}
}
func TestRunPersistsToolLoop(t *testing.T) {
var requests atomic.Int32
// TestRun checks that Run saves the task as a user message in the
// conversation and runs inference to completion.
func TestRun(t *testing.T) {
var (
ctx = t.Context()
log = slog.New(slog.DiscardHandler)
)
inference := httptest.NewServer(http.HandlerFunc(func(
w http.ResponseWriter, r *http.Request,
) {
var request llm.ChatRequest
if err := json.UnmarshalRead(r.Body, &request); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
w.Header().Set("Content-Type", "application/json")
switch requests.Add(1) {
case 1:
if len(request.Messages) != 2 ||
request.Messages[0].Role != llm.RoleSystem ||
request.Messages[1].Role != llm.RoleUser {
http.Error(w, "unexpected initial history", http.StatusBadRequest)
return
}
fmt.Fprint(w, `{
"choices":[{
"message":{
"role":"assistant",
"tool_calls":[{
"id":"call_test",
"type":"function",
"function":{"name":"check_db","arguments":"{}"}
}]
},
"finish_reason":"tool_calls"
}]
}`)
case 2:
if len(request.Messages) != 4 ||
request.Messages[2].Role != llm.RoleAssistant ||
request.Messages[3].Role != llm.RoleTool ||
request.Messages[3].ToolCallID != "call_test" {
http.Error(w, "unexpected tool history", http.StatusBadRequest)
return
}
fmt.Fprint(w, `{
"choices":[{
"message":{"role":"assistant","content":"finished"},
"finish_reason":"stop"
}]
}`)
default:
http.Error(w, "unexpected inference request", http.StatusBadRequest)
}
fmt.Fprint(w, `{
"choices":[{
"message":{"role":"assistant","content":"finished"},
"finish_reason":"stop"
}]
}`)
}))
defer inference.Close()
var (
ctx = t.Context()
log = slog.New(slog.NewTextHandler(io.Discard, nil))
databasePath = filepath.Join(t.TempDir(), "job.db")
)
database, err := db.Connect(ctx, db.Config{Path: databasePath})
database, err := db.Connect(ctx, db.Config{
Path: filepath.Join(t.TempDir(), "job.db"),
})
if err != nil {
t.Fatalf("db.Connect: %v", err)
}
defer database.Close()
if err := database.ApplySchema(ctx, slog.Default()); err != nil {
if err := database.ApplySchema(ctx, log); err != nil {
t.Fatalf("ApplySchema: %v", err)
}
t.Setenv("ODIDERE_JOB_TEST_DATABASE", databasePath)
executable, err := os.Executable()
saved, err := db.New(database.Writer).InsertConversation(
ctx, db.InsertConversationParams{UserId: "default"},
)
if err != nil {
t.Fatalf("os.Executable: %v", err)
t.Fatalf("InsertConversation: %v", err)
}
registry, err := tool.NewRegistry([]config.ToolConfig{{
Name: "check_db",
Description: "Verify job persistence",
Command: executable,
Arguments: []string{"-test.run=^TestJobPersistenceHelper$"},
Parameters: map[string]any{"type": "object"},
}})
registry, err := tool.NewRegistry(nil)
if err != nil {
t.Fatalf("tool.NewRegistry: %v", err)
}
client, err := llm.NewClient(llm.Config{
Provider: "test",
URL: inference.URL,
}, log)
c, err := conversation.New(
conversation.Config{Id: saved.Id}, database, log, registry,
)
if err != nil {
t.Fatalf("conversation.New: %v", err)
}
client, err := llm.NewClient(
llm.Config{Provider: "test", URL: inference.URL}, log,
)
if err != nil {
t.Fatalf("llm.NewClient: %v", err)
}
job, err := New(
config.JobConfig{
Name: "daily report",
Task: "prepare the report",
Schedule: "0 9 * * *",
SystemMessage: "system",
Model: config.ModelConfig{Model: "test-model"},
UserId: "default",
},
database.Writer,
client,
log,
registry,
)
j, err := New(config.JobConfig{
Name: "daily report",
Task: "prepare the report",
Schedule: "0 9 * * *",
Model: config.ModelConfig{Model: "test-model"},
UserId: "default",
}, client, log)
if err != nil {
t.Fatalf("New: %v", err)
}
result := job.Run(ctx)
if result.Error != nil {
t.Fatalf("Run: %v", result.Error)
}
if result.ConversationID == "" {
t.Fatal("conversation ID is empty")
}
if requests.Load() != 2 {
t.Fatalf("inference requests = %d, want 2", requests.Load())
if err := j.Run(ctx, c); err != nil {
t.Fatalf("Run: %v", err)
}
q := db.New(database.Reader)
conversation, err := q.SelectConversationId(ctx, result.ConversationID)
if err != nil {
t.Fatalf("SelectConversationId: %v", err)
}
if conversation.UserId != "default" {
t.Errorf("conversation user = %q, want default", conversation.UserId)
}
if conversation.Title == nil || *conversation.Title != "daily report" {
t.Errorf("conversation title = %v, want daily report", conversation.Title)
}
stored, err := q.SelectMessagesConversationId(ctx, result.ConversationID)
rows, err := db.New(database.Reader).
SelectMessagesConversationId(ctx, saved.Id)
if err != nil {
t.Fatalf("SelectMessagesConversationId: %v", err)
}
wantRoles := []db.MessageRole{
db.MessageRoleSystem,
db.MessageRoleUser,
db.MessageRoleAssistant,
db.MessageRoleTool,
db.MessageRoleAssistant,
}
if len(stored) != len(wantRoles) {
t.Fatalf("stored messages = %d, want %d", len(stored), len(wantRoles))
}
for i, message := range stored {
if message.Role != db.MessageRoleAssistant {
if message.Metadata != nil {
t.Errorf("message %d metadata = %v, want nil", i, message.Metadata)
}
continue
}
if message.Metadata == nil {
t.Fatalf("assistant %d has no metadata", i)
}
var metadata llm.Metadata
if err := json.Unmarshal(message.Metadata, &metadata); err != nil {
t.Fatal(err)
}
if metadata.Provider == nil || *metadata.Provider != "test" ||
metadata.RequestedModel == nil || *metadata.RequestedModel != "test-model" {
t.Errorf("assistant %d metadata = %s", i, string(message.Metadata))
}
wantReason := "stop"
if i == 2 {
wantReason = "tool_calls"
}
if metadata.FinishReason == nil || *metadata.FinishReason != wantReason {
t.Errorf("assistant %d finish reason = %v", i, metadata.FinishReason)
}
}
for i, want := range wantRoles {
if stored[i].Sequence != i || stored[i].Role != want {
t.Errorf(
"stored message %d = sequence %d, role %q; want %d, %q",
i, stored[i].Sequence, stored[i].Role, i, want,
)
}
}
}
func TestJobPersistenceHelper(t *testing.T) {
path := os.Getenv("ODIDERE_JOB_TEST_DATABASE")
if path == "" {
return
}
database, err := sql.Open("sqlite3", "file:"+path+"?mode=ro")
if err != nil {
t.Fatalf("sql.Open: %v", err)
}
defer database.Close()
var role string
if err := database.QueryRow(
"SELECT role FROM messages ORDER BY created_at DESC, sequence DESC LIMIT 1",
).Scan(&role); err != nil {
t.Fatalf("select latest message: %v", err)
}
if role != string(db.MessageRoleAssistant) {
t.Fatalf("latest message role = %q, want assistant", role)
}
}
func TestCallTool(t *testing.T) {
database, client, _ := testDependencies(t)
registry, err := tool.NewRegistry([]config.ToolConfig{
{
Name: "finish",
Description: "finish",
Command: "sh",
Arguments: []string{"-c", "echo done"},
},
{
Name: "fail",
Description: "fail",
Command: "sh",
Arguments: []string{"-c", "echo oops; exit 3"},
},
{
Name: "missing",
Description: "missing",
Command: "nonexistent_binary_xyz",
},
{
Name: "interrupt",
Description: "interrupt",
Command: "sh",
// exec replaces the shell, so the timeout kills the
// process that holds the output pipe.
Arguments: []string{"-c", "echo partial; exec sleep 10"},
Timeout: 100 * time.Millisecond,
},
})
if err != nil {
t.Fatalf("tool.NewRegistry: %v", err)
}
log := slog.New(slog.DiscardHandler)
job, err := New(
config.JobConfig{
Name: "call tool",
Task: "call a tool",
Schedule: "0 9 * * *",
UserId: "default",
},
database.Writer,
client,
log,
registry,
)
if err != nil {
t.Fatalf("New: %v", err)
}
conversation, err := job.createConversation(t.Context(), nil)
if err != nil {
t.Fatalf("createConversation: %v", err)
}
tests := []struct {
name string
tool string
// want lists text the message must contain, in this order.
want []string
// warn reports whether the message must carry
// interruptedWarning.
warn bool
// status is the execution row's status
// after callTool returns.
status db.ExecutionStatus
}{
{
name: "finished",
tool: "finish",
want: []string{"▸ stdout:\ndone\n", "┗ exit: 0"},
status: db.ExecutionStatusCompleted,
},
{
name: "non-zero exit",
tool: "fail",
want: []string{"▸ stdout:\noops\n", "┗ exit: 3"},
status: db.ExecutionStatusCompleted,
},
{
name: "never started",
tool: "missing",
want: []string{"executable file not found"},
status: db.ExecutionStatusCompleted,
},
{
name: "interrupted",
tool: "interrupt",
want: []string{
"timed out",
interruptedWarning,
"▸ stdout:\npartial\n",
},
warn: true,
status: db.ExecutionStatusRunning,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
call := llm.ToolCall{
Function: llm.Function{Arguments: "{}", Name: tt.tool},
ID: "call_" + tt.tool,
Type: llm.ToolTypeFunction,
}
assistant, err := job.insertMessage(
t.Context(), conversation.Id, llm.Message{
Role: llm.RoleAssistant,
ToolCalls: []llm.ToolCall{call},
},
&llm.Metadata{ID: new("assistant-completion")},
nil,
)
if err != nil {
t.Fatalf("insertMessage: %v", err)
}
message, err := job.callTool(
t.Context(), log, conversation.Id, assistant.Id, call,
)
if err != nil {
t.Fatalf("callTool: %v", err)
}
if message.Role != llm.RoleTool ||
message.ToolCallID != call.ID ||
message.Name != tt.tool {
t.Errorf(
"message = role %q, tool call %q, name %q; "+
"want %q, %q, %q",
message.Role, message.ToolCallID, message.Name,
llm.RoleTool, call.ID, tt.tool,
)
}
content := message.ContentParts[0].Text
rest := content
for _, want := range tt.want {
i := strings.Index(rest, want)
if i < 0 {
t.Fatalf(
"content missing %q, or out of order:\n%s",
want, content,
)
}
rest = rest[i+len(want):]
}
if got := strings.Contains(
content, interruptedWarning,
); got != tt.warn {
t.Errorf(
"content has warning = %t, want %t:\n%s",
got, tt.warn, content,
)
}
stored, err := db.New(database.Reader).
SelectMessagesConversationId(
t.Context(), conversation.Id,
)
if err != nil {
t.Fatalf("SelectMessagesConversationId: %v", err)
}
want, err := json.Marshal(*message)
if err != nil {
t.Fatalf("json.Marshal: %v", err)
}
wantMetadata := `{"choice_index":null,"created":null,` +
`"finish_reason":null,"id":"assistant-completion",` +
`"logprobs":null,"model":null,"provider":null,` +
`"requested_model":null,"service_tier":null,` +
`"system_fingerprint":null,"timings":null,"usage":null}`
if assistant.Metadata == nil ||
string(assistant.Metadata) != wantMetadata {
t.Errorf("assistant metadata = %v", assistant.Metadata)
}
last := stored[len(stored)-1]
if last.Metadata != nil {
t.Errorf("tool metadata = %v, want nil", last.Metadata)
}
if last.Role != db.MessageRoleTool ||
string(last.Content) != string(want) {
t.Errorf(
"last stored message = role %q, %s; "+
"want %q, %s",
last.Role, last.Content,
db.MessageRoleTool, want,
)
}
executions, err := db.New(database.Reader).
SelectExecutionsMessageId(t.Context(), assistant.Id)
if err != nil {
t.Fatalf("SelectExecutionsMessageId: %v", err)
}
if len(executions) != 1 ||
executions[0].ToolCallId != call.ID ||
executions[0].Status != tt.status {
t.Errorf(
"executions = %+v, want one %q row for %q",
executions, tt.status, call.ID,
)
}
})
}
}
// TestCallToolUnsavedRow checks that callTool does not run the tool when
// its executions row cannot be saved. The row references the assistant
// message, so a message ID that does not exist makes the save fail.
func TestCallToolUnsavedRow(t *testing.T) {
database, client, _ := testDependencies(t)
marker := filepath.Join(t.TempDir(), "ran")
t.Setenv("ODIDERE_CALL_TOOL_MARKER", marker)
registry, err := tool.NewRegistry([]config.ToolConfig{{
Name: "touch",
Description: "touch",
Command: "sh",
Arguments: []string{"-c", `touch "$ODIDERE_CALL_TOOL_MARKER"`},
}})
if err != nil {
t.Fatalf("tool.NewRegistry: %v", err)
}
log := slog.New(slog.DiscardHandler)
job, err := New(
config.JobConfig{
Name: "unsaved row",
Task: "call a tool",
Schedule: "0 9 * * *",
UserId: "default",
},
database.Writer,
client,
log,
registry,
)
if err != nil {
t.Fatalf("New: %v", err)
}
conversation, err := job.createConversation(t.Context(), nil)
if err != nil {
t.Fatalf("createConversation: %v", err)
}
message, err := job.callTool(
t.Context(), log, conversation.Id, "missing",
llm.ToolCall{
Function: llm.Function{
Arguments: "{}", Name: "touch",
},
ID: "call_touch",
Type: llm.ToolTypeFunction,
},
)
if err == nil {
t.Fatal("callTool: expected error")
}
if message != nil {
t.Errorf("message = %+v, want nil", message)
}
if _, err := os.Stat(marker); err == nil {
t.Error("tool ran without a saved row")
}
executions, err := db.New(database.Reader).
SelectExecutionsMessageId(t.Context(), "missing")
if err != nil {
t.Fatalf("SelectExecutionsMessageId: %v", err)
}
if len(executions) != 0 {
t.Errorf("executions = %+v, want none", executions)
if len(rows) != 2 ||
rows[0].Role != db.MessageRoleUser ||
rows[1].Role != db.MessageRoleAssistant {
t.Fatalf("saved rows = %v, want user then assistant", rows)
}
}