Merge pull request #2 from trycortexai/stream-response-support

feat: add streamResponse

Author: zo0r

examples/hono/src/index.ts

@@ -4,19 +4,19 @@ import dotenv from 'dotenv';
 import {Hono, type Context} from 'hono';
 import {env} from 'hono/adapter';
 
-import {stepStream, workflowStream} from './stream.js';
+import {stepStreamResponse, worklflowStreamResponse} from './stream.js';
 
 dotenv.config();
 
 const PORT = 8085;
 
 const app = new Hono();
 
-const createCortex = (c: Context) => {
+const createCortex = (ctx: Context) => {
   const {CORTEX_API_KEY, BASE_URL} = env<{
     CORTEX_API_KEY: string;
     BASE_URL: string;
-  }>(c);
+  }>(ctx);
 
   const cortex = new Cortex({
     apiKey: CORTEX_API_KEY,
@@ -30,17 +30,30 @@ app.get('/', c => {
   return c.text('Cortex SDK');
 });
 
-app.get('/stream/step', c => {
-  const cortex = createCortex(c);
-  stepStream(cortex);
-  return c.text('Stream Step');
+app.get('/stream/step', async ctx => {
+  const cortex = createCortex(ctx);
+
+  const response = await stepStreamResponse(cortex);
+
+  if (response instanceof Response) {
+    return response;
+  }
+
+  return ctx.text(response);
 });
 
-app.get('/stream/workflow', c => {
-  const {WORKFLOW_ID} = env<{WORKFLOW_ID: string}>(c);
-  const cortex = createCortex(c);
-  workflowStream(cortex, WORKFLOW_ID);
-  return c.text('Stream Workflow');
+app.get('/stream/workflow', async ctx => {
+  const {WORKFLOW_ID} = env<{WORKFLOW_ID: string}>(ctx);
+
+  const cortex = createCortex(ctx);
+
+  const response = await worklflowStreamResponse(cortex, WORKFLOW_ID);
+
+  if (response instanceof Response) {
+    return response;
+  }
+
+  return ctx.text(response);
 });
 
 console.log(`Server is running on http://localhost:${PORT}`);

examples/hono/src/stream.ts

@@ -4,9 +4,9 @@ import {
   type StepModelOutputSchema,
 } from '@cortex-ai/sdk';
 
-export const stepStream = async (cortex: Cortex) => {
+export const stepStreamResponse = async (cortex: Cortex) => {
   try {
-    const result = await cortex.apps.runs.step.stream(
+    const result = await cortex.apps.runs.step.streamResponse(
       {
         step: {
           type: 'model',
@@ -24,45 +24,48 @@ export const stepStream = async (cortex: Cortex) => {
         },
       },
       {
-        onStream: partial => {
+        onStream: (partial, event) => {
           const stepOutput = partial.output as
             | StepModelOutputSchema
             | undefined;
 
-          console.log('Chunk Message:', stepOutput?.message);
+          console.log(event, stepOutput?.message);
         },
       },
     );
 
-    const stepOutput = result.output as StepModelOutputSchema | undefined;
-
-    console.log('\n\nFull Message: ', stepOutput?.message);
+    return result;
   } catch (error) {
     console.log('Error', error);
+    return error?.toString() ?? 'Error';
   }
 };
 
-export const workflowStream = async (cortex: Cortex, workflowId: string) => {
+export const worklflowStreamResponse = async (
+  cortex: Cortex,
+  workflowId: string,
+) => {
   try {
-    const result = await cortex.apps.workflows.runs.stream(
+    const result = await cortex.apps.workflows.runs.streamResponse(
       workflowId,
       {
         input: {
           message: 'Hello, World!',
         },
       },
       {
-        onStream: partial => {
+        onStream: (partial, event) => {
           const stepOutput = partial.output
             ?.MODEL as CastRunStepOutputSchema<StepModelOutputSchema>;
 
-          console.log('chunk', stepOutput?.output?.message);
+          console.log(event, stepOutput?.output?.message);
         },
       },
     );
 
-    console.log('\n\nresult: ', JSON.stringify(result, null, 2));
+    return result;
   } catch (error) {
     console.log('Error', error);
+    return error?.toString() ?? 'Error';
   }
 };

openapi.utils.ts

@@ -86,6 +86,7 @@ type MethodOptions = {
   body?: string;
   query?: string;
   returns: string;
+  finalReturns?: string;
   description?: string;
 };
 
@@ -144,6 +145,13 @@ export const generateAPIMethods = (schema: OpenAPI3): string => {
           name: 'stream',
           stream: true,
         });
+
+        methodsToPush.push({
+          ...apiMethod,
+          name: 'streamResponse',
+          finalReturns: 'Response',
+          stream: true,
+        });
       }
 
       //
@@ -379,6 +387,7 @@ const createMethod = ({
   body,
   stream,
   returns,
+  finalReturns,
 }: MethodOptions) => {
   const params = Object.entries(parameters);
 
@@ -412,12 +421,12 @@ const createMethod = ({
 
   const parsedMethod = `(${allParameters
     .map(([key, value]) => `${convertToCamelCase(key)}: ${value}`)
-    .join(', ')}): Promise<${returns}> => {
+    .join(', ')}): Promise<${finalReturns ?? returns}> => {
   return callAPI(${encodeParam(name)}, {
     method: ${encodeParam(method)},
     endpoint: ${encodeParam(endpoint)},
     ${paramsString}${body ? `body,` : ''}options,
-  }) as Promise<${returns}>;
+  }) as Promise<${finalReturns ?? returns}>;
 }`;
 
   return parsedMethod;

src/classes/Cortex.ts

@@ -1,15 +1,11 @@
 import {APP_LESS_PARAM} from '../constants/api';
 import {STREAM_PARSERS} from '../constants/stream';
-import {
-  APIMethodRequest,
-  APIMethods,
-  createAPI,
-  paths,
-} from '../generated/openapi';
+import {APIMethodRequest, APIMethods, createAPI} from '../generated/openapi';
 import {APIFetchClient, ClientOptions, ErrorResponse} from '../types/api';
 import {createAPIFetchClient, readSSE} from '../utils/api';
 import {getObjectProperty} from '../utils/object';
 import {CortexAPIError} from './CortexAPIError';
+import {HttpStream} from './HttpStream';
 
 export class Cortex {
   private readonly apiMethods: APIMethods;
@@ -64,7 +60,11 @@ export class Cortex {
     const isPagination = name === 'list';
 
     const streamParser = isStream
-      ? STREAM_PARSERS[endpoint as keyof paths]
+      ? (STREAM_PARSERS[endpoint as keyof typeof STREAM_PARSERS] as (
+          stream: unknown,
+          event: string,
+          data: unknown,
+        ) => unknown)
       : null;
 
     if (isStream && isPagination) {
@@ -120,17 +120,35 @@ export class Cortex {
       }
 
       let resultStream: unknown = null;
+      const httpStream = name === 'streamResponse' ? new HttpStream() : null;
 
-      await readSSE(response, (event, data) => {
+      const sse = readSSE(response, (event, data) => {
         if (!event) {
           throw new Error(`Event is missing in stream for ${endpoint}`);
         }
 
+        httpStream?.writeSSE({
+          event,
+          data: data as object,
+        });
+
         resultStream = streamParser?.(resultStream, event, data);
 
         onStream?.(resultStream, event, data);
       });
 
+      if (httpStream) {
+        (async () => {
+          await sse;
+
+          httpStream.close();
+        })();
+
+        return httpStream.getResponse();
+      }
+
+      await sse;
+
       return resultStream;
     }
 

src/classes/HttpStream.ts

@@ -0,0 +1,86 @@
+export interface SSEMessage {
+  data: string | object;
+  event?: string;
+  id?: string;
+}
+
+export class HttpStream {
+  private _writer: WritableStreamDefaultWriter<Uint8Array>;
+
+  private _readable: ReadableStream;
+
+  textEncoder: InstanceType<typeof TextEncoder>;
+
+  constructor() {
+    const {readable, writable} = new TransformStream();
+
+    this._readable = readable;
+    this._writer = writable.getWriter();
+    this.textEncoder = new TextEncoder();
+
+    //
+  }
+
+  getResponse() {
+    return new Response(this._readable, {
+      headers: {
+        'Transfer-Encoding': 'chunked',
+        'Content-Type': 'text/event-stream',
+        'Cache-Control': 'no-cache',
+        Connection: 'keep-alive',
+      },
+    });
+  }
+
+  async write(input: Uint8Array | string) {
+    try {
+      if (typeof input === 'string') {
+        input = this.textEncoder.encode(input);
+      }
+      await this._writer.write(input);
+    } catch (e) {
+      if (e) {
+        console.warn('Error writing to stream', e?.toString());
+      }
+    }
+    return this;
+  }
+
+  async writeln(input: string) {
+    await this.write(input + '\n');
+    return this;
+  }
+
+  async close() {
+    try {
+      await this._writer.close();
+    } catch (e) {
+      console.warn('Error closing stream', e?.toString());
+    }
+  }
+
+  async writeSSE(message: SSEMessage) {
+    const parsedData =
+      message.data && typeof message.data === 'object'
+        ? JSON.stringify(message.data)
+        : message.data;
+
+    const data = parsedData
+      .split('\n')
+      .map(line => {
+        return `data: ${line}`;
+      })
+      .join('\n');
+
+    const sseData =
+      [
+        message.event && `event: ${message.event}`,
+        data,
+        message.id && `id: ${message.id}`,
+      ]
+        .filter(Boolean)
+        .join('\n') + '\n\n';
+
+    await this.write(sseData);
+  }
+}

src/constants/stream.ts

@@ -2,14 +2,12 @@ import {paths, RunOutputSchema, RunSchema} from '../generated/openapi';
 import {CastRunStepOutputSchema} from '../types/runs';
 import {handleStepStream} from '../utils/runs';
 
-export const STREAM_PARSERS: Partial<
-  Record<keyof paths, (current: any, event: any, data: any) => void>
-> = {
+export const STREAM_PARSERS = {
   '/apps/{app_id}/workflows/{workflow_id}/runs': (
     run: RunSchema,
     event: 'run' | 'step' | 'chunk',
     data: unknown,
-  ) => {
+  ): RunSchema => {
     switch (event) {
       case 'run': {
         return data as RunSchema;
@@ -37,19 +35,23 @@ export const STREAM_PARSERS: Partial<
         };
       }
     }
+
+    return run;
   },
 
   '/apps/{app_id}/runs/step': (
     step: CastRunStepOutputSchema,
     event: 'step' | 'chunk',
     data: unknown,
-  ) => {
+  ): CastRunStepOutputSchema => {
     switch (event) {
       case 'step': {
         return data as CastRunStepOutputSchema;
       }
       case 'chunk': {
-        if (!step) return;
+        if (!step) {
+          break;
+        }
 
         step = handleStepStream(
           event,
@@ -60,5 +62,14 @@ export const STREAM_PARSERS: Partial<
         return step;
       }
     }
+
+    return step;
   },
+} satisfies Partial<
+  Record<keyof paths, (current: any, event: any, data: any) => void>
+>;
+
+export const STREAM_PARSERS_TYPES = {
+  step: STREAM_PARSERS['/apps/{app_id}/runs/step'],
+  run: STREAM_PARSERS['/apps/{app_id}/workflows/{workflow_id}/runs'],
 };