Skip to content

Commit a9ccd7b

Browse files
committed
chore(config): add model config yaml
1 parent ffb9801 commit a9ccd7b

File tree

6 files changed

+73
-16
lines changed

6 files changed

+73
-16
lines changed

config.yaml.full

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ model:
55
name: doubao-seed-1-6-250615
66
api_base: https://ark.cn-beijing.volces.com/api/v3/
77
api_key:
8+
extra_headers: '{"x-is-encrypted": "true"}'
9+
extra_body: '{"caching": {"type": "enabled"}, "thinking": {"type": "disabled"}, "expire_at": 1}' # expire after 1 hour
810
# [optional] for llm-as-a-judge a evaluation
911
judge:
1012
name: doubao-seed-1-6-250615

config.yaml.simple

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,6 @@ model:
33
provider: openai
44
name: doubao-seed-1-6-250615
55
api_base: https://ark.cn-beijing.volces.com/api/v3/
6-
api_key:
6+
api_key:
7+
extra_headers: '{"x-is-encrypted": "true"}'
8+
extra_body: '{"caching": {"type": "enabled"}, "thinking": {"type": "disabled"}, "expire_at": 1}' # expire after 1 hour

docs/docs/agent.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ Agent 中主要包括如下属性:
1515
| model_provider | str | Agent 中内置模型的提供商,默认从环境变量中获取 |
1616
| model_api_base | str | Agent 中内置模型的 API Base,默认从环境变量中获取 |
1717
| model_api_key | str | Agent 中内置模型的 API Key,默认从环境变量中获取 |
18+
| model_extra_config | dict | Agent 进行模型请求时的额外参数,Key 值为属性名,Value 值为属性值 |
1819
| tools | list | Function call 中的工具列表,既可以是本地工具,也可以是 MCP 工具 |
1920
| sub_agents | list | 子 Agent 列表,用于多 Agent 之间交互 |
2021
| knowledgebase | KnowledgeBase | 知识库,后端支持本地内存(local)和数据库(opensearch、viking、redis、mysql),通常设置为一个能够检索的向量数据库 |

tests/test_agent.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
from google.adk.tools import load_memory
1616

1717
from veadk import Agent
18-
from veadk.consts import DEFAULT_MODEL_EXTRA_HEADERS
18+
from veadk.consts import DEFAULT_MODEL_EXTRA_CONFIG
1919
from veadk.knowledgebase import KnowledgeBase
2020
from veadk.memory.long_term_memory import LongTermMemory
2121
from veadk.tools import load_knowledgebase_tool
@@ -27,14 +27,17 @@ def test_agent():
2727
long_term_memory = LongTermMemory(backend="local")
2828
tracer = OpentelemetryTracer()
2929

30-
model_extra_headers = {"test-header": "test-value"}
30+
extra_config = {
31+
"extra_headers": {"thinking": "test"},
32+
"extra_body": {"content": "test"},
33+
}
3134

3235
agent = Agent(
3336
model_name="test_model_name",
3437
model_provider="test_model_provider",
3538
model_api_key="test_model_api_key",
3639
model_api_base="test_model_api_base",
37-
model_extra_headers=model_extra_headers,
40+
model_extra_config=extra_config,
3841
tools=[],
3942
sub_agents=[],
4043
knowledgebase=knowledgebase,
@@ -43,10 +46,13 @@ def test_agent():
4346
serve_url="",
4447
)
4548

46-
model_extra_headers |= DEFAULT_MODEL_EXTRA_HEADERS
47-
4849
assert agent.model.model == f"{agent.model_provider}/{agent.model_name}"
49-
assert agent.model_extra_headers == model_extra_headers
50+
51+
expected_config = DEFAULT_MODEL_EXTRA_CONFIG.copy()
52+
expected_config["extra_headers"] |= extra_config["extra_headers"]
53+
expected_config["extra_body"] |= extra_config["extra_body"]
54+
55+
assert agent.model_extra_config == expected_config
5056

5157
assert agent.knowledgebase == knowledgebase
5258
assert agent.knowledgebase.backend == "local"

veadk/agent.py

Lines changed: 41 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,10 @@
2828

2929
from veadk.config import getenv
3030
from veadk.consts import (
31-
DEFAULT_MODEL_AGENT_PROVIDER,
3231
DEFAULT_MODEL_AGENT_API_BASE,
3332
DEFAULT_MODEL_AGENT_NAME,
34-
DEFAULT_MODEL_EXTRA_HEADERS,
33+
DEFAULT_MODEL_AGENT_PROVIDER,
34+
DEFAULT_MODEL_EXTRA_CONFIG,
3535
)
3636
from veadk.evaluation import EvalSetRecorder
3737
from veadk.knowledgebase import KnowledgeBase
@@ -47,6 +47,29 @@
4747
logger = get_logger(__name__)
4848

4949

50+
def _parse_user_model_extra_config() -> dict:
51+
import json
52+
53+
extra_headers = getenv("MODEL_AGENT_EXTRA_HEADERS", {}, allow_false_values=True)
54+
extra_body = getenv("MODEL_AGENT_EXTRA_BODY", {}, allow_false_values=True)
55+
56+
if extra_headers:
57+
extra_headers = json.loads(extra_headers)
58+
if extra_body:
59+
extra_body = json.loads(extra_body)
60+
if extra_body.get("expire_at", 0):
61+
import time
62+
63+
extra_body["expire_at"] = (
64+
int(time.time()) + int(extra_body["expire_at"]) * 3600
65+
)
66+
67+
return {
68+
"extra_headers": extra_headers,
69+
"extra_body": extra_body,
70+
}
71+
72+
5073
class Agent(LlmAgent):
5174
"""LLM-based Agent with Volcengine capabilities."""
5275

@@ -74,7 +97,7 @@ class Agent(LlmAgent):
7497
model_api_key: str = Field(default_factory=lambda: getenv("MODEL_AGENT_API_KEY"))
7598
"""The api key of the model for agent running."""
7699

77-
model_extra_config: dict = Field(default_factory=dict)
100+
model_extra_config: dict = Field(default_factory=_parse_user_model_extra_config)
78101
"""The extra config to include in the model requests."""
79102

80103
tools: list[ToolUnion] = []
@@ -101,11 +124,21 @@ class Agent(LlmAgent):
101124
def model_post_init(self, __context: Any) -> None:
102125
super().model_post_init(None) # for sub_agents init
103126

104-
# add model request source (veadk) in extra headers
105-
if self.model_extra_config and "extra_headers" in self.model_extra_config:
106-
self.model_extra_config["extra_headers"] |= DEFAULT_MODEL_EXTRA_HEADERS
107-
else:
108-
self.model_extra_config["extra_headers"] = DEFAULT_MODEL_EXTRA_HEADERS
127+
# combine user model config with VeADK defaults
128+
headers = DEFAULT_MODEL_EXTRA_CONFIG["extra_headers"].copy()
129+
body = DEFAULT_MODEL_EXTRA_CONFIG["extra_body"].copy()
130+
131+
if self.model_extra_config:
132+
user_headers = self.model_extra_config.get("extra_headers", {})
133+
user_body = self.model_extra_config.get("extra_body", {})
134+
135+
headers |= user_headers
136+
body |= user_body
137+
138+
self.model_extra_config |= {
139+
"extra_headers": headers,
140+
"extra_body": body,
141+
}
109142

110143
if not self.model:
111144
self.model = LiteLlm(

veadk/consts.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,25 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import time
16+
1517
from veadk.version import VERSION
1618

1719
DEFAULT_MODEL_AGENT_NAME = "doubao-seed-1-6-250615"
1820
DEFAULT_MODEL_AGENT_PROVIDER = "openai"
1921
DEFAULT_MODEL_AGENT_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
20-
DEFAULT_MODEL_EXTRA_HEADERS = {"veadk-source": "veadk", "veadk-version": VERSION}
22+
DEFAULT_MODEL_EXTRA_CONFIG = {
23+
"extra_headers": {
24+
"x-is-encrypted": "true",
25+
"veadk-source": "veadk",
26+
"veadk-version": VERSION,
27+
},
28+
"extra_body": {
29+
"caching": {"type": "enabled"},
30+
"thinking": {"type": "disabled"},
31+
"expire_at": int(time.time()) + 3600, # expire after 1 hour
32+
},
33+
}
2134

2235
DEFAULT_APMPLUS_OTEL_EXPORTER_ENDPOINT = "http://apmplus-cn-beijing.volces.com:4317"
2336
DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME = "veadk_tracing"

0 commit comments

Comments
 (0)