config.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301
  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 IngestApiConfig:
  102. """External knowledge ingest API configuration.
  103. Empty url means real ingest is disabled; dry-run paths do not need it.
  104. """
  105. url: str = ""
  106. token: str = ""
  107. timeout_seconds: float = 30.0
  108. retry_delays_seconds: tuple[float, ...] = (1.0, 3.0)
  109. @classmethod
  110. def from_env(cls, env_file: str | Path = ".env") -> "IngestApiConfig":
  111. file_env = load_env_file(env_file)
  112. return cls(
  113. url=env_value("CK_INGEST_API_URL", file_env),
  114. token=env_value("CK_INGEST_TOKEN", file_env),
  115. timeout_seconds=float(env_value("CK_INGEST_TIMEOUT_SECONDS", file_env, "30")),
  116. retry_delays_seconds=env_float_tuple(
  117. "CK_INGEST_RETRY_DELAYS_SECONDS",
  118. file_env,
  119. "1,3",
  120. ),
  121. )
  122. @dataclass
  123. class Settings:
  124. """全流程配置聚合。M1 只用到 pg;其余字段供 M2+ 使用。"""
  125. pg: PgConfig
  126. # 爬虫(M2):小红书/微信等走 aiddit;抖音单独走 piaoquantv。
  127. aiddit_crawler_base_url: str
  128. crawler_timeout: int
  129. # OpenRouter / Gemini(遗留视频链路和兼容烟测)
  130. openrouter_model: str
  131. openrouter_timeout_seconds: float
  132. openrouter_base_url: str
  133. openrouter_api_key: str
  134. # 文本步 LLM(闸门/拆分/解构)
  135. llm_model: str
  136. # 卡片 / 抽帧
  137. max_cards: int
  138. frames_dir: str
  139. douyin_ratio: str # 视频下载偏好码率(控成本/体积),如 540p
  140. data_dir: str # 媒体本地落盘根目录;空字符串=关闭媒体保存(解耦开关)
  141. # query 搜索 + OSS 转存(默认值必填:tests/test_video_extract.py 用关键字构造 Settings 且不传这些)
  142. oss_upload_url: str = "http://crawler-upload-v2.aiddit.com/crawler/oss/upload_stream"
  143. oss_upload_timeout_seconds: int = 60
  144. video_oss_retry_delays_seconds: tuple[float, ...] = (5.0, 10.0)
  145. search_default_limit: int = 5
  146. search_content_type: str = "图文"
  147. search_sort_type: str = "综合"
  148. # 抖音走独立后端(piaoquantv,绕开 aiddit 限流):host / account_id / cookie_batch 单独配
  149. piaoquantv_douyin_base_url: str = "http://crawapi.piaoquantv.com"
  150. piaoquantv_douyin_account_id: str = "7450041106378522636"
  151. piaoquantv_douyin_cookie_batch: str = "default"
  152. # Aliyun Bailian / Qwen(旧解构图文读懂 + 文本闸门/拆解)。
  153. # 带默认值,保护旧测试里直接 Settings(...) 的关键字构造。
  154. bailian_api_key: str = ""
  155. bailian_base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1"
  156. bailian_vl_model: str = "qwen-vl-plus"
  157. bailian_text_model: str = "qwen-plus"
  158. bailian_timeout_seconds: float = 120.0
  159. demo_data_dir: str = ""
  160. @classmethod
  161. def from_env(cls, env_file: str | Path = ".env") -> "Settings":
  162. file_env = load_env_file(env_file)
  163. return cls(
  164. pg=PgConfig.from_env(env_file),
  165. aiddit_crawler_base_url=env_value(
  166. "AIDDIT_CRAWLER_BASE_URL", file_env,
  167. "http://crawler.aiddit.com",
  168. ),
  169. crawler_timeout=int(env_value(
  170. "AIDDIT_CRAWLER_TIMEOUT_SECONDS", file_env,
  171. "30",
  172. )),
  173. openrouter_model=env_value(
  174. "OPENROUTER_MODEL", file_env,
  175. "google/gemini-3-flash-preview",
  176. ),
  177. openrouter_timeout_seconds=float(env_value(
  178. "OPENROUTER_TIMEOUT_SECONDS", file_env,
  179. "90",
  180. )),
  181. openrouter_base_url=env_value(
  182. "OPENROUTER_BASE_URL", file_env, "https://openrouter.ai/api/v1"
  183. ),
  184. openrouter_api_key=env_value("OPENROUTER_API_KEY", file_env),
  185. bailian_api_key=env_value("ALIYUN_BAILIAN_API_KEY", file_env),
  186. bailian_base_url=env_value(
  187. "ALIYUN_BAILIAN_BASE_URL", file_env,
  188. "https://dashscope.aliyuncs.com/compatible-mode/v1",
  189. ),
  190. bailian_vl_model=env_value(
  191. "ALIYUN_BAILIAN_VL_MODEL", file_env,
  192. env_value("ALIYUN_BAILIAN_MODEL", file_env, "qwen-vl-plus"),
  193. ),
  194. bailian_text_model=env_value(
  195. "ALIYUN_BAILIAN_TEXT_MODEL", file_env,
  196. env_value("ALIYUN_BAILIAN_MODEL", file_env, "qwen-plus"),
  197. ),
  198. bailian_timeout_seconds=float(env_value(
  199. "ALIYUN_BAILIAN_TIMEOUT_SECONDS", file_env,
  200. "120",
  201. )),
  202. # 文本步(筛选/拆分/解构)默认走百炼 Qwen;可用 CK_STAGE_MODEL 覆盖
  203. llm_model=env_value(
  204. "CK_STAGE_MODEL", file_env,
  205. env_value("ALIYUN_BAILIAN_TEXT_MODEL", file_env,
  206. env_value("ALIYUN_BAILIAN_MODEL", file_env, "qwen-plus")),
  207. ),
  208. max_cards=int(env_value("CK_MAX_CARDS", file_env, "12")),
  209. frames_dir=env_value("CK_FRAMES_DIR", file_env, "runtime/frames"),
  210. douyin_ratio=env_value("CK_DOUYIN_RATIO", file_env, "540p"),
  211. data_dir=env_value("CK_DATA_DIR", file_env, "data"),
  212. demo_data_dir=env_value(
  213. "CK_DEMO_DATA_DIR", file_env,
  214. str(Path(env_value("CK_DATA_DIR", file_env, "data")) / "demo"),
  215. ),
  216. oss_upload_url=env_value(
  217. "CRAWLER_OSS_UPLOAD_URL", file_env,
  218. "http://crawler-upload-v2.aiddit.com/crawler/oss/upload_stream",
  219. ),
  220. oss_upload_timeout_seconds=int(
  221. env_value(
  222. "CRAWLER_OSS_UPLOAD_TIMEOUT_SECONDS", file_env,
  223. "60",
  224. )
  225. ),
  226. video_oss_retry_delays_seconds=env_float_tuple(
  227. "CK_VIDEO_OSS_RETRY_DELAYS_SECONDS",
  228. file_env,
  229. "5,10",
  230. ),
  231. search_default_limit=int(env_value(
  232. "CREATION_SEARCH_DEFAULT_LIMIT", file_env,
  233. "5",
  234. )),
  235. search_content_type=env_value(
  236. "CREATION_SEARCH_CONTENT_TYPE", file_env,
  237. "图文",
  238. ),
  239. search_sort_type=env_value(
  240. "CREATION_SEARCH_SORT_TYPE", file_env,
  241. "综合",
  242. ),
  243. piaoquantv_douyin_base_url=env_value(
  244. "PIAOQUANTV_DOUYIN_BASE_URL", file_env,
  245. "http://crawapi.piaoquantv.com",
  246. ),
  247. piaoquantv_douyin_account_id=env_value(
  248. "PIAOQUANTV_DOUYIN_ACCOUNT_ID", file_env,
  249. "7450041106378522636",
  250. ),
  251. piaoquantv_douyin_cookie_batch=env_value(
  252. "PIAOQUANTV_DOUYIN_COOKIE_BATCH", file_env,
  253. "default",
  254. ),
  255. )
  256. @property
  257. def crawler_base_url(self) -> str:
  258. return self.aiddit_crawler_base_url
  259. @property
  260. def douyin_base_url(self) -> str:
  261. return self.piaoquantv_douyin_base_url
  262. @property
  263. def douyin_account_id(self) -> str:
  264. return self.piaoquantv_douyin_account_id
  265. @property
  266. def douyin_cookie_batch(self) -> str:
  267. return self.piaoquantv_douyin_cookie_batch
  268. @property
  269. def video_model(self) -> str:
  270. return self.openrouter_model