suggest question values on record completed

app.py

@@ -1,74 +1,151 @@
+import json
+import traceback
 from queue import Queue
+from threading import Thread
+from typing import List
 
 import argilla as rg
 import gradio as gr
+from gradio_client import Client
 
 client = rg.Argilla()
 
-incoming_events = Queue()
+completed_record_events = Queue()
 
-def check_incoming_events():
-    """
-    This function is called every 5 seconds to check if there are any incoming
-    events and send data to update the JSON component.
-    """
-    events = []
-    while not incoming_events.empty():
-        events.append(incoming_events.get())
 
-    return {"events": events}
+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"]))
+
+    dataset_name = "stanfordnlp_imdb"
+    dataset = client.datasets(dataset_name) or rg.Dataset.from_hub(
+        "stanfordnlp/imdb",
+        name=dataset_name,
+        settings=settings,
+        client=client,
+        split="train[:1000]"
+    )
+
+    return 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}).")
-    json_component = gr.JSON(label="Incoming argilla events:")
-    gr.Timer(5, active=True).tick(check_incoming_events, outputs=json_component)
-
+    gr.Markdown("### Record Events")
+    gr.Markdown("#### Records are processed in background and suggestions are added.")
 
 server, _, _ = demo.launch(prevent_thread_lock=True, app_kwargs={"docs_url": "/docs"})
 
 # Set up the webhook listeners
 rg.set_webhook_server(server)
 
-# Delete all existing webhooks
 for webhook in client.webhooks:
-    webhook.delete()
+    webhook.enabled = False
+    webhook.update()
+
 
 # Create a webhook for record events
+@rg.webhook_listener(events="record.completed")
+async def record_events(record: rg.Record, type: str, **kwargs):
+    print("Received event", type)
+
+    completed_record_events.put(record)
+
+
+dataset = build_dataset(client)
+
+
+def add_record_suggestions_on_response_created():
+    print("Starting thread")
+
+    completed_records_filter = rg.Filter(("status", "==", "completed"))
+    pending_records_filter = rg.Filter(("status", "==", "pending"))
+
+    while True:
+        try:
+            record: rg.Record = completed_record_events.get()
+
+            if dataset.id != record.dataset.id:
+                continue
+
+            # Prepare predict data
+
+            field = dataset.settings.fields["text"]
+            question = dataset.settings.questions["sentiment"]
+
+            examples = list(
+                dataset.records(
+                    query=rg.Query(filter=completed_records_filter),
+                    limit=5,
+                )
+            )
+
+            some_pending_records = list(
+                dataset.records(
+                    query=rg.Query(filter=pending_records_filter),
+                    limit=5,
+                )
+            )
+
+            if not some_pending_records:
+                continue
+
+            some_pending_records = parse_pending_records(some_pending_records, field, question, examples)
+            dataset.records.log(some_pending_records)
+
+        except Exception:
+            print(traceback.format_exc())
+            continue
 
-@rg.webhook_listener(
-    events=["record.created", "record.updated", "record.completed"],
-    raw_event=True # Using raw events until PR https://github.com/argilla-io/argilla/pull/5500 is merged
-)
-async def record_events(event:dict):
-    print("Received event", event)
 
-    incoming_events.put(event)
+def parse_pending_records(
+    records: List[rg.Record],
+    field: rg.Field,
+    question,
+    example_records: List[rg.Record]
+) -> List[rg.Record]:
+    gradio_client = Client("davidberenstein1957/distilabel-argilla-labeller")
 
-# Create a webhook for dataset events
-@rg.webhook_listener(events=["dataset.created", "dataset.updated", "dataset.published"])
-async def dataset_events(type: str, dataset: rg.Dataset | None = None, **kwargs):
-    print(f"Received event {type} for dataset {dataset.id}")
+    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],
+    }
 
-    incoming_events.put((type, dataset))
+    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]
+    }
 
-# Create a webhook for response events
-@rg.webhook_listener(
-    events=["response.created", "response.updated"],
-    raw_event=True # Using raw events until PR https://github.com/argilla-io/argilla/pull/5500 is merged
-)
-async def response_events(event: dict):
-    print("Received event", event)
+    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"],
+            )
+        )
 
-    incoming_events.put(event)
+    return records
 
-@rg.webhook_listener(events=["record.deleted", "dataset.deleted", "response.deleted"])
-async def deleted_events(type: str, data: dict, **kwargs):
-    print(f"Received event {type} for resource {data}")
 
-    incoming_events.put((type, data))
+thread = Thread(target=add_record_suggestions_on_response_created)
 
+thread.start()
 demo.block_thread()