Skip to content
Merged

Dev #33

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
11 changes: 7 additions & 4 deletions src/agentstr/agents/agentstr.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,19 +107,22 @@ def checkpointer(self):
checkpointer = None
if self.database.conn_str.startswith("postgres"):
key_manager = os.getenv("AGENT_VAULT_KEY_MANAGER")
key_manager_prefix = os.getenv("AGENT_VAULT_KEY_MANAGER_PREFIX")
key_manager_prefix = os.getenv("AGENT_VAULT_KEY_MANAGER_PREFIX", f"AGENTSTR-{self.name}-".upper().replace(' ', '-'))
if key_manager:
try:
from agent_vault.langgraph import async_insecure_postgres_saver, async_secure_postgres_saver
from agent_vault.utils.key_manager import AWSSecretsManagerKeyManager, AzureKeyVaultKeyManager
from agent_vault.utils.key_manager import AWSParameterStoreKeyManager, AWSSecretsManagerKeyManager, AzureKeyVaultKeyManager
except ImportError:
raise ValueError("agent_vault is not installed")
if key_manager == "none":
checkpointer = async_insecure_postgres_saver(self.database.conn_str)
elif key_manager == "aws":
checkpointer = async_secure_postgres_saver(self.database.conn_str, AWSSecretsManagerKeyManager(prefix=key_manager_prefix))
checkpointer = async_secure_postgres_saver(self.database.conn_str, AWSParameterStoreKeyManager(prefix=key_manager_prefix))
elif key_manager == "azure":
checkpointer = async_secure_postgres_saver(self.database.conn_str, AzureKeyVaultKeyManager(prefix=key_manager_prefix))
key_vault_url = os.getenv("AZURE_KEY_VAULT_URL")
if not key_vault_url:
raise ValueError("Agent Vault Azure Key Manager requested by AZURE_KEY_VAULT_URL environment variable is not set")
checkpointer = async_secure_postgres_saver(self.database.conn_str, AzureKeyVaultKeyManager(vault_url=key_vault_url, prefix=key_manager_prefix))
else:
raise ValueError(f"Unsupported key manager: {key_manager}")
else:
Expand Down