config.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274
  1. """配置加载:从 .env 文件或环境变量读取,风格对齐 ContentFindAgentNew。
  2. 只读环境变量,不硬编码密钥。PG 是 M1 唯一必填项;其余环节配置随里程碑推进逐步使用。
  3. """
  4. from __future__ import annotations
  5. import os
  6. from dataclasses import dataclass
  7. from pathlib import Path
  8. # 我们自己的过程库 schema(与 OPEN_AIGC_PG_SCHEMA=public 区分开)
  9. DEFAULT_PG_SCHEMA = "creation_knowledge"
  10. def load_env_file(env_path: str | Path = ".env") -> dict[str, str]:
  11. """解析 .env 为 dict;文件不存在则返回空。与 crawapi_http._load_env_file 同款。"""
  12. path = Path(env_path)
  13. if not path.exists():
  14. return {}
  15. env: dict[str, str] = {}
  16. for line in path.read_text(encoding="utf-8").splitlines():
  17. stripped = line.strip()
  18. if not stripped or stripped.startswith("#") or "=" not in stripped:
  19. continue
  20. key, value = stripped.split("=", 1)
  21. env[key.strip()] = value.strip().strip('"').strip("'")
  22. return env
  23. def env_value(
  24. key: str,
  25. file_env: dict[str, str],
  26. default: str | None = None,
  27. required: bool = False,
  28. ) -> str:
  29. """取值优先级:os.environ > .env 文件 > default。required 缺失时报错。"""
  30. value = os.getenv(key) or file_env.get(key) or default
  31. if required and not value:
  32. raise RuntimeError(f"missing required env: {key}")
  33. return value or ""
  34. def env_float_tuple(
  35. key: str,
  36. file_env: dict[str, str],
  37. default: str,
  38. ) -> tuple[float, ...]:
  39. raw = env_value(key, file_env, default)
  40. values: list[float] = []
  41. for part in raw.split(","):
  42. text = part.strip()
  43. if not text:
  44. continue
  45. values.append(float(text))
  46. return tuple(values)
  47. @dataclass
  48. class PgConfig:
  49. """过程库(Greenplum open_aigc,schema creation_knowledge)连接配置。"""
  50. host: str
  51. port: int
  52. user: str
  53. password: str
  54. database: str
  55. schema: str = DEFAULT_PG_SCHEMA
  56. timeout: int = 10
  57. @classmethod
  58. def from_env(cls, env_file: str | Path = ".env") -> "PgConfig":
  59. file_env = load_env_file(env_file)
  60. return cls(
  61. host=env_value("OPEN_AIGC_PG_HOST", file_env, required=True),
  62. port=int(env_value("OPEN_AIGC_PG_PORT", file_env, "5432")),
  63. user=env_value("OPEN_AIGC_PG_USER", file_env, required=True),
  64. password=env_value("OPEN_AIGC_PG_PASSWORD", file_env, required=True),
  65. database=env_value(
  66. "OPEN_AIGC_PG_DB_NAME", file_env, "open_aigc"
  67. ),
  68. schema=env_value("CK_PG_SCHEMA", file_env, DEFAULT_PG_SCHEMA),
  69. )
  70. @dataclass
  71. class CreationDbConfig:
  72. """Formal creation-knowledge state database, backed by cloud PostgreSQL."""
  73. host: str
  74. port: int
  75. user: str
  76. password: str
  77. database: str
  78. schema: str = DEFAULT_PG_SCHEMA
  79. timeout: int = 10
  80. application_name: str = "creation-knowledge"
  81. pool_min: int = 1
  82. pool_max: int = 10
  83. @classmethod
  84. def from_env(cls, env_file: str | Path = ".env") -> "CreationDbConfig":
  85. file_env = load_env_file(env_file)
  86. return cls(
  87. host=env_value("CK_DB_HOST", file_env, required=True),
  88. port=int(env_value("CK_DB_PORT", file_env, "5432")),
  89. user=env_value("CK_DB_USER", file_env, required=True),
  90. password=env_value("CK_DB_PASSWORD", file_env, required=True),
  91. database=env_value(
  92. "CK_DB_NAME",
  93. file_env,
  94. "creation_knowledge_prod",
  95. ),
  96. schema=env_value("CK_DB_SCHEMA", file_env, DEFAULT_PG_SCHEMA),
  97. pool_min=int(env_value("CK_DB_POOL_MIN", file_env, "1")),
  98. pool_max=int(env_value("CK_DB_POOL_MAX", file_env, "10")),
  99. )
  100. @dataclass
  101. class Settings:
  102. """全流程配置聚合。M1 只用到 pg;其余字段供 M2+ 使用。"""
  103. pg: PgConfig
  104. # 爬虫(M2):小红书/微信等走 aiddit;抖音单独走 piaoquantv。
  105. aiddit_crawler_base_url: str
  106. crawler_timeout: int
  107. # OpenRouter / Gemini(遗留视频链路和兼容烟测)
  108. openrouter_model: str
  109. openrouter_timeout_seconds: float
  110. openrouter_base_url: str
  111. openrouter_api_key: str
  112. # 文本步 LLM(闸门/拆分/解构)
  113. llm_model: str
  114. # 卡片 / 抽帧
  115. max_cards: int
  116. frames_dir: str
  117. douyin_ratio: str # 视频下载偏好码率(控成本/体积),如 540p
  118. data_dir: str # 媒体本地落盘根目录;空字符串=关闭媒体保存(解耦开关)
  119. # query 搜索 + OSS 转存(默认值必填:tests/test_video_extract.py 用关键字构造 Settings 且不传这些)
  120. oss_upload_url: str = "http://crawler-upload-v2.aiddit.com/crawler/oss/upload_stream"
  121. oss_upload_timeout_seconds: int = 60
  122. video_oss_retry_delays_seconds: tuple[float, ...] = (5.0, 10.0)
  123. search_default_limit: int = 5
  124. search_content_type: str = "图文"
  125. search_sort_type: str = "综合"
  126. # 抖音走独立后端(piaoquantv,绕开 aiddit 限流):host / account_id / cookie_batch 单独配
  127. piaoquantv_douyin_base_url: str = "http://crawapi.piaoquantv.com"
  128. piaoquantv_douyin_account_id: str = "7450041106378522636"
  129. piaoquantv_douyin_cookie_batch: str = "default"
  130. # Aliyun Bailian / Qwen(旧解构图文读懂 + 文本闸门/拆解)。
  131. # 带默认值,保护旧测试里直接 Settings(...) 的关键字构造。
  132. bailian_api_key: str = ""
  133. bailian_base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1"
  134. bailian_vl_model: str = "qwen-vl-plus"
  135. bailian_text_model: str = "qwen-plus"
  136. bailian_timeout_seconds: float = 120.0
  137. demo_data_dir: str = ""
  138. @classmethod
  139. def from_env(cls, env_file: str | Path = ".env") -> "Settings":
  140. file_env = load_env_file(env_file)
  141. return cls(
  142. pg=PgConfig.from_env(env_file),
  143. aiddit_crawler_base_url=env_value(
  144. "AIDDIT_CRAWLER_BASE_URL", file_env,
  145. "http://crawler.aiddit.com",
  146. ),
  147. crawler_timeout=int(env_value(
  148. "AIDDIT_CRAWLER_TIMEOUT_SECONDS", file_env,
  149. "30",
  150. )),
  151. openrouter_model=env_value(
  152. "OPENROUTER_MODEL", file_env,
  153. "google/gemini-3-flash-preview",
  154. ),
  155. openrouter_timeout_seconds=float(env_value(
  156. "OPENROUTER_TIMEOUT_SECONDS", file_env,
  157. "90",
  158. )),
  159. openrouter_base_url=env_value(
  160. "OPENROUTER_BASE_URL", file_env, "https://openrouter.ai/api/v1"
  161. ),
  162. openrouter_api_key=env_value("OPENROUTER_API_KEY", file_env),
  163. bailian_api_key=env_value("ALIYUN_BAILIAN_API_KEY", file_env),
  164. bailian_base_url=env_value(
  165. "ALIYUN_BAILIAN_BASE_URL", file_env,
  166. "https://dashscope.aliyuncs.com/compatible-mode/v1",
  167. ),
  168. bailian_vl_model=env_value(
  169. "ALIYUN_BAILIAN_VL_MODEL", file_env,
  170. env_value("ALIYUN_BAILIAN_MODEL", file_env, "qwen-vl-plus"),
  171. ),
  172. bailian_text_model=env_value(
  173. "ALIYUN_BAILIAN_TEXT_MODEL", file_env,
  174. env_value("ALIYUN_BAILIAN_MODEL", file_env, "qwen-plus"),
  175. ),
  176. bailian_timeout_seconds=float(env_value(
  177. "ALIYUN_BAILIAN_TIMEOUT_SECONDS", file_env,
  178. "120",
  179. )),
  180. # 文本步(筛选/拆分/解构)默认走百炼 Qwen;可用 CK_STAGE_MODEL 覆盖
  181. llm_model=env_value(
  182. "CK_STAGE_MODEL", file_env,
  183. env_value("ALIYUN_BAILIAN_TEXT_MODEL", file_env,
  184. env_value("ALIYUN_BAILIAN_MODEL", file_env, "qwen-plus")),
  185. ),
  186. max_cards=int(env_value("CK_MAX_CARDS", file_env, "12")),
  187. frames_dir=env_value("CK_FRAMES_DIR", file_env, "runtime/frames"),
  188. douyin_ratio=env_value("CK_DOUYIN_RATIO", file_env, "540p"),
  189. data_dir=env_value("CK_DATA_DIR", file_env, "data"),
  190. demo_data_dir=env_value(
  191. "CK_DEMO_DATA_DIR", file_env,
  192. str(Path(env_value("CK_DATA_DIR", file_env, "data")) / "demo"),
  193. ),
  194. oss_upload_url=env_value(
  195. "CRAWLER_OSS_UPLOAD_URL", file_env,
  196. "http://crawler-upload-v2.aiddit.com/crawler/oss/upload_stream",
  197. ),
  198. oss_upload_timeout_seconds=int(
  199. env_value(
  200. "CRAWLER_OSS_UPLOAD_TIMEOUT_SECONDS", file_env,
  201. "60",
  202. )
  203. ),
  204. video_oss_retry_delays_seconds=env_float_tuple(
  205. "CK_VIDEO_OSS_RETRY_DELAYS_SECONDS",
  206. file_env,
  207. "5,10",
  208. ),
  209. search_default_limit=int(env_value(
  210. "CREATION_SEARCH_DEFAULT_LIMIT", file_env,
  211. "5",
  212. )),
  213. search_content_type=env_value(
  214. "CREATION_SEARCH_CONTENT_TYPE", file_env,
  215. "图文",
  216. ),
  217. search_sort_type=env_value(
  218. "CREATION_SEARCH_SORT_TYPE", file_env,
  219. "综合",
  220. ),
  221. piaoquantv_douyin_base_url=env_value(
  222. "PIAOQUANTV_DOUYIN_BASE_URL", file_env,
  223. "http://crawapi.piaoquantv.com",
  224. ),
  225. piaoquantv_douyin_account_id=env_value(
  226. "PIAOQUANTV_DOUYIN_ACCOUNT_ID", file_env,
  227. "7450041106378522636",
  228. ),
  229. piaoquantv_douyin_cookie_batch=env_value(
  230. "PIAOQUANTV_DOUYIN_COOKIE_BATCH", file_env,
  231. "default",
  232. ),
  233. )
  234. @property
  235. def crawler_base_url(self) -> str:
  236. return self.aiddit_crawler_base_url
  237. @property
  238. def douyin_base_url(self) -> str:
  239. return self.piaoquantv_douyin_base_url
  240. @property
  241. def douyin_account_id(self) -> str:
  242. return self.piaoquantv_douyin_account_id
  243. @property
  244. def douyin_cookie_batch(self) -> str:
  245. return self.piaoquantv_douyin_cookie_batch
  246. @property
  247. def video_model(self) -> str:
  248. return self.openrouter_model