promptfoo

1
type Probability = 'NEGLIGIBLE' | 'LOW' | 'MEDIUM' | 'HIGH';
2

3
interface SafetyRating {
4
  category:
5
    | 'HARM_CATEGORY_HARASSMENT'
6
    | 'HARM_CATEGORY_HATE_SPEECH'
7
    | 'HARM_CATEGORY_SEXUALLY_EXPLICIT'
8
    | 'HARM_CATEGORY_DANGEROUS_CONTENT';
9
  probability: Probability;
10
}
11

12
interface PartText {
13
  text: string;
14
}
15

16
interface PartFunctionCall {
17
  functionCall: {
18
    name: string;
19
    args: Record<string, string>;
20
  };
21
}
22

23
type Part = PartText | PartFunctionCall;
24

25
interface Content {
26
  parts: Part[];
27
  role?: 'model';
28
}
29

30
interface Candidate {
31
  content: Content;
32
  finishReason?: 'FINISH_REASON_STOP' | 'STOP';
33
  safetyRatings: SafetyRating[];
34
}
35

36
interface GeminiUsageMetadata {
37
  promptTokenCount: number;
38
  candidatesTokenCount?: number;
39
  totalTokenCount: number;
40
}
41

42
export interface GeminiErrorResponse {
43
  error: {
44
    code: number;
45
    message: string;
46
    status: string;
47
  };
48
}
49

50
export interface GeminiResponseData {
51
  candidates: Candidate[];
52
  usageMetadata?: GeminiUsageMetadata;
53
}
54

55
export type GeminiApiResponse = (GeminiResponseData | GeminiErrorResponse)[];
56

57
export interface Palm2ApiResponse {
58
  error?: {
59
    code: string;
60
    message: string;
61
  };
62
  predictions?: [
63
    {
64
      candidates: [
65
        {
66
          content: string;
67
        },
68
      ];
69
    },
70
  ];
71
}
72

73
export function maybeCoerceToGeminiFormat(contents: any) {
74
  let coerced = false;
75
  if (Array.isArray(contents) && typeof contents[0].content === 'string') {
76
    // This looks like an OpenAI chat prompt.  Convert it to a compatible format
77
    contents = {
78
      role: 'user',
79
      parts: {
80
        text: contents.map((item) => item.content).join(''),
81
      },
82
    };
83
    coerced = true;
84
  }
85
  return { contents, coerced };
86
}
87