Update AI RAG response parsing.

Author: dnwpark

gel/ai/core.py

@@ -102,6 +102,35 @@ def _make_rag_request(
             stream=stream,
         )
 
+    @staticmethod
+    def _parse_rag_response(
+        resp: typing.Any
+    )-> str:
+        data: dict[str, typing.Any] = resp.json()
+
+        if text := data.get("text"):
+            if isinstance(text, str):
+                return text
+            else:
+                raise RuntimeError(
+                    f"Expected data.text to be a string, but got "
+                    f"{type(text).__name__}: {text}"
+                )
+
+        elif response := data.get("response"):
+            if isinstance(text, str):
+                return text
+            else:
+                raise RuntimeError(
+                    f"Expected data.text to be a string, but got "
+                    f"{type(response).__name__}: {response}"
+                )
+
+        raise RuntimeError(
+            f"Expected response to include a non-empty string for either the "
+            f"'text' or 'response' key, but got: {data}"
+        )
+
 
 class RAGClient(BaseRAGClient):
     client: httpx.Client
@@ -120,7 +149,7 @@ def query_rag(
             ).to_httpx_request()
         )
         resp.raise_for_status()
-        return resp.json()["response"]
+        return BaseRAGClient._parse_rag_response(resp)
 
     def stream_rag(
         self, message: str, context: typing.Optional[types.QueryContext] = None
@@ -163,7 +192,7 @@ async def query_rag(
             ).to_httpx_request()
         )
         resp.raise_for_status()
-        return resp.json()["response"]
+        return BaseRAGClient._parse_rag_response(resp)
 
     async def stream_rag(
         self, message: str, context: typing.Optional[types.QueryContext] = None