feat(lightspeed): add api client (#2020)

feat(lightspeed): Add api client

Author: divyanshiGupta

plugins/lightspeed/package.json

@@ -4,7 +4,6 @@
   "main": "src/index.ts",
   "types": "src/index.ts",
   "license": "Apache-2.0",
-  "private": true,
   "publishConfig": {
     "access": "public",
     "main": "dist/index.esm.js",

plugins/lightspeed/src/api/LightspeedProxyClient.test.ts

@@ -0,0 +1,107 @@
+import { ConfigApi, IdentityApi } from '@backstage/core-plugin-api';
+
+import OpenAI from 'openai';
+
+import { LightspeedProxyClient, Options } from './LightspeedProxyClient';
+
+jest.mock('openai');
+
+describe('LightspeedProxyClient', () => {
+  let configApi: jest.Mocked<ConfigApi>;
+  let identityApi: jest.Mocked<IdentityApi>;
+  let lightspeedClient: LightspeedProxyClient;
+  let mockCreateChatCompletion: jest.Mock;
+
+  beforeEach(() => {
+    configApi = {
+      getString: jest.fn().mockReturnValue('http://localhost:7007'),
+    } as unknown as jest.Mocked<ConfigApi>;
+
+    identityApi = {
+      getCredentials: jest.fn().mockResolvedValue({ token: 'mock-token' }),
+    } as unknown as jest.Mocked<IdentityApi>;
+
+    const options: Options = { configApi, identityApi };
+    lightspeedClient = new LightspeedProxyClient(options);
+
+    mockCreateChatCompletion = jest.fn().mockResolvedValue({} as any);
+    const mockOpenAI = new OpenAI({});
+    mockOpenAI.chat = {
+      completions: {
+        create: mockCreateChatCompletion,
+      },
+    } as any;
+
+    (OpenAI as unknown as jest.Mock).mockImplementation(() => mockOpenAI);
+  });
+
+  describe('constructor', () => {
+    it('should initialize the client with the correct baseURL', () => {
+      expect(OpenAI).toHaveBeenCalledWith({
+        baseURL: 'http://localhost:7007/api/proxy/lightspeed/api',
+        apiKey: 'random-key',
+        dangerouslyAllowBrowser: true,
+      });
+    });
+  });
+
+  describe('getUserAuthorization', () => {
+    it('should return the idToken from identityApi', async () => {
+      const token = await lightspeedClient.getUserAuthorization();
+      expect(token).toBe('mock-token');
+      expect(identityApi.getCredentials).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  describe('createChatCompletions', () => {
+    it('should call openAIApi.chat.completions.create with correct parameters', async () => {
+      const prompt = 'Test prompt';
+      await lightspeedClient.createChatCompletions(prompt);
+
+      expect(mockCreateChatCompletion).toHaveBeenCalledWith(
+        {
+          messages: [
+            {
+              role: 'system',
+              content:
+                'You are a helpful assistant that can answer question in Red Hat Developer Hub.',
+            },
+            { role: 'user', content: prompt },
+          ],
+          model: 'llama3',
+          stream: true,
+        },
+        {
+          headers: {
+            Authorization: 'Bearer mock-token',
+          },
+        },
+      );
+    });
+
+    it('should not include Authorization header if idToken is not available', async () => {
+      identityApi.getCredentials.mockResolvedValueOnce({ token: undefined });
+
+      const prompt = 'Test prompt';
+      await lightspeedClient.createChatCompletions(prompt);
+
+      expect(mockCreateChatCompletion).toHaveBeenCalledWith(
+        {
+          messages: [
+            {
+              role: 'system',
+              content:
+                'You are a helpful assistant that can answer question in Red Hat Developer Hub.',
+            },
+            { role: 'user', content: prompt },
+          ],
+          model: 'llama3',
+          stream: true,
+        },
+        {
+          headers: {},
+        },
+      );
+    });
+  });
+});

plugins/lightspeed/src/api/LightspeedProxyClient.ts

@@ -0,0 +1,70 @@
+import {
+  ConfigApi,
+  createApiRef,
+  IdentityApi,
+} from '@backstage/core-plugin-api';
+
+import OpenAI from 'openai';
+import { Stream } from 'openai/streaming';
+
+export type LightspeedAPI = {
+  createChatCompletions: (
+    prompt: string,
+  ) => Promise<Stream<OpenAI.Chat.Completions.ChatCompletionChunk>>;
+};
+
+export const lightspeedApiRef = createApiRef<LightspeedAPI>({
+  id: 'plugin.lightspeed.service',
+});
+
+export type Options = {
+  configApi: ConfigApi;
+  identityApi: IdentityApi;
+};
+
+export class LightspeedProxyClient implements LightspeedAPI {
+  // @ts-ignore
+  private readonly configApi: ConfigApi;
+  private readonly identityApi: IdentityApi;
+  private readonly openAIApi: OpenAI;
+
+  constructor(options: Options) {
+    this.configApi = options.configApi;
+    this.identityApi = options.identityApi;
+    this.openAIApi = new OpenAI({
+      baseURL: `${this.configApi.getString('backend.baseUrl')}/api/proxy/lightspeed/api`,
+
+      // required but ignored
+      apiKey: 'random-key',
+      dangerouslyAllowBrowser: true,
+    });
+  }
+
+  async getUserAuthorization() {
+    const { token: idToken } = await this.identityApi.getCredentials();
+    return idToken;
+  }
+
+  async createChatCompletions(prompt: string) {
+    const idToken = await this.getUserAuthorization();
+    return await this.openAIApi.chat.completions.create(
+      {
+        messages: [
+          {
+            role: 'system',
+            content:
+              'You are a helpful assistant that can answer question in Red Hat Developer Hub.',
+          },
+          { role: 'user', content: prompt },
+        ],
+        model: 'llama3',
+        stream: true,
+      },
+      {
+        headers: {
+          ...(idToken && { Authorization: `Bearer ${idToken}` }),
+        },
+      },
+    );
+  }
+}

plugins/lightspeed/src/components/LightspeedPage.tsx

@@ -1,12 +1,12 @@
 import React from 'react';
 
 import { Content, Header, HeaderLabel, Page } from '@backstage/core-components';
-import { configApiRef, useApi } from '@backstage/core-plugin-api';
+import { useApi } from '@backstage/core-plugin-api';
 
 import { Paper } from '@material-ui/core';
 import { createStyles, makeStyles } from '@material-ui/core/styles';
-import OpenAI from 'openai';
 
+import { lightspeedApiRef } from '../api/LightspeedProxyClient';
 import { LightspeedInput } from './LightspeedInput';
 import { SystemMessage, UserMessage } from './Message';
 
@@ -38,51 +38,31 @@ const useStyles = makeStyles(() =>
 export const LightspeedPage = () => {
   const classes = useStyles();
 
-  const configApi = useApi(configApiRef);
+  const lightspeedApi = useApi(lightspeedApiRef);
 
   const [, setChunkIndex] = React.useState(0);
   const [prompts, setPrompts] = React.useState<string[]>([]);
   const [completions, setCompletions] = React.useState<{
     [key: string]: string;
   }>({});
-  const backendUrl = configApi.getString('backend.baseUrl');
-  const openai = new OpenAI({
-    baseURL: `${backendUrl}/api/proxy/lightspeed/api`,
-
-    // required but ignored
-    apiKey: 'random-key',
-    dangerouslyAllowBrowser: true,
-  });
 
   const handleInputPrompt = React.useCallback(
     async (prompt: string) => {
       setPrompts(p => [...p, prompt]);
       setChunkIndex(0);
 
-      const result = await openai.chat.completions.create({
-        messages: [
-          {
-            role: 'system',
-            content:
-              'You are a helpful assistant that can answer question in Red Hat Developer Hub.',
-          },
-          { role: 'user', content: prompt },
-        ],
-        model: 'llama3',
-        stream: true,
-      });
+      const result = await lightspeedApi.createChatCompletions(prompt);
 
       for await (const chunk of result) {
         setChunkIndex(index => index + 1);
         setCompletions(c => {
-          // console.log('string ---', s);
           c[prompt] =
             `${c[prompt] || ''}${chunk.choices[0]?.delta?.content || ''}`;
           return c;
         });
       }
     },
-    [openai.chat.completions],
+    [lightspeedApi],
   );
 
   return (

plugins/lightspeed/src/plugin.ts

@@ -1,15 +1,33 @@
 import {
+  configApiRef,
+  createApiFactory,
   createPlugin,
   createRoutableExtension,
+  identityApiRef,
 } from '@backstage/core-plugin-api';
 
+import {
+  lightspeedApiRef,
+  LightspeedProxyClient,
+} from './api/LightspeedProxyClient';
 import { rootRouteRef } from './routes';
 
 export const lightspeedPlugin = createPlugin({
   id: 'lightspeed',
   routes: {
     root: rootRouteRef,
   },
+  apis: [
+    createApiFactory({
+      api: lightspeedApiRef,
+      deps: {
+        configApi: configApiRef,
+        identityApi: identityApiRef,
+      },
+      factory: ({ configApi, identityApi }) =>
+        new LightspeedProxyClient({ configApi, identityApi }),
+    }),
+  ],
 });
 
 export const LightspeedPage = lightspeedPlugin.provide(