feat: :sparkles: make follow-up edits

1 files changed, 5 insertions, 5 deletions

diff --git a/continuedev/src/continuedev/core/main.py b/continuedev/src/continuedev/core/main.py
index ace1ad60..63a3e6a9 100644
--- a/continuedev/src/continuedev/core/main.py
+++ b/continuedev/src/continuedev/core/main.py
@@ -1,5 +1,5 @@
 import json
-from typing import Coroutine, Dict, List, Literal, Optional, Union
+from typing import Any, Coroutine, Dict, List, Literal, Optional, Union
 
 from pydantic import BaseModel, validator
 from pydantic.schema import schema
@@ -401,13 +401,13 @@ class Validator(Step):
 
 
 class Context:
-    key_value: Dict[str, str] = {}
+    key_value: Dict[str, Any] = {}
 
-    def set(self, key: str, value: str):
+    def set(self, key: str, value: Any):
         self.key_value[key] = value
 
-    def get(self, key: str) -> str:
-        return self.key_value[key]
+    def get(self, key: str) -> Any:
+        return self.key_value.get(key, None)
 
 
 class ContinueCustomException(Exception):