feat: support extra headers for model (#103)

Author: yaozheng-fang

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"

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
@@ -73,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."""
 
@@ -96,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

veadk/consts.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from veadk.version import VERSION
+
 DEFAULT_MODEL_AGENT_NAME = "doubao-seed-1-6-250615"
 DEFALUT_MODEL_AGENT_PROVIDER = "openai"
 DEFAULT_MODEL_AGENT_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
+DEFAULT_MODEL_EXTRA_HEADERS = {"veadk-source": "veadk", "veadk-version": VERSION}

veadk/tracing/telemetry/telemetry.py

@@ -64,7 +64,7 @@ def set_common_attributes(
         )
         return
 
-    if isinstance(invocation_context.agent, Agent):
+    if isinstance(invocation_context.agent, Agent) and invocation_context.agent.tracers:
         try:
             from veadk.tracing.telemetry.opentelemetry_tracer import OpentelemetryTracer