Merge pull request #994 from narengogi/feat/enhanced-grounding-support

changes to support cached streaming in grounding requests

Author: VisargD

src/providers/google-vertex-ai/types.ts

@@ -1,4 +1,4 @@
-import { ChatCompletionResponse } from '../types';
+import { ChatCompletionResponse, GroundingMetadata } from '../types';
 
 export interface GoogleErrorResponse {
   error: {
@@ -46,24 +46,7 @@ export interface GoogleResponseCandidate {
     category: string;
     probability: string;
   }[];
-  groundingMetadata?: {
-    webSearchQueries?: string[];
-    searchEntryPoint?: {
-      renderedContent: string;
-    };
-    groundingSupports?: Array<{
-      segment: {
-        startIndex: number;
-        endIndex: number;
-        text: string;
-      };
-      groundingChunkIndices: number[];
-      confidenceScores: number[];
-    }>;
-    retrievalMetadata?: {
-      webDynamicRetrievalScore: number;
-    };
-  };
+  groundingMetadata?: GroundingMetadata;
 }
 
 export interface GoogleGenerateContentResponse {

src/providers/google/chatComplete.ts

@@ -19,6 +19,7 @@ import {
 import {
   ChatCompletionResponse,
   ErrorResponse,
+  GroundingMetadata,
   Logprobs,
   ProviderConfig,
 } from '../types';
@@ -452,24 +453,7 @@ interface GoogleResponseCandidate {
     category: string;
     probability: string;
   }[];
-  groundingMetadata?: {
-    webSearchQueries?: string[];
-    searchEntryPoint?: {
-      renderedContent: string;
-    };
-    groundingSupports?: Array<{
-      segment: {
-        startIndex: number;
-        endIndex: number;
-        text: string;
-      };
-      groundingChunkIndices: number[];
-      confidenceScores: number[];
-    }>;
-    retrievalMetadata?: {
-      webDynamicRetrievalScore: number;
-    };
-  };
+  groundingMetadata?: GroundingMetadata;
 }
 
 interface GoogleGenerateContentResponse {

src/providers/openai/chatComplete.ts

@@ -146,13 +146,14 @@ export const OpenAIChatCompleteJSONToStreamResponseTransform: (
   provider: string
 ) => Array<string> = (response, provider) => {
   const streamChunkArray: Array<string> = [];
-  const { id, model, system_fingerprint, choices } = response;
+  const { id, model, system_fingerprint, choices, citations } = response;
 
   const {
     prompt_tokens,
     completion_tokens,
     cache_read_input_tokens,
     cache_creation_input_tokens,
+    num_search_queries,
   } = response.usage || {};
 
   let total_tokens;
@@ -179,7 +180,9 @@ export const OpenAIChatCompleteJSONToStreamResponseTransform: (
         cache_read_input_tokens,
         cache_creation_input_tokens,
       }),
+      ...(num_search_queries && { num_search_queries }),
     },
+    ...(citations && { citations }),
   };
 
   for (const [index, choice] of choices.entries()) {
@@ -262,6 +265,9 @@ export const OpenAIChatCompleteJSONToStreamResponseTransform: (
                   role: 'assistant',
                   content: word,
                 },
+                ...(choice.groundingMetadata && {
+                  groundingMetadata: choice.groundingMetadata,
+                }),
               },
             ],
           })}\n\n`

src/providers/types.ts

@@ -152,6 +152,7 @@ export interface CResponse extends BaseResponse {
      */
     cache_read_input_tokens?: number;
     cache_creation_input_tokens?: number;
+    num_search_queries?: number;
   };
 }
 
@@ -168,6 +169,25 @@ export interface CompletionResponse extends CResponse {
   }[];
 }
 
+export interface GroundingMetadata {
+  webSearchQueries?: string[];
+  searchEntryPoint?: {
+    renderedContent: string;
+  };
+  groundingSupports?: Array<{
+    segment: {
+      startIndex: number;
+      endIndex: number;
+      text: string;
+    };
+    groundingChunkIndices: number[];
+    confidenceScores: number[];
+  }>;
+  retrievalMetadata?: {
+    webDynamicRetrievalScore: number;
+  };
+}
+
 /**
  * The structure of a choice in a chat completion response.
  * @interface
@@ -177,6 +197,7 @@ export interface ChatChoice {
   message: Message;
   finish_reason: string;
   logprobs?: object | null;
+  groundingMetadata?: GroundingMetadata;
 }
 
 export interface Logprobs {
@@ -197,6 +218,7 @@ export interface Logprobs {
 export interface ChatCompletionResponse extends CResponse {
   choices: ChatChoice[];
   provider?: string;
+  citations?: string[];
 }
 
 /**