(modelName: string, apiKey: string)
| 158 | } |
| 159 | |
| 160 | function buildGeminiProvider(modelName: string, apiKey: string): ResolvedProvider['buildRequest'] { |
| 161 | return (inputs, inputType) => ({ |
| 162 | apiUrl: `https://generativelanguage.googleapis.com/v1beta/models/${modelName}:batchEmbedContents`, |
| 163 | headers: { |
| 164 | 'Content-Type': 'application/json', |
| 165 | 'x-goog-api-key': apiKey, |
| 166 | }, |
| 167 | body: { |
| 168 | requests: inputs.map((text) => ({ |
| 169 | model: `models/${modelName}`, |
| 170 | content: { parts: [{ text }] }, |
| 171 | taskType: inputType === 'query' ? 'RETRIEVAL_QUERY' : 'RETRIEVAL_DOCUMENT', |
| 172 | outputDimensionality: EMBEDDING_DIMENSIONS, |
| 173 | })), |
| 174 | }, |
| 175 | parse: (json) => { |
| 176 | const data = json as { embeddings: Array<{ values: number[] }> } |
| 177 | return data.embeddings.map((item) => l2Normalize(item.values)) |
| 178 | }, |
| 179 | }) |
| 180 | } |
| 181 | |
| 182 | /** |
| 183 | * Returns the embedding model to use for new knowledge bases. |
no test coverage detected