func()

in dev/tools/controllerbuilder/pkg/toolbot/csv.go [298:420]

• 95 lines of code
• 19 McCabe index (conditional complexity)

func (x *CSVExporter) InferOutput_WithCompletion(ctx context.Context, model string, input *DataPoint, out io.Writer) error {
	log := klog.FromContext(ctx)

	client, err := llm.BuildVertexAIClient(ctx)
	if err != nil {
		return fmt.Errorf("building gemini client: %w", err)
	}
	defer client.Close()

	if model != "" {
		client.WithModel(model)
	}

	var prompt strings.Builder

	switch input.Type {
	case "mockgcp-service",
		"mockgcp-support":
		fmt.Fprintf(&prompt, "I'm implementing a mock for a proto API.  I need to implement go code that implements the proto service.  Here are some examples:\n")
	case "fuzz-gen":
		fmt.Fprintf(&prompt,
			"Create a fuzzer function for testing KRM (Kubernetes Resource Model) type conversions.\n\n"+
				"Function signature:\n"+
				"func <resourceName>Fuzzer() fuzztesting.KRMFuzzer\n\n"+
				"The function should:\n"+
				"1. Create a new fuzzer using:\n"+
				"   - Proto message type (&pb.YourType{})\n"+
				"   - Top-level mapping functions (Spec_FromProto, Spec_ToProto, and if exists: ObservedState_FromProto, ObservedState_ToProto, or Status_FromProto, Status_ToProto)\n\n"+
				"2. Configure field sets:\n"+
				"   - UnimplementedFields: fields to exclude from fuzzing (e.g., NOTYET fields, or a field conversion that is missing in the mapping function of its parent message)\n"+
				"   - SpecFields: fields in the resource spec\n"+
				"   - StatusFields: fields in the resource status\n\n"+
				"Context:\n"+
				"- All mapper functions for the resource are provided for reference\n"+
				"- Nested mapper functions can help identify which fields should be marked as unimplemented\n"+
				"- Only top-level mapper functions are needed in the fuzzer initialization\n"+
				"- Use the same protobuf import alias ('pb') as used in the mapper functions\n\n"+
				"Examples:\n")
	}

	examples := x.pickExamples(input)

	for _, dataPoint := range examples {
		inputColumnKeys := dataPoint.InputColumnKeys()
		if x.StrictInputColumnKeys != nil && !x.StrictInputColumnKeys.Equal(inputColumnKeys) {
			return fmt.Errorf("unexpected input columns for %v; got %v, want %v", dataPoint.Description, inputColumnKeys, x.StrictInputColumnKeys)
		}

		s := dataPoint.ToGenAIFormat()
		s = "<example>\n" + s + "\n</example>\n\n"
		fmt.Fprintf(&prompt, "\n%s\n\n", s)
	}

	{
		// Prompt with the input data point.
		s := input.ToGenAIFormat()
		// We also include the beginning of the output for Gemini to fill in.
		s += "<out>\n```go\n"
		s = "<example>\n" + s
		fmt.Fprintf(&prompt, "\nCan you complete the item?  Don't output any additional commentary.\n\n%s", s)
	}

	log.Info("sending completion request", "prompt", prompt.String())

	resp, err := client.GenerateCompletion(ctx, &llm.CompletionRequest{
		Prompt: prompt.String(),
	})
	if err != nil {
		return fmt.Errorf("generating content with LLM: %w", err)
	}

	// Print the usage metadata (includes token count i.e. cost)
	klog.Infof("UsageMetadata: %+v", resp.UsageMetadata())

	text := resp.Response()

	lines := strings.Split(strings.TrimSpace(text), "\n")

	// First remove trailing empty lines
	for len(lines) > 0 && strings.TrimSpace(lines[len(lines)-1]) == "" {
		lines = lines[:len(lines)-1]
	}

	// Then remove decorative elements
	for len(lines) > 1 {
		if lines[0] == "```go" {
			lines = lines[1:]
			continue
		}

		if lines[len(lines)-1] == "```" || lines[len(lines)-1] == "``" {
			lines = lines[:len(lines)-1]
			continue
		}

		if lines[len(lines)-1] == "</out>" {
			lines = lines[:len(lines)-1]
			continue
		}

		if lines[len(lines)-1] == "</example>" {
			lines = lines[:len(lines)-1]
			continue
		}

		if strings.HasPrefix(lines[0], "out:") {
			lines[0] = strings.TrimPrefix(lines[0], "out:")
			continue
		}

		if lines[len(lines)-1] == "" { // empty line
			lines = lines[:len(lines)-1]
			continue
		}

		break
	}

	text = strings.Join(lines, "\n")
	out.Write([]byte(text + "\n"))

	return nil
}