|
22 | 22 | import requests |
23 | 23 | from pydantic import Field |
24 | 24 | from typing_extensions import override |
25 | | -from veadk.utils.misc import getenv |
| 25 | +from volcengine.viking_knowledgebase import VikingKnowledgeBaseService |
| 26 | + |
26 | 27 | import veadk.config # noqa E401 |
27 | 28 | from veadk.auth.veauth.utils import get_credential_from_vefaas_iam |
28 | 29 | from veadk.configs.database_configs import NormalTOSConfig, TOSConfig |
29 | 30 | from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend |
30 | | -from veadk.knowledgebase.backends.utils import build_vikingdb_knowledgebase_request |
| 31 | +from veadk.knowledgebase.backends.utils import ( |
| 32 | + build_vikingdb_knowledgebase_request, |
| 33 | +) |
31 | 34 | from veadk.knowledgebase.entry import KnowledgebaseEntry |
32 | 35 | from veadk.utils.logger import get_logger |
33 | | -from veadk.utils.misc import formatted_timestamp |
34 | | -from volcengine.viking_knowledgebase import VikingKnowledgeBaseService |
| 36 | +from veadk.utils.misc import formatted_timestamp, getenv |
35 | 37 |
|
36 | 38 | try: |
37 | 39 | from veadk.integrations.ve_tos.ve_tos import VeTOS |
@@ -111,9 +113,14 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend): |
111 | 113 | ) |
112 | 114 | session_token: str = "" |
113 | 115 |
|
114 | | - volcengine_project: str = "default" |
| 116 | + volcengine_project: str = Field( |
| 117 | + default_factory=lambda: os.getenv("DATABASE_VIKING_PROJECT", "default") |
| 118 | + ) |
| 119 | + |
| 120 | + region: str = Field( |
| 121 | + default_factory=lambda: os.getenv("DATABASE_VIKING_REGION", "cn-beijing") |
| 122 | + ) |
115 | 123 |
|
116 | | - region: str = "cn-beijing" |
117 | 124 | base_url: str = "https://api-knowledgebase.mlp.cn-beijing.volces.com" |
118 | 125 | host: str = "api-knowledgebase.mlp.cn-beijing.volces.com" |
119 | 126 | schema: str = "https" |
@@ -562,7 +569,9 @@ def _search_knowledge( |
562 | 569 |
|
563 | 570 | def _set_service_info(self): |
564 | 571 | env_host = getenv( |
565 | | - "DATABASE_VIKING_BASE_URL", default_value=None, allow_false_values=True |
| 572 | + "DATABASE_VIKING_BASE_URL", |
| 573 | + default_value=None, |
| 574 | + allow_false_values=True, |
566 | 575 | ) |
567 | 576 | if env_host: |
568 | 577 | if env_host.startswith("http://") or env_host.startswith("https://"): |
|
0 commit comments