claxcode 0.1.1__py3-none-any.whl

claxcode/core/context.py

@@ -0,0 +1,114 @@
+"""Context window management."""
+
+from typing import Any
+
+from claxcode.config import settings
+
+
+def estimate_tokens(text: str) -> int:
+    """Very rough token estimation (1 token ~= 4 chars for English/CJK mix)."""
+    return len(text) // 3 + 1
+
+
+class ContextManager:
+    """Manages conversation history with compression support."""
+
+    def __init__(self, system_prompt: str = ""):
+        self.system_prompt = system_prompt
+        self.messages: list[dict[str, Any]] = []
+        if system_prompt:
+            self.messages.append({"role": "system", "content": system_prompt})
+
+    def add_message(self, role: str, content: str) -> None:
+        self.messages.append({"role": role, "content": content})
+        self._maybe_compress()
+
+    def add_tool_result(self, tool_call_id: str, name: str, output: str) -> None:
+        """Add a tool result message (OpenAI format)."""
+        self.messages.append(
+            {
+                "role": "tool",
+                "tool_call_id": tool_call_id,
+                "content": output,
+            }
+        )
+        self._maybe_compress()
+
+    def add_assistant_tool_calls(self, tool_calls: list[dict[str, Any]]) -> None:
+        self.messages.append(
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": tool_calls,
+            }
+        )
+        self._maybe_compress()
+
+    def get_messages(self) -> list[dict[str, Any]]:
+        return self.messages.copy()
+
+    def _total_tokens(self) -> int:
+        text = "\n".join(str(m.get("content", "")) for m in self.messages)
+        return estimate_tokens(text) + estimate_tokens(self.system_prompt)
+
+    def over_threshold(self) -> bool:
+        """Whether the estimated token count exceeds the compression threshold."""
+        threshold = int(settings.max_context_tokens * settings.context_compression_threshold)
+        return self._total_tokens() > threshold
+
+    def pop_oldest(self, count: int) -> list[dict[str, Any]]:
+        """Remove and return up to `count` oldest non-system messages.
+
+        Orphaned tool results at the cut boundary are dropped as well —
+        providers reject a tool result without its assistant tool call.
+        """
+        end = min(1 + count, len(self.messages))
+        removed = self.messages[1:end]
+        del self.messages[1:end]
+        while len(self.messages) > 1 and self.messages[1].get("role") == "tool":
+            removed.append(self.messages.pop(1))
+        return removed
+
+    def insert_summary(self, summary: str) -> None:
+        """Insert a compression summary right after the system prompt."""
+        self.messages.insert(1, {"role": "system", "content": summary})
+
+    def restore(self, messages: list[dict[str, Any]]) -> None:
+        """Restore persisted messages after the system prompt.
+
+        Leading tool results are dropped — without their assistant tool call
+        (e.g. lost to compression in a previous session) providers reject them.
+        Trailing assistant tool_calls without matching tool results are also
+        dropped, e.g. when a previous run crashed between tool_calls and results.
+        """
+        self.messages.extend(messages)
+        while len(self.messages) > 1 and self.messages[1].get("role") == "tool":
+            self.messages.pop(1)
+
+        i = 1
+        while i < len(self.messages):
+            msg = self.messages[i]
+            if msg.get("role") == "assistant" and msg.get("tool_calls"):
+                needed = {tc["id"] for tc in msg["tool_calls"]}
+                j = i + 1
+                found: set[str] = set()
+                while j < len(self.messages) and self.messages[j].get("role") == "tool":
+                    found.add(self.messages[j].get("tool_call_id"))
+                    j += 1
+                if needed != found:
+                    # Drop the orphan assistant and any trailing tool messages.
+                    del self.messages[i:j]
+                    continue
+            i += 1
+
+    def _maybe_compress(self) -> None:
+        """Emergency fallback: hard-drop oldest messages while over the threshold."""
+        while self.over_threshold() and len(self.messages) > 3:
+            self.messages.pop(1)  # keep system at index 0
+        while len(self.messages) > 1 and self.messages[1].get("role") == "tool":
+            self.messages.pop(1)
+
+    def clear(self) -> None:
+        self.messages = []
+        if self.system_prompt:
+            self.messages.append({"role": "system", "content": self.system_prompt})

claxcode/core/models.py

@@ -0,0 +1,446 @@
+"""LLM abstraction layer supporting multiple providers and custom models."""
+
+import json
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+
+from claxcode.config import settings
+
+
+class ToolCallChunk:
+    def __init__(self):
+        self.id: str = ""
+        self.name: str = ""
+        self.arguments: str = ""
+
+
+class LLMProvider(ABC):
+    """Abstract LLM provider."""
+
+    @abstractmethod
+    async def stream_chat(
+        self,
+        messages: list[dict[str, str]],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> AsyncIterator[str | dict[str, Any]]:
+        """Yield text chunks or tool-call dicts."""
+        ...
+
+
+class DummyProvider(LLMProvider):
+    """Fallback provider when no LLM package is installed or config is invalid."""
+
+    def __init__(self, reason: str = ""):
+        self.reason = reason
+
+    async def stream_chat(
+        self,
+        messages: list[dict[str, str]],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> AsyncIterator[str | dict[str, Any]]:
+        if self.reason:
+            yield f"[DummyProvider] {self.reason}"
+        else:
+            yield "[DummyProvider] 未配置 LLM。请设置 LLM_PROVIDER 和 LLM_API_KEY。"
+
+
+class OpenAIProvider(LLMProvider):
+    """OpenAI-compatible API provider using httpx (no openai package required)."""
+
+    def __init__(self):
+        self.api_key = settings.llm_api_key or "sk-test"
+        self.base_url = (settings.llm_base_url or "https://api.openai.com").rstrip("/")
+        self.model = settings.llm_model
+        self.temperature = settings.llm_temperature
+        self.max_tokens = settings.llm_max_tokens
+
+    async def stream_chat(
+        self,
+        messages: list[dict[str, str]],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> AsyncIterator[str | dict[str, Any]]:
+        url = f"{self.base_url}/v1/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        payload: dict[str, Any] = {
+            "model": self.model,
+            "messages": messages,
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+            "stream": True,
+        }
+        if tools:
+            payload["tools"] = tools
+
+        if settings.debug:
+            print("[LLM PAYLOAD]", json.dumps(payload, ensure_ascii=False, indent=2))
+
+        async with httpx.AsyncClient(timeout=120) as client:
+            async with client.stream("POST", url, headers=headers, json=payload) as response:
+                if response.status_code >= 400:
+                    body = await response.aread()
+                    print(f"[LLM ERROR {response.status_code}] {body.decode('utf-8', errors='replace')}")
+                response.raise_for_status()
+                current_tool: ToolCallChunk | None = None
+
+                async for line in response.aiter_lines():
+                    if not line:
+                        continue
+                    if line.startswith("data: "):
+                        line = line[6:]
+                    if line == "[DONE]":
+                        break
+                    try:
+                        data = json.loads(line)
+                    except json.JSONDecodeError:
+                        continue
+
+                    delta = data.get("choices", [{}])[0].get("delta", {})
+                    content = delta.get("content", "")
+                    if content:
+                        yield content
+
+                    # Tool calls
+                    tcs = delta.get("tool_calls", [])
+                    for tc in tcs:
+                        if tc.get("id"):
+                            if current_tool is None:
+                                current_tool = ToolCallChunk()
+                            current_tool.id = tc["id"]
+                        func = tc.get("function", {})
+                        if func.get("name"):
+                            if current_tool is None:
+                                current_tool = ToolCallChunk()
+                            current_tool.name = func["name"]
+                        if func.get("arguments"):
+                            if current_tool is None:
+                                current_tool = ToolCallChunk()
+                            current_tool.arguments += func["arguments"]
+
+                    finish_reason = data.get("choices", [{}])[0].get("finish_reason")
+                    if finish_reason and current_tool and current_tool.name:
+                        try:
+                            args = json.loads(current_tool.arguments or "{}")
+                        except json.JSONDecodeError:
+                            # Some providers return malformed JSON; fall back to raw string
+                            args = {"_raw": current_tool.arguments}
+                        yield {
+                            "type": "tool_call",
+                            "id": current_tool.id,
+                            "name": current_tool.name,
+                            "arguments": args,
+                        }
+                        current_tool = None
+
+
+class AnthropicProvider(LLMProvider):
+    def __init__(self):
+        try:
+            from anthropic import AsyncAnthropic
+        except ImportError as exc:
+            raise ImportError(
+                "AnthropicProvider 需要 anthropic 包。请安装: pip install anthropic"
+            ) from exc
+
+        self.client = AsyncAnthropic(api_key=settings.llm_api_key or "sk-test")
+        self.model = settings.llm_model
+        self.temperature = settings.llm_temperature
+        self.max_tokens = settings.llm_max_tokens
+
+    async def stream_chat(
+        self,
+        messages: list[dict[str, str]],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> AsyncIterator[str | dict[str, Any]]:
+        system = ""
+        anthropic_messages = []
+        for m in messages:
+            if m["role"] == "system":
+                system = m["content"]
+            else:
+                anthropic_messages.append({"role": m["role"], "content": m["content"]})
+
+        kwargs: dict[str, Any] = {
+            "model": self.model,
+            "messages": anthropic_messages,
+            "max_tokens": self.max_tokens,
+            "temperature": self.temperature,
+            "stream": True,
+        }
+        if system:
+            kwargs["system"] = system
+        if tools:
+            kwargs["tools"] = tools
+
+        async with self.client.messages.stream(**kwargs) as stream:
+            current_tool: ToolCallChunk | None = None
+            async for event in stream:
+                if event.type == "content_block_delta":
+                    if hasattr(event.delta, "text") and event.delta.text:
+                        yield event.delta.text
+                    if hasattr(event.delta, "partial_json"):
+                        if current_tool is None:
+                            current_tool = ToolCallChunk()
+                        current_tool.arguments += event.delta.partial_json
+                elif event.type == "content_block_start":
+                    if hasattr(event.content_block, "name"):
+                        if current_tool is None:
+                            current_tool = ToolCallChunk()
+                        current_tool.name = event.content_block.name
+                        current_tool.id = getattr(event.content_block, "id", "")
+                elif event.type == "message_stop":
+                    if current_tool and current_tool.name:
+                        yield {
+                            "type": "tool_call",
+                            "id": current_tool.id,
+                            "name": current_tool.name,
+                            "arguments": json.loads(current_tool.arguments or "{}"),
+                        }
+                        current_tool = None
+
+
+class GenericOpenAIProvider(OpenAIProvider):
+    """For local models with OpenAI-compatible API (e.g., vLLM, LM Studio, DashScope)."""
+
+    pass
+
+
+class OllamaProvider(LLMProvider):
+    """Native Ollama API provider."""
+
+    def __init__(self):
+        self.base_url = settings.llm_base_url or "http://localhost:11434"
+        self.model = settings.llm_model
+        self.temperature = settings.llm_temperature
+        self.max_tokens = settings.llm_max_tokens
+        self.options = self._parse_options()
+
+    def _parse_options(self) -> dict[str, Any]:
+        """Parse extra Ollama options from custom config."""
+        opts: dict[str, Any] = {}
+        try:
+            custom = json.loads(settings.llm_custom_config or "{}")
+            opts.update(custom.get("ollama_options", {}))
+        except json.JSONDecodeError:
+            pass
+        return opts
+
+    async def stream_chat(
+        self,
+        messages: list[dict[str, str]],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> AsyncIterator[str | dict[str, Any]]:
+        url = f"{self.base_url.rstrip('/')}/api/chat"
+        payload: dict[str, Any] = {
+            "model": self.model,
+            "messages": messages,
+            "stream": True,
+            "options": {
+                "temperature": self.temperature,
+                "num_predict": self.max_tokens,
+                **self.options,
+            },
+        }
+        if tools:
+            payload["tools"] = [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": t["function"]["name"],
+                        "description": t["function"]["description"],
+                        "parameters": t["function"]["parameters"],
+                    },
+                }
+                for t in tools
+            ]
+
+        async with httpx.AsyncClient(timeout=120) as client:
+            async with client.stream("POST", url, json=payload) as response:
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    if not line:
+                        continue
+                    try:
+                        data = json.loads(line)
+                    except json.JSONDecodeError:
+                        continue
+                    if data.get("done"):
+                        break
+                    msg = data.get("message", {})
+                    content = msg.get("content", "")
+                    if content:
+                        yield content
+                    tool_calls = msg.get("tool_calls", [])
+                    for tc in tool_calls:
+                        func = tc.get("function", {})
+                        yield {
+                            "type": "tool_call",
+                            "id": tc.get("id", ""),
+                            "name": func.get("name", ""),
+                            "arguments": func.get("arguments", {}),
+                        }
+
+
+class CustomHTTPProvider(LLMProvider):
+    """Generic HTTP provider for custom LLM APIs.
+
+    Uses llm_custom_config to define request/response format:
+    {
+        "request_template": {
+            "url": "http://localhost:8000/v1/chat",
+            "method": "POST",
+            "headers": {"Authorization": "Bearer {api_key}"},
+            "body": {
+                "model": "{model}",
+                "messages": "{messages}",
+                "temperature": "{temperature}",
+                "max_tokens": "{max_tokens}",
+                "stream": true
+            }
+        },
+        "response_parser": {
+            "content_path": "choices.0.delta.content",
+            "tool_calls_path": "choices.0.delta.tool_calls"
+        }
+    }
+    """
+
+    def __init__(self):
+        self.model = settings.llm_model
+        self.temperature = settings.llm_temperature
+        self.max_tokens = settings.llm_max_tokens
+        self.api_key = settings.llm_api_key
+        self.custom_config = self._load_config()
+
+    def _load_config(self) -> dict[str, Any]:
+        try:
+            return json.loads(settings.llm_custom_config or "{}")
+        except json.JSONDecodeError:
+            return {}
+
+    def _render_template(self, template: Any, **values: Any) -> Any:
+        if isinstance(template, str):
+            result = template
+            for key, val in values.items():
+                result = result.replace(f"{{{key}}}", json.dumps(val) if isinstance(val, (list, dict)) else str(val))
+            return result
+        if isinstance(template, dict):
+            return {k: self._render_template(v, **values) for k, v in template.items()}
+        if isinstance(template, list):
+            return [self._render_template(item, **values) for item in template]
+        return template
+
+    async def stream_chat(
+        self,
+        messages: list[dict[str, str]],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> AsyncIterator[str | dict[str, Any]]:
+        cfg = self.custom_config.get("request_template", {})
+        url = cfg.get("url", "http://localhost:8000/v1/chat")
+        method = cfg.get("method", "POST")
+        headers = self._render_template(cfg.get("headers", {}), api_key=self.api_key)
+        body = self._render_template(
+            cfg.get("body", {}),
+            model=self.model,
+            messages=messages,
+            temperature=self.temperature,
+            max_tokens=self.max_tokens,
+            tools=tools or [],
+        )
+
+        async with httpx.AsyncClient(timeout=120) as client:
+            async with client.stream(method, url, headers=headers, json=body) as response:
+                response.raise_for_status()
+                parser = self.custom_config.get("response_parser", {})
+                content_path = parser.get("content_path", "choices.0.delta.content")
+
+                async for line in response.aiter_lines():
+                    if not line:
+                        continue
+                    if line.startswith("data: "):
+                        line = line[6:]
+                    if line == "[DONE]":
+                        break
+                    try:
+                        data = json.loads(line)
+                    except json.JSONDecodeError:
+                        continue
+
+                    content = self._get_path(data, content_path)
+                    if content:
+                        yield content
+
+                    tool_path = parser.get("tool_calls_path")
+                    if tool_path:
+                        tcs = self._get_path(data, tool_path)
+                        if tcs:
+                            for tc in tcs:
+                                yield {
+                                    "type": "tool_call",
+                                    "id": tc.get("id", ""),
+                                    "name": tc.get("function", {}).get("name", ""),
+                                    "arguments": tc.get("function", {}).get("arguments", {}),
+                                }
+
+    def _get_path(self, data: dict, path: str) -> Any:
+        parts = path.split(".")
+        current: Any = data
+        for part in parts:
+            if isinstance(current, dict):
+                current = current.get(part)
+            elif isinstance(current, list) and part.isdigit():
+                current = current[int(part)] if int(part) < len(current) else None
+            else:
+                return None
+            if current is None:
+                return None
+        return current
+
+
+_CUSTOM_PROVIDERS: dict[str, type[LLMProvider]] = {}
+
+
+def register_custom_provider(name: str, provider_cls: type[LLMProvider]) -> None:
+    """Register a custom LLM provider class at runtime.
+
+    Example:
+        from claxcode.core.models import register_custom_provider, LLMProvider
+
+        class MyProvider(LLMProvider):
+            async def stream_chat(self, messages, tools=None):
+                yield "Hello from my model!"
+
+        register_custom_provider("my_model", MyProvider)
+    """
+    _CUSTOM_PROVIDERS[name.lower()] = provider_cls
+
+
+def get_provider() -> LLMProvider:
+    provider = settings.llm_provider.lower()
+    if provider in _CUSTOM_PROVIDERS:
+        return _CUSTOM_PROVIDERS[provider]()
+    if provider == "openai":
+        return OpenAIProvider()
+    if provider == "anthropic":
+        try:
+            return AnthropicProvider()
+        except ImportError as exc:
+            return DummyProvider(
+                reason=f"加载提供商 '{provider}' 失败。{exc}. "
+                "请安装所需的包或选择其他 LLM_PROVIDER。"
+            )
+    if provider == "generic_openai":
+        return GenericOpenAIProvider()
+    if provider == "ollama":
+        return OllamaProvider()
+    if provider == "custom_http":
+        return CustomHTTPProvider()
+    return DummyProvider(
+        reason=f"未知的提供商 '{provider}'. "
+        "支持的提供商: openai, anthropic, generic_openai, ollama, custom_http。"
+    )

claxcode/extensions/__init__.py

@@ -0,0 +1 @@
+"""Extension system for ClaxCode."""

claxcode/extensions/hooks.py

@@ -0,0 +1,131 @@
+"""Lifecycle hooks — run shell commands before/after tool execution.
+
+Config is read from `~/.claxcode/hooks.json` (global) and
+`<project>/.claxcode/hooks.json` (project). Format:
+
+{
+  "PreToolUse":  [{"match": "run_shell|write_file", "command": "python guard.py"}],
+  "PostToolUse": [{"match": "", "command": "echo done"}]
+}
+
+`match` is a regex tested against the tool name (empty or "*" matches all).
+The hook command receives a JSON payload on stdin
+({"tool": ..., "arguments": {...}, "result": ...}) plus the env vars
+CLAXCODE_TOOL_NAME / CLAXCODE_TOOL_ARGS, and runs with the config's
+project directory as cwd. A PreToolUse hook blocks the tool call when it
+exits non-zero; its stdout becomes the block reason. Hooks that fail to
+spawn or time out never block the main flow.
+"""
+
+import asyncio
+import json
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+HOOK_TIMEOUT = 30
+HOOK_EVENTS = ("PreToolUse", "PostToolUse")
+
+
+@dataclass
+class Hook:
+    """A single shell hook from a hooks.json config."""
+
+    match: str
+    command: str
+    cwd: Path
+
+    def matches(self, tool_name: str) -> bool:
+        if not self.match or self.match == "*":
+            return True
+        try:
+            return re.search(self.match, tool_name) is not None
+        except re.error:
+            return self.match == tool_name
+
+
+class HookManager:
+    """Load and execute PreToolUse/PostToolUse shell hooks."""
+
+    def __init__(self, project_path: Path | None = None):
+        self._hooks: dict[str, list[Hook]] = {event: [] for event in HOOK_EVENTS}
+        for config_path in self._config_paths(project_path):
+            self._load_config(config_path)
+
+    def _config_paths(self, project_path: Path | None) -> list[Path]:
+        paths = [Path.home() / ".claxcode" / "hooks.json"]
+        if project_path is not None:
+            paths.append(project_path / ".claxcode" / "hooks.json")
+        return paths
+
+    def _load_config(self, path: Path) -> None:
+        if not path.exists():
+            return
+        try:
+            data = json.loads(path.read_text(encoding="utf-8"))
+        except (json.JSONDecodeError, OSError):
+            return
+        cwd = path.parent.parent if path.parent.name == ".claxcode" else path.parent
+        for event in HOOK_EVENTS:
+            for entry in data.get(event, []):
+                command = entry.get("command", "")
+                if command:
+                    self._hooks[event].append(Hook(entry.get("match", ""), command, cwd))
+
+    async def run_pre(self, tool_name: str, arguments: dict[str, Any]) -> tuple[bool, str]:
+        """Run PreToolUse hooks. Returns (allowed, reason)."""
+        for hook in self._hooks["PreToolUse"]:
+            if not hook.matches(tool_name):
+                continue
+            code, output = await self._run_command(hook, tool_name, arguments)
+            if code != 0:
+                return False, output or f"钩子命令退出码 {code}"
+        return True, ""
+
+    async def run_post(self, tool_name: str, arguments: dict[str, Any], result: str) -> str:
+        """Run PostToolUse hooks. Returns combined hook output (may be empty)."""
+        outputs = []
+        for hook in self._hooks["PostToolUse"]:
+            if not hook.matches(tool_name):
+                continue
+            _, output = await self._run_command(hook, tool_name, arguments, result)
+            if output:
+                outputs.append(output)
+        return "\n".join(outputs)
+
+    async def _run_command(
+        self,
+        hook: Hook,
+        tool_name: str,
+        arguments: dict[str, Any],
+        result: str = "",
+    ) -> tuple[int, str]:
+        payload = {"tool": tool_name, "arguments": arguments, "result": result}
+        env = {
+            **os.environ,
+            "CLAXCODE_TOOL_NAME": tool_name,
+            "CLAXCODE_TOOL_ARGS": json.dumps(arguments, ensure_ascii=False),
+        }
+        try:
+            proc = await asyncio.create_subprocess_shell(
+                hook.command,
+                stdin=asyncio.subprocess.PIPE,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.STDOUT,
+                cwd=str(hook.cwd),
+                env=env,
+            )
+        except OSError:
+            return 0, ""  # 钩子无法启动时不阻塞主流程
+
+        try:
+            stdout, _ = await asyncio.wait_for(
+                proc.communicate(json.dumps(payload, ensure_ascii=False).encode()),
+                timeout=HOOK_TIMEOUT,
+            )
+        except TimeoutError:
+            proc.kill()
+            return 0, ""
+        return proc.returncode or 0, stdout.decode("utf-8", errors="replace").strip()