feat(api): add streaming to kg question (#101)

writer · Nov 8, 2024 · 0279ffc · 0279ffc
1 parent 8d775fb
commit 0279ffc
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 6 deletions.
diff --git a/.stats.yml b/.stats.yml
@@ -1,2 +1,2 @@
 configured_endpoints: 21
-openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/writerai%2Fwriter-c350270f059b9d4c51d9bca0c34c3e3018e544f5daaf2d18aacc55a0561573f8.yml
+openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/writerai%2Fwriter-9c49d7cc176d84ba91accd64aa4a8c91a44cdce83e9ea59da8b259fa19ccf9f6.yml
diff --git a/src/index.ts b/src/index.ts
@@ -46,6 +46,8 @@ import {
   GraphDeleteResponse,
   GraphListParams,
   GraphQuestionParams,
+  GraphQuestionParamsNonStreaming,
+  GraphQuestionParamsStreaming,
   GraphRemoveFileFromGraphResponse,
   GraphUpdateParams,
   GraphUpdateResponse,
@@ -294,6 +296,8 @@ export declare namespace Writer {
     type GraphListParams as GraphListParams,
     type GraphAddFileToGraphParams as GraphAddFileToGraphParams,
     type GraphQuestionParams as GraphQuestionParams,
+    type GraphQuestionParamsNonStreaming as GraphQuestionParamsNonStreaming,
+    type GraphQuestionParamsStreaming as GraphQuestionParamsStreaming,
   };
 
   export {

diff --git a/src/resources/graphs.ts b/src/resources/graphs.ts
@@ -2,9 +2,12 @@
 
 import { APIResource } from '../resource';
 import { isRequestOptions } from '../core';
+import { APIPromise } from '../core';
 import * as Core from '../core';
+import * as GraphsAPI from './graphs';
 import * as FilesAPI from './files';
 import { CursorPage, type CursorPageParams } from '../pagination';
+import { Stream } from '../streaming';
 
 export class Graphs extends APIResource {
   /**
@@ -85,8 +88,19 @@ export class Graphs extends APIResource {
   /**
    * Ask a question to specified Knowledge Graphs.
    */
-  question(body: GraphQuestionParams, options?: Core.RequestOptions): Core.APIPromise<Question> {
-    return this._client.post('/v1/graphs/question', { body, ...options });
+  question(body: GraphQuestionParamsNonStreaming, options?: Core.RequestOptions): APIPromise<Question>;
+  question(body: GraphQuestionParamsStreaming, options?: Core.RequestOptions): APIPromise<Stream<Question>>;
+  question(
+    body: GraphQuestionParamsBase,
+    options?: Core.RequestOptions,
+  ): APIPromise<Stream<Question> | Question>;
+  question(
+    body: GraphQuestionParams,
+    options?: Core.RequestOptions,
+  ): APIPromise<Question> | APIPromise<Stream<Question>> {
+    return this._client.post('/v1/graphs/question', { body, ...options, stream: body.stream ?? false }) as
+      | APIPromise<Question>
+      | APIPromise<Stream<Question>>;
   }
 
   /**
@@ -322,7 +336,9 @@ export interface GraphAddFileToGraphParams {
   file_id: string;
 }
 
-export interface GraphQuestionParams {
+export type GraphQuestionParams = GraphQuestionParamsNonStreaming | GraphQuestionParamsStreaming;
+
+export interface GraphQuestionParamsBase {
   /**
    * The unique identifiers of the Knowledge Graphs to be queried.
    */
@@ -346,6 +362,29 @@ export interface GraphQuestionParams {
   subqueries: boolean;
 }
 
+export namespace GraphQuestionParams {
+  export type GraphQuestionParamsNonStreaming = GraphsAPI.GraphQuestionParamsNonStreaming;
+  export type GraphQuestionParamsStreaming = GraphsAPI.GraphQuestionParamsStreaming;
+}
+
+export interface GraphQuestionParamsNonStreaming extends GraphQuestionParamsBase {
+  /**
+   * Determines whether the model's output should be streamed. If true, the output is
+   * generated and sent incrementally, which can be useful for real-time
+   * applications.
+   */
+  stream: false;
+}
+
+export interface GraphQuestionParamsStreaming extends GraphQuestionParamsBase {
+  /**
+   * Determines whether the model's output should be streamed. If true, the output is
+   * generated and sent incrementally, which can be useful for real-time
+   * applications.
+   */
+  stream: true;
+}
+
 Graphs.GraphsCursorPage = GraphsCursorPage;
 
 export declare namespace Graphs {
@@ -362,5 +401,7 @@ export declare namespace Graphs {
     type GraphListParams as GraphListParams,
     type GraphAddFileToGraphParams as GraphAddFileToGraphParams,
     type GraphQuestionParams as GraphQuestionParams,
+    type GraphQuestionParamsNonStreaming as GraphQuestionParamsNonStreaming,
+    type GraphQuestionParamsStreaming as GraphQuestionParamsStreaming,
   };
 }
diff --git a/src/resources/index.ts b/src/resources/index.ts
@@ -45,6 +45,8 @@ export {
   type GraphListParams,
   type GraphAddFileToGraphParams,
   type GraphQuestionParams,
+  type GraphQuestionParamsNonStreaming,
+  type GraphQuestionParamsStreaming,
 } from './graphs';
 export { Models, type ModelListResponse } from './models';
 export {

diff --git a/tests/api-resources/graphs.test.ts b/tests/api-resources/graphs.test.ts
@@ -134,7 +134,7 @@ describe('resource graphs', () => {
     const responsePromise = client.graphs.question({
       graph_ids: ['182bd5e5-6e1a-4fe4-a799-aa6d9a6ab26e'],
       question: 'question',
-      stream: true,
+      stream: false,
       subqueries: true,
     });
     const rawResponse = await responsePromise.asResponse();
@@ -150,7 +150,7 @@ describe('resource graphs', () => {
     const response = await client.graphs.question({
       graph_ids: ['182bd5e5-6e1a-4fe4-a799-aa6d9a6ab26e'],
       question: 'question',
-      stream: true,
+      stream: false,
       subqueries: true,
     });
   });