Skip to content

Commit

Permalink
♻️ refactor: refactor the Google Gen AI (lobehub#4484)
Browse files Browse the repository at this point in the history
  • Loading branch information
@arvinxx
arvinxx authored Oct 25, 2024
1 parent 8f83863 commit b890e8d
Show file tree
Hide file tree
Showing 5 changed files with 21 additions and 151 deletions.
2 changes: 1 addition & 1 deletion package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@
"@clerk/themes": "^2.1.37",
"@codesandbox/sandpack-react": "^2.19.9",
"@cyntler/react-doc-viewer": "^1.17.0",
"@google/generative-ai": "^0.16.1",
"@google/generative-ai": "^0.21.0",
"@huggingface/inference": "^2.8.1",
"@icons-pack/react-simple-icons": "9.6.0",
"@khmyznikov/pwa-install": "^0.3.9",
Expand Down
16 changes: 8 additions & 8 deletions src/app/(main)/chat/(workspace)/_layout/Desktop/Portal.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,11 @@ import { rgba } from 'polished';
import { PropsWithChildren, memo } from 'react';
import { Flexbox } from 'react-layout-kit';

import { CHAT_DOCK_TOOL_UI_WIDTH, CHAT_DOCK_WIDTH, MAX_WIDTH } from '@/const/layoutTokens';
import {
CHAT_PORTAL_MAX_WIDTH,
CHAT_PORTAL_TOOL_UI_WIDTH,
CHAT_PORTAL_WIDTH,
} from '@/const/layoutTokens';
import { useChatStore } from '@/store/chat';
import { chatPortalSelectors } from '@/store/chat/slices/portal/selectors';

Expand All @@ -25,12 +29,8 @@ const useStyles = createStyles(({ css, token, isDarkMode }) => ({
`,
panel: css`
overflow: hidden;
height: 100%;
margin: 4px;
background: ${isDarkMode ? rgba(token.colorBgElevated, 0.8) : token.colorBgElevated};
border-radius: 8px;
`,
}));

Expand All @@ -53,8 +53,8 @@ const PortalPanel = memo(({ children }: PropsWithChildren) => {
}}
expand
hanlderStyle={{ display: 'none' }}
maxWidth={MAX_WIDTH}
minWidth={showArtifactUI || showToolUI ? CHAT_DOCK_TOOL_UI_WIDTH : CHAT_DOCK_WIDTH}
maxWidth={CHAT_PORTAL_MAX_WIDTH}
minWidth={showArtifactUI || showToolUI ? CHAT_PORTAL_TOOL_UI_WIDTH : CHAT_PORTAL_WIDTH}
mode={md ? 'fixed' : 'float'}
placement={'right'}
showHandlerWhenUnexpand={false}
Expand All @@ -65,7 +65,7 @@ const PortalPanel = memo(({ children }: PropsWithChildren) => {
flex: 'none',
height: '100%',
maxHeight: '100vh',
minWidth: CHAT_DOCK_WIDTH,
minWidth: CHAT_PORTAL_WIDTH,
}}
>
<Flexbox className={styles.panel}>{children}</Flexbox>
Expand Down
5 changes: 3 additions & 2 deletions src/const/layoutTokens.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,9 @@ export const CHAT_TEXTAREA_HEIGHT = 160;
export const CHAT_TEXTAREA_HEIGHT_MOBILE = 108;
export const CHAT_SIDEBAR_WIDTH = 280;

export const CHAT_DOCK_WIDTH = 400;
export const CHAT_DOCK_TOOL_UI_WIDTH = 600;
export const CHAT_PORTAL_WIDTH = 400;
export const CHAT_PORTAL_MAX_WIDTH = 1280;
export const CHAT_PORTAL_TOOL_UI_WIDTH = 600;

export const MARKET_SIDEBAR_WIDTH = 400;
export const FOLDER_WIDTH = 270;
Expand Down
95 changes: 4 additions & 91 deletions src/libs/agent-runtime/google/index.test.ts
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
// @vitest-environment edge-runtime
import { FunctionDeclarationSchemaType, FunctionDeclarationsTool } from '@google/generative-ai';
import { JSONSchema7 } from 'json-schema';
import { FunctionDeclarationsTool } from '@google/generative-ai';
import OpenAI from 'openai';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';

Expand Down Expand Up @@ -479,103 +478,17 @@ describe('LobeGoogleAI', () => {
name: 'testTool',
description: 'A test tool',
parameters: {
type: FunctionDeclarationSchemaType.OBJECT,
type: 'object',
properties: {
param1: { type: FunctionDeclarationSchemaType.STRING },
param2: { type: FunctionDeclarationSchemaType.NUMBER },
param1: { type: 'string' },
param2: { type: 'number' },
},
required: ['param1'],
},
});
});
});

describe('convertSchemaObject', () => {
it('should correctly convert object schema', () => {
const schema: JSONSchema7 = {
type: 'object',
properties: {
prop1: { type: 'string' },
prop2: { type: 'number' },
},
};

const converted = instance['convertSchemaObject'](schema);

expect(converted).toEqual({
type: FunctionDeclarationSchemaType.OBJECT,
properties: {
prop1: { type: FunctionDeclarationSchemaType.STRING },
prop2: { type: FunctionDeclarationSchemaType.NUMBER },
},
});
});

it('should correctly convert nested schema', () => {
const schema: JSONSchema7 = {
type: 'object',
properties: {
nested: {
type: 'array',
items: {
type: 'object',
properties: {
prop: { type: 'string' },
},
},
},
},
};

const converted = instance['convertSchemaObject'](schema);

expect(converted).toEqual({
type: FunctionDeclarationSchemaType.OBJECT,
properties: {
nested: {
type: FunctionDeclarationSchemaType.ARRAY,
items: {
type: FunctionDeclarationSchemaType.OBJECT,
properties: {
prop: { type: FunctionDeclarationSchemaType.STRING },
},
},
},
},
});
});

it('should correctly convert array schema', () => {
const schema: JSONSchema7 = {
type: 'array',
items: { type: 'string' },
};
const converted = instance['convertSchemaObject'](schema);
expect(converted).toEqual({
type: FunctionDeclarationSchemaType.ARRAY,
items: { type: FunctionDeclarationSchemaType.STRING },
});
});

it('should correctly convert string schema', () => {
const schema: JSONSchema7 = { type: 'string' };
const converted = instance['convertSchemaObject'](schema);
expect(converted).toEqual({ type: FunctionDeclarationSchemaType.STRING });
});

it('should correctly convert number schema', () => {
const schema: JSONSchema7 = { type: 'number' };
const converted = instance['convertSchemaObject'](schema);
expect(converted).toEqual({ type: FunctionDeclarationSchemaType.NUMBER });
});

it('should correctly convert boolean schema', () => {
const schema: JSONSchema7 = { type: 'boolean' };
const converted = instance['convertSchemaObject'](schema);
expect(converted).toEqual({ type: FunctionDeclarationSchemaType.BOOLEAN });
});
});

describe('convertOAIMessagesToGoogleMessage', () => {
it('should correctly convert assistant message', async () => {
const message: OpenAIChatMessage = {
Expand Down
54 changes: 5 additions & 49 deletions src/libs/agent-runtime/google/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,11 @@ import {
Content,
FunctionCallPart,
FunctionDeclaration,
FunctionDeclarationSchemaProperty,
FunctionDeclarationSchemaType,
Tool as GoogleFunctionCallTool,
GoogleGenerativeAI,
Part,
SchemaType,
} from '@google/generative-ai';
import { JSONSchema7 } from 'json-schema';
import { transform } from 'lodash-es';

import { imageUrlToBase64 } from '@/utils/imageToBase64';
import { safeParseJSON } from '@/utils/safeParseJSON';
Expand Down Expand Up @@ -190,13 +187,12 @@ export class LobeGoogleAI implements LobeRuntimeAI {
};
};

// convert messages from the Vercel AI SDK Format to the format
// that is expected by the Google GenAI SDK
// convert messages from the OpenAI format to Google GenAI SDK
private buildGoogleMessages = async (
messages: OpenAIChatMessage[],
model: string,
): Promise<Content[]> => {
// if the model is gemini-1.0 we don't need to pair messages
// if the model is gemini-1.0 we need to pair messages
if (model.startsWith('gemini-1.0')) {
const contents: Content[] = [];
let lastRole = 'model';
Expand Down Expand Up @@ -298,52 +294,12 @@ export class LobeGoogleAI implements LobeRuntimeAI {
name: functionDeclaration.name,
parameters: {
description: parameters?.description,
properties: transform(parameters?.properties, (result, value, key: string) => {
result[key] = this.convertSchemaObject(value as JSONSchema7);
}),
properties: parameters?.properties,
required: parameters?.required,
type: FunctionDeclarationSchemaType.OBJECT,
type: SchemaType.OBJECT,
},
};
};

private convertSchemaObject(schema: JSONSchema7): FunctionDeclarationSchemaProperty {
switch (schema.type) {
default:
case 'object': {
return {
...schema,
properties: Object.fromEntries(
Object.entries(schema.properties || {}).map(([key, value]) => [
key,
this.convertSchemaObject(value as JSONSchema7),
]),
),
type: FunctionDeclarationSchemaType.OBJECT,
} as any;
}

case 'array': {
return {
...schema,
items: this.convertSchemaObject(schema.items as JSONSchema7),
type: FunctionDeclarationSchemaType.ARRAY,
} as any;
}

case 'string': {
return { ...schema, type: FunctionDeclarationSchemaType.STRING } as any;
}

case 'number': {
return { ...schema, type: FunctionDeclarationSchemaType.NUMBER } as any;
}

case 'boolean': {
return { ...schema, type: FunctionDeclarationSchemaType.BOOLEAN } as any;
}
}
}
}

export default LobeGoogleAI;
Expand Down

0 comments on commit b890e8d

Please sign in to comment.