feat: Add BraveSearchAPI document loader (#3486)

2026-06-28 13:00:56 +03:00 · 2024-11-08 14:59:22 +02:00
parent 275c9e3a1e
commit ccc0088d0d
2 changed files with 149 additions and 0 deletions
@@ -0,0 +1,141 @@
+import { omit } from 'lodash'
+import { ICommonObject, IDocument, INode, INodeData, INodeParams } from '../../../src/Interface'
+import { TextSplitter } from 'langchain/text_splitter'
+import { BraveSearch } from '@langchain/community/tools/brave_search'
+import { getCredentialData, getCredentialParam } from '../../../src/utils'
+import { Document } from '@langchain/core/documents'
+
+class BraveSearchAPI_DocumentLoaders implements INode {
+    label: string
+    name: string
+    version: number
+    description: string
+    type: string
+    icon: string
+    category: string
+    baseClasses: string[]
+    credential: INodeParams
+    inputs: INodeParams[]
+
+    constructor() {
+        this.label = 'BraveSearch API Document Loader'
+        this.name = 'braveSearchApiLoader'
+        this.version = 1.0
+        this.type = 'Document'
+        this.icon = 'brave.svg'
+        this.category = 'Document Loaders'
+        this.description = 'Load and process data from BraveSearch results'
+        this.baseClasses = [this.type]
+        this.credential = {
+            label: 'Connect Credential',
+            name: 'credential',
+            type: 'credential',
+            optional: false,
+            credentialNames: ['braveSearchApi']
+        }
+        this.inputs = [
+            {
+                label: 'Query',
+                name: 'query',
+                type: 'string'
+            },
+            {
+                label: 'Text Splitter',
+                name: 'textSplitter',
+                type: 'TextSplitter',
+                optional: true
+            },
+            {
+                label: 'Additional Metadata',
+                name: 'metadata',
+                type: 'json',
+                description: 'Additional metadata to be added to the extracted documents',
+                optional: true,
+                additionalParams: true
+            },
+            {
+                label: 'Omit Metadata Keys',
+                name: 'omitMetadataKeys',
+                type: 'string',
+                rows: 4,
+                description:
+                    'Each document loader comes with a default set of metadata keys that are extracted from the document. You can use this field to omit some of the default metadata keys. The value should be a list of keys, separated by comma. Use * to omit all metadata keys except the ones you specify in the Additional Metadata field',
+                placeholder: 'key1, key2, key3.nestedKey1',
+                optional: true,
+                additionalParams: true
+            }
+        ]
+    }
+
+    async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
+        const textSplitter = nodeData.inputs?.textSplitter as TextSplitter
+        const query = nodeData.inputs?.query as string
+        const metadata = nodeData.inputs?.metadata
+        const _omitMetadataKeys = nodeData.inputs?.omitMetadataKeys as string
+
+        let omitMetadataKeys: string[] = []
+        if (_omitMetadataKeys) {
+            omitMetadataKeys = _omitMetadataKeys.split(',').map((key) => key.trim())
+        }
+
+        const credentialData = await getCredentialData(nodeData.credential ?? '', options)
+        const braveApiKey = getCredentialParam('braveApiKey', credentialData, nodeData)
+        const loader = new BraveSearch({ apiKey: braveApiKey })
+
+        let docs: IDocument[] = []
+        const searchResults = await loader._call(query) // Use _call method for search
+        const parsedResults = JSON.parse(searchResults) // Parse the JSON string to get documents
+
+        // Format the results to match the expected Document structure
+        docs = parsedResults.map(
+            (result: any) =>
+                new Document({
+                    pageContent: result.snippet, // Assuming snippet is the content
+                    metadata: {
+                        title: result.title,
+                        link: result.link
+                    }
+                })
+        )
+
+        if (textSplitter) {
+            docs = await textSplitter.splitDocuments(docs)
+        }
+
+        if (metadata) {
+            const parsedMetadata = typeof metadata === 'object' ? metadata : JSON.parse(metadata)
+            docs = docs.map((doc) => ({
+                ...doc,
+                metadata:
+                    _omitMetadataKeys === '*'
+                        ? {
+                              ...parsedMetadata
+                          }
+                        : omit(
+                              {
+                                  ...doc.metadata,
+                                  ...parsedMetadata
+                              },
+                              omitMetadataKeys
+                          )
+            }))
+        } else {
+            docs = docs.map((doc) => ({
+                ...doc,
+                metadata:
+                    _omitMetadataKeys === '*'
+                        ? {}
+                        : omit(
+                              {
+                                  ...doc.metadata
+                              },
+                              omitMetadataKeys
+                          )
+            }))
+        }
+
+        return docs
+    }
+}
+
+module.exports = { nodeClass: BraveSearchAPI_DocumentLoaders }