axi-cli 0.0.6__tar.gz → 0.0.8__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.
Files changed (29) hide show
  1. {axi_cli-0.0.6 → axi_cli-0.0.8}/CHANGELOG.md +26 -0
  2. {axi_cli-0.0.6 → axi_cli-0.0.8}/PKG-INFO +3 -2
  3. {axi_cli-0.0.6 → axi_cli-0.0.8}/README.md +2 -1
  4. {axi_cli-0.0.6 → axi_cli-0.0.8}/pyproject.toml +5 -1
  5. {axi_cli-0.0.6 → axi_cli-0.0.8}/skills/axi-creator/SKILL.md +37 -4
  6. {axi_cli-0.0.6 → axi_cli-0.0.8}/skills/axi-use/SKILL.md +4 -3
  7. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/__init__.py +2 -1
  8. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/cli.py +152 -30
  9. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/config.py +25 -1
  10. axi_cli-0.0.8/src/axi/context.py +36 -0
  11. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/daemon/client.py +29 -13
  12. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/daemon/protocol.py +6 -7
  13. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/daemon/server.py +3 -23
  14. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/mcp_serve.py +18 -3
  15. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/models.py +20 -2
  16. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/providers/mcp.py +34 -20
  17. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/providers/native.py +9 -0
  18. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/registry.py +45 -22
  19. axi_cli-0.0.8/src/axi/search/__init__.py +6 -0
  20. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/search/cache.py +5 -2
  21. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/search/hybrid.py +2 -2
  22. axi_cli-0.0.6/src/axi/search/__init__.py +0 -7
  23. {axi_cli-0.0.6 → axi_cli-0.0.8}/.gitignore +0 -0
  24. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/executor.py +0 -0
  25. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/providers/__init__.py +0 -0
  26. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/search/bm25.py +0 -0
  27. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/search/embedding.py +0 -0
  28. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/search/regex.py +0 -0
  29. {axi_cli-0.0.6 → axi_cli-0.0.8}/src/axi/search/tokenize.py +0 -0
@@ -1,5 +1,31 @@
1
1
  # Changelog
2
2
 
3
+ ## [0.0.8] - 2026-07-07
4
+
5
+ ### Added
6
+ - `axi.env(name, default=None)`:native 工具的统一变量读取入口。CLI / stdio 形态下读进程环境变量;经 `axi mcp --transport http` 导出为共享 HTTP server 时,每个客户端通过 `Axi-*` 请求头传入的变量按请求隔离、优先于进程环境变量(`Axi-Foo-Bar` → `FOO_BAR`:去前缀、`-`→`_`、转大写)
7
+ - `mcpServers` 新增 `headers` 字段:连接需认证的 HTTP MCP server 时随每个请求附带(如 `Authorization: Bearer xxx`);stdio 型凭据继续走 `env`,两种 transport 各用各的通道
8
+
9
+ ### Fixed
10
+ - MCP 连接建立中途失败(如 `initialize` 报错、远端拒绝认证)不再泄漏已打开的资源(httpx 客户端 / 流 / session):改用 `AsyncExitStack.pop_all()` 惯用法,全部成功才移交所有权,失败由 with 自动清理
11
+
12
+ ## [0.0.7] - 2026-07-04
13
+
14
+ ### Added
15
+ - `axi doctor` 自检命令:一次性检查配置、daemon、MCP server 连接、embedding、native 工具来源;配置了 MCP server 时拉起 daemon 验证连接,有问题以非零码退出并在 `issues` 里给出可执行的下一步。`native_tools.from_entry_points` 列出所有经 Python entry_points 自动注入的工具及来源模块,便于审计
16
+ - daemon 请求超时可配置:`daemon.requestTimeout`(默认 120 秒,原硬编码 30 秒)+ `AXI_REQUEST_TIMEOUT` 环境变量(优先),超时信息指向调高方式
17
+ - native 工具加载记录来源日志(模块 / entry_point + server + 工具名),便于审计自动注入
18
+
19
+ ### Fixed
20
+ - 并发多个 axi 进程同时发现 daemon 未运行会各自拉起、互相 unlink socket / 覆盖 PID:改用启动文件锁串行化,只有持锁进程 spawn,其余复查后复用
21
+ - CLI `--key value` 解析忽略工具 schema,把 string 字段的 `false` / `42` / `null` 误转成布尔 / 数字 / 空值:现在按目标字段类型决定是否 JSON 解析(string 字段保留原文,其余照常解析)
22
+
23
+ ### Changed
24
+ - **Breaking**:搜索结果移除 `score` 字段——native 与 MCP 来自两套独立索引、各自归一化,跨源分数不可比,暴露只会误导 Agent;结果按相关性降序返回,顺序本身即信号
25
+ - 搜索栈(bm25s / jieba / numpy)延迟到首次 `search` / `grep` 才 import,`describe` / `run` / `list` 不再付这笔冷启动成本;各引擎按工具集版本号独立重建
26
+ - embedding 缓存从当前目录 `.axi/cache/embeddings.json` 移到 `~/.axi/cache/<配置路径哈希>.json`,按配置隔离、不再依赖 cwd
27
+ - `allowed_types` / `server_tool_counts` 抽为公共函数,消除 cli / daemon / registry 间的重复实现
28
+
3
29
  ## [0.0.6] - 2026-07-03
4
30
 
5
31
  ### Added
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: axi-cli
3
- Version: 0.0.6
3
+ Version: 0.0.8
4
4
  Summary: Agent eXecution Interface - unified tool layer for AI Agents
5
5
  Requires-Python: >=3.12
6
6
  Requires-Dist: bm25s>=0.3.3
@@ -97,6 +97,7 @@ axi run jina/jina_search --query "hello world" --count 3
97
97
  | `axi describe <tool>` | 查看工具完整 schema |
98
98
  | `axi run <tool> --key value` | 执行工具(也支持 `-j '{...}'` 传 JSON) |
99
99
  | `axi mcp` | 将 axi 工具导出为 MCP server(stdio / HTTP) |
100
+ | `axi doctor` | 自检:配置、daemon、MCP 连接、embedding、native 工具来源;有问题非零退出 |
100
101
  | `axi daemon start` | 手动启动 daemon(通常无需手动,CLI 会自动拉起) |
101
102
  | `axi daemon status` | 查看 daemon 状态(PID、运行时长、空闲时长、工具数量等) |
102
103
  | `axi daemon stop` | 手动停止 daemon |
@@ -169,7 +170,7 @@ axi 提供两种搜索方式:
169
170
  - **`axi search`** — 混合搜索,默认 BM25 关键词搜索(bm25s + jieba 分词,支持中英文)。配置 `search.embedding` 后启用 BM25 + Embedding 混合搜索(RRF 融合排序,分数归一化 0-1)
170
171
  - **`axi grep`** — 正则表达式搜索,按工具名和描述匹配
171
172
 
172
- 搜索结果包含 `score` 字段,便于 Agent 判断相关性。
173
+ 结果按相关性排序返回,顺序本身即信号(不含分数:native MCP 来自两套索引、分数不可比)。
173
174
 
174
175
  ## Daemon 模式
175
176
 
@@ -80,6 +80,7 @@ axi run jina/jina_search --query "hello world" --count 3
80
80
  | `axi describe <tool>` | 查看工具完整 schema |
81
81
  | `axi run <tool> --key value` | 执行工具(也支持 `-j '{...}'` 传 JSON) |
82
82
  | `axi mcp` | 将 axi 工具导出为 MCP server(stdio / HTTP) |
83
+ | `axi doctor` | 自检:配置、daemon、MCP 连接、embedding、native 工具来源;有问题非零退出 |
83
84
  | `axi daemon start` | 手动启动 daemon(通常无需手动,CLI 会自动拉起) |
84
85
  | `axi daemon status` | 查看 daemon 状态(PID、运行时长、空闲时长、工具数量等) |
85
86
  | `axi daemon stop` | 手动停止 daemon |
@@ -152,7 +153,7 @@ axi 提供两种搜索方式:
152
153
  - **`axi search`** — 混合搜索,默认 BM25 关键词搜索(bm25s + jieba 分词,支持中英文)。配置 `search.embedding` 后启用 BM25 + Embedding 混合搜索(RRF 融合排序,分数归一化 0-1)
153
154
  - **`axi grep`** — 正则表达式搜索,按工具名和描述匹配
154
155
 
155
- 搜索结果包含 `score` 字段,便于 Agent 判断相关性。
156
+ 结果按相关性排序返回,顺序本身即信号(不含分数:native MCP 来自两套索引、分数不可比)。
156
157
 
157
158
  ## Daemon 模式
158
159
 
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "axi-cli"
3
- version = "0.0.6"
3
+ version = "0.0.8"
4
4
  description = "Agent eXecution Interface - unified tool layer for AI Agents"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.12"
@@ -45,3 +45,7 @@ dev = [
45
45
  "pytest-asyncio>=1.3.0",
46
46
  "ruff>=0.15.8",
47
47
  ]
48
+
49
+ [tool.pytest.ini_options]
50
+ # 只从 tests/ 收集,隔离本地 scratch 的 test/(已 gitignore,不随包分发)
51
+ testpaths = ["tests"]
@@ -64,6 +64,21 @@ def get_can_data(
64
64
  - 复杂输入用 `BaseModel` 嵌套;axi 会把它展开成 object schema
65
65
  - **填 `output_example`**:MCP 协议没有 output schema,原生工具是 axi 独有的优势——Agent 能在 `axi describe` 里直接看到返回形状,不用试跑一遍
66
66
 
67
+ ### 凭据与配置变量:用 `axi.env`,别用 `os.environ`
68
+
69
+ 工具里需要 API Key 等变量时,一律通过 `axi.env(name, default=None)` 读取:
70
+
71
+ ```python
72
+ import axi
73
+
74
+ @tool(description="调用上游 API")
75
+ def call_api(query: str) -> dict:
76
+ key = axi.env("API_KEY") # CLI/stdio 下读环境变量;HTTP 导出时读客户端 Axi-* header
77
+ ...
78
+ ```
79
+
80
+ CLI 和 stdio 形态下它就是环境变量;但当工具经 `axi mcp --transport http` 导出为共享 HTTP server 时,每个客户端可通过 `Axi-*` 请求头传自己的变量(见下文"导出为 MCP server"),`axi.env` 会优先取当前请求的值。直接 `os.environ` 的工具在 HTTP 形态下拿不到客户端变量。
81
+
67
82
  ### 注册:两种方式,二选一
68
83
 
69
84
  **方式 1:本项目内用 → `axi.json` 的 `nativeTools`**
@@ -104,13 +119,14 @@ smartlink = "smartlink_axi.tools"
104
119
  "env": {"API_KEY": "xxx"}
105
120
  },
106
121
  "remote": {
107
- "url": "http://localhost:8000/mcp" // 或 HTTP streaming
122
+ "url": "http://localhost:8000/mcp", // 或 HTTP streaming
123
+ "headers": {"Authorization": "Bearer xxx"} // 可选:远端要求认证时随每个请求附带
108
124
  }
109
125
  }
110
126
  }
111
127
  ```
112
128
 
113
- `command` / `url` 二选一。保存后直接 `axi list` 就能看到新工具——daemon 会按需自动拉起连接。
129
+ `command` / `url` 二选一;凭据传递按 transport 选字段——stdio 用 `env`(注入子进程环境),HTTP 用 `headers`(附到每个请求)。保存后直接 `axi list` 就能看到新工具——daemon 会按需自动拉起连接。
114
130
 
115
131
  ---
116
132
 
@@ -125,7 +141,7 @@ axi describe <server>/<tool> # input_schema 字段、类型、req
125
141
  axi run <server>/<tool> --key value # 跑通,返回 {"status":"success","data":...}
126
142
  ```
127
143
 
128
- **搜不到比跑不通更常见**——description 太短、只写了类名或方法名都会导致 BM25/Embedding 召回差。看着 `axi search` 返回列表里有没有你的工具,score 是不是高位。
144
+ **搜不到比跑不通更常见**——description 太短、只写了类名或方法名都会导致 BM25/Embedding 召回差。用几个真实查询跑 `axi search`,看你的工具有没有出现、排得够不够靠前。
129
145
 
130
146
  ---
131
147
 
@@ -155,6 +171,23 @@ axi mcp --server a,b --transport http --port 8321 # HTTP 长驻服务
155
171
  }
156
172
  ```
157
173
 
174
+ ### HTTP 形态下给每个客户端传变量:`Axi-*` header
175
+
176
+ stdio 是一客户端一进程,变量走 `env` 即可;HTTP 是共享进程,进程环境变量在启动时定死,**每个客户端的变量改走请求头**。约定:`Axi-<名字>` 映射为变量名(去前缀、`-`→`_`、转大写),工具内用 `axi.env` 读取,请求级值优先于进程环境变量:
177
+
178
+ ```json
179
+ {
180
+ "mcpServers": {
181
+ "my-tools": {
182
+ "url": "http://your-host:8321/mcp",
183
+ "headers": { "Axi-API-KEY": "sk_xxx", "Axi-REGION": "cn" }
184
+ }
185
+ }
186
+ }
187
+ ```
188
+
189
+ 工具里 `axi.env("API_KEY")` 即得 `sk_xxx`,不同客户端互不可见。注意:这解决的是"客户端带变量进来",不是服务端准入控制——公网暴露时另需网关层鉴权。
190
+
158
191
  ---
159
192
 
160
193
  ## 反模式
@@ -163,7 +196,7 @@ axi mcp --server a,b --transport http --port 8321 # HTTP 长驻服务
163
196
  - 所有参数都写 `str` — 失去类型约束和枚举提示,Agent 要靠 describe 里的自然语言猜
164
197
  - `description` 只写一个词("搜索"、"查询")— BM25 召回差、Embedding 语义稀薄
165
198
  - 返回不稳定结构(同一工具有时返 list 有时返 dict)— 下游 Agent 没法写聚合逻辑
166
- - 把敏感配置(API Key)硬编码在 `@tool` 函数里 — 走环境变量或 `mcpServers[...].env`
199
+ - 把敏感配置(API Key)硬编码在 `@tool` 函数里 — `axi.env` 读取(MCP server 的凭据走 `mcpServers[...].env` / `headers`)
167
200
  - entry_points 的 value 写成文件路径 `"pkg/tools.py"` — 必须是可 import 的模块路径 `"pkg.tools"`
168
201
 
169
202
  ---
@@ -1,6 +1,6 @@
1
1
  ---
2
2
  name: axi-use
3
- description: 通过 `axi` CLI(或在 Python 中使用 `from axi import tool`)发现并调用工具。当任务需要执行外部工具(当前拥有plane任务管理工具)请使用此skill。核心工作流程为:`axi search <query>` → `axi describe <tool>` → `axi run <tool> --key value`;当任务涉及多次调用之间的循环、聚合或后处理时,请优先使用 Python 进行处理。
3
+ description: 通过 `axi` CLI(或在 Python 中使用 `from axi import tool`)发现并调用工具。当任务需要执行外部工具(当前拥有<changeme>)请使用此skill。核心工作流程为:`axi search <query>` → `axi describe <tool>` → `axi run <tool> --key value`;当任务涉及多次调用之间的循环、聚合或后处理时,请优先使用 Python 进行处理。
4
4
  ---
5
5
 
6
6
  # axi — Agent 使用手册
@@ -28,8 +28,9 @@ axi 是当前环境里的工具调用层。有工具需求时,**先 `axi searc
28
28
 
29
29
  ```bash
30
30
  axi search "web" --top-k 3
31
- # [{"name":"jina/jina_search","description":"Search the web. ...","source":"mcp","score":1.0},
32
- # {"name":"jina/jina_reader","description":"Read and extract content from web page.","source":"mcp","score":0.92}]
31
+ # [{"name":"jina/jina_search","description":"Search the web. ...","source":"mcp"},
32
+ # {"name":"jina/jina_reader","description":"Read and extract content from web page.","source":"mcp"}]
33
+ # 结果按相关性排序(越靠前越相关),不含分数
33
34
 
34
35
  axi grep "^jina/" --limit 5
35
36
  # 精确列出 jina server 下所有工具
@@ -4,9 +4,10 @@ import logging
4
4
  from typing import Any, Callable
5
5
 
6
6
  from axi.cli import get_executor, get_registry
7
+ from axi.context import env
7
8
  from axi.providers.native import register_tool
8
9
 
9
- __all__ = ["tool"]
10
+ __all__ = ["tool", "env"]
10
11
 
11
12
  # 包级 logger:各子模块通过 logging.getLogger(__name__) 自动继承
12
13
  _logger = logging.getLogger("axi")
@@ -2,13 +2,12 @@
2
2
 
3
3
  import json
4
4
  import logging
5
- from collections import Counter
6
5
  from enum import Enum
7
6
 
8
7
  import typer
9
8
  from pydantic import BaseModel
10
9
 
11
- from axi.config import app_config
10
+ from axi.config import CONFIG_PATH, app_config
12
11
  from axi.daemon.client import (
13
12
  LOG_PATH,
14
13
  daemon_request,
@@ -18,7 +17,7 @@ from axi.daemon.client import (
18
17
  )
19
18
  from axi.daemon.protocol import DaemonRequest
20
19
  from axi.executor import Executor
21
- from axi.models import RunResult, SearchResult
20
+ from axi.models import RunResult, SearchResult, allowed_types
22
21
  from axi.registry import (
23
22
  AmbiguousToolError,
24
23
  Registry,
@@ -128,16 +127,10 @@ def daemon_status() -> None:
128
127
  _output_json({"status": "error", "error": resp.error})
129
128
  raise typer.Exit(code=1)
130
129
 
131
- data = resp.data
132
- # 原生工具按 server 统计
133
- native_server_tools = dict(
134
- Counter(meta.server or "unknown" for meta in _registry.list_all())
135
- )
136
-
137
130
  result = {
138
131
  "status": "running",
139
- **data,
140
- "native_tools": native_server_tools,
132
+ **resp.data,
133
+ "native_tools": _registry.server_tool_counts(),
141
134
  }
142
135
  _output_json(result)
143
136
 
@@ -159,9 +152,9 @@ def _search_and_merge(
159
152
  elif resp.status == "error":
160
153
  logger.warning("Daemon search failed: %s", resp.error)
161
154
 
162
- combined = [r.model_dump(exclude_none=True) for r in local_results] + (
163
- mcp_results or []
164
- )
155
+ # native(本地 registry)与 MCP(daemon)是两套独立索引,分数不可比,
156
+ # 无法跨源真正排序;各自已取回 top_k,直接并列拼接(本地在前)。
157
+ combined = [r.model_dump() for r in local_results] + mcp_results
165
158
  typer.echo(json.dumps(combined, ensure_ascii=False))
166
159
 
167
160
 
@@ -286,6 +279,14 @@ def _resolve_tool(name: str) -> dict:
286
279
  return {"error": resp.error or f"Tool not found: {name}"}
287
280
 
288
281
 
282
+ def _daemon_input_schema(tool_name: str) -> dict:
283
+ """从 daemon 取 MCP 工具的 input_schema;取不到返回 {}(解析退回 JSON 猜类型)。"""
284
+ resp = daemon_request(DaemonRequest(method="describe", tool_name=tool_name))
285
+ if resp.status == "success" and isinstance(resp.data, dict):
286
+ return resp.data.get("input_schema") or {}
287
+ return {}
288
+
289
+
289
290
  @app.command()
290
291
  def describe(
291
292
  tool_name: str = typer.Argument(help="工具完整名称(逗号分隔多个)"),
@@ -319,25 +320,32 @@ def run(
319
320
 
320
321
  json_str, args = _extract_option(args, "--json", "-j")
321
322
 
323
+ # 先解析工具一次:本地命中 → 原生(进程内执行,schema 就在 meta 上);
324
+ # 未命中 → 走 daemon(MCP)。避免解析参数和执行各 resolve 一遍。
325
+ try:
326
+ meta = _registry.resolve(tool_name)
327
+ except AmbiguousToolError as e:
328
+ _output_json({"error": str(e)})
329
+ raise typer.Exit(code=1)
330
+ except ToolNotFoundError:
331
+ meta = None
332
+
322
333
  if json_str:
323
334
  try:
324
335
  parsed = json.loads(json_str)
325
336
  except json.JSONDecodeError as e:
326
337
  _output_json(RunResult.fail(f"Invalid JSON argument: {e}"))
327
338
  raise typer.Exit(code=1)
339
+ elif not args:
340
+ parsed = {}
328
341
  else:
329
- parsed = _parse_params(args)
342
+ # 按目标 schema 解析 --key value(string 字段保留原文);MCP schema 问 daemon
343
+ schema = meta.input_schema if meta else _daemon_input_schema(tool_name)
344
+ parsed = _parse_params(args, schema)
330
345
 
331
- try:
332
- meta = _registry.resolve(tool_name)
333
- result = _executor.run(meta.full_name, parsed)
334
- _output_json(result)
346
+ if meta is not None:
347
+ _output_json(_executor.run(meta.full_name, parsed))
335
348
  return
336
- except AmbiguousToolError as e:
337
- _output_json({"error": str(e)})
338
- raise typer.Exit(code=1)
339
- except ToolNotFoundError:
340
- pass # 本地未找到,继续尝试 daemon
341
349
 
342
350
  resp = daemon_request(
343
351
  DaemonRequest(method="call_tool", tool_name=tool_name, params=parsed)
@@ -380,6 +388,98 @@ def mcp_command(
380
388
  )
381
389
 
382
390
 
391
+ @app.command()
392
+ def doctor() -> None:
393
+ """自检:配置、daemon、MCP 连接、embedding、native 工具来源。
394
+
395
+ 有问题时以非零码退出并在 ``issues`` 里给出可执行的下一步。
396
+ 配置了 MCP server 时会拉起 daemon 以验证连接。
397
+ """
398
+ import importlib.metadata
399
+
400
+ from axi.providers.native import NATIVE_TOOLS_ENTRY_POINT_GROUP
401
+
402
+ issues: list[str] = []
403
+
404
+ # native 工具及其来源(entry_points 自动发现 = 任何已装包都可能注入工具,需可见)
405
+ try:
406
+ eps = list(
407
+ importlib.metadata.entry_points(group=NATIVE_TOOLS_ENTRY_POINT_GROUP)
408
+ )
409
+ except Exception:
410
+ eps = []
411
+ server_counts = _registry.server_tool_counts()
412
+ native = {
413
+ "total": sum(server_counts.values()),
414
+ "servers": server_counts,
415
+ "from_config": [e.module for e in app_config.native_tools],
416
+ "from_entry_points": [{"name": ep.name, "value": ep.value} for ep in eps],
417
+ }
418
+
419
+ # embedding
420
+ emb = app_config.search.embedding
421
+ if emb.provider:
422
+ embedding = {
423
+ "provider": emb.provider,
424
+ "api_key": "present" if emb.api_key else "missing",
425
+ "model": emb.model,
426
+ }
427
+ if not emb.api_key:
428
+ env_name = "JINA_API_KEY" if emb.provider == "jina" else "OPENAI_API_KEY"
429
+ issues.append(
430
+ f"Embedding provider '{emb.provider}' set but no API key. Put it in "
431
+ f"axi.json search.embedding.apiKey or the {env_name} env var."
432
+ )
433
+ else:
434
+ embedding = {"provider": None, "note": "BM25-only, no semantic search"}
435
+
436
+ # daemon + MCP 连接(仅在配置了 MCP server 时拉起 daemon 验证)
437
+ configured = set(app_config.mcp_servers)
438
+ mcp_servers: list[dict] = []
439
+ if configured:
440
+ resp = daemon_request(DaemonRequest(method="status"))
441
+ if resp.status == "success":
442
+ connected = resp.data.get("server_tools", {})
443
+ daemon = {
444
+ "running": True,
445
+ "pid": resp.data.get("pid"),
446
+ "uptime_seconds": resp.data.get("uptime_seconds"),
447
+ }
448
+ for name in sorted(configured):
449
+ if name in connected:
450
+ mcp_servers.append(
451
+ {
452
+ "server": name,
453
+ "status": "connected",
454
+ "tools": connected[name],
455
+ }
456
+ )
457
+ else:
458
+ mcp_servers.append({"server": name, "status": "not_connected"})
459
+ issues.append(
460
+ f"MCP server '{name}' configured but not connected. "
461
+ f"Check daemon log: {LOG_PATH}"
462
+ )
463
+ else:
464
+ daemon = {"running": False, "error": resp.error}
465
+ issues.append(f"Daemon unavailable: {resp.error}. Check log: {LOG_PATH}")
466
+ else:
467
+ daemon = {"running": is_daemon_running()}
468
+
469
+ report = {
470
+ "ok": not issues,
471
+ "config": {"path": str(CONFIG_PATH), "exists": CONFIG_PATH.exists()},
472
+ "daemon": daemon,
473
+ "native_tools": native,
474
+ "mcp_servers": mcp_servers,
475
+ "embedding": embedding,
476
+ "issues": issues,
477
+ }
478
+ _output_json(report)
479
+ if issues:
480
+ raise typer.Exit(code=1)
481
+
482
+
383
483
  # ── 参数解析辅助函数 ──────────────────────────────────────────────
384
484
 
385
485
 
@@ -397,7 +497,33 @@ def _extract_option(args: list[str], *names: str) -> tuple[str | None, list[str]
397
497
  return value, remaining
398
498
 
399
499
 
400
- def _parse_params(params: list[str]) -> dict:
500
+ def _schema_field_types(schema: dict, key: str) -> set[str]:
501
+ """取字段在 schema 里声明的类型集合;字段不存在或无 schema 返回空集。"""
502
+ props = schema.get("properties")
503
+ if not isinstance(props, dict):
504
+ return set()
505
+ prop = props.get(key)
506
+ return allowed_types(prop) if isinstance(prop, dict) else set()
507
+
508
+
509
+ def _coerce_cli_value(key: str, value: str, schema: dict) -> object:
510
+ """按目标字段类型决定是否 JSON 解析。
511
+
512
+ 字段声明为 string → 原样保留,避免 ``true``/``42``/``null`` 被解析成
513
+ 非字符串类型(``--title false`` 应是字符串 "false")。其余(含未知字段、
514
+ 无 schema)尝试 JSON 解析,让 ``--count 42``→42、``--flag true``→True、
515
+ ``--data {...}``→dict 正常工作,失败落回原字符串。
516
+ """
517
+ if "string" in _schema_field_types(schema, key):
518
+ return value
519
+ try:
520
+ return json.loads(value)
521
+ except (json.JSONDecodeError, ValueError):
522
+ return value
523
+
524
+
525
+ def _parse_params(params: list[str], schema: dict | None = None) -> dict:
526
+ schema = schema or {}
401
527
  parsed: dict = {}
402
528
  i = 0
403
529
  while i < len(params):
@@ -405,11 +531,7 @@ def _parse_params(params: list[str]) -> dict:
405
531
  if arg.startswith("--"):
406
532
  key = arg[2:]
407
533
  if i + 1 < len(params) and not params[i + 1].startswith("--"):
408
- value = params[i + 1]
409
- try:
410
- parsed[key] = json.loads(value)
411
- except (json.JSONDecodeError, ValueError):
412
- parsed[key] = value
534
+ parsed[key] = _coerce_cli_value(key, params[i + 1], schema)
413
535
  i += 2
414
536
  else:
415
537
  parsed[key] = True
@@ -1,5 +1,6 @@
1
1
  """axi 配置中心:Pydantic 模型化,统一加载 axi.json,全局共享。"""
2
2
 
3
+ import hashlib
3
4
  import json
4
5
  import logging
5
6
  import os
@@ -10,9 +11,13 @@ from pydantic import BaseModel, Field, model_validator
10
11
 
11
12
  logger = logging.getLogger(__name__)
12
13
 
14
+ AXI_HOME = Path("~/.axi").expanduser()
13
15
  CONFIG_PATH = (
14
- Path(os.environ.get("AXI_CONFIG", "~/.axi/axi.json")).expanduser().resolve()
16
+ Path(os.environ.get("AXI_CONFIG", AXI_HOME / "axi.json")).expanduser().resolve()
15
17
  )
18
+ # 每份配置的隔离键:daemon socket/pid/log 与 embedding 缓存都按它分目录,
19
+ # 多项目(多份 AXI_CONFIG)互不干扰。
20
+ CONFIG_HASH = hashlib.sha256(str(CONFIG_PATH).encode()).hexdigest()[:12]
16
21
 
17
22
 
18
23
  # ── 子配置模型 ──────────────────────────────────────────────
@@ -77,6 +82,24 @@ class DaemonConfig(BaseModel):
77
82
  idle_timeout_minutes: int = Field(
78
83
  default=30, alias="idleTimeoutMinutes", description="空闲自动关闭(分钟)"
79
84
  )
85
+ request_timeout: float = Field(
86
+ default=120,
87
+ alias="requestTimeout",
88
+ description="CLI 等待 daemon 单个请求的超时(秒),含工具调用",
89
+ )
90
+
91
+ @model_validator(mode="before")
92
+ @classmethod
93
+ def override_with_env(cls, values: dict) -> dict:
94
+ # 环境变量 AXI_REQUEST_TIMEOUT 优先于 axi.json(长工具调用可临时调高)
95
+ if isinstance(values, dict):
96
+ env = os.environ.get("AXI_REQUEST_TIMEOUT")
97
+ if env:
98
+ try:
99
+ values["requestTimeout"] = float(env)
100
+ except ValueError:
101
+ logger.warning("Ignoring invalid AXI_REQUEST_TIMEOUT=%r", env)
102
+ return values
80
103
 
81
104
 
82
105
  class MCPServerConfig(BaseModel):
@@ -86,6 +109,7 @@ class MCPServerConfig(BaseModel):
86
109
  args: list[str] = Field(default_factory=list)
87
110
  env: dict[str, str] | None = None
88
111
  url: str | None = None
112
+ headers: dict[str, str] | None = None
89
113
 
90
114
 
91
115
  class NativeToolEntry(BaseModel):
@@ -0,0 +1,36 @@
1
+ """请求级变量:native 工具的统一变量读取入口。
2
+
3
+ stdio / CLI 形态下变量来自进程环境变量;HTTP MCP 形态下 server 是共享进程,
4
+ 每个客户端的变量通过 ``Axi-*`` header 随请求传入,按请求隔离。
5
+ """
6
+
7
+ import os
8
+ from collections.abc import Mapping
9
+ from contextvars import ContextVar
10
+
11
+ _request_env: ContextVar[dict[str, str]] = ContextVar("axi_request_env", default={})
12
+
13
+ _ENV_HEADER_PREFIX = "axi-"
14
+
15
+
16
+ def set_request_env(headers: Mapping[str, str] | None) -> None:
17
+ """从 ``Axi-*`` header 提取当前请求的变量:``Axi-Foo-Bar`` → ``FOO_BAR``。
18
+
19
+ 由 mcp_serve 在每次 native 工具调用前调用;``None``(stdio 无请求)置空。
20
+ ASGI 规范保证 header 名已小写。
21
+ """
22
+ _request_env.set(
23
+ {
24
+ name.removeprefix(_ENV_HEADER_PREFIX).replace("-", "_").upper(): value
25
+ for name, value in headers.items()
26
+ if name.startswith(_ENV_HEADER_PREFIX)
27
+ }
28
+ if headers
29
+ else {}
30
+ )
31
+
32
+
33
+ def env(name: str, default: str | None = None) -> str | None:
34
+ """读取变量:请求级(HTTP header)优先,其次进程环境变量。"""
35
+ value = _request_env.get().get(name)
36
+ return value if value is not None else os.environ.get(name, default)
@@ -1,14 +1,16 @@
1
1
  """daemon 客户端:CLI 侧通过 Unix socket 与 daemon 通信。"""
2
2
 
3
3
  import asyncio
4
+ import fcntl
4
5
  import logging
5
6
  import os
6
7
  import subprocess
7
8
  import sys
8
9
  import time
9
10
 
10
- from axi.config import CONFIG_PATH
11
+ from axi.config import CONFIG_PATH, app_config
11
12
  from axi.daemon.protocol import (
13
+ LOCK_PATH,
12
14
  LOG_PATH,
13
15
  SOCKET_DIR,
14
16
  SOCKET_PATH,
@@ -21,7 +23,6 @@ logger = logging.getLogger(__name__)
21
23
 
22
24
  _DAEMON_START_POLL_RETRIES = 30
23
25
  _DAEMON_START_POLL_INTERVAL = 0.1 # seconds
24
- _DAEMON_REQUEST_TIMEOUT = 30 # seconds
25
26
 
26
27
 
27
28
  def is_daemon_running() -> bool:
@@ -38,13 +39,8 @@ def is_daemon_running() -> bool:
38
39
  return False
39
40
 
40
41
 
41
- def ensure_daemon() -> bool:
42
- """确保 daemon 已启动。未运行时自动启动,返回是否就绪。"""
43
- if is_daemon_running():
44
- return True
45
-
46
- os.makedirs(SOCKET_DIR, exist_ok=True)
47
-
42
+ def _spawn_and_wait() -> bool:
43
+ """拉起 daemon 子进程并轮询直到就绪。调用方需持有启动锁。"""
48
44
  with open(LOG_PATH, "a") as log_file:
49
45
  subprocess.Popen(
50
46
  [sys.executable, "-m", "axi.daemon.server"],
@@ -62,6 +58,26 @@ def ensure_daemon() -> bool:
62
58
  return False
63
59
 
64
60
 
61
+ def ensure_daemon() -> bool:
62
+ """确保 daemon 已启动。未运行时自动启动,返回是否就绪。
63
+
64
+ 并行的多个 axi 进程可能同时发现 daemon 未运行。用文件锁串行化启动:
65
+ 只有拿到锁的进程 spawn,其余进程阻塞在锁上,拿到锁后复查发现已就绪即返回。
66
+ 否则会多进程各拉起一个 daemon、互相 unlink socket、覆盖 PID 文件。
67
+ """
68
+ if is_daemon_running():
69
+ return True
70
+
71
+ os.makedirs(SOCKET_DIR, exist_ok=True)
72
+
73
+ with open(LOCK_PATH, "w") as lock:
74
+ fcntl.flock(lock, fcntl.LOCK_EX)
75
+ # 锁内复查:可能刚才另一个进程已经把 daemon 拉起来了
76
+ if is_daemon_running():
77
+ return True
78
+ return _spawn_and_wait()
79
+
80
+
65
81
  def send_request(req: DaemonRequest) -> DaemonResponse:
66
82
  """向 daemon 发送请求并获取响应。"""
67
83
  return asyncio.run(_send(req))
@@ -86,16 +102,16 @@ async def _send(req: DaemonRequest) -> DaemonResponse:
86
102
  writer.write(req.model_dump_json().encode() + b"\n")
87
103
  await writer.drain()
88
104
 
89
- line = await asyncio.wait_for(
90
- reader.readline(), timeout=_DAEMON_REQUEST_TIMEOUT
91
- )
105
+ timeout = app_config.daemon.request_timeout
106
+ line = await asyncio.wait_for(reader.readline(), timeout=timeout)
92
107
  if not line:
93
108
  return DaemonResponse.fail("Daemon connection closed unexpectedly")
94
109
 
95
110
  return DaemonResponse.model_validate_json(line)
96
111
  except asyncio.TimeoutError:
97
112
  return DaemonResponse.fail(
98
- f"Daemon request timed out after {_DAEMON_REQUEST_TIMEOUT}s"
113
+ f"Daemon request timed out after {timeout}s. Raise it via "
114
+ "AXI_REQUEST_TIMEOUT or axi.json daemon.requestTimeout."
99
115
  )
100
116
  finally:
101
117
  writer.close()
@@ -1,20 +1,19 @@
1
1
  """daemon 通信协议:JSON 行格式的请求/响应。"""
2
2
 
3
- import hashlib
4
3
  import os
5
4
  from typing import Any, Literal
6
5
 
7
6
  from pydantic import BaseModel, Field
8
7
 
9
- from axi.config import CONFIG_PATH
8
+ from axi.config import AXI_HOME, CONFIG_HASH
10
9
  from axi.models import ResultEnvelope
11
10
 
12
11
  # 每份配置对应一个独立 daemon,socket/pid/log 按配置路径 hash 隔离
13
- SOCKET_DIR = os.path.expanduser("~/.axi/daemons")
14
- _CONFIG_HASH = hashlib.sha256(str(CONFIG_PATH).encode()).hexdigest()[:12]
15
- SOCKET_PATH = os.path.join(SOCKET_DIR, f"{_CONFIG_HASH}.sock")
16
- PID_PATH = os.path.join(SOCKET_DIR, f"{_CONFIG_HASH}.pid")
17
- LOG_PATH = os.path.join(SOCKET_DIR, f"{_CONFIG_HASH}.log")
12
+ SOCKET_DIR = str(AXI_HOME / "daemons")
13
+ SOCKET_PATH = os.path.join(SOCKET_DIR, f"{CONFIG_HASH}.sock")
14
+ PID_PATH = os.path.join(SOCKET_DIR, f"{CONFIG_HASH}.pid")
15
+ LOG_PATH = os.path.join(SOCKET_DIR, f"{CONFIG_HASH}.log")
16
+ LOCK_PATH = os.path.join(SOCKET_DIR, f"{CONFIG_HASH}.lock")
18
17
 
19
18
 
20
19
  class DaemonRequest(BaseModel):
@@ -5,7 +5,6 @@ import logging
5
5
  import os
6
6
  import signal
7
7
  import time
8
- from collections import Counter
9
8
  from typing import Any
10
9
 
11
10
  import jsonschema
@@ -19,7 +18,7 @@ from axi.daemon.protocol import (
19
18
  DaemonResponse,
20
19
  DaemonStatus,
21
20
  )
22
- from axi.models import ToolSource
21
+ from axi.models import ToolSource, allowed_types
23
22
  from axi.providers.mcp import MCPProvider
24
23
  from axi.registry import Registry, ToolResolveError
25
24
 
@@ -28,20 +27,6 @@ logger = logging.getLogger(__name__)
28
27
  _IDLE_EXEMPT_METHODS = frozenset({"status", "shutdown"})
29
28
 
30
29
 
31
- def _allowed_types(prop: dict) -> set[str]:
32
- """提取属性 schema 允许的原始类型,覆盖 type 为字符串/列表及 anyOf 分支。"""
33
- types: set[str] = set()
34
- t = prop.get("type")
35
- if isinstance(t, str):
36
- types.add(t)
37
- elif isinstance(t, list):
38
- types.update(x for x in t if isinstance(x, str))
39
- for branch in prop.get("anyOf", []):
40
- if isinstance(branch, dict):
41
- types |= _allowed_types(branch)
42
- return types
43
-
44
-
45
30
  def _coerce_to_schema(params: dict, schema: dict) -> dict:
46
31
  """按 schema 顶层类型做轻量转换,弥合 CLI 参数解析的类型猜测与严格校验的间隙。
47
32
 
@@ -57,7 +42,7 @@ def _coerce_to_schema(params: dict, schema: dict) -> dict:
57
42
  prop = props.get(key)
58
43
  if not isinstance(prop, dict):
59
44
  continue
60
- allowed = _allowed_types(prop)
45
+ allowed = allowed_types(prop)
61
46
  if (
62
47
  "string" in allowed
63
48
  and not ({"integer", "number"} & allowed)
@@ -288,11 +273,6 @@ class DaemonServer:
288
273
  idle = now - self._last_activity
289
274
  idle_remaining = max(0.0, self._idle_timeout - idle)
290
275
 
291
- # 按 server 统计工具数量
292
- server_tools = dict(
293
- Counter(t.server or "unknown" for t in self.registry.list_all())
294
- )
295
-
296
276
  status = DaemonStatus(
297
277
  pid=os.getpid(),
298
278
  config_path=str(CONFIG_PATH),
@@ -300,7 +280,7 @@ class DaemonServer:
300
280
  idle_seconds=int(idle),
301
281
  idle_timeout_seconds=int(self._idle_timeout),
302
282
  idle_remaining_seconds=int(idle_remaining),
303
- server_tools=server_tools,
283
+ server_tools=self.registry.server_tool_counts(),
304
284
  )
305
285
  return DaemonResponse.success(status.model_dump())
306
286
 
@@ -18,6 +18,7 @@ import mcp.types as types
18
18
  from mcp.server.lowlevel import Server
19
19
 
20
20
  from axi.config import app_config
21
+ from axi.context import set_request_env
21
22
  from axi.daemon.client import daemon_request
22
23
  from axi.daemon.protocol import LOG_PATH, DaemonRequest
23
24
  from axi.executor import Executor
@@ -87,10 +88,18 @@ def build_registry(
87
88
 
88
89
 
89
90
  async def _execute(
90
- registry: Registry, executor: Executor, name: str, params: dict[str, Any]
91
+ server: Server,
92
+ registry: Registry,
93
+ executor: Executor,
94
+ name: str,
95
+ params: dict[str, Any],
91
96
  ) -> RunResult:
92
97
  """按来源分发执行:native 进程内,MCP 转发 daemon。
93
98
 
99
+ native 执行前把当前 HTTP 请求的 Axi-* header 写入请求级变量
100
+ (stdio 无请求对象,置空);MCP 工具在 daemon 进程执行,读不到
101
+ contextvar,不捕获。
102
+
94
103
  Executor.run 与 send_request 内部都会 asyncio.run,
95
104
  必须放到 worker 线程执行以避免嵌套事件循环。
96
105
  """
@@ -100,6 +109,8 @@ async def _execute(
100
109
  return RunResult.fail(str(e))
101
110
 
102
111
  if meta.source == ToolSource.NATIVE:
112
+ request = server.request_context.request
113
+ set_request_env(request.headers if request else None)
103
114
  return await asyncio.to_thread(executor.run, meta.full_name, params)
104
115
 
105
116
  # daemon_request 内含 ensure_daemon:daemon idle 自杀后长驻 serve 进程可自动重启它
@@ -174,7 +185,7 @@ def _build_flat_server(registry: Registry, executor: Executor) -> Server:
174
185
  raise ValueError(
175
186
  f"Unknown tool: {name}. Call tools/list to see the available tools."
176
187
  )
177
- result = await _execute(registry, executor, full_name, arguments)
188
+ result = await _execute(server, registry, executor, full_name, arguments)
178
189
  if result.status == "error":
179
190
  raise RuntimeError(result.error)
180
191
  return _text(result.data)
@@ -269,7 +280,11 @@ def _build_meta_server(registry: Registry, executor: Executor) -> Server:
269
280
  return _text(out[0] if len(out) == 1 else out)
270
281
  if name == "run":
271
282
  result = await _execute(
272
- registry, executor, arguments["name"], arguments.get("params") or {}
283
+ server,
284
+ registry,
285
+ executor,
286
+ arguments["name"],
287
+ arguments.get("params") or {},
273
288
  )
274
289
  return _text(result.model_dump(exclude_none=True))
275
290
  raise ValueError(
@@ -6,6 +6,20 @@ from typing import Any, Literal, Self
6
6
  from pydantic import BaseModel, Field, model_validator
7
7
 
8
8
 
9
+ def allowed_types(prop: dict) -> set[str]:
10
+ """提取属性 schema 允许的原始类型,覆盖 type 为字符串/列表及 anyOf 分支(递归)。"""
11
+ types: set[str] = set()
12
+ t = prop.get("type")
13
+ if isinstance(t, str):
14
+ types.add(t)
15
+ elif isinstance(t, list):
16
+ types.update(x for x in t if isinstance(x, str))
17
+ for branch in prop.get("anyOf", []):
18
+ if isinstance(branch, dict):
19
+ types |= allowed_types(branch)
20
+ return types
21
+
22
+
9
23
  class ToolSource(str, Enum):
10
24
  """工具来源类型。"""
11
25
 
@@ -66,9 +80,13 @@ class RunResult(ResultEnvelope):
66
80
 
67
81
 
68
82
  class SearchResult(BaseModel):
69
- """搜索结果条目。"""
83
+ """搜索结果条目。
84
+
85
+ 不含相关性分数:native(本地)与 MCP(daemon)来自两套独立索引、
86
+ 各自归一化,跨源分数不可比,暴露出来只会误导 Agent。结果按相关性
87
+ 排序即可,顺序本身就是信号。
88
+ """
70
89
 
71
90
  name: str = Field(description="工具完整名称")
72
91
  description: str = Field(description="工具描述")
73
92
  source: ToolSource = Field(description="工具来源")
74
- score: float | None = Field(default=None, description="相关性分数")
@@ -8,6 +8,7 @@ import threading
8
8
  from contextlib import AsyncExitStack
9
9
  from typing import Any, Self
10
10
 
11
+ import httpx
11
12
  from mcp import ClientSession, StdioServerParameters
12
13
  from mcp.client.stdio import stdio_client
13
14
  from mcp.client.streamable_http import streamable_http_client
@@ -49,27 +50,40 @@ class MCPConnection:
49
50
  self._exit_stack: AsyncExitStack | None = None
50
51
 
51
52
  async def connect(self) -> None:
52
- """建立连接并初始化 session"""
53
- self._exit_stack = AsyncExitStack()
54
-
55
- if self.config.url:
56
- read_stream, write_stream, _ = await self._exit_stack.enter_async_context(
57
- streamable_http_client(self.config.url)
58
- )
59
- else:
60
- server_params = StdioServerParameters(
61
- command=self.config.command,
62
- args=self.config.args,
63
- env=self.config.env,
64
- )
65
- devnull = self._exit_stack.enter_context(open(os.devnull, "w"))
66
- read_stream, write_stream = await self._exit_stack.enter_async_context(
67
- stdio_client(server_params, errlog=devnull)
53
+ """建立连接并初始化 session。中途失败时 with 退出自动清理已打开的资源。"""
54
+ async with AsyncExitStack() as stack:
55
+ if self.config.url:
56
+ http_client = None
57
+ if self.config.headers:
58
+ # 超时与 SDK 默认客户端(create_mcp_http_client)一致:
59
+ # 读超时 300s 是 SSE 长调用的存活窗口,不能用整体 30s
60
+ http_client = await stack.enter_async_context(
61
+ httpx.AsyncClient(
62
+ headers=self.config.headers,
63
+ timeout=httpx.Timeout(30, read=300),
64
+ follow_redirects=True,
65
+ )
66
+ )
67
+ read_stream, write_stream, _ = await stack.enter_async_context(
68
+ streamable_http_client(self.config.url, http_client=http_client)
69
+ )
70
+ else:
71
+ server_params = StdioServerParameters(
72
+ command=self.config.command,
73
+ args=self.config.args,
74
+ env=self.config.env,
75
+ )
76
+ devnull = stack.enter_context(open(os.devnull, "w"))
77
+ read_stream, write_stream = await stack.enter_async_context(
78
+ stdio_client(server_params, errlog=devnull)
79
+ )
80
+ session = await stack.enter_async_context(
81
+ ClientSession(read_stream, write_stream)
68
82
  )
69
- self.session = await self._exit_stack.enter_async_context(
70
- ClientSession(read_stream, write_stream)
71
- )
72
- await self.session.initialize()
83
+ await session.initialize()
84
+ # 全部成功才接管资源所有权;此前任何一步抛出都由 with 清理
85
+ self._exit_stack = stack.pop_all()
86
+ self.session = session
73
87
 
74
88
  async def list_tools(self) -> list[ToolMeta]:
75
89
  """获取该 server 的所有工具。"""
@@ -183,6 +183,15 @@ def _load_native_entry(
183
183
  for name in new_names:
184
184
  registry.set_server(name, server_name)
185
185
  _rekey_native_tool(name, f"{server_name}/{name}")
186
+ # 记录来源,便于审计:entry_points 自动发现意味着任何已安装的包都可能注入工具
187
+ logger.info(
188
+ "Loaded %d native tool(s) from '%s' (%s) under server '%s': %s",
189
+ len(new_names),
190
+ module,
191
+ source,
192
+ server_name,
193
+ ", ".join(new_names),
194
+ )
186
195
 
187
196
 
188
197
  def load_native_tool_modules() -> None:
@@ -2,12 +2,17 @@
2
2
 
3
3
  from __future__ import annotations
4
4
 
5
+ from collections import Counter
6
+ from typing import TYPE_CHECKING
7
+
5
8
  from axi.config import SearchConfig
6
9
  from axi.models import SearchResult, ToolMeta
7
10
  from axi.search.cache import EmbeddingCache
8
11
  from axi.search.embedding import EmbeddingProvider, create_embedding_provider
9
- from axi.search.hybrid import HybridSearch
10
- from axi.search.regex import RegexSearch
12
+
13
+ if TYPE_CHECKING:
14
+ from axi.search.hybrid import HybridSearch
15
+ from axi.search.regex import RegexSearch
11
16
 
12
17
 
13
18
  def split_names(value: str) -> list[str]:
@@ -38,14 +43,19 @@ class Registry:
38
43
  weight_embedding: float = 0.7,
39
44
  ) -> None:
40
45
  self._tools: dict[str, ToolMeta] = {}
41
- self._regex = RegexSearch()
42
- self._hybrid = HybridSearch(
43
- embedding_provider=embedding_provider,
44
- embedding_cache=embedding_cache,
45
- weight_bm25=weight_bm25,
46
- weight_embedding=weight_embedding,
47
- )
48
- self._dirty = True
46
+ # 搜索引擎按需构建:bm25s/jieba/numpy 是重 import,describe/run/list
47
+ # 完全用不到,延迟到 search()/grep() 首次调用时才付这笔冷启动成本。
48
+ self._embedding_provider = embedding_provider
49
+ self._embedding_cache = embedding_cache
50
+ self._weight_bm25 = weight_bm25
51
+ self._weight_embedding = weight_embedding
52
+ self._regex: RegexSearch | None = None
53
+ self._hybrid: HybridSearch | None = None
54
+ # 工具集每变一次 _version +1;各引擎记住自己上次 build 的 version,
55
+ # 只在落后时重建——晚建的引擎不再拖着已是最新的另一引擎陪跑一次。
56
+ self._version = 0
57
+ self._regex_version = -1
58
+ self._hybrid_version = -1
49
59
 
50
60
  @classmethod
51
61
  def from_search_config(cls, cfg: SearchConfig) -> Registry:
@@ -61,7 +71,7 @@ class Registry:
61
71
  def register(self, meta: ToolMeta) -> None:
62
72
  """注册一个工具。"""
63
73
  self._tools[meta.full_name] = meta
64
- self._dirty = True
74
+ self._version += 1
65
75
 
66
76
  def get(self, full_name: str) -> ToolMeta | None:
67
77
  """按完整名称获取工具元数据。"""
@@ -104,6 +114,10 @@ class Registry:
104
114
  """列出所有工具的 full_name。"""
105
115
  return list(self._tools.keys())
106
116
 
117
+ def server_tool_counts(self) -> dict[str, int]:
118
+ """按 server 统计工具数量(无 server 归入 "unknown")。"""
119
+ return dict(Counter(m.server or "unknown" for m in self._tools.values()))
120
+
107
121
  def set_server(self, full_name: str, server: str) -> None:
108
122
  """更新工具的 server 名,同时更新注册 key。"""
109
123
  meta = self._tools.pop(full_name, None)
@@ -111,22 +125,31 @@ class Registry:
111
125
  return
112
126
  updated = meta.model_copy(update={"server": server})
113
127
  self._tools[updated.full_name] = updated
114
- self._dirty = True
128
+ self._version += 1
115
129
 
116
130
  def search(self, query: str, top_k: int = 5) -> list[SearchResult]:
117
131
  """混合搜索工具(BM25 + Embedding)。"""
118
- self._rebuild_index()
132
+ if self._hybrid is None:
133
+ from axi.search.hybrid import HybridSearch
134
+
135
+ self._hybrid = HybridSearch(
136
+ embedding_provider=self._embedding_provider,
137
+ embedding_cache=self._embedding_cache,
138
+ weight_bm25=self._weight_bm25,
139
+ weight_embedding=self._weight_embedding,
140
+ )
141
+ if self._hybrid_version != self._version:
142
+ self._hybrid.build(list(self._tools.values()))
143
+ self._hybrid_version = self._version
119
144
  return self._hybrid.search(query, top_k=top_k)
120
145
 
121
146
  def grep(self, pattern: str, top_k: int = 10) -> list[SearchResult]:
122
147
  """正则表达式搜索工具。"""
123
- self._rebuild_index()
124
- return self._regex.search(pattern, top_k=top_k)
148
+ if self._regex is None:
149
+ from axi.search.regex import RegexSearch
125
150
 
126
- def _rebuild_index(self) -> None:
127
- if not self._dirty:
128
- return
129
- tools = list(self._tools.values())
130
- self._regex.build(tools)
131
- self._hybrid.build(tools)
132
- self._dirty = False
151
+ self._regex = RegexSearch()
152
+ if self._regex_version != self._version:
153
+ self._regex.build(list(self._tools.values()))
154
+ self._regex_version = self._version
155
+ return self._regex.search(pattern, top_k=top_k)
@@ -0,0 +1,6 @@
1
+ """搜索模块。
2
+
3
+ 子模块按需直接 import(``from axi.search.bm25 import BM25Search``)。
4
+ 此处不做包级 re-export,避免 ``import axi.search.cache`` 这类轻量导入
5
+ 连带把 bm25s/jieba/numpy 拖进来——describe/run/list 完全用不到搜索栈。
6
+ """
@@ -7,10 +7,13 @@ import json
7
7
  import logging
8
8
  from pathlib import Path
9
9
 
10
+ from axi.config import AXI_HOME, CONFIG_HASH
11
+
10
12
  logger = logging.getLogger(__name__)
11
13
 
12
- CACHE_DIR = Path(".axi/cache")
13
- CACHE_FILE = CACHE_DIR / "embeddings.json"
14
+ # 缓存随配置隔离,落在 ~/.axi 下(与 daemon 文件同处),不再依赖 cwd
15
+ CACHE_DIR = AXI_HOME / "cache"
16
+ CACHE_FILE = CACHE_DIR / f"{CONFIG_HASH}.json"
14
17
 
15
18
 
16
19
  def content_hash(text: str) -> str:
@@ -37,14 +37,14 @@ def _rrf_fuse(
37
37
  def _to_results(
38
38
  tools: list[ToolMeta], scored: list[tuple[int, float]]
39
39
  ) -> list[SearchResult]:
40
+ # scored 已按分数降序;分数只用于排序,不进入结果(跨源不可比)。
40
41
  return [
41
42
  SearchResult(
42
43
  name=tools[idx].full_name,
43
44
  description=tools[idx].description,
44
45
  source=tools[idx].source,
45
- score=round(score, 4),
46
46
  )
47
- for idx, score in scored
47
+ for idx, _score in scored
48
48
  ]
49
49
 
50
50
 
@@ -1,7 +0,0 @@
1
- """搜索模块。"""
2
-
3
- from axi.search.bm25 import BM25Search
4
- from axi.search.hybrid import HybridSearch
5
- from axi.search.regex import RegexSearch
6
-
7
- __all__ = ["BM25Search", "HybridSearch", "RegexSearch"]
File without changes
File without changes
File without changes
File without changes