veadk-python 1.0.2__py3-none-any.whl → 1.0.3__py3-none-any.whl

veadk/a2a/registry_client.py

@@ -22,6 +22,7 @@ import os
 import re
 import time
 import uuid
+from collections.abc import Mapping
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from typing import Any
@@ -30,6 +31,7 @@ from urllib.parse import quote, urlparse, urlunparse
 import requests
 
 from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
+from veadk.utils.auth import VE_TIP_TOKEN_HEADER
 
 DEFAULT_ENDPOINT = "http://volcengineapi.byted.org/"
 DEFAULT_VERSION = "2025-10-30"
@@ -66,6 +68,8 @@ class AgentKitA2ARegistryConfig:
     top_k: int = DEFAULT_TOP_K
     timeout_ms: int = DEFAULT_TIMEOUT_MS
     poll_interval_ms: int = DEFAULT_POLL_INTERVAL_MS
+    upstream_tip_token: str = ""
+    upstream_authorization: str = ""
 
 
 @dataclass(frozen=True)
@@ -97,9 +101,41 @@ def registry_config_from_env() -> AgentKitA2ARegistryConfig:
         poll_interval_ms=_int_env(
             "REGISTRY_POLL_INTERVAL_MS", DEFAULT_POLL_INTERVAL_MS, minimum=100
         ),
+        upstream_tip_token=_first_env(
+            [
+                "REGISTRY_UPSTREAM_TIP_TOKEN",
+                "AGENTKIT_UPSTREAM_TIP_TOKEN",
+                "A2A_REGISTRY_UPSTREAM_TIP_TOKEN",
+                "VE_TIP_TOKEN",
+                "X_VE_TIP_TOKEN",
+                "TIP_TOKEN",
+            ]
+        ),
     )
 
 
+def registry_tip_token_from_headers(headers: Mapping[str, str]) -> str:
+    """Extract an inbound TIP token from HTTP headers."""
+
+    normalized = {str(key).lower(): str(value) for key, value in headers.items()}
+    return normalized.get(VE_TIP_TOKEN_HEADER.lower(), "").strip()
+
+
+def registry_authorization_from_headers(headers: Mapping[str, str]) -> str:
+    """Extract a bearer Authorization header for downstream A2A calls."""
+
+    normalized = {str(key).lower(): str(value) for key, value in headers.items()}
+    for header_name in (
+        "authorization",
+        "x-forwarded-authorization",
+        "x-original-authorization",
+    ):
+        value = _normalize_bearer_authorization(normalized.get(header_name, ""))
+        if value:
+            return value
+    return ""
+
+
 def search_agent_cards(
     prompt: str,
     top_k: int | None = None,
@@ -250,6 +286,12 @@ def _resolve_config(
                 or DEFAULT_POLL_INTERVAL_MS
             ),
         ),
+        upstream_tip_token=(
+            config.upstream_tip_token or env_config.upstream_tip_token
+        ).strip(),
+        upstream_authorization=_normalize_bearer_authorization(
+            config.upstream_authorization or env_config.upstream_authorization
+        ),
     )
 
 
@@ -534,15 +576,25 @@ def _send_message(
     if task_id:
         message["taskId"] = task_id
 
+    headers: dict[str, str] = {}
     try:
+        headers = _agent_auth_headers(card, config)
         return _a2a_jsonrpc(
             card["url"],
             "message/send",
             {"message": message, "configuration": {"blocking": False}},
-            _agent_auth_headers(card, config),
+            headers,
             config,
         )
     except RegistryError as exc:
+        if _should_retry_with_m2m(exc, card, config, headers):
+            return _a2a_jsonrpc(
+                card["url"],
+                "message/send",
+                {"message": message, "configuration": {"blocking": False}},
+                _agent_auth_headers(card, config, force_m2m=True),
+                config,
+            )
         if exc.code in {
             "A2A_HTTP_FAILED",
             "A2A_RESPONSE_PARSE_FAILED",
@@ -563,13 +615,25 @@ def _poll_card(
     started: float | None = None,
 ) -> dict[str, Any]:
     started = started or time.monotonic()
-    a2a_result = _a2a_jsonrpc(
-        card["url"],
-        "tasks/get",
-        {"id": task_id.strip(), "historyLength": max(0, int(history_length))},
-        _agent_auth_headers(card, config),
-        config,
-    )
+    headers = _agent_auth_headers(card, config)
+    try:
+        a2a_result = _a2a_jsonrpc(
+            card["url"],
+            "tasks/get",
+            {"id": task_id.strip(), "historyLength": max(0, int(history_length))},
+            headers,
+            config,
+        )
+    except RegistryError as exc:
+        if not _should_retry_with_m2m(exc, card, config, headers):
+            raise
+        a2a_result = _a2a_jsonrpc(
+            card["url"],
+            "tasks/get",
+            {"id": task_id.strip(), "historyLength": max(0, int(history_length))},
+            _agent_auth_headers(card, config, force_m2m=True),
+            config,
+        )
     state = _task_state(a2a_result)
     is_terminal = state in TERMINAL_STATES
     payload: dict[str, Any] = {
@@ -815,11 +879,18 @@ def _sanitize_get_agent_result(
 
 
 def _agent_auth_headers(
-    card: dict[str, Any], config: AgentKitA2ARegistryConfig | None = None
+    card: dict[str, Any],
+    config: AgentKitA2ARegistryConfig | None = None,
+    *,
+    force_m2m: bool = False,
 ) -> dict[str, str]:
+    resolved_config = _resolve_config(config)
     security = card.get("security") or []
     schemes = card.get("securitySchemes") or {}
     headers: dict[str, str] = {}
+    upstream_authorization = _normalize_bearer_authorization(
+        resolved_config.upstream_authorization
+    )
 
     for requirement in security:
         if not isinstance(requirement, dict):
@@ -837,9 +908,17 @@ def _agent_auth_headers(
                 if isinstance(token, str) and token:
                     headers[header_name] = token
             elif scheme_type == "oauth2":
-                headers["Authorization"] = "Bearer " + _oauth2_client_credentials_token(
-                    scheme, _resolve_config(config)
-                )
+                if upstream_authorization and not force_m2m:
+                    headers["Authorization"] = upstream_authorization
+                else:
+                    headers["Authorization"] = (
+                        "Bearer "
+                        + _oauth2_client_credentials_token(scheme, resolved_config)
+                    )
+
+    tip_token = resolved_config.upstream_tip_token
+    if tip_token:
+        headers[VE_TIP_TOKEN_HEADER] = tip_token
 
     if security and not headers:
         raise RegistryError(
@@ -849,6 +928,38 @@ def _agent_auth_headers(
     return headers
 
 
+def _should_retry_with_m2m(
+    exc: RegistryError,
+    card: dict[str, Any],
+    config: AgentKitA2ARegistryConfig | None,
+    headers: dict[str, str],
+) -> bool:
+    if exc.code != "A2A_HTTP_FAILED" or exc.diagnostics.get("status_code") != 401:
+        return False
+
+    resolved_config = _resolve_config(config)
+    upstream_authorization = _normalize_bearer_authorization(
+        resolved_config.upstream_authorization
+    )
+    if not upstream_authorization:
+        return False
+    if headers.get("Authorization") != upstream_authorization:
+        return False
+    return _has_oauth2_security(card)
+
+
+def _has_oauth2_security(card: dict[str, Any]) -> bool:
+    schemes = card.get("securitySchemes") or {}
+    for requirement in card.get("security") or []:
+        if not isinstance(requirement, dict):
+            continue
+        for scheme_name in requirement:
+            scheme = schemes.get(scheme_name) or {}
+            if str(scheme.get("type") or "").lower() == "oauth2":
+                return True
+    return False
+
+
 def _oauth2_client_credentials_token(
     scheme: dict[str, Any], config: AgentKitA2ARegistryConfig
 ) -> str:
@@ -957,6 +1068,16 @@ def _clean_config_url(value: Any) -> str:
     return cleaned
 
 
+def _normalize_bearer_authorization(value: Any) -> str:
+    cleaned = _clean_config_url(value)
+    if not cleaned:
+        return ""
+    parts = cleaned.split(None, 1)
+    if len(parts) == 2 and parts[0].lower() == "bearer" and parts[1].strip():
+        return f"Bearer {parts[1].strip()}"
+    return ""
+
+
 def _user_pool_id_from_token_url(token_url: str) -> str:
     host = urlparse(token_url).netloc
     match = re.match(r"userpool-([^.]+)\.userpool\.auth\.id\.", host)

veadk/cloud/harness_app/app.py

@@ -62,6 +62,10 @@ from google.adk.utils.context_utils import Aclosing
 from typing_extensions import override
 
 from veadk import Agent
+from veadk.a2a.registry_client import (
+    registry_authorization_from_headers,
+    registry_tip_token_from_headers,
+)
 from veadk.a2a.utils.agent_to_a2a import to_a2a
 from veadk.cloud.harness_app.agent import agent, short_term_memory
 from veadk.cloud.harness_app.harness_plugins import (
@@ -229,6 +233,8 @@ class HarnessApp:
             )
 
             try:
+                tip_token = registry_tip_token_from_headers(http_request.headers)
+                auth_header = registry_authorization_from_headers(http_request.headers)
                 header_plugins = build_harness_plugins_from_headers(
                     http_request.headers
                 )
@@ -269,6 +275,8 @@ class HarnessApp:
                             request.prompt,
                             request.harness,
                             download_dir=Path(work_dir),
+                            registry_tip_token=tip_token,
+                            registry_authorization=auth_header,
                         )
                         runner = Runner(
                             agent=agent,
@@ -283,11 +291,15 @@ class HarnessApp:
                             run_config=run_config,
                         )
                 elif needs_scoped_runner:
-                    run_agent = (
-                        spawn_harness_run_agent(self.agent, request.prompt)
-                        if has_registry
-                        else self.agent
-                    )
+                    if has_registry:
+                        run_agent = spawn_harness_run_agent(
+                            self.agent,
+                            request.prompt,
+                            registry_tip_token=tip_token,
+                            registry_authorization=auth_header,
+                        )
+                    else:
+                        run_agent = self.agent
                     runner = Runner(
                         agent=run_agent,
                         short_term_memory=self.short_term_memory,
@@ -359,7 +371,7 @@ class HarnessApp:
                 break
 
         @self.app.post("/run_sse")
-        async def run_sse(req: HarnessRunAgentRequest):
+        async def run_sse(req: HarnessRunAgentRequest, http_request: Request):
             if (
                 req.harness is None
                 and not has_a2a_registry_config(self.agent)
@@ -367,8 +379,11 @@ class HarnessApp:
             ):
                 # No override -> exactly ADK's default /run_sse.
                 return await adk_run_sse(req)
+            tip_token = registry_tip_token_from_headers(http_request.headers)
+            auth_header = registry_authorization_from_headers(http_request.headers)
             return StreamingResponse(
-                self._run_sse_events(req), media_type="text/event-stream"
+                self._run_sse_events(req, tip_token, auth_header),
+                media_type="text/event-stream",
             )
 
         # Move ours to the front so it wins (Starlette matches the first route),
@@ -381,7 +396,12 @@ class HarnessApp:
                 routes.insert(0, routes.pop(i))
                 break
 
-    async def _run_sse_events(self, req: "HarnessRunAgentRequest"):
+    async def _run_sse_events(
+        self,
+        req: "HarnessRunAgentRequest",
+        tip_token: str = "",
+        auth_header: str = "",
+    ):
         """Yield SSE ``data:`` lines for a run, spawning the agent on override."""
         run_config = RunConfig(
             streaming_mode=StreamingMode.SSE if req.streaming else StreamingMode.NONE
@@ -400,13 +420,20 @@ class HarnessApp:
                         prompt,
                         req.harness,
                         download_dir=Path(work_dir_ctx.name),
+                        registry_tip_token=tip_token,
+                        registry_authorization=auth_header,
                     )
                 except (SkillLoadError, ToolLoadError) as e:
                     logger.error(f"Once-time override failed to load: {e}")
                     yield f"data: {json.dumps({'error': str(e)})}\n\n"
                     return
             elif has_a2a_registry_config(self.agent):
-                agent = spawn_harness_run_agent(self.agent, prompt)
+                agent = spawn_harness_run_agent(
+                    self.agent,
+                    prompt,
+                    registry_tip_token=tip_token,
+                    registry_authorization=auth_header,
+                )
             else:
                 agent = self.agent
 

veadk/cloud/harness_app/utils.py

@@ -527,6 +527,30 @@ def _apply_registry_overrides(
     setattr(agent, _REGISTRY_CONFIG_ATTR, overridden_config)
 
 
+def _apply_registry_request_auth(
+    agent: Agent, tip_token: str = "", authorization: str = ""
+) -> None:
+    cleaned_tip_token = (tip_token or "").strip()
+    cleaned_authorization = (authorization or "").strip()
+    if not cleaned_tip_token and not cleaned_authorization:
+        return
+
+    from veadk.tools.builtin_tools.a2a_registry import build_a2a_registry_tools
+
+    config = getattr(agent, _REGISTRY_CONFIG_ATTR, None)
+    if config is None:
+        return
+
+    updated_config = replace(
+        config,
+        upstream_tip_token=cleaned_tip_token or config.upstream_tip_token,
+        upstream_authorization=cleaned_authorization or config.upstream_authorization,
+    )
+    _remove_a2a_registry_tools(agent)
+    agent.tools.extend(build_a2a_registry_tools(updated_config))
+    setattr(agent, _REGISTRY_CONFIG_ATTR, updated_config)
+
+
 def has_a2a_registry_config(agent: Agent) -> bool:
     """Return whether ``agent`` has an AgentKit A2A registry configured."""
 
@@ -604,6 +628,8 @@ def spawn_harness_run_agent(
     prompt: str,
     overrides: HarnessOverrides | None = None,
     download_dir: Path | None = None,
+    registry_tip_token: str = "",
+    registry_authorization: str = "",
 ) -> Agent:
     """Clone a harness agent for one run and attach per-turn dynamic tools."""
 
@@ -612,5 +638,6 @@ def spawn_harness_run_agent(
     else:
         cloned = base_agent.clone(update={})
 
+    _apply_registry_request_auth(cloned, registry_tip_token, registry_authorization)
     _add_dynamic_a2a_agent_tools(cloned, prompt)
     return cloned

veadk/config.py

@@ -21,6 +21,7 @@ from pydantic import BaseModel, Field
 from veadk.configs.auth_configs import VeIdentityConfig
 from veadk.configs.model_configs import RealtimeModelConfig
 from veadk.configs.database_configs import (
+    MilvusConfig,
     MysqlConfig,
     OpensearchConfig,
     RedisConfig,
@@ -79,6 +80,7 @@ class VeADKConfig(BaseModel):
     opensearch: OpensearchConfig = Field(default_factory=OpensearchConfig)
     mysql: MysqlConfig = Field(default_factory=MysqlConfig)
     redis: RedisConfig = Field(default_factory=RedisConfig)
+    milvus: MilvusConfig = Field(default_factory=MilvusConfig)
     viking_knowledgebase: VikingKnowledgebaseConfig = Field(
         default_factory=VikingKnowledgebaseConfig
     )

veadk/configs/database_configs.py

@@ -15,6 +15,7 @@
 import os
 from functools import cached_property
 
+from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from veadk.consts import DEFAULT_TOS_BUCKET_NAME
@@ -95,6 +96,26 @@ class RedisConfig(BaseSettings):
     """STS token for Redis auth, not supported yet."""
 
 
+class MilvusConfig(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="DATABASE_MILVUS_")
+
+    uri: str = ""
+
+    token: str = ""
+
+    user: str = ""
+
+    password: str = ""
+
+    db_name: str = "default"
+
+    overwrite: bool = False
+
+    timeout: float | None = None
+
+    output_fields: list[str] | str = Field(default_factory=list)
+
+
 class Mem0Config(BaseSettings):
     model_config = SettingsConfigDict(env_prefix="DATABASE_MEM0_")
 
@@ -108,6 +129,16 @@ class Mem0Config(BaseSettings):
     base_url: str = ""  # "https://api.mem0.ai/v1"
 
 
+class OpenVikingConfig(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="DATABASE_OPENVIKING_")
+
+    url: str = ""
+    """OpenViking server base URL"""
+
+    api_key: str = ""
+    """OpenViking service owner user key"""
+
+
 class VikingKnowledgebaseConfig(BaseSettings):
     model_config = SettingsConfigDict(env_prefix="DATABASE_VIKING_")
 

veadk/knowledgebase/backends/milvus_backend.py

@@ -0,0 +1,201 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# 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.
+
+import json
+import re
+from typing import Any
+
+from pydantic import Field
+from typing_extensions import override
+
+import veadk.config  # noqa E401
+from veadk.configs.database_configs import MilvusConfig
+from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig
+from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
+
+
+class MilvusKnowledgeBackend(BaseKnowledgebaseBackend):
+    """Milvus-based backend for knowledgebase.
+
+    Milvus backend stores embedded chunks in a Milvus collection through
+    LlamaIndex's Milvus vector store. ``index`` maps directly to the Milvus
+    collection name.
+    """
+
+    milvus_config: MilvusConfig = Field(default_factory=MilvusConfig)
+    """Milvus connection config."""
+
+    embedding_config: EmbeddingModelConfig | NormalEmbeddingModelConfig = Field(
+        default_factory=EmbeddingModelConfig
+    )
+    """Embedding model configs."""
+
+    def model_post_init(self, __context: Any) -> None:
+        self.precheck_index_naming()
+        self._precheck_milvus_uri()
+
+        try:
+            from llama_index.core import StorageContext, VectorStoreIndex
+            from llama_index.vector_stores.milvus import MilvusVectorStore
+            from veadk.models.ark_embedding import create_embedding_model
+        except ImportError as e:
+            raise ImportError(
+                "Please install VeADK extensions\npip install veadk-python[extensions]"
+            ) from e
+
+        self._embed_model = create_embedding_model(
+            model_name=self.embedding_config.name,
+            api_key=self.embedding_config.api_key,
+            api_base=self.embedding_config.api_base,
+        )
+
+        vector_store_kwargs: dict[str, Any] = {
+            "uri": self.milvus_config.uri,
+            "collection_name": self.index,
+            "dim": self.embedding_config.dim,
+            "overwrite": self.milvus_config.overwrite,
+        }
+
+        token = self._resolve_token()
+        if token:
+            vector_store_kwargs["token"] = token
+
+        if self.milvus_config.db_name:
+            vector_store_kwargs["db_name"] = self.milvus_config.db_name
+
+        if self.milvus_config.timeout is not None:
+            vector_store_kwargs["timeout"] = self.milvus_config.timeout
+
+        output_fields = self._resolve_output_fields()
+        if output_fields:
+            vector_store_kwargs["output_fields"] = output_fields
+
+        self._vector_store = MilvusVectorStore(**vector_store_kwargs)
+        self._storage_context = StorageContext.from_defaults(
+            vector_store=self._vector_store
+        )
+        self._vector_index = VectorStoreIndex(
+            nodes=[],
+            storage_context=self._storage_context,
+            embed_model=self._embed_model,
+        )
+
+    @override
+    def precheck_index_naming(self) -> None:
+        if not isinstance(self.index, str) or not self.index:
+            raise ValueError("Milvus collection name must not be empty.")
+        if len(self.index) > 255:
+            raise ValueError("Milvus collection name is too long.")
+        if not re.fullmatch(r"^[A-Za-z_][A-Za-z0-9_]*$", self.index):
+            raise ValueError(
+                "Milvus collection name must start with a letter or underscore "
+                "and contain only letters, numbers, and underscores."
+            )
+
+    def _precheck_milvus_uri(self) -> None:
+        if not self.milvus_config.uri:
+            raise ValueError(
+                "Milvus uri must be configured via DATABASE_MILVUS_URI or "
+                "MilvusConfig(uri=...)."
+            )
+
+    @override
+    def add_from_directory(self, directory: str) -> bool:
+        from llama_index.core import SimpleDirectoryReader
+
+        documents = SimpleDirectoryReader(input_dir=directory).load_data()
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def add_from_files(self, files: list[str]) -> bool:
+        from llama_index.core import SimpleDirectoryReader
+
+        documents = SimpleDirectoryReader(input_files=files).load_data()
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def add_from_text(self, text: str | list[str]) -> bool:
+        from llama_index.core import Document
+
+        if isinstance(text, str):
+            documents = [Document(text=text)]
+        else:
+            documents = [Document(text=t) for t in text]
+        nodes = self._split_documents(documents)
+        self._vector_index.insert_nodes(nodes)
+        return True
+
+    @override
+    def search(self, query: str, top_k: int = 5) -> list[str]:
+        self._ensure_collection_loaded()
+        _retriever = self._vector_index.as_retriever(similarity_top_k=top_k)
+        retrieved_nodes = _retriever.retrieve(query)
+        return [node.text for node in retrieved_nodes]
+
+    def _resolve_token(self) -> str | None:
+        if self.milvus_config.token:
+            return self.milvus_config.token
+        if self.milvus_config.user and self.milvus_config.password:
+            return f"{self.milvus_config.user}:{self.milvus_config.password}"
+        return None
+
+    def _resolve_output_fields(self) -> list[str]:
+        output_fields = self.milvus_config.output_fields
+        if isinstance(output_fields, str):
+            output_fields = output_fields.strip()
+            if not output_fields:
+                return []
+            if output_fields.startswith("["):
+                parsed_output_fields = json.loads(output_fields)
+                if not isinstance(parsed_output_fields, list) or not all(
+                    isinstance(field, str) for field in parsed_output_fields
+                ):
+                    raise ValueError("Milvus output_fields must be a list of strings.")
+                output_fields = parsed_output_fields
+            else:
+                output_fields = output_fields.split(",")
+
+        return [field.strip() for field in output_fields if field.strip()]
+
+    def _ensure_collection_loaded(self) -> None:
+        vector_store = getattr(self, "_vector_store", None)
+        client = getattr(vector_store, "client", None)
+        collection_name = getattr(vector_store, "collection_name", self.index)
+        if client is None or not collection_name:
+            return
+
+        try:
+            load_state = client.get_load_state(collection_name=collection_name)
+        except Exception:
+            client.load_collection(collection_name=collection_name)
+            return
+
+        state = load_state.get("state") if isinstance(load_state, dict) else load_state
+        if "loaded" not in str(state).lower():
+            client.load_collection(collection_name=collection_name)
+
+    def _split_documents(self, documents: list[Any]) -> list[Any]:
+        """Split document into chunks."""
+        from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+
+        nodes = []
+        for document in documents:
+            splitter = get_llama_index_splitter(document.metadata.get("file_path", ""))
+            _nodes = splitter.get_nodes_from_documents([document])
+            nodes.extend(_nodes)
+        return nodes