|
1 |
| -import { OpenAI, OpenAIEmbedding, Settings } from "llamaindex"; |
2 |
| - |
3 |
| -import { |
4 |
| - DefaultAzureCredential, |
5 |
| - getBearerTokenProvider, |
6 |
| -} from "@azure/identity"; |
7 |
| -import { OllamaEmbedding } from "llamaindex/embeddings/OllamaEmbedding"; |
8 |
| -import { Ollama } from "llamaindex/llm/ollama"; |
| 1 | +import { Settings } from "llamaindex"; |
| 2 | +import { setupProvider } from "./provider"; |
9 | 3 |
|
10 | 4 | const CHUNK_SIZE = 512;
|
11 | 5 | const CHUNK_OVERLAP = 20;
|
12 |
| -const AZURE_COGNITIVE_SERVICES_SCOPE = |
13 |
| - "https://cognitiveservices.azure.com/.default"; |
14 | 6 |
|
15 | 7 | export const initSettings = async () => {
|
16 | 8 | console.log(`Using '${process.env.MODEL_PROVIDER}' model provider`);
|
17 | 9 |
|
18 |
| - // if provider is OpenAI, MODEL must be set |
19 |
| - if (process.env.MODEL_PROVIDER === 'openai' && process.env.OPENAI_API_TYPE !== 'AzureOpenAI' && !process.env.MODEL) { |
20 |
| - throw new Error("'MODEL' env variable must be set."); |
21 |
| - } |
22 |
| - |
23 |
| - // if provider is Azure OpenAI, AZURE_DEPLOYMENT_NAME must be set |
24 |
| - if (process.env.MODEL_PROVIDER === 'openai' && process.env.OPENAI_API_TYPE === 'AzureOpenAI' && !process.env.AZURE_DEPLOYMENT_NAME) { |
25 |
| - throw new Error("'AZURE_DEPLOYMENT_NAME' env variables must be set."); |
| 10 | + if (!process.env.MODEL || !process.env.EMBEDDING_MODEL) { |
| 11 | + throw new Error("'MODEL' and 'EMBEDDING_MODEL' env variables must be set."); |
26 | 12 | }
|
27 | 13 |
|
28 |
| - if (!process.env.EMBEDDING_MODEL) { |
29 |
| - throw new Error("'EMBEDDING_MODEL' env variable must be set."); |
30 |
| - } |
31 |
| - |
32 |
| - switch (process.env.MODEL_PROVIDER) { |
33 |
| - case "ollama": |
34 |
| - initOllama(); |
35 |
| - break; |
36 |
| - case "openai": |
37 |
| - if (process.env.OPENAI_API_TYPE === "AzureOpenAI") { |
38 |
| - await initAzureOpenAI(); |
39 |
| - } else { |
40 |
| - initOpenAI(); |
41 |
| - } |
42 |
| - break; |
43 |
| - default: |
44 |
| - throw new Error( |
45 |
| - `Model provider '${process.env.MODEL_PROVIDER}' not supported.`, |
46 |
| - ); |
47 |
| - } |
48 | 14 | Settings.chunkSize = CHUNK_SIZE;
|
49 | 15 | Settings.chunkOverlap = CHUNK_OVERLAP;
|
50 |
| -}; |
51 |
| - |
52 |
| -function initOpenAI() { |
53 |
| - Settings.llm = new OpenAI({ |
54 |
| - model: process.env.MODEL ?? "gpt-3.5-turbo", |
55 |
| - maxTokens: 512, |
56 |
| - }); |
57 |
| - Settings.embedModel = new OpenAIEmbedding({ |
58 |
| - model: process.env.EMBEDDING_MODEL, |
59 |
| - dimensions: process.env.EMBEDDING_DIM |
60 |
| - ? parseInt(process.env.EMBEDDING_DIM) |
61 |
| - : undefined, |
62 |
| - }); |
63 |
| -} |
64 |
| - |
65 |
| -async function initAzureOpenAI() { |
66 |
| - const credential = new DefaultAzureCredential(); |
67 |
| - const azureADTokenProvider = getBearerTokenProvider( |
68 |
| - credential, |
69 |
| - AZURE_COGNITIVE_SERVICES_SCOPE, |
70 |
| - ); |
71 |
| - |
72 |
| - const azure = { |
73 |
| - azureADTokenProvider, |
74 |
| - deployment: process.env.AZURE_OPENAI_DEPLOYMENT ?? "gpt-35-turbo", |
75 |
| - }; |
76 |
| - |
77 |
| - // configure LLM model |
78 |
| - Settings.llm = new OpenAI({ |
79 |
| - azure, |
80 |
| - }) as any; |
81 |
| - |
82 |
| - // configure embedding model |
83 |
| - azure.deployment = process.env.EMBEDDING_MODEL as string; |
84 |
| - Settings.embedModel = new OpenAIEmbedding({ |
85 |
| - azure, |
86 |
| - model: process.env.EMBEDDING_MODEL, |
87 |
| - dimensions: process.env.EMBEDDING_DIM |
88 |
| - ? parseInt(process.env.EMBEDDING_DIM) |
89 |
| - : undefined, |
90 |
| - }); |
91 |
| -} |
92 |
| - |
93 |
| -function initOllama() { |
94 |
| - const config = { |
95 |
| - host: process.env.OLLAMA_BASE_URL ?? "http://127.0.0.1:11434", |
96 |
| - }; |
97 |
| - Settings.llm = new Ollama({ |
98 |
| - model: process.env.MODEL ?? "", |
99 |
| - config, |
100 |
| - }); |
101 |
| - Settings.embedModel = new OllamaEmbedding({ |
102 |
| - model: process.env.EMBEDDING_MODEL ?? "", |
103 |
| - config, |
104 |
| - }); |
105 |
| -} |
106 | 16 |
|
| 17 | + setupProvider(); |
| 18 | +}; |
0 commit comments