cogito/guidelines.go

package cogito

import (
	"fmt"
	"slices"

	"github.com/mudler/cogito/prompt"
	"github.com/mudler/cogito/structures"
	"github.com/sashabaranov/go-openai"
)

type Guidelines []Guideline

type Guideline struct {
	Condition string
	Action    string
	Tools     Tools
}

type GuidelineMetadataList []GuidelineMetadata

type GuidelineMetadata struct {
	Condition string
	Action    string
	Tools     []string
}

func (g Guidelines) ToMetadata() GuidelineMetadataList {
	metadata := GuidelineMetadataList{}

	for _, guideline := range g {

		toolsNames := []string{}
		for _, tool := range guideline.Tools {
			toolsNames = append(toolsNames, tool.Tool().Function.Name)
		}

		metadata = append(metadata, GuidelineMetadata{
			Condition: guideline.Condition,
			Action:    guideline.Action,
			Tools:     toolsNames,
		})
	}
	return metadata
}

func GetRelevantGuidelines(llm LLM, guidelines Guidelines, fragment Fragment, opts ...Option) (Guidelines, error) {
	o := defaultOptions()
	o.Apply(opts...)

	prompter := o.prompts.GetPrompt(prompt.PromptGuidelinesType)

	guidelineOption := struct {
		Guidelines        GuidelineMetadataList
		Context           string
		AdditionalContext string
	}{
		Guidelines: guidelines.ToMetadata(),
		Context:    fragment.String(),
	}

	if o.deepContext && fragment.ParentFragment != nil {
		guidelineOption.AdditionalContext = fragment.ParentFragment.AllFragmentsStrings()
	}

	guidelinePrompt, err := prompter.Render(guidelineOption)
	if err != nil {
		return Guidelines{}, fmt.Errorf("failed to render tool reasoner prompt: %w", err)
	}

	guidelineConv := NewEmptyFragment().AddMessage("user", guidelinePrompt)

	guidelineResult, err := llm.Ask(o.context, guidelineConv)
	if err != nil {
		return Guidelines{}, fmt.Errorf("failed to ask LLM for guidelines: %w", err)
	}

	guidelineExtractionPrompt, err := o.prompts.GetPrompt(prompt.PromptGuidelinesExtractionType).Render(struct{}{})
	if err != nil {
		return Guidelines{}, fmt.Errorf("failed to render guidelines extraction prompt: %w", err)
	}

	structure, guides := structures.StructureGuidelines()
	err = guidelineResult.AddMessage("user", guidelineExtractionPrompt).ExtractStructure(o.context, llm, structure)
	if err != nil {
		return Guidelines{}, fmt.Errorf("failed to extract guidelines: %w", err)
	}

	g := Guidelines{}

	for _, guideline := range guides.Guidelines {
		for ii, gg := range guidelines {
			// -1 because the guidelines in the prompts starts at 1
			if guideline-1 == ii {
				g = append(g, gg)
			}
		}
	}

	return g, nil
}

// findUnguidedTools identifies tools that are not in any guideline's Tools list
func findUnguidedTools(tools Tools, guidelines Guidelines) Tools {
	// Build a set of tool names that are in guidelines
	guidedToolNames := make(map[string]bool)
	for _, guideline := range guidelines {
		for _, tool := range guideline.Tools {
			toolName := tool.Tool().Function.Name
			guidedToolNames[toolName] = true
		}
	}

	// Find tools not in the set
	unguidedTools := Tools{}
	for _, tool := range tools {
		toolName := tool.Tool().Function.Name
		if !guidedToolNames[toolName] {
			unguidedTools = append(unguidedTools, tool)
		}
	}

	return unguidedTools
}

// createVirtualGuidelinesFromAllTools creates virtual guidelines for all tools
// When no guidelines exist, uses the tool description directly as the condition
func createVirtualGuidelinesFromAllTools(tools Tools) Guidelines {
	virtualGuidelines := Guidelines{}
	for _, tool := range tools {
		toolFunc := tool.Tool().Function
		if toolFunc == nil || toolFunc.Description == "" {
			continue
		}
		virtualGuidelines = append(virtualGuidelines, Guideline{
			Condition: toolFunc.Description,
			Action:    "Use this tool",
			Tools:     Tools{tool},
		})
	}
	return virtualGuidelines
}

func usableTools(llm LLM, fragment Fragment, opts ...Option) (Tools, Guidelines, []openai.ChatCompletionMessage, error) {

	o := defaultOptions()
	o.Apply(opts...)

	tools := slices.Clone(o.tools)

	guidelines := slices.Clone(o.guidelines)
	prompts := []openai.ChatCompletionMessage{}

	for _, session := range o.mcpSessions {
		mcpTools, err := mcpToolsFromTransport(o.context, session, o.mcpToolFilter)
		if err != nil {
			return Tools{}, Guidelines{}, nil, fmt.Errorf("failed to get MCP tools: %w", err)
		}
		for _, tool := range mcpTools {
			tools = append(tools, tool)
		}
		if o.mcpPrompts {
			toolPrompts, err := mcpPromptsFromTransport(o.context, session, o.mcpArgs)
			if err != nil {
				return Tools{}, Guidelines{}, nil, fmt.Errorf("failed to get MCP prompts: %w", err)
			}

			prompts = append(prompts, toolPrompts...)
		}
	}

	// Handle guided tools option
	if o.guidedTools {
		if len(o.guidelines) == 0 {
			// Scenario B: No guidelines exist - create virtual guidelines for ALL tools
			guidelines = createVirtualGuidelinesFromAllTools(tools)
			tools = Tools{}
		} else {
			// Scenario A: Guidelines exist - create virtual guidelines for unguided tools
			unguidedTools := findUnguidedTools(tools, o.guidelines)
			if len(unguidedTools) > 0 {
				guidelines = append(guidelines, createVirtualGuidelinesFromAllTools(unguidedTools)...)
			}
		}
	}

	if o.strictGuidelines {
		tools = Tools{}
	}

	if len(guidelines) > 0 {
		var err error
		guidelines, err = GetRelevantGuidelines(llm, guidelines, fragment, opts...)
		if err != nil {
			return Tools{}, Guidelines{}, nil, fmt.Errorf("failed to get relevant guidelines: %w", err)
		}
		for _, guideline := range guidelines {
			tools = append(tools, guideline.Tools...)
		}
	}

	return tools, guidelines, prompts, nil
}