diff --git a/veadk/agent.py b/veadk/agent.py index 7d947252..b3120e3f 100644 --- a/veadk/agent.py +++ b/veadk/agent.py @@ -74,8 +74,8 @@ 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.""" + model_extra_config: dict = Field(default_factory=dict) + """The extra config to include in the model requests.""" tools: list[ToolUnion] = [] """The tools provided to agent.""" @@ -101,17 +101,21 @@ class Agent(LlmAgent): def model_post_init(self, __context: Any) -> None: super().model_post_init(None) # for sub_agents init - self.model_extra_headers |= DEFAULT_MODEL_EXTRA_HEADERS + # add model request source (veadk) in extra headers + if self.model_extra_config and "extra_headers" in self.model_extra_config: + self.model_extra_config["extra_headers"] |= DEFAULT_MODEL_EXTRA_HEADERS + else: + self.model_extra_config["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, + **self.model_extra_config, ) logger.debug( - f"LiteLLM client created with extra headers: {self.model_extra_headers}" + f"LiteLLM client created with config: {self.model_extra_config}" ) else: logger.warning( @@ -133,7 +137,7 @@ def model_post_init(self, __context: Any) -> None: logger.info(f"{self.__class__.__name__} `{self.name}` init done.") logger.debug( - f"Agent: {self.model_dump(include={'name', 'model_name', 'model_api_base', 'tools', 'serve_url'})}" + f"Agent: {self.model_dump(include={'name', 'model_name', 'model_api_base', 'tools'})}" ) async def _run( diff --git a/veadk/tracing/telemetry/attributes/extractors/tool_attributes_extractors.py b/veadk/tracing/telemetry/attributes/extractors/tool_attributes_extractors.py index 1ebb5d8a..388646a9 100644 --- a/veadk/tracing/telemetry/attributes/extractors/tool_attributes_extractors.py +++ b/veadk/tracing/telemetry/attributes/extractors/tool_attributes_extractors.py @@ -26,7 +26,6 @@ def tool_gen_ai_operation_name(params: ToolAttributesParams) -> ExtractorRespons def tool_gen_ai_tool_message(params: ToolAttributesParams) -> ExtractorResponse: tool_input = { - "id": "123", "role": "tool", "content": json.dumps( {