agenttest/scripted_language_model.go
Rama main · réplica de lectura del repositorio original.
package agenttest
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
"sync"
"github.com/yeomyeonggeori/blueprotocol/model"
)
type ScriptedLanguageModelOptions struct {
ActionResponses []string
ChatResponsesBySchema map[string][]string
StructuredResponsesBySchema map[string][]string
DefaultResponsesBySchema map[string]string
ProviderName string
ModelName string
}
type ScriptedLanguageModel struct {
mutex sync.Mutex
actionResponses []string
chatResponsesBySchema map[string][]string
structuredResponsesBySchema map[string][]string
defaultResponsesBySchema map[string]string
requests []model.StructuredResponseRequest
providerName string
modelName string
actionsServed int
onActionServed func(actionsServed int)
}
type scriptedChatCompleter struct {
languageModel *ScriptedLanguageModel
}
func NewScriptedLanguageModel(options ScriptedLanguageModelOptions) *ScriptedLanguageModel {
return &ScriptedLanguageModel{
actionResponses: append([]string{}, options.ActionResponses...),
chatResponsesBySchema: copyResponseQueues(options.ChatResponsesBySchema),
structuredResponsesBySchema: copyResponseQueues(options.StructuredResponsesBySchema),
defaultResponsesBySchema: mergeDefaultResponses(options.DefaultResponsesBySchema),
providerName: firstNonEmpty(options.ProviderName, "test"),
modelName: firstNonEmpty(options.ModelName, "scripted"),
}
}
func (languageModel *ScriptedLanguageModel) TextChatCompleter() (model.ChatCompleter, bool) {
languageModel.mutex.Lock()
defer languageModel.mutex.Unlock()
if len(languageModel.chatResponsesBySchema) == 0 {
return nil, false
}
return scriptedChatCompleter{languageModel: languageModel}, true
}
func (languageModel *ScriptedLanguageModel) ChatCompleter() model.ChatCompleter {
return scriptedChatCompleter{languageModel: languageModel}
}
func (completer scriptedChatCompleter) GenerateChatCompletion(_ context.Context, request model.ChatCompletionRequest) (model.ChatCompletionResponse, error) {
languageModel := completer.languageModel
languageModel.mutex.Lock()
defer languageModel.mutex.Unlock()
schemaName := strings.TrimSpace(request.SchemaName)
languageModel.requests = append(languageModel.requests, structuredRequestFromChat(request))
if schemaName == "bluecollar_agent_turn_action" {
response, errorValue := languageModel.popActionResponse()
if errorValue != nil {
return model.ChatCompletionResponse{}, errorValue
}
return languageModel.actionChatResponse(request, response)
}
responses := languageModel.chatResponsesBySchema[schemaName]
if len(responses) == 0 {
return model.ChatCompletionResponse{}, fmt.Errorf("scripted language model has no %s chat response", request.SchemaName)
}
languageModel.chatResponsesBySchema[schemaName] = responses[1:]
return model.ChatCompletionResponse{
FinishReason: "stop",
ProviderName: languageModel.providerName,
ModelName: languageModel.modelName,
SelectedBackend: "device",
Message: model.ChatCompletionMessage{
Role: "assistant",
Content: responses[0],
},
}, nil
}
func (languageModel *ScriptedLanguageModel) actionChatResponse(request model.ChatCompletionRequest, content string) (model.ChatCompletionResponse, error) {
var actionDocument struct {
Action string `json:"action"`
ToolName string `json:"toolName"`
ToolInput json.RawMessage `json:"toolInput"`
}
if errorValue := json.Unmarshal([]byte(content), &actionDocument); errorValue != nil {
return model.ChatCompletionResponse{}, errorValue
}
toolName := strings.TrimSpace(actionDocument.Action)
arguments := json.RawMessage(content)
if toolName == "continue" {
toolName = strings.TrimSpace(actionDocument.ToolName)
arguments = actionDocument.ToolInput
}
if toolName == "" || len(arguments) == 0 {
return model.ChatCompletionResponse{}, errors.New("scripted agent action is incomplete")
}
if !chatRequestHasTool(request, toolName) {
return model.ChatCompletionResponse{}, fmt.Errorf("scripted agent action tool %q is not exposed; available tools: %s", toolName, strings.Join(chatRequestToolNames(request), ", "))
}
arguments = removeActionDiscriminator(arguments)
return languageModel.toolCallResponse(toolName, arguments), nil
}
func chatRequestHasTool(request model.ChatCompletionRequest, toolName string) bool {
for _, availableToolName := range chatRequestToolNames(request) {
if availableToolName == toolName {
return true
}
}
return false
}
func chatRequestToolNames(request model.ChatCompletionRequest) []string {
toolNames := make([]string, 0, len(request.Tools))
for _, tool := range request.Tools {
toolNames = append(toolNames, tool.Function.Name)
}
return toolNames
}
func (languageModel *ScriptedLanguageModel) toolCallResponse(toolName string, arguments json.RawMessage) model.ChatCompletionResponse {
return model.ChatCompletionResponse{
FinishReason: "tool_calls",
ProviderName: languageModel.providerName,
ModelName: languageModel.modelName,
SelectedBackend: "device",
Message: model.ChatCompletionMessage{
Role: "assistant",
ToolCalls: []model.ChatCompletionToolCall{{
ID: "scripted-call",
Type: "function",
Function: model.ChatCompletionToolCallFunction{
Name: toolName,
Arguments: string(arguments),
},
}},
},
}
}
func removeActionDiscriminator(document json.RawMessage) json.RawMessage {
var values map[string]json.RawMessage
if json.Unmarshal(document, &values) != nil {
return document
}
delete(values, "action")
normalizedDocument, errorValue := json.Marshal(values)
if errorValue != nil {
return document
}
return normalizedDocument
}
func structuredRequestFromChat(request model.ChatCompletionRequest) model.StructuredResponseRequest {
messages := make([]model.Message, 0, len(request.Messages))
for _, message := range request.Messages {
messages = append(messages, model.Message{Role: message.Role, Content: message.Content})
}
return model.StructuredResponseRequest{
Messages: messages,
StructuredOutputSchema: model.StructuredOutputSchema{
Name: request.SchemaName,
},
GenerationOptions: request.GenerationOptions,
}
}
func NewActionScriptedLanguageModel(actionResponses ...string) *ScriptedLanguageModel {
return NewScriptedLanguageModel(ScriptedLanguageModelOptions{ActionResponses: actionResponses})
}
func (languageModel *ScriptedLanguageModel) UseActionServedObserver(observer func(actionsServed int)) {
languageModel.mutex.Lock()
defer languageModel.mutex.Unlock()
languageModel.actionsServed = 0
languageModel.onActionServed = observer
}
func (languageModel *ScriptedLanguageModel) EnqueueActionResponses(actionResponses ...string) {
languageModel.mutex.Lock()
defer languageModel.mutex.Unlock()
languageModel.actionResponses = append(languageModel.actionResponses, actionResponses...)
}
func (languageModel *ScriptedLanguageModel) SetActionResponses(actionResponses ...string) {
languageModel.mutex.Lock()
defer languageModel.mutex.Unlock()
languageModel.actionResponses = append([]string{}, actionResponses...)
}
func (languageModel *ScriptedLanguageModel) EnqueueStructuredResponses(schemaName string, responses ...string) {
languageModel.mutex.Lock()
defer languageModel.mutex.Unlock()
trimmedSchemaName := strings.TrimSpace(schemaName)
if trimmedSchemaName == "" {
return
}
languageModel.structuredResponsesBySchema[trimmedSchemaName] = append(languageModel.structuredResponsesBySchema[trimmedSchemaName], responses...)
}
func (languageModel *ScriptedLanguageModel) PendingResponseCounts() map[string]int {
languageModel.mutex.Lock()
defer languageModel.mutex.Unlock()
pendingCounts := map[string]int{}
if len(languageModel.actionResponses) > 0 {
pendingCounts["bluecollar_agent_turn_action"] = len(languageModel.actionResponses)
}
for schemaName, responses := range languageModel.structuredResponsesBySchema {
if len(responses) > 0 {
pendingCounts[schemaName] += len(responses)
}
}
for schemaName, responses := range languageModel.chatResponsesBySchema {
if len(responses) > 0 {
pendingCounts[schemaName] += len(responses)
}
}
return pendingCounts
}
func (languageModel *ScriptedLanguageModel) RequestCount() int {
languageModel.mutex.Lock()
defer languageModel.mutex.Unlock()
return len(languageModel.requests)
}
func (languageModel *ScriptedLanguageModel) Requests() []model.StructuredResponseRequest {
return languageModel.RequestsSince(0)
}
func (languageModel *ScriptedLanguageModel) RequestsSince(startIndex int) []model.StructuredResponseRequest {
languageModel.mutex.Lock()
defer languageModel.mutex.Unlock()
if startIndex < 0 || startIndex > len(languageModel.requests) {
startIndex = 0
}
return append([]model.StructuredResponseRequest{}, languageModel.requests[startIndex:]...)
}
func (languageModel *ScriptedLanguageModel) GenerateResponse(context.Context, string) (string, error) {
return "", nil
}
func (languageModel *ScriptedLanguageModel) GenerateStructuredResponse(_ context.Context, request model.StructuredResponseRequest) (model.StructuredResponse, error) {
languageModel.mutex.Lock()
defer languageModel.mutex.Unlock()
languageModel.requests = append(languageModel.requests, request)
schemaName := strings.TrimSpace(request.StructuredOutputSchema.Name)
if response, isFound := languageModel.popStructuredResponse(request); isFound {
return languageModel.structuredResponse(response), nil
}
if schemaName == "bluecollar_agent_turn_action" {
response, errorValue := languageModel.popActionResponse()
if errorValue != nil {
return model.StructuredResponse{}, errorValue
}
return languageModel.structuredResponse(response), nil
}
response := strings.TrimSpace(languageModel.defaultResponsesBySchema[schemaName])
if response == "" {
if schemaName == "bluecollar_contract_skill_arbitration" {
response, errorValue := defaultContractSkillArbitrationResponse(request)
if errorValue != nil {
return model.StructuredResponse{}, errorValue
}
return languageModel.structuredResponse(response), nil
}
if schemaName == "approval_question" {
return languageModel.structuredResponse(defaultApprovalQuestionResponse(request)), nil
}
return model.StructuredResponse{}, noScriptedResponseError{schemaName: schemaName}
}
if schemaName == "bluecollar_skill_search_queries" && response == `{"queries":[]}` {
return languageModel.structuredResponse(defaultSkillSearchQueriesResponse(request)), nil
}
return languageModel.structuredResponse(response), nil
}
func defaultContractSkillArbitrationResponse(request model.StructuredResponseRequest) (string, error) {
var schema struct {
Properties map[string]struct {
Items struct {
Enum []string `json:"enum"`
} `json:"items"`
} `json:"properties"`
}
if errorValue := json.Unmarshal([]byte(request.StructuredOutputSchema.Document), &schema); errorValue != nil {
return "", fmt.Errorf("decode contract skill arbitration schema: %w", errorValue)
}
skillNames := schema.Properties["selectedSkillNames"].Items.Enum
selectedSkillNames := append([]string{}, skillNames...)
rejectedSkillNames := []string{}
expectedEvidence := []string{}
document, errorValue := json.Marshal(map[string]any{
"selectedSkillNames": selectedSkillNames,
"rejectedSkillNames": rejectedSkillNames,
"requiredNextToolNames": []string{},
"expectedEvidence": expectedEvidence,
"unmetPreconditions": []string{},
"reason": "scripted test default",
})
if errorValue != nil {
return "", fmt.Errorf("encode contract skill arbitration fixture: %w", errorValue)
}
return string(document), nil
}
func defaultSkillSearchQueriesResponse(request model.StructuredResponseRequest) string {
prompt := ""
for index := len(request.Messages) - 1; index >= 0; index-- {
if request.Messages[index].Role == "user" {
prompt = strings.TrimSpace(request.Messages[index].Content)
break
}
}
if prompt == "" {
return `{"queries":[]}`
}
document, errorValue := json.Marshal(map[string]any{"queries": []map[string]string{{"description": prompt}}})
if errorValue != nil {
return `{"queries":[]}`
}
return string(document)
}
func (languageModel *ScriptedLanguageModel) popStructuredResponse(request model.StructuredResponseRequest) (string, bool) {
schemaName := strings.TrimSpace(request.StructuredOutputSchema.Name)
responses := languageModel.structuredResponsesBySchema[schemaName]
if len(responses) == 0 {
return "", false
}
languageModel.structuredResponsesBySchema[schemaName] = responses[1:]
return responses[0], true
}
func (languageModel *ScriptedLanguageModel) popActionResponse() (string, error) {
for index, response := range languageModel.actionResponses {
if !isActionResponse(response) {
continue
}
languageModel.actionResponses = append(languageModel.actionResponses[:index], languageModel.actionResponses[index+1:]...)
languageModel.actionsServed++
if languageModel.onActionServed != nil {
languageModel.onActionServed(languageModel.actionsServed)
}
return response, nil
}
return "", fmt.Errorf("scripted language model action response queue is empty")
}
func (languageModel *ScriptedLanguageModel) structuredResponse(content string) model.StructuredResponse {
return model.StructuredResponse{ProviderName: languageModel.providerName, ModelName: languageModel.modelName, Content: content}
}
func isActionResponse(content string) bool {
document := map[string]any{}
if errorValue := json.Unmarshal([]byte(content), &document); errorValue != nil {
return false
}
return stringMapValue(document, "action") != ""
}
func copyResponseQueues(responseQueues map[string][]string) map[string][]string {
copiedResponseQueues := map[string][]string{}
for schemaName, responses := range responseQueues {
copiedResponseQueues[strings.TrimSpace(schemaName)] = append([]string{}, responses...)
}
return copiedResponseQueues
}
func mergeDefaultResponses(defaultResponses map[string]string) map[string]string {
mergedResponses := map[string]string{
"bluecollar_skill_search_queries": `{"queries":[]}`,
"bluecollar_execution_plan": `{"originalInstruction":"scripted test request","summary":"scripted test request","targets":[],"schedule":"","startAt":"","endAt":"","cadence":"","externalSend":false,"thirdPartyExternalSend":false,"repeated":false,"highFrequency":false,"destructive":false,"permissionChange":false,"publicDeploy":false,"paidAction":false,"missingInformation":[],"continuationInstruction":"scripted test request"}`,
"bluecollar_confirmation_message": `{"reply":"Understood. I will proceed once it is approved."}`,
}
for schemaName, response := range defaultResponses {
mergedResponses[strings.TrimSpace(schemaName)] = response
}
return mergedResponses
}
type approvalQuestionContextDocument struct {
OriginalRequest string `json:"originalRequest"`
ModelDraft string `json:"modelDraft"`
ActionDetails map[string]json.RawMessage `json:"actionDetails"`
}
func defaultApprovalQuestionResponse(request model.StructuredResponseRequest) string {
contextDocument := approvalQuestionContextFromRequest(request)
target := firstNonEmpty(actionDetail(contextDocument, "resolvedTarget"), actionDetail(contextDocument, "personHint"), actionDetail(contextDocument, "channelName"))
content := firstNonEmpty(actionDetail(contextDocument, "message"), actionDetail(contextDocument, "subject"), actionDetail(contextDocument, "body"), actionDetail(contextDocument, "title"), actionDetail(contextDocument, "summary"), actionDetail(contextDocument, "reason"), actionDetail(contextDocument, "approvalReason"))
question := defaultApprovalQuestionFromContext(contextDocument, target, content)
document, errorValue := json.Marshal(map[string]string{"question": question})
if errorValue != nil {
return `{"question":"should this be approved?"}`
}
return string(document)
}
func actionDetail(contextDocument approvalQuestionContextDocument, name string) string {
var value string
if json.Unmarshal(contextDocument.ActionDetails[name], &value) != nil {
return ""
}
return strings.TrimSpace(value)
}
func defaultApprovalQuestionFromContext(contextDocument approvalQuestionContextDocument, target string, content string) string {
if target != "" && content != "" {
return target + "should this be sent to?\n\n" + content
}
if draftQuestion := approvalQuestionFromDraft(contextDocument.ModelDraft); draftQuestion != "" {
return draftQuestion
}
if requestQuestion := approvalQuestionFromDraft(contextDocument.OriginalRequest); requestQuestion != "" {
return requestQuestion
}
if content != "" {
return content + "\n\nshould we proceed?"
}
return "should this be approved?"
}
func approvalQuestionContextFromRequest(request model.StructuredResponseRequest) approvalQuestionContextDocument {
for index := len(request.Messages) - 1; index >= 0; index-- {
var contextDocument approvalQuestionContextDocument
if json.Unmarshal([]byte(request.Messages[index].Content), &contextDocument) == nil {
return contextDocument
}
}
return approvalQuestionContextDocument{}
}
func approvalQuestionFromDraft(value string) string {
trimmedValue := strings.Trim(strings.TrimSpace(value), ".。")
if trimmedValue == "" {
return ""
}
if strings.HasSuffix(trimmedValue, "?") || strings.HasSuffix(trimmedValue, "？") {
return trimmedValue
}
if strings.HasSuffix(trimmedValue, "doing it") {
return strings.TrimSuffix(trimmedValue, "doing it") + "shall we?"
}
if strings.HasSuffix(trimmedValue, "please do it") {
return strings.TrimSuffix(trimmedValue, "please do it") + "shall we?"
}
return trimmedValue + "?"
}
func stringMapValue(document map[string]any, key string) string {
value, isString := document[key].(string)
if !isString {
return ""
}
return strings.TrimSpace(value)
}
func firstNonEmpty(values ...string) string {
for _, value := range values {
if strings.TrimSpace(value) != "" {
return strings.TrimSpace(value)
}
}
return ""
}
var ErrNoScriptedResponse = errors.New("scripted language model has no response")
type noScriptedResponseError struct {
schemaName string
}
func (errorValue noScriptedResponseError) Error() string {
return "scripted language model has no " + errorValue.schemaName + " response"
}
func (errorValue noScriptedResponseError) Is(target error) bool {
return target == ErrNoScriptedResponse
}