kagent-adk 0.7.11__py3-none-any.whl

kagent/adk/artifacts/stage_artifacts_tool.py

@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+import logging
+import mimetypes
+from pathlib import Path
+from typing import Any, List
+
+from google.adk.tools import BaseTool, ToolContext
+from google.genai import types
+from typing_extensions import override
+
+from .session_path import get_session_path
+
+logger = logging.getLogger("kagent_adk." + __name__)
+
+# Maximum file size for staging (100 MB)
+MAX_ARTIFACT_SIZE_BYTES = 100 * 1024 * 1024
+
+
+class StageArtifactsTool(BaseTool):
+    """A tool to stage artifacts from the artifact service to the local filesystem.
+
+    This tool enables working with user-uploaded files by staging them from the
+    artifact store to a local working directory where they can be accessed by
+    scripts, commands, and other tools.
+
+    Workflow:
+    1. Stage: Copy artifacts from artifact store to local 'uploads/' directory
+    2. Access: Use the staged files in commands, scripts, or other processing
+    """
+
+    def __init__(self):
+        super().__init__(
+            name="stage_artifacts",
+            description=(
+                "Stage artifacts from the artifact store to a local filesystem path, "
+                "making them available for processing by tools and scripts.\n\n"
+                "WORKFLOW:\n"
+                "1. When a user uploads a file, it's stored as an artifact with a name\n"
+                "2. Use this tool to copy the artifact to your local 'uploads/' directory\n"
+                "3. Then reference the staged file path in commands or scripts\n\n"
+                "USAGE EXAMPLE:\n"
+                "- stage_artifacts(artifact_names=['data.csv'])\n"
+                "  Returns: 'Successfully staged 1 file(s): uploads/data.csv (1.2 MB)'\n"
+                "- Then use: bash('python scripts/process.py uploads/data.csv')\n\n"
+                "PARAMETERS:\n"
+                "- artifact_names: List of artifact names to stage (required)\n"
+                "- destination_path: Target directory within session (default: 'uploads/')\n\n"
+                "BEST PRACTICES:\n"
+                "- Always stage artifacts before using them\n"
+                "- Use default 'uploads/' destination for consistency\n"
+                "- Stage all artifacts at the start of your workflow\n"
+                "- Check returned paths to confirm successful staging"
+            ),
+        )
+
+    def _get_declaration(self) -> types.FunctionDeclaration | None:
+        return types.FunctionDeclaration(
+            name=self.name,
+            description=self.description,
+            parameters=types.Schema(
+                type=types.Type.OBJECT,
+                properties={
+                    "artifact_names": types.Schema(
+                        type=types.Type.ARRAY,
+                        description=(
+                            "List of artifact names to stage. These are artifact identifiers "
+                            "provided by the system when files are uploaded. "
+                            "The tool will copy each artifact from the artifact store to the destination directory."
+                        ),
+                        items=types.Schema(type=types.Type.STRING),
+                    ),
+                    "destination_path": types.Schema(
+                        type=types.Type.STRING,
+                        description=(
+                            "Relative path within the session directory to save the files. "
+                            "Default is 'uploads/' where user-uploaded files are conventionally stored. "
+                            "Path must be within the session directory for security. "
+                            "Useful for organizing different types of artifacts (e.g., 'uploads/input/', 'uploads/processed/')."
+                        ),
+                        default="uploads/",
+                    ),
+                },
+                required=["artifact_names"],
+            ),
+        )
+
+    @override
+    async def run_async(self, *, args: dict[str, Any], tool_context: ToolContext) -> str:
+        artifact_names: List[str] = args.get("artifact_names", [])
+        destination_path_str: str = args.get("destination_path", "uploads/")
+
+        if not tool_context._invocation_context.artifact_service:
+            return "Error: Artifact service is not available in this context."
+
+        try:
+            staging_root = get_session_path(session_id=tool_context.session.id)
+            destination_dir = (staging_root / destination_path_str).resolve()
+
+            # Security: Ensure the destination is within the staging path
+            if staging_root not in destination_dir.parents and destination_dir != staging_root:
+                return f"Error: Invalid destination path '{destination_path_str}'."
+
+            destination_dir.mkdir(parents=True, exist_ok=True)
+
+            staged_files = []
+            for name in artifact_names:
+                artifact = await tool_context.load_artifact(name)
+                if artifact is None or artifact.inline_data is None:
+                    logger.warning('Artifact "%s" not found or has no data, skipping', name)
+                    continue
+
+                # Check file size
+                data_size = len(artifact.inline_data.data)
+                if data_size > MAX_ARTIFACT_SIZE_BYTES:
+                    size_mb = data_size / (1024 * 1024)
+                    logger.warning(f'Artifact "{name}" exceeds size limit: {size_mb:.1f} MB')
+                    continue
+
+                # Use artifact name as filename (frontend should provide meaningful names)
+                # If name has no extension, try to infer from MIME type
+                filename = self._ensure_proper_extension(name, artifact.inline_data.mime_type)
+                output_file = destination_dir / filename
+
+                # Write file to disk
+                output_file.write_bytes(artifact.inline_data.data)
+
+                relative_path = output_file.relative_to(staging_root)
+                size_kb = data_size / 1024
+                staged_files.append(f"{relative_path} ({size_kb:.1f} KB)")
+
+                logger.info(f"Staged artifact: {name} -> {relative_path} ({size_kb:.1f} KB)")
+
+            if not staged_files:
+                return "No valid artifacts were staged."
+
+            return f"Successfully staged {len(staged_files)} file(s):\n" + "\n".join(
+                f"  • {file}" for file in staged_files
+            )
+
+        except Exception as e:
+            logger.error("Error staging artifacts: %s", e, exc_info=True)
+            return f"An error occurred while staging artifacts: {e}"
+
+    def _ensure_proper_extension(self, filename: str, mime_type: str) -> str:
+        """Ensure filename has proper extension based on MIME type.
+
+        If filename already has an extension, keep it.
+        If not, add extension based on MIME type.
+
+        Args:
+            filename: Original filename from artifact
+            mime_type: MIME type of the file
+
+        Returns:
+            Filename with proper extension
+        """
+        if not filename or not mime_type:
+            return filename
+
+        # If filename already has an extension, use it
+        if Path(filename).suffix:
+            return filename
+
+        # Try to infer extension from MIME type
+        extension = mimetypes.guess_extension(mime_type)
+        if extension:
+            return f"{filename}{extension}"
+
+        return filename

kagent/adk/cli.py

@@ -0,0 +1,249 @@
+import asyncio
+import importlib
+import json
+import logging
+import os
+from typing import Annotated, Optional
+
+import typer
+import uvicorn
+from a2a.types import AgentCard
+from agentsts.adk import ADKSTSIntegration, ADKTokenPropagationPlugin
+from google.adk.agents import BaseAgent
+from google.adk.cli.utils.agent_loader import AgentLoader
+
+from kagent.core import KAgentConfig, configure_logging, configure_tracing
+
+from . import AgentConfig, KAgentApp
+from .skill_fetcher import fetch_skill
+from .tools import add_skills_tool_to_agent
+
+logger = logging.getLogger(__name__)
+logging.getLogger("google_adk.google.adk.tools.base_authenticated_tool").setLevel(logging.ERROR)
+
+app = typer.Typer()
+
+
+kagent_url_override = os.getenv("KAGENT_URL")
+sts_well_known_uri = os.getenv("STS_WELL_KNOWN_URI")
+propagate_token = os.getenv("KAGENT_PROPAGATE_TOKEN")
+
+
+def create_sts_integration() -> Optional[ADKTokenPropagationPlugin]:
+    if sts_well_known_uri or propagate_token:
+        sts_integration = None
+        if sts_well_known_uri:
+            sts_integration = ADKSTSIntegration(sts_well_known_uri)
+        return ADKTokenPropagationPlugin(sts_integration)
+
+
+def maybe_add_skills(root_agent: BaseAgent):
+    skills_directory = os.getenv("KAGENT_SKILLS_FOLDER", None)
+    if skills_directory:
+        logger.info(f"Adding skills from directory: {skills_directory}")
+        add_skills_tool_to_agent(skills_directory, root_agent)
+
+
+@app.command()
+def static(
+    host: str = "127.0.0.1",
+    port: int = 8080,
+    workers: int = 1,
+    filepath: str = "/config",
+    reload: Annotated[bool, typer.Option("--reload")] = False,
+):
+    app_cfg = KAgentConfig()
+
+    with open(os.path.join(filepath, "config.json"), "r") as f:
+        config = json.load(f)
+    agent_config = AgentConfig.model_validate(config)
+    with open(os.path.join(filepath, "agent-card.json"), "r") as f:
+        agent_card = json.load(f)
+    agent_card = AgentCard.model_validate(agent_card)
+    plugins = None
+    sts_integration = create_sts_integration()
+    if sts_integration:
+        plugins = [sts_integration]
+
+    def root_agent_factory() -> BaseAgent:
+        root_agent = agent_config.to_agent(app_cfg.name, sts_integration)
+
+        maybe_add_skills(root_agent)
+
+        return root_agent
+
+    kagent_app = KAgentApp(
+        root_agent_factory,
+        agent_card,
+        app_cfg.url,
+        app_cfg.app_name,
+        plugins=plugins,
+        stream=agent_config.stream if agent_config.stream is not None else False,
+    )
+
+    server = kagent_app.build()
+    configure_tracing(server)
+
+    uvicorn.run(
+        server,
+        host=host,
+        port=port,
+        workers=workers,
+        reload=reload,
+    )
+
+
+@app.command()
+def pull_skills(
+    skills: Annotated[list[str], typer.Argument()],
+    insecure: Annotated[
+        bool,
+        typer.Option("--insecure", help="Allow insecure connections to registries"),
+    ] = False,
+):
+    skill_dir = os.environ.get("KAGENT_SKILLS_FOLDER", ".")
+    logger.info("Pulling skills")
+    for skill in skills:
+        fetch_skill(skill, skill_dir, insecure)
+
+
+def add_to_agent(sts_integration: ADKTokenPropagationPlugin, agent: BaseAgent):
+    """
+    Add the plugin to an ADK LLM agent by updating its MCP toolset
+    Call this once when setting up the agent; do not call it at runtime.
+    """
+    from google.adk.agents import LlmAgent
+    from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
+
+    if not isinstance(agent, LlmAgent):
+        return
+
+    if not agent.tools:
+        return
+
+    for tool in agent.tools:
+        if isinstance(tool, McpToolset):
+            mcp_toolset = tool
+            mcp_toolset._header_provider = sts_integration.header_provider
+            logger.debug("Updated tool connection params to include access token from STS server")
+
+
+@app.command()
+def run(
+    name: Annotated[str, typer.Argument(help="The name of the agent to run")],
+    working_dir: str = ".",
+    host: str = "127.0.0.1",
+    port: int = 8080,
+    workers: int = 1,
+    local: Annotated[
+        bool, typer.Option("--local", help="Run with in-memory session service (for local development)")
+    ] = False,
+):
+    app_cfg = KAgentConfig()
+
+    plugins = None
+    sts_integration = create_sts_integration()
+    if sts_integration:
+        plugins = [sts_integration]
+
+    agent_loader = AgentLoader(agents_dir=working_dir)
+
+    def root_agent_factory() -> BaseAgent:
+        root_agent = agent_loader.load_agent(name)
+
+        if sts_integration:
+            add_to_agent(sts_integration, root_agent)
+
+        maybe_add_skills(root_agent)
+
+        return root_agent
+
+    # Load agent config to get stream setting
+    agent_config = None
+    config_path = os.path.join(working_dir, name, "config.json")
+    try:
+        with open(config_path, "r") as f:
+            config = json.load(f)
+        agent_config = AgentConfig.model_validate(config)
+    except FileNotFoundError:
+        logger.debug(f"No config.json found at {config_path}, using defaults")
+
+    with open(os.path.join(working_dir, name, "agent-card.json"), "r") as f:
+        agent_card = json.load(f)
+    agent_card = AgentCard.model_validate(agent_card)
+
+    # Attempt to import optional user-defined lifespan(app) from the agent package
+    lifespan = None
+    try:
+        module_candidate = importlib.import_module(name)
+        if hasattr(module_candidate, "lifespan"):
+            lifespan = module_candidate.lifespan
+    except Exception:
+        logger.exception(f"Failed to load agent module '{name}' for lifespan")
+
+    kagent_app = KAgentApp(
+        root_agent_factory,
+        agent_card,
+        app_cfg.url,
+        app_cfg.app_name,
+        lifespan=lifespan,
+        plugins=plugins,
+        stream=agent_config.stream if agent_config and agent_config.stream is not None else False,
+    )
+
+    if local:
+        logger.info("Running in local mode with InMemorySessionService")
+        server = kagent_app.build(local=True)
+    else:
+        server = kagent_app.build()
+
+    configure_tracing(server)
+
+    uvicorn.run(
+        server,
+        host=host,
+        port=port,
+        workers=workers,
+    )
+
+
+async def test_agent(agent_config: AgentConfig, agent_card: AgentCard, task: str):
+    app_cfg = KAgentConfig(url="http://fake-url.example.com", name="test-agent", namespace="kagent")
+    plugins = None
+    sts_integration = create_sts_integration()
+    if sts_integration:
+        plugins = [sts_integration]
+
+    def root_agent_factory() -> BaseAgent:
+        root_agent = agent_config.to_agent(app_cfg.name, sts_integration)
+        maybe_add_skills(root_agent)
+        return root_agent
+
+    app = KAgentApp(root_agent_factory, agent_card, app_cfg.url, app_cfg.app_name, plugins=plugins)
+    await app.test(task)
+
+
+@app.command()
+def test(
+    task: Annotated[str, typer.Option("--task", help="The task to test the agent with")],
+    filepath: Annotated[str, typer.Option("--filepath", help="The path to the agent config file")],
+):
+    with open(os.path.join(filepath, "config.json"), "r") as f:
+        content = f.read()
+        config = json.loads(content)
+
+    with open(os.path.join(filepath, "agent-card.json"), "r") as f:
+        agent_card = json.load(f)
+    agent_card = AgentCard.model_validate(agent_card)
+    agent_config = AgentConfig.model_validate(config)
+    asyncio.run(test_agent(agent_config, agent_card, task))
+
+
+def run_cli():
+    configure_logging()
+    logger.info("Starting KAgent")
+    app()
+
+
+if __name__ == "__main__":
+    run_cli()

kagent/adk/converters/error_mappings.py

@@ -0,0 +1,60 @@
+"""Error code to user-friendly message mappings for ADK events.
+
+This module provides mappings from Google GenAI finish reasons to user-friendly
+error messages, excluding STOP which is a normal completion reason.
+"""
+
+from typing import Dict, Optional
+
+from google.genai import types as genai_types
+
+# Error code to user-friendly message mappings
+# Based on Google GenAI types.py FinishReason enum (excluding STOP)
+ERROR_CODE_MESSAGES: Dict[str, str] = {
+    # Length and token limits
+    genai_types.FinishReason.MAX_TOKENS: "Response was truncated due to maximum token limit. Try asking a shorter question or breaking it into parts.",
+    # Safety and content filtering
+    genai_types.FinishReason.SAFETY: "Response was blocked due to safety concerns. Please rephrase your request to avoid potentially harmful content.",
+    genai_types.FinishReason.RECITATION: "Response was blocked due to unauthorized citations. Please rephrase your request.",
+    genai_types.FinishReason.BLOCKLIST: "Response was blocked due to restricted terminology. Please rephrase your request using different words.",
+    genai_types.FinishReason.PROHIBITED_CONTENT: "Response was blocked due to prohibited content. Please rephrase your request.",
+    genai_types.FinishReason.SPII: "Response was blocked due to sensitive personal information concerns. Please avoid including personal details.",
+    # Function calling errors
+    genai_types.FinishReason.MALFORMED_FUNCTION_CALL: "The agent generated an invalid function call. This may be due to complex input data. Try rephrasing your request or breaking it into simpler steps.",
+    # Generic fallback
+    genai_types.FinishReason.OTHER: "An unexpected error occurred during processing. Please try again or rephrase your request.",
+}
+
+# Normal completion reasons that should not be treated as errors
+NORMAL_COMPLETION_REASONS = {
+    genai_types.FinishReason.STOP,  # Normal completion
+}
+
+# Default error message when no specific mapping exists
+DEFAULT_ERROR_MESSAGE = "An error occurred during processing"
+
+
+def _get_error_message(error_code: Optional[str]) -> str:
+    """Get a user-friendly error message for the given error code.
+
+    Args:
+        error_code: The error code from the ADK event (e.g., finish_reason)
+
+    Returns:
+        User-friendly error message string
+    """
+
+    # Return mapped message or default
+    return ERROR_CODE_MESSAGES.get(error_code, DEFAULT_ERROR_MESSAGE)
+
+
+def _is_normal_completion(error_code: Optional[str]) -> bool:
+    """Check if the error code represents normal completion rather than an error.
+
+    Args:
+        error_code: The error code to check
+
+    Returns:
+        True if this is a normal completion reason, False otherwise
+    """
+    return error_code in NORMAL_COMPLETION_REASONS