UNPKG

@langchain/community

Version:
1 lines 9.33 kB
{"version":3,"file":"supabase.cjs","names":["BaseRetriever","Document"],"sources":["../../src/retrievers/supabase.ts"],"sourcesContent":["import type { SupabaseClient } from \"@supabase/supabase-js\";\nimport type { EmbeddingsInterface } from \"@langchain/core/embeddings\";\nimport { Document } from \"@langchain/core/documents\";\nimport {\n BaseRetriever,\n type BaseRetrieverInput,\n} from \"@langchain/core/retrievers\";\nimport {\n CallbackManagerForRetrieverRun,\n Callbacks,\n} from \"@langchain/core/callbacks/manager\";\n\ninterface SearchEmbeddingsParams {\n query_embedding: number[];\n match_count: number; // int\n filter?: Record<string, unknown>; // jsonb\n}\n\ninterface SearchKeywordParams {\n query_text: string;\n match_count: number; // int\n}\n\ninterface SearchResponseRow {\n id: number;\n content: string;\n metadata: object;\n similarity: number;\n}\n\ntype SearchResult = [Document, number, number];\n\nexport interface SupabaseLibArgs extends BaseRetrieverInput {\n client: SupabaseClient;\n /**\n * The table name on Supabase. Defaults to \"documents\".\n */\n tableName?: string;\n /**\n * The name of the Similarity search function on Supabase. Defaults to \"match_documents\".\n */\n similarityQueryName?: string;\n /**\n * The name of the Keyword search function on Supabase. Defaults to \"kw_match_documents\".\n */\n keywordQueryName?: string;\n /**\n * The number of documents to return from the similarity search. Defaults to 2.\n */\n similarityK?: number;\n /**\n * The number of documents to return from the keyword search. Defaults to 2.\n */\n keywordK?: number;\n}\n\nexport interface SupabaseHybridSearchParams {\n query: string;\n similarityK: number;\n keywordK: number;\n}\n\n/**\n * Class for performing hybrid search operations on a Supabase database.\n * It extends the `BaseRetriever` class and implements methods for\n * similarity search, keyword search, and hybrid search.\n */\nexport class SupabaseHybridSearch extends BaseRetriever {\n static lc_name() {\n return \"SupabaseHybridSearch\";\n }\n\n lc_namespace = [\"langchain\", \"retrievers\", \"supabase\"];\n\n similarityK: number;\n\n query: string;\n\n keywordK: number;\n\n similarityQueryName: string;\n\n client: SupabaseClient;\n\n tableName: string;\n\n keywordQueryName: string;\n\n embeddings: EmbeddingsInterface;\n\n constructor(embeddings: EmbeddingsInterface, args: SupabaseLibArgs) {\n super(args);\n this.embeddings = embeddings;\n this.client = args.client;\n this.tableName = args.tableName || \"documents\";\n this.similarityQueryName = args.similarityQueryName || \"match_documents\";\n this.keywordQueryName = args.keywordQueryName || \"kw_match_documents\";\n this.similarityK = args.similarityK || 2;\n this.keywordK = args.keywordK || 2;\n }\n\n /**\n * Performs a similarity search on the Supabase database using the\n * provided query and returns the top 'k' similar documents.\n * @param query The query to use for the similarity search.\n * @param k The number of top similar documents to return.\n * @param _callbacks Optional callbacks to pass to the embedQuery method.\n * @returns A promise that resolves to an array of search results. Each result is a tuple containing a Document, its similarity score, and its ID.\n */\n protected async similaritySearch(\n query: string,\n k: number,\n _callbacks?: Callbacks // implement passing to embedQuery later\n ): Promise<SearchResult[]> {\n const embeddedQuery = await this.embeddings.embedQuery(query);\n\n const matchDocumentsParams: SearchEmbeddingsParams = {\n query_embedding: embeddedQuery,\n match_count: k,\n };\n\n if (Object.keys(this.metadata ?? {}).length > 0) {\n matchDocumentsParams.filter = this.metadata;\n }\n\n const { data: searches, error } = await this.client.rpc(\n this.similarityQueryName,\n matchDocumentsParams\n );\n\n if (error) {\n throw new Error(\n `Error searching for documents: ${error.code} ${error.message} ${error.details}`\n );\n }\n\n return (searches as SearchResponseRow[]).map((resp) => [\n new Document({\n metadata: resp.metadata,\n pageContent: resp.content,\n }),\n resp.similarity,\n resp.id,\n ]);\n }\n\n /**\n * Performs a keyword search on the Supabase database using the provided\n * query and returns the top 'k' documents that match the keywords.\n * @param query The query to use for the keyword search.\n * @param k The number of top documents to return that match the keywords.\n * @returns A promise that resolves to an array of search results. Each result is a tuple containing a Document, its similarity score multiplied by 10, and its ID.\n */\n protected async keywordSearch(\n query: string,\n k: number\n ): Promise<SearchResult[]> {\n const kwMatchDocumentsParams: SearchKeywordParams = {\n query_text: query,\n match_count: k,\n };\n\n const { data: searches, error } = await this.client.rpc(\n this.keywordQueryName,\n kwMatchDocumentsParams\n );\n\n if (error) {\n throw new Error(\n `Error searching for documents: ${error.code} ${error.message} ${error.details}`\n );\n }\n\n return (searches as SearchResponseRow[]).map((resp) => [\n new Document({\n metadata: resp.metadata,\n pageContent: resp.content,\n }),\n resp.similarity * 10,\n resp.id,\n ]);\n }\n\n /**\n * Combines the results of the `similaritySearch` and `keywordSearch`\n * methods and returns the top 'k' documents based on a combination of\n * similarity and keyword matching.\n * @param query The query to use for the hybrid search.\n * @param similarityK The number of top similar documents to return.\n * @param keywordK The number of top documents to return that match the keywords.\n * @param callbacks Optional callbacks to pass to the similaritySearch method.\n * @returns A promise that resolves to an array of search results. Each result is a tuple containing a Document, its combined score, and its ID.\n */\n protected async hybridSearch(\n query: string,\n similarityK: number,\n keywordK: number,\n callbacks?: Callbacks\n ): Promise<SearchResult[]> {\n const similarity_search = this.similaritySearch(\n query,\n similarityK,\n callbacks\n );\n\n const keyword_search = this.keywordSearch(query, keywordK);\n\n return Promise.all([similarity_search, keyword_search])\n .then((results) => results.flat())\n .then((results) => {\n const picks = new Map<number, SearchResult>();\n\n results.forEach((result) => {\n const id = result[2];\n const nextScore = result[1];\n const prevScore = picks.get(id)?.[1];\n\n if (prevScore === undefined || nextScore > prevScore) {\n picks.set(id, result);\n }\n });\n\n return Array.from(picks.values());\n })\n .then((results) => results.sort((a, b) => b[1] - a[1]));\n }\n\n async _getRelevantDocuments(\n query: string,\n runManager?: CallbackManagerForRetrieverRun\n ): Promise<Document[]> {\n const searchResults = await this.hybridSearch(\n query,\n this.similarityK,\n this.keywordK,\n runManager?.getChild(\"hybrid_search\")\n );\n\n return searchResults.map(([doc]) => doc);\n }\n}\n"],"mappings":";;;;;;;;;;;AAmEA,IAAa,uBAAb,cAA0CA,2BAAAA,cAAc;CACtD,OAAO,UAAU;AACf,SAAO;;CAGT,eAAe;EAAC;EAAa;EAAc;EAAW;CAEtD;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA,YAAY,YAAiC,MAAuB;AAClE,QAAM,KAAK;AACX,OAAK,aAAa;AAClB,OAAK,SAAS,KAAK;AACnB,OAAK,YAAY,KAAK,aAAa;AACnC,OAAK,sBAAsB,KAAK,uBAAuB;AACvD,OAAK,mBAAmB,KAAK,oBAAoB;AACjD,OAAK,cAAc,KAAK,eAAe;AACvC,OAAK,WAAW,KAAK,YAAY;;;;;;;;;;CAWnC,MAAgB,iBACd,OACA,GACA,YACyB;EAGzB,MAAM,uBAA+C;GACnD,iBAHoB,MAAM,KAAK,WAAW,WAAW,MAAM;GAI3D,aAAa;GACd;AAED,MAAI,OAAO,KAAK,KAAK,YAAY,EAAE,CAAC,CAAC,SAAS,EAC5C,sBAAqB,SAAS,KAAK;EAGrC,MAAM,EAAE,MAAM,UAAU,UAAU,MAAM,KAAK,OAAO,IAClD,KAAK,qBACL,qBACD;AAED,MAAI,MACF,OAAM,IAAI,MACR,kCAAkC,MAAM,KAAK,GAAG,MAAM,QAAQ,GAAG,MAAM,UACxE;AAGH,SAAQ,SAAiC,KAAK,SAAS;GACrD,IAAIC,0BAAAA,SAAS;IACX,UAAU,KAAK;IACf,aAAa,KAAK;IACnB,CAAC;GACF,KAAK;GACL,KAAK;GACN,CAAC;;;;;;;;;CAUJ,MAAgB,cACd,OACA,GACyB;EACzB,MAAM,yBAA8C;GAClD,YAAY;GACZ,aAAa;GACd;EAED,MAAM,EAAE,MAAM,UAAU,UAAU,MAAM,KAAK,OAAO,IAClD,KAAK,kBACL,uBACD;AAED,MAAI,MACF,OAAM,IAAI,MACR,kCAAkC,MAAM,KAAK,GAAG,MAAM,QAAQ,GAAG,MAAM,UACxE;AAGH,SAAQ,SAAiC,KAAK,SAAS;GACrD,IAAIA,0BAAAA,SAAS;IACX,UAAU,KAAK;IACf,aAAa,KAAK;IACnB,CAAC;GACF,KAAK,aAAa;GAClB,KAAK;GACN,CAAC;;;;;;;;;;;;CAaJ,MAAgB,aACd,OACA,aACA,UACA,WACyB;EACzB,MAAM,oBAAoB,KAAK,iBAC7B,OACA,aACA,UACD;EAED,MAAM,iBAAiB,KAAK,cAAc,OAAO,SAAS;AAE1D,SAAO,QAAQ,IAAI,CAAC,mBAAmB,eAAe,CAAC,CACpD,MAAM,YAAY,QAAQ,MAAM,CAAC,CACjC,MAAM,YAAY;GACjB,MAAM,wBAAQ,IAAI,KAA2B;AAE7C,WAAQ,SAAS,WAAW;IAC1B,MAAM,KAAK,OAAO;IAClB,MAAM,YAAY,OAAO;IACzB,MAAM,YAAY,MAAM,IAAI,GAAG,GAAG;AAElC,QAAI,cAAc,KAAA,KAAa,YAAY,UACzC,OAAM,IAAI,IAAI,OAAO;KAEvB;AAEF,UAAO,MAAM,KAAK,MAAM,QAAQ,CAAC;IACjC,CACD,MAAM,YAAY,QAAQ,MAAM,GAAG,MAAM,EAAE,KAAK,EAAE,GAAG,CAAC;;CAG3D,MAAM,sBACJ,OACA,YACqB;AAQrB,UAPsB,MAAM,KAAK,aAC/B,OACA,KAAK,aACL,KAAK,UACL,YAAY,SAAS,gBAAgB,CACtC,EAEoB,KAAK,CAAC,SAAS,IAAI"}