Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[project]
name = "veadk-python"
version = "0.5.11"
version = "0.5.12"
description = "Volcengine agent development kit, integrations with Volcengine cloud services."
readme = "README.md"
requires-python = ">=3.10"
Expand Down
63 changes: 38 additions & 25 deletions veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,30 +116,55 @@ def _create_collection(self) -> None:
logger.debug(f"Create collection with response {response}")
return response

def _get_client(self) -> VikingDBMemoryClient:
if not (self.volcengine_access_key and self.volcengine_secret_key):
def _get_ak_sk_sts(self) -> tuple[str, str, str]:
ak = ""
sk = ""
sts_token = ""

if self.volcengine_access_key and self.volcengine_secret_key:
ak = self.volcengine_access_key
sk = self.volcengine_secret_key
sts_token = self.session_token
logger.debug(
"Get volcengine credential from Environments for VikingMEMBackend"
)

else:
cred = get_credential_from_vefaas_iam()
self.volcengine_access_key = cred.access_key_id
self.volcengine_secret_key = cred.secret_access_key
self.session_token = cred.session_token
ak = cred.access_key_id
sk = cred.secret_access_key
sts_token = cred.session_token
logger.debug(
"Get volcengine credential from VeFaaS IAM file for VikingMEMBackend."
)

return ak, sk, sts_token

def _get_client(self) -> VikingDBMemoryClient:
ak, sk, sts_token = self._get_ak_sk_sts()
return VikingDBMemoryClient(
ak=self.volcengine_access_key,
sk=self.volcengine_secret_key,
sts_token=self.session_token,
ak=ak,
sk=sk,
sts_token=sts_token,
region=self.region,
)

def _get_sdk_client(self) -> VikingMem:
client = self._get_client()
ak, sk, sts_token = self._get_ak_sk_sts()
client = VikingDBMemoryClient(
ak=ak,
sk=sk,
sts_token=sts_token,
region=self.region,
)
return VikingMem(
host=client.get_host(),
region=self.region,
auth=IAM(
ak=self.volcengine_access_key,
sk=self.volcengine_secret_key,
ak=ak,
sk=sk,
),
sts_token=self.session_token,
sts_token=sts_token,
)

@override
Expand Down Expand Up @@ -237,19 +262,7 @@ def get_user_profile(self, user_id: str) -> str:
f"Get user profile for user_id={user_id} from Viking Memory with mem_id={mem_id}"
)

ak = ""
sk = ""
sts_token = ""
if self.volcengine_access_key and self.volcengine_secret_key:
ak = self.volcengine_access_key
sk = self.volcengine_secret_key
sts_token = self.session_token
else:
cred = get_credential_from_vefaas_iam()
ak = cred.access_key_id
sk = cred.secret_access_key
sts_token = cred.session_token

ak, sk, sts_token = self._get_ak_sk_sts()
response = ve_request(
request_body={
"filter": {
Expand Down
2 changes: 1 addition & 1 deletion veadk/version.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,4 +12,4 @@
# See the License for the specific language governing permissions and
# limitations under the License.

VERSION = "0.5.11"
VERSION = "0.5.12"