From 6d0037d94b5682466dabd4d006cb109d7456a70f Mon Sep 17 00:00:00 2001
From: DarkSky <darksky2048@gmail.com>
Date: Wed, 20 Mar 2024 19:02:38 +0800
Subject: [PATCH] feat: improve logging & error handle

---
 .../backend/server/src/config/affine.self.ts  |  5 +++
 .../server/src/plugins/copilot/provider.ts    | 35 ++++++++++++-------
 packages/backend/server/src/plugins/index.ts  |  1 +
 packages/backend/server/src/schema.gql        |  1 +
 4 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/packages/backend/server/src/config/affine.self.ts b/packages/backend/server/src/config/affine.self.ts
index cecc33fbbc08e..5cc19b37f7529 100644
--- a/packages/backend/server/src/config/affine.self.ts
+++ b/packages/backend/server/src/config/affine.self.ts
@@ -38,6 +38,11 @@ if (env.R2_OBJECT_STORAGE_ACCOUNT_ID) {
   }`;
 }
 
+AFFiNE.plugins.use('copilot', {
+  openai: {
+    apiKey: 'test',
+  },
+});
 AFFiNE.plugins.use('redis');
 AFFiNE.plugins.use('payment', {
   stripe: {
diff --git a/packages/backend/server/src/plugins/copilot/provider.ts b/packages/backend/server/src/plugins/copilot/provider.ts
index 2bbaf9f13146b..24bf67ba63f55 100644
--- a/packages/backend/server/src/plugins/copilot/provider.ts
+++ b/packages/backend/server/src/plugins/copilot/provider.ts
@@ -1,6 +1,6 @@
 import assert from 'node:assert';
 
-import { Injectable } from '@nestjs/common';
+import { Injectable, Logger } from '@nestjs/common';
 
 import { Config } from '../../fundamentals';
 import {
@@ -27,7 +27,7 @@ interface CopilotProviderDefinition<C extends CopilotProviderConfig> {
 // registered provider factory
 const COPILOT_PROVIDER = new Map<
   CopilotProviderType,
-  (config: Config) => CopilotProvider
+  (config: Config, logger: Logger) => CopilotProvider
 >();
 
 // map of capabilities to providers
@@ -42,26 +42,36 @@ const ASSERT_CONFIG = new Map<CopilotProviderType, (config: Config) => void>();
 export function registerCopilotProvider<
   C extends CopilotProviderConfig = CopilotProviderConfig,
 >(provider: CopilotProviderDefinition<C>) {
-  const factory = (config: Config) => {
-    assert(config.plugins.copilot);
-    assert(config.plugins.copilot[provider.type]);
+  const type = provider.type;
+
+  const factory = (config: Config, logger: Logger) => {
+    const providerConfig = config.plugins.copilot?.[type];
+    if (!provider.assetsConfig(providerConfig as C)) {
+      throw new Error(
+        `Invalid configuration for copilot provider ${type}: ${providerConfig}`
+      );
+    }
+    const instance = new provider(providerConfig as C);
+    logger.log(
+      `Copilot provider ${type} registered, capabilities: ${provider.capabilities.join(', ')}`
+    );
 
-    return new provider(config.plugins.copilot[provider.type] as C);
+    return instance;
   };
   // register the provider
-  COPILOT_PROVIDER.set(provider.type, factory);
+  COPILOT_PROVIDER.set(type, factory);
   // register the provider capabilities
   for (const capability of provider.capabilities) {
     const providers = PROVIDER_CAPABILITY_MAP.get(capability) || [];
-    if (!providers.includes(provider.type)) {
-      providers.push(provider.type);
+    if (!providers.includes(type)) {
+      providers.push(type);
     }
     PROVIDER_CAPABILITY_MAP.set(capability, providers);
   }
   // register the provider config assertion
-  ASSERT_CONFIG.set(provider.type, (config: Config) => {
+  ASSERT_CONFIG.set(type, (config: Config) => {
     assert(config.plugins.copilot);
-    const providerConfig = config.plugins.copilot[provider.type];
+    const providerConfig = config.plugins.copilot[type];
     if (!providerConfig) return false;
     return provider.assetsConfig(providerConfig as C);
   });
@@ -78,6 +88,7 @@ export function assertProvidersConfigs(config: Config) {
 
 @Injectable()
 export class CopilotProviderService {
+  private readonly logger = new Logger(CopilotProviderService.name);
   constructor(private readonly config: Config) {}
 
   private readonly cachedProviders = new Map<
@@ -93,7 +104,7 @@ export class CopilotProviderService {
       throw new Error(`Unknown copilot provider type: ${provider}`);
     }
 
-    return providerFactory(this.config);
+    return providerFactory(this.config, this.logger);
   }
 
   getProvider(provider: CopilotProviderType): CopilotProvider {
diff --git a/packages/backend/server/src/plugins/index.ts b/packages/backend/server/src/plugins/index.ts
index 42ea147ad3be6..9d82b90c10128 100644
--- a/packages/backend/server/src/plugins/index.ts
+++ b/packages/backend/server/src/plugins/index.ts
@@ -1,3 +1,4 @@
+import './copilot';
 import './gcloud';
 import './oauth';
 import './payment';
diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql
index 2d18df9e2759b..36820c7a89cd2 100644
--- a/packages/backend/server/src/schema.gql
+++ b/packages/backend/server/src/schema.gql
@@ -262,6 +262,7 @@ enum ServerDeploymentType {
 }
 
 enum ServerFeature {
+  Copilot
   OAuth
   Payment
 }