dg-dataagent 0.1.0__py3-none-any.whl

dataagent/__init__.py

@@ -0,0 +1,37 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""dataagent 包入口。
+
+注意：这里保持轻量，避免在 import dataagent 时就触发大量依赖初始化/循环导入。
+需要 DataAgent 等对象时，通过属性懒加载获取。
+"""
+
+from __future__ import annotations
+
+__version__ = "0.1.0"
+__author__ = "DataAgent Team"
+
+from typing import Any
+
+
+def __getattr__(name: str) -> Any:
+    if name in ("DataAgent", "AgentBuilder", "load_agent_from_config", "BaseDataAgent"):
+        from dataagent.interface.sdk import AgentBuilder, BaseDataAgent, DataAgent, load_agent_from_config
+
+        return {
+            "DataAgent": DataAgent,
+            "AgentBuilder": AgentBuilder,
+            "BaseDataAgent": BaseDataAgent,
+            "load_agent_from_config": load_agent_from_config,
+        }[name]
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

dataagent/__main__.py

@@ -0,0 +1,16 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+from dataagent.interface.cli.main import main
+
+if __name__ == "__main__":
+    main()

dataagent/a2a_server/__init__.py

@@ -0,0 +1,25 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""DataAgent A2A 1.0 Server — exposes DataAgent agents as A2A-compliant services."""
+
+from dataagent.a2a_server.agent_card import build_agent_card
+from dataagent.a2a_server.agent_executor import DataAgentExecutor
+from dataagent.a2a_server.server import create_a2a_server, run_a2a_server
+
+__all__ = [
+    "DataAgentExecutor",
+    "build_agent_card",
+    "create_a2a_server",
+    "run_a2a_server",
+]

dataagent/a2a_server/agent_card.py

@@ -0,0 +1,85 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Build A2A 1.0 AgentCard from a DataAgent."""
+
+from a2a.types.a2a_pb2 import (
+    AgentCapabilities,
+    AgentCard,
+    AgentInterface,
+    AgentSkill,
+)
+
+from dataagent.interface.sdk.agent import DataAgent
+
+
+def build_agent_card(
+    agent: DataAgent,
+    host: str = "0.0.0.0",
+    port: int = 9999,
+    jsonrpc_path: str = "/a2a/jsonrpc",
+    rest_path: str = "/a2a/rest",
+) -> AgentCard:
+    """Create an A2A 1.0 AgentCard for the given DataAgent.
+
+    Args:
+        agent: The DataAgent instance to expose.
+        host: Server host for the interface URLs.
+        port: Server port for the interface URLs.
+        jsonrpc_path: JSON-RPC route path.
+        rest_path: REST route path.
+
+    Returns:
+        A2A 1.0 AgentCard with supported_interfaces.
+    """
+    agent_name = agent.name() or "DataAgent"
+    agent_desc = agent.description() or "DataAgent data analysis agent"
+    agent_ver = agent.version() or "0.1.0"
+
+    # AgentCard URL should use 127.0.0.1 for local access (host may be 0.0.0.0 for binding)
+    card_host = "127.0.0.1" if host in ("0.0.0.0", "::") else host
+    base_url = f"http://{card_host}:{port}"
+
+    # Build skills from agent info and known capabilities
+    skills = [
+        AgentSkill(
+            id="chat",
+            name="Chat",
+            description="Interactive conversational data analysis",
+            tags=["data-analysis", "chat"],
+            input_modes=["text/plain"],
+            output_modes=["text/plain"],
+        ),
+    ]
+
+    card = AgentCard(
+        name=agent_name,
+        description=agent_desc,
+        version=agent_ver,
+        default_input_modes=["text/plain"],
+        default_output_modes=["text/plain"],
+        capabilities=AgentCapabilities(streaming=True),
+        skills=skills,
+        supported_interfaces=[
+            AgentInterface(
+                protocol_binding="JSONRPC",
+                url=f"{base_url}{jsonrpc_path}",
+                protocol_version="1.0",
+            ),
+            AgentInterface(
+                protocol_binding="HTTP+JSON",
+                url=f"{base_url}{rest_path}",
+                protocol_version="1.0",
+            ),
+        ],
+    )
+    return card

dataagent/a2a_server/agent_executor.py

@@ -0,0 +1,310 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""DataAgentExecutor — bridges A2A 1.0 requests to DataAgent."""
+
+import asyncio
+import uuid
+from typing import Any
+
+from a2a.helpers import (
+    new_task,
+    new_text_artifact,
+    new_text_message,
+    new_text_status_update_event,
+)
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.events import EventQueue
+from a2a.types.a2a_pb2 import (
+    Role,
+    TaskArtifactUpdateEvent,
+    TaskState,
+    TaskStatus,
+    TaskStatusUpdateEvent,
+)
+from loguru import logger
+
+from dataagent.interface.sdk.agent import DataAgent
+
+
+class DataAgentExecutor(AgentExecutor):
+    """Bridges A2A 1.0 protocol requests to DataAgent execution."""
+
+    def __init__(self, agent: DataAgent):
+        self._agent = agent
+        self._cancellation_events: dict[str, asyncio.Event] = {}
+        self._run_counters: dict[str, int] = {}
+        self._run_lock = asyncio.Lock()
+        self._session_locks: dict[str, asyncio.Lock] = {}
+        self._session_locks_guard = asyncio.Lock()
+
+    async def execute(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+    ) -> None:
+        """Handle an A2A request by invoking the DataAgent.
+
+        Publishes task status updates and artifacts to the event queue,
+        supporting both sync and streaming (SSE) responses.
+        """
+        task = context.current_task or new_task(
+            task_id=context.task_id,
+            context_id=context.context_id,
+            state=TaskState.TASK_STATE_SUBMITTED,
+        )
+        await event_queue.enqueue_event(task)
+
+        user_text = _extract_text_from_context(context)
+        if not user_text:
+            await self._fail_with_message(event_queue, context, "No user message provided.")
+            return
+
+        cancel_event = asyncio.Event()
+        self._cancellation_events[context.task_id] = cancel_event
+
+        session_id = context.context_id or context.task_id or f"a2a-{uuid.uuid4().hex[:12]}"
+        session_lock = await self._get_session_lock(session_id)
+
+        try:
+            async with session_lock:
+                await event_queue.enqueue_event(
+                    new_text_status_update_event(
+                        task_id=context.task_id,
+                        context_id=context.context_id,
+                        state=TaskState.TASK_STATE_WORKING,
+                        text="Processing request...",
+                    )
+                )
+
+                response = await self._execute_agent_chat(user_text, session_id, context.task_id)
+
+                if cancel_event.is_set():
+                    await event_queue.enqueue_event(
+                        new_text_status_update_event(
+                            task_id=context.task_id,
+                            context_id=context.context_id,
+                            state=TaskState.TASK_STATE_CANCELED,
+                            text="Task was canceled.",
+                        )
+                    )
+                    return
+
+                result_text = _extract_final_answer(response)
+
+                has_error = self._has_error_response(response)
+                if has_error:
+                    error_msg = self._extract_error_message(response)
+                    logger.error(f"Agent execution failed: {error_msg}")
+                    await event_queue.enqueue_event(
+                        new_text_status_update_event(
+                            task_id=context.task_id,
+                            context_id=context.context_id,
+                            state=TaskState.TASK_STATE_FAILED,
+                            text=str(error_msg) if error_msg else "Agent execution failed",
+                        )
+                    )
+                    return
+
+                logger.info("Agent execution completed successfully")
+
+                await event_queue.enqueue_event(
+                    TaskArtifactUpdateEvent(
+                        task_id=context.task_id,
+                        context_id=context.context_id,
+                        artifact=new_text_artifact(name="dataagent_result", text=result_text),
+                    )
+                )
+                await event_queue.enqueue_event(
+                    new_text_status_update_event(
+                        task_id=context.task_id,
+                        context_id=context.context_id,
+                        state=TaskState.TASK_STATE_COMPLETED,
+                        text=result_text,
+                    )
+                )
+
+        except asyncio.CancelledError:
+            await event_queue.enqueue_event(
+                new_text_status_update_event(
+                    task_id=context.task_id,
+                    context_id=context.context_id,
+                    state=TaskState.TASK_STATE_CANCELED,
+                    text="Task was canceled.",
+                )
+            )
+        except Exception as e:
+            logger.error(f"DataAgent execution failed: {e}")
+            await event_queue.enqueue_event(
+                new_text_status_update_event(
+                    task_id=context.task_id,
+                    context_id=context.context_id,
+                    state=TaskState.TASK_STATE_FAILED,
+                    text=f"Error: {str(e)}",
+                )
+            )
+        finally:
+            self._cancellation_events.pop(context.task_id, None)
+
+    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
+        """Request cancellation of an ongoing task."""
+        cancel_event = self._cancellation_events.get(context.task_id)
+        if cancel_event:
+            cancel_event.set()
+        await event_queue.enqueue_event(
+            new_text_status_update_event(
+                task_id=context.task_id,
+                context_id=context.context_id,
+                state=TaskState.TASK_STATE_CANCELED,
+                text="Task was canceled.",
+            )
+        )
+
+    async def _fail_with_message(
+        self,
+        event_queue: EventQueue,
+        context: RequestContext,
+        message: str,
+    ) -> None:
+        """Publish a failed status event with the given message."""
+        await event_queue.enqueue_event(
+            TaskStatusUpdateEvent(
+                task_id=context.task_id,
+                context_id=context.context_id,
+                status=TaskStatus(
+                    state=TaskState.TASK_STATE_FAILED,
+                    message=new_text_message(
+                        text=message,
+                        role=Role.ROLE_AGENT,
+                        task_id=context.task_id,
+                        context_id=context.context_id,
+                    ),
+                ),
+            )
+        )
+
+    async def _execute_agent_chat(
+        self,
+        user_text: str,
+        session_id: str,
+        task_id: str,
+    ) -> Any:
+        """Acquire session lock, increment run counter, and execute agent chat."""
+        async with self._run_lock:
+            run_id = self._run_counters.get(session_id, 0)
+            self._run_counters[session_id] = run_id + 1
+        initial_state = {"session_id": session_id, "run_id": run_id}
+
+        response = await self._agent.chat(user_query=user_text, session_id=session_id, initial_state=initial_state)
+
+        logger.info(f"[DEBUG] _agent.chat() returned type: {type(response)}")
+        if isinstance(response, dict):
+            logger.info(f"[DEBUG] _agent.chat() response keys: {list(response.keys())}")
+            logger.info(f"[DEBUG] _agent.chat() response: {response}")
+
+        return response
+
+    def _has_error_response(self, response: Any) -> bool:
+        """Check if the response contains an error.
+
+        Errors can be in two places:
+        1. Top-level: response.get("error") - from DataAgent.chat() exception handling
+        2. Nested: response["messages"][-1].additional_kwargs.get("error") - from FlexAgent planner
+        """
+        if not isinstance(response, dict):
+            return False
+
+        if "error" in response:
+            return True
+
+        if response.get("messages"):
+            last_msg = response["messages"][-1]
+            if isinstance(last_msg, dict):
+                return bool(last_msg.get("additional_kwargs", {}).get("error"))
+            return bool(getattr(last_msg, "additional_kwargs", {}).get("error"))
+
+        return False
+
+    def _extract_error_message(self, response: Any) -> str:
+        """Extract the error message from a response."""
+        if not isinstance(response, dict):
+            return str(response)
+
+        if "error" in response:
+            return str(response.get("error"))
+
+        if response.get("messages"):
+            last_msg = response["messages"][-1]
+            if isinstance(last_msg, dict):
+                return str(last_msg.get("content", str(response)))
+            return str(getattr(last_msg, "content", str(response)))
+
+        return str(response)
+
+    async def _get_session_lock(self, session_id: str) -> asyncio.Lock:
+        """Get or create a per-session Lock for FIFO task serialization."""
+        async with self._session_locks_guard:
+            lock = self._session_locks.get(session_id)
+            if lock is None:
+                lock = asyncio.Lock()
+                self._session_locks[session_id] = lock
+            return lock
+
+
+def _extract_text_from_context(context: RequestContext) -> str:
+    """Extract text content from the request context message."""
+    # RequestContext.get_user_input() already handles text extraction
+    user_text = context.get_user_input()
+    if user_text:
+        return user_text
+
+    # Fallback: extract from context.message directly
+    message = context.message
+    if message is None:
+        return ""
+
+    parts_text = []
+    for part in message.parts:
+        if part.text:
+            parts_text.append(part.text)
+        elif part.data:
+            try:
+                import json
+
+                from google.protobuf.json_format import MessageToDict
+
+                data_dict = MessageToDict(part.data)
+                parts_text.append(json.dumps(data_dict, ensure_ascii=False))
+            except Exception:
+                parts_text.append(str(part.data))
+
+    return "".join(parts_text).strip()
+
+
+def _extract_final_answer(response: Any) -> str:
+    """Extract the final answer text from a DataAgent response."""
+    if isinstance(response, dict):
+        for key in ("final_answer", "answer", "result", "response"):
+            if key in response:
+                val = response[key]
+                if isinstance(val, str):
+                    return val
+                return str(val)
+
+        messages = response.get("messages", [])
+        if messages:
+            last_msg = messages[-1]
+            content = getattr(last_msg, "content", str(last_msg))
+            return str(content)
+
+        return str(response)
+    return str(response)

dataagent/a2a_server/server.py

@@ -0,0 +1,169 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Assemble and run an A2A 1.0 server for a DataAgent."""
+
+from a2a.server.agent_execution import AgentExecutor
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.request_handlers.response_helpers import agent_card_to_dict
+from a2a.server.routes import (
+    create_jsonrpc_routes,
+    create_rest_routes,
+)
+from a2a.server.tasks import InMemoryTaskStore
+from a2a.types.a2a_pb2 import AgentCard
+from a2a.utils.constants import AGENT_CARD_WELL_KNOWN_PATH
+from starlette.applications import Starlette
+from starlette.middleware import Middleware
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.routing import Route
+
+
+def _make_bearer_middleware(token: str, exclude_paths: tuple[str, ...]):
+    """Create a Bearer token authentication middleware."""
+    from starlette.middleware.base import BaseHTTPMiddleware
+
+    class BearerAuthMiddleware(BaseHTTPMiddleware):
+        """Validate Bearer token on every request, skipping excluded paths."""
+
+        async def dispatch(self, request: Request, call_next):
+            """Check Authorization header against the expected token.
+
+            Returns 401 if the token is missing or invalid, otherwise
+            forwards the request to the next middleware/route handler.
+            """
+            if request.url.path in exclude_paths:
+                return await call_next(request)
+
+            auth = request.headers.get("Authorization", "")
+            if not auth.startswith("Bearer ") or auth[7:] != token:
+                return JSONResponse(
+                    status_code=401,
+                    content={"detail": "Unauthorized: invalid or missing Bearer token"},
+                )
+            return await call_next(request)
+
+    return Middleware(BearerAuthMiddleware)
+
+
+def _create_dynamic_agent_card_routes(
+    agent_card: AgentCard,
+    jsonrpc_path: str,
+    rest_path: str,
+) -> list:
+    """Create agent card routes that use the requestʼs Host header for interface URLs.
+
+    Clients from any network can discover the correct URLs instead of always
+    receiving 127.0.0.1.
+    """
+
+    async def _get_agent_card(request: Request) -> JSONResponse:
+        base_card = agent_card_to_dict(agent_card)
+        host = request.headers.get("Host", "127.0.0.1:9999")
+        scheme = request.headers.get("X-Forwarded-Proto", request.url.scheme or "http")
+        base_url = f"{scheme}://{host}"
+
+        for iface in base_card.get("supportedInterfaces", []):
+            protocol = iface.get("protocolBinding", "")
+            if protocol == "JSONRPC":
+                iface["url"] = f"{base_url}{jsonrpc_path}"
+            elif protocol == "HTTP+JSON":
+                iface["url"] = f"{base_url}{rest_path}"
+
+        return JSONResponse(base_card)
+
+    return [
+        Route(path=AGENT_CARD_WELL_KNOWN_PATH, endpoint=_get_agent_card, methods=["GET"]),
+    ]
+
+
+def create_a2a_server(
+    agent_card: AgentCard,
+    executor: AgentExecutor,
+    jsonrpc_path: str = "/a2a/jsonrpc",
+    rest_path: str = "/a2a/rest",
+    enable_v0_3_compat: bool = False,
+    auth_token: str | None = None,
+    auth_exclude_paths: tuple[str, ...] = ("/.well-known/agent-card.json",),
+) -> Starlette:
+    """Create a Starlette application with A2A 1.0 routes.
+
+    Args:
+        agent_card: The A2A AgentCard for this server.
+        executor: The AgentExecutor that handles requests.
+        jsonrpc_path: JSON-RPC route path (default: /a2a/jsonrpc).
+        rest_path: REST route path (default: /a2a/rest).
+        enable_v0_3_compat: Enable backward compatibility with A2A v0.3 clients.
+        auth_token: Bearer token for authentication. If None, no auth is required.
+        auth_exclude_paths: Paths to exclude from auth checks (default: AgentCard discovery).
+
+    Returns:
+        A configured Starlette application.
+    """
+    task_store = InMemoryTaskStore()
+
+    request_handler = DefaultRequestHandler(
+        agent_executor=executor,
+        task_store=task_store,
+        agent_card=agent_card,
+    )
+
+    routes: list = []
+    routes.extend(_create_dynamic_agent_card_routes(agent_card, jsonrpc_path, rest_path))
+    routes.extend(create_jsonrpc_routes(request_handler, rpc_url=jsonrpc_path, enable_v0_3_compat=enable_v0_3_compat))
+    routes.extend(create_rest_routes(request_handler, path_prefix=rest_path, enable_v0_3_compat=enable_v0_3_compat))
+
+    middleware: list = []
+    if auth_token:
+        middleware.append(_make_bearer_middleware(auth_token, exclude_paths=auth_exclude_paths))
+
+    app = Starlette(routes=routes, middleware=middleware)
+    return app
+
+
+def run_a2a_server(
+    agent_card: AgentCard,
+    executor: AgentExecutor,
+    host: str = "",
+    port: int = 9999,
+    jsonrpc_path: str = "/a2a/jsonrpc",
+    rest_path: str = "/a2a/rest",
+    auth_token: str | None = None,
+    auth_exclude_paths: tuple[str, ...] = ("/.well-known/agent-card.json",),
+) -> None:
+    """Run the A2A server using uvicorn.
+
+    Args:
+        agent_card: The A2A AgentCard.
+        executor: The AgentExecutor.
+        host: Server host.
+        port: Server port.
+        jsonrpc_path: JSON-RPC route path.
+        rest_path: REST route path.
+        auth_token: Bearer token for authentication. If None, no auth is required.
+        auth_exclude_paths: Paths excluded from auth checks.
+    """
+    import uvicorn
+
+    if not host:
+        raise ValueError("Host must be specified")
+    app = create_a2a_server(
+        agent_card=agent_card,
+        executor=executor,
+        jsonrpc_path=jsonrpc_path,
+        rest_path=rest_path,
+        auth_token=auth_token,
+        auth_exclude_paths=auth_exclude_paths,
+    )
+
+    uvicorn.run(app, host=host, port=port, log_level="info")

dataagent/a2a_server/task_store.py

@@ -0,0 +1,29 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Task store utilities for the A2A server.
+
+For MVP we use the in-memory store from a2a-sdk.
+For production, a PostgreSQL-backed store can be implemented here
+using a2a-sdk[sql] extras.
+"""
+
+from a2a.server.tasks import InMemoryTaskStore
+
+
+def create_task_store() -> InMemoryTaskStore:
+    """Create an in-memory task store (MVP).
+
+    Returns:
+        An InMemoryTaskStore instance.
+    """
+    return InMemoryTaskStore()

dataagent/actions/__init__.py

@@ -0,0 +1,12 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================