Include function call history in restored threads.

This should discourage the model from lying about doing
things the second time they are requested in a thread, by
demonstrating to it how it handled the request the last time.

Signed-off-by: Katharine Berry <[email protected]>

Author: Katharine

service/assistant/functions/functions.go

@@ -46,6 +46,8 @@ type Registration struct {
 	Cb CallbackFunction
 	// A function that summarises the provided input
 	Thought ThoughtFunction
+	// Whether to redact the function output in chat history.
+	RedactOutputInChatHistory bool
 	// An instance of the object used to hold the function's parameters. This is what will be passed to
 	// either Fn or Cb, and it will also be processed to pass to the model - including the comments.
 	InputType interface{}
@@ -231,3 +233,12 @@ func GetFunctionDefinitionsForCapabilities(capabilities []string) []*genai.Funct
 	}
 	return definitions
 }
+func GetFunctionRegistration(fn string) *Registration {
+	if realFunction, ok := functionAliases[fn]; ok {
+		fn = realFunction
+	}
+	if reg, ok := functionMap[fn]; ok {
+		return &reg
+	}
+	return nil
+}

service/assistant/functions/wikipedia.go

@@ -74,9 +74,10 @@ func init() {
 				Required: []string{"wiki", "article_name"},
 			},
 		},
-		Fn:        queryWiki,
-		Thought:   queryWikiThought,
-		InputType: WikiRequest{},
+		Fn:                        queryWiki,
+		Thought:                   queryWikiThought,
+		RedactOutputInChatHistory: true,
+		InputType:                 WikiRequest{},
 	})
 }
 

service/assistant/session.go

@@ -380,20 +380,38 @@ func (ps *PromptSession) Run(ctx context.Context) {
 }
 
 type SerializedMessage struct {
-	Role    string `json:"role"`
-	Content string `json:"content"`
+	Role             string                  `json:"role"`
+	Content          string                  `json:"content"`
+	FunctionCall     *genai.FunctionCall     `json:"functionCall,omitempty"`
+	FunctionResponse *genai.FunctionResponse `json:"functionResponse,omitempty"`
 }
 
 func (ps *PromptSession) storeThread(ctx context.Context, messages []*genai.Content) error {
 	ctx, span := beeline.StartSpan(ctx, "store_thread")
 	defer span.Send()
 	var toStore []SerializedMessage
 	for _, m := range messages {
-		if len(m.Parts) != 0 && (m.Role == "user" || m.Role == "model") && len(strings.TrimSpace(m.Parts[0].Text)) > 0 {
-			toStore = append(toStore, SerializedMessage{
-				Content: m.Parts[0].Text,
-				Role:    m.Role,
-			})
+		if len(m.Parts) != 0 {
+			if m.Role == "user" || m.Role == "model" {
+				sm := SerializedMessage{
+					Role:         m.Role,
+					Content:      m.Parts[0].Text,
+					FunctionCall: m.Parts[0].FunctionCall,
+				}
+				if sm.FunctionCall != nil || len(strings.TrimSpace(m.Parts[0].Text)) > 0 {
+					toStore = append(toStore, sm)
+				}
+			} else if m.Role == "function" && m.Parts[0].FunctionResponse != nil {
+				fr := *m.Parts[0].FunctionResponse
+				fnInfo := functions.GetFunctionRegistration(fr.Name)
+				if fnInfo != nil && fnInfo.RedactOutputInChatHistory {
+					fr.Response = map[string]any{"redacted": "redacted to reduce context size, call again if necessary"}
+				}
+				toStore = append(toStore, SerializedMessage{
+					Role:             m.Role,
+					FunctionResponse: &fr,
+				})
+			}
 		}
 	}
 	j, err := json.Marshal(toStore)
@@ -421,7 +439,7 @@ func (ps *PromptSession) restoreThread(ctx context.Context, oldThreadId string)
 	var result []*genai.Content
 	for _, m := range messages {
 		result = append(result, &genai.Content{
-			Parts: []*genai.Part{{Text: m.Content}},
+			Parts: []*genai.Part{{Text: m.Content, FunctionCall: m.FunctionCall, FunctionResponse: m.FunctionResponse}},
 			Role:  m.Role,
 		})
 	}