eric/audita
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 6 Wiki Activity

Preserve cache settings in repair conversation replay

Browse Source
This commit is contained in:
Eric Rakestraw
2026-11-11 12:26:27 +00:00
parent 8267eba6a5
commit 45b42a2689
4 changed files with 37 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
internal/framework/llm/adapter.go
View File

@@ -123,7 +123,11 @@ func (c *RoleClient) CompleteStructured(parent context.Context, req structuredou
baseID, baseVars = exchange.PromptID, exchange.Variables
suffix = append(suffix, pk.RenderedMessage{Role: pk.RoleAssistant, Content: exchange.FinalContent})
for _, message := range req.Replay.Turns {
suffix = append(suffix, pk.RenderedMessage{Role: message.Role, Content: message.Content})
native := pk.RenderedMessage{Role: message.Role, Content: message.Content}
if cache := message.CacheControl; cache != nil {
native.CacheControl = &pk.CacheControl{Type: pk.CacheControlType(cache.Type), TTL: cache.TTL}
}
suffix = append(suffix, native)
}
suffix = append(suffix, instruction.Messages...)
}
@@ -316,6 +320,9 @@ func convertMessages(messages []pk.RenderedMessage) []structuredoutput.Message {
result := make([]structuredoutput.Message, len(messages))
for i, message := range messages {
result[i] = structuredoutput.Message{Role: string(message.Role), Content: message.Content}
if cache := message.CacheControl; cache != nil {
result[i].CacheControl = &structuredoutput.CacheControl{Type: string(cache.Type), TTL: cache.TTL}
}
}
return result
}

26
internal/framework/llm/conversation_test.go
View File

@@ -157,7 +157,7 @@ func TestRoleClientCapturesAndReplaysRepairConversation(t *testing.T) {
dir := t.TempDir()
for name, content := range map[string]string{
"config.yml": "version: 3\n",
"prompts.yml": "profiles:\n - id: proposal\n endpoint: https://example.invalid/v1\n model: proposal\n - id: repair\n endpoint: https://example.invalid/v1\n model: repair\n",
"prompts.yml": "profiles:\n - id: proposal\n endpoint: https://example.invalid/v1\n model: proposal\n - id: repair\n endpoint: https://example.invalid/v1\n model: repair\ncollections:\n audita:\n prompts:\n - id: modules.grammar.repair\n messages:\n - role: user\n content: '{{ .RepairRequestJSON }}'\n cache_control: {type: ephemeral, ttl: 1h}\n output:\n format: json\n validation_mode: json_schema\n schema_path: correction_repair.v1.json\n",
} {
if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0644); err != nil {
t.Fatal(err)
@@ -167,13 +167,17 @@ func TestRoleClientCapturesAndReplaysRepairConversation(t *testing.T) {
if err != nil {
t.Fatal(err)
}
var messages [][]structuredoutput.Message
type wireMessage struct {
Role string
Content json.RawMessage
}
var messages [][]wireMessage
var models []string
outputs := []string{`{"corrections":[]}`, `{"proposal_ref":1,"action":"revise","correction":{"id":1,"original_text":"a","corrected_text":"b","confidence":0.8}}`, `{"proposal_ref":1,"action":"withdraw","correction":null}`, `{"proposal_ref":1,"action":"withdraw","correction":null}`}
transport := transportFunc(func(r *http.Request) (*http.Response, error) {
var req struct {
Model string `json:"model"`
Messages []structuredoutput.Message `json:"messages"`
Model string `json:"model"`
Messages []wireMessage `json:"messages"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Fatal(err)
@@ -218,15 +222,21 @@ func TestRoleClientCapturesAndReplaysRepairConversation(t *testing.T) {
if !reflect.DeepEqual(models, []string{"proposal", "repair", "repair"}) {
t.Fatalf("models = %v", models)
}
prefix := initial.Exchange.Messages
if len(messages) != 3 || !reflect.DeepEqual(messages[0], prefix) || !reflect.DeepEqual(messages[1][:len(prefix)], prefix) || !reflect.DeepEqual(messages[2][:len(prefix)], prefix) {
prefix := messages[0]
if len(messages) != 3 || !reflect.DeepEqual(messages[1][:len(prefix)], prefix) || !reflect.DeepEqual(messages[2][:len(prefix)], prefix) {
t.Fatalf("prefixes: %+v", messages)
}
if len(messages[1]) != len(prefix)+2 || len(messages[2]) != len(prefix)+4 || messages[2][len(prefix)+2].Content != outputs[1] {
encodedRevision, _ := json.Marshal(outputs[1])
if len(messages[1]) != len(prefix)+2 || len(messages[2]) != len(prefix)+4 || string(messages[2][len(prefix)+2].Content) != string(encodedRevision) {
t.Fatalf("turns: %+v", messages)
}
first := messages[1][len(prefix)+1]
if first.Role != "user" || !strings.Contains(string(first.Content), `"ttl":"1h"`) ||
!strings.Contains(string(first.Content), `"text":"first"`) || !reflect.DeepEqual(first, messages[2][len(prefix)+1]) {
t.Fatalf("cached instruction changed during replay: first=%s replay=%s", first.Content, messages[2][len(prefix)+1].Content)
}
_, err = repair.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarRepair, Schema: responseschema.CorrectionRepairKey, Variables: map[string]string{"RepairRequestJSON": "independent"}, Replay: &structuredoutput.Replay{Exchange: initial.Exchange}}, &map[string]any{})
if err != nil || len(messages) != 4 || len(messages[3]) != len(prefix)+2 || !strings.Contains(messages[3][len(prefix)+1].Content, "independent") {
if err != nil || len(messages) != 4 || len(messages[3]) != len(prefix)+2 || !strings.Contains(string(messages[3][len(prefix)+1].Content), "independent") {
t.Fatalf("independent branch: %v %+v", err, messages)
}
bad := *initial.Exchange