feat: add payload code

app.py

@@ -16,7 +16,9 @@ completed_record_events = Queue()
 def build_dataset(client: rg.Argilla) -> rg.Dataset:
     settings = rg.Settings.from_hub("stanfordnlp/imdb")
 
-    settings.questions.add(rg.LabelQuestion(name="sentiment", labels=["negative", "positive"]))
+    settings.questions.add(
+        rg.LabelQuestion(name="sentiment", labels=["negative", "positive"])
+    )
 
     dataset_name = "stanfordnlp_imdb"
     dataset = client.datasets(dataset_name) or rg.Dataset.from_hub(
@@ -24,7 +26,7 @@ def build_dataset(client: rg.Argilla) -> rg.Dataset:
         name=dataset_name,
         settings=settings,
         client=client,
-        split="train[:1000]"
+        split="train[:1000]",
     )
 
     return dataset
@@ -33,7 +35,9 @@ def build_dataset(client: rg.Argilla) -> rg.Dataset:
 with gr.Blocks() as demo:
     argilla_server = client.http_client.base_url
     gr.Markdown("## Argilla Events")
-    gr.Markdown(f"This demo shows the incoming events from the [Argilla Server]({argilla_server}).")
+    gr.Markdown(
+        f"This demo shows the incoming events from the [Argilla Server]({argilla_server})."
+    )
     gr.Markdown("### Record Events")
     gr.Markdown("#### Records are processed in background and suggestions are added.")
 
@@ -72,7 +76,6 @@ def add_record_suggestions_on_response_created():
                 continue
 
             # Prepare predict data
-
             field = dataset.settings.fields["text"]
             question = dataset.settings.questions["sentiment"]
 
@@ -93,7 +96,9 @@ def add_record_suggestions_on_response_created():
             if not some_pending_records:
                 continue
 
-            some_pending_records = parse_pending_records(some_pending_records, field, question, examples)
+            some_pending_records = parse_pending_records(
+                some_pending_records, field, question, examples
+            )
             dataset.records.log(some_pending_records)
 
         except Exception:
@@ -105,42 +110,33 @@ def parse_pending_records(
     records: List[rg.Record],
     field: rg.Field,
     question,
-    example_records: List[rg.Record]
+    example_records: List[rg.Record],
 ) -> List[rg.Record]:
-    gradio_client = Client("davidberenstein1957/distilabel-argilla-labeller")
-
-    payload = {
-        "records": [record.to_dict() for record in records],
-        "fields": [field.serialize()],
-        "question": question.serialize(),
-        "example_records": [record.to_dict() for record in example_records],
-    }
-
-    print("Sending payload")
-    # response = gradio_client.predict(
-    #     records=json.dumps(payload["records"]),
-    #     example_records=json.dumps(payload["example_records"]),
-    #     fields=json.dumps(payload["fields"]),
-    #     question=json.dumps(payload["question"]),
-    #     api_name="/predict"
-    # )
-    # print(response)
-    # response = json.loads(response)
-    # print("Response ", response)
-
-    response = {
-        "results": [{"value": "positive", "score": None, "agent": "mock"} for _ in records]
-    }
-
-    for record, suggestion in zip(records, response["results"]):
-        record.suggestions.add(
-            rg.Suggestion(
-                question_name=question.name,
-                value=suggestion["value"],
-                score=suggestion["score"],
-                agent=suggestion["agent"],
+    try:
+        gradio_client = Client("davidberenstein1957/distilabel-argilla-labeller")
+
+        payload = {
+            "records": [record.to_dict() for record in records],
+            "fields": [field.serialize()],
+            "question": question.serialize(),
+            "example_records": [record.to_dict() for record in example_records],
+            "api_name": "/predict",
+        }
+
+        response = gradio_client.predict(**payload)
+        response = json.loads(response) if isinstance(response, str) else response
+
+        for record, suggestion in zip(records, response["results"]):
+            record.suggestions.add(
+                rg.Suggestion(
+                    question_name=question.name,
+                    value=suggestion["value"],
+                    score=suggestion["score"],
+                    agent=suggestion["agent"],
+                )
             )
-        )
+    except Exception:
+        print(traceback.format_exc())
 
     return records