Add job pkg metadata
This commit is contained in:
+17
-4
@@ -8,6 +8,7 @@ package job
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"database/sql"
|
"database/sql"
|
||||||
|
"encoding/json/jsontext"
|
||||||
"encoding/json/v2"
|
"encoding/json/v2"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
@@ -254,7 +255,7 @@ func (j *Job) callTool(
|
|||||||
context.WithoutCancel(ctx), resultSaveTimeout,
|
context.WithoutCancel(ctx), resultSaveTimeout,
|
||||||
)
|
)
|
||||||
defer saveCancel()
|
defer saveCancel()
|
||||||
if _, err := j.insertMessage(saveCtx, cid, message, update); err != nil {
|
if _, err := j.insertMessage(saveCtx, cid, message, nil, update); err != nil {
|
||||||
log.ErrorContext(
|
log.ErrorContext(
|
||||||
ctx,
|
ctx,
|
||||||
"failed to persist tool result",
|
"failed to persist tool result",
|
||||||
@@ -370,8 +371,8 @@ func (j *Job) createConversation(
|
|||||||
// Cron returns the parsed cron schedule.
|
// Cron returns the parsed cron schedule.
|
||||||
func (j *Job) Cron() cron.Schedule { return j.cron }
|
func (j *Job) Cron() cron.Schedule { return j.cron }
|
||||||
|
|
||||||
// insertMessage stores one message in a conversation, after every
|
// insertMessage stores a message and its optional inference metadata in the
|
||||||
// message already stored there.
|
// 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
|
// If executionUpdate is non-nil, its status change and the message commit in
|
||||||
// the same transaction. A nil update leaves execution statuses unchanged.
|
// the same transaction. A nil update leaves execution statuses unchanged.
|
||||||
// The caller is responsible for identifying the execution and choosing a
|
// The caller is responsible for identifying the execution and choosing a
|
||||||
@@ -380,12 +381,23 @@ func (j *Job) insertMessage(
|
|||||||
ctx context.Context,
|
ctx context.Context,
|
||||||
cid string,
|
cid string,
|
||||||
message llm.Message,
|
message llm.Message,
|
||||||
|
metadata *llm.Metadata,
|
||||||
executionUpdate *db.UpdateExecutionStatusParams,
|
executionUpdate *db.UpdateExecutionStatusParams,
|
||||||
) (*db.Message, error) {
|
) (*db.Message, error) {
|
||||||
content, err := json.Marshal(message)
|
content, err := json.Marshal(message)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("marshal message: %w", err)
|
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)
|
tx, err := j.db.BeginTx(ctx, nil)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -401,6 +413,7 @@ func (j *Job) insertMessage(
|
|||||||
m, err := q.InsertMessage(ctx, db.InsertMessageParams{
|
m, err := q.InsertMessage(ctx, db.InsertMessageParams{
|
||||||
Content: string(content),
|
Content: string(content),
|
||||||
ConversationId: cid,
|
ConversationId: cid,
|
||||||
|
Metadata: metadataJSON,
|
||||||
Role: db.MessageRole(message.Role),
|
Role: db.MessageRole(message.Role),
|
||||||
Sequence: sequence,
|
Sequence: sequence,
|
||||||
})
|
})
|
||||||
@@ -507,7 +520,7 @@ func (j *Job) Run(ctx context.Context) Result {
|
|||||||
// Persist the assistant message before acting on it,
|
// Persist the assistant message before acting on it,
|
||||||
// so an interruption leaves a committed record.
|
// so an interruption leaves a committed record.
|
||||||
message := result.Message
|
message := result.Message
|
||||||
tm, err = j.insertMessage(ctx, c.Id, message, nil)
|
tm, err = j.insertMessage(ctx, c.Id, message, &result.Metadata, nil)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -450,6 +450,32 @@ func TestRunPersistsToolLoop(t *testing.T) {
|
|||||||
if len(stored) != len(wantRoles) {
|
if len(stored) != len(wantRoles) {
|
||||||
t.Fatalf("stored messages = %d, want %d", 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 {
|
for i, want := range wantRoles {
|
||||||
if stored[i].Sequence != i || stored[i].Role != want {
|
if stored[i].Sequence != i || stored[i].Role != want {
|
||||||
t.Errorf(
|
t.Errorf(
|
||||||
@@ -590,6 +616,7 @@ func TestCallTool(t *testing.T) {
|
|||||||
Role: llm.RoleAssistant,
|
Role: llm.RoleAssistant,
|
||||||
ToolCalls: []llm.ToolCall{call},
|
ToolCalls: []llm.ToolCall{call},
|
||||||
},
|
},
|
||||||
|
&llm.Metadata{ID: new("assistant-completion")},
|
||||||
nil,
|
nil,
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -644,7 +671,19 @@ func TestCallTool(t *testing.T) {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("json.Marshal: %v", err)
|
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]
|
last := stored[len(stored)-1]
|
||||||
|
if last.Metadata != nil {
|
||||||
|
t.Errorf("tool metadata = %v, want nil", last.Metadata)
|
||||||
|
}
|
||||||
if last.Role != db.MessageRoleTool ||
|
if last.Role != db.MessageRoleTool ||
|
||||||
last.Content != string(want) {
|
last.Content != string(want) {
|
||||||
t.Errorf(
|
t.Errorf(
|
||||||
|
|||||||
Reference in New Issue
Block a user