__init__.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200
  1. from typing import Any, Optional
  2. from pydantic import Field, NonNegativeInt, PositiveInt, computed_field
  3. from pydantic_settings import BaseSettings
  4. from configs.middleware.cache.redis_config import RedisConfig
  5. from configs.middleware.storage.aliyun_oss_storage_config import AliyunOSSStorageConfig
  6. from configs.middleware.storage.amazon_s3_storage_config import S3StorageConfig
  7. from configs.middleware.storage.azure_blob_storage_config import AzureBlobStorageConfig
  8. from configs.middleware.storage.google_cloud_storage_config import GoogleCloudStorageConfig
  9. from configs.middleware.storage.oci_storage_config import OCIStorageConfig
  10. from configs.middleware.storage.tencent_cos_storage_config import TencentCloudCOSStorageConfig
  11. from configs.middleware.vdb.analyticdb_config import AnalyticdbConfig
  12. from configs.middleware.vdb.chroma_config import ChromaConfig
  13. from configs.middleware.vdb.milvus_config import MilvusConfig
  14. from configs.middleware.vdb.opensearch_config import OpenSearchConfig
  15. from configs.middleware.vdb.oracle_config import OracleConfig
  16. from configs.middleware.vdb.pgvector_config import PGVectorConfig
  17. from configs.middleware.vdb.pgvectors_config import PGVectoRSConfig
  18. from configs.middleware.vdb.qdrant_config import QdrantConfig
  19. from configs.middleware.vdb.relyt_config import RelytConfig
  20. from configs.middleware.vdb.tencent_vector_config import TencentVectorDBConfig
  21. from configs.middleware.vdb.tidb_vector_config import TiDBVectorConfig
  22. from configs.middleware.vdb.weaviate_config import WeaviateConfig
  23. class StorageConfig(BaseSettings):
  24. STORAGE_TYPE: str = Field(
  25. description='storage type,'
  26. ' default to `local`,'
  27. ' available values are `local`, `s3`, `azure-blob`, `aliyun-oss`, `google-storage`.',
  28. default='local',
  29. )
  30. STORAGE_LOCAL_PATH: str = Field(
  31. description='local storage path',
  32. default='storage',
  33. )
  34. class VectorStoreConfig(BaseSettings):
  35. VECTOR_STORE: Optional[str] = Field(
  36. description='vector store type',
  37. default=None,
  38. )
  39. class KeywordStoreConfig(BaseSettings):
  40. KEYWORD_STORE: str = Field(
  41. description='keyword store type',
  42. default='jieba',
  43. )
  44. class DatabaseConfig:
  45. DB_HOST: str = Field(
  46. description='db host',
  47. default='localhost',
  48. )
  49. DB_PORT: PositiveInt = Field(
  50. description='db port',
  51. default=5432,
  52. )
  53. DB_USERNAME: str = Field(
  54. description='db username',
  55. default='postgres',
  56. )
  57. DB_PASSWORD: str = Field(
  58. description='db password',
  59. default='',
  60. )
  61. DB_DATABASE: str = Field(
  62. description='db database',
  63. default='dify',
  64. )
  65. DB_CHARSET: str = Field(
  66. description='db charset',
  67. default='',
  68. )
  69. DB_EXTRAS: str = Field(
  70. description='db extras options. Example: keepalives_idle=60&keepalives=1',
  71. default='',
  72. )
  73. SQLALCHEMY_DATABASE_URI_SCHEME: str = Field(
  74. description='db uri scheme',
  75. default='postgresql',
  76. )
  77. @computed_field
  78. @property
  79. def SQLALCHEMY_DATABASE_URI(self) -> str:
  80. db_extras = (
  81. f"{self.DB_EXTRAS}&client_encoding={self.DB_CHARSET}"
  82. if self.DB_CHARSET
  83. else self.DB_EXTRAS
  84. ).strip("&")
  85. db_extras = f"?{db_extras}" if db_extras else ""
  86. return (f"{self.SQLALCHEMY_DATABASE_URI_SCHEME}://"
  87. f"{self.DB_USERNAME}:{self.DB_PASSWORD}@{self.DB_HOST}:{self.DB_PORT}/{self.DB_DATABASE}"
  88. f"{db_extras}")
  89. SQLALCHEMY_POOL_SIZE: NonNegativeInt = Field(
  90. description='pool size of SqlAlchemy',
  91. default=30,
  92. )
  93. SQLALCHEMY_MAX_OVERFLOW: NonNegativeInt = Field(
  94. description='max overflows for SqlAlchemy',
  95. default=10,
  96. )
  97. SQLALCHEMY_POOL_RECYCLE: NonNegativeInt = Field(
  98. description='SqlAlchemy pool recycle',
  99. default=3600,
  100. )
  101. SQLALCHEMY_POOL_PRE_PING: bool = Field(
  102. description='whether to enable pool pre-ping in SqlAlchemy',
  103. default=False,
  104. )
  105. SQLALCHEMY_ECHO: bool | str = Field(
  106. description='whether to enable SqlAlchemy echo',
  107. default=False,
  108. )
  109. @computed_field
  110. @property
  111. def SQLALCHEMY_ENGINE_OPTIONS(self) -> dict[str, Any]:
  112. return {
  113. 'pool_size': self.SQLALCHEMY_POOL_SIZE,
  114. 'max_overflow': self.SQLALCHEMY_MAX_OVERFLOW,
  115. 'pool_recycle': self.SQLALCHEMY_POOL_RECYCLE,
  116. 'pool_pre_ping': self.SQLALCHEMY_POOL_PRE_PING,
  117. 'connect_args': {'options': '-c timezone=UTC'},
  118. }
  119. class CeleryConfig(DatabaseConfig):
  120. CELERY_BACKEND: str = Field(
  121. description='Celery backend, available values are `database`, `redis`',
  122. default='database',
  123. )
  124. CELERY_BROKER_URL: Optional[str] = Field(
  125. description='CELERY_BROKER_URL',
  126. default=None,
  127. )
  128. @computed_field
  129. @property
  130. def CELERY_RESULT_BACKEND(self) -> str | None:
  131. return 'db+{}'.format(self.SQLALCHEMY_DATABASE_URI) \
  132. if self.CELERY_BACKEND == 'database' else self.CELERY_BROKER_URL
  133. @computed_field
  134. @property
  135. def BROKER_USE_SSL(self) -> bool:
  136. return self.CELERY_BROKER_URL.startswith('rediss://') if self.CELERY_BROKER_URL else False
  137. class MiddlewareConfig(
  138. # place the configs in alphabet order
  139. CeleryConfig,
  140. DatabaseConfig,
  141. KeywordStoreConfig,
  142. RedisConfig,
  143. # configs of storage and storage providers
  144. StorageConfig,
  145. AliyunOSSStorageConfig,
  146. AzureBlobStorageConfig,
  147. GoogleCloudStorageConfig,
  148. TencentCloudCOSStorageConfig,
  149. S3StorageConfig,
  150. OCIStorageConfig,
  151. # configs of vdb and vdb providers
  152. VectorStoreConfig,
  153. AnalyticdbConfig,
  154. ChromaConfig,
  155. MilvusConfig,
  156. OpenSearchConfig,
  157. OracleConfig,
  158. PGVectorConfig,
  159. PGVectoRSConfig,
  160. QdrantConfig,
  161. RelytConfig,
  162. TencentVectorDBConfig,
  163. TiDBVectorConfig,
  164. WeaviateConfig,
  165. ):
  166. pass