googleai: test GenerateFromSinglePrompt convenience function (#538)

tmc · Jan 20, 2024 · 19b5f8e · 19b5f8e
1 parent 715b8c2
commit 19b5f8e
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/llms/googleai/googleai_llm_test.go b/llms/googleai/googleai_llm_test.go
@@ -48,6 +48,17 @@ func TestMultiContentText(t *testing.T) {
 	assert.Regexp(t, "(?i)dog|canid|canine", c1.Content)
 }
 
+func TestGenerateFromSinglePrompt(t *testing.T) {
+	t.Parallel()
+	llm := newClient(t)
+
+	prompt := "name all the planets in the solar system"
+	rsp, err := llms.GenerateFromSinglePrompt(context.Background(), llm, prompt)
+	require.NoError(t, err)
+
+	assert.Regexp(t, "(?i)jupiter", rsp)
+}
+
 func TestMultiContentTextStream(t *testing.T) {
 	t.Parallel()
 	llm := newClient(t)