-
Notifications
You must be signed in to change notification settings - Fork 23
/
env_manager.py
87 lines (79 loc) · 2.49 KB
/
env_manager.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
import os
from dotenv import load_dotenv
from logger import logger
from translation import (
BaseTranslationClass,
BhashiniTranslationClass,
DhruvaTranslationClass,
GoogleCloudTranslationClass
)
from storage import (
BaseStorageClass,
AwsS3BucketClass,
GcpBucketClass,
OciBucketClass
)
from llm import (
BaseChatClient,
AzureChatClient,
OpenAIChatClient,
OllamaChatClient
)
from vectorstores import (
BaseVectorStore,
MarqoVectorStore
)
class EnvironmentManager():
"""
Class for initializing functions respective to the env variable provided
"""
def __init__(self):
load_dotenv()
self.indexes = {
"llm": {
"class": {
"openai": OpenAIChatClient,
"azure": AzureChatClient,
"ollama": OllamaChatClient
},
"env_key": "LLM_TYPE"
},
"translate": {
"class": {
"bhashini": BhashiniTranslationClass,
"google": GoogleCloudTranslationClass,
"dhruva": DhruvaTranslationClass
},
"env_key": "TRANSLATION_TYPE"
},
"storage": {
"class": {
"oci": OciBucketClass,
"gcp": GcpBucketClass,
"aws": AwsS3BucketClass
},
"env_key": "BUCKET_TYPE"
},
"vectorstore": {
"class": {
"marqo": MarqoVectorStore
},
"env_key": "VECTOR_STORE_TYPE"
}
}
def create_instance(self, env_key):
env_var = self.indexes[env_key]["env_key"]
type_value = os.getenv(env_var)
if type_value is None:
raise ValueError(
f"Missing credentials. Please pass the `{env_var}` environment variable"
)
logger.info(f"Init {env_key} class for: {type_value}")
return self.indexes[env_key]["class"].get(type_value)()
env_class = EnvironmentManager()
# create instances of functions
logger.info(f"Initializing required classes for components")
llm_class: BaseChatClient = env_class.create_instance("llm")
translate_class: BaseTranslationClass = env_class.create_instance("translate")
storage_class: BaseStorageClass = env_class.create_instance("storage")
vectorstore_class: BaseVectorStore = env_class.create_instance("vectorstore")