From 7eb19e08168eb1c59222521c01d6cf44ca3069cd Mon Sep 17 00:00:00 2001
From: Christina Holland <chholland@google.com>
Date: Fri, 5 Apr 2024 16:17:30 -0700
Subject: [PATCH 1/2] add totalBillableCharacters

---
 packages/vertexai/src/types/responses.ts | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/packages/vertexai/src/types/responses.ts b/packages/vertexai/src/types/responses.ts
index bdba5d28f23..865f88fcae8 100644
--- a/packages/vertexai/src/types/responses.ts
+++ b/packages/vertexai/src/types/responses.ts
@@ -148,5 +148,13 @@ export interface SafetyRating {
  * @public
  */
 export interface CountTokensResponse {
+  /**
+   * The total number of tokens counted across all instances from the request.
+   */
   totalTokens: number;
+  /**
+   * The total number of billable characters counted across all instances
+   * from the request.
+   */
+  totalBillableCharacters?: number;
 }

From 07c02a595bee9e2323d8ced97b492f50b42a1c80 Mon Sep 17 00:00:00 2001
From: Christina Holland <chholland@google.com>
Date: Mon, 8 Apr 2024 12:03:40 -0700
Subject: [PATCH 2/2] add more missing types

---
 packages/vertexai/src/types/enums.ts     | 29 +++++++++++++
 packages/vertexai/src/types/requests.ts  |  5 ++-
 packages/vertexai/src/types/responses.ts | 53 +++++++++++++++++++++++-
 3 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/packages/vertexai/src/types/enums.ts b/packages/vertexai/src/types/enums.ts
index cdb8e4bca79..817450a069b 100644
--- a/packages/vertexai/src/types/enums.ts
+++ b/packages/vertexai/src/types/enums.ts
@@ -56,6 +56,18 @@ export enum HarmBlockThreshold {
   BLOCK_NONE = 'BLOCK_NONE'
 }
 
+/**
+ * @public
+ */
+export enum HarmBlockMethod {
+  // The harm block method is unspecified.
+  HARM_BLOCK_METHOD_UNSPECIFIED = 'HARM_BLOCK_METHOD_UNSPECIFIED',
+  // The harm block method uses both probability and severity scores.
+  SEVERITY = 'SEVERITY',
+  // The harm block method uses the probability score.
+  PROBABILITY = 'PROBABILITY'
+}
+
 /**
  * Probability that a prompt or candidate matches a harm category.
  * @public
@@ -73,6 +85,23 @@ export enum HarmProbability {
   HIGH = 'HIGH'
 }
 
+/**
+ * Harm severity levels.
+ * @public
+ */
+export enum HarmSeverity {
+  // Harm severity unspecified.
+  HARM_SEVERITY_UNSPECIFIED = 'HARM_SEVERITY_UNSPECIFIED',
+  // Negligible level of harm severity.
+  HARM_SEVERITY_NEGLIGIBLE = 'HARM_SEVERITY_NEGLIGIBLE',
+  // Low level of harm severity.
+  HARM_SEVERITY_LOW = 'HARM_SEVERITY_LOW',
+  // Medium level of harm severity.
+  HARM_SEVERITY_MEDIUM = 'HARM_SEVERITY_MEDIUM',
+  // High level of harm severity.
+  HARM_SEVERITY_HIGH = 'HARM_SEVERITY_HIGH'
+}
+
 /**
  * Reason that a prompt was blocked.
  * @public
diff --git a/packages/vertexai/src/types/requests.ts b/packages/vertexai/src/types/requests.ts
index 0ab2f14b8c1..7083341c1cd 100644
--- a/packages/vertexai/src/types/requests.ts
+++ b/packages/vertexai/src/types/requests.ts
@@ -16,7 +16,7 @@
  */
 
 import { Content } from './content';
-import { HarmBlockThreshold, HarmCategory } from './enums';
+import { HarmBlockMethod, HarmBlockThreshold, HarmCategory } from './enums';
 
 /**
  * Base parameters for a number of methods.
@@ -52,6 +52,7 @@ export interface GenerateContentRequest extends BaseParams {
 export interface SafetySetting {
   category: HarmCategory;
   threshold: HarmBlockThreshold;
+  method: HarmBlockMethod;
 }
 
 /**
@@ -65,6 +66,8 @@ export interface GenerationConfig {
   temperature?: number;
   topP?: number;
   topK?: number;
+  presencePenalty?: number;
+  frequencyPenalty?: number;
 }
 
 /**
diff --git a/packages/vertexai/src/types/responses.ts b/packages/vertexai/src/types/responses.ts
index 865f88fcae8..375bd340e42 100644
--- a/packages/vertexai/src/types/responses.ts
+++ b/packages/vertexai/src/types/responses.ts
@@ -20,7 +20,8 @@ import {
   BlockReason,
   FinishReason,
   HarmCategory,
-  HarmProbability
+  HarmProbability,
+  HarmSeverity
 } from './enums';
 
 /**
@@ -101,6 +102,7 @@ export interface GenerateContentCandidate {
   finishMessage?: string;
   safetyRatings?: SafetyRating[];
   citationMetadata?: CitationMetadata;
+  groundingMetadata?: GroundingMetadata;
 }
 
 /**
@@ -124,6 +126,51 @@ export interface Citation {
   publicationDate?: Date;
 }
 
+/**
+ * Metadata returned to client when grounding is enabled.
+ * @public
+ */
+export interface GroundingMetadata {
+  webSearchQueries?: string[];
+  retrievalQueries?: string[];
+  groundingAttributions: GroundingAttribution[];
+}
+
+/**
+ * @public
+ */
+export interface GroundingAttribution {
+  segment: Segment;
+  confidenceScore?: number;
+  web?: WebAttribution;
+  retrievedContext?: RetrievedContextAttribution;
+}
+
+/**
+ * @public
+ */
+export interface Segment {
+  partIndex: number;
+  startIndex: number;
+  endIndex: number;
+}
+
+/**
+ * @public
+ */
+export interface WebAttribution {
+  uri: string;
+  title: string;
+}
+
+/**
+ * @public
+ */
+export interface RetrievedContextAttribution {
+  uri: string;
+  title: string;
+}
+
 /**
  * Protobuf google.type.Date
  * @public
@@ -141,6 +188,10 @@ export interface Date {
 export interface SafetyRating {
   category: HarmCategory;
   probability: HarmProbability;
+  severity: HarmSeverity;
+  probabilityScore: number;
+  severityScore: number;
+  blocked: boolean;
 }
 
 /**