From 96bd5ccfd142fae9fe73f720e18d2d5b5656ea1d Mon Sep 17 00:00:00 2001 From: Grant Linville Date: Mon, 16 Sep 2024 10:28:28 -0400 Subject: [PATCH 1/3] feat: add credential management Signed-off-by: Grant Linville --- gptscript/credentials.py | 74 ++++++++++++++++++++++++++++++++++++++++ gptscript/gptscript.py | 33 ++++++++++++++++++ tests/test_gptscript.py | 18 ++++++++++ 3 files changed, 125 insertions(+) create mode 100644 gptscript/credentials.py diff --git a/gptscript/credentials.py b/gptscript/credentials.py new file mode 100644 index 0000000..42e506f --- /dev/null +++ b/gptscript/credentials.py @@ -0,0 +1,74 @@ +import json +from datetime import datetime, timezone +from enum import Enum + + +def is_timezone_aware(dt: datetime): + return dt.tzinfo is not None and dt.tzinfo.utcoffset(dt) is not None + + +class CredentialType(Enum): + Tool = "tool", + ModelProvider = "modelProvider" + + +class Credential: + def __init__(self, + context: str = "default", + toolName: str = "", + type: CredentialType = CredentialType.Tool, + env: dict[str, str] = None, + ephemeral: bool = False, + expiresAt: datetime = None, + refreshToken: str = "", + ): + self.context = context + self.toolName = toolName + self.type = type + self.env = env + self.ephemeral = ephemeral + self.expiresAt = expiresAt + self.refreshToken = refreshToken + + if self.env is None: + self.env = {} + + def to_json(self): + datetime_str = "" + + if self.expiresAt is not None: + system_tz = datetime.now().astimezone().tzinfo + + if not is_timezone_aware(self.expiresAt): + self.expiresAt = self.expiresAt.replace(tzinfo=system_tz) + datetime_str = self.expiresAt.isoformat() + + # For UTC only, replace the "+00:00" with "Z" + if self.expiresAt.tzinfo == timezone.utc: + datetime_str = datetime_str.replace("+00:00", "Z") + + req = { + "context": self.context, + "toolName": self.toolName, + "type": self.type.value[0], + "env": self.env, + "ephemeral": self.ephemeral, + "refreshToken": self.refreshToken, + } + + if datetime_str != "": + req["expiresAt"] = datetime_str + + return json.dumps(req) + +class CredentialRequest: + def __init__(self, + content: str = "", + allContexts: bool = False, + context: str = "default", + name: str = "", + ): + self.content = content + self.allContexts = allContexts + self.context = context + self.name = name diff --git a/gptscript/gptscript.py b/gptscript/gptscript.py index c29a155..b27079d 100644 --- a/gptscript/gptscript.py +++ b/gptscript/gptscript.py @@ -9,6 +9,7 @@ import requests from gptscript.confirm import AuthResponse +from gptscript.credentials import Credential from gptscript.frame import RunFrame, CallFrame, PromptFrame, Program from gptscript.opts import GlobalOptions from gptscript.prompt import PromptResponse @@ -183,6 +184,38 @@ async def list_models(self, providers: list[str] = None, credential_overrides: l {"providers": providers, "credentialOverrides": credential_overrides} )).split("\n") + async def list_credentials(self, context: str = "default", all_contexts: bool = False) -> list[Credential] | str: + res = await self._run_basic_command( + "credentials", + {"context": context, "allContexts": all_contexts} + ) + if res.startswith("an error occurred:"): + return res + + return [Credential(**c) for c in json.loads(res)] + + async def create_credential(self, cred: Credential) -> str: + return await self._run_basic_command( + "credentials/create", + {"content": cred.to_json()} + ) + + async def reveal_credential(self, context: str = "default", name: str = "") -> Credential | str: + res = await self._run_basic_command( + "credentials/reveal", + {"context": context, "name": name} + ) + if res.startswith("an error occurred:"): + return res + + return Credential(**json.loads(res)) + + async def delete_credential(self, context: str = "default", name: str = "") -> str: + return await self._run_basic_command( + "credentials/delete", + {"context": context, "name": name} + ) + def _get_command(): if os.getenv("GPTSCRIPT_BIN") is not None: diff --git a/tests/test_gptscript.py b/tests/test_gptscript.py index 1eea8a4..17177d0 100644 --- a/tests/test_gptscript.py +++ b/tests/test_gptscript.py @@ -8,6 +8,7 @@ import pytest from gptscript.confirm import AuthResponse +from gptscript.credentials import Credential from gptscript.exec_utils import get_env from gptscript.frame import RunEventType, CallFrame, RunFrame, RunState, PromptFrame from gptscript.gptscript import GPTScript @@ -683,3 +684,20 @@ async def test_parse_with_metadata_then_run(gptscript): tools = await gptscript.parse(cwd + "/tests/fixtures/parse-with-metadata.gpt") run = gptscript.evaluate(tools[0]) assert "200" == await run.text(), "Expect file to have correct output" + +@pytest.mark.asyncio +async def test_credentials(gptscript): + name = "test-" + str(os.urandom(4).hex()) + res = await gptscript.create_credential(Credential(toolName=name, env={"TEST": "test"})) + assert not res.startswith("an error occurred"), "Unexpected error creating credential: " + res + + res = await gptscript.list_credentials() + assert not str(res).startswith("an error occurred"), "Unexpected error listing credentials: " + str(res) + assert len(res) > 0, "Expected at least one credential" + + res = await gptscript.reveal_credential(name=name) + assert not str(res).startswith("an error occurred"), "Unexpected error revealing credential: " + res + assert res.env["TEST"] == "test", "Unexpected credential value: " + str(res) + + res = await gptscript.delete_credential(name=name) + assert not res.startswith("an error occurred"), "Unexpected error deleting credential: " + res From e5ed8843763a2b7957dea14821587c957edc2c1d Mon Sep 17 00:00:00 2001 From: Grant Linville Date: Wed, 18 Sep 2024 10:34:40 -0400 Subject: [PATCH 2/3] update for stacked contexts Signed-off-by: Grant Linville --- gptscript/credentials.py | 8 ++++++-- gptscript/gptscript.py | 18 ++++++++++++------ 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/gptscript/credentials.py b/gptscript/credentials.py index 42e506f..09a4b64 100644 --- a/gptscript/credentials.py +++ b/gptscript/credentials.py @@ -1,6 +1,7 @@ import json from datetime import datetime, timezone from enum import Enum +from typing import List def is_timezone_aware(dt: datetime): @@ -65,10 +66,13 @@ class CredentialRequest: def __init__(self, content: str = "", allContexts: bool = False, - context: str = "default", + contexts: List[str] = None, name: str = "", ): + if contexts is None: + contexts = ["default"] + self.content = content self.allContexts = allContexts - self.context = context + self.contexts = contexts self.name = name diff --git a/gptscript/gptscript.py b/gptscript/gptscript.py index b27079d..946f746 100644 --- a/gptscript/gptscript.py +++ b/gptscript/gptscript.py @@ -4,7 +4,7 @@ from subprocess import Popen, PIPE from sys import executable from time import sleep -from typing import Any, Callable, Awaitable +from typing import Any, Callable, Awaitable, List import requests @@ -184,10 +184,13 @@ async def list_models(self, providers: list[str] = None, credential_overrides: l {"providers": providers, "credentialOverrides": credential_overrides} )).split("\n") - async def list_credentials(self, context: str = "default", all_contexts: bool = False) -> list[Credential] | str: + async def list_credentials(self, contexts: List[str] = None, all_contexts: bool = False) -> list[Credential] | str: + if contexts is None: + contexts = ["default"] + res = await self._run_basic_command( "credentials", - {"context": context, "allContexts": all_contexts} + {"context": contexts, "allContexts": all_contexts} ) if res.startswith("an error occurred:"): return res @@ -200,10 +203,13 @@ async def create_credential(self, cred: Credential) -> str: {"content": cred.to_json()} ) - async def reveal_credential(self, context: str = "default", name: str = "") -> Credential | str: + async def reveal_credential(self, contexts: List[str] = None, name: str = "") -> Credential | str: + if contexts is None: + contexts = ["default"] + res = await self._run_basic_command( "credentials/reveal", - {"context": context, "name": name} + {"context": contexts, "name": name} ) if res.startswith("an error occurred:"): return res @@ -213,7 +219,7 @@ async def reveal_credential(self, context: str = "default", name: str = "") -> C async def delete_credential(self, context: str = "default", name: str = "") -> str: return await self._run_basic_command( "credentials/delete", - {"context": context, "name": name} + {"context": [context], "name": name} ) From 69519e02872c56d94ad9d71244fcd8495e4f0fc9 Mon Sep 17 00:00:00 2001 From: Grant Linville Date: Wed, 18 Sep 2024 14:01:49 -0400 Subject: [PATCH 3/3] update for stacked contexts Signed-off-by: Grant Linville --- gptscript/credentials.py | 23 ++++++++++++++++++++--- gptscript/gptscript.py | 6 +++--- tests/test_gptscript.py | 10 +++++++++- 3 files changed, 32 insertions(+), 7 deletions(-) diff --git a/gptscript/credentials.py b/gptscript/credentials.py index 09a4b64..5662ff0 100644 --- a/gptscript/credentials.py +++ b/gptscript/credentials.py @@ -9,15 +9,15 @@ def is_timezone_aware(dt: datetime): class CredentialType(Enum): - Tool = "tool", - ModelProvider = "modelProvider" + tool = "tool", + modelProvider = "modelProvider" class Credential: def __init__(self, context: str = "default", toolName: str = "", - type: CredentialType = CredentialType.Tool, + type: CredentialType = CredentialType.tool, env: dict[str, str] = None, ephemeral: bool = False, expiresAt: datetime = None, @@ -76,3 +76,20 @@ def __init__(self, self.allContexts = allContexts self.contexts = contexts self.name = name + +def to_credential(c) -> Credential: + expiresAt = c["expiresAt"] + if expiresAt is not None: + if expiresAt.endswith("Z"): + expiresAt = expiresAt.replace("Z", "+00:00") + expiresAt = datetime.fromisoformat(expiresAt) + + return Credential( + context=c["context"], + toolName=c["toolName"], + type=CredentialType[c["type"]], + env=c["env"], + ephemeral=c.get("ephemeral", False), + expiresAt=expiresAt, + refreshToken=c["refreshToken"], + ) diff --git a/gptscript/gptscript.py b/gptscript/gptscript.py index 946f746..bd5c9d6 100644 --- a/gptscript/gptscript.py +++ b/gptscript/gptscript.py @@ -9,7 +9,7 @@ import requests from gptscript.confirm import AuthResponse -from gptscript.credentials import Credential +from gptscript.credentials import Credential, to_credential from gptscript.frame import RunFrame, CallFrame, PromptFrame, Program from gptscript.opts import GlobalOptions from gptscript.prompt import PromptResponse @@ -195,7 +195,7 @@ async def list_credentials(self, contexts: List[str] = None, all_contexts: bool if res.startswith("an error occurred:"): return res - return [Credential(**c) for c in json.loads(res)] + return [to_credential(cred) for cred in json.loads(res)] async def create_credential(self, cred: Credential) -> str: return await self._run_basic_command( @@ -214,7 +214,7 @@ async def reveal_credential(self, contexts: List[str] = None, name: str = "") -> if res.startswith("an error occurred:"): return res - return Credential(**json.loads(res)) + return to_credential(json.loads(res)) async def delete_credential(self, context: str = "default", name: str = "") -> str: return await self._run_basic_command( diff --git a/tests/test_gptscript.py b/tests/test_gptscript.py index 17177d0..3347dc7 100644 --- a/tests/test_gptscript.py +++ b/tests/test_gptscript.py @@ -4,6 +4,8 @@ import os import platform import subprocess +from datetime import datetime, timedelta, timezone +from time import sleep import pytest @@ -688,12 +690,18 @@ async def test_parse_with_metadata_then_run(gptscript): @pytest.mark.asyncio async def test_credentials(gptscript): name = "test-" + str(os.urandom(4).hex()) - res = await gptscript.create_credential(Credential(toolName=name, env={"TEST": "test"})) + now = datetime.now() + res = await gptscript.create_credential(Credential(toolName=name, env={"TEST": "test"}, expiresAt=now + timedelta(seconds=5))) assert not res.startswith("an error occurred"), "Unexpected error creating credential: " + res + sleep(5) + res = await gptscript.list_credentials() assert not str(res).startswith("an error occurred"), "Unexpected error listing credentials: " + str(res) assert len(res) > 0, "Expected at least one credential" + for cred in res: + if cred.toolName == name: + assert cred.expiresAt < datetime.now(timezone.utc), "Expected credential to have expired" res = await gptscript.reveal_credential(name=name) assert not str(res).startswith("an error occurred"), "Unexpected error revealing credential: " + res