Merge pull request #1719 from benjidotsh/rag-ai/improve-new-backend-s…

…upport fix(rag-ai): improve support for the new backend system
RoadieHQ · Nov 14, 2024 · e3f9d8c · e3f9d8c
2 parents 07ce45b + 14057d9
commit e3f9d8c
Show file tree

Hide file tree

Showing 8 changed files with 39 additions and 9 deletions.
diff --git a/.changeset/silent-years-tap.md b/.changeset/silent-years-tap.md
@@ -0,0 +1,7 @@
+---
+'@roadiehq/rag-ai-backend-retrieval-augmenter': patch
+'@roadiehq/rag-ai-backend-embeddings-aws': patch
+'@roadiehq/rag-ai-node': patch
+---
+
+Improve support for new backend system
diff --git a/plugins/backend/rag-ai-backend-embeddings-aws/package.json b/plugins/backend/rag-ai-backend-embeddings-aws/package.json
@@ -44,6 +44,7 @@
     "@aws-sdk/client-bedrock-runtime": "^3.602.0",
     "@aws-sdk/types": "^3.468.0",
     "@backstage/backend-common": "^0.25.0",
+    "@backstage/backend-plugin-api": "^1.0.1",
     "@backstage/catalog-client": "^1.7.1",
     "@backstage/config": "^1.2.0",
     "@langchain/aws": "^0.0.10",

diff --git a/plugins/backend/rag-ai-backend-embeddings-aws/src/RoadieBedrockAugmenter.ts b/plugins/backend/rag-ai-backend-embeddings-aws/src/RoadieBedrockAugmenter.ts
@@ -13,7 +13,6 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { TokenManager } from '@backstage/backend-common';
 import { BedrockEmbeddings } from '@langchain/aws';
 import { AwsCredentialIdentity, Provider } from '@aws-sdk/types';
 import {
@@ -34,7 +33,6 @@ export class RoadieBedrockAugmenter extends DefaultVectorAugmentationIndexer {
         region: string;
       };
       bedrockConfig: BedrockConfig;
-      tokenManager: TokenManager;
     },
   ) {
     const embeddingsConfig = {

diff --git a/plugins/backend/rag-ai-backend-embeddings-aws/src/index.ts b/plugins/backend/rag-ai-backend-embeddings-aws/src/index.ts
@@ -25,10 +25,11 @@ import { PluginEndpointDiscovery } from '@backstage/backend-common';
 import { AwsCredentialIdentity, Provider } from '@aws-sdk/types';
 import { Config } from '@backstage/config';
 import { AugmentationOptions } from '@roadiehq/rag-ai-backend-retrieval-augmenter';
+import { LoggerService } from '@backstage/backend-plugin-api';
 
 export interface RoadieBedrockEmbeddingsConfig {
-  logger: Logger;
-  tokenManager: TokenManager;
+  logger: Logger | LoggerService;
+  tokenManager?: TokenManager;
   vectorStore: RoadieVectorStore;
   catalogApi: CatalogApi;
   discovery: PluginEndpointDiscovery;

diff --git a/plugins/backend/rag-ai-backend-retrieval-augmenter/src/indexing/types.ts b/plugins/backend/rag-ai-backend-retrieval-augmenter/src/indexing/types.ts
@@ -28,7 +28,7 @@ export type AugmentationOptions = {
 
 export interface RoadieEmbeddingsConfig {
   logger: LoggerService;
-  tokenManager: TokenManager;
+  tokenManager?: TokenManager;
   vectorStore: RoadieVectorStore;
   catalogApi: CatalogApi;
   discovery: PluginEndpointDiscovery;

diff --git a/plugins/backend/rag-ai-backend/src/plugin.ts b/plugins/backend/rag-ai-backend/src/plugin.ts
@@ -29,6 +29,7 @@ import {
   modelExtensionPoint,
 } from '@roadiehq/rag-ai-node';
 import { BaseLLM } from '@langchain/core/language_models/llms';
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 
 /**
  * Rag AI backend plugin
@@ -40,7 +41,7 @@ export const ragAiPlugin = createBackendPlugin({
   register(env) {
     let augmentationIndexer: AugmentationIndexer | undefined;
     let retrievalPipeline: RetrievalPipeline | undefined;
-    let model: BaseLLM | undefined;
+    let model: BaseLLM | BaseChatModel | undefined;
 
     env.registerExtensionPoint(augmentationIndexerExtensionPoint, {
       setAugmentationIndexer(indexer) {

diff --git a/plugins/backend/rag-ai-node/src/extensions.ts b/plugins/backend/rag-ai-node/src/extensions.ts
@@ -16,6 +16,7 @@
 
 import { AugmentationIndexer, RetrievalPipeline } from './types';
 import { createExtensionPoint } from '@backstage/backend-plugin-api';
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { BaseLLM } from '@langchain/core/language_models/llms';
 
 export interface AugmentationIndexerExtensionPoint {
@@ -37,7 +38,7 @@ export const retrievalPipelineExtensionPoint =
   });
 
 export interface ModelExtensionPoint {
-  setBaseLLM(baseLLM: BaseLLM): void;
+  setBaseLLM(baseLLM: BaseLLM | BaseChatModel): void;
 }
 
 export const modelExtensionPoint = createExtensionPoint<ModelExtensionPoint>({

diff --git a/yarn.lock b/yarn.lock
@@ -14589,13 +14589,20 @@
   resolved "https://registry.yarnpkg.com/@types/range-parser/-/range-parser-1.2.7.tgz#50ae4353eaaddc04044279812f52c8c65857dbcb"
   integrity sha512-hKormJbkJqzQGhziax5PItDUTMAM9uE2XXQmM37dyd4hVM+5aVl7oVxMVUiVQn2oCQFN/LKCZdvSM0pFRqbSmQ==
 
-"@types/react-dom@*", "@types/react-dom@<18.0.0", "@types/react-dom@^18", "@types/react-dom@^18.0.0":
+"@types/react-dom@*", "@types/react-dom@^18.0.0":
   version "18.3.0"
   resolved "https://registry.yarnpkg.com/@types/react-dom/-/react-dom-18.3.0.tgz#0cbc818755d87066ab6ca74fbedb2547d74a82b0"
   integrity sha512-EhwApuTmMBmXuFOikhQLIBUn6uFg81SwLMOAUgodJF14SOBOCMdU04gDoYi0WOJJHD144TL32z4yDqCW3dnkQg==
   dependencies:
     "@types/react" "*"
 
+"@types/react-dom@<18.0.0":
+  version "17.0.25"
+  resolved "https://registry.yarnpkg.com/@types/react-dom/-/react-dom-17.0.25.tgz#e0e5b3571e1069625b3a3da2b279379aa33a0cb5"
+  integrity sha512-urx7A7UxkZQmThYA4So0NelOVjx3V4rNFVJwp0WZlbIK5eM4rNJDiN3R/E9ix0MBh6kAEojk/9YL+Te6D9zHNA==
+  dependencies:
+    "@types/react" "^17"
+
 "@types/react-redux@^7.1.20":
   version "7.1.33"
   resolved "https://registry.yarnpkg.com/@types/react-redux/-/react-redux-7.1.33.tgz#53c5564f03f1ded90904e3c90f77e4bd4dc20b15"
@@ -14627,14 +14634,23 @@
   dependencies:
     "@types/react" "*"
 
-"@types/react@*", "@types/react@^16.13.1 || ^17.0.0", "@types/react@^16.13.1 || ^17.0.0 || ^18.0.0", "@types/react@^18":
+"@types/react@*", "@types/react@^16.13.1 || ^17.0.0 || ^18.0.0":
   version "18.3.11"
   resolved "https://registry.yarnpkg.com/@types/react/-/react-18.3.11.tgz#9d530601ff843ee0d7030d4227ea4360236bd537"
   integrity sha512-r6QZ069rFTjrEYgFdOck1gK7FLVsgJE7tTz0pQBczlBNUhBNk0MQH4UbnFSwjpQLMkLzgqvBBa+qGpLje16eTQ==
   dependencies:
     "@types/prop-types" "*"
     csstype "^3.0.2"
 
+"@types/react@^16.13.1 || ^17.0.0", "@types/react@^17":
+  version "17.0.83"
+  resolved "https://registry.yarnpkg.com/@types/react/-/react-17.0.83.tgz#b477c56387b74279281149dcf5ba2a1e2216d131"
+  integrity sha512-l0m4ArKJvmFtR4e8UmKrj1pB4tUgOhJITf+mADyF/p69Ts1YAR/E+G9XEM0mHXKVRa1dQNHseyyDNzeuAXfXQw==
+  dependencies:
+    "@types/prop-types" "*"
+    "@types/scheduler" "^0.16"
+    csstype "^3.0.2"
+
 "@types/recharts@^1.8.14":
   version "1.8.29"
   resolved "https://registry.yarnpkg.com/@types/recharts/-/recharts-1.8.29.tgz#5e117521a65bf015b808350b45b65553ff5011f3"
@@ -14682,6 +14698,11 @@
   dependencies:
     htmlparser2 "^4.1.0"
 
+"@types/scheduler@^0.16":
+  version "0.16.8"
+  resolved "https://registry.yarnpkg.com/@types/scheduler/-/scheduler-0.16.8.tgz#ce5ace04cfeabe7ef87c0091e50752e36707deff"
+  integrity sha512-WZLiwShhwLRmeV6zH+GkbOFT6Z6VklCItrDioxUnv+u4Ll+8vKeFySoFyK/0ctcRpOmwAicELfmys1sDc/Rw+A==
+
 "@types/[email protected]", "@types/semver@^7.5.0":
   version "7.5.8"
   resolved "https://registry.yarnpkg.com/@types/semver/-/semver-7.5.8.tgz#8268a8c57a3e4abd25c165ecd36237db7948a55e"