Merge branch 'main' into feat/tealtiger-governance-middleware

Author: nagasatish007

ag2/tools/search/perplexity.py

@@ -4,7 +4,7 @@
 
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass, field
-from typing import Annotated, Any, Literal, TypeAlias
+from typing import Annotated, Any, Final, Literal, TypeAlias
 
 import httpx
 from perplexity import AsyncPerplexity
@@ -18,6 +18,7 @@
 from ag2.tools.builtin._resolve import resolve_variable
 from ag2.tools.final import Toolkit, tool
 from ag2.tools.final.function_tool import FunctionTool
+from ag2.version import __version__
 
 SonarModel: TypeAlias = Literal[
     "sonar",
@@ -30,6 +31,8 @@
 SearchContextSize: TypeAlias = Literal["low", "medium", "high"]
 RecencyFilter: TypeAlias = Literal["hour", "day", "week", "month", "year"]
 
+_PPLX_INTEGRATION_HEADERS: Final[dict[str, str]] = {"X-Pplx-Integration": f"ag2/{__version__}"}
+
 
 @dataclass(slots=True)
 class PerplexitySearchResult:
@@ -101,6 +104,10 @@ def __init__(
         self._proxy = proxy
         self._verify = verify
         self._timeout = timeout
+        client_kwargs["default_headers"] = {
+            **(client_kwargs.get("default_headers") or {}),
+            **_PPLX_INTEGRATION_HEADERS,
+        }
         self._client_kwargs = client_kwargs
 
         super().__init__(

test/tools/search/test_perplexity.py

@@ -12,7 +12,7 @@
 
 pytest.importorskip("perplexity")
 
-from ag2 import Agent, Context, DataInput, ImageInput, Variable
+from ag2 import Agent, Context, DataInput, ImageInput, Variable, __version__
 from ag2.events import ModelResponse, ToolCallEvent, ToolCallsEvent, ToolResultsEvent
 from ag2.testing import TestConfig, TrackingConfig
 from ag2.tools.search.perplexity import (
@@ -238,7 +238,11 @@ async def test_all_params_forwarded(self) -> None:
     async def test_client_kwargs_forwarded_to_sdk(self) -> None:
         custom_url = "https://custom.perplexity.example"
         route = respx.post(f"{custom_url}/search").mock(return_value=httpx.Response(200, json=_search_response()))
-        toolkit = PerplexitySearchToolkit(api_key="test", base_url=custom_url)
+        toolkit = PerplexitySearchToolkit(
+            api_key="test",
+            base_url=custom_url,
+            default_headers={"X-Trace-Id": "abc-123"},
+        )
 
         agent = Agent(
             "a",
@@ -248,6 +252,8 @@ async def test_client_kwargs_forwarded_to_sdk(self) -> None:
         await agent.ask("search")
 
         assert route.called
+        assert route.calls.last.request.headers["X-Pplx-Integration"] == f"ag2/{__version__}"
+        assert route.calls.last.request.headers["X-Trace-Id"] == "abc-123"
 
     @respx.mock
     async def test_custom_tool_name_in_agent(self) -> None:
@@ -393,7 +399,11 @@ async def test_client_kwargs_forwarded_to_sdk(self) -> None:
         route = respx.post(f"{custom_url}/chat/completions").mock(
             return_value=httpx.Response(200, json=_chat_response())
         )
-        toolkit = PerplexitySearchToolkit(api_key="test", base_url=custom_url)
+        toolkit = PerplexitySearchToolkit(
+            api_key="test",
+            base_url=custom_url,
+            default_headers={"X-Trace-Id": "abc-123"},
+        )
 
         agent = Agent(
             "a",
@@ -403,6 +413,8 @@ async def test_client_kwargs_forwarded_to_sdk(self) -> None:
         await agent.ask("answer")
 
         assert route.called
+        assert route.calls.last.request.headers["X-Pplx-Integration"] == f"ag2/{__version__}"
+        assert route.calls.last.request.headers["X-Trace-Id"] == "abc-123"
 
     @respx.mock
     async def test_custom_tool_name_in_agent(self) -> None: