-
Notifications
You must be signed in to change notification settings - Fork 2.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
community[minor]: Add Dria retriever (#4302)
* initial builds and setups * examples, integration & docs * small doc fix
- Loading branch information
Showing
8 changed files
with
252 additions
and
3 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,39 @@ | ||
| --- | ||
| hide_table_of_contents: true | ||
| --- | ||
|
|
||
| # Dria Retriever | ||
|
|
||
| The [Dria](https://dria.co/profile) retriever allows an agent to perform a text-based search across a comprehensive knowledge hub. | ||
|
|
||
| ## Setup | ||
|
|
||
| To use Dria retriever, first install Dria JS client: | ||
|
|
||
| ```bash npm2yarn | ||
| npm install dria | ||
| ``` | ||
|
|
||
| You need to provide two things to the retriever: | ||
|
|
||
| - **API Key**: you can get yours at your [profile page](https://dria.co/profile) when you create an account. | ||
| - **Contract ID**: accessible at the top of the page when viewing a knowledge or in its URL. | ||
| For example, the Bitcoin whitepaper is uploaded on Dria at https://dria.co/knowledge/2KxNbEb040GKQ1DSDNDsA-Fsj_BlQIEAlzBNuiapBR0, so its contract ID is `2KxNbEb040GKQ1DSDNDsA-Fsj_BlQIEAlzBNuiapBR0`. | ||
| Contract ID can be omitted during instantiation, and later be set via `dria.contractId = "your-contract"` | ||
|
|
||
| Dria retriever exposes the underlying [Dria client](https://npmjs.com/package/dria) as well, refer to the [Dria documentation](https://github.com/firstbatchxyz/dria-js-client?tab=readme-ov-file#usage) to learn more about the client. | ||
|
|
||
| ## Usage | ||
|
|
||
| import IntegrationInstallTooltip from "@mdx_components/integration_install_tooltip.mdx"; | ||
|
|
||
| <IntegrationInstallTooltip></IntegrationInstallTooltip> | ||
|
|
||
| ```bash npm2yarn | ||
| npm install dria @langchain/community | ||
| ``` | ||
|
|
||
| import CodeBlock from "@theme/CodeBlock"; | ||
| import Example from "@examples/retrievers/dria.ts"; | ||
|
|
||
| <CodeBlock language="typescript">{Example}</CodeBlock> |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,14 @@ | ||
| import { DriaRetriever } from "@langchain/community/retrievers/dria"; | ||
|
|
||
| // contract of TypeScript Handbook v4.9 uploaded to Dria | ||
| // https://dria.co/knowledge/-B64DjhUtCwBdXSpsRytlRQCu-bie-vSTvTIT8Ap3g0 | ||
| const contractId = "-B64DjhUtCwBdXSpsRytlRQCu-bie-vSTvTIT8Ap3g0"; | ||
|
|
||
| const retriever = new DriaRetriever({ | ||
| contractId, // a knowledge to connect to | ||
| apiKey: "DRIA_API_KEY", // if not provided, will check env for `DRIA_API_KEY` | ||
| topK: 15, // optional: default value is 10 | ||
| }); | ||
|
|
||
| const docs = await retriever.getRelevantDocuments("What is a union type?"); | ||
| console.log(docs); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,122 @@ | ||
| import { | ||
| BaseRetriever, | ||
| type BaseRetrieverInput, | ||
| } from "@langchain/core/retrievers"; | ||
| import { Document } from "@langchain/core/documents"; | ||
| import { getEnvironmentVariable } from "@langchain/core/utils/env"; | ||
| import type { DriaParams, SearchOptions as DriaSearchOptions } from "dria"; | ||
| import { Dria } from "dria"; | ||
|
|
||
| /** | ||
| * Configurations for Dria retriever. | ||
| * | ||
| * - `contractId`: a Dria knowledge's contract ID. | ||
| * - `apiKey`: a Dria API key; if omitted, the retriever will check for `DRIA_API_KEY` environment variable. | ||
| * | ||
| * The retrieval can be configured with the following options: | ||
| * | ||
| * - `topK`: number of results to return, max 20. (default: 10) | ||
| * - `rerank`: re-rank the results from most to least semantically relevant to the given search query. (default: true) | ||
| * - `level`: level of detail for the search, must be an integer from 0 to 5 (inclusive). (default: 1) | ||
| * - `field`: CSV field name, only relevant for the CSV files. | ||
| */ | ||
| export interface DriaRetrieverArgs | ||
| extends DriaParams, | ||
| BaseRetrieverInput, | ||
| DriaSearchOptions {} | ||
|
|
||
| /** | ||
| * Class for retrieving documents from knowledge uploaded to Dria. | ||
| * | ||
| * @example | ||
| * ```typescript | ||
| * // contract of TypeScript Handbook v4.9 uploaded to Dria | ||
| * const contractId = "-B64DjhUtCwBdXSpsRytlRQCu-bie-vSTvTIT8Ap3g0"; | ||
| * const retriever = new DriaRetriever({ contractId }); | ||
| * | ||
| * const docs = await retriever.getRelevantDocuments("What is a union type?"); | ||
| * console.log(docs); | ||
| * ``` | ||
| */ | ||
| export class DriaRetriever extends BaseRetriever { | ||
| static lc_name() { | ||
| return "DriaRetriever"; | ||
| } | ||
|
|
||
| lc_namespace = ["langchain", "retrievers", "dria"]; | ||
|
|
||
| get lc_secrets() { | ||
| return { apiKey: "DRIA_API_KEY" }; | ||
| } | ||
|
|
||
| get lc_aliases() { | ||
| return { apiKey: "api_key" }; | ||
| } | ||
|
|
||
| apiKey: string; | ||
|
|
||
| public driaClient: Dria; | ||
|
|
||
| private searchOptions: DriaSearchOptions; | ||
|
|
||
| constructor(fields: DriaRetrieverArgs) { | ||
| super(fields); | ||
|
|
||
| const apiKey = fields.apiKey ?? getEnvironmentVariable("DRIA_API_KEY"); | ||
| if (!apiKey) throw new Error("Missing DRIA_API_KEY."); | ||
| this.apiKey = apiKey; | ||
|
|
||
| this.searchOptions = { | ||
| topK: fields.topK, | ||
| field: fields.field, | ||
| rerank: fields.rerank, | ||
| level: fields.level, | ||
| }; | ||
|
|
||
| this.driaClient = new Dria({ | ||
| contractId: fields.contractId, | ||
| apiKey: this.apiKey, | ||
| }); | ||
| } | ||
|
|
||
| /** | ||
| * Currently connected knowledge on Dria. | ||
| * | ||
| * Retriever will use this contract ID while retrieving documents, | ||
| * and will throw an error if `undefined`. | ||
| * | ||
| * In the case that this is `undefined`, the user is expected to | ||
| * set contract ID manually, such as after creating a new knowledge & inserting | ||
| * data there with the Dria client. | ||
| */ | ||
| get contractId(): string | undefined { | ||
| return this.driaClient.contractId; | ||
| } | ||
|
|
||
| set contractId(value: string) { | ||
| this.driaClient.contractId = value; | ||
| } | ||
|
|
||
| /** | ||
| * Retrieves documents from Dria with respect to the configured contract ID, based on | ||
| * the given query string. | ||
| * | ||
| * @param query The query string | ||
| * @returns A promise that resolves to an array of documents, with page content as text, | ||
| * along with `id` and the relevance `score` within the metadata. | ||
| */ | ||
| async _getRelevantDocuments(query: string): Promise<Document[]> { | ||
| const docs = await this.driaClient.search(query, this.searchOptions); | ||
| return docs.map( | ||
| (d) => | ||
| new Document({ | ||
| // dria.search returns a string within the metadata as the content | ||
| pageContent: d.metadata, | ||
| metadata: { | ||
| id: d.id, | ||
| score: d.score, | ||
| }, | ||
| }) | ||
| ); | ||
| } | ||
| } |
16 changes: 16 additions & 0 deletions
16
libs/langchain-community/src/retrievers/tests/dria.int.test.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,16 @@ | ||
| import { test, expect } from "@jest/globals"; | ||
| import { DriaRetriever } from "../dria.js"; | ||
|
|
||
| test.skip("DriaRetriever", async () => { | ||
| // contract of TypeScript Handbook v4.9 uploaded to Dria | ||
| // https://dria.co/knowledge/-B64DjhUtCwBdXSpsRytlRQCu-bie-vSTvTIT8Ap3g0 | ||
| const contractId = "-B64DjhUtCwBdXSpsRytlRQCu-bie-vSTvTIT8Ap3g0"; | ||
| const topK = 10; | ||
|
|
||
| const retriever = new DriaRetriever({ contractId, topK }); | ||
|
|
||
| const docs = await retriever.getRelevantDocuments("What is a union type?"); | ||
| expect(docs.length).toBe(topK); | ||
|
|
||
| console.log(docs[0].pageContent); | ||
| }); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters