chore: enhance Python SDK.

e2e/src/e2e/api.py

@@ -11,15 +11,18 @@
 limitations under the License.
 """
 
-import flamepy
 from typing import Optional
+from dataclasses import dataclass
 
-class TestContext(flamepy.Request):
+@dataclass
+class TestContext:
     common_data: Optional[str] = None
 
-class TestRequest(flamepy.Request):
+@dataclass
+class TestRequest:
     input: Optional[str] = None
 
-class TestResponse(flamepy.Response):
+@dataclass
+class TestResponse:
     output: Optional[str] = None
     common_data: Optional[str] = None

e2e/src/e2e/service.py

@@ -13,21 +13,16 @@
 
 import flamepy
 
-from e2e import TestRequest, TestResponse, TestContext
+from e2e.api import TestRequest, TestResponse
 
 instance = flamepy.FlameInstance()
 
-sys_context = None
-
 @instance.entrypoint
 def e2e_service_entrypoint(req: TestRequest) -> TestResponse:
-    return TestResponse(output=req.input, common_data=sys_context)
+    cxt = instance.context()
+    data = cxt.common_data if cxt is not None else None
 
-@instance.context
-def e2e_service_context(ctx: TestContext = None):
-    global sys_context
-    if ctx is not None:
-        sys_context = ctx.common_data
+    return TestResponse(output=req.input, common_data=data)
 
 if __name__ == "__main__":
-    instance.run()
+    instance.run()

e2e/tests/test_session.py

@@ -36,8 +36,7 @@ def __init__(self, expected_output):
     def on_update(self, task):
         self.latest_state = task.state
         if task.state == flamepy.TaskState.SUCCEED:
-            resp = TestResponse.from_json(task.output)
-            assert resp.output == self.expected_output
+            assert task.output.output == self.expected_output, f"Task output: {task.output.output}, Expected: {self.expected_output}"
         elif task.state == flamepy.TaskState.FAILED:
             for event in task.events:
                 if event.code == flamepy.TaskState.FAILED:
@@ -99,8 +98,7 @@ async def test_invoke_task_without_common_data():
 
     input = random_string()
 
-    resp = await session.invoke(TestRequest(input=input))
-    output = TestResponse.from_json(resp)
+    output = await session.invoke(TestRequest(input=input))
     assert output.output == input
     assert output.common_data is None
 
@@ -122,8 +120,7 @@ async def test_invoke_task_with_common_data():
     assert ssn_list[0].application == FLM_TEST_APP
     assert ssn_list[0].state == SessionState.OPEN
 
-    resp = await session.invoke(TestRequest(input=input))
-    output = TestResponse.from_json(resp)
+    output = await session.invoke(TestRequest(input=input))
     assert output.output == input
     assert output.common_data == sys_context
 

sdk/python/example/instance.py

@@ -32,13 +32,6 @@ class Summary(flamepy.Response):
 You are a helpful assistant.
 """
 
-
-@ins.context
-def sys_context(sp: SysPrompt):
-    global sys_prompt
-    sys_prompt = sp.prompt
-
-
 @ins.entrypoint
 def summarize_blog(bl: Blog) -> Summary:
     global sys_prompt

sdk/python/src/flamepy/__init__.py

@@ -52,8 +52,6 @@
     Application,
     FlameContext,
     TaskInformer,
-    Request,
-    Response,
 )
 
 from .client import (
@@ -113,8 +111,6 @@
     "Application",
     "FlameContext",
     "TaskInformer",
-    "Request",
-    "Response",
     # Client classes
     "Connection",
     "connect",

sdk/python/src/flamepy/cache.py

@@ -20,7 +20,7 @@ class Object(BaseModel):
     """Object."""
 
     version: int
-    data: list 
+    data: list
 
 class ObjectMetadata(BaseModel):
     """Object metadata."""
@@ -45,8 +45,12 @@ async def put_object(session_id: str, data: bytes) -> "DataExpr":
 
     return DataExpr(source=DataSource.REMOTE, url=metadata.endpoint, data=data, version=metadata.version)
 
+
 async def get_object(de: DataExpr) -> "DataExpr":
     """Get an object from the cache."""
+    if de is None:
+        return None
+
     if de.source != DataSource.REMOTE:
         return de
 
@@ -63,6 +67,9 @@ async def get_object(de: DataExpr) -> "DataExpr":
 
 async def update_object(de: DataExpr) -> "DataExpr":
     """Update an object in the cache."""
+    if de is None:
+        return None
+
     if de.source != DataSource.REMOTE:
         return de