Merge remote-tracking branch 'upstream/main' into feature/toshandler

Author: ZQlQZ

.github/workflows/publish-tag-to-pypi.yaml

@@ -16,7 +16,7 @@ jobs:
         with:
           python-version: "3.x"
       - name: Install build tool
-        run: python3 -m pip install build --user
+        run: python3 -m pip install build --user --upgrade
       - name: Build wheel and source distribution
         run: python3 -m build
       - name: Upload distributables as artifact
@@ -39,7 +39,7 @@ jobs:
           path: dist/
 
       - name: Install Twine
-        run: python3 -m pip install --user twine
+        run: python3 -m pip install --user --upgrade twine packaging
 
       - name: Publish to PyPI
         env:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "veadk-python"
-version = "0.2.2"
+version = "0.2.4"
 description = "Volcengine agent development kit, integrations with Volcengine cloud services."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -18,19 +18,22 @@ dependencies = [
     "google-adk>=1.10.0", # For basic agent architecture
     "litellm>=1.74.3", # For model inference
     "loguru>=0.7.3", # For better logging
-    "openinference-instrumentation-google-adk>=0.1.1", # For OpenInference instrumentation
     "opentelemetry-exporter-otlp>=1.35.0",
     "opentelemetry-instrumentation-logging>=0.56b0",
     "wrapt>=1.17.2", # For patching built-in functions
-    "openai<1.100" # For fix https://github.com/BerriAI/litellm/issues/13710
+    "openai<1.100", # For fix https://github.com/BerriAI/litellm/issues/13710
+    "volcengine-python-sdk==4.0.3", # For Volcengine API
+    "agent-pilot-sdk>=0.0.9",    # Prompt optimization by Volcengine AgentPilot/PromptPilot toolkits
+    "fastmcp>=2.11.3",    # For running MCP
+    "cookiecutter>=2.6.0", # For cloud deploy
+    "opensearch-py==2.8.0"    # For OpenSearch database
 ]
 
 [project.scripts]
 veadk = "veadk.cli.cli:veadk"
 
 [project.optional-dependencies]
 database = [
-    "opensearch-py==2.8.0",         # For OpenSearch database
     "redis>=6.2.0",                 # For Redis database
     "pymysql>=1.1.1",               # For MySQL database
     "volcengine>=1.0.193",          # For Viking DB
@@ -41,12 +44,7 @@ eval = [
     "deepeval>=3.2.6",              # For DeepEval-based evaluation
     "google-adk[eval]",             # For Google ADK-based evaluation
 ]
-cli = [
-    "volcengine-python-sdk==4.0.3", # For Volcengine API
-    "agent-pilot-sdk>=0.0.9",       # Prompt optimization by Volcengine AgentPilot/PromptPilot toolkits
-    "fastmcp>=2.11.3",              # For running MCP
-    "cookiecutter>=2.6.0", # For cloud deploy
-]
+cli = []
 dev = [
     "pre-commit>=4.2.0",            # Format checking
     "pytest>=8.4.1",
@@ -77,4 +75,4 @@ include-package-data = true
 [tool.ruff]
 exclude = [
     "veadk/integrations/ve_faas/template/*"
-]
+]

tests/test_agent.py

@@ -15,6 +15,7 @@
 from google.adk.tools import load_memory
 
 from veadk import Agent
+from veadk.consts import DEFAULT_MODEL_EXTRA_HEADERS
 from veadk.knowledgebase import KnowledgeBase
 from veadk.memory.long_term_memory import LongTermMemory
 from veadk.tools import load_knowledgebase_tool
@@ -26,11 +27,14 @@ def test_agent():
     long_term_memory = LongTermMemory(backend="local")
     tracer = OpentelemetryTracer()
 
+    model_extra_headers = {"test-header": "test-value"}
+
     agent = Agent(
         model_name="test_model_name",
         model_provider="test_model_provider",
         model_api_key="test_model_api_key",
         model_api_base="test_model_api_base",
+        model_extra_headers=model_extra_headers,
         tools=[],
         sub_agents=[],
         knowledgebase=knowledgebase,
@@ -39,7 +43,10 @@ def test_agent():
         serve_url="",
     )
 
+    model_extra_headers |= DEFAULT_MODEL_EXTRA_HEADERS
+
     assert agent.model.model == f"{agent.model_provider}/{agent.model_name}"
+    assert agent.model_extra_headers == model_extra_headers
 
     assert agent.knowledgebase == knowledgebase
     assert agent.knowledgebase.backend == "local"
@@ -48,7 +55,3 @@ def test_agent():
 
     assert agent.long_term_memory.backend == "local"
     assert load_memory in agent.tools
-
-    assert tracer.tracer_hook_before_model in agent.before_model_callback
-    assert tracer.tracer_hook_after_model in agent.after_model_callback
-    assert tracer.tracer_hook_after_tool in agent.after_tool_callback

tests/test_evaluator.py

@@ -129,7 +129,7 @@ def test_evaluator():
     with open(eval_set_file_path, "w") as f:
         json.dump(EVAL_SET_DATA, f)
 
-    base_evaluator.generate_eval_data(file_path=eval_set_file_path)
+    base_evaluator.build_eval_set(file_path=eval_set_file_path)
 
     assert len(base_evaluator.invocation_list) == 1
     assert len(base_evaluator.invocation_list[0].invocations) == 1
@@ -149,7 +149,7 @@ def test_tracing_file_to_evalset():
     with open(tracing_file_path, "w") as f:
         json.dump(TRACE_SET_DATA, f)
 
-    base_evaluator.generate_eval_data(file_path=tracing_file_path)
+    base_evaluator.build_eval_set(file_path=tracing_file_path)
 
     assert len(base_evaluator.invocation_list) == 1
     assert len(base_evaluator.invocation_list[0].invocations) == 1

tests/test_tracing.py

@@ -82,7 +82,7 @@ async def test_tracing():
     exporters = init_exporters()
     tracer = OpentelemetryTracer(exporters=exporters)
 
-    assert len(tracer.exporters) == 4  # with extra 1 built-in exporters
+    assert len(tracer.exporters) == 3
 
     # TODO: Ensure the tracing provider is set correctly after loading SDK
 
@@ -98,7 +98,7 @@ async def test_tracing_with_global_provider():
     #
     tracer = OpentelemetryTracer(exporters=exporters)
 
-    assert len(tracer.exporters) == 4  # with extra 1 built-in exporters
+    assert len(tracer.exporters) == 3  # with extra 1 built-in exporters
 
 
 @pytest.mark.asyncio
@@ -113,4 +113,4 @@ async def test_tracing_with_apmplus_global_provider():
     tracer = OpentelemetryTracer(exporters=exporters)
 
     # apmplus exporter won't init again
-    assert len(tracer.exporters) == 3  # with extra 1 built-in exporters
+    assert len(tracer.exporters) == 2  # with extra 1 built-in exporters

veadk/agent.py

@@ -31,6 +31,7 @@
     DEFALUT_MODEL_AGENT_PROVIDER,
     DEFAULT_MODEL_AGENT_API_BASE,
     DEFAULT_MODEL_AGENT_NAME,
+    DEFAULT_MODEL_EXTRA_HEADERS,
 )
 from veadk.evaluation import EvalSetRecorder
 from veadk.knowledgebase import KnowledgeBase
@@ -40,6 +41,7 @@
 from veadk.tracing.base_tracer import BaseTracer
 from veadk.utils.logger import get_logger
 from veadk.utils.patches import patch_asyncio
+from veadk.version import VERSION
 
 patch_asyncio()
 logger = get_logger(__name__)
@@ -72,6 +74,9 @@ class Agent(LlmAgent):
     model_api_key: str = Field(default_factory=lambda: getenv("MODEL_AGENT_API_KEY"))
     """The api key of the model for agent running."""
 
+    model_extra_headers: dict = Field(default_factory=dict)
+    """The extra headers to include in the model requests."""
+
     tools: list[ToolUnion] = []
     """The tools provided to agent."""
 
@@ -95,11 +100,23 @@ class Agent(LlmAgent):
 
     def model_post_init(self, __context: Any) -> None:
         super().model_post_init(None)  # for sub_agents init
-        self.model = LiteLlm(
-            model=f"{self.model_provider}/{self.model_name}",
-            api_key=self.model_api_key,
-            api_base=self.model_api_base,
-        )
+
+        self.model_extra_headers |= DEFAULT_MODEL_EXTRA_HEADERS
+
+        if not self.model:
+            self.model = LiteLlm(
+                model=f"{self.model_provider}/{self.model_name}",
+                api_key=self.model_api_key,
+                api_base=self.model_api_base,
+                extra_headers=self.model_extra_headers,
+            )
+            logger.debug(
+                f"LiteLLM client created with extra headers: {self.model_extra_headers}"
+            )
+        else:
+            logger.warning(
+                "You are trying to use your own LiteLLM client, some default request headers may be missing."
+            )
 
         if self.knowledgebase:
             from veadk.tools import load_knowledgebase_tool
@@ -112,9 +129,7 @@ def model_post_init(self, __context: Any) -> None:
 
             self.tools.append(load_memory)
 
-        if self.tracers:
-            for tracer in self.tracers:
-                tracer.do_hooks(self)
+        logger.info(f"VeADK version: {VERSION}")
 
         logger.info(f"{self.__class__.__name__} `{self.name}` init done.")
         logger.debug(
@@ -215,9 +230,6 @@ async def run(
             session_service=session_service,
             memory_service=self.long_term_memory,
         )
-        if getattr(self, "tracers", None):
-            for tracer in self.tracers:
-                tracer.set_app_name(app_name)
 
         logger.info(f"Begin to process prompt {prompt}")
         # run

veadk/agents/loop_agent.py

@@ -19,7 +19,6 @@
 from pydantic import ConfigDict, Field
 from typing_extensions import Any
 
-from veadk.agent import Agent
 from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
 from veadk.tracing.base_tracer import BaseTracer
 from veadk.utils.logger import get_logger
@@ -50,21 +49,7 @@ class LoopAgent(GoogleADKLoopAgent):
     tracers: list[BaseTracer] = []
     """The tracers provided to agent."""
 
-    def set_sub_agents_tracer(self, tracer) -> None:
-        from veadk.agents.parallel_agent import ParallelAgent
-        from veadk.agents.sequential_agent import SequentialAgent
-
-        for sub_agent in self.sub_agents:
-            if isinstance(sub_agent, Agent):
-                tracer.do_hooks(sub_agent)
-            elif isinstance(sub_agent, (SequentialAgent, LoopAgent, ParallelAgent)):
-                sub_agent.set_sub_agents_tracer(tracer)
-
     def model_post_init(self, __context: Any) -> None:
         super().model_post_init(None)  # for sub_agents init
 
-        if self.tracers:
-            for tracer in self.tracers:
-                self.set_sub_agents_tracer(tracer)
-
         logger.info(f"{self.__class__.__name__} `{self.name}` init done.")

veadk/agents/parallel_agent.py

@@ -19,7 +19,6 @@
 from pydantic import ConfigDict, Field
 from typing_extensions import Any
 
-from veadk.agent import Agent
 from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
 from veadk.tracing.base_tracer import BaseTracer
 from veadk.utils.logger import get_logger
@@ -50,24 +49,12 @@ class ParallelAgent(GoogleADKParallelAgent):
     tracers: list[BaseTracer] = []
     """The tracers provided to agent."""
 
-    def set_sub_agents_tracer(self, tracer) -> None:
-        from veadk.agents.loop_agent import LoopAgent
-        from veadk.agents.sequential_agent import SequentialAgent
-
-        for sub_agent in self.sub_agents:
-            if isinstance(sub_agent, Agent):
-                tracer.do_hooks(sub_agent)
-            elif isinstance(sub_agent, (SequentialAgent, LoopAgent, ParallelAgent)):
-                sub_agent.set_sub_agents_tracer(tracer)
-
     def model_post_init(self, __context: Any) -> None:
         super().model_post_init(None)  # for sub_agents init
 
         if self.tracers:
             logger.warning(
                 "Enable tracing in ParallelAgent may cause OpenTelemetry context error. Issue see https://github.com/google/adk-python/issues/1670"
             )
-            for tracer in self.tracers:
-                self.set_sub_agents_tracer(tracer)
 
         logger.info(f"{self.__class__.__name__} `{self.name}` init done.")

veadk/agents/sequential_agent.py

@@ -19,7 +19,6 @@
 from pydantic import ConfigDict, Field
 from typing_extensions import Any
 
-from veadk.agent import Agent
 from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
 from veadk.tracing.base_tracer import BaseTracer
 from veadk.utils.logger import get_logger
@@ -50,21 +49,7 @@ class SequentialAgent(GoogleADKSequentialAgent):
     tracers: list[BaseTracer] = []
     """The tracers provided to agent."""
 
-    def set_sub_agents_tracer(self, tracer) -> None:
-        from veadk.agents.loop_agent import LoopAgent
-        from veadk.agents.parallel_agent import ParallelAgent
-
-        for sub_agent in self.sub_agents:
-            if isinstance(sub_agent, Agent):
-                tracer.do_hooks(sub_agent)
-            elif isinstance(sub_agent, (SequentialAgent, LoopAgent, ParallelAgent)):
-                sub_agent.set_sub_agents_tracer(tracer)
-
     def model_post_init(self, __context: Any) -> None:
         super().model_post_init(None)  # for sub_agents init
 
-        if self.tracers:
-            for tracer in self.tracers:
-                self.set_sub_agents_tracer(tracer)
-
         logger.info(f"{self.__class__.__name__} `{self.name}` init done.")

veadk/cli/cli_web.py

@@ -102,6 +102,32 @@ def _get_memory(
     return short_term_memory, long_term_memory
 
 
+def patch_adkwebserver_disable_openapi():
+    """
+    Monkey patch AdkWebServer.get_fast_api to remove openapi.json route.
+    """
+    import google.adk.cli.adk_web_server
+    from fastapi.routing import APIRoute
+    from starlette.routing import Route
+
+    original_get_fast_api = google.adk.cli.adk_web_server.AdkWebServer.get_fast_api_app
+
+    def wrapped_get_fast_api(self, *args, **kwargs):
+        app = original_get_fast_api(self, *args, **kwargs)
+
+        paths = ["/openapi.json", "/docs", "/redoc"]
+        new_routes = []
+        for route in app.router.routes:
+            if isinstance(route, (APIRoute, Route)) and route.path in paths:
+                continue
+            new_routes.append(route)
+        app.router.routes = new_routes
+
+        return app
+
+    google.adk.cli.adk_web_server.AdkWebServer.get_fast_api_app = wrapped_get_fast_api
+
+
 @click.command()
 @click.option("--host", default="127.0.0.1", help="Host to run the web server on")
 def web(host: str) -> None:
@@ -145,6 +171,7 @@ def init_for_veadk(
     import google.adk.cli.adk_web_server
 
     google.adk.cli.adk_web_server.AdkWebServer.__init__ = init_for_veadk
+    patch_adkwebserver_disable_openapi()
 
     import google.adk.cli.cli_tools_click as cli_tools_click