axi-cli 0.0.3__py3-none-any.whl

axi/__init__.py

@@ -0,0 +1,81 @@
+"""axi — Agent eXecution Interface."""
+
+import logging
+from typing import Any, Callable
+
+from axi.cli import get_executor, get_registry
+from axi.providers.native import register_tool
+
+__all__ = ["tool"]
+
+# 包级 logger：各子模块通过 logging.getLogger(__name__) 自动继承
+_logger = logging.getLogger("axi")
+_logger.addHandler(logging.NullHandler())
+
+
+def tool(
+    name: str | None = None,
+    description: str | None = None,
+    output_example: Any | None = None,
+) -> Callable:
+    """装饰器：注册一个 Python 函数为 axi 工具。
+
+    用法：
+        @tool(name="query_orders", description="按区域查询订单")
+        def query_orders(region: str, limit: int = 10) -> dict:
+            ...
+
+    也可以作为函数调用获取 PTC 可调用对象：
+        query = tool("query_orders")
+        result = query(region="cn")
+    """
+    # 当作为 PTC 调用时：tool("tool_name") 返回可调用对象
+    if name is not None and description is None and output_example is None:
+        # 先查本地原生工具
+        registry = get_registry()
+        meta = registry.get(name)
+        if meta is not None:
+            executor = get_executor()
+
+            def _native_caller(**kwargs: Any) -> Any:
+                result = executor.run(name, kwargs)  # type: ignore[arg-type]
+                if result.status == "error":
+                    raise RuntimeError(result.error)
+                return result.data
+
+            return _native_caller
+
+        # 再查 daemon（MCP 工具）
+        return _make_daemon_caller(name)
+
+    # 当作为装饰器时
+    def decorator(func: Callable) -> Callable:
+        meta = register_tool(
+            func,
+            name=name,
+            description=description,
+            output_example=output_example,
+        )
+        get_registry().register(meta)
+        return func
+
+    return decorator
+
+
+def _make_daemon_caller(tool_name: str) -> Callable:
+    """创建通过 daemon 调用 MCP 工具的 PTC 函数。"""
+    from axi.daemon.client import ensure_daemon, send_request
+    from axi.daemon.protocol import DaemonRequest
+
+    def _daemon_caller(**kwargs: Any) -> Any:
+        if not ensure_daemon():
+            raise RuntimeError("Daemon is not running. Start it with: axi daemon start")
+
+        resp = send_request(
+            DaemonRequest(method="call_tool", tool_name=tool_name, params=kwargs)
+        )
+        if resp.status == "error":
+            raise RuntimeError(resp.error)
+        return resp.data
+
+    return _daemon_caller

axi/cli.py

@@ -0,0 +1,378 @@
+"""axi CLI 入口：Typer app。"""
+
+import json
+import logging
+from collections import Counter
+
+import typer
+from pydantic import BaseModel
+
+from axi.config import app_config
+from axi.daemon.client import ensure_daemon, is_daemon_running, send_request
+from axi.daemon.protocol import DaemonRequest, DaemonResponse
+from axi.executor import Executor
+from axi.models import RunResult, SearchResult
+from axi.registry import AmbiguousToolError, Registry, ToolNotFoundError
+
+logger = logging.getLogger(__name__)
+
+app = typer.Typer(
+    name="axi",
+    help="Agent eXecution Interface - unified tool layer for AI Agents",
+    no_args_is_help=True,
+    rich_markup_mode="rich" if app_config.cli.rich else None,
+)
+
+daemon_app = typer.Typer(help="管理 axi daemon")
+app.add_typer(daemon_app, name="daemon")
+
+# 全局实例（原生工具用）
+_registry = Registry()
+_executor = Executor(_registry)
+
+
+@app.callback(invoke_without_command=True)
+def main_callback(ctx: typer.Context) -> None:
+    """CLI 启动时加载 axi.json 中配置的原生工具模块。"""
+    from axi.providers.mcp import load_native_tool_modules
+
+    load_native_tool_modules()
+    if ctx.invoked_subcommand is None:
+        raise typer.Exit()
+
+
+def get_registry() -> Registry:
+    return _registry
+
+
+def get_executor() -> Executor:
+    return _executor
+
+
+def _output_json(data: object) -> None:
+    """统一 JSON 输出。"""
+    if isinstance(data, BaseModel):
+        d = data.model_dump(exclude_none=True)
+    elif isinstance(data, list):
+        d = [
+            item.model_dump(exclude_none=True) if isinstance(item, BaseModel) else item
+            for item in data
+        ]
+    else:
+        d = data
+    typer.echo(json.dumps(d, ensure_ascii=False))
+
+
+def _daemon_request(req: DaemonRequest) -> DaemonResponse:
+    """向 daemon 发送请求。如果 daemon 未运行则自动启动。"""
+    if not ensure_daemon():
+        return DaemonResponse.fail(
+            "Daemon is not running. Start it with: axi daemon start"
+        )
+    return send_request(req)
+
+
+# ── daemon 管理命令 ──────────────────────────────────────────────
+
+
+@daemon_app.command("start")
+def daemon_start() -> None:
+    """启动 daemon。"""
+    if is_daemon_running():
+        typer.echo("Daemon is already running.")
+        return
+
+    if ensure_daemon():
+        typer.echo("Daemon started.")
+    else:
+        typer.echo("Failed to start daemon.")
+        raise typer.Exit(code=1)
+
+
+@daemon_app.command("stop")
+def daemon_stop() -> None:
+    """停止 daemon。"""
+    if not is_daemon_running():
+        typer.echo("Daemon is not running.")
+        return
+
+    try:
+        resp = send_request(DaemonRequest(method="shutdown"))
+        if resp.status == "error":
+            typer.echo(f"Failed to stop daemon: {resp.error}", err=True)
+            raise typer.Exit(code=1)
+    except OSError as e:
+        typer.echo(f"Failed to connect to daemon: {e}", err=True)
+        raise typer.Exit(code=1)
+    typer.echo("Daemon stopped.")
+
+
+@daemon_app.command("status")
+def daemon_status() -> None:
+    """查看 daemon 状态。"""
+    if not is_daemon_running():
+        _output_json({"status": "stopped"})
+        return
+
+    resp = send_request(DaemonRequest(method="status"))
+    if resp.status == "error":
+        _output_json({"status": "error", "error": resp.error})
+        raise typer.Exit(code=1)
+
+    data = resp.data
+    # 原生工具按 server 统计
+    native_server_tools = dict(
+        Counter(meta.server or "unknown" for meta in _registry.list_all())
+    )
+
+    result = {
+        "status": "running",
+        **data,
+        "native_tools": native_server_tools,
+    }
+    _output_json(result)
+
+
+# ── 核心命令 ──────────────────────────────────────────────
+
+
+def _search_and_merge(
+    local_results: list[SearchResult],
+    daemon_method: str,
+    query: str,
+    top_k: int,
+) -> None:
+    """执行 daemon 搜索，合并本地和远程结果后输出 JSON。"""
+    mcp_results: list[dict] = []
+    resp = _daemon_request(
+        DaemonRequest(method=daemon_method, query=query, top_k=top_k)
+    )
+    if resp.status == "success" and resp.data:
+        mcp_results = resp.data
+    elif resp.status == "error":
+        logger.warning("Daemon search failed: %s", resp.error)
+
+    combined = [r.model_dump(exclude_none=True) for r in local_results] + (
+        mcp_results or []
+    )
+    typer.echo(json.dumps(combined, ensure_ascii=False))
+
+
+@app.command()
+def search(
+    query: str = typer.Argument(help="搜索关键词"),
+    top_k: int = typer.Option(5, "--top-k", "-k", help="返回结果数量"),
+) -> None:
+    """混合搜索工具（BM25 + Embedding）。"""
+    local_results = _registry.search(query, top_k=top_k)
+    _search_and_merge(local_results, "search", query, top_k)
+
+
+@app.command()
+def grep(
+    pattern: str = typer.Argument(help="正则表达式"),
+    limit: int = typer.Option(10, "--limit", "-l", help="返回结果数量"),
+) -> None:
+    """正则表达式搜索工具。"""
+    try:
+        local_results = _registry.grep(pattern, top_k=limit)
+    except ValueError as e:
+        _output_json({"error": str(e)})
+        raise typer.Exit(code=1)
+    _search_and_merge(local_results, "grep", pattern, limit)
+
+
+def _collect_tool_groups() -> dict[str | None, list[dict]]:
+    """收集所有工具并按 server 分组。"""
+    from axi.providers.mcp import MCPProvider
+
+    groups: dict[str | None, list[dict]] = {}
+
+    # 本地原生工具
+    for meta in _registry.list_all():
+        groups.setdefault(meta.server, []).append(
+            {"name": meta.name, "description": meta.description}
+        )
+
+    # MCP 工具（通过 daemon）
+    mcp_tools: list[dict] = []
+    resp = _daemon_request(DaemonRequest(method="list_tools"))
+    if resp.status == "success" and resp.data:
+        mcp_tools = resp.data
+        for t in mcp_tools:
+            groups.setdefault(t.get("server"), []).append(
+                {"name": t["name"], "description": t.get("description", "")}
+            )
+
+    # daemon 未返回工具时，至少从配置列出 server
+    if not mcp_tools:
+        provider = MCPProvider()
+        for cfg in provider.load_config():
+            if cfg.server not in groups:
+                groups[cfg.server] = []
+
+    return groups
+
+
+def _filter_groups(
+    groups: dict[str | None, list[dict]], server_name: str
+) -> dict[str, list[dict]]:
+    """按逗号分隔的 server 名过滤分组。未找到时抛 typer.Exit。"""
+    names = [n.strip() for n in server_name.split(",") if n.strip()]
+    filtered = {k: v for k, v in groups.items() if k in names}
+    if not filtered:
+        missing = [n for n in names if n not in groups]
+        _output_json({"error": f"Server not found: {', '.join(missing)}"})
+        raise typer.Exit(code=1)
+    return filtered
+
+
+@app.command("list")
+def list_tools(
+    server_name: str | None = typer.Argument(
+        None, help="只列出指定 server 的工具（逗号分隔多个）"
+    ),
+) -> None:
+    """列出所有 server 及其工具。"""
+    groups = _collect_tool_groups()
+
+    if server_name is not None:
+        filtered = _filter_groups(groups, server_name)
+        if len(filtered) == 1:
+            key = next(iter(filtered))
+            typer.echo(
+                json.dumps({"server": key, "tools": filtered[key]}, ensure_ascii=False)
+            )
+        else:
+            typer.echo(
+                json.dumps(
+                    [{"server": k, "tools": v} for k, v in filtered.items()],
+                    ensure_ascii=False,
+                )
+            )
+        return
+
+    # 全部列出（只显示工具名）
+    result = [
+        {"server": key, "tools": [t["name"] for t in tools]}
+        for key, tools in groups.items()
+    ]
+    typer.echo(json.dumps(result, ensure_ascii=False))
+
+
+def _resolve_tool(name: str) -> dict:
+    """解析单个工具，返回工具详情 dict 或 error dict。"""
+    try:
+        meta = _registry.resolve(name)
+        return meta.model_dump(exclude_none=True)
+    except AmbiguousToolError as e:
+        return {"error": str(e)}
+    except ToolNotFoundError:
+        pass  # 本地未找到，继续尝试 daemon
+
+    resp = _daemon_request(DaemonRequest(method="describe", tool_name=name))
+    if resp.status == "success":
+        return resp.data
+    return {"error": resp.error or f"Tool not found: {name}"}
+
+
+@app.command()
+def describe(
+    tool_name: str = typer.Argument(help="工具完整名称（逗号分隔多个）"),
+) -> None:
+    """查看工具详情。"""
+    names = [n.strip() for n in tool_name.split(",") if n.strip()]
+    if len(names) == 1:
+        result = _resolve_tool(names[0])
+        if "error" in result:
+            _output_json(result)
+            raise typer.Exit(code=1)
+        _output_json(result)
+        return
+    results = [_resolve_tool(n) for n in names]
+    typer.echo(json.dumps(results, ensure_ascii=False))
+
+
+@app.command(
+    context_settings={"allow_extra_args": True, "allow_interspersed_args": False},
+)
+def run(
+    ctx: typer.Context,
+    tool_name: str = typer.Argument(help="工具完整名称"),
+) -> None:
+    """执行工具。参数支持 --key value 或 --json '{...}' 格式。"""
+    args = ctx.args
+
+    if "--help" in args or "-h" in args:
+        typer.echo(ctx.get_help())
+        raise typer.Exit()
+
+    json_str, args = _extract_option(args, "--json", "-j")
+
+    if json_str:
+        try:
+            parsed = json.loads(json_str)
+        except json.JSONDecodeError as e:
+            _output_json(RunResult.fail(f"Invalid JSON argument: {e}"))
+            raise typer.Exit(code=1)
+    else:
+        parsed = _parse_params(args)
+
+    try:
+        meta = _registry.resolve(tool_name)
+        result = _executor.run(meta.full_name, parsed)
+        _output_json(result)
+        return
+    except AmbiguousToolError as e:
+        _output_json({"error": str(e)})
+        raise typer.Exit(code=1)
+    except ToolNotFoundError:
+        pass  # 本地未找到，继续尝试 daemon
+
+    resp = _daemon_request(
+        DaemonRequest(method="call_tool", tool_name=tool_name, params=parsed)
+    )
+    if resp.status == "success":
+        _output_json(RunResult.success(resp.data))
+    else:
+        _output_json(RunResult.fail(resp.error or "Unknown error"))
+
+
+# ── 参数解析辅助函数 ──────────────────────────────────────────────
+
+
+def _extract_option(args: list[str], *names: str) -> tuple[str | None, list[str]]:
+    remaining = []
+    value = None
+    i = 0
+    while i < len(args):
+        if args[i] in names and i + 1 < len(args):
+            value = args[i + 1]
+            i += 2
+        else:
+            remaining.append(args[i])
+            i += 1
+    return value, remaining
+
+
+def _parse_params(params: list[str]) -> dict:
+    parsed: dict = {}
+    i = 0
+    while i < len(params):
+        arg = params[i]
+        if arg.startswith("--"):
+            key = arg[2:]
+            if i + 1 < len(params) and not params[i + 1].startswith("--"):
+                value = params[i + 1]
+                try:
+                    parsed[key] = json.loads(value)
+                except (json.JSONDecodeError, ValueError):
+                    parsed[key] = value
+                i += 2
+            else:
+                parsed[key] = True
+                i += 1
+        else:
+            logger.warning("Ignoring unrecognized argument: %s", arg)
+            i += 1
+    return parsed

axi/config.py

@@ -0,0 +1,140 @@
+"""axi 配置中心：Pydantic 模型化，统一加载 axi.json，全局共享。"""
+
+import json
+import logging
+import os
+from pathlib import Path
+
+from pydantic import BaseModel, Field, model_validator
+
+logger = logging.getLogger(__name__)
+
+CONFIG_PATH = Path(os.environ.get("AXI_CONFIG", "axi.json"))
+
+
+# ── 子配置模型 ──────────────────────────────────────────────
+
+
+class CliConfig(BaseModel):
+    """CLI 显示配置。"""
+
+    rich: bool = Field(default=False, description="启用 Rich 格式化输出")
+
+    @model_validator(mode="before")
+    @classmethod
+    def override_with_env(cls, values: dict) -> dict:
+        if isinstance(values, dict):
+            env = os.environ.get("AXI_RICH", "").lower()
+            if env in ("1", "true"):
+                values["rich"] = True
+            elif env in ("0", "false"):
+                values["rich"] = False
+        return values
+
+
+class EmbeddingConfig(BaseModel):
+    """Embedding 搜索配置。"""
+
+    provider: str | None = Field(default=None, description="jina 或 openai")
+    api_key: str | None = Field(default=None, alias="apiKey", description="API 密钥")
+    model: str | None = Field(default=None, description="模型名称")
+    base_url: str | None = Field(
+        default=None, alias="baseUrl", description="自定义端点"
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def override_with_env(cls, values: dict) -> dict:
+        if isinstance(values, dict) and not values.get("apiKey"):
+            provider = values.get("provider", "")
+            if provider == "jina":
+                values["apiKey"] = os.environ.get("JINA_API_KEY")
+            elif provider == "openai":
+                values["apiKey"] = os.environ.get("OPENAI_API_KEY")
+        return values
+
+
+class SearchWeightsConfig(BaseModel):
+    """混合搜索 RRF 融合权重。"""
+
+    bm25: float = Field(default=0.3, description="BM25 权重")
+    embedding: float = Field(default=0.7, description="Embedding 权重")
+
+
+class SearchConfig(BaseModel):
+    """搜索引擎配置。"""
+
+    embedding: EmbeddingConfig = Field(default_factory=EmbeddingConfig)
+    weights: SearchWeightsConfig = Field(default_factory=SearchWeightsConfig)
+
+
+class DaemonConfig(BaseModel):
+    """Daemon 进程配置。"""
+
+    idle_timeout_minutes: int = Field(
+        default=30, alias="idleTimeoutMinutes", description="空闲自动关闭（分钟）"
+    )
+
+
+class MCPServerConfig(BaseModel):
+    """单个 MCP server 的配置。"""
+
+    command: str | None = None
+    args: list[str] = Field(default_factory=list)
+    env: dict[str, str] | None = None
+    url: str | None = None
+
+
+class NativeToolEntry(BaseModel):
+    """原生工具模块声明。"""
+
+    module: str = Field(description="Python 模块路径或文件路径")
+    name: str | None = Field(default=None, description="server 名，省略时自动推导")
+
+
+# ── 主配置 ──────────────────────────────────────────────
+
+
+class AxiConfig(BaseModel):
+    """axi 主配置。"""
+
+    cli: CliConfig = Field(default_factory=CliConfig)
+    mcp_servers: dict[str, MCPServerConfig] = Field(
+        default_factory=dict, alias="mcpServers"
+    )
+    native_tools: list[NativeToolEntry] = Field(
+        default_factory=list, alias="nativeTools"
+    )
+    search: SearchConfig = Field(default_factory=SearchConfig)
+    daemon: DaemonConfig = Field(default_factory=DaemonConfig)
+
+
+# ── 加载 ──────────────────────────────────────────────
+
+
+def load_config(path: Path) -> AxiConfig:
+    """读取并解析配置文件。找不到文件则返回默认配置，格式错误则报错退出。"""
+    if not path.exists():
+        return AxiConfig()
+    with open(path) as f:
+        try:
+            raw = json.load(f)
+        except json.JSONDecodeError as e:
+            raise SystemExit(f"Error: Malformed config file {path}: {e}")
+    try:
+        return AxiConfig.model_validate(raw)
+    except Exception as e:
+        raise SystemExit(f"Error: Invalid config in {path}: {e}")
+
+
+def _load_app_config() -> AxiConfig:
+    """延迟加载配置，捕获异常并输出友好信息。"""
+    try:
+        return load_config(CONFIG_PATH)
+    except SystemExit:
+        raise
+    except Exception as e:
+        raise SystemExit(f"Error: Failed to load config: {e}")
+
+
+app_config: AxiConfig = _load_app_config()

axi/daemon/client.py

@@ -0,0 +1,94 @@
+"""daemon 客户端：CLI 侧通过 Unix socket 与 daemon 通信。"""
+
+import asyncio
+import logging
+import os
+import subprocess
+import sys
+import time
+
+from axi.daemon.protocol import (
+    SOCKET_DIR,
+    SOCKET_PATH,
+    PID_PATH,
+    DaemonRequest,
+    DaemonResponse,
+)
+
+logger = logging.getLogger(__name__)
+
+DAEMON_LOG_PATH = os.path.join(SOCKET_DIR, "daemon.log")
+
+_DAEMON_START_POLL_RETRIES = 30
+_DAEMON_START_POLL_INTERVAL = 0.1  # seconds
+_DAEMON_REQUEST_TIMEOUT = 30  # seconds
+
+
+def is_daemon_running() -> bool:
+    """检查 daemon 是否在运行。"""
+    if not os.path.exists(PID_PATH):
+        return False
+
+    try:
+        with open(PID_PATH) as f:
+            pid = int(f.read().strip())
+        os.kill(pid, 0)
+        return os.path.exists(SOCKET_PATH)
+    except (OSError, ValueError):
+        return False
+
+
+def ensure_daemon() -> bool:
+    """确保 daemon 已启动。未运行时自动启动，返回是否就绪。"""
+    if is_daemon_running():
+        return True
+
+    os.makedirs(SOCKET_DIR, exist_ok=True)
+
+    with open(DAEMON_LOG_PATH, "a") as log_file:
+        subprocess.Popen(
+            [sys.executable, "-m", "axi.daemon.server"],
+            stdout=log_file,
+            stderr=log_file,
+            start_new_session=True,
+        )
+
+    for _ in range(_DAEMON_START_POLL_RETRIES):
+        time.sleep(_DAEMON_START_POLL_INTERVAL)
+        if is_daemon_running():
+            return True
+    logger.error("Daemon failed to start. Check log: %s", DAEMON_LOG_PATH)
+    return False
+
+
+def send_request(req: DaemonRequest) -> DaemonResponse:
+    """向 daemon 发送请求并获取响应。"""
+    return asyncio.run(_send(req))
+
+
+async def _send(req: DaemonRequest) -> DaemonResponse:
+    try:
+        reader, writer = await asyncio.open_unix_connection(SOCKET_PATH)
+    except OSError as e:
+        return DaemonResponse.fail(
+            f"Cannot connect to daemon: {e}. Try: axi daemon stop && axi daemon start"
+        )
+
+    try:
+        writer.write(req.model_dump_json().encode() + b"\n")
+        await writer.drain()
+
+        line = await asyncio.wait_for(
+            reader.readline(), timeout=_DAEMON_REQUEST_TIMEOUT
+        )
+        if not line:
+            return DaemonResponse.fail("Daemon connection closed unexpectedly")
+
+        return DaemonResponse.model_validate_json(line)
+    except asyncio.TimeoutError:
+        return DaemonResponse.fail(
+            f"Daemon request timed out after {_DAEMON_REQUEST_TIMEOUT}s"
+        )
+    finally:
+        writer.close()
+        await writer.wait_closed()