config.py 9.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249
  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. @dataclass
  35. class PgConfig:
  36. """过程库(Greenplum open_aigc,schema creation_knowledge)连接配置。"""
  37. host: str
  38. port: int
  39. user: str
  40. password: str
  41. database: str
  42. schema: str = DEFAULT_PG_SCHEMA
  43. timeout: int = 10
  44. @classmethod
  45. def from_env(cls, env_file: str | Path = ".env") -> "PgConfig":
  46. file_env = load_env_file(env_file)
  47. return cls(
  48. host=env_value("OPEN_AIGC_PG_HOST", file_env, required=True),
  49. port=int(env_value("OPEN_AIGC_PG_PORT", file_env, "5432")),
  50. user=env_value("OPEN_AIGC_PG_USER", file_env, required=True),
  51. password=env_value("OPEN_AIGC_PG_PASSWORD", file_env, required=True),
  52. database=env_value(
  53. "OPEN_AIGC_PG_DB_NAME", file_env, "open_aigc"
  54. ),
  55. schema=env_value("CK_PG_SCHEMA", file_env, DEFAULT_PG_SCHEMA),
  56. )
  57. @dataclass
  58. class CreationDbConfig:
  59. """Formal creation-knowledge state database, backed by cloud PostgreSQL."""
  60. host: str
  61. port: int
  62. user: str
  63. password: str
  64. database: str
  65. schema: str = DEFAULT_PG_SCHEMA
  66. timeout: int = 10
  67. application_name: str = "creation-knowledge"
  68. @classmethod
  69. def from_env(cls, env_file: str | Path = ".env") -> "CreationDbConfig":
  70. file_env = load_env_file(env_file)
  71. return cls(
  72. host=env_value("CK_DB_HOST", file_env, required=True),
  73. port=int(env_value("CK_DB_PORT", file_env, "5432")),
  74. user=env_value("CK_DB_USER", file_env, required=True),
  75. password=env_value("CK_DB_PASSWORD", file_env, required=True),
  76. database=env_value(
  77. "CK_DB_NAME",
  78. file_env,
  79. "creation_knowledge_prod",
  80. ),
  81. schema=env_value("CK_DB_SCHEMA", file_env, DEFAULT_PG_SCHEMA),
  82. )
  83. @dataclass
  84. class Settings:
  85. """全流程配置聚合。M1 只用到 pg;其余字段供 M2+ 使用。"""
  86. pg: PgConfig
  87. # 爬虫(M2):小红书/微信等走 aiddit;抖音单独走 piaoquantv。
  88. aiddit_crawler_base_url: str
  89. crawler_timeout: int
  90. # OpenRouter / Gemini(遗留视频链路和兼容烟测)
  91. openrouter_model: str
  92. openrouter_timeout_seconds: float
  93. openrouter_base_url: str
  94. openrouter_api_key: str
  95. # 文本步 LLM(闸门/拆分/解构)
  96. llm_model: str
  97. # 卡片 / 抽帧
  98. max_cards: int
  99. frames_dir: str
  100. douyin_ratio: str # 视频下载偏好码率(控成本/体积),如 540p
  101. data_dir: str # 媒体本地落盘根目录;空字符串=关闭媒体保存(解耦开关)
  102. # query 搜索 + OSS 转存(默认值必填:tests/test_video_extract.py 用关键字构造 Settings 且不传这些)
  103. oss_upload_url: str = "http://crawler-upload-v2.aiddit.com/crawler/oss/upload_stream"
  104. oss_upload_timeout_seconds: int = 60
  105. search_default_limit: int = 5
  106. search_content_type: str = "图文"
  107. search_sort_type: str = "综合"
  108. # 抖音走独立后端(piaoquantv,绕开 aiddit 限流):host / account_id / cookie_batch 单独配
  109. piaoquantv_douyin_base_url: str = "http://crawapi.piaoquantv.com"
  110. piaoquantv_douyin_account_id: str = "7450041106378522636"
  111. piaoquantv_douyin_cookie_batch: str = "default"
  112. # Aliyun Bailian / Qwen(旧解构图文读懂 + 文本闸门/拆解)。
  113. # 带默认值,保护旧测试里直接 Settings(...) 的关键字构造。
  114. bailian_api_key: str = ""
  115. bailian_base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1"
  116. bailian_vl_model: str = "qwen-vl-plus"
  117. bailian_text_model: str = "qwen-plus"
  118. bailian_timeout_seconds: float = 120.0
  119. demo_data_dir: str = ""
  120. @classmethod
  121. def from_env(cls, env_file: str | Path = ".env") -> "Settings":
  122. file_env = load_env_file(env_file)
  123. return cls(
  124. pg=PgConfig.from_env(env_file),
  125. aiddit_crawler_base_url=env_value(
  126. "AIDDIT_CRAWLER_BASE_URL", file_env,
  127. "http://crawler.aiddit.com",
  128. ),
  129. crawler_timeout=int(env_value(
  130. "AIDDIT_CRAWLER_TIMEOUT_SECONDS", file_env,
  131. "30",
  132. )),
  133. openrouter_model=env_value(
  134. "OPENROUTER_MODEL", file_env,
  135. "google/gemini-3-flash-preview",
  136. ),
  137. openrouter_timeout_seconds=float(env_value(
  138. "OPENROUTER_TIMEOUT_SECONDS", file_env,
  139. "90",
  140. )),
  141. openrouter_base_url=env_value(
  142. "OPENROUTER_BASE_URL", file_env, "https://openrouter.ai/api/v1"
  143. ),
  144. openrouter_api_key=env_value("OPENROUTER_API_KEY", file_env),
  145. bailian_api_key=env_value("ALIYUN_BAILIAN_API_KEY", file_env),
  146. bailian_base_url=env_value(
  147. "ALIYUN_BAILIAN_BASE_URL", file_env,
  148. "https://dashscope.aliyuncs.com/compatible-mode/v1",
  149. ),
  150. bailian_vl_model=env_value(
  151. "ALIYUN_BAILIAN_VL_MODEL", file_env,
  152. env_value("ALIYUN_BAILIAN_MODEL", file_env, "qwen-vl-plus"),
  153. ),
  154. bailian_text_model=env_value(
  155. "ALIYUN_BAILIAN_TEXT_MODEL", file_env,
  156. env_value("ALIYUN_BAILIAN_MODEL", file_env, "qwen-plus"),
  157. ),
  158. bailian_timeout_seconds=float(env_value(
  159. "ALIYUN_BAILIAN_TIMEOUT_SECONDS", file_env,
  160. "120",
  161. )),
  162. # 文本步(筛选/拆分/解构)默认走百炼 Qwen;可用 CK_STAGE_MODEL 覆盖
  163. llm_model=env_value(
  164. "CK_STAGE_MODEL", file_env,
  165. env_value("ALIYUN_BAILIAN_TEXT_MODEL", file_env,
  166. env_value("ALIYUN_BAILIAN_MODEL", file_env, "qwen-plus")),
  167. ),
  168. max_cards=int(env_value("CK_MAX_CARDS", file_env, "12")),
  169. frames_dir=env_value("CK_FRAMES_DIR", file_env, "runtime/frames"),
  170. douyin_ratio=env_value("CK_DOUYIN_RATIO", file_env, "540p"),
  171. data_dir=env_value("CK_DATA_DIR", file_env, "data"),
  172. demo_data_dir=env_value(
  173. "CK_DEMO_DATA_DIR", file_env,
  174. str(Path(env_value("CK_DATA_DIR", file_env, "data")) / "demo"),
  175. ),
  176. oss_upload_url=env_value(
  177. "CRAWLER_OSS_UPLOAD_URL", file_env,
  178. "http://crawler-upload-v2.aiddit.com/crawler/oss/upload_stream",
  179. ),
  180. oss_upload_timeout_seconds=int(
  181. env_value(
  182. "CRAWLER_OSS_UPLOAD_TIMEOUT_SECONDS", file_env,
  183. "60",
  184. )
  185. ),
  186. search_default_limit=int(env_value(
  187. "CREATION_SEARCH_DEFAULT_LIMIT", file_env,
  188. "5",
  189. )),
  190. search_content_type=env_value(
  191. "CREATION_SEARCH_CONTENT_TYPE", file_env,
  192. "图文",
  193. ),
  194. search_sort_type=env_value(
  195. "CREATION_SEARCH_SORT_TYPE", file_env,
  196. "综合",
  197. ),
  198. piaoquantv_douyin_base_url=env_value(
  199. "PIAOQUANTV_DOUYIN_BASE_URL", file_env,
  200. "http://crawapi.piaoquantv.com",
  201. ),
  202. piaoquantv_douyin_account_id=env_value(
  203. "PIAOQUANTV_DOUYIN_ACCOUNT_ID", file_env,
  204. "7450041106378522636",
  205. ),
  206. piaoquantv_douyin_cookie_batch=env_value(
  207. "PIAOQUANTV_DOUYIN_COOKIE_BATCH", file_env,
  208. "default",
  209. ),
  210. )
  211. @property
  212. def crawler_base_url(self) -> str:
  213. return self.aiddit_crawler_base_url
  214. @property
  215. def douyin_base_url(self) -> str:
  216. return self.piaoquantv_douyin_base_url
  217. @property
  218. def douyin_account_id(self) -> str:
  219. return self.piaoquantv_douyin_account_id
  220. @property
  221. def douyin_cookie_batch(self) -> str:
  222. return self.piaoquantv_douyin_cookie_batch
  223. @property
  224. def video_model(self) -> str:
  225. return self.openrouter_model