krons 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

krons/agent/providers/claude_code.py

@@ -0,0 +1,276 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from typing import Any
+
+from pydantic import BaseModel
+
+from krons.agent.third_party.claude_code import (
+    ClaudeChunk,
+    ClaudeCodeRequest,
+    ClaudeSession,
+    stream_claude_code_cli,
+)
+from krons.resource.backend import NormalizedResponseModel
+from krons.resource.endpoint import Endpoint, EndpointConfig
+
+__all__ = (
+    "ClaudeCodeEndpoint",
+    "create_claude_code_config",
+)
+
+
+def create_claude_code_config(
+    name: str | None = None,
+) -> dict:
+    """Factory for Claude Code CLI endpoint config.
+
+    Args:
+        name: Config name (default: "claude_code_cli")
+
+    Returns:
+        Config dict for ClaudeCodeEndpoint
+    """
+    return {
+        "name": name or "claude_code_cli",
+        "provider": "claude_code",
+        "base_url": "internal",
+        "endpoint": "query_cli",
+        "api_key": "dummy-key",
+        "request_options": ClaudeCodeRequest,
+        "timeout": 3600,  # 1 hour max (EndpointConfig limit)
+    }
+
+
+class ClaudeCodeEndpoint(Endpoint):
+    """Claude Code CLI endpoint for local AI agent execution.
+
+    Usage:
+        endpoint = ClaudeCodeEndpoint()
+        response = await endpoint.call({
+            "messages": [{"role": "user", "content": "List files"}]
+        })
+    """
+
+    def __init__(
+        self,
+        config: dict | EndpointConfig | None = None,
+        circuit_breaker: Any | None = None,
+        **kwargs,
+    ):
+        """Initialize with Claude Code config."""
+        if config is None:
+            config = create_claude_code_config()
+        elif isinstance(config, EndpointConfig):
+            config = config.model_dump()
+        if not isinstance(config, dict):
+            raise ValueError(
+                "Provided config must be a dict or EndpointConfig instance"
+            )
+
+        super().__init__(config=config, circuit_breaker=circuit_breaker, **kwargs)
+
+    def create_payload(
+        self, request: dict | BaseModel, extra_headers: dict | None = None, **kwargs
+    ) -> tuple[dict, dict]:
+        """Create payload for Claude Code CLI.
+
+        Args:
+            request: Request dict with 'messages' or ClaudeCodeRequest
+            extra_headers: Ignored (CLI doesn't use HTTP headers)
+            **kwargs: Additional arguments merged into request
+
+        Returns:
+            (payload_dict, headers_dict) - headers always empty for CLI
+        """
+        from krons.utils.fuzzy import to_dict
+
+        # Convert request to dict if BaseModel
+        request_dict = to_dict(request) if isinstance(request, BaseModel) else request
+
+        # Merge config kwargs, request, and call kwargs (ignore extra_headers for CLI)
+        req_dict = {**self.config.kwargs, **request_dict, **kwargs}
+
+        # Extract messages (required)
+        messages = req_dict.pop("messages", None)
+        if not messages:
+            raise ValueError(
+                f"'messages' required for Claude Code endpoint. Got keys: {list(req_dict.keys())}"
+            )
+
+        # Create ClaudeCodeRequest object
+        # Pass messages to messages key so validator converts to prompt string
+        req_obj = ClaudeCodeRequest(messages=messages, **req_dict)  # type: ignore[arg-type]
+
+        return {"request": req_obj}, {}
+
+    async def call(
+        self,
+        request: dict | BaseModel,
+        skip_payload_creation: bool = False,
+        extra_headers: dict | None = None,
+        **kwargs,
+    ) -> NormalizedResponseModel:
+        """Execute Claude Code CLI and return normalized response.
+
+        Overrides parent to handle tuple return from create_payload.
+
+        Args:
+            request: Request parameters or Pydantic model.
+            skip_payload_creation: Bypass create_payload validation.
+            extra_headers: Ignored (CLI doesn't use HTTP headers).
+            **kwargs: Extra CLI arguments.
+
+        Returns:
+            NormalizedResponseModel wrapping the CLI response.
+        """
+        if skip_payload_creation:
+            # If already a ClaudeCodeRequest, wrap it; otherwise treat as dict
+            if isinstance(request, ClaudeCodeRequest):
+                payload = {"request": request}
+            elif isinstance(request, dict) and "request" in request:
+                # Already in payload format
+                payload = request
+            else:
+                # Assume it's a raw dict that needs to be converted
+                req_dict = (
+                    request if isinstance(request, dict) else request.model_dump()
+                )
+                messages = req_dict.pop("messages", None)
+                if messages:
+                    payload = {
+                        "request": ClaudeCodeRequest(messages=messages, **req_dict)
+                    }
+                else:
+                    raise ValueError("'messages' required for Claude Code endpoint")
+        else:
+            payload, _ = self.create_payload(request, extra_headers, **kwargs)
+
+        raw_response = await self._call(payload, {}, **kwargs)
+        return self.normalize_response(raw_response)
+
+    async def stream(  # type: ignore[override]
+        self, request: dict | BaseModel, **kwargs: Any
+    ) -> AsyncIterator[ClaudeChunk | dict | ClaudeSession]:
+        """Stream Claude Code CLI response chunks.
+
+        Yields:
+            ClaudeChunk, dict (system messages), or ClaudeSession (final)
+        """
+        payload, _ = self.create_payload(request, **kwargs)
+        request_obj: ClaudeCodeRequest = payload["request"]
+
+        async for chunk in stream_claude_code_cli(request_obj):
+            yield chunk
+
+    async def _call(
+        self,
+        payload: dict,
+        headers: dict,
+        **kwargs,
+    ) -> dict:
+        """Execute Claude Code CLI and return raw result chunk + session data.
+
+        Returns:
+            Dict with:
+            - raw_result: Raw "result" chunk from Claude Code CLI
+            - session: Organized session data from ClaudeSession
+        """
+        from krons.utils.fuzzy import to_dict
+
+        request: ClaudeCodeRequest = payload["request"]
+        session = ClaudeSession()
+        system: dict | None = None
+
+        # 1. Stream the Claude Code response
+        async for chunk in stream_claude_code_cli(request, session, **kwargs):
+            if isinstance(chunk, dict):
+                if chunk.get("type") == "done":
+                    break
+                system = chunk
+
+        # 2. Auto-finish if requested and not already finished
+        if request.auto_finish and not isinstance(
+            session.chunks[-1] if session.chunks else None, ClaudeSession
+        ):
+            req2 = request.model_copy(deep=True)
+            req2.prompt = "Please provide the final result message only"
+            req2.max_turns = 1
+            req2.continue_conversation = True
+            if system:
+                req2.resume = system.get("session_id")
+
+            async for chunk in stream_claude_code_cli(req2, session, **kwargs):
+                if isinstance(chunk, ClaudeSession):
+                    break
+
+        # 3. Use session.result directly (intermediate chunks are conversation flow, not final output)
+        # Don't concatenate chunk.text with session.result - causes duplication for JSON responses
+
+        # 4. Populate summary if requested
+        if request.cli_include_summary:
+            session.populate_summary()
+
+        # 5. Extract raw "result" chunk from session.chunks
+        raw_result_chunk = {}
+        for chunk in session.chunks:
+            if chunk.type == "result":
+                raw_result_chunk = chunk.raw
+                break
+
+        # 6. Return both raw CLI result and organized session data
+        return {
+            "raw_result": raw_result_chunk,
+            "session": to_dict(session, recursive=True),
+        }
+
+    def normalize_response(
+        self, raw_response: dict[str, Any]
+    ) -> NormalizedResponseModel:
+        """Normalize Claude Code response to standard format.
+
+        Args:
+            raw_response: Dict with:
+                - raw_result: Raw "result" chunk from Claude Code CLI
+                - session: Organized ClaudeSession data
+
+        Returns:
+            NormalizedResponseModel with:
+            - status: "success" or "error"
+            - data: Final result text
+            - raw_response: Actual raw CLI "result" chunk
+            - metadata: Organized session data
+        """
+        # Extract session data (our organized structure)
+        session = raw_response.get("session", {})
+        # Extract actual raw CLI result chunk
+        raw_cli_result = raw_response.get("raw_result", {})
+
+        # Extract text result from session
+        text = session.get("result", "")
+
+        # Extract metadata from organized session
+        metadata: dict[str, Any] = {
+            "session_id": session.get("session_id"),
+            "model": session.get("model"),
+            "usage": session.get("usage", {}),
+            "total_cost_usd": session.get("total_cost_usd"),
+            "num_turns": session.get("num_turns"),
+            "duration_ms": session.get("duration_ms"),
+            "duration_api_ms": session.get("duration_api_ms"),
+            "is_error": session.get("is_error", False),
+            "tool_uses": session.get("tool_uses", []),
+            "tool_results": session.get("tool_results", []),
+            "thinking_log": session.get("thinking_log", []),
+            "summary": session.get("summary"),  # Always include (None if not present)
+        }
+
+        return NormalizedResponseModel(
+            status="error" if session.get("is_error") else "success",
+            data=text,
+            raw_response=raw_cli_result,  # Use actual raw CLI result chunk
+            metadata=metadata,
+        )

krons/agent/providers/gemini.py

@@ -0,0 +1,268 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from typing import Any
+
+from pydantic import BaseModel
+
+from krons.agent.third_party.gemini_models import (
+    GeminiChunk,
+    GeminiCodeRequest,
+    GeminiSession,
+    stream_gemini_cli,
+)
+from krons.resource.backend import NormalizedResponseModel
+from krons.resource.endpoint import Endpoint, EndpointConfig
+
+__all__ = (
+    "GeminiCodeEndpoint",
+    "create_gemini_code_config",
+)
+
+
+def create_gemini_code_config(
+    name: str | None = None,
+    model: str | None = "gemini-2.5-pro",
+) -> dict:
+    """Factory for Gemini CLI endpoint config.
+
+    Args:
+        name: Config name (default: "gemini_code_cli")
+        model: Default model to use (gemini-2.5-pro, gemini-2.5-flash, gemini-3-pro)
+
+    Returns:
+        Config dict for GeminiCodeEndpoint
+    """
+    return {
+        "name": name or "gemini_code_cli",
+        "provider": "gemini_code",
+        "base_url": "internal",
+        "endpoint": "query_cli",
+        "api_key": "dummy-key",  # CLI uses Google OAuth, not API key
+        "request_options": GeminiCodeRequest,
+        "timeout": 3600,  # 1 hour max
+        "kwargs": {"model": model} if model else {},
+    }
+
+
+class GeminiCodeEndpoint(Endpoint):
+    """Gemini CLI endpoint for local AI agent execution.
+
+    Usage:
+        endpoint = GeminiCodeEndpoint()
+        response = await endpoint.call({
+            "messages": [{"role": "user", "content": "List files"}]
+        })
+
+        # Or with direct prompt:
+        response = await endpoint.call({
+            "prompt": "Analyze this codebase",
+            "yolo": True  # Auto-approve actions
+        })
+    """
+
+    def __init__(
+        self,
+        config: dict | EndpointConfig | None = None,
+        circuit_breaker: Any | None = None,
+        **kwargs,
+    ):
+        """Initialize with Gemini CLI config."""
+        if config is None:
+            config = create_gemini_code_config(
+                name=kwargs.pop("name", None),
+                model=kwargs.pop("model", "gemini-2.5-pro"),
+            )
+        elif isinstance(config, EndpointConfig):
+            config = config.model_dump()
+        if not isinstance(config, dict):
+            raise ValueError(
+                "Provided config must be a dict or EndpointConfig instance"
+            )
+
+        super().__init__(config=config, circuit_breaker=circuit_breaker, **kwargs)
+
+    def create_payload(
+        self, request: dict | BaseModel, extra_headers: dict | None = None, **kwargs
+    ) -> tuple[dict, dict]:
+        """Create payload for Gemini CLI.
+
+        Args:
+            request: Request dict with 'messages' or 'prompt', or GeminiCodeRequest
+            extra_headers: Ignored (CLI doesn't use HTTP headers)
+            **kwargs: Additional arguments merged into request
+
+        Returns:
+            (payload_dict, headers_dict) - headers always empty for CLI
+        """
+        from krons.utils.fuzzy import to_dict
+
+        # Convert request to dict if BaseModel
+        request_dict = to_dict(request) if isinstance(request, BaseModel) else request
+
+        # Merge config kwargs, request, and call kwargs
+        req_dict = {**self.config.kwargs, **request_dict, **kwargs}
+
+        # Extract prompt or messages (one is required)
+        prompt = req_dict.get("prompt")
+        messages = req_dict.get("messages")
+
+        if not prompt and not messages:
+            raise ValueError(
+                f"'prompt' or 'messages' required for Gemini CLI endpoint. "
+                f"Got keys: {list(req_dict.keys())}"
+            )
+
+        # Create GeminiCodeRequest object
+        req_obj = GeminiCodeRequest(**req_dict)
+
+        return {"request": req_obj}, {}
+
+    async def call(
+        self,
+        request: dict | BaseModel,
+        skip_payload_creation: bool = False,
+        extra_headers: dict | None = None,
+        **kwargs,
+    ) -> NormalizedResponseModel:
+        """Execute Gemini CLI and return normalized response.
+
+        Overrides parent to handle tuple return from create_payload.
+
+        Args:
+            request: Request parameters or Pydantic model.
+            skip_payload_creation: Bypass create_payload validation.
+            extra_headers: Ignored (CLI doesn't use HTTP headers).
+            **kwargs: Extra CLI arguments.
+
+        Returns:
+            NormalizedResponseModel wrapping the CLI response.
+        """
+        if skip_payload_creation:
+            # If already a GeminiCodeRequest, wrap it; otherwise treat as dict
+            if isinstance(request, GeminiCodeRequest):
+                payload = {"request": request}
+            elif isinstance(request, dict) and "request" in request:
+                # Already in payload format
+                payload = request
+            else:
+                # Assume it's a raw dict that needs to be converted
+                req_dict = (
+                    request if isinstance(request, dict) else request.model_dump()
+                )
+                payload = {"request": GeminiCodeRequest(**req_dict)}
+        else:
+            payload, _ = self.create_payload(request, extra_headers, **kwargs)
+
+        raw_response = await self._call(payload, {}, **kwargs)
+        return self.normalize_response(raw_response)
+
+    async def stream(  # type: ignore[override]
+        self, request: dict | BaseModel, **kwargs: Any
+    ) -> AsyncIterator[GeminiChunk | dict | GeminiSession]:
+        """Stream Gemini CLI response chunks.
+
+        Yields:
+            GeminiChunk, dict (system messages), or GeminiSession (final)
+        """
+        payload, _ = self.create_payload(request, **kwargs)
+        request_obj: GeminiCodeRequest = payload["request"]
+
+        async for chunk in stream_gemini_cli(request_obj):
+            yield chunk
+
+    async def _call(
+        self,
+        payload: dict,
+        headers: dict,
+        **kwargs,
+    ) -> dict:
+        """Execute Gemini CLI and return raw result + session data.
+
+        Returns:
+            Dict with:
+            - raw_result: Raw result chunk from Gemini CLI
+            - session: Organized session data from GeminiSession
+        """
+        from krons.utils.fuzzy import to_dict
+
+        request: GeminiCodeRequest = payload["request"]
+        session = GeminiSession()
+
+        # Stream the Gemini response
+        async for chunk in stream_gemini_cli(request, session, **kwargs):
+            if isinstance(chunk, dict):
+                if chunk.get("type") == "done":
+                    break
+            elif isinstance(chunk, GeminiSession):
+                break
+
+        # Use session.result if available (from final "result" event)
+        # Only fall back to combining chunk texts if no result was set
+        if not session.result:
+            texts = []
+            for chunk in session.chunks:
+                if chunk.text is not None:
+                    texts.append(chunk.text)
+            session.result = "\n".join(texts)
+
+        # Populate summary if requested
+        if request.cli_include_summary:
+            session.populate_summary()
+
+        # Extract raw result chunk
+        raw_result_chunk = {}
+        for chunk in session.chunks:
+            if chunk.type in ("result", "response"):
+                raw_result_chunk = chunk.raw
+                break
+
+        return {
+            "raw_result": raw_result_chunk,
+            "session": to_dict(session, recursive=True),
+        }
+
+    def normalize_response(
+        self, raw_response: dict[str, Any]
+    ) -> NormalizedResponseModel:
+        """Normalize Gemini CLI response to standard format.
+
+        Args:
+            raw_response: Dict with:
+                - raw_result: Raw result chunk from Gemini CLI
+                - session: Organized GeminiSession data
+
+        Returns:
+            NormalizedResponseModel with:
+            - status: "success" or "error"
+            - data: Final result text
+            - raw_response: Actual raw CLI result chunk
+            - metadata: Organized session data
+        """
+        session = raw_response.get("session", {})
+        raw_cli_result = raw_response.get("raw_result", {})
+
+        text = session.get("result", "")
+
+        metadata: dict[str, Any] = {
+            "session_id": session.get("session_id"),
+            "model": session.get("model"),
+            "usage": session.get("usage", {}),
+            "total_cost_usd": session.get("total_cost_usd"),
+            "num_turns": session.get("num_turns"),
+            "duration_ms": session.get("duration_ms"),
+            "is_error": session.get("is_error", False),
+            "tool_uses": session.get("tool_uses", []),
+            "tool_results": session.get("tool_results", []),
+            "summary": session.get("summary"),
+        }
+
+        return NormalizedResponseModel(
+            status="error" if session.get("is_error") else "success",
+            data=text,
+            raw_response=raw_cli_result,
+            metadata=metadata,
+        )

krons/agent/providers/match.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+import warnings
+
+from krons.resource.endpoint import Endpoint
+
+KNOWN_PROVIDERS = frozenset(
+    {
+        "anthropic",
+        "openai",
+        "groq",
+        "openrouter",
+        "nvidia_nim",
+        "claude_code",
+        "gemini_code",
+    }
+)
+
+
+def match_endpoint(
+    provider: str,
+    endpoint: str,
+    **kwargs,
+) -> Endpoint:
+    """Match provider and endpoint to appropriate Endpoint class.
+
+    Args:
+        provider: Provider name (e.g., "anthropic", "openai", "claude_code")
+        endpoint: Endpoint name (e.g., "messages", "chat/completions", "query_cli")
+        **kwargs: Additional kwargs passed to Endpoint constructor
+
+    Returns:
+        Endpoint instance configured for the provider
+
+    Example:
+        >>> endpoint = match_endpoint("anthropic", "messages", api_key="...")
+        >>> endpoint = match_endpoint("openai", "chat/completions")
+        >>> endpoint = match_endpoint("claude_code", "query_cli")
+    """
+    if provider == "anthropic" and ("messages" in endpoint or "chat" in endpoint):
+        from .anthropic_messages import AnthropicMessagesEndpoint
+
+        return AnthropicMessagesEndpoint(None, **kwargs)
+
+    if provider == "claude_code":
+        from .claude_code import ClaudeCodeEndpoint
+
+        return ClaudeCodeEndpoint(None, **kwargs)
+
+    if provider == "gemini_code":
+        from .gemini import GeminiCodeEndpoint
+
+        return GeminiCodeEndpoint(None, **kwargs)
+
+    if (
+        provider in ("openai", "groq", "openrouter", "nvidia_nim")
+        and "chat" in endpoint
+    ):
+        from .oai_chat import OAIChatEndpoint
+
+        return OAIChatEndpoint(None, provider=provider, endpoint=endpoint, **kwargs)
+
+    # OpenAI-compatible fallback with warning for unknown providers
+    if provider not in KNOWN_PROVIDERS:
+        warnings.warn(
+            f"Unknown provider '{provider}', falling back to OpenAI-compatible endpoint. "
+            f"Known providers: {sorted(KNOWN_PROVIDERS)}",
+            UserWarning,
+            stacklevel=2,
+        )
+
+    from .oai_chat import OAIChatEndpoint
+
+    return OAIChatEndpoint(None, provider=provider, endpoint=endpoint, **kwargs)