feat(agents): add streamlit agent (#87)

Signed-off-by: Tomas Pilar <[email protected]>

migrations/Migration20241121182852.ts

@@ -0,0 +1,30 @@
+/**
+ * Copyright 2024 IBM Corp.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { Migration } from '@mikro-orm/migrations-mongodb';
+
+import { Assistant } from '@/assistants/assistant.entity';
+import { Agent } from '@/runs/execution/constants';
+
+export class Migration20241121182852 extends Migration {
+  async up(): Promise<void> {
+    await this.getCollection(Assistant).updateMany(
+      {},
+      { $set: { agent: Agent.BEE } },
+      { session: this.ctx }
+    );
+  }
+}

package.json

@@ -48,7 +48,7 @@
     "@zilliz/milvus2-sdk-node": "^2.4.4",
     "ajv": "^8.17.1",
     "axios": "^1.7.7",
-    "bee-agent-framework": "0.0.41",
+    "bee-agent-framework": "0.0.42",
     "bee-observe-connector": "0.0.5",
     "bullmq": "5.8.1",
     "cache-manager": "^5.7.6",

seeders/DatabaseSeeder.ts

@@ -36,6 +36,7 @@ import {
   PROJECT_ID_DEFAULT
 } from '@/config';
 import { redactProjectKeyValue } from '@/administration/helpers';
+import { Agent } from '@/runs/execution/constants';
 
 const USER_EXTERNAL_ID = 'test';
 const PROJECT_API_KEY = `${API_KEY_PREFIX}testkey`;
@@ -89,8 +90,9 @@ export class DatabaseSeeder extends Seeder {
       project: ref(project),
       redactedValue: redactProjectKeyValue(PROJECT_API_KEY)
     });
-    const assistant = new Assistant({
+    const beeAssistant = new Assistant({
       model: getDefaultModel(),
+      agent: Agent.BEE,
       tools: [
         {
           type: 'system',
@@ -117,7 +119,20 @@ export class DatabaseSeeder extends Seeder {
         $ui_icon: 'Bee'
       }
     });
-    em.persist([assistant, projectApiKey]);
+    const streamlitAssistant = new Assistant({
+      model: getDefaultModel(),
+      agent: Agent.STREAMLIT,
+      tools: [],
+      name: 'Builder Assistant',
+      project: ref(project),
+      createdBy: ref(projectUser),
+      description: 'An example streamlit agent, tailored for building Streamlit applications.',
+      metadata: {
+        $ui_color: 'white',
+        $ui_icon: 'Bee'
+      }
+    });
+    em.persist([beeAssistant, streamlitAssistant, projectApiKey]);
     await em.flush();
     process.env.IN_SEEDER = undefined;
   }

src/assistants/assistant.entity.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { Embedded, Entity, Property } from '@mikro-orm/core';
+import { Embedded, Entity, Enum, Property } from '@mikro-orm/core';
 
 import { ProjectScopedEntity, ProjectScopedEntityInput } from '@/common/project-scoped.entity';
 import { AnyToolUsage } from '@/tools/entities/tool-usages/tool-usage.entity.js';
@@ -28,13 +28,17 @@ import { AnyToolResource } from '@/tools/entities/tool-resources/tool-resource.e
 import { FileSearchResource } from '@/tools/entities/tool-resources/file-search-resources.entity.js';
 import { UserResource } from '@/tools/entities/tool-resources/user-resource.entity';
 import { SystemResource } from '@/tools/entities/tool-resources/system-resource.entity';
+import { Agent } from '@/runs/execution/constants';
 
 @Entity()
 export class Assistant extends ProjectScopedEntity {
   getIdPrefix(): string {
     return 'asst';
   }
 
+  @Enum(() => Agent)
+  agent: Agent;
+
   // Union must be defined in alphabetical order, otherwise Mikro-ORM won't discovered the auto-created virtual polymorphic entity
   @Embedded({ object: true })
   tools!: (CodeInterpreterUsage | FileSearchUsage | FunctionUsage | SystemUsage | UserUsage)[];
@@ -71,6 +75,7 @@ export class Assistant extends ProjectScopedEntity {
     tools,
     toolResources,
     model,
+    agent,
     topP,
     temperature,
     systemPromptOverwrite,
@@ -84,6 +89,7 @@ export class Assistant extends ProjectScopedEntity {
     this.tools = tools;
     this.toolResources = toolResources;
     this.model = model;
+    this.agent = agent;
     this.topP = topP;
     this.temperature = temperature;
     this.systemPromptOverwrite = systemPromptOverwrite;
@@ -93,5 +99,8 @@ export class Assistant extends ProjectScopedEntity {
 export type AssistantInput = ProjectScopedEntityInput & {
   tools: AnyToolUsage[];
   toolResources?: AnyToolResource[];
-} & Pick<Assistant, 'instructions' | 'name' | 'description' | 'model' | 'systemPromptOverwrite'> &
+} & Pick<
+    Assistant,
+    'instructions' | 'name' | 'description' | 'model' | 'agent' | 'systemPromptOverwrite'
+  > &
   Partial<Pick<Assistant, 'topP' | 'temperature'>>;

src/assistants/assistants.service.ts

@@ -42,6 +42,7 @@ import { Tool, ToolType } from '@/tools/entities/tool/tool.entity.js';
 import { getUpdatedValue } from '@/utils/update.js';
 import { createDeleteResponse } from '@/utils/delete.js';
 import { getDefaultModel } from '@/runs/execution/factory';
+import { Agent } from '@/runs/execution/constants.js';
 
 export function toDto(assistant: Loaded<Assistant>): AssistantDto {
   return {
@@ -55,6 +56,7 @@ export function toDto(assistant: Loaded<Assistant>): AssistantDto {
     metadata: assistant.metadata ?? {},
     created_at: dayjs(assistant.createdAt).unix(),
     model: assistant.model,
+    agent: assistant.agent,
     top_p: assistant.topP,
     temperature: assistant.temperature,
     system_prompt: assistant.systemPromptOverwrite
@@ -70,9 +72,23 @@ export async function createAssistant({
   metadata,
   top_p,
   model,
+  agent,
   temperature,
   system_prompt_overwrite
 }: AssistantCreateBody): Promise<AssistantCreateResponse> {
+  if (agent === Agent.STREAMLIT) {
+    if (toolsParam)
+      throw new APIError({
+        code: APIErrorCode.INVALID_INPUT,
+        message: 'Tools are currently not supported by Streamlit agent'
+      });
+    if (tool_resources)
+      throw new APIError({
+        code: APIErrorCode.INVALID_INPUT,
+        message: 'Tool resouces are currently not supported by Streamlit agent'
+      });
+  }
+
   const customToolIds = toolsParam.flatMap((toolUsage) =>
     toolUsage.type === ToolType.USER ? toolUsage.user.tool.id : []
   );
@@ -99,6 +115,7 @@ export async function createAssistant({
     metadata,
     topP: top_p ?? undefined,
     model: model ?? getDefaultModel(),
+    agent,
     temperature: temperature ?? undefined,
     systemPromptOverwrite: system_prompt_overwrite ?? undefined
   });
@@ -131,6 +148,20 @@ export async function updateAssistant({
   const assistant = await ORM.em.getRepository(Assistant).findOneOrFail({
     id: assistant_id
   });
+
+  if (assistant.agent === Agent.STREAMLIT) {
+    if (tools)
+      throw new APIError({
+        code: APIErrorCode.INVALID_INPUT,
+        message: 'Tools are currently not supported by Streamlit agent'
+      });
+    if (tool_resources)
+      throw new APIError({
+        code: APIErrorCode.INVALID_INPUT,
+        message: 'Tool resouces are currently not supported by Streamlit agent'
+      });
+  }
+
   assistant.name = getUpdatedValue(name, assistant.name);
   assistant.description = getUpdatedValue(description, assistant.description);
   assistant.instructions = getUpdatedValue(instructions, assistant.instructions);
@@ -157,10 +188,15 @@ export async function listAssistants({
   before,
   order,
   order_by,
+  agent,
   search
 }: AssistantsListQuery): Promise<AssistantsListResponse> {
   const where: FilterQuery<Assistant> = {};
 
+  if (agent) {
+    where.agent = agent;
+  }
+
   if (search) {
     const regexp = new RegExp(search, 'i');
     where.$or = [{ description: regexp }, { name: regexp }];

src/assistants/dtos/assistant-create.ts

@@ -21,6 +21,7 @@ import { assistantSchema } from './assistant.js';
 import { toolUsageSchema } from '@/tools/dtos/tools-usage.js';
 import { metadataSchema } from '@/schema.js';
 import { toolResourcesSchema } from '@/tools/dtos/tool-resources.js';
+import { Agent } from '@/runs/execution/constants.js';
 
 export const assistantCreateBodySchema = {
   type: 'object',
@@ -55,6 +56,11 @@ export const assistantCreateBodySchema = {
     model: {
       type: 'string'
     },
+    agent: {
+      type: 'string',
+      enum: Object.values(Agent),
+      default: Agent.BEE
+    },
     top_p: {
       type: 'number',
       nullable: true,

src/assistants/dtos/assistant.ts

@@ -19,6 +19,7 @@ import { FromSchema, JSONSchema } from 'json-schema-to-ts';
 import { toolUsageSchema } from '@/tools/dtos/tools-usage.js';
 import { metadataSchema } from '@/schema.js';
 import { toolResourcesSchema } from '@/tools/dtos/tool-resources.js';
+import { Agent } from '@/runs/execution/constants';
 
 export const assistantSchema = {
   type: 'object',
@@ -31,7 +32,8 @@ export const assistantSchema = {
     'description',
     'metadata',
     'created_at',
-    'model'
+    'model',
+    'agent'
   ],
   properties: {
     id: { type: 'string' },
@@ -62,6 +64,10 @@ export const assistantSchema = {
     model: {
       type: 'string'
     },
+    agent: {
+      type: 'string',
+      enum: Object.values(Agent)
+    },
     top_p: {
       type: 'number',
       nullable: true

src/assistants/dtos/assistants-list.ts

@@ -19,6 +19,7 @@ import { FromSchema, JSONSchema } from 'json-schema-to-ts';
 import { assistantSchema } from './assistant.js';
 
 import { createPaginationQuerySchema, withPagination } from '@/schema.js';
+import { Agent } from '@/runs/execution/constants.js';
 
 export const assistantsListQuerySchema = {
   type: 'object',
@@ -29,6 +30,10 @@ export const assistantsListQuerySchema = {
       type: 'boolean',
       default: false
     },
+    agent: {
+      type: 'string',
+      enum: Object.values(Agent)
+    },
     search: {
       type: 'string',
       nullable: true

src/runs/execution/constants.ts

@@ -17,6 +17,12 @@
 export const RUN_EXPIRATION_MILLISECONDS = 10 * 60 * 1000;
 export const STATUS_POLL_INTERVAL = 5 * 1000;
 
+export const Agent = {
+  BEE: 'bee',
+  STREAMLIT: 'streamlit'
+} as const;
+export type Agent = (typeof Agent)[keyof typeof Agent];
+
 export const LLMBackend = {
   OLLAMA: 'ollama',
   IBM_VLLM: 'ibm-vllm',