From 7b1cd0e7ec97569264fd3952e303c173ab989390 Mon Sep 17 00:00:00 2001
From: gautamgambhir97 <140384949+gautamgambhir97@users.noreply.github.com>
Date: Fri, 6 Sep 2024 10:34:57 +0530
Subject: [PATCH] fix(docs): fix on query handler example issue (#923)

---
 pages/guides/agents/intermediate/handlers.mdx | 32 ++++++++++++-------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/pages/guides/agents/intermediate/handlers.mdx b/pages/guides/agents/intermediate/handlers.mdx
index 56d94e944..a70610ff6 100644
--- a/pages/guides/agents/intermediate/handlers.mdx
+++ b/pages/guides/agents/intermediate/handlers.mdx
@@ -326,9 +326,10 @@ Let's explore the Proxy code script step-by-step:
     ```py copy
     import json
 
-    from fastapi import FastAPI
+    from fastapi import FastAPI, Request
     from uagents import Model
     from uagents.query import query
+    from uagents.envelope import Envelope
 
     AGENT_ADDRESS = "agent1qt6ehs6kqdgtrsduuzslqnrzwkrcn3z0cfvwsdj22s27kvatrxu8sy3vag0"
 
@@ -340,9 +341,11 @@ Let's explore the Proxy code script step-by-step:
 
     ```py copy
     async def agent_query(req):
-        response = await query(destination=AGENT_ADDRESS, message=req, timeout=15.0)
-        data = json.loads(response.decode_payload())
-        return data["text"]
+        response = await query(destination=AGENT_ADDRESS, message=req, timeout=15)
+        if isinstance(response, Envelope):
+            data = json.loads(response.decode_payload())
+            return data["text"]
+        return response
     ```
 
 5. Initialize a `FastAPI` app:
@@ -363,9 +366,10 @@ Let's explore the Proxy code script step-by-step:
 
     ```py copy
     @app.post("/endpoint")
-    async def make_agent_call(req: TestRequest):
+    async def make_agent_call(req: Request):
+        model = TestRequest.parse_obj(await req.json())
         try:
-            res = await agent_query(req)
+            res = await agent_query(model)
             return f"successful call - agent response: {res}"
         except Exception:
             return "unsuccessful agent call"
@@ -378,9 +382,10 @@ The overall script should look as follows:
 ```py copy filename="on_query.py"
 import json
 
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from uagents import Model
 from uagents.query import query
+from uagents.envelope import Envelope 
 
 AGENT_ADDRESS = "agent1qt6ehs6kqdgtrsduuzslqnrzwkrcn3z0cfvwsdj22s27kvatrxu8sy3vag0"
 
@@ -388,9 +393,11 @@ class TestRequest(Model):
     message: str
 
 async def agent_query(req):
-    response = await query(destination=AGENT_ADDRESS, message=req, timeout=15.0)
-    data = json.loads(response.decode_payload())
-    return data["text"]
+    response = await query(destination=AGENT_ADDRESS, message=req, timeout=15)
+    if isinstance(response, Envelope):
+        data = json.loads(response.decode_payload())
+        return data["text"]
+    return response
 
 app = FastAPI()
 
@@ -399,9 +406,10 @@ def read_root():
     return "Hello from the Agent controller"
 
 @app.post("/endpoint")
-async def make_agent_call(req: TestRequest):
+async def make_agent_call(req: Request):
+    model = TestRequest.parse_obj(await req.json())
     try:
-        res = await agent_query(req)
+        res = await agent_query(model)
         return f"successful call - agent response: {res}"
     except Exception:
         return "unsuccessful agent call"