aigroup-econ-mcp 0.3.5__tar.gz → 0.3.6__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of aigroup-econ-mcp might be problematic. Click here for more details.
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/PKG-INFO +1 -1
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/pyproject.toml +1 -1
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/__init__.py +1 -1
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/server.py +205 -91
- aigroup_econ_mcp-0.3.6/src/aigroup_econ_mcp/tools/timeout.py +283 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/.gitignore +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/LICENSE +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/README.md +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/cli.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/config.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/__init__.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/base.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/cache.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/machine_learning.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/monitoring.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/optimized_example.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/panel_data.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/regression.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/statistics.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/time_series.py +0 -0
- {aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/validation.py +0 -0
|
@@ -9,6 +9,9 @@ from collections.abc import AsyncIterator
|
|
|
9
9
|
from contextlib import asynccontextmanager
|
|
10
10
|
from dataclasses import dataclass
|
|
11
11
|
import importlib.metadata
|
|
12
|
+
import json
|
|
13
|
+
import os
|
|
14
|
+
from pathlib import Path
|
|
12
15
|
|
|
13
16
|
import pandas as pd
|
|
14
17
|
import numpy as np
|
|
@@ -21,6 +24,55 @@ from mcp.server.fastmcp import FastMCP, Context
|
|
|
21
24
|
from mcp.server.session import ServerSession
|
|
22
25
|
from mcp.types import CallToolResult, TextContent
|
|
23
26
|
|
|
27
|
+
# 导入超时管理器
|
|
28
|
+
from .tools.timeout import TimeoutManager
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
# 性能配置加载函数
|
|
32
|
+
def load_performance_config() -> Dict[str, Any]:
|
|
33
|
+
"""加载性能配置"""
|
|
34
|
+
config_paths = [
|
|
35
|
+
Path("config/performance_config.json"),
|
|
36
|
+
Path("performance_config.json"),
|
|
37
|
+
Path.cwd() / "config" / "performance_config.json"
|
|
38
|
+
]
|
|
39
|
+
|
|
40
|
+
for config_path in config_paths:
|
|
41
|
+
if config_path.exists():
|
|
42
|
+
try:
|
|
43
|
+
with open(config_path, 'r', encoding='utf-8') as f:
|
|
44
|
+
config = json.load(f)
|
|
45
|
+
print(f"✅ 性能配置已加载: {config_path}")
|
|
46
|
+
return config
|
|
47
|
+
except Exception as e:
|
|
48
|
+
print(f"⚠️ 加载性能配置失败 {config_path}: {e}")
|
|
49
|
+
|
|
50
|
+
# 返回默认配置
|
|
51
|
+
default_config = {
|
|
52
|
+
"performance": {
|
|
53
|
+
"cache": {"enabled": True, "ttl": 3600, "max_size": 1000},
|
|
54
|
+
"monitoring": {"enabled": True, "memory_check_interval": 1.0},
|
|
55
|
+
"timeouts": {
|
|
56
|
+
"simple_models": 30,
|
|
57
|
+
"complex_models": 120,
|
|
58
|
+
"machine_learning": 300,
|
|
59
|
+
"time_series": 180
|
|
60
|
+
}
|
|
61
|
+
},
|
|
62
|
+
"models": {
|
|
63
|
+
"descriptive_statistics": {"timeout": 30},
|
|
64
|
+
"ols_regression": {"timeout": 60},
|
|
65
|
+
"time_series_analysis": {"timeout": 120},
|
|
66
|
+
"panel_data_analysis": {"timeout": 180},
|
|
67
|
+
"var_model": {"timeout": 300},
|
|
68
|
+
"garch_model": {"timeout": 240},
|
|
69
|
+
"random_forest": {"timeout": 300},
|
|
70
|
+
"gradient_boosting": {"timeout": 300}
|
|
71
|
+
}
|
|
72
|
+
}
|
|
73
|
+
print("⚠️ 使用默认性能配置")
|
|
74
|
+
return default_config
|
|
75
|
+
|
|
24
76
|
|
|
25
77
|
# 数据模型定义 - 使用Pydantic实现结构化输出
|
|
26
78
|
class DescriptiveStatsResult(BaseModel):
|
|
@@ -70,6 +122,8 @@ class AppContext:
|
|
|
70
122
|
"""应用上下文,包含共享资源"""
|
|
71
123
|
config: Dict[str, Any]
|
|
72
124
|
version: str = importlib.metadata.version("aigroup-econ-mcp")
|
|
125
|
+
performance_config: Dict[str, Any] = None
|
|
126
|
+
timeout_manager: Any = None
|
|
73
127
|
|
|
74
128
|
|
|
75
129
|
# 服务器图标(已移除,因为MCP库不再支持Icon类)
|
|
@@ -86,8 +140,19 @@ async def lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
|
|
|
86
140
|
"data_types": ["cross_section", "time_series", "panel"]
|
|
87
141
|
}
|
|
88
142
|
|
|
143
|
+
# 加载性能配置
|
|
144
|
+
performance_config = load_performance_config()
|
|
145
|
+
|
|
146
|
+
# 初始化超时管理器
|
|
147
|
+
timeout_manager = TimeoutManager(performance_config.get("timeouts", {}))
|
|
148
|
+
|
|
89
149
|
try:
|
|
90
|
-
yield AppContext(
|
|
150
|
+
yield AppContext(
|
|
151
|
+
config=config,
|
|
152
|
+
version=importlib.metadata.version("aigroup-econ-mcp"),
|
|
153
|
+
performance_config=performance_config,
|
|
154
|
+
timeout_manager=timeout_manager
|
|
155
|
+
)
|
|
91
156
|
finally:
|
|
92
157
|
# 清理资源
|
|
93
158
|
pass
|
|
@@ -202,51 +267,44 @@ async def descriptive_statistics(
|
|
|
202
267
|
await ctx.info(f"开始计算描述性统计,处理 {len(data)} 个变量")
|
|
203
268
|
|
|
204
269
|
try:
|
|
205
|
-
#
|
|
206
|
-
|
|
207
|
-
|
|
270
|
+
# 获取超时管理器
|
|
271
|
+
timeout_manager = ctx.session.context.timeout_manager
|
|
272
|
+
model_timeout = ctx.session.context.performance_config.get("models", {}).get("descriptive_statistics", {}).get("timeout", 30)
|
|
208
273
|
|
|
209
|
-
|
|
210
|
-
|
|
211
|
-
|
|
212
|
-
|
|
213
|
-
|
|
214
|
-
|
|
215
|
-
# 基础统计量 - 修复:返回所有变量的综合统计
|
|
216
|
-
result = DescriptiveStatsResult(
|
|
217
|
-
count=len(df),
|
|
218
|
-
mean=df.mean().mean(), # 所有变量的均值
|
|
219
|
-
std=df.std().mean(), # 所有变量的标准差均值
|
|
220
|
-
min=df.min().min(), # 所有变量的最小值
|
|
221
|
-
max=df.max().max(), # 所有变量的最大值
|
|
222
|
-
median=df.median().mean(), # 所有变量的中位数均值
|
|
223
|
-
skewness=df.skew().mean(), # 所有变量的偏度均值
|
|
224
|
-
kurtosis=df.kurtosis().mean() # 所有变量的峰度均值
|
|
274
|
+
# 使用超时控制执行计算
|
|
275
|
+
result = await timeout_manager.execute_with_timeout(
|
|
276
|
+
model_name="descriptive_statistics",
|
|
277
|
+
timeout_seconds=model_timeout,
|
|
278
|
+
func=_compute_descriptive_stats,
|
|
279
|
+
data=data
|
|
225
280
|
)
|
|
226
281
|
|
|
227
|
-
|
|
228
|
-
correlation_matrix = df.corr().round(4)
|
|
229
|
-
|
|
230
|
-
await ctx.info(f"描述性统计计算完成,样本大小: {len(df)}")
|
|
282
|
+
await ctx.info(f"描述性统计计算完成,样本大小: {result['count']}")
|
|
231
283
|
|
|
232
284
|
return CallToolResult(
|
|
233
285
|
content=[
|
|
234
286
|
TextContent(
|
|
235
287
|
type="text",
|
|
236
288
|
text=f"描述性统计结果:\n"
|
|
237
|
-
f"均值: {result
|
|
238
|
-
f"标准差: {result
|
|
239
|
-
f"最小值: {result
|
|
240
|
-
f"最大值: {result
|
|
241
|
-
f"中位数: {result
|
|
242
|
-
f"偏度: {result
|
|
243
|
-
f"峰度: {result
|
|
244
|
-
f"相关系数矩阵:\n{correlation_matrix
|
|
289
|
+
f"均值: {result['mean']:.4f}\n"
|
|
290
|
+
f"标准差: {result['std']:.4f}\n"
|
|
291
|
+
f"最小值: {result['min']:.4f}\n"
|
|
292
|
+
f"最大值: {result['max']:.4f}\n"
|
|
293
|
+
f"中位数: {result['median']:.4f}\n"
|
|
294
|
+
f"偏度: {result['skewness']:.4f}\n"
|
|
295
|
+
f"峰度: {result['kurtosis']:.4f}\n\n"
|
|
296
|
+
f"相关系数矩阵:\n{result['correlation_matrix']}"
|
|
245
297
|
)
|
|
246
298
|
],
|
|
247
|
-
structuredContent=result
|
|
299
|
+
structuredContent=result
|
|
248
300
|
)
|
|
249
301
|
|
|
302
|
+
except TimeoutError:
|
|
303
|
+
await ctx.error("描述性统计计算超时,请减少数据量或调整配置")
|
|
304
|
+
return CallToolResult(
|
|
305
|
+
content=[TextContent(type="text", text="错误: 计算超时,请减少数据量或调整配置")],
|
|
306
|
+
isError=True
|
|
307
|
+
)
|
|
250
308
|
except Exception as e:
|
|
251
309
|
await ctx.error(f"计算描述性统计时出错: {str(e)}")
|
|
252
310
|
return CallToolResult(
|
|
@@ -255,6 +313,37 @@ async def descriptive_statistics(
|
|
|
255
313
|
)
|
|
256
314
|
|
|
257
315
|
|
|
316
|
+
def _compute_descriptive_stats(data: Dict[str, List[float]]) -> Dict[str, Any]:
|
|
317
|
+
"""计算描述性统计(同步函数,用于超时控制)"""
|
|
318
|
+
# 数据验证
|
|
319
|
+
if not data:
|
|
320
|
+
raise ValueError("数据不能为空")
|
|
321
|
+
|
|
322
|
+
df = pd.DataFrame(data)
|
|
323
|
+
|
|
324
|
+
# 检查数据一致性
|
|
325
|
+
if len(df.columns) == 0:
|
|
326
|
+
raise ValueError("至少需要一个变量")
|
|
327
|
+
|
|
328
|
+
# 基础统计量
|
|
329
|
+
result = {
|
|
330
|
+
"count": len(df),
|
|
331
|
+
"mean": float(df.mean().mean()), # 所有变量的均值
|
|
332
|
+
"std": float(df.std().mean()), # 所有变量的标准差均值
|
|
333
|
+
"min": float(df.min().min()), # 所有变量的最小值
|
|
334
|
+
"max": float(df.max().max()), # 所有变量的最大值
|
|
335
|
+
"median": float(df.median().mean()), # 所有变量的中位数均值
|
|
336
|
+
"skewness": float(df.skew().mean()), # 所有变量的偏度均值
|
|
337
|
+
"kurtosis": float(df.kurtosis().mean()) # 所有变量的峰度均值
|
|
338
|
+
}
|
|
339
|
+
|
|
340
|
+
# 计算相关系数矩阵
|
|
341
|
+
correlation_matrix = df.corr().round(4)
|
|
342
|
+
result["correlation_matrix"] = correlation_matrix.to_string()
|
|
343
|
+
|
|
344
|
+
return result
|
|
345
|
+
|
|
346
|
+
|
|
258
347
|
@mcp.tool()
|
|
259
348
|
async def ols_regression(
|
|
260
349
|
ctx: Context[ServerSession, AppContext],
|
|
@@ -342,59 +431,18 @@ async def ols_regression(
|
|
|
342
431
|
await ctx.info(f"开始OLS回归分析,样本大小: {len(y_data)},自变量数量: {len(x_data[0]) if x_data else 0}")
|
|
343
432
|
|
|
344
433
|
try:
|
|
345
|
-
#
|
|
346
|
-
|
|
347
|
-
|
|
348
|
-
if not x_data:
|
|
349
|
-
raise ValueError("自变量数据不能为空")
|
|
350
|
-
if len(y_data) != len(x_data):
|
|
351
|
-
raise ValueError(f"因变量和自变量的观测数量不一致: y_data={len(y_data)}, x_data={len(x_data)}")
|
|
352
|
-
|
|
353
|
-
# 准备数据
|
|
354
|
-
X = np.array(x_data)
|
|
355
|
-
y = np.array(y_data)
|
|
356
|
-
|
|
357
|
-
# 添加常数项
|
|
358
|
-
X_with_const = sm.add_constant(X)
|
|
359
|
-
|
|
360
|
-
# 拟合模型
|
|
361
|
-
model = sm.OLS(y, X_with_const).fit()
|
|
362
|
-
|
|
363
|
-
# 构建系数字典
|
|
364
|
-
conf_int = model.conf_int()
|
|
365
|
-
coefficients = {}
|
|
434
|
+
# 获取超时管理器
|
|
435
|
+
timeout_manager = ctx.session.context.timeout_manager
|
|
436
|
+
model_timeout = ctx.session.context.performance_config.get("models", {}).get("ols_regression", {}).get("timeout", 60)
|
|
366
437
|
|
|
367
|
-
#
|
|
368
|
-
|
|
369
|
-
|
|
370
|
-
|
|
371
|
-
|
|
372
|
-
|
|
373
|
-
|
|
374
|
-
|
|
375
|
-
if i == 0:
|
|
376
|
-
var_name = "const"
|
|
377
|
-
else:
|
|
378
|
-
var_name = feature_names[i-1]
|
|
379
|
-
|
|
380
|
-
coefficients[var_name] = {
|
|
381
|
-
"coef": float(coef), # 转换numpy.float64为float
|
|
382
|
-
"std_err": float(model.bse[i]),
|
|
383
|
-
"t_value": float(model.tvalues[i]),
|
|
384
|
-
"p_value": float(model.pvalues[i]),
|
|
385
|
-
"ci_lower": float(conf_int[i][0]),
|
|
386
|
-
"ci_upper": float(conf_int[i][1])
|
|
387
|
-
}
|
|
388
|
-
|
|
389
|
-
# 构建结果
|
|
390
|
-
result = OLSRegressionResult(
|
|
391
|
-
rsquared=float(model.rsquared),
|
|
392
|
-
rsquared_adj=float(model.rsquared_adj),
|
|
393
|
-
f_statistic=float(model.fvalue),
|
|
394
|
-
f_pvalue=float(model.f_pvalue),
|
|
395
|
-
aic=float(model.aic),
|
|
396
|
-
bic=float(model.bic),
|
|
397
|
-
coefficients=coefficients
|
|
438
|
+
# 使用超时控制执行计算
|
|
439
|
+
result = await timeout_manager.execute_with_timeout(
|
|
440
|
+
model_name="ols_regression",
|
|
441
|
+
timeout_seconds=model_timeout,
|
|
442
|
+
func=_compute_ols_regression,
|
|
443
|
+
y_data=y_data,
|
|
444
|
+
x_data=x_data,
|
|
445
|
+
feature_names=feature_names
|
|
398
446
|
)
|
|
399
447
|
|
|
400
448
|
await ctx.info("OLS回归分析完成")
|
|
@@ -404,16 +452,22 @@ async def ols_regression(
|
|
|
404
452
|
TextContent(
|
|
405
453
|
type="text",
|
|
406
454
|
text=f"OLS回归分析结果:\n"
|
|
407
|
-
f"R² = {result
|
|
408
|
-
f"调整R² = {result
|
|
409
|
-
f"F统计量 = {result
|
|
410
|
-
f"AIC = {result
|
|
411
|
-
f"回归系数:\n{
|
|
455
|
+
f"R² = {result['rsquared']:.4f}\n"
|
|
456
|
+
f"调整R² = {result['rsquared_adj']:.4f}\n"
|
|
457
|
+
f"F统计量 = {result['f_statistic']:.4f} (p = {result['f_pvalue']:.4f})\n"
|
|
458
|
+
f"AIC = {result['aic']:.2f}, BIC = {result['bic']:.2f}\n\n"
|
|
459
|
+
f"回归系数:\n{result['summary_table']}"
|
|
412
460
|
)
|
|
413
461
|
],
|
|
414
|
-
structuredContent=result
|
|
462
|
+
structuredContent=result
|
|
415
463
|
)
|
|
416
464
|
|
|
465
|
+
except TimeoutError:
|
|
466
|
+
await ctx.error("OLS回归分析超时,请减少数据量或调整配置")
|
|
467
|
+
return CallToolResult(
|
|
468
|
+
content=[TextContent(type="text", text="错误: 计算超时,请减少数据量或调整配置")],
|
|
469
|
+
isError=True
|
|
470
|
+
)
|
|
417
471
|
except Exception as e:
|
|
418
472
|
await ctx.error(f"OLS回归分析出错: {str(e)}")
|
|
419
473
|
return CallToolResult(
|
|
@@ -422,6 +476,66 @@ async def ols_regression(
|
|
|
422
476
|
)
|
|
423
477
|
|
|
424
478
|
|
|
479
|
+
def _compute_ols_regression(y_data: List[float], x_data: List[List[float]], feature_names: Optional[List[str]] = None) -> Dict[str, Any]:
|
|
480
|
+
"""计算OLS回归(同步函数,用于超时控制)"""
|
|
481
|
+
# 数据验证
|
|
482
|
+
if not y_data:
|
|
483
|
+
raise ValueError("因变量数据不能为空")
|
|
484
|
+
if not x_data:
|
|
485
|
+
raise ValueError("自变量数据不能为空")
|
|
486
|
+
if len(y_data) != len(x_data):
|
|
487
|
+
raise ValueError(f"因变量和自变量的观测数量不一致: y_data={len(y_data)}, x_data={len(x_data)}")
|
|
488
|
+
|
|
489
|
+
# 准备数据
|
|
490
|
+
X = np.array(x_data)
|
|
491
|
+
y = np.array(y_data)
|
|
492
|
+
|
|
493
|
+
# 添加常数项
|
|
494
|
+
X_with_const = sm.add_constant(X)
|
|
495
|
+
|
|
496
|
+
# 拟合模型
|
|
497
|
+
model = sm.OLS(y, X_with_const).fit()
|
|
498
|
+
|
|
499
|
+
# 构建系数字典
|
|
500
|
+
conf_int = model.conf_int()
|
|
501
|
+
coefficients = {}
|
|
502
|
+
|
|
503
|
+
# 处理feature_names
|
|
504
|
+
if feature_names is None:
|
|
505
|
+
feature_names = [f"x{i+1}" for i in range(X.shape[1])]
|
|
506
|
+
elif len(feature_names) != X.shape[1]:
|
|
507
|
+
feature_names = [f"x{i+1}" for i in range(X.shape[1])]
|
|
508
|
+
|
|
509
|
+
for i, coef in enumerate(model.params):
|
|
510
|
+
if i == 0:
|
|
511
|
+
var_name = "const"
|
|
512
|
+
else:
|
|
513
|
+
var_name = feature_names[i-1]
|
|
514
|
+
|
|
515
|
+
coefficients[var_name] = {
|
|
516
|
+
"coef": float(coef),
|
|
517
|
+
"std_err": float(model.bse[i]),
|
|
518
|
+
"t_value": float(model.tvalues[i]),
|
|
519
|
+
"p_value": float(model.pvalues[i]),
|
|
520
|
+
"ci_lower": float(conf_int[i][0]),
|
|
521
|
+
"ci_upper": float(conf_int[i][1])
|
|
522
|
+
}
|
|
523
|
+
|
|
524
|
+
# 构建结果
|
|
525
|
+
result = {
|
|
526
|
+
"rsquared": float(model.rsquared),
|
|
527
|
+
"rsquared_adj": float(model.rsquared_adj),
|
|
528
|
+
"f_statistic": float(model.fvalue),
|
|
529
|
+
"f_pvalue": float(model.f_pvalue),
|
|
530
|
+
"aic": float(model.aic),
|
|
531
|
+
"bic": float(model.bic),
|
|
532
|
+
"coefficients": coefficients,
|
|
533
|
+
"summary_table": str(model.summary().tables[1])
|
|
534
|
+
}
|
|
535
|
+
|
|
536
|
+
return result
|
|
537
|
+
|
|
538
|
+
|
|
425
539
|
@mcp.tool()
|
|
426
540
|
async def hypothesis_testing(
|
|
427
541
|
ctx: Context[ServerSession, AppContext],
|
|
@@ -0,0 +1,283 @@
|
|
|
1
|
+
"""
|
|
2
|
+
超时控制模块
|
|
3
|
+
为复杂计算任务提供超时控制和资源管理
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
import asyncio
|
|
7
|
+
import signal
|
|
8
|
+
import threading
|
|
9
|
+
import time
|
|
10
|
+
from typing import Any, Callable, Optional, TypeVar, Union
|
|
11
|
+
from functools import wraps
|
|
12
|
+
from contextlib import contextmanager
|
|
13
|
+
import warnings
|
|
14
|
+
|
|
15
|
+
T = TypeVar('T')
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class TimeoutError(Exception):
|
|
19
|
+
"""超时错误"""
|
|
20
|
+
pass
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class TimeoutManager:
|
|
24
|
+
"""
|
|
25
|
+
超时管理器
|
|
26
|
+
提供同步和异步的超时控制
|
|
27
|
+
"""
|
|
28
|
+
|
|
29
|
+
def __init__(self):
|
|
30
|
+
self._timeout_config = {}
|
|
31
|
+
|
|
32
|
+
def set_timeout_config(self, config: dict) -> None:
|
|
33
|
+
"""设置超时配置"""
|
|
34
|
+
self._timeout_config = config
|
|
35
|
+
|
|
36
|
+
def get_timeout_for_function(self, function_name: str, default: int = 60) -> int:
|
|
37
|
+
"""获取函数的超时时间"""
|
|
38
|
+
# 从配置中获取超时时间
|
|
39
|
+
if function_name in self._timeout_config:
|
|
40
|
+
return self._timeout_config[function_name]
|
|
41
|
+
|
|
42
|
+
# 根据函数类型返回默认超时
|
|
43
|
+
if function_name.startswith(('descriptive_', 'correlation_')):
|
|
44
|
+
return 30
|
|
45
|
+
elif function_name.startswith(('ols_', 'hypothesis_')):
|
|
46
|
+
return 60
|
|
47
|
+
elif function_name.startswith(('time_series_', 'panel_')):
|
|
48
|
+
return 120
|
|
49
|
+
elif function_name.startswith(('var_', 'vecm_', 'garch_')):
|
|
50
|
+
return 180
|
|
51
|
+
elif function_name.startswith(('random_forest_', 'gradient_boosting_')):
|
|
52
|
+
return 300
|
|
53
|
+
else:
|
|
54
|
+
return default
|
|
55
|
+
|
|
56
|
+
async def execute_with_timeout(self, model_name: str, timeout_seconds: int, func: callable, *args, **kwargs):
|
|
57
|
+
"""使用超时执行函数"""
|
|
58
|
+
try:
|
|
59
|
+
if asyncio.iscoroutinefunction(func):
|
|
60
|
+
# 异步函数
|
|
61
|
+
return await asyncio.wait_for(
|
|
62
|
+
func(*args, **kwargs),
|
|
63
|
+
timeout=timeout_seconds
|
|
64
|
+
)
|
|
65
|
+
else:
|
|
66
|
+
# 同步函数 - 在线程池中执行
|
|
67
|
+
loop = asyncio.get_event_loop()
|
|
68
|
+
return await loop.run_in_executor(
|
|
69
|
+
None,
|
|
70
|
+
lambda: self._execute_sync_with_timeout(func, timeout_seconds, *args, **kwargs)
|
|
71
|
+
)
|
|
72
|
+
except asyncio.TimeoutError:
|
|
73
|
+
raise TimeoutError(f"模型 '{model_name}' 执行超时 ({timeout_seconds}秒)")
|
|
74
|
+
|
|
75
|
+
def _execute_sync_with_timeout(self, func: callable, timeout_seconds: int, *args, **kwargs):
|
|
76
|
+
"""同步函数超时执行"""
|
|
77
|
+
import threading
|
|
78
|
+
import queue
|
|
79
|
+
|
|
80
|
+
result_queue = queue.Queue()
|
|
81
|
+
exception_queue = queue.Queue()
|
|
82
|
+
|
|
83
|
+
def worker():
|
|
84
|
+
try:
|
|
85
|
+
result = func(*args, **kwargs)
|
|
86
|
+
result_queue.put(result)
|
|
87
|
+
except Exception as e:
|
|
88
|
+
exception_queue.put(e)
|
|
89
|
+
|
|
90
|
+
thread = threading.Thread(target=worker)
|
|
91
|
+
thread.daemon = True
|
|
92
|
+
thread.start()
|
|
93
|
+
thread.join(timeout_seconds)
|
|
94
|
+
|
|
95
|
+
if thread.is_alive():
|
|
96
|
+
raise TimeoutError(f"同步函数执行超时 ({timeout_seconds}秒)")
|
|
97
|
+
|
|
98
|
+
if not exception_queue.empty():
|
|
99
|
+
raise exception_queue.get()
|
|
100
|
+
|
|
101
|
+
return result_queue.get()
|
|
102
|
+
|
|
103
|
+
@contextmanager
|
|
104
|
+
def timeout_context(self, seconds: int):
|
|
105
|
+
"""
|
|
106
|
+
同步超时上下文管理器
|
|
107
|
+
|
|
108
|
+
Args:
|
|
109
|
+
seconds: 超时时间(秒)
|
|
110
|
+
"""
|
|
111
|
+
def timeout_handler(signum, frame):
|
|
112
|
+
raise TimeoutError(f"操作超时 ({seconds}秒)")
|
|
113
|
+
|
|
114
|
+
# 设置信号处理(仅适用于Unix系统)
|
|
115
|
+
original_handler = signal.signal(signal.SIGALRM, timeout_handler)
|
|
116
|
+
signal.alarm(seconds)
|
|
117
|
+
|
|
118
|
+
try:
|
|
119
|
+
yield
|
|
120
|
+
finally:
|
|
121
|
+
# 取消警报
|
|
122
|
+
signal.alarm(0)
|
|
123
|
+
signal.signal(signal.SIGALRM, original_handler)
|
|
124
|
+
|
|
125
|
+
async def async_timeout_context(self, seconds: int):
|
|
126
|
+
"""
|
|
127
|
+
异步超时上下文管理器
|
|
128
|
+
|
|
129
|
+
Args:
|
|
130
|
+
seconds: 超时时间(秒)
|
|
131
|
+
"""
|
|
132
|
+
try:
|
|
133
|
+
await asyncio.wait_for(asyncio.sleep(0), timeout=seconds)
|
|
134
|
+
except asyncio.TimeoutError:
|
|
135
|
+
raise TimeoutError(f"异步操作超时 ({seconds}秒)")
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
def timeout(seconds: int = 60):
|
|
139
|
+
"""
|
|
140
|
+
同步函数超时装饰器
|
|
141
|
+
|
|
142
|
+
Args:
|
|
143
|
+
seconds: 超时时间(秒)
|
|
144
|
+
"""
|
|
145
|
+
def decorator(func):
|
|
146
|
+
@wraps(func)
|
|
147
|
+
def wrapper(*args, **kwargs):
|
|
148
|
+
manager = TimeoutManager()
|
|
149
|
+
|
|
150
|
+
# 在Windows上使用线程实现超时
|
|
151
|
+
if hasattr(signal, 'SIGALRM'):
|
|
152
|
+
# Unix系统使用信号
|
|
153
|
+
with manager.timeout_context(seconds):
|
|
154
|
+
return func(*args, **kwargs)
|
|
155
|
+
else:
|
|
156
|
+
# Windows系统使用线程
|
|
157
|
+
result = [None]
|
|
158
|
+
exception = [None]
|
|
159
|
+
|
|
160
|
+
def target():
|
|
161
|
+
try:
|
|
162
|
+
result[0] = func(*args, **kwargs)
|
|
163
|
+
except Exception as e:
|
|
164
|
+
exception[0] = e
|
|
165
|
+
|
|
166
|
+
thread = threading.Thread(target=target)
|
|
167
|
+
thread.daemon = True
|
|
168
|
+
thread.start()
|
|
169
|
+
thread.join(seconds)
|
|
170
|
+
|
|
171
|
+
if thread.is_alive():
|
|
172
|
+
raise TimeoutError(f"函数 {func.__name__} 执行超时 ({seconds}秒)")
|
|
173
|
+
|
|
174
|
+
if exception[0]:
|
|
175
|
+
raise exception[0]
|
|
176
|
+
|
|
177
|
+
return result[0]
|
|
178
|
+
|
|
179
|
+
return wrapper
|
|
180
|
+
return decorator
|
|
181
|
+
|
|
182
|
+
|
|
183
|
+
def async_timeout(seconds: int = 60):
|
|
184
|
+
"""
|
|
185
|
+
异步函数超时装饰器
|
|
186
|
+
|
|
187
|
+
Args:
|
|
188
|
+
seconds: 超时时间(秒)
|
|
189
|
+
"""
|
|
190
|
+
def decorator(func):
|
|
191
|
+
@wraps(func)
|
|
192
|
+
async def wrapper(*args, **kwargs):
|
|
193
|
+
try:
|
|
194
|
+
return await asyncio.wait_for(
|
|
195
|
+
func(*args, **kwargs),
|
|
196
|
+
timeout=seconds
|
|
197
|
+
)
|
|
198
|
+
except asyncio.TimeoutError:
|
|
199
|
+
raise TimeoutError(f"异步函数 {func.__name__} 执行超时 ({seconds}秒)")
|
|
200
|
+
|
|
201
|
+
return wrapper
|
|
202
|
+
return decorator
|
|
203
|
+
|
|
204
|
+
|
|
205
|
+
class ResourceMonitor:
|
|
206
|
+
"""
|
|
207
|
+
资源监控器
|
|
208
|
+
监控内存和CPU使用情况
|
|
209
|
+
"""
|
|
210
|
+
|
|
211
|
+
def __init__(self):
|
|
212
|
+
self._start_time = None
|
|
213
|
+
self._peak_memory = 0
|
|
214
|
+
|
|
215
|
+
@contextmanager
|
|
216
|
+
def monitor_resources(self):
|
|
217
|
+
"""监控资源使用的上下文管理器"""
|
|
218
|
+
import psutil
|
|
219
|
+
import os
|
|
220
|
+
|
|
221
|
+
process = psutil.Process(os.getpid())
|
|
222
|
+
self._start_time = time.time()
|
|
223
|
+
initial_memory = process.memory_info().rss
|
|
224
|
+
|
|
225
|
+
try:
|
|
226
|
+
yield
|
|
227
|
+
finally:
|
|
228
|
+
current_memory = process.memory_info().rss
|
|
229
|
+
memory_increase = (current_memory - initial_memory) / 1024 / 1024 # MB
|
|
230
|
+
|
|
231
|
+
execution_time = time.time() - self._start_time
|
|
232
|
+
|
|
233
|
+
# 记录资源使用情况
|
|
234
|
+
if memory_increase > 100: # 超过100MB
|
|
235
|
+
warnings.warn(
|
|
236
|
+
f"高内存使用警告: 内存增加 {memory_increase:.2f}MB, "
|
|
237
|
+
f"执行时间: {execution_time:.2f}秒"
|
|
238
|
+
)
|
|
239
|
+
|
|
240
|
+
|
|
241
|
+
# 全局超时管理器实例
|
|
242
|
+
global_timeout_manager = TimeoutManager()
|
|
243
|
+
|
|
244
|
+
|
|
245
|
+
# 便捷装饰器
|
|
246
|
+
def with_timeout(seconds: int = 60):
|
|
247
|
+
"""便捷超时装饰器,自动选择同步或异步"""
|
|
248
|
+
def decorator(func):
|
|
249
|
+
if asyncio.iscoroutinefunction(func):
|
|
250
|
+
return async_timeout(seconds)(func)
|
|
251
|
+
else:
|
|
252
|
+
return timeout(seconds)(func)
|
|
253
|
+
return decorator
|
|
254
|
+
|
|
255
|
+
|
|
256
|
+
def econometric_timeout(function_name: str = None):
|
|
257
|
+
"""
|
|
258
|
+
计量经济学专用超时装饰器
|
|
259
|
+
根据函数类型自动设置合适的超时时间
|
|
260
|
+
"""
|
|
261
|
+
def decorator(func):
|
|
262
|
+
name = function_name or func.__name__
|
|
263
|
+
timeout_seconds = global_timeout_manager.get_timeout_for_function(name)
|
|
264
|
+
|
|
265
|
+
if asyncio.iscoroutinefunction(func):
|
|
266
|
+
return async_timeout(timeout_seconds)(func)
|
|
267
|
+
else:
|
|
268
|
+
return timeout(timeout_seconds)(func)
|
|
269
|
+
|
|
270
|
+
return decorator
|
|
271
|
+
|
|
272
|
+
|
|
273
|
+
# 导出主要类和函数
|
|
274
|
+
__all__ = [
|
|
275
|
+
"TimeoutError",
|
|
276
|
+
"TimeoutManager",
|
|
277
|
+
"timeout",
|
|
278
|
+
"async_timeout",
|
|
279
|
+
"with_timeout",
|
|
280
|
+
"econometric_timeout",
|
|
281
|
+
"ResourceMonitor",
|
|
282
|
+
"global_timeout_manager"
|
|
283
|
+
]
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/machine_learning.py
RENAMED
|
File without changes
|
|
File without changes
|
{aigroup_econ_mcp-0.3.5 → aigroup_econ_mcp-0.3.6}/src/aigroup_econ_mcp/tools/optimized_example.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|