From 9ef266cee3e09ab66a1d9f14dec0ba8ce251c87c Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Fri, 1 Jul 2022 20:34:47 +0200
Subject: [PATCH] Add builtin deserialization

---
 src/great_ai/great_ai/remote/call_remote_great_ai.py | 12 +++++++-----
 .../great_ai/remote/call_remote_great_ai_async.py    | 11 ++++++++---
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/src/great_ai/great_ai/remote/call_remote_great_ai.py b/src/great_ai/great_ai/remote/call_remote_great_ai.py
index 449c64b..cd03251 100644
--- a/src/great_ai/great_ai/remote/call_remote_great_ai.py
+++ b/src/great_ai/great_ai/remote/call_remote_great_ai.py
@@ -1,5 +1,7 @@
 import asyncio
-from typing import Any, Mapping
+from typing import Any, Mapping, Optional, TypeVar
+
+from pydantic import BaseModel
 
 from ...utilities import get_logger
 from ..views import Trace
@@ -7,10 +9,10 @@ from .call_remote_great_ai_async import call_remote_great_ai_async
 
 logger = get_logger("call_remote_great_ai")
 
-
+T = TypeVar("T", BaseModel)
 def call_remote_great_ai(
-    base_uri: str, data: Mapping[str, Any], retry_count: int = 4
-) -> Trace:
+    base_uri: str, data: Mapping[str, Any], retry_count: int = 4, model_class:Optional[Type[T]]=None
+) -> Trace[T]:
     try:
         asyncio.get_running_loop()
         raise Exception(
@@ -20,7 +22,7 @@ def call_remote_great_ai(
         pass
 
     future = call_remote_great_ai_async(
-        base_uri=base_uri, data=data, retry_count=retry_count
+        base_uri=base_uri, data=data, retry_count=retry_count,model_class=model_class
     )
 
     return asyncio.run(future)
diff --git a/src/great_ai/great_ai/remote/call_remote_great_ai_async.py b/src/great_ai/great_ai/remote/call_remote_great_ai_async.py
index 78f059a..95dbbb2 100644
--- a/src/great_ai/great_ai/remote/call_remote_great_ai_async.py
+++ b/src/great_ai/great_ai/remote/call_remote_great_ai_async.py
@@ -1,4 +1,6 @@
-from typing import Any, Mapping, Optional
+from typing import Any, Mapping, Optional, Type, TypeVar
+
+from pydantic import BaseModel
 
 from ..views import Trace
 from .http_client import HttpClient
@@ -6,10 +8,11 @@ from .remote_call_error import RemoteCallError
 
 http: Optional[HttpClient] = None
 
+T = TypeVar("T", BaseModel)
 
 async def call_remote_great_ai_async(
-    base_uri: str, data: Mapping[str, Any], retry_count: int = 4
-) -> Trace:
+    base_uri: str, data: Mapping[str, Any], retry_count: int = 4, model_class:Optional[Type[T]]=None
+) -> Trace[T]:
     global http
     if http is None:
         http = HttpClient()
@@ -23,6 +26,8 @@ async def call_remote_great_ai_async(
     )
 
     try:
+        if model_class is not None:
+            response['output'] = model_class.parse_obj(response['output'])
         return Trace.parse_obj(response)
     except Exception:
         raise RemoteCallError("Could not parse response")