fix(kb-deep-analyzer): 子线程环境变量传递
- 在 _analyze_chunks_concurrent 中使用 _worker_env 保存环境变量 - 添加 init_worker() 初始化函数 - ThreadPoolExecutor 使用 initializer=init_worker 传递环境变量
This commit is contained in:
@@ -702,6 +702,18 @@ class DeepAnalyzer:
|
|||||||
"""
|
"""
|
||||||
import concurrent.futures
|
import concurrent.futures
|
||||||
|
|
||||||
|
# 确保环境变量在子线程中可用
|
||||||
|
_worker_env = {
|
||||||
|
'AI_API_KEY': os.environ.get('AI_API_KEY', ''),
|
||||||
|
'AI_PROVIDER': os.environ.get('AI_PROVIDER', ''),
|
||||||
|
'AI_MODEL': os.environ.get('AI_MODEL', ''),
|
||||||
|
'AI_BASE_URL': os.environ.get('AI_BASE_URL', ''),
|
||||||
|
}
|
||||||
|
|
||||||
|
def init_worker():
|
||||||
|
"""初始化 worker 线程的环境变量"""
|
||||||
|
os.environ.update(_worker_env)
|
||||||
|
|
||||||
result = AnalysisResult()
|
result = AnalysisResult()
|
||||||
CHUNK_MAX_CHARS = 3000
|
CHUNK_MAX_CHARS = 3000
|
||||||
total_chars = len(text)
|
total_chars = len(text)
|
||||||
@@ -778,7 +790,7 @@ class DeepAnalyzer:
|
|||||||
batch_tasks = [(i, batch_chunks[j], batch_times[j][0], batch_times[j][1])
|
batch_tasks = [(i, batch_chunks[j], batch_times[j][0], batch_times[j][1])
|
||||||
for j, i in enumerate(range(batch_idx, batch_end))]
|
for j, i in enumerate(range(batch_idx, batch_end))]
|
||||||
|
|
||||||
with concurrent.futures.ThreadPoolExecutor(max_workers=batch_size) as executor:
|
with concurrent.futures.ThreadPoolExecutor(max_workers=batch_size, initializer=init_worker) as executor:
|
||||||
futures = {executor.submit(self._analyze_chunk, task[1], task[0] + 1, metadata): task[0]
|
futures = {executor.submit(self._analyze_chunk, task[1], task[0] + 1, metadata): task[0]
|
||||||
for task in batch_tasks}
|
for task in batch_tasks}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user