Skip to content

Commit 4f3b53d

Browse files
OhYeeOhYee
committed
feat: 添加凭证支持到模型服务，增强 API 密钥管理
Change-Id: Ib6831f5a3978b3a47843c3c9f22c1b63c2255c32 Signed-off-by: OhYee <oyohyee@oyohyee.com>
1 parent dcfdabe commit 4f3b53d

File tree

4 files changed

+192
-4
lines changed

4 files changed

+192
-4
lines changed

‎agentrun/model/__model_service_async_template.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,15 @@ def model_info(self, config: Optional[Config] = None) -> BaseInfo:
208208
assert self.provider_settings is not None
209209
assert self.provider_settings.base_url is not None
210210

211+
api_key = self.provider_settings.api_key or ""
212+
if not api_key and self.credential_name:
213+
from agentrun.credential import Credential
214+
215+
credential = Credential.get_by_name(
216+
self.credential_name, config=cfg
217+
)
218+
api_key = credential.credential_secret or ""
219+
211220
default_model = (
212221
self.provider_settings.model_names[0]
213222
if self.provider_settings.model_names is not None
@@ -216,7 +225,7 @@ def model_info(self, config: Optional[Config] = None) -> BaseInfo:
216225
)
217226

218227
return BaseInfo(
219-
api_key=self.provider_settings.api_key or "",
228+
api_key=api_key,
220229
base_url=self.provider_settings.base_url,
221230
model=default_model,
222231
headers=cfg.get_headers(),

‎agentrun/model/model_service.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,15 @@ def model_info(self, config: Optional[Config] = None) -> BaseInfo:
379379
assert self.provider_settings is not None
380380
assert self.provider_settings.base_url is not None
381381

382+
api_key = self.provider_settings.api_key or ""
383+
if not api_key and self.credential_name:
384+
from agentrun.credential import Credential
385+
386+
credential = Credential.get_by_name(
387+
self.credential_name, config=cfg
388+
)
389+
api_key = credential.credential_secret or ""
390+
382391
default_model = (
383392
self.provider_settings.model_names[0]
384393
if self.provider_settings.model_names is not None
@@ -387,7 +396,7 @@ def model_info(self, config: Optional[Config] = None) -> BaseInfo:
387396
)
388397

389398
return BaseInfo(
390-
api_key=self.provider_settings.api_key or "",
399+
api_key=api_key,
391400
base_url=self.provider_settings.base_url,
392401
model=default_model,
393402
headers=cfg.get_headers(),

‎tests/e2e/__test_model_async_template.py‎

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@
3939
)
4040
from agentrun.utils.config import Config
4141
from agentrun.utils.exception import (
42-
ClientError,
4342
ResourceAlreadyExistError,
4443
ResourceNotExistError,
4544
)
@@ -264,6 +263,64 @@ async def test_model_service_invoke_async(self, model_service_name: str):
264263

265264
ms.delete()
266265

266+
async def test_model_service_with_credential_async(
267+
self, model_service_name: str
268+
):
269+
# 创建 Credential
270+
from agentrun.credential import (
271+
Credential,
272+
CredentialConfig,
273+
CredentialCreateInput,
274+
)
275+
276+
cr = await Credential.create_async(
277+
CredentialCreateInput(
278+
credential_name=f"{model_service_name}-credential",
279+
enabled=True,
280+
credential_config=CredentialConfig.outbound_llm_api_key(
281+
api_key=api_key,
282+
provider="openai",
283+
),
284+
)
285+
)
286+
287+
# 创建 model service
288+
ms = await ModelService.create_async(
289+
ModelServiceCreateInput(
290+
model_service_name=model_service_name,
291+
description="原始描述",
292+
model_type=ModelType.LLM,
293+
provider="openai",
294+
credential_name=cr.credential_name,
295+
provider_settings=ProviderSettings(
296+
base_url=base_url,
297+
model_names=model_names,
298+
),
299+
)
300+
)
301+
ms.wait_until_ready_or_failed()
302+
303+
result = ms.completions(
304+
messages=[
305+
{
306+
"role": "system",
307+
"content": "你是一个回音壁，会原封不动返回用户的输入",
308+
},
309+
{"role": "user", "content": "你好！"},
310+
{"role": "assistant", "content": "你好！"},
311+
{"role": "user", "content": "今天天气怎么样？"},
312+
],
313+
stream=False,
314+
)
315+
assert isinstance(result, ModelResponse)
316+
assert (
317+
pydash.get(result, "choices[0].message.content")
318+
== "今天天气怎么样？"
319+
)
320+
321+
await ms.delete_async()
322+
await cr.delete_async()
323+
267324

268325
class TestModelProxy:
269326
"""ModelProxy 模块 E2E 测试"""

‎tests/e2e/test_model.py‎

Lines changed: 114 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,6 @@
4949
)
5050
from agentrun.utils.config import Config
5151
from agentrun.utils.exception import (
52-
ClientError,
5352
ResourceAlreadyExistError,
5453
ResourceNotExistError,
5554
)
@@ -474,6 +473,120 @@ def test_model_service_invoke(self, model_service_name: str):
474473

475474
ms.delete()
476475

476+
async def test_model_service_with_credential_async(
477+
self, model_service_name: str
478+
):
479+
# 创建 Credential
480+
from agentrun.credential import (
481+
Credential,
482+
CredentialConfig,
483+
CredentialCreateInput,
484+
)
485+
486+
cr = await Credential.create_async(
487+
CredentialCreateInput(
488+
credential_name=f"{model_service_name}-credential",
489+
enabled=True,
490+
credential_config=CredentialConfig.outbound_llm_api_key(
491+
api_key=api_key,
492+
provider="openai",
493+
),
494+
)
495+
)
496+
497+
# 创建 model service
498+
ms = await ModelService.create_async(
499+
ModelServiceCreateInput(
500+
model_service_name=model_service_name,
501+
description="原始描述",
502+
model_type=ModelType.LLM,
503+
provider="openai",
504+
credential_name=cr.credential_name,
505+
provider_settings=ProviderSettings(
506+
base_url=base_url,
507+
model_names=model_names,
508+
),
509+
)
510+
)
511+
ms.wait_until_ready_or_failed()
512+
513+
result = ms.completions(
514+
messages=[
515+
{
516+
"role": "system",
517+
"content": "你是一个回音壁，会原封不动返回用户的输入",
518+
},
519+
{"role": "user", "content": "你好！"},
520+
{"role": "assistant", "content": "你好！"},
521+
{"role": "user", "content": "今天天气怎么样？"},
522+
],
523+
stream=False,
524+
)
525+
assert isinstance(result, ModelResponse)
526+
assert (
527+
pydash.get(result, "choices[0].message.content")
528+
== "今天天气怎么样？"
529+
)
530+
531+
await ms.delete_async()
532+
await cr.delete_async()
533+
534+
def test_model_service_with_credential(self, model_service_name: str):
535+
# 创建 Credential
536+
from agentrun.credential import (
537+
Credential,
538+
CredentialConfig,
539+
CredentialCreateInput,
540+
)
541+
542+
cr = Credential.create(
543+
CredentialCreateInput(
544+
credential_name=f"{model_service_name}-credential",
545+
enabled=True,
546+
credential_config=CredentialConfig.outbound_llm_api_key(
547+
api_key=api_key,
548+
provider="openai",
549+
),
550+
)
551+
)
552+
553+
# 创建 model service
554+
ms = ModelService.create(
555+
ModelServiceCreateInput(
556+
model_service_name=model_service_name,
557+
description="原始描述",
558+
model_type=ModelType.LLM,
559+
provider="openai",
560+
credential_name=cr.credential_name,
561+
provider_settings=ProviderSettings(
562+
base_url=base_url,
563+
model_names=model_names,
564+
),
565+
)
566+
)
567+
ms.wait_until_ready_or_failed()
568+
569+
result = ms.completions(
570+
messages=[
571+
{
572+
"role": "system",
573+
"content": "你是一个回音壁，会原封不动返回用户的输入",
574+
},
575+
{"role": "user", "content": "你好！"},
576+
{"role": "assistant", "content": "你好！"},
577+
{"role": "user", "content": "今天天气怎么样？"},
578+
],
579+
stream=False,
580+
)
581+
assert isinstance(result, ModelResponse)
582+
assert (
583+
pydash.get(result, "choices[0].message.content")
584+
== "今天天气怎么样？"
585+
)
586+
587+
ms.delete()
588+
cr.delete()
589+
477590

478591
class TestModelProxy:
479592
"""ModelProxy 模块 E2E 测试"""

0 commit comments

Comments
 (0)