adam-community 1.0.28__py3-none-any.whl → 1.0.29__py3-none-any.whl
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.
- adam_community/__init__.py +8 -6
- adam_community/__version__.py +1 -1
- adam_community/util/__init__.py +73 -0
- adam_community/util/api.py +178 -0
- adam_community/util/cmd.py +166 -0
- adam_community/util/config.py +8 -0
- adam_community/util/markdown.py +8 -0
- adam_community/util/rag.py +126 -0
- adam_community/util/state.py +195 -0
- {adam_community-1.0.28.dist-info → adam_community-1.0.29.dist-info}/METADATA +78 -1
- {adam_community-1.0.28.dist-info → adam_community-1.0.29.dist-info}/RECORD +14 -7
- {adam_community-1.0.28.dist-info → adam_community-1.0.29.dist-info}/WHEEL +0 -0
- {adam_community-1.0.28.dist-info → adam_community-1.0.29.dist-info}/entry_points.txt +0 -0
- {adam_community-1.0.28.dist-info → adam_community-1.0.29.dist-info}/top_level.txt +0 -0
adam_community/__init__.py
CHANGED
|
@@ -1,15 +1,17 @@
|
|
|
1
|
-
from .util import messageSend, knowledgeSearch, completionCreate, runCmd
|
|
1
|
+
from .util import messageSend, knowledgeSearch, completionCreate, runCmd, execCmd, CmdResult
|
|
2
2
|
from .tool import Tool
|
|
3
3
|
from .cli.parser import parse_python_file, parse_directory, convert_python_type_to_json_schema
|
|
4
4
|
from .__version__ import __version__
|
|
5
5
|
|
|
6
6
|
__all__ = [
|
|
7
|
-
'Tool',
|
|
8
|
-
'messageSend',
|
|
9
|
-
'knowledgeSearch',
|
|
10
|
-
'completionCreate',
|
|
7
|
+
'Tool',
|
|
8
|
+
'messageSend',
|
|
9
|
+
'knowledgeSearch',
|
|
10
|
+
'completionCreate',
|
|
11
11
|
'runCmd',
|
|
12
|
+
'execCmd',
|
|
13
|
+
'CmdResult',
|
|
12
14
|
'parse_python_file',
|
|
13
|
-
'parse_directory',
|
|
15
|
+
'parse_directory',
|
|
14
16
|
'convert_python_type_to_json_schema'
|
|
15
17
|
]
|
adam_community/__version__.py
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
__version__ = "1.0.
|
|
1
|
+
__version__ = "1.0.29"
|
|
@@ -0,0 +1,73 @@
|
|
|
1
|
+
# Re-export all public APIs to maintain backward compatibility
|
|
2
|
+
# Usage: from adam_community.util import execCmd, runCmd, ...
|
|
3
|
+
|
|
4
|
+
from .config import (
|
|
5
|
+
ADAM_API_HOST,
|
|
6
|
+
ADAM_API_TOKEN,
|
|
7
|
+
ADAM_TASK_ID,
|
|
8
|
+
ADAM_USER_ID,
|
|
9
|
+
CONDA_ENV,
|
|
10
|
+
ADAM_TASK_DIR,
|
|
11
|
+
)
|
|
12
|
+
|
|
13
|
+
from .cmd import (
|
|
14
|
+
runCmd,
|
|
15
|
+
execCmd,
|
|
16
|
+
CmdResult,
|
|
17
|
+
)
|
|
18
|
+
|
|
19
|
+
from .api import (
|
|
20
|
+
retry_on_exception,
|
|
21
|
+
_make_http_request,
|
|
22
|
+
messageSend,
|
|
23
|
+
completionCreate,
|
|
24
|
+
DynamicObject,
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
from .rag import (
|
|
28
|
+
RAG,
|
|
29
|
+
knowledgeSearch,
|
|
30
|
+
)
|
|
31
|
+
|
|
32
|
+
from .state import (
|
|
33
|
+
_StatesManager,
|
|
34
|
+
setState,
|
|
35
|
+
getState,
|
|
36
|
+
trackPath,
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
from .markdown import (
|
|
40
|
+
markdown_color,
|
|
41
|
+
markdown_terminal,
|
|
42
|
+
)
|
|
43
|
+
|
|
44
|
+
__all__ = [
|
|
45
|
+
# config
|
|
46
|
+
'ADAM_API_HOST',
|
|
47
|
+
'ADAM_API_TOKEN',
|
|
48
|
+
'ADAM_TASK_ID',
|
|
49
|
+
'ADAM_USER_ID',
|
|
50
|
+
'CONDA_ENV',
|
|
51
|
+
'ADAM_TASK_DIR',
|
|
52
|
+
# cmd
|
|
53
|
+
'runCmd',
|
|
54
|
+
'execCmd',
|
|
55
|
+
'CmdResult',
|
|
56
|
+
# api
|
|
57
|
+
'retry_on_exception',
|
|
58
|
+
'_make_http_request',
|
|
59
|
+
'messageSend',
|
|
60
|
+
'completionCreate',
|
|
61
|
+
'DynamicObject',
|
|
62
|
+
# rag
|
|
63
|
+
'RAG',
|
|
64
|
+
'knowledgeSearch',
|
|
65
|
+
# state
|
|
66
|
+
'_StatesManager',
|
|
67
|
+
'setState',
|
|
68
|
+
'getState',
|
|
69
|
+
'trackPath',
|
|
70
|
+
# markdown
|
|
71
|
+
'markdown_color',
|
|
72
|
+
'markdown_terminal',
|
|
73
|
+
]
|
|
@@ -0,0 +1,178 @@
|
|
|
1
|
+
import time
|
|
2
|
+
import urllib.request
|
|
3
|
+
import urllib.error
|
|
4
|
+
import json
|
|
5
|
+
import ssl
|
|
6
|
+
from functools import wraps
|
|
7
|
+
|
|
8
|
+
from .config import ADAM_API_HOST, ADAM_API_TOKEN, ADAM_TASK_ID
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def retry_on_exception(max_retries=3, delay=5):
|
|
12
|
+
"""
|
|
13
|
+
重试装饰器,在发生异常时最多重试指定次数
|
|
14
|
+
|
|
15
|
+
Args:
|
|
16
|
+
max_retries: 最大重试次数
|
|
17
|
+
delay: 重试间隔时间(秒)
|
|
18
|
+
"""
|
|
19
|
+
def decorator(func):
|
|
20
|
+
@wraps(func)
|
|
21
|
+
def wrapper(*args, **kwargs):
|
|
22
|
+
last_exception = None
|
|
23
|
+
for attempt in range(max_retries):
|
|
24
|
+
try:
|
|
25
|
+
return func(*args, **kwargs)
|
|
26
|
+
except Exception as e:
|
|
27
|
+
last_exception = e
|
|
28
|
+
if attempt < max_retries - 1:
|
|
29
|
+
time.sleep(delay)
|
|
30
|
+
continue
|
|
31
|
+
raise last_exception
|
|
32
|
+
return None
|
|
33
|
+
return wrapper
|
|
34
|
+
return decorator
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def _make_http_request(url, data, return_json=True):
|
|
38
|
+
"""
|
|
39
|
+
通用的 HTTP POST 请求函数
|
|
40
|
+
|
|
41
|
+
Args:
|
|
42
|
+
url: 请求的 URL
|
|
43
|
+
data: 请求数据(字典格式)
|
|
44
|
+
return_json: 是否将响应解析为 JSON,False 则返回文本
|
|
45
|
+
|
|
46
|
+
Returns:
|
|
47
|
+
响应数据(JSON 或文本)
|
|
48
|
+
"""
|
|
49
|
+
if not ADAM_API_TOKEN:
|
|
50
|
+
raise ValueError("ADAM_API_TOKEN environment variable is not set")
|
|
51
|
+
|
|
52
|
+
headers = {
|
|
53
|
+
"Authorization": f"Bearer {ADAM_API_TOKEN}",
|
|
54
|
+
"Content-Type": "application/json"
|
|
55
|
+
}
|
|
56
|
+
|
|
57
|
+
try:
|
|
58
|
+
# 将请求数据转换为JSON字符串并编码
|
|
59
|
+
json_data = json.dumps(data).encode('utf-8')
|
|
60
|
+
|
|
61
|
+
# 创建请求对象
|
|
62
|
+
req = urllib.request.Request(url, data=json_data, headers=headers, method='POST')
|
|
63
|
+
|
|
64
|
+
context = ssl._create_unverified_context()
|
|
65
|
+
|
|
66
|
+
# 发送请求
|
|
67
|
+
with urllib.request.urlopen(req, context=context) as response:
|
|
68
|
+
# 检查响应状态码
|
|
69
|
+
if response.status >= 400:
|
|
70
|
+
raise Exception(f"HTTP错误: {response.status}")
|
|
71
|
+
|
|
72
|
+
# 读取响应
|
|
73
|
+
response_data = response.read().decode('utf-8')
|
|
74
|
+
|
|
75
|
+
# 根据需要返回 JSON 或文本
|
|
76
|
+
if return_json:
|
|
77
|
+
return json.loads(response_data)
|
|
78
|
+
else:
|
|
79
|
+
return response_data
|
|
80
|
+
|
|
81
|
+
except urllib.error.HTTPError as e:
|
|
82
|
+
raise Exception(f"HTTP请求失败 - HTTP错误 {e.code}: {e.reason}")
|
|
83
|
+
except urllib.error.URLError as e:
|
|
84
|
+
raise Exception(f"HTTP请求失败 - 网络错误: {str(e)}")
|
|
85
|
+
except json.JSONDecodeError as e:
|
|
86
|
+
if return_json:
|
|
87
|
+
raise Exception(f"HTTP请求失败 - JSON解析错误: {str(e)}")
|
|
88
|
+
else:
|
|
89
|
+
raise Exception(f"HTTP请求失败 - 响应解码错误: {str(e)}")
|
|
90
|
+
except Exception as e:
|
|
91
|
+
raise Exception(f"HTTP请求失败: {str(e)}")
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
def messageSend(message):
|
|
95
|
+
"""
|
|
96
|
+
发送消息给用户
|
|
97
|
+
"""
|
|
98
|
+
if not ADAM_API_HOST:
|
|
99
|
+
raise ValueError("ADAM_API_HOST environment variable is not set")
|
|
100
|
+
if not ADAM_TASK_ID:
|
|
101
|
+
raise ValueError("ADAM_TASK_ID environment variable is not set")
|
|
102
|
+
|
|
103
|
+
url = f"{ADAM_API_HOST}/api/task/create_message"
|
|
104
|
+
|
|
105
|
+
request_data = {
|
|
106
|
+
"task_id": ADAM_TASK_ID,
|
|
107
|
+
"message": message,
|
|
108
|
+
"role": "tool"
|
|
109
|
+
}
|
|
110
|
+
|
|
111
|
+
try:
|
|
112
|
+
return _make_http_request(url, request_data, return_json=True)
|
|
113
|
+
except Exception as e:
|
|
114
|
+
raise Exception(f"发送消息失败: {str(e)}")
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
class DynamicObject:
|
|
118
|
+
"""
|
|
119
|
+
动态对象类,用于处理 JSON 响应
|
|
120
|
+
|
|
121
|
+
这个类可以动态地将字典转换为对象属性,支持嵌套的字典结构
|
|
122
|
+
同时保持原始数据的访问能力
|
|
123
|
+
"""
|
|
124
|
+
def __init__(self, data):
|
|
125
|
+
self._raw_data = data
|
|
126
|
+
if isinstance(data, dict):
|
|
127
|
+
for key, value in data.items():
|
|
128
|
+
if isinstance(value, (dict, list)):
|
|
129
|
+
setattr(self, key, self._convert_value(value))
|
|
130
|
+
else:
|
|
131
|
+
setattr(self, key, value)
|
|
132
|
+
|
|
133
|
+
def _convert_value(self, value):
|
|
134
|
+
"""递归转换嵌套的数据结构"""
|
|
135
|
+
if isinstance(value, dict):
|
|
136
|
+
return DynamicObject(value)
|
|
137
|
+
elif isinstance(value, list):
|
|
138
|
+
return [self._convert_value(item) for item in value]
|
|
139
|
+
return value
|
|
140
|
+
|
|
141
|
+
def __getattr__(self, name):
|
|
142
|
+
"""处理未定义的属性访问"""
|
|
143
|
+
return self._raw_data.get(name)
|
|
144
|
+
|
|
145
|
+
def __getitem__(self, key):
|
|
146
|
+
"""支持字典式访问"""
|
|
147
|
+
return self._raw_data.get(key)
|
|
148
|
+
|
|
149
|
+
def __str__(self):
|
|
150
|
+
"""返回可读的字符串表示"""
|
|
151
|
+
return f"DynamicObject({self._raw_data})"
|
|
152
|
+
|
|
153
|
+
def __repr__(self):
|
|
154
|
+
return self.__str__()
|
|
155
|
+
|
|
156
|
+
def to_dict(self):
|
|
157
|
+
"""将对象转换回字典"""
|
|
158
|
+
return self._raw_data
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
@retry_on_exception(max_retries=3)
|
|
162
|
+
def completionCreate(request_params):
|
|
163
|
+
"""
|
|
164
|
+
创建 openai 的代理函数
|
|
165
|
+
|
|
166
|
+
Returns:
|
|
167
|
+
DynamicObject: 一个动态对象,可以像访问属性一样访问 API 响应的所有字段
|
|
168
|
+
"""
|
|
169
|
+
if not ADAM_API_HOST:
|
|
170
|
+
raise ValueError("ADAM_API_HOST environment variable is not set")
|
|
171
|
+
|
|
172
|
+
url = f"{ADAM_API_HOST}/api/chat/completions"
|
|
173
|
+
|
|
174
|
+
try:
|
|
175
|
+
response = _make_http_request(url, request_params, return_json=True)
|
|
176
|
+
return DynamicObject(response)
|
|
177
|
+
except Exception as e:
|
|
178
|
+
raise Exception(f"调用聊天补全接口失败: {str(e)}")
|
|
@@ -0,0 +1,166 @@
|
|
|
1
|
+
import subprocess
|
|
2
|
+
import sys
|
|
3
|
+
import os
|
|
4
|
+
import time
|
|
5
|
+
import threading
|
|
6
|
+
from dataclasses import dataclass
|
|
7
|
+
from typing import Callable, Optional
|
|
8
|
+
from subprocess import CalledProcessError, TimeoutExpired
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@dataclass
|
|
12
|
+
class CmdResult:
|
|
13
|
+
"""命令执行结果"""
|
|
14
|
+
stdout: str
|
|
15
|
+
stderr: str
|
|
16
|
+
returncode: int
|
|
17
|
+
duration: float
|
|
18
|
+
command: str
|
|
19
|
+
timed_out: bool = False
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def execCmd(
|
|
23
|
+
cmd: str,
|
|
24
|
+
*,
|
|
25
|
+
timeout: Optional[float] = None,
|
|
26
|
+
cwd: Optional[str] = None,
|
|
27
|
+
env: Optional[dict] = None,
|
|
28
|
+
shell: str = "/bin/bash",
|
|
29
|
+
echo: bool = False,
|
|
30
|
+
on_stdout: Optional[Callable[[str], None]] = None,
|
|
31
|
+
on_stderr: Optional[Callable[[str], None]] = None,
|
|
32
|
+
) -> CmdResult:
|
|
33
|
+
"""
|
|
34
|
+
执行 shell 命令
|
|
35
|
+
|
|
36
|
+
Args:
|
|
37
|
+
cmd: 要执行的命令
|
|
38
|
+
timeout: 超时秒数,None 表示无限制
|
|
39
|
+
cwd: 工作目录
|
|
40
|
+
env: 环境变量(合并到当前环境)
|
|
41
|
+
shell: shell 路径,默认 /bin/bash
|
|
42
|
+
echo: 是否实时打印到控制台
|
|
43
|
+
on_stdout: stdout 回调函数,签名 (line: str) -> None
|
|
44
|
+
on_stderr: stderr 回调函数,签名 (line: str) -> None
|
|
45
|
+
|
|
46
|
+
Returns:
|
|
47
|
+
CmdResult: 包含 stdout, stderr, returncode, duration, command, timed_out
|
|
48
|
+
|
|
49
|
+
Raises:
|
|
50
|
+
subprocess.CalledProcessError: 命令返回非零退出码
|
|
51
|
+
subprocess.TimeoutExpired: 命令执行超时
|
|
52
|
+
"""
|
|
53
|
+
# 合并环境变量
|
|
54
|
+
process_env = os.environ.copy()
|
|
55
|
+
if env:
|
|
56
|
+
process_env.update(env)
|
|
57
|
+
|
|
58
|
+
start_time = time.time()
|
|
59
|
+
stdout_lines = []
|
|
60
|
+
stderr_lines = []
|
|
61
|
+
|
|
62
|
+
process = subprocess.Popen(
|
|
63
|
+
cmd,
|
|
64
|
+
shell=True,
|
|
65
|
+
executable=shell,
|
|
66
|
+
stdout=subprocess.PIPE,
|
|
67
|
+
stderr=subprocess.PIPE,
|
|
68
|
+
universal_newlines=True,
|
|
69
|
+
bufsize=1,
|
|
70
|
+
cwd=cwd,
|
|
71
|
+
env=process_env,
|
|
72
|
+
)
|
|
73
|
+
|
|
74
|
+
def read_stream(stream, lines_buffer, callback, is_stderr=False):
|
|
75
|
+
"""读取流并处理输出"""
|
|
76
|
+
for line in iter(stream.readline, ''):
|
|
77
|
+
line = line.rstrip('\n\r')
|
|
78
|
+
lines_buffer.append(line)
|
|
79
|
+
if echo:
|
|
80
|
+
if is_stderr:
|
|
81
|
+
print(line, file=sys.stderr)
|
|
82
|
+
else:
|
|
83
|
+
print(line)
|
|
84
|
+
if callback:
|
|
85
|
+
callback(line)
|
|
86
|
+
stream.close()
|
|
87
|
+
|
|
88
|
+
# 使用线程并发读取 stdout 和 stderr
|
|
89
|
+
stdout_thread = threading.Thread(
|
|
90
|
+
target=read_stream,
|
|
91
|
+
args=(process.stdout, stdout_lines, on_stdout, False)
|
|
92
|
+
)
|
|
93
|
+
stderr_thread = threading.Thread(
|
|
94
|
+
target=read_stream,
|
|
95
|
+
args=(process.stderr, stderr_lines, on_stderr, True)
|
|
96
|
+
)
|
|
97
|
+
|
|
98
|
+
stdout_thread.start()
|
|
99
|
+
stderr_thread.start()
|
|
100
|
+
|
|
101
|
+
# 等待进程完成或超时
|
|
102
|
+
timed_out = False
|
|
103
|
+
try:
|
|
104
|
+
process.wait(timeout=timeout)
|
|
105
|
+
except subprocess.TimeoutExpired:
|
|
106
|
+
timed_out = True
|
|
107
|
+
process.kill()
|
|
108
|
+
# 等待线程完成读取剩余输出
|
|
109
|
+
stdout_thread.join(timeout=1)
|
|
110
|
+
stderr_thread.join(timeout=1)
|
|
111
|
+
|
|
112
|
+
duration = time.time() - start_time
|
|
113
|
+
stdout_str = '\n'.join(stdout_lines)
|
|
114
|
+
stderr_str = '\n'.join(stderr_lines)
|
|
115
|
+
|
|
116
|
+
raise TimeoutExpired(
|
|
117
|
+
cmd=cmd,
|
|
118
|
+
timeout=timeout,
|
|
119
|
+
output=stdout_str,
|
|
120
|
+
stderr=stderr_str,
|
|
121
|
+
)
|
|
122
|
+
|
|
123
|
+
# 等待线程完成
|
|
124
|
+
stdout_thread.join()
|
|
125
|
+
stderr_thread.join()
|
|
126
|
+
|
|
127
|
+
duration = time.time() - start_time
|
|
128
|
+
stdout_str = '\n'.join(stdout_lines)
|
|
129
|
+
stderr_str = '\n'.join(stderr_lines)
|
|
130
|
+
|
|
131
|
+
# 检查返回码
|
|
132
|
+
if process.returncode != 0:
|
|
133
|
+
raise CalledProcessError(
|
|
134
|
+
returncode=process.returncode,
|
|
135
|
+
cmd=cmd,
|
|
136
|
+
output=stdout_str,
|
|
137
|
+
stderr=stderr_str,
|
|
138
|
+
)
|
|
139
|
+
|
|
140
|
+
return CmdResult(
|
|
141
|
+
stdout=stdout_str,
|
|
142
|
+
stderr=stderr_str,
|
|
143
|
+
returncode=process.returncode,
|
|
144
|
+
duration=duration,
|
|
145
|
+
command=cmd,
|
|
146
|
+
timed_out=False,
|
|
147
|
+
)
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
def runCmd(cmd):
|
|
151
|
+
"""
|
|
152
|
+
执行命令,实时输出执行结果(向后兼容)
|
|
153
|
+
|
|
154
|
+
内部调用 execCmd,失败时直接退出进程。
|
|
155
|
+
推荐使用 execCmd 以获得更好的错误处理能力。
|
|
156
|
+
|
|
157
|
+
Args:
|
|
158
|
+
cmd: 要执行的命令
|
|
159
|
+
|
|
160
|
+
Returns:
|
|
161
|
+
CmdResult: 命令执行结果
|
|
162
|
+
"""
|
|
163
|
+
if os.getenv("ADAM_OUTPUT_RAW"):
|
|
164
|
+
return cmd
|
|
165
|
+
|
|
166
|
+
return execCmd(cmd, echo=True)
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
import os
|
|
2
|
+
|
|
3
|
+
ADAM_API_HOST = os.getenv('ADAM_API_HOST', 'https://sidereus-ai.com')
|
|
4
|
+
ADAM_API_TOKEN = os.getenv('ADAM_API_TOKEN')
|
|
5
|
+
ADAM_TASK_ID = os.getenv('ADAM_TASK_ID')
|
|
6
|
+
ADAM_USER_ID = os.getenv('ADAM_USER_ID')
|
|
7
|
+
CONDA_ENV = os.getenv('CONDA_ENV')
|
|
8
|
+
ADAM_TASK_DIR = os.getenv('ADAM_TASK_DIR')
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
def markdown_color(content, color):
|
|
2
|
+
return f'<span style="color: {color}">{content}</span>'
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
def markdown_terminal(content, conda_env="base", user="Adam", workdir=""):
|
|
6
|
+
user = markdown_color(f"{user}@Adam", "green")
|
|
7
|
+
workdir = markdown_color(f":~/{workdir}", "blue")
|
|
8
|
+
return f'({conda_env}) {user}{workdir}$ {content}'
|
|
@@ -0,0 +1,126 @@
|
|
|
1
|
+
import json
|
|
2
|
+
import logging
|
|
3
|
+
from subprocess import run, CalledProcessError
|
|
4
|
+
|
|
5
|
+
from .api import retry_on_exception
|
|
6
|
+
|
|
7
|
+
logger = logging.getLogger(__name__)
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def _build_akb_query_command(query: str, collection: str) -> str:
|
|
11
|
+
"""
|
|
12
|
+
构建 akb 查询命令,包含目录存在性检查
|
|
13
|
+
|
|
14
|
+
Args:
|
|
15
|
+
query: 查询字符串
|
|
16
|
+
collection: 知识库名称
|
|
17
|
+
|
|
18
|
+
Returns:
|
|
19
|
+
完整的 bash 命令字符串
|
|
20
|
+
"""
|
|
21
|
+
collection_dir = f"/share/programs/akb/database/{collection}"
|
|
22
|
+
|
|
23
|
+
return f"""if [ ! -d "{collection_dir}" ]; then
|
|
24
|
+
echo "知识库目录不存在: {collection_dir}"
|
|
25
|
+
exit 1
|
|
26
|
+
fi
|
|
27
|
+
akb simple-query \\
|
|
28
|
+
-i {collection_dir}/{collection}.index \\
|
|
29
|
+
-m /share/programs/BAAI/bge-m3 \\
|
|
30
|
+
-p {collection_dir}/{collection}.parquet \\
|
|
31
|
+
-l 5 \\
|
|
32
|
+
-f json -q "{query}\""""
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class RAG:
|
|
36
|
+
"""
|
|
37
|
+
搜寻RAG知识
|
|
38
|
+
|
|
39
|
+
:param str query: 需要运行的命令
|
|
40
|
+
:param str collection: 搜寻的知识库名称。必须是下面之一:"DSDP","MPNN","PySCF","RFdiffusion","gaussian","protenix","GPU4PySCF","MolecularDynamics","RDkit","OpenBabel"
|
|
41
|
+
"""
|
|
42
|
+
|
|
43
|
+
def call(self, query: str, collection: str):
|
|
44
|
+
cmd = _build_akb_query_command(query, collection)
|
|
45
|
+
logger.info(f"搜索知识库 {collection}: {query}")
|
|
46
|
+
|
|
47
|
+
try:
|
|
48
|
+
result = run(cmd, shell='/bin/bash', check=True, text=True, capture_output=True)
|
|
49
|
+
except CalledProcessError as e:
|
|
50
|
+
logger.error(e.stderr.strip())
|
|
51
|
+
return e.stderr.strip()
|
|
52
|
+
|
|
53
|
+
try:
|
|
54
|
+
parsed_data = json.loads(result.stdout)
|
|
55
|
+
# 如果返回的是数组,直接使用
|
|
56
|
+
if isinstance(parsed_data, list):
|
|
57
|
+
return "\n".join(parsed_data)
|
|
58
|
+
# 如果返回的是对象且包含 text 字段,使用 text 字段
|
|
59
|
+
elif isinstance(parsed_data, dict) and "text" in parsed_data:
|
|
60
|
+
text_array = parsed_data["text"]
|
|
61
|
+
if isinstance(text_array, list):
|
|
62
|
+
return "\n".join(text_array)
|
|
63
|
+
else:
|
|
64
|
+
return str(text_array)
|
|
65
|
+
else:
|
|
66
|
+
# 其他情况返回原始输出
|
|
67
|
+
return result.stdout
|
|
68
|
+
except Exception as e:
|
|
69
|
+
logger.error(e)
|
|
70
|
+
return result.stdout
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
@retry_on_exception(max_retries=3)
|
|
74
|
+
def knowledgeSearch(query_info, messages_prev, project_name, collection_name, max_messages=4):
|
|
75
|
+
"""
|
|
76
|
+
知识库检索函数
|
|
77
|
+
"""
|
|
78
|
+
if not collection_name:
|
|
79
|
+
raise ValueError("collection_name 参数是必需的")
|
|
80
|
+
|
|
81
|
+
try:
|
|
82
|
+
# 使用本地RAG实现
|
|
83
|
+
rag = RAG()
|
|
84
|
+
result = rag.call(query_info, collection_name)
|
|
85
|
+
|
|
86
|
+
# 构造返回格式,保持与原有API格式一致
|
|
87
|
+
response_data = {
|
|
88
|
+
"code": 0,
|
|
89
|
+
"data": {
|
|
90
|
+
"collection_name": collection_name,
|
|
91
|
+
"count": 1,
|
|
92
|
+
"result_list": [{
|
|
93
|
+
"chunk_id": 0,
|
|
94
|
+
"chunk_source": "document",
|
|
95
|
+
"chunk_title": result.split("\n")[0] if result else "",
|
|
96
|
+
"chunk_type": "text",
|
|
97
|
+
"content": result,
|
|
98
|
+
"doc_info": {
|
|
99
|
+
"create_time": 0,
|
|
100
|
+
"doc_id": "local_doc",
|
|
101
|
+
"doc_name": f"{collection_name}_knowledge",
|
|
102
|
+
"doc_type": "text",
|
|
103
|
+
"source": "local"
|
|
104
|
+
},
|
|
105
|
+
"score": 1.0
|
|
106
|
+
}],
|
|
107
|
+
"token_usage": {
|
|
108
|
+
"embedding_token_usage": {
|
|
109
|
+
"completion_tokens": 0,
|
|
110
|
+
"prompt_tokens": 0,
|
|
111
|
+
"total_tokens": 0
|
|
112
|
+
},
|
|
113
|
+
"rerank_token_usage": 0
|
|
114
|
+
}
|
|
115
|
+
},
|
|
116
|
+
"message": "success",
|
|
117
|
+
"request_id": "local_request"
|
|
118
|
+
}
|
|
119
|
+
|
|
120
|
+
return json.dumps(response_data, ensure_ascii=False)
|
|
121
|
+
except Exception as e:
|
|
122
|
+
error_response = {
|
|
123
|
+
"code": 1,
|
|
124
|
+
"message": f"知识库搜索失败: {str(e)}"
|
|
125
|
+
}
|
|
126
|
+
return json.dumps(error_response, ensure_ascii=False)
|
|
@@ -0,0 +1,195 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import json
|
|
3
|
+
import time
|
|
4
|
+
|
|
5
|
+
from .config import ADAM_TASK_DIR
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class _StatesManager:
|
|
9
|
+
"""内部状态管理类"""
|
|
10
|
+
_tool_name = "tool"
|
|
11
|
+
|
|
12
|
+
def _get_states_file(self) -> str:
|
|
13
|
+
"""获取 states.json 文件路径"""
|
|
14
|
+
# 1. 优先使用环境变量
|
|
15
|
+
task_dir = ADAM_TASK_DIR
|
|
16
|
+
if task_dir:
|
|
17
|
+
return os.path.join(task_dir, ".slurm", "states.json")
|
|
18
|
+
|
|
19
|
+
# 2. 尝试当前目录的 .slurm
|
|
20
|
+
current_dir = os.getcwd()
|
|
21
|
+
local_states = os.path.join(current_dir, ".slurm", "states.json")
|
|
22
|
+
if os.path.exists(local_states):
|
|
23
|
+
return local_states
|
|
24
|
+
|
|
25
|
+
# 3. 在当前目录创建
|
|
26
|
+
return local_states
|
|
27
|
+
|
|
28
|
+
def _read(self) -> dict:
|
|
29
|
+
"""读取 states.json,保留所有来源的数据"""
|
|
30
|
+
states_file = self._get_states_file()
|
|
31
|
+
if not os.path.exists(states_file):
|
|
32
|
+
return {
|
|
33
|
+
"files": {"updated_at": None, "sources": {}},
|
|
34
|
+
"states": {"source": None, "updated_at": None, "data": {}}
|
|
35
|
+
}
|
|
36
|
+
try:
|
|
37
|
+
with open(states_file, 'r', encoding='utf-8') as f:
|
|
38
|
+
return json.load(f)
|
|
39
|
+
except (json.JSONDecodeError, IOError):
|
|
40
|
+
return {
|
|
41
|
+
"files": {"updated_at": None, "sources": {}},
|
|
42
|
+
"states": {"source": None, "updated_at": None, "data": {}}
|
|
43
|
+
}
|
|
44
|
+
|
|
45
|
+
def _write(self, data: dict):
|
|
46
|
+
"""写入 states.json"""
|
|
47
|
+
states_file = self._get_states_file()
|
|
48
|
+
os.makedirs(os.path.dirname(states_file), exist_ok=True)
|
|
49
|
+
with open(states_file, 'w', encoding='utf-8') as f:
|
|
50
|
+
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
51
|
+
|
|
52
|
+
def _set_nested(self, data: dict, key: str, value):
|
|
53
|
+
"""设置嵌套值,支持 dot notation"""
|
|
54
|
+
keys = key.split('.')
|
|
55
|
+
current = data
|
|
56
|
+
for k in keys[:-1]:
|
|
57
|
+
current = current.setdefault(k, {})
|
|
58
|
+
current[keys[-1]] = value
|
|
59
|
+
|
|
60
|
+
def _get_nested(self, data: dict, key: str, default=None):
|
|
61
|
+
"""获取嵌套值,支持 dot notation"""
|
|
62
|
+
keys = key.split('.')
|
|
63
|
+
current = data
|
|
64
|
+
for k in keys:
|
|
65
|
+
if isinstance(current, dict):
|
|
66
|
+
current = current.get(k)
|
|
67
|
+
else:
|
|
68
|
+
return default
|
|
69
|
+
if current is None:
|
|
70
|
+
return default
|
|
71
|
+
return current
|
|
72
|
+
|
|
73
|
+
def set(self, key: str, value):
|
|
74
|
+
"""
|
|
75
|
+
设置状态
|
|
76
|
+
|
|
77
|
+
特殊 key:
|
|
78
|
+
"files" - 记录文件列表,会自动补充到 files.sources[tool_name]
|
|
79
|
+
其他 - 记录到 states.data
|
|
80
|
+
"""
|
|
81
|
+
data = self._read()
|
|
82
|
+
|
|
83
|
+
if key == "files":
|
|
84
|
+
data["files"]["sources"][self._tool_name] = {
|
|
85
|
+
"updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
|
86
|
+
"items": value
|
|
87
|
+
}
|
|
88
|
+
else:
|
|
89
|
+
self._set_nested(data["states"]["data"], key, value)
|
|
90
|
+
|
|
91
|
+
self._write(data)
|
|
92
|
+
|
|
93
|
+
def get(self, key: str):
|
|
94
|
+
"""
|
|
95
|
+
获取状态
|
|
96
|
+
|
|
97
|
+
Returns:
|
|
98
|
+
状态值,不存在返回 None
|
|
99
|
+
"""
|
|
100
|
+
data = self._read()
|
|
101
|
+
|
|
102
|
+
if key == "files":
|
|
103
|
+
# 合并所有来源的文件列表
|
|
104
|
+
merged = {}
|
|
105
|
+
for source, content in data["files"]["sources"].items():
|
|
106
|
+
for item in content.get("items", []):
|
|
107
|
+
path = item.get("path")
|
|
108
|
+
if not path:
|
|
109
|
+
continue
|
|
110
|
+
mtime = item.get("mtime", 0)
|
|
111
|
+
if path not in merged or mtime > merged[path]["mtime"]:
|
|
112
|
+
merged[path] = item
|
|
113
|
+
result = list(merged.values())
|
|
114
|
+
# 按 mtime 降序排序
|
|
115
|
+
result.sort(key=lambda x: x.get("mtime", 0), reverse=True)
|
|
116
|
+
return result
|
|
117
|
+
|
|
118
|
+
return self._get_nested(data["states"]["data"], key)
|
|
119
|
+
|
|
120
|
+
def cleanup(self):
|
|
121
|
+
"""清理当前工具的所有记录"""
|
|
122
|
+
data = self._read()
|
|
123
|
+
data["files"]["sources"].pop(self._tool_name, None)
|
|
124
|
+
self._write(data)
|
|
125
|
+
|
|
126
|
+
def trackPath(self, path: str, max_items: int = 30):
|
|
127
|
+
"""
|
|
128
|
+
追踪文件/目录,自动检测文件信息并记录
|
|
129
|
+
|
|
130
|
+
Args:
|
|
131
|
+
path: 文件或目录路径
|
|
132
|
+
max_items: 最大记录数量,先进先出
|
|
133
|
+
"""
|
|
134
|
+
if not os.path.exists(path):
|
|
135
|
+
return
|
|
136
|
+
|
|
137
|
+
# 检测文件信息
|
|
138
|
+
try:
|
|
139
|
+
is_dir = os.path.isdir(path)
|
|
140
|
+
mtime = int(os.path.getmtime(path))
|
|
141
|
+
# 从环境变量获取 task_dir
|
|
142
|
+
task_dir = os.getenv('ADAM_TASK_DIR')
|
|
143
|
+
if task_dir:
|
|
144
|
+
rel_path = os.path.relpath(path, task_dir)
|
|
145
|
+
else:
|
|
146
|
+
rel_path = path
|
|
147
|
+
except (OSError, ValueError):
|
|
148
|
+
return
|
|
149
|
+
|
|
150
|
+
# 获取当前文件列表
|
|
151
|
+
data = self._read()
|
|
152
|
+
source = self._tool_name
|
|
153
|
+
items = data["files"]["sources"].get(source, {}).get("items", [])
|
|
154
|
+
|
|
155
|
+
# 检查是否已存在相同路径,存在则移除(后面会加到末尾,保持最新追踪的在最后)
|
|
156
|
+
items = [item for item in items if item.get("path") != rel_path]
|
|
157
|
+
|
|
158
|
+
# 添加到末尾(最新的位置)
|
|
159
|
+
items.append({"path": rel_path, "is_dir": is_dir, "mtime": mtime})
|
|
160
|
+
|
|
161
|
+
# 先进先出,保留最多 max_items 条
|
|
162
|
+
if len(items) > max_items:
|
|
163
|
+
items = items[-max_items:]
|
|
164
|
+
|
|
165
|
+
# 写回
|
|
166
|
+
data["files"]["sources"][source] = {
|
|
167
|
+
"updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
|
168
|
+
"items": items
|
|
169
|
+
}
|
|
170
|
+
self._write(data)
|
|
171
|
+
|
|
172
|
+
|
|
173
|
+
# 全局实例
|
|
174
|
+
_states = _StatesManager()
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
def setState(key: str, value):
|
|
178
|
+
"""设置状态"""
|
|
179
|
+
return _states.set(key, value)
|
|
180
|
+
|
|
181
|
+
|
|
182
|
+
def getState(key: str):
|
|
183
|
+
"""获取状态"""
|
|
184
|
+
return _states.get(key)
|
|
185
|
+
|
|
186
|
+
|
|
187
|
+
def trackPath(path: str, max_items: int = 30):
|
|
188
|
+
"""
|
|
189
|
+
追踪文件/目录,自动检测文件信息并记录
|
|
190
|
+
|
|
191
|
+
Args:
|
|
192
|
+
path: 文件或目录路径
|
|
193
|
+
max_items: 最大记录数量,先进先出,默认 30
|
|
194
|
+
"""
|
|
195
|
+
return _states.trackPath(path, max_items)
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.2
|
|
2
2
|
Name: adam_community
|
|
3
|
-
Version: 1.0.
|
|
3
|
+
Version: 1.0.29
|
|
4
4
|
Summary: Adam Community Tools and Utilities
|
|
5
5
|
Home-page: https://github.com/yourusername/adam-community
|
|
6
6
|
Author: Adam Community
|
|
@@ -97,6 +97,83 @@ classes = parse_directory(Path("./"))
|
|
|
97
97
|
success, errors, zip_name = build_package(Path("./"))
|
|
98
98
|
```
|
|
99
99
|
|
|
100
|
+
### 命令执行
|
|
101
|
+
|
|
102
|
+
#### execCmd - 复杂执行
|
|
103
|
+
|
|
104
|
+
`execCmd` 是新的命令执行函数,支持异常处理、超时控制、实时输出等特性。
|
|
105
|
+
|
|
106
|
+
```python
|
|
107
|
+
from adam_community import execCmd, CmdResult
|
|
108
|
+
from subprocess import CalledProcessError, TimeoutExpired
|
|
109
|
+
|
|
110
|
+
# 基本用法
|
|
111
|
+
result = execCmd("python train.py")
|
|
112
|
+
print(result.stdout) # 标准输出
|
|
113
|
+
print(result.stderr) # 错误输出
|
|
114
|
+
print(result.returncode) # 退出码
|
|
115
|
+
print(result.duration) # 执行耗时(秒)
|
|
116
|
+
|
|
117
|
+
# 异常处理
|
|
118
|
+
try:
|
|
119
|
+
result = execCmd("python train.py", timeout=3600)
|
|
120
|
+
except TimeoutExpired as e:
|
|
121
|
+
print(f"超时: {e.output}")
|
|
122
|
+
except CalledProcessError as e:
|
|
123
|
+
print(f"失败 (exit {e.returncode}): {e.stderr}")
|
|
124
|
+
|
|
125
|
+
# 实时输出到控制台
|
|
126
|
+
result = execCmd("python train.py", echo=True)
|
|
127
|
+
|
|
128
|
+
# 自定义回调(如写日志、发送到前端)
|
|
129
|
+
result = execCmd(
|
|
130
|
+
"python train.py",
|
|
131
|
+
on_stdout=lambda line: logger.info(f"[OUT] {line}"),
|
|
132
|
+
on_stderr=lambda line: logger.error(f"[ERR] {line}"),
|
|
133
|
+
)
|
|
134
|
+
|
|
135
|
+
# 指定工作目录和环境变量
|
|
136
|
+
result = execCmd(
|
|
137
|
+
"python train.py",
|
|
138
|
+
cwd="/workspace/project",
|
|
139
|
+
env={"CUDA_VISIBLE_DEVICES": "0,1"}, # 合并到当前环境
|
|
140
|
+
)
|
|
141
|
+
```
|
|
142
|
+
|
|
143
|
+
**参数说明:**
|
|
144
|
+
|
|
145
|
+
| 参数 | 类型 | 默认值 | 说明 |
|
|
146
|
+
|------|------|--------|------|
|
|
147
|
+
| `cmd` | str | - | 要执行的命令 |
|
|
148
|
+
| `timeout` | float | None | 超时秒数,None 表示无限制 |
|
|
149
|
+
| `cwd` | str | None | 工作目录 |
|
|
150
|
+
| `env` | dict | None | 环境变量(合并到当前环境) |
|
|
151
|
+
| `shell` | str | /bin/bash | shell 路径 |
|
|
152
|
+
| `echo` | bool | False | 是否实时打印到控制台 |
|
|
153
|
+
| `on_stdout` | Callable | None | stdout 回调 `(line: str) -> None` |
|
|
154
|
+
| `on_stderr` | Callable | None | stderr 回调 `(line: str) -> None` |
|
|
155
|
+
|
|
156
|
+
**返回值 `CmdResult`:**
|
|
157
|
+
|
|
158
|
+
| 字段 | 类型 | 说明 |
|
|
159
|
+
|------|------|------|
|
|
160
|
+
| `stdout` | str | 完整标准输出 |
|
|
161
|
+
| `stderr` | str | 完整错误输出 |
|
|
162
|
+
| `returncode` | int | 退出码 |
|
|
163
|
+
| `duration` | float | 执行耗时(秒) |
|
|
164
|
+
| `command` | str | 原始命令 |
|
|
165
|
+
| `timed_out` | bool | 是否超时 |
|
|
166
|
+
|
|
167
|
+
#### runCmd - 最简执行
|
|
168
|
+
|
|
169
|
+
`runCmd` 是最简化版命令执行函数,实时输出到控制台,失败时直接退出进程。
|
|
170
|
+
|
|
171
|
+
```python
|
|
172
|
+
from adam_community import runCmd
|
|
173
|
+
|
|
174
|
+
runCmd("echo 'Hello, World!'")
|
|
175
|
+
```
|
|
176
|
+
|
|
100
177
|
### States Management(任务状态管理)
|
|
101
178
|
|
|
102
179
|
用于在任务执行过程中记录和读取状态,与服务端共享 `states.json` 文件。
|
|
@@ -1,5 +1,5 @@
|
|
|
1
|
-
adam_community/__init__.py,sha256=
|
|
2
|
-
adam_community/__version__.py,sha256=
|
|
1
|
+
adam_community/__init__.py,sha256=Sw3CfnL8xXYAzTXg6egENLO5IYksBLe_ivVNU20W514,476
|
|
2
|
+
adam_community/__version__.py,sha256=BBmi5TxVi3pzemafZwdH6Fv0y3TGRe8rQn6QWN0sJDw,23
|
|
3
3
|
adam_community/tool.py,sha256=F6jxRU3urqTfgjLIZSW-hVWyj0FpNwvY65jOODXI19w,4954
|
|
4
4
|
adam_community/util.py,sha256=SDXnWLHLfJxEJ5WBnwNq_FeHNwXxGK_F36dpI_hVkNc,17766
|
|
5
5
|
adam_community/cli/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
|
|
@@ -26,8 +26,15 @@ adam_community/cli/templates/long_description.md.j2,sha256=Rj6hcuNzEL0Sp17GQVCRJ
|
|
|
26
26
|
adam_community/cli/templates/long_description_en.md.j2,sha256=xSbahwGarXlWopZqHw7lrcv1dQuvwj2ChhZv5pJmUy4,1725
|
|
27
27
|
adam_community/cli/templates/rag_python.py.j2,sha256=YJL7-WIx-Dumt7lHuUGxl3Rbaw0kpkh8hpcCJ5lz9lA,2494
|
|
28
28
|
adam_community/cli/templates/toolbox_python.py.j2,sha256=EOnmsJUvQRrcO7K7c88kI42gMmcM0Z4ab46qwOJXbH8,4192
|
|
29
|
-
adam_community
|
|
30
|
-
adam_community
|
|
31
|
-
adam_community
|
|
32
|
-
adam_community
|
|
33
|
-
adam_community
|
|
29
|
+
adam_community/util/__init__.py,sha256=HJOq2_G5hTl9GM2MmRtj7cD60oBB-4ikU5pbOjoDTZ4,1178
|
|
30
|
+
adam_community/util/api.py,sha256=1msEvAY-cG2vYyrMH7OybVXYeAO1il8axtGSS0SSWWk,5367
|
|
31
|
+
adam_community/util/cmd.py,sha256=LghuiDaPoikvnRA1G19SgczF8M7hZ3hSTeHKt4vIZKY,4264
|
|
32
|
+
adam_community/util/config.py,sha256=On5t0kl1kD34tld5X30LwKqGsX0_WBG7vYVR7XTaT94,286
|
|
33
|
+
adam_community/util/markdown.py,sha256=iT9zX3c0mO3VLrJEhPyqe285aQoLtfxirTbHlL8Jnek,333
|
|
34
|
+
adam_community/util/rag.py,sha256=B7mju5YQmuvJddcsYXha4vkkJdvYDwVKrqmIQ3T_tmk,4147
|
|
35
|
+
adam_community/util/state.py,sha256=NRZ2mIgj_-Ykwbun1XsyfJSWv4BPvEgLaHQHbSNIwZ0,6073
|
|
36
|
+
adam_community-1.0.29.dist-info/METADATA,sha256=C6KbAAiVt4Rs9iGGbAK35NhnaFI1iXKPnpGrKxwpRmY,5891
|
|
37
|
+
adam_community-1.0.29.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
|
|
38
|
+
adam_community-1.0.29.dist-info/entry_points.txt,sha256=4I7yRkn7cHwPY8-fWQLeAvKjc24zUy8Z65VsZNs0Wos,56
|
|
39
|
+
adam_community-1.0.29.dist-info/top_level.txt,sha256=MS8jbePXKZChih9kGizNVX0I1MFZFGWBMCIW_r86qhU,15
|
|
40
|
+
adam_community-1.0.29.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|
|
File without changes
|