1- import { LanguageModel } from '../types' ;
1+ import { EmbeddingModel , LanguageModel } from '../types' ;
22import { InvalidModelIdError } from './invalid-model-id-error' ;
33import { NoSuchModelError } from './no-such-model-error' ;
44import { NoSuchProviderError } from './no-such-provider-error' ;
@@ -19,86 +19,113 @@ The model id is then passed to the provider function to get the model.
1919@returns {LanguageModel } The language model associated with the id.
2020 */
2121 languageModel ( id : string ) : LanguageModel ;
22+
23+ /**
24+ Returns the text embedding model with the given id in the format `providerId:modelId`.
25+ The model id is then passed to the provider function to get the model.
26+
27+ @param {string } id - The id of the model to return.
28+
29+ @throws {NoSuchModelError } If no model with the given id exists.
30+ @throws {NoSuchProviderError } If no provider with the given id exists.
31+
32+ @returns {LanguageModel } The language model associated with the id.
33+ */
34+ textEmbeddingModel ( id : string ) : EmbeddingModel < string > ;
2235} ;
2336
2437/**
2538 * @deprecated Use `experimental_ProviderRegistry` instead.
2639 */
2740export type experimental_ModelRegistry = experimental_ProviderRegistry ;
2841
42+ /**
43+ * Provider for language and text embedding models. Compatible with the
44+ * provider registry.
45+ */
46+ interface Provider {
47+ /**
48+ * Returns a language model with the given id.
49+ */
50+ languageModel ?: ( modelId : string ) => LanguageModel ;
51+
52+ /**
53+ * Returns a text embedding model with the given id.
54+ */
55+ textEmbedding ?: ( modelId : string ) => EmbeddingModel < string > ;
56+ }
57+
2958/**
3059 * Creates a registry for the given providers.
3160 */
3261export function experimental_createProviderRegistry (
33- providers : Record < string , ( id : string ) => LanguageModel > ,
62+ providers : Record < string , Provider > ,
3463) : experimental_ProviderRegistry {
3564 const registry = new DefaultProviderRegistry ( ) ;
3665
3766 for ( const [ id , provider ] of Object . entries ( providers ) ) {
38- registry . registerLanguageModelProvider ( { id, provider } ) ;
67+ registry . registerProvider ( { id, provider } ) ;
3968 }
4069
4170 return registry ;
4271}
4372
44- class DefaultProviderRegistry implements experimental_ProviderRegistry {
45- // Mapping of provider id to provider
46- private providers : Record < string , ( id : string ) => LanguageModel > = { } ;
73+ /**
74+ * @deprecated Use `experimental_createProviderRegistry` instead.
75+ */
76+ export const experimental_createModelRegistry =
77+ experimental_createProviderRegistry ;
4778
48- /**
49- Registers a language model provider with a given id.
79+ class DefaultProviderRegistry implements experimental_ProviderRegistry {
80+ private providers : Record < string , Provider > = { } ;
5081
51- @param {string } id - The id of the provider.
52- @param {(id: string) => LanguageModel } provider - The provider function to register.
53- */
54- registerLanguageModelProvider ( {
55- id,
56- provider,
57- } : {
58- id : string ;
59- provider : ( id : string ) => LanguageModel ;
60- } ) : void {
82+ registerProvider ( { id, provider } : { id : string ; provider : Provider } ) : void {
6183 this . providers [ id ] = provider ;
6284 }
6385
64- /**
65- Returns the language model with the given id.
66- The id can either be a registered model id or use a provider prefix.
67- Provider ids are separated from the model id by a colon: `providerId:modelId`.
68- The model id is then passed to the provider function to get the model.
86+ private getProvider ( id : string ) : Provider {
87+ const provider = this . providers [ id ] ;
6988
70- @param {string } id - The id of the model to return.
89+ if ( provider == null ) {
90+ throw new NoSuchProviderError ( {
91+ providerId : id ,
92+ availableProviders : Object . keys ( this . providers ) ,
93+ } ) ;
94+ }
7195
72- @throws { NoSuchModelError } If no model with the given id exists.
73- @throws { NoSuchProviderError } If no provider with the given id exists.
96+ return provider ;
97+ }
7498
75- @returns {LanguageModel } The language model associated with the id.
76- */
77- languageModel ( id : string ) : LanguageModel {
99+ private splitId ( id : string ) : [ string , string ] {
78100 if ( ! id . includes ( ':' ) ) {
79101 throw new InvalidModelIdError ( { id } ) ;
80102 }
81103
82- const [ providerId , modelId ] = id . split ( ':' ) ;
104+ return id . split ( ':' ) as [ string , string ] ;
105+ }
83106
84- const provider = this . providers [ providerId ] ;
107+ languageModel ( id : string ) : LanguageModel {
108+ const [ providerId , modelId ] = this . splitId ( id ) ;
109+ const model = this . getProvider ( providerId ) . languageModel ?.( modelId ) ;
85110
86- if ( ! provider ) {
87- throw new NoSuchProviderError ( { providerId } ) ;
111+ if ( model == null ) {
112+ throw new NoSuchModelError ( { modelId : id , modelType : 'language model' } ) ;
88113 }
89114
90- const model = provider ( modelId ) ;
115+ return model ;
116+ }
117+
118+ textEmbeddingModel ( id : string ) : EmbeddingModel < string > {
119+ const [ providerId , modelId ] = this . splitId ( id ) ;
120+ const model = this . getProvider ( providerId ) . textEmbedding ?.( modelId ) ;
91121
92- if ( ! model ) {
93- throw new NoSuchModelError ( { modelId : id } ) ;
122+ if ( model == null ) {
123+ throw new NoSuchModelError ( {
124+ modelId : id ,
125+ modelType : 'text embedding model' ,
126+ } ) ;
94127 }
95128
96129 return model ;
97130 }
98131}
99-
100- /**
101- * @deprecated Use `experimental_createProviderRegistry` instead.
102- */
103- export const experimental_createModelRegistry =
104- experimental_createProviderRegistry ;
0 commit comments