diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..be86c8426eedf3224e93d8677e53de92c95abb76 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text
+potato/static/vendor/font-awesome-6.7.2/webfonts/fa-brands-400.ttf filter=lfs diff=lfs merge=lfs -text
+potato/static/vendor/font-awesome-6.7.2/webfonts/fa-brands-400.woff2 filter=lfs diff=lfs merge=lfs -text
+potato/static/vendor/font-awesome-6.7.2/webfonts/fa-solid-900.ttf filter=lfs diff=lfs merge=lfs -text
+potato/static/vendor/font-awesome-6.7.2/webfonts/fa-solid-900.woff2 filter=lfs diff=lfs merge=lfs -text
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000000000000000000000000000000000000..c33bd688429da992c0e840006a809d8a17f34035
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,41 @@
+FROM python:3.11-slim
+
+# Create non-root user (HF Spaces requires UID 1000)
+RUN useradd -m -u 1000 potato
+
+# Install system dependencies
+RUN apt-get update && \
+ apt-get install -y --no-install-recommends git && \
+ rm -rf /var/lib/apt/lists/*
+
+# Set working directory
+WORKDIR /app
+
+# Copy requirements first for layer caching
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt gunicorn
+
+# Copy the application source
+COPY . .
+
+# Install potato
+RUN pip install --no-cache-dir -e .
+
+# Copy entrypoint
+COPY entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+
+# Create directories for output
+RUN mkdir -p /app/annotation_output && \
+ chown -R potato:potato /app
+
+# Switch to non-root user
+USER potato
+
+# HuggingFace Spaces expects port 7860
+EXPOSE 7860
+
+ENV POTATO_CONFIG=config.yaml
+ENV PORT=7860
+
+ENTRYPOINT ["/entrypoint.sh"]
diff --git a/README.md b/README.md
index e92fb527741aae294b49e20ec05565fc20efb14b..2d02b2d9f6eb25711c58b9f7f3fe7c1fd7e488da 100644
--- a/README.md
+++ b/README.md
@@ -1,10 +1,38 @@
---
-title: Ner Span
-emoji: ๐
-colorFrom: yellow
-colorTo: purple
+title: Potato โ Span Labeling (NER)
+emoji: ๐ฅ
+colorFrom: green
+colorTo: blue
sdk: docker
+app_port: 7860
pinned: false
+license: apache-2.0
+tags:
+ - annotation
+ - potato
+ - span
+ - ner
---
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Potato โ Span Labeling (NER)
+
+Highlight-and-label text spans for NER / extraction.
+
+A live demo of [Potato](https://github.com/davidjurgens/potato), the free, self-hosted annotation platform for NLP,
+agentic, and GenAI research โ configured entirely through YAML.
+
+## Try it out
+
+1. Enter any username to log in (no password required).
+2. Read the item shown in the main panel.
+3. Annotate using the schemes on the right.
+4. Click **Next** to continue.
+
+> Annotations in this demo are ephemeral. To collect and keep data, deploy your own
+> Space โ see the [deployment guide](https://github.com/davidjurgens/potato/blob/master/deployment/huggingface-spaces/deploy.md).
+
+## About Potato
+
+Potato supports 20+ annotation types โ text, spans, images, audio, video, documents,
+and agent traces โ with AI-assisted labeling, quality control, and adjudication.
+[Learn more on GitHub](https://github.com/davidjurgens/potato) ยท [Browse all demos](https://github.com/davidjurgens/potato/blob/master/docs/data-export/potato_on_huggingface.md).
diff --git a/annotation_output/.gitkeep b/annotation_output/.gitkeep
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/config-talk.yaml b/config-talk.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..9cf46513638bec5ea6895a90f2b09cae56b466a8
--- /dev/null
+++ b/config-talk.yaml
@@ -0,0 +1,24 @@
+port: 9001
+annotation_task_name: Certainty Highlighting
+task_dir: .
+output_annotation_dir: annotation_output/talk-span/
+output_annotation_format: json
+annotation_codebook_url: ''
+data_files:
+- data/talk-certainty.csv
+item_properties:
+ id_key: id
+ text_key: text
+user_config:
+ allow_all_users: true
+ users: []
+alert_time_each_instance: 10000000
+annotation_schemes:
+- annotation_type: span
+ name: certainty
+ description: Highlight which phrases make the sentence more or less certain
+ labels:
+ - certain
+ - uncertain
+ sequential_key_binding: true
+site_dir: default
diff --git a/config.yaml b/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..38168dcdd8a9f051d982ac8994766ca78ef07507
--- /dev/null
+++ b/config.yaml
@@ -0,0 +1,24 @@
+port: 9001
+annotation_task_name: Simple Highlighting Example
+task_dir: .
+output_annotation_dir: annotation_output/simple-span-labeling/
+output_annotation_format: json
+annotation_codebook_url: ''
+data_files:
+- data/toy-example.csv
+item_properties:
+ id_key: id
+ text_key: text
+user_config:
+ allow_all_users: true
+ users: []
+alert_time_each_instance: 10000000
+annotation_schemes:
+- annotation_type: span
+ name: certainty
+ description: Highlight which phrases make the sentence more or less certain
+ labels:
+ - certain
+ - uncertain
+ sequential_key_binding: true
+site_dir: default
diff --git a/data/talk-certainty.csv b/data/talk-certainty.csv
new file mode 100644
index 0000000000000000000000000000000000000000..f7d780ad8f6e8b27e992b3dec67b3566b0568374
--- /dev/null
+++ b/data/talk-certainty.csv
@@ -0,0 +1,5 @@
+id,text
+cert_1,"The data may suggest a weak correlation, but the evidence is far from conclusive."
+cert_2,"We are confident that the new treatment significantly improves recovery times."
+cert_3,"It is unclear whether the model will generalize to previously unseen domains."
+cert_4,"These findings strongly indicate that sleep quality predicts next-day performance."
diff --git a/data/toy-example.csv b/data/toy-example.csv
new file mode 100644
index 0000000000000000000000000000000000000000..f6db583cf231ce2c080eb0732d72bc7e57dff91f
--- /dev/null
+++ b/data/toy-example.csv
@@ -0,0 +1,1001 @@
+id,text
+item_1,Aliquam eius modi ut.
+item_2,Numquam dolorem labore voluptatem dolore etincidunt neque.
+item_3,Consectetur adipisci ipsum neque est.
+item_4,Quaerat velit aliquam voluptatem etincidunt tempora.
+item_5,Dolorem aliquam dolor ut sit porro.
+item_6,Dolore dolorem velit sit consectetur etincidunt magnam.
+item_7,Aliquam tempora labore porro quaerat consectetur.
+item_8,Dolore modi sit modi dolor numquam eius numquam.
+item_9,Est tempora dolor dolor.
+item_10,Labore aliquam est est aliquam.
+item_11,Ut ipsum adipisci magnam.
+item_12,Eius labore velit dolore amet.
+item_13,Dolore est magnam eius numquam numquam modi.
+item_14,Etincidunt dolore dolorem quaerat numquam labore numquam est.
+item_15,Sit amet sed quaerat eius ut dolorem labore.
+item_16,Quiquia magnam aliquam velit sit quaerat neque.
+item_17,Amet aliquam tempora modi tempora dolorem non.
+item_18,Eius amet labore aliquam dolorem aliquam quiquia.
+item_19,Est dolore consectetur magnam porro dolorem.
+item_20,Dolorem quisquam sed dolore.
+item_21,Quiquia modi tempora etincidunt numquam modi.
+item_22,Quaerat non quiquia quiquia.
+item_23,Quiquia eius adipisci eius amet etincidunt est.
+item_24,Adipisci labore dolorem ut quisquam dolore.
+item_25,Tempora sit amet dolorem tempora numquam dolorem.
+item_26,Ipsum quisquam dolorem ut quisquam numquam.
+item_27,Amet aliquam quiquia tempora quisquam sed magnam consectetur.
+item_28,Dolore neque quaerat ut etincidunt.
+item_29,Neque quaerat amet eius neque dolorem.
+item_30,Labore quiquia sit dolorem dolorem.
+item_31,Dolorem quiquia porro magnam quiquia tempora sit voluptatem.
+item_32,Adipisci sed dolor eius.
+item_33,Ut etincidunt tempora sit tempora eius dolor.
+item_34,Quiquia numquam tempora sit ut adipisci.
+item_35,Modi sed sit eius dolore quiquia.
+item_36,Ut quisquam sit ipsum sed modi consectetur.
+item_37,Sit aliquam ipsum ipsum.
+item_38,Labore voluptatem aliquam ipsum sit etincidunt.
+item_39,Dolore ipsum neque quisquam dolorem.
+item_40,Dolore eius tempora sit consectetur.
+item_41,Dolore non non est.
+item_42,Dolore neque eius labore quiquia dolorem numquam voluptatem.
+item_43,Est est eius labore numquam magnam quaerat.
+item_44,Magnam est aliquam quisquam magnam.
+item_45,Non non sed sed magnam velit eius labore.
+item_46,Dolorem etincidunt magnam quiquia numquam.
+item_47,Magnam labore sed dolore non non aliquam.
+item_48,Eius tempora etincidunt velit velit amet sit quiquia.
+item_49,Adipisci est adipisci sit.
+item_50,Sed ipsum velit neque.
+item_51,Numquam numquam non consectetur.
+item_52,Modi consectetur porro est quiquia etincidunt sit.
+item_53,Porro quiquia voluptatem etincidunt dolorem consectetur porro.
+item_54,Eius adipisci non quisquam.
+item_55,Consectetur adipisci est etincidunt modi neque est.
+item_56,Sed est voluptatem est numquam.
+item_57,Neque quaerat magnam amet etincidunt adipisci dolorem dolor.
+item_58,Quaerat ipsum labore dolorem magnam sit dolor.
+item_59,Ut quisquam labore eius.
+item_60,Labore ut dolorem amet.
+item_61,Amet quiquia neque non velit quiquia labore porro.
+item_62,Voluptatem consectetur modi dolor quaerat quaerat voluptatem non.
+item_63,Velit eius dolore sed.
+item_64,Adipisci sit ut sit tempora quiquia.
+item_65,Amet modi neque quisquam.
+item_66,Labore modi dolor sit voluptatem consectetur modi modi.
+item_67,Dolorem labore quiquia porro voluptatem amet.
+item_68,Quisquam voluptatem dolor tempora.
+item_69,Dolor ut velit adipisci magnam neque.
+item_70,Quiquia non porro amet numquam dolore sed.
+item_71,Modi aliquam numquam dolorem labore.
+item_72,Aliquam ut voluptatem non non.
+item_73,Quisquam magnam dolor sit quaerat quisquam non eius.
+item_74,Consectetur dolor consectetur non numquam magnam voluptatem porro.
+item_75,Numquam dolor ut adipisci eius quaerat modi.
+item_76,Amet non ut ut.
+item_77,Quisquam modi magnam non modi.
+item_78,Ut aliquam etincidunt ut aliquam.
+item_79,Numquam dolore non neque dolore voluptatem numquam porro.
+item_80,Modi dolor sed quiquia.
+item_81,Consectetur ut voluptatem modi neque quisquam amet.
+item_82,Dolore neque labore dolor sed dolor quisquam dolorem.
+item_83,Neque quisquam quisquam porro porro labore quisquam.
+item_84,Etincidunt quiquia porro amet dolore.
+item_85,Sit dolore amet amet.
+item_86,Dolore etincidunt aliquam voluptatem sit sit eius labore.
+item_87,Quiquia amet magnam ut non.
+item_88,Labore adipisci numquam tempora neque quisquam sed.
+item_89,Velit modi magnam consectetur.
+item_90,Quiquia quaerat dolore modi porro dolor dolor.
+item_91,Adipisci dolorem etincidunt neque dolor modi porro eius.
+item_92,Neque sit numquam aliquam eius.
+item_93,Tempora modi tempora tempora magnam.
+item_94,Dolore modi adipisci sit adipisci ipsum tempora quiquia.
+item_95,Non quiquia quisquam tempora voluptatem ut.
+item_96,Dolor dolor etincidunt quisquam consectetur etincidunt est.
+item_97,Magnam eius amet aliquam quaerat.
+item_98,Dolore adipisci magnam porro non dolore.
+item_99,Consectetur adipisci ut adipisci velit.
+item_100,Quiquia quisquam velit ipsum eius.
+item_101,Dolorem dolore dolore porro amet amet.
+item_102,Magnam dolore etincidunt sit quaerat numquam dolore.
+item_103,Sed eius consectetur magnam consectetur numquam voluptatem quiquia.
+item_104,Tempora dolore ut ipsum dolorem aliquam non.
+item_105,Labore amet sit dolorem numquam amet quisquam labore.
+item_106,Neque quiquia numquam etincidunt.
+item_107,Amet amet consectetur dolore ipsum dolorem.
+item_108,Eius labore modi eius magnam magnam modi.
+item_109,Ut non quiquia porro.
+item_110,Consectetur porro quiquia dolore quisquam quaerat amet.
+item_111,Ut est sit labore.
+item_112,Velit aliquam modi ipsum.
+item_113,Quaerat dolorem est consectetur neque amet.
+item_114,Voluptatem dolorem sed consectetur dolore.
+item_115,Etincidunt aliquam eius modi.
+item_116,Modi quaerat numquam velit consectetur quaerat.
+item_117,Ut est etincidunt amet.
+item_118,Non velit voluptatem sed adipisci dolor modi aliquam.
+item_119,Etincidunt quaerat ipsum ipsum consectetur.
+item_120,Sed aliquam porro velit ut.
+item_121,Non ipsum ipsum dolorem etincidunt sit dolor tempora.
+item_122,Sed quiquia velit dolore numquam numquam voluptatem neque.
+item_123,Velit adipisci consectetur non amet sed numquam.
+item_124,Etincidunt quaerat etincidunt sit sed adipisci.
+item_125,Est dolore adipisci velit dolorem dolorem magnam.
+item_126,Sed magnam modi sed neque quaerat.
+item_127,Labore adipisci quisquam dolor.
+item_128,Dolor eius est labore etincidunt est neque dolore.
+item_129,Velit sit sit non quaerat voluptatem.
+item_130,Etincidunt ipsum tempora porro labore.
+item_131,Magnam ut etincidunt voluptatem sed etincidunt dolor.
+item_132,Etincidunt tempora dolorem etincidunt ipsum neque.
+item_133,Consectetur magnam sit quisquam non etincidunt sed.
+item_134,Est quisquam est magnam.
+item_135,Dolor consectetur est porro dolorem.
+item_136,Dolorem eius ut neque aliquam.
+item_137,Quisquam ipsum adipisci voluptatem tempora.
+item_138,Quaerat velit dolore neque velit.
+item_139,Non labore neque dolorem dolorem dolore quiquia.
+item_140,Non quiquia neque quisquam quisquam dolore.
+item_141,Eius numquam ipsum velit dolorem.
+item_142,Magnam sit non dolore modi quaerat dolore.
+item_143,Velit neque aliquam magnam.
+item_144,Etincidunt sed labore ipsum.
+item_145,Ut ut sed dolor ipsum amet tempora quaerat.
+item_146,Porro quisquam dolore adipisci labore est non.
+item_147,Sed sit magnam numquam.
+item_148,Ipsum eius aliquam etincidunt.
+item_149,Quisquam modi est quaerat.
+item_150,Quiquia porro quiquia quaerat.
+item_151,Amet numquam voluptatem porro modi porro est.
+item_152,Ipsum neque non ipsum sit voluptatem numquam.
+item_153,Etincidunt magnam numquam numquam sed consectetur quiquia.
+item_154,Aliquam aliquam quisquam numquam magnam ut est porro.
+item_155,Dolore consectetur magnam sit consectetur aliquam.
+item_156,Aliquam porro amet ipsum ut magnam.
+item_157,Aliquam non quiquia aliquam.
+item_158,Modi aliquam non porro amet ut ipsum.
+item_159,Adipisci quiquia sit quisquam.
+item_160,Consectetur magnam quiquia ipsum consectetur modi.
+item_161,Tempora quaerat dolore labore etincidunt.
+item_162,Eius dolor amet amet dolorem magnam quisquam.
+item_163,Quaerat quisquam magnam quaerat ipsum porro.
+item_164,Est numquam labore voluptatem magnam.
+item_165,Ut dolor neque sit eius est etincidunt.
+item_166,Modi amet magnam aliquam velit.
+item_167,Velit labore neque est ipsum amet dolor ut.
+item_168,Sed porro sit sit quaerat neque dolor.
+item_169,Quaerat modi modi dolor sed est labore.
+item_170,Dolorem tempora dolor neque sit aliquam.
+item_171,Ipsum neque numquam neque consectetur quiquia.
+item_172,Consectetur aliquam voluptatem porro amet ut numquam dolor.
+item_173,Quiquia modi ipsum neque ut eius quaerat non.
+item_174,Modi dolore eius est.
+item_175,Aliquam est ipsum sed eius labore.
+item_176,Numquam adipisci dolorem sed.
+item_177,Labore consectetur eius ipsum ipsum est velit sed.
+item_178,Modi dolorem quaerat modi.
+item_179,Ut porro numquam quaerat eius numquam.
+item_180,Amet ipsum dolor dolorem quiquia est.
+item_181,Voluptatem sit quiquia labore magnam.
+item_182,Est ipsum porro magnam non tempora magnam amet.
+item_183,Adipisci magnam labore eius quaerat.
+item_184,Dolore sit eius labore sed dolorem.
+item_185,Amet dolor sit est.
+item_186,Velit tempora aliquam non tempora neque.
+item_187,Dolorem quisquam neque magnam ipsum quisquam neque.
+item_188,Aliquam magnam modi est neque porro.
+item_189,Porro aliquam dolorem voluptatem velit dolor magnam.
+item_190,Consectetur modi consectetur est.
+item_191,Tempora etincidunt magnam amet porro sed.
+item_192,Labore modi velit etincidunt quisquam numquam velit neque.
+item_193,Dolorem tempora numquam non quiquia dolore porro.
+item_194,Dolore labore dolorem velit magnam dolore.
+item_195,Adipisci dolor sit quaerat neque dolor est.
+item_196,Tempora ipsum dolore velit voluptatem.
+item_197,Ut sed quaerat quaerat sit sit.
+item_198,Porro quiquia aliquam adipisci consectetur.
+item_199,Dolorem dolore amet dolorem dolorem.
+item_200,Velit aliquam velit adipisci modi ut modi labore.
+item_201,Tempora etincidunt tempora sed dolorem quiquia ut amet.
+item_202,Etincidunt aliquam etincidunt etincidunt labore dolore.
+item_203,Numquam magnam modi sit magnam adipisci dolor quiquia.
+item_204,Quisquam etincidunt quisquam etincidunt labore.
+item_205,Ut labore sed adipisci sit neque.
+item_206,Magnam consectetur etincidunt magnam voluptatem amet voluptatem neque.
+item_207,Ipsum neque aliquam amet.
+item_208,Eius magnam non aliquam.
+item_209,Eius neque amet adipisci non magnam sit dolor.
+item_210,Modi etincidunt aliquam neque voluptatem.
+item_211,Quaerat etincidunt magnam sed.
+item_212,Ut est velit ipsum aliquam.
+item_213,Tempora sed quiquia modi labore.
+item_214,Neque sed ut modi.
+item_215,Ipsum quiquia numquam tempora porro.
+item_216,Tempora tempora numquam non neque modi dolore est.
+item_217,Numquam dolore porro est porro.
+item_218,Porro sed quiquia tempora consectetur voluptatem non.
+item_219,Non consectetur quiquia quaerat.
+item_220,Velit etincidunt est labore ut est voluptatem.
+item_221,Adipisci tempora tempora quiquia.
+item_222,Voluptatem eius dolore consectetur.
+item_223,Dolorem aliquam dolorem velit dolorem aliquam labore quiquia.
+item_224,Sed magnam tempora neque eius eius ipsum dolore.
+item_225,Tempora sit adipisci dolore.
+item_226,Quaerat consectetur amet labore quiquia magnam dolor modi.
+item_227,Eius non quiquia dolor labore consectetur non sed.
+item_228,Dolorem dolor modi ut.
+item_229,Adipisci dolore amet amet quisquam dolor sed.
+item_230,Dolore dolor consectetur amet magnam dolor porro.
+item_231,Sed ipsum aliquam numquam.
+item_232,Aliquam sed magnam voluptatem voluptatem consectetur est dolore.
+item_233,Ipsum quisquam non dolorem tempora etincidunt velit etincidunt.
+item_234,Sit sed labore non dolorem quisquam quiquia consectetur.
+item_235,Dolore modi quiquia ipsum sit.
+item_236,Velit aliquam dolorem quaerat.
+item_237,Ipsum quisquam amet non quisquam dolor.
+item_238,Dolorem est adipisci neque quaerat.
+item_239,Quisquam amet magnam ipsum.
+item_240,Sit quiquia amet quiquia quiquia ut aliquam magnam.
+item_241,Voluptatem aliquam porro velit modi.
+item_242,Sit quisquam ut sed.
+item_243,Neque ipsum eius dolore.
+item_244,Labore aliquam velit neque amet ut.
+item_245,Modi eius modi porro est ipsum modi.
+item_246,Modi numquam voluptatem dolor numquam neque consectetur non.
+item_247,Aliquam voluptatem sit sit dolor non modi velit.
+item_248,Dolorem eius amet dolorem.
+item_249,Modi est adipisci quaerat dolore.
+item_250,Etincidunt etincidunt sit consectetur amet aliquam etincidunt amet.
+item_251,Neque ipsum porro dolorem.
+item_252,Adipisci modi modi voluptatem.
+item_253,Quaerat labore aliquam adipisci labore velit adipisci tempora.
+item_254,Sit neque tempora modi dolor dolore dolor eius.
+item_255,Porro eius eius aliquam amet dolorem.
+item_256,Modi ipsum dolorem dolor dolore magnam.
+item_257,Aliquam dolore consectetur consectetur neque est est amet.
+item_258,Amet labore voluptatem labore ipsum quisquam quisquam.
+item_259,Aliquam amet adipisci aliquam ipsum dolorem magnam.
+item_260,Quaerat eius ut velit dolorem ipsum quisquam.
+item_261,Neque porro quiquia adipisci tempora.
+item_262,Non etincidunt voluptatem quaerat sit dolore velit quaerat.
+item_263,Etincidunt sit quaerat dolorem.
+item_264,Voluptatem eius dolorem labore tempora dolore.
+item_265,Ut tempora velit sed neque dolore.
+item_266,Aliquam numquam quiquia neque quisquam quisquam.
+item_267,Labore non numquam dolorem est dolorem modi.
+item_268,Velit dolor amet etincidunt modi neque labore modi.
+item_269,Porro est voluptatem consectetur dolor.
+item_270,Dolorem quaerat neque est labore dolore eius dolore.
+item_271,Magnam amet sit quiquia dolorem.
+item_272,Dolore voluptatem velit velit.
+item_273,Eius sit magnam dolor.
+item_274,Ut eius dolorem ut dolor dolor.
+item_275,Numquam magnam adipisci sit.
+item_276,Labore ipsum etincidunt tempora.
+item_277,Aliquam velit sit eius magnam quisquam modi ut.
+item_278,Aliquam adipisci est modi non porro amet porro.
+item_279,Tempora quiquia eius velit tempora.
+item_280,Aliquam etincidunt sit adipisci modi velit modi etincidunt.
+item_281,Non neque sed neque quiquia aliquam adipisci ut.
+item_282,Sit ipsum aliquam dolor sit velit ut etincidunt.
+item_283,Est sed porro ipsum modi dolor.
+item_284,Sed numquam modi eius velit aliquam amet.
+item_285,Numquam magnam quisquam numquam aliquam voluptatem.
+item_286,Sit dolorem dolorem numquam.
+item_287,Voluptatem labore aliquam velit consectetur.
+item_288,Consectetur dolore ipsum quiquia numquam ipsum.
+item_289,Est voluptatem ut adipisci labore.
+item_290,Dolorem velit consectetur est.
+item_291,Adipisci sed consectetur ipsum ipsum ut non adipisci.
+item_292,Modi aliquam aliquam ut dolor sed.
+item_293,Dolor neque adipisci sed dolore quaerat neque adipisci.
+item_294,Adipisci consectetur aliquam quiquia labore.
+item_295,Tempora neque est neque tempora.
+item_296,Voluptatem voluptatem velit tempora.
+item_297,Dolorem labore aliquam dolor quisquam sit dolore.
+item_298,Magnam numquam numquam est ipsum non quisquam adipisci.
+item_299,Etincidunt etincidunt amet amet.
+item_300,Amet voluptatem eius consectetur labore modi.
+item_301,Numquam ut est quaerat adipisci sit voluptatem amet.
+item_302,Ipsum ipsum numquam adipisci.
+item_303,Ipsum velit dolore numquam.
+item_304,Velit eius numquam neque.
+item_305,Adipisci sed eius etincidunt tempora modi quiquia.
+item_306,Labore eius porro aliquam.
+item_307,Dolor quiquia voluptatem non sed quisquam.
+item_308,Dolorem velit velit sit dolore porro neque.
+item_309,Dolore non est amet aliquam.
+item_310,Velit porro dolorem quaerat consectetur sed porro.
+item_311,Numquam numquam modi numquam velit quaerat.
+item_312,Dolor est aliquam quiquia.
+item_313,Magnam consectetur consectetur dolor dolorem aliquam dolor.
+item_314,Quisquam aliquam est dolor neque adipisci eius ut.
+item_315,Ipsum neque etincidunt numquam ipsum tempora.
+item_316,Amet dolore tempora etincidunt sed eius.
+item_317,Ut sed numquam numquam ipsum non non.
+item_318,Numquam velit sit est adipisci neque.
+item_319,Magnam aliquam dolor amet dolore amet amet.
+item_320,Sit neque consectetur velit voluptatem porro.
+item_321,Eius consectetur aliquam dolore velit quisquam labore eius.
+item_322,Velit neque quaerat numquam est modi dolore.
+item_323,Magnam neque tempora etincidunt sit numquam velit.
+item_324,Velit velit sit aliquam.
+item_325,Labore porro modi dolore.
+item_326,Etincidunt ut quaerat eius est est.
+item_327,Magnam eius eius numquam tempora.
+item_328,Adipisci velit voluptatem neque aliquam porro dolor voluptatem.
+item_329,Labore est sed velit.
+item_330,Tempora tempora porro voluptatem porro.
+item_331,Etincidunt sed dolore eius aliquam.
+item_332,Adipisci labore dolor magnam.
+item_333,Neque etincidunt aliquam quisquam labore quisquam amet.
+item_334,Eius consectetur consectetur voluptatem non.
+item_335,Non magnam dolor amet sit.
+item_336,Ipsum consectetur aliquam quisquam eius quaerat eius tempora.
+item_337,Non velit magnam sed eius neque.
+item_338,Non ut labore voluptatem magnam adipisci voluptatem quiquia.
+item_339,Quiquia dolore ut voluptatem labore voluptatem.
+item_340,Non non est dolore quisquam ut dolorem voluptatem.
+item_341,Sed etincidunt etincidunt voluptatem amet non etincidunt ut.
+item_342,Amet labore adipisci dolore quiquia dolorem velit.
+item_343,Magnam modi quiquia sit modi adipisci.
+item_344,Magnam velit est dolor est tempora.
+item_345,Magnam magnam quisquam aliquam.
+item_346,Eius est consectetur velit amet amet.
+item_347,Neque voluptatem magnam labore ut.
+item_348,Dolore dolorem numquam tempora adipisci ipsum adipisci dolorem.
+item_349,Sit velit ut voluptatem.
+item_350,Labore adipisci non ut dolor.
+item_351,Tempora dolorem quisquam sit.
+item_352,Adipisci quiquia voluptatem voluptatem ut non.
+item_353,Numquam quisquam ut neque velit.
+item_354,Porro sit sit non quisquam voluptatem voluptatem neque.
+item_355,Quaerat ipsum numquam voluptatem ipsum dolor adipisci.
+item_356,Magnam quiquia consectetur neque quisquam sed ipsum.
+item_357,Sed eius eius velit adipisci est tempora quaerat.
+item_358,Neque dolorem velit ipsum eius sed quiquia.
+item_359,Est ipsum sit numquam non non eius.
+item_360,Porro amet modi ut quaerat adipisci voluptatem.
+item_361,Velit tempora voluptatem numquam.
+item_362,Quaerat eius labore tempora sed.
+item_363,Labore quiquia adipisci ipsum ut non.
+item_364,Etincidunt quiquia etincidunt consectetur dolore eius ipsum est.
+item_365,Porro aliquam quisquam aliquam ut.
+item_366,Magnam dolorem eius ipsum.
+item_367,Ut eius adipisci velit voluptatem ipsum est.
+item_368,Modi est tempora neque.
+item_369,Dolorem dolor velit magnam adipisci.
+item_370,Etincidunt dolore amet sed consectetur non adipisci quisquam.
+item_371,Ut numquam labore eius quiquia dolor velit ut.
+item_372,Numquam est numquam est eius.
+item_373,Eius dolorem amet dolor ut voluptatem neque neque.
+item_374,Quiquia etincidunt magnam dolore quisquam dolore.
+item_375,Etincidunt porro adipisci tempora dolore amet adipisci.
+item_376,Quiquia est voluptatem quiquia neque tempora.
+item_377,Etincidunt quaerat quiquia est non dolore.
+item_378,Sit non etincidunt ut porro quiquia amet.
+item_379,Non dolor etincidunt est.
+item_380,Velit non etincidunt consectetur.
+item_381,Labore adipisci dolore etincidunt modi.
+item_382,Dolor quaerat labore quaerat dolorem neque etincidunt.
+item_383,Etincidunt ut numquam quaerat.
+item_384,Ipsum etincidunt quaerat adipisci voluptatem magnam.
+item_385,Consectetur etincidunt sit consectetur quiquia ipsum.
+item_386,Voluptatem velit dolorem sed.
+item_387,Dolore modi porro labore quisquam dolore.
+item_388,Eius adipisci velit consectetur porro porro tempora quiquia.
+item_389,Adipisci labore est amet sit sit tempora sed.
+item_390,Porro non velit labore dolore quisquam numquam labore.
+item_391,Non sed aliquam magnam eius tempora eius.
+item_392,Labore quiquia sit aliquam.
+item_393,Est adipisci magnam quisquam sit etincidunt quiquia sed.
+item_394,Consectetur velit modi eius magnam non neque tempora.
+item_395,Aliquam eius magnam eius amet eius porro.
+item_396,Consectetur quisquam quisquam amet quaerat eius.
+item_397,Tempora consectetur tempora voluptatem dolore magnam consectetur.
+item_398,Adipisci quiquia labore dolorem tempora.
+item_399,Dolore consectetur ut adipisci ipsum magnam modi etincidunt.
+item_400,Sit eius modi dolore adipisci voluptatem amet sit.
+item_401,Ipsum adipisci dolore quisquam.
+item_402,Aliquam labore quaerat ipsum.
+item_403,Dolor etincidunt modi eius sit tempora.
+item_404,Ipsum non dolore eius modi sit adipisci quiquia.
+item_405,Velit ipsum etincidunt quaerat dolore.
+item_406,Dolore tempora neque est.
+item_407,Quaerat consectetur sed amet quaerat quisquam ut neque.
+item_408,Modi porro numquam ipsum ipsum neque.
+item_409,Dolor ut quiquia neque adipisci est.
+item_410,Ut consectetur est amet numquam modi porro.
+item_411,Ut modi magnam sed magnam voluptatem sed sit.
+item_412,Porro porro magnam sit aliquam adipisci consectetur velit.
+item_413,Ut sed sed eius sed voluptatem.
+item_414,Voluptatem non aliquam numquam consectetur numquam adipisci.
+item_415,Quaerat eius numquam etincidunt tempora.
+item_416,Velit neque modi est magnam etincidunt.
+item_417,Non adipisci voluptatem modi.
+item_418,Ipsum neque consectetur eius magnam voluptatem aliquam.
+item_419,Magnam dolor eius eius non neque.
+item_420,Sit etincidunt etincidunt quiquia labore eius eius.
+item_421,Dolore quiquia magnam labore consectetur voluptatem dolore quiquia.
+item_422,Adipisci modi ut etincidunt adipisci porro quiquia.
+item_423,Eius consectetur velit dolor.
+item_424,Numquam consectetur numquam quaerat quaerat.
+item_425,Sed tempora numquam sed modi consectetur.
+item_426,Velit ut consectetur ut dolore sit sit.
+item_427,Quisquam eius sed dolor neque neque.
+item_428,Dolore numquam dolorem amet dolorem quaerat.
+item_429,Non sed sed quaerat modi quiquia.
+item_430,Eius aliquam non consectetur quiquia.
+item_431,Quaerat sit etincidunt voluptatem.
+item_432,Consectetur porro non sed porro consectetur.
+item_433,Eius ipsum est porro neque porro.
+item_434,Porro non dolorem ipsum ut velit.
+item_435,Amet porro ipsum ipsum labore est.
+item_436,Sed sed aliquam porro adipisci modi.
+item_437,Velit quisquam neque tempora labore.
+item_438,Sit adipisci etincidunt consectetur magnam.
+item_439,Sit tempora tempora velit.
+item_440,Eius dolorem velit dolor.
+item_441,Sit quaerat aliquam magnam quisquam quiquia.
+item_442,Numquam porro ipsum dolor.
+item_443,Eius voluptatem labore porro quisquam magnam modi.
+item_444,Neque quisquam velit voluptatem dolorem quisquam consectetur amet.
+item_445,Consectetur adipisci voluptatem consectetur aliquam est ut non.
+item_446,Neque quaerat etincidunt labore etincidunt amet dolorem.
+item_447,Voluptatem sed velit sed non velit porro magnam.
+item_448,Labore numquam eius aliquam est.
+item_449,Sed etincidunt modi numquam sit quaerat adipisci.
+item_450,Neque non ut ut dolore.
+item_451,Non etincidunt non amet etincidunt adipisci quiquia.
+item_452,Porro dolorem quisquam non quiquia.
+item_453,Voluptatem numquam voluptatem quiquia amet tempora ut.
+item_454,Consectetur tempora aliquam velit dolore modi.
+item_455,Adipisci neque labore quiquia non quisquam quisquam numquam.
+item_456,Voluptatem quisquam ipsum ut dolor ipsum.
+item_457,Sit quiquia ut quisquam ut.
+item_458,Labore modi dolore magnam etincidunt aliquam.
+item_459,Non aliquam consectetur quisquam est numquam.
+item_460,Tempora voluptatem tempora consectetur.
+item_461,Porro numquam etincidunt voluptatem.
+item_462,Magnam tempora eius labore.
+item_463,Numquam neque quisquam magnam.
+item_464,Dolor dolore neque voluptatem dolorem voluptatem consectetur porro.
+item_465,Non eius consectetur sit ut.
+item_466,Numquam modi consectetur labore.
+item_467,Neque consectetur ipsum porro.
+item_468,Aliquam adipisci dolorem voluptatem quisquam sit dolor.
+item_469,Ut velit quaerat consectetur sit labore.
+item_470,Dolore eius etincidunt modi adipisci.
+item_471,Adipisci dolor dolor quisquam non amet voluptatem.
+item_472,Aliquam magnam dolorem quaerat dolore.
+item_473,Aliquam dolore quiquia adipisci aliquam amet non quiquia.
+item_474,Eius quisquam modi modi.
+item_475,Quisquam quisquam adipisci sit sed dolore aliquam.
+item_476,Sed ut quaerat quiquia voluptatem.
+item_477,Ipsum adipisci dolore quisquam voluptatem numquam sed sit.
+item_478,Est quisquam aliquam numquam voluptatem sit porro.
+item_479,Aliquam ut aliquam ut neque.
+item_480,Ut magnam non est adipisci.
+item_481,Ut porro est neque tempora modi.
+item_482,Quisquam dolor non aliquam tempora quiquia modi numquam.
+item_483,Ut eius aliquam adipisci dolore labore.
+item_484,Amet porro modi neque voluptatem neque quiquia aliquam.
+item_485,Adipisci sed velit dolorem modi velit numquam.
+item_486,Dolorem dolor est ut ut ut dolor.
+item_487,Non numquam dolorem consectetur adipisci.
+item_488,Eius dolorem non dolorem numquam sed quisquam.
+item_489,Ut sed velit ut numquam etincidunt modi modi.
+item_490,Voluptatem tempora est quisquam.
+item_491,Dolor ipsum dolor neque dolor amet.
+item_492,Consectetur sit sit tempora voluptatem.
+item_493,Quiquia sit tempora quaerat aliquam.
+item_494,Labore labore aliquam consectetur numquam.
+item_495,Velit quisquam sit sit aliquam modi aliquam.
+item_496,Velit est etincidunt adipisci adipisci dolore.
+item_497,Porro ipsum ut voluptatem.
+item_498,Consectetur dolor modi sed.
+item_499,Dolor quisquam aliquam consectetur quaerat etincidunt.
+item_500,Magnam neque quiquia non numquam est neque labore.
+item_501,Est magnam numquam amet quaerat dolore.
+item_502,Dolorem tempora labore sit.
+item_503,Quaerat numquam velit sit.
+item_504,Magnam eius amet amet quisquam.
+item_505,Etincidunt modi sit neque ipsum sed labore.
+item_506,Amet etincidunt quaerat sed non dolore.
+item_507,Velit numquam adipisci dolor etincidunt dolorem.
+item_508,Sed labore sit quiquia ipsum voluptatem quiquia.
+item_509,Neque consectetur velit labore ut velit sed amet.
+item_510,Numquam velit neque dolore porro adipisci voluptatem.
+item_511,Tempora non quisquam amet eius non aliquam.
+item_512,Magnam quisquam porro eius.
+item_513,Est velit porro velit quisquam porro tempora.
+item_514,Quisquam porro modi dolore labore est.
+item_515,Dolorem aliquam dolore neque.
+item_516,Quaerat ut neque sed magnam aliquam etincidunt quiquia.
+item_517,Dolore sed sed neque magnam sit voluptatem dolor.
+item_518,Dolore etincidunt dolore adipisci.
+item_519,Modi labore quaerat numquam voluptatem non sed.
+item_520,Sed quisquam labore sed.
+item_521,Ipsum ipsum neque etincidunt neque.
+item_522,Tempora tempora tempora adipisci dolore sit.
+item_523,Eius amet consectetur quisquam non aliquam non.
+item_524,Quaerat dolorem ipsum magnam aliquam magnam amet non.
+item_525,Labore dolorem magnam dolor.
+item_526,Quaerat dolor ipsum sed sit voluptatem porro.
+item_527,Tempora ut labore tempora porro numquam adipisci non.
+item_528,Consectetur numquam porro ipsum ipsum velit etincidunt.
+item_529,Non eius dolorem quisquam amet sed.
+item_530,Tempora sed quaerat quiquia.
+item_531,Neque sit tempora magnam.
+item_532,Consectetur amet aliquam ut numquam est.
+item_533,Consectetur etincidunt quisquam dolor.
+item_534,Eius dolor aliquam adipisci.
+item_535,Velit modi sit quisquam.
+item_536,Ut amet labore numquam adipisci.
+item_537,Voluptatem adipisci sit sed ipsum numquam etincidunt.
+item_538,Sed neque quisquam ipsum ipsum porro.
+item_539,Labore sit dolorem ut porro.
+item_540,Ut adipisci voluptatem tempora velit est amet.
+item_541,Est modi quiquia quaerat sit quaerat.
+item_542,Eius est quaerat dolorem voluptatem magnam sit labore.
+item_543,Labore ipsum voluptatem est velit magnam ut.
+item_544,Numquam modi dolore est.
+item_545,Est dolor neque dolor voluptatem etincidunt.
+item_546,Tempora neque voluptatem modi quiquia adipisci.
+item_547,Non etincidunt adipisci voluptatem quisquam.
+item_548,Est amet eius porro adipisci.
+item_549,Etincidunt adipisci etincidunt dolor porro labore.
+item_550,Magnam etincidunt porro dolorem velit ut numquam.
+item_551,Est quiquia sit dolorem.
+item_552,Quaerat aliquam neque labore modi est velit modi.
+item_553,Dolore labore est voluptatem.
+item_554,Quaerat quaerat voluptatem tempora adipisci.
+item_555,Quiquia dolorem aliquam quisquam modi non.
+item_556,Aliquam labore tempora adipisci non consectetur.
+item_557,Aliquam etincidunt quaerat amet.
+item_558,Porro non voluptatem eius.
+item_559,Sed dolore dolor dolorem sit ut tempora.
+item_560,Dolorem quiquia quaerat modi dolore tempora.
+item_561,Aliquam sit etincidunt amet.
+item_562,Dolore est dolor ipsum voluptatem adipisci tempora.
+item_563,Velit quisquam tempora aliquam.
+item_564,Neque dolor tempora numquam aliquam modi porro.
+item_565,Sed velit magnam non voluptatem modi.
+item_566,Ut consectetur voluptatem modi.
+item_567,Ipsum magnam non sed porro.
+item_568,Dolor ipsum amet adipisci numquam ut.
+item_569,Amet aliquam neque magnam.
+item_570,Eius velit consectetur ipsum quaerat etincidunt est.
+item_571,Amet sit consectetur consectetur amet.
+item_572,Velit magnam dolorem quiquia non dolor sed.
+item_573,Consectetur eius ipsum sed aliquam porro.
+item_574,Sit voluptatem etincidunt voluptatem adipisci.
+item_575,Ut labore ipsum amet.
+item_576,Quiquia non sit amet velit sit.
+item_577,Porro dolorem adipisci labore adipisci.
+item_578,Non aliquam ut dolore neque sit est dolore.
+item_579,Dolore velit non sit adipisci.
+item_580,Quiquia amet quiquia adipisci.
+item_581,Est magnam eius ipsum.
+item_582,Quiquia magnam adipisci voluptatem amet modi tempora porro.
+item_583,Adipisci ut quaerat non dolorem magnam numquam amet.
+item_584,Adipisci quiquia velit sit dolore adipisci.
+item_585,Amet modi dolorem magnam quisquam modi.
+item_586,Eius sit modi sed.
+item_587,Numquam dolorem aliquam quaerat consectetur.
+item_588,Quiquia modi quiquia non.
+item_589,Quisquam porro magnam sed dolorem dolor velit.
+item_590,Non ipsum porro eius dolorem.
+item_591,Non quaerat quiquia neque non adipisci velit numquam.
+item_592,Quisquam est quaerat eius.
+item_593,Sit est quaerat sed.
+item_594,Dolore sed aliquam ut modi.
+item_595,Voluptatem sit est quisquam dolorem.
+item_596,Numquam numquam non adipisci sit consectetur dolore.
+item_597,Voluptatem velit consectetur etincidunt quisquam est.
+item_598,Eius est non numquam.
+item_599,Dolore quaerat dolor neque sit.
+item_600,Est labore voluptatem quiquia magnam dolorem dolorem.
+item_601,Sed quisquam aliquam quiquia.
+item_602,Est adipisci magnam modi magnam.
+item_603,Labore consectetur sit eius magnam adipisci.
+item_604,Adipisci aliquam sit porro eius.
+item_605,Aliquam quiquia eius quisquam sed.
+item_606,Voluptatem est est labore.
+item_607,Dolor dolore adipisci ut adipisci amet dolore.
+item_608,Aliquam magnam quiquia labore tempora ipsum.
+item_609,Dolor dolorem tempora consectetur eius quaerat dolorem.
+item_610,Adipisci labore quaerat ipsum quisquam sed.
+item_611,Aliquam adipisci amet ut quiquia amet.
+item_612,Consectetur porro ut modi voluptatem quisquam est.
+item_613,Porro quisquam eius dolore.
+item_614,Sed sed neque etincidunt sit non dolorem quaerat.
+item_615,Adipisci dolore dolor adipisci dolor modi dolorem consectetur.
+item_616,Numquam labore dolore aliquam dolorem voluptatem quisquam.
+item_617,Tempora porro est tempora dolore dolor.
+item_618,Voluptatem velit dolorem velit voluptatem non velit numquam.
+item_619,Ut quisquam dolorem adipisci numquam quaerat amet.
+item_620,Neque magnam voluptatem modi aliquam est quisquam etincidunt.
+item_621,Eius quiquia adipisci labore.
+item_622,Numquam velit voluptatem dolor non etincidunt porro quaerat.
+item_623,Aliquam magnam adipisci modi labore.
+item_624,Amet ipsum adipisci etincidunt sed magnam neque aliquam.
+item_625,Adipisci amet sit modi quisquam amet sit.
+item_626,Dolore quiquia porro voluptatem.
+item_627,Labore quaerat tempora dolorem tempora voluptatem sit porro.
+item_628,Porro numquam amet quisquam quiquia ut.
+item_629,Ut voluptatem amet eius ut est.
+item_630,Aliquam sit quisquam etincidunt dolore etincidunt.
+item_631,Modi dolore ut modi adipisci.
+item_632,Ipsum ipsum etincidunt neque adipisci.
+item_633,Tempora etincidunt adipisci dolorem velit dolor quiquia quaerat.
+item_634,Velit dolor porro etincidunt numquam.
+item_635,Velit dolorem voluptatem est.
+item_636,Etincidunt etincidunt dolor quiquia modi dolorem dolorem quisquam.
+item_637,Est consectetur quisquam ipsum dolore.
+item_638,Consectetur ipsum ut numquam sed ut porro.
+item_639,Ut modi dolorem aliquam velit velit.
+item_640,Velit eius eius quisquam numquam modi magnam.
+item_641,Quiquia non ipsum dolore.
+item_642,Est sit sit ut.
+item_643,Modi dolore amet neque est amet numquam.
+item_644,Aliquam adipisci numquam dolore voluptatem voluptatem adipisci.
+item_645,Quiquia sed porro non numquam quaerat numquam.
+item_646,Adipisci magnam amet etincidunt.
+item_647,Ut est non sit magnam etincidunt.
+item_648,Ipsum etincidunt tempora dolore quiquia.
+item_649,Ipsum adipisci velit numquam numquam sit est.
+item_650,Tempora est quaerat est sit eius etincidunt etincidunt.
+item_651,Numquam est aliquam magnam ut dolorem eius.
+item_652,Velit tempora est etincidunt.
+item_653,Quiquia porro quisquam porro neque quiquia etincidunt.
+item_654,Eius quiquia labore sit consectetur modi ut.
+item_655,Magnam adipisci dolore amet quiquia dolorem dolor.
+item_656,Eius quaerat eius labore voluptatem quaerat aliquam.
+item_657,Dolore eius numquam modi velit.
+item_658,Aliquam sit amet ipsum ipsum modi.
+item_659,Porro consectetur eius aliquam.
+item_660,Numquam dolor etincidunt dolorem neque.
+item_661,Ipsum est amet consectetur neque porro dolore.
+item_662,Quiquia tempora dolorem dolor.
+item_663,Aliquam voluptatem neque magnam neque tempora.
+item_664,Non quisquam quaerat velit ut dolore numquam.
+item_665,Aliquam amet eius dolore ut est.
+item_666,Adipisci quaerat eius etincidunt aliquam.
+item_667,Numquam quisquam sit etincidunt ut quaerat etincidunt quisquam.
+item_668,Labore modi dolor quiquia.
+item_669,Dolorem sit voluptatem neque sit.
+item_670,Porro tempora adipisci magnam porro consectetur.
+item_671,Numquam neque aliquam numquam.
+item_672,Porro magnam quaerat tempora.
+item_673,Adipisci amet dolorem quaerat dolorem tempora modi.
+item_674,Quisquam tempora dolorem dolor amet quiquia adipisci quisquam.
+item_675,Etincidunt dolorem non amet amet dolorem.
+item_676,Quiquia eius modi adipisci est sit ut.
+item_677,Neque modi voluptatem dolore.
+item_678,Dolor amet quaerat numquam quisquam.
+item_679,Sed porro consectetur velit.
+item_680,Modi dolor dolorem quiquia dolor quisquam adipisci.
+item_681,Modi labore consectetur voluptatem ut est.
+item_682,Modi quiquia dolore ipsum voluptatem dolorem.
+item_683,Labore adipisci non sed ut neque aliquam eius.
+item_684,Eius sed labore eius magnam.
+item_685,Amet dolorem tempora quiquia.
+item_686,Sed voluptatem magnam magnam labore.
+item_687,Porro adipisci tempora quisquam numquam magnam adipisci.
+item_688,Amet quiquia modi etincidunt neque.
+item_689,Dolor aliquam labore adipisci est neque tempora.
+item_690,Porro dolorem voluptatem porro dolor.
+item_691,Eius adipisci quaerat consectetur dolore eius eius eius.
+item_692,Modi voluptatem est quaerat quisquam voluptatem velit etincidunt.
+item_693,Est magnam aliquam modi voluptatem ipsum.
+item_694,Sed magnam quiquia dolore.
+item_695,Sed tempora neque eius modi labore porro.
+item_696,Tempora dolore ut tempora.
+item_697,Aliquam neque est sed velit voluptatem adipisci.
+item_698,Quisquam neque sit ipsum est.
+item_699,Dolor magnam quaerat amet ut est.
+item_700,Neque est sit quaerat sed magnam.
+item_701,Magnam aliquam adipisci voluptatem dolore amet ipsum.
+item_702,Dolorem dolorem aliquam dolore numquam quaerat voluptatem sed.
+item_703,Neque est quaerat eius aliquam eius est quaerat.
+item_704,Dolore eius ipsum numquam.
+item_705,Non etincidunt est dolor est.
+item_706,Modi magnam dolore voluptatem dolore.
+item_707,Dolore quaerat ipsum quisquam neque velit.
+item_708,Magnam est voluptatem porro velit.
+item_709,Quisquam quiquia voluptatem dolorem tempora porro consectetur.
+item_710,Tempora tempora est eius magnam quaerat.
+item_711,Quiquia tempora consectetur est quiquia dolore dolorem numquam.
+item_712,Aliquam neque ipsum porro aliquam.
+item_713,Quiquia adipisci adipisci numquam.
+item_714,Etincidunt consectetur dolore tempora quisquam.
+item_715,Est etincidunt tempora quaerat dolorem velit neque non.
+item_716,Numquam etincidunt adipisci dolore dolore consectetur labore.
+item_717,Non tempora tempora sed est.
+item_718,Voluptatem tempora labore consectetur voluptatem consectetur.
+item_719,Quisquam voluptatem aliquam quaerat modi adipisci.
+item_720,Dolor neque dolorem eius sit etincidunt.
+item_721,Modi labore aliquam adipisci modi est.
+item_722,Neque ipsum amet tempora tempora ipsum.
+item_723,Magnam amet sit aliquam dolor.
+item_724,Sit ut voluptatem quaerat.
+item_725,Dolorem sit dolorem sed.
+item_726,Velit non modi magnam.
+item_727,Dolor labore dolore dolore amet eius.
+item_728,Tempora magnam ut voluptatem modi voluptatem.
+item_729,Adipisci etincidunt porro dolorem eius non est sed.
+item_730,Est magnam consectetur ipsum.
+item_731,Ipsum modi consectetur ut adipisci porro.
+item_732,Dolorem quiquia aliquam sed tempora neque sit etincidunt.
+item_733,Dolore quiquia ut ipsum consectetur.
+item_734,Dolorem neque porro magnam quiquia numquam dolorem amet.
+item_735,Quiquia quisquam tempora tempora.
+item_736,Consectetur voluptatem est amet voluptatem voluptatem voluptatem porro.
+item_737,Quisquam dolor dolorem dolore velit dolor amet est.
+item_738,Magnam porro numquam eius dolorem dolor tempora.
+item_739,Velit numquam eius labore sit numquam consectetur dolor.
+item_740,Ut neque magnam etincidunt quaerat adipisci labore.
+item_741,Neque neque ut dolore neque.
+item_742,Voluptatem ut numquam quaerat porro.
+item_743,Neque velit adipisci amet est consectetur.
+item_744,Dolorem ut est consectetur.
+item_745,Neque ut etincidunt sed numquam.
+item_746,Sit quaerat ut eius modi consectetur magnam.
+item_747,Tempora adipisci etincidunt porro ipsum quaerat tempora.
+item_748,Neque eius ipsum est dolorem numquam labore est.
+item_749,Est adipisci dolor quiquia labore modi consectetur dolore.
+item_750,Porro voluptatem est sit dolor.
+item_751,Sed amet eius sed.
+item_752,Numquam dolorem neque dolorem.
+item_753,Porro aliquam sit ipsum eius etincidunt.
+item_754,Ut non tempora ut.
+item_755,Aliquam quaerat neque eius modi magnam quaerat.
+item_756,Est dolor aliquam dolorem numquam.
+item_757,Est etincidunt adipisci quiquia voluptatem.
+item_758,Dolor quaerat modi ut velit magnam adipisci voluptatem.
+item_759,Non tempora dolorem etincidunt est sed.
+item_760,Numquam amet etincidunt eius etincidunt.
+item_761,Ut modi dolorem magnam neque porro adipisci sed.
+item_762,Numquam sed etincidunt aliquam sit.
+item_763,Labore numquam tempora ipsum.
+item_764,Eius aliquam dolorem sit numquam porro magnam quisquam.
+item_765,Etincidunt etincidunt adipisci dolorem dolore.
+item_766,Quisquam sed adipisci sed velit.
+item_767,Dolore aliquam tempora ut sed ipsum tempora quiquia.
+item_768,Modi sit velit etincidunt.
+item_769,Porro tempora consectetur tempora ipsum.
+item_770,Ipsum quiquia consectetur neque velit numquam ut quisquam.
+item_771,Dolore dolorem sit labore voluptatem sit.
+item_772,Quaerat voluptatem neque quisquam dolore non amet tempora.
+item_773,Eius non sed etincidunt magnam.
+item_774,Dolore dolor porro neque consectetur porro velit.
+item_775,Numquam quisquam tempora adipisci.
+item_776,Ipsum porro amet quaerat numquam dolore tempora.
+item_777,Voluptatem amet consectetur labore.
+item_778,Aliquam dolor voluptatem est ut.
+item_779,Adipisci dolore consectetur dolor quiquia.
+item_780,Dolore non etincidunt numquam ut tempora.
+item_781,Sed dolorem quiquia sit dolore.
+item_782,Sit ut dolor adipisci quisquam sed.
+item_783,Non neque quiquia quisquam sit etincidunt non.
+item_784,Aliquam sed tempora adipisci sed.
+item_785,Adipisci non porro dolor eius voluptatem porro.
+item_786,Dolorem consectetur magnam porro etincidunt.
+item_787,Tempora eius sed porro sed neque.
+item_788,Modi dolorem quaerat etincidunt ipsum modi tempora aliquam.
+item_789,Dolorem adipisci aliquam neque eius adipisci labore.
+item_790,Consectetur quisquam dolore dolor sed dolorem dolorem.
+item_791,Sit amet etincidunt dolore ut quiquia aliquam.
+item_792,Dolore dolore non numquam quisquam dolore sed.
+item_793,Neque modi magnam adipisci etincidunt voluptatem ut tempora.
+item_794,Numquam consectetur etincidunt dolorem.
+item_795,Eius quiquia quiquia porro neque.
+item_796,Aliquam numquam velit adipisci consectetur.
+item_797,Quaerat consectetur sed ut quiquia quiquia.
+item_798,Ipsum sit etincidunt dolorem etincidunt dolorem etincidunt.
+item_799,Velit voluptatem dolore neque quisquam magnam quiquia.
+item_800,Amet est amet quiquia neque non tempora.
+item_801,Voluptatem consectetur modi quiquia velit modi.
+item_802,Neque aliquam etincidunt aliquam sit.
+item_803,Est quaerat dolore labore etincidunt eius voluptatem.
+item_804,Ut dolorem est porro.
+item_805,Quaerat etincidunt aliquam labore est velit quaerat.
+item_806,Neque magnam adipisci adipisci labore dolore dolore.
+item_807,Aliquam porro quiquia adipisci labore amet.
+item_808,Quisquam dolore consectetur quisquam neque magnam aliquam.
+item_809,Adipisci consectetur amet sed quiquia quisquam ipsum.
+item_810,Velit sit adipisci ut quiquia voluptatem.
+item_811,Dolorem adipisci dolor dolorem.
+item_812,Voluptatem quiquia numquam dolor dolore magnam.
+item_813,Tempora numquam dolore sed dolor.
+item_814,Tempora modi numquam porro velit quisquam.
+item_815,Velit quiquia voluptatem amet ipsum neque quisquam consectetur.
+item_816,Aliquam eius eius quiquia porro voluptatem.
+item_817,Est numquam magnam dolorem.
+item_818,Ipsum neque quiquia etincidunt etincidunt dolore amet.
+item_819,Numquam sed dolorem est adipisci non.
+item_820,Dolor est modi non.
+item_821,Quisquam amet sed non.
+item_822,Etincidunt numquam tempora amet sit voluptatem dolorem quisquam.
+item_823,Porro porro ipsum non numquam dolore.
+item_824,Amet ut eius porro aliquam.
+item_825,Numquam voluptatem voluptatem consectetur.
+item_826,Quisquam quisquam neque dolore adipisci sed.
+item_827,Aliquam adipisci quaerat ipsum.
+item_828,Quaerat est modi dolorem ut.
+item_829,Dolor sit dolor quiquia quaerat.
+item_830,Modi numquam quaerat magnam.
+item_831,Voluptatem amet dolor quisquam voluptatem etincidunt consectetur aliquam.
+item_832,Modi amet ut dolorem tempora.
+item_833,Modi labore quisquam velit dolore.
+item_834,Numquam sit quaerat ut quisquam.
+item_835,Tempora modi amet adipisci quiquia.
+item_836,Eius ut ipsum etincidunt.
+item_837,Magnam dolore modi quiquia aliquam tempora.
+item_838,Consectetur voluptatem consectetur labore sed neque dolore.
+item_839,Ipsum magnam labore est.
+item_840,Ut consectetur dolorem sed labore.
+item_841,Quiquia non amet ut aliquam est ipsum est.
+item_842,Dolorem numquam dolorem etincidunt.
+item_843,Porro magnam est etincidunt sit magnam quisquam dolorem.
+item_844,Quiquia quiquia neque modi ipsum ut dolore.
+item_845,Porro sit numquam amet ipsum amet non voluptatem.
+item_846,Etincidunt etincidunt dolore dolorem.
+item_847,Sit dolorem ut adipisci eius sed adipisci.
+item_848,Quaerat dolor velit voluptatem numquam etincidunt est.
+item_849,Magnam sed est quiquia.
+item_850,Aliquam labore ipsum amet.
+item_851,Est numquam neque dolore.
+item_852,Ut amet quisquam ipsum non.
+item_853,Consectetur aliquam voluptatem modi non ut.
+item_854,Ut sed modi quaerat.
+item_855,Dolor voluptatem dolore dolor.
+item_856,Aliquam etincidunt est ipsum consectetur numquam sed ipsum.
+item_857,Adipisci etincidunt adipisci velit ut etincidunt sit porro.
+item_858,Quaerat quaerat ut numquam quisquam ipsum quaerat aliquam.
+item_859,Adipisci eius sit est voluptatem amet.
+item_860,Dolorem labore adipisci dolore neque.
+item_861,Dolore dolor velit amet quisquam.
+item_862,Etincidunt modi sed adipisci ipsum etincidunt ut adipisci.
+item_863,Ipsum consectetur sit neque aliquam porro.
+item_864,Ut modi tempora tempora voluptatem quiquia.
+item_865,Sit dolor est quiquia aliquam modi.
+item_866,Eius voluptatem magnam quisquam sed.
+item_867,Voluptatem aliquam labore numquam dolore dolore aliquam quaerat.
+item_868,Numquam voluptatem amet voluptatem adipisci tempora.
+item_869,Quisquam sed neque quisquam voluptatem.
+item_870,Dolor dolore aliquam ipsum quisquam neque ut consectetur.
+item_871,Eius quiquia aliquam porro velit modi modi consectetur.
+item_872,Eius tempora aliquam etincidunt.
+item_873,Labore aliquam eius sit modi dolor.
+item_874,Quaerat ut dolor eius dolor quiquia.
+item_875,Eius ipsum est numquam dolore neque aliquam adipisci.
+item_876,Amet sed amet magnam amet labore.
+item_877,Est aliquam dolorem ipsum.
+item_878,Sed ipsum quaerat quisquam voluptatem est.
+item_879,Numquam quiquia neque etincidunt eius aliquam velit porro.
+item_880,Tempora quisquam dolorem ipsum.
+item_881,Non ut modi porro est neque consectetur.
+item_882,Consectetur aliquam sed est.
+item_883,Eius ipsum est ut.
+item_884,Dolore etincidunt numquam numquam modi dolor.
+item_885,Non sed dolorem non modi sed sit.
+item_886,Consectetur quaerat porro modi labore non magnam modi.
+item_887,Adipisci magnam neque quaerat aliquam quisquam dolore.
+item_888,Dolorem est velit modi velit.
+item_889,Magnam aliquam voluptatem quiquia numquam.
+item_890,Ut non eius quisquam neque.
+item_891,Neque adipisci porro est labore consectetur dolor.
+item_892,Dolor sed ipsum voluptatem numquam numquam.
+item_893,Quaerat numquam consectetur porro.
+item_894,Sit consectetur eius voluptatem eius quisquam porro labore.
+item_895,Labore aliquam aliquam ipsum quaerat magnam.
+item_896,Ut est dolor dolor non est est labore.
+item_897,Labore dolor eius sit consectetur quisquam quaerat.
+item_898,Aliquam numquam velit dolor adipisci.
+item_899,Adipisci sit numquam consectetur tempora.
+item_900,Aliquam non ut numquam neque consectetur etincidunt.
+item_901,Dolorem neque amet velit adipisci ipsum.
+item_902,Porro dolor sed amet non quaerat sed.
+item_903,Consectetur sit eius quaerat etincidunt.
+item_904,Adipisci dolor numquam aliquam ipsum quisquam etincidunt.
+item_905,Voluptatem est dolor quaerat non voluptatem dolorem.
+item_906,Neque tempora dolorem dolore porro dolore sed dolore.
+item_907,Ipsum eius dolorem ipsum dolore amet quisquam.
+item_908,Ipsum porro ipsum ipsum amet.
+item_909,Labore quiquia sed adipisci.
+item_910,Quaerat aliquam est adipisci neque sit.
+item_911,Consectetur consectetur numquam quiquia numquam non ipsum neque.
+item_912,Est ut ut labore modi.
+item_913,Porro labore magnam sed ipsum.
+item_914,Etincidunt ipsum eius ipsum voluptatem quiquia sed tempora.
+item_915,Amet dolore quaerat etincidunt tempora labore.
+item_916,Modi porro consectetur est aliquam dolorem amet ipsum.
+item_917,Labore magnam neque dolorem quisquam dolor dolorem.
+item_918,Quisquam sed amet amet.
+item_919,Dolore porro etincidunt voluptatem modi sed.
+item_920,Est velit numquam ipsum sit est dolore.
+item_921,Ipsum dolore quisquam est velit.
+item_922,Aliquam eius sit quiquia consectetur amet ut labore.
+item_923,Numquam neque ut non ut dolor eius sed.
+item_924,Dolor est voluptatem etincidunt aliquam tempora ut.
+item_925,Ut quaerat aliquam numquam magnam.
+item_926,Adipisci est ipsum ut quaerat velit aliquam amet.
+item_927,Sit labore dolor quiquia dolor dolor est.
+item_928,Dolor amet porro dolor etincidunt etincidunt numquam.
+item_929,Modi adipisci non dolorem quisquam.
+item_930,Non consectetur consectetur neque dolor sed.
+item_931,Dolore dolore adipisci neque modi amet voluptatem.
+item_932,Consectetur dolor quaerat est labore.
+item_933,Dolore aliquam labore quisquam sed.
+item_934,Consectetur quaerat amet sed.
+item_935,Numquam numquam quiquia ut dolore non consectetur.
+item_936,Quaerat ut aliquam quiquia adipisci aliquam.
+item_937,Quisquam porro tempora neque sit ipsum quisquam.
+item_938,Aliquam dolore quaerat voluptatem ut numquam neque porro.
+item_939,Tempora quiquia dolore magnam magnam consectetur neque eius.
+item_940,Quiquia ut modi eius tempora magnam modi.
+item_941,Tempora modi velit dolor adipisci neque.
+item_942,Dolorem magnam dolorem non labore quiquia quiquia dolor.
+item_943,Quisquam ipsum ut ipsum.
+item_944,Quiquia quaerat sit quisquam quisquam.
+item_945,Non ut porro eius velit amet sed.
+item_946,Sed ut est tempora quaerat.
+item_947,Quisquam porro sit ut dolore dolorem dolore.
+item_948,Magnam dolor non velit quiquia quiquia.
+item_949,Est magnam consectetur quiquia velit magnam magnam.
+item_950,Sed quisquam adipisci quaerat non sit velit.
+item_951,Tempora neque modi non.
+item_952,Ut dolore etincidunt adipisci labore.
+item_953,Sit ut ipsum tempora quiquia voluptatem non non.
+item_954,Porro eius ut eius.
+item_955,Ut non labore est velit neque labore consectetur.
+item_956,Ipsum etincidunt ipsum sed.
+item_957,Quaerat numquam consectetur velit.
+item_958,Sit numquam labore tempora.
+item_959,Velit adipisci ipsum etincidunt adipisci dolore.
+item_960,Ut labore numquam dolore numquam adipisci tempora.
+item_961,Non dolore quiquia porro ipsum.
+item_962,Quiquia quisquam ipsum voluptatem voluptatem neque.
+item_963,Etincidunt est quaerat neque amet eius.
+item_964,Sed eius est labore etincidunt neque neque.
+item_965,Sed quisquam adipisci voluptatem tempora dolor tempora.
+item_966,Labore numquam porro ut ipsum eius.
+item_967,Quiquia est eius ut est ut dolorem.
+item_968,Labore labore ipsum ut sed.
+item_969,Dolor dolor voluptatem quiquia dolorem modi aliquam.
+item_970,Quaerat velit velit neque quaerat.
+item_971,Ipsum neque ut magnam.
+item_972,Ipsum quiquia ut neque.
+item_973,Aliquam quisquam non quaerat dolor modi sit.
+item_974,Numquam dolorem quaerat dolorem tempora voluptatem neque sed.
+item_975,Porro amet voluptatem modi.
+item_976,Aliquam dolor aliquam ipsum non dolorem.
+item_977,Modi eius non non.
+item_978,Amet voluptatem voluptatem dolor consectetur.
+item_979,Eius magnam quisquam quaerat numquam sit.
+item_980,Est tempora neque tempora velit voluptatem.
+item_981,Velit eius consectetur non.
+item_982,Etincidunt non etincidunt sit etincidunt.
+item_983,Sed sed sed est.
+item_984,Sed sed ut neque etincidunt numquam.
+item_985,Quiquia amet velit sit sit modi dolore dolorem.
+item_986,Adipisci sed dolor magnam quiquia.
+item_987,Tempora eius aliquam dolor sit porro aliquam.
+item_988,Magnam consectetur sed quiquia.
+item_989,Aliquam labore sit sed dolor ipsum neque quisquam.
+item_990,Consectetur amet neque consectetur amet.
+item_991,Amet amet neque dolore labore adipisci.
+item_992,Adipisci porro numquam est neque neque velit.
+item_993,Dolor velit amet aliquam velit quaerat.
+item_994,Sed numquam ut magnam numquam voluptatem adipisci est.
+item_995,Non modi dolore porro ut quaerat adipisci.
+item_996,Voluptatem modi quiquia ipsum etincidunt dolor.
+item_997,Non velit dolore adipisci dolore est.
+item_998,Consectetur consectetur porro sed aliquam amet tempora.
+item_999,Porro magnam magnam modi modi dolor.
+item_1000,Consectetur labore porro dolor numquam adipisci ut modi.
diff --git a/entrypoint.sh b/entrypoint.sh
new file mode 100644
index 0000000000000000000000000000000000000000..140f2f12b5e8e1f5a7c8a10f92572611411743f0
--- /dev/null
+++ b/entrypoint.sh
@@ -0,0 +1,30 @@
+#!/bin/bash
+set -e
+
+# Configuration
+CONFIG_FILE="${POTATO_CONFIG:-config.yaml}"
+PORT="${PORT:-7860}"
+WORKERS="${GUNICORN_WORKERS:-2}"
+THREADS="${GUNICORN_THREADS:-4}"
+TIMEOUT="${GUNICORN_TIMEOUT:-120}"
+
+echo "Starting Potato Demo Space..."
+echo " Config: ${CONFIG_FILE}"
+echo " Port: ${PORT}"
+echo " Workers: ${WORKERS}"
+
+# Validate config exists
+if [ ! -f "${CONFIG_FILE}" ]; then
+ echo "ERROR: Config file not found: ${CONFIG_FILE}"
+ exit 1
+fi
+
+# Start with gunicorn using the factory pattern
+exec gunicorn \
+ --bind "0.0.0.0:${PORT}" \
+ --workers "${WORKERS}" \
+ --threads "${THREADS}" \
+ --timeout "${TIMEOUT}" \
+ --access-logfile - \
+ --error-logfile - \
+ "potato.flask_server:create_app('${CONFIG_FILE}')"
diff --git a/layouts/task_layout.html b/layouts/task_layout.html
new file mode 100644
index 0000000000000000000000000000000000000000..1fa2e3bcffa1123144c804ca7a523ba44e5eb37d
--- /dev/null
+++ b/layouts/task_layout.html
@@ -0,0 +1,48 @@
+
+
+
+
+
+
+
diff --git a/potato/__init__.py b/potato/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a8e89cfdaffad4d00849cff05a16624488cd83b1
--- /dev/null
+++ b/potato/__init__.py
@@ -0,0 +1,45 @@
+"""
+Potato Annotation Platform
+
+A flexible, web-based platform for text annotation tasks.
+
+This package provides a comprehensive annotation system with the following features:
+- Multi-phase annotation workflows (consent, instructions, training, annotation, post-study)
+- Support for various annotation types (labels, spans, text, likert scales, best-worst scaling)
+- User authentication and session management
+- Active learning capabilities
+- Admin dashboard for monitoring progress
+- Configurable assignment strategies
+- Multi-language and multi-task support
+
+Main Components:
+- flask_server: Core Flask application and server logic
+- routes: HTTP route handlers and request processing
+- user_state_management: User progress tracking and state persistence
+- item_state_management: Data item management and assignment
+- authentificaton: User authentication backends
+- admin: Admin dashboard functionality
+- activelearning: Active learning algorithms and model training
+
+Usage:
+ from potato.flask_server import create_app
+ app = create_app()
+ app.run()
+"""
+
+from .flask_server import create_app
+
+__version__ = "2.6.0"
+__author__ = "Potato Annotation Platform Team"
+__description__ = "A flexible, web-based platform for text annotation tasks"
+
+
+def __getattr__(name):
+ """Lazy imports for optional heavy dependencies."""
+ if name == "load_as_dataset":
+ from .datasets_integration import load_as_dataset
+ return load_as_dataset
+ if name == "load_annotations":
+ from .datasets_integration import load_annotations
+ return load_annotations
+ raise AttributeError(f"module 'potato' has no attribute {name!r}")
\ No newline at end of file
diff --git a/potato/__main__.py b/potato/__main__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a023e5eb55074b783311863e90bd6b9ae226598c
--- /dev/null
+++ b/potato/__main__.py
@@ -0,0 +1,3 @@
+from potato.flask_server import main
+
+main()
diff --git a/potato/active_learning_manager.py b/potato/active_learning_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..023024201b664b17f7bb3e603bc3d7ed00a25fc8
--- /dev/null
+++ b/potato/active_learning_manager.py
@@ -0,0 +1,1623 @@
+"""
+Enhanced Active Learning Manager with Database Persistence
+
+This module provides a comprehensive active learning system with optional
+database persistence, model saving, LLM integration, and multiple query
+strategies including uncertainty sampling, diversity sampling, BADGE, BALD,
+and hybrid combinations.
+
+References:
+ [1] Ash et al. (2020) "Deep Batch Active Learning by Diverse, Uncertain
+ Gradient Lower Bounds" (BADGE). ICLR 2020.
+ [2] Houlsby et al. (2011) "Bayesian Active Learning for Classification
+ and Preference Learning" (BALD).
+ [3] Bayer et al. (2024) "ActiveLLM: Large Language Model-Based Active
+ Learning for Textual Few-Shot Scenarios". TACL.
+ [4] Yuan et al. (2024) "Hide and Seek in Noise Labels: Noise-Robust
+ Collaborative Active Learning" (NoiseAL). ACL 2024.
+ [5] Mavromatis et al. (2024) "CoverICL: Selective Annotation for
+ In-Context Learning via Active Graph Coverage". EMNLP 2024.
+"""
+
+import threading
+import logging
+import time
+import os
+import pickle
+import json
+from typing import Dict, List, Optional, Tuple, Any, Union
+from collections import defaultdict, Counter
+import dataclasses
+from dataclasses import dataclass, field, asdict
+from enum import Enum
+import random
+import queue
+from datetime import datetime
+from abc import ABC, abstractmethod
+
+from sklearn.pipeline import Pipeline
+from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
+from sklearn.linear_model import LogisticRegression
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.svm import SVC
+from sklearn.metrics import accuracy_score, classification_report
+import numpy as np
+
+from potato.item_state_management import ItemStateManager, get_item_state_manager
+from potato.user_state_management import get_user_state_manager
+
+
+logger = logging.getLogger(__name__)
+
+
+class ResolutionStrategy(Enum):
+ """Strategies for resolving multiple annotations per instance."""
+ MAJORITY_VOTE = "majority_vote"
+ RANDOM = "random"
+ CONSENSUS = "consensus"
+ WEIGHTED_AVERAGE = "weighted_average"
+
+
+# ---------------------------------------------------------------------------
+# SentenceTransformerVectorizer
+# ---------------------------------------------------------------------------
+
+class SentenceTransformerVectorizer:
+ """sklearn-compatible wrapper for sentence-transformers.
+
+ Uses dense embeddings from pre-trained transformer models instead of
+ bag-of-words features. Produces 384-dim vectors (for default model)
+ that capture semantic meaning, enabling better classification with
+ fewer training examples.
+
+ The ``sentence-transformers`` package is an **optional** dependency and
+ is only imported when this vectorizer is actually used.
+ """
+
+ def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
+ self.model_name = model_name
+ self._model = None
+
+ def fit(self, X, y=None):
+ from sentence_transformers import SentenceTransformer
+ self._model = SentenceTransformer(self.model_name)
+ return self
+
+ def transform(self, X):
+ if self._model is None:
+ raise RuntimeError("SentenceTransformerVectorizer has not been fitted yet")
+ return self._model.encode(list(X), show_progress_bar=False)
+
+ def fit_transform(self, X, y=None):
+ self.fit(X, y)
+ return self.transform(X)
+
+
+# ---------------------------------------------------------------------------
+# Query Strategies
+# ---------------------------------------------------------------------------
+
+class QueryStrategy(ABC):
+ """Base class for active learning query strategies."""
+
+ @abstractmethod
+ def rank(self, texts: List[str], model, vectorizer,
+ annotated_texts: Optional[List[str]] = None) -> List[Tuple[int, float]]:
+ """Return list of (index, score) sorted by selection priority (highest first)."""
+
+
+class UncertaintySampling(QueryStrategy):
+ """Select instances where classifier is least confident.
+
+ Selects x* = argmax_x (1 - max_y P(y|x)), i.e., instances where the
+ model's best guess has lowest confidence.
+ """
+
+ def rank(self, texts, model, vectorizer, annotated_texts=None):
+ try:
+ features = vectorizer.transform(texts)
+ probas = model.predict_proba(features)
+ # Score = 1 - max_prob (higher = more uncertain = higher priority)
+ scores = 1.0 - np.max(probas, axis=1)
+ ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
+ return ranked
+ except Exception as e:
+ logger.warning(f"UncertaintySampling failed: {e}")
+ return [(i, 0.5) for i in range(len(texts))]
+
+
+class DiversitySampling(QueryStrategy):
+ """Select instances that maximize feature-space coverage.
+
+ Uses cosine distance from already-annotated instances in the vectorized
+ feature space. Ensures the training set covers the full data distribution
+ rather than over-sampling one region.
+ """
+
+ def rank(self, texts, model, vectorizer, annotated_texts=None):
+ from sklearn.metrics.pairwise import cosine_distances
+
+ try:
+ features = vectorizer.transform(texts)
+ if hasattr(features, 'toarray'):
+ features = features.toarray()
+
+ if annotated_texts:
+ annotated_features = vectorizer.transform(annotated_texts)
+ if hasattr(annotated_features, 'toarray'):
+ annotated_features = annotated_features.toarray()
+ # Score = min cosine distance to any annotated instance
+ distances = cosine_distances(features, annotated_features)
+ scores = np.min(distances, axis=1)
+ else:
+ # No annotated texts yet: use distance from centroid
+ centroid = np.mean(features, axis=0, keepdims=True)
+ scores = cosine_distances(features, centroid).ravel()
+
+ ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
+ return ranked
+ except Exception as e:
+ logger.warning(f"DiversitySampling failed: {e}")
+ return [(i, 0.5) for i in range(len(texts))]
+
+
+class BadgeStrategy(QueryStrategy):
+ """BADGE approximation: uncertainty-weighted diversity.
+
+ Inspired by Ash et al. (2020) [Ref 1]. Full BADGE uses gradient embeddings
+ from neural networks. Our approximation:
+ 1. Weight feature vectors by (1 - max_prob) as uncertainty proxy
+ 2. Run k-means++ initialization on weighted vectors to select
+ diverse-uncertain instances.
+ """
+
+ def rank(self, texts, model, vectorizer, annotated_texts=None):
+ try:
+ features = vectorizer.transform(texts)
+ if hasattr(features, 'toarray'):
+ features = features.toarray()
+
+ probas = model.predict_proba(features)
+ uncertainty = 1.0 - np.max(probas, axis=1)
+
+ # Weight features by uncertainty
+ weighted = features * uncertainty[:, np.newaxis]
+
+ # Use k-means++ initialization to select diverse-uncertain points
+ from sklearn.cluster import kmeans_plusplus
+ n_clusters = min(len(texts), max(1, len(texts) // 2))
+ _, indices = kmeans_plusplus(weighted, n_clusters=n_clusters,
+ random_state=42)
+
+ # Build score: selected centroids get highest scores
+ scores = np.zeros(len(texts))
+ for rank_pos, idx in enumerate(indices):
+ scores[idx] = len(indices) - rank_pos # highest for first-selected
+
+ # For non-selected, use uncertainty as tiebreaker
+ for i in range(len(texts)):
+ if scores[i] == 0:
+ scores[i] = uncertainty[i] * 0.01
+
+ ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
+ return ranked
+ except Exception as e:
+ logger.warning(f"BadgeStrategy failed, falling back to uncertainty: {e}")
+ return UncertaintySampling().rank(texts, model, vectorizer, annotated_texts)
+
+
+class BaldStrategy(QueryStrategy):
+ """BALD: Bayesian Active Learning by Disagreement.
+
+ Based on Houlsby et al. (2011) [Ref 2]. Trains an ensemble of classifiers
+ with different random seeds/bootstrap samples. Selects instances with
+ highest mutual information: H[y|x] - E_theta[H[y|x,theta]], i.e.,
+ where the ensemble disagrees most.
+ """
+
+ def __init__(self, n_estimators: int = 5, bootstrap_fraction: float = 0.8):
+ self.n_estimators = n_estimators
+ self.bootstrap_fraction = bootstrap_fraction
+
+ def rank(self, texts, model, vectorizer, annotated_texts=None):
+ try:
+ features = vectorizer.transform(texts)
+ if hasattr(features, 'toarray'):
+ features = features.toarray()
+
+ probas = model.predict_proba(features)
+ # Average entropy
+ avg_proba = probas
+ entropy_avg = -np.sum(avg_proba * np.log(avg_proba + 1e-10), axis=1)
+
+ # For a single model, we approximate BALD by using dropout-like noise
+ # or by comparing with uniform. Since we store the ensemble models
+ # on the manager, we just use the single model's entropy here and
+ # the ensemble version is handled in ActiveLearningManager._train_bald_ensemble
+ scores = entropy_avg
+ ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
+ return ranked
+ except Exception as e:
+ logger.warning(f"BaldStrategy failed: {e}")
+ return [(i, 0.5) for i in range(len(texts))]
+
+ def rank_with_ensemble(self, texts, ensemble_models, vectorizer):
+ """Rank using actual ensemble disagreement (mutual information)."""
+ try:
+ features = vectorizer.transform(texts)
+ if hasattr(features, 'toarray'):
+ features = features.toarray()
+
+ all_probas = []
+ for m in ensemble_models:
+ all_probas.append(m.predict_proba(features))
+
+ all_probas = np.array(all_probas) # (n_estimators, n_samples, n_classes)
+
+ # Mean prediction across ensemble
+ mean_proba = np.mean(all_probas, axis=0) # (n_samples, n_classes)
+
+ # H[y|x] - entropy of mean prediction
+ entropy_mean = -np.sum(mean_proba * np.log(mean_proba + 1e-10), axis=1)
+
+ # E_theta[H[y|x,theta]] - mean of individual entropies
+ individual_entropies = -np.sum(all_probas * np.log(all_probas + 1e-10), axis=2)
+ mean_entropy = np.mean(individual_entropies, axis=0)
+
+ # Mutual information = H[y|x] - E[H[y|x,theta]]
+ mutual_info = entropy_mean - mean_entropy
+
+ ranked = sorted(enumerate(mutual_info), key=lambda x: x[1], reverse=True)
+ return ranked
+ except Exception as e:
+ logger.warning(f"BaldStrategy ensemble ranking failed: {e}")
+ return [(i, 0.5) for i in range(len(texts))]
+
+
+class HybridStrategy(QueryStrategy):
+ """Weighted combination of uncertainty and diversity scores.
+
+ Combines strategies with configurable weights. Default: 0.7 uncertainty +
+ 0.3 diversity.
+ """
+
+ def __init__(self, weights: Optional[Dict[str, float]] = None):
+ self.weights = weights or {"uncertainty": 0.7, "diversity": 0.3}
+
+ def rank(self, texts, model, vectorizer, annotated_texts=None):
+ try:
+ strategies = {}
+ if self.weights.get("uncertainty", 0) > 0:
+ strategies["uncertainty"] = UncertaintySampling()
+ if self.weights.get("diversity", 0) > 0:
+ strategies["diversity"] = DiversitySampling()
+
+ # Collect raw scores from each strategy
+ all_scores = {}
+ for name, strategy in strategies.items():
+ rankings = strategy.rank(texts, model, vectorizer, annotated_texts)
+ score_map = {idx: score for idx, score in rankings}
+ all_scores[name] = score_map
+
+ # Normalize each strategy's scores to [0, 1]
+ for name in all_scores:
+ vals = list(all_scores[name].values())
+ min_val, max_val = min(vals), max(vals)
+ rng = max_val - min_val if max_val > min_val else 1.0
+ all_scores[name] = {
+ idx: (s - min_val) / rng for idx, s in all_scores[name].items()
+ }
+
+ # Weighted combination
+ combined = {}
+ for i in range(len(texts)):
+ combined[i] = sum(
+ self.weights.get(name, 0) * all_scores.get(name, {}).get(i, 0)
+ for name in self.weights
+ )
+
+ ranked = sorted(combined.items(), key=lambda x: x[1], reverse=True)
+ return ranked
+ except Exception as e:
+ logger.warning(f"HybridStrategy failed: {e}")
+ return UncertaintySampling().rank(texts, model, vectorizer, annotated_texts)
+
+
+# Strategy registry
+STRATEGY_REGISTRY = {
+ "uncertainty": UncertaintySampling,
+ "diversity": DiversitySampling,
+ "badge": BadgeStrategy,
+ "bald": BaldStrategy,
+ "hybrid": HybridStrategy,
+}
+
+
+def create_query_strategy(config: 'ActiveLearningConfig') -> QueryStrategy:
+ """Create a query strategy from config."""
+ strategy_name = config.query_strategy
+ if strategy_name == "hybrid":
+ return HybridStrategy(weights=config.hybrid_weights)
+ elif strategy_name == "bald":
+ params = config.bald_params
+ return BaldStrategy(
+ n_estimators=params.get("n_estimators", 5),
+ bootstrap_fraction=params.get("bootstrap_fraction", 0.8),
+ )
+ elif strategy_name in STRATEGY_REGISTRY:
+ return STRATEGY_REGISTRY[strategy_name]()
+ else:
+ logger.warning(f"Unknown strategy '{strategy_name}', falling back to uncertainty")
+ return UncertaintySampling()
+
+
+# ---------------------------------------------------------------------------
+# ICLClassifier wrapper (Phase 5A)
+# ---------------------------------------------------------------------------
+
+class ICLClassifier:
+ """Wraps ICLLabeler as an sklearn-compatible classifier for ensemble use.
+
+ Enables combining LLM-based ICL predictions with traditional classifier
+ predictions in a hybrid ensemble for active learning scoring.
+ """
+
+ def __init__(self, icl_labeler, schema_name: str, label_names: List[str]):
+ self.icl_labeler = icl_labeler
+ self.schema_name = schema_name
+ self.label_names = label_names
+ self.classes_ = np.array(label_names)
+
+ def predict_proba(self, texts: List[str]) -> np.ndarray:
+ """Get label probabilities from LLM via ICL."""
+ n_classes = len(self.label_names)
+ probas = np.full((len(texts), n_classes), 1.0 / n_classes)
+
+ for i, text in enumerate(texts):
+ try:
+ prediction = self.icl_labeler.label_instance(
+ instance_id=f"_al_query_{i}",
+ schema_name=self.schema_name,
+ instance_text=text,
+ )
+ if prediction and prediction.predicted_label in self.label_names:
+ idx = self.label_names.index(prediction.predicted_label)
+ conf = prediction.confidence_score
+ # Distribute: conf to predicted label, (1-conf)/(n-1) to others
+ remaining = (1.0 - conf) / max(1, n_classes - 1)
+ probas[i] = remaining
+ probas[i, idx] = conf
+ except Exception:
+ pass # Keep uniform distribution
+
+ return probas
+
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+@dataclass
+class ActiveLearningConfig:
+ """Enhanced configuration for active learning."""
+ enabled: bool = False
+ classifier_name: str = "sklearn.linear_model.LogisticRegression"
+ classifier_kwargs: Dict[str, Any] = None
+ vectorizer_name: str = "sklearn.feature_extraction.text.TfidfVectorizer"
+ vectorizer_kwargs: Dict[str, Any] = None
+ min_annotations_per_instance: int = 1
+ min_instances_for_training: int = 10
+ max_instances_to_reorder: Optional[int] = None
+ resolution_strategy: ResolutionStrategy = ResolutionStrategy.MAJORITY_VOTE
+ random_sample_percent: float = 0.2
+ update_frequency: int = 5
+ schema_names: List[str] = None
+
+ # Classifier/vectorizer passthrough params (Phase 1C)
+ classifier_params: Dict[str, Any] = field(default_factory=dict)
+ vectorizer_params: Dict[str, Any] = field(default_factory=dict)
+
+ # Probability calibration (Phase 1D)
+ calibrate_probabilities: bool = True
+
+ # Query strategy (Phase 2)
+ query_strategy: str = "uncertainty"
+ hybrid_weights: Dict[str, float] = field(
+ default_factory=lambda: {"uncertainty": 0.7, "diversity": 0.3}
+ )
+ bald_params: Dict[str, Any] = field(
+ default_factory=lambda: {"n_estimators": 5, "bootstrap_fraction": 0.8}
+ )
+
+ # Cold-start (Phase 3)
+ cold_start_strategy: str = "random"
+ cold_start_batch_size: int = 20
+
+ # ICL ensemble (Phase 5)
+ use_icl_ensemble: bool = False
+ icl_ensemble_params: Dict[str, Any] = field(default_factory=lambda: {
+ "initial_icl_weight": 0.7,
+ "final_icl_weight": 0.2,
+ "transition_instances": 100,
+ })
+
+ # Annotation routing (Phase 5D)
+ annotation_routing: bool = False
+ routing_thresholds: Dict[str, float] = field(default_factory=lambda: {
+ "auto_label_min_confidence": 0.9,
+ "show_suggestion_below": 0.5,
+ })
+ verification_sample_rate: float = 0.2
+
+ # Database persistence
+ database_enabled: bool = False
+ database_config: Dict[str, Any] = None
+
+ # Model persistence
+ model_persistence_enabled: bool = False
+ model_save_directory: Optional[str] = None
+ model_retention_count: int = 2
+
+ # LLM integration
+ llm_enabled: bool = False
+ llm_config: Dict[str, Any] = None
+
+ def __post_init__(self):
+ if self.classifier_kwargs is None:
+ self.classifier_kwargs = {}
+ if self.vectorizer_kwargs is None:
+ self.vectorizer_kwargs = {}
+ if self.schema_names is None:
+ self.schema_names = []
+ if self.database_config is None:
+ self.database_config = {}
+ if self.llm_config is None:
+ self.llm_config = {}
+ # Merge classifier_params into classifier_kwargs
+ if self.classifier_params:
+ self.classifier_kwargs.update(self.classifier_params)
+ # Merge vectorizer_params into vectorizer_kwargs
+ if self.vectorizer_params:
+ self.vectorizer_kwargs.update(self.vectorizer_params)
+
+
+@dataclass
+class TrainingMetrics:
+ """Metrics for a training run."""
+ schema_name: str
+ training_time: float
+ accuracy: float
+ instance_count: int
+ timestamp: datetime
+ model_file_path: Optional[str] = None
+ confidence_distribution: Dict[str, float] = None
+ error_message: Optional[str] = None
+
+
+class ModelPersistence:
+ """Handles model saving and loading with metadata."""
+
+ def __init__(self, save_directory: str, retention_count: int = 2):
+ self.save_directory = save_directory
+ self.retention_count = retention_count
+ self.logger = logging.getLogger(__name__)
+
+ # Ensure directory exists
+ os.makedirs(save_directory, exist_ok=True)
+
+ def save_model(self, model: Pipeline, schema_name: str, instance_count: int) -> str:
+ """Save a trained model with metadata."""
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+ filename = f"{schema_name}_{instance_count}_{timestamp}.pkl"
+ filepath = os.path.join(self.save_directory, filename)
+
+ try:
+ # Save the complete model (including vectorizer)
+ with open(filepath, 'wb') as f:
+ pickle.dump(model, f)
+
+ self.logger.info(f"Saved model to {filepath}")
+
+ # Clean up old models
+ self._cleanup_old_models(schema_name)
+
+ return filepath
+ except Exception as e:
+ self.logger.error(f"Failed to save model: {e}")
+ raise
+
+ def load_model(self, filepath: str) -> Optional[Pipeline]:
+ """Load a saved model."""
+ try:
+ with open(filepath, 'rb') as f:
+ model = pickle.load(f)
+
+ # TODO: Add schema validation here in the future
+ # This is a placeholder for future schema validation enhancement
+
+ self.logger.info(f"Loaded model from {filepath}")
+ return model
+ except Exception as e:
+ self.logger.error(f"Failed to load model from {filepath}: {e}")
+ return None
+
+ def _cleanup_old_models(self, schema_name: str):
+ """Clean up old models based on retention policy."""
+ try:
+ # Find all model files for this schema
+ model_files = []
+
+ for filename in os.listdir(self.save_directory):
+ if filename.startswith(f"{schema_name}_") and filename.endswith(".pkl"):
+ filepath = os.path.join(self.save_directory, filename)
+ model_files.append((filepath, os.path.getmtime(filepath)))
+
+ # Sort by modification time (newest first)
+ model_files.sort(key=lambda x: x[1], reverse=True)
+
+ # Remove old models beyond retention count
+ for filepath, _ in model_files[self.retention_count:]:
+ try:
+ os.remove(filepath)
+ self.logger.info(f"Removed old model: {filepath}")
+ except Exception as e:
+ self.logger.warning(f"Failed to remove old model {filepath}: {e}")
+
+ except Exception as e:
+ self.logger.error(f"Error during model cleanup: {e}")
+
+
+class DatabaseStateManager:
+ """Manages database persistence for active learning state."""
+
+ def __init__(self, config: Dict[str, Any]):
+ self.config = config
+ self.logger = logging.getLogger(__name__)
+ self.connection = None
+ self._initialize_database()
+
+ def _initialize_database(self):
+ """Initialize database connection and create tables."""
+ try:
+ # Use the same database system as main Potato application
+ if self.config.get('type') == 'mysql':
+ self._init_mysql_connection()
+ else:
+ self._init_file_based_connection()
+
+ self._create_tables()
+ self.logger.info("Active learning database initialized successfully")
+ except Exception as e:
+ self.logger.error(f"Failed to initialize database: {e}")
+ raise
+
+ def _init_mysql_connection(self):
+ """Initialize MySQL connection."""
+ # TODO: Implement MySQL connection
+ pass
+
+ def _init_file_based_connection(self):
+ """Initialize file-based database connection."""
+ # TODO: Implement file-based database
+ pass
+
+ def _create_tables(self):
+ """Create database tables for active learning."""
+ # TODO: Implement table creation
+ pass
+
+ def save_training_metrics(self, metrics: TrainingMetrics):
+ """Save training metrics to database."""
+ # TODO: Implement metrics saving
+ pass
+
+ def get_training_history(self, schema_name: Optional[str] = None) -> List[TrainingMetrics]:
+ """Get training history from database."""
+ # TODO: Implement history retrieval
+ return []
+
+ def save_schema_cycling_state(self, current_schema: str, schema_order: List[str]):
+ """Save current schema cycling state."""
+ # TODO: Implement state saving
+ pass
+
+ def get_schema_cycling_state(self) -> Tuple[str, List[str]]:
+ """Get current schema cycling state."""
+ # TODO: Implement state retrieval
+ return "", []
+
+
+class SchemaCycler:
+ """Manages cycling through multiple annotation schemes."""
+
+ def __init__(self, schema_names: List[str], database_manager: Optional[DatabaseStateManager] = None):
+ self.schema_names = self._validate_schemas(schema_names)
+ self.database_manager = database_manager
+ self.current_index = 0
+ self.logger = logging.getLogger(__name__)
+ self._lock = threading.Lock()
+
+ # Load state from database if available
+ if self.database_manager:
+ self._load_state()
+
+ def _validate_schemas(self, schema_names: List[str]) -> List[str]:
+ """Validate and filter schema names."""
+ valid_schemas = []
+
+ for schema in schema_names:
+ # Exclude text and span annotation schemes
+ if schema in ['text', 'span']:
+ raise ValueError(f"Text and span annotation schemes are not supported for active learning: {schema}")
+ valid_schemas.append(schema)
+
+ return valid_schemas
+
+ def _load_state(self):
+ """Load cycling state from database."""
+ try:
+ current_schema, schema_order = self.database_manager.get_schema_cycling_state()
+ with self._lock:
+ if current_schema in self.schema_names:
+ self.current_index = self.schema_names.index(current_schema)
+ except Exception as e:
+ self.logger.warning(f"Failed to load schema cycling state: {e}")
+
+ def get_current_schema(self) -> Optional[str]:
+ """Get the current schema for training."""
+ if not self.schema_names:
+ return None
+ with self._lock:
+ return self.schema_names[self.current_index]
+
+ def advance_schema(self):
+ """Advance to the next schema in the cycle."""
+ if not self.schema_names:
+ return
+
+ with self._lock:
+ self.current_index = (self.current_index + 1) % len(self.schema_names)
+ current_schema = self.schema_names[self.current_index]
+
+ # Save state to database if available
+ if self.database_manager:
+ try:
+ self.database_manager.save_schema_cycling_state(
+ current_schema,
+ self.schema_names
+ )
+ except Exception as e:
+ self.logger.warning(f"Failed to save schema cycling state: {e}")
+
+ def get_schema_order(self) -> List[str]:
+ """Get the current schema cycling order."""
+ return self.schema_names.copy()
+
+
+class ActiveLearningManager:
+ """
+ Manages active learning operations including classifier training and instance reordering.
+
+ This class provides thread-safe operations for:
+ - Training classifiers on annotated data
+ - Predicting confidence scores for unlabeled instances
+ - Reordering instances based on configurable query strategies
+ - Cold-start LLM-based instance selection
+ - ICL/classifier ensemble for improved ranking
+ - Noise-aware annotation routing
+ - Managing training state and progress
+ - Database persistence and model saving
+ """
+
+ def __init__(self, config: ActiveLearningConfig):
+ self.config = config
+ self.logger = logging.getLogger(__name__)
+
+ # Thread safety
+ self._lock = threading.RLock()
+ self._training_queue = queue.Queue()
+ self._training_thread = None
+ self._stop_training = threading.Event()
+
+ # State tracking
+ self._last_training_time = 0
+ self._training_count = 0
+ self._models = {} # schema_name -> trained_model
+ self._vectorizers = {} # schema_name -> fitted vectorizer
+ self._bald_ensembles = {} # schema_name -> list of classifiers
+ self._last_annotation_count = 0
+ self._training_metrics = [] # List of TrainingMetrics
+ self._annotated_texts = {} # schema_name -> list of annotated texts
+
+ # Query strategy
+ self._query_strategy = create_query_strategy(config)
+
+ # Database and persistence
+ self.database_manager = None
+ self.model_persistence = None
+ self.schema_cycler = None
+
+ # Initialize components
+ self._initialize_components()
+
+ # Start training thread if enabled
+ if self.config.enabled:
+ self._start_training_thread()
+
+ def _initialize_components(self):
+ """Initialize database, model persistence, and schema cycler."""
+ # Initialize database manager if enabled
+ if self.config.database_enabled:
+ try:
+ self.database_manager = DatabaseStateManager(self.config.database_config)
+ except Exception as e:
+ self.logger.error(f"Failed to initialize database manager: {e}")
+ # Continue without database persistence
+
+ # Initialize model persistence if enabled
+ if self.config.model_persistence_enabled and self.config.model_save_directory:
+ try:
+ self.model_persistence = ModelPersistence(
+ self.config.model_save_directory,
+ self.config.model_retention_count
+ )
+ except Exception as e:
+ self.logger.error(f"Failed to initialize model persistence: {e}")
+ # Continue without model persistence
+
+ # Initialize schema cycler
+ try:
+ self.schema_cycler = SchemaCycler(self.config.schema_names, self.database_manager)
+ except Exception as e:
+ self.logger.error(f"Failed to initialize schema cycler: {e}")
+ raise # Schema cycler is critical
+
+ def _start_training_thread(self):
+ """Start the background training thread."""
+ if self._training_thread is None or not self._training_thread.is_alive():
+ self._training_thread = threading.Thread(target=self._training_worker, daemon=True)
+ self._training_thread.start()
+ self.logger.info("Active learning training thread started")
+
+ def _training_worker(self):
+ """Background worker for training classifiers."""
+ while not self._stop_training.is_set():
+ try:
+ # Wait for training request
+ training_request = self._training_queue.get(timeout=1.0)
+ if training_request is None: # Shutdown signal
+ break
+
+ self._perform_training()
+ self._training_queue.task_done()
+
+ except queue.Empty:
+ continue
+ except Exception as e:
+ self.logger.error(f"Error in training worker: {e}")
+
+ def _perform_training(self):
+ """Perform the actual classifier training."""
+ with self._lock:
+ try:
+ self.logger.info("Starting active learning classifier training")
+ start_time = time.time()
+
+ # Get current schema for training
+ current_schema = self.schema_cycler.get_current_schema()
+ if not current_schema:
+ self.logger.warning("No schema available for training")
+ return
+
+ # Get current annotation state
+ item_manager = get_item_state_manager()
+ user_manager = get_user_state_manager()
+
+ # Collect training data
+ training_data = self._collect_training_data(item_manager, user_manager, current_schema)
+
+ if not training_data:
+ self.logger.warning(f"No training data available for schema {current_schema}")
+ # If in cold-start phase, try LLM-based reordering
+ if self.config.cold_start_strategy == "llm" and self.config.llm_enabled:
+ self._cold_start_reorder(item_manager)
+ return
+
+ # Train classifier
+ model, metrics = self._train_classifier(training_data, current_schema)
+
+ if model:
+ self._models[current_schema] = model
+ self._annotated_texts[current_schema] = training_data["texts"]
+
+ # Save model if persistence is enabled
+ if self.model_persistence:
+ try:
+ model_path = self.model_persistence.save_model(
+ model, current_schema, len(training_data["texts"])
+ )
+ metrics.model_file_path = model_path
+ except Exception as e:
+ self.logger.error(f"Failed to save model: {e}")
+
+ # Save metrics to database if available
+ if self.database_manager:
+ try:
+ self.database_manager.save_training_metrics(metrics)
+ except Exception as e:
+ self.logger.error(f"Failed to save metrics: {e}")
+
+ # Reorder instances
+ self._reorder_instances(item_manager, current_schema)
+
+ # Advance to next schema
+ self.schema_cycler.advance_schema()
+
+ self._training_count += 1
+ self._last_training_time = time.time()
+
+ training_duration = time.time() - start_time
+ self.logger.info(f"Active learning training completed for schema {current_schema} "
+ f"(run #{self._training_count}, duration: {training_duration:.2f}s)")
+ else:
+ self.logger.warning(f"Failed to train model for schema {current_schema}")
+ # Try cold-start if not enough data
+ if (self.config.cold_start_strategy == "llm"
+ and self.config.llm_enabled
+ and len(training_data.get("texts", [])) < self.config.min_instances_for_training):
+ self._cold_start_reorder(item_manager)
+
+ except Exception as e:
+ self.logger.error(f"Error during training: {e}")
+ # Continue without failing the entire system
+
+ def _collect_training_data(self, item_manager: ItemStateManager, user_manager, schema_name: str) -> Dict:
+ """Collect training data for a specific schema."""
+ training_data = {"texts": [], "labels": [], "instance_ids": []}
+
+ # Get all user states
+ user_states = user_manager.get_all_users()
+ self.logger.debug(f"Found {len(user_states)} user states")
+
+ # Collect annotations per instance
+ instance_annotations = defaultdict(list)
+
+ for user_state in user_states:
+ user_annotations = user_state.get_all_annotations()
+ self.logger.debug(f"User {user_state.user_id} has {len(user_annotations)} annotations")
+ for instance_id, annotations in user_annotations.items():
+ # Check if the schema exists in the labels section
+ if 'labels' in annotations:
+ labels_dict = annotations['labels']
+ # Handle Label objects as keys
+ for label_obj, value in labels_dict.items():
+ if hasattr(label_obj, 'get_schema') and label_obj.get_schema() == schema_name:
+ instance_annotations[instance_id].append({
+ "label": label_obj.get_name(),
+ "value": value,
+ "user": user_state.user_id
+ })
+
+ self.logger.debug(f"Collected annotations for {len(instance_annotations)} instances")
+
+ # Filter instances with sufficient annotations
+ for instance_id, annotations in instance_annotations.items():
+ if len(annotations) >= self.config.min_annotations_per_instance:
+ # Resolve multiple annotations
+ resolved_label = self._resolve_annotations(annotations)
+ if resolved_label:
+ item = item_manager.get_item(instance_id)
+ if item:
+ text = item.get_text()
+ training_data["texts"].append(text)
+ training_data["labels"].append(resolved_label)
+ training_data["instance_ids"].append(instance_id)
+
+ self.logger.debug(f"Training data collected: {len(training_data['texts'])} texts, {len(training_data['labels'])} labels")
+ return training_data
+
+ def _resolve_annotations(self, annotations: List[Dict]) -> Optional[str]:
+ """Resolve multiple annotations using the configured strategy."""
+ if not annotations:
+ return None
+
+ if self.config.resolution_strategy == ResolutionStrategy.MAJORITY_VOTE:
+ return self._majority_vote(annotations)
+ elif self.config.resolution_strategy == ResolutionStrategy.RANDOM:
+ return self._random_selection(annotations)
+ elif self.config.resolution_strategy == ResolutionStrategy.CONSENSUS:
+ return self._consensus_resolution(annotations)
+ else:
+ return self._majority_vote(annotations) # Default fallback
+
+ def _majority_vote(self, annotations: List[Dict]) -> str:
+ """Resolve annotations using majority vote with random tie-breaking."""
+ label_counts = Counter(ann["label"] for ann in annotations)
+ max_count = max(label_counts.values())
+ # Find all labels with the maximum count (handles ties)
+ tied_labels = [label for label, count in label_counts.items() if count == max_count]
+ # Break ties randomly
+ return random.choice(tied_labels)
+
+ def _random_selection(self, annotations: List[Dict]) -> str:
+ """Resolve annotations by random selection."""
+ return random.choice(annotations)["label"]
+
+ def _consensus_resolution(self, annotations: List[Dict]) -> Optional[str]:
+ """Resolve annotations by consensus (all must agree)."""
+ labels = [ann["label"] for ann in annotations]
+ if len(set(labels)) == 1:
+ return labels[0]
+ return None
+
+ def _train_classifier(self, training_data: Dict, schema_name: str) -> Tuple[Optional[Pipeline], TrainingMetrics]:
+ """Train a classifier for a specific schema."""
+ start_time = time.time()
+
+ if len(training_data["texts"]) < self.config.min_instances_for_training:
+ error_msg = f"Insufficient training data for schema {schema_name}: {len(training_data['texts'])} < {self.config.min_instances_for_training}"
+ self.logger.warning(error_msg)
+ return None, TrainingMetrics(
+ schema_name=schema_name,
+ training_time=time.time() - start_time,
+ accuracy=0.0,
+ instance_count=len(training_data["texts"]),
+ timestamp=datetime.now(),
+ error_message=error_msg
+ )
+
+ # Check for sufficient label diversity
+ unique_labels = set(training_data["labels"])
+ if len(unique_labels) < 2:
+ error_msg = f"Insufficient label diversity for schema {schema_name}: {len(unique_labels)} unique labels"
+ self.logger.warning(error_msg)
+ return None, TrainingMetrics(
+ schema_name=schema_name,
+ training_time=time.time() - start_time,
+ accuracy=0.0,
+ instance_count=len(training_data["texts"]),
+ timestamp=datetime.now(),
+ error_message=error_msg
+ )
+
+ try:
+ # Create and train classifier
+ classifier = self._create_classifier()
+ vectorizer = self._create_vectorizer()
+
+ pipeline = Pipeline([
+ ("vectorizer", vectorizer),
+ ("classifier", classifier)
+ ])
+
+ pipeline.fit(training_data["texts"], training_data["labels"])
+
+ # Apply probability calibration if enabled
+ if self.config.calibrate_probabilities and hasattr(classifier, 'predict_proba'):
+ num_samples = len(training_data["texts"])
+ if num_samples >= 5:
+ try:
+ from sklearn.calibration import CalibratedClassifierCV
+ cv_folds = min(3, num_samples // 2)
+ if cv_folds >= 2:
+ calibrated = CalibratedClassifierCV(
+ pipeline, cv=cv_folds, method='isotonic'
+ )
+ calibrated.fit(training_data["texts"], training_data["labels"])
+ pipeline = calibrated
+ self.logger.debug(f"Applied probability calibration with {cv_folds}-fold CV")
+ except Exception as e:
+ self.logger.warning(f"Calibration failed, using uncalibrated model: {e}")
+
+ # Store vectorizer separately for strategy use
+ self._vectorizers[schema_name] = pipeline.named_steps.get("vectorizer", vectorizer) if hasattr(pipeline, 'named_steps') else vectorizer
+
+ # Train BALD ensemble if needed
+ if self.config.query_strategy == "bald":
+ self._train_bald_ensemble(training_data, schema_name)
+
+ # Calculate accuracy
+ predictions = pipeline.predict(training_data["texts"])
+ accuracy = accuracy_score(training_data["labels"], predictions)
+
+ # Calculate confidence distribution
+ confidence_distribution = self._calculate_confidence_distribution(pipeline, training_data["texts"])
+
+ training_time = time.time() - start_time
+
+ metrics = TrainingMetrics(
+ schema_name=schema_name,
+ training_time=training_time,
+ accuracy=accuracy,
+ instance_count=len(training_data["texts"]),
+ timestamp=datetime.now(),
+ confidence_distribution=confidence_distribution
+ )
+
+ self.logger.info(f"Trained classifier for schema {schema_name} with {len(training_data['texts'])} instances, "
+ f"accuracy: {accuracy:.3f}, time: {training_time:.2f}s")
+
+ return pipeline, metrics
+
+ except Exception as e:
+ error_msg = f"Error training classifier for schema {schema_name}: {e}"
+ self.logger.error(error_msg)
+ return None, TrainingMetrics(
+ schema_name=schema_name,
+ training_time=time.time() - start_time,
+ accuracy=0.0,
+ instance_count=len(training_data["texts"]),
+ timestamp=datetime.now(),
+ error_message=error_msg
+ )
+
+ def _train_bald_ensemble(self, training_data: Dict, schema_name: str):
+ """Train an ensemble of classifiers for BALD strategy."""
+ params = self.config.bald_params
+ n_estimators = params.get("n_estimators", 5)
+ bootstrap_fraction = params.get("bootstrap_fraction", 0.8)
+
+ texts = training_data["texts"]
+ labels = training_data["labels"]
+ n_samples = len(texts)
+ bootstrap_size = max(2, int(n_samples * bootstrap_fraction))
+
+ ensemble = []
+ for i in range(n_estimators):
+ indices = np.random.choice(n_samples, size=bootstrap_size, replace=True)
+ boot_texts = [texts[j] for j in indices]
+ boot_labels = [labels[j] for j in indices]
+
+ # Need at least 2 classes
+ if len(set(boot_labels)) < 2:
+ continue
+
+ clf = self._create_classifier()
+ vec = self._create_vectorizer()
+ pipe = Pipeline([("vectorizer", vec), ("classifier", clf)])
+ pipe.fit(boot_texts, boot_labels)
+ ensemble.append(pipe)
+
+ if ensemble:
+ self._bald_ensembles[schema_name] = ensemble
+ self.logger.info(f"Trained BALD ensemble with {len(ensemble)} models for {schema_name}")
+
+ def _calculate_confidence_distribution(self, pipeline, texts: List[str]) -> Dict[str, float]:
+ """Calculate confidence score distribution."""
+ try:
+ probas = pipeline.predict_proba(texts)
+ max_confidences = np.max(probas, axis=1)
+
+ # Create histogram bins
+ bins = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
+ hist, _ = np.histogram(max_confidences, bins=bins)
+
+ # Convert to percentages
+ total = len(max_confidences)
+ distribution = {}
+ for i, count in enumerate(hist):
+ bin_label = f"{bins[i]:.1f}-{bins[i+1]:.1f}"
+ distribution[bin_label] = (count / total) * 100 if total > 0 else 0
+
+ return distribution
+ except Exception as e:
+ self.logger.warning(f"Failed to calculate confidence distribution: {e}")
+ return {}
+
+ def _create_classifier(self):
+ """Create classifier instance based on configuration."""
+ kwargs = dict(self.config.classifier_kwargs)
+
+ if self.config.classifier_name == "sklearn.linear_model.LogisticRegression":
+ return LogisticRegression(**kwargs)
+ elif self.config.classifier_name == "sklearn.ensemble.RandomForestClassifier":
+ return RandomForestClassifier(**kwargs)
+ elif self.config.classifier_name == "sklearn.svm.SVC":
+ kwargs.setdefault("probability", True)
+ return SVC(**kwargs)
+ else:
+ # Try to import dynamically
+ try:
+ module_name, class_name = self.config.classifier_name.rsplit('.', 1)
+ module = __import__(module_name, fromlist=[class_name])
+ classifier_class = getattr(module, class_name)
+ return classifier_class(**kwargs)
+ except Exception as e:
+ self.logger.error(f"Failed to create classifier {self.config.classifier_name}: {e}")
+ return LogisticRegression() # Fallback
+
+ def _create_vectorizer(self):
+ """Create vectorizer instance based on configuration."""
+ kwargs = dict(self.config.vectorizer_kwargs)
+
+ if self.config.vectorizer_name == "sklearn.feature_extraction.text.CountVectorizer":
+ return CountVectorizer(**kwargs)
+ elif self.config.vectorizer_name == "sklearn.feature_extraction.text.TfidfVectorizer":
+ return TfidfVectorizer(**kwargs)
+ elif self.config.vectorizer_name == "sentence-transformers":
+ model_name = kwargs.pop("model_name", "all-MiniLM-L6-v2")
+ return SentenceTransformerVectorizer(model_name=model_name)
+ else:
+ # Try to import dynamically
+ try:
+ module_name, class_name = self.config.vectorizer_name.rsplit('.', 1)
+ module = __import__(module_name, fromlist=[class_name])
+ vectorizer_class = getattr(module, class_name)
+ return vectorizer_class(**kwargs)
+ except Exception as e:
+ self.logger.error(f"Failed to create vectorizer {self.config.vectorizer_name}: {e}")
+ return TfidfVectorizer() # Fallback
+
+ def _reorder_instances(self, item_manager: ItemStateManager, schema_name: str):
+ """Reorder instances based on the configured query strategy."""
+ if schema_name not in self._models:
+ self.logger.warning(f"No trained model available for schema {schema_name}")
+ return
+
+ # Get unlabeled instances
+ unlabeled_instances = []
+ unlabeled_texts = []
+ for instance_id in item_manager.get_instance_ids():
+ if not item_manager.get_annotators_for_item(instance_id):
+ item = item_manager.get_item(instance_id)
+ if item:
+ unlabeled_instances.append(instance_id)
+ unlabeled_texts.append(item.get_text())
+
+ if not unlabeled_texts:
+ self.logger.info("No unlabeled instances to reorder")
+ return
+
+ # Limit number of instances to process
+ if self.config.max_instances_to_reorder:
+ limit = self.config.max_instances_to_reorder
+ unlabeled_instances = unlabeled_instances[:limit]
+ unlabeled_texts = unlabeled_texts[:limit]
+
+ model = self._models[schema_name]
+ annotated = self._annotated_texts.get(schema_name, [])
+
+ # Get rankings from strategy
+ if (self.config.query_strategy == "bald"
+ and schema_name in self._bald_ensembles
+ and isinstance(self._query_strategy, BaldStrategy)):
+ vectorizer = self._vectorizers.get(schema_name)
+ if vectorizer:
+ rankings = self._query_strategy.rank_with_ensemble(
+ unlabeled_texts, self._bald_ensembles[schema_name], vectorizer
+ )
+ else:
+ rankings = self._query_strategy.rank(unlabeled_texts, model, model, annotated)
+ else:
+ # Extract vectorizer and classifier from pipeline for strategy use
+ vectorizer = self._vectorizers.get(schema_name)
+ classifier = model
+ if vectorizer:
+ rankings = self._query_strategy.rank(
+ unlabeled_texts, classifier, vectorizer, annotated
+ )
+ else:
+ # Fallback: use confidence scores directly
+ instance_scores = self._calculate_confidence_scores(
+ unlabeled_instances, item_manager, schema_name
+ )
+ sorted_instances = sorted(instance_scores, key=lambda x: x[1])
+ self._apply_reordering(sorted_instances, item_manager)
+ return
+
+ # ICL ensemble blending (Phase 5B)
+ if self.config.use_icl_ensemble:
+ rankings = self._blend_icl_scores(
+ rankings, unlabeled_texts, schema_name
+ )
+
+ # Map rankings back to instance IDs
+ sorted_instances = [
+ (unlabeled_instances[idx], score) for idx, score in rankings
+ if idx < len(unlabeled_instances)
+ ]
+
+ # Apply reordering with random sampling
+ self._apply_reordering(sorted_instances, item_manager)
+
+ def _blend_icl_scores(self, rankings: List[Tuple[int, float]],
+ texts: List[str], schema_name: str) -> List[Tuple[int, float]]:
+ """Blend query strategy scores with ICL predictions."""
+ try:
+ from potato.ai.icl_labeler import get_icl_labeler
+ icl_labeler = get_icl_labeler()
+ if icl_labeler is None or not icl_labeler.has_enough_examples(schema_name):
+ return rankings
+
+ # Determine interpolation weight based on annotation count
+ params = self.config.icl_ensemble_params
+ initial_w = params.get("initial_icl_weight", 0.7)
+ final_w = params.get("final_icl_weight", 0.2)
+ transition = params.get("transition_instances", 100)
+
+ annotated_count = len(self._annotated_texts.get(schema_name, []))
+ progress = min(1.0, annotated_count / max(1, transition))
+ icl_weight = initial_w + (final_w - initial_w) * progress
+ strategy_weight = 1.0 - icl_weight
+
+ # Get ICL confidence for each text
+ icl_scores = {}
+ for idx, text in enumerate(texts):
+ try:
+ pred = icl_labeler.label_instance(
+ instance_id=f"_al_blend_{idx}",
+ schema_name=schema_name,
+ instance_text=text,
+ )
+ if pred:
+ # Lower confidence = higher priority (more uncertain)
+ icl_scores[idx] = 1.0 - pred.confidence_score
+ else:
+ icl_scores[idx] = 0.5
+ except Exception:
+ icl_scores[idx] = 0.5
+
+ # Normalize strategy scores
+ strategy_map = {idx: score for idx, score in rankings}
+ s_vals = list(strategy_map.values())
+ s_min, s_max = min(s_vals), max(s_vals)
+ s_rng = s_max - s_min if s_max > s_min else 1.0
+
+ # Normalize ICL scores
+ i_vals = list(icl_scores.values())
+ i_min, i_max = min(i_vals), max(i_vals)
+ i_rng = i_max - i_min if i_max > i_min else 1.0
+
+ blended = []
+ for idx, score in rankings:
+ norm_s = (score - s_min) / s_rng
+ norm_i = (icl_scores.get(idx, 0.5) - i_min) / i_rng
+ combined = strategy_weight * norm_s + icl_weight * norm_i
+ blended.append((idx, combined))
+
+ blended.sort(key=lambda x: x[1], reverse=True)
+ return blended
+
+ except ImportError:
+ return rankings
+ except Exception as e:
+ self.logger.warning(f"ICL blending failed: {e}")
+ return rankings
+
+ def _cold_start_reorder(self, item_manager: ItemStateManager):
+ """LLM-based cold-start instance selection (Phase 3A).
+
+ Based on Bayer et al. (2024) ActiveLLM approach. Before enough
+ annotations exist for classifier training, use LLM to estimate
+ which instances are most informative by finding those where LLM
+ confidence is moderate (on the decision boundary).
+ """
+ try:
+ from potato.ai.llm_active_learning import create_llm_active_learning
+
+ llm = create_llm_active_learning(self.config.llm_config)
+
+ # Sample candidate instances
+ all_ids = list(item_manager.get_instance_ids())
+ unannotated = [
+ iid for iid in all_ids
+ if not item_manager.get_annotators_for_item(iid)
+ ]
+
+ if not unannotated:
+ return
+
+ batch_size = min(self.config.cold_start_batch_size, len(unannotated))
+ candidates = random.sample(unannotated, batch_size)
+
+ instances = []
+ for iid in candidates:
+ item = item_manager.get_item(iid)
+ if item:
+ instances.append({"id": iid, "text": item.get_text()})
+
+ if not instances:
+ return
+
+ # Get LLM predictions
+ schema_name = self.schema_cycler.get_current_schema() if self.schema_cycler else None
+ predictions = llm.predict_instances(
+ instances=instances,
+ annotation_instructions="Rate your confidence in labeling this text.",
+ schema_name=schema_name or "default",
+ label_options=["positive", "negative", "neutral"],
+ )
+
+ # Select instances with moderate confidence (decision boundary)
+ moderate = []
+ other = []
+ for pred in predictions:
+ if 0.4 <= pred.confidence_score <= 0.7:
+ moderate.append((pred.instance_id, pred.confidence_score))
+ else:
+ other.append((pred.instance_id, pred.confidence_score))
+
+ # Moderate-confidence first, then others, interleaved with random
+ reordered = [iid for iid, _ in moderate] + [iid for iid, _ in other]
+
+ # Add remaining unannotated instances not in the sample
+ sampled_set = set(candidates)
+ remaining = [iid for iid in unannotated if iid not in sampled_set]
+ random.shuffle(remaining)
+ reordered.extend(remaining)
+
+ item_manager.reorder_instances(reordered)
+ self.logger.info(f"Cold-start LLM reordering: {len(moderate)} moderate-confidence, "
+ f"{len(other)} other, {len(remaining)} remaining")
+
+ except Exception as e:
+ self.logger.warning(f"Cold-start LLM reordering failed: {e}")
+
+ def _route_annotation(self, instance_id: str, instance_text: str,
+ schema_name: str) -> Dict[str, Any]:
+ """Noise-aware annotation routing (Phase 5D).
+
+ Based on Yuan et al. (2024) NoiseAL approach. Routes instances
+ between LLM auto-labeling and human annotation based on LLM
+ confidence levels.
+
+ Returns:
+ Dict with 'route' ('human'|'auto'), optional 'suggestion',
+ and optional 'auto_label'.
+ """
+ if not self.config.annotation_routing:
+ return {"route": "human"}
+
+ thresholds = self.config.routing_thresholds
+ auto_min = thresholds.get("auto_label_min_confidence", 0.9)
+ suggest_below = thresholds.get("show_suggestion_below", 0.5)
+
+ try:
+ from potato.ai.icl_labeler import get_icl_labeler
+ icl_labeler = get_icl_labeler()
+ if icl_labeler is None or not icl_labeler.has_enough_examples(schema_name):
+ return {"route": "human"}
+
+ prediction = icl_labeler.label_instance(
+ instance_id=instance_id,
+ schema_name=schema_name,
+ instance_text=instance_text,
+ )
+
+ if prediction is None:
+ return {"route": "human"}
+
+ confidence = prediction.confidence_score
+
+ if confidence >= auto_min:
+ # High confidence: auto-label with periodic verification
+ should_verify = random.random() < self.config.verification_sample_rate
+ return {
+ "route": "auto",
+ "auto_label": prediction.predicted_label,
+ "confidence": confidence,
+ "needs_verification": should_verify,
+ }
+ elif confidence < suggest_below:
+ # Low confidence: route to human with LLM suggestion
+ return {
+ "route": "human",
+ "suggestion": prediction.predicted_label,
+ "confidence": confidence,
+ }
+ else:
+ # Medium confidence: route to human (most informative)
+ return {"route": "human"}
+
+ except ImportError:
+ return {"route": "human"}
+ except Exception as e:
+ self.logger.warning(f"Annotation routing failed for {instance_id}: {e}")
+ return {"route": "human"}
+
+ def _calculate_confidence_scores(self, instance_ids: List[str], item_manager: ItemStateManager, schema_name: str) -> List[Tuple[str, float]]:
+ """Calculate confidence scores for instances."""
+ instance_scores = []
+ model = self._models[schema_name]
+
+ for instance_id in instance_ids:
+ item = item_manager.get_item(instance_id)
+ if not item:
+ continue
+
+ text = item.get_text()
+
+ try:
+ # Get prediction probabilities
+ probas = model.predict_proba([text])[0]
+ confidence = np.max(probas)
+ instance_scores.append((instance_id, confidence))
+ except Exception as e:
+ self.logger.warning(f"Error predicting for instance {instance_id}: {e}")
+ # Default to low confidence for failed predictions
+ instance_scores.append((instance_id, 0.1))
+
+ return instance_scores
+
+ def _apply_reordering(self, sorted_instances: List[Tuple[str, float]], item_manager: ItemStateManager):
+ """Apply the new ordering to the item manager."""
+ # Extract instance IDs in new order
+ new_order = [instance_id for instance_id, _ in sorted_instances]
+
+ if not new_order:
+ return
+
+ # Apply random sampling
+ random_count = int(len(new_order) * self.config.random_sample_percent)
+ if random_count > 0 and random_count <= len(new_order):
+ random_instances = random.sample(new_order, random_count)
+ else:
+ random_instances = []
+
+ # Interleave active learning and random instances
+ final_order = []
+ al_idx = 0
+ rand_idx = 0
+
+ while al_idx < len(new_order) or rand_idx < len(random_instances):
+ if al_idx < len(new_order):
+ final_order.append(new_order[al_idx])
+ al_idx += 1
+ if rand_idx < len(random_instances):
+ final_order.append(random_instances[rand_idx])
+ rand_idx += 1
+
+ # Update item manager ordering
+ item_manager.reorder_instances(final_order)
+ self.logger.info(f"Reordered {len(final_order)} instances")
+
+ def check_and_trigger_training(self):
+ """Check if training should be triggered and queue it if needed."""
+ if not self.config.enabled:
+ self.logger.debug("Active learning is disabled")
+ return
+
+ with self._lock:
+ # Count current annotations
+ user_manager = get_user_state_manager()
+ current_annotation_count = sum(
+ len(user_state.get_all_annotations())
+ for user_state in user_manager.get_all_users()
+ )
+
+ self.logger.debug(f"Current annotation count: {current_annotation_count}, last count: {self._last_annotation_count}, update_frequency: {self.config.update_frequency}")
+
+ # Check if we should trigger training
+ if (current_annotation_count - self._last_annotation_count) >= self.config.update_frequency:
+ self._training_queue.put("train")
+ self._last_annotation_count = current_annotation_count
+ self.logger.info(f"Queued active learning training (annotations: {current_annotation_count})")
+ else:
+ self.logger.debug("Not enough new annotations to trigger training")
+
+ def force_training(self):
+ """Force immediate training (for testing purposes)."""
+ if not self.config.enabled:
+ self.logger.debug("Active learning is disabled")
+ return
+
+ self.logger.info("Forcing immediate active learning training")
+ self._training_queue.put("train")
+
+ def get_stats(self) -> Dict[str, Any]:
+ """Get active learning statistics."""
+ with self._lock:
+ stats = {
+ "enabled": self.config.enabled,
+ "training_count": self._training_count,
+ "last_training_time": self._last_training_time,
+ "models_trained": list(self._models.keys()),
+ "current_schema": self.schema_cycler.get_current_schema() if self.schema_cycler else None,
+ "schema_order": self.schema_cycler.get_schema_order() if self.schema_cycler else [],
+ "database_enabled": self.config.database_enabled,
+ "model_persistence_enabled": self.config.model_persistence_enabled,
+ "llm_enabled": self.config.llm_enabled,
+ "query_strategy": self.config.query_strategy,
+ "calibrate_probabilities": self.config.calibrate_probabilities,
+ "cold_start_strategy": self.config.cold_start_strategy,
+ "use_icl_ensemble": self.config.use_icl_ensemble,
+ "annotation_routing": self.config.annotation_routing,
+ }
+
+ # Add training metrics if available
+ if self.database_manager:
+ try:
+ stats["training_history"] = [
+ asdict(metrics) for metrics in self.database_manager.get_training_history()
+ ]
+ except Exception as e:
+ self.logger.warning(f"Failed to get training history: {e}")
+ stats["training_history"] = []
+
+ return stats
+
+ def shutdown(self):
+ """Shutdown the active learning manager."""
+ self._stop_training.set()
+ if self._training_thread and self._training_thread.is_alive():
+ self._training_queue.put(None) # Shutdown signal
+ self._training_thread.join(timeout=5.0)
+ self.logger.info("Active learning manager shutdown complete")
+
+
+# Global singleton instance
+ACTIVE_LEARNING_MANAGER = None
+
+
+def parse_active_learning_config(config_data: Dict[str, Any]) -> Optional[ActiveLearningConfig]:
+ """Build an ``ActiveLearningConfig`` from a Potato project config dict.
+
+ Returns None when active learning is not enabled. Maps the keys under the
+ ``active_learning:`` section onto the dataclass fields (unknown keys are
+ ignored), and defaults ``schema_names`` to the project's labelable
+ annotation schemes when not given.
+ """
+ al_dict = (config_data or {}).get("active_learning", {}) or {}
+ if not al_dict.get("enabled"):
+ return None
+
+ valid_fields = {f.name for f in dataclasses.fields(ActiveLearningConfig)}
+ kwargs = {k: v for k, v in al_dict.items() if k in valid_fields}
+
+ # Honor the nested `active_learning.llm:` block (LLM cold-start / ICL).
+ # The dataclass uses flat fields (llm_enabled / llm_config), so translate.
+ llm_block = al_dict.get("llm")
+ if isinstance(llm_block, dict):
+ kwargs.setdefault("llm_enabled", bool(llm_block.get("enabled", False)))
+ kwargs.setdefault("llm_config", llm_block)
+
+ # YAML parses sequences as lists, but sklearn's vectorizers require a tuple
+ # for ngram_range (e.g. (1, 2)). Coerce it so training doesn't fail.
+ vec_params = kwargs.get("vectorizer_params")
+ if isinstance(vec_params, dict) and isinstance(vec_params.get("ngram_range"), list):
+ vec_params = dict(vec_params)
+ vec_params["ngram_range"] = tuple(vec_params["ngram_range"])
+ kwargs["vectorizer_params"] = vec_params
+
+ # resolution_strategy may arrive as a string; coerce to the enum.
+ rs = kwargs.get("resolution_strategy")
+ if isinstance(rs, str):
+ try:
+ kwargs["resolution_strategy"] = ResolutionStrategy(rs)
+ except ValueError:
+ kwargs.pop("resolution_strategy", None)
+
+ # Default schema_names to the labelable schemes in the project.
+ if not kwargs.get("schema_names"):
+ schemes = config_data.get("annotation_schemes", []) or []
+ kwargs["schema_names"] = [
+ s.get("name") for s in schemes
+ if s.get("name") and s.get("annotation_type") in (
+ "radio", "multiselect", "likert", "select"
+ )
+ ]
+
+ return ActiveLearningConfig(**kwargs)
+
+
+def init_active_learning_manager(config: ActiveLearningConfig) -> ActiveLearningManager:
+ """Initialize the global active learning manager."""
+ global ACTIVE_LEARNING_MANAGER
+
+ if ACTIVE_LEARNING_MANAGER is None:
+ ACTIVE_LEARNING_MANAGER = ActiveLearningManager(config)
+
+ return ACTIVE_LEARNING_MANAGER
+
+
+def get_active_learning_manager() -> Optional[ActiveLearningManager]:
+ """Get the global active learning manager."""
+ return ACTIVE_LEARNING_MANAGER
+
+
+def clear_active_learning_manager():
+ """Clear the global active learning manager (for testing)."""
+ global ACTIVE_LEARNING_MANAGER
+ if ACTIVE_LEARNING_MANAGER:
+ ACTIVE_LEARNING_MANAGER.shutdown()
+ ACTIVE_LEARNING_MANAGER = None
diff --git a/potato/adjudication.py b/potato/adjudication.py
new file mode 100644
index 0000000000000000000000000000000000000000..afc731915162ef94a69756d044a597bc85ae16d4
--- /dev/null
+++ b/potato/adjudication.py
@@ -0,0 +1,1224 @@
+"""
+Adjudication Module
+
+This module provides a comprehensive adjudication system where designated users
+review items with multiple annotations, resolve disagreements, and produce
+gold-standard final decisions.
+
+Adjudication is NOT a phase โ it's a parallel workflow accessible via a dedicated
+/adjudicate route, available to users with adjudicator privileges. This avoids
+disrupting the existing phase progression system.
+
+Key Components:
+- AdjudicationConfig: Configuration dataclass for adjudication settings
+- AdjudicationItem: Represents an item eligible for adjudication with all annotations
+- AdjudicationDecision: Represents an adjudicator's final decision on an item
+- AdjudicationManager: Singleton manager for the adjudication workflow
+
+The workflow:
+1. Annotators complete annotations via /annotate (existing workflow)
+2. AdjudicationManager monitors annotation counts and agreement
+3. Items are flagged when criteria are met (min annotations, low agreement)
+4. Adjudicators review items via /adjudicate and submit decisions
+5. Final dataset CLI merges unanimous + adjudicated decisions
+"""
+
+import json
+import logging
+import math
+import os
+import threading
+from collections import Counter, defaultdict
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Dict, List, Optional, Any, Set
+
+logger = logging.getLogger(__name__)
+
+# Singleton instance
+_ADJUDICATION_MANAGER = None
+_ADJUDICATION_LOCK = threading.Lock()
+
+
+@dataclass
+class AdjudicationConfig:
+ """Configuration for adjudication features."""
+ enabled: bool = False
+ adjudicator_users: List[str] = field(default_factory=list)
+
+ # Trigger criteria
+ min_annotations: int = 2
+ require_fully_annotated: bool = False
+ agreement_threshold: float = 0.75
+ show_all_items: bool = False
+
+ # Display options
+ show_annotator_names: bool = True
+ show_timing_data: bool = True
+ show_agreement_scores: bool = True
+ fast_decision_warning_ms: int = 2000
+
+ # Adjudicator metadata fields
+ require_confidence: bool = True
+ require_notes_on_override: bool = False
+ error_taxonomy: List[str] = field(default_factory=lambda: [
+ "ambiguous_text", "guideline_gap", "annotator_error",
+ "edge_case", "subjective_disagreement", "other"
+ ])
+
+ # Similarity (Phase 3, optional)
+ similarity_enabled: bool = False
+ similarity_model: str = "all-MiniLM-L6-v2"
+ similarity_top_k: int = 5
+ similarity_precompute: bool = True
+
+ # Output
+ output_subdir: str = "adjudication"
+
+
+@dataclass
+class AdjudicationItem:
+ """Represents an item eligible for adjudication with all annotator data."""
+ instance_id: str
+ annotations: Dict[str, Dict[str, Any]] # user_id -> {schema: {label: value}}
+ span_annotations: Dict[str, List[Dict]] # user_id -> [span_dict, ...]
+ behavioral_data: Dict[str, Dict] # user_id -> {total_time_ms, ...}
+ agreement_scores: Dict[str, float] # schema_name -> agreement score
+ overall_agreement: float
+ num_annotators: int
+ status: str = "pending" # pending, in_progress, completed, skipped
+ assigned_adjudicator: Optional[str] = None
+ mace_predictions: Dict[str, Any] = field(default_factory=dict) # schema -> predicted label
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary for JSON output."""
+ result = {
+ "instance_id": self.instance_id,
+ "annotations": self.annotations,
+ "span_annotations": self.span_annotations,
+ "behavioral_data": self.behavioral_data,
+ "agreement_scores": self.agreement_scores,
+ "overall_agreement": self.overall_agreement,
+ "num_annotators": self.num_annotators,
+ "status": self.status,
+ "assigned_adjudicator": self.assigned_adjudicator,
+ }
+ if self.mace_predictions:
+ result["mace_predictions"] = self.mace_predictions
+ return result
+
+
+@dataclass
+class AdjudicationDecision:
+ """Represents an adjudicator's final decision on an item."""
+ instance_id: str
+ adjudicator_id: str
+ timestamp: str # ISO format string
+ label_decisions: Dict[str, Any] # schema -> value
+ span_decisions: List[Dict] # list of span dicts
+ source: Dict[str, str] # schema -> "annotator_X" | "adjudicator" | "merged"
+ confidence: str # "high", "medium", "low"
+ notes: str
+ error_taxonomy: List[str]
+ guideline_update_flag: bool = False
+ guideline_update_notes: str = ""
+ time_spent_ms: int = 0
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary for JSON output."""
+ return {
+ "instance_id": self.instance_id,
+ "adjudicator_id": self.adjudicator_id,
+ "timestamp": self.timestamp,
+ "label_decisions": self.label_decisions,
+ "span_decisions": self.span_decisions,
+ "source": self.source,
+ "confidence": self.confidence,
+ "notes": self.notes,
+ "error_taxonomy": self.error_taxonomy,
+ "guideline_update_flag": self.guideline_update_flag,
+ "guideline_update_notes": self.guideline_update_notes,
+ "time_spent_ms": self.time_spent_ms,
+ }
+
+ @classmethod
+ def from_dict(cls, d: Dict[str, Any]) -> "AdjudicationDecision":
+ """Deserialize from dictionary."""
+ return cls(
+ instance_id=d["instance_id"],
+ adjudicator_id=d["adjudicator_id"],
+ timestamp=d["timestamp"],
+ label_decisions=d.get("label_decisions", {}),
+ span_decisions=d.get("span_decisions", []),
+ source=d.get("source", {}),
+ confidence=d.get("confidence", "medium"),
+ notes=d.get("notes", ""),
+ error_taxonomy=d.get("error_taxonomy", []),
+ guideline_update_flag=d.get("guideline_update_flag", False),
+ guideline_update_notes=d.get("guideline_update_notes", ""),
+ time_spent_ms=d.get("time_spent_ms", 0),
+ )
+
+
+class AdjudicationManager:
+ """
+ Manages the adjudication workflow including queue building, agreement
+ computation, decision storage, and final dataset generation.
+
+ Follows the singleton pattern used by QualityControlManager.
+ """
+
+ def __init__(self, config: Dict[str, Any]):
+ """
+ Initialize the adjudication manager.
+
+ Args:
+ config: The full application configuration dictionary
+ """
+ self.config = config
+ self.logger = logging.getLogger(__name__)
+ self._lock = threading.RLock()
+
+ # Parse configuration
+ self.adj_config = self._parse_config(config)
+
+ # Queue and decisions
+ self.queue: Dict[str, AdjudicationItem] = {} # instance_id -> AdjudicationItem
+ self.decisions: Dict[str, AdjudicationDecision] = {} # instance_id -> decision
+ self._queue_built = False
+
+ # Load any previously saved decisions
+ self._load_decisions()
+
+ # Initialize similarity engine (Phase 3)
+ self.similarity_engine = None
+ if self.adj_config.similarity_enabled:
+ from potato.similarity import init_similarity_engine
+ self.similarity_engine = init_similarity_engine(config, self.adj_config)
+ if (self.similarity_engine and self.similarity_engine.enabled
+ and self.adj_config.similarity_precompute):
+ self._precompute_similarities()
+
+ self.logger.info(
+ f"AdjudicationManager initialized: enabled={self.adj_config.enabled}, "
+ f"adjudicators={self.adj_config.adjudicator_users}"
+ )
+
+ def _parse_config(self, config: Dict[str, Any]) -> AdjudicationConfig:
+ """Parse adjudication configuration from the main config."""
+ adj = AdjudicationConfig()
+
+ adj_config = config.get("adjudication", {})
+ if not adj_config or not adj_config.get("enabled", False):
+ return adj
+
+ adj.enabled = True
+ adj.adjudicator_users = adj_config.get("adjudicator_users", [])
+ adj.min_annotations = adj_config.get("min_annotations", 2)
+ adj.require_fully_annotated = adj_config.get("require_fully_annotated", False)
+ adj.agreement_threshold = adj_config.get("agreement_threshold", 0.75)
+ adj.show_all_items = adj_config.get("show_all_items", False)
+ adj.show_annotator_names = adj_config.get("show_annotator_names", True)
+ adj.show_timing_data = adj_config.get("show_timing_data", True)
+ adj.show_agreement_scores = adj_config.get("show_agreement_scores", True)
+ adj.fast_decision_warning_ms = adj_config.get("fast_decision_warning_ms", 2000)
+ adj.require_confidence = adj_config.get("require_confidence", True)
+ adj.require_notes_on_override = adj_config.get("require_notes_on_override", False)
+
+ if "error_taxonomy" in adj_config:
+ adj.error_taxonomy = adj_config["error_taxonomy"]
+
+ # Similarity settings
+ sim_config = adj_config.get("similarity", {})
+ if sim_config.get("enabled", False):
+ adj.similarity_enabled = True
+ adj.similarity_model = sim_config.get("model", "all-MiniLM-L6-v2")
+ adj.similarity_top_k = sim_config.get("top_k", 5)
+ adj.similarity_precompute = sim_config.get("precompute_on_start", True)
+
+ adj.output_subdir = adj_config.get("output_subdir", "adjudication")
+
+ return adj
+
+ def is_adjudicator(self, username: str) -> bool:
+ """Check if a user is an authorized adjudicator."""
+ if not self.adj_config.enabled:
+ return False
+ return username in self.adj_config.adjudicator_users
+
+ def build_queue(self) -> List[AdjudicationItem]:
+ """
+ Scan all user annotations and build the adjudication queue.
+
+ Items become eligible when they have enough annotations and
+ agreement is below the threshold.
+
+ Returns:
+ List of AdjudicationItem objects
+ """
+ from potato.user_state_management import get_user_state_manager
+ from potato.item_state_management import get_item_state_manager
+
+ with self._lock:
+ usm = get_user_state_manager()
+ ism = get_item_state_manager()
+
+ # Get all annotation schemes from config
+ annotation_schemes = self.config.get("annotation_schemes", [])
+ scheme_names = [s.get("name", "") for s in annotation_schemes]
+
+ # Iterate over all items
+ for instance_id, item in ism.instance_id_to_instance.items():
+ instance_id_str = str(instance_id)
+
+ # Skip if already decided
+ if instance_id_str in self.decisions:
+ if instance_id_str not in self.queue:
+ continue
+ # Mark as completed if decision exists
+ self.queue[instance_id_str].status = "completed"
+ continue
+
+ # Get all annotators for this item
+ annotators = ism.instance_annotators.get(instance_id, set())
+ # Filter out adjudicators from annotator list
+ annotators = {
+ u for u in annotators
+ if u not in self.adj_config.adjudicator_users
+ }
+
+ if len(annotators) < self.adj_config.min_annotations:
+ continue
+
+ # Check if we require fully annotated items
+ if self.adj_config.require_fully_annotated:
+ max_per_item = ism.max_annotations_per_item
+ if max_per_item > 0 and len(annotators) < max_per_item:
+ continue
+
+ # Collect annotations from all annotators
+ item_annotations = {}
+ item_spans = {}
+ item_behavioral = {}
+
+ for user_id in annotators:
+ user_state = usm.get_user_state(user_id)
+ if not user_state:
+ continue
+
+ # Get label annotations
+ label_annots = user_state.instance_id_to_label_to_value.get(
+ instance_id_str, {}
+ )
+ if label_annots:
+ item_annotations[user_id] = self._serialize_labels(label_annots)
+
+ # Get span annotations
+ span_annots = user_state.instance_id_to_span_to_value.get(
+ instance_id_str, {}
+ )
+ if span_annots:
+ item_spans[user_id] = self._serialize_spans(span_annots)
+
+ # Get behavioral data
+ bd = user_state.instance_id_to_behavioral_data.get(
+ instance_id_str, {}
+ )
+ if bd:
+ item_behavioral[user_id] = self._serialize_behavioral(bd)
+
+ if not item_annotations and not item_spans:
+ continue
+
+ # Compute agreement scores
+ agreement_scores = self._compute_agreement(
+ item_annotations, scheme_names
+ )
+ overall = self._compute_overall_agreement(agreement_scores)
+
+ # Filter by agreement threshold
+ if not self.adj_config.show_all_items:
+ if overall >= self.adj_config.agreement_threshold:
+ continue
+
+ # Preserve existing status if already in queue
+ existing = self.queue.get(instance_id_str)
+ status = existing.status if existing else "pending"
+ assigned = existing.assigned_adjudicator if existing else None
+
+ # Enrich with MACE predictions if available
+ mace_preds = {}
+ try:
+ from potato.mace_manager import get_mace_manager
+ mace_mgr = get_mace_manager()
+ if mace_mgr and mace_mgr.results:
+ for sname in scheme_names:
+ pred = mace_mgr.get_prediction(instance_id_str, sname)
+ if pred is not None:
+ mace_preds[sname] = pred
+ except Exception:
+ pass # MACE is optional
+
+ self.queue[instance_id_str] = AdjudicationItem(
+ instance_id=instance_id_str,
+ annotations=item_annotations,
+ span_annotations=item_spans,
+ behavioral_data=item_behavioral,
+ agreement_scores=agreement_scores,
+ overall_agreement=overall,
+ num_annotators=len(annotators),
+ status=status,
+ assigned_adjudicator=assigned,
+ mace_predictions=mace_preds,
+ )
+
+ self._queue_built = True
+ return list(self.queue.values())
+
+ def try_enqueue_item(self, instance_id: str) -> bool:
+ """
+ Evaluate a single item and, if it qualifies, add it to the queue.
+
+ Called when an overlap-sample item saturates so that low-agreement
+ items show up in the adjudication queue without needing a full
+ ``build_queue()`` rescan. Returns True if the item ended up in the
+ queue, False otherwise.
+ """
+ if not self.adj_config.enabled:
+ return False
+
+ from potato.user_state_management import get_user_state_manager
+ from potato.item_state_management import get_item_state_manager
+
+ usm = get_user_state_manager()
+ ism = get_item_state_manager()
+ if usm is None or ism is None:
+ return False
+
+ with self._lock:
+ instance_id_str = str(instance_id)
+ if instance_id_str in self.decisions:
+ return False
+ item = ism.instance_id_to_instance.get(instance_id)
+ if item is None:
+ return False
+
+ annotators = {
+ u for u in ism.instance_annotators.get(instance_id, set())
+ if u not in self.adj_config.adjudicator_users
+ }
+ if len(annotators) < self.adj_config.min_annotations:
+ return False
+
+ scheme_names = [s.get("name", "") for s in self.config.get("annotation_schemes", [])]
+ item_annotations: Dict[str, Any] = {}
+ item_spans: Dict[str, Any] = {}
+ item_behavioral: Dict[str, Any] = {}
+ for user_id in annotators:
+ ustate = usm.get_user_state(user_id)
+ if not ustate:
+ continue
+ la = ustate.instance_id_to_label_to_value.get(instance_id_str, {})
+ if la:
+ item_annotations[user_id] = self._serialize_labels(la)
+ sa = ustate.instance_id_to_span_to_value.get(instance_id_str, {})
+ if sa:
+ item_spans[user_id] = self._serialize_spans(sa)
+ bd = ustate.instance_id_to_behavioral_data.get(instance_id_str, {})
+ if bd:
+ item_behavioral[user_id] = self._serialize_behavioral(bd)
+
+ if not item_annotations and not item_spans:
+ return False
+
+ agreement_scores = self._compute_agreement(item_annotations, scheme_names)
+ overall = self._compute_overall_agreement(agreement_scores)
+ if not self.adj_config.show_all_items:
+ if overall >= self.adj_config.agreement_threshold:
+ return False
+
+ existing = self.queue.get(instance_id_str)
+ self.queue[instance_id_str] = AdjudicationItem(
+ instance_id=instance_id_str,
+ annotations=item_annotations,
+ span_annotations=item_spans,
+ behavioral_data=item_behavioral,
+ agreement_scores=agreement_scores,
+ overall_agreement=overall,
+ num_annotators=len(annotators),
+ status=existing.status if existing else "pending",
+ assigned_adjudicator=existing.assigned_adjudicator if existing else None,
+ )
+ self.logger.info(
+ "Auto-routed item %s into adjudication queue (overall agreement=%.3f, "
+ "threshold=%.3f, annotators=%d)",
+ instance_id_str, overall, self.adj_config.agreement_threshold, len(annotators),
+ )
+ return True
+
+ def _serialize_labels(self, label_data: Dict) -> Dict[str, Any]:
+ """Convert label annotation data to serializable dict."""
+ result = {}
+ for key, value in label_data.items():
+ # Key might be a Label object or a string
+ if hasattr(key, 'get_schema'):
+ schema = key.get_schema()
+ name = key.get_name()
+ if schema not in result:
+ result[schema] = {}
+ result[schema][name] = value
+ elif isinstance(key, str):
+ result[key] = value
+ else:
+ result[str(key)] = value
+ return result
+
+ def _serialize_spans(self, span_data: Dict) -> List[Dict]:
+ """Convert span annotation data to serializable list."""
+ spans = []
+ for key, value in span_data.items():
+ if hasattr(key, 'get_schema'):
+ spans.append({
+ "schema": key.get_schema(),
+ "name": key.get_name(),
+ "title": key.get_title() if hasattr(key, 'get_title') else "",
+ "start": key.get_start(),
+ "end": key.get_end(),
+ "id": key.get_id(),
+ "target_field": key.get_target_field() if hasattr(key, 'get_target_field') else None,
+ })
+ elif isinstance(value, dict):
+ spans.append(value)
+ return spans
+
+ def _serialize_behavioral(self, bd) -> Dict:
+ """Convert behavioral data to serializable dict."""
+ if hasattr(bd, 'to_dict'):
+ return bd.to_dict()
+ elif isinstance(bd, dict):
+ return bd
+ return {}
+
+ def _compute_agreement(
+ self, item_annotations: Dict[str, Dict], scheme_names: List[str]
+ ) -> Dict[str, float]:
+ """
+ Compute per-schema agreement for an item.
+
+ Uses simple percentage agreement (proportion of annotators who chose
+ the most common label). For more sophisticated metrics, simpledorff
+ can be used but requires multiple items.
+
+ Returns:
+ Dict mapping schema_name to agreement score (0.0 - 1.0)
+ """
+ agreement_scores = {}
+
+ for schema in scheme_names:
+ values = []
+ for user_id, user_annots in item_annotations.items():
+ if schema in user_annots:
+ val = user_annots[schema]
+ # Normalize to comparable form
+ if isinstance(val, dict):
+ # Radio stores {label: label} (value is the label string)
+ # and multiselect stores {label: value/true}. A label is
+ # "selected" when its value is present/truthy. The old
+ # filter (v is True / == "true" / == 1) dropped radio's
+ # string value, collapsing every annotator to an empty
+ # frozenset -> a spurious 1.0 agreement even on total
+ # disagreement.
+ falsey = (False, None, "", "false", "False", 0, "0")
+ selected = frozenset(
+ k for k, v in val.items() if v not in falsey
+ )
+ values.append(selected)
+ else:
+ values.append(val)
+
+ if len(values) < 2:
+ continue
+
+ # Compute pairwise agreement (percentage)
+ agree_count = 0
+ total_pairs = 0
+ for i in range(len(values)):
+ for j in range(i + 1, len(values)):
+ total_pairs += 1
+ if values[i] == values[j]:
+ agree_count += 1
+
+ agreement_scores[schema] = (
+ agree_count / total_pairs if total_pairs > 0 else 1.0
+ )
+
+ return agreement_scores
+
+ def _compute_overall_agreement(self, agreement_scores: Dict[str, float]) -> float:
+ """Compute overall agreement as the mean of per-schema scores."""
+ if not agreement_scores:
+ return 1.0
+ return sum(agreement_scores.values()) / len(agreement_scores)
+
+ def get_queue(
+ self,
+ adjudicator_id: Optional[str] = None,
+ filter_status: Optional[str] = None,
+ ) -> List[AdjudicationItem]:
+ """
+ Get the adjudication queue, optionally filtered by status.
+
+ Args:
+ adjudicator_id: Optional adjudicator to filter by assignment
+ filter_status: Optional status filter ("pending", "completed", etc.)
+
+ Returns:
+ List of AdjudicationItem objects
+ """
+ with self._lock:
+ if not self._queue_built:
+ self.build_queue()
+
+ items = list(self.queue.values())
+
+ if filter_status:
+ items = [i for i in items if i.status == filter_status]
+
+ # Sort: pending first, then by agreement (lowest first)
+ items.sort(key=lambda x: (
+ 0 if x.status == "pending" else 1 if x.status == "in_progress" else 2,
+ x.overall_agreement,
+ ))
+
+ return items
+
+ def get_item(self, instance_id: str) -> Optional[AdjudicationItem]:
+ """
+ Get full item data for adjudication.
+
+ Args:
+ instance_id: The instance ID to retrieve
+
+ Returns:
+ AdjudicationItem or None if not in queue
+ """
+ with self._lock:
+ if not self._queue_built:
+ self.build_queue()
+ return self.queue.get(str(instance_id))
+
+ def get_item_text(self, instance_id: str) -> str:
+ """Get the text content for an item."""
+ from potato.item_state_management import get_item_state_manager
+
+ ism = get_item_state_manager()
+ item = ism.instance_id_to_instance.get(instance_id)
+ if item:
+ # Use text_key from config if available
+ text_key = self.config.get("item_properties", {}).get("text_key", "text")
+ data = item.get_data()
+ if isinstance(data, dict) and text_key in data:
+ return data[text_key]
+ return item.get_text()
+ return ""
+
+ def get_item_data(self, instance_id: str) -> Dict[str, Any]:
+ """Get the full raw data for an item."""
+ from potato.item_state_management import get_item_state_manager
+
+ ism = get_item_state_manager()
+ item = ism.instance_id_to_instance.get(instance_id)
+ if item:
+ data = item.get_data()
+ if isinstance(data, dict):
+ return data
+ return {"text": str(data)}
+ return {}
+
+ def get_next_item(self, adjudicator_id: str) -> Optional[AdjudicationItem]:
+ """Get the next pending item for an adjudicator."""
+ items = self.get_queue(filter_status="pending")
+ if items:
+ return items[0]
+ return None
+
+ def skip_item(self, instance_id: str, adjudicator_id: str) -> bool:
+ """Mark an item as skipped."""
+ with self._lock:
+ item = self.queue.get(str(instance_id))
+ if item:
+ item.status = "skipped"
+ return True
+ return False
+
+ def submit_decision(self, decision: AdjudicationDecision) -> bool:
+ """
+ Submit an adjudication decision.
+
+ Args:
+ decision: The AdjudicationDecision to save
+
+ Returns:
+ True if successful
+ """
+ with self._lock:
+ instance_id = str(decision.instance_id)
+ self.decisions[instance_id] = decision
+
+ # Update queue status
+ if instance_id in self.queue:
+ self.queue[instance_id].status = "completed"
+ self.queue[instance_id].assigned_adjudicator = decision.adjudicator_id
+
+ # Persist to disk
+ self._save_decisions()
+
+ self.logger.info(
+ f"Adjudication decision saved for {instance_id} "
+ f"by {decision.adjudicator_id}"
+ )
+ return True
+
+ def get_stats(self) -> Dict[str, Any]:
+ """Get adjudication progress statistics."""
+ with self._lock:
+ if not self._queue_built:
+ self.build_queue()
+
+ total = len(self.queue)
+ completed = sum(
+ 1 for i in self.queue.values() if i.status == "completed"
+ )
+ pending = sum(
+ 1 for i in self.queue.values() if i.status == "pending"
+ )
+ skipped = sum(
+ 1 for i in self.queue.values() if i.status == "skipped"
+ )
+ in_progress = sum(
+ 1 for i in self.queue.values() if i.status == "in_progress"
+ )
+
+ avg_agreement = 0.0
+ if self.queue:
+ avg_agreement = sum(
+ i.overall_agreement for i in self.queue.values()
+ ) / len(self.queue)
+
+ # Per-adjudicator stats
+ adjudicator_stats = defaultdict(lambda: {"completed": 0, "total_time_ms": 0})
+ for decision in self.decisions.values():
+ adj_id = decision.adjudicator_id
+ adjudicator_stats[adj_id]["completed"] += 1
+ adjudicator_stats[adj_id]["total_time_ms"] += decision.time_spent_ms
+
+ return {
+ "total": total,
+ "completed": completed,
+ "pending": pending,
+ "skipped": skipped,
+ "in_progress": in_progress,
+ "completion_rate": completed / total if total > 0 else 0.0,
+ "avg_agreement": avg_agreement,
+ "adjudicator_stats": dict(adjudicator_stats),
+ }
+
+ def get_decision(self, instance_id: str) -> Optional[AdjudicationDecision]:
+ """Get the decision for an item, if one exists."""
+ return self.decisions.get(str(instance_id))
+
+ # ------------------------------------------------------------------
+ # Phase 3: Similarity integration
+ # ------------------------------------------------------------------
+
+ def _precompute_similarities(self) -> None:
+ """Precompute embeddings for all items in the item state manager."""
+ if not self.similarity_engine or not self.similarity_engine.enabled:
+ return
+
+ from potato.item_state_management import get_item_state_manager
+
+ try:
+ ism = get_item_state_manager()
+ item_texts = {}
+ for instance_id, item in ism.instance_id_to_instance.items():
+ text = self.get_item_text(str(instance_id))
+ if text:
+ item_texts[str(instance_id)] = text
+
+ if item_texts:
+ count = self.similarity_engine.precompute_embeddings(item_texts)
+ self.logger.info(f"Precomputed {count} similarity embeddings")
+ except Exception as e:
+ self.logger.error(f"Error precomputing similarities: {e}")
+
+ def get_similar_items(
+ self, instance_id: str, include_metadata: bool = True
+ ) -> List[Dict[str, Any]]:
+ """
+ Get similar items for a given instance, enriched with queue metadata.
+
+ Args:
+ instance_id: The reference instance ID
+ include_metadata: Whether to include decision/consensus data
+
+ Returns:
+ List of dicts with instance_id, similarity, and optional metadata
+ """
+ if not self.similarity_engine or not self.similarity_engine.enabled:
+ return []
+
+ similar = self.similarity_engine.find_similar(instance_id)
+ results = []
+
+ for other_id, score in similar:
+ entry = {
+ "instance_id": other_id,
+ "similarity": round(score, 4),
+ "text_preview": self.similarity_engine.text_cache.get(
+ other_id, ""
+ ),
+ }
+
+ if include_metadata:
+ queue_item = self.queue.get(other_id)
+ decision = self.decisions.get(other_id)
+
+ entry["in_queue"] = queue_item is not None
+ entry["status"] = queue_item.status if queue_item else None
+ entry["overall_agreement"] = (
+ queue_item.overall_agreement if queue_item else None
+ )
+
+ if decision:
+ entry["decision"] = "completed"
+ entry["consensus_label"] = None
+ else:
+ entry["decision"] = None
+ if queue_item:
+ entry["consensus_label"] = self._get_consensus_label(
+ queue_item
+ )
+ else:
+ entry["consensus_label"] = None
+
+ results.append(entry)
+
+ return results
+
+ def _get_consensus_label(self, item: AdjudicationItem) -> Optional[str]:
+ """
+ Get the majority/consensus label for an item across the first schema.
+
+ Args:
+ item: The AdjudicationItem
+
+ Returns:
+ The most common label value as a string, or None
+ """
+ if not item.annotations:
+ return None
+
+ # Use the first schema that has values
+ for user_annots in item.annotations.values():
+ for schema_name in user_annots:
+ # Collect all values for this schema
+ values = []
+ for ua in item.annotations.values():
+ val = ua.get(schema_name)
+ if val is not None:
+ if isinstance(val, dict):
+ # Multiselect: use frozenset representation
+ selected = sorted(
+ k for k, v in val.items()
+ if v is True or v == "true" or v == 1
+ )
+ values.append(", ".join(selected) if selected else str(val))
+ else:
+ values.append(str(val))
+
+ if values:
+ counter = Counter(values)
+ return counter.most_common(1)[0][0]
+
+ return None
+
+ # ------------------------------------------------------------------
+ # Phase 3: Behavioral signal analysis
+ # ------------------------------------------------------------------
+
+ def get_annotator_signals(
+ self, user_id: str, instance_id: str
+ ) -> Dict[str, Any]:
+ """
+ Compute per-annotator quality signals for a specific item.
+
+ Returns:
+ Dict with user_id, instance_id, flags list, and metrics dict
+ """
+ flags = []
+ metrics = {}
+
+ instance_id = str(instance_id)
+ item = self.queue.get(instance_id)
+ if not item:
+ return {"user_id": user_id, "instance_id": instance_id,
+ "flags": [], "metrics": {}}
+
+ # Get behavioral data for this user on this item
+ bd = item.behavioral_data.get(user_id, {})
+ if hasattr(bd, 'to_dict'):
+ bd = bd.to_dict()
+
+ total_time = bd.get("total_time_ms", 0)
+ metrics["total_time_ms"] = total_time
+
+ # 1. Speed z-score vs user's typical time
+ user_times = self._get_user_times(user_id)
+ if len(user_times) >= 3 and total_time > 0:
+ mean_time = sum(user_times) / len(user_times)
+ std_time = math.sqrt(
+ sum((t - mean_time) ** 2 for t in user_times) / len(user_times)
+ )
+ if std_time > 0:
+ z_score = (total_time - mean_time) / std_time
+ metrics["speed_z_score"] = round(z_score, 2)
+ if z_score < -2.0:
+ flags.append({
+ "type": "unusually_fast",
+ "severity": "high",
+ "message": f"Annotation time ({total_time}ms) is {abs(z_score):.1f} std devs below average"
+ })
+
+ # 2. Fast decision warning
+ fast_threshold = self.adj_config.fast_decision_warning_ms
+ if fast_threshold > 0 and 0 < total_time < fast_threshold:
+ flags.append({
+ "type": "fast_decision",
+ "severity": "medium",
+ "message": f"Decision made in {total_time}ms (below {fast_threshold}ms threshold)"
+ })
+
+ # 3. Annotation change count
+ raw_changes = bd.get("annotation_changes", [])
+ change_count = len(raw_changes) if isinstance(raw_changes, list) else int(raw_changes or 0)
+ metrics["annotation_changes"] = change_count
+ if change_count > 5:
+ flags.append({
+ "type": "excessive_changes",
+ "severity": "medium",
+ "message": f"Made {change_count} annotation changes on this item"
+ })
+
+ # 4. Historical agreement rate with consensus
+ agreement_rate = self._compute_user_agreement_rate(user_id)
+ if agreement_rate is not None:
+ metrics["agreement_rate"] = round(agreement_rate, 3)
+ if agreement_rate < 0.4:
+ flags.append({
+ "type": "low_agreement",
+ "severity": "high",
+ "message": f"Agreement rate with consensus: {agreement_rate:.0%}"
+ })
+
+ # 5. Similar item consistency
+ if self.similarity_engine and self.similarity_engine.enabled:
+ inconsistencies = self._check_similar_item_consistency(
+ user_id, instance_id
+ )
+ metrics["similar_item_inconsistencies"] = inconsistencies
+ if inconsistencies > 0:
+ flags.append({
+ "type": "similar_item_inconsistency",
+ "severity": "medium",
+ "message": f"Different label on {inconsistencies} similar item(s)"
+ })
+
+ return {
+ "user_id": user_id,
+ "instance_id": instance_id,
+ "flags": flags,
+ "metrics": metrics,
+ }
+
+ def _get_user_times(self, user_id: str) -> List[float]:
+ """Collect all annotation times for a user across queue items."""
+ times = []
+ for item in self.queue.values():
+ bd = item.behavioral_data.get(user_id, {})
+ if hasattr(bd, 'to_dict'):
+ bd = bd.to_dict()
+ t = bd.get("total_time_ms", 0)
+ if t > 0:
+ times.append(t)
+ return times
+
+ def _compute_user_agreement_rate(self, user_id: str) -> Optional[float]:
+ """
+ Compute how often a user agrees with the consensus across all items.
+
+ Returns:
+ Float 0-1 or None if insufficient data (needs >= 3 items)
+ """
+ agree_count = 0
+ total_count = 0
+
+ for item in self.queue.values():
+ if user_id not in item.annotations:
+ continue
+
+ consensus = self._get_consensus_label(item)
+ if consensus is None:
+ continue
+
+ user_annots = item.annotations[user_id]
+ # Check the first schema
+ for schema_name, val in user_annots.items():
+ if isinstance(val, dict):
+ selected = sorted(
+ k for k, v in val.items()
+ if v is True or v == "true" or v == 1
+ )
+ user_label = ", ".join(selected) if selected else str(val)
+ else:
+ user_label = str(val)
+
+ if user_label == consensus:
+ agree_count += 1
+ total_count += 1
+ break # Only check first schema
+
+ if total_count < 3:
+ return None
+
+ return agree_count / total_count
+
+ def _check_similar_item_consistency(
+ self, user_id: str, instance_id: str
+ ) -> int:
+ """
+ Check if user's label on similar items (>0.8 similarity) is consistent.
+
+ Returns:
+ Count of similar items where user's label differs
+ """
+ if not self.similarity_engine:
+ return 0
+
+ similar = self.similarity_engine.find_similar(instance_id)
+ if not similar:
+ return 0
+
+ # Get user's label on the current item
+ item = self.queue.get(instance_id)
+ if not item or user_id not in item.annotations:
+ return 0
+
+ user_annots = item.annotations[user_id]
+ current_label = None
+ current_schema = None
+ for schema_name, val in user_annots.items():
+ current_schema = schema_name
+ if isinstance(val, dict):
+ selected = sorted(
+ k for k, v in val.items()
+ if v is True or v == "true" or v == 1
+ )
+ current_label = ", ".join(selected) if selected else str(val)
+ else:
+ current_label = str(val)
+ break
+
+ if current_label is None:
+ return 0
+
+ inconsistencies = 0
+ for other_id, score in similar:
+ if score < 0.8:
+ break # Results are sorted by score desc
+
+ other_item = self.queue.get(other_id)
+ if not other_item or user_id not in other_item.annotations:
+ continue
+
+ other_annots = other_item.annotations[user_id]
+ other_val = other_annots.get(current_schema)
+ if other_val is None:
+ continue
+
+ if isinstance(other_val, dict):
+ selected = sorted(
+ k for k, v in other_val.items()
+ if v is True or v == "true" or v == 1
+ )
+ other_label = ", ".join(selected) if selected else str(other_val)
+ else:
+ other_label = str(other_val)
+
+ if other_label != current_label:
+ inconsistencies += 1
+
+ return inconsistencies
+
+ def _get_output_dir(self) -> str:
+ """Get the adjudication output directory."""
+ output_dir = self.config.get("output_annotation_dir", "annotation_output")
+ adj_dir = os.path.join(output_dir, self.adj_config.output_subdir)
+ os.makedirs(adj_dir, exist_ok=True)
+ return adj_dir
+
+ def _save_decisions(self) -> None:
+ """Persist all decisions to disk."""
+ try:
+ adj_dir = self._get_output_dir()
+ decisions_file = os.path.join(adj_dir, "decisions.json")
+
+ data = {
+ "decisions": [d.to_dict() for d in self.decisions.values()],
+ "last_updated": datetime.now().isoformat(),
+ }
+
+ with open(decisions_file, "w", encoding="utf-8") as f:
+ json.dump(data, f, indent=2)
+
+ except Exception as e:
+ self.logger.error(f"Failed to save adjudication decisions: {e}")
+
+ def _load_decisions(self) -> None:
+ """Load previously saved decisions from disk."""
+ try:
+ output_dir = self.config.get("output_annotation_dir", "annotation_output")
+ adj_dir = os.path.join(output_dir, self.adj_config.output_subdir)
+ decisions_file = os.path.join(adj_dir, "decisions.json")
+
+ if not os.path.exists(decisions_file):
+ return
+
+ with open(decisions_file, "r", encoding="utf-8") as f:
+ data = json.load(f)
+
+ for d in data.get("decisions", []):
+ decision = AdjudicationDecision.from_dict(d)
+ self.decisions[decision.instance_id] = decision
+
+ self.logger.info(
+ f"Loaded {len(self.decisions)} previous adjudication decisions"
+ )
+
+ except Exception as e:
+ self.logger.warning(f"Failed to load adjudication decisions: {e}")
+
+ def generate_final_dataset(self) -> List[Dict[str, Any]]:
+ """
+ Generate the final dataset by merging unanimous agreements
+ and adjudication decisions.
+
+ Returns:
+ List of item dicts with final labels and provenance
+ """
+ from potato.user_state_management import get_user_state_manager
+ from potato.item_state_management import get_item_state_manager
+
+ usm = get_user_state_manager()
+ ism = get_item_state_manager()
+
+ annotation_schemes = self.config.get("annotation_schemes", [])
+ scheme_names = [s.get("name", "") for s in annotation_schemes]
+ results = []
+
+ for instance_id, item in ism.instance_id_to_instance.items():
+ instance_id_str = str(instance_id)
+ result = {
+ "instance_id": instance_id_str,
+ "item_data": item.get_data() if hasattr(item, 'get_data') else {},
+ }
+
+ # Check if we have an adjudication decision
+ decision = self.decisions.get(instance_id_str)
+ if decision:
+ result["labels"] = decision.label_decisions
+ result["spans"] = decision.span_decisions
+ result["source"] = "adjudicated"
+ result["adjudicator"] = decision.adjudicator_id
+ result["confidence"] = decision.confidence
+ result["provenance"] = decision.source
+ results.append(result)
+ continue
+
+ # Check for unanimous agreement
+ annotators = ism.instance_annotators.get(instance_id, set())
+ annotators = {
+ u for u in annotators
+ if u not in self.adj_config.adjudicator_users
+ }
+
+ if len(annotators) < 2:
+ continue
+
+ # Collect annotations
+ annotations = {}
+ for user_id in annotators:
+ user_state = usm.get_user_state(user_id)
+ if not user_state:
+ continue
+ labels = user_state.instance_id_to_label_to_value.get(
+ instance_id_str, {}
+ )
+ if labels:
+ annotations[user_id] = self._serialize_labels(labels)
+
+ if not annotations:
+ continue
+
+ # Check for unanimity per schema
+ unanimous_labels = {}
+ is_unanimous = True
+ for schema in scheme_names:
+ values = []
+ for user_annots in annotations.values():
+ if schema in user_annots:
+ values.append(json.dumps(user_annots[schema], sort_keys=True))
+
+ if len(values) < 2:
+ continue
+
+ if len(set(values)) == 1:
+ unanimous_labels[schema] = json.loads(values[0])
+ else:
+ is_unanimous = False
+
+ if is_unanimous and unanimous_labels:
+ result["labels"] = unanimous_labels
+ result["source"] = "unanimous"
+ result["num_annotators"] = len(annotators)
+ results.append(result)
+ else:
+ result["labels"] = {}
+ result["source"] = "unresolved"
+ result["num_annotators"] = len(annotators)
+ results.append(result)
+
+ return results
+
+
+def init_adjudication_manager(config: Dict[str, Any]) -> Optional[AdjudicationManager]:
+ """Initialize the singleton AdjudicationManager."""
+ global _ADJUDICATION_MANAGER
+
+ with _ADJUDICATION_LOCK:
+ if _ADJUDICATION_MANAGER is None:
+ _ADJUDICATION_MANAGER = AdjudicationManager(config)
+
+ return _ADJUDICATION_MANAGER
+
+
+def get_adjudication_manager() -> Optional[AdjudicationManager]:
+ """Get the singleton AdjudicationManager instance."""
+ return _ADJUDICATION_MANAGER
+
+
+def clear_adjudication_manager():
+ """Clear the singleton (for testing)."""
+ global _ADJUDICATION_MANAGER
+ with _ADJUDICATION_LOCK:
+ _ADJUDICATION_MANAGER = None
diff --git a/potato/adjudication_export.py b/potato/adjudication_export.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1ce83f7f9a71aae460d3da85622b31e7f7a01cf
--- /dev/null
+++ b/potato/adjudication_export.py
@@ -0,0 +1,162 @@
+"""
+Adjudication Export CLI
+
+Generate final datasets by merging unanimous agreements and adjudication decisions.
+
+Usage:
+ python -m potato.adjudication_export --config config.yaml --output final_dataset.jsonl
+ python -m potato.adjudication_export --config config.yaml --output final.csv --format csv
+ python -m potato.adjudication_export --config config.yaml --output final.json --format json
+"""
+
+import argparse
+import csv
+import json
+import os
+import sys
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+ parser = argparse.ArgumentParser(
+ description="Export adjudicated dataset from Potato annotation project"
+ )
+ parser.add_argument(
+ "--config", required=True,
+ help="Path to the Potato config YAML file"
+ )
+ parser.add_argument(
+ "--output", required=True,
+ help="Output file path"
+ )
+ parser.add_argument(
+ "--format", choices=["jsonl", "json", "csv"], default="jsonl",
+ help="Output format (default: jsonl)"
+ )
+ parser.add_argument(
+ "--include-unresolved", action="store_true",
+ help="Include items without adjudication or consensus"
+ )
+ parser.add_argument(
+ "--verbose", "-v", action="store_true",
+ help="Verbose output"
+ )
+
+ args = parser.parse_args()
+
+ if args.verbose:
+ logging.basicConfig(level=logging.DEBUG)
+ else:
+ logging.basicConfig(level=logging.INFO)
+
+ # Load config
+ from potato.server_utils.config_module import init_config, config
+ try:
+ init_config(args.config)
+ except Exception as e:
+ print(f"Error loading config: {e}", file=sys.stderr)
+ sys.exit(1)
+
+ # Initialize state managers
+ from potato.item_state_management import init_item_state_manager
+ from potato.user_state_management import init_user_state_manager
+
+ init_user_state_manager(config)
+ init_item_state_manager(config)
+
+ # Load data (this loads items and user annotations from disk)
+ # We need a minimal load - just items and user states
+ from potato.flask_server import load_instance_data, load_user_data
+ load_instance_data(config)
+ load_user_data(config)
+
+ # Initialize adjudication manager
+ from potato.adjudication import init_adjudication_manager
+ adj_mgr = init_adjudication_manager(config)
+
+ if not adj_mgr or not adj_mgr.adj_config.enabled:
+ print("Adjudication is not enabled in this config.", file=sys.stderr)
+ sys.exit(1)
+
+ # Build queue to compute agreements
+ adj_mgr.build_queue()
+
+ # Generate final dataset
+ results = adj_mgr.generate_final_dataset()
+
+ # Filter unresolved if not requested
+ if not args.include_unresolved:
+ results = [r for r in results if r.get("source") != "unresolved"]
+
+ # Write output
+ output_path = args.output
+ fmt = args.format
+
+ if fmt == "jsonl":
+ with open(output_path, "w") as f:
+ for item in results:
+ f.write(json.dumps(item) + "\n")
+
+ elif fmt == "json":
+ with open(output_path, "w") as f:
+ json.dump(results, f, indent=2)
+
+ elif fmt == "csv":
+ if not results:
+ print("No results to export.", file=sys.stderr)
+ sys.exit(0)
+
+ # Flatten for CSV
+ fieldnames = set()
+ flat_results = []
+ for item in results:
+ flat = {
+ "instance_id": item["instance_id"],
+ "source": item.get("source", ""),
+ }
+ # Flatten labels
+ labels = item.get("labels", {})
+ for schema, value in labels.items():
+ if isinstance(value, dict):
+ flat[schema] = json.dumps(value)
+ else:
+ flat[schema] = value
+
+ # Add provenance fields
+ if "adjudicator" in item:
+ flat["adjudicator"] = item["adjudicator"]
+ if "confidence" in item:
+ flat["confidence"] = item["confidence"]
+ if "num_annotators" in item:
+ flat["num_annotators"] = item["num_annotators"]
+
+ fieldnames.update(flat.keys())
+ flat_results.append(flat)
+
+ # Sort fieldnames for consistent output
+ fieldnames = sorted(fieldnames)
+
+ with open(output_path, "w", newline="") as f:
+ writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
+ writer.writeheader()
+ writer.writerows(flat_results)
+
+ # Summary
+ total = len(results)
+ unanimous = sum(1 for r in results if r.get("source") == "unanimous")
+ adjudicated = sum(1 for r in results if r.get("source") == "adjudicated")
+ unresolved = sum(1 for r in results if r.get("source") == "unresolved")
+
+ print(f"\nExport complete: {output_path}")
+ print(f" Total items: {total}")
+ print(f" Unanimous: {unanimous}")
+ print(f" Adjudicated: {adjudicated}")
+ if args.include_unresolved:
+ print(f" Unresolved: {unresolved}")
+ print(f" Format: {fmt}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/potato/admin.py b/potato/admin.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ebc8733ac29b6c54f8546db1b24ccf45c7579d2
--- /dev/null
+++ b/potato/admin.py
@@ -0,0 +1,2253 @@
+"""
+Admin Dashboard Module
+
+This module provides comprehensive admin functionality for the annotation platform,
+including dashboard data generation, timing analysis, and configuration management.
+
+The admin dashboard offers:
+- Real-time overview of annotation progress and statistics
+- Detailed annotator performance metrics and timing analysis
+- Instance-level annotation tracking and disagreement analysis
+- Configuration management and system state monitoring
+- Question and annotation scheme analysis
+- User progress tracking and completion statistics
+- Comprehensive annotation history tracking and suspicious activity detection
+- Performance metrics and quality assurance monitoring
+- Session tracking and behavioral analysis
+
+Key Components:
+- AdminDashboard: Main class for admin functionality
+- AnnotatorTimingData: Data class for annotator timing information
+- InstanceData: Data class for instance information and statistics
+- Dashboard data generation and analysis functions
+- Configuration update and management functions
+- AnnotationHistoryAnalyzer: Advanced history analysis and suspicious activity detection
+
+The dashboard provides insights into:
+- Overall annotation progress and completion rates
+- Individual annotator performance and efficiency
+- Annotation quality through disagreement analysis
+- System configuration and operational status
+- Real-time monitoring of active annotation sessions
+- Fine-grained annotation timing and behavioral patterns
+- Suspicious activity detection and quality assurance
+- Session-based performance analysis
+
+Access Control:
+- Admin access is controlled via API key authentication
+- Debug mode allows admin access without API key
+- All admin endpoints require proper authentication
+"""
+
+import json
+import logging
+import datetime
+from typing import Dict, List, Optional, Tuple, Any
+from collections import defaultdict, Counter
+from dataclasses import dataclass
+from flask import request, jsonify, session
+
+from potato.flask_server import (
+ config, logger, get_user_state_manager, get_item_state_manager,
+ get_users, get_total_annotations
+)
+from potato.annotation_history import AnnotationHistoryManager, AnnotationAction
+from potato.quality_control import get_quality_control_manager
+
+@dataclass
+class AnnotatorTimingData:
+ """
+ Data class for annotator timing information.
+
+ This class encapsulates timing metrics for individual annotators,
+ including total annotations, working time, and performance statistics.
+ Now enhanced with annotation history tracking and suspicious activity detection.
+ """
+ user_id: str
+ total_annotations: int
+ total_seconds: int
+ average_seconds_per_annotation: float
+ last_activity: Optional[datetime.datetime]
+ current_instance_time: Optional[int]
+ annotations_per_hour: float
+ phase: str
+ has_assignments: bool
+ remaining_assignments: bool
+
+ # Annotation history metrics
+ total_actions: int
+ average_action_time_ms: float
+ fastest_action_time_ms: int
+ slowest_action_time_ms: int
+ actions_per_minute: float
+ suspicious_score: float
+ suspicious_level: str
+ fast_actions_count: int
+ burst_actions_count: int
+ session_start_time: Optional[datetime.datetime]
+ current_session_duration_minutes: Optional[float]
+ recent_actions_count: int # Actions in last 5 minutes
+
+ # Training metrics
+ training_completed: bool
+ training_correct_answers: int
+ training_total_attempts: int
+ training_pass_rate: float
+ training_current_question: int
+ training_total_questions: int
+
+@dataclass
+class InstanceData:
+ """
+ Data class for instance information.
+
+ This class encapsulates information about annotation instances,
+ including annotation counts, disagreement scores, and annotator lists.
+ """
+ id: str
+ text: str
+ displayed_text: str
+ annotation_count: int
+ completion_percentage: float
+ most_frequent_label: Optional[str]
+ label_disagreement: float
+ annotators: List[str]
+ num_ai_instance: int
+ average_time_per_annotation: Optional[float]
+
+class AdminDashboard:
+ """
+ Main class for admin dashboard functionality.
+
+ This class provides comprehensive admin features including dashboard
+ data generation, timing analysis, configuration management, and
+ system monitoring capabilities.
+ """
+
+ def __init__(self):
+ """Initialize the admin dashboard."""
+ self.logger = logging.getLogger(__name__)
+
+ def check_admin_access(self) -> bool:
+ """
+ Check if the current request has admin access via API key.
+
+ Validates against all key sources (config, env var, auto-generated file)
+ and accepts keys from X-API-Key header or session.
+
+ Returns:
+ bool: True if admin access is granted, False otherwise
+ """
+ from potato.server_utils.admin_key import validate_admin_api_key
+ api_key = request.headers.get('X-API-Key') or session.get('admin_api_key')
+ return validate_admin_api_key(api_key, config)
+
+ def get_dashboard_overview(self) -> Dict[str, Any]:
+ """
+ Get comprehensive dashboard overview data.
+
+ This method generates a complete overview of the annotation system,
+ including user statistics, annotation progress, and system configuration.
+
+ Returns:
+ Dict containing overview statistics with the following structure:
+ - overview: User counts, annotation counts, completion percentages
+ - config: System configuration and settings
+
+ Side Effects:
+ - Logs errors if data generation fails
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ usm = get_user_state_manager()
+ ism = get_item_state_manager()
+
+ # Get all users and their states
+ users = get_users()
+ total_annotations = get_total_annotations()
+
+ # Calculate user statistics
+ active_users = 0
+ completed_users = 0
+ total_working_time = 0
+
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if user_state:
+ if user_state.get_phase().value == "ANNOTATION":
+ active_users += 1
+ elif user_state.get_phase().value == "DONE":
+ completed_users += 1
+
+ # Get timing data
+ timing_data = self._get_annotator_timing_data(username)
+ if timing_data:
+ total_working_time += timing_data.total_seconds
+
+ # Get item statistics
+ items = ism.items()
+ items_with_annotations = 0
+ total_assignments = 0
+
+ for item in items:
+ item_id = item.get_id()
+ annotators = ism.get_annotators_for_item(item_id)
+ if annotators:
+ items_with_annotations += 1
+ total_assignments += len(annotators)
+
+ # Calculate completion percentages
+ total_items = len(items)
+ completion_percentage = (items_with_annotations / total_items * 100) if total_items > 0 else 0
+
+ # Format total working time
+ hours = total_working_time // 3600
+ minutes = (total_working_time % 3600) // 60
+ formatted_time = f"{hours}h {minutes}m"
+
+ return {
+ "overview": {
+ "total_users": len(users),
+ "active_users": active_users,
+ "completed_users": completed_users,
+ "total_annotations": total_annotations,
+ "total_items": total_items,
+ "items_with_annotations": items_with_annotations,
+ "completion_percentage": round(completion_percentage, 1),
+ "total_assignments": total_assignments,
+ "total_working_time": formatted_time,
+ "average_annotations_per_item": round(total_annotations / total_items, 1) if total_items > 0 else 0
+ },
+ "config": {
+ "annotation_task_name": config.get("annotation_task_name", "Unknown"),
+ "max_annotations_per_user": config.get("max_annotations_per_user", "Unlimited"),
+ "max_annotations_per_item": config.get("max_annotations_per_item", "Unlimited"),
+ "assignment_strategy": config.get("assignment_strategy", "fixed_order"),
+ "debug_mode": config.get("debug", False)
+ }
+ }
+
+ except Exception as e:
+ self.logger.error(f"Error getting dashboard overview: {e}")
+ return {"error": f"Failed to get dashboard overview: {str(e)}"}, 500
+
+ def get_annotators_data(self) -> Dict[str, Any]:
+ """
+ Get detailed annotator data including timing information.
+
+ Returns:
+ Dict containing annotator data with timing analysis
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ usm = get_user_state_manager()
+ users = get_users()
+ annotators_data = []
+
+
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if user_state:
+ timing_data = self._get_annotator_timing_data(username)
+ if timing_data:
+ annotators_data.append({
+ "user_id": timing_data.user_id,
+ "total_annotations": timing_data.total_annotations,
+ "completion_percentage": self._calculate_completion_percentage(timing_data.user_id),
+ "total_seconds": timing_data.total_seconds,
+ "average_seconds_per_annotation": timing_data.average_seconds_per_annotation,
+ "annotations_per_hour": timing_data.annotations_per_hour,
+ "phase": timing_data.phase,
+ "has_assignments": timing_data.has_assignments,
+ "remaining_assignments": timing_data.remaining_assignments,
+ "max_assignments": user_state.get_max_assignments(),
+ "last_activity": timing_data.last_activity.isoformat() if timing_data.last_activity else None,
+ "current_instance_time": timing_data.current_instance_time,
+
+ # NEW: Annotation history metrics
+ "total_actions": timing_data.total_actions,
+ "average_action_time_ms": timing_data.average_action_time_ms,
+ "fastest_action_time_ms": timing_data.fastest_action_time_ms if timing_data.fastest_action_time_ms != float('inf') else None,
+ "slowest_action_time_ms": timing_data.slowest_action_time_ms,
+ "actions_per_minute": timing_data.actions_per_minute,
+ "suspicious_score": timing_data.suspicious_score,
+ "suspicious_level": timing_data.suspicious_level,
+ "fast_actions_count": timing_data.fast_actions_count,
+ "burst_actions_count": timing_data.burst_actions_count,
+ "session_start_time": timing_data.session_start_time.isoformat() if timing_data.session_start_time else None,
+ "current_session_duration_minutes": timing_data.current_session_duration_minutes,
+ "recent_actions_count": timing_data.recent_actions_count,
+
+ # Training metrics
+ "training_completed": timing_data.training_completed,
+ "training_correct_answers": timing_data.training_correct_answers,
+ "training_total_attempts": timing_data.training_total_attempts,
+ "training_pass_rate": round(timing_data.training_pass_rate, 2),
+ "training_current_question": timing_data.training_current_question,
+ "training_total_questions": timing_data.training_total_questions
+ })
+
+ # Sort by suspicious score (highest first)
+ annotators_data.sort(key=lambda x: x["suspicious_score"], reverse=True)
+
+ return {
+ "total_annotators": len(annotators_data),
+ "annotators": annotators_data,
+ "summary": {
+ "high_suspicious_count": len([a for a in annotators_data if a["suspicious_level"] in ["High", "Very High"]]),
+ "medium_suspicious_count": len([a for a in annotators_data if a["suspicious_level"] == "Medium"]),
+ "low_suspicious_count": len([a for a in annotators_data if a["suspicious_level"] == "Low"]),
+ "normal_count": len([a for a in annotators_data if a["suspicious_level"] == "Normal"]),
+ "average_suspicious_score": sum(a["suspicious_score"] for a in annotators_data) / len(annotators_data) if annotators_data else 0
+ }
+ }
+
+ except Exception as e:
+ self.logger.error(f"Error getting annotators data: {e}")
+ return {"error": f"Failed to get annotators data: {str(e)}"}, 500
+
+ def get_annotation_history_data(self, user_id: Optional[str] = None,
+ instance_id: Optional[str] = None,
+ minutes: Optional[int] = None) -> Dict[str, Any]:
+ """
+ Get detailed annotation history data with filtering options.
+
+ Args:
+ user_id: Optional user ID to filter by
+ instance_id: Optional instance ID to filter by
+ minutes: Optional time window in minutes
+
+ Returns:
+ Dict containing annotation history data
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ usm = get_user_state_manager()
+
+ if user_id:
+ # Get history for specific user
+ user_state = usm.get_user_state(user_id)
+ if not user_state:
+ return {"error": f"User {user_id} not found"}, 404
+
+ actions = user_state.get_annotation_history(instance_id)
+ if minutes:
+ actions = user_state.get_recent_actions(minutes)
+
+ return self._format_annotation_history(actions, user_id)
+ else:
+ # Get history for all users
+ all_actions = []
+ users = get_users()
+
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if user_state:
+ user_actions = user_state.get_annotation_history(instance_id)
+ if minutes:
+ user_actions = user_state.get_recent_actions(minutes)
+ all_actions.extend(user_actions)
+
+ return self._format_annotation_history(all_actions, "all_users")
+
+ except Exception as e:
+ self.logger.error(f"Error getting annotation history data: {e}")
+ return {"error": f"Failed to get annotation history data: {str(e)}"}, 500
+
+ def get_suspicious_activity_data(self) -> Dict[str, Any]:
+ """
+ Get comprehensive suspicious activity analysis.
+
+ Returns:
+ Dict containing suspicious activity data
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ usm = get_user_state_manager()
+ users = get_users()
+ suspicious_data = []
+
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if user_state:
+ suspicious_actions = user_state.get_suspicious_activity()
+ if suspicious_actions:
+ suspicious_data.append({
+ "user_id": username,
+ "suspicious_actions_count": len(suspicious_actions),
+ "suspicious_actions": [
+ {
+ "action_id": action.action_id,
+ "timestamp": action.timestamp.isoformat(),
+ "instance_id": action.instance_id,
+ "action_type": action.action_type,
+ "schema_name": action.schema_name,
+ "label_name": action.label_name,
+ "server_processing_time_ms": action.server_processing_time_ms,
+ "session_id": action.session_id
+ }
+ for action in suspicious_actions[:10] # Limit to 10 most recent
+ ]
+ })
+
+ return {
+ "total_users_with_suspicious_activity": len(suspicious_data),
+ "suspicious_activity": suspicious_data
+ }
+
+ except Exception as e:
+ self.logger.error(f"Error getting suspicious activity data: {e}")
+ return {"error": f"Failed to get suspicious activity data: {str(e)}"}, 500
+
+ def get_instances_data(self, page: int = 1, page_size: int = 25,
+ sort_by: str = "annotation_count", sort_order: str = "desc",
+ filter_completion: Optional[str] = None) -> Dict[str, Any]:
+ """
+ Get paginated instances data with sorting and filtering.
+
+ Args:
+ page: Page number (1-based)
+ page_size: Number of instances per page
+ sort_by: Field to sort by (annotation_count, completion_percentage, disagreement, id)
+ sort_order: Sort order (asc, desc)
+ filter_completion: Filter by completion status (completed, incomplete, all)
+
+ Returns:
+ Dict containing paginated instances data
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ ism = get_item_state_manager()
+ items = ism.items()
+
+ # Convert items to InstanceData objects
+ instances_data = []
+ for item in items:
+ item_id = item.get_id()
+ annotators = ism.get_annotators_for_item(item_id)
+ annotation_count = len(annotators) if annotators else 0
+
+ # Calculate completion percentage
+ max_annotations = config.get("max_annotations_per_item", -1)
+ if max_annotations > 0:
+ completion_percentage = min(100, (annotation_count / max_annotations) * 100)
+ else:
+ completion_percentage = 100 if annotation_count > 0 else 0
+
+ # Calculate most frequent label and disagreement
+ most_frequent_label, disagreement = self._calculate_label_statistics(item_id)
+
+ # Calculate average time per annotation
+ avg_time = self._calculate_average_time_per_annotation(item_id)
+
+ instance_data = InstanceData(
+ id=item_id,
+ text=item.get_text(),
+ displayed_text=item.get_displayed_text(),
+ annotation_count=annotation_count,
+ completion_percentage=completion_percentage,
+ most_frequent_label=most_frequent_label,
+ label_disagreement=disagreement,
+ annotators=list(annotators) if annotators else [],
+ average_time_per_annotation=avg_time,
+ num_ai_instance=self._calculate_total_instance_ai(item_id)
+ )
+ instances_data.append(instance_data)
+
+ # Apply filters
+ if filter_completion == "completed":
+ instances_data = [i for i in instances_data if i.completion_percentage >= 100]
+ elif filter_completion == "incomplete":
+ instances_data = [i for i in instances_data if i.completion_percentage < 100]
+
+ # Apply sorting
+ reverse = sort_order.lower() == "desc"
+ if sort_by == "annotation_count":
+ instances_data.sort(key=lambda x: x.annotation_count, reverse=reverse)
+ elif sort_by == "completion_percentage":
+ instances_data.sort(key=lambda x: x.completion_percentage, reverse=reverse)
+ elif sort_by == "disagreement":
+ instances_data.sort(key=lambda x: x.label_disagreement, reverse=reverse)
+ elif sort_by == "id":
+ instances_data.sort(key=lambda x: x.id, reverse=reverse)
+ elif sort_by == "average_time":
+ instances_data.sort(key=lambda x: x.average_time_per_annotation or 0, reverse=reverse)
+
+ # Apply pagination
+ total_instances = len(instances_data)
+ start_idx = (page - 1) * page_size
+ end_idx = start_idx + page_size
+ paginated_instances = instances_data[start_idx:end_idx]
+
+ # Convert to serializable format
+ serialized_instances = []
+ for instance in paginated_instances:
+ serialized_instances.append({
+ "id": instance.id,
+ "text": instance.text[:100] + "..." if len(instance.text) > 100 else instance.text,
+ "displayed_text": instance.displayed_text[:100] + "..." if len(instance.displayed_text) > 100 else instance.displayed_text,
+ "annotation_count": instance.annotation_count,
+ "completion_percentage": round(instance.completion_percentage, 1),
+ "most_frequent_label": instance.most_frequent_label,
+ "label_disagreement": round(instance.label_disagreement, 2),
+ "annotators": instance.annotators,
+ "num_ai_instance": instance.num_ai_instance,
+ "average_time_per_annotation": self._format_seconds(instance.average_time_per_annotation) if instance.average_time_per_annotation else None
+ })
+
+ return {
+ "instances": serialized_instances,
+ "pagination": {
+ "page": page,
+ "page_size": page_size,
+ "total_instances": total_instances,
+ "total_pages": (total_instances + page_size - 1) // page_size,
+ "has_next": end_idx < total_instances,
+ "has_prev": page > 1
+ },
+ "summary": {
+ "completed_instances": len([i for i in instances_data if i.completion_percentage >= 100]),
+ "incomplete_instances": len([i for i in instances_data if i.completion_percentage < 100]),
+ "average_annotations_per_instance": round(sum(i.annotation_count for i in instances_data) / len(instances_data), 1) if instances_data else 0,
+ "average_disagreement": round(sum(i.label_disagreement for i in instances_data) / len(instances_data), 2) if instances_data else 0
+ }
+ }
+
+ except Exception as e:
+ self.logger.error(f"Error getting instances data: {e}")
+ return {"error": f"Failed to get instances data: {str(e)}"}, 500
+
+ def update_config(self, config_updates: Dict[str, Any]) -> Dict[str, Any]:
+ """
+ Update system configuration.
+
+ Args:
+ config_updates: Dictionary of configuration updates
+
+ Returns:
+ Dict containing update result
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ # Validate and apply updates
+ updated_fields = []
+
+ for key, value in config_updates.items():
+ if key in ["max_annotations_per_user", "max_annotations_per_item"]:
+ if isinstance(value, int) and value >= -1:
+ config[key] = value
+ updated_fields.append(key)
+ else:
+ return {"error": f"Invalid value for {key}: must be integer >= -1"}, 400
+
+ elif key == "assignment_strategy":
+ valid_strategies = ["random", "fixed_order", "least_annotated", "max_diversity", "active_learning", "llm_confidence"]
+ if value in valid_strategies:
+ config[key] = value
+ updated_fields.append(key)
+ else:
+ return {"error": f"Invalid assignment strategy: {value}"}, 400
+
+ return {
+ "status": "success",
+ "message": f"Updated configuration fields: {', '.join(updated_fields)}",
+ "updated_fields": updated_fields
+ }
+
+ except Exception as e:
+ self.logger.error(f"Error updating config: {e}")
+ return {"error": f"Failed to update config: {str(e)}"}, 500
+
+ def get_questions_data(self) -> Dict[str, Any]:
+ """
+ Get aggregate analysis data for each annotation schema/question.
+
+ Returns:
+ Dict containing questions data with visualizations for different annotation types
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ ism = get_item_state_manager()
+ annotation_schemes = config.get("annotation_schemes", [])
+ questions_data = []
+
+ users = get_users()
+
+ for scheme in annotation_schemes:
+ scheme_name = scheme.get("name", "Unknown")
+ annotation_type = scheme.get("annotation_type", "unknown")
+
+ all_annotations = []
+ item_annotations = {}
+
+ for item in ism.items():
+ item_id = item.get_id()
+ item_annotations[item_id] = []
+
+ for username in users:
+ user_state = get_user_state_manager().get_user_state(username)
+ if user_state:
+ label_annotations = user_state.get_label_annotations(item_id)
+ for label, value in label_annotations.items():
+ label_schema = None
+ label_name = None
+ if hasattr(label, 'get_schema'):
+ label_schema = label.get_schema()
+ label_name = label.get_name()
+ elif hasattr(label, 'schema'):
+ label_schema = label.schema
+ label_name = getattr(label, 'name', None)
+ elif isinstance(label, str):
+ label_schema = label
+
+ if label_schema == scheme_name:
+ normalized_value = label_name if label_name else value
+
+ if annotation_type in ["radio", "select"]:
+ normalized_value = self._normalize_categorical_value(normalized_value)
+ elif annotation_type == "multiselect" and isinstance(normalized_value, list):
+ normalized_value = [
+ normalized_label
+ for normalized_label in (
+ self._normalize_categorical_value(v) for v in normalized_value
+ )
+ if normalized_label is not None
+ ]
+
+ if normalized_value is not None:
+ all_annotations.append(normalized_value)
+ item_annotations[item_id].append(normalized_value)
+
+ analysis = self._analyze_annotation_scheme(
+ annotation_type, scheme, all_annotations, item_annotations
+ )
+
+ questions_data.append({
+ "name": scheme_name,
+ "type": annotation_type,
+ "description": scheme.get("description", ""),
+ "total_annotations": len(all_annotations),
+ "items_with_annotations": len([item_id for item_id, annotations in item_annotations.items() if annotations]),
+ "analysis": analysis
+ })
+
+ return {
+ "questions": questions_data,
+ "summary": {
+ "total_questions": len(questions_data),
+ "total_annotations": sum(q["total_annotations"] for q in questions_data),
+ "question_types": list(set(q["type"] for q in questions_data))
+ }
+ }
+
+ except Exception as e:
+ self.logger.error(f"Error getting questions data: {e}")
+ return {"error": f"Failed to get questions data: {str(e)}"}, 500
+
+ def _analyze_annotation_scheme(self, annotation_type: str, scheme: dict,
+ all_annotations: list, item_annotations: dict) -> dict:
+ """
+ Analyze annotations based on their type and generate appropriate visualizations.
+ """
+ if not all_annotations:
+ return {"error": "No annotations found"}
+
+ analysis = {
+ "type": annotation_type,
+ "total_count": len(all_annotations)
+ }
+
+ if annotation_type in ["radio", "select"]:
+ normalized_annotations = [
+ normalized for normalized in
+ (self._normalize_categorical_value(annotation) for annotation in all_annotations)
+ if normalized is not None
+ ]
+ if not normalized_annotations:
+ return {"error": "No annotations found"}
+
+ label_counts = Counter(normalized_annotations)
+ raw_labels = scheme.get("labels", [])
+ labels = [
+ normalized for normalized in
+ (self._normalize_categorical_value(label) for label in raw_labels)
+ if normalized is not None
+ ]
+
+ analysis.update({
+ "visualization_type": "histogram",
+ "data": {
+ "labels": labels,
+ "counts": [label_counts.get(label, 0) for label in labels],
+ "percentages": [round(label_counts.get(label, 0) / len(normalized_annotations) * 100, 1)
+ for label in labels]
+ },
+ "most_common": label_counts.most_common(1)[0] if label_counts else None,
+ "agreement_score": self._calculate_agreement_score(item_annotations)
+ })
+ elif annotation_type == "multiselect":
+ # Multi-label data - show label frequency and co-occurrence
+ label_counts = Counter()
+ co_occurrence = defaultdict(int)
+ labels = scheme.get("labels", [])
+
+ for annotations in item_annotations.values():
+ if isinstance(annotations, list):
+ # Count individual labels
+ for annotation in annotations:
+ if isinstance(annotation, list):
+ for label in annotation:
+ label_counts[label] += 1
+
+ # Count co-occurrences
+ for i, annotation1 in enumerate(annotations):
+ if isinstance(annotation1, list):
+ for j, annotation2 in enumerate(annotations):
+ if i != j and isinstance(annotation2, list):
+ for label1 in annotation1:
+ for label2 in annotation2:
+ if label1 < label2:
+ co_occurrence[(label1, label2)] += 1
+
+ analysis.update({
+ "visualization_type": "multiselect_analysis",
+ "data": {
+ "labels": labels,
+ "counts": [label_counts.get(label, 0) for label in labels],
+ "percentages": [round(label_counts.get(label, 0) / len(item_annotations) * 100, 1)
+ for label in labels],
+ "co_occurrence": dict(co_occurrence)
+ },
+ "most_common": label_counts.most_common(3) if label_counts else [],
+ "average_labels_per_item": round(sum(len(ann) if isinstance(ann, list) else 1
+ for anns in item_annotations.values()
+ for ann in anns) / len(all_annotations), 2)
+ })
+
+ elif annotation_type in ["likert", "number", "slider"]:
+ # Numeric data - show distribution and statistics
+ numeric_values = []
+ for value in all_annotations:
+ try:
+ if isinstance(value, (int, float)):
+ numeric_values.append(float(value))
+ elif isinstance(value, str) and value.replace('.', '').replace('-', '').isdigit():
+ numeric_values.append(float(value))
+ except (ValueError, TypeError):
+ continue
+
+ if numeric_values:
+ analysis.update({
+ "visualization_type": "distribution",
+ "data": {
+ "values": numeric_values,
+ "bins": self._create_histogram_bins(numeric_values, scheme),
+ "statistics": {
+ "mean": round(sum(numeric_values) / len(numeric_values), 2),
+ "median": round(sorted(numeric_values)[len(numeric_values)//2], 2),
+ "min": min(numeric_values),
+ "max": max(numeric_values),
+ "std": round((sum((x - sum(numeric_values)/len(numeric_values))**2
+ for x in numeric_values) / len(numeric_values))**0.5, 2)
+ }
+ },
+ "range": scheme.get("min", 0) if "min" in scheme else None,
+ "max": scheme.get("max", 10) if "max" in scheme else None
+ })
+ else:
+ analysis["error"] = "No valid numeric values found"
+
+ elif annotation_type == "text":
+ # Text data - show length distribution and common patterns
+ text_lengths = []
+ word_counts = []
+ common_words = Counter()
+
+ for value in all_annotations:
+ if isinstance(value, str) and value.strip():
+ text_lengths.append(len(value))
+ words = value.lower().split()
+ word_counts.append(len(words))
+ common_words.update(words)
+
+ if text_lengths:
+ analysis.update({
+ "visualization_type": "text_analysis",
+ "data": {
+ "lengths": text_lengths,
+ "word_counts": word_counts,
+ "common_words": common_words.most_common(10),
+ "statistics": {
+ "avg_length": round(sum(text_lengths) / len(text_lengths), 1),
+ "avg_words": round(sum(word_counts) / len(word_counts), 1),
+ "min_length": min(text_lengths),
+ "max_length": max(text_lengths),
+ "empty_responses": len([v for v in all_annotations
+ if not isinstance(v, str) or not v.strip()])
+ }
+ }
+ })
+ else:
+ analysis["error"] = "No valid text responses found"
+
+ elif annotation_type == "span":
+ # Span data - show coverage and overlap statistics
+ span_counts = []
+ total_spans = 0
+
+ for annotations in item_annotations.values():
+ if isinstance(annotations, list):
+ for annotation in annotations:
+ if isinstance(annotation, list):
+ span_counts.append(len(annotation))
+ total_spans += len(annotation)
+
+ if span_counts:
+ analysis.update({
+ "visualization_type": "span_analysis",
+ "data": {
+ "span_counts": span_counts,
+ "total_spans": total_spans,
+ "statistics": {
+ "avg_spans_per_item": round(sum(span_counts) / len(span_counts), 2),
+ "items_with_spans": len([c for c in span_counts if c > 0]),
+ "max_spans": max(span_counts) if span_counts else 0,
+ "min_spans": min(span_counts) if span_counts else 0
+ }
+ }
+ })
+ else:
+ analysis["error"] = "No valid span annotations found"
+
+ else:
+ analysis["error"] = f"Unsupported annotation type: {annotation_type}"
+
+ return analysis
+
+ def _calculate_agreement_score(self, item_annotations: dict) -> float:
+ """Calculate agreement score for categorical annotations."""
+ if not item_annotations:
+ return 0.0
+
+ agreement_scores = []
+ for annotations in item_annotations.values():
+ if len(annotations) > 1:
+ # Calculate percentage of most common annotation
+ counter = Counter(annotations)
+ most_common_count = counter.most_common(1)[0][1]
+ agreement_scores.append(most_common_count / len(annotations))
+
+ return round(sum(agreement_scores) / len(agreement_scores) * 100, 1) if agreement_scores else 0.0
+
+ def _create_histogram_bins(self, values: list, scheme: dict) -> dict:
+ """Create histogram bins for numeric data."""
+ if not values:
+ return {"bins": [], "counts": []}
+
+ min_val = scheme.get("min", min(values))
+ max_val = scheme.get("max", max(values))
+
+ # Create 10 bins
+ bin_size = (max_val - min_val) / 10
+ bins = [min_val + i * bin_size for i in range(11)]
+ counts = [0] * 10
+
+ for value in values:
+ bin_index = min(int((value - min_val) / bin_size), 9)
+ counts[bin_index] += 1
+
+ return {
+ "bins": [round(b, 2) for b in bins],
+ "counts": counts
+ }
+
+ def _get_annotator_timing_data(self, user_id: str) -> Optional[AnnotatorTimingData]:
+ """
+ Get timing data for a specific annotator.
+
+ Args:
+ user_id: The user ID to get timing data for
+
+ Returns:
+ AnnotatorTimingData object or None if user not found
+ """
+ try:
+ usm = get_user_state_manager()
+ user_state = usm.get_user_state(user_id)
+
+ if not user_state:
+ return None
+
+ # Get basic user info
+ total_annotations = len(user_state.get_all_annotations())
+ phase = str(user_state.get_phase())
+ has_assignments = user_state.has_assignments()
+ remaining_assignments = user_state.has_remaining_assignments()
+
+ # Calculate timing data
+ total_seconds = 0
+ instance_times = []
+
+ for instance_id, behavioral_data in user_state.instance_id_to_behavioral_data.items():
+ instance_seconds = None
+ # Handle both BehavioralData objects and plain dicts
+ if hasattr(behavioral_data, 'total_time_ms'):
+ # BehavioralData object (loaded from JSON)
+ if behavioral_data.total_time_ms:
+ instance_seconds = behavioral_data.total_time_ms / 1000.0
+ elif isinstance(behavioral_data, dict):
+ # Plain dict (runtime data)
+ if behavioral_data.get("total_time_ms"):
+ instance_seconds = behavioral_data["total_time_ms"] / 1000.0
+ elif behavioral_data.get("time_string"):
+ parsed_time = user_state.parse_time_string(behavioral_data["time_string"])
+ if parsed_time:
+ instance_seconds = parsed_time["total_seconds"]
+ if instance_seconds is not None:
+ total_seconds += instance_seconds
+ instance_times.append(instance_seconds)
+
+ # Calculate averages
+ average_seconds_per_annotation = total_seconds / total_annotations if total_annotations > 0 else 0
+ annotations_per_hour = (total_annotations * 3600) / total_seconds if total_seconds > 0 else 0
+
+ # Get current instance time (if any)
+ current_instance_time = None
+ current_instance = user_state.get_current_instance()
+ if current_instance:
+ current_instance_id = current_instance.get_id()
+ current_behavioral = user_state.instance_id_to_behavioral_data.get(current_instance_id)
+ if current_behavioral:
+ if hasattr(current_behavioral, 'total_time_ms'):
+ if current_behavioral.total_time_ms:
+ current_instance_time = current_behavioral.total_time_ms / 1000.0
+ elif isinstance(current_behavioral, dict):
+ if current_behavioral.get("total_time_ms"):
+ current_instance_time = current_behavioral["total_time_ms"] / 1000.0
+ elif current_behavioral.get("time_string"):
+ parsed_current = user_state.parse_time_string(current_behavioral["time_string"])
+ if parsed_current:
+ current_instance_time = parsed_current["total_seconds"]
+
+ # Estimate last activity (for now, use current time - this could be enhanced)
+ last_activity = datetime.datetime.now()
+
+ # NEW: Get annotation history metrics
+ performance_metrics = user_state.get_performance_metrics()
+ suspicious_analysis = AnnotationHistoryManager.detect_suspicious_activity(
+ user_state.get_annotation_history()
+ )
+ recent_actions = user_state.get_recent_actions(5) # Last 5 minutes
+
+ # Calculate session duration
+ current_session_duration_minutes = None
+ if user_state.session_start_time:
+ duration = datetime.datetime.now() - user_state.session_start_time
+ current_session_duration_minutes = duration.total_seconds() / 60
+
+ # Get training statistics
+ training_state = user_state.get_training_state()
+ training_completed = training_state.is_passed() if training_state else False
+ training_correct_answers = training_state.get_correct_answer_count() if training_state else 0
+ training_total_attempts = training_state.get_total_attempts() if training_state else 0
+ training_pass_rate = (training_correct_answers / training_total_attempts * 100) if training_total_attempts > 0 else 0
+ training_current_question = training_state.get_current_question_index() if training_state else 0
+ training_total_questions = len(training_state.get_training_instances()) if training_state else 0
+
+ return AnnotatorTimingData(
+ user_id=user_id,
+ total_annotations=total_annotations,
+ total_seconds=total_seconds,
+ average_seconds_per_annotation=average_seconds_per_annotation,
+ last_activity=last_activity,
+ current_instance_time=current_instance_time,
+ annotations_per_hour=annotations_per_hour,
+ phase=phase,
+ has_assignments=has_assignments,
+ remaining_assignments=remaining_assignments,
+
+ # NEW: Annotation history metrics
+ total_actions=performance_metrics.get('total_actions', 0),
+ average_action_time_ms=performance_metrics.get('average_action_time_ms', 0.0),
+ fastest_action_time_ms=performance_metrics.get('fastest_action_time_ms', 0),
+ slowest_action_time_ms=performance_metrics.get('slowest_action_time_ms', 0),
+ actions_per_minute=performance_metrics.get('actions_per_minute', 0.0),
+ suspicious_score=suspicious_analysis.get('suspicious_score', 0.0),
+ suspicious_level=suspicious_analysis.get('suspicious_level', 'Normal'),
+ fast_actions_count=suspicious_analysis.get('fast_actions_count', 0),
+ burst_actions_count=suspicious_analysis.get('burst_actions_count', 0),
+ session_start_time=user_state.session_start_time,
+ current_session_duration_minutes=current_session_duration_minutes,
+ recent_actions_count=len(recent_actions),
+
+ # Training metrics
+ training_completed=training_completed,
+ training_correct_answers=training_correct_answers,
+ training_total_attempts=training_total_attempts,
+ training_pass_rate=training_pass_rate,
+ training_current_question=training_current_question,
+ training_total_questions=training_total_questions
+ )
+
+ except Exception as e:
+ self.logger.error(f"Error getting timing data for user {user_id}: {e}")
+ return None
+
+ def _extract_behavioral_total_seconds(self, behavioral_data: Any, user_state=None) -> Optional[float]:
+ """Extract total annotation time in seconds from behavioral data objects or legacy dicts."""
+ if not behavioral_data:
+ return None
+
+ if hasattr(behavioral_data, 'total_time_ms') and behavioral_data.total_time_ms is not None:
+ return behavioral_data.total_time_ms / 1000.0
+
+ if isinstance(behavioral_data, dict):
+ total_time_ms = behavioral_data.get("total_time_ms")
+ if total_time_ms is not None:
+ return total_time_ms / 1000.0
+
+ time_string = behavioral_data.get("time_string")
+ if time_string and user_state and hasattr(user_state, 'parse_time_string'):
+ parsed_time = user_state.parse_time_string(time_string)
+ if parsed_time:
+ return parsed_time.get("total_seconds")
+
+ return None
+
+ def _extract_behavioral_ai_count(self, behavioral_data: Any) -> int:
+ """Extract AI usage count from behavioral data objects or legacy dicts."""
+ if not behavioral_data:
+ return 0
+
+ if hasattr(behavioral_data, 'ai_usage'):
+ ai_usage = behavioral_data.ai_usage or []
+ return len(ai_usage)
+
+ if isinstance(behavioral_data, dict):
+ ai_usage = behavioral_data.get("ai_usage", []) or []
+ return len(ai_usage)
+
+ return 0
+
+ def _calculate_total_instance_ai(self, instance_id: str) -> int:
+ """
+ Calculate total AI assistance events for an instance across all users.
+
+ Args:
+ instance_id: The instance ID to analyze
+
+ Returns:
+ Total number of AI usage events recorded for the instance
+ """
+ try:
+ usm = get_user_state_manager()
+ users = get_users()
+
+ total_ai = 0
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if not user_state:
+ continue
+
+ behavioral_data = user_state.instance_id_to_behavioral_data.get(instance_id)
+ total_ai += self._extract_behavioral_ai_count(behavioral_data)
+
+ return total_ai
+
+ except Exception as e:
+ self.logger.error(f"Error calculating AI statistics for instance {instance_id}: {e}")
+ return 0
+
+ def _calculate_average_time_per_annotation(self, instance_id: str) -> Optional[float]:
+ """
+ Calculate average time per annotation for an instance.
+
+ Args:
+ instance_id: The instance ID to analyze
+
+ Returns:
+ Average time in seconds or None if no data
+ """
+ try:
+ usm = get_user_state_manager()
+ users = get_users()
+
+ total_time = 0
+ annotation_count = 0
+
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if user_state:
+ behavioral_data = user_state.instance_id_to_behavioral_data.get(instance_id)
+ total_seconds = self._extract_behavioral_total_seconds(behavioral_data, user_state)
+ if total_seconds is not None:
+ total_time += total_seconds
+ annotation_count += 1
+
+ return total_time / annotation_count if annotation_count > 0 else None
+
+ except Exception as e:
+ self.logger.error(f"Error calculating average time for instance {instance_id}: {e}")
+ return None
+
+ def _calculate_completion_percentage(self, user_id: str) -> float:
+ """
+ Calculate completion percentage for a user.
+
+ Args:
+ user_id: The user ID to calculate completion for
+
+ Returns:
+ Completion percentage (0-100)
+ """
+ try:
+ usm = get_user_state_manager()
+ user_state = usm.get_user_state(user_id)
+
+ if not user_state:
+ return 0.0
+
+ total_assignments = user_state.get_assigned_instance_count()
+ completed_assignments = len(user_state.get_all_annotations())
+
+ if total_assignments == 0:
+ return 0.0
+
+ return (completed_assignments / total_assignments) * 100
+
+ except Exception as e:
+ self.logger.error(f"Error calculating completion percentage for user {user_id}: {e}")
+ return 0.0
+
+ def _format_seconds(self, seconds: Optional[float]) -> Optional[str]:
+ """
+ Format seconds into a human-readable string.
+
+ Args:
+ seconds: Number of seconds to format
+
+ Returns:
+ Formatted time string or None if input is None
+ """
+ if seconds is None:
+ return None
+
+ if seconds < 60:
+ return f"{int(seconds)}s"
+ elif seconds < 3600:
+ minutes = int(seconds // 60)
+ remaining_seconds = int(seconds % 60)
+ return f"{minutes}m {remaining_seconds}s"
+ else:
+ hours = int(seconds // 3600)
+ remaining_minutes = int((seconds % 3600) // 60)
+ return f"{hours}h {remaining_minutes}m"
+
+ def _format_annotation_history(self, actions: List[AnnotationAction], context: str) -> Dict[str, Any]:
+ """
+ Format annotation history data for API response.
+
+ Args:
+ actions: List of annotation actions
+ context: Context string (user_id or "all_users")
+
+ Returns:
+ Formatted annotation history data
+ """
+ if not actions:
+ return {
+ "context": context,
+ "total_actions": 0,
+ "actions": [],
+ "summary": {
+ "action_types": {},
+ "time_distribution": {},
+ "performance_metrics": {}
+ }
+ }
+
+ # Calculate summary statistics
+ action_types = Counter(action.action_type for action in actions)
+ time_distribution = self._calculate_time_distribution(actions)
+ performance_metrics = AnnotationHistoryManager.calculate_performance_metrics(actions)
+
+ # Format actions for response
+ formatted_actions = []
+ for action in actions[-100:]: # Limit to 100 most recent
+ formatted_actions.append({
+ "action_id": action.action_id,
+ "timestamp": action.timestamp.isoformat(),
+ "user_id": action.user_id,
+ "instance_id": action.instance_id,
+ "action_type": action.action_type,
+ "schema_name": action.schema_name,
+ "label_name": action.label_name,
+ "old_value": action.old_value,
+ "new_value": action.new_value,
+ "span_data": action.span_data,
+ "session_id": action.session_id,
+ "client_timestamp": action.client_timestamp.isoformat() if action.client_timestamp else None,
+ "server_processing_time_ms": action.server_processing_time_ms,
+ "metadata": action.metadata
+ })
+
+ return {
+ "context": context,
+ "total_actions": len(actions),
+ "actions": formatted_actions,
+ "summary": {
+ "action_types": dict(action_types),
+ "time_distribution": time_distribution,
+ "performance_metrics": performance_metrics
+ }
+ }
+
+ def _calculate_time_distribution(self, actions: List[AnnotationAction]) -> Dict[str, int]:
+ """
+ Calculate time distribution of actions.
+
+ Args:
+ actions: List of annotation actions
+
+ Returns:
+ Dictionary with time distribution data
+ """
+ if not actions:
+ return {}
+
+ # Group by hour of day
+ hour_distribution = defaultdict(int)
+ for action in actions:
+ hour = action.timestamp.hour
+ hour_distribution[f"{hour:02d}:00"] += 1
+
+ return dict(hour_distribution)
+
+ def get_crowdsourcing_data(self) -> Dict[str, Any]:
+ """
+ Get crowdsourcing platform statistics (MTurk, Prolific).
+
+ This method analyzes user data to provide statistics about workers
+ from crowdsourcing platforms like Amazon Mechanical Turk and Prolific.
+
+ Returns:
+ Dict containing crowdsourcing statistics with the following structure:
+ - summary: Overall counts of crowdsourcing workers
+ - prolific: Prolific-specific statistics
+ - mturk: MTurk-specific statistics
+ - workers: List of individual worker data
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ from potato.authentication import UserAuthenticator
+
+ usm = get_user_state_manager()
+ users = get_users()
+
+ # Initialize counters
+ prolific_workers = []
+ mturk_workers = []
+ other_workers = []
+
+ # Track unique study/HIT IDs
+ prolific_study_ids = set()
+ mturk_hit_ids = set()
+
+ # Get user authenticator to access stored user data
+ try:
+ user_auth = UserAuthenticator.get_instance()
+ user_data_store = getattr(user_auth.auth_backend, 'user_data', {})
+ except (ValueError, AttributeError):
+ user_data_store = {}
+
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if not user_state:
+ continue
+
+ # Get timing data for the user
+ timing_data = self._get_annotator_timing_data(username)
+
+ # Get stored user data (from authentication)
+ stored_data = user_data_store.get(username, {})
+
+ # Determine platform based on stored data
+ prolific_session_id = stored_data.get('prolific_session_id')
+ prolific_study_id = stored_data.get('prolific_study_id')
+ mturk_assignment_id = stored_data.get('mturk_assignment_id')
+ mturk_hit_id = stored_data.get('mturk_hit_id')
+
+ worker_info = {
+ "worker_id": username,
+ "total_annotations": timing_data.total_annotations if timing_data else 0,
+ "phase": timing_data.phase if timing_data else "unknown",
+ "total_seconds": timing_data.total_seconds if timing_data else 0,
+ "annotations_per_hour": timing_data.annotations_per_hour if timing_data else 0,
+ "completion_percentage": self._calculate_completion_percentage(username),
+ "suspicious_level": timing_data.suspicious_level if timing_data else "Normal",
+ }
+
+ # Check for Prolific workers
+ if prolific_session_id or prolific_study_id or username.startswith('P'):
+ worker_info["platform"] = "prolific"
+ worker_info["session_id"] = prolific_session_id
+ worker_info["study_id"] = prolific_study_id
+ prolific_workers.append(worker_info)
+ if prolific_study_id:
+ prolific_study_ids.add(prolific_study_id)
+
+ # Check for MTurk workers
+ elif mturk_assignment_id or mturk_hit_id or username.startswith('A'):
+ worker_info["platform"] = "mturk"
+ worker_info["assignment_id"] = mturk_assignment_id
+ worker_info["hit_id"] = mturk_hit_id
+ mturk_workers.append(worker_info)
+ if mturk_hit_id:
+ mturk_hit_ids.add(mturk_hit_id)
+
+ else:
+ worker_info["platform"] = "other"
+ other_workers.append(worker_info)
+
+ # Calculate summary statistics
+ all_workers = prolific_workers + mturk_workers + other_workers
+
+ def calc_stats(workers):
+ if not workers:
+ return {
+ "count": 0,
+ "total_annotations": 0,
+ "total_time_seconds": 0,
+ "avg_annotations_per_worker": 0,
+ "avg_time_per_worker_minutes": 0,
+ "completed_count": 0,
+ "in_progress_count": 0,
+ }
+ total_annotations = sum(w["total_annotations"] for w in workers)
+ total_time = sum(w["total_seconds"] for w in workers)
+ completed = len([w for w in workers if w["phase"] == "Phase.DONE"])
+ in_progress = len([w for w in workers if w["phase"] == "Phase.ANNOTATION"])
+ return {
+ "count": len(workers),
+ "total_annotations": total_annotations,
+ "total_time_seconds": total_time,
+ "avg_annotations_per_worker": round(total_annotations / len(workers), 1) if workers else 0,
+ "avg_time_per_worker_minutes": round(total_time / len(workers) / 60, 1) if workers else 0,
+ "completed_count": completed,
+ "in_progress_count": in_progress,
+ }
+
+ return {
+ "summary": {
+ "total_workers": len(all_workers),
+ "prolific_workers": len(prolific_workers),
+ "mturk_workers": len(mturk_workers),
+ "other_workers": len(other_workers),
+ "prolific_studies": len(prolific_study_ids),
+ "mturk_hits": len(mturk_hit_ids),
+ },
+ "prolific": {
+ "stats": calc_stats(prolific_workers),
+ "study_ids": list(prolific_study_ids),
+ "workers": prolific_workers,
+ },
+ "mturk": {
+ "stats": calc_stats(mturk_workers),
+ "hit_ids": list(mturk_hit_ids),
+ "workers": mturk_workers,
+ },
+ "other": {
+ "stats": calc_stats(other_workers),
+ "workers": other_workers,
+ },
+ }
+
+ except Exception as e:
+ self.logger.error(f"Error getting crowdsourcing data: {e}")
+ return {"error": f"Failed to get crowdsourcing data: {str(e)}"}, 500
+
+
+ def get_agreement_metrics(self) -> Dict[str, Any]:
+ """
+ Get inter-annotator agreement metrics using Krippendorff's alpha.
+
+ This leverages the existing agreement.py module for calculations.
+
+ Returns:
+ Dict containing agreement metrics by schema and overall
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ import simpledorff
+ from simpledorff.metrics import nominal_metric, interval_metric
+ import pandas as pd
+
+ agreement_config = config.get("agreement_metrics", {})
+ min_overlap = agreement_config.get("min_overlap", 2)
+
+ ism = get_item_state_manager()
+ usm = get_user_state_manager()
+ annotation_schemes = config.get("annotation_schemes", [])
+ users = get_users()
+
+ metrics = {
+ "enabled": agreement_config.get("enabled", True),
+ "overall": {},
+ "by_schema": {},
+ "warnings": []
+ }
+
+ for scheme in annotation_schemes:
+ schema_name = scheme.get("name", "Unknown")
+ annotation_type = scheme.get("annotation_type", "unknown")
+
+ # Collect annotations per item for this schema
+ annotations_by_item = {}
+
+ for item in ism.items():
+ item_id = item.get_id()
+ item_annotations = []
+
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if not user_state:
+ continue
+
+ # Get annotations for this item
+ all_annotations = user_state.get_all_annotations()
+ if item_id not in all_annotations:
+ continue
+
+ instance_annotations = all_annotations[item_id]
+ labels = instance_annotations.get("labels", {})
+
+ # Find annotation for this schema
+ for label, value in labels.items():
+ label_schema = None
+ if hasattr(label, 'schema'):
+ label_schema = label.schema
+ elif hasattr(label, 'get_schema'):
+ label_schema = label.get_schema()
+
+ if label_schema == schema_name:
+ item_annotations.append({
+ "user": username,
+ "value": value
+ })
+
+ if item_annotations:
+ annotations_by_item[item_id] = item_annotations
+
+ # Filter items with minimum overlap
+ valid_items = {
+ item_id: annots
+ for item_id, annots in annotations_by_item.items()
+ if len(annots) >= min_overlap
+ }
+
+ if not valid_items:
+ metrics["by_schema"][schema_name] = {
+ "error": f"No items with {min_overlap}+ annotators",
+ "items_count": len(annotations_by_item)
+ }
+ continue
+
+ # Format for simpledorff
+ try:
+ reliability_data = []
+ for item_id, annots in valid_items.items():
+ for annot in annots:
+ reliability_data.append({
+ "unit": item_id,
+ "annotator": annot["user"],
+ "annotation": self._normalize_annotation_value(annot["value"])
+ })
+
+ df = pd.DataFrame(reliability_data)
+
+ # Choose metric based on annotation type
+ if annotation_type in ["likert", "slider", "number"]:
+ metric_fn = interval_metric
+ metric_name = "interval"
+ else:
+ metric_fn = nominal_metric
+ metric_name = "nominal"
+
+ # Calculate alpha
+ alpha = simpledorff.calculate_krippendorffs_alpha(
+ df,
+ experiment_col="unit",
+ annotator_col="annotator",
+ class_col="annotation",
+ metric_fn=metric_fn
+ )
+
+ schema_metrics = {
+ "krippendorff_alpha": round(alpha, 4),
+ "metric_type": metric_name,
+ "items_evaluated": len(valid_items),
+ "total_annotations": len(reliability_data),
+ "interpretation": self._interpret_alpha(alpha)
+ }
+
+ # Cohen's kappa (pairwise) and Fleiss' kappa apply to
+ # categorical schemas; skip for interval-metric data where
+ # Krippendorff alpha is the appropriate measure.
+ if metric_name == "nominal":
+ try:
+ from potato.agreement import (
+ cohen_kappa_pairwise, fleiss_kappa,
+ )
+ schema_metrics["cohen_kappa"] = cohen_kappa_pairwise(df)
+ schema_metrics["fleiss_kappa"] = fleiss_kappa(df)
+ except Exception as e:
+ self.logger.error(f"Error calculating kappas for {schema_name}: {e}")
+ schema_metrics["kappa_error"] = str(e)
+
+ metrics["by_schema"][schema_name] = schema_metrics
+
+ except Exception as e:
+ self.logger.error(f"Error calculating alpha for {schema_name}: {e}")
+ metrics["by_schema"][schema_name] = {
+ "error": str(e),
+ "items_count": len(valid_items)
+ }
+
+ # Calculate overall metrics
+ alphas = [
+ m["krippendorff_alpha"]
+ for m in metrics["by_schema"].values()
+ if "krippendorff_alpha" in m
+ ]
+ if alphas:
+ avg_alpha = sum(alphas) / len(alphas)
+ metrics["overall"] = {
+ "average_krippendorff_alpha": round(avg_alpha, 4),
+ "schemas_evaluated": len(alphas),
+ "interpretation": self._interpret_alpha(avg_alpha)
+ }
+
+ cohen_means = [
+ m["cohen_kappa"]["mean_kappa"]
+ for m in metrics["by_schema"].values()
+ if isinstance(m.get("cohen_kappa"), dict)
+ and m["cohen_kappa"].get("mean_kappa") is not None
+ ]
+ if cohen_means:
+ metrics["overall"]["average_cohen_kappa"] = round(
+ sum(cohen_means) / len(cohen_means), 4
+ )
+
+ fleiss_values = [
+ m["fleiss_kappa"]["kappa"]
+ for m in metrics["by_schema"].values()
+ if isinstance(m.get("fleiss_kappa"), dict)
+ and m["fleiss_kappa"].get("kappa") is not None
+ ]
+ if fleiss_values:
+ metrics["overall"]["average_fleiss_kappa"] = round(
+ sum(fleiss_values) / len(fleiss_values), 4
+ )
+
+ return metrics
+
+ except ImportError as e:
+ self.logger.error(f"simpledorff not installed: {e}")
+ return {
+ "enabled": False,
+ "error": "simpledorff library not installed. Run: pip install simpledorff"
+ }
+ except Exception as e:
+ self.logger.error(f"Error getting agreement metrics: {e}")
+ return {"error": f"Failed to get agreement metrics: {str(e)}"}, 500
+
+ def _interpret_alpha(self, alpha: float) -> str:
+ """Human-readable interpretation of Krippendorff's alpha."""
+ if alpha >= 0.8:
+ return "Good agreement"
+ elif alpha >= 0.67:
+ return "Tentative agreement"
+ elif alpha >= 0.33:
+ return "Low agreement"
+ else:
+ return "Poor agreement"
+
+ def _normalize_annotation_value(self, value: Any) -> Any:
+ """Normalize annotation value for comparison."""
+ if isinstance(value, list):
+ return tuple(sorted(str(v) for v in value))
+ elif isinstance(value, bool):
+ return str(value).lower()
+ return str(value)
+
+ def _normalize_categorical_value(self, value: Any) -> Optional[str]:
+ """Normalize categorical annotation values and label definitions into readable strings."""
+ if value is None:
+ return None
+
+ if isinstance(value, str):
+ return value
+
+ if isinstance(value, dict):
+ for key in ("name", "label", "value", "id", "text"):
+ candidate = value.get(key)
+ if isinstance(candidate, str) and candidate:
+ return candidate
+ return json.dumps(value, sort_keys=True)
+
+ if isinstance(value, (int, float, bool)):
+ return str(value)
+
+ return str(value)
+
+ def get_code_cooccurrence_matrix(self, schema_filter: Optional[str] = None,
+ min_count: int = 1) -> Dict[str, Any]:
+ """
+ Compute pairwise code co-occurrence across instances.
+
+ Two codes co-occur on an instance when at least one annotator applied
+ each to that instance. Pairs are de-duplicated within an instance
+ (multiple annotators applying the same pair count once).
+
+ Args:
+ schema_filter: If set, restrict to codes belonging to this schema.
+ min_count: Skip pairs with co-occurrence below this threshold.
+
+ Returns:
+ Dict with `codes` (sorted code list), `pairs`
+ ({code_a, code_b, count}), and `n_instances` for context.
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ ism = get_item_state_manager()
+ usm = get_user_state_manager()
+ users = get_users()
+
+ codes_per_instance: Dict[str, set] = {}
+ for item in ism.items():
+ instance_id = item.get_id()
+ codes: set = set()
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if not user_state:
+ continue
+ all_anns = user_state.get_all_annotations()
+ if instance_id not in all_anns:
+ continue
+ instance_anns = all_anns[instance_id]
+ labels = instance_anns.get("labels", {}) or {}
+ for label, value in labels.items():
+ schema_name = self._schema_for_label_key(label)
+ if schema_filter and schema_name != schema_filter:
+ continue
+ for code in self._labels_from_value(value):
+ codes.add(f"{schema_name}::{code}")
+ spans = instance_anns.get("spans", {}) or {}
+ for schema_name, span_list in spans.items():
+ if schema_filter and schema_name != schema_filter:
+ continue
+ for span in span_list or []:
+ code = span.get("label") or span.get("annotation")
+ if code:
+ codes.add(f"{schema_name}::{code}")
+ if codes:
+ codes_per_instance[instance_id] = codes
+
+ pair_counts: Dict[Tuple[str, str], int] = {}
+ for codes in codes_per_instance.values():
+ sorted_codes = sorted(codes)
+ for i in range(len(sorted_codes)):
+ for j in range(i + 1, len(sorted_codes)):
+ key = (sorted_codes[i], sorted_codes[j])
+ pair_counts[key] = pair_counts.get(key, 0) + 1
+
+ pairs = [
+ {"code_a": a, "code_b": b, "count": c}
+ for (a, b), c in pair_counts.items() if c >= min_count
+ ]
+ pairs.sort(key=lambda x: x["count"], reverse=True)
+
+ all_codes = sorted({
+ code for codes in codes_per_instance.values() for code in codes
+ })
+
+ return {
+ "codes": all_codes,
+ "pairs": pairs,
+ "n_instances": len(codes_per_instance),
+ "n_pairs": len(pairs),
+ "schema_filter": schema_filter,
+ "min_count": min_count,
+ }
+ except Exception as e:
+ self.logger.error(f"Error computing co-occurrence: {e}")
+ return {"error": f"Failed to compute co-occurrence: {str(e)}"}, 500
+
+ def get_code_crosstab(self, attribute_key: str,
+ schema_filter: Optional[str] = None) -> Dict[str, Any]:
+ """
+ Compute a codes-by-instance-attribute crosstab.
+
+ Each instance contributes one row to the count of (code, attribute_value);
+ multiple annotators applying the same code count once per instance.
+
+ Args:
+ attribute_key: Name of the item-metadata field to use as the column axis
+ (e.g. "site", "condition", "language").
+ schema_filter: If set, restrict to codes belonging to this schema.
+
+ Returns:
+ Dict with `codes` (row labels), `values` (column labels),
+ `cells` ({code, value, count}), and totals.
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+ if not attribute_key:
+ return {"error": "attribute_key is required"}, 400
+
+ try:
+ ism = get_item_state_manager()
+ usm = get_user_state_manager()
+ users = get_users()
+
+ cell_counts: Dict[Tuple[str, str], int] = {}
+ values_seen: set = set()
+ codes_seen: set = set()
+ n_instances_with_attr = 0
+
+ # When the attribute is not on the instance itself, fall
+ # back to the case-level attribute (cases group instances by
+ # participant/respondent; the attribute may live on the
+ # case). No-op for projects that don't use cases.
+ cb_task_dir = config.get("task_dir", ".")
+ cb_project = config.get("annotation_task_name") or "default"
+
+ for item in ism.items():
+ instance_id = item.get_id()
+ item_data = self._get_item_data(item)
+ attr_value = item_data.get(attribute_key)
+ if attr_value is None or attr_value == "":
+ try:
+ from potato.cases import attribute_for_instance
+ attr_value = attribute_for_instance(
+ cb_task_dir, cb_project, instance_id,
+ attribute_key)
+ except Exception:
+ attr_value = None
+ if attr_value is None or attr_value == "":
+ continue
+ attr_value = str(attr_value)
+ values_seen.add(attr_value)
+ n_instances_with_attr += 1
+
+ codes: set = set()
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if not user_state:
+ continue
+ all_anns = user_state.get_all_annotations()
+ if instance_id not in all_anns:
+ continue
+ instance_anns = all_anns[instance_id]
+ labels = instance_anns.get("labels", {}) or {}
+ for label, value in labels.items():
+ schema_name = self._schema_for_label_key(label)
+ if schema_filter and schema_name != schema_filter:
+ continue
+ for code in self._labels_from_value(value):
+ codes.add(f"{schema_name}::{code}")
+ spans = instance_anns.get("spans", {}) or {}
+ for schema_name, span_list in spans.items():
+ if schema_filter and schema_name != schema_filter:
+ continue
+ for span in span_list or []:
+ code = span.get("label") or span.get("annotation")
+ if code:
+ codes.add(f"{schema_name}::{code}")
+
+ for code in codes:
+ codes_seen.add(code)
+ key = (code, attr_value)
+ cell_counts[key] = cell_counts.get(key, 0) + 1
+
+ cells = [
+ {"code": code, "value": value, "count": count}
+ for (code, value), count in cell_counts.items()
+ ]
+ cells.sort(key=lambda x: (x["code"], x["value"]))
+
+ return {
+ "codes": sorted(codes_seen),
+ "values": sorted(values_seen),
+ "cells": cells,
+ "n_instances": n_instances_with_attr,
+ "attribute_key": attribute_key,
+ "schema_filter": schema_filter,
+ }
+ except Exception as e:
+ self.logger.error(f"Error computing crosstab: {e}")
+ return {"error": f"Failed to compute crosstab: {str(e)}"}, 500
+
+ @staticmethod
+ def _schema_for_label_key(label_key) -> str:
+ """Extract schema name from a label key in user_state annotations."""
+ if hasattr(label_key, "schema"):
+ return label_key.schema
+ if hasattr(label_key, "get_schema"):
+ return label_key.get_schema()
+ return str(label_key)
+
+ @staticmethod
+ def _labels_from_value(value) -> List[str]:
+ """Pull individual code names out of an annotation value blob."""
+ if value is None or value == "":
+ return []
+ if isinstance(value, dict):
+ return [k for k, v in value.items() if v]
+ if isinstance(value, list):
+ return [str(x) for x in value]
+ return [str(value)]
+
+ @staticmethod
+ def _get_item_data(item) -> dict:
+ """Return the raw data dict for an ItemStateManager item."""
+ for attr in ("data", "_data", "item_data"):
+ data = getattr(item, attr, None)
+ if isinstance(data, dict):
+ return data
+ if hasattr(item, "to_dict"):
+ try:
+ return item.to_dict()
+ except Exception:
+ pass
+ return {}
+
+ def get_quality_control_data(self) -> Dict[str, Any]:
+ """
+ Get quality control metrics (attention checks, gold standards, pre-annotation).
+
+ Returns:
+ Dict containing quality control metrics
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ qc_manager = get_quality_control_manager()
+
+ if not qc_manager:
+ return {
+ "enabled": False,
+ "message": "Quality control not configured"
+ }
+
+ metrics = qc_manager.get_quality_metrics()
+ return {
+ "enabled": True,
+ **metrics
+ }
+
+ except Exception as e:
+ self.logger.error(f"Error getting quality control data: {e}")
+ return {"error": f"Failed to get quality control data: {str(e)}"}, 500
+
+ def _behavioral_sequence(self, value: Any) -> list:
+ """Normalize behavioral list-like values to a safe list."""
+ if value is None:
+ return []
+ if isinstance(value, list):
+ return value
+ if isinstance(value, tuple):
+ return list(value)
+ return [value]
+
+ def _behavioral_field(self, payload: Any, field_name: str, default: Any = None) -> Any:
+ """Read a field from either a dict or an object used in behavioral analytics."""
+ if isinstance(payload, dict):
+ return payload.get(field_name, default)
+ return getattr(payload, field_name, default)
+
+ def get_behavioral_analytics_data(self) -> Dict[str, Any]:
+ """
+ Get comprehensive behavioral analytics data for all annotators.
+
+ Returns:
+ Dict containing behavioral analytics metrics including:
+ - Per-user statistics (time, interactions, AI usage)
+ - Aggregate statistics
+ - Quality indicators
+ - AI assistance analysis
+ """
+ if not self.check_admin_access():
+ return {"error": "Admin access required"}, 403
+
+ try:
+ usm = get_user_state_manager()
+ users = get_users()
+
+ user_stats = []
+ ai_usage_total = {'requests': 0, 'accepts': 0, 'rejects': 0, 'decision_times': []}
+ all_times = []
+ interaction_counts = Counter()
+ change_sources = Counter()
+ total_interactions = 0
+ total_changes = 0
+ total_ai_requests = 0
+ users_with_fast_annotations = 0
+ users_with_low_interaction = 0
+ users_with_no_changes = 0
+
+ for user_id in users:
+ user_state = usm.get_user_state(user_id)
+ if not user_state:
+ continue
+
+ behavioral_data = user_state.instance_id_to_behavioral_data
+ if not behavioral_data:
+ continue
+
+ user_times = []
+ user_interactions = 0
+ user_changes = 0
+ user_ai_requests = 0
+ user_ai_accepts = 0
+ user_fast_count = 0
+ user_low_interaction_count = 0
+ user_no_scroll_count = 0
+ user_no_change_count = 0
+
+ for instance_id, bd in behavioral_data.items():
+ time_ms = self._behavioral_field(bd, 'total_time_ms', 0) or 0
+ time_sec = time_ms / 1000
+ user_times.append(time_sec)
+ all_times.append(time_sec)
+
+ if time_sec < 5:
+ user_fast_count += 1
+
+ interactions = self._behavioral_sequence(self._behavioral_field(bd, 'interactions', []))
+ user_interactions += len(interactions)
+ total_interactions += len(interactions)
+ if len(interactions) < 3:
+ user_low_interaction_count += 1
+
+ for event in interactions:
+ event_type = self._behavioral_field(event, 'event_type', 'unknown')
+ interaction_counts[event_type] += 1
+
+ scroll = self._behavioral_field(bd, 'scroll_depth_max', 0) or 0
+ if scroll < 25:
+ user_no_scroll_count += 1
+
+ changes = self._behavioral_sequence(self._behavioral_field(bd, 'annotation_changes', []))
+ user_changes += len(changes)
+ total_changes += len(changes)
+ if len(changes) == 0:
+ user_no_change_count += 1
+
+ for change in changes:
+ source = self._behavioral_field(change, 'source', 'user')
+ change_sources[source] += 1
+
+ ai_events = self._behavioral_sequence(self._behavioral_field(bd, 'ai_usage', []))
+ for ai in ai_events:
+ user_ai_requests += 1
+ total_ai_requests += 1
+ ai_usage_total['requests'] += 1
+
+ accepted = self._behavioral_field(ai, 'suggestion_accepted', None)
+ if accepted:
+ user_ai_accepts += 1
+ ai_usage_total['accepts'] += 1
+ else:
+ ai_usage_total['rejects'] += 1
+
+ decision_time = self._behavioral_field(ai, 'time_to_decision_ms', None)
+ if isinstance(decision_time, (int, float)):
+ ai_usage_total['decision_times'].append(decision_time)
+
+ total_instances = len(behavioral_data)
+ if total_instances > 0:
+ fast_rate = user_fast_count / total_instances
+ low_interaction_rate = user_low_interaction_count / total_instances
+ no_scroll_rate = user_no_scroll_count / total_instances
+ no_change_rate = user_no_change_count / total_instances
+ suspicion_score = fast_rate * 0.3 + low_interaction_rate * 0.35 + no_scroll_rate * 0.2 + no_change_rate * 0.15
+
+ if user_fast_count > 0:
+ users_with_fast_annotations += 1
+ if user_low_interaction_count > 0:
+ users_with_low_interaction += 1
+ if user_no_change_count > 0:
+ users_with_no_changes += 1
+
+ user_stats.append({
+ 'user_id': user_id,
+ 'total_instances': total_instances,
+ 'total_time_sec': sum(user_times),
+ 'avg_time_sec': sum(user_times) / len(user_times) if user_times else 0,
+ 'min_time_sec': min(user_times) if user_times else 0,
+ 'max_time_sec': max(user_times) if user_times else 0,
+ 'total_interactions': user_interactions,
+ 'avg_interactions': user_interactions / total_instances,
+ 'total_changes': user_changes,
+ 'avg_changes': user_changes / total_instances,
+ 'ai_requests': user_ai_requests,
+ 'ai_accepts': user_ai_accepts,
+ 'ai_accept_rate': (user_ai_accepts / user_ai_requests) if user_ai_requests > 0 else None,
+ 'fast_annotation_rate': fast_rate,
+ 'low_interaction_rate': low_interaction_rate,
+ 'no_scroll_rate': no_scroll_rate,
+ 'no_change_rate': no_change_rate,
+ 'suspicion_score': suspicion_score,
+ 'quality_flag': 'SUSPICIOUS' if suspicion_score > 0.5 else 'WARNING' if suspicion_score > 0.3 else 'OK'
+ })
+
+ aggregate = {
+ 'total_users_with_data': len(user_stats),
+ 'total_instances': sum(u['total_instances'] for u in user_stats),
+ 'total_time_minutes': sum(u['total_time_sec'] for u in user_stats) / 60,
+ 'avg_time_per_instance': sum(all_times) / len(all_times) if all_times else 0,
+ 'median_time_per_instance': sorted(all_times)[len(all_times)//2] if all_times else 0,
+ }
+ aggregate_stats = {
+ 'total_users': len(user_stats),
+ 'total_instances': aggregate['total_instances'],
+ 'avg_time_per_instance_sec': aggregate['avg_time_per_instance'],
+ 'total_interactions': total_interactions,
+ 'total_changes': total_changes,
+ 'total_ai_requests': total_ai_requests,
+ }
+
+ ai_summary = {
+ 'total_requests': ai_usage_total['requests'],
+ 'total_accepts': ai_usage_total['accepts'],
+ 'total_rejects': ai_usage_total['rejects'],
+ 'accept_rate': (ai_usage_total['accepts'] / ai_usage_total['requests']) if ai_usage_total['requests'] > 0 else 0,
+ 'avg_decision_time_ms': sum(ai_usage_total['decision_times']) / len(ai_usage_total['decision_times']) if ai_usage_total['decision_times'] else 0
+ }
+
+ flagged_users = [u for u in user_stats if u['quality_flag'] == 'SUSPICIOUS']
+ warning_users = [u for u in user_stats if u['quality_flag'] == 'WARNING']
+ total_users_with_data = len(user_stats)
+ quality_summary = {
+ 'total_flagged': len(flagged_users),
+ 'total_warnings': len(warning_users),
+ 'flagged_user_ids': [u['user_id'] for u in flagged_users],
+ 'warning_user_ids': [u['user_id'] for u in warning_users],
+ 'high_suspicion_users': len(flagged_users),
+ 'fast_annotation_rate': (users_with_fast_annotations / total_users_with_data) if total_users_with_data > 0 else 0,
+ 'low_interaction_rate': (users_with_low_interaction / total_users_with_data) if total_users_with_data > 0 else 0,
+ 'no_change_rate': (users_with_no_changes / total_users_with_data) if total_users_with_data > 0 else 0,
+ }
+
+ return {
+ 'aggregate': aggregate,
+ 'aggregate_stats': aggregate_stats,
+ 'ai_usage': ai_summary,
+ 'quality_summary': quality_summary,
+ 'interaction_types': dict(interaction_counts.most_common(20)),
+ 'change_sources': dict(change_sources),
+ 'users': sorted(user_stats, key=lambda x: -x['suspicion_score'])
+ }
+
+ except Exception as e:
+ self.logger.error(f"Error getting behavioral analytics data: {e}")
+ import traceback
+ traceback.print_exc()
+ return {"error": f"Failed to get behavioral analytics data: {str(e)}"}, 500
+
+ def get_adjudication_overview(self) -> Dict[str, Any]:
+ """
+ Get an overview of adjudication status for the admin dashboard.
+
+ Returns:
+ Dict with queue stats, adjudicator stats, error taxonomy,
+ guideline flags, disagreement patterns, and similarity stats.
+ """
+ from potato.adjudication import get_adjudication_manager
+
+ adj_mgr = get_adjudication_manager()
+ if not adj_mgr or not adj_mgr.adj_config.enabled:
+ return {"enabled": False, "message": "Adjudication not configured"}
+
+ try:
+ # Queue stats
+ queue_stats = adj_mgr.get_stats()
+
+ # Error taxonomy frequency
+ error_counts = Counter()
+ guideline_flag_count = 0
+ for decision in adj_mgr.decisions.values():
+ for tag in decision.error_taxonomy:
+ error_counts[tag] += 1
+ if decision.guideline_update_flag:
+ guideline_flag_count += 1
+
+ # Per-adjudicator stats with avg time
+ adjudicator_details = {}
+ for adj_id, stats in queue_stats.get("adjudicator_stats", {}).items():
+ completed = stats.get("completed", 0)
+ total_time = stats.get("total_time_ms", 0)
+ adjudicator_details[adj_id] = {
+ "completed": completed,
+ "total_time_ms": total_time,
+ "avg_time_ms": (
+ round(total_time / completed) if completed > 0 else 0
+ ),
+ }
+
+ # Disagreement patterns
+ disagreement_patterns = self._analyze_disagreement_patterns(adj_mgr)
+
+ # Similarity engine stats
+ similarity_stats = {}
+ if adj_mgr.similarity_engine:
+ similarity_stats = adj_mgr.similarity_engine.get_stats()
+
+ return {
+ "enabled": True,
+ "queue_stats": queue_stats,
+ "adjudicator_details": adjudicator_details,
+ "error_taxonomy_counts": dict(error_counts.most_common()),
+ "guideline_flag_count": guideline_flag_count,
+ "disagreement_patterns": disagreement_patterns,
+ "similarity_stats": similarity_stats,
+ }
+
+ except Exception as e:
+ self.logger.error(f"Error getting adjudication overview: {e}")
+ return {"enabled": True, "error": str(e)}
+
+ def _analyze_disagreement_patterns(self, adj_mgr) -> List[Dict[str, Any]]:
+ """
+ Analyze per-schema disagreement patterns across the queue.
+
+ Returns:
+ List of dicts sorted by worst agreement first, with schema name
+ and average agreement score.
+ """
+ from collections import defaultdict
+
+ schema_scores = defaultdict(list)
+
+ for item in adj_mgr.queue.values():
+ for schema_name, score in item.agreement_scores.items():
+ schema_scores[schema_name].append(score)
+
+ patterns = []
+ for schema_name, scores in schema_scores.items():
+ avg = sum(scores) / len(scores) if scores else 1.0
+ patterns.append({
+ "schema": schema_name,
+ "avg_agreement": round(avg, 3),
+ "num_items": len(scores),
+ })
+
+ patterns.sort(key=lambda x: x["avg_agreement"])
+ return patterns
+
+ # ========================================================================
+ # MACE Competence Estimation
+ # ========================================================================
+
+ def get_mace_overview(self) -> Dict[str, Any]:
+ """Get MACE competence estimation overview for admin dashboard.
+
+ Returns:
+ Dict with competence scores, schema summaries, and config.
+ """
+ from potato.mace_manager import get_mace_manager
+
+ mace_mgr = get_mace_manager()
+ if not mace_mgr or not mace_mgr.mace_config.enabled:
+ return {"enabled": False, "message": "MACE not configured"}
+
+ return mace_mgr.get_results_summary()
+
+ def get_mace_predictions(
+ self, schema: str, instance_id: str = None
+ ) -> Dict[str, Any]:
+ """Get MACE predicted labels for a schema, optionally filtered by instance.
+
+ Args:
+ schema: Schema name to get predictions for.
+ instance_id: Optional specific instance to filter.
+
+ Returns:
+ Dict with predictions and entropy data.
+ """
+ from potato.mace_manager import get_mace_manager
+
+ mace_mgr = get_mace_manager()
+ if not mace_mgr or not mace_mgr.mace_config.enabled:
+ return {"error": "MACE not configured"}
+
+ return mace_mgr.get_predictions_for_schema(schema, instance_id)
+
+ def _calculate_label_statistics(self, instance_id: str) -> Tuple[Optional[str], float]:
+ """
+ Calculate most frequent label and disagreement for an instance.
+
+ Args:
+ instance_id: The instance ID to analyze
+
+ Returns:
+ Tuple of (most_frequent_label, disagreement_score)
+ """
+ try:
+ usm = get_user_state_manager()
+ users = get_users()
+
+ all_labels = []
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if user_state:
+ annotations = user_state.get_all_annotations()
+ if instance_id in annotations:
+ instance_annotations = annotations[instance_id]
+ if "labels" in instance_annotations:
+ for label, value in instance_annotations["labels"].items():
+ if hasattr(label, 'label_name'):
+ all_labels.append(label.label_name)
+ else:
+ all_labels.append(str(value))
+
+ if not all_labels:
+ return None, 0.0
+
+ label_counts = Counter(all_labels)
+ most_frequent_label = label_counts.most_common(1)[0][0]
+ total_annotations = len(all_labels)
+ most_frequent_count = label_counts[most_frequent_label]
+ disagreement = 1 - (most_frequent_count / total_annotations)
+
+ return most_frequent_label, disagreement
+
+ except Exception as e:
+ self.logger.error(f"Error calculating label statistics for instance {instance_id}: {e}")
+ return None, 0.0
+
+
+# Global instance
+admin_dashboard = AdminDashboard()
diff --git a/potato/agent_proxy/__init__.py b/potato/agent_proxy/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1cafb201aab49a6197ee08237c2102547ecfaf46
--- /dev/null
+++ b/potato/agent_proxy/__init__.py
@@ -0,0 +1,44 @@
+"""
+Agent Proxy Package
+
+Provides agent proxy implementations for live agent interaction during annotation.
+Proxies communicate with AI agent backends (echo, HTTP, OpenAI) and return
+responses to the annotation interface.
+
+Usage:
+ from potato.agent_proxy import AgentProxyFactory
+
+ proxy = AgentProxyFactory.create(config)
+ context = proxy.start_session("Book a flight to Paris")
+ response = proxy.send_message("Hello", context)
+"""
+
+from .base import AgentMessage, AgentResponse, BaseAgentProxy, AgentProxyFactory
+from .session import (
+ AgentSession,
+ AgentSessionManager,
+ init_agent_session_manager,
+ get_agent_session_manager,
+ clear_agent_session_manager,
+)
+from .sandbox import SafetySandbox, SandboxViolation
+
+# Import proxy implementations to trigger registration
+from . import echo_proxy
+from . import http_proxy
+from . import openai_proxy
+from . import coding_proxy # subprocess_coding + docker_coding
+
+__all__ = [
+ "AgentMessage",
+ "AgentResponse",
+ "BaseAgentProxy",
+ "AgentProxyFactory",
+ "AgentSession",
+ "AgentSessionManager",
+ "init_agent_session_manager",
+ "get_agent_session_manager",
+ "clear_agent_session_manager",
+ "SafetySandbox",
+ "SandboxViolation",
+]
diff --git a/potato/agent_proxy/base.py b/potato/agent_proxy/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..c03ebda9246d3499b4518681d9e05edb74d46754
--- /dev/null
+++ b/potato/agent_proxy/base.py
@@ -0,0 +1,138 @@
+"""
+Agent Proxy Base Module
+
+Provides the abstract base class and data structures for agent proxies,
+plus a factory registry for creating proxy instances from configuration.
+
+Agent proxies allow annotators to interact with AI agents live during
+annotation tasks. Each proxy type (echo, http, openai) handles
+communication with a specific kind of agent backend.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, Any, Optional, List
+import logging
+import time
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AgentMessage:
+ """A single message in an agent conversation."""
+ role: str # "user", "agent", "system", "error"
+ content: str
+ timestamp: float = field(default_factory=time.time)
+ metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class AgentResponse:
+ """Response from an agent proxy after sending a message."""
+ message: AgentMessage
+ done: bool = False
+ error: Optional[str] = None
+
+
+class BaseAgentProxy(ABC):
+ """
+ Abstract base class for agent proxies.
+
+ Subclasses implement communication with specific agent backends
+ (echo for testing, HTTP for generic REST APIs, OpenAI for chat completions).
+ """
+
+ proxy_type: str = ""
+
+ def __init__(self, config: dict):
+ self.config = config
+ self._initialize()
+
+ @abstractmethod
+ def _initialize(self):
+ """Set up connections, validate config. Called by __init__."""
+ pass
+
+ @abstractmethod
+ def start_session(self, task_description: str) -> dict:
+ """
+ Start a new interaction session.
+
+ Args:
+ task_description: The task the annotator should accomplish with the agent.
+
+ Returns:
+ Proxy-specific session context dict (stored in AgentSession.proxy_context).
+ """
+ pass
+
+ @abstractmethod
+ def send_message(self, message: str, session_context: dict) -> AgentResponse:
+ """
+ Send a message to the agent and get a blocking response.
+
+ Args:
+ message: The user's message text.
+ session_context: The proxy-specific context from start_session.
+
+ Returns:
+ AgentResponse with the agent's reply.
+ """
+ pass
+
+ def end_session(self, session_context: dict):
+ """
+ Clean up session resources. Override if needed.
+
+ Args:
+ session_context: The proxy-specific context from start_session.
+ """
+ pass
+
+
+class AgentProxyFactory:
+ """Factory registry for creating agent proxy instances."""
+
+ _proxies: Dict[str, type] = {}
+
+ @classmethod
+ def register(cls, proxy_type: str, proxy_class: type):
+ """Register a proxy type."""
+ cls._proxies[proxy_type] = proxy_class
+ logger.debug(f"Registered agent proxy type: {proxy_type}")
+
+ @classmethod
+ def create(cls, config: dict) -> BaseAgentProxy:
+ """
+ Create an agent proxy from configuration.
+
+ Args:
+ config: The full config dict. Reads from config["agent_proxy"].
+
+ Returns:
+ Configured BaseAgentProxy instance.
+
+ Raises:
+ ValueError: If proxy type is unknown or missing.
+ """
+ agent_config = config.get("agent_proxy", {})
+ proxy_type = agent_config.get("type")
+
+ if not proxy_type:
+ raise ValueError("agent_proxy.type is required")
+
+ if proxy_type not in cls._proxies:
+ supported = ", ".join(sorted(cls._proxies.keys()))
+ raise ValueError(
+ f"Unknown agent proxy type: '{proxy_type}'. "
+ f"Supported types: {supported}"
+ )
+
+ proxy_class = cls._proxies[proxy_type]
+ return proxy_class(agent_config)
+
+ @classmethod
+ def get_supported_types(cls) -> List[str]:
+ """Get list of registered proxy type names."""
+ return sorted(cls._proxies.keys())
diff --git a/potato/agent_proxy/coding_proxy.py b/potato/agent_proxy/coding_proxy.py
new file mode 100644
index 0000000000000000000000000000000000000000..883df590eea468aa585a6be0011736ee0115a17c
--- /dev/null
+++ b/potato/agent_proxy/coding_proxy.py
@@ -0,0 +1,466 @@
+"""
+Coding-agent proxies โ LLM plans + sandboxed code execution.
+
+Two implementations behind a shared base class:
+
+- :class:`SubprocessCodingAgentProxy` (default, ``type: subprocess_coding``)
+ runs each Python or shell action in a per-session temp workspace via
+ ``subprocess.run`` with a per-step timeout and an output cap. Light
+ isolation โ suitable for trusted-input research workflows. The
+ workspace is sandboxed (separate cwd) but **not** a security boundary;
+ malicious code can still touch the host filesystem outside ``cwd``.
+
+- :class:`DockerCodingAgentProxy` (``type: docker_coding``) runs each
+ action inside an ephemeral Docker container with ``--network=none``,
+ ``--memory``, ``--cpus``, ``--read-only`` and a writable workspace
+ bind-mounted at ``/work``. Real isolation โ survives untrusted code.
+ Requires the ``docker`` Python package and a running Docker daemon.
+
+Both inherit per-step / per-session / rate-limit enforcement from the
+existing :mod:`potato.agent_proxy.sandbox` framework via the standard
+``send_message`` flow in ``routes.py:agent_chat_send``.
+
+Configuration shape (both proxies):
+
+ agent_proxy:
+ type: subprocess_coding | docker_coding
+ llm:
+ endpoint_type: ollama
+ model: llama3.2:3b
+ base_url: http://localhost:11434
+ temperature: 0.2
+ max_tokens: 800
+ execution:
+ per_step_timeout: 8 # seconds
+ max_output_chars: 4000
+ starter_files: {} # {filename: contents} written into workspace
+ docker: # only for docker_coding
+ image: python:3.11-slim
+ memory: 512m
+ cpus: 1.0
+ network: none # "none" or "bridge"
+ sandbox: { max_steps: 20, ... } # standard agent-proxy sandbox knobs
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+import shutil
+import subprocess
+import tempfile
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+from .base import AgentMessage, AgentProxyFactory, AgentResponse, BaseAgentProxy
+
+logger = logging.getLogger(__name__)
+
+
+_PLANNER_SYSTEM_PROMPT = (
+ "You are an autonomous coding agent. The user will give you a coding task. "
+ "Each turn, decide a SINGLE next action and respond with ONLY a JSON object "
+ "of the form: {\"thought\": str, \"action\": {\"type\": str, \"code\": str}}. "
+ "The 'type' must be one of:\n"
+ " - \"python\": run the contents of 'code' as a python script in the workspace\n"
+ " - \"shell\": run 'code' as a bash command in the workspace\n"
+ " - \"finish\": stop and return your final answer in 'code' (which is then shown "
+ "to the user as your conclusion -- no execution happens)\n"
+ "Keep each action small and focused. Use 'finish' as soon as the task is done."
+)
+
+
+@dataclass
+class _ExecResult:
+ stdout: str
+ stderr: str
+ exit_code: Optional[int]
+ timed_out: bool = False
+ error: Optional[str] = None
+
+
+class CodingAgentProxy(BaseAgentProxy):
+ """Shared planner/executor scaffold for coding agents.
+
+ Subclasses implement :meth:`_execute` to run an action in their
+ chosen sandbox. Each call to :meth:`send_message`:
+
+ 1. Appends the user message to the running history.
+ 2. Asks an LLM (configurable endpoint) for a JSON ``{thought, action}``.
+ 3. Hands ``action`` to the subclass for execution.
+ 4. Returns a single reply combining thought + tool output.
+ """
+
+ def _initialize(self):
+ llm_cfg = self.config.get("llm") or {}
+ self.llm_endpoint_type = llm_cfg.get("endpoint_type", "ollama")
+ self.llm_model = llm_cfg.get("model")
+ self.llm_base_url = llm_cfg.get("base_url")
+ self.llm_temperature = llm_cfg.get("temperature", 0.2)
+ self.llm_max_tokens = llm_cfg.get("max_tokens", 800)
+ # OpenAI-compatible servers (vLLM etc.) ignore the key but the SDK
+ # requires a non-empty string. Ollama needs none. Without forwarding
+ # this the planner silently failed with "planner_unavailable".
+ self.llm_api_key = llm_cfg.get("api_key")
+ # Last endpoint init / call error, surfaced to the user instead of
+ # an opaque "planner unavailable" message.
+ self._llm_error: Optional[str] = None
+
+ execution_cfg = self.config.get("execution") or {}
+ self.per_step_timeout = execution_cfg.get("per_step_timeout", 8)
+ self.max_output_chars = execution_cfg.get("max_output_chars", 4000)
+ self.starter_files: Dict[str, str] = execution_cfg.get("starter_files", {}) or {}
+
+ self._llm = None # lazy
+
+ # ------------------------------------------------------------------
+ # LLM lazy-init
+ # ------------------------------------------------------------------
+
+ def _get_llm(self):
+ if self._llm is not None:
+ return self._llm
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ ai_cfg: Dict[str, Any] = {
+ "model": self.llm_model,
+ "max_tokens": self.llm_max_tokens,
+ "temperature": self.llm_temperature,
+ }
+ if self.llm_base_url:
+ ai_cfg["base_url"] = self.llm_base_url
+ # Forward the key for OpenAI-compatible endpoints; vLLM ignores
+ # its value but the OpenAI SDK rejects an empty one. Fall back to
+ # env then a non-empty placeholder so local servers just work.
+ ai_cfg["api_key"] = (
+ self.llm_api_key
+ or os.environ.get("OPENAI_API_KEY")
+ or os.environ.get("ANTHROPIC_API_KEY")
+ or "EMPTY"
+ )
+
+ self._llm = AIEndpointFactory.create_endpoint({
+ "ai_support": {
+ "enabled": True,
+ "endpoint_type": self.llm_endpoint_type,
+ "ai_config": ai_cfg,
+ }
+ })
+ self._llm_error = None
+ except Exception as e:
+ logger.warning("CodingAgentProxy: planner LLM init failed: %s", e)
+ self._llm_error = f"{type(e).__name__}: {e}"
+ self._llm = None
+ return self._llm
+
+ # ------------------------------------------------------------------
+ # Lifecycle
+ # ------------------------------------------------------------------
+
+ def start_session(self, task_description: str) -> dict:
+ workspace = tempfile.mkdtemp(prefix="potato_coding_agent_")
+ for filename, contents in self.starter_files.items():
+ target = os.path.join(workspace, filename)
+ os.makedirs(os.path.dirname(target) or workspace, exist_ok=True)
+ with open(target, "w") as f:
+ f.write(contents)
+ history = [
+ {"role": "system", "content": _PLANNER_SYSTEM_PROMPT},
+ {
+ "role": "system",
+ "content": f"Workspace: {workspace}\nTask: {task_description}",
+ },
+ ]
+ return {
+ "workspace": workspace,
+ "history": history,
+ "step": 0,
+ "finished": False,
+ }
+
+ def end_session(self, session_context: dict):
+ workspace = session_context.get("workspace") if session_context else None
+ if workspace and os.path.isdir(workspace):
+ shutil.rmtree(workspace, ignore_errors=True)
+
+ # ------------------------------------------------------------------
+ # Per-turn flow
+ # ------------------------------------------------------------------
+
+ def send_message(self, message: str, session_context: dict) -> AgentResponse:
+ if session_context.get("finished"):
+ return AgentResponse(
+ message=AgentMessage(role="agent", content="(session already finished)"),
+ done=True,
+ )
+
+ history: List[Dict[str, str]] = session_context.setdefault("history", [])
+ history.append({"role": "user", "content": message})
+ session_context["step"] = session_context.get("step", 0) + 1
+
+ plan = self._plan_next_action(history)
+ if plan is None:
+ detail = self._llm_error or "no response from planner LLM"
+ reply = (
+ f"Planner LLM unavailable ({self.llm_endpoint_type}): "
+ f"{detail}"
+ )
+ history.append({"role": "assistant", "content": reply})
+ session_context["finished"] = True
+ return AgentResponse(
+ message=AgentMessage(role="error", content=reply), error="planner_unavailable",
+ )
+
+ thought = plan.get("thought", "")
+ action = plan.get("action") or {}
+ atype = (action.get("type") or "").strip().lower()
+ code = action.get("code") or ""
+
+ if atype == "finish":
+ reply = self._format_finish_reply(thought, code)
+ history.append({"role": "assistant", "content": reply})
+ session_context["finished"] = True
+ return AgentResponse(
+ message=AgentMessage(role="agent", content=reply), done=True,
+ )
+
+ if atype not in ("python", "shell"):
+ reply = (
+ f"{thought}\n\n[invalid action type {atype!r}; try 'python', "
+ "'shell', or 'finish']"
+ )
+ history.append({"role": "assistant", "content": reply})
+ return AgentResponse(message=AgentMessage(role="agent", content=reply))
+
+ result = self._execute(atype, code, session_context)
+ reply = self._format_exec_reply(thought, atype, code, result)
+ history.append({"role": "assistant", "content": reply})
+ return AgentResponse(message=AgentMessage(role="agent", content=reply))
+
+ # ------------------------------------------------------------------
+ # Planning
+ # ------------------------------------------------------------------
+
+ def _plan_next_action(self, history: List[Dict[str, str]]) -> Optional[Dict[str, Any]]:
+ endpoint = self._get_llm()
+ if endpoint is None:
+ return None
+ try:
+ if hasattr(endpoint, "chat_query"):
+ raw = endpoint.chat_query(history)
+ else:
+ flat = "\n".join(f'{m["role"]}: {m["content"]}' for m in history)
+ raw = endpoint.query(flat + "\nassistant:", None)
+ except Exception as e:
+ logger.warning("Planner LLM call failed: %s", e)
+ self._llm_error = f"{type(e).__name__}: {e}"
+ return None
+
+ if isinstance(raw, dict):
+ return raw # already parsed JSON
+ text = str(raw or "").strip()
+ if not text:
+ return None
+ try:
+ return json.loads(text)
+ except json.JSONDecodeError:
+ match = re.search(r"\{.*\}", text, flags=re.DOTALL)
+ if match:
+ try:
+ return json.loads(match.group(0))
+ except json.JSONDecodeError:
+ pass
+ # Last-ditch: treat the whole text as a finish reply.
+ return {"thought": "", "action": {"type": "finish", "code": text[:500]}}
+
+ # ------------------------------------------------------------------
+ # Reply formatting
+ # ------------------------------------------------------------------
+
+ def _format_exec_reply(
+ self, thought: str, atype: str, code: str, result: _ExecResult
+ ) -> str:
+ parts: List[str] = []
+ if thought:
+ parts.append(thought.strip())
+ parts.append(f"```{atype}\n{code.strip()}\n```")
+ out_block: List[str] = []
+ if result.timed_out:
+ out_block.append(f"[timeout after {self.per_step_timeout}s]")
+ if result.error:
+ out_block.append(f"[error: {result.error}]")
+ if result.exit_code is not None:
+ out_block.append(f"[exit={result.exit_code}]")
+ if result.stdout:
+ out_block.append("stdout:\n" + self._truncate(result.stdout))
+ if result.stderr:
+ out_block.append("stderr:\n" + self._truncate(result.stderr))
+ if not out_block:
+ out_block.append("(no output)")
+ parts.append("\n".join(out_block))
+ return "\n\n".join(parts)
+
+ def _format_finish_reply(self, thought: str, final_text: str) -> str:
+ parts = []
+ if thought:
+ parts.append(thought.strip())
+ if final_text:
+ parts.append(final_text.strip())
+ if not parts:
+ parts.append("(done)")
+ return "\n\n".join(parts)
+
+ def _truncate(self, text: str) -> str:
+ if len(text) <= self.max_output_chars:
+ return text
+ cut = self.max_output_chars
+ return text[:cut] + f"\n[...truncated {len(text) - cut} chars]"
+
+ # ------------------------------------------------------------------
+ # Sandbox-specific execution
+ # ------------------------------------------------------------------
+
+ @abstractmethod
+ def _execute(
+ self, action_type: str, code: str, session_context: dict
+ ) -> _ExecResult:
+ """Execute ``code`` (one of 'python' / 'shell') in the sandbox."""
+
+
+class SubprocessCodingAgentProxy(CodingAgentProxy):
+ """Local subprocess-based execution.
+
+ NOT a security boundary -- the per-step timeout + tempdir cwd is the
+ only protection. Use ``DockerCodingAgentProxy`` for untrusted input.
+ """
+
+ proxy_type = "subprocess_coding"
+
+ def _execute(
+ self, action_type: str, code: str, session_context: dict
+ ) -> _ExecResult:
+ workspace = session_context["workspace"]
+ env = self._build_env()
+ try:
+ if action_type == "python":
+ script_path = os.path.join(workspace, "_action.py")
+ with open(script_path, "w") as f:
+ f.write(code)
+ proc = subprocess.run(
+ ["python", script_path],
+ cwd=workspace,
+ env=env,
+ capture_output=True,
+ text=True,
+ timeout=self.per_step_timeout,
+ )
+ else: # shell
+ proc = subprocess.run(
+ ["bash", "-c", code],
+ cwd=workspace,
+ env=env,
+ capture_output=True,
+ text=True,
+ timeout=self.per_step_timeout,
+ )
+ except subprocess.TimeoutExpired as e:
+ return _ExecResult(
+ stdout=(e.stdout or "") if isinstance(e.stdout, str) else "",
+ stderr=(e.stderr or "") if isinstance(e.stderr, str) else "",
+ exit_code=None,
+ timed_out=True,
+ )
+ except FileNotFoundError as e:
+ return _ExecResult(stdout="", stderr="", exit_code=None, error=str(e))
+ return _ExecResult(
+ stdout=proc.stdout or "",
+ stderr=proc.stderr or "",
+ exit_code=proc.returncode,
+ )
+
+ def _build_env(self) -> Dict[str, str]:
+ # Strip env down to a minimal set so subprocess code can't
+ # accidentally exfiltrate the host's secrets via env vars.
+ keep = {"PATH", "HOME", "LANG", "LC_ALL"}
+ return {k: v for k, v in os.environ.items() if k in keep}
+
+
+class DockerCodingAgentProxy(CodingAgentProxy):
+ """Ephemeral-container execution. Real isolation; requires Docker."""
+
+ proxy_type = "docker_coding"
+
+ def _initialize(self):
+ super()._initialize()
+ docker_cfg = self.config.get("docker") or {}
+ self.docker_image = docker_cfg.get("image", "python:3.11-slim")
+ self.docker_memory = docker_cfg.get("memory", "512m")
+ self.docker_cpus = str(docker_cfg.get("cpus", 1.0))
+ self.docker_network = docker_cfg.get("network", "none")
+ self._docker = None # lazy
+ # Sanity check: warn if docker CLI isn't on PATH
+ if shutil.which("docker") is None:
+ logger.warning(
+ "DockerCodingAgentProxy: 'docker' CLI not found on PATH. "
+ "Container execution will fail at runtime."
+ )
+
+ def _execute(
+ self, action_type: str, code: str, session_context: dict
+ ) -> _ExecResult:
+ workspace = session_context["workspace"]
+ # Materialise the code as a file in the workspace so the container
+ # can run it without inline injection through `-c`.
+ if action_type == "python":
+ target = os.path.join(workspace, "_action.py")
+ with open(target, "w") as f:
+ f.write(code)
+ container_cmd = ["python", "/work/_action.py"]
+ else: # shell
+ target = os.path.join(workspace, "_action.sh")
+ with open(target, "w") as f:
+ f.write(code)
+ os.chmod(target, 0o755)
+ container_cmd = ["bash", "/work/_action.sh"]
+
+ cmd = [
+ "docker", "run", "--rm",
+ f"--network={self.docker_network}",
+ f"--memory={self.docker_memory}",
+ f"--cpus={self.docker_cpus}",
+ "--read-only",
+ "--tmpfs", "/tmp:exec,size=64m",
+ "-v", f"{workspace}:/work",
+ "-w", "/work",
+ self.docker_image,
+ ] + container_cmd
+ try:
+ proc = subprocess.run(
+ cmd,
+ capture_output=True,
+ text=True,
+ timeout=self.per_step_timeout + 5, # docker pull/start overhead
+ )
+ except subprocess.TimeoutExpired as e:
+ return _ExecResult(
+ stdout=(e.stdout or "") if isinstance(e.stdout, str) else "",
+ stderr=(e.stderr or "") if isinstance(e.stderr, str) else "",
+ exit_code=None,
+ timed_out=True,
+ )
+ except FileNotFoundError as e:
+ return _ExecResult(stdout="", stderr="", exit_code=None, error=str(e))
+ return _ExecResult(
+ stdout=proc.stdout or "",
+ stderr=proc.stderr or "",
+ exit_code=proc.returncode,
+ )
+
+
+# Register both with the factory so configs can refer to them by name.
+AgentProxyFactory.register("subprocess_coding", SubprocessCodingAgentProxy)
+AgentProxyFactory.register("docker_coding", DockerCodingAgentProxy)
diff --git a/potato/agent_proxy/echo_proxy.py b/potato/agent_proxy/echo_proxy.py
new file mode 100644
index 0000000000000000000000000000000000000000..c6d44e4f3e229536090db500502f8fcee095f74f
--- /dev/null
+++ b/potato/agent_proxy/echo_proxy.py
@@ -0,0 +1,55 @@
+"""
+Echo Agent Proxy
+
+A testing/demo proxy that returns responses from a configurable list.
+Cycles through responses in order, wrapping around when exhausted.
+
+Configuration:
+ agent_proxy:
+ type: echo
+ responses:
+ - "I understand your request."
+ - "Working on it now."
+ - "Here's what I found."
+"""
+
+import logging
+
+from .base import BaseAgentProxy, AgentMessage, AgentResponse, AgentProxyFactory
+
+logger = logging.getLogger(__name__)
+
+
+class EchoProxy(BaseAgentProxy):
+ """Test proxy that returns canned responses in order."""
+
+ proxy_type = "echo"
+
+ def _initialize(self):
+ self.responses = self.config.get("responses", [
+ "I understand.",
+ "Working on it.",
+ "Done!",
+ ])
+
+ def start_session(self, task_description: str) -> dict:
+ return {"response_index": 0, "task_description": task_description}
+
+ def send_message(self, message: str, session_context: dict) -> AgentResponse:
+ idx = session_context.get("response_index", 0)
+ response_text = self.responses[idx % len(self.responses)]
+ session_context["response_index"] = idx + 1
+
+ return AgentResponse(
+ message=AgentMessage(
+ role="agent",
+ content=response_text,
+ )
+ )
+
+ def end_session(self, session_context: dict):
+ pass
+
+
+# Register with factory
+AgentProxyFactory.register("echo", EchoProxy)
diff --git a/potato/agent_proxy/http_proxy.py b/potato/agent_proxy/http_proxy.py
new file mode 100644
index 0000000000000000000000000000000000000000..59315a0e5b46b1ed438adc567365d44d75486441
--- /dev/null
+++ b/potato/agent_proxy/http_proxy.py
@@ -0,0 +1,108 @@
+"""
+Generic HTTP Agent Proxy
+
+POSTs to any REST endpoint with configurable field mapping.
+Supports sending full conversation history and custom headers.
+
+Configuration:
+ agent_proxy:
+ type: http
+ url: "http://localhost:8080/chat"
+ headers:
+ Authorization: "Bearer YOUR_KEY"
+ message_key: "message" # key in request body for user message
+ response_key: "response" # key in response JSON for agent reply
+ session_id_key: "session_id" # key in request/response for session tracking
+ send_history: false # whether to send full conversation history
+ history_key: "messages" # key for history array in request body
+"""
+
+import logging
+import uuid
+
+import requests
+
+from .base import BaseAgentProxy, AgentMessage, AgentResponse, AgentProxyFactory
+
+logger = logging.getLogger(__name__)
+
+
+class GenericHTTPProxy(BaseAgentProxy):
+ """Generic REST API proxy with configurable field mapping."""
+
+ proxy_type = "http"
+
+ def _initialize(self):
+ self.url = self.config.get("url")
+ if not self.url:
+ raise ValueError("http proxy requires 'url' in agent_proxy config")
+
+ self.headers = self.config.get("headers", {})
+ self.message_key = self.config.get("message_key", "message")
+ self.response_key = self.config.get("response_key", "response")
+ self.session_id_key = self.config.get("session_id_key", "session_id")
+ self.send_history = self.config.get("send_history", False)
+ self.history_key = self.config.get("history_key", "messages")
+ self.timeout = self.config.get("sandbox", {}).get(
+ "request_timeout_seconds", 60
+ )
+
+ def start_session(self, task_description: str) -> dict:
+ return {
+ "session_id": str(uuid.uuid4()),
+ "task_description": task_description,
+ "history": [],
+ }
+
+ def send_message(self, message: str, session_context: dict) -> AgentResponse:
+ payload = {
+ self.message_key: message,
+ self.session_id_key: session_context["session_id"],
+ }
+
+ if self.send_history:
+ payload[self.history_key] = session_context.get("history", [])
+
+ try:
+ resp = requests.post(
+ self.url,
+ json=payload,
+ headers=self.headers,
+ timeout=self.timeout,
+ )
+ resp.raise_for_status()
+ data = resp.json()
+
+ response_text = data.get(self.response_key, "")
+ if not response_text and isinstance(data, str):
+ response_text = data
+
+ # Update history
+ session_context.setdefault("history", []).append(
+ {"role": "user", "content": message}
+ )
+ session_context["history"].append(
+ {"role": "agent", "content": response_text}
+ )
+
+ return AgentResponse(
+ message=AgentMessage(role="agent", content=str(response_text))
+ )
+
+ except requests.Timeout:
+ return AgentResponse(
+ message=AgentMessage(role="error", content="Agent request timed out."),
+ error="timeout",
+ )
+ except requests.RequestException as e:
+ logger.error(f"HTTP proxy request failed: {e}")
+ return AgentResponse(
+ message=AgentMessage(
+ role="error", content=f"Agent communication error: {e}"
+ ),
+ error=str(e),
+ )
+
+
+# Register with factory
+AgentProxyFactory.register("http", GenericHTTPProxy)
diff --git a/potato/agent_proxy/openai_proxy.py b/potato/agent_proxy/openai_proxy.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae5a42136b79a7e01f26e7f6c009634e8d205261
--- /dev/null
+++ b/potato/agent_proxy/openai_proxy.py
@@ -0,0 +1,105 @@
+"""
+OpenAI Chat Completions Agent Proxy
+
+Uses the OpenAI SDK to communicate with chat completion models.
+Maintains conversation history in session context for multi-turn dialogue.
+
+Configuration:
+ agent_proxy:
+ type: openai
+ api_key: "${OPENAI_API_KEY}" # or set OPENAI_API_KEY env var
+ model: "gpt-4o"
+ system_prompt: "You are a helpful travel agent."
+ temperature: 0.7
+ max_tokens: 1024
+"""
+
+import logging
+import os
+
+from .base import BaseAgentProxy, AgentMessage, AgentResponse, AgentProxyFactory
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIChatProxy(BaseAgentProxy):
+ """OpenAI Chat Completions proxy."""
+
+ proxy_type = "openai"
+
+ def _initialize(self):
+ api_key = self.config.get("api_key", "")
+ # Support environment variable references like ${OPENAI_API_KEY}
+ if api_key.startswith("${") and api_key.endswith("}"):
+ env_var = api_key[2:-1]
+ api_key = os.environ.get(env_var, "")
+
+ if not api_key:
+ api_key = os.environ.get("OPENAI_API_KEY", "")
+
+ if not api_key:
+ raise ValueError(
+ "OpenAI proxy requires api_key in config or OPENAI_API_KEY env var"
+ )
+
+ try:
+ import openai
+ self.client = openai.OpenAI(api_key=api_key)
+ except ImportError:
+ raise ImportError(
+ "openai package is required for the OpenAI proxy. "
+ "Install with: pip install openai"
+ )
+
+ self.model = self.config.get("model", "gpt-4o")
+ self.system_prompt = self.config.get("system_prompt", "")
+ self.temperature = self.config.get("temperature", 0.7)
+ self.max_tokens = self.config.get("max_tokens", 1024)
+ self.timeout = self.config.get("sandbox", {}).get(
+ "request_timeout_seconds", 60
+ )
+
+ def start_session(self, task_description: str) -> dict:
+ messages = []
+ if self.system_prompt:
+ messages.append({"role": "system", "content": self.system_prompt})
+ # Include task description as system context
+ messages.append({
+ "role": "system",
+ "content": f"The user's task: {task_description}",
+ })
+ return {"messages": messages}
+
+ def send_message(self, message: str, session_context: dict) -> AgentResponse:
+ messages = session_context.get("messages", [])
+ messages.append({"role": "user", "content": message})
+
+ try:
+ response = self.client.chat.completions.create(
+ model=self.model,
+ messages=messages,
+ temperature=self.temperature,
+ max_tokens=self.max_tokens,
+ timeout=self.timeout,
+ )
+
+ content = response.choices[0].message.content or ""
+ messages.append({"role": "assistant", "content": content})
+ session_context["messages"] = messages
+
+ return AgentResponse(
+ message=AgentMessage(role="agent", content=content)
+ )
+
+ except Exception as e:
+ logger.error(f"OpenAI proxy error: {e}")
+ return AgentResponse(
+ message=AgentMessage(
+ role="error", content=f"Agent error: {e}"
+ ),
+ error=str(e),
+ )
+
+
+# Register with factory
+AgentProxyFactory.register("openai", OpenAIChatProxy)
diff --git a/potato/agent_proxy/sandbox.py b/potato/agent_proxy/sandbox.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae7b0f0763fe56d8044feb77020f1d49a8a4300d
--- /dev/null
+++ b/potato/agent_proxy/sandbox.py
@@ -0,0 +1,76 @@
+"""
+Agent Proxy Safety Sandbox
+
+Enforces limits on agent interactions: step counts, session timeouts,
+rate limits, and request timeouts. Prevents runaway or abusive sessions.
+"""
+
+import time
+import threading
+import logging
+from collections import defaultdict
+from typing import Dict, List
+
+logger = logging.getLogger(__name__)
+
+
+class SandboxViolation(Exception):
+ """Raised when a safety limit is exceeded."""
+ pass
+
+
+class SafetySandbox:
+ """Enforces safety limits on agent interactions."""
+
+ def __init__(self, config: dict):
+ sandbox_config = config.get("sandbox", {})
+ self.max_steps = sandbox_config.get("max_steps", 20)
+ self.max_session_seconds = sandbox_config.get("max_session_seconds", 600)
+ self.rate_limit_per_minute = sandbox_config.get("rate_limit_per_minute", 10)
+ self.request_timeout = sandbox_config.get("request_timeout_seconds", 60)
+
+ # Sliding window rate limit tracking: user_id -> list of timestamps
+ self._rate_windows: Dict[str, List[float]] = defaultdict(list)
+ self._lock = threading.Lock()
+
+ def check_step_limit(self, current_steps: int):
+ """Raise SandboxViolation if step limit reached."""
+ if current_steps >= self.max_steps:
+ raise SandboxViolation(
+ f"Step limit reached ({self.max_steps}). "
+ f"Please finish the conversation."
+ )
+
+ def check_session_timeout(self, session_start: float):
+ """Raise SandboxViolation if session has timed out."""
+ elapsed = time.time() - session_start
+ if elapsed > self.max_session_seconds:
+ raise SandboxViolation(
+ f"Session timeout ({self.max_session_seconds}s). "
+ f"Please finish the conversation."
+ )
+
+ def check_rate_limit(self, user_id: str):
+ """Raise SandboxViolation if user is sending too fast."""
+ now = time.time()
+ window_start = now - 60.0
+
+ with self._lock:
+ # Remove old entries outside the 1-minute window
+ timestamps = self._rate_windows[user_id]
+ self._rate_windows[user_id] = [
+ t for t in timestamps if t > window_start
+ ]
+
+ if len(self._rate_windows[user_id]) >= self.rate_limit_per_minute:
+ raise SandboxViolation(
+ f"Rate limit exceeded ({self.rate_limit_per_minute}/min). "
+ f"Please wait before sending another message."
+ )
+
+ # Record this request
+ self._rate_windows[user_id].append(now)
+
+ def get_request_timeout(self) -> float:
+ """Get the timeout in seconds for proxy HTTP requests."""
+ return self.request_timeout
diff --git a/potato/agent_proxy/session.py b/potato/agent_proxy/session.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d5e83a7d438ca3b29cb371ff1b49d6c40f1da2d
--- /dev/null
+++ b/potato/agent_proxy/session.py
@@ -0,0 +1,119 @@
+"""
+Agent Session Manager
+
+Thread-safe singleton that tracks active agent interaction sessions.
+Each session maps a (user_id, instance_id) pair to an AgentSession
+containing the proxy, conversation history, and step count.
+
+Follows the same singleton pattern as ItemStateManager and UserStateManager.
+"""
+
+import threading
+import time
+import logging
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple
+
+from .base import AgentMessage, BaseAgentProxy
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AgentSession:
+ """An active agent interaction session."""
+ user_id: str
+ instance_id: str
+ proxy: BaseAgentProxy
+ task_description: str
+ proxy_context: dict = field(default_factory=dict)
+ messages: List[AgentMessage] = field(default_factory=list)
+ step_count: int = 0
+ started_at: float = field(default_factory=time.time)
+ finished: bool = False
+
+
+class AgentSessionManager:
+ """Thread-safe manager for active agent sessions."""
+
+ def __init__(self, config: dict):
+ self.config = config
+ self._sessions: Dict[Tuple[str, str], AgentSession] = {}
+ self._lock = threading.RLock()
+
+ def create_session(
+ self,
+ user_id: str,
+ instance_id: str,
+ proxy: BaseAgentProxy,
+ task_description: str,
+ ) -> AgentSession:
+ """Create a new session for a user/instance pair."""
+ with self._lock:
+ key = (user_id, instance_id)
+ if key in self._sessions and not self._sessions[key].finished:
+ logger.warning(
+ f"Session already exists for {key}, returning existing"
+ )
+ return self._sessions[key]
+
+ proxy_context = proxy.start_session(task_description)
+ session = AgentSession(
+ user_id=user_id,
+ instance_id=instance_id,
+ proxy=proxy,
+ task_description=task_description,
+ proxy_context=proxy_context,
+ )
+ self._sessions[key] = session
+ logger.debug(f"Created agent session for {key}")
+ return session
+
+ def get_session(
+ self, user_id: str, instance_id: str
+ ) -> Optional[AgentSession]:
+ """Get an active session, or None if not found."""
+ with self._lock:
+ return self._sessions.get((user_id, instance_id))
+
+ def remove_session(self, user_id: str, instance_id: str):
+ """Remove a session and clean up proxy resources."""
+ with self._lock:
+ key = (user_id, instance_id)
+ session = self._sessions.pop(key, None)
+ if session:
+ try:
+ session.proxy.end_session(session.proxy_context)
+ except Exception as e:
+ logger.warning(f"Error ending proxy session for {key}: {e}")
+ logger.debug(f"Removed agent session for {key}")
+
+
+# Singleton management
+_AGENT_SESSION_MANAGER: Optional[AgentSessionManager] = None
+_AGENT_SESSION_MANAGER_LOCK = threading.Lock()
+
+
+def init_agent_session_manager(config: dict) -> AgentSessionManager:
+ """Initialize the singleton AgentSessionManager."""
+ global _AGENT_SESSION_MANAGER
+ if _AGENT_SESSION_MANAGER is None:
+ with _AGENT_SESSION_MANAGER_LOCK:
+ if _AGENT_SESSION_MANAGER is None:
+ _AGENT_SESSION_MANAGER = AgentSessionManager(config)
+ return _AGENT_SESSION_MANAGER
+
+
+def get_agent_session_manager() -> AgentSessionManager:
+ """Get the singleton AgentSessionManager."""
+ global _AGENT_SESSION_MANAGER
+ if _AGENT_SESSION_MANAGER is None:
+ raise ValueError("AgentSessionManager has not been initialized yet!")
+ return _AGENT_SESSION_MANAGER
+
+
+def clear_agent_session_manager():
+ """Clear the singleton instance (for testing)."""
+ global _AGENT_SESSION_MANAGER
+ with _AGENT_SESSION_MANAGER_LOCK:
+ _AGENT_SESSION_MANAGER = None
diff --git a/potato/agent_runner.py b/potato/agent_runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..760a6e85fbf8c878b85feb0051c5846101bb13f6
--- /dev/null
+++ b/potato/agent_runner.py
@@ -0,0 +1,1008 @@
+"""
+Live Agent Runner
+
+Manages an AI agent that browses the web via Playwright, controlled by an LLM.
+Annotators can observe, pause, instruct, or take over the agent in real time.
+
+The agent loop runs in a background thread with its own asyncio event loop.
+Communication with Flask routes happens through thread-safe state and queues.
+"""
+
+import asyncio
+import base64
+import json
+import logging
+import os
+import threading
+import time
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from queue import Queue, Empty
+from typing import Any, Callable, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class AgentState(Enum):
+ """States of the agent lifecycle."""
+ IDLE = "idle"
+ RUNNING = "running"
+ PAUSED = "paused"
+ TAKEOVER = "takeover"
+ COMPLETED = "completed"
+ ERROR = "error"
+
+
+@dataclass
+class AgentStep:
+ """A single step in the agent's execution."""
+ step_index: int
+ screenshot_path: str
+ action: Dict[str, Any]
+ thought: str
+ observation: str
+ timestamp: float
+ url: str = ""
+ viewport: Optional[Dict[str, int]] = None
+ coordinates: Optional[Dict[str, int]] = None
+ element: Optional[Dict[str, Any]] = None
+ annotator_instruction: Optional[str] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ d = {
+ "step_index": self.step_index,
+ "screenshot_url": self.screenshot_path,
+ "action_type": self.action.get("type", "unknown"),
+ "action": self.action,
+ "thought": self.thought,
+ "observation": self.observation,
+ "timestamp": self.timestamp,
+ "url": self.url,
+ }
+ if self.viewport:
+ d["viewport"] = self.viewport
+ if self.coordinates:
+ d["coordinates"] = self.coordinates
+ if self.element:
+ d["element"] = self.element
+ if self.annotator_instruction:
+ d["annotator_instruction"] = self.annotator_instruction
+ return d
+
+
+@dataclass
+class AgentConfig:
+ """Configuration for the agent runner."""
+ max_steps: int = 30
+ step_delay: float = 1.0
+ viewport_width: int = 1280
+ viewport_height: int = 720
+ system_prompt: str = ""
+ model: str = "claude-sonnet-4-20250514"
+ api_key: str = ""
+ max_tokens: int = 4096
+ temperature: float = 0.3
+ endpoint_type: str = "anthropic_vision"
+ history_window: int = 5 # Number of recent steps to include in LLM context
+ timeout: int = 60 # Per-request timeout in seconds
+
+ base_url: str = "" # For Ollama: server URL
+
+ @classmethod
+ def from_config(cls, config: Dict[str, Any]) -> "AgentConfig":
+ """Create AgentConfig from a live_agent YAML config dict."""
+ ai_config = config.get("ai_config", {})
+ viewport = config.get("viewport", {})
+ endpoint_type = config.get("endpoint_type", "anthropic_vision")
+
+ # API key: Ollama doesn't need one; OpenAI-compatible servers
+ # (e.g. vLLM) ignore it but the SDK requires a non-empty string.
+ if endpoint_type == "ollama_vision":
+ api_key = ai_config.get("api_key", "")
+ default_model = "gemma3:4b"
+ elif endpoint_type == "openai_vision":
+ api_key = ai_config.get("api_key", os.environ.get("OPENAI_API_KEY", "EMPTY"))
+ default_model = "" # must be set explicitly (e.g. served model id)
+ else:
+ api_key = ai_config.get("api_key", os.environ.get("ANTHROPIC_API_KEY", ""))
+ default_model = "claude-sonnet-4-20250514"
+
+ return cls(
+ max_steps=config.get("max_steps", 30),
+ step_delay=config.get("step_delay", 1.0),
+ viewport_width=viewport.get("width", 1280),
+ viewport_height=viewport.get("height", 720),
+ system_prompt=config.get("system_prompt", DEFAULT_SYSTEM_PROMPT),
+ model=ai_config.get("model", default_model),
+ api_key=api_key,
+ max_tokens=ai_config.get("max_tokens", 4096),
+ temperature=ai_config.get("temperature", 0.3),
+ endpoint_type=endpoint_type,
+ history_window=config.get("history_window", 5),
+ timeout=ai_config.get("timeout", 60),
+ base_url=ai_config.get("base_url", "http://localhost:11434"),
+ )
+
+
+DEFAULT_SYSTEM_PROMPT = """You are a web browsing agent. You can see screenshots of web pages and take actions to complete tasks.
+
+For each step, analyze the current screenshot and respond with a JSON object:
+{
+ "thought": "Your reasoning about what you see and what to do next",
+ "action": {
+ "type": "click|type|scroll|navigate|wait|done",
+ // For click: "x": 100, "y": 200
+ // For type: "text": "hello world"
+ // For scroll: "direction": "up|down", "amount": 300
+ // For navigate: "url": "https://..."
+ // For wait: (no extra fields)
+ // For done: "summary": "Task completed because..."
+ }
+}
+
+Always respond with valid JSON only. No markdown, no extra text."""
+
+
+class AgentRunner:
+ """
+ Runs an AI agent that browses the web via Playwright.
+
+ The agent loop:
+ 1. Takes a screenshot
+ 2. Sends it to the LLM with context/history
+ 3. Parses the LLM response for an action
+ 4. Executes the action via Playwright
+ 5. Emits events to all listeners (for SSE)
+ 6. Repeats until done, error, or max_steps
+
+ Thread-safe control methods allow pause/resume/instruct/takeover.
+ """
+
+ def __init__(self, session_id: str, config: AgentConfig, screenshot_dir: str):
+ self.session_id = session_id
+ self.config = config
+ self.screenshot_dir = screenshot_dir
+
+ # State
+ self._state = AgentState.IDLE
+ self._state_lock = threading.Lock()
+ self._steps: List[AgentStep] = []
+ self._error: Optional[str] = None
+
+ # Control
+ self._pause_event = threading.Event()
+ self._pause_event.set() # Not paused initially
+ self._stop_flag = threading.Event()
+ self._instruction_queue: Queue = Queue()
+ self._takeover_actions: Queue = Queue()
+
+ # Listeners for SSE
+ self._listeners: List[Callable] = []
+ self._listeners_lock = threading.Lock()
+
+ # Annotator interactions log
+ self._interactions: List[Dict[str, Any]] = []
+
+ # Playwright session (set during run)
+ self._playwright_session = None
+ self._llm_client = None
+
+ # Background thread
+ self._thread: Optional[threading.Thread] = None
+
+ @property
+ def state(self) -> AgentState:
+ with self._state_lock:
+ return self._state
+
+ @state.setter
+ def state(self, new_state: AgentState):
+ with self._state_lock:
+ old_state = self._state
+ self._state = new_state
+ self._emit_event("state_change", {
+ "old_state": old_state.value,
+ "new_state": new_state.value,
+ "timestamp": time.time(),
+ })
+
+ @property
+ def steps(self) -> List[AgentStep]:
+ return list(self._steps)
+
+ @property
+ def step_count(self) -> int:
+ return len(self._steps)
+
+ @property
+ def error(self) -> Optional[str]:
+ return self._error
+
+ # --- Control methods (thread-safe) ---
+
+ def pause(self):
+ """Pause the agent loop after the current step completes."""
+ if self.state == AgentState.RUNNING:
+ self._pause_event.clear()
+ self.state = AgentState.PAUSED
+ logger.info(f"[{self.session_id}] Agent paused")
+
+ def resume(self):
+ """Resume a paused agent."""
+ if self.state == AgentState.PAUSED:
+ self.state = AgentState.RUNNING
+ self._pause_event.set()
+ logger.info(f"[{self.session_id}] Agent resumed")
+
+ def inject_instruction(self, instruction: str):
+ """Send an instruction to the agent (processed at next step)."""
+ self._instruction_queue.put(instruction)
+ self._interactions.append({
+ "type": "instruction",
+ "text": instruction,
+ "timestamp": time.time(),
+ "step_index": self.step_count,
+ })
+ self._emit_event("instruction_received", {"instruction": instruction})
+ logger.info(f"[{self.session_id}] Instruction injected: {instruction[:100]}")
+
+ def enter_takeover(self):
+ """Switch to manual takeover mode."""
+ if self.state in (AgentState.RUNNING, AgentState.PAUSED):
+ self._pause_event.clear() # Pause the agent loop
+ self.state = AgentState.TAKEOVER
+ self._interactions.append({
+ "type": "takeover_start",
+ "timestamp": time.time(),
+ "step_index": self.step_count,
+ })
+ logger.info(f"[{self.session_id}] Takeover mode entered")
+
+ def exit_takeover(self):
+ """Exit manual takeover and resume the agent."""
+ if self.state == AgentState.TAKEOVER:
+ self._interactions.append({
+ "type": "takeover_end",
+ "timestamp": time.time(),
+ "step_index": self.step_count,
+ })
+ self.state = AgentState.RUNNING
+ self._pause_event.set()
+ logger.info(f"[{self.session_id}] Takeover mode exited")
+
+ def submit_manual_action(self, action: Dict[str, Any]):
+ """Submit a manual action during takeover mode."""
+ if self.state == AgentState.TAKEOVER:
+ self._takeover_actions.put(action)
+
+ def stop(self):
+ """Stop the agent loop."""
+ self._stop_flag.set()
+ self._pause_event.set() # Unblock if paused
+ logger.info(f"[{self.session_id}] Stop requested")
+
+ # --- Listener management ---
+
+ def add_listener(self, callback: Callable):
+ """Add an SSE listener callback."""
+ with self._listeners_lock:
+ self._listeners.append(callback)
+
+ def remove_listener(self, callback: Callable):
+ """Remove an SSE listener callback."""
+ with self._listeners_lock:
+ self._listeners = [l for l in self._listeners if l is not callback]
+
+ def _emit_event(self, event_type: str, data: Dict[str, Any]):
+ """Emit an event to all listeners."""
+ event = {"type": event_type, "data": data, "session_id": self.session_id}
+ with self._listeners_lock:
+ for listener in self._listeners:
+ try:
+ listener(event)
+ except Exception as e:
+ logger.warning(f"Listener error: {e}")
+
+ # --- Main agent loop ---
+
+ def start(self, task_description: str, start_url: str):
+ """Start the agent in a background thread."""
+ if self.state != AgentState.IDLE:
+ raise RuntimeError(f"Cannot start agent in state {self.state}")
+
+ self._thread = threading.Thread(
+ target=self._run_thread,
+ args=(task_description, start_url),
+ daemon=True,
+ name=f"agent-{self.session_id}",
+ )
+ self._thread.start()
+
+ def _run_thread(self, task_description: str, start_url: str):
+ """Thread target: runs the async agent loop."""
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ try:
+ loop.run_until_complete(self._run_async(task_description, start_url))
+ except Exception as e:
+ logger.error(f"[{self.session_id}] Agent thread error: {e}")
+ self._error = str(e)
+ self.state = AgentState.ERROR
+ self._emit_event("error", {"message": str(e)})
+ finally:
+ loop.close()
+
+ async def _run_async(self, task_description: str, start_url: str):
+ """Async agent loop."""
+ from potato.web_playwright import PlaywrightSession
+
+ self.state = AgentState.RUNNING
+
+ # Initialize Playwright
+ self._playwright_session = PlaywrightSession(
+ width=self.config.viewport_width,
+ height=self.config.viewport_height,
+ )
+ started = await self._playwright_session.start(start_url)
+ if not started:
+ raise RuntimeError("Failed to start Playwright browser session")
+
+ # Initialize LLM client
+ self._init_llm_client()
+
+ self._emit_event("started", {
+ "task": task_description,
+ "start_url": start_url,
+ "max_steps": self.config.max_steps,
+ })
+
+ try:
+ for step_index in range(self.config.max_steps):
+ # Check stop flag
+ if self._stop_flag.is_set():
+ logger.info(f"[{self.session_id}] Stopped by user")
+ break
+
+ # Wait if paused (blocks until resume/stop)
+ while not self._pause_event.is_set():
+ if self._stop_flag.is_set():
+ break
+ # Handle takeover actions while paused in takeover mode
+ if self.state == AgentState.TAKEOVER:
+ await self._process_takeover_actions()
+ await asyncio.sleep(0.1)
+
+ if self._stop_flag.is_set():
+ break
+
+ # Check for injected instructions
+ instruction = None
+ try:
+ instruction = self._instruction_queue.get_nowait()
+ except Empty:
+ pass
+
+ # Execute one agent step
+ step = await self._agent_step(
+ step_index, task_description, instruction
+ )
+ self._steps.append(step)
+
+ # Check if agent decided it's done
+ if step.action.get("type") == "done":
+ logger.info(f"[{self.session_id}] Agent completed task")
+ break
+
+ # Step delay
+ if self.config.step_delay > 0:
+ await asyncio.sleep(self.config.step_delay)
+
+ self.state = AgentState.COMPLETED
+ self._emit_event("complete", {
+ "total_steps": len(self._steps),
+ "final_url": (await self._playwright_session.get_state()).get("url", ""),
+ })
+
+ finally:
+ await self._playwright_session.stop()
+ self._playwright_session = None
+
+ async def _agent_step(
+ self,
+ step_index: int,
+ task_description: str,
+ instruction: Optional[str] = None,
+ ) -> AgentStep:
+ """Execute a single agent step: screenshot โ LLM โ action โ emit."""
+
+ # 1. Take screenshot
+ screenshot_bytes = await self._playwright_session.screenshot()
+ if not screenshot_bytes:
+ raise RuntimeError("Failed to capture screenshot")
+
+ screenshot_path = os.path.join(
+ self.screenshot_dir, f"step_{step_index:03d}.png"
+ )
+ os.makedirs(os.path.dirname(screenshot_path), exist_ok=True)
+ with open(screenshot_path, "wb") as f:
+ f.write(screenshot_bytes)
+
+ # 2. Get page state
+ page_state = await self._playwright_session.get_state()
+
+ # 3. Emit thinking event
+ self._emit_event("thinking", {
+ "step_index": step_index,
+ "screenshot_url": screenshot_path,
+ "url": page_state.get("url", ""),
+ })
+
+ # 4. Build messages and query LLM
+ screenshot_b64 = base64.b64encode(screenshot_bytes).decode("utf-8")
+ messages = self._build_llm_messages(
+ screenshot_b64, task_description, instruction
+ )
+ llm_response = self._query_llm(messages)
+
+ # 5. Parse action from response
+ thought, action = self._parse_action(llm_response)
+
+ # 6. Execute action
+ observation = await self._execute_action(action)
+
+ # 7. Build step
+ step = AgentStep(
+ step_index=step_index,
+ screenshot_path=screenshot_path,
+ action=action,
+ thought=thought,
+ observation=observation,
+ timestamp=time.time(),
+ url=page_state.get("url", ""),
+ viewport=page_state.get("viewport"),
+ coordinates=_extract_coordinates(action),
+ annotator_instruction=instruction,
+ )
+
+ # 8. Emit step event
+ self._emit_event("step", step.to_dict())
+
+ return step
+
+ def _build_llm_messages(
+ self,
+ screenshot_b64: str,
+ task_description: str,
+ instruction: Optional[str] = None,
+ ) -> List[Dict[str, Any]]:
+ """Build message list for the LLM vision API."""
+ messages = []
+
+ # System message
+ system_prompt = self.config.system_prompt or DEFAULT_SYSTEM_PROMPT
+ messages.append({"role": "system", "content": system_prompt})
+
+ # Task description
+ task_msg = f"Task: {task_description}"
+ if instruction:
+ task_msg += f"\n\nAnnotator instruction: {instruction}"
+
+ # Include recent step history
+ history_steps = self._steps[-self.config.history_window:]
+ if history_steps:
+ history_parts = []
+ for s in history_steps:
+ entry = f"Step {s.step_index}: thought='{s.thought}', action={json.dumps(s.action)}, observation='{s.observation}'"
+ history_parts.append(entry)
+ task_msg += "\n\nRecent history:\n" + "\n".join(history_parts)
+
+ messages.append({"role": "user", "content": task_msg})
+
+ # Current screenshot (as a separate user message with image)
+ messages.append({
+ "role": "user",
+ "content": [
+ {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": "image/png",
+ "data": screenshot_b64,
+ },
+ },
+ {
+ "type": "text",
+ "text": f"Current page screenshot (step {len(self._steps)}). What action should I take next?",
+ },
+ ],
+ })
+
+ return messages
+
+ def _init_llm_client(self):
+ """Initialize the LLM client based on endpoint_type."""
+ if self.config.endpoint_type == "anthropic_vision":
+ try:
+ import anthropic
+ except ImportError:
+ raise RuntimeError(
+ "anthropic package required. Install with: pip install anthropic"
+ )
+ api_key = self.config.api_key or os.environ.get("ANTHROPIC_API_KEY")
+ if not api_key:
+ raise RuntimeError(
+ "Anthropic API key required. Set in config or ANTHROPIC_API_KEY env var."
+ )
+ self._llm_client = anthropic.Anthropic(
+ api_key=api_key, timeout=self.config.timeout
+ )
+ elif self.config.endpoint_type == "ollama_vision":
+ try:
+ import ollama
+ except ImportError:
+ raise RuntimeError(
+ "ollama package required. Install with: pip install ollama"
+ )
+ host = self.config.base_url or "http://localhost:11434"
+ self._llm_client = ollama.Client(
+ host=host, timeout=self.config.timeout
+ )
+ # Verify connectivity
+ try:
+ self._llm_client.list()
+ logger.info(f"Connected to Ollama at {host}, model: {self.config.model}")
+ except Exception as e:
+ raise RuntimeError(f"Failed to connect to Ollama at {host}: {e}")
+ elif self.config.endpoint_type == "openai_vision":
+ try:
+ from openai import OpenAI
+ except ImportError:
+ raise RuntimeError(
+ "openai package required. Install with: pip install openai"
+ )
+ base_url = self.config.base_url or "https://api.openai.com/v1"
+ self._llm_client = OpenAI(
+ base_url=base_url,
+ api_key=self.config.api_key or "EMPTY",
+ timeout=self.config.timeout,
+ )
+ try:
+ self._llm_client.models.list()
+ logger.info(
+ f"Connected to OpenAI-compatible endpoint at {base_url}, "
+ f"model: {self.config.model}"
+ )
+ except Exception as e:
+ # Non-fatal: some servers gate /models; the chat call will
+ # surface a real error if the endpoint is truly unreachable.
+ logger.warning(
+ f"Could not list models at {base_url} ({e}); continuing."
+ )
+ else:
+ raise RuntimeError(
+ f"Unsupported endpoint_type: {self.config.endpoint_type}. "
+ f"Supported: 'anthropic_vision', 'ollama_vision', 'openai_vision'."
+ )
+
+ def _query_llm(self, messages: List[Dict[str, Any]]) -> str:
+ """Send messages to the LLM and return the text response."""
+ if self.config.endpoint_type == "anthropic_vision":
+ return self._query_anthropic(messages)
+ elif self.config.endpoint_type == "ollama_vision":
+ return self._query_ollama(messages)
+ elif self.config.endpoint_type == "openai_vision":
+ return self._query_openai(messages)
+ raise RuntimeError(f"Unsupported endpoint type: {self.config.endpoint_type}")
+
+ def _query_openai(self, messages: List[Dict[str, Any]]) -> str:
+ """Query an OpenAI-compatible vision endpoint (OpenAI, vLLM, etc.).
+
+ Converts the internal Anthropic-style message blocks into OpenAI
+ chat-completions format (image blocks become ``image_url`` data
+ URIs). Requests a JSON object response when the server supports it,
+ falling back gracefully if it does not.
+ """
+ oai_messages = []
+ for msg in messages:
+ role = msg["role"]
+ content = msg.get("content", "")
+ if isinstance(content, str):
+ oai_messages.append({"role": role, "content": content})
+ continue
+ parts = []
+ for block in content:
+ if not isinstance(block, dict):
+ continue
+ if block.get("type") == "text":
+ parts.append({"type": "text", "text": block.get("text", "")})
+ elif block.get("type") == "image":
+ src = block.get("source", {})
+ if src.get("type") == "base64":
+ media = src.get("media_type", "image/png")
+ parts.append({
+ "type": "image_url",
+ "image_url": {
+ "url": f"data:{media};base64,{src['data']}"
+ },
+ })
+ oai_messages.append({"role": role, "content": parts})
+
+ kwargs = {
+ "model": self.config.model,
+ "messages": oai_messages,
+ "max_tokens": self.config.max_tokens,
+ "temperature": self.config.temperature,
+ }
+
+ def _is_rate_limit(exc) -> bool:
+ if getattr(exc, "status_code", None) == 429:
+ return True
+ s = str(exc).lower()
+ return ("429" in s or "rate limit" in s or "quota" in s
+ or "resource_exhausted" in s)
+
+ def _create(use_rf: bool):
+ if use_rf:
+ return self._llm_client.chat.completions.create(
+ response_format={"type": "json_object"}, **kwargs)
+ return self._llm_client.chat.completions.create(**kwargs)
+
+ # Transient 429s (per-minute rate/token bursts) are common mid-run
+ # even on paid tiers; back off and retry instead of failing the
+ # whole agent session.
+ backoffs = [5, 15, 30, 30, 30]
+ use_rf = True
+ attempt = 0
+ while True:
+ try:
+ resp = _create(use_rf)
+ break
+ except Exception as e:
+ if _is_rate_limit(e):
+ if attempt >= len(backoffs):
+ raise
+ wait = backoffs[attempt]
+ attempt += 1
+ logger.warning(
+ f"[{self.session_id}] LLM 429/rate-limited; "
+ f"retry {attempt}/{len(backoffs)} in {wait}s"
+ )
+ self._emit_event("thinking", {
+ "text": f"Rate-limited by the model API; "
+ f"waiting {wait}s before retryingโฆ"
+ })
+ time.sleep(wait)
+ continue
+ if use_rf:
+ # Server may not support response_format; drop it once.
+ use_rf = False
+ continue
+ raise
+ return resp.choices[0].message.content or ""
+
+ def _query_anthropic(self, messages: List[Dict[str, Any]]) -> str:
+ """Query Anthropic Claude with vision support."""
+ # Separate system message
+ system = ""
+ api_messages = []
+ for msg in messages:
+ if msg["role"] == "system":
+ system = msg["content"]
+ else:
+ api_messages.append(msg)
+
+ kwargs = {
+ "model": self.config.model,
+ "max_tokens": self.config.max_tokens,
+ "temperature": self.config.temperature,
+ "messages": api_messages,
+ }
+ if system:
+ kwargs["system"] = system
+
+ response = self._llm_client.messages.create(**kwargs)
+ return response.content[0].text
+
+ def _query_ollama(self, messages: List[Dict[str, Any]]) -> str:
+ """Query Ollama vision model.
+
+ Converts Anthropic-format messages to Ollama format:
+ - System messages are prepended to the prompt text
+ - Multiple user messages are merged into a single message
+ - Content blocks with images use Ollama's 'images' key
+ """
+ # Extract text and images from Anthropic-format messages
+ all_text_parts = []
+ all_images = []
+ for msg in messages:
+ content = msg.get("content", "")
+ if msg["role"] == "system":
+ if isinstance(content, str) and content:
+ all_text_parts.insert(0, content)
+ continue
+ if isinstance(content, list):
+ for block in content:
+ if isinstance(block, dict):
+ if block.get("type") == "text":
+ all_text_parts.append(block["text"])
+ elif block.get("type") == "image":
+ source = block.get("source", {})
+ if source.get("type") == "base64":
+ all_images.append(source["data"])
+ elif isinstance(content, str) and content:
+ all_text_parts.append(content)
+
+ ollama_msg = {
+ "role": "user",
+ "content": "\n\n".join(all_text_parts),
+ }
+ if all_images:
+ ollama_msg["images"] = all_images
+
+ options = {
+ "temperature": self.config.temperature,
+ "num_predict": self.config.max_tokens,
+ }
+
+ # Use Ollama's format schema to force structured JSON output
+ agent_schema = {
+ "type": "object",
+ "properties": {
+ "thought": {"type": "string"},
+ "action": {
+ "type": "object",
+ "properties": {
+ "type": {"type": "string"},
+ "x": {"type": "integer"},
+ "y": {"type": "integer"},
+ "text": {"type": "string"},
+ "url": {"type": "string"},
+ "direction": {"type": "string"},
+ "amount": {"type": "integer"},
+ "summary": {"type": "string"},
+ },
+ "required": ["type"],
+ },
+ },
+ "required": ["thought", "action"],
+ }
+
+ response = self._llm_client.chat(
+ model=self.config.model,
+ messages=[ollama_msg],
+ options=options,
+ format=agent_schema,
+ )
+
+ # Extract content from response (handle both dict and Pydantic model)
+ message = (
+ response.get("message")
+ if hasattr(response, "get")
+ else getattr(response, "message", None)
+ )
+ if message is None:
+ raise RuntimeError("No message in Ollama response")
+
+ content = (
+ message.get("content")
+ if hasattr(message, "get")
+ else getattr(message, "content", None)
+ )
+
+ # Some models (e.g. qwen3-vl) put responses in 'thinking' field
+ # and leave content empty. Extract the agent JSON from thinking.
+ if not content:
+ thinking = (
+ message.get("thinking")
+ if hasattr(message, "get")
+ else getattr(message, "thinking", None)
+ )
+ if thinking:
+ content = _extract_agent_json(thinking)
+
+ return content or ""
+
+ def _parse_action(self, llm_response: str) -> tuple:
+ """Parse thought and action from LLM JSON response.
+
+ Returns:
+ (thought, action_dict)
+ """
+ # Try to extract JSON from response
+ text = llm_response.strip()
+
+ # Handle markdown code blocks
+ if "```json" in text:
+ import re
+ match = re.search(r"```json\s*([\s\S]*?)\s*```", text)
+ if match:
+ text = match.group(1).strip()
+ elif "```" in text:
+ import re
+ match = re.search(r"```\s*([\s\S]*?)\s*```", text)
+ if match:
+ text = match.group(1).strip()
+
+ try:
+ parsed = json.loads(text)
+ except json.JSONDecodeError:
+ logger.warning(f"Failed to parse LLM response as JSON: {text[:200]}")
+ return text, {"type": "wait"}
+
+ thought = parsed.get("thought", "")
+ action = parsed.get("action", {"type": "wait"})
+
+ # Validate action has a type
+ if "type" not in action:
+ action["type"] = "wait"
+
+ return thought, action
+
+ async def _execute_action(self, action: Dict[str, Any]) -> str:
+ """Execute an action via Playwright and return observation."""
+ action_type = action.get("type", "wait")
+ pw = self._playwright_session
+
+ try:
+ if action_type == "click":
+ x = int(action.get("x", 0))
+ y = int(action.get("y", 0))
+ success = await pw.click(x, y)
+ return f"Clicked at ({x}, {y})" if success else f"Click failed at ({x}, {y})"
+
+ elif action_type == "type":
+ text = action.get("text", "")
+ # Handle control characters via keyboard.press
+ if text == "\b":
+ success = await pw.page.keyboard.press("Backspace") or True
+ return "Pressed Backspace"
+ elif text == "\n":
+ success = await pw.page.keyboard.press("Enter") or True
+ return "Pressed Enter"
+ elif text == "\t":
+ success = await pw.page.keyboard.press("Tab") or True
+ return "Pressed Tab"
+ else:
+ success = await pw.type_text(text)
+ return f"Typed '{text}'" if success else f"Type failed: '{text}'"
+
+ elif action_type == "scroll":
+ direction = action.get("direction", "down")
+ amount = int(action.get("amount", 300))
+ dy = amount if direction == "down" else -amount
+ success = await pw.scroll(0, dy)
+ return f"Scrolled {direction} by {amount}px" if success else "Scroll failed"
+
+ elif action_type == "navigate":
+ url = action.get("url", "")
+ success = await pw.navigate(url)
+ return f"Navigated to {url}" if success else f"Navigation failed: {url}"
+
+ elif action_type == "wait":
+ await asyncio.sleep(1)
+ return "Waited 1 second"
+
+ elif action_type == "done":
+ summary = action.get("summary", "Task completed")
+ return summary
+
+ else:
+ logger.warning(f"Unknown action type: {action_type}")
+ return f"Unknown action: {action_type}"
+
+ except Exception as e:
+ logger.error(f"Action execution error: {e}")
+ return f"Error executing {action_type}: {e}"
+
+ async def _process_takeover_actions(self):
+ """Process manual actions submitted during takeover mode."""
+ try:
+ action = self._takeover_actions.get_nowait()
+ except Empty:
+ return
+
+ pw = self._playwright_session
+ if not pw:
+ return
+
+ observation = await self._execute_action(action)
+
+ # Take screenshot after manual action
+ screenshot_bytes = await pw.screenshot()
+ step_index = len(self._steps)
+ screenshot_path = os.path.join(
+ self.screenshot_dir, f"step_{step_index:03d}_manual.png"
+ )
+ if screenshot_bytes:
+ with open(screenshot_path, "wb") as f:
+ f.write(screenshot_bytes)
+
+ page_state = await pw.get_state()
+
+ step = AgentStep(
+ step_index=step_index,
+ screenshot_path=screenshot_path,
+ action={**action, "_manual": True},
+ thought="[Manual takeover action]",
+ observation=observation,
+ timestamp=time.time(),
+ url=page_state.get("url", ""),
+ viewport=page_state.get("viewport"),
+ coordinates=_extract_coordinates(action),
+ )
+ self._steps.append(step)
+ self._emit_event("step", step.to_dict())
+
+ # --- Trace export ---
+
+ def get_trace(self) -> Dict[str, Any]:
+ """Export the session as a web_agent_trace-compatible dict."""
+ return {
+ "steps": [s.to_dict() for s in self._steps],
+ "task_description": "", # Set by caller
+ "session_id": self.session_id,
+ "agent_config": {
+ "model": self.config.model,
+ "endpoint_type": self.config.endpoint_type,
+ "max_steps": self.config.max_steps,
+ },
+ "annotator_interactions": self._interactions,
+ "state": self.state.value,
+ "total_steps": len(self._steps),
+ }
+
+ def get_state_summary(self) -> Dict[str, Any]:
+ """Get a summary of current state for API responses."""
+ return {
+ "session_id": self.session_id,
+ "state": self.state.value,
+ "step_count": len(self._steps),
+ "error": self._error,
+ "has_instructions_pending": not self._instruction_queue.empty(),
+ }
+
+
+def _extract_agent_json(text: str) -> str:
+ """Extract the last valid JSON object containing 'thought' or 'action' from text.
+
+ Some models (qwen3-vl) put their chain-of-thought in the thinking field
+ with the actual JSON answer embedded in the text. This function finds
+ that JSON, skipping any example/template JSON from the prompt.
+ """
+ import re
+
+ # Find all JSON-like blocks (balanced braces)
+ candidates = []
+ depth = 0
+ start = None
+ for i, ch in enumerate(text):
+ if ch == "{":
+ if depth == 0:
+ start = i
+ depth += 1
+ elif ch == "}":
+ depth -= 1
+ if depth == 0 and start is not None:
+ candidates.append(text[start : i + 1])
+ start = None
+
+ # Try each candidate (last first โ most likely to be the final answer)
+ for candidate in reversed(candidates):
+ try:
+ parsed = json.loads(candidate)
+ if isinstance(parsed, dict) and ("thought" in parsed or "action" in parsed):
+ return candidate
+ except (json.JSONDecodeError, ValueError):
+ continue
+
+ # Fallback: try greedy regex for any JSON
+ match = re.search(r"\{[^{}]*\}", text)
+ return match.group(0) if match else ""
+
+
+def _extract_coordinates(action: Dict[str, Any]) -> Optional[Dict[str, int]]:
+ """Extract x, y coordinates from an action if present."""
+ if "x" in action and "y" in action:
+ return {"x": int(action["x"]), "y": int(action["y"])}
+ return None
diff --git a/potato/agent_runner_manager.py b/potato/agent_runner_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..b1903506951ab6c4ad2bf06166b96c9dfdad517e
--- /dev/null
+++ b/potato/agent_runner_manager.py
@@ -0,0 +1,226 @@
+"""
+Agent Runner Session Manager
+
+Singleton that manages active AgentRunner sessions.
+Keyed by "{user_id}:{instance_id}" for per-user, per-instance isolation.
+Includes TTL-based cleanup and max concurrent session limits.
+"""
+
+import atexit
+import logging
+import threading
+import time
+from typing import Dict, Optional
+
+from potato.agent_runner import AgentConfig, AgentRunner, AgentState
+
+logger = logging.getLogger(__name__)
+
+# Default limits
+DEFAULT_MAX_SESSIONS = 10
+DEFAULT_SESSION_TTL = 3600 # 1 hour
+
+
+class AgentRunnerManager:
+ """
+ Manages active AgentRunner sessions with lifecycle control.
+
+ Thread-safe singleton. Sessions are keyed by "{user_id}:{instance_id}".
+ """
+
+ _instance = None
+ _lock = threading.Lock()
+
+ def __init__(
+ self,
+ max_sessions: int = DEFAULT_MAX_SESSIONS,
+ session_ttl: int = DEFAULT_SESSION_TTL,
+ ):
+ self._sessions: Dict[str, AgentRunner] = {}
+ self._session_created: Dict[str, float] = {}
+ self._session_meta: Dict[str, Dict] = {}
+ self._lock = threading.Lock()
+ self.max_sessions = max_sessions
+ self.session_ttl = session_ttl
+
+ # Start cleanup thread
+ self._cleanup_stop = threading.Event()
+ self._cleanup_thread = threading.Thread(
+ target=self._cleanup_loop, daemon=True, name="agent-cleanup"
+ )
+ self._cleanup_thread.start()
+
+ @classmethod
+ def get_instance(cls, **kwargs) -> "AgentRunnerManager":
+ """Get or create the singleton instance."""
+ if cls._instance is None:
+ with cls._lock:
+ if cls._instance is None:
+ cls._instance = cls(**kwargs)
+ return cls._instance
+
+ @classmethod
+ def clear_instance(cls):
+ """Clear the singleton (for testing)."""
+ with cls._lock:
+ if cls._instance is not None:
+ cls._instance.shutdown()
+ cls._instance = None
+
+ def create_session(
+ self,
+ user_id: str,
+ instance_id: str,
+ config: AgentConfig,
+ screenshot_dir: str,
+ ) -> AgentRunner:
+ """
+ Create a new agent session.
+
+ Args:
+ user_id: Annotator user ID
+ instance_id: Annotation instance ID
+ config: Agent configuration
+ screenshot_dir: Directory to store screenshots
+
+ Returns:
+ AgentRunner instance
+
+ Raises:
+ RuntimeError: If max sessions reached or session already exists
+ """
+ session_key = f"{user_id}:{instance_id}"
+
+ with self._lock:
+ # Clean up expired sessions first
+ self._cleanup_expired_locked()
+
+ # Check for existing active session
+ if session_key in self._sessions:
+ existing = self._sessions[session_key]
+ if existing.state in (AgentState.RUNNING, AgentState.PAUSED, AgentState.TAKEOVER):
+ raise RuntimeError(
+ f"Active session already exists for {session_key}. "
+ f"Stop it first."
+ )
+ # Old completed/error session โ remove it
+ del self._sessions[session_key]
+ del self._session_created[session_key]
+ if session_key in self._session_meta:
+ del self._session_meta[session_key]
+
+ # Check capacity
+ active_count = sum(
+ 1
+ for s in self._sessions.values()
+ if s.state in (AgentState.RUNNING, AgentState.PAUSED, AgentState.TAKEOVER)
+ )
+ if active_count >= self.max_sessions:
+ raise RuntimeError(
+ f"Maximum concurrent sessions ({self.max_sessions}) reached"
+ )
+
+ import uuid
+ session_id = str(uuid.uuid4())[:12]
+ runner = AgentRunner(session_id, config, screenshot_dir)
+
+ self._sessions[session_key] = runner
+ self._session_created[session_key] = time.time()
+ self._session_meta[session_key] = {
+ "user_id": user_id,
+ "instance_id": instance_id,
+ "session_id": session_id,
+ }
+
+ logger.info(
+ f"Created agent session {session_id} for {session_key}"
+ )
+ return runner
+
+ def get_session(self, session_id: str) -> Optional[AgentRunner]:
+ """Get a session by its session_id."""
+ with self._lock:
+ for runner in self._sessions.values():
+ if runner.session_id == session_id:
+ return runner
+ return None
+
+ def get_session_by_key(self, user_id: str, instance_id: str) -> Optional[AgentRunner]:
+ """Get a session by user_id and instance_id."""
+ session_key = f"{user_id}:{instance_id}"
+ with self._lock:
+ return self._sessions.get(session_key)
+
+ def remove_session(self, session_id: str):
+ """Remove a session by session_id."""
+ with self._lock:
+ key_to_remove = None
+ for key, runner in self._sessions.items():
+ if runner.session_id == session_id:
+ key_to_remove = key
+ break
+ if key_to_remove:
+ runner = self._sessions.pop(key_to_remove)
+ self._session_created.pop(key_to_remove, None)
+ self._session_meta.pop(key_to_remove, None)
+ runner.stop()
+ logger.info(f"Removed agent session {session_id}")
+
+ def list_sessions(self) -> list:
+ """List all active sessions."""
+ with self._lock:
+ result = []
+ for key, runner in self._sessions.items():
+ meta = self._session_meta.get(key, {})
+ result.append({
+ "session_id": runner.session_id,
+ "user_id": meta.get("user_id"),
+ "instance_id": meta.get("instance_id"),
+ "state": runner.state.value,
+ "step_count": runner.step_count,
+ "created": self._session_created.get(key),
+ })
+ return result
+
+ def _cleanup_expired_locked(self):
+ """Remove expired sessions. Must be called with self._lock held."""
+ now = time.time()
+ expired_keys = []
+ for key, created_at in self._session_created.items():
+ if now - created_at > self.session_ttl:
+ runner = self._sessions.get(key)
+ if runner and runner.state in (AgentState.COMPLETED, AgentState.ERROR, AgentState.IDLE):
+ expired_keys.append(key)
+ elif runner and now - created_at > self.session_ttl * 2:
+ # Force-stop sessions that have been running too long
+ runner.stop()
+ expired_keys.append(key)
+
+ for key in expired_keys:
+ self._sessions.pop(key, None)
+ self._session_created.pop(key, None)
+ self._session_meta.pop(key, None)
+ logger.info(f"Cleaned up expired session: {key}")
+
+ def _cleanup_loop(self):
+ """Background cleanup thread."""
+ while not self._cleanup_stop.is_set():
+ self._cleanup_stop.wait(60) # Check every 60 seconds
+ if self._cleanup_stop.is_set():
+ break
+ with self._lock:
+ self._cleanup_expired_locked()
+
+ def shutdown(self):
+ """Stop all sessions and cleanup thread."""
+ self._cleanup_stop.set()
+ with self._lock:
+ for key, runner in self._sessions.items():
+ try:
+ runner.stop()
+ except Exception as e:
+ logger.warning(f"Error stopping session {key}: {e}")
+ self._sessions.clear()
+ self._session_created.clear()
+ self._session_meta.clear()
+ logger.info("AgentRunnerManager shut down")
diff --git a/potato/agreement.py b/potato/agreement.py
new file mode 100644
index 0000000000000000000000000000000000000000..8fcf7108dcebbf64caf6fc58954528e6155c8053
--- /dev/null
+++ b/potato/agreement.py
@@ -0,0 +1,278 @@
+"""
+Inter-Annotator Agreement Calculation Module
+
+This module provides functionality for calculating inter-annotator agreement metrics,
+including Krippendorff's alpha, Cohen's kappa (pairwise), and Fleiss' kappa
+(N raters), from annotation data. It supports both rating agreement (interval
+metric) and skip agreement (nominal metric) calculations.
+
+The module processes annotation files in JSON format and outputs agreement statistics
+along with a CSV file containing the processed annotation data.
+"""
+
+import argparse
+from itertools import combinations
+import simpledorff
+from simpledorff.metrics import *
+import ujson
+import pandas as pd
+
+from collections import defaultdict
+import numpy as np
+
+
+def get_nans(shape):
+ """
+ Create a numpy array filled with NaN values.
+
+ Args:
+ shape: The shape of the array to create
+
+ Returns:
+ numpy.ndarray: Array filled with NaN values
+ """
+ ar = np.empty(shape)
+ ar[:] = np.NaN
+ return ar
+
+
+def cohen_kappa_pairwise(reliability_df):
+ """
+ Compute Cohen's kappa for every pair of annotators and return aggregate stats.
+
+ Cohen's kappa is defined for exactly two raters. With N>2 raters we compute
+ kappa for each pair on the items they both rated, then return the mean and the
+ per-pair breakdown. Pairs that share fewer than 2 items are skipped.
+
+ Args:
+ reliability_df: long-format DataFrame with columns
+ unit (item id), annotator (user), annotation (label value).
+
+ Returns:
+ dict with keys: mean_kappa (float | None), pairs (list of
+ {annotator_a, annotator_b, kappa, n_items}), n_pairs_evaluated,
+ n_pairs_skipped.
+ """
+ from sklearn.metrics import cohen_kappa_score
+
+ annotators = sorted(reliability_df["annotator"].unique())
+ pairs = []
+ skipped = 0
+
+ for a, b in combinations(annotators, 2):
+ a_rows = reliability_df[reliability_df["annotator"] == a].set_index("unit")["annotation"]
+ b_rows = reliability_df[reliability_df["annotator"] == b].set_index("unit")["annotation"]
+ shared = a_rows.index.intersection(b_rows.index)
+ if len(shared) < 2:
+ skipped += 1
+ continue
+
+ y_a = a_rows.loc[shared].astype(str).tolist()
+ y_b = b_rows.loc[shared].astype(str).tolist()
+ try:
+ kappa = float(cohen_kappa_score(y_a, y_b))
+ except Exception:
+ skipped += 1
+ continue
+ pairs.append({
+ "annotator_a": a,
+ "annotator_b": b,
+ "kappa": round(kappa, 4),
+ "n_items": int(len(shared)),
+ })
+
+ mean_kappa = (sum(p["kappa"] for p in pairs) / len(pairs)) if pairs else None
+ return {
+ "mean_kappa": round(mean_kappa, 4) if mean_kappa is not None else None,
+ "pairs": pairs,
+ "n_pairs_evaluated": len(pairs),
+ "n_pairs_skipped": skipped,
+ }
+
+
+def fleiss_kappa(reliability_df):
+ """
+ Compute Fleiss' kappa for N raters over a categorical label set.
+
+ Fleiss' kappa assumes the same number of ratings per item but tolerates
+ different rater identities per item. Items with fewer than 2 ratings are
+ dropped; the remaining items are padded by repeating their available
+ ratings up to the per-item rater count (`n_raters = max ratings per item`).
+ When per-item rater counts vary widely the metric is approximate; we report
+ `n_raters` and `n_items_evaluated` so the caller can judge.
+
+ Args:
+ reliability_df: long-format DataFrame with columns
+ unit (item id), annotator (user), annotation (label value).
+
+ Returns:
+ dict with keys: kappa (float | None), n_items_evaluated (int),
+ n_raters (int), n_categories (int), interpretation (str).
+ """
+ if reliability_df.empty:
+ return {"kappa": None, "n_items_evaluated": 0, "n_raters": 0,
+ "n_categories": 0, "interpretation": "No data"}
+
+ df = reliability_df.copy()
+ df["annotation"] = df["annotation"].astype(str)
+
+ counts_by_item = df.groupby(["unit", "annotation"]).size().unstack(fill_value=0)
+ items_with_ratings = counts_by_item.sum(axis=1)
+ counts_by_item = counts_by_item.loc[items_with_ratings >= 2]
+
+ if counts_by_item.empty:
+ return {"kappa": None, "n_items_evaluated": 0, "n_raters": 0,
+ "n_categories": int(df["annotation"].nunique()),
+ "interpretation": "No items with >=2 raters"}
+
+ n_raters = int(counts_by_item.sum(axis=1).max())
+ n_items = int(counts_by_item.shape[0])
+ n_categories = int(counts_by_item.shape[1])
+
+ matrix = counts_by_item.to_numpy(dtype=float)
+ row_sums = matrix.sum(axis=1, keepdims=True)
+ row_sums[row_sums == 0] = 1.0
+ matrix = matrix * (n_raters / row_sums)
+
+ p_j = matrix.sum(axis=0) / (n_items * n_raters)
+ if n_raters < 2:
+ return {"kappa": None, "n_items_evaluated": n_items, "n_raters": n_raters,
+ "n_categories": n_categories,
+ "interpretation": "Need >=2 raters per item"}
+ p_i = (np.sum(matrix ** 2, axis=1) - n_raters) / (n_raters * (n_raters - 1))
+ p_bar = float(p_i.mean())
+ p_e = float(np.sum(p_j ** 2))
+
+ if p_e >= 1.0:
+ kappa = 1.0 if p_bar >= 1.0 else 0.0
+ else:
+ kappa = (p_bar - p_e) / (1 - p_e)
+
+ return {
+ "kappa": round(float(kappa), 4),
+ "n_items_evaluated": n_items,
+ "n_raters": n_raters,
+ "n_categories": n_categories,
+ "interpretation": interpret_kappa(kappa),
+ }
+
+
+def interpret_kappa(kappa):
+ """Landis & Koch (1977) interpretation bands for kappa-family metrics."""
+ if kappa is None:
+ return "No agreement computable"
+ if kappa < 0:
+ return "Worse than chance"
+ if kappa < 0.21:
+ return "Slight"
+ if kappa < 0.41:
+ return "Fair"
+ if kappa < 0.61:
+ return "Moderate"
+ if kappa < 0.81:
+ return "Substantial"
+ return "Almost perfect"
+
+
+def flatten(annotations):
+ """
+ Flatten annotation data structure for processing.
+
+ Converts a list of annotation dictionaries into a format where each
+ annotation is a dictionary mapping user IDs to their labels.
+
+ Args:
+ annotations: List of annotation dictionaries
+
+ Returns:
+ list: Flattened annotation data structure
+
+ Example:
+ Input: [{"user": "user1", "label": "positive"}, {"user": "user2", "label": "negative"}]
+ Output: [{"user1": "positive", "user2": "negative"}]
+ """
+ return [{a["user"]: a["label"] for a in ann} for ann in annotations]
+
+
+def main(args):
+ """
+ Main function for calculating inter-annotator agreement.
+
+ This function processes annotation data from a JSON file, calculates
+ Krippendorff's alpha for both rating agreement and skip agreement,
+ and outputs the results along with a CSV file of the processed data.
+
+ Args:
+ args: Command line arguments containing file paths
+
+ Side Effects:
+ - Reads annotation data from input file
+ - Prints agreement statistics to console
+ - Writes processed data to output CSV file
+
+ The function processes the first 385 annotations by default and handles
+ missing annotations and skipped items appropriately.
+ """
+ # Load annotation data from JSON file
+ with open(args.file, "r") as f:
+ annotations = [ujson.loads(line)["annotations"] for line in f]
+
+ # Extract unique user IDs from all annotations
+ users = set([a["user"] for ann in annotations for a in ann])
+ annotations = flatten(annotations)
+
+ # Limit to first 385 annotations (configurable limit)
+ annotations = annotations[:385]
+
+ # Create data matrix for agreement calculation
+ # Each row represents a user, each column represents an annotation
+ # -1 values indicate skipped annotations, NaN indicates missing annotations
+ data = [
+ [np.nan if user not in a or int(a[user]) == -1 else int(a[user]) for a in annotations]
+ for user in users
+ ]
+
+ # Create skip data matrix (boolean indicating if annotation was skipped)
+ skip_data = [
+ [np.nan if user not in a else int(a[user]) < 0 for a in annotations] for user in users
+ ]
+
+ # Calculate statistics for each user
+ labeled = ~np.isnan(data)
+ skipped = [
+ [False if user not in a else int(a[user]) < 0 for a in annotations] for user in users
+ ]
+
+ # Print summary statistics
+ print("calculating over:")
+ for user, skip in zip(labeled, skipped):
+ print("labeled:", sum(user))
+ print("skipped:", sum(skip))
+
+ # Count instances where all users provided annotations
+ print(np.all(labeled, axis=0).sum())
+
+ # Calculate and print Krippendorff's alpha for rating agreement
+ # Uses interval metric for continuous rating scales
+ print("rating agreement:")
+ print(simpledorff.calculate_krippendorffs_alpha(pd.DataFrame(data),metric_fn=interval_metric))
+
+ # Calculate and print Krippendorff's alpha for skip agreement
+ # Uses nominal metric for binary skip/no-skip decisions
+ print("skip agreement:")
+ print(simpledorff.calculate_krippendorffs_alpha(pd.DataFrame(data),metric_fn=nominal_metric))
+
+ # Write processed data to CSV file
+ with open(args.outfile, "w") as f:
+ for row in zip(*data):
+ f.write(",".join([str(a) for a in row]) + "\n")
+
+
+if __name__ == "__main__":
+ # Set up command line argument parsing
+ parser = argparse.ArgumentParser(
+ description="Calculate Krippendorf's alpha from given JSON file of annotations"
+ )
+ parser.add_argument("file", help="path to JSON file")
+ parser.add_argument("outfile", help="write path to CSV")
+ main(parser.parse_args())
diff --git a/potato/ai/__init__.py b/potato/ai/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f15af8c8a1ce1da66103652416ec60c2924fb47
--- /dev/null
+++ b/potato/ai/__init__.py
@@ -0,0 +1 @@
+from .ai_help_wrapper import generate_ai_help_html
\ No newline at end of file
diff --git a/potato/ai/ai_cache.py b/potato/ai/ai_cache.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c3222feb43d1b4755d8dfd89626480364705f6f
--- /dev/null
+++ b/potato/ai/ai_cache.py
@@ -0,0 +1,1473 @@
+from __future__ import annotations
+import json
+import logging
+import os
+from typing import Any, Dict, Union
+import requests
+from tqdm import tqdm
+import time
+from concurrent.futures import ThreadPoolExecutor
+import threading
+from builtins import open
+from potato.server_utils.config_module import config
+
+logger = logging.getLogger(__name__)
+
+from potato.item_state_management import get_item_state_manager
+from potato.ai.ai_endpoint import (
+ AIEndpointFactory,
+ Annotation_Type,
+ AnnotationInput,
+ ImageData,
+ VisualAnnotationInput,
+ ModelCapabilities,
+)
+from potato.ai.ollama_endpoint import OllamaEndpoint
+from potato.ai.openrouter_endpoint import OpenRouterEndpoint
+from potato.ai.ai_prompt import ModelManager, get_ai_prompt
+
+
+AICACHEMANAGER = None
+
+
+def _get_scheme_field(annotation_id: int, field: str, default=None):
+ """Safely get a field from an annotation scheme with a clear error message."""
+ schemes = config.get("annotation_schemes", [])
+ if annotation_id >= len(schemes):
+ raise ValueError(
+ f"AI cache: annotation_id {annotation_id} out of range "
+ f"(only {len(schemes)} scheme(s) configured)"
+ )
+ scheme = schemes[annotation_id]
+ if default is not None:
+ return scheme.get(field, default)
+ if field not in scheme:
+ scheme_name = scheme.get("name", f"index {annotation_id}")
+ scheme_type = scheme.get("annotation_type", "unknown")
+ raise ValueError(
+ f"AI cache: annotation scheme '{scheme_name}' (type '{scheme_type}') "
+ f"missing required field '{field}'"
+ )
+ return scheme[field]
+
+
+def _get_instance_text(instance_id: int) -> str:
+ """Get the text content from an instance using the configured text_key."""
+ item = get_item_state_manager().items()[instance_id]
+ item_data = item.get_data()
+
+ # Get the configured text_key
+ text_key = config.get("item_properties", {}).get("text_key", "text")
+
+ # Try the configured text_key first
+ if text_key in item_data:
+ return item_data[text_key]
+
+ # Fall back to common keys
+ for key in ['text', 'content', 'message']:
+ if key in item_data:
+ return item_data[key]
+
+ # Last resort: return any string value
+ for value in item_data.values():
+ if isinstance(value, str):
+ return value
+
+ return str(item_data)
+
+def _is_image_url(text: str) -> bool:
+ """Check if text appears to be an image URL."""
+ if not isinstance(text, str):
+ return False
+ text_lower = text.lower()
+ # Check for image extensions
+ image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp']
+ if any(ext in text_lower for ext in image_extensions):
+ return True
+ # Check for common image hosting services
+ image_hosts = ['unsplash.com', 'imgur.com', 'flickr.com', 'picsum.photos']
+ if any(host in text_lower for host in image_hosts):
+ return True
+ # Check if URL starts with http and might be an image
+ if text_lower.startswith(('http://', 'https://')) and 'image' in text_lower:
+ return True
+ return False
+
+def _get_image_data_from_url(url: str) -> ImageData:
+ """Download image from URL and return as ImageData.
+
+ Includes SSRF protection to prevent fetching from private/internal IPs.
+ """
+ import base64
+ import ipaddress
+ import socket
+ from urllib.parse import urlparse
+
+ # SSRF protection: validate URL scheme and resolve hostname
+ try:
+ parsed = urlparse(url)
+ if parsed.scheme not in ('http', 'https'):
+ logger.warning(f"Blocked non-HTTP image URL: {url[:100]}")
+ return None
+
+ hostname = parsed.hostname
+ if hostname:
+ addr_info = socket.getaddrinfo(hostname, None)
+ for info in addr_info:
+ ip_str = info[4][0]
+ try:
+ ip = ipaddress.ip_address(ip_str)
+ if ip.is_private or ip.is_loopback or ip.is_link_local:
+ logger.warning(
+ f"Blocked image URL resolving to private IP: "
+ f"{hostname} -> {ip_str}"
+ )
+ return None
+ except ValueError:
+ pass
+ except Exception as e:
+ logger.warning(f"Failed to validate image URL {url[:100]}: {e}")
+ return None
+
+ try:
+ response = requests.get(url, timeout=30)
+ response.raise_for_status()
+ b64_data = base64.b64encode(response.content).decode('utf-8')
+ # Determine mime type from content-type header or URL
+ content_type = response.headers.get('content-type', 'image/jpeg')
+ return ImageData(source='base64', data=b64_data, mime_type=content_type)
+ except Exception as e:
+ logger.error(f"Failed to download image from {url}: {e}")
+ return None
+
+def init_ai_cache_manager():
+ global AICACHEMANAGER
+ if AICACHEMANAGER is None:
+ AICACHEMANAGER = AiCacheManager()
+
+ return AICACHEMANAGER
+
+def get_ai_cache_manager():
+ """Get the AI cache manager instance. Returns None if not initialized (AI support disabled)."""
+ global AICACHEMANAGER
+ return AICACHEMANAGER
+
+def clear_ai_cache_manager():
+ """Clear the AI cache manager singleton. Used for testing."""
+ global AICACHEMANAGER
+ AICACHEMANAGER = None
+
+class AiCacheManager:
+ def __init__(self):
+ ai_support = config["ai_support"]
+ if not ai_support["enabled"]:
+ return
+ cache_config = ai_support.get("cache_config", {})
+ ai_config = ai_support.get("ai_config", {})
+ include = ai_config.get("include") or {}
+ special_include = include.get("special_include", None)
+ self.include_all = include.get("all", False)
+ self.special_includes = {}
+
+ self.model_manager = ModelManager()
+ self.model_manager.load_models_module()
+
+ if special_include:
+ for page_key, page_value in special_include.items():
+ # Convert string keys to integers for easier lookup
+ page_index = int(page_key)
+ self.special_includes[page_index] = {}
+ for annotation_id, annotation_types in page_value.items():
+ annotation_id_int = int(annotation_id)
+ self.special_includes[page_index][annotation_id_int] = annotation_types
+
+ # Disk cache configuration.
+ # F-028: tolerate a partial/absent ai_cache config (e.g. AI support
+ # enabled for ICL with no disk_cache block) instead of crashing boot
+ # with KeyError: 'disk_cache'.
+ disk_cache_cfg = cache_config.get("disk_cache", {}) if isinstance(cache_config, dict) else {}
+ self.disk_cache_enabled = disk_cache_cfg.get("enabled", False)
+
+ disk_cache_path = disk_cache_cfg.get("path")
+ if self.disk_cache_enabled and not disk_cache_path:
+ raise Exception("You have enable disk cache, but you did not specific the path!")
+ self.disk_persistence_path = disk_cache_path
+
+ # Validate cache path stays within task directory
+ if self.disk_persistence_path:
+ task_dir = os.path.abspath(config.get("task_dir", "."))
+ cache_abs = os.path.abspath(
+ os.path.join(task_dir, self.disk_persistence_path)
+ if not os.path.isabs(self.disk_persistence_path)
+ else self.disk_persistence_path
+ )
+ if not cache_abs.startswith(task_dir + os.sep) and cache_abs != task_dir:
+ raise ValueError(
+ f"Cache path '{self.disk_persistence_path}' resolves to "
+ f"'{cache_abs}' which is outside the task directory "
+ f"'{task_dir}'. Path traversal is not allowed."
+ )
+
+ # Prefetch configuration โ clamp to sane ranges.
+ # F-028: default to no prefetch when the prefetch block is absent
+ # (e.g. cache_config: {enabled: false}) instead of KeyError on boot.
+ prefetch_cfg = cache_config.get("prefetch", {}) if isinstance(cache_config, dict) else {}
+ self.warm_up_page_count = max(0, min(int(prefetch_cfg.get("warm_up_page_count", 0)), 10000))
+ self.prefetch_page_count_on_next = max(0, min(int(prefetch_cfg.get("on_next", 0)), 10000))
+ self.prefetch_page_count_on_prev = max(0, min(int(prefetch_cfg.get("on_prev", 0)), 10000))
+
+ # Option highlighting configuration
+ option_highlighting = ai_support.get("option_highlighting", {})
+ self.option_highlighting_enabled = option_highlighting.get("enabled", False)
+ self.option_highlighting_top_k = option_highlighting.get("top_k", 3)
+ self.option_highlighting_dim_opacity = option_highlighting.get("dim_opacity", 0.4)
+ self.option_highlighting_auto_apply = option_highlighting.get("auto_apply", True)
+ self.option_highlighting_schemas = option_highlighting.get("schemas", None) # None means all
+ # Prefetch count for option highlighting โ clamp to sane range
+ self.option_highlighting_prefetch_count = max(0, min(
+ int(option_highlighting.get("prefetch_count", 20)), 10000
+ ))
+
+ # Threading
+ self.in_progress = {}
+ self.lock = threading.RLock()
+ self.executor = ThreadPoolExecutor(max_workers=20)
+
+ AIEndpointFactory.register_endpoint("ollama", OllamaEndpoint)
+ AIEndpointFactory.register_endpoint("open_router", OpenRouterEndpoint)
+
+ # Register visual AI endpoints
+ try:
+ from potato.ai.yolo_endpoint import YOLOEndpoint
+ AIEndpointFactory.register_endpoint("yolo", YOLOEndpoint)
+ except ImportError:
+ logger.debug("YOLO endpoint not available (ultralytics not installed)")
+
+ try:
+ from potato.ai.ollama_vision_endpoint import OllamaVisionEndpoint
+ AIEndpointFactory.register_endpoint("ollama_vision", OllamaVisionEndpoint)
+ except ImportError:
+ logger.debug("Ollama Vision endpoint not available")
+
+ try:
+ from potato.ai.openai_vision_endpoint import OpenAIVisionEndpoint
+ AIEndpointFactory.register_endpoint("openai_vision", OpenAIVisionEndpoint)
+ except ImportError:
+ logger.debug("OpenAI Vision endpoint not available")
+
+ try:
+ from potato.ai.anthropic_vision_endpoint import AnthropicVisionEndpoint
+ AIEndpointFactory.register_endpoint("anthropic_vision", AnthropicVisionEndpoint)
+ except ImportError:
+ logger.debug("Anthropic Vision endpoint not available")
+
+ # Degrade gracefully if the AI backend (e.g. a local Ollama/vLLM server)
+ # is unreachable at boot: log a warning and serve the task with AI
+ # support disabled rather than aborting server startup.
+ try:
+ self.ai_endpoint = AIEndpointFactory.create_endpoint(config)
+ except Exception as e:
+ logger.warning(
+ "AI endpoint unavailable at startup (%s). Continuing with AI "
+ "support disabled. Check that your AI backend is running.", e
+ )
+ self.ai_endpoint = None
+
+ # Create visual endpoint if different from main endpoint
+ self.visual_endpoint = None
+ visual_endpoint_type = config.get("ai_support", {}).get("visual_endpoint_type")
+ if visual_endpoint_type and visual_endpoint_type != config.get("ai_support", {}).get("endpoint_type"):
+ visual_config = {
+ "ai_support": {
+ "enabled": True,
+ "endpoint_type": visual_endpoint_type,
+ "ai_config": config.get("ai_support", {}).get("visual_ai_config", config.get("ai_support", {}).get("ai_config", {}))
+ }
+ }
+ try:
+ self.visual_endpoint = AIEndpointFactory.create_endpoint(visual_config)
+ except Exception as e:
+ logger.warning(
+ "Visual AI endpoint unavailable at startup (%s). Continuing "
+ "without visual AI support.", e
+ )
+ self.visual_endpoint = None
+
+ annotation_scheme = config.get("annotation_schemes")
+ self.annotations = []
+ for scheme in annotation_scheme:
+ self.annotations.append(scheme)
+
+ # Check if main endpoint supports vision
+ self.endpoint_supports_vision = hasattr(self.ai_endpoint, 'query_with_image')
+ logger.info(f"AI endpoint supports vision: {self.endpoint_supports_vision}")
+
+ # Initialize cache
+ if self.disk_cache_enabled:
+ self.load_cache_from_disk()
+ self.start_warmup()
+
+ def _validate_assistant_compatibility(
+ self, instance_id: int, annotation_id: int, ai_assistant: str
+ ) -> tuple:
+ """
+ Validate that the AI assistant is compatible with the input type and model capabilities.
+
+ Args:
+ instance_id: The instance/item index
+ annotation_id: The annotation scheme index
+ ai_assistant: Type of assistance ('hint', 'keyword', 'rationale', 'detection', etc.)
+
+ Returns:
+ Tuple of (is_valid: bool, error_message: str)
+ If valid, error_message is empty string.
+ """
+ try:
+ text = _get_instance_text(instance_id)
+ is_image = _is_image_url(text)
+
+ # Determine which endpoint to use
+ if is_image and self.visual_endpoint:
+ endpoint = self.visual_endpoint
+ elif is_image and self.endpoint_supports_vision:
+ endpoint = self.ai_endpoint
+ else:
+ endpoint = self.ai_endpoint
+
+ # Get capabilities from endpoint
+ capabilities = getattr(endpoint, 'CAPABILITIES', None)
+
+ if capabilities is None:
+ # No capabilities declared - allow all (backward compatibility)
+ logger.debug(f"Endpoint {type(endpoint).__name__} has no CAPABILITIES, allowing {ai_assistant}")
+ return True, ""
+
+ # Check if the assistant type is supported
+ if not capabilities.supports_assistant(ai_assistant, is_image):
+ input_type = "image" if is_image else "text"
+ return False, (
+ f"Model {type(endpoint).__name__} does not support '{ai_assistant}' "
+ f"for {input_type} content"
+ )
+
+ return True, ""
+
+ except Exception as e:
+ logger.warning(f"Error validating assistant compatibility: {e}")
+ # On validation error, allow the request (fail open for now)
+ return True, ""
+
+ def get_endpoint_capabilities(self, for_image: bool = False) -> ModelCapabilities:
+ """
+ Get the capabilities of the appropriate endpoint for the given input type.
+
+ Args:
+ for_image: Whether the input is an image
+
+ Returns:
+ ModelCapabilities instance, or a default permissive one if not declared
+ """
+ if for_image and self.visual_endpoint:
+ endpoint = self.visual_endpoint
+ elif for_image and self.endpoint_supports_vision:
+ endpoint = self.ai_endpoint
+ else:
+ endpoint = self.ai_endpoint
+
+ capabilities = getattr(endpoint, 'CAPABILITIES', None)
+ if capabilities is None:
+ # Return permissive defaults for backward compatibility
+ return ModelCapabilities(
+ text_generation=True,
+ vision_input=for_image,
+ bounding_box_output=False,
+ text_classification=True,
+ image_classification=for_image,
+ rationale_generation=True,
+ keyword_extraction=not for_image,
+ )
+ return capabilities
+
+ def _get_ai_with_vision_support(self, text: str, prompt: str, output_format) -> str:
+ """
+ Get AI response, using vision if text is an image URL and endpoint supports it.
+ """
+ # Check if we should use vision
+ if self.endpoint_supports_vision and _is_image_url(text):
+ logger.debug(f"Using vision query for image URL: {text[:50]}...")
+ image_data = _get_image_data_from_url(text)
+ if image_data:
+ try:
+ return self.ai_endpoint.query_with_image(prompt, image_data, output_format)
+ except Exception as e:
+ logger.error(f"Vision query failed: {e}")
+ # Fall back to text query
+
+ # Fall back to regular text query
+ return self.ai_endpoint.query(prompt, output_format)
+
+ def start_warmup(self):
+ self.start_prefetch(0, self.warm_up_page_count)
+
+ # Also prefetch option highlights if enabled
+ if self.option_highlighting_enabled:
+ self.start_option_highlight_prefetch(0, self.warm_up_page_count)
+
+ total = len(self.in_progress)
+ desc = "Preloading the AI"
+
+ progress_bar = tqdm(total=total, desc=desc, unit="item")
+
+ def count_completed():
+ return total - len(self.in_progress)
+
+ prev_done = 0
+ while self.in_progress:
+ current_done = count_completed()
+ progress_bar.update(current_done - prev_done)
+ prev_done = current_done
+ time.sleep(0.2)
+
+ final_done = count_completed()
+ if final_done > prev_done:
+ progress_bar.update(final_done - prev_done)
+
+ progress_bar.close()
+
+ def load_disk_cache_data(self, file_path: str) -> Dict[str, Any]:
+ """loads the cache JSON from disk and returns a dictionary of stringified keys to values."""
+ try:
+ with open(file_path, 'r', encoding='utf-8') as f:
+ return json.load(f)
+ except Exception as e:
+ logger.error(f"Error loading disk cache: {e}")
+ return {}
+
+ def load_cache_from_disk(self):
+ """Initializes disk cache file if it doesn't exist."""
+ if not self.disk_cache_enabled or not self.disk_persistence_path:
+ return
+
+ if os.path.exists(self.disk_persistence_path):
+ data = self.load_disk_cache_data(self.disk_persistence_path)
+ logger.info(f"Disk cache initialized with {len(data)} items")
+ else:
+ try:
+ # Create parent directory if it doesn't exist
+ os.makedirs(os.path.dirname(self.disk_persistence_path), exist_ok=True)
+ with open(self.disk_persistence_path, 'w', encoding='utf-8') as file:
+ json.dump({}, file)
+ logger.info(f"Initialized empty disk cache at {self.disk_persistence_path}")
+ except Exception as e:
+ logger.error(f"Failed to create disk cache: {e}")
+
+ def save_cache_to_disk(self, key, value):
+ """saves a single key-value pair to disk cache using atomic write."""
+ if not self.disk_cache_enabled or not self.disk_persistence_path:
+ return
+
+ try:
+ os.makedirs(os.path.dirname(self.disk_persistence_path), exist_ok=True)
+
+ # Load existing disk data first
+ existing_disk_data = {}
+ if os.path.exists(self.disk_persistence_path):
+ existing_disk_data = self.load_disk_cache_data(self.disk_persistence_path)
+
+ # Add the new key-value pair
+ existing_disk_data[str(key)] = value
+
+ temp_path = self.disk_persistence_path + ".tmp"
+ with open(temp_path, 'w', encoding='utf-8') as f:
+ json.dump(existing_disk_data, f, indent=2, ensure_ascii=False)
+ os.rename(temp_path, self.disk_persistence_path)
+ except Exception as e:
+ logger.error(f"Error saving cache to disk: {e}")
+
+ def add_to_cache(self, key, value):
+ """inserts a key-value into the disk cache."""
+ with self.lock:
+ if self.disk_cache_enabled:
+ self.save_cache_to_disk(key, value)
+
+ def get_from_cache(self, key):
+ """Tries to retrieve the item from disk cache."""
+ with self.lock:
+ # Try disk cache
+ if self.disk_cache_enabled and self.disk_persistence_path and os.path.exists(self.disk_persistence_path):
+ try:
+ disk_data = self.load_disk_cache_data(self.disk_persistence_path)
+ key_str = str(key)
+ if key_str in disk_data:
+ return disk_data[key_str]
+ except Exception as e:
+ logger.error(f"Error reading from disk: {e}")
+ return None
+
+ def generate_likert(self, instance_id: int, annotation_id: int, ai_assistant: str) -> str:
+ from string import Template
+ annotation_type = _get_scheme_field(annotation_id, "annotation_type")
+ description = _get_scheme_field(annotation_id, "description")
+ text = _get_instance_text(instance_id)
+ min_label = _get_scheme_field(annotation_id, "min_label")
+ max_label = _get_scheme_field(annotation_id, "max_label")
+ size = _get_scheme_field(annotation_id, "size")
+
+ ai_prompt = get_ai_prompt()
+ output_format = self.model_manager.get_model_class_by_name(ai_prompt[annotation_type].get(ai_assistant).get("output_format"))
+
+ # Check if we should use vision endpoint for image-based content
+ if self.endpoint_supports_vision and _is_image_url(text):
+ logger.debug(f"Using vision for likert {ai_assistant} on image: {text[:50]}...")
+ image_data = _get_image_data_from_url(text)
+ if image_data:
+ # Build vision-specific prompts based on ai_assistant type
+ if ai_assistant == "hint":
+ prompt = f"""Look at this image and help with the following annotation task:
+
+Task: {description}
+Rating scale: {size} points, from "{min_label}" (1) to "{max_label}" ({size})
+
+Please analyze the image and suggest an appropriate rating with a brief explanation.
+Respond in JSON format: {{"hint": "", "suggestive_choice": ""}}"""
+ elif ai_assistant == "rationale":
+ prompt = f"""Look at this image and explain the reasoning for different rating choices:
+
+Task: {description}
+Rating scale: {size} points, from "{min_label}" (1) to "{max_label}" ({size})
+
+For each possible rating, explain what visual evidence in the image would support that rating.
+Respond in JSON format: {{"rationales": [{{"label": "", "reasoning": ""}}]}}"""
+ elif ai_assistant == "keyword":
+ prompt = f"""Look at this image and identify visual features relevant to the rating task:
+
+Task: {description}
+Rating scale: {size} points, from "{min_label}" (1) to "{max_label}" ({size})
+
+Identify key visual elements that would influence the rating.
+Respond in JSON format: {{"keywords": ["", ""]}}"""
+ else:
+ prompt = f"Analyze this image for: {description}"
+
+ try:
+ return self.ai_endpoint.query_with_image(prompt, image_data, output_format)
+ except Exception as e:
+ logger.error(f"Vision query failed for likert {ai_assistant}: {e}")
+
+ # Fall back to standard text-based generation
+ data = AnnotationInput(
+ ai_assistant=ai_assistant,
+ annotation_type=annotation_type,
+ text=text,
+ description=description,
+ min_label=min_label,
+ max_label=max_label,
+ size=size
+ )
+ res = self.ai_endpoint.get_ai(data, output_format)
+ return res
+
+ def generate_multiselect(self, instance_id: int, annotation_id: int, ai_assistant: str) -> str:
+ annotation_type = _get_scheme_field(annotation_id, "annotation_type")
+ description = _get_scheme_field(annotation_id, "description")
+ labels = _get_scheme_field(annotation_id, "labels")
+ text = _get_instance_text(instance_id)
+
+ ai_prompt = get_ai_prompt()
+ output_format = self.model_manager.get_model_class_by_name(ai_prompt[annotation_type].get(ai_assistant).get("output_format"))
+
+ # Check if we should use vision endpoint for image-based content
+ if self.endpoint_supports_vision and _is_image_url(text):
+ logger.debug(f"Using vision for multiselect {ai_assistant} on image: {text[:50]}...")
+ image_data = _get_image_data_from_url(text)
+ if image_data:
+ # Format labels for the prompt
+ label_names = [l.get('name', l) if isinstance(l, dict) else l for l in labels]
+ labels_str = ', '.join(f'"{name}"' for name in label_names)
+
+ # Build vision-specific prompts based on ai_assistant type
+ if ai_assistant == "hint":
+ prompt = f"""Look at this image and help with the following annotation task:
+
+Task: {description}
+Available options (select all that apply): {labels_str}
+
+Please analyze the image and suggest which options apply.
+Respond in JSON format: {{"hint": "", "suggestive_choices": ["", ""]}}"""
+ elif ai_assistant == "rationale":
+ prompt = f"""Look at this image and explain the reasoning for each option:
+
+Task: {description}
+Available options: {labels_str}
+
+For each option, explain what visual evidence supports or contradicts it.
+Respond in JSON format: {{"rationales": [{{"label": "", "reasoning": ""}}]}}"""
+ elif ai_assistant == "keyword":
+ prompt = f"""Look at this image and identify visual features for each option:
+
+Task: {description}
+Available options: {labels_str}
+
+For each option, identify visual cues that indicate its presence.
+Respond in JSON format: {{"label_keywords": [{{"label": "", "keywords": ["", ""]}}]}}"""
+ else:
+ prompt = f"Analyze this image for: {description}. Options: {labels_str}"
+
+ try:
+ return self.ai_endpoint.query_with_image(prompt, image_data, output_format)
+ except Exception as e:
+ logger.error(f"Vision query failed for multiselect {ai_assistant}: {e}")
+
+ # Fall back to standard text-based generation
+ data = AnnotationInput(
+ ai_assistant=ai_assistant,
+ annotation_type=annotation_type,
+ text=text,
+ description=description,
+ labels=labels
+ )
+ res = self.ai_endpoint.get_ai(data, output_format)
+ return res
+
+ def generate_radio(self, instance_id: int, annotation_id: int, ai_assistant: str) -> str:
+ annotation_type = _get_scheme_field(annotation_id, "annotation_type")
+ description = _get_scheme_field(annotation_id, "description")
+ text = _get_instance_text(instance_id)
+ labels = _get_scheme_field(annotation_id, "labels")
+
+ ai_prompt = get_ai_prompt()
+ output_format = self.model_manager.get_model_class_by_name(ai_prompt[annotation_type].get(ai_assistant).get("output_format"))
+
+ # Check if we should use vision endpoint for image-based content
+ if self.endpoint_supports_vision and _is_image_url(text):
+ logger.debug(f"Using vision for radio {ai_assistant} on image: {text[:50]}...")
+ image_data = _get_image_data_from_url(text)
+ if image_data:
+ # Format labels for the prompt
+ label_names = [l.get('name', l) if isinstance(l, dict) else l for l in labels]
+ labels_str = ', '.join(f'"{name}"' for name in label_names)
+
+ # Build vision-specific prompts based on ai_assistant type
+ if ai_assistant == "hint":
+ prompt = f"""Look at this image and help with the following annotation task:
+
+Task: {description}
+Available options: {labels_str}
+
+Please analyze the image and suggest the most appropriate option.
+Respond in JSON format: {{"hint": "", "suggestive_choice": ""}}"""
+ elif ai_assistant == "rationale":
+ prompt = f"""Look at this image and explain the reasoning for each option:
+
+Task: {description}
+Available options: {labels_str}
+
+For each option, explain what visual evidence in the image supports or contradicts it.
+Respond in JSON format: {{"rationales": [{{"label": "", "reasoning": ""}}]}}"""
+ elif ai_assistant == "keyword":
+ prompt = f"""Look at this image and identify visual features for each option:
+
+Task: {description}
+Available options: {labels_str}
+
+For each option, identify visual cues that would indicate its presence.
+Respond in JSON format: {{"label_keywords": [{{"label": "", "keywords": ["", ""]}}]}}"""
+ else:
+ prompt = f"Analyze this image for: {description}. Options: {labels_str}"
+
+ try:
+ return self.ai_endpoint.query_with_image(prompt, image_data, output_format)
+ except Exception as e:
+ logger.error(f"Vision query failed for radio {ai_assistant}: {e}")
+
+ # Fall back to standard text-based generation
+ data = AnnotationInput(
+ ai_assistant=ai_assistant,
+ annotation_type=annotation_type,
+ text=text,
+ description=description,
+ labels=labels
+ )
+ res = self.ai_endpoint.get_ai(data, output_format)
+ return res
+
+ def generate_number(self, instance_id: int, annotation_id: int, ai_assistant: str) -> str:
+ annotation_type = _get_scheme_field(annotation_id, "annotation_type")
+ description = _get_scheme_field(annotation_id, "description")
+ text = _get_instance_text(instance_id)
+
+ data = AnnotationInput(
+ ai_assistant=ai_assistant,
+ annotation_type=annotation_type,
+ text=text,
+ description=description,
+ )
+ ai_prompt = get_ai_prompt();
+ output_format = self.model_manager.get_model_class_by_name(ai_prompt[annotation_type].get(ai_assistant).get("output_format"))
+ res = self.ai_endpoint.get_ai(data, output_format)
+ return res
+
+ def generate_select(self, instance_id: int, annotation_id: int, ai_assistant: str) -> str:
+ annotation_type = _get_scheme_field(annotation_id, "annotation_type")
+ description = _get_scheme_field(annotation_id, "description")
+ labels = _get_scheme_field(annotation_id, "labels")
+ text = _get_instance_text(instance_id)
+
+
+ data = AnnotationInput(
+ ai_assistant=ai_assistant,
+ annotation_type=annotation_type,
+ text=text,
+ description=description,
+ labels=labels
+ )
+ ai_prompt = get_ai_prompt();
+ output_format = self.model_manager.get_model_class_by_name(ai_prompt[annotation_type].get(ai_assistant).get("output_format"))
+ res = self.ai_endpoint.get_ai(data, output_format)
+ return res
+
+ def generate_slider(self, instance_id: int, annotation_id: int, ai_assistant: str) -> str:
+ annotation_type = _get_scheme_field(annotation_id, "annotation_type")
+ description = _get_scheme_field(annotation_id, "description")
+ min_value = _get_scheme_field(annotation_id, "min_value")
+ max_value = _get_scheme_field(annotation_id, "max_value")
+ step = _get_scheme_field(annotation_id, "step", default=1)
+ text = _get_instance_text(instance_id)
+
+ data = AnnotationInput(
+ ai_assistant=ai_assistant,
+ annotation_type=annotation_type,
+ text=text,
+ description=description,
+ min_value=min_value,
+ max_value=max_value,
+ step=step
+ )
+
+ ai_prompt = get_ai_prompt();
+ output_format = self.model_manager.get_model_class_by_name(ai_prompt[annotation_type].get(ai_assistant).get("output_format"))
+ res = self.ai_endpoint.get_ai(data, output_format)
+ return res
+
+ def generate_span(self, instance_id: int, annotation_id: int, ai_assistant: str) -> str:
+ annotation_type = _get_scheme_field(annotation_id, "annotation_type")
+ description = _get_scheme_field(annotation_id, "description")
+ labels = _get_scheme_field(annotation_id, "labels")
+ text = _get_instance_text(instance_id)
+
+ data = AnnotationInput(
+ ai_assistant=ai_assistant,
+ annotation_type=annotation_type,
+ text=text,
+ description=description,
+ labels=labels
+ )
+ ai_prompt = get_ai_prompt();
+ logger.debug(f"Generating span annotation with labels: {labels}")
+ output_format = self.model_manager.get_model_class_by_name(ai_prompt[annotation_type].get(ai_assistant).get("output_format"))
+ res = self.ai_endpoint.get_ai(data, output_format)
+ return res
+
+ def generate_textbox(self, instance_id: int, annotation_id: int, ai_assistant: str) -> str:
+ logger.debug(f"Generating textbox for annotation_id: {annotation_id}")
+ annotation_type = _get_scheme_field(annotation_id, "annotation_type")
+ description = _get_scheme_field(annotation_id, "description")
+ text = _get_instance_text(instance_id)
+
+ data = AnnotationInput(
+ ai_assistant=ai_assistant,
+ annotation_type=annotation_type,
+ text=text,
+ description=description,
+ )
+ ai_prompt = get_ai_prompt();
+ output_format = self.model_manager.get_model_class_by_name(ai_prompt[annotation_type].get(ai_assistant).get("output_format"))
+ res = self.ai_endpoint.get_ai(data, output_format)
+ return res
+
+ def generate_image_annotation(self, instance_id: int, annotation_id: int, ai_assistant: str) -> Dict:
+ """Generate AI assistance for image annotation tasks.
+
+ Args:
+ instance_id: The instance/item index
+ annotation_id: The annotation scheme index
+ ai_assistant: Type of assistance ('detection', 'classification', 'hint', 'pre_annotate', etc.)
+
+ Returns:
+ Dict with AI suggestions (detections, classifications, hints, etc.)
+ """
+ logger.debug(f"Generating image annotation for instance={instance_id}, annotation={annotation_id}, assistant={ai_assistant}")
+
+ annotation_type = _get_scheme_field(annotation_id, "annotation_type")
+ description = _get_scheme_field(annotation_id, "description", default="")
+ labels = _get_scheme_field(annotation_id, "labels", default=[])
+
+ # Extract label names if labels are dicts
+ if labels and isinstance(labels[0], dict):
+ labels = [l.get("name", str(l)) for l in labels]
+
+ # Get image URL from item data
+ item_data = get_item_state_manager().items()[instance_id].get_data()
+ image_url = self._extract_image_url(item_data)
+
+ if not image_url:
+ return {"error": "No image URL found in instance data"}
+
+ # Determine which endpoint to use
+ endpoint = self._get_visual_endpoint()
+ if not endpoint:
+ return {"error": "No visual AI endpoint configured"}
+
+ # Check if endpoint supports visual queries
+ if not hasattr(endpoint, 'query_with_image'):
+ # Fall back to text-based hint
+ return self._generate_text_hint_for_visual(instance_id, annotation_id, ai_assistant)
+
+ # Prepare image data
+ image_data = self._prepare_image_data(image_url)
+
+ # Get confidence threshold from config
+ confidence_threshold = _get_scheme_field(annotation_id, "ai_support", default={}).get(
+ "confidence_threshold", 0.5
+ )
+
+ # Build VisualAnnotationInput
+ data = VisualAnnotationInput(
+ ai_assistant=ai_assistant,
+ annotation_type=annotation_type,
+ task_type=ai_assistant, # detection, classification, hint, etc.
+ image_data=image_data,
+ description=description,
+ labels=labels,
+ confidence_threshold=confidence_threshold
+ )
+
+ # Get output format from prompt config
+ ai_prompt = get_ai_prompt()
+ prompt_config = ai_prompt.get(annotation_type, {}).get(ai_assistant, {})
+ output_format_name = prompt_config.get("output_format", "visual_detection")
+ output_format = self.model_manager.get_model_class_by_name(output_format_name)
+
+ # Query the visual endpoint
+ result = endpoint.get_visual_ai(data, output_format)
+ return result
+
+ def generate_video_annotation(self, instance_id: int, annotation_id: int, ai_assistant: str) -> Dict:
+ """Generate AI assistance for video annotation tasks.
+
+ Args:
+ instance_id: The instance/item index
+ annotation_id: The annotation scheme index
+ ai_assistant: Type of assistance ('scene_detection', 'frame_classification', etc.)
+
+ Returns:
+ Dict with AI suggestions (segments, keyframes, etc.)
+ """
+ logger.debug(f"Generating video annotation for instance={instance_id}, annotation={annotation_id}, assistant={ai_assistant}")
+
+ annotation_type = _get_scheme_field(annotation_id, "annotation_type")
+ description = _get_scheme_field(annotation_id, "description", default="")
+ labels = _get_scheme_field(annotation_id, "labels", default=[])
+
+ # Extract label names if labels are dicts
+ if labels and isinstance(labels[0], dict):
+ labels = [l.get("name", str(l)) for l in labels]
+
+ # Get video URL from item data
+ item_data = get_item_state_manager().items()[instance_id].get_data()
+ video_url = self._extract_video_url(item_data)
+
+ if not video_url:
+ return {"error": "No video URL found in instance data"}
+
+ # Determine which endpoint to use
+ endpoint = self._get_visual_endpoint()
+ if not endpoint:
+ return {"error": "No visual AI endpoint configured"}
+
+ # Check if endpoint supports visual queries
+ if not hasattr(endpoint, 'query_with_image'):
+ return self._generate_text_hint_for_visual(instance_id, annotation_id, ai_assistant)
+
+ # Extract video frames
+ try:
+ frames = endpoint.extract_video_frames(video_url)
+ video_metadata = endpoint.get_video_metadata(video_url)
+ except Exception as e:
+ logger.error(f"Failed to extract video frames: {e}")
+ return {"error": f"Failed to process video: {str(e)}"}
+
+ # Build VisualAnnotationInput
+ data = VisualAnnotationInput(
+ ai_assistant=ai_assistant,
+ annotation_type=annotation_type,
+ task_type=ai_assistant,
+ image_data=frames, # List of frame images
+ description=description,
+ labels=labels,
+ video_metadata=video_metadata
+ )
+
+ # Get output format
+ ai_prompt = get_ai_prompt()
+ prompt_config = ai_prompt.get(annotation_type, {}).get(ai_assistant, {})
+ output_format_name = prompt_config.get("output_format", "video_scene_detection")
+ output_format = self.model_manager.get_model_class_by_name(output_format_name)
+
+ # Query the visual endpoint
+ result = endpoint.get_visual_ai(data, output_format)
+ return result
+
+ def _get_visual_endpoint(self):
+ """Get the appropriate endpoint for visual tasks."""
+ # Use dedicated visual endpoint if configured
+ if self.visual_endpoint:
+ return self.visual_endpoint
+
+ # Check if main endpoint supports vision
+ if hasattr(self.ai_endpoint, 'query_with_image'):
+ return self.ai_endpoint
+
+ # Try to find a visual endpoint from registered types
+ visual_types = ['yolo', 'ollama_vision', 'openai_vision', 'anthropic_vision']
+ for vtype in visual_types:
+ if vtype in AIEndpointFactory._endpoints:
+ try:
+ visual_config = {
+ "ai_support": {
+ "enabled": True,
+ "endpoint_type": vtype,
+ "ai_config": config.get("ai_support", {}).get("ai_config", {})
+ }
+ }
+ return AIEndpointFactory.create_endpoint(visual_config)
+ except Exception as e:
+ logger.debug(f"Could not create {vtype} endpoint: {e}")
+ continue
+
+ return None
+
+ def _extract_image_url(self, item_data: Dict) -> str:
+ """Extract image URL from item data.
+
+ Looks for common field names that might contain image URLs.
+ """
+ # Common field names for images
+ image_fields = ['image', 'image_url', 'img', 'img_url', 'url', 'path', 'file', 'src']
+
+ for field in image_fields:
+ if field in item_data:
+ value = item_data[field]
+ if isinstance(value, str) and (
+ value.startswith(('http://', 'https://', '/')) or
+ value.endswith(('.jpg', '.jpeg', '.png', '.gif', '.webp'))
+ ):
+ return value
+
+ # Check 'text' field for URL (common in simple configs)
+ if 'text' in item_data:
+ text = item_data['text']
+ if isinstance(text, str) and (
+ text.startswith(('http://', 'https://')) and
+ any(ext in text.lower() for ext in ['.jpg', '.jpeg', '.png', '.gif', '.webp'])
+ ):
+ return text
+
+ return None
+
+ def _extract_video_url(self, item_data: Dict) -> str:
+ """Extract video URL from item data."""
+ # Common field names for videos
+ video_fields = ['video', 'video_url', 'url', 'path', 'file', 'src', 'media']
+
+ for field in video_fields:
+ if field in item_data:
+ value = item_data[field]
+ if isinstance(value, str) and (
+ value.startswith(('http://', 'https://', '/')) or
+ value.endswith(('.mp4', '.webm', '.ogg', '.avi', '.mov'))
+ ):
+ return value
+
+ # Check 'text' field for URL
+ if 'text' in item_data:
+ text = item_data['text']
+ if isinstance(text, str) and (
+ text.startswith(('http://', 'https://')) and
+ any(ext in text.lower() for ext in ['.mp4', '.webm', '.ogg', '.avi', '.mov'])
+ ):
+ return text
+
+ return None
+
+ def _prepare_image_data(self, image_url: str) -> ImageData:
+ """Prepare ImageData from URL or path."""
+ if image_url.startswith(('http://', 'https://')):
+ return ImageData(source="url", data=image_url)
+ else:
+ # Local file path - encode as base64
+ from potato.ai.visual_ai_endpoint import BaseVisualAIEndpoint
+ return BaseVisualAIEndpoint.encode_image_to_base64(image_url)
+
+ def _generate_text_hint_for_visual(self, instance_id: int, annotation_id: int, ai_assistant: str) -> Dict:
+ """Generate text-based hint when visual endpoint is not available."""
+ description = config["annotation_schemes"][annotation_id].get("description", "")
+ labels = config["annotation_schemes"][annotation_id].get("labels", [])
+
+ if labels and isinstance(labels[0], dict):
+ labels = [l.get("name", str(l)) for l in labels]
+
+ return {
+ "hint": f"Review the {'image' if 'image' in config['annotation_schemes'][annotation_id]['annotation_type'] else 'video'} carefully. "
+ f"Look for: {', '.join(labels) if labels else 'relevant content'}. "
+ f"Task: {description}",
+ "suggestive_choice": ""
+ }
+
+ def is_option_highlighting_enabled_for_scheme(self, annotation_id: int) -> bool:
+ """Check if option highlighting is enabled for a specific annotation scheme."""
+ if not self.option_highlighting_enabled:
+ return False
+
+ scheme = config["annotation_schemes"][annotation_id]
+ annotation_type = scheme.get("annotation_type", "")
+ scheme_name = scheme.get("name", "")
+
+ # Only applicable to discrete option types
+ discrete_types = ["radio", "multiselect", "likert", "select"]
+ if annotation_type not in discrete_types:
+ return False
+
+ # Check if schemas filter is set
+ if self.option_highlighting_schemas is not None:
+ if scheme_name not in self.option_highlighting_schemas:
+ return False
+
+ return True
+
+ def get_option_highlighting_config(self) -> Dict:
+ """Get the option highlighting configuration for the frontend."""
+ return {
+ "enabled": self.option_highlighting_enabled,
+ "top_k": self.option_highlighting_top_k,
+ "dim_opacity": self.option_highlighting_dim_opacity,
+ "auto_apply": self.option_highlighting_auto_apply,
+ "schemas": self.option_highlighting_schemas,
+ "prefetch_count": self.option_highlighting_prefetch_count,
+ }
+
+ def generate_option_highlights(self, instance_id: int, annotation_id: int) -> Dict:
+ """Generate option highlighting suggestions for an annotation.
+
+ Args:
+ instance_id: The instance/item index
+ annotation_id: The annotation scheme index
+
+ Returns:
+ Dict with highlighted options and configuration:
+ {
+ "highlighted": ["option1", "option2"],
+ "top_k": 3,
+ "confidence": 0.85
+ }
+ """
+ from string import Template
+
+ if not self.is_option_highlighting_enabled_for_scheme(annotation_id):
+ return {"error": "Option highlighting not enabled for this scheme"}
+
+ annotation_type = _get_scheme_field(annotation_id, "annotation_type", default="")
+ description = _get_scheme_field(annotation_id, "description", default="")
+ labels = _get_scheme_field(annotation_id, "labels", default=[])
+
+ # Extract label names
+ if labels and isinstance(labels[0], dict):
+ label_names = [l.get("name", str(l)) for l in labels]
+ else:
+ label_names = [str(l) for l in labels]
+
+ # For likert scales, generate label names from min/max labels
+ if annotation_type == "likert":
+ size = scheme.get("size", 5)
+ min_label = scheme.get("min_label", "1")
+ max_label = scheme.get("max_label", str(size))
+ label_names = [f"{i+1} ({min_label if i == 0 else max_label if i == size-1 else ''})" for i in range(size)]
+ # Clean up empty parentheses
+ label_names = [l.replace(" ()", "") for l in label_names]
+
+ text = _get_instance_text(instance_id)
+ top_k = min(self.option_highlighting_top_k, len(label_names))
+
+ # Get prompt template
+ ai_prompt = get_ai_prompt()
+ prompt_config = ai_prompt.get("option_highlight", {}).get("option_highlight", {})
+
+ if not prompt_config:
+ return {"error": "Option highlight prompt not configured"}
+
+ prompt_template = prompt_config.get("prompt", "")
+ output_format_name = prompt_config.get("output_format", "option_highlight")
+ output_format = self.model_manager.get_model_class_by_name(output_format_name)
+
+ # Build the prompt with clear delimiters to mitigate prompt injection.
+ # The user content is wrapped in XML-style tags so the LLM can
+ # distinguish between instructions and untrusted data.
+ delimited_text = (
+ f"\n{text}\n "
+ )
+ template = Template(prompt_template)
+ prompt = template.safe_substitute(
+ text=delimited_text,
+ description=description,
+ labels=", ".join(label_names),
+ top_k=top_k
+ )
+
+ # Query the AI endpoint
+ try:
+ result = self.ai_endpoint.query(prompt, output_format)
+ logger.debug(f"Option highlight raw result: {result}")
+
+ # Parse the result
+ if isinstance(result, str):
+ import json as json_module
+ try:
+ # Try to parse JSON from the response
+ result = json_module.loads(result)
+ except json_module.JSONDecodeError:
+ # Try to extract JSON from markdown code block
+ if "```json" in result:
+ json_start = result.find("```json") + 7
+ json_end = result.find("```", json_start)
+ result = json_module.loads(result[json_start:json_end].strip())
+ elif "```" in result:
+ json_start = result.find("```") + 3
+ json_end = result.find("```", json_start)
+ result = json_module.loads(result[json_start:json_end].strip())
+ else:
+ return {"error": f"Could not parse response: {result[:100]}"}
+
+ highlighted = result.get("highlighted_options", [])
+ confidence = result.get("confidence", None)
+
+ # Validate highlighted options against available labels
+ valid_highlighted = [opt for opt in highlighted if opt in label_names]
+
+ return {
+ "highlighted": valid_highlighted[:top_k],
+ "top_k": top_k,
+ "confidence": confidence
+ }
+
+ except Exception as e:
+ logger.error(f"Error generating option highlights: {e}")
+ return {"error": str(e)}
+
+ def get_option_highlights(self, instance_id: int, annotation_id: int) -> Dict:
+ """Get option highlights from cache or generate them.
+
+ Args:
+ instance_id: The instance/item index
+ annotation_id: The annotation scheme index
+
+ Returns:
+ Dict with highlighted options
+ """
+ key = (instance_id, annotation_id, "option_highlight")
+
+ # Try cache first
+ if self.disk_cache_enabled:
+ cached = self.get_from_cache(key)
+ if cached is not None:
+ logger.debug(f"Option highlight cache hit for {key}")
+ return cached
+
+ # Generate
+ result = self.generate_option_highlights(instance_id, annotation_id)
+
+ # Cache if successful
+ if "error" not in result and self.disk_cache_enabled:
+ self.add_to_cache(key, result)
+
+ return result
+
+ def start_option_highlight_prefetch(self, page_id: int, prefetch_amount: int = None):
+ """Prefetch option highlights for upcoming items.
+
+ Args:
+ page_id: Current page/instance index
+ prefetch_amount: Number of items to prefetch (uses config default if None)
+ """
+ if not self.option_highlighting_enabled or not self.disk_cache_enabled:
+ return
+
+ if prefetch_amount is None:
+ prefetch_amount = self.option_highlighting_prefetch_count
+
+ ism = get_item_state_manager()
+ with self.lock:
+ # Calculate range
+ if prefetch_amount >= 0:
+ start_idx = page_id
+ end_idx = min(start_idx + prefetch_amount, len(ism.items()))
+ else:
+ start_idx = max(page_id + prefetch_amount, 0)
+ end_idx = page_id
+
+ keys = []
+ for i in range(start_idx, end_idx):
+ for annotation_id, scheme in enumerate(config["annotation_schemes"]):
+ if self.is_option_highlighting_enabled_for_scheme(annotation_id):
+ key = (i, annotation_id, "option_highlight")
+ # Check if not already cached or in progress
+ if self.get_from_cache(key) is None and key not in self.in_progress:
+ keys.append(key)
+
+ # Submit prefetch jobs
+ for key in keys:
+ instance_id, annotation_id, _ = key
+ future = self.executor.submit(self.generate_option_highlights, instance_id, annotation_id)
+ self.in_progress[key] = future
+
+ def callback(fut, cache_key=key):
+ with self.lock:
+ try:
+ result = fut.result()
+ if "error" not in result:
+ self.add_to_cache(cache_key, result)
+ except Exception as e:
+ logger.error(f"Option highlight prefetch failed for {cache_key}: {e}")
+ self.in_progress.pop(cache_key, None)
+
+ future.add_done_callback(callback)
+
+ if keys:
+ logger.debug(f"Started option highlight prefetch for {len(keys)} items")
+
+ def get_include_all(self):
+ return self.include_all
+
+ def get_special_include(self, page_number_int, annotation_id_int):
+ logger.debug(f"get_special_include: page={page_number_int}, annotation_id={annotation_id_int}")
+ if not self.special_includes.get(page_number_int):
+ return None
+ elif not self.special_includes.get(page_number_int).get(annotation_id_int):
+ return None
+ return self.special_includes.get(page_number_int).get(annotation_id_int)
+
+ def start_prefetch(self, page_id, prefetch_amount):
+ """Prefetches a fixed number of upcoming items to warm the cache."""
+ if not config.get("ai_support", {}).get("enabled") or not self.disk_cache_enabled:
+ return
+
+ ism = get_item_state_manager()
+ with self.lock:
+ # Calculate range bounds
+ if prefetch_amount >= 0:
+ start_idx = page_id
+ end_idx = min(start_idx + prefetch_amount, len(ism.items()))
+ else:
+ start_idx = max(page_id - prefetch_amount, 0)
+ end_idx = page_id
+
+ logger.debug(f"Prefetch range: start_idx={start_idx}, end_idx={end_idx}")
+ keys = []
+
+ for i in range(start_idx, end_idx):
+ # Check if this page should be included
+ if not self.should_include_page(i):
+ continue
+
+ # Process each annotation scheme for this page
+ for annotation_id, scheme in enumerate(config["annotation_schemes"]):
+ if not self.should_include_scheme(i, annotation_id):
+ continue
+
+ annotation_type = scheme["annotation_type"]
+ ai_prompt = get_ai_prompt()
+
+
+ if not ai_prompt[annotation_type]:
+ raise Exception(f"{annotation_type} is not defined in ai_prompt")
+
+ # Generate keys for this page/scheme combination
+ scheme_keys = self.get_keys_for_scheme(i, annotation_type, annotation_id, ai_prompt)
+ keys.extend(scheme_keys)
+
+ if keys:
+ self.prefetch(keys)
+
+ def should_include_page(self, page_index):
+ """Determine if a page should be included based on include_all and special_includes."""
+ if self.include_all:
+ return True
+ return page_index in self.special_includes
+
+ def should_include_scheme(self, page_index, annotation_id):
+ """Determine if a scheme should be included for a given page."""
+ if self.include_all:
+ return True
+
+ # Check if page is in special_includes and scheme is specified
+ if page_index in self.special_includes:
+ page_includes = self.special_includes[page_index]
+ # Handle both list and dict formats for page_includes
+ if isinstance(page_includes, dict):
+ return annotation_id in page_includes
+ elif isinstance(page_includes, list):
+ return annotation_id in page_includes
+
+ return False
+
+ def get_keys_for_scheme(self, page_index, annotation_type, annotation_id, ai_prompt):
+ """Get all keys for a specific page combination."""
+ keys = []
+
+ # Check if this page/annotation has specific overrides in special_includes
+ if (page_index in self.special_includes and
+ isinstance(self.special_includes[page_index], dict) and
+ annotation_id in self.special_includes[page_index]):
+
+ # Use special_includes (overrides include_all setting)
+ specified_keys = self.special_includes[page_index][annotation_id]
+ for key in specified_keys:
+ keys.append((page_index, annotation_id, key))
+ elif self.include_all:
+ # No specific override, so include all available keys for this annotation type
+ for key in ai_prompt[annotation_type]:
+ keys.append((page_index, annotation_id, key))
+ # If include_all is False and no special_include entry, return empty keys
+
+ return keys
+
+ def prefetch(self, keys: list):
+ """checks if keys are already cached and asynchronously generates missing ones"""
+ with self.lock:
+ for key in keys:
+ if self.get_from_cache(key) is None and key not in self.in_progress:
+ # i, annotation_id, annotation_type, ai_prompt
+ instance_id, annotation_id, ai_assistant = key
+
+ future = self.executor.submit(self.compute_help, instance_id, annotation_id, ai_assistant)
+ self.in_progress[key] = future
+ def callback(fut, cache_key=key):
+ with self.lock:
+ try:
+ result = fut.result()
+ self.add_to_cache(cache_key, result)
+ except Exception as e:
+ logger.error(f"Prefetch failed for key {cache_key}: {e}")
+ self.in_progress.pop(cache_key, None)
+
+ future.add_done_callback(callback)
+
+ def get_ai_help(self, instance_id: int, annotation_id: int, ai_assistant: str) -> str:
+ """retrieves AI help either from cache, waits for in-progress, or computes on-demand."""
+ key = (instance_id, annotation_id, ai_assistant)
+
+ # Check if caching is enabled for this help type
+ if not self.disk_cache_enabled:
+ return self.compute_help(instance_id, annotation_id, ai_assistant)
+
+ # Try to get from cache if caching is enabled
+ cached_value = self.get_from_cache(key)
+ if cached_value is not None:
+ logger.debug(f"Cache hit for key: {key}")
+ return cached_value
+
+ with self.lock:
+ if key in self.in_progress:
+ future = self.in_progress[key]
+ else:
+ future = self.executor.submit(self.compute_help, instance_id, annotation_id, ai_assistant)
+ self.in_progress[key] = future
+ try:
+ result = future.result(timeout=60)
+ # Don't cache error responses
+ is_error_response = (
+ isinstance(result, str) and
+ (result.startswith("Unable to generate") or
+ result.startswith("Error:") or
+ "error" in result.lower()[:50])
+ )
+ if self.disk_cache_enabled and not is_error_response:
+ self.add_to_cache(key, result)
+ elif is_error_response:
+ logger.warning(f"Not caching error response for key {key}: {result[:100]}")
+ with self.lock:
+ self.in_progress.pop(key, None)
+ return result
+ except Exception as e:
+ logger.error(f"Error computing help for key {key}: {e}")
+ with self.lock:
+ self.in_progress.pop(key, None)
+ return f"Error: {str(e)}"
+
+ def compute_help(self, instance_id: int, annotation_id: int, ai_assistant: str):
+ # Validate that the assistant type is compatible with the model and input
+ is_valid, error_message = self._validate_assistant_compatibility(
+ instance_id, annotation_id, ai_assistant
+ )
+ if not is_valid:
+ logger.warning(f"Assistant compatibility check failed: {error_message}")
+ return {"error": error_message}
+
+ annotation_type_str = config["annotation_schemes"][annotation_id]["annotation_type"]
+ annotation_type = Annotation_Type(annotation_type_str)
+ if annotation_type == Annotation_Type.LIKERT:
+ return self.generate_likert(instance_id, annotation_id, ai_assistant)
+ elif annotation_type == Annotation_Type.RADIO:
+ return self.generate_radio(instance_id, annotation_id, ai_assistant)
+ elif annotation_type == Annotation_Type.MULTISELECT:
+ return self.generate_multiselect(instance_id, annotation_id, ai_assistant)
+ elif annotation_type == Annotation_Type.NUMBER:
+ return self.generate_number(instance_id, annotation_id, ai_assistant)
+ elif annotation_type == Annotation_Type.SELECT:
+ return self.generate_select(instance_id, annotation_id, ai_assistant)
+ elif annotation_type == Annotation_Type.SLIDER:
+ return self.generate_slider(instance_id, annotation_id, ai_assistant)
+ elif annotation_type == Annotation_Type.SPAN:
+ return self.generate_span(instance_id, annotation_id, ai_assistant)
+ elif annotation_type == Annotation_Type.TEXTBOX:
+ return self.generate_textbox(instance_id, annotation_id, ai_assistant)
+ elif annotation_type == Annotation_Type.IMAGE_ANNOTATION:
+ return self.generate_image_annotation(instance_id, annotation_id, ai_assistant)
+ elif annotation_type == Annotation_Type.VIDEO_ANNOTATION:
+ return self.generate_video_annotation(instance_id, annotation_id, ai_assistant)
+ else:
+ raise ValueError(f"Unknown annotation type: {annotation_type}")
+
+ def get_cache_stats(self) -> Dict[str, int]:
+ """returns statistics on disk cache and in-progress cache entries."""
+ with self.lock:
+ disk_count = 0
+ if self.disk_cache_enabled and self.disk_persistence_path and os.path.exists(self.disk_persistence_path):
+ try:
+ disk_data = self.load_disk_cache_data(self.disk_persistence_path)
+ disk_count = len(disk_data)
+ except:
+ pass
+
+ return {
+ 'disk_cache_enabled': self.disk_cache_enabled,
+ 'cached_items_disk': disk_count,
+ 'in_progress_items': len(self.in_progress)
+ }
+
+ def clear_cache(self):
+ """clears disk cache and cancels any ongoing generation."""
+ with self.lock:
+ for future in self.in_progress.values():
+ future.cancel()
+ self.in_progress.clear()
+
+ if self.disk_cache_enabled and self.disk_persistence_path and os.path.exists(self.disk_persistence_path):
+ try:
+ os.remove(self.disk_persistence_path)
+ logger.info("Disk cache file removed")
+ except Exception as e:
+ logger.error(f"Error removing disk cache file: {e}")
+ logger.info("Cache cleared")
+
+
+
diff --git a/potato/ai/ai_endpoint.py b/potato/ai/ai_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..a967207d27d2c1b8014af397306768746f34c1dd
--- /dev/null
+++ b/potato/ai/ai_endpoint.py
@@ -0,0 +1,688 @@
+
+"""
+Unified AI endpoint interface for various LLM providers.
+
+This module provides a common interface for interacting with different LLM providers
+including OpenAI, Anthropic, Hugging Face, Ollama, and VLLM endpoints.
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+import logging
+from abc import ABC, abstractmethod
+import os
+from typing import Dict, Any, Optional, List, Type, Union
+import json
+from string import Template
+
+from pydantic import BaseModel
+
+from .ai_prompt import get_ai_prompt
+
+logger = logging.getLogger(__name__)
+
+
+class Annotation_Type(Enum):
+ RADIO = "radio"
+ LIKERT = "likert"
+ NUMBER = "number"
+ TEXTBOX = "text"
+ MULTISELECT = "multiselect"
+ SPAN = "span"
+ SELECT = "select"
+ SLIDER = "slider"
+ IMAGE_ANNOTATION = "image_annotation"
+ VIDEO_ANNOTATION = "video_annotation"
+
+
+@dataclass
+class ImageData:
+ """Data structure for image input to visual AI endpoints."""
+ source: str # 'url' | 'base64'
+ data: str # The URL or base64-encoded image data
+ width: Optional[int] = None
+ height: Optional[int] = None
+ mime_type: Optional[str] = None # e.g., 'image/jpeg', 'image/png'
+
+
+@dataclass
+class VisualAnnotationInput:
+ """Input data structure for visual annotation AI assistance."""
+ ai_assistant: str # 'detection', 'classification', 'hint', 'pre_annotate', etc.
+ annotation_type: str # 'image_annotation' | 'video_annotation'
+ task_type: str # Specific task: 'detection', 'classification', 'scene_detection', etc.
+ image_data: Union[ImageData, List[ImageData]] # Single image or list of frames
+ description: str # Task description from annotation scheme
+ labels: Optional[List[str]] = None # Available labels for the task
+ video_metadata: Optional[Dict[str, Any]] = field(default_factory=dict) # fps, duration for video
+ region: Optional[Dict[str, float]] = None # Selected region for classification (x, y, width, height)
+ confidence_threshold: float = 0.5 # Minimum confidence for detections
+
+
+@dataclass
+class AnnotationInput:
+ ai_assistant: str
+ annotation_type: Annotation_Type
+ text: str
+ description: str
+ min_label: Optional[str] = ""
+ max_label: Optional[str] = ""
+ size: Optional[int] = -1
+ labels: Optional[List[str]] = None
+ min_value: Optional[int] = -1
+ max_value: Optional[int] = -1
+ step: Optional[int] = -1
+
+
+@dataclass
+class ModelCapabilities:
+ """
+ Declares what operations an AI endpoint can perform.
+
+ This dataclass is used to define the capabilities of different AI endpoints,
+ enabling the system to automatically filter AI assistant buttons and validate
+ requests based on what each model can actually do.
+
+ Attributes:
+ text_generation: Can generate text (hints, rationales, descriptions)
+ vision_input: Can process images as input
+ bounding_box_output: Can output precise coordinate detections
+ text_classification: Can classify text into categories
+ image_classification: Can classify images into categories
+ rationale_generation: Can generate explanations/rationales for labels
+ keyword_extraction: Can extract keywords from text (not applicable to images)
+ """
+ text_generation: bool = False
+ vision_input: bool = False
+ bounding_box_output: bool = False
+ text_classification: bool = False
+ image_classification: bool = False
+ rationale_generation: bool = False
+ keyword_extraction: bool = False
+
+ def supports_assistant(self, assistant_type: str, has_image_input: bool = False) -> bool:
+ """
+ Check if model supports a specific AI assistant type.
+
+ Args:
+ assistant_type: The type of AI assistant ('hint', 'keyword', 'rationale',
+ 'detection', 'pre_annotate', 'classification')
+ has_image_input: Whether the current content is an image
+
+ Returns:
+ True if the model supports this assistant type for the given input type
+ """
+ if assistant_type == "hint":
+ # Hints require text generation; for images, also need vision
+ if has_image_input:
+ return self.text_generation and self.vision_input
+ return self.text_generation
+
+ elif assistant_type == "keyword":
+ # Keywords require keyword extraction AND text input (not images)
+ # Keyword highlighting doesn't make sense for images
+ return self.keyword_extraction and not has_image_input
+
+ elif assistant_type == "rationale":
+ # Rationales require rationale generation; for images, also need vision
+ if has_image_input:
+ return self.rationale_generation and self.vision_input
+ return self.rationale_generation
+
+ elif assistant_type in ("detection", "detect", "pre_annotate"):
+ # Detection requires vision and bounding box output
+ return self.bounding_box_output and self.vision_input
+
+ elif assistant_type == "classification":
+ # Classification depends on input type
+ if has_image_input:
+ return self.image_classification and self.vision_input
+ return self.text_classification
+
+ # Unknown assistant type - default to False for safety
+ return False
+
+ def get_supported_assistants(self, has_image_input: bool = False) -> List[str]:
+ """
+ Get list of assistant types supported for the given input type.
+
+ Args:
+ has_image_input: Whether the current content is an image
+
+ Returns:
+ List of supported assistant type names
+ """
+ all_types = ["hint", "keyword", "rationale", "detection", "pre_annotate", "classification"]
+ return [t for t in all_types if self.supports_assistant(t, has_image_input)]
+
+
+class AIEndpointError(Exception):
+ """Base exception for AI endpoint errors."""
+ pass
+
+
+class AIEndpointConfigError(AIEndpointError):
+ """Exception raised for configuration errors."""
+ pass
+
+
+class AIEndpointRequestError(AIEndpointError):
+ """Exception raised for request/API errors."""
+ pass
+
+
+class BaseAIEndpoint(ABC):
+ """
+ Abstract base class for AI endpoints.
+
+ All AI endpoint implementations should inherit from this class
+ and implement the required methods.
+ """
+
+ def __init__(self, config: Dict[str, Any]):
+ """
+ Initialize the AI endpoint with configuration.
+
+ Args:
+ config: Configuration dictionary containing endpoint-specific settings
+ """
+ self.config = config
+ self.description = config.get("description", "")
+ self.annotation_type = config.get("annotation_type", "")
+ self.ai_config = config.get("ai_config", {})
+
+ # Model configuration
+ self.model = self.ai_config.get("model", self._get_default_model())
+ self.max_tokens = self.ai_config.get("max_tokens", 100)
+ self.temperature = self.ai_config.get("temperature", 0.1)
+
+ # prompt
+ self.prompts = get_ai_prompt()
+
+ # Initialize the client
+ self._initialize_client()
+
+ @abstractmethod
+ def _initialize_client(self) -> None:
+ """Initialize the client for the specific AI provider."""
+ pass
+
+ @abstractmethod
+ def _get_default_model(self) -> str:
+ """Get the default model name for this provider."""
+ pass
+
+ @abstractmethod
+ def query(self, prompt: str, output_format: Type[BaseModel]):
+ """
+ Send a query to the AI model and return the response.
+
+ Args:
+ prompt: The prompt to send to the model
+
+ Returns:
+ The model's response as a string
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ pass
+
+ def chat_query_with_image(
+ self,
+ messages: List[Dict[str, Any]],
+ images: Optional[List["ImageData"]] = None,
+ ) -> str:
+ """
+ Send a multi-turn chat with interleaved images to the AI model.
+
+ Messages may contain content blocks (text + image) instead of plain strings.
+ Used by the live agent runner for vision-based agent loops.
+
+ Default implementation raises NotImplementedError โ only vision-capable
+ endpoints should override this.
+
+ Args:
+ messages: List of message dicts. 'content' may be a string or a list
+ of content blocks (e.g., {"type": "text", "text": "..."} or
+ {"type": "image", "source": {...}}).
+ images: Optional list of ImageData to include (alternative to inline images).
+
+ Returns:
+ The model's response as a plain text string.
+
+ Raises:
+ NotImplementedError: If the endpoint doesn't support vision.
+ """
+ raise NotImplementedError(
+ f"{self.__class__.__name__} does not support chat_query_with_image. "
+ f"Use a vision-capable endpoint (e.g., anthropic_vision)."
+ )
+
+ def chat_query(self, messages: List[Dict[str, str]]) -> str:
+ """
+ Send a multi-turn chat conversation to the AI model.
+
+ Default implementation flattens messages into a single prompt and calls query().
+ Subclasses should override with native multi-turn support.
+
+ Args:
+ messages: List of message dicts with 'role' and 'content' keys.
+ Roles: 'system', 'user', 'assistant'
+
+ Returns:
+ The model's response as a plain text string.
+ """
+ # Flatten messages into a single prompt
+ parts = []
+ for msg in messages:
+ role = msg.get("role", "user")
+ content = msg.get("content", "")
+ if role == "system":
+ parts.append(f"System: {content}")
+ elif role == "assistant":
+ parts.append(f"Assistant: {content}")
+ else:
+ parts.append(f"User: {content}")
+ prompt = "\n\n".join(parts) + "\n\nAssistant:"
+
+ try:
+ result = self.query(prompt)
+ # query() may return parsed JSON or a string; ensure we return a string
+ if isinstance(result, dict):
+ return result.get("response", result.get("content", str(result)))
+ return str(result)
+ except Exception as e:
+ raise AIEndpointRequestError(f"Chat query failed: {e}")
+
+ def parseStringToJson(self, response_content: str) -> str:
+ """
+ Parse structured output from any LLM response, with robust fallbacks.
+
+ Handles common issues across all endpoint types (ollama, vllm, openai, etc.):
+ 1. Clean JSON responses -> direct parse
+ 2. JSON wrapped in markdown code blocks (```json ... ```)
+ 3. JSON embedded in surrounding prose text
+ 4. Truncated JSON (max_tokens exceeded) -> extract complete key-value pairs
+ 5. Plain text with no JSON structure -> return as {"response": text}
+
+ Args:
+ response_content: Raw response string from the LLM
+
+ Returns:
+ Parsed dict or the raw string if JSON extraction succeeds
+
+ Raises:
+ ValueError: Only if response is completely empty
+ """
+ import re
+
+ # Handle empty or None content
+ if not response_content:
+ raise ValueError("Empty response content received from AI endpoint")
+
+ # If it's already a dict, return it
+ if isinstance(response_content, dict):
+ return response_content
+
+ # Convert to string if needed
+ content_str = str(response_content).strip()
+ if not content_str:
+ raise ValueError("Empty response content received from AI endpoint")
+
+ # Strategy 0: Strip thinking/reasoning blocks that wrap the actual output
+ # Many models (qwen3, deepseek, etc.) produce ... blocks
+ cleaned = content_str
+ for tag in ['think', 'thinking', 'thought', 'inner_monologue']:
+ cleaned = re.sub(
+ rf'<{tag}>[\s\S]*?{tag}>\s*',
+ '', cleaned, flags=re.IGNORECASE
+ ).strip()
+ if cleaned and cleaned != content_str:
+ content_str = cleaned
+
+ # Strategy 1: Try direct JSON parse
+ try:
+ return json.loads(content_str)
+ except json.JSONDecodeError:
+ pass
+
+ # Strategy 2: Extract from markdown code blocks
+ for pattern in [
+ r'```json\s*([\s\S]*?)\s*```',
+ r'```\s*([\s\S]*?)\s*```',
+ ]:
+ match = re.search(pattern, content_str)
+ if match:
+ try:
+ return json.loads(match.group(1).strip())
+ except json.JSONDecodeError:
+ pass
+
+ # Strategy 3: Extract from tool call format
+ # Some models return: {"name": "...", "arguments": {...}}
+ # or function_call blocks
+ for pattern in [
+ r'\s*([\s\S]*?)\s* ',
+ r'\s*([\s\S]*?)\s* ',
+ r'\s*([\s\S]*?)\s* ',
+ r'\s*([\s\S]*?)\s* ',
+ r'\s*([\s\S]*?)\s* ',
+ ]:
+ match = re.search(pattern, content_str, re.IGNORECASE)
+ if match:
+ try:
+ parsed = json.loads(match.group(1).strip())
+ # If it's a tool call wrapper, extract the arguments
+ if isinstance(parsed, dict) and 'arguments' in parsed:
+ return parsed['arguments']
+ return parsed
+ except json.JSONDecodeError:
+ pass
+
+ # Strategy 4: Find a JSON object anywhere in the text
+ # Greedy match for the outermost { ... }
+ match = re.search(r'\{[\s\S]*\}', content_str)
+ if match:
+ try:
+ return json.loads(match.group(0))
+ except json.JSONDecodeError:
+ pass
+
+ # Strategy 4: Salvage truncated JSON
+ # Extract complete "key": "value" and "key": number pairs
+ salvaged = self._salvage_key_value_pairs(content_str)
+ if salvaged:
+ logger.warning(
+ f"Salvaged {len(salvaged)} fields from truncated/malformed response"
+ )
+ return salvaged
+
+ # Strategy 6: Parse XML-style output
+ # Some models (especially larger ones) produce XML like:
+ # joy 90
+ # or joy
+ xml_result = self._parse_xml_to_dict(content_str)
+ if xml_result:
+ logger.info(
+ f"Parsed {len(xml_result)} fields from XML-style response"
+ )
+ return xml_result
+
+ # Strategy 7: Return raw text wrapped in a dict
+ logger.warning(
+ f"Could not parse JSON or XML from response ({len(content_str)} chars), "
+ f"returning as raw text"
+ )
+ return {"response": content_str}
+
+ @staticmethod
+ def _salvage_key_value_pairs(text: str) -> Optional[dict]:
+ """Extract key-value pairs from truncated or malformed JSON.
+
+ Handles cases where max_tokens cuts off a response mid-field, e.g.:
+ {"label": "joy", "confidence": 90, "reasoning": "The text expres...
+
+ Returns:
+ Dict of extracted key-value pairs, or None if nothing found.
+ """
+ import re
+ result = {}
+
+ # Extract "key": "value" pairs (string values)
+ for match in re.finditer(r'"(\w+)"\s*:\s*"([^"]*)"', text):
+ result[match.group(1)] = match.group(2)
+
+ # Extract "key": number pairs
+ for match in re.finditer(r'"(\w+)"\s*:\s*(-?\d+(?:\.\d+)?)\b', text):
+ key = match.group(1)
+ if key not in result:
+ try:
+ val = float(match.group(2))
+ result[key] = int(val) if val == int(val) else val
+ except ValueError:
+ pass
+
+ # Extract "key": true/false/null
+ for match in re.finditer(r'"(\w+)"\s*:\s*(true|false|null)\b', text):
+ key = match.group(1)
+ if key not in result:
+ val_str = match.group(2)
+ result[key] = (
+ True if val_str == 'true'
+ else False if val_str == 'false'
+ else None
+ )
+
+ return result if result else None
+
+ @staticmethod
+ def _parse_xml_to_dict(text: str) -> Optional[dict]:
+ """Extract key-value pairs from XML-style LLM output.
+
+ Handles patterns like:
+ - joy 90
+ - joy
+ - Mixed XML with text: "The emotion is joy "
+
+ Returns:
+ Dict of tag->content pairs, or None if no XML tags found.
+ """
+ import re
+ result = {}
+
+ # Find all content pairs (non-nested simple tags)
+ for match in re.finditer(
+ r'<(\w+)>([^<]*)\1>', text, re.IGNORECASE
+ ):
+ tag = match.group(1).lower()
+ value = match.group(2).strip()
+
+ # Skip wrapper tags that contain other tags
+ if tag in ('response', 'output', 'result', 'answer', 'root'):
+ continue
+
+ # Try to parse numeric values
+ try:
+ if '.' in value:
+ result[tag] = float(value)
+ else:
+ result[tag] = int(value)
+ except ValueError:
+ # Boolean
+ if value.lower() in ('true', 'false'):
+ result[tag] = value.lower() == 'true'
+ else:
+ result[tag] = value
+
+ return result if result else None
+
+ def get_ai(self, data: AnnotationInput, output_format) -> str:
+ """
+ Get a hint for annotating the given text.
+
+ Args:
+ text: The text to get a hint for
+
+ Returns:
+ A helpful hint for annotation
+ """
+
+ try:
+ # Check if annotation type exists (comparing string against enum values)
+ valid_types = [e.value for e in Annotation_Type]
+ if data.annotation_type not in valid_types:
+ logger.warning(f"Annotation type '{data.annotation_type}' not found")
+ return "Unable to generate suggestion - annotation type not configured"
+
+ # Check if ai_assistant exists
+ ai_prompt = get_ai_prompt()
+ if data.ai_assistant not in ai_prompt[data.annotation_type]:
+ logger.warning(f"'ai_assistant' not found for {data.annotation_type}")
+ return "Unable to generate suggestion - prompt not configured"
+
+ template_str = self.prompts.get(data.annotation_type).get(data.ai_assistant).get("prompt")
+ template = Template(template_str)
+ prompt = template.substitute(
+ text=data.text,
+ description=data.description,
+ min_label=data.min_label,
+ max_label=data.max_label,
+ size=data.size,
+ labels=data.labels,
+ min_value=data.min_value,
+ max_value=data.max_value,
+ step=data.step
+ )
+ return self.query(prompt, output_format)
+ except Exception as e:
+ logger.error(f"[get_ai] AnnotationInput: {data}")
+ logger.error(f"[get_ai] Error for {data.annotation_type}/{data.ai_assistant}: {type(e).__name__}: {e}")
+ import traceback
+ logger.error(f"[get_ai] Traceback:\n{traceback.format_exc()}")
+ return "Unable to generate hint at this time."
+
+ def health_check(self) -> bool:
+ """
+ Check if the AI endpoint is healthy and accessible.
+
+ Returns:
+ True if the endpoint is healthy, False otherwise
+ """
+ try:
+ # Simple test query
+ test_response = self.query("Hello")
+ return bool(test_response and test_response.strip())
+ except Exception as e:
+ logger.error(f"Health check failed: {e}")
+ return False
+
+
+class AIEndpointFactory:
+ """
+ Factory class for creating AI endpoint instances.
+ """
+
+ _endpoints = { }
+
+ @classmethod
+ def register_endpoint(cls, endpoint_type: str, endpoint_class: type):
+ """Register a new endpoint type."""
+ cls._endpoints[endpoint_type] = endpoint_class
+
+ @classmethod
+ def create_endpoint(cls, config: Dict[str, Any]) -> Optional[BaseAIEndpoint]:
+ """
+ Create an AI endpoint instance based on configuration.
+
+ Args:
+ config: Configuration dictionary containing ai_support settings
+
+ Returns:
+ An AI endpoint instance or None if AI support is disabled
+
+ Raises:
+ AIEndpointConfigError: If the configuration is invalid
+ """
+ if not config.get("ai_support", {}).get("enabled", False):
+ return None
+
+ ai_support = config["ai_support"]
+ endpoint_type = ai_support.get("endpoint_type")
+
+ if not endpoint_type:
+ raise AIEndpointConfigError("endpoint_type is required when ai_support is enabled")
+
+ if endpoint_type not in cls._endpoints:
+ raise AIEndpointConfigError(f"Unknown endpoint type: {endpoint_type}")
+
+ # Prepare endpoint configuration
+ endpoint_config = {
+ "ai_config": ai_support.get("ai_config", {})
+ }
+
+ try:
+ endpoint_class = cls._endpoints[endpoint_type]
+ return endpoint_class(endpoint_config)
+ except Exception as e:
+ raise AIEndpointConfigError(f"Failed to create {endpoint_type} endpoint: {e}")
+
+
+# Legacy function for backward compatibility
+def get_ai_endpoint(config: dict):
+ """
+ Get an AI endpoint instance (legacy function).
+
+ This function is maintained for backward compatibility.
+ New code should use AIEndpointFactory.create_endpoint().
+ """
+ return AIEndpointFactory.create_endpoint(config)
+
+
+# Register built-in endpoints
+try:
+ from .ollama_endpoint import OllamaEndpoint
+ AIEndpointFactory.register_endpoint("ollama", OllamaEndpoint)
+except ImportError:
+ logger.debug("Ollama endpoint not available")
+
+try:
+ from .openai_endpoint import OpenAIEndpoint
+ AIEndpointFactory.register_endpoint("openai", OpenAIEndpoint)
+except ImportError:
+ logger.debug("OpenAI endpoint not available")
+
+try:
+ from .huggingface_endpoint import HuggingfaceEndpoint
+ AIEndpointFactory.register_endpoint("huggingface", HuggingfaceEndpoint)
+except ImportError:
+ logger.debug("Hugging Face endpoint not available")
+
+try:
+ from .gemini_endpoint import GeminiEndpoint
+ AIEndpointFactory.register_endpoint("gemini", GeminiEndpoint)
+except ImportError:
+ logger.debug("Gemini endpoint not available")
+
+try:
+ from .anthropic_endpoint import AnthropicEndpoint
+ AIEndpointFactory.register_endpoint("anthropic", AnthropicEndpoint)
+except ImportError:
+ logger.debug("Anthropic endpoint not available")
+
+try:
+ from .vllm_endpoint import VLLMEndpoint
+ AIEndpointFactory.register_endpoint("vllm", VLLMEndpoint)
+except ImportError:
+ logger.debug("VLLM endpoint not available")
+
+# Register visual AI endpoints
+try:
+ from .yolo_endpoint import YOLOEndpoint
+ AIEndpointFactory.register_endpoint("yolo", YOLOEndpoint)
+except ImportError:
+ logger.debug("YOLO endpoint not available (ultralytics not installed)")
+
+try:
+ from .ollama_vision_endpoint import OllamaVisionEndpoint
+ AIEndpointFactory.register_endpoint("ollama_vision", OllamaVisionEndpoint)
+except ImportError:
+ logger.debug("Ollama Vision endpoint not available")
+
+try:
+ from .openai_vision_endpoint import OpenAIVisionEndpoint
+ AIEndpointFactory.register_endpoint("openai_vision", OpenAIVisionEndpoint)
+except ImportError:
+ logger.debug("OpenAI Vision endpoint not available")
+
+try:
+ from .anthropic_vision_endpoint import AnthropicVisionEndpoint
+ AIEndpointFactory.register_endpoint("anthropic_vision", AnthropicVisionEndpoint)
+except ImportError:
+ logger.debug("Anthropic Vision endpoint not available")
+
+try:
+ from .openrouter_endpoint import OpenRouterEndpoint
+ AIEndpointFactory.register_endpoint("openrouter", OpenRouterEndpoint)
+except ImportError:
+ logger.debug("OpenRouter endpoint not available")
diff --git a/potato/ai/ai_help_wrapper.py b/potato/ai/ai_help_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1f260c5cfceb76a9db54204f5101526c1002c91
--- /dev/null
+++ b/potato/ai/ai_help_wrapper.py
@@ -0,0 +1,203 @@
+from flask import render_template_string
+from typing import Optional, Dict, Any, List
+from potato.ai.ai_cache import get_ai_cache_manager, _is_image_url, _get_instance_text
+from potato.ai.ai_prompt import get_ai_prompt
+from potato.server_utils.config_module import config
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Global instance
+DYNAMICAIHELP = None
+
+def init_dynamic_ai_help():
+ import logging
+ logger = logging.getLogger(__name__)
+ logger.info(f"[init_dynamic_ai_help] Called. ai_support.enabled={config.get('ai_support', {}).get('enabled', False)}")
+
+ if not config["ai_support"]["enabled"]:
+ logger.info("[init_dynamic_ai_help] AI support disabled, returning")
+ return
+ global DYNAMICAIHELP
+ if DYNAMICAIHELP is None:
+ DYNAMICAIHELP = DynamicAIHelp()
+ logger.info(f"[init_dynamic_ai_help] Created DYNAMICAIHELP instance: {id(DYNAMICAIHELP)}")
+ else:
+ logger.info(f"[init_dynamic_ai_help] DYNAMICAIHELP already exists: {id(DYNAMICAIHELP)}")
+
+ return DYNAMICAIHELP
+
+def get_dynamic_ai_help():
+ global DYNAMICAIHELP
+ return DYNAMICAIHELP
+
+class DynamicAIHelp:
+ def __init__(self):
+ self.template = """
+ {% if ai_assistant %}
+ {{ ai_assistant | safe }}
+ {% elif error_message %}
+ {{ error_message }}
+ {% endif %}
+ """
+
+ def get_empty_wrapper(self):
+ return f''
+
+ def generate_ai_assistant(self, ai_prompts, annotation_type, ai_assistant):
+ str_html = f''
+ img_url = ai_prompts[annotation_type].get(ai_assistant).get("img")
+ if img_url:
+ # Use empty alt since the button already has a text label
+ str_html += f'
'
+ name = ai_prompts[annotation_type].get(ai_assistant).get("name", ai_assistant.capitalize())
+ str_html += f'
{name} '
+ str_html += "
"
+ return str_html
+
+ def _filter_assistants_by_capability(
+ self, ai_cache_manager, assistant_keys: List[str], is_image_content: bool
+ ) -> List[str]:
+ """
+ Filter assistant types based on model capabilities.
+
+ Args:
+ ai_cache_manager: The AI cache manager instance
+ assistant_keys: List of assistant type keys to filter
+ is_image_content: Whether the current content is an image
+
+ Returns:
+ Filtered list of assistant keys that the model supports
+ """
+ # Get capabilities from the cache manager
+ capabilities = ai_cache_manager.get_endpoint_capabilities(for_image=is_image_content)
+
+ filtered_keys = []
+ for key in assistant_keys:
+ if capabilities.supports_assistant(key, is_image_content):
+ filtered_keys.append(key)
+ else:
+ logger.debug(
+ f"[get_ai_help_data] Skipping '{key}' button - "
+ f"not supported for {'image' if is_image_content else 'text'} content"
+ )
+
+ return filtered_keys
+
+ def get_ai_help_data(self, instance: int, annotation_id: int, annotation_type: str) -> Dict[str, Any]:
+ """Get current AI help configuration with the new prompt structure"""
+ try:
+ context = {
+ 'ai_assistant': None,
+ 'error_message': None,
+ }
+ ai_prompts = get_ai_prompt()
+ logger.debug(f"[get_ai_help_data] ai_prompts keys: {list(ai_prompts.keys()) if ai_prompts else 'None'}")
+
+ if not ai_prompts:
+ context["error_message"] = f'No AI prompt configured'
+ logger.debug("[get_ai_help_data] No AI prompts configured")
+ return context
+ elif annotation_type not in ai_prompts:
+ context["error_message"] = f'annotation type {annotation_type} does not exist in ai_prompts'
+ logger.debug(f"[get_ai_help_data] annotation type {annotation_type} not in prompts")
+ return context
+
+ ai_cache_manager = get_ai_cache_manager()
+ logger.debug(f"[get_ai_help_data] ai_cache_manager: {ai_cache_manager is not None}")
+
+ if ai_cache_manager is None:
+ context["error_message"] = "AI cache manager not initialized"
+ logger.debug("[get_ai_help_data] AI cache manager is None")
+ return context
+
+ ai_assistant_html_parts = []
+
+ # Determine if content is an image for capability-based filtering
+ is_image_content = False
+ try:
+ text = _get_instance_text(instance)
+ is_image_content = _is_image_url(text)
+ if is_image_content:
+ logger.debug(f"[get_ai_help_data] Content is an image URL")
+ except Exception as e:
+ logger.debug(f"[get_ai_help_data] Could not determine if content is image: {e}")
+
+ # Check if user specified specific assistant types
+ special_include_types = ai_cache_manager.get_special_include(instance, annotation_id)
+ logger.debug(f"[get_ai_help_data] special_include_types: {special_include_types}")
+
+ if special_include_types:
+ # Generate HTML for specific included keys
+ logger.debug(f"[get_ai_help_data] Using special include types: {special_include_types}")
+ # Filter by capability
+ valid_keys = [k for k in special_include_types if k in ai_prompts[annotation_type]]
+ filtered_keys = self._filter_assistants_by_capability(
+ ai_cache_manager, valid_keys, is_image_content
+ )
+ for key in filtered_keys:
+ ai_assistant_html_parts.append(self.generate_ai_assistant(ai_prompts, annotation_type, key))
+
+ elif ai_cache_manager.get_include_all():
+ # Generate HTML for all keys in the annotation type
+ all_keys = list(ai_prompts[annotation_type].keys())
+ logger.debug(f"[get_ai_help_data] include_all=True, available keys: {all_keys}")
+ # Filter by capability
+ filtered_keys = self._filter_assistants_by_capability(
+ ai_cache_manager, all_keys, is_image_content
+ )
+ logger.debug(f"[get_ai_help_data] After capability filter: {filtered_keys}")
+ for key in filtered_keys:
+ ai_assistant_html_parts.append(self.generate_ai_assistant(ai_prompts, annotation_type, key))
+ else:
+ logger.debug("[get_ai_help_data] No special includes and include_all=False")
+
+ # Combine all HTML parts
+ ai_assistant_html = '| '.join(ai_assistant_html_parts) if ai_assistant_html_parts else None
+ logger.debug(f"[get_ai_help_data] ai_assistant_html_parts count: {len(ai_assistant_html_parts)}")
+
+ if ai_assistant_html:
+ context['ai_assistant'] = ai_assistant_html
+
+ logger.debug(f"[get_ai_help_data] Final context: ai_assistant={'set' if context['ai_assistant'] else 'None'}, error={'set' if context['error_message'] else 'None'}")
+ return context
+ except Exception as e:
+ logger.error(f"[get_ai_help_data] Exception: {e}", exc_info=True)
+ return {
+ 'ai_assistant': None,
+ 'error_message': f'Error loading AI help: {str(e)}',
+ }
+
+ def render(self, instance: int, annotation_id: int, annotation_type) -> str:
+ """Render AI help HTML with current data"""
+ context = self.get_ai_help_data(instance, annotation_id, annotation_type)
+ context.update({
+ 'instance': instance,
+ 'annotation_id': annotation_id
+ })
+ return render_template_string(self.template, **context)
+
+def generate_ai_help_html(instance: int, annotation_id: int, annotation_type: str) -> Optional[str]:
+ """
+ Generates dynamic AI help HTML using template rendering.
+ Now works with the new prompt structure: {annotation_type: {prompt: ..., outputformat: ...}}
+ """
+ import logging
+ logger = logging.getLogger(__name__)
+
+ if DYNAMICAIHELP is None:
+ logger.debug("[generate_ai_help_html] DYNAMICAIHELP is None - AI support not enabled")
+ return "" # AI support not enabled
+
+ result = DYNAMICAIHELP.render(instance, annotation_id, annotation_type)
+ logger.debug(f"[generate_ai_help_html] Rendered result: '{result[:100] if result else 'empty'}...'")
+ return result
+
+def get_ai_wrapper():
+ import logging
+ logger = logging.getLogger(__name__)
+ helper = get_dynamic_ai_help()
+ logger.debug(f"[get_ai_wrapper] DYNAMICAIHELP is {'set' if helper else 'None'}")
+ result = helper.get_empty_wrapper() if helper else ""
+ logger.debug(f"[get_ai_wrapper] Returning: '{result[:50] if result else 'empty'}...'")
+ return result
\ No newline at end of file
diff --git a/potato/ai/ai_prompt.py b/potato/ai/ai_prompt.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d55a5f8987bd38679fa263d772c77089e5e477b
--- /dev/null
+++ b/potato/ai/ai_prompt.py
@@ -0,0 +1,94 @@
+import importlib
+import json
+import os
+from pathlib import Path
+from typing import Optional, Type
+from pydantic import BaseModel
+from potato.server_utils.config_module import config
+ANNOTATIONS = None
+
+class ModelManager:
+ def __init__(self):
+ self.models_module = None
+
+ def load_models_module(self):
+ """Load the models module if not already loaded"""
+ if self.models_module is None:
+ # absolute pathing
+ module_path = config.get("ai_support").get("model_module")
+ if module_path:
+ file_path = Path(module_path)
+ if not file_path.exists():
+ raise FileNotFoundError(f"Model module file not found: {file_path}")
+ module_name = file_path.stem
+
+ spec = importlib.util.spec_from_file_location(module_name, file_path)
+ self.models_module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(self.models_module)
+
+ else:
+ default_path = Path(__file__).resolve().parent / "prompt" / "models_module.py"
+
+ if not default_path.exists():
+ raise FileNotFoundError(f"Default model module file not found: {default_path}")
+
+ module_name = default_path.stem
+
+ spec = importlib.util.spec_from_file_location(module_name, default_path)
+ self.models_module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(self.models_module)
+
+ return self.models_module
+
+ def get_model_class_by_name(self, name: str) -> Optional[Type[BaseModel]]:
+ """
+ Return a Pydantic model class based on the provided name.
+ """
+ models_module = self.load_models_module()
+ return models_module.CLASS_REGISTRY.get(name)
+
+
+def init_ai_prompt(config):
+ global ANNOTATIONS
+ if not config["ai_support"]["enabled"]:
+ return
+ try:
+ annotation_paths = config.get("ai_support", {}).get("annotation_path")
+
+ ANNOTATIONS = {}
+
+ if annotation_paths:
+ # Load files from specified paths
+ for key, path in annotation_paths.items():
+ if path and os.path.exists(path):
+ with open(path, "r", encoding="utf-8") as f:
+ ANNOTATIONS[key] = json.load(f)
+ else:
+ raise Exception(f"File path for annotations does not exist: {path}")
+ else:
+ # Load all JSON files from default directory (parent/prompt)
+ default_path = Path(__file__).resolve().parent / "prompt"
+
+ if default_path.exists() and default_path.is_dir():
+ # Find all JSON files in the directory
+ json_files = list(default_path.glob("*.json"))
+
+ if not json_files:
+ raise Exception(f"No JSON files found in default directory: {default_path}")
+
+ # Load each JSON file, using filename (without extension) as key
+ for file_path in json_files:
+ key = file_path.stem
+ with open(file_path, "r", encoding="utf-8") as f:
+ ANNOTATIONS[key] = json.load(f)
+ else:
+ raise Exception(f"Default annotation directory does not exist: {default_path}")
+
+ except json.JSONDecodeError as e:
+ raise ValueError(f"Invalid JSON in annotation file: {e}")
+ except Exception as e:
+ raise RuntimeError(f"Unexpected error loading AI prompt: {e}")
+
+def get_ai_prompt():
+ global ANNOTATIONS
+ return ANNOTATIONS
\ No newline at end of file
diff --git a/potato/ai/anthropic_endpoint.py b/potato/ai/anthropic_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..96fd349073a8e2c0f7598726f060fff3f3463462
--- /dev/null
+++ b/potato/ai/anthropic_endpoint.py
@@ -0,0 +1,118 @@
+"""
+Anthropic AI endpoint implementation.
+
+This module provides integration with Anthropic's Claude API for LLM inference.
+"""
+
+from typing import Dict, List
+import anthropic
+from .ai_endpoint import BaseAIEndpoint, AIEndpointRequestError, ModelCapabilities
+
+DEFAULT_MODEL = "claude-3-5-sonnet-20241022"
+DEFAULT_HINT_PROMPT = '''
+ You are assisting a user with an annotation task.
+ The annotation instruction is : {description}
+ The annotation task type is: {annotation_type}
+ The sentence (or item) to annotate is : {text}
+ Your goal is to generate a short, helpful hint that guides the annotator in how to think about the input โ **without providing the answer**.
+
+ The hint should:
+ - Highlight key aspects of the input relevant to the task
+ - Encourage thoughtful reasoning or observation
+ - Point to subtle features (tone, wording, structure, implication) that matter for the annotation
+ - Be specific and informative, not vague or generic
+ '''
+
+DEFAULT_KEYWORD_PROMPT = '''
+ You are assisting a user with an annotation task.
+ The annotation instruction is : {description}
+ The annotation task type is: {annotation_type}
+ The sentence (or item) to annotate is : {text}
+ Your goal is : Print out just a sequence of keywords, not sentences, in the text that most relate to the task. Do not explain your answer. Do not print out the entire text. If no part of the text relates to the task, print the empty string.
+ '''
+
+class AnthropicEndpoint(BaseAIEndpoint):
+ """Anthropic Claude endpoint for cloud-based LLM inference."""
+
+ # Capabilities declaration for text-based Anthropic Claude models
+ CAPABILITIES = ModelCapabilities(
+ text_generation=True,
+ vision_input=False,
+ bounding_box_output=False,
+ text_classification=True,
+ image_classification=False,
+ rationale_generation=True,
+ keyword_extraction=True,
+ )
+
+ def _initialize_client(self) -> None:
+ """Initialize the Anthropic client."""
+ api_key = self.ai_config.get("api_key", "")
+ if not api_key:
+ raise AIEndpointRequestError("Anthropic API key is required")
+
+ # Default timeout of 30 seconds, configurable via ai_config
+ timeout = self.ai_config.get("timeout", 30)
+ self.client = anthropic.Anthropic(api_key=api_key, timeout=timeout)
+
+ def _get_default_model(self) -> str:
+ """Get the default Anthropic model."""
+ return DEFAULT_MODEL
+
+ def _get_default_hint_prompt(self) -> str:
+ """Get the default hint prompt for Anthropic."""
+ return DEFAULT_HINT_PROMPT
+
+ def _get_default_keyword_prompt(self) -> str:
+ """Get the default keyword prompt for Anthropic."""
+ return DEFAULT_KEYWORD_PROMPT
+
+ def query(self, prompt: str) -> str:
+ """
+ Send a query to Anthropic Claude and return the response.
+
+ Args:
+ prompt: The prompt to send to the model
+
+ Returns:
+ The model's response as a string
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ try:
+ response = self.client.messages.create(
+ model=self.model,
+ max_tokens=self.max_tokens,
+ temperature=self.temperature,
+ messages=[{"role": "user", "content": prompt}]
+ )
+ return response.content[0].text
+ except Exception as e:
+ raise AIEndpointRequestError(f"Anthropic request failed: {e}")
+
+ def chat_query(self, messages: List[Dict[str, str]]) -> str:
+ """Send a multi-turn chat to Anthropic using native messages API."""
+ try:
+ # Extract system message if present
+ system_text = ""
+ chat_messages = []
+ for msg in messages:
+ if msg["role"] == "system":
+ system_text = msg["content"]
+ else:
+ chat_messages.append({"role": msg["role"], "content": msg["content"]})
+
+ kwargs = {
+ "model": self.model,
+ "max_tokens": self.max_tokens,
+ "temperature": self.temperature,
+ "messages": chat_messages,
+ }
+ if system_text:
+ kwargs["system"] = system_text
+
+ response = self.client.messages.create(**kwargs)
+ return response.content[0].text
+ except Exception as e:
+ raise AIEndpointRequestError(f"Anthropic chat request failed: {e}")
\ No newline at end of file
diff --git a/potato/ai/anthropic_vision_endpoint.py b/potato/ai/anthropic_vision_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8de262c9998b607943559bb902a94e4bc247a26
--- /dev/null
+++ b/potato/ai/anthropic_vision_endpoint.py
@@ -0,0 +1,405 @@
+"""
+Anthropic Vision AI Endpoint
+
+This module provides integration with Anthropic's Claude models for visual
+analysis using the image content block format.
+"""
+
+import base64
+import logging
+from typing import Any, Dict, List, Type, Union
+
+from pydantic import BaseModel
+
+from .ai_endpoint import AIEndpointRequestError, ImageData, ModelCapabilities
+from .visual_ai_endpoint import BaseVisualAIEndpoint
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = "claude-sonnet-4-20250514"
+
+# Supported image types for Claude
+SUPPORTED_MEDIA_TYPES = [
+ "image/jpeg",
+ "image/png",
+ "image/gif",
+ "image/webp"
+]
+
+
+class AnthropicVisionEndpoint(BaseVisualAIEndpoint):
+ """
+ Anthropic Vision endpoint for Claude models with vision capabilities.
+
+ Uses the image content block format for multimodal inputs.
+
+ Configuration options:
+ - model: Model to use (default: claude-sonnet-4-20250514)
+ - api_key: Anthropic API key (can also use ANTHROPIC_API_KEY env var)
+ - max_tokens: Maximum response tokens (default: 1024)
+ - temperature: Sampling temperature (default: 0.1)
+ """
+
+ # Capabilities declaration for Anthropic Claude vision models
+ # Claude models can understand images and generate detailed reasoning but bboxes are approximate
+ CAPABILITIES = ModelCapabilities(
+ text_generation=True,
+ vision_input=True,
+ bounding_box_output=False, # Claude bboxes are approximate, not precise
+ text_classification=True,
+ image_classification=True,
+ rationale_generation=True,
+ keyword_extraction=False, # Keywords don't apply to images
+ )
+
+ def _initialize_client(self) -> None:
+ """Initialize the Anthropic client."""
+ try:
+ import anthropic
+ except ImportError:
+ raise AIEndpointRequestError(
+ "anthropic package is required. Install it with: pip install anthropic"
+ )
+
+ import os
+
+ api_key = self.ai_config.get("api_key") or os.environ.get("ANTHROPIC_API_KEY")
+ if not api_key:
+ raise AIEndpointRequestError(
+ "Anthropic API key is required. Set it in config or ANTHROPIC_API_KEY env var."
+ )
+
+ timeout = self.ai_config.get("timeout", 60)
+
+ self.client = anthropic.Anthropic(api_key=api_key, timeout=timeout)
+ logger.info(f"Anthropic Vision client initialized with model: {self.model}")
+
+ def _get_default_model(self) -> str:
+ """Get the default Anthropic model."""
+ return DEFAULT_MODEL
+
+ def query(self, prompt: str, output_format: Type[BaseModel]) -> Any:
+ """
+ Standard text query without images.
+
+ Args:
+ prompt: Text prompt
+ output_format: Pydantic model for structured output
+
+ Returns:
+ Parsed response
+ """
+ try:
+ # Add JSON instruction to prompt
+ json_prompt = f"""{prompt}
+
+Please respond with valid JSON matching this schema:
+{output_format.model_json_schema()}"""
+
+ response = self.client.messages.create(
+ model=self.model,
+ max_tokens=self.max_tokens,
+ messages=[{"role": "user", "content": json_prompt}],
+ )
+
+ content = response.content[0].text
+ return self.parseStringToJson(content)
+
+ except Exception as e:
+ raise AIEndpointRequestError(f"Anthropic query failed: {e}")
+
+ def query_with_image(
+ self,
+ prompt: str,
+ image_data: Union[ImageData, List[ImageData]],
+ output_format: Type[BaseModel]
+ ) -> Any:
+ """
+ Send a query with image(s) to Claude vision model.
+
+ Args:
+ prompt: Text prompt describing what to analyze
+ image_data: Single ImageData or list of ImageData
+ output_format: Pydantic model for structured output
+
+ Returns:
+ Parsed response according to output_format
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ try:
+ # Prepare images
+ images = [image_data] if isinstance(image_data, ImageData) else image_data
+
+ # Build content array with images first, then text
+ content = []
+
+ for img in images:
+ image_block = self._build_image_block(img)
+ content.append(image_block)
+
+ # Add JSON instruction to prompt
+ json_prompt = f"""{prompt}
+
+Please respond with valid JSON matching this schema:
+{output_format.model_json_schema()}
+
+Only return the JSON object, no other text."""
+
+ content.append({"type": "text", "text": json_prompt})
+
+ # Make request
+ response = self.client.messages.create(
+ model=self.model,
+ max_tokens=self.max_tokens,
+ messages=[{"role": "user", "content": content}],
+ )
+
+ response_content = response.content[0].text
+ logger.debug(f"Anthropic vision response: {response_content[:500] if response_content else 'empty'}")
+
+ return self.parseStringToJson(response_content)
+
+ except AIEndpointRequestError:
+ raise
+ except Exception as e:
+ logger.error(f"Anthropic vision query failed: {e}")
+ import traceback
+ logger.error(traceback.format_exc())
+ raise AIEndpointRequestError(f"Anthropic vision query failed: {e}")
+
+ def chat_query_with_image(
+ self,
+ messages: List[Dict[str, Any]],
+ images: Any = None,
+ ) -> str:
+ """
+ Multi-turn chat with interleaved images for vision-based agent loops.
+
+ Messages may have 'content' as a string (text only) or a list of
+ content blocks (text + image dicts in Anthropic format).
+
+ Args:
+ messages: List of message dicts with 'role' and 'content'.
+ images: Unused (images are inline in messages).
+
+ Returns:
+ The model's response as a plain text string.
+ """
+ try:
+ system = ""
+ api_messages = []
+
+ for msg in messages:
+ if msg["role"] == "system":
+ system = msg["content"] if isinstance(msg["content"], str) else str(msg["content"])
+ else:
+ api_messages.append({
+ "role": msg["role"],
+ "content": msg["content"],
+ })
+
+ kwargs = {
+ "model": self.model,
+ "max_tokens": self.max_tokens,
+ "temperature": self.temperature,
+ "messages": api_messages,
+ }
+ if system:
+ kwargs["system"] = system
+
+ response = self.client.messages.create(**kwargs)
+ return response.content[0].text
+
+ except Exception as e:
+ logger.error(f"Anthropic vision chat query failed: {e}")
+ raise AIEndpointRequestError(f"Anthropic vision chat query failed: {e}")
+
+ def _build_image_block(self, image_data: ImageData) -> Dict[str, Any]:
+ """
+ Build image content block for Anthropic API.
+
+ Args:
+ image_data: ImageData object
+
+ Returns:
+ Dict with type: "image" and source content
+ """
+ if image_data.source == "url":
+ # Claude supports URL sources directly
+ return {
+ "type": "image",
+ "source": {
+ "type": "url",
+ "url": image_data.data
+ }
+ }
+
+ elif image_data.source == "base64":
+ # Determine media type
+ media_type = image_data.mime_type or "image/jpeg"
+
+ # Validate media type
+ if media_type not in SUPPORTED_MEDIA_TYPES:
+ logger.warning(f"Media type {media_type} may not be supported. Using image/jpeg.")
+ media_type = "image/jpeg"
+
+ return {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": media_type,
+ "data": image_data.data
+ }
+ }
+
+ else:
+ raise AIEndpointRequestError(f"Unknown image source: {image_data.source}")
+
+ def analyze_image(
+ self,
+ image_path_or_url: str,
+ prompt: str,
+ output_format: Type[BaseModel] = None
+ ) -> Any:
+ """
+ Convenience method for analyzing a single image.
+
+ Args:
+ image_path_or_url: Path to image file or URL
+ prompt: Analysis prompt
+ output_format: Optional output format model
+
+ Returns:
+ Analysis result
+ """
+ # Prepare image data
+ if image_path_or_url.startswith(("http://", "https://")):
+ # Claude can use URLs directly
+ image_data = self.create_url_image_data(image_path_or_url)
+ else:
+ image_data = self.encode_image_to_base64(image_path_or_url)
+
+ # Use a generic format if not specified
+ if output_format is None:
+ from .prompt.models_module import GeneralHintFormat
+ output_format = GeneralHintFormat
+
+ return self.query_with_image(prompt, image_data, output_format)
+
+ def detect_objects(
+ self,
+ image_path_or_url: str,
+ labels: List[str] = None
+ ) -> Dict[str, Any]:
+ """
+ Detect objects in an image and return bounding boxes.
+
+ Args:
+ image_path_or_url: Path to image file or URL
+ labels: Optional list of labels to detect
+
+ Returns:
+ Dict with detections list
+ """
+ from .prompt.models_module import VisualDetectionFormat
+
+ labels_str = ", ".join(labels) if labels else "all visible objects"
+
+ prompt = f"""Analyze this image and detect objects. For each object, provide:
+1. The label (from: {labels_str})
+2. A bounding box with normalized coordinates (0-1 range)
+3. Confidence score (0-1)
+
+Return a JSON object with this exact structure:
+{{
+ "detections": [
+ {{
+ "label": "object_name",
+ "bbox": {{"x": 0.1, "y": 0.2, "width": 0.3, "height": 0.4}},
+ "confidence": 0.95
+ }}
+ ]
+}}
+
+Important:
+- Coordinates are normalized (0-1) where x,y is the top-left corner
+- x increases left to right, y increases top to bottom
+- width and height are also normalized (0-1)
+- Only include objects you can clearly identify
+- Estimate bounding boxes as accurately as possible"""
+
+ # Prepare image
+ if image_path_or_url.startswith(("http://", "https://")):
+ image_data = self.create_url_image_data(image_path_or_url)
+ else:
+ image_data = self.encode_image_to_base64(image_path_or_url)
+
+ return self.query_with_image(prompt, image_data, VisualDetectionFormat)
+
+ def get_annotation_hint(
+ self,
+ image_path_or_url: str,
+ task_description: str,
+ labels: List[str]
+ ) -> Dict[str, Any]:
+ """
+ Get a hint for annotating an image without revealing exact locations.
+
+ Args:
+ image_path_or_url: Path to image file or URL
+ task_description: Description of the annotation task
+ labels: Available labels
+
+ Returns:
+ Dict with hint text and optional suggested label
+ """
+ labels_str = ", ".join(labels)
+
+ prompt = f"""You are helping an annotator with this task: {task_description}
+
+Available labels: {labels_str}
+
+Provide a helpful hint that guides the annotator without giving away the exact answer.
+The hint should:
+1. Point out relevant features to consider
+2. Suggest what to look for
+3. Not explicitly state the answer or exact locations
+
+Return JSON:
+{{
+ "hint": "Your helpful hint here",
+ "suggested_focus": "What area or aspect to focus on"
+}}"""
+
+ # Prepare image
+ if image_path_or_url.startswith(("http://", "https://")):
+ image_data = self.create_url_image_data(image_path_or_url)
+ else:
+ image_data = self.encode_image_to_base64(image_path_or_url)
+
+ class HintFormat(BaseModel):
+ hint: str
+ suggested_focus: str
+
+ return self.query_with_image(prompt, image_data, HintFormat)
+
+ def health_check(self) -> bool:
+ """
+ Check if the Anthropic API is accessible.
+
+ Returns:
+ True if API is reachable, False otherwise
+ """
+ try:
+ # Simple test message
+ self.client.messages.create(
+ model=self.model,
+ max_tokens=10,
+ messages=[{"role": "user", "content": "Hello"}]
+ )
+ return True
+ except Exception as e:
+ logger.error(f"Anthropic health check failed: {e}")
+ return False
diff --git a/potato/ai/gemini_endpoint.py b/potato/ai/gemini_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..5a9d145f7cf6915a8f430ec0d36faee23fe50415
--- /dev/null
+++ b/potato/ai/gemini_endpoint.py
@@ -0,0 +1,58 @@
+"""
+Google Gemini AI endpoint implementation.
+
+This module provides integration with Google's Gemini API for LLM inference.
+"""
+
+from google import genai
+from .ai_endpoint import BaseAIEndpoint, AIEndpointRequestError
+
+DEFAULT_MODEL = "gemini-2.0-flash-exp"
+
+
+class GeminiEndpoint(BaseAIEndpoint):
+ """Google Gemini endpoint for cloud-based LLM inference."""
+
+ def _initialize_client(self) -> None:
+ """Initialize the Gemini client."""
+ api_key = self.ai_config.get("api_key", "")
+ if not api_key:
+ raise AIEndpointRequestError("Gemini API key is required")
+
+ # Default timeout of 30 seconds, configurable via ai_config
+ timeout = self.ai_config.get("timeout", 30)
+ self.client = genai.Client(
+ api_key=api_key,
+ http_options={'timeout': timeout}
+ )
+
+ def _get_default_model(self) -> str:
+ """Get the default Gemini model."""
+ return DEFAULT_MODEL
+
+ def query(self, prompt: str, prompt_format: dict) -> str:
+ """
+ Send a query to Gemini and return the response.
+
+ Args:
+ prompt: The prompt to send to the model
+
+ Returns:
+ The model's response as a string
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ try:
+ response = self.client.models.generate_content(
+ model=self.model,
+ contents=prompt,
+ generation_config={
+ 'max_output_tokens': self.max_tokens,
+ 'temperature': self.temperature,
+ 'response_schema': prompt_format.model_json_schema(),
+ }
+ )
+ return response.text
+ except Exception as e:
+ raise AIEndpointRequestError(f"Gemini request failed: {e}")
diff --git a/potato/ai/huggingface_endpoint.py b/potato/ai/huggingface_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..aa496c337b51a7c980df9f4b844aa786bc24ee68
--- /dev/null
+++ b/potato/ai/huggingface_endpoint.py
@@ -0,0 +1,62 @@
+"""
+Hugging Face AI endpoint implementation.
+
+This module provides integration with Hugging Face's Inference API for LLM inference.
+"""
+
+from huggingface_hub import InferenceClient
+from .ai_endpoint import BaseAIEndpoint, AIEndpointRequestError
+
+DEFAULT_MODEL = "meta-llama/Llama-3.2-3B-Instruct"
+
+class HuggingfaceEndpoint(BaseAIEndpoint):
+ """Hugging Face endpoint for cloud-based LLM inference."""
+
+ def _initialize_client(self) -> None:
+ """Initialize the Hugging Face client."""
+ api_key = self.ai_config.get("api_key", "")
+ if not api_key:
+ raise AIEndpointRequestError("Hugging Face API key is required")
+
+ # Default timeout of 30 seconds, configurable via ai_config
+ timeout = self.ai_config.get("timeout", 30)
+ self.client = InferenceClient(
+ model=self.model,
+ token=api_key,
+ timeout=timeout
+ )
+
+ def _get_default_model(self) -> str:
+ """Get the default Hugging Face model."""
+ return DEFAULT_MODEL
+
+ def query(self, prompt: str, output_format: dict) -> str:
+ """
+ Send a query to Hugging Face and return the response.
+
+ Args:
+ prompt: The prompt to send to the model
+
+ Returns:
+ The model's response as a string
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ try:
+ response = self.client.chat_completion(
+ messages=[{"role": "user", "content": prompt}],
+ max_tokens=self.max_tokens,
+ temperature=self.temperature,
+ response_format= {
+ "type": "json_schema",
+ "json_schema": {
+ "name": "output_format",
+ "schema": output_format.model_json_schema(),
+ "strict": True,
+ }
+ }
+ )
+ return response.choices[0].message.content
+ except Exception as e:
+ raise AIEndpointRequestError(f"Hugging Face request failed: {e}")
diff --git a/potato/ai/icl_labeler.py b/potato/ai/icl_labeler.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d4fb1900c8736840a8c46cb40e3d3d5c0efe8f8
--- /dev/null
+++ b/potato/ai/icl_labeler.py
@@ -0,0 +1,1110 @@
+"""
+In-Context Learning (ICL) Labeler Module
+
+This module provides AI-assisted labeling using high-confidence human annotations
+as in-context examples to prompt an LLM to label remaining data.
+
+Key features:
+- Identifies high-confidence examples where annotators agree
+- Uses examples as in-context demonstrations for LLM labeling
+- Tracks LLM confidence scores on predictions
+- Routes subset of LLM labels to humans for verification
+- Calculates and reports LLM accuracy based on verification
+"""
+
+import json
+import logging
+import os
+import random
+import threading
+import time
+from collections import Counter, defaultdict
+from dataclasses import dataclass, field, asdict
+from datetime import datetime
+from typing import Dict, List, Optional, Any, Tuple, Set
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class HighConfidenceExample:
+ """A human-annotated example suitable for in-context learning."""
+ instance_id: str
+ text: str
+ schema_name: str
+ label: str
+ agreement_score: float # Proportion of annotators who chose this label
+ annotator_count: int
+ timestamp: datetime = field(default_factory=datetime.now)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'instance_id': self.instance_id,
+ 'text': self.text,
+ 'schema_name': self.schema_name,
+ 'label': self.label,
+ 'agreement_score': self.agreement_score,
+ 'annotator_count': self.annotator_count,
+ 'timestamp': self.timestamp.isoformat()
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'HighConfidenceExample':
+ """Deserialize from dictionary."""
+ timestamp = data.get('timestamp')
+ if isinstance(timestamp, str):
+ timestamp = datetime.fromisoformat(timestamp)
+ elif timestamp is None:
+ timestamp = datetime.now()
+
+ return cls(
+ instance_id=data['instance_id'],
+ text=data['text'],
+ schema_name=data['schema_name'],
+ label=data['label'],
+ agreement_score=data['agreement_score'],
+ annotator_count=data['annotator_count'],
+ timestamp=timestamp
+ )
+
+
+@dataclass
+class ICLPrediction:
+ """Record of an LLM prediction using in-context learning."""
+ instance_id: str
+ schema_name: str
+ predicted_label: str
+ confidence_score: float # 0.0-1.0
+ timestamp: datetime = field(default_factory=datetime.now)
+
+ # In-context examples used
+ example_instance_ids: List[str] = field(default_factory=list)
+
+ # Verification tracking
+ verification_status: str = 'pending' # 'pending', 'verified_correct', 'verified_incorrect'
+ verified_by: Optional[str] = None
+ verified_at: Optional[datetime] = None
+ human_label: Optional[str] = None # Human's label if verified
+
+ # LLM metadata
+ model_name: str = ""
+ reasoning: str = ""
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'instance_id': self.instance_id,
+ 'schema_name': self.schema_name,
+ 'predicted_label': self.predicted_label,
+ 'confidence_score': self.confidence_score,
+ 'timestamp': self.timestamp.isoformat(),
+ 'example_instance_ids': self.example_instance_ids,
+ 'verification_status': self.verification_status,
+ 'verified_by': self.verified_by,
+ 'verified_at': self.verified_at.isoformat() if self.verified_at else None,
+ 'human_label': self.human_label,
+ 'model_name': self.model_name,
+ 'reasoning': self.reasoning
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'ICLPrediction':
+ """Deserialize from dictionary."""
+ timestamp = data.get('timestamp')
+ if isinstance(timestamp, str):
+ timestamp = datetime.fromisoformat(timestamp)
+ elif timestamp is None:
+ timestamp = datetime.now()
+
+ verified_at = data.get('verified_at')
+ if isinstance(verified_at, str):
+ verified_at = datetime.fromisoformat(verified_at)
+
+ return cls(
+ instance_id=data['instance_id'],
+ schema_name=data['schema_name'],
+ predicted_label=data['predicted_label'],
+ confidence_score=data['confidence_score'],
+ timestamp=timestamp,
+ example_instance_ids=data.get('example_instance_ids', []),
+ verification_status=data.get('verification_status', 'pending'),
+ verified_by=data.get('verified_by'),
+ verified_at=verified_at,
+ human_label=data.get('human_label'),
+ model_name=data.get('model_name', ''),
+ reasoning=data.get('reasoning', '')
+ )
+
+
+class ICLLabeler:
+ """
+ Manages in-context learning based labeling using high-confidence human annotations.
+
+ Workflow:
+ 1. Monitors annotation progress for high-confidence examples
+ 2. Periodically refreshes pool of high-confidence examples
+ 3. Uses examples to prompt LLM for labeling unlabeled instances
+ 4. Routes some LLM-labeled instances for human verification (blind)
+ 5. Tracks accuracy metrics
+ """
+
+ _instance = None
+ _lock = threading.RLock()
+
+ def __new__(cls, *args, **kwargs):
+ """Singleton pattern."""
+ if cls._instance is None:
+ with cls._lock:
+ if cls._instance is None:
+ cls._instance = super().__new__(cls)
+ cls._instance._initialized = False
+ return cls._instance
+
+ def __init__(self, config: Optional[Dict[str, Any]] = None):
+ """
+ Initialize the ICLLabeler.
+
+ Args:
+ config: Configuration dictionary with settings
+ """
+ if self._initialized:
+ return
+
+ self.config = config or {}
+ self._ai_endpoint = None
+
+ # Get ICL labeling config
+ icl_config = self.config.get('icl_labeling', {})
+
+ # Example selection config
+ example_config = icl_config.get('example_selection', {})
+ self.min_agreement_threshold = example_config.get('min_agreement_threshold', 0.8)
+ self.min_annotators_per_instance = example_config.get('min_annotators_per_instance', 2)
+ self.max_examples_per_schema = example_config.get('max_examples_per_schema', 10)
+ self.example_refresh_interval = example_config.get('refresh_interval_seconds', 300)
+
+ # LLM labeling config
+ llm_config = icl_config.get('llm_labeling', {})
+ self.batch_size = llm_config.get('batch_size', 20)
+ self.trigger_threshold = llm_config.get('trigger_threshold', 5)
+ self.confidence_threshold = llm_config.get('confidence_threshold', 0.7)
+ self.batch_interval = llm_config.get('batch_interval_seconds', 600)
+
+ # Limits to prevent labeling entire dataset at once
+ # This allows iterative improvement - verify accuracy before labeling more
+ self.max_total_labels = llm_config.get('max_total_labels', None) # Max instances to label total
+ self.max_unlabeled_ratio = llm_config.get('max_unlabeled_ratio', 0.5) # Max % of unlabeled to label
+ self.pause_on_low_accuracy = llm_config.get('pause_on_low_accuracy', True)
+ self.min_accuracy_threshold = llm_config.get('min_accuracy_threshold', 0.7) # Pause if accuracy below
+
+ # Verification config
+ verification_config = icl_config.get('verification', {})
+ self.verification_enabled = verification_config.get('enabled', True)
+ self.verification_sample_rate = verification_config.get('sample_rate', 0.2)
+ self.verification_strategy = verification_config.get('selection_strategy', 'low_confidence')
+
+ # Persistence config
+ persistence_config = icl_config.get('persistence', {})
+ self.predictions_file = persistence_config.get('predictions_file', 'icl_predictions.json')
+
+ # State
+ self.schema_to_examples: Dict[str, List[HighConfidenceExample]] = {}
+ self.predictions: Dict[str, Dict[str, ICLPrediction]] = {} # instance_id -> schema -> prediction
+ self.verification_queue: List[Tuple[str, str]] = [] # [(instance_id, schema_name), ...]
+ self.labeled_instance_ids: Set[str] = set() # Instances labeled by LLM
+
+ self.last_example_refresh: Optional[datetime] = None
+ self.last_batch_run: Optional[datetime] = None
+
+ # Background worker
+ self._worker_thread: Optional[threading.Thread] = None
+ self._stop_worker = threading.Event()
+
+ self._initialized = True
+ logger.info("ICLLabeler initialized")
+
+ def _get_ai_endpoint(self):
+ """Get or create AI endpoint from config (reuses ai_support config)."""
+ if self._ai_endpoint is None:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+ self._ai_endpoint = AIEndpointFactory.create_endpoint(self.config)
+ return self._ai_endpoint
+
+ def _get_annotation_schemes(self) -> List[Dict[str, Any]]:
+ """Get annotation schemes from config."""
+ return self.config.get('annotation_schemes', [])
+
+ def _get_text_key(self) -> str:
+ """Get the text key from item_properties."""
+ return self.config.get('item_properties', {}).get('text_key', 'text')
+
+ # === High-Confidence Example Collection ===
+
+ def refresh_high_confidence_examples(self) -> Dict[str, List[HighConfidenceExample]]:
+ """
+ Scan annotations and identify high-confidence examples.
+
+ Returns:
+ Dictionary mapping schema name to list of high-confidence examples
+ """
+ from potato.flask_server import get_users, get_user_state, get_item_state_manager
+
+ with self._lock:
+ new_examples: Dict[str, List[HighConfidenceExample]] = defaultdict(list)
+
+ try:
+ ism = get_item_state_manager()
+ if ism is None:
+ logger.warning("ItemStateManager not available")
+ return new_examples
+
+ text_key = self._get_text_key()
+ schemas = self._get_annotation_schemes()
+ schema_names = [s.get('name') for s in schemas if s.get('name')]
+
+ # Collect all annotations per instance
+ instance_annotations: Dict[str, Dict[str, List[Tuple[str, Any]]]] = defaultdict(
+ lambda: defaultdict(list)
+ ) # instance_id -> schema_name -> [(user_id, value), ...]
+
+ for username in get_users():
+ user_state = get_user_state(username)
+ if not user_state:
+ continue
+
+ all_annotations = user_state.get_all_annotations()
+ for instance_id, instance_data in all_annotations.items():
+ if 'labels' not in instance_data:
+ continue
+
+ for label, value in instance_data['labels'].items():
+ schema_name = label.get_schema() if hasattr(label, 'get_schema') else str(label)
+ if schema_name in schema_names:
+ instance_annotations[instance_id][schema_name].append((username, value))
+
+ # Find high-confidence examples
+ for instance_id, schema_data in instance_annotations.items():
+ for schema_name, annotations in schema_data.items():
+ annotator_count = len(annotations)
+
+ if annotator_count < self.min_annotators_per_instance:
+ continue
+
+ # Count votes per label
+ label_counts = Counter(value for _, value in annotations)
+ most_common_label, most_common_count = label_counts.most_common(1)[0]
+
+ # Calculate agreement
+ agreement_score = most_common_count / annotator_count
+
+ if agreement_score >= self.min_agreement_threshold:
+ # Get instance text
+ item = ism.get_item(instance_id)
+ instance_data = item.get_data() if item else None
+ if instance_data is None:
+ continue
+
+ text = instance_data.get(text_key, '')
+ if not text:
+ continue
+
+ example = HighConfidenceExample(
+ instance_id=instance_id,
+ text=text,
+ schema_name=schema_name,
+ label=str(most_common_label),
+ agreement_score=agreement_score,
+ annotator_count=annotator_count
+ )
+ new_examples[schema_name].append(example)
+
+ # Select examples using coverage-based selection (CoverICL-inspired)
+ # or fall back to agreement-score sorting
+ for schema_name in new_examples:
+ candidates = new_examples[schema_name]
+ if len(candidates) > self.max_examples_per_schema:
+ selected = self._select_diverse_examples(
+ candidates, self.max_examples_per_schema
+ )
+ new_examples[schema_name] = selected
+ else:
+ new_examples[schema_name].sort(
+ key=lambda x: x.agreement_score, reverse=True
+ )
+
+ self.schema_to_examples = dict(new_examples)
+ self.last_example_refresh = datetime.now()
+
+ total_examples = sum(len(examples) for examples in new_examples.values())
+ logger.info(f"Refreshed high-confidence examples: {total_examples} examples across {len(new_examples)} schemas")
+
+ except Exception as e:
+ logger.error(f"Error refreshing examples: {e}")
+
+ return self.schema_to_examples
+
+ def get_examples_for_schema(self, schema_name: str) -> List[HighConfidenceExample]:
+ """Get high-confidence examples for a specific schema."""
+ return self.schema_to_examples.get(schema_name, [])
+
+ def has_enough_examples(self, schema_name: str) -> bool:
+ """Check if we have enough examples to start labeling."""
+ return len(self.get_examples_for_schema(schema_name)) >= self.trigger_threshold
+
+ def _select_diverse_examples(
+ self,
+ candidates: List[HighConfidenceExample],
+ k: int,
+ ) -> List[HighConfidenceExample]:
+ """CoverICL-inspired coverage-based example selection.
+
+ Uses greedy facility location to select examples that maximize
+ coverage of the instance embedding space, ensuring diverse and
+ representative ICL demonstrations.
+
+ Inspired by Mavromatis et al. (2024) CoverICL. Falls back to
+ agreement-score sorting if vectorization fails.
+
+ Args:
+ candidates: Pool of high-confidence examples
+ k: Number of examples to select
+
+ Returns:
+ List of selected examples maximizing coverage
+ """
+ if len(candidates) <= k:
+ return candidates
+
+ try:
+ from sklearn.feature_extraction.text import TfidfVectorizer
+ from sklearn.metrics.pairwise import cosine_distances
+
+ texts = [c.text for c in candidates]
+ vectorizer = TfidfVectorizer(max_features=5000)
+ features = vectorizer.fit_transform(texts).toarray()
+
+ # Greedy facility location: iteratively pick the candidate that
+ # maximizes the minimum distance to already-selected examples
+ n = len(candidates)
+ selected_indices = []
+
+ # Start with the candidate that has highest agreement score
+ # (quality-weighted seed)
+ agreements = [c.agreement_score for c in candidates]
+ first = int(max(range(n), key=lambda i: agreements[i]))
+ selected_indices.append(first)
+
+ dist_matrix = cosine_distances(features)
+
+ for _ in range(k - 1):
+ # For each unselected candidate, compute min distance to selected set
+ best_idx = -1
+ best_score = -1.0
+
+ for i in range(n):
+ if i in selected_indices:
+ continue
+ min_dist = min(dist_matrix[i][j] for j in selected_indices)
+ # Weight by agreement score for quality-aware selection
+ score = min_dist * candidates[i].agreement_score
+ if score > best_score:
+ best_score = score
+ best_idx = i
+
+ if best_idx >= 0:
+ selected_indices.append(best_idx)
+ else:
+ break
+
+ selected = [candidates[i] for i in selected_indices]
+ logger.debug(f"CoverICL selection: {len(selected)} diverse examples from {n} candidates")
+ return selected
+
+ except Exception as e:
+ logger.warning(f"Coverage-based selection failed, using agreement sorting: {e}")
+ candidates.sort(key=lambda x: x.agreement_score, reverse=True)
+ return candidates[:k]
+
+ # === LLM Labeling ===
+
+ def label_instance(
+ self,
+ instance_id: str,
+ schema_name: str,
+ instance_text: str
+ ) -> Optional[ICLPrediction]:
+ """
+ Label a single instance using in-context learning.
+
+ Args:
+ instance_id: The instance to label
+ schema_name: The annotation schema to use
+ instance_text: The text to label
+
+ Returns:
+ ICLPrediction if successful, None otherwise
+ """
+ from potato.ai.icl_prompt_builder import ICLPromptBuilder
+
+ examples = self.get_examples_for_schema(schema_name)
+ if not examples:
+ logger.warning(f"No examples available for schema {schema_name}")
+ return None
+
+ # Get schema info
+ schemas = self._get_annotation_schemes()
+ schema_info = next((s for s in schemas if s.get('name') == schema_name), None)
+ if not schema_info:
+ logger.warning(f"Schema {schema_name} not found in config")
+ return None
+
+ endpoint = self._get_ai_endpoint()
+ if endpoint is None:
+ logger.warning("AI endpoint not available")
+ return None
+
+ try:
+ # Build prompt
+ prompt_builder = ICLPromptBuilder()
+ prompt = prompt_builder.build_prompt(
+ schema=schema_info,
+ examples=examples,
+ target_text=instance_text
+ )
+
+ # Query LLM
+ from pydantic import BaseModel
+
+ class ICLResponse(BaseModel):
+ label: str
+ confidence: float
+ reasoning: str = ""
+
+ response = endpoint.query(prompt, ICLResponse)
+
+ # Parse response
+ if isinstance(response, str):
+ response_data = json.loads(response)
+ elif hasattr(response, 'model_dump'):
+ response_data = response.model_dump()
+ else:
+ response_data = response
+
+ predicted_label = response_data.get('label', '')
+ confidence = float(response_data.get('confidence', 0.5))
+ reasoning = response_data.get('reasoning', '')
+
+ # Validate label against schema
+ valid_labels = self._get_valid_labels(schema_info)
+ if valid_labels and predicted_label not in valid_labels:
+ # Try fuzzy matching
+ predicted_label = self._fuzzy_match_label(predicted_label, valid_labels)
+ if predicted_label is None:
+ logger.warning(f"LLM returned invalid label for {instance_id}")
+ return None
+
+ # Create prediction
+ prediction = ICLPrediction(
+ instance_id=instance_id,
+ schema_name=schema_name,
+ predicted_label=predicted_label,
+ confidence_score=min(1.0, max(0.0, confidence)),
+ example_instance_ids=[e.instance_id for e in examples],
+ model_name=endpoint.model if hasattr(endpoint, 'model') else '',
+ reasoning=reasoning
+ )
+
+ # Store prediction
+ with self._lock:
+ if instance_id not in self.predictions:
+ self.predictions[instance_id] = {}
+ self.predictions[instance_id][schema_name] = prediction
+ self.labeled_instance_ids.add(instance_id)
+
+ # Maybe add to verification queue
+ if self.verification_enabled and random.random() < self.verification_sample_rate:
+ self.verification_queue.append((instance_id, schema_name))
+
+ logger.debug(f"Labeled {instance_id} with {predicted_label} (confidence: {confidence:.2f})")
+ return prediction
+
+ except Exception as e:
+ logger.error(f"Error labeling instance {instance_id}: {e}")
+ return None
+
+ def _get_valid_labels(self, schema_info: Dict[str, Any]) -> List[str]:
+ """Extract valid labels from schema info."""
+ labels = schema_info.get('labels', [])
+ valid_labels = []
+ for label in labels:
+ if isinstance(label, str):
+ valid_labels.append(label)
+ elif isinstance(label, dict):
+ valid_labels.append(label.get('name', str(label)))
+ return valid_labels
+
+ def _fuzzy_match_label(self, predicted: str, valid_labels: List[str]) -> Optional[str]:
+ """Try to match predicted label to a valid label."""
+ predicted_lower = predicted.lower().strip()
+ for label in valid_labels:
+ if label.lower().strip() == predicted_lower:
+ return label
+ return None
+
+ def should_pause_labeling(self) -> Tuple[bool, str]:
+ """
+ Check if labeling should be paused based on limits and accuracy.
+
+ Returns:
+ Tuple of (should_pause, reason)
+ """
+ # Check if max total labels reached
+ if self.max_total_labels is not None:
+ current_count = len(self.labeled_instance_ids)
+ if current_count >= self.max_total_labels:
+ return True, f"Reached max_total_labels limit ({self.max_total_labels})"
+
+ # Check accuracy threshold
+ if self.pause_on_low_accuracy:
+ metrics = self.get_accuracy_metrics()
+ total_verified = metrics.get('total_verified', 0)
+ accuracy = metrics.get('accuracy')
+
+ # Only check accuracy if we have enough verifications
+ min_verifications = 10
+ if total_verified >= min_verifications and accuracy is not None:
+ if accuracy < self.min_accuracy_threshold:
+ return True, f"Accuracy ({accuracy:.1%}) below threshold ({self.min_accuracy_threshold:.1%})"
+
+ return False, ""
+
+ def get_remaining_label_capacity(self) -> int:
+ """
+ Get how many more instances can be labeled.
+
+ Returns:
+ Number of instances that can still be labeled, or -1 for unlimited
+ """
+ from potato.item_state_management import get_item_state_manager
+ from potato.user_state_management import get_user_state_manager
+
+ try:
+ ism = get_item_state_manager()
+ except ValueError:
+ # ISM not initialized yet
+ return 0
+ if ism is None:
+ return 0
+
+ # Count unlabeled instances (not labeled by humans or LLM)
+ unlabeled_count = 0
+ usm = get_user_state_manager()
+ for instance_id in ism.instance_id_ordering:
+ if instance_id in self.labeled_instance_ids:
+ continue
+
+ has_human_annotation = False
+ for username in usm.get_all_users():
+ user_state = usm.get_user_state(username)
+ if user_state:
+ all_annotations = user_state.get_all_annotations()
+ if instance_id in all_annotations:
+ has_human_annotation = True
+ break
+
+ if not has_human_annotation:
+ unlabeled_count += 1
+
+ current_llm_labels = len(self.labeled_instance_ids)
+
+ # Calculate max based on ratio
+ max_from_ratio = int(unlabeled_count * self.max_unlabeled_ratio)
+
+ # Calculate max based on total limit
+ if self.max_total_labels is not None:
+ max_from_total = self.max_total_labels - current_llm_labels
+ return min(max_from_ratio, max_from_total)
+
+ return max_from_ratio
+
+ def batch_label_instances(self, schema_name: str) -> List[ICLPrediction]:
+ """
+ Label multiple unlabeled instances for a schema.
+
+ Respects configured limits to prevent labeling entire dataset at once.
+
+ Returns:
+ List of successful predictions
+ """
+ from potato.flask_server import get_item_state_manager, get_users, get_user_state
+
+ # Check if we should pause labeling
+ should_pause, reason = self.should_pause_labeling()
+ if should_pause:
+ logger.info(f"Labeling paused: {reason}")
+ return []
+
+ if not self.has_enough_examples(schema_name):
+ logger.info(f"Not enough examples for schema {schema_name}")
+ return []
+
+ ism = get_item_state_manager()
+ if ism is None:
+ return []
+
+ # Check remaining capacity
+ remaining_capacity = self.get_remaining_label_capacity()
+ if remaining_capacity <= 0:
+ logger.info("No remaining label capacity")
+ return []
+
+ # Limit batch size to remaining capacity
+ effective_batch_size = min(self.batch_size, remaining_capacity)
+
+ text_key = self._get_text_key()
+ predictions = []
+
+ # Find unlabeled instances
+ unlabeled_ids = []
+ for instance_id in ism.instance_id_ordering:
+ # Skip if already labeled by LLM
+ if instance_id in self.labeled_instance_ids:
+ continue
+
+ # Skip if already annotated by humans
+ has_human_annotation = False
+ for username in get_users():
+ user_state = get_user_state(username)
+ if user_state:
+ all_annotations = user_state.get_all_annotations()
+ if instance_id in all_annotations:
+ has_human_annotation = True
+ break
+
+ if not has_human_annotation:
+ unlabeled_ids.append(instance_id)
+
+ if len(unlabeled_ids) >= effective_batch_size:
+ break
+
+ # Label instances
+ for instance_id in unlabeled_ids:
+ item = ism.get_item(instance_id)
+ instance_data = item.get_data() if item else None
+ if instance_data is None:
+ continue
+
+ text = instance_data.get(text_key, '')
+ if not text:
+ continue
+
+ prediction = self.label_instance(instance_id, schema_name, text)
+ if prediction and prediction.confidence_score >= self.confidence_threshold:
+ predictions.append(prediction)
+
+ self.last_batch_run = datetime.now()
+ logger.info(f"Batch labeled {len(predictions)} instances for schema {schema_name}")
+
+ return predictions
+
+ # === Verification Workflow ===
+
+ def get_pending_verifications(self, count: int = 1) -> List[Tuple[str, str]]:
+ """
+ Get instances pending human verification.
+
+ Args:
+ count: Number of verification tasks to return
+
+ Returns:
+ List of (instance_id, schema_name) tuples
+ """
+ with self._lock:
+ if self.verification_strategy == 'low_confidence':
+ # Sort by confidence ascending
+ pending = [
+ (inst_id, schema)
+ for inst_id, schema in self.verification_queue
+ if (inst_id in self.predictions and
+ schema in self.predictions[inst_id] and
+ self.predictions[inst_id][schema].verification_status == 'pending')
+ ]
+ pending.sort(
+ key=lambda x: self.predictions[x[0]][x[1]].confidence_score
+ )
+ return pending[:count]
+
+ elif self.verification_strategy == 'random':
+ pending = [
+ (inst_id, schema)
+ for inst_id, schema in self.verification_queue
+ if (inst_id in self.predictions and
+ schema in self.predictions[inst_id] and
+ self.predictions[inst_id][schema].verification_status == 'pending')
+ ]
+ random.shuffle(pending)
+ return pending[:count]
+
+ else: # mixed
+ pending = [
+ (inst_id, schema)
+ for inst_id, schema in self.verification_queue
+ if (inst_id in self.predictions and
+ schema in self.predictions[inst_id] and
+ self.predictions[inst_id][schema].verification_status == 'pending')
+ ]
+ # 50% low confidence, 50% random
+ pending.sort(
+ key=lambda x: self.predictions[x[0]][x[1]].confidence_score
+ )
+ half = count // 2
+ low_conf = pending[:half]
+ rest = pending[half:]
+ random.shuffle(rest)
+ return low_conf + rest[:count - half]
+
+ def record_verification(
+ self,
+ instance_id: str,
+ schema_name: str,
+ human_label: str,
+ verified_by: str
+ ) -> bool:
+ """
+ Record human verification of an LLM prediction.
+
+ Args:
+ instance_id: The verified instance
+ schema_name: The schema verified
+ human_label: The human's label
+ verified_by: Username of verifier
+
+ Returns:
+ True if verification recorded successfully
+ """
+ with self._lock:
+ if instance_id not in self.predictions:
+ logger.warning(f"No prediction found for instance {instance_id}")
+ return False
+
+ if schema_name not in self.predictions[instance_id]:
+ logger.warning(f"No prediction found for schema {schema_name}")
+ return False
+
+ prediction = self.predictions[instance_id][schema_name]
+ prediction.human_label = human_label
+ prediction.verified_by = verified_by
+ prediction.verified_at = datetime.now()
+
+ if prediction.predicted_label == human_label:
+ prediction.verification_status = 'verified_correct'
+ else:
+ prediction.verification_status = 'verified_incorrect'
+
+ # Remove from verification queue
+ try:
+ self.verification_queue.remove((instance_id, schema_name))
+ except ValueError:
+ pass
+
+ logger.info(
+ f"Verification recorded for {instance_id}: "
+ f"predicted={prediction.predicted_label}, human={human_label}, "
+ f"status={prediction.verification_status}"
+ )
+
+ return True
+
+ # === Accuracy Tracking ===
+
+ def get_accuracy_metrics(self, schema_name: Optional[str] = None) -> Dict[str, Any]:
+ """
+ Calculate accuracy metrics from verified predictions.
+
+ Args:
+ schema_name: Optional schema to filter by
+
+ Returns:
+ Dictionary with accuracy metrics
+ """
+ with self._lock:
+ verified_correct = 0
+ verified_incorrect = 0
+ pending = 0
+ total_predictions = 0
+
+ confidence_correct = []
+ confidence_incorrect = []
+
+ for inst_id, schemas in self.predictions.items():
+ for s_name, prediction in schemas.items():
+ if schema_name and s_name != schema_name:
+ continue
+
+ total_predictions += 1
+
+ if prediction.verification_status == 'verified_correct':
+ verified_correct += 1
+ confidence_correct.append(prediction.confidence_score)
+ elif prediction.verification_status == 'verified_incorrect':
+ verified_incorrect += 1
+ confidence_incorrect.append(prediction.confidence_score)
+ else:
+ pending += 1
+
+ total_verified = verified_correct + verified_incorrect
+ accuracy = verified_correct / total_verified if total_verified > 0 else None
+
+ avg_confidence_correct = (
+ sum(confidence_correct) / len(confidence_correct)
+ if confidence_correct else None
+ )
+ avg_confidence_incorrect = (
+ sum(confidence_incorrect) / len(confidence_incorrect)
+ if confidence_incorrect else None
+ )
+
+ return {
+ 'total_predictions': total_predictions,
+ 'verified_correct': verified_correct,
+ 'verified_incorrect': verified_incorrect,
+ 'pending_verification': pending,
+ 'total_verified': total_verified,
+ 'accuracy': accuracy,
+ 'avg_confidence_correct': avg_confidence_correct,
+ 'avg_confidence_incorrect': avg_confidence_incorrect,
+ 'schema_name': schema_name
+ }
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get overall ICL labeler status."""
+ with self._lock:
+ total_examples = sum(len(ex) for ex in self.schema_to_examples.values())
+ examples_by_schema = {
+ schema: len(examples)
+ for schema, examples in self.schema_to_examples.items()
+ }
+
+ # Check labeling status
+ should_pause, pause_reason = self.should_pause_labeling()
+ remaining_capacity = self.get_remaining_label_capacity()
+
+ return {
+ 'enabled': self.config.get('icl_labeling', {}).get('enabled', False),
+ 'total_examples': total_examples,
+ 'examples_by_schema': examples_by_schema,
+ 'total_predictions': sum(
+ len(schemas) for schemas in self.predictions.values()
+ ),
+ 'labeled_instances': len(self.labeled_instance_ids),
+ 'verification_queue_size': len(self.verification_queue),
+ 'last_example_refresh': (
+ self.last_example_refresh.isoformat()
+ if self.last_example_refresh else None
+ ),
+ 'last_batch_run': (
+ self.last_batch_run.isoformat()
+ if self.last_batch_run else None
+ ),
+ 'worker_running': (
+ self._worker_thread is not None and
+ self._worker_thread.is_alive()
+ ),
+ 'accuracy_metrics': self.get_accuracy_metrics(),
+ # Labeling limits status
+ 'labeling_paused': should_pause,
+ 'pause_reason': pause_reason,
+ 'remaining_label_capacity': remaining_capacity,
+ 'max_total_labels': self.max_total_labels,
+ 'max_unlabeled_ratio': self.max_unlabeled_ratio,
+ 'min_accuracy_threshold': self.min_accuracy_threshold
+ }
+
+ # === Background Worker ===
+
+ def start_background_worker(self) -> None:
+ """Start the background worker thread."""
+ if self._worker_thread is not None and self._worker_thread.is_alive():
+ logger.warning("Background worker already running")
+ return
+
+ self._stop_worker.clear()
+ self._worker_thread = threading.Thread(
+ target=self._worker_loop,
+ name="ICLLabelerWorker",
+ daemon=True
+ )
+ self._worker_thread.start()
+ logger.info("Started ICL labeler background worker")
+
+ def stop_background_worker(self) -> None:
+ """Stop the background worker thread."""
+ if self._worker_thread is None:
+ return
+
+ self._stop_worker.set()
+ self._worker_thread.join(timeout=5.0)
+ self._worker_thread = None
+ logger.info("Stopped ICL labeler background worker")
+
+ def _worker_loop(self) -> None:
+ """Main loop for the background worker."""
+ logger.info(
+ f"ICL background worker started, "
+ f"example_refresh={self.example_refresh_interval}s, "
+ f"batch_interval={self.batch_interval}s"
+ )
+
+ last_example_refresh = 0
+ last_batch = 0
+
+ while not self._stop_worker.is_set():
+ try:
+ current_time = time.time()
+
+ # Refresh examples periodically
+ if current_time - last_example_refresh >= self.example_refresh_interval:
+ self.refresh_high_confidence_examples()
+ last_example_refresh = current_time
+
+ # Run batch labeling periodically
+ if current_time - last_batch >= self.batch_interval:
+ schemas = self._get_annotation_schemes()
+ for schema in schemas:
+ schema_name = schema.get('name')
+ if schema_name and self.has_enough_examples(schema_name):
+ predictions = self.batch_label_instances(schema_name)
+ if predictions:
+ self.save_state()
+ last_batch = current_time
+
+ except Exception as e:
+ logger.error(f"ICL background worker error: {e}")
+
+ # Wait for next interval or stop signal
+ self._stop_worker.wait(min(self.example_refresh_interval, self.batch_interval) / 2)
+
+ # === Persistence ===
+
+ def save_state(self) -> None:
+ """Save current state to disk."""
+ task_dir = self.config.get('output_annotation_dir', '')
+ if not task_dir:
+ return
+
+ filepath = os.path.join(task_dir, self.predictions_file)
+
+ try:
+ with self._lock:
+ state = {
+ 'predictions': {
+ inst_id: {
+ schema: pred.to_dict()
+ for schema, pred in schemas.items()
+ }
+ for inst_id, schemas in self.predictions.items()
+ },
+ 'examples': {
+ schema: [ex.to_dict() for ex in examples]
+ for schema, examples in self.schema_to_examples.items()
+ },
+ 'verification_queue': self.verification_queue,
+ 'labeled_instance_ids': list(self.labeled_instance_ids),
+ 'last_example_refresh': (
+ self.last_example_refresh.isoformat()
+ if self.last_example_refresh else None
+ ),
+ 'last_batch_run': (
+ self.last_batch_run.isoformat()
+ if self.last_batch_run else None
+ )
+ }
+
+ # Atomic write
+ temp_path = filepath + '.tmp'
+ with open(temp_path, 'w') as f:
+ json.dump(state, f, indent=2)
+ os.replace(temp_path, filepath)
+
+ logger.debug(f"Saved ICL state to {filepath}")
+
+ except Exception as e:
+ logger.error(f"Error saving ICL state: {e}")
+
+ def load_state(self) -> None:
+ """Load state from disk."""
+ task_dir = self.config.get('output_annotation_dir', '')
+ if not task_dir:
+ return
+
+ filepath = os.path.join(task_dir, self.predictions_file)
+
+ if not os.path.exists(filepath):
+ return
+
+ try:
+ with open(filepath, 'r') as f:
+ state = json.load(f)
+
+ with self._lock:
+ # Load predictions
+ self.predictions = {}
+ for inst_id, schemas in state.get('predictions', {}).items():
+ self.predictions[inst_id] = {
+ schema: ICLPrediction.from_dict(pred_data)
+ for schema, pred_data in schemas.items()
+ }
+
+ # Load examples
+ self.schema_to_examples = {}
+ for schema, examples in state.get('examples', {}).items():
+ self.schema_to_examples[schema] = [
+ HighConfidenceExample.from_dict(ex) for ex in examples
+ ]
+
+ # Load other state
+ self.verification_queue = [
+ tuple(item) for item in state.get('verification_queue', [])
+ ]
+ self.labeled_instance_ids = set(state.get('labeled_instance_ids', []))
+
+ if state.get('last_example_refresh'):
+ self.last_example_refresh = datetime.fromisoformat(
+ state['last_example_refresh']
+ )
+ if state.get('last_batch_run'):
+ self.last_batch_run = datetime.fromisoformat(
+ state['last_batch_run']
+ )
+
+ logger.info(f"Loaded ICL state from {filepath}")
+
+ except Exception as e:
+ logger.error(f"Error loading ICL state: {e}")
+
+
+# Module-level singleton access
+_icl_labeler: Optional[ICLLabeler] = None
+
+
+def init_icl_labeler(config: Dict[str, Any]) -> ICLLabeler:
+ """Initialize the global ICL labeler."""
+ global _icl_labeler
+ _icl_labeler = ICLLabeler(config)
+ _icl_labeler.load_state()
+ return _icl_labeler
+
+
+def get_icl_labeler() -> Optional[ICLLabeler]:
+ """Get the global ICL labeler instance."""
+ return _icl_labeler
+
+
+def clear_icl_labeler() -> None:
+ """Clear the global ICL labeler (for testing)."""
+ global _icl_labeler
+ if _icl_labeler:
+ _icl_labeler.stop_background_worker()
+ _icl_labeler = None
+ ICLLabeler._instance = None
diff --git a/potato/ai/icl_prompt_builder.py b/potato/ai/icl_prompt_builder.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c54f8e1c876ca2fae1ee649d79cf1e537206ea1
--- /dev/null
+++ b/potato/ai/icl_prompt_builder.py
@@ -0,0 +1,315 @@
+"""
+In-Context Learning Prompt Builder
+
+This module builds effective prompts for in-context learning based labeling.
+It formats high-confidence examples and target instances into prompts that
+elicit accurate label predictions with confidence scores.
+"""
+
+import json
+import logging
+import re
+from typing import Dict, List, Any, Tuple, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from potato.ai.icl_labeler import HighConfidenceExample
+
+logger = logging.getLogger(__name__)
+
+
+class ICLPromptBuilder:
+ """
+ Builds effective prompts for in-context learning.
+
+ The prompt structure:
+ 1. System instructions explaining the task
+ 2. Schema description and available labels
+ 3. High-confidence examples with their labels
+ 4. Target text to label
+ 5. Output format instructions (JSON with label, confidence, reasoning)
+ """
+
+ def __init__(self, max_example_length: int = 500, max_target_length: int = 1000):
+ """
+ Initialize the prompt builder.
+
+ Args:
+ max_example_length: Maximum characters per example text
+ max_target_length: Maximum characters for target text
+ """
+ self.max_example_length = max_example_length
+ self.max_target_length = max_target_length
+
+ def build_prompt(
+ self,
+ schema: Dict[str, Any],
+ examples: List['HighConfidenceExample'],
+ target_text: str
+ ) -> str:
+ """
+ Build a complete ICL prompt.
+
+ Args:
+ schema: Annotation schema dictionary with name, description, labels
+ examples: List of high-confidence examples
+ target_text: The text to be labeled
+
+ Returns:
+ Complete prompt string
+ """
+ parts = []
+
+ # System instructions
+ parts.append(self._build_system_prompt(schema))
+
+ # Examples section
+ if examples:
+ parts.append("\n## Examples\n")
+ parts.append("Here are examples of correctly labeled texts:\n")
+ for i, example in enumerate(examples, 1):
+ parts.append(self._format_example(example, i))
+
+ # Target text section
+ parts.append("\n## Your Task\n")
+ parts.append("Now label the following text:\n")
+ parts.append(f'Text: "{self._truncate_text(target_text, self.max_target_length)}"\n')
+
+ # Output format instructions
+ parts.append(self._build_output_instructions(schema))
+
+ return "\n".join(parts)
+
+ def _build_system_prompt(self, schema: Dict[str, Any]) -> str:
+ """Build the system/instruction portion of the prompt."""
+ schema_name = schema.get('name', 'unknown')
+ description = schema.get('description', 'Label the text according to the schema.')
+ labels = self._get_labels_from_schema(schema)
+ annotation_type = schema.get('annotation_type', 'radio')
+
+ prompt = f"""You are an expert annotation assistant. Your task is to label text according to a specific annotation schema.
+
+## Schema: {schema_name}
+
+**Description:** {description}
+
+**Available Labels:** {', '.join(labels)}
+"""
+
+ # Add type-specific instructions
+ if annotation_type == 'radio':
+ prompt += "\n**Task Type:** Single-choice classification. Select exactly ONE label.\n"
+ elif annotation_type == 'multiselect':
+ prompt += "\n**Task Type:** Multi-label classification. Select ALL applicable labels.\n"
+ elif annotation_type == 'likert':
+ prompt += "\n**Task Type:** Rating scale. Choose the most appropriate rating.\n"
+
+ return prompt
+
+ def _format_example(self, example: 'HighConfidenceExample', index: int) -> str:
+ """Format a single example for the prompt."""
+ truncated_text = self._truncate_text(example.text, self.max_example_length)
+
+ return f"""
+### Example {index}
+Text: "{truncated_text}"
+Label: **{example.label}**
+(Agreement: {example.agreement_score:.0%} from {example.annotator_count} annotators)
+"""
+
+ def _build_output_instructions(self, schema: Dict[str, Any]) -> str:
+ """Build instructions for the expected output format."""
+ labels = self._get_labels_from_schema(schema)
+ labels_json = json.dumps(labels)
+
+ return f"""
+## Output Format
+
+Respond with a JSON object containing:
+- `label`: Your chosen label (must be one of: {labels_json})
+- `confidence`: Your confidence score from 0.0 to 1.0
+ - 1.0 = Absolutely certain
+ - 0.7-0.9 = High confidence
+ - 0.5-0.7 = Moderate confidence
+ - 0.3-0.5 = Low confidence
+ - 0.0-0.3 = Very uncertain
+- `reasoning`: Brief explanation for your choice (1-2 sentences)
+
+**Important:**
+- Only use labels from the provided list
+- Be honest about your confidence - reflect your actual certainty, not a fixed value
+- Base your decision on the examples and schema description
+- Use the full 0.0โ1.0 range: reserve 0.9+ for near-certain cases, use 0.4โ0.6 when genuinely unsure
+
+Example response (the confidence value here is illustrative only โ yours should reflect actual certainty):
+```json
+{{"label": "example_label", "confidence": 0.72, "reasoning": "The text shows clear indicators of..."}}
+```
+
+Now provide your response as JSON:
+"""
+
+ def _get_labels_from_schema(self, schema: Dict[str, Any]) -> List[str]:
+ """Extract label names from schema definition."""
+ labels = schema.get('labels', [])
+ result = []
+ for label in labels:
+ if isinstance(label, str):
+ result.append(label)
+ elif isinstance(label, dict):
+ result.append(label.get('name', str(label)))
+ return result
+
+ def _truncate_text(self, text: str, max_length: int) -> str:
+ """Truncate text to max length, preserving word boundaries."""
+ if len(text) <= max_length:
+ return text
+
+ truncated = text[:max_length]
+ # Try to break at word boundary
+ last_space = truncated.rfind(' ')
+ if last_space > max_length * 0.8:
+ truncated = truncated[:last_space]
+
+ return truncated + "..."
+
+ def parse_response(
+ self,
+ response: str,
+ schema: Dict[str, Any]
+ ) -> Tuple[Optional[str], float, str]:
+ """
+ Parse the LLM response to extract label, confidence, and reasoning.
+
+ Args:
+ response: Raw response from LLM
+ schema: Schema for validation
+
+ Returns:
+ Tuple of (label, confidence, reasoning) or (None, 0.0, "") on failure
+ """
+ try:
+ # Try to parse as JSON directly
+ data = self._extract_json(response)
+ if data:
+ label = data.get('label', '')
+ confidence = float(data.get('confidence', 0.5))
+ reasoning = data.get('reasoning', '')
+
+ # Validate label
+ valid_labels = self._get_labels_from_schema(schema)
+ if label in valid_labels:
+ return label, min(1.0, max(0.0, confidence)), reasoning
+
+ # Try fuzzy matching
+ matched = self._fuzzy_match_label(label, valid_labels)
+ if matched:
+ return matched, min(1.0, max(0.0, confidence)), reasoning
+
+ # Fallback: try to extract label from text
+ return self._extract_label_from_text(response, schema)
+
+ except Exception as e:
+ logger.warning(f"Error parsing response: {e}")
+ return None, 0.0, ""
+
+ def _extract_json(self, text: str) -> Optional[Dict[str, Any]]:
+ """Extract JSON from text, handling markdown code blocks."""
+ # Try direct parse
+ try:
+ return json.loads(text)
+ except json.JSONDecodeError:
+ pass
+
+ # Try to find JSON in code blocks
+ json_patterns = [
+ r'```json\s*(.*?)\s*```',
+ r'```\s*(.*?)\s*```',
+ r'\{[^{}]*"label"[^{}]*\}'
+ ]
+
+ for pattern in json_patterns:
+ match = re.search(pattern, text, re.DOTALL)
+ if match:
+ try:
+ json_str = match.group(1) if match.lastindex else match.group(0)
+ return json.loads(json_str)
+ except (json.JSONDecodeError, IndexError):
+ continue
+
+ return None
+
+ def _fuzzy_match_label(self, label: str, valid_labels: List[str]) -> Optional[str]:
+ """Try to match label with case-insensitive comparison."""
+ label_lower = label.lower().strip()
+ for valid in valid_labels:
+ if valid.lower().strip() == label_lower:
+ return valid
+ return None
+
+ def _extract_label_from_text(
+ self,
+ text: str,
+ schema: Dict[str, Any]
+ ) -> Tuple[Optional[str], float, str]:
+ """Fallback: try to extract label directly from text."""
+ valid_labels = self._get_labels_from_schema(schema)
+ text_lower = text.lower()
+
+ for label in valid_labels:
+ # Look for label mentioned in text
+ if label.lower() in text_lower:
+ return label, 0.5, "Extracted from response text (low confidence)"
+
+ return None, 0.0, ""
+
+
+class MultiSelectPromptBuilder(ICLPromptBuilder):
+ """
+ Specialized prompt builder for multi-select (multi-label) tasks.
+ """
+
+ def _build_output_instructions(self, schema: Dict[str, Any]) -> str:
+ """Build output instructions for multi-select."""
+ labels = self._get_labels_from_schema(schema)
+ labels_json = json.dumps(labels)
+
+ return f"""
+## Output Format
+
+Respond with a JSON object containing:
+- `labels`: Array of selected labels (from: {labels_json})
+- `confidence`: Your overall confidence score from 0.0 to 1.0 โ reflect actual certainty, not a fixed value
+- `reasoning`: Brief explanation for your choices
+
+Example response (the confidence value here is illustrative only โ yours should reflect actual certainty):
+```json
+{{"labels": ["label1", "label2"], "confidence": 0.65, "reasoning": "The text exhibits both..."}}
+```
+
+Now provide your response as JSON:
+"""
+
+ def parse_response(
+ self,
+ response: str,
+ schema: Dict[str, Any]
+ ) -> Tuple[Optional[List[str]], float, str]:
+ """Parse multi-select response."""
+ try:
+ data = self._extract_json(response)
+ if data:
+ labels = data.get('labels', [])
+ confidence = float(data.get('confidence', 0.5))
+ reasoning = data.get('reasoning', '')
+
+ valid_labels = self._get_labels_from_schema(schema)
+ validated = [l for l in labels if l in valid_labels]
+
+ if validated:
+ return validated, min(1.0, max(0.0, confidence)), reasoning
+
+ return None, 0.0, ""
+
+ except Exception as e:
+ logger.warning(f"Error parsing multi-select response: {e}")
+ return None, 0.0, ""
diff --git a/potato/ai/judge.py b/potato/ai/judge.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f5d7ccb3e86b5d4bc7daffa4238905f753101bb
--- /dev/null
+++ b/potato/ai/judge.py
@@ -0,0 +1,265 @@
+"""
+LLM-as-Judge service for human-alignment.
+
+Produces a judge verdict (label + confidence + reasoning) for an annotation
+instance, given the schema (labels + description + an editable rubric) and,
+optionally, few-shot examples drawn from high-agreement human labels. The
+verdicts are compared against human labels elsewhere
+(``potato/server_utils/judge_alignment.py``) to measure and calibrate
+humanโjudge agreement (Cohen's ฮบ).
+
+This deliberately does NOT reuse ``ICLLabeler`` as the judge โ ICL auto-labels
+from inter-annotator agreement, which would leak the gold labels we are trying
+to measure the judge against. We only borrow ICLLabeler's *example selection*
+for few-shot calibration, and we always exclude the instance being judged from
+its own example set.
+
+The judge call goes through the same ``AIEndpointFactory`` / ``BaseAIEndpoint``
+machinery as every other AI feature (mirrors ``icl_labeler.label_instance``),
+so it works with any configured provider.
+"""
+
+import hashlib
+import json
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class JudgePrediction:
+ """A single LLM-judge verdict for one instance + schema."""
+
+ instance_id: str
+ schema_name: str
+ predicted_label: str
+ confidence: float # 0.0โ1.0
+ reasoning: str = ""
+ model_name: str = ""
+ prompt_version: str = ""
+ examples_used: List[str] = field(default_factory=list)
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ "instance_id": self.instance_id,
+ "schema_name": self.schema_name,
+ "predicted_label": self.predicted_label,
+ "confidence": self.confidence,
+ "reasoning": self.reasoning,
+ "model_name": self.model_name,
+ "prompt_version": self.prompt_version,
+ "examples_used": self.examples_used,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "JudgePrediction":
+ return cls(
+ instance_id=data["instance_id"],
+ schema_name=data["schema_name"],
+ predicted_label=data.get("predicted_label", ""),
+ confidence=float(data.get("confidence", 0.0)),
+ reasoning=data.get("reasoning", ""),
+ model_name=data.get("model_name", ""),
+ prompt_version=data.get("prompt_version", ""),
+ examples_used=data.get("examples_used", []),
+ )
+
+
+def extract_labels(schema_info: Dict[str, Any]) -> List[str]:
+ """Return the allowed label names for a categorical schema.
+
+ Supports ``radio``/``select``/``multiselect`` (``labels`` list of
+ str|dict) and ``likert`` (1..size). Returns ``[]`` for unsupported types.
+ """
+ atype = schema_info.get("annotation_type", "")
+ if atype == "likert":
+ size = int(schema_info.get("size", 5))
+ return [str(i) for i in range(1, size + 1)]
+ labels = schema_info.get("labels", [])
+ out = []
+ for lab in labels:
+ if isinstance(lab, dict):
+ out.append(str(lab.get("name", "")))
+ else:
+ out.append(str(lab))
+ return [x for x in out if x]
+
+
+def compute_prompt_version(rubric: str, schema_name: str, few_shot: bool) -> str:
+ """Stable short hash identifying this judge configuration.
+
+ Editing the rubric (or toggling few-shot) yields a new version so the admin
+ report can track ฮบ across prompt versions.
+ """
+ basis = f"{schema_name}โ{int(bool(few_shot))}โ{rubric or ''}"
+ return "v_" + hashlib.sha1(basis.encode("utf-8")).hexdigest()[:10]
+
+
+class JudgeService:
+ """Builds judge prompts and queries the configured AI endpoint."""
+
+ def __init__(self, config: Dict[str, Any]):
+ self.config = config or {}
+ self.judge_config = self.config.get("judge_alignment", {}) or {}
+ self._endpoint = None
+ self._endpoint_initialized = False
+
+ # ----- endpoint -------------------------------------------------------
+
+ def _get_endpoint(self):
+ if not self._endpoint_initialized:
+ self._endpoint_initialized = True
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+ # The judge endpoint config lives under judge_alignment, but
+ # fall back to the task's ai_support so a single endpoint can
+ # serve both. Shape mirrors what AIEndpointFactory expects.
+ ai_support = self.judge_config.get("ai_support") or self.config.get("ai_support")
+ if not ai_support:
+ logger.warning("Judge: no ai_support / judge_alignment.ai_support configured")
+ return None
+ self._endpoint = AIEndpointFactory.create_endpoint({"ai_support": ai_support})
+ except Exception as e:
+ logger.error(f"Judge: failed to create endpoint: {e}")
+ self._endpoint = None
+ return self._endpoint
+
+ # ----- prompt ---------------------------------------------------------
+
+ def _schema_judge_config(self, schema_name: str) -> Dict[str, Any]:
+ per_schema = self.judge_config.get("schemas", {}) or {}
+ return per_schema.get(schema_name, {}) or {}
+
+ def get_rubric(self, schema_info: Dict[str, Any]) -> str:
+ """Editable rubric for a schema; falls back to its description."""
+ sc = self._schema_judge_config(schema_info.get("name", ""))
+ return sc.get("rubric") or schema_info.get("description", "") or ""
+
+ def build_prompt(
+ self,
+ schema_info: Dict[str, Any],
+ instance_text: str,
+ few_shot_examples: Optional[List[Dict[str, str]]] = None,
+ ) -> str:
+ """Compose the judge prompt.
+
+ few_shot_examples: list of {"text": ..., "label": ...} gold exemplars
+ (already excluding the target instance).
+ """
+ labels = extract_labels(schema_info)
+ rubric = self.get_rubric(schema_info)
+ parts = [
+ "You are an expert evaluator acting as an impartial judge.",
+ "Assign exactly one label to the item below, following the rubric.",
+ "",
+ f"Task: {schema_info.get('description', '')}".rstrip(),
+ f"Rubric: {rubric}".rstrip(),
+ "Allowed labels: " + ", ".join(labels) if labels else "",
+ ]
+ if few_shot_examples:
+ parts.append("\nExamples (item โ correct label):")
+ for ex in few_shot_examples:
+ parts.append(f"- {_truncate(ex.get('text', ''))} โ {ex.get('label', '')}")
+ parts.append("\nItem to judge:")
+ parts.append(_truncate(instance_text, 4000))
+ parts.append(
+ '\nRespond as JSON: {"label": , '
+ '"confidence": <0.0-1.0>, "reasoning": }.'
+ )
+ return "\n".join(p for p in parts if p != "")
+
+ # ----- judging --------------------------------------------------------
+
+ def judge_instance(
+ self,
+ instance_id: str,
+ schema_info: Dict[str, Any],
+ instance_text: str,
+ few_shot_examples: Optional[List[Dict[str, str]]] = None,
+ prompt_version: Optional[str] = None,
+ ) -> Optional[JudgePrediction]:
+ """Query the judge for one instance. Returns None on failure."""
+ endpoint = self._get_endpoint()
+ if endpoint is None:
+ return None
+
+ schema_name = schema_info.get("name", "")
+ valid_labels = extract_labels(schema_info)
+ rubric = self.get_rubric(schema_info)
+ if prompt_version is None:
+ prompt_version = compute_prompt_version(
+ rubric, schema_name, bool(few_shot_examples)
+ )
+
+ prompt = self.build_prompt(schema_info, instance_text, few_shot_examples)
+
+ try:
+ from pydantic import BaseModel
+
+ class JudgeVerdict(BaseModel):
+ label: str
+ confidence: float = 0.5
+ reasoning: str = ""
+
+ response = endpoint.query(prompt, JudgeVerdict)
+ if isinstance(response, str):
+ data = json.loads(response)
+ elif hasattr(response, "model_dump"):
+ data = response.model_dump()
+ elif hasattr(response, "dict"):
+ data = response.dict()
+ else:
+ data = response or {}
+ except Exception as e:
+ logger.error(f"Judge: query/parse failed for {instance_id}/{schema_name}: {e}")
+ return None
+
+ predicted = str(data.get("label", "")).strip()
+ try:
+ confidence = float(data.get("confidence", 0.5))
+ except (TypeError, ValueError):
+ confidence = 0.5
+ confidence = min(1.0, max(0.0, confidence))
+ reasoning = str(data.get("reasoning", ""))
+
+ if valid_labels and predicted not in valid_labels:
+ matched = _fuzzy_match_label(predicted, valid_labels)
+ if matched is None:
+ logger.warning(
+ f"Judge: invalid label '{predicted}' for {instance_id}/{schema_name}"
+ )
+ return None
+ predicted = matched
+
+ return JudgePrediction(
+ instance_id=instance_id,
+ schema_name=schema_name,
+ predicted_label=predicted,
+ confidence=confidence,
+ reasoning=reasoning,
+ model_name=getattr(endpoint, "model", ""),
+ prompt_version=prompt_version,
+ examples_used=[e.get("id", "") for e in (few_shot_examples or []) if e.get("id")],
+ )
+
+
+def _truncate(text: str, limit: int = 300) -> str:
+ text = str(text or "")
+ return text if len(text) <= limit else text[:limit] + "โฆ"
+
+
+def _fuzzy_match_label(predicted: str, valid_labels: List[str]) -> Optional[str]:
+ """Case-insensitive / prefix match a model label to an allowed label."""
+ if not predicted:
+ return None
+ low = predicted.lower().strip()
+ for lab in valid_labels:
+ if lab.lower() == low:
+ return lab
+ for lab in valid_labels:
+ ll = lab.lower()
+ if ll in low or low in ll:
+ return lab
+ return None
diff --git a/potato/ai/llm_active_learning.py b/potato/ai/llm_active_learning.py
new file mode 100644
index 0000000000000000000000000000000000000000..589862acade15123ba93d3fca0bd75d95b730805
--- /dev/null
+++ b/potato/ai/llm_active_learning.py
@@ -0,0 +1,733 @@
+"""
+LLM Integration for Active Learning
+
+This module provides LLM-based active learning capabilities using VLLM endpoints.
+It implements confidence-based instance selection and prediction using large language
+models, with support for multiple confidence elicitation methods:
+
+- **logprobs**: Extract token-level log probabilities from VLLM/OpenAI-compatible
+ endpoints for calibrated confidence scores.
+- **verbalized**: Ask the LLM to self-report confidence on a 1-10 scale (default).
+- **consistency**: Query the same instance N times with temperature > 0 and use
+ agreement rate as confidence (works with any endpoint).
+
+References:
+ Tian et al. (2023) "Just Ask for Calibration: Strategies for Eliciting
+ Calibrated Confidence Scores from Language Models Fine-Tuned with Human
+ Feedback." EMNLP 2023.
+
+ Xiong et al. (2024) "Can LLMs Express Their Uncertainty? An Empirical
+ Evaluation of Confidence Elicitation in LLMs." ICLR 2024.
+"""
+
+import logging
+import math
+import time
+import json
+import requests
+from collections import Counter
+from typing import Dict, List, Optional, Tuple, Any
+from dataclasses import dataclass, field
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import numpy as np
+
+from potato.active_learning_manager import TrainingMetrics
+
+
+def _loads_lenient(content: str):
+ """json.loads that tolerates markdown code fences and surrounding prose.
+
+ Many models (e.g. Gemma on vLLM) wrap JSON in ```json ... ``` fences even
+ when response_format=json_object is requested, which breaks a naive
+ json.loads(). Strip fences and, failing that, extract the first {...} block.
+ """
+ import re
+ if content is None:
+ raise json.JSONDecodeError("empty content", "", 0)
+ s = content.strip()
+ # Strip a leading ```json / ``` fence and trailing ```
+ s = re.sub(r"^```(?:json|JSON)?\s*", "", s)
+ s = re.sub(r"\s*```$", "", s).strip()
+ try:
+ return json.loads(s)
+ except json.JSONDecodeError:
+ # Fall back to the first balanced-looking {...} object in the text.
+ m = re.search(r"\{.*\}", s, re.DOTALL)
+ if m:
+ return json.loads(m.group(0))
+ raise
+
+
+@dataclass
+class LLMPrediction:
+ """Result of an LLM prediction."""
+ instance_id: str
+ predicted_label: str
+ confidence_score: float
+ raw_response: str
+ error_message: Optional[str] = None
+ confidence_method: str = "verbalized"
+
+
+@dataclass
+class LLMConfig:
+ """Configuration for LLM integration."""
+ endpoint_url: str
+ model_name: str
+ max_tokens: int = 512
+ temperature: float = 0.1
+ timeout: int = 30
+ batch_size: int = 10
+ retry_attempts: int = 3
+ retry_delay: float = 1.0
+ max_instances_per_request: int = 5
+ confidence_method: str = "verbalized" # logprobs | verbalized | consistency
+ consistency_samples: int = 3
+
+
+class LLMActiveLearning:
+ """
+ LLM-based active learning implementation.
+
+ This class provides methods for:
+ - Querying LLMs for predictions and confidence scores
+ - Batch processing of instances
+ - Error handling and retry logic
+ - Integration with the active learning pipeline
+ """
+
+ def __init__(self, config: LLMConfig):
+ self.config = config
+ self.logger = logging.getLogger(__name__)
+ self.session = requests.Session()
+
+ # Configure session
+ self.session.timeout = config.timeout
+
+ # Test connection on initialization
+ self._test_connection()
+
+ def _test_connection(self):
+ """Test the connection to the LLM endpoint."""
+ try:
+ test_payload = {
+ "model": self.config.model_name,
+ "messages": [{"role": "user", "content": "Hello"}],
+ "max_tokens": 10,
+ "temperature": 0.1
+ }
+
+ response = self.session.post(
+ self.config.endpoint_url,
+ json=test_payload,
+ timeout=5
+ )
+
+ if response.status_code == 200:
+ self.logger.info(f"Successfully connected to LLM endpoint: {self.config.endpoint_url}")
+ else:
+ self.logger.warning(f"LLM endpoint returned status {response.status_code}: {response.text}")
+
+ except Exception as e:
+ self.logger.error(f"Failed to connect to LLM endpoint: {e}")
+ # Don't raise - allow fallback to traditional methods
+
+ def predict_instances(self, instances: List[Dict[str, Any]],
+ annotation_instructions: str,
+ schema_name: str,
+ label_options: List[str]) -> List[LLMPrediction]:
+ """
+ Predict labels and confidence scores for instances using LLM.
+
+ Args:
+ instances: List of instances to predict
+ annotation_instructions: Instructions for the annotation task
+ schema_name: Name of the annotation schema
+ label_options: Available label options
+
+ Returns:
+ List of LLM predictions with confidence scores
+ """
+ if not instances:
+ return []
+
+ self.logger.info(f"Starting LLM prediction for {len(instances)} instances")
+
+ # Create prompts for each instance
+ prompts = self._create_prompts(instances, annotation_instructions, schema_name, label_options)
+
+ # Process in batches
+ all_predictions = []
+
+ for i in range(0, len(prompts), self.config.batch_size):
+ batch_prompts = prompts[i:i + self.config.batch_size]
+ batch_instances = instances[i:i + self.config.batch_size]
+
+ batch_predictions = self._process_batch(batch_prompts, batch_instances)
+ all_predictions.extend(batch_predictions)
+
+ # Small delay between batches to avoid overwhelming the endpoint
+ if i + self.config.batch_size < len(prompts):
+ time.sleep(0.1)
+
+ self.logger.info(f"Completed LLM prediction for {len(all_predictions)} instances")
+ return all_predictions
+
+ def _create_prompts(self, instances: List[Dict[str, Any]],
+ annotation_instructions: str,
+ schema_name: str,
+ label_options: List[str]) -> List[str]:
+ """Create prompts for LLM prediction."""
+ prompts = []
+
+ # Create the base prompt template
+ base_prompt = self._create_base_prompt(annotation_instructions, schema_name, label_options)
+
+ for instance in instances:
+ # Extract text content
+ text_content = self._extract_text_content(instance)
+
+ # Create instance-specific prompt
+ prompt = f"{base_prompt}\n\nText to annotate:\n{text_content}\n\nPlease provide your prediction and confidence score."
+
+ prompts.append(prompt)
+
+ return prompts
+
+ def _create_base_prompt(self, annotation_instructions: str,
+ schema_name: str,
+ label_options: List[str]) -> str:
+ """Create the base prompt for LLM prediction."""
+ prompt = f"""You are an expert annotator for a text classification task.
+
+Task: {annotation_instructions}
+
+Schema: {schema_name}
+
+Available labels: {', '.join(label_options)}
+
+For each text, please:
+1. Analyze the text carefully
+2. Choose the most appropriate label from the available options
+3. Provide a confidence score from 1 to 10 (where 1 = very uncertain, 10 = very confident)
+
+Please respond in the following JSON format:
+{{
+ "label": "chosen_label",
+ "confidence": confidence_score,
+ "reasoning": "brief explanation of your choice"
+}}
+
+Example response:
+{{
+ "label": "{label_options[0] if label_options else 'example'}",
+ "confidence": 8,
+ "reasoning": "The text clearly expresses positive sentiment based on the language used."
+}}"""
+
+ return prompt
+
+ def _extract_text_content(self, instance: Dict[str, Any]) -> str:
+ """Extract text content from an instance."""
+ # Try common text field names
+ text_fields = ['text', 'content', 'message', 'sentence', 'document']
+
+ for field in text_fields:
+ if field in instance:
+ content = instance[field]
+ if isinstance(content, str):
+ return content
+ elif isinstance(content, dict):
+ # Handle nested text fields
+ for nested_field in text_fields:
+ if nested_field in content:
+ return str(content[nested_field])
+
+ # Fallback: convert the entire instance to string
+ return str(instance)
+
+ def _process_batch(self, prompts: List[str], instances: List[Dict[str, Any]]) -> List[LLMPrediction]:
+ """Process a batch of prompts."""
+ predictions = []
+
+ # Use ThreadPoolExecutor for parallel processing within the batch
+ with ThreadPoolExecutor(max_workers=min(len(prompts), 5)) as executor:
+ future_to_index = {
+ executor.submit(self._predict_single, prompt, instances[i]): i
+ for i, prompt in enumerate(prompts)
+ }
+
+ for future in as_completed(future_to_index):
+ index = future_to_index[future]
+ try:
+ prediction = future.result()
+ predictions.append(prediction)
+ except Exception as e:
+ self.logger.error(f"Error processing instance {index}: {e}")
+ # Create error prediction
+ error_prediction = LLMPrediction(
+ instance_id=instances[index].get('id', f'instance_{index}'),
+ predicted_label='',
+ confidence_score=0.1,
+ raw_response='',
+ error_message=str(e)
+ )
+ predictions.append(error_prediction)
+
+ return predictions
+
+ def _predict_single(self, prompt: str, instance: Dict[str, Any]) -> LLMPrediction:
+ """Make a single prediction using the LLM.
+
+ Dispatches to the appropriate confidence method:
+ - logprobs: Extract token-level log probabilities
+ - consistency: Query N times, use agreement rate
+ - verbalized (default): Parse self-reported confidence from JSON
+ """
+ method = self.config.confidence_method
+
+ if method == "consistency":
+ return self._predict_consistency(prompt, instance)
+ elif method == "logprobs":
+ return self._predict_with_logprobs(prompt, instance)
+ else:
+ return self._predict_verbalized(prompt, instance)
+
+ def _predict_verbalized(self, prompt: str, instance: Dict[str, Any]) -> LLMPrediction:
+ """Original verbalized confidence method (1-10 scale)."""
+ instance_id = instance.get('id', 'unknown')
+
+ for attempt in range(self.config.retry_attempts):
+ try:
+ payload = {
+ "model": self.config.model_name,
+ "messages": [
+ {"role": "system", "content": "You are a helpful assistant that provides structured JSON responses."},
+ {"role": "user", "content": prompt}
+ ],
+ "max_tokens": self.config.max_tokens,
+ "temperature": self.config.temperature,
+ "response_format": {"type": "json_object"}
+ }
+
+ response = self.session.post(
+ self.config.endpoint_url,
+ json=payload,
+ timeout=self.config.timeout
+ )
+
+ if response.status_code == 200:
+ result = response.json()
+
+ if 'choices' in result and len(result['choices']) > 0:
+ content = result['choices'][0]['message']['content']
+
+ try:
+ parsed_response = _loads_lenient(content)
+ predicted_label = parsed_response.get('label', '')
+ confidence_score = parsed_response.get('confidence', 1)
+
+ if not isinstance(confidence_score, (int, float)):
+ confidence_score = 1
+ else:
+ confidence_score = max(1, min(10, confidence_score)) / 10.0
+
+ return LLMPrediction(
+ instance_id=instance_id,
+ predicted_label=predicted_label,
+ confidence_score=confidence_score,
+ raw_response=content,
+ confidence_method="verbalized"
+ )
+
+ except json.JSONDecodeError as e:
+ self.logger.warning(f"Failed to parse JSON response for instance {instance_id}: {e}")
+ return self._extract_from_raw_response(content, instance_id)
+
+ else:
+ raise Exception(f"Invalid response format: {result}")
+
+ else:
+ raise Exception(f"HTTP {response.status_code}: {response.text}")
+
+ except Exception as e:
+ self.logger.warning(f"Attempt {attempt + 1} failed for instance {instance_id}: {e}")
+
+ if attempt < self.config.retry_attempts - 1:
+ time.sleep(self.config.retry_delay * (attempt + 1))
+ else:
+ return LLMPrediction(
+ instance_id=instance_id,
+ predicted_label='',
+ confidence_score=0.1,
+ raw_response='',
+ error_message=f"All attempts failed: {e}",
+ confidence_method="verbalized"
+ )
+
+ return LLMPrediction(
+ instance_id=instance_id,
+ predicted_label='',
+ confidence_score=0.1,
+ raw_response='',
+ error_message="Unknown error",
+ confidence_method="verbalized"
+ )
+
+ def _predict_with_logprobs(self, prompt: str, instance: Dict[str, Any]) -> LLMPrediction:
+ """Extract confidence from token-level log probabilities.
+
+ Requests logprobs=True from VLLM/OpenAI-compatible endpoints and
+ computes confidence as exp(mean_logprob) over the label tokens.
+ Falls back to verbalized confidence if logprobs unavailable.
+ """
+ instance_id = instance.get('id', 'unknown')
+
+ for attempt in range(self.config.retry_attempts):
+ try:
+ payload = {
+ "model": self.config.model_name,
+ "messages": [
+ {"role": "system", "content": "You are a helpful assistant that provides structured JSON responses."},
+ {"role": "user", "content": prompt}
+ ],
+ "max_tokens": self.config.max_tokens,
+ "temperature": self.config.temperature,
+ "response_format": {"type": "json_object"},
+ "logprobs": True,
+ "top_logprobs": 5,
+ }
+
+ response = self.session.post(
+ self.config.endpoint_url,
+ json=payload,
+ timeout=self.config.timeout
+ )
+
+ if response.status_code == 200:
+ result = response.json()
+
+ if 'choices' not in result or len(result['choices']) == 0:
+ raise Exception(f"Invalid response format: {result}")
+
+ choice = result['choices'][0]
+ content = choice['message']['content']
+
+ # Parse label from JSON content
+ try:
+ parsed = _loads_lenient(content)
+ except json.JSONDecodeError:
+ return self._extract_from_raw_response(content, instance_id)
+
+ predicted_label = parsed.get('label', '')
+
+ # Try to extract logprobs
+ logprobs_data = choice.get('logprobs', {})
+ token_logprobs = logprobs_data.get('content', [])
+
+ if token_logprobs:
+ # Compute mean logprob across all tokens
+ log_probs = [
+ t['logprob'] for t in token_logprobs
+ if 'logprob' in t and t['logprob'] is not None
+ ]
+ if log_probs:
+ mean_logprob = sum(log_probs) / len(log_probs)
+ confidence_score = min(1.0, max(0.0, math.exp(mean_logprob)))
+ else:
+ # No valid logprobs, fall back to verbalized
+ confidence_score = parsed.get('confidence', 5)
+ if isinstance(confidence_score, (int, float)):
+ confidence_score = max(1, min(10, confidence_score)) / 10.0
+ else:
+ confidence_score = 0.5
+ else:
+ # Endpoint didn't return logprobs, fall back to verbalized
+ self.logger.debug(f"No logprobs returned for {instance_id}, using verbalized")
+ confidence_score = parsed.get('confidence', 5)
+ if isinstance(confidence_score, (int, float)):
+ confidence_score = max(1, min(10, confidence_score)) / 10.0
+ else:
+ confidence_score = 0.5
+
+ return LLMPrediction(
+ instance_id=instance_id,
+ predicted_label=predicted_label,
+ confidence_score=confidence_score,
+ raw_response=content,
+ confidence_method="logprobs" if token_logprobs else "verbalized"
+ )
+
+ else:
+ raise Exception(f"HTTP {response.status_code}: {response.text}")
+
+ except Exception as e:
+ self.logger.warning(f"Logprobs attempt {attempt + 1} failed for {instance_id}: {e}")
+ if attempt < self.config.retry_attempts - 1:
+ time.sleep(self.config.retry_delay * (attempt + 1))
+ else:
+ return LLMPrediction(
+ instance_id=instance_id,
+ predicted_label='',
+ confidence_score=0.1,
+ raw_response='',
+ error_message=f"All logprob attempts failed: {e}",
+ confidence_method="logprobs"
+ )
+
+ return LLMPrediction(
+ instance_id=instance_id,
+ predicted_label='',
+ confidence_score=0.1,
+ raw_response='',
+ error_message="Unknown error",
+ confidence_method="logprobs"
+ )
+
+ def _predict_consistency(self, prompt: str, instance: Dict[str, Any]) -> LLMPrediction:
+ """Consistency-based confidence: query N times, use agreement rate.
+
+ Works with any endpoint (including Anthropic, Ollama) that doesn't
+ support logprobs. Confidence = fraction of samples that agree on
+ the most common label.
+ """
+ instance_id = instance.get('id', 'unknown')
+ n_samples = self.config.consistency_samples
+
+ labels = []
+ raw_responses = []
+
+ for _ in range(n_samples):
+ try:
+ payload = {
+ "model": self.config.model_name,
+ "messages": [
+ {"role": "system", "content": "You are a helpful assistant that provides structured JSON responses."},
+ {"role": "user", "content": prompt}
+ ],
+ "max_tokens": self.config.max_tokens,
+ "temperature": max(0.5, self.config.temperature), # Need some randomness
+ "response_format": {"type": "json_object"}
+ }
+
+ response = self.session.post(
+ self.config.endpoint_url,
+ json=payload,
+ timeout=self.config.timeout
+ )
+
+ if response.status_code == 200:
+ result = response.json()
+ if 'choices' in result and len(result['choices']) > 0:
+ content = result['choices'][0]['message']['content']
+ raw_responses.append(content)
+ try:
+ parsed = _loads_lenient(content)
+ labels.append(parsed.get('label', ''))
+ except json.JSONDecodeError:
+ pass
+
+ except Exception as e:
+ self.logger.debug(f"Consistency sample failed for {instance_id}: {e}")
+
+ if not labels:
+ return LLMPrediction(
+ instance_id=instance_id,
+ predicted_label='',
+ confidence_score=0.1,
+ raw_response='',
+ error_message="All consistency samples failed",
+ confidence_method="consistency"
+ )
+
+ # Most common label
+ label_counts = Counter(labels)
+ predicted_label, count = label_counts.most_common(1)[0]
+ confidence_score = count / len(labels)
+
+ return LLMPrediction(
+ instance_id=instance_id,
+ predicted_label=predicted_label,
+ confidence_score=confidence_score,
+ raw_response=raw_responses[0] if raw_responses else '',
+ confidence_method="consistency"
+ )
+
+ def _extract_from_raw_response(self, raw_response: str, instance_id: str) -> LLMPrediction:
+ """Extract prediction from raw response when JSON parsing fails."""
+ try:
+ # Try to find label and confidence in the raw text
+ lines = raw_response.lower().split('\n')
+
+ predicted_label = ''
+ confidence_score = 0.1
+
+ for line in lines:
+ if 'label' in line and ':' in line:
+ label_part = line.split(':', 1)[1].strip().strip('"\'')
+ if label_part:
+ predicted_label = label_part
+
+ if 'confidence' in line and ':' in line:
+ conf_part = line.split(':', 1)[1].strip()
+ try:
+ conf_value = float(conf_part)
+ confidence_score = max(0.1, min(1.0, conf_value / 10.0))
+ except ValueError:
+ pass
+
+ return LLMPrediction(
+ instance_id=instance_id,
+ predicted_label=predicted_label,
+ confidence_score=confidence_score,
+ raw_response=raw_response
+ )
+
+ except Exception as e:
+ self.logger.error(f"Failed to extract from raw response for instance {instance_id}: {e}")
+ return LLMPrediction(
+ instance_id=instance_id,
+ predicted_label='',
+ confidence_score=0.1,
+ raw_response=raw_response,
+ error_message=f"Failed to extract prediction: {e}"
+ )
+
+ def calculate_confidence_distribution(self, predictions: List[LLMPrediction]) -> Dict[str, float]:
+ """Calculate confidence score distribution from predictions."""
+ if not predictions:
+ return {}
+
+ # Filter out predictions with errors
+ valid_predictions = [p for p in predictions if p.error_message is None]
+
+ if not valid_predictions:
+ return {}
+
+ confidence_scores = [p.confidence_score for p in valid_predictions]
+
+ # Create histogram bins
+ bins = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
+ hist, _ = np.histogram(confidence_scores, bins=bins)
+
+ # Convert to percentages
+ total = len(confidence_scores)
+ distribution = {}
+ for i, count in enumerate(hist):
+ bin_label = f"{bins[i]:.1f}-{bins[i+1]:.1f}"
+ distribution[bin_label] = (count / total) * 100 if total > 0 else 0
+
+ return distribution
+
+ def get_prediction_stats(self, predictions: List[LLMPrediction]) -> Dict[str, Any]:
+ """Get statistics about the predictions."""
+ if not predictions:
+ return {
+ "total_predictions": 0,
+ "successful_predictions": 0,
+ "error_rate": 0.0,
+ "average_confidence": 0.0,
+ "confidence_distribution": {}
+ }
+
+ total = len(predictions)
+ successful = len([p for p in predictions if p.error_message is None])
+ error_rate = (total - successful) / total if total > 0 else 0.0
+
+ valid_predictions = [p for p in predictions if p.error_message is None]
+ average_confidence = np.mean([p.confidence_score for p in valid_predictions]) if valid_predictions else 0.0
+
+ confidence_distribution = self.calculate_confidence_distribution(predictions)
+
+ return {
+ "total_predictions": total,
+ "successful_predictions": successful,
+ "error_rate": error_rate,
+ "average_confidence": average_confidence,
+ "confidence_distribution": confidence_distribution
+ }
+
+
+class MockLLMActiveLearning(LLMActiveLearning):
+ """
+ Mock LLM implementation for testing and development.
+
+ This class provides realistic mock responses for testing active learning
+ without requiring an actual LLM endpoint.
+ """
+
+ def __init__(self, config: LLMConfig):
+ super().__init__(config)
+ self.logger.info("Using Mock LLM for active learning")
+
+ # Mock response patterns
+ self._mock_responses = [
+ {"label": "positive", "confidence": 8, "reasoning": "Clear positive sentiment"},
+ {"label": "negative", "confidence": 7, "reasoning": "Negative tone detected"},
+ {"label": "neutral", "confidence": 6, "reasoning": "Balanced perspective"},
+ {"label": "positive", "confidence": 9, "reasoning": "Very positive language"},
+ {"label": "negative", "confidence": 5, "reasoning": "Somewhat negative"},
+ {"label": "neutral", "confidence": 4, "reasoning": "Mixed signals"},
+ {"label": "positive", "confidence": 3, "reasoning": "Uncertain positive"},
+ {"label": "negative", "confidence": 8, "reasoning": "Clearly negative"},
+ {"label": "neutral", "confidence": 7, "reasoning": "Neutral stance"},
+ {"label": "positive", "confidence": 6, "reasoning": "Moderately positive"}
+ ]
+ self._response_index = 0
+
+ def _test_connection(self):
+ """Mock connection test."""
+ self.logger.info("Mock LLM connection test successful")
+
+ def _predict_single(self, prompt: str, instance: Dict[str, Any]) -> LLMPrediction:
+ """Make a mock prediction."""
+ instance_id = instance.get('id', 'unknown')
+
+ # Simulate processing time
+ time.sleep(0.1)
+
+ # Get next mock response
+ mock_response = self._mock_responses[self._response_index % len(self._mock_responses)]
+ self._response_index += 1
+
+ # Add some randomness to confidence scores
+ confidence_variation = np.random.normal(0, 0.1)
+ confidence_score = max(0.1, min(1.0, (mock_response['confidence'] / 10.0) + confidence_variation))
+
+ return LLMPrediction(
+ instance_id=instance_id,
+ predicted_label=mock_response['label'],
+ confidence_score=confidence_score,
+ raw_response=json.dumps(mock_response)
+ )
+
+
+def create_llm_active_learning(config: Dict[str, Any]) -> LLMActiveLearning:
+ """
+ Factory function to create LLM active learning instance.
+
+ Args:
+ config: LLM configuration dictionary
+
+ Returns:
+ LLMActiveLearning: Configured LLM active learning instance
+ """
+ llm_config = LLMConfig(
+ endpoint_url=config.get('endpoint_url', ''),
+ model_name=config.get('model_name', ''),
+ max_tokens=config.get('max_tokens', 512),
+ temperature=config.get('temperature', 0.1),
+ timeout=config.get('timeout', 30),
+ batch_size=config.get('batch_size', 10),
+ retry_attempts=config.get('retry_attempts', 3),
+ retry_delay=config.get('retry_delay', 1.0),
+ max_instances_per_request=config.get('max_instances_per_request', 5),
+ confidence_method=config.get('confidence_method', 'verbalized'),
+ consistency_samples=config.get('consistency_samples', 3),
+ )
+
+ # Use mock implementation for testing or when endpoint is not available
+ if config.get('use_mock', False) or not llm_config.endpoint_url:
+ return MockLLMActiveLearning(llm_config)
+ else:
+ return LLMActiveLearning(llm_config)
\ No newline at end of file
diff --git a/potato/ai/ollama_endpoint.py b/potato/ai/ollama_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..efffafa80974ec6d5963d0a8265d0efb55c6ce8c
--- /dev/null
+++ b/potato/ai/ollama_endpoint.py
@@ -0,0 +1,160 @@
+"""
+Ollama AI endpoint implementation.
+
+This module provides integration with Ollama for local LLM inference.
+"""
+
+import json
+from typing import Dict, List, Optional, Type
+import ollama
+from pydantic import BaseModel
+from .ai_endpoint import BaseAIEndpoint, AIEndpointRequestError, ModelCapabilities
+import re
+
+
+DEFAULT_MODEL = "llama3.2"
+
+
+class OllamaEndpoint(BaseAIEndpoint):
+ """Ollama endpoint for local LLM inference."""
+
+ # Capabilities declaration for text-based Ollama models
+ CAPABILITIES = ModelCapabilities(
+ text_generation=True,
+ vision_input=False,
+ bounding_box_output=False,
+ text_classification=True,
+ image_classification=False,
+ rationale_generation=True,
+ keyword_extraction=True,
+ )
+
+ def _initialize_client(self) -> None:
+ """Initialize the Ollama client."""
+ # Default timeout of 60 seconds for local inference (can be slower)
+ timeout = self.ai_config.get("timeout", 60)
+ host = self.ai_config.get("base_url", "http://localhost:11434")
+
+ # Create client with timeout
+ self.client = ollama.Client(host=host, timeout=timeout)
+
+ # Check if Ollama is available
+ try:
+ self.client.list()
+ except Exception as e:
+ raise AIEndpointRequestError(f"Failed to connect to Ollama: {e}")
+
+ def _get_default_model(self) -> str:
+ """Get the default Ollama model."""
+ return DEFAULT_MODEL
+
+ def query(self, prompt: str, output_format: Type[BaseModel]) -> str:
+ """
+ Send a query to Ollama and return the response.
+
+ Args:
+ prompt: The prompt to send to the model
+
+ Returns:
+ The model's response as a string
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ import logging
+ logger = logging.getLogger(__name__)
+
+ try:
+ logger.debug(f"[Ollama] Querying model: {self.model}")
+ logger.debug(f"[Ollama] Prompt (first 200 chars): {prompt[:200]}...")
+
+ options = {
+ 'temperature': self.temperature,
+ 'num_predict': self.max_tokens
+ }
+
+ # Think mode: configurable via ai_config['think'], defaults to False
+ # for fast structured output (thinking wastes tokens on JSON tasks)
+ think = self.ai_config.get('think', False)
+
+ response = self.client.chat(
+ model=self.model,
+ messages=[{'role': 'user', 'content': prompt}],
+ options=options,
+ format=output_format.model_json_schema(),
+ think=think,
+ )
+
+ # Log full response structure for debugging
+ logger.debug(f"[Ollama] Response type: {type(response)}")
+ logger.debug(f"[Ollama] Full response: {response}")
+
+ # Get the message object - handle both dict and object access
+ message = response.get('message') if hasattr(response, 'get') else getattr(response, 'message', None)
+ if message is None:
+ raise AIEndpointRequestError("No message in Ollama response")
+
+ # Get content - handle both dict and object access
+ content = message.get('content') if hasattr(message, 'get') else getattr(message, 'content', None)
+
+ # Some models put response in 'thinking' field - check that too
+ if not content and hasattr(message, 'thinking') and message.thinking:
+ logger.warning("[Ollama] Content empty but thinking field has data - model may need think=False")
+ # Try to extract JSON from thinking field as fallback
+ thinking_text = message.thinking
+ # Look for JSON in the thinking text
+ import re
+ json_match = re.search(r'\{[^{}]*\}', thinking_text)
+ if json_match:
+ content = json_match.group(0)
+ logger.debug(f"[Ollama] Extracted JSON from thinking: {content}")
+
+ logger.debug(f"[Ollama] Content type: {type(content)}")
+ logger.debug(f"[Ollama] Content value: {repr(content)[:200] if content else 'EMPTY'}")
+
+ # If content is already a dict (structured output), return it directly
+ if isinstance(content, dict):
+ logger.debug("[Ollama] Content is already a dict, returning directly")
+ return content
+
+ # Parse response using the base class's robust parser
+ # (handles truncated JSON, markdown blocks, plain text, etc.)
+ if content:
+ logger.debug(f"[Ollama] Response content (first 500 chars): {str(content)[:500]}")
+ return self.parseStringToJson(content)
+ else:
+ raise AIEndpointRequestError("Empty content from Ollama - try a different model or disable thinking mode")
+ except Exception as e:
+ logger.error(f"[Ollama] Request failed: {e}")
+ raise AIEndpointRequestError(f"Ollama request failed: {e}")
+
+ def chat_query(self, messages: List[Dict[str, str]]) -> str:
+ """Send a multi-turn chat to Ollama using native chat API."""
+ import logging
+ logger = logging.getLogger(__name__)
+
+ try:
+ options = {
+ 'temperature': self.temperature,
+ 'num_predict': self.max_tokens,
+ }
+
+ think = self.ai_config.get('think', False)
+ response = self.client.chat(
+ model=self.model,
+ messages=messages,
+ options=options,
+ think=think,
+ )
+
+ message = response.get('message') if hasattr(response, 'get') else getattr(response, 'message', None)
+ if message is None:
+ raise AIEndpointRequestError("No message in Ollama chat response")
+
+ content = message.get('content') if hasattr(message, 'get') else getattr(message, 'content', None)
+ return content or ""
+ except Exception as e:
+ logger.error(f"[Ollama] Chat request failed: {e}")
+ raise AIEndpointRequestError(f"Ollama chat request failed: {e}")
+
+
diff --git a/potato/ai/ollama_vision_endpoint.py b/potato/ai/ollama_vision_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..bfc0a57405fce8b797ef397ae7df6dbfe0a719b6
--- /dev/null
+++ b/potato/ai/ollama_vision_endpoint.py
@@ -0,0 +1,313 @@
+"""
+Ollama Vision AI Endpoint
+
+This module provides integration with Ollama vision models for local
+visual AI inference. Supports LLaVA, Llama 3.2 Vision, BakLLaVA, and Qwen-VL models.
+"""
+
+import base64
+import json
+import logging
+from typing import Any, Dict, List, Type, Union
+
+from pydantic import BaseModel
+
+from .ai_endpoint import AIEndpointRequestError, ImageData, ModelCapabilities
+from .visual_ai_endpoint import BaseVisualAIEndpoint
+
+logger = logging.getLogger(__name__)
+
+# Default vision model
+DEFAULT_MODEL = "llava:latest"
+
+# Models known to support vision (used only for warning suppression, not to restrict usage)
+# Any Ollama model can be used - this list just prevents "may not support vision" warnings
+# for models we know are vision-capable
+VISION_MODELS = [
+ # LLaVA family
+ "llava",
+ "llava-llama3",
+ "llava-phi3",
+ "bakllava",
+ # Llama Vision
+ "llama3.2-vision",
+ # Qwen Vision-Language
+ "qwen2.5-vl",
+ "qwen2-vl",
+ "qwen3-vl",
+ # Other vision models
+ "moondream",
+ "minicpm-v",
+ "gemma3", # Gemma 3 has vision capabilities
+ "gemma4", # Gemma 4 family (e.g. gemma4:e4b) supports vision
+]
+
+
+class OllamaVisionEndpoint(BaseVisualAIEndpoint):
+ """
+ Ollama Vision endpoint for multimodal local inference.
+
+ Supports vision-capable models like LLaVA, Llama 3.2 Vision, BakLLaVA.
+ Images are sent as base64 in the 'images' field.
+
+ Configuration options:
+ - model: Vision model to use (default: llava:latest)
+ - base_url: Ollama server URL (default: http://localhost:11434)
+ - timeout: Request timeout in seconds (default: 120)
+ - max_tokens: Maximum response tokens (default: 500)
+ - temperature: Sampling temperature (default: 0.1)
+ """
+
+ # Capabilities declaration for vision-capable Ollama models (LLaVA, Qwen-VL, etc.)
+ # Note: VLLMs can generate text about images but cannot do precise bounding box detection
+ # Keyword extraction is disabled because it doesn't apply to image content
+ CAPABILITIES = ModelCapabilities(
+ text_generation=True,
+ vision_input=True,
+ bounding_box_output=False, # VLLMs are not reliable for precise bbox coordinates
+ text_classification=True,
+ image_classification=True,
+ rationale_generation=True,
+ keyword_extraction=False, # Keywords don't apply to images
+ )
+
+ def _initialize_client(self) -> None:
+ """Initialize the Ollama client."""
+ try:
+ import ollama
+ except ImportError:
+ raise AIEndpointRequestError(
+ "ollama package is required. Install it with: pip install ollama"
+ )
+
+ timeout = self.ai_config.get("timeout", 120) # Vision models can be slower
+ host = self.ai_config.get("base_url", "http://localhost:11434")
+
+ self.client = ollama.Client(host=host, timeout=timeout)
+
+ # Verify connection and model availability
+ try:
+ models = self.client.list()
+ logger.info(f"Connected to Ollama at {host}")
+
+ # Check if the specified model is a known vision model
+ # This is just informational - any model can be used
+ model_lower = self.model.lower()
+ is_known_vision_model = any(vm in model_lower for vm in VISION_MODELS)
+ if not is_known_vision_model:
+ logger.info(
+ f"Model '{self.model}' not in known vision models list. "
+ f"This is fine if it supports vision - proceeding anyway."
+ )
+
+ except Exception as e:
+ raise AIEndpointRequestError(f"Failed to connect to Ollama: {e}")
+
+ def _get_default_model(self) -> str:
+ """Get the default vision model."""
+ return DEFAULT_MODEL
+
+ def query(self, prompt: str, output_format: Type[BaseModel]) -> Any:
+ """
+ Standard text query (falls back to text-only mode).
+
+ For vision tasks, use query_with_image() instead.
+ """
+ try:
+ options = {
+ 'temperature': self.temperature,
+ 'num_predict': self.max_tokens
+ }
+
+ response = self.client.chat(
+ model=self.model,
+ messages=[{'role': 'user', 'content': prompt}],
+ options=options,
+ format=output_format.model_json_schema(),
+ think=False,
+ )
+
+ message = response.get('message') if hasattr(response, 'get') else getattr(response, 'message', None)
+ if message is None:
+ raise AIEndpointRequestError("No message in Ollama response")
+
+ content = message.get('content') if hasattr(message, 'get') else getattr(message, 'content', None)
+
+ if isinstance(content, dict):
+ return content
+
+ if content:
+ return self.parseStringToJson(content)
+ else:
+ raise AIEndpointRequestError("Empty content from Ollama")
+
+ except Exception as e:
+ raise AIEndpointRequestError(f"Ollama query failed: {e}")
+
+ def query_with_image(
+ self,
+ prompt: str,
+ image_data: Union[ImageData, List[ImageData]],
+ output_format: Type[BaseModel]
+ ) -> Any:
+ """
+ Send a query with image(s) to Ollama vision model.
+
+ Args:
+ prompt: Text prompt describing what to analyze
+ image_data: Single ImageData or list of ImageData
+ output_format: Pydantic model for structured output
+
+ Returns:
+ Parsed response according to output_format
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ try:
+ # Prepare images
+ images = [image_data] if isinstance(image_data, ImageData) else image_data
+
+ # Convert to base64 if needed
+ image_base64_list = []
+ for img in images:
+ b64_data = self._get_base64_image(img)
+ image_base64_list.append(b64_data)
+
+ # Build message with images
+ options = {
+ 'temperature': self.temperature,
+ 'num_predict': self.max_tokens
+ }
+
+ # Ollama expects images as a list of base64 strings
+ response = self.client.chat(
+ model=self.model,
+ messages=[{
+ 'role': 'user',
+ 'content': prompt,
+ 'images': image_base64_list
+ }],
+ options=options,
+ format=output_format.model_json_schema(),
+ think=False,
+ )
+
+ logger.debug(f"Ollama vision response type: {type(response)}")
+
+ # Extract content from response
+ message = response.get('message') if hasattr(response, 'get') else getattr(response, 'message', None)
+ if message is None:
+ raise AIEndpointRequestError("No message in Ollama vision response")
+
+ content = message.get('content') if hasattr(message, 'get') else getattr(message, 'content', None)
+
+ logger.debug(f"Ollama vision content type: {type(content)}")
+
+ # Parse response
+ if isinstance(content, dict):
+ return content
+
+ if content:
+ return self.parseStringToJson(content)
+ else:
+ raise AIEndpointRequestError("Empty content from Ollama vision model")
+
+ except AIEndpointRequestError:
+ raise
+ except Exception as e:
+ logger.error(f"Ollama vision query failed: {e}")
+ import traceback
+ logger.error(traceback.format_exc())
+ raise AIEndpointRequestError(f"Ollama vision query failed: {e}")
+
+ def _get_base64_image(self, image_data: ImageData) -> str:
+ """
+ Get base64-encoded image data.
+
+ Args:
+ image_data: ImageData object
+
+ Returns:
+ Base64-encoded image string (without data URL prefix)
+ """
+ if image_data.source == "base64":
+ # Already base64, just return the data
+ return image_data.data
+
+ elif image_data.source == "url":
+ # Download and convert to base64
+ downloaded = self.download_image_to_base64(image_data.data)
+ return downloaded.data
+
+ else:
+ raise AIEndpointRequestError(f"Unknown image source: {image_data.source}")
+
+ def analyze_image(
+ self,
+ image_path_or_url: str,
+ prompt: str,
+ output_format: Type[BaseModel] = None
+ ) -> Any:
+ """
+ Convenience method for analyzing a single image.
+
+ Args:
+ image_path_or_url: Path to image file or URL
+ prompt: Analysis prompt
+ output_format: Optional output format model
+
+ Returns:
+ Analysis result
+ """
+ # Prepare image data
+ if image_path_or_url.startswith(("http://", "https://")):
+ image_data = self.download_image_to_base64(image_path_or_url)
+ else:
+ image_data = self.encode_image_to_base64(image_path_or_url)
+
+ # Use a generic format if not specified
+ if output_format is None:
+ from .prompt.models_module import GeneralHintFormat
+ output_format = GeneralHintFormat
+
+ return self.query_with_image(prompt, image_data, output_format)
+
+ def describe_image(self, image_path_or_url: str) -> str:
+ """
+ Get a natural language description of an image.
+
+ Args:
+ image_path_or_url: Path to image file or URL
+
+ Returns:
+ Text description of the image
+ """
+ # Use a simple model that returns text
+ class DescriptionFormat(BaseModel):
+ description: str
+
+ result = self.analyze_image(
+ image_path_or_url,
+ "Describe this image in detail. What objects, people, or scenes do you see?",
+ DescriptionFormat
+ )
+
+ if isinstance(result, dict) and "description" in result:
+ return result["description"]
+ return str(result)
+
+ def health_check(self) -> bool:
+ """
+ Check if the Ollama vision model is available.
+
+ Returns:
+ True if model is ready, False otherwise
+ """
+ try:
+ # Try to list models
+ self.client.list()
+ return True
+ except Exception as e:
+ logger.error(f"Ollama vision health check failed: {e}")
+ return False
diff --git a/potato/ai/openai_endpoint.py b/potato/ai/openai_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..976bd861caba8ded2e07a7f93b74971f12944f7a
--- /dev/null
+++ b/potato/ai/openai_endpoint.py
@@ -0,0 +1,94 @@
+"""
+OpenAI AI endpoint implementation.
+
+This module provides integration with OpenAI's API for LLM inference.
+"""
+
+import os
+from typing import Dict, List
+from openai import OpenAI
+from .ai_endpoint import BaseAIEndpoint, AIEndpointRequestError, ModelCapabilities
+
+DEFAULT_MODEL = "gpt-4o-mini"
+
+
+class OpenAIEndpoint(BaseAIEndpoint):
+ """OpenAI endpoint for cloud-based LLM inference."""
+
+ # Capabilities declaration for text-based OpenAI models
+ CAPABILITIES = ModelCapabilities(
+ text_generation=True,
+ vision_input=False,
+ bounding_box_output=False,
+ text_classification=True,
+ image_classification=False,
+ rationale_generation=True,
+ keyword_extraction=True,
+ )
+
+ def _initialize_client(self) -> None:
+ """Initialize the OpenAI client."""
+ # OpenAI-compatible servers (vLLM, llama.cpp, etc.) ignore the key
+ # but the SDK rejects an empty string, so accept a placeholder.
+ api_key = self.ai_config.get("api_key") or os.environ.get(
+ "OPENAI_API_KEY", ""
+ )
+ base_url = self.ai_config.get("base_url")
+ if not api_key:
+ if base_url:
+ api_key = "EMPTY" # non-empty placeholder for local servers
+ else:
+ raise AIEndpointRequestError("OpenAI API key is required")
+
+ # Default timeout of 30 seconds, configurable via ai_config
+ timeout = self.ai_config.get("timeout", 30)
+ client_kwargs = {"api_key": api_key, "timeout": timeout}
+ # Honor a custom base_url so this endpoint can target any
+ # OpenAI-compatible server (previously ignored -> always hit
+ # api.openai.com even when a local base_url was configured).
+ if base_url:
+ client_kwargs["base_url"] = base_url
+ self.client = OpenAI(**client_kwargs)
+
+ def _get_default_model(self) -> str:
+ """Get the default OpenAI model."""
+ return DEFAULT_MODEL
+
+ def query(self, prompt: str, output_format: dict) -> str:
+ """
+ Send a query to OpenAI and return the response.
+
+ Args:
+ prompt: The prompt to send to the model
+
+ Returns:
+ The model's response as a string
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ try:
+ response = self.client.chat.completions.create(
+ model=self.model,
+ messages=[{"role": "user", "content": prompt}],
+ max_tokens=self.max_tokens,
+ temperature=self.temperature,
+ text_format=output_format.model_json_schema(),
+ )
+ return response.choices[0].message.content
+ except Exception as e:
+ raise AIEndpointRequestError(f"OpenAI request failed: {e}")
+
+ def chat_query(self, messages: List[Dict[str, str]]) -> str:
+ """Send a multi-turn chat to OpenAI using native messages API."""
+ try:
+ response = self.client.chat.completions.create(
+ model=self.model,
+ messages=messages,
+ max_tokens=self.max_tokens,
+ temperature=self.temperature,
+ )
+ return response.choices[0].message.content
+ except Exception as e:
+ raise AIEndpointRequestError(f"OpenAI chat request failed: {e}")
+
diff --git a/potato/ai/openai_vision_endpoint.py b/potato/ai/openai_vision_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..bab3d057461d6f4c6daf3247c14feede752edf35
--- /dev/null
+++ b/potato/ai/openai_vision_endpoint.py
@@ -0,0 +1,324 @@
+"""
+OpenAI Vision AI Endpoint
+
+This module provides integration with OpenAI's vision models (GPT-4o, GPT-4o-mini)
+for visual analysis and annotation assistance.
+"""
+
+import base64
+import logging
+from typing import Any, Dict, List, Type, Union
+
+from pydantic import BaseModel
+
+from .ai_endpoint import AIEndpointRequestError, ImageData, ModelCapabilities
+from .visual_ai_endpoint import BaseVisualAIEndpoint
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = "gpt-4o"
+
+
+class OpenAIVisionEndpoint(BaseVisualAIEndpoint):
+ """
+ OpenAI Vision endpoint for GPT-4o and GPT-4o-mini vision capabilities.
+
+ Supports both URL and base64 image inputs using the image_url content type.
+
+ Configuration options:
+ - model: Model to use (gpt-4o, gpt-4o-mini) (default: gpt-4o)
+ - api_key: OpenAI API key (can also use OPENAI_API_KEY env var)
+ - max_tokens: Maximum response tokens (default: 1000)
+ - temperature: Sampling temperature (default: 0.1)
+ - detail: Image detail level - 'low', 'high', or 'auto' (default: auto)
+ """
+
+ # Capabilities declaration for OpenAI vision models (GPT-4o, GPT-4o-mini)
+ # These models can understand images and generate text but bounding boxes are approximate
+ CAPABILITIES = ModelCapabilities(
+ text_generation=True,
+ vision_input=True,
+ bounding_box_output=False, # GPT-4V bboxes are approximate, not precise
+ text_classification=True,
+ image_classification=True,
+ rationale_generation=True,
+ keyword_extraction=False, # Keywords don't apply to images
+ )
+
+ def _initialize_client(self) -> None:
+ """Initialize the OpenAI client."""
+ try:
+ import openai
+ except ImportError:
+ raise AIEndpointRequestError(
+ "openai package is required. Install it with: pip install openai"
+ )
+
+ import os
+
+ api_key = self.ai_config.get("api_key") or os.environ.get("OPENAI_API_KEY")
+ if not api_key:
+ raise AIEndpointRequestError(
+ "OpenAI API key is required. Set it in config or OPENAI_API_KEY env var."
+ )
+
+ timeout = self.ai_config.get("timeout", 60)
+ self.detail = self.ai_config.get("detail", "auto")
+
+ self.client = openai.OpenAI(api_key=api_key, timeout=timeout)
+ logger.info(f"OpenAI Vision client initialized with model: {self.model}")
+
+ def _get_default_model(self) -> str:
+ """Get the default OpenAI vision model."""
+ return DEFAULT_MODEL
+
+ def query(self, prompt: str, output_format: Type[BaseModel]) -> Any:
+ """
+ Standard text query without images.
+
+ Args:
+ prompt: Text prompt
+ output_format: Pydantic model for structured output
+
+ Returns:
+ Parsed response
+ """
+ try:
+ response = self.client.chat.completions.create(
+ model=self.model,
+ messages=[{"role": "user", "content": prompt}],
+ max_tokens=self.max_tokens,
+ temperature=self.temperature,
+ response_format={"type": "json_object"},
+ )
+
+ content = response.choices[0].message.content
+ return self.parseStringToJson(content)
+
+ except Exception as e:
+ raise AIEndpointRequestError(f"OpenAI query failed: {e}")
+
+ def query_with_image(
+ self,
+ prompt: str,
+ image_data: Union[ImageData, List[ImageData]],
+ output_format: Type[BaseModel]
+ ) -> Any:
+ """
+ Send a query with image(s) to OpenAI vision model.
+
+ Args:
+ prompt: Text prompt describing what to analyze
+ image_data: Single ImageData or list of ImageData
+ output_format: Pydantic model for structured output
+
+ Returns:
+ Parsed response according to output_format
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ try:
+ # Prepare images
+ images = [image_data] if isinstance(image_data, ImageData) else image_data
+
+ # Build content array with text and images
+ content = [{"type": "text", "text": prompt}]
+
+ for img in images:
+ image_content = self._build_image_content(img)
+ content.append(image_content)
+
+ # Make request
+ response = self.client.chat.completions.create(
+ model=self.model,
+ messages=[{"role": "user", "content": content}],
+ max_tokens=self.max_tokens,
+ temperature=self.temperature,
+ response_format={"type": "json_object"},
+ )
+
+ response_content = response.choices[0].message.content
+ logger.debug(f"OpenAI vision response: {response_content[:500] if response_content else 'empty'}")
+
+ return self.parseStringToJson(response_content)
+
+ except AIEndpointRequestError:
+ raise
+ except Exception as e:
+ logger.error(f"OpenAI vision query failed: {e}")
+ import traceback
+ logger.error(traceback.format_exc())
+ raise AIEndpointRequestError(f"OpenAI vision query failed: {e}")
+
+ def _build_image_content(self, image_data: ImageData) -> Dict[str, Any]:
+ """
+ Build image content block for OpenAI API.
+
+ Args:
+ image_data: ImageData object
+
+ Returns:
+ Dict with type: "image_url" and image_url content
+ """
+ if image_data.source == "url":
+ # Direct URL reference
+ return {
+ "type": "image_url",
+ "image_url": {
+ "url": image_data.data,
+ "detail": self.detail
+ }
+ }
+
+ elif image_data.source == "base64":
+ # Data URL format
+ mime_type = image_data.mime_type or "image/jpeg"
+ data_url = f"data:{mime_type};base64,{image_data.data}"
+
+ return {
+ "type": "image_url",
+ "image_url": {
+ "url": data_url,
+ "detail": self.detail
+ }
+ }
+
+ else:
+ raise AIEndpointRequestError(f"Unknown image source: {image_data.source}")
+
+ def analyze_image(
+ self,
+ image_path_or_url: str,
+ prompt: str,
+ output_format: Type[BaseModel] = None
+ ) -> Any:
+ """
+ Convenience method for analyzing a single image.
+
+ Args:
+ image_path_or_url: Path to image file or URL
+ prompt: Analysis prompt
+ output_format: Optional output format model
+
+ Returns:
+ Analysis result
+ """
+ # Prepare image data - use URL directly if possible
+ if image_path_or_url.startswith(("http://", "https://")):
+ image_data = self.create_url_image_data(image_path_or_url)
+ else:
+ image_data = self.encode_image_to_base64(image_path_or_url)
+
+ # Use a generic format if not specified
+ if output_format is None:
+ from .prompt.models_module import GeneralHintFormat
+ output_format = GeneralHintFormat
+
+ return self.query_with_image(prompt, image_data, output_format)
+
+ def detect_objects(
+ self,
+ image_path_or_url: str,
+ labels: List[str] = None
+ ) -> Dict[str, Any]:
+ """
+ Detect objects in an image and return bounding boxes.
+
+ Args:
+ image_path_or_url: Path to image file or URL
+ labels: Optional list of labels to detect
+
+ Returns:
+ Dict with detections list
+ """
+ from .prompt.models_module import VisualDetectionFormat
+
+ labels_str = ", ".join(labels) if labels else "all visible objects"
+
+ prompt = f"""Analyze this image and detect objects. For each object, provide:
+1. The label (from: {labels_str})
+2. A bounding box with normalized coordinates (0-1 range)
+3. Confidence score (0-1)
+
+Return JSON with this structure:
+{{
+ "detections": [
+ {{
+ "label": "object_name",
+ "bbox": {{"x": 0.1, "y": 0.2, "width": 0.3, "height": 0.4}},
+ "confidence": 0.95
+ }}
+ ]
+}}
+
+Coordinates are normalized (0-1) where x,y is the top-left corner.
+Only include objects you can clearly identify with confidence > 0.5."""
+
+ # Prepare image
+ if image_path_or_url.startswith(("http://", "https://")):
+ image_data = self.create_url_image_data(image_path_or_url)
+ else:
+ image_data = self.encode_image_to_base64(image_path_or_url)
+
+ return self.query_with_image(prompt, image_data, VisualDetectionFormat)
+
+ def describe_region(
+ self,
+ image_path_or_url: str,
+ region: Dict[str, float],
+ labels: List[str] = None
+ ) -> Dict[str, Any]:
+ """
+ Describe or classify a specific region in an image.
+
+ Args:
+ image_path_or_url: Path to image file or URL
+ region: Dict with x, y, width, height (normalized 0-1)
+ labels: Optional list of possible labels
+
+ Returns:
+ Classification result with suggested label and confidence
+ """
+ labels_str = ", ".join(labels) if labels else "any appropriate category"
+
+ prompt = f"""Look at the region marked in this image:
+- Region: x={region['x']:.2f}, y={region['y']:.2f}, width={region['width']:.2f}, height={region['height']:.2f}
+(Coordinates are normalized 0-1, where 0,0 is top-left)
+
+Classify what you see in this region from these options: {labels_str}
+
+Return JSON:
+{{
+ "suggested_label": "label_name",
+ "confidence": 0.85,
+ "reasoning": "Brief explanation"
+}}"""
+
+ # Prepare image
+ if image_path_or_url.startswith(("http://", "https://")):
+ image_data = self.create_url_image_data(image_path_or_url)
+ else:
+ image_data = self.encode_image_to_base64(image_path_or_url)
+
+ class RegionClassificationFormat(BaseModel):
+ suggested_label: str
+ confidence: float
+ reasoning: str
+
+ return self.query_with_image(prompt, image_data, RegionClassificationFormat)
+
+ def health_check(self) -> bool:
+ """
+ Check if the OpenAI API is accessible.
+
+ Returns:
+ True if API is reachable, False otherwise
+ """
+ try:
+ # Simple models list check
+ self.client.models.list()
+ return True
+ except Exception as e:
+ logger.error(f"OpenAI health check failed: {e}")
+ return False
diff --git a/potato/ai/openrouter_endpoint.py b/potato/ai/openrouter_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..72ab3830fcadc30db3242abc60ea94b3cbe09a14
--- /dev/null
+++ b/potato/ai/openrouter_endpoint.py
@@ -0,0 +1,93 @@
+"""
+OpenRouter endpoint implementation.
+This module provides integration with OpenRouter's API for LLM inference.
+"""
+import requests
+from .ai_endpoint import BaseAIEndpoint, AIEndpointRequestError
+
+DEFAULT_MODEL = "openai/gpt-4o-mini"
+
+class OpenRouterEndpoint(BaseAIEndpoint):
+ """OpenRouter endpoint for cloud-based LLM inference."""
+
+ # Models that support structured output
+ STRUCTURED_OUTPUT_MODELS = {
+ "openai/gpt-4o",
+ "openai/gpt-4o-mini",
+ "openai/gpt-4-turbo",
+ "anthropic/claude-3-5-sonnet",
+ "deepseek/deepseek-r1:free"
+ }
+
+ def _initialize_client(self) -> None:
+ """Initialize the OpenAI client."""
+ api_key = self.ai_config.get("api_key", "")
+ if not api_key:
+ raise AIEndpointRequestError("OpenRouter API key is required")
+
+ def _get_default_model(self) -> str:
+ """Get the default OpenAI model."""
+ return DEFAULT_MODEL
+
+ def supports_structured_output(self) -> bool:
+ """Check if the current model supports structured output."""
+ model = self.model or DEFAULT_MODEL
+ return any(model.startswith(prefix.split('/')[0]) or model in self.STRUCTURED_OUTPUT_MODELS
+ for prefix in self.STRUCTURED_OUTPUT_MODELS)
+
+ def query(self, prompt: str, output_format: dict) -> str:
+ """
+ Send a query to OpenRouter and return the response.
+
+ Args:
+ prompt: The prompt to send to the model (as messages list or string)
+ output_format: Pydantic model for structured output
+
+ Returns:
+ The model's response as a string
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ try:
+ url = "https://openrouter.ai/api/v1/chat/completions"
+ headers = {
+ "Authorization": f"Bearer {self.ai_config.get('api_key')}",
+ "Content-Type": "application/json"
+ }
+
+ messages = [{"role": "user", "content": prompt}]
+ schema = output_format.model_json_schema()
+
+ body = {
+ "model": self.model or DEFAULT_MODEL,
+ "max_tokens": self.max_tokens,
+ "temperature": self.temperature,
+ }
+
+ # Handle structured output based on model support
+ if self.supports_structured_output():
+ body["messages"] = messages
+ body["response_format"] = {
+ "type": "json_schema",
+ "json_schema": {
+ "name": "response",
+ "schema": schema,
+ "strict": True
+ }
+ }
+ else:
+ # If model does not support structured format, just send raw prompt
+ body["messages"] = messages
+
+ r = requests.post(url, headers=headers, json=body)
+
+ if r.status_code >= 400:
+ raise AIEndpointRequestError(f"OpenRouter error {r.status_code}: {r.text}")
+
+ data = r.json()
+ if self.supports_structured_output():
+ return self.parseStringToJson(data["choices"][0]["message"]["content"])
+ return data["choices"][0]["message"]["content"]
+ except Exception as e:
+ raise AIEndpointRequestError(f"OpenRouter request failed: {e}")
\ No newline at end of file
diff --git a/potato/ai/prompt/image_annotation.json b/potato/ai/prompt/image_annotation.json
new file mode 100644
index 0000000000000000000000000000000000000000..d537d38ef35610f8316bf84c47e1ac8426ec6f25
--- /dev/null
+++ b/potato/ai/prompt/image_annotation.json
@@ -0,0 +1,44 @@
+{
+ "detect": {
+ "prompt": "TASK: Detect objects in this image that match the specified labels.\n\nDescription: ${description}\nLabels to detect: ${labels}\nMinimum confidence: ${confidence_threshold}\n\nFor each detected object, provide:\n1. The label (must be from the provided labels list)\n2. Bounding box with normalized coordinates (0-1 range)\n3. Confidence score (0-1)\n\nCoordinate system:\n- x,y is the top-left corner of the box\n- x increases left to right (0 = left edge, 1 = right edge)\n- y increases top to bottom (0 = top edge, 1 = bottom edge)\n- width and height are also normalized (0-1)\n\nReturn JSON:\n{\n \"detections\": [\n {\n \"label\": \"object_name\",\n \"bbox\": {\"x\": 0.1, \"y\": 0.2, \"width\": 0.3, \"height\": 0.4},\n \"confidence\": 0.95\n }\n ]\n}\n\nOnly include objects you can clearly identify with confidence >= ${confidence_threshold}.",
+ "output_format": "visual_detection",
+ "img": "/static/ai_assistant_img/detect.svg",
+ "name": "Detect"
+ },
+ "detection": {
+ "prompt": "TASK: Detect objects in this image that match the specified labels.\n\nDescription: ${description}\nLabels to detect: ${labels}\nMinimum confidence: ${confidence_threshold}\n\nFor each detected object, provide:\n1. The label (must be from the provided labels list)\n2. Bounding box with normalized coordinates (0-1 range)\n3. Confidence score (0-1)\n\nCoordinate system:\n- x,y is the top-left corner of the box\n- x increases left to right (0 = left edge, 1 = right edge)\n- y increases top to bottom (0 = top edge, 1 = bottom edge)\n- width and height are also normalized (0-1)\n\nReturn JSON:\n{\n \"detections\": [\n {\n \"label\": \"object_name\",\n \"bbox\": {\"x\": 0.1, \"y\": 0.2, \"width\": 0.3, \"height\": 0.4},\n \"confidence\": 0.95\n }\n ]\n}\n\nOnly include objects you can clearly identify with confidence >= ${confidence_threshold}.",
+ "output_format": "visual_detection",
+ "img": "/static/ai_assistant_img/detect.svg",
+ "name": "Detect"
+ },
+ "pre_annotate": {
+ "prompt": "TASK: Pre-annotate this image by detecting all objects that match the available labels.\n\nDescription: ${description}\nAvailable labels: ${labels}\n\nDetect ALL instances of objects matching these labels. Be thorough - it's better to include uncertain detections (the human annotator will review them).\n\nFor each detection provide:\n1. Label from the available list\n2. Bounding box (normalized 0-1 coordinates)\n3. Confidence score\n\nReturn JSON:\n{\n \"detections\": [\n {\n \"label\": \"label_name\",\n \"bbox\": {\"x\": 0.0, \"y\": 0.0, \"width\": 0.0, \"height\": 0.0},\n \"confidence\": 0.0\n }\n ]\n}\n\nInclude all potential detections, even with lower confidence. The annotator will verify.",
+ "output_format": "visual_detection",
+ "img": "/static/ai_assistant_img/auto.svg",
+ "name": "Auto"
+ },
+ "classification": {
+ "prompt": "TASK: Classify the specified region in this image.\n\nDescription: ${description}\nRegion: ${region}\nAvailable labels: ${labels}\n\nLook at the indicated region and determine which label best describes what you see there.\n\nReturn JSON:\n{\n \"suggested_label\": \"label_name\",\n \"confidence\": 0.85,\n \"reasoning\": \"Brief explanation of why this label fits\"\n}",
+ "output_format": "visual_classification",
+ "img": "/static/ai_assistant_img/classify.svg",
+ "name": "Classify"
+ },
+ "hint": {
+ "prompt": "TASK: Provide a helpful hint for annotating this image WITHOUT revealing exact answers.\n\nAnnotation task: ${description}\nAvailable labels: ${labels}\n\nProvide guidance that helps the annotator without giving away:\n- Exact object locations\n- Specific label assignments\n\nGood hints:\n- Point out relevant visual features\n- Suggest areas that deserve attention\n- Note potential challenges or ambiguities\n- Remind about edge cases\n\nReturn JSON:\n{\n \"hint\": \"Your helpful guidance here\",\n \"suggestive_choice\": \"optional_focus_area\"\n}",
+ "output_format": "default_hint",
+ "img": "/static/ai_assistant_img/blub.svg",
+ "name": "Hint"
+ },
+ "keyword": {
+ "prompt": "TASK: Identify visual keywords/features associated with each label in this image.\n\nAnnotation task: ${description}\nAvailable labels: ${labels}\n\nFor each label, identify visual cues or features that would indicate its presence.\n\nReturn JSON:\n{\n \"label_keywords\": [\n {\n \"label\": \"label_name\",\n \"keywords\": [\"visual_feature_1\", \"visual_feature_2\"]\n }\n ]\n}",
+ "output_format": "default_keyword",
+ "img": "/static/ai_assistant_img/highlight.svg",
+ "name": "Keywords"
+ },
+ "rationale": {
+ "prompt": "TASK: Provide rationale for how each label might apply to this image.\n\nAnnotation task: ${description}\nAvailable labels: ${labels}\n\nFor each label, explain what evidence in the image supports or contradicts its application.\n\nReturn JSON:\n{\n \"rationales\": [\n {\n \"label\": \"label_name\",\n \"reasoning\": \"Explanation of visual evidence for/against this label\"\n }\n ]\n}",
+ "output_format": "default_rationale",
+ "img": "/static/ai_assistant_img/question.svg",
+ "name": "Rationale"
+ }
+}
diff --git a/potato/ai/prompt/likert.json b/potato/ai/prompt/likert.json
new file mode 100644
index 0000000000000000000000000000000000000000..b7d861fe3c9737ca603eb1880cd97fecfa1fc46d
--- /dev/null
+++ b/potato/ai/prompt/likert.json
@@ -0,0 +1,20 @@
+{
+ "hint": {
+ "prompt": "TASK: Generate annotation guidance for Likert scale task.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Scale: ${min_label} (1) to ${max_label} (${size})\n- Scale points: 1, 2, 3, ..., ${size}\n\nINSTRUCTIONS:\n1. Analyze the text for features relevant to the annotation task\n2. Generate a helpful hint that guides thinking WITHOUT revealing the answer\n3. Suggest a scale position based on your analysis\n\nHINT REQUIREMENTS:\n- Focus on specific textual evidence (word choice, tone, structure)\n- Point out subtle indicators the annotator should notice\n- Be concrete and actionable, not generic\n- Guide analytical thinking without bias toward any scale position",
+ "output_format": "default_hint",
+ "img": "/static/ai_assistant_img/blub.svg"
+ },
+ "keyword": {
+ "prompt": "TASK: Extract key words/phrases that guide Likert scale annotation decisions.\n\nINPUT DETAILS:\n- Text: \"${text}\"\n- Annotation task: ${description}\n- Scale: ${min_label} (1) to ${max_label} (${size})\n- Scale range: 1, 2, 3, ..., ${size}\n\nOBJECTIVE: Identify 3-5 most significant words/phrases that directly indicate scale positioning. Output as JSON array format.\n\nSELECTION CRITERIA:\n- Words that signal intensity/degree (extremely, slightly, moderately)\n- Sentiment markers (positive/negative indicators)\n- Qualifying language (hedges, certainty markers)\n- Context-specific terminology relevant to the annotation task\n- Structural indicators (but, however, although)\n\nPRIORITIZE:\n1. Words with clear scale implications\n2. Phrases that distinguish between scale levels\n3. Contextual clues that affect interpretation\n4. Intensity modifiers and qualifiers",
+ "output_format": "default_keyword",
+ "img": "/static/ai_assistant_img/highlight.svg"
+ },
+ "rationale": {
+ "name": "Rationale",
+ "prompt": "TASK: Generate rationales for different positions on the Likert scale.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Scale: ${min_label} (1) to ${max_label} (${size})\n\nINSTRUCTIONS:\nProvide rationales for different scale positions (low, middle, high). Explain what evidence in the text could support each position. Be balanced and objective.\n\nOUTPUT FORMAT:\nReturn a JSON object with \"rationales\" array containing objects with \"label\" (scale position descriptor) and \"reasoning\" fields.\n\nEXAMPLE OUTPUT:\n{\"rationales\": [{\"label\": \"low (1-2)\", \"reasoning\": \"The negative tone and words like 'disappointing' suggest a low rating\"}, {\"label\": \"middle (3)\", \"reasoning\": \"Mixed signals with both positive and negative elements\"}, {\"label\": \"high (4-5)\", \"reasoning\": \"Strong positive language like 'excellent' supports a high rating\"}]}",
+ "output_format": "default_rationale",
+ "img": "/static/ai_assistant_img/question.svg"
+ }
+}
+
+
diff --git a/potato/ai/prompt/models_module.py b/potato/ai/prompt/models_module.py
new file mode 100644
index 0000000000000000000000000000000000000000..17d74d6e803632a9496c37f329dbc1bf4eef60d2
--- /dev/null
+++ b/potato/ai/prompt/models_module.py
@@ -0,0 +1,258 @@
+from typing import Optional, Type, Union, Dict, List
+from pydantic import BaseModel
+
+
+class GeneralHintFormat(BaseModel):
+ hint: str
+ suggestive_choice: Union[str, int]
+
+
+class LabelKeywords(BaseModel):
+ """Keywords/phrases associated with a specific label."""
+ label: str
+ keywords: List[str]
+
+
+class GeneralKeywordFormat(BaseModel):
+ """Simplified keyword format: list of label -> keywords mappings.
+
+ Example output:
+ {
+ "label_keywords": [
+ {"label": "positive", "keywords": ["great", "love it", "excellent"]},
+ {"label": "negative", "keywords": ["terrible", "awful"]}
+ ]
+ }
+ """
+ label_keywords: List[LabelKeywords]
+
+
+class GeneralRandomFormat(BaseModel):
+ """Deprecated: Use GeneralRationaleFormat instead."""
+ random: str
+
+
+class LabelRationale(BaseModel):
+ """Rationale/reasoning for why a specific label might apply."""
+ label: str
+ reasoning: str
+
+
+class GeneralRationaleFormat(BaseModel):
+ """Rationale format: explanations for how each label might apply to the text.
+
+ Example output:
+ {
+ "rationales": [
+ {"label": "positive", "reasoning": "The phrase 'excellent quality' suggests satisfaction"},
+ {"label": "negative", "reasoning": "The mention of 'delayed shipping' indicates frustration"}
+ ]
+ }
+ """
+ rationales: List[LabelRationale]
+
+
+# ============================================================================
+# Visual Annotation Output Formats
+# ============================================================================
+
+class BoundingBox(BaseModel):
+ """Normalized bounding box coordinates (0-1 range).
+
+ x, y: top-left corner position
+ width, height: box dimensions
+ All values are normalized to image dimensions (0-1).
+ """
+ x: float
+ y: float
+ width: float
+ height: float
+
+
+class Detection(BaseModel):
+ """Single object detection result.
+
+ Example:
+ {
+ "label": "person",
+ "bbox": {"x": 0.1, "y": 0.2, "width": 0.3, "height": 0.5},
+ "confidence": 0.95
+ }
+ """
+ label: str
+ bbox: BoundingBox
+ confidence: float
+
+
+class VisualDetectionFormat(BaseModel):
+ """Object detection results for an image.
+
+ Example output:
+ {
+ "detections": [
+ {"label": "car", "bbox": {"x": 0.1, "y": 0.2, "width": 0.3, "height": 0.2}, "confidence": 0.92},
+ {"label": "person", "bbox": {"x": 0.5, "y": 0.3, "width": 0.1, "height": 0.4}, "confidence": 0.87}
+ ]
+ }
+ """
+ detections: List[Detection]
+
+
+class VisualClassificationFormat(BaseModel):
+ """Classification result for an image or region.
+
+ Example output:
+ {
+ "suggested_label": "cat",
+ "confidence": 0.89,
+ "reasoning": "The image shows a feline with pointed ears and whiskers"
+ }
+ """
+ suggested_label: str
+ confidence: float
+ reasoning: Optional[str] = None
+
+
+class VideoSegment(BaseModel):
+ """Temporal segment in a video.
+
+ Times are in seconds.
+ """
+ start_time: float
+ end_time: float
+ suggested_label: str
+ confidence: float
+ description: Optional[str] = None
+
+
+class VideoSceneDetectionFormat(BaseModel):
+ """Scene/segment detection results for a video.
+
+ Example output:
+ {
+ "segments": [
+ {"start_time": 0.0, "end_time": 5.5, "suggested_label": "intro", "confidence": 0.9},
+ {"start_time": 5.5, "end_time": 15.0, "suggested_label": "action", "confidence": 0.85}
+ ]
+ }
+ """
+ segments: List[VideoSegment]
+
+
+class VideoKeyframe(BaseModel):
+ """Keyframe annotation for a video.
+
+ timestamp: Time in seconds
+ """
+ timestamp: float
+ suggested_label: str
+ confidence: float
+ reason: Optional[str] = None
+
+
+class VideoKeyframeDetectionFormat(BaseModel):
+ """Keyframe detection results for a video.
+
+ Example output:
+ {
+ "keyframes": [
+ {"timestamp": 2.5, "suggested_label": "scene_change", "confidence": 0.95, "reason": "Major visual transition"},
+ {"timestamp": 8.0, "suggested_label": "action_peak", "confidence": 0.82, "reason": "Key moment in action"}
+ ]
+ }
+ """
+ keyframes: List[VideoKeyframe]
+
+
+class TrackPosition(BaseModel):
+ """Object position in a single frame for tracking."""
+ frame_index: int
+ bbox: BoundingBox
+ confidence: float
+
+
+class ObjectTrack(BaseModel):
+ """Tracked object across multiple frames."""
+ track_id: int
+ label: str
+ positions: List[TrackPosition]
+
+
+class VideoTrackingSuggestionFormat(BaseModel):
+ """Object tracking suggestions for a video.
+
+ Example output:
+ {
+ "tracks": [
+ {
+ "track_id": 1,
+ "label": "person",
+ "positions": [
+ {"frame_index": 0, "bbox": {"x": 0.1, "y": 0.2, "width": 0.15, "height": 0.3}, "confidence": 0.9},
+ {"frame_index": 1, "bbox": {"x": 0.12, "y": 0.22, "width": 0.15, "height": 0.3}, "confidence": 0.88}
+ ]
+ }
+ ]
+ }
+ """
+ tracks: List[ObjectTrack]
+
+
+class FrameDetections(BaseModel):
+ """Detections for a single video frame."""
+ frame_index: int
+ detections: List[Detection]
+
+
+class MultiFrameDetectionFormat(BaseModel):
+ """Detection results across multiple video frames.
+
+ Used when running detection on sampled video frames.
+ """
+ frames: List[FrameDetections]
+
+
+# ============================================================================
+# Class Registry
+# ============================================================================
+
+# ============================================================================
+# Option Highlighting Output Format
+# ============================================================================
+
+class OptionHighlightFormat(BaseModel):
+ """LLM response for option highlighting.
+
+ Used to identify the most likely correct options for a discrete annotation task.
+ The highlighted options are shown at full opacity while others are dimmed.
+
+ Example output:
+ {
+ "highlighted_options": ["positive", "neutral"],
+ "confidence": 0.85
+ }
+ """
+ highlighted_options: List[str] # Top-k most likely option names/values
+ confidence: Optional[float] = None # Optional overall confidence score (0-1)
+
+
+CLASS_REGISTRY = {
+ # Text annotation formats
+ "default_hint": GeneralHintFormat,
+ "default_keyword": GeneralKeywordFormat,
+ "default_random": GeneralRandomFormat, # Keep for backwards compatibility
+ "default_rationale": GeneralRationaleFormat,
+
+ # Option highlighting format
+ "option_highlight": OptionHighlightFormat,
+
+ # Visual annotation formats - Image
+ "visual_detection": VisualDetectionFormat,
+ "visual_classification": VisualClassificationFormat,
+
+ # Visual annotation formats - Video
+ "video_scene_detection": VideoSceneDetectionFormat,
+ "video_keyframe_detection": VideoKeyframeDetectionFormat,
+ "video_tracking_suggestion": VideoTrackingSuggestionFormat,
+ "multi_frame_detection": MultiFrameDetectionFormat,
+}
\ No newline at end of file
diff --git a/potato/ai/prompt/multiselect.json b/potato/ai/prompt/multiselect.json
new file mode 100644
index 0000000000000000000000000000000000000000..44fe0a74f7f96fb43209e8674dfd9653843fbf31
--- /dev/null
+++ b/potato/ai/prompt/multiselect.json
@@ -0,0 +1,20 @@
+{
+ "hint": {
+ "prompt": "TASK: Generate guidance for multiple label selection.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nINSTRUCTIONS:\n1. Analyze text for features that may correspond to multiple labels\n2. Guide toward identifying ALL applicable categories\n3. Focus on overlapping characteristics and comprehensive analysis\n\nHINT REQUIREMENTS:\n- Identify indicators for each potential label category\n- Point out that multiple selections may be appropriate\n- Guide systematic evaluation of all label options\n- Highlight overlapping or complementary features",
+ "output_format": "default_hint",
+ "img": "/static/ai_assistant_img/blub.svg"
+ },
+ "keyword": {
+ "prompt": "TASK: Extract words/phrases that indicate multiple applicable labels.\n\nINPUT DETAILS:\n- Text: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nOBJECTIVE: Identify terms that support selection of multiple labels.\n\nSELECTION CRITERIA:\n- Words that indicate multiple categories simultaneously\n- Terms supporting different label aspects\n- Overlapping category indicators\n- Comprehensive feature markers\n- Multi-faceted descriptors",
+ "output_format": "default_keyword",
+ "img": "/static/ai_assistant_img/highlight.svg"
+ },
+ "rationale": {
+ "name": "Rationale",
+ "prompt": "TASK: Generate rationales explaining why each label might apply to this text.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nCRITICAL REQUIREMENT:\nYou MUST provide a rationale for EVERY label listed above. Your output must have exactly one entry per label.\n\nINSTRUCTIONS:\nFor EACH available label (ALL of them), provide a brief rationale explaining what evidence in the text could support selecting that label. Since multiple labels can be selected, focus on independent evidence for each label. If a label doesn't apply, explain why.\n\nOUTPUT FORMAT:\nReturn a JSON object with \"rationales\" array containing one object per label, each with \"label\" and \"reasoning\" fields.\n\nEXAMPLE OUTPUT:\n{\"rationales\": [{\"label\": \"category1\", \"reasoning\": \"The text mentions X which relates to this category\"}, {\"label\": \"category2\", \"reasoning\": \"The phrase Y suggests this also applies\"}, {\"label\": \"category3\", \"reasoning\": \"No direct evidence for this label in the text\"}]}",
+ "output_format": "default_rationale",
+ "img": "/static/ai_assistant_img/question.svg"
+ }
+}
+
+
diff --git a/potato/ai/prompt/number.json b/potato/ai/prompt/number.json
new file mode 100644
index 0000000000000000000000000000000000000000..5d74053f71ed5d6a4aba08d1749c7ef45708f3f2
--- /dev/null
+++ b/potato/ai/prompt/number.json
@@ -0,0 +1,18 @@
+{
+ "hint": {
+ "prompt": "TASK: Generate guidance for numerical value extraction/estimation.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n\nINSTRUCTIONS:\n1. Analyze the text for numerical clues or quantifiable elements\n2. Guide the annotator toward identifying the correct numerical value\n3. Focus on mathematical, statistical, or countable aspects\n\nHINT REQUIREMENTS:\n- Identify numerical indicators, quantities, or measurable elements\n- Point out calculation methods or counting strategies\n- Highlight context that affects numerical interpretation\n- Guide toward systematic analysis approach",
+ "output_format": "default_hint",
+ "img": "/static/ai_assistant_img/blub.svg"
+ },
+ "keyword": {
+ "prompt": "TASK: Extract words/phrases that contain or indicate numerical information.\n\nINPUT DETAILS:\n- Text: \"${text}\"\n- Annotation task: ${description}\n\nOBJECTIVE: Identify words/phrases that directly relate to the numerical answer.\n\nSELECTION CRITERIA:\n- Explicit numbers, quantities, or measurements\n- Words indicating amount, frequency, or degree\n- Mathematical or statistical terminology\n- Comparative language (more, less, double, half)\n- Time references, percentages, ratios",
+ "output_format": "default_keyword",
+ "img": "/static/ai_assistant_img/highlight.svg"
+ },
+ "rationale": {
+ "name": "Rationale",
+ "prompt": "TASK: Generate rationales for different approaches to determining the numerical answer.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n\nINSTRUCTIONS:\nProvide different rationales or approaches for determining the numerical value. Consider different interpretations or calculation methods if applicable.\n\nOUTPUT FORMAT:\nReturn a JSON object with \"rationales\" array containing objects with \"label\" (approach name) and \"reasoning\" fields.\n\nEXAMPLE OUTPUT:\n{\"rationales\": [{\"label\": \"literal count\", \"reasoning\": \"Counting explicit mentions gives X\"}, {\"label\": \"inclusive interpretation\", \"reasoning\": \"Including implicit references increases the count to Y\"}]}",
+ "output_format": "default_rationale",
+ "img": "/static/ai_assistant_img/question.svg"
+ }
+}
diff --git a/potato/ai/prompt/option_highlight.json b/potato/ai/prompt/option_highlight.json
new file mode 100644
index 0000000000000000000000000000000000000000..27730b7025af0b5c6452bdc4cf3dbca4e0efe7a5
--- /dev/null
+++ b/potato/ai/prompt/option_highlight.json
@@ -0,0 +1,8 @@
+{
+ "option_highlight": {
+ "prompt": "TASK: Identify the most likely correct options for an annotation task.\n\nINPUT DETAILS:\n- Content to annotate: \"${text}\"\n- Annotation task: ${description}\n- Available options: ${labels}\n- Number of options to highlight: ${top_k}\n\nINSTRUCTIONS:\n1. Analyze the content carefully in the context of the annotation task\n2. Consider which ${top_k} options are most likely to be correct based on:\n - Direct evidence in the content\n - Contextual clues and tone\n - Domain knowledge relevant to the task\n3. Select exactly ${top_k} options (or fewer if fewer options exist)\n\nIMPORTANT:\n- Return ONLY the option names/values exactly as they appear in the available options list\n- Do not modify, paraphrase, or abbreviate the option names\n- If you're uncertain, choose options that seem most plausible given the content\n\nOUTPUT FORMAT:\nReturn a JSON object with:\n- \"highlighted_options\": array of ${top_k} option names from the available options\n- \"confidence\": your confidence score from 0.0 to 1.0\n\nEXAMPLE (if options are: positive, negative, neutral):\n{\"highlighted_options\": [\"positive\", \"neutral\"], \"confidence\": 0.75}",
+ "output_format": "option_highlight",
+ "img": "/static/ai_assistant_img/highlight.svg",
+ "name": "Option Highlight"
+ }
+}
diff --git a/potato/ai/prompt/radio.json b/potato/ai/prompt/radio.json
new file mode 100644
index 0000000000000000000000000000000000000000..1793b7872edfaf7e7c07b8295cb97cfc77569c0e
--- /dev/null
+++ b/potato/ai/prompt/radio.json
@@ -0,0 +1,18 @@
+{
+ "hint": {
+ "prompt": "TASK: Generate annotation guidance for single-choice selection.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nINSTRUCTIONS:\n1. Analyze the text for features that distinguish between the available labels\n2. Generate a helpful hint that guides classification WITHOUT revealing the answer\n3. Focus on decision-making criteria between options\n\nHINT REQUIREMENTS:\n- Identify key textual indicators that differentiate between label options\n- Point out distinguishing features (style, content, context)\n- Guide analytical thinking toward the classification criteria\n- Be specific about what to examine, not generic",
+ "output_format": "default_hint",
+ "img": "/static/ai_assistant_img/blub.svg"
+ },
+ "keyword": {
+ "prompt": "TASK: Identify words or short phrases in the text that relate to each label.\n\nINPUT DETAILS:\n- Text: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nINSTRUCTIONS:\nFor each label, find words or short phrases from the text that indicate or relate to that label. Only include words/phrases that actually appear in the text. Return 1-5 keywords per label. If no words relate to a label, return an empty list for that label.\n\nEXAMPLE OUTPUT FORMAT:\n{\"label_keywords\": [{\"label\": \"positive\", \"keywords\": [\"great\", \"love it\"]}, {\"label\": \"negative\", \"keywords\": [\"terrible\"]}]}",
+ "output_format": "default_keyword",
+ "img": "/static/ai_assistant_img/highlight.svg"
+ },
+ "rationale": {
+ "name": "Rationale",
+ "prompt": "TASK: Generate rationales explaining why each label might apply to this text.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nCRITICAL REQUIREMENT:\nYou MUST provide a rationale for EVERY label listed above. Your output must have exactly one entry per label.\n\nINSTRUCTIONS:\nFor EACH available label (ALL of them), provide a brief rationale explaining what evidence in the text could support choosing that label. Be balanced and objective - present the case for each label fairly. If a label doesn't strongly apply, explain what would need to be present.\n\nOUTPUT FORMAT:\nReturn a JSON object with \"rationales\" array containing one object per label, each with \"label\" and \"reasoning\" fields.\n\nEXAMPLE (if labels are: positive, negative, neutral, mixed):\n{\"rationales\": [{\"label\": \"positive\", \"reasoning\": \"The phrase 'excellent service' suggests satisfaction\"}, {\"label\": \"negative\", \"reasoning\": \"No clear negative indicators present\"}, {\"label\": \"neutral\", \"reasoning\": \"The tone is more emotional than neutral\"}, {\"label\": \"mixed\", \"reasoning\": \"Would need both positive and negative elements\"}]}",
+ "output_format": "default_rationale",
+ "img": "/static/ai_assistant_img/question.svg"
+ }
+}
diff --git a/potato/ai/prompt/select.json b/potato/ai/prompt/select.json
new file mode 100644
index 0000000000000000000000000000000000000000..99f03d66a83c6b26c01b1f10549c0bf1898bb9de
--- /dev/null
+++ b/potato/ai/prompt/select.json
@@ -0,0 +1,18 @@
+{
+ "hint": {
+ "prompt": "TASK: Generate guidance for single selection from dropdown options.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nINSTRUCTIONS:\n1. Analyze text features that distinguish between dropdown options\n2. Generate guidance for selecting the most appropriate single option\n3. Focus on decision criteria between available choices\n\nHINT REQUIREMENTS:\n- Identify key distinguishing features between label options\n- Point out decision-making criteria for selection\n- Guide toward systematic evaluation of options\n- Highlight most relevant textual evidence",
+ "output_format": "default_hint",
+ "img": "/static/ai_assistant_img/blub.svg"
+ },
+ "keyword": {
+ "prompt": "TASK: Extract words/phrases that indicate the correct dropdown selection.\n\nINPUT DETAILS:\n- Text: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nOBJECTIVE: Identify terms that point toward the most appropriate label selection.\n\nSELECTION CRITERIA:\n- Strong indicators for specific label options\n- Discriminating features between choices\n- Context clues that favor one option over others\n- Classification markers\n- Decision-supporting evidence",
+ "output_format": "default_keyword",
+ "img": "/static/ai_assistant_img/highlight.svg"
+ },
+ "rationale": {
+ "name": "Rationale",
+ "prompt": "TASK: Generate rationales explaining why each label might apply to this text.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nCRITICAL REQUIREMENT:\nYou MUST provide a rationale for EVERY label listed above. Your output must have exactly one entry per label.\n\nINSTRUCTIONS:\nFor EACH available label (ALL of them), provide a brief rationale explaining what evidence in the text could support choosing that label. Be balanced and objective - present the case for each label fairly. If a label doesn't apply well, explain why.\n\nOUTPUT FORMAT:\nReturn a JSON object with \"rationales\" array containing one object per label, each with \"label\" and \"reasoning\" fields.\n\nEXAMPLE OUTPUT:\n{\"rationales\": [{\"label\": \"option1\", \"reasoning\": \"Evidence supporting this choice...\"}, {\"label\": \"option2\", \"reasoning\": \"Evidence supporting this choice...\"}, {\"label\": \"option3\", \"reasoning\": \"This label is less applicable because...\"}]}",
+ "output_format": "default_rationale",
+ "img": "/static/ai_assistant_img/question.svg"
+ }
+}
diff --git a/potato/ai/prompt/slider.json b/potato/ai/prompt/slider.json
new file mode 100644
index 0000000000000000000000000000000000000000..65a532320f6e0550741b2b001f2e1afefada9bdb
--- /dev/null
+++ b/potato/ai/prompt/slider.json
@@ -0,0 +1,18 @@
+{
+ "hint": {
+ "prompt": "TASK: Generate guidance for numerical range selection.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n- Range: ${min_value} to ${max_value}\n- Step size: ${step}\n\nINSTRUCTIONS:\n1. Analyze text for features that indicate position on the numerical scale\n2. Guide toward identifying the appropriate range value\n3. Focus on intensity, degree, or quantity indicators\n\nHINT REQUIREMENTS:\n- Identify intensity markers and degree indicators\n- Point out quantitative or qualitative measures\n- Guide toward scale positioning logic\n- Highlight comparative or relative indicators",
+ "output_format": "default_hint",
+ "img": "/static/ai_assistant_img/blub.svg"
+ },
+ "keyword": {
+ "prompt": "TASK: Extract words/phrases that indicate scale positioning.\n\nINPUT DETAILS:\n- Text: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n- Range: ${min_value} to ${max_value}\n- Step size: ${step}\n\nOBJECTIVE: Identify terms that indicate where on the numerical scale the value should be positioned.\n\nSELECTION CRITERIA:\n- Intensity modifiers (extremely, moderately, slightly)\n- Quantitative indicators\n- Comparative terms (more, less, higher, lower)\n- Degree markers and qualifiers\n- Scale-relevant descriptors",
+ "output_format": "default_keyword",
+ "img": "/static/ai_assistant_img/highlight.svg"
+ },
+ "rationale": {
+ "name": "Rationale",
+ "prompt": "TASK: Generate rationales for different positions on the numerical scale.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Range: ${min_value} to ${max_value}\n\nINSTRUCTIONS:\nProvide rationales for different positions on the scale (low, middle, high range). Explain what evidence in the text could support each position. Be balanced and objective.\n\nOUTPUT FORMAT:\nReturn a JSON object with \"rationales\" array containing objects with \"label\" (scale range descriptor) and \"reasoning\" fields.\n\nEXAMPLE OUTPUT:\n{\"rationales\": [{\"label\": \"low range\", \"reasoning\": \"Weak indicators and hedging language suggest a lower value\"}, {\"label\": \"mid range\", \"reasoning\": \"Moderate language without strong indicators\"}, {\"label\": \"high range\", \"reasoning\": \"Strong intensity markers and superlatives support a higher value\"}]}",
+ "output_format": "default_rationale",
+ "img": "/static/ai_assistant_img/question.svg"
+ }
+}
diff --git a/potato/ai/prompt/span.json b/potato/ai/prompt/span.json
new file mode 100644
index 0000000000000000000000000000000000000000..85e12764a46db79a1692f9c13c47fb5526c48928
--- /dev/null
+++ b/potato/ai/prompt/span.json
@@ -0,0 +1,18 @@
+{
+ "hint": {
+ "prompt": "TASK: Generate guidance for text span highlighting/selection.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nINSTRUCTIONS:\n1. Guide toward identifying specific text portions that need highlighting\n2. Focus on boundary detection and span selection criteria\n3. Help distinguish between different types of spans to highlight\n\nHINT REQUIREMENTS:\n- Identify markers that indicate span boundaries (start/end points)\n- Point out different types of spans corresponding to available labels\n- Guide toward precise text selection (not too broad/narrow)\n- Highlight contextual clues for span classification",
+ "output_format": "default_hint",
+ "img": "/static/ai_assistant_img/blub.svg"
+ },
+ "keyword": {
+ "prompt": "TASK: Extract boundary markers and span indicators.\n\nINPUT DETAILS:\n- Text: '${text}'\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nOBJECTIVE: Identify text spans that mark boundaries or indicate span types. Spans can be single words, multiple words, phrases, sentences, or any continuous text segment.\n\nSELECTION CRITERIA:\n- Boundary markers (punctuation, transitional words, phrases)\n- Type indicators for different span categories\n- Start/end signals for text segments (can be single words or multi-word expressions)\n- Classification markers within spans (words, phrases, or full sentences)\n- Structural elements that define spans (can span multiple words)\n\nSPAN LENGTH EXAMPLES:\n- Single character: ',' or '.' (punctuation)\n- Single word: 'however' or 'therefore'\n- Multiple words: 'on the other hand' or 'in conclusion'\n- Phrase: 'despite the challenges'\n- Sentence: 'This marks a significant transition.'\n- Paragraph or longer: Any continuous text segment\n\nREQUIRED OUTPUT FORMAT \n- 'label': One of the available labels (${labels})\n- 'start': Zero-indexed character position where span begins\n- 'end': Zero-indexed character position where span ends (exclusive)\n- 'text': Exact text extracted from the input (can be any length)\n- 'reasoning': Brief explanation for the classification\n\nCHARACTER INDEXING RULES:\n- First character is at index 0\n- 'end' is exclusive: text = input.substring(start, end)\n- Single word example: In 'Hello world', 'Hello' = {start: 0, end: 5}\n- Multi-word example: In 'Hello world', 'Hello world' = {start: 0, end: 11}\n- Phrase example: In 'on the other hand', entire phrase = {start: 0, end: 17}\n\nEXAMPLE OUTPUT:\n[{\n 'label': 'boundary_marker', 'start': 5, 'end': 6, 'text': ',', 'reasoning': 'Punctuation marking clause boundary'\n}, {\n 'label': 'transition_phrase', 'start': 10, 'end': 28, 'text': 'on the other hand', 'reasoning': 'Multi-word transitional phrase indicating contrast'\n}, {\n 'label': 'conclusion_sentence', 'start': 30, 'end': 65, 'text': 'This marks a significant transition.', 'reasoning': 'Complete sentence signaling major shift'\n}]\n",
+ "output_format": "default_keyword",
+ "img": "/static/ai_assistant_img/highlight.svg"
+ },
+ "rationale": {
+ "name": "Rationale",
+ "prompt": "TASK: Generate rationales explaining why each span label might apply to portions of this text.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n- Available labels: ${labels}\n\nCRITICAL REQUIREMENT:\nYou MUST provide a rationale for EVERY label listed above. Count the labels and ensure your output has exactly that many rationale entries.\n\nINSTRUCTIONS:\nFor EACH available label (ALL of them, no exceptions), provide a brief rationale explaining what types of text spans in this document could be tagged with that label. Even if a label doesn't seem to apply, explain what would need to be present for it to apply.\n\nOUTPUT FORMAT:\nReturn a JSON object with \"rationales\" array containing one object per label, each with \"label\" and \"reasoning\" fields.\n\nEXAMPLE (if labels are: positive, negative, product, quality):\n{\"rationales\": [{\"label\": \"positive\", \"reasoning\": \"Words expressing satisfaction or approval\"}, {\"label\": \"negative\", \"reasoning\": \"Words expressing dissatisfaction or criticism\"}, {\"label\": \"product\", \"reasoning\": \"Names of items being discussed\"}, {\"label\": \"quality\", \"reasoning\": \"Descriptions of attributes or characteristics\"}]}",
+ "output_format": "default_rationale",
+ "img": "/static/ai_assistant_img/question.svg"
+ }
+}
diff --git a/potato/ai/prompt/text.json b/potato/ai/prompt/text.json
new file mode 100644
index 0000000000000000000000000000000000000000..29c125523801f87bd2848c35180917ef8037f14a
--- /dev/null
+++ b/potato/ai/prompt/text.json
@@ -0,0 +1,18 @@
+{
+ "hint": {
+ "prompt": "TASK: Generate guidance for open-ended text extraction/generation.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n\nINSTRUCTIONS:\n1. Analyze what type of textual response is required\n2. Guide toward identifying relevant content for extraction/summarization\n3. Focus on text selection or generation criteria\n\nHINT REQUIREMENTS:\n- Identify the type of text response needed (summary, extraction, classification reason, etc.)\n- Point out relevant sections or themes in the source text\n- Guide toward appropriate level of detail and focus\n- Highlight key content areas to address",
+ "output_format": "default_hint",
+ "img": "/static/ai_assistant_img/blub.svg"
+ },
+ "keyword": {
+ "prompt": "TASK: Extract words/phrases most relevant for the text response.\n\nINPUT DETAILS:\n- Text: \"${text}\"\n- Annotation task: ${description}\n\nOBJECTIVE: Identify key terms that should inform the text response.\n\nSELECTION CRITERIA:\n- Central themes or main concepts\n- Keywords relevant to the annotation task\n- Important phrases that capture key information\n- Terms that represent core ideas or arguments\n- Context-specific vocabulary",
+ "output_format": "default_keyword",
+ "img": "/static/ai_assistant_img/highlight.svg"
+ },
+ "rationale": {
+ "name": "Rationale",
+ "prompt": "TASK: Generate rationales for different approaches to the text response.\n\nINPUT DETAILS:\n- Text to annotate: \"${text}\"\n- Annotation task: ${description}\n\nINSTRUCTIONS:\nProvide different perspectives or approaches for crafting the text response. Consider different angles, levels of detail, or interpretations.\n\nOUTPUT FORMAT:\nReturn a JSON object with \"rationales\" array containing objects with \"label\" (approach name) and \"reasoning\" fields.\n\nEXAMPLE OUTPUT:\n{\"rationales\": [{\"label\": \"summary focus\", \"reasoning\": \"Focus on the main argument about X\"}, {\"label\": \"detail focus\", \"reasoning\": \"Emphasize the specific examples and evidence provided\"}]}",
+ "output_format": "default_rationale",
+ "img": "/static/ai_assistant_img/question.svg"
+ }
+}
diff --git a/potato/ai/prompt/video_annotation.json b/potato/ai/prompt/video_annotation.json
new file mode 100644
index 0000000000000000000000000000000000000000..db71fc2cff8215b9cf0e8f992f3b2220444097a0
--- /dev/null
+++ b/potato/ai/prompt/video_annotation.json
@@ -0,0 +1,38 @@
+{
+ "scene_detection": {
+ "prompt": "TASK: Detect distinct scenes or segments in this video.\n\nDescription: ${description}\nVideo duration: ${duration} seconds\nNumber of frames analyzed: ${num_frames}\nAvailable labels: ${labels}\n\nAnalyze the provided frames (sampled from the video) and identify distinct scenes or segments. For each segment, provide:\n1. Start and end times (in seconds)\n2. A label from the available list\n3. Confidence score\n4. Brief description of why this segment was identified\n\nReturn JSON:\n{\n \"segments\": [\n {\n \"start_time\": 0.0,\n \"end_time\": 5.5,\n \"suggested_label\": \"label_name\",\n \"confidence\": 0.85,\n \"description\": \"Brief reason for this segment\"\n }\n ]\n}\n\nEnsure segments cover the entire video duration and don't overlap.",
+ "output_format": "video_scene_detection",
+ "img": "/static/ai_assistant_img/scene.svg",
+ "name": "Scenes"
+ },
+ "frame_classification": {
+ "prompt": "TASK: Classify this video frame.\n\nDescription: ${description}\nCurrent time: approximately frame ${num_frames}\nAvailable labels: ${labels}\n\nAnalyze this frame and determine which label best describes the current content.\n\nReturn JSON:\n{\n \"suggested_label\": \"label_name\",\n \"confidence\": 0.85,\n \"reasoning\": \"Brief explanation\"\n}",
+ "output_format": "visual_classification",
+ "img": "/static/ai_assistant_img/classify.svg",
+ "name": "Classify Frame"
+ },
+ "keyframe_detection": {
+ "prompt": "TASK: Identify keyframes (significant moments) in this video.\n\nDescription: ${description}\nVideo duration: ${duration} seconds\nNumber of frames analyzed: ${num_frames}\nAvailable labels: ${labels}\n\nAnalyze the frames and identify significant moments that would make good keyframes for annotation. Keyframes should represent:\n- Scene transitions\n- Important actions or events\n- Representative moments of each segment\n\nReturn JSON:\n{\n \"keyframes\": [\n {\n \"timestamp\": 2.5,\n \"suggested_label\": \"label_name\",\n \"confidence\": 0.9,\n \"reason\": \"Why this is a significant moment\"\n }\n ]\n}",
+ "output_format": "video_keyframe_detection",
+ "img": "/static/ai_assistant_img/keyframe.svg",
+ "name": "Keyframes"
+ },
+ "tracking_suggestion": {
+ "prompt": "TASK: Suggest object positions for tracking across frames.\n\nDescription: ${description}\nNumber of frames: ${num_frames}\nLabels to track: ${labels}\n\nAnalyze these sequential frames and identify the position of objects across time. For each object, provide its bounding box in each frame where it's visible.\n\nReturn JSON:\n{\n \"tracks\": [\n {\n \"track_id\": 1,\n \"label\": \"object_label\",\n \"positions\": [\n {\n \"frame_index\": 0,\n \"bbox\": {\"x\": 0.1, \"y\": 0.2, \"width\": 0.15, \"height\": 0.2},\n \"confidence\": 0.9\n }\n ]\n }\n ]\n}\n\nNote: frame_index corresponds to the order of frames provided (0-indexed).",
+ "output_format": "video_tracking_suggestion",
+ "img": "/static/ai_assistant_img/track.svg",
+ "name": "Track"
+ },
+ "hint": {
+ "prompt": "TASK: Provide a helpful hint for annotating this video WITHOUT revealing exact answers.\n\nAnnotation task: ${description}\nVideo duration: ${duration} seconds\nAvailable labels: ${labels}\n\nProvide guidance that helps the annotator without giving away:\n- Exact segment boundaries\n- Specific label assignments\n\nGood hints:\n- Note patterns or transitions to watch for\n- Suggest areas of the video that need attention\n- Point out potential challenges or ambiguities\n\nReturn JSON:\n{\n \"hint\": \"Your helpful guidance here\",\n \"suggestive_choice\": \"optional_focus_area\"\n}",
+ "output_format": "default_hint",
+ "img": "/static/ai_assistant_img/blub.svg",
+ "name": "Hint"
+ },
+ "pre_annotate": {
+ "prompt": "TASK: Pre-annotate this video with scene segments.\n\nDescription: ${description}\nVideo duration: ${duration} seconds\nNumber of frames analyzed: ${num_frames}\nAvailable labels: ${labels}\n\nAutomatically segment the entire video and assign labels. Be thorough - the annotator will review and adjust your suggestions.\n\nReturn JSON:\n{\n \"segments\": [\n {\n \"start_time\": 0.0,\n \"end_time\": 5.5,\n \"suggested_label\": \"label_name\",\n \"confidence\": 0.85,\n \"description\": \"Brief description of this segment\"\n }\n ]\n}\n\nInclude all segments, even uncertain ones. Ensure complete coverage of video duration.",
+ "output_format": "video_scene_detection",
+ "img": "/static/ai_assistant_img/auto.svg",
+ "name": "Auto"
+ }
+}
diff --git a/potato/ai/visual_ai_endpoint.py b/potato/ai/visual_ai_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..8c43e5aee665f23858896edd6b73a43d2421e561
--- /dev/null
+++ b/potato/ai/visual_ai_endpoint.py
@@ -0,0 +1,521 @@
+"""
+Base Visual AI Endpoint
+
+Abstract base class for AI endpoints that work with images and videos.
+Provides common utilities for image encoding, video frame extraction,
+and visual annotation tasks.
+"""
+
+import base64
+import logging
+import os
+import tempfile
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Type, Union
+
+from pydantic import BaseModel
+
+from .ai_endpoint import BaseAIEndpoint, ImageData, VisualAnnotationInput, AIEndpointRequestError
+
+logger = logging.getLogger(__name__)
+
+
+class BaseVisualAIEndpoint(BaseAIEndpoint, ABC):
+ """
+ Abstract base class for visual AI endpoints.
+
+ Extends BaseAIEndpoint with capabilities for processing images and videos.
+ Subclasses should implement query_with_image() for provider-specific image handling.
+ """
+
+ def __init__(self, config: Dict[str, Any]):
+ """
+ Initialize the visual AI endpoint.
+
+ Args:
+ config: Configuration dictionary containing endpoint-specific settings
+ """
+ super().__init__(config)
+
+ # Visual-specific configuration
+ self.max_image_size = self.ai_config.get("max_image_size", 4096) # Max dimension in pixels
+ self.default_video_fps = self.ai_config.get("default_video_fps", 1) # Frames per second for sampling
+ self.max_frames = self.ai_config.get("max_frames", 10) # Max frames for video analysis
+
+ @abstractmethod
+ def query_with_image(
+ self,
+ prompt: str,
+ image_data: Union[ImageData, List[ImageData]],
+ output_format: Type[BaseModel]
+ ) -> Any:
+ """
+ Send a query with image(s) to the AI model.
+
+ Args:
+ prompt: The text prompt describing what to analyze
+ image_data: Single ImageData or list of ImageData for multiple frames
+ output_format: Pydantic model for structured output
+
+ Returns:
+ The model's response parsed according to output_format
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ pass
+
+ def get_visual_ai(
+ self,
+ data: VisualAnnotationInput,
+ output_format: Type[BaseModel]
+ ) -> Any:
+ """
+ Get AI assistance for visual annotation.
+
+ This is the main entry point for visual annotation tasks.
+ It builds the prompt from templates and calls query_with_image().
+
+ Args:
+ data: VisualAnnotationInput containing task details and image data
+ output_format: Pydantic model for structured output
+
+ Returns:
+ AI response (detections, classifications, hints, etc.)
+ """
+ try:
+ from .ai_prompt import get_ai_prompt
+ from string import Template
+
+ ai_prompt = get_ai_prompt()
+
+ # Check if annotation type and ai_assistant exist in prompts
+ if data.annotation_type not in ai_prompt:
+ logger.warning(f"No prompts found for annotation type: {data.annotation_type}")
+ return {"error": f"No prompts configured for {data.annotation_type}"}
+
+ if data.ai_assistant not in ai_prompt[data.annotation_type]:
+ logger.warning(f"No prompt found for ai_assistant: {data.ai_assistant}")
+ return {"error": f"No prompt configured for {data.ai_assistant}"}
+
+ prompt_config = ai_prompt[data.annotation_type][data.ai_assistant]
+ template_str = prompt_config.get("prompt", "")
+
+ # Build template variables
+ template_vars = {
+ "description": data.description,
+ "labels": ", ".join(data.labels) if data.labels else "any objects",
+ "task_type": data.task_type,
+ "confidence_threshold": data.confidence_threshold,
+ }
+
+ # Add video-specific variables
+ if data.video_metadata:
+ template_vars.update({
+ "duration": data.video_metadata.get("duration", 0),
+ "fps": data.video_metadata.get("fps", 30),
+ "num_frames": len(data.image_data) if isinstance(data.image_data, list) else 1,
+ })
+
+ # Add region info for classification
+ if data.region:
+ template_vars["region"] = f"x={data.region.get('x', 0):.2f}, y={data.region.get('y', 0):.2f}, width={data.region.get('width', 1):.2f}, height={data.region.get('height', 1):.2f}"
+
+ # Substitute template variables
+ template = Template(template_str)
+ prompt = template.safe_substitute(template_vars)
+
+ logger.debug(f"Visual AI prompt: {prompt[:200]}...")
+
+ return self.query_with_image(prompt, data.image_data, output_format)
+
+ except Exception as e:
+ logger.error(f"Error in get_visual_ai: {type(e).__name__}: {e}")
+ import traceback
+ logger.error(f"Traceback:\n{traceback.format_exc()}")
+ return {"error": f"Failed to get visual AI assistance: {str(e)}"}
+
+ @staticmethod
+ def encode_image_to_base64(image_path: str) -> ImageData:
+ """
+ Read an image file and encode it as base64.
+
+ Args:
+ image_path: Path to the image file
+
+ Returns:
+ ImageData with base64-encoded image
+
+ Raises:
+ AIEndpointRequestError: If the file cannot be read
+ """
+ try:
+ import mimetypes
+
+ # Determine MIME type
+ mime_type, _ = mimetypes.guess_type(image_path)
+ if not mime_type:
+ # Default to JPEG if unknown
+ mime_type = "image/jpeg"
+
+ with open(image_path, "rb") as f:
+ image_bytes = f.read()
+
+ encoded = base64.b64encode(image_bytes).decode("utf-8")
+
+ # Try to get dimensions using PIL if available
+ width, height = None, None
+ try:
+ from PIL import Image
+ with Image.open(image_path) as img:
+ width, height = img.size
+ except ImportError:
+ logger.debug("PIL not available, skipping dimension extraction")
+ except Exception as e:
+ logger.debug(f"Could not extract dimensions: {e}")
+
+ return ImageData(
+ source="base64",
+ data=encoded,
+ width=width,
+ height=height,
+ mime_type=mime_type
+ )
+
+ except Exception as e:
+ raise AIEndpointRequestError(f"Failed to encode image: {e}")
+
+ @staticmethod
+ def download_image_to_base64(url: str, timeout: int = 30) -> ImageData:
+ """
+ Download an image from URL and encode as base64.
+
+ Args:
+ url: URL of the image
+ timeout: Request timeout in seconds
+
+ Returns:
+ ImageData with base64-encoded image
+
+ Raises:
+ AIEndpointRequestError: If the download fails
+ """
+ try:
+ import requests
+
+ response = requests.get(url, timeout=timeout)
+ response.raise_for_status()
+
+ # Get MIME type from content-type header
+ content_type = response.headers.get("Content-Type", "image/jpeg")
+ if ";" in content_type:
+ content_type = content_type.split(";")[0].strip()
+
+ encoded = base64.b64encode(response.content).decode("utf-8")
+
+ # Try to get dimensions
+ width, height = None, None
+ try:
+ from PIL import Image
+ import io
+ img = Image.open(io.BytesIO(response.content))
+ width, height = img.size
+ img.close()
+ except ImportError:
+ logger.debug("PIL not available, skipping dimension extraction")
+ except Exception as e:
+ logger.debug(f"Could not extract dimensions: {e}")
+
+ return ImageData(
+ source="base64",
+ data=encoded,
+ width=width,
+ height=height,
+ mime_type=content_type
+ )
+
+ except Exception as e:
+ raise AIEndpointRequestError(f"Failed to download image from {url}: {e}")
+
+ @staticmethod
+ def create_url_image_data(url: str) -> ImageData:
+ """
+ Create an ImageData object for a URL without downloading.
+
+ Some APIs accept image URLs directly. Use this when you don't
+ need to download the image first.
+
+ Args:
+ url: URL of the image
+
+ Returns:
+ ImageData with URL reference
+ """
+ return ImageData(
+ source="url",
+ data=url,
+ mime_type=None
+ )
+
+ def extract_video_frames(
+ self,
+ video_path_or_url: str,
+ fps: Optional[float] = None,
+ max_frames: Optional[int] = None,
+ start_time: float = 0,
+ end_time: Optional[float] = None
+ ) -> List[ImageData]:
+ """
+ Extract frames from a video file or URL.
+
+ Args:
+ video_path_or_url: Path to video file or URL
+ fps: Frames per second to sample (default: self.default_video_fps)
+ max_frames: Maximum number of frames to extract (default: self.max_frames)
+ start_time: Start time in seconds
+ end_time: End time in seconds (None for entire video)
+
+ Returns:
+ List of ImageData objects containing base64-encoded frames
+
+ Raises:
+ AIEndpointRequestError: If video processing fails
+ """
+ try:
+ import cv2
+ except ImportError:
+ raise AIEndpointRequestError(
+ "OpenCV (cv2) is required for video frame extraction. "
+ "Install it with: pip install opencv-python"
+ )
+
+ fps = fps or self.default_video_fps
+ max_frames = max_frames or self.max_frames
+
+ temp_file = None
+ video_path = video_path_or_url
+
+ try:
+ # If URL, download to temp file
+ if video_path_or_url.startswith(("http://", "https://")):
+ import requests
+
+ response = requests.get(video_path_or_url, stream=True, timeout=60)
+ response.raise_for_status()
+
+ # Create temp file with appropriate extension
+ suffix = ".mp4"
+ if "." in video_path_or_url.split("/")[-1]:
+ suffix = "." + video_path_or_url.split(".")[-1].split("?")[0]
+
+ temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+ for chunk in response.iter_content(chunk_size=8192):
+ temp_file.write(chunk)
+ temp_file.close()
+ video_path = temp_file.name
+
+ # Open video
+ cap = cv2.VideoCapture(video_path)
+ if not cap.isOpened():
+ raise AIEndpointRequestError(f"Could not open video: {video_path_or_url}")
+
+ # Get video properties
+ video_fps = cap.get(cv2.CAP_PROP_FPS)
+ total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+ duration = total_frames / video_fps if video_fps > 0 else 0
+
+ if end_time is None:
+ end_time = duration
+
+ # Calculate frame interval
+ frame_interval = int(video_fps / fps) if fps < video_fps else 1
+ start_frame = int(start_time * video_fps)
+ end_frame = int(min(end_time, duration) * video_fps)
+
+ frames: List[ImageData] = []
+ current_frame = start_frame
+
+ cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
+
+ while current_frame < end_frame and len(frames) < max_frames:
+ cap.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
+ ret, frame = cap.read()
+
+ if not ret:
+ break
+
+ # Encode frame as JPEG
+ _, buffer = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
+ encoded = base64.b64encode(buffer).decode("utf-8")
+
+ height, width = frame.shape[:2]
+
+ frames.append(ImageData(
+ source="base64",
+ data=encoded,
+ width=width,
+ height=height,
+ mime_type="image/jpeg"
+ ))
+
+ current_frame += frame_interval
+
+ cap.release()
+
+ logger.info(f"Extracted {len(frames)} frames from video")
+ return frames
+
+ except AIEndpointRequestError:
+ raise
+ except Exception as e:
+ raise AIEndpointRequestError(f"Failed to extract video frames: {e}")
+ finally:
+ # Clean up temp file
+ if temp_file and os.path.exists(temp_file.name):
+ try:
+ os.unlink(temp_file.name)
+ except Exception:
+ pass
+
+ def get_video_metadata(self, video_path_or_url: str) -> Dict[str, Any]:
+ """
+ Get metadata from a video file or URL.
+
+ Args:
+ video_path_or_url: Path to video file or URL
+
+ Returns:
+ Dictionary with fps, duration, width, height, total_frames
+
+ Raises:
+ AIEndpointRequestError: If metadata extraction fails
+ """
+ try:
+ import cv2
+ except ImportError:
+ raise AIEndpointRequestError(
+ "OpenCV (cv2) is required for video metadata extraction. "
+ "Install it with: pip install opencv-python"
+ )
+
+ temp_file = None
+ video_path = video_path_or_url
+
+ try:
+ # If URL, download to temp file
+ if video_path_or_url.startswith(("http://", "https://")):
+ import requests
+
+ response = requests.get(video_path_or_url, stream=True, timeout=60)
+ response.raise_for_status()
+
+ suffix = ".mp4"
+ if "." in video_path_or_url.split("/")[-1]:
+ suffix = "." + video_path_or_url.split(".")[-1].split("?")[0]
+
+ temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+ for chunk in response.iter_content(chunk_size=8192):
+ temp_file.write(chunk)
+ temp_file.close()
+ video_path = temp_file.name
+
+ cap = cv2.VideoCapture(video_path)
+ if not cap.isOpened():
+ raise AIEndpointRequestError(f"Could not open video: {video_path_or_url}")
+
+ fps = cap.get(cv2.CAP_PROP_FPS)
+ total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+ duration = total_frames / fps if fps > 0 else 0
+
+ cap.release()
+
+ return {
+ "fps": fps,
+ "duration": duration,
+ "width": width,
+ "height": height,
+ "total_frames": total_frames
+ }
+
+ except AIEndpointRequestError:
+ raise
+ except Exception as e:
+ raise AIEndpointRequestError(f"Failed to get video metadata: {e}")
+ finally:
+ if temp_file and os.path.exists(temp_file.name):
+ try:
+ os.unlink(temp_file.name)
+ except Exception:
+ pass
+
+ def resize_image(
+ self,
+ image_data: ImageData,
+ max_dimension: Optional[int] = None
+ ) -> ImageData:
+ """
+ Resize an image to fit within max dimensions.
+
+ Args:
+ image_data: ImageData to resize
+ max_dimension: Maximum width/height (default: self.max_image_size)
+
+ Returns:
+ Resized ImageData (or original if already within limits)
+ """
+ try:
+ from PIL import Image
+ import io
+ except ImportError:
+ logger.warning("PIL not available, cannot resize image")
+ return image_data
+
+ max_dimension = max_dimension or self.max_image_size
+
+ try:
+ # Decode image
+ if image_data.source == "base64":
+ img_bytes = base64.b64decode(image_data.data)
+ else:
+ # URL - need to download first
+ import requests
+ response = requests.get(image_data.data, timeout=30)
+ img_bytes = response.content
+
+ img = Image.open(io.BytesIO(img_bytes))
+ width, height = img.size
+
+ # Check if resize needed
+ if width <= max_dimension and height <= max_dimension:
+ return image_data
+
+ # Calculate new dimensions
+ if width > height:
+ new_width = max_dimension
+ new_height = int(height * (max_dimension / width))
+ else:
+ new_height = max_dimension
+ new_width = int(width * (max_dimension / height))
+
+ # Resize
+ img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+ # Re-encode
+ buffer = io.BytesIO()
+ img_format = "JPEG" if image_data.mime_type in [None, "image/jpeg"] else "PNG"
+ img.save(buffer, format=img_format, quality=85)
+ encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+ return ImageData(
+ source="base64",
+ data=encoded,
+ width=new_width,
+ height=new_height,
+ mime_type=f"image/{img_format.lower()}"
+ )
+
+ except Exception as e:
+ logger.warning(f"Failed to resize image: {e}")
+ return image_data
diff --git a/potato/ai/vllm_endpoint.py b/potato/ai/vllm_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..df27efc06196a836973eb728709637a9c2402487
--- /dev/null
+++ b/potato/ai/vllm_endpoint.py
@@ -0,0 +1,201 @@
+"""
+VLLM AI endpoint implementation.
+
+This module provides integration with VLLM for local LLM inference.
+"""
+
+import requests
+import json
+from .ai_endpoint import BaseAIEndpoint, AIEndpointRequestError
+
+DEFAULT_MODEL = "meta-llama/Llama-3.2-3B-Instruct"
+DEFAULT_HINT_PROMPT = '''
+ You are assisting a user with an annotation task.
+ The annotation instruction is : {description}
+ The annotation task type is: {annotation_type}
+ The sentence (or item) to annotate is : {text}
+ Your goal is to generate a short, helpful hint that guides the annotator in how to think about the input โ **without providing the answer**.
+
+ The hint should:
+ - Highlight key aspects of the input relevant to the task
+ - Encourage thoughtful reasoning or observation
+ - Point to subtle features (tone, wording, structure, implication) that matter for the annotation
+ - Be specific and informative, not vague or generic
+ '''
+
+DEFAULT_KEYWORD_PROMPT = '''
+ You are assisting a user with an annotation task.
+ The annotation instruction is : {description}
+ The annotation task type is: {annotation_type}
+ The sentence (or item) to annotate is : {text}
+ Your goal is : Print out just a sequence of keywords, not sentences, in the text that most relate to the task. Do not explain your answer. Do not print out the entire text. If no part of the text relates to the task, print the empty string.
+ '''
+
+class VLLMEndpoint(BaseAIEndpoint):
+ """VLLM endpoint for local LLM inference."""
+
+ def _initialize_client(self) -> None:
+ """Initialize the VLLM client."""
+ raw_base = self.ai_config.get("base_url", "http://localhost:8000")
+ # Accept either the server root or an OpenAI-style ".../v1" base_url.
+ # We append "/v1/chat/completions" ourselves, so a trailing "/v1"
+ # (or "/") would otherwise double up ("/v1/v1/...") and the health
+ # probe would hit "/v1/health" (404 on vLLM, whose health is at root).
+ self.base_url = raw_base.rstrip("/")
+ if self.base_url.endswith("/v1"):
+ self.base_url = self.base_url[: -len("/v1")]
+ self.api_key = self.ai_config.get("api_key", "")
+ # Default timeout of 30 seconds, configurable via ai_config
+ self.timeout = self.ai_config.get("timeout", 30)
+
+ # Liveness probe: prefer /health, but some deployments gate it, so
+ # fall back to the OpenAI-compatible /v1/models listing before
+ # declaring the server unreachable.
+ last_err = None
+ for probe in (f"{self.base_url}/health", f"{self.base_url}/v1/models"):
+ try:
+ resp = requests.get(probe, timeout=5)
+ if resp.status_code == 200:
+ last_err = None
+ break
+ last_err = f"{probe} -> HTTP {resp.status_code}"
+ except requests.exceptions.RequestException as e:
+ last_err = f"{probe} -> {e}"
+ if last_err is not None:
+ raise AIEndpointRequestError(
+ f"Failed to reach VLLM server at {self.base_url}: {last_err}"
+ )
+
+ def _get_default_model(self) -> str:
+ """Get the default VLLM model."""
+ return DEFAULT_MODEL
+
+ def _get_default_hint_prompt(self) -> str:
+ """Get the default hint prompt for VLLM."""
+ return DEFAULT_HINT_PROMPT
+
+ def _get_default_keyword_prompt(self) -> str:
+ """Get the default keyword prompt for VLLM."""
+ return DEFAULT_KEYWORD_PROMPT
+
+ def query(self, prompt: str, output_format=None) -> str:
+ """
+ Send a query to VLLM and return the response.
+
+ Args:
+ prompt: The prompt to send to the model
+ output_format: Optional Pydantic model class for structured output.
+ When provided, the schema is sent as guided_json for
+ constrained generation (vLLM native feature).
+
+ Returns:
+ The model's response as a string (or parsed dict for structured output)
+
+ Raises:
+ AIEndpointRequestError: If the request fails
+ """
+ import logging
+ logger = logging.getLogger(__name__)
+
+ try:
+ headers = {"Content-Type": "application/json"}
+ if self.api_key:
+ headers["Authorization"] = f"Bearer {self.api_key}"
+
+ # Think mode: configurable via ai_config['think']
+ # For qwen3/qwen3.5 on vLLM, this controls enable_thinking
+ think = self.ai_config.get('think', False)
+
+ payload = {
+ "model": self.model,
+ "messages": [{"role": "user", "content": prompt}],
+ "max_tokens": self.max_tokens,
+ "temperature": self.temperature,
+ "stream": False,
+ "chat_template_kwargs": {"enable_thinking": think},
+ }
+
+ # Add structured output via guided_json if schema provided
+ if output_format is not None and hasattr(output_format, 'model_json_schema'):
+ payload["guided_json"] = output_format.model_json_schema()
+
+ response = requests.post(
+ f"{self.base_url}/v1/chat/completions",
+ headers=headers,
+ json=payload,
+ timeout=self.timeout
+ )
+
+ if response.status_code != 200:
+ raise AIEndpointRequestError(
+ f"VLLM request failed with status {response.status_code}: "
+ f"{response.text[:500]}"
+ )
+
+ result = response.json()
+ message = result["choices"][0]["message"]
+ content = message.get("content") or ""
+
+ # When thinking is enabled, content may be empty while reasoning
+ # has the thinking. Check if content has the actual answer.
+ if not content and message.get("reasoning"):
+ reasoning = message["reasoning"]
+ logger.debug(
+ f"[vLLM] Content empty, reasoning present "
+ f"({len(reasoning)} chars). Model may need more tokens."
+ )
+ # Try to extract JSON from reasoning as last resort
+ content = reasoning
+
+ if content:
+ return self.parseStringToJson(content)
+
+ raise AIEndpointRequestError(
+ "Empty content from vLLM - model may need more max_tokens "
+ "or thinking mode disabled"
+ )
+
+ except requests.exceptions.RequestException as e:
+ raise AIEndpointRequestError(f"VLLM request failed: {e}")
+ except (KeyError, IndexError) as e:
+ raise AIEndpointRequestError(f"Invalid VLLM response format: {e}")
+
+ def chat_query(self, messages, **kwargs) -> str:
+ """Send a multi-turn chat to vLLM."""
+ import logging
+ logger = logging.getLogger(__name__)
+
+ try:
+ headers = {"Content-Type": "application/json"}
+ if self.api_key:
+ headers["Authorization"] = f"Bearer {self.api_key}"
+
+ think = self.ai_config.get('think', False)
+
+ payload = {
+ "model": self.model,
+ "messages": messages,
+ "max_tokens": self.max_tokens,
+ "temperature": self.temperature,
+ "stream": False,
+ "chat_template_kwargs": {"enable_thinking": think},
+ }
+
+ response = requests.post(
+ f"{self.base_url}/v1/chat/completions",
+ headers=headers,
+ json=payload,
+ timeout=self.timeout
+ )
+
+ if response.status_code != 200:
+ raise AIEndpointRequestError(
+ f"VLLM chat failed: {response.status_code}"
+ )
+
+ result = response.json()
+ content = result["choices"][0]["message"].get("content") or ""
+ return content
+
+ except requests.exceptions.RequestException as e:
+ raise AIEndpointRequestError(f"VLLM chat failed: {e}")
\ No newline at end of file
diff --git a/potato/ai/yolo_endpoint.py b/potato/ai/yolo_endpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..6efa48020cb2613ba5c7cb4365f82ba6fce01738
--- /dev/null
+++ b/potato/ai/yolo_endpoint.py
@@ -0,0 +1,370 @@
+"""
+YOLO AI Endpoint for Object Detection
+
+This module provides integration with YOLO models (via ultralytics) for
+local object detection inference. Supports YOLOv8 and YOLO-World models.
+"""
+
+import base64
+import logging
+import os
+import tempfile
+from typing import Any, Dict, List, Optional, Type, Union
+
+from pydantic import BaseModel
+
+from .ai_endpoint import AIEndpointRequestError, ImageData, VisualAnnotationInput, ModelCapabilities
+from .visual_ai_endpoint import BaseVisualAIEndpoint
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = "yolov8m.pt"
+
+
+class YOLOEndpoint(BaseVisualAIEndpoint):
+ """
+ YOLO endpoint for object detection using ultralytics.
+
+ Supports:
+ - YOLOv8 models (yolov8n, yolov8s, yolov8m, yolov8l, yolov8x)
+ - YOLO-World models for open-vocabulary detection
+ - Custom trained models
+
+ Configuration options:
+ - model: Model name or path (default: yolov8m.pt)
+ - confidence_threshold: Minimum detection confidence (default: 0.5)
+ - iou_threshold: IOU threshold for NMS (default: 0.45)
+ - device: Device to run on (default: auto - uses GPU if available)
+ - classes: List of class indices to detect (optional)
+ """
+
+ # Capabilities declaration for YOLO detection models
+ # YOLO excels at object detection with precise bounding boxes but cannot generate text
+ CAPABILITIES = ModelCapabilities(
+ text_generation=False, # YOLO doesn't generate text
+ vision_input=True,
+ bounding_box_output=True, # YOLO's primary strength
+ text_classification=False,
+ image_classification=True, # Can classify detected objects
+ rationale_generation=False, # Cannot explain reasoning
+ keyword_extraction=False, # Not applicable
+ )
+
+ def _initialize_client(self) -> None:
+ """Initialize the YOLO model."""
+ try:
+ from ultralytics import YOLO
+ except ImportError:
+ raise AIEndpointRequestError(
+ "ultralytics is required for YOLO detection. "
+ "Install it with: pip install ultralytics"
+ )
+
+ model_name = self.model
+ self.confidence_threshold = self.ai_config.get("confidence_threshold", 0.5)
+ self.iou_threshold = self.ai_config.get("iou_threshold", 0.45)
+ self.device = self.ai_config.get("device", None) # None = auto-detect
+ self.classes = self.ai_config.get("classes", None) # None = all classes
+
+ # YOLO-World specific: custom vocabulary
+ self.custom_classes = self.ai_config.get("custom_classes", None)
+
+ try:
+ logger.info(f"Loading YOLO model: {model_name}")
+ self.yolo_model = YOLO(model_name)
+
+ # Set custom classes for YOLO-World models
+ if self.custom_classes and hasattr(self.yolo_model, "set_classes"):
+ logger.info(f"Setting custom classes: {self.custom_classes}")
+ self.yolo_model.set_classes(self.custom_classes)
+
+ logger.info(f"YOLO model loaded successfully")
+
+ except Exception as e:
+ raise AIEndpointRequestError(f"Failed to load YOLO model: {e}")
+
+ def _get_default_model(self) -> str:
+ """Get the default YOLO model."""
+ return DEFAULT_MODEL
+
+ def query(self, prompt: str, output_format: Type[BaseModel]) -> Any:
+ """
+ Standard query method - not typically used for YOLO.
+
+ YOLO doesn't process text prompts in the traditional sense.
+ Use query_with_image() instead.
+ """
+ logger.warning("YOLO endpoint doesn't support text-only queries. Use query_with_image().")
+ return {"error": "YOLO requires image input. Use query_with_image() instead."}
+
+ def query_with_image(
+ self,
+ prompt: str,
+ image_data: Union[ImageData, List[ImageData]],
+ output_format: Type[BaseModel]
+ ) -> Any:
+ """
+ Run YOLO detection on image(s).
+
+ Args:
+ prompt: Text description (used for filtering labels if provided)
+ image_data: Single ImageData or list of ImageData
+ output_format: Pydantic model for output (typically VisualDetectionFormat)
+
+ Returns:
+ Detection results with normalized bounding boxes
+
+ Raises:
+ AIEndpointRequestError: If detection fails
+ """
+ try:
+ # Handle single image or list
+ images = [image_data] if isinstance(image_data, ImageData) else image_data
+
+ all_detections = []
+
+ for idx, img_data in enumerate(images):
+ detections = self._detect_single_image(img_data, prompt)
+ all_detections.append({
+ "frame_index": idx,
+ "detections": detections
+ })
+
+ # If single image, return flat detections
+ if len(images) == 1:
+ return {"detections": all_detections[0]["detections"]}
+
+ # For multiple images (video frames), return per-frame results
+ return {"frames": all_detections}
+
+ except AIEndpointRequestError:
+ raise
+ except Exception as e:
+ logger.error(f"YOLO detection failed: {e}")
+ raise AIEndpointRequestError(f"YOLO detection failed: {e}")
+
+ def _detect_single_image(self, image_data: ImageData, prompt: str = "") -> List[Dict[str, Any]]:
+ """
+ Run detection on a single image.
+
+ Args:
+ image_data: ImageData to process
+ prompt: Optional text for label filtering
+
+ Returns:
+ List of detection dictionaries
+ """
+ import numpy as np
+
+ # Convert image data to format YOLO can process
+ img = self._prepare_image(image_data)
+
+ # Parse prompt for label filtering
+ filter_labels = self._parse_prompt_for_labels(prompt)
+ logger.debug(f"Filter labels from prompt: {filter_labels}")
+
+ # Run inference
+ logger.debug(f"Running YOLO inference with conf={self.confidence_threshold}, iou={self.iou_threshold}")
+ results = self.yolo_model(
+ img,
+ conf=self.confidence_threshold,
+ iou=self.iou_threshold,
+ device=self.device,
+ classes=self.classes,
+ verbose=False
+ )
+
+ detections = []
+
+ for result in results:
+ if result.boxes is None:
+ logger.debug("No boxes in result")
+ continue
+
+ boxes = result.boxes
+ img_height, img_width = result.orig_shape
+ logger.debug(f"Image size: {img_width}x{img_height}, found {len(boxes)} boxes")
+
+ for i in range(len(boxes)):
+ # Get box coordinates (xyxy format)
+ box = boxes.xyxy[i].cpu().numpy()
+ confidence = float(boxes.conf[i].cpu().numpy())
+ class_id = int(boxes.cls[i].cpu().numpy())
+
+ # Get class name
+ class_name = result.names.get(class_id, f"class_{class_id}")
+ logger.debug(f"Detection {i}: class={class_name}, conf={confidence:.3f}")
+
+ # Filter by label if specified
+ if filter_labels and class_name.lower() not in [l.lower() for l in filter_labels]:
+ logger.debug(f" -> Filtered out (not in {filter_labels})")
+ continue
+
+ # Normalize coordinates to 0-1 range
+ x1, y1, x2, y2 = box
+ normalized_box = {
+ "x": float(x1 / img_width),
+ "y": float(y1 / img_height),
+ "width": float((x2 - x1) / img_width),
+ "height": float((y2 - y1) / img_height)
+ }
+
+ detections.append({
+ "label": class_name,
+ "bbox": normalized_box,
+ "confidence": round(confidence, 4)
+ })
+
+ return detections
+
+ def _prepare_image(self, image_data: ImageData) -> Any:
+ """
+ Prepare image data for YOLO processing.
+
+ Args:
+ image_data: ImageData object
+
+ Returns:
+ Image in format suitable for YOLO (numpy array or PIL Image)
+ """
+ try:
+ from PIL import Image
+ import io
+ except ImportError:
+ raise AIEndpointRequestError("PIL is required for image processing")
+
+ if image_data.source == "base64":
+ # Decode base64 to PIL Image
+ img_bytes = base64.b64decode(image_data.data)
+ img = Image.open(io.BytesIO(img_bytes))
+ return img
+
+ elif image_data.source == "url":
+ # Download and convert to PIL Image
+ import requests
+ response = requests.get(image_data.data, timeout=30)
+ response.raise_for_status()
+ img = Image.open(io.BytesIO(response.content))
+ return img
+
+ else:
+ raise AIEndpointRequestError(f"Unknown image source type: {image_data.source}")
+
+ def _parse_prompt_for_labels(self, prompt: str) -> List[str]:
+ """
+ Extract label names from prompt for filtering.
+
+ Looks for patterns like:
+ - "Labels to detect: person, car, dog"
+ - "Available labels: person, car, dog"
+ - "detect: person, car, dog"
+ - "labels: [person, car]"
+
+ Args:
+ prompt: Text prompt
+
+ Returns:
+ List of label names to filter by (empty for no filtering)
+ """
+ if not prompt:
+ return []
+
+ # Look for explicit label specifications
+ import re
+
+ # Pattern: "Labels to detect: label1, label2" or "detect: label1, label2"
+ match = re.search(r"(?:labels to detect|available labels|detect|labels)[:\s]+([^\n]+)", prompt, re.IGNORECASE)
+ if match:
+ labels_str = match.group(1)
+ # Clean up and split - handle comma-separated, possibly with "..."
+ labels_str = labels_str.replace("[", "").replace("]", "").replace("...", "")
+ # Split by comma and clean up
+ labels = [l.strip() for l in labels_str.split(",")]
+ # Filter out empty strings and common non-label words
+ stop_words = {'any', 'objects', 'in', 'this', 'image', 'that', 'match', 'the', 'specified', ''}
+ labels = [l for l in labels if l.lower() not in stop_words and l]
+ logger.debug(f"Parsed labels from prompt: {labels}")
+ return labels
+
+ # Pattern: "find person and car" - extract nouns around "and"
+ match = re.search(r"(?:find|detect|identify)\s+(.+)", prompt, re.IGNORECASE)
+ if match:
+ labels_str = match.group(1)
+ # Split by "and" and commas
+ labels = re.split(r'\s+and\s+|,\s*', labels_str)
+ labels = [l.strip() for l in labels if l.strip()]
+ if labels:
+ logger.debug(f"Parsed labels from natural language prompt: {labels}")
+ return labels
+
+ return []
+
+ def set_custom_classes(self, classes: List[str]) -> None:
+ """
+ Set custom classes for YOLO-World models.
+
+ Args:
+ classes: List of class names to detect
+ """
+ if hasattr(self.yolo_model, "set_classes"):
+ self.yolo_model.set_classes(classes)
+ self.custom_classes = classes
+ logger.info(f"Updated YOLO-World classes: {classes}")
+ else:
+ logger.warning("Model does not support custom classes (not YOLO-World)")
+
+ def detect(
+ self,
+ image_path_or_data: Union[str, ImageData],
+ confidence_threshold: Optional[float] = None,
+ labels: Optional[List[str]] = None
+ ) -> List[Dict[str, Any]]:
+ """
+ Convenience method for direct detection.
+
+ Args:
+ image_path_or_data: Path to image file or ImageData
+ confidence_threshold: Override default confidence threshold
+ labels: Labels to filter by
+
+ Returns:
+ List of detections
+ """
+ # Prepare image data
+ if isinstance(image_path_or_data, str):
+ if image_path_or_data.startswith(("http://", "https://")):
+ image_data = self.create_url_image_data(image_path_or_data)
+ else:
+ image_data = self.encode_image_to_base64(image_path_or_data)
+ else:
+ image_data = image_path_or_data
+
+ # Temporarily override confidence if specified
+ original_conf = self.confidence_threshold
+ if confidence_threshold is not None:
+ self.confidence_threshold = confidence_threshold
+
+ try:
+ prompt = f"detect: {', '.join(labels)}" if labels else ""
+ return self._detect_single_image(image_data, prompt)
+ finally:
+ self.confidence_threshold = original_conf
+
+ def health_check(self) -> bool:
+ """
+ Check if the YOLO model is loaded and working.
+
+ Returns:
+ True if model is ready, False otherwise
+ """
+ try:
+ # Create a small test image
+ import numpy as np
+ test_img = np.zeros((100, 100, 3), dtype=np.uint8)
+
+ # Run inference
+ self.yolo_model(test_img, verbose=False)
+ return True
+ except Exception as e:
+ logger.error(f"YOLO health check failed: {e}")
+ return False
diff --git a/potato/annotation_history.py b/potato/annotation_history.py
new file mode 100644
index 0000000000000000000000000000000000000000..0053f8200e61bf38009e8be1805f629ff0628a4a
--- /dev/null
+++ b/potato/annotation_history.py
@@ -0,0 +1,286 @@
+"""
+Annotation History Module
+
+This module provides comprehensive tracking of all annotation actions with fine-grained
+timestamp metadata. It enables performance analysis, quality assurance, and future
+undo functionality.
+
+Key Components:
+- AnnotationAction: Dataclass representing a single annotation action
+- AnnotationHistoryManager: Utility class for creating and analyzing annotation actions
+- Performance metrics calculation and suspicious activity detection
+"""
+
+import uuid
+import datetime
+import logging
+from dataclasses import dataclass, asdict
+from typing import Optional, Dict, Any, List
+import json
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AnnotationAction:
+ """
+ Represents a single annotation action with full metadata.
+
+ This class captures all information about an annotation change, including
+ timing data, user context, and action details for comprehensive tracking.
+ """
+ action_id: str # UUID for unique identification
+ timestamp: datetime.datetime # Precise timestamp
+ user_id: str
+ instance_id: str
+ action_type: str # 'add_label', 'update_label', 'delete_label', 'add_span', 'update_span', 'delete_span'
+ schema_name: str
+ label_name: str
+ old_value: Optional[Any] # Previous value (for updates/deletes)
+ new_value: Optional[Any] # New value (for adds/updates)
+ span_data: Optional[Dict] # For span annotations (start, end, text)
+ session_id: str # Browser session identifier
+ client_timestamp: Optional[datetime.datetime] # Frontend timestamp
+ server_processing_time_ms: int # Server processing time
+ metadata: Dict[str, Any] # Additional metadata (browser info, etc.)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to dictionary for serialization"""
+ data = asdict(self)
+ data['timestamp'] = self.timestamp.isoformat()
+ if self.client_timestamp:
+ data['client_timestamp'] = self.client_timestamp.isoformat()
+ return data
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'AnnotationAction':
+ """Create from dictionary"""
+ data['timestamp'] = datetime.datetime.fromisoformat(data['timestamp'])
+ if data.get('client_timestamp'):
+ data['client_timestamp'] = datetime.datetime.fromisoformat(data['client_timestamp'])
+ return cls(**data)
+
+ def __str__(self) -> str:
+ """String representation for logging"""
+ return f"AnnotationAction({self.action_type}: {self.schema_name}:{self.label_name} = {self.new_value})"
+
+
+class AnnotationHistoryManager:
+ """
+ Manages annotation history and provides analytics.
+
+ This class provides utilities for creating annotation actions, calculating
+ performance metrics, and detecting suspicious activity patterns.
+ """
+
+ @staticmethod
+ def create_action(
+ user_id: str,
+ instance_id: str,
+ action_type: str,
+ schema_name: str,
+ label_name: str,
+ old_value: Optional[Any],
+ new_value: Optional[Any],
+ span_data: Optional[Dict] = None,
+ session_id: str = None,
+ client_timestamp: Optional[datetime.datetime] = None,
+ server_processing_time_ms: int = 0,
+ metadata: Optional[Dict] = None
+ ) -> AnnotationAction:
+ """
+ Create a new annotation action with current timestamp.
+
+ Args:
+ user_id: The user performing the action
+ instance_id: The instance being annotated
+ action_type: Type of action (add_label, update_label, etc.)
+ schema_name: Name of the annotation schema
+ label_name: Name of the specific label
+ old_value: Previous value (for updates/deletes)
+ new_value: New value (for adds/updates)
+ span_data: Span annotation data (start, end, text)
+ session_id: Browser session identifier
+ client_timestamp: Frontend timestamp
+ server_processing_time_ms: Server processing time in milliseconds
+ metadata: Additional metadata
+
+ Returns:
+ AnnotationAction object with current timestamp
+ """
+ return AnnotationAction(
+ action_id=str(uuid.uuid4()),
+ timestamp=datetime.datetime.now(),
+ user_id=user_id,
+ instance_id=instance_id,
+ action_type=action_type,
+ schema_name=schema_name,
+ label_name=label_name,
+ old_value=old_value,
+ new_value=new_value,
+ span_data=span_data,
+ session_id=session_id or "unknown",
+ client_timestamp=client_timestamp,
+ server_processing_time_ms=server_processing_time_ms,
+ metadata=metadata or {}
+ )
+
+ @staticmethod
+ def calculate_performance_metrics(actions: List[AnnotationAction]) -> Dict[str, Any]:
+ """
+ Calculate performance metrics from action history.
+
+ Args:
+ actions: List of annotation actions to analyze
+
+ Returns:
+ Dictionary containing performance metrics
+ """
+ if not actions:
+ return {
+ 'total_actions': 0,
+ 'average_action_time_ms': 0,
+ 'fastest_action_time_ms': 0,
+ 'slowest_action_time_ms': 0,
+ 'actions_per_minute': 0,
+ 'total_processing_time_ms': 0
+ }
+
+ processing_times = [a.server_processing_time_ms for a in actions]
+ total_time = sum(processing_times)
+
+ # Calculate actions per minute
+ if len(actions) > 1:
+ time_span = (actions[-1].timestamp - actions[0].timestamp).total_seconds() / 60
+ actions_per_minute = len(actions) / time_span if time_span > 0 else 0
+ else:
+ actions_per_minute = 0
+
+ return {
+ 'total_actions': len(actions),
+ 'average_action_time_ms': total_time / len(actions),
+ 'fastest_action_time_ms': min(processing_times),
+ 'slowest_action_time_ms': max(processing_times),
+ 'actions_per_minute': actions_per_minute,
+ 'total_processing_time_ms': total_time
+ }
+
+ @staticmethod
+ def detect_suspicious_activity(actions: List[AnnotationAction],
+ fast_threshold_ms: int = 500,
+ burst_threshold_seconds: int = 2) -> Dict[str, Any]:
+ """
+ Detect potentially suspicious annotation activity.
+
+ Args:
+ actions: List of annotation actions to analyze
+ fast_threshold_ms: Threshold for considering an action "too fast"
+ burst_threshold_seconds: Threshold for burst activity detection
+
+ Returns:
+ Dictionary containing suspicious activity analysis
+ """
+ if not actions:
+ return {
+ 'suspicious_actions': [],
+ 'fast_actions_count': 0,
+ 'burst_actions_count': 0,
+ 'suspicious_score': 0
+ }
+
+ suspicious_actions = []
+ fast_actions = []
+ burst_actions = []
+
+ # Detect fast actions
+ for action in actions:
+ if action.server_processing_time_ms < fast_threshold_ms:
+ fast_actions.append(action)
+ suspicious_actions.append(action)
+
+ # Detect burst activity (multiple actions in quick succession)
+ for i in range(1, len(actions)):
+ time_diff = (actions[i].timestamp - actions[i-1].timestamp).total_seconds()
+ if time_diff < burst_threshold_seconds:
+ burst_actions.append(actions[i])
+ if actions[i] not in suspicious_actions:
+ suspicious_actions.append(actions[i])
+
+ # Calculate suspicious score (0-100)
+ total_actions = len(actions)
+ fast_percentage = (len(fast_actions) / total_actions) * 100 if total_actions > 0 else 0
+ burst_percentage = (len(burst_actions) / total_actions) * 100 if total_actions > 0 else 0
+
+ suspicious_score = min(100, (fast_percentage * 0.6) + (burst_percentage * 0.4))
+
+ return {
+ 'suspicious_actions': suspicious_actions,
+ 'fast_actions_count': len(fast_actions),
+ 'burst_actions_count': len(burst_actions),
+ 'fast_actions_percentage': fast_percentage,
+ 'burst_actions_percentage': burst_percentage,
+ 'suspicious_score': suspicious_score,
+ 'suspicious_level': _get_suspicious_level(suspicious_score)
+ }
+
+ @staticmethod
+ def get_actions_by_time_range(actions: List[AnnotationAction],
+ start_time: datetime.datetime,
+ end_time: datetime.datetime) -> List[AnnotationAction]:
+ """
+ Filter actions by time range.
+
+ Args:
+ actions: List of annotation actions
+ start_time: Start of time range
+ end_time: End of time range
+
+ Returns:
+ Filtered list of actions within the time range
+ """
+ return [action for action in actions
+ if start_time <= action.timestamp <= end_time]
+
+ @staticmethod
+ def get_actions_by_instance(actions: List[AnnotationAction],
+ instance_id: str) -> List[AnnotationAction]:
+ """
+ Filter actions by instance ID.
+
+ Args:
+ actions: List of annotation actions
+ instance_id: Instance ID to filter by
+
+ Returns:
+ Filtered list of actions for the specified instance
+ """
+ return [action for action in actions if action.instance_id == instance_id]
+
+ @staticmethod
+ def get_actions_by_type(actions: List[AnnotationAction],
+ action_type: str) -> List[AnnotationAction]:
+ """
+ Filter actions by action type.
+
+ Args:
+ actions: List of annotation actions
+ action_type: Action type to filter by
+
+ Returns:
+ Filtered list of actions of the specified type
+ """
+ return [action for action in actions if action.action_type == action_type]
+
+
+def _get_suspicious_level(score: float) -> str:
+ """Convert suspicious score to level description."""
+ if score < 10:
+ return "Normal"
+ elif score < 30:
+ return "Low"
+ elif score < 60:
+ return "Medium"
+ elif score < 80:
+ return "High"
+ else:
+ return "Very High"
\ No newline at end of file
diff --git a/potato/archive/activelearning_old.py b/potato/archive/activelearning_old.py
new file mode 100644
index 0000000000000000000000000000000000000000..75dfb1fbb2405036fb9b7cbcc139dd3d527ec7cc
--- /dev/null
+++ b/potato/archive/activelearning_old.py
@@ -0,0 +1,225 @@
+"""
+Active Learning Module
+
+This module provides active learning capabilities for the annotation platform.
+It implements machine learning algorithms to intelligently select which instances
+should be annotated next, based on model confidence and disagreement scores.
+
+The active learning system:
+1. Trains classifiers on existing annotations
+2. Predicts confidence scores for unlabeled instances
+3. Reorders instances to prioritize those with low confidence
+4. Maintains a balance between active learning and random sampling
+
+This helps reduce the total number of annotations needed while maintaining
+high-quality results by focusing on the most informative instances.
+"""
+
+def actively_learn():
+ """
+ Main active learning function that reorders instances based on model predictions.
+
+ This function implements the core active learning algorithm:
+ 1. Collects all current annotations from users
+ 2. Resolves multiple annotations per instance using a specified strategy
+ 3. Trains classifiers for each annotation scheme
+ 4. Predicts confidence scores for unlabeled instances
+ 5. Reorders instances to prioritize low-confidence predictions
+ 6. Updates user assignment queues while preserving existing annotations
+
+ Side Effects:
+ - Trains machine learning models on current annotations
+ - Reorders instance assignments for all users
+ - Updates active learning state tracking
+ - Logs training progress and statistics
+
+ The function maintains a balance between active learning selection and
+ random sampling to ensure diversity in the training data.
+ """
+ global user_to_annotation_state
+ global instance_id_to_data
+
+ # Check if active learning is configured
+ if "active_learning_config" not in config:
+ logger.warning(
+ "the server is trying to do active learning " + "but this hasn't been configured"
+ )
+ return
+
+ al_config = config["active_learning_config"]
+
+ # Skip if the user doesn't want us to do active learning
+ if "enable_active_learning" in al_config and not al_config["enable_active_learning"]:
+ return
+
+ # Validate required configuration parameters
+ if "classifier_name" not in al_config:
+ raise Exception('active learning enabled but no classifier is set with "classifier_name"')
+
+ if "vectorizer_name" not in al_config:
+ raise Exception('active learning enabled but no vectorizer is set with "vectorizer_name"')
+
+ if "resolution_strategy" not in al_config:
+ raise Exception("active learning enabled but resolution_strategy is not set")
+
+ # This specifies which schema we need to use in active learning (separate
+ # classifiers for each). If the user doesn't specify these, we use all of
+ # them.
+ schema_used = []
+ if "active_learning_schema" in al_config:
+ schema_used = al_config["active_learning_schema"]
+
+ # Get configuration parameters for classifiers and vectorizers
+ cls_kwargs = al_config.get("classifier_kwargs", {})
+ cls_kwargs = al_config.get("classifier_kwargs", {})
+ vectorizer_kwargs = al_config.get("vectorizer_kwargs", {})
+ strategy = al_config["resolution_strategy"]
+
+ # Collect all the current labels from all users
+ # This creates a mapping from instance ID to list of annotations
+ instance_to_labels = defaultdict(list)
+ for uas in user_to_annotation_state.values():
+ for iid, annotation in uas.instance_id_to_labeling.items():
+ instance_to_labels[iid].append(annotation)
+
+ # Resolve all the multiple-annotations to a single one using the provided
+ # strategy to get training data
+ # This handles cases where multiple users have annotated the same instance
+ instance_to_label = {}
+ schema_seen = set()
+ for iid, annotations in instance_to_labels.items():
+ resolved = resolve(annotations, strategy)
+
+ # Prune to just the schema we care about for active learning
+ if len(schema_used) > 0:
+ resolved = {k: resolved[k] for k in schema_used}
+
+ for s in resolved:
+ schema_seen.add(s)
+ instance_to_label[iid] = resolved
+
+ # Construct a dataframe for easy processing
+ texts = []
+ # We'll train one classifier for each scheme
+ scheme_to_labels = defaultdict(list)
+ text_key = config["item_properties"]["text_key"]
+ for iid, schema_to_label in instance_to_label.items():
+ # get the text content for this instance
+ text = instance_id_to_data[iid][text_key]
+ texts.append(text)
+ for s in schema_seen:
+ # In some cases where the user has not selected anything but somehow
+ # this is considered annotated, we include some dummy label
+ label = schema_to_label.get(s, "DUMMY:NONE")
+
+ # HACK: this needs to get fixed for multilabel data and possibly
+ # number data
+ label = list(label.keys())[0]
+ scheme_to_labels[s].append(label)
+
+ scheme_to_classifier = {}
+
+ # Train a classifier for each annotation scheme
+ for scheme, labels in scheme_to_labels.items():
+
+ # Sanity check we have more than 1 label
+ # Active learning requires at least 2 different labels to work
+ label_counts = Counter(labels)
+ if len(label_counts) < 2:
+ logger.warning(
+ (
+ "In the current data, data labeled with %s has only a"
+ + "single unique label, which is insufficient for "
+ + "active learning; skipping..."
+ )
+ % scheme
+ )
+ continue
+
+ # Instantiate the classifier and the tokenizer
+ cls = get_class(al_config["classifier_name"])(**cls_kwargs)
+ vectorizer = get_class(al_config["vectorizer_name"])(**vectorizer_kwargs)
+
+ # Train the classifier using a pipeline
+ clf = Pipeline([("vectorizer", vectorizer), ("classifier", cls)])
+ logger.info("training classifier for %s..." % scheme)
+ clf.fit(texts, labels)
+ logger.info("done training classifier for %s" % scheme)
+ scheme_to_classifier[scheme] = clf
+
+ # Get the remaining unlabeled instances and start predicting
+ unlabeled_ids = [iid for iid in instance_id_to_data if iid not in instance_to_label]
+ random.shuffle(unlabeled_ids)
+
+ # Calculate the percentage of instances to keep random
+ # This ensures we don't bias too heavily toward active learning
+ perc_random = al_config["random_sample_percent"] / 100
+
+ # Split to keep some of the data random
+ # This maintains diversity in the training data
+ random_ids = unlabeled_ids[int(len(unlabeled_ids) * perc_random) :]
+ unlabeled_ids = unlabeled_ids[: int(len(unlabeled_ids) * perc_random)]
+ remaining_ids = []
+
+ # Cap how much inference we need to do (important for big datasets)
+ # This prevents the system from becoming too slow with large datasets
+ if "max_inferred_predictions" in al_config:
+ max_insts = al_config["max_inferred_predictions"]
+ remaining_ids = unlabeled_ids[max_insts:]
+ unlabeled_ids = unlabeled_ids[:max_insts]
+
+ # For each scheme, use its classifier to label the data
+ # This generates confidence scores for each unlabeled instance
+ scheme_to_predictions = {}
+ unlabeled_texts = [instance_id_to_data[iid][text_key] for iid in unlabeled_ids]
+ for scheme, clf in scheme_to_classifier.items():
+ logger.info("Inferring labels for %s" % scheme)
+ preds = clf.predict_proba(unlabeled_texts)
+ scheme_to_predictions[scheme] = preds
+
+ # Figure out which of the instances to prioritize, keeping the specified
+ # ratio of random-vs-AL-selected instances.
+ # We select instances with the lowest confidence scores (highest uncertainty)
+ ids_and_confidence = []
+ logger.info("Scoring items by model confidence")
+ for i, iid in enumerate(tqdm(unlabeled_ids)):
+ most_confident_pred = 0
+ mp_scheme = None
+ for scheme, all_preds in scheme_to_predictions.items():
+
+ preds = all_preds[i, :]
+ mp = max(preds)
+ if mp > most_confident_pred:
+ most_confident_pred = mp
+ mp_scheme = scheme
+ ids_and_confidence.append((iid, most_confident_pred, mp_scheme))
+
+ # Sort by confidence (lowest first for active learning)
+ # This prioritizes instances where the model is least confident
+ ids_and_confidence = sorted(ids_and_confidence, key=lambda x: x[1])
+
+ # Re-order all of the unlabeled instances
+ # Interleave active learning selections with random selections
+ new_id_order = []
+ id_to_selection_type = {}
+ for (al, rand_id) in zip_longest(ids_and_confidence, random_ids, fillvalue=None):
+ if al:
+ new_id_order.append(al[0])
+ id_to_selection_type[al[0]] = "%s Classifier" % al[2]
+ if rand_id:
+ new_id_order.append(rand_id)
+ id_to_selection_type[rand_id] = "Random"
+
+ # These are the IDs that weren't in the random sample or that we didn't
+ # reorder with active learning
+ new_id_order.extend(remaining_ids)
+
+ # Update each user's ordering, preserving the order for any item that has
+ # any annotation so that it stays in the front of the users' queues even if
+ # they haven't gotten to it yet (but others have)
+ # This ensures that partially annotated instances remain accessible
+ already_annotated = list(instance_to_labels.keys())
+ for annotation_state in user_to_annotation_state.values():
+ annotation_state.reorder_remaining_instances(new_id_order, already_annotated)
+
+ logger.info("Finished reordering instances")
\ No newline at end of file
diff --git a/potato/auth_backends/__init__.py b/potato/auth_backends/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d26bb3705311bfafadcbafc19d83f34363b0128c
--- /dev/null
+++ b/potato/auth_backends/__init__.py
@@ -0,0 +1,5 @@
+"""Authentication backends for Potato annotation platform."""
+
+from potato.auth_backends.oauth_backend import OAuthBackend
+
+__all__ = ["OAuthBackend"]
diff --git a/potato/auth_backends/oauth_backend.py b/potato/auth_backends/oauth_backend.py
new file mode 100644
index 0000000000000000000000000000000000000000..aa23b9028dbc48fabe9bbb60010d5d6936269138
--- /dev/null
+++ b/potato/auth_backends/oauth_backend.py
@@ -0,0 +1,297 @@
+"""
+OAuth 2.0 / OpenID Connect Authentication Backend
+
+Supports Google, GitHub, and generic OIDC providers for single sign-on.
+Uses Authlib for OAuth flow management, token exchange, and OIDC discovery.
+"""
+
+import logging
+import os
+from typing import Any, Dict, List, Optional
+from urllib.parse import urljoin
+
+from authlib.integrations.flask_client import OAuth
+
+from potato.authentication import AuthBackend
+
+logger = logging.getLogger(__name__)
+
+# Well-known OAuth provider configurations
+PROVIDER_CONFIGS = {
+ "google": {
+ "display_name": "Google",
+ "server_metadata_url": "https://accounts.google.com/.well-known/openid-configuration",
+ "client_kwargs": {"scope": "openid email profile"},
+ "icon_class": "fab fa-google",
+ "button_class": "oauth-btn-google",
+ },
+ "github": {
+ "display_name": "GitHub",
+ "api_base_url": "https://api.github.com/",
+ "access_token_url": "https://github.com/login/oauth/access_token",
+ "authorize_url": "https://github.com/login/oauth/authorize",
+ "client_kwargs": {"scope": "user:email"},
+ "icon_class": "fab fa-github",
+ "button_class": "oauth-btn-github",
+ },
+ "huggingface": {
+ "display_name": "HuggingFace",
+ "server_metadata_url": "https://huggingface.co/.well-known/openid-configuration",
+ "client_kwargs": {"scope": "openid profile email"},
+ "icon_class": "fas fa-robot",
+ "button_class": "oauth-btn-huggingface",
+ },
+}
+
+
+class OAuthBackend(AuthBackend):
+ """Authentication backend using OAuth 2.0 / OIDC providers.
+
+ This backend delegates authentication to external identity providers
+ (Google, GitHub, or any OIDC-compliant provider). After a user
+ authenticates with the provider, they are mapped to a Potato user
+ identity based on the configured identity field (email, username, etc.).
+ """
+
+ def __init__(self, auth_config: Dict[str, Any]):
+ """Initialize the OAuth backend from authentication config.
+
+ Args:
+ auth_config: The 'authentication' section of the Potato YAML config.
+ Must contain 'providers' dict with at least one provider.
+ """
+ self.users = {} # username -> user profile data
+ self.providers_config = auth_config.get("providers", {})
+ self.user_identity_field = auth_config.get("user_identity_field", "email")
+ self.auto_register = auth_config.get("auto_register", True)
+ self.allow_local_login = auth_config.get("allow_local_login", False)
+ self._oauth = None # Initialized later via init_oauth()
+ self._provider_metadata = {} # Cached display info per provider
+
+ if not self.providers_config:
+ raise ValueError(
+ "OAuth authentication requires at least one provider in "
+ "'authentication.providers'"
+ )
+
+ # Build provider metadata for login page rendering
+ for name, pconfig in self.providers_config.items():
+ well_known = PROVIDER_CONFIGS.get(name, {})
+ self._provider_metadata[name] = {
+ "name": name,
+ "display_name": pconfig.get(
+ "display_name", well_known.get("display_name", name.title())
+ ),
+ "icon_class": well_known.get("icon_class", "fas fa-sign-in-alt"),
+ "button_class": well_known.get("button_class", "oauth-btn-generic"),
+ }
+
+ logger.info(
+ "OAuth backend initialized with %d provider(s)",
+ len(self.providers_config),
+ )
+
+ def init_oauth(self, app):
+ """Register OAuth providers with the Flask app.
+
+ Must be called after the Flask app is created but before any
+ OAuth routes are used.
+
+ Args:
+ app: The Flask application instance.
+ """
+ self._oauth = OAuth(app)
+
+ for name, pconfig in self.providers_config.items():
+ well_known = PROVIDER_CONFIGS.get(name, {})
+
+ # Build the registration kwargs
+ reg_kwargs = {
+ "client_id": self._resolve_env(pconfig.get("client_id", "")),
+ "client_secret": self._resolve_env(pconfig.get("client_secret", "")),
+ }
+
+ # For well-known providers, use built-in config
+ if name == "google":
+ reg_kwargs["server_metadata_url"] = well_known["server_metadata_url"]
+ reg_kwargs["client_kwargs"] = well_known["client_kwargs"]
+ elif name == "github":
+ reg_kwargs["api_base_url"] = well_known["api_base_url"]
+ reg_kwargs["access_token_url"] = well_known["access_token_url"]
+ reg_kwargs["authorize_url"] = well_known["authorize_url"]
+ # GitHub needs read:org scope if allowed_org is set
+ scope = "user:email"
+ if pconfig.get("allowed_org"):
+ scope = "user:email read:org"
+ reg_kwargs["client_kwargs"] = {"scope": scope}
+ elif name == "oidc" or "discovery_url" in pconfig:
+ # Generic OIDC provider
+ discovery_url = pconfig.get("discovery_url", "")
+ if not discovery_url:
+ raise ValueError(
+ f"OIDC provider '{name}' requires 'discovery_url'"
+ )
+ reg_kwargs["server_metadata_url"] = discovery_url
+ scopes = pconfig.get("scopes", ["openid", "email", "profile"])
+ reg_kwargs["client_kwargs"] = {"scope": " ".join(scopes)}
+
+ self._oauth.register(name, **reg_kwargs)
+ logger.info("Registered OAuth provider: %s", name)
+
+ def get_oauth_client(self, provider_name: str):
+ """Get the Authlib OAuth client for a provider.
+
+ Args:
+ provider_name: The provider key (e.g. 'google', 'github', 'oidc').
+
+ Returns:
+ The Authlib OAuth client, or None if provider not found.
+ """
+ if self._oauth is None:
+ logger.error("OAuth not initialized โ call init_oauth(app) first")
+ return None
+ client = getattr(self._oauth, provider_name, None)
+ if client is None:
+ logger.error("Unknown OAuth provider: %s", provider_name)
+ return client
+
+ def get_login_providers(self) -> List[Dict[str, str]]:
+ """Return display metadata for all configured providers.
+
+ Used by the login template to render SSO buttons.
+
+ Returns:
+ List of dicts with keys: name, display_name, icon_class, button_class, login_url
+ """
+ providers = []
+ for name, meta in self._provider_metadata.items():
+ providers.append({
+ **meta,
+ "login_url": f"/auth/login/{name}",
+ })
+ return providers
+
+ def extract_user_id(self, profile: Dict[str, Any], provider_name: str = None) -> str:
+ """Extract the Potato user identity from an OAuth profile.
+
+ Args:
+ profile: The user profile dict from the OAuth provider.
+ provider_name: The provider name (for fallback logic).
+
+ Returns:
+ The user identity string.
+
+ Raises:
+ ValueError: If no usable identity field is found.
+ """
+ field = self.user_identity_field
+
+ # Try the configured field first
+ if field == "email" and profile.get("email"):
+ return profile["email"]
+ if field == "username":
+ # GitHub uses 'login', others may use 'preferred_username'
+ for key in ("login", "preferred_username", "username"):
+ if profile.get(key):
+ return profile[key]
+ if field == "sub" and profile.get("sub"):
+ return str(profile["sub"])
+ if field == "name" and profile.get("name"):
+ return profile["name"]
+
+ # Fallback chain: email -> login -> sub
+ if profile.get("email"):
+ return profile["email"]
+ if profile.get("login"):
+ return profile["login"]
+ if profile.get("sub"):
+ return str(profile["sub"])
+
+ raise ValueError(
+ f"Cannot extract user identity from OAuth profile. "
+ f"Configured field '{field}' not found, and no fallback available. "
+ f"Profile keys: {list(profile.keys())}"
+ )
+
+ def check_restrictions(self, provider_name: str, profile: Dict[str, Any]) -> tuple:
+ """Check domain/org restrictions for a provider.
+
+ Args:
+ provider_name: The provider key.
+ profile: The user profile dict.
+
+ Returns:
+ Tuple of (allowed: bool, reason: str). reason is empty if allowed.
+ """
+ pconfig = self.providers_config.get(provider_name, {})
+
+ # Check Google domain restriction
+ allowed_domain = pconfig.get("allowed_domain")
+ if allowed_domain:
+ email = profile.get("email", "")
+ domain = email.split("@")[-1] if "@" in email else ""
+ if domain.lower() != allowed_domain.lower():
+ return False, (
+ f"Access restricted to {allowed_domain} accounts. "
+ f"Your email domain ({domain}) is not allowed."
+ )
+
+ # GitHub org restriction is checked separately via API (see routes)
+ # We store the config here for routes to use
+ # Note: allowed_org check requires an API call with the user's token,
+ # which happens in the route handler, not here.
+
+ return True, ""
+
+ # --- AuthBackend interface ---
+
+ def authenticate(self, username: str, password: Optional[str]) -> bool:
+ """Authenticate an OAuth user.
+
+ For OAuth, authentication happens via the provider redirect flow,
+ not via username/password. This method returns True if the user
+ exists (was previously authenticated via OAuth).
+ """
+ return username in self.users
+
+ def add_user(self, username: str, password: Optional[str], **kwargs) -> str:
+ """Register an OAuth-authenticated user."""
+ if username in self.users:
+ # Update profile data
+ self.users[username].update(kwargs)
+ return "Success"
+ self.users[username] = kwargs
+ return "Success"
+
+ def is_valid_username(self, username: str) -> bool:
+ """Check if a username was registered via OAuth."""
+ return username in self.users
+
+ def update_password(self, username: str, new_password: str) -> bool:
+ """Not supported for OAuth - passwords are managed by providers."""
+ raise NotImplementedError("Password management is handled by OAuth providers")
+
+ def get_all_users(self) -> list:
+ """Return all registered OAuth usernames."""
+ return list(self.users.keys())
+
+ def get_allowed_org(self, provider_name: str) -> Optional[str]:
+ """Get the allowed_org restriction for a provider, if any."""
+ return self.providers_config.get(provider_name, {}).get("allowed_org")
+
+ # --- Helpers ---
+
+ @staticmethod
+ def _resolve_env(value: str) -> str:
+ """Resolve ${ENV_VAR} references in a config value."""
+ if not isinstance(value, str):
+ return value
+ if value.startswith("${") and value.endswith("}"):
+ env_name = value[2:-1]
+ resolved = os.environ.get(env_name, "")
+ if not resolved:
+ logger.warning(
+ "Environment variable referenced in OAuth config is not set"
+ )
+ return resolved
+ return value
diff --git a/potato/authentication.py b/potato/authentication.py
new file mode 100644
index 0000000000000000000000000000000000000000..56fe0949d20f9cea38893f190e0ed851c06a299c
--- /dev/null
+++ b/potato/authentication.py
@@ -0,0 +1,751 @@
+"""
+Authentication System Module
+
+This module provides a comprehensive authentication system for the Potato annotation platform.
+It supports multiple authentication backends including in-memory storage, database storage,
+and third-party SSO providers like Clerk.
+
+The system is designed to be extensible and supports both password-based and passwordless
+authentication modes. It includes user management, session validation, and secure
+password handling.
+
+Key Features:
+- Multiple authentication backends (in-memory, database, Clerk SSO)
+- Password hashing with PBKDF2 and per-user salts
+- Passwordless authentication support
+- User registration and management
+- Password reset with secure tokens
+- Session-based authentication
+- Configurable authentication requirements
+"""
+
+import os
+import json
+import logging
+import hashlib
+import hmac
+import secrets
+import sqlite3
+import requests
+import threading
+import time
+from abc import ABC, abstractmethod
+from typing import Optional, Dict, Any, List, Union
+
+logger = logging.getLogger(__name__)
+
+# Global singleton instance of the user authenticator with thread-safe lock
+USER_AUTHENTICATOR_SINGLETON = None
+_USER_AUTHENTICATOR_LOCK = threading.Lock()
+
+# Format for per-user salt storage: "<32-char-hex-salt>$"
+_SALT_HASH_SEPARATOR = "$"
+
+
+def _is_salted_hash(value: str) -> bool:
+ """Check if a stored password value is in the per-user salt$hash format."""
+ if not value or _SALT_HASH_SEPARATOR not in value:
+ return False
+ parts = value.split(_SALT_HASH_SEPARATOR, 1)
+ # salt is 32 hex chars (16 bytes), hash is 64 hex chars (32 bytes sha256)
+ return len(parts) == 2 and len(parts[0]) == 32 and len(parts[1]) == 64
+
+
+def _hash_password_with_salt(password: str, salt: str = None) -> str:
+ """Hash a password with a per-user salt using PBKDF2.
+
+ Args:
+ password: The plain text password to hash
+ salt: Hex-encoded salt string. If None, generates a new random salt.
+
+ Returns:
+ str: The combined "salt$hash" string
+ """
+ if not password:
+ return ""
+ if salt is None:
+ salt = secrets.token_hex(16)
+ hash_value = hashlib.pbkdf2_hmac(
+ 'sha256',
+ password.encode('utf-8'),
+ salt.encode('utf-8'),
+ 100000
+ ).hex()
+ return f"{salt}{_SALT_HASH_SEPARATOR}{hash_value}"
+
+
+def _verify_password(password: str, stored: str) -> bool:
+ """Verify a password against a stored salt$hash value using constant-time comparison."""
+ if not password or not stored:
+ return False
+ if not _is_salted_hash(stored):
+ return False
+ salt, expected_hash = stored.split(_SALT_HASH_SEPARATOR, 1)
+ actual_hash = hashlib.pbkdf2_hmac(
+ 'sha256',
+ password.encode('utf-8'),
+ salt.encode('utf-8'),
+ 100000
+ ).hex()
+ return hmac.compare_digest(expected_hash, actual_hash)
+
+
+class AuthBackend(ABC):
+ """
+ Abstract base class for authentication backends.
+
+ This class defines the interface that all authentication backends must implement.
+ It provides a consistent API for user authentication, registration, and validation
+ regardless of the underlying storage mechanism.
+ """
+ @abstractmethod
+ def authenticate(self, username: str, password: Optional[str]) -> bool:
+ """Authenticate a user against this backend."""
+ pass
+
+ @abstractmethod
+ def add_user(self, username: str, password: Optional[str], **kwargs) -> str:
+ """Add a user to this backend. Returns status message."""
+ pass
+
+ @abstractmethod
+ def is_valid_username(self, username: str) -> bool:
+ """Check if a username exists in this backend."""
+ pass
+
+ @abstractmethod
+ def update_password(self, username: str, new_password: str) -> bool:
+ """Update a user's password. Returns True on success."""
+ pass
+
+ @abstractmethod
+ def get_all_users(self) -> List[str]:
+ """Return list of all usernames."""
+ pass
+
+ def add_user_prehashed(self, username: str, hashed_password: str, **kwargs) -> str:
+ """Load user with already-hashed password (for file loading). Override in subclasses."""
+ raise NotImplementedError("This backend does not support loading pre-hashed passwords")
+
+
+class InMemoryAuthBackend(AuthBackend):
+ """
+ Authentication backend that stores users in memory with per-user salts.
+
+ Password storage format: "salt$hash" where salt is 32 hex chars and hash is 64 hex chars.
+ """
+ def __init__(self):
+ self.users = {} # username -> "salt$hash"
+ self.user_data = {} # username -> additional data
+
+ def authenticate(self, username: str, password: Optional[str]) -> bool:
+ if username not in self.users:
+ return False
+ if password is None: # Passwordless login
+ return True
+ return _verify_password(password, self.users[username])
+
+ def add_user(self, username: str, password: Optional[str], **kwargs) -> str:
+ if username in self.users:
+ return "Duplicate user"
+ self.users[username] = _hash_password_with_salt(password) if password else ""
+ self.user_data[username] = kwargs
+ return "Success"
+
+ def add_user_prehashed(self, username: str, hashed_password: str, **kwargs) -> str:
+ """Store a user with an already-hashed password (salt$hash format)."""
+ if username in self.users:
+ return "Duplicate user"
+ self.users[username] = hashed_password
+ self.user_data[username] = kwargs
+ return "Success"
+
+ def is_valid_username(self, username: str) -> bool:
+ return username in self.users
+
+ def update_password(self, username: str, new_password: str) -> bool:
+ if username not in self.users:
+ return False
+ self.users[username] = _hash_password_with_salt(new_password)
+ return True
+
+ def get_all_users(self) -> List[str]:
+ return list(self.users.keys())
+
+
+class DatabaseAuthBackend(AuthBackend):
+ """
+ Authentication backend using SQLite (stdlib) or PostgreSQL (psycopg2).
+
+ Connection string formats:
+ sqlite:///path/to/db.db (relative or absolute)
+ postgresql://user:pass@host/dbname
+ """
+ def __init__(self, db_connection_string: str):
+ self.db_connection_string = db_connection_string
+ self._lock = threading.Lock()
+ self._db_type = None # 'sqlite' or 'postgresql'
+ self._connection = None
+
+ if db_connection_string.startswith("sqlite:///"):
+ self._db_type = "sqlite"
+ self._init_sqlite(db_connection_string[len("sqlite:///"):])
+ elif db_connection_string.startswith("postgresql://"):
+ self._db_type = "postgresql"
+ self._init_postgresql(db_connection_string)
+ else:
+ raise ValueError(
+ f"Unsupported database URL: {db_connection_string}. "
+ "Use sqlite:///path/to/db or postgresql://user:pass@host/dbname"
+ )
+
+ logger.info(f"Database auth backend initialized ({self._db_type})")
+
+ def _init_sqlite(self, db_path: str):
+ """Initialize SQLite database."""
+ # Create parent directories if needed
+ db_dir = os.path.dirname(db_path)
+ if db_dir:
+ os.makedirs(db_dir, exist_ok=True)
+
+ self._connection = sqlite3.connect(db_path, check_same_thread=False)
+ self._connection.execute("PRAGMA journal_mode=WAL")
+ self._connection.execute("""
+ CREATE TABLE IF NOT EXISTS users (
+ username TEXT PRIMARY KEY,
+ password_hash TEXT NOT NULL,
+ email TEXT,
+ created_at TEXT DEFAULT (datetime('now')),
+ updated_at TEXT DEFAULT (datetime('now'))
+ )
+ """)
+ self._connection.commit()
+
+ def _init_postgresql(self, connection_string: str):
+ """Initialize PostgreSQL database."""
+ try:
+ import psycopg2
+ except ImportError:
+ raise ImportError(
+ "psycopg2 is required for PostgreSQL authentication backend. "
+ "Install it with: pip install psycopg2-binary"
+ )
+ self._connection = psycopg2.connect(connection_string)
+ self._connection.autocommit = True
+ with self._connection.cursor() as cur:
+ cur.execute("""
+ CREATE TABLE IF NOT EXISTS users (
+ username TEXT PRIMARY KEY,
+ password_hash TEXT NOT NULL,
+ email TEXT,
+ created_at TIMESTAMP DEFAULT NOW(),
+ updated_at TIMESTAMP DEFAULT NOW()
+ )
+ """)
+
+ def _execute(self, query: str, params: tuple = (), fetch: str = None):
+ """Thread-safe query execution.
+
+ Args:
+ query: SQL query with ? placeholders (auto-converted to %s for PostgreSQL)
+ params: Query parameters
+ fetch: None, 'one', or 'all'
+
+ Returns:
+ Query result based on fetch parameter
+ """
+ with self._lock:
+ if self._db_type == "postgresql":
+ query = query.replace("?", "%s")
+
+ if self._db_type == "sqlite":
+ cursor = self._connection.cursor()
+ cursor.execute(query, params)
+ if fetch == "one":
+ result = cursor.fetchone()
+ elif fetch == "all":
+ result = cursor.fetchall()
+ else:
+ self._connection.commit()
+ result = None
+ cursor.close()
+ return result
+ else:
+ with self._connection.cursor() as cur:
+ cur.execute(query, params)
+ if fetch == "one":
+ return cur.fetchone()
+ elif fetch == "all":
+ return cur.fetchall()
+ return None
+
+ def authenticate(self, username: str, password: Optional[str]) -> bool:
+ row = self._execute(
+ "SELECT password_hash FROM users WHERE username = ?",
+ (username,), fetch="one"
+ )
+ if not row:
+ return False
+ if password is None: # Passwordless login
+ return True
+ return _verify_password(password, row[0])
+
+ def add_user(self, username: str, password: Optional[str], **kwargs) -> str:
+ existing = self._execute(
+ "SELECT 1 FROM users WHERE username = ?",
+ (username,), fetch="one"
+ )
+ if existing:
+ return "Duplicate user"
+
+ hashed = _hash_password_with_salt(password) if password else ""
+ email = kwargs.get("email", "")
+ self._execute(
+ "INSERT INTO users (username, password_hash, email) VALUES (?, ?, ?)",
+ (username, hashed, email)
+ )
+ return "Success"
+
+ def add_user_prehashed(self, username: str, hashed_password: str, **kwargs) -> str:
+ """Store a user with an already-hashed password."""
+ existing = self._execute(
+ "SELECT 1 FROM users WHERE username = ?",
+ (username,), fetch="one"
+ )
+ if existing:
+ return "Duplicate user"
+
+ email = kwargs.get("email", "")
+ self._execute(
+ "INSERT INTO users (username, password_hash, email) VALUES (?, ?, ?)",
+ (username, hashed_password, email)
+ )
+ return "Success"
+
+ def is_valid_username(self, username: str) -> bool:
+ row = self._execute(
+ "SELECT 1 FROM users WHERE username = ?",
+ (username,), fetch="one"
+ )
+ return row is not None
+
+ def update_password(self, username: str, new_password: str) -> bool:
+ if not self.is_valid_username(username):
+ return False
+ hashed = _hash_password_with_salt(new_password)
+ if self._db_type == "sqlite":
+ self._execute(
+ "UPDATE users SET password_hash = ?, updated_at = datetime('now') WHERE username = ?",
+ (hashed, username)
+ )
+ else:
+ self._execute(
+ "UPDATE users SET password_hash = ?, updated_at = NOW() WHERE username = ?",
+ (hashed, username)
+ )
+ return True
+
+ def get_all_users(self) -> List[str]:
+ rows = self._execute("SELECT username FROM users", fetch="all")
+ return [r[0] for r in rows]
+
+ def close(self):
+ """Close the database connection."""
+ if self._connection:
+ self._connection.close()
+ self._connection = None
+
+
+class ClerkAuthBackend(AuthBackend):
+ """
+ Authentication backend that uses Clerk for SSO.
+ """
+ def __init__(self, api_key: str, frontend_api: str):
+ self.api_key = api_key
+ self.frontend_api = frontend_api
+ self.users = {} # Cache of known users
+ logger.info("Clerk SSO backend initialized")
+
+ def authenticate(self, username: str, token: Optional[str]) -> bool:
+ if not token:
+ return False
+ try:
+ headers = {
+ "Authorization": f"Bearer {self.api_key}",
+ "Content-Type": "application/json"
+ }
+ response = requests.get(
+ f"https://api.clerk.dev/v1/sessions/{token}",
+ headers=headers
+ )
+ if response.status_code == 200:
+ user_data = response.json()
+ self.users[username] = user_data
+ return True
+ return False
+ except Exception as e:
+ logger.error(f"Error authenticating with Clerk: {str(e)}")
+ return False
+
+ def add_user(self, username: str, password: Optional[str], **kwargs) -> str:
+ return "User management happens through Clerk dashboard"
+
+ def is_valid_username(self, username: str) -> bool:
+ return username in self.users
+
+ def update_password(self, username: str, new_password: str) -> bool:
+ raise NotImplementedError("Password management is handled by Clerk")
+
+ def get_all_users(self) -> List[str]:
+ return list(self.users.keys())
+
+
+class UserAuthenticator:
+ """
+ A class for maintaining state on which users are allowed to use the system.
+
+ This class provides a unified interface for user authentication and management
+ regardless of the underlying backend. It supports multiple authentication methods
+ and can be configured for passwordless operation.
+ """
+
+ def __init__(self, user_config_path, auth_method="in_memory", auth_config=None):
+ self.allow_all_users = True
+ self.user_config_path = user_config_path
+ self.user_config_path_explicit = False # Set to True if path was explicitly configured
+ self.authorized_users = []
+ self.userlist = []
+ self.usernames = set()
+ self.users = {}
+ self.required_user_info_keys = ["username", "password"]
+ self.require_password = True
+ self.auth_method = auth_method
+ self.auth_config = auth_config or {}
+ self.auth_backend = self._initialize_backend(auth_method, auth_config)
+
+ # Token management for password reset
+ self._reset_tokens = {} # token -> {username, expires}
+ self._token_lock = threading.Lock()
+
+ # Track load outcomes so init_from_config can warn on a silently empty
+ # (e.g. wrong-format) user file. F-036.
+ self.users_loaded_from_file = 0
+ self.user_file_parse_errors = 0
+
+ # Load users from config file if it exists
+ if os.path.isfile(self.user_config_path):
+ logger.info(f"Loading users from {self.user_config_path}")
+ before = len(self.users)
+ with open(self.user_config_path, "rt", encoding="utf-8") as f:
+ for lineno, line in enumerate(f.readlines(), start=1):
+ line = line.strip()
+ if not line:
+ continue
+ # Tolerate a malformed line instead of aborting the whole
+ # load (and crashing server boot) on one bad row.
+ try:
+ single_user = json.loads(line)
+ except (ValueError, TypeError) as e:
+ self.user_file_parse_errors += 1
+ logger.error(
+ f"User file {self.user_config_path} line {lineno}: "
+ f"not valid JSON ({e}); skipping. Expected JSONL โ "
+ f'one object per line, e.g. {{"username": "alice", "password": "x"}}'
+ )
+ continue
+ # Detect salt$hash format in password field
+ password_val = single_user.get("password", "") if isinstance(single_user, dict) else ""
+ if password_val and _is_salted_hash(password_val):
+ self._add_user_prehashed(single_user)
+ else:
+ self.add_single_user(single_user)
+ self.users_loaded_from_file = len(self.users) - before
+
+ def _initialize_backend(self, auth_method: str, auth_config: dict = None) -> AuthBackend:
+ if auth_method == "in_memory":
+ return InMemoryAuthBackend()
+ elif auth_method == "database":
+ db_url = (auth_config or {}).get("database_url") or \
+ os.environ.get("POTATO_DB_CONNECTION", "sqlite:///potato_users.db")
+ return DatabaseAuthBackend(db_url)
+ elif auth_method == "clerk":
+ api_key = os.environ.get("CLERK_API_KEY", "")
+ frontend_api = os.environ.get("CLERK_FRONTEND_API", "")
+ if not api_key:
+ logger.error("CLERK_API_KEY environment variable is not set")
+ raise ValueError("CLERK_API_KEY must be set for Clerk authentication")
+ return ClerkAuthBackend(api_key, frontend_api)
+ elif auth_method == "oauth":
+ from potato.auth_backends.oauth_backend import OAuthBackend
+ if not auth_config:
+ raise ValueError("OAuth authentication requires an 'authentication' config section with 'providers'")
+ return OAuthBackend(auth_config)
+ else:
+ logger.error(f"Unknown authentication method: {auth_method}")
+ raise ValueError(f"Unknown authentication method: {auth_method}")
+
+ @staticmethod
+ def init_from_config(config: dict) -> "UserAuthenticator":
+ """Initialize the UserAuthenticator from a configuration dictionary (singleton)."""
+ global USER_AUTHENTICATOR_SINGLETON
+
+ if USER_AUTHENTICATOR_SINGLETON is None:
+ with _USER_AUTHENTICATOR_LOCK:
+ if USER_AUTHENTICATOR_SINGLETON is None:
+ auth_method = config.get("authentication", {}).get("method", "in_memory")
+ user_config_path = config.get("authentication", {}).get("user_config_path", None)
+ require_password = config.get("require_password", True)
+
+ path_explicit = user_config_path is not None
+
+ if user_config_path is None:
+ config_dir = os.path.dirname(config['output_annotation_dir'])
+ user_config_path = os.path.join(config_dir, "user_config.json")
+ else:
+ # Don't raise if file doesn't exist โ it will be created on first registration
+ if not os.path.isfile(user_config_path):
+ logger.info(f"user_config_path '{user_config_path}' does not exist yet; will be created on first registration")
+
+ logger.debug(f"User config path: {user_config_path}")
+
+ auth_config = config.get("authentication", {})
+
+ USER_AUTHENTICATOR_SINGLETON = UserAuthenticator(user_config_path, auth_method, auth_config)
+ USER_AUTHENTICATOR_SINGLETON.require_password = require_password
+ USER_AUTHENTICATOR_SINGLETON.user_config_path_explicit = path_explicit
+
+ # F-036: a user file was explicitly configured and exists, but
+ # produced zero usable users (e.g. wrong format / all rows
+ # invalid). With closed enrolment this is a silently broken
+ # deployment โ nobody can log in. Warn prominently.
+ _auth = USER_AUTHENTICATOR_SINGLETON
+ if (path_explicit and os.path.isfile(user_config_path)
+ and _auth.users_loaded_from_file == 0):
+ allow_all = config.get("user_config", {}).get("allow_all_users", False)
+ logger.warning(
+ "user_config_path '%s' was configured but loaded 0 users "
+ "(%d malformed line(s)). Expected JSONL โ one object per "
+ 'line, e.g. {"username": "alice", "password": "x"}. %s',
+ user_config_path, _auth.user_file_parse_errors,
+ ("Open registration is on, so new users can still self-register."
+ if allow_all else
+ "allow_all_users is false, so NO ONE will be able to log in."),
+ )
+
+ logger.info(f"Initialized UserAuthenticator with method: {auth_method}, require_password: {require_password}")
+
+ return USER_AUTHENTICATOR_SINGLETON
+
+ @staticmethod
+ def get_instance():
+ global USER_AUTHENTICATOR_SINGLETON
+ if USER_AUTHENTICATOR_SINGLETON is None:
+ raise ValueError("UserAuthenticator not initialized; call init_from_config first")
+ return USER_AUTHENTICATOR_SINGLETON
+
+ @staticmethod
+ def authenticate(username: str, password: Optional[str]) -> bool:
+ authenticator = UserAuthenticator.get_instance()
+
+ if not authenticator.auth_backend.is_valid_username(username):
+ logger.warning(f"Authentication failed: user '{username}' does not exist")
+ return False
+
+ if not authenticator.require_password:
+ logger.debug(f"Passwordless authentication for user: {username}")
+ return authenticator.auth_backend.authenticate(username, None)
+
+ return authenticator.auth_backend.authenticate(username, password)
+
+ def add_user(self, username, password: Optional[str], **kwargs):
+ """Add a user to the authentication system."""
+ if not self.require_password:
+ logger.debug(f"Passwordless mode - allowing any user: {username}")
+ elif self.allow_all_users == False and not self.is_authorized_user(username):
+ return "Unauthorized user"
+
+ result = self.auth_backend.add_user(username, password, **kwargs)
+ if result == "Success":
+ user_data = {"username": username}
+ user_data.update(kwargs)
+ self.users[username] = user_data
+ self.userlist.append(username)
+ return result
+
+ def _add_user_prehashed(self, single_user):
+ """Add a user with an already-hashed password (loaded from file)."""
+ username = single_user["username"]
+ hashed_password = single_user.get("password", "")
+
+ result = self.auth_backend.add_user_prehashed(
+ username,
+ hashed_password,
+ **{k: v for k, v in single_user.items() if k not in ["username", "password"]}
+ )
+
+ if result == "Success":
+ self.users[username] = single_user
+ self.userlist.append(username)
+
+ return result
+
+ def add_single_user(self, single_user):
+ """Add a single user to the full user dict."""
+ if not self.require_password:
+ logger.debug(f"Passwordless mode - allowing any user: {single_user['username']}")
+ elif self.allow_all_users == False and not self.is_authorized_user(single_user["username"]):
+ return "Unauthorized user"
+
+ if not self.require_password:
+ required_keys = ["username"]
+ else:
+ required_keys = self.required_user_info_keys
+
+ for key in required_keys:
+ if key not in single_user:
+ logger.error(f"Missing {key} in user info")
+ return f"Missing {key} in user info"
+
+ result = self.auth_backend.add_user(
+ single_user["username"],
+ single_user.get("password"),
+ **{k: v for k, v in single_user.items() if k not in ["username", "password"]}
+ )
+
+ if result == "Success":
+ self.users[single_user["username"]] = single_user
+ self.userlist.append(single_user["username"])
+
+ return result
+
+ def update_password(self, username: str, new_password: str) -> bool:
+ """Update a user's password via the backend."""
+ result = self.auth_backend.update_password(username, new_password)
+ if result and username in self.users:
+ # Update the stored user dict with the new hash for save_user_config
+ if isinstance(self.users[username], dict):
+ self.users[username]["password"] = self.auth_backend.users[username] \
+ if hasattr(self.auth_backend, 'users') else _hash_password_with_salt(new_password)
+ return result
+
+ def save_user_config(self):
+ """Save user config to file.
+
+ Saves when:
+ - auth_method is in_memory AND user_config_path was explicitly configured
+ - auth_method is not in_memory and not database (other file-based methods)
+
+ Skips when:
+ - auth_method is database (DB handles its own persistence)
+ - auth_method is in_memory with auto-generated default path (preserve old behavior)
+ """
+ if self.auth_method == "database":
+ logger.debug("User config not saved - using database authentication (DB handles persistence)")
+ return
+
+ if self.auth_method == "in_memory" and not self.user_config_path_explicit:
+ logger.debug("User config not saved - using in_memory with default path")
+ return
+
+ if self.user_config_path:
+ with open(self.user_config_path, "wt", encoding="utf-8") as f:
+ for k in self.userlist:
+ user_data = self.users.get(k, {})
+ if isinstance(user_data, dict):
+ # Ensure password field contains the hashed value
+ output = dict(user_data)
+ if hasattr(self.auth_backend, 'users') and k in self.auth_backend.users:
+ output["password"] = self.auth_backend.users[k]
+ f.write(json.dumps(output) + "\n")
+ else:
+ f.write(json.dumps({"username": k}) + "\n")
+ logger.info(f"User info file saved at: {self.user_config_path}")
+ else:
+ logger.warning("WARNING: user_config_path not specified, user registration info are not saved")
+
+ # --- Token-based password reset ---
+
+ def create_reset_token(self, username: str, ttl_hours: int = 24) -> Optional[str]:
+ """Create a password reset token for a user.
+
+ Args:
+ username: The username to create a token for
+ ttl_hours: Token validity in hours (default 24)
+
+ Returns:
+ The token string, or None if user doesn't exist
+ """
+ if not self.auth_backend.is_valid_username(username):
+ return None
+
+ token = secrets.token_urlsafe(32)
+ expires = time.time() + (ttl_hours * 3600)
+
+ with self._token_lock:
+ # Invalidate any existing tokens for this user
+ self._reset_tokens = {
+ t: v for t, v in self._reset_tokens.items()
+ if v["username"] != username
+ }
+ self._reset_tokens[token] = {
+ "username": username,
+ "expires": expires
+ }
+
+ return token
+
+ def validate_reset_token(self, token: str) -> Optional[str]:
+ """Validate a reset token and return the username, or None if invalid/expired."""
+ with self._token_lock:
+ # Clean expired tokens
+ now = time.time()
+ self._reset_tokens = {
+ t: v for t, v in self._reset_tokens.items()
+ if v["expires"] > now
+ }
+
+ if token not in self._reset_tokens:
+ return None
+ return self._reset_tokens[token]["username"]
+
+ def consume_reset_token(self, token: str) -> Optional[str]:
+ """Validate, delete, and return the username for a reset token. Single-use."""
+ with self._token_lock:
+ now = time.time()
+ self._reset_tokens = {
+ t: v for t, v in self._reset_tokens.items()
+ if v["expires"] > now
+ }
+
+ if token not in self._reset_tokens:
+ return None
+ username = self._reset_tokens[token]["username"]
+ del self._reset_tokens[token]
+ return username
+
+ # --- End token management ---
+
+ def is_authorized_user(self, username):
+ return username in self.authorized_users
+
+ def is_valid_username(self, username):
+ return self.auth_backend.is_valid_username(username)
+
+ def is_valid_password(self, username, password):
+ return self.authenticate(username, password)
+
+ def get_clerk_frontend_api(self) -> str:
+ if self.auth_method == "clerk" and isinstance(self.auth_backend, ClerkAuthBackend):
+ return self.auth_backend.frontend_api
+ return ""
+
+ def get_oauth_backend(self):
+ if self.auth_method == "oauth":
+ return self.auth_backend
+ return None
+
+ def get_login_providers(self) -> list:
+ oauth_backend = self.get_oauth_backend()
+ if oauth_backend:
+ return oauth_backend.get_login_providers()
+ return []
diff --git a/potato/bws_scoring.py b/potato/bws_scoring.py
new file mode 100644
index 0000000000000000000000000000000000000000..0c78c7f0644f667ed8bb14dde1185bbc24f96323
--- /dev/null
+++ b/potato/bws_scoring.py
@@ -0,0 +1,451 @@
+"""
+Best-Worst Scaling Score Estimation
+
+Computes item scores from BWS annotations using three methods:
+1. Counting: score = (best_count - worst_count) / appearances (no dependencies)
+2. Bradley-Terry: pairwise comparison model via choix (requires choix)
+3. Plackett-Luce: partial ranking model via choix (requires choix)
+
+Usage as library:
+ from potato.bws_scoring import BwsScorer
+ scorer = BwsScorer(annotations, pool_items, id_key)
+ scores = scorer.counting()
+
+Usage as CLI:
+ python -m potato.bws_scoring --config config.yaml --method counting
+"""
+
+import argparse
+import csv
+import json
+import logging
+import os
+import sys
+from typing import Any, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+class BwsScorer:
+ """Compute BWS scores from annotations."""
+
+ def __init__(
+ self,
+ annotations: List[Dict[str, Any]],
+ pool_items: List[Dict[str, Any]],
+ id_key: str,
+ text_key: str = "text",
+ ):
+ """
+ Args:
+ annotations: List of annotation dicts, each with:
+ - "instance_id": tuple instance ID (e.g. "bws_tuple_0001")
+ - "bws_items": list of {source_id, text, position}
+ - "best": position label (e.g. "B")
+ - "worst": position label (e.g. "D")
+ - "annotator": username
+ pool_items: Original pool items
+ id_key: Key for item IDs in pool_items
+ text_key: Key for item text in pool_items
+ """
+ self.annotations = annotations
+ self.pool_items = pool_items
+ self.id_key = id_key
+ self.text_key = text_key
+
+ # Build item index
+ self.item_ids = [str(item[id_key]) for item in pool_items]
+ self.item_texts = {
+ str(item[id_key]): str(item.get(text_key, ""))
+ for item in pool_items
+ }
+ self.item_id_to_idx = {iid: idx for idx, iid in enumerate(self.item_ids)}
+
+ def _resolve_annotation(
+ self, ann: Dict[str, Any]
+ ) -> Optional[Tuple[str, str, List[str]]]:
+ """Resolve an annotation to (best_source_id, worst_source_id, all_source_ids).
+
+ Returns None if annotation is incomplete.
+ """
+ best_pos = ann.get("best")
+ worst_pos = ann.get("worst")
+ bws_items = ann.get("bws_items", [])
+
+ if not best_pos or not worst_pos or not bws_items:
+ return None
+
+ pos_to_id = {item["position"]: item["source_id"] for item in bws_items}
+ best_id = pos_to_id.get(best_pos)
+ worst_id = pos_to_id.get(worst_pos)
+
+ if not best_id or not worst_id:
+ return None
+
+ all_ids = [item["source_id"] for item in bws_items]
+ return best_id, worst_id, all_ids
+
+ def counting(self) -> Dict[str, Dict[str, Any]]:
+ """Counting method: score = (best_count - worst_count) / appearances.
+
+ Returns dict mapping item_id to {score, best_count, worst_count, appearances, text}.
+ """
+ best_counts = {iid: 0 for iid in self.item_ids}
+ worst_counts = {iid: 0 for iid in self.item_ids}
+ appearances = {iid: 0 for iid in self.item_ids}
+
+ for ann in self.annotations:
+ resolved = self._resolve_annotation(ann)
+ if not resolved:
+ continue
+
+ best_id, worst_id, all_ids = resolved
+ for iid in all_ids:
+ if iid in appearances:
+ appearances[iid] += 1
+ if best_id in best_counts:
+ best_counts[best_id] += 1
+ if worst_id in worst_counts:
+ worst_counts[worst_id] += 1
+
+ scores = {}
+ for iid in self.item_ids:
+ app = appearances[iid]
+ if app > 0:
+ score = (best_counts[iid] - worst_counts[iid]) / app
+ else:
+ score = 0.0
+
+ scores[iid] = {
+ "score": score,
+ "best_count": best_counts[iid],
+ "worst_count": worst_counts[iid],
+ "appearances": app,
+ "text": self.item_texts.get(iid, ""),
+ }
+
+ return scores
+
+ def bradley_terry(self) -> Dict[str, Dict[str, Any]]:
+ """Bradley-Terry model via choix.
+
+ Converts each BWS annotation to pairwise comparisons:
+ - Best item beats every other item (K-1 comparisons)
+ - Every item beats the worst item (K-1 comparisons)
+ """
+ try:
+ import choix
+ except ImportError:
+ raise ImportError(
+ "Bradley-Terry scoring requires the 'choix' package. "
+ "Install it with: pip install choix"
+ )
+
+ n_items = len(self.item_ids)
+ comparisons = []
+
+ for ann in self.annotations:
+ resolved = self._resolve_annotation(ann)
+ if not resolved:
+ continue
+
+ best_id, worst_id, all_ids = resolved
+ best_idx = self.item_id_to_idx.get(best_id)
+ worst_idx = self.item_id_to_idx.get(worst_id)
+
+ if best_idx is None or worst_idx is None:
+ continue
+
+ # Best beats all others
+ for iid in all_ids:
+ idx = self.item_id_to_idx.get(iid)
+ if idx is not None and idx != best_idx:
+ comparisons.append((best_idx, idx))
+
+ # All others beat worst
+ for iid in all_ids:
+ idx = self.item_id_to_idx.get(iid)
+ if idx is not None and idx != worst_idx:
+ comparisons.append((idx, worst_idx))
+
+ if not comparisons:
+ return {
+ iid: {"score": 0.0, "text": self.item_texts.get(iid, "")}
+ for iid in self.item_ids
+ }
+
+ params = choix.ilsr_pairwise(n_items, comparisons, alpha=0.01)
+
+ scores = {}
+ for iid in self.item_ids:
+ idx = self.item_id_to_idx[iid]
+ scores[iid] = {
+ "score": float(params[idx]),
+ "text": self.item_texts.get(iid, ""),
+ }
+
+ return scores
+
+ def plackett_luce(self) -> Dict[str, Dict[str, Any]]:
+ """Plackett-Luce model via choix.
+
+ Converts BWS to partial rankings:
+ Each annotation yields top-1 (best) selections, processed via ilsr_top1.
+ """
+ try:
+ import choix
+ except ImportError:
+ raise ImportError(
+ "Plackett-Luce scoring requires the 'choix' package. "
+ "Install it with: pip install choix"
+ )
+
+ n_items = len(self.item_ids)
+ # Use pairwise comparisons to approximate partial rankings
+ # Best > middle items, middle items > worst
+ comparisons = []
+
+ for ann in self.annotations:
+ resolved = self._resolve_annotation(ann)
+ if not resolved:
+ continue
+
+ best_id, worst_id, all_ids = resolved
+ best_idx = self.item_id_to_idx.get(best_id)
+ worst_idx = self.item_id_to_idx.get(worst_id)
+
+ if best_idx is None or worst_idx is None:
+ continue
+
+ middle_ids = [
+ iid for iid in all_ids if iid != best_id and iid != worst_id
+ ]
+
+ # Best beats all middle items
+ for iid in middle_ids:
+ idx = self.item_id_to_idx.get(iid)
+ if idx is not None:
+ comparisons.append((best_idx, idx))
+
+ # All middle items beat worst
+ for iid in middle_ids:
+ idx = self.item_id_to_idx.get(iid)
+ if idx is not None:
+ comparisons.append((idx, worst_idx))
+
+ # Best beats worst
+ comparisons.append((best_idx, worst_idx))
+
+ if not comparisons:
+ return {
+ iid: {"score": 0.0, "text": self.item_texts.get(iid, "")}
+ for iid in self.item_ids
+ }
+
+ params = choix.ilsr_pairwise(n_items, comparisons, alpha=0.01)
+
+ scores = {}
+ for iid in self.item_ids:
+ idx = self.item_id_to_idx[iid]
+ scores[iid] = {
+ "score": float(params[idx]),
+ "text": self.item_texts.get(iid, ""),
+ }
+
+ return scores
+
+ def score(self, method: str = "counting") -> Dict[str, Dict[str, Any]]:
+ """Compute scores using the specified method."""
+ if method == "counting":
+ return self.counting()
+ elif method == "bradley_terry":
+ return self.bradley_terry()
+ elif method == "plackett_luce":
+ return self.plackett_luce()
+ else:
+ raise ValueError(
+ f"Unknown scoring method: {method}. "
+ "Use 'counting', 'bradley_terry', or 'plackett_luce'."
+ )
+
+
+def write_scores(
+ scores: Dict[str, Dict[str, Any]],
+ output_path: str,
+) -> None:
+ """Write scores to a TSV file.
+
+ Output columns: item_id, text, score, best_count, worst_count, appearances, rank
+ """
+ # Sort by score descending
+ sorted_items = sorted(scores.items(), key=lambda x: x[1]["score"], reverse=True)
+
+ os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else ".", exist_ok=True)
+
+ with open(output_path, "w", newline="") as f:
+ writer = csv.writer(f, delimiter="\t")
+ writer.writerow(
+ ["item_id", "text", "score", "best_count", "worst_count", "appearances", "rank"]
+ )
+ for rank, (item_id, data) in enumerate(sorted_items, 1):
+ writer.writerow([
+ item_id,
+ data.get("text", ""),
+ f"{data['score']:.6f}",
+ data.get("best_count", ""),
+ data.get("worst_count", ""),
+ data.get("appearances", ""),
+ rank,
+ ])
+
+ logger.info(f"Wrote BWS scores to {output_path}")
+
+
+def collect_annotations_from_output(
+ output_dir: str, bws_schema_name: str, config: dict
+) -> List[Dict[str, Any]]:
+ """Collect BWS annotations from Potato's output directory.
+
+ Reads annotation files and reconstructs BWS annotation records.
+ """
+ annotations = []
+ pool_items_by_tuple = {}
+
+ # Get pool items from config
+ bws_pool = config.get("_bws_pool_items", [])
+ id_key = config["item_properties"]["id_key"]
+
+ # We need to read the saved annotations from the output dir
+ # Potato saves annotations as {output_dir}/{annotator}.jsonl
+ if not os.path.isdir(output_dir):
+ logger.warning(f"Output directory not found: {output_dir}")
+ return annotations
+
+ for fname in os.listdir(output_dir):
+ if not fname.endswith(".jsonl"):
+ continue
+
+ annotator = fname.replace(".jsonl", "")
+ fpath = os.path.join(output_dir, fname)
+
+ with open(fpath, "r") as f:
+ for line in f:
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ record = json.loads(line)
+ except json.JSONDecodeError:
+ continue
+
+ instance_id = record.get("id")
+ ann_data = record.get("annotation", {})
+
+ # Look for BWS schema annotations
+ best_val = None
+ worst_val = None
+ for schema_name, schema_ann in ann_data.items():
+ if schema_name == bws_schema_name:
+ best_val = schema_ann.get("best")
+ worst_val = schema_ann.get("worst")
+ break
+
+ if not best_val or not worst_val:
+ continue
+
+ # Get BWS items from the instance data
+ bws_items = record.get("_bws_items", [])
+
+ annotations.append({
+ "instance_id": instance_id,
+ "bws_items": bws_items,
+ "best": best_val,
+ "worst": worst_val,
+ "annotator": annotator,
+ })
+
+ return annotations
+
+
+def main():
+ """CLI entry point for BWS scoring."""
+ parser = argparse.ArgumentParser(
+ description="Compute BWS scores from Potato annotation output"
+ )
+ parser.add_argument(
+ "--config", required=True, help="Path to Potato config YAML file"
+ )
+ parser.add_argument(
+ "--method",
+ default="counting",
+ choices=["counting", "bradley_terry", "plackett_luce"],
+ help="Scoring method (default: counting)",
+ )
+ parser.add_argument(
+ "--output",
+ default=None,
+ help="Output TSV file path (default: {output_dir}/bws_scores.tsv)",
+ )
+ args = parser.parse_args()
+
+ logging.basicConfig(level=logging.INFO)
+
+ # Load config
+ import yaml
+
+ with open(args.config, "r") as f:
+ config = yaml.safe_load(f)
+
+ output_dir = config.get("output_annotation_dir", "annotation_output")
+ id_key = config["item_properties"]["id_key"]
+ text_key = config["item_properties"]["text_key"]
+
+ # Find BWS schema name
+ bws_schema_name = None
+ for scheme in config.get("annotation_schemes", []):
+ if scheme.get("annotation_type") == "bws":
+ bws_schema_name = scheme["name"]
+ break
+
+ if not bws_schema_name:
+ print("Error: No BWS annotation scheme found in config", file=sys.stderr)
+ sys.exit(1)
+
+ # Load pool items from data files
+ pool_items = []
+ for data_file in config.get("data_files", []):
+ if isinstance(data_file, dict):
+ data_file = data_file.get("path")
+ if not data_file:
+ continue
+
+ with open(data_file, "r") as f:
+ if data_file.endswith(".json"):
+ pool_items.extend(json.load(f))
+ else:
+ for line in f:
+ line = line.strip()
+ if line:
+ pool_items.append(json.loads(line))
+
+ # Collect annotations
+ annotations = collect_annotations_from_output(output_dir, bws_schema_name, config)
+
+ if not annotations:
+ print("No BWS annotations found in output directory", file=sys.stderr)
+ sys.exit(1)
+
+ print(f"Found {len(annotations)} BWS annotations for {len(pool_items)} pool items")
+
+ # Score
+ scorer = BwsScorer(annotations, pool_items, id_key, text_key)
+ scores = scorer.score(args.method)
+
+ # Write output
+ output_path = args.output or os.path.join(output_dir, "bws_scores.tsv")
+ write_scores(scores, output_path)
+ print(f"Scores written to {output_path}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/potato/bws_tuple_generator.py b/potato/bws_tuple_generator.py
new file mode 100644
index 0000000000000000000000000000000000000000..b3c97f0f9814e992e9989945435c1f7990eb8250
--- /dev/null
+++ b/potato/bws_tuple_generator.py
@@ -0,0 +1,130 @@
+"""
+Best-Worst Scaling Tuple Generator
+
+Generates tuples of K items from a pool for Best-Worst Scaling annotation.
+Each tuple is a synthetic instance containing references to K original pool items.
+Annotators select the "best" and "worst" item from each tuple.
+
+Key features:
+- Reproducible via random seed
+- Configurable tuple size and number of tuples
+- Auto-calculates num_tuples based on Louviere's guideline (2 * tuple_size appearances per item)
+- Each item appears in multiple tuples; no item repeats within a single tuple
+"""
+
+import logging
+import math
+import random
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+# Position labels: A, B, C, ... Z
+POSITION_LABELS = [chr(ord('A') + i) for i in range(26)]
+
+
+class BwsTupleGenerator:
+ """Generate BWS tuples from a pool of items."""
+
+ def __init__(
+ self,
+ pool_items: List[Dict[str, Any]],
+ id_key: str,
+ text_key: str,
+ tuple_size: int = 4,
+ num_tuples: Optional[int] = None,
+ seed: int = 42,
+ min_item_appearances: Optional[int] = None,
+ ):
+ self.pool_items = pool_items
+ self.id_key = id_key
+ self.text_key = text_key
+ self.tuple_size = tuple_size
+ self.seed = seed
+ self.min_item_appearances = min_item_appearances
+ self._num_tuples = num_tuples
+
+ def validate(self):
+ """Validate configuration before generation."""
+ if self.tuple_size < 2:
+ raise ValueError(f"tuple_size must be >= 2, got {self.tuple_size}")
+ if self.tuple_size > len(self.pool_items):
+ raise ValueError(
+ f"tuple_size ({self.tuple_size}) exceeds pool size ({len(self.pool_items)})"
+ )
+ if self._num_tuples is not None and self._num_tuples < 1:
+ raise ValueError(f"num_tuples must be >= 1, got {self._num_tuples}")
+
+ def _calculate_num_tuples(self) -> int:
+ """Auto-calculate number of tuples.
+
+ Uses Louviere's guideline: each item should appear at least
+ 2 * tuple_size times across all tuples.
+ """
+ min_appearances = self.min_item_appearances
+ if min_appearances is None:
+ min_appearances = 2 * self.tuple_size
+
+ pool_size = len(self.pool_items)
+ # Each tuple uses tuple_size items, so on average each item appears
+ # (num_tuples * tuple_size) / pool_size times.
+ # We need: (num_tuples * tuple_size) / pool_size >= min_appearances
+ num_tuples = math.ceil(pool_size * min_appearances / self.tuple_size)
+ return max(num_tuples, 1)
+
+ def generate(self) -> List[Dict[str, Any]]:
+ """Generate tuple instances from pool items.
+
+ Returns list of synthetic item dicts, each with:
+ - id_key: "bws_tuple_001"
+ - "_bws_items": list of {source_id, text, position} dicts
+ - "_bws_tuple_size": int
+ - text_key: "" (empty โ BWS JS handles display)
+ """
+ self.validate()
+
+ num_tuples = self._num_tuples if self._num_tuples else self._calculate_num_tuples()
+ rng = random.Random(self.seed)
+
+ logger.info(
+ f"Generating {num_tuples} BWS tuples of size {self.tuple_size} "
+ f"from pool of {len(self.pool_items)} items (seed={self.seed})"
+ )
+
+ tuples = []
+ for i in range(num_tuples):
+ sampled = rng.sample(self.pool_items, self.tuple_size)
+
+ bws_items = []
+ for pos_idx, item in enumerate(sampled):
+ bws_items.append({
+ "source_id": str(item[self.id_key]),
+ "text": str(item.get(self.text_key, "")),
+ "position": POSITION_LABELS[pos_idx],
+ })
+
+ tuple_id = f"bws_tuple_{i + 1:04d}"
+ tuple_instance = {
+ self.id_key: tuple_id,
+ self.text_key: "",
+ "_bws_items": bws_items,
+ "_bws_tuple_size": self.tuple_size,
+ }
+ tuples.append(tuple_instance)
+
+ # Log coverage statistics
+ item_counts = {}
+ for t in tuples:
+ for bws_item in t["_bws_items"]:
+ sid = bws_item["source_id"]
+ item_counts[sid] = item_counts.get(sid, 0) + 1
+
+ min_count = min(item_counts.values()) if item_counts else 0
+ max_count = max(item_counts.values()) if item_counts else 0
+ avg_count = sum(item_counts.values()) / len(item_counts) if item_counts else 0
+ logger.info(
+ f"BWS tuple coverage: min={min_count}, max={max_count}, avg={avg_count:.1f} "
+ f"appearances per item"
+ )
+
+ return tuples
diff --git a/potato/cases/__init__.py b/potato/cases/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..83395fa9d7d18074276a31233dcd4c539508a884
--- /dev/null
+++ b/potato/cases/__init__.py
@@ -0,0 +1,41 @@
+"""
+Cases (universal annotation feature).
+
+Groups instances into units of analysis (interview participant,
+respondent, document set) backed by the universal `project.sqlite`.
+Top-level `cases:` config; QDA auto-detects from
+`participant_id`/`respondent_id`/`case_id`. The crosstab reads
+case-level attributes when set.
+
+Layers:
+- `store` โ SQLite persistence (cases / case_attributes / case_documents).
+- `service` โ get-or-create, auto-detection, attribute accessors.
+"""
+
+from .service import (
+ DEFAULT_CASE_KEYS,
+ assign_instance,
+ attribute_for_instance,
+ attributes,
+ auto_detect,
+ case_for_instance,
+ cases_enabled,
+ get_or_create_case,
+ init_cases_from_config,
+ list_cases,
+ set_attribute,
+)
+
+__all__ = [
+ "DEFAULT_CASE_KEYS",
+ "get_or_create_case",
+ "list_cases",
+ "set_attribute",
+ "attributes",
+ "assign_instance",
+ "case_for_instance",
+ "attribute_for_instance",
+ "auto_detect",
+ "cases_enabled",
+ "init_cases_from_config",
+]
diff --git a/potato/cases/api.py b/potato/cases/api.py
new file mode 100644
index 0000000000000000000000000000000000000000..4b1f9e8b538c0816bc58ae94aa6dca987ef9f80f
--- /dev/null
+++ b/potato/cases/api.py
@@ -0,0 +1,78 @@
+"""
+Cases REST API (universal).
+
+Blueprint mounted at /api/cases. Read-only: list cases (with
+attributes) and resolve the case for an instance. Cases are
+created/assigned by auto-detection at server start, not via this API.
+Enabled when `cases` is configured or under QDA mode.
+"""
+
+from __future__ import annotations
+
+import logging
+from functools import wraps
+
+from flask import Blueprint, jsonify, request, session
+
+from potato.cases import (
+ attributes,
+ case_for_instance,
+ cases_enabled,
+ list_cases,
+)
+
+logger = logging.getLogger(__name__)
+
+cases_bp = Blueprint("cases", __name__, url_prefix="/api/cases")
+
+
+def _config() -> dict:
+ from potato.server_utils.config_module import config
+ return config
+
+
+def _ctx():
+ config = _config()
+ if not cases_enabled(config):
+ return None, ("disabled",)
+ if not session.get("username"):
+ return None, ("unauth",)
+ return {
+ "task_dir": config.get("task_dir", "."),
+ "project": config.get("annotation_task_name") or "default",
+ }, None
+
+
+def cases_view(view):
+ @wraps(view)
+ def wrapper(*args, **kwargs):
+ ctx, err = _ctx()
+ if err == ("disabled",):
+ return jsonify({
+ "error": "Cases are not enabled in this deployment.",
+ "hint": "Set cases.enabled: true (on by default in "
+ "qda_mode).",
+ }), 503
+ if err == ("unauth",):
+ return jsonify({"error": "Not authenticated"}), 401
+ return view(ctx, *args, **kwargs)
+ return wrapper
+
+
+@cases_bp.route("", methods=["GET"])
+@cases_view
+def list_all(ctx):
+ cases = list_cases(ctx["task_dir"], ctx["project"])
+ for c in cases:
+ c["attributes"] = attributes(ctx["task_dir"], c["id"])
+ return jsonify({"cases": cases})
+
+
+@cases_bp.route("/instance/", methods=["GET"])
+@cases_view
+def for_instance(ctx, instance_id):
+ case = case_for_instance(ctx["task_dir"], ctx["project"], instance_id)
+ if case is None:
+ return jsonify({"case": None})
+ case["attributes"] = attributes(ctx["task_dir"], case["id"])
+ return jsonify({"case": case})
diff --git a/potato/cases/service.py b/potato/cases/service.py
new file mode 100644
index 0000000000000000000000000000000000000000..0c9db192a211b27033f448ab34178f01967f7b7a
--- /dev/null
+++ b/potato/cases/service.py
@@ -0,0 +1,171 @@
+"""
+Cases service.
+
+get-or-create semantics, QDA auto-detection from item metadata, and the
+attribute accessors the crosstab uses. Single write path so detection
+and manual case creation share one audit trail.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional, Sequence
+
+from potato.cases import store
+
+logger = logging.getLogger(__name__)
+
+# Item-data keys QDA scans, in priority order, when no explicit
+# `cases.key` is configured.
+DEFAULT_CASE_KEYS = ("case_id", "participant_id", "respondent_id")
+
+
+def get_or_create_case(
+ task_dir: str, *, project: str, name: str,
+ created_by: str = "cases",
+) -> Dict[str, Any]:
+ existing = store.find_case(task_dir, project, name)
+ if existing is not None:
+ return existing
+ return store.insert_case(
+ task_dir, project=project, name=name, created_by=created_by)
+
+
+def list_cases(task_dir: str, project: str) -> List[Dict[str, Any]]:
+ return store.list_cases(task_dir, project)
+
+
+def set_attribute(
+ task_dir: str, case_id: str, key: str, value: Optional[str]
+) -> None:
+ store.set_attribute(task_dir, case_id, key, value)
+
+
+def attributes(task_dir: str, case_id: str) -> Dict[str, Any]:
+ return store.attributes(task_dir, case_id)
+
+
+def assign_instance(
+ task_dir: str, *, project: str, instance_id: str, case_id: str
+) -> None:
+ store.assign_instance(
+ task_dir, project=project, instance_id=instance_id,
+ case_id=case_id)
+
+
+def case_for_instance(
+ task_dir: str, project: str, instance_id: str
+) -> Optional[Dict[str, Any]]:
+ return store.case_for_instance(task_dir, project, instance_id)
+
+
+def attribute_for_instance(
+ task_dir: str, project: str, instance_id: str, key: str
+) -> Optional[str]:
+ """Resolve a case-level attribute for the instance's case. Used by
+ the crosstab so codes can be tabulated by participant-level
+ metadata that does not live on each instance."""
+ case = store.case_for_instance(task_dir, project, instance_id)
+ if case is None:
+ return None
+ return store.attributes(task_dir, case["id"]).get(key)
+
+
+def _detect_key(item: Dict[str, Any], keys: Sequence[str]) -> Optional[str]:
+ for k in keys:
+ v = item.get(k)
+ if v is not None and str(v).strip() != "":
+ return str(v)
+ return None
+
+
+def auto_detect(
+ task_dir: str,
+ *,
+ project: str,
+ items: Sequence[Dict[str, Any]],
+ case_key: Optional[str] = None,
+ attribute_keys: Optional[Sequence[str]] = None,
+) -> Dict[str, int]:
+ """Group items into cases by `case_key` (or the first present of
+ DEFAULT_CASE_KEYS) and lift `attribute_keys` onto the case. Returns
+ {"cases": n, "assigned": m}. Idempotent (get-or-create + upsert)."""
+ keys = (case_key,) if case_key else DEFAULT_CASE_KEYS
+ attr_keys = list(attribute_keys or [])
+ cases_seen: Dict[str, str] = {}
+ assigned = 0
+
+ for item in items:
+ if not isinstance(item, dict):
+ continue
+ name = _detect_key(item, keys)
+ if name is None:
+ continue
+ instance_id = str(
+ item.get("id") or item.get("instance_id") or "")
+ if not instance_id:
+ continue
+
+ cid = cases_seen.get(name)
+ if cid is None:
+ case = get_or_create_case(
+ task_dir, project=project, name=name,
+ created_by="cases-autodetect")
+ cid = case["id"]
+ cases_seen[name] = cid
+
+ store.assign_instance(
+ task_dir, project=project, instance_id=instance_id,
+ case_id=cid)
+ assigned += 1
+
+ for ak in attr_keys:
+ if ak in item and item[ak] is not None:
+ store.set_attribute(task_dir, cid, ak, str(item[ak]))
+
+ if cases_seen:
+ logger.info(
+ "Cases auto-detect: %d case(s), %d instance(s) assigned "
+ "for project %r", len(cases_seen), assigned, project)
+ return {"cases": len(cases_seen), "assigned": assigned}
+
+
+def cases_enabled(config: Dict[str, Any]) -> bool:
+ """Cases run when explicitly enabled, or implicitly under QDA mode
+ (unless `cases.enabled: false` opts out)."""
+ cases_cfg = config.get("cases") or {}
+ if cases_cfg.get("enabled") is True:
+ return True
+ if cases_cfg.get("enabled") is False:
+ return False
+ return bool((config.get("qda_mode") or {}).get("enabled"))
+
+
+def init_cases_from_config(config: Dict[str, Any]) -> Dict[str, int]:
+ """Server-start entry point: auto-detect cases from loaded items.
+ No-op (returns zeros) when cases are disabled or auto_detect is off.
+ """
+ if not cases_enabled(config):
+ return {"cases": 0, "assigned": 0}
+ cases_cfg = config.get("cases") or {}
+ if cases_cfg.get("auto_detect") is False:
+ return {"cases": 0, "assigned": 0}
+
+ task_dir = config.get("task_dir", ".")
+ project = config.get("annotation_task_name") or "default"
+ case_key = cases_cfg.get("key")
+ attribute_keys = cases_cfg.get("attributes") or []
+
+ from potato.item_state_management import get_item_state_manager
+ ism = get_item_state_manager()
+ items: List[Dict[str, Any]] = []
+ for iid in ism.get_instance_ids():
+ data = ism.get_item(iid).get_data()
+ if isinstance(data, dict):
+ row = dict(data)
+ row.setdefault("id", str(iid))
+ items.append(row)
+
+ return auto_detect(
+ task_dir, project=project, items=items,
+ case_key=case_key, attribute_keys=attribute_keys)
diff --git a/potato/cases/store.py b/potato/cases/store.py
new file mode 100644
index 0000000000000000000000000000000000000000..07041fa2f2f7c786905d4de59869dffce41e1bdc
--- /dev/null
+++ b/potato/cases/store.py
@@ -0,0 +1,161 @@
+"""
+Cases storage (universal).
+
+SQLite-backed CRUD over `cases`, `case_attributes`, and
+`case_documents` in `/project.sqlite`. A *case* groups
+instances that belong to the same unit of analysis (an interview
+participant, a respondent, a document set). Universal โ usable in
+standard annotation, solo mode, and QDA mode; QDA auto-detects cases
+from `participant_id`/`respondent_id`/`case_id` in the item data.
+
+No business rules here (the service layer owns get-or-create,
+auto-detection, and attribute lifting). One instance belongs to at most
+one case (PK on `project, instance_id`).
+"""
+
+from __future__ import annotations
+
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+from potato.persistence import Migration, get_db, register_migration
+
+_CASES_MIGRATION = Migration(
+ name="0001_cases",
+ sql="""
+ CREATE TABLE IF NOT EXISTS cases (
+ id TEXT PRIMARY KEY,
+ project TEXT NOT NULL,
+ name TEXT NOT NULL,
+ created_by TEXT NOT NULL,
+ created_at REAL NOT NULL,
+ updated_at REAL NOT NULL,
+ UNIQUE (project, name)
+ );
+ CREATE INDEX IF NOT EXISTS idx_cases_project ON cases (project);
+
+ CREATE TABLE IF NOT EXISTS case_attributes (
+ case_id TEXT NOT NULL,
+ key TEXT NOT NULL,
+ value TEXT,
+ PRIMARY KEY (case_id, key)
+ );
+
+ CREATE TABLE IF NOT EXISTS case_documents (
+ project TEXT NOT NULL,
+ instance_id TEXT NOT NULL,
+ case_id TEXT NOT NULL,
+ PRIMARY KEY (project, instance_id)
+ );
+ CREATE INDEX IF NOT EXISTS idx_case_docs_case
+ ON case_documents (case_id);
+ """,
+)
+
+register_migration(_CASES_MIGRATION)
+
+
+def _db(task_dir: str):
+ register_migration(_CASES_MIGRATION)
+ return get_db(task_dir)
+
+
+# ---- cases ---------------------------------------------------------------
+
+def insert_case(
+ task_dir: str, *, project: str, name: str, created_by: str,
+ case_id: Optional[str] = None,
+) -> Dict[str, Any]:
+ cid = case_id or uuid.uuid4().hex
+ now = time.time()
+ conn = _db(task_dir)
+ conn.execute(
+ """INSERT INTO cases
+ (id, project, name, created_by, created_at, updated_at)
+ VALUES (?, ?, ?, ?, ?, ?)""",
+ (cid, project, name, created_by, now, now),
+ )
+ conn.commit()
+ return get_case(task_dir, cid)
+
+
+def get_case(task_dir: str, case_id: str) -> Optional[Dict[str, Any]]:
+ row = _db(task_dir).execute(
+ "SELECT * FROM cases WHERE id = ?", (case_id,)
+ ).fetchone()
+ return dict(row) if row else None
+
+
+def find_case(
+ task_dir: str, project: str, name: str
+) -> Optional[Dict[str, Any]]:
+ row = _db(task_dir).execute(
+ "SELECT * FROM cases WHERE project = ? AND name = ?",
+ (project, name),
+ ).fetchone()
+ return dict(row) if row else None
+
+
+def list_cases(task_dir: str, project: str) -> List[Dict[str, Any]]:
+ rows = _db(task_dir).execute(
+ "SELECT * FROM cases WHERE project = ? ORDER BY name ASC",
+ (project,),
+ ).fetchall()
+ return [dict(r) for r in rows]
+
+
+# ---- attributes ----------------------------------------------------------
+
+def set_attribute(
+ task_dir: str, case_id: str, key: str, value: Optional[str]
+) -> None:
+ conn = _db(task_dir)
+ conn.execute(
+ """INSERT OR REPLACE INTO case_attributes (case_id, key, value)
+ VALUES (?, ?, ?)""",
+ (case_id, key, None if value is None else str(value)),
+ )
+ conn.commit()
+
+
+def attributes(task_dir: str, case_id: str) -> Dict[str, Any]:
+ rows = _db(task_dir).execute(
+ "SELECT key, value FROM case_attributes WHERE case_id = ?",
+ (case_id,),
+ ).fetchall()
+ return {r["key"]: r["value"] for r in rows}
+
+
+# ---- documents (instance <-> case) --------------------------------------
+
+def assign_instance(
+ task_dir: str, *, project: str, instance_id: str, case_id: str
+) -> None:
+ conn = _db(task_dir)
+ conn.execute(
+ """INSERT OR REPLACE INTO case_documents
+ (project, instance_id, case_id) VALUES (?, ?, ?)""",
+ (project, instance_id, case_id),
+ )
+ conn.commit()
+
+
+def case_for_instance(
+ task_dir: str, project: str, instance_id: str
+) -> Optional[Dict[str, Any]]:
+ row = _db(task_dir).execute(
+ """SELECT c.* FROM case_documents d
+ JOIN cases c ON c.id = d.case_id
+ WHERE d.project = ? AND d.instance_id = ?""",
+ (project, instance_id),
+ ).fetchone()
+ return dict(row) if row else None
+
+
+def instances_for_case(task_dir: str, case_id: str) -> List[str]:
+ rows = _db(task_dir).execute(
+ "SELECT instance_id FROM case_documents WHERE case_id = ?",
+ (case_id,),
+ ).fetchall()
+ return [r["instance_id"] for r in rows]
diff --git a/potato/chat_manager.py b/potato/chat_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..445f9ae87a52f9bb51f41bb73c1e8bb95279b672
--- /dev/null
+++ b/potato/chat_manager.py
@@ -0,0 +1,188 @@
+"""
+Chat Manager for LLM-based annotator assistance.
+
+Provides a singleton ChatManager that handles multi-turn conversations
+between annotators and an LLM, with context about the current annotation task.
+"""
+
+import logging
+import time
+from typing import Any, Dict, List, Optional
+
+from potato.ai.ai_endpoint import AIEndpointFactory, BaseAIEndpoint, AIEndpointConfigError
+
+logger = logging.getLogger(__name__)
+
+# Singleton instance
+_chat_manager: Optional["ChatManager"] = None
+
+DEFAULT_SYSTEM_PROMPT = """You are an annotation assistant for the task "{task_name}".
+
+Task description: {task_description}
+Available labels: {annotation_labels}
+
+The annotator is currently looking at this text:
+---
+{instance_text}
+---
+
+Help the annotator think through how to annotate this instance. You may:
+- Highlight relevant aspects of the text
+- Explain what to look for given the task description
+- Clarify label definitions if asked
+
+Do NOT tell the annotator which label to choose. Your role is to help them reason through the decision themselves."""
+
+
+class ChatManager:
+ """Manages LLM chat interactions for annotator assistance."""
+
+ def __init__(self, config: Dict[str, Any]):
+ chat_config = config.get("chat_support", {})
+ self.enabled = chat_config.get("enabled", False)
+
+ # UI settings
+ ui_config = chat_config.get("ui", {})
+ self.title = ui_config.get("title", "Ask AI")
+ self.placeholder = ui_config.get("placeholder", "Ask about this annotation...")
+ self.sidebar_width = ui_config.get("sidebar_width", 380)
+ self.max_history_per_instance = ui_config.get("max_history_per_instance", 50)
+
+ # System prompt template
+ prompt_config = chat_config.get("system_prompt", {})
+ self.system_prompt_template = prompt_config.get("template", DEFAULT_SYSTEM_PROMPT)
+
+ # Extract task-level info for system prompt
+ self.task_name = config.get("annotation_task_name", "Annotation Task")
+ self.task_description = config.get("annotation_task_description", "")
+
+ # Build label summary from annotation schemes
+ labels = []
+ for scheme in config.get("annotation_schemes", []):
+ name = scheme.get("name", "")
+ scheme_labels = scheme.get("labels", [])
+ if scheme_labels:
+ labels.append(f"{name}: {', '.join(str(l) for l in scheme_labels)}")
+ elif scheme.get("description"):
+ labels.append(f"{name}: {scheme['description']}")
+ self.annotation_labels = "; ".join(labels) if labels else "See task description"
+
+ # Create the LLM endpoint
+ self.endpoint: Optional[BaseAIEndpoint] = None
+ if self.enabled:
+ self._init_endpoint(config)
+
+ def _init_endpoint(self, config: Dict[str, Any]):
+ """Initialize the LLM endpoint from chat_support config."""
+ chat_config = config["chat_support"]
+
+ # Build a config dict that AIEndpointFactory expects
+ # (it looks for ai_support.enabled, ai_support.endpoint_type, ai_support.ai_config)
+ endpoint_config = {
+ "ai_support": {
+ "enabled": True,
+ "endpoint_type": chat_config.get("endpoint_type"),
+ "ai_config": chat_config.get("ai_config", {}),
+ }
+ }
+
+ try:
+ self.endpoint = AIEndpointFactory.create_endpoint(endpoint_config)
+ if self.endpoint:
+ logger.info(f"Chat endpoint initialized: {chat_config.get('endpoint_type')}")
+ else:
+ logger.error("Chat endpoint creation returned None")
+ self.enabled = False
+ except AIEndpointConfigError as e:
+ logger.error(f"Failed to initialize chat endpoint: {e}")
+ self.enabled = False
+
+ def build_system_prompt(self, instance_text: str, instance_id: str) -> str:
+ """Build the system prompt with current instance context."""
+ try:
+ return self.system_prompt_template.format(
+ task_name=self.task_name,
+ task_description=self.task_description,
+ annotation_labels=self.annotation_labels,
+ instance_text=instance_text,
+ instance_id=instance_id,
+ )
+ except KeyError as e:
+ logger.warning(f"System prompt template variable not found: {e}, using default")
+ return DEFAULT_SYSTEM_PROMPT.format(
+ task_name=self.task_name,
+ task_description=self.task_description,
+ annotation_labels=self.annotation_labels,
+ instance_text=instance_text,
+ instance_id=instance_id,
+ )
+
+ def send_message(
+ self,
+ user_message: str,
+ instance_text: str,
+ instance_id: str,
+ history: List[Dict[str, str]],
+ ) -> Dict[str, Any]:
+ """
+ Send a user message and get an LLM response.
+
+ Args:
+ user_message: The annotator's message
+ instance_text: Current instance text for context
+ instance_id: Current instance ID
+ history: Previous messages as list of {role, content} dicts
+
+ Returns:
+ Dict with 'content' (str) and 'response_time_ms' (int)
+ """
+ if not self.endpoint:
+ return {"content": "Chat support is not configured.", "response_time_ms": 0}
+
+ system_prompt = self.build_system_prompt(instance_text, instance_id)
+
+ # Build messages array: system + history + new user message
+ messages = [{"role": "system", "content": system_prompt}]
+ messages.extend(history)
+ messages.append({"role": "user", "content": user_message})
+
+ start_time = time.time()
+ try:
+ response_text = self.endpoint.chat_query(messages)
+ elapsed_ms = int((time.time() - start_time) * 1000)
+ return {"content": response_text, "response_time_ms": elapsed_ms}
+ except Exception as e:
+ elapsed_ms = int((time.time() - start_time) * 1000)
+ logger.error(f"Chat query failed: {e}")
+ return {
+ "content": "Sorry, I encountered an error. Please try again.",
+ "response_time_ms": elapsed_ms,
+ }
+
+ def get_ui_config(self) -> Dict[str, Any]:
+ """Return UI configuration for the frontend."""
+ return {
+ "enabled": self.enabled,
+ "title": self.title,
+ "placeholder": self.placeholder,
+ "sidebar_width": self.sidebar_width,
+ "max_history_per_instance": self.max_history_per_instance,
+ }
+
+
+def init_chat_manager(config: Dict[str, Any]) -> ChatManager:
+ """Initialize the global ChatManager singleton."""
+ global _chat_manager
+ _chat_manager = ChatManager(config)
+ return _chat_manager
+
+
+def get_chat_manager() -> Optional[ChatManager]:
+ """Get the ChatManager singleton. Returns None if not initialized."""
+ return _chat_manager
+
+
+def clear_chat_manager():
+ """Clear the ChatManager singleton. Used for testing."""
+ global _chat_manager
+ _chat_manager = None
diff --git a/potato/cli.py b/potato/cli.py
new file mode 100644
index 0000000000000000000000000000000000000000..8642dba87bf0b602f9b8e0dc63099b5839b6599b
--- /dev/null
+++ b/potato/cli.py
@@ -0,0 +1,33 @@
+#!/usr/bin/env python
+"""
+Command Line Interface for Potato Annotation Platform
+
+This module provides the main CLI entry point for running the Potato annotation server.
+It serves as a bridge between the command line and the Flask server application.
+
+The CLI can be invoked directly or through the potato command after installation.
+"""
+
+from potato.flask_server import main
+from potato import *
+
+def potato():
+ """
+ Main CLI entry point for the Potato annotation platform.
+
+ This function serves as the primary interface for starting the annotation server
+ from the command line. It delegates to the main() function in flask_server.py
+ which handles argument parsing, configuration loading, and server startup.
+
+ Side Effects:
+ - Initializes the Flask application
+ - Loads configuration from files
+ - Starts the web server on the configured port
+ - Sets up logging and error handling
+ """
+ main()
+
+
+if __name__ == '__main__':
+ # Direct script execution - start the Potato annotation server
+ potato()
\ No newline at end of file
diff --git a/potato/codebook/__init__.py b/potato/codebook/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d4e4a3c48de4569f6b4ebdb32758c2865d3a476
--- /dev/null
+++ b/potato/codebook/__init__.py
@@ -0,0 +1,84 @@
+"""
+Codebook (universal annotation feature).
+
+A mutable, optionally nested set of codes (labels) backed by the
+universal `project.sqlite`. Opt-in per scheme via
+`annotation_schemes[].codebook: true`; `codebook_mode`
+(fixed|extensible|open) governs whether annotators may add codes on the
+fly. Universal โ usable in standard annotation, solo mode (human + LLM
+co-edit), and QDA mode.
+
+Layers:
+- `store` โ SQLite persistence (codes + annotation_codes).
+- `codebook` โ read-model: tree + flat label list for the schema bridge.
+- `service` โ the single audited mutation path (create/rename/recolor/
+ move_under/delete + annotation links); fires change listeners (ICL
+ prompt-cache invalidation).
+"""
+
+# Import order matters: .service -> .store registers the 0001_codebook
+# migration (CREATE TABLE codes) which .revision's ALTER depends on.
+from .service import (
+ CodebookCycleError,
+ CodebookError,
+ CodeNotFound,
+ DuplicateCodeError,
+ apply_code,
+ clear_change_listeners,
+ codes_on,
+ create_code,
+ delete_code,
+ merge_codes,
+ move_under,
+ recolor_code,
+ register_change_listener,
+ remove_code,
+ rename_code,
+ split_code,
+)
+from .codebook import Codebook
+from .similar import derive_code_name, similar_code_names
+from . import revision
+from . import changelog
+from .changelog import propose_change
+from .revision import (
+ all_stale_instances,
+ codes_added_since,
+ current_revision,
+ instance_revision,
+ record_annotation,
+ stale_instances,
+ touch_instances,
+)
+
+__all__ = [
+ "Codebook",
+ "CodebookError",
+ "CodeNotFound",
+ "DuplicateCodeError",
+ "CodebookCycleError",
+ "create_code",
+ "rename_code",
+ "recolor_code",
+ "move_under",
+ "delete_code",
+ "merge_codes",
+ "split_code",
+ "touch_instances",
+ "apply_code",
+ "remove_code",
+ "codes_on",
+ "register_change_listener",
+ "clear_change_listeners",
+ "revision",
+ "current_revision",
+ "record_annotation",
+ "instance_revision",
+ "stale_instances",
+ "all_stale_instances",
+ "codes_added_since",
+ "derive_code_name",
+ "similar_code_names",
+ "changelog",
+ "propose_change",
+]
diff --git a/potato/codebook/api.py b/potato/codebook/api.py
new file mode 100644
index 0000000000000000000000000000000000000000..bbce8ca26cdac85560e7548f524e85ed8bcaee71
--- /dev/null
+++ b/potato/codebook/api.py
@@ -0,0 +1,511 @@
+"""
+Codebook REST API (universal).
+
+Blueprint mounted at /api/codebook. Read access is open to any
+authenticated annotator when the codebook is enabled; write access is
+governed by the effective ``codebook_mode``:
+
+- ``fixed`` โ no API mutations (config/CLI only)
+- ``extensible`` โ authenticated users may *add* codes
+- ``open`` โ authenticated users may add / rename / recolor /
+ move / delete
+
+Adjudicators are privileged and may always mutate (any mode but
+``fixed``-locked still applies โ fixed means locked for everyone here;
+use ``potato codebook`` / config to change a fixed codebook).
+
+The single mutation path is the codebook service, so human and LLM
+edits (solo mode) share one audit trail (``created_by``).
+"""
+
+from __future__ import annotations
+
+import logging
+from functools import wraps
+
+from flask import Blueprint, jsonify, request, session
+
+from potato.codebook import (
+ CodebookError,
+ CodeNotFound,
+ DuplicateCodeError,
+ Codebook,
+ codes_added_since,
+ create_code,
+ current_revision,
+ delete_code,
+ instance_revision,
+ move_under,
+ recolor_code,
+ rename_code,
+ stale_instances,
+)
+from potato.codebook.store import ROOT
+
+logger = logging.getLogger(__name__)
+
+codebook_bp = Blueprint("codebook", __name__, url_prefix="/api/codebook")
+
+
+def _config() -> dict:
+ from potato.server_utils.config_module import config
+ return config
+
+
+def codebook_enabled(config: dict) -> bool:
+ """On when a codebook scheme/config is present, or under qda/solo."""
+ if config.get("codebook_mode") is not None:
+ return True
+ cb = config.get("codebook")
+ if isinstance(cb, dict) and cb.get("enabled") is not None:
+ return bool(cb.get("enabled"))
+ for s in config.get("annotation_schemes") or []:
+ if isinstance(s, dict) and s.get("codebook"):
+ return True
+ return bool(
+ (config.get("qda_mode") or {}).get("enabled")
+ or (config.get("solo_mode") or {}).get("enabled")
+ )
+
+
+def _is_privileged(username: str) -> bool:
+ try:
+ from potato.adjudication import get_adjudication_manager
+ adj = get_adjudication_manager()
+ return bool(adj and adj.is_adjudicator(username))
+ except Exception:
+ return False
+
+
+def _ctx():
+ config = _config()
+ if not codebook_enabled(config):
+ return None, ("disabled",)
+ username = session.get("username")
+ if not username:
+ return None, ("unauth",)
+ from potato.server_utils.config_module import get_codebook_mode
+ return {
+ "task_dir": config.get("task_dir", "."),
+ "project": config.get("annotation_task_name") or "default",
+ "username": username,
+ "privileged": _is_privileged(username),
+ "mode": get_codebook_mode(config),
+ }, None
+
+
+def codebook_view(view):
+ @wraps(view)
+ def wrapper(*args, **kwargs):
+ ctx, err = _ctx()
+ if err == ("disabled",):
+ return jsonify({
+ "error": "Codebook is not enabled in this deployment.",
+ "hint": "Add `codebook: true` to a scheme, or set "
+ "codebook_mode (on by default in qda/solo).",
+ }), 503
+ if err == ("unauth",):
+ return jsonify({"error": "Not authenticated"}), 401
+ return view(ctx, *args, **kwargs)
+ return wrapper
+
+
+def _can_mutate(ctx, *, need_open: bool) -> bool:
+ if ctx["mode"] == "fixed":
+ return False
+ if ctx["privileged"]:
+ return True
+ if ctx["mode"] == "open":
+ return True
+ # extensible: add allowed, structural edits are not
+ return not need_open
+
+
+def _handle(fn):
+ try:
+ return fn()
+ except CodeNotFound as e:
+ return jsonify({"error": str(e)}), 404
+ except DuplicateCodeError as e:
+ return jsonify({"error": str(e)}), 409
+ except CodebookError as e:
+ return jsonify({"error": str(e)}), 400
+
+
+def _codebook_scheme_names() -> list:
+ """Names of schemes opted into the codebook โ the forms the tray
+ refreshes in place after an add."""
+ cfg = _config()
+ return [s.get("name") for s in (cfg.get("annotation_schemes") or [])
+ if isinstance(s, dict) and s.get("codebook") and s.get("name")]
+
+
+def _instance_index_map() -> dict:
+ """instance_id -> 0-based position, so the review worklist can jump
+ via the existing index-based navigateToInstance()."""
+ try:
+ from potato.item_state_management import get_item_state_manager
+ ids = get_item_state_manager().get_instance_ids()
+ return {str(iid): i for i, iid in enumerate(ids)}
+ except Exception:
+ return {}
+
+
+@codebook_bp.route("", methods=["GET"])
+@codebook_view
+def get_codebook(ctx):
+ cb = Codebook.load(ctx["task_dir"], ctx["project"])
+ return jsonify({
+ "mode": ctx["mode"],
+ "labels": cb.labels(),
+ "tree": cb.as_tree(),
+ "revision": current_revision(ctx["task_dir"], ctx["project"]),
+ "schemes": _codebook_scheme_names(),
+ "invivo_key": str(
+ _config().get("codebook_invivo_key") or "i")[:1].lower(),
+ "can_add": _can_mutate(ctx, need_open=False),
+ "can_edit": _can_mutate(ctx, need_open=True),
+ })
+
+
+@codebook_bp.route("/version", methods=["GET"])
+@codebook_view
+def version(ctx):
+ """Lightweight revision poll. The client checks this on each
+ navigation and only re-fetches the full codebook (GET /api/codebook)
+ when the revision has moved โ instead of downloading the whole tree
+ on every page load."""
+ return jsonify({
+ "revision": current_revision(ctx["task_dir"], ctx["project"])})
+
+
+@codebook_bp.route("/provenance", methods=["GET"])
+@codebook_view
+def provenance(ctx):
+ """Is one instance stale for the current annotator (labeled before
+ later code additions)? Powers the dismissible revisit banner."""
+ instance_id = request.args.get("instance_id")
+ if not instance_id:
+ return jsonify({"error": "instance_id is required"}), 400
+ cur = current_revision(ctx["task_dir"], ctx["project"])
+ ann = instance_revision(
+ ctx["task_dir"], ctx["project"], instance_id, ctx["username"])
+ added = ([] if ann is None or ann >= cur
+ else codes_added_since(ctx["task_dir"], ctx["project"], ann))
+ return jsonify({
+ "instance_id": instance_id,
+ "annotated_revision": ann,
+ "current_revision": cur,
+ "stale": bool(added),
+ "codes_added_since": added,
+ })
+
+
+@codebook_bp.route("/stale", methods=["GET"])
+@codebook_view
+def stale(ctx):
+ """The current annotator's review worklist: their instances labeled
+ under an older revision, each with the codes added since."""
+ items = stale_instances(
+ ctx["task_dir"], ctx["project"], ctx["username"])
+ idx = _instance_index_map()
+ for it in items:
+ it["index"] = idx.get(str(it["instance_id"]))
+ return jsonify({"stale": items, "count": len(items)})
+
+
+def _admin_or_adjudicator() -> bool:
+ try:
+ from potato.admin import admin_dashboard
+ if admin_dashboard.check_admin_access():
+ return True
+ except Exception:
+ pass
+ username = session.get("username")
+ if username:
+ try:
+ from potato.adjudication import get_adjudication_manager
+ adj = get_adjudication_manager()
+ if adj and adj.is_adjudicator(username):
+ return True
+ except Exception:
+ pass
+ return False
+
+
+@codebook_bp.route("/admin/stale", methods=["GET"])
+def admin_stale():
+ """Project-wide stale instances (all users) for oversight. Admin
+ API key or adjudicator only."""
+ from potato.server_utils.config_module import config as _cfg
+ if not codebook_enabled(_cfg):
+ return jsonify({"error": "Codebook not enabled"}), 503
+ if not _admin_or_adjudicator():
+ return jsonify({
+ "error": "Admin or adjudicator access required"}), 403
+ from potato.codebook.revision import all_stale_instances
+ task_dir = _cfg.get("task_dir", ".")
+ project = _cfg.get("annotation_task_name") or "default"
+ items = all_stale_instances(task_dir, project)
+ return jsonify({"stale": items, "count": len(items)})
+
+
+def _admin_ctx():
+ """(task_dir, project, username, None) or (None,None,None, resp).
+ Mirrors admin_stale's gate for the Phase 2 (C) retroactive ops."""
+ from potato.server_utils.config_module import config as _cfg
+ if not codebook_enabled(_cfg):
+ return None, None, None, (
+ jsonify({"error": "Codebook not enabled"}), 503)
+ if not _admin_or_adjudicator():
+ return None, None, None, (
+ jsonify({"error": "Admin or adjudicator access required"}),
+ 403)
+ return (_cfg.get("task_dir", "."),
+ _cfg.get("annotation_task_name") or "default",
+ session.get("username") or "admin", None)
+
+
+@codebook_bp.route("/admin/merge", methods=["POST"])
+def admin_merge():
+ """Fold src into dst retroactively (append-only). Admin only."""
+ td, project, user, err = _admin_ctx()
+ if err:
+ return err
+ from potato.codebook import merge_codes
+ data = request.get_json(silent=True) or {}
+ src_id = (data.get("src_id") or "").strip()
+ dst_id = (data.get("dst_id") or "").strip()
+ if not src_id or not dst_id:
+ return jsonify({"error": "src_id and dst_id are required"}), 400
+ return _handle(lambda: jsonify(merge_codes(
+ td, project=project, src_id=src_id, dst_id=dst_id,
+ actor=user, actor_kind="human")))
+
+
+@codebook_bp.route("/admin/split", methods=["POST"])
+def admin_split():
+ """Split a code by annotator retroactively. Admin only."""
+ td, project, user, err = _admin_ctx()
+ if err:
+ return err
+ from potato.codebook import split_code
+ data = request.get_json(silent=True) or {}
+ src_id = (data.get("src_id") or "").strip()
+ annotator = (data.get("annotator") or "").strip()
+ if not src_id or not annotator:
+ return jsonify({
+ "error": "src_id and annotator are required"}), 400
+ return _handle(lambda: jsonify(split_code(
+ td, project=project, src_id=src_id, annotator=annotator,
+ new_name=(data.get("new_name") or "").strip() or None,
+ target_id=(data.get("target_id") or "").strip() or None,
+ actor=user, actor_kind="human")))
+
+
+@codebook_bp.route("/admin/changes", methods=["GET"])
+def admin_changes():
+ """Full change-log for the before->after delta view. Admin only."""
+ td, project, _user, err = _admin_ctx()
+ if err:
+ return err
+ from potato.codebook import changelog
+ rows = changelog.all_changes(td, project)
+ return jsonify({"changes": rows, "count": len(rows)})
+
+
+@codebook_bp.route("/proposals", methods=["POST"])
+@codebook_view
+def submit_proposal(ctx):
+ """Producer contract: a model/agent stages a codebook edit for human
+ confirmation. `actor_kind=="model"` is the machine path (no admin
+ gate โ it only QUEUES; nothing changes until an admin confirms).
+ A human-submitted proposal still requires edit rights."""
+ data = request.get_json(silent=True) or {}
+ op = (data.get("op") or "").strip()
+ payload = data.get("payload") or {}
+ actor_kind = (data.get("actor_kind") or "model").strip()
+ if op not in ("merge", "split", "rename", "recolor", "move",
+ "delete"):
+ return jsonify({"error": f"unsupported op {op!r}"}), 400
+ if actor_kind != "model" and not _can_mutate(ctx, need_open=True):
+ return jsonify({
+ "error": "Proposing edits requires edit rights"}), 403
+ from potato.codebook import changelog
+ prop = changelog.record_proposal(
+ task_dir=ctx["task_dir"], project=ctx["project"], op=op,
+ payload=payload, actor=ctx["username"], actor_kind=actor_kind)
+ return jsonify({"proposal": prop}), 201
+
+
+@codebook_bp.route("/admin/proposals", methods=["GET"])
+def admin_list_proposals():
+ td, project, _user, err = _admin_ctx()
+ if err:
+ return err
+ from potato.codebook import changelog
+ items = changelog.list_proposals(td, project, status="pending")
+ return jsonify({"proposals": items, "count": len(items)})
+
+
+def _apply_proposed(td, project, op, payload, actor):
+ """Dispatch a confirmed proposal through the audited service path."""
+ from potato.codebook import (
+ merge_codes, split_code, rename_code, recolor_code,
+ move_under, delete_code)
+ if op == "merge":
+ return merge_codes(
+ td, project=project, src_id=payload["src_id"],
+ dst_id=payload["dst_id"], actor=actor, actor_kind="model")
+ if op == "split":
+ return split_code(
+ td, project=project, src_id=payload["src_id"],
+ annotator=payload["annotator"],
+ new_name=payload.get("new_name"),
+ target_id=payload.get("target_id"),
+ actor=actor, actor_kind="model")
+ if op == "rename":
+ return rename_code(
+ td, payload["code_id"], new_name=payload["new_name"],
+ project=project, actor=actor, actor_kind="model")
+ if op == "recolor":
+ return recolor_code(
+ td, payload["code_id"], color=payload["color"],
+ project=project, actor=actor, actor_kind="model")
+ if op == "move":
+ return move_under(
+ td, payload["code_id"],
+ new_parent_id=payload.get("parent_id") or "",
+ project=project, actor=actor, actor_kind="model")
+ if op == "delete":
+ return delete_code(
+ td, payload["code_id"], project=project,
+ actor=actor, actor_kind="model")
+ raise CodebookError(f"unsupported op {op!r}")
+
+
+@codebook_bp.route("/admin/proposals//confirm", methods=["POST"])
+def admin_confirm_proposal(pid):
+ td, project, user, err = _admin_ctx()
+ if err:
+ return err
+ from potato.codebook import changelog
+ prop = changelog.get_proposal(td, pid)
+ if not prop or prop["project"] != project:
+ return jsonify({"error": "proposal not found"}), 404
+ if prop["status"] != "pending":
+ return jsonify({
+ "error": f"proposal already {prop['status']}"}), 409
+
+ def _do():
+ result = _apply_proposed(
+ td, project, prop["op"], prop["payload"], user)
+ cid = changelog.log_change(
+ td, project=project, op="llm_confirmed",
+ old_value=prop["op"], new_value=str(result),
+ actor=user, actor_kind="model",
+ revision=current_revision(td, project))
+ changelog.set_proposal_status(
+ td, pid, status="confirmed", decided_by=user,
+ change_id=result.get("change_id") or cid)
+ return jsonify({"confirmed": True, "result": result})
+
+ return _handle(_do)
+
+
+@codebook_bp.route("/admin/proposals//reject", methods=["POST"])
+def admin_reject_proposal(pid):
+ td, project, user, err = _admin_ctx()
+ if err:
+ return err
+ from potato.codebook import changelog
+ prop = changelog.get_proposal(td, pid)
+ if not prop or prop["project"] != project:
+ return jsonify({"error": "proposal not found"}), 404
+ if prop["status"] != "pending":
+ return jsonify({
+ "error": f"proposal already {prop['status']}"}), 409
+ cid = changelog.log_change(
+ td, project=project, op="llm_rejected",
+ old_value=prop["op"], new_value=None, actor=user,
+ actor_kind="model", revision=0)
+ changelog.set_proposal_status(
+ td, pid, status="rejected", decided_by=user, change_id=cid)
+ return jsonify({"rejected": True})
+
+
+@codebook_bp.route("/similar", methods=["GET"])
+@codebook_view
+def similar(ctx):
+ """Soft suggest-on-create: existing codes that closely match a
+ proposed name (Phase 2 #1). Read-only โ drives a non-blocking
+ "Use ยซXยป?" prompt before the in-vivo / on-the-fly add commits."""
+ from potato.codebook.similar import similar_code_names
+ name = (request.args.get("name") or "").strip()
+ if not name:
+ return jsonify({"name": name, "matches": []})
+ cb = Codebook.load(ctx["task_dir"], ctx["project"])
+ return jsonify({
+ "name": name,
+ "matches": similar_code_names(cb.labels(), name),
+ })
+
+
+@codebook_bp.route("", methods=["POST"])
+@codebook_view
+def add_code(ctx):
+ if not _can_mutate(ctx, need_open=False):
+ return jsonify({
+ "error": f"Adding codes is not allowed (codebook_mode="
+ f"{ctx['mode']})."}), 403
+ data = request.get_json(silent=True) or {}
+ name = (data.get("name") or "").strip()
+ if not name:
+ return jsonify({"error": "name is required"}), 400
+ return _handle(lambda: jsonify({"code": create_code(
+ ctx["task_dir"], project=ctx["project"], name=name,
+ created_by=ctx["username"], color=data.get("color"),
+ parent_id=data.get("parent_id") or ROOT,
+ )}))
+
+
+@codebook_bp.route("/", methods=["PATCH"])
+@codebook_view
+def edit_code(ctx, code_id):
+ if not _can_mutate(ctx, need_open=True):
+ return jsonify({
+ "error": f"Editing codes requires codebook_mode=open "
+ f"(current: {ctx['mode']})."}), 403
+ data = request.get_json(silent=True) or {}
+
+ def _do():
+ result = None
+ if "name" in data:
+ result = rename_code(
+ ctx["task_dir"], code_id,
+ new_name=data["name"], project=ctx["project"])
+ if "color" in data:
+ result = recolor_code(
+ ctx["task_dir"], code_id,
+ color=data["color"], project=ctx["project"])
+ if "parent_id" in data:
+ result = move_under(
+ ctx["task_dir"], code_id,
+ new_parent_id=data["parent_id"] or ROOT,
+ project=ctx["project"])
+ if result is None:
+ return jsonify({"error": "nothing to update"}), 400
+ return jsonify({"code": result})
+
+ return _handle(_do)
+
+
+@codebook_bp.route("/", methods=["DELETE"])
+@codebook_view
+def remove_code(ctx, code_id):
+ if not _can_mutate(ctx, need_open=True):
+ return jsonify({
+ "error": f"Deleting codes requires codebook_mode=open "
+ f"(current: {ctx['mode']})."}), 403
+ return _handle(lambda: jsonify({"deleted": delete_code(
+ ctx["task_dir"], code_id, project=ctx["project"])}))
diff --git a/potato/codebook/changelog.py b/potato/codebook/changelog.py
new file mode 100644
index 0000000000000000000000000000000000000000..a785c287576f605510261d49d10ff232b43e3e51
--- /dev/null
+++ b/potato/codebook/changelog.py
@@ -0,0 +1,232 @@
+"""
+Codebook change-provenance overlay (Phase 2 C).
+
+A **separate** audit/overlay layer for retroactive codebook edits
+(merge / split / rename / recolor / move / delete) and the LLM
+propose -> human-confirm flow. It is deliberately NOT part of the
+`codes` records: `Codebook.labels()` / `as_tree()` feed the ICL prompt
+verbatim, so authorship / change history must never join into them
+(open-question #2 resolution).
+
+Two tables (own migration, universal project.sqlite):
+- ``codebook_change`` โ append-only event log: every retroactive op,
+ who/what/old->new/when, with ``actor_kind`` (human|model). Generalises
+ the additive-only ``revision.codes_added_since`` so the review
+ worklist / banner can say "X merged into Y", not just "N codes added".
+- ``codebook_proposal`` โ pending model-proposed edits awaiting human
+ confirmation (status pending|confirmed|rejected).
+
+Plus the temporal columns that make retroactive edits **append-only**:
+- ``annotation_codes.invalidated_at`` / ``invalidated_by_change`` โ
+ a superseded link is marked, never DELETEd (NULL = live). NOTE:
+ ``started_at``/``ended_at`` already mean *elapsed time on a span /
+ agentic trace* โ they are NOT validity and must not be reused.
+- ``codes.archived_at`` โ a merged-away source code is archived (leaves
+ the palette + ICL prompt) but its row and history survive.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+from potato.persistence import Migration, get_db, register_migration
+
+_CHANGE_MIGRATION = Migration(
+ name="0003_codebook_change_provenance",
+ sql="""
+ ALTER TABLE annotation_codes ADD COLUMN invalidated_at REAL;
+ ALTER TABLE annotation_codes ADD COLUMN invalidated_by_change TEXT;
+ ALTER TABLE codes ADD COLUMN archived_at REAL;
+
+ CREATE TABLE IF NOT EXISTS codebook_change (
+ id TEXT PRIMARY KEY,
+ project TEXT NOT NULL,
+ code_id TEXT,
+ related_code_id TEXT,
+ op TEXT NOT NULL,
+ old_value TEXT,
+ new_value TEXT,
+ actor TEXT NOT NULL,
+ actor_kind TEXT NOT NULL DEFAULT 'human',
+ created_at REAL NOT NULL,
+ revision INTEGER NOT NULL DEFAULT 0
+ );
+ CREATE INDEX IF NOT EXISTS idx_cbchange_proj
+ ON codebook_change (project, created_at);
+
+ CREATE TABLE IF NOT EXISTS codebook_proposal (
+ id TEXT PRIMARY KEY,
+ project TEXT NOT NULL,
+ op TEXT NOT NULL,
+ payload TEXT NOT NULL,
+ status TEXT NOT NULL DEFAULT 'pending',
+ actor TEXT NOT NULL,
+ actor_kind TEXT NOT NULL DEFAULT 'model',
+ created_at REAL NOT NULL,
+ decided_by TEXT,
+ decided_at REAL,
+ change_id TEXT
+ );
+ CREATE INDEX IF NOT EXISTS idx_cbproposal_proj
+ ON codebook_proposal (project, status, created_at);
+ """,
+)
+
+# Defensive ordering: 0003 ALTERs `annotation_codes` and `codes`, so the
+# 0001 CREATE and both 0002 ALTERs must register first regardless of
+# import path (mirrors revision.py โ import the Migration objects
+# directly, not via package side effects, to avoid circular imports).
+from potato.codebook.store import _CODEBOOK_MIGRATION as _CB_MIG
+from potato.codebook.revision import (
+ _REVISION_MIGRATION as _REV_MIG,
+ _CODES_REV_MIGRATION as _CODES_REV_MIG,
+)
+
+register_migration(_CB_MIG)
+register_migration(_REV_MIG)
+register_migration(_CODES_REV_MIG)
+register_migration(_CHANGE_MIGRATION)
+
+
+def _db(task_dir: str):
+ register_migration(_CHANGE_MIGRATION)
+ return get_db(task_dir)
+
+
+# ---- change log ----------------------------------------------------------
+
+def log_change(
+ task_dir: str, *, project: str, op: str, actor: str,
+ actor_kind: str = "human", code_id: Optional[str] = None,
+ related_code_id: Optional[str] = None,
+ old_value: Optional[str] = None, new_value: Optional[str] = None,
+ revision: int = 0,
+) -> str:
+ """Append one immutable change-log row; return its id."""
+ cid = uuid.uuid4().hex
+ conn = _db(task_dir)
+ conn.execute(
+ """INSERT INTO codebook_change
+ (id, project, code_id, related_code_id, op, old_value,
+ new_value, actor, actor_kind, created_at, revision)
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+ (cid, project, code_id, related_code_id, op, old_value,
+ new_value, actor, actor_kind, time.time(), revision),
+ )
+ conn.commit()
+ return cid
+
+
+def changes_since(
+ task_dir: str, project: str, revision: int
+) -> List[Dict[str, Any]]:
+ """Change-log rows recorded after `revision` โ the non-additive
+ counterpart to revision.codes_added_since, for the worklist/banner
+ and the before->after delta view."""
+ rows = _db(task_dir).execute(
+ """SELECT id, code_id, related_code_id, op, old_value, new_value,
+ actor, actor_kind, created_at, revision
+ FROM codebook_change
+ WHERE project = ? AND revision > ?
+ ORDER BY created_at ASC""",
+ (project, revision),
+ ).fetchall()
+ return [dict(r) for r in rows]
+
+
+def all_changes(
+ task_dir: str, project: str
+) -> List[Dict[str, Any]]:
+ rows = _db(task_dir).execute(
+ """SELECT id, code_id, related_code_id, op, old_value, new_value,
+ actor, actor_kind, created_at, revision
+ FROM codebook_change
+ WHERE project = ?
+ ORDER BY created_at ASC""",
+ (project,),
+ ).fetchall()
+ return [dict(r) for r in rows]
+
+
+# ---- proposals (model -> human-confirm) ---------------------------------
+
+def record_proposal(
+ task_dir: str, *, project: str, op: str, payload: Dict[str, Any],
+ actor: str, actor_kind: str = "model",
+) -> Dict[str, Any]:
+ pid = uuid.uuid4().hex
+ now = time.time()
+ conn = _db(task_dir)
+ conn.execute(
+ """INSERT INTO codebook_proposal
+ (id, project, op, payload, status, actor, actor_kind,
+ created_at)
+ VALUES (?, ?, ?, ?, 'pending', ?, ?, ?)""",
+ (pid, project, op, json.dumps(payload), actor, actor_kind, now),
+ )
+ conn.commit()
+ return get_proposal(task_dir, pid)
+
+
+def get_proposal(task_dir: str, proposal_id: str) -> Optional[Dict]:
+ row = _db(task_dir).execute(
+ "SELECT * FROM codebook_proposal WHERE id = ?", (proposal_id,)
+ ).fetchone()
+ if not row:
+ return None
+ d = dict(row)
+ try:
+ d["payload"] = json.loads(d["payload"])
+ except Exception:
+ pass
+ return d
+
+
+def list_proposals(
+ task_dir: str, project: str, status: str = "pending"
+) -> List[Dict[str, Any]]:
+ rows = _db(task_dir).execute(
+ """SELECT * FROM codebook_proposal
+ WHERE project = ? AND status = ?
+ ORDER BY created_at ASC""",
+ (project, status),
+ ).fetchall()
+ out = []
+ for r in rows:
+ d = dict(r)
+ try:
+ d["payload"] = json.loads(d["payload"])
+ except Exception:
+ pass
+ out.append(d)
+ return out
+
+
+def set_proposal_status(
+ task_dir: str, proposal_id: str, *, status: str,
+ decided_by: str, change_id: Optional[str] = None,
+) -> bool:
+ conn = _db(task_dir)
+ cur = conn.execute(
+ """UPDATE codebook_proposal
+ SET status = ?, decided_by = ?, decided_at = ?, change_id = ?
+ WHERE id = ? AND status = 'pending'""",
+ (status, decided_by, time.time(), change_id, proposal_id),
+ )
+ conn.commit()
+ return cur.rowcount > 0
+
+
+def propose_change(
+ task_dir: str, *, project: str, op: str, payload: Dict[str, Any],
+ actor: str, actor_kind: str = "model",
+) -> Dict[str, Any]:
+ """Thin entry point for in-process (e.g. solo-mode LLM) producers
+ that aren't going through the HTTP API. The contract is identical to
+ POST /api/codebook/proposals."""
+ return record_proposal(
+ task_dir, project=project, op=op, payload=payload,
+ actor=actor, actor_kind=actor_kind)
diff --git a/potato/codebook/codebook.py b/potato/codebook/codebook.py
new file mode 100644
index 0000000000000000000000000000000000000000..0a9b788d144576de4cd56b22c1aa6d585c2a0a4b
--- /dev/null
+++ b/potato/codebook/codebook.py
@@ -0,0 +1,89 @@
+"""
+Codebook read-model.
+
+A thin, read-only view over the `codes` table for one project: builds
+the code tree and the flat label list the schema-loader bridge needs.
+Mutations go through `service.py` (the single audited write path).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from potato.codebook import store
+
+
+class Codebook:
+ """In-memory snapshot of a project's codebook.
+
+ Construct via `Codebook.load(task_dir, project)`. Cheap to rebuild;
+ callers reload after a mutation rather than mutating the snapshot.
+ """
+
+ def __init__(self, project: str, codes: List[Dict[str, Any]]):
+ self.project = project
+ self._codes = codes
+ self._by_id: Dict[str, Dict[str, Any]] = {c["id"]: c for c in codes}
+
+ @classmethod
+ def load(cls, task_dir: str, project: str) -> "Codebook":
+ # Archived codes (e.g. merged away in Phase 2 C) must not reach
+ # the label list / tree โ those feed the ICL prompt and the live
+ # forms. `.get` keeps this tolerant of pre-0003 schemas.
+ codes = [c for c in store.list_codes(task_dir, project)
+ if not c.get("archived_at")]
+ return cls(project, codes)
+
+ def __len__(self) -> int:
+ return len(self._codes)
+
+ def is_empty(self) -> bool:
+ return not self._codes
+
+ def get(self, code_id: str) -> Optional[Dict[str, Any]]:
+ return self._by_id.get(code_id)
+
+ def children(self, parent_id: str = store.ROOT) -> List[Dict[str, Any]]:
+ kids = [c for c in self._codes if c["parent_id"] == parent_id]
+ kids.sort(key=lambda c: (c["sort_order"], c["name"]))
+ return kids
+
+ def labels(self) -> List[str]:
+ """Flat list of code names in tree order โ the legacy label list
+ the radio/multiselect/span loaders consume when a scheme opts in
+ via `codebook: true`."""
+ out: List[str] = []
+
+ def walk(parent: str) -> None:
+ for c in self.children(parent):
+ out.append(c["name"])
+ walk(c["id"])
+
+ walk(store.ROOT)
+ return out
+
+ def label_to_id(self) -> Dict[str, str]:
+ """Map code name -> code id (first occurrence in tree order).
+ Lets the annotation pipeline store a parallel `code_id` while
+ keeping the legacy `label` string."""
+ mapping: Dict[str, str] = {}
+ for name in self.labels():
+ if name not in mapping:
+ match = next(
+ (c for c in self._codes if c["name"] == name), None)
+ if match:
+ mapping[name] = match["id"]
+ return mapping
+
+ def as_tree(self) -> List[Dict[str, Any]]:
+ """Nested [{id,name,color,children:[...]}] for the codebook UI."""
+
+ def node(c: Dict[str, Any]) -> Dict[str, Any]:
+ return {
+ "id": c["id"],
+ "name": c["name"],
+ "color": c["color"],
+ "children": [node(k) for k in self.children(c["id"])],
+ }
+
+ return [node(c) for c in self.children(store.ROOT)]
diff --git a/potato/codebook/revision.py b/potato/codebook/revision.py
new file mode 100644
index 0000000000000000000000000000000000000000..5b857da572dba2552a4651c6e982f04d8068a048
--- /dev/null
+++ b/potato/codebook/revision.py
@@ -0,0 +1,223 @@
+"""
+Codebook revision provenance.
+
+A per-project monotonic ``codebook_revision`` counter, bumped on **any**
+codebook change (create / rename / recolor / move / delete) โ all
+codebook edits are server-persisted and every one advances the
+revision. Every saved annotation is stamped with the revision in
+effect, so analysts can condition on the codebook state an annotation
+was made against, and the UI can softly flag instances labeled under an
+older codebook revision.
+
+`codes.created_revision` records the revision a code first appeared in,
+so the review worklist can show *which* codes were added since a given
+instance was labeled (precision: a niche new code only resurfaces
+instances that predate it).
+
+Tables (own migrations, universal project.sqlite):
+- ``codebook_revision(project PK, revision, updated_at)``
+- ``annotation_provenance(project, instance_id, username, revision,
+ updated_at)`` โ PK(project, instance_id, username)
+- ``codes.created_revision`` (added via ALTER; default 0 = pre-feature)
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Dict, List, Optional
+
+from potato.persistence import Migration, get_db, register_migration
+
+_REVISION_MIGRATION = Migration(
+ name="0002_codebook_revision",
+ sql="""
+ CREATE TABLE IF NOT EXISTS codebook_revision (
+ project TEXT PRIMARY KEY,
+ revision INTEGER NOT NULL DEFAULT 0,
+ updated_at REAL NOT NULL
+ );
+ CREATE TABLE IF NOT EXISTS annotation_provenance (
+ project TEXT NOT NULL,
+ instance_id TEXT NOT NULL,
+ username TEXT NOT NULL,
+ revision INTEGER NOT NULL,
+ updated_at REAL NOT NULL,
+ PRIMARY KEY (project, instance_id, username)
+ );
+ CREATE INDEX IF NOT EXISTS idx_provenance_stale
+ ON annotation_provenance (project, username, revision);
+ """,
+)
+
+# Separate migration: add created_revision to the codes table (created
+# by store.py's 0001_codebook, which registers first at import).
+_CODES_REV_MIGRATION = Migration(
+ name="0002_codes_created_revision",
+ sql="""
+ ALTER TABLE codes ADD COLUMN created_revision INTEGER NOT NULL
+ DEFAULT 0;
+ """,
+)
+
+# Defensive: guarantee the codes table migration (0001_codebook) is
+# registered before this module's ALTER, regardless of import path.
+from potato.codebook.store import _CODEBOOK_MIGRATION as _CB_MIG
+
+register_migration(_CB_MIG)
+register_migration(_REVISION_MIGRATION)
+register_migration(_CODES_REV_MIGRATION)
+
+
+def _db(task_dir: str):
+ register_migration(_REVISION_MIGRATION)
+ register_migration(_CODES_REV_MIGRATION)
+ return get_db(task_dir)
+
+
+def current_revision(task_dir: str, project: str) -> int:
+ row = _db(task_dir).execute(
+ "SELECT revision FROM codebook_revision WHERE project = ?",
+ (project,),
+ ).fetchone()
+ return int(row["revision"]) if row else 0
+
+
+def bump_revision(task_dir: str, project: str) -> int:
+ """Increment (or initialise) the project's revision; return the new
+ value. Called only for option-set-changing codebook ops."""
+ conn = _db(task_dir)
+ now = time.time()
+ conn.execute(
+ """INSERT INTO codebook_revision (project, revision, updated_at)
+ VALUES (?, 1, ?)
+ ON CONFLICT(project) DO UPDATE SET
+ revision = revision + 1,
+ updated_at = excluded.updated_at""",
+ (project, now),
+ )
+ conn.commit()
+ return current_revision(task_dir, project)
+
+
+def record_annotation(
+ task_dir: str, project: str, instance_id: str, username: str
+) -> int:
+ """Stamp (project, instance_id, username) with the current revision
+ at annotation save time. Idempotent upsert; returns the revision."""
+ rev = current_revision(task_dir, project)
+ conn = _db(task_dir)
+ conn.execute(
+ """INSERT INTO annotation_provenance
+ (project, instance_id, username, revision, updated_at)
+ VALUES (?, ?, ?, ?, ?)
+ ON CONFLICT(project, instance_id, username) DO UPDATE SET
+ revision = excluded.revision,
+ updated_at = excluded.updated_at""",
+ (project, instance_id, username, rev, time.time()),
+ )
+ conn.commit()
+ return rev
+
+
+def instance_revision(
+ task_dir: str, project: str, instance_id: str, username: str
+) -> Optional[int]:
+ row = _db(task_dir).execute(
+ """SELECT revision FROM annotation_provenance
+ WHERE project = ? AND instance_id = ? AND username = ?""",
+ (project, instance_id, username),
+ ).fetchone()
+ return int(row["revision"]) if row else None
+
+
+def stale_instances(
+ task_dir: str, project: str, username: str
+) -> List[Dict[str, object]]:
+ """This user's annotated instances whose stamped revision is behind
+ the current one, with the count of codes added since."""
+ cur = current_revision(task_dir, project)
+ if cur <= 0:
+ return []
+ rows = _db(task_dir).execute(
+ """SELECT instance_id, revision FROM annotation_provenance
+ WHERE project = ? AND username = ? AND revision < ?
+ ORDER BY revision ASC, instance_id ASC""",
+ (project, username, cur),
+ ).fetchall()
+ out: List[Dict[str, object]] = []
+ for r in rows:
+ out.append({
+ "instance_id": r["instance_id"],
+ "annotated_revision": int(r["revision"]),
+ "current_revision": cur,
+ "codes_added_since": codes_added_since(
+ task_dir, project, int(r["revision"])),
+ })
+ return out
+
+
+def all_stale_instances(
+ task_dir: str, project: str
+) -> List[Dict[str, object]]:
+ """Every (instance, user) annotated under an older revision โ
+ project-wide, for admin oversight."""
+ cur = current_revision(task_dir, project)
+ if cur <= 0:
+ return []
+ rows = _db(task_dir).execute(
+ """SELECT instance_id, username, revision
+ FROM annotation_provenance
+ WHERE project = ? AND revision < ?
+ ORDER BY revision ASC, instance_id ASC, username ASC""",
+ (project, cur),
+ ).fetchall()
+ return [{
+ "instance_id": r["instance_id"],
+ "username": r["username"],
+ "annotated_revision": int(r["revision"]),
+ "current_revision": cur,
+ "codes_added_since": codes_added_since(
+ task_dir, project, int(r["revision"])),
+ } for r in rows]
+
+
+def touch_instances(
+ task_dir: str, project: str, instance_ids: List[str]
+) -> int:
+ """Re-flag specific instances as stale after a retroactive codebook
+ edit (merge/split/rename) that affected exactly them. Sets their
+ stamped revision to current-1 so `stale_instances` resurfaces them โ
+ soft and dismissible, never a hard re-label gate (Phase 2 (B)
+ policy). Only lowers a revision (never un-stales an already-older
+ row), and never raises one above current."""
+ if not instance_ids:
+ return 0
+ cur = current_revision(task_dir, project)
+ if cur <= 0:
+ return 0
+ target = cur - 1
+ qs = ",".join("?" * len(instance_ids))
+ conn = _db(task_dir)
+ cur_ = conn.execute(
+ f"""UPDATE annotation_provenance
+ SET revision = ?, updated_at = ?
+ WHERE project = ? AND instance_id IN ({qs})
+ AND revision > ?""",
+ [target, time.time(), project, *instance_ids, target],
+ )
+ conn.commit()
+ return cur_.rowcount
+
+
+def codes_added_since(
+ task_dir: str, project: str, revision: int
+) -> List[str]:
+ """Names of codes created after `revision` โ the precise set that
+ could change a label made at that revision."""
+ rows = _db(task_dir).execute(
+ """SELECT name FROM codes
+ WHERE project = ? AND created_revision > ?
+ ORDER BY created_revision ASC, name ASC""",
+ (project, revision),
+ ).fetchall()
+ return [r["name"] for r in rows]
diff --git a/potato/codebook/schema_bridge.py b/potato/codebook/schema_bridge.py
new file mode 100644
index 0000000000000000000000000000000000000000..e1ea01be65e94ded859e4acc4251de4255e05825
--- /dev/null
+++ b/potato/codebook/schema_bridge.py
@@ -0,0 +1,104 @@
+"""
+Schema-loader codebook bridge.
+
+When an annotation scheme opts in with ``codebook: true``, its label
+list is sourced from the project's mutable codebook instead of (only)
+the static YAML ``labels``. Applied once at server start, before
+front-end generation, so every downstream generator
+(radio/multiselect/span/hierarchical_multiselect) keeps reading
+``scheme["labels"]`` unchanged.
+
+Legacy preservation: a config's existing YAML ``labels`` seed the
+codebook the first time (so old configs keep working and the codebook
+starts populated); thereafter the database is the source of truth.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List
+
+from potato.codebook import create_code
+from potato.codebook.codebook import Codebook
+from potato.codebook.service import DuplicateCodeError
+
+logger = logging.getLogger(__name__)
+
+
+def _label_name(entry: Any) -> str:
+ if isinstance(entry, str):
+ return entry
+ if isinstance(entry, dict):
+ return str(entry.get("name") or entry.get("label") or "").strip()
+ return str(entry).strip()
+
+
+def _project_of(config: Dict[str, Any]) -> str:
+ return config.get("annotation_task_name") or "default"
+
+
+def _seed_from_yaml(
+ task_dir: str, project: str, yaml_labels: List[Any]
+) -> None:
+ for entry in yaml_labels or []:
+ name = _label_name(entry)
+ if not name:
+ continue
+ try:
+ create_code(
+ task_dir, project=project, name=name,
+ created_by="config")
+ except DuplicateCodeError:
+ pass # idempotent: re-seeding an existing code is fine
+
+
+def apply_codebook_to_schemes(config: Dict[str, Any]) -> None:
+ """Mutate ``config['annotation_schemes']`` in place: for every
+ scheme with ``codebook: true``, point ``labels`` at the codebook
+ (seeding it from the scheme's YAML labels on first run)."""
+ schemes = config.get("annotation_schemes") or []
+ task_dir = config.get("task_dir", ".")
+ project = _project_of(config)
+
+ for scheme in schemes:
+ if not isinstance(scheme, dict) or not scheme.get("codebook"):
+ continue
+
+ cb = Codebook.load(task_dir, project)
+ if cb.is_empty():
+ _seed_from_yaml(task_dir, project, scheme.get("labels"))
+ cb = Codebook.load(task_dir, project)
+
+ names = cb.labels()
+ if names:
+ scheme["labels"] = names
+ logger.info(
+ "Codebook bridge: scheme %r now sources %d label(s) "
+ "from the project codebook",
+ scheme.get("name"), len(names))
+
+
+def _icl_sync_listener(task_dir: str, project: str) -> None:
+ """Codebook change listener: refresh the *live* server config's
+ scheme labels so ICL prompts (built fresh from ``schema['labels']``
+ each call) are restricted to the codebook's current set. Refreshing
+ the source the prompt is built from *is* the prompt-cache
+ invalidation โ there is no separate persistent ICL prompt cache.
+ """
+ try:
+ from potato.server_utils import config_module
+ cfg = config_module.config
+ except Exception:
+ return
+ if not cfg:
+ return
+ if (cfg.get("annotation_task_name") or "default") != project:
+ return
+ apply_codebook_to_schemes(cfg)
+
+
+def install_codebook_icl_sync() -> None:
+ """Register the ICL-sync listener (idempotent). Called at server
+ init alongside the other mode initializers."""
+ from potato.codebook.service import register_change_listener
+ register_change_listener(_icl_sync_listener)
diff --git a/potato/codebook/service.py b/potato/codebook/service.py
new file mode 100644
index 0000000000000000000000000000000000000000..7abaccdbc63b754934147b5ecd778756e3321a6b
--- /dev/null
+++ b/potato/codebook/service.py
@@ -0,0 +1,362 @@
+"""
+Codebook service โ the single, audited mutation path.
+
+All codebook writes (human *or* LLM, in standard / solo / QDA mode) go
+through here so they share one audit trail (`created_by`), one set of
+invariants (no duplicate siblings, no cycles, recursive delete), and one
+change-notification hook (used by ICL to invalidate its prompt cache โ
+registered via `register_change_listener` to avoid a hard import edge).
+
+Phase 1 ops: create / rename / recolor / move_under / delete.
+merge / split are Phase 2.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, List, Optional
+
+from potato.codebook import store
+from potato.codebook.codebook import Codebook
+
+logger = logging.getLogger(__name__)
+
+
+class CodebookError(Exception):
+ """Base for codebook mutation errors."""
+
+
+class CodeNotFound(CodebookError):
+ pass
+
+
+class DuplicateCodeError(CodebookError):
+ pass
+
+
+class CodebookCycleError(CodebookError):
+ pass
+
+
+# Change listeners: called (task_dir, project) after any successful
+# mutation. ICL registers one to invalidate its prompt cache. Kept as a
+# registry so codebook has no import dependency on the ICL/AI layer.
+_CHANGE_LISTENERS: List[Callable[[str, str], None]] = []
+
+
+def register_change_listener(fn: Callable[[str, str], None]) -> None:
+ if fn not in _CHANGE_LISTENERS:
+ _CHANGE_LISTENERS.append(fn)
+
+
+def clear_change_listeners() -> None:
+ """Tests only โ the registry is process-global."""
+ _CHANGE_LISTENERS.clear()
+
+
+def _notify(task_dir: str, project: str) -> None:
+ for fn in list(_CHANGE_LISTENERS):
+ try:
+ fn(task_dir, project)
+ except Exception: # a listener must never break a mutation
+ logger.exception("codebook change listener failed")
+
+
+def _require(task_dir: str, code_id: str) -> Dict[str, Any]:
+ code = store.get_code(task_dir, code_id)
+ if code is None:
+ raise CodeNotFound(f"Code {code_id} not found")
+ return code
+
+
+def create_code(
+ task_dir: str,
+ *,
+ project: str,
+ name: str,
+ created_by: str,
+ color: Optional[str] = None,
+ parent_id: str = store.ROOT,
+ code_id: Optional[str] = None,
+) -> Dict[str, Any]:
+ name = (name or "").strip()
+ if not name:
+ raise CodebookError("Code name must not be empty")
+ if parent_id != store.ROOT and store.get_code(task_dir, parent_id) is None:
+ raise CodeNotFound(f"Parent code {parent_id} not found")
+ if store.find_code(task_dir, project, parent_id, name) is not None:
+ raise DuplicateCodeError(
+ f"A code named {name!r} already exists at this level")
+ siblings = store.children_of(task_dir, project, parent_id)
+ # A new code changes the option set -> bump the project revision and
+ # stamp the code with the revision it first appeared in.
+ from potato.codebook import revision
+ new_rev = revision.bump_revision(task_dir, project)
+ code = store.insert_code(
+ task_dir, project=project, name=name, created_by=created_by,
+ color=color, parent_id=parent_id, sort_order=len(siblings),
+ code_id=code_id, created_revision=new_rev,
+ )
+ _notify(task_dir, project)
+ return code
+
+
+def _restamp(task_dir: str, project: str, code_ids: List[str]) -> None:
+ """Re-flag exactly the instances whose live links touch `code_ids`
+ so the (B) review worklist resurfaces them (soft, dismissible)."""
+ from potato.codebook import revision
+ affected: List[str] = []
+ seen = set()
+ for cid in code_ids:
+ for aid in store.affected_annotation_ids(task_dir, project, cid):
+ if aid not in seen:
+ seen.add(aid)
+ affected.append(aid)
+ revision.touch_instances(task_dir, project, affected)
+
+
+def rename_code(
+ task_dir: str, code_id: str, *, new_name: str, project: str,
+ actor: str = "system", actor_kind: str = "human",
+) -> Dict[str, Any]:
+ new_name = (new_name or "").strip()
+ if not new_name:
+ raise CodebookError("Code name must not be empty")
+ code = _require(task_dir, code_id)
+ old_name = code["name"]
+ clash = store.find_code(
+ task_dir, project, code["parent_id"], new_name)
+ if clash is not None and clash["id"] != code_id:
+ raise DuplicateCodeError(
+ f"A code named {new_name!r} already exists at this level")
+ updated = store.update_code(task_dir, code_id, name=new_name)
+ # Any codebook change bumps the revision (provenance: an instance
+ # labeled before this change is flagged stale on revisit).
+ from potato.codebook import revision
+ from potato.codebook import changelog
+ new_rev = revision.bump_revision(task_dir, project)
+ changelog.log_change(
+ task_dir, project=project, op="rename", code_id=code_id,
+ old_value=old_name, new_value=new_name, actor=actor,
+ actor_kind=actor_kind, revision=new_rev)
+ _restamp(task_dir, project, [code_id])
+ _notify(task_dir, project)
+ return updated
+
+
+def recolor_code(
+ task_dir: str, code_id: str, *, color: str, project: str,
+ actor: str = "system", actor_kind: str = "human",
+) -> Dict[str, Any]:
+ code = _require(task_dir, code_id)
+ updated = store.update_code(task_dir, code_id, color=color)
+ from potato.codebook import revision
+ from potato.codebook import changelog
+ new_rev = revision.bump_revision(task_dir, project)
+ changelog.log_change(
+ task_dir, project=project, op="recolor", code_id=code_id,
+ old_value=code.get("color"), new_value=color, actor=actor,
+ actor_kind=actor_kind, revision=new_rev)
+ _restamp(task_dir, project, [code_id])
+ _notify(task_dir, project)
+ return updated
+
+
+def _subtree_ids(task_dir: str, project: str, root_id: str) -> List[str]:
+ cb = Codebook.load(task_dir, project)
+ out: List[str] = []
+
+ def walk(cid: str) -> None:
+ out.append(cid)
+ for kid in cb.children(cid):
+ walk(kid["id"])
+
+ walk(root_id)
+ return out
+
+
+def move_under(
+ task_dir: str, code_id: str, *, new_parent_id: str, project: str,
+ actor: str = "system", actor_kind: str = "human",
+) -> Dict[str, Any]:
+ code = _require(task_dir, code_id)
+ if new_parent_id == code_id:
+ raise CodebookCycleError("A code cannot be its own parent")
+ if new_parent_id != store.ROOT:
+ if store.get_code(task_dir, new_parent_id) is None:
+ raise CodeNotFound(f"Parent code {new_parent_id} not found")
+ if new_parent_id in _subtree_ids(task_dir, project, code_id):
+ raise CodebookCycleError(
+ "Cannot move a code under one of its own descendants")
+ clash = store.find_code(
+ task_dir, project, new_parent_id, code["name"])
+ if clash is not None and clash["id"] != code_id:
+ raise DuplicateCodeError(
+ f"A code named {code['name']!r} already exists at the target")
+ siblings = store.children_of(task_dir, project, new_parent_id)
+ old_parent = code["parent_id"]
+ updated = store.update_code(
+ task_dir, code_id,
+ parent_id=new_parent_id, sort_order=len(siblings))
+ from potato.codebook import revision
+ from potato.codebook import changelog
+ new_rev = revision.bump_revision(task_dir, project)
+ changelog.log_change(
+ task_dir, project=project, op="move", code_id=code_id,
+ old_value=old_parent, new_value=new_parent_id, actor=actor,
+ actor_kind=actor_kind, revision=new_rev)
+ _restamp(task_dir, project, [code_id])
+ _notify(task_dir, project)
+ return updated
+
+
+def delete_code(
+ task_dir: str, code_id: str, *, project: str,
+ actor: str = "system", actor_kind: str = "human",
+) -> int:
+ """Delete a code and its entire subtree (and annotation links).
+ Returns the number of code rows removed."""
+ code = _require(task_dir, code_id)
+ ids = _subtree_ids(task_dir, project, code_id)
+ # Capture affected instances BEFORE the (existing) hard delete so
+ # the worklist can still resurface them.
+ from potato.codebook import revision
+ from potato.codebook import changelog
+ affected: List[str] = []
+ seen = set()
+ for cid in ids:
+ for aid in store.affected_annotation_ids(task_dir, project, cid):
+ if aid not in seen:
+ seen.add(aid)
+ affected.append(aid)
+ n = store.delete_codes(task_dir, ids)
+ # Removing a code also changes the option set.
+ new_rev = revision.bump_revision(task_dir, project)
+ changelog.log_change(
+ task_dir, project=project, op="delete", code_id=code_id,
+ old_value=code["name"], new_value=None, actor=actor,
+ actor_kind=actor_kind, revision=new_rev)
+ revision.touch_instances(task_dir, project, affected)
+ _notify(task_dir, project)
+ return n
+
+
+# ---- annotation <-> code links (audited, same notify path) -------------
+
+def apply_code(
+ task_dir: str,
+ *,
+ project: str,
+ annotation_id: str,
+ code_id: str,
+ created_by: str,
+ started_at: Optional[float] = None,
+ ended_at: Optional[float] = None,
+) -> None:
+ _require(task_dir, code_id)
+ store.link_annotation(
+ task_dir, project=project, annotation_id=annotation_id,
+ code_id=code_id, created_by=created_by,
+ started_at=started_at, ended_at=ended_at)
+
+
+def remove_code(
+ task_dir: str, *, annotation_id: str, code_id: str
+) -> bool:
+ return store.unlink_annotation(task_dir, annotation_id, code_id)
+
+
+def codes_on(task_dir: str, annotation_id: str) -> List[Dict[str, Any]]:
+ return store.codes_for_annotation(task_dir, annotation_id)
+
+
+# ---- Phase 2 (C): retroactive merge / split (append-only) --------------
+
+def merge_codes(
+ task_dir: str, *, project: str, src_id: str, dst_id: str,
+ actor: str = "system", actor_kind: str = "human",
+) -> Dict[str, Any]:
+ """Fold `src` into `dst`: every live annotation link to src is
+ re-pointed at dst (idempotent if the annotation already had dst),
+ src's links are invalidated (not deleted), and src is archived (it
+ leaves the palette/ICL prompt but its row + history survive).
+ Affected instances are softly re-flagged for review."""
+ if src_id == dst_id:
+ raise CodebookError("Cannot merge a code into itself")
+ src = _require(task_dir, src_id)
+ dst = _require(task_dir, dst_id)
+ from potato.codebook import revision, changelog
+
+ affected = store.affected_annotation_ids(task_dir, project, src_id)
+ new_rev = revision.bump_revision(task_dir, project)
+ change_id = changelog.log_change(
+ task_dir, project=project, op="merge", code_id=src_id,
+ related_code_id=dst_id, old_value=src["name"],
+ new_value=dst["name"], actor=actor, actor_kind=actor_kind,
+ revision=new_rev)
+ for aid in affected:
+ link = store.get_link(task_dir, aid, src_id) or {}
+ store.set_link_live(
+ task_dir, project=project, annotation_id=aid,
+ code_id=dst_id, created_by=link.get("created_by", actor),
+ started_at=link.get("started_at"),
+ ended_at=link.get("ended_at"))
+ store.invalidate_links(
+ task_dir, project=project, code_id=src_id, change_id=change_id)
+ store.archive_code(task_dir, src_id)
+ revision.touch_instances(task_dir, project, affected)
+ _notify(task_dir, project)
+ return {"merged": len(affected), "src_id": src_id,
+ "dst_id": dst_id, "change_id": change_id}
+
+
+def split_code(
+ task_dir: str, *, project: str, src_id: str, annotator: str,
+ new_name: Optional[str] = None, target_id: Optional[str] = None,
+ actor: str = "system", actor_kind: str = "human",
+) -> Dict[str, Any]:
+ """Split `src` BY ANNOTATOR: move just `annotator`'s live links from
+ src to a target code (existing `target_id`, or a new code named
+ `new_name`). src stays live for other annotators; it is archived
+ only if it ends up with no live links and no children."""
+ src = _require(task_dir, src_id)
+ if not annotator:
+ raise CodebookError("An annotator must be given to split by")
+ from potato.codebook import revision, changelog
+
+ if target_id:
+ target = _require(task_dir, target_id)
+ elif new_name:
+ target = create_code(
+ task_dir, project=project, name=new_name,
+ created_by=actor, parent_id=src["parent_id"])
+ else:
+ raise CodebookError("Provide either target_id or new_name")
+
+ affected = store.affected_annotation_ids(
+ task_dir, project, src_id, created_by=annotator)
+ new_rev = revision.bump_revision(task_dir, project)
+ change_id = changelog.log_change(
+ task_dir, project=project, op="split", code_id=src_id,
+ related_code_id=target["id"], old_value=src["name"],
+ new_value=f"{target['name']} [{annotator}]", actor=actor,
+ actor_kind=actor_kind, revision=new_rev)
+ for aid in affected:
+ link = store.get_link(task_dir, aid, src_id) or {}
+ store.set_link_live(
+ task_dir, project=project, annotation_id=aid,
+ code_id=target["id"], created_by=annotator,
+ started_at=link.get("started_at"),
+ ended_at=link.get("ended_at"))
+ store.invalidate_links(
+ task_dir, project=project, code_id=src_id,
+ change_id=change_id, created_by=annotator)
+ # Archive src only if nothing live remains and it has no children.
+ remaining = store.affected_annotation_ids(task_dir, project, src_id)
+ children = Codebook.load(task_dir, project).children(src_id)
+ if not remaining and not children:
+ store.archive_code(task_dir, src_id)
+ revision.touch_instances(task_dir, project, affected)
+ _notify(task_dir, project)
+ return {"moved": len(affected), "src_id": src_id,
+ "target_id": target["id"], "change_id": change_id}
diff --git a/potato/codebook/similar.py b/potato/codebook/similar.py
new file mode 100644
index 0000000000000000000000000000000000000000..2b3aa502e4a5fe591f96dfe69400c71ad937a3bd
--- /dev/null
+++ b/potato/codebook/similar.py
@@ -0,0 +1,72 @@
+"""
+Soft suggest-on-create helpers (Phase 2 #1 resolution).
+
+When an annotator (or, in solo mode, the LLM) is about to add a code โ
+especially via in-vivo coding where the name is derived from a text
+selection โ near-duplicates proliferate fast ("cost", "costs", "cost
+concerns"). Rather than block or silently merge, we *suggest*: surface
+existing codes that closely match the proposed name so the annotator
+can reuse one. Non-destructive and adjudicator-free, so it works in
+solo mode too.
+
+Pure functions, no I/O โ unit-testable in isolation.
+"""
+
+from __future__ import annotations
+
+import re
+from difflib import SequenceMatcher, get_close_matches
+from typing import List
+
+_WS = re.compile(r"\s+")
+# Conservative default: 0.78 catches "cost concern" ~ "cost concerns"
+# and case/space variants without dragging in merely topical names.
+DEFAULT_CUTOFF = 0.78
+MAX_CODE_NAME = 60
+
+
+def _norm(s: str) -> str:
+ return _WS.sub(" ", str(s or "")).strip().lower()
+
+
+def derive_code_name(text: str, cap: int = MAX_CODE_NAME) -> str:
+ """Propose a code name from a raw text selection: collapse
+ whitespace, trim, and cap length at a word boundary when possible.
+ Mirrors the client-side derivation (codebook.js) โ keep in sync."""
+ s = _WS.sub(" ", str(text or "")).strip()
+ if len(s) <= cap:
+ return s
+ head = s[:cap].rsplit(" ", 1)[0]
+ return (head or s[:cap]).strip()
+
+
+def similar_code_names(
+ names: List[str], proposed: str,
+ cutoff: float = DEFAULT_CUTOFF, n: int = 5,
+) -> List[str]:
+ """Existing code names that closely match `proposed` (normalized),
+ ordered best-first, returned in their ORIGINAL casing. An exact
+ normalized match is always included first."""
+ p = _norm(proposed)
+ if not p:
+ return []
+ norm_to_orig = {}
+ for original in names:
+ norm_to_orig.setdefault(_norm(original), original)
+ keys = list(norm_to_orig.keys())
+
+ ordered: List[str] = []
+ if p in norm_to_orig: # exact (normalized) hit
+ ordered.append(norm_to_orig[p])
+
+ for key in get_close_matches(p, keys, n=n, cutoff=cutoff):
+ orig = norm_to_orig[key]
+ if orig not in ordered:
+ ordered.append(orig)
+
+ # Stable best-first ordering by similarity ratio.
+ ordered.sort(
+ key=lambda o: SequenceMatcher(None, p, _norm(o)).ratio(),
+ reverse=True,
+ )
+ return ordered[:n]
diff --git a/potato/codebook/store.py b/potato/codebook/store.py
new file mode 100644
index 0000000000000000000000000000000000000000..984aa036a2027a8b621d22b8f0d90748ac6c1626
--- /dev/null
+++ b/potato/codebook/store.py
@@ -0,0 +1,358 @@
+"""
+Codebook storage (universal).
+
+SQLite-backed CRUD over `codes` and `annotation_codes` in
+`/project.sqlite` via the universal persistence layer. No
+business rules live here (no cycle checks, no permissions, no cache
+invalidation) โ that is the service layer's job. This module only
+persists rows.
+
+A *code* is a (possibly nested) label in a project's codebook. An
+*annotation_code* links a stored annotation to a code, optionally with a
+time span (`started_at`/`ended_at`) for temporal / agentic-trace coding.
+Universal: usable in standard annotation, solo mode, and QDA mode.
+
+Design notes:
+- `parent_id` is TEXT NOT NULL DEFAULT '' where '' means "root". Using a
+ sentinel instead of NULL lets `UNIQUE(project, parent_id, name)`
+ actually prevent duplicate sibling names at the root too (SQLite
+ treats NULLs as distinct in UNIQUE constraints).
+- No SQL foreign keys (consistent with the memos store): the service
+ layer enforces parent existence, cycle-freedom, and recursive delete.
+"""
+
+from __future__ import annotations
+
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+from potato.persistence import Migration, get_db, register_migration
+
+ROOT = "" # sentinel parent_id for top-level codes
+
+_CODEBOOK_MIGRATION = Migration(
+ name="0001_codebook",
+ sql="""
+ CREATE TABLE IF NOT EXISTS codes (
+ id TEXT PRIMARY KEY,
+ project TEXT NOT NULL,
+ name TEXT NOT NULL,
+ color TEXT,
+ parent_id TEXT NOT NULL DEFAULT '',
+ sort_order INTEGER NOT NULL DEFAULT 0,
+ created_by TEXT NOT NULL,
+ created_at REAL NOT NULL,
+ updated_at REAL NOT NULL,
+ UNIQUE (project, parent_id, name)
+ );
+ CREATE INDEX IF NOT EXISTS idx_codes_project ON codes (project);
+ CREATE INDEX IF NOT EXISTS idx_codes_parent
+ ON codes (project, parent_id);
+
+ CREATE TABLE IF NOT EXISTS annotation_codes (
+ annotation_id TEXT NOT NULL,
+ code_id TEXT NOT NULL,
+ project TEXT NOT NULL,
+ created_by TEXT NOT NULL,
+ started_at REAL,
+ ended_at REAL,
+ PRIMARY KEY (annotation_id, code_id)
+ );
+ CREATE INDEX IF NOT EXISTS idx_anncodes_code
+ ON annotation_codes (project, code_id);
+ CREATE INDEX IF NOT EXISTS idx_anncodes_ann
+ ON annotation_codes (annotation_id);
+ """,
+)
+
+register_migration(_CODEBOOK_MIGRATION)
+
+
+def _db(task_dir: str):
+ """Connection guaranteeing the codebook migration is registered.
+
+ register_migration is idempotent, so this is a no-op normally; it
+ makes the store robust if a test helper (clear_migrations) wiped the
+ process-global registry before this task_dir's first get_db().
+ """
+ register_migration(_CODEBOOK_MIGRATION)
+ return get_db(task_dir)
+
+
+def _ensure_temporal_schema() -> None:
+ """Guarantee the Phase 2 (C) append-only columns/tables exist before
+ any link read/write that depends on `invalidated_at`. Lazy import
+ avoids a module-load cycle (changelog imports this module). The
+ 0003 migration is additive (nullable cols + new tables), so this is
+ safe even for callers that only registered 0001."""
+ from potato.codebook.changelog import _CHANGE_MIGRATION
+ from potato.codebook.revision import (
+ _REVISION_MIGRATION, _CODES_REV_MIGRATION)
+ register_migration(_REVISION_MIGRATION)
+ register_migration(_CODES_REV_MIGRATION)
+ register_migration(_CHANGE_MIGRATION)
+
+
+# ---- codes ---------------------------------------------------------------
+
+def insert_code(
+ task_dir: str,
+ *,
+ project: str,
+ name: str,
+ created_by: str,
+ color: Optional[str] = None,
+ parent_id: str = ROOT,
+ sort_order: int = 0,
+ code_id: Optional[str] = None,
+ created_revision: int = 0,
+) -> Dict[str, Any]:
+ """Insert one code row and return it. `code_id` lets the init CLI
+ supply a deterministic id; otherwise a random uuid4 is used.
+ `created_revision` records the codebook revision the code first
+ appeared in (for provenance / the review worklist)."""
+ cid = code_id or uuid.uuid4().hex
+ now = time.time()
+ conn = _db(task_dir)
+ conn.execute(
+ """INSERT INTO codes
+ (id, project, name, color, parent_id, sort_order,
+ created_by, created_at, updated_at, created_revision)
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+ (cid, project, name, color, parent_id, sort_order,
+ created_by, now, now, created_revision),
+ )
+ conn.commit()
+ return get_code(task_dir, cid)
+
+
+def get_code(task_dir: str, code_id: str) -> Optional[Dict[str, Any]]:
+ row = _db(task_dir).execute(
+ "SELECT * FROM codes WHERE id = ?", (code_id,)
+ ).fetchone()
+ return dict(row) if row else None
+
+
+def find_code(
+ task_dir: str, project: str, parent_id: str, name: str
+) -> Optional[Dict[str, Any]]:
+ row = _db(task_dir).execute(
+ """SELECT * FROM codes
+ WHERE project = ? AND parent_id = ? AND name = ?""",
+ (project, parent_id, name),
+ ).fetchone()
+ return dict(row) if row else None
+
+
+def list_codes(task_dir: str, project: str) -> List[Dict[str, Any]]:
+ rows = _db(task_dir).execute(
+ """SELECT * FROM codes WHERE project = ?
+ ORDER BY parent_id ASC, sort_order ASC, name ASC""",
+ (project,),
+ ).fetchall()
+ return [dict(r) for r in rows]
+
+
+def children_of(
+ task_dir: str, project: str, parent_id: str
+) -> List[Dict[str, Any]]:
+ rows = _db(task_dir).execute(
+ """SELECT * FROM codes WHERE project = ? AND parent_id = ?
+ ORDER BY sort_order ASC, name ASC""",
+ (project, parent_id),
+ ).fetchall()
+ return [dict(r) for r in rows]
+
+
+def update_code(
+ task_dir: str,
+ code_id: str,
+ *,
+ name: Optional[str] = None,
+ color: Optional[str] = None,
+ parent_id: Optional[str] = None,
+ sort_order: Optional[int] = None,
+) -> Optional[Dict[str, Any]]:
+ sets, params = [], []
+ if name is not None:
+ sets.append("name = ?"); params.append(name)
+ if color is not None:
+ sets.append("color = ?"); params.append(color)
+ if parent_id is not None:
+ sets.append("parent_id = ?"); params.append(parent_id)
+ if sort_order is not None:
+ sets.append("sort_order = ?"); params.append(sort_order)
+ if not sets:
+ return get_code(task_dir, code_id)
+ sets.append("updated_at = ?"); params.append(time.time())
+ params.append(code_id)
+ conn = _db(task_dir)
+ conn.execute(f"UPDATE codes SET {', '.join(sets)} WHERE id = ?", params)
+ conn.commit()
+ return get_code(task_dir, code_id)
+
+
+def delete_codes(task_dir: str, code_ids: List[str]) -> int:
+ """Delete the given codes and their annotation_codes links. The
+ service computes the full subtree; this just executes the delete."""
+ if not code_ids:
+ return 0
+ qs = ",".join("?" * len(code_ids))
+ conn = _db(task_dir)
+ conn.execute(
+ f"DELETE FROM annotation_codes WHERE code_id IN ({qs})", code_ids)
+ cur = conn.execute(
+ f"DELETE FROM codes WHERE id IN ({qs})", code_ids)
+ conn.commit()
+ return cur.rowcount
+
+
+# ---- annotation_codes ----------------------------------------------------
+
+def link_annotation(
+ task_dir: str,
+ *,
+ project: str,
+ annotation_id: str,
+ code_id: str,
+ created_by: str,
+ started_at: Optional[float] = None,
+ ended_at: Optional[float] = None,
+) -> None:
+ conn = _db(task_dir)
+ conn.execute(
+ """INSERT OR REPLACE INTO annotation_codes
+ (annotation_id, code_id, project, created_by,
+ started_at, ended_at)
+ VALUES (?, ?, ?, ?, ?, ?)""",
+ (annotation_id, code_id, project, created_by,
+ started_at, ended_at),
+ )
+ conn.commit()
+
+
+def unlink_annotation(
+ task_dir: str, annotation_id: str, code_id: str
+) -> bool:
+ conn = _db(task_dir)
+ cur = conn.execute(
+ """DELETE FROM annotation_codes
+ WHERE annotation_id = ? AND code_id = ?""",
+ (annotation_id, code_id),
+ )
+ conn.commit()
+ return cur.rowcount > 0
+
+
+def codes_for_annotation(
+ task_dir: str, annotation_id: str
+) -> List[Dict[str, Any]]:
+ # THE single load-bearing temporal reader: only LIVE links, and
+ # never an archived (e.g. merged-away) code.
+ _ensure_temporal_schema()
+ rows = _db(task_dir).execute(
+ """SELECT ac.code_id, ac.started_at, ac.ended_at,
+ ac.created_by, c.name, c.color, c.parent_id
+ FROM annotation_codes ac
+ JOIN codes c ON c.id = ac.code_id
+ WHERE ac.annotation_id = ?
+ AND ac.invalidated_at IS NULL
+ AND c.archived_at IS NULL
+ ORDER BY c.name ASC""",
+ (annotation_id,),
+ ).fetchall()
+ return [dict(r) for r in rows]
+
+
+# ---- Phase 2 (C): append-only retroactive primitives --------------------
+
+def affected_annotation_ids(
+ task_dir: str, project: str, code_id: str,
+ created_by: Optional[str] = None,
+) -> List[str]:
+ """annotation_ids with a LIVE link to `code_id` (optionally only
+ those created by `created_by` โ the split-by-annotator selector)."""
+ _ensure_temporal_schema()
+ q = ("SELECT DISTINCT annotation_id FROM annotation_codes "
+ "WHERE project = ? AND code_id = ? AND invalidated_at IS NULL")
+ p: List[Any] = [project, code_id]
+ if created_by is not None:
+ q += " AND created_by = ?"
+ p.append(created_by)
+ rows = _db(task_dir).execute(q, p).fetchall()
+ return [r["annotation_id"] for r in rows]
+
+
+def get_link(
+ task_dir: str, annotation_id: str, code_id: str
+) -> Optional[Dict[str, Any]]:
+ _ensure_temporal_schema()
+ row = _db(task_dir).execute(
+ """SELECT * FROM annotation_codes
+ WHERE annotation_id = ? AND code_id = ?""",
+ (annotation_id, code_id),
+ ).fetchone()
+ return dict(row) if row else None
+
+
+def invalidate_links(
+ task_dir: str, *, project: str, code_id: str, change_id: str,
+ created_by: Optional[str] = None,
+) -> int:
+ """Mark live links to `code_id` superseded (append-only โ never
+ DELETE). Optional `created_by` scopes to one annotator (split)."""
+ _ensure_temporal_schema()
+ q = ("UPDATE annotation_codes SET invalidated_at = ?, "
+ "invalidated_by_change = ? "
+ "WHERE project = ? AND code_id = ? AND invalidated_at IS NULL")
+ p: List[Any] = [time.time(), change_id, project, code_id]
+ if created_by is not None:
+ q += " AND created_by = ?"
+ p.append(created_by)
+ conn = _db(task_dir)
+ cur = conn.execute(q, p)
+ conn.commit()
+ return cur.rowcount
+
+
+def set_link_live(
+ task_dir: str, *, project: str, annotation_id: str, code_id: str,
+ created_by: str, started_at: Optional[float] = None,
+ ended_at: Optional[float] = None,
+) -> None:
+ """Make (annotation_id, code_id) a LIVE link. Idempotent against the
+ PK(annotation_id, code_id): if the row exists (live or invalidated)
+ it is reactivated rather than duplicated/clobbered โ this is how a
+ merge stays correct when the annotation is already on the target."""
+ _ensure_temporal_schema()
+ conn = _db(task_dir)
+ cur = conn.execute(
+ """UPDATE annotation_codes
+ SET invalidated_at = NULL, invalidated_by_change = NULL
+ WHERE annotation_id = ? AND code_id = ?""",
+ (annotation_id, code_id),
+ )
+ if cur.rowcount == 0:
+ conn.execute(
+ """INSERT INTO annotation_codes
+ (annotation_id, code_id, project, created_by,
+ started_at, ended_at)
+ VALUES (?, ?, ?, ?, ?, ?)""",
+ (annotation_id, code_id, project, created_by,
+ started_at, ended_at),
+ )
+ conn.commit()
+
+
+def archive_code(task_dir: str, code_id: str) -> bool:
+ """Soft-archive a code (merged away): leaves the live palette + ICL
+ prompt but the row and its history survive (append-only)."""
+ _ensure_temporal_schema()
+ conn = _db(task_dir)
+ cur = conn.execute(
+ "UPDATE codes SET archived_at = ?, updated_at = ? WHERE id = ?",
+ (time.time(), time.time(), code_id),
+ )
+ conn.commit()
+ return cur.rowcount > 0
diff --git a/potato/codebook_cli.py b/potato/codebook_cli.py
new file mode 100644
index 0000000000000000000000000000000000000000..8bfeadc942d22dc2e3238d263d5ff330788b6e9b
--- /dev/null
+++ b/potato/codebook_cli.py
@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+"""
+`potato codebook ` โ initialise / migrate a project's
+codebook from its YAML config.
+
+For every annotation scheme with ``codebook: true`` it ensures a code
+exists for each YAML label. Codes get **deterministic** ids
+(``uuid5`` over ``project | parent_id | name``) so re-running is a
+no-op and the same config always yields the same ids across machines
+(important when annotation rows carry a parallel ``code_id``).
+
+Idempotent: existing codes are left untouched; only missing ones are
+created. Safe to run repeatedly and in CI.
+
+Usage:
+ potato codebook path/to/config.yaml
+ potato codebook path/to/config.yaml --dry-run
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import uuid
+from typing import Any, Dict, List
+
+import yaml
+
+from potato.codebook import create_code
+from potato.codebook.codebook import Codebook
+from potato.codebook.service import DuplicateCodeError
+from potato.codebook.store import ROOT
+
+# Stable namespace so ids are reproducible across machines/runs.
+_NS = uuid.UUID("6b9b1f6e-1c2d-5a4b-9e3f-c0deb00c0de5")
+
+
+def deterministic_code_id(project: str, parent_id: str, name: str) -> str:
+ return uuid.uuid5(_NS, f"{project}\x1f{parent_id}\x1f{name}").hex
+
+
+def _label_name(entry: Any) -> str:
+ if isinstance(entry, str):
+ return entry
+ if isinstance(entry, dict):
+ return str(entry.get("name") or entry.get("label") or "").strip()
+ return str(entry).strip()
+
+
+def _resolve_task_dir(config_file: str, config: Dict[str, Any]) -> str:
+ base = os.path.dirname(os.path.abspath(config_file))
+ return os.path.normpath(os.path.join(base, config.get("task_dir", ".")))
+
+
+def init_codebook(config_file: str, *, dry_run: bool = False) -> Dict[str, int]:
+ """Seed missing codes for every codebook-enabled scheme.
+
+ Returns {"created": n, "existing": m}.
+ """
+ with open(config_file, "rt", encoding="utf-8") as fh:
+ config = yaml.safe_load(fh) or {}
+
+ task_dir = _resolve_task_dir(config_file, config)
+ project = config.get("annotation_task_name") or "default"
+ schemes: List[Dict[str, Any]] = config.get("annotation_schemes") or []
+
+ created = existing = 0
+ for scheme in schemes:
+ if not isinstance(scheme, dict) or not scheme.get("codebook"):
+ continue
+ cb = Codebook.load(task_dir, project)
+ present = set(cb.labels())
+ for entry in scheme.get("labels") or []:
+ name = _label_name(entry)
+ if not name:
+ continue
+ if name in present:
+ existing += 1
+ continue
+ if dry_run:
+ created += 1
+ present.add(name)
+ continue
+ cid = deterministic_code_id(project, ROOT, name)
+ try:
+ create_code(task_dir, project=project, name=name,
+ created_by="codebook-cli", code_id=cid)
+ created += 1
+ present.add(name)
+ except DuplicateCodeError:
+ existing += 1
+
+ return {"created": created, "existing": existing}
+
+
+def main(argv=None) -> int:
+ parser = argparse.ArgumentParser(
+ prog="potato codebook",
+ description="Initialise a project codebook from its YAML config.")
+ parser.add_argument("config_file", help="Path to the project config.yaml")
+ parser.add_argument(
+ "--dry-run", action="store_true",
+ help="Report what would be created without writing.")
+ args = parser.parse_args(argv)
+
+ if not os.path.isfile(args.config_file):
+ print(f"Config not found: {args.config_file}", file=sys.stderr)
+ return 2
+
+ result = init_codebook(args.config_file, dry_run=args.dry_run)
+ verb = "Would create" if args.dry_run else "Created"
+ print(f"{verb} {result['created']} code(s); "
+ f"{result['existing']} already present.")
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/potato/coding_agent_backend.py b/potato/coding_agent_backend.py
new file mode 100644
index 0000000000000000000000000000000000000000..b17e955da1417d281c4d0f8bc12815463749c223
--- /dev/null
+++ b/potato/coding_agent_backend.py
@@ -0,0 +1,324 @@
+"""
+Coding Agent Backend Abstraction
+
+Defines the interface for coding agent backends and common event types.
+Backends implement the agent loop (LLM + tool execution) and yield
+events that the CodingAgentRunner consumes.
+
+Available backends:
+- anthropic_tool_use: Custom agent loop using Anthropic API
+- ollama_tool_use: Custom agent loop using Ollama (fully local, no API key)
+- openai_tool_use: Custom agent loop using any OpenAI-compatible server
+ (OpenAI, vLLM, llama.cpp, ...) with tool calling
+- claude_sdk: Claude Agent SDK (subprocess with JSON-lines IPC)
+- subprocess: Generic CLI agent (Phase 4)
+- opencode: OpenCode SDK (Phase 4)
+"""
+
+import logging
+import os
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, Iterator, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class CodingAgentEventType(str, Enum):
+ """Event types emitted by coding agent backends."""
+ THINKING = "thinking"
+ TOOL_CALL_START = "tool_call_start"
+ TOOL_CALL_END = "tool_call_end"
+ TURN_END = "turn_end"
+ ERROR = "error"
+ COMPLETE = "complete"
+
+
+@dataclass
+class CodingAgentEvent:
+ """Single event from a coding agent backend."""
+ event_type: CodingAgentEventType
+ timestamp: float = 0.0
+ data: Dict[str, Any] = field(default_factory=dict)
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ "event_type": self.event_type.value,
+ "timestamp": self.timestamp,
+ "data": self.data,
+ }
+
+
+# Tool definitions for custom tool-use backends
+CODING_TOOLS = [
+ {
+ "name": "Read",
+ "description": "Read a file from the filesystem. Returns the file contents.",
+ "input_schema": {
+ "type": "object",
+ "properties": {
+ "file_path": {"type": "string", "description": "Absolute or relative path to the file"},
+ },
+ "required": ["file_path"],
+ },
+ },
+ {
+ "name": "Edit",
+ "description": "Replace a specific string in a file with a new string.",
+ "input_schema": {
+ "type": "object",
+ "properties": {
+ "file_path": {"type": "string", "description": "Path to the file to edit"},
+ "old_string": {"type": "string", "description": "The exact text to find and replace"},
+ "new_string": {"type": "string", "description": "The replacement text"},
+ },
+ "required": ["file_path", "old_string", "new_string"],
+ },
+ },
+ {
+ "name": "Write",
+ "description": "Create or overwrite a file with the given content.",
+ "input_schema": {
+ "type": "object",
+ "properties": {
+ "file_path": {"type": "string", "description": "Path to the file to write"},
+ "content": {"type": "string", "description": "The full file content"},
+ },
+ "required": ["file_path", "content"],
+ },
+ },
+ {
+ "name": "Bash",
+ "description": "Execute a bash command and return its output.",
+ "input_schema": {
+ "type": "object",
+ "properties": {
+ "command": {"type": "string", "description": "The command to execute"},
+ },
+ "required": ["command"],
+ },
+ },
+ {
+ "name": "Grep",
+ "description": "Search for a pattern in files. Returns matching lines with file paths.",
+ "input_schema": {
+ "type": "object",
+ "properties": {
+ "pattern": {"type": "string", "description": "Regex pattern to search for"},
+ "path": {"type": "string", "description": "Directory or file to search in"},
+ },
+ "required": ["pattern"],
+ },
+ },
+ {
+ "name": "Glob",
+ "description": "Find files matching a glob pattern.",
+ "input_schema": {
+ "type": "object",
+ "properties": {
+ "pattern": {"type": "string", "description": "Glob pattern (e.g. '**/*.py')"},
+ },
+ "required": ["pattern"],
+ },
+ },
+]
+
+# Ollama-compatible tool format (OpenAI function calling style)
+CODING_TOOLS_OLLAMA = [
+ {
+ "type": "function",
+ "function": {
+ "name": t["name"],
+ "description": t["description"],
+ "parameters": t["input_schema"],
+ },
+ }
+ for t in CODING_TOOLS
+]
+
+
+def execute_tool(tool_name: str, tool_input: dict, working_dir: str) -> str:
+ """Execute a coding tool in the working directory.
+
+ Args:
+ tool_name: Tool name (Read, Edit, Write, Bash, Grep, Glob)
+ tool_input: Tool input parameters
+ working_dir: Working directory for file operations
+
+ Returns:
+ Tool output as a string
+ """
+ import glob as glob_module
+ import subprocess
+
+ try:
+ if tool_name == "Read":
+ file_path = tool_input["file_path"]
+ abs_path = os.path.join(working_dir, file_path) if not os.path.isabs(file_path) else file_path
+ with open(abs_path, "r", encoding="utf-8", errors="replace") as f:
+ return f.read()
+
+ elif tool_name == "Edit":
+ file_path = tool_input["file_path"]
+ abs_path = os.path.join(working_dir, file_path) if not os.path.isabs(file_path) else file_path
+ old_string = tool_input["old_string"]
+ new_string = tool_input["new_string"]
+ with open(abs_path, "r", encoding="utf-8") as f:
+ content = f.read()
+ if old_string not in content:
+ return f"Error: old_string not found in {file_path}"
+ content = content.replace(old_string, new_string, 1)
+ with open(abs_path, "w", encoding="utf-8") as f:
+ f.write(content)
+ return "Edit applied successfully."
+
+ elif tool_name == "Write":
+ file_path = tool_input["file_path"]
+ abs_path = os.path.join(working_dir, file_path) if not os.path.isabs(file_path) else file_path
+ os.makedirs(os.path.dirname(abs_path), exist_ok=True)
+ with open(abs_path, "w", encoding="utf-8") as f:
+ f.write(tool_input["content"])
+ return f"File written: {file_path}"
+
+ elif tool_name == "Bash":
+ command = tool_input["command"]
+ result = subprocess.run(
+ command, shell=True, capture_output=True, text=True,
+ cwd=working_dir, timeout=60,
+ )
+ output = result.stdout
+ if result.stderr:
+ output += "\n" + result.stderr
+ if result.returncode != 0:
+ output += f"\n[exit code: {result.returncode}]"
+ return output.strip() or "(no output)"
+
+ elif tool_name == "Grep":
+ pattern = tool_input["pattern"]
+ path = tool_input.get("path", ".")
+ abs_path = os.path.join(working_dir, path) if not os.path.isabs(path) else path
+ result = subprocess.run(
+ ["grep", "-rn", pattern, abs_path],
+ capture_output=True, text=True, cwd=working_dir, timeout=30,
+ )
+ return result.stdout.strip() or "(no matches)"
+
+ elif tool_name == "Glob":
+ pattern = tool_input["pattern"]
+ matches = sorted(glob_module.glob(
+ os.path.join(working_dir, pattern), recursive=True
+ ))
+ # Make paths relative to working_dir
+ rel_matches = [os.path.relpath(m, working_dir) for m in matches]
+ return "\n".join(rel_matches) or "(no matches)"
+
+ else:
+ return f"Unknown tool: {tool_name}"
+
+ except FileNotFoundError as e:
+ return f"Error: File not found: {e}"
+ except PermissionError as e:
+ return f"Error: Permission denied: {e}"
+ except subprocess.TimeoutExpired:
+ return "Error: Command timed out (60s limit)"
+ except Exception as e:
+ return f"Error: {type(e).__name__}: {e}"
+
+
+class CodingAgentBackend(ABC):
+ """Abstract interface for coding agent backends."""
+
+ @abstractmethod
+ def start(self, task: str, working_dir: str, system_prompt: str = "") -> None:
+ """Start the agent with a task description."""
+ ...
+
+ @abstractmethod
+ def get_events(self) -> Iterator[CodingAgentEvent]:
+ """Yield events as the agent works. Blocks until next event or completion."""
+ ...
+
+ @abstractmethod
+ def pause(self) -> None:
+ """Pause the agent between tool executions."""
+ ...
+
+ @abstractmethod
+ def resume(self) -> None:
+ """Resume a paused agent."""
+ ...
+
+ @abstractmethod
+ def inject_instruction(self, text: str) -> None:
+ """Send an instruction to the agent (appended as user message)."""
+ ...
+
+ @abstractmethod
+ def stop(self) -> None:
+ """Stop the agent."""
+ ...
+
+ @abstractmethod
+ def get_conversation_history(self) -> List[Dict]:
+ """Get the full conversation history."""
+ ...
+
+ @abstractmethod
+ def get_state(self) -> str:
+ """Get the current state: running, paused, completed, error."""
+ ...
+
+ def truncate_history(self, to_step: int) -> None:
+ """Truncate conversation history to the given step (for rollback)."""
+ pass # Optional, backends that support rollback override this
+
+
+# Backend registry
+BACKEND_REGISTRY: Dict[str, type] = {}
+
+
+def register_backend(name: str, cls: type) -> None:
+ """Register a backend implementation."""
+ BACKEND_REGISTRY[name] = cls
+
+
+def create_backend(backend_type: str, config: dict) -> CodingAgentBackend:
+ """Create a backend instance from config."""
+ if backend_type not in BACKEND_REGISTRY:
+ available = ", ".join(sorted(BACKEND_REGISTRY.keys()))
+ raise ValueError(
+ f"Unknown backend type '{backend_type}'. Available: {available}"
+ )
+ cls = BACKEND_REGISTRY[backend_type]
+ return cls(config)
+
+
+def _register_builtin_backends():
+ """Register built-in backends. Called on import."""
+ try:
+ from .coding_agent_backends.anthropic_backend import AnthropicToolUseBackend
+ register_backend("anthropic_tool_use", AnthropicToolUseBackend)
+ except ImportError:
+ logger.debug("Anthropic backend not available (missing anthropic package)")
+
+ try:
+ from .coding_agent_backends.ollama_backend import OllamaToolUseBackend
+ register_backend("ollama_tool_use", OllamaToolUseBackend)
+ except ImportError:
+ logger.debug("Ollama backend not available")
+
+ try:
+ from .coding_agent_backends.openai_backend import OpenAIToolUseBackend
+ register_backend("openai_tool_use", OpenAIToolUseBackend)
+ except ImportError:
+ logger.debug("OpenAI backend not available (missing openai package)")
+
+ try:
+ from .coding_agent_backends.claude_sdk_backend import ClaudeSDKBackend
+ register_backend("claude_sdk", ClaudeSDKBackend)
+ except ImportError:
+ logger.debug("Claude SDK backend not available (missing claude-agent-sdk)")
+
+
+_register_builtin_backends()
diff --git a/potato/coding_agent_backends/__init__.py b/potato/coding_agent_backends/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f7166633af5c71f96ea6c647d844ebe9e67c4363
--- /dev/null
+++ b/potato/coding_agent_backends/__init__.py
@@ -0,0 +1 @@
+# Coding Agent Backends
diff --git a/potato/coding_agent_backends/anthropic_backend.py b/potato/coding_agent_backends/anthropic_backend.py
new file mode 100644
index 0000000000000000000000000000000000000000..820d7802d012efa94b37a409bc58f61de0c63194
--- /dev/null
+++ b/potato/coding_agent_backends/anthropic_backend.py
@@ -0,0 +1,287 @@
+"""
+Anthropic Tool Use Backend
+
+Custom agent loop using the Anthropic Messages API with tool_use.
+Defines Read/Edit/Write/Bash/Grep/Glob tools and executes them
+in the working directory.
+"""
+
+import json
+import logging
+import time
+import threading
+from typing import Any, Dict, Iterator, List, Optional
+
+from ..coding_agent_backend import (
+ CodingAgentBackend,
+ CodingAgentEvent,
+ CodingAgentEventType,
+ CODING_TOOLS,
+ execute_tool,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AnthropicToolUseBackend(CodingAgentBackend):
+ """Agent loop using Anthropic Messages API with tool_use."""
+
+ def __init__(self, config: dict):
+ self._config = config
+ ai = config.get("ai_config", {})
+ self._model = ai.get("model", "claude-sonnet-4-20250514")
+ self._api_key = ai.get("api_key", "")
+ self._max_tokens = ai.get("max_tokens", 8192)
+ self._temperature = ai.get("temperature", 0.3)
+ self._max_turns = config.get("max_turns", 50)
+
+ self._state = "idle"
+ self._working_dir = ""
+ self._messages: List[Dict] = []
+ self._system_prompt = ""
+ self._events: list = []
+ self._event_idx = 0
+ self._pause_event = threading.Event()
+ self._pause_event.set() # Not paused initially
+ self._stop_flag = False
+ self._instruction_queue: list = []
+ self._lock = threading.Lock()
+ self._client = None
+
+ def start(self, task: str, working_dir: str, system_prompt: str = "") -> None:
+ self._working_dir = working_dir
+ self._system_prompt = system_prompt or (
+ "You are a coding agent. You have access to tools for reading, "
+ "editing, and creating files, running bash commands, and searching code. "
+ "Use these tools to complete the task. When you are done, stop calling tools "
+ "and summarize what you did."
+ )
+ self._messages = [{"role": "user", "content": task}]
+ self._state = "running"
+ self._stop_flag = False
+ self._events = []
+ self._event_idx = 0
+
+ # Initialize Anthropic client
+ try:
+ import anthropic
+ self._client = anthropic.Anthropic(api_key=self._api_key or None)
+ except ImportError:
+ self._emit(CodingAgentEventType.ERROR, {"message": "anthropic package not installed"})
+ self._state = "error"
+ return
+
+ # Run the agent loop in a thread
+ thread = threading.Thread(target=self._run_loop, daemon=True)
+ thread.start()
+
+ def _run_loop(self):
+ """Main agent loop: query LLM โ execute tools โ repeat."""
+ turn_index = 0
+ try:
+ while not self._stop_flag and turn_index < self._max_turns:
+ # Check for pause
+ self._pause_event.wait()
+ if self._stop_flag:
+ break
+
+ # Check for injected instructions
+ with self._lock:
+ if self._instruction_queue:
+ instruction = self._instruction_queue.pop(0)
+ self._messages.append({"role": "user", "content": instruction})
+
+ # Convert tools to Anthropic format
+ tools = [
+ {"name": t["name"], "description": t["description"],
+ "input_schema": t["input_schema"]}
+ for t in CODING_TOOLS
+ ]
+
+ # Query LLM
+ self._emit(CodingAgentEventType.THINKING, {
+ "turn_index": turn_index,
+ "text": "Thinking...",
+ })
+
+ try:
+ response = self._client.messages.create(
+ model=self._model,
+ max_tokens=self._max_tokens,
+ temperature=self._temperature,
+ system=self._system_prompt,
+ messages=self._messages,
+ tools=tools,
+ )
+ except Exception as e:
+ self._emit(CodingAgentEventType.ERROR, {"message": str(e)})
+ self._state = "error"
+ return
+
+ # Process response
+ reasoning_parts = []
+ tool_calls = []
+ tool_use_blocks = []
+
+ for block in response.content:
+ if block.type == "text":
+ reasoning_parts.append(block.text)
+ self._emit(CodingAgentEventType.THINKING, {
+ "turn_index": turn_index,
+ "text": block.text,
+ })
+ elif block.type == "tool_use":
+ tool_use_blocks.append(block)
+
+ # Add assistant message to history
+ self._messages.append({
+ "role": "assistant",
+ "content": [b.model_dump() for b in response.content],
+ })
+
+ # Execute tool calls
+ tool_results = []
+ for block in tool_use_blocks:
+ if self._stop_flag:
+ break
+
+ # Check for pause between tool executions
+ self._pause_event.wait()
+ if self._stop_flag:
+ break
+
+ tool_name = block.name
+ tool_input = block.input if isinstance(block.input, dict) else {}
+
+ self._emit(CodingAgentEventType.TOOL_CALL_START, {
+ "turn_index": turn_index,
+ "tool": tool_name,
+ "input": tool_input,
+ })
+
+ # Execute the tool
+ output = execute_tool(tool_name, tool_input, self._working_dir)
+
+ # Classify output type
+ output_type = self._classify_output_type(tool_name)
+
+ tc = {
+ "tool": tool_name,
+ "input": tool_input,
+ "output": output,
+ "output_type": output_type,
+ }
+ tool_calls.append(tc)
+
+ self._emit(CodingAgentEventType.TOOL_CALL_END, {
+ "turn_index": turn_index,
+ "tool_index": len(tool_calls) - 1,
+ **tc,
+ })
+
+ tool_results.append({
+ "type": "tool_result",
+ "tool_use_id": block.id,
+ "content": output,
+ })
+
+ # Add tool results to history
+ if tool_results:
+ self._messages.append({
+ "role": "user",
+ "content": tool_results,
+ })
+
+ # Emit turn_end
+ self._emit(CodingAgentEventType.TURN_END, {
+ "turn_index": turn_index,
+ "content": "\n".join(reasoning_parts),
+ "tool_calls": tool_calls,
+ })
+
+ turn_index += 1
+
+ # If no tool calls, the agent is done
+ if not tool_use_blocks or response.stop_reason == "end_turn":
+ break
+
+ self._state = "completed"
+ self._emit(CodingAgentEventType.COMPLETE, {"total_turns": turn_index})
+
+ except Exception as e:
+ logger.exception("Agent loop error")
+ self._state = "error"
+ self._emit(CodingAgentEventType.ERROR, {"message": str(e)})
+
+ def _classify_output_type(self, tool_name: str) -> str:
+ name = tool_name.lower()
+ if name in ("bash", "terminal", "shell"):
+ return "terminal"
+ if name in ("edit", "replace"):
+ return "diff"
+ if name in ("read", "grep", "glob", "search", "write"):
+ return "code"
+ return "generic"
+
+ def _emit(self, event_type: CodingAgentEventType, data: dict):
+ event = CodingAgentEvent(
+ event_type=event_type,
+ timestamp=time.time(),
+ data=data,
+ )
+ with self._lock:
+ self._events.append(event)
+
+ def get_events(self) -> Iterator[CodingAgentEvent]:
+ while True:
+ with self._lock:
+ if self._event_idx < len(self._events):
+ event = self._events[self._event_idx]
+ self._event_idx += 1
+ yield event
+ if event.event_type in (CodingAgentEventType.COMPLETE, CodingAgentEventType.ERROR):
+ return
+ continue
+ # No new events, wait a bit
+ if self._state in ("completed", "error"):
+ return
+ time.sleep(0.1)
+
+ def pause(self) -> None:
+ self._pause_event.clear()
+ self._state = "paused"
+
+ def resume(self) -> None:
+ self._state = "running"
+ self._pause_event.set()
+
+ def inject_instruction(self, text: str) -> None:
+ with self._lock:
+ self._instruction_queue.append(text)
+
+ def stop(self) -> None:
+ self._stop_flag = True
+ self._pause_event.set() # Unblock if paused
+ self._state = "completed"
+
+ def get_conversation_history(self) -> List[Dict]:
+ with self._lock:
+ return list(self._messages)
+
+ def get_state(self) -> str:
+ return self._state
+
+ def truncate_history(self, to_step: int) -> None:
+ """Truncate conversation to the given turn index."""
+ with self._lock:
+ # Keep initial user message + 2 messages per turn (assistant + tool_results)
+ keep = 1 + (to_step * 2)
+ self._messages = self._messages[:keep]
+ # Also truncate events
+ new_events = []
+ for e in self._events:
+ ti = e.data.get("turn_index", -1)
+ if ti < to_step or ti == -1:
+ new_events.append(e)
+ self._events = new_events
+ self._event_idx = min(self._event_idx, len(self._events))
diff --git a/potato/coding_agent_backends/claude_sdk_backend.py b/potato/coding_agent_backends/claude_sdk_backend.py
new file mode 100644
index 0000000000000000000000000000000000000000..83e3ee16800b6fc0adc1eac6665c0ad08cbe8ab0
--- /dev/null
+++ b/potato/coding_agent_backends/claude_sdk_backend.py
@@ -0,0 +1,272 @@
+"""
+Claude Agent SDK Backend
+
+Spawns Claude Code as a subprocess via the official Agent SDK.
+Communicates via JSON-lines over stdin/stdout.
+Inherits CLAUDE.md, hooks, MCP servers from the working directory.
+
+Requires: pip install claude-agent-sdk
+"""
+
+import json
+import logging
+import subprocess
+import threading
+import time
+from typing import Any, Dict, Iterator, List, Optional
+
+from ..coding_agent_backend import (
+ CodingAgentBackend,
+ CodingAgentEvent,
+ CodingAgentEventType,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ClaudeSDKBackend(CodingAgentBackend):
+ """Backend using Claude Agent SDK (subprocess with JSON-lines IPC)."""
+
+ def __init__(self, config: dict):
+ self._config = config
+ self._state = "idle"
+ self._working_dir = ""
+ self._events: list = []
+ self._event_idx = 0
+ self._lock = threading.Lock()
+ self._process: Optional[subprocess.Popen] = None
+ self._stop_flag = False
+ self._pause_event = threading.Event()
+ self._pause_event.set()
+ self._messages: List[Dict] = []
+
+ def start(self, task: str, working_dir: str, system_prompt: str = "") -> None:
+ self._working_dir = working_dir
+ self._state = "running"
+ self._stop_flag = False
+ self._events = []
+ self._event_idx = 0
+
+ thread = threading.Thread(target=self._run_sdk, args=(task,), daemon=True)
+ thread.start()
+
+ def _run_sdk(self, task: str):
+ """Run Claude Code via the Agent SDK."""
+ try:
+ # Try to use the Agent SDK
+ try:
+ from claude_agent_sdk import query
+ self._run_with_sdk(task)
+ return
+ except ImportError:
+ pass
+
+ # Fallback: spawn claude CLI directly
+ self._run_with_cli(task)
+
+ except Exception as e:
+ logger.exception("Claude SDK backend error")
+ self._state = "error"
+ self._emit(CodingAgentEventType.ERROR, {"message": str(e)})
+
+ def _run_with_sdk(self, task: str):
+ """Use the claude-agent-sdk Python package."""
+ import asyncio
+ from claude_agent_sdk import query
+
+ async def _run():
+ turn_index = 0
+ current_reasoning = []
+ current_tool_calls = []
+
+ async for message in query(prompt=task, options={"cwd": self._working_dir}):
+ if self._stop_flag:
+ break
+
+ self._pause_event.wait()
+ if self._stop_flag:
+ break
+
+ msg_type = message.get("type", "")
+
+ if msg_type == "assistant":
+ # Assistant reasoning
+ content = message.get("message", {}).get("content", [])
+ for block in content:
+ if isinstance(block, dict):
+ if block.get("type") == "text":
+ text = block.get("text", "")
+ current_reasoning.append(text)
+ self._emit(CodingAgentEventType.THINKING, {
+ "turn_index": turn_index,
+ "text": text,
+ })
+ elif block.get("type") == "tool_use":
+ tool_name = block.get("name", "unknown")
+ tool_input = block.get("input", {})
+ self._emit(CodingAgentEventType.TOOL_CALL_START, {
+ "turn_index": turn_index,
+ "tool": tool_name,
+ "input": tool_input,
+ })
+
+ elif msg_type == "result":
+ # Tool result
+ result = message.get("result", "")
+ tool_name = message.get("tool_name", "")
+ tool_input = message.get("tool_input", {})
+ output_type = self._classify_output_type(tool_name)
+
+ tc = {
+ "tool": tool_name,
+ "input": tool_input,
+ "output": str(result),
+ "output_type": output_type,
+ }
+ current_tool_calls.append(tc)
+
+ self._emit(CodingAgentEventType.TOOL_CALL_END, {
+ "turn_index": turn_index,
+ "tool_index": len(current_tool_calls) - 1,
+ **tc,
+ })
+
+ # Check if turn ended (no more tool calls pending)
+ if msg_type in ("assistant",) and not message.get("message", {}).get("content", []):
+ if current_reasoning or current_tool_calls:
+ self._emit(CodingAgentEventType.TURN_END, {
+ "turn_index": turn_index,
+ "content": "\n".join(current_reasoning),
+ "tool_calls": current_tool_calls,
+ })
+ turn_index += 1
+ current_reasoning = []
+ current_tool_calls = []
+
+ # Final turn
+ if current_reasoning or current_tool_calls:
+ self._emit(CodingAgentEventType.TURN_END, {
+ "turn_index": turn_index,
+ "content": "\n".join(current_reasoning),
+ "tool_calls": current_tool_calls,
+ })
+ turn_index += 1
+
+ self._state = "completed"
+ self._emit(CodingAgentEventType.COMPLETE, {"total_turns": turn_index})
+
+ asyncio.run(_run())
+
+ def _run_with_cli(self, task: str):
+ """Fallback: spawn claude CLI as subprocess."""
+ try:
+ self._process = subprocess.Popen(
+ ["claude", "--bare", "-p", task],
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ cwd=self._working_dir,
+ text=True,
+ )
+ except FileNotFoundError:
+ self._emit(CodingAgentEventType.ERROR, {
+ "message": "claude CLI not found. Install Claude Code or claude-agent-sdk."
+ })
+ self._state = "error"
+ return
+
+ turn_index = 0
+ output_lines = []
+
+ for line in self._process.stdout:
+ if self._stop_flag:
+ break
+
+ line = line.strip()
+ if not line:
+ continue
+
+ # Try to parse as JSON-lines
+ try:
+ msg = json.loads(line)
+ msg_type = msg.get("type", "")
+
+ if msg_type in ("text", "content"):
+ text = msg.get("text", msg.get("content", ""))
+ self._emit(CodingAgentEventType.THINKING, {
+ "turn_index": turn_index,
+ "text": text,
+ })
+ elif msg_type == "tool_use":
+ self._emit(CodingAgentEventType.TOOL_CALL_START, {
+ "turn_index": turn_index,
+ "tool": msg.get("name", ""),
+ "input": msg.get("input", {}),
+ })
+ elif msg_type == "tool_result":
+ self._emit(CodingAgentEventType.TOOL_CALL_END, {
+ "turn_index": turn_index,
+ "tool": msg.get("name", ""),
+ "input": msg.get("input", {}),
+ "output": msg.get("output", ""),
+ "output_type": self._classify_output_type(msg.get("name", "")),
+ })
+ except json.JSONDecodeError:
+ # Plain text output
+ output_lines.append(line)
+
+ self._process.wait()
+ self._state = "completed"
+ self._emit(CodingAgentEventType.COMPLETE, {"total_turns": turn_index + 1})
+
+ def _classify_output_type(self, tool_name: str) -> str:
+ name = (tool_name or "").lower()
+ if name in ("bash", "terminal", "shell"):
+ return "terminal"
+ if name in ("edit", "replace"):
+ return "diff"
+ return "code"
+
+ def _emit(self, event_type: CodingAgentEventType, data: dict):
+ event = CodingAgentEvent(event_type=event_type, timestamp=time.time(), data=data)
+ with self._lock:
+ self._events.append(event)
+
+ def get_events(self) -> Iterator[CodingAgentEvent]:
+ while True:
+ with self._lock:
+ if self._event_idx < len(self._events):
+ event = self._events[self._event_idx]
+ self._event_idx += 1
+ yield event
+ if event.event_type in (CodingAgentEventType.COMPLETE, CodingAgentEventType.ERROR):
+ return
+ continue
+ if self._state in ("completed", "error"):
+ return
+ time.sleep(0.1)
+
+ def pause(self) -> None:
+ self._pause_event.clear()
+ self._state = "paused"
+
+ def resume(self) -> None:
+ self._state = "running"
+ self._pause_event.set()
+
+ def inject_instruction(self, text: str) -> None:
+ # Claude SDK doesn't support mid-session instruction injection
+ # This would require restarting with appended context
+ logger.warning("inject_instruction not fully supported in Claude SDK backend")
+
+ def stop(self) -> None:
+ self._stop_flag = True
+ self._pause_event.set()
+ if self._process and self._process.poll() is None:
+ self._process.terminate()
+ self._state = "completed"
+
+ def get_conversation_history(self) -> List[Dict]:
+ return list(self._messages)
+
+ def get_state(self) -> str:
+ return self._state
diff --git a/potato/coding_agent_backends/ollama_backend.py b/potato/coding_agent_backends/ollama_backend.py
new file mode 100644
index 0000000000000000000000000000000000000000..a77efeda33459824a88a30edab6872302505f525
--- /dev/null
+++ b/potato/coding_agent_backends/ollama_backend.py
@@ -0,0 +1,256 @@
+"""
+Ollama Tool Use Backend
+
+Custom agent loop using Ollama's API with local models.
+No API key required โ fully local execution.
+Uses Ollama's tool/function calling support.
+"""
+
+import json
+import logging
+import time
+import threading
+from typing import Any, Dict, Iterator, List, Optional
+
+import requests as http_requests
+
+from ..coding_agent_backend import (
+ CodingAgentBackend,
+ CodingAgentEvent,
+ CodingAgentEventType,
+ CODING_TOOLS_OLLAMA,
+ CODING_TOOLS,
+ execute_tool,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class OllamaToolUseBackend(CodingAgentBackend):
+ """Agent loop using Ollama API with tool/function calling."""
+
+ def __init__(self, config: dict):
+ self._config = config
+ ai = config.get("ai_config", {})
+ self._model = ai.get("model", "qwen2.5-coder:14b")
+ self._base_url = ai.get("base_url", "http://localhost:11434")
+ self._max_tokens = ai.get("max_tokens", 8192)
+ self._temperature = ai.get("temperature", 0.3)
+ self._max_turns = config.get("max_turns", 50)
+
+ self._state = "idle"
+ self._working_dir = ""
+ self._messages: List[Dict] = []
+ self._system_prompt = ""
+ self._events: list = []
+ self._event_idx = 0
+ self._pause_event = threading.Event()
+ self._pause_event.set()
+ self._stop_flag = False
+ self._instruction_queue: list = []
+ self._lock = threading.Lock()
+
+ def start(self, task: str, working_dir: str, system_prompt: str = "") -> None:
+ self._working_dir = working_dir
+ self._system_prompt = system_prompt or (
+ "You are a coding agent. You have access to tools for reading, "
+ "editing, and creating files, running bash commands, and searching code. "
+ "Use these tools to complete the task. When you are done, stop calling tools "
+ "and summarize what you did."
+ )
+ self._messages = [
+ {"role": "system", "content": self._system_prompt},
+ {"role": "user", "content": task},
+ ]
+ self._state = "running"
+ self._stop_flag = False
+ self._events = []
+ self._event_idx = 0
+
+ thread = threading.Thread(target=self._run_loop, daemon=True)
+ thread.start()
+
+ def _run_loop(self):
+ """Main agent loop using Ollama's chat API with tools."""
+ turn_index = 0
+ try:
+ while not self._stop_flag and turn_index < self._max_turns:
+ self._pause_event.wait()
+ if self._stop_flag:
+ break
+
+ # Check for injected instructions
+ with self._lock:
+ if self._instruction_queue:
+ instruction = self._instruction_queue.pop(0)
+ self._messages.append({"role": "user", "content": instruction})
+
+ self._emit(CodingAgentEventType.THINKING, {
+ "turn_index": turn_index,
+ "text": "Thinking...",
+ })
+
+ # Query Ollama
+ try:
+ resp = http_requests.post(
+ f"{self._base_url}/api/chat",
+ json={
+ "model": self._model,
+ "messages": self._messages,
+ "tools": CODING_TOOLS_OLLAMA,
+ "stream": False,
+ "options": {
+ "num_predict": self._max_tokens,
+ "temperature": self._temperature,
+ },
+ },
+ timeout=120,
+ )
+ resp.raise_for_status()
+ result = resp.json()
+ except Exception as e:
+ self._emit(CodingAgentEventType.ERROR, {"message": f"Ollama error: {e}"})
+ self._state = "error"
+ return
+
+ message = result.get("message", {})
+ content = message.get("content", "")
+ tool_calls_raw = message.get("tool_calls", [])
+
+ # Emit thinking
+ if content:
+ self._emit(CodingAgentEventType.THINKING, {
+ "turn_index": turn_index,
+ "text": content,
+ })
+
+ # Add assistant message to history
+ self._messages.append(message)
+
+ # Execute tool calls
+ tool_calls = []
+ for tc_raw in tool_calls_raw:
+ if self._stop_flag:
+ break
+ self._pause_event.wait()
+ if self._stop_flag:
+ break
+
+ func = tc_raw.get("function", {})
+ tool_name = func.get("name", "unknown")
+ tool_input = func.get("arguments", {})
+ if isinstance(tool_input, str):
+ try:
+ tool_input = json.loads(tool_input)
+ except json.JSONDecodeError:
+ tool_input = {"command": tool_input}
+
+ self._emit(CodingAgentEventType.TOOL_CALL_START, {
+ "turn_index": turn_index,
+ "tool": tool_name,
+ "input": tool_input,
+ })
+
+ output = execute_tool(tool_name, tool_input, self._working_dir)
+ output_type = self._classify_output_type(tool_name)
+
+ tc = {
+ "tool": tool_name,
+ "input": tool_input,
+ "output": output,
+ "output_type": output_type,
+ }
+ tool_calls.append(tc)
+
+ self._emit(CodingAgentEventType.TOOL_CALL_END, {
+ "turn_index": turn_index,
+ "tool_index": len(tool_calls) - 1,
+ **tc,
+ })
+
+ # Add tool result to messages (Ollama format)
+ self._messages.append({
+ "role": "tool",
+ "content": output,
+ })
+
+ # Emit turn_end
+ self._emit(CodingAgentEventType.TURN_END, {
+ "turn_index": turn_index,
+ "content": content,
+ "tool_calls": tool_calls,
+ })
+
+ turn_index += 1
+
+ # If no tool calls, agent is done
+ if not tool_calls_raw:
+ break
+
+ self._state = "completed"
+ self._emit(CodingAgentEventType.COMPLETE, {"total_turns": turn_index})
+
+ except Exception as e:
+ logger.exception("Ollama agent loop error")
+ self._state = "error"
+ self._emit(CodingAgentEventType.ERROR, {"message": str(e)})
+
+ def _classify_output_type(self, tool_name: str) -> str:
+ name = tool_name.lower()
+ if name in ("bash", "terminal", "shell"):
+ return "terminal"
+ if name in ("edit", "replace"):
+ return "diff"
+ return "code"
+
+ def _emit(self, event_type: CodingAgentEventType, data: dict):
+ event = CodingAgentEvent(event_type=event_type, timestamp=time.time(), data=data)
+ with self._lock:
+ self._events.append(event)
+
+ def get_events(self) -> Iterator[CodingAgentEvent]:
+ while True:
+ with self._lock:
+ if self._event_idx < len(self._events):
+ event = self._events[self._event_idx]
+ self._event_idx += 1
+ yield event
+ if event.event_type in (CodingAgentEventType.COMPLETE, CodingAgentEventType.ERROR):
+ return
+ continue
+ if self._state in ("completed", "error"):
+ return
+ time.sleep(0.1)
+
+ def pause(self) -> None:
+ self._pause_event.clear()
+ self._state = "paused"
+
+ def resume(self) -> None:
+ self._state = "running"
+ self._pause_event.set()
+
+ def inject_instruction(self, text: str) -> None:
+ with self._lock:
+ self._instruction_queue.append(text)
+
+ def stop(self) -> None:
+ self._stop_flag = True
+ self._pause_event.set()
+ self._state = "completed"
+
+ def get_conversation_history(self) -> List[Dict]:
+ with self._lock:
+ return list(self._messages)
+
+ def get_state(self) -> str:
+ return self._state
+
+ def truncate_history(self, to_step: int) -> None:
+ with self._lock:
+ # Keep system + initial user + 2 messages per turn
+ keep = 2 + (to_step * 2)
+ self._messages = self._messages[:keep]
+ new_events = [e for e in self._events if e.data.get("turn_index", -1) < to_step or e.data.get("turn_index", -1) == -1]
+ self._events = new_events
+ self._event_idx = min(self._event_idx, len(self._events))
diff --git a/potato/coding_agent_backends/openai_backend.py b/potato/coding_agent_backends/openai_backend.py
new file mode 100644
index 0000000000000000000000000000000000000000..fd5abf2dc115d24ac633e835d38525d28a84d1d4
--- /dev/null
+++ b/potato/coding_agent_backends/openai_backend.py
@@ -0,0 +1,315 @@
+"""
+OpenAI Tool Use Backend
+
+Custom coding-agent loop using any OpenAI-compatible chat-completions
+server (OpenAI, vLLM, llama.cpp, etc.) with function/tool calling.
+
+vLLM ignores the API key but the OpenAI SDK rejects an empty string, so
+a non-empty placeholder is substituted for local servers. A configured
+base_url is honored and normalized to the ".../v1" form the SDK expects
+(accepts either the server root or an explicit "/v1" base_url).
+"""
+
+import json
+import logging
+import time
+import threading
+from typing import Dict, Iterator, List
+
+from ..coding_agent_backend import (
+ CodingAgentBackend,
+ CodingAgentEvent,
+ CodingAgentEventType,
+ CODING_TOOLS,
+ execute_tool,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _to_openai_tools(tools: list) -> list:
+ """CODING_TOOLS is in Anthropic shape ({name, description,
+ input_schema}); the OpenAI/vLLM API needs
+ {type:"function", function:{name, description, parameters}}."""
+ converted = []
+ for t in tools:
+ if t.get("type") == "function" and "function" in t:
+ converted.append(t) # already OpenAI shape
+ continue
+ converted.append({
+ "type": "function",
+ "function": {
+ "name": t["name"],
+ "description": t.get("description", ""),
+ "parameters": t.get("input_schema", {"type": "object", "properties": {}}),
+ },
+ })
+ return converted
+
+
+def _normalize_base_url(raw: str) -> str:
+ """The OpenAI SDK appends '/chat/completions' to base_url, so it must
+ end at the '/v1' root. Accept either the server root or a '/v1' URL."""
+ if not raw:
+ return raw
+ u = raw.rstrip("/")
+ if not u.endswith("/v1"):
+ u = u + "/v1"
+ return u
+
+
+class OpenAIToolUseBackend(CodingAgentBackend):
+ """Agent loop using an OpenAI-compatible API with tool calling."""
+
+ def __init__(self, config: dict):
+ self._config = config
+ ai = config.get("ai_config", {})
+ self._model = ai.get("model", "gpt-4o-mini")
+ self._base_url = _normalize_base_url(ai.get("base_url", "")) or None
+ # vLLM/local servers ignore the key; SDK requires non-empty.
+ import os
+ self._api_key = (
+ ai.get("api_key")
+ or os.environ.get("OPENAI_API_KEY")
+ or "EMPTY"
+ )
+ self._max_tokens = ai.get("max_tokens", 8192)
+ self._temperature = ai.get("temperature", 0.3)
+ self._timeout = ai.get("timeout", 120)
+ self._max_turns = config.get("max_turns", 50)
+ self._tools = _to_openai_tools(CODING_TOOLS)
+
+ self._state = "idle"
+ self._working_dir = ""
+ self._messages: List[Dict] = []
+ self._system_prompt = ""
+ self._events: list = []
+ self._event_idx = 0
+ self._pause_event = threading.Event()
+ self._pause_event.set()
+ self._stop_flag = False
+ self._instruction_queue: list = []
+ self._lock = threading.Lock()
+ self._client = None
+
+ def _get_client(self):
+ if self._client is not None:
+ return self._client
+ from openai import OpenAI
+
+ kwargs = {"api_key": self._api_key, "timeout": self._timeout}
+ if self._base_url:
+ kwargs["base_url"] = self._base_url
+ self._client = OpenAI(**kwargs)
+ return self._client
+
+ def start(self, task: str, working_dir: str, system_prompt: str = "") -> None:
+ self._working_dir = working_dir
+ self._system_prompt = system_prompt or (
+ "You are a coding agent. You have access to tools for reading, "
+ "editing, and creating files, running bash commands, and searching code. "
+ "Use these tools to complete the task. When you are done, stop calling tools "
+ "and summarize what you did."
+ )
+ self._messages = [
+ {"role": "system", "content": self._system_prompt},
+ {"role": "user", "content": task},
+ ]
+ self._state = "running"
+ self._stop_flag = False
+ self._events = []
+ self._event_idx = 0
+
+ thread = threading.Thread(target=self._run_loop, daemon=True)
+ thread.start()
+
+ def _run_loop(self):
+ """Main agent loop using the OpenAI chat API with tools."""
+ turn_index = 0
+ try:
+ client = self._get_client()
+ while not self._stop_flag and turn_index < self._max_turns:
+ self._pause_event.wait()
+ if self._stop_flag:
+ break
+
+ with self._lock:
+ if self._instruction_queue:
+ instruction = self._instruction_queue.pop(0)
+ self._messages.append({"role": "user", "content": instruction})
+
+ self._emit(CodingAgentEventType.THINKING, {
+ "turn_index": turn_index,
+ "text": "Thinking...",
+ })
+
+ try:
+ resp = client.chat.completions.create(
+ model=self._model,
+ messages=self._messages,
+ tools=self._tools,
+ tool_choice="auto",
+ max_tokens=self._max_tokens,
+ temperature=self._temperature,
+ )
+ except Exception as e:
+ # Includes models/servers that don't support tools --
+ # surface a clear error instead of stalling the loop.
+ self._emit(CodingAgentEventType.ERROR, {
+ "message": f"OpenAI-compatible request failed: {e}"
+ })
+ self._state = "error"
+ return
+
+ choice = resp.choices[0].message
+ content = choice.content or ""
+ tool_calls_raw = choice.tool_calls or []
+
+ if content:
+ self._emit(CodingAgentEventType.THINKING, {
+ "turn_index": turn_index,
+ "text": content,
+ })
+
+ # Append the assistant message verbatim (must include
+ # tool_calls so the following tool messages pair by id).
+ try:
+ assistant_msg = choice.model_dump(exclude_none=True)
+ except Exception:
+ assistant_msg = {"role": "assistant", "content": content}
+ self._messages.append(assistant_msg)
+
+ tool_calls = []
+ for tc_raw in tool_calls_raw:
+ if self._stop_flag:
+ break
+ self._pause_event.wait()
+ if self._stop_flag:
+ break
+
+ fn = tc_raw.function
+ tool_name = fn.name or "unknown"
+ raw_args = fn.arguments
+ if isinstance(raw_args, str):
+ try:
+ tool_input = json.loads(raw_args) if raw_args else {}
+ except json.JSONDecodeError:
+ tool_input = {"command": raw_args}
+ elif isinstance(raw_args, dict):
+ tool_input = raw_args
+ else:
+ tool_input = {}
+
+ self._emit(CodingAgentEventType.TOOL_CALL_START, {
+ "turn_index": turn_index,
+ "tool": tool_name,
+ "input": tool_input,
+ })
+
+ output = execute_tool(tool_name, tool_input, self._working_dir)
+ output_type = self._classify_output_type(tool_name)
+
+ tc = {
+ "tool": tool_name,
+ "input": tool_input,
+ "output": output,
+ "output_type": output_type,
+ }
+ tool_calls.append(tc)
+
+ self._emit(CodingAgentEventType.TOOL_CALL_END, {
+ "turn_index": turn_index,
+ "tool_index": len(tool_calls) - 1,
+ **tc,
+ })
+
+ # OpenAI requires the tool result to reference the
+ # originating tool_call_id.
+ self._messages.append({
+ "role": "tool",
+ "tool_call_id": tc_raw.id,
+ "content": output,
+ })
+
+ self._emit(CodingAgentEventType.TURN_END, {
+ "turn_index": turn_index,
+ "content": content,
+ "tool_calls": tool_calls,
+ })
+
+ turn_index += 1
+
+ if not tool_calls_raw:
+ break
+
+ self._state = "completed"
+ self._emit(CodingAgentEventType.COMPLETE, {"total_turns": turn_index})
+
+ except Exception as e:
+ logger.exception("OpenAI agent loop error")
+ self._state = "error"
+ self._emit(CodingAgentEventType.ERROR, {"message": str(e)})
+
+ def _classify_output_type(self, tool_name: str) -> str:
+ name = tool_name.lower()
+ if name in ("bash", "terminal", "shell"):
+ return "terminal"
+ if name in ("edit", "replace"):
+ return "diff"
+ return "code"
+
+ def _emit(self, event_type: CodingAgentEventType, data: dict):
+ event = CodingAgentEvent(event_type=event_type, timestamp=time.time(), data=data)
+ with self._lock:
+ self._events.append(event)
+
+ def get_events(self) -> Iterator[CodingAgentEvent]:
+ while True:
+ with self._lock:
+ if self._event_idx < len(self._events):
+ event = self._events[self._event_idx]
+ self._event_idx += 1
+ yield event
+ if event.event_type in (CodingAgentEventType.COMPLETE, CodingAgentEventType.ERROR):
+ return
+ continue
+ if self._state in ("completed", "error"):
+ return
+ time.sleep(0.1)
+
+ def pause(self) -> None:
+ self._pause_event.clear()
+ self._state = "paused"
+
+ def resume(self) -> None:
+ self._state = "running"
+ self._pause_event.set()
+
+ def inject_instruction(self, text: str) -> None:
+ with self._lock:
+ self._instruction_queue.append(text)
+
+ def stop(self) -> None:
+ self._stop_flag = True
+ self._pause_event.set()
+ self._state = "completed"
+
+ def get_conversation_history(self) -> List[Dict]:
+ with self._lock:
+ return list(self._messages)
+
+ def get_state(self) -> str:
+ return self._state
+
+ def truncate_history(self, to_step: int) -> None:
+ with self._lock:
+ # Best-effort: keep system + initial user, then drop events
+ # for turns >= to_step. (Messages vary per turn with tool
+ # calls; keep them since OpenAI needs tool_call_id pairing.)
+ new_events = [
+ e for e in self._events
+ if e.data.get("turn_index", -1) < to_step
+ or e.data.get("turn_index", -1) == -1
+ ]
+ self._events = new_events
+ self._event_idx = min(self._event_idx, len(self._events))
diff --git a/potato/coding_agent_branch.py b/potato/coding_agent_branch.py
new file mode 100644
index 0000000000000000000000000000000000000000..a11be10366a2b13a8249107cb8b1d3a085f560a5
--- /dev/null
+++ b/potato/coding_agent_branch.py
@@ -0,0 +1,209 @@
+"""
+Coding Agent Branch Manager
+
+Manages alternative trajectory branches for coding agent sessions.
+Each branch is backed by a git branch, enabling independent file states
+and conversation histories.
+
+Branch model:
+ main โโโโโโโโโโโโโโโโโโโโ (original trajectory)
+ โ
+ โโโ branch-1 โโโโโโโโโ (replayed with new instructions)
+ โ
+ โโโ branch-2 โโโโโโ (edited action)
+"""
+
+import logging
+import os
+import subprocess
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TrajectoryBranch:
+ """A single branch in the trajectory tree."""
+ branch_id: str
+ parent_branch_id: Optional[str]
+ branch_point_step: Optional[int] # step where this diverges from parent
+ turns: List[Dict[str, Any]]
+ git_branch: str
+ status: str = "active" # active, completed, abandoned
+ created_at: float = 0.0
+ instructions: Optional[str] = None
+ edited_actions: Optional[List[Dict]] = None
+
+ def to_dict(self) -> dict:
+ return {
+ "branch_id": self.branch_id,
+ "parent_branch_id": self.parent_branch_id,
+ "branch_point_step": self.branch_point_step,
+ "turns": self.turns,
+ "git_branch": self.git_branch,
+ "status": self.status,
+ "created_at": self.created_at,
+ "instructions": self.instructions,
+ "edited_actions": self.edited_actions,
+ "turn_count": len(self.turns),
+ }
+
+
+class BranchManager:
+ """Manages trajectory branches for a coding agent session."""
+
+ def __init__(self, session_id: str, working_dir: str):
+ self._session_id = session_id
+ self._working_dir = os.path.abspath(working_dir)
+ self._branches: Dict[str, TrajectoryBranch] = {}
+ self._active_branch_id: Optional[str] = None
+
+ # Create the main branch
+ main = TrajectoryBranch(
+ branch_id="main",
+ parent_branch_id=None,
+ branch_point_step=None,
+ turns=[],
+ git_branch=f"potato-agent-{session_id[:12]}",
+ created_at=time.time(),
+ )
+ self._branches["main"] = main
+ self._active_branch_id = "main"
+
+ @property
+ def active_branch(self) -> TrajectoryBranch:
+ return self._branches[self._active_branch_id]
+
+ @property
+ def active_branch_id(self) -> str:
+ return self._active_branch_id
+
+ def create_branch(self, parent_branch_id: str, branch_point_step: int,
+ instructions: Optional[str] = None,
+ edited_actions: Optional[List[Dict]] = None) -> TrajectoryBranch:
+ """Create a new branch from a parent at a given step.
+
+ Args:
+ parent_branch_id: ID of the parent branch
+ branch_point_step: Step index where the branch diverges
+ instructions: Optional user instructions for the new branch
+ edited_actions: Optional modified tool calls to execute
+
+ Returns:
+ The new TrajectoryBranch
+ """
+ parent = self._branches.get(parent_branch_id)
+ if not parent:
+ raise ValueError(f"Parent branch '{parent_branch_id}' not found")
+
+ branch_id = f"branch-{len(self._branches)}"
+ git_branch = f"potato-agent-{self._session_id[:8]}-{branch_id}"
+
+ # Create git branch from parent's state at branch_point_step
+ try:
+ # First, ensure we're on the parent branch
+ self._run_git("checkout", parent.git_branch)
+
+ # Find the commit at branch_point_step
+ # We use git log to find commits with [potato] step=N
+ log = self._run_git("log", "--oneline", "--all")
+ target_commit = None
+ for line in log.strip().split("\n"):
+ if f"step={branch_point_step}" in line:
+ target_commit = line.split()[0]
+ break
+
+ if target_commit:
+ self._run_git("checkout", "-b", git_branch, target_commit)
+ else:
+ # Fallback: branch from current HEAD
+ self._run_git("checkout", "-b", git_branch)
+ logger.warning(f"Could not find commit for step {branch_point_step}, branching from HEAD")
+
+ except subprocess.CalledProcessError as e:
+ logger.error(f"Failed to create git branch: {e}")
+ # Create branch without git backing
+ git_branch = parent.git_branch
+
+ # Copy turns up to branch point
+ branch_turns = list(parent.turns[:branch_point_step + 1])
+
+ branch = TrajectoryBranch(
+ branch_id=branch_id,
+ parent_branch_id=parent_branch_id,
+ branch_point_step=branch_point_step,
+ turns=branch_turns,
+ git_branch=git_branch,
+ created_at=time.time(),
+ instructions=instructions,
+ edited_actions=edited_actions,
+ )
+ self._branches[branch_id] = branch
+ self._active_branch_id = branch_id
+
+ logger.info(f"Created branch {branch_id} from {parent_branch_id} at step {branch_point_step}")
+ return branch
+
+ def switch_branch(self, branch_id: str) -> bool:
+ """Switch to a different branch."""
+ if branch_id not in self._branches:
+ return False
+
+ branch = self._branches[branch_id]
+
+ try:
+ self._run_git("checkout", branch.git_branch)
+ except subprocess.CalledProcessError as e:
+ logger.warning(f"Failed to switch git branch: {e}")
+
+ self._active_branch_id = branch_id
+ logger.info(f"Switched to branch {branch_id}")
+ return True
+
+ def add_turn_to_active(self, turn: Dict[str, Any]) -> None:
+ """Add a turn to the active branch."""
+ self.active_branch.turns.append(turn)
+
+ def get_branch(self, branch_id: str) -> Optional[TrajectoryBranch]:
+ return self._branches.get(branch_id)
+
+ def list_branches(self) -> List[dict]:
+ return [b.to_dict() for b in self._branches.values()]
+
+ def get_branch_tree(self) -> dict:
+ """Return tree structure for UI rendering."""
+ tree = {}
+ for bid, branch in self._branches.items():
+ tree[bid] = {
+ "branch_id": bid,
+ "parent": branch.parent_branch_id,
+ "branch_point": branch.branch_point_step,
+ "turns": len(branch.turns),
+ "status": branch.status,
+ "instructions": branch.instructions,
+ "is_active": bid == self._active_branch_id,
+ }
+ return tree
+
+ def save_all(self) -> dict:
+ """Serialize all branches for trace export."""
+ return {
+ bid: branch.to_dict()
+ for bid, branch in self._branches.items()
+ }
+
+ def _run_git(self, *args) -> str:
+ result = subprocess.run(
+ ["git"] + list(args),
+ cwd=self._working_dir,
+ capture_output=True, text=True, timeout=30,
+ )
+ if result.returncode != 0:
+ raise subprocess.CalledProcessError(
+ result.returncode, ["git"] + list(args),
+ output=result.stdout, stderr=result.stderr,
+ )
+ return result.stdout
diff --git a/potato/coding_agent_checkpoint.py b/potato/coding_agent_checkpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..531da4bc062cf59021d70da185b218b954acc596
--- /dev/null
+++ b/potato/coding_agent_checkpoint.py
@@ -0,0 +1,287 @@
+"""
+Coding Agent Checkpoint Manager
+
+Git-based checkpointing for coding agent sessions. Creates lightweight
+commits after each file-modifying tool call, enabling rollback to any
+previous step.
+
+Uses a dedicated git branch (potato-agent-) to avoid
+interfering with the user's branches.
+"""
+
+import logging
+import os
+import subprocess
+import time
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Checkpoint:
+ """A snapshot of the working directory state."""
+ checkpoint_id: str # git commit hash
+ step_index: int
+ tool_name: str
+ description: str
+ timestamp: float
+ files_changed: List[str] = field(default_factory=list)
+
+ def to_dict(self) -> dict:
+ return {
+ "checkpoint_id": self.checkpoint_id,
+ "step_index": self.step_index,
+ "tool_name": self.tool_name,
+ "description": self.description,
+ "timestamp": self.timestamp,
+ "files_changed": self.files_changed,
+ }
+
+
+class CheckpointManager:
+ """Manages git-based checkpoints for a coding agent session."""
+
+ def __init__(self, working_dir: str, session_id: str):
+ self._working_dir = os.path.abspath(working_dir)
+ self._session_id = session_id
+ self._branch_name = f"potato-agent-{session_id[:12]}"
+ self._checkpoints: List[Checkpoint] = []
+ self._initialized = False
+
+ @property
+ def checkpoints(self) -> List[Checkpoint]:
+ return list(self._checkpoints)
+
+ def init(self) -> bool:
+ """Initialize git repo and create session branch.
+
+ Returns True if initialization succeeded.
+ """
+ if self._initialized:
+ return True
+
+ # Ensure git repo exists
+ if not self._is_git_repo():
+ try:
+ self._run_git("init")
+ self._run_git("add", "-A")
+ self._run_git("commit", "--allow-empty", "-m", "[potato] init")
+ except Exception as e:
+ logger.warning(f"Failed to init git repo: {e}")
+ return False
+
+ # Create session branch from current HEAD
+ try:
+ current_branch = self._run_git("rev-parse", "--abbrev-ref", "HEAD").strip()
+ self._run_git("checkout", "-b", self._branch_name)
+ except subprocess.CalledProcessError:
+ # Branch might already exist (session restart)
+ try:
+ self._run_git("checkout", self._branch_name)
+ except subprocess.CalledProcessError as e:
+ logger.warning(f"Failed to create/checkout session branch: {e}")
+ return False
+
+ # Create initial checkpoint
+ try:
+ self._run_git("add", "-A")
+ self._run_git("commit", "--allow-empty", "-m",
+ f"[potato] session start {self._session_id[:8]}")
+ commit_hash = self._get_head_hash()
+ self._checkpoints.append(Checkpoint(
+ checkpoint_id=commit_hash,
+ step_index=-1,
+ tool_name="init",
+ description="Session start",
+ timestamp=time.time(),
+ ))
+ except Exception as e:
+ logger.warning(f"Failed to create initial checkpoint: {e}")
+
+ self._initialized = True
+ logger.info(f"CheckpointManager initialized on branch {self._branch_name}")
+ return True
+
+ def create_checkpoint(self, step_index: int, tool_name: str,
+ description: str = "") -> Optional[str]:
+ """Create a checkpoint after a tool execution.
+
+ Returns the commit hash, or None if no changes to commit.
+ """
+ if not self._initialized:
+ if not self.init():
+ return None
+
+ try:
+ # Stage all changes
+ self._run_git("add", "-A")
+
+ # Check if there are changes to commit
+ status = self._run_git("status", "--porcelain")
+ if not status.strip():
+ # No changes, but still record the checkpoint for rollback
+ commit_hash = self._get_head_hash()
+ self._checkpoints.append(Checkpoint(
+ checkpoint_id=commit_hash,
+ step_index=step_index,
+ tool_name=tool_name,
+ description=description or f"Step {step_index}: {tool_name}",
+ timestamp=time.time(),
+ ))
+ return commit_hash
+
+ # Get list of changed files
+ changed = [
+ line.split(None, 1)[-1].strip()
+ for line in status.strip().split("\n")
+ if line.strip()
+ ]
+
+ # Commit
+ msg = f"[potato] step={step_index} tool={tool_name}"
+ if description:
+ msg += f" {description}"
+ self._run_git("commit", "-m", msg)
+
+ commit_hash = self._get_head_hash()
+ checkpoint = Checkpoint(
+ checkpoint_id=commit_hash,
+ step_index=step_index,
+ tool_name=tool_name,
+ description=description or f"Step {step_index}: {tool_name}",
+ timestamp=time.time(),
+ files_changed=changed,
+ )
+ self._checkpoints.append(checkpoint)
+
+ logger.debug(f"Created checkpoint {commit_hash[:8]} at step {step_index}")
+ return commit_hash
+
+ except Exception as e:
+ logger.warning(f"Failed to create checkpoint: {e}")
+ return None
+
+ def rollback_to(self, step_index: int) -> bool:
+ """Rollback to the checkpoint at the given step index.
+
+ Returns True if rollback succeeded.
+ """
+ # Find the checkpoint
+ target = None
+ for cp in self._checkpoints:
+ if cp.step_index == step_index:
+ target = cp
+ break
+ if cp.step_index <= step_index:
+ target = cp # Use the latest checkpoint at or before step_index
+
+ if not target:
+ logger.warning(f"No checkpoint found at or before step {step_index}")
+ return False
+
+ try:
+ self._run_git("reset", "--hard", target.checkpoint_id)
+
+ # Truncate checkpoint list
+ self._checkpoints = [
+ cp for cp in self._checkpoints
+ if cp.step_index <= step_index
+ ]
+
+ logger.info(f"Rolled back to step {step_index} (commit {target.checkpoint_id[:8]})")
+ return True
+
+ except Exception as e:
+ logger.error(f"Rollback failed: {e}")
+ return False
+
+ def get_diff_between(self, from_step: int, to_step: int) -> str:
+ """Get the git diff between two checkpoints."""
+ from_cp = self._find_checkpoint(from_step)
+ to_cp = self._find_checkpoint(to_step)
+ if not from_cp or not to_cp:
+ return ""
+
+ try:
+ return self._run_git("diff", from_cp.checkpoint_id, to_cp.checkpoint_id)
+ except Exception:
+ return ""
+
+ def get_diff_since(self, step_index: int) -> str:
+ """Get the diff from a checkpoint to current HEAD."""
+ cp = self._find_checkpoint(step_index)
+ if not cp:
+ return ""
+ try:
+ return self._run_git("diff", cp.checkpoint_id, "HEAD")
+ except Exception:
+ return ""
+
+ def get_file_at(self, step_index: int, file_path: str) -> Optional[str]:
+ """Get file contents at a specific checkpoint."""
+ cp = self._find_checkpoint(step_index)
+ if not cp:
+ return None
+ try:
+ return self._run_git("show", f"{cp.checkpoint_id}:{file_path}")
+ except Exception:
+ return None
+
+ def list_checkpoints(self) -> List[dict]:
+ """Return checkpoint metadata as list of dicts."""
+ return [cp.to_dict() for cp in self._checkpoints]
+
+ def cleanup(self) -> None:
+ """Clean up the session branch."""
+ if not self._initialized:
+ return
+
+ try:
+ # Switch back to the original branch
+ branches = self._run_git("branch", "--list").strip().split("\n")
+ main_branch = None
+ for b in branches:
+ name = b.strip().lstrip("* ")
+ if name and name != self._branch_name:
+ main_branch = name
+ break
+
+ if main_branch:
+ self._run_git("checkout", main_branch)
+ self._run_git("branch", "-D", self._branch_name)
+ logger.info(f"Cleaned up session branch {self._branch_name}")
+ except Exception as e:
+ logger.warning(f"Failed to clean up session branch: {e}")
+
+ def _find_checkpoint(self, step_index: int) -> Optional[Checkpoint]:
+ for cp in self._checkpoints:
+ if cp.step_index == step_index:
+ return cp
+ return None
+
+ def _is_git_repo(self) -> bool:
+ try:
+ self._run_git("rev-parse", "--git-dir")
+ return True
+ except (subprocess.CalledProcessError, FileNotFoundError):
+ return False
+
+ def _get_head_hash(self) -> str:
+ return self._run_git("rev-parse", "HEAD").strip()
+
+ def _run_git(self, *args) -> str:
+ result = subprocess.run(
+ ["git"] + list(args),
+ cwd=self._working_dir,
+ capture_output=True,
+ text=True,
+ timeout=30,
+ )
+ if result.returncode != 0:
+ raise subprocess.CalledProcessError(
+ result.returncode, ["git"] + list(args),
+ output=result.stdout, stderr=result.stderr,
+ )
+ return result.stdout
diff --git a/potato/coding_agent_runner.py b/potato/coding_agent_runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..befe3d0b7611a1485baa7911262e182e35ee8f17
--- /dev/null
+++ b/potato/coding_agent_runner.py
@@ -0,0 +1,440 @@
+"""
+Coding Agent Runner
+
+Manages the lifecycle of a coding agent session. Mirrors AgentRunner
+but adapted for terminal-based coding agents (no Playwright).
+
+State machine: IDLE โ RUNNING โ PAUSED โ COMPLETED โ ERROR
+Communication: SSE listener pattern (same as AgentRunner)
+"""
+
+import json
+import logging
+import os
+import threading
+import time
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Callable, Dict, List, Optional
+
+from .coding_agent_backend import (
+ CodingAgentBackend,
+ CodingAgentEvent,
+ CodingAgentEventType,
+ create_backend,
+)
+from .coding_agent_branch import BranchManager
+from .coding_agent_checkpoint import CheckpointManager
+from .coding_agent_sandbox import SandboxManager
+
+logger = logging.getLogger(__name__)
+
+
+class CodingAgentState(str, Enum):
+ IDLE = "idle"
+ RUNNING = "running"
+ PAUSED = "paused"
+ COMPLETED = "completed"
+ ERROR = "error"
+
+
+@dataclass
+class CodingAgentConfig:
+ """Configuration for a coding agent session."""
+ backend_type: str = "ollama_tool_use"
+ ai_config: Dict[str, Any] = field(default_factory=dict)
+ working_dir: str = "."
+ max_turns: int = 50
+ system_prompt: str = ""
+ sandbox_mode: str = "worktree"
+
+ @classmethod
+ def from_config(cls, config: dict) -> "CodingAgentConfig":
+ """Create from YAML config dict."""
+ live_config = config.get("live_coding_agent", {})
+ return cls(
+ backend_type=live_config.get("backend_type", "ollama_tool_use"),
+ ai_config=live_config.get("ai_config", {}),
+ working_dir=live_config.get("working_dir", "."),
+ max_turns=live_config.get("max_turns", 50),
+ system_prompt=live_config.get("system_prompt", ""),
+ sandbox_mode=live_config.get("sandbox_mode", "worktree"),
+ )
+
+
+class CodingAgentRunner:
+ """Manages a coding agent session with SSE event broadcasting."""
+
+ def __init__(self, session_id: str, config: CodingAgentConfig,
+ trace_dir: str = ""):
+ self.session_id = session_id
+ self.config = config
+ self.trace_dir = trace_dir
+
+ self._state = CodingAgentState.IDLE
+ self._state_lock = threading.Lock()
+ self._listeners: List[Callable] = []
+ self._listener_lock = threading.Lock()
+
+ self._backend: Optional[CodingAgentBackend] = None
+ self._sandbox: Optional[SandboxManager] = None
+ self._checkpoint_mgr: Optional[CheckpointManager] = None
+ self._branch_mgr: Optional[BranchManager] = None
+ self._structured_turns: List[Dict] = []
+ self._task_description = ""
+ self._started_at = 0.0
+ self._event_thread: Optional[threading.Thread] = None
+
+ @property
+ def state(self) -> CodingAgentState:
+ with self._state_lock:
+ return self._state
+
+ def _set_state(self, new_state: CodingAgentState):
+ with self._state_lock:
+ old = self._state
+ self._state = new_state
+ self._emit_event("state_change", {"old_state": old.value, "new_state": new_state.value})
+
+ # --- Listener/SSE pattern (mirrors AgentRunner) ---
+
+ def add_listener(self, callback: Callable) -> None:
+ with self._listener_lock:
+ self._listeners.append(callback)
+
+ def remove_listener(self, callback: Callable) -> None:
+ with self._listener_lock:
+ self._listeners = [l for l in self._listeners if l is not callback]
+
+ def _emit_event(self, event_type: str, data: dict):
+ with self._listener_lock:
+ for listener in self._listeners:
+ try:
+ listener(event_type, data)
+ except Exception:
+ pass
+
+ # --- Control methods ---
+
+ def start(self, task_description: str) -> None:
+ """Start the coding agent session."""
+ if self.state != CodingAgentState.IDLE:
+ raise RuntimeError(f"Cannot start in state {self.state}")
+
+ self._task_description = task_description
+ self._started_at = time.time()
+ self._structured_turns = []
+
+ # Create sandbox
+ self._sandbox = SandboxManager(
+ mode=self.config.sandbox_mode,
+ base_dir=os.path.abspath(self.config.working_dir),
+ )
+ working_dir = self._sandbox.create(self.session_id)
+
+ # Initialize checkpoint and branch managers
+ self._checkpoint_mgr = CheckpointManager(working_dir, self.session_id)
+ self._checkpoint_mgr.init()
+ self._branch_mgr = BranchManager(self.session_id, working_dir)
+
+ # Create backend
+ backend_config = {
+ "ai_config": self.config.ai_config,
+ "max_turns": self.config.max_turns,
+ }
+ self._backend = create_backend(self.config.backend_type, backend_config)
+
+ # Start the backend
+ self._backend.start(task_description, working_dir, self.config.system_prompt)
+ self._set_state(CodingAgentState.RUNNING)
+
+ # Start event consumer thread
+ self._event_thread = threading.Thread(target=self._consume_events, daemon=True)
+ self._event_thread.start()
+
+ self._emit_event("started", {
+ "session_id": self.session_id,
+ "task": task_description,
+ "backend": self.config.backend_type,
+ })
+
+ def pause(self) -> None:
+ if self._backend:
+ self._backend.pause()
+ self._set_state(CodingAgentState.PAUSED)
+
+ def resume(self) -> None:
+ if self._backend:
+ self._backend.resume()
+ self._set_state(CodingAgentState.RUNNING)
+
+ def inject_instruction(self, instruction: str) -> None:
+ if self._backend:
+ self._backend.inject_instruction(instruction)
+ self._emit_event("instruction_received", {"instruction": instruction})
+
+ def stop(self) -> None:
+ if self._backend:
+ self._backend.stop()
+ self._set_state(CodingAgentState.COMPLETED)
+ self._save_trace()
+
+ # --- Event consumption ---
+
+ def _consume_events(self):
+ """Consume events from the backend and broadcast via SSE."""
+ if not self._backend:
+ return
+
+ current_turn: Optional[Dict] = None
+
+ try:
+ for event in self._backend.get_events():
+ et = event.event_type
+ data = event.data
+
+ if et == CodingAgentEventType.THINKING:
+ self._emit_event("thinking", data)
+
+ elif et == CodingAgentEventType.TOOL_CALL_START:
+ self._emit_event("tool_call_start", data)
+
+ elif et == CodingAgentEventType.TOOL_CALL_END:
+ self._emit_event("tool_call", data)
+
+ # Create checkpoint for file-modifying tools
+ tool_name = data.get("tool", "")
+ if tool_name.lower() in ("edit", "write", "bash", "create", "replace"):
+ turn_idx = data.get("turn_index", len(self._structured_turns))
+ if self._checkpoint_mgr:
+ cp_id = self._checkpoint_mgr.create_checkpoint(
+ turn_idx, tool_name,
+ f"{tool_name} at step {turn_idx}",
+ )
+ if cp_id:
+ self._emit_event("checkpoint", {
+ "step_index": turn_idx,
+ "checkpoint_id": cp_id,
+ "tool": tool_name,
+ })
+
+ elif et == CodingAgentEventType.TURN_END:
+ # Accumulate into structured_turns
+ turn = {
+ "role": "assistant",
+ "content": data.get("content", ""),
+ "tool_calls": data.get("tool_calls", []),
+ }
+ self._structured_turns.append(turn)
+ self._emit_event("turn_end", {
+ "turn_index": data.get("turn_index", len(self._structured_turns) - 1),
+ **turn,
+ })
+
+ elif et == CodingAgentEventType.ERROR:
+ self._set_state(CodingAgentState.ERROR)
+ self._emit_event("error", data)
+
+ elif et == CodingAgentEventType.COMPLETE:
+ self._set_state(CodingAgentState.COMPLETED)
+ self._emit_event("complete", {
+ "total_turns": len(self._structured_turns),
+ })
+ self._save_trace()
+
+ except Exception as e:
+ logger.exception("Error consuming backend events")
+ self._set_state(CodingAgentState.ERROR)
+ self._emit_event("error", {"message": str(e)})
+
+ # --- Checkpoint + Rollback ---
+
+ def rollback_to_step(self, step_index: int) -> bool:
+ """Rollback files and conversation to the given step.
+
+ Pauses the agent, restores files, truncates history.
+ Returns True on success.
+ """
+ if self._backend:
+ self._backend.pause()
+ self._set_state(CodingAgentState.PAUSED)
+
+ # Rollback files
+ if self._checkpoint_mgr:
+ if not self._checkpoint_mgr.rollback_to(step_index):
+ return False
+
+ # Truncate structured turns
+ self._structured_turns = self._structured_turns[:step_index + 1]
+
+ # Truncate backend conversation history
+ if self._backend:
+ self._backend.truncate_history(step_index + 1)
+
+ self._emit_event("rollback", {
+ "step_index": step_index,
+ "remaining_turns": len(self._structured_turns),
+ })
+ return True
+
+ def get_checkpoints(self) -> List[dict]:
+ """Return list of checkpoint metadata."""
+ if self._checkpoint_mgr:
+ return self._checkpoint_mgr.list_checkpoints()
+ return []
+
+ def replay_from_step(self, step_index: int,
+ instructions: Optional[str] = None,
+ edited_actions: Optional[List[Dict]] = None) -> Optional[str]:
+ """Create a new branch from step_index and resume the agent.
+
+ Args:
+ step_index: Step to branch from
+ instructions: Optional new instructions to inject
+ edited_actions: Optional modified tool calls to execute first
+
+ Returns:
+ The new branch_id, or None on failure.
+ """
+ # Pause current execution
+ if self._backend:
+ self._backend.pause()
+ self._set_state(CodingAgentState.PAUSED)
+
+ # Create a new branch
+ if not self._branch_mgr:
+ return None
+
+ active_id = self._branch_mgr.active_branch_id
+ try:
+ branch = self._branch_mgr.create_branch(
+ active_id, step_index,
+ instructions=instructions,
+ edited_actions=edited_actions,
+ )
+ except Exception as e:
+ logger.error(f"Failed to create branch: {e}")
+ return None
+
+ # Rollback files to the branch point
+ if self._checkpoint_mgr:
+ self._checkpoint_mgr.rollback_to(step_index)
+
+ # Truncate structured turns and backend history
+ self._structured_turns = list(branch.turns)
+ if self._backend:
+ self._backend.truncate_history(step_index + 1)
+
+ # Execute edited actions if provided
+ if edited_actions:
+ from .coding_agent_backend import execute_tool
+ working_dir = self._sandbox.working_dir if self._sandbox else self.config.working_dir
+ for action in edited_actions:
+ tool_name = action.get("tool", "")
+ tool_input = action.get("input", {})
+ output = execute_tool(tool_name, tool_input, working_dir)
+ action["output"] = output
+
+ # Inject instructions if provided
+ if instructions and self._backend:
+ self._backend.inject_instruction(instructions)
+
+ # Resume the agent
+ if self._backend:
+ self._backend.resume()
+ self._set_state(CodingAgentState.RUNNING)
+
+ self._emit_event("branch_created", {
+ "branch_id": branch.branch_id,
+ "parent_branch": active_id,
+ "branch_point": step_index,
+ "instructions": instructions,
+ })
+
+ return branch.branch_id
+
+ def get_branches(self) -> List[dict]:
+ """List all branches."""
+ if self._branch_mgr:
+ return self._branch_mgr.list_branches()
+ return []
+
+ def switch_branch(self, branch_id: str) -> bool:
+ """Switch to a different branch."""
+ if not self._branch_mgr:
+ return False
+ if self._branch_mgr.switch_branch(branch_id):
+ branch = self._branch_mgr.get_branch(branch_id)
+ if branch:
+ self._structured_turns = list(branch.turns)
+ return True
+ return False
+
+ def get_diff_since_step(self, step_index: int) -> str:
+ """Get diff from a step to current state."""
+ if self._checkpoint_mgr:
+ return self._checkpoint_mgr.get_diff_since(step_index)
+ return ""
+
+ # --- Trace export ---
+
+ def get_trace(self) -> Dict[str, Any]:
+ """Get the full trace in CodingTraceDisplay format."""
+ trace = {
+ "session_id": self.session_id,
+ "task_description": self._task_description,
+ "structured_turns": list(self._structured_turns),
+ "backend": self.config.backend_type,
+ "model": self.config.ai_config.get("model", ""),
+ "started_at": self._started_at,
+ "sandbox_mode": self.config.sandbox_mode,
+ }
+ # Include branches if any were created
+ if self._branch_mgr and len(self._branch_mgr.list_branches()) > 1:
+ trace["branches"] = self._branch_mgr.save_all()
+ # Include checkpoints
+ if self._checkpoint_mgr:
+ trace["checkpoints"] = self._checkpoint_mgr.list_checkpoints()
+ return trace
+
+ def get_structured_turns(self) -> List[Dict]:
+ return list(self._structured_turns)
+
+ def get_state_summary(self) -> Dict[str, Any]:
+ return {
+ "session_id": self.session_id,
+ "state": self.state.value,
+ "task": self._task_description,
+ "turns": len(self._structured_turns),
+ "backend": self.config.backend_type,
+ }
+
+ def _save_trace(self):
+ """Save trace to disk."""
+ if not self.trace_dir:
+ return
+
+ os.makedirs(self.trace_dir, exist_ok=True)
+ trace_path = os.path.join(self.trace_dir, "trace.json")
+ try:
+ with open(trace_path, "w", encoding="utf-8") as f:
+ json.dump(self.get_trace(), f, indent=2, ensure_ascii=False)
+ logger.info(f"Saved coding agent trace to {trace_path}")
+ except Exception as e:
+ logger.error(f"Failed to save trace: {e}")
+
+ # --- Cleanup ---
+
+ def cleanup(self):
+ """Clean up resources."""
+ if self._backend:
+ try:
+ self._backend.stop()
+ except Exception:
+ pass
+ if self._sandbox:
+ try:
+ self._sandbox.cleanup()
+ except Exception:
+ pass
diff --git a/potato/coding_agent_runner_manager.py b/potato/coding_agent_runner_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..882eec26bc37b38081ba16eda5d8c3500fb4dab5
--- /dev/null
+++ b/potato/coding_agent_runner_manager.py
@@ -0,0 +1,112 @@
+"""
+Coding Agent Runner Manager
+
+Singleton manager for CodingAgentRunner sessions.
+Mirrors AgentRunnerManager pattern.
+"""
+
+import logging
+import threading
+import time
+import uuid
+from typing import Dict, List, Optional
+
+from .coding_agent_runner import CodingAgentRunner, CodingAgentConfig, CodingAgentState
+
+logger = logging.getLogger(__name__)
+
+
+class CodingAgentRunnerManager:
+ """Singleton manager for coding agent sessions."""
+
+ _instance = None
+ _lock = threading.Lock()
+
+ def __init__(self, max_sessions: int = 10, session_ttl: int = 3600):
+ self._sessions: Dict[str, CodingAgentRunner] = {}
+ self._session_keys: Dict[str, str] = {} # user:instance -> session_id
+ self._max_sessions = max_sessions
+ self._session_ttl = session_ttl
+ self._cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True)
+ self._cleanup_thread.start()
+
+ @classmethod
+ def get_instance(cls, **kwargs) -> "CodingAgentRunnerManager":
+ if cls._instance is None:
+ with cls._lock:
+ if cls._instance is None:
+ cls._instance = cls(**kwargs)
+ return cls._instance
+
+ @classmethod
+ def clear_instance(cls):
+ with cls._lock:
+ if cls._instance:
+ for runner in cls._instance._sessions.values():
+ runner.cleanup()
+ cls._instance = None
+
+ def create_session(self, user_id: str, instance_id: str,
+ config: CodingAgentConfig, trace_dir: str = "") -> CodingAgentRunner:
+ """Create a new coding agent session."""
+ key = f"{user_id}:{instance_id}"
+
+ # Check for existing active session
+ if key in self._session_keys:
+ existing = self._sessions.get(self._session_keys[key])
+ if existing and existing.state in (CodingAgentState.RUNNING, CodingAgentState.PAUSED):
+ return existing
+
+ if len(self._sessions) >= self._max_sessions:
+ self._evict_oldest()
+
+ session_id = str(uuid.uuid4())
+ runner = CodingAgentRunner(session_id, config, trace_dir)
+
+ self._sessions[session_id] = runner
+ self._session_keys[key] = session_id
+
+ logger.info(f"Created coding agent session {session_id} for {key}")
+ return runner
+
+ def get_session(self, session_id: str) -> Optional[CodingAgentRunner]:
+ return self._sessions.get(session_id)
+
+ def get_session_by_key(self, user_id: str, instance_id: str) -> Optional[CodingAgentRunner]:
+ key = f"{user_id}:{instance_id}"
+ sid = self._session_keys.get(key)
+ if sid:
+ return self._sessions.get(sid)
+ return None
+
+ def remove_session(self, session_id: str) -> None:
+ runner = self._sessions.pop(session_id, None)
+ if runner:
+ runner.cleanup()
+ # Remove from keys
+ self._session_keys = {
+ k: v for k, v in self._session_keys.items() if v != session_id
+ }
+
+ def list_sessions(self) -> List[Dict]:
+ return [r.get_state_summary() for r in self._sessions.values()]
+
+ def _evict_oldest(self):
+ """Remove the oldest completed/error session."""
+ for sid, runner in sorted(self._sessions.items()):
+ if runner.state in (CodingAgentState.COMPLETED, CodingAgentState.ERROR):
+ self.remove_session(sid)
+ return
+
+ def _cleanup_loop(self):
+ """Background cleanup of expired sessions."""
+ while True:
+ time.sleep(60)
+ expired = []
+ for sid, runner in list(self._sessions.items()):
+ if runner.state in (CodingAgentState.COMPLETED, CodingAgentState.ERROR):
+ if time.time() - runner._started_at > self._session_ttl:
+ expired.append(sid)
+ for sid in expired:
+ self.remove_session(sid)
+ logger.debug(f"Cleaned up expired session {sid}")
diff --git a/potato/coding_agent_sandbox.py b/potato/coding_agent_sandbox.py
new file mode 100644
index 0000000000000000000000000000000000000000..1b46f6b3787142095713ac32a01c1da0e45833d6
--- /dev/null
+++ b/potato/coding_agent_sandbox.py
@@ -0,0 +1,158 @@
+"""
+Coding Agent Sandbox Manager
+
+Manages isolated working directories for coding agent sessions.
+Supports three modes:
+- worktree: git worktree (lightweight copy, requires git repo)
+- docker: Docker container with mounted workspace
+- direct: No isolation (works directly in working_dir)
+"""
+
+import logging
+import os
+import shutil
+import subprocess
+import uuid
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+class SandboxManager:
+ """Manages sandboxed working directories for agent sessions."""
+
+ def __init__(self, mode: str = "worktree", base_dir: str = "."):
+ """Initialize the sandbox manager.
+
+ Args:
+ mode: Sandbox mode โ "worktree", "docker", or "direct"
+ base_dir: Base directory for creating sandboxes
+ """
+ if mode not in ("worktree", "docker", "direct"):
+ raise ValueError(f"Invalid sandbox mode: {mode}. Must be worktree, docker, or direct.")
+ self._mode = mode
+ self._base_dir = os.path.abspath(base_dir)
+ self._sandbox_dir: Optional[str] = None
+ self._session_id: Optional[str] = None
+ self._worktree_branch: Optional[str] = None
+
+ @property
+ def working_dir(self) -> str:
+ """The working directory for the agent."""
+ return self._sandbox_dir or self._base_dir
+
+ @property
+ def mode(self) -> str:
+ return self._mode
+
+ def create(self, session_id: str) -> str:
+ """Create a sandbox for the given session.
+
+ Returns:
+ The working directory path.
+ """
+ self._session_id = session_id
+
+ if self._mode == "worktree":
+ return self._create_worktree(session_id)
+ elif self._mode == "docker":
+ return self._create_docker(session_id)
+ else: # direct
+ self._sandbox_dir = self._base_dir
+ return self._base_dir
+
+ def cleanup(self) -> None:
+ """Clean up the sandbox."""
+ if self._mode == "worktree":
+ self._cleanup_worktree()
+ elif self._mode == "docker":
+ self._cleanup_docker()
+ # direct mode: nothing to clean up
+
+ def _create_worktree(self, session_id: str) -> str:
+ """Create a git worktree for isolation."""
+ # Check if base_dir is a git repo
+ try:
+ subprocess.run(
+ ["git", "rev-parse", "--git-dir"],
+ cwd=self._base_dir, capture_output=True, check=True,
+ )
+ except (subprocess.CalledProcessError, FileNotFoundError):
+ logger.warning(
+ f"Directory {self._base_dir} is not a git repo. "
+ f"Falling back to direct mode."
+ )
+ self._mode = "direct"
+ self._sandbox_dir = self._base_dir
+ return self._base_dir
+
+ # Create worktree in a temp location
+ branch_name = f"potato-agent-{session_id[:8]}"
+ worktree_dir = os.path.join(
+ os.path.dirname(self._base_dir),
+ f".potato-sandbox-{session_id[:8]}",
+ )
+
+ try:
+ # Create a new branch from HEAD
+ subprocess.run(
+ ["git", "branch", branch_name, "HEAD"],
+ cwd=self._base_dir, capture_output=True, check=True,
+ )
+
+ # Create worktree
+ subprocess.run(
+ ["git", "worktree", "add", worktree_dir, branch_name],
+ cwd=self._base_dir, capture_output=True, check=True,
+ )
+
+ self._sandbox_dir = worktree_dir
+ self._worktree_branch = branch_name
+ logger.info(f"Created git worktree sandbox at {worktree_dir}")
+ return worktree_dir
+
+ except subprocess.CalledProcessError as e:
+ logger.warning(f"Failed to create worktree: {e}. Falling back to direct mode.")
+ self._mode = "direct"
+ self._sandbox_dir = self._base_dir
+ return self._base_dir
+
+ def _cleanup_worktree(self) -> None:
+ """Remove the git worktree and branch."""
+ if not self._sandbox_dir or self._sandbox_dir == self._base_dir:
+ return
+
+ try:
+ # Remove worktree
+ subprocess.run(
+ ["git", "worktree", "remove", self._sandbox_dir, "--force"],
+ cwd=self._base_dir, capture_output=True,
+ )
+ logger.info(f"Removed worktree at {self._sandbox_dir}")
+ except Exception as e:
+ logger.warning(f"Failed to remove worktree: {e}")
+ # Manual cleanup
+ if os.path.exists(self._sandbox_dir):
+ shutil.rmtree(self._sandbox_dir, ignore_errors=True)
+
+ # Clean up the branch
+ if self._worktree_branch:
+ try:
+ subprocess.run(
+ ["git", "branch", "-D", self._worktree_branch],
+ cwd=self._base_dir, capture_output=True,
+ )
+ except Exception:
+ pass
+
+ def _create_docker(self, session_id: str) -> str:
+ """Create a Docker container for maximum isolation."""
+ # For Phase 4 โ placeholder
+ logger.warning("Docker sandbox not yet implemented. Using direct mode.")
+ self._mode = "direct"
+ self._sandbox_dir = self._base_dir
+ return self._base_dir
+
+ def _cleanup_docker(self) -> None:
+ """Remove Docker container."""
+ pass # Phase 4
diff --git a/potato/create_task_cli.py b/potato/create_task_cli.py
new file mode 100644
index 0000000000000000000000000000000000000000..18e177287347708b9fc0d388816de36ba89bd6c1
--- /dev/null
+++ b/potato/create_task_cli.py
@@ -0,0 +1,163 @@
+"""
+Interactive Task Creation CLI Module
+
+Provides an interactive command-line interface for creating new annotation tasks.
+Guides users through configuration and generates a YAML config file.
+"""
+
+import os
+
+import click
+import yaml
+
+
+# Common annotation types shown in the interactive prompt.
+# The full set of supported types is in the schema registry.
+COMMON_ANNOTATION_TYPES = [
+ "radio",
+ "multiselect",
+ "text",
+ "likert",
+ "slider",
+ "span",
+ "bws",
+ "pairwise",
+ "ranking",
+]
+
+
+def create_task_cli():
+ """
+ Interactive task creation wizard.
+
+ Walks the user through server setup, data files, annotation schemes,
+ and output settings, then writes a YAML config file.
+ """
+ click.echo("Welcome to the Potato annotation task creation wizard!")
+ click.echo("This will generate a .yaml config file for your annotation task.")
+ click.echo("You can edit the file afterwards to add advanced features.\n")
+
+ click.echo("โ Server setup โ\n")
+
+ task_name = click.prompt("Task name (shown to annotators)", default="My Annotation Task")
+ port = click.prompt("Server port", default=8000, type=click.IntRange(1, 65535))
+
+ click.echo("\nโ Data files โ\n")
+
+ data_files = []
+ fname = click.prompt("Path to your data file (JSONL)")
+ data_files.append(fname)
+ while click.confirm("Add another data file?", default=False):
+ fname = click.prompt("Path to data file")
+ data_files.append(fname)
+
+ id_key = click.prompt("Which field in the data is the item ID?", default="id")
+ text_key = click.prompt("Which field is the item text?", default="text")
+
+ click.echo("\nโ Annotation schemes โ\n")
+
+ annotation_schemes = []
+ while True:
+ atype = click.prompt(
+ "Annotation type",
+ type=click.Choice(COMMON_ANNOTATION_TYPES, case_sensitive=False),
+ )
+ name = click.prompt("Internal name for this scheme (used in output)")
+ desc = click.prompt("Description/question shown to annotators")
+
+ scheme = {
+ "annotation_type": atype,
+ "name": name,
+ "description": desc,
+ }
+
+ if atype in ("radio", "multiselect"):
+ labels = []
+ click.echo("Enter labels one at a time (empty line to finish):")
+ while True:
+ label = click.prompt(" Label", default="", show_default=False)
+ if not label:
+ break
+ labels.append(label)
+ scheme["labels"] = labels
+
+ elif atype == "likert":
+ size = click.prompt("Scale size", default=5, type=int)
+ min_label = click.prompt("Label for minimum", default="Strongly Disagree")
+ max_label = click.prompt("Label for maximum", default="Strongly Agree")
+ scheme["size"] = size
+ scheme["min_label"] = min_label
+ scheme["max_label"] = max_label
+
+ elif atype == "slider":
+ scheme["min_value"] = click.prompt("Minimum value", default=0, type=int)
+ scheme["max_value"] = click.prompt("Maximum value", default=100, type=int)
+
+ elif atype == "bws":
+ size = click.prompt("Tuple size (items shown at once)", default=4, type=int)
+ scheme["size"] = size
+
+ elif atype == "ranking":
+ labels = []
+ click.echo("Enter options to rank (empty line to finish):")
+ while True:
+ label = click.prompt(" Option", default="", show_default=False)
+ if not label:
+ break
+ labels.append(label)
+ scheme["labels"] = labels
+
+ annotation_schemes.append(scheme)
+
+ if not click.confirm("\nAdd another annotation scheme?", default=False):
+ break
+
+ click.echo("\nโ Output settings โ\n")
+
+ output_dir = click.prompt("Output directory for annotations", default="annotation_output")
+ codebook_url = click.prompt("Annotation codebook URL (optional)", default="")
+
+ auto_export = click.confirm("Auto-export annotations in CSV/JSONL?", default=False)
+ export_format = None
+ if auto_export:
+ export_format = click.prompt(
+ "Export format",
+ type=click.Choice(["csv", "tsv", "jsonl"], case_sensitive=False),
+ default="csv",
+ )
+
+ # Build config
+ config = {
+ "annotation_task_name": task_name,
+ "port": port,
+ "data_files": data_files,
+ "item_properties": {
+ "id_key": id_key,
+ "text_key": text_key,
+ },
+ "annotation_schemes": annotation_schemes,
+ "output_annotation_dir": output_dir,
+ "annotation_codebook_url": codebook_url,
+ "user_config": {
+ "allow_all_users": True,
+ "users": [],
+ },
+ }
+
+ if export_format:
+ config["export_annotation_format"] = export_format
+
+ click.echo("\nโ Save config โ\n")
+
+ config_file = click.prompt("Path for the config file", default="config.yaml")
+
+ if os.path.exists(config_file):
+ if not click.confirm(f"{config_file} already exists. Overwrite?", default=False):
+ click.echo("Aborted.")
+ return
+
+ with open(config_file, "w", encoding="utf-8") as f:
+ yaml.dump(config, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
+
+ click.echo(f"\nConfig written to {config_file}")
+ click.echo(f"Start annotating with: potato start {config_file}")
diff --git a/potato/data_sources/__init__.py b/potato/data_sources/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..aca583333ede182c984f95e38a301998542402db
--- /dev/null
+++ b/potato/data_sources/__init__.py
@@ -0,0 +1,78 @@
+"""
+Data Sources Module
+
+This module provides extensible data loading from various sources including
+local files, URLs, cloud storage (Google Drive, Dropbox, S3), databases,
+and other remote sources.
+
+The module follows a singleton pattern for the DataSourceManager, similar
+to other managers in Potato (ItemStateManager, UserStateManager).
+
+Example usage:
+ from potato.data_sources import (
+ init_data_source_manager,
+ get_data_source_manager,
+ clear_data_source_manager
+ )
+
+ # Initialize with config
+ manager = init_data_source_manager(config)
+
+ # Load initial data
+ manager.load_initial_data()
+
+ # Get manager later
+ manager = get_data_source_manager()
+
+ # Load more data incrementally
+ manager.load_more(source_id, count=500)
+"""
+
+from potato.data_sources.manager import (
+ init_data_source_manager,
+ get_data_source_manager,
+ clear_data_source_manager,
+ DataSourceManager,
+)
+
+from potato.data_sources.base import (
+ DataSource,
+ SourceType,
+ SourceConfig,
+)
+
+from potato.data_sources.credentials import (
+ CredentialManager,
+ substitute_env_vars,
+)
+
+from potato.data_sources.cache_manager import (
+ CacheManager,
+ CacheEntry,
+)
+
+from potato.data_sources.partial_reader import (
+ PartialReader,
+ PartialReadState,
+)
+
+__all__ = [
+ # Manager functions
+ "init_data_source_manager",
+ "get_data_source_manager",
+ "clear_data_source_manager",
+ "DataSourceManager",
+ # Base classes
+ "DataSource",
+ "SourceType",
+ "SourceConfig",
+ # Credential management
+ "CredentialManager",
+ "substitute_env_vars",
+ # Cache management
+ "CacheManager",
+ "CacheEntry",
+ # Partial/incremental loading
+ "PartialReader",
+ "PartialReadState",
+]
diff --git a/potato/data_sources/base.py b/potato/data_sources/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..4fcc4d620fc9cccf30bd7c3747d0ed1065d7a676
--- /dev/null
+++ b/potato/data_sources/base.py
@@ -0,0 +1,276 @@
+"""
+Base classes and types for data sources.
+
+This module defines the abstract base class for all data sources and
+common types used throughout the data sources subsystem.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, Iterator, List, Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class SourceType(Enum):
+ """Enumeration of supported data source types."""
+ FILE = "file"
+ URL = "url"
+ GOOGLE_DRIVE = "google_drive"
+ DROPBOX = "dropbox"
+ S3 = "s3"
+ HUGGINGFACE = "huggingface"
+ GOOGLE_SHEETS = "google_sheets"
+ DATABASE = "database"
+
+
+@dataclass
+class SourceConfig:
+ """
+ Configuration for a data source.
+
+ This dataclass holds the parsed configuration for a single data source,
+ including type-specific settings and common options.
+
+ Attributes:
+ source_type: The type of data source
+ source_id: Unique identifier for this source (auto-generated if not provided)
+ config: The raw configuration dictionary for this source
+ enabled: Whether this source is enabled
+ """
+ source_type: SourceType
+ source_id: str
+ config: Dict[str, Any]
+ enabled: bool = True
+
+ @classmethod
+ def from_dict(cls, config_dict: Dict[str, Any], index: int = 0) -> "SourceConfig":
+ """
+ Create a SourceConfig from a configuration dictionary.
+
+ Args:
+ config_dict: Dictionary containing source configuration
+ index: Index in the sources list (used for auto-generated ID)
+
+ Returns:
+ SourceConfig instance
+
+ Raises:
+ ValueError: If the source type is invalid or required fields are missing
+ """
+ type_str = config_dict.get("type")
+ if not type_str:
+ raise ValueError("Data source configuration must include 'type' field")
+
+ try:
+ source_type = SourceType(type_str)
+ except ValueError:
+ valid_types = [t.value for t in SourceType]
+ raise ValueError(
+ f"Invalid data source type '{type_str}'. "
+ f"Valid types are: {', '.join(valid_types)}"
+ )
+
+ # Generate source_id if not provided
+ source_id = config_dict.get("id") or config_dict.get("source_id")
+ if not source_id:
+ # Generate ID based on type and index or key identifying info
+ if source_type == SourceType.FILE:
+ path = config_dict.get("path", "")
+ source_id = f"file_{index}_{path.replace('/', '_').replace('.', '_')}"
+ elif source_type == SourceType.URL:
+ url = config_dict.get("url", "")
+ # Use last part of URL path as identifier
+ source_id = f"url_{index}_{url.split('/')[-1][:30]}"
+ else:
+ source_id = f"{source_type.value}_{index}"
+
+ enabled = config_dict.get("enabled", True)
+
+ return cls(
+ source_type=source_type,
+ source_id=source_id,
+ config=config_dict,
+ enabled=enabled
+ )
+
+
+class DataSource(ABC):
+ """
+ Abstract base class for all data sources.
+
+ Each data source implementation must provide methods for:
+ - Identifying the source
+ - Checking availability
+ - Reading items (with optional partial reading support)
+ - Reporting total item count
+
+ Thread Safety:
+ Implementations should be thread-safe for concurrent read operations.
+ Write operations (if any) should be protected by appropriate locking.
+ """
+
+ def __init__(self, config: SourceConfig):
+ """
+ Initialize the data source.
+
+ Args:
+ config: Source configuration
+ """
+ self._config = config
+ self._source_id = config.source_id
+ self._raw_config = config.config
+
+ @property
+ def source_id(self) -> str:
+ """Get the unique identifier for this source."""
+ return self._source_id
+
+ @property
+ def source_type(self) -> SourceType:
+ """Get the type of this source."""
+ return self._config.source_type
+
+ @property
+ def config(self) -> Dict[str, Any]:
+ """Get the raw configuration dictionary."""
+ return self._raw_config
+
+ @abstractmethod
+ def get_source_id(self) -> str:
+ """
+ Get the unique identifier for this data source.
+
+ Returns:
+ String identifier unique within the DataSourceManager
+ """
+ pass
+
+ @abstractmethod
+ def is_available(self) -> bool:
+ """
+ Check if the data source is available and accessible.
+
+ This method should verify that:
+ - Required dependencies are installed
+ - Credentials are valid (if applicable)
+ - The source location exists and is readable
+
+ Returns:
+ True if the source is ready to read from
+ """
+ pass
+
+ @abstractmethod
+ def read_items(
+ self,
+ start: int = 0,
+ count: Optional[int] = None
+ ) -> Iterator[Dict[str, Any]]:
+ """
+ Read items from the data source.
+
+ This method yields dictionaries representing annotation items.
+ For sources that support partial reading, the start and count
+ parameters allow fetching specific ranges of items.
+
+ Args:
+ start: Index of the first item to read (0-based)
+ count: Maximum number of items to read (None = all remaining)
+
+ Yields:
+ Dictionary containing item data with at least id_key field
+
+ Raises:
+ RuntimeError: If the source is not available
+ IOError: If reading fails
+ """
+ pass
+
+ @abstractmethod
+ def get_total_count(self) -> Optional[int]:
+ """
+ Get the total number of items in the source.
+
+ Returns:
+ Total item count, or None if unknown (e.g., streaming source)
+ """
+ pass
+
+ @abstractmethod
+ def supports_partial_reading(self) -> bool:
+ """
+ Check if this source supports reading partial ranges.
+
+ Partial reading allows loading data in chunks, which is useful
+ for large datasets or incremental annotation workflows.
+
+ Returns:
+ True if read_items() supports start/count parameters
+ """
+ pass
+
+ def validate_config(self) -> List[str]:
+ """
+ Validate the source configuration.
+
+ Override this method to add source-specific validation.
+
+ Returns:
+ List of validation error messages (empty if valid)
+ """
+ return []
+
+ def get_status(self) -> Dict[str, Any]:
+ """
+ Get the current status of this data source.
+
+ Returns:
+ Dictionary containing status information:
+ - source_id: The source identifier
+ - source_type: The source type
+ - available: Whether the source is available
+ - total_count: Total item count (or None)
+ - supports_partial: Whether partial reading is supported
+ """
+ return {
+ "source_id": self.source_id,
+ "source_type": self.source_type.value,
+ "available": self.is_available(),
+ "total_count": self.get_total_count(),
+ "supports_partial": self.supports_partial_reading(),
+ }
+
+ def refresh(self) -> bool:
+ """
+ Refresh the data source (re-fetch from remote, re-validate, etc.).
+
+ Override this method for sources that cache data or need
+ periodic refresh. The default implementation does nothing.
+
+ Returns:
+ True if refresh was successful
+ """
+ return True
+
+ def close(self) -> None:
+ """
+ Close the data source and release any resources.
+
+ Override this method for sources that hold resources like
+ database connections or file handles.
+ """
+ pass
+
+ def __enter__(self) -> "DataSource":
+ """Context manager entry."""
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+ """Context manager exit."""
+ self.close()
+
+ def __repr__(self) -> str:
+ return f"{self.__class__.__name__}(source_id={self.source_id!r})"
diff --git a/potato/data_sources/cache_manager.py b/potato/data_sources/cache_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..4526a7ca489d44677428fd7105588ec09e0afb89
--- /dev/null
+++ b/potato/data_sources/cache_manager.py
@@ -0,0 +1,525 @@
+"""
+Cache manager for remote data sources.
+
+This module provides caching functionality for downloaded remote files,
+including TTL-based expiration, ETag support for HTTP caching, and
+thread-safe operations.
+"""
+
+import hashlib
+import json
+import logging
+import os
+import shutil
+import threading
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CacheEntry:
+ """
+ Represents a cached file entry.
+
+ Attributes:
+ source_id: Identifier of the data source
+ source_url: Original URL or path
+ cache_path: Local path to cached file
+ etag: HTTP ETag for cache validation (optional)
+ last_modified: HTTP Last-Modified header value (optional)
+ created_at: Unix timestamp when cached
+ expires_at: Unix timestamp when cache expires
+ file_size: Size of cached file in bytes
+ content_type: MIME type of cached content
+ metadata: Additional metadata about the cached content
+ """
+ source_id: str
+ source_url: str
+ cache_path: str
+ etag: Optional[str] = None
+ last_modified: Optional[str] = None
+ created_at: float = field(default_factory=time.time)
+ expires_at: Optional[float] = None
+ file_size: int = 0
+ content_type: Optional[str] = None
+ metadata: Dict[str, Any] = field(default_factory=dict)
+
+ def is_expired(self) -> bool:
+ """Check if this cache entry has expired."""
+ if self.expires_at is None:
+ return False
+ return time.time() > self.expires_at
+
+ def is_valid(self) -> bool:
+ """Check if the cached file exists and hasn't expired."""
+ if not os.path.exists(self.cache_path):
+ return False
+ if self.is_expired():
+ return False
+ return True
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to dictionary for JSON serialization."""
+ return {
+ "source_id": self.source_id,
+ "source_url": self.source_url,
+ "cache_path": self.cache_path,
+ "etag": self.etag,
+ "last_modified": self.last_modified,
+ "created_at": self.created_at,
+ "expires_at": self.expires_at,
+ "file_size": self.file_size,
+ "content_type": self.content_type,
+ "metadata": self.metadata,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "CacheEntry":
+ """Create from dictionary."""
+ return cls(
+ source_id=data["source_id"],
+ source_url=data["source_url"],
+ cache_path=data["cache_path"],
+ etag=data.get("etag"),
+ last_modified=data.get("last_modified"),
+ created_at=data.get("created_at", time.time()),
+ expires_at=data.get("expires_at"),
+ file_size=data.get("file_size", 0),
+ content_type=data.get("content_type"),
+ metadata=data.get("metadata", {}),
+ )
+
+
+class CacheManager:
+ """
+ Manages a local file cache for remote data sources.
+
+ This class provides thread-safe caching of downloaded files with:
+ - TTL-based expiration
+ - ETag and Last-Modified support for conditional requests
+ - Automatic cache directory management
+ - Persistent cache index for restart recovery
+
+ Attributes:
+ cache_dir: Path to the cache directory
+ ttl_seconds: Default time-to-live for cached files
+ max_size_mb: Maximum total cache size in megabytes
+ """
+
+ DEFAULT_TTL = 3600 # 1 hour
+ DEFAULT_MAX_SIZE_MB = 500
+ INDEX_FILENAME = "_cache_index.json"
+
+ def __init__(
+ self,
+ cache_dir: str,
+ ttl_seconds: int = DEFAULT_TTL,
+ max_size_mb: int = DEFAULT_MAX_SIZE_MB
+ ):
+ """
+ Initialize the cache manager.
+
+ Args:
+ cache_dir: Directory to store cached files
+ ttl_seconds: Default TTL for cached entries
+ max_size_mb: Maximum total cache size
+ """
+ self.cache_dir = Path(cache_dir)
+ self.ttl_seconds = ttl_seconds
+ self.max_size_bytes = max_size_mb * 1024 * 1024
+
+ self._entries: Dict[str, CacheEntry] = {}
+ self._lock = threading.RLock()
+
+ # Create cache directory if it doesn't exist
+ self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+ # Load existing cache index
+ self._load_index()
+
+ logger.info(
+ f"CacheManager initialized: dir={cache_dir}, "
+ f"ttl={ttl_seconds}s, max_size={max_size_mb}MB"
+ )
+
+ def _load_index(self) -> None:
+ """Load the cache index from disk."""
+ index_path = self.cache_dir / self.INDEX_FILENAME
+ if not index_path.exists():
+ return
+
+ try:
+ with open(index_path, 'r', encoding='utf-8') as f:
+ data = json.load(f)
+
+ for entry_data in data.get("entries", []):
+ entry = CacheEntry.from_dict(entry_data)
+ # Only load if cached file still exists
+ if os.path.exists(entry.cache_path):
+ self._entries[entry.source_id] = entry
+ else:
+ logger.debug(f"Cached file missing, skipping: {entry.cache_path}")
+
+ logger.debug(f"Loaded {len(self._entries)} cache entries from index")
+
+ except (json.JSONDecodeError, KeyError) as e:
+ logger.warning(f"Failed to load cache index: {e}")
+ self._entries = {}
+
+ def _save_index(self) -> None:
+ """Save the cache index to disk."""
+ index_path = self.cache_dir / self.INDEX_FILENAME
+
+ data = {
+ "version": 1,
+ "entries": [entry.to_dict() for entry in self._entries.values()]
+ }
+
+ try:
+ with open(index_path, 'w', encoding='utf-8') as f:
+ json.dump(data, f, indent=2)
+ except IOError as e:
+ logger.error(f"Failed to save cache index: {e}")
+
+ def _generate_cache_key(self, source_id: str, url: str) -> str:
+ """Generate a unique cache key for a source."""
+ hash_input = f"{source_id}:{url}"
+ return hashlib.sha256(hash_input.encode()).hexdigest()[:32]
+
+ def _generate_cache_path(self, source_id: str, url: str, extension: str = "") -> Path:
+ """Generate the cache file path for a source."""
+ cache_key = self._generate_cache_key(source_id, url)
+ filename = f"{cache_key}{extension}"
+ return self.cache_dir / filename
+
+ def get(self, source_id: str) -> Optional[CacheEntry]:
+ """
+ Get a cache entry by source ID.
+
+ Args:
+ source_id: The source identifier
+
+ Returns:
+ CacheEntry if found and valid, None otherwise
+ """
+ with self._lock:
+ entry = self._entries.get(source_id)
+ if entry and entry.is_valid():
+ return entry
+ return None
+
+ def get_if_valid(
+ self,
+ source_id: str,
+ etag: Optional[str] = None,
+ last_modified: Optional[str] = None
+ ) -> Optional[CacheEntry]:
+ """
+ Get a cache entry if it's still valid.
+
+ For HTTP sources, this checks ETag and Last-Modified headers
+ for cache validation.
+
+ Args:
+ source_id: The source identifier
+ etag: Current ETag from server (for validation)
+ last_modified: Current Last-Modified from server
+
+ Returns:
+ CacheEntry if cache hit, None if miss or stale
+ """
+ with self._lock:
+ entry = self._entries.get(source_id)
+ if not entry:
+ return None
+
+ # Check if file exists
+ if not os.path.exists(entry.cache_path):
+ del self._entries[source_id]
+ self._save_index()
+ return None
+
+ # Check TTL expiration
+ if entry.is_expired():
+ return None
+
+ # If ETag provided, validate it matches
+ if etag and entry.etag and entry.etag != etag:
+ return None
+
+ # If Last-Modified provided, validate it
+ if last_modified and entry.last_modified and entry.last_modified != last_modified:
+ return None
+
+ return entry
+
+ def put(
+ self,
+ source_id: str,
+ source_url: str,
+ data: bytes,
+ etag: Optional[str] = None,
+ last_modified: Optional[str] = None,
+ content_type: Optional[str] = None,
+ ttl_seconds: Optional[int] = None,
+ metadata: Optional[Dict[str, Any]] = None
+ ) -> CacheEntry:
+ """
+ Store data in the cache.
+
+ Args:
+ source_id: Unique identifier for this source
+ source_url: Original URL of the data
+ data: The data to cache
+ etag: HTTP ETag header value
+ last_modified: HTTP Last-Modified header value
+ content_type: MIME type of the content
+ ttl_seconds: Time-to-live (uses default if not specified)
+ metadata: Additional metadata to store
+
+ Returns:
+ The created CacheEntry
+ """
+ # Determine file extension from content type
+ extension = self._extension_from_content_type(content_type, source_url)
+ cache_path = self._generate_cache_path(source_id, source_url, extension)
+
+ with self._lock:
+ # Write data to cache file
+ try:
+ with open(cache_path, 'wb') as f:
+ f.write(data)
+ except IOError as e:
+ logger.error(f"Failed to write cache file: {e}")
+ raise
+
+ # Create cache entry
+ ttl = ttl_seconds if ttl_seconds is not None else self.ttl_seconds
+ entry = CacheEntry(
+ source_id=source_id,
+ source_url=source_url,
+ cache_path=str(cache_path),
+ etag=etag,
+ last_modified=last_modified,
+ created_at=time.time(),
+ expires_at=time.time() + ttl if ttl > 0 else None,
+ file_size=len(data),
+ content_type=content_type,
+ metadata=metadata or {},
+ )
+
+ self._entries[source_id] = entry
+ self._save_index()
+
+ # Check cache size and cleanup if needed
+ self._enforce_size_limit()
+
+ logger.debug(f"Cached {len(data)} bytes for {source_id}")
+ return entry
+
+ def put_file(
+ self,
+ source_id: str,
+ source_url: str,
+ file_path: str,
+ etag: Optional[str] = None,
+ last_modified: Optional[str] = None,
+ content_type: Optional[str] = None,
+ ttl_seconds: Optional[int] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ move: bool = False
+ ) -> CacheEntry:
+ """
+ Store a file in the cache.
+
+ Args:
+ source_id: Unique identifier for this source
+ source_url: Original URL of the data
+ file_path: Path to the file to cache
+ etag: HTTP ETag header value
+ last_modified: HTTP Last-Modified header value
+ content_type: MIME type of the content
+ ttl_seconds: Time-to-live
+ metadata: Additional metadata
+ move: If True, move the file instead of copying
+
+ Returns:
+ The created CacheEntry
+ """
+ extension = self._extension_from_content_type(content_type, source_url)
+ cache_path = self._generate_cache_path(source_id, source_url, extension)
+
+ with self._lock:
+ try:
+ if move:
+ shutil.move(file_path, cache_path)
+ else:
+ shutil.copy2(file_path, cache_path)
+ except IOError as e:
+ logger.error(f"Failed to cache file: {e}")
+ raise
+
+ file_size = os.path.getsize(cache_path)
+
+ ttl = ttl_seconds if ttl_seconds is not None else self.ttl_seconds
+ entry = CacheEntry(
+ source_id=source_id,
+ source_url=source_url,
+ cache_path=str(cache_path),
+ etag=etag,
+ last_modified=last_modified,
+ created_at=time.time(),
+ expires_at=time.time() + ttl if ttl > 0 else None,
+ file_size=file_size,
+ content_type=content_type,
+ metadata=metadata or {},
+ )
+
+ self._entries[source_id] = entry
+ self._save_index()
+ self._enforce_size_limit()
+
+ return entry
+
+ def invalidate(self, source_id: str) -> bool:
+ """
+ Invalidate (remove) a cache entry.
+
+ Args:
+ source_id: The source identifier
+
+ Returns:
+ True if entry was removed, False if not found
+ """
+ with self._lock:
+ entry = self._entries.pop(source_id, None)
+ if entry:
+ try:
+ if os.path.exists(entry.cache_path):
+ os.remove(entry.cache_path)
+ except IOError as e:
+ logger.warning(f"Failed to remove cache file: {e}")
+
+ self._save_index()
+ return True
+ return False
+
+ def clear(self) -> int:
+ """
+ Clear all cache entries.
+
+ Returns:
+ Number of entries cleared
+ """
+ with self._lock:
+ count = len(self._entries)
+
+ for entry in self._entries.values():
+ try:
+ if os.path.exists(entry.cache_path):
+ os.remove(entry.cache_path)
+ except IOError as e:
+ logger.warning(f"Failed to remove cache file: {e}")
+
+ self._entries.clear()
+ self._save_index()
+
+ logger.info(f"Cleared {count} cache entries")
+ return count
+
+ def cleanup_expired(self) -> int:
+ """
+ Remove all expired cache entries.
+
+ Returns:
+ Number of entries removed
+ """
+ with self._lock:
+ expired = [
+ source_id
+ for source_id, entry in self._entries.items()
+ if entry.is_expired()
+ ]
+
+ for source_id in expired:
+ self.invalidate(source_id)
+
+ if expired:
+ logger.debug(f"Cleaned up {len(expired)} expired cache entries")
+
+ return len(expired)
+
+ def get_stats(self) -> Dict[str, Any]:
+ """
+ Get cache statistics.
+
+ Returns:
+ Dictionary with cache statistics
+ """
+ with self._lock:
+ total_size = sum(e.file_size for e in self._entries.values())
+ expired_count = sum(1 for e in self._entries.values() if e.is_expired())
+
+ return {
+ "cache_dir": str(self.cache_dir),
+ "entry_count": len(self._entries),
+ "total_size_bytes": total_size,
+ "total_size_mb": round(total_size / (1024 * 1024), 2),
+ "max_size_mb": self.max_size_bytes // (1024 * 1024),
+ "expired_count": expired_count,
+ "ttl_seconds": self.ttl_seconds,
+ }
+
+ def _enforce_size_limit(self) -> None:
+ """Remove oldest entries if cache exceeds size limit."""
+ total_size = sum(e.file_size for e in self._entries.values())
+
+ if total_size <= self.max_size_bytes:
+ return
+
+ # Sort by creation time (oldest first)
+ sorted_entries = sorted(
+ self._entries.items(),
+ key=lambda x: x[1].created_at
+ )
+
+ removed = 0
+ for source_id, entry in sorted_entries:
+ if total_size <= self.max_size_bytes:
+ break
+
+ total_size -= entry.file_size
+ self.invalidate(source_id)
+ removed += 1
+
+ if removed:
+ logger.info(f"Removed {removed} cache entries to enforce size limit")
+
+ def _extension_from_content_type(
+ self,
+ content_type: Optional[str],
+ url: str
+ ) -> str:
+ """Determine file extension from content type or URL."""
+ # Try content type first
+ if content_type:
+ type_to_ext = {
+ "application/json": ".json",
+ "text/csv": ".csv",
+ "text/tab-separated-values": ".tsv",
+ "text/plain": ".txt",
+ }
+ for mime, ext in type_to_ext.items():
+ if content_type.startswith(mime):
+ return ext
+
+ # Fall back to URL extension
+ url_path = url.split('?')[0] # Remove query string
+ if '.' in url_path.split('/')[-1]:
+ return '.' + url_path.split('.')[-1]
+
+ return "" # No extension
diff --git a/potato/data_sources/credentials.py b/potato/data_sources/credentials.py
new file mode 100644
index 0000000000000000000000000000000000000000..45762051c6dc6f5bb7f18bbc09e0ca0e83aed030
--- /dev/null
+++ b/potato/data_sources/credentials.py
@@ -0,0 +1,342 @@
+"""
+Credential management for data sources.
+
+This module provides secure credential handling including:
+- Environment variable substitution in configuration values
+- Support for .env files
+- Service account and API key management
+- Credential validation without logging sensitive values
+"""
+
+import logging
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+# Pattern to match environment variable references: ${VAR_NAME}
+ENV_VAR_PATTERN = re.compile(r'\$\{([A-Za-z_][A-Za-z0-9_]*)\}')
+
+
+def substitute_env_vars(value: Any, env_file: Optional[str] = None) -> Any:
+ """
+ Substitute environment variable references in a configuration value.
+
+ Supports the ${VAR_NAME} syntax for referencing environment variables.
+ If an environment variable is not set, the reference is left unchanged
+ and a warning is logged.
+
+ Args:
+ value: The value to process (string, dict, list, or other)
+ env_file: Optional path to .env file to load additional variables
+
+ Returns:
+ The value with environment variables substituted
+
+ Examples:
+ >>> os.environ['API_KEY'] = 'secret123'
+ >>> substitute_env_vars('Bearer ${API_KEY}')
+ 'Bearer secret123'
+
+ >>> substitute_env_vars({'auth': '${TOKEN}'})
+ {'auth': ''}
+ """
+ # Load .env file if specified and exists
+ if env_file:
+ load_env_file(env_file)
+
+ return _substitute_recursive(value)
+
+
+def _substitute_recursive(value: Any) -> Any:
+ """Recursively substitute environment variables in nested structures."""
+ if isinstance(value, str):
+ return _substitute_in_string(value)
+ elif isinstance(value, dict):
+ return {k: _substitute_recursive(v) for k, v in value.items()}
+ elif isinstance(value, list):
+ return [_substitute_recursive(v) for v in value]
+ else:
+ return value
+
+
+def _substitute_in_string(value: str) -> str:
+ """Substitute environment variables in a string value."""
+ def replacer(match):
+ var_name = match.group(1)
+ env_value = os.environ.get(var_name)
+ if env_value is None:
+ logger.warning(
+ f"Environment variable '{var_name}' is not set. "
+ f"The reference ${{{var_name}}} will be left unchanged."
+ )
+ return match.group(0) # Return original ${VAR_NAME}
+ return env_value
+
+ return ENV_VAR_PATTERN.sub(replacer, value)
+
+
+def load_env_file(env_file: str) -> int:
+ """
+ Load environment variables from a .env file.
+
+ The file format is:
+ VAR_NAME=value
+ # Comments are ignored
+ ANOTHER_VAR="quoted value"
+
+ Variables are added to os.environ but do not override existing values.
+
+ Args:
+ env_file: Path to the .env file
+
+ Returns:
+ Number of variables loaded
+
+ Raises:
+ FileNotFoundError: If the env_file does not exist
+ """
+ env_path = Path(env_file)
+ if not env_path.exists():
+ raise FileNotFoundError(f"Environment file not found: {env_file}")
+
+ count = 0
+ with open(env_path, 'r', encoding='utf-8') as f:
+ for line_no, line in enumerate(f, 1):
+ line = line.strip()
+
+ # Skip empty lines and comments
+ if not line or line.startswith('#'):
+ continue
+
+ # Parse VAR=value format
+ if '=' not in line:
+ logger.warning(
+ f"Invalid line {line_no} in {env_file}: missing '=' separator"
+ )
+ continue
+
+ key, _, value = line.partition('=')
+ key = key.strip()
+ value = value.strip()
+
+ # Remove quotes if present
+ if (value.startswith('"') and value.endswith('"')) or \
+ (value.startswith("'") and value.endswith("'")):
+ value = value[1:-1]
+
+ # Only set if not already in environment (don't override)
+ if key not in os.environ:
+ os.environ[key] = value
+ count += 1
+ logger.debug(f"Loaded environment variable: {key}")
+ else:
+ logger.debug(
+ f"Skipping {key} from {env_file}: already set in environment"
+ )
+
+ logger.info(f"Loaded {count} environment variables from {env_file}")
+ return count
+
+
+@dataclass
+class CredentialManager:
+ """
+ Manages credentials for data source authentication.
+
+ This class provides a centralized way to handle credentials including:
+ - Environment variable substitution
+ - Loading from .env files
+ - Validating required credentials
+ - Masking credentials in logs
+
+ Attributes:
+ env_substitution: Whether to perform env var substitution
+ env_file: Path to optional .env file
+ """
+
+ env_substitution: bool = True
+ env_file: Optional[str] = None
+ _env_loaded: bool = False
+
+ def __post_init__(self):
+ """Load .env file if configured."""
+ if self.env_file and not self._env_loaded:
+ try:
+ load_env_file(self.env_file)
+ self._env_loaded = True
+ except FileNotFoundError:
+ logger.warning(f"Environment file not found: {self.env_file}")
+
+ @classmethod
+ def from_config(cls, config: Dict[str, Any]) -> "CredentialManager":
+ """
+ Create a CredentialManager from configuration.
+
+ Args:
+ config: Configuration dictionary containing:
+ - credentials.env_substitution: bool (default True)
+ - credentials.env_file: str (optional path to .env file)
+
+ Returns:
+ Configured CredentialManager instance
+ """
+ cred_config = config.get("credentials", {})
+ return cls(
+ env_substitution=cred_config.get("env_substitution", True),
+ env_file=cred_config.get("env_file")
+ )
+
+ def process_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
+ """
+ Process a configuration dictionary, substituting environment variables.
+
+ Args:
+ config: Configuration dictionary to process
+
+ Returns:
+ Configuration with environment variables substituted
+ """
+ if not self.env_substitution:
+ return config
+
+ return substitute_env_vars(config, self.env_file)
+
+ def get_credential(
+ self,
+ config: Dict[str, Any],
+ key: str,
+ required: bool = True
+ ) -> Optional[str]:
+ """
+ Get a credential value from configuration.
+
+ This method retrieves a credential, performing environment variable
+ substitution if enabled.
+
+ Args:
+ config: Configuration dictionary
+ key: Key to look up
+ required: Whether to raise an error if missing
+
+ Returns:
+ The credential value, or None if not found and not required
+
+ Raises:
+ ValueError: If required credential is missing
+ """
+ value = config.get(key)
+ if value is None:
+ if required:
+ raise ValueError(f"Required credential '{key}' is not configured")
+ return None
+
+ if self.env_substitution and isinstance(value, str):
+ value = _substitute_in_string(value)
+
+ # Check if substitution failed (still contains ${...})
+ if isinstance(value, str) and ENV_VAR_PATTERN.search(value):
+ unresolved = ENV_VAR_PATTERN.findall(value)
+ if required:
+ raise ValueError(
+ f"Credential '{key}' contains unresolved environment variables: "
+ f"{', '.join(unresolved)}"
+ )
+ logger.warning(
+ f"Credential '{key}' has unresolved env vars: {unresolved}"
+ )
+
+ return value
+
+ def validate_credentials(
+ self,
+ config: Dict[str, Any],
+ required_keys: list
+ ) -> list:
+ """
+ Validate that required credentials are present and resolved.
+
+ Args:
+ config: Configuration dictionary
+ required_keys: List of required credential keys
+
+ Returns:
+ List of validation error messages (empty if valid)
+ """
+ errors = []
+
+ for key in required_keys:
+ try:
+ value = self.get_credential(config, key, required=True)
+ if not value:
+ errors.append(f"Credential '{key}' is empty")
+ except ValueError as e:
+ errors.append(str(e))
+
+ return errors
+
+ @staticmethod
+ def mask_credential(value: str, show_chars: int = 4) -> str:
+ """
+ Mask a credential value for safe logging.
+
+ Args:
+ value: The credential value to mask
+ show_chars: Number of characters to show at the end
+
+ Returns:
+ Masked value like '***abc123'
+ """
+ if not value or len(value) <= show_chars:
+ return '***'
+
+ return '***' + value[-show_chars:]
+
+ def get_service_account_credentials(
+ self,
+ config: Dict[str, Any],
+ credentials_file_key: str = "credentials_file"
+ ) -> Optional[Dict[str, Any]]:
+ """
+ Load service account credentials from a JSON file.
+
+ Args:
+ config: Configuration dictionary
+ credentials_file_key: Key containing path to credentials file
+
+ Returns:
+ Parsed credentials dictionary, or None if not configured
+
+ Raises:
+ FileNotFoundError: If credentials file doesn't exist
+ ValueError: If credentials file is invalid JSON
+ """
+ import json
+
+ cred_file = config.get(credentials_file_key)
+ if not cred_file:
+ return None
+
+ # Substitute env vars in the path
+ if self.env_substitution:
+ cred_file = _substitute_in_string(cred_file)
+
+ cred_path = Path(cred_file)
+ if not cred_path.exists():
+ raise FileNotFoundError(
+ f"Service account credentials file not found: {cred_file}"
+ )
+
+ try:
+ with open(cred_path, 'r', encoding='utf-8') as f:
+ credentials = json.load(f)
+ except json.JSONDecodeError as e:
+ raise ValueError(
+ f"Invalid JSON in credentials file {cred_file}: {e}"
+ )
+
+ logger.debug(f"Loaded service account credentials from {cred_file}")
+ return credentials
diff --git a/potato/data_sources/manager.py b/potato/data_sources/manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..b09e5ad41e9787bbb9c8cba40188eb7c1e228687
--- /dev/null
+++ b/potato/data_sources/manager.py
@@ -0,0 +1,538 @@
+"""
+Data Source Manager
+
+This module provides the central manager for all data sources, implementing
+the singleton pattern for thread-safe access across the application.
+"""
+
+import logging
+import threading
+from typing import Any, Dict, Iterator, List, Optional, Type, TYPE_CHECKING
+
+from potato.data_sources.base import DataSource, SourceConfig, SourceType
+from potato.data_sources.credentials import CredentialManager
+from potato.data_sources.cache_manager import CacheManager
+from potato.data_sources.partial_reader import PartialReader, PartialLoadingConfig
+
+if TYPE_CHECKING:
+ from potato.item_state_management import ItemStateManager
+
+logger = logging.getLogger(__name__)
+
+# Singleton instance with thread-safe initialization
+DATA_SOURCE_MANAGER: Optional["DataSourceManager"] = None
+_MANAGER_LOCK = threading.Lock()
+
+
+# Registry of source type implementations
+_SOURCE_REGISTRY: Dict[SourceType, Type[DataSource]] = {}
+
+
+def register_source_type(source_type: SourceType, source_class: Type[DataSource]) -> None:
+ """
+ Register a data source implementation.
+
+ Args:
+ source_type: The SourceType enum value
+ source_class: The DataSource subclass
+ """
+ _SOURCE_REGISTRY[source_type] = source_class
+ logger.debug(f"Registered source type: {source_type.value} -> {source_class.__name__}")
+
+
+def get_source_class(source_type: SourceType) -> Optional[Type[DataSource]]:
+ """
+ Get the source class for a given type.
+
+ Args:
+ source_type: The SourceType to look up
+
+ Returns:
+ The DataSource subclass, or None if not registered
+ """
+ return _SOURCE_REGISTRY.get(source_type)
+
+
+def get_registered_types() -> List[str]:
+ """Get list of registered source type names."""
+ return [t.value for t in _SOURCE_REGISTRY.keys()]
+
+
+class DataSourceManager:
+ """
+ Central manager for all data sources.
+
+ This class provides:
+ - Registration and lifecycle management of data sources
+ - Credential management with environment variable substitution
+ - Caching for remote sources
+ - Partial/incremental loading coordination
+ - Thread-safe access to all sources
+
+ Attributes:
+ config: The application configuration
+ credential_manager: Handles credential resolution
+ cache_manager: Manages cached remote files
+ partial_reader: Coordinates incremental loading
+ """
+
+ def __init__(
+ self,
+ config: Dict[str, Any],
+ item_state_manager: "ItemStateManager"
+ ):
+ """
+ Initialize the data source manager.
+
+ Args:
+ config: Application configuration dictionary
+ item_state_manager: The ItemStateManager for adding items
+ """
+ self._config = config
+ self._item_state_manager = item_state_manager
+ self._sources: Dict[str, DataSource] = {}
+ self._lock = threading.RLock()
+
+ # Initialize sub-managers
+ self.credential_manager = CredentialManager.from_config(config)
+
+ # Set up cache manager if caching is enabled
+ cache_config = config.get("data_cache", {})
+ if cache_config.get("enabled", True):
+ cache_dir = cache_config.get(
+ "cache_dir",
+ ".potato_cache/data_sources"
+ )
+ # Resolve relative to task_dir
+ task_dir = config.get("task_dir", ".")
+ if not cache_dir.startswith("/"):
+ import os
+ cache_dir = os.path.join(task_dir, cache_dir)
+
+ self.cache_manager = CacheManager(
+ cache_dir=cache_dir,
+ ttl_seconds=cache_config.get("ttl_seconds", 3600),
+ max_size_mb=cache_config.get("max_size_mb", 500)
+ )
+ else:
+ self.cache_manager = None
+
+ # Set up partial reader if incremental loading is configured
+ partial_config = PartialLoadingConfig.from_dict(config)
+ if partial_config.enabled:
+ output_dir = config.get("output_annotation_dir", ".")
+ self.partial_reader = PartialReader(partial_config, output_dir)
+ else:
+ self.partial_reader = None
+
+ # Get item property keys
+ item_props = config.get("item_properties", {})
+ self._id_key = item_props.get("id_key", "id")
+ self._text_key = item_props.get("text_key", "text")
+
+ # Initialize sources from configuration
+ self._init_sources()
+
+ logger.info(f"DataSourceManager initialized with {len(self._sources)} sources")
+
+ def _init_sources(self) -> None:
+ """Initialize data sources from configuration."""
+ data_sources = self._config.get("data_sources", [])
+
+ for index, source_dict in enumerate(data_sources):
+ try:
+ # Process credentials in the source config
+ processed_config = self.credential_manager.process_config(source_dict)
+
+ # Parse source configuration
+ source_config = SourceConfig.from_dict(processed_config, index)
+
+ if not source_config.enabled:
+ logger.debug(f"Skipping disabled source: {source_config.source_id}")
+ continue
+
+ # Get the source class for this type
+ source_class = get_source_class(source_config.source_type)
+ if not source_class:
+ logger.warning(
+ f"No implementation for source type: {source_config.source_type.value}. "
+ f"Available types: {get_registered_types()}"
+ )
+ continue
+
+ # Create the source instance
+ source = source_class(source_config)
+
+ # Validate configuration
+ errors = source.validate_config()
+ if errors:
+ logger.error(
+ f"Invalid configuration for source {source_config.source_id}: "
+ f"{'; '.join(errors)}"
+ )
+ continue
+
+ # Check availability
+ if not source.is_available():
+ logger.warning(
+ f"Source {source_config.source_id} is not available. "
+ f"Check dependencies and credentials."
+ )
+ # Still register the source, but log the warning
+ # It may become available later
+
+ self._sources[source_config.source_id] = source
+ logger.info(
+ f"Initialized source: {source_config.source_id} "
+ f"(type={source_config.source_type.value})"
+ )
+
+ except Exception as e:
+ logger.error(f"Failed to initialize source at index {index}: {e}")
+
+ def get_source(self, source_id: str) -> Optional[DataSource]:
+ """
+ Get a data source by ID.
+
+ Args:
+ source_id: The source identifier
+
+ Returns:
+ The DataSource instance, or None if not found
+ """
+ with self._lock:
+ return self._sources.get(source_id)
+
+ def get_all_sources(self) -> Dict[str, DataSource]:
+ """
+ Get all registered sources.
+
+ Returns:
+ Dictionary mapping source_id to DataSource
+ """
+ with self._lock:
+ return dict(self._sources)
+
+ def list_sources(self) -> List[Dict[str, Any]]:
+ """
+ List all sources with their status.
+
+ Returns:
+ List of source status dictionaries
+ """
+ with self._lock:
+ statuses = []
+ for source in self._sources.values():
+ status = source.get_status()
+
+ # Add partial loading state if available
+ if self.partial_reader:
+ state = self.partial_reader.get_state(source.source_id)
+ if state:
+ status["items_loaded"] = state.items_loaded
+ status["is_complete"] = state.is_complete
+ status["last_loaded_at"] = state.last_loaded_at
+
+ statuses.append(status)
+
+ return statuses
+
+ def load_initial_data(self) -> int:
+ """
+ Load initial data from all sources.
+
+ If partial loading is enabled, loads only the initial_count items
+ from each source. Otherwise, loads all data.
+
+ Returns:
+ Total number of items loaded
+ """
+ total_loaded = 0
+
+ with self._lock:
+ for source_id, source in self._sources.items():
+ try:
+ count = self._load_from_source(source, is_initial=True)
+ total_loaded += count
+ logger.info(f"Loaded {count} items from {source_id}")
+ except Exception as e:
+ logger.error(f"Failed to load from {source_id}: {e}")
+
+ return total_loaded
+
+ def load_more(
+ self,
+ source_id: str,
+ count: Optional[int] = None
+ ) -> int:
+ """
+ Load more items from a specific source.
+
+ Args:
+ source_id: The source to load from
+ count: Number of items to load (uses batch_size if not specified)
+
+ Returns:
+ Number of items loaded
+
+ Raises:
+ ValueError: If source_id is not found
+ """
+ with self._lock:
+ source = self._sources.get(source_id)
+ if not source:
+ raise ValueError(f"Unknown source: {source_id}")
+
+ return self._load_from_source(source, is_initial=False, count=count)
+
+ def _load_from_source(
+ self,
+ source: DataSource,
+ is_initial: bool = True,
+ count: Optional[int] = None
+ ) -> int:
+ """
+ Load items from a source into the ItemStateManager.
+
+ Args:
+ source: The data source
+ is_initial: Whether this is the initial load
+ count: Number of items to load (overrides config)
+
+ Returns:
+ Number of items loaded
+ """
+ source_id = source.source_id
+
+ # Check if source is complete
+ if self.partial_reader:
+ state = self.partial_reader.get_state(source_id)
+ if state and state.is_complete:
+ logger.debug(f"Source {source_id} is already complete")
+ return 0
+
+ # Determine how many items to load and from what position
+ if self.partial_reader and self.partial_reader.config.enabled:
+ start = self.partial_reader.get_start_position(source_id)
+ if count is None:
+ count = self.partial_reader.get_load_count(source_id, is_initial)
+ else:
+ start = 0
+ count = None # Load all
+
+ # Check if source supports partial reading
+ if start > 0 and not source.supports_partial_reading():
+ logger.warning(
+ f"Source {source_id} does not support partial reading, "
+ f"cannot continue from position {start}"
+ )
+ return 0
+
+ # Load items
+ items_loaded = 0
+ is_complete = False
+
+ try:
+ for item in source.read_items(start=start, count=count):
+ # Validate ID key exists
+ if self._id_key not in item:
+ logger.warning(
+ f"Missing id_key '{self._id_key}' in item from {source_id}"
+ )
+ continue
+
+ instance_id = str(item[self._id_key])
+
+ # Check for duplicates
+ if self._item_state_manager.has_item(instance_id):
+ logger.debug(f"Skipping duplicate ID: {instance_id}")
+ continue
+
+ # Add item to state manager
+ try:
+ self._item_state_manager.add_item(instance_id, item)
+ items_loaded += 1
+ except ValueError as e:
+ logger.warning(f"Failed to add item {instance_id}: {e}")
+
+ # Check if we loaded fewer items than requested (source exhausted)
+ if count is not None and items_loaded < count:
+ is_complete = True
+
+ except StopIteration:
+ is_complete = True
+
+ # Update partial reader state
+ if self.partial_reader:
+ total_estimate = source.get_total_count()
+ self.partial_reader.update_state(
+ source_id=source_id,
+ items_added=items_loaded,
+ is_complete=is_complete,
+ total_estimate=total_estimate
+ )
+
+ return items_loaded
+
+ def refresh_source(self, source_id: str) -> bool:
+ """
+ Refresh a data source (re-fetch from remote).
+
+ Args:
+ source_id: The source to refresh
+
+ Returns:
+ True if refresh was successful
+
+ Raises:
+ ValueError: If source_id is not found
+ """
+ with self._lock:
+ source = self._sources.get(source_id)
+ if not source:
+ raise ValueError(f"Unknown source: {source_id}")
+
+ # Invalidate cache
+ if self.cache_manager:
+ self.cache_manager.invalidate(source_id)
+
+ # Reset partial reader state
+ if self.partial_reader:
+ self.partial_reader.reset_state(source_id)
+
+ return source.refresh()
+
+ def check_auto_load(
+ self,
+ annotated_count: int,
+ total_loaded: int
+ ) -> Dict[str, int]:
+ """
+ Check if any sources should auto-load more data.
+
+ Args:
+ annotated_count: Total number of annotated items
+ total_loaded: Total number of loaded items
+
+ Returns:
+ Dictionary mapping source_id to items loaded (for sources that triggered)
+ """
+ if not self.partial_reader or not self.partial_reader.config.auto_load_enabled:
+ return {}
+
+ results = {}
+
+ with self._lock:
+ for source_id, source in self._sources.items():
+ if self.partial_reader.should_load_more(
+ source_id, annotated_count, total_loaded
+ ):
+ try:
+ loaded = self._load_from_source(source, is_initial=False)
+ if loaded > 0:
+ results[source_id] = loaded
+ logger.info(
+ f"Auto-loaded {loaded} items from {source_id}"
+ )
+ except Exception as e:
+ logger.error(f"Auto-load failed for {source_id}: {e}")
+
+ return results
+
+ def clear_cache(self) -> int:
+ """
+ Clear the download cache for all sources.
+
+ Returns:
+ Number of cache entries cleared
+ """
+ if self.cache_manager:
+ return self.cache_manager.clear()
+ return 0
+
+ def get_stats(self) -> Dict[str, Any]:
+ """
+ Get comprehensive statistics.
+
+ Returns:
+ Dictionary with source and manager statistics
+ """
+ stats = {
+ "source_count": len(self._sources),
+ "sources": self.list_sources(),
+ }
+
+ if self.cache_manager:
+ stats["cache"] = self.cache_manager.get_stats()
+
+ if self.partial_reader:
+ stats["partial_loading"] = self.partial_reader.get_stats()
+
+ return stats
+
+ def close(self) -> None:
+ """Close all sources and release resources."""
+ with self._lock:
+ for source in self._sources.values():
+ try:
+ source.close()
+ except Exception as e:
+ logger.warning(f"Error closing source {source.source_id}: {e}")
+
+ self._sources.clear()
+
+
+def init_data_source_manager(config: Dict[str, Any]) -> Optional[DataSourceManager]:
+ """
+ Initialize the global DataSourceManager singleton.
+
+ This function creates the manager if data_sources is configured in
+ the configuration. Thread-safe initialization using double-checked
+ locking pattern.
+
+ Args:
+ config: Application configuration dictionary
+
+ Returns:
+ The DataSourceManager instance, or None if not configured
+ """
+ global DATA_SOURCE_MANAGER
+
+ # Check if data_sources is configured
+ if "data_sources" not in config or not config["data_sources"]:
+ return None
+
+ # Double-checked locking
+ if DATA_SOURCE_MANAGER is None:
+ with _MANAGER_LOCK:
+ if DATA_SOURCE_MANAGER is None:
+ from potato.item_state_management import get_item_state_manager
+ ism = get_item_state_manager()
+ DATA_SOURCE_MANAGER = DataSourceManager(config, ism)
+
+ return DATA_SOURCE_MANAGER
+
+
+def get_data_source_manager() -> Optional[DataSourceManager]:
+ """
+ Get the global DataSourceManager singleton.
+
+ Returns:
+ The DataSourceManager instance, or None if not initialized
+ """
+ return DATA_SOURCE_MANAGER
+
+
+def clear_data_source_manager() -> None:
+ """
+ Clear the global DataSourceManager singleton.
+
+ This function closes all sources and clears the singleton instance.
+ Thread-safe. Used primarily for testing.
+ """
+ global DATA_SOURCE_MANAGER
+
+ with _MANAGER_LOCK:
+ if DATA_SOURCE_MANAGER is not None:
+ DATA_SOURCE_MANAGER.close()
+ DATA_SOURCE_MANAGER = None
diff --git a/potato/data_sources/partial_reader.py b/potato/data_sources/partial_reader.py
new file mode 100644
index 0000000000000000000000000000000000000000..802854dc8367969d82f59296de40fc04f17c112d
--- /dev/null
+++ b/potato/data_sources/partial_reader.py
@@ -0,0 +1,407 @@
+"""
+Partial/incremental loading for data sources.
+
+This module provides functionality for loading data in chunks, enabling:
+- Initial loading of first K items
+- Batch loading of additional items as annotation progresses
+- Auto-loading when annotation reaches a threshold
+- State persistence for resumption after restart
+"""
+
+import json
+import logging
+import os
+import threading
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from potato.data_sources.base import DataSource
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PartialReadState:
+ """
+ Tracks the read state for a data source.
+
+ This dataclass maintains information about how much data has been
+ read from a source, enabling incremental loading.
+
+ Attributes:
+ source_id: Identifier of the data source
+ items_loaded: Number of items loaded so far
+ total_estimate: Estimated total items (None if unknown)
+ file_position: Byte position for file-based sources
+ line_number: Line number for line-based sources
+ is_complete: Whether all data has been loaded
+ last_loaded_at: Unix timestamp of last load
+ metadata: Additional source-specific state
+ """
+ source_id: str
+ items_loaded: int = 0
+ total_estimate: Optional[int] = None
+ file_position: int = 0
+ line_number: int = 0
+ is_complete: bool = False
+ last_loaded_at: Optional[float] = None
+ metadata: Dict[str, Any] = field(default_factory=dict)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to dictionary for JSON serialization."""
+ return {
+ "source_id": self.source_id,
+ "items_loaded": self.items_loaded,
+ "total_estimate": self.total_estimate,
+ "file_position": self.file_position,
+ "line_number": self.line_number,
+ "is_complete": self.is_complete,
+ "last_loaded_at": self.last_loaded_at,
+ "metadata": self.metadata,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "PartialReadState":
+ """Create from dictionary."""
+ return cls(
+ source_id=data["source_id"],
+ items_loaded=data.get("items_loaded", 0),
+ total_estimate=data.get("total_estimate"),
+ file_position=data.get("file_position", 0),
+ line_number=data.get("line_number", 0),
+ is_complete=data.get("is_complete", False),
+ last_loaded_at=data.get("last_loaded_at"),
+ metadata=data.get("metadata", {}),
+ )
+
+
+@dataclass
+class PartialLoadingConfig:
+ """
+ Configuration for partial/incremental loading.
+
+ Attributes:
+ enabled: Whether partial loading is enabled
+ initial_count: Number of items to load initially
+ batch_size: Number of items per incremental load
+ auto_load_threshold: Auto-load when this fraction is annotated (0.0-1.0)
+ auto_load_enabled: Whether auto-loading is enabled
+ """
+ enabled: bool = False
+ initial_count: int = 1000
+ batch_size: int = 500
+ auto_load_threshold: float = 0.8
+ auto_load_enabled: bool = True
+
+ @classmethod
+ def from_dict(cls, config: Dict[str, Any]) -> "PartialLoadingConfig":
+ """Create from configuration dictionary."""
+ partial_config = config.get("partial_loading", {})
+ return cls(
+ enabled=partial_config.get("enabled", False),
+ initial_count=partial_config.get("initial_count", 1000),
+ batch_size=partial_config.get("batch_size", 500),
+ auto_load_threshold=partial_config.get("auto_load_threshold", 0.8),
+ auto_load_enabled=partial_config.get("auto_load_enabled", True),
+ )
+
+ def validate(self) -> List[str]:
+ """Validate the configuration."""
+ errors = []
+
+ if self.initial_count < 1:
+ errors.append("partial_loading.initial_count must be at least 1")
+
+ if self.batch_size < 1:
+ errors.append("partial_loading.batch_size must be at least 1")
+
+ if not 0.0 <= self.auto_load_threshold <= 1.0:
+ errors.append(
+ "partial_loading.auto_load_threshold must be between 0.0 and 1.0"
+ )
+
+ return errors
+
+
+class PartialReader:
+ """
+ Manages partial/incremental loading of data sources.
+
+ This class coordinates loading data in chunks across multiple sources,
+ tracking state for each source and providing auto-loading when
+ annotation progress reaches a threshold.
+
+ Thread Safety:
+ All public methods are thread-safe. Internal state is protected
+ by a lock.
+
+ Attributes:
+ config: Partial loading configuration
+ state_file: Path to state persistence file
+ """
+
+ STATE_FILENAME = ".data_source_state.json"
+
+ def __init__(
+ self,
+ config: PartialLoadingConfig,
+ output_dir: str
+ ):
+ """
+ Initialize the partial reader.
+
+ Args:
+ config: Partial loading configuration
+ output_dir: Directory for state persistence
+ """
+ self.config = config
+ self.output_dir = Path(output_dir)
+ self.state_file = self.output_dir / self.STATE_FILENAME
+
+ self._states: Dict[str, PartialReadState] = {}
+ self._lock = threading.RLock()
+
+ # Load existing state
+ self._load_state()
+
+ logger.info(
+ f"PartialReader initialized: initial={config.initial_count}, "
+ f"batch={config.batch_size}, auto_threshold={config.auto_load_threshold}"
+ )
+
+ def _load_state(self) -> None:
+ """Load persisted state from disk."""
+ if not self.state_file.exists():
+ return
+
+ try:
+ with open(self.state_file, 'r', encoding='utf-8') as f:
+ data = json.load(f)
+
+ for state_data in data.get("sources", []):
+ state = PartialReadState.from_dict(state_data)
+ self._states[state.source_id] = state
+
+ logger.debug(f"Loaded partial read state for {len(self._states)} sources")
+
+ except (json.JSONDecodeError, KeyError) as e:
+ logger.warning(f"Failed to load partial read state: {e}")
+ self._states = {}
+
+ def _save_state(self) -> None:
+ """Save state to disk."""
+ # Ensure output directory exists
+ self.output_dir.mkdir(parents=True, exist_ok=True)
+
+ data = {
+ "version": 1,
+ "sources": [state.to_dict() for state in self._states.values()]
+ }
+
+ try:
+ with open(self.state_file, 'w', encoding='utf-8') as f:
+ json.dump(data, f, indent=2)
+ except IOError as e:
+ logger.error(f"Failed to save partial read state: {e}")
+
+ def get_state(self, source_id: str) -> Optional[PartialReadState]:
+ """
+ Get the current state for a source.
+
+ Args:
+ source_id: The source identifier
+
+ Returns:
+ PartialReadState or None if not tracked
+ """
+ with self._lock:
+ return self._states.get(source_id)
+
+ def get_or_create_state(self, source_id: str) -> PartialReadState:
+ """
+ Get or create state for a source.
+
+ Args:
+ source_id: The source identifier
+
+ Returns:
+ PartialReadState for the source
+ """
+ with self._lock:
+ if source_id not in self._states:
+ self._states[source_id] = PartialReadState(source_id=source_id)
+ self._save_state()
+ return self._states[source_id]
+
+ def update_state(
+ self,
+ source_id: str,
+ items_added: int,
+ file_position: Optional[int] = None,
+ line_number: Optional[int] = None,
+ is_complete: bool = False,
+ total_estimate: Optional[int] = None
+ ) -> PartialReadState:
+ """
+ Update the state after loading items.
+
+ Args:
+ source_id: The source identifier
+ items_added: Number of items added in this batch
+ file_position: New file position (for file sources)
+ line_number: New line number (for line-based sources)
+ is_complete: Whether all data has been loaded
+ total_estimate: Updated total estimate
+
+ Returns:
+ Updated PartialReadState
+ """
+ import time
+
+ with self._lock:
+ state = self.get_or_create_state(source_id)
+ state.items_loaded += items_added
+ state.last_loaded_at = time.time()
+
+ if file_position is not None:
+ state.file_position = file_position
+ if line_number is not None:
+ state.line_number = line_number
+ if total_estimate is not None:
+ state.total_estimate = total_estimate
+ if is_complete:
+ state.is_complete = True
+
+ self._save_state()
+ return state
+
+ def mark_complete(self, source_id: str) -> None:
+ """
+ Mark a source as completely loaded.
+
+ Args:
+ source_id: The source identifier
+ """
+ with self._lock:
+ state = self.get_or_create_state(source_id)
+ state.is_complete = True
+ self._save_state()
+
+ def should_load_more(
+ self,
+ source_id: str,
+ annotated_count: int,
+ total_loaded: int
+ ) -> bool:
+ """
+ Check if more data should be loaded based on annotation progress.
+
+ This method implements the auto-load threshold logic.
+
+ Args:
+ source_id: The source identifier
+ annotated_count: Number of items annotated
+ total_loaded: Total number of items loaded
+
+ Returns:
+ True if more data should be loaded
+ """
+ if not self.config.auto_load_enabled:
+ return False
+
+ with self._lock:
+ state = self._states.get(source_id)
+ if state and state.is_complete:
+ return False
+
+ if total_loaded == 0:
+ return False
+
+ progress = annotated_count / total_loaded
+ return progress >= self.config.auto_load_threshold
+
+ def get_load_count(
+ self,
+ source_id: str,
+ is_initial: bool = False
+ ) -> int:
+ """
+ Get the number of items to load.
+
+ Args:
+ source_id: The source identifier
+ is_initial: Whether this is the initial load
+
+ Returns:
+ Number of items to load
+ """
+ if is_initial:
+ return self.config.initial_count
+ return self.config.batch_size
+
+ def get_start_position(self, source_id: str) -> int:
+ """
+ Get the starting position for the next load.
+
+ Args:
+ source_id: The source identifier
+
+ Returns:
+ Number of items already loaded (start position for next batch)
+ """
+ with self._lock:
+ state = self._states.get(source_id)
+ if state:
+ return state.items_loaded
+ return 0
+
+ def reset_state(self, source_id: str) -> None:
+ """
+ Reset the state for a source.
+
+ Args:
+ source_id: The source identifier
+ """
+ with self._lock:
+ if source_id in self._states:
+ del self._states[source_id]
+ self._save_state()
+
+ def clear_all_state(self) -> None:
+ """Clear state for all sources."""
+ with self._lock:
+ self._states.clear()
+ self._save_state()
+
+ def get_all_states(self) -> Dict[str, PartialReadState]:
+ """
+ Get all source states.
+
+ Returns:
+ Dictionary mapping source_id to PartialReadState
+ """
+ with self._lock:
+ return dict(self._states)
+
+ def get_stats(self) -> Dict[str, Any]:
+ """
+ Get statistics about partial loading.
+
+ Returns:
+ Dictionary with loading statistics
+ """
+ with self._lock:
+ total_loaded = sum(s.items_loaded for s in self._states.values())
+ complete_count = sum(1 for s in self._states.values() if s.is_complete)
+
+ return {
+ "enabled": self.config.enabled,
+ "initial_count": self.config.initial_count,
+ "batch_size": self.config.batch_size,
+ "auto_load_threshold": self.config.auto_load_threshold,
+ "sources_tracked": len(self._states),
+ "sources_complete": complete_count,
+ "total_items_loaded": total_loaded,
+ }
diff --git a/potato/data_sources/sources/__init__.py b/potato/data_sources/sources/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..23b78a64b14b67ef01ab93de97913044c458bc9f
--- /dev/null
+++ b/potato/data_sources/sources/__init__.py
@@ -0,0 +1,82 @@
+"""
+Data source implementations.
+
+This module registers all available data source types. Source implementations
+are lazy-loaded to avoid importing optional dependencies.
+
+To add a new source type:
+1. Create a new module (e.g., my_source.py) with a DataSource subclass
+2. Import and register it in this file
+3. Add any required dependencies to the requirements.txt (as optional)
+"""
+
+import logging
+
+from potato.data_sources.base import SourceType
+from potato.data_sources.manager import register_source_type
+
+logger = logging.getLogger(__name__)
+
+
+def _register_all_sources() -> None:
+ """Register all available source implementations."""
+
+ # Local file source - always available
+ try:
+ from potato.data_sources.sources.local_source import LocalFileSource
+ register_source_type(SourceType.FILE, LocalFileSource)
+ except ImportError as e:
+ logger.debug(f"LocalFileSource not available: {e}")
+
+ # URL source - always available (uses standard library)
+ try:
+ from potato.data_sources.sources.url_source import URLSource
+ register_source_type(SourceType.URL, URLSource)
+ except ImportError as e:
+ logger.debug(f"URLSource not available: {e}")
+
+ # Google Drive source - requires google-api-python-client
+ try:
+ from potato.data_sources.sources.gdrive_source import GoogleDriveSource
+ register_source_type(SourceType.GOOGLE_DRIVE, GoogleDriveSource)
+ except ImportError as e:
+ logger.debug(f"GoogleDriveSource not available: {e}")
+
+ # Dropbox source - requires dropbox
+ try:
+ from potato.data_sources.sources.dropbox_source import DropboxSource
+ register_source_type(SourceType.DROPBOX, DropboxSource)
+ except ImportError as e:
+ logger.debug(f"DropboxSource not available: {e}")
+
+ # S3 source - requires boto3
+ try:
+ from potato.data_sources.sources.s3_source import S3Source
+ register_source_type(SourceType.S3, S3Source)
+ except ImportError as e:
+ logger.debug(f"S3Source not available: {e}")
+
+ # HuggingFace source - requires datasets
+ try:
+ from potato.data_sources.sources.huggingface_source import HuggingFaceSource
+ register_source_type(SourceType.HUGGINGFACE, HuggingFaceSource)
+ except ImportError as e:
+ logger.debug(f"HuggingFaceSource not available: {e}")
+
+ # Google Sheets source - requires google-api-python-client
+ try:
+ from potato.data_sources.sources.gsheets_source import GoogleSheetsSource
+ register_source_type(SourceType.GOOGLE_SHEETS, GoogleSheetsSource)
+ except ImportError as e:
+ logger.debug(f"GoogleSheetsSource not available: {e}")
+
+ # Database source - requires sqlalchemy
+ try:
+ from potato.data_sources.sources.database_source import DatabaseSource
+ register_source_type(SourceType.DATABASE, DatabaseSource)
+ except ImportError as e:
+ logger.debug(f"DatabaseSource not available: {e}")
+
+
+# Register sources when module is imported
+_register_all_sources()
diff --git a/potato/data_sources/sources/database_source.py b/potato/data_sources/sources/database_source.py
new file mode 100644
index 0000000000000000000000000000000000000000..60a5a09b85a1b2226a643dfac779cf7cecea156c
--- /dev/null
+++ b/potato/data_sources/sources/database_source.py
@@ -0,0 +1,317 @@
+"""
+SQL Database data source.
+
+This module provides data loading from SQL databases using SQLAlchemy,
+supporting PostgreSQL, MySQL, SQLite, and other databases.
+"""
+
+import logging
+import re
+from typing import Any, Dict, Iterator, List, Optional
+from urllib.parse import quote_plus
+
+from potato.data_sources.base import DataSource, SourceConfig
+
+logger = logging.getLogger(__name__)
+
+
+class DatabaseSource(DataSource):
+ """
+ Data source for SQL databases.
+
+ Loads data from SQL databases using SQLAlchemy, supporting:
+ - PostgreSQL, MySQL, SQLite
+ - Custom SQL queries or simple table select
+ - Connection via connection string or individual parameters
+ - Incremental loading via OFFSET/LIMIT
+
+ Configuration with connection string:
+ type: database
+ connection_string: "${DATABASE_URL}"
+ query: "SELECT id, text, metadata FROM items WHERE status = 'pending'"
+
+ Configuration with individual parameters:
+ type: database
+ dialect: postgresql # postgresql, mysql, sqlite
+ host: "localhost"
+ port: 5432
+ database: "annotations"
+ username: "${DB_USER}"
+ password: "${DB_PASSWORD}"
+ table: "items" # Simple table select
+ id_column: "id"
+ text_column: "text"
+
+ Note: Requires SQLAlchemy and appropriate database driver:
+ pip install sqlalchemy psycopg2-binary # PostgreSQL
+ pip install sqlalchemy pymysql # MySQL
+ """
+
+ # Check for optional dependencies
+ _HAS_SQLALCHEMY = None
+
+ @classmethod
+ def _check_dependencies(cls) -> bool:
+ """Check if SQLAlchemy is available."""
+ if cls._HAS_SQLALCHEMY is None:
+ try:
+ import sqlalchemy
+ cls._HAS_SQLALCHEMY = True
+ except ImportError:
+ cls._HAS_SQLALCHEMY = False
+ return cls._HAS_SQLALCHEMY
+
+ # Pattern for safe SQL identifiers (table/column names)
+ # Allows: word chars, dots for schema.table, backticks/brackets for quoted identifiers
+ _SAFE_IDENTIFIER_RE = re.compile(r'\A[\w][\w.$]*\Z', re.ASCII)
+
+ @staticmethod
+ def _validate_identifier(name: str) -> str:
+ """
+ Validate a SQL identifier (table or column name) against injection.
+
+ Only allows alphanumeric characters, underscores, dots (for schema.table),
+ and dollar signs. Rejects anything else to prevent SQL injection.
+
+ Raises:
+ ValueError: If the identifier contains unsafe characters
+ """
+ if not name or not DatabaseSource._SAFE_IDENTIFIER_RE.match(name):
+ raise ValueError(
+ f"Invalid SQL identifier: '{name}'. "
+ f"Only alphanumeric characters, underscores, dots, and "
+ f"dollar signs are allowed."
+ )
+ return name
+
+ # Dialect to driver mapping
+ DIALECT_DRIVERS = {
+ 'postgresql': 'postgresql+psycopg2',
+ 'postgres': 'postgresql+psycopg2',
+ 'mysql': 'mysql+pymysql',
+ 'sqlite': 'sqlite',
+ 'mssql': 'mssql+pyodbc',
+ }
+
+ def __init__(self, config: SourceConfig):
+ """Initialize the database source."""
+ super().__init__(config)
+
+ # Connection options
+ self._connection_string = config.config.get("connection_string", "")
+ self._dialect = config.config.get("dialect", "")
+ self._host = config.config.get("host", "localhost")
+ self._port = config.config.get("port")
+ self._database = config.config.get("database", "")
+ self._username = config.config.get("username", "")
+ self._password = config.config.get("password", "")
+
+ # Query options
+ self._query = config.config.get("query", "")
+ self._table = config.config.get("table", "")
+ self._id_column = config.config.get("id_column", "id")
+ self._text_column = config.config.get("text_column", "text")
+
+ # Connection pooling options
+ self._pool_size = config.config.get("pool_size", 5)
+ self._pool_timeout = config.config.get("pool_timeout", 30)
+
+ self._engine = None
+ self._total_count: Optional[int] = None
+
+ def get_source_id(self) -> str:
+ """Get unique identifier."""
+ return self._source_id
+
+ def validate_config(self) -> List[str]:
+ """Validate source configuration."""
+ errors = []
+
+ # Must have connection string OR individual parameters
+ if not self._connection_string:
+ if not self._dialect:
+ errors.append(
+ "Either 'connection_string' or 'dialect' is required"
+ )
+ elif self._dialect not in self.DIALECT_DRIVERS:
+ errors.append(
+ f"Unknown dialect '{self._dialect}'. "
+ f"Supported: {', '.join(self.DIALECT_DRIVERS.keys())}"
+ )
+
+ if not self._database and self._dialect != 'sqlite':
+ errors.append("'database' is required")
+
+ # Must have query OR table
+ if not self._query and not self._table:
+ errors.append("Either 'query' or 'table' is required")
+
+ # Validate table name if provided (prevent SQL injection)
+ if self._table:
+ try:
+ self._validate_identifier(self._table)
+ except ValueError as e:
+ errors.append(str(e))
+
+ return errors
+
+ def is_available(self) -> bool:
+ """Check if the source is available."""
+ if not self._check_dependencies():
+ logger.warning(
+ "SQLAlchemy not installed. "
+ "Install with: pip install sqlalchemy"
+ )
+ return False
+
+ return True
+
+ def _build_connection_string(self) -> str:
+ """Build connection string from individual parameters."""
+ if self._connection_string:
+ return self._connection_string
+
+ driver = self.DIALECT_DRIVERS.get(self._dialect, self._dialect)
+
+ if self._dialect == 'sqlite':
+ return f"sqlite:///{self._database}"
+
+ # Build URL with credentials
+ if self._username:
+ userpass = self._username
+ if self._password:
+ userpass += f":{quote_plus(self._password)}"
+ userpass += "@"
+ else:
+ userpass = ""
+
+ host_port = self._host
+ if self._port:
+ host_port += f":{self._port}"
+
+ return f"{driver}://{userpass}{host_port}/{self._database}"
+
+ def _get_engine(self):
+ """Get or create the SQLAlchemy engine."""
+ if self._engine:
+ return self._engine
+
+ from sqlalchemy import create_engine
+
+ connection_string = self._build_connection_string()
+
+ # Create engine with connection pooling
+ engine_kwargs = {}
+ if self._dialect != 'sqlite':
+ engine_kwargs = {
+ 'pool_size': self._pool_size,
+ 'pool_timeout': self._pool_timeout,
+ 'pool_pre_ping': True, # Enable connection health checks
+ }
+
+ self._engine = create_engine(connection_string, **engine_kwargs)
+ return self._engine
+
+ def _build_query(self, offset: int = 0, limit: Optional[int] = None) -> str:
+ """Build the SQL query with optional pagination."""
+ if self._query:
+ base_query = self._query.rstrip(';')
+ else:
+ # Validate table name to prevent SQL injection
+ safe_table = self._validate_identifier(self._table)
+ base_query = f"SELECT * FROM {safe_table}"
+
+ # Add pagination using validated integer values
+ if limit is not None or offset > 0:
+ if limit is not None:
+ base_query += f" LIMIT {int(limit)}"
+ if offset > 0:
+ base_query += f" OFFSET {int(offset)}"
+
+ return base_query
+
+ def _row_to_dict(self, row, columns: List[str]) -> Dict[str, Any]:
+ """Convert a database row to a dictionary."""
+ item = {}
+ for i, col in enumerate(columns):
+ value = row[i]
+ # Handle special types
+ if hasattr(value, 'isoformat'): # datetime
+ value = value.isoformat()
+ elif hasattr(value, 'tobytes'): # memoryview/bytes
+ value = value.tobytes().decode('utf-8', errors='replace')
+ item[col] = value
+ return item
+
+ def read_items(
+ self,
+ start: int = 0,
+ count: Optional[int] = None
+ ) -> Iterator[Dict[str, Any]]:
+ """Read items from the database."""
+ from sqlalchemy import text
+
+ engine = self._get_engine()
+ query = self._build_query(offset=start, limit=count)
+
+ with engine.connect() as connection:
+ result = connection.execute(text(query))
+
+ # Get column names
+ columns = list(result.keys())
+
+ for row in result:
+ item = self._row_to_dict(row, columns)
+ yield item
+
+ def get_total_count(self) -> Optional[int]:
+ """Get total number of items."""
+ if self._total_count is not None:
+ return self._total_count
+
+ from sqlalchemy import text
+
+ try:
+ engine = self._get_engine()
+
+ if self._query:
+ # Wrap query in count (query is admin-provided from YAML config)
+ count_query = f"SELECT COUNT(*) FROM ({self._query.rstrip(';')}) AS subquery"
+ else:
+ # Validate table name to prevent SQL injection
+ safe_table = self._validate_identifier(self._table)
+ count_query = f"SELECT COUNT(*) FROM {safe_table}"
+
+ with engine.connect() as connection:
+ result = connection.execute(text(count_query))
+ self._total_count = result.scalar()
+ return self._total_count
+
+ except Exception as e:
+ logger.error(f"Error getting count: {e}")
+ return None
+
+ def supports_partial_reading(self) -> bool:
+ """Database sources support efficient partial reading via OFFSET/LIMIT."""
+ return True
+
+ def refresh(self) -> bool:
+ """Refresh by clearing cached count."""
+ self._total_count = None
+ return True
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get source status."""
+ status = super().get_status()
+ status["dialect"] = self._dialect
+ status["database"] = self._database
+ status["table"] = self._table
+ status["has_custom_query"] = bool(self._query)
+ return status
+
+ def close(self) -> None:
+ """Close the database connection."""
+ if self._engine:
+ self._engine.dispose()
+ self._engine = None
+ self._total_count = None
diff --git a/potato/data_sources/sources/dropbox_source.py b/potato/data_sources/sources/dropbox_source.py
new file mode 100644
index 0000000000000000000000000000000000000000..07318c69cbb7960a5ac587240d8b908fe9119482
--- /dev/null
+++ b/potato/data_sources/sources/dropbox_source.py
@@ -0,0 +1,294 @@
+"""
+Dropbox data source.
+
+This module provides data loading from Dropbox files,
+supporting both public share links and authenticated access.
+"""
+
+import json
+import logging
+import re
+from typing import Any, Dict, Iterator, List, Optional
+from urllib.parse import urlparse, parse_qs
+
+from potato.data_sources.base import DataSource, SourceConfig
+
+logger = logging.getLogger(__name__)
+
+
+def convert_share_link(url: str) -> str:
+ """
+ Convert a Dropbox share link to a direct download URL.
+
+ Args:
+ url: Dropbox share link
+
+ Returns:
+ Direct download URL
+
+ Examples:
+ https://www.dropbox.com/s/xxx/file.json?dl=0
+ -> https://www.dropbox.com/s/xxx/file.json?dl=1
+ """
+ parsed = urlparse(url)
+
+ # Check if it's a Dropbox URL
+ if 'dropbox.com' not in parsed.netloc:
+ raise ValueError(f"Not a Dropbox URL: {url}")
+
+ # Convert dl=0 to dl=1 for direct download
+ if 'dl=0' in url:
+ return url.replace('dl=0', 'dl=1')
+ elif 'dl=1' in url:
+ return url
+ else:
+ # Add dl=1 parameter
+ separator = '&' if '?' in url else '?'
+ return f"{url}{separator}dl=1"
+
+
+class DropboxSource(DataSource):
+ """
+ Data source for Dropbox files.
+
+ Supports both public share links (no authentication required)
+ and private files with access token authentication.
+
+ Configuration for public files:
+ type: dropbox
+ url: "https://www.dropbox.com/s/xxx/file.jsonl?dl=0"
+
+ Configuration for private files:
+ type: dropbox
+ path: "/path/to/file.jsonl" # Path in Dropbox
+ access_token: "${DROPBOX_TOKEN}"
+
+ Supported formats: JSON, JSONL, CSV, TSV
+ """
+
+ # Check for optional dependencies
+ _HAS_DROPBOX = None
+
+ @classmethod
+ def _check_dependencies(cls) -> bool:
+ """Check if Dropbox SDK is available."""
+ if cls._HAS_DROPBOX is None:
+ try:
+ import dropbox
+ cls._HAS_DROPBOX = True
+ except ImportError:
+ cls._HAS_DROPBOX = False
+ return cls._HAS_DROPBOX
+
+ def __init__(self, config: SourceConfig):
+ """Initialize the Dropbox source."""
+ super().__init__(config)
+
+ self._url = config.config.get("url", "")
+ self._path = config.config.get("path", "")
+ self._access_token = config.config.get("access_token")
+
+ self._cached_data: Optional[List[Dict]] = None
+ self._client = None
+
+ def get_source_id(self) -> str:
+ """Get unique identifier."""
+ return self._source_id
+
+ def validate_config(self) -> List[str]:
+ """Validate source configuration."""
+ errors = []
+
+ if not self._url and not self._path:
+ errors.append(
+ "Either 'url' or 'path' is required for Dropbox source"
+ )
+ return errors
+
+ # If path is provided, token is required
+ if self._path and not self._access_token:
+ errors.append(
+ "'access_token' is required when using 'path' for private files"
+ )
+
+ # Validate URL format if provided
+ if self._url:
+ try:
+ convert_share_link(self._url)
+ except ValueError as e:
+ errors.append(str(e))
+
+ return errors
+
+ def is_available(self) -> bool:
+ """Check if the source is available."""
+ # For authenticated access, check dependencies
+ if self._access_token:
+ if not self._check_dependencies():
+ logger.warning(
+ "Dropbox SDK not installed. "
+ "Install with: pip install dropbox"
+ )
+ return False
+
+ return True
+
+ def _get_client(self):
+ """Get or create the Dropbox client."""
+ if self._client:
+ return self._client
+
+ if not self._access_token:
+ return None
+
+ import dropbox
+ self._client = dropbox.Dropbox(self._access_token)
+ return self._client
+
+ def _fetch_public_file(self, url: str) -> bytes:
+ """Fetch a public file using direct download URL."""
+ import urllib.request
+ import urllib.error
+
+ download_url = convert_share_link(url)
+
+ request = urllib.request.Request(download_url)
+ request.add_header('User-Agent', 'Potato-Annotation-Tool/1.0')
+
+ try:
+ with urllib.request.urlopen(request, timeout=60) as response:
+ return response.read()
+ except urllib.error.HTTPError as e:
+ if e.code == 404:
+ raise ValueError("File not found or link has expired")
+ raise RuntimeError(f"HTTP error {e.code}: {e.reason}")
+ except urllib.error.URLError as e:
+ raise RuntimeError(f"URL error: {e.reason}")
+
+ def _fetch_authenticated_file(self, path: str) -> bytes:
+ """Fetch a file using authenticated API access."""
+ client = self._get_client()
+ if not client:
+ raise RuntimeError("No access token configured")
+
+ import dropbox
+
+ try:
+ # Ensure path starts with /
+ if not path.startswith('/'):
+ path = '/' + path
+
+ metadata, response = client.files_download(path)
+ logger.debug(f"Downloaded: {metadata.name} ({metadata.size} bytes)")
+ return response.content
+
+ except dropbox.exceptions.ApiError as e:
+ if e.error.is_path():
+ raise ValueError(f"File not found: {path}")
+ raise RuntimeError(f"Dropbox API error: {e}")
+
+ def _fetch_data(self) -> List[Dict[str, Any]]:
+ """Fetch and parse data from Dropbox."""
+ # Fetch the file
+ if self._url:
+ content = self._fetch_public_file(self._url)
+ else:
+ content = self._fetch_authenticated_file(self._path)
+
+ # Decode and parse
+ text = content.decode('utf-8')
+ return self._parse_content(text)
+
+ def _parse_content(self, text: str) -> List[Dict[str, Any]]:
+ """Parse file content."""
+ # Try JSON array first
+ try:
+ data = json.loads(text)
+ if isinstance(data, list):
+ return data
+ elif isinstance(data, dict):
+ return [data]
+ except json.JSONDecodeError:
+ pass
+
+ # Try JSONL
+ items = []
+ lines = text.strip().split('\n')
+ for line in lines:
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ item = json.loads(line)
+ if isinstance(item, list):
+ items.extend(item)
+ else:
+ items.append(item)
+ except json.JSONDecodeError:
+ pass
+
+ if items:
+ return items
+
+ # Try CSV
+ import csv
+ from io import StringIO
+
+ try:
+ reader = csv.DictReader(StringIO(text))
+ items = [dict(row) for row in reader]
+ if items:
+ return items
+ except Exception:
+ pass
+
+ raise ValueError("Could not parse file content as JSON, JSONL, or CSV")
+
+ def read_items(
+ self,
+ start: int = 0,
+ count: Optional[int] = None
+ ) -> Iterator[Dict[str, Any]]:
+ """Read items from Dropbox file."""
+ if self._cached_data is None:
+ self._cached_data = self._fetch_data()
+
+ items = self._cached_data[start:]
+ if count is not None:
+ items = items[:count]
+
+ yield from items
+
+ def get_total_count(self) -> Optional[int]:
+ """Get total number of items."""
+ if self._cached_data is None:
+ try:
+ self._cached_data = self._fetch_data()
+ except Exception as e:
+ logger.error(f"Error fetching data: {e}")
+ return None
+
+ return len(self._cached_data)
+
+ def supports_partial_reading(self) -> bool:
+ """Partial reading is supported after initial fetch."""
+ return True
+
+ def refresh(self) -> bool:
+ """Refresh by clearing cached data."""
+ self._cached_data = None
+ return True
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get source status."""
+ status = super().get_status()
+ status["url"] = self._url
+ status["path"] = self._path
+ status["authenticated"] = self._access_token is not None
+ status["cached"] = self._cached_data is not None
+ return status
+
+ def close(self) -> None:
+ """Close the source."""
+ self._client = None
+ self._cached_data = None
diff --git a/potato/data_sources/sources/gdrive_source.py b/potato/data_sources/sources/gdrive_source.py
new file mode 100644
index 0000000000000000000000000000000000000000..a0acdf27d3add286e2897034bf6f2091bd45501b
--- /dev/null
+++ b/potato/data_sources/sources/gdrive_source.py
@@ -0,0 +1,349 @@
+"""
+Google Drive data source.
+
+This module provides data loading from Google Drive files,
+supporting both public share links and authenticated access
+via service account credentials.
+"""
+
+import io
+import json
+import logging
+import re
+from typing import Any, Dict, Iterator, List, Optional
+
+from potato.data_sources.base import DataSource, SourceConfig
+
+logger = logging.getLogger(__name__)
+
+# Patterns for extracting file ID from various Google Drive URL formats
+GDRIVE_URL_PATTERNS = [
+ # https://drive.google.com/file/d/FILE_ID/view
+ re.compile(r'drive\.google\.com/file/d/([a-zA-Z0-9_-]+)'),
+ # https://drive.google.com/open?id=FILE_ID
+ re.compile(r'drive\.google\.com/open\?id=([a-zA-Z0-9_-]+)'),
+ # https://docs.google.com/document/d/FILE_ID/edit
+ re.compile(r'docs\.google\.com/\w+/d/([a-zA-Z0-9_-]+)'),
+ # https://drive.google.com/uc?id=FILE_ID
+ re.compile(r'drive\.google\.com/uc\?.*id=([a-zA-Z0-9_-]+)'),
+]
+
+
+def extract_file_id(url_or_id: str) -> str:
+ """
+ Extract Google Drive file ID from a URL or return the ID directly.
+
+ Args:
+ url_or_id: Either a Google Drive URL or a file ID
+
+ Returns:
+ The file ID
+
+ Raises:
+ ValueError: If the URL format is not recognized
+ """
+ # Check if it's already a file ID (no slashes or dots)
+ if not ('/' in url_or_id or '.' in url_or_id):
+ return url_or_id
+
+ # Try each URL pattern
+ for pattern in GDRIVE_URL_PATTERNS:
+ match = pattern.search(url_or_id)
+ if match:
+ return match.group(1)
+
+ raise ValueError(
+ f"Could not extract Google Drive file ID from: {url_or_id}. "
+ f"Please provide a valid Google Drive URL or file ID."
+ )
+
+
+class GoogleDriveSource(DataSource):
+ """
+ Data source for Google Drive files.
+
+ Supports both public share links (no authentication required)
+ and private files with service account credentials.
+
+ Configuration for public files:
+ type: google_drive
+ url: "https://drive.google.com/file/d/xxx/view?usp=sharing"
+
+ Configuration for private files:
+ type: google_drive
+ file_id: "xxx" # Or use url
+ credentials_file: "credentials/gdrive_service_account.json"
+
+ Supported formats: JSON, JSONL, CSV, TSV
+ """
+
+ # Check for optional dependencies
+ _HAS_GOOGLE_API = None
+
+ @classmethod
+ def _check_dependencies(cls) -> bool:
+ """Check if Google API dependencies are available."""
+ if cls._HAS_GOOGLE_API is None:
+ try:
+ from google.oauth2 import service_account
+ from googleapiclient.discovery import build
+ cls._HAS_GOOGLE_API = True
+ except ImportError:
+ cls._HAS_GOOGLE_API = False
+ return cls._HAS_GOOGLE_API
+
+ def __init__(self, config: SourceConfig):
+ """Initialize the Google Drive source."""
+ super().__init__(config)
+
+ self._url = config.config.get("url", "")
+ self._file_id = config.config.get("file_id", "")
+ self._credentials_file = config.config.get("credentials_file")
+
+ # Resolve file ID from URL if provided
+ if self._url and not self._file_id:
+ self._file_id = extract_file_id(self._url)
+
+ self._cached_data: Optional[List[Dict]] = None
+ self._service = None
+
+ def get_source_id(self) -> str:
+ """Get unique identifier."""
+ return self._source_id
+
+ def validate_config(self) -> List[str]:
+ """Validate source configuration."""
+ errors = []
+
+ if not self._url and not self._file_id:
+ errors.append("Either 'url' or 'file_id' is required for Google Drive source")
+ return errors
+
+ # Try to extract file ID
+ try:
+ if self._url and not self._file_id:
+ extract_file_id(self._url)
+ except ValueError as e:
+ errors.append(str(e))
+
+ return errors
+
+ def is_available(self) -> bool:
+ """Check if the source is available."""
+ # For authenticated access, check dependencies and credentials
+ if self._credentials_file:
+ if not self._check_dependencies():
+ logger.warning(
+ "Google API dependencies not installed. "
+ "Install with: pip install google-api-python-client google-auth"
+ )
+ return False
+
+ import os
+ if not os.path.exists(self._credentials_file):
+ logger.warning(
+ f"Credentials file not found: {self._credentials_file}"
+ )
+ return False
+
+ return True
+
+ def _get_service(self):
+ """Get or create the Google Drive API service."""
+ if self._service:
+ return self._service
+
+ if not self._credentials_file:
+ return None
+
+ from google.oauth2 import service_account
+ from googleapiclient.discovery import build
+
+ credentials = service_account.Credentials.from_service_account_file(
+ self._credentials_file,
+ scopes=['https://www.googleapis.com/auth/drive.readonly']
+ )
+
+ self._service = build('drive', 'v3', credentials=credentials)
+ return self._service
+
+ def _fetch_public_file(self, file_id: str) -> bytes:
+ """Fetch a public file using direct download URL."""
+ import urllib.request
+ import urllib.error
+
+ # Construct direct download URL
+ # This works for publicly shared files
+ download_url = f"https://drive.google.com/uc?export=download&id={file_id}"
+
+ request = urllib.request.Request(download_url)
+ request.add_header('User-Agent', 'Potato-Annotation-Tool/1.0')
+
+ try:
+ with urllib.request.urlopen(request, timeout=60) as response:
+ content = response.read()
+
+ # Check for virus scan warning (large files)
+ if b'Google Drive - Virus scan warning' in content:
+ # Extract confirmation token and retry
+ import re
+ confirm_match = re.search(
+ rb'confirm=([0-9A-Za-z_-]+)',
+ content
+ )
+ if confirm_match:
+ confirm_token = confirm_match.group(1).decode()
+ download_url = (
+ f"https://drive.google.com/uc?export=download"
+ f"&confirm={confirm_token}&id={file_id}"
+ )
+ request = urllib.request.Request(download_url)
+ request.add_header('User-Agent', 'Potato-Annotation-Tool/1.0')
+ with urllib.request.urlopen(request, timeout=60) as response2:
+ content = response2.read()
+
+ return content
+
+ except urllib.error.HTTPError as e:
+ if e.code == 404:
+ raise ValueError(
+ f"File not found. Make sure the file is publicly shared."
+ )
+ raise RuntimeError(f"HTTP error {e.code}: {e.reason}")
+ except urllib.error.URLError as e:
+ raise RuntimeError(f"URL error: {e.reason}")
+
+ def _fetch_authenticated_file(self, file_id: str) -> bytes:
+ """Fetch a file using authenticated API access."""
+ service = self._get_service()
+ if not service:
+ raise RuntimeError(
+ "No credentials configured for authenticated access"
+ )
+
+ from googleapiclient.http import MediaIoBaseDownload
+
+ # Get file metadata
+ file_metadata = service.files().get(fileId=file_id).execute()
+ logger.debug(f"Fetching file: {file_metadata.get('name')}")
+
+ # Download file content
+ request = service.files().get_media(fileId=file_id)
+ buffer = io.BytesIO()
+ downloader = MediaIoBaseDownload(buffer, request)
+
+ done = False
+ while not done:
+ status, done = downloader.next_chunk()
+ if status:
+ logger.debug(f"Download progress: {int(status.progress() * 100)}%")
+
+ buffer.seek(0)
+ return buffer.read()
+
+ def _fetch_data(self) -> List[Dict[str, Any]]:
+ """Fetch and parse data from Google Drive."""
+ file_id = self._file_id
+
+ # Fetch the file
+ if self._credentials_file:
+ content = self._fetch_authenticated_file(file_id)
+ else:
+ content = self._fetch_public_file(file_id)
+
+ # Decode and parse
+ text = content.decode('utf-8')
+ return self._parse_content(text)
+
+ def _parse_content(self, text: str) -> List[Dict[str, Any]]:
+ """Parse file content."""
+ # Try JSON array first
+ try:
+ data = json.loads(text)
+ if isinstance(data, list):
+ return data
+ elif isinstance(data, dict):
+ return [data]
+ except json.JSONDecodeError:
+ pass
+
+ # Try JSONL
+ items = []
+ lines = text.strip().split('\n')
+ for line in lines:
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ item = json.loads(line)
+ if isinstance(item, list):
+ items.extend(item)
+ else:
+ items.append(item)
+ except json.JSONDecodeError:
+ pass
+
+ if items:
+ return items
+
+ # Try CSV
+ import csv
+ from io import StringIO
+
+ try:
+ reader = csv.DictReader(StringIO(text))
+ items = [dict(row) for row in reader]
+ if items:
+ return items
+ except Exception:
+ pass
+
+ raise ValueError("Could not parse file content as JSON, JSONL, or CSV")
+
+ def read_items(
+ self,
+ start: int = 0,
+ count: Optional[int] = None
+ ) -> Iterator[Dict[str, Any]]:
+ """Read items from Google Drive file."""
+ if self._cached_data is None:
+ self._cached_data = self._fetch_data()
+
+ items = self._cached_data[start:]
+ if count is not None:
+ items = items[:count]
+
+ yield from items
+
+ def get_total_count(self) -> Optional[int]:
+ """Get total number of items."""
+ if self._cached_data is None:
+ try:
+ self._cached_data = self._fetch_data()
+ except Exception as e:
+ logger.error(f"Error fetching data: {e}")
+ return None
+
+ return len(self._cached_data)
+
+ def supports_partial_reading(self) -> bool:
+ """Partial reading is supported after initial fetch."""
+ return True
+
+ def refresh(self) -> bool:
+ """Refresh by clearing cached data."""
+ self._cached_data = None
+ return True
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get source status."""
+ status = super().get_status()
+ status["file_id"] = self._file_id
+ status["authenticated"] = self._credentials_file is not None
+ status["cached"] = self._cached_data is not None
+ return status
+
+ def close(self) -> None:
+ """Close the source."""
+ self._service = None
+ self._cached_data = None
diff --git a/potato/data_sources/sources/gsheets_source.py b/potato/data_sources/sources/gsheets_source.py
new file mode 100644
index 0000000000000000000000000000000000000000..c093af9a9e35603faf8e952007d994faf26fd0b6
--- /dev/null
+++ b/potato/data_sources/sources/gsheets_source.py
@@ -0,0 +1,234 @@
+"""
+Google Sheets data source.
+
+This module provides data loading from Google Sheets spreadsheets,
+supporting service account authentication.
+"""
+
+import logging
+from typing import Any, Dict, Iterator, List, Optional
+
+from potato.data_sources.base import DataSource, SourceConfig
+
+logger = logging.getLogger(__name__)
+
+
+class GoogleSheetsSource(DataSource):
+ """
+ Data source for Google Sheets.
+
+ Loads data from Google Sheets using the Sheets API with
+ service account authentication.
+
+ Configuration:
+ type: google_sheets
+ spreadsheet_id: "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs74OgvE2upms"
+ sheet_name: "Sheet1" # Optional: sheet name (default: first sheet)
+ range: "A:Z" # Optional: range to read
+ credentials_file: "credentials/service_account.json"
+
+ # Header options
+ header_row: 1 # Row containing headers (1-indexed)
+ skip_rows: 0 # Rows to skip after header
+
+ Note: Requires google-api-python-client:
+ pip install google-api-python-client google-auth
+ """
+
+ # Check for optional dependencies
+ _HAS_GOOGLE_API = None
+
+ @classmethod
+ def _check_dependencies(cls) -> bool:
+ """Check if Google API dependencies are available."""
+ if cls._HAS_GOOGLE_API is None:
+ try:
+ from google.oauth2 import service_account
+ from googleapiclient.discovery import build
+ cls._HAS_GOOGLE_API = True
+ except ImportError:
+ cls._HAS_GOOGLE_API = False
+ return cls._HAS_GOOGLE_API
+
+ def __init__(self, config: SourceConfig):
+ """Initialize the Google Sheets source."""
+ super().__init__(config)
+
+ self._spreadsheet_id = config.config.get("spreadsheet_id", "")
+ self._sheet_name = config.config.get("sheet_name")
+ self._range = config.config.get("range", "A:Z")
+ self._credentials_file = config.config.get("credentials_file")
+
+ self._header_row = config.config.get("header_row", 1)
+ self._skip_rows = config.config.get("skip_rows", 0)
+
+ self._service = None
+ self._cached_data: Optional[List[Dict]] = None
+
+ def get_source_id(self) -> str:
+ """Get unique identifier."""
+ return self._source_id
+
+ def validate_config(self) -> List[str]:
+ """Validate source configuration."""
+ errors = []
+
+ if not self._spreadsheet_id:
+ errors.append("'spreadsheet_id' is required for Google Sheets source")
+
+ if not self._credentials_file:
+ errors.append("'credentials_file' is required for Google Sheets source")
+
+ return errors
+
+ def is_available(self) -> bool:
+ """Check if the source is available."""
+ if not self._check_dependencies():
+ logger.warning(
+ "Google API dependencies not installed. "
+ "Install with: pip install google-api-python-client google-auth"
+ )
+ return False
+
+ import os
+ if self._credentials_file and not os.path.exists(self._credentials_file):
+ logger.warning(f"Credentials file not found: {self._credentials_file}")
+ return False
+
+ return True
+
+ def _get_service(self):
+ """Get or create the Sheets API service."""
+ if self._service:
+ return self._service
+
+ from google.oauth2 import service_account
+ from googleapiclient.discovery import build
+
+ credentials = service_account.Credentials.from_service_account_file(
+ self._credentials_file,
+ scopes=['https://www.googleapis.com/auth/spreadsheets.readonly']
+ )
+
+ self._service = build('sheets', 'v4', credentials=credentials)
+ return self._service
+
+ def _fetch_data(self) -> List[Dict[str, Any]]:
+ """Fetch and parse data from Google Sheets."""
+ service = self._get_service()
+
+ # Construct the range
+ if self._sheet_name:
+ range_notation = f"'{self._sheet_name}'!{self._range}"
+ else:
+ range_notation = self._range
+
+ try:
+ result = service.spreadsheets().values().get(
+ spreadsheetId=self._spreadsheet_id,
+ range=range_notation,
+ valueRenderOption='UNFORMATTED_VALUE',
+ dateTimeRenderOption='FORMATTED_STRING'
+ ).execute()
+
+ values = result.get('values', [])
+
+ if not values:
+ logger.warning("No data found in spreadsheet")
+ return []
+
+ # Extract headers
+ header_index = self._header_row - 1 # Convert to 0-indexed
+ if header_index >= len(values):
+ raise ValueError(
+ f"Header row {self._header_row} is beyond data range"
+ )
+
+ headers = values[header_index]
+
+ # Clean up headers
+ headers = [str(h).strip() if h else f"column_{i}"
+ for i, h in enumerate(headers)]
+
+ # Extract data rows
+ data_start = header_index + 1 + self._skip_rows
+ data_rows = values[data_start:]
+
+ # Convert to list of dictionaries
+ items = []
+ for row_index, row in enumerate(data_rows):
+ # Skip empty rows
+ if not row or all(cell == '' or cell is None for cell in row):
+ continue
+
+ # Pad row if shorter than headers
+ while len(row) < len(headers):
+ row.append('')
+
+ item = {headers[i]: row[i] for i in range(len(headers))}
+
+ # Add row number as fallback ID if no 'id' column
+ if 'id' not in item:
+ item['_row_number'] = data_start + row_index + 1
+
+ items.append(item)
+
+ logger.info(
+ f"Loaded {len(items)} rows from spreadsheet "
+ f"(sheet={self._sheet_name or 'first'}, "
+ f"columns={len(headers)})"
+ )
+
+ return items
+
+ except Exception as e:
+ raise RuntimeError(f"Failed to fetch spreadsheet data: {e}")
+
+ def read_items(
+ self,
+ start: int = 0,
+ count: Optional[int] = None
+ ) -> Iterator[Dict[str, Any]]:
+ """Read items from Google Sheets."""
+ if self._cached_data is None:
+ self._cached_data = self._fetch_data()
+
+ items = self._cached_data[start:]
+ if count is not None:
+ items = items[:count]
+
+ yield from items
+
+ def get_total_count(self) -> Optional[int]:
+ """Get total number of rows."""
+ if self._cached_data is None:
+ try:
+ self._cached_data = self._fetch_data()
+ except Exception as e:
+ logger.error(f"Error fetching data: {e}")
+ return None
+
+ return len(self._cached_data)
+
+ def supports_partial_reading(self) -> bool:
+ """Partial reading is supported after initial fetch."""
+ return True
+
+ def refresh(self) -> bool:
+ """Refresh by clearing cached data."""
+ self._cached_data = None
+ return True
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get source status."""
+ status = super().get_status()
+ status["spreadsheet_id"] = self._spreadsheet_id
+ status["sheet_name"] = self._sheet_name
+ status["range"] = self._range
+ status["cached"] = self._cached_data is not None
+ return status
+
+ def close(self) -> None:
+ """Close the source."""
+ self._service = None
+ self._cached_data = None
diff --git a/potato/data_sources/sources/huggingface_source.py b/potato/data_sources/sources/huggingface_source.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d10720c1b55edea507e0194d1a7e2243ebc2010
--- /dev/null
+++ b/potato/data_sources/sources/huggingface_source.py
@@ -0,0 +1,245 @@
+"""
+Hugging Face Datasets data source.
+
+This module provides data loading from Hugging Face Hub datasets,
+supporting both public and private datasets.
+"""
+
+import logging
+from typing import Any, Dict, Iterator, List, Optional
+
+from potato.data_sources.base import DataSource, SourceConfig
+
+logger = logging.getLogger(__name__)
+
+
+class HuggingFaceSource(DataSource):
+ """
+ Data source for Hugging Face Hub datasets.
+
+ Loads data from Hugging Face's datasets library, supporting:
+ - Public datasets from the Hub
+ - Private datasets with authentication token
+ - Specific splits (train, validation, test)
+ - Dataset subsets/configurations
+
+ Configuration:
+ type: huggingface
+ dataset: "squad" # Required: dataset name
+ split: "train" # Optional: split name (default: train)
+ subset: null # Optional: dataset subset/config
+ token: "${HF_TOKEN}" # Optional: for private datasets
+
+ # Field mapping
+ id_field: "id" # Field to use as item ID
+ text_field: "context" # Field to use as text
+
+ Note: Requires the 'datasets' library: pip install datasets
+ """
+
+ # Check for optional dependencies
+ _HAS_DATASETS = None
+
+ @classmethod
+ def _check_dependencies(cls) -> bool:
+ """Check if datasets library is available."""
+ if cls._HAS_DATASETS is None:
+ try:
+ import datasets
+ cls._HAS_DATASETS = True
+ except ImportError:
+ cls._HAS_DATASETS = False
+ return cls._HAS_DATASETS
+
+ def __init__(self, config: SourceConfig):
+ """Initialize the HuggingFace source."""
+ super().__init__(config)
+
+ self._dataset_name = config.config.get("dataset", "")
+ self._split = config.config.get("split", "train")
+ self._subset = config.config.get("subset")
+ self._token = config.config.get("token")
+
+ # Field mapping for converting HF dataset to Potato items
+ self._id_field = config.config.get("id_field", "id")
+ self._text_field = config.config.get("text_field", "text")
+ self._include_fields = config.config.get("include_fields") # List or None
+
+ self._dataset = None
+ self._cached_items: Optional[List[Dict]] = None
+
+ def get_source_id(self) -> str:
+ """Get unique identifier."""
+ return self._source_id
+
+ def validate_config(self) -> List[str]:
+ """Validate source configuration."""
+ errors = []
+
+ if not self._dataset_name:
+ errors.append("'dataset' is required for HuggingFace source")
+
+ return errors
+
+ def is_available(self) -> bool:
+ """Check if the source is available."""
+ if not self._check_dependencies():
+ logger.warning(
+ "datasets library not installed. "
+ "Install with: pip install datasets"
+ )
+ return False
+
+ return True
+
+ def _load_dataset(self):
+ """Load the HuggingFace dataset."""
+ if self._dataset is not None:
+ return self._dataset
+
+ from datasets import load_dataset
+
+ load_kwargs = {
+ 'path': self._dataset_name,
+ 'split': self._split,
+ }
+
+ if self._subset:
+ load_kwargs['name'] = self._subset
+
+ if self._token:
+ load_kwargs['token'] = self._token
+
+ try:
+ self._dataset = load_dataset(**load_kwargs)
+ logger.info(
+ f"Loaded HuggingFace dataset: {self._dataset_name} "
+ f"(split={self._split}, {len(self._dataset)} examples)"
+ )
+ return self._dataset
+
+ except Exception as e:
+ raise RuntimeError(f"Failed to load dataset: {e}")
+
+ def _convert_example(self, example: Dict, index: int) -> Dict[str, Any]:
+ """Convert a HuggingFace example to a Potato item."""
+ item = {}
+
+ # Handle ID field
+ if self._id_field in example:
+ item['id'] = str(example[self._id_field])
+ else:
+ # Generate ID from index
+ item['id'] = f"{self._dataset_name}_{self._split}_{index}"
+
+ # Handle text field
+ if self._text_field in example:
+ item['text'] = example[self._text_field]
+
+ # Include specified fields or all fields
+ if self._include_fields:
+ for field in self._include_fields:
+ if field in example:
+ item[field] = example[field]
+ else:
+ # Include all fields from the example
+ for key, value in example.items():
+ if key not in item:
+ # Convert non-serializable types
+ item[key] = self._serialize_value(value)
+
+ return item
+
+ def _serialize_value(self, value: Any) -> Any:
+ """Convert a value to a JSON-serializable format."""
+ import numpy as np
+
+ if isinstance(value, (str, int, float, bool, type(None))):
+ return value
+ elif isinstance(value, (list, tuple)):
+ return [self._serialize_value(v) for v in value]
+ elif isinstance(value, dict):
+ return {k: self._serialize_value(v) for k, v in value.items()}
+ elif isinstance(value, np.ndarray):
+ return value.tolist()
+ elif hasattr(value, 'item'): # numpy scalar
+ return value.item()
+ else:
+ return str(value)
+
+ def _fetch_data(self) -> List[Dict[str, Any]]:
+ """Fetch and convert all data from the dataset."""
+ dataset = self._load_dataset()
+
+ items = []
+ for index, example in enumerate(dataset):
+ item = self._convert_example(example, index)
+ items.append(item)
+
+ return items
+
+ def read_items(
+ self,
+ start: int = 0,
+ count: Optional[int] = None
+ ) -> Iterator[Dict[str, Any]]:
+ """Read items from the HuggingFace dataset."""
+ # Use cached items if available
+ if self._cached_items is not None:
+ items = self._cached_items[start:]
+ if count is not None:
+ items = items[:count]
+ yield from items
+ return
+
+ # Load dataset
+ dataset = self._load_dataset()
+
+ # For partial reading, slice the dataset
+ end_index = None
+ if count is not None:
+ end_index = start + count
+
+ items_yielded = 0
+ for index, example in enumerate(dataset):
+ if index < start:
+ continue
+ if end_index is not None and index >= end_index:
+ break
+
+ item = self._convert_example(example, index)
+ yield item
+ items_yielded += 1
+
+ def get_total_count(self) -> Optional[int]:
+ """Get total number of items in the dataset."""
+ try:
+ dataset = self._load_dataset()
+ return len(dataset)
+ except Exception as e:
+ logger.error(f"Error getting dataset count: {e}")
+ return None
+
+ def supports_partial_reading(self) -> bool:
+ """HuggingFace datasets support efficient partial reading."""
+ return True
+
+ def refresh(self) -> bool:
+ """Refresh by reloading the dataset."""
+ self._dataset = None
+ self._cached_items = None
+ return True
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get source status."""
+ status = super().get_status()
+ status["dataset"] = self._dataset_name
+ status["split"] = self._split
+ status["subset"] = self._subset
+ status["loaded"] = self._dataset is not None
+ return status
+
+ def close(self) -> None:
+ """Close the source."""
+ self._dataset = None
+ self._cached_items = None
diff --git a/potato/data_sources/sources/local_source.py b/potato/data_sources/sources/local_source.py
new file mode 100644
index 0000000000000000000000000000000000000000..2dcc65984446dae88807142e406c4eec69c64676
--- /dev/null
+++ b/potato/data_sources/sources/local_source.py
@@ -0,0 +1,298 @@
+"""
+Local file data source.
+
+This module provides data loading from local files, supporting
+JSON, JSONL, CSV, and TSV formats with partial reading support.
+"""
+
+import csv
+import json
+import logging
+import os
+from typing import Any, Dict, Iterator, List, Optional
+
+from potato.data_sources.base import DataSource, SourceConfig
+
+logger = logging.getLogger(__name__)
+
+
+class LocalFileSource(DataSource):
+ """
+ Data source for local files.
+
+ Supports reading from JSON, JSONL, CSV, and TSV files with
+ optional partial reading for large files.
+
+ Configuration:
+ type: file
+ path: "data/annotations.jsonl" # Required: path to file
+
+ Supported formats:
+ - .json: JSON array or object per line
+ - .jsonl: JSON Lines (one JSON object per line)
+ - .csv: Comma-separated values
+ - .tsv: Tab-separated values
+ """
+
+ SUPPORTED_EXTENSIONS = ('.json', '.jsonl', '.csv', '.tsv')
+
+ def __init__(self, config: SourceConfig):
+ """
+ Initialize the local file source.
+
+ Args:
+ config: Source configuration
+ """
+ super().__init__(config)
+
+ self._path = config.config.get("path", "")
+ self._resolved_path: Optional[str] = None
+ self._total_count: Optional[int] = None
+ self._file_positions: Dict[int, int] = {} # line_number -> file_position
+
+ def get_source_id(self) -> str:
+ """Get unique identifier for this source."""
+ return self._source_id
+
+ def _resolve_path(self) -> str:
+ """Resolve the file path, validating relative paths stay within the task directory."""
+ if self._resolved_path:
+ return self._resolved_path
+
+ path = self._path
+ task_dir = os.path.abspath(self._raw_config.get("task_dir", "."))
+
+ # If path is relative, resolve against task_dir and validate containment
+ if not os.path.isabs(path):
+ resolved = os.path.abspath(os.path.join(task_dir, path))
+
+ # Ensure the resolved path is within the task directory
+ if not resolved.startswith(task_dir + os.sep) and resolved != task_dir:
+ raise ValueError(
+ f"Path '{self._path}' resolves to '{resolved}' which is "
+ f"outside the task directory '{task_dir}'. "
+ f"Path traversal is not allowed."
+ )
+ else:
+ # Absolute paths are used as-is (admin-provided via config)
+ resolved = os.path.abspath(path)
+
+ self._resolved_path = resolved
+ return self._resolved_path
+
+ def is_available(self) -> bool:
+ """Check if the file exists and is readable."""
+ try:
+ path = self._resolve_path()
+ if not os.path.exists(path):
+ logger.warning(f"File does not exist: {path}")
+ return False
+ if not os.path.isfile(path):
+ logger.warning(f"Path is not a file: {path}")
+ return False
+ if not os.access(path, os.R_OK):
+ logger.warning(f"File is not readable: {path}")
+ return False
+ return True
+ except Exception as e:
+ logger.error(f"Error checking file availability: {e}")
+ return False
+
+ def validate_config(self) -> List[str]:
+ """Validate source configuration."""
+ errors = []
+
+ if not self._path:
+ errors.append("'path' is required for file source")
+ return errors
+
+ # Check extension
+ ext = os.path.splitext(self._path)[1].lower()
+ if ext not in self.SUPPORTED_EXTENSIONS:
+ errors.append(
+ f"Unsupported file extension '{ext}'. "
+ f"Supported: {', '.join(self.SUPPORTED_EXTENSIONS)}"
+ )
+
+ return errors
+
+ def read_items(
+ self,
+ start: int = 0,
+ count: Optional[int] = None
+ ) -> Iterator[Dict[str, Any]]:
+ """
+ Read items from the file.
+
+ Args:
+ start: Index of first item to read (0-based)
+ count: Maximum number of items to read
+
+ Yields:
+ Item dictionaries
+ """
+ path = self._resolve_path()
+ ext = os.path.splitext(path)[1].lower()
+
+ if ext in ('.json', '.jsonl'):
+ yield from self._read_json_items(path, start, count)
+ elif ext == '.csv':
+ yield from self._read_csv_items(path, start, count, delimiter=',')
+ elif ext == '.tsv':
+ yield from self._read_csv_items(path, start, count, delimiter='\t')
+ else:
+ raise ValueError(f"Unsupported file format: {ext}")
+
+ def _read_json_items(
+ self,
+ path: str,
+ start: int,
+ count: Optional[int]
+ ) -> Iterator[Dict[str, Any]]:
+ """Read items from JSON/JSONL file."""
+ ext = os.path.splitext(path)[1].lower()
+
+ with open(path, 'r', encoding='utf-8') as f:
+ if ext == '.json':
+ # Try to parse as JSON array first
+ content = f.read()
+ try:
+ data = json.loads(content)
+ if isinstance(data, list):
+ # JSON array
+ items = data
+ elif isinstance(data, dict):
+ # Single object
+ items = [data]
+ else:
+ raise ValueError(f"Unexpected JSON type: {type(data)}")
+
+ # Apply start/count
+ items = items[start:]
+ if count is not None:
+ items = items[:count]
+
+ yield from items
+ return
+
+ except json.JSONDecodeError:
+ # Fall back to JSONL parsing
+ pass
+
+ # Reset file position for JSONL parsing
+ f.seek(0)
+
+ items_yielded = 0
+ current_line = 0
+
+ for line_no, line in enumerate(f):
+ line = line.strip()
+ if not line:
+ continue
+
+ # Skip lines before start
+ if current_line < start:
+ current_line += 1
+ continue
+
+ # Check count limit
+ if count is not None and items_yielded >= count:
+ break
+
+ try:
+ item = json.loads(line)
+ if isinstance(item, list):
+ # Line contains array - expand
+ for sub_item in item:
+ if count is not None and items_yielded >= count:
+ break
+ yield sub_item
+ items_yielded += 1
+ else:
+ yield item
+ items_yielded += 1
+ except json.JSONDecodeError as e:
+ logger.warning(f"Invalid JSON at line {line_no + 1}: {e}")
+
+ current_line += 1
+
+ def _read_csv_items(
+ self,
+ path: str,
+ start: int,
+ count: Optional[int],
+ delimiter: str
+ ) -> Iterator[Dict[str, Any]]:
+ """Read items from CSV/TSV file."""
+ with open(path, 'r', encoding='utf-8', newline='') as f:
+ reader = csv.DictReader(f, delimiter=delimiter)
+
+ items_yielded = 0
+ current_row = 0
+
+ for row in reader:
+ # Skip rows before start
+ if current_row < start:
+ current_row += 1
+ continue
+
+ # Check count limit
+ if count is not None and items_yielded >= count:
+ break
+
+ yield dict(row)
+ items_yielded += 1
+ current_row += 1
+
+ def get_total_count(self) -> Optional[int]:
+ """Get total number of items in the file."""
+ if self._total_count is not None:
+ return self._total_count
+
+ if not self.is_available():
+ return None
+
+ try:
+ path = self._resolve_path()
+ ext = os.path.splitext(path)[1].lower()
+
+ count = 0
+ if ext in ('.json', '.jsonl'):
+ with open(path, 'r', encoding='utf-8') as f:
+ content = f.read()
+ try:
+ data = json.loads(content)
+ if isinstance(data, list):
+ count = len(data)
+ else:
+ count = 1
+ except json.JSONDecodeError:
+ # JSONL - count non-empty lines
+ for line in content.split('\n'):
+ if line.strip():
+ count += 1
+
+ elif ext in ('.csv', '.tsv'):
+ delimiter = ',' if ext == '.csv' else '\t'
+ with open(path, 'r', encoding='utf-8', newline='') as f:
+ reader = csv.reader(f, delimiter=delimiter)
+ next(reader, None) # Skip header
+ count = sum(1 for _ in reader)
+
+ self._total_count = count
+ return count
+
+ except Exception as e:
+ logger.error(f"Error counting items: {e}")
+ return None
+
+ def supports_partial_reading(self) -> bool:
+ """Local files support partial reading."""
+ return True
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get source status."""
+ status = super().get_status()
+ status["path"] = self._path
+ status["resolved_path"] = self._resolve_path() if self.is_available() else None
+ return status
diff --git a/potato/data_sources/sources/s3_source.py b/potato/data_sources/sources/s3_source.py
new file mode 100644
index 0000000000000000000000000000000000000000..8a7a4ac8d8061d228523a3f7608b45d924fee04d
--- /dev/null
+++ b/potato/data_sources/sources/s3_source.py
@@ -0,0 +1,335 @@
+"""
+Amazon S3 data source.
+
+This module provides data loading from Amazon S3 buckets,
+supporting various authentication methods.
+"""
+
+import ipaddress
+import json
+import logging
+import socket
+from typing import Any, Dict, Iterator, List, Optional
+from urllib.parse import urlparse
+
+from potato.data_sources.base import DataSource, SourceConfig
+
+logger = logging.getLogger(__name__)
+
+
+class S3Source(DataSource):
+ """
+ Data source for Amazon S3 buckets.
+
+ Supports loading data from S3 with multiple authentication options:
+ - AWS credentials file (~/.aws/credentials)
+ - Environment variables (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)
+ - Explicit credentials in config
+ - S3-compatible storage (MinIO, etc.)
+
+ Configuration:
+ type: s3
+ bucket: "my-annotation-data" # Required
+ key: "datasets/items.jsonl" # Required
+ region: "us-east-1" # Optional, default us-east-1
+
+ # Optional: explicit credentials (prefer env vars)
+ access_key_id: "${AWS_ACCESS_KEY_ID}"
+ secret_access_key: "${AWS_SECRET_ACCESS_KEY}"
+
+ # Optional: for S3-compatible storage
+ endpoint_url: "https://minio.example.com"
+
+ Supported formats: JSON, JSONL, CSV, TSV
+ """
+
+ # Check for optional dependencies
+ _HAS_BOTO3 = None
+
+ @classmethod
+ def _check_dependencies(cls) -> bool:
+ """Check if boto3 is available."""
+ if cls._HAS_BOTO3 is None:
+ try:
+ import boto3
+ cls._HAS_BOTO3 = True
+ except ImportError:
+ cls._HAS_BOTO3 = False
+ return cls._HAS_BOTO3
+
+ def __init__(self, config: SourceConfig):
+ """Initialize the S3 source."""
+ super().__init__(config)
+
+ self._bucket = config.config.get("bucket", "")
+ self._key = config.config.get("key", "")
+ self._region = config.config.get("region", "us-east-1")
+ self._access_key_id = config.config.get("access_key_id")
+ self._secret_access_key = config.config.get("secret_access_key")
+ self._endpoint_url = config.config.get("endpoint_url")
+
+ self._cached_data: Optional[List[Dict]] = None
+ self._client = None
+
+ def get_source_id(self) -> str:
+ """Get unique identifier."""
+ return self._source_id
+
+ def validate_config(self) -> List[str]:
+ """Validate source configuration."""
+ errors = []
+
+ if not self._bucket:
+ errors.append("'bucket' is required for S3 source")
+
+ if not self._key:
+ errors.append("'key' is required for S3 source")
+
+ # Check that both access key and secret are provided together
+ if self._access_key_id and not self._secret_access_key:
+ errors.append(
+ "'secret_access_key' is required when 'access_key_id' is provided"
+ )
+ if self._secret_access_key and not self._access_key_id:
+ errors.append(
+ "'access_key_id' is required when 'secret_access_key' is provided"
+ )
+
+ # SSRF protection: validate endpoint_url does not point to private IPs
+ if self._endpoint_url:
+ try:
+ parsed = urlparse(self._endpoint_url)
+ if parsed.scheme not in ('http', 'https'):
+ errors.append(
+ f"Invalid endpoint_url scheme '{parsed.scheme}'. "
+ f"Only http/https allowed."
+ )
+ hostname = parsed.hostname
+ if hostname:
+ try:
+ addr_info = socket.getaddrinfo(hostname, None)
+ for info in addr_info:
+ ip_str = info[4][0]
+ try:
+ ip = ipaddress.ip_address(ip_str)
+ if ip.is_loopback or ip.is_link_local:
+ errors.append(
+ f"endpoint_url host '{hostname}' resolves "
+ f"to blocked IP {ip_str}. Loopback and "
+ f"link-local addresses are not allowed."
+ )
+ except ValueError:
+ pass
+ except socket.gaierror:
+ # Can't resolve at validation time โ will fail at connect
+ pass
+ except Exception as e:
+ errors.append(f"Invalid endpoint_url: {e}")
+
+ return errors
+
+ def is_available(self) -> bool:
+ """Check if the source is available."""
+ if not self._check_dependencies():
+ logger.warning(
+ "boto3 not installed. Install with: pip install boto3"
+ )
+ return False
+
+ return True
+
+ def _get_client(self):
+ """Get or create the S3 client."""
+ if self._client:
+ return self._client
+
+ import boto3
+
+ # Build client configuration
+ client_kwargs = {
+ 'region_name': self._region,
+ }
+
+ if self._endpoint_url:
+ client_kwargs['endpoint_url'] = self._endpoint_url
+
+ if self._access_key_id and self._secret_access_key:
+ client_kwargs['aws_access_key_id'] = self._access_key_id
+ client_kwargs['aws_secret_access_key'] = self._secret_access_key
+
+ self._client = boto3.client('s3', **client_kwargs)
+ return self._client
+
+ def _fetch_data(self) -> List[Dict[str, Any]]:
+ """Fetch and parse data from S3."""
+ client = self._get_client()
+
+ try:
+ response = client.get_object(Bucket=self._bucket, Key=self._key)
+ content = response['Body'].read()
+ content_type = response.get('ContentType', '')
+
+ logger.debug(
+ f"Downloaded s3://{self._bucket}/{self._key} "
+ f"({len(content)} bytes, {content_type})"
+ )
+
+ # Decode and parse
+ text = content.decode('utf-8')
+ return self._parse_content(text, content_type)
+
+ except client.exceptions.NoSuchKey:
+ raise ValueError(
+ f"Object not found: s3://{self._bucket}/{self._key}"
+ )
+ except client.exceptions.NoSuchBucket:
+ raise ValueError(f"Bucket not found: {self._bucket}")
+ except Exception as e:
+ raise RuntimeError(f"S3 error: {e}")
+
+ def _parse_content(
+ self,
+ text: str,
+ content_type: str = ""
+ ) -> List[Dict[str, Any]]:
+ """Parse file content based on content type or key extension."""
+ key_lower = self._key.lower()
+
+ # Determine format
+ is_json = 'json' in content_type or key_lower.endswith('.json')
+ is_jsonl = 'ndjson' in content_type or key_lower.endswith('.jsonl')
+ is_csv = 'csv' in content_type or key_lower.endswith('.csv')
+ is_tsv = 'tab' in content_type or key_lower.endswith('.tsv')
+
+ # Try JSON array first
+ if is_json or is_jsonl:
+ try:
+ data = json.loads(text)
+ if isinstance(data, list):
+ return data
+ elif isinstance(data, dict):
+ return [data]
+ except json.JSONDecodeError:
+ pass
+
+ # Try JSONL
+ if is_jsonl or is_json:
+ items = []
+ for line in text.strip().split('\n'):
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ item = json.loads(line)
+ if isinstance(item, list):
+ items.extend(item)
+ else:
+ items.append(item)
+ except json.JSONDecodeError:
+ pass
+
+ if items:
+ return items
+
+ # Try CSV/TSV
+ if is_csv or is_tsv:
+ import csv
+ from io import StringIO
+
+ delimiter = '\t' if is_tsv else ','
+ reader = csv.DictReader(StringIO(text), delimiter=delimiter)
+ return [dict(row) for row in reader]
+
+ # Auto-detect: try JSON, then JSONL, then CSV
+ try:
+ data = json.loads(text)
+ if isinstance(data, list):
+ return data
+ elif isinstance(data, dict):
+ return [data]
+ except json.JSONDecodeError:
+ pass
+
+ # Try JSONL
+ items = []
+ for line in text.strip().split('\n'):
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ item = json.loads(line)
+ if isinstance(item, list):
+ items.extend(item)
+ else:
+ items.append(item)
+ except json.JSONDecodeError:
+ pass
+
+ if items:
+ return items
+
+ # Try CSV as last resort
+ import csv
+ from io import StringIO
+
+ try:
+ reader = csv.DictReader(StringIO(text))
+ items = [dict(row) for row in reader]
+ if items:
+ return items
+ except Exception:
+ pass
+
+ raise ValueError(
+ f"Could not parse content from s3://{self._bucket}/{self._key}"
+ )
+
+ def read_items(
+ self,
+ start: int = 0,
+ count: Optional[int] = None
+ ) -> Iterator[Dict[str, Any]]:
+ """Read items from S3."""
+ if self._cached_data is None:
+ self._cached_data = self._fetch_data()
+
+ items = self._cached_data[start:]
+ if count is not None:
+ items = items[:count]
+
+ yield from items
+
+ def get_total_count(self) -> Optional[int]:
+ """Get total number of items."""
+ if self._cached_data is None:
+ try:
+ self._cached_data = self._fetch_data()
+ except Exception as e:
+ logger.error(f"Error fetching data: {e}")
+ return None
+
+ return len(self._cached_data)
+
+ def supports_partial_reading(self) -> bool:
+ """Partial reading is supported after initial fetch."""
+ return True
+
+ def refresh(self) -> bool:
+ """Refresh by clearing cached data."""
+ self._cached_data = None
+ return True
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get source status."""
+ status = super().get_status()
+ status["bucket"] = self._bucket
+ status["key"] = self._key
+ status["region"] = self._region
+ status["endpoint_url"] = self._endpoint_url
+ status["cached"] = self._cached_data is not None
+ return status
+
+ def close(self) -> None:
+ """Close the source."""
+ self._client = None
+ self._cached_data = None
diff --git a/potato/data_sources/sources/url_source.py b/potato/data_sources/sources/url_source.py
new file mode 100644
index 0000000000000000000000000000000000000000..5b3f1f1d700f54d8dc849b1f1d1f30b445d4cda5
--- /dev/null
+++ b/potato/data_sources/sources/url_source.py
@@ -0,0 +1,379 @@
+"""
+URL data source.
+
+This module provides data loading from HTTP/HTTPS URLs with security
+protections against SSRF attacks.
+"""
+
+import ipaddress
+import json
+import logging
+import os
+import socket
+import tempfile
+from typing import Any, Dict, Iterator, List, Optional
+from urllib.parse import urlparse
+
+from potato.data_sources.base import DataSource, SourceConfig
+
+logger = logging.getLogger(__name__)
+
+# Default limits
+DEFAULT_MAX_SIZE_MB = 100
+DEFAULT_TIMEOUT_SECONDS = 30
+
+# Private IP ranges to block (SSRF protection)
+PRIVATE_IP_RANGES = [
+ ipaddress.ip_network("10.0.0.0/8"),
+ ipaddress.ip_network("172.16.0.0/12"),
+ ipaddress.ip_network("192.168.0.0/16"),
+ ipaddress.ip_network("127.0.0.0/8"),
+ ipaddress.ip_network("169.254.0.0/16"),
+ ipaddress.ip_network("::1/128"),
+ ipaddress.ip_network("fc00::/7"),
+ ipaddress.ip_network("fe80::/10"),
+]
+
+
+def is_private_ip(ip_str: str) -> bool:
+ """Check if an IP address is in a private range."""
+ try:
+ ip = ipaddress.ip_address(ip_str)
+ for network in PRIVATE_IP_RANGES:
+ if ip in network:
+ return True
+ return False
+ except ValueError:
+ return False
+
+
+def resolve_and_validate_url(url: str, block_private_ips: bool = True) -> tuple:
+ """
+ Validate a URL and resolve it, checking for SSRF vulnerabilities.
+
+ Returns the validated URL and a list of validated (non-private) IP addresses
+ that can be used for IP-pinned connections, preventing DNS rebinding attacks.
+
+ Args:
+ url: The URL to validate
+ block_private_ips: Whether to block private/internal IPs
+
+ Returns:
+ Tuple of (validated_url, list_of_validated_ips)
+
+ Raises:
+ ValueError: If the URL is invalid or points to a blocked IP
+ """
+ parsed = urlparse(url)
+
+ # Only allow http/https
+ if parsed.scheme not in ('http', 'https'):
+ raise ValueError(f"Invalid URL scheme '{parsed.scheme}'. Only http/https allowed.")
+
+ if not parsed.netloc:
+ raise ValueError("Invalid URL: missing host")
+
+ # Extract hostname (without port)
+ hostname = parsed.hostname
+ if not hostname:
+ raise ValueError("Invalid URL: missing hostname")
+
+ validated_ips = []
+
+ if block_private_ips:
+ # Resolve hostname to IP
+ try:
+ # Get all IP addresses for the hostname
+ addr_info = socket.getaddrinfo(hostname, None)
+ for info in addr_info:
+ ip = info[4][0]
+ if is_private_ip(ip):
+ raise ValueError(
+ f"URL host '{hostname}' resolves to private IP {ip}. "
+ f"Access to private networks is not allowed."
+ )
+ validated_ips.append(ip)
+ except socket.gaierror as e:
+ raise ValueError(f"Could not resolve hostname '{hostname}': {e}")
+
+ return url, validated_ips
+
+
+class URLSource(DataSource):
+ """
+ Data source for HTTP/HTTPS URLs.
+
+ Supports fetching data from remote URLs with:
+ - SSRF protection (blocks private IPs)
+ - Custom headers for authentication
+ - Size limits and timeouts
+ - Content-type validation
+ - Caching integration
+
+ Configuration:
+ type: url
+ url: "https://example.com/data.jsonl" # Required
+ headers: # Optional custom headers
+ Authorization: "Bearer ${API_TOKEN}"
+ max_size_mb: 100 # Optional size limit
+ timeout_seconds: 30 # Optional request timeout
+ block_private_ips: true # Optional SSRF protection
+
+ Supported content types:
+ - application/json, application/x-ndjson
+ - text/csv, text/tab-separated-values
+ - application/x-jsonlines
+ """
+
+ def __init__(self, config: SourceConfig):
+ """Initialize the URL source."""
+ super().__init__(config)
+
+ self._url = config.config.get("url", "")
+ self._headers = config.config.get("headers", {})
+ self._max_size_bytes = config.config.get(
+ "max_size_mb", DEFAULT_MAX_SIZE_MB
+ ) * 1024 * 1024
+ self._timeout = config.config.get("timeout_seconds", DEFAULT_TIMEOUT_SECONDS)
+ self._block_private_ips = config.config.get("block_private_ips", True)
+ self._allowed_domains = config.config.get("allowed_domains")
+
+ # Cached data
+ self._cached_data: Optional[List[Dict]] = None
+ self._content_type: Optional[str] = None
+
+ def get_source_id(self) -> str:
+ """Get unique identifier."""
+ return self._source_id
+
+ def validate_config(self) -> List[str]:
+ """Validate source configuration."""
+ errors = []
+
+ if not self._url:
+ errors.append("'url' is required for URL source")
+ return errors
+
+ try:
+ parsed = urlparse(self._url)
+ if parsed.scheme not in ('http', 'https'):
+ errors.append(
+ f"Invalid URL scheme '{parsed.scheme}'. Only http/https allowed."
+ )
+
+ if not parsed.netloc:
+ errors.append("Invalid URL: missing host")
+
+ # Check domain allowlist if configured
+ if self._allowed_domains:
+ hostname = parsed.hostname
+ if hostname and hostname not in self._allowed_domains:
+ errors.append(
+ f"Domain '{hostname}' is not in allowed domains list"
+ )
+
+ except Exception as e:
+ errors.append(f"Invalid URL: {e}")
+
+ return errors
+
+ def is_available(self) -> bool:
+ """Check if the URL is accessible."""
+ try:
+ resolve_and_validate_url(self._url, self._block_private_ips)
+ return True
+ except ValueError as e:
+ logger.warning(f"URL not available: {e}")
+ return False
+ except Exception as e:
+ logger.warning(f"Error checking URL availability: {e}")
+ return False
+
+ def _fetch_data(self) -> List[Dict[str, Any]]:
+ """Fetch and parse data from the URL."""
+ import urllib.request
+ import urllib.error
+
+ # Resolve and validate URL; returns validated IPs for post-connection check.
+ # We validate immediately before the fetch to minimize TOCTOU window.
+ _, validated_ips = resolve_and_validate_url(
+ self._url, self._block_private_ips
+ )
+
+ # Build request with headers
+ request = urllib.request.Request(self._url)
+ for key, value in self._headers.items():
+ request.add_header(key, value)
+
+ # Add User-Agent if not specified
+ if 'User-Agent' not in self._headers:
+ request.add_header('User-Agent', 'Potato-Annotation-Tool/1.0')
+
+ try:
+ with urllib.request.urlopen(request, timeout=self._timeout) as response:
+ # Post-connection SSRF check: verify the connected IP is not private.
+ # This guards against DNS rebinding between validation and connect.
+ if self._block_private_ips:
+ try:
+ sock = None
+ # Navigate to the underlying socket
+ fp = getattr(response, 'fp', None)
+ raw = getattr(fp, 'raw', None) if fp else None
+ sock = getattr(raw, '_sock', None) if raw else None
+ if sock is not None:
+ peer = sock.getpeername()
+ if peer and is_private_ip(peer[0]):
+ raise ValueError(
+ f"Connection resolved to private IP {peer[0]}. "
+ f"Possible DNS rebinding attack."
+ )
+ except (AttributeError, OSError):
+ # If we can't inspect the socket, the pre-connect
+ # validation still provides the primary protection
+ pass
+
+ # Check content length
+ content_length = response.headers.get('Content-Length')
+ if content_length:
+ size = int(content_length)
+ if size > self._max_size_bytes:
+ raise ValueError(
+ f"Response size {size} exceeds limit {self._max_size_bytes}"
+ )
+
+ # Read with size limit
+ data = b""
+ chunk_size = 8192
+ while True:
+ chunk = response.read(chunk_size)
+ if not chunk:
+ break
+ data += chunk
+ if len(data) > self._max_size_bytes:
+ raise ValueError(
+ f"Response exceeded size limit of "
+ f"{self._max_size_bytes / (1024*1024):.1f}MB"
+ )
+
+ self._content_type = response.headers.get('Content-Type', '')
+
+ # Parse based on content type
+ return self._parse_content(data, self._content_type)
+
+ except urllib.error.HTTPError as e:
+ raise RuntimeError(f"HTTP error {e.code}: {e.reason}")
+ except urllib.error.URLError as e:
+ raise RuntimeError(f"URL error: {e.reason}")
+
+ def _parse_content(
+ self,
+ data: bytes,
+ content_type: str
+ ) -> List[Dict[str, Any]]:
+ """Parse content based on content type or URL extension."""
+ text = data.decode('utf-8')
+
+ # Determine format from content type or URL
+ url_path = urlparse(self._url).path.lower()
+
+ if any(ct in content_type for ct in ['json', 'ndjson', 'jsonlines']):
+ return self._parse_json(text)
+ elif url_path.endswith('.json') or url_path.endswith('.jsonl'):
+ return self._parse_json(text)
+ elif 'csv' in content_type or url_path.endswith('.csv'):
+ return self._parse_csv(text, ',')
+ elif 'tab-separated' in content_type or url_path.endswith('.tsv'):
+ return self._parse_csv(text, '\t')
+ else:
+ # Try JSON first, fall back to JSONL
+ try:
+ return self._parse_json(text)
+ except json.JSONDecodeError:
+ raise ValueError(
+ f"Could not parse content. "
+ f"Content-Type: {content_type}, URL: {self._url}"
+ )
+
+ def _parse_json(self, text: str) -> List[Dict[str, Any]]:
+ """Parse JSON or JSONL content."""
+ # Try as JSON array first
+ try:
+ data = json.loads(text)
+ if isinstance(data, list):
+ return data
+ elif isinstance(data, dict):
+ return [data]
+ else:
+ raise ValueError(f"Unexpected JSON type: {type(data)}")
+ except json.JSONDecodeError:
+ pass
+
+ # Parse as JSONL
+ items = []
+ for line_no, line in enumerate(text.split('\n'), 1):
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ item = json.loads(line)
+ if isinstance(item, list):
+ items.extend(item)
+ else:
+ items.append(item)
+ except json.JSONDecodeError as e:
+ logger.warning(f"Invalid JSON at line {line_no}: {e}")
+
+ return items
+
+ def _parse_csv(self, text: str, delimiter: str) -> List[Dict[str, Any]]:
+ """Parse CSV/TSV content."""
+ import csv
+ from io import StringIO
+
+ reader = csv.DictReader(StringIO(text), delimiter=delimiter)
+ return [dict(row) for row in reader]
+
+ def read_items(
+ self,
+ start: int = 0,
+ count: Optional[int] = None
+ ) -> Iterator[Dict[str, Any]]:
+ """Read items from the URL."""
+ # Fetch data if not cached
+ if self._cached_data is None:
+ self._cached_data = self._fetch_data()
+
+ # Apply start/count
+ items = self._cached_data[start:]
+ if count is not None:
+ items = items[:count]
+
+ yield from items
+
+ def get_total_count(self) -> Optional[int]:
+ """Get total number of items."""
+ if self._cached_data is None:
+ try:
+ self._cached_data = self._fetch_data()
+ except Exception as e:
+ logger.error(f"Error fetching data for count: {e}")
+ return None
+
+ return len(self._cached_data)
+
+ def supports_partial_reading(self) -> bool:
+ """URL source supports partial reading after fetch."""
+ return True
+
+ def refresh(self) -> bool:
+ """Refresh by clearing cached data."""
+ self._cached_data = None
+ return True
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get source status."""
+ status = super().get_status()
+ status["url"] = self._url
+ status["cached"] = self._cached_data is not None
+ status["content_type"] = self._content_type
+ return status
diff --git a/potato/database/__init__.py b/potato/database/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7b84aaef66a2fae0e0ef4ea9db07f1576fc7d6f2
--- /dev/null
+++ b/potato/database/__init__.py
@@ -0,0 +1,11 @@
+"""
+Database module for Potato annotation platform.
+
+This module provides database connectivity and management for user state persistence.
+It supports both MySQL and file-based storage backends.
+"""
+
+from .connection import DatabaseManager
+from .mysql_user_state import MysqlUserState
+
+__all__ = ['DatabaseManager', 'MysqlUserState']
\ No newline at end of file
diff --git a/potato/database/connection.py b/potato/database/connection.py
new file mode 100644
index 0000000000000000000000000000000000000000..c3d58beb6e20d7c90938259cc124aa2b2048b91c
--- /dev/null
+++ b/potato/database/connection.py
@@ -0,0 +1,254 @@
+"""
+Database connection management for Potato annotation platform.
+
+This module provides connection pooling and management for MySQL database operations.
+"""
+
+import mysql.connector
+from mysql.connector import pooling
+import logging
+from contextlib import contextmanager
+from typing import Optional, Dict, Any
+
+logger = logging.getLogger(__name__)
+
+
+class DatabaseManager:
+ """
+ Manages database connections and provides connection pooling for MySQL.
+
+ This class handles the creation and management of database connections,
+ including connection pooling for better performance and resource management.
+ """
+
+ def __init__(self, config: Dict[str, Any]):
+ """
+ Initialize the database manager with configuration.
+
+ Args:
+ config: Configuration dictionary containing database settings
+ """
+ self.config = config
+ self.pool = None
+ self._create_connection_pool()
+
+ def _create_connection_pool(self):
+ """Create the MySQL connection pool."""
+ db_config = self.config.get('database', {})
+
+ # Validate required database configuration
+ required_fields = ['host', 'database', 'username', 'password']
+ for field in required_fields:
+ if field not in db_config:
+ raise ValueError(f"Missing required database field: {field}")
+
+ pool_config = {
+ 'host': db_config.get('host', 'localhost'),
+ 'port': db_config.get('port', 3306),
+ 'database': db_config['database'],
+ 'user': db_config['username'],
+ 'password': db_config['password'],
+ 'charset': db_config.get('charset', 'utf8mb4'),
+ 'pool_name': 'potato_pool',
+ 'pool_size': db_config.get('pool_size', 10),
+ 'pool_reset_session': True,
+ 'autocommit': False, # We'll handle transactions explicitly
+ 'raise_on_warnings': True
+ }
+
+ try:
+ self.pool = pooling.MySQLConnectionPool(**pool_config)
+ logger.info(f"Created MySQL connection pool with {pool_config['pool_size']} connections")
+ except mysql.connector.Error as e:
+ logger.error(f"Failed to create database connection pool: {e}")
+ raise
+
+ @contextmanager
+ def get_connection(self):
+ """
+ Get a database connection from the pool.
+
+ Yields:
+ mysql.connector.connection.MySQLConnection: Database connection
+
+ Raises:
+ mysql.connector.Error: If connection cannot be established
+ """
+ connection = None
+ try:
+ connection = self.pool.get_connection()
+ yield connection
+ except mysql.connector.Error as e:
+ logger.error(f"Database connection error: {e}")
+ if connection:
+ connection.rollback()
+ raise
+ finally:
+ if connection:
+ try:
+ connection.close()
+ except mysql.connector.Error as e:
+ logger.warning(f"Error closing connection: {e}")
+
+ def test_connection(self) -> bool:
+ """
+ Test the database connection.
+
+ Returns:
+ bool: True if connection is successful, False otherwise
+ """
+ try:
+ with self.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("SELECT 1")
+ result = cursor.fetchone()
+ return result[0] == 1
+ except Exception as e:
+ logger.error(f"Database connection test failed: {e}")
+ return False
+
+ def create_tables(self):
+ """Create all required database tables if they don't exist."""
+ with self.get_connection() as conn:
+ cursor = conn.cursor()
+
+ # Create user_states table
+ cursor.execute("""
+ CREATE TABLE IF NOT EXISTS user_states (
+ id INT AUTO_INCREMENT PRIMARY KEY,
+ user_id VARCHAR(255) NOT NULL UNIQUE,
+ current_phase VARCHAR(50) NOT NULL,
+ current_page VARCHAR(255),
+ current_instance_index INT DEFAULT -1,
+ max_assignments INT DEFAULT -1,
+ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+ updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
+ INDEX idx_user_id (user_id)
+ ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+ """)
+
+ # Create user_instance_assignments table
+ cursor.execute("""
+ CREATE TABLE IF NOT EXISTS user_instance_assignments (
+ id INT AUTO_INCREMENT PRIMARY KEY,
+ user_id VARCHAR(255) NOT NULL,
+ instance_id VARCHAR(255) NOT NULL,
+ assignment_order INT NOT NULL,
+ assigned_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+ UNIQUE KEY unique_user_instance (user_id, instance_id),
+ INDEX idx_user_order (user_id, assignment_order),
+ FOREIGN KEY (user_id) REFERENCES user_states(user_id) ON DELETE CASCADE
+ ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+ """)
+
+ # Create label_annotations table
+ cursor.execute("""
+ CREATE TABLE IF NOT EXISTS label_annotations (
+ id INT AUTO_INCREMENT PRIMARY KEY,
+ user_id VARCHAR(255) NOT NULL,
+ instance_id VARCHAR(255) NOT NULL,
+ schema_name VARCHAR(255) NOT NULL,
+ label_name VARCHAR(255) NOT NULL,
+ label_value TEXT,
+ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+ updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
+ UNIQUE KEY unique_annotation (user_id, instance_id, schema_name, label_name),
+ INDEX idx_user_instance (user_id, instance_id),
+ FOREIGN KEY (user_id) REFERENCES user_states(user_id) ON DELETE CASCADE
+ ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+ """)
+
+ # Create span_annotations table
+ cursor.execute("""
+ CREATE TABLE IF NOT EXISTS span_annotations (
+ id INT AUTO_INCREMENT PRIMARY KEY,
+ user_id VARCHAR(255) NOT NULL,
+ instance_id VARCHAR(255) NOT NULL,
+ schema_name VARCHAR(255) NOT NULL,
+ span_name VARCHAR(255) NOT NULL,
+ span_title VARCHAR(255),
+ start_pos INT NOT NULL,
+ end_pos INT NOT NULL,
+ kb_id VARCHAR(255) DEFAULT NULL,
+ kb_source VARCHAR(255) DEFAULT NULL,
+ kb_label VARCHAR(512) DEFAULT NULL,
+ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+ INDEX idx_user_instance (user_id, instance_id),
+ FOREIGN KEY (user_id) REFERENCES user_states(user_id) ON DELETE CASCADE
+ ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+ """)
+
+ # Create phase_annotations table
+ cursor.execute("""
+ CREATE TABLE IF NOT EXISTS phase_annotations (
+ id INT AUTO_INCREMENT PRIMARY KEY,
+ user_id VARCHAR(255) NOT NULL,
+ phase_name VARCHAR(50) NOT NULL,
+ page_name VARCHAR(255) NOT NULL,
+ schema_name VARCHAR(255) NOT NULL,
+ label_name VARCHAR(255) NOT NULL,
+ label_value TEXT,
+ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+ updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
+ UNIQUE KEY unique_phase_annotation (user_id, phase_name, page_name, schema_name, label_name),
+ FOREIGN KEY (user_id) REFERENCES user_states(user_id) ON DELETE CASCADE
+ ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+ """)
+
+ # Create behavioral_data table
+ cursor.execute("""
+ CREATE TABLE IF NOT EXISTS behavioral_data (
+ id INT AUTO_INCREMENT PRIMARY KEY,
+ user_id VARCHAR(255) NOT NULL,
+ instance_id VARCHAR(255) NOT NULL,
+ data_key VARCHAR(255) NOT NULL,
+ data_value TEXT,
+ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+ INDEX idx_user_instance (user_id, instance_id),
+ FOREIGN KEY (user_id) REFERENCES user_states(user_id) ON DELETE CASCADE
+ ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+ """)
+
+ # Create ai_hints table
+ cursor.execute("""
+ CREATE TABLE IF NOT EXISTS ai_hints (
+ id INT AUTO_INCREMENT PRIMARY KEY,
+ user_id VARCHAR(255) NOT NULL,
+ instance_id VARCHAR(255) NOT NULL,
+ hint_text TEXT NOT NULL,
+ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+ UNIQUE KEY unique_hint (user_id, instance_id),
+ FOREIGN KEY (user_id) REFERENCES user_states(user_id) ON DELETE CASCADE
+ ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+ """)
+
+ conn.commit()
+ logger.info("Database tables created successfully")
+
+ def drop_tables(self):
+ """Drop all database tables (for testing)."""
+ with self.get_connection() as conn:
+ cursor = conn.cursor()
+
+ # Drop tables in reverse dependency order
+ tables = [
+ 'ai_hints',
+ 'behavioral_data',
+ 'phase_annotations',
+ 'span_annotations',
+ 'label_annotations',
+ 'user_instance_assignments',
+ 'user_states'
+ ]
+
+ for table in tables:
+ cursor.execute(f"DROP TABLE IF EXISTS {table}")
+
+ conn.commit()
+ logger.info("Database tables dropped successfully")
+
+ def close(self):
+ """Close the database connection pool."""
+ if self.pool:
+ self.pool.close()
+ logger.info("Database connection pool closed")
\ No newline at end of file
diff --git a/potato/database/mysql_user_state.py b/potato/database/mysql_user_state.py
new file mode 100644
index 0000000000000000000000000000000000000000..47cdb5115bd3131d0947d2e4c4e89e4f0af16ee6
--- /dev/null
+++ b/potato/database/mysql_user_state.py
@@ -0,0 +1,701 @@
+"""
+MySQL-backed UserState implementation for Potato annotation platform.
+
+This module provides a database-backed implementation of the UserState interface,
+storing all user state data in MySQL tables for persistence and scalability.
+"""
+
+import logging
+import threading
+from typing import Dict, List, Set, Any, Optional, Tuple
+from collections import defaultdict
+
+from potato.user_state_management import UserState
+from potato.phase import UserPhase
+from potato.item_state_management import Item, Label, SpanAnnotation
+from .connection import DatabaseManager
+
+logger = logging.getLogger(__name__)
+
+
+class MysqlUserState(UserState):
+ """
+ MySQL-backed implementation of UserState.
+
+ This class stores all user state data in MySQL tables, providing
+ persistence and scalability for annotation workflows.
+
+ Backend feature parity gap: link annotations (SpanLink) and event
+ annotations (EventAnnotation) are only supported by InMemoryUserState.
+ The MySQL schema has no link_annotations or event_annotations tables,
+ and this class does not implement add_link_annotation /
+ add_event_annotation / get_*_annotations. Configurations that use
+ span_link or event_annotation schemas will fail with AttributeError
+ on the MySQL backend.
+ """
+
+ def __init__(self, user_id: str, db_manager: DatabaseManager, max_assignments: int = -1):
+ """
+ Initialize the MySQL user state.
+
+ Args:
+ user_id: Unique identifier for the user
+ db_manager: Database manager instance
+ max_assignments: Maximum number of assignments for this user
+ """
+ self.user_id = user_id
+ self.db_manager = db_manager
+ self.max_assignments = max_assignments
+
+ # Thread-safe cache lock
+ self._cache_lock = threading.Lock()
+
+ # Ensure user exists in database
+ self._ensure_user_exists()
+
+ # Cache for performance (protected by _cache_lock)
+ self._instance_ordering_cache = None
+ self._current_phase_cache = None
+ self._current_page_cache = None
+ self._current_instance_index_cache = None
+
+ def _ensure_user_exists(self):
+ """Create user record if it doesn't exist."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ INSERT IGNORE INTO user_states
+ (user_id, current_phase, current_page, current_instance_index, max_assignments)
+ VALUES (%s, %s, %s, %s, %s)
+ """, (self.user_id, 'LOGIN', None, -1, self.max_assignments))
+ conn.commit()
+
+ def _invalidate_cache(self):
+ """Invalidate cached data (thread-safe)."""
+ with self._cache_lock:
+ self._instance_ordering_cache = None
+ self._current_phase_cache = None
+ self._current_page_cache = None
+ self._current_instance_index_cache = None
+
+ def advance_to_phase(self, phase: UserPhase, page: str) -> None:
+ """Advance the user to a new phase and page."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ UPDATE user_states
+ SET current_phase = %s, current_page = %s
+ WHERE user_id = %s
+ """, (str(phase), page, self.user_id))
+ conn.commit()
+
+ self._invalidate_cache()
+
+ def assign_instance(self, item: Item) -> None:
+ """Assign an instance to the user for annotation."""
+ instance_id = item.get_id()
+
+ # Check if already assigned
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT COUNT(*) FROM user_instance_assignments
+ WHERE user_id = %s AND instance_id = %s
+ """, (self.user_id, instance_id))
+
+ result = cursor.fetchone()
+ if result is not None and result[0] > 0:
+ return # Already assigned
+
+ # Get next assignment order
+ cursor.execute("""
+ SELECT COALESCE(MAX(assignment_order), -1) + 1
+ FROM user_instance_assignments
+ WHERE user_id = %s
+ """, (self.user_id,))
+ result = cursor.fetchone()
+ next_order = result[0] if result is not None else 0
+
+ # Insert assignment
+ cursor.execute("""
+ INSERT INTO user_instance_assignments (user_id, instance_id, assignment_order)
+ VALUES (%s, %s, %s)
+ """, (self.user_id, instance_id, next_order))
+
+ # Update current instance index if this is the first assignment
+ cursor.execute("""
+ SELECT current_instance_index FROM user_states WHERE user_id = %s
+ """, (self.user_id,))
+ result = cursor.fetchone()
+ current_index = result[0] if result is not None else -1
+
+ if current_index == -1:
+ cursor.execute("""
+ UPDATE user_states SET current_instance_index = 0 WHERE user_id = %s
+ """, (self.user_id,))
+
+ conn.commit()
+
+ self._invalidate_cache()
+
+ def assign_instance_at_index(self, item: Item, index: int) -> bool:
+ """Insert ``item`` at ``index`` in the user's assignment ordering.
+
+ Used by quality-control injection (attention checks, gold standards).
+ Returns False if the item is already assigned. Raises IndexError if
+ ``index`` is outside [0, current_assignment_count].
+ """
+ instance_id = item.get_id()
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT COUNT(*) FROM user_instance_assignments
+ WHERE user_id = %s AND instance_id = %s
+ """, (self.user_id, instance_id))
+ already = cursor.fetchone()
+ if already is not None and already[0] > 0:
+ return False
+
+ cursor.execute("""
+ SELECT COUNT(*) FROM user_instance_assignments WHERE user_id = %s
+ """, (self.user_id,))
+ count_result = cursor.fetchone()
+ current_count = count_result[0] if count_result is not None else 0
+ if index < 0 or index > current_count:
+ raise IndexError(
+ f"assign_instance_at_index: index {index} out of range "
+ f"[0, {current_count}]"
+ )
+
+ # Shift later orders up to make room for the insert.
+ cursor.execute("""
+ UPDATE user_instance_assignments
+ SET assignment_order = assignment_order + 1
+ WHERE user_id = %s AND assignment_order >= %s
+ """, (self.user_id, index))
+ cursor.execute("""
+ INSERT INTO user_instance_assignments (user_id, instance_id, assignment_order)
+ VALUES (%s, %s, %s)
+ """, (self.user_id, instance_id, index))
+
+ # Rebalance the user's cursor.
+ cursor.execute("""
+ SELECT current_instance_index FROM user_states WHERE user_id = %s
+ """, (self.user_id,))
+ current_index_result = cursor.fetchone()
+ current_index = current_index_result[0] if current_index_result is not None else -1
+ if current_index == -1:
+ new_index = 0
+ elif current_index >= index:
+ new_index = current_index + 1
+ else:
+ new_index = current_index
+ cursor.execute("""
+ UPDATE user_states SET current_instance_index = %s WHERE user_id = %s
+ """, (new_index, self.user_id))
+ conn.commit()
+
+ self._invalidate_cache()
+ return True
+
+ def unassign_instance(self, instance_id: str) -> bool:
+ """Remove an instance assignment from the user."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT assignment_order FROM user_instance_assignments
+ WHERE user_id = %s AND instance_id = %s
+ """, (self.user_id, instance_id))
+ result = cursor.fetchone()
+ if result is None:
+ return False
+
+ removed_order = result[0]
+ cursor.execute("""
+ DELETE FROM user_instance_assignments
+ WHERE user_id = %s AND instance_id = %s
+ """, (self.user_id, instance_id))
+ cursor.execute("""
+ UPDATE user_instance_assignments
+ SET assignment_order = assignment_order - 1
+ WHERE user_id = %s AND assignment_order > %s
+ """, (self.user_id, removed_order))
+
+ # get_current_instance_index opens its own connection, so under
+ # READ COMMITTED (MySQL default) it reads the pre-DELETE value,
+ # which is what the index-rebalance math below needs. The COUNT(*)
+ # that follows runs on this outer cursor and sees the DELETE.
+ current_index = self.get_current_instance_index()
+ cursor.execute("""
+ SELECT COUNT(*) FROM user_instance_assignments WHERE user_id = %s
+ """, (self.user_id,))
+ count_result = cursor.fetchone()
+ assignment_count = count_result[0] if count_result is not None else 0
+
+ if assignment_count == 0:
+ new_index = -1
+ elif current_index > removed_order:
+ new_index = current_index - 1
+ elif current_index == removed_order:
+ new_index = min(removed_order, assignment_count - 1)
+ else:
+ new_index = min(current_index, assignment_count - 1)
+
+ cursor.execute("""
+ UPDATE user_states SET current_instance_index = %s WHERE user_id = %s
+ """, (new_index, self.user_id))
+ conn.commit()
+
+ self._invalidate_cache()
+ return True
+
+ def get_current_instance(self) -> Optional[Item]:
+ """Get the current instance the user is annotating."""
+ current_index = self.get_current_instance_index()
+ if current_index < 0:
+ return None
+
+ instance_ordering = self._get_instance_ordering()
+ if current_index >= len(instance_ordering):
+ return None
+
+ instance_id = instance_ordering[current_index]
+ from potato.item_state_management import get_item_state_manager
+ return get_item_state_manager().get_item(instance_id)
+
+ def get_current_instance_index(self) -> int:
+ """Get the current instance index."""
+ if self._current_instance_index_cache is not None:
+ return self._current_instance_index_cache
+
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT current_instance_index FROM user_states WHERE user_id = %s
+ """, (self.user_id,))
+ result = cursor.fetchone()
+ self._current_instance_index_cache = result[0] if result else -1
+ return self._current_instance_index_cache
+
+ def get_user_id(self) -> str:
+ """Get the user ID."""
+ return self.user_id
+
+ def goto_prev_instance(self) -> bool:
+ """Move to the previous instance."""
+ current_index = self.get_current_instance_index()
+ if current_index > 0:
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ UPDATE user_states SET current_instance_index = %s WHERE user_id = %s
+ """, (current_index - 1, self.user_id))
+ conn.commit()
+
+ self._invalidate_cache()
+ return True
+ return False
+
+ def goto_next_instance(self) -> bool:
+ """Move to the next instance."""
+ current_index = self.get_current_instance_index()
+ instance_ordering = self._get_instance_ordering()
+
+ if current_index < len(instance_ordering) - 1:
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ UPDATE user_states SET current_instance_index = %s WHERE user_id = %s
+ """, (current_index + 1, self.user_id))
+ conn.commit()
+
+ self._invalidate_cache()
+ return True
+ return False
+
+ def go_to_index(self, instance_index: int) -> None:
+ """Move to a specific instance index."""
+ instance_ordering = self._get_instance_ordering()
+ if 0 <= instance_index < len(instance_ordering):
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ UPDATE user_states SET current_instance_index = %s WHERE user_id = %s
+ """, (instance_index, self.user_id))
+ conn.commit()
+
+ self._invalidate_cache()
+
+ def get_all_annotations(self) -> Dict[str, Dict[str, Any]]:
+ """Get all annotations for this user."""
+ annotations = {}
+
+ # Get label annotations
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT instance_id, schema_name, label_name, label_value
+ FROM label_annotations
+ WHERE user_id = %s
+ """, (self.user_id,))
+
+ for row in cursor.fetchall():
+ instance_id, schema_name, label_name, label_value = row
+ if instance_id not in annotations:
+ annotations[instance_id] = {"labels": {}, "spans": {}}
+
+ if schema_name not in annotations[instance_id]["labels"]:
+ annotations[instance_id]["labels"][schema_name] = {}
+
+ annotations[instance_id]["labels"][schema_name][label_name] = label_value
+
+ # Get span annotations
+ cursor.execute("""
+ SELECT instance_id, schema_name, span_name, span_title, start_pos, end_pos
+ FROM span_annotations
+ WHERE user_id = %s
+ """, (self.user_id,))
+
+ for row in cursor.fetchall():
+ instance_id, schema_name, span_name, span_title, start_pos, end_pos = row
+ if instance_id not in annotations:
+ annotations[instance_id] = {"labels": {}, "spans": {}}
+
+ if schema_name not in annotations[instance_id]["spans"]:
+ annotations[instance_id]["spans"][schema_name] = {}
+
+ annotations[instance_id]["spans"][schema_name][span_name] = {
+ "title": span_title,
+ "start": start_pos,
+ "end": end_pos
+ }
+
+ return annotations
+
+ def get_label_annotations(self, instance_id: str) -> Dict[Label, Any]:
+ """Get label annotations for a specific instance."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT schema_name, label_name, label_value
+ FROM label_annotations
+ WHERE user_id = %s AND instance_id = %s
+ """, (self.user_id, instance_id))
+
+ annotations = {}
+ for row in cursor.fetchall():
+ schema_name, label_name, label_value = row
+ label = Label(schema_name, label_name)
+ annotations[label] = label_value
+
+ return annotations
+
+ def get_span_annotations(self, instance_id: str) -> Dict[SpanAnnotation, Any]:
+ """Get span annotations for a specific instance."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT schema_name, span_name, span_title, start_pos, end_pos,
+ kb_id, kb_source, kb_label
+ FROM span_annotations
+ WHERE user_id = %s AND instance_id = %s
+ """, (self.user_id, instance_id))
+
+ annotations = {}
+ for row in cursor.fetchall():
+ schema_name, span_name, span_title, start_pos, end_pos, \
+ kb_id, kb_source, kb_label = row
+ span = SpanAnnotation(
+ schema_name, span_name, span_title, start_pos, end_pos,
+ kb_id=kb_id, kb_source=kb_source, kb_label=kb_label,
+ )
+ annotations[span] = True # Span annotations are boolean
+
+ return annotations
+
+ def get_current_phase_and_page(self) -> Tuple[UserPhase, Optional[str]]:
+ """Get the current phase and page."""
+ if self._current_phase_cache is not None and self._current_page_cache is not None:
+ return self._current_phase_cache, self._current_page_cache
+
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT current_phase, current_page FROM user_states WHERE user_id = %s
+ """, (self.user_id,))
+ result = cursor.fetchone()
+
+ if result:
+ phase_str, page = result
+ phase = UserPhase.fromstr(phase_str)
+ self._current_phase_cache = phase
+ self._current_page_cache = page
+ return phase, page
+ else:
+ return UserPhase.LOGIN, None
+
+ def get_annotation_count(self) -> int:
+ """Get the number of annotated instances."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT COUNT(DISTINCT instance_id) FROM label_annotations WHERE user_id = %s
+ UNION
+ SELECT COUNT(DISTINCT instance_id) FROM span_annotations WHERE user_id = %s
+ """, (self.user_id, self.user_id))
+
+ results = cursor.fetchall()
+ return sum(result[0] for result in results)
+
+ def get_assigned_instance_count(self) -> int:
+ """Get the number of assigned instances."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT COUNT(*) FROM user_instance_assignments WHERE user_id = %s
+ """, (self.user_id,))
+ result = cursor.fetchone()
+ return result[0] if result is not None else 0
+
+ def get_assigned_instance_ids(self) -> Set[str]:
+ """Get the set of assigned instance IDs."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT instance_id FROM user_instance_assignments
+ WHERE user_id = %s ORDER BY assignment_order
+ """, (self.user_id,))
+ return {row[0] for row in cursor.fetchall()}
+
+ def add_label_annotation(self, instance_id: str, label: Label, value: Any) -> None:
+ """Add a label annotation."""
+ phase, page = self.get_current_phase_and_page()
+
+ if phase == UserPhase.ANNOTATION:
+ # Store in label_annotations table
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ INSERT INTO label_annotations
+ (user_id, instance_id, schema_name, label_name, label_value)
+ VALUES (%s, %s, %s, %s, %s)
+ ON DUPLICATE KEY UPDATE label_value = VALUES(label_value)
+ """, (self.user_id, instance_id, label.get_schema(),
+ label.get_name(), str(value)))
+ conn.commit()
+ else:
+ # Store in phase_annotations table
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ INSERT INTO phase_annotations
+ (user_id, phase_name, page_name, schema_name, label_name, label_value)
+ VALUES (%s, %s, %s, %s, %s, %s)
+ ON DUPLICATE KEY UPDATE label_value = VALUES(label_value)
+ """, (self.user_id, str(phase), page, label.get_schema(),
+ label.get_name(), str(value)))
+ conn.commit()
+
+ def add_span_annotation(self, instance_id: str, span: SpanAnnotation, value: Any) -> None:
+ """Add a span annotation."""
+ phase, page = self.get_current_phase_and_page()
+
+ if phase == UserPhase.ANNOTATION:
+ # Store in span_annotations table
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ INSERT INTO span_annotations
+ (user_id, instance_id, schema_name, span_name, span_title, start_pos, end_pos,
+ kb_id, kb_source, kb_label)
+ VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+ ON DUPLICATE KEY UPDATE
+ span_title = VALUES(span_title),
+ start_pos = VALUES(start_pos),
+ end_pos = VALUES(end_pos),
+ kb_id = VALUES(kb_id),
+ kb_source = VALUES(kb_source),
+ kb_label = VALUES(kb_label)
+ """, (self.user_id, instance_id, span.get_schema(), span.get_name(),
+ span.get_title(), span.get_start(), span.get_end(),
+ getattr(span, 'kb_id', None), getattr(span, 'kb_source', None),
+ getattr(span, 'kb_label', None)))
+ conn.commit()
+ else:
+ # For non-annotation phases, store in phase_annotations as JSON
+ span_data = {
+ "title": span.get_title(),
+ "start": span.get_start(),
+ "end": span.get_end()
+ }
+ import json
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ INSERT INTO phase_annotations
+ (user_id, phase_name, page_name, schema_name, label_name, label_value)
+ VALUES (%s, %s, %s, %s, %s, %s)
+ ON DUPLICATE KEY UPDATE label_value = VALUES(label_value)
+ """, (self.user_id, str(phase), page, span.get_schema(),
+ span.get_name(), json.dumps(span_data)))
+ conn.commit()
+
+ def get_annotated_instance_ids(self) -> Set[str]:
+ """Get the set of annotated instance IDs."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT DISTINCT instance_id FROM label_annotations WHERE user_id = %s
+ UNION
+ SELECT DISTINCT instance_id FROM span_annotations WHERE user_id = %s
+ """, (self.user_id, self.user_id))
+ return {row[0] for row in cursor.fetchall()}
+
+ def has_annotated(self, instance_id: str) -> bool:
+ """Check if the user has annotated a specific instance."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT COUNT(*) FROM label_annotations WHERE user_id = %s AND instance_id = %s
+ UNION
+ SELECT COUNT(*) FROM span_annotations WHERE user_id = %s AND instance_id = %s
+ """, (self.user_id, instance_id, self.user_id, instance_id))
+
+ results = cursor.fetchall()
+ return any(result[0] > 0 for result in results)
+
+ def clear_all_annotations(self) -> None:
+ """Clear all annotations for this user."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("DELETE FROM label_annotations WHERE user_id = %s", (self.user_id,))
+ cursor.execute("DELETE FROM span_annotations WHERE user_id = %s", (self.user_id,))
+ cursor.execute("DELETE FROM phase_annotations WHERE user_id = %s", (self.user_id,))
+ cursor.execute("DELETE FROM behavioral_data WHERE user_id = %s", (self.user_id,))
+ cursor.execute("DELETE FROM ai_hints WHERE user_id = %s", (self.user_id,))
+ conn.commit()
+
+ def clear_instance_annotations(self, instance_id: str) -> None:
+ """Clear all annotations for one instance."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute(
+ "DELETE FROM label_annotations WHERE user_id = %s AND instance_id = %s",
+ (self.user_id, instance_id),
+ )
+ cursor.execute(
+ "DELETE FROM span_annotations WHERE user_id = %s AND instance_id = %s",
+ (self.user_id, instance_id),
+ )
+ cursor.execute(
+ "DELETE FROM behavioral_data WHERE user_id = %s AND instance_id = %s",
+ (self.user_id, instance_id),
+ )
+ cursor.execute(
+ "DELETE FROM ai_hints WHERE user_id = %s AND instance_id = %s",
+ (self.user_id, instance_id),
+ )
+ conn.commit()
+
+ def has_assignments(self) -> bool:
+ """Check if the user has any assignments."""
+ return self.get_assigned_instance_count() > 0
+
+ def has_remaining_assignments(self) -> bool:
+ """Check if the user has remaining assignments."""
+ from potato.item_state_management import get_item_state_manager
+
+ has_available_items = get_item_state_manager().has_unlabeled_items_for_user(self)
+
+ if self.max_assignments >= 0:
+ return self.get_annotation_count() < self.max_assignments and has_available_items
+
+ return has_available_items
+
+ def set_max_assignments(self, max_assignments: int) -> None:
+ """Set the maximum number of assignments."""
+ self.max_assignments = max_assignments
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ UPDATE user_states SET max_assignments = %s WHERE user_id = %s
+ """, (max_assignments, self.user_id))
+ conn.commit()
+
+ def get_max_assignments(self) -> int:
+ """Get the maximum number of assignments."""
+ return self.max_assignments
+
+ def hint_exists(self, instance_id: str) -> bool:
+ """Check if a hint exists for an instance."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT COUNT(*) FROM ai_hints WHERE user_id = %s AND instance_id = %s
+ """, (self.user_id, instance_id))
+ result = cursor.fetchone()
+ return result is not None and result[0] > 0
+
+ def get_hint(self, instance_id: str) -> Optional[str]:
+ """Get the hint for an instance."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT hint_text FROM ai_hints WHERE user_id = %s AND instance_id = %s
+ """, (self.user_id, instance_id))
+ result = cursor.fetchone()
+ return result[0] if result else None
+
+ def cache_hint(self, instance_id: str, hint: str) -> None:
+ """Cache a hint for an instance."""
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ INSERT INTO ai_hints (user_id, instance_id, hint_text)
+ VALUES (%s, %s, %s)
+ ON DUPLICATE KEY UPDATE hint_text = VALUES(hint_text)
+ """, (self.user_id, instance_id, hint))
+ conn.commit()
+
+ def _get_instance_ordering(self) -> List[str]:
+ """Get the ordered list of assigned instance IDs."""
+ if self._instance_ordering_cache is not None:
+ return self._instance_ordering_cache
+
+ with self.db_manager.get_connection() as conn:
+ cursor = conn.cursor()
+ cursor.execute("""
+ SELECT instance_id FROM user_instance_assignments
+ WHERE user_id = %s ORDER BY assignment_order
+ """, (self.user_id,))
+
+ self._instance_ordering_cache = [row[0] for row in cursor.fetchall()]
+ return self._instance_ordering_cache
+
+ def is_at_end_index(self) -> bool:
+ """Check if the user is at the end of their assignments."""
+ current_index = self.get_current_instance_index()
+ instance_ordering = self._get_instance_ordering()
+ return current_index == len(instance_ordering) - 1
+
+ def go_back(self) -> bool:
+ """Move back to the previous instance."""
+ return self.goto_prev_instance()
+
+ def go_forward(self) -> bool:
+ """Move forward to the next instance."""
+ return self.goto_next_instance()
+
+ def get_current_instance_id(self) -> Optional[str]:
+ """Get the ID of the current instance."""
+ current_instance = self.get_current_instance()
+ return current_instance.get_id() if current_instance else None
+
+ def get_labels(self) -> Dict[str, Dict[str, str]]:
+ """Get all labels (deprecated, use get_all_annotations)."""
+ annotations = self.get_all_annotations()
+ labels = {}
+ for instance_id, data in annotations.items():
+ labels[instance_id] = data.get("labels", {})
+ return labels
+
diff --git a/potato/datasets_integration.py b/potato/datasets_integration.py
new file mode 100644
index 0000000000000000000000000000000000000000..0524c42bd8f075927f9746d1e49cec9b710a02a7
--- /dev/null
+++ b/potato/datasets_integration.py
@@ -0,0 +1,117 @@
+"""
+HuggingFace Datasets Integration
+
+Convenience API for loading Potato annotations as HuggingFace Datasets
+or pandas DataFrames โ no Hub round-trip required.
+
+Requires: pip install datasets>=2.14.0
+
+Usage:
+ from potato import load_as_dataset, load_annotations
+
+ # Load as HuggingFace DatasetDict
+ ds = load_as_dataset("path/to/config.yaml")
+ print(ds["annotations"][0])
+
+ # Load as pandas DataFrame
+ df = load_annotations("path/to/config.yaml")
+ print(df.head())
+"""
+
+import logging
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+def load_as_dataset(config_path: str,
+ include_spans: bool = True,
+ include_items: bool = True):
+ """
+ Load Potato annotations as a HuggingFace DatasetDict.
+
+ Reads the config file, loads annotations from the output directory,
+ and returns an in-memory DatasetDict with up to three splits:
+ 'annotations', 'spans', and 'items'.
+
+ Args:
+ config_path: Path to the Potato YAML config file
+ include_spans: Include a 'spans' split (default True)
+ include_items: Include an 'items' split (default True)
+
+ Returns:
+ datasets.DatasetDict with annotation data
+
+ Raises:
+ ImportError: If the 'datasets' package is not installed
+ FileNotFoundError: If config_path does not exist
+ ValueError: If no annotations are found
+ """
+ try:
+ from datasets import DatasetDict # noqa: F401
+ except ImportError:
+ raise ImportError(
+ "The 'datasets' package is required for load_as_dataset(). "
+ "Install with: pip install datasets>=2.14.0"
+ )
+
+ from potato.export.cli import build_export_context
+ from potato.export.huggingface_exporter import HuggingFaceExporter
+
+ context = build_export_context(config_path)
+ exporter = HuggingFaceExporter()
+
+ return exporter.build_dataset_dict(
+ context,
+ include_spans=include_spans,
+ include_items=include_items,
+ )
+
+
+def load_annotations(config_path: str):
+ """
+ Load Potato annotations as a pandas DataFrame.
+
+ Reads the config file, loads annotations from the output directory,
+ and returns a flattened DataFrame with one row per (instance, user)
+ annotation pair.
+
+ Args:
+ config_path: Path to the Potato YAML config file
+
+ Returns:
+ pandas.DataFrame with columns: instance_id, user_id, and one
+ column per annotation schema
+
+ Raises:
+ FileNotFoundError: If config_path does not exist
+ ValueError: If no annotations are found
+ """
+ import json
+ import pandas as pd
+
+ from potato.export.cli import build_export_context
+
+ context = build_export_context(config_path)
+
+ if not context.annotations:
+ raise ValueError(
+ f"No annotations found for config: {config_path}"
+ )
+
+ schema_map = {s["name"]: s for s in context.schemas}
+ rows = []
+ for ann in context.annotations:
+ row = {
+ "instance_id": ann.get("instance_id", ""),
+ "user_id": ann.get("user_id", ""),
+ }
+ labels = ann.get("labels", {})
+ for schema_name, value in labels.items():
+ if isinstance(value, (dict, list)):
+ row[schema_name] = json.dumps(value, ensure_ascii=False)
+ else:
+ row[schema_name] = value
+ rows.append(row)
+
+ return pd.DataFrame(rows)
diff --git a/potato/directory_watcher.py b/potato/directory_watcher.py
new file mode 100644
index 0000000000000000000000000000000000000000..c37470758e1acdc59aa30dddc683c2686f9b94f4
--- /dev/null
+++ b/potato/directory_watcher.py
@@ -0,0 +1,581 @@
+"""
+Directory Watcher Module
+
+This module provides functionality for loading annotation instances from a directory
+and optionally watching that directory for new or modified files. When watching is
+enabled, a background thread periodically scans the directory and dynamically loads
+new instances or updates existing ones.
+
+The module supports the same file formats as the standard data_files configuration:
+JSON, JSONL, CSV, and TSV.
+
+Configuration:
+ data_directory: str - Path to the directory containing data files
+ watch_data_directory: bool - Whether to watch for changes (default: False)
+ watch_poll_interval: float - Seconds between directory scans (default: 5.0)
+ data_directory_encoding: str - File encoding for directory files (default: "utf-8")
+
+Example config:
+ data_directory: "./data/incoming"
+ watch_data_directory: true
+ watch_poll_interval: 10.0
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import threading
+import glob
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Set, Tuple, TYPE_CHECKING
+
+try:
+ import pandas as pd
+ HAS_PANDAS = True
+except ImportError:
+ HAS_PANDAS = False
+
+if TYPE_CHECKING:
+ from potato.item_state_management import ItemStateManager
+
+logger = logging.getLogger(__name__)
+
+# Singleton instance with thread-safe initialization
+DIRECTORY_WATCHER: Optional['DirectoryWatcher'] = None
+_DIRECTORY_WATCHER_LOCK = threading.Lock()
+
+
+@dataclass
+class FileState:
+ """
+ Tracks the state of a watched file.
+
+ Attributes:
+ file_path: Absolute path to the file
+ last_modified: Last modification time (os.path.getmtime)
+ file_size: File size in bytes
+ instance_ids: Set of instance IDs loaded from this file
+ last_error: Last error message if processing failed, None otherwise
+ last_processed: Timestamp of last successful processing
+ """
+ file_path: str
+ last_modified: float = 0.0
+ file_size: int = 0
+ instance_ids: Set[str] = field(default_factory=set)
+ last_error: Optional[str] = None
+ last_processed: Optional[float] = None
+
+
+class DirectoryWatcher:
+ """
+ Watches a directory for new or modified data files and loads them as annotation instances.
+
+ This class provides two modes of operation:
+ 1. Static loading: Load all files from a directory at startup (load_directory())
+ 2. Dynamic watching: Continuously monitor for changes (start_watching())
+
+ The watcher tracks which instances came from which file, enabling proper handling
+ of file modifications (updating existing instances rather than creating duplicates).
+
+ Thread Safety:
+ All public methods are thread-safe. The internal state is protected by
+ a reentrant lock (_lock) to allow safe concurrent access from the main
+ application thread and the background watching thread.
+
+ Attributes:
+ data_directory: Path to the directory to watch
+ poll_interval: Seconds between directory scans
+ id_key: Key in data items containing the unique instance ID
+ text_key: Key in data items containing the text to annotate
+ """
+
+ # Supported file extensions
+ SUPPORTED_EXTENSIONS = ('.json', '.jsonl', '.csv', '.tsv')
+
+ def __init__(self, config: dict, item_state_manager: 'ItemStateManager'):
+ """
+ Initialize the directory watcher.
+
+ Args:
+ config: Configuration dictionary containing:
+ - data_directory: Path to watch
+ - watch_poll_interval: Seconds between scans (default: 5.0)
+ - data_directory_encoding: File encoding (default: "utf-8")
+ - item_properties.id_key: Key for instance IDs
+ - item_properties.text_key: Key for text content
+ item_state_manager: The ItemStateManager instance to add items to
+
+ Raises:
+ ValueError: If data_directory is not configured or doesn't exist
+ """
+ self.data_directory = config.get("data_directory")
+ if not self.data_directory:
+ raise ValueError("data_directory must be configured")
+
+ # Resolve relative paths based on task_dir if available
+ if not os.path.isabs(self.data_directory):
+ task_dir = config.get("task_dir", "")
+ if task_dir:
+ self.data_directory = os.path.join(task_dir, self.data_directory)
+ self.data_directory = os.path.abspath(self.data_directory)
+
+ if not os.path.isdir(self.data_directory):
+ raise ValueError(f"data_directory does not exist or is not a directory: {self.data_directory}")
+
+ self.poll_interval = config.get("watch_poll_interval", 5.0)
+ self.encoding = config.get("data_directory_encoding", "utf-8")
+ self.id_key = config["item_properties"]["id_key"]
+ self.text_key = config["item_properties"]["text_key"]
+
+ self._item_state_manager = item_state_manager
+
+ # File tracking state
+ self._file_states: Dict[str, FileState] = {}
+ self._instance_to_file: Dict[str, str] = {} # instance_id -> file_path
+
+ # Threading
+ self._lock = threading.RLock()
+ self._stop_event = threading.Event()
+ self._watch_thread: Optional[threading.Thread] = None
+
+ logger.info(f"DirectoryWatcher initialized for: {self.data_directory}")
+
+ def load_directory(self) -> int:
+ """
+ Load all supported files from the data directory.
+
+ This method performs an initial scan of the directory and loads all
+ instances from supported file formats. It should be called once at
+ startup before start_watching().
+
+ Returns:
+ int: Total number of instances loaded
+
+ Side Effects:
+ - Populates ItemStateManager with loaded instances
+ - Updates internal file tracking state
+ """
+ total_added = 0
+
+ with self._lock:
+ files = self._scan_directory()
+ logger.info(f"Found {len(files)} supported files in {self.data_directory}")
+
+ for file_path in files:
+ try:
+ added, updated = self._process_file(file_path)
+ total_added += added
+ if added > 0 or updated > 0:
+ logger.info(f"Loaded {file_path}: {added} added, {updated} updated")
+ except Exception as e:
+ logger.error(f"Error loading {file_path}: {e}")
+
+ logger.info(f"Directory load complete: {total_added} total instances loaded")
+ return total_added
+
+ def start_watching(self) -> None:
+ """
+ Start the background directory watching thread.
+
+ The watching thread will periodically scan the directory for new or
+ modified files and process them. Use stop() to terminate the thread.
+
+ Note:
+ This method is idempotent - calling it multiple times has no effect
+ if the thread is already running.
+ """
+ with self._lock:
+ if self._watch_thread is not None and self._watch_thread.is_alive():
+ logger.warning("Directory watcher thread is already running")
+ return
+
+ self._stop_event.clear()
+ self._watch_thread = threading.Thread(
+ target=self._watch_loop,
+ name="DirectoryWatcher",
+ daemon=True
+ )
+ self._watch_thread.start()
+ logger.info(f"Directory watching started (poll interval: {self.poll_interval}s)")
+
+ def stop(self) -> None:
+ """
+ Stop the directory watching thread gracefully.
+
+ This method signals the watching thread to stop and waits for it
+ to terminate (up to 5 seconds). It's safe to call this method
+ even if watching was never started.
+ """
+ self._stop_event.set()
+
+ if self._watch_thread is not None and self._watch_thread.is_alive():
+ self._watch_thread.join(timeout=5.0)
+ if self._watch_thread.is_alive():
+ logger.warning("Directory watcher thread did not stop gracefully")
+ else:
+ logger.info("Directory watcher stopped")
+
+ self._watch_thread = None
+
+ def get_stats(self) -> dict:
+ """
+ Get statistics about the directory watcher state.
+
+ Returns:
+ dict: Statistics including:
+ - data_directory: Path being watched
+ - is_watching: Whether the watch thread is running
+ - poll_interval: Seconds between scans
+ - files_tracked: Number of files being tracked
+ - total_instances: Total instances loaded from this directory
+ - files: List of file states with details
+ """
+ with self._lock:
+ return {
+ "data_directory": self.data_directory,
+ "is_watching": self._watch_thread is not None and self._watch_thread.is_alive(),
+ "poll_interval": self.poll_interval,
+ "files_tracked": len(self._file_states),
+ "total_instances": len(self._instance_to_file),
+ "files": [
+ {
+ "path": fs.file_path,
+ "last_modified": fs.last_modified,
+ "instance_count": len(fs.instance_ids),
+ "last_error": fs.last_error
+ }
+ for fs in self._file_states.values()
+ ]
+ }
+
+ def force_rescan(self) -> Tuple[int, int]:
+ """
+ Force an immediate rescan of the directory.
+
+ This method can be called to trigger an immediate check for changes
+ without waiting for the next poll interval.
+
+ Returns:
+ Tuple[int, int]: (total_added, total_updated) counts
+ """
+ return self._scan_and_process()
+
+ def _watch_loop(self) -> None:
+ """
+ Main watching loop that runs in the background thread.
+
+ This loop periodically scans the directory for changes and processes
+ any new or modified files. It continues until stop() is called.
+ """
+ logger.debug("Directory watch loop started")
+
+ while not self._stop_event.is_set():
+ try:
+ added, updated = self._scan_and_process()
+ if added > 0 or updated > 0:
+ logger.info(f"Directory scan: {added} instances added, {updated} updated")
+ except Exception as e:
+ logger.error(f"Error in directory watch loop: {e}", exc_info=True)
+
+ # Wait for the poll interval or until stopped
+ self._stop_event.wait(timeout=self.poll_interval)
+
+ logger.debug("Directory watch loop ended")
+
+ def _scan_and_process(self) -> Tuple[int, int]:
+ """
+ Scan for changed files and process them.
+
+ Returns:
+ Tuple[int, int]: (total_added, total_updated) counts
+ """
+ total_added = 0
+ total_updated = 0
+
+ with self._lock:
+ current_files = set(self._scan_directory())
+ tracked_files = set(self._file_states.keys())
+
+ # Find new and potentially modified files
+ for file_path in current_files:
+ try:
+ stat = os.stat(file_path)
+ current_mtime = stat.st_mtime
+ current_size = stat.st_size
+ except OSError as e:
+ logger.warning(f"Cannot stat file {file_path}: {e}")
+ continue
+
+ # Check if file is new or modified
+ if file_path not in self._file_states:
+ # New file
+ added, updated = self._process_file(file_path)
+ total_added += added
+ total_updated += updated
+ else:
+ # Check if modified
+ fs = self._file_states[file_path]
+ if current_mtime > fs.last_modified or current_size != fs.file_size:
+ logger.debug(f"File modified: {file_path}")
+ added, updated = self._process_file(file_path)
+ total_added += added
+ total_updated += updated
+
+ # Note: We don't remove instances when files are deleted - this preserves
+ # annotations that may have been made on those instances.
+ removed_files = tracked_files - current_files
+ for file_path in removed_files:
+ logger.info(f"File removed (instances preserved): {file_path}")
+ # Keep the file state but mark that the file is gone
+ if file_path in self._file_states:
+ self._file_states[file_path].last_error = "File removed from directory"
+
+ return total_added, total_updated
+
+ def _scan_directory(self) -> List[str]:
+ """
+ Scan the data directory for supported files.
+
+ Returns:
+ List[str]: List of absolute paths to supported files
+ """
+ files = []
+ for ext in self.SUPPORTED_EXTENSIONS:
+ pattern = os.path.join(self.data_directory, f"*{ext}")
+ files.extend(glob.glob(pattern))
+ return sorted(files)
+
+ def _process_file(self, file_path: str) -> Tuple[int, int]:
+ """
+ Process a single data file, adding or updating instances.
+
+ Args:
+ file_path: Absolute path to the file to process
+
+ Returns:
+ Tuple[int, int]: (added_count, updated_count)
+
+ Side Effects:
+ - Updates ItemStateManager with new/updated instances
+ - Updates file tracking state
+ """
+ added_count = 0
+ updated_count = 0
+
+ try:
+ instances = self._parse_file(file_path)
+ stat = os.stat(file_path)
+
+ # Get or create file state
+ if file_path not in self._file_states:
+ self._file_states[file_path] = FileState(file_path=file_path)
+
+ fs = self._file_states[file_path]
+ new_instance_ids: Set[str] = set()
+
+ for instance_data in instances:
+ # Validate ID key exists
+ if self.id_key not in instance_data:
+ logger.warning(f"Missing id_key '{self.id_key}' in {file_path}, skipping instance")
+ continue
+
+ instance_id = str(instance_data[self.id_key])
+ new_instance_ids.add(instance_id)
+
+ # Check if text_key is missing (warning only)
+ if self.text_key not in instance_data:
+ logger.warning(f"Missing text_key '{self.text_key}' for instance {instance_id}")
+
+ # Add or update the instance
+ if self._item_state_manager.has_item(instance_id):
+ # Update existing instance
+ if self._item_state_manager.update_item(instance_id, instance_data):
+ updated_count += 1
+ logger.debug(f"Updated instance: {instance_id}")
+ else:
+ # Add new instance
+ try:
+ self._item_state_manager.add_item(instance_id, instance_data)
+ self._instance_to_file[instance_id] = file_path
+ added_count += 1
+ logger.debug(f"Added instance: {instance_id}")
+ except ValueError as e:
+ logger.error(f"Failed to add instance {instance_id}: {e}")
+
+ # Update file state
+ fs.last_modified = stat.st_mtime
+ fs.file_size = stat.st_size
+ fs.instance_ids = new_instance_ids
+ fs.last_error = None
+ fs.last_processed = stat.st_mtime
+
+ except Exception as e:
+ logger.error(f"Error processing file {file_path}: {e}")
+ if file_path in self._file_states:
+ self._file_states[file_path].last_error = str(e)
+ else:
+ self._file_states[file_path] = FileState(
+ file_path=file_path,
+ last_error=str(e)
+ )
+
+ return added_count, updated_count
+
+ def _parse_file(self, file_path: str) -> List[dict]:
+ """
+ Parse a data file and return a list of instance dictionaries.
+
+ Args:
+ file_path: Absolute path to the file
+
+ Returns:
+ List[dict]: List of instance data dictionaries
+
+ Raises:
+ ValueError: If file format is unsupported or parsing fails
+ """
+ ext = os.path.splitext(file_path)[1].lower()
+
+ if ext in ('.json', '.jsonl'):
+ return self._parse_json_file(file_path)
+ elif ext == '.csv':
+ return self._parse_csv_file(file_path, separator=',')
+ elif ext == '.tsv':
+ return self._parse_csv_file(file_path, separator='\t')
+ else:
+ raise ValueError(f"Unsupported file format: {ext}")
+
+ def _parse_json_file(self, file_path: str) -> List[dict]:
+ """
+ Parse a JSON or JSONL file.
+
+ Supports both:
+ - JSONL format: One JSON object per line
+ - JSON format: Single JSON array or object per line
+
+ Args:
+ file_path: Path to the JSON/JSONL file
+
+ Returns:
+ List[dict]: List of parsed instance dictionaries
+ """
+ instances = []
+
+ with open(file_path, 'rt', encoding=self.encoding) as f:
+ for line_no, line in enumerate(f, 1):
+ line = line.strip()
+ if not line:
+ continue
+
+ try:
+ item = json.loads(line)
+
+ # Handle both single objects and arrays
+ if isinstance(item, list):
+ instances.extend(item)
+ else:
+ instances.append(item)
+
+ except json.JSONDecodeError as e:
+ raise ValueError(
+ f"Invalid JSON at line {line_no} in {file_path}: {e}"
+ ) from e
+
+ return instances
+
+ def _parse_csv_file(self, file_path: str, separator: str) -> List[dict]:
+ """
+ Parse a CSV or TSV file.
+
+ Args:
+ file_path: Path to the CSV/TSV file
+ separator: Column separator (',' for CSV, '\t' for TSV)
+
+ Returns:
+ List[dict]: List of row dictionaries
+
+ Raises:
+ ImportError: If pandas is not available
+ ValueError: If required columns are missing
+ """
+ if not HAS_PANDAS:
+ raise ImportError(
+ "pandas is required for CSV/TSV file support. "
+ "Install it with: pip install pandas"
+ )
+
+ df = pd.read_csv(file_path, sep=separator, encoding=self.encoding)
+
+ # Validate ID column exists
+ if self.id_key not in df.columns:
+ raise ValueError(f"ID column '{self.id_key}' not found in {file_path}")
+
+ # Convert ID column to string
+ df[self.id_key] = df[self.id_key].astype(str)
+
+ # Convert text column to string if present
+ if self.text_key in df.columns:
+ df[self.text_key] = df[self.text_key].astype(str)
+
+ return df.to_dict('records')
+
+
+def init_directory_watcher(config: dict) -> Optional[DirectoryWatcher]:
+ """
+ Initialize the global DirectoryWatcher singleton if data_directory is configured.
+
+ This function creates a DirectoryWatcher instance if the configuration includes
+ a data_directory setting. The watcher is initialized but not started - call
+ load_directory() and optionally start_watching() after initialization.
+
+ Args:
+ config: Configuration dictionary
+
+ Returns:
+ DirectoryWatcher: The initialized watcher, or None if not configured
+
+ Note:
+ Thread-safe initialization using double-checked locking pattern.
+ """
+ global DIRECTORY_WATCHER
+
+ # Check if data_directory is configured
+ if "data_directory" not in config:
+ return None
+
+ # Double-checked locking for thread safety
+ if DIRECTORY_WATCHER is None:
+ with _DIRECTORY_WATCHER_LOCK:
+ if DIRECTORY_WATCHER is None:
+ from potato.item_state_management import get_item_state_manager
+ ism = get_item_state_manager()
+ DIRECTORY_WATCHER = DirectoryWatcher(config, ism)
+
+ return DIRECTORY_WATCHER
+
+
+def get_directory_watcher() -> Optional[DirectoryWatcher]:
+ """
+ Get the global DirectoryWatcher singleton instance.
+
+ Returns:
+ DirectoryWatcher: The singleton instance, or None if not initialized
+ """
+ return DIRECTORY_WATCHER
+
+
+def clear_directory_watcher() -> None:
+ """
+ Clear the global DirectoryWatcher singleton (for testing).
+
+ This function stops any running watch thread and clears the global instance.
+ Thread-safe.
+ """
+ global DIRECTORY_WATCHER
+
+ with _DIRECTORY_WATCHER_LOCK:
+ if DIRECTORY_WATCHER is not None:
+ DIRECTORY_WATCHER.stop()
+ DIRECTORY_WATCHER = None
diff --git a/potato/diversity_manager.py b/potato/diversity_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..fc291b1faaf63b994475821fcdd73763b450888f
--- /dev/null
+++ b/potato/diversity_manager.py
@@ -0,0 +1,805 @@
+"""
+Diversity Manager Module
+
+Provides embedding-based clustering and round-robin sampling to maximize diversity
+in annotation item ordering. Uses sentence-transformers for embeddings and k-means
+clustering, then samples items from different clusters to ensure annotators see
+diverse content rather than similar items in sequence.
+
+Key Components:
+- DiversityConfig: Configuration dataclass for diversity ordering
+- ClusterState: Per-user cluster tracking for round-robin sampling
+- DiversityManager: Main class for embeddings, clustering, and diverse ordering
+- Singleton management: init/get/clear pattern matching other managers
+"""
+
+import json
+import logging
+import os
+import queue
+import threading
+import time
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple
+
+logger = logging.getLogger(__name__)
+
+# Availability is probed WITHOUT importing the heavy stack: importing
+# sentence_transformers eagerly pulls in transformers + torch (~4s and several
+# hundred MB of RSS), and this module is imported at server boot (flask_server.py)
+# even for tasks that never use diversity ordering. We detect the packages via
+# importlib.util.find_spec and defer the real imports to first use (model load /
+# clustering). numpy stays eager โ it is light and used throughout this module.
+import importlib.util
+
+try:
+ import numpy as np
+except ImportError: # numpy is a core dependency; absence disables this module
+ np = None
+
+_SENTENCE_TRANSFORMERS_AVAILABLE = (
+ np is not None
+ and importlib.util.find_spec("sentence_transformers") is not None
+ and importlib.util.find_spec("sklearn") is not None
+)
+
+# Singleton
+_DIVERSITY_MANAGER: Optional['DiversityManager'] = None
+_DIVERSITY_LOCK = threading.Lock()
+
+
+@dataclass
+class DiversityConfig:
+ """Configuration for diversity-based ordering."""
+ enabled: bool = False
+ model_name: str = "all-MiniLM-L6-v2"
+ num_clusters: int = 10
+ items_per_cluster: int = 20
+ auto_clusters: bool = True
+ prefill_count: int = 100
+ batch_size: int = 32
+ cache_dir: Optional[str] = None
+ custom_embedding_function: Optional[Callable[[str], Any]] = None
+ recluster_threshold: float = 1.0
+ preserve_visited: bool = True
+ trigger_ai_prefetch: bool = True
+
+
+@dataclass
+class ClusterState:
+ """Per-user cluster tracking for round-robin sampling."""
+ sampled_clusters: Set[int] = field(default_factory=set)
+ cluster_sample_counts: Dict[int, int] = field(default_factory=dict)
+ current_cluster_index: int = 0
+ visited_instance_ids: Set[str] = field(default_factory=set)
+ skipped_instance_ids: Set[str] = field(default_factory=set)
+ last_recluster_time: Optional[datetime] = None
+
+
+class DiversityManager:
+ """
+ Manages embedding-based clustering for diversity-aware item ordering.
+
+ This class provides:
+ - Sentence-transformer embeddings with configurable model
+ - K-means clustering with auto-calculated cluster count
+ - Round-robin cluster sampling for diverse ordering
+ - Async embedding of new items after annotation
+ - Re-clustering when user has sampled all clusters
+ - Order preservation for annotated, visited, and skipped items
+ - AI cache prefetch integration after reordering
+ - Embedding persistence across server restarts
+ """
+
+ def __init__(self, config: DiversityConfig, app_config: Dict[str, Any]):
+ """
+ Initialize the diversity manager.
+
+ Args:
+ config: DiversityConfig instance with diversity settings
+ app_config: Full application configuration dictionary
+ """
+ self.config = config
+ self.app_config = app_config
+ self.logger = logging.getLogger(__name__)
+ self._lock = threading.RLock()
+
+ # Core state
+ self.enabled = False
+ self.model = None
+ self.embeddings: Dict[str, Any] = {} # instance_id -> numpy array
+ self.cluster_labels: Dict[str, int] = {} # instance_id -> cluster_id
+ self.cluster_members: Dict[int, List[str]] = {} # cluster_id -> [instance_ids]
+ self.user_cluster_states: Dict[str, ClusterState] = {} # user_id -> state
+ self.num_clusters: int = config.num_clusters
+
+ # Threading for async operations
+ self._embedding_executor = ThreadPoolExecutor(max_workers=4)
+ self._pending_futures: Dict[str, Future] = {}
+
+ # Check dependencies
+ if not _SENTENCE_TRANSFORMERS_AVAILABLE:
+ self.logger.warning(
+ "sentence-transformers or scikit-learn not installed. "
+ "Diversity ordering disabled. Install with: "
+ "pip install sentence-transformers scikit-learn"
+ )
+ return
+
+ if not config.enabled:
+ self.logger.info("Diversity ordering disabled in config")
+ return
+
+ # Initialize model
+ try:
+ if config.custom_embedding_function:
+ self.logger.info("Using custom embedding function")
+ self._embed_function = config.custom_embedding_function
+ else:
+ self.logger.info(f"Loading sentence-transformer model: {config.model_name}")
+ from sentence_transformers import SentenceTransformer # lazy: heavy import
+ self.model = SentenceTransformer(config.model_name)
+ self._embed_function = self._embed_with_model
+
+ self.enabled = True
+ self._load_cache()
+ self.logger.info(
+ f"Diversity manager ready: model={config.model_name}, "
+ f"cached_embeddings={len(self.embeddings)}"
+ )
+ except Exception as e:
+ self.logger.error(f"Failed to initialize diversity manager: {e}")
+ self.enabled = False
+
+ def _embed_with_model(self, texts: List[str]) -> Any:
+ """Embed texts using the loaded sentence-transformer model."""
+ return self.model.encode(texts, show_progress_bar=False)
+
+ def _get_cache_dir(self) -> str:
+ """Get the cache directory path."""
+ if self.config.cache_dir:
+ cache_dir = self.config.cache_dir
+ else:
+ output_dir = self.app_config.get("output_annotation_dir", "annotation_output")
+ cache_dir = os.path.join(output_dir, ".diversity_cache")
+ os.makedirs(cache_dir, exist_ok=True)
+ return cache_dir
+
+ def _save_cache(self) -> None:
+ """Save embeddings and cluster labels to disk."""
+ if not self.enabled:
+ return
+
+ try:
+ import numpy as np
+
+ cache_dir = self._get_cache_dir()
+
+ # Save embeddings as numpy .npz (safe, unlike pickle)
+ emb_path = os.path.join(cache_dir, "embeddings.npz")
+ if self.embeddings:
+ ids = list(self.embeddings.keys())
+ vectors = np.array([self.embeddings[iid] for iid in ids])
+ np.savez(emb_path, ids=np.array(ids), vectors=vectors)
+ else:
+ # Save empty arrays
+ np.savez(emb_path, ids=np.array([]), vectors=np.array([]))
+
+ # Also remove legacy pickle file if it exists
+ legacy_pkl = os.path.join(cache_dir, "embeddings.pkl")
+ if os.path.exists(legacy_pkl):
+ os.remove(legacy_pkl)
+
+ # Save cluster labels as JSON
+ labels_path = os.path.join(cache_dir, "cluster_labels.json")
+ with open(labels_path, "w") as f:
+ json.dump(self.cluster_labels, f)
+
+ self.logger.debug(f"Saved diversity cache: {len(self.embeddings)} embeddings")
+ except Exception as e:
+ self.logger.error(f"Failed to save diversity cache: {e}")
+
+ def _load_cache(self) -> None:
+ """Load cached embeddings and cluster labels from disk."""
+ try:
+ import numpy as np
+
+ cache_dir = self._get_cache_dir()
+
+ # Load embeddings from numpy .npz (safe format)
+ emb_path = os.path.join(cache_dir, "embeddings.npz")
+ if os.path.exists(emb_path):
+ data = np.load(emb_path, allow_pickle=False)
+ ids = data["ids"]
+ vectors = data["vectors"]
+ if len(ids) > 0 and len(vectors) > 0:
+ self.embeddings = {
+ str(iid): vec for iid, vec in zip(ids, vectors)
+ }
+ else:
+ self.embeddings = {}
+ elif os.path.exists(os.path.join(cache_dir, "embeddings.pkl")):
+ # Legacy pickle file โ refuse to load (security risk)
+ self.logger.warning(
+ "Found legacy embeddings.pkl cache file. "
+ "Refusing to load pickle files due to security risks. "
+ "Embeddings will be recomputed and saved in safe .npz format."
+ )
+
+ labels_path = os.path.join(cache_dir, "cluster_labels.json")
+ if os.path.exists(labels_path):
+ with open(labels_path, "r") as f:
+ self.cluster_labels = json.load(f)
+ # Rebuild cluster_members from labels
+ self._rebuild_cluster_members()
+
+ if self.embeddings:
+ self.logger.info(f"Loaded {len(self.embeddings)} cached embeddings")
+ except Exception as e:
+ self.logger.warning(f"Failed to load diversity cache: {e}")
+ self.embeddings = {}
+ self.cluster_labels = {}
+
+ def _rebuild_cluster_members(self) -> None:
+ """Rebuild cluster_members dict from cluster_labels."""
+ self.cluster_members = {}
+ for instance_id, cluster_id in self.cluster_labels.items():
+ if cluster_id not in self.cluster_members:
+ self.cluster_members[cluster_id] = []
+ self.cluster_members[cluster_id].append(instance_id)
+
+ def compute_embedding(self, text: str) -> Optional[Any]:
+ """
+ Compute embedding for a single text.
+
+ Args:
+ text: Text to embed
+
+ Returns:
+ Numpy array embedding, or None if failed
+ """
+ if not self.enabled:
+ return None
+
+ try:
+ embeddings = self._embed_function([text])
+ return embeddings[0]
+ except Exception as e:
+ self.logger.error(f"Error computing embedding: {e}")
+ return None
+
+ def compute_embeddings_batch(
+ self,
+ texts: Dict[str, str],
+ callback: Optional[Callable[[str, Any], None]] = None
+ ) -> int:
+ """
+ Batch-encode texts and store embeddings.
+
+ Args:
+ texts: Mapping of instance_id to text content
+ callback: Optional callback(instance_id, embedding) after each batch
+
+ Returns:
+ Number of new embeddings computed
+ """
+ if not self.enabled:
+ return 0
+
+ with self._lock:
+ # Filter out already cached items
+ new_items = {
+ iid: text for iid, text in texts.items()
+ if iid not in self.embeddings
+ }
+
+ if not new_items:
+ return 0
+
+ try:
+ ids = list(new_items.keys())
+ text_list = list(new_items.values())
+
+ # Process in batches
+ total_computed = 0
+ batch_size = self.config.batch_size
+
+ for i in range(0, len(text_list), batch_size):
+ batch_ids = ids[i:i + batch_size]
+ batch_texts = text_list[i:i + batch_size]
+
+ vecs = self._embed_function(batch_texts)
+
+ for j, iid in enumerate(batch_ids):
+ self.embeddings[iid] = vecs[j]
+ if callback:
+ callback(iid, vecs[j])
+
+ total_computed += len(batch_ids)
+
+ self._save_cache()
+ self.logger.info(f"Computed {total_computed} new embeddings")
+ return total_computed
+
+ except Exception as e:
+ self.logger.error(f"Error computing embeddings batch: {e}")
+ return 0
+
+ def start_async_embedding(self, instance_id: str, text: str) -> Optional[Future]:
+ """
+ Start async embedding computation for a single item.
+
+ Args:
+ instance_id: The instance ID
+ text: Text content to embed
+
+ Returns:
+ Future for the embedding computation
+ """
+ if not self.enabled:
+ return None
+
+ with self._lock:
+ if instance_id in self.embeddings:
+ return None # Already computed
+
+ if instance_id in self._pending_futures:
+ return self._pending_futures[instance_id]
+
+ def compute():
+ emb = self.compute_embedding(text)
+ if emb is not None:
+ with self._lock:
+ self.embeddings[instance_id] = emb
+ self._save_cache()
+ return emb
+
+ future = self._embedding_executor.submit(compute)
+ self._pending_futures[instance_id] = future
+
+ def cleanup(f):
+ with self._lock:
+ self._pending_futures.pop(instance_id, None)
+
+ future.add_done_callback(cleanup)
+ return future
+
+ def cluster_items(self, force: bool = False) -> bool:
+ """
+ Cluster items using k-means on embeddings.
+
+ Args:
+ force: Force re-clustering even if already clustered
+
+ Returns:
+ True if clustering was performed
+ """
+ if not self.enabled or not _SENTENCE_TRANSFORMERS_AVAILABLE:
+ return False
+
+ with self._lock:
+ if not self.embeddings:
+ self.logger.warning("No embeddings available for clustering")
+ return False
+
+ if self.cluster_labels and not force:
+ self.logger.debug("Items already clustered, skipping")
+ return False
+
+ try:
+ ids = list(self.embeddings.keys())
+ vectors = np.array([self.embeddings[iid] for iid in ids])
+
+ # Calculate number of clusters
+ if self.config.auto_clusters:
+ n_items = len(ids)
+ target_size = self.config.items_per_cluster
+ self.num_clusters = max(2, min(n_items // target_size, n_items // 2))
+ else:
+ self.num_clusters = min(self.config.num_clusters, len(ids))
+
+ self.logger.info(f"Clustering {len(ids)} items into {self.num_clusters} clusters")
+
+ from sklearn.cluster import KMeans # lazy: heavy import
+ kmeans = KMeans(n_clusters=self.num_clusters, random_state=42, n_init=10)
+ labels = kmeans.fit_predict(vectors)
+
+ # Store results
+ self.cluster_labels = {}
+ self.cluster_members = {}
+
+ for i, iid in enumerate(ids):
+ cluster_id = int(labels[i])
+ self.cluster_labels[iid] = cluster_id
+ if cluster_id not in self.cluster_members:
+ self.cluster_members[cluster_id] = []
+ self.cluster_members[cluster_id].append(iid)
+
+ self._save_cache()
+
+ # Log cluster sizes
+ sizes = [len(m) for m in self.cluster_members.values()]
+ self.logger.info(
+ f"Clustering complete: {self.num_clusters} clusters, "
+ f"sizes range {min(sizes)}-{max(sizes)}, avg {sum(sizes)/len(sizes):.1f}"
+ )
+ return True
+
+ except Exception as e:
+ self.logger.error(f"Clustering failed: {e}")
+ return False
+
+ def get_user_cluster_state(self, user_id: str) -> ClusterState:
+ """Get or create cluster state for a user."""
+ with self._lock:
+ if user_id not in self.user_cluster_states:
+ self.user_cluster_states[user_id] = ClusterState()
+ return self.user_cluster_states[user_id]
+
+ def _get_next_cluster(self, user_id: str, available_clusters: Set[int]) -> Optional[int]:
+ """
+ Get the next cluster for round-robin sampling.
+
+ Args:
+ user_id: User identifier
+ available_clusters: Set of clusters with available items
+
+ Returns:
+ Next cluster ID to sample from, or None if no clusters available
+ """
+ if not available_clusters:
+ return None
+
+ state = self.get_user_cluster_state(user_id)
+
+ # Sort clusters for deterministic ordering
+ sorted_clusters = sorted(available_clusters)
+
+ # Find the next cluster that hasn't been fully sampled
+ for _ in range(len(sorted_clusters)):
+ # Cycle through clusters
+ idx = state.current_cluster_index % len(sorted_clusters)
+ cluster_id = sorted_clusters[idx]
+
+ state.current_cluster_index = (state.current_cluster_index + 1) % len(sorted_clusters)
+
+ if cluster_id in available_clusters:
+ state.sampled_clusters.add(cluster_id)
+ state.cluster_sample_counts[cluster_id] = state.cluster_sample_counts.get(cluster_id, 0) + 1
+ return cluster_id
+
+ return None
+
+ def get_next_diverse_item(
+ self,
+ user_id: str,
+ available_ids: Set[str]
+ ) -> Optional[str]:
+ """
+ Get the next item for a user using round-robin cluster sampling.
+
+ Args:
+ user_id: User identifier
+ available_ids: Set of available instance IDs
+
+ Returns:
+ Next instance ID to assign, or None if no items available
+ """
+ if not self.enabled or not self.cluster_labels:
+ return None
+
+ with self._lock:
+ # Find clusters with available items
+ available_by_cluster: Dict[int, List[str]] = {}
+ for iid in available_ids:
+ if iid in self.cluster_labels:
+ cluster_id = self.cluster_labels[iid]
+ if cluster_id not in available_by_cluster:
+ available_by_cluster[cluster_id] = []
+ available_by_cluster[cluster_id].append(iid)
+
+ if not available_by_cluster:
+ # No clustered items available
+ return list(available_ids)[0] if available_ids else None
+
+ # Get next cluster using round-robin
+ next_cluster = self._get_next_cluster(user_id, set(available_by_cluster.keys()))
+
+ if next_cluster is None:
+ return None
+
+ # Return first available item from the cluster
+ items = available_by_cluster.get(next_cluster, [])
+ return items[0] if items else None
+
+ def generate_diverse_ordering(
+ self,
+ user_id: str,
+ available_ids: List[str],
+ preserve_ids: Set[str]
+ ) -> List[str]:
+ """
+ Generate a diverse ordering of items with order preservation.
+
+ Items in preserve_ids will maintain their original positions.
+ Remaining items are reordered using round-robin cluster sampling.
+
+ Args:
+ user_id: User identifier
+ available_ids: List of available instance IDs in current order
+ preserve_ids: Set of instance IDs that should keep their positions
+
+ Returns:
+ New ordering of instance IDs
+ """
+ if not self.enabled or not self.cluster_labels:
+ return available_ids
+
+ with self._lock:
+ state = self.get_user_cluster_state(user_id)
+
+ # Combine all items to preserve
+ all_preserve = preserve_ids | state.visited_instance_ids
+ if self.config.preserve_visited:
+ all_preserve |= state.skipped_instance_ids
+
+ # Separate into preserved (keep position) and reorderable
+ preserved_positions: List[Tuple[int, str]] = []
+ reorderable: Set[str] = set()
+
+ for i, iid in enumerate(available_ids):
+ if iid in all_preserve:
+ preserved_positions.append((i, iid))
+ else:
+ reorderable.add(iid)
+
+ # Generate diverse order for reorderable items
+ diverse_order: List[str] = []
+ remaining = reorderable.copy()
+
+ while remaining:
+ next_item = self.get_next_diverse_item(user_id, remaining)
+ if next_item:
+ diverse_order.append(next_item)
+ remaining.discard(next_item)
+ else:
+ # Fallback: append remaining items
+ diverse_order.extend(sorted(remaining))
+ break
+
+ # Merge preserved items back at their original positions using
+ # a slot-based approach. Pre-allocate slots for preserved items,
+ # then fill remaining slots with diverse items in order.
+ total_len = len(diverse_order) + len(preserved_positions)
+ result = [None] * total_len
+
+ # Place preserved items in their original slots
+ for orig_idx, iid in preserved_positions:
+ slot = min(orig_idx, total_len - 1)
+ result[slot] = iid
+
+ # Fill remaining slots with diverse items in order
+ diverse_iter = iter(diverse_order)
+ for i in range(total_len):
+ if result[i] is None:
+ try:
+ result[i] = next(diverse_iter)
+ except StopIteration:
+ break
+
+ # Remove any remaining None slots (shouldn't happen normally)
+ result = [x for x in result if x is not None]
+
+ return result
+
+ def should_recluster(self, user_id: str) -> bool:
+ """
+ Check if reclustering should be triggered for a user.
+
+ Returns True when user has sampled from all clusters (based on threshold).
+
+ Args:
+ user_id: User identifier
+
+ Returns:
+ True if reclustering is needed
+ """
+ if not self.enabled or not self.cluster_members:
+ return False
+
+ with self._lock:
+ state = self.get_user_cluster_state(user_id)
+
+ total_clusters = len(self.cluster_members)
+ sampled_clusters = len(state.sampled_clusters)
+
+ if total_clusters == 0:
+ return False
+
+ coverage = sampled_clusters / total_clusters
+ return coverage >= self.config.recluster_threshold
+
+ def trigger_recluster(self, user_id: str) -> bool:
+ """
+ Trigger reclustering and reset user's cluster state.
+
+ Args:
+ user_id: User identifier
+
+ Returns:
+ True if reclustering was performed
+ """
+ with self._lock:
+ # Reset user's cluster sampling state
+ state = self.get_user_cluster_state(user_id)
+ state.sampled_clusters.clear()
+ state.cluster_sample_counts.clear()
+ state.current_cluster_index = 0
+ state.last_recluster_time = datetime.now()
+
+ # Force reclustering
+ result = self.cluster_items(force=True)
+
+ if result:
+ self.logger.info(f"Reclustered items for user {user_id}")
+
+ return result
+
+ def mark_item_visited(self, user_id: str, instance_id: str) -> None:
+ """Mark an item as visited by a user."""
+ with self._lock:
+ state = self.get_user_cluster_state(user_id)
+ state.visited_instance_ids.add(instance_id)
+
+ def mark_item_skipped(self, user_id: str, instance_id: str) -> None:
+ """Mark an item as skipped by a user."""
+ with self._lock:
+ state = self.get_user_cluster_state(user_id)
+ state.skipped_instance_ids.add(instance_id)
+
+ def on_annotation_complete(
+ self,
+ user_id: str,
+ instance_id: str,
+ text: str
+ ) -> None:
+ """
+ Handle annotation completion - compute embedding if needed.
+
+ Args:
+ user_id: User identifier
+ instance_id: Instance that was annotated
+ text: Text content of the instance
+ """
+ with self._lock:
+ state = self.get_user_cluster_state(user_id)
+ state.visited_instance_ids.add(instance_id)
+
+ # Start async embedding if not already computed
+ if instance_id not in self.embeddings:
+ self.start_async_embedding(instance_id, text)
+
+ def apply_to_user_ordering(
+ self,
+ user_id: str,
+ available_ids: List[str],
+ annotated_ids: Set[str]
+ ) -> List[str]:
+ """
+ Apply diversity ordering to a user's available items.
+
+ This is the main entry point for integrating with ItemStateManager.
+
+ Args:
+ user_id: User identifier
+ available_ids: Available instance IDs in current order
+ annotated_ids: Instance IDs the user has already annotated
+
+ Returns:
+ Diversely ordered list of instance IDs
+ """
+ if not self.enabled:
+ return available_ids
+
+ return self.generate_diverse_ordering(user_id, available_ids, annotated_ids)
+
+ def get_stats(self) -> Dict[str, Any]:
+ """Get diversity manager statistics."""
+ with self._lock:
+ cluster_sizes = {}
+ if self.cluster_members:
+ cluster_sizes = {
+ cid: len(members)
+ for cid, members in self.cluster_members.items()
+ }
+
+ return {
+ "available": _SENTENCE_TRANSFORMERS_AVAILABLE,
+ "enabled": self.enabled,
+ "model": self.config.model_name,
+ "embedding_count": len(self.embeddings),
+ "cluster_count": len(self.cluster_members),
+ "cluster_sizes": cluster_sizes,
+ "num_users": len(self.user_cluster_states),
+ "pending_embeddings": len(self._pending_futures),
+ }
+
+ def shutdown(self) -> None:
+ """Shutdown the diversity manager."""
+ self._embedding_executor.shutdown(wait=False)
+ self.logger.info("Diversity manager shutdown complete")
+
+
+def parse_diversity_config(config_data: Dict[str, Any]) -> DiversityConfig:
+ """
+ Parse diversity_ordering section from config into DiversityConfig.
+
+ Args:
+ config_data: Full application configuration
+
+ Returns:
+ DiversityConfig instance
+ """
+ dc = config_data.get("diversity_ordering", {})
+
+ return DiversityConfig(
+ enabled=dc.get("enabled", False),
+ model_name=dc.get("model_name", "all-MiniLM-L6-v2"),
+ num_clusters=dc.get("num_clusters", 10),
+ items_per_cluster=dc.get("items_per_cluster", 20),
+ auto_clusters=dc.get("auto_clusters", True),
+ prefill_count=dc.get("prefill_count", 100),
+ batch_size=dc.get("batch_size", 32),
+ cache_dir=dc.get("cache_dir"),
+ recluster_threshold=dc.get("recluster_threshold", 1.0),
+ preserve_visited=dc.get("preserve_visited", True),
+ trigger_ai_prefetch=dc.get("trigger_ai_prefetch", True),
+ )
+
+
+def init_diversity_manager(
+ config_data: Dict[str, Any]
+) -> Optional[DiversityManager]:
+ """
+ Initialize the singleton DiversityManager.
+
+ Args:
+ config_data: Full application configuration
+
+ Returns:
+ DiversityManager instance, or None if disabled
+ """
+ global _DIVERSITY_MANAGER
+
+ with _DIVERSITY_LOCK:
+ if _DIVERSITY_MANAGER is None:
+ diversity_config = parse_diversity_config(config_data)
+
+ # Check if diversity clustering is the assignment strategy
+ assignment_strategy = config_data.get("assignment_strategy", "")
+ if isinstance(assignment_strategy, dict):
+ assignment_strategy = assignment_strategy.get("name", "")
+
+ # Enable if strategy is diversity_clustering, even if diversity_ordering.enabled is false
+ if assignment_strategy == "diversity_clustering" and not diversity_config.enabled:
+ diversity_config.enabled = True
+
+ _DIVERSITY_MANAGER = DiversityManager(diversity_config, config_data)
+
+ return _DIVERSITY_MANAGER
+
+
+def get_diversity_manager() -> Optional[DiversityManager]:
+ """Get the singleton DiversityManager instance."""
+ return _DIVERSITY_MANAGER
+
+
+def clear_diversity_manager() -> None:
+ """Clear the singleton (for testing)."""
+ global _DIVERSITY_MANAGER
+ with _DIVERSITY_LOCK:
+ if _DIVERSITY_MANAGER is not None:
+ _DIVERSITY_MANAGER.shutdown()
+ _DIVERSITY_MANAGER = None
diff --git a/potato/embedding_visualization.py b/potato/embedding_visualization.py
new file mode 100644
index 0000000000000000000000000000000000000000..f2a18e8176bc740a773c4a0a84ec3083354920c3
--- /dev/null
+++ b/potato/embedding_visualization.py
@@ -0,0 +1,764 @@
+"""
+Embedding Visualization Module
+
+Provides 2D visualization of text/image embeddings for the admin dashboard,
+enabling interactive exploration and prioritization of annotation items.
+
+Key Components:
+- EmbeddingVisualizationManager: Main class for embedding visualization
+- UMAP dimensionality reduction for 2D projection
+- Label coloring via MACE or majority vote
+- Interactive selection and queue reordering
+
+The visualization allows admins to:
+- See clustering patterns in the data
+- Identify annotated vs unannotated items
+- Select regions to prioritize for annotation
+- Interleave multiple selections for diverse sampling
+"""
+
+import logging
+import threading
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Set, Tuple
+import hashlib
+import json
+
+logger = logging.getLogger(__name__)
+
+# Guarded imports for optional dependencies
+try:
+ import numpy as np
+ _NUMPY_AVAILABLE = True
+except ImportError:
+ _NUMPY_AVAILABLE = False
+ np = None
+
+try:
+ import umap
+ _UMAP_AVAILABLE = True
+except ImportError:
+ _UMAP_AVAILABLE = False
+ umap = None
+
+# Singleton
+_EMBEDDING_VIZ_MANAGER: Optional['EmbeddingVisualizationManager'] = None
+_EMBEDDING_VIZ_LOCK = threading.Lock()
+
+
+@dataclass
+class EmbeddingVizConfig:
+ """Configuration for embedding visualization."""
+ enabled: bool = True
+ sample_size: int = 1000
+ include_all_annotated: bool = True
+ embedding_model: str = "all-MiniLM-L6-v2"
+ image_embedding_model: str = "clip-ViT-B-32"
+ umap_n_neighbors: int = 15
+ umap_min_dist: float = 0.1
+ umap_metric: str = "cosine"
+ label_source: str = "mace" # "mace" or "majority"
+
+
+@dataclass
+class VisualizationPoint:
+ """A single point in the visualization."""
+ instance_id: str
+ x: float
+ y: float
+ label: Optional[str] = None
+ label_source: Optional[str] = None
+ preview: str = ""
+ preview_type: str = "text" # "text" or "image"
+ annotated: bool = False
+ annotation_count: int = 0
+
+
+@dataclass
+class VisualizationData:
+ """Complete visualization data for the scatter plot."""
+ points: List[VisualizationPoint] = field(default_factory=list)
+ labels: List[Optional[str]] = field(default_factory=list)
+ label_colors: Dict[Optional[str], str] = field(default_factory=dict)
+ stats: Dict[str, Any] = field(default_factory=dict)
+
+
+# Default color palette for labels
+DEFAULT_COLORS = [
+ "#22c55e", # green
+ "#ef4444", # red
+ "#3b82f6", # blue
+ "#eab308", # yellow
+ "#8b5cf6", # purple
+ "#f97316", # orange
+ "#06b6d4", # cyan
+ "#ec4899", # pink
+ "#14b8a6", # teal
+ "#f59e0b", # amber
+]
+
+UNANNOTATED_COLOR = "#94a3b8" # slate gray
+
+
+class EmbeddingVisualizationManager:
+ """
+ Manages embedding visualization for the admin dashboard.
+
+ This class provides:
+ - 2D UMAP projections of text/image embeddings
+ - Label coloring via MACE or majority vote
+ - Interactive selection and queue reordering
+ - Caching with invalidation on new annotations
+ """
+
+ def __init__(self, config: EmbeddingVizConfig, app_config: Dict[str, Any]):
+ """
+ Initialize the embedding visualization manager.
+
+ Args:
+ config: EmbeddingVizConfig instance
+ app_config: Full application configuration dictionary
+ """
+ self.config = config
+ self.app_config = app_config
+ self.logger = logging.getLogger(__name__)
+ self._lock = threading.RLock()
+
+ # State
+ self.enabled = False
+ self._projection_cache: Optional[Dict[str, Tuple[float, float]]] = None
+ self._cache_hash: Optional[str] = None
+ self._label_cache: Dict[str, Optional[str]] = {}
+
+ # Check dependencies
+ if not _NUMPY_AVAILABLE:
+ self.logger.warning(
+ "numpy not available. Embedding visualization disabled."
+ )
+ return
+
+ if not _UMAP_AVAILABLE:
+ self.logger.warning(
+ "umap-learn not installed. Embedding visualization disabled. "
+ "Install with: pip install umap-learn"
+ )
+ return
+
+ if not config.enabled:
+ self.logger.info("Embedding visualization disabled in config")
+ return
+
+ self.enabled = True
+ self.logger.info("Embedding visualization manager initialized")
+
+ def _get_diversity_manager(self):
+ """Get the DiversityManager singleton."""
+ from potato.diversity_manager import get_diversity_manager
+ return get_diversity_manager()
+
+ def _get_item_state_manager(self):
+ """Get the ItemStateManager singleton."""
+ from potato.item_state_management import get_item_state_manager
+ return get_item_state_manager()
+
+ def _get_user_state_manager(self):
+ """Get the UserStateManager singleton."""
+ from potato.user_state_management import get_user_state_manager
+ return get_user_state_manager()
+
+ def _compute_embedding_hash(self, embeddings: Dict[str, Any]) -> str:
+ """Compute a hash of embedding IDs for cache invalidation."""
+ sorted_ids = sorted(embeddings.keys())
+ return hashlib.md5(",".join(sorted_ids).encode()).hexdigest()
+
+ def compute_umap_projection(
+ self,
+ embeddings: Dict[str, Any],
+ force: bool = False
+ ) -> Dict[str, Tuple[float, float]]:
+ """
+ Compute UMAP 2D projection of embeddings.
+
+ Args:
+ embeddings: Dict mapping instance_id to embedding vector
+ force: Force recomputation even if cached
+
+ Returns:
+ Dict mapping instance_id to (x, y) coordinates
+ """
+ if not self.enabled or not embeddings:
+ return {}
+
+ with self._lock:
+ # Check cache
+ current_hash = self._compute_embedding_hash(embeddings)
+ if not force and self._projection_cache and self._cache_hash == current_hash:
+ self.logger.debug("Using cached UMAP projection")
+ return self._projection_cache
+
+ try:
+ self.logger.info(f"Computing UMAP projection for {len(embeddings)} embeddings")
+
+ # Convert to numpy array
+ instance_ids = list(embeddings.keys())
+ vectors = np.array([embeddings[iid] for iid in instance_ids])
+
+ # Ensure we have enough samples for UMAP
+ n_samples = len(vectors)
+ n_neighbors = min(self.config.umap_n_neighbors, n_samples - 1)
+ if n_neighbors < 2:
+ self.logger.warning(f"Not enough samples ({n_samples}) for UMAP")
+ return {}
+
+ # Run UMAP
+ reducer = umap.UMAP(
+ n_neighbors=n_neighbors,
+ min_dist=self.config.umap_min_dist,
+ metric=self.config.umap_metric,
+ n_components=2,
+ random_state=42
+ )
+ projection = reducer.fit_transform(vectors)
+
+ # Build result dict
+ result = {}
+ for i, instance_id in enumerate(instance_ids):
+ result[instance_id] = (float(projection[i, 0]), float(projection[i, 1]))
+
+ # Cache result
+ self._projection_cache = result
+ self._cache_hash = current_hash
+
+ self.logger.info(f"UMAP projection complete: {len(result)} points")
+ return result
+
+ except Exception as e:
+ self.logger.error(f"UMAP projection failed: {e}")
+ return {}
+
+ def get_labels_for_instances(
+ self,
+ instance_ids: List[str],
+ source: str = "mace"
+ ) -> Dict[str, Optional[str]]:
+ """
+ Get predicted labels for instances.
+
+ Args:
+ instance_ids: List of instance IDs
+ source: Label source - "mace" or "majority"
+
+ Returns:
+ Dict mapping instance_id to label (or None if unannotated)
+ """
+ result = {}
+
+ if source == "mace":
+ result = self._get_mace_labels(instance_ids)
+ else:
+ result = self._get_majority_labels(instance_ids)
+
+ return result
+
+ def _get_mace_labels(self, instance_ids: List[str]) -> Dict[str, Optional[str]]:
+ """Get MACE predicted labels for instances."""
+ result = {iid: None for iid in instance_ids}
+
+ try:
+ from potato.mace_manager import get_mace_manager
+
+ mace_mgr = get_mace_manager()
+ if not mace_mgr or not mace_mgr.mace_config.enabled:
+ self.logger.debug("MACE not available, falling back to majority")
+ return self._get_majority_labels(instance_ids)
+
+ # Get predictions from all schemas
+ summary = mace_mgr.get_results_summary()
+ if "error" in summary or not summary.get("enabled"):
+ return self._get_majority_labels(instance_ids)
+
+ # Use first schema's predictions (most common case)
+ schemas = summary.get("schemas", {})
+ if not schemas:
+ return self._get_majority_labels(instance_ids)
+
+ # Get first schema with predictions
+ for schema_name, schema_data in schemas.items():
+ predictions = schema_data.get("predictions", {})
+ label_names = schema_data.get("label_names", [])
+
+ for instance_id in instance_ids:
+ if instance_id in predictions:
+ pred_idx = predictions[instance_id]
+ if isinstance(pred_idx, int) and pred_idx < len(label_names):
+ result[instance_id] = label_names[pred_idx]
+ break # Use first schema only
+
+ except ImportError:
+ self.logger.debug("MACE manager not available")
+ except Exception as e:
+ self.logger.error(f"Error getting MACE labels: {e}")
+
+ return result
+
+ def _get_majority_labels(self, instance_ids: List[str]) -> Dict[str, Optional[str]]:
+ """Get majority vote labels for instances."""
+ from collections import Counter
+
+ result = {iid: None for iid in instance_ids}
+
+ try:
+ usm = self._get_user_state_manager()
+ if not usm:
+ return result
+
+ # Get annotation schemes
+ annotation_schemes = self.app_config.get("annotation_schemes", [])
+ if not annotation_schemes:
+ return result
+
+ # Use first categorical schema
+ target_schema = None
+ for scheme in annotation_schemes:
+ if scheme.get("annotation_type") in ["radio", "select", "multiselect"]:
+ target_schema = scheme.get("name")
+ break
+
+ if not target_schema:
+ return result
+
+ # Count labels per instance
+ from potato.flask_server import get_users
+ users = get_users()
+
+ for instance_id in instance_ids:
+ labels = []
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if not user_state:
+ continue
+
+ annotations = user_state.get_all_annotations()
+ if instance_id not in annotations:
+ continue
+
+ instance_annot = annotations[instance_id]
+ label_annotations = instance_annot.get("labels", {})
+
+ for label, value in label_annotations.items():
+ label_schema = None
+ label_name = None
+
+ if hasattr(label, 'schema'):
+ label_schema = label.schema
+ label_name = getattr(label, 'name', None)
+ elif hasattr(label, 'get_schema'):
+ label_schema = label.get_schema()
+ label_name = label.get_name() if hasattr(label, 'get_name') else None
+
+ if label_schema == target_schema and label_name:
+ labels.append(label_name)
+
+ if labels:
+ counter = Counter(labels)
+ result[instance_id] = counter.most_common(1)[0][0]
+
+ except Exception as e:
+ self.logger.error(f"Error getting majority labels: {e}")
+
+ return result
+
+ def _assign_label_colors(self, unique_labels: List[Optional[str]]) -> Dict[Optional[str], str]:
+ """Assign consistent colors to labels."""
+ colors = {}
+ color_idx = 0
+
+ for label in unique_labels:
+ if label is None:
+ colors[None] = UNANNOTATED_COLOR
+ else:
+ colors[label] = DEFAULT_COLORS[color_idx % len(DEFAULT_COLORS)]
+ color_idx += 1
+
+ return colors
+
+ def get_visualization_data(self, force_refresh: bool = False) -> VisualizationData:
+ """
+ Get complete visualization data for the scatter plot.
+
+ Args:
+ force_refresh: Force recomputation of projections
+
+ Returns:
+ VisualizationData with points, labels, and colors
+ """
+ if not self.enabled:
+ return VisualizationData(
+ stats={"error": "Embedding visualization not enabled"}
+ )
+
+ with self._lock:
+ dm = self._get_diversity_manager()
+ ism = self._get_item_state_manager()
+
+ if not dm or not dm.enabled:
+ return VisualizationData(
+ stats={"error": "Diversity manager not available. Enable diversity_ordering in config."}
+ )
+
+ if not dm.embeddings:
+ return VisualizationData(
+ stats={"error": "No embeddings available. Ensure items have been loaded."}
+ )
+
+ # Get embeddings (possibly sampled)
+ all_embedding_ids = set(dm.embeddings.keys())
+ annotated_ids = set()
+
+ # Find annotated instances
+ if ism:
+ for instance_id in all_embedding_ids:
+ annotators = ism.get_annotators_for_item(instance_id)
+ if annotators:
+ annotated_ids.add(instance_id)
+
+ # Sample if needed
+ sample_ids = self._sample_instances(
+ all_embedding_ids,
+ annotated_ids,
+ self.config.sample_size,
+ self.config.include_all_annotated
+ )
+
+ # Get embeddings for sampled instances
+ sampled_embeddings = {
+ iid: dm.embeddings[iid]
+ for iid in sample_ids
+ if iid in dm.embeddings
+ }
+
+ # Compute UMAP projection
+ projection = self.compute_umap_projection(sampled_embeddings, force=force_refresh)
+ if not projection:
+ return VisualizationData(
+ stats={"error": "UMAP projection failed"}
+ )
+
+ # Get labels
+ labels = self.get_labels_for_instances(
+ list(projection.keys()),
+ source=self.config.label_source
+ )
+
+ # Build points
+ points = []
+ unique_labels = set()
+
+ for instance_id, (x, y) in projection.items():
+ label = labels.get(instance_id)
+ unique_labels.add(label)
+
+ # Get preview text
+ preview = ""
+ preview_type = "text"
+ if ism:
+ item = ism.get_instance_by_id(instance_id)
+ if item:
+ text = item.get_text()
+ if text:
+ preview = text[:200] + "..." if len(text) > 200 else text
+ # Check for image
+ if hasattr(item, 'get_image_path'):
+ img_path = item.get_image_path()
+ if img_path:
+ preview = img_path
+ preview_type = "image"
+
+ annotation_count = 0
+ if ism:
+ annotators = ism.get_annotators_for_item(instance_id)
+ annotation_count = len(annotators) if annotators else 0
+
+ points.append(VisualizationPoint(
+ instance_id=instance_id,
+ x=x,
+ y=y,
+ label=label,
+ label_source=self.config.label_source if label else None,
+ preview=preview,
+ preview_type=preview_type,
+ annotated=instance_id in annotated_ids,
+ annotation_count=annotation_count
+ ))
+
+ # Assign colors
+ label_colors = self._assign_label_colors(list(unique_labels))
+
+ # Build stats
+ stats = {
+ "total_instances": len(all_embedding_ids),
+ "visualized_instances": len(points),
+ "annotated_instances": len(annotated_ids),
+ "unannotated_instances": len(all_embedding_ids) - len(annotated_ids),
+ "label_source": self.config.label_source,
+ "unique_labels": len([l for l in unique_labels if l is not None])
+ }
+
+ return VisualizationData(
+ points=points,
+ labels=sorted([l for l in unique_labels if l is not None]) + [None],
+ label_colors=label_colors,
+ stats=stats
+ )
+
+ def _sample_instances(
+ self,
+ all_ids: Set[str],
+ annotated_ids: Set[str],
+ sample_size: int,
+ include_all_annotated: bool
+ ) -> Set[str]:
+ """
+ Sample instances for visualization.
+
+ Args:
+ all_ids: All available instance IDs
+ annotated_ids: IDs that have been annotated
+ sample_size: Maximum number of instances to include
+ include_all_annotated: Always include all annotated instances
+
+ Returns:
+ Set of instance IDs to visualize
+ """
+ if len(all_ids) <= sample_size:
+ return all_ids
+
+ result = set()
+
+ if include_all_annotated:
+ result.update(annotated_ids)
+
+ # Sample remaining from unannotated
+ remaining_needed = sample_size - len(result)
+ if remaining_needed > 0:
+ unannotated = all_ids - annotated_ids
+ if len(unannotated) <= remaining_needed:
+ result.update(unannotated)
+ else:
+ # Random sample
+ import random
+ sampled = random.sample(list(unannotated), remaining_needed)
+ result.update(sampled)
+
+ return result
+
+ def reorder_instances(
+ self,
+ selections: List[Dict[str, Any]],
+ interleave: bool = True
+ ) -> Dict[str, Any]:
+ """
+ Reorder the annotation queue based on selections.
+
+ Args:
+ selections: List of selection groups, each with:
+ - instance_ids: List of selected instance IDs
+ - priority: Priority number (lower = higher priority)
+ interleave: Whether to interleave selections (default True)
+
+ Returns:
+ Dict with success status and reordering info
+ """
+ if not selections:
+ return {"success": False, "error": "No selections provided"}
+
+ ism = self._get_item_state_manager()
+ if not ism:
+ return {"success": False, "error": "ItemStateManager not available"}
+
+ try:
+ # Build new order
+ if interleave:
+ new_order = self._interleave_selections(selections)
+ else:
+ # Concatenate by priority
+ sorted_selections = sorted(selections, key=lambda s: s.get("priority", 999))
+ new_order = []
+ for sel in sorted_selections:
+ new_order.extend(sel.get("instance_ids", []))
+
+ # Deduplicate while preserving order
+ seen = set()
+ deduped_order = []
+ for iid in new_order:
+ if iid not in seen:
+ seen.add(iid)
+ deduped_order.append(iid)
+
+ # Apply reordering
+ ism.reorder_instances(deduped_order)
+
+ # Build preview of new order (first 10)
+ preview = deduped_order[:10]
+
+ return {
+ "success": True,
+ "reordered_count": len(deduped_order),
+ "new_order_preview": preview
+ }
+
+ except Exception as e:
+ self.logger.error(f"Error reordering instances: {e}")
+ return {"success": False, "error": str(e)}
+
+ def _interleave_selections(self, selections: List[Dict[str, Any]]) -> List[str]:
+ """
+ Interleave instances from multiple selections by priority.
+
+ Example:
+ selections = [
+ {"instance_ids": ["a", "b", "c"], "priority": 1},
+ {"instance_ids": ["x", "y"], "priority": 2}
+ ]
+ Result: ["a", "x", "b", "y", "c"]
+
+ Lower priority number = higher priority (comes first in each round)
+
+ Args:
+ selections: List of selection dicts with instance_ids and priority
+
+ Returns:
+ List of interleaved instance IDs
+ """
+ # Sort by priority
+ sorted_selections = sorted(selections, key=lambda s: s.get("priority", 999))
+
+ # Create iterators
+ iterators = [iter(s.get("instance_ids", [])) for s in sorted_selections]
+
+ result = []
+ while iterators:
+ exhausted = []
+ for i, it in enumerate(iterators):
+ try:
+ result.append(next(it))
+ except StopIteration:
+ exhausted.append(i)
+
+ # Remove exhausted iterators (in reverse to maintain indices)
+ for i in reversed(exhausted):
+ iterators.pop(i)
+
+ return result
+
+ def invalidate_cache(self) -> None:
+ """Invalidate the projection cache."""
+ with self._lock:
+ self._projection_cache = None
+ self._cache_hash = None
+ self._label_cache = {}
+ self.logger.info("Embedding visualization cache invalidated")
+
+ def get_stats(self) -> Dict[str, Any]:
+ """Get visualization manager statistics."""
+ dm = self._get_diversity_manager()
+
+ return {
+ "enabled": self.enabled,
+ "umap_available": _UMAP_AVAILABLE,
+ "numpy_available": _NUMPY_AVAILABLE,
+ "embeddings_available": dm.enabled if dm else False,
+ "embedding_count": len(dm.embeddings) if dm and dm.embeddings else 0,
+ "cache_valid": self._projection_cache is not None,
+ "config": {
+ "sample_size": self.config.sample_size,
+ "include_all_annotated": self.config.include_all_annotated,
+ "label_source": self.config.label_source,
+ "umap_n_neighbors": self.config.umap_n_neighbors,
+ "umap_min_dist": self.config.umap_min_dist,
+ }
+ }
+
+ def to_json(self) -> Dict[str, Any]:
+ """Convert visualization data to JSON-serializable format."""
+ data = self.get_visualization_data()
+
+ points_json = []
+ for p in data.points:
+ points_json.append({
+ "instance_id": p.instance_id,
+ "x": p.x,
+ "y": p.y,
+ "label": p.label,
+ "label_source": p.label_source,
+ "preview": p.preview,
+ "preview_type": p.preview_type,
+ "annotated": p.annotated,
+ "annotation_count": p.annotation_count
+ })
+
+ return {
+ "points": points_json,
+ "labels": data.labels,
+ "label_colors": data.label_colors,
+ "stats": data.stats
+ }
+
+
+def parse_embedding_viz_config(config_data: Dict[str, Any]) -> EmbeddingVizConfig:
+ """
+ Parse embedding_visualization section from config.
+
+ Args:
+ config_data: Full application configuration
+
+ Returns:
+ EmbeddingVizConfig instance
+ """
+ ev = config_data.get("embedding_visualization", {})
+
+ return EmbeddingVizConfig(
+ enabled=ev.get("enabled", True),
+ sample_size=ev.get("sample_size", 1000),
+ include_all_annotated=ev.get("include_all_annotated", True),
+ embedding_model=ev.get("embedding_model", "all-MiniLM-L6-v2"),
+ image_embedding_model=ev.get("image_embedding_model", "clip-ViT-B-32"),
+ umap_n_neighbors=ev.get("umap", {}).get("n_neighbors", 15),
+ umap_min_dist=ev.get("umap", {}).get("min_dist", 0.1),
+ umap_metric=ev.get("umap", {}).get("metric", "cosine"),
+ label_source=ev.get("label_source", "mace"),
+ )
+
+
+def init_embedding_viz_manager(
+ config_data: Dict[str, Any]
+) -> Optional[EmbeddingVisualizationManager]:
+ """
+ Initialize the singleton EmbeddingVisualizationManager.
+
+ Args:
+ config_data: Full application configuration
+
+ Returns:
+ EmbeddingVisualizationManager instance, or None if disabled
+ """
+ global _EMBEDDING_VIZ_MANAGER
+
+ with _EMBEDDING_VIZ_LOCK:
+ if _EMBEDDING_VIZ_MANAGER is None:
+ viz_config = parse_embedding_viz_config(config_data)
+ _EMBEDDING_VIZ_MANAGER = EmbeddingVisualizationManager(viz_config, config_data)
+
+ return _EMBEDDING_VIZ_MANAGER
+
+
+def get_embedding_viz_manager() -> Optional[EmbeddingVisualizationManager]:
+ """Get the singleton EmbeddingVisualizationManager instance."""
+ return _EMBEDDING_VIZ_MANAGER
+
+
+def clear_embedding_viz_manager() -> None:
+ """Clear the singleton (for testing)."""
+ global _EMBEDDING_VIZ_MANAGER
+ with _EMBEDDING_VIZ_LOCK:
+ _EMBEDDING_VIZ_MANAGER = None
diff --git a/potato/expertise_manager.py b/potato/expertise_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..6a3f1407ba4def1fe4bb85ccc2ebf63f43acf913
--- /dev/null
+++ b/potato/expertise_manager.py
@@ -0,0 +1,529 @@
+"""
+Dynamic Category Expertise Manager
+
+This module provides dynamic category-based assignment where annotator expertise
+is determined by agreement with other annotators, rather than gold labels.
+
+Key features:
+- Tracks per-user, per-category expertise scores based on agreement
+- Background worker periodically recalculates expertise from annotation agreement
+- Probabilistic routing: all categories possible, weighted by expertise
+- Expertise increases when annotator agrees with consensus, decreases otherwise
+"""
+
+import logging
+import threading
+import time
+from collections import defaultdict, Counter
+from dataclasses import dataclass, field
+from typing import Dict, Set, List, Optional, Any, Tuple
+from enum import Enum
+
+logger = logging.getLogger(__name__)
+
+
+class AgreementMethod(Enum):
+ """Methods for calculating agreement/consensus."""
+ MAJORITY_VOTE = "majority_vote" # Simple majority
+ SUPER_MAJORITY = "super_majority" # 2/3 or more agree
+ UNANIMOUS = "unanimous" # All must agree
+
+
+@dataclass
+class CategoryExpertise:
+ """Tracks expertise metrics for a single category."""
+ category: str
+ agreements: int = 0 # Times user agreed with consensus
+ disagreements: int = 0 # Times user disagreed with consensus
+ total_evaluated: int = 0 # Total instances evaluated for this category
+ expertise_score: float = 0.5 # 0.0 to 1.0, starts neutral
+
+ def update_score(self, agreed: bool, learning_rate: float = 0.1) -> None:
+ """
+ Update expertise score based on agreement.
+
+ Uses exponential moving average to smooth updates.
+ """
+ self.total_evaluated += 1
+ if agreed:
+ self.agreements += 1
+ # Increase score, but cap at 1.0
+ self.expertise_score = min(1.0,
+ self.expertise_score + learning_rate * (1.0 - self.expertise_score))
+ else:
+ self.disagreements += 1
+ # Decrease score, but floor at 0.0
+ self.expertise_score = max(0.0,
+ self.expertise_score - learning_rate * self.expertise_score)
+
+ def get_accuracy(self) -> float:
+ """Get raw agreement accuracy."""
+ if self.total_evaluated == 0:
+ return 0.5 # Neutral if no data
+ return self.agreements / self.total_evaluated
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'category': self.category,
+ 'agreements': self.agreements,
+ 'disagreements': self.disagreements,
+ 'total_evaluated': self.total_evaluated,
+ 'expertise_score': self.expertise_score
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'CategoryExpertise':
+ """Deserialize from dictionary."""
+ return cls(
+ category=data['category'],
+ agreements=data.get('agreements', 0),
+ disagreements=data.get('disagreements', 0),
+ total_evaluated=data.get('total_evaluated', 0),
+ expertise_score=data.get('expertise_score', 0.5)
+ )
+
+
+@dataclass
+class UserExpertiseProfile:
+ """Complete expertise profile for a user."""
+ user_id: str
+ category_expertise: Dict[str, CategoryExpertise] = field(default_factory=dict)
+ evaluated_instances: Set[str] = field(default_factory=set) # Instances already evaluated
+ last_updated: float = 0.0 # Timestamp of last update
+
+ def get_expertise(self, category: str) -> CategoryExpertise:
+ """Get or create expertise for a category."""
+ if category not in self.category_expertise:
+ self.category_expertise[category] = CategoryExpertise(category=category)
+ return self.category_expertise[category]
+
+ def get_expertise_score(self, category: str) -> float:
+ """Get expertise score for a category (0.5 if unknown)."""
+ if category in self.category_expertise:
+ return self.category_expertise[category].expertise_score
+ return 0.5 # Neutral for unknown categories
+
+ def get_all_expertise_scores(self) -> Dict[str, float]:
+ """Get all category expertise scores."""
+ return {cat: exp.expertise_score for cat, exp in self.category_expertise.items()}
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'user_id': self.user_id,
+ 'category_expertise': {
+ cat: exp.to_dict() for cat, exp in self.category_expertise.items()
+ },
+ 'evaluated_instances': list(self.evaluated_instances),
+ 'last_updated': self.last_updated
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'UserExpertiseProfile':
+ """Deserialize from dictionary."""
+ profile = cls(user_id=data['user_id'])
+ profile.category_expertise = {
+ cat: CategoryExpertise.from_dict(exp_data)
+ for cat, exp_data in data.get('category_expertise', {}).items()
+ }
+ profile.evaluated_instances = set(data.get('evaluated_instances', []))
+ profile.last_updated = data.get('last_updated', 0.0)
+ return profile
+
+
+class ExpertiseManager:
+ """
+ Manages dynamic category expertise for all users.
+
+ This class:
+ - Tracks expertise scores per user per category
+ - Calculates agreement with consensus for completed instances
+ - Updates expertise scores based on agreement
+ - Provides weighted probabilities for category assignment
+ """
+
+ _instance = None
+ _lock = threading.RLock()
+
+ def __new__(cls, *args, **kwargs):
+ """Singleton pattern."""
+ if cls._instance is None:
+ with cls._lock:
+ if cls._instance is None:
+ cls._instance = super().__new__(cls)
+ cls._instance._initialized = False
+ return cls._instance
+
+ def __init__(self, config: Optional[Dict[str, Any]] = None):
+ """
+ Initialize the ExpertiseManager.
+
+ Args:
+ config: Configuration dictionary with settings
+ """
+ if self._initialized:
+ return
+
+ self.config = config or {}
+ self.user_profiles: Dict[str, UserExpertiseProfile] = {}
+
+ # Configuration options
+ dynamic_config = self.config.get('category_assignment', {}).get('dynamic', {})
+ self.min_annotations_for_consensus = dynamic_config.get('min_annotations_for_consensus', 2)
+ self.agreement_method = AgreementMethod(
+ dynamic_config.get('agreement_method', 'majority_vote')
+ )
+ self.learning_rate = dynamic_config.get('learning_rate', 0.1)
+ self.update_interval_seconds = dynamic_config.get('update_interval_seconds', 60)
+ self.base_probability = dynamic_config.get('base_probability', 0.1) # Min probability for any category
+
+ # Background worker
+ self._worker_thread: Optional[threading.Thread] = None
+ self._stop_worker = threading.Event()
+
+ self._initialized = True
+ logger.info("ExpertiseManager initialized")
+
+ def get_user_profile(self, user_id: str) -> UserExpertiseProfile:
+ """Get or create expertise profile for a user."""
+ with self._lock:
+ if user_id not in self.user_profiles:
+ self.user_profiles[user_id] = UserExpertiseProfile(user_id=user_id)
+ return self.user_profiles[user_id]
+
+ def calculate_consensus(
+ self,
+ instance_id: str,
+ category: str,
+ schema_name: str
+ ) -> Optional[Tuple[Any, int]]:
+ """
+ Calculate the consensus annotation for an instance.
+
+ Args:
+ instance_id: The instance ID
+ category: The category of the instance
+ schema_name: The annotation schema to check
+
+ Returns:
+ Tuple of (consensus_value, num_annotators) or None if not enough annotations
+ """
+ # Import here to avoid circular imports
+ from potato.flask_server import get_users, get_user_state
+
+ annotations = []
+ for username in get_users():
+ user_state = get_user_state(username)
+ if user_state:
+ all_annotations = user_state.get_all_annotations()
+ if instance_id in all_annotations:
+ instance_annotations = all_annotations[instance_id]
+ if 'labels' in instance_annotations:
+ for label, value in instance_annotations['labels'].items():
+ if label.get_schema() == schema_name:
+ annotations.append((username, value))
+
+ if len(annotations) < self.min_annotations_for_consensus:
+ return None
+
+ # Extract just the values
+ values = [v for _, v in annotations]
+ counter = Counter(values)
+ most_common_value, most_common_count = counter.most_common(1)[0]
+
+ # Check agreement method
+ if self.agreement_method == AgreementMethod.MAJORITY_VOTE:
+ if most_common_count > len(values) / 2:
+ return (most_common_value, len(annotations))
+ elif self.agreement_method == AgreementMethod.SUPER_MAJORITY:
+ if most_common_count >= len(values) * 2 / 3:
+ return (most_common_value, len(annotations))
+ elif self.agreement_method == AgreementMethod.UNANIMOUS:
+ if most_common_count == len(values):
+ return (most_common_value, len(annotations))
+
+ return None # No clear consensus
+
+ def update_user_expertise(
+ self,
+ user_id: str,
+ instance_id: str,
+ category: str,
+ user_annotation: Any,
+ consensus_value: Any
+ ) -> bool:
+ """
+ Update a user's expertise based on agreement with consensus.
+
+ Args:
+ user_id: The user ID
+ instance_id: The instance ID
+ category: The category of the instance
+ user_annotation: The user's annotation value
+ consensus_value: The consensus annotation value
+
+ Returns:
+ True if user agreed with consensus, False otherwise
+ """
+ with self._lock:
+ profile = self.get_user_profile(user_id)
+
+ # Skip if already evaluated
+ eval_key = f"{instance_id}:{category}"
+ if eval_key in profile.evaluated_instances:
+ return False
+
+ # Determine agreement
+ agreed = (user_annotation == consensus_value)
+
+ # Update expertise
+ expertise = profile.get_expertise(category)
+ expertise.update_score(agreed, self.learning_rate)
+
+ # Mark as evaluated
+ profile.evaluated_instances.add(eval_key)
+ profile.last_updated = time.time()
+
+ logger.debug(
+ f"Updated expertise for {user_id} in {category}: "
+ f"agreed={agreed}, new_score={expertise.expertise_score:.3f}"
+ )
+
+ return agreed
+
+ def evaluate_all_instances(self) -> Dict[str, int]:
+ """
+ Evaluate all completed instances and update expertise scores.
+
+ This is called by the background worker periodically.
+
+ Returns:
+ Dictionary mapping user_id to number of new evaluations
+ """
+ # Import here to avoid circular imports
+ from potato.flask_server import get_users, get_user_state, get_item_state_manager
+
+ updates_per_user: Dict[str, int] = defaultdict(int)
+
+ try:
+ ism = get_item_state_manager()
+ if ism is None:
+ return updates_per_user
+
+ # Get all annotation schemes from config
+ annotation_schemes = self.config.get('annotation_schemes', [])
+ schema_names = [s.get('name') for s in annotation_schemes if s.get('name')]
+
+ if not schema_names:
+ return updates_per_user
+
+ # Use first schema for agreement calculation
+ primary_schema = schema_names[0]
+
+ # Get all instances with categories
+ for instance_id in ism.instance_id_ordering:
+ categories = ism.get_categories_for_instance(instance_id)
+ if not categories:
+ continue
+
+ for category in categories:
+ # Try to get consensus
+ consensus_result = self.calculate_consensus(
+ instance_id, category, primary_schema
+ )
+
+ if consensus_result is None:
+ continue # Not enough annotations yet
+
+ consensus_value, num_annotators = consensus_result
+
+ # Update each user who annotated this instance
+ for username in get_users():
+ user_state = get_user_state(username)
+ if not user_state:
+ continue
+
+ all_annotations = user_state.get_all_annotations()
+ if instance_id not in all_annotations:
+ continue
+
+ instance_annotations = all_annotations[instance_id]
+ if 'labels' not in instance_annotations:
+ continue
+
+ # Find user's annotation for this schema
+ user_value = None
+ for label, value in instance_annotations['labels'].items():
+ if label.get_schema() == primary_schema:
+ user_value = value
+ break
+
+ if user_value is not None:
+ self.update_user_expertise(
+ username, instance_id, category,
+ user_value, consensus_value
+ )
+ updates_per_user[username] += 1
+
+ except Exception as e:
+ logger.error(f"Error evaluating instances: {e}")
+
+ return updates_per_user
+
+ def get_category_probabilities(
+ self,
+ user_id: str,
+ available_categories: Set[str]
+ ) -> Dict[str, float]:
+ """
+ Calculate assignment probabilities for each category.
+
+ Uses expertise scores to weight categories, but ensures all categories
+ have at least base_probability chance.
+
+ Args:
+ user_id: The user ID
+ available_categories: Set of categories with available instances
+
+ Returns:
+ Dictionary mapping category to probability (sums to 1.0)
+ """
+ if not available_categories:
+ return {}
+
+ profile = self.get_user_profile(user_id)
+
+ # Get raw scores for each category
+ raw_scores = {}
+ for category in available_categories:
+ score = profile.get_expertise_score(category)
+ # Ensure minimum probability
+ raw_scores[category] = max(self.base_probability, score)
+
+ # Normalize to probabilities
+ total = sum(raw_scores.values())
+ if total == 0:
+ # Equal probability if no scores
+ prob = 1.0 / len(available_categories)
+ return {cat: prob for cat in available_categories}
+
+ return {cat: score / total for cat, score in raw_scores.items()}
+
+ def select_category_probabilistically(
+ self,
+ user_id: str,
+ available_categories: Set[str],
+ random_instance=None
+ ) -> Optional[str]:
+ """
+ Select a category using weighted random selection.
+
+ Args:
+ user_id: The user ID
+ available_categories: Set of categories with available instances
+ random_instance: Optional random.Random instance for reproducibility
+
+ Returns:
+ Selected category name, or None if no categories available
+ """
+ import random as random_module
+
+ if not available_categories:
+ return None
+
+ probs = self.get_category_probabilities(user_id, available_categories)
+ if not probs:
+ return None
+
+ rng = random_instance or random_module
+ categories = list(probs.keys())
+ weights = [probs[cat] for cat in categories]
+
+ # Use random.choices for weighted selection
+ selected = rng.choices(categories, weights=weights, k=1)[0]
+ return selected
+
+ def start_background_worker(self) -> None:
+ """Start the background worker thread for periodic expertise updates."""
+ if self._worker_thread is not None and self._worker_thread.is_alive():
+ logger.warning("Background worker already running")
+ return
+
+ self._stop_worker.clear()
+ self._worker_thread = threading.Thread(
+ target=self._background_worker_loop,
+ name="ExpertiseWorker",
+ daemon=True
+ )
+ self._worker_thread.start()
+ logger.info("Started expertise background worker")
+
+ def stop_background_worker(self) -> None:
+ """Stop the background worker thread."""
+ if self._worker_thread is None:
+ return
+
+ self._stop_worker.set()
+ self._worker_thread.join(timeout=5.0)
+ self._worker_thread = None
+ logger.info("Stopped expertise background worker")
+
+ def _background_worker_loop(self) -> None:
+ """Main loop for the background worker."""
+ logger.info(f"Background worker started, interval={self.update_interval_seconds}s")
+
+ while not self._stop_worker.is_set():
+ try:
+ updates = self.evaluate_all_instances()
+ if updates:
+ total_updates = sum(updates.values())
+ logger.info(f"Background worker: {total_updates} expertise updates")
+ except Exception as e:
+ logger.error(f"Background worker error: {e}")
+
+ # Wait for next interval or stop signal
+ self._stop_worker.wait(self.update_interval_seconds)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize all expertise data."""
+ with self._lock:
+ return {
+ 'user_profiles': {
+ user_id: profile.to_dict()
+ for user_id, profile in self.user_profiles.items()
+ }
+ }
+
+ def from_dict(self, data: Dict[str, Any]) -> None:
+ """Load expertise data from dictionary."""
+ with self._lock:
+ self.user_profiles = {
+ user_id: UserExpertiseProfile.from_dict(profile_data)
+ for user_id, profile_data in data.get('user_profiles', {}).items()
+ }
+
+
+# Module-level singleton access
+_expertise_manager: Optional[ExpertiseManager] = None
+
+
+def init_expertise_manager(config: Dict[str, Any]) -> ExpertiseManager:
+ """Initialize the global expertise manager."""
+ global _expertise_manager
+ _expertise_manager = ExpertiseManager(config)
+ return _expertise_manager
+
+
+def get_expertise_manager() -> Optional[ExpertiseManager]:
+ """Get the global expertise manager instance."""
+ return _expertise_manager
+
+
+def clear_expertise_manager() -> None:
+ """Clear the global expertise manager (for testing)."""
+ global _expertise_manager
+ if _expertise_manager:
+ _expertise_manager.stop_background_worker()
+ _expertise_manager = None
+ ExpertiseManager._instance = None
diff --git a/potato/export/__init__.py b/potato/export/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e56e6d28708a240a63eed8779210ea2fbf174d77
--- /dev/null
+++ b/potato/export/__init__.py
@@ -0,0 +1,23 @@
+"""
+Potato Export System
+
+Pluggable export framework for converting Potato annotations into
+standard formats (COCO, YOLO, Pascal VOC, CoNLL-2003, CoNLL-U, etc.).
+
+Usage:
+ from potato.export.registry import export_registry
+
+ # List available exporters
+ exporters = export_registry.list_exporters()
+
+ # Export annotations
+ result = export_registry.export("coco", context, output_path)
+
+CLI:
+ python -m potato.export --config config.yaml --format coco --output ./out/
+"""
+
+from .base import BaseExporter, ExportContext, ExportResult
+from .registry import export_registry
+
+__all__ = ["BaseExporter", "ExportContext", "ExportResult", "export_registry"]
diff --git a/potato/export/__main__.py b/potato/export/__main__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e072c17bf7b767ae88f5692d3e53077986dfb2b1
--- /dev/null
+++ b/potato/export/__main__.py
@@ -0,0 +1,4 @@
+"""Allow running export as a module: python -m potato.export"""
+from .cli import main
+
+main()
diff --git a/potato/export/agent_eval_exporter.py b/potato/export/agent_eval_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..c7973360fe969cf06117054aa83fb07aa3295f06
--- /dev/null
+++ b/potato/export/agent_eval_exporter.py
@@ -0,0 +1,403 @@
+"""
+Agent Evaluation Exporter
+
+Exports annotations in a structured format optimized for agent evaluation,
+producing per-trace aggregated scores, error distributions, and per-step
+assessment summaries.
+
+Output format:
+{
+ "summary": {
+ "total_traces": 10,
+ "total_annotators": 3,
+ "schemas_evaluated": ["task_success", "efficiency", ...]
+ },
+ "per_trace": [
+ {
+ "trace_id": "trace_001",
+ "annotations": {
+ "task_success": {"distribution": {"success": 2, "partial": 1}, "majority": "success"},
+ "efficiency": {"mean": 4.2, "std": 0.5, "values": [4, 5, 4]},
+ "mast_errors": {"counts": {"no_errors": 3}, "total_annotations": 3}
+ },
+ "annotator_count": 3
+ }
+ ],
+ "aggregate": {
+ "task_success": {"success_rate": 0.7, "partial_rate": 0.2, "failure_rate": 0.1},
+ "efficiency": {"overall_mean": 3.8, "overall_std": 0.9},
+ "mast_errors": {"total_distribution": {"no_errors": 25, "step_repetition": 3, ...}}
+ }
+}
+"""
+
+import csv
+import io
+import json
+import logging
+import os
+from collections import Counter, defaultdict
+from typing import Dict, List, Any, Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+
+logger = logging.getLogger(__name__)
+
+
+class AgentEvalExporter(BaseExporter):
+ """
+ Exporter for agent trace evaluation annotations.
+
+ Produces structured JSON output optimized for evaluation dashboards
+ and leaderboard computation.
+ """
+
+ format_name = "agent_eval"
+ description = "Agent evaluation export with aggregated scores and error distributions"
+ file_extensions = [".json"]
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ files_written = []
+ warnings = []
+
+ try:
+ # Group annotations by trace
+ trace_annotations = self._group_by_trace(context.annotations)
+
+ # Get schema info
+ schema_map = {s["name"]: s for s in context.schemas}
+
+ # Compute per-trace aggregations
+ per_trace_results = []
+ for trace_id, annotations in sorted(trace_annotations.items()):
+ trace_result = self._aggregate_trace(trace_id, annotations, schema_map)
+ per_trace_results.append(trace_result)
+
+ # Compute global aggregations
+ aggregate = self._compute_aggregate(per_trace_results, schema_map)
+
+ # Build summary
+ all_annotators = set()
+ for anns in trace_annotations.values():
+ for ann in anns:
+ all_annotators.add(ann.get("user_id", "unknown"))
+
+ summary = {
+ "total_traces": len(trace_annotations),
+ "total_annotators": len(all_annotators),
+ "annotators": sorted(all_annotators),
+ "schemas_evaluated": sorted(schema_map.keys()),
+ }
+
+ # Build output
+ output = {
+ "summary": summary,
+ "per_trace": per_trace_results,
+ "aggregate": aggregate,
+ }
+
+ # Write output
+ os.makedirs(output_path, exist_ok=True)
+ output_file = os.path.join(output_path, "agent_evaluation.json")
+ with open(output_file, "w", encoding="utf-8") as f:
+ json.dump(output, f, indent=2, ensure_ascii=False)
+ files_written.append(output_file)
+
+ # Also write a per-trace CSV for easy analysis
+ csv_file = os.path.join(output_path, "agent_evaluation_summary.csv")
+ self._write_summary_csv(csv_file, per_trace_results, schema_map)
+ files_written.append(csv_file)
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=files_written,
+ warnings=warnings,
+ stats={
+ "total_traces": len(trace_annotations),
+ "total_annotations": sum(len(a) for a in trace_annotations.values()),
+ "total_annotators": len(all_annotators),
+ },
+ )
+
+ except Exception as e:
+ logger.error(f"Agent eval export failed: {e}")
+ return ExportResult(
+ success=False,
+ format_name=self.format_name,
+ errors=[str(e)],
+ )
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ if not context.annotations:
+ return False, "No annotations to export"
+ if not context.schemas:
+ return False, "No annotation schemas defined"
+ return True, ""
+
+ def _group_by_trace(self, annotations: List[dict]) -> Dict[str, List[dict]]:
+ """Group annotations by instance (trace) ID."""
+ grouped = defaultdict(list)
+ for ann in annotations:
+ trace_id = ann.get("instance_id", "unknown")
+ grouped[trace_id].append(ann)
+ return dict(grouped)
+
+ def _aggregate_trace(self, trace_id: str, annotations: List[dict],
+ schema_map: Dict[str, dict]) -> dict:
+ """Aggregate annotations for a single trace."""
+ result = {
+ "trace_id": trace_id,
+ "annotator_count": len(set(a.get("user_id", "unknown") for a in annotations)),
+ "annotations": {},
+ }
+
+ # Group by schema
+ schema_values = defaultdict(list)
+ for ann in annotations:
+ labels = ann.get("labels", {})
+ for schema_name, value in labels.items():
+ schema_values[schema_name].append(value)
+
+ # Aggregate each schema
+ for schema_name, values in schema_values.items():
+ schema_config = schema_map.get(schema_name, {})
+ schema_type = schema_config.get("annotation_type", "")
+
+ if schema_type in ("radio", "select"):
+ result["annotations"][schema_name] = self._aggregate_categorical(values)
+ elif schema_type in ("likert", "slider", "number"):
+ result["annotations"][schema_name] = self._aggregate_numeric(values)
+ elif schema_type == "multiselect":
+ result["annotations"][schema_name] = self._aggregate_multiselect(values)
+ elif schema_type == "multirate":
+ result["annotations"][schema_name] = self._aggregate_multirate(values)
+ elif schema_type == "text":
+ result["annotations"][schema_name] = {"responses": values}
+ else:
+ result["annotations"][schema_name] = {"values": values}
+
+ return result
+
+ def _aggregate_categorical(self, values: List) -> dict:
+ """Aggregate categorical (radio/select) annotations."""
+ # Flatten nested dicts - values might be {"label": "value"} or just strings
+ flat_values = []
+ for v in values:
+ if isinstance(v, dict):
+ # Take the key with the highest value (convert to float for comparison)
+ if v:
+ def _sort_key(k):
+ try:
+ return float(v[k])
+ except (ValueError, TypeError):
+ return 0
+ flat_values.append(max(v.keys(), key=_sort_key))
+ else:
+ flat_values.append(str(v))
+
+ distribution = dict(Counter(flat_values))
+ majority = max(distribution, key=distribution.get) if distribution else ""
+
+ return {
+ "distribution": distribution,
+ "majority": majority,
+ "agreement": max(distribution.values()) / len(flat_values) if flat_values else 0,
+ }
+
+ def _aggregate_numeric(self, values: List) -> dict:
+ """Aggregate numeric (likert/slider) annotations."""
+ numeric_values = []
+ for v in values:
+ if isinstance(v, (int, float)):
+ numeric_values.append(float(v))
+ elif isinstance(v, str):
+ try:
+ numeric_values.append(float(v))
+ except ValueError:
+ pass
+ elif isinstance(v, dict):
+ # Try to extract numeric value
+ for val in v.values():
+ try:
+ numeric_values.append(float(val))
+ except (ValueError, TypeError):
+ pass
+
+ if not numeric_values:
+ return {"mean": None, "values": values}
+
+ mean = sum(numeric_values) / len(numeric_values)
+ # Use sample standard deviation (N-1) when N > 1, population (N) when N == 1
+ n = len(numeric_values)
+ variance = sum((x - mean) ** 2 for x in numeric_values) / max(n - 1, 1)
+ std = variance ** 0.5
+
+ return {
+ "mean": round(mean, 3),
+ "std": round(std, 3),
+ "min": min(numeric_values),
+ "max": max(numeric_values),
+ "values": numeric_values,
+ }
+
+ def _aggregate_multiselect(self, values: List) -> dict:
+ """Aggregate multiselect annotations."""
+ counts = Counter()
+ total = 0
+ for v in values:
+ total += 1
+ if isinstance(v, dict):
+ for label, selected in v.items():
+ if selected:
+ counts[label] += 1
+ elif isinstance(v, list):
+ for label in v:
+ counts[label] += 1
+
+ return {
+ "counts": dict(counts),
+ "total_annotations": total,
+ }
+
+ def _aggregate_multirate(self, values: List) -> dict:
+ """Aggregate multirate annotations."""
+ item_ratings = defaultdict(list)
+ for v in values:
+ if isinstance(v, dict):
+ for item_name, rating in v.items():
+ try:
+ item_ratings[item_name].append(float(rating))
+ except (ValueError, TypeError):
+ item_ratings[item_name].append(rating)
+
+ result = {}
+ for item_name, ratings in item_ratings.items():
+ numeric = [r for r in ratings if isinstance(r, (int, float))]
+ if numeric:
+ result[item_name] = {
+ "mean": round(sum(numeric) / len(numeric), 3),
+ "values": ratings,
+ }
+ else:
+ result[item_name] = {"values": ratings}
+
+ return {"per_item": result}
+
+ def _compute_aggregate(self, per_trace_results: List[dict],
+ schema_map: Dict[str, dict]) -> dict:
+ """Compute aggregate statistics across all traces."""
+ aggregate = {}
+
+ for schema_name, schema_config in schema_map.items():
+ schema_type = schema_config.get("annotation_type", "")
+
+ if schema_type in ("radio", "select"):
+ aggregate[schema_name] = self._aggregate_categorical_global(
+ per_trace_results, schema_name
+ )
+ elif schema_type in ("likert", "slider", "number"):
+ aggregate[schema_name] = self._aggregate_numeric_global(
+ per_trace_results, schema_name
+ )
+ elif schema_type == "multiselect":
+ aggregate[schema_name] = self._aggregate_multiselect_global(
+ per_trace_results, schema_name
+ )
+
+ return aggregate
+
+ def _aggregate_categorical_global(self, results: List[dict], schema_name: str) -> dict:
+ """Compute global rates for categorical annotations."""
+ all_majorities = []
+ total_dist = Counter()
+
+ for result in results:
+ ann = result.get("annotations", {}).get(schema_name, {})
+ if "majority" in ann:
+ all_majorities.append(ann["majority"])
+ if "distribution" in ann:
+ for label, count in ann["distribution"].items():
+ total_dist[label] += count
+
+ # Compute rates
+ total = sum(total_dist.values())
+ rates = {}
+ for label, count in total_dist.items():
+ rates[f"{label}_rate"] = round(count / total, 3) if total > 0 else 0
+
+ return {
+ "rates": rates,
+ "total_distribution": dict(total_dist),
+ "majority_distribution": dict(Counter(all_majorities)),
+ }
+
+ def _aggregate_numeric_global(self, results: List[dict], schema_name: str) -> dict:
+ """Compute global stats for numeric annotations."""
+ all_means = []
+ for result in results:
+ ann = result.get("annotations", {}).get(schema_name, {})
+ if ann.get("mean") is not None:
+ all_means.append(ann["mean"])
+
+ if not all_means:
+ return {"overall_mean": None}
+
+ overall_mean = sum(all_means) / len(all_means)
+ n = len(all_means)
+ variance = sum((x - overall_mean) ** 2 for x in all_means) / max(n - 1, 1)
+
+ return {
+ "overall_mean": round(overall_mean, 3),
+ "overall_std": round(variance ** 0.5, 3),
+ "num_traces": len(all_means),
+ }
+
+ def _aggregate_multiselect_global(self, results: List[dict], schema_name: str) -> dict:
+ """Compute global counts for multiselect annotations."""
+ total_counts = Counter()
+ for result in results:
+ ann = result.get("annotations", {}).get(schema_name, {})
+ for label, count in ann.get("counts", {}).items():
+ total_counts[label] += count
+
+ return {"total_distribution": dict(total_counts)}
+
+ def _write_summary_csv(self, csv_path: str, per_trace_results: List[dict],
+ schema_map: Dict[str, dict]) -> None:
+ """Write a summary CSV with one row per trace."""
+ if not per_trace_results:
+ return
+
+ # Collect all column names
+ columns = ["trace_id", "annotator_count"]
+ for result in per_trace_results:
+ for schema_name in result.get("annotations", {}):
+ schema_type = schema_map.get(schema_name, {}).get("annotation_type", "")
+ if schema_type in ("radio", "select"):
+ col = f"{schema_name}_majority"
+ if col not in columns:
+ columns.append(col)
+ elif schema_type in ("likert", "slider", "number"):
+ col = f"{schema_name}_mean"
+ if col not in columns:
+ columns.append(col)
+
+ # Write CSV using csv module for proper escaping
+ with open(csv_path, "w", encoding="utf-8", newline="") as f:
+ writer = csv.writer(f)
+ writer.writerow(columns)
+ for result in per_trace_results:
+ row = [result["trace_id"], str(result["annotator_count"])]
+ for col in columns[2:]:
+ schema_name = col.rsplit("_", 1)[0]
+ ann = result.get("annotations", {}).get(schema_name, {})
+ if col.endswith("_majority"):
+ row.append(str(ann.get("majority", "")))
+ elif col.endswith("_mean"):
+ row.append(str(ann.get("mean", "")))
+ else:
+ row.append("")
+ writer.writerow(row)
diff --git a/potato/export/base.py b/potato/export/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..a8a9bfb5b76fddff6ecef90d98737dcd14704bba
--- /dev/null
+++ b/potato/export/base.py
@@ -0,0 +1,108 @@
+"""
+Export Base Classes
+
+Defines the abstract base class for exporters and data structures
+for passing annotation data through the export pipeline.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, List, Any, Optional, Tuple
+
+
+@dataclass
+class ExportContext:
+ """
+ Container for all data needed by an exporter.
+
+ Attributes:
+ config: Full Potato YAML configuration dictionary
+ annotations: Flattened list of annotation records, each containing:
+ - instance_id: str
+ - user_id: str
+ - labels: dict mapping schema_name -> {label: value}
+ - spans: dict mapping schema_name -> list of span dicts
+ - links: dict mapping schema_name -> list of link dicts
+ items: Mapping of instance_id -> item data dict (original data)
+ schemas: List of annotation_scheme configuration dicts
+ output_dir: Base output directory path
+ """
+ config: dict
+ annotations: List[dict]
+ items: Dict[str, dict]
+ schemas: List[dict]
+ output_dir: str
+ phase_responses: List[dict] = field(default_factory=list)
+
+
+@dataclass
+class ExportResult:
+ """
+ Result of an export operation.
+
+ Attributes:
+ success: Whether the export completed successfully
+ format_name: Name of the export format used
+ files_written: List of file paths that were created
+ warnings: Non-fatal issues encountered during export
+ errors: Fatal errors that prevented full export
+ stats: Summary statistics (e.g., num_images, num_annotations)
+ """
+ success: bool
+ format_name: str
+ files_written: List[str] = field(default_factory=list)
+ warnings: List[str] = field(default_factory=list)
+ errors: List[str] = field(default_factory=list)
+ stats: Dict[str, Any] = field(default_factory=dict)
+
+
+class BaseExporter(ABC):
+ """
+ Abstract base class for annotation exporters.
+
+ Subclasses must implement:
+ - export(): Perform the actual export
+ - can_export(): Check if the context is compatible with this format
+ """
+
+ format_name: str = ""
+ description: str = ""
+ file_extensions: List[str] = []
+
+ @abstractmethod
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ """
+ Export annotations to the target format.
+
+ Args:
+ context: ExportContext containing all annotation data
+ output_path: Directory or file path for output
+ options: Format-specific options
+
+ Returns:
+ ExportResult with status and written file paths
+ """
+ ...
+
+ @abstractmethod
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ """
+ Check whether this exporter can handle the given context.
+
+ Args:
+ context: ExportContext to validate
+
+ Returns:
+ Tuple of (can_export: bool, reason: str).
+ If can_export is False, reason explains why.
+ """
+ ...
+
+ def get_format_info(self) -> dict:
+ """Return metadata about this export format."""
+ return {
+ "format_name": self.format_name,
+ "description": self.description,
+ "file_extensions": self.file_extensions,
+ }
diff --git a/potato/export/cli.py b/potato/export/cli.py
new file mode 100644
index 0000000000000000000000000000000000000000..d045a983aa44b9a065ed954a7dddd143c4d99b3a
--- /dev/null
+++ b/potato/export/cli.py
@@ -0,0 +1,393 @@
+"""
+Export CLI
+
+Command-line interface for exporting Potato annotations to various formats.
+
+Usage:
+ python -m potato.export --config config.yaml --format coco --output ./out/
+ python -m potato.export --config config.yaml --format conll_2003 --output ./out/
+ python -m potato.export --list-formats
+"""
+
+import argparse
+import json
+import os
+import sys
+import logging
+import glob
+
+import yaml
+
+from .base import ExportContext
+from .registry import export_registry
+
+logger = logging.getLogger(__name__)
+
+
+def load_annotations_from_output_dir(output_dir: str, schemas: list) -> list:
+ """
+ Load user annotations from the Potato output directory.
+
+ Reads user_state.json files from each user subdirectory
+ and flattens annotations into a list of records.
+
+ Args:
+ output_dir: Path to the annotation output directory
+ schemas: List of annotation scheme configs
+
+ Returns:
+ List of annotation dicts
+ """
+ annotations = []
+
+ if not os.path.isdir(output_dir):
+ logger.warning(f"Output directory not found: {output_dir}")
+ return annotations
+
+ for user_dir in sorted(os.listdir(output_dir)):
+ user_path = os.path.join(output_dir, user_dir)
+ if not os.path.isdir(user_path):
+ continue
+
+ state_file = os.path.join(user_path, "user_state.json")
+ if not os.path.exists(state_file):
+ continue
+
+ with open(state_file, "r") as f:
+ user_state = json.load(f)
+
+ user_id = user_state.get("user_id", user_dir)
+
+ # Extract label annotations
+ label_data = user_state.get("instance_id_to_label_to_value", {})
+ span_data = user_state.get("instance_id_to_span_to_value", {})
+
+ # Collect all instance IDs
+ all_instances = set(label_data.keys()) | set(span_data.keys())
+
+ for instance_id in all_instances:
+ # Labels may be stored as a list of [[{schema, name}, value], ...]
+ # or as a dict of {schema_name: {label_name: value}}.
+ # Normalize to dict format.
+ raw_labels = label_data.get(instance_id, {})
+ if isinstance(raw_labels, list):
+ labels_dict = {}
+ for entry in raw_labels:
+ if isinstance(entry, (list, tuple)) and len(entry) == 2:
+ label_obj, value = entry
+ if isinstance(label_obj, dict):
+ schema = label_obj.get("schema", "")
+ name = label_obj.get("name", "")
+ else:
+ schema, name = str(label_obj), ""
+ labels_dict.setdefault(schema, {})[name] = value
+ raw_labels = labels_dict
+
+ record = {
+ "instance_id": instance_id,
+ "user_id": user_id,
+ "labels": raw_labels,
+ "spans": {},
+ "links": {},
+ "image_annotations": {},
+ }
+
+ # Process span data
+ instance_spans = span_data.get(instance_id, {})
+ for schema_name, span_list in instance_spans.items():
+ if isinstance(span_list, list):
+ record["spans"][schema_name] = span_list
+ elif isinstance(span_list, dict):
+ # Span data might be stored as a dict of span_id -> span_obj
+ record["spans"][schema_name] = list(span_list.values())
+
+ # Extract image annotations from labels
+ # Image annotations are stored as JSON strings in label values
+ for schema_name, label_dict in record["labels"].items():
+ schema_config = _find_schema(schemas, schema_name)
+ if schema_config and schema_config.get("annotation_type") == "image_annotation":
+ # Image annotation data is stored in the label value
+ for label_key, value in label_dict.items():
+ if isinstance(value, str):
+ try:
+ parsed = json.loads(value)
+ if isinstance(parsed, list):
+ record["image_annotations"][schema_name] = parsed
+ except (json.JSONDecodeError, TypeError):
+ pass
+ elif isinstance(value, list):
+ record["image_annotations"][schema_name] = value
+
+ annotations.append(record)
+
+ return annotations
+
+
+def load_phase_responses_from_output_dir(output_dir: str) -> list:
+ """
+ Load phase/surveyflow responses from the Potato output directory.
+
+ Reads phase_to_page_to_label_to_value from each user's user_state.json
+ and flattens into a list of records.
+
+ Returns:
+ List of dicts with keys: user_id, phase, page, schema, label_name, value
+ """
+ responses = []
+
+ if not os.path.isdir(output_dir):
+ return responses
+
+ for user_dir in sorted(os.listdir(output_dir)):
+ user_path = os.path.join(output_dir, user_dir)
+ if not os.path.isdir(user_path):
+ continue
+
+ state_file = os.path.join(user_path, "user_state.json")
+ if not os.path.exists(state_file):
+ continue
+
+ with open(state_file, "r") as f:
+ user_state = json.load(f)
+
+ user_id = user_state.get("user_id", user_dir)
+ phase_data = user_state.get("phase_to_page_to_label_to_value", {})
+
+ for phase, pages in phase_data.items():
+ for page, label_values in pages.items():
+ # label_values is a list of [[{schema, name}, value], ...]
+ if isinstance(label_values, list):
+ for entry in label_values:
+ if isinstance(entry, (list, tuple)) and len(entry) == 2:
+ label_obj, value = entry
+ if isinstance(label_obj, dict):
+ schema = label_obj.get("schema", "")
+ label_name = label_obj.get("name", "")
+ else:
+ schema, label_name = str(label_obj), ""
+ responses.append({
+ "user_id": user_id,
+ "phase": phase,
+ "page": page,
+ "schema": schema,
+ "label_name": label_name,
+ "value": value,
+ })
+ elif isinstance(label_values, dict):
+ for label_obj, value in label_values.items():
+ responses.append({
+ "user_id": user_id,
+ "phase": phase,
+ "page": page,
+ "schema": str(label_obj),
+ "label_name": "",
+ "value": value,
+ })
+
+ return responses
+
+
+def load_items_from_data_files(config: dict, config_dir: str) -> dict:
+ """
+ Load item data from the data files specified in config.
+
+ Args:
+ config: Full Potato configuration dict
+ config_dir: Directory containing the config file
+
+ Returns:
+ Dict mapping instance_id -> item data
+ """
+ items = {}
+ item_props = config.get("item_properties", {})
+ id_key = item_props.get("id_key", "id")
+
+ data_files = config.get("data_files", [])
+ if isinstance(data_files, str):
+ data_files = [data_files]
+
+ task_dir = config.get("task_dir", ".")
+ base_dir = os.path.normpath(os.path.join(config_dir, task_dir))
+
+ for data_file_entry in data_files:
+ if isinstance(data_file_entry, dict):
+ path = data_file_entry.get("path", "")
+ else:
+ path = str(data_file_entry)
+
+ if not os.path.isabs(path):
+ path = os.path.join(base_dir, path)
+
+ if not os.path.exists(path):
+ logger.warning(f"Data file not found: {path}")
+ continue
+
+ with open(path, "r") as f:
+ for line_num, line in enumerate(f, 1):
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ item = json.loads(line)
+ item_id = str(item.get(id_key, f"item_{line_num}"))
+ items[item_id] = item
+ except json.JSONDecodeError:
+ # Try CSV/TSV
+ logger.debug(f"Line {line_num} in {path} is not JSON, skipping")
+
+ return items
+
+
+def _find_schema(schemas: list, name: str) -> dict:
+ """Find a schema config by name."""
+ for s in schemas:
+ if s.get("name") == name:
+ return s
+ return {}
+
+
+def build_export_context(config_path: str) -> ExportContext:
+ """
+ Build an ExportContext from a Potato config file.
+
+ Args:
+ config_path: Path to YAML config file
+
+ Returns:
+ ExportContext ready for export
+ """
+ config_path = os.path.abspath(config_path)
+ config_dir = os.path.dirname(config_path)
+
+ with open(config_path, "r") as f:
+ config = yaml.safe_load(f)
+
+ schemas = config.get("annotation_schemes", [])
+
+ # Determine output directory
+ task_dir = config.get("task_dir", ".")
+ base_dir = os.path.normpath(os.path.join(config_dir, task_dir))
+ output_annotation_dir = config.get(
+ "output_annotation_dir",
+ os.path.join(base_dir, "annotation_output")
+ )
+ if not os.path.isabs(output_annotation_dir):
+ output_annotation_dir = os.path.join(base_dir, output_annotation_dir)
+
+ items = load_items_from_data_files(config, config_dir)
+ annotations = load_annotations_from_output_dir(output_annotation_dir, schemas)
+ phase_responses = load_phase_responses_from_output_dir(output_annotation_dir)
+
+ return ExportContext(
+ config=config,
+ annotations=annotations,
+ items=items,
+ schemas=schemas,
+ output_dir=output_annotation_dir,
+ phase_responses=phase_responses,
+ )
+
+
+def main():
+ parser = argparse.ArgumentParser(
+ description="Export Potato annotations to standard formats"
+ )
+ parser.add_argument(
+ "--config", "-c",
+ help="Path to Potato YAML config file",
+ )
+ parser.add_argument(
+ "--format", "-f",
+ help="Export format (e.g., coco, yolo, pascal_voc, conll_2003, conll_u)",
+ )
+ parser.add_argument(
+ "--output", "-o",
+ help="Output directory",
+ default="./export_output",
+ )
+ parser.add_argument(
+ "--list-formats",
+ action="store_true",
+ help="List available export formats and exit",
+ )
+ parser.add_argument(
+ "--option",
+ action="append",
+ default=[],
+ help="Format-specific option as key=value (can be repeated)",
+ )
+ parser.add_argument(
+ "--verbose", "-v",
+ action="store_true",
+ help="Enable verbose logging",
+ )
+
+ args = parser.parse_args()
+
+ logging.basicConfig(
+ level=logging.DEBUG if args.verbose else logging.INFO,
+ format="%(levelname)s: %(message)s",
+ )
+
+ if args.list_formats:
+ formats = export_registry.list_exporters()
+ if not formats:
+ print("No export formats registered.")
+ else:
+ print("Available export formats:\n")
+ for fmt in formats:
+ exts = ", ".join(fmt["file_extensions"])
+ print(f" {fmt['format_name']:15s} {fmt['description']}")
+ print(f" {'':15s} Extensions: {exts}")
+ print()
+ return
+
+ if not args.config:
+ parser.error("--config is required (unless using --list-formats)")
+ if not args.format:
+ parser.error("--format is required (unless using --list-formats)")
+
+ if not os.path.exists(args.config):
+ print(f"Error: Config file not found: {args.config}", file=sys.stderr)
+ sys.exit(1)
+
+ # Parse options
+ options = {}
+ for opt in args.option:
+ if "=" in opt:
+ k, v = opt.split("=", 1)
+ options[k.strip()] = v.strip()
+
+ # Build context
+ print(f"Loading config from: {args.config}")
+ context = build_export_context(args.config)
+ print(f"Loaded {len(context.items)} items, {len(context.annotations)} annotations")
+
+ # Export
+ print(f"Exporting to {args.format} format...")
+ result = export_registry.export(args.format, context, args.output, options)
+
+ if result.success:
+ print(f"\nExport successful!")
+ print(f"Files written:")
+ for f in result.files_written:
+ print(f" {f}")
+ if result.stats:
+ print(f"\nStatistics:")
+ for k, v in result.stats.items():
+ print(f" {k}: {v}")
+ else:
+ print(f"\nExport failed!", file=sys.stderr)
+ for err in result.errors:
+ print(f" ERROR: {err}", file=sys.stderr)
+
+ if result.warnings:
+ print(f"\nWarnings:")
+ for w in result.warnings:
+ print(f" WARNING: {w}")
+
+ sys.exit(0 if result.success else 1)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/potato/export/coco_exporter.py b/potato/export/coco_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..706b8dd94e97882830c4365b1a1a5a743e42d4e9
--- /dev/null
+++ b/potato/export/coco_exporter.py
@@ -0,0 +1,184 @@
+"""
+COCO JSON Exporter
+
+Exports image annotations to COCO format with images[], annotations[],
+and categories[] arrays. Supports bbox, polygon/freeform segmentation.
+"""
+
+import json
+import os
+import logging
+from typing import Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+from .cv_utils import (
+ build_category_mapping,
+ polygon_to_bbox,
+ polygon_area,
+ flatten_polygon,
+ extract_image_annotations,
+ get_image_dimensions,
+ get_image_filename,
+ decode_rle,
+ rle_to_coco_rle,
+ rle_bbox,
+ rle_area,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class COCOExporter(BaseExporter):
+ format_name = "coco"
+ description = "COCO JSON format for object detection and segmentation"
+ file_extensions = [".json"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ has_image_schema = any(
+ s.get("annotation_type") == "image_annotation"
+ for s in context.schemas
+ )
+ if not has_image_schema:
+ return False, "No image_annotation schema found in config"
+ return True, ""
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ warnings = []
+ annotation_id_counter = 1
+
+ category_map = build_category_mapping(context.annotations, context.schemas)
+ # COCO uses 1-indexed category IDs
+ coco_categories = [
+ {"id": idx + 1, "name": name, "supercategory": ""}
+ for name, idx in sorted(category_map.items(), key=lambda kv: kv[1])
+ ]
+
+ coco_images = []
+ coco_annotations = []
+ image_id_map = {} # instance_id -> image_id
+ image_id_counter = 1
+
+ for ann in context.annotations:
+ instance_id = ann.get("instance_id", "")
+ item = context.items.get(instance_id, {})
+ img_anns = extract_image_annotations(ann)
+ if not img_anns:
+ continue
+
+ # Assign image ID (deduplicate by instance_id)
+ if instance_id not in image_id_map:
+ image_id = image_id_counter
+ image_id_counter += 1
+ image_id_map[instance_id] = image_id
+
+ width, height = get_image_dimensions(item)
+ file_name = get_image_filename(item) or instance_id
+
+ coco_images.append({
+ "id": image_id,
+ "file_name": file_name,
+ "width": width,
+ "height": height,
+ })
+ else:
+ image_id = image_id_map[instance_id]
+
+ for schema_name, objects in img_anns:
+ for obj in objects:
+ obj_type = obj.get("type", "")
+ label = obj.get("label", "")
+
+ if label not in category_map:
+ warnings.append(
+ f"Unknown label '{label}' in {instance_id}, skipping"
+ )
+ continue
+
+ cat_id = category_map[label] + 1 # 1-indexed for COCO
+
+ coco_ann = {
+ "id": annotation_id_counter,
+ "image_id": image_id,
+ "category_id": cat_id,
+ "iscrowd": 0,
+ }
+ annotation_id_counter += 1
+
+ if obj_type == "bbox":
+ x = obj.get("x", 0)
+ y = obj.get("y", 0)
+ w = obj.get("width", 0)
+ h = obj.get("height", 0)
+ coco_ann["bbox"] = [x, y, w, h]
+ coco_ann["area"] = w * h
+ coco_ann["segmentation"] = []
+
+ elif obj_type in ("polygon", "freeform"):
+ points = obj.get("points", [])
+ if not points:
+ warnings.append(
+ f"Empty points for {obj_type} in {instance_id}"
+ )
+ continue
+ flat = flatten_polygon(points)
+ coco_ann["segmentation"] = [flat]
+ bx, by, bw, bh = polygon_to_bbox(points)
+ coco_ann["bbox"] = [bx, by, bw, bh]
+ coco_ann["area"] = polygon_area(points)
+
+ elif obj_type == "mask":
+ rle = obj.get("rle", {})
+ if not rle.get("counts"):
+ warnings.append(
+ f"Empty RLE mask in {instance_id}"
+ )
+ continue
+ size = rle.get("size", [])
+ mask_h = size[0] if len(size) >= 2 else height
+ mask_w = size[1] if len(size) >= 2 else width
+ coco_rle = rle_to_coco_rle(rle, mask_w, mask_h)
+ decoded = decode_rle(rle, mask_w, mask_h)
+ coco_ann["segmentation"] = coco_rle
+ coco_ann["bbox"] = rle_bbox(decoded, mask_w, mask_h)
+ coco_ann["area"] = rle_area(decoded)
+ coco_ann["iscrowd"] = 1
+
+ elif obj_type == "landmark":
+ warnings.append(
+ f"Landmark annotation in {instance_id} skipped "
+ f"(not standard in COCO detection format)"
+ )
+ continue
+
+ else:
+ warnings.append(
+ f"Unknown annotation type '{obj_type}' in {instance_id}"
+ )
+ continue
+
+ coco_annotations.append(coco_ann)
+
+ coco_output = {
+ "images": coco_images,
+ "annotations": coco_annotations,
+ "categories": coco_categories,
+ }
+
+ os.makedirs(output_path, exist_ok=True)
+ out_file = os.path.join(output_path, "annotations.json")
+ with open(out_file, "w") as f:
+ json.dump(coco_output, f, indent=2)
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=[out_file],
+ warnings=warnings,
+ stats={
+ "num_images": len(coco_images),
+ "num_annotations": len(coco_annotations),
+ "num_categories": len(coco_categories),
+ },
+ )
diff --git a/potato/export/codebook_exporter.py b/potato/export/codebook_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..bb0ed52d44885ffafedccd09c30d66a63a592341
--- /dev/null
+++ b/potato/export/codebook_exporter.py
@@ -0,0 +1,157 @@
+"""
+Codebook Exporter
+
+Exports the project codebook (label/code taxonomy) to a CSV file with one row per
+code. Designed for qualitative-research workflows where the codebook is a
+deliverable in its own right.
+
+Output columns:
+ schema annotation_scheme name
+ annotation_type schema type (radio, multiselect, span, hierarchical_multiselect)
+ code label name
+ parent parent code (for hierarchical schemas)
+ description description / tooltip from the schema config
+ color color hex if defined
+ n_uses number of times this code was applied across all annotators
+"""
+
+import csv
+import logging
+import os
+from typing import Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+
+logger = logging.getLogger(__name__)
+
+
+# Schemas that contribute codes to a codebook export.
+CODEBOOK_SCHEMA_TYPES = {
+ "radio", "multiselect", "select", "likert",
+ "span", "hierarchical_multiselect", "tree_annotation",
+}
+
+
+class CodebookExporter(BaseExporter):
+ format_name = "codebook"
+ description = "Project codebook (CSV) with code names, hierarchy, and use counts"
+ file_extensions = [".csv"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ has_codeable_schema = any(
+ s.get("annotation_type") in CODEBOOK_SCHEMA_TYPES
+ for s in context.schemas
+ )
+ if not has_codeable_schema:
+ return False, "No codeable schema (radio/multiselect/span/etc.) in config"
+ return True, ""
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ os.makedirs(output_path, exist_ok=True)
+ out_file = os.path.join(output_path, "codebook.csv")
+
+ use_counts = self._count_label_uses(context)
+
+ rows = []
+ for scheme in context.schemas:
+ atype = scheme.get("annotation_type")
+ if atype not in CODEBOOK_SCHEMA_TYPES:
+ continue
+ schema_name = scheme.get("name", "")
+ for code_row in self._iter_codes(scheme):
+ code_row["schema"] = schema_name
+ code_row["annotation_type"] = atype
+ code_row["n_uses"] = use_counts.get(
+ (schema_name, code_row["code"]), 0
+ )
+ rows.append(code_row)
+
+ fieldnames = [
+ "schema", "annotation_type", "code", "parent",
+ "description", "color", "n_uses",
+ ]
+ with open(out_file, "w", newline="", encoding="utf-8") as f:
+ writer = csv.DictWriter(f, fieldnames=fieldnames)
+ writer.writeheader()
+ for r in rows:
+ writer.writerow({k: r.get(k, "") for k in fieldnames})
+
+ logger.info(f"Codebook exported to {out_file}: {len(rows)} codes")
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=[out_file],
+ stats={"codes_exported": len(rows)},
+ )
+
+ @staticmethod
+ def _iter_codes(scheme):
+ """Yield {code, parent, description, color} dicts for a schema."""
+ atype = scheme.get("annotation_type")
+
+ if atype == "hierarchical_multiselect":
+ yield from CodebookExporter._iter_hierarchical(scheme.get("labels", []), parent="")
+ return
+ if atype == "tree_annotation":
+ yield from CodebookExporter._iter_hierarchical(scheme.get("labels", []), parent="")
+ return
+
+ labels = scheme.get("labels", [])
+ for label in labels:
+ if isinstance(label, dict):
+ name = label.get("name", "")
+ yield {
+ "code": name,
+ "parent": "",
+ "description": label.get("description") or label.get("tooltip", ""),
+ "color": label.get("color", ""),
+ }
+ else:
+ yield {"code": str(label), "parent": "", "description": "", "color": ""}
+
+ @staticmethod
+ def _iter_hierarchical(nodes, parent):
+ if not isinstance(nodes, list):
+ return
+ for node in nodes:
+ if isinstance(node, dict):
+ name = node.get("name", "")
+ yield {
+ "code": name,
+ "parent": parent,
+ "description": node.get("description") or node.get("tooltip", ""),
+ "color": node.get("color", ""),
+ }
+ children = node.get("children") or node.get("labels") or []
+ yield from CodebookExporter._iter_hierarchical(children, parent=name)
+ else:
+ yield {"code": str(node), "parent": parent, "description": "", "color": ""}
+
+ @staticmethod
+ def _count_label_uses(context):
+ counts = {}
+ for ann in context.annotations:
+ labels = ann.get("labels", {}) or {}
+ for schema_name, schema_payload in labels.items():
+ names = []
+ if isinstance(schema_payload, dict):
+ names = [k for k, v in schema_payload.items() if v]
+ elif isinstance(schema_payload, list):
+ names = [str(x) for x in schema_payload]
+ elif schema_payload not in (None, ""):
+ names = [str(schema_payload)]
+ for n in names:
+ key = (schema_name, n)
+ counts[key] = counts.get(key, 0) + 1
+
+ spans = ann.get("spans", {}) or {}
+ for schema_name, span_list in spans.items():
+ for span in span_list or []:
+ label = span.get("label") or span.get("annotation")
+ if label:
+ key = (schema_name, label)
+ counts[key] = counts.get(key, 0) + 1
+
+ return counts
diff --git a/potato/export/coding_eval_exporter.py b/potato/export/coding_eval_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..f08d7e223a970d6ed3f58aa1d61650abca62bea6
--- /dev/null
+++ b/potato/export/coding_eval_exporter.py
@@ -0,0 +1,252 @@
+"""
+Coding Agent Evaluation Exporter
+
+Exports coding agent annotations in formats useful for training:
+- PRM (Process Reward Model): per-step reward signals
+- DPO/RLHF preference format: chosen/rejected trace pairs
+- SWE-bench compatible evaluation results
+- Code review format: structured review data
+"""
+
+import json
+import os
+import logging
+from typing import Dict, List, Any, Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+
+logger = logging.getLogger(__name__)
+
+
+class CodingEvalExporter(BaseExporter):
+ """Export coding agent annotations for ML training pipelines."""
+
+ format_name = "coding_eval"
+ description = "Coding agent evaluation data (PRM, DPO, SWE-bench, code review)"
+ file_extensions = [".jsonl", ".json"]
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ export_types = options.get("types", ["prm", "preference", "swebench", "code_review"])
+ files_written = []
+ warnings = []
+ stats = {}
+
+ os.makedirs(output_path, exist_ok=True)
+
+ if "prm" in export_types:
+ path, count = self._export_prm(context, output_path)
+ if path:
+ files_written.append(path)
+ stats["prm_instances"] = count
+
+ if "preference" in export_types:
+ path, count = self._export_preference(context, output_path)
+ if path:
+ files_written.append(path)
+ stats["preference_pairs"] = count
+
+ if "swebench" in export_types:
+ path, count = self._export_swebench(context, output_path)
+ if path:
+ files_written.append(path)
+ stats["swebench_results"] = count
+
+ if "code_review" in export_types:
+ path, count = self._export_code_review(context, output_path)
+ if path:
+ files_written.append(path)
+ stats["code_reviews"] = count
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=files_written,
+ warnings=warnings,
+ stats=stats,
+ )
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ if not context.annotations:
+ return False, "No annotations to export"
+
+ # Check for relevant schema types
+ schema_types = {s.get("annotation_type") for s in context.schemas}
+ relevant = schema_types & {"process_reward", "code_review", "pairwise", "radio"}
+ if not relevant:
+ return False, "No coding evaluation schemas found (process_reward, code_review, pairwise, radio)"
+
+ return True, ""
+
+ def _export_prm(self, context: ExportContext, output_dir: str) -> Tuple[Optional[str], int]:
+ """Export PRM training data."""
+ output_path = os.path.join(output_dir, "prm_training_data.jsonl")
+ count = 0
+
+ with open(output_path, "w") as f:
+ for ann in context.annotations:
+ instance_id = ann.get("instance_id", "")
+ labels = ann.get("labels", {})
+
+ for schema_name, value in labels.items():
+ if not isinstance(value, dict):
+ continue
+ label_val = value.get("label", "")
+ if not isinstance(label_val, str):
+ continue
+
+ try:
+ parsed = json.loads(label_val)
+ except (json.JSONDecodeError, TypeError):
+ continue
+
+ if not isinstance(parsed, dict) or "steps" not in parsed:
+ continue
+
+ steps = parsed["steps"]
+ if not isinstance(steps, list):
+ continue
+
+ record = {
+ "instance_id": instance_id,
+ "annotator": ann.get("user_id", ""),
+ "steps": [
+ {"index": s.get("index", i), "reward": s.get("reward", 0)}
+ for i, s in enumerate(steps)
+ ],
+ }
+ if "mode" in parsed:
+ record["mode"] = parsed["mode"]
+
+ f.write(json.dumps(record, ensure_ascii=False) + "\n")
+ count += 1
+
+ if count == 0:
+ os.remove(output_path)
+ return None, 0
+
+ logger.info(f"Exported {count} PRM records to {output_path}")
+ return output_path, count
+
+ def _export_preference(self, context: ExportContext, output_dir: str) -> Tuple[Optional[str], int]:
+ """Export DPO/RLHF preference pairs from pairwise annotations."""
+ output_path = os.path.join(output_dir, "preference_pairs.jsonl")
+ count = 0
+
+ with open(output_path, "w") as f:
+ for ann in context.annotations:
+ instance_id = ann.get("instance_id", "")
+ labels = ann.get("labels", {})
+
+ for schema_name, value in labels.items():
+ if not isinstance(value, dict):
+ continue
+
+ label_val = value.get("label", "")
+ # Pairwise annotations store "A" or "B"
+ if label_val not in ("A", "B", "a", "b"):
+ continue
+
+ # Get the instance data to extract prompt
+ item_data = context.items.get(instance_id, {})
+ prompt = item_data.get("task_description", item_data.get("text", ""))
+
+ record = {
+ "instance_id": instance_id,
+ "prompt": prompt,
+ "chosen": label_val.upper(),
+ "annotator": ann.get("user_id", ""),
+ }
+ f.write(json.dumps(record, ensure_ascii=False) + "\n")
+ count += 1
+
+ if count == 0:
+ os.remove(output_path)
+ return None, 0
+
+ logger.info(f"Exported {count} preference pairs to {output_path}")
+ return output_path, count
+
+ def _export_swebench(self, context: ExportContext, output_dir: str) -> Tuple[Optional[str], int]:
+ """Export SWE-bench compatible evaluation results."""
+ output_path = os.path.join(output_dir, "swebench_results.jsonl")
+ count = 0
+
+ with open(output_path, "w") as f:
+ for ann in context.annotations:
+ instance_id = ann.get("instance_id", "")
+ labels = ann.get("labels", {})
+
+ # Look for task_success or similar radio annotation
+ resolved = None
+ for schema_name, value in labels.items():
+ if not isinstance(value, dict):
+ continue
+ label_val = value.get("label", "")
+ if label_val in ("success", "resolved", "correct"):
+ resolved = True
+ elif label_val in ("failure", "unresolved", "incorrect"):
+ resolved = False
+ elif label_val in ("partial", "partially_resolved"):
+ resolved = False # SWE-bench is binary
+
+ if resolved is not None:
+ record = {
+ "instance_id": instance_id,
+ "resolved": resolved,
+ "annotator": ann.get("user_id", ""),
+ }
+ f.write(json.dumps(record, ensure_ascii=False) + "\n")
+ count += 1
+
+ if count == 0:
+ os.remove(output_path)
+ return None, 0
+
+ logger.info(f"Exported {count} SWE-bench results to {output_path}")
+ return output_path, count
+
+ def _export_code_review(self, context: ExportContext, output_dir: str) -> Tuple[Optional[str], int]:
+ """Export structured code review data."""
+ output_path = os.path.join(output_dir, "code_reviews.jsonl")
+ count = 0
+
+ with open(output_path, "w") as f:
+ for ann in context.annotations:
+ instance_id = ann.get("instance_id", "")
+ labels = ann.get("labels", {})
+
+ for schema_name, value in labels.items():
+ if not isinstance(value, dict):
+ continue
+ label_val = value.get("label", "")
+ if not isinstance(label_val, str):
+ continue
+
+ try:
+ parsed = json.loads(label_val)
+ except (json.JSONDecodeError, TypeError):
+ continue
+
+ if not isinstance(parsed, dict):
+ continue
+
+ # Check for code review structure
+ if "verdict" in parsed or "comments" in parsed:
+ record = {
+ "instance_id": instance_id,
+ "annotator": ann.get("user_id", ""),
+ "verdict": parsed.get("verdict", ""),
+ "comments": parsed.get("comments", []),
+ "file_ratings": parsed.get("file_ratings", {}),
+ }
+ f.write(json.dumps(record, ensure_ascii=False) + "\n")
+ count += 1
+
+ if count == 0:
+ os.remove(output_path)
+ return None, 0
+
+ logger.info(f"Exported {count} code reviews to {output_path}")
+ return output_path, count
diff --git a/potato/export/conll_2003_exporter.py b/potato/export/conll_2003_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..b6ef534b1d7f8bbedfc7562a425c78f24c23afdd
--- /dev/null
+++ b/potato/export/conll_2003_exporter.py
@@ -0,0 +1,135 @@
+"""
+CoNLL-2003 Exporter
+
+Exports span annotations to CoNLL-2003 format:
+- Tab-separated columns: WORD POS CHUNK NER
+- Blank lines between sentences
+- -DOCSTART- markers between documents
+"""
+
+import os
+import logging
+from typing import Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+from .nlp_utils import tokenize_text, char_spans_to_bio_tags, group_sentences
+
+logger = logging.getLogger(__name__)
+
+
+class CoNLL2003Exporter(BaseExporter):
+ format_name = "conll_2003"
+ description = "CoNLL-2003 NER format (WORD POS CHUNK NER)"
+ file_extensions = [".conll", ".txt"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ has_span_schema = any(
+ s.get("annotation_type") == "span"
+ for s in context.schemas
+ )
+ if not has_span_schema:
+ return False, "No span annotation schema found in config"
+ return True, ""
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ warnings = []
+
+ tokenization = options.get("tokenization", "whitespace")
+ pos_column = options.get("pos_column", "_")
+ chunk_column = options.get("chunk_column", "_")
+ # Which span schema to export (defaults to first span schema)
+ schema_name = options.get("schema_name")
+ if not schema_name:
+ for s in context.schemas:
+ if s.get("annotation_type") == "span":
+ schema_name = s.get("name")
+ break
+
+ os.makedirs(output_path, exist_ok=True)
+ out_file = os.path.join(output_path, "annotations.conll")
+
+ lines = []
+ total_tokens = 0
+ total_entities = 0
+
+ # Get text key from config
+ item_props = context.config.get("item_properties", {})
+ text_key = item_props.get("text_key", "text")
+
+ # Group annotations by instance to handle multiple annotators
+ instance_annotations = {}
+ for ann in context.annotations:
+ iid = ann.get("instance_id", "")
+ if iid not in instance_annotations:
+ instance_annotations[iid] = ann
+ # If multiple annotators, use first one (could be configurable)
+
+ for instance_id, ann in instance_annotations.items():
+ item = context.items.get(instance_id, {})
+ text = item.get(text_key, "")
+ if not text:
+ # Try alternative text fields
+ for alt_key in ("text", "sentence", "content"):
+ if alt_key in item:
+ text = item[alt_key]
+ break
+
+ if not text:
+ warnings.append(f"No text found for {instance_id}")
+ continue
+
+ # Handle text that's a list
+ if isinstance(text, list):
+ text = " ".join(str(t) for t in text)
+
+ # Tokenize
+ tokens = tokenize_text(text, method=tokenization)
+ if not tokens:
+ continue
+
+ # Get spans for this instance
+ spans = []
+ for span_schema, span_list in ann.get("spans", {}).items():
+ if schema_name and span_schema != schema_name:
+ continue
+ for sp in span_list:
+ spans.append({
+ "start": sp.get("start", 0),
+ "end": sp.get("end", 0),
+ "label": sp.get("name") or sp.get("label", "ENTITY"),
+ })
+
+ bio_tags = char_spans_to_bio_tags(tokens, spans)
+ total_tokens += len(tokens)
+ total_entities += sum(1 for t in bio_tags if t.startswith("B-"))
+
+ # Doc separator
+ lines.append("-DOCSTART- -X- -X- O")
+ lines.append("")
+
+ # Group into sentences
+ sentences = group_sentences(tokens, text)
+
+ for sentence_indices in sentences:
+ for idx in sentence_indices:
+ tok = tokens[idx]
+ tag = bio_tags[idx]
+ lines.append(f"{tok['token']}\t{pos_column}\t{chunk_column}\t{tag}")
+ lines.append("") # Blank line between sentences
+
+ with open(out_file, "w") as f:
+ f.write("\n".join(lines))
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=[out_file],
+ warnings=warnings,
+ stats={
+ "num_documents": len(instance_annotations),
+ "num_tokens": total_tokens,
+ "num_entities": total_entities,
+ },
+ )
diff --git a/potato/export/conll_u_exporter.py b/potato/export/conll_u_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..2078410f2238c788f466a0084c1526a1aee1b1a2
--- /dev/null
+++ b/potato/export/conll_u_exporter.py
@@ -0,0 +1,170 @@
+"""
+CoNLL-U Exporter
+
+Exports span annotations to CoNLL-U format:
+- 10 columns: ID FORM LEMMA UPOS XPOS FEATS HEAD DEPREL DEPS MISC
+- NER annotations placed in MISC column as SpaceAfter/NER features
+- Blank lines between sentences
+- Comment lines with sent_id and text
+"""
+
+import os
+import logging
+from typing import Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+from .nlp_utils import tokenize_text, char_spans_to_bio_tags, group_sentences
+
+logger = logging.getLogger(__name__)
+
+
+class CoNLLUExporter(BaseExporter):
+ format_name = "conll_u"
+ description = "CoNLL-U format (Universal Dependencies compatible, NER in MISC)"
+ file_extensions = [".conllu"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ has_span_schema = any(
+ s.get("annotation_type") == "span"
+ for s in context.schemas
+ )
+ if not has_span_schema:
+ return False, "No span annotation schema found in config"
+ return True, ""
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ warnings = []
+
+ tokenization = options.get("tokenization", "whitespace")
+ schema_name = options.get("schema_name")
+ if not schema_name:
+ for s in context.schemas:
+ if s.get("annotation_type") == "span":
+ schema_name = s.get("name")
+ break
+
+ os.makedirs(output_path, exist_ok=True)
+ out_file = os.path.join(output_path, "annotations.conllu")
+
+ lines = []
+ total_tokens = 0
+ total_entities = 0
+ sent_counter = 0
+
+ item_props = context.config.get("item_properties", {})
+ text_key = item_props.get("text_key", "text")
+
+ # Deduplicate by instance
+ instance_annotations = {}
+ for ann in context.annotations:
+ iid = ann.get("instance_id", "")
+ if iid not in instance_annotations:
+ instance_annotations[iid] = ann
+
+ for instance_id, ann in instance_annotations.items():
+ item = context.items.get(instance_id, {})
+ text = item.get(text_key, "")
+ if not text:
+ for alt_key in ("text", "sentence", "content"):
+ if alt_key in item:
+ text = item[alt_key]
+ break
+
+ if not text:
+ warnings.append(f"No text found for {instance_id}")
+ continue
+
+ if isinstance(text, list):
+ text = " ".join(str(t) for t in text)
+
+ tokens = tokenize_text(text, method=tokenization)
+ if not tokens:
+ continue
+
+ # Get spans
+ spans = []
+ for span_schema, span_list in ann.get("spans", {}).items():
+ if schema_name and span_schema != schema_name:
+ continue
+ for sp in span_list:
+ spans.append({
+ "start": sp.get("start", 0),
+ "end": sp.get("end", 0),
+ "label": sp.get("name") or sp.get("label", "ENTITY"),
+ })
+
+ bio_tags = char_spans_to_bio_tags(tokens, spans)
+ total_tokens += len(tokens)
+ total_entities += sum(1 for t in bio_tags if t.startswith("B-"))
+
+ sentences = group_sentences(tokens, text)
+
+ for sentence_indices in sentences:
+ sent_counter += 1
+ sent_tokens = [tokens[i] for i in sentence_indices]
+ sent_tags = [bio_tags[i] for i in sentence_indices]
+
+ # Reconstruct sentence text
+ if sent_tokens:
+ s_start = sent_tokens[0]["start"]
+ s_end = sent_tokens[-1]["end"]
+ sent_text = text[s_start:s_end]
+ else:
+ sent_text = ""
+
+ lines.append(f"# sent_id = {instance_id}-s{sent_counter}")
+ lines.append(f"# text = {sent_text}")
+
+ for tok_num, (tok, ner_tag) in enumerate(
+ zip(sent_tokens, sent_tags), start=1
+ ):
+ # Build MISC field
+ misc_parts = []
+
+ # SpaceAfter=No if no space before next token
+ if tok_num < len(sent_tokens):
+ next_tok = sent_tokens[tok_num] # 0-indexed next
+ if tok["end"] == next_tok["start"]:
+ misc_parts.append("SpaceAfter=No")
+
+ # NER tag
+ if ner_tag != "O":
+ misc_parts.append(f"NER={ner_tag}")
+
+ misc = "|".join(misc_parts) if misc_parts else "_"
+
+ # 10-column CoNLL-U format
+ # ID FORM LEMMA UPOS XPOS FEATS HEAD DEPREL DEPS MISC
+ cols = [
+ str(tok_num), # ID
+ tok["token"], # FORM
+ "_", # LEMMA
+ "_", # UPOS
+ "_", # XPOS
+ "_", # FEATS
+ "_", # HEAD
+ "_", # DEPREL
+ "_", # DEPS
+ misc, # MISC
+ ]
+ lines.append("\t".join(cols))
+
+ lines.append("") # Blank line between sentences
+
+ with open(out_file, "w") as f:
+ f.write("\n".join(lines))
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=[out_file],
+ warnings=warnings,
+ stats={
+ "num_documents": len(instance_annotations),
+ "num_sentences": sent_counter,
+ "num_tokens": total_tokens,
+ "num_entities": total_entities,
+ },
+ )
diff --git a/potato/export/cv_utils.py b/potato/export/cv_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..88bca65e110b032705f3459bb5ce1e53ec0f3a74
--- /dev/null
+++ b/potato/export/cv_utils.py
@@ -0,0 +1,382 @@
+"""
+CV Export Utilities
+
+Shared helper functions for computer vision export formats (COCO, YOLO, VOC).
+"""
+
+from typing import Dict, List, Tuple, Any, Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def build_category_mapping(annotations: List[dict], schemas: List[dict]) -> Dict[str, int]:
+ """
+ Build a mapping from label names to integer category IDs.
+
+ Extracts labels from image_annotation schemas first (preserving config order),
+ then discovers any additional labels from annotations.
+
+ Args:
+ annotations: List of annotation records
+ schemas: List of annotation_scheme config dicts
+
+ Returns:
+ Dict mapping label name -> integer ID (starting from 1 for COCO, 0-indexed for YOLO)
+ """
+ labels = []
+ seen = set()
+
+ # First, collect labels from schema configs (preserves defined order)
+ for schema in schemas:
+ if schema.get("annotation_type") == "image_annotation":
+ for label_def in schema.get("labels", []):
+ name = label_def if isinstance(label_def, str) else label_def.get("name", "")
+ if name and name not in seen:
+ labels.append(name)
+ seen.add(name)
+
+ # Then discover any labels in annotation data not already in config
+ for ann in annotations:
+ for schema_name, img_annotations in ann.get("image_annotations", {}).items():
+ if not isinstance(img_annotations, list):
+ continue
+ for obj in img_annotations:
+ label = obj.get("label", "")
+ if label and label not in seen:
+ labels.append(label)
+ seen.add(label)
+
+ return {name: idx for idx, name in enumerate(labels)}
+
+
+def polygon_to_bbox(points: List[List[float]]) -> Tuple[float, float, float, float]:
+ """
+ Compute axis-aligned bounding box from a polygon.
+
+ Args:
+ points: List of [x, y] coordinate pairs
+
+ Returns:
+ Tuple of (x_min, y_min, width, height)
+ """
+ if not points:
+ return (0, 0, 0, 0)
+
+ xs = [p[0] for p in points]
+ ys = [p[1] for p in points]
+ x_min = min(xs)
+ y_min = min(ys)
+ return (x_min, y_min, max(xs) - x_min, max(ys) - y_min)
+
+
+def polygon_area(points: List[List[float]]) -> float:
+ """
+ Compute the area of a polygon using the shoelace formula.
+
+ Args:
+ points: List of [x, y] coordinate pairs
+
+ Returns:
+ Absolute area of the polygon
+ """
+ n = len(points)
+ if n < 3:
+ return 0.0
+ area = 0.0
+ for i in range(n):
+ j = (i + 1) % n
+ area += points[i][0] * points[j][1]
+ area -= points[j][0] * points[i][1]
+ return abs(area) / 2.0
+
+
+def normalize_bbox(x: float, y: float, w: float, h: float,
+ img_w: float, img_h: float) -> Tuple[float, float, float, float]:
+ """
+ Normalize bounding box coordinates to [0, 1] range.
+
+ Args:
+ x, y: Top-left corner coordinates
+ w, h: Width and height
+ img_w, img_h: Image dimensions
+
+ Returns:
+ Tuple of (center_x, center_y, width, height) normalized to [0, 1]
+ """
+ if img_w <= 0 or img_h <= 0:
+ return (0, 0, 0, 0)
+ cx = max(0.0, min(1.0, (x + w / 2) / img_w))
+ cy = max(0.0, min(1.0, (y + h / 2) / img_h))
+ nw = max(0.0, min(1.0, w / img_w))
+ nh = max(0.0, min(1.0, h / img_h))
+ return (cx, cy, nw, nh)
+
+
+def flatten_polygon(points: List[List[float]]) -> List[float]:
+ """
+ Flatten a list of [x, y] points into a flat coordinate list [x1, y1, x2, y2, ...].
+
+ This is the format used by COCO segmentation.
+
+ Args:
+ points: List of [x, y] coordinate pairs
+
+ Returns:
+ Flat list of coordinates
+ """
+ result = []
+ for p in points:
+ result.extend(p[:2])
+ return result
+
+
+def extract_image_annotations(annotation: dict) -> List[Tuple[str, List[dict]]]:
+ """
+ Extract image annotation objects from an annotation record.
+
+ Args:
+ annotation: Single annotation record with image_annotations field
+
+ Returns:
+ List of (schema_name, annotation_objects) tuples
+ """
+ results = []
+ for schema_name, objects in annotation.get("image_annotations", {}).items():
+ if isinstance(objects, list) and objects:
+ results.append((schema_name, objects))
+ return results
+
+
+def get_image_dimensions(item: dict, default_width: int = 0,
+ default_height: int = 0) -> Tuple[int, int]:
+ """
+ Extract image dimensions from item metadata.
+
+ Checks common field names for image width/height.
+
+ Args:
+ item: Item data dict
+ default_width: Fallback width
+ default_height: Fallback height
+
+ Returns:
+ Tuple of (width, height)
+ """
+ # Check common field patterns
+ width = default_width
+ for w_key in ("image_width", "width", "img_width", "w"):
+ if w_key in item:
+ try:
+ width = int(item[w_key])
+ except (ValueError, TypeError):
+ pass
+ break
+
+ height = default_height
+ for h_key in ("image_height", "height", "img_height", "h"):
+ if h_key in item:
+ try:
+ height = int(item[h_key])
+ except (ValueError, TypeError):
+ pass
+ break
+
+ return (width, height)
+
+
+def get_image_filename(item: dict) -> Optional[str]:
+ """
+ Extract image filename from item data.
+
+ Args:
+ item: Item data dict
+
+ Returns:
+ Image filename/path string or None
+ """
+ for key in ("image", "image_path", "image_url", "file_name", "filename", "img"):
+ if key in item and item[key]:
+ return str(item[key])
+ return None
+
+
+# ---------------------------------------------------------------------------
+# RLE mask utilities (Potato RLE <-> COCO RLE conversion)
+# ---------------------------------------------------------------------------
+
+
+def decode_rle(rle: dict, width: int, height: int) -> List[int]:
+ """
+ Decode Potato RLE-encoded mask to a flat binary array (row-major order).
+
+ Potato RLE stores counts alternating between 0-pixels and 1-pixels,
+ starting with 0s, in row-major (left-to-right, top-to-bottom) order.
+
+ Args:
+ rle: Dict with 'counts' (list of ints) and 'size' [height, width]
+ width: Image width
+ height: Image height
+
+ Returns:
+ Flat list of 0/1 values in row-major order
+ """
+ counts = rle.get("counts", [])
+ total = width * height
+ mask = [0] * total
+ pos = 0
+ val = 0
+ for count in counts:
+ for _ in range(count):
+ if pos < total:
+ mask[pos] = val
+ pos += 1
+ val = 1 - val
+ return mask
+
+
+def rle_bbox(mask: List[int], width: int, height: int) -> List[float]:
+ """
+ Compute axis-aligned bounding box [x, y, w, h] from a flat binary mask.
+
+ Args:
+ mask: Flat list of 0/1 values (row-major)
+ width: Image width
+ height: Image height
+
+ Returns:
+ [x_min, y_min, bbox_width, bbox_height] or [0, 0, 0, 0] if empty
+ """
+ x_min, y_min = width, height
+ x_max, y_max = -1, -1
+ for i, val in enumerate(mask):
+ if val:
+ y = i // width
+ x = i % width
+ if x < x_min:
+ x_min = x
+ if x > x_max:
+ x_max = x
+ if y < y_min:
+ y_min = y
+ if y > y_max:
+ y_max = y
+ if x_max < 0:
+ return [0, 0, 0, 0]
+ return [float(x_min), float(y_min),
+ float(x_max - x_min + 1), float(y_max - y_min + 1)]
+
+
+def rle_area(mask: List[int]) -> int:
+ """
+ Compute mask area as the count of foreground pixels.
+
+ Args:
+ mask: Flat list of 0/1 values
+
+ Returns:
+ Number of 1-pixels
+ """
+ return sum(mask)
+
+
+def _column_major_rle_counts(mask_2d: List[List[int]], height: int,
+ width: int) -> List[int]:
+ """
+ Read a 2D mask in column-major order and compute RLE counts.
+
+ Counts alternate between 0-pixels and 1-pixels, starting with 0s.
+
+ Args:
+ mask_2d: 2D list [height][width] of 0/1 values
+ height: Image height
+ width: Image width
+
+ Returns:
+ List of integer run counts in column-major order
+ """
+ counts: List[int] = []
+ current_val = 0
+ current_run = 0
+
+ for x in range(width):
+ for y in range(height):
+ pixel = mask_2d[y][x]
+ if pixel == current_val:
+ current_run += 1
+ else:
+ counts.append(current_run)
+ current_val = 1 - current_val
+ current_run = 1
+ counts.append(current_run)
+ return counts
+
+
+def _encode_coco_rle_string(counts: List[int]) -> str:
+ """
+ Encode RLE integer counts as a COCO compressed ASCII string.
+
+ Implements the exact algorithm from pycocotools maskApi.c rleToString():
+ - Delta encoding for i > 2: x = counts[i] - counts[i-2]
+ - Each value encoded as 6-bit groups (5 data bits + 1 continuation bit)
+ - Each group offset by 48 to produce printable ASCII
+ - Signed values supported via arithmetic right shift
+
+ Args:
+ counts: List of integer run counts
+
+ Returns:
+ Encoded ASCII string
+ """
+ chars = []
+ for i, cnt in enumerate(counts):
+ # Delta encoding: for i > 2, encode difference from counts[i-2]
+ x = cnt - counts[i - 2] if i > 2 else cnt
+ while True:
+ c = x & 0x1F
+ x >>= 5
+ # If bit 4 set, sign bit is 1 โ more groups unless x is all-ones (-1)
+ # If bit 4 clear, sign bit is 0 โ more groups unless x is all-zeros (0)
+ if c & 0x10:
+ more = (x != -1)
+ else:
+ more = (x != 0)
+ if more:
+ c |= 0x20
+ chars.append(chr(c + 48))
+ if not more:
+ break
+ return "".join(chars)
+
+
+def rle_to_coco_rle(rle: dict, width: int, height: int) -> Dict[str, Any]:
+ """
+ Convert Potato RLE to COCO RLE format.
+
+ Potato RLE is row-major; COCO RLE is column-major with compressed
+ ASCII string encoding.
+
+ Args:
+ rle: Potato RLE dict with 'counts' and 'size'
+ width: Image width
+ height: Image height
+
+ Returns:
+ COCO RLE dict {"counts": "encoded_string", "size": [height, width]}
+ """
+ # Decode to flat row-major mask
+ flat = decode_rle(rle, width, height)
+
+ # Reshape to 2D
+ mask_2d = []
+ for y in range(height):
+ row = flat[y * width:(y + 1) * width]
+ mask_2d.append(row)
+
+ # Compute column-major RLE counts
+ col_counts = _column_major_rle_counts(mask_2d, height, width)
+
+ # Encode as COCO compressed string
+ encoded = _encode_coco_rle_string(col_counts)
+
+ return {"counts": encoded, "size": [height, width]}
diff --git a/potato/export/eaf_exporter.py b/potato/export/eaf_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..d864e751f8acb3669cbf61b9dd417e36f09c0d59
--- /dev/null
+++ b/potato/export/eaf_exporter.py
@@ -0,0 +1,432 @@
+"""
+EAF Exporter
+
+Exports tiered annotations to ELAN Annotation Format (EAF) XML files.
+EAF is the native format for ELAN (https://archive.mpi.nl/tla/elan),
+a tool widely used for linguistic annotation of audio/video data.
+
+The EAF format supports:
+- Time-aligned annotations with millisecond precision
+- Hierarchical tier structures with parent-child relationships
+- Multiple linguistic types with different constraints
+- Media references (audio/video files)
+"""
+
+import logging
+import os
+import xml.etree.ElementTree as ET
+from datetime import datetime
+from typing import Dict, List, Any, Optional, Tuple, Set
+
+from .base import BaseExporter, ExportContext, ExportResult
+
+logger = logging.getLogger(__name__)
+
+
+class EAFExporter(BaseExporter):
+ """
+ Exports tiered annotations to ELAN Annotation Format (EAF).
+
+ This exporter creates valid EAF 3.0 XML files that can be opened
+ directly in ELAN for review or further annotation.
+ """
+
+ format_name = "eaf"
+ description = "ELAN Annotation Format (EAF) for linguistic annotation"
+ file_extensions = [".eaf"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ """
+ Check if the context contains tiered_annotation schema.
+
+ Args:
+ context: ExportContext to validate
+
+ Returns:
+ Tuple of (can_export, reason)
+ """
+ for schema in context.schemas:
+ if schema.get("annotation_type") == "tiered_annotation":
+ return True, ""
+
+ return False, "No tiered_annotation schema found in configuration"
+
+ def export(
+ self,
+ context: ExportContext,
+ output_path: str,
+ options: Optional[dict] = None
+ ) -> ExportResult:
+ """
+ Export annotations to EAF format.
+
+ Args:
+ context: ExportContext with annotation data
+ output_path: Directory path for output files
+ options: Optional settings:
+ - author: Author name for EAF header
+ - include_empty_tiers: Whether to include tiers with no annotations
+
+ Returns:
+ ExportResult with status and file paths
+ """
+ options = options or {}
+ files_written = []
+ warnings = []
+ stats = {"instances": 0, "annotations": 0, "tiers": 0}
+
+ # Create output directory
+ os.makedirs(output_path, exist_ok=True)
+
+ # Find tiered_annotation schemas
+ tiered_schemas = [
+ s for s in context.schemas
+ if s.get("annotation_type") == "tiered_annotation"
+ ]
+
+ for instance_id, item in context.items.items():
+ # Get annotations for this instance
+ instance_annotations = [
+ a for a in context.annotations
+ if a.get("instance_id") == instance_id
+ ]
+
+ for schema in tiered_schemas:
+ schema_name = schema.get("name", "tiered")
+
+ # Get tiered annotation data for this schema
+ tiered_data = None
+ for ann in instance_annotations:
+ if schema_name in ann.get("labels", {}):
+ try:
+ import json
+ raw_value = ann["labels"][schema_name]
+ if isinstance(raw_value, str):
+ tiered_data = json.loads(raw_value)
+ elif isinstance(raw_value, dict):
+ tiered_data = raw_value
+ except (json.JSONDecodeError, TypeError):
+ pass
+ break
+
+ if not tiered_data:
+ continue
+
+ # Get media URL
+ source_field = schema.get("source_field", "audio_url")
+ media_url = item.get(source_field, "")
+
+ # Generate EAF XML
+ root = self._create_eaf_document(
+ schema,
+ tiered_data,
+ media_url,
+ options
+ )
+
+ # Write to file
+ safe_id = "".join(c if c.isalnum() or c in "-_" else "_" for c in instance_id)
+ filename = f"{safe_id}_{schema_name}.eaf"
+ filepath = os.path.join(output_path, filename)
+
+ tree = ET.ElementTree(root)
+ tree.write(filepath, encoding="utf-8", xml_declaration=True)
+
+ files_written.append(filepath)
+ stats["instances"] += 1
+
+ # Count annotations
+ annotations = tiered_data.get("annotations", {})
+ for tier_anns in annotations.values():
+ stats["annotations"] += len(tier_anns)
+ stats["tiers"] = len(schema.get("tiers", []))
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=files_written,
+ warnings=warnings,
+ stats=stats
+ )
+
+ def _create_eaf_document(
+ self,
+ schema: dict,
+ tiered_data: dict,
+ media_url: str,
+ options: dict
+ ) -> ET.Element:
+ """
+ Create the EAF XML document structure.
+
+ Args:
+ schema: The tiered_annotation schema configuration
+ tiered_data: The annotation data
+ media_url: URL/path to the media file
+ options: Export options
+
+ Returns:
+ ET.Element root of the EAF document
+ """
+ # Root element
+ root = ET.Element("ANNOTATION_DOCUMENT")
+ root.set("xmlns:xsi", "http://www.w3.org/2001/XMLSchema-instance")
+ root.set("xsi:noNamespaceSchemaLocation",
+ "http://www.mpi.nl/tools/elan/EAFv3.0.xsd")
+ root.set("DATE", datetime.now().isoformat())
+ root.set("FORMAT", "3.0")
+ root.set("VERSION", "3.0")
+ root.set("AUTHOR", options.get("author", "Potato Annotation Tool"))
+
+ # Header
+ header = ET.SubElement(root, "HEADER")
+ header.set("MEDIA_FILE", "")
+ header.set("TIME_UNITS", "milliseconds")
+
+ if media_url:
+ media_descriptor = ET.SubElement(header, "MEDIA_DESCRIPTOR")
+ media_descriptor.set("MEDIA_URL", media_url)
+ media_descriptor.set("MIME_TYPE", self._get_mime_type(media_url))
+ media_descriptor.set("RELATIVE_MEDIA_URL", "")
+
+ # Property elements
+ props = [
+ ("lastUsedAnnotationId", "0"),
+ ]
+ for name, value in props:
+ prop = ET.SubElement(header, "PROPERTY")
+ prop.set("NAME", name)
+ prop.text = value
+
+ # TIME_ORDER - collect all unique time slots
+ time_order = ET.SubElement(root, "TIME_ORDER")
+ time_slots = tiered_data.get("time_slots", {})
+
+ if not time_slots:
+ # Generate from annotations
+ time_slots = self._generate_time_slots(tiered_data.get("annotations", {}))
+
+ # Sort and add time slots
+ slot_items = sorted(time_slots.items(), key=lambda x: x[1])
+ for slot_id, time_ms in slot_items:
+ ts = ET.SubElement(time_order, "TIME_SLOT")
+ ts.set("TIME_SLOT_ID", slot_id)
+ ts.set("TIME_VALUE", str(int(time_ms)))
+
+ # Create reverse mapping for looking up slot IDs
+ time_to_slot = {v: k for k, v in time_slots.items()}
+ initial_slot_count = len(time_to_slot)
+
+ # TIER elements
+ tiers = schema.get("tiers", [])
+ annotations = tiered_data.get("annotations", {})
+ annotation_id_counter = [0] # Use list to allow mutation in nested function
+
+ for tier_def in tiers:
+ tier_el = self._create_tier_element(
+ root,
+ tier_def,
+ annotations.get(tier_def["name"], []),
+ time_to_slot,
+ annotation_id_counter
+ )
+
+ # Add any dynamically created time slots to TIME_ORDER
+ if len(time_to_slot) > initial_slot_count:
+ slot_to_time = {v: k for k, v in time_to_slot.items()}
+ new_slots = sorted(
+ [(sid, slot_to_time[sid]) for sid in slot_to_time if sid not in time_slots],
+ key=lambda x: x[1]
+ )
+ for slot_id, time_ms in new_slots:
+ ts = ET.SubElement(time_order, "TIME_SLOT")
+ ts.set("TIME_SLOT_ID", slot_id)
+ ts.set("TIME_VALUE", str(int(time_ms)))
+
+ # LINGUISTIC_TYPE elements
+ self._create_linguistic_types(root, tiers)
+
+ # CONSTRAINT elements (for dependent tiers)
+ self._create_constraints(root)
+
+ return root
+
+ def _create_tier_element(
+ self,
+ root: ET.Element,
+ tier_def: dict,
+ tier_annotations: List[dict],
+ time_to_slot: dict,
+ annotation_id_counter: List[int]
+ ) -> ET.Element:
+ """
+ Create a TIER element with its annotations.
+
+ Args:
+ root: The root EAF element
+ tier_def: Tier definition from schema
+ tier_annotations: List of annotations for this tier
+ time_to_slot: Mapping of time values to slot IDs
+ annotation_id_counter: Counter for generating annotation IDs
+
+ Returns:
+ The created TIER element
+ """
+ tier_el = ET.SubElement(root, "TIER")
+ tier_el.set("TIER_ID", tier_def["name"])
+
+ # Determine linguistic type
+ if tier_def.get("tier_type") == "dependent":
+ constraint_type = tier_def.get("constraint_type", "included_in")
+ ling_type = f"default-lt-{constraint_type}"
+ tier_el.set("PARENT_REF", tier_def.get("parent_tier", ""))
+ else:
+ ling_type = "default-lt"
+
+ tier_el.set("LINGUISTIC_TYPE_REF", tier_def.get("linguistic_type", ling_type))
+ tier_el.set("DEFAULT_LOCALE", "en")
+
+ # Add annotations
+ for ann in sorted(tier_annotations, key=lambda a: a.get("start_time", 0)):
+ annotation_id_counter[0] += 1
+ ann_id = f"a{annotation_id_counter[0]}"
+
+ annotation_el = ET.SubElement(tier_el, "ANNOTATION")
+
+ if tier_def.get("tier_type") == "independent":
+ # ALIGNABLE_ANNOTATION for independent tiers
+ alignable = ET.SubElement(annotation_el, "ALIGNABLE_ANNOTATION")
+ alignable.set("ANNOTATION_ID", ann_id)
+
+ start_time = int(ann.get("start_time", 0))
+ end_time = int(ann.get("end_time", 0))
+
+ start_slot = self._get_or_create_slot(time_to_slot, start_time)
+ end_slot = self._get_or_create_slot(time_to_slot, end_time)
+
+ alignable.set("TIME_SLOT_REF1", start_slot)
+ alignable.set("TIME_SLOT_REF2", end_slot)
+
+ value_el = ET.SubElement(alignable, "ANNOTATION_VALUE")
+ value_el.text = ann.get("value") or ann.get("label", "")
+
+ else:
+ # REF_ANNOTATION for dependent tiers
+ ref_ann = ET.SubElement(annotation_el, "REF_ANNOTATION")
+ ref_ann.set("ANNOTATION_ID", ann_id)
+
+ # Reference parent annotation
+ parent_id = ann.get("parent_id", "")
+ if parent_id:
+ # Convert internal ID to EAF annotation ID
+ # For simplicity, use annotation reference
+ ref_ann.set("ANNOTATION_REF", parent_id)
+
+ value_el = ET.SubElement(ref_ann, "ANNOTATION_VALUE")
+ value_el.text = ann.get("value") or ann.get("label", "")
+
+ return tier_el
+
+ def _get_or_create_slot(
+ self,
+ time_to_slot: dict,
+ time_ms: int
+ ) -> str:
+ """Get existing slot ID or create a new mapping."""
+ if time_ms in time_to_slot:
+ return time_to_slot[time_ms]
+
+ # Create new slot
+ slot_id = f"ts{len(time_to_slot) + 1}"
+ time_to_slot[time_ms] = slot_id
+ return slot_id
+
+ def _generate_time_slots(
+ self,
+ annotations: Dict[str, List[dict]]
+ ) -> Dict[str, int]:
+ """Generate time slots from annotations."""
+ times: Set[int] = set()
+
+ for tier_anns in annotations.values():
+ for ann in tier_anns:
+ if ann.get("start_time") is not None:
+ times.add(int(ann["start_time"]))
+ if ann.get("end_time") is not None:
+ times.add(int(ann["end_time"]))
+
+ return {
+ f"ts{i+1}": time
+ for i, time in enumerate(sorted(times))
+ }
+
+ def _create_linguistic_types(
+ self,
+ root: ET.Element,
+ tiers: List[dict]
+ ) -> None:
+ """Create LINGUISTIC_TYPE elements for all tier types."""
+ # Default linguistic type for independent tiers
+ lt = ET.SubElement(root, "LINGUISTIC_TYPE")
+ lt.set("LINGUISTIC_TYPE_ID", "default-lt")
+ lt.set("TIME_ALIGNABLE", "true")
+ lt.set("GRAPHIC_REFERENCES", "false")
+
+ # Linguistic types for each constraint type
+ constraint_types = set()
+ for tier in tiers:
+ if tier.get("tier_type") == "dependent":
+ constraint_types.add(tier.get("constraint_type", "included_in"))
+
+ for constraint in constraint_types:
+ lt = ET.SubElement(root, "LINGUISTIC_TYPE")
+ lt.set("LINGUISTIC_TYPE_ID", f"default-lt-{constraint}")
+ lt.set("CONSTRAINTS", self._constraint_to_elan(constraint))
+ lt.set("TIME_ALIGNABLE", "true" if constraint in ("time_subdivision", "included_in") else "false")
+ lt.set("GRAPHIC_REFERENCES", "false")
+
+ def _constraint_to_elan(self, constraint_type: str) -> str:
+ """Map constraint type to ELAN constraint stereotype."""
+ mapping = {
+ "time_subdivision": "Time_Subdivision",
+ "included_in": "Included_In",
+ "symbolic_association": "Symbolic_Association",
+ "symbolic_subdivision": "Symbolic_Subdivision",
+ }
+ return mapping.get(constraint_type, "Included_In")
+
+ def _create_constraints(self, root: ET.Element) -> None:
+ """Create CONSTRAINT elements for the standard ELAN constraint types."""
+ constraints = [
+ ("Time_Subdivision", "Time subdivision of parent annotation's time interval, no time gaps allowed within this interval"),
+ ("Symbolic_Subdivision", "Symbolic subdivision of a parent annotation. Annotations refer to the same time interval as the parent"),
+ ("Symbolic_Association", "1-1 association with a parent annotation"),
+ ("Included_In", "Time alignable annotations within the parent annotation's time interval, gaps are allowed"),
+ ]
+
+ for constraint_id, description in constraints:
+ constraint = ET.SubElement(root, "CONSTRAINT")
+ constraint.set("DESCRIPTION", description)
+ constraint.set("STEREOTYPE", constraint_id)
+
+ def _get_mime_type(self, url: str) -> str:
+ """Determine MIME type from file extension."""
+ url_lower = url.lower()
+ if url_lower.endswith((".mp4", ".m4v")):
+ return "video/mp4"
+ elif url_lower.endswith((".webm",)):
+ return "video/webm"
+ elif url_lower.endswith((".avi",)):
+ return "video/avi"
+ elif url_lower.endswith((".mov",)):
+ return "video/quicktime"
+ elif url_lower.endswith((".wav",)):
+ return "audio/wav"
+ elif url_lower.endswith((".mp3",)):
+ return "audio/mpeg"
+ elif url_lower.endswith((".ogg", ".oga")):
+ return "audio/ogg"
+ elif url_lower.endswith((".flac",)):
+ return "audio/flac"
+ else:
+ return "audio/x-wav" # Default
diff --git a/potato/export/huggingface_exporter.py b/potato/export/huggingface_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..622a422ea66171224f4a08a59c6ed8748a184e4e
--- /dev/null
+++ b/potato/export/huggingface_exporter.py
@@ -0,0 +1,333 @@
+"""
+HuggingFace Hub Exporter
+
+Pushes annotations as a HuggingFace Dataset to the Hub, making them
+available for download via `datasets.load_dataset()`.
+
+Requires: pip install huggingface_hub>=0.20.0 datasets>=2.14.0
+
+Usage:
+ python -m potato.export \\
+ --config config.yaml \\
+ --format huggingface \\
+ --output your-org/my-annotations \\
+ --option token=hf_xxx \\
+ --option private=true
+"""
+
+import json
+import logging
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+
+logger = logging.getLogger(__name__)
+
+
+def _check_deps():
+ """Try to import HF dependencies and return them, or raise ImportError."""
+ from datasets import Dataset, DatasetDict
+ from huggingface_hub import DatasetCard, DatasetCardData
+ return Dataset, DatasetDict, DatasetCard, DatasetCardData
+
+
+class HuggingFaceExporter(BaseExporter):
+ """
+ Exports annotations to HuggingFace Hub as a Dataset.
+
+ The output_path parameter is used as the repo_id (e.g., "your-org/dataset-name").
+ Produces a DatasetDict with an 'annotations' split, plus optional 'spans' and 'items'.
+ """
+
+ format_name = "huggingface"
+ description = "Push annotations to HuggingFace Hub as a Dataset"
+ file_extensions = [] # No local files โ pushes to Hub
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ try:
+ _check_deps()
+ except ImportError:
+ return False, (
+ "huggingface_hub and datasets are required for HuggingFace export. "
+ "Install with: pip install huggingface_hub>=0.20.0 datasets>=2.14.0"
+ )
+
+ if not context.annotations:
+ return False, "No annotations to export"
+
+ return True, ""
+
+ def build_dataset_dict(self, context: ExportContext,
+ include_spans: bool = True,
+ include_items: bool = True) -> "DatasetDict":
+ """
+ Build a DatasetDict from an ExportContext without pushing to Hub.
+
+ Args:
+ context: ExportContext with annotations, items, schemas
+ include_spans: Include a 'spans' split
+ include_items: Include an 'items' split
+
+ Returns:
+ datasets.DatasetDict with annotations/spans/items splits
+
+ Raises:
+ ImportError: If datasets library is not installed
+ ValueError: If no data to build
+ """
+ Dataset, DatasetDict, _, _ = _check_deps()
+
+ schema_map = {s["name"]: s for s in context.schemas}
+ splits = {}
+
+ # 1. Annotations split
+ ann_rows = self._build_annotation_rows(context.annotations, schema_map)
+ if ann_rows:
+ splits["annotations"] = Dataset.from_list(ann_rows)
+
+ # 2. Spans split (optional)
+ if include_spans:
+ span_rows = self._build_span_rows(context.annotations)
+ if span_rows:
+ splits["spans"] = Dataset.from_list(span_rows)
+
+ # 3. Items split (optional)
+ if include_items and context.items:
+ item_rows = self._build_item_rows(context.items)
+ if item_rows:
+ splits["items"] = Dataset.from_list(item_rows)
+
+ if not splits:
+ raise ValueError("No data to build โ annotations list is empty")
+
+ return DatasetDict(splits)
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ warnings_list = []
+
+ try:
+ _, _, DatasetCard, DatasetCardData = _check_deps()
+ except ImportError as e:
+ return ExportResult(
+ success=False,
+ format_name=self.format_name,
+ errors=[str(e)],
+ )
+
+ # Parse options
+ repo_id = output_path # e.g., "your-org/my-annotations"
+ token = options.get("token") or os.environ.get("HF_TOKEN")
+ private = options.get("private", False)
+ commit_message = options.get("commit_message", "Upload annotations from Potato")
+ include_items = options.get("include_items", True)
+ include_spans = options.get("include_spans", True)
+
+ # Normalize string booleans from CLI
+ if isinstance(private, str):
+ private = private.lower() not in ("false", "0", "no")
+ if isinstance(include_items, str):
+ include_items = include_items.lower() not in ("false", "0", "no")
+ if isinstance(include_spans, str):
+ include_spans = include_spans.lower() not in ("false", "0", "no")
+
+ if not repo_id or "/" not in repo_id:
+ return ExportResult(
+ success=False,
+ format_name=self.format_name,
+ errors=[
+ f"output_path must be a HuggingFace repo ID "
+ f"(e.g., 'your-org/dataset-name'), got: '{repo_id}'"
+ ],
+ )
+
+ try:
+ dataset_dict = self.build_dataset_dict(
+ context,
+ include_spans=include_spans,
+ include_items=include_items,
+ )
+
+ dataset_dict.push_to_hub(
+ repo_id,
+ token=token,
+ private=private,
+ commit_message=commit_message,
+ )
+
+ # Compute stats by rebuilding row counts (avoids depending on
+ # DatasetDict internals for len/keys).
+ schema_map = {s["name"]: s for s in context.schemas}
+ ann_rows = self._build_annotation_rows(context.annotations, schema_map)
+ span_rows = self._build_span_rows(context.annotations) if include_spans else []
+ item_rows = self._build_item_rows(context.items) if include_items and context.items else []
+
+ # Generate and push dataset card
+ try:
+ card_content = self._build_dataset_card(
+ context, repo_id, ann_rows, schema_map
+ )
+ card = DatasetCard(card_content)
+ card.push_to_hub(repo_id, token=token)
+ except Exception as e:
+ warnings_list.append(f"Dataset card push failed: {e}")
+ logger.warning("Failed to push dataset card: %s", e)
+
+ # Build splits list based on what was actually included
+ splits_list = []
+ if ann_rows:
+ splits_list.append("annotations")
+ if span_rows:
+ splits_list.append("spans")
+ if item_rows:
+ splits_list.append("items")
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ warnings=warnings_list,
+ stats={
+ "repo_id": repo_id,
+ "annotation_rows": len(ann_rows),
+ "span_rows": len(span_rows),
+ "item_rows": len(item_rows),
+ "splits": splits_list,
+ "private": private,
+ },
+ )
+
+ except ValueError as e:
+ return ExportResult(
+ success=False,
+ format_name=self.format_name,
+ errors=[str(e)],
+ )
+ except Exception as e:
+ logger.error("HuggingFace Hub export failed: %s", e)
+ return ExportResult(
+ success=False,
+ format_name=self.format_name,
+ errors=[str(e)],
+ )
+
+ def _build_annotation_rows(self, annotations: List[dict],
+ schema_map: Dict[str, dict]) -> List[dict]:
+ """Build flat row dicts for the annotations dataset."""
+ rows = []
+ for ann in annotations:
+ row = {
+ "instance_id": ann.get("instance_id", ""),
+ "user_id": ann.get("user_id", ""),
+ }
+
+ labels = ann.get("labels", {})
+ for schema_name, value in labels.items():
+ # Serialize complex values as JSON strings for schema flexibility
+ if isinstance(value, (dict, list)):
+ row[schema_name] = json.dumps(value, ensure_ascii=False)
+ else:
+ row[schema_name] = value
+
+ rows.append(row)
+ return rows
+
+ def _build_span_rows(self, annotations: List[dict]) -> List[dict]:
+ """Build flat row dicts for the spans dataset."""
+ rows = []
+ for ann in annotations:
+ instance_id = ann.get("instance_id", "")
+ user_id = ann.get("user_id", "")
+ spans = ann.get("spans", {})
+
+ for schema_name, span_list in spans.items():
+ if not isinstance(span_list, list):
+ continue
+ for span in span_list:
+ if not isinstance(span, dict):
+ continue
+ rows.append({
+ "instance_id": instance_id,
+ "user_id": user_id,
+ "schema_name": schema_name,
+ "start": span.get("start"),
+ "end": span.get("end"),
+ "label": span.get("label", ""),
+ "text": span.get("text", ""),
+ })
+ return rows
+
+ def _build_item_rows(self, items: Dict[str, dict]) -> List[dict]:
+ """Build flat row dicts for the items dataset."""
+ rows = []
+ for item_id, item_data in items.items():
+ row = {"item_id": item_id}
+ if isinstance(item_data, dict):
+ for key, val in item_data.items():
+ if isinstance(val, (dict, list)):
+ row[key] = json.dumps(val, ensure_ascii=False)
+ else:
+ row[key] = val
+ rows.append(row)
+ return rows
+
+ def _build_dataset_card(self, context: ExportContext, repo_id: str,
+ ann_rows: List[dict],
+ schema_map: Dict[str, dict]) -> str:
+ """Build a DatasetCard markdown string with task metadata."""
+ schema_descriptions = []
+ for name, schema in schema_map.items():
+ ann_type = schema.get("annotation_type", "unknown")
+ desc = schema.get("description", "")
+ labels = schema.get("labels", [])
+ label_str = ", ".join(labels[:10]) if labels else "N/A"
+ if len(labels) > 10:
+ label_str += f" (+{len(labels) - 10} more)"
+ schema_descriptions.append(
+ f"- **{name}** ({ann_type}): {desc}\n Labels: {label_str}"
+ )
+
+ schemas_section = "\n".join(schema_descriptions) if schema_descriptions else "N/A"
+
+ card = f"""---
+annotations_creators:
+- crowdsourced
+language_creators:
+- expert-generated
+source_datasets: []
+task_categories:
+- text-classification
+tags:
+- potato-annotation
+---
+
+# {repo_id.split('/')[-1]}
+
+Annotations exported from [Potato](https://github.com/davidjurgens/potato) annotation tool.
+
+## Dataset Structure
+
+### Splits
+
+- **annotations**: {len(ann_rows)} annotation records (one per instance-annotator pair)
+
+### Annotation Schemas
+
+{schemas_section}
+
+## Usage
+
+```python
+from datasets import load_dataset
+
+ds = load_dataset("{repo_id}")
+print(ds["annotations"][0])
+```
+
+## Export Details
+
+- Exported by: Potato annotation platform
+- Format: HuggingFace Datasets
+"""
+ return card
diff --git a/potato/export/mask_exporter.py b/potato/export/mask_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b67858e31bae9e16e9f577976157962035fb0b6
--- /dev/null
+++ b/potato/export/mask_exporter.py
@@ -0,0 +1,136 @@
+"""
+Mask Exporter
+
+Exports segmentation mask annotations as PNG binary images.
+Each label gets a separate PNG where filled pixels are the label color
+and background is transparent.
+
+Requires: numpy and Pillow (PIL)
+"""
+
+import os
+import logging
+from typing import Optional, Tuple, List
+
+from .base import BaseExporter, ExportContext, ExportResult
+from .cv_utils import (
+ extract_image_annotations,
+ get_image_dimensions,
+ get_image_filename,
+ build_category_mapping,
+ decode_rle,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class MaskExporter(BaseExporter):
+ format_name = "mask_png"
+ description = "Segmentation masks as PNG images (requires Pillow)"
+ file_extensions = [".png"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ # Check for Pillow
+ try:
+ from PIL import Image
+ except ImportError:
+ return False, "Pillow (PIL) is required for mask export. Install with: pip install Pillow"
+
+ has_image_schema = any(
+ s.get("annotation_type") == "image_annotation"
+ for s in context.schemas
+ )
+ if not has_image_schema:
+ return False, "No image_annotation schema found in config"
+
+ return True, ""
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ from PIL import Image
+
+ options = options or {}
+ warnings = []
+ files_written = []
+
+ os.makedirs(output_path, exist_ok=True)
+ category_map = build_category_mapping(context.annotations, context.schemas)
+
+ # Assign colors to categories
+ default_colors = [
+ (255, 0, 0), (0, 255, 0), (0, 0, 255),
+ (255, 255, 0), (255, 0, 255), (0, 255, 255),
+ (128, 0, 0), (0, 128, 0), (0, 0, 128),
+ (128, 128, 0),
+ ]
+ category_colors = {}
+ for name, idx in category_map.items():
+ category_colors[name] = default_colors[idx % len(default_colors)]
+
+ masks_exported = 0
+
+ for ann in context.annotations:
+ instance_id = ann.get("instance_id", "")
+ item = context.items.get(instance_id, {})
+ img_anns = extract_image_annotations(ann)
+ if not img_anns:
+ continue
+
+ width, height = get_image_dimensions(item)
+ if width <= 0 or height <= 0:
+ # Try to get from mask RLE size
+ for _, objects in img_anns:
+ for obj in objects:
+ if obj.get("type") == "mask" and "rle" in obj:
+ size = obj["rle"].get("size", [])
+ if len(size) == 2:
+ height, width = size
+ break
+ if width > 0:
+ break
+
+ if width <= 0 or height <= 0:
+ warnings.append(f"No dimensions for {instance_id}, skipping masks")
+ continue
+
+ file_name = get_image_filename(item) or instance_id
+ raw_stem = os.path.splitext(os.path.basename(file_name))[0]
+ stem = "".join(c if c.isalnum() or c in "-_." else "_" for c in raw_stem)
+
+ for schema_name, objects in img_anns:
+ for obj in objects:
+ if obj.get("type") != "mask":
+ continue
+
+ label = obj.get("label", "unknown")
+ rle = obj.get("rle", {})
+ if not rle.get("counts"):
+ continue
+
+ mask_data = decode_rle(rle, width, height)
+ color = category_colors.get(label, (255, 255, 255))
+
+ # Create RGBA image
+ img = Image.new("RGBA", (width, height), (0, 0, 0, 0))
+ pixels = img.load()
+
+ for i, val in enumerate(mask_data):
+ if val:
+ y = i // width
+ x = i % width
+ if x < width and y < height:
+ pixels[x, y] = (color[0], color[1], color[2], 200)
+
+ safe_label = "".join(c if c.isalnum() or c in "-_." else "_" for c in label)
+ mask_file = os.path.join(output_path, f"{stem}_{safe_label}_mask.png")
+ img.save(mask_file)
+ files_written.append(mask_file)
+ masks_exported += 1
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=files_written,
+ warnings=warnings,
+ stats={"num_masks": masks_exported},
+ )
diff --git a/potato/export/nlp_utils.py b/potato/export/nlp_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..a7d9817626f19337f36b19bbb6d86c30a418048e
--- /dev/null
+++ b/potato/export/nlp_utils.py
@@ -0,0 +1,182 @@
+"""
+NLP Export Utilities
+
+Shared helpers for NLP export formats (CoNLL-2003, CoNLL-U).
+Provides tokenization and BIO tag alignment.
+"""
+
+from typing import List, Dict, Tuple, Optional
+import logging
+import re
+
+logger = logging.getLogger(__name__)
+
+
+def tokenize_text(text: str, method: str = "whitespace") -> List[Dict]:
+ """
+ Tokenize text into tokens with character offsets.
+
+ Args:
+ text: Input text string
+ method: Tokenization method. Options:
+ - "whitespace": Split on whitespace (default)
+ - "word_punct": Split on word boundaries and punctuation
+
+ Returns:
+ List of dicts with keys: token, start, end
+ """
+ if not text:
+ return []
+
+ if method == "word_punct":
+ tokens = []
+ for match in re.finditer(r'\S+', text):
+ raw = match.group()
+ raw_start = match.start()
+ # Split punctuation from word boundaries
+ sub_tokens = re.finditer(r'[\w]+|[^\w\s]', raw)
+ for sub in sub_tokens:
+ tokens.append({
+ "token": sub.group(),
+ "start": raw_start + sub.start(),
+ "end": raw_start + sub.end(),
+ })
+ return tokens
+
+ # Default: whitespace tokenization
+ tokens = []
+ for match in re.finditer(r'\S+', text):
+ tokens.append({
+ "token": match.group(),
+ "start": match.start(),
+ "end": match.end(),
+ })
+ return tokens
+
+
+def char_spans_to_bio_tags(
+ tokens: List[Dict],
+ spans: List[Dict],
+ scheme: str = "BIO"
+) -> List[str]:
+ """
+ Convert character-level spans to token-level BIO tags.
+
+ Handles:
+ - Multi-token entities
+ - Tokens partially inside spans (included if majority overlap)
+ - Overlapping spans (longest match wins)
+
+ Args:
+ tokens: List of token dicts with keys: token, start, end
+ spans: List of span dicts with keys: start, end, label (or name)
+ scheme: Tagging scheme - "BIO" (default) or "BIOES"
+
+ Returns:
+ List of BIO tag strings, one per token (e.g., ["O", "B-PER", "I-PER"])
+ """
+ if not tokens:
+ return []
+
+ tags = ["O"] * len(tokens)
+
+ if not spans:
+ return tags
+
+ # Sort spans by length (longest first) so longest match wins on overlap
+ sorted_spans = sorted(
+ spans,
+ key=lambda s: (s.get("end", 0) - s.get("start", 0)),
+ reverse=True,
+ )
+
+ # Track which tokens are already assigned
+ assigned = [False] * len(tokens)
+
+ for span in sorted_spans:
+ span_start = span.get("start", 0)
+ span_end = span.get("end", 0)
+ label = span.get("label") or span.get("name", "ENTITY")
+
+ if span_start >= span_end:
+ continue
+
+ # Find tokens that overlap with this span
+ span_tokens = []
+ for i, tok in enumerate(tokens):
+ if assigned[i]:
+ continue
+ # Calculate overlap
+ overlap_start = max(tok["start"], span_start)
+ overlap_end = min(tok["end"], span_end)
+ overlap = max(0, overlap_end - overlap_start)
+ tok_len = tok["end"] - tok["start"]
+ if tok_len > 0 and overlap > 0:
+ # Include token if overlap covers majority of the token
+ if overlap >= tok_len / 2:
+ span_tokens.append(i)
+
+ if not span_tokens:
+ continue
+
+ # Assign BIO tags
+ for j, tok_idx in enumerate(span_tokens):
+ if j == 0:
+ tags[tok_idx] = f"B-{label}"
+ else:
+ tags[tok_idx] = f"I-{label}"
+ assigned[tok_idx] = True
+
+ # Apply BIOES if requested
+ if scheme == "BIOES" and span_tokens:
+ if len(span_tokens) == 1:
+ tags[span_tokens[0]] = f"S-{label}"
+ else:
+ tags[span_tokens[-1]] = f"E-{label}"
+
+ return tags
+
+
+def group_sentences(tokens: List[Dict], text: str) -> List[List[int]]:
+ """
+ Group token indices into sentences based on sentence-ending punctuation.
+
+ Args:
+ tokens: List of token dicts
+ text: Original text
+
+ Returns:
+ List of lists of token indices, one list per sentence
+ """
+ if not tokens:
+ return []
+
+ sentences = []
+ current = []
+
+ for i, tok in enumerate(tokens):
+ current.append(i)
+ # Sentence boundary: token ends with sentence-final punctuation
+ # and is followed by whitespace + uppercase or end of text
+ token_text = tok["token"]
+ ends_with_sent_punct = (
+ token_text in (".", "!", "?", "...", "ใ")
+ or token_text.endswith(".")
+ or token_text.endswith("!")
+ or token_text.endswith("?")
+ )
+ if ends_with_sent_punct:
+ # Check if next token starts a new sentence (uppercase or end)
+ if i + 1 >= len(tokens):
+ sentences.append(current)
+ current = []
+ else:
+ next_tok = tokens[i + 1]["token"]
+ if next_tok and next_tok[0].isupper():
+ sentences.append(current)
+ current = []
+
+ if current:
+ sentences.append(current)
+
+ return sentences
diff --git a/potato/export/parquet_exporter.py b/potato/export/parquet_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..4c8bd7a1e86ea163be86ca0351f05ed49ed695f8
--- /dev/null
+++ b/potato/export/parquet_exporter.py
@@ -0,0 +1,253 @@
+"""
+Parquet Exporter
+
+Exports annotations as Apache Parquet files via PyArrow, producing columnar
+tables suitable for analysis with pandas, DuckDB, Spark, or HuggingFace Datasets.
+
+Output files:
+ annotations.parquet - One row per (instance_id, user_id) with flattened schema columns
+ spans.parquet - One row per span annotation (if span schemas exist)
+ items.parquet - One row per item with original data fields
+"""
+
+import logging
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+
+logger = logging.getLogger(__name__)
+
+
+def _check_pyarrow():
+ """Try to import pyarrow and return (pa, pq) or raise ImportError."""
+ import pyarrow as pa
+ import pyarrow.parquet as pq
+ return pa, pq
+
+
+class ParquetExporter(BaseExporter):
+ """
+ Exports annotations as Parquet files.
+
+ Produces up to three tables:
+ - annotations.parquet: one row per (instance_id, user_id), flat columns per schema
+ - spans.parquet: one row per span annotation
+ - items.parquet: one row per item (original data)
+ """
+
+ format_name = "parquet"
+ description = "Apache Parquet columnar format for large-scale analysis (pandas, DuckDB, Spark)"
+ file_extensions = [".parquet"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ try:
+ _check_pyarrow()
+ except ImportError:
+ return False, "pyarrow is required for Parquet export. Install with: pip install pyarrow>=12.0.0"
+
+ if not context.annotations:
+ return False, "No annotations to export"
+
+ return True, ""
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ files_written = []
+ warnings = []
+
+ try:
+ pa, pq = _check_pyarrow()
+ except ImportError as e:
+ return ExportResult(
+ success=False,
+ format_name=self.format_name,
+ errors=[str(e)],
+ )
+
+ compression = options.get("compression", "snappy")
+ include_items = options.get("include_items", True)
+ include_spans = options.get("include_spans", True)
+ row_group_size = options.get("row_group_size", None)
+
+ # Normalize string booleans from CLI
+ if isinstance(include_items, str):
+ include_items = include_items.lower() not in ("false", "0", "no")
+ if isinstance(include_spans, str):
+ include_spans = include_spans.lower() not in ("false", "0", "no")
+
+ try:
+ os.makedirs(output_path, exist_ok=True)
+ schema_map = {s["name"]: s for s in context.schemas}
+
+ # 1. Write annotations.parquet
+ ann_path = os.path.join(output_path, "annotations.parquet")
+ ann_rows = self._build_annotation_rows(context.annotations, schema_map)
+ if ann_rows:
+ table = pa.Table.from_pylist(ann_rows)
+ write_kwargs = {"compression": compression}
+ if row_group_size is not None:
+ write_kwargs["row_group_size"] = int(row_group_size)
+ pq.write_table(table, ann_path, **write_kwargs)
+ files_written.append(ann_path)
+
+ # 2. Write spans.parquet
+ if include_spans:
+ span_rows = self._build_span_rows(context.annotations)
+ if span_rows:
+ span_path = os.path.join(output_path, "spans.parquet")
+ span_table = pa.Table.from_pylist(span_rows)
+ pq.write_table(span_table, span_path, compression=compression)
+ files_written.append(span_path)
+
+ # 3. Write items.parquet
+ if include_items and context.items:
+ items_path = os.path.join(output_path, "items.parquet")
+ item_rows = self._build_item_rows(context.items)
+ if item_rows:
+ items_table = pa.Table.from_pylist(item_rows)
+ pq.write_table(items_table, items_path, compression=compression)
+ files_written.append(items_path)
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=files_written,
+ warnings=warnings,
+ stats={
+ "annotation_rows": len(ann_rows),
+ "span_rows": len(span_rows) if include_spans else 0,
+ "item_rows": len(item_rows) if include_items and context.items else 0,
+ "compression": compression,
+ },
+ )
+
+ except Exception as e:
+ logger.error(f"Parquet export failed: {e}")
+ return ExportResult(
+ success=False,
+ format_name=self.format_name,
+ files_written=files_written,
+ errors=[str(e)],
+ )
+
+ def _build_annotation_rows(self, annotations: List[dict],
+ schema_map: Dict[str, dict]) -> List[dict]:
+ """Build flat row dicts for the annotations table."""
+ rows = []
+ for ann in annotations:
+ row = {
+ "instance_id": ann.get("instance_id", ""),
+ "user_id": ann.get("user_id", ""),
+ }
+
+ labels = ann.get("labels", {})
+ for schema_name, value in labels.items():
+ schema_config = schema_map.get(schema_name, {})
+ schema_type = schema_config.get("annotation_type", "")
+ row[schema_name] = self._flatten_value(value, schema_type)
+
+ rows.append(row)
+ return rows
+
+ def _flatten_value(self, value: Any, schema_type: str) -> Any:
+ """Flatten an annotation value to a Parquet-compatible type."""
+ if schema_type in ("radio", "select"):
+ return self._flatten_categorical(value)
+ elif schema_type in ("likert", "slider", "number"):
+ return self._flatten_numeric(value)
+ elif schema_type == "multiselect":
+ return self._flatten_multiselect(value)
+ elif schema_type == "text":
+ if isinstance(value, str):
+ return value
+ return str(value) if value is not None else None
+ else:
+ # Generic fallback
+ if isinstance(value, dict):
+ return self._flatten_categorical(value)
+ return value
+
+ def _flatten_categorical(self, value: Any) -> Optional[str]:
+ """Extract the selected label from a categorical annotation."""
+ if isinstance(value, dict):
+ if not value:
+ return None
+ # Return the key with the highest value
+ def _sort_key(k):
+ try:
+ return float(value[k])
+ except (ValueError, TypeError):
+ return 0
+ return max(value.keys(), key=_sort_key)
+ if isinstance(value, str):
+ return value
+ return str(value) if value is not None else None
+
+ def _flatten_numeric(self, value: Any) -> Optional[float]:
+ """Extract a numeric value."""
+ if isinstance(value, (int, float)):
+ return float(value)
+ if isinstance(value, str):
+ try:
+ return float(value)
+ except ValueError:
+ return None
+ if isinstance(value, dict):
+ # Try to extract numeric from dict values
+ for v in value.values():
+ try:
+ return float(v)
+ except (ValueError, TypeError):
+ pass
+ return None
+
+ def _flatten_multiselect(self, value: Any) -> Optional[List[str]]:
+ """Extract selected labels from a multiselect annotation."""
+ if isinstance(value, dict):
+ return [label for label, selected in value.items() if selected]
+ if isinstance(value, list):
+ return [str(v) for v in value]
+ return None
+
+ def _build_span_rows(self, annotations: List[dict]) -> List[dict]:
+ """Build flat row dicts for the spans table."""
+ rows = []
+ for ann in annotations:
+ instance_id = ann.get("instance_id", "")
+ user_id = ann.get("user_id", "")
+ spans = ann.get("spans", {})
+
+ for schema_name, span_list in spans.items():
+ if not isinstance(span_list, list):
+ continue
+ for span in span_list:
+ if not isinstance(span, dict):
+ continue
+ rows.append({
+ "instance_id": instance_id,
+ "user_id": user_id,
+ "schema_name": schema_name,
+ "start": span.get("start"),
+ "end": span.get("end"),
+ "label": span.get("label", ""),
+ "text": span.get("text", ""),
+ })
+ return rows
+
+ def _build_item_rows(self, items: Dict[str, dict]) -> List[dict]:
+ """Build flat row dicts for the items table."""
+ rows = []
+ for item_id, item_data in items.items():
+ row = {"item_id": item_id}
+ if isinstance(item_data, dict):
+ for key, val in item_data.items():
+ # Convert non-primitive types to strings for Parquet compatibility
+ if isinstance(val, (dict, list)):
+ import json
+ row[key] = json.dumps(val, ensure_ascii=False)
+ else:
+ row[key] = val
+ rows.append(row)
+ return rows
diff --git a/potato/export/pascal_voc_exporter.py b/potato/export/pascal_voc_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d91dd710994a61bd1aefc661c0d78d15b140a88
--- /dev/null
+++ b/potato/export/pascal_voc_exporter.py
@@ -0,0 +1,148 @@
+"""
+Pascal VOC XML Exporter
+
+Exports image annotations to Pascal VOC format:
+- One XML file per image with structure
+"""
+
+import os
+import logging
+from xml.etree.ElementTree import Element, SubElement, ElementTree, indent
+from typing import Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+from .cv_utils import (
+ build_category_mapping,
+ polygon_to_bbox,
+ extract_image_annotations,
+ get_image_dimensions,
+ get_image_filename,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class PascalVOCExporter(BaseExporter):
+ format_name = "pascal_voc"
+ description = "Pascal VOC XML format for object detection"
+ file_extensions = [".xml"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ has_image_schema = any(
+ s.get("annotation_type") == "image_annotation"
+ for s in context.schemas
+ )
+ if not has_image_schema:
+ return False, "No image_annotation schema found in config"
+ return True, ""
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ warnings = []
+ files_written = []
+
+ os.makedirs(output_path, exist_ok=True)
+
+ # Group annotations by instance_id to produce one XML per image
+ image_objects = {} # instance_id -> list of object dicts
+
+ for ann in context.annotations:
+ instance_id = ann.get("instance_id", "")
+ img_anns = extract_image_annotations(ann)
+ if not img_anns:
+ continue
+
+ if instance_id not in image_objects:
+ image_objects[instance_id] = []
+
+ for schema_name, objects in img_anns:
+ for obj in objects:
+ image_objects[instance_id].append(obj)
+
+ for instance_id, objects in image_objects.items():
+ item = context.items.get(instance_id, {})
+ width, height = get_image_dimensions(item)
+ file_name = get_image_filename(item) or instance_id
+ raw_stem = os.path.splitext(os.path.basename(file_name))[0]
+ stem = "".join(c if c.isalnum() or c in "-_." else "_" for c in raw_stem)
+
+ root = Element("annotation")
+
+ folder_elem = SubElement(root, "folder")
+ folder_elem.text = "images"
+
+ filename_elem = SubElement(root, "filename")
+ filename_elem.text = os.path.basename(file_name)
+
+ size_elem = SubElement(root, "size")
+ SubElement(size_elem, "width").text = str(width)
+ SubElement(size_elem, "height").text = str(height)
+ SubElement(size_elem, "depth").text = str(item.get("depth", 3))
+
+ SubElement(root, "segmented").text = "0"
+
+ for obj in objects:
+ obj_type = obj.get("type", "")
+ label = obj.get("label", "")
+
+ if obj_type == "landmark":
+ warnings.append(
+ f"Landmark in {instance_id} skipped "
+ f"(not supported in Pascal VOC)"
+ )
+ continue
+
+ if obj_type == "bbox":
+ xmin = obj.get("x", 0)
+ ymin = obj.get("y", 0)
+ xmax = xmin + obj.get("width", 0)
+ ymax = ymin + obj.get("height", 0)
+
+ elif obj_type in ("polygon", "freeform"):
+ points = obj.get("points", [])
+ if not points:
+ continue
+ bx, by, bw, bh = polygon_to_bbox(points)
+ xmin = bx
+ ymin = by
+ xmax = bx + bw
+ ymax = by + bh
+ warnings.append(
+ f"{obj_type} in {instance_id} converted to enclosing bbox"
+ )
+
+ else:
+ warnings.append(
+ f"Unknown type '{obj_type}' in {instance_id}"
+ )
+ continue
+
+ obj_elem = SubElement(root, "object")
+ SubElement(obj_elem, "name").text = label
+ SubElement(obj_elem, "pose").text = "Unspecified"
+ SubElement(obj_elem, "truncated").text = "0"
+ SubElement(obj_elem, "difficult").text = "0"
+
+ bndbox = SubElement(obj_elem, "bndbox")
+ SubElement(bndbox, "xmin").text = str(int(round(xmin)))
+ SubElement(bndbox, "ymin").text = str(int(round(ymin)))
+ SubElement(bndbox, "xmax").text = str(int(round(xmax)))
+ SubElement(bndbox, "ymax").text = str(int(round(ymax)))
+
+ xml_file = os.path.join(output_path, f"{stem}.xml")
+ tree = ElementTree(root)
+ indent(tree, space=" ")
+ tree.write(xml_file, encoding="unicode", xml_declaration=True)
+ files_written.append(xml_file)
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=files_written,
+ warnings=warnings,
+ stats={
+ "num_images": len(image_objects),
+ "num_objects": sum(len(v) for v in image_objects.values()),
+ },
+ )
diff --git a/potato/export/quotation_report_exporter.py b/potato/export/quotation_report_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f2c20127f871f5a1776a7b838abcf343f54a44c
--- /dev/null
+++ b/potato/export/quotation_report_exporter.py
@@ -0,0 +1,150 @@
+"""
+Quotation Report Exporter
+
+Exports one row per coded text span ("quotation") with full provenance, suitable
+for qualitative-research deliverables and audit trails.
+
+Output columns:
+ schema annotation_scheme name
+ code label applied to the span
+ text quoted span text
+ start character offset (inclusive)
+ end character offset (exclusive)
+ field display field key (for multi-field instances)
+ instance_id source item id
+ source_doc text_key value or item id, for cross-reference
+ coder user id
+
+Option:
+ include_memos=true Also append one row per memo (schema="(memo)",
+ code=, text=, offsets from
+ the memo anchor when span-anchored). Useful for
+ audit trails and qualitative deliverables.
+"""
+
+import csv
+import logging
+import os
+from typing import Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+
+logger = logging.getLogger(__name__)
+
+
+class QuotationReportExporter(BaseExporter):
+ format_name = "quotation_report"
+ description = "Per-span CSV report with code, text, offsets, source, and coder"
+ file_extensions = [".csv"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ has_span_schema = any(
+ s.get("annotation_type") == "span" for s in context.schemas
+ )
+ if not has_span_schema:
+ return False, "No span annotation schema found in config"
+ return True, ""
+
+ @staticmethod
+ def _truthy(v) -> bool:
+ return str(v).strip().lower() in ("1", "true", "yes", "on")
+
+ def _memo_rows(self, context: ExportContext, default_text_key: str) -> list:
+ """One row per memo. Reads the universal project.sqlite via the
+ memos store. No-op (returns []) when the DB or memos package is
+ absent โ never breaks the core span export."""
+ task_dir = context.config.get("task_dir", ".")
+ db_path = os.path.join(task_dir, "project.sqlite")
+ if not os.path.exists(db_path):
+ return []
+ try:
+ from potato.memos import store as memo_store
+ except Exception:
+ return []
+ project = context.config.get("annotation_task_name") or "default"
+ out = []
+ for instance_id, item in (context.items or {}).items():
+ try:
+ memos = memo_store.list_for_instance(task_dir, project, instance_id)
+ except Exception as e:
+ logger.warning(f"Skipping memos for {instance_id}: {e}")
+ continue
+ source_doc = (item or {}).get(default_text_key, "") or instance_id
+ for m in memos:
+ anchor = m.get("anchor") or {}
+ out.append({
+ "schema": "(memo)",
+ "code": m.get("visibility", ""),
+ "text": m.get("body", ""),
+ "start": anchor.get("start", ""),
+ "end": anchor.get("end", ""),
+ "field": anchor.get("field", ""),
+ "instance_id": instance_id,
+ "source_doc": (source_doc[:200]
+ if isinstance(source_doc, str) else source_doc),
+ "coder": m.get("created_by", ""),
+ })
+ logger.info(f"Quotation report including {len(out)} memo rows")
+ return out
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ os.makedirs(output_path, exist_ok=True)
+ out_file = os.path.join(output_path, "quotations.csv")
+
+ item_props = context.config.get("item_properties", {})
+ default_text_key = item_props.get("text_key", "text")
+
+ rows = []
+ for ann in context.annotations:
+ instance_id = ann.get("instance_id", "")
+ coder = ann.get("user_id", "")
+ item = context.items.get(instance_id, {}) or {}
+ source_doc = item.get(default_text_key, "") or instance_id
+
+ spans = ann.get("spans", {}) or {}
+ for schema_name, span_list in spans.items():
+ for span in span_list or []:
+ label = (
+ span.get("label")
+ or span.get("annotation")
+ or span.get("category")
+ or ""
+ )
+ text = span.get("text", "")
+ start = span.get("start") if span.get("start") is not None else span.get("start_offset", "")
+ end = span.get("end") if span.get("end") is not None else span.get("end_offset", "")
+ field = span.get("field") or span.get("target_field") or ""
+ rows.append({
+ "schema": schema_name,
+ "code": label,
+ "text": text,
+ "start": start,
+ "end": end,
+ "field": field,
+ "instance_id": instance_id,
+ "source_doc": (source_doc[:200] if isinstance(source_doc, str) else source_doc),
+ "coder": coder,
+ })
+
+ if self._truthy(options.get("include_memos")):
+ rows.extend(self._memo_rows(context, default_text_key))
+
+ fieldnames = [
+ "schema", "code", "text", "start", "end", "field",
+ "instance_id", "source_doc", "coder",
+ ]
+ with open(out_file, "w", newline="", encoding="utf-8") as f:
+ writer = csv.DictWriter(f, fieldnames=fieldnames)
+ writer.writeheader()
+ for r in rows:
+ writer.writerow(r)
+
+ logger.info(f"Quotation report exported to {out_file}: {len(rows)} quotations")
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=[out_file],
+ stats={"quotations_exported": len(rows)},
+ )
diff --git a/potato/export/registry.py b/potato/export/registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..27721c0451c636329ec53f1e0d6e2569fbf86ce2
--- /dev/null
+++ b/potato/export/registry.py
@@ -0,0 +1,166 @@
+"""
+Export Registry
+
+Centralized registry for export format handlers, following the same
+pattern as SchemaRegistry and DisplayRegistry.
+
+Usage:
+ from potato.export.registry import export_registry
+
+ # List all exporters
+ exporters = export_registry.list_exporters()
+
+ # Export annotations
+ result = export_registry.export("coco", context, output_path)
+"""
+
+import logging
+from typing import Dict, List, Optional, Any
+
+from .base import BaseExporter, ExportContext, ExportResult
+
+logger = logging.getLogger(__name__)
+
+
+class ExportRegistry:
+ """
+ Centralized registry for annotation export formats.
+
+ Provides methods to register, retrieve, and invoke exporters.
+ """
+
+ def __init__(self):
+ self._exporters: Dict[str, BaseExporter] = {}
+ logger.debug("ExportRegistry initialized")
+
+ def register(self, exporter: BaseExporter) -> None:
+ """
+ Register an exporter instance.
+
+ Args:
+ exporter: BaseExporter subclass instance
+
+ Raises:
+ ValueError: If an exporter with the same format_name is already registered
+ """
+ name = exporter.format_name
+ if not name:
+ raise ValueError("Exporter must have a non-empty format_name")
+ if name in self._exporters:
+ raise ValueError(f"Exporter '{name}' is already registered")
+
+ self._exporters[name] = exporter
+ logger.debug(f"Registered exporter: {name}")
+
+ def get(self, name: str) -> Optional[BaseExporter]:
+ """Get an exporter by format name."""
+ return self._exporters.get(name)
+
+ def export(self, format_name: str, context: ExportContext,
+ output_path: str, options: Optional[dict] = None) -> ExportResult:
+ """
+ Export annotations using the named format.
+
+ Args:
+ format_name: Export format identifier (e.g., "coco", "yolo")
+ context: ExportContext with annotation data
+ output_path: Output directory or file path
+ options: Format-specific options
+
+ Returns:
+ ExportResult
+
+ Raises:
+ ValueError: If format is not registered or cannot handle the context
+ """
+ exporter = self.get(format_name)
+ if not exporter:
+ supported = ", ".join(sorted(self._exporters.keys()))
+ raise ValueError(
+ f"Unknown export format: '{format_name}'. "
+ f"Supported formats: {supported}"
+ )
+
+ can, reason = exporter.can_export(context)
+ if not can:
+ return ExportResult(
+ success=False,
+ format_name=format_name,
+ errors=[f"Cannot export: {reason}"],
+ )
+
+ return exporter.export(context, output_path, options)
+
+ def list_exporters(self) -> List[Dict[str, Any]]:
+ """List all registered exporters with metadata."""
+ return [
+ exporter.get_format_info()
+ for exporter in sorted(self._exporters.values(),
+ key=lambda e: e.format_name)
+ ]
+
+ def get_supported_formats(self) -> List[str]:
+ """Get sorted list of supported format names."""
+ return sorted(self._exporters.keys())
+
+ def is_registered(self, name: str) -> bool:
+ """Check if a format is registered."""
+ return name in self._exporters
+
+
+# Global registry instance
+export_registry = ExportRegistry()
+
+
+def _register_builtin_exporters():
+ """Register all built-in exporters. Called on import."""
+ from .coco_exporter import COCOExporter
+ from .yolo_exporter import YOLOExporter
+ from .pascal_voc_exporter import PascalVOCExporter
+ from .conll_2003_exporter import CoNLL2003Exporter
+ from .conll_u_exporter import CoNLLUExporter
+ from .mask_exporter import MaskExporter
+ from .eaf_exporter import EAFExporter
+ from .textgrid_exporter import TextGridExporter
+ from .agent_eval_exporter import AgentEvalExporter
+ from .coding_eval_exporter import CodingEvalExporter
+ from .trajectory_correction_exporter import TrajectoryCorrectionExporter
+ from .parquet_exporter import ParquetExporter
+ from .tabular_exporter import CSVExporter, TSVExporter, JSONLExporter
+ from .codebook_exporter import CodebookExporter
+ from .quotation_report_exporter import QuotationReportExporter
+
+ exporters = [
+ COCOExporter(),
+ YOLOExporter(),
+ PascalVOCExporter(),
+ CoNLL2003Exporter(),
+ CoNLLUExporter(),
+ MaskExporter(),
+ EAFExporter(),
+ TextGridExporter(),
+ AgentEvalExporter(),
+ CodingEvalExporter(),
+ TrajectoryCorrectionExporter(),
+ ParquetExporter(),
+ CSVExporter(),
+ TSVExporter(),
+ JSONLExporter(),
+ CodebookExporter(),
+ QuotationReportExporter(),
+ ]
+
+ # Optional exporters with external dependencies
+ try:
+ from .huggingface_exporter import HuggingFaceExporter
+ exporters.append(HuggingFaceExporter())
+ except ImportError:
+ logger.debug("HuggingFace exporter not available (missing dependencies)")
+
+ for exporter in exporters:
+ export_registry.register(exporter)
+
+ logger.debug(f"Registered {len(exporters)} built-in exporters")
+
+
+_register_builtin_exporters()
diff --git a/potato/export/tabular_exporter.py b/potato/export/tabular_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..630547eb48be707b61b62da347db6ff67649e2ca
--- /dev/null
+++ b/potato/export/tabular_exporter.py
@@ -0,0 +1,240 @@
+"""
+Tabular Exporters (CSV, TSV, JSONL)
+
+Exports annotations to flat tabular formats suitable for analysis in
+spreadsheets, pandas, or streaming pipelines.
+"""
+
+import csv
+import json
+import os
+import logging
+from typing import Optional, Tuple, List
+
+from .base import BaseExporter, ExportContext, ExportResult
+
+logger = logging.getLogger(__name__)
+
+
+def _flatten_annotation(ann: dict) -> dict:
+ """Flatten a single annotation record into a flat dict for tabular output."""
+ row = {
+ "instance_id": ann.get("instance_id", ""),
+ "user_id": ann.get("user_id", ""),
+ }
+ # Flatten labels: schema_name.label_name = value
+ for schema_name, labels in ann.get("labels", {}).items():
+ if isinstance(labels, dict):
+ for label_name, value in labels.items():
+ col = f"{schema_name}.{label_name}" if label_name else schema_name
+ row[col] = value if not isinstance(value, (dict, list)) else json.dumps(value)
+ else:
+ row[schema_name] = labels if not isinstance(labels, (dict, list)) else json.dumps(labels)
+
+ # Flatten spans as JSON strings
+ for schema_name, spans in ann.get("spans", {}).items():
+ row[f"{schema_name}._spans"] = json.dumps(spans)
+
+ return row
+
+
+class CSVExporter(BaseExporter):
+ """Export annotations to CSV format."""
+
+ format_name = "csv"
+ description = "Comma-separated values (one row per user-instance annotation)"
+ file_extensions = [".csv"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ if not context.annotations:
+ return False, "No annotations to export"
+ return True, ""
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ return _write_delimited(context, output_path, "csv", ",")
+
+
+class TSVExporter(BaseExporter):
+ """Export annotations to TSV format."""
+
+ format_name = "tsv"
+ description = "Tab-separated values (one row per user-instance annotation)"
+ file_extensions = [".tsv"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ if not context.annotations:
+ return False, "No annotations to export"
+ return True, ""
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ return _write_delimited(context, output_path, "tsv", "\t")
+
+
+class JSONLExporter(BaseExporter):
+ """Export annotations to JSONL format (one JSON object per line)."""
+
+ format_name = "jsonl"
+ description = "JSON Lines (one JSON object per user-instance annotation)"
+ file_extensions = [".jsonl"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ if not context.annotations:
+ return False, "No annotations to export"
+ return True, ""
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ os.makedirs(output_path, exist_ok=True)
+ out_file = os.path.join(output_path, "annotations.jsonl")
+
+ with open(out_file, "w", encoding="utf-8") as f:
+ for ann in context.annotations:
+ record = {
+ "instance_id": ann.get("instance_id", ""),
+ "user_id": ann.get("user_id", ""),
+ "labels": ann.get("labels", {}),
+ "spans": ann.get("spans", {}),
+ "links": ann.get("links", {}),
+ }
+ f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+ files_written = [out_file]
+ phase_file = _write_phase_jsonl(context, output_path)
+ if phase_file:
+ files_written.append(phase_file)
+
+ warnings = []
+ excl = _phase_exclusion_warning(context)
+ if excl:
+ warnings.append(excl)
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=files_written,
+ warnings=warnings,
+ stats={
+ "num_records": len(context.annotations),
+ "num_phase_responses": len(context.phase_responses) if phase_file else 0,
+ "num_phase_responses_excluded": (
+ len(context.phase_responses) if not phase_file else 0),
+ },
+ )
+
+
+def _should_include_phase_data(context: ExportContext) -> bool:
+ """Check if phase response export is enabled."""
+ return (
+ bool(context.phase_responses)
+ and context.config.get("export_include_phase_data", False)
+ )
+
+
+def _phase_exclusion_warning(context: ExportContext) -> Optional[str]:
+ """Return a warning when phase/survey responses exist but are NOT exported.
+
+ Phase-response export is opt-in via ``export_include_phase_data``. Without this
+ warning a survey/consent/instrument study would export with all phase responses
+ silently missing and the stats reporting ``num_phase_responses: 0`` (F-047),
+ making it look like no survey data was ever collected.
+ """
+ if context.phase_responses and not context.config.get("export_include_phase_data", False):
+ return (
+ f"{len(context.phase_responses)} phase/survey responses were found but "
+ f"NOT exported. Set 'export_include_phase_data: true' in your config to "
+ f"write them to a phase_responses file."
+ )
+ return None
+
+
+def _write_phase_delimited(context: ExportContext, output_path: str,
+ fmt_name: str, delimiter: str) -> Optional[str]:
+ """Write phase responses as a separate delimited file. Returns file path or None."""
+ if not _should_include_phase_data(context):
+ return None
+
+ out_file = os.path.join(output_path, f"phase_responses.{fmt_name}")
+ columns = ["user_id", "phase", "page", "schema", "label_name", "value"]
+
+ with open(out_file, "w", newline="", encoding="utf-8") as f:
+ writer = csv.DictWriter(f, fieldnames=columns, delimiter=delimiter,
+ extrasaction="ignore")
+ writer.writeheader()
+ for row in context.phase_responses:
+ writer.writerow(row)
+
+ return out_file
+
+
+def _write_phase_jsonl(context: ExportContext, output_path: str) -> Optional[str]:
+ """Write phase responses as a JSONL file. Returns file path or None."""
+ if not _should_include_phase_data(context):
+ return None
+
+ out_file = os.path.join(output_path, "phase_responses.jsonl")
+
+ with open(out_file, "w", encoding="utf-8") as f:
+ for row in context.phase_responses:
+ f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+ return out_file
+
+
+def _write_delimited(context: ExportContext, output_path: str,
+ fmt_name: str, delimiter: str) -> ExportResult:
+ """Write annotations as a delimited file (CSV or TSV)."""
+ os.makedirs(output_path, exist_ok=True)
+ out_file = os.path.join(output_path, f"annotations.{fmt_name}")
+
+ # Flatten all annotations to collect the full set of columns
+ rows = [_flatten_annotation(ann) for ann in context.annotations]
+
+ if not rows:
+ return ExportResult(
+ success=True,
+ format_name=fmt_name,
+ files_written=[out_file],
+ stats={"num_records": 0},
+ )
+
+ # Collect all column names preserving order (instance_id, user_id first)
+ columns = ["instance_id", "user_id"]
+ seen = set(columns)
+ for row in rows:
+ for key in row:
+ if key not in seen:
+ columns.append(key)
+ seen.add(key)
+
+ with open(out_file, "w", newline="", encoding="utf-8") as f:
+ writer = csv.DictWriter(f, fieldnames=columns, delimiter=delimiter,
+ extrasaction="ignore")
+ writer.writeheader()
+ for row in rows:
+ writer.writerow(row)
+
+ files_written = [out_file]
+ phase_file = _write_phase_delimited(context, output_path, fmt_name, delimiter)
+ if phase_file:
+ files_written.append(phase_file)
+
+ warnings = []
+ excl = _phase_exclusion_warning(context)
+ if excl:
+ warnings.append(excl)
+
+ return ExportResult(
+ success=True,
+ format_name=fmt_name,
+ files_written=files_written,
+ warnings=warnings,
+ stats={
+ "num_records": len(rows),
+ "num_columns": len(columns),
+ "num_phase_responses": len(context.phase_responses) if phase_file else 0,
+ "num_phase_responses_excluded": (
+ len(context.phase_responses) if not phase_file else 0),
+ },
+ )
diff --git a/potato/export/textgrid_exporter.py b/potato/export/textgrid_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..043fac7b040a7142b57221caeec1dd3dcf6526a6
--- /dev/null
+++ b/potato/export/textgrid_exporter.py
@@ -0,0 +1,342 @@
+"""
+TextGrid Exporter
+
+Exports tiered annotations to Praat TextGrid format.
+TextGrid is the native format for Praat (https://www.fon.hum.uva.nl/praat/),
+a tool widely used for phonetic analysis and annotation.
+
+The TextGrid format supports:
+- Interval tiers (segments with start/end times)
+- Point tiers (single-point annotations)
+- Multiple tiers with independent time alignments
+
+Note: TextGrid doesn't natively support hierarchical relationships between
+tiers, so the export flattens the hierarchy while preserving all annotations.
+"""
+
+import logging
+import os
+from typing import Dict, List, Any, Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+
+logger = logging.getLogger(__name__)
+
+
+class TextGridExporter(BaseExporter):
+ """
+ Exports tiered annotations to Praat TextGrid format.
+
+ This exporter creates TextGrid files that can be opened in Praat
+ for phonetic analysis or further annotation.
+ """
+
+ format_name = "textgrid"
+ description = "Praat TextGrid format for phonetic annotation"
+ file_extensions = [".TextGrid"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ """
+ Check if the context contains tiered_annotation schema.
+
+ Args:
+ context: ExportContext to validate
+
+ Returns:
+ Tuple of (can_export, reason)
+ """
+ for schema in context.schemas:
+ if schema.get("annotation_type") == "tiered_annotation":
+ return True, ""
+
+ return False, "No tiered_annotation schema found in configuration"
+
+ def export(
+ self,
+ context: ExportContext,
+ output_path: str,
+ options: Optional[dict] = None
+ ) -> ExportResult:
+ """
+ Export annotations to TextGrid format.
+
+ Args:
+ context: ExportContext with annotation data
+ output_path: Directory path for output files
+ options: Optional settings:
+ - format: "long" (default) or "short" TextGrid format
+ - fill_gaps: Whether to fill gaps between annotations
+
+ Returns:
+ ExportResult with status and file paths
+ """
+ options = options or {}
+ files_written = []
+ warnings = []
+ stats = {"instances": 0, "annotations": 0, "tiers": 0}
+
+ # Create output directory
+ os.makedirs(output_path, exist_ok=True)
+
+ use_short_format = options.get("format", "long") == "short"
+
+ # Find tiered_annotation schemas
+ tiered_schemas = [
+ s for s in context.schemas
+ if s.get("annotation_type") == "tiered_annotation"
+ ]
+
+ for instance_id, item in context.items.items():
+ # Get annotations for this instance
+ instance_annotations = [
+ a for a in context.annotations
+ if a.get("instance_id") == instance_id
+ ]
+
+ for schema in tiered_schemas:
+ schema_name = schema.get("name", "tiered")
+
+ # Get tiered annotation data for this schema
+ tiered_data = None
+ for ann in instance_annotations:
+ if schema_name in ann.get("labels", {}):
+ try:
+ import json
+ raw_value = ann["labels"][schema_name]
+ if isinstance(raw_value, str):
+ tiered_data = json.loads(raw_value)
+ elif isinstance(raw_value, dict):
+ tiered_data = raw_value
+ except (json.JSONDecodeError, TypeError):
+ pass
+ break
+
+ if not tiered_data:
+ continue
+
+ # Generate TextGrid content
+ content = self._create_textgrid(
+ schema,
+ tiered_data,
+ use_short_format
+ )
+
+ # Write to file
+ safe_id = "".join(c if c.isalnum() or c in "-_" else "_" for c in instance_id)
+ filename = f"{safe_id}_{schema_name}.TextGrid"
+ filepath = os.path.join(output_path, filename)
+
+ with open(filepath, 'w', encoding='utf-8') as f:
+ f.write(content)
+
+ files_written.append(filepath)
+ stats["instances"] += 1
+
+ # Count annotations
+ annotations = tiered_data.get("annotations", {})
+ for tier_anns in annotations.values():
+ stats["annotations"] += len(tier_anns)
+ stats["tiers"] = len(schema.get("tiers", []))
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=files_written,
+ warnings=warnings,
+ stats=stats
+ )
+
+ def _create_textgrid(
+ self,
+ schema: dict,
+ tiered_data: dict,
+ use_short_format: bool = False
+ ) -> str:
+ """
+ Create TextGrid file content.
+
+ Args:
+ schema: The tiered_annotation schema configuration
+ tiered_data: The annotation data
+ use_short_format: Whether to use short TextGrid format
+
+ Returns:
+ TextGrid file content as string
+ """
+ tiers = schema.get("tiers", [])
+ annotations = tiered_data.get("annotations", {})
+
+ # Calculate time bounds
+ min_time = 0.0
+ max_time = self._get_max_time(annotations)
+
+ if max_time == 0:
+ max_time = 1.0 # Default duration if no annotations
+
+ if use_short_format:
+ return self._create_short_textgrid(tiers, annotations, min_time, max_time)
+ else:
+ return self._create_long_textgrid(tiers, annotations, min_time, max_time)
+
+ def _create_long_textgrid(
+ self,
+ tiers: List[dict],
+ annotations: Dict[str, List[dict]],
+ min_time: float,
+ max_time: float
+ ) -> str:
+ """Create long format TextGrid (more readable)."""
+ lines = []
+ lines.append('File type = "ooTextFile"')
+ lines.append('Object class = "TextGrid"')
+ lines.append('')
+ lines.append(f'xmin = {min_time}')
+ lines.append(f'xmax = {max_time}')
+ lines.append('tiers? ')
+ lines.append(f'size = {len(tiers)}')
+ lines.append('item []:')
+
+ for i, tier_def in enumerate(tiers, 1):
+ tier_name = tier_def["name"]
+ tier_anns = annotations.get(tier_name, [])
+
+ # Sort annotations by start time
+ sorted_anns = sorted(tier_anns, key=lambda a: a.get("start_time", 0))
+
+ # Fill gaps to create complete intervals
+ intervals = self._create_intervals(sorted_anns, min_time, max_time)
+
+ lines.append(f' item [{i}]:')
+ lines.append(' class = "IntervalTier"')
+ lines.append(f' name = "{self._escape_text(tier_name)}"')
+ lines.append(f' xmin = {min_time}')
+ lines.append(f' xmax = {max_time}')
+ lines.append(f' intervals: size = {len(intervals)}')
+
+ for j, interval in enumerate(intervals, 1):
+ lines.append(f' intervals [{j}]:')
+ lines.append(f' xmin = {interval["start"]}')
+ lines.append(f' xmax = {interval["end"]}')
+ lines.append(f' text = "{self._escape_text(interval["text"])}"')
+
+ return '\n'.join(lines)
+
+ def _create_short_textgrid(
+ self,
+ tiers: List[dict],
+ annotations: Dict[str, List[dict]],
+ min_time: float,
+ max_time: float
+ ) -> str:
+ """Create short format TextGrid (more compact)."""
+ lines = []
+ lines.append('File type = "ooTextFile"')
+ lines.append('Object class = "TextGrid"')
+ lines.append('')
+ lines.append(str(min_time))
+ lines.append(str(max_time))
+ lines.append('')
+ lines.append(str(len(tiers)))
+
+ for tier_def in tiers:
+ tier_name = tier_def["name"]
+ tier_anns = annotations.get(tier_name, [])
+
+ # Sort and create intervals
+ sorted_anns = sorted(tier_anns, key=lambda a: a.get("start_time", 0))
+ intervals = self._create_intervals(sorted_anns, min_time, max_time)
+
+ lines.append('"IntervalTier"')
+ lines.append(f'"{self._escape_text(tier_name)}"')
+ lines.append(str(min_time))
+ lines.append(str(max_time))
+ lines.append(str(len(intervals)))
+
+ for interval in intervals:
+ lines.append(str(interval["start"]))
+ lines.append(str(interval["end"]))
+ lines.append(f'"{self._escape_text(interval["text"])}"')
+
+ return '\n'.join(lines)
+
+ def _create_intervals(
+ self,
+ annotations: List[dict],
+ min_time: float,
+ max_time: float
+ ) -> List[dict]:
+ """
+ Create a complete list of intervals, filling gaps with empty intervals.
+
+ Args:
+ annotations: Sorted list of annotations
+ min_time: Start time of the TextGrid
+ max_time: End time of the TextGrid
+
+ Returns:
+ List of interval dicts with start, end, and text
+ """
+ intervals = []
+ current_time = min_time
+
+ for ann in annotations:
+ start_sec = ann.get("start_time", 0) / 1000.0 # Convert ms to seconds
+ end_sec = ann.get("end_time", 0) / 1000.0
+ text = ann.get("value") or ann.get("label", "")
+
+ # Add empty interval for gap
+ if start_sec > current_time + 0.0001: # Small tolerance
+ intervals.append({
+ "start": current_time,
+ "end": start_sec,
+ "text": ""
+ })
+
+ # Add annotation interval
+ intervals.append({
+ "start": start_sec,
+ "end": end_sec,
+ "text": text
+ })
+ current_time = end_sec
+
+ # Add final empty interval if needed
+ if current_time < max_time - 0.0001:
+ intervals.append({
+ "start": current_time,
+ "end": max_time,
+ "text": ""
+ })
+
+ # If no intervals at all, create one empty interval
+ if not intervals:
+ intervals.append({
+ "start": min_time,
+ "end": max_time,
+ "text": ""
+ })
+
+ return intervals
+
+ def _get_max_time(self, annotations: Dict[str, List[dict]]) -> float:
+ """Get the maximum end time from all annotations in seconds."""
+ max_time = 0.0
+
+ for tier_anns in annotations.values():
+ for ann in tier_anns:
+ end_time = ann.get("end_time", 0)
+ if end_time:
+ max_time = max(max_time, end_time / 1000.0) # Convert ms to seconds
+
+ return max_time
+
+ def _escape_text(self, text: str) -> str:
+ """Escape special characters for TextGrid format."""
+ if not text:
+ return ""
+ # Escape quotes and backslashes
+ text = text.replace('\\', '\\\\')
+ text = text.replace('"', '\\"')
+ # Remove or replace newlines
+ text = text.replace('\n', ' ').replace('\r', '')
+ return text
diff --git a/potato/export/trajectory_correction_exporter.py b/potato/export/trajectory_correction_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..9887d75bcc6ab125eedfdce9853529e8146136f3
--- /dev/null
+++ b/potato/export/trajectory_correction_exporter.py
@@ -0,0 +1,216 @@
+"""
+Trajectory Correction Exporter
+
+Turns ``trajectory_edit`` annotations (human-corrected agent traces) into
+training-ready data:
+
+- ``trajectory_corrections.json`` โ full records with the original trace, the
+ reconstructed corrected trace, and per-field edit details.
+- ``trajectory_sft.jsonl`` โ one record per *edited* trace:
+ ``{"prompt": , "completion": }`` (SFT target).
+- ``trajectory_dpo.jsonl`` โ one record per *edited* trace:
+ ``{"prompt": , "chosen": , "rejected": }``.
+
+Unedited annotations are counted but never produce SFT/DPO records (no point
+training on an unchanged trajectory); the count of skipped/unedited traces is
+reported in ``stats`` and ``warnings`` so coverage is never silently dropped.
+"""
+
+import copy
+import json
+import logging
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+
+logger = logging.getLogger(__name__)
+
+
+class TrajectoryCorrectionExporter(BaseExporter):
+ """Exporter for trajectory_edit annotations โ SFT/DPO training data."""
+
+ format_name = "trajectory_correction"
+ description = "Corrected agent trajectories as SFT targets and DPO preference pairs"
+ file_extensions = [".json", ".jsonl"]
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ warnings: List[str] = []
+
+ try:
+ trajedit_schemas = {
+ s["name"]: s for s in context.schemas
+ if s.get("annotation_type") == "trajectory_edit"
+ }
+ if not trajedit_schemas:
+ return ExportResult(
+ success=False, format_name=self.format_name,
+ errors=["No trajectory_edit schemas defined"],
+ )
+
+ records: List[dict] = []
+ n_unedited = 0
+ n_unparseable = 0
+
+ for ann in context.annotations:
+ instance_id = ann.get("instance_id", "")
+ labels = ann.get("labels", {}) or {}
+ item = context.items.get(instance_id, {}) or {}
+
+ for schema_name, value in labels.items():
+ scheme = trajedit_schemas.get(schema_name)
+ if scheme is None:
+ continue
+
+ label_val = value.get("label", "") if isinstance(value, dict) else value
+ if not isinstance(label_val, str) or not label_val.strip():
+ continue
+ try:
+ correction = json.loads(label_val)
+ except (ValueError, TypeError):
+ n_unparseable += 1
+ continue
+
+ record = self._build_record(
+ instance_id, ann.get("user_id", ""), scheme, item, correction
+ )
+ if record is None:
+ continue
+ records.append(record)
+ if record["n_edits"] == 0:
+ n_unedited += 1
+
+ edited_records = [r for r in records if r["n_edits"] > 0]
+ if n_unedited:
+ warnings.append(
+ f"{n_unedited} annotation(s) had no edits โ included in "
+ f"corrections.json but excluded from SFT/DPO output."
+ )
+
+ os.makedirs(output_path, exist_ok=True)
+ files_written = []
+
+ corrections_file = os.path.join(output_path, "trajectory_corrections.json")
+ with open(corrections_file, "w", encoding="utf-8") as f:
+ json.dump({"records": records, "n_total": len(records),
+ "n_edited": len(edited_records),
+ "n_unedited": n_unedited}, f, indent=2, ensure_ascii=False)
+ files_written.append(corrections_file)
+
+ sft_file = os.path.join(output_path, "trajectory_sft.jsonl")
+ with open(sft_file, "w", encoding="utf-8") as f:
+ for r in edited_records:
+ f.write(json.dumps({
+ "prompt": r["task"],
+ "completion": r["corrected_trace"],
+ "trace_id": r["trace_id"],
+ }, ensure_ascii=False) + "\n")
+ files_written.append(sft_file)
+
+ dpo_file = os.path.join(output_path, "trajectory_dpo.jsonl")
+ with open(dpo_file, "w", encoding="utf-8") as f:
+ for r in edited_records:
+ f.write(json.dumps({
+ "prompt": r["task"],
+ "chosen": r["corrected_trace"],
+ "rejected": r["original_trace"],
+ "trace_id": r["trace_id"],
+ }, ensure_ascii=False) + "\n")
+ files_written.append(dpo_file)
+
+ return ExportResult(
+ success=True, format_name=self.format_name,
+ files_written=files_written, warnings=warnings,
+ stats={
+ "total_corrections": len(records),
+ "edited_traces": len(edited_records),
+ "unedited_traces": n_unedited,
+ "unparseable": n_unparseable,
+ },
+ )
+
+ except Exception as e:
+ logger.error(f"Trajectory correction export failed: {e}")
+ return ExportResult(
+ success=False, format_name=self.format_name, errors=[str(e)],
+ )
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ if not context.annotations:
+ return False, "No annotations to export"
+ if not any(s.get("annotation_type") == "trajectory_edit" for s in context.schemas):
+ return False, "No trajectory_edit schema defined"
+ return True, ""
+
+ def _build_record(self, instance_id: str, user_id: str, scheme: dict,
+ item: dict, correction: dict) -> Optional[dict]:
+ """Reconstruct the corrected trace from the original + per-field edits."""
+ steps_key = scheme.get("steps_key", "steps")
+ final_answer_key = scheme.get("final_answer_key", "final_answer")
+
+ original_steps = item.get(steps_key, [])
+ if not isinstance(original_steps, list):
+ original_steps = []
+
+ corrected_steps = copy.deepcopy(original_steps)
+ edits = correction.get("steps", []) or []
+ applied_edits = []
+
+ for e in edits:
+ if not e.get("edited"):
+ continue
+ idx = e.get("step_index")
+ field = e.get("field")
+ edited_text = e.get("edited_text", "")
+ if idx is None or field is None:
+ continue
+ if 0 <= idx < len(corrected_steps):
+ step = corrected_steps[idx]
+ if isinstance(step, dict):
+ step[field] = edited_text
+ elif isinstance(step, str) and field == scheme.get("step_text_key", "action"):
+ corrected_steps[idx] = edited_text
+ applied_edits.append({
+ "step_index": idx, "field": field,
+ "original_text": e.get("original_text", ""),
+ "edited_text": edited_text,
+ "edit_distance_chars": e.get("edit_distance_chars", 0),
+ "edit_distance_words": e.get("edit_distance_words", 0),
+ "reason": e.get("reason", ""),
+ })
+
+ # Final answer correction (optional)
+ original_final = item.get(final_answer_key)
+ corrected_final = original_final
+ fa = correction.get("final_answer")
+ if isinstance(fa, dict) and fa.get("edited"):
+ corrected_final = fa.get("edited_text", original_final)
+ applied_edits.append({
+ "step_index": None, "field": final_answer_key,
+ "original_text": fa.get("original_text", original_final),
+ "edited_text": corrected_final,
+ "edit_distance_chars": fa.get("edit_distance_chars", 0),
+ "edit_distance_words": fa.get("edit_distance_words", 0),
+ "reason": "",
+ })
+
+ task = item.get("task_description", item.get("task", item.get("text", "")))
+
+ original_trace = {"task": task, "steps": original_steps}
+ corrected_trace = {"task": task, "steps": corrected_steps}
+ if original_final is not None or corrected_final is not None:
+ original_trace[final_answer_key] = original_final
+ corrected_trace[final_answer_key] = corrected_final
+
+ return {
+ "trace_id": instance_id,
+ "annotator": user_id,
+ "task": task,
+ "original_trace": original_trace,
+ "corrected_trace": corrected_trace,
+ "edits": applied_edits,
+ "n_edits": len(applied_edits),
+ "total_edit_distance": correction.get("total_edit_distance", 0),
+ }
diff --git a/potato/export/yolo_exporter.py b/potato/export/yolo_exporter.py
new file mode 100644
index 0000000000000000000000000000000000000000..c24050700ca96638a761d929956971dea982b2d5
--- /dev/null
+++ b/potato/export/yolo_exporter.py
@@ -0,0 +1,173 @@
+"""
+YOLO Exporter
+
+Exports image annotations to YOLO format:
+- One .txt file per image with lines: class_id cx cy w h (normalized 0-1)
+- classes.txt listing class names
+- data.yaml for Ultralytics compatibility
+"""
+
+import os
+import logging
+from typing import Optional, Tuple
+
+from .base import BaseExporter, ExportContext, ExportResult
+from .cv_utils import (
+ build_category_mapping,
+ polygon_to_bbox,
+ normalize_bbox,
+ extract_image_annotations,
+ get_image_dimensions,
+ get_image_filename,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class YOLOExporter(BaseExporter):
+ format_name = "yolo"
+ description = "YOLO format for object detection (Ultralytics compatible)"
+ file_extensions = [".txt", ".yaml"]
+
+ def can_export(self, context: ExportContext) -> Tuple[bool, str]:
+ has_image_schema = any(
+ s.get("annotation_type") == "image_annotation"
+ for s in context.schemas
+ )
+ if not has_image_schema:
+ return False, "No image_annotation schema found in config"
+
+ # Check that we can get image dimensions
+ missing_dims = []
+ for ann in context.annotations:
+ instance_id = ann.get("instance_id", "")
+ item = context.items.get(instance_id, {})
+ img_anns = extract_image_annotations(ann)
+ if img_anns:
+ w, h = get_image_dimensions(item)
+ if w <= 0 or h <= 0:
+ missing_dims.append(instance_id)
+
+ if missing_dims:
+ return (
+ False,
+ f"YOLO requires image dimensions. Missing for: "
+ f"{', '.join(missing_dims[:5])}"
+ f"{'...' if len(missing_dims) > 5 else ''}"
+ )
+ return True, ""
+
+ def export(self, context: ExportContext, output_path: str,
+ options: Optional[dict] = None) -> ExportResult:
+ options = options or {}
+ warnings = []
+ files_written = []
+
+ category_map = build_category_mapping(context.annotations, context.schemas)
+ labels_dir = os.path.join(output_path, "labels")
+ os.makedirs(labels_dir, exist_ok=True)
+
+ # Track which images have been written (handle multiple annotators)
+ image_labels = {} # filename_stem -> list of label lines
+
+ for ann in context.annotations:
+ instance_id = ann.get("instance_id", "")
+ item = context.items.get(instance_id, {})
+ img_anns = extract_image_annotations(ann)
+ if not img_anns:
+ continue
+
+ img_w, img_h = get_image_dimensions(item)
+ if img_w <= 0 or img_h <= 0:
+ warnings.append(f"Skipping {instance_id}: no image dimensions")
+ continue
+
+ file_name = get_image_filename(item) or instance_id
+ raw_stem = os.path.splitext(os.path.basename(file_name))[0]
+ stem = "".join(c if c.isalnum() or c in "-_." else "_" for c in raw_stem)
+
+ if stem not in image_labels:
+ image_labels[stem] = []
+
+ for schema_name, objects in img_anns:
+ for obj in objects:
+ obj_type = obj.get("type", "")
+ label = obj.get("label", "")
+
+ if label not in category_map:
+ warnings.append(f"Unknown label '{label}' in {instance_id}")
+ continue
+
+ class_id = category_map[label]
+
+ if obj_type == "bbox":
+ x = obj.get("x", 0)
+ y = obj.get("y", 0)
+ w = obj.get("width", 0)
+ h = obj.get("height", 0)
+ cx, cy, nw, nh = normalize_bbox(x, y, w, h, img_w, img_h)
+ image_labels[stem].append(
+ f"{class_id} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}"
+ )
+
+ elif obj_type in ("polygon", "freeform"):
+ points = obj.get("points", [])
+ if not points:
+ continue
+ bx, by, bw, bh = polygon_to_bbox(points)
+ cx, cy, nw, nh = normalize_bbox(bx, by, bw, bh, img_w, img_h)
+ warnings.append(
+ f"{obj_type} in {instance_id} converted to enclosing bbox"
+ )
+ image_labels[stem].append(
+ f"{class_id} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}"
+ )
+
+ elif obj_type == "landmark":
+ warnings.append(
+ f"Landmark in {instance_id} skipped (not supported in YOLO)"
+ )
+
+ else:
+ warnings.append(
+ f"Unknown type '{obj_type}' in {instance_id}"
+ )
+
+ # Write label files
+ for stem, lines in image_labels.items():
+ label_file = os.path.join(labels_dir, f"{stem}.txt")
+ with open(label_file, "w") as f:
+ f.write("\n".join(lines))
+ if lines:
+ f.write("\n")
+ files_written.append(label_file)
+
+ # Write classes.txt
+ sorted_labels = sorted(category_map.items(), key=lambda kv: kv[1])
+ classes_file = os.path.join(output_path, "classes.txt")
+ with open(classes_file, "w") as f:
+ for name, _ in sorted_labels:
+ f.write(f"{name}\n")
+ files_written.append(classes_file)
+
+ # Write data.yaml for Ultralytics
+ data_yaml = os.path.join(output_path, "data.yaml")
+ with open(data_yaml, "w") as f:
+ f.write(f"path: {output_path}\n")
+ f.write("train: images/train\n")
+ f.write("val: images/val\n")
+ f.write(f"nc: {len(sorted_labels)}\n")
+ f.write(f"names: [{', '.join(repr(n) for n, _ in sorted_labels)}]\n")
+ files_written.append(data_yaml)
+
+ return ExportResult(
+ success=True,
+ format_name=self.format_name,
+ files_written=files_written,
+ warnings=warnings,
+ stats={
+ "num_images": len(image_labels),
+ "num_annotations": sum(len(v) for v in image_labels.values()),
+ "num_classes": len(sorted_labels),
+ },
+ )
diff --git a/potato/filter_by_annotation.py b/potato/filter_by_annotation.py
new file mode 100644
index 0000000000000000000000000000000000000000..9184fe33bbd327202aba5becd43f0d34ccbdc2e8
--- /dev/null
+++ b/potato/filter_by_annotation.py
@@ -0,0 +1,365 @@
+"""
+Filter Data by Prior Annotations
+
+This module provides functionality to filter data items based on prior annotation
+decisions. This is particularly useful for workflows like:
+
+1. Triage -> Full Annotation: Filter items that were "accepted" in triage
+2. Quality Control: Filter items that passed quality checks
+3. Multi-phase Annotation: Chain annotation tasks together
+
+Usage (CLI):
+ python -m potato.filter_by_annotation \\
+ --annotations annotation_output/ \\
+ --data data/items.json \\
+ --schema data_quality \\
+ --value accept \\
+ --output accepted_items.json
+
+Usage (Python):
+ from potato.filter_by_annotation import filter_items_by_annotation
+
+ filtered = filter_items_by_annotation(
+ annotation_dir="annotation_output/",
+ data_file="data/items.json",
+ schema_name="data_quality",
+ filter_value="accept",
+ id_key="id"
+ )
+"""
+
+import argparse
+import json
+import os
+import logging
+from pathlib import Path
+from typing import List, Dict, Any, Optional, Set, Union
+
+logger = logging.getLogger(__name__)
+
+
+def load_annotations_from_dir(annotation_dir: str) -> Dict[str, Dict[str, Any]]:
+ """
+ Load all annotations from an annotation output directory.
+
+ Args:
+ annotation_dir: Path to annotation_output directory
+
+ Returns:
+ Dict mapping instance_id -> {schema_name -> value}
+ """
+ annotations = {}
+ annotation_path = Path(annotation_dir)
+
+ if not annotation_path.exists():
+ logger.warning(f"Annotation directory does not exist: {annotation_dir}")
+ return annotations
+
+ # Look for user_state.json files in user subdirectories
+ for user_dir in annotation_path.iterdir():
+ if not user_dir.is_dir():
+ continue
+
+ state_file = user_dir / "user_state.json"
+ if not state_file.exists():
+ continue
+
+ try:
+ with open(state_file, 'r', encoding='utf-8') as f:
+ user_state = json.load(f)
+ except (json.JSONDecodeError, IOError) as e:
+ logger.warning(f"Failed to load {state_file}: {e}")
+ continue
+
+ # Extract label annotations
+ instance_labels = user_state.get("instance_id_to_label_to_value", {})
+
+ for instance_id, label_list in instance_labels.items():
+ if instance_id not in annotations:
+ annotations[instance_id] = {}
+
+ # label_list is a list of [label_dict, value] pairs
+ for label_entry in label_list:
+ if isinstance(label_entry, (list, tuple)) and len(label_entry) >= 2:
+ label_dict, value = label_entry[0], label_entry[1]
+ schema = label_dict.get("schema", "")
+ name = label_dict.get("name", "")
+
+ # For triage, the "name" is the decision (accept/reject/skip)
+ # Store both the name and raw value
+ if schema:
+ annotations[instance_id][schema] = {
+ "name": name,
+ "value": value
+ }
+
+ return annotations
+
+
+def load_data_file(data_file: str) -> List[Dict[str, Any]]:
+ """
+ Load data from a JSON or JSONL file.
+
+ Args:
+ data_file: Path to data file
+
+ Returns:
+ List of data items
+ """
+ data_path = Path(data_file)
+
+ if not data_path.exists():
+ raise FileNotFoundError(f"Data file not found: {data_file}")
+
+ with open(data_path, 'r', encoding='utf-8') as f:
+ content = f.read().strip()
+
+ # Try JSON array first
+ try:
+ data = json.loads(content)
+ if isinstance(data, list):
+ return data
+ elif isinstance(data, dict):
+ return [data]
+ except json.JSONDecodeError:
+ pass
+
+ # Try JSONL (newline-delimited JSON)
+ items = []
+ for line in content.split('\n'):
+ line = line.strip()
+ if line:
+ try:
+ items.append(json.loads(line))
+ except json.JSONDecodeError:
+ continue
+
+ return items
+
+
+def filter_items_by_annotation(
+ annotation_dir: str,
+ data_file: str,
+ schema_name: str,
+ filter_value: Union[str, List[str]],
+ id_key: str = "id",
+ invert: bool = False
+) -> List[Dict[str, Any]]:
+ """
+ Filter data items based on prior annotation decisions.
+
+ Args:
+ annotation_dir: Path to annotation_output directory
+ data_file: Path to original data file
+ schema_name: Name of the annotation schema to filter by
+ filter_value: Value(s) to filter for (e.g., "accept" or ["accept", "maybe"])
+ id_key: Key in data items containing the instance ID
+ invert: If True, return items that DON'T match the filter
+
+ Returns:
+ List of filtered data items
+ """
+ # Normalize filter_value to a set
+ if isinstance(filter_value, str):
+ filter_values = {filter_value}
+ else:
+ filter_values = set(filter_value)
+
+ # Load annotations
+ annotations = load_annotations_from_dir(annotation_dir)
+ logger.info(f"Loaded annotations for {len(annotations)} instances")
+
+ # Load data
+ data_items = load_data_file(data_file)
+ logger.info(f"Loaded {len(data_items)} data items")
+
+ # Filter items
+ filtered = []
+ for item in data_items:
+ instance_id = str(item.get(id_key, ""))
+
+ if not instance_id:
+ logger.warning(f"Item missing id_key '{id_key}': {item}")
+ continue
+
+ # Check if this instance has the annotation we're looking for
+ instance_annotations = annotations.get(instance_id, {})
+ schema_annotation = instance_annotations.get(schema_name, {})
+
+ # Get the annotation value (check both 'name' and 'value' fields)
+ anno_value = schema_annotation.get("name") or schema_annotation.get("value")
+
+ matches = anno_value in filter_values
+
+ if invert:
+ matches = not matches
+
+ if matches:
+ filtered.append(item)
+
+ logger.info(f"Filtered to {len(filtered)} items (schema={schema_name}, value={filter_values})")
+ return filtered
+
+
+def get_annotation_summary(annotation_dir: str, schema_name: str) -> Dict[str, int]:
+ """
+ Get a summary of annotation value counts for a schema.
+
+ Args:
+ annotation_dir: Path to annotation_output directory
+ schema_name: Name of the annotation schema
+
+ Returns:
+ Dict mapping value -> count
+ """
+ annotations = load_annotations_from_dir(annotation_dir)
+
+ counts = {}
+ for instance_id, schemas in annotations.items():
+ if schema_name in schemas:
+ value = schemas[schema_name].get("name") or schemas[schema_name].get("value")
+ if value:
+ counts[value] = counts.get(value, 0) + 1
+
+ return counts
+
+
+def main():
+ """CLI entry point."""
+ parser = argparse.ArgumentParser(
+ description="Filter data items based on prior annotation decisions",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+Examples:
+ # Filter for accepted items from triage
+ python -m potato.filter_by_annotation \\
+ --annotations annotation_output/ \\
+ --data data/items.json \\
+ --schema data_quality \\
+ --value accept \\
+ --output accepted_items.json
+
+ # Filter for multiple values
+ python -m potato.filter_by_annotation \\
+ --annotations annotation_output/ \\
+ --data data/items.json \\
+ --schema data_quality \\
+ --value accept maybe \\
+ --output filtered_items.json
+
+ # Show annotation summary
+ python -m potato.filter_by_annotation \\
+ --annotations annotation_output/ \\
+ --schema data_quality \\
+ --summary
+ """
+ )
+
+ parser.add_argument(
+ "--annotations", "-a",
+ required=True,
+ help="Path to annotation_output directory"
+ )
+ parser.add_argument(
+ "--data", "-d",
+ help="Path to original data file (JSON or JSONL)"
+ )
+ parser.add_argument(
+ "--schema", "-s",
+ required=True,
+ help="Name of the annotation schema to filter by"
+ )
+ parser.add_argument(
+ "--value", "-v",
+ nargs="+",
+ help="Value(s) to filter for (e.g., 'accept' or 'accept maybe')"
+ )
+ parser.add_argument(
+ "--output", "-o",
+ help="Output file path for filtered data"
+ )
+ parser.add_argument(
+ "--id-key",
+ default="id",
+ help="Key in data items containing the instance ID (default: 'id')"
+ )
+ parser.add_argument(
+ "--invert",
+ action="store_true",
+ help="Invert filter: return items that DON'T match"
+ )
+ parser.add_argument(
+ "--summary",
+ action="store_true",
+ help="Show annotation value summary instead of filtering"
+ )
+ parser.add_argument(
+ "--format",
+ choices=["json", "jsonl"],
+ default="json",
+ help="Output format (default: json)"
+ )
+ parser.add_argument(
+ "--verbose", "-V",
+ action="store_true",
+ help="Enable verbose logging"
+ )
+
+ args = parser.parse_args()
+
+ # Setup logging
+ logging.basicConfig(
+ level=logging.DEBUG if args.verbose else logging.INFO,
+ format="%(levelname)s: %(message)s"
+ )
+
+ # Summary mode
+ if args.summary:
+ counts = get_annotation_summary(args.annotations, args.schema)
+ if counts:
+ print(f"\nAnnotation summary for schema '{args.schema}':")
+ print("-" * 40)
+ total = sum(counts.values())
+ for value, count in sorted(counts.items(), key=lambda x: -x[1]):
+ pct = 100 * count / total
+ print(f" {value}: {count} ({pct:.1f}%)")
+ print("-" * 40)
+ print(f" Total: {total}")
+ else:
+ print(f"No annotations found for schema '{args.schema}'")
+ return
+
+ # Filter mode
+ if not args.data:
+ parser.error("--data is required for filtering (use --summary for summary mode)")
+ if not args.value:
+ parser.error("--value is required for filtering")
+ if not args.output:
+ parser.error("--output is required for filtering")
+
+ # Filter items
+ filtered = filter_items_by_annotation(
+ annotation_dir=args.annotations,
+ data_file=args.data,
+ schema_name=args.schema,
+ filter_value=args.value,
+ id_key=args.id_key,
+ invert=args.invert
+ )
+
+ # Write output
+ output_path = Path(args.output)
+ output_path.parent.mkdir(parents=True, exist_ok=True)
+
+ with open(output_path, 'w', encoding='utf-8') as f:
+ if args.format == "jsonl":
+ for item in filtered:
+ f.write(json.dumps(item) + "\n")
+ else:
+ json.dump(filtered, f, indent=2)
+
+ print(f"Wrote {len(filtered)} items to {args.output}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/potato/flask_server.py b/potato/flask_server.py
new file mode 100644
index 0000000000000000000000000000000000000000..d673d7f51ab4c02e58d481c89abaef71804ec9ab
--- /dev/null
+++ b/potato/flask_server.py
@@ -0,0 +1,4074 @@
+"""
+Flask Server Driver
+
+This module provides the main Flask server implementation for the annotation platform.
+Features include:
+- User authentication and session management
+- Annotation state tracking
+- Multi-phase annotation workflow
+- Survey flow support
+- Data loading and persistence
+- AI augmentation support
+- Active learning integration
+- Admin dashboard functionality
+
+The server handles:
+1. Data loading from various file formats (JSON, CSV, TSV, JSONL)
+2. User session management and authentication
+3. Annotation submission and validation
+4. Phase progression and workflow management
+5. AI hint generation and integration
+6. Active learning model training and instance reordering
+7. Admin dashboard data generation
+8. Configuration management and validation
+
+Key Components:
+- Flask application setup and configuration
+- Data loading and preprocessing
+- User state initialization
+- Annotation scheme processing
+- Template rendering and customization
+- Session timeout management
+- Error handling and logging
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+
+import logging
+import os
+import sys
+import random
+import json
+import re
+from collections import deque, defaultdict, Counter, OrderedDict
+from itertools import zip_longest
+import string
+import threading
+import yaml
+from datetime import datetime, timedelta
+from typing import List, Dict, Any
+
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+import simpledorff
+from simpledorff.metrics import nominal_metric, interval_metric
+
+import flask
+from flask import Flask, session, render_template, request, redirect, url_for, jsonify, make_response
+from bs4 import BeautifulSoup
+import shutil
+
+from dataclasses import dataclass
+
+# Get current working directory and program directory
+cur_working_dir = os.getcwd() #get the current working dir
+cur_program_dir = os.path.dirname(os.path.abspath(__file__)) #get the current program dir (for the case of pypi, it will be the path where potato is installed)
+flask_templates_dir = os.path.join(cur_program_dir,'templates') #get the dir where the flask templates are saved
+base_html_dir = os.path.join(cur_program_dir,'base_htmls') #get the dir where the the base_html templates files are saved
+
+#insert the current program dir into sys path
+sys.path.insert(0, cur_program_dir)
+
+from potato.item_state_management import ItemStateManager, Item, Label, SpanAnnotation
+from potato.item_state_management import get_item_state_manager, init_item_state_manager
+from potato.user_state_management import UserStateManager, UserState, get_user_state_manager, init_user_state_manager
+from potato.authentication import UserAuthenticator
+from potato.phase import UserPhase
+from potato.expertise_manager import init_expertise_manager, get_expertise_manager, clear_expertise_manager
+from potato.quality_control import (
+ init_quality_control_manager, get_quality_control_manager, clear_quality_control_manager
+)
+from potato.adjudication import (
+ init_adjudication_manager, get_adjudication_manager, clear_adjudication_manager
+)
+from potato.diversity_manager import (
+ init_diversity_manager, get_diversity_manager, clear_diversity_manager
+)
+from potato.knowledge_base import init_kb_manager
+
+from potato.solo_mode import init_solo_mode_manager, get_solo_mode_manager
+from potato.solo_mode.routes import solo_mode_bp
+
+from potato.qda_mode import init_qda_mode_manager
+
+from potato.create_task_cli import create_task_cli
+from potato.server_utils.arg_utils import arguments
+from potato.server_utils.config_module import init_config, config
+from potato.server_utils.schemas.span import render_span_annotations
+from potato.server_utils.prolific_apis import ProlificStudy
+from potato.server_utils.mturk_apis import init_mturk_hit, get_mturk_hit
+from potato.server_utils.json import easy_json
+from potato.server_utils.instance_display import InstanceDisplayRenderer, get_instance_display_renderer
+
+# This allows us to create an AI endpoint for the system to interact with as needed (if configured)
+from potato.ai.ai_endpoint import get_ai_endpoint
+
+# AI support initialization
+from potato.ai.ai_prompt import init_ai_prompt
+from potato.ai.ai_cache import init_ai_cache_manager, get_ai_cache_manager
+from potato.ai.ai_help_wrapper import init_dynamic_ai_help
+
+# Initialize Flask app
+app = Flask(__name__)
+app.register_blueprint(solo_mode_bp)
+from potato.judge_calibration.routes import judge_calibration_bp
+app.register_blueprint(judge_calibration_bp)
+# Note: qda_mode_bp is registered on the *served* app in
+# potato.routes.configure_routes (the module-level `app` here is discarded
+# and rebuilt by create_app), so it is intentionally not registered here.
+
+# Web agent recording and proxy blueprints (registered lazily in configure_app
+# only when web_agent display types are configured)
+
+# Secret key will be set in configure_app() from config
+
+# Use centralized logging configuration
+from potato.logging_config import get_logger, setup_logging
+logger = get_logger(__name__)
+
+_FRONTEND_TEMPLATE_TEXT_CACHE: dict[str, tuple[float, str]] = {}
+
+# Set random seed for reproducible behavior
+random.seed(0)
+
+# Global variables for file management and user tracking
+domain_file_path = ""
+file_list = []
+file_list_size = 0
+default_port = 8000
+user_dict = {}
+
+file_to_read_from = ""
+
+# User story position tracking and response queue management
+user_story_pos = defaultdict(lambda: 0, dict())
+user_response_dicts_queue = defaultdict(deque)
+
+# path to save user information
+USER_CONFIG_PATH = "user_config.json"
+DEFAULT_LABELS_PER_INSTANCE = 3
+
+# Hacky nonsense - schema label to color mapping
+schema_label_to_color = {}
+
+# Global Prolific study instance for API integration
+PROLIFIC_STUDY_INSTANCE = None
+
+# Keyword Highlights File Data
+@dataclass(frozen=True)
+class HighlightSchema:
+ """
+ Data class for highlight schema information.
+
+ This class represents a highlight schema with a label and schema name.
+ It's used for organizing highlight data and ensuring consistent
+ color assignments across the annotation interface.
+ """
+ label: str
+ schema: str
+
+ def __hash__(self):
+ return hash((self.label, self.schema))
+
+# Global emphasis corpus to schemas mapping
+emphasis_corpus_to_schemas = defaultdict(set)
+
+# Keyword highlight patterns loaded from TSV file
+# List of dicts: {pattern: str, regex: compiled_regex, label: str, schema: str}
+keyword_highlight_patterns = []
+
+# Keyword highlight settings (probabilities, random word config)
+# These control randomization and caching behavior for keyword highlights
+keyword_highlight_settings = {
+ 'keyword_probability': 1.0, # Probability of showing a matched keyword (0.0-1.0)
+ 'random_word_probability': 0.0, # Probability of highlighting random words (disabled by default)
+ 'random_word_label': 'distractor', # Label for random word highlights
+ 'random_word_schema': 'keyword', # Schema for random word highlights
+}
+
+def get_keyword_highlight_patterns():
+ """Get the current keyword highlight patterns list."""
+ logger.debug(f"[get_keyword_highlight_patterns] Returning {len(keyword_highlight_patterns)} patterns")
+ return keyword_highlight_patterns
+
+def get_keyword_highlight_settings():
+ """Get the current keyword highlight settings."""
+ return keyword_highlight_settings
+
+# Response Highlight Class
+@dataclass(frozen=True)
+class SuggestedResponse:
+ """
+ Data class for suggested response information.
+
+ This class represents a suggested response with a name and label.
+ It's used for AI-generated suggestions and pre-filled annotation values.
+ """
+ name: str
+ label: str
+
+ def __hash__(self):
+ return hash((self.name, self.label))
+
+# Color palette for annotation interface
+COLOR_PALETTE = [
+ "rgb(179,226,205)",
+ "rgb(253,205,172)",
+ "rgb(203,213,232)",
+ "rgb(244,202,228)",
+ "rgb(230,245,201)",
+ "rgb(255,242,174)",
+ "rgb(241,226,204)",
+ "rgb(204,204,204)",
+ "rgb(102, 197, 204)",
+ "rgb(246, 207, 113)",
+ "rgb(248, 156, 116)",
+ "rgb(220, 176, 242)",
+ "rgb(135, 197, 95)",
+ "rgb(158, 185, 243)",
+ "rgb(254, 136, 177)",
+ "rgb(201, 219, 116)",
+ "rgb(139, 224, 164)",
+ "rgb(180, 151, 231)",
+ "rgb(179, 179, 179)",
+]
+
+# Mapping the base html template str to the real file
+# REMOVED: template_dict is no longer needed since we use hardcoded template paths
+
+class ActiveLearningState:
+ """
+ A class for maintaining state on active learning.
+
+ This class tracks active learning selection types and update rounds
+ to ensure proper coordination between active learning cycles and
+ user assignment updates.
+ """
+
+ def __init__(self):
+ """Initialize the active learning state tracker."""
+ self.id_to_selection_type = {}
+ self.id_to_update_round = {}
+ self.cur_round = 0
+
+ def update_selection_types(self, id_to_selection_type):
+ """
+ Update the selection types for active learning.
+
+ Args:
+ id_to_selection_type: Dictionary mapping instance IDs to selection types
+ """
+ self.cur_round += 1
+
+ for iid, st in id_to_selection_type.items():
+ self.id_to_selection_type[iid] = st
+ self.id_to_update_round[iid] = self.cur_round
+
+# Set session timeout duration (e.g., 30 minutes)
+SESSION_TIMEOUT = timedelta(minutes=1)
+
+
+def _read_cached_template_text(path: str) -> str:
+ """Read a generated template file with a lightweight mtime cache."""
+ if not path:
+ return ""
+
+ try:
+ mtime = os.path.getmtime(path)
+ except OSError:
+ return ""
+
+ cached = _FRONTEND_TEMPLATE_TEXT_CACHE.get(path)
+ if cached and cached[0] == mtime:
+ return cached[1]
+
+ try:
+ with open(path, "rt", encoding="utf-8") as f:
+ text = f.read()
+ except OSError:
+ return ""
+
+ _FRONTEND_TEMPLATE_TEXT_CACHE[path] = (mtime, text)
+ return text
+
+
+def _resolve_generated_template_path(html_file: str) -> str:
+ """Resolve ``config['site_file']`` (a bare filename) to the absolute path of
+ the generated template on disk.
+
+ The generated template lives under ``/generated/``, but
+ ``config['site_file']`` is stored as just the filename (Jinja resolves it via
+ its template search path). The server ``chdir``s into ``task_dir`` at startup,
+ so reading the bare name with ``open()`` would fail and silently disable every
+ page-template-gated frontend asset. Resolving against the (absolute) site_dir
+ makes asset detection work regardless of the process CWD.
+ """
+ if not html_file:
+ return html_file
+ if os.path.isabs(html_file) and os.path.exists(html_file):
+ return html_file
+ site_dir = config.get("site_dir") or ""
+ candidate = os.path.join(site_dir, "generated", html_file)
+ if os.path.exists(candidate):
+ return candidate
+ return html_file
+
+
+# Authoritative mapping from asset key to the HTML markers that trigger loading.
+# Tests verify these markers appear in the actual schema/display generators,
+# so adding a new generator or renaming a CSS class will cause a test failure
+# rather than a silent asset-loading miss.
+FRONTEND_ASSET_MARKERS: dict[str, tuple[str, ...]] = {
+ "image_annotation": ("image-annotation-container",),
+ "audio_annotation": ("audio-annotation-container",),
+ "video_annotation": ("video-annotation-container",),
+ "span_link": ("span-link-container",),
+ "event_annotation": ("event-annotation-container",),
+ "coreference": ('data-annotation-type="coreference"', "coref-chain-panel"),
+ "conversation_tree": ("conv-tree",),
+ "tracking": ("tracking-panel", "tracking-overlay", "tracking-controls-group"),
+ "triage": ('class="annotation-form triage"', 'data-annotation-type="triage"', "triage-container"),
+ "tiered_annotation": ("tiered-annotation-container",),
+ "document_bbox": ("document-bbox-mode", "document-bbox-container", "document-bbox-canvas"),
+ "pdf_bbox": ("pdf-bbox-mode", "pdf-bbox-container", "pdf-bbox-canvas"),
+ "web_agent_viewer": ('class="web-agent-viewer"', 'class="live-agent-viewer"'),
+ "web_agent_playback": ('data-auto-playback="true"',),
+ "web_agent_recorder": ("web-agent-recorder",),
+ "live_coding_agent": ("live-coding-agent-viewer",),
+}
+
+
+def _detect_frontend_assets_for_page(html_file: str, display_html: str = "") -> dict[str, bool]:
+ """
+ Detect which frontend assets are needed for the current page only.
+
+ This avoids loading every specialized bundle just because some other phase
+ in the overall task config happens to use it.
+ """
+ page_html = _read_cached_template_text(_resolve_generated_template_path(html_file))
+ combined_html = f"{page_html}\n{display_html or ''}"
+
+ def has_any(*markers: str) -> bool:
+ return any(marker in combined_html for marker in markers)
+
+ detected = {key: has_any(*markers) for key, markers in FRONTEND_ASSET_MARKERS.items()}
+
+ # segmentation_tools is an alias โ loaded whenever image_annotation is present
+ detected["segmentation_tools"] = detected["image_annotation"]
+ # span_link also loads when coreference is present
+ detected["span_link"] = detected["span_link"] or detected["coreference"]
+
+ return detected
+
+
+def _apply_annotation_filter(items: list, filter_config: dict, id_key: str) -> list:
+ """
+ Filter items based on prior annotation decisions.
+
+ This enables chaining annotation tasks, e.g., triage -> full annotation.
+
+ Args:
+ items: List of data items to filter
+ filter_config: Configuration dict with:
+ - annotation_dir: Path to annotation_output directory
+ - schema: Name of the annotation schema to filter by
+ - value: Value(s) to filter for (string or list)
+ - invert: If True, return items that DON'T match (optional)
+ id_key: Key in items containing the instance ID
+
+ Returns:
+ Filtered list of items
+ """
+ from potato.filter_by_annotation import load_annotations_from_dir
+
+ annotation_dir = filter_config.get("annotation_dir")
+ schema_name = filter_config.get("schema")
+ filter_value = filter_config.get("value")
+ invert = filter_config.get("invert", False)
+
+ if not annotation_dir:
+ logger.warning("filter_by_prior_annotation missing 'annotation_dir', skipping filter")
+ return items
+ if not schema_name:
+ logger.warning("filter_by_prior_annotation missing 'schema', skipping filter")
+ return items
+ if not filter_value:
+ logger.warning("filter_by_prior_annotation missing 'value', skipping filter")
+ return items
+
+ # Normalize filter_value to a set
+ if isinstance(filter_value, str):
+ filter_values = {filter_value}
+ else:
+ filter_values = set(filter_value)
+
+ # Load prior annotations
+ annotations = load_annotations_from_dir(annotation_dir)
+ logger.debug(f"Loaded prior annotations for {len(annotations)} instances")
+
+ # Filter items
+ filtered = []
+ for item in items:
+ instance_id = str(item.get(id_key, ""))
+ if not instance_id:
+ continue
+
+ # Check if this instance has the annotation we're looking for
+ instance_annotations = annotations.get(instance_id, {})
+ schema_annotation = instance_annotations.get(schema_name, {})
+
+ # Get the annotation value
+ anno_value = schema_annotation.get("name") or schema_annotation.get("value")
+ matches = anno_value in filter_values
+
+ if invert:
+ matches = not matches
+
+ if matches:
+ filtered.append(item)
+
+ return filtered
+
+
+def load_instance_data(config: dict):
+ """
+ Load instance data from the files specified in the config.
+
+ This function reads annotation data from various file formats (JSON, CSV, TSV, JSONL)
+ and populates the ItemStateManager with the data. It handles different data structures
+ and validates that required fields are present.
+
+ Supports multiple data loading modes:
+ 1. data_files: List of local file paths (traditional mode)
+ 2. data_sources: Extended sources including URLs, cloud storage, databases
+ 3. data_directory: Watch a directory for files (handled separately)
+
+ Args:
+ config: Configuration dictionary containing data file paths and item properties
+
+ Side Effects:
+ - Populates ItemStateManager with loaded data
+ - Validates data structure and required fields
+ - Logs loading progress and statistics
+
+ Raises:
+ Exception: If file format is unsupported or required fields are missing
+ """
+ ism = get_item_state_manager()
+
+ # Where to look in the JSON item object for the text to annotate
+ text_key = config["item_properties"]["text_key"]
+ id_key = config["item_properties"]["id_key"]
+
+ # Check if data_sources is configured (new extended data loading)
+ if config.get("data_sources"):
+ _load_from_data_sources(config, ism, id_key, text_key)
+ return
+
+ data_files = config.get("data_files", [])
+ if not data_files:
+ # No data_files, might use data_directory which is handled elsewhere
+ logger.debug("No data_files configured, skipping file-based loading")
+ return
+
+ logger.debug("Loading data from %d files" % (len(data_files)))
+
+ for data_file_entry in data_files:
+ # Support both string paths and dict configs
+ if isinstance(data_file_entry, dict):
+ data_fname = data_file_entry.get("path")
+ filter_config = data_file_entry.get("filter_by_prior_annotation")
+ encoding = data_file_entry.get("encoding", "utf-8")
+ else:
+ data_fname = data_file_entry
+ filter_config = None
+ encoding = "utf-8"
+
+ if not data_fname:
+ logger.warning(f"Skipping data_files entry with no path: {data_file_entry}")
+ continue
+ fmt = data_fname.split(".")[-1]
+ if fmt not in ["csv", "tsv", "json", "jsonl", "parquet"]:
+ raise Exception("Unsupported input file format %s for %s" % (fmt, data_fname))
+
+ logger.debug("Reading data from " + data_fname)
+
+ if fmt in ["json", "jsonl"]:
+ # Handle JSON and JSONL formats
+ # Try parsing as a JSON array first, fall back to JSON Lines
+ with open(data_fname, "rt", encoding=encoding) as f:
+ raw = f.read()
+
+ items = None
+ if fmt == "json":
+ try:
+ parsed = json.loads(raw)
+ if isinstance(parsed, list):
+ items = parsed
+ logger.debug(f"Parsed {data_fname} as JSON array with {len(items)} items")
+ except json.JSONDecodeError:
+ pass # Fall through to JSON Lines parsing
+
+ if items is None:
+ # Parse as JSON Lines (one JSON object per line)
+ items = []
+ for line_no, line in enumerate(raw.splitlines()):
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ items.append(json.loads(line))
+ except json.JSONDecodeError as e:
+ raise ValueError(
+ f"Invalid JSON at line {line_no+1} in {data_fname}: {e}"
+ ) from e
+
+ # Apply filter_by_prior_annotation if configured
+ if filter_config:
+ items = _apply_annotation_filter(items, filter_config, id_key)
+ logger.info(f"Filtered to {len(items)} items based on prior annotations")
+
+ for item_no, item in enumerate(items):
+ if not isinstance(item, dict):
+ raise ValueError(f"Expected JSON object at item {item_no+1} in {data_fname}, got {type(item).__name__}")
+
+ # Validate that the ID key exists in the item
+ if id_key not in item:
+ raise KeyError(f"ID key '{id_key}' not found in item {item_no+1}")
+
+ instance_id = str(item[id_key]) # Ensure ID is string
+
+ # Check for duplicate IDs
+ if ism.has_item(instance_id):
+ raise ValueError(f"Duplicate instance ID '{instance_id}' found at item {item_no+1}")
+
+ # Validate text key exists if required
+ if text_key not in item:
+ logger.warning(f"Text key '{text_key}' not found in item with ID '{instance_id}'")
+
+ ism.add_item(instance_id, item)
+
+ line_no = len(items)
+ elif fmt == "parquet":
+ import pyarrow.parquet as pq
+
+ table = pq.read_table(data_fname)
+ df = table.to_pandas()
+
+ if id_key not in df.columns:
+ raise KeyError(f"ID column '{id_key}' not found in file {data_fname}")
+ if text_key not in df.columns:
+ logger.warning(f"Text column '{text_key}' not found in file {data_fname}")
+
+ df[id_key] = df[id_key].astype(str)
+
+ if df[id_key].duplicated().any():
+ dupes = df[id_key][df[id_key].duplicated()].tolist()
+ raise ValueError(f"Duplicate instance IDs found in {data_fname}: {dupes}")
+
+ existing_dupes = [id for id in df[id_key] if ism.has_item(id)]
+ if existing_dupes:
+ raise ValueError(f"Instance IDs in {data_fname} conflict with existing IDs: {existing_dupes}")
+
+ if text_key in df.columns:
+ df = df.astype({text_key: str})
+
+ items = df.to_dict('records')
+
+ if filter_config:
+ items = _apply_annotation_filter(items, filter_config, id_key)
+ logger.info(f"Filtered to {len(items)} items based on prior annotations")
+
+ for item in items:
+ instance_id = item[id_key]
+ ism.add_item(instance_id, item)
+
+ line_no = len(items)
+ else:
+ sep = "," if fmt == "csv" else "\t"
+
+ # Validate required columns exist
+ df = pd.read_csv(data_fname, sep=sep, encoding=encoding)
+ if id_key not in df.columns:
+ raise KeyError(f"ID column '{id_key}' not found in file {data_fname}")
+ if text_key not in df.columns:
+ logger.warning(f"Text column '{text_key}' not found in file {data_fname}")
+
+ # Convert ID column to string to ensure consistent typing
+ df[id_key] = df[id_key].astype(str)
+
+ # Check for duplicate IDs in the dataframe
+ if df[id_key].duplicated().any():
+ dupes = df[id_key][df[id_key].duplicated()].tolist()
+ raise ValueError(f"Duplicate instance IDs found in {data_fname}: {dupes}")
+
+ # Check for duplicate IDs with existing items
+ existing_dupes = [id for id in df[id_key] if ism.has_item(id)]
+ if existing_dupes:
+ raise ValueError(f"Instance IDs in {data_fname} conflict with existing IDs: {existing_dupes}")
+
+ # Load data with proper type conversion
+ df = df.astype({id_key: str})
+ if text_key in df.columns:
+ df = df.astype({text_key: str})
+
+ # Convert to list of dicts for filtering
+ items = df.to_dict('records')
+
+ # Apply filter_by_prior_annotation if configured
+ if filter_config:
+ items = _apply_annotation_filter(items, filter_config, id_key)
+ logger.info(f"Filtered to {len(items)} items based on prior annotations")
+
+ # Add items to state manager
+ for item in items:
+ instance_id = item[id_key]
+ ism.add_item(instance_id, item)
+
+ line_no = len(items)
+
+ # If the admin didn't specify a subset, have the user annotate all instances
+ # (or unlimited when a dynamic source can add more at runtime โ see F-037).
+ max_annotations_per_user = _default_max_annotations_per_user(config, ism)
+ get_user_state_manager().set_max_annotations_per_user(max_annotations_per_user)
+
+ logger.debug("Loaded %d instances from %s" % (line_no, data_fname))
+
+ # If BWS config is present, generate tuples from pool items
+ bws_config = config.get("bws_config")
+ if bws_config:
+ from potato.bws_tuple_generator import BwsTupleGenerator
+
+ # Collect all loaded pool items
+ pool_items = [item.get_data() for item in ism.items()]
+
+ # Store pool items for scoring later
+ config["_bws_pool_items"] = [dict(item) for item in pool_items]
+
+ generator = BwsTupleGenerator(
+ pool_items=pool_items,
+ id_key=id_key,
+ text_key=text_key,
+ tuple_size=bws_config.get("tuple_size", 4),
+ num_tuples=bws_config.get("num_tuples"),
+ seed=bws_config.get("seed", 42),
+ min_item_appearances=bws_config.get("min_item_appearances"),
+ )
+ generator.validate()
+ tuples = generator.generate()
+
+ # Clear pool items and replace with generated tuples
+ ism.clear()
+ for t in tuples:
+ ism.add_item(str(t[id_key]), t)
+
+ # Update max annotations per user for the new tuple count
+ max_annotations_per_user = config.get(
+ "max_annotations_per_user", len(ism.get_instance_ids())
+ )
+ get_user_state_manager().set_max_annotations_per_user(max_annotations_per_user)
+
+ logger.info(f"BWS: Replaced {len(pool_items)} pool items with {len(tuples)} tuples")
+
+ # If IBWS config is present, initialize iterative BWS manager and generate round 1 tuples
+ ibws_config = config.get("ibws_config")
+ if ibws_config:
+ from potato.ibws_manager import init_ibws_manager
+
+ # Collect all loaded pool items
+ pool_items = [item.get_data() for item in ism.items()]
+
+ # Store pool items for scoring
+ config["_bws_pool_items"] = [dict(item) for item in pool_items]
+
+ # Initialize IBWS manager
+ ibws_mgr = init_ibws_manager(config, pool_items, id_key, text_key)
+
+ # Generate round 1 tuples
+ round1_tuples = ibws_mgr.generate_round_tuples()
+
+ # Clear pool items and replace with round 1 tuples
+ ism.clear()
+ for t in round1_tuples:
+ ism.add_item(str(t[id_key]), t)
+
+ # Set unlimited annotations โ IBWS manager controls completion
+ get_user_state_manager().set_max_annotations_per_user(-1)
+
+ logger.info(
+ f"IBWS: Initialized with {len(pool_items)} pool items, "
+ f"generated {len(round1_tuples)} round-1 tuples"
+ )
+
+ # For each item, render the text to display in the UI ahead of time.
+ _render_displayed_text(text_key)
+
+
+def _default_max_annotations_per_user(config: dict, ism) -> int:
+ """
+ Resolve the default per-user annotation quota.
+
+ When ``max_annotations_per_user`` is explicitly configured, honor it.
+ Otherwise the historical default is "annotate everything" = the instance
+ count. But that count is frozen at load time, so when a DYNAMIC data source
+ can add items at runtime (trace ingestion, directory watching), freezing the
+ cap means later-added items exceed every user's quota and are never assigned
+ to any annotator (F-037). In that case default to unlimited (-1) instead, so
+ the live ``remaining_instance_ids`` pool stays fully assignable.
+ """
+ configured = config.get("max_annotations_per_user")
+ if configured is not None:
+ return configured
+ dynamic_source = bool(
+ (config.get("trace_ingestion") or {}).get("enabled")
+ or config.get("watch_data_directory")
+ )
+ if dynamic_source:
+ return -1
+ return len(ism.get_instance_ids())
+
+
+def _render_displayed_text(text_key: str) -> None:
+ """
+ Render the displayed text for all items.
+
+ This processes the text_key field to generate the displayed_text
+ that will be shown in the annotation UI.
+
+ Args:
+ text_key: The key in item data containing the text to display
+ """
+ for item in get_item_state_manager().items():
+ item_data = item.get_data()
+
+ # Validate text key exists before rendering
+ if text_key in item_data:
+ item_data["displayed_text"] = get_displayed_text(item_data[text_key])
+ else:
+ item_data["displayed_text"] = ""
+ logger.warning(f"No text found for item {item.get_id()}, using empty string")
+
+
+def _load_from_data_sources(config: dict, ism, id_key: str, text_key: str) -> None:
+ """
+ Load data using the extended DataSourceManager.
+
+ This function initializes the DataSourceManager and loads data from
+ configured sources (URLs, cloud storage, databases, etc.).
+
+ Args:
+ config: Application configuration
+ ism: ItemStateManager instance
+ id_key: Key for item IDs
+ text_key: Key for text content
+ """
+ # Import and register source implementations
+ from potato.data_sources import init_data_source_manager, get_data_source_manager
+ import potato.data_sources.sources # This registers all source types
+
+ # Initialize the data source manager
+ manager = init_data_source_manager(config)
+
+ if not manager:
+ logger.warning("DataSourceManager initialization failed")
+ return
+
+ # Load initial data from all sources
+ total_loaded = manager.load_initial_data()
+ logger.info(f"Loaded {total_loaded} items from data sources")
+
+ # Set max annotations per user (unlimited when a dynamic source โ e.g. a
+ # watched directory or trace ingestion โ can add items at runtime; F-037).
+ max_annotations_per_user = _default_max_annotations_per_user(config, ism)
+ get_user_state_manager().set_max_annotations_per_user(max_annotations_per_user)
+
+ # Render displayed text for all loaded items
+ _render_displayed_text(text_key)
+
+
+def load_user_data(config: dict):
+
+ user_data_dir = config['output_annotation_dir']
+ usm = get_user_state_manager()
+
+ # Check if the output directory exists
+ if not os.path.exists(user_data_dir):
+ os.makedirs(user_data_dir)
+ logger.info("Created output directory: %s" % user_data_dir)
+ return
+
+ # For each user's directory, load in their state
+ user_dirs = [d for d in os.listdir(user_data_dir) if os.path.isdir(os.path.join(user_data_dir, d))]
+
+ for user_dir in user_dirs:
+ try:
+ usm.load_user_state(os.path.join(user_data_dir, user_dir))
+ except ValueError as e:
+ # Skip directories that don't have valid user state files
+ logger.warning("Skipping invalid user directory %s: %s" % (user_dir, str(e)))
+ continue
+
+ # Rebuild instance_annotators from loaded user state so that
+ # adjudication build_queue() (and other code that relies on
+ # ism.instance_annotators) works with pre-loaded annotation data.
+ ism = get_item_state_manager()
+ for user_id in usm.get_user_ids():
+ user_state = usm.get_user_state(user_id)
+ if user_state:
+ for instance_id in user_state.instance_id_to_label_to_value:
+ if instance_id in ism.instance_id_to_instance:
+ ism.register_annotator(instance_id, user_id)
+ for instance_id in user_state.instance_id_to_span_to_value:
+ if instance_id in ism.instance_id_to_instance:
+ ism.register_annotator(instance_id, user_id)
+
+ logger.info("Loaded user data for %d users" % len(usm.get_user_ids()))
+
+def load_training_data(config: dict) -> None:
+ """
+ Load training data from the training data file specified in the config.
+
+ This function loads training instances with correct answers and explanations
+ for the training phase. It validates the training data format and stores
+ the training instances for use during the training phase.
+
+ Args:
+ config: Configuration dictionary containing training settings
+
+ Side Effects:
+ - Stores training instances in global training data storage
+ - Validates training data format and consistency
+ - Logs loading progress and statistics
+
+ Raises:
+ Exception: If training data file is not found or invalid
+ """
+ if 'training' not in config or not config['training'].get('enabled', False):
+ logger.debug("Training not enabled, skipping training data loading")
+ return
+
+ training_config = config['training']
+ data_file = training_config.get('data_file')
+
+ if not data_file:
+ logger.warning("Training enabled but no data_file specified")
+ return
+
+ # Resolve the training data file path
+ try:
+ training_data_path = get_abs_or_rel_path(data_file, config)
+ except FileNotFoundError:
+ logger.error(f"Training data file not found: {data_file}")
+ raise Exception(f"Training data file not found: {data_file}")
+
+ logger.debug(f"Loading training data from {training_data_path}")
+
+ try:
+ with open(training_data_path, 'r', encoding='utf-8') as f:
+ training_data = json.load(f)
+ except (json.JSONDecodeError, UnicodeDecodeError) as e:
+ logger.error(f"Invalid training data file format: {e}")
+ raise Exception(f"Invalid training data file format: {e}")
+
+ if not isinstance(training_data, dict):
+ raise Exception("Training data must be a JSON object")
+
+ if 'training_instances' not in training_data:
+ raise Exception("Training data must contain 'training_instances' field")
+
+ training_instances = training_data['training_instances']
+ if not isinstance(training_instances, list):
+ raise Exception("training_instances must be a list")
+
+ if not training_instances:
+ raise Exception("training_instances cannot be empty")
+
+ # Validate training data against annotation schemes
+ annotation_schemes = training_config.get('annotation_schemes', config.get('annotation_schemes', []))
+
+ # Handle both string references and full scheme dictionaries
+ scheme_names = set()
+ for scheme in annotation_schemes:
+ if isinstance(scheme, str):
+ # String reference to existing scheme
+ scheme_names.add(scheme)
+ elif isinstance(scheme, dict) and 'name' in scheme:
+ # Full scheme dictionary
+ scheme_names.add(scheme['name'])
+ else:
+ logger.warning(f"Invalid annotation scheme format: {scheme}")
+
+ # Convert training instances to Item objects and store them
+ global training_items
+ training_items = []
+
+ for instance in training_instances:
+ # Validate required fields
+ if 'id' not in instance or 'text' not in instance or 'correct_answers' not in instance:
+ raise Exception(f"Training instance missing required fields: {instance}")
+
+ # Validate correct_answers correspond to annotation schemes
+ for scheme_name in instance['correct_answers'].keys():
+ if scheme_name not in scheme_names:
+ logger.warning(f"Training instance {instance['id']} contains unknown scheme: {scheme_name}")
+
+ # Normalize category field (can be string or list)
+ category_value = instance.get('category')
+ if category_value is not None:
+ if isinstance(category_value, str):
+ categories = [category_value]
+ elif isinstance(category_value, list):
+ categories = [c for c in category_value if isinstance(c, str) and c.strip()]
+ else:
+ logger.warning(f"Training instance {instance['id']} has invalid category type: {type(category_value)}")
+ categories = []
+ else:
+ categories = []
+
+ # Create Item object for training instance
+ item_data = {
+ 'id': instance['id'],
+ 'text': instance['text'],
+ 'correct_answers': instance['correct_answers'],
+ 'explanation': instance.get('explanation', ''),
+ 'displayed_text': get_displayed_text(instance['text']),
+ 'categories': categories # Store normalized categories list
+ }
+
+ training_item = Item(instance['id'], item_data)
+ training_items.append(training_item)
+
+ logger.info(f"Loaded {len(training_items)} training instances")
+ logger.debug(f"Training instances: {[item.get_id() for item in training_items]}")
+
+
+def get_training_instances() -> List[Item]:
+ """
+ Get the loaded training instances.
+
+ Returns:
+ List of training Item objects
+ """
+ global training_items
+ # Bridge the __main__ vs potato.flask_server module split. When the server is
+ # launched from source via `python potato/flask_server.py start ...`, this file
+ # executes in the `__main__` namespace, so load_training_data's
+ # `global training_items` binds `__main__.training_items`. But routes.py and
+ # user_state_management.py call this function via
+ # `from potato.flask_server import get_training_instances`, which is the
+ # *potato.flask_server* copy whose own global is never set โ previously
+ # leaving the training phase with "No training instance available". Scan the
+ # candidate module namespaces so the data is found regardless of launch style.
+ for _mod_name in (__name__, 'potato.flask_server', '__main__'):
+ _mod = sys.modules.get(_mod_name)
+ _items = getattr(_mod, 'training_items', None) if _mod is not None else None
+ if _items:
+ return _items
+ return training_items if 'training_items' in globals() else []
+
+
+def get_training_correct_answers(instance_id: str) -> Dict[str, Any]:
+ """
+ Get the correct answers for a training instance.
+
+ Args:
+ instance_id: The ID of the training instance
+
+ Returns:
+ Dictionary of correct answers for the instance
+ """
+ training_items = get_training_instances()
+ for item in training_items:
+ if item.get_id() == instance_id:
+ return item.get_data().get('correct_answers', {})
+ return {}
+
+
+def get_training_explanation(instance_id: str) -> str:
+ """
+ Get the explanation for a training instance.
+
+ Args:
+ instance_id: The ID of the training instance
+
+ Returns:
+ Explanation string for the instance
+ """
+ training_items = get_training_instances()
+ for item in training_items:
+ if item.get_id() == instance_id:
+ return item.get_data().get('explanation', '')
+ return ''
+
+
+def get_training_instance_categories(instance_id: str) -> List[str]:
+ """
+ Get the categories for a training instance.
+
+ Args:
+ instance_id: The ID of the training instance
+
+ Returns:
+ List of category names (empty list if no categories)
+ """
+ training_items = get_training_instances()
+ for item in training_items:
+ if item.get_id() == instance_id:
+ return item.get_data().get('categories', [])
+ return []
+
+
+# =============================================================================
+# Prolific Integration Functions
+# =============================================================================
+
+def init_prolific_study(config: dict) -> None:
+ """
+ Initialize the Prolific study instance from config.
+
+ This function reads the Prolific configuration and initializes the
+ ProlificStudy API wrapper for tracking participants and managing
+ study status.
+
+ Args:
+ config: The application configuration dictionary
+
+ Side Effects:
+ - Sets global PROLIFIC_STUDY_INSTANCE
+ - May start workload checker thread
+ """
+ global PROLIFIC_STUDY_INSTANCE
+
+ prolific_config = config.get('prolific', {})
+ if not prolific_config:
+ logger.debug("No Prolific configuration found")
+ return
+
+ # Check for config file path
+ config_file_path = prolific_config.get('config_file_path')
+ if config_file_path:
+ # Load Prolific config from file
+ import yaml
+ prolific_config_path = get_abs_or_rel_path(config_file_path, config)
+ if os.path.exists(prolific_config_path):
+ with open(prolific_config_path, 'r', encoding='utf-8') as f:
+ prolific_settings = yaml.safe_load(f)
+ logger.info(f"Loaded Prolific config from {prolific_config_path}")
+ else:
+ logger.warning(f"Prolific config file not found: {prolific_config_path}")
+ return
+ else:
+ # Use inline config
+ prolific_settings = prolific_config
+
+ # Validate required fields
+ token = prolific_settings.get('token')
+ study_id = prolific_settings.get('study_id')
+
+ if not token or not study_id:
+ logger.warning("Prolific config missing 'token' or 'study_id'")
+ return
+
+ # Get optional settings
+ max_concurrent_sessions = prolific_settings.get('max_concurrent_sessions', 30)
+ workload_checker_period = prolific_settings.get('workload_checker_period', 60)
+
+ # Get saving directory for submission data
+ saving_dir = config.get('output_annotation_dir', 'annotation_output')
+
+ try:
+ PROLIFIC_STUDY_INSTANCE = ProlificStudy(
+ token=token,
+ study_id=study_id,
+ saving_dir=saving_dir,
+ max_concurrent_sessions=max_concurrent_sessions,
+ workload_checker_period=workload_checker_period
+ )
+ logger.info(f"Initialized Prolific study: {study_id}")
+ logger.info(f"Study info: {PROLIFIC_STUDY_INSTANCE.get_basic_study_info()}")
+
+ except Exception as e:
+ logger.error(f"Failed to initialize Prolific study: {e}")
+ PROLIFIC_STUDY_INSTANCE = None
+
+
+def get_prolific_study() -> 'ProlificStudy':
+ """
+ Get the global Prolific study instance.
+
+ Returns:
+ ProlificStudy instance if configured, None otherwise
+ """
+ global PROLIFIC_STUDY_INSTANCE
+ # Same __main__ vs potato.flask_server split as get_training_instances (F-044):
+ # init_prolific_study reassigns this global, so under `python flask_server.py`
+ # it lands on __main__ while routes.py reads the potato.flask_server copy (None).
+ # Scan candidate namespaces so a configured study is found regardless of launch.
+ for _mod_name in (__name__, 'potato.flask_server', '__main__'):
+ _mod = sys.modules.get(_mod_name)
+ _inst = getattr(_mod, 'PROLIFIC_STUDY_INSTANCE', None) if _mod is not None else None
+ if _inst is not None:
+ return _inst
+ return PROLIFIC_STUDY_INSTANCE
+
+
+def _prefill_diversity_embeddings(dm, config: dict) -> None:
+ """
+ Prefill embeddings for diversity ordering with progress bar.
+
+ Args:
+ dm: DiversityManager instance
+ config: Application configuration
+ """
+ from tqdm import tqdm
+
+ ism = get_item_state_manager()
+ text_key = config.get("item_properties", {}).get("text_key", "text")
+
+ # Collect texts for prefill
+ items = list(ism.items())[:dm.config.prefill_count]
+ texts = {}
+
+ for item in items:
+ item_data = item.get_data()
+ text = item_data.get(text_key, item.get_text())
+ texts[item.get_id()] = text
+
+ if not texts:
+ return
+
+ print(f"Prefilling {len(texts)} embeddings for diversity ordering...")
+
+ # Track progress with tqdm
+ completed = [0]
+
+ def on_complete(iid, emb):
+ completed[0] += 1
+
+ with tqdm(total=len(texts), desc="Computing embeddings", unit="item") as pbar:
+ # Compute in batches
+ batch_size = dm.config.batch_size
+ ids = list(texts.keys())
+
+ for i in range(0, len(ids), batch_size):
+ batch_ids = ids[i:i + batch_size]
+ batch_texts = {iid: texts[iid] for iid in batch_ids}
+ dm.compute_embeddings_batch(batch_texts, callback=on_complete)
+ pbar.update(len(batch_ids))
+
+ # Run clustering after prefill
+ if dm.cluster_items():
+ stats = dm.get_stats()
+ logger.info(
+ f"Clustered {stats['embedding_count']} items into "
+ f"{stats['cluster_count']} clusters"
+ )
+
+
+def load_all_data(config: dict):
+ '''Loads instance and annotation data from the files specified in the config.'''
+ load_annotation_schematic_data(config)
+ load_instance_data(config)
+ # Stamp per-item annotator caps for the overlap sample (must run before
+ # user_data so that initial assignments see the heterogeneous caps).
+ try:
+ from potato.server_utils.overlap_sampler import apply_overlap_sample
+ sampled = apply_overlap_sample(get_item_state_manager(), config)
+ if sampled:
+ logger.info("Overlap sampling stamped %d items", len(sampled))
+ except Exception as exc:
+ logger.warning("Overlap sampling skipped due to error: %s", exc)
+ load_user_data(config)
+ load_phase_data(config)
+ load_highlights_data(config)
+ load_training_data(config)
+ init_prolific_study(config)
+ init_mturk_hit(config)
+
+ logger.debug(f"STATES: {get_user_state_manager().phase_type_to_name_to_page}")
+
+def load_annotation_schematic_data(config: dict) -> None:
+ # Lazy import - only when this function is called
+ from server_utils.front_end import generate_annotation_html_template
+
+ # No longer need to swap in template paths - they are hardcoded in front_end.py
+
+ task_dir = config["task_dir"]
+ # Swap in the right file paths if the user specified the default templates
+ if config["site_dir"] == "default" or True:
+ templates_dir = os.path.join(cur_program_dir, 'templates')
+ if not os.path.exists(templates_dir):
+ # make the directory
+ os.makedirs(templates_dir)
+ config["site_dir"] = templates_dir
+
+ # Creates the templates we'll use in flask by mashing annotation
+ # specification on top of the proto-templates
+ html_template_fname = generate_annotation_html_template(config)
+
+ # Register that we have an annotation phase. Theoretically, we always
+ # should have this, but perhaps there will be some future case where
+ # annotation is not the primary task.
+ #
+ # NOTE: We don't have any HTML for this yet...
+ usm = get_user_state_manager()
+ usm.add_phase(UserPhase.ANNOTATION, config['annotation_task_name'],
+ html_template_fname)
+
+def load_highlights_data(config: dict) -> None:
+ """
+ Load keyword highlights from a TSV file specified in the config.
+
+ The TSV file should have columns: Word, Label, Schema
+ - Word: The keyword or phrase to highlight (supports * wildcards)
+ - Label: The annotation label associated with this keyword
+ - Schema: The annotation schema name
+
+ Wildcards are converted to regex patterns:
+ - 'word*' matches 'word', 'words', 'wording', etc.
+ - '*word' matches 'sword', 'keyword', etc.
+ - 'word' matches exactly 'word' (case-insensitive, word boundaries)
+
+ Also loads keyword_highlight_settings from config:
+ - keyword_probability: Probability of showing matched keywords (default: 1.0)
+ - random_word_probability: Probability of highlighting random words (default: 0.0)
+ - random_word_label: Label for random highlights (default: 'distractor')
+ - random_word_schema: Schema for random highlights (default: 'keyword')
+ """
+ # IMPORTANT: When running as __main__, we need to modify the list in the
+ # package module (potato.flask_server) so that routes.py can see the changes.
+ # This is because Python treats __main__ and potato.flask_server as different modules.
+ import sys
+ if __name__ == '__main__' and 'potato.flask_server' in sys.modules:
+ # Use the package module's list instead of __main__'s list
+ pkg_module = sys.modules['potato.flask_server']
+ patterns_list = pkg_module.keyword_highlight_patterns
+ emphasis_map = pkg_module.emphasis_corpus_to_schemas
+ settings_dict = pkg_module.keyword_highlight_settings
+ else:
+ global keyword_highlight_patterns, emphasis_corpus_to_schemas, keyword_highlight_settings
+ patterns_list = keyword_highlight_patterns
+ emphasis_map = emphasis_corpus_to_schemas
+ settings_dict = keyword_highlight_settings
+
+ # Load keyword highlight settings from config (with defaults)
+ config_settings = config.get('keyword_highlight_settings', {})
+ settings_dict['keyword_probability'] = config_settings.get('keyword_probability', 1.0)
+ settings_dict['random_word_probability'] = config_settings.get('random_word_probability', 0.0)
+ settings_dict['random_word_label'] = config_settings.get('random_word_label', 'distractor')
+ settings_dict['random_word_schema'] = config_settings.get('random_word_schema', 'keyword')
+ logger.debug(f"Loaded keyword highlight settings: {settings_dict}")
+
+ keyword_highlights_file = config.get("keyword_highlights_file")
+ if not keyword_highlights_file:
+ logger.debug("No keyword_highlights_file specified in config")
+ return
+
+ # Note: CWD is already set to task_dir by config_module.py,
+ # so we just need to convert to absolute path from CWD
+ # (don't prepend task_dir again, or we'll double the path)
+ keyword_highlights_file = os.path.realpath(keyword_highlights_file)
+
+ if not os.path.exists(keyword_highlights_file):
+ logger.warning(f"Keyword highlights file not found: {keyword_highlights_file}")
+ return
+
+ logger.info(f"Loading keyword highlights from: {keyword_highlights_file}")
+
+ # Clear the existing list in place (don't reassign) so that modules
+ # that imported keyword_highlight_patterns see the updated contents
+ patterns_list.clear()
+
+ try:
+ with open(keyword_highlights_file, 'r', encoding='utf-8') as f:
+ import csv
+ reader = csv.DictReader(f, delimiter='\t')
+
+ for row in reader:
+ word = row.get('Word', '').strip()
+ label = row.get('Label', '').strip()
+ schema = row.get('Schema', '').strip()
+
+ if not word:
+ continue
+
+ # Convert wildcard pattern to regex
+ # Escape special regex characters except *
+ escaped = re.escape(word).replace(r'\*', r'\w*')
+
+ # Add word boundary markers for exact matching
+ # If pattern starts with wildcard, don't require word boundary at start
+ # If pattern ends with wildcard, don't require word boundary at end
+ if word.startswith('*'):
+ pattern = escaped
+ else:
+ pattern = r'\b' + escaped
+
+ if word.endswith('*'):
+ pattern = pattern
+ else:
+ pattern = pattern + r'\b'
+
+ try:
+ compiled_regex = re.compile(pattern, re.IGNORECASE)
+ patterns_list.append({
+ 'pattern': word,
+ 'regex': compiled_regex,
+ 'label': label,
+ 'schema': schema
+ })
+
+ # Also populate the emphasis corpus for backward compatibility
+ emphasis_map[word].add(HighlightSchema(label=label, schema=schema))
+
+ except re.error as e:
+ logger.warning(f"Invalid regex pattern for keyword '{word}': {e}")
+ continue
+
+ logger.info(f"Loaded {len(patterns_list)} keyword highlight patterns")
+
+ except Exception as e:
+ logger.error(f"Error loading keyword highlights file: {e}")
+ patterns_list.clear()
+
+def load_phase_data(config: dict) -> None:
+ # Lazy import - only when this function is called
+ from server_utils.front_end import generate_html_from_schematic
+
+ global logger
+
+ if "phases" not in config or not config["phases"]:
+ return
+
+ phases = config["phases"]
+
+ # Handle both list and dictionary formats for phases
+ if isinstance(phases, list):
+ # If phases is a list, use the order as defined in the list
+ phase_order = [phase["name"] for phase in phases]
+ # Convert list to dict for easier access
+ phases_dict = {phase["name"]: phase for phase in phases}
+ else:
+ # Original dictionary format
+ if "order" in phases:
+ phase_order = phases["order"]
+ else:
+ phase_order = [k for k in phases.keys() if k != "order"]
+ phases_dict = phases
+
+ logger.debug(f"[PHASE LOAD] phases: {phases}")
+ logger.debug(f"[PHASE LOAD] phase_order: {phase_order}")
+
+ logger.debug("Loading %d phases in order: %s" % (len(phase_order), phase_order))
+
+ for phase_name in phase_order:
+ try:
+ # Skip 'annotation' โ it's handled by the main annotation flow,
+ # not the phase loader. It can appear in the order for sequencing
+ # but doesn't need a phase dict entry.
+ if phase_name not in phases_dict:
+ if phase_name == "annotation":
+ logger.debug(f"Skipping phase '{phase_name}' in loader (handled by main annotation flow)")
+ else:
+ logger.warning(f"Phase '{phase_name}' in order but not defined in phases config, skipping")
+ continue
+
+ phase = phases_dict[phase_name]
+
+ # Handle new format with annotation_schemes directly in phase
+ if "annotation_schemes" in phase:
+ phase_labeling_schemes = phase["annotation_schemes"]
+ # Determine phase type by checking all annotation schemes
+ if phase_labeling_schemes:
+ display_only_count = sum(
+ 1 for s in phase_labeling_schemes
+ if s.get("annotation_type") == "pure_display"
+ )
+ interactive_count = len(phase_labeling_schemes) - display_only_count
+
+ if display_only_count > 0 and interactive_count > 0:
+ logger.warning(
+ f"Phase '{phase_name}' has mixed scheme types: "
+ f"{display_only_count} display-only and {interactive_count} interactive. "
+ f"Treating as ANNOTATION phase."
+ )
+ phase_type = UserPhase.ANNOTATION
+ elif display_only_count == len(phase_labeling_schemes):
+ phase_type = UserPhase.INSTRUCTIONS
+ else:
+ phase_type = UserPhase.ANNOTATION
+ else:
+ phase_type = UserPhase.ANNOTATION
+ else:
+ # File-based phase. Prefer an explicit `type`; otherwise
+ # infer it from the phase name when the name is itself a
+ # canonical phase (e.g. a phase literally named `consent`
+ # or `prestudy`). This makes the documented `phases`
+ # config work even without a `type` field, while still
+ # requiring an explicit `type` for custom-named phases.
+ explicit_type = phase.get("type") if isinstance(phase, dict) else None
+ if explicit_type:
+ phase_type = UserPhase.fromstr(explicit_type)
+ else:
+ try:
+ phase_type = UserPhase.fromstr(phase_name)
+ except ValueError:
+ logger.error(
+ f"Phase '{phase_name}' has no 'type' and its name is "
+ f"not a canonical phase type"
+ )
+ raise Exception(
+ "Phase %s does not have a 'type' and its name is not "
+ "a canonical phase type (one of: consent, prestudy, "
+ "instructions, training, annotation, poststudy). Add "
+ "a 'type:' field to this phase." % phase_name
+ )
+
+ # Instructions phase with an HTML file: register the HTML
+ # directly as a template rather than parsing it as annotation
+ # schemes.
+ if phase_type == UserPhase.INSTRUCTIONS and "file" in phase and phase['file']:
+ phase_file = get_abs_or_rel_path(phase['file'], config)
+ if phase_file.endswith(('.html', '.htm')):
+ logger.debug(f"Instructions phase '{phase_name}' using HTML file: {phase_file}")
+ # Read the HTML and write it as a generated template
+ with open(phase_file, 'rt', encoding='utf-8') as f:
+ instructions_html = f.read()
+
+ # Wrap in a minimal page template for consistency
+ cur_program_dir = os.path.dirname(os.path.abspath(__file__))
+ from server_utils.front_end import get_html
+ html_template_file = os.path.join(cur_program_dir, 'templates', 'base_template_v2.html')
+ header_file = os.path.join(cur_program_dir, 'templates', 'header.html')
+ html_template = get_html(html_template_file, config)
+ header = get_html(header_file, config)
+ html_template = html_template.replace("{{ HEADER }}", header)
+ html_template = html_template.replace("{{ TASK_LAYOUT }}", instructions_html)
+ html_template = html_template.replace("{{annotation_codebook}}", "")
+ html_template = html_template.replace("{{annotation_task_name}}",
+ config.get("annotation_task_name", ""))
+ html_template = html_template.replace("{{keybindings}}", "")
+ html_template = html_template.replace("{{statistics_nav}}", "")
+
+ # Inject project-level base CSS
+ from server_utils.front_end import load_project_base_css_html
+ try:
+ project_css = load_project_base_css_html(config)
+ except FileNotFoundError:
+ project_css = ""
+ html_template = html_template.replace("{{ PROJECT_BASE_CSS }}", project_css)
+
+ site_name = (
+ "_".join(config["annotation_task_name"].split(" "))
+ + "-" + "%s.html" % phase_name
+ )
+ generated_dir = os.path.join(config["site_dir"], "generated")
+ if not os.path.exists(generated_dir):
+ os.makedirs(generated_dir)
+ output_html_fname = os.path.join(generated_dir, site_name)
+ with open(output_html_fname, "wt", encoding="utf-8") as outf:
+ outf.write(html_template)
+
+ user_state_manager = get_user_state_manager()
+ user_state_manager.add_phase(phase_type, phase_name, site_name)
+ logger.debug(f"Registered instructions phase {phase_name} with HTML {site_name}")
+ continue
+
+ # Training and annotation phases can work without a file
+ # They use the main annotation schemes from the config.
+ # Training files typically contain training data items (with
+ # gold_label), not annotation schemes โ so always use the main
+ # annotation schemes for the training phase layout.
+ if phase_type in [UserPhase.TRAINING, UserPhase.ANNOTATION]:
+ phase_labeling_schemes = config.get('annotation_schemes', [])
+ logger.debug(f"Phase {phase_name} using main annotation schemes")
+ else:
+ # Other phases (prestudy, poststudy, etc.)
+ # Support instrument/instruments keys for standard survey instruments
+ phase_labeling_schemes = []
+
+ # Handle single instrument reference
+ if "instrument" in phase:
+ from potato.survey_instruments import get_instrument_questions
+ inst_id = phase["instrument"]
+ logger.debug(f"Phase {phase_name} loading instrument: {inst_id}")
+ phase_labeling_schemes = get_instrument_questions(inst_id)
+
+ # Handle multiple instruments
+ elif "instruments" in phase:
+ from potato.survey_instruments import get_instrument_questions
+ for inst_id in phase["instruments"]:
+ logger.debug(f"Phase {phase_name} loading instrument: {inst_id}")
+ phase_labeling_schemes.extend(get_instrument_questions(inst_id))
+
+ # Handle file reference (can be combined with instrument)
+ if "file" in phase and phase['file']:
+ phase_scheme_fname = get_abs_or_rel_path(phase['file'], config)
+ logger.debug(f"Resolved phase file for {phase_name}: {phase_scheme_fname}")
+ file_schemes = get_phase_annotation_schemes(phase_scheme_fname)
+ if phase_labeling_schemes:
+ # Append file schemes after instrument schemes
+ phase_labeling_schemes.extend(file_schemes)
+ else:
+ phase_labeling_schemes = file_schemes
+
+ # Require at least one source of questions
+ if not phase_labeling_schemes:
+ logger.error(f"Phase {phase_name} requires 'instrument', 'instruments', or 'file'")
+ raise Exception(
+ f"Phase {phase_name} requires 'instrument', 'instruments', or 'file' "
+ "to specify its annotation schemes"
+ )
+
+ # Survey/consent/instrument labels are author-written prose
+ # (often full sentences, any language), not machine identifiers,
+ # so don't title-case them by default. Without this, a label like
+ # "Ja, natรผrlich mรถchte ich teilnehmen" was humanized to
+ # "Ja, Natรผrlich Mรถchte Ich Teilnehmen" (F-049). An individual
+ # survey scheme can still opt back in with humanize_labels: true.
+ # (The TRAINING/ANNOTATION branch above reuses the main schemes
+ # and keeps their existing humanization behavior.)
+ for _survey_scheme in phase_labeling_schemes:
+ if isinstance(_survey_scheme, dict):
+ _survey_scheme.setdefault("humanize_labels", False)
+
+ # Use the default templates unless specified in the phase config
+ # Note: Template paths are now hardcoded in front_end.py
+ # Only handle custom task_layout if specified
+ task_layout_file = None
+ if 'task_layout' in phase:
+ task_layout_file = phase['task_layout']
+
+ try:
+ phase_html_fname = generate_html_from_schematic(
+ phase_labeling_schemes,
+ False, False,
+ phase_name, config,
+ task_layout_file)
+ except KeyError as e:
+ logger.error(f"Error generating HTML for phase {phase_name}: {e}")
+ raise Exception("Error generating HTML for phase %s: %s" \
+ % (phase_name, str(e)))
+
+ # Register the HTML so it's easy to find later
+ user_state_manager = get_user_state_manager()
+ user_state_manager.add_phase(phase_type, phase_name, phase_html_fname)
+ logger.debug(f"Registered phase {phase_name} as {phase_type} with HTML {phase_html_fname}")
+
+ except Exception as e:
+ logger.error(f"Failed to load phase '{phase_name}': {e}")
+ continue
+
+ user_state_manager = get_user_state_manager()
+ logger.debug(f"[PHASE LOAD] phase_type_to_name_to_page: {user_state_manager.phase_type_to_name_to_page}")
+
+
+def get_phase_annotation_schemes(filename: str) -> list[dict]:
+ '''Returns the annotation schemes for a phase from a file.'''
+
+ schemes = []
+ if not os.path.exists(filename):
+ raise Exception("Phase labeling schemes file %s does not exist" % filename)
+
+ if filename.endswith(".json"):
+ with open(filename, "rt", encoding="utf-8") as f:
+ schemes = json.load(f)
+ # Allow users to have specified a single scheme in the JSON file
+ if type(schemes) != list:
+ schemes = [schemes]
+ elif filename.endswith(".jsonl"):
+ with open(filename, 'rt', encoding='utf-8') as f:
+ for line_no, line in enumerate(f):
+ line = line.strip()
+ if not line: # Skip empty lines
+ continue
+ try:
+ schemes.append(json.loads(line))
+ except json.JSONDecodeError as e:
+ raise ValueError(
+ f"Invalid JSON at line {line_no+1} in {filename}: {e}"
+ ) from e
+ elif filename.endswith(".yaml") or filename.endswith(".yml"):
+ with open(filename, 'rt', encoding='utf-8') as f:
+ schemes = yaml.safe_load(f)
+ else:
+ raise Exception("Unknown file format for phase labeling schemes file %s" % filename)
+ return schemes
+
+def get_abs_or_rel_path(fname: str, config: dict) -> str:
+ """
+ Returns the path to the fname if it exists as specified, or if not, attempts to find
+ the file in the relative paths from the config file.
+ """
+ import os
+ logger = globals().get('logger', None)
+ if logger:
+ logger.debug(f"get_abs_or_rel_path: input fname={fname}")
+ if os.path.exists(fname):
+ if logger:
+ logger.debug(f"get_abs_or_rel_path: found file at {fname}")
+ return fname
+
+ # See if we can find the file in the same directory as the config file
+ dname = os.path.dirname(config["__config_file__"]) if "__config_file__" in config else os.getcwd()
+ rel_path = os.path.join(dname, fname)
+ if logger:
+ logger.debug(f"get_abs_or_rel_path: trying {rel_path}")
+ if os.path.exists(rel_path):
+ if logger:
+ logger.debug(f"get_abs_or_rel_path: found file at {rel_path}")
+ return rel_path
+
+ # See if we can locate the file in the current working directory
+ cwd = os.getcwd()
+ rel_path = os.path.join(cwd, fname)
+ if logger:
+ logger.debug(f"get_abs_or_rel_path: trying {rel_path}")
+ if os.path.exists(rel_path):
+ if logger:
+ logger.debug(f"get_abs_or_rel_path: found file at {rel_path}")
+ return rel_path
+
+ # See if we can figure it out from the real path directory
+ real_path = os.path.abspath(dname)
+ dir_path = os.path.dirname(real_path)
+ fname2 = os.path.join(dir_path, fname)
+ if logger:
+ logger.debug(f"get_abs_or_rel_path: trying {fname2}")
+ if not os.path.exists(fname2):
+ if logger:
+ logger.error(f"File not found: {fname2}")
+ raise FileNotFoundError("File not found: %s" % fname2)
+ return fname2
+
+def get_displayed_text(text):
+ """Render the text to display to the user in the annotation interface.
+
+ Handles both string and list inputs. When text is a list (for dialogue
+ or pairwise comparisons), it formats the list items according to list_as_text config.
+
+ Supported prefix types:
+ - alphabet: A. B. C. prefixes
+ - number: 1. 2. 3. prefixes
+ - bullet: โข prefixes
+ - none: No prefix (use for dialogue with speaker names in text)
+
+ Additional options:
+ - horizontal: Display items side-by-side (for pairwise comparison)
+ - alternating_shading: Shade every other turn (for dialogue readability)
+ """
+ import re
+
+ # Handle dict inputs (for tree structures, agent traces, complex data)
+ # Convert to JSON string for display โ the actual rendering is handled by
+ # display types (conversation_tree, web_agent_trace, live_agent, etc.)
+ if isinstance(text, dict):
+ import json as _json
+ return _json.dumps(text, ensure_ascii=False, indent=2)
+
+ # Handle list inputs (for dialogue or pairwise comparisons with list_as_text config)
+ if isinstance(text, list):
+ list_config = config.get("list_as_text", {})
+ prefix_type = list_config.get("text_list_prefix_type", "alphabet")
+ horizontal = list_config.get("horizontal", False)
+ alternating_shading = list_config.get("alternating_shading", False)
+
+ formatted_items = []
+ for i, item in enumerate(text):
+ # Generate prefix based on type
+ if prefix_type == "alphabet":
+ prefix = f"{chr(ord('A') + i)}. "
+ elif prefix_type == "number":
+ prefix = f"{i + 1}. "
+ elif prefix_type == "bullet":
+ prefix = "โข "
+ elif prefix_type == "none":
+ prefix = ""
+ else:
+ # Default to alphabet for unknown types
+ prefix = f"{chr(ord('A') + i)}. "
+
+ # Recursively process each item
+ processed_item = get_displayed_text(item) if isinstance(item, str) else str(item)
+
+ # Apply alternating shading for dialogue readability
+ if alternating_shading:
+ shade_class = "dialogue-turn-even" if i % 2 == 0 else "dialogue-turn-odd"
+
+ # Try to extract speaker name (text before first colon)
+ speaker_match = re.match(r'^([^:]+):\s*(.*)$', processed_item, re.DOTALL)
+ if speaker_match:
+ speaker_name = speaker_match.group(1).strip()
+ speaker_text = speaker_match.group(2).strip()
+ # Generate a consistent color index based on speaker name
+ speaker_hash = sum(ord(c) for c in speaker_name) % 6
+ # Use span with display:block style (spans are in sanitizer allowlist)
+ formatted_items.append(
+ f''
+ f'{speaker_name}: '
+ f'{prefix}{speaker_text} '
+ )
+ else:
+ # No speaker detected, use simple format
+ formatted_items.append(
+ f'{prefix}{processed_item} '
+ )
+ else:
+ formatted_items.append(f"{prefix}{processed_item}")
+
+ # Join based on layout type
+ if horizontal:
+ # Horizontal layout for pairwise comparison
+ cell_width = 100 // len(formatted_items) if formatted_items else 100
+ cells = [
+ f'{item} '
+ for item in formatted_items
+ ]
+ text = '' + ''.join(cells) + ' '
+ elif alternating_shading:
+ # Already wrapped in divs, join without extra breaks
+ text = ''.join(formatted_items)
+ else:
+ # Vertical layout with double line breaks
+ text = " ".join(formatted_items)
+ return text
+
+ # Normalize text for consistent positioning (matches client-side normalization)
+ # Remove control characters but preserve all Unicode (fixes issue #114)
+ text = re.sub(r'[\x00-\x1F\x7F]', lambda m: m.group() if m.group() == '\n' else '', text)
+ text = re.sub(r'[ \t]+', ' ', text) # Normalize horizontal whitespace only
+ text = text.strip()
+
+ if config.get("highlight_linebreaks", False):
+ text = text.replace("\n", " ")
+
+ return text
+
+# Core functions used by routes.py
+
+def init_user_state(username):
+ """
+ Initialize the state for a user, returning the user state object.
+ """
+ usm = get_user_state_manager()
+ usm.add_user(username)
+
+ # Store the session creation time
+ session['created_at'] = datetime.now()
+
+ return usm.get_user_state(username)
+
+def is_session_valid() -> bool:
+ """
+ Check if the current session is valid based on the creation time.
+ """
+ if 'created_at' not in session:
+ return False
+ return datetime.now() - session['created_at'] < SESSION_TIMEOUT
+
+@app.before_request
+def before_request():
+ """
+ Check session validity before processing any request.
+ Only enforce session validation for protected routes.
+ """
+ # Skip session validation in debug mode
+ if config.get("debug", False):
+ return None
+
+ # Allow unauthenticated access to these endpoints
+ allowed_paths = [
+ '/', '/auth', '/register', '/static/', '/favicon.ico', '/robots.txt', '/health', '/api/', '/api/instance/', '/api/instances', '/api/config', '/api/status', '/api/heartbeat'
+ ]
+ path = request.path
+ if any(path == allowed or path.startswith(allowed) for allowed in allowed_paths):
+ return None
+
+ if not is_session_valid():
+ session.clear() # Clear the session
+ return redirect(url_for('home')) # Redirect to home page (login/register)
+
+def get_users():
+ """
+ Returns the list of users that have logged in.
+ """
+ return get_user_state_manager().get_user_ids()
+
+def get_user_state(username):
+ """
+ Returns the user state object for the given username.
+ """
+ return get_user_state_manager().get_user_state(username)
+
+def move_to_prev_instance(user_id) -> bool:
+ '''Moves the user back to the previous instance and returns True if successful'''
+ user_state = get_user_state(user_id)
+ return user_state.go_back()
+
+def move_to_next_instance(user_id) -> bool:
+ '''Moves the user forward to the next instance and returns True if successful'''
+ logger.debug(f"=== MOVE_TO_NEXT_INSTANCE START ===")
+ logger.debug(f"User ID: {user_id}")
+
+ user_state = get_user_state(user_id)
+ logger.debug(f"Before navigation - current_instance_index: {user_state.get_current_instance_index()}")
+ logger.debug(f"Before navigation - instance_id_ordering: {user_state.instance_id_ordering}")
+
+ # If the user is at the end of the list, try to assign instances to the user
+ if user_state.is_at_end_index():
+ logger.debug(f"User {user_id} is at the end of the list, assigning new instances")
+ num_assigned = get_item_state_manager().assign_instances_to_user(user_state)
+ logger.debug(f"Assigned {num_assigned} new instances to user {user_id}")
+
+ result = user_state.go_forward()
+ logger.debug(f"After navigation - current_instance_index: {user_state.get_current_instance_index()}")
+ logger.debug(f"Navigation result: {result}")
+
+ logger.debug(f"=== MOVE_TO_NEXT_INSTANCE END ===")
+ return result
+
+def go_to_id(user_id: str, instance_index: int):
+ '''Causes the user's view to change to the Item at the given index.'''
+ user_state = get_user_state(user_id)
+ user_state.go_to_index(int(instance_index))
+
+def get_current_page_html(config, username):
+ """
+ Returns the HTML for the current page that the user is on.
+
+ For phase pages (consent, instructions, etc.), this provides minimal
+ context variables needed by the shared template structure.
+ """
+ user_state = get_user_state(username)
+ phase, page = user_state.get_current_phase_and_page()
+
+ is_annotation_page = phase == UserPhase.ANNOTATION
+
+ usm = get_user_state_manager()
+ html_fname = usm.get_phase_html_fname(phase, page)
+
+ # Provide context variables needed by the template
+ # For phase pages, many annotation-specific fields can be empty/default
+ context = {
+ 'username': username,
+ 'annotation_task_name': config.get('annotation_task_name', ''),
+ 'annotation_codebook_url': _sanitize_codebook_url(config.get('annotation_codebook_url', '')),
+ 'debug_mode': config.get('debug', False),
+ 'ui_debug': config.get('ui_debug', False),
+ 'server_debug': config.get('server_debug', False),
+ 'debug_phase': config.get('debug_phase', None),
+ 'instance': '',
+ 'instance_plain_text': '',
+ 'instance_id': '',
+ 'instance_index': 0,
+ 'finished': 0,
+ 'total_count': user_state.get_assigned_instance_count() if hasattr(user_state, 'get_assigned_instance_count') else 0,
+ 'ui_config': config.get('ui_config', {}),
+ 'is_annotation_page': is_annotation_page,
+ 'annotation_instructions': config.get('annotation_instructions', ''),
+ 'annotation_status': 'unlabeled',
+ 'instance_has_annotations': False,
+ 'can_go_back': usm.can_user_go_back(username),
+ 'jumping_to_id_disabled': config.get('jumping_to_id_disabled', False),
+ }
+ rendered_html = render_template(html_fname, **context)
+ soup = BeautifulSoup(rendered_html, "html.parser")
+
+ phase_annotations = user_state.phase_to_page_to_label_to_value.get(phase, {}).get(page, {})
+ for label_obj, value in phase_annotations.items():
+ schema = label_obj.get_schema()
+ label = label_obj.get_name()
+ name = schema + ":::" + label
+
+ input_fields = soup.find_all(["input", "select", "textarea"], {"name": name})
+ if not input_fields:
+ input_fields = soup.find_all(["input"], {"schema": schema, "label_name": label})
+
+ for input_field in input_fields:
+ if input_field is None:
+ continue
+
+ if input_field.get('type') == 'checkbox' or input_field.get('type') == 'radio':
+ if value:
+ if input_field.get('type') == 'radio':
+ if input_field.get('value') == value:
+ input_field['checked'] = True
+ else:
+ input_field['checked'] = True
+
+ if input_field.get('type') == 'text':
+ if isinstance(value, str):
+ input_field['value'] = value
+
+ if input_field.get('type') == 'number':
+ input_field['value'] = str(value)
+
+ if input_field.name == 'textarea':
+ if isinstance(value, str):
+ input_field.string = value
+
+ if input_field.name == 'select':
+ if isinstance(value, str):
+ options = input_field.find_all("option", {"value": value})
+ if options:
+ options[0]["selected"] = "selected"
+
+ return str(soup)
+
+def _sanitize_codebook_url(url: str) -> str:
+ """Sanitize codebook URL to prevent javascript: and other dangerous protocols."""
+ if not url:
+ return ""
+ stripped = url.strip()
+ # Block dangerous URL schemes
+ lower = stripped.lower().replace('\t', '').replace('\n', '').replace('\r', '')
+ for scheme in ('javascript:', 'vbscript:', 'data:'):
+ if lower.startswith(scheme):
+ logger.warning(f"Blocked dangerous scheme in annotation_codebook_url: {scheme}")
+ return ""
+ return stripped
+
+
+def _scheme_is_required(scheme: dict) -> bool:
+ """Check if an annotation scheme is marked as required."""
+ if scheme.get("required") is True:
+ return True
+
+ lr = scheme.get("label_requirement", {})
+ if lr is True:
+ return True
+ if isinstance(lr, dict) and lr.get("required") is True:
+ return True
+ return False
+
+
+def _scheme_has_required_annotation(user_state, instance_id: str, scheme: dict) -> bool:
+ """Check whether a required scheme has any annotation value for an instance."""
+ schema_name = scheme.get("name", "")
+
+ label_annotations = user_state.instance_id_to_label_to_value.get(instance_id, {})
+ for label_key, value in label_annotations.items():
+ if hasattr(label_key, "get_schema") and label_key.get_schema() == schema_name:
+ if value:
+ return True
+
+ span_annotations = user_state.instance_id_to_span_to_value.get(instance_id, {})
+ if span_annotations:
+ if isinstance(span_annotations, dict) and span_annotations:
+ return True
+ if isinstance(span_annotations, list) and len(span_annotations) > 0:
+ return True
+
+ return False
+
+
+def _instance_meets_required_annotation_rules(user_state, instance_id: str) -> list:
+ """Return the names of required schemes that are still unsatisfied."""
+ unsatisfied = []
+ for scheme in config.get("annotation_schemes", []):
+ if _scheme_is_required(scheme) and not _scheme_has_required_annotation(user_state, instance_id, scheme):
+ unsatisfied.append(scheme.get("name", "unknown"))
+ return unsatisfied
+
+
+def _is_user_adjudicator(username: str) -> bool:
+ """Check if a user is an authorized adjudicator."""
+ adj_mgr = get_adjudication_manager()
+ if adj_mgr and adj_mgr.adj_config.enabled:
+ return adj_mgr.is_adjudicator(username)
+ return False
+
+
+def render_page_with_annotations(username: str):
+ '''
+ When annotating, shows the current instance to the user with any annotations
+ they may have made. This method is called when the user is in the annotation
+ phase and is currently annotating.
+ '''
+
+ # Hacky nonsense
+ global emphasis_corpus_to_schemas
+
+ user_state = get_user_state_manager().get_user_state(username)
+ phase, page = user_state.get_current_phase_and_page()
+
+ is_annotation_page = phase == UserPhase.ANNOTATION
+
+ item = user_state.get_current_instance()
+ if item is None:
+ logger.warning(
+ f"User {username} has no valid current instance after loading state"
+ )
+ if not user_state.has_remaining_assignments():
+ get_user_state_manager().advance_phase(username)
+ return redirect(url_for("home"))
+
+ instance_id = item.get_id()
+
+ # Extract pre-annotation data if quality control is enabled
+ pre_annotation_data = None
+ qc_manager = get_quality_control_manager()
+ if qc_manager:
+ pre_annotation_data = qc_manager.extract_pre_annotations(instance_id, item.get_data())
+
+ # LLM-judge inline suggestion (judge โ human alignment). Reads a persisted
+ # judge prediction for this instance/schema; optionally computes on demand.
+ judge_prediction = _get_inline_judge_prediction(instance_id, item)
+
+ # Signal-based triage: why was this item prioritized in the queue?
+ triage_info = _get_triage_info(item)
+
+ # DEBUG: Add detailed logging
+ logger.debug(f"=== RENDER_PAGE_WITH_ANNOTATIONS START ===")
+ logger.debug(f"Username: {username}")
+ logger.debug(f"User state current_instance_index: {user_state.get_current_instance_index()}")
+ logger.debug(f"User state instance_id_ordering: {user_state.instance_id_ordering}")
+ logger.debug(f"Current instance ID: {instance_id}")
+
+ # print('instance_id: ', instance_id)
+
+ # directly display the prepared displayed_text
+ item_data = item.get_data() if hasattr(item, "get_data") else {}
+ text_key = config.get("item_properties", {}).get("text_key", "text")
+ raw_text = None
+ if isinstance(item_data, dict):
+ raw_text = item_data.get("displayed_text")
+ if raw_text is None:
+ raw_text = item_data.get(text_key, item_data.get("text"))
+
+ if raw_text is None:
+ raw_text = item.get_displayed_text() if hasattr(item, "get_displayed_text") else item.get_text()
+
+ text = raw_text if "displayed_text" in (item_data or {}) else get_displayed_text(raw_text)
+ # print('displayed_text: ', text)
+
+ # Save the original plain text BEFORE any span rendering
+ # This is needed for the frontend to calculate correct span positions
+ # The data-original-text attribute must contain plain text (no HTML span tags)
+ # while the DOM content contains the rendered HTML with span highlights
+ # Strip HTML tags to get actual plain text for position calculations
+ import re as re_module
+ original_plain_text = re_module.sub(r'<[^>]+>', '', text)
+ # Also normalize whitespace
+ original_plain_text = re_module.sub(r'\s+', ' ', original_plain_text).strip()
+
+ var_elems = {
+ "instance": { "text": text },
+ "emphasis": list(emphasis_corpus_to_schemas)
+ }
+
+ # Include full instance data for dynamic schemas (extractive_qa, text_edit,
+ # error_span, card_sort, conjoint) that need fields beyond text_key
+ if item_data and isinstance(item_data, dict):
+ var_elems["instance_data"] = {
+ k: v for k, v in item_data.items()
+ if isinstance(v, (str, int, float, bool, list))
+ }
+
+ # also save the displayed text in the metadata dict
+ # instance_id_to_data[instance_id]['displayed_text'] = text
+
+ # If the user has labeled spans within this instance before, replace the
+ # current instance text with pre-annotated mark-up. We do this here before
+ # the render_template call so that we can directly insert the span-marked-up
+ # HTML into the template.
+ #
+ # NOTE: This currently requires a very tight (and kludgy) binding between
+ # the UI code for how Potato represents span annotations and how the
+ # back-end displays these. Future work when we are better programmers will
+ # pass this info to client side for rendering, rather than doing
+ # pre-rendering here. This also means that any changes to the UI code for
+ # rendering need to be updated here too.
+ #
+ # NOTE2: We have to this here to account for any keyword highlighting before
+ # the instance text gets marked up in the post-processing below
+ span_annotations = get_span_annotations_for_user_on(username, instance_id)
+ if span_annotations is not None and len(span_annotations) > 0:
+ # Mark up the instance text where the annotated spans were
+ text = render_span_annotations(text, span_annotations)
+
+ # If the admin has specified that certain keywords need to be highlighted,
+ # post-process the selected instance so that it now also has colored span
+ # overlays for keywords. This also include label suggestions for the user.
+ #
+ # NOTE: this code is probably going to break the span annotation's
+ # understanding of the instance. Need to check this...
+ schema_content_to_prefill = []
+
+ #prepare label suggestions
+ label_suggestion_json = get_label_suggestions(item, config, schema_content_to_prefill)
+
+ var_elems["suggestions"] = list(label_suggestion_json)
+
+ # Pass BWS items data to frontend JS
+ if config.get("bws_config") or config.get("ibws_config"):
+ var_elems["bws_items"] = item.get_data().get("_bws_items", [])
+ # Fill in the kwargs that the user wanted us to include when rendering the page
+ kwargs = {}
+ for kw in config["item_properties"].get("kwargs", []):
+ if kw in item.get_data():
+ kwargs[kw] = item.get_data()[kw]
+
+ all_statistics = get_user_state(username).generate_user_statistics()
+
+ # TODO: Display plots for agreement scores instead of only the overall score
+ # in the statistics sidebar
+ # all_statistics['Agreement'] = get_agreement_score('all', 'all', return_type='overall_average')
+ # print(all_statistics)
+
+ # Set the html file as surveyflow pages when the instance is a not an
+ # annotation page (survey pages, prestudy pass or fail page)
+ html_file = config["site_file"]
+
+ var_elems_html = "".join(
+ map(lambda item : (
+ f''
+ ), var_elems.items())
+ )
+
+ custom_js = ""
+ if config["customjs"] and config.get("customjs_hostname"):
+ custom_js = (
+ f''
+ )
+ elif config["customjs"]:
+ custom_js = (
+ ''
+ )
+ else:
+ custom_js = (
+ ''
+ )
+
+ # Shea: Test for AI suggestion
+ # ai_hints = get_ai_hints(text)
+
+ # Flask will fill in the things we need into the HTML template we've created,
+ # replacing {{variable_name}} with the associated text for keyword arguments
+
+ # Calculate progress counter values
+ # Get the number of completed annotations and remaining assignable items
+ finished_count = get_user_state(username).get_annotation_count()
+ remaining_count = get_item_state_manager().get_total_assignable_items_for_user(get_user_state(username))
+ # Total = finished + remaining (so counter shows "X / Total" not "X / Remaining")
+ total_count = finished_count + remaining_count
+
+ # Cap total by max_assignments if set (so progress shows "3/6" not "3/100")
+ max_assignments = get_user_state(username).get_max_assignments()
+ if max_assignments >= 0:
+ total_count = min(total_count, max_assignments)
+
+ # Determine annotation status for the status badge (three-state)
+ annotation_status = "unlabeled"
+ if user_state.has_annotated(instance_id):
+ unsatisfied = _instance_meets_required_annotation_rules(user_state, instance_id)
+ annotation_status = "labeled" if not unsatisfied else "in_progress"
+ instance_has_annotations = (annotation_status != "unlabeled")
+
+ # Get UI configuration from config
+ ui_config = config.get("ui", {})
+ annotation_schemes = config.get("annotation_schemes", [])
+
+ # Add layout configuration to ui_config for JavaScript access
+ if config.get("layout"):
+ ui_config = dict(ui_config) # Make a copy to avoid modifying the original
+ ui_config["layout"] = config["layout"]
+
+ # Detect if any annotation scheme is video_annotation type
+ # This is used to customize the display (show "Video to Annotate:" instead of "Text to Annotate:")
+ has_video_annotation = any(
+ scheme.get("annotation_type") == "video_annotation"
+ for scheme in annotation_schemes
+ )
+
+ # Detect if any annotation scheme is audio_annotation type (or tiered_annotation with audio media)
+ # This is used to customize the display (hide "Text to Annotate:" for audio-focused tasks)
+ has_audio_annotation = any(
+ scheme.get("annotation_type") == "audio_annotation"
+ or (scheme.get("annotation_type") == "tiered_annotation" and scheme.get("media_type") == "audio")
+ for scheme in annotation_schemes
+ )
+
+ # Detect if any annotation scheme is image_annotation type
+ # This is used to customize the display (show "Image to Annotate:" instead of "Text to Annotate:")
+ has_image_annotation = any(
+ scheme.get("annotation_type") == "image_annotation"
+ for scheme in annotation_schemes
+ )
+
+ # Initialize display_html before it's referenced by _detect_frontend_assets_for_page
+ display_html = ""
+
+ frontend_assets = _detect_frontend_assets_for_page(html_file, display_html)
+
+ # Check if AI support is enabled (for conditional loading of visual_ai_assistant.js)
+ ai_enabled = config.get("ai_support", {}).get("enabled", False)
+
+ # Check if agent proxy is configured (for conditional loading of agent-chat.js/css)
+ agent_proxy_enabled = "agent_proxy" in config
+
+ # Check if chat support is enabled (for conditional loading of llm-chat-sidebar assets)
+ chat_enabled = config.get("chat_support", {}).get("enabled", False)
+
+ # Check if live agent is enabled (for conditional loading of live-agent assets)
+ live_agent_enabled = bool(config.get("live_agent"))
+
+ # Get pre-annotation configuration
+ pre_annotation_config = {}
+ if qc_manager:
+ pre_annotation_config = qc_manager.get_pre_annotation_config()
+
+ # Check if instance_display is configured (new explicit display mode)
+ has_instance_display = "instance_display" in config
+ display_html = ""
+ display_template_vars = {}
+
+ if has_instance_display:
+ try:
+ display_renderer = get_instance_display_renderer(config)
+ display_template_vars = display_renderer.get_template_variables(item.get_data())
+ display_html = display_template_vars.get("display_html", "")
+ logger.debug(f"Instance display rendered: {len(display_html)} chars")
+ except Exception as e:
+ logger.error(f"Error rendering instance display: {e}")
+ has_instance_display = False # Fall back to legacy mode
+
+ frontend_assets = _detect_frontend_assets_for_page(html_file, display_html)
+
+ # Get IBWS round info if active
+ ibws_round_info = None
+ if config.get("ibws_config"):
+ from potato.ibws_manager import get_ibws_manager
+ ibws_mgr = get_ibws_manager()
+ if ibws_mgr:
+ ibws_round_info = ibws_mgr.get_round_info()
+
+ rendered_html = render_template(
+ html_file,
+ username=username,
+ # This is what instance the user is currently on (may contain span HTML)
+ instance=text,
+ # Original plain text without span HTML (for data-original-text attribute)
+ instance_plain_text=original_plain_text,
+ instance_obj=item,
+ # Full record dict so schemas like process_reward / trajectory_eval
+ # can bind to structured fields (e.g. structured_turns) via the
+ # [data-instance-json] element.
+ instance_record=item.get_data(),
+ instance_id=instance_id,
+ instance_index=user_state.get_current_instance_index(),
+ finished=get_user_state(username).get_annotation_count(),
+ total_count=total_count,
+ alert_time_each_instance=config.get("alert_time_each_instance", 10000000),
+ statistics_nav=all_statistics,
+ var_elems=var_elems_html,
+ custom_js=custom_js,
+ # Pass annotation schemes to the template
+ annotation_schemes=annotation_schemes,
+ annotation_task_name=config["annotation_task_name"],
+ debug=config.get("debug", False),
+ ui_config=ui_config,
+ has_video_annotation=has_video_annotation,
+ has_audio_annotation=has_audio_annotation,
+ has_image_annotation=has_image_annotation,
+ ai_enabled=ai_enabled,
+ # Pre-annotation data for model predictions
+ pre_annotations=pre_annotation_data,
+ pre_annotation_config=pre_annotation_config,
+ # LLM-judge inline suggestion (judge โ human alignment)
+ judge_prediction=judge_prediction,
+ # Signal-based triage badge (why this item was prioritized)
+ triage_info=triage_info,
+ # Instance display (new explicit display mode)
+ has_instance_display=has_instance_display,
+ display_html=display_html,
+ display_fields=display_template_vars.get("display_fields", {}),
+ display_raw=display_template_vars.get("display_raw", {}),
+ span_targets=display_template_vars.get("span_targets", []),
+ multi_span_mode=display_template_vars.get("multi_span_mode", False),
+ frontend_assets=frontend_assets,
+ # Agent proxy (for conditional loading of agent-chat assets)
+ agent_proxy_enabled=agent_proxy_enabled,
+ # Chat support (for conditional loading of llm-chat-sidebar assets)
+ chat_enabled=chat_enabled,
+ # Live agent (for conditional loading of live-agent assets)
+ live_agent_enabled=live_agent_enabled,
+ # Annotation instructions (collapsible banner)
+ annotation_instructions=config.get("annotation_instructions", ""),
+ # Adjudication: show link for adjudicators
+ is_adjudicator=_is_user_adjudicator(username),
+ annotation_codebook_url=_sanitize_codebook_url(config.get("annotation_codebook_url", "")),
+ # Annotation status indicator (three-state: labeled/in_progress/unlabeled)
+ annotation_status=annotation_status,
+ instance_has_annotations=instance_has_annotations,
+ # if this is an annotation page
+ is_annotation_page=is_annotation_page,
+ # IBWS round info (for round banner)
+ ibws_round_info=ibws_round_info,
+ # Hide back button when on first instance with no previous phase
+ can_go_back=get_user_state_manager().can_user_go_back(username),
+ # Hide jump-to-ID navigation controls when disabled
+ jumping_to_id_disabled=config.get("jumping_to_id_disabled", False),
+ # ai=ai_hints,
+ **kwargs
+ )
+
+ # Parse the page so we can programmatically reset the annotation state
+ # to what it was before
+ soup = BeautifulSoup(rendered_html, "html.parser")
+
+ # If the user has annotated this before, walk the DOM and fill out what they
+ # did
+ annotations = get_annotations_for_user_on(username, instance_id)
+
+ # If no annotations yet, check for pre-annotations (model predictions).
+ # NOTE: get_annotations_for_user_on returns an empty dict {} (not None) for a
+ # user with no annotations, so guard on falsiness rather than `is None`.
+ if not annotations and pre_annotation_data:
+ logger.debug(f"Applying pre-annotations for instance {instance_id}")
+ scheme_dict = {}
+ annotations = defaultdict(dict)
+ for it in config['annotation_schemes']:
+ if it['annotation_type'] in ['radio', 'multiselect']:
+ it['label2value'] = {(l if type(l) == str else l['name']):str(i+1) for i,l in enumerate(it['labels'])}
+ scheme_dict[it['name']] = it
+
+ for schema_name, predicted_value in pre_annotation_data.items():
+ if schema_name not in scheme_dict:
+ logger.debug(f"Pre-annotation schema {schema_name} not found in annotation schemes")
+ continue
+
+ scheme = scheme_dict[schema_name]
+ if scheme['annotation_type'] in ['radio', 'multiselect']:
+ # predicted_value should be a label name. Store the LABEL NAME as
+ # the value (not the label2value index): the renderer below checks
+ # a radio when input.value == value, and the radio's value
+ # attribute is the label name. This matches how a returning
+ # user's restored annotations are stored.
+ if isinstance(predicted_value, str) and predicted_value in scheme.get('label2value', {}):
+ annotations[schema_name][predicted_value] = predicted_value
+ elif isinstance(predicted_value, list):
+ # Multi-select: multiple values
+ for val in predicted_value:
+ if val in scheme.get('label2value', {}):
+ annotations[schema_name][val] = val
+ elif scheme['annotation_type'] in ['text']:
+ if "labels" not in scheme:
+ annotations[schema_name]['text_box'] = str(predicted_value)
+ elif scheme['annotation_type'] in ['likert', 'slider', 'number']:
+ annotations[schema_name]['slider'] = str(predicted_value)
+ else:
+ logger.debug(f"Pre-annotation not yet supported for {scheme['annotation_type']}")
+
+ # convert the label suggestions into annotations for front-end rendering
+ # (empty dict, like None, means "no user annotations yet")
+ if not annotations and schema_content_to_prefill:
+ scheme_dict = {}
+ annotations = defaultdict(dict)
+ for it in config['annotation_schemes']:
+ if it['annotation_type'] in ['radio', 'multiselect']:
+ it['label2value'] = {(l if type(l) == str else l['name']):str(i+1) for i,l in enumerate(it['labels'])}
+ scheme_dict[it['name']] = it
+ for s in schema_content_to_prefill:
+ if scheme_dict[s['name']]['annotation_type'] in ['radio', 'multiselect']:
+ # Store the label NAME as value so the renderer matches the
+ # radio/checkbox input's value attribute (not the index).
+ annotations[s['name']][s['label']] = s['label']
+ elif scheme_dict[s['name']]['annotation_type'] in ['text']:
+ if "labels" not in scheme_dict[s['name']]:
+ annotations[s['name']]['text_box'] = s['label']
+ else:
+ logger.warning('Label suggestions not supported for annotation_type %s, please submit a github issue to get support' % scheme_dict[s['name']]['annotation_type'])
+ logger.debug(f"annotations: {annotations}")
+ if annotations is not None:
+ # Reset the state
+ for schema_name, label_dict in annotations.items():
+ # this needs to be fixed, there is a chance that we get incorrect type
+ if not isinstance(label_dict, dict):
+ logger.warning(f"Skipping {schema_name}: Expected dict but got {type(label_dict)} -> {label_dict}")
+ continue
+
+ for label_name, value in label_dict.items():
+ schema = schema_name
+ label = label_name
+ name = schema + ":::" + label
+
+ # Find all the input, select, and textarea tags with this name
+ # (which was annotated) and figure out which one to fill in
+ input_fields = soup.find_all(["input", "select", "textarea"], {"name": name})
+
+ # For radio buttons, the name attribute is just the schema (not schema:::label)
+ # because all radio buttons in a group must have the same name for HTML mutual exclusivity
+ # So we also search by schema and label_name attributes
+ if not input_fields:
+ input_fields = soup.find_all(
+ ["input"],
+ {"schema": schema, "label_name": label}
+ )
+
+ # For image/audio/video annotation data, the hidden input has name=schema_name
+ # and the label is "_data"
+ if not input_fields and label == "_data":
+ input_fields = soup.find_all(
+ ["input"],
+ {"name": schema, "class": "annotation-data-input"}
+ )
+ logger.debug(f"Looking for annotation-data-input with name={schema}, found {len(input_fields)}")
+
+ for input_field in input_fields:
+
+ if input_field is None:
+ logger.debug(f"No input for {name}")
+ continue
+
+ # If it's a range input (slider, soft_label, vas, range_slider, etc.),
+ # set the value attribute so loadAnnotations() reads it back
+ if input_field.get('type') == 'range':
+ input_field['value'] = value
+ continue
+
+ if input_field.get('type') == 'checkbox' or input_field.get('type') == 'radio':
+ if value:
+ # For radio buttons, only check if the value matches
+ # (multiple radios share the same schema/label_name but have different values)
+ if input_field.get('type') == 'radio':
+ if input_field.get('value') == value:
+ input_field['checked'] = True
+ else:
+ # For checkboxes, set checked
+ input_field['checked'] = True
+
+ # Handle text inputs - set value attribute
+ if input_field.get('type') == 'text':
+ if isinstance(value, str):
+ input_field['value'] = value
+
+ # Handle number inputs - set value attribute
+ if input_field.get('type') == 'number':
+ input_field['value'] = str(value)
+
+ # Handle textareas - set content between tags (not value attribute)
+ # Textareas don't have a type attribute, check tag name instead
+ if input_field.name == 'textarea':
+ if isinstance(value, str):
+ input_field.string = value
+
+ # Handle hidden inputs for image/audio/video annotation data
+ if input_field.get('type') == 'hidden':
+ if isinstance(value, str):
+ input_field['value'] = value
+ # Mark this input as server-set to distinguish from browser-cached values
+ input_field['data-server-set'] = 'true'
+ logger.debug(f"Set hidden input {name} value (length: {len(value)}) with server-set flag")
+
+ # Handle select elements - set the 'selected' attribute on matching option
+ if input_field.name == 'select':
+ if isinstance(value, str):
+ # Find the option with the matching value and set it as selected
+ options = input_field.find_all("option", {"value": value})
+ if options:
+ options[0]["selected"] = "selected"
+ logger.debug(f"Set select {name} option to {value}")
+ else:
+ logger.debug(f"No option found with value {value} for select {name}")
+
+ if False:
+ # If it's not a text area, let's see if this is the button
+ # that was checked, and if so mark it as checked
+ if input_field.name != "textarea" and input_field.has_attr("value") and input_field.get("value") != value:
+ continue
+ else:
+ input_field["checked"] = True
+ input_field["value"] = value
+
+ # Set the input value for textarea input
+ #if input_field.name == "textarea" and isinstance(value, str):
+ # input_field.string = value
+
+ # Find the right option and set it as selected if the current
+ # annotation schema is a select box
+ if label == "select-one":
+ option = input_field.findChildren("option", {"value": value})[0]
+ option["selected"] = "selected"
+
+ # randomize the order of options for schemas that support it
+ selected_schemas_for_option_randomization = []
+ for it in config['annotation_schemes']:
+ if it.get('option_randomization') and it['annotation_type'] in ('multirate', 'radio', 'multiselect', 'select'):
+ selected_schemas_for_option_randomization.append(it['description'])
+
+
+ soup = randomize_options(soup, selected_schemas_for_option_randomization,
+ map_user_id_to_digit(username))
+
+ # If the admin has turned on AI hints, add them to the page
+ soup = add_ai_hints(soup, instance_id)
+
+ rendered_html = str(soup)
+
+ # Filter options per instance based on dynamic_options config
+ dynamic_option_schemes = [
+ s for s in config.get('annotation_schemes', [])
+ if s.get('dynamic_options') and s['annotation_type'] in ('radio', 'multiselect', 'select')
+ ]
+ if dynamic_option_schemes:
+ soup = filter_dynamic_options(soup, dynamic_option_schemes, item.get_data())
+ rendered_html = str(soup)
+
+ # Populate dynamic multirate options from instance data
+ has_dynamic_multirate = any(
+ scheme.get('options_from_data')
+ for scheme in config.get('annotation_schemes', [])
+ if scheme.get('annotation_type') == 'multirate'
+ )
+ if has_dynamic_multirate:
+ from potato.server_utils.schemas.multirate import populate_dynamic_multirate
+ rendered_html = populate_dynamic_multirate(rendered_html, item.get_data())
+
+ return rendered_html
+
+
+def _get_inline_judge_prediction(instance_id, item):
+ """Return a judge suggestion dict for the inline display, or None.
+
+ Gated by ``judge_alignment.inline.enabled``. Prefers a persisted prediction
+ (admin pre-runs the batch); computes on demand only if
+ ``judge_alignment.inline.compute_on_demand`` is set. Shape mirrors solo
+ mode's ``llm_prediction``: {label, confidence, reasoning, schema,
+ prompt_version, running}.
+ """
+ ja = config.get("judge_alignment", {}) or {}
+ inline = ja.get("inline", {}) or {}
+ if not inline.get("enabled"):
+ return None
+ try:
+ from potato.server_utils import judge_alignment as ja_mod
+ schemas = ja_mod.judge_scoped_schemas(config)
+ # Optional inline schema allow-list.
+ allow = set(inline.get("schemas", []) or [])
+ if allow:
+ schemas = [s for s in schemas if s.get("name") in allow]
+ if not schemas:
+ return None
+ schema_info = schemas[0]
+ schema_name = schema_info.get("name")
+
+ # 1) persisted prediction for the latest prompt version
+ preds = ja_mod.load_predictions(config)
+ version = ja_mod.latest_prompt_version(config)
+ pred = (preds.get(version, {}) or {}).get(f"{instance_id}::{schema_name}") if version else None
+
+ # 2) optional on-demand compute
+ if pred is None and inline.get("compute_on_demand"):
+ from potato.ai.judge import JudgeService
+ svc = JudgeService(config)
+ jp = svc.judge_instance(instance_id, schema_info, item.get_text())
+ if jp is not None:
+ ja_mod.save_prediction(config, jp)
+ pred = jp.to_dict()
+
+ if not pred:
+ return None
+
+ running = ja_mod.running_agreement(config, schema_name)
+ return {
+ "label": pred.get("predicted_label"),
+ "confidence": pred.get("confidence", 0.0),
+ "reasoning": pred.get("reasoning", ""),
+ "schema": schema_name,
+ "prompt_version": pred.get("prompt_version", ""),
+ "running": running,
+ }
+ except Exception as e:
+ logger.warning(f"Inline judge prediction failed for {instance_id}: {e}")
+ return None
+
+
+def _get_triage_info(item):
+ """Return a triage badge dict for the inline display, or None.
+
+ Gated by ``triage.show_badge`` (default true when triage is enabled). Reads
+ the priority/reason stored on the item's metadata by the triage scorer at
+ load/ingestion time. Only returns something when the item carries a reason
+ (i.e. a rule flagged it), so unflagged items show no banner.
+ """
+ triage_cfg = config.get("triage", {}) or {}
+ if not triage_cfg.get("enabled"):
+ return None
+ if not triage_cfg.get("show_badge", True):
+ return None
+ try:
+ reason = item.get_metadata("triage_reason")
+ if not reason:
+ return None
+ return {
+ "reason": reason,
+ "rule": item.get_metadata("triage_rule"),
+ "priority": item.get_metadata("triage_priority"),
+ }
+ except Exception as e:
+ logger.warning(f"Triage info failed for {item.get_id()}: {e}")
+ return None
+
+
+def get_label_suggestions(item, config, schema_content_to_prefill) -> set[SuggestedResponse]:
+
+ label_suggestions_json = set()
+ if 'label_suggestions' in item.get_data():
+ suggestions = item.get_data()['label_suggestions']
+ for schema in config['annotation_schemes']:
+ if schema['name'] not in suggestions:
+ continue
+ suggested_labels = suggestions[schema['name']]
+ if type(suggested_labels) == str:
+ suggested_labels = [suggested_labels]
+ elif type(suggested_labels) == list:
+ suggested_labels = suggested_labels
+ else:
+ logger.warning("Unsupported suggested label type %s, please check your input data" % type(suggested_labels))
+ continue
+
+ if not schema.get('label_suggestions') in ['highlight', 'prefill']:
+ logger.warning('The style of suggested labels is not defined, please check your configuration file.')
+ continue
+
+ label_suggestion = schema['label_suggestions']
+ for s in suggested_labels:
+ if label_suggestion == 'highlight':
+ #bad suggestion -- TODO make chance configurable
+ if random.randrange(0, 3) == 2:
+ label_suggestions_json.add(SuggestedResponse(schema['name'], random.choice(schema['labels'])))
+ continue
+
+ label_suggestions_json.add(SuggestedResponse(schema['name'], s))
+ elif label_suggestion == 'prefill':
+ schema_content_to_prefill.append({'name':schema['name'], 'label':s})
+ return label_suggestions_json
+
+def add_ai_hints(soup: BeautifulSoup, instance_id: str) -> BeautifulSoup:
+ """
+ Adds AI-generated hints to the page, if enabled. This is a hook for adding hints to the
+ page based on the instance that the user is currently annotating.
+ """
+
+ return soup
+
+# Shea: a function to get some suggestions from AI
+def ai_hints(text: str) -> str:
+ """
+ Returns the AI hints for the given instance.
+ """
+ import requests
+ logger.debug(f"AI hints text: {text}")
+ schemes = config.get("annotation_schemes", [])
+ if not schemes:
+ logger.warning("Cannot generate AI hints: no annotation_schemes configured")
+ return ""
+ description = schemes[0].get("description", "")
+ annotation_type = schemes[0].get("annotation_type", "")
+ logger.debug(f"AI hints description: {description}")
+ prompt = f'''You are assisting a user with an annotation task. Here is the annotation instruction: {description}
+ Here is the annotation task type: {annotation_type}
+ Here is the sentence (or item) to annotate: {text}
+ Based on the instruction, task type, and the given sentence, generate a short, helpful hint that guides the user on how to approach this annotation.
+ Also, give a short reason of your answer and the relevant part(keyword or text).
+ The hint should not provide the label or answer directly, but should highlight what the user might consider or look for.'''
+
+ try:
+ response = requests.post(
+ 'http://localhost:11434/api/generate',
+ json={
+ # 'model': 'llama3.2',
+ 'model': 'qwen3:0.6b',
+ 'prompt': prompt,
+ 'stream': False
+ },
+ timeout=5 # Add timeout to prevent hanging
+ )
+ result = response.json()['response']
+ logger.debug(f"AI hints response: {result}")
+ return result
+ except requests.exceptions.ConnectionError:
+ logger.warning("AI hints service not available (Ollama not running)")
+ return "AI hints are currently unavailable. Please proceed with manual annotation."
+ except requests.exceptions.Timeout:
+ logger.warning("AI hints service timeout")
+ return "AI hints service is slow to respond. Please proceed with manual annotation."
+ except Exception as e:
+ logger.error(f"Error getting AI hints: {e}")
+ return "AI hints are currently unavailable. Please proceed with manual annotation."
+
+
+
+def render_page_with_annotations_WEIRD(username):
+ """
+ Renders the annotation page with the current instance and any existing annotations.
+ """
+ user_state = get_user_state(username)
+ instance_id = user_state.get_current_instance_id()
+
+ # Get the annotations for this instance
+ annotations = get_annotations_for_user_on(username, instance_id)
+ span_annotations = get_span_annotations_for_user_on(username, instance_id)
+
+ # Get the instance data
+ item = get_item_state_manager().get_item(instance_id)
+ item_data = item.get_data()
+
+ # Get the HTML template
+ phase, page = user_state.get_current_phase_and_page()
+ html_fname = get_user_state_manager().get_phase_html_fname(phase, page)
+
+ # Get user progress information
+ progress = user_state.get_progress()
+
+ # Get UI configuration from config
+ ui_config = config.get("ui", {})
+
+ # Add layout configuration to ui_config for JavaScript access
+ if config.get("layout"):
+ ui_config = dict(ui_config) # Make a copy to avoid modifying the original
+ ui_config["layout"] = config["layout"]
+
+ return render_template(
+ html_fname,
+ instance_id=instance_id,
+ instance_data=item_data,
+ instance_record=item_data,
+ annotations=annotations,
+ span_annotations=span_annotations,
+ progress=progress,
+ username=username,
+ ui_config=ui_config,
+ annotation_codebook_url=_sanitize_codebook_url(config.get("annotation_codebook_url", "")),
+ )
+
+def randomize_options(soup, legend_names, seed):
+ random.seed(seed)
+
+ # Find all fieldsets in the soup
+ fieldsets = soup.find_all('fieldset')
+ if not fieldsets:
+ logger.debug("No fieldsets found.")
+ return soup
+
+ # Initialize a variable to track whether the legend is found
+ legend_found = False
+
+ # Iterate through each fieldset
+ for fieldset in fieldsets:
+ # Find the legend within the current fieldset
+ legend = fieldset.find('legend')
+ if legend and legend.string in legend_names:
+ # Legend found, set the flag
+ legend_found = True
+
+ # Determine the parent form's annotation type
+ parent_form = fieldset.find_parent('form')
+ annotation_type = parent_form.get('data-annotation-type', '') if parent_form else ''
+
+ if annotation_type == 'multirate':
+ # Multirate: shuffle rows in table (skip header row)
+ table = fieldset.find('table')
+ if not table:
+ logger.debug("Table not found within the fieldset.")
+ continue
+ tr_elements = table.find_all('tr')[1:]
+ random.shuffle(tr_elements)
+ for tr in tr_elements:
+ table.append(tr)
+
+ elif annotation_type == 'radio':
+ # Radio: shuffle elements
+ options_container = fieldset.find('div', class_='shadcn-radio-options')
+ if not options_container:
+ options_container = fieldset
+ option_divs = options_container.find_all('div', class_='shadcn-radio-option', recursive=False)
+ if option_divs:
+ random.shuffle(option_divs)
+ for div in option_divs:
+ options_container.append(div)
+
+ elif annotation_type == 'multiselect':
+ # Multiselect: shuffle checkbox option divs within the grid
+ grid = fieldset.find('div', class_='shadcn-multiselect-grid')
+ if not grid:
+ grid = fieldset
+ option_divs = grid.find_all('div', class_='shadcn-multiselect-option', recursive=False)
+ if option_divs:
+ random.shuffle(option_divs)
+ for div in option_divs:
+ grid.append(div)
+
+ elif annotation_type == 'select':
+ # Select: shuffle
elements (skip first if it's a placeholder)
+ select_el = fieldset.find('select')
+ if select_el:
+ options = select_el.find_all('option')
+ # Keep placeholder (first option with empty value) in place
+ placeholder = None
+ shuffleable = []
+ for opt in options:
+ if not placeholder and (opt.get('value', '') == '' or opt.get('disabled') is not None):
+ placeholder = opt
+ else:
+ shuffleable.append(opt)
+ random.shuffle(shuffleable)
+ # Clear and re-insert
+ select_el.clear()
+ if placeholder:
+ select_el.append(placeholder)
+ for opt in shuffleable:
+ select_el.append(opt)
+ else:
+ logger.debug(f"Unsupported annotation type for randomization: {annotation_type}")
+
+ # Check if any legend was found
+ if not legend_found:
+ logger.debug("No matching legends found within any fieldset.")
+
+ return soup
+
+def filter_dynamic_options(soup, schemes, instance_data):
+ """
+ Filter annotation options per instance based on dynamic_options config.
+
+ Each scheme can specify a `dynamic_options_field` that references a field
+ in the instance data containing a list of visible option labels.
+ Options not in the list are removed from the DOM.
+
+ Args:
+ soup: BeautifulSoup page object
+ schemes: List of annotation scheme dicts with dynamic_options enabled
+ instance_data: The current instance's data dictionary
+ """
+ for scheme in schemes:
+ field_name = scheme.get('dynamic_options_field', 'visible_labels')
+ visible_labels = instance_data.get(field_name)
+ if visible_labels is None:
+ continue # No filtering for this instance
+
+ if isinstance(visible_labels, str):
+ visible_labels = [visible_labels]
+
+ visible_set = set(visible_labels)
+ schema_name = scheme['name']
+ annotation_type = scheme['annotation_type']
+
+ # Find the form for this schema
+ form = soup.find('form', {'data-schema-name': schema_name})
+ if not form:
+ form = soup.find('form', id=schema_name)
+ if not form:
+ continue
+
+ if annotation_type == 'radio':
+ for option_div in form.find_all('div', class_='shadcn-radio-option'):
+ input_el = option_div.find('input', type='radio')
+ if input_el and input_el.get('value') not in visible_set:
+ option_div.decompose()
+
+ elif annotation_type == 'multiselect':
+ for option_div in form.find_all('div', class_='shadcn-multiselect-option'):
+ input_el = option_div.find('input', type='checkbox')
+ if input_el and input_el.get('value') not in visible_set:
+ option_div.decompose()
+
+ elif annotation_type == 'select':
+ select_el = form.find('select')
+ if select_el:
+ for option in select_el.find_all('option'):
+ val = option.get('value', '')
+ # Keep placeholder options (empty value or disabled)
+ if val == '' or option.get('disabled') is not None:
+ continue
+ if val not in visible_set:
+ option.decompose()
+
+ return soup
+
+
+def map_user_id_to_digit(user_id_str):
+ # Convert the user_id_str to an integer using a hash function
+ user_id_hash = hash(user_id_str)
+
+ # Map the hashed value to a single-digit integer using modulus
+ digit = abs(user_id_hash) % 9 + 1 # Add 1 to avoid 0
+
+ return digit
+
+def get_total_annotations():
+ """
+ Returns the total number of unique annotations done across all users.
+ """
+ total = 0
+ for username in get_users():
+ user_state = get_user_state(username)
+ total += user_state.get_annotation_count()
+
+ return total
+
+def update_annotation_state(username, form):
+ """
+ DEPRECATED: This function is no longer called during navigation.
+
+ Annotations are now saved in real-time via /updateinstance endpoint when users
+ interact with checkboxes, radio buttons, etc. This ensures proper timing tracking
+ for behavioral data analysis.
+
+ This function is kept for backward compatibility but should not be used in new code.
+ Use add_label_annotation() via /updateinstance instead.
+
+ Original purpose: Parses the state of the HTML form (what the user did to the
+ instance) and updates the state of the instance's annotations accordingly.
+ """
+
+ # Get what the user has already annotated, which might include this instance too
+ user_state = get_user_state(username)
+
+ # Jiaxin: the instance_id are changed to the user's local instance cursor
+ instance_id = user_state.get_current_instance_id()
+
+ schema_to_label_to_value = defaultdict(dict)
+
+ behavioral_data_dict = {}
+
+ did_change = False
+ for key in form:
+
+ # look for behavioral information regarding time, click, ...
+ if key[:9] == "behavior_":
+ behavioral_data_dict[key[9:]] = form[key]
+ continue
+
+ # Look for the marker that indicates an annotation label.
+ #
+ # NOTE: The span annotation uses radio buttons as well to figure out
+ # which label. These inputs are labeled with "span_label" so we can skip
+ # them as being actual annotatins (the spans are saved below though).
+ if ":::" in key and "span_label" not in key:
+
+ cols = key.split(":::")
+ annotation_schema = cols[0]
+ annotation_label = cols[1]
+ annotation_value = form[key]
+
+ # skip the input when it is an empty string (from a text-box)
+ if annotation_value == "":
+ continue
+
+ schema_to_label_to_value[annotation_schema][annotation_label] = annotation_value
+
+
+ # Span annotations are a bit funkier since we're getting raw HTML that
+ # we need to post-process on the server side.
+ span_annotations = None # Changed from [] to None to preserve existing spans during navigation
+ if "span-annotation" in form:
+ span_annotation_html = form["span-annotation"]
+ span_text, span_annotations = parse_html_span_annotation(span_annotation_html)
+
+ did_change = user_state.set_annotation(
+ instance_id, schema_to_label_to_value, span_annotations, behavioral_data_dict
+ )
+ # update the behavioral information regarding time only when the annotations are changed
+ if did_change:
+ # Include keyword highlight state in behavioral data for research tracking
+ keyword_state = user_state.get_keyword_highlight_state(instance_id)
+ if keyword_state:
+ behavioral_data_dict['keyword_highlights_shown'] = keyword_state.get('highlights', [])
+ user_state.instance_id_to_behavioral_data[instance_id] = behavioral_data_dict
+ return did_change
+
+
+def get_annotations_for_user_on(username, instance_id):
+ """
+ Returns the label-based annotations made by this user on the instance.
+
+ Handles two data formats:
+ 1. Label objects as keys: {Label("schema", "label"): value}
+ - Created by add_label_annotation() via /updateinstance endpoint
+ 2. Nested string dicts: {"schema": {"label": value}}
+ - Created by set_annotation() via /annotate navigation
+ """
+ # Normalize instance_id to string for consistent key lookup
+ instance_id = str(instance_id)
+
+ user_state = get_user_state(username)
+ logger.debug(f"instance_id: {instance_id}")
+ raw_annotations = user_state.get_label_annotations(instance_id)
+
+ # Process the raw annotations into the expected format
+ processed_annotations = {}
+ for label, value in raw_annotations.items():
+ # Check for Label object - the Label class uses 'schema' and 'name' attributes
+ # with get_schema() and get_name() getter methods
+ if hasattr(label, 'get_schema') and hasattr(label, 'get_name'):
+ # Format 1: Label object as key (from add_label_annotation via /updateinstance)
+ schema_name = label.get_schema()
+ label_name = label.get_name()
+ if schema_name not in processed_annotations:
+ processed_annotations[schema_name] = {}
+ processed_annotations[schema_name][label_name] = value
+ elif isinstance(label, str) and isinstance(value, dict):
+ # Format 2: Nested dict format {"schema": {"label": value}}
+ # (legacy format from set_annotation, kept for backward compatibility)
+ schema_name = label
+ if schema_name not in processed_annotations:
+ processed_annotations[schema_name] = {}
+ for label_name, label_value in value.items():
+ processed_annotations[schema_name][label_name] = label_value
+ else:
+ # Unknown format - log and skip
+ logger.warning(f"Skipping unknown annotation format: key={label}, value={value}")
+ continue
+
+ return processed_annotations
+
+
+def get_span_annotations_for_user_on(username, instance_id):
+ """
+ Returns the span annotations made by this user on the instance.
+ """
+ logger.debug(f"=== GET_SPAN_ANNOTATIONS_FOR_USER_ON START ===")
+ logger.debug(f"Username: {username}")
+ logger.debug(f"Instance ID: {instance_id}")
+
+ # Normalize instance_id to string for consistent key lookup
+ instance_id = str(instance_id)
+ logger.debug(f"Normalized Instance ID: {instance_id}")
+
+ user_state = get_user_state(username)
+ logger.debug(f"User state: {user_state}")
+
+ if not user_state:
+ logger.warning(f"User state not found for user: {username}")
+ return []
+
+ # DEBUG: Check if this instance has any span annotations at all
+ if hasattr(user_state, 'instance_id_to_span_to_value'):
+ logger.debug(f"User state instance_id_to_span_to_value keys: {list(user_state.instance_id_to_span_to_value.keys())}")
+
+ if instance_id in user_state.instance_id_to_span_to_value:
+ instance_spans = user_state.instance_id_to_span_to_value[instance_id]
+ logger.debug(f"Spans for instance {instance_id}: {instance_spans}")
+
+ # DEBUG: Show each span in detail
+ for span, value in instance_spans.items():
+ logger.debug(f"Span: {span}, Value: {value}")
+ if hasattr(span, 'get_schema'):
+ logger.debug(f" Schema: {span.get_schema()}")
+ logger.debug(f" Name: {span.get_name()}")
+ logger.debug(f" Start: {span.get_start()}")
+ logger.debug(f" End: {span.get_end()}")
+ logger.debug(f" ID: {span.get_id()}")
+ else:
+ logger.debug(f"No spans found for instance {instance_id}")
+
+ span_annotations_dict = user_state.get_span_annotations(instance_id)
+ logger.debug(f"Raw span annotations from user state: {span_annotations_dict}")
+
+ # Convert dictionary to list of SpanAnnotation objects
+ span_annotations = list(span_annotations_dict.keys()) if span_annotations_dict else []
+ logger.debug(f"Converted to list: {span_annotations}")
+
+ # Log details of each span
+ for span in span_annotations:
+ logger.debug(f"[DEBUG SPAN] schema={span.get_schema()} label={span.get_name()} start={span.get_start()} end={span.get_end()} id={span.get_id()}")
+
+ logger.debug(f"=== GET_SPAN_ANNOTATIONS_FOR_USER_ON END ===")
+ return span_annotations
+
+def parse_html_span_annotation(html):
+ """
+ Parses the HTML for span annotations and returns the text and a list of spans.
+ """
+ soup = BeautifulSoup(html, "html.parser")
+ spans = []
+ for span in soup.find_all("span", {"data-annotation": True}):
+ spans.append({
+ "text": span.get_text(),
+ "label": span["data-label"],
+ "start": int(span["data-start"]),
+ "end": int(span["data-end"])
+ })
+ return soup.get_text(), spans
+
+def validate_annotation(annotation):
+ """
+ Validates that the annotation is properly formatted.
+ """
+ # Simple validation for now - can be expanded as needed
+ return isinstance(annotation, dict)
+
+# Configure the Flask application
+def configure_app(flask_app):
+ """
+ Configure the Flask application instance
+
+ Args:
+ flask_app: The Flask application instance
+
+ Returns:
+ The configured Flask application instance
+ """
+ global app
+ app = flask_app
+
+ # Set application configuration
+ # Use a random secret key if sessions shouldn't persist, otherwise use the configured one
+ if config.get("persist_sessions", False):
+ secret_key = config.get("secret_key") or os.environ.get("POTATO_SECRET_KEY")
+ if not secret_key:
+ raise ValueError(
+ "persist_sessions is enabled but no secret_key is configured. "
+ "Set 'secret_key' in your config file or POTATO_SECRET_KEY environment variable."
+ )
+ app.secret_key = secret_key
+ else:
+ # Generate a random secret key to ensure sessions don't persist between restarts
+ import secrets
+ app.secret_key = secrets.token_hex(32)
+
+ app.permanent_session_lifetime = timedelta(days=config.get("session_lifetime_days", 2))
+
+ # Configure routes from the routes module
+ from routes import configure_routes
+ configure_routes(app, config)
+
+ # Conditionally register web agent blueprints only when needed
+ _register_web_agent_blueprints_if_needed(app, config)
+
+ return app
+
+
+def _register_web_agent_blueprints_if_needed(flask_app, config):
+ """Register web agent blueprints only if web_agent display types are configured."""
+ needs_web_agent = False
+ instance_display = config.get("instance_display", {})
+ fields = instance_display.get("fields", [])
+ for field in fields:
+ if isinstance(field, dict):
+ field_type = field.get("type", "")
+ if field_type in ("web_agent_trace", "web_agent_recorder"):
+ needs_web_agent = True
+ break
+
+ if needs_web_agent:
+ from potato.routes_web_agent import web_agent_bp
+ from potato.web_proxy import web_proxy_bp
+ flask_app.register_blueprint(web_agent_bp)
+ flask_app.register_blueprint(web_proxy_bp)
+ logger.info("Registered web agent blueprints (web_agent_trace/recorder display type detected)")
+
+ # Check for live_agent display type
+ needs_live_agent = False
+ for field in fields:
+ if isinstance(field, dict) and field.get("type") == "live_agent":
+ needs_live_agent = True
+ break
+
+ if needs_live_agent:
+ from potato.routes_live_agent import live_agent_bp
+ flask_app.register_blueprint(live_agent_bp)
+ # Store live_agent config on the app for route access
+ live_agent_config = config.get("live_agent", {})
+ flask_app.config["live_agent"] = live_agent_config
+ flask_app.config["live_agent_enabled"] = True
+ logger.info("Registered live agent blueprint (live_agent display type detected)")
+
+ # Register cleanup on app shutdown
+ import atexit
+ def _cleanup_agent_sessions():
+ try:
+ from potato.agent_runner_manager import AgentRunnerManager
+ AgentRunnerManager.clear_instance()
+ except Exception as e:
+ logger.warning(f"Failed to clean up agent sessions: {e}")
+ atexit.register(_cleanup_agent_sessions)
+
+ # Check for live_coding_agent display type
+ needs_live_coding_agent = False
+ for field in fields:
+ if isinstance(field, dict) and field.get("type") == "live_coding_agent":
+ needs_live_coding_agent = True
+ break
+
+ if needs_live_coding_agent:
+ from potato.routes_live_coding_agent import live_coding_agent_bp
+ flask_app.register_blueprint(live_coding_agent_bp)
+ flask_app.config["live_coding_agent_enabled"] = True
+ logger.info("Registered live coding agent blueprint (live_coding_agent display type detected)")
+
+ import atexit
+ def _cleanup_coding_agent_sessions():
+ try:
+ from potato.coding_agent_runner_manager import CodingAgentRunnerManager
+ CodingAgentRunnerManager.clear_instance()
+ except Exception as e:
+ logger.warning(f"Failed to clean up coding agent sessions: {e}")
+ atexit.register(_cleanup_coding_agent_sessions)
+
+ # Check for trace_ingestion config
+ trace_ingestion_config = config.get("trace_ingestion", {})
+ if trace_ingestion_config.get("enabled", False):
+ from potato.routes_trace_ingestion import trace_ingestion_bp
+ flask_app.register_blueprint(trace_ingestion_bp)
+ flask_app.config["trace_ingestion"] = trace_ingestion_config
+ logger.info("Registered trace ingestion blueprint")
+
+ # Start Langfuse poller if configured. Guard against `sources:` being
+ # present-but-null in YAML (e.g. when all entries are commented out),
+ # which yields None rather than an empty list.
+ sources = trace_ingestion_config.get("sources") or []
+ for source in sources:
+ if source.get("type") == "langfuse":
+ from potato.trace_ingestion.langfuse_poller import LangfusePoller
+ poller = LangfusePoller(
+ api_url=source.get("api_url", "https://cloud.langfuse.com"),
+ public_key=source.get("public_key", ""),
+ secret_key=source.get("secret_key", source.get("api_key", "")),
+ poll_interval=source.get("poll_interval", 30),
+ )
+ poller.start()
+ logger.info(f"Started Langfuse poller (interval={source.get('poll_interval', 30)}s)")
+
+ import atexit
+ atexit.register(poller.stop)
+
+# Function to create and initialize the Flask application
+def create_app(config_file=None):
+ """
+ Create and configure the Flask application.
+
+ When *config_file* is provided (e.g. from a gunicorn factory call like
+ ``gunicorn "potato.flask_server:create_app('config.yaml')"``), this
+ function also performs the full server initialization (config loading,
+ state managers, data loading, etc.) that ``run_server()`` normally does.
+ This makes it compatible with WSGI servers that use the factory pattern.
+
+ Args:
+ config_file: Optional path to a YAML config file. When provided,
+ ``init_config`` and ``_initialize_from_config`` are called
+ automatically so the app is ready to serve requests.
+
+ Returns:
+ The configured Flask application instance
+ """
+ global app
+
+ # If a config file was provided, perform full initialization first.
+ # This is the code path used by gunicorn / WSGI factory calls.
+ if config_file is not None:
+ _initialize_from_config(config_file)
+
+ # Initialize the app with explicit static folder configuration
+ static_folder = os.path.join(cur_program_dir, 'static')
+ app = Flask(__name__, static_folder=static_folder)
+ _apply_url_prefix_from_env(app)
+ _apply_proxy_fix_from_env(app)
+
+ # Configure Jinja2 to look in both main templates and generated templates directories
+ real_templates_dir = os.path.join(cur_program_dir, 'templates')
+ generated_templates_dir = os.path.join(real_templates_dir, 'generated')
+
+ # Ensure the generated directory exists
+ if not os.path.exists(generated_templates_dir):
+ os.makedirs(generated_templates_dir, exist_ok=True)
+
+ # Add the generated directory to the template search path
+ from jinja2 import ChoiceLoader, FileSystemLoader
+ app.jinja_loader = ChoiceLoader([
+ FileSystemLoader(real_templates_dir),
+ FileSystemLoader(generated_templates_dir)
+ ])
+
+ # Register HTML sanitization filters for XSS protection
+ from potato.server_utils.html_sanitizer import register_jinja_filters
+ register_jinja_filters(app)
+
+ # Configure the app
+ configure_app(app)
+
+ # Add context processor for debug settings and common config values
+ @app.context_processor
+ def inject_template_context():
+ """Inject debug settings and common config values into all templates."""
+ from potato.logging_config import is_ui_debug_enabled, is_server_debug_enabled
+
+ # Build ui_lang dict with defaults, overridden by config
+ ui_lang_defaults = {
+ # Navigation & controls
+ 'next_button': 'Next',
+ 'previous_button': 'Previous',
+ 'submit_button': 'Submit',
+ 'go_button': 'Go',
+ 'retry_button': 'Retry',
+ 'logout': 'Logout',
+ 'jump_prev_unannotated': 'Previous unannotated',
+ 'jump_next_unannotated': 'Next unannotated',
+ # Status indicators
+ 'labeled_badge': 'Labeled',
+ 'in_progress_badge': 'In Progress',
+ 'not_labeled_badge': 'Not labeled',
+ 'progress_label': 'Progress',
+ 'loading': 'Loading annotation interface...',
+ 'error_heading': 'Error',
+ # Annotation interface
+ 'adjudicate': 'Adjudicate',
+ 'codebook': 'Codebook',
+ 'instructions_heading': 'Instructions',
+ 'text_to_annotate': 'Text to Annotate:',
+ 'video_to_annotate': 'Video to Annotate:',
+ 'audio_to_annotate': 'Audio to Annotate:',
+ # Login / registration page
+ 'login_title': 'Annotation Platform',
+ 'login_subtitle_password': 'Sign in to continue',
+ 'login_subtitle_username': 'Enter your username to continue',
+ 'sign_in_tab': 'Sign In',
+ 'register_tab': 'Register',
+ 'username_label': 'Username',
+ 'password_label': 'Password',
+ 'sign_in_button': 'Sign In',
+ 'continue_button': 'Continue',
+ 'register_button': 'Register',
+ 'forgot_password': 'Forgot Password?',
+ 'username_placeholder': 'Enter your username',
+ 'choose_username_placeholder': 'Choose a username',
+ 'create_password_placeholder': 'Create a password',
+ 'sign_in_with': 'Sign in with',
+ 'or_divider': 'or',
+ # Footer
+ 'powered_by': 'Powered by',
+ 'cite_us': 'Cite Us',
+ # Language / direction
+ 'html_lang': 'en',
+ 'html_dir': 'ltr',
+ }
+ ui_lang_config = config.get('ui_language', {})
+ ui_lang = {**ui_lang_defaults, **ui_lang_config}
+
+ # Load project-level base CSS if configured
+ from potato.server_utils.front_end import load_project_base_css_html, resolve_header_logo_src
+ try:
+ project_base_css = load_project_base_css_html(config)
+ except FileNotFoundError:
+ project_base_css = ""
+ logger.warning("base_css file configured but not found")
+
+ # Resolve header logo (cached as data URL at startup)
+ header_logo_url = resolve_header_logo_src(config)
+
+ return {
+ 'ui_debug': is_ui_debug_enabled(),
+ 'server_debug': is_server_debug_enabled(),
+ 'debug_mode': config.get('debug', False),
+ 'debug_phase': config.get('debug_phase'),
+ # Add common config values needed by templates
+ 'annotation_task_name': config.get('annotation_task_name', 'Annotation Task'),
+ 'annotation_codebook_url': _sanitize_codebook_url(config.get('annotation_codebook_url', '')),
+ # Multilingual UI strings
+ 'ui_lang': ui_lang,
+ # Project-level base CSS
+ 'PROJECT_BASE_CSS': project_base_css,
+ # Header logo
+ 'header_logo_url': header_logo_url,
+ # Deployment URL prefix for client-side fetch/beacon/media URLs.
+ #
+ # Both proxy mechanisms converge on the WSGI SCRIPT_NAME: ProxyFix
+ # sets it from X-Forwarded-Prefix, and StaticPrefixMiddleware sets it
+ # from POTATO_URL_PREFIX. request.script_root surfaces that value, so
+ # it is the single source of truth and works in BOTH modes (including
+ # a real WSGI mount). We fall back to the env var defensively. When no
+ # proxy is involved script_root is "" and this is a no-op.
+ 'url_prefix': request.script_root or _normalize_url_prefix(
+ os.environ.get("POTATO_URL_PREFIX", "")
+ ),
+ # Custom footer HTML (e.g., promotional banner for HF Spaces)
+ 'custom_footer_html': config.get('custom_footer_html', ''),
+ }
+
+ return app
+
+
+def _env_flag_enabled(name: str) -> bool:
+ return os.environ.get(name, "").strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _normalize_url_prefix(prefix: str) -> str:
+ prefix = (prefix or "").strip().strip("/")
+ return f"/{prefix}" if prefix else ""
+
+
+def _apply_url_prefix_from_env(flask_app):
+ """
+ Force URL generation to include a deployment prefix when the reverse proxy
+ cannot send X-Forwarded-Prefix.
+
+ This is for deployments where nginx exposes Potato at /app1/ and strips
+ that prefix before proxying to Flask. Setting POTATO_URL_PREFIX=/app1 makes
+ url_for('static', ...) emit /app1/static/... while Flask still receives
+ backend paths such as /static/styles.css.
+ """
+ prefix = _normalize_url_prefix(os.environ.get("POTATO_URL_PREFIX", ""))
+ if not prefix:
+ return
+
+ class StaticPrefixMiddleware:
+ def __init__(self, wrapped_app, script_name):
+ self.wrapped_app = wrapped_app
+ self.script_name = script_name
+
+ def __call__(self, environ, start_response):
+ if not environ.get("SCRIPT_NAME"):
+ environ["SCRIPT_NAME"] = self.script_name
+ return self.wrapped_app(environ, start_response)
+
+ flask_app.wsgi_app = StaticPrefixMiddleware(flask_app.wsgi_app, prefix)
+
+
+def _apply_proxy_fix_from_env(flask_app):
+ """
+ Enable reverse-proxy prefix handling when explicitly requested.
+
+ Deployments mounted below a path such as /round1 need Flask to see the
+ forwarded prefix so url_for('static', ...) emits /round1/static/... instead
+ of /static/.... Without that, the annotation shell renders but CSS/JS 404s.
+ """
+ if not _env_flag_enabled("POTATO_PROXY_FIX"):
+ return
+
+ from werkzeug.middleware.proxy_fix import ProxyFix
+
+ flask_app.wsgi_app = ProxyFix(
+ flask_app.wsgi_app,
+ x_for=int(os.environ.get("POTATO_PROXY_FIX_X_FOR", "1")),
+ x_proto=int(os.environ.get("POTATO_PROXY_FIX_X_PROTO", "1")),
+ x_host=int(os.environ.get("POTATO_PROXY_FIX_X_HOST", "1")),
+ x_prefix=int(os.environ.get("POTATO_PROXY_FIX_X_PREFIX", "1")),
+ )
+
+
+def _initialize_from_config(config_file):
+ """
+ Perform full server initialization from a config file path.
+
+ This is used by ``create_app(config_file)`` for WSGI/gunicorn deployments
+ where ``run_server()`` is not called. It mirrors the initialization steps
+ in ``run_server()`` but constructs a minimal ``args`` namespace instead of
+ parsing sys.argv.
+ """
+ import types
+
+ # Build a minimal args namespace that init_config expects
+ args = types.SimpleNamespace(
+ config_file=config_file,
+ port=None,
+ verbose=False,
+ very_verbose=False,
+ debug=False,
+ debug_log=None,
+ debug_phase=None,
+ customjs=None,
+ customjs_hostname=None,
+ persist_sessions=False,
+ require_password=None,
+ mode="start",
+ )
+
+ # Initialize configuration
+ init_config(args)
+
+ # Handle require_no_password
+ if config.get("require_no_password", False):
+ config["require_password"] = False
+
+ # For URL-direct login, disable password requirement
+ login_config = config.get("login", {})
+ if login_config.get("type") in ["url_direct", "prolific"]:
+ config["require_password"] = False
+
+ # Set random seed default
+ if "random_seed" not in config:
+ config["random_seed"] = 1234
+
+ # Set up logging
+ setup_logging(
+ verbose=config.get("verbose", False),
+ debug=config.get("debug", False),
+ debug_log=config.get("debug_log"),
+ log_dir=config.get("output_annotation_dir"),
+ )
+
+ # Ensure directories exist
+ task_dir = config.get("task_dir", ".")
+ if not os.path.exists(task_dir):
+ os.makedirs(task_dir)
+
+ output_annotation_dir = config.get("output_annotation_dir", "annotation_output")
+ if not os.path.exists(output_annotation_dir):
+ os.makedirs(output_annotation_dir)
+
+ # Initialize authenticator
+ UserAuthenticator.init_from_config(config)
+
+ # Initialize state managers (singletons โ safe to call if already initialized)
+ init_user_state_manager(config)
+ init_item_state_manager(config)
+
+ # Initialize AI support if enabled
+ if config.get("ai_support", {}).get("enabled", False):
+ init_ai_prompt(config)
+ init_dynamic_ai_help()
+
+ # Load data
+ load_all_data(config)
+
+ # Initialize AI cache after data is loaded
+ if config.get("ai_support", {}).get("enabled", False):
+ init_ai_cache_manager()
+
+ # Initialize quality control if enabled
+ qc_enabled = (
+ config.get("attention_checks", {}).get("enabled", False)
+ or config.get("gold_standards", {}).get("enabled", False)
+ or config.get("pre_annotation", {}).get("enabled", False)
+ )
+ if qc_enabled:
+ qc_task_dir = config.get(
+ "task_dir", os.path.dirname(config.get("config_file", ""))
+ )
+ init_quality_control_manager(config, qc_task_dir)
+
+ # Initialize adjudication if configured
+ if config.get("adjudication", {}).get("enabled", False):
+ init_adjudication_manager(config)
+
+ # Initialize knowledge base manager
+ init_kb_manager(config)
+
+ # Initialize WaveformService for audio annotation
+ _init_waveform_service(config)
+
+ # Initialize webhook emitter if configured
+ if config.get("webhooks", {}).get("enabled", False):
+ from potato.webhooks import init_webhook_emitter
+ init_webhook_emitter(config)
+
+ # Initialize Solo Mode if enabled (parity with run_server() โ the
+ # WSGI/gunicorn factory path must initialize it too, otherwise the
+ # /solo routes exist but the manager is never created).
+ if config.get("solo_mode", {}).get("enabled", False):
+ logger.info("Initializing Solo Mode...")
+ init_solo_mode_manager(config)
+ logger.info("Solo Mode initialized successfully")
+
+ # Initialize QDA Mode if enabled (parity with run_server()).
+ if config.get("qda_mode", {}).get("enabled", False):
+ logger.info("Initializing QDA Mode...")
+ init_qda_mode_manager(config)
+ logger.info("QDA Mode initialized successfully")
+
+ # Initialize Judge Calibration if enabled (parity with run_server()).
+ if config.get("judge_calibration", {}).get("enabled", False):
+ logger.info("Initializing Judge Calibration...")
+ from potato.judge_calibration import init_judge_calibration_manager
+ init_judge_calibration_manager(config)
+ logger.info("Judge Calibration initialized successfully")
+
+ # Keep ICL prompts restricted to the codebook's current set: a
+ # change listener re-syncs live scheme labels on any codebook edit.
+ try:
+ from potato.codebook.schema_bridge import install_codebook_icl_sync
+ install_codebook_icl_sync()
+ except Exception as e:
+ logger.warning(f"Codebook ICL sync not installed: {e}")
+
+ # Auto-detect cases from item metadata (no-op unless cases enabled
+ # or QDA mode is on).
+ try:
+ from potato.cases import init_cases_from_config
+ init_cases_from_config(config)
+ except Exception as e:
+ logger.warning(f"Cases auto-detect skipped: {e}")
+
+ # Build the universal search index (no-op if search disabled).
+ try:
+ from potato.search import init_search_from_item_state
+ init_search_from_item_state(config)
+ except Exception as e:
+ logger.warning(f"Search index init skipped: {e}")
+
+ logger.info("Server initialization complete (WSGI factory mode)")
+
+
+def _init_waveform_service(config: dict) -> None:
+ """
+ Initialize the WaveformService for audio annotation if the config
+ includes audio_annotation schemes.
+
+ Args:
+ config: The application configuration dictionary
+ """
+ # Check if any audio_annotation schemes are configured
+ has_audio_annotation = False
+ annotation_schemes = config.get('annotation_schemes', [])
+ for scheme in annotation_schemes:
+ if scheme.get('annotation_type') == 'audio_annotation':
+ has_audio_annotation = True
+ break
+
+ if not has_audio_annotation:
+ logger.debug("No audio_annotation schemes found, skipping WaveformService initialization")
+ return
+
+ # Get waveform configuration
+ audio_config = config.get('audio_annotation', {})
+ task_dir = config.get('task_dir', '.')
+
+ # Default cache directory
+ cache_dir = audio_config.get('waveform_cache_dir')
+ if not cache_dir:
+ cache_dir = os.path.join(task_dir, 'waveform_cache')
+
+ # Make cache_dir absolute if relative
+ if not os.path.isabs(cache_dir):
+ cache_dir = os.path.join(task_dir, cache_dir)
+
+ # Get other configuration options
+ look_ahead = audio_config.get('waveform_look_ahead', 5)
+ cache_max_size = audio_config.get('waveform_cache_max_size', 100)
+ client_fallback_max_duration = audio_config.get('client_fallback_max_duration', 1800)
+
+ try:
+ from potato.server_utils.waveform_service import init_waveform_service, get_waveform_service
+
+ waveform_service = init_waveform_service(
+ cache_dir=cache_dir,
+ look_ahead=look_ahead,
+ cache_max_size=cache_max_size,
+ client_fallback_max_duration=client_fallback_max_duration
+ )
+
+ if waveform_service.is_available:
+ logger.info(f"WaveformService initialized with audiowaveform tool (cache: {cache_dir})")
+ else:
+ logger.warning("WaveformService initialized but audiowaveform tool not available. "
+ "Client-side waveform generation will be used as fallback.")
+
+ # Register cleanup handler
+ import atexit
+ def cleanup_waveform_service():
+ service = get_waveform_service()
+ if service:
+ service.stop_background_precompute()
+ logger.info("WaveformService background precompute stopped")
+ atexit.register(cleanup_waveform_service)
+
+ except Exception as e:
+ logger.error(f"Failed to initialize WaveformService: {e}")
+ logger.warning("Audio annotation will use client-side waveform generation only")
+
+
+def run_server(args):
+ """
+ Run the Flask server with the given arguments.
+ """
+
+ # Initialize configuration
+ init_config(args)
+
+ # Apply command line flags that override config settings
+ if args.require_password is not None:
+ # Command line flag takes precedence over config file
+ config["require_password"] = args.require_password
+ logger.debug(f"Password requirement set from command line: {args.require_password}")
+
+ # Handle require_no_password (inverse of require_password) for backwards compatibility
+ # This is commonly used in Prolific/MTurk configs
+ if config.get("require_no_password", False):
+ config["require_password"] = False
+ logger.debug("Password requirement disabled via require_no_password config")
+
+ # For URL-direct login, automatically disable password requirement
+ login_config = config.get('login', {})
+ if login_config.get('type') in ['url_direct', 'prolific']:
+ config["require_password"] = False
+ logger.debug(f"Password requirement disabled for {login_config.get('type')} login type")
+
+ # Override port from command line if specified
+ if args.port is not None:
+ config["port"] = args.port
+ logger.debug(f"Port set from command line: {args.port}")
+
+ # Apply persist_sessions flag from command line
+ config["persist_sessions"] = args.persist_sessions
+ logger.debug(f"Session persistence set from command line: {args.persist_sessions}")
+
+ # --- Add support for random seed ---
+ # Admins can set 'random_seed' in config YAML to control assignment randomness (default 1234)
+ if "random_seed" not in config:
+ config["random_seed"] = 1234
+ logger.info(f"Assignment random seed set to: {config['random_seed']}")
+ # -----------------------------------
+
+ # Set up centralized logging with appropriate verbosity
+ setup_logging(
+ verbose=config.get("verbose", False),
+ debug=config.get("debug", False) or config.get("very_verbose", False),
+ debug_log=config.get("debug_log"),
+ log_dir=config.get("output_annotation_dir"),
+ )
+
+ # Log debug phase setting if specified
+ if config.get("debug_phase"):
+ logger.info(f"Debug phase set to: {config['debug_phase']}")
+
+ # Ensure that the task directory exists
+ task_dir = config["task_dir"]
+ if not os.path.exists(task_dir):
+ os.makedirs(task_dir)
+
+ # Ensure that the output annotation directory exists
+ output_annotation_dir = config["output_annotation_dir"]
+ if not os.path.exists(output_annotation_dir):
+ os.makedirs(output_annotation_dir)
+
+ # Initialize authenticator
+ UserAuthenticator.init_from_config(config)
+
+ init_user_state_manager(config)
+ init_item_state_manager(config)
+
+ # Initialize AI prompt and wrapper BEFORE load_all_data() because
+ # template generation needs get_ai_wrapper() to return the AI help div
+ if config.get("ai_support", {}).get("enabled", False):
+ logger.info("Initializing AI prompt and wrapper...")
+ init_ai_prompt(config)
+ init_dynamic_ai_help()
+
+ load_all_data(config)
+
+ # Initialize AI cache manager AFTER load_all_data() because
+ # it needs item_state_manager to be fully initialized for warmup
+ if config.get("ai_support", {}).get("enabled", False):
+ logger.info("Initializing AI cache manager...")
+ init_ai_cache_manager()
+ logger.info("AI support initialized successfully")
+
+ # Initialize chat manager if enabled
+ if config.get("chat_support", {}).get("enabled", False):
+ logger.info("Initializing Chat Manager...")
+ from potato.chat_manager import init_chat_manager
+ init_chat_manager(config)
+ logger.info("Chat support initialized successfully")
+
+ # Initialize Solo Mode if enabled
+ if config.get("solo_mode", {}).get("enabled", False):
+ logger.info("Initializing Solo Mode...")
+ init_solo_mode_manager(config)
+ logger.info("Solo Mode initialized successfully")
+
+ # Initialize QDA Mode if enabled
+ if config.get("qda_mode", {}).get("enabled", False):
+ logger.info("Initializing QDA Mode...")
+ init_qda_mode_manager(config)
+ logger.info("QDA Mode initialized successfully")
+
+ # Initialize Judge Calibration if enabled
+ if config.get("judge_calibration", {}).get("enabled", False):
+ logger.info("Initializing Judge Calibration...")
+ from potato.judge_calibration import init_judge_calibration_manager
+ init_judge_calibration_manager(config)
+ logger.info("Judge Calibration initialized successfully")
+
+ # Keep ICL prompts restricted to the codebook's current set: a
+ # change listener re-syncs live scheme labels on any codebook edit.
+ try:
+ from potato.codebook.schema_bridge import install_codebook_icl_sync
+ install_codebook_icl_sync()
+ except Exception as e:
+ logger.warning(f"Codebook ICL sync not installed: {e}")
+
+ # Auto-detect cases from item metadata (no-op unless cases enabled
+ # or QDA mode is on).
+ try:
+ from potato.cases import init_cases_from_config
+ init_cases_from_config(config)
+ except Exception as e:
+ logger.warning(f"Cases auto-detect skipped: {e}")
+
+ # Build the universal search index (no-op if search disabled).
+ try:
+ from potato.search import init_search_from_item_state
+ init_search_from_item_state(config)
+ except Exception as e:
+ logger.warning(f"Search index init skipped: {e}")
+
+ # Initialize diversity manager if diversity_clustering strategy is used
+ # or if diversity_ordering is explicitly enabled
+ assignment_strategy = config.get("assignment_strategy", "")
+ if isinstance(assignment_strategy, dict):
+ assignment_strategy = assignment_strategy.get("name", "")
+ diversity_enabled = (
+ assignment_strategy == "diversity_clustering" or
+ config.get("diversity_ordering", {}).get("enabled", False)
+ )
+ if diversity_enabled:
+ logger.info("Initializing diversity manager...")
+ dm = init_diversity_manager(config)
+ if dm and dm.enabled:
+ # Prefill embeddings for first N items
+ _prefill_diversity_embeddings(dm, config)
+ logger.info("Diversity manager initialized successfully")
+
+ # Initialize embedding visualization manager (requires diversity manager)
+ from potato.embedding_visualization import init_embedding_viz_manager
+ viz_manager = init_embedding_viz_manager(config)
+ if viz_manager and viz_manager.enabled:
+ logger.info("Embedding visualization manager initialized")
+ else:
+ logger.debug(
+ "Embedding visualization not enabled. "
+ "Install umap-learn: pip install umap-learn"
+ )
+ else:
+ logger.warning(
+ "Diversity ordering requested but manager not enabled. "
+ "Install sentence-transformers and scikit-learn: "
+ "pip install sentence-transformers scikit-learn"
+ )
+
+ # Initialize active learning manager if enabled. The manager trains a
+ # classifier on annotations in a background thread and reorders the
+ # unlabeled pool by query strategy (uncertainty/BADGE/BALD/hybrid). Without
+ # this, `assignment_strategy: active_learning` falls back to random order.
+ if config.get('active_learning', {}).get('enabled', False):
+ try:
+ from potato.active_learning_manager import (
+ parse_active_learning_config, init_active_learning_manager,
+ )
+ al_cfg = parse_active_learning_config(config)
+ if al_cfg:
+ init_active_learning_manager(al_cfg)
+ logger.info(
+ "Active learning manager initialized (query_strategy=%s, "
+ "update_frequency=%s, schemas=%s)",
+ al_cfg.query_strategy, al_cfg.update_frequency, al_cfg.schema_names,
+ )
+ except Exception as e:
+ logger.warning(
+ "Active learning requested but could not initialize (%s). "
+ "Continuing without active-learning reordering.", e
+ )
+
+ # Initialize quality control manager if any QC features are enabled
+ qc_enabled = (
+ config.get('attention_checks', {}).get('enabled', False) or
+ config.get('gold_standards', {}).get('enabled', False) or
+ config.get('pre_annotation', {}).get('enabled', False)
+ )
+ if qc_enabled:
+ task_dir = config.get('task_dir', os.path.dirname(config.get('config_file', '')))
+ init_quality_control_manager(config, task_dir)
+ logger.info("Quality control manager initialized")
+
+ # Initialize adjudication manager if configured
+ if config.get('adjudication', {}).get('enabled', False):
+ init_adjudication_manager(config)
+ logger.info("Adjudication manager initialized")
+
+ # Initialize MACE competence estimation if configured
+ if config.get('mace', {}).get('enabled', False):
+ from potato.mace_manager import init_mace_manager
+ init_mace_manager(config)
+ logger.info("MACE manager initialized")
+
+ # Initialize knowledge base manager for entity linking
+ init_kb_manager(config)
+ logger.info("Knowledge base manager initialized")
+
+ # Initialize agent session manager if agent_proxy is configured
+ if "agent_proxy" in config:
+ from potato.agent_proxy import init_agent_session_manager
+ init_agent_session_manager(config)
+ logger.info(f"Agent session manager initialized (proxy type: {config['agent_proxy'].get('type', 'unknown')})")
+
+ # Initialize ExpertiseManager for dynamic category assignment
+ category_assignment = config.get('category_assignment', {})
+ dynamic_config = category_assignment.get('dynamic', {})
+ if dynamic_config.get('enabled', False):
+ expertise_manager = init_expertise_manager(config)
+ expertise_manager.start_background_worker()
+ logger.info("Dynamic category expertise enabled with background worker")
+
+ # Register cleanup handler for expertise manager
+ import atexit
+ def cleanup_expertise_manager():
+ em = get_expertise_manager()
+ if em:
+ em.stop_background_worker()
+ logger.info("Expertise manager background worker stopped")
+ atexit.register(cleanup_expertise_manager)
+
+ # Initialize ICL labeler for AI-assisted labeling if configured
+ icl_config = config.get('icl_labeling', {})
+ if icl_config.get('enabled', False):
+ from potato.ai.icl_labeler import init_icl_labeler, get_icl_labeler
+ icl_labeler = init_icl_labeler(config)
+ icl_labeler.start_background_worker()
+ logger.info("ICL (In-Context Learning) labeler enabled with background worker")
+
+ # Register cleanup handler for ICL labeler
+ import atexit
+ def cleanup_icl_labeler():
+ labeler = get_icl_labeler()
+ if labeler:
+ labeler.stop_background_worker()
+ labeler.save_state()
+ logger.info("ICL labeler background worker stopped and state saved")
+ atexit.register(cleanup_icl_labeler)
+
+ # Initialize directory watcher if configured
+ if "data_directory" in config:
+ from potato.directory_watcher import init_directory_watcher, get_directory_watcher
+ dw = init_directory_watcher(config)
+ if dw:
+ # Load all files from the directory
+ count = dw.load_directory()
+ logger.info(f"Loaded {count} instances from data_directory: {config['data_directory']}")
+
+ # Start watching if enabled
+ if config.get("watch_data_directory", False):
+ dw.start_watching()
+ logger.info(f"Directory watching enabled (poll interval: {config.get('watch_poll_interval', 5.0)}s)")
+
+ # Register cleanup handler
+ import atexit
+ def cleanup_directory_watcher():
+ watcher = get_directory_watcher()
+ if watcher:
+ watcher.stop()
+ logger.info("Directory watcher stopped")
+ atexit.register(cleanup_directory_watcher)
+
+ # Initialize webhook emitter if configured
+ if config.get('webhooks', {}).get('enabled', False):
+ from potato.webhooks import init_webhook_emitter, get_webhook_emitter
+ init_webhook_emitter(config)
+ logger.info("Webhook emitter initialized")
+
+ import atexit
+ def cleanup_webhook_emitter():
+ emitter = get_webhook_emitter()
+ if emitter:
+ emitter.stop()
+ logger.info("Webhook emitter stopped")
+ atexit.register(cleanup_webhook_emitter)
+
+ # Initialize HuggingFace CommitScheduler for live backup if configured
+ hf_backup = config.get('huggingface_backup', {})
+ if hf_backup.get('enabled', False):
+ try:
+ from huggingface_hub import CommitScheduler
+ task_dir = config.get('task_dir', '.')
+ output_dir = os.path.join(
+ task_dir,
+ config.get('output_annotation_dir', 'annotation_output')
+ )
+ hf_token = hf_backup.get('token') or os.environ.get('HF_TOKEN')
+ scheduler = CommitScheduler(
+ repo_id=hf_backup['repo_id'],
+ folder_path=output_dir,
+ token=hf_token,
+ private=hf_backup.get('private', True),
+ every=hf_backup.get('schedule_minutes', 5),
+ )
+ logger.info("HuggingFace CommitScheduler initialized: %s (every %d min)",
+ hf_backup['repo_id'], hf_backup.get('schedule_minutes', 5))
+ except ImportError:
+ logger.warning("huggingface_hub not installed, skipping CommitScheduler. "
+ "Install with: pip install huggingface_hub>=0.20.0")
+ except Exception as e:
+ logger.error("Failed to initialize HuggingFace CommitScheduler: %s", e)
+
+ # Initialize WaveformService for audio annotation if configured
+ _init_waveform_service(config)
+
+ # Log password requirement status
+ logger.info(f"Password authentication required: {config.get('require_password', True)}")
+
+ # Create and configure the Flask app
+ app = create_app()
+
+ # Initialize OAuth with Flask app if using OAuth authentication
+ # (must happen after create_app() since OAuth needs the Flask app instance)
+ auth_method = config.get("authentication", {}).get("method", "in_memory")
+ if auth_method == "oauth":
+ authenticator = UserAuthenticator.get_instance()
+ oauth_backend = authenticator.get_oauth_backend()
+ if oauth_backend:
+ oauth_backend.init_oauth(app)
+ logger.info("OAuth providers initialized with Flask app")
+
+ # Run the Flask app
+ host = config.get("host", "0.0.0.0")
+ port = config.get("port", 8000)
+ # Use threaded=True so background LLM calls (solo mode refinement,
+ # edge case synthesis, etc.) don't block the HTTP server.
+ app.run(host=host, port=port, debug=config.get("debug", False),
+ use_reloader=False, threaded=True)
+
+
+# Define the main entry point for the Flask server
+def main():
+ """
+ Main entry point for the Flask server
+
+ This function initializes the application, loads data, and runs the server.
+ """
+ # Parse command line arguments
+ args = arguments()
+
+ if args.mode == 'start':
+ logger.info("Starting server mode")
+ run_server(args)
+ elif args.mode == 'reset-password':
+ logger.info("Starting password reset")
+ from potato.password_reset import cli_reset_password
+ cli_reset_password(args)
+ return
+ elif args.mode == 'migrate':
+ logger.info("Starting config migration")
+ from potato.migrate_cli import main as migrate_main
+ # Pass arguments to migrate CLI
+ migrate_args = [args.config_file]
+ if args.to_v2:
+ migrate_args.append("--to-v2")
+ if args.output_file:
+ migrate_args.extend(["--output", args.output_file])
+ if args.in_place:
+ migrate_args.append("--in-place")
+ if args.dry_run:
+ migrate_args.append("--dry-run")
+ if args.quiet:
+ migrate_args.append("--quiet")
+ sys.exit(migrate_main(migrate_args))
+ elif args.mode == 'codebook':
+ logger.info("Starting codebook initialization")
+ from potato.codebook_cli import main as codebook_main
+ sys.exit(codebook_main([args.config_file]))
+
+ logger.info("Annotation platform shutdown complete")
+
+
+# Main entry point
+if __name__ == "__main__":
+ main()
+
diff --git a/potato/format_handlers/__init__.py b/potato/format_handlers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8a2bf0c77a41c548d0db02955d575fd422427824
--- /dev/null
+++ b/potato/format_handlers/__init__.py
@@ -0,0 +1,48 @@
+"""
+Format Handlers Module
+
+Provides a pluggable system for parsing various document formats (PDF, DOCX, Markdown,
+spreadsheets, source code) and extracting annotatable content with coordinate mappings.
+
+Usage:
+ from potato.format_handlers import format_handler_registry, FormatOutput
+
+ # Auto-detect and extract content from a file
+ output = format_handler_registry.extract("document.pdf")
+
+ # Access extracted content
+ text = output.text
+ html = output.rendered_html
+ coords = output.coordinate_map
+
+ # List supported formats
+ formats = format_handler_registry.get_supported_formats()
+"""
+
+from .base import BaseFormatHandler, FormatOutput
+from .registry import format_handler_registry, FormatHandlerRegistry
+from .coordinate_mapping import (
+ CoordinateMapper,
+ CharacterCoordinate,
+ PDFCoordinate,
+ SpreadsheetCoordinate,
+ DocumentCoordinate,
+ CodeCoordinate,
+ BoundingBoxCoordinate,
+)
+
+__all__ = [
+ # Core classes
+ "BaseFormatHandler",
+ "FormatOutput",
+ "FormatHandlerRegistry",
+ "format_handler_registry",
+ # Coordinate types
+ "CoordinateMapper",
+ "CharacterCoordinate",
+ "PDFCoordinate",
+ "SpreadsheetCoordinate",
+ "DocumentCoordinate",
+ "CodeCoordinate",
+ "BoundingBoxCoordinate",
+]
diff --git a/potato/format_handlers/base.py b/potato/format_handlers/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..0593db82304b72972ffc059480b8736770823392
--- /dev/null
+++ b/potato/format_handlers/base.py
@@ -0,0 +1,217 @@
+"""
+Base Format Handler
+
+Provides the abstract base class for format handlers and the FormatOutput
+dataclass that represents extracted content from documents.
+
+Usage:
+ from potato.format_handlers.base import BaseFormatHandler, FormatOutput
+
+ class MyFormatHandler(BaseFormatHandler):
+ format_name = "my_format"
+ supported_extensions = [".myf"]
+
+ def extract(self, file_path, options=None):
+ # Parse file and return FormatOutput
+ return FormatOutput(
+ text="extracted text",
+ rendered_html="rendered content
",
+ coordinate_map={...},
+ metadata={...}
+ )
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, List, Any, Optional
+from pathlib import Path
+
+
+@dataclass
+class FormatOutput:
+ """
+ Represents the extracted content from a document.
+
+ Attributes:
+ text: Plain text extracted from the document (for annotation)
+ rendered_html: HTML representation for display in the annotation UI
+ coordinate_map: Mapping from character offsets to format-specific coordinates
+ metadata: Additional document metadata (pages, structure, etc.)
+ format_name: Name of the format that produced this output
+ source_path: Path to the original source file
+ """
+ text: str
+ rendered_html: str
+ coordinate_map: Dict[str, Any] = field(default_factory=dict)
+ metadata: Dict[str, Any] = field(default_factory=dict)
+ format_name: str = ""
+ source_path: str = ""
+
+ def get_format_coords(self, start: int, end: int) -> Optional[Dict[str, Any]]:
+ """
+ Get format-specific coordinates for a character range.
+
+ Args:
+ start: Start character offset (inclusive)
+ end: End character offset (exclusive)
+
+ Returns:
+ Dictionary with format-specific coordinates, or None if not available
+ """
+ if not self.coordinate_map:
+ return None
+
+ # Look up coordinates using the mapping
+ # Implementation varies by format type
+ if "get_coords_for_range" in self.coordinate_map:
+ return self.coordinate_map["get_coords_for_range"](start, end)
+
+ return None
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to dictionary for serialization."""
+ return {
+ "text": self.text,
+ "rendered_html": self.rendered_html,
+ "metadata": self.metadata,
+ "format_name": self.format_name,
+ "source_path": self.source_path,
+ }
+
+
+class BaseFormatHandler(ABC):
+ """
+ Abstract base class for document format handlers.
+
+ Subclasses must implement the `extract` method and define class
+ attributes for format identification.
+
+ Class Attributes:
+ format_name: Unique identifier for this format (e.g., "pdf", "docx")
+ supported_extensions: List of file extensions this handler supports
+ description: Human-readable description of this format handler
+ requires_dependencies: List of optional dependencies needed
+ """
+
+ format_name: str = ""
+ supported_extensions: List[str] = []
+ description: str = ""
+ requires_dependencies: List[str] = []
+
+ @abstractmethod
+ def extract(self, file_path: str, options: Optional[Dict[str, Any]] = None) -> FormatOutput:
+ """
+ Extract annotatable content from a document.
+
+ Args:
+ file_path: Path to the document file
+ options: Optional configuration for extraction:
+ - extraction_mode: How to extract text (e.g., 'text', 'ocr', 'hybrid')
+ - preserve_layout: Whether to preserve document layout
+ - max_pages: Maximum pages to process (for paged documents)
+ - encoding: Text encoding to use
+
+ Returns:
+ FormatOutput with extracted text, rendered HTML, and coordinate mappings
+
+ Raises:
+ FileNotFoundError: If the file doesn't exist
+ ValueError: If the file format is not supported
+ ImportError: If required dependencies are not installed
+ """
+ pass
+
+ def can_handle(self, file_path: str) -> bool:
+ """
+ Check if this handler can process the given file.
+
+ Args:
+ file_path: Path to the file
+
+ Returns:
+ True if this handler supports the file's extension
+ """
+ ext = Path(file_path).suffix.lower()
+ return ext in self.supported_extensions
+
+ def check_dependencies(self) -> List[str]:
+ """
+ Check if required dependencies are installed.
+
+ Returns:
+ List of missing dependency names (empty if all installed)
+ """
+ missing = []
+ for dep in self.requires_dependencies:
+ try:
+ __import__(dep.replace("-", "_"))
+ except ImportError:
+ missing.append(dep)
+ return missing
+
+ def validate_file(self, file_path: str) -> List[str]:
+ """
+ Validate that a file can be processed.
+
+ Args:
+ file_path: Path to the file
+
+ Returns:
+ List of error messages (empty if valid)
+ """
+ errors = []
+ path = Path(file_path)
+
+ if not path.exists():
+ errors.append(f"File not found: {file_path}")
+ return errors
+
+ if not path.is_file():
+ errors.append(f"Not a file: {file_path}")
+ return errors
+
+ if not self.can_handle(file_path):
+ errors.append(
+ f"Unsupported extension '{path.suffix}'. "
+ f"Supported: {', '.join(self.supported_extensions)}"
+ )
+
+ missing_deps = self.check_dependencies()
+ if missing_deps:
+ errors.append(
+ f"Missing dependencies for {self.format_name}: "
+ f"{', '.join(missing_deps)}. "
+ f"Install with: pip install {' '.join(missing_deps)}"
+ )
+
+ return errors
+
+ def get_default_options(self) -> Dict[str, Any]:
+ """
+ Get default extraction options for this handler.
+
+ Override in subclasses to provide format-specific defaults.
+
+ Returns:
+ Dictionary of default option values
+ """
+ return {
+ "preserve_layout": False,
+ "max_pages": None,
+ "encoding": "utf-8",
+ }
+
+ def merge_options(self, options: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+ """
+ Merge user options with defaults.
+
+ Args:
+ options: User-provided options
+
+ Returns:
+ Merged options dictionary
+ """
+ merged = self.get_default_options()
+ if options:
+ merged.update(options)
+ return merged
diff --git a/potato/format_handlers/code_handler.py b/potato/format_handlers/code_handler.py
new file mode 100644
index 0000000000000000000000000000000000000000..c03d9a8c21bddd40bfb263677df33ac3dbf57186
--- /dev/null
+++ b/potato/format_handlers/code_handler.py
@@ -0,0 +1,380 @@
+"""
+Source Code Format Handler
+
+Parses source code files with syntax highlighting and line/column
+coordinate mapping for code annotation tasks.
+
+Usage:
+ from potato.format_handlers.code_handler import CodeHandler
+
+ handler = CodeHandler()
+ output = handler.extract("script.py", {
+ "show_line_numbers": True,
+ "highlight_syntax": True,
+ })
+"""
+
+from typing import Dict, List, Any, Optional
+from pathlib import Path
+import html
+import logging
+import re
+
+from .base import BaseFormatHandler, FormatOutput
+from .coordinate_mapping import CoordinateMapper, CodeCoordinate
+
+logger = logging.getLogger(__name__)
+
+# Check if Pygments is available
+try:
+ from pygments import highlight
+ from pygments.lexers import get_lexer_by_name, get_lexer_for_filename, guess_lexer
+ from pygments.formatters import HtmlFormatter
+ from pygments.token import Token
+ PYGMENTS_AVAILABLE = True
+except ImportError:
+ PYGMENTS_AVAILABLE = False
+
+
+# Common source code extensions and their languages
+LANGUAGE_MAP = {
+ ".py": "python",
+ ".js": "javascript",
+ ".jsx": "jsx",
+ ".ts": "typescript",
+ ".tsx": "tsx",
+ ".java": "java",
+ ".c": "c",
+ ".cpp": "cpp",
+ ".h": "c",
+ ".hpp": "cpp",
+ ".cs": "csharp",
+ ".go": "go",
+ ".rs": "rust",
+ ".rb": "ruby",
+ ".php": "php",
+ ".swift": "swift",
+ ".kt": "kotlin",
+ ".scala": "scala",
+ ".r": "r",
+ ".R": "r",
+ ".sql": "sql",
+ ".sh": "bash",
+ ".bash": "bash",
+ ".zsh": "zsh",
+ ".ps1": "powershell",
+ ".yaml": "yaml",
+ ".yml": "yaml",
+ ".json": "json",
+ ".xml": "xml",
+ ".html": "html",
+ ".css": "css",
+ ".scss": "scss",
+ ".less": "less",
+ ".lua": "lua",
+ ".pl": "perl",
+ ".m": "matlab",
+ ".jl": "julia",
+ ".hs": "haskell",
+ ".ml": "ocaml",
+ ".ex": "elixir",
+ ".exs": "elixir",
+ ".erl": "erlang",
+ ".clj": "clojure",
+ ".lisp": "lisp",
+ ".vim": "vim",
+ ".dockerfile": "docker",
+ ".tf": "terraform",
+ ".proto": "protobuf",
+ ".graphql": "graphql",
+}
+
+
+class CodeHandler(BaseFormatHandler):
+ """
+ Handler for source code files.
+
+ Provides syntax highlighting via Pygments and line/column coordinate
+ mapping for code annotation.
+ """
+
+ format_name = "code"
+ supported_extensions = list(LANGUAGE_MAP.keys())
+ description = "Source code with syntax highlighting and line/column mapping"
+ requires_dependencies = ["pygments"]
+
+ def get_default_options(self) -> Dict[str, Any]:
+ """Get default extraction options."""
+ return {
+ "highlight_syntax": True,
+ "show_line_numbers": True,
+ "language": None, # Auto-detect from extension
+ "tab_size": 4,
+ "max_lines": None,
+ "start_line": 1,
+ "extract_structure": True, # Extract function/class names
+ }
+
+ def extract(
+ self,
+ file_path: str,
+ options: Optional[Dict[str, Any]] = None
+ ) -> FormatOutput:
+ """
+ Parse and render a source code file.
+
+ Args:
+ file_path: Path to the source code file
+ options: Extraction options:
+ - highlight_syntax: Apply syntax highlighting
+ - show_line_numbers: Include line numbers in output
+ - language: Override language detection
+ - tab_size: Spaces per tab for rendering
+ - max_lines: Limit number of lines
+ - extract_structure: Extract function/class definitions
+
+ Returns:
+ FormatOutput with code text, highlighted HTML, and coordinates
+ """
+ opts = self.merge_options(options)
+ path = Path(file_path)
+
+ # Read source file
+ source_text = path.read_text(encoding="utf-8")
+
+ # Expand tabs if needed
+ if opts.get("tab_size"):
+ source_text = source_text.expandtabs(opts["tab_size"])
+
+ # Build line index and coordinates
+ lines = source_text.split("\n")
+ mapper = CoordinateMapper()
+ line_offsets = []
+ current_offset = 0
+
+ # Apply line limits
+ start_line = opts.get("start_line", 1) - 1 # Convert to 0-indexed
+ max_lines = opts.get("max_lines")
+ end_line = min(len(lines), start_line + max_lines) if max_lines else len(lines)
+
+ # Build coordinate mappings for each line
+ for line_num, line in enumerate(lines):
+ line_start = current_offset
+ line_end = current_offset + len(line)
+ line_offsets.append((line_start, line_end))
+
+ # Only map lines within our range
+ if start_line <= line_num < end_line:
+ mapper.add_mapping(
+ line_start,
+ line_end,
+ CodeCoordinate(
+ line=line_num + 1, # 1-indexed
+ column=1,
+ )
+ )
+
+ current_offset = line_end + 1 # +1 for newline
+
+ # Extract text for the requested range
+ if start_line > 0 or max_lines:
+ display_lines = lines[start_line:end_line]
+ display_text = "\n".join(display_lines)
+ else:
+ display_text = source_text
+
+ # Detect language
+ language = opts.get("language")
+ if not language:
+ ext = path.suffix.lower()
+ language = LANGUAGE_MAP.get(ext, "text")
+
+ # Render HTML
+ if opts.get("highlight_syntax") and PYGMENTS_AVAILABLE:
+ rendered_html = self._render_highlighted(
+ display_text, language, opts, start_line + 1
+ )
+ else:
+ rendered_html = self._render_plain(
+ display_text, opts, start_line + 1
+ )
+
+ # Extract code structure
+ structure = []
+ if opts.get("extract_structure"):
+ structure = self._extract_structure(source_text, language)
+
+ metadata = {
+ "format": "code",
+ "source_file": str(file_path),
+ "language": language,
+ "line_count": len(lines),
+ "char_count": len(source_text),
+ "displayed_lines": (start_line + 1, end_line),
+ "structure": structure,
+ }
+
+ coord_dict = mapper.to_dict()
+ coord_dict["get_coords_for_range"] = mapper.get_coords_for_range
+
+ return FormatOutput(
+ text=display_text,
+ rendered_html=rendered_html,
+ coordinate_map=coord_dict,
+ metadata=metadata,
+ format_name=self.format_name,
+ source_path=str(file_path),
+ )
+
+ def _render_highlighted(
+ self,
+ code: str,
+ language: str,
+ opts: Dict[str, Any],
+ start_line: int
+ ) -> str:
+ """
+ Render code with Pygments syntax highlighting.
+ """
+ try:
+ lexer = get_lexer_by_name(language, stripall=False)
+ except Exception:
+ try:
+ lexer = guess_lexer(code)
+ except Exception:
+ lexer = get_lexer_by_name("text")
+
+ # Configure formatter
+ formatter_opts = {
+ "cssclass": "code-highlight",
+ "linenos": opts.get("show_line_numbers", True),
+ "linenostart": start_line,
+ "lineanchors": "line",
+ "anchorlinenos": True,
+ }
+
+ if opts.get("show_line_numbers"):
+ formatter_opts["linenos"] = "table"
+
+ formatter = HtmlFormatter(**formatter_opts)
+ highlighted = highlight(code, lexer, formatter)
+
+ # Wrap in container
+ return f'{highlighted}
'
+
+ def _render_plain(
+ self,
+ code: str,
+ opts: Dict[str, Any],
+ start_line: int
+ ) -> str:
+ """
+ Render code as plain text with optional line numbers.
+ """
+ lines = code.split("\n")
+ html_parts = []
+
+ html_parts.append('')
+ html_parts.append('
')
+
+ for i, line in enumerate(lines, start=start_line):
+ escaped_line = html.escape(line) or " "
+
+ if opts.get("show_line_numbers"):
+ html_parts.append(
+ f''
+ f'{i} '
+ f'{escaped_line} '
+ f' '
+ )
+ else:
+ html_parts.append(
+ f''
+ f'{escaped_line} '
+ f' '
+ )
+
+ html_parts.append('
')
+ html_parts.append('
')
+
+ return "\n".join(html_parts)
+
+ def _extract_structure(
+ self,
+ code: str,
+ language: str
+ ) -> List[Dict[str, Any]]:
+ """
+ Extract code structure (functions, classes) using pattern matching.
+
+ This is a simplified extraction that works for common languages.
+ For production use, consider using language-specific parsers.
+ """
+ structure = []
+
+ # Pattern definitions for common languages
+ patterns = {
+ "python": {
+ "function": r"^\s*(?:async\s+)?def\s+(\w+)\s*\(",
+ "class": r"^\s*class\s+(\w+)\s*[:\(]",
+ },
+ "javascript": {
+ "function": r"(?:function\s+(\w+)|(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s*)?\([^)]*\)\s*=>|(?:const|let|var)\s+(\w+)\s*=\s*function)",
+ "class": r"class\s+(\w+)",
+ },
+ "java": {
+ "function": r"(?:public|private|protected|static|\s)+[\w<>\[\]]+\s+(\w+)\s*\([^\)]*\)\s*(?:throws\s+[\w,\s]+)?\s*\{",
+ "class": r"(?:public|private|protected)?\s*class\s+(\w+)",
+ },
+ "go": {
+ "function": r"func\s+(?:\([^)]+\)\s+)?(\w+)\s*\(",
+ "class": r"type\s+(\w+)\s+struct",
+ },
+ "rust": {
+ "function": r"(?:pub\s+)?fn\s+(\w+)",
+ "class": r"(?:pub\s+)?struct\s+(\w+)",
+ },
+ }
+
+ # Get patterns for this language (or use generic)
+ lang_patterns = patterns.get(language, {})
+
+ lines = code.split("\n")
+ for line_num, line in enumerate(lines, start=1):
+ # Check for functions
+ if "function" in lang_patterns:
+ match = re.search(lang_patterns["function"], line)
+ if match:
+ name = next((g for g in match.groups() if g), None)
+ if name:
+ structure.append({
+ "type": "function",
+ "name": name,
+ "line": line_num,
+ })
+
+ # Check for classes/structs
+ if "class" in lang_patterns:
+ match = re.search(lang_patterns["class"], line)
+ if match:
+ name = match.group(1)
+ structure.append({
+ "type": "class",
+ "name": name,
+ "line": line_num,
+ })
+
+ return structure
+
+ def get_language(self, file_path: str) -> str:
+ """
+ Detect the programming language from a file path.
+
+ Args:
+ file_path: Path to the source file
+
+ Returns:
+ Language identifier
+ """
+ ext = Path(file_path).suffix.lower()
+ return LANGUAGE_MAP.get(ext, "text")
diff --git a/potato/format_handlers/coordinate_mapping.py b/potato/format_handlers/coordinate_mapping.py
new file mode 100644
index 0000000000000000000000000000000000000000..ddd8c1b40310d80cc6c30684ed41fbcfee73acc8
--- /dev/null
+++ b/potato/format_handlers/coordinate_mapping.py
@@ -0,0 +1,504 @@
+"""
+Coordinate Mapping Utilities
+
+Provides data structures and utilities for mapping character offsets
+to format-specific coordinates (page/bbox for PDF, row/col for spreadsheets, etc.).
+
+Usage:
+ from potato.format_handlers.coordinate_mapping import (
+ CoordinateMapper,
+ PDFCoordinate,
+ SpreadsheetCoordinate,
+ )
+
+ # Build a coordinate map during extraction
+ mapper = CoordinateMapper()
+ mapper.add_mapping(0, 100, PDFCoordinate(page=1, bbox=[10, 20, 200, 30]))
+
+ # Look up coordinates for a span
+ coords = mapper.get_coords_for_range(50, 75)
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Any, Optional, Tuple, Union
+from bisect import bisect_left, bisect_right
+import json
+
+
+@dataclass
+class CharacterCoordinate:
+ """Base coordinate type representing a character position."""
+ offset: int
+ format_type: str = "character"
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {"format": self.format_type, "offset": self.offset}
+
+
+@dataclass
+class PDFCoordinate:
+ """
+ Coordinate for PDF documents.
+
+ Attributes:
+ page: Page number (1-indexed)
+ bbox: Bounding box [x0, y0, x1, y1] in PDF points
+ line: Optional line number on the page
+ """
+ page: int
+ bbox: List[float] = field(default_factory=list)
+ line: Optional[int] = None
+ format_type: str = "pdf"
+
+ def to_dict(self) -> Dict[str, Any]:
+ result = {
+ "format": self.format_type,
+ "page": self.page,
+ }
+ if self.bbox:
+ result["bbox"] = self.bbox
+ if self.line is not None:
+ result["line"] = self.line
+ return result
+
+
+@dataclass
+class SpreadsheetCoordinate:
+ """
+ Coordinate for spreadsheet documents.
+
+ Attributes:
+ row: Row number (0-indexed internally, displayed as 1-indexed)
+ col: Column number (0-indexed internally)
+ cell_ref: Cell reference in A1 notation (e.g., "B5")
+ sheet: Sheet name (for multi-sheet documents)
+ """
+ row: int
+ col: Optional[int] = None
+ cell_ref: Optional[str] = None
+ sheet: Optional[str] = None
+ format_type: str = "spreadsheet"
+
+ def to_dict(self) -> Dict[str, Any]:
+ result = {
+ "format": self.format_type,
+ "row": self.row + 1, # Convert to 1-indexed for output
+ }
+ if self.col is not None:
+ result["col"] = self.col + 1
+ if self.cell_ref:
+ result["cell_ref"] = self.cell_ref
+ if self.sheet:
+ result["sheet"] = self.sheet
+ return result
+
+
+@dataclass
+class DocumentCoordinate:
+ """
+ Coordinate for document formats (DOCX, Markdown).
+
+ Attributes:
+ paragraph_id: Unique identifier for the paragraph
+ local_offset: Character offset within the paragraph
+ section: Optional section name/number
+ heading_level: If in a heading, its level (1-6)
+ """
+ paragraph_id: str
+ local_offset: int = 0
+ section: Optional[str] = None
+ heading_level: Optional[int] = None
+ format_type: str = "document"
+
+ def to_dict(self) -> Dict[str, Any]:
+ result = {
+ "format": self.format_type,
+ "paragraph_id": self.paragraph_id,
+ "local_offset": self.local_offset,
+ }
+ if self.section:
+ result["section"] = self.section
+ if self.heading_level:
+ result["heading_level"] = self.heading_level
+ return result
+
+
+@dataclass
+class CodeCoordinate:
+ """
+ Coordinate for source code files.
+
+ Attributes:
+ line: Line number (1-indexed)
+ column: Column number (1-indexed)
+ function_name: Name of the containing function (if any)
+ class_name: Name of the containing class (if any)
+ """
+ line: int
+ column: int = 1
+ function_name: Optional[str] = None
+ class_name: Optional[str] = None
+ format_type: str = "code"
+
+ def to_dict(self) -> Dict[str, Any]:
+ result = {
+ "format": self.format_type,
+ "line": self.line,
+ "column": self.column,
+ }
+ if self.function_name:
+ result["function_name"] = self.function_name
+ if self.class_name:
+ result["class_name"] = self.class_name
+ return result
+
+
+@dataclass
+class BoundingBoxCoordinate:
+ """
+ Coordinate for bounding box annotations on images/PDF pages.
+
+ Attributes:
+ page: Page number (1-indexed, for PDFs)
+ bbox: Bounding box [x, y, width, height] in normalized coordinates (0-1)
+ bbox_pixels: Optional bounding box in pixel coordinates
+ label: Label/class for the bounding box
+ confidence: Optional confidence score (0-1)
+ """
+ page: int
+ bbox: List[float] # [x, y, width, height] normalized 0-1
+ bbox_pixels: Optional[List[float]] = None # [x, y, width, height] in pixels
+ label: Optional[str] = None
+ confidence: Optional[float] = None
+ format_type: str = "bounding_box"
+
+ def to_dict(self) -> Dict[str, Any]:
+ result = {
+ "format": self.format_type,
+ "page": self.page,
+ "bbox": self.bbox,
+ }
+ if self.bbox_pixels:
+ result["bbox_pixels"] = self.bbox_pixels
+ if self.label:
+ result["label"] = self.label
+ if self.confidence is not None:
+ result["confidence"] = self.confidence
+ return result
+
+ @classmethod
+ def from_pixel_coords(
+ cls,
+ page: int,
+ x: float,
+ y: float,
+ width: float,
+ height: float,
+ page_width: float,
+ page_height: float,
+ label: Optional[str] = None
+ ) -> "BoundingBoxCoordinate":
+ """
+ Create a BoundingBoxCoordinate from pixel coordinates.
+
+ Args:
+ page: Page number (1-indexed)
+ x, y, width, height: Bounding box in pixels
+ page_width, page_height: Page dimensions in pixels
+ label: Optional label for the box
+
+ Returns:
+ BoundingBoxCoordinate with normalized coords
+ """
+ return cls(
+ page=page,
+ bbox=[
+ x / page_width,
+ y / page_height,
+ width / page_width,
+ height / page_height
+ ],
+ bbox_pixels=[x, y, width, height],
+ label=label
+ )
+
+ def to_pixel_coords(
+ self,
+ page_width: float,
+ page_height: float
+ ) -> List[float]:
+ """
+ Convert normalized coordinates to pixels.
+
+ Args:
+ page_width: Page width in pixels
+ page_height: Page height in pixels
+
+ Returns:
+ [x, y, width, height] in pixels
+ """
+ return [
+ self.bbox[0] * page_width,
+ self.bbox[1] * page_height,
+ self.bbox[2] * page_width,
+ self.bbox[3] * page_height
+ ]
+
+
+# Type alias for any coordinate type
+Coordinate = Union[
+ CharacterCoordinate,
+ PDFCoordinate,
+ SpreadsheetCoordinate,
+ DocumentCoordinate,
+ CodeCoordinate,
+ BoundingBoxCoordinate,
+]
+
+
+@dataclass
+class CoordinateMapping:
+ """
+ Maps a character range to format-specific coordinates.
+
+ Attributes:
+ start: Start character offset (inclusive)
+ end: End character offset (exclusive)
+ coordinate: Format-specific coordinate data
+ """
+ start: int
+ end: int
+ coordinate: Coordinate
+
+ def contains(self, offset: int) -> bool:
+ """Check if an offset falls within this mapping."""
+ return self.start <= offset < self.end
+
+ def overlaps(self, start: int, end: int) -> bool:
+ """Check if a range overlaps with this mapping."""
+ return self.start < end and start < self.end
+
+
+class CoordinateMapper:
+ """
+ Manages mappings from character offsets to format-specific coordinates.
+
+ Provides efficient lookup of coordinates for character ranges.
+ """
+
+ def __init__(self):
+ self._mappings: List[CoordinateMapping] = []
+ self._sorted = True
+ self._start_offsets: List[int] = [] # For binary search
+
+ def add_mapping(
+ self,
+ start: int,
+ end: int,
+ coordinate: Coordinate
+ ) -> None:
+ """
+ Add a mapping from character range to coordinate.
+
+ Args:
+ start: Start character offset (inclusive)
+ end: End character offset (exclusive)
+ coordinate: Format-specific coordinate data
+ """
+ mapping = CoordinateMapping(start=start, end=end, coordinate=coordinate)
+ self._mappings.append(mapping)
+ self._sorted = False
+
+ def _ensure_sorted(self) -> None:
+ """Sort mappings by start offset if needed."""
+ if not self._sorted:
+ self._mappings.sort(key=lambda m: m.start)
+ self._start_offsets = [m.start for m in self._mappings]
+ self._sorted = True
+
+ def get_coordinate_at(self, offset: int) -> Optional[Coordinate]:
+ """
+ Get the coordinate at a specific character offset.
+
+ Args:
+ offset: Character offset
+
+ Returns:
+ Coordinate if found, None otherwise
+ """
+ self._ensure_sorted()
+
+ # Binary search for the potential mapping
+ idx = bisect_right(self._start_offsets, offset) - 1
+ if idx >= 0 and self._mappings[idx].contains(offset):
+ return self._mappings[idx].coordinate
+ return None
+
+ def get_coords_for_range(
+ self,
+ start: int,
+ end: int
+ ) -> Optional[Dict[str, Any]]:
+ """
+ Get format coordinates for a character range.
+
+ For spans crossing multiple coordinates (e.g., multiple PDF pages),
+ returns the coordinates of the first character.
+
+ Args:
+ start: Start character offset (inclusive)
+ end: End character offset (exclusive)
+
+ Returns:
+ Dictionary with format-specific coordinates, or None if not found
+ """
+ coord = self.get_coordinate_at(start)
+ if coord:
+ return coord.to_dict()
+ return None
+
+ def get_all_coords_for_range(
+ self,
+ start: int,
+ end: int
+ ) -> List[Dict[str, Any]]:
+ """
+ Get all coordinates that overlap with a character range.
+
+ Useful for spans crossing multiple structural elements.
+
+ Args:
+ start: Start character offset (inclusive)
+ end: End character offset (exclusive)
+
+ Returns:
+ List of coordinate dictionaries
+ """
+ self._ensure_sorted()
+ coords = []
+
+ for mapping in self._mappings:
+ if mapping.overlaps(start, end):
+ coords.append(mapping.coordinate.to_dict())
+ elif mapping.start >= end:
+ break # No more overlapping mappings
+
+ return coords
+
+ def get_mapping_count(self) -> int:
+ """Get the number of mappings stored."""
+ return len(self._mappings)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """
+ Export mappings as a dictionary.
+
+ Returns:
+ Dictionary representation of all mappings
+ """
+ self._ensure_sorted()
+ return {
+ "mappings": [
+ {
+ "start": m.start,
+ "end": m.end,
+ "coordinate": m.coordinate.to_dict()
+ }
+ for m in self._mappings
+ ]
+ }
+
+ def to_json(self) -> str:
+ """Export mappings as JSON string."""
+ return json.dumps(self.to_dict(), indent=2)
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "CoordinateMapper":
+ """
+ Create a CoordinateMapper from a dictionary.
+
+ Args:
+ data: Dictionary with mappings
+
+ Returns:
+ New CoordinateMapper instance
+ """
+ mapper = cls()
+ for m in data.get("mappings", []):
+ coord_data = m["coordinate"]
+ format_type = coord_data.get("format", "character")
+
+ # Recreate coordinate object based on format type
+ if format_type == "pdf":
+ coord = PDFCoordinate(
+ page=coord_data["page"],
+ bbox=coord_data.get("bbox", []),
+ line=coord_data.get("line"),
+ )
+ elif format_type == "spreadsheet":
+ coord = SpreadsheetCoordinate(
+ row=coord_data["row"] - 1, # Convert back to 0-indexed
+ col=coord_data.get("col", 1) - 1 if coord_data.get("col") else None,
+ cell_ref=coord_data.get("cell_ref"),
+ sheet=coord_data.get("sheet"),
+ )
+ elif format_type == "document":
+ coord = DocumentCoordinate(
+ paragraph_id=coord_data["paragraph_id"],
+ local_offset=coord_data.get("local_offset", 0),
+ section=coord_data.get("section"),
+ heading_level=coord_data.get("heading_level"),
+ )
+ elif format_type == "code":
+ coord = CodeCoordinate(
+ line=coord_data["line"],
+ column=coord_data.get("column", 1),
+ function_name=coord_data.get("function_name"),
+ class_name=coord_data.get("class_name"),
+ )
+ elif format_type == "bounding_box":
+ coord = BoundingBoxCoordinate(
+ page=coord_data["page"],
+ bbox=coord_data["bbox"],
+ bbox_pixels=coord_data.get("bbox_pixels"),
+ label=coord_data.get("label"),
+ confidence=coord_data.get("confidence"),
+ )
+ else:
+ coord = CharacterCoordinate(offset=coord_data.get("offset", m["start"]))
+
+ mapper.add_mapping(m["start"], m["end"], coord)
+
+ return mapper
+
+
+def get_column_letter(col_idx: int) -> str:
+ """
+ Convert a 0-indexed column number to Excel column letter.
+
+ Args:
+ col_idx: 0-indexed column number
+
+ Returns:
+ Column letter (A, B, ..., Z, AA, AB, ...)
+ """
+ result = ""
+ col_idx += 1 # Convert to 1-indexed
+ while col_idx > 0:
+ col_idx, remainder = divmod(col_idx - 1, 26)
+ result = chr(65 + remainder) + result
+ return result
+
+
+def get_cell_reference(row: int, col: int) -> str:
+ """
+ Get A1-style cell reference.
+
+ Args:
+ row: 0-indexed row number
+ col: 0-indexed column number
+
+ Returns:
+ Cell reference like "A1", "B5", "AA100"
+ """
+ return f"{get_column_letter(col)}{row + 1}"
diff --git a/potato/format_handlers/docx_handler.py b/potato/format_handlers/docx_handler.py
new file mode 100644
index 0000000000000000000000000000000000000000..22e3c107982bb45916163210290acfe2e6ade671
--- /dev/null
+++ b/potato/format_handlers/docx_handler.py
@@ -0,0 +1,345 @@
+"""
+DOCX Format Handler
+
+Extracts text and structure from Word documents using python-docx and mammoth.
+Supports paragraph-level coordinate mapping for span annotations.
+
+Usage:
+ from potato.format_handlers.docx_handler import DocxHandler
+
+ handler = DocxHandler()
+ output = handler.extract("document.docx", {
+ "preserve_styles": True,
+ "include_headers": True,
+ })
+"""
+
+from typing import Dict, List, Any, Optional
+from pathlib import Path
+import html
+import logging
+import uuid
+
+from .base import BaseFormatHandler, FormatOutput
+from .coordinate_mapping import CoordinateMapper, DocumentCoordinate
+
+logger = logging.getLogger(__name__)
+
+# Check if dependencies are available
+try:
+ import docx
+ from docx.document import Document
+ from docx.oxml.ns import qn
+ DOCX_AVAILABLE = True
+except ImportError:
+ DOCX_AVAILABLE = False
+ docx = None
+
+try:
+ import mammoth
+ MAMMOTH_AVAILABLE = True
+except ImportError:
+ MAMMOTH_AVAILABLE = False
+ mammoth = None
+
+
+class DocxHandler(BaseFormatHandler):
+ """
+ Handler for Word documents (.docx).
+
+ Uses python-docx for text extraction with structure preservation
+ and mammoth for HTML conversion.
+ """
+
+ format_name = "docx"
+ supported_extensions = [".docx"]
+ description = "Word document extraction with paragraph/section mapping"
+ requires_dependencies = ["python-docx", "mammoth"]
+
+ def get_default_options(self) -> Dict[str, Any]:
+ """Get default extraction options."""
+ return {
+ "preserve_styles": True,
+ "include_headers": True,
+ "include_footers": False,
+ "include_tables": True,
+ "paragraph_separator": "\n\n",
+ "use_mammoth_html": True, # Use mammoth for rich HTML conversion
+ }
+
+ def extract(
+ self,
+ file_path: str,
+ options: Optional[Dict[str, Any]] = None
+ ) -> FormatOutput:
+ """
+ Extract text and structure from a Word document.
+
+ Args:
+ file_path: Path to the .docx file
+ options: Extraction options:
+ - preserve_styles: Keep heading levels and formatting
+ - include_headers: Include document headers
+ - include_footers: Include document footers
+ - include_tables: Extract table content
+
+ Returns:
+ FormatOutput with extracted text, HTML, and coordinate mappings
+ """
+ if not DOCX_AVAILABLE:
+ raise ImportError(
+ "python-docx is required for DOCX extraction. "
+ "Install with: pip install python-docx"
+ )
+
+ opts = self.merge_options(options)
+ mapper = CoordinateMapper()
+
+ # Open document
+ doc = docx.Document(file_path)
+
+ # Extract text with structure
+ text_parts = []
+ html_parts = []
+ current_offset = 0
+
+ metadata = {
+ "format": "docx",
+ "source_file": str(file_path),
+ "paragraphs": [],
+ "sections": [],
+ "tables": [],
+ }
+
+ # Check if mammoth is available for rich HTML
+ if opts.get("use_mammoth_html") and MAMMOTH_AVAILABLE:
+ rendered_html = self._extract_with_mammoth(file_path, opts)
+ else:
+ rendered_html = None
+
+ html_parts.append('')
+
+ # Extract headers if requested
+ if opts.get("include_headers"):
+ for section in doc.sections:
+ header = section.header
+ if header and header.paragraphs:
+ header_text = "\n".join(p.text for p in header.paragraphs if p.text.strip())
+ if header_text:
+ para_id = f"header_{uuid.uuid4().hex[:8]}"
+ text_parts.append(header_text)
+ text_parts.append("\n\n")
+
+ html_parts.append(f'')
+
+ mapper.add_mapping(
+ current_offset,
+ current_offset + len(header_text),
+ DocumentCoordinate(
+ paragraph_id=para_id,
+ local_offset=0,
+ section="header",
+ )
+ )
+ current_offset += len(header_text) + 2
+ metadata["sections"].append({"type": "header", "id": para_id})
+
+ # Extract main body
+ current_section = None
+ for i, para in enumerate(doc.paragraphs):
+ para_text = para.text
+ if not para_text.strip():
+ continue
+
+ para_id = f"p_{i}_{uuid.uuid4().hex[:8]}"
+
+ # Detect heading level
+ heading_level = None
+ if para.style and para.style.name:
+ style_name = para.style.name.lower()
+ if style_name.startswith("heading"):
+ try:
+ heading_level = int(style_name.replace("heading", "").strip())
+ except ValueError:
+ pass
+
+ # Update section tracking
+ if heading_level:
+ current_section = para_text.strip()
+
+ # Build text
+ start_offset = current_offset
+ text_parts.append(para_text)
+ end_offset = current_offset + len(para_text)
+
+ # Add paragraph separator
+ text_parts.append(opts["paragraph_separator"])
+ current_offset = end_offset + len(opts["paragraph_separator"])
+
+ # Build HTML
+ css_class = "docx-paragraph"
+ if heading_level:
+ css_class = f"docx-heading docx-h{heading_level}"
+ html_tag = f"h{min(heading_level, 6)}"
+ else:
+ html_tag = "p"
+
+ html_parts.append(
+ f'<{html_tag} class="{css_class}" '
+ f'data-para-id="{para_id}" '
+ f'data-start="{start_offset}" '
+ f'data-end="{end_offset}">'
+ f'{html.escape(para_text)}'
+ f'{html_tag}>'
+ )
+
+ # Add coordinate mapping
+ mapper.add_mapping(
+ start_offset,
+ end_offset,
+ DocumentCoordinate(
+ paragraph_id=para_id,
+ local_offset=0,
+ section=current_section,
+ heading_level=heading_level,
+ )
+ )
+
+ # Track paragraph metadata
+ metadata["paragraphs"].append({
+ "id": para_id,
+ "start": start_offset,
+ "end": end_offset,
+ "heading_level": heading_level,
+ "section": current_section,
+ "char_count": len(para_text),
+ })
+
+ # Extract tables if requested
+ if opts.get("include_tables"):
+ for t_idx, table in enumerate(doc.tables):
+ table_id = f"table_{t_idx}_{uuid.uuid4().hex[:8]}"
+ table_text, table_html = self._extract_table(
+ table, table_id, current_offset
+ )
+
+ if table_text:
+ text_parts.append("\n")
+ text_parts.append(table_text)
+ text_parts.append("\n")
+ html_parts.append(table_html)
+
+ mapper.add_mapping(
+ current_offset,
+ current_offset + len(table_text),
+ DocumentCoordinate(
+ paragraph_id=table_id,
+ section=current_section,
+ )
+ )
+ current_offset += len(table_text) + 2
+ metadata["tables"].append({"id": table_id})
+
+ html_parts.append('
')
+
+ full_text = "".join(text_parts)
+
+ # Use mammoth HTML if available, otherwise use our generated HTML
+ if rendered_html:
+ final_html = rendered_html
+ else:
+ final_html = "\n".join(html_parts)
+
+ coord_dict = mapper.to_dict()
+ coord_dict["get_coords_for_range"] = mapper.get_coords_for_range
+
+ return FormatOutput(
+ text=full_text,
+ rendered_html=final_html,
+ coordinate_map=coord_dict,
+ metadata=metadata,
+ format_name=self.format_name,
+ source_path=str(file_path),
+ )
+
+ def _extract_with_mammoth(self, file_path: str, opts: Dict[str, Any]) -> str:
+ """
+ Use mammoth for rich HTML conversion.
+ """
+ if not MAMMOTH_AVAILABLE:
+ return None
+
+ try:
+ with open(file_path, "rb") as f:
+ result = mammoth.convert_to_html(f)
+ html_content = result.value
+
+ # Wrap in container
+ return f'{html_content}
'
+ except Exception as e:
+ logger.warning(f"Mammoth conversion failed: {e}")
+ return None
+
+ def _extract_table(
+ self,
+ table,
+ table_id: str,
+ base_offset: int
+ ) -> tuple:
+ """
+ Extract text and HTML from a table.
+
+ Returns:
+ Tuple of (text, html)
+ """
+ text_rows = []
+ html_parts = []
+
+ html_parts.append(f'')
+
+ for row_idx, row in enumerate(table.rows):
+ row_texts = []
+ html_parts.append('')
+
+ for cell_idx, cell in enumerate(row.cells):
+ cell_text = cell.text.strip()
+ row_texts.append(cell_text)
+ html_parts.append(f'{html.escape(cell_text)} ')
+
+ html_parts.append(' ')
+ text_rows.append("\t".join(row_texts))
+
+ html_parts.append('
')
+
+ return "\n".join(text_rows), "\n".join(html_parts)
+
+ def extract_metadata(self, file_path: str) -> Dict[str, Any]:
+ """
+ Extract document metadata (author, title, etc.).
+
+ Args:
+ file_path: Path to the .docx file
+
+ Returns:
+ Dictionary of metadata properties
+ """
+ if not DOCX_AVAILABLE:
+ raise ImportError("python-docx is required")
+
+ doc = docx.Document(file_path)
+ core_props = doc.core_properties
+
+ return {
+ "author": core_props.author,
+ "title": core_props.title,
+ "subject": core_props.subject,
+ "keywords": core_props.keywords,
+ "created": str(core_props.created) if core_props.created else None,
+ "modified": str(core_props.modified) if core_props.modified else None,
+ "last_modified_by": core_props.last_modified_by,
+ "revision": core_props.revision,
+ "category": core_props.category,
+ "comments": core_props.comments,
+ }
diff --git a/potato/format_handlers/markdown_handler.py b/potato/format_handlers/markdown_handler.py
new file mode 100644
index 0000000000000000000000000000000000000000..2e2b6ac2b1e2e3560d9f7155076724e6bc59c805
--- /dev/null
+++ b/potato/format_handlers/markdown_handler.py
@@ -0,0 +1,245 @@
+"""
+Markdown Format Handler
+
+Parses Markdown files and extracts text with source line mapping.
+Supports syntax highlighting for code blocks.
+
+Usage:
+ from potato.format_handlers.markdown_handler import MarkdownHandler
+
+ handler = MarkdownHandler()
+ output = handler.extract("document.md", {
+ "highlight_code": True,
+ "gfm": True, # GitHub Flavored Markdown
+ })
+"""
+
+from typing import Dict, List, Any, Optional
+from pathlib import Path
+import html
+import logging
+import re
+import uuid
+
+from .base import BaseFormatHandler, FormatOutput
+from .coordinate_mapping import CoordinateMapper, CodeCoordinate, DocumentCoordinate
+
+logger = logging.getLogger(__name__)
+
+# Check if dependencies are available
+try:
+ import mistune
+ MISTUNE_AVAILABLE = True
+except ImportError:
+ MISTUNE_AVAILABLE = False
+ mistune = None
+
+try:
+ from pygments import highlight
+ from pygments.lexers import get_lexer_by_name, guess_lexer
+ from pygments.formatters import HtmlFormatter
+ PYGMENTS_AVAILABLE = True
+except ImportError:
+ PYGMENTS_AVAILABLE = False
+
+
+class MarkdownHandler(BaseFormatHandler):
+ """
+ Handler for Markdown files.
+
+ Uses mistune for parsing and Pygments for syntax highlighting.
+ Maintains line/column coordinate mappings.
+ """
+
+ format_name = "markdown"
+ supported_extensions = [".md", ".markdown", ".mdown", ".mkd"]
+ description = "Markdown parsing with line/column mapping and syntax highlighting"
+ requires_dependencies = ["mistune"]
+
+ def get_default_options(self) -> Dict[str, Any]:
+ """Get default extraction options."""
+ return {
+ "highlight_code": True,
+ "gfm": True, # GitHub Flavored Markdown
+ "include_raw_blocks": False,
+ "preserve_line_breaks": True,
+ }
+
+ def extract(
+ self,
+ file_path: str,
+ options: Optional[Dict[str, Any]] = None
+ ) -> FormatOutput:
+ """
+ Parse and render a Markdown file.
+
+ Args:
+ file_path: Path to the Markdown file
+ options: Extraction options:
+ - highlight_code: Syntax highlight code blocks
+ - gfm: Use GitHub Flavored Markdown extensions
+ - preserve_line_breaks: Keep original line structure
+
+ Returns:
+ FormatOutput with text, rendered HTML, and coordinate mappings
+ """
+ if not MISTUNE_AVAILABLE:
+ raise ImportError(
+ "mistune is required for Markdown extraction. "
+ "Install with: pip install mistune"
+ )
+
+ opts = self.merge_options(options)
+
+ # Read source file
+ path = Path(file_path)
+ source_text = path.read_text(encoding="utf-8")
+
+ # Build line index for coordinate mapping
+ line_offsets = self._build_line_index(source_text)
+
+ # Parse and render
+ mapper = CoordinateMapper()
+ rendered_html = self._render_markdown(source_text, opts)
+
+ # Build coordinate mappings for each line
+ for line_num, (start, end) in enumerate(line_offsets, start=1):
+ mapper.add_mapping(
+ start,
+ end,
+ CodeCoordinate(line=line_num, column=1)
+ )
+
+ metadata = {
+ "format": "markdown",
+ "source_file": str(file_path),
+ "line_count": len(line_offsets),
+ "char_count": len(source_text),
+ "headings": self._extract_headings(source_text),
+ }
+
+ coord_dict = mapper.to_dict()
+ coord_dict["get_coords_for_range"] = mapper.get_coords_for_range
+
+ return FormatOutput(
+ text=source_text,
+ rendered_html=rendered_html,
+ coordinate_map=coord_dict,
+ metadata=metadata,
+ format_name=self.format_name,
+ source_path=str(file_path),
+ )
+
+ def _build_line_index(self, text: str) -> List[tuple]:
+ """
+ Build an index of line start/end offsets.
+
+ Returns:
+ List of (start, end) tuples for each line
+ """
+ lines = []
+ start = 0
+
+ for line in text.split("\n"):
+ end = start + len(line)
+ lines.append((start, end))
+ start = end + 1 # +1 for newline
+
+ return lines
+
+ def _render_markdown(self, text: str, opts: Dict[str, Any]) -> str:
+ """
+ Render Markdown to HTML using mistune.
+ """
+ # Create custom renderer with code highlighting
+ if opts.get("highlight_code") and PYGMENTS_AVAILABLE:
+ renderer = HighlightRenderer()
+ else:
+ renderer = None
+
+ # Configure mistune
+ if opts.get("gfm"):
+ # Use plugins for GFM features
+ md = mistune.create_markdown(
+ renderer=renderer,
+ plugins=['strikethrough', 'table', 'task_lists']
+ )
+ else:
+ md = mistune.create_markdown(renderer=renderer)
+
+ html_content = md(text)
+
+ # Wrap in container
+ return f'{html_content}
'
+
+ def _extract_headings(self, text: str) -> List[Dict[str, Any]]:
+ """
+ Extract heading structure from Markdown source.
+ """
+ headings = []
+ heading_pattern = re.compile(r'^(#{1,6})\s+(.+)$', re.MULTILINE)
+
+ for match in heading_pattern.finditer(text):
+ level = len(match.group(1))
+ title = match.group(2).strip()
+ headings.append({
+ "level": level,
+ "title": title,
+ "offset": match.start(),
+ "line": text[:match.start()].count("\n") + 1,
+ })
+
+ return headings
+
+ def extract_toc(self, file_path: str) -> List[Dict[str, Any]]:
+ """
+ Extract table of contents from a Markdown file.
+
+ Args:
+ file_path: Path to the Markdown file
+
+ Returns:
+ List of heading entries with level, title, and line number
+ """
+ path = Path(file_path)
+ source_text = path.read_text(encoding="utf-8")
+ return self._extract_headings(source_text)
+
+
+class HighlightRenderer(mistune.HTMLRenderer if MISTUNE_AVAILABLE else object):
+ """
+ Custom Mistune renderer with Pygments syntax highlighting.
+ """
+
+ def __init__(self):
+ if MISTUNE_AVAILABLE:
+ super().__init__()
+ self.formatter = HtmlFormatter(cssclass="highlight") if PYGMENTS_AVAILABLE else None
+
+ def block_code(self, code: str, info: str = None) -> str:
+ """
+ Render a code block with syntax highlighting.
+ """
+ if not PYGMENTS_AVAILABLE or not self.formatter:
+ escaped = html.escape(code)
+ lang_attr = f' class="language-{info}"' if info else ''
+ return f'{escaped} \n'
+
+ try:
+ if info:
+ lexer = get_lexer_by_name(info, stripall=True)
+ else:
+ lexer = guess_lexer(code)
+ except Exception:
+ # Fall back to plain text
+ escaped = html.escape(code)
+ return f'{escaped} \n'
+
+ return highlight(code, lexer, self.formatter)
+
+ def codespan(self, text: str) -> str:
+ """
+ Render inline code.
+ """
+ escaped = html.escape(text)
+ return f'{escaped}'
diff --git a/potato/format_handlers/pdf_handler.py b/potato/format_handlers/pdf_handler.py
new file mode 100644
index 0000000000000000000000000000000000000000..62f4e2e55f30da51bd70bf2d7bb96f51c7bd849a
--- /dev/null
+++ b/potato/format_handlers/pdf_handler.py
@@ -0,0 +1,417 @@
+"""
+PDF Format Handler
+
+Extracts text and layout information from PDF files using pdfplumber.
+Supports text extraction with character-level position mapping.
+
+Usage:
+ from potato.format_handlers.pdf_handler import PDFHandler
+
+ handler = PDFHandler()
+ output = handler.extract("document.pdf", {
+ "extraction_mode": "text", # or "layout"
+ "max_pages": 10,
+ })
+
+ # Access extracted content
+ text = output.text
+ html = output.rendered_html
+ coords = output.coordinate_map
+"""
+
+from typing import Dict, List, Any, Optional
+from pathlib import Path
+import html
+import logging
+
+from .base import BaseFormatHandler, FormatOutput
+from .coordinate_mapping import CoordinateMapper, PDFCoordinate
+
+logger = logging.getLogger(__name__)
+
+# Check if pdfplumber is available
+try:
+ import pdfplumber
+ PDFPLUMBER_AVAILABLE = True
+except ImportError:
+ PDFPLUMBER_AVAILABLE = False
+ pdfplumber = None
+
+
+class PDFHandler(BaseFormatHandler):
+ """
+ Handler for PDF documents.
+
+ Uses pdfplumber for text extraction with position information.
+ Generates HTML representation suitable for span annotation.
+ """
+
+ format_name = "pdf"
+ supported_extensions = [".pdf"]
+ description = "PDF document text extraction with page/position mapping"
+ requires_dependencies = ["pdfplumber"]
+
+ def get_default_options(self) -> Dict[str, Any]:
+ """Get default extraction options."""
+ return {
+ "extraction_mode": "text", # "text" or "layout"
+ "preserve_layout": False,
+ "max_pages": None,
+ "include_page_breaks": True,
+ "page_separator": "\n\n--- Page {page} ---\n\n",
+ "extract_tables": False,
+ "x_tolerance": 3, # Horizontal tolerance for word grouping
+ "y_tolerance": 3, # Vertical tolerance for line grouping
+ }
+
+ def extract(
+ self,
+ file_path: str,
+ options: Optional[Dict[str, Any]] = None
+ ) -> FormatOutput:
+ """
+ Extract text and layout from a PDF file.
+
+ Args:
+ file_path: Path to the PDF file
+ options: Extraction options:
+ - extraction_mode: "text" (plain) or "layout" (preserve layout)
+ - max_pages: Maximum pages to process (None for all)
+ - include_page_breaks: Include page separators in text
+ - page_separator: Format string for page breaks ({page} replaced)
+ - extract_tables: Also extract table structures
+
+ Returns:
+ FormatOutput with extracted text, HTML, and coordinate mappings
+ """
+ if not PDFPLUMBER_AVAILABLE:
+ raise ImportError(
+ "pdfplumber is required for PDF extraction. "
+ "Install with: pip install pdfplumber"
+ )
+
+ opts = self.merge_options(options)
+ mapper = CoordinateMapper()
+
+ text_parts = []
+ html_parts = []
+ current_offset = 0
+
+ metadata = {
+ "format": "pdf",
+ "pages": [],
+ "total_pages": 0,
+ "source_file": str(file_path),
+ }
+
+ html_parts.append('')
+
+ with pdfplumber.open(file_path) as pdf:
+ metadata["total_pages"] = len(pdf.pages)
+ max_pages = opts.get("max_pages") or len(pdf.pages)
+
+ for page_num, page in enumerate(pdf.pages[:max_pages], start=1):
+ page_text, page_html, page_coords = self._extract_page(
+ page, page_num, opts, current_offset
+ )
+
+ # Add page coordinates to mapper
+ for coord_info in page_coords:
+ mapper.add_mapping(
+ coord_info["start"],
+ coord_info["end"],
+ PDFCoordinate(
+ page=page_num,
+ bbox=coord_info.get("bbox", []),
+ line=coord_info.get("line"),
+ )
+ )
+
+ # Add page separator
+ if page_num > 1 and opts.get("include_page_breaks"):
+ separator = opts["page_separator"].format(page=page_num)
+ text_parts.append(separator)
+ current_offset += len(separator)
+
+ text_parts.append(page_text)
+ html_parts.append(page_html)
+ current_offset += len(page_text)
+
+ # Page metadata
+ page_meta = {
+ "page_number": page_num,
+ "width": float(page.width),
+ "height": float(page.height),
+ "char_count": len(page_text),
+ }
+ metadata["pages"].append(page_meta)
+
+ html_parts.append('
')
+
+ full_text = "".join(text_parts)
+ full_html = "\n".join(html_parts)
+
+ # Create output with coordinate lookup function
+ coord_dict = mapper.to_dict()
+ coord_dict["get_coords_for_range"] = mapper.get_coords_for_range
+
+ return FormatOutput(
+ text=full_text,
+ rendered_html=full_html,
+ coordinate_map=coord_dict,
+ metadata=metadata,
+ format_name=self.format_name,
+ source_path=str(file_path),
+ )
+
+ def _extract_page(
+ self,
+ page,
+ page_num: int,
+ opts: Dict[str, Any],
+ base_offset: int
+ ) -> tuple:
+ """
+ Extract text and HTML from a single page.
+
+ Returns:
+ Tuple of (text, html, coordinate_mappings)
+ """
+ extraction_mode = opts.get("extraction_mode", "text")
+
+ if extraction_mode == "layout":
+ return self._extract_page_layout(page, page_num, opts, base_offset)
+ else:
+ return self._extract_page_text(page, page_num, opts, base_offset)
+
+ def _extract_page_text(
+ self,
+ page,
+ page_num: int,
+ opts: Dict[str, Any],
+ base_offset: int
+ ) -> tuple:
+ """
+ Extract text with word-level coordinate mapping.
+ """
+ text_parts = []
+ html_parts = []
+ coords = []
+ current_offset = base_offset
+
+ # Extract words with their positions
+ words = page.extract_words(
+ x_tolerance=opts.get("x_tolerance", 3),
+ y_tolerance=opts.get("y_tolerance", 3),
+ )
+
+ html_parts.append(f'')
+
+ if not words:
+ # Fall back to full text extraction if no words found
+ text = page.extract_text() or ""
+ text_parts.append(text)
+ html_parts.append(f'{html.escape(text)} ')
+
+ if text:
+ coords.append({
+ "start": current_offset,
+ "end": current_offset + len(text),
+ "bbox": [0, 0, float(page.width), float(page.height)],
+ })
+ else:
+ # Process words with positions
+ current_line_top = None
+ line_words = []
+
+ for word in words:
+ word_top = word["top"]
+
+ # Check if this is a new line
+ if current_line_top is None:
+ current_line_top = word_top
+ elif abs(word_top - current_line_top) > opts.get("y_tolerance", 3):
+ # Flush current line
+ if line_words:
+ line_text, line_html, line_coords = self._process_line(
+ line_words, current_offset
+ )
+ text_parts.append(line_text)
+ text_parts.append("\n")
+ html_parts.append(line_html)
+ html_parts.append(" ")
+ coords.extend(line_coords)
+ current_offset += len(line_text) + 1 # +1 for newline
+
+ line_words = []
+ current_line_top = word_top
+
+ line_words.append(word)
+
+ # Process final line
+ if line_words:
+ line_text, line_html, line_coords = self._process_line(
+ line_words, current_offset
+ )
+ text_parts.append(line_text)
+ html_parts.append(line_html)
+ coords.extend(line_coords)
+
+ html_parts.append('
')
+
+ return "".join(text_parts), "\n".join(html_parts), coords
+
+ def _process_line(
+ self,
+ words: List[Dict],
+ base_offset: int
+ ) -> tuple:
+ """
+ Process a line of words into text, HTML, and coordinates.
+ """
+ text_parts = []
+ html_parts = []
+ coords = []
+ current_offset = base_offset
+
+ for i, word in enumerate(words):
+ word_text = word["text"]
+
+ # Add space between words
+ if i > 0:
+ text_parts.append(" ")
+ current_offset += 1
+
+ start = current_offset
+ end = start + len(word_text)
+
+ text_parts.append(word_text)
+ html_parts.append(
+ f''
+ f'{html.escape(word_text)} '
+ )
+
+ # Store coordinate mapping
+ coords.append({
+ "start": start,
+ "end": end,
+ "bbox": [
+ float(word["x0"]),
+ float(word["top"]),
+ float(word["x1"]),
+ float(word["bottom"]),
+ ],
+ })
+
+ current_offset = end
+
+ return "".join(text_parts), " ".join(html_parts), coords
+
+ def _extract_page_layout(
+ self,
+ page,
+ page_num: int,
+ opts: Dict[str, Any],
+ base_offset: int
+ ) -> tuple:
+ """
+ Extract text preserving visual layout.
+ """
+ # Use extract_text with layout preservation
+ text = page.extract_text(layout=True) or ""
+
+ html_parts = []
+ html_parts.append(f'')
+ html_parts.append(f'
{html.escape(text)} ')
+ html_parts.append('
')
+
+ # For layout mode, we map the entire page
+ coords = [{
+ "start": base_offset,
+ "end": base_offset + len(text),
+ "bbox": [0, 0, float(page.width), float(page.height)],
+ }]
+
+ return text, "\n".join(html_parts), coords
+
+ def get_page_count(self, file_path: str) -> int:
+ """
+ Get the number of pages in a PDF.
+
+ Args:
+ file_path: Path to the PDF file
+
+ Returns:
+ Number of pages
+ """
+ if not PDFPLUMBER_AVAILABLE:
+ raise ImportError("pdfplumber is required")
+
+ with pdfplumber.open(file_path) as pdf:
+ return len(pdf.pages)
+
+ def extract_page(
+ self,
+ file_path: str,
+ page_number: int,
+ options: Optional[Dict[str, Any]] = None
+ ) -> FormatOutput:
+ """
+ Extract a single page from a PDF.
+
+ Args:
+ file_path: Path to the PDF file
+ page_number: Page number (1-indexed)
+ options: Extraction options
+
+ Returns:
+ FormatOutput for the single page
+ """
+ if not PDFPLUMBER_AVAILABLE:
+ raise ImportError("pdfplumber is required")
+
+ opts = self.merge_options(options)
+ opts["max_pages"] = page_number # Process up to this page
+ opts["include_page_breaks"] = False
+
+ # Extract only the requested page
+ mapper = CoordinateMapper()
+
+ with pdfplumber.open(file_path) as pdf:
+ if page_number < 1 or page_number > len(pdf.pages):
+ raise ValueError(
+ f"Page {page_number} out of range (1-{len(pdf.pages)})"
+ )
+
+ page = pdf.pages[page_number - 1]
+ page_text, page_html, page_coords = self._extract_page(
+ page, page_number, opts, 0
+ )
+
+ for coord_info in page_coords:
+ mapper.add_mapping(
+ coord_info["start"],
+ coord_info["end"],
+ PDFCoordinate(
+ page=page_number,
+ bbox=coord_info.get("bbox", []),
+ )
+ )
+
+ coord_dict = mapper.to_dict()
+ coord_dict["get_coords_for_range"] = mapper.get_coords_for_range
+
+ return FormatOutput(
+ text=page_text,
+ rendered_html=page_html,
+ coordinate_map=coord_dict,
+ metadata={
+ "format": "pdf",
+ "page_number": page_number,
+ "total_pages": len(pdf.pages),
+ },
+ format_name=self.format_name,
+ source_path=str(file_path),
+ )
diff --git a/potato/format_handlers/registry.py b/potato/format_handlers/registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d498773913d98e9fd3326e44f13d963ed422075
--- /dev/null
+++ b/potato/format_handlers/registry.py
@@ -0,0 +1,317 @@
+"""
+Format Handler Registry
+
+Provides a centralized registry for managing format handlers.
+Supports auto-detection of formats from file extensions.
+
+Usage:
+ from potato.format_handlers.registry import format_handler_registry
+
+ # Extract content from a file (auto-detect format)
+ output = format_handler_registry.extract("document.pdf")
+
+ # Extract with specific options
+ output = format_handler_registry.extract(
+ "document.pdf",
+ options={"extraction_mode": "text", "max_pages": 10}
+ )
+
+ # List supported formats
+ formats = format_handler_registry.get_supported_formats()
+
+ # Check if a file is supported
+ if format_handler_registry.can_handle("document.pdf"):
+ output = format_handler_registry.extract("document.pdf")
+"""
+
+from typing import Dict, List, Any, Optional, Type
+from pathlib import Path
+import logging
+
+from .base import BaseFormatHandler, FormatOutput
+
+logger = logging.getLogger(__name__)
+
+
+class FormatHandlerRegistry:
+ """
+ Centralized registry for format handlers.
+
+ Provides methods to register, retrieve, and use format handlers.
+ Supports both built-in handlers and custom plugins.
+ """
+
+ def __init__(self):
+ self._handlers: Dict[str, BaseFormatHandler] = {}
+ self._extension_map: Dict[str, str] = {} # Extension -> format_name
+ logger.debug("FormatHandlerRegistry initialized")
+
+ def register(self, handler: BaseFormatHandler) -> None:
+ """
+ Register a format handler.
+
+ Args:
+ handler: BaseFormatHandler instance to register
+
+ Raises:
+ ValueError: If a handler for this format is already registered
+ """
+ name = handler.format_name
+
+ if name in self._handlers:
+ raise ValueError(f"Format handler '{name}' is already registered")
+
+ self._handlers[name] = handler
+
+ # Map extensions to this handler
+ for ext in handler.supported_extensions:
+ ext_lower = ext.lower()
+ if ext_lower in self._extension_map:
+ existing = self._extension_map[ext_lower]
+ logger.warning(
+ f"Extension '{ext}' already mapped to '{existing}', "
+ f"overriding with '{name}'"
+ )
+ self._extension_map[ext_lower] = name
+
+ logger.debug(
+ f"Registered format handler: {name} "
+ f"(extensions: {handler.supported_extensions})"
+ )
+
+ def unregister(self, format_name: str) -> bool:
+ """
+ Unregister a format handler.
+
+ Args:
+ format_name: Name of the format to unregister
+
+ Returns:
+ True if handler was unregistered, False if not found
+ """
+ if format_name not in self._handlers:
+ return False
+
+ handler = self._handlers[format_name]
+
+ # Remove extension mappings
+ for ext in handler.supported_extensions:
+ ext_lower = ext.lower()
+ if self._extension_map.get(ext_lower) == format_name:
+ del self._extension_map[ext_lower]
+
+ del self._handlers[format_name]
+ logger.debug(f"Unregistered format handler: {format_name}")
+ return True
+
+ def get_handler(self, format_name: str) -> Optional[BaseFormatHandler]:
+ """
+ Get a handler by format name.
+
+ Args:
+ format_name: The format name (e.g., "pdf", "docx")
+
+ Returns:
+ BaseFormatHandler if found, None otherwise
+ """
+ return self._handlers.get(format_name)
+
+ def get_handler_for_file(self, file_path: str) -> Optional[BaseFormatHandler]:
+ """
+ Get the appropriate handler for a file based on its extension.
+
+ Args:
+ file_path: Path to the file
+
+ Returns:
+ BaseFormatHandler if a matching handler exists, None otherwise
+ """
+ ext = Path(file_path).suffix.lower()
+ format_name = self._extension_map.get(ext)
+ if format_name:
+ return self._handlers.get(format_name)
+ return None
+
+ def detect_format(self, file_path: str) -> Optional[str]:
+ """
+ Detect the format of a file based on its extension.
+
+ Args:
+ file_path: Path to the file
+
+ Returns:
+ Format name if detected, None otherwise
+ """
+ ext = Path(file_path).suffix.lower()
+ return self._extension_map.get(ext)
+
+ def can_handle(self, file_path: str) -> bool:
+ """
+ Check if any registered handler can process the file.
+
+ Args:
+ file_path: Path to the file
+
+ Returns:
+ True if a handler is available
+ """
+ return self.get_handler_for_file(file_path) is not None
+
+ def extract(
+ self,
+ file_path: str,
+ format_name: Optional[str] = None,
+ options: Optional[Dict[str, Any]] = None
+ ) -> FormatOutput:
+ """
+ Extract content from a file.
+
+ Args:
+ file_path: Path to the file
+ format_name: Optional format override (auto-detect if not specified)
+ options: Optional extraction options
+
+ Returns:
+ FormatOutput with extracted content
+
+ Raises:
+ ValueError: If no handler is available for the file
+ FileNotFoundError: If the file doesn't exist
+ """
+ # Determine handler to use
+ if format_name:
+ handler = self.get_handler(format_name)
+ if not handler:
+ raise ValueError(
+ f"No handler registered for format '{format_name}'. "
+ f"Available formats: {', '.join(self.get_supported_formats())}"
+ )
+ else:
+ handler = self.get_handler_for_file(file_path)
+ if not handler:
+ ext = Path(file_path).suffix
+ raise ValueError(
+ f"No handler available for extension '{ext}'. "
+ f"Supported extensions: {', '.join(self.get_supported_extensions())}"
+ )
+
+ # Validate file
+ errors = handler.validate_file(file_path)
+ if errors:
+ raise ValueError(f"File validation failed: {'; '.join(errors)}")
+
+ # Extract content
+ logger.info(f"Extracting content from '{file_path}' using {handler.format_name} handler")
+ return handler.extract(file_path, options)
+
+ def get_supported_formats(self) -> List[str]:
+ """
+ Get list of all supported format names.
+
+ Returns:
+ Sorted list of format names
+ """
+ return sorted(self._handlers.keys())
+
+ def get_supported_extensions(self) -> List[str]:
+ """
+ Get list of all supported file extensions.
+
+ Returns:
+ Sorted list of extensions
+ """
+ return sorted(self._extension_map.keys())
+
+ def list_handlers(self) -> List[Dict[str, Any]]:
+ """
+ List all registered handlers with their metadata.
+
+ Returns:
+ List of handler information dictionaries
+ """
+ result = []
+ for name, handler in sorted(self._handlers.items()):
+ missing_deps = handler.check_dependencies()
+ result.append({
+ "name": name,
+ "description": handler.description,
+ "extensions": handler.supported_extensions,
+ "requires": handler.requires_dependencies,
+ "available": len(missing_deps) == 0,
+ "missing_dependencies": missing_deps,
+ })
+ return result
+
+ def is_registered(self, format_name: str) -> bool:
+ """
+ Check if a format is registered.
+
+ Args:
+ format_name: The format name
+
+ Returns:
+ True if registered
+ """
+ return format_name in self._handlers
+
+
+# Global registry instance
+format_handler_registry = FormatHandlerRegistry()
+
+
+def _register_builtin_handlers() -> None:
+ """
+ Register all built-in format handlers.
+ Called automatically when this module is imported.
+ """
+ # Import handlers here to avoid circular imports
+ # and to make dependencies optional
+ handlers_to_register = []
+
+ # PDF Handler
+ try:
+ from .pdf_handler import PDFHandler
+ handlers_to_register.append(PDFHandler())
+ except ImportError as e:
+ logger.debug(f"PDF handler not available: {e}")
+
+ # DOCX Handler
+ try:
+ from .docx_handler import DocxHandler
+ handlers_to_register.append(DocxHandler())
+ except ImportError as e:
+ logger.debug(f"DOCX handler not available: {e}")
+
+ # Markdown Handler
+ try:
+ from .markdown_handler import MarkdownHandler
+ handlers_to_register.append(MarkdownHandler())
+ except ImportError as e:
+ logger.debug(f"Markdown handler not available: {e}")
+
+ # Spreadsheet Handler
+ try:
+ from .spreadsheet_handler import SpreadsheetHandler
+ handlers_to_register.append(SpreadsheetHandler())
+ except ImportError as e:
+ logger.debug(f"Spreadsheet handler not available: {e}")
+
+ # Code Handler
+ try:
+ from .code_handler import CodeHandler
+ handlers_to_register.append(CodeHandler())
+ except ImportError as e:
+ logger.debug(f"Code handler not available: {e}")
+
+ # Register all available handlers
+ for handler in handlers_to_register:
+ try:
+ format_handler_registry.register(handler)
+ except Exception as e:
+ logger.warning(f"Failed to register {handler.format_name} handler: {e}")
+
+ logger.debug(f"Registered {len(handlers_to_register)} format handlers")
+
+
+# Auto-register built-in handlers on import
+_register_builtin_handlers()
diff --git a/potato/format_handlers/spreadsheet_handler.py b/potato/format_handlers/spreadsheet_handler.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb03985a074e4e7a5c1bca4738942d426d67816c
--- /dev/null
+++ b/potato/format_handlers/spreadsheet_handler.py
@@ -0,0 +1,396 @@
+"""
+Spreadsheet Format Handler
+
+Extracts data from spreadsheet files (Excel, CSV, TSV) with row/cell
+coordinate mapping for annotation.
+
+Usage:
+ from potato.format_handlers.spreadsheet_handler import SpreadsheetHandler
+
+ handler = SpreadsheetHandler()
+ output = handler.extract("data.xlsx", {
+ "annotation_mode": "row", # or "cell"
+ "max_rows": 1000,
+ })
+"""
+
+from typing import Dict, List, Any, Optional
+from pathlib import Path
+import html
+import logging
+import csv
+
+from .base import BaseFormatHandler, FormatOutput
+from .coordinate_mapping import (
+ CoordinateMapper,
+ SpreadsheetCoordinate,
+ get_cell_reference,
+)
+
+logger = logging.getLogger(__name__)
+
+# Check if dependencies are available
+try:
+ import openpyxl
+ OPENPYXL_AVAILABLE = True
+except ImportError:
+ OPENPYXL_AVAILABLE = False
+ openpyxl = None
+
+try:
+ import pandas as pd
+ PANDAS_AVAILABLE = True
+except ImportError:
+ PANDAS_AVAILABLE = False
+ pd = None
+
+
+class SpreadsheetHandler(BaseFormatHandler):
+ """
+ Handler for spreadsheet files.
+
+ Supports Excel (.xlsx, .xls) via openpyxl and CSV/TSV via pandas or stdlib.
+ Provides row-based and cell-based annotation modes.
+ """
+
+ format_name = "spreadsheet"
+ supported_extensions = [".csv", ".tsv", ".xlsx", ".xls"]
+ description = "Spreadsheet extraction with row/cell coordinate mapping"
+ requires_dependencies = ["openpyxl"]
+
+ def get_default_options(self) -> Dict[str, Any]:
+ """Get default extraction options."""
+ return {
+ "annotation_mode": "row", # "row", "cell", or "range"
+ "max_rows": 1000,
+ "header_row": 0, # Row index for headers (None for no headers)
+ "sheet_name": None, # Sheet to extract (None for first/active)
+ "skip_empty_rows": True,
+ "text_columns": None, # Columns to include (None for all)
+ "row_separator": "\n",
+ "cell_separator": "\t",
+ }
+
+ def extract(
+ self,
+ file_path: str,
+ options: Optional[Dict[str, Any]] = None
+ ) -> FormatOutput:
+ """
+ Extract data from a spreadsheet file.
+
+ Args:
+ file_path: Path to the spreadsheet file
+ options: Extraction options:
+ - annotation_mode: "row" (annotate rows) or "cell" (annotate cells)
+ - max_rows: Maximum rows to process
+ - header_row: Row index for column headers
+ - sheet_name: Sheet to extract (for Excel files)
+
+ Returns:
+ FormatOutput with extracted text, HTML table, and coordinate mappings
+ """
+ opts = self.merge_options(options)
+ path = Path(file_path)
+ ext = path.suffix.lower()
+
+ # Load data based on file type
+ if ext in [".xlsx", ".xls"]:
+ if not OPENPYXL_AVAILABLE:
+ raise ImportError(
+ "openpyxl is required for Excel extraction. "
+ "Install with: pip install openpyxl"
+ )
+ data, headers, sheet_name = self._load_excel(file_path, opts)
+ else:
+ # CSV/TSV
+ data, headers = self._load_csv(file_path, opts)
+ sheet_name = None
+
+ # Process data
+ mapper = CoordinateMapper()
+ text_parts = []
+ html_parts = []
+ current_offset = 0
+
+ metadata = {
+ "format": "spreadsheet",
+ "source_file": str(file_path),
+ "file_type": ext[1:], # Remove dot
+ "row_count": len(data),
+ "column_count": len(headers) if headers else (len(data[0]) if data else 0),
+ "headers": headers,
+ "sheet_name": sheet_name,
+ "annotation_mode": opts["annotation_mode"],
+ }
+
+ # Build HTML table
+ html_parts.append(
+ f''
+ )
+ html_parts.append('
')
+
+ # Header row
+ if headers:
+ html_parts.append('')
+ for col_idx, header in enumerate(headers):
+ html_parts.append(f'{html.escape(str(header))} ')
+ html_parts.append(' ')
+
+ # Data rows
+ html_parts.append('')
+
+ row_separator = opts["row_separator"]
+ cell_separator = opts["cell_separator"]
+
+ for row_idx, row in enumerate(data):
+ row_start = current_offset
+ row_texts = []
+
+ html_parts.append(
+ f''
+ )
+
+ for col_idx, cell_value in enumerate(row):
+ cell_text = str(cell_value) if cell_value is not None else ""
+ cell_start = current_offset
+ cell_end = cell_start + len(cell_text)
+
+ row_texts.append(cell_text)
+
+ # Add cell coordinate mapping
+ cell_ref = get_cell_reference(row_idx, col_idx)
+
+ if opts["annotation_mode"] == "cell":
+ mapper.add_mapping(
+ cell_start,
+ cell_end,
+ SpreadsheetCoordinate(
+ row=row_idx,
+ col=col_idx,
+ cell_ref=cell_ref,
+ sheet=sheet_name,
+ )
+ )
+
+ # Build cell HTML
+ data_attrs = (
+ f'data-row="{row_idx}" '
+ f'data-col="{col_idx}" '
+ f'data-cell-ref="{cell_ref}" '
+ f'data-start="{cell_start}" '
+ f'data-end="{cell_end}"'
+ )
+ html_parts.append(
+ f''
+ f'{html.escape(cell_text)} '
+ )
+
+ current_offset = cell_end
+ if col_idx < len(row) - 1:
+ current_offset += len(cell_separator)
+
+ html_parts.append(' ')
+
+ # Build row text
+ row_text = cell_separator.join(row_texts)
+ text_parts.append(row_text)
+
+ row_end = current_offset
+
+ # Add row coordinate mapping
+ if opts["annotation_mode"] == "row":
+ mapper.add_mapping(
+ row_start,
+ row_end,
+ SpreadsheetCoordinate(
+ row=row_idx,
+ sheet=sheet_name,
+ )
+ )
+
+ # Add row separator
+ current_offset += len(row_separator)
+
+ html_parts.append(' ')
+ html_parts.append('
')
+ html_parts.append('
')
+
+ full_text = row_separator.join(text_parts)
+ full_html = "\n".join(html_parts)
+
+ coord_dict = mapper.to_dict()
+ coord_dict["get_coords_for_range"] = mapper.get_coords_for_range
+
+ return FormatOutput(
+ text=full_text,
+ rendered_html=full_html,
+ coordinate_map=coord_dict,
+ metadata=metadata,
+ format_name=self.format_name,
+ source_path=str(file_path),
+ )
+
+ def _load_excel(
+ self,
+ file_path: str,
+ opts: Dict[str, Any]
+ ) -> tuple:
+ """
+ Load data from an Excel file using openpyxl.
+
+ Returns:
+ Tuple of (data_rows, headers, sheet_name)
+ """
+ wb = openpyxl.load_workbook(file_path, read_only=True, data_only=True)
+
+ # Select sheet
+ sheet_name = opts.get("sheet_name")
+ if sheet_name:
+ if sheet_name not in wb.sheetnames:
+ raise ValueError(
+ f"Sheet '{sheet_name}' not found. "
+ f"Available: {', '.join(wb.sheetnames)}"
+ )
+ ws = wb[sheet_name]
+ else:
+ ws = wb.active
+ sheet_name = ws.title
+
+ # Read data
+ data = []
+ headers = None
+ header_row = opts.get("header_row")
+ max_rows = opts.get("max_rows", 1000)
+ skip_empty = opts.get("skip_empty_rows", True)
+ text_columns = opts.get("text_columns")
+
+ for row_idx, row in enumerate(ws.iter_rows(max_row=max_rows + (1 if header_row is not None else 0))):
+ row_values = [cell.value for cell in row]
+
+ # Filter columns if specified
+ if text_columns:
+ row_values = [row_values[i] for i in text_columns if i < len(row_values)]
+
+ # Check for empty rows
+ if skip_empty and all(v is None or str(v).strip() == "" for v in row_values):
+ continue
+
+ # Handle header row
+ if header_row is not None and row_idx == header_row:
+ headers = [str(v) if v else f"Column_{i}" for i, v in enumerate(row_values)]
+ continue
+
+ data.append(row_values)
+
+ if len(data) >= max_rows:
+ break
+
+ wb.close()
+ return data, headers, sheet_name
+
+ def _load_csv(
+ self,
+ file_path: str,
+ opts: Dict[str, Any]
+ ) -> tuple:
+ """
+ Load data from a CSV/TSV file.
+
+ Returns:
+ Tuple of (data_rows, headers)
+ """
+ path = Path(file_path)
+ ext = path.suffix.lower()
+
+ # Determine delimiter
+ delimiter = "\t" if ext == ".tsv" else ","
+
+ data = []
+ headers = None
+ header_row = opts.get("header_row")
+ max_rows = opts.get("max_rows", 1000)
+ skip_empty = opts.get("skip_empty_rows", True)
+ text_columns = opts.get("text_columns")
+
+ # Try pandas first if available
+ if PANDAS_AVAILABLE:
+ try:
+ df = pd.read_csv(
+ file_path,
+ delimiter=delimiter,
+ header=header_row,
+ nrows=max_rows,
+ skip_blank_lines=skip_empty,
+ usecols=text_columns,
+ )
+ headers = df.columns.tolist()
+ data = df.values.tolist()
+ return data, headers
+ except Exception as e:
+ logger.debug(f"Pandas read failed, falling back to csv: {e}")
+
+ # Fall back to stdlib csv
+ with open(file_path, "r", encoding="utf-8", newline="") as f:
+ reader = csv.reader(f, delimiter=delimiter)
+
+ for row_idx, row in enumerate(reader):
+ # Filter columns if specified
+ if text_columns:
+ row = [row[i] for i in text_columns if i < len(row)]
+
+ # Check for empty rows
+ if skip_empty and all(v.strip() == "" for v in row):
+ continue
+
+ # Handle header row
+ if header_row is not None and row_idx == header_row:
+ headers = [v if v else f"Column_{i}" for i, v in enumerate(row)]
+ continue
+
+ data.append(row)
+
+ if len(data) >= max_rows:
+ break
+
+ return data, headers
+
+ def get_sheet_names(self, file_path: str) -> List[str]:
+ """
+ Get list of sheet names in an Excel file.
+
+ Args:
+ file_path: Path to Excel file
+
+ Returns:
+ List of sheet names
+ """
+ if not OPENPYXL_AVAILABLE:
+ raise ImportError("openpyxl is required")
+
+ wb = openpyxl.load_workbook(file_path, read_only=True)
+ names = wb.sheetnames
+ wb.close()
+ return names
+
+ def extract_sheet(
+ self,
+ file_path: str,
+ sheet_name: str,
+ options: Optional[Dict[str, Any]] = None
+ ) -> FormatOutput:
+ """
+ Extract a specific sheet from an Excel file.
+
+ Args:
+ file_path: Path to Excel file
+ sheet_name: Name of sheet to extract
+ options: Extraction options
+
+ Returns:
+ FormatOutput for the specified sheet
+ """
+ opts = self.merge_options(options) if options else self.get_default_options()
+ opts["sheet_name"] = sheet_name
+ return self.extract(file_path, opts)
diff --git a/potato/hierarchy.py b/potato/hierarchy.py
new file mode 100644
index 0000000000000000000000000000000000000000..98fa7cb94a6f4e222d6a8956e1aa537a7663fc42
--- /dev/null
+++ b/potato/hierarchy.py
@@ -0,0 +1,814 @@
+"""
+Hierarchical Annotation Framework
+
+Provides a general framework for managing parent-child relationships between
+annotations. This is a type-agnostic system that works with any annotation type
+(temporal segments, text spans, image regions, etc.).
+
+The framework consists of:
+- ConstraintType: Rules for how children relate to parents
+- TierDefinition: Definition of an annotation tier
+- HierarchyDefinition: Complete hierarchy configuration
+- HierarchyManager: Validates constraints and manages relationships
+
+This module supports ELAN-style tiered annotation as the primary use case,
+but the design is general enough to extend to other hierarchical annotation
+scenarios (discourse > paragraph > sentence, scene > object > part, etc.).
+"""
+
+import logging
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import List, Dict, Optional, Any, Set, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+class ConstraintType(Enum):
+ """
+ Defines how child annotations relate to their parent annotations.
+
+ These constraint types follow ELAN's linguistic type stereotypes:
+ - TIME_SUBDIVISION: Children partition the parent's time span with no gaps
+ - INCLUDED_IN: Children are within parent bounds but may have gaps
+ - SYMBOLIC_ASSOCIATION: Children linked to parent without own time alignment
+ - SYMBOLIC_SUBDIVISION: Children subdivide parent symbolically (no time)
+ - NONE: No constraints (independent tier)
+ """
+ TIME_SUBDIVISION = "time_subdivision"
+ INCLUDED_IN = "included_in"
+ SYMBOLIC_ASSOCIATION = "symbolic_association"
+ SYMBOLIC_SUBDIVISION = "symbolic_subdivision"
+ NONE = "none"
+
+ @classmethod
+ def from_string(cls, value: Optional[str]) -> "ConstraintType":
+ """Convert a string value to ConstraintType, defaulting to NONE."""
+ if value is None:
+ return cls.NONE
+ try:
+ return cls(value.lower())
+ except ValueError:
+ logger.warning(f"Unknown constraint type '{value}', defaulting to NONE")
+ return cls.NONE
+
+
+@dataclass
+class TierDefinition:
+ """
+ Definition of an annotation tier.
+
+ Attributes:
+ name: Unique identifier for this tier
+ tier_type: "independent" (time-aligned) or "dependent" (references parent)
+ parent_tier: Name of parent tier (required if tier_type is "dependent")
+ constraint_type: How child annotations relate to parent
+ description: Human-readable description
+ labels: List of label definitions (name, color, tooltip, etc.)
+ linguistic_type: ELAN linguistic type name (for EAF export)
+ """
+ name: str
+ tier_type: str = "independent" # "independent" | "dependent"
+ parent_tier: Optional[str] = None
+ constraint_type: ConstraintType = ConstraintType.NONE
+ description: str = ""
+ labels: List[Dict[str, Any]] = field(default_factory=list)
+ linguistic_type: Optional[str] = None
+
+ def __post_init__(self):
+ """Normalize and validate tier definition."""
+ # Normalize tier_type
+ self.tier_type = self.tier_type.lower() if self.tier_type else "independent"
+ if self.tier_type not in ("independent", "dependent"):
+ logger.warning(f"Invalid tier_type '{self.tier_type}', defaulting to 'independent'")
+ self.tier_type = "independent"
+
+ # Convert constraint_type if it's a string
+ if isinstance(self.constraint_type, str):
+ self.constraint_type = ConstraintType.from_string(self.constraint_type)
+
+ # Dependent tiers should have a constraint type
+ if self.tier_type == "dependent" and self.constraint_type == ConstraintType.NONE:
+ self.constraint_type = ConstraintType.INCLUDED_IN
+ logger.debug(f"Tier '{self.name}' is dependent but has no constraint, defaulting to INCLUDED_IN")
+
+ @property
+ def is_independent(self) -> bool:
+ """Check if this is an independent (time-aligned) tier."""
+ return self.tier_type == "independent"
+
+ @property
+ def is_dependent(self) -> bool:
+ """Check if this is a dependent tier (references parent)."""
+ return self.tier_type == "dependent"
+
+ @property
+ def is_time_aligned(self) -> bool:
+ """Check if annotations on this tier have their own time alignment."""
+ if self.is_independent:
+ return True
+ return self.constraint_type in (
+ ConstraintType.TIME_SUBDIVISION,
+ ConstraintType.INCLUDED_IN
+ )
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "TierDefinition":
+ """Create a TierDefinition from a configuration dictionary."""
+ return cls(
+ name=data.get("name", ""),
+ tier_type=data.get("tier_type", "independent"),
+ parent_tier=data.get("parent_tier"),
+ constraint_type=ConstraintType.from_string(data.get("constraint_type")),
+ description=data.get("description", ""),
+ labels=data.get("labels", []),
+ linguistic_type=data.get("linguistic_type"),
+ )
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to a dictionary for serialization."""
+ result = {
+ "name": self.name,
+ "tier_type": self.tier_type,
+ "description": self.description,
+ "labels": self.labels,
+ }
+ if self.parent_tier:
+ result["parent_tier"] = self.parent_tier
+ if self.constraint_type != ConstraintType.NONE:
+ result["constraint_type"] = self.constraint_type.value
+ if self.linguistic_type:
+ result["linguistic_type"] = self.linguistic_type
+ return result
+
+
+@dataclass
+class HierarchyDefinition:
+ """
+ Complete hierarchy configuration defining all tiers and their relationships.
+
+ Attributes:
+ tiers: Ordered list of tier definitions (order = display order)
+ """
+ tiers: List[TierDefinition] = field(default_factory=list)
+
+ def __post_init__(self):
+ """Convert dict tiers to TierDefinition objects."""
+ converted = []
+ for tier in self.tiers:
+ if isinstance(tier, dict):
+ converted.append(TierDefinition.from_dict(tier))
+ elif isinstance(tier, TierDefinition):
+ converted.append(tier)
+ else:
+ raise ValueError(f"Invalid tier type: {type(tier)}")
+ self.tiers = converted
+
+ def get_tier(self, name: str) -> Optional[TierDefinition]:
+ """Get a tier definition by name."""
+ for tier in self.tiers:
+ if tier.name == name:
+ return tier
+ return None
+
+ def get_tier_index(self, name: str) -> int:
+ """Get the index of a tier by name (-1 if not found)."""
+ for i, tier in enumerate(self.tiers):
+ if tier.name == name:
+ return i
+ return -1
+
+ def get_children(self, tier_name: str) -> List[TierDefinition]:
+ """Get all tiers that have the given tier as their parent."""
+ return [t for t in self.tiers if t.parent_tier == tier_name]
+
+ def get_descendants(self, tier_name: str) -> List[TierDefinition]:
+ """Get all tiers descended from the given tier (recursive)."""
+ descendants = []
+ for child in self.get_children(tier_name):
+ descendants.append(child)
+ descendants.extend(self.get_descendants(child.name))
+ return descendants
+
+ def get_ancestors(self, tier_name: str) -> List[TierDefinition]:
+ """Get all ancestor tiers of the given tier (up to root)."""
+ ancestors = []
+ tier = self.get_tier(tier_name)
+ while tier and tier.parent_tier:
+ parent = self.get_tier(tier.parent_tier)
+ if parent:
+ ancestors.append(parent)
+ tier = parent
+ else:
+ break
+ return ancestors
+
+ def get_root_tiers(self) -> List[TierDefinition]:
+ """Get all independent (root) tiers."""
+ return [t for t in self.tiers if t.is_independent]
+
+ def get_tier_names(self) -> List[str]:
+ """Get list of all tier names in display order."""
+ return [t.name for t in self.tiers]
+
+ def validate_structure(self) -> List[str]:
+ """
+ Validate the hierarchy structure.
+
+ Returns:
+ List of error messages (empty if valid)
+ """
+ errors = []
+ tier_names: Set[str] = set()
+
+ # Check for duplicate names and build name set
+ for tier in self.tiers:
+ if not tier.name:
+ errors.append("Tier name cannot be empty")
+ continue
+ if tier.name in tier_names:
+ errors.append(f"Duplicate tier name: '{tier.name}'")
+ tier_names.add(tier.name)
+
+ # Validate parent references and dependent tier requirements
+ for tier in self.tiers:
+ if tier.is_dependent:
+ if not tier.parent_tier:
+ errors.append(
+ f"Tier '{tier.name}' is dependent but has no parent_tier specified"
+ )
+ elif tier.parent_tier not in tier_names:
+ errors.append(
+ f"Tier '{tier.name}' references unknown parent '{tier.parent_tier}'"
+ )
+ elif tier.parent_tier == tier.name:
+ errors.append(
+ f"Tier '{tier.name}' cannot be its own parent"
+ )
+
+ # Check for cycles
+ cycle_errors = self._detect_cycles()
+ errors.extend(cycle_errors)
+
+ return errors
+
+ def _detect_cycles(self) -> List[str]:
+ """Detect cycles in the tier hierarchy using DFS."""
+ errors = []
+ visited: Set[str] = set()
+ rec_stack: Set[str] = set()
+
+ def dfs(tier_name: str, path: List[str]) -> bool:
+ """Returns True if a cycle is detected."""
+ visited.add(tier_name)
+ rec_stack.add(tier_name)
+
+ tier = self.get_tier(tier_name)
+ if tier and tier.parent_tier:
+ if tier.parent_tier in rec_stack:
+ # Found a cycle
+ cycle_path = path + [tier_name, tier.parent_tier]
+ errors.append(
+ f"Cycle detected in tier hierarchy: {' -> '.join(cycle_path)}"
+ )
+ return True
+ elif tier.parent_tier not in visited:
+ if dfs(tier.parent_tier, path + [tier_name]):
+ return True
+
+ rec_stack.remove(tier_name)
+ return False
+
+ # Run DFS from each unvisited tier
+ for tier in self.tiers:
+ if tier.name not in visited:
+ dfs(tier.name, [])
+
+ return errors
+
+ @classmethod
+ def from_config(cls, tiers_config: List[Dict[str, Any]]) -> "HierarchyDefinition":
+ """Create a HierarchyDefinition from a configuration list."""
+ return cls(tiers=tiers_config)
+
+ def to_config(self) -> List[Dict[str, Any]]:
+ """Convert to a configuration list for serialization."""
+ return [tier.to_dict() for tier in self.tiers]
+
+
+@dataclass
+class Annotation:
+ """
+ Represents a single annotation on a tier.
+
+ Attributes:
+ id: Unique identifier for this annotation
+ tier: Name of the tier this annotation belongs to
+ start_time: Start time in milliseconds (for time-aligned tiers)
+ end_time: End time in milliseconds (for time-aligned tiers)
+ label: The label/value of this annotation
+ parent_id: ID of parent annotation (for dependent tiers)
+ value: Optional additional text/value content
+ metadata: Optional additional metadata
+ """
+ id: str
+ tier: str
+ label: str = ""
+ start_time: Optional[float] = None
+ end_time: Optional[float] = None
+ parent_id: Optional[str] = None
+ value: Optional[str] = None
+ metadata: Dict[str, Any] = field(default_factory=dict)
+
+ @property
+ def duration(self) -> Optional[float]:
+ """Get duration in milliseconds (if time-aligned)."""
+ if self.start_time is not None and self.end_time is not None:
+ return self.end_time - self.start_time
+ return None
+
+ def overlaps(self, other: "Annotation") -> bool:
+ """Check if this annotation overlaps with another."""
+ if self.start_time is None or self.end_time is None:
+ return False
+ if other.start_time is None or other.end_time is None:
+ return False
+ return (self.start_time < other.end_time and
+ self.end_time > other.start_time)
+
+ def contains(self, other: "Annotation") -> bool:
+ """Check if this annotation fully contains another."""
+ if self.start_time is None or self.end_time is None:
+ return False
+ if other.start_time is None or other.end_time is None:
+ return False
+ return (self.start_time <= other.start_time and
+ self.end_time >= other.end_time)
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "Annotation":
+ """Create an Annotation from a dictionary."""
+ return cls(
+ id=data.get("id", ""),
+ tier=data.get("tier", ""),
+ label=data.get("label", ""),
+ start_time=data.get("start_time"),
+ end_time=data.get("end_time"),
+ parent_id=data.get("parent_id"),
+ value=data.get("value"),
+ metadata=data.get("metadata", {}),
+ )
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to a dictionary for serialization."""
+ result = {
+ "id": self.id,
+ "tier": self.tier,
+ "label": self.label,
+ }
+ if self.start_time is not None:
+ result["start_time"] = self.start_time
+ if self.end_time is not None:
+ result["end_time"] = self.end_time
+ if self.parent_id:
+ result["parent_id"] = self.parent_id
+ if self.value:
+ result["value"] = self.value
+ if self.metadata:
+ result["metadata"] = self.metadata
+ return result
+
+
+@dataclass
+class ValidationResult:
+ """Result of a constraint validation check."""
+ valid: bool
+ error: str = ""
+ warnings: List[str] = field(default_factory=list)
+
+
+class HierarchyManager:
+ """
+ Validates and manages hierarchical annotations.
+
+ This class is responsible for:
+ - Validating that annotations satisfy tier constraints
+ - Finding parent annotations for dependent tiers
+ - Managing the annotation collection for a single instance
+ """
+
+ def __init__(self, definition: HierarchyDefinition):
+ """
+ Initialize with a hierarchy definition.
+
+ Args:
+ definition: HierarchyDefinition describing the tier structure
+ """
+ self.definition = definition
+ self._annotations: Dict[str, List[Annotation]] = {}
+
+ # Validate the hierarchy structure
+ errors = definition.validate_structure()
+ if errors:
+ error_msg = "; ".join(errors)
+ raise ValueError(f"Invalid hierarchy definition: {error_msg}")
+
+ # Initialize annotation lists for each tier
+ for tier in definition.tiers:
+ self._annotations[tier.name] = []
+
+ @property
+ def annotations(self) -> Dict[str, List[Annotation]]:
+ """Get all annotations organized by tier."""
+ return self._annotations
+
+ def get_tier_annotations(self, tier_name: str) -> List[Annotation]:
+ """Get all annotations for a specific tier."""
+ return self._annotations.get(tier_name, [])
+
+ def get_annotation(self, annotation_id: str) -> Optional[Annotation]:
+ """Find an annotation by ID across all tiers."""
+ for annotations in self._annotations.values():
+ for ann in annotations:
+ if ann.id == annotation_id:
+ return ann
+ return None
+
+ def find_parent_annotation(
+ self,
+ tier_name: str,
+ start_time: float,
+ end_time: float
+ ) -> Optional[Annotation]:
+ """
+ Find a parent annotation that contains the given time range.
+
+ Args:
+ tier_name: Name of the child tier
+ start_time: Start time of proposed annotation
+ end_time: End time of proposed annotation
+
+ Returns:
+ Parent annotation if found, None otherwise
+ """
+ tier = self.definition.get_tier(tier_name)
+ if not tier or not tier.parent_tier:
+ return None
+
+ parent_annotations = self._annotations.get(tier.parent_tier, [])
+ for parent in parent_annotations:
+ if parent.start_time is None or parent.end_time is None:
+ continue
+ if parent.start_time <= start_time and parent.end_time >= end_time:
+ return parent
+
+ return None
+
+ def find_overlapping_parent(
+ self,
+ tier_name: str,
+ start_time: float,
+ end_time: float
+ ) -> Optional[Annotation]:
+ """
+ Find a parent annotation that overlaps the given time range.
+
+ Useful for INCLUDED_IN constraint where child doesn't need to be
+ fully contained but just overlap.
+
+ Args:
+ tier_name: Name of the child tier
+ start_time: Start time of proposed annotation
+ end_time: End time of proposed annotation
+
+ Returns:
+ Overlapping parent annotation if found, None otherwise
+ """
+ tier = self.definition.get_tier(tier_name)
+ if not tier or not tier.parent_tier:
+ return None
+
+ parent_annotations = self._annotations.get(tier.parent_tier, [])
+ for parent in parent_annotations:
+ if parent.start_time is None or parent.end_time is None:
+ continue
+ if parent.start_time < end_time and parent.end_time > start_time:
+ return parent
+
+ return None
+
+ def validate_annotation(
+ self,
+ tier_name: str,
+ start_time: Optional[float],
+ end_time: Optional[float],
+ parent_annotation: Optional[Annotation] = None
+ ) -> ValidationResult:
+ """
+ Validate that a proposed annotation satisfies tier constraints.
+
+ Args:
+ tier_name: Name of the tier for this annotation
+ start_time: Start time (None for symbolic annotations)
+ end_time: End time (None for symbolic annotations)
+ parent_annotation: Parent annotation (for dependent tiers)
+
+ Returns:
+ ValidationResult with valid flag and any error message
+ """
+ tier = self.definition.get_tier(tier_name)
+ if not tier:
+ return ValidationResult(
+ valid=False,
+ error=f"Unknown tier: '{tier_name}'"
+ )
+
+ # Independent tiers have no parent constraints
+ if tier.is_independent:
+ return self._validate_time_range(start_time, end_time)
+
+ # Dependent tiers require a parent
+ if not parent_annotation:
+ return ValidationResult(
+ valid=False,
+ error=f"Dependent tier '{tier_name}' requires a parent annotation"
+ )
+
+ # Validate based on constraint type
+ constraint = tier.constraint_type
+
+ if constraint == ConstraintType.TIME_SUBDIVISION:
+ return self._validate_time_subdivision(
+ start_time, end_time, parent_annotation
+ )
+
+ elif constraint == ConstraintType.INCLUDED_IN:
+ return self._validate_included_in(
+ start_time, end_time, parent_annotation
+ )
+
+ elif constraint == ConstraintType.SYMBOLIC_ASSOCIATION:
+ # No time constraints - just need valid parent
+ return ValidationResult(valid=True)
+
+ elif constraint == ConstraintType.SYMBOLIC_SUBDIVISION:
+ # No time constraints - just need valid parent
+ return ValidationResult(valid=True)
+
+ return ValidationResult(valid=True)
+
+ def _validate_time_range(
+ self,
+ start_time: Optional[float],
+ end_time: Optional[float]
+ ) -> ValidationResult:
+ """Validate that start/end times are valid."""
+ if start_time is None or end_time is None:
+ return ValidationResult(valid=True) # Symbolic annotation
+
+ if start_time < 0:
+ return ValidationResult(
+ valid=False,
+ error="Start time cannot be negative"
+ )
+
+ if end_time < start_time:
+ return ValidationResult(
+ valid=False,
+ error="End time must be >= start time"
+ )
+
+ return ValidationResult(valid=True)
+
+ def _validate_time_subdivision(
+ self,
+ start_time: Optional[float],
+ end_time: Optional[float],
+ parent: Annotation
+ ) -> ValidationResult:
+ """Validate TIME_SUBDIVISION constraint."""
+ if start_time is None or end_time is None:
+ return ValidationResult(
+ valid=False,
+ error="Time subdivision requires time-aligned annotation"
+ )
+
+ if parent.start_time is None or parent.end_time is None:
+ return ValidationResult(
+ valid=False,
+ error="Parent must be time-aligned for time subdivision"
+ )
+
+ # Child must be within parent bounds
+ if start_time < parent.start_time:
+ return ValidationResult(
+ valid=False,
+ error=f"Start time ({start_time}ms) is before parent start ({parent.start_time}ms)"
+ )
+
+ if end_time > parent.end_time:
+ return ValidationResult(
+ valid=False,
+ error=f"End time ({end_time}ms) is after parent end ({parent.end_time}ms)"
+ )
+
+ return ValidationResult(valid=True)
+
+ def _validate_included_in(
+ self,
+ start_time: Optional[float],
+ end_time: Optional[float],
+ parent: Annotation
+ ) -> ValidationResult:
+ """Validate INCLUDED_IN constraint."""
+ if start_time is None or end_time is None:
+ return ValidationResult(
+ valid=False,
+ error="Included-in requires time-aligned annotation"
+ )
+
+ if parent.start_time is None or parent.end_time is None:
+ return ValidationResult(
+ valid=False,
+ error="Parent must be time-aligned for included-in"
+ )
+
+ # Child must be within parent bounds (same as subdivision)
+ if start_time < parent.start_time or end_time > parent.end_time:
+ return ValidationResult(
+ valid=False,
+ error="Annotation must be within parent time bounds"
+ )
+
+ return ValidationResult(valid=True)
+
+ def add_annotation(self, annotation: Annotation) -> ValidationResult:
+ """
+ Add an annotation after validating it.
+
+ Args:
+ annotation: The annotation to add
+
+ Returns:
+ ValidationResult indicating success or failure
+ """
+ tier = self.definition.get_tier(annotation.tier)
+ if not tier:
+ return ValidationResult(
+ valid=False,
+ error=f"Unknown tier: '{annotation.tier}'"
+ )
+
+ # Find parent if needed
+ parent = None
+ if tier.is_dependent and annotation.parent_id:
+ parent = self.get_annotation(annotation.parent_id)
+ if not parent:
+ return ValidationResult(
+ valid=False,
+ error=f"Parent annotation '{annotation.parent_id}' not found"
+ )
+ elif tier.is_dependent and annotation.start_time is not None:
+ # Try to find parent by time range
+ parent = self.find_parent_annotation(
+ annotation.tier,
+ annotation.start_time,
+ annotation.end_time or annotation.start_time
+ )
+ if parent:
+ annotation.parent_id = parent.id
+
+ # Validate the annotation
+ result = self.validate_annotation(
+ annotation.tier,
+ annotation.start_time,
+ annotation.end_time,
+ parent
+ )
+
+ if not result.valid:
+ return result
+
+ # Add to the appropriate tier
+ self._annotations[annotation.tier].append(annotation)
+ return ValidationResult(valid=True)
+
+ def remove_annotation(
+ self,
+ annotation_id: str,
+ cascade: bool = True
+ ) -> Tuple[bool, List[str]]:
+ """
+ Remove an annotation and optionally its dependents.
+
+ Args:
+ annotation_id: ID of annotation to remove
+ cascade: If True, also remove child annotations
+
+ Returns:
+ Tuple of (success, list of removed annotation IDs)
+ """
+ annotation = self.get_annotation(annotation_id)
+ if not annotation:
+ return False, []
+
+ removed_ids = [annotation_id]
+
+ # Remove from its tier
+ tier_annotations = self._annotations.get(annotation.tier, [])
+ self._annotations[annotation.tier] = [
+ a for a in tier_annotations if a.id != annotation_id
+ ]
+
+ # Cascade delete children if requested
+ if cascade:
+ tier = self.definition.get_tier(annotation.tier)
+ if tier:
+ for child_tier in self.definition.get_children(tier.name):
+ child_annotations = self._annotations.get(child_tier.name, [])
+ for child in child_annotations[:]: # Copy to allow modification
+ if child.parent_id == annotation_id:
+ _, child_removed = self.remove_annotation(
+ child.id, cascade=True
+ )
+ removed_ids.extend(child_removed)
+
+ return True, removed_ids
+
+ def clear(self) -> None:
+ """Clear all annotations."""
+ for tier_name in self._annotations:
+ self._annotations[tier_name] = []
+
+ def load_annotations(
+ self,
+ annotations_data: Dict[str, List[Dict[str, Any]]]
+ ) -> List[str]:
+ """
+ Load annotations from serialized data.
+
+ Args:
+ annotations_data: Dict mapping tier names to list of annotation dicts
+
+ Returns:
+ List of error messages (empty if successful)
+ """
+ errors = []
+ self.clear()
+
+ # Load in tier order to ensure parents exist before children
+ for tier in self.definition.tiers:
+ tier_data = annotations_data.get(tier.name, [])
+ for ann_data in tier_data:
+ annotation = Annotation.from_dict(ann_data)
+ annotation.tier = tier.name # Ensure tier is set
+ result = self.add_annotation(annotation)
+ if not result.valid:
+ errors.append(f"Tier '{tier.name}': {result.error}")
+
+ return errors
+
+ def serialize(self) -> Dict[str, List[Dict[str, Any]]]:
+ """
+ Serialize all annotations for storage.
+
+ Returns:
+ Dict mapping tier names to list of annotation dicts
+ """
+ return {
+ tier_name: [ann.to_dict() for ann in annotations]
+ for tier_name, annotations in self._annotations.items()
+ }
+
+ def generate_time_slots(self) -> Dict[str, int]:
+ """
+ Generate ELAN-style time slots from all annotations.
+
+ Returns:
+ Dict mapping slot ID to time in milliseconds
+ """
+ times: Set[int] = set()
+
+ for annotations in self._annotations.values():
+ for ann in annotations:
+ if ann.start_time is not None:
+ times.add(int(ann.start_time))
+ if ann.end_time is not None:
+ times.add(int(ann.end_time))
+
+ # Generate slot IDs in chronological order
+ return {
+ f"ts{i+1}": time
+ for i, time in enumerate(sorted(times))
+ }
+
+ def get_time_slot_id(
+ self,
+ time_ms: int,
+ time_slots: Dict[str, int]
+ ) -> Optional[str]:
+ """Find the time slot ID for a given time value."""
+ for slot_id, slot_time in time_slots.items():
+ if slot_time == time_ms:
+ return slot_id
+ return None
diff --git a/potato/ibws_manager.py b/potato/ibws_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..6f9973d24f9bde702fe0972ded9db9be899ea062
--- /dev/null
+++ b/potato/ibws_manager.py
@@ -0,0 +1,393 @@
+"""
+Iterative Best-Worst Scaling (IBWS) Manager
+
+Implements the IBWS algorithm from "Baby Bear: Seeking a Just Right Rating Scale
+for Scalar Annotations" (arxiv 2408.09765). IBWS extends standard BWS with a
+Quicksort-like adaptive loop:
+
+1. Round 1: Generate tuples from full pool, annotators select best/worst
+2. Score items, partition into upper/middle/lower buckets
+3. Round N: Generate tuples WITHIN each bucket, annotate, partition again
+4. Stop when all buckets are terminal (< tuple_size items) or max_rounds reached
+
+Output: Ordinal ranking from bucket positions + within-bucket scores.
+
+Usage:
+ from potato.ibws_manager import get_ibws_manager, init_ibws_manager
+
+ mgr = init_ibws_manager(config, pool_items, id_key, text_key)
+ round1_tuples = mgr.get_current_round_tuples()
+
+ # After annotations are complete for current round:
+ if mgr.check_round_complete(ism, bws_schema_name):
+ new_tuples = mgr.advance_round(ism, bws_schema_name)
+ # Add new_tuples to ISM
+"""
+
+import logging
+import math
+import threading
+from typing import Any, Dict, List, Optional, Tuple
+
+from potato.bws_scoring import BwsScorer
+from potato.bws_tuple_generator import BwsTupleGenerator
+
+logger = logging.getLogger(__name__)
+
+# Singleton instance
+_ibws_manager = None
+_ibws_lock = threading.Lock()
+
+
+def init_ibws_manager(config: dict, pool_items: List[Dict[str, Any]],
+ id_key: str, text_key: str) -> "IBWSManager":
+ """Initialize the global IBWS manager singleton."""
+ global _ibws_manager
+ with _ibws_lock:
+ _ibws_manager = IBWSManager(config, pool_items, id_key, text_key)
+ return _ibws_manager
+
+
+def get_ibws_manager() -> Optional["IBWSManager"]:
+ """Get the global IBWS manager (None if not initialized)."""
+ return _ibws_manager
+
+
+def clear_ibws_manager():
+ """Clear the global IBWS manager (for testing)."""
+ global _ibws_manager
+ with _ibws_lock:
+ _ibws_manager = None
+
+
+class IBWSManager:
+ """Manages iterative BWS rounds, partitioning, and tuple generation."""
+
+ def __init__(self, config: dict, pool_items: List[Dict[str, Any]],
+ id_key: str, text_key: str):
+ self._lock = threading.RLock()
+
+ self.id_key = id_key
+ self.text_key = text_key
+
+ ibws_config = config["ibws_config"]
+ self.tuple_size = ibws_config.get("tuple_size", 4)
+ self.max_rounds = ibws_config.get("max_rounds", None) # None = auto
+ self.seed = ibws_config.get("seed", 42)
+ self.scoring_method = ibws_config.get("scoring_method", "counting")
+ self.tuples_per_item_per_round = ibws_config.get("tuples_per_item_per_round", 2)
+
+ # Store original pool items with their IDs
+ self.pool_items = list(pool_items)
+ self.pool_item_map = {str(item[id_key]): item for item in pool_items}
+
+ # Partition state: list of buckets, each bucket is a list of item IDs
+ # Start with one bucket containing all items
+ self.current_round = 0 # 0 = not started, 1 = round 1 active, etc.
+ self.buckets: List[List[str]] = [[str(item[id_key]) for item in pool_items]]
+ self.terminal_buckets: List[List[str]] = [] # Buckets too small to partition further
+
+ # Track tuples for each round: round_num -> list of tuple IDs
+ self.round_tuples: Dict[int, List[str]] = {}
+
+ # Track all generated tuples' data for scoring
+ self.tuple_data: Dict[str, Dict[str, Any]] = {}
+
+ # Completed flag
+ self.completed = False
+
+ def generate_round_tuples(self) -> List[Dict[str, Any]]:
+ """Generate tuples for the next round from current active buckets.
+
+ Returns list of tuple instance dicts ready for ISM.add_item().
+ """
+ with self._lock:
+ self.current_round += 1
+ round_num = self.current_round
+
+ all_tuples = []
+ tuple_ids = []
+
+ new_buckets = []
+ for bucket_idx, bucket_item_ids in enumerate(self.buckets):
+ if len(bucket_item_ids) < self.tuple_size:
+ # Terminal bucket โ too few items to form a tuple
+ self.terminal_buckets.append(bucket_item_ids)
+ continue
+
+ new_buckets.append(bucket_item_ids)
+
+ # Build pool items for this bucket
+ bucket_pool = [self.pool_item_map[iid] for iid in bucket_item_ids
+ if iid in self.pool_item_map]
+
+ if len(bucket_pool) < self.tuple_size:
+ self.terminal_buckets.append(bucket_item_ids)
+ continue
+
+ # Calculate tuples needed for this bucket
+ min_appearances = self.tuples_per_item_per_round * self.tuple_size
+ num_tuples = max(1, math.ceil(
+ len(bucket_pool) * self.tuples_per_item_per_round / self.tuple_size
+ ))
+
+ prefix = f"ibws_r{round_num}_b{bucket_idx}"
+ generator = BwsTupleGenerator(
+ pool_items=bucket_pool,
+ id_key=self.id_key,
+ text_key=self.text_key,
+ tuple_size=self.tuple_size,
+ num_tuples=num_tuples,
+ seed=self.seed + round_num * 1000 + bucket_idx,
+ min_item_appearances=min_appearances,
+ )
+
+ tuples = generator.generate()
+
+ # Rename tuple IDs with our prefix
+ for i, t in enumerate(tuples):
+ new_id = f"{prefix}_{i + 1:04d}"
+ t[self.id_key] = new_id
+ t["_ibws_round"] = round_num
+ t["_ibws_bucket"] = bucket_idx
+ self.tuple_data[new_id] = t
+ tuple_ids.append(new_id)
+
+ all_tuples.extend(tuples)
+
+ # Update active buckets (excluding those that became terminal)
+ self.buckets = new_buckets
+ self.round_tuples[round_num] = tuple_ids
+
+ if not all_tuples:
+ # All buckets are terminal
+ self.completed = True
+
+ logger.info(
+ f"IBWS round {round_num}: Generated {len(all_tuples)} tuples "
+ f"across {len(self.buckets)} active buckets "
+ f"({len(self.terminal_buckets)} terminal)"
+ )
+
+ return all_tuples
+
+ def check_round_complete(self, ism, bws_schema_name: str) -> bool:
+ """Check if all tuples in the current round have been annotated.
+
+ Uses ISM's instance_annotators tracking to see if each tuple
+ has at least one annotator.
+
+ Args:
+ ism: ItemStateManager instance
+ bws_schema_name: Name of the BWS annotation schema
+
+ Returns:
+ True if all current round tuples have at least one annotation
+ """
+ with self._lock:
+ if self.completed or self.current_round == 0:
+ return False
+
+ round_tuple_ids = self.round_tuples.get(self.current_round, [])
+ if not round_tuple_ids:
+ return False
+
+ # Check that every tuple in this round has at least one annotator
+ for tuple_id in round_tuple_ids:
+ annotators = ism.instance_annotators.get(tuple_id, set())
+ if not annotators:
+ return False
+
+ return True
+
+ def advance_round(self, ism, usm, bws_schema_name: str) -> List[Dict[str, Any]]:
+ """Score current round, partition buckets, generate next round tuples.
+
+ Args:
+ ism: ItemStateManager instance
+ usm: UserStateManager instance
+ bws_schema_name: Name of the BWS annotation schema
+
+ Returns:
+ List of new tuple instance dicts for the next round (empty if done)
+ """
+ with self._lock:
+ if self.completed:
+ return []
+
+ if self.max_rounds and self.current_round >= self.max_rounds:
+ self.completed = True
+ logger.info(f"IBWS: Reached max_rounds ({self.max_rounds}), stopping")
+ return []
+
+ # Score current round and partition each active bucket
+ new_buckets = []
+ for bucket_idx, bucket_item_ids in enumerate(self.buckets):
+ if len(bucket_item_ids) < self.tuple_size:
+ self.terminal_buckets.append(bucket_item_ids)
+ continue
+
+ # Collect annotations for tuples that contain items from this bucket
+ annotations = self._collect_bucket_annotations(
+ bucket_item_ids, ism, usm, bws_schema_name
+ )
+
+ if not annotations:
+ # No annotations โ can't partition, keep bucket as-is
+ new_buckets.append(bucket_item_ids)
+ continue
+
+ # Score items in this bucket
+ bucket_pool = [self.pool_item_map[iid] for iid in bucket_item_ids
+ if iid in self.pool_item_map]
+ scorer = BwsScorer(annotations, bucket_pool, self.id_key, self.text_key)
+ scores = scorer.score(self.scoring_method)
+
+ # Partition into upper/middle/lower thirds
+ upper, middle, lower = self._partition_bucket(bucket_item_ids, scores)
+
+ for sub_bucket in [upper, middle, lower]:
+ if sub_bucket:
+ new_buckets.append(sub_bucket)
+
+ self.buckets = new_buckets
+
+ # Check if all remaining buckets are terminal
+ active_count = sum(1 for b in self.buckets if len(b) >= self.tuple_size)
+ if active_count == 0:
+ # Move remaining small buckets to terminal
+ for b in self.buckets:
+ if len(b) < self.tuple_size:
+ self.terminal_buckets.append(b)
+ self.buckets = []
+ self.completed = True
+ logger.info("IBWS: All buckets terminal, annotation complete")
+ return []
+
+ # Generate tuples for the next round
+ return self.generate_round_tuples()
+
+ def _collect_bucket_annotations(self, bucket_item_ids, ism, usm,
+ bws_schema_name: str) -> List[Dict[str, Any]]:
+ """Collect BWS annotations for tuples containing items from a bucket."""
+ bucket_id_set = set(bucket_item_ids)
+ annotations = []
+
+ # Look at tuples from the current round
+ round_tuple_ids = self.round_tuples.get(self.current_round, [])
+
+ for tuple_id in round_tuple_ids:
+ tuple_info = self.tuple_data.get(tuple_id)
+ if not tuple_info:
+ continue
+
+ bws_items = tuple_info.get("_bws_items", [])
+ # Check if this tuple's items overlap with our bucket
+ tuple_source_ids = {item["source_id"] for item in bws_items}
+ if not tuple_source_ids.intersection(bucket_id_set):
+ continue
+
+ # Collect annotations from all users for this tuple
+ for user_state in usm.get_all_users():
+ username = user_state.get_user_id()
+ label_store = getattr(user_state, 'instance_id_to_label_to_value', {})
+
+ if tuple_id not in label_store:
+ continue
+
+ labels = label_store[tuple_id]
+ best_val = None
+ worst_val = None
+ for label_obj, value in labels.items():
+ if label_obj.get_schema() == bws_schema_name:
+ if label_obj.get_name() == "best":
+ best_val = value
+ elif label_obj.get_name() == "worst":
+ worst_val = value
+
+ if best_val and worst_val:
+ annotations.append({
+ "instance_id": tuple_id,
+ "bws_items": bws_items,
+ "best": best_val,
+ "worst": worst_val,
+ "annotator": username,
+ })
+
+ return annotations
+
+ def _partition_bucket(self, item_ids: List[str],
+ scores: Dict[str, Dict[str, Any]]) -> Tuple[List[str], List[str], List[str]]:
+ """Partition a bucket into upper/middle/lower thirds by score.
+
+ Uses equal-thirds of sorted list (not score thresholds) for balanced partitions.
+ """
+ # Sort by score descending
+ sorted_ids = sorted(
+ item_ids,
+ key=lambda iid: scores.get(iid, {}).get("score", 0.0),
+ reverse=True
+ )
+
+ n = len(sorted_ids)
+ third = n // 3
+
+ # Handle remainder: distribute extra items to middle
+ upper = sorted_ids[:third]
+ lower = sorted_ids[n - third:] if third > 0 else []
+ middle = sorted_ids[third:n - third] if third > 0 else sorted_ids
+
+ return upper, middle, lower
+
+ def get_round_info(self) -> Dict[str, Any]:
+ """Get current round information for UI display."""
+ with self._lock:
+ total_tuples_this_round = len(self.round_tuples.get(self.current_round, []))
+ active_buckets = len([b for b in self.buckets if len(b) >= self.tuple_size])
+ terminal_count = len(self.terminal_buckets)
+ total_items = len(self.pool_items)
+
+ # Items in terminal buckets (already ranked)
+ terminal_items = sum(len(b) for b in self.terminal_buckets)
+
+ return {
+ "current_round": self.current_round,
+ "max_rounds": self.max_rounds,
+ "total_tuples_this_round": total_tuples_this_round,
+ "active_buckets": active_buckets,
+ "terminal_buckets": terminal_count,
+ "total_items": total_items,
+ "terminal_items": terminal_items,
+ "completed": self.completed,
+ }
+
+ def get_final_ranking(self) -> List[Dict[str, Any]]:
+ """Produce final ordinal ranking from bucket positions + within-bucket scores.
+
+ Returns list of dicts sorted by rank:
+ [{"item_id": str, "rank": int, "bucket_position": int, "text": str}, ...]
+ """
+ with self._lock:
+ # Combine terminal buckets (ordered by when they became terminal = higher quality)
+ # and any remaining active buckets
+ all_buckets = list(self.terminal_buckets) + list(self.buckets)
+
+ ranking = []
+ rank = 1
+ for bucket_position, bucket in enumerate(all_buckets):
+ for item_id in bucket:
+ item = self.pool_item_map.get(item_id, {})
+ ranking.append({
+ "item_id": item_id,
+ "rank": rank,
+ "bucket_position": bucket_position,
+ "text": str(item.get(self.text_key, "")),
+ })
+ rank += 1
+
+ return ranking
+
+ def is_completed(self) -> bool:
+ """Check if IBWS has completed all rounds."""
+ with self._lock:
+ return self.completed
diff --git a/potato/integrations/__init__.py b/potato/integrations/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4c35b7a3d1de736492209520b10ee8de2817523b
--- /dev/null
+++ b/potato/integrations/__init__.py
@@ -0,0 +1,9 @@
+"""
+Potato Integrations
+
+Third-party integrations for connecting external ML frameworks
+to Potato's annotation platform.
+
+Available integrations:
+- LangChain callback handler: Auto-send LangChain traces to Potato
+"""
diff --git a/potato/integrations/langchain_callback.py b/potato/integrations/langchain_callback.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d5dbe0917715dd5b851b4a8bc77b278703409dd
--- /dev/null
+++ b/potato/integrations/langchain_callback.py
@@ -0,0 +1,393 @@
+"""
+LangChain Callback Handler for Potato
+
+Automatically sends LangChain agent traces to a Potato instance
+for human evaluation and annotation.
+
+Requires: pip install langchain-core>=0.1.0
+
+Usage:
+ from potato.integrations.langchain_callback import PotatoCallbackHandler
+
+ handler = PotatoCallbackHandler(potato_url="http://localhost:8000")
+ chain.invoke({"input": "..."}, config={"callbacks": [handler]})
+"""
+
+import json
+import logging
+import threading
+import time
+import uuid
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+def _safe_serialize(obj: Any) -> Any:
+ """Convert an object to a JSON-safe representation."""
+ if obj is None or isinstance(obj, (str, int, float, bool)):
+ return obj
+ if isinstance(obj, dict):
+ return {str(k): _safe_serialize(v) for k, v in obj.items()}
+ if isinstance(obj, (list, tuple)):
+ return [_safe_serialize(v) for v in obj]
+ # Fall back to string representation
+ try:
+ return str(obj)
+ except Exception:
+ return ""
+
+
+class PotatoCallbackHandler:
+ """
+ LangChain callback handler that sends completed traces to Potato.
+
+ Collects run events (chain, LLM, tool starts/ends), tracks
+ parent-child relationships, and POSTs the full trace to Potato's
+ webhook endpoint when the root chain completes.
+
+ The payload uses the LangSmith format expected by
+ ``POST /api/traces/langsmith``.
+
+ Args:
+ potato_url: Base URL of the Potato server (e.g., ``http://localhost:8000``)
+ api_key: API key for authenticating with Potato's webhook endpoint
+ endpoint: Webhook path (default ``/api/traces/langsmith``)
+ send_timeout: HTTP timeout in seconds for the POST request (default 10)
+ metadata: Extra metadata dict attached to every trace
+ """
+
+ def __init__(
+ self,
+ potato_url: str,
+ api_key: str = "",
+ endpoint: str = "/api/traces/langsmith",
+ send_timeout: int = 10,
+ metadata: Optional[Dict[str, Any]] = None,
+ ):
+ self.potato_url = potato_url.rstrip("/")
+ self.api_key = api_key
+ self.endpoint = endpoint
+ self.send_timeout = send_timeout
+ self.extra_metadata = metadata or {}
+
+ # Run tracking โ protected by lock for thread-safety
+ self._lock = threading.Lock()
+ self._runs: Dict[str, dict] = {} # run_id -> run dict
+ self._root_run_id: Optional[str] = None
+ self._pending_sends: List[threading.Thread] = []
+
+ # ------------------------------------------------------------------
+ # LangChain BaseCallbackHandler interface
+ # ------------------------------------------------------------------
+
+ def on_chain_start(
+ self,
+ serialized: Dict[str, Any],
+ inputs: Dict[str, Any],
+ *,
+ run_id: uuid.UUID,
+ parent_run_id: Optional[uuid.UUID] = None,
+ tags: Optional[List[str]] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ **kwargs: Any,
+ ) -> None:
+ self._start_run(
+ run_id=str(run_id),
+ parent_run_id=str(parent_run_id) if parent_run_id else None,
+ run_type="chain",
+ name=serialized.get("name", serialized.get("id", ["unknown"])[-1]
+ if isinstance(serialized.get("id"), list) else "chain"),
+ inputs=inputs,
+ tags=tags,
+ metadata=metadata,
+ )
+
+ def on_chain_end(
+ self,
+ outputs: Dict[str, Any],
+ *,
+ run_id: uuid.UUID,
+ parent_run_id: Optional[uuid.UUID] = None,
+ **kwargs: Any,
+ ) -> None:
+ self._end_run(str(run_id), outputs=outputs)
+
+ def on_chain_error(
+ self,
+ error: BaseException,
+ *,
+ run_id: uuid.UUID,
+ parent_run_id: Optional[uuid.UUID] = None,
+ **kwargs: Any,
+ ) -> None:
+ self._end_run(str(run_id), error=error)
+
+ def on_llm_start(
+ self,
+ serialized: Dict[str, Any],
+ prompts: List[str],
+ *,
+ run_id: uuid.UUID,
+ parent_run_id: Optional[uuid.UUID] = None,
+ tags: Optional[List[str]] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ **kwargs: Any,
+ ) -> None:
+ self._start_run(
+ run_id=str(run_id),
+ parent_run_id=str(parent_run_id) if parent_run_id else None,
+ run_type="llm",
+ name=serialized.get("name", "llm"),
+ inputs={"prompts": prompts},
+ tags=tags,
+ metadata=metadata,
+ )
+
+ def on_llm_end(
+ self,
+ response: Any,
+ *,
+ run_id: uuid.UUID,
+ parent_run_id: Optional[uuid.UUID] = None,
+ **kwargs: Any,
+ ) -> None:
+ output = {}
+ if hasattr(response, "generations") and response.generations:
+ texts = []
+ for gen_list in response.generations:
+ for gen in gen_list:
+ texts.append(gen.text if hasattr(gen, "text") else str(gen))
+ output = {"text": "\n".join(texts)}
+ self._end_run(str(run_id), outputs=output)
+
+ def on_llm_error(
+ self,
+ error: BaseException,
+ *,
+ run_id: uuid.UUID,
+ parent_run_id: Optional[uuid.UUID] = None,
+ **kwargs: Any,
+ ) -> None:
+ self._end_run(str(run_id), error=error)
+
+ def on_tool_start(
+ self,
+ serialized: Dict[str, Any],
+ input_str: str,
+ *,
+ run_id: uuid.UUID,
+ parent_run_id: Optional[uuid.UUID] = None,
+ tags: Optional[List[str]] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ **kwargs: Any,
+ ) -> None:
+ self._start_run(
+ run_id=str(run_id),
+ parent_run_id=str(parent_run_id) if parent_run_id else None,
+ run_type="tool",
+ name=serialized.get("name", "tool"),
+ inputs={"input": input_str},
+ tags=tags,
+ metadata=metadata,
+ )
+
+ def on_tool_end(
+ self,
+ output: str,
+ *,
+ run_id: uuid.UUID,
+ parent_run_id: Optional[uuid.UUID] = None,
+ **kwargs: Any,
+ ) -> None:
+ self._end_run(str(run_id), outputs={"output": output})
+
+ def on_tool_error(
+ self,
+ error: BaseException,
+ *,
+ run_id: uuid.UUID,
+ parent_run_id: Optional[uuid.UUID] = None,
+ **kwargs: Any,
+ ) -> None:
+ self._end_run(str(run_id), error=error)
+
+ # Retriever callbacks
+ def on_retriever_start(
+ self,
+ serialized: Dict[str, Any],
+ query: str,
+ *,
+ run_id: uuid.UUID,
+ parent_run_id: Optional[uuid.UUID] = None,
+ tags: Optional[List[str]] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ **kwargs: Any,
+ ) -> None:
+ self._start_run(
+ run_id=str(run_id),
+ parent_run_id=str(parent_run_id) if parent_run_id else None,
+ run_type="retriever",
+ name=serialized.get("name", "retriever"),
+ inputs={"query": query},
+ tags=tags,
+ metadata=metadata,
+ )
+
+ def on_retriever_end(
+ self,
+ documents: Any,
+ *,
+ run_id: uuid.UUID,
+ parent_run_id: Optional[uuid.UUID] = None,
+ **kwargs: Any,
+ ) -> None:
+ output = {"documents": _safe_serialize(documents)}
+ self._end_run(str(run_id), outputs=output)
+
+ # Text callbacks (no-ops โ captured by LLM callbacks)
+ def on_text(self, text: str, **kwargs: Any) -> None:
+ pass
+
+ # ------------------------------------------------------------------
+ # Internal helpers
+ # ------------------------------------------------------------------
+
+ def _start_run(
+ self,
+ run_id: str,
+ parent_run_id: Optional[str],
+ run_type: str,
+ name: str,
+ inputs: Any,
+ tags: Optional[List[str]] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ ) -> None:
+ run = {
+ "id": run_id,
+ "parent_run_id": parent_run_id,
+ "run_type": run_type,
+ "name": name,
+ "inputs": _safe_serialize(inputs),
+ "outputs": {},
+ "status": "running",
+ "start_time": time.time(),
+ "end_time": None,
+ "tags": tags or [],
+ "metadata": metadata or {},
+ }
+
+ with self._lock:
+ self._runs[run_id] = run
+ if parent_run_id is None:
+ self._root_run_id = run_id
+
+ def _end_run(
+ self,
+ run_id: str,
+ outputs: Optional[Dict[str, Any]] = None,
+ error: Optional[BaseException] = None,
+ ) -> None:
+ with self._lock:
+ run = self._runs.get(run_id)
+ if run is None:
+ return
+
+ run["end_time"] = time.time()
+ run["latency"] = run["end_time"] - run["start_time"]
+
+ if error:
+ run["status"] = "error"
+ run["outputs"] = {"error": str(error)}
+ else:
+ run["status"] = "completed"
+ run["outputs"] = _safe_serialize(outputs or {})
+
+ # If this is the root run, send the trace
+ is_root = run_id == self._root_run_id
+
+ if is_root:
+ self._send_trace()
+
+ def _build_payload(self) -> dict:
+ """Build a LangSmith-format payload from collected runs."""
+ with self._lock:
+ runs = list(self._runs.values())
+ root_id = self._root_run_id
+
+ # Find the root run for metadata
+ root_run = None
+ for r in runs:
+ if r["id"] == root_id:
+ root_run = r
+ break
+
+ payload = {
+ "runs": [
+ {
+ "id": r["id"],
+ "parent_run_id": r["parent_run_id"],
+ "run_type": r["run_type"],
+ "name": r["name"],
+ "inputs": r["inputs"],
+ "outputs": r["outputs"],
+ "status": r["status"],
+ "latency": r.get("latency"),
+ "tags": r.get("tags", []),
+ }
+ for r in runs
+ ],
+ }
+
+ if root_run:
+ payload["project_name"] = root_run.get("name", "langchain")
+
+ return payload
+
+ def _send_trace(self) -> None:
+ """POST the trace to Potato in a background thread."""
+ payload = self._build_payload()
+
+ def _do_send():
+ try:
+ url = f"{self.potato_url}{self.endpoint}"
+ headers = {"Content-Type": "application/json"}
+ if self.api_key:
+ headers["Authorization"] = f"Bearer {self.api_key}"
+
+ resp = requests.post(
+ url,
+ json=payload,
+ headers=headers,
+ timeout=self.send_timeout,
+ )
+ if resp.status_code < 300:
+ logger.info("Trace sent to Potato: %s", resp.json())
+ else:
+ logger.warning(
+ "Potato returned %s: %s", resp.status_code, resp.text
+ )
+ except Exception as e:
+ logger.error("Failed to send trace to Potato: %s", e)
+
+ thread = threading.Thread(target=_do_send, daemon=True)
+ with self._lock:
+ self._pending_sends.append(thread)
+ thread.start()
+
+ def flush(self, timeout: float = 30.0) -> None:
+ """Block until all pending sends complete (or timeout)."""
+ with self._lock:
+ threads = list(self._pending_sends)
+ for t in threads:
+ t.join(timeout=timeout)
+ with self._lock:
+ self._pending_sends = [t for t in self._pending_sends if t.is_alive()]
+
+ def reset(self) -> None:
+ """Clear all collected runs (for reuse across multiple chains)."""
+ with self._lock:
+ self._runs.clear()
+ self._root_run_id = None
diff --git a/potato/interaction_tracking.py b/potato/interaction_tracking.py
new file mode 100644
index 0000000000000000000000000000000000000000..17844b812f55966f69cee415c950079244b1e412
--- /dev/null
+++ b/potato/interaction_tracking.py
@@ -0,0 +1,399 @@
+"""
+Interaction tracking data structures and utilities for behavioral analysis.
+
+This module provides dataclasses for tracking user interactions during annotation,
+including clicks, focus changes, navigation, AI assistance usage, and annotation changes.
+All data is designed to be serializable for persistence and later analysis.
+"""
+from dataclasses import dataclass, field, asdict
+from typing import Dict, List, Any, Optional
+import time
+
+
+@dataclass
+class InteractionEvent:
+ """
+ A single user interaction with the annotation interface.
+
+ Attributes:
+ event_type: Type of interaction ("click", "focus_in", "focus_out",
+ "navigation", "save", "scroll", "keypress", etc.)
+ timestamp: Server-side Unix timestamp when event was recorded
+ target: Element identifier (e.g., "label:positive", "nav:next", "schema:sentiment")
+ instance_id: The annotation instance this event occurred on
+ client_timestamp: Client-side timestamp in milliseconds (for latency analysis)
+ metadata: Additional context (position, value changes, duration, etc.)
+ """
+ event_type: str
+ timestamp: float
+ target: str
+ instance_id: str
+ client_timestamp: Optional[float] = None
+ metadata: Dict[str, Any] = field(default_factory=dict)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to dictionary for JSON serialization."""
+ return {
+ 'event_type': self.event_type,
+ 'timestamp': self.timestamp,
+ 'target': self.target,
+ 'instance_id': self.instance_id,
+ 'client_timestamp': self.client_timestamp,
+ 'metadata': self.metadata,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'InteractionEvent':
+ """Reconstruct from serialized dictionary."""
+ return cls(
+ event_type=data.get('event_type', ''),
+ timestamp=data.get('timestamp', 0),
+ target=data.get('target', ''),
+ instance_id=data.get('instance_id', ''),
+ client_timestamp=data.get('client_timestamp'),
+ metadata=data.get('metadata', {}),
+ )
+
+
+@dataclass
+class AIUsageEvent:
+ """
+ Tracks AI assistance usage for an annotation instance.
+
+ Captures the full lifecycle of an AI assistance request:
+ request -> response -> user decision (accept/reject/ignore)
+
+ Attributes:
+ request_timestamp: When AI assistance was requested
+ schema_name: Which annotation schema the AI assisted with
+ suggestions_shown: List of labels/values the AI suggested
+ response_timestamp: When the AI response was received
+ suggestion_accepted: The value the user accepted (None if rejected/ignored)
+ final_annotation: What the user ultimately annotated for this schema
+ time_to_decision_ms: Milliseconds from response to user action
+ """
+ request_timestamp: float
+ schema_name: str
+ suggestions_shown: List[str] = field(default_factory=list)
+ response_timestamp: Optional[float] = None
+ suggestion_accepted: Optional[str] = None
+ final_annotation: Optional[str] = None
+ time_to_decision_ms: Optional[int] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to dictionary for JSON serialization."""
+ return {
+ 'request_timestamp': self.request_timestamp,
+ 'schema_name': self.schema_name,
+ 'suggestions_shown': self.suggestions_shown,
+ 'response_timestamp': self.response_timestamp,
+ 'suggestion_accepted': self.suggestion_accepted,
+ 'final_annotation': self.final_annotation,
+ 'time_to_decision_ms': self.time_to_decision_ms,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'AIUsageEvent':
+ """Reconstruct from serialized dictionary."""
+ return cls(
+ request_timestamp=data.get('request_timestamp', 0),
+ schema_name=data.get('schema_name', ''),
+ suggestions_shown=data.get('suggestions_shown', []),
+ response_timestamp=data.get('response_timestamp'),
+ suggestion_accepted=data.get('suggestion_accepted'),
+ final_annotation=data.get('final_annotation'),
+ time_to_decision_ms=data.get('time_to_decision_ms'),
+ )
+
+
+@dataclass
+class AnnotationChange:
+ """
+ Records a single change to an annotation.
+
+ Attributes:
+ timestamp: When the change occurred
+ schema_name: Which schema was modified
+ label_name: Which label was affected (if applicable)
+ action: Type of change ("select", "deselect", "update", "clear")
+ old_value: Previous value (if any)
+ new_value: New value after the change
+ source: What triggered the change ("user", "ai_accept", "prefill", "keyboard")
+ """
+ timestamp: float
+ schema_name: str
+ action: str
+ label_name: Optional[str] = None
+ old_value: Any = None
+ new_value: Any = None
+ source: str = "user"
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to dictionary for JSON serialization."""
+ return {
+ 'timestamp': self.timestamp,
+ 'schema_name': self.schema_name,
+ 'label_name': self.label_name,
+ 'action': self.action,
+ 'old_value': self.old_value,
+ 'new_value': self.new_value,
+ 'source': self.source,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'AnnotationChange':
+ """Reconstruct from serialized dictionary."""
+ return cls(
+ timestamp=data.get('timestamp', 0),
+ schema_name=data.get('schema_name', ''),
+ label_name=data.get('label_name'),
+ action=data.get('action', ''),
+ old_value=data.get('old_value'),
+ new_value=data.get('new_value'),
+ source=data.get('source', 'user'),
+ )
+
+
+@dataclass
+class ChatMessage:
+ """
+ A single chat message between an annotator and the LLM assistant.
+
+ Attributes:
+ role: 'user' or 'assistant'
+ content: The message text
+ timestamp: Unix timestamp when the message was sent/received
+ instance_id: The annotation instance this message relates to
+ response_time_ms: Milliseconds for the LLM to respond (assistant messages only)
+ """
+ role: str
+ content: str
+ timestamp: float
+ instance_id: str
+ response_time_ms: Optional[int] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to dictionary for JSON serialization."""
+ return {
+ 'role': self.role,
+ 'content': self.content,
+ 'timestamp': self.timestamp,
+ 'instance_id': self.instance_id,
+ 'response_time_ms': self.response_time_ms,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'ChatMessage':
+ """Reconstruct from serialized dictionary."""
+ return cls(
+ role=data.get('role', ''),
+ content=data.get('content', ''),
+ timestamp=data.get('timestamp', 0),
+ instance_id=data.get('instance_id', ''),
+ response_time_ms=data.get('response_time_ms'),
+ )
+
+
+@dataclass
+class BehavioralData:
+ """
+ Complete behavioral data for an annotation instance session.
+
+ Aggregates all tracking data for a single instance annotation session,
+ including timing, interactions, AI usage, and annotation changes.
+
+ Attributes:
+ instance_id: The annotation instance ID
+ session_start: Unix timestamp when user first loaded this instance
+ session_end: Unix timestamp when user navigated away or saved
+ total_time_ms: Total milliseconds spent on this instance
+ interactions: List of all interaction events
+ ai_usage: List of AI assistance usage events
+ annotation_changes: List of annotation modifications
+ navigation_history: List of navigation events to/from this instance
+ focus_time_by_element: Milliseconds spent focused on each element
+ scroll_depth_max: Maximum scroll percentage reached (0-100)
+ keyword_highlights_shown: Keyword highlights displayed (from randomization feature)
+ """
+ instance_id: str
+ session_start: float = field(default_factory=time.time)
+ session_end: Optional[float] = None
+ total_time_ms: int = 0
+ interactions: List[InteractionEvent] = field(default_factory=list)
+ ai_usage: List[AIUsageEvent] = field(default_factory=list)
+ annotation_changes: List[AnnotationChange] = field(default_factory=list)
+ navigation_history: List[Dict[str, Any]] = field(default_factory=list)
+ focus_time_by_element: Dict[str, int] = field(default_factory=dict)
+ scroll_depth_max: float = 0.0
+ keyword_highlights_shown: List[Dict[str, Any]] = field(default_factory=list)
+ chat_history: List[ChatMessage] = field(default_factory=list)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to dictionary for JSON serialization."""
+ return {
+ 'instance_id': self.instance_id,
+ 'session_start': self.session_start,
+ 'session_end': self.session_end,
+ 'total_time_ms': self.total_time_ms,
+ 'interactions': [
+ e.to_dict() if hasattr(e, 'to_dict') else e
+ for e in self.interactions
+ ],
+ 'ai_usage': [
+ e.to_dict() if hasattr(e, 'to_dict') else e
+ for e in self.ai_usage
+ ],
+ 'annotation_changes': [
+ e.to_dict() if hasattr(e, 'to_dict') else e
+ for e in self.annotation_changes
+ ],
+ 'navigation_history': self.navigation_history,
+ 'focus_time_by_element': self.focus_time_by_element,
+ 'scroll_depth_max': self.scroll_depth_max,
+ 'keyword_highlights_shown': self.keyword_highlights_shown,
+ 'chat_history': [
+ e.to_dict() if hasattr(e, 'to_dict') else e
+ for e in self.chat_history
+ ],
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'BehavioralData':
+ """
+ Reconstruct from serialized dictionary.
+
+ Handles both raw dictionaries and properly typed objects.
+ """
+ bd = cls(instance_id=data.get('instance_id', ''))
+ bd.session_start = data.get('session_start', 0)
+ bd.session_end = data.get('session_end')
+ bd.total_time_ms = data.get('total_time_ms', 0)
+
+ # Reconstruct interactions
+ interactions = data.get('interactions', [])
+ bd.interactions = [
+ InteractionEvent.from_dict(e) if isinstance(e, dict) else e
+ for e in interactions
+ ]
+
+ # Reconstruct AI usage events
+ ai_usage = data.get('ai_usage', [])
+ bd.ai_usage = [
+ AIUsageEvent.from_dict(e) if isinstance(e, dict) else e
+ for e in ai_usage
+ ]
+
+ # Reconstruct annotation changes
+ changes = data.get('annotation_changes', [])
+ bd.annotation_changes = [
+ AnnotationChange.from_dict(e) if isinstance(e, dict) else e
+ for e in changes
+ ]
+
+ bd.navigation_history = data.get('navigation_history', [])
+ bd.focus_time_by_element = data.get('focus_time_by_element', {})
+ bd.scroll_depth_max = data.get('scroll_depth_max', 0.0)
+ bd.keyword_highlights_shown = data.get('keyword_highlights_shown', [])
+
+ # Reconstruct chat history
+ chat_history = data.get('chat_history', [])
+ bd.chat_history = [
+ ChatMessage.from_dict(e) if isinstance(e, dict) else e
+ for e in chat_history
+ ]
+
+ return bd
+
+ def add_interaction(self, event_type: str, target: str,
+ client_timestamp: Optional[float] = None,
+ metadata: Optional[Dict[str, Any]] = None) -> None:
+ """Add an interaction event with current timestamp."""
+ self.interactions.append(InteractionEvent(
+ event_type=event_type,
+ timestamp=time.time(),
+ target=target,
+ instance_id=self.instance_id,
+ client_timestamp=client_timestamp,
+ metadata=metadata or {},
+ ))
+
+ def add_ai_request(self, schema_name: str) -> AIUsageEvent:
+ """Record an AI assistance request and return the event for later update."""
+ event = AIUsageEvent(
+ request_timestamp=time.time(),
+ schema_name=schema_name,
+ )
+ self.ai_usage.append(event)
+ return event
+
+ def add_annotation_change(self, schema_name: str, action: str,
+ label_name: Optional[str] = None,
+ old_value: Any = None, new_value: Any = None,
+ source: str = "user") -> None:
+ """Record an annotation change."""
+ self.annotation_changes.append(AnnotationChange(
+ timestamp=time.time(),
+ schema_name=schema_name,
+ label_name=label_name,
+ action=action,
+ old_value=old_value,
+ new_value=new_value,
+ source=source,
+ ))
+
+ def add_navigation(self, action: str, from_instance: Optional[str] = None,
+ to_instance: Optional[str] = None) -> None:
+ """Record a navigation event."""
+ self.navigation_history.append({
+ 'action': action,
+ 'from_instance': from_instance,
+ 'to_instance': to_instance,
+ 'timestamp': time.time(),
+ })
+
+ def update_focus_time(self, element: str, duration_ms: int) -> None:
+ """Add time spent focused on an element."""
+ current = self.focus_time_by_element.get(element, 0)
+ self.focus_time_by_element[element] = current + duration_ms
+
+ def update_scroll_depth(self, depth: float) -> None:
+ """Update maximum scroll depth if new depth is greater."""
+ if depth > self.scroll_depth_max:
+ self.scroll_depth_max = depth
+
+ def finalize_session(self) -> None:
+ """Mark session as ended and calculate total time."""
+ self.session_end = time.time()
+ self.total_time_ms = int((self.session_end - self.session_start) * 1000)
+
+
+def create_behavioral_data(instance_id: str) -> BehavioralData:
+ """Factory function to create new behavioral data for an instance."""
+ return BehavioralData(instance_id=instance_id)
+
+
+def get_or_create_behavioral_data(
+ behavioral_data_dict: Dict[str, Any],
+ instance_id: str
+) -> BehavioralData:
+ """
+ Get existing behavioral data or create new one.
+
+ Args:
+ behavioral_data_dict: Dictionary mapping instance_id to BehavioralData
+ instance_id: The instance to get/create data for
+
+ Returns:
+ BehavioralData object for the instance
+ """
+ if instance_id not in behavioral_data_dict:
+ behavioral_data_dict[instance_id] = create_behavioral_data(instance_id)
+
+ bd = behavioral_data_dict[instance_id]
+
+ # Handle case where dict contains raw dict instead of BehavioralData
+ if isinstance(bd, dict):
+ bd = BehavioralData.from_dict(bd)
+ behavioral_data_dict[instance_id] = bd
+
+ return bd
diff --git a/potato/item_state_management.py b/potato/item_state_management.py
new file mode 100644
index 0000000000000000000000000000000000000000..243ae72c8d88cb6cce44adecfe8b9452e5344e91
--- /dev/null
+++ b/potato/item_state_management.py
@@ -0,0 +1,2409 @@
+"""
+Item State Management Module
+
+This module provides the core data structures and management logic for annotation items
+in the Potato platform. It handles item storage, assignment strategies, and tracking
+of annotation progress across users.
+
+The module includes:
+- Item class: Represents individual annotation items with metadata
+- Label class: Represents annotation labels with schema information
+- SpanAnnotation class: Represents text span annotations with position data
+- AssignmentStrategy enum: Defines different strategies for assigning items to users
+- ItemStateManager: Main class for managing item state and assignments
+
+The system supports multiple assignment strategies including random, fixed order,
+active learning, and diversity-based assignment to optimize annotation efficiency.
+"""
+
+from __future__ import annotations
+
+# Need to import UserState as a type hint for the ItemStateManager
+from typing import TYPE_CHECKING, Dict, Set, List, Optional
+if TYPE_CHECKING:
+ from potato.user_state_management import UserState
+
+from enum import Enum
+from collections import OrderedDict, deque, Counter, defaultdict
+import random
+import uuid
+import logging
+import threading
+import json
+import os
+
+# Singleton instance of the ItemStateManager with thread-safe lock
+ITEM_STATE_MANAGER = None
+_ITEM_STATE_MANAGER_LOCK = threading.Lock()
+
+def init_item_state_manager(config: dict) -> ItemStateManager:
+ """
+ Initialize the singleton ItemStateManager instance.
+
+ This function creates the global ItemStateManager that will be shared
+ across all users. It's designed to be called once during application startup.
+ Thread-safe initialization using double-checked locking pattern.
+
+ Args:
+ config: Configuration dictionary containing item management settings
+
+ Returns:
+ ItemStateManager: The initialized singleton instance
+
+ Note:
+ TODO: make the manager type configurable between in-memory and DB-backed.
+ The DB back-end is for when we have a ton of data and don't want it sitting in
+ memory all the time (or where some external process is going to be adding new items)
+ """
+ global ITEM_STATE_MANAGER
+
+ # Double-checked locking for thread safety
+ if ITEM_STATE_MANAGER is None:
+ with _ITEM_STATE_MANAGER_LOCK:
+ # Check again inside the lock
+ if ITEM_STATE_MANAGER is None:
+ ITEM_STATE_MANAGER = ItemStateManager(config)
+
+ return ITEM_STATE_MANAGER
+
+def clear_item_state_manager():
+ """
+ Clear the singleton item state manager instance (for testing).
+
+ This function is primarily used for testing purposes to reset the
+ global state between test runs. Thread-safe.
+ """
+ global ITEM_STATE_MANAGER
+ with _ITEM_STATE_MANAGER_LOCK:
+ ITEM_STATE_MANAGER = None
+
+def get_item_state_manager() -> ItemStateManager:
+ """
+ Get the singleton ItemStateManager instance.
+
+ Returns:
+ ItemStateManager: The singleton instance
+
+ Raises:
+ ValueError: If the manager has not been initialized
+
+ Note:
+ TODO: make the manager type configurable between in-memory and DB-backed.
+ The DB back-end is for when we have a ton of data and don't want it sitting in
+ memory all the time (or where some external process is going to be adding new items)
+ """
+ global ITEM_STATE_MANAGER
+
+ if ITEM_STATE_MANAGER is None:
+ raise ValueError("Item State Manager has not been initialized yet!")
+
+ return ITEM_STATE_MANAGER
+
+class Item:
+ """
+ A class for maintaining state on items that are being annotated.
+
+ The state of the annotations themselves are stored in the UserState class.
+ The item itself is largely immutable but can be updated with metadata.
+ """
+
+ def __init__(self, item_id, item_data):
+ """
+ Initialize an annotation item.
+
+ Args:
+ item_id: Unique identifier for this item
+ item_data: Dictionary containing the item's data (text, context, etc.)
+ """
+ self.item_id = item_id
+ self.item_data = item_data
+ self.metadata = {}
+
+ # This data structure keeps the label-based annotations the user has
+ # completed so far
+ self.labels = {}
+
+ # This data structure keeps the span-based annotations the user has
+ # completed so far
+ self.span_annotations = {}
+
+ def __getattr__(self, name):
+ """Expose raw data fields as attributes for template access.
+
+ F-029: dynamic-label / video_as_label schemas reference instance data
+ fields in Jinja as ``{{instance_obj.[i]}}`` (e.g.
+ ``instance_obj.gifs[0]``). Without this, only fields whose name happens
+ to collide with a real Item attribute resolve, and others (gifs,
+ gifs_path, โฆ) raise UndefinedError -> 500 at render time.
+
+ __getattr__ is only invoked when normal attribute lookup fails, so real
+ attributes (item_id, item_data, metadata, labels, span_annotations) are
+ never shadowed. Names absent from the data dict still raise
+ AttributeError so Jinja's Undefined handling and copy/pickle behave.
+ """
+ if name.startswith("__") and name.endswith("__"):
+ raise AttributeError(name)
+ data = self.__dict__.get("item_data")
+ if isinstance(data, dict) and name in data:
+ return data[name]
+ raise AttributeError(name)
+
+ def add_metadata(self, metadata_name: str, metadata_value: str):
+ """Add metadata to this item"""
+ self.metadata[metadata_name] = metadata_value
+
+ def get_id(self):
+ """Get the item's unique identifier"""
+ return self.item_id
+
+ def get_data(self):
+ """Get the item's raw data dictionary"""
+ return self.item_data
+
+ def get_text(self):
+ """
+ Get the text content from the item data.
+
+ This method intelligently extracts text from various data structures,
+ trying common keys first, then falling back to string conversion.
+
+ Returns:
+ str: The text content for annotation
+ """
+ if isinstance(self.item_data, dict):
+ # Try to get text from common keys
+ for key in ['text', 'content', 'message', 'title']:
+ if key in self.item_data:
+ return self.item_data[key]
+ # If no text key found, return the first string value
+ for value in self.item_data.values():
+ if isinstance(value, str):
+ return value
+ elif isinstance(self.item_data, str):
+ return self.item_data
+ return str(self.item_data)
+
+ def get_displayed_text(self):
+ """Get the displayed text (same as get_text for now)"""
+ return self.get_text()
+
+ def get_metadata(self, metadata_name: str):
+ """Get metadata value by name"""
+ return self.metadata.get(metadata_name, None)
+
+ def __str__(self):
+ return f"Item(id:{self.item_id}, data:{self.item_data}, metadata:{self.metadata})"
+
+class Label:
+ """
+ A utility class for representing a single label in any annotation scheme.
+
+ Labels may have a integer value (likert), a string value (text), or a boolean value (binary).
+ Span annotations are represented with a different class.
+ """
+ def __init__(self, schema: str, name: str):
+ """
+ Initialize a label.
+
+ Args:
+ schema: The annotation scheme this label belongs to
+ name: The label name/value
+ """
+ self.schema = schema
+ self.name = name
+
+ def get_schema(self):
+ """Get the schema this label belongs to"""
+ return self.schema
+
+ def get_name(self):
+ """Get the label name/value"""
+ return self.name
+
+ def __str__(self):
+ return f"Label(schema:{self.schema}, name:{self.name})"
+
+ def __eq__(self, other):
+ """Check if two labels are equal"""
+ return self.schema == other.schema and self.name == other.name
+
+ def __hash__(self):
+ """Generate hash for label (enables use in sets/dicts)"""
+ return hash((self.schema, self.name))
+
+class SpanAnnotation:
+ """
+ A utility class for representing a single span annotation in any annotation scheme.
+
+ Spans are represented by a start and end index, as well as a label.
+ Optionally includes format-specific coordinates (e.g., PDF page/bbox, spreadsheet row/col).
+
+ Discontinuous spans are supported via the additional_parts parameter, which stores
+ a list of non-contiguous text ranges that are all part of the same annotation.
+ For example, "New" and "York" in "New and exciting York" can be annotated as
+ a single LOCATION entity with additional_parts.
+ """
+ def __init__(self, schema: str, name: str, title: str, start: int, end: int,
+ id: str = None, annotation_id: str = None, target_field: str = None,
+ format_coords: dict = None, additional_parts: list = None,
+ kb_id: str = None, kb_source: str = None, kb_label: str = None):
+ """
+ Initialize a span annotation.
+
+ Args:
+ schema: The annotation scheme this span belongs to
+ name: The span label name
+ title: The span title/description
+ start: Start character index (inclusive)
+ end: End character index (exclusive)
+ id: Optional custom ID for the span
+ annotation_id: Alternative parameter name for ID (for compatibility)
+ target_field: The display field this span targets (for multi-span mode)
+ format_coords: Optional format-specific coordinates (e.g., PDF page/bbox,
+ spreadsheet row/col). Structure depends on source format:
+ - PDF: {"format": "pdf", "page": 1, "bbox": [x0, y0, x1, y1]}
+ - Spreadsheet: {"format": "spreadsheet", "row": 1, "col": 2, "cell_ref": "B1"}
+ - Code: {"format": "code", "line": 10, "column": 5}
+ - Document: {"format": "document", "paragraph_id": "p_0", "local_offset": 0}
+ additional_parts: Optional list of additional text ranges for discontinuous spans.
+ Each part is a dict with {"start": int, "end": int, "text": str}.
+ Used for entities that span non-contiguous text, e.g.,
+ "New" and "York" in "New and exciting York".
+ kb_id: Optional knowledge base entity ID (e.g., "Q937" for Wikidata)
+ kb_source: Optional knowledge base source name (e.g., "wikidata", "umls")
+ kb_label: Optional human-readable label from the knowledge base
+ """
+ self.schema = schema
+ self.start = start
+ self.title = title
+ self.end = end
+ self.name = name
+ self.target_field = target_field # For multi-span support
+ self.format_coords = format_coords # Format-specific coordinates
+ self.additional_parts = additional_parts or [] # For discontinuous spans
+ self.kb_id = kb_id # Knowledge base entity ID
+ self.kb_source = kb_source # Knowledge base source (e.g., "wikidata", "umls")
+ self.kb_label = kb_label # Human-readable KB entity label
+ # Accept both id and annotation_id for compatibility
+ _id = id if id is not None else annotation_id
+ if _id is not None:
+ self._id = _id
+ else:
+ # Generate a unique ID if none provided
+ self._id = f"span_{uuid.uuid4().hex}"
+
+ def get_schema(self):
+ """Get the schema this span belongs to"""
+ return self.schema
+
+ def get_start(self):
+ """Get the start character index"""
+ return self.start
+
+ def get_end(self):
+ """Get the end character index"""
+ return self.end
+
+ def get_name(self):
+ """Get the span label name"""
+ return self.name
+
+ def get_title(self):
+ """Get the span title/description"""
+ return self.title
+
+ def get_id(self):
+ """Get the span's unique identifier"""
+ return self._id
+
+ def get_target_field(self):
+ """Get the target field key (for multi-span mode)"""
+ return self.target_field
+
+ def get_format_coords(self):
+ """Get format-specific coordinates (for document format support)"""
+ return self.format_coords
+
+ def set_format_coords(self, coords: dict):
+ """Set format-specific coordinates"""
+ self.format_coords = coords
+
+ def get_kb_id(self):
+ """Get the knowledge base entity ID"""
+ return self.kb_id
+
+ def set_kb_id(self, kb_id: str):
+ """Set the knowledge base entity ID"""
+ self.kb_id = kb_id
+
+ def get_kb_source(self):
+ """Get the knowledge base source name"""
+ return self.kb_source
+
+ def set_kb_source(self, kb_source: str):
+ """Set the knowledge base source name"""
+ self.kb_source = kb_source
+
+ def get_kb_label(self):
+ """Get the knowledge base entity label"""
+ return self.kb_label
+
+ def set_kb_label(self, kb_label: str):
+ """Set the knowledge base entity label"""
+ self.kb_label = kb_label
+
+ def has_entity_link(self) -> bool:
+ """Check if this span has a knowledge base entity link"""
+ return bool(self.kb_id and self.kb_source)
+
+ def set_entity_link(self, kb_id: str, kb_source: str, kb_label: str = None):
+ """
+ Set the knowledge base entity link for this span.
+
+ Args:
+ kb_id: Knowledge base entity ID (e.g., "Q937")
+ kb_source: Knowledge base source (e.g., "wikidata")
+ kb_label: Optional human-readable label
+ """
+ self.kb_id = kb_id
+ self.kb_source = kb_source
+ self.kb_label = kb_label
+
+ def clear_entity_link(self):
+ """Remove the knowledge base entity link from this span"""
+ self.kb_id = None
+ self.kb_source = None
+ self.kb_label = None
+
+ def get_additional_parts(self):
+ """Get additional parts for discontinuous spans"""
+ return self.additional_parts
+
+ def add_part(self, start: int, end: int, text: str = None):
+ """
+ Add an additional part to this discontinuous span.
+
+ Args:
+ start: Start character index (inclusive)
+ end: End character index (exclusive)
+ text: Optional text content of this part
+ """
+ part = {"start": start, "end": end}
+ if text is not None:
+ part["text"] = text
+ self.additional_parts.append(part)
+ # Keep parts sorted by start position
+ self.additional_parts.sort(key=lambda p: p["start"])
+
+ def remove_part(self, start: int, end: int):
+ """
+ Remove a part from this discontinuous span.
+
+ Args:
+ start: Start character index of the part to remove
+ end: End character index of the part to remove
+ """
+ self.additional_parts = [
+ p for p in self.additional_parts
+ if not (p["start"] == start and p["end"] == end)
+ ]
+
+ def is_discontinuous(self) -> bool:
+ """Check if this span has multiple parts (is discontinuous)"""
+ return len(self.additional_parts) > 0
+
+ def get_all_parts(self) -> list:
+ """
+ Get all parts of this span (primary + additional) as a sorted list.
+
+ Returns:
+ List of dicts with {"start": int, "end": int, "text": str (optional)}
+ """
+ primary = {"start": self.start, "end": self.end}
+ all_parts = [primary] + self.additional_parts
+ return sorted(all_parts, key=lambda p: p["start"])
+
+ def to_dict(self) -> dict:
+ """Convert span annotation to dictionary for serialization."""
+ result = {
+ "schema": self.schema,
+ "name": self.name,
+ "title": self.title,
+ "start": self.start,
+ "end": self.end,
+ "id": self._id,
+ }
+ if self.target_field:
+ result["target_field"] = self.target_field
+ if self.format_coords:
+ result["format_coords"] = self.format_coords
+ if self.additional_parts:
+ result["additional_parts"] = self.additional_parts
+ if self.kb_id:
+ result["kb_id"] = self.kb_id
+ if self.kb_source:
+ result["kb_source"] = self.kb_source
+ if self.kb_label:
+ result["kb_label"] = self.kb_label
+ return result
+
+ def __str__(self):
+ field_str = f", target_field:{self.target_field}" if self.target_field else ""
+ coords_str = f", format_coords:{self.format_coords}" if self.format_coords else ""
+ parts_str = f", additional_parts:{len(self.additional_parts)}" if self.additional_parts else ""
+ kb_str = f", kb:{self.kb_source}:{self.kb_id}" if self.kb_id else ""
+ return f"SpanAnnotation(schema:{self.schema}, name:{self.name}, start:{self.start}, end:{self.end}, id:{self._id}{field_str}{coords_str}{parts_str}{kb_str})"
+
+ def __eq__(self, other):
+ """Check if two span annotations are equal"""
+ if not isinstance(other, SpanAnnotation):
+ return False
+ # Convert additional_parts to comparable format (list of tuples)
+ self_parts = tuple((p["start"], p["end"]) for p in self.additional_parts)
+ other_parts = tuple((p["start"], p["end"]) for p in other.additional_parts)
+ return (
+ self.schema == other.schema
+ and self.name == other.name
+ and self.title == other.title
+ and self.start == other.start
+ and self.end == other.end
+ and self.target_field == other.target_field
+ and self_parts == other_parts
+ # Note: format_coords not included in equality check
+ # as they are derived from position, not essential identity
+ )
+
+ def __hash__(self):
+ """Generate hash for span annotation (enables use in sets/dicts)"""
+ # Include additional_parts in hash for discontinuous spans
+ parts_hash = tuple((p["start"], p["end"]) for p in self.additional_parts)
+ return hash((self.schema, self.name, self.title, self.start, self.end, self.target_field, parts_hash))
+
+
+class SpanLink:
+ """
+ A utility class for representing a link/relationship between spans.
+
+ SpanLinks connect two or more spans to represent relationships like
+ "PERSON works_for ORGANIZATION" or multi-way relationships.
+ """
+ def __init__(self, schema: str, link_type: str, span_ids: List[str],
+ direction: str = "undirected", id: str = None, properties: dict = None):
+ """
+ Initialize a span link.
+
+ Args:
+ schema: The annotation scheme this link belongs to
+ link_type: The type of relationship (e.g., "WORKS_FOR", "KNOWS")
+ span_ids: Ordered list of span IDs that are connected by this link
+ direction: "directed" or "undirected" - for directed links, order matters
+ id: Optional custom ID for the link
+ properties: Optional dictionary of additional properties
+ """
+ self.schema = schema
+ self.link_type = link_type
+ self.span_ids = span_ids # Ordered list for directed links
+ self.direction = direction # "directed", "undirected"
+ self.properties = properties or {}
+ self._id = id if id else f"link_{uuid.uuid4().hex}"
+
+ def get_schema(self) -> str:
+ """Get the schema this link belongs to"""
+ return self.schema
+
+ def get_link_type(self) -> str:
+ """Get the link type/relationship name"""
+ return self.link_type
+
+ def get_span_ids(self) -> List[str]:
+ """Get the ordered list of span IDs connected by this link"""
+ return self.span_ids
+
+ def get_direction(self) -> str:
+ """Get whether this link is directed or undirected"""
+ return self.direction
+
+ def get_id(self) -> str:
+ """Get the link's unique identifier"""
+ return self._id
+
+ def get_properties(self) -> dict:
+ """Get additional properties for this link"""
+ return self.properties
+
+ def is_directed(self) -> bool:
+ """Check if this link is directed"""
+ return self.direction == "directed"
+
+ def to_dict(self) -> dict:
+ """Convert the span link to a dictionary for serialization"""
+ return {
+ "id": self._id,
+ "schema": self.schema,
+ "link_type": self.link_type,
+ "span_ids": self.span_ids,
+ "direction": self.direction,
+ "properties": self.properties
+ }
+
+ @classmethod
+ def from_dict(cls, data: dict) -> 'SpanLink':
+ """Create a SpanLink from a dictionary"""
+ return cls(
+ schema=data["schema"],
+ link_type=data["link_type"],
+ span_ids=data["span_ids"],
+ direction=data.get("direction", "undirected"),
+ id=data.get("id"),
+ properties=data.get("properties", {})
+ )
+
+ def __str__(self):
+ return f"SpanLink(schema:{self.schema}, type:{self.link_type}, spans:{self.span_ids}, direction:{self.direction}, id:{self._id})"
+
+ def __eq__(self, other):
+ """Check if two span links are equal"""
+ return (
+ isinstance(other, SpanLink)
+ and self.schema == other.schema
+ and self.link_type == other.link_type
+ and self.span_ids == other.span_ids
+ and self.direction == other.direction
+ )
+
+ def __hash__(self):
+ """Generate hash for span link (enables use in sets/dicts)"""
+ return hash((self.schema, self.link_type, tuple(self.span_ids), self.direction))
+
+
+class EventAnnotation:
+ """
+ A utility class for representing an N-ary event annotation.
+
+ Events consist of:
+ - A trigger span: The word/phrase indicating the event (e.g., "attacked", "hired")
+ - Argument spans: Entities with typed roles (e.g., attacker, target, weapon)
+
+ This enables information extraction tasks where events have multiple participants
+ with specific semantic roles.
+ """
+ def __init__(self, schema: str, event_type: str, trigger_span_id: str,
+ arguments: List[Dict[str, str]], id: str = None, properties: dict = None):
+ """
+ Initialize an event annotation.
+
+ Args:
+ schema: The annotation scheme this event belongs to
+ event_type: The type of event (e.g., "ATTACK", "HIRE")
+ trigger_span_id: ID of the span that triggers/indicates the event
+ arguments: List of argument dicts, each with:
+ - role: The semantic role (e.g., "attacker", "target")
+ - span_id: ID of the span filling this role
+ id: Optional custom ID for the event
+ properties: Optional dictionary of additional properties
+ """
+ self.schema = schema
+ self.event_type = event_type
+ self.trigger_span_id = trigger_span_id
+ self.arguments = arguments # [{role: "attacker", span_id: "..."}, ...]
+ self.properties = properties or {}
+ self._id = id if id else f"event_{uuid.uuid4().hex}"
+
+ def get_schema(self) -> str:
+ """Get the schema this event belongs to"""
+ return self.schema
+
+ def get_event_type(self) -> str:
+ """Get the event type"""
+ return self.event_type
+
+ def get_trigger_span_id(self) -> str:
+ """Get the trigger span ID"""
+ return self.trigger_span_id
+
+ def get_arguments(self) -> List[Dict[str, str]]:
+ """Get the list of arguments with their roles"""
+ return self.arguments
+
+ def get_argument_by_role(self, role: str) -> Optional[Dict[str, str]]:
+ """Get the argument for a specific role, or None if not found"""
+ for arg in self.arguments:
+ if arg.get('role') == role:
+ return arg
+ return None
+
+ def get_all_span_ids(self) -> List[str]:
+ """Get all span IDs involved in this event (trigger + arguments)"""
+ span_ids = [self.trigger_span_id]
+ for arg in self.arguments:
+ if 'span_id' in arg:
+ span_ids.append(arg['span_id'])
+ return span_ids
+
+ def get_id(self) -> str:
+ """Get the event's unique identifier"""
+ return self._id
+
+ def get_properties(self) -> dict:
+ """Get additional properties for this event"""
+ return self.properties
+
+ def to_dict(self) -> dict:
+ """Convert the event annotation to a dictionary for serialization"""
+ return {
+ "id": self._id,
+ "schema": self.schema,
+ "event_type": self.event_type,
+ "trigger_span_id": self.trigger_span_id,
+ "arguments": self.arguments,
+ "properties": self.properties
+ }
+
+ @classmethod
+ def from_dict(cls, data: dict) -> 'EventAnnotation':
+ """Create an EventAnnotation from a dictionary"""
+ return cls(
+ schema=data["schema"],
+ event_type=data["event_type"],
+ trigger_span_id=data["trigger_span_id"],
+ arguments=data.get("arguments", []),
+ id=data.get("id"),
+ properties=data.get("properties", {})
+ )
+
+ def __str__(self):
+ args_str = ", ".join(f"{a.get('role', '?')}:{a.get('span_id', '?')}" for a in self.arguments)
+ return f"EventAnnotation(schema:{self.schema}, type:{self.event_type}, trigger:{self.trigger_span_id}, args:[{args_str}], id:{self._id})"
+
+ def __eq__(self, other):
+ """Check if two event annotations are equal"""
+ return (
+ isinstance(other, EventAnnotation)
+ and self.schema == other.schema
+ and self.event_type == other.event_type
+ and self.trigger_span_id == other.trigger_span_id
+ and self.arguments == other.arguments
+ )
+
+ def __hash__(self):
+ """Generate hash for event annotation (enables use in sets/dicts)"""
+ # Convert arguments to a hashable tuple representation
+ args_tuple = tuple((a.get('role', ''), a.get('span_id', '')) for a in self.arguments)
+ return hash((self.schema, self.event_type, self.trigger_span_id, args_tuple))
+
+
+class AssignmentStrategy(Enum):
+ """
+ Enumeration of strategies for assigning items to users.
+
+ Different strategies optimize for different goals:
+ - RANDOM: Maximizes diversity and reduces bias
+ - FIXED_ORDER: Ensures consistent ordering across users
+ - ACTIVE_LEARNING: Prioritizes items with high uncertainty
+ - LLM_CONFIDENCE: Uses AI model confidence for prioritization
+ - MAX_DIVERSITY: Prioritizes items with high disagreement
+ - LEAST_ANNOTATED: Prioritizes items with fewest annotations
+ - CATEGORY_BASED: Assigns items matching user's qualified categories
+ - DIVERSITY_CLUSTERING: Samples items round-robin from embedding clusters
+ - BATCH: Restricts assignment to configured annotator/item cohorts
+ - PRIORITY: Serves items by a triage signal (errors/low-score first)
+ """
+ RANDOM = 'random'
+ FIXED_ORDER = 'fixed_order'
+ ACTIVE_LEARNING = 'active_learning'
+ LLM_CONFIDENCE = 'llm_confidence'
+ MAX_DIVERSITY = 'max_diversity'
+ LEAST_ANNOTATED = 'least_annotated'
+ CATEGORY_BASED = 'category_based'
+ DIVERSITY_CLUSTERING = 'diversity_clustering'
+ BATCH = 'batch'
+ PRIORITY = 'priority'
+
+ def fromstr(phase: str) -> AssignmentStrategy:
+ """
+ Convert a string representation to an AssignmentStrategy enum value.
+
+ Args:
+ phase: String representation of the strategy (case-insensitive)
+
+ Returns:
+ AssignmentStrategy: The corresponding enum value
+
+ Raises:
+ ValueError: If the string doesn't match any known strategy
+ """
+ phase = phase.lower()
+ if phase == "random":
+ return AssignmentStrategy.RANDOM
+ elif phase == "fixed_order":
+ return AssignmentStrategy.FIXED_ORDER
+ elif phase == "active_learning":
+ return AssignmentStrategy.ACTIVE_LEARNING
+ elif phase == "llm_confidence":
+ return AssignmentStrategy.LLM_CONFIDENCE
+ elif phase == "max_diversity":
+ return AssignmentStrategy.MAX_DIVERSITY
+ elif phase == "least_annotated":
+ return AssignmentStrategy.LEAST_ANNOTATED
+ elif phase == "category_based":
+ return AssignmentStrategy.CATEGORY_BASED
+ elif phase == "diversity_clustering":
+ return AssignmentStrategy.DIVERSITY_CLUSTERING
+ elif phase == "batch":
+ return AssignmentStrategy.BATCH
+ elif phase == "priority":
+ return AssignmentStrategy.PRIORITY
+ else:
+ raise ValueError(f"Unknown phase: {phase}")
+
+
+class ItemStateManager:
+ """
+ A class for maintaining state on the ordering and metadata of items that are being annotated.
+
+ This class aims to be a singleton that is shared across all users and provides the functionality
+ of determining which item is next to be annotated.
+ The state of the annotations themselves are stored in the UserState class.
+ """
+
+ def __init__(self, config: dict):
+ """
+ Initialize the item state manager.
+
+ Args:
+ config: Configuration dictionary containing item management settings
+ """
+ # Cache the config for later
+ self.config = config
+ self.logger = logging.getLogger(__name__)
+
+ # Thread-safe lock for concurrent access to item data
+ self._lock = threading.RLock()
+
+ # This data structure keeps the ordering of the items that are being annotated
+ # and a mapping from item ID to the Item object
+ self.instance_id_to_instance = OrderedDict()
+
+ self.instance_id_ordering = []
+
+ # Load max annotations per item from config. Prefers the canonical
+ # `num_annotators_per_item` (int or dict.default) and falls back to the
+ # deprecated `max_annotations_per_item` for backwards compatibility.
+ try:
+ from potato.server_utils.config_module import resolve_num_annotators_per_item
+ self.max_annotations_per_item = resolve_num_annotators_per_item(config)
+ except ImportError:
+ self.max_annotations_per_item = config.get('max_annotations_per_item', -1)
+
+ # Adaptive boost: dynamically increase the per-item cap when early
+ # annotators disagree. Parsed from num_annotators_per_item.adaptive.
+ nap = config.get('num_annotators_per_item')
+ adaptive_cfg = (nap.get('adaptive') if isinstance(nap, dict) else None) or {}
+ self.adaptive_boost = {
+ 'enabled': bool(adaptive_cfg.get('enabled', False)),
+ 'threshold': float(adaptive_cfg.get('disagreement_threshold', 0.5)),
+ 'boost_to': int(adaptive_cfg.get('boost_to', 0)),
+ }
+
+ # Minimum coverage floor; forces continued assignment even if some
+ # users have not yet rated an item.
+ self.min_annotations_per_item = None
+ if isinstance(nap, dict) and nap.get('min') is not None:
+ self.min_annotations_per_item = int(nap['min'])
+ elif config.get('min_annotators_per_instance') is not None:
+ self.min_annotations_per_item = int(config['min_annotators_per_instance'])
+
+ # Track which annotators have worked on each item
+ self.instance_annotators = defaultdict(set)
+
+ # Track assignment timestamps for stale reclamation: {instance_id: {username: timestamp}}
+ self.assignment_timestamps = defaultdict(dict)
+
+ # Instance reclamation config
+ reclaim_config = config.get('instance_reclaim', {})
+ self.reclaim_enabled = reclaim_config.get('enabled', False)
+ self.reclaim_timeout_hours = reclaim_config.get('timeout_hours', 24)
+ self.reclaim_config = reclaim_config
+
+ # Queue of remaining instances to be assigned
+ self.remaining_instance_ids = deque()
+
+ # NOTE: We use an extra set to keep track of completed instances to allow for
+ # O(1) tests of whether an item needs to be removed from the remaining list
+ self.completed_instance_ids = set()
+
+ # Initialize item annotation counts for tracking
+ self.item_annotation_counts = defaultdict(int)
+
+ # Signal-based triage: score each item into a priority from the
+ # `triage` config block (errors / thumbs-down / low score first). The
+ # scorer runs in add_item, so it covers both statically loaded data and
+ # traces ingested at runtime. None when triage is disabled.
+ try:
+ from potato.server_utils.triage import build_scorer
+ self.triage_scorer = build_scorer(config)
+ except ImportError:
+ self.triage_scorer = None
+
+ # Load how we want to assign items to users
+ if 'assignment_strategy' in config:
+ strat = config['assignment_strategy']
+ if isinstance(strat, str):
+ self.assignment_strategy = AssignmentStrategy.fromstr(strat)
+ elif isinstance(strat, dict):
+ self.assignment_strategy = AssignmentStrategy.fromstr(strat['name'])
+ else:
+ raise ValueError("Invalid assignment_strategy in config")
+ elif self.triage_scorer is not None:
+ # Triage enabled without an explicit strategy: prioritize the queue
+ # by the triage signal (the whole point of enabling triage).
+ self.assignment_strategy = AssignmentStrategy.PRIORITY
+ else:
+ self.assignment_strategy = AssignmentStrategy.FIXED_ORDER
+
+ # Set up random seed for assignment strategies
+ self.random_seed = config.get('random_seed', 1234)
+ self.random = random.Random(self.random_seed)
+ self.logger.info(f"ItemStateManager initialized with random_seed={self.random_seed}")
+
+ # Category-based assignment support
+ item_properties = config.get('item_properties', {})
+ self.category_key = item_properties.get('category_key', None)
+
+ # Maps category name to set of instance IDs in that category
+ self.category_to_instance_ids: Dict[str, Set[str]] = defaultdict(set)
+
+ # Maps instance ID to its set of categories
+ self.instance_id_to_categories: Dict[str, Set[str]] = {}
+
+ # Instances with no category
+ self.uncategorized_instance_ids: Set[str] = set()
+
+ # Category assignment fallback behavior (loaded from category_assignment config)
+ category_assignment_config = config.get('category_assignment', {})
+ self.category_fallback = category_assignment_config.get('fallback', 'uncategorized')
+
+ # Dynamic expertise mode - uses probabilistic routing based on annotator agreement
+ dynamic_config = category_assignment_config.get('dynamic', {})
+ self.dynamic_expertise_enabled = dynamic_config.get('enabled', False)
+
+ # Batch assignment restricts items to explicit annotator cohorts. Groups
+ # can be configured centrally, or each item can carry an annotator list.
+ self.batch_assignment_config = config.get('batch_assignment') or {}
+ self.batch_assignment_annotator_key = self.batch_assignment_config.get(
+ 'annotator_key',
+ 'assigned_annotators',
+ )
+ self.batch_assignment_id_key = config.get('item_properties', {}).get('id_key', 'id')
+ self.batch_user_to_instance_ids: Dict[str, List[str]] = defaultdict(list)
+ self._load_batch_assignment_groups()
+
+ def _load_batch_assignment_groups(self) -> None:
+ """Index configured batch groups by annotator while preserving item order."""
+ groups = self.batch_assignment_config.get('groups') or []
+ for group in groups:
+ if not isinstance(group, dict):
+ continue
+ users = group.get('annotators', group.get('users', []))
+ instance_ids = group.get(
+ 'instances',
+ group.get('items', group.get('instance_ids', [])),
+ )
+ file_entry = group.get(
+ 'instances_file',
+ group.get('items_file', group.get('instance_ids_file')),
+ )
+ if file_entry:
+ instance_ids = list(instance_ids) if isinstance(instance_ids, list) else []
+ instance_ids.extend(self._load_batch_instance_ids_from_file(file_entry))
+ if isinstance(users, str):
+ users = [users]
+ if isinstance(instance_ids, str):
+ instance_ids = [instance_ids]
+ users = [str(user) for user in users if isinstance(user, str) and user]
+ instance_ids = [
+ str(instance_id)
+ for instance_id in instance_ids
+ if isinstance(instance_id, str) and instance_id
+ ]
+ for user_id in users:
+ for instance_id in instance_ids:
+ if instance_id not in self.batch_user_to_instance_ids[user_id]:
+ self.batch_user_to_instance_ids[user_id].append(instance_id)
+
+ def _resolve_batch_file_path(self, path: str) -> str:
+ """Resolve a batch assignment file path relative to task_dir."""
+ task_dir = self.config.get('task_dir', '.')
+ raw_path = path if os.path.isabs(path) else os.path.join(task_dir, path)
+ from potato.server_utils.config_module import validate_path_security
+
+ return validate_path_security(os.path.normpath(raw_path), task_dir)
+
+ def _instance_id_from_batch_record(self, record, source: str, index: int) -> str:
+ """Extract an instance ID from a supported batch assignment record."""
+ if isinstance(record, str):
+ return record
+ if isinstance(record, dict):
+ if self.batch_assignment_id_key not in record:
+ raise KeyError(
+ f"ID key '{self.batch_assignment_id_key}' not found in "
+ f"batch assignment file {source} at item {index + 1}"
+ )
+ return str(record[self.batch_assignment_id_key])
+ raise ValueError(
+ f"Expected object or string in batch assignment file {source} "
+ f"at item {index + 1}, got {type(record).__name__}"
+ )
+
+ def _load_batch_instance_ids_from_file(self, file_entry) -> List[str]:
+ """
+ Load batch instance IDs from a supported Potato data file format.
+
+ Supports the same local formats as ``data_files``: JSON, JSONL, CSV,
+ TSV, and Parquet. JSON may be a list of item objects, a list of ID
+ strings, or a mapping with ``instances``/``items``/``instance_ids``.
+ """
+ if isinstance(file_entry, dict):
+ raw_path = file_entry.get('path')
+ encoding = file_entry.get('encoding', 'utf-8')
+ else:
+ raw_path = file_entry
+ encoding = 'utf-8'
+
+ if not raw_path or not isinstance(raw_path, str):
+ raise ValueError("batch assignment file entry must define a path")
+
+ path = self._resolve_batch_file_path(raw_path)
+ fmt = path.rsplit('.', 1)[-1].lower()
+ if fmt not in {'csv', 'tsv', 'json', 'jsonl', 'parquet'}:
+ raise ValueError(f"Unsupported batch assignment file format {fmt} for {path}")
+
+ if fmt == 'json':
+ with open(path, 'rt', encoding=encoding) as f:
+ parsed = json.load(f)
+ if isinstance(parsed, dict):
+ for key in ('instances', 'items', 'instance_ids'):
+ if key in parsed:
+ parsed = parsed[key]
+ break
+ if not isinstance(parsed, list):
+ raise ValueError(f"Expected JSON list in batch assignment file {path}")
+ return [
+ self._instance_id_from_batch_record(record, path, idx)
+ for idx, record in enumerate(parsed)
+ ]
+
+ if fmt == 'jsonl':
+ instance_ids = []
+ with open(path, 'rt', encoding=encoding) as f:
+ for line_no, line in enumerate(f):
+ line = line.strip()
+ if not line:
+ continue
+ record = json.loads(line)
+ instance_ids.append(
+ self._instance_id_from_batch_record(record, path, line_no)
+ )
+ return instance_ids
+
+ if fmt == 'parquet':
+ import pyarrow.parquet as pq
+
+ table = pq.read_table(path)
+ df = table.to_pandas()
+ else:
+ import pandas as pd
+
+ sep = ',' if fmt == 'csv' else '\t'
+ df = pd.read_csv(path, sep=sep, encoding=encoding)
+
+ if self.batch_assignment_id_key not in df.columns:
+ raise KeyError(
+ f"ID column '{self.batch_assignment_id_key}' not found in "
+ f"batch assignment file {path}"
+ )
+ return [str(value) for value in df[self.batch_assignment_id_key].tolist()]
+
+ def _item_batch_annotators(self, instance_id: str) -> Set[str]:
+ """Return annotators allowed by an item's batch annotator key."""
+ item = self.instance_id_to_instance.get(instance_id)
+ if item is None:
+ return set()
+ item_data = item.get_data()
+ if not isinstance(item_data, dict):
+ return set()
+ annotators = item_data.get(self.batch_assignment_annotator_key)
+ if annotators is None:
+ return set()
+ if isinstance(annotators, str):
+ annotators = [annotators]
+ if not isinstance(annotators, list):
+ return set()
+ return {value for value in annotators if isinstance(value, str) and value}
+
+ def _batch_candidate_ids_for_user(self, user_id: str) -> List[str]:
+ """Return batch-eligible item IDs for a user in deterministic order."""
+ candidate_ids = []
+ seen = set()
+
+ for instance_id in self.batch_user_to_instance_ids.get(user_id, []):
+ if instance_id not in seen:
+ candidate_ids.append(instance_id)
+ seen.add(instance_id)
+
+ for instance_id in self.remaining_instance_ids:
+ if instance_id in seen:
+ continue
+ if user_id in self._item_batch_annotators(instance_id):
+ candidate_ids.append(instance_id)
+ seen.add(instance_id)
+
+ return candidate_ids
+
+ def has_item(self, instance_id: str) -> bool:
+ """Returns True if the item is in the state manager"""
+ return instance_id in self.instance_id_to_instance
+
+ def add_item(self, instance_id: str, instance_data: dict):
+ """
+ Adds a new instance to be annotated to the state manager (thread-safe).
+
+ Args:
+ instance_id: Unique identifier for the item
+ instance_data: Dictionary containing the item's data
+
+ Raises:
+ ValueError: If an item with the same ID already exists
+ """
+ with self._lock:
+ item = Item(instance_id, instance_data)
+ if instance_id in self.instance_id_to_instance:
+ raise ValueError(f"Duplicate Item ID! Item with ID {instance_id} already exists in the state manager")
+
+ self.instance_id_to_instance[instance_id] = item
+ self.instance_id_ordering.append(instance_id)
+ self.remaining_instance_ids.append(instance_id)
+
+ # Signal-based triage: store the quality-signal priority + reason on
+ # the item so the PRIORITY assignment strategy and the inline badge
+ # can read it. Runs for both static and runtime-ingested items.
+ if self.triage_scorer is not None:
+ try:
+ score = self.triage_scorer.score(instance_data)
+ for k, v in score.to_metadata().items():
+ item.add_metadata(k, v)
+ except Exception as e:
+ self.logger.warning(f"Triage scoring failed for {instance_id}: {e}")
+
+ # Index categories for this item
+ self._index_item_categories(instance_id, instance_data)
+
+ def update_item(self, instance_id: str, instance_data: dict) -> bool:
+ """
+ Update an existing instance's data (thread-safe).
+
+ This method updates the item_data for an existing instance while preserving
+ all existing annotations (labels, span_annotations) and metadata. This is
+ useful for dynamic data loading scenarios where file contents change.
+
+ Args:
+ instance_id: Unique identifier for the item to update
+ instance_data: New data dictionary for the item
+
+ Returns:
+ bool: True if the item was updated, False if the item doesn't exist
+ """
+ with self._lock:
+ if instance_id not in self.instance_id_to_instance:
+ return False
+ item = self.instance_id_to_instance[instance_id]
+ # Update item_data while preserving labels, span_annotations, and metadata
+ item.item_data = instance_data
+ return True
+
+ def add_items(self, instances: dict[str, dict]):
+ """
+ Given a dictionary of instance IDs to instance data, add them to the state manager.
+
+ Args:
+ instances: Dictionary mapping instance IDs to instance data dictionaries
+ """
+ for iid, instance_data in instances.items():
+ self.add_item(iid, instance_data)
+
+ # =========================================================================
+ # Category Indexing Methods
+ # =========================================================================
+
+ def _index_item_categories(self, instance_id: str, instance_data: dict) -> None:
+ """
+ Extract and index categories for an item.
+
+ Categories can be specified as a string or list of strings in the data.
+ If no category_key is configured or the item has no category, it is
+ added to uncategorized_instance_ids.
+
+ Args:
+ instance_id: The ID of the item
+ instance_data: The item's data dictionary
+ """
+ if not self.category_key:
+ # No category key configured, all items are uncategorized
+ self.uncategorized_instance_ids.add(instance_id)
+ self.instance_id_to_categories[instance_id] = set()
+ return
+
+ category_value = instance_data.get(self.category_key)
+
+ if category_value is None:
+ # Item has no category
+ self.uncategorized_instance_ids.add(instance_id)
+ self.instance_id_to_categories[instance_id] = set()
+ return
+
+ # Normalize to list
+ if isinstance(category_value, str):
+ # Treat empty/whitespace-only strings as uncategorized
+ if category_value.strip():
+ categories = [category_value]
+ else:
+ categories = []
+ elif isinstance(category_value, list):
+ categories = [c for c in category_value if isinstance(c, str) and c.strip()]
+ else:
+ self.logger.warning(
+ f"Item {instance_id} has invalid category value type: {type(category_value)}. "
+ f"Expected string or list of strings."
+ )
+ self.uncategorized_instance_ids.add(instance_id)
+ self.instance_id_to_categories[instance_id] = set()
+ return
+
+ if not categories:
+ # Empty category list
+ self.uncategorized_instance_ids.add(instance_id)
+ self.instance_id_to_categories[instance_id] = set()
+ return
+
+ # Index the categories
+ category_set = set(categories)
+ self.instance_id_to_categories[instance_id] = category_set
+
+ for category in category_set:
+ self.category_to_instance_ids[category].add(instance_id)
+
+ def get_instances_by_category(self, category: str) -> Set[str]:
+ """
+ Get all instance IDs that belong to a specific category.
+
+ Args:
+ category: The category name
+
+ Returns:
+ Set of instance IDs in that category
+ """
+ with self._lock:
+ return self.category_to_instance_ids.get(category, set()).copy()
+
+ def get_instances_by_categories(self, categories: Set[str]) -> Set[str]:
+ """
+ Get all instance IDs that belong to any of the specified categories.
+
+ Args:
+ categories: Set of category names
+
+ Returns:
+ Set of instance IDs in any of those categories
+ """
+ with self._lock:
+ result = set()
+ for category in categories:
+ result.update(self.category_to_instance_ids.get(category, set()))
+ return result
+
+ def get_categories_for_instance(self, instance_id: str) -> Set[str]:
+ """
+ Get all categories that an instance belongs to.
+
+ Args:
+ instance_id: The instance ID
+
+ Returns:
+ Set of category names (empty set if uncategorized)
+ """
+ with self._lock:
+ return self.instance_id_to_categories.get(instance_id, set()).copy()
+
+ def get_uncategorized_instances(self) -> Set[str]:
+ """
+ Get all instance IDs that have no category.
+
+ Returns:
+ Set of uncategorized instance IDs
+ """
+ with self._lock:
+ return self.uncategorized_instance_ids.copy()
+
+ def get_all_categories(self) -> Set[str]:
+ """
+ Get all unique category names in the system.
+
+ Returns:
+ Set of all category names
+ """
+ with self._lock:
+ return set(self.category_to_instance_ids.keys())
+
+ def get_category_counts(self) -> Dict[str, int]:
+ """
+ Get the count of instances per category.
+
+ Returns:
+ Dictionary mapping category names to instance counts
+ """
+ with self._lock:
+ return {cat: len(ids) for cat, ids in self.category_to_instance_ids.items()}
+
+ # =========================================================================
+ # Assignment Methods
+ # =========================================================================
+
+ def _get_annotator_cap_for_item(self, instance_id: str) -> int:
+ """
+ Resolve the annotator cap for a single item.
+
+ Looks up the per-item override (set by overlap sampling or adaptive boost
+ via ``Item.metadata['required_annotations']``) and falls back to the
+ global ``max_annotations_per_item`` when no override is set.
+
+ Returns -1 to mean unlimited (matching the legacy convention).
+ """
+ item = self.instance_id_to_instance.get(instance_id)
+ if item is not None:
+ per_item = item.get_metadata('required_annotations')
+ if per_item is not None:
+ try:
+ return int(per_item)
+ except (TypeError, ValueError):
+ pass
+ return self.max_annotations_per_item
+
+ def _item_is_saturated(self, instance_id: str) -> bool:
+ """True if the item has reached its annotator cap (per-item or global)."""
+ cap = self._get_annotator_cap_for_item(instance_id)
+ return cap >= 0 and len(self.instance_annotators[instance_id]) >= cap
+
+ def has_unlabeled_items_for_user(self, user_state: 'UserState') -> bool:
+ """Check whether any items remain for this user to annotate (read-only)."""
+ if self.assignment_strategy == AssignmentStrategy.BATCH:
+ user_id = getattr(user_state, 'user_id', None)
+ if not user_id:
+ return False
+ for iid in self._batch_candidate_ids_for_user(str(user_id)):
+ if iid not in self.remaining_instance_ids:
+ continue
+ if self._item_is_saturated(iid):
+ continue
+ if not user_state.has_annotated(iid):
+ return True
+ return False
+
+ for iid in self.remaining_instance_ids:
+ if self._item_is_saturated(iid):
+ continue
+ if not user_state.has_annotated(iid):
+ return True
+ return False
+
+ def _count_assignable_batch_items(self, user_state: 'UserState') -> int:
+ """Count batch-eligible items that can still be newly assigned."""
+ user_id = getattr(user_state, 'user_id', None)
+ if not user_id:
+ return 0
+ already_assigned = user_state.get_assigned_instance_ids()
+ count = 0
+ for iid in self._batch_candidate_ids_for_user(str(user_id)):
+ if iid not in self.remaining_instance_ids:
+ continue
+ if self._item_is_saturated(iid):
+ continue
+ if iid in already_assigned:
+ continue
+ if user_state.has_annotated(iid):
+ continue
+ count += 1
+ return count
+
+ def assign_instances_to_user(self, user_state: UserState) -> int:
+ """
+ Assigns a set of instances to a user based on the current state of the system
+ and returns the number of instances assigned.
+
+ This method implements various assignment strategies to optimize annotation
+ efficiency and quality. The strategy used depends on the configuration.
+
+ If ICL verification is enabled with mix_with_regular_assignments, this method
+ may include verification tasks from the ICL labeler's queue. These appear as
+ regular annotation tasks (blind labeling) so users don't know they're verifying
+ LLM predictions.
+
+ Args:
+ user_state: The user state object to assign instances to
+
+ Returns:
+ int: Number of instances assigned to the user
+
+ Side Effects:
+ - Updates user_state with new instance assignments
+ - Updates internal tracking of item assignments
+ - May modify remaining_instance_ids queue
+ """
+ self.logger.debug(f"Assigning instances to user {getattr(user_state, 'user_id', None)} with strategy {self.assignment_strategy} and random_seed={self.random_seed}")
+
+ # Snapshot existing assignments so we can record timestamps for new ones
+ existing_assignments = set(user_state.get_assigned_instance_ids()) if self.reclaim_enabled else None
+
+ result = self._assign_instances_to_user_inner(user_state)
+
+ # Record timestamps for newly assigned instances
+ if self.reclaim_enabled and existing_assignments is not None:
+ import time
+ username = getattr(user_state, 'user_id', None)
+ if username:
+ now = time.time()
+ new_assignments = set(user_state.get_assigned_instance_ids()) - existing_assignments
+ for iid in new_assignments:
+ self.assignment_timestamps[iid][username] = now
+
+ return result
+
+ def _assign_instances_to_user_inner(self, user_state: 'UserState') -> int:
+ """Inner assignment logic called by assign_instances_to_user."""
+
+ # Check if we should assign a verification task from ICL labeling
+ verification_assigned = self._maybe_assign_icl_verification(user_state)
+ if verification_assigned:
+ # Return early if we assigned a verification task
+ return verification_assigned
+
+ # Reclaim stale assignments before assigning new ones
+ if self.reclaim_enabled:
+ self._reclaim_stale_assignments()
+
+ # Decline to assign new items to users that have completed the maximum
+ if not user_state.has_remaining_assignments():
+ return 0
+
+ # Determine how many instances to assign
+ current_assignments = user_state.get_assigned_instance_count()
+ max_assignments = user_state.get_max_assignments()
+
+ if max_assignments > 0:
+ remaining_capacity = max_assignments - current_assignments
+ if remaining_capacity <= 0:
+ return 0
+ # For fixed_order strategy, assign all remaining capacity at once
+ # For other strategies, use the original incremental logic
+ if self.assignment_strategy in {AssignmentStrategy.FIXED_ORDER, AssignmentStrategy.BATCH}:
+ instances_to_assign = remaining_capacity
+ else:
+ # If user has less than 3 assignments, assign up to 3 more (or remaining capacity)
+ if current_assignments < 3:
+ instances_to_assign = min(3, remaining_capacity)
+ else:
+ # Otherwise, assign one at a time
+ instances_to_assign = 1
+ else:
+ # Batch assignment should still hand out the whole eligible batch
+ # when per-user quota is unlimited.
+ if self.assignment_strategy == AssignmentStrategy.BATCH:
+ instances_to_assign = self._count_assignable_batch_items(user_state)
+ else:
+ # No maximum, assign one at a time
+ instances_to_assign = 1
+
+ # TODO: add strategy for assigning instances to users:
+ #
+ # 1) Random assignment (up to max per item/user)
+ # 2) Dynamic assignment based on user performance
+ # 3) Dynamic assignment based on item difficulty
+ # 4) Dynamic assignment based on item diversity
+ # 5) Dynamic assignment based on active learning model uncertainty
+ #
+ # NOTE: This method should probably be where we periodically
+ # check for item re-assignment where some items are assigned
+ # for a long time but never get annotated
+ #
+ # FOR NOW, just assign all instances to the user
+ if self.assignment_strategy == AssignmentStrategy.RANDOM:
+ # Random assignment strategy
+ unlabeled_items = []
+ for iid in self.remaining_instance_ids:
+ annotation_count = len(self.instance_annotators[iid])
+ cap = self._get_annotator_cap_for_item(iid)
+ self.logger.debug(f"[ASSIGNMENT] Considering {iid}: annotation_count={annotation_count}, cap={cap}")
+ # Always skip items that have reached max annotations, but do not remove here
+ if cap >= 0 and annotation_count >= cap:
+ self.logger.debug(f"[ASSIGNMENT] Skipping {iid}: reached annotation cap")
+ continue
+ if not user_state.has_annotated(iid):
+ unlabeled_items.append(iid)
+ else:
+ self.logger.debug(f"User {getattr(user_state, 'user_id', None)} already annotated {iid}, skipping.")
+ self.logger.debug(f"Unlabeled items for user: {unlabeled_items}")
+ if not unlabeled_items:
+ self.logger.info(f"No unlabeled items available for user {getattr(user_state, 'user_id', None)}")
+ return 0
+ to_assign = self.random.sample(unlabeled_items, min(instances_to_assign, len(unlabeled_items)))
+ self.logger.debug(f"Randomly assigning items {to_assign} to user {getattr(user_state, 'user_id', None)}")
+ for item_id in to_assign:
+ user_state.assign_instance(self.instance_id_to_instance[item_id])
+ return len(to_assign)
+ elif self.assignment_strategy == AssignmentStrategy.LEAST_ANNOTATED:
+ # Least annotated strategy: prioritize items with fewest annotations
+ candidates = []
+ for iid in list(self.remaining_instance_ids):
+ annotation_count = len(self.instance_annotators[iid])
+ cap = self._get_annotator_cap_for_item(iid)
+ if cap >= 0 and annotation_count >= cap:
+ if iid in self.remaining_instance_ids:
+ self.remaining_instance_ids.remove(iid)
+ continue
+ if iid not in user_state.get_assigned_instance_ids():
+ candidates.append((iid, annotation_count))
+ if not candidates:
+ return 0
+ # Sort by annotation count (fewest first), then by id for determinism
+ candidates.sort(key=lambda x: (x[1], x[0]))
+ assigned = 0
+ for item_id, _ in candidates[:instances_to_assign]:
+ user_state.assign_instance(self.instance_id_to_instance[item_id])
+ assigned += 1
+ return assigned
+ elif self.assignment_strategy == AssignmentStrategy.FIXED_ORDER:
+ # Fixed order assignment strategy
+ assigned = 0
+ for iid in list(self.remaining_instance_ids):
+ if self._item_is_saturated(iid):
+ if iid in self.remaining_instance_ids:
+ self.remaining_instance_ids.remove(iid)
+ continue
+ if iid not in user_state.get_assigned_instance_ids():
+ user_state.assign_instance(self.instance_id_to_instance[iid])
+ assigned += 1
+ if assigned >= instances_to_assign:
+ break
+ return assigned
+ elif self.assignment_strategy == AssignmentStrategy.MAX_DIVERSITY:
+ # Maximum diversity assignment strategy
+ unlabeled_items = []
+ for iid in list(self.remaining_instance_ids):
+ if self._item_is_saturated(iid):
+ if iid in self.remaining_instance_ids:
+ self.remaining_instance_ids.remove(iid)
+ continue
+ if not user_state.has_annotated(iid):
+ unlabeled_items.append(iid)
+ if not unlabeled_items:
+ return 0
+ # Calculate disagreement scores for each item
+ item_disagreement_scores = {}
+ for iid in unlabeled_items:
+ disagreement_score = self._calculate_disagreement_score(iid)
+ item_disagreement_scores[iid] = disagreement_score
+ # Sort by disagreement score (highest first)
+ sorted_items = sorted(item_disagreement_scores.keys(), key=lambda x: item_disagreement_scores[x], reverse=True)
+ assigned = 0
+ for item_id in sorted_items[:instances_to_assign]:
+ user_state.assign_instance(self.instance_id_to_instance[item_id])
+ assigned += 1
+ return assigned
+ elif self.assignment_strategy == AssignmentStrategy.ACTIVE_LEARNING:
+ # Active learning: serve items in the current pool order. The
+ # ActiveLearningManager (when enabled) reorders remaining_instance_ids
+ # by query strategy (uncertainty/BADGE/BALD/hybrid) after each
+ # retrain, so taking items in-order surfaces the most informative
+ # unlabeled instances first. Before any retrain (cold start) this is
+ # simply the configured order.
+ unlabeled_items = []
+ for iid in list(self.remaining_instance_ids):
+ if self._item_is_saturated(iid):
+ if iid in self.remaining_instance_ids:
+ self.remaining_instance_ids.remove(iid)
+ continue
+ if not user_state.has_annotated(iid):
+ unlabeled_items.append(iid)
+ if not unlabeled_items:
+ return 0
+ to_assign = unlabeled_items[:instances_to_assign]
+ self.logger.debug(f"Active learning: assigning items {to_assign} to user {getattr(user_state, 'user_id', None)}")
+ for item_id in to_assign:
+ user_state.assign_instance(self.instance_id_to_instance[item_id])
+ return len(to_assign)
+ elif self.assignment_strategy == AssignmentStrategy.LLM_CONFIDENCE:
+ # LLM confidence assignment strategy (currently falls back to random)
+ unlabeled_items = []
+ for iid in list(self.remaining_instance_ids):
+ if self._item_is_saturated(iid):
+ if iid in self.remaining_instance_ids:
+ self.remaining_instance_ids.remove(iid)
+ continue
+ if not user_state.has_annotated(iid):
+ unlabeled_items.append(iid)
+ if not unlabeled_items:
+ return 0
+ to_assign = self.random.sample(unlabeled_items, min(instances_to_assign, len(unlabeled_items)))
+ self.logger.debug(f"LLM confidence (random fallback): assigning items {to_assign} to user {getattr(user_state, 'user_id', None)}")
+ for item_id in to_assign:
+ user_state.assign_instance(self.instance_id_to_instance[item_id])
+ return len(to_assign)
+ elif self.assignment_strategy == AssignmentStrategy.CATEGORY_BASED:
+ # Category-based assignment strategy
+ user_id = getattr(user_state, 'user_id', None)
+
+ # Check if dynamic expertise mode is enabled
+ if self.dynamic_expertise_enabled:
+ return self._assign_category_based_dynamic(user_state, instances_to_assign)
+
+ # Standard category-based assignment using qualification
+ # Assigns instances from categories the user has qualified for
+ qualified_categories = user_state.get_qualified_categories()
+
+ self.logger.debug(f"Category-based assignment for user {user_id}, qualified categories: {qualified_categories}")
+
+ # Get candidate instances from qualified categories
+ candidate_ids = set()
+ if qualified_categories:
+ for category in qualified_categories:
+ candidate_ids.update(self.category_to_instance_ids.get(category, set()))
+
+ # If no candidates from categories, apply fallback behavior
+ if not candidate_ids:
+ self.logger.debug(f"No category matches for user {user_id}, using fallback: {self.category_fallback}")
+ if self.category_fallback == 'uncategorized':
+ candidate_ids = self.uncategorized_instance_ids.copy()
+ elif self.category_fallback == 'random':
+ candidate_ids = set(self.remaining_instance_ids)
+ # 'none' fallback means no assignment
+
+ # Filter candidates: not already annotated by user, not completed
+ unlabeled_items = []
+ for iid in candidate_ids:
+ # Skip if item is not in remaining (already completed)
+ if iid not in self.remaining_instance_ids:
+ continue
+ # Skip if item has reached max annotations
+ if self._item_is_saturated(iid):
+ continue
+ # Skip if user already annotated this item
+ if not user_state.has_annotated(iid):
+ unlabeled_items.append(iid)
+
+ self.logger.debug(f"Category-based: {len(unlabeled_items)} unlabeled items available for user {user_id}")
+
+ if not unlabeled_items:
+ return 0
+
+ # Randomly sample from eligible items (can be combined with other sub-strategies in future)
+ to_assign = self.random.sample(unlabeled_items, min(instances_to_assign, len(unlabeled_items)))
+ self.logger.debug(f"Category-based: assigning items {to_assign} to user {user_id}")
+
+ for item_id in to_assign:
+ user_state.assign_instance(self.instance_id_to_instance[item_id])
+
+ return len(to_assign)
+ elif self.assignment_strategy == AssignmentStrategy.DIVERSITY_CLUSTERING:
+ # Diversity clustering assignment strategy
+ from potato.diversity_manager import get_diversity_manager
+ dm = get_diversity_manager()
+
+ if dm and dm.enabled:
+ # Get user's annotated items for preservation
+ annotated_ids = set(user_state.get_annotated_instance_ids()) if hasattr(user_state, 'get_annotated_instance_ids') else set()
+
+ # Get available items (respecting per-item / global annotator caps)
+ available_ids = []
+ for iid in self.remaining_instance_ids:
+ # Skip if item has reached annotation limit
+ if self._item_is_saturated(iid):
+ continue
+ # Skip if user already annotated
+ if user_state.has_annotated(iid):
+ continue
+ available_ids.append(iid)
+
+ if not available_ids:
+ return 0
+
+ # Get user_id for diversity manager
+ user_id = getattr(user_state, 'user_id', 'anonymous')
+
+ # Generate diverse ordering with preserved positions
+ diverse_order = dm.apply_to_user_ordering(
+ user_id, available_ids, annotated_ids
+ )
+
+ # Assign items from diverse order
+ assigned = 0
+ for item_id in diverse_order[:instances_to_assign]:
+ user_state.assign_instance(self.instance_id_to_instance[item_id])
+ assigned += 1
+
+ return assigned
+ else:
+ # Fallback to random if diversity manager unavailable
+ self.logger.debug("Diversity manager not available, falling back to random")
+ return self._assign_random_fallback(user_state, instances_to_assign)
+ elif self.assignment_strategy == AssignmentStrategy.BATCH:
+ return self._assign_batch(user_state, instances_to_assign)
+ elif self.assignment_strategy == AssignmentStrategy.PRIORITY:
+ # Signal-based triage: serve items by their stored triage priority
+ # (errors / thumbs-down / low score first). Mirrors MAX_DIVERSITY but
+ # sorts by the ingestion-time signal in item metadata instead of
+ # disagreement. Ties broken by the original queue order for
+ # determinism. `triage.order: asc` flips to lowest-priority-first.
+ candidates = []
+ for iid in list(self.remaining_instance_ids):
+ if self._item_is_saturated(iid):
+ if iid in self.remaining_instance_ids:
+ self.remaining_instance_ids.remove(iid)
+ continue
+ if iid in user_state.get_assigned_instance_ids():
+ continue
+ if user_state.has_annotated(iid):
+ continue
+ candidates.append(iid)
+ if not candidates:
+ return 0
+ order_index = {iid: i for i, iid in enumerate(self.instance_id_ordering)}
+
+ def _priority_of(iid):
+ p = self.instance_id_to_instance[iid].get_metadata("triage_priority")
+ return p if p is not None else 0
+
+ ascending = bool(self.triage_scorer) and self.triage_scorer.order == "asc"
+ # Sort by priority (desc by default), tie-break by original order asc.
+ candidates.sort(key=lambda iid: (
+ _priority_of(iid) if ascending else -_priority_of(iid),
+ order_index.get(iid, 0),
+ ))
+ assigned = 0
+ for iid in candidates[:instances_to_assign]:
+ user_state.assign_instance(self.instance_id_to_instance[iid])
+ assigned += 1
+ return assigned
+ else:
+ # Default fallback to fixed order
+ self.logger.warning(f"Unknown assignment strategy: {self.assignment_strategy}, falling back to fixed order")
+ assigned = 0
+ for iid in list(self.remaining_instance_ids):
+ if self._item_is_saturated(iid):
+ if iid in self.remaining_instance_ids:
+ self.remaining_instance_ids.remove(iid)
+ continue
+ if iid not in user_state.get_assigned_instance_ids():
+ user_state.assign_instance(self.instance_id_to_instance[iid])
+ assigned += 1
+ if assigned >= instances_to_assign:
+ break
+ return assigned
+
+ def _assign_batch(self, user_state: 'UserState', instances_to_assign: int) -> int:
+ """
+ Assign only items explicitly batched with this annotator.
+
+ This strategy supports round-to-round cohort preservation: a round-2
+ item can either be listed in ``batch_assignment.groups`` for a cohort,
+ or carry an annotator list in item data using ``annotator_key``.
+ """
+ user_id = getattr(user_state, 'user_id', None)
+ if not user_id:
+ return 0
+
+ assigned = 0
+ already_assigned = user_state.get_assigned_instance_ids()
+
+ for item_id in self._batch_candidate_ids_for_user(str(user_id)):
+ if item_id not in self.instance_id_to_instance:
+ self.logger.warning(
+ "Batch assignment references unknown item %s for user %s",
+ item_id,
+ user_id,
+ )
+ continue
+ if item_id not in self.remaining_instance_ids:
+ continue
+ if self._item_is_saturated(item_id):
+ if item_id in self.remaining_instance_ids:
+ self.remaining_instance_ids.remove(item_id)
+ continue
+ if item_id in already_assigned:
+ continue
+ if user_state.has_annotated(item_id):
+ continue
+
+ user_state.assign_instance(self.instance_id_to_instance[item_id])
+ already_assigned.add(item_id)
+ assigned += 1
+ if assigned >= instances_to_assign:
+ break
+
+ return assigned
+
+ def _assign_category_based_dynamic(self, user_state: 'UserState', instances_to_assign: int) -> int:
+ """
+ Dynamic category-based assignment using probabilistic routing.
+
+ In dynamic mode, users can receive instances from ALL categories, but are
+ more likely to get instances from categories they have demonstrated expertise
+ in (based on agreement with other annotators).
+
+ Args:
+ user_state: The user state to assign instances to
+ instances_to_assign: Number of instances to assign
+
+ Returns:
+ int: Number of instances actually assigned
+ """
+ from potato.expertise_manager import get_expertise_manager
+
+ user_id = getattr(user_state, 'user_id', None)
+ expertise_manager = get_expertise_manager()
+
+ if not expertise_manager:
+ # Fallback to random if expertise manager not available
+ self.logger.warning("ExpertiseManager not available, falling back to random assignment")
+ return self._assign_random_fallback(user_state, instances_to_assign)
+
+ assigned_count = 0
+
+ for _ in range(instances_to_assign):
+ # Find categories with available (unlabeled) instances for this user
+ available_categories = set()
+ category_to_eligible_items: Dict[str, List[str]] = {}
+
+ for category, instance_ids in self.category_to_instance_ids.items():
+ eligible_items = []
+ for iid in instance_ids:
+ # Skip if item is not in remaining (already completed)
+ if iid not in self.remaining_instance_ids:
+ continue
+ # Skip if item has reached max annotations
+ if self._item_is_saturated(iid):
+ continue
+ # Skip if user already annotated this item
+ if not user_state.has_annotated(iid):
+ eligible_items.append(iid)
+
+ if eligible_items:
+ available_categories.add(category)
+ category_to_eligible_items[category] = eligible_items
+
+ # Also check uncategorized instances
+ uncategorized_eligible = []
+ for iid in self.uncategorized_instance_ids:
+ if iid not in self.remaining_instance_ids:
+ continue
+ if self._item_is_saturated(iid):
+ continue
+ if not user_state.has_annotated(iid):
+ uncategorized_eligible.append(iid)
+
+ if not available_categories and not uncategorized_eligible:
+ # No more instances available
+ break
+
+ # Use ExpertiseManager to probabilistically select a category
+ if available_categories:
+ selected_category = expertise_manager.select_category_probabilistically(
+ user_id,
+ available_categories,
+ random_instance=self.random
+ )
+ else:
+ selected_category = None
+
+ # Get an instance from the selected category (or uncategorized)
+ if selected_category and selected_category in category_to_eligible_items:
+ eligible_items = category_to_eligible_items[selected_category]
+ selected_item = self.random.choice(eligible_items)
+ elif uncategorized_eligible:
+ selected_item = self.random.choice(uncategorized_eligible)
+ else:
+ break
+
+ # Assign the selected instance
+ user_state.assign_instance(self.instance_id_to_instance[selected_item])
+ assigned_count += 1
+
+ self.logger.debug(
+ f"Dynamic category assignment: assigned {selected_item} "
+ f"(category={selected_category}) to user {user_id}"
+ )
+
+ return assigned_count
+
+ def _assign_random_fallback(self, user_state: 'UserState', instances_to_assign: int) -> int:
+ """Fallback to random assignment when expertise manager is not available."""
+ unlabeled_items = []
+ for iid in self.remaining_instance_ids:
+ if self._item_is_saturated(iid):
+ continue
+ if not user_state.has_annotated(iid):
+ unlabeled_items.append(iid)
+
+ if not unlabeled_items:
+ return 0
+
+ to_assign = self.random.sample(unlabeled_items, min(instances_to_assign, len(unlabeled_items)))
+ for item_id in to_assign:
+ user_state.assign_instance(self.instance_id_to_instance[item_id])
+
+ return len(to_assign)
+
+ def _calculate_disagreement_score(self, instance_id: str) -> float:
+ """
+ Calculate a disagreement score in [0, 1] for an instance.
+
+ Walks every user who has annotated this instance and, per schema,
+ computes the ratio of distinct labels (or distinct span-set
+ signatures) to the number of annotators. The result is the maximum
+ ratio across schemas โ items with at least one disagreeing schema
+ get a high score.
+
+ Returns 0.0 when fewer than two annotators have rated the item or
+ when no UserStateManager is available (e.g., during tests that
+ exercise the item manager in isolation).
+ """
+ try:
+ from potato.user_state_management import get_user_state_manager
+ usm = get_user_state_manager()
+ except (ImportError, ValueError):
+ return 0.0
+ if usm is None:
+ return 0.0
+
+ annotators = list(self.instance_annotators.get(instance_id, ()))
+ if len(annotators) < 2:
+ return 0.0
+
+ # Aggregate per-schema labels + spans
+ schema_to_user_value: Dict[str, Dict[str, tuple]] = defaultdict(dict)
+ for uid in annotators:
+ ustate = usm.get_user_state(uid) if hasattr(usm, "get_user_state") else None
+ if ustate is None:
+ continue
+ for schema, labels in (ustate.get_label_annotations(instance_id) or {}).items():
+ # Represent as a tuple of label names so dict-style equality works.
+ value = tuple(sorted(
+ str(getattr(l, "name", l.get("name") if isinstance(l, dict) else l))
+ for l in labels
+ ))
+ schema_to_user_value[schema][uid] = value
+ for schema, spans in (ustate.get_span_annotations(instance_id) or {}).items():
+ value = tuple(sorted(
+ (int(getattr(s, "start", s["start"] if isinstance(s, dict) else 0)),
+ int(getattr(s, "end", s["end"] if isinstance(s, dict) else 0)),
+ str(getattr(s, "name", s.get("name") if isinstance(s, dict) else "")))
+ for s in spans
+ ))
+ schema_to_user_value[f"__span__{schema}"][uid] = value
+
+ if not schema_to_user_value:
+ return 0.0
+
+ scores = []
+ for values_by_user in schema_to_user_value.values():
+ if len(values_by_user) < 2:
+ continue
+ distinct = len(set(values_by_user.values()))
+ scores.append((distinct - 1) / max(1, len(values_by_user) - 1))
+ return max(scores) if scores else 0.0
+
+ def _reclaim_stale_assignments(self):
+ """
+ Reclaim instances from users who were assigned them but never annotated
+ within the configured timeout period. Reclaimed instances are returned
+ to the remaining_instance_ids pool.
+ """
+ import time
+ if not self.reclaim_enabled:
+ return
+
+ cutoff = time.time() - (self.reclaim_timeout_hours * 3600)
+ usm = None
+ reclaimed_count = 0
+
+ for iid in list(self.assignment_timestamps.keys()):
+ for username in list(self.assignment_timestamps[iid].keys()):
+ timestamp = self.assignment_timestamps[iid][username]
+ if timestamp > cutoff:
+ continue # Not stale yet
+
+ # Check if the user has actually annotated this instance
+ if usm is None:
+ from potato.user_state_management import get_user_state_manager
+ usm = get_user_state_manager()
+
+ user_state = usm.get_user_state(username)
+ if user_state and user_state.has_annotated(iid):
+ # User annotated it โ remove from tracking, not stale
+ del self.assignment_timestamps[iid][username]
+ continue
+
+ # Stale: reclaim the instance
+ self.logger.info(f"Reclaiming stale instance {iid} from user {username} "
+ f"(assigned {(time.time() - timestamp)/3600:.1f} hours ago)")
+
+ reclaimed = False
+ if user_state:
+ reclaimed = self._reclaim_unannotated_assignment(
+ user_state,
+ iid,
+ reason="stale_assignment",
+ )
+ else:
+ if iid not in self.completed_instance_ids and iid not in self.remaining_instance_ids:
+ self.remaining_instance_ids.append(iid)
+ self.instance_annotators[iid].discard(username)
+ if iid in self.assignment_timestamps:
+ self.assignment_timestamps[iid].pop(username, None)
+ reclaimed = True
+
+ if reclaimed:
+ reclaimed_count += 1
+
+ # Clean up empty entries
+ if iid in self.assignment_timestamps and not self.assignment_timestamps[iid]:
+ del self.assignment_timestamps[iid]
+
+ if reclaimed_count > 0:
+ self.logger.info(f"Reclaimed {reclaimed_count} stale instance assignments")
+
+ def _reclaim_unannotated_assignment(
+ self,
+ user_state: 'UserState',
+ instance_id: str,
+ reason: str = "assignment_reclaim",
+ ) -> bool:
+ """Reclaim one unannotated assignment from a user."""
+ user_id = getattr(user_state, 'user_id', None) or user_state.get_user_id()
+
+ if user_state.has_annotated(instance_id):
+ return False
+
+ unassigned = user_state.unassign_instance(instance_id)
+ if not unassigned:
+ return False
+
+ if instance_id not in self.completed_instance_ids and instance_id not in self.remaining_instance_ids:
+ self.remaining_instance_ids.append(instance_id)
+
+ self.instance_annotators[instance_id].discard(user_id)
+ if instance_id in self.assignment_timestamps:
+ self.assignment_timestamps[instance_id].pop(user_id, None)
+ if not self.assignment_timestamps[instance_id]:
+ del self.assignment_timestamps[instance_id]
+
+ self.logger.info(
+ "Reclaimed unannotated assignment %s from user %s (%s)",
+ instance_id,
+ user_id,
+ reason,
+ )
+ return True
+
+ def _clear_completed_assignment(
+ self,
+ user_state: 'UserState',
+ instance_id: str,
+ reason: str = "assignment_reclaim",
+ ) -> bool:
+ """Clear a completed annotation and release that assignment."""
+ user_id = getattr(user_state, 'user_id', None) or user_state.get_user_id()
+
+ if not user_state.has_annotated(instance_id):
+ return self._reclaim_unannotated_assignment(user_state, instance_id, reason)
+
+ user_state.clear_instance_annotations(instance_id)
+
+ unassigned = user_state.unassign_instance(instance_id)
+ if not unassigned:
+ return False
+
+ had_annotator_credit = user_id in self.instance_annotators[instance_id]
+ self.instance_annotators[instance_id].discard(user_id)
+ if had_annotator_credit and self.item_annotation_counts[instance_id] > 0:
+ self.item_annotation_counts[instance_id] -= 1
+
+ if instance_id in self.completed_instance_ids:
+ if not self._item_is_saturated(instance_id):
+ self.completed_instance_ids.discard(instance_id)
+
+ if instance_id not in self.completed_instance_ids and instance_id not in self.remaining_instance_ids:
+ self.remaining_instance_ids.append(instance_id)
+
+ if instance_id in self.assignment_timestamps:
+ self.assignment_timestamps[instance_id].pop(user_id, None)
+ if not self.assignment_timestamps[instance_id]:
+ del self.assignment_timestamps[instance_id]
+
+ self.logger.info(
+ "Cleared and reclaimed completed assignment %s from user %s (%s)",
+ instance_id,
+ user_id,
+ reason,
+ )
+ return True
+
+ def should_preserve_completed_annotations(self, reason: str = "assignment_reclaim") -> bool:
+ """Return whether completed annotations should survive a reclaim reason."""
+ default = self.reclaim_config.get('preserve_completed_annotations', True)
+
+ prolific_statuses = {
+ "prolific_returned": "RETURNED",
+ "prolific_timed_out": "TIMED-OUT",
+ "prolific_rejected": "REJECTED",
+ }
+ if reason in prolific_statuses:
+ prolific_config = self.reclaim_config.get("prolific", {})
+ if not isinstance(prolific_config, dict):
+ return default
+
+ prolific_default = prolific_config.get('preserve_completed_annotations', default)
+ status_policies = prolific_config.get("status_policies", {})
+ if not isinstance(status_policies, dict):
+ status_policies = {}
+ status_policy = status_policies.get(prolific_statuses[reason], {})
+ if isinstance(status_policy, dict):
+ return status_policy.get('preserve_completed_annotations', prolific_default)
+ return prolific_default
+
+ reason_to_section = {
+ "quality_control_block": ("quality_control",),
+ "stale_assignment": ("stale",),
+ "manual_admin_reclaim": ("manual",),
+ "prolific_dropped": ("prolific",),
+ }
+
+ path = reason_to_section.get(reason)
+ if not path:
+ return default
+
+ value = self.reclaim_config
+ for key in path:
+ if not isinstance(value, dict) or key not in value:
+ return default
+ value = value[key]
+
+ if isinstance(value, dict):
+ return value.get('preserve_completed_annotations', default)
+ return default
+
+ def reclaim_unannotated_assignments_for_user(
+ self,
+ user_state: 'UserState',
+ reason: str = "assignment_reclaim",
+ preserve_completed_annotations: Optional[bool] = None,
+ ) -> List[str]:
+ """Reclaim assignments for a user according to the retention policy."""
+ if preserve_completed_annotations is None:
+ preserve_completed_annotations = self.should_preserve_completed_annotations(reason)
+
+ reclaimed = []
+ with self._lock:
+ for instance_id in list(user_state.get_assigned_instance_ids()):
+ if preserve_completed_annotations:
+ did_reclaim = self._reclaim_unannotated_assignment(user_state, instance_id, reason)
+ else:
+ did_reclaim = self._clear_completed_assignment(user_state, instance_id, reason)
+
+ if did_reclaim:
+ reclaimed.append(instance_id)
+ return reclaimed
+
+ def reclaim_unannotated_assignments_for_users(
+ self,
+ user_ids: List[str],
+ reason: str = "assignment_reclaim",
+ preserve_completed_annotations: Optional[bool] = None,
+ ) -> Dict[str, List[str]]:
+ """Reclaim assignments for several users according to the retention policy."""
+ from potato.user_state_management import get_user_state_manager
+
+ usm = get_user_state_manager()
+ reclaimed_by_user = {}
+ for user_id in user_ids:
+ user_state = usm.get_user_state(user_id)
+ if not user_state:
+ continue
+
+ reclaimed = self.reclaim_unannotated_assignments_for_user(
+ user_state,
+ reason=reason,
+ preserve_completed_annotations=preserve_completed_annotations,
+ )
+ if not reclaimed:
+ continue
+
+ reclaimed_by_user[user_id] = reclaimed
+ try:
+ usm.save_user_state(user_state)
+ except Exception as e:
+ self.logger.warning(
+ "Could not persist reclaimed assignments for user %s: %s",
+ user_id,
+ e,
+ )
+
+ return reclaimed_by_user
+
+ def _maybe_assign_icl_verification(self, user_state: 'UserState') -> int:
+ """
+ Maybe assign an ICL verification task to the user.
+
+ This implements "blind labeling" - the user receives an instance that was
+ already labeled by the LLM, but they don't know it's a verification task.
+ After they annotate it, we compare their label to the LLM's prediction.
+
+ Args:
+ user_state: The user state to potentially assign to
+
+ Returns:
+ int: Number of verification instances assigned (0 or 1)
+ """
+ # Check if ICL labeling is enabled
+ icl_config = self.config.get('icl_labeling', {})
+ if not icl_config.get('enabled', False):
+ return 0
+
+ # Check if verification is enabled with mixed assignments
+ verification_config = icl_config.get('verification', {})
+ if not verification_config.get('enabled', True):
+ return 0
+ if not verification_config.get('mix_with_regular_assignments', True):
+ return 0
+
+ # Probabilistic check - only assign verification ~20% of the time
+ # This ensures users still get regular tasks most of the time
+ verification_mix_rate = verification_config.get('assignment_mix_rate', 0.2)
+ if self.random.random() > verification_mix_rate:
+ return 0
+
+ try:
+ from potato.ai.icl_labeler import get_icl_labeler
+ icl_labeler = get_icl_labeler()
+ if icl_labeler is None:
+ return 0
+
+ # Get pending verifications that this user hasn't already annotated
+ pending = icl_labeler.get_pending_verifications(count=5)
+ user_id = getattr(user_state, 'user_id', None)
+
+ for instance_id, schema_name in pending:
+ # Skip if user already annotated this instance
+ if user_state.has_annotated(instance_id):
+ continue
+
+ # Skip if instance is already assigned to user
+ if instance_id in user_state.get_assigned_instance_ids():
+ continue
+
+ # Check if instance exists in our manager
+ if instance_id not in self.instance_id_to_instance:
+ continue
+
+ # Assign the verification instance
+ item = self.instance_id_to_instance[instance_id]
+ user_state.assign_instance(item)
+
+ # Mark this as a verification task in the user's metadata
+ # This is stored privately so we can record verification after annotation
+ user_state.mark_instance_as_verification(instance_id, schema_name)
+
+ self.logger.info(
+ f"Assigned ICL verification task {instance_id} to user {user_id}"
+ )
+ return 1
+
+ except ImportError:
+ # ICL labeler module not available
+ pass
+ except Exception as e:
+ self.logger.warning(f"Error assigning ICL verification: {e}")
+
+ return 0
+
+ def generate_id_order_mapping(self):
+ """Generate a mapping from instance IDs to their order"""
+ return {iid: idx for idx, iid in enumerate(self.instance_id_ordering)}
+
+ def get_next_instance_id(self, user_state: UserState) -> str:
+ """
+ Get the next instance ID for a user based on the assignment strategy.
+
+ Args:
+ user_state: The user state to get the next instance for
+
+ Returns:
+ str: The next instance ID, or None if no more instances available
+ """
+ # This method would implement the logic to determine which instance
+ # should be next for a given user based on the assignment strategy
+ # For now, it's a placeholder
+ return None
+
+ def get_instance_ids(self) -> list[str]:
+ """Get all instance IDs in the manager"""
+ return list(self.instance_id_to_instance.keys())
+
+ def get_item(self, instance_id: str) -> Item:
+ """Get an item by its ID"""
+ return self.instance_id_to_instance[instance_id]
+
+ def get_annotators_for_item(self, instance_id: str) -> set[str]:
+ """Get the set of annotators who have worked on this item"""
+ return self.instance_annotators[instance_id]
+
+ def get_total_assignable_items_for_user(self, user_state: UserState) -> int:
+ """
+ Get the total number of items that can be assigned to a user.
+
+ This takes into account:
+ - Items the user hasn't already annotated
+ - Items that haven't reached their annotation limit
+ - Items that are still available for assignment
+
+ Args:
+ user_state: The user state to check assignments for
+
+ Returns:
+ int: Number of items that can be assigned
+ """
+ count = 0
+ for iid in self.remaining_instance_ids:
+ # Check if item has reached annotation limit (per-item override or global)
+ if self._item_is_saturated(iid):
+ continue
+ # Check if user has already annotated this item
+ if user_state.has_annotated(iid):
+ continue
+ count += 1
+ return count
+
+ def items(self) -> list[Item]:
+ """Get all items in the manager"""
+ return list(self.instance_id_to_instance.values())
+
+ def register_annotator(self, instance_id: str, user_id: str):
+ """
+ Register that a user has annotated an instance.
+
+ This method updates the tracking of which users have worked on which
+ items, and may trigger cleanup of completed items.
+
+ Args:
+ instance_id: The ID of the instance that was annotated
+ user_id: The ID of the user who did the annotation
+
+ Side Effects:
+ - Updates instance_annotators tracking
+ - May remove items from remaining_instance_ids if they reach limits
+ - Updates item_annotation_counts
+ """
+ # Add user to the set of annotators for this item
+ self.instance_annotators[instance_id].add(user_id)
+
+ # Update annotation count
+ self.item_annotation_counts[instance_id] += 1
+
+ # Adaptive boost: when an item has begun gathering annotations under
+ # the default cap and shows real disagreement, raise its per-item cap
+ # so we keep gathering votes. The boost is one-shot per item.
+ if self.adaptive_boost.get('enabled') and self.adaptive_boost.get('boost_to', 0) >= 2:
+ current_count = len(self.instance_annotators[instance_id])
+ current_cap = self._get_annotator_cap_for_item(instance_id)
+ boost_to = self.adaptive_boost['boost_to']
+ if 2 <= current_count and current_cap >= 0 and current_cap < boost_to:
+ threshold = self.adaptive_boost.get('threshold', 0.5)
+ if self._calculate_disagreement_score(instance_id) >= threshold:
+ item = self.instance_id_to_instance.get(instance_id)
+ if item is not None:
+ item.add_metadata('required_annotations', boost_to)
+ self.logger.info(
+ "Adaptive boost: %s raised to %d annotators (current=%d)",
+ instance_id, boost_to, current_count,
+ )
+ if instance_id in self.completed_instance_ids:
+ self.completed_instance_ids.discard(instance_id)
+ if instance_id not in self.remaining_instance_ids:
+ self.remaining_instance_ids.append(instance_id)
+
+ # Check if this item has reached its annotation limit
+ if self._item_is_saturated(instance_id):
+ # Remove from remaining instances if it's there
+ if instance_id in self.remaining_instance_ids:
+ self.remaining_instance_ids.remove(instance_id)
+ # Mark as completed
+ self.completed_instance_ids.add(instance_id)
+ # If this was an overlap-sample item (cap >= 2), check whether
+ # annotators disagreed enough to route the item to adjudication.
+ cap = self._get_annotator_cap_for_item(instance_id)
+ if cap >= 2:
+ try:
+ from potato.adjudication import get_adjudication_manager
+ adj_mgr = get_adjudication_manager()
+ if adj_mgr is not None:
+ adj_mgr.try_enqueue_item(instance_id)
+ except Exception as exc:
+ self.logger.debug("Adjudication auto-route skipped: %s", exc)
+
+ def update_annotation_count(self, instance_id: str, delta=1):
+ """
+ Update the annotation count for an instance.
+
+ Args:
+ instance_id: The ID of the instance to update
+ delta: The change in annotation count (default: +1)
+ """
+ self.item_annotation_counts[instance_id] += delta
+
+ def reorder_instances(self, new_order: List[str]):
+ """
+ Reorder instances based on active learning predictions.
+
+ Args:
+ new_order: List of instance IDs in the new desired order
+
+ Note:
+ This method preserves instances that are not in the new_order list
+ by appending them to the end of the ordering.
+ """
+ # Create a set of instances in the new order for efficient lookup
+ new_order_set = set(new_order)
+
+ # Filter out instances that don't exist in our manager
+ valid_new_order = [instance_id for instance_id in new_order if instance_id in self.instance_id_to_instance]
+
+ # Find instances that are not in the new order
+ remaining_instances = [instance_id for instance_id in self.instance_id_ordering if instance_id not in new_order_set]
+
+ # Combine the new order with remaining instances
+ self.instance_id_ordering = valid_new_order + remaining_instances
+
+ # Update the remaining_instance_ids queue to match the new ordering
+ self.remaining_instance_ids.clear()
+ for instance_id in self.instance_id_ordering:
+ if instance_id not in self.completed_instance_ids:
+ self.remaining_instance_ids.append(instance_id)
+
+ self.logger.info(f"Reordered {len(valid_new_order)} instances, {len(remaining_instances)} instances preserved")
+
+ def clear(self):
+ """Clear all data from the manager (for testing)"""
+ self.instance_id_to_instance.clear()
+ self.instance_id_ordering.clear()
+ self.remaining_instance_ids.clear()
+ self.completed_instance_ids.clear()
+ self.instance_annotators.clear()
+ self.item_annotation_counts.clear()
diff --git a/potato/judge_calibration/__init__.py b/potato/judge_calibration/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d322c603a011d8b81f1811270545d8e430e5fec2
--- /dev/null
+++ b/potato/judge_calibration/__init__.py
@@ -0,0 +1,44 @@
+"""
+Judge Calibration mode.
+
+A lightweight LLM-as-judge auto-labeling + blind human calibration workflow.
+Researchers configure a judge prompt and a set of LLMs, each sampled k times to
+label (up to a cap of) the data; human(s) then blind-label a sample and Potato
+produces a calibration report (accuracy vs human gold, IAA human<->LLM and
+LLM<->LLM, ECE/reliability, per-LLM confusion matrices) plus a file of every
+LLM's labels.
+
+This is a deliberately simpler cousin of ``solo_mode`` (no refinement loops,
+edge-case synthesis or disagreement-resolution UI) and is distinct from the
+single-judge ``judge_alignment`` feature (which shows suggestions inline and
+draws few-shot from gold labels).
+"""
+
+from .config import JudgeCalibrationConfig, parse_judge_calibration_config
+from .phase import JCPhase, JCPhaseController
+from .aggregation import ModelItemResult, aggregate
+from .storage import ResultStore
+from .generation import LLMGenerationThread, build_prompt, parse_sample
+from .manager import (
+ JudgeCalibrationManager,
+ init_judge_calibration_manager,
+ get_judge_calibration_manager,
+ clear_judge_calibration_manager,
+)
+
+__all__ = [
+ "JudgeCalibrationConfig",
+ "parse_judge_calibration_config",
+ "JCPhase",
+ "JCPhaseController",
+ "ModelItemResult",
+ "aggregate",
+ "ResultStore",
+ "LLMGenerationThread",
+ "build_prompt",
+ "parse_sample",
+ "JudgeCalibrationManager",
+ "init_judge_calibration_manager",
+ "get_judge_calibration_manager",
+ "clear_judge_calibration_manager",
+]
diff --git a/potato/judge_calibration/aggregation.py b/potato/judge_calibration/aggregation.py
new file mode 100644
index 0000000000000000000000000000000000000000..df8586b323d327b0fe553a11691a46ae9c2c8716
--- /dev/null
+++ b/potato/judge_calibration/aggregation.py
@@ -0,0 +1,211 @@
+"""
+Aggregation of k judge samples into a single prediction + confidence.
+
+For each (model, item, schema) we draw ``k`` samples from the LLM. This module
+reduces those raw samples to a modal prediction and an empirical confidence
+(the vote fraction = share of the k samples agreeing with the modal answer).
+
+Confidence is intentionally computed over *all* k draws including failures:
+a ``None`` sample (parse error / invalid label) counts toward the denominator
+so the confidence honestly reflects how often the judge produced a usable,
+consistent answer.
+
+Per-schema reducers:
+- radio / likert : modal label; confidence = count(modal) / k
+- multiselect : per-label vote fraction; predicted set = labels with
+ fraction >= ``multiselect_threshold``; confidence = mean
+ fraction over the predicted set (or 1 - mean fraction over
+ rejected labels when the set is empty)
+- span : delegated to span aggregation (Phase 7); not handled here.
+"""
+
+from collections import Counter
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+
+def span_iou(a: Tuple[int, int], b: Tuple[int, int]) -> float:
+ """Character-offset IoU of two [start, end) spans. 0 if disjoint."""
+ s1, e1 = a
+ s2, e2 = b
+ inter = max(0, min(e1, e2) - max(s1, s2))
+ if inter == 0:
+ return 0.0
+ union = (e1 - s1) + (e2 - s2) - inter
+ return inter / union if union > 0 else 0.0
+
+
+@dataclass
+class ModelItemResult:
+ """A single judge model's aggregated verdict for one (item, schema)."""
+ model: str
+ instance_id: str
+ schema_name: str
+ annotation_type: str
+ modal_label: Any # str | int | list[str] | None
+ confidence: float # 0.0 - 1.0 (vote fraction)
+ k: int
+ samples: List[Any] = field(default_factory=list) # raw per-draw values (None = failed)
+ # Per-label vote fractions for multiselect (label -> fraction); empty otherwise.
+ per_label_confidence: Dict[str, float] = field(default_factory=dict)
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ "model": self.model,
+ "instance_id": self.instance_id,
+ "schema_name": self.schema_name,
+ "annotation_type": self.annotation_type,
+ "modal_label": self.modal_label,
+ "confidence": self.confidence,
+ "k": self.k,
+ "samples": self.samples,
+ "per_label_confidence": self.per_label_confidence,
+ }
+
+ @classmethod
+ def from_dict(cls, d: Dict[str, Any]) -> "ModelItemResult":
+ return cls(
+ model=d["model"],
+ instance_id=d["instance_id"],
+ schema_name=d["schema_name"],
+ annotation_type=d.get("annotation_type", "radio"),
+ modal_label=d.get("modal_label"),
+ confidence=float(d.get("confidence", 0.0)),
+ k=int(d.get("k", 0)),
+ samples=d.get("samples", []),
+ per_label_confidence=d.get("per_label_confidence", {}),
+ )
+
+
+def _aggregate_categorical(samples: List[Any], k: int):
+ """Modal label + vote fraction for single-label schemas (radio/likert).
+
+ None samples count toward k (the denominator) but are never the modal
+ label unless every draw failed.
+ """
+ valid = [s for s in samples if s is not None]
+ if not valid:
+ return None, 0.0
+ counts = Counter(str(s) for s in valid)
+ modal_str, modal_count = counts.most_common(1)[0]
+ # Recover the original (typed) value for the modal label.
+ modal_value = next(s for s in valid if str(s) == modal_str)
+ confidence = modal_count / k if k else 0.0
+ return modal_value, confidence
+
+
+def _aggregate_multiselect(samples: List[Any], k: int, threshold: float):
+ """Per-label vote fraction; predicted set = labels with fraction >= threshold.
+
+ Each sample is a list (possibly empty) of selected label names. None
+ samples count toward k.
+ """
+ per_label: Dict[str, int] = Counter()
+ for s in samples:
+ if s is None:
+ continue
+ for lab in s:
+ per_label[str(lab)] += 1
+ per_label_conf = {lab: cnt / k for lab, cnt in per_label.items()} if k else {}
+ predicted = sorted([lab for lab, frac in per_label_conf.items() if frac >= threshold])
+ if predicted:
+ confidence = sum(per_label_conf[lab] for lab in predicted) / len(predicted)
+ elif per_label_conf:
+ # No label cleared the bar: confidence in the (empty) prediction is how
+ # strongly the judges agreed to *exclude* the labels they saw.
+ confidence = 1.0 - (sum(per_label_conf.values()) / len(per_label_conf))
+ else:
+ confidence = 1.0 # every draw selected nothing -> confident empty set
+ return predicted, confidence, per_label_conf
+
+
+def _aggregate_span(samples: List[Any], k: int, cluster_threshold: float, keep_threshold: float):
+ """Cluster spans across k samples (EXPERIMENTAL).
+
+ Each non-None sample is a list of span dicts {start, end, label}. Spans are
+ greedily clustered when they share a label and overlap (IoU >=
+ cluster_threshold). A cluster's support is the number of distinct samples
+ that contributed to it; confidence = support / k. Clusters with confidence
+ >= keep_threshold are kept; the representative span is the modal exact
+ (start, end) within the cluster.
+
+ Returns (modal_spans, mean_confidence) where modal_spans is a list of
+ {start, end, label, confidence}.
+ """
+ clusters: List[Dict[str, Any]] = [] # each: {label, rep:(s,e), members:[(s,e,sample_idx)], samples:set}
+ for idx, sample in enumerate(samples):
+ if not sample:
+ continue
+ for sp in sample:
+ try:
+ s, e, lab = int(sp["start"]), int(sp["end"]), str(sp["label"])
+ except (KeyError, TypeError, ValueError):
+ continue
+ if e <= s:
+ continue
+ placed = False
+ for c in clusters:
+ if c["label"] == lab and span_iou(c["rep"], (s, e)) >= cluster_threshold:
+ c["members"].append((s, e, idx))
+ c["samples"].add(idx)
+ placed = True
+ break
+ if not placed:
+ clusters.append({"label": lab, "rep": (s, e),
+ "members": [(s, e, idx)], "samples": {idx}})
+
+ modal_spans = []
+ confidences = []
+ for c in clusters:
+ support = len(c["samples"])
+ confidence = support / k if k else 0.0
+ if confidence < keep_threshold:
+ continue
+ # representative = modal exact (start,end) among members
+ offset_counts = Counter((s, e) for s, e, _ in c["members"])
+ (rs, re), _ = offset_counts.most_common(1)[0]
+ modal_spans.append({"start": rs, "end": re, "label": c["label"],
+ "confidence": round(confidence, 6)})
+ confidences.append(confidence)
+
+ modal_spans.sort(key=lambda d: (d["start"], d["end"], d["label"]))
+ mean_conf = sum(confidences) / len(confidences) if confidences else 0.0
+ return modal_spans, mean_conf
+
+
+def aggregate(
+ model: str,
+ instance_id: str,
+ schema_name: str,
+ annotation_type: str,
+ samples: List[Any],
+ k: int,
+ multiselect_threshold: float = 0.5,
+ span_cluster_threshold: float = 0.5,
+ span_keep_threshold: float = 0.5,
+) -> ModelItemResult:
+ """Reduce raw samples to a ModelItemResult for the given schema type."""
+ per_label_conf: Dict[str, float] = {}
+ if annotation_type == "multiselect":
+ modal, confidence, per_label_conf = _aggregate_multiselect(
+ samples, k, multiselect_threshold
+ )
+ elif annotation_type == "span":
+ modal, confidence = _aggregate_span(
+ samples, k, span_cluster_threshold, span_keep_threshold
+ )
+ else:
+ # radio, likert, select, and any other single-label categorical type
+ modal, confidence = _aggregate_categorical(samples, k)
+
+ return ModelItemResult(
+ model=model,
+ instance_id=instance_id,
+ schema_name=schema_name,
+ annotation_type=annotation_type,
+ modal_label=modal,
+ confidence=round(confidence, 6),
+ k=k,
+ samples=samples,
+ per_label_confidence={kk: round(v, 6) for kk, v in per_label_conf.items()},
+ )
diff --git a/potato/judge_calibration/calibration.py b/potato/judge_calibration/calibration.py
new file mode 100644
index 0000000000000000000000000000000000000000..e81d03f0a3da556ad78b5540aef716154e139c8a
--- /dev/null
+++ b/potato/judge_calibration/calibration.py
@@ -0,0 +1,97 @@
+"""
+Calibration metrics for judge predictions.
+
+Pure functions (no Potato imports) computing Expected Calibration Error (ECE),
+reliability-diagram bins, and the Brier score from per-prediction
+(confidence, correctness) pairs. ``confidence`` is the k-sample vote fraction;
+``correct`` is 1 if the modal prediction matched the human gold label else 0.
+
+These did not exist anywhere in Potato โ implemented here and covered by
+known-value fixtures in tests/unit/test_jc_calibration.py.
+"""
+
+from typing import Dict, List, Sequence
+
+
+def _bin_index(conf: float, n_bins: int) -> int:
+ """Equal-width bin index in [0, n_bins-1]; conf==1.0 lands in the last bin."""
+ if conf <= 0:
+ return 0
+ if conf >= 1:
+ return n_bins - 1
+ return min(n_bins - 1, int(conf * n_bins))
+
+
+def reliability_bins(
+ confidences: Sequence[float],
+ correctness: Sequence[int],
+ n_bins: int = 10,
+) -> List[Dict[str, float]]:
+ """Return per-bin stats for a reliability diagram.
+
+ Each bin dict has: bin_lo, bin_hi, count, mean_confidence, accuracy.
+ Empty bins are included (count=0) so the diagram has a consistent x-axis.
+ """
+ if len(confidences) != len(correctness):
+ raise ValueError("confidences and correctness must be the same length")
+ width = 1.0 / n_bins
+ sums_conf = [0.0] * n_bins
+ sums_corr = [0.0] * n_bins
+ counts = [0] * n_bins
+ for c, y in zip(confidences, correctness):
+ b = _bin_index(float(c), n_bins)
+ sums_conf[b] += float(c)
+ sums_corr[b] += float(y)
+ counts[b] += 1
+
+ bins = []
+ for b in range(n_bins):
+ n = counts[b]
+ bins.append({
+ "bin_lo": round(b * width, 6),
+ "bin_hi": round((b + 1) * width, 6),
+ "count": n,
+ "mean_confidence": round(sums_conf[b] / n, 6) if n else 0.0,
+ "accuracy": round(sums_corr[b] / n, 6) if n else 0.0,
+ })
+ return bins
+
+
+def expected_calibration_error(
+ confidences: Sequence[float],
+ correctness: Sequence[int],
+ n_bins: int = 10,
+) -> float:
+ """ECE = sum over bins of (n_b/N) * |accuracy_b - mean_confidence_b|."""
+ n_total = len(confidences)
+ if n_total == 0:
+ return 0.0
+ ece = 0.0
+ for b in reliability_bins(confidences, correctness, n_bins):
+ if b["count"] == 0:
+ continue
+ ece += (b["count"] / n_total) * abs(b["accuracy"] - b["mean_confidence"])
+ return round(ece, 6)
+
+
+def brier_score(confidences: Sequence[float], correctness: Sequence[int]) -> float:
+ """Mean squared error between confidence and correctness (lower is better)."""
+ n = len(confidences)
+ if n == 0:
+ return 0.0
+ return round(sum((float(c) - float(y)) ** 2 for c, y in zip(confidences, correctness)) / n, 6)
+
+
+def calibration_report(
+ confidences: Sequence[float],
+ correctness: Sequence[int],
+ n_bins: int = 10,
+) -> Dict:
+ """Bundle ECE, Brier, and reliability bins for one model."""
+ return {
+ "ece": expected_calibration_error(confidences, correctness, n_bins),
+ "brier": brier_score(confidences, correctness),
+ "n_bins": n_bins,
+ "reliability_bins": reliability_bins(confidences, correctness, n_bins),
+ "n": len(confidences),
+ }
diff --git a/potato/judge_calibration/config.py b/potato/judge_calibration/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..d64f021b571bfa9c167ce872b2c94f37c49b0d9b
--- /dev/null
+++ b/potato/judge_calibration/config.py
@@ -0,0 +1,209 @@
+"""
+Judge Calibration configuration.
+
+Parses the ``judge_calibration:`` block of an application config into a typed
+``JudgeCalibrationConfig``. This is a deliberately lightweight cousin of
+``solo_mode`` โ there are no refinement loops, edge-case synthesis or
+disagreement-resolution settings. The workflow is: pick N judge LLMs, sample
+each ``k`` times over (up to) a capped number of items, then have human(s)
+*blind*-label a sample and produce a calibration report.
+
+LLM endpoint settings reuse ``solo_mode.config.ModelConfig`` (and its
+``to_endpoint_config()``), which already covers every provider Potato supports
+(openai/anthropic/ollama/vllm/gemini/...) plus env-var API-key expansion.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+import logging
+import os
+
+# Reuse the solo-mode model dataclass + parser so endpoint handling stays in
+# one place (env-var expansion, to_endpoint_config(), base_url aliasing, ...).
+from potato.solo_mode.config import ModelConfig, _parse_model_config
+
+logger = logging.getLogger(__name__)
+
+
+_VALID_SAMPLING_STRATEGIES = ("random", "stratified", "all")
+_VALID_GOLD = ("single", "majority")
+
+
+@dataclass
+class SamplingConfig:
+ """How the human calibration subset is drawn from the LLM-labeled items."""
+ strategy: str = "random" # random | stratified | all
+ stratify_by: Optional[str] = None # item-data field to stratify on
+ sample_size: int = 200 # number of items humans blind-label
+ seed: int = 42
+
+
+@dataclass
+class HumanConfig:
+ """Who provides the human ground-truth labels."""
+ num_raters: int = 1 # 1 = solo researcher; N adds human-human IAA
+ gold: str = "single" # single | majority (used when num_raters > 1)
+
+
+@dataclass
+class CalibrationConfig:
+ """Calibration (ECE / reliability diagram) settings."""
+ n_bins: int = 10
+
+
+@dataclass
+class OutputConfig:
+ """Where outputs are written and under what filenames."""
+ dir: str = "judge_calibration_output"
+ labels_file: str = "llm_labels.jsonl"
+ report_json: str = "report.json"
+ report_html: str = "report.html"
+
+
+@dataclass
+class JudgeCalibrationConfig:
+ """Top-level configuration for Judge Calibration mode."""
+ enabled: bool = False
+
+ # The judge instruction shown to every model. Supports ``{text}``,
+ # ``{labels}`` and ``{description}`` substitution (filled per-item /
+ # per-schema at generation time).
+ prompt: str = ""
+
+ # Judge LLMs. Each is sampled ``k_samples`` times per item.
+ models: List[ModelConfig] = field(default_factory=list)
+
+ k_samples: int = 5
+
+ # Cap on how many items the LLMs label. Exactly one of max_items / fraction
+ # is honored (max_items wins if both set); None/None means "all items".
+ max_items: Optional[int] = None
+ fraction: Optional[float] = None
+
+ sampling: SamplingConfig = field(default_factory=SamplingConfig)
+ human: HumanConfig = field(default_factory=HumanConfig)
+
+ # Annotation scheme names to evaluate. Empty = all categorical schemes.
+ schemas: List[str] = field(default_factory=list)
+
+ calibration: CalibrationConfig = field(default_factory=CalibrationConfig)
+ output: OutputConfig = field(default_factory=OutputConfig)
+
+ state_dir: str = "judge_calibration_state"
+
+ def validate(self) -> List[str]:
+ """Return a list of human-readable error strings (empty if valid)."""
+ errors: List[str] = []
+ if not self.enabled:
+ return errors
+
+ if not self.models:
+ errors.append("judge_calibration.models is required (at least one model)")
+ for i, m in enumerate(self.models):
+ if not m.model:
+ errors.append(f"judge_calibration.models[{i}].model is required")
+ if not m.endpoint_type:
+ errors.append(f"judge_calibration.models[{i}].endpoint_type is required")
+
+ if self.k_samples < 1:
+ errors.append("judge_calibration.k_samples must be >= 1")
+
+ if self.fraction is not None and not (0 < self.fraction <= 1):
+ errors.append("judge_calibration.fraction must be in (0, 1]")
+ if self.max_items is not None and self.max_items < 1:
+ errors.append("judge_calibration.max_items must be >= 1")
+
+ if self.sampling.strategy not in _VALID_SAMPLING_STRATEGIES:
+ errors.append(
+ f"judge_calibration.sampling.strategy must be one of "
+ f"{_VALID_SAMPLING_STRATEGIES}, got '{self.sampling.strategy}'"
+ )
+ if self.sampling.sample_size < 1:
+ errors.append("judge_calibration.sampling.sample_size must be >= 1")
+
+ if self.human.num_raters < 1:
+ errors.append("judge_calibration.human.num_raters must be >= 1")
+ if self.human.gold not in _VALID_GOLD:
+ errors.append(
+ f"judge_calibration.human.gold must be one of {_VALID_GOLD}, "
+ f"got '{self.human.gold}'"
+ )
+
+ if self.calibration.n_bins < 1:
+ errors.append("judge_calibration.calibration.n_bins must be >= 1")
+
+ # Confidence degeneracy guard: with k>1 samples but temperature 0 the
+ # samples never vary, so the vote-fraction confidence is always 1.0 and
+ # the calibration report is meaningless. Warn (don't fail).
+ if self.k_samples > 1:
+ for i, m in enumerate(self.models):
+ if m.temperature == 0:
+ logger.warning(
+ "judge_calibration.models[%d] (%s) has temperature=0 with "
+ "k_samples=%d; samples will be identical and confidence "
+ "will always be 1.0. Set temperature > 0 for meaningful "
+ "calibration.",
+ i, m.model, self.k_samples,
+ )
+
+ return errors
+
+
+def _parse_sampling(data: Dict[str, Any]) -> SamplingConfig:
+ data = data or {}
+ return SamplingConfig(
+ strategy=data.get("strategy", "random"),
+ stratify_by=data.get("stratify_by"),
+ sample_size=int(data.get("sample_size", 200)),
+ seed=int(data.get("seed", 42)),
+ )
+
+
+def _parse_human(data: Dict[str, Any]) -> HumanConfig:
+ data = data or {}
+ return HumanConfig(
+ num_raters=int(data.get("num_raters", 1)),
+ gold=data.get("gold", "single"),
+ )
+
+
+def _parse_output(data: Dict[str, Any]) -> OutputConfig:
+ data = data or {}
+ files = data.get("files", {}) or {}
+ return OutputConfig(
+ dir=data.get("dir", "judge_calibration_output"),
+ labels_file=files.get("labels", "llm_labels.jsonl"),
+ report_json=files.get("report_json", "report.json"),
+ report_html=files.get("report_html", "report.html"),
+ )
+
+
+def parse_judge_calibration_config(config_data: Dict[str, Any]) -> JudgeCalibrationConfig:
+ """Parse the ``judge_calibration`` section into a JudgeCalibrationConfig."""
+ jc = config_data.get("judge_calibration", {})
+ if not jc:
+ return JudgeCalibrationConfig(enabled=False)
+
+ models = [_parse_model_config(m) for m in jc.get("models", [])]
+
+ state_dir = jc.get("state_dir")
+ if not state_dir:
+ output_dir = config_data.get("output_annotation_dir", "annotation_output")
+ state_dir = os.path.join(output_dir, ".judge_calibration")
+
+ cal_data = jc.get("calibration", {}) or {}
+
+ return JudgeCalibrationConfig(
+ enabled=jc.get("enabled", False),
+ prompt=jc.get("prompt", ""),
+ models=models,
+ k_samples=int(jc.get("k_samples", 5)),
+ max_items=jc.get("max_items"),
+ fraction=jc.get("fraction"),
+ sampling=_parse_sampling(jc.get("sampling", {})),
+ human=_parse_human(jc.get("human", {})),
+ schemas=list(jc.get("schemas", []) or []),
+ calibration=CalibrationConfig(n_bins=int(cal_data.get("n_bins", 10))),
+ output=_parse_output(jc.get("output", {})),
+ state_dir=state_dir,
+ )
diff --git a/potato/judge_calibration/gamma.py b/potato/judge_calibration/gamma.py
new file mode 100644
index 0000000000000000000000000000000000000000..637fff55ff540514707c646361cca72f186af783
--- /dev/null
+++ b/potato/judge_calibration/gamma.py
@@ -0,0 +1,203 @@
+"""
+Local ฮณ (Gamma) inter-annotator agreement for spans โ EXPERIMENTAL.
+
+This is a self-contained, dependency-light reimplementation of the *core ideas*
+of the ฮณ measure (Mathet, Widlรถcher & Mรฉtivier, 2015, "The Unified and
+Holistic Method Gamma (ฮณ) for Inter-Annotator Agreement Measure and
+Alignment"). It is NOT a bit-exact reproduction of the canonical
+``pygamma-agreement`` package โ notably it computes ฮณ **pairwise** (then
+averages over annotator pairs) rather than solving the full multi-annotator
+continuum alignment, and its chance model is a simpler positional shuffle.
+
+What it keeps faithfully:
+ ฮณ = 1 โ (observed disorder) / (expected disorder)
+where *disorder* is the average dissimilarity of the best alignment between two
+annotators' unit sets, and *expected disorder* is that same quantity computed
+over random ("chance") annotators obtained by relocating each unit to a random
+position in the continuum (keeping its length and category).
+
+Dissimilarity of two units u, v (Mathet positional + categorical):
+ d_pos(u,v) = ((|start_uโstart_v| + |end_uโend_v|) / (len_u + len_v))ยฒ
+ d_cat(u,v) = 0 if same label else 1
+ d(u,v) = ฮฑยทd_pos + ฮฒยทd_cat
+A unit left unaligned costs ``delta_empty`` (so two units align only when
+d(u,v) < 2ยทdelta_empty).
+
+For the canonical, peer-reviewed implementation use ``pygamma-agreement``; this
+module exists so Potato can report a ฮณ-style number with no extra dependency.
+"""
+
+import logging
+from itertools import combinations
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import numpy as np
+from scipy.optimize import linear_sum_assignment
+
+logger = logging.getLogger(__name__)
+
+# A "unit" is (start, end, label).
+Unit = Tuple[int, int, str]
+
+
+def unit_dissimilarity(u: Unit, v: Unit, alpha: float = 1.0, beta: float = 1.0) -> float:
+ """Mathet positional + categorical dissimilarity between two units."""
+ su, eu, lu = u
+ sv, ev, lv = v
+ len_u, len_v = (eu - su), (ev - sv)
+ denom = len_u + len_v
+ d_pos = (((abs(su - sv) + abs(eu - ev)) / denom) ** 2) if denom > 0 else 0.0
+ d_cat = 0.0 if lu == lv else 1.0
+ return alpha * d_pos + beta * d_cat
+
+
+def _alignment_cost(
+ units_a: List[Unit], units_b: List[Unit],
+ delta_empty: float, alpha: float, beta: float,
+) -> Tuple[float, int]:
+ """Optimal alignment cost between two unit sets (Hungarian, with empties).
+
+ Returns (total_cost, n_alignments) where n_alignments counts every aligned
+ pair that involves at least one real unit (realโreal and realโempty).
+ """
+ na, nb = len(units_a), len(units_b)
+ if na == 0 and nb == 0:
+ return 0.0, 0
+ if na == 0:
+ return delta_empty * nb, nb
+ if nb == 0:
+ return delta_empty * na, na
+
+ # Square cost matrix with empty slots:
+ # rows: [real A] + [empty-for-B] cols: [real B] + [empty-for-A]
+ big = delta_empty * 1e6
+ n = na + nb
+ cost = np.full((n, n), 0.0)
+ for i in range(na):
+ for j in range(nb):
+ cost[i, j] = unit_dissimilarity(units_a[i], units_b[j], alpha, beta)
+ # A_i -> its own empty (col nb + i); other A-empties are forbidden
+ for i in range(na):
+ for j in range(na):
+ cost[i, nb + j] = delta_empty if i == j else big
+ # empty-for-B (row na + j) -> B_j; others forbidden
+ for j in range(nb):
+ for i in range(nb):
+ cost[na + j, i] = delta_empty if i == j else big
+ # empty-empty quadrant already 0
+
+ rows, cols = linear_sum_assignment(cost)
+ total = 0.0
+ n_align = 0
+ for r, c in zip(rows, cols):
+ is_real_r = r < na
+ is_real_c = c < nb
+ if not is_real_r and not is_real_c:
+ continue # empty-empty, ignore
+ total += cost[r, c]
+ n_align += 1
+ return total, n_align
+
+
+def _pair_disorder(
+ a_by_iid: Dict[str, List[Unit]], b_by_iid: Dict[str, List[Unit]],
+ delta_empty: float, alpha: float, beta: float,
+) -> Optional[float]:
+ """Average alignment dissimilarity between two annotators over shared items."""
+ shared = set(a_by_iid) & set(b_by_iid)
+ total, n = 0.0, 0
+ for iid in shared:
+ c, k = _alignment_cost(a_by_iid[iid], b_by_iid[iid], delta_empty, alpha, beta)
+ total += c
+ n += k
+ if n == 0:
+ return None
+ return total / n
+
+
+def _random_annotator(
+ by_iid: Dict[str, List[Unit]], lengths: Dict[str, int], rng: np.random.RandomState,
+) -> Dict[str, List[Unit]]:
+ """Relocate each unit to a random position in its continuum (keep len+label)."""
+ out: Dict[str, List[Unit]] = {}
+ for iid, units in by_iid.items():
+ L = lengths.get(iid, 0)
+ new_units = []
+ for (s, e, lab) in units:
+ ln = e - s
+ if L <= ln:
+ new_units.append((0, max(ln, 1), lab))
+ else:
+ ns = int(rng.randint(0, L - ln + 1))
+ new_units.append((ns, ns + ln, lab))
+ out[iid] = new_units
+ return out
+
+
+def _expected_pair_disorder(
+ a_by_iid, b_by_iid, lengths, delta_empty, alpha, beta, n_samples, rng,
+) -> Optional[float]:
+ vals = []
+ for _ in range(n_samples):
+ ra = _random_annotator(a_by_iid, lengths, rng)
+ rb = _random_annotator(b_by_iid, lengths, rng)
+ d = _pair_disorder(ra, rb, delta_empty, alpha, beta)
+ if d is not None:
+ vals.append(d)
+ if not vals:
+ return None
+ return float(np.mean(vals))
+
+
+def gamma_agreement(
+ raters: Dict[str, Dict[str, List[Unit]]],
+ lengths: Dict[str, int],
+ is_llm: Callable[[str], bool],
+ delta_empty: float = 1.0,
+ alpha: float = 1.0,
+ beta: float = 1.0,
+ n_samples: int = 30,
+ seed: int = 0,
+) -> Dict[str, Any]:
+ """Pairwise ฮณ averaged over annotator pairs, partitioned by rater kind.
+
+ Args:
+ raters: {rater_name: {instance_id: [(start, end, label), ...]}}
+ lengths: {instance_id: continuum length (chars)}; missing -> max end seen.
+ is_llm: predicate marking a rater name as an LLM (vs human).
+ """
+ names = sorted(raters)
+ if len(names) < 2:
+ return {"gamma": None, "n_pairs": 0}
+
+ rng = np.random.RandomState(seed)
+ partitioned = {"human_llm": [], "llm_llm": [], "human_human": []}
+ all_g = []
+
+ for a, b in combinations(names, 2):
+ obs = _pair_disorder(raters[a], raters[b], delta_empty, alpha, beta)
+ if obs is None:
+ continue
+ exp = _expected_pair_disorder(
+ raters[a], raters[b], lengths, delta_empty, alpha, beta, n_samples, rng)
+ if not exp: # None or 0 -> ฮณ undefined
+ continue
+ g = 1.0 - obs / exp
+ all_g.append(g)
+ kind = ("llm_llm" if is_llm(a) and is_llm(b)
+ else "human_human" if not is_llm(a) and not is_llm(b)
+ else "human_llm")
+ partitioned[kind].append(g)
+
+ def _mean(xs):
+ return round(float(np.mean(xs)), 4) if xs else None
+
+ return {
+ "gamma": _mean(all_g),
+ "mean_human_llm": _mean(partitioned["human_llm"]),
+ "mean_llm_llm": _mean(partitioned["llm_llm"]),
+ "mean_human_human": _mean(partitioned["human_human"]),
+ "n_pairs": len(all_g),
+ "n_samples": n_samples,
+ "approximate": True,
+ }
diff --git a/potato/judge_calibration/generation.py b/potato/judge_calibration/generation.py
new file mode 100644
index 0000000000000000000000000000000000000000..e57c017919ac3a53b37d1114b23ccd2a099e2612
--- /dev/null
+++ b/potato/judge_calibration/generation.py
@@ -0,0 +1,314 @@
+"""
+Background LLM judge generation.
+
+For every (model, item, schema) the judge LLM is queried ``k`` times; the raw
+draws are aggregated (see ``aggregation.py``) into a modal prediction + a
+vote-fraction confidence and persisted to the ``ResultStore``.
+
+Endpoints are built with ``AIEndpointFactory`` via the reused solo-mode
+``ModelConfig.to_endpoint_config()``. Label parsing reuses
+``ai.judge.extract_labels`` / ``_fuzzy_match_label`` so model output is mapped
+onto the schema's allowed label space exactly as the single-judge feature does.
+"""
+
+import json
+import logging
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+from pydantic import BaseModel
+
+from potato.ai.judge import extract_labels, _fuzzy_match_label
+from potato.judge_calibration.aggregation import aggregate, ModelItemResult
+
+logger = logging.getLogger(__name__)
+
+_MAX_WORKERS = 8
+_SAVE_EVERY = 25 # persist the result store after this many completed units
+
+
+# ----- output schemas -----------------------------------------------------
+
+class _SingleLabel(BaseModel):
+ label: str = ""
+ reasoning: str = ""
+
+
+class _MultiLabel(BaseModel):
+ labels: List[str] = []
+ reasoning: str = ""
+
+
+class _SpanItem(BaseModel):
+ start: int = 0
+ end: int = 0
+ label: str = ""
+
+
+class _SpanList(BaseModel):
+ spans: List[_SpanItem] = []
+ reasoning: str = ""
+
+
+def _output_model(annotation_type: str):
+ if annotation_type == "multiselect":
+ return _MultiLabel
+ if annotation_type == "span":
+ return _SpanList
+ return _SingleLabel
+
+
+# ----- prompt -------------------------------------------------------------
+
+def build_prompt(judge_prompt: str, schema_info: Dict[str, Any], text: str) -> str:
+ """Compose the judge prompt for one item + schema.
+
+ The user's ``judge_prompt`` may contain ``{text}``, ``{labels}`` and
+ ``{description}`` placeholders; if present they are substituted, otherwise
+ the standard scaffold (labels + item + JSON instruction) is appended.
+ """
+ labels = extract_labels(schema_info)
+ description = schema_info.get("description", "") or ""
+ labels_str = ", ".join(labels)
+ annotation_type = schema_info.get("annotation_type", "radio")
+
+ head = judge_prompt or "You are an impartial expert annotator."
+ try:
+ head = head.format(text=text, labels=labels_str, description=description)
+ substituted = head != (judge_prompt or "")
+ except (KeyError, IndexError, ValueError):
+ substituted = False
+
+ parts = [head, ""]
+ if description:
+ parts.append(f"Task: {description}")
+ if labels:
+ parts.append("Allowed labels: " + labels_str)
+ if annotation_type == "span":
+ # Spans need the exact text + character-offset instructions.
+ parts.append("\nItem to judge (label spans by 0-based character offsets):")
+ parts.append(_truncate(text, 4000))
+ parts.append(
+ '\nRespond as JSON: {"spans": [{"start": , '
+ '"end": , "label": }], '
+ '"reasoning": }. Return an empty list if no span applies.'
+ )
+ return "\n".join(p for p in parts if p != "")
+ if not substituted:
+ parts.append("\nItem to judge:")
+ parts.append(_truncate(text, 4000))
+ if annotation_type == "multiselect":
+ parts.append(
+ '\nRespond as JSON: {"labels": [], '
+ '"reasoning": }.'
+ )
+ else:
+ parts.append(
+ '\nRespond as JSON: {"label": , '
+ '"reasoning": }.'
+ )
+ return "\n".join(p for p in parts if p != "")
+
+
+def _truncate(text: str, limit: int = 4000) -> str:
+ text = str(text or "")
+ return text if len(text) <= limit else text[:limit] + "โฆ"
+
+
+# ----- response parsing ---------------------------------------------------
+
+def _extract_data(endpoint, response) -> Dict[str, Any]:
+ if isinstance(response, str):
+ try:
+ return json.loads(endpoint.parseStringToJson(response))
+ except Exception:
+ try:
+ return json.loads(response)
+ except Exception:
+ return {}
+ if hasattr(response, "model_dump"):
+ return response.model_dump()
+ if hasattr(response, "dict"):
+ return response.dict()
+ return response or {}
+
+
+def parse_sample(endpoint, response, annotation_type: str, valid_labels: List[str]):
+ """Parse one model response into a label value (or None on failure).
+
+ Returns a label name (str) for single-label schemas, a sorted list of
+ matched label names for multiselect, or None if nothing usable was parsed.
+ """
+ data = _extract_data(endpoint, response)
+ if not isinstance(data, dict):
+ return None
+
+ if annotation_type == "multiselect":
+ raw = data.get("labels", [])
+ if not isinstance(raw, list):
+ raw = [raw]
+ matched = []
+ for item in raw:
+ m = _fuzzy_match_label(str(item).strip(), valid_labels) if valid_labels else str(item).strip()
+ if m:
+ matched.append(m)
+ return sorted(set(matched)) # empty list is a valid "selected nothing"
+
+ if annotation_type == "span":
+ raw = data.get("spans", [])
+ if not isinstance(raw, list):
+ return None
+ spans = []
+ for sp in raw:
+ if not isinstance(sp, dict):
+ continue
+ try:
+ start, end = int(sp.get("start")), int(sp.get("end"))
+ except (TypeError, ValueError):
+ continue
+ if end <= start:
+ continue
+ lab = str(sp.get("label", "")).strip()
+ m = _fuzzy_match_label(lab, valid_labels) if valid_labels else lab
+ if not m:
+ continue
+ spans.append({"start": start, "end": end, "label": m})
+ return spans # empty list is a valid "no spans"
+
+ raw = str(data.get("label", "")).strip()
+ if not raw:
+ return None
+ if valid_labels and raw not in valid_labels:
+ matched = _fuzzy_match_label(raw, valid_labels)
+ return matched # may be None
+ return raw
+
+
+# ----- generation thread --------------------------------------------------
+
+class LLMGenerationThread(threading.Thread):
+ """Runs judge generation for all models/items/schemas in the background."""
+
+ def __init__(
+ self,
+ config, # JudgeCalibrationConfig
+ work_items: List[Tuple[str, str]], # (instance_id, text)
+ schema_infos: List[Dict[str, Any]], # schema dicts to evaluate
+ result_store, # ResultStore
+ on_progress: Optional[Callable[[Dict[str, Any]], None]] = None,
+ on_complete: Optional[Callable[[], None]] = None,
+ multiselect_threshold: float = 0.5,
+ ):
+ super().__init__(daemon=True, name="jc-generation")
+ self.config = config
+ self.work_items = work_items
+ self.schema_infos = schema_infos
+ self.store = result_store
+ self.on_progress = on_progress
+ self.on_complete = on_complete
+ self.multiselect_threshold = multiselect_threshold
+ self._stop_event = threading.Event()
+ self.error: Optional[str] = None
+
+ self.total_units = len(work_items) * len(schema_infos) * max(1, len(config.models))
+ self.done_units = 0
+ self._lock = threading.Lock()
+
+ def stop(self) -> None:
+ self._stop_event.set()
+
+ # -- one (model, item, schema) unit: k samples + aggregate + store --
+ def _label_unit(self, endpoint, model_name: str, instance_id: str, text: str,
+ schema_info: Dict[str, Any]) -> ModelItemResult:
+ annotation_type = schema_info.get("annotation_type", "radio")
+ schema_name = schema_info.get("name", "")
+ valid_labels = extract_labels(schema_info)
+ output_model = _output_model(annotation_type)
+ prompt = build_prompt(self.config.prompt, schema_info, text)
+
+ samples: List[Any] = []
+ for _ in range(self.config.k_samples):
+ if self._stop_event.is_set():
+ break
+ try:
+ response = endpoint.query(prompt, output_model)
+ samples.append(parse_sample(endpoint, response, annotation_type, valid_labels))
+ except Exception as e:
+ logger.warning("JC query failed (%s/%s/%s): %s",
+ model_name, instance_id, schema_name, e)
+ samples.append(None)
+
+ return aggregate(
+ model=model_name,
+ instance_id=instance_id,
+ schema_name=schema_name,
+ annotation_type=annotation_type,
+ samples=samples,
+ k=self.config.k_samples,
+ multiselect_threshold=self.multiselect_threshold,
+ )
+
+ def _bump_progress(self, model_name: str) -> None:
+ with self._lock:
+ self.done_units += 1
+ done = self.done_units
+ if self.on_progress and (done % 5 == 0 or done == self.total_units):
+ self.on_progress({
+ "done_units": done,
+ "total_units": self.total_units,
+ "current_model": model_name,
+ })
+
+ def run(self) -> None:
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+ for model in self.config.models:
+ if self._stop_event.is_set():
+ break
+ try:
+ endpoint = AIEndpointFactory.create_endpoint(model.to_endpoint_config())
+ except Exception as e:
+ self.error = f"Failed to create endpoint for {model.model}: {e}"
+ logger.error(self.error)
+ continue
+ if endpoint is None:
+ logger.error("JC: endpoint None for %s", model.model)
+ continue
+
+ # Build the list of units still needing generation (resume-safe).
+ units: List[Tuple[str, str, Dict[str, Any]]] = []
+ for instance_id, text in self.work_items:
+ for schema_info in self.schema_infos:
+ if self.store.has(model.model, instance_id, schema_info.get("name", "")):
+ with self._lock:
+ self.done_units += 1
+ continue
+ units.append((instance_id, text, schema_info))
+
+ completed = 0
+ with ThreadPoolExecutor(max_workers=_MAX_WORKERS) as ex:
+ futures = {
+ ex.submit(self._label_unit, endpoint, model.model, iid, text, si): (iid, si)
+ for (iid, text, si) in units
+ }
+ for fut in as_completed(futures):
+ if self._stop_event.is_set():
+ break
+ try:
+ result = fut.result()
+ self.store.upsert(result, save=False)
+ except Exception as e:
+ logger.warning("JC unit failed: %s", e)
+ completed += 1
+ if completed % _SAVE_EVERY == 0:
+ self.store._save()
+ self._bump_progress(model.model)
+
+ self.store._save()
+
+ if self.on_complete and not self._stop_event.is_set():
+ self.on_complete()
+ except Exception as e:
+ self.error = str(e)
+ logger.exception("JC generation thread crashed: %s", e)
diff --git a/potato/judge_calibration/manager.py b/potato/judge_calibration/manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..51aaaff73b75efdbcc0dcd92707229ad071ce24d
--- /dev/null
+++ b/potato/judge_calibration/manager.py
@@ -0,0 +1,286 @@
+"""
+Judge Calibration manager โ the singleton orchestrator.
+
+Owns the parsed config, the phase state machine and the LLM result store, and
+drives the run lifecycle:
+
+ SETUP -> GENERATING (background thread) -> HUMAN_CALIBRATION -> REPORT -> COMPLETED
+
+Items and annotation schemes are resolved lazily (at ``start_generation``)
+because data is loaded after the manager is constructed at app startup.
+"""
+
+import logging
+import threading
+from typing import Any, Dict, List, Optional, Tuple
+
+from potato.judge_calibration.config import (
+ JudgeCalibrationConfig,
+ parse_judge_calibration_config,
+)
+from potato.judge_calibration.phase import JCPhase, JCPhaseController
+from potato.judge_calibration.storage import ResultStore
+from potato.judge_calibration.generation import LLMGenerationThread
+
+logger = logging.getLogger(__name__)
+
+# Generation is schema-type agnostic (aggregation handles each), but we only
+# feed it types whose label space extract_labels() understands.
+SUPPORTED_GENERATION_TYPES = {"radio", "select", "likert", "multiselect", "span"}
+
+
+class JudgeCalibrationManager:
+ def __init__(self, app_config: Dict[str, Any]):
+ self.app_config = app_config or {}
+ self.config: JudgeCalibrationConfig = parse_judge_calibration_config(self.app_config)
+ self.phase = JCPhaseController(self.config.state_dir)
+ self.store = ResultStore(self.config.state_dir)
+ self._lock = threading.RLock()
+ self._gen_thread: Optional[LLMGenerationThread] = None
+ self._progress: Dict[str, Any] = {
+ "done_units": 0,
+ "total_units": 0,
+ "current_model": None,
+ }
+
+ # Resume any prior run.
+ self.phase.load_state()
+ self.store.load()
+
+ # ----- schema / item resolution --------------------------------------
+
+ def get_schema_infos(self) -> List[Dict[str, Any]]:
+ """Annotation schemes to evaluate (filtered by config.schemas + type)."""
+ schemes = self.app_config.get("annotation_schemes", []) or []
+ wanted = set(self.config.schemas)
+ out = []
+ for s in schemes:
+ if not isinstance(s, dict):
+ continue
+ name = s.get("name")
+ atype = s.get("annotation_type")
+ if wanted and name not in wanted:
+ continue
+ if atype not in SUPPORTED_GENERATION_TYPES:
+ if wanted and name in wanted:
+ logger.warning(
+ "judge_calibration: schema '%s' has unsupported type '%s'; skipping",
+ name, atype,
+ )
+ continue
+ out.append(s)
+ return out
+
+ def _resolve_cap(self, total: int) -> int:
+ if self.config.max_items is not None:
+ return min(total, max(1, self.config.max_items))
+ if self.config.fraction is not None:
+ return max(1, int(round(total * self.config.fraction)))
+ return total
+
+ def gather_work_items(self) -> List[Tuple[str, str]]:
+ """(instance_id, text) pairs the LLMs will label, honoring the cap."""
+ from potato.item_state_management import get_item_state_manager
+
+ ism = get_item_state_manager()
+ all_ids = list(ism.instance_id_to_instance.keys())
+ cap = self._resolve_cap(len(all_ids))
+ selected = all_ids[:cap]
+
+ text_key = (self.app_config.get("item_properties", {}) or {}).get("text_key")
+ items: List[Tuple[str, str]] = []
+ for iid in selected:
+ item = ism.instance_id_to_instance[iid]
+ data = item.get_data()
+ if isinstance(data, dict) and text_key and text_key in data:
+ text = data[text_key]
+ else:
+ text = item.get_text()
+ items.append((iid, str(text)))
+ return items
+
+ # ----- lifecycle -----------------------------------------------------
+
+ def update_config(self, overrides: Dict[str, Any]) -> List[str]:
+ """Merge wizard overrides into the judge_calibration config + re-parse.
+
+ Returns a list of validation errors (empty if the new config is valid).
+ The config is updated regardless so the wizard can show the errors.
+ """
+ with self._lock:
+ jc = dict(self.app_config.get("judge_calibration", {}) or {})
+ jc.update(overrides or {})
+ jc["enabled"] = True
+ self.app_config["judge_calibration"] = jc
+ self.config = parse_judge_calibration_config(self.app_config)
+ return self.config.validate()
+
+ def is_generating(self) -> bool:
+ with self._lock:
+ return self._gen_thread is not None and self._gen_thread.is_alive()
+
+ def start_generation(self, force_restart: bool = False) -> bool:
+ """Kick off (or resume) background LLM generation. Returns True if started."""
+ with self._lock:
+ if self.is_generating():
+ return False
+ if not self.config.models:
+ raise ValueError("judge_calibration: no models configured")
+
+ schema_infos = self.get_schema_infos()
+ if not schema_infos:
+ raise ValueError("judge_calibration: no supported schemas to evaluate")
+ work_items = self.gather_work_items()
+ if not work_items:
+ raise ValueError("judge_calibration: no items to label")
+
+ if force_restart:
+ self.store.clear()
+ self.phase.reset()
+
+ if self.phase.get_current_phase() != JCPhase.GENERATING:
+ # SETUP -> GENERATING (force covers re-runs from later phases).
+ self.phase.transition_to(JCPhase.GENERATING, reason="start", force=True)
+
+ # Record run metadata for the report.
+ self.phase.set_phase_data("n_models", len(self.config.models))
+ self.phase.set_phase_data("n_schemas", len(schema_infos))
+ self.phase.set_phase_data("n_items", len(work_items))
+
+ self._progress = {
+ "done_units": 0,
+ "total_units": len(work_items) * len(schema_infos) * len(self.config.models),
+ "current_model": None,
+ }
+
+ self._gen_thread = LLMGenerationThread(
+ config=self.config,
+ work_items=work_items,
+ schema_infos=schema_infos,
+ result_store=self.store,
+ on_progress=self._on_progress,
+ on_complete=self._on_generation_complete,
+ )
+ self._gen_thread.start()
+ logger.info("judge_calibration: generation started (%d units)",
+ self._progress["total_units"])
+ return True
+
+ def _on_progress(self, progress: Dict[str, Any]) -> None:
+ with self._lock:
+ self._progress.update(progress)
+
+ def _on_generation_complete(self) -> None:
+ try:
+ self.phase.transition_to(JCPhase.HUMAN_CALIBRATION, reason="generation complete")
+ # Draw the human calibration subset now that all items are labeled.
+ self.select_calibration_sample()
+ except Exception as e:
+ logger.error("judge_calibration: failed to advance phase: %s", e)
+
+ # ----- calibration sample + report -----------------------------------
+
+ def _stratum_fn(self):
+ """Build a stratum mapping for stratified sampling, or None.
+
+ If ``sampling.stratify_by`` names an item-data field, stratify on that.
+ Otherwise stratify on the first model's modal label (stratify-by-label)
+ for the first evaluated schema.
+ """
+ if self.config.sampling.strategy != "stratified":
+ return None
+
+ field = self.config.sampling.stratify_by
+ if field:
+ from potato.item_state_management import get_item_state_manager
+ ism = get_item_state_manager()
+
+ def by_field(iid):
+ item = ism.instance_id_to_instance.get(iid)
+ data = item.get_data() if item else None
+ return data.get(field) if isinstance(data, dict) else None
+ return by_field
+
+ # stratify-by-label: use the first model + first schema modal label
+ schema_infos = self.get_schema_infos()
+ if not schema_infos or not self.config.models:
+ return None
+ first_schema = schema_infos[0].get("name")
+ first_model = self.config.models[0].model
+
+ def by_label(iid):
+ r = self.store.get(first_model, iid, first_schema)
+ return r.modal_label if r else None
+ return by_label
+
+ def select_calibration_sample(self) -> List[str]:
+ """Pick (and persist) the human calibration subset from labeled items."""
+ from potato.judge_calibration.sampler import select_calibration_sample as _select
+
+ labeled = self.store.labeled_instance_ids()
+ sample = _select(labeled, self.config.sampling, stratum_of=self._stratum_fn())
+ self.phase.set_phase_data("calibration_sample", sample)
+ logger.info("judge_calibration: calibration sample = %d items", len(sample))
+ return sample
+
+ def get_calibration_sample(self) -> List[str]:
+ return self.phase.get_phase_data("calibration_sample", []) or []
+
+ def build_report(self) -> Dict[str, Any]:
+ """Compute metrics and write report files. Advances phase to COMPLETED."""
+ from potato.judge_calibration import report as report_module
+ from potato.judge_calibration.phase import JCPhase
+
+ with self._lock:
+ if self.phase.get_current_phase() not in (JCPhase.REPORT, JCPhase.COMPLETED):
+ self.phase.transition_to(JCPhase.REPORT, reason="build report", force=True)
+ result = report_module.build_report(self)
+ with self._lock:
+ if self.phase.get_current_phase() != JCPhase.COMPLETED:
+ self.phase.transition_to(JCPhase.COMPLETED, reason="report complete")
+ return result
+
+ def get_progress(self) -> Dict[str, Any]:
+ with self._lock:
+ error = self._gen_thread.error if self._gen_thread else None
+ return {
+ "phase": self.phase.get_current_phase().to_str(),
+ "generating": self.is_generating(),
+ "results": self.store.count(),
+ "error": error,
+ **self._progress,
+ }
+
+ def get_status(self) -> Dict[str, Any]:
+ with self._lock:
+ return {
+ "enabled": self.config.enabled,
+ "phase": self.phase.get_current_phase().to_str(),
+ "progress": self.get_progress(),
+ "n_models": len(self.config.models),
+ "models": [m.model for m in self.config.models],
+ }
+
+
+# ----- singleton ----------------------------------------------------------
+
+_manager: Optional[JudgeCalibrationManager] = None
+_manager_lock = threading.Lock()
+
+
+def init_judge_calibration_manager(app_config: Dict[str, Any]) -> JudgeCalibrationManager:
+ global _manager
+ with _manager_lock:
+ _manager = JudgeCalibrationManager(app_config)
+ logger.info("Judge Calibration manager initialized")
+ return _manager
+
+
+def get_judge_calibration_manager() -> Optional[JudgeCalibrationManager]:
+ return _manager
+
+
+def clear_judge_calibration_manager() -> None:
+ global _manager
+ with _manager_lock:
+ _manager = None
diff --git a/potato/judge_calibration/metrics.py b/potato/judge_calibration/metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..c7c40e171e4eb889c8461bff6ed7f41e07f898c5
--- /dev/null
+++ b/potato/judge_calibration/metrics.py
@@ -0,0 +1,606 @@
+"""
+Metrics for judge calibration.
+
+Pure(ish) functions that take plain label dictionaries and produce the report
+numbers: per-model accuracy/precision/recall/F1 vs human gold, inter-annotator
+agreement (Cohen's kappa pairwise โ partitioned into human<->LLM / LLM<->LLM /
+human<->human, plus Fleiss' kappa and Krippendorff's alpha over all raters),
+per-model confusion matrices, calibration (ECE/Brier/reliability), and โ for
+likert โ mean absolute error.
+
+IAA reuses ``potato.agreement`` (cohen_kappa_pairwise / fleiss_kappa /
+interpret_kappa). Krippendorff uses simpledorff (nominal for radio, interval
+for likert โ this build of simpledorff ships nominal_metric + interval_metric
+only; interval is used as the ordinal proxy).
+
+Inputs (all keyed by instance id):
+ llm_modal: {model_name: {iid: label}}
+ llm_conf: {model_name: {iid: confidence}}
+ human_labels: {human_id: {iid: label}}
+Single-label schemas (radio/select/likert) are supported here; multiselect and
+span get their own handling in later phases.
+"""
+
+import logging
+from collections import Counter
+from typing import Any, Dict, List, Optional
+
+from potato.agreement import cohen_kappa_pairwise, fleiss_kappa, interpret_kappa
+from potato.judge_calibration.calibration import calibration_report
+
+logger = logging.getLogger(__name__)
+
+LLM_PREFIX = "llm::"
+HUMAN_PREFIX = "human::"
+
+
+# ----- gold resolution ----------------------------------------------------
+
+def _majority(labels: List[Any]) -> Optional[Any]:
+ """Most common label; ties broken by sorted string order (deterministic)."""
+ if not labels:
+ return None
+ counts = Counter(str(l) for l in labels)
+ top = max(counts.values())
+ winners = sorted(k for k, v in counts.items() if v == top)
+ winner_str = winners[0]
+ return next(l for l in labels if str(l) == winner_str)
+
+
+def resolve_gold(human_labels: Dict[str, Dict[str, Any]], gold_strategy: str) -> Dict[str, Any]:
+ """Per-instance human gold label."""
+ humans = sorted(human_labels.keys())
+ if not humans:
+ return {}
+ if gold_strategy == "single" or len(humans) == 1:
+ if len(humans) > 1:
+ logger.info("judge_calibration: gold=single with %d humans; using '%s'",
+ len(humans), humans[0])
+ return dict(human_labels[humans[0]])
+
+ # majority across humans
+ all_iids = set()
+ for h in humans:
+ all_iids.update(human_labels[h].keys())
+ gold = {}
+ for iid in all_iids:
+ votes = [human_labels[h][iid] for h in humans if iid in human_labels[h]]
+ m = _majority(votes)
+ if m is not None:
+ gold[iid] = m
+ return gold
+
+
+# ----- classification metrics --------------------------------------------
+
+def _classification_metrics(y_true: List[str], y_pred: List[str], valid_labels: List[str]) -> Dict[str, Any]:
+ from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix
+
+ labels = valid_labels or sorted(set(y_true) | set(y_pred))
+ acc = float(accuracy_score(y_true, y_pred)) if y_true else 0.0
+ p, r, f1, _ = precision_recall_fscore_support(
+ y_true, y_pred, labels=labels, average="macro", zero_division=0
+ )
+ pw, rw, f1w, _ = precision_recall_fscore_support(
+ y_true, y_pred, labels=labels, average="weighted", zero_division=0
+ )
+ cm = confusion_matrix(y_true, y_pred, labels=labels)
+ confusion = {
+ gold: {pred: int(cm[i][j]) for j, pred in enumerate(labels)}
+ for i, gold in enumerate(labels)
+ }
+ return {
+ "accuracy": round(acc, 6),
+ "precision_macro": round(float(p), 6),
+ "recall_macro": round(float(r), 6),
+ "f1_macro": round(float(f1), 6),
+ "precision_weighted": round(float(pw), 6),
+ "recall_weighted": round(float(rw), 6),
+ "f1_weighted": round(float(f1w), 6),
+ "confusion_matrix": confusion,
+ "labels": labels,
+ "n": len(y_true),
+ }
+
+
+def _mae(y_true: List[str], y_pred: List[str]) -> Optional[float]:
+ """Mean absolute error for numeric (likert) labels; None if non-numeric."""
+ try:
+ diffs = [abs(float(t) - float(p)) for t, p in zip(y_true, y_pred)]
+ except (TypeError, ValueError):
+ return None
+ return round(sum(diffs) / len(diffs), 6) if diffs else None
+
+
+# ----- IAA ----------------------------------------------------------------
+
+def _build_reliability_df(
+ llm_modal: Dict[str, Dict[str, Any]],
+ human_labels: Dict[str, Dict[str, Any]],
+):
+ import pandas as pd
+
+ rows = []
+ for model, preds in llm_modal.items():
+ for iid, label in preds.items():
+ if label is None:
+ continue
+ rows.append({"unit": iid, "annotator": LLM_PREFIX + model, "annotation": str(label)})
+ for human, preds in human_labels.items():
+ for iid, label in preds.items():
+ if label is None:
+ continue
+ rows.append({"unit": iid, "annotator": HUMAN_PREFIX + human, "annotation": str(label)})
+ return pd.DataFrame(rows, columns=["unit", "annotator", "annotation"])
+
+
+def _pair_kind(a: str, b: str) -> str:
+ a_llm, b_llm = a.startswith(LLM_PREFIX), b.startswith(LLM_PREFIX)
+ if a_llm and b_llm:
+ return "llm_llm"
+ if (not a_llm) and (not b_llm):
+ return "human_human"
+ return "human_llm"
+
+
+def compute_iaa(
+ llm_modal: Dict[str, Dict[str, Any]],
+ human_labels: Dict[str, Dict[str, Any]],
+ ordinal: bool = False,
+) -> Dict[str, Any]:
+ df = _build_reliability_df(llm_modal, human_labels)
+ if df.empty:
+ return {"cohen": {}, "fleiss": {}, "krippendorff": None}
+
+ cohen = cohen_kappa_pairwise(df)
+ # Partition pairs by rater kind.
+ partitioned = {"human_llm": [], "llm_llm": [], "human_human": []}
+ for pair in cohen.get("pairs", []):
+ kind = _pair_kind(pair["annotator_a"], pair["annotator_b"])
+ partitioned[kind].append(pair)
+
+ def _mean(pairs):
+ ks = [p["kappa"] for p in pairs]
+ return round(sum(ks) / len(ks), 4) if ks else None
+
+ fleiss = fleiss_kappa(df)
+
+ krippendorff = None
+ try:
+ import simpledorff
+ from simpledorff.metrics import nominal_metric, interval_metric
+ kdf = df.copy()
+ metric_fn = nominal_metric
+ if ordinal:
+ try:
+ kdf["annotation"] = kdf["annotation"].astype(float)
+ metric_fn = interval_metric
+ except (TypeError, ValueError):
+ metric_fn = nominal_metric
+ alpha = simpledorff.calculate_krippendorffs_alpha_for_df(
+ kdf, experiment_col="unit", annotator_col="annotator",
+ class_col="annotation", metric_fn=metric_fn,
+ )
+ krippendorff = {
+ "alpha": round(float(alpha), 4),
+ "metric": "interval" if (ordinal and metric_fn is interval_metric) else "nominal",
+ "interpretation": interpret_kappa(float(alpha)),
+ }
+ except Exception as e:
+ logger.warning("judge_calibration: krippendorff failed: %s", e)
+
+ return {
+ "cohen": {
+ "mean_kappa": cohen.get("mean_kappa"),
+ "mean_human_llm": _mean(partitioned["human_llm"]),
+ "mean_llm_llm": _mean(partitioned["llm_llm"]),
+ "mean_human_human": _mean(partitioned["human_human"]),
+ "pairs": cohen.get("pairs", []),
+ },
+ "fleiss": fleiss,
+ "krippendorff": krippendorff,
+ }
+
+
+# ----- multiselect --------------------------------------------------------
+
+def _jaccard(a, b) -> float:
+ sa, sb = set(a or []), set(b or [])
+ if not sa and not sb:
+ return 1.0
+ union = sa | sb
+ return len(sa & sb) / len(union) if union else 1.0
+
+
+def _mean_pairwise_jaccard(
+ rater_labels: Dict[str, Dict[str, Any]],
+) -> Dict[str, Any]:
+ """Mean pairwise Jaccard agreement, partitioned by rater kind."""
+ from itertools import combinations
+
+ raters = sorted(rater_labels.keys())
+ partitioned = {"human_llm": [], "llm_llm": [], "human_human": []}
+ all_scores = []
+ for a, b in combinations(raters, 2):
+ shared = set(rater_labels[a]) & set(rater_labels[b])
+ if not shared:
+ continue
+ score = sum(_jaccard(rater_labels[a][i], rater_labels[b][i]) for i in shared) / len(shared)
+ all_scores.append(score)
+ partitioned[_pair_kind(a, b)].append(score)
+
+ def _mean(xs):
+ return round(sum(xs) / len(xs), 4) if xs else None
+
+ return {
+ "mean_jaccard": _mean(all_scores),
+ "mean_human_llm": _mean(partitioned["human_llm"]),
+ "mean_llm_llm": _mean(partitioned["llm_llm"]),
+ "mean_human_human": _mean(partitioned["human_human"]),
+ }
+
+
+def compute_multiselect_report(
+ schema_name: str,
+ valid_labels: List[str],
+ llm_modal: Dict[str, Dict[str, Any]],
+ llm_conf: Dict[str, Dict[str, float]],
+ human_labels: Dict[str, Dict[str, Any]],
+ gold_strategy: str = "single",
+ n_bins: int = 10,
+) -> Dict[str, Any]:
+ """Per-label P/R/F1 + mean Jaccard + set-match calibration for multiselect."""
+ from sklearn.metrics import precision_recall_fscore_support
+ from sklearn.preprocessing import MultiLabelBinarizer
+ from potato.judge_calibration.calibration import calibration_report
+
+ # gold: per-instance set (majority = union of labels appearing in >half of raters)
+ humans = sorted(human_labels.keys())
+ gold: Dict[str, Any] = {}
+ if humans:
+ if gold_strategy == "single" or len(humans) == 1:
+ gold = {iid: set(v) for iid, v in human_labels[humans[0]].items()}
+ else:
+ all_iids = set().union(*[set(human_labels[h]) for h in humans])
+ for iid in all_iids:
+ votes = Counter()
+ raters_for = 0
+ for h in humans:
+ if iid in human_labels[h]:
+ raters_for += 1
+ for lab in human_labels[h][iid]:
+ votes[lab] += 1
+ gold[iid] = {lab for lab, c in votes.items() if c > raters_for / 2}
+
+ mlb = MultiLabelBinarizer(classes=valid_labels)
+ mlb.fit([valid_labels])
+
+ per_model = {}
+ for model, preds in llm_modal.items():
+ overlap = sorted(iid for iid in preds if iid in gold)
+ if not overlap:
+ per_model[model] = {"n": 0}
+ continue
+ y_true = mlb.transform([sorted(gold[iid]) for iid in overlap])
+ y_pred = mlb.transform([sorted(preds[iid] or []) for iid in overlap])
+ p, r, f1, _ = precision_recall_fscore_support(
+ y_true, y_pred, average="macro", zero_division=0)
+ pmi, rmi, f1mi, _ = precision_recall_fscore_support(
+ y_true, y_pred, average="micro", zero_division=0)
+ jacc = sum(_jaccard(preds[iid], gold[iid]) for iid in overlap) / len(overlap)
+
+ conf = llm_conf.get(model, {})
+ confidences = [float(conf.get(iid, 0.0)) for iid in overlap]
+ correctness = [1 if set(preds[iid] or []) == set(gold[iid]) else 0 for iid in overlap]
+
+ per_model[model] = {
+ "precision_macro": round(float(p), 6),
+ "recall_macro": round(float(r), 6),
+ "f1_macro": round(float(f1), 6),
+ "f1_micro": round(float(f1mi), 6),
+ "mean_jaccard": round(jacc, 6),
+ "exact_match_accuracy": round(sum(correctness) / len(correctness), 6),
+ "calibration": calibration_report(confidences, correctness, n_bins),
+ "labels": valid_labels,
+ "n": len(overlap),
+ }
+
+ # IAA via mean pairwise Jaccard across all raters
+ rater_labels = {LLM_PREFIX + m: {i: set(v or []) for i, v in d.items()}
+ for m, d in llm_modal.items()}
+ for h, d in human_labels.items():
+ rater_labels[HUMAN_PREFIX + h] = {i: set(v or []) for i, v in d.items()}
+
+ return {
+ "schema": schema_name,
+ "annotation_type": "multiselect",
+ "n_gold": len(gold),
+ "gold_strategy": gold_strategy,
+ "per_model": per_model,
+ "iaa": {"jaccard": _mean_pairwise_jaccard(rater_labels)},
+ }
+
+
+# ----- span (EXPERIMENTAL) ------------------------------------------------
+
+def _match_spans(pred: List[dict], gold: List[dict], iou_threshold: float):
+ """Greedy IoU matching of predicted spans to gold spans (same label).
+
+ Returns (tp, fp, fn, matched_ious, matched_pred_flags). Each predicted span
+ is a dict with start/end/label (+ optional confidence); matched_pred_flags
+ is a list aligned to `pred` indicating which predicted spans matched a gold.
+ """
+ from potato.judge_calibration.aggregation import span_iou
+
+ used_gold = set()
+ matched_ious = []
+ matched_pred = [False] * len(pred)
+ for pi, p in enumerate(pred):
+ best_iou, best_gi = 0.0, None
+ for gi, g in enumerate(gold):
+ if gi in used_gold or g.get("label") != p.get("label"):
+ continue
+ iou = span_iou((p["start"], p["end"]), (g["start"], g["end"]))
+ if iou >= iou_threshold and iou > best_iou:
+ best_iou, best_gi = iou, gi
+ if best_gi is not None:
+ used_gold.add(best_gi)
+ matched_ious.append(best_iou)
+ matched_pred[pi] = True
+ tp = len(matched_ious)
+ fp = len(pred) - tp
+ fn = len(gold) - tp
+ return tp, fp, fn, matched_ious, matched_pred
+
+
+def _prf(tp: int, fp: int, fn: int) -> Dict[str, float]:
+ precision = tp / (tp + fp) if (tp + fp) else 0.0
+ recall = tp / (tp + fn) if (tp + fn) else 0.0
+ f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
+ return {"precision": round(precision, 6), "recall": round(recall, 6), "f1": round(f1, 6)}
+
+
+def _pairwise_span_f1(a: Dict[str, List[dict]], b: Dict[str, List[dict]], iou_threshold: float):
+ """Symmetric span-F1 agreement between two raters over shared instances."""
+ shared = set(a) & set(b)
+ if not shared:
+ return None
+ tp = fp = fn = 0
+ for iid in shared:
+ t, f_p, f_n, _, _ = _match_spans(a[iid], b[iid], iou_threshold)
+ tp += t; fp += f_p; fn += f_n
+ return _prf(tp, fp, fn)["f1"]
+
+
+def _to_units(spans: List[dict]):
+ """Convert span dicts to (start, end, label) tuples."""
+ return [(int(s["start"]), int(s["end"]), str(s["label"])) for s in spans]
+
+
+def _segment_label(seg_lo: int, seg_hi: int, units) -> str:
+ """Label of the span covering an atomic segment, or 'O' (outside)."""
+ for (s, e, lab) in units:
+ if s <= seg_lo and e >= seg_hi:
+ return lab
+ return "O"
+
+
+def compute_span_token_iaa(rater_units: Dict[str, Dict[str, list]]) -> Dict[str, Any]:
+ """Chance-corrected span agreement via atomic-segment projection.
+
+ Each instance is cut at every span boundary any annotator drew; each atomic
+ segment gets that annotator's label (or 'O'). Standard Cohen/Fleiss ฮบ and
+ Krippendorff ฮฑ (nominal) then run over the (instance:segment, annotator,
+ label) table. Only segments inside the union of annotated regions are
+ considered, which limits โ but does not eliminate โ 'O'-inflation.
+ """
+ import pandas as pd
+
+ rows = []
+ # union of all instance ids
+ all_iids = set()
+ for d in rater_units.values():
+ all_iids.update(d.keys())
+
+ for iid in all_iids:
+ present = {name: d[iid] for name, d in rater_units.items() if iid in d}
+ boundaries = sorted({b for units in present.values() for (s, e, _) in units for b in (s, e)})
+ if len(boundaries) < 2:
+ continue # no annotated extent -> nothing to compare
+ for k in range(len(boundaries) - 1):
+ lo, hi = boundaries[k], boundaries[k + 1]
+ if hi <= lo:
+ continue
+ for name, units in present.items():
+ rows.append({"unit": f"{iid}:{k}", "annotator": name,
+ "annotation": _segment_label(lo, hi, units)})
+
+ if not rows:
+ return {"cohen": {}, "fleiss": {}, "krippendorff": None, "note": "no overlapping segments"}
+
+ df = pd.DataFrame(rows, columns=["unit", "annotator", "annotation"])
+ cohen = cohen_kappa_pairwise(df)
+ partitioned = {"human_llm": [], "llm_llm": [], "human_human": []}
+ for pair in cohen.get("pairs", []):
+ partitioned[_pair_kind(pair["annotator_a"], pair["annotator_b"])].append(pair["kappa"])
+
+ def _mean(xs):
+ return round(sum(xs) / len(xs), 4) if xs else None
+
+ krippendorff = None
+ try:
+ import simpledorff
+ from simpledorff.metrics import nominal_metric
+ alpha = simpledorff.calculate_krippendorffs_alpha_for_df(
+ df, experiment_col="unit", annotator_col="annotator",
+ class_col="annotation", metric_fn=nominal_metric)
+ krippendorff = {"alpha": round(float(alpha), 4), "interpretation": interpret_kappa(float(alpha))}
+ except Exception as e:
+ logger.warning("judge_calibration: span token krippendorff failed: %s", e)
+
+ return {
+ "cohen": {
+ "mean_kappa": cohen.get("mean_kappa"),
+ "mean_human_llm": _mean(partitioned["human_llm"]),
+ "mean_llm_llm": _mean(partitioned["llm_llm"]),
+ "mean_human_human": _mean(partitioned["human_human"]),
+ },
+ "fleiss": fleiss_kappa(df),
+ "krippendorff": krippendorff,
+ "n_segments": df["unit"].nunique(),
+ }
+
+
+def compute_span_report(
+ schema_name: str,
+ valid_labels: List[str],
+ llm_spans: Dict[str, Dict[str, List[dict]]],
+ human_spans: Dict[str, Dict[str, List[dict]]],
+ gold_strategy: str = "single",
+ iou_threshold: float = 0.5,
+ n_bins: int = 10,
+ instance_lengths: Optional[Dict[str, int]] = None,
+ gamma_samples: int = 30,
+) -> Dict[str, Any]:
+ """EXPERIMENTAL span metrics: IoU-matched P/R/F1, mean IoU, span calibration."""
+ from itertools import combinations
+ from potato.judge_calibration.calibration import calibration_report
+
+ humans = sorted(human_spans.keys())
+ gold: Dict[str, List[dict]] = {}
+ if humans:
+ if gold_strategy == "majority" and len(humans) > 1:
+ logger.info("judge_calibration: span gold=majority not supported; using single human '%s'", humans[0])
+ gold = dict(human_spans[humans[0]])
+
+ per_model = {}
+ for model, by_iid in llm_spans.items():
+ overlap = sorted(iid for iid in by_iid if iid in gold)
+ tp = fp = fn = 0
+ all_ious = []
+ confidences, correctness = [], []
+ for iid in overlap:
+ pred = by_iid[iid]
+ t, f_p, f_n, ious, matched = _match_spans(pred, gold[iid], iou_threshold)
+ tp += t; fp += f_p; fn += f_n
+ all_ious.extend(ious)
+ for p, ok in zip(pred, matched):
+ confidences.append(float(p.get("confidence", 0.0)))
+ correctness.append(1 if ok else 0)
+ block = _prf(tp, fp, fn)
+ block["mean_iou"] = round(sum(all_ious) / len(all_ious), 6) if all_ious else 0.0
+ block["tp"] = tp
+ block["fp"] = fp
+ block["fn"] = fn
+ block["n_instances"] = len(overlap)
+ block["calibration"] = calibration_report(confidences, correctness, n_bins)
+ per_model[model] = block
+
+ # IAA: mean pairwise span-F1 across all raters, partitioned by kind.
+ rater_spans = {LLM_PREFIX + m: d for m, d in llm_spans.items()}
+ for h, d in human_spans.items():
+ rater_spans[HUMAN_PREFIX + h] = d
+ partitioned = {"human_llm": [], "llm_llm": [], "human_human": []}
+ all_scores = []
+ for a, b in combinations(sorted(rater_spans), 2):
+ score = _pairwise_span_f1(rater_spans[a], rater_spans[b], iou_threshold)
+ if score is None:
+ continue
+ all_scores.append(score)
+ partitioned[_pair_kind(a, b)].append(score)
+
+ def _mean(xs):
+ return round(sum(xs) / len(xs), 4) if xs else None
+
+ # Chance-corrected agreement: token/segment ฮบ-ฮฑ and a local ฮณ.
+ rater_units = {name: {iid: _to_units(sp) for iid, sp in d.items()}
+ for name, d in rater_spans.items()}
+ token_iaa = compute_span_token_iaa(rater_units)
+
+ # Continuum lengths for ฮณ's chance model: actual text length if provided,
+ # else the max span end seen per instance.
+ lengths = dict(instance_lengths or {})
+ for d in rater_units.values():
+ for iid, units in d.items():
+ max_end = max((e for (_, e, _) in units), default=0)
+ lengths[iid] = max(lengths.get(iid, 0), max_end)
+
+ gamma = None
+ try:
+ from potato.judge_calibration.gamma import gamma_agreement
+ gamma = gamma_agreement(
+ rater_units, lengths,
+ is_llm=lambda n: n.startswith(LLM_PREFIX),
+ n_samples=gamma_samples,
+ )
+ except Exception as e:
+ logger.warning("judge_calibration: gamma agreement failed: %s", e)
+
+ return {
+ "schema": schema_name,
+ "annotation_type": "span",
+ "experimental": True,
+ "n_gold": len(gold),
+ "gold_strategy": "single",
+ "iou_threshold": iou_threshold,
+ "per_model": per_model,
+ "iaa": {
+ "span_f1": {
+ "mean": _mean(all_scores),
+ "mean_human_llm": _mean(partitioned["human_llm"]),
+ "mean_llm_llm": _mean(partitioned["llm_llm"]),
+ "mean_human_human": _mean(partitioned["human_human"]),
+ },
+ "token_kappa": token_iaa,
+ "gamma": gamma,
+ },
+ }
+
+
+# ----- top-level per-schema report ---------------------------------------
+
+def compute_schema_report(
+ schema_name: str,
+ annotation_type: str,
+ valid_labels: List[str],
+ llm_modal: Dict[str, Dict[str, Any]],
+ llm_conf: Dict[str, Dict[str, float]],
+ human_labels: Dict[str, Dict[str, Any]],
+ gold_strategy: str = "single",
+ n_bins: int = 10,
+) -> Dict[str, Any]:
+ """Build the full metric block for one single-label schema."""
+ ordinal = annotation_type == "likert"
+ gold = resolve_gold(human_labels, gold_strategy)
+
+ per_model: Dict[str, Any] = {}
+ for model, preds in llm_modal.items():
+ # overlap = instances with both a model prediction and a gold label
+ overlap = sorted(
+ iid for iid, lab in preds.items()
+ if lab is not None and iid in gold
+ )
+ y_true = [str(gold[iid]) for iid in overlap]
+ y_pred = [str(preds[iid]) for iid in overlap]
+
+ block = _classification_metrics(y_true, y_pred, valid_labels)
+ if ordinal:
+ block["mae"] = _mae(y_true, y_pred)
+
+ # calibration: correct vs gold; confidence = vote fraction
+ conf = llm_conf.get(model, {})
+ confidences = [float(conf.get(iid, 0.0)) for iid in overlap]
+ correctness = [1 if str(preds[iid]) == str(gold[iid]) else 0 for iid in overlap]
+ block["calibration"] = calibration_report(confidences, correctness, n_bins)
+ per_model[model] = block
+
+ iaa = compute_iaa(llm_modal, human_labels, ordinal=ordinal)
+
+ return {
+ "schema": schema_name,
+ "annotation_type": annotation_type,
+ "n_gold": len(gold),
+ "gold_strategy": gold_strategy,
+ "per_model": per_model,
+ "iaa": iaa,
+ }
diff --git a/potato/judge_calibration/phase.py b/potato/judge_calibration/phase.py
new file mode 100644
index 0000000000000000000000000000000000000000..541c395d02f661ba50fee251419073834faebb43
--- /dev/null
+++ b/potato/judge_calibration/phase.py
@@ -0,0 +1,179 @@
+"""
+Judge Calibration phase state machine.
+
+A minimal linear workflow (no branching loops, unlike solo mode):
+
+ SETUP -> GENERATING -> HUMAN_CALIBRATION -> REPORT -> COMPLETED
+
+State persists atomically to ``/phase_state.json`` so a run resumes
+across server restarts.
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum, auto
+from typing import Any, Dict, List, Optional, Set
+import json
+import logging
+import os
+import threading
+
+logger = logging.getLogger(__name__)
+
+
+class JCPhase(Enum):
+ """Judge Calibration workflow phases."""
+ SETUP = auto() # Config loaded; nothing generated yet
+ GENERATING = auto() # LLMs labeling in the background
+ HUMAN_CALIBRATION = auto() # Human(s) blind-labeling the sample
+ REPORT = auto() # Metrics/report being built
+ COMPLETED = auto() # Report available
+
+ @classmethod
+ def from_str(cls, s: str) -> "JCPhase":
+ return cls[s.upper().replace("-", "_")]
+
+ def to_str(self) -> str:
+ return self.name.lower().replace("_", "-")
+
+
+PHASE_TRANSITIONS: Dict[JCPhase, Set[JCPhase]] = {
+ JCPhase.SETUP: {JCPhase.GENERATING},
+ JCPhase.GENERATING: {JCPhase.HUMAN_CALIBRATION, JCPhase.SETUP},
+ JCPhase.HUMAN_CALIBRATION: {JCPhase.REPORT},
+ JCPhase.REPORT: {JCPhase.COMPLETED, JCPhase.HUMAN_CALIBRATION},
+ JCPhase.COMPLETED: {JCPhase.REPORT}, # allow re-running the report
+}
+
+
+@dataclass
+class JCPhaseState:
+ """Serializable phase state for a calibration run."""
+ current_phase: JCPhase = JCPhase.SETUP
+ phase_data: Dict[str, Any] = field(default_factory=dict)
+ history: List[Dict[str, str]] = field(default_factory=list)
+ started_at: Optional[datetime] = None
+ completed_at: Optional[datetime] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ "current_phase": self.current_phase.to_str(),
+ "phase_data": self.phase_data,
+ "history": self.history,
+ "started_at": self.started_at.isoformat() if self.started_at else None,
+ "completed_at": self.completed_at.isoformat() if self.completed_at else None,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "JCPhaseState":
+ return cls(
+ current_phase=JCPhase.from_str(data.get("current_phase", "setup")),
+ phase_data=data.get("phase_data", {}),
+ history=data.get("history", []),
+ started_at=(
+ datetime.fromisoformat(data["started_at"])
+ if data.get("started_at") else None
+ ),
+ completed_at=(
+ datetime.fromisoformat(data["completed_at"])
+ if data.get("completed_at") else None
+ ),
+ )
+
+
+class JCPhaseController:
+ """Phase state machine with atomic JSON persistence."""
+
+ _STATE_FILE = "phase_state.json"
+
+ def __init__(self, state_dir: Optional[str] = None):
+ self._lock = threading.RLock()
+ self.state = JCPhaseState()
+ self.state_dir = state_dir
+
+ def get_current_phase(self) -> JCPhase:
+ with self._lock:
+ return self.state.current_phase
+
+ def is_phase(self, phase: JCPhase) -> bool:
+ return self.get_current_phase() == phase
+
+ def can_transition_to(self, target: JCPhase) -> bool:
+ with self._lock:
+ return target in PHASE_TRANSITIONS.get(self.state.current_phase, set())
+
+ def transition_to(self, target: JCPhase, reason: str = "", force: bool = False) -> bool:
+ with self._lock:
+ current = self.state.current_phase
+ if not force and not self.can_transition_to(target):
+ raise ValueError(
+ f"Invalid phase transition: {current.to_str()} -> {target.to_str()}. "
+ f"Allowed: {[p.to_str() for p in PHASE_TRANSITIONS.get(current, set())]}"
+ )
+ self.state.history.append({
+ "from": current.to_str(),
+ "to": target.to_str(),
+ "timestamp": datetime.now().isoformat(),
+ "reason": reason,
+ })
+ self.state.current_phase = target
+ if current == JCPhase.SETUP and self.state.started_at is None:
+ self.state.started_at = datetime.now()
+ if target == JCPhase.COMPLETED:
+ self.state.completed_at = datetime.now()
+ logger.info("JC phase: %s -> %s (%s)", current.to_str(), target.to_str(), reason or "none")
+ self._save_state()
+ return True
+
+ def get_phase_data(self, key: str, default: Any = None) -> Any:
+ with self._lock:
+ return self.state.phase_data.get(key, default)
+
+ def set_phase_data(self, key: str, value: Any) -> None:
+ with self._lock:
+ self.state.phase_data[key] = value
+ self._save_state()
+
+ def reset(self) -> None:
+ with self._lock:
+ self.state = JCPhaseState()
+ self._save_state()
+
+ def get_status(self) -> Dict[str, Any]:
+ with self._lock:
+ return {
+ "current_phase": self.state.current_phase.to_str(),
+ "started_at": self.state.started_at.isoformat() if self.state.started_at else None,
+ "completed_at": self.state.completed_at.isoformat() if self.state.completed_at else None,
+ "history": self.state.history,
+ }
+
+ def _save_state(self) -> None:
+ if not self.state_dir:
+ return
+ try:
+ os.makedirs(self.state_dir, exist_ok=True)
+ path = os.path.join(self.state_dir, self._STATE_FILE)
+ tmp = path + ".tmp"
+ with open(tmp, "w") as f:
+ json.dump(self.state.to_dict(), f, indent=2)
+ os.replace(tmp, path)
+ except Exception as e:
+ logger.error("Error saving JC phase state: %s", e)
+
+ def load_state(self) -> bool:
+ if not self.state_dir:
+ return False
+ path = os.path.join(self.state_dir, self._STATE_FILE)
+ if not os.path.exists(path):
+ return False
+ try:
+ with open(path) as f:
+ data = json.load(f)
+ with self._lock:
+ self.state = JCPhaseState.from_dict(data)
+ logger.info("Loaded JC phase state: %s", self.state.current_phase.to_str())
+ return True
+ except Exception as e:
+ logger.error("Error loading JC phase state: %s", e)
+ return False
diff --git a/potato/judge_calibration/report.py b/potato/judge_calibration/report.py
new file mode 100644
index 0000000000000000000000000000000000000000..174ccc30d73c94fb418f74585eb396dbd1837870
--- /dev/null
+++ b/potato/judge_calibration/report.py
@@ -0,0 +1,353 @@
+"""
+Judge Calibration report assembly + output files.
+
+Pulls together the LLM results (from the ResultStore) and the human blind
+labels (from UserState) into the metric inputs, runs ``metrics`` per schema,
+and writes three artifacts under the configured output dir:
+
+- ``llm_labels.jsonl`` : every LLM's label on every labeled item (the deliverable)
+- ``report.json`` : the structured metrics report
+- ``report.html`` : a human-readable summary
+
+Human labels live entirely in UserState (never as pseudo-users) so they're read
+through the normal user-state API here.
+"""
+
+import json
+import logging
+import os
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Tuple
+
+from potato.judge_calibration.metrics import (
+ compute_schema_report,
+ compute_multiselect_report,
+ compute_span_report,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _is_selected(value: Any) -> bool:
+ """Whether a stored label value counts as 'chosen'."""
+ if value is None or value is False:
+ return False
+ return str(value).strip().lower() not in ("", "false", "0", "none")
+
+
+def extract_human_label(label_dict, schema_name: str, annotation_type: str):
+ """Pull a human's label for one schema from a flat {Label: value} dict.
+
+ Returns a label name (str) for single-label schemas, a sorted list for
+ multiselect, or None if the schema was not answered.
+ """
+ chosen = [
+ lab.name for lab, val in label_dict.items()
+ if getattr(lab, "schema", None) == schema_name and _is_selected(val)
+ ]
+ if annotation_type == "multiselect":
+ return sorted(chosen)
+ return chosen[0] if chosen else None
+
+
+def collect_metric_inputs(
+ store, schema_info: Dict[str, Any], restrict_ids: Optional[set] = None,
+) -> Tuple[Dict[str, Dict[str, Any]], Dict[str, Dict[str, float]], Dict[str, Dict[str, Any]]]:
+ """Assemble (llm_modal, llm_conf, human_labels) for one schema."""
+ schema_name = schema_info.get("name")
+ annotation_type = schema_info.get("annotation_type", "radio")
+
+ llm_modal: Dict[str, Dict[str, Any]] = {}
+ llm_conf: Dict[str, Dict[str, float]] = {}
+ for r in store.all_results():
+ if r.schema_name != schema_name:
+ continue
+ if restrict_ids is not None and r.instance_id not in restrict_ids:
+ continue
+ llm_modal.setdefault(r.model, {})[r.instance_id] = r.modal_label
+ llm_conf.setdefault(r.model, {})[r.instance_id] = r.confidence
+
+ human_labels: Dict[str, Dict[str, Any]] = {}
+ try:
+ from potato.user_state_management import get_user_state_manager
+ usm = get_user_state_manager()
+ users = usm.get_all_users() if usm else []
+ except Exception as e:
+ logger.warning("judge_calibration: could not load user states: %s", e)
+ users = []
+
+ for user in users:
+ uid = user.get_user_id()
+ for iid in user.get_annotated_instance_ids():
+ if restrict_ids is not None and iid not in restrict_ids:
+ continue
+ label_dict = user.get_label_annotations(iid)
+ if not label_dict:
+ continue
+ lab = extract_human_label(label_dict, schema_name, annotation_type)
+ if lab is None or (isinstance(lab, list) and not lab):
+ continue
+ human_labels.setdefault(uid, {})[iid] = lab
+
+ return llm_modal, llm_conf, human_labels
+
+
+def extract_human_spans(span_dict, schema_name: str) -> List[dict]:
+ """Pull a human's spans for one schema from a flat {SpanAnnotation: value} dict."""
+ out = []
+ for sp in span_dict:
+ if getattr(sp, "schema", None) != schema_name:
+ continue
+ out.append({"start": sp.start, "end": sp.end, "label": sp.name})
+ return out
+
+
+def collect_span_inputs(store, schema_info, restrict_ids=None):
+ """Assemble (llm_spans, human_spans) for one span schema.
+
+ llm_spans[model][iid] = list of {start, end, label, confidence}
+ human_spans[hid][iid] = list of {start, end, label}
+ """
+ schema_name = schema_info.get("name")
+
+ llm_spans: Dict[str, Dict[str, List[dict]]] = {}
+ for r in store.all_results():
+ if r.schema_name != schema_name:
+ continue
+ if restrict_ids is not None and r.instance_id not in restrict_ids:
+ continue
+ # modal_label for span is a list of span dicts (may be empty)
+ llm_spans.setdefault(r.model, {})[r.instance_id] = r.modal_label or []
+
+ human_spans: Dict[str, Dict[str, List[dict]]] = {}
+ try:
+ from potato.user_state_management import get_user_state_manager
+ usm = get_user_state_manager()
+ users = usm.get_all_users() if usm else []
+ except Exception as e:
+ logger.warning("judge_calibration: could not load user states: %s", e)
+ users = []
+
+ for user in users:
+ uid = user.get_user_id()
+ for iid in user.get_annotated_instance_ids():
+ if restrict_ids is not None and iid not in restrict_ids:
+ continue
+ span_dict = user.get_span_annotations(iid)
+ if not span_dict:
+ continue
+ spans = extract_human_spans(span_dict, schema_name)
+ if spans:
+ human_spans.setdefault(uid, {})[iid] = spans
+ return llm_spans, human_spans
+
+
+def _instance_text_lengths(manager) -> Dict[str, int]:
+ """Map instance_id -> text length (chars) for ฮณ's continuum model."""
+ lengths: Dict[str, int] = {}
+ try:
+ from potato.item_state_management import get_item_state_manager
+ ism = get_item_state_manager()
+ text_key = (manager.app_config.get("item_properties", {}) or {}).get("text_key")
+ for iid, item in ism.instance_id_to_instance.items():
+ data = item.get_data()
+ if isinstance(data, dict) and text_key and text_key in data:
+ lengths[iid] = len(str(data[text_key]))
+ else:
+ lengths[iid] = len(str(item.get_text()))
+ except Exception as e:
+ logger.warning("judge_calibration: could not compute text lengths: %s", e)
+ return lengths
+
+
+def build_report(manager) -> Dict[str, Any]:
+ """Compute the report for all evaluated schemas and write output files."""
+ config = manager.config
+ schema_infos = manager.get_schema_infos()
+ out_dir = config.output.dir
+ os.makedirs(out_dir, exist_ok=True)
+
+ # --- per-LLM labels file (all labeled items) ---
+ labels_path = os.path.join(out_dir, config.output.labels_file)
+ with open(labels_path, "w") as f:
+ for r in manager.store.all_results():
+ f.write(json.dumps(r.to_dict()) + "\n")
+
+ # --- restrict metrics to the human calibration sample if one was drawn ---
+ sample_ids = manager.phase.get_phase_data("calibration_sample")
+ restrict = set(sample_ids) if sample_ids else None
+
+ from potato.ai.judge import extract_labels
+
+ schema_reports = {}
+ for schema_info in schema_infos:
+ name = schema_info.get("name")
+ atype = schema_info.get("annotation_type", "radio")
+ valid_labels = extract_labels(schema_info)
+ if atype == "span":
+ llm_spans, human_spans = collect_span_inputs(
+ manager.store, schema_info, restrict_ids=restrict
+ )
+ schema_reports[name] = compute_span_report(
+ schema_name=name,
+ valid_labels=valid_labels,
+ llm_spans=llm_spans,
+ human_spans=human_spans,
+ gold_strategy=config.human.gold,
+ n_bins=config.calibration.n_bins,
+ instance_lengths=_instance_text_lengths(manager),
+ )
+ continue
+ llm_modal, llm_conf, human_labels = collect_metric_inputs(
+ manager.store, schema_info, restrict_ids=restrict
+ )
+ if atype == "multiselect":
+ schema_reports[name] = compute_multiselect_report(
+ schema_name=name,
+ valid_labels=valid_labels,
+ llm_modal=llm_modal,
+ llm_conf=llm_conf,
+ human_labels=human_labels,
+ gold_strategy=config.human.gold,
+ n_bins=config.calibration.n_bins,
+ )
+ else:
+ schema_reports[name] = compute_schema_report(
+ schema_name=name,
+ annotation_type=atype,
+ valid_labels=valid_labels,
+ llm_modal=llm_modal,
+ llm_conf=llm_conf,
+ human_labels=human_labels,
+ gold_strategy=config.human.gold,
+ n_bins=config.calibration.n_bins,
+ )
+
+ report = {
+ "generated_at": datetime.now().isoformat(),
+ "models": manager.store.models(),
+ "n_models": len(config.models),
+ "k_samples": config.k_samples,
+ "n_labeled_items": len(manager.store.labeled_instance_ids()),
+ "n_calibration_sample": len(restrict) if restrict else None,
+ "human": {"num_raters": config.human.num_raters, "gold": config.human.gold},
+ "schemas": schema_reports,
+ }
+
+ json_path = os.path.join(out_dir, config.output.report_json)
+ with open(json_path, "w") as f:
+ json.dump(report, f, indent=2)
+
+ html_path = os.path.join(out_dir, config.output.report_html)
+ with open(html_path, "w") as f:
+ f.write(render_html(report))
+
+ logger.info("judge_calibration: wrote report to %s", out_dir)
+ return report
+
+
+# ----- HTML rendering -----------------------------------------------------
+
+def render_html(report: Dict[str, Any]) -> str:
+ """Render a compact, self-contained HTML summary of the report."""
+ parts = [
+ " ",
+ " ",
+ "Judge Calibration Report ",
+ # Self-contained (portable/emailable) but aligned to Potato's brand:
+ # Outfit-first font stack, violet accent, accessible contrast, dark mode.
+ "",
+ "Judge Calibration Report ",
+ f"Generated {report.get('generated_at','')} ยท "
+ f"{report.get('n_models',0)} model(s), k={report.get('k_samples','?')} samples ยท "
+ f"{report.get('n_labeled_items',0)} items labeled"
+ + (f" ยท {report['n_calibration_sample']} in calibration sample"
+ if report.get('n_calibration_sample') else "")
+ + "
",
+ ]
+
+ for name, sr in (report.get("schemas") or {}).items():
+ parts.append(f"Schema: {name} ({sr.get('annotation_type','')}) ")
+ if sr.get("skipped"):
+ parts.append(f"{sr['skipped']}
")
+ continue
+ parts.append(f"{sr.get('n_gold',0)} human gold labels "
+ f"(gold = {sr.get('gold_strategy','single')})
")
+
+ # per-model metrics table
+ parts.append("Model Acc F1 (macro) "
+ "ECE Brier n ")
+ if sr.get("experimental"):
+ parts.append("โ Experimental โ span aggregation and "
+ "IoU matching are heuristic.
")
+
+ for model, m in (sr.get("per_model") or {}).items():
+ cal = m.get("calibration", {})
+ mae = f" ยท MAE {m['mae']}" if m.get("mae") is not None else ""
+ acc = m.get("accuracy", m.get("exact_match_accuracy", ""))
+ extra = f" ยท Jaccard {m['mean_jaccard']}" if m.get("mean_jaccard") is not None else ""
+ if m.get("mean_iou") is not None and "mean_jaccard" not in m:
+ extra = f" ยท IoU {m['mean_iou']}"
+ f1 = m.get("f1_macro", m.get("f1", ""))
+ n = m.get("n", m.get("n_instances", ""))
+ parts.append(
+ f"{model} "
+ f"{acc} "
+ f"{f1}{mae}{extra} "
+ f"{cal.get('ece','')} "
+ f"{cal.get('brier','')} "
+ f"{n} "
+ )
+ parts.append("
")
+
+ # IAA
+ iaa = sr.get("iaa", {})
+ parts.append("Agreement Value ")
+ if "span_f1" in iaa:
+ j = iaa["span_f1"]
+ parts.append(f"Span F1 (humanโLLM) {j.get('mean_human_llm')} ")
+ parts.append(f"Span F1 (LLMโLLM) {j.get('mean_llm_llm')} ")
+ parts.append(f"Span F1 (humanโhuman) {j.get('mean_human_human')} ")
+ tk = (iaa.get("token_kappa") or {}).get("cohen", {}) or {}
+ if tk:
+ parts.append(f"Token ฮบ (humanโLLM, chance-corrected) {tk.get('mean_human_llm')} ")
+ parts.append(f"Token ฮบ (LLMโLLM) {tk.get('mean_llm_llm')} ")
+ tkr = (iaa.get("token_kappa") or {}).get("krippendorff") or {}
+ if tkr:
+ parts.append(f"Token Krippendorff ฮฑ {tkr.get('alpha')} ")
+ g = iaa.get("gamma") or {}
+ if g and g.get("gamma") is not None:
+ parts.append(f"ฮณ (Gamma, approx.) overall {g.get('gamma')} ")
+ parts.append(f"ฮณ (humanโLLM) {g.get('mean_human_llm')} ")
+ elif "jaccard" in iaa:
+ j = iaa["jaccard"]
+ parts.append(f"Jaccard (humanโLLM) {j.get('mean_human_llm')} ")
+ parts.append(f"Jaccard (LLMโLLM) {j.get('mean_llm_llm')} ")
+ parts.append(f"Jaccard (humanโhuman) {j.get('mean_human_human')} ")
+ else:
+ cohen = iaa.get("cohen", {})
+ kripp = iaa.get("krippendorff") or {}
+ parts.append(f"Cohen ฮบ (humanโLLM) {cohen.get('mean_human_llm')} ")
+ parts.append(f"Cohen ฮบ (LLMโLLM) {cohen.get('mean_llm_llm')} ")
+ parts.append(f"Cohen ฮบ (humanโhuman) {cohen.get('mean_human_human')} ")
+ parts.append(f"Fleiss ฮบ (all raters) {iaa.get('fleiss',{}).get('kappa')} ")
+ parts.append(f"Krippendorff ฮฑ ({kripp.get('metric','')}) {kripp.get('alpha')} ")
+ parts.append("
")
+
+ parts.append("")
+ return "".join(parts)
diff --git a/potato/judge_calibration/routes.py b/potato/judge_calibration/routes.py
new file mode 100644
index 0000000000000000000000000000000000000000..77d1298abbc8ad89521b4ab53d5f5acb2bc19748
--- /dev/null
+++ b/potato/judge_calibration/routes.py
@@ -0,0 +1,153 @@
+"""
+Judge Calibration routes.
+
+Admin-gated endpoints for the calibration wizard:
+- GET /judge_calibration/admin -> wizard (prefilled from config)
+- POST /judge_calibration/run -> apply overrides + start generation
+- GET /judge_calibration/progress -> generation progress (JSON, polled)
+- POST /judge_calibration/report -> build the report
+- GET /judge_calibration/report -> rendered report.html
+- GET /judge_calibration/status -> status (JSON)
+
+Human blind-labeling happens through Potato's standard ``/annotate`` flow โ
+LLM labels live in a separate store and are never injected into the annotation
+UI, so blindness is structural (no special UI needed here).
+
+All endpoints require a valid admin API key (X-API-Key header or session),
+matching the solo_mode pattern; debug mode bypasses the check.
+"""
+
+import logging
+import os
+from functools import wraps
+
+from flask import Blueprint, Response, jsonify, render_template, request, session
+
+from potato.judge_calibration.manager import get_judge_calibration_manager
+
+logger = logging.getLogger(__name__)
+
+judge_calibration_bp = Blueprint("judge_calibration", __name__, url_prefix="/judge_calibration")
+
+
+def _enabled_required(f):
+ @wraps(f)
+ def wrapper(*args, **kwargs):
+ if get_judge_calibration_manager() is None:
+ return jsonify({"error": "Judge Calibration not enabled"}), 400
+ return f(*args, **kwargs)
+ return wrapper
+
+
+def admin_required(f):
+ """Require a valid admin API key (X-API-Key header or session)."""
+ @wraps(f)
+ def wrapper(*args, **kwargs):
+ from potato.server_utils.admin_key import validate_admin_api_key
+ from potato.flask_server import config as _config
+
+ api_key = request.headers.get("X-API-Key") or session.get("admin_api_key")
+ if not validate_admin_api_key(api_key, _config):
+ return jsonify({"error": "Admin authentication required"}), 403
+ return f(*args, **kwargs)
+ return wrapper
+
+
+def _config_for_wizard(cfg):
+ """Serialize the current config for prefilling the wizard form."""
+ return {
+ "prompt": cfg.prompt,
+ "k_samples": cfg.k_samples,
+ "max_items": cfg.max_items,
+ "fraction": cfg.fraction,
+ "models": [
+ {
+ "endpoint_type": m.endpoint_type,
+ "model": m.model,
+ "base_url": m.base_url,
+ "temperature": m.temperature,
+ }
+ for m in cfg.models
+ ],
+ "sampling": {
+ "strategy": cfg.sampling.strategy,
+ "stratify_by": cfg.sampling.stratify_by,
+ "sample_size": cfg.sampling.sample_size,
+ "seed": cfg.sampling.seed,
+ },
+ "human": {"num_raters": cfg.human.num_raters, "gold": cfg.human.gold},
+ "schemas": cfg.schemas,
+ "calibration": {"n_bins": cfg.calibration.n_bins},
+ }
+
+
+@judge_calibration_bp.route("/admin", methods=["GET"])
+@admin_required
+@_enabled_required
+def admin():
+ manager = get_judge_calibration_manager()
+ return render_template(
+ "judge_calibration/wizard.html",
+ config=_config_for_wizard(manager.config),
+ status=manager.get_status(),
+ )
+
+
+@judge_calibration_bp.route("/status", methods=["GET"])
+@admin_required
+@_enabled_required
+def status():
+ return jsonify(get_judge_calibration_manager().get_status())
+
+
+@judge_calibration_bp.route("/progress", methods=["GET"])
+@admin_required
+@_enabled_required
+def progress():
+ return jsonify(get_judge_calibration_manager().get_progress())
+
+
+@judge_calibration_bp.route("/run", methods=["POST"])
+@admin_required
+@_enabled_required
+def run():
+ manager = get_judge_calibration_manager()
+ overrides = request.get_json(silent=True) or {}
+ force = bool(overrides.pop("force_restart", False))
+ errors = manager.update_config(overrides)
+ if errors:
+ return jsonify({"error": "Invalid configuration", "errors": errors}), 400
+ try:
+ started = manager.start_generation(force_restart=force)
+ except ValueError as e:
+ return jsonify({"error": str(e)}), 400
+ if not started:
+ return jsonify({"error": "Generation already in progress"}), 409
+ return jsonify({"started": True, "progress": manager.get_progress()})
+
+
+@judge_calibration_bp.route("/report", methods=["POST"])
+@admin_required
+@_enabled_required
+def build_report():
+ manager = get_judge_calibration_manager()
+ if manager.is_generating():
+ return jsonify({"error": "Generation still in progress"}), 409
+ try:
+ report = manager.build_report()
+ except Exception as e:
+ logger.exception("judge_calibration: report build failed")
+ return jsonify({"error": str(e)}), 500
+ return jsonify({"built": True, "report": report})
+
+
+@judge_calibration_bp.route("/report", methods=["GET"])
+@admin_required
+@_enabled_required
+def view_report():
+ manager = get_judge_calibration_manager()
+ html_path = os.path.join(manager.config.output.dir, manager.config.output.report_html)
+ if os.path.exists(html_path):
+ with open(html_path) as f:
+ return Response(f.read(), mimetype="text/html")
+ return render_template("judge_calibration/status.html", status=manager.get_status())
diff --git a/potato/judge_calibration/sampler.py b/potato/judge_calibration/sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..e5a77eeac677187df155ed1533f2625955940cc4
--- /dev/null
+++ b/potato/judge_calibration/sampler.py
@@ -0,0 +1,80 @@
+"""
+Calibration-sample selection.
+
+Given the LLM-labeled instance ids, pick the subset that human(s) will
+*blind*-label for calibration. Mirrors the deterministic stratify+seed logic of
+``server_utils/overlap_sampler.py`` but returns a plain id list (we don't want
+the overlap mechanism's per-item annotator-cap stamping here).
+
+Strategies:
+- random : uniform random sample of ``sample_size`` ids
+- stratified : proportional allocation across strata defined by an item-data
+ field (or, for the calibration use-case, the modal LLM label)
+- all : every labeled id (ignores sample_size)
+"""
+
+import logging
+import random as _random
+from collections import defaultdict
+from typing import Any, Callable, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+def select_calibration_sample(
+ instance_ids: List[str],
+ sampling_cfg, # SamplingConfig
+ stratum_of: Optional[Callable[[str], Any]] = None,
+) -> List[str]:
+ """Return a deterministic subset of ``instance_ids`` to show humans.
+
+ Args:
+ instance_ids: candidate ids (the LLM-labeled items).
+ sampling_cfg: a SamplingConfig (strategy / sample_size / seed / stratify_by).
+ stratum_of: optional callable mapping an id -> a stratum key. Required
+ for the 'stratified' strategy; ignored otherwise. (The manager
+ wires this to either an item-data field or the modal LLM label.)
+ """
+ ids = sorted(set(instance_ids))
+ if not ids:
+ return []
+
+ strategy = sampling_cfg.strategy
+ if strategy == "all":
+ return ids
+
+ sample_size = min(sampling_cfg.sample_size, len(ids))
+ seed = sampling_cfg.seed
+
+ if strategy == "random" or stratum_of is None:
+ if strategy == "stratified" and stratum_of is None:
+ logger.warning("judge_calibration: stratified sampling requested but no "
+ "stratum mapping available; falling back to random")
+ rng = _random.Random(seed)
+ shuffled = list(ids)
+ rng.shuffle(shuffled)
+ return sorted(shuffled[:sample_size])
+
+ # stratified: proportional allocation, deterministic per-stratum shuffle.
+ strata: Dict[Any, List[str]] = defaultdict(list)
+ for iid in ids:
+ key = stratum_of(iid)
+ strata[key if key is not None else "__none__"].append(iid)
+
+ selected: List[str] = []
+ total = len(ids)
+ for key in sorted(strata.keys(), key=str):
+ bucket = sorted(strata[key])
+ rng_local = _random.Random(f"{seed}:{key}")
+ rng_local.shuffle(bucket)
+ # proportional quota (at least 1 per non-empty stratum)
+ quota = max(1, round(sample_size * len(bucket) / total))
+ selected.extend(bucket[:quota])
+
+ # Trim/pad to exactly sample_size deterministically.
+ selected = sorted(set(selected))
+ if len(selected) > sample_size:
+ rng = _random.Random(seed)
+ rng.shuffle(selected)
+ selected = sorted(selected[:sample_size])
+ return selected
diff --git a/potato/judge_calibration/storage.py b/potato/judge_calibration/storage.py
new file mode 100644
index 0000000000000000000000000000000000000000..3698311d4f8cbd9e167f67ca2e1930c4ae86b54d
--- /dev/null
+++ b/potato/judge_calibration/storage.py
@@ -0,0 +1,115 @@
+"""
+Persistence for Judge Calibration LLM results.
+
+LLM verdicts are stored in a dedicated JSON file under ``state_dir`` โ NOT as
+pseudo-users in the annotation store. This keeps them entirely out of the
+human annotation data path (guaranteeing humans never see them) and avoids
+polluting assignment/quota logic.
+
+The on-disk shape is a flat list of ``ModelItemResult`` dicts keyed implicitly
+by (model, instance_id, schema_name). ``ResultStore`` provides idempotent
+upsert so an interrupted GENERATING phase can resume without duplicating work.
+"""
+
+import json
+import logging
+import os
+import threading
+from typing import Dict, List, Optional, Tuple
+
+from potato.judge_calibration.aggregation import ModelItemResult
+
+logger = logging.getLogger(__name__)
+
+
+class ResultStore:
+ """In-memory store of ModelItemResults with atomic JSON persistence."""
+
+ _RESULTS_FILE = "llm_results.json"
+
+ def __init__(self, state_dir: Optional[str] = None):
+ self._lock = threading.RLock()
+ self.state_dir = state_dir
+ # (model, instance_id, schema_name) -> ModelItemResult
+ self._results: Dict[Tuple[str, str, str], ModelItemResult] = {}
+
+ @staticmethod
+ def _key(model: str, instance_id: str, schema_name: str) -> Tuple[str, str, str]:
+ return (model, instance_id, schema_name)
+
+ def upsert(self, result: ModelItemResult, save: bool = True) -> None:
+ with self._lock:
+ self._results[self._key(result.model, result.instance_id, result.schema_name)] = result
+ if save:
+ self._save()
+
+ def upsert_many(self, results: List[ModelItemResult]) -> None:
+ with self._lock:
+ for r in results:
+ self._results[self._key(r.model, r.instance_id, r.schema_name)] = r
+ self._save()
+
+ def has(self, model: str, instance_id: str, schema_name: str) -> bool:
+ with self._lock:
+ return self._key(model, instance_id, schema_name) in self._results
+
+ def get(self, model: str, instance_id: str, schema_name: str) -> Optional[ModelItemResult]:
+ with self._lock:
+ return self._results.get(self._key(model, instance_id, schema_name))
+
+ def all_results(self) -> List[ModelItemResult]:
+ with self._lock:
+ return list(self._results.values())
+
+ def models(self) -> List[str]:
+ with self._lock:
+ return sorted({r.model for r in self._results.values()})
+
+ def labeled_instance_ids(self) -> List[str]:
+ """Instance ids that have at least one model result."""
+ with self._lock:
+ return sorted({r.instance_id for r in self._results.values()})
+
+ def count(self) -> int:
+ with self._lock:
+ return len(self._results)
+
+ def clear(self) -> None:
+ with self._lock:
+ self._results = {}
+ self._save()
+
+ # ----- persistence ----------------------------------------------------
+
+ def _save(self) -> None:
+ if not self.state_dir:
+ return
+ try:
+ os.makedirs(self.state_dir, exist_ok=True)
+ path = os.path.join(self.state_dir, self._RESULTS_FILE)
+ tmp = path + ".tmp"
+ with open(tmp, "w") as f:
+ json.dump([r.to_dict() for r in self._results.values()], f)
+ os.replace(tmp, path)
+ except Exception as e:
+ logger.error("Error saving JC results: %s", e)
+
+ def load(self) -> bool:
+ if not self.state_dir:
+ return False
+ path = os.path.join(self.state_dir, self._RESULTS_FILE)
+ if not os.path.exists(path):
+ return False
+ try:
+ with open(path) as f:
+ data = json.load(f)
+ with self._lock:
+ self._results = {}
+ for d in data:
+ r = ModelItemResult.from_dict(d)
+ self._results[self._key(r.model, r.instance_id, r.schema_name)] = r
+ logger.info("Loaded %d JC results", len(self._results))
+ return True
+ except Exception as e:
+ logger.error("Error loading JC results: %s", e)
+ return False
diff --git a/potato/knowledge_base.py b/potato/knowledge_base.py
new file mode 100644
index 0000000000000000000000000000000000000000..07e70be2ac59e3967a7b8c82c529e825a4043522
--- /dev/null
+++ b/potato/knowledge_base.py
@@ -0,0 +1,962 @@
+"""
+Knowledge Base Client Module
+
+This module provides clients for querying external knowledge bases to support
+entity linking in span annotations. Supported knowledge bases include:
+
+- Wikidata: Open knowledge graph with millions of entities
+- UMLS: Unified Medical Language System (requires API key)
+- Custom REST APIs: Generic interface for custom knowledge bases
+
+Usage:
+ from potato.knowledge_base import get_kb_client, KnowledgeBaseConfig
+
+ # Configure and get a client
+ config = KnowledgeBaseConfig(
+ name="wikidata",
+ kb_type="wikidata",
+ language="en"
+ )
+ client = get_kb_client(config)
+
+ # Search for entities
+ results = client.search("Einstein", limit=10)
+
+ # Get entity details
+ entity = client.get_entity("Q937")
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Any
+import logging
+import json
+import urllib.parse
+
+logger = logging.getLogger(__name__)
+
+# Try to import requests, but make it optional
+try:
+ import requests
+ REQUESTS_AVAILABLE = True
+except ImportError:
+ REQUESTS_AVAILABLE = False
+ logger.warning("requests library not available. KB clients will not function.")
+
+
+@dataclass
+class KnowledgeBaseConfig:
+ """
+ Configuration for a knowledge base client.
+
+ Attributes:
+ name: Unique identifier for this KB configuration
+ kb_type: Type of knowledge base ("wikidata", "umls", "rest")
+ api_key: Optional API key for authenticated services
+ base_url: Base URL for REST APIs
+ language: Language code for results (default: "en")
+ timeout: Request timeout in seconds
+ extra_params: Additional parameters for the API
+ """
+ name: str
+ kb_type: str
+ api_key: Optional[str] = None
+ base_url: Optional[str] = None
+ language: str = "en"
+ timeout: int = 10
+ extra_params: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class KBEntity:
+ """
+ Represents an entity from a knowledge base.
+
+ Attributes:
+ entity_id: Unique identifier in the KB (e.g., "Q937" for Wikidata)
+ kb_source: Name of the knowledge base (e.g., "wikidata", "umls")
+ label: Human-readable label/name
+ description: Short description of the entity
+ aliases: Alternative names for the entity
+ entity_type: Type/class of the entity (if available)
+ url: URL to the entity page in the KB
+ extra_data: Additional data from the KB
+ """
+ entity_id: str
+ kb_source: str
+ label: str
+ description: str = ""
+ aliases: List[str] = field(default_factory=list)
+ entity_type: Optional[str] = None
+ url: Optional[str] = None
+ extra_data: Dict[str, Any] = field(default_factory=dict)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to dictionary for JSON serialization."""
+ return {
+ "entity_id": self.entity_id,
+ "kb_source": self.kb_source,
+ "label": self.label,
+ "description": self.description,
+ "aliases": self.aliases,
+ "entity_type": self.entity_type,
+ "url": self.url,
+ "extra_data": self.extra_data
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "KBEntity":
+ """Create entity from dictionary."""
+ return cls(
+ entity_id=data.get("entity_id", ""),
+ kb_source=data.get("kb_source", ""),
+ label=data.get("label", ""),
+ description=data.get("description", ""),
+ aliases=data.get("aliases", []),
+ entity_type=data.get("entity_type"),
+ url=data.get("url"),
+ extra_data=data.get("extra_data", {})
+ )
+
+
+class KnowledgeBaseClient(ABC):
+ """
+ Abstract base class for knowledge base clients.
+
+ Subclasses must implement search() and get_entity() methods.
+ """
+
+ def __init__(self, config: KnowledgeBaseConfig):
+ """
+ Initialize the KB client with configuration.
+
+ Args:
+ config: KnowledgeBaseConfig with connection settings
+ """
+ self.config = config
+ self.name = config.name
+
+ @abstractmethod
+ def search(self, query: str, limit: int = 10, entity_type: Optional[str] = None) -> List[KBEntity]:
+ """
+ Search the knowledge base for entities matching the query.
+
+ Args:
+ query: Search query string
+ limit: Maximum number of results to return
+ entity_type: Optional filter for entity type
+
+ Returns:
+ List of KBEntity objects matching the query
+ """
+ pass
+
+ @abstractmethod
+ def get_entity(self, entity_id: str) -> Optional[KBEntity]:
+ """
+ Get detailed information about a specific entity.
+
+ Args:
+ entity_id: The unique identifier for the entity
+
+ Returns:
+ KBEntity object if found, None otherwise
+ """
+ pass
+
+ def is_available(self) -> bool:
+ """
+ Check if the KB service is available.
+
+ Returns:
+ True if the service can be reached, False otherwise
+ """
+ return REQUESTS_AVAILABLE
+
+
+class WikidataClient(KnowledgeBaseClient):
+ """
+ Client for querying Wikidata, the free knowledge base.
+
+ Wikidata provides millions of entities across all domains with
+ multilingual labels, descriptions, and structured data.
+
+ API documentation: https://www.wikidata.org/w/api.php
+ """
+
+ WIKIDATA_API = "https://www.wikidata.org/w/api.php"
+ WIKIDATA_ENTITY_URL = "https://www.wikidata.org/wiki/"
+
+ # User-Agent header required by Wikimedia API policy
+ # https://meta.wikimedia.org/wiki/User-Agent_policy
+ HEADERS = {
+ "User-Agent": "Potato-Annotation-Tool/1.0 (https://github.com/davidjurgens/potato; annotation@example.com)"
+ }
+
+ def search(self, query: str, limit: int = 10, entity_type: Optional[str] = None) -> List[KBEntity]:
+ """
+ Search Wikidata for entities matching the query.
+
+ Args:
+ query: Search query string
+ limit: Maximum number of results (max 50)
+ entity_type: Optional Wikidata item type (e.g., "Q5" for humans)
+
+ Returns:
+ List of KBEntity objects
+ """
+ if not REQUESTS_AVAILABLE:
+ logger.error("requests library not available for Wikidata search")
+ return []
+
+ if not query or not query.strip():
+ return []
+
+ params = {
+ "action": "wbsearchentities",
+ "format": "json",
+ "language": self.config.language,
+ "uselang": self.config.language,
+ "search": query.strip(),
+ "limit": min(limit, 50), # Wikidata max is 50
+ "type": "item"
+ }
+
+ try:
+ response = requests.get(
+ self.WIKIDATA_API,
+ params=params,
+ headers=self.HEADERS,
+ timeout=self.config.timeout
+ )
+ response.raise_for_status()
+ data = response.json()
+
+ results = []
+ for item in data.get("search", []):
+ entity = KBEntity(
+ entity_id=item.get("id", ""),
+ kb_source="wikidata",
+ label=item.get("label", ""),
+ description=item.get("description", ""),
+ aliases=item.get("aliases", []),
+ url=f"{self.WIKIDATA_ENTITY_URL}{item.get('id', '')}"
+ )
+ results.append(entity)
+
+ # Filter by entity_type if specified (requires additional API call)
+ if entity_type and results:
+ results = self._filter_by_type(results, entity_type)
+
+ return results
+
+ except requests.RequestException as e:
+ logger.error(f"Wikidata search error: {e}")
+ return []
+ except json.JSONDecodeError as e:
+ logger.error(f"Wikidata JSON parse error: {e}")
+ return []
+
+ def _filter_by_type(self, entities: List[KBEntity], entity_type: str) -> List[KBEntity]:
+ """
+ Filter entities by their instance_of (P31) property.
+
+ Args:
+ entities: List of entities to filter
+ entity_type: Wikidata item ID for the type (e.g., "Q5" for human)
+
+ Returns:
+ Filtered list of entities
+ """
+ if not entities:
+ return []
+
+ # Get entity IDs
+ entity_ids = [e.entity_id for e in entities]
+
+ # Query for P31 (instance of) claims
+ params = {
+ "action": "wbgetentities",
+ "format": "json",
+ "ids": "|".join(entity_ids),
+ "props": "claims",
+ "languages": self.config.language
+ }
+
+ try:
+ response = requests.get(
+ self.WIKIDATA_API,
+ params=params,
+ headers=self.HEADERS,
+ timeout=self.config.timeout
+ )
+ response.raise_for_status()
+ data = response.json()
+
+ # Build set of matching entity IDs
+ matching_ids = set()
+ for eid, entity_data in data.get("entities", {}).items():
+ claims = entity_data.get("claims", {})
+ p31_claims = claims.get("P31", [])
+
+ for claim in p31_claims:
+ mainsnak = claim.get("mainsnak", {})
+ datavalue = mainsnak.get("datavalue", {})
+ value = datavalue.get("value", {})
+ if value.get("id") == entity_type:
+ matching_ids.add(eid)
+ break
+
+ # Filter entities
+ return [e for e in entities if e.entity_id in matching_ids]
+
+ except requests.RequestException as e:
+ logger.warning(f"Wikidata type filter error: {e}")
+ return entities # Return unfiltered on error
+
+ def get_entity(self, entity_id: str) -> Optional[KBEntity]:
+ """
+ Get detailed information about a Wikidata entity.
+
+ Args:
+ entity_id: Wikidata item ID (e.g., "Q937")
+
+ Returns:
+ KBEntity with full details, or None if not found
+ """
+ if not REQUESTS_AVAILABLE:
+ logger.error("requests library not available for Wikidata get_entity")
+ return None
+
+ if not entity_id:
+ return None
+
+ params = {
+ "action": "wbgetentities",
+ "format": "json",
+ "ids": entity_id,
+ "props": "labels|descriptions|aliases|claims|sitelinks",
+ "languages": self.config.language
+ }
+
+ try:
+ response = requests.get(
+ self.WIKIDATA_API,
+ params=params,
+ headers=self.HEADERS,
+ timeout=self.config.timeout
+ )
+ response.raise_for_status()
+ data = response.json()
+
+ entities = data.get("entities", {})
+ if entity_id not in entities or "missing" in entities.get(entity_id, {}):
+ return None
+
+ entity_data = entities[entity_id]
+
+ # Extract label
+ labels = entity_data.get("labels", {})
+ label = labels.get(self.config.language, {}).get("value", entity_id)
+
+ # Extract description
+ descriptions = entity_data.get("descriptions", {})
+ description = descriptions.get(self.config.language, {}).get("value", "")
+
+ # Extract aliases
+ aliases_data = entity_data.get("aliases", {})
+ aliases = [a.get("value", "") for a in aliases_data.get(self.config.language, [])]
+
+ # Extract type from P31 (instance of)
+ entity_type = None
+ claims = entity_data.get("claims", {})
+ p31_claims = claims.get("P31", [])
+ if p31_claims:
+ first_claim = p31_claims[0]
+ mainsnak = first_claim.get("mainsnak", {})
+ datavalue = mainsnak.get("datavalue", {})
+ value = datavalue.get("value", {})
+ entity_type = value.get("id")
+
+ # Get Wikipedia URL if available
+ sitelinks = entity_data.get("sitelinks", {})
+ wiki_key = f"{self.config.language}wiki"
+ url = None
+ if wiki_key in sitelinks:
+ wiki_title = sitelinks[wiki_key].get("title", "")
+ if wiki_title:
+ url = f"https://{self.config.language}.wikipedia.org/wiki/{urllib.parse.quote(wiki_title.replace(' ', '_'), safe='')}"
+
+ if not url:
+ url = f"{self.WIKIDATA_ENTITY_URL}{entity_id}"
+
+ return KBEntity(
+ entity_id=entity_id,
+ kb_source="wikidata",
+ label=label,
+ description=description,
+ aliases=aliases,
+ entity_type=entity_type,
+ url=url,
+ extra_data={
+ "claims_count": len(claims),
+ "sitelinks_count": len(sitelinks)
+ }
+ )
+
+ except requests.RequestException as e:
+ logger.error(f"Wikidata get_entity error: {e}")
+ return None
+ except json.JSONDecodeError as e:
+ logger.error(f"Wikidata JSON parse error: {e}")
+ return None
+
+
+class UMLSClient(KnowledgeBaseClient):
+ """
+ Client for querying UMLS (Unified Medical Language System).
+
+ UMLS is a comprehensive medical terminology database maintained by
+ the National Library of Medicine. Requires an API key.
+
+ API documentation: https://documentation.uts.nlm.nih.gov/rest/home.html
+ """
+
+ UMLS_API = "https://uts-ws.nlm.nih.gov/rest"
+
+ def __init__(self, config: KnowledgeBaseConfig):
+ super().__init__(config)
+ if not config.api_key:
+ logger.warning("UMLS client initialized without API key. Searches will fail.")
+
+ def search(self, query: str, limit: int = 10, entity_type: Optional[str] = None) -> List[KBEntity]:
+ """
+ Search UMLS for concepts matching the query.
+
+ Args:
+ query: Search query string
+ limit: Maximum number of results
+ entity_type: Optional semantic type filter (e.g., "Disease or Syndrome")
+
+ Returns:
+ List of KBEntity objects
+ """
+ if not REQUESTS_AVAILABLE:
+ logger.error("requests library not available for UMLS search")
+ return []
+
+ if not self.config.api_key:
+ logger.error("UMLS API key required")
+ return []
+
+ if not query or not query.strip():
+ return []
+
+ params = {
+ "apiKey": self.config.api_key,
+ "string": query.strip(),
+ "pageSize": min(limit, 50),
+ "returnIdType": "code"
+ }
+
+ if entity_type:
+ params["sabs"] = entity_type
+
+ try:
+ response = requests.get(
+ f"{self.UMLS_API}/search/current",
+ params=params,
+ timeout=self.config.timeout
+ )
+ response.raise_for_status()
+ data = response.json()
+
+ results = []
+ result_data = data.get("result", {})
+ for item in result_data.get("results", []):
+ cui = item.get("ui", "")
+ entity = KBEntity(
+ entity_id=cui,
+ kb_source="umls",
+ label=item.get("name", ""),
+ description=item.get("rootSource", ""),
+ url=f"https://uts.nlm.nih.gov/uts/umls/concept/{cui}",
+ extra_data={
+ "root_source": item.get("rootSource", ""),
+ "uri": item.get("uri", "")
+ }
+ )
+ results.append(entity)
+
+ return results
+
+ except requests.RequestException as e:
+ logger.error(f"UMLS search error: {e}")
+ return []
+ except json.JSONDecodeError as e:
+ logger.error(f"UMLS JSON parse error: {e}")
+ return []
+
+ def get_entity(self, entity_id: str) -> Optional[KBEntity]:
+ """
+ Get detailed information about a UMLS concept.
+
+ Args:
+ entity_id: UMLS CUI (Concept Unique Identifier)
+
+ Returns:
+ KBEntity with concept details, or None if not found
+ """
+ if not REQUESTS_AVAILABLE:
+ logger.error("requests library not available for UMLS get_entity")
+ return None
+
+ if not self.config.api_key:
+ logger.error("UMLS API key required")
+ return None
+
+ if not entity_id:
+ return None
+
+ try:
+ # Get concept info
+ response = requests.get(
+ f"{self.UMLS_API}/content/current/CUI/{entity_id}",
+ params={"apiKey": self.config.api_key},
+ timeout=self.config.timeout
+ )
+ response.raise_for_status()
+ data = response.json()
+
+ result = data.get("result", {})
+ if not result:
+ return None
+
+ # Get semantic types
+ semantic_types = []
+ for st in result.get("semanticTypes", []):
+ semantic_types.append(st.get("name", ""))
+
+ # Get definitions โ validate URL to prevent SSRF via API response
+ definitions_url = result.get("definitions", "")
+ description = ""
+ if definitions_url:
+ # Only follow definitions URLs pointing to the legitimate UMLS API
+ UMLS_API_BASE = "https://uts-ws.nlm.nih.gov/"
+ if not definitions_url.startswith(UMLS_API_BASE):
+ logger.warning(
+ f"Ignoring definitions URL from non-UMLS domain: "
+ f"{definitions_url[:100]}"
+ )
+ definitions_url = ""
+
+ if definitions_url:
+ try:
+ def_response = requests.get(
+ definitions_url,
+ params={"apiKey": self.config.api_key},
+ timeout=self.config.timeout
+ )
+ def_response.raise_for_status()
+ def_data = def_response.json()
+ defs = def_data.get("result", [])
+ if defs:
+ description = defs[0].get("value", "")
+ except requests.RequestException:
+ pass
+
+ return KBEntity(
+ entity_id=entity_id,
+ kb_source="umls",
+ label=result.get("name", ""),
+ description=description,
+ entity_type=semantic_types[0] if semantic_types else None,
+ url=f"https://uts.nlm.nih.gov/uts/umls/concept/{entity_id}",
+ extra_data={
+ "semantic_types": semantic_types,
+ "atom_count": result.get("atomCount", 0),
+ "relation_count": result.get("relationCount", 0)
+ }
+ )
+
+ except requests.RequestException as e:
+ logger.error(f"UMLS get_entity error: {e}")
+ return None
+ except json.JSONDecodeError as e:
+ logger.error(f"UMLS JSON parse error: {e}")
+ return None
+
+
+class RESTClient(KnowledgeBaseClient):
+ """
+ Generic REST API client for custom knowledge bases.
+
+ Supports configurable endpoints for search and entity lookup.
+ Expected response format can be customized via configuration.
+
+ Configuration example:
+ KnowledgeBaseConfig(
+ name="my_kb",
+ kb_type="rest",
+ base_url="https://api.example.com",
+ extra_params={
+ "search_endpoint": "/search",
+ "entity_endpoint": "/entity/{entity_id}",
+ "search_query_param": "q",
+ "results_path": "data.results",
+ "entity_id_field": "id",
+ "label_field": "name",
+ "description_field": "description"
+ }
+ )
+ """
+
+ def __init__(self, config: KnowledgeBaseConfig):
+ super().__init__(config)
+ if not config.base_url:
+ raise ValueError("REST client requires base_url in configuration")
+
+ # Default field mappings
+ self.search_endpoint = config.extra_params.get("search_endpoint", "/search")
+ self.entity_endpoint = config.extra_params.get("entity_endpoint", "/entity/{entity_id}")
+ self.search_query_param = config.extra_params.get("search_query_param", "q")
+ self.results_path = config.extra_params.get("results_path", "results")
+ self.entity_id_field = config.extra_params.get("entity_id_field", "id")
+ self.label_field = config.extra_params.get("label_field", "label")
+ self.description_field = config.extra_params.get("description_field", "description")
+ self.aliases_field = config.extra_params.get("aliases_field", "aliases")
+ self.type_field = config.extra_params.get("type_field", "type")
+ self.url_field = config.extra_params.get("url_field", "url")
+
+ def _get_nested_value(self, data: Dict, path: str) -> Any:
+ """
+ Get a nested value from a dictionary using dot notation.
+
+ Args:
+ data: Dictionary to search
+ path: Dot-separated path (e.g., "data.results")
+
+ Returns:
+ Value at the path, or None if not found
+ """
+ parts = path.split(".")
+ value = data
+ for part in parts:
+ if isinstance(value, dict) and part in value:
+ value = value[part]
+ else:
+ return None
+ return value
+
+ def search(self, query: str, limit: int = 10, entity_type: Optional[str] = None) -> List[KBEntity]:
+ """
+ Search the REST API for entities.
+
+ Args:
+ query: Search query string
+ limit: Maximum number of results
+ entity_type: Optional entity type filter
+
+ Returns:
+ List of KBEntity objects
+ """
+ if not REQUESTS_AVAILABLE:
+ logger.error("requests library not available for REST search")
+ return []
+
+ if not query or not query.strip():
+ return []
+
+ url = f"{self.config.base_url.rstrip('/')}{self.search_endpoint}"
+ params = {
+ self.search_query_param: query.strip(),
+ "limit": limit
+ }
+
+ if entity_type:
+ params["type"] = entity_type
+
+ if self.config.api_key:
+ params["api_key"] = self.config.api_key
+
+ # Add any extra parameters
+ for key, value in self.config.extra_params.items():
+ if key.startswith("param_"):
+ params[key[6:]] = value
+
+ try:
+ response = requests.get(
+ url,
+ params=params,
+ timeout=self.config.timeout
+ )
+ response.raise_for_status()
+ data = response.json()
+
+ # Extract results using configured path
+ results_data = self._get_nested_value(data, self.results_path)
+ if not results_data or not isinstance(results_data, list):
+ return []
+
+ results = []
+ for item in results_data[:limit]:
+ entity = KBEntity(
+ entity_id=str(item.get(self.entity_id_field, "")),
+ kb_source=self.config.name,
+ label=item.get(self.label_field, ""),
+ description=item.get(self.description_field, ""),
+ aliases=item.get(self.aliases_field, []) or [],
+ entity_type=item.get(self.type_field),
+ url=item.get(self.url_field)
+ )
+ results.append(entity)
+
+ return results
+
+ except requests.RequestException as e:
+ logger.error(f"REST search error: {e}")
+ return []
+ except json.JSONDecodeError as e:
+ logger.error(f"REST JSON parse error: {e}")
+ return []
+
+ def get_entity(self, entity_id: str) -> Optional[KBEntity]:
+ """
+ Get entity details from the REST API.
+
+ Args:
+ entity_id: Entity identifier
+
+ Returns:
+ KBEntity if found, None otherwise
+ """
+ if not REQUESTS_AVAILABLE:
+ logger.error("requests library not available for REST get_entity")
+ return None
+
+ if not entity_id:
+ return None
+
+ endpoint = self.entity_endpoint.replace("{entity_id}", entity_id)
+ url = f"{self.config.base_url.rstrip('/')}{endpoint}"
+
+ params = {}
+ if self.config.api_key:
+ params["api_key"] = self.config.api_key
+
+ try:
+ response = requests.get(
+ url,
+ params=params,
+ timeout=self.config.timeout
+ )
+ response.raise_for_status()
+ item = response.json()
+
+ return KBEntity(
+ entity_id=str(item.get(self.entity_id_field, entity_id)),
+ kb_source=self.config.name,
+ label=item.get(self.label_field, ""),
+ description=item.get(self.description_field, ""),
+ aliases=item.get(self.aliases_field, []) or [],
+ entity_type=item.get(self.type_field),
+ url=item.get(self.url_field)
+ )
+
+ except requests.RequestException as e:
+ logger.error(f"REST get_entity error: {e}")
+ return None
+ except json.JSONDecodeError as e:
+ logger.error(f"REST JSON parse error: {e}")
+ return None
+
+
+# Registry of available KB clients
+KB_CLIENT_REGISTRY: Dict[str, type] = {
+ "wikidata": WikidataClient,
+ "umls": UMLSClient,
+ "rest": RESTClient
+}
+
+
+def get_kb_client(config: KnowledgeBaseConfig) -> KnowledgeBaseClient:
+ """
+ Factory function to create a KB client based on configuration.
+
+ Args:
+ config: KnowledgeBaseConfig specifying the KB type and settings
+
+ Returns:
+ KnowledgeBaseClient instance
+
+ Raises:
+ ValueError: If kb_type is not supported
+ """
+ kb_type = config.kb_type.lower()
+ if kb_type not in KB_CLIENT_REGISTRY:
+ supported = ", ".join(KB_CLIENT_REGISTRY.keys())
+ raise ValueError(f"Unsupported KB type '{kb_type}'. Supported types: {supported}")
+
+ client_class = KB_CLIENT_REGISTRY[kb_type]
+ return client_class(config)
+
+
+def register_kb_client(kb_type: str, client_class: type):
+ """
+ Register a custom KB client class.
+
+ Args:
+ kb_type: Type identifier for the client
+ client_class: KnowledgeBaseClient subclass
+ """
+ if not issubclass(client_class, KnowledgeBaseClient):
+ raise TypeError("client_class must be a subclass of KnowledgeBaseClient")
+ KB_CLIENT_REGISTRY[kb_type.lower()] = client_class
+
+
+# KB manager for handling multiple configured knowledge bases
+class KnowledgeBaseManager:
+ """
+ Manager for multiple knowledge base configurations.
+
+ Maintains a registry of configured KBs and provides unified
+ search across multiple sources.
+ """
+
+ def __init__(self):
+ self._clients: Dict[str, KnowledgeBaseClient] = {}
+ self._configs: Dict[str, KnowledgeBaseConfig] = {}
+
+ def configure_from_yaml(self, kb_config: Dict[str, Any]) -> None:
+ """
+ Configure knowledge bases from YAML configuration.
+
+ Expected format:
+ entity_linking:
+ enabled: true
+ knowledge_bases:
+ - name: "wikidata"
+ type: "wikidata"
+ language: "en"
+ - name: "umls"
+ type: "umls"
+ api_key: "${UMLS_API_KEY}"
+
+ Args:
+ kb_config: Dictionary from YAML configuration
+ """
+ if not kb_config.get("enabled", False):
+ return
+
+ for kb in kb_config.get("knowledge_bases", []):
+ config = KnowledgeBaseConfig(
+ name=kb.get("name", ""),
+ kb_type=kb.get("type", ""),
+ api_key=kb.get("api_key"),
+ base_url=kb.get("base_url"),
+ language=kb.get("language", "en"),
+ timeout=kb.get("timeout", 10),
+ extra_params=kb.get("extra_params", {})
+ )
+
+ try:
+ client = get_kb_client(config)
+ self._clients[config.name] = client
+ self._configs[config.name] = config
+ logger.info(f"Configured KB client: {config.name} ({config.kb_type})")
+ except Exception as e:
+ logger.error(f"Failed to configure KB '{config.name}': {e}")
+
+ def add_client(self, name: str, client: KnowledgeBaseClient) -> None:
+ """Add a pre-configured client."""
+ self._clients[name] = client
+
+ def get_client(self, name: str) -> Optional[KnowledgeBaseClient]:
+ """Get a configured KB client by name."""
+ return self._clients.get(name)
+
+ def get_config(self, name: str) -> Optional[KnowledgeBaseConfig]:
+ """Get KB configuration by name."""
+ return self._configs.get(name)
+
+ def list_clients(self) -> List[str]:
+ """List names of all configured KB clients."""
+ return list(self._clients.keys())
+
+ def search_all(self, query: str, limit: int = 10) -> Dict[str, List[KBEntity]]:
+ """
+ Search all configured knowledge bases.
+
+ Args:
+ query: Search query
+ limit: Max results per KB
+
+ Returns:
+ Dictionary mapping KB name to list of results
+ """
+ results = {}
+ for name, client in self._clients.items():
+ try:
+ results[name] = client.search(query, limit=limit)
+ except Exception as e:
+ logger.error(f"Search failed for KB '{name}': {e}")
+ results[name] = []
+ return results
+
+ def search(self, query: str, kb_name: str, limit: int = 10) -> List[KBEntity]:
+ """
+ Search a specific knowledge base.
+
+ Args:
+ query: Search query
+ kb_name: Name of the KB to search
+ limit: Max results
+
+ Returns:
+ List of KBEntity results
+ """
+ client = self.get_client(kb_name)
+ if not client:
+ logger.warning(f"KB '{kb_name}' not configured")
+ return []
+
+ try:
+ return client.search(query, limit=limit)
+ except Exception as e:
+ logger.error(f"Search failed for KB '{kb_name}': {e}")
+ return []
+
+
+# Global KB manager instance
+_kb_manager: Optional[KnowledgeBaseManager] = None
+
+
+def get_kb_manager() -> KnowledgeBaseManager:
+ """Get or create the global KB manager."""
+ global _kb_manager
+ if _kb_manager is None:
+ _kb_manager = KnowledgeBaseManager()
+ return _kb_manager
+
+
+def init_kb_manager(config: Dict[str, Any]) -> KnowledgeBaseManager:
+ """
+ Initialize the KB manager from configuration.
+
+ Args:
+ config: Full application config dictionary
+
+ Returns:
+ Configured KnowledgeBaseManager
+ """
+ global _kb_manager
+ _kb_manager = KnowledgeBaseManager()
+
+ # Look for entity_linking config in annotation_schemes
+ for scheme in config.get("annotation_schemes", []):
+ if scheme.get("annotation_type") == "span":
+ entity_linking = scheme.get("entity_linking", {})
+ if entity_linking:
+ _kb_manager.configure_from_yaml(entity_linking)
+
+ return _kb_manager
diff --git a/potato/logging_config.py b/potato/logging_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..4ba07d63810513283e101f1ce73c431039a355cc
--- /dev/null
+++ b/potato/logging_config.py
@@ -0,0 +1,265 @@
+"""
+Centralized Logging Configuration for Potato
+
+This module provides a unified logging configuration for the entire Potato
+annotation platform. It ensures consistent log formatting, appropriate log
+levels, and optional file logging across all modules.
+
+Usage:
+ from potato.logging_config import setup_logging, get_logger
+
+ # At application startup (in flask_server.py):
+ setup_logging(verbose=config.get('verbose'), debug=config.get('debug'))
+
+ # In any module:
+ logger = get_logger(__name__)
+ logger.info("Something happened")
+
+Debug Log Modes:
+ --debug-log=all Enable debug logging for both UI and server
+ --debug-log=ui Enable debug logging for UI/frontend only
+ --debug-log=server Enable debug logging for server/backend only
+ --debug-log=none Disable all debug logging
+"""
+
+import logging
+import os
+import sys
+from logging.handlers import RotatingFileHandler
+from typing import Optional
+
+
+# Default log format
+DEFAULT_FORMAT = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+DEFAULT_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+# Verbose format includes more details
+VERBOSE_FORMAT = "%(asctime)s [%(levelname)s] %(name)s (%(filename)s:%(lineno)d): %(message)s"
+
+# Module loggers that should be configured
+POTATO_LOGGERS = [
+ "potato",
+ "potato.flask_server",
+ "potato.routes",
+ "potato.admin",
+ "potato.authentication",
+ "potato.user_state_management",
+ "potato.item_state_management",
+ "potato.active_learning_manager",
+ "potato.directory_watcher",
+ "potato.agreement",
+ "potato.ai",
+ "potato.ai.ai_endpoint",
+ "potato.ai.icl_labeler",
+ "potato.server_utils",
+ "potato.server_utils.config_module",
+ "potato.server_utils.front_end",
+ "potato.database",
+]
+
+# Track if logging has been set up
+_logging_initialized = False
+
+# Track debug log settings for UI
+_ui_debug_enabled = False
+_server_debug_enabled = False
+
+
+def setup_logging(
+ verbose: bool = False,
+ debug: bool = False,
+ debug_log: Optional[str] = None,
+ log_file: Optional[str] = None,
+ log_dir: Optional[str] = None,
+ max_bytes: int = 10 * 1024 * 1024, # 10 MB
+ backup_count: int = 5,
+) -> None:
+ """
+ Set up logging for the entire Potato application.
+
+ This function configures the root logger and all Potato module loggers
+ with consistent formatting and appropriate log levels.
+
+ Args:
+ verbose: If True, set log level to DEBUG and use verbose format
+ debug: If True, set log level to DEBUG (same as verbose)
+ debug_log: Selective debug logging mode:
+ - 'all': Enable debug for both UI and server
+ - 'ui': Enable debug for UI/frontend only
+ - 'server': Enable debug for server/backend only
+ - 'none': Disable all debug logging
+ - None: Use verbose/debug flags as before
+ log_file: Optional path to a log file. If provided, logs will be
+ written to this file in addition to console.
+ log_dir: Optional directory for log files. If log_file is not provided
+ but log_dir is, a default log file will be created there.
+ max_bytes: Maximum size of each log file before rotation (default 10MB)
+ backup_count: Number of backup log files to keep (default 5)
+ """
+ global _logging_initialized, _ui_debug_enabled, _server_debug_enabled
+
+ # Handle selective debug logging
+ if debug_log:
+ if debug_log == 'all':
+ _ui_debug_enabled = True
+ _server_debug_enabled = True
+ elif debug_log == 'ui':
+ _ui_debug_enabled = True
+ _server_debug_enabled = False
+ elif debug_log == 'server':
+ _ui_debug_enabled = False
+ _server_debug_enabled = True
+ elif debug_log == 'none':
+ _ui_debug_enabled = False
+ _server_debug_enabled = False
+ else:
+ # Default behavior based on debug/verbose flags
+ _ui_debug_enabled = debug or verbose
+ _server_debug_enabled = debug or verbose
+
+ # Determine server log level
+ if _server_debug_enabled:
+ log_level = logging.DEBUG
+ log_format = VERBOSE_FORMAT
+ elif verbose or debug:
+ # If debug_log explicitly disabled server but debug flag is on,
+ # still use INFO level for server
+ log_level = logging.INFO
+ log_format = DEFAULT_FORMAT
+ else:
+ log_level = logging.INFO
+ log_format = DEFAULT_FORMAT
+
+ # Create formatter
+ formatter = logging.Formatter(log_format, datefmt=DEFAULT_DATE_FORMAT)
+
+ # Configure root logger
+ root_logger = logging.getLogger()
+ root_logger.setLevel(log_level)
+
+ # Remove existing handlers to avoid duplicates on reinitialization
+ for handler in root_logger.handlers[:]:
+ root_logger.removeHandler(handler)
+
+ # Add console handler
+ console_handler = logging.StreamHandler(sys.stdout)
+ console_handler.setLevel(log_level)
+ console_handler.setFormatter(formatter)
+ root_logger.addHandler(console_handler)
+
+ # Add file handler if requested
+ if log_file or log_dir:
+ if not log_file and log_dir:
+ os.makedirs(log_dir, exist_ok=True)
+ log_file = os.path.join(log_dir, "potato.log")
+
+ file_handler = RotatingFileHandler(
+ log_file,
+ maxBytes=max_bytes,
+ backupCount=backup_count,
+ )
+ file_handler.setLevel(log_level)
+ file_handler.setFormatter(formatter)
+ root_logger.addHandler(file_handler)
+
+ # Configure all Potato loggers
+ for logger_name in POTATO_LOGGERS:
+ module_logger = logging.getLogger(logger_name)
+ module_logger.setLevel(log_level)
+ # Don't add handlers - they inherit from root
+ module_logger.propagate = True
+
+ # Reduce noise from third-party libraries
+ logging.getLogger("werkzeug").setLevel(logging.WARNING)
+ logging.getLogger("urllib3").setLevel(logging.WARNING)
+
+ _logging_initialized = True
+
+ # Log that logging has been configured
+ logger = logging.getLogger("potato")
+ logger.debug(f"Logging initialized with level={logging.getLevelName(log_level)}")
+
+
+def get_logger(name: str) -> logging.Logger:
+ """
+ Get a logger for the specified module.
+
+ This is a convenience function that ensures the logger is properly
+ configured even if setup_logging() hasn't been called yet.
+
+ Args:
+ name: The name of the logger, typically __name__
+
+ Returns:
+ A configured Logger instance
+ """
+ if not _logging_initialized:
+ # Set up basic logging if not initialized
+ # This ensures logging works even before setup_logging() is called
+ logging.basicConfig(
+ format=DEFAULT_FORMAT,
+ datefmt=DEFAULT_DATE_FORMAT,
+ level=logging.INFO,
+ )
+
+ return logging.getLogger(name)
+
+
+def set_log_level(level: int) -> None:
+ """
+ Change the log level for all Potato loggers at runtime.
+
+ Args:
+ level: The logging level (e.g., logging.DEBUG, logging.INFO)
+ """
+ root_logger = logging.getLogger()
+ root_logger.setLevel(level)
+
+ for handler in root_logger.handlers:
+ handler.setLevel(level)
+
+ for logger_name in POTATO_LOGGERS:
+ logging.getLogger(logger_name).setLevel(level)
+
+
+def get_log_level() -> int:
+ """
+ Get the current log level.
+
+ Returns:
+ The current logging level
+ """
+ return logging.getLogger("potato").level
+
+
+def is_ui_debug_enabled() -> bool:
+ """
+ Check if UI/frontend debug logging is enabled.
+
+ Returns:
+ True if UI debug logging is enabled
+ """
+ return _ui_debug_enabled
+
+
+def is_server_debug_enabled() -> bool:
+ """
+ Check if server/backend debug logging is enabled.
+
+ Returns:
+ True if server debug logging is enabled
+ """
+ return _server_debug_enabled
+
+
+def get_debug_log_settings() -> dict:
+ """
+ Get the current debug log settings for passing to frontend.
+
+ Returns:
+ Dict with 'ui_debug' and 'server_debug' boolean flags
+ """
+ return {
+ 'ui_debug': _ui_debug_enabled,
+ 'server_debug': _server_debug_enabled,
+ }
diff --git a/potato/mace.py b/potato/mace.py
new file mode 100644
index 0000000000000000000000000000000000000000..68709da78fc4f7ae4e3408e5f20e1e3cc7dc4b90
--- /dev/null
+++ b/potato/mace.py
@@ -0,0 +1,305 @@
+"""
+MACE (Multi-Annotator Competence Estimation) algorithm implementation.
+
+Implements the Variational Bayes EM algorithm from:
+ Hovy, D., Berg-Kirkpatrick, T., Vaswani, A., & Hovy, E. (2013).
+ Learning Whom to Trust with MACE. NAACL-HLT.
+
+Each annotator is modeled as either "knowing" (produces correct labels) or
+"guessing" (random strategy). The algorithm jointly estimates:
+1. True labels for each item (posterior distribution over categories)
+2. Annotator competence scores โ P(knowing) per annotator (0.0โ1.0)
+
+This module has no Potato dependencies and can be used standalone.
+"""
+
+import logging
+
+import numpy as np
+from scipy.special import digamma
+
+logger = logging.getLogger(__name__)
+
+# Small constant to prevent log(0)
+EPS = 1e-10
+
+
+class MACEAlgorithm:
+ """Pure MACE implementation using Variational Bayes EM.
+
+ Args:
+ num_annotators: Number of annotators (columns in annotations matrix).
+ num_labels: Number of possible label categories.
+ num_instances: Number of items (rows in annotations matrix).
+ alpha: Beta prior parameter for spamming (competence). Default 0.5.
+ beta: Dirichlet prior parameter for guessing strategy. Default 0.5.
+ num_restarts: Number of random restarts to find best solution. Default 10.
+ num_iters: Number of EM iterations per restart. Default 50.
+ seed: Random seed for reproducibility. None for non-deterministic.
+ """
+
+ def __init__(
+ self,
+ num_annotators,
+ num_labels,
+ num_instances,
+ alpha=0.5,
+ beta=0.5,
+ num_restarts=10,
+ num_iters=50,
+ seed=None,
+ ):
+ self.num_annotators = num_annotators
+ self.num_labels = num_labels
+ self.num_instances = num_instances
+ self.alpha = alpha
+ self.beta = beta
+ self.num_restarts = num_restarts
+ self.num_iters = num_iters
+ self.rng = np.random.RandomState(seed)
+
+ def fit(self, annotations):
+ """Run MACE on an annotation matrix.
+
+ Args:
+ annotations: np.ndarray of shape (num_instances, num_annotators).
+ Values are label indices 0..num_labels-1, or -1 for missing.
+
+ Returns:
+ tuple: (predicted_labels, competence, marginals, log_likelihood)
+ - predicted_labels: np.ndarray of shape (num_instances,), argmax label per item
+ - competence: np.ndarray of shape (num_annotators,), P(knowing) per annotator
+ - marginals: np.ndarray of shape (num_instances, num_labels), posterior over labels
+ - log_likelihood: float, log-likelihood of the best restart
+ """
+ best_ll = -np.inf
+ best_result = None
+
+ for restart in range(self.num_restarts):
+ spamming, theta = self._initialize()
+
+ for iteration in range(self.num_iters):
+ # E-step: compute posterior over true labels
+ marginals = self._e_step(annotations, spamming, theta)
+
+ # M-step: update spamming and theta via variational update
+ spamming, theta = self._m_step(annotations, marginals)
+
+ ll = self._log_likelihood(annotations, marginals, spamming, theta)
+
+ if ll > best_ll:
+ best_ll = ll
+ best_result = (marginals, spamming, theta)
+
+ marginals, spamming, theta = best_result
+
+ # Decode: argmax of marginals
+ predicted_labels = np.argmax(marginals, axis=1)
+
+ # Competence: E[spamming[:,0]] = P(knowing) per annotator
+ # spamming[:,0] is the "knowing" component, spamming[:,1] is "guessing"
+ competence = spamming[:, 0] / (spamming[:, 0] + spamming[:, 1])
+
+ return predicted_labels, competence, marginals, best_ll
+
+ def _initialize(self):
+ """Random initialization of parameters.
+
+ Returns:
+ tuple: (spamming, theta)
+ - spamming: np.ndarray shape (num_annotators, 2), Beta variational params
+ Column 0 = "knowing" mass, Column 1 = "guessing" mass
+ - theta: np.ndarray shape (num_annotators, num_labels), Dirichlet params
+ for guessing strategy per annotator
+ """
+ # Initialize spamming from Beta(alpha, alpha) prior
+ # Add random perturbation to break symmetry
+ spamming = np.zeros((self.num_annotators, 2))
+ spamming[:, 0] = self.alpha + self.rng.random(self.num_annotators)
+ spamming[:, 1] = self.alpha + self.rng.random(self.num_annotators)
+
+ # Initialize theta (guessing strategy) from Dirichlet(beta,...,beta)
+ theta = np.zeros((self.num_annotators, self.num_labels))
+ for j in range(self.num_annotators):
+ theta[j] = self.beta + self.rng.random(self.num_labels)
+
+ return spamming, theta
+
+ def _e_step(self, annotations, spamming, theta):
+ """Compute posterior P(true_label=k | observations) for each item.
+
+ Uses the current spamming and theta parameters to compute the
+ expected true label distribution via Bayes rule.
+
+ Args:
+ annotations: np.ndarray shape (num_instances, num_annotators), -1 = missing
+ spamming: np.ndarray shape (num_annotators, 2)
+ theta: np.ndarray shape (num_annotators, num_labels)
+
+ Returns:
+ marginals: np.ndarray shape (num_instances, num_labels)
+ """
+ marginals = np.zeros((self.num_instances, self.num_labels))
+
+ # Precompute expected log parameters using digamma
+ # E[log spamming_j] for knowing vs guessing
+ e_log_s = digamma(spamming) - digamma(spamming.sum(axis=1, keepdims=True))
+ # e_log_s[:, 0] = E[log P(knowing)]
+ # e_log_s[:, 1] = E[log P(guessing)]
+
+ # E[log theta_j_k] for each annotator's guessing distribution
+ e_log_theta = digamma(theta) - digamma(theta.sum(axis=1, keepdims=True))
+
+ for k in range(self.num_labels):
+ log_prob = np.zeros(self.num_instances)
+
+ for j in range(self.num_annotators):
+ # Mask for instances where annotator j provided a label
+ observed = annotations[:, j] >= 0
+ if not np.any(observed):
+ continue
+
+ label_j = annotations[observed, j].astype(int)
+
+ # P(observation | knowing, true_label=k):
+ # = 1 if label_j == k, else 0
+ # In log space: log(P(knowing) * I(label==k) + P(guessing) * theta[j,label])
+ # We use the variational decomposition:
+ # log P(x_ij | T_i=k) = log(exp(E[log s_j0]) * I(a_ij=k)
+ # + exp(E[log s_j1]) * exp(E[log theta_j,a_ij]))
+
+ # For numerical stability, compute in log-sum-exp form
+ knowing_term = np.full(observed.sum(), -np.inf)
+ match = label_j == k
+ knowing_term[match] = e_log_s[j, 0]
+
+ # Guessing term: P(guessing) * theta[j, observed_label]
+ guessing_term = e_log_s[j, 1] + e_log_theta[j, label_j]
+
+ # log-sum-exp of knowing and guessing
+ max_term = np.maximum(knowing_term, guessing_term)
+ log_sum = max_term + np.log(
+ np.exp(knowing_term - max_term) + np.exp(guessing_term - max_term) + EPS
+ )
+
+ log_prob[observed] += log_sum
+
+ marginals[:, k] = log_prob
+
+ # Normalize to probabilities (softmax over labels)
+ max_marginals = marginals.max(axis=1, keepdims=True)
+ marginals = np.exp(marginals - max_marginals)
+ row_sums = marginals.sum(axis=1, keepdims=True)
+ row_sums = np.maximum(row_sums, EPS)
+ marginals /= row_sums
+
+ return marginals
+
+ def _m_step(self, annotations, marginals):
+ """Variational M-step: update spamming and theta using expected counts.
+
+ Args:
+ annotations: np.ndarray shape (num_instances, num_annotators), -1 = missing
+ marginals: np.ndarray shape (num_instances, num_labels)
+
+ Returns:
+ tuple: (spamming, theta) updated parameters
+ """
+ spamming = np.zeros((self.num_annotators, 2))
+ theta = np.zeros((self.num_annotators, self.num_labels))
+
+ for j in range(self.num_annotators):
+ observed = annotations[:, j] >= 0
+ if not np.any(observed):
+ # No observations for this annotator โ use prior
+ spamming[j, 0] = self.alpha
+ spamming[j, 1] = self.alpha
+ theta[j] = self.beta
+ continue
+
+ label_j = annotations[observed, j].astype(int)
+ marginals_j = marginals[observed]
+
+ # Expected count of "knowing" for annotator j:
+ # Sum over instances where annotator's label matches true label
+ # weighted by P(true_label=k)
+ knowing_count = 0.0
+ guessing_count = 0.0
+
+ for i_idx in range(len(label_j)):
+ k = label_j[i_idx]
+ p_correct = marginals_j[i_idx, k]
+ knowing_count += p_correct
+ guessing_count += (1.0 - p_correct)
+
+ spamming[j, 0] = self.alpha + knowing_count
+ spamming[j, 1] = self.alpha + guessing_count
+
+ # Update theta: expected count of guessing label k
+ # When guessing, the annotator produces label a_ij with probability theta[j, a_ij]
+ # The expected count of guessing-and-producing-label-k is:
+ # sum_i (1 - P(knowing_ij)) * I(a_ij = k)
+ for k in range(self.num_labels):
+ mask = label_j == k
+ if np.any(mask):
+ # Weight by P(guessing) โ 1 - P(correct)
+ theta[j, k] = self.beta + np.sum(1.0 - marginals_j[mask, k])
+ else:
+ theta[j, k] = self.beta
+
+ return spamming, theta
+
+ def _log_likelihood(self, annotations, marginals, spamming, theta):
+ """Compute log-likelihood of the data given current parameters.
+
+ Args:
+ annotations: np.ndarray shape (num_instances, num_annotators)
+ marginals: np.ndarray shape (num_instances, num_labels)
+ spamming: np.ndarray shape (num_annotators, 2)
+ theta: np.ndarray shape (num_annotators, num_labels)
+
+ Returns:
+ float: log-likelihood value
+ """
+ ll = 0.0
+
+ # Normalize spamming and theta to probabilities for likelihood
+ s_norm = spamming / spamming.sum(axis=1, keepdims=True)
+ t_norm = theta / theta.sum(axis=1, keepdims=True)
+
+ for i in range(self.num_instances):
+ for k in range(self.num_labels):
+ if marginals[i, k] < EPS:
+ continue
+
+ log_p = 0.0
+ for j in range(self.num_annotators):
+ if annotations[i, j] < 0:
+ continue
+ a = int(annotations[i, j])
+
+ # P(a_ij | T_i=k) = s_j * I(a==k) + (1-s_j) * theta_j_a
+ p_knowing = s_norm[j, 0] * (1.0 if a == k else 0.0)
+ p_guessing = s_norm[j, 1] * t_norm[j, a]
+ p = p_knowing + p_guessing
+ log_p += np.log(max(p, EPS))
+
+ ll += marginals[i, k] * log_p
+
+ return ll
+
+ @staticmethod
+ def entropy(marginals):
+ """Compute entropy of label distributions per item.
+
+ Higher entropy = more uncertainty about the true label.
+
+ Args:
+ marginals: np.ndarray shape (num_instances, num_labels)
+
+ Returns:
+ np.ndarray shape (num_instances,), entropy per item
+ """
+ # Clip to avoid log(0)
+ p = np.clip(marginals, EPS, 1.0)
+ return -np.sum(p * np.log(p), axis=1)
diff --git a/potato/mace_manager.py b/potato/mace_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d4d4920b4ef9853233d6bb2e6eb064f093d4449
--- /dev/null
+++ b/potato/mace_manager.py
@@ -0,0 +1,596 @@
+"""
+MACE Manager โ integration layer between MACE algorithm and Potato data model.
+
+Extracts annotation data from Potato's state managers, converts to MACE format,
+runs the algorithm, and stores/caches results. Follows the singleton pattern
+used by SimilarityEngine and other Potato managers.
+
+Supports:
+- Radio, likert, select: single categorical annotation per item
+- Multiselect: per-option binary MACE (each checkbox = separate yes/no run)
+"""
+
+import json
+import logging
+import os
+import threading
+import time
+from dataclasses import dataclass, field, asdict
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+from potato.mace import MACEAlgorithm
+
+logger = logging.getLogger(__name__)
+
+# Categorical annotation types that MACE can process
+CATEGORICAL_TYPES = {"radio", "likert", "select", "multiselect"}
+
+
+@dataclass
+class MACEConfig:
+ """Configuration for MACE competence estimation."""
+
+ enabled: bool = False
+ trigger_every_n: int = 10
+ min_annotations_per_item: int = 3
+ min_items: int = 5
+ num_restarts: int = 10
+ num_iters: int = 50
+ alpha: float = 0.5
+ beta: float = 0.5
+ output_subdir: str = "mace"
+ cache_results: bool = True
+
+ @classmethod
+ def from_dict(cls, d):
+ """Create MACEConfig from a dictionary, ignoring unknown keys."""
+ known = {f.name for f in cls.__dataclass_fields__.values()}
+ return cls(**{k: v for k, v in d.items() if k in known})
+
+
+@dataclass
+class MACEResult:
+ """Result of a single MACE run for one schema (or schema+option for multiselect)."""
+
+ schema_name: str
+ competence_scores: Dict[str, float] # user_id -> P(knowing)
+ predicted_labels: Dict[str, Any] # instance_id -> predicted label
+ label_entropy: Dict[str, float] # instance_id -> entropy
+ label_mapping: Dict[int, str] # index -> original label value
+ num_annotators: int
+ num_instances: int
+ timestamp: str
+ log_likelihood: float
+ option_name: Optional[str] = None # For multiselect per-option
+
+ def to_dict(self):
+ return {
+ "schema_name": self.schema_name,
+ "competence_scores": self.competence_scores,
+ "predicted_labels": self.predicted_labels,
+ "label_entropy": self.label_entropy,
+ "label_mapping": self.label_mapping,
+ "num_annotators": self.num_annotators,
+ "num_instances": self.num_instances,
+ "timestamp": self.timestamp,
+ "log_likelihood": self.log_likelihood,
+ "option_name": self.option_name,
+ }
+
+ @classmethod
+ def from_dict(cls, d):
+ return cls(**d)
+
+
+class MACEManager:
+ """Manages MACE computation, caching, and result access.
+
+ Args:
+ config: Full Potato configuration dictionary.
+ """
+
+ def __init__(self, config: dict):
+ self.config = config
+ self.mace_config = MACEConfig.from_dict(config.get("mace", {}))
+ self._lock = threading.Lock()
+ self._last_trigger_count = 0
+ self.results: Dict[str, MACEResult] = {} # key -> MACEResult
+
+ # Determine output directory
+ output_dir = config.get("output_annotation_dir", "annotation_output")
+ self._output_dir = os.path.join(output_dir, self.mace_config.output_subdir)
+
+ # Load cached results if configured
+ if self.mace_config.cache_results:
+ self._load_cache()
+
+ def _result_key(self, schema_name: str, option_name: Optional[str] = None) -> str:
+ """Generate a unique key for a MACE result."""
+ if option_name:
+ return f"{schema_name}::{option_name}"
+ return schema_name
+
+ def check_and_run(self, total_annotations: int) -> bool:
+ """Check if it's time to run MACE and trigger if so.
+
+ Args:
+ total_annotations: Current total annotation count across all users.
+
+ Returns:
+ True if MACE was run, False otherwise.
+ """
+ if not self.mace_config.enabled:
+ return False
+
+ trigger_n = self.mace_config.trigger_every_n
+ if trigger_n <= 0:
+ return False
+
+ # Check if we've crossed the next threshold
+ if total_annotations - self._last_trigger_count >= trigger_n:
+ self._last_trigger_count = total_annotations
+ try:
+ self.run_all_schemas()
+ return True
+ except Exception as e:
+ logger.error(f"MACE run failed: {e}", exc_info=True)
+ return False
+
+ return False
+
+ def run_all_schemas(self, _usm=None, _ism=None) -> Dict[str, MACEResult]:
+ """Run MACE for all eligible annotation schemas.
+
+ Args:
+ _usm: Optional UserStateManager override (for testing).
+ _ism: Optional ItemStateManager override (for testing).
+
+ Returns:
+ Dict mapping result keys to MACEResult objects.
+ """
+ if _usm is None:
+ from potato.user_state_management import get_user_state_manager
+ _usm = get_user_state_manager()
+ if _ism is None:
+ from potato.item_state_management import get_item_state_manager
+ _ism = get_item_state_manager()
+
+ usm = _usm
+ ism = _ism
+
+ annotation_schemes = self.config.get("annotation_schemes", [])
+ new_results = {}
+
+ for scheme in annotation_schemes:
+ schema_type = scheme.get("annotation_type", "")
+ schema_name = scheme.get("name", "")
+
+ if schema_type not in CATEGORICAL_TYPES:
+ continue
+
+ if not schema_name:
+ continue
+
+ if schema_type == "multiselect":
+ # Per-option binary MACE
+ labels = scheme.get("labels", [])
+ for option in labels:
+ option_name = option if isinstance(option, str) else str(option)
+ result = self._run_for_schema(
+ usm, ism, schema_name, schema_type,
+ binary_option=option_name
+ )
+ if result:
+ key = self._result_key(schema_name, option_name)
+ new_results[key] = result
+ else:
+ # Standard categorical MACE
+ result = self._run_for_schema(usm, ism, schema_name, schema_type)
+ if result:
+ key = self._result_key(schema_name)
+ new_results[key] = result
+
+ with self._lock:
+ self.results.update(new_results)
+
+ # Save to disk
+ if self.mace_config.cache_results and new_results:
+ self._save_cache()
+
+ if new_results:
+ logger.info(
+ f"MACE completed: {len(new_results)} schema(s) processed"
+ )
+
+ return new_results
+
+ def _run_for_schema(
+ self, usm, ism, schema_name: str, schema_type: str,
+ binary_option: Optional[str] = None
+ ) -> Optional[MACEResult]:
+ """Run MACE for a single schema (or schema+option).
+
+ Args:
+ usm: UserStateManager instance
+ ism: ItemStateManager instance
+ schema_name: Name of the annotation schema
+ schema_type: Type of annotation (radio, likert, select, multiselect)
+ binary_option: For multiselect, the specific option to create binary MACE for
+
+ Returns:
+ MACEResult or None if insufficient data.
+ """
+ # Collect all annotations: {instance_id: {user_id: label_value}}
+ annotations_by_item = {}
+ all_annotators = set()
+
+ user_ids = usm.get_user_ids()
+ for user_id in user_ids:
+ user_state = usm.get_user_state(user_id)
+ if not user_state:
+ continue
+
+ for instance_id, label_dict in user_state.instance_id_to_label_to_value.items():
+ annotation_value = self._extract_annotation(
+ label_dict, schema_name, schema_type, binary_option
+ )
+ if annotation_value is not None:
+ if instance_id not in annotations_by_item:
+ annotations_by_item[instance_id] = {}
+ annotations_by_item[instance_id][user_id] = annotation_value
+ all_annotators.add(user_id)
+
+ # Filter items with enough annotators
+ min_annots = self.mace_config.min_annotations_per_item
+ eligible_items = {
+ iid: annots for iid, annots in annotations_by_item.items()
+ if len(annots) >= min_annots
+ }
+
+ if len(eligible_items) < self.mace_config.min_items:
+ logger.debug(
+ f"MACE skip {schema_name}: only {len(eligible_items)} eligible items "
+ f"(need {self.mace_config.min_items})"
+ )
+ return None
+
+ # Build label mapping
+ if binary_option:
+ # Binary: 0 = False/No, 1 = True/Yes
+ all_values = {"0", "1"}
+ else:
+ all_values = set()
+ for annots in eligible_items.values():
+ all_values.update(str(v) for v in annots.values())
+
+ sorted_values = sorted(all_values)
+ value_to_idx = {v: i for i, v in enumerate(sorted_values)}
+ idx_to_value = {i: v for v, i in value_to_idx.items()}
+ num_labels = len(sorted_values)
+
+ if num_labels < 2:
+ logger.debug(f"MACE skip {schema_name}: only {num_labels} unique label(s)")
+ return None
+
+ # Build annotator mapping
+ sorted_annotators = sorted(all_annotators)
+ annotator_to_idx = {u: i for i, u in enumerate(sorted_annotators)}
+ num_annotators = len(sorted_annotators)
+
+ # Build item mapping
+ sorted_items = sorted(eligible_items.keys())
+ num_instances = len(sorted_items)
+
+ # Build annotation matrix
+ matrix = np.full((num_instances, num_annotators), -1, dtype=int)
+ for i, iid in enumerate(sorted_items):
+ for uid, val in eligible_items[iid].items():
+ j = annotator_to_idx[uid]
+ str_val = str(val)
+ if str_val in value_to_idx:
+ matrix[i, j] = value_to_idx[str_val]
+
+ # Run MACE
+ mace = MACEAlgorithm(
+ num_annotators=num_annotators,
+ num_labels=num_labels,
+ num_instances=num_instances,
+ alpha=self.mace_config.alpha,
+ beta=self.mace_config.beta,
+ num_restarts=self.mace_config.num_restarts,
+ num_iters=self.mace_config.num_iters,
+ seed=42,
+ )
+
+ predicted_indices, competence, marginals, log_lik = mace.fit(matrix)
+ entropies = MACEAlgorithm.entropy(marginals)
+
+ # Map results back to original IDs
+ competence_scores = {
+ sorted_annotators[j]: float(competence[j])
+ for j in range(num_annotators)
+ }
+ predicted_labels = {
+ sorted_items[i]: idx_to_value[int(predicted_indices[i])]
+ for i in range(num_instances)
+ }
+ label_entropy = {
+ sorted_items[i]: float(entropies[i])
+ for i in range(num_instances)
+ }
+
+ return MACEResult(
+ schema_name=schema_name,
+ competence_scores=competence_scores,
+ predicted_labels=predicted_labels,
+ label_entropy=label_entropy,
+ label_mapping={int(k): v for k, v in idx_to_value.items()},
+ num_annotators=num_annotators,
+ num_instances=num_instances,
+ timestamp=time.strftime("%Y-%m-%dT%H:%M:%S"),
+ log_likelihood=float(log_lik),
+ option_name=binary_option,
+ )
+
+ def _extract_annotation(
+ self, label_dict, schema_name: str, schema_type: str,
+ binary_option: Optional[str] = None
+ ) -> Optional[str]:
+ """Extract a single annotation value from a label dictionary.
+
+ Args:
+ label_dict: Dict mapping Label objects to values for one instance+user
+ schema_name: Target schema name
+ schema_type: Annotation type (radio, likert, select, multiselect)
+ binary_option: For multiselect, the specific option name
+
+ Returns:
+ The annotation value as a string, or None if not found.
+ """
+ # Known falsy values that indicate "not selected"
+ _FALSY = (False, "false", "False", 0, "0", "", None)
+
+ if schema_type == "multiselect" and binary_option:
+ # Look for the specific option's Label
+ for label, value in label_dict.items():
+ if (label.get_schema() == schema_name
+ and label.get_name() == binary_option):
+ # Convert to binary: "1" for checked, "0" for unchecked
+ if value not in _FALSY:
+ return "1"
+ return "0"
+ return None
+ else:
+ # Radio/likert/select: find the label with a truthy (non-falsy) value.
+ # The value may be True, "true", or the label name itself (e.g. "positive").
+ for label, value in label_dict.items():
+ if label.get_schema() != schema_name:
+ continue
+ if value not in _FALSY:
+ return label.get_name()
+ return None
+
+ def get_competence(self, user_id: str) -> Dict[str, float]:
+ """Get competence scores for a user across all schemas.
+
+ Args:
+ user_id: The user/annotator ID.
+
+ Returns:
+ Dict mapping schema (or schema::option) keys to competence scores.
+ """
+ scores = {}
+ with self._lock:
+ for key, result in self.results.items():
+ if user_id in result.competence_scores:
+ scores[key] = result.competence_scores[user_id]
+ return scores
+
+ def get_prediction(self, instance_id: str, schema: str) -> Optional[str]:
+ """Get MACE predicted label for an instance and schema.
+
+ Args:
+ instance_id: The instance/item ID.
+ schema: The schema name.
+
+ Returns:
+ Predicted label string, or None if no prediction available.
+ """
+ with self._lock:
+ result = self.results.get(schema)
+ if result and instance_id in result.predicted_labels:
+ return result.predicted_labels[instance_id]
+ return None
+
+ def get_results_summary(self) -> dict:
+ """Get a summary of all MACE results for the admin API.
+
+ Returns:
+ Dict with schema results, overall stats, and per-user competence.
+ """
+ with self._lock:
+ if not self.results:
+ return {
+ "enabled": self.mace_config.enabled,
+ "has_results": False,
+ "schemas": [],
+ "annotator_competence": {},
+ }
+
+ schemas = []
+ all_competence = {}
+
+ for key, result in self.results.items():
+ schemas.append({
+ "key": key,
+ "schema_name": result.schema_name,
+ "option_name": result.option_name,
+ "num_annotators": result.num_annotators,
+ "num_instances": result.num_instances,
+ "log_likelihood": result.log_likelihood,
+ "timestamp": result.timestamp,
+ "label_mapping": result.label_mapping,
+ })
+
+ # Aggregate competence across schemas
+ for uid, score in result.competence_scores.items():
+ if uid not in all_competence:
+ all_competence[uid] = {}
+ all_competence[uid][key] = score
+
+ # Compute average competence per annotator
+ annotator_competence = {}
+ for uid, schema_scores in all_competence.items():
+ scores = list(schema_scores.values())
+ annotator_competence[uid] = {
+ "scores": schema_scores,
+ "average": sum(scores) / len(scores) if scores else 0.0,
+ }
+
+ return {
+ "enabled": self.mace_config.enabled,
+ "has_results": True,
+ "schemas": schemas,
+ "annotator_competence": annotator_competence,
+ "config": {
+ "trigger_every_n": self.mace_config.trigger_every_n,
+ "min_annotations_per_item": self.mace_config.min_annotations_per_item,
+ "min_items": self.mace_config.min_items,
+ "num_restarts": self.mace_config.num_restarts,
+ "num_iters": self.mace_config.num_iters,
+ },
+ }
+
+ def get_predictions_for_schema(
+ self, schema: str, instance_id: Optional[str] = None
+ ) -> dict:
+ """Get predictions with optional filtering.
+
+ Args:
+ schema: Schema name to get predictions for.
+ instance_id: Optional specific instance to get.
+
+ Returns:
+ Dict with predictions and entropy data.
+ """
+ with self._lock:
+ result = self.results.get(schema)
+ if not result:
+ return {"error": f"No MACE results for schema '{schema}'"}
+
+ if instance_id:
+ pred = result.predicted_labels.get(instance_id)
+ ent = result.label_entropy.get(instance_id)
+ if pred is None:
+ return {"error": f"No prediction for instance '{instance_id}'"}
+ return {
+ "instance_id": instance_id,
+ "predicted_label": pred,
+ "entropy": ent,
+ "label_mapping": result.label_mapping,
+ }
+
+ return {
+ "schema_name": result.schema_name,
+ "option_name": result.option_name,
+ "predicted_labels": result.predicted_labels,
+ "label_entropy": result.label_entropy,
+ "label_mapping": result.label_mapping,
+ "num_instances": result.num_instances,
+ }
+
+ def count_total_annotations(self, _usm=None) -> int:
+ """Count total annotations across all users.
+
+ Args:
+ _usm: Optional UserStateManager override (for testing).
+
+ Returns:
+ Total number of annotated instances across all users.
+ """
+ if _usm is None:
+ from potato.user_state_management import get_user_state_manager
+ _usm = get_user_state_manager()
+
+ usm = _usm
+ total = 0
+ for user_id in usm.get_user_ids():
+ user_state = usm.get_user_state(user_id)
+ if user_state:
+ total += len(user_state.instance_id_to_label_to_value)
+ return total
+
+ def _save_cache(self):
+ """Save current results to disk."""
+ try:
+ os.makedirs(self._output_dir, exist_ok=True)
+ cache_path = os.path.join(self._output_dir, "mace_results.json")
+ with self._lock:
+ data = {key: result.to_dict() for key, result in self.results.items()}
+ with open(cache_path, "w") as f:
+ json.dump(data, f, indent=2)
+ logger.debug(f"MACE results saved to {cache_path}")
+ except Exception as e:
+ logger.error(f"Failed to save MACE cache: {e}")
+
+ def _load_cache(self):
+ """Load cached results from disk if available."""
+ cache_path = os.path.join(self._output_dir, "mace_results.json")
+ if not os.path.exists(cache_path):
+ return
+
+ try:
+ with open(cache_path, "r") as f:
+ data = json.load(f)
+ for key, result_dict in data.items():
+ self.results[key] = MACEResult.from_dict(result_dict)
+ logger.info(f"Loaded {len(self.results)} cached MACE results")
+ except Exception as e:
+ logger.warning(f"Failed to load MACE cache: {e}")
+
+
+# ============================================================================
+# Singleton management
+# ============================================================================
+
+_MACE_MANAGER: Optional[MACEManager] = None
+_MACE_LOCK = threading.Lock()
+
+
+def init_mace_manager(config: dict) -> Optional[MACEManager]:
+ """Initialize the MACE manager singleton.
+
+ Args:
+ config: Full Potato configuration dictionary.
+
+ Returns:
+ MACEManager instance, or None if not enabled.
+ """
+ global _MACE_MANAGER
+ with _MACE_LOCK:
+ if _MACE_MANAGER is None:
+ mace_config = config.get("mace", {})
+ if mace_config.get("enabled", False):
+ _MACE_MANAGER = MACEManager(config)
+ logger.info("MACE manager initialized")
+ else:
+ logger.debug("MACE not enabled in config")
+ return _MACE_MANAGER
+
+
+def get_mace_manager() -> Optional[MACEManager]:
+ """Get the MACE manager singleton.
+
+ Returns:
+ MACEManager instance, or None if not initialized.
+ """
+ return _MACE_MANAGER
+
+
+def clear_mace_manager():
+ """Clear the MACE manager singleton (for testing)."""
+ global _MACE_MANAGER
+ with _MACE_LOCK:
+ _MACE_MANAGER = None
diff --git a/potato/memos/__init__.py b/potato/memos/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..94d69456d94e8595027acd43a96677c157e839ea
--- /dev/null
+++ b/potato/memos/__init__.py
@@ -0,0 +1,33 @@
+"""
+Memos (universal annotation feature).
+
+Free-text notes an annotator attaches to an instance or to a text
+selection within an instance. Universal โ available in standard
+annotation, solo mode, and QDA mode; gated by `annotation_ui.memos`
+(default off in standard mode; on for qda_mode/solo_mode).
+
+Layers:
+- `store` โ SQLite persistence over the universal `project.sqlite`.
+- `service` โ visibility + permission rules (private/shared; admins
+ always read; author-only edit; author/admin delete).
+"""
+
+from .service import (
+ MemoError,
+ MemoNotFound,
+ MemoPermissionError,
+ create_memo,
+ delete_memo,
+ list_visible,
+ update_memo,
+)
+
+__all__ = [
+ "MemoError",
+ "MemoNotFound",
+ "MemoPermissionError",
+ "create_memo",
+ "list_visible",
+ "update_memo",
+ "delete_memo",
+]
diff --git a/potato/memos/api.py b/potato/memos/api.py
new file mode 100644
index 0000000000000000000000000000000000000000..bc825123c930372c409f415ac5fcd624ed1c4657
--- /dev/null
+++ b/potato/memos/api.py
@@ -0,0 +1,151 @@
+"""
+Memos REST API (universal).
+
+Blueprint mounted at /api/memos. Visibility/permission enforcement lives
+in the service layer; this layer handles auth (logged-in user), the
+feature gate (annotation_ui.memos), request parsing, and error mapping.
+
+Privilege tier ("always read / may delete others"): adjudicators, via the
+adjudication manager. Admin-dashboard memo moderation is a later follow-up.
+"""
+
+from __future__ import annotations
+
+import logging
+from functools import wraps
+
+from flask import Blueprint, jsonify, request, session
+
+from . import (
+ MemoError,
+ MemoNotFound,
+ MemoPermissionError,
+ create_memo,
+ delete_memo,
+ list_visible,
+ update_memo,
+)
+
+logger = logging.getLogger(__name__)
+
+memos_bp = Blueprint("memos", __name__, url_prefix="/api/memos")
+
+
+def _config() -> dict:
+ from potato.server_utils.config_module import config
+ return config
+
+
+def memos_enabled(config: dict) -> bool:
+ """Default: off in standard mode; on when qda_mode or solo_mode is on.
+ Explicit annotation_ui.memos always wins."""
+ ui = config.get("annotation_ui") or {}
+ if isinstance(ui, dict) and "memos" in ui:
+ return bool(ui["memos"])
+ return bool(
+ (config.get("qda_mode") or {}).get("enabled")
+ or (config.get("solo_mode") or {}).get("enabled")
+ )
+
+
+def _default_visibility(config: dict) -> str:
+ ui = config.get("annotation_ui") or {}
+ v = ui.get("visibility") if isinstance(ui, dict) else None
+ return v if v in ("private", "shared") else "private"
+
+
+def _is_privileged(username: str) -> bool:
+ try:
+ from potato.adjudication import get_adjudication_manager
+ adj = get_adjudication_manager()
+ return bool(adj and adj.is_adjudicator(username))
+ except Exception:
+ return False
+
+
+def _ctx():
+ """(task_dir, project, username, is_privileged) or None if not usable."""
+ config = _config()
+ if not memos_enabled(config):
+ return None, None, None, None, ("memos_disabled",)
+ username = session.get("username")
+ if not username:
+ return None, None, None, None, ("unauthenticated",)
+ task_dir = config.get("task_dir", ".")
+ project = config.get("annotation_task_name") or "default"
+ return task_dir, project, username, _is_privileged(username), None
+
+
+def memos_required(view):
+ @wraps(view)
+ def wrapper(*args, **kwargs):
+ task_dir, project, username, priv, err = _ctx()
+ if err == ("memos_disabled",):
+ return jsonify({
+ "error": "Memos are not enabled in this deployment.",
+ "hint": "Set annotation_ui.memos: true (on by default in "
+ "qda_mode/solo_mode).",
+ }), 503
+ if err == ("unauthenticated",):
+ return jsonify({"error": "Not authenticated"}), 401
+ return view(task_dir, project, username, priv, *args, **kwargs)
+ return wrapper
+
+
+def _handle(fn):
+ """Map service exceptions to HTTP codes."""
+ try:
+ return fn()
+ except MemoNotFound as e:
+ return jsonify({"error": str(e)}), 404
+ except MemoPermissionError as e:
+ return jsonify({"error": str(e)}), 403
+ except MemoError as e:
+ return jsonify({"error": str(e)}), 400
+
+
+@memos_bp.route("", methods=["GET"])
+@memos_required
+def list_memos(task_dir, project, username, priv):
+ instance_id = request.args.get("instance_id")
+ if not instance_id:
+ return jsonify({"error": "instance_id is required"}), 400
+ memos = list_visible(
+ task_dir, project=project, instance_id=instance_id,
+ requester=username, is_privileged=priv,
+ )
+ return jsonify({"memos": memos})
+
+
+@memos_bp.route("", methods=["POST"])
+@memos_required
+def post_memo(task_dir, project, username, priv):
+ data = request.get_json(silent=True) or {}
+ instance_id = data.get("instance_id")
+ if not instance_id:
+ return jsonify({"error": "instance_id is required"}), 400
+ visibility = data.get("visibility") or _default_visibility(_config())
+ return _handle(lambda: jsonify({"memo": create_memo(
+ task_dir, project=project, instance_id=instance_id,
+ body=data.get("body", ""), created_by=username,
+ anchor=data.get("anchor"), visibility=visibility,
+ )}))
+
+
+@memos_bp.route("/", methods=["PATCH"])
+@memos_required
+def patch_memo(task_dir, project, username, priv, memo_id):
+ data = request.get_json(silent=True) or {}
+ return _handle(lambda: jsonify({"memo": update_memo(
+ task_dir, memo_id, requester=username, is_privileged=priv,
+ body=data.get("body"), visibility=data.get("visibility"),
+ )}))
+
+
+@memos_bp.route("/", methods=["DELETE"])
+@memos_required
+def remove_memo(task_dir, project, username, priv, memo_id):
+ def _do():
+ delete_memo(task_dir, memo_id, requester=username, is_privileged=priv)
+ return jsonify({"ok": True})
+ return _handle(_do)
diff --git a/potato/memos/service.py b/potato/memos/service.py
new file mode 100644
index 0000000000000000000000000000000000000000..45fefdd7dc7bf0b841db3e2de44197a952ce950a
--- /dev/null
+++ b/potato/memos/service.py
@@ -0,0 +1,130 @@
+"""
+Memo service (universal) โ visibility + permission rules over the store.
+
+Visibility model (decided 2026-05-18):
+- ``private`` (default): visible to the author and to admins/adjudicators.
+- ``shared``: visible to the author, admins/adjudicators, AND peer
+ annotators on the same project.
+- Admins/adjudicators can ALWAYS read every memo regardless of setting.
+
+Permissions:
+- Read: per the visibility rule above.
+- Edit (body/visibility): author only.
+- Delete: author OR an admin/adjudicator (moderation).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from . import store
+
+VALID_VISIBILITY = ("private", "shared")
+
+
+class MemoError(Exception):
+ """Base error for memo operations (maps to 4xx at the API layer)."""
+
+
+class MemoNotFound(MemoError):
+ pass
+
+
+class MemoPermissionError(MemoError):
+ pass
+
+
+def _can_read(memo: Dict[str, Any], requester: str, is_privileged: bool) -> bool:
+ return (
+ is_privileged
+ or memo["created_by"] == requester
+ or memo["visibility"] == "shared"
+ )
+
+
+def create_memo(
+ task_dir: str,
+ *,
+ project: str,
+ instance_id: str,
+ body: str,
+ created_by: str,
+ anchor: Optional[Dict[str, Any]] = None,
+ visibility: str = "private",
+) -> Dict[str, Any]:
+ body = (body or "").strip()
+ if not body:
+ raise MemoError("Memo body must not be empty")
+ if visibility not in VALID_VISIBILITY:
+ raise MemoError(
+ f"visibility must be one of {VALID_VISIBILITY} (got {visibility!r})"
+ )
+ if anchor is not None:
+ if not isinstance(anchor, dict) or "start" not in anchor or "end" not in anchor:
+ raise MemoError("anchor must be {start, end[, field]} or null")
+ return store.create(
+ task_dir, project=project, instance_id=instance_id, body=body,
+ created_by=created_by, anchor=anchor, visibility=visibility,
+ )
+
+
+def list_visible(
+ task_dir: str,
+ *,
+ project: str,
+ instance_id: str,
+ requester: str,
+ is_privileged: bool = False,
+) -> List[Dict[str, Any]]:
+ """Memos on an instance the requester is allowed to see."""
+ return [
+ m for m in store.list_for_instance(task_dir, project, instance_id)
+ if _can_read(m, requester, is_privileged)
+ ]
+
+
+def _load_or_404(task_dir: str, memo_id: str) -> Dict[str, Any]:
+ memo = store.get(task_dir, memo_id)
+ if memo is None:
+ raise MemoNotFound(f"Memo {memo_id} not found")
+ return memo
+
+
+def update_memo(
+ task_dir: str,
+ memo_id: str,
+ *,
+ requester: str,
+ is_privileged: bool = False,
+ body: Optional[str] = None,
+ visibility: Optional[str] = None,
+) -> Dict[str, Any]:
+ memo = _load_or_404(task_dir, memo_id)
+ if memo["created_by"] != requester:
+ raise MemoPermissionError("Only the memo author may edit it")
+ if body is not None and not body.strip():
+ raise MemoError("Memo body must not be empty")
+ if visibility is not None and visibility not in VALID_VISIBILITY:
+ raise MemoError(
+ f"visibility must be one of {VALID_VISIBILITY} (got {visibility!r})"
+ )
+ return store.update(
+ task_dir, memo_id,
+ body=body.strip() if body is not None else None,
+ visibility=visibility,
+ )
+
+
+def delete_memo(
+ task_dir: str,
+ memo_id: str,
+ *,
+ requester: str,
+ is_privileged: bool = False,
+) -> None:
+ memo = _load_or_404(task_dir, memo_id)
+ if memo["created_by"] != requester and not is_privileged:
+ raise MemoPermissionError(
+ "Only the author or an admin/adjudicator may delete this memo"
+ )
+ store.delete(task_dir, memo_id)
diff --git a/potato/memos/store.py b/potato/memos/store.py
new file mode 100644
index 0000000000000000000000000000000000000000..e44d14430333ee6efc27777ca0d3a6ac6c58dc33
--- /dev/null
+++ b/potato/memos/store.py
@@ -0,0 +1,143 @@
+"""
+Memo storage (universal).
+
+SQLite-backed CRUD over the `memos` table in `/project.sqlite`
+via the universal persistence layer. No visibility/permission logic lives
+here โ that is the service layer's job. This module only persists rows.
+
+A memo is a free-text note an annotator attaches to an instance, or to a
+text selection within an instance (offset-anchored). Universal: usable in
+standard annotation, solo mode, and QDA mode.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+from potato.persistence import Migration, get_db, register_migration
+
+_MEMOS_MIGRATION = Migration(
+ name="0001_memos",
+ sql="""
+ CREATE TABLE IF NOT EXISTS memos (
+ id TEXT PRIMARY KEY,
+ project TEXT NOT NULL,
+ instance_id TEXT NOT NULL,
+ anchor TEXT, -- NULL = instance-level; else JSON {start,end,field}
+ body TEXT NOT NULL,
+ created_by TEXT NOT NULL,
+ created_at REAL NOT NULL,
+ updated_at REAL NOT NULL,
+ visibility TEXT NOT NULL DEFAULT 'private'
+ CHECK (visibility IN ('private', 'shared'))
+ );
+ CREATE INDEX IF NOT EXISTS idx_memos_instance
+ ON memos (project, instance_id);
+ CREATE INDEX IF NOT EXISTS idx_memos_author
+ ON memos (project, created_by);
+ """,
+)
+
+# Registered at import so the table exists on the first get_db() call.
+register_migration(_MEMOS_MIGRATION)
+
+
+def _db(task_dir: str):
+ """Connection for the memos store, guaranteeing the migration is
+ registered first. register_migration is idempotent, so this is a
+ no-op in normal operation; it makes the store robust if a test
+ helper (clear_migrations) wiped the process-global registry before
+ the first get_db() for this task_dir opens the connection."""
+ register_migration(_MEMOS_MIGRATION)
+ return get_db(task_dir)
+
+
+def _row_to_dict(row) -> Dict[str, Any]:
+ d = dict(row)
+ d["anchor"] = json.loads(d["anchor"]) if d["anchor"] else None
+ return d
+
+
+def create(
+ task_dir: str,
+ *,
+ project: str,
+ instance_id: str,
+ body: str,
+ created_by: str,
+ anchor: Optional[Dict[str, Any]] = None,
+ visibility: str = "private",
+) -> Dict[str, Any]:
+ """Insert a memo row and return it as a dict."""
+ memo_id = uuid.uuid4().hex
+ now = time.time()
+ conn = _db(task_dir)
+ conn.execute(
+ """INSERT INTO memos
+ (id, project, instance_id, anchor, body, created_by,
+ created_at, updated_at, visibility)
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+ (
+ memo_id, project, instance_id,
+ json.dumps(anchor) if anchor else None,
+ body, created_by, now, now, visibility,
+ ),
+ )
+ conn.commit()
+ return get(task_dir, memo_id)
+
+
+def get(task_dir: str, memo_id: str) -> Optional[Dict[str, Any]]:
+ row = _db(task_dir).execute(
+ "SELECT * FROM memos WHERE id = ?", (memo_id,)
+ ).fetchone()
+ return _row_to_dict(row) if row else None
+
+
+def list_for_instance(
+ task_dir: str, project: str, instance_id: str
+) -> List[Dict[str, Any]]:
+ """All memos on an instance (no visibility filtering โ service does that)."""
+ rows = _db(task_dir).execute(
+ """SELECT * FROM memos
+ WHERE project = ? AND instance_id = ?
+ ORDER BY created_at ASC""",
+ (project, instance_id),
+ ).fetchall()
+ return [_row_to_dict(r) for r in rows]
+
+
+def update(
+ task_dir: str,
+ memo_id: str,
+ *,
+ body: Optional[str] = None,
+ visibility: Optional[str] = None,
+) -> Optional[Dict[str, Any]]:
+ """Patch body and/or visibility; bumps updated_at. No-op fields ignored."""
+ sets, params = [], []
+ if body is not None:
+ sets.append("body = ?")
+ params.append(body)
+ if visibility is not None:
+ sets.append("visibility = ?")
+ params.append(visibility)
+ if not sets:
+ return get(task_dir, memo_id)
+ sets.append("updated_at = ?")
+ params.append(time.time())
+ params.append(memo_id)
+ conn = _db(task_dir)
+ conn.execute(f"UPDATE memos SET {', '.join(sets)} WHERE id = ?", params)
+ conn.commit()
+ return get(task_dir, memo_id)
+
+
+def delete(task_dir: str, memo_id: str) -> bool:
+ conn = _db(task_dir)
+ cur = conn.execute("DELETE FROM memos WHERE id = ?", (memo_id,))
+ conn.commit()
+ return cur.rowcount > 0
diff --git a/potato/migrate_cli.py b/potato/migrate_cli.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c0507968cdf24717a2f6979d9ed40f432f0c1d6
--- /dev/null
+++ b/potato/migrate_cli.py
@@ -0,0 +1,533 @@
+#!/usr/bin/env python
+"""
+Config Migration Tool for Potato
+
+This module provides utilities to migrate Potato configuration files
+from older formats to the current v2 format.
+
+Usage:
+ potato migrate config.yaml --to-v2
+ potato migrate config.yaml --to-v2 --output new_config.yaml
+ potato migrate config.yaml --to-v2 --in-place
+ potato migrate config.yaml --to-v2 --dry-run
+"""
+
+import argparse
+import logging
+import os
+import sys
+import copy
+from typing import Dict, Any, List, Tuple
+from pathlib import Path
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+class MigrationRule:
+ """Base class for migration rules."""
+
+ def __init__(self, name: str, description: str):
+ self.name = name
+ self.description = description
+
+ def applies(self, config: Dict[str, Any]) -> bool:
+ """Check if this rule applies to the config."""
+ raise NotImplementedError
+
+ def migrate(self, config: Dict[str, Any]) -> Tuple[Dict[str, Any], List[str]]:
+ """
+ Apply the migration rule.
+
+ Returns:
+ Tuple of (migrated_config, list of changes made)
+ """
+ raise NotImplementedError
+
+
+def _collect_all_schemes(config: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """Return every annotation scheme dict in a config.
+
+ Gathers schemes from the top-level ``annotation_schemes``, every phase's
+ ``annotation_schemes`` (phases may be a list or a dict keyed by phase name,
+ where the ``order`` key is not a phase), and ``training.annotation_schemes``.
+ Mutating the returned dicts mutates the config in place.
+ """
+ schemes: List[Dict[str, Any]] = []
+
+ top = config.get("annotation_schemes")
+ if isinstance(top, list):
+ schemes.extend(s for s in top if isinstance(s, dict))
+
+ phases = config.get("phases")
+ if isinstance(phases, list):
+ for phase in phases:
+ if isinstance(phase, dict) and isinstance(phase.get("annotation_schemes"), list):
+ schemes.extend(s for s in phase["annotation_schemes"] if isinstance(s, dict))
+ elif isinstance(phases, dict):
+ for phase_name, phase in phases.items():
+ if phase_name == "order" or not isinstance(phase, dict):
+ continue
+ if isinstance(phase.get("annotation_schemes"), list):
+ schemes.extend(s for s in phase["annotation_schemes"] if isinstance(s, dict))
+
+ training = config.get("training")
+ if isinstance(training, dict) and isinstance(training.get("annotation_schemes"), list):
+ schemes.extend(s for s in training["annotation_schemes"] if isinstance(s, dict))
+
+ return schemes
+
+
+class HighlightToSpanRule(MigrationRule):
+ """Rename the legacy ``highlight`` annotation type to ``span``.
+
+ The span annotation type was named ``highlight`` in v1.x. In v2 it is
+ ``span`` and ``highlight`` is rejected at boot with "Unknown annotation
+ type: highlight". This is the single most common v1->v2 breaking change
+ (MIGRATION.md), so the migrator must apply it or the "migrated" config
+ still fails to start.
+ """
+
+ def __init__(self):
+ super().__init__(
+ "highlight_to_span",
+ "Rename annotation_type: highlight to span"
+ )
+
+ def applies(self, config: Dict[str, Any]) -> bool:
+ return any(
+ scheme.get("annotation_type") == "highlight"
+ for scheme in _collect_all_schemes(config)
+ )
+
+ def migrate(self, config: Dict[str, Any]) -> Tuple[Dict[str, Any], List[str]]:
+ config = copy.deepcopy(config)
+ changes = []
+ for scheme in _collect_all_schemes(config):
+ if scheme.get("annotation_type") == "highlight":
+ scheme["annotation_type"] = "span"
+ changes.append(
+ f"Renamed annotation_type: highlight -> span "
+ f"in scheme '{scheme.get('name', 'unknown')}'"
+ )
+ return config, changes
+
+
+class TextareaToMultilineRule(MigrationRule):
+ """Migrate textarea.on to multiline format."""
+
+ def __init__(self):
+ super().__init__(
+ "textarea_to_multiline",
+ "Convert textarea.on to multiline format in textbox schemas"
+ )
+
+ def applies(self, config: Dict[str, Any]) -> bool:
+ """Check if any annotation scheme uses old textarea format."""
+ schemes = self._get_all_schemes(config)
+ for scheme in schemes:
+ if scheme.get("annotation_type") == "text":
+ if "textarea" in scheme and isinstance(scheme["textarea"], dict):
+ if scheme["textarea"].get("on"):
+ return True
+ return False
+
+ def migrate(self, config: Dict[str, Any]) -> Tuple[Dict[str, Any], List[str]]:
+ """Convert textarea.on to multiline."""
+ config = copy.deepcopy(config)
+ changes = []
+
+ schemes = self._get_all_schemes(config)
+ for scheme in schemes:
+ if scheme.get("annotation_type") == "text":
+ if "textarea" in scheme and isinstance(scheme["textarea"], dict):
+ textarea = scheme["textarea"]
+ if textarea.get("on"):
+ # Convert to new format
+ scheme["multiline"] = True
+ if "rows" in textarea:
+ scheme["rows"] = textarea["rows"]
+ if "cols" in textarea:
+ scheme["cols"] = textarea["cols"]
+
+ # Remove old textarea config
+ del scheme["textarea"]
+
+ changes.append(
+ f"Converted textarea.on to multiline in schema '{scheme.get('name', 'unknown')}'"
+ )
+
+ return config, changes
+
+ def _get_all_schemes(self, config: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """Get all annotation schemes from config."""
+ schemes = []
+
+ # Top-level annotation_schemes
+ if "annotation_schemes" in config:
+ schemes.extend(config["annotation_schemes"])
+
+ # Phase-level annotation_schemes
+ if "phases" in config:
+ phases = config["phases"]
+ if isinstance(phases, list):
+ for phase in phases:
+ if "annotation_schemes" in phase:
+ schemes.extend(phase["annotation_schemes"])
+ elif isinstance(phases, dict):
+ for phase_name, phase in phases.items():
+ if phase_name != "order" and isinstance(phase, dict):
+ if "annotation_schemes" in phase:
+ schemes.extend(phase["annotation_schemes"])
+
+ # Training annotation_schemes
+ if "training" in config and isinstance(config["training"], dict):
+ if "annotation_schemes" in config["training"]:
+ for scheme in config["training"]["annotation_schemes"]:
+ if isinstance(scheme, dict):
+ schemes.append(scheme)
+
+ return schemes
+
+
+class LegacyUserConfigRule(MigrationRule):
+ """Migrate legacy user_config format."""
+
+ def __init__(self):
+ super().__init__(
+ "legacy_user_config",
+ "Migrate legacy user_config to login format"
+ )
+
+ def applies(self, config: Dict[str, Any]) -> bool:
+ """Check if config uses legacy user_config without login."""
+ has_user_config = "user_config" in config
+ has_login = "login" in config
+
+ # If has user_config but no login, and user_config has old format
+ if has_user_config and not has_login:
+ user_config = config["user_config"]
+ # Check for patterns that suggest old format
+ if "allow_all_users" in user_config:
+ return True
+ return False
+
+ def migrate(self, config: Dict[str, Any]) -> Tuple[Dict[str, Any], List[str]]:
+ """Migrate user_config to login format."""
+ config = copy.deepcopy(config)
+ changes = []
+
+ if "user_config" in config and "login" not in config:
+ user_config = config["user_config"]
+
+ # Determine login type based on user_config
+ if user_config.get("allow_all_users", False):
+ config["login"] = {
+ "type": "open",
+ }
+ changes.append("Added login.type: open (from allow_all_users: true)")
+ elif user_config.get("users"):
+ config["login"] = {
+ "type": "password",
+ }
+ changes.append("Added login.type: password (user list detected)")
+
+ # user_config is still valid, just add login section
+ changes.append("Note: user_config is still valid, login section added for clarity")
+
+ return config, changes
+
+
+class LegacyOutputFormatRule(MigrationRule):
+ """Suggest modern output format options."""
+
+ def __init__(self):
+ super().__init__(
+ "output_format",
+ "Suggest modern output format options"
+ )
+
+ def applies(self, config: Dict[str, Any]) -> bool:
+ """Check if config uses legacy output format."""
+ output_format = config.get("output_annotation_format", "")
+ return output_format in ["csv", "tsv"]
+
+ def migrate(self, config: Dict[str, Any]) -> Tuple[Dict[str, Any], List[str]]:
+ """Add note about JSON format being recommended."""
+ config = copy.deepcopy(config)
+ changes = []
+
+ output_format = config.get("output_annotation_format", "")
+ if output_format in ["csv", "tsv"]:
+ changes.append(
+ f"Note: output_annotation_format is '{output_format}'. "
+ f"Consider using 'json' for richer annotation data (spans, metadata)."
+ )
+
+ return config, changes
+
+
+class DeprecatedSiteConfigRule(MigrationRule):
+ """Handle deprecated site configuration options."""
+
+ def __init__(self):
+ super().__init__(
+ "deprecated_site_config",
+ "Migrate deprecated site configuration options"
+ )
+
+ def applies(self, config: Dict[str, Any]) -> bool:
+ """Check for deprecated site config options."""
+ # site_dir: "default" is still valid but could note about auto-generation
+ return config.get("site_dir") == "default"
+
+ def migrate(self, config: Dict[str, Any]) -> Tuple[Dict[str, Any], List[str]]:
+ """Add note about site auto-generation."""
+ config = copy.deepcopy(config)
+ changes = []
+
+ if config.get("site_dir") == "default":
+ changes.append(
+ "Note: site_dir: default uses auto-generated templates. "
+ "This is the recommended approach for v2."
+ )
+
+ return config, changes
+
+
+class LegacyLabelRequirementRule(MigrationRule):
+ """Migrate legacy label_requirement format."""
+
+ def __init__(self):
+ super().__init__(
+ "legacy_label_requirement",
+ "Ensure label_requirement uses modern format"
+ )
+
+ def applies(self, config: Dict[str, Any]) -> bool:
+ """Check for old label_requirement format."""
+ schemes = self._get_all_schemes(config)
+ for scheme in schemes:
+ if "label_requirement" in scheme:
+ lr = scheme["label_requirement"]
+ # Check if it's a simple boolean instead of dict
+ if isinstance(lr, bool):
+ return True
+ return False
+
+ def migrate(self, config: Dict[str, Any]) -> Tuple[Dict[str, Any], List[str]]:
+ """Convert label_requirement boolean to dict format."""
+ config = copy.deepcopy(config)
+ changes = []
+
+ schemes = self._get_all_schemes(config)
+ for scheme in schemes:
+ if "label_requirement" in scheme:
+ lr = scheme["label_requirement"]
+ if isinstance(lr, bool):
+ scheme["label_requirement"] = {"required": lr}
+ changes.append(
+ f"Converted label_requirement: {lr} to label_requirement.required: {lr} "
+ f"in schema '{scheme.get('name', 'unknown')}'"
+ )
+
+ return config, changes
+
+ def _get_all_schemes(self, config: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """Get all annotation schemes from config."""
+ schemes = []
+ if "annotation_schemes" in config:
+ schemes.extend(config["annotation_schemes"])
+ if "phases" in config:
+ phases = config["phases"]
+ if isinstance(phases, list):
+ for phase in phases:
+ if "annotation_schemes" in phase:
+ schemes.extend(phase["annotation_schemes"])
+ elif isinstance(phases, dict):
+ for phase_name, phase in phases.items():
+ if phase_name != "order" and isinstance(phase, dict):
+ if "annotation_schemes" in phase:
+ schemes.extend(phase["annotation_schemes"])
+ return schemes
+
+
+# All migration rules in order of application
+MIGRATION_RULES = [
+ HighlightToSpanRule(),
+ TextareaToMultilineRule(),
+ LegacyLabelRequirementRule(),
+ LegacyUserConfigRule(),
+ LegacyOutputFormatRule(),
+ DeprecatedSiteConfigRule(),
+]
+
+
+def migrate_config(config: Dict[str, Any], rules: List[MigrationRule] = None) -> Tuple[Dict[str, Any], List[str]]:
+ """
+ Apply all migration rules to a configuration.
+
+ Args:
+ config: The configuration dictionary to migrate
+ rules: Optional list of rules to apply (defaults to all rules)
+
+ Returns:
+ Tuple of (migrated_config, list of all changes)
+ """
+ if rules is None:
+ rules = MIGRATION_RULES
+
+ all_changes = []
+ current_config = copy.deepcopy(config)
+
+ for rule in rules:
+ if rule.applies(current_config):
+ current_config, changes = rule.migrate(current_config)
+ if changes:
+ all_changes.append(f"\n[{rule.name}] {rule.description}:")
+ all_changes.extend([f" - {change}" for change in changes])
+
+ return current_config, all_changes
+
+
+def load_yaml(file_path: str) -> Dict[str, Any]:
+ """Load a YAML configuration file."""
+ with open(file_path, 'r', encoding='utf-8') as f:
+ return yaml.safe_load(f)
+
+
+def save_yaml(config: Dict[str, Any], file_path: str) -> None:
+ """Save a configuration to a YAML file."""
+ with open(file_path, 'w', encoding='utf-8') as f:
+ yaml.dump(config, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
+
+
+def format_yaml(config: Dict[str, Any]) -> str:
+ """Format configuration as YAML string."""
+ return yaml.dump(config, default_flow_style=False, allow_unicode=True, sort_keys=False)
+
+
+def migrate_arguments():
+ """Create argument parser for migrate command."""
+ parser = argparse.ArgumentParser(
+ description="Migrate Potato configuration files to v2 format",
+ prog="potato migrate"
+ )
+
+ parser.add_argument(
+ "config_file",
+ help="Path to the configuration file to migrate"
+ )
+
+ parser.add_argument(
+ "--to-v2",
+ action="store_true",
+ dest="to_v2",
+ help="Migrate to v2 format (required)",
+ required=True
+ )
+
+ parser.add_argument(
+ "--output", "-o",
+ dest="output_file",
+ help="Output file path (default: print to stdout)"
+ )
+
+ parser.add_argument(
+ "--in-place", "-i",
+ action="store_true",
+ dest="in_place",
+ help="Modify the config file in place"
+ )
+
+ parser.add_argument(
+ "--dry-run",
+ action="store_true",
+ dest="dry_run",
+ help="Show what changes would be made without applying them"
+ )
+
+ parser.add_argument(
+ "--quiet", "-q",
+ action="store_true",
+ dest="quiet",
+ help="Suppress informational output"
+ )
+
+ return parser
+
+
+def main(args=None):
+ """Main entry point for the migrate command."""
+ parser = migrate_arguments()
+
+ if args is None:
+ args = parser.parse_args()
+ else:
+ args = parser.parse_args(args)
+
+ # Validate arguments
+ if args.in_place and args.output_file:
+ print("Error: Cannot use both --in-place and --output together", file=sys.stderr)
+ return 1
+
+ # Check config file exists
+ if not os.path.exists(args.config_file):
+ print(f"Error: Configuration file not found: {args.config_file}", file=sys.stderr)
+ return 1
+
+ # Load configuration
+ try:
+ config = load_yaml(args.config_file)
+ except yaml.YAMLError as e:
+ print(f"Error: Invalid YAML in configuration file: {e}", file=sys.stderr)
+ return 1
+ except Exception as e:
+ print(f"Error: Failed to read configuration file: {e}", file=sys.stderr)
+ return 1
+
+ if config is None:
+ print("Error: Configuration file is empty", file=sys.stderr)
+ return 1
+
+ # Apply migrations
+ migrated_config, changes = migrate_config(config)
+
+ # Report changes
+ if not args.quiet:
+ if changes:
+ print("Migration changes:", file=sys.stderr)
+ for change in changes:
+ print(change, file=sys.stderr)
+ print("", file=sys.stderr)
+ else:
+ print("No migrations needed - config is already up to date.", file=sys.stderr)
+
+ # Handle dry-run
+ if args.dry_run:
+ if not args.quiet:
+ print("Dry run - no changes written.", file=sys.stderr)
+ if changes:
+ print("\nMigrated configuration would be:", file=sys.stderr)
+ print(format_yaml(migrated_config))
+ return 0
+
+ # Output the migrated config
+ if args.in_place:
+ save_yaml(migrated_config, args.config_file)
+ if not args.quiet:
+ print(f"Updated {args.config_file} in place.", file=sys.stderr)
+ elif args.output_file:
+ save_yaml(migrated_config, args.output_file)
+ if not args.quiet:
+ print(f"Wrote migrated config to {args.output_file}", file=sys.stderr)
+ else:
+ # Print to stdout
+ print(format_yaml(migrated_config))
+
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/potato/password_reset.py b/potato/password_reset.py
new file mode 100644
index 0000000000000000000000000000000000000000..6147e4f0759ada4bcb2c0bc99e25d68dc71fc0b1
--- /dev/null
+++ b/potato/password_reset.py
@@ -0,0 +1,61 @@
+"""
+CLI password reset utility for the Potato annotation platform.
+
+Usage:
+ python potato/flask_server.py reset-password --username
+"""
+
+import getpass
+import logging
+import sys
+
+from potato.server_utils.config_module import init_config, config
+from potato.authentication import UserAuthenticator
+
+logger = logging.getLogger(__name__)
+
+
+def cli_reset_password(args):
+ """Reset a user's password from the command line.
+
+ Args:
+ args: Parsed command-line arguments (must include config_file, optionally username)
+ """
+ # Initialize config (loads YAML, sets up paths)
+ init_config(args)
+
+ # Initialize authenticator from config
+ authenticator = UserAuthenticator.init_from_config(config)
+
+ # Get username
+ username = args.username
+ if not username:
+ username = input("Username: ").strip()
+ if not username:
+ print("Error: Username is required.")
+ sys.exit(1)
+
+ # Check that user exists
+ if not authenticator.is_valid_username(username):
+ print(f"Error: User '{username}' does not exist.")
+ sys.exit(1)
+
+ # Get new password
+ new_password = getpass.getpass("New password: ")
+ confirm_password = getpass.getpass("Confirm password: ")
+
+ if new_password != confirm_password:
+ print("Error: Passwords do not match.")
+ sys.exit(1)
+
+ if not new_password:
+ print("Error: Password cannot be empty.")
+ sys.exit(1)
+
+ # Update password
+ if authenticator.update_password(username, new_password):
+ authenticator.save_user_config()
+ print(f"Password for '{username}' has been reset successfully.")
+ else:
+ print(f"Error: Failed to reset password for '{username}'.")
+ sys.exit(1)
diff --git a/potato/persistence/__init__.py b/potato/persistence/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d0743988d5726d6d04487ff8be0a72474760b63
--- /dev/null
+++ b/potato/persistence/__init__.py
@@ -0,0 +1,34 @@
+"""
+Universal Persistence Layer
+
+SQLite-backed persistence for features that need project-scoped state:
+memos, search index, codebook, cases, queries, smart codes.
+
+A single `/project.sqlite` file holds all tables. Modules register
+their schemas via the migration registry; on each `get_db()` call any pending
+migrations run inside a transaction.
+
+This package is universal โ not gated to QDA Mode. Both standard annotation
+projects and QDA Mode projects can read/write to the same DB; QDA Mode just
+populates more tables.
+"""
+
+from .sqlite import (
+ Migration,
+ get_db,
+ close_db,
+ clear_db_cache,
+ clear_migrations,
+ register_migration,
+ registered_migrations,
+)
+
+__all__ = [
+ "Migration",
+ "get_db",
+ "close_db",
+ "clear_db_cache",
+ "clear_migrations",
+ "register_migration",
+ "registered_migrations",
+]
diff --git a/potato/persistence/sqlite.py b/potato/persistence/sqlite.py
new file mode 100644
index 0000000000000000000000000000000000000000..0350e9cac68da8d0954f12207cbf2edad865c5de
--- /dev/null
+++ b/potato/persistence/sqlite.py
@@ -0,0 +1,202 @@
+"""
+Universal SQLite Helper
+
+Provides a process-wide cache of WAL-mode SQLite connections keyed by task_dir,
+plus a migration registry so each feature module declares the tables it owns.
+
+Design choices:
+- One database file per project: `/project.sqlite`.
+- WAL journal mode for concurrent reads alongside writes.
+- `foreign_keys = ON` enforced on every connection.
+- Migrations are idempotent and tracked in a `schema_migrations` table.
+- Modules register migrations at import time via `register_migration()`;
+ pending migrations run on first `get_db(task_dir)` call.
+- Thread-safe: a per-process lock guards the cache and the migration runner.
+
+Usage:
+ from potato.persistence import register_migration, Migration, get_db
+
+ register_migration(Migration(
+ name="0001_memos",
+ sql=\"""CREATE TABLE IF NOT EXISTS memos (
+ id TEXT PRIMARY KEY,
+ ...
+ );\""",
+ ))
+
+ conn = get_db(task_dir)
+ conn.execute("INSERT INTO memos ...")
+ conn.commit()
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import sqlite3
+import threading
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class Migration:
+ """A single, idempotent schema migration.
+
+ Attributes:
+ name: Unique identifier (e.g. "0001_memos"). Used as the migration
+ key in the `schema_migrations` table โ running twice is a no-op.
+ sql: SQL statement(s). Multi-statement scripts allowed; runs via
+ ``connection.executescript``.
+ """
+ name: str
+ sql: str
+
+
+_MIGRATIONS: List[Migration] = []
+_MIGRATION_NAMES: set = set()
+_REGISTRY_LOCK = threading.Lock()
+
+_DB_CACHE: Dict[str, sqlite3.Connection] = {}
+_DB_CACHE_LOCK = threading.Lock()
+
+_SCHEMA_MIGRATIONS_DDL = """
+CREATE TABLE IF NOT EXISTS schema_migrations (
+ name TEXT PRIMARY KEY,
+ applied_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
+);
+"""
+
+
+def register_migration(migration: Migration) -> None:
+ """Register a migration so it runs on the next `get_db()` call.
+
+ Re-registering a migration with the same name is a no-op; this lets
+ modules call `register_migration` unconditionally at import time without
+ worrying about double-imports.
+ """
+ with _REGISTRY_LOCK:
+ if migration.name in _MIGRATION_NAMES:
+ return
+ _MIGRATIONS.append(migration)
+ _MIGRATION_NAMES.add(migration.name)
+ logger.debug(f"Registered migration: {migration.name}")
+
+
+def registered_migrations() -> List[Migration]:
+ """Return a copy of the current migration registry (in registration order)."""
+ with _REGISTRY_LOCK:
+ return list(_MIGRATIONS)
+
+
+def clear_migrations() -> None:
+ """Reset the process-global migration registry. Tests only.
+
+ The registry is process-global and migrations are normally registered
+ once at import time. A test that registers an ad-hoc migration would
+ otherwise leak it into every later test's `get_db()` in the same
+ pytest process. Call this (and `clear_db_cache()`) for isolation.
+ """
+ with _REGISTRY_LOCK:
+ _MIGRATIONS.clear()
+ _MIGRATION_NAMES.clear()
+
+
+def get_db(task_dir: str) -> sqlite3.Connection:
+ """Return the cached WAL-mode SQLite connection for this project.
+
+ On first call for a given task_dir, opens `/project.sqlite`,
+ sets WAL + foreign_keys, and runs any pending migrations.
+
+ Connections are cached per task_dir and reused across requests. They
+ are NOT thread-local โ SQLite connections created with
+ `check_same_thread=False` are safe for serialized access from multiple
+ threads, which matches Flask's per-request threading model.
+ """
+ abs_dir = os.path.abspath(task_dir)
+ with _DB_CACHE_LOCK:
+ existing = _DB_CACHE.get(abs_dir)
+ if existing is not None:
+ # Apply any migrations registered AFTER this connection was
+ # first opened. Idempotent and cheap (one indexed SELECT when
+ # nothing is pending). Without this, a feature whose module
+ # is imported/registered later than another feature's first
+ # get_db() call would never get its tables created.
+ _run_pending_migrations(existing)
+ return existing
+ os.makedirs(abs_dir, exist_ok=True)
+ db_path = os.path.join(abs_dir, "project.sqlite")
+ conn = sqlite3.connect(
+ db_path,
+ check_same_thread=False,
+ detect_types=sqlite3.PARSE_DECLTYPES,
+ isolation_level=None, # autocommit; callers manage transactions
+ )
+ conn.row_factory = sqlite3.Row
+ conn.execute("PRAGMA journal_mode = WAL")
+ conn.execute("PRAGMA foreign_keys = ON")
+ conn.execute("PRAGMA synchronous = NORMAL")
+ _run_pending_migrations(conn)
+ _DB_CACHE[abs_dir] = conn
+ logger.info(f"Opened SQLite project DB: {db_path}")
+ return conn
+
+
+def close_db(task_dir: str) -> None:
+ """Close and evict the cached connection for one task_dir."""
+ abs_dir = os.path.abspath(task_dir)
+ with _DB_CACHE_LOCK:
+ conn = _DB_CACHE.pop(abs_dir, None)
+ if conn is not None:
+ try:
+ conn.close()
+ except sqlite3.Error as e:
+ logger.warning(f"Error closing DB for {abs_dir}: {e}")
+
+
+def clear_db_cache() -> None:
+ """Close every cached connection. Primarily for tests."""
+ with _DB_CACHE_LOCK:
+ connections = list(_DB_CACHE.values())
+ _DB_CACHE.clear()
+ for conn in connections:
+ try:
+ conn.close()
+ except sqlite3.Error:
+ pass
+
+
+def _run_pending_migrations(conn: sqlite3.Connection) -> None:
+ """Apply migrations that haven't been recorded in schema_migrations yet."""
+ conn.executescript(_SCHEMA_MIGRATIONS_DDL)
+ applied = {
+ row["name"]
+ for row in conn.execute("SELECT name FROM schema_migrations").fetchall()
+ }
+ with _REGISTRY_LOCK:
+ pending = [m for m in _MIGRATIONS if m.name not in applied]
+
+ if not pending:
+ return
+
+ # Note on atomicity: Python's `executescript()` issues an implicit COMMIT
+ # at the start, which dissolves any transaction (BEGIN/COMMIT *or*
+ # SAVEPOINT/RELEASE) we wrap around it. So we don't wrap. The convention
+ # is that every migration uses idempotent DDL (`CREATE TABLE IF NOT
+ # EXISTS`, `CREATE INDEX IF NOT EXISTS`, etc.); if `executescript()`
+ # raises, the migration record is *not* inserted, and the migration will
+ # retry cleanly on the next `get_db()` call. `INSERT OR IGNORE` makes the
+ # success path safe against double-application from a benign race.
+ for migration in pending:
+ try:
+ conn.executescript(migration.sql)
+ conn.execute(
+ "INSERT OR IGNORE INTO schema_migrations (name) VALUES (?)",
+ (migration.name,),
+ )
+ logger.info(f"Applied migration: {migration.name}")
+ except sqlite3.Error as e:
+ logger.error(f"Migration {migration.name} failed: {e}")
+ raise
diff --git a/potato/phase.py b/potato/phase.py
new file mode 100644
index 0000000000000000000000000000000000000000..a83ebd1e360a7000242cf788b3bfc8662fc1ded4
--- /dev/null
+++ b/potato/phase.py
@@ -0,0 +1,91 @@
+"""
+User Phase Management System
+
+This module defines the phases that users progress through during the annotation process.
+Each phase represents a distinct stage in the annotation workflow, from initial login
+to completion of the annotation task.
+
+The phase system supports multi-page phases where users may need to complete multiple
+pages within a single phase type (e.g., multiple instruction pages).
+"""
+
+from __future__ import annotations
+from enum import Enum
+
+class UserPhase(Enum):
+ """
+ Enumeration of user phases in the annotation workflow.
+
+ Each phase represents a distinct stage that users must complete before
+ proceeding to the next phase. The phases are processed in order:
+ LOGIN -> CONSENT -> PRESTUDY -> INSTRUCTIONS -> TRAINING -> ANNOTATION -> POSTSTUDY -> DONE
+
+ Attributes:
+ LOGIN: Initial authentication phase
+ CONSENT: User consent and agreement phase
+ PRESTUDY: Pre-study questions or screening phase
+ INSTRUCTIONS: Task instructions and guidelines phase
+ TRAINING: Practice/training examples phase
+ ANNOTATION: Main annotation task phase
+ POSTSTUDY: Post-study questions or feedback phase
+ DONE: Completion phase
+ """
+ LOGIN = 'login'
+ CONSENT = 'consent'
+ PRESTUDY = 'prestudy'
+ INSTRUCTIONS = 'instructions'
+ TRAINING = 'training'
+ ANNOTATION = 'annotation'
+ POSTSTUDY = 'poststudy'
+ DONE = 'done'
+
+ #@classmethod
+ #def list(cls):
+ # return list(map(lambda c: c.value, cls))
+
+ def fromstr(phase: str) -> UserPhase:
+ """
+ Convert a string representation to a UserPhase enum value.
+
+ This method provides a safe way to convert string inputs (e.g., from
+ configuration files or API requests) to UserPhase enum values.
+
+ Args:
+ phase: String representation of the phase (case-insensitive)
+
+ Returns:
+ UserPhase: The corresponding enum value
+
+ Raises:
+ ValueError: If the string doesn't match any known phase
+
+ Example:
+ >>> UserPhase.fromstr("annotation")
+
+ """
+ phase = phase.lower()
+ if phase == "login":
+ return UserPhase.LOGIN
+ elif phase == "consent":
+ return UserPhase.CONSENT
+ elif phase == "prestudy":
+ return UserPhase.PRESTUDY
+ elif phase == "instructions":
+ return UserPhase.INSTRUCTIONS
+ elif phase == "training":
+ return UserPhase.TRAINING
+ elif phase == "annotation":
+ return UserPhase.ANNOTATION
+ elif phase == "poststudy":
+ return UserPhase.POSTSTUDY
+ else:
+ raise ValueError(f"Unknown phase: {phase}")
+
+ def __str__(self) -> str:
+ """
+ Return the string representation of the phase.
+
+ Returns:
+ str: The phase name as a string
+ """
+ return self.value
\ No newline at end of file
diff --git a/potato/preview_cli.py b/potato/preview_cli.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a46c8565c9782584cd60b3961b60f98ab1a694e
--- /dev/null
+++ b/potato/preview_cli.py
@@ -0,0 +1,499 @@
+#!/usr/bin/env python3
+"""
+Potato Preview CLI
+
+A command-line tool for previewing annotation task configurations.
+Helps administrators validate configs and see how schemas will render
+without running the full server.
+
+Usage:
+ potato preview config.yaml # Summary output (default)
+ potato preview config.yaml --format html # HTML output
+ potato preview config.yaml --format json # JSON output
+
+ # Or run as module:
+ python -m potato.preview_cli config.yaml
+"""
+
+import argparse
+import json
+import os
+import sys
+import yaml
+import logging
+from typing import Dict, Any, List, Tuple, Optional
+
+# Set up logging
+logging.basicConfig(level=logging.WARNING)
+logger = logging.getLogger(__name__)
+
+
+def load_config(config_path: str) -> Dict[str, Any]:
+ """
+ Load and parse a YAML configuration file.
+
+ Args:
+ config_path: Path to the configuration file
+
+ Returns:
+ Parsed configuration dictionary
+
+ Raises:
+ FileNotFoundError: If config file doesn't exist
+ yaml.YAMLError: If config is invalid YAML
+ """
+ if not os.path.exists(config_path):
+ raise FileNotFoundError(f"Configuration file not found: {config_path}")
+
+ with open(config_path, 'r', encoding='utf-8') as f:
+ config = yaml.safe_load(f)
+
+ if not isinstance(config, dict):
+ raise ValueError("Configuration must be a YAML object (dictionary)")
+
+ return config
+
+
+def validate_config(config: Dict[str, Any]) -> List[str]:
+ """
+ Validate configuration and return list of issues.
+
+ Args:
+ config: Configuration dictionary
+
+ Returns:
+ List of validation error/warning messages
+ """
+ issues = []
+
+ # Required fields
+ required = ['annotation_task_name', 'item_properties', 'task_dir', 'output_annotation_dir']
+ for field in required:
+ if field not in config:
+ issues.append(f"ERROR: Missing required field '{field}'")
+
+ # Data source validation
+ has_data_files = config.get('data_files') and len(config.get('data_files', [])) > 0
+ has_data_directory = bool(config.get('data_directory'))
+ if not has_data_files and not has_data_directory:
+ issues.append("ERROR: Must have either 'data_files' or 'data_directory'")
+
+ # Annotation schemes validation
+ has_schemes = 'annotation_schemes' in config
+ has_phases = 'phases' in config and config['phases']
+
+ if not has_schemes and not has_phases:
+ issues.append("ERROR: Must have either 'annotation_schemes' or 'phases'")
+
+ if has_schemes and has_phases:
+ # Check for potential conflict
+ if isinstance(config['phases'], list):
+ phases_with_schemes = [p.get('name', f'phase[{i}]')
+ for i, p in enumerate(config['phases'])
+ if 'annotation_schemes' in p]
+ else:
+ phases_with_schemes = [name for name, p in config['phases'].items()
+ if name != 'order' and isinstance(p, dict) and 'annotation_schemes' in p]
+
+ if phases_with_schemes:
+ issues.append(f"ERROR: Both top-level and phase-level annotation_schemes found in: {', '.join(phases_with_schemes)}")
+
+ return issues
+
+
+def get_annotation_schemes(config: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """
+ Extract all annotation schemes from config.
+
+ Args:
+ config: Configuration dictionary
+
+ Returns:
+ List of annotation scheme dictionaries
+ """
+ schemes = []
+
+ if 'annotation_schemes' in config:
+ schemes.extend(config['annotation_schemes'])
+
+ if 'phases' in config and config['phases']:
+ phases = config['phases']
+ if isinstance(phases, list):
+ for phase in phases:
+ if 'annotation_schemes' in phase:
+ schemes.extend(phase['annotation_schemes'])
+ else:
+ for name, phase in phases.items():
+ if name != 'order' and isinstance(phase, dict) and 'annotation_schemes' in phase:
+ schemes.extend(phase['annotation_schemes'])
+
+ return schemes
+
+
+def detect_keybinding_conflicts(schemes: List[Dict[str, Any]]) -> List[str]:
+ """
+ Detect keyboard shortcut conflicts across all schemes.
+
+ Args:
+ schemes: List of annotation scheme dictionaries
+
+ Returns:
+ List of conflict warning messages
+ """
+ conflicts = []
+ global_keys = {} # key -> (schema_name, label)
+
+ for scheme in schemes:
+ schema_name = scheme.get('name', 'unknown')
+ labels = scheme.get('labels', [])
+
+ for i, label_data in enumerate(labels):
+ key_value = None
+
+ # Check for explicit key_value
+ if isinstance(label_data, dict) and 'key_value' in label_data:
+ key_value = str(label_data['key_value'])
+ label_name = label_data.get('name', f'label[{i}]')
+ elif scheme.get('sequential_key_binding') and len(labels) <= 10:
+ key_value = str((i + 1) % 10)
+ label_name = label_data if isinstance(label_data, str) else label_data.get('name', f'label[{i}]')
+ else:
+ continue
+
+ if key_value:
+ key_id = f"{key_value}"
+ if key_id in global_keys:
+ prev_schema, prev_label = global_keys[key_id]
+ if prev_schema != schema_name: # Only warn for cross-schema conflicts
+ conflicts.append(
+ f"WARNING: Key '{key_value}' used by both "
+ f"'{prev_schema}:{prev_label}' and '{schema_name}:{label_name}'"
+ )
+ else:
+ global_keys[key_id] = (schema_name, label_name)
+
+ return conflicts
+
+
+def generate_preview_html(schemes: List[Dict[str, Any]]) -> str:
+ """
+ Generate HTML preview for annotation schemes.
+
+ Args:
+ schemes: List of annotation scheme dictionaries
+
+ Returns:
+ HTML string with rendered schemes
+ """
+ from potato.server_utils.schemas.registry import schema_registry
+
+ html_parts = []
+ all_keybindings = []
+
+ html_parts.append("""
+
+
+
+ Annotation Preview
+
+
+
+
+
+
Annotation Preview
+""")
+
+ for idx, scheme in enumerate(schemes):
+ scheme_name = scheme.get('name', 'unknown')
+ scheme_type = scheme.get('annotation_type', 'unknown')
+
+ html_parts.append(f"""
+
+
{scheme_name}
+
Type: {scheme_type}
+
+""")
+
+ try:
+ # Set annotation_id before generating (required by schema generators)
+ scheme["annotation_id"] = idx
+ html, keybindings = schema_registry.generate(scheme)
+ html_parts.append(html)
+ all_keybindings.extend(keybindings)
+ except Exception as e:
+ html_parts.append(f'
Error generating preview: {str(e)}
')
+
+ html_parts.append("
")
+
+ # Add keybindings summary
+ if all_keybindings:
+ html_parts.append("
Keyboard Shortcuts Key Action ")
+ for key, action in all_keybindings:
+ html_parts.append(f"{key} {action} ")
+ html_parts.append("
")
+
+ html_parts.append("
")
+
+ return "\n".join(html_parts)
+
+
+def generate_preview_json(config: Dict[str, Any], schemes: List[Dict[str, Any]], issues: List[str]) -> str:
+ """
+ Generate JSON preview output.
+
+ Args:
+ config: Full configuration dictionary
+ schemes: List of annotation schemes
+ issues: List of validation issues
+
+ Returns:
+ JSON string with preview data
+ """
+ from potato.server_utils.schemas.registry import schema_registry
+
+ result = {
+ "task_name": config.get('annotation_task_name', 'Unknown'),
+ "validation_issues": issues,
+ "schema_count": len(schemes),
+ "schemas": []
+ }
+
+ for idx, scheme in enumerate(schemes):
+ schema_info = {
+ "name": scheme.get('name'),
+ "type": scheme.get('annotation_type'),
+ "description": scheme.get('description'),
+ "labels": None,
+ "keybindings": [],
+ "error": None
+ }
+
+ # Extract labels
+ if 'labels' in scheme:
+ labels = scheme['labels']
+ schema_info['labels'] = [
+ l if isinstance(l, str) else l.get('name', str(l))
+ for l in labels
+ ]
+
+ # Try to generate and get keybindings
+ try:
+ # Set annotation_id before generating (required by schema generators)
+ scheme["annotation_id"] = idx
+ _, keybindings = schema_registry.generate(scheme)
+ schema_info['keybindings'] = [{"key": k, "action": a} for k, a in keybindings]
+ except Exception as e:
+ schema_info['error'] = str(e)
+
+ result['schemas'].append(schema_info)
+
+ return json.dumps(result, indent=2)
+
+
+def generate_preview_summary(config: Dict[str, Any], schemes: List[Dict[str, Any]],
+ issues: List[str], conflicts: List[str]) -> str:
+ """
+ Generate text summary preview.
+
+ Args:
+ config: Full configuration dictionary
+ schemes: List of annotation schemes
+ issues: List of validation issues
+ conflicts: List of keybinding conflicts
+
+ Returns:
+ Text summary string
+ """
+ lines = []
+ lines.append("=" * 60)
+ lines.append(f"ANNOTATION TASK PREVIEW")
+ lines.append("=" * 60)
+ lines.append(f"Task Name: {config.get('annotation_task_name', 'Unknown')}")
+ lines.append(f"Task Directory: {config.get('task_dir', 'Not set')}")
+ lines.append("")
+
+ # Validation issues
+ if issues:
+ lines.append("VALIDATION ISSUES:")
+ for issue in issues:
+ lines.append(f" {issue}")
+ lines.append("")
+ else:
+ lines.append("Validation: PASSED")
+ lines.append("")
+
+ # Keybinding conflicts
+ if conflicts:
+ lines.append("KEYBINDING CONFLICTS:")
+ for conflict in conflicts:
+ lines.append(f" {conflict}")
+ lines.append("")
+
+ # Schema summary
+ lines.append(f"ANNOTATION SCHEMAS ({len(schemes)} total):")
+ lines.append("-" * 40)
+
+ from potato.server_utils.schemas.registry import schema_registry
+
+ for idx, scheme in enumerate(schemes):
+ name = scheme.get('name', 'unknown')
+ ann_type = scheme.get('annotation_type', 'unknown')
+ desc = scheme.get('description', '')[:50]
+
+ lines.append(f" [{ann_type}] {name}")
+ if desc:
+ lines.append(f" {desc}...")
+
+ # Count labels if present
+ if 'labels' in scheme:
+ label_count = len(scheme['labels'])
+ lines.append(f" Labels: {label_count}")
+
+ # Try to get keybindings
+ try:
+ # Set annotation_id before generating (required by schema generators)
+ scheme["annotation_id"] = idx
+ _, keybindings = schema_registry.generate(scheme)
+ if keybindings:
+ lines.append(f" Keybindings: {len(keybindings)}")
+ except Exception as e:
+ lines.append(f" ERROR: {str(e)}")
+
+ lines.append("")
+
+ lines.append("=" * 60)
+ return "\n".join(lines)
+
+
+def generate_layout_html(schemes: List[Dict[str, Any]]) -> str:
+ """
+ Generate just the task layout HTML snippet (no wrapper page).
+
+ This outputs the HTML that would go inside {{ TASK_LAYOUT }} in the
+ annotation template, allowing admins to prototype and debug their
+ task layout without running the full server.
+
+ Args:
+ schemes: List of annotation scheme dictionaries
+
+ Returns:
+ HTML string with the annotation schema div and all schema forms
+ """
+ from potato.server_utils.schemas.registry import schema_registry
+
+ html_parts = []
+ html_parts.append('')
+
+ for idx, scheme in enumerate(schemes):
+ # Set annotation_id before generating (required by schema generators)
+ scheme["annotation_id"] = idx
+ try:
+ html, _ = schema_registry.generate(scheme)
+ html_parts.append(html)
+ except Exception as e:
+ schema_name = scheme.get('name', 'unknown')
+ html_parts.append(f'')
+
+ html_parts.append('
')
+ return "\n".join(html_parts)
+
+
+def main():
+ """Main entry point for preview CLI."""
+ parser = argparse.ArgumentParser(
+ description="Preview annotation task configuration",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+Examples:
+ python -m potato.preview_cli config.yaml # Summary output
+ python -m potato.preview_cli config.yaml --format html # Full HTML page preview
+ python -m potato.preview_cli config.yaml --format json # JSON output
+ python -m potato.preview_cli config.yaml --layout-only # Just the task layout HTML snippet
+
+ # Save HTML to file:
+ python -m potato.preview_cli config.yaml --format html > preview.html
+
+ # Get just the annotation schema div for embedding:
+ python -m potato.preview_cli config.yaml --layout-only > task_layout.html
+"""
+ )
+
+ parser.add_argument(
+ 'config_file',
+ help='Path to YAML configuration file'
+ )
+ parser.add_argument(
+ '--format', '-f',
+ choices=['summary', 'html', 'json'],
+ default='summary',
+ help='Output format (default: summary)'
+ )
+ parser.add_argument(
+ '--layout-only', '-l',
+ action='store_true',
+ help='Output only the task layout HTML snippet (no wrapper page). This is the HTML that goes inside {{ TASK_LAYOUT }}.'
+ )
+ parser.add_argument(
+ '--verbose', '-v',
+ action='store_true',
+ help='Enable verbose output'
+ )
+
+ args = parser.parse_args()
+
+ if args.verbose:
+ logging.getLogger().setLevel(logging.DEBUG)
+
+ try:
+ # Load configuration
+ config = load_config(args.config_file)
+
+ # Validate
+ issues = validate_config(config)
+
+ # Get schemes
+ schemes = get_annotation_schemes(config)
+ if not schemes:
+ print("WARNING: No annotation schemes found in configuration", file=sys.stderr)
+
+ # Detect conflicts
+ conflicts = detect_keybinding_conflicts(schemes)
+
+ # Generate output
+ if args.layout_only:
+ # Output just the task layout HTML snippet
+ print(generate_layout_html(schemes))
+ elif args.format == 'html':
+ print(generate_preview_html(schemes))
+ elif args.format == 'json':
+ print(generate_preview_json(config, schemes, issues))
+ else: # summary
+ print(generate_preview_summary(config, schemes, issues, conflicts))
+
+ # Exit with error code if there are issues (skip for layout-only mode)
+ if not args.layout_only:
+ error_count = len([i for i in issues if i.startswith('ERROR')])
+ sys.exit(1 if error_count > 0 else 0)
+
+ except FileNotFoundError as e:
+ print(f"Error: {e}", file=sys.stderr)
+ sys.exit(1)
+ except yaml.YAMLError as e:
+ print(f"Error: Invalid YAML in configuration file: {e}", file=sys.stderr)
+ sys.exit(1)
+ except Exception as e:
+ print(f"Error: {e}", file=sys.stderr)
+ if args.verbose:
+ import traceback
+ traceback.print_exc()
+ sys.exit(1)
+
+
+if __name__ == '__main__':
+ main()
diff --git a/potato/qda_mode/__init__.py b/potato/qda_mode/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..dfb95562a283fda4696ec98b02ea55ff1291d83f
--- /dev/null
+++ b/potato/qda_mode/__init__.py
@@ -0,0 +1,44 @@
+"""
+QDA Mode โ Qualitative Data Analysis workflow for Potato.
+
+QDA Mode is a distinct mode of Potato, parallel to `potato/solo_mode/`. When
+enabled, it composes a set of universal features (memos, search, cases,
+queries, label-explorer) with QDA-Mode-only additions (codebook, smart codes,
+optional network editor) and tunes their defaults for qualitative-coding
+workflows.
+
+Activation:
+ yaml
+ qda_mode:
+ enabled: true
+
+The standard annotation surface is unaffected when `qda_mode.enabled` is false
+or the block is absent.
+
+Architecture:
+ See `internal/qda-implementation-plan.md` and
+ `internal/qda-redesign-design.md` for the full design rationale and the
+ phased rollout plan. The role of `qda_mode/` is composition and defaults;
+ universal features live in their own top-level packages
+ (`potato/memos/`, `potato/search/`, `potato/cases/`, `potato/queries/`,
+ `potato/analytics/`, `potato/media_sync/`).
+"""
+
+from .config import QDAModeConfig, parse_qda_mode_config
+from .manager import (
+ QDAModeManager,
+ init_qda_mode_manager,
+ get_qda_mode_manager,
+ clear_qda_mode_manager,
+)
+from .routes import qda_mode_bp
+
+__all__ = [
+ "QDAModeConfig",
+ "parse_qda_mode_config",
+ "QDAModeManager",
+ "init_qda_mode_manager",
+ "get_qda_mode_manager",
+ "clear_qda_mode_manager",
+ "qda_mode_bp",
+]
diff --git a/potato/qda_mode/config.py b/potato/qda_mode/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..c6a0f1ced790de583e1d33f0f7ef72c656735ddb
--- /dev/null
+++ b/potato/qda_mode/config.py
@@ -0,0 +1,103 @@
+"""
+QDA Mode Configuration
+
+Parses and validates the `qda_mode:` block in a Potato YAML config. Features
+populated in later phases (codebook, queries, cases) live under sub-blocks
+here; Phase 0 ships memos + search.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MemosConfig:
+ """QDA Mode memo defaults. Memo storage itself is universal
+ (`potato/memos/`); this block only controls QDA-Mode-specific defaults
+ such as whether the sidebar is visible by default."""
+ enabled: bool = True
+ show_sidebar_by_default: bool = True
+
+
+@dataclass
+class CodebookConfig:
+ """Codebook is the QDA-Mode-only mutable runtime label set. Lands in
+ Phase 1; this scaffolding is here so configs can already declare it."""
+ enabled: bool = True
+ mode: str = "open" # 'open' | 'extensible' | 'fixed'
+
+
+@dataclass
+class QDAModeConfig:
+ """Top-level QDA Mode config.
+
+ Sub-blocks for codebook, cases, queries, etc. land in subsequent phases;
+ they're declared here as Optional so configs can write forward-compatible
+ YAML now without breaking when those features ship.
+ """
+ enabled: bool = False
+ memos: MemosConfig = field(default_factory=MemosConfig)
+ codebook: Optional[CodebookConfig] = None
+ # Sub-blocks added in later phases (typed as Any to keep this stable):
+ # cases, queries, smart_codes, network, media_sync
+
+ # Free-form passthrough for forward compatibility โ anything we haven't
+ # explicitly modeled gets preserved here for later phases to consume.
+ extras: Dict[str, Any] = field(default_factory=dict)
+
+ def validate(self) -> List[str]:
+ """Return a list of validation errors; empty means valid."""
+ errors: List[str] = []
+ if self.codebook and self.codebook.mode not in {"open", "extensible", "fixed"}:
+ errors.append(
+ f"qda_mode.codebook.mode must be one of "
+ f"'open' | 'extensible' | 'fixed' (got {self.codebook.mode!r})"
+ )
+ return errors
+
+
+_KNOWN_TOP_LEVEL_KEYS = {"enabled", "memos", "codebook"}
+
+
+def parse_qda_mode_config(config_data: Dict[str, Any]) -> QDAModeConfig:
+ """Parse the `qda_mode:` block out of a full Potato config dict.
+
+ Returns an empty (disabled) QDAModeConfig if the block is missing.
+ Unknown keys are preserved in `extras` so forward-compatible YAML
+ (e.g. a `queries:` block written before Phase 3 ships) doesn't error.
+ """
+ raw = config_data.get("qda_mode") or {}
+ if not isinstance(raw, dict):
+ logger.warning(
+ "qda_mode block must be a mapping; got %r โ treating as disabled",
+ type(raw).__name__,
+ )
+ return QDAModeConfig()
+
+ memos_raw = raw.get("memos") or {}
+ memos = MemosConfig(
+ enabled=bool(memos_raw.get("enabled", True)),
+ show_sidebar_by_default=bool(memos_raw.get("show_sidebar_by_default", True)),
+ ) if isinstance(memos_raw, dict) else MemosConfig()
+
+ codebook = None
+ codebook_raw = raw.get("codebook")
+ if isinstance(codebook_raw, dict):
+ codebook = CodebookConfig(
+ enabled=bool(codebook_raw.get("enabled", True)),
+ mode=str(codebook_raw.get("mode", "open")),
+ )
+
+ extras = {k: v for k, v in raw.items() if k not in _KNOWN_TOP_LEVEL_KEYS}
+
+ return QDAModeConfig(
+ enabled=bool(raw.get("enabled", False)),
+ memos=memos,
+ codebook=codebook,
+ extras=extras,
+ )
diff --git a/potato/qda_mode/manager.py b/potato/qda_mode/manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..f05f1663752f329276e02ecf4722194c3b6db385
--- /dev/null
+++ b/potato/qda_mode/manager.py
@@ -0,0 +1,93 @@
+"""
+QDA Mode Manager
+
+Singleton orchestrator for QDA Mode. Currently minimal โ owns the parsed
+config and the task_dir reference; later phases attach codebook, smart-codes,
+network, and other QDA-only state here.
+
+The pattern mirrors potato/solo_mode/manager.py: a module-level singleton
+guarded by a lock, with init/get/clear helpers used by flask_server.py.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any, Dict, Optional
+
+from .config import QDAModeConfig, parse_qda_mode_config
+
+logger = logging.getLogger(__name__)
+
+
+_QDA_MODE_MANAGER: Optional["QDAModeManager"] = None
+_QDA_MODE_LOCK = threading.Lock()
+
+
+class QDAModeManager:
+ """Holds parsed QDA Mode state for the running server."""
+
+ def __init__(self, qda_config: QDAModeConfig, full_config: Dict[str, Any]):
+ self.config = qda_config
+ self.task_dir: str = full_config.get("task_dir") or "."
+ self._full_config = full_config
+
+ def shutdown(self) -> None:
+ """Release any held resources. No-op today; later phases close DBs etc."""
+ logger.info("QDAModeManager shutdown")
+
+ def __repr__(self) -> str:
+ return (
+ f"QDAModeManager(enabled={self.config.enabled}, "
+ f"task_dir={self.task_dir!r}, "
+ f"memos={self.config.memos.enabled}, "
+ f"codebook={self.config.codebook is not None})"
+ )
+
+
+def init_qda_mode_manager(config_data: Dict[str, Any]) -> Optional[QDAModeManager]:
+ """Initialize the singleton QDAModeManager from a full Potato config.
+
+ Returns None (and leaves the singleton unset) when QDA Mode is disabled.
+ Calling twice is a no-op โ the second call returns the existing singleton.
+ """
+ global _QDA_MODE_MANAGER
+
+ with _QDA_MODE_LOCK:
+ if _QDA_MODE_MANAGER is not None:
+ return _QDA_MODE_MANAGER
+
+ qda_config = parse_qda_mode_config(config_data)
+ if not qda_config.enabled:
+ logger.info("QDA Mode disabled in config")
+ return None
+
+ errors = qda_config.validate()
+ if errors:
+ for err in errors:
+ logger.error(f"QDA Mode config error: {err}")
+ # Fail loud: an explicitly-enabled but misconfigured qda_mode
+ # block must abort startup, not silently boot with QDA off
+ # (the same silent-failure class as the legacy phase-type bug).
+ from potato.server_utils.config_module import ConfigValidationError
+ raise ConfigValidationError(
+ "Invalid qda_mode configuration: " + "; ".join(errors)
+ )
+
+ _QDA_MODE_MANAGER = QDAModeManager(qda_config, config_data)
+ logger.info(f"QDA Mode initialized: {_QDA_MODE_MANAGER!r}")
+ return _QDA_MODE_MANAGER
+
+
+def get_qda_mode_manager() -> Optional[QDAModeManager]:
+ """Get the singleton QDAModeManager, or None if QDA Mode is disabled."""
+ return _QDA_MODE_MANAGER
+
+
+def clear_qda_mode_manager() -> None:
+ """Clear the singleton. Primarily for tests."""
+ global _QDA_MODE_MANAGER
+ with _QDA_MODE_LOCK:
+ if _QDA_MODE_MANAGER is not None:
+ _QDA_MODE_MANAGER.shutdown()
+ _QDA_MODE_MANAGER = None
diff --git a/potato/qda_mode/routes.py b/potato/qda_mode/routes.py
new file mode 100644
index 0000000000000000000000000000000000000000..9de880695a0b01019e8cb0ec3970c58cc2472056
--- /dev/null
+++ b/potato/qda_mode/routes.py
@@ -0,0 +1,98 @@
+"""
+QDA Mode Routes
+
+Flask blueprint for QDA-Mode-specific endpoints. Mounted at `/qda` so it
+never collides with the universal `/admin/api/*` or `/api/*` namespaces.
+
+In Phase 0 this is intentionally minimal: a status endpoint to verify the
+mode is wired correctly. Memos and search endpoints land in their own
+universal packages (`potato/memos/api.py`, `potato/search/api.py`) โ not
+here โ because those features are not QDA-Mode-only.
+
+Later phases add codebook, smart-codes, and network endpoints here.
+"""
+
+from __future__ import annotations
+
+import logging
+from functools import wraps
+
+from flask import Blueprint, jsonify
+
+from .manager import get_qda_mode_manager
+from potato.codebook import Codebook
+from potato.cases import list_cases as _list_cases
+
+logger = logging.getLogger(__name__)
+
+
+qda_mode_bp = Blueprint("qda_mode", __name__, url_prefix="/qda")
+
+
+def qda_mode_required(view):
+ """Decorator: return 503 if QDA Mode is not enabled.
+
+ Mirrors potato/solo_mode/routes.py:solo_mode_required, but returns 503
+ (Service Unavailable) instead of 400 โ the endpoint exists in the URL
+ space, the mode just isn't active in this deployment. Clients can tell
+ the difference between "wrong request" (400) and "wrong deployment"
+ (503).
+ """
+ @wraps(view)
+ def wrapper(*args, **kwargs):
+ manager = get_qda_mode_manager()
+ if manager is None:
+ return jsonify({
+ "error": "QDA Mode not enabled in this deployment.",
+ "hint": "Set qda_mode.enabled: true in your config.yaml.",
+ }), 503
+ return view(*args, **kwargs)
+ return wrapper
+
+
+@qda_mode_bp.route("/status", methods=["GET"])
+def qda_mode_status():
+ """Report whether QDA Mode is enabled and the resolved config summary.
+
+ Always returns 200 โ admins use this to confirm the mode is wired
+ correctly. When disabled, returns ``{"enabled": false}`` so the UI
+ can decide whether to render QDA panels.
+ """
+ manager = get_qda_mode_manager()
+ if manager is None:
+ return jsonify({"enabled": False})
+ cfg = manager.config
+ return jsonify({
+ "enabled": True,
+ "memos": {
+ "enabled": cfg.memos.enabled,
+ "show_sidebar_by_default": cfg.memos.show_sidebar_by_default,
+ },
+ "codebook": (
+ None if cfg.codebook is None else
+ {"enabled": cfg.codebook.enabled, "mode": cfg.codebook.mode}
+ ),
+ "task_dir": manager.task_dir,
+ })
+
+
+@qda_mode_bp.route("/codebook", methods=["GET"])
+@qda_mode_required
+def qda_codebook():
+ """QDA project codebook (tree + flat labels).
+
+ First real ``/qda/*`` route behind ``qda_mode_required`` โ returns
+ 503 on a served app where QDA Mode is not enabled (F5/T3). The
+ universal codebook CRUD API lives at ``/api/codebook``; this is the
+ QDA-scoped read view bound to the active QDA project.
+ """
+ manager = get_qda_mode_manager()
+ from potato.server_utils.config_module import config
+ project = config.get("annotation_task_name") or "default"
+ cb = Codebook.load(manager.task_dir, project)
+ return jsonify({
+ "project": project,
+ "labels": cb.labels(),
+ "tree": cb.as_tree(),
+ "cases": _list_cases(manager.task_dir, project),
+ })
diff --git a/potato/quality_control.py b/potato/quality_control.py
new file mode 100644
index 0000000000000000000000000000000000000000..6a83e50b79c0ccffeefe49fa0c630f9a1a647b7d
--- /dev/null
+++ b/potato/quality_control.py
@@ -0,0 +1,1030 @@
+"""
+Quality Control Module
+
+This module provides comprehensive quality control features for annotation projects:
+- Attention Checks: Inject known-answer items to verify annotator engagement
+- Gold Standards: Compare annotations against expert-labeled items for accuracy tracking
+- Pre-annotation Support: Pre-fill forms with model predictions
+
+The module integrates with ItemStateManager for item injection and UserStateManager
+for tracking results.
+"""
+
+import json
+import logging
+import random
+import threading
+from datetime import datetime
+from typing import Dict, List, Optional, Any, Tuple, Set
+from dataclasses import dataclass, field
+from pathlib import Path
+from collections import defaultdict
+
+logger = logging.getLogger(__name__)
+
+# Singleton instance
+_QUALITY_CONTROL_MANAGER = None
+_QUALITY_CONTROL_LOCK = threading.Lock()
+
+
+@dataclass
+class AttentionCheckResult:
+ """Result of an attention check evaluation."""
+ item_id: str
+ user_id: str
+ passed: bool
+ expected: Dict[str, Any]
+ actual: Dict[str, Any]
+ timestamp: datetime = field(default_factory=datetime.now)
+ response_time_seconds: Optional[float] = None
+
+
+@dataclass
+class GoldStandardResult:
+ """Result of a gold standard evaluation."""
+ item_id: str
+ user_id: str
+ correct: bool
+ gold_label: Dict[str, Any]
+ user_response: Dict[str, Any]
+ explanation: Optional[str] = None
+ timestamp: datetime = field(default_factory=datetime.now)
+
+
+@dataclass
+class QualityControlConfig:
+ """Configuration for quality control features."""
+ # Attention checks config
+ attention_checks_enabled: bool = False
+ attention_items_file: Optional[str] = None
+ attention_frequency: Optional[int] = None # Insert one every N items
+ attention_probability: Optional[float] = None # OR probability per item
+ attention_min_response_time: float = 0.0 # Minimum seconds (flag fast responses)
+ attention_warn_threshold: int = 2
+ attention_warn_message: str = "Please read items carefully before answering."
+ attention_block_threshold: int = 5
+ attention_block_message: str = "You have been blocked due to too many incorrect attention check responses."
+
+ # Gold standards config
+ gold_standards_enabled: bool = False
+ gold_items_file: Optional[str] = None
+ gold_mode: str = "mixed" # training, mixed, separate
+ gold_frequency: Optional[int] = None
+ gold_min_accuracy: float = 0.7
+ gold_evaluation_count: int = 10
+ gold_show_correct_answer: bool = False # Default to silent (admin-only tracking)
+ gold_show_explanation: bool = False # Default to silent (admin-only tracking)
+
+ # Auto-promotion: items become gold standards when annotators agree
+ gold_auto_promote_enabled: bool = False
+ gold_auto_promote_min_annotators: int = 3 # Minimum annotators before checking
+ gold_auto_promote_agreement: float = 1.0 # Agreement threshold (1.0 = unanimous)
+
+ # Pre-annotation config
+ pre_annotation_enabled: bool = False
+ pre_annotation_field: str = "predictions"
+ pre_annotation_allow_modification: bool = True
+ pre_annotation_show_confidence: bool = False
+ pre_annotation_highlight_threshold: float = 0.7
+
+
+class QualityControlManager:
+ """
+ Manages quality control features including attention checks, gold standards,
+ and pre-annotation support.
+ """
+
+ def __init__(self, config: Dict[str, Any], base_dir: str):
+ """
+ Initialize the quality control manager.
+
+ Args:
+ config: The full application configuration
+ base_dir: Base directory for resolving file paths
+ """
+ self.config = config
+ self.base_dir = base_dir
+ self.logger = logging.getLogger(__name__)
+ self._lock = threading.RLock()
+
+ # Parse configuration
+ self.qc_config = self._parse_config(config)
+
+ # Attention check data
+ self.attention_items: List[Dict] = []
+ self.attention_expected: Dict[str, Dict[str, Any]] = {} # item_id -> expected_answer
+ self.attention_results: Dict[str, List[AttentionCheckResult]] = defaultdict(list) # user_id -> results
+ self.user_items_since_attention: Dict[str, int] = defaultdict(int) # user_id -> count
+ self.user_items_since_gold: Dict[str, int] = defaultdict(int) # user_id -> count
+
+ # Gold standard data
+ self.gold_items: List[Dict] = []
+ self.gold_labels: Dict[str, Dict[str, Any]] = {} # item_id -> gold_label
+ self.gold_explanations: Dict[str, str] = {} # item_id -> explanation
+ self.gold_results: Dict[str, List[GoldStandardResult]] = defaultdict(list) # user_id -> results
+
+ # Pre-annotation data (stored per-item)
+ self.pre_annotations: Dict[str, Dict[str, Any]] = {} # item_id -> pre_annotation_data
+
+ # Auto-promotion tracking
+ self.item_annotations: Dict[str, Dict[str, Any]] = defaultdict(dict) # item_id -> {user_id -> response}
+ self.promoted_gold_items: List[Dict] = [] # Items promoted to gold via consensus
+ self.promoted_gold_labels: Dict[str, Dict[str, Any]] = {} # Promoted item_id -> consensus_label
+
+ # Load data files if configured
+ self._load_attention_checks()
+ self._load_gold_standards()
+
+ self.logger.info(f"QualityControlManager initialized: "
+ f"attention_checks={self.qc_config.attention_checks_enabled}, "
+ f"gold_standards={self.qc_config.gold_standards_enabled}, "
+ f"pre_annotation={self.qc_config.pre_annotation_enabled}")
+
+ def _parse_config(self, config: Dict[str, Any]) -> QualityControlConfig:
+ """Parse quality control configuration from the main config."""
+ qc = QualityControlConfig()
+
+ # Parse attention checks config
+ attn_config = config.get('attention_checks', {})
+ if attn_config.get('enabled', False):
+ qc.attention_checks_enabled = True
+ qc.attention_items_file = attn_config.get('items_file')
+ qc.attention_frequency = attn_config.get('frequency')
+ qc.attention_probability = attn_config.get('probability')
+ qc.attention_min_response_time = attn_config.get('min_response_time', 0.0)
+
+ failure_handling = attn_config.get('failure_handling', {})
+ qc.attention_warn_threshold = failure_handling.get('warn_threshold', 2)
+ qc.attention_warn_message = failure_handling.get('warn_message', qc.attention_warn_message)
+ qc.attention_block_threshold = failure_handling.get('block_threshold', 5)
+ qc.attention_block_message = failure_handling.get('block_message', qc.attention_block_message)
+
+ # Parse gold standards config
+ gold_config = config.get('gold_standards', {})
+ if gold_config.get('enabled', False):
+ qc.gold_standards_enabled = True
+ qc.gold_items_file = gold_config.get('items_file')
+ qc.gold_mode = gold_config.get('mode', 'mixed')
+ qc.gold_frequency = gold_config.get('frequency')
+
+ accuracy_config = gold_config.get('accuracy', {})
+ qc.gold_min_accuracy = accuracy_config.get('min_threshold', 0.7)
+ qc.gold_evaluation_count = accuracy_config.get('evaluation_count', 10)
+
+ feedback_config = gold_config.get('feedback', {})
+ qc.gold_show_correct_answer = feedback_config.get('show_correct_answer', False)
+ qc.gold_show_explanation = feedback_config.get('show_explanation', False)
+
+ # Auto-promotion settings
+ auto_promote = gold_config.get('auto_promote', {})
+ if auto_promote.get('enabled', False):
+ qc.gold_auto_promote_enabled = True
+ qc.gold_auto_promote_min_annotators = auto_promote.get('min_annotators', 3)
+ qc.gold_auto_promote_agreement = auto_promote.get('agreement_threshold', 1.0)
+
+ # Parse pre-annotation config
+ pre_config = config.get('pre_annotation', {})
+ if pre_config.get('enabled', False):
+ qc.pre_annotation_enabled = True
+ qc.pre_annotation_field = pre_config.get('field', 'predictions')
+ qc.pre_annotation_allow_modification = pre_config.get('allow_modification', True)
+ qc.pre_annotation_show_confidence = pre_config.get('show_confidence', False)
+ qc.pre_annotation_highlight_threshold = pre_config.get('highlight_low_confidence', 0.7)
+
+ return qc
+
+ def _load_json_or_jsonl(self, file_path: str) -> List[Dict[str, Any]]:
+ """Load items from a JSON array or JSONL file."""
+ with open(file_path, 'r', encoding='utf-8') as f:
+ raw = f.read()
+
+ try:
+ parsed = json.loads(raw)
+ if isinstance(parsed, list):
+ return parsed
+ except json.JSONDecodeError:
+ pass
+
+ items = []
+ for line_no, line in enumerate(raw.splitlines()):
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ items.append(json.loads(line))
+ except json.JSONDecodeError as e:
+ raise ValueError(
+ f"Invalid JSON at line {line_no + 1} in {file_path}: {e}"
+ ) from e
+
+ return items
+
+ def _load_attention_checks(self) -> None:
+ """Load attention check items from file."""
+ if not self.qc_config.attention_checks_enabled:
+ return
+
+ if not self.qc_config.attention_items_file:
+ self.logger.warning("Attention checks enabled but no items_file specified")
+ return
+
+ file_path = Path(self.base_dir) / self.qc_config.attention_items_file
+ if not file_path.exists():
+ self.logger.warning(f"Attention checks file not found: {file_path}")
+ return
+
+ try:
+ items = self._load_json_or_jsonl(str(file_path))
+
+ for item in items:
+ if 'id' not in item or 'expected_answer' not in item:
+ self.logger.warning(f"Attention check item missing required fields: {item}")
+ continue
+
+ self.attention_items.append(item)
+ self.attention_expected[item['id']] = item['expected_answer']
+
+ self.logger.info(f"Loaded {len(self.attention_items)} attention check items")
+
+ except (json.JSONDecodeError, ValueError) as e:
+ self.logger.error(f"Failed to parse attention checks file: {e}")
+ except Exception as e:
+ self.logger.error(f"Failed to load attention checks: {e}")
+
+ def _load_gold_standards(self) -> None:
+ """Load gold standard items from file."""
+ if not self.qc_config.gold_standards_enabled:
+ return
+
+ if not self.qc_config.gold_items_file:
+ self.logger.warning("Gold standards enabled but no items_file specified")
+ return
+
+ file_path = Path(self.base_dir) / self.qc_config.gold_items_file
+ if not file_path.exists():
+ self.logger.warning(f"Gold standards file not found: {file_path}")
+ return
+
+ try:
+ items = self._load_json_or_jsonl(str(file_path))
+
+ for item in items:
+ if 'id' not in item or 'gold_label' not in item:
+ self.logger.warning(f"Gold standard item missing required fields: {item}")
+ continue
+
+ self.gold_items.append(item)
+ self.gold_labels[item['id']] = item['gold_label']
+ if 'explanation' in item:
+ self.gold_explanations[item['id']] = item['explanation']
+
+ self.logger.info(f"Loaded {len(self.gold_items)} gold standard items")
+
+ except (json.JSONDecodeError, ValueError) as e:
+ self.logger.error(f"Failed to parse gold standards file: {e}")
+ except Exception as e:
+ self.logger.error(f"Failed to load gold standards: {e}")
+
+ # =========================================================================
+ # Attention Check Methods
+ # =========================================================================
+
+ def is_attention_check(self, item_id: str) -> bool:
+ """Check if an item is an attention check."""
+ return item_id in self.attention_expected
+
+ def should_inject_attention_check(self, user_id: str) -> bool:
+ """
+ Determine if an attention check should be injected for this user.
+
+ Args:
+ user_id: The user ID
+
+ Returns:
+ True if an attention check should be injected
+ """
+ if not self.qc_config.attention_checks_enabled or not self.attention_items:
+ return False
+
+ with self._lock:
+ # Frequency-based injection
+ if self.qc_config.attention_frequency:
+ items_since = self.user_items_since_attention.get(user_id, 0)
+ return items_since >= self.qc_config.attention_frequency
+
+ # Probability-based injection
+ if self.qc_config.attention_probability:
+ return random.random() < self.qc_config.attention_probability
+
+ return False
+
+ def get_attention_check_item(self, user_id: str) -> Optional[Dict]:
+ """
+ Get a random attention check item for a user.
+
+ Args:
+ user_id: The user ID
+
+ Returns:
+ An attention check item dict, or None if none available
+ """
+ if not self.attention_items:
+ return None
+
+ with self._lock:
+ # Get items this user hasn't seen yet
+ seen_ids = {r.item_id for r in self.attention_results.get(user_id, [])}
+ available = [item for item in self.attention_items if item['id'] not in seen_ids]
+
+ if not available:
+ # Recycle items if all have been seen
+ available = self.attention_items
+
+ selected = random.choice(available)
+ # Reset counter
+ self.user_items_since_attention[user_id] = 0
+ return selected
+
+ def record_regular_item(self, user_id: str) -> None:
+ """Record that a user annotated a regular (non-attention-check) item."""
+ with self._lock:
+ self.user_items_since_attention[user_id] = self.user_items_since_attention.get(user_id, 0) + 1
+ self.user_items_since_gold[user_id] = self.user_items_since_gold.get(user_id, 0) + 1
+
+ def validate_attention_response(
+ self,
+ user_id: str,
+ item_id: str,
+ response: Dict[str, Any],
+ response_time_seconds: Optional[float] = None
+ ) -> Optional[Dict[str, Any]]:
+ """
+ Validate a response to an attention check.
+
+ Args:
+ user_id: The user ID
+ item_id: The attention check item ID
+ response: The user's response (schema_name -> value)
+ response_time_seconds: Time taken to respond
+
+ Returns:
+ Dict with validation result if this is an attention check, None otherwise.
+ Result includes: passed, warning (optional), blocked (optional), message (optional)
+ """
+ if item_id not in self.attention_expected:
+ return None
+
+ expected = self.attention_expected[item_id]
+ passed = self._compare_responses(expected, response)
+
+ # Check for suspiciously fast response
+ if (response_time_seconds is not None and
+ self.qc_config.attention_min_response_time > 0 and
+ response_time_seconds < self.qc_config.attention_min_response_time):
+ self.logger.warning(f"User {user_id} responded to attention check {item_id} "
+ f"in {response_time_seconds:.1f}s (min: {self.qc_config.attention_min_response_time}s)")
+ # Still record the result but log the fast response
+
+ # Record result โ at most one result per (user, item) pair.
+ # Re-saves on the same item replace the prior result rather than
+ # appending, so a single attention-check item can only count once.
+ result = AttentionCheckResult(
+ item_id=item_id,
+ user_id=user_id,
+ passed=passed,
+ expected=expected,
+ actual=response,
+ response_time_seconds=response_time_seconds
+ )
+
+ with self._lock:
+ user_results = self.attention_results[user_id]
+ previous_failures = len([r for r in user_results if not r.passed])
+
+ # Replace any existing result for this (user, item) pair
+ self.attention_results[user_id] = [
+ r for r in user_results if r.item_id != item_id
+ ]
+ self.attention_results[user_id].append(result)
+
+ failures = len([r for r in self.attention_results[user_id] if not r.passed])
+
+ # Check thresholds โ only log when crossing for the first time
+ response_data = {"passed": passed}
+
+ if failures >= self.qc_config.attention_block_threshold:
+ response_data["blocked"] = True
+ response_data["message"] = self.qc_config.attention_block_message
+ if previous_failures < self.qc_config.attention_block_threshold:
+ self.logger.warning(f"User {user_id} blocked after {failures} attention check failures")
+ elif failures >= self.qc_config.attention_warn_threshold and not passed:
+ response_data["warning"] = True
+ response_data["message"] = self.qc_config.attention_warn_message
+ if previous_failures < self.qc_config.attention_warn_threshold:
+ self.logger.info(f"User {user_id} warned after {failures} attention check failures")
+
+ return response_data
+
+ def get_attention_check_stats(self, user_id: str) -> Dict[str, Any]:
+ """Get attention check statistics for a user."""
+ with self._lock:
+ results = self.attention_results.get(user_id, [])
+ passed = len([r for r in results if r.passed])
+ failed = len([r for r in results if not r.passed])
+ total = passed + failed
+
+ return {
+ "total": total,
+ "passed": passed,
+ "failed": failed,
+ "pass_rate": passed / total if total > 0 else 0.0
+ }
+
+ # =========================================================================
+ # Gold Standard Methods
+ # =========================================================================
+
+ def is_gold_standard(self, item_id: str) -> bool:
+ """Check if an item is a gold standard."""
+ return item_id in self.gold_labels
+
+ def should_inject_gold_standard(self, user_id: str, items_since_last: Optional[int] = None) -> bool:
+ """
+ Determine if a gold standard should be injected.
+
+ Args:
+ user_id: The user ID
+ items_since_last: Deprecated override; if omitted, uses the gold-specific counter
+
+ Returns:
+ True if a gold standard should be injected
+ """
+ if not self.qc_config.gold_standards_enabled or not self.gold_items:
+ return False
+
+ if self.qc_config.gold_mode == 'training':
+ return False
+
+ if items_since_last is None:
+ items_since_last = self.user_items_since_gold.get(user_id, 0)
+
+ if self.qc_config.gold_frequency:
+ return items_since_last >= self.qc_config.gold_frequency
+
+ return False
+
+ def get_gold_standard_item(self, user_id: str) -> Optional[Dict]:
+ """
+ Get a gold standard item for a user.
+
+ Args:
+ user_id: The user ID
+
+ Returns:
+ A gold standard item dict, or None if none available
+ """
+ if not self.gold_items:
+ return None
+
+ with self._lock:
+ # Get items this user hasn't seen yet
+ seen_ids = {r.item_id for r in self.gold_results.get(user_id, [])}
+ available = [item for item in self.gold_items if item['id'] not in seen_ids]
+
+ if not available:
+ # Recycle items if all have been seen
+ available = self.gold_items
+
+ self.user_items_since_gold[user_id] = 0
+ return random.choice(available)
+
+ def validate_gold_response(
+ self,
+ user_id: str,
+ item_id: str,
+ response: Dict[str, Any]
+ ) -> Optional[Dict[str, Any]]:
+ """
+ Validate a response against a gold standard.
+
+ By default, gold standards are "silent" - results are recorded for admin
+ review but no feedback is shown to users. This can be changed via config
+ for training scenarios.
+
+ Args:
+ user_id: The user ID
+ item_id: The gold standard item ID
+ response: The user's response
+
+ Returns:
+ Dict with validation result if this is a gold standard, None otherwise.
+ By default only contains 'recorded': True to indicate silent recording.
+ If feedback is enabled, also includes: correct, gold_label, explanation.
+ """
+ if item_id not in self.gold_labels:
+ return None
+
+ gold = self.gold_labels[item_id]
+ correct = self._compare_responses(gold, response)
+ explanation = self.gold_explanations.get(item_id)
+
+ # Record result (always happens, regardless of feedback settings)
+ result = GoldStandardResult(
+ item_id=item_id,
+ user_id=user_id,
+ correct=correct,
+ gold_label=gold,
+ user_response=response,
+ explanation=explanation
+ )
+
+ with self._lock:
+ self.gold_results[user_id].append(result)
+
+ # By default, gold standards are silent (recorded but no feedback to user)
+ # Only include feedback if explicitly enabled in config
+ show_feedback = (self.qc_config.gold_show_correct_answer or
+ self.qc_config.gold_show_explanation)
+
+ if not show_feedback:
+ # Silent recording - just indicate it was recorded, no user feedback
+ return {"recorded": True, "silent": True}
+
+ # Build response with feedback (only if enabled)
+ response_data = {"correct": correct, "recorded": True}
+
+ if self.qc_config.gold_show_correct_answer:
+ response_data["gold_label"] = gold
+
+ if self.qc_config.gold_show_explanation and explanation:
+ response_data["explanation"] = explanation
+
+ # Check accuracy threshold (only show warning if feedback is enabled)
+ accuracy_data = self.get_gold_accuracy(user_id)
+ if accuracy_data["total"] >= self.qc_config.gold_evaluation_count:
+ if accuracy_data["accuracy"] < self.qc_config.gold_min_accuracy:
+ response_data["accuracy_warning"] = True
+ response_data["current_accuracy"] = accuracy_data["accuracy"]
+ response_data["required_accuracy"] = self.qc_config.gold_min_accuracy
+
+ return response_data
+
+ def get_gold_accuracy(self, user_id: str) -> Dict[str, Any]:
+ """Get gold standard accuracy for a user."""
+ with self._lock:
+ results = self.gold_results.get(user_id, [])
+ correct = len([r for r in results if r.correct])
+ total = len(results)
+
+ return {
+ "total": total,
+ "correct": correct,
+ "accuracy": correct / total if total > 0 else 0.0,
+ "results": [
+ {
+ "item_id": r.item_id,
+ "correct": r.correct,
+ "timestamp": r.timestamp.isoformat()
+ }
+ for r in results
+ ]
+ }
+
+ # =========================================================================
+ # Gold Standard Auto-Promotion Methods
+ # =========================================================================
+
+ def record_item_annotation(
+ self,
+ item_id: str,
+ user_id: str,
+ response: Dict[str, Any]
+ ) -> Optional[Dict[str, Any]]:
+ """
+ Record an annotation for potential gold standard auto-promotion.
+
+ When enough annotators agree on a label, the item is automatically
+ promoted to the gold standard pool.
+
+ Args:
+ item_id: The item ID
+ user_id: The user ID
+ response: The user's annotation response
+
+ Returns:
+ Dict with promotion info if item was promoted, None otherwise
+ """
+ if not self.qc_config.gold_auto_promote_enabled:
+ return None
+
+ # Don't track items that are already gold standards
+ if self.is_gold_standard(item_id):
+ return None
+
+ with self._lock:
+ # Record this annotation
+ self.item_annotations[item_id][user_id] = response
+
+ # Check if we have enough annotators
+ annotations = self.item_annotations[item_id]
+ if len(annotations) < self.qc_config.gold_auto_promote_min_annotators:
+ return None
+
+ # Check agreement
+ promotion_result = self._check_and_promote(item_id, annotations)
+ return promotion_result
+
+ def _check_and_promote(
+ self,
+ item_id: str,
+ annotations: Dict[str, Dict[str, Any]]
+ ) -> Optional[Dict[str, Any]]:
+ """
+ Check if annotations meet agreement threshold and promote if so.
+
+ Args:
+ item_id: The item ID
+ annotations: Dict of user_id -> response
+
+ Returns:
+ Promotion info if promoted, None otherwise
+ """
+ # Extract responses and check agreement for each schema
+ consensus_label = {}
+ all_schemas_agree = True
+
+ # Group by schema
+ schema_responses: Dict[str, List[Any]] = defaultdict(list)
+ for user_id, response in annotations.items():
+ for schema, value in response.items():
+ schema_responses[schema].append(value)
+
+ # Check agreement for each schema
+ for schema, values in schema_responses.items():
+ # Calculate agreement ratio
+ from collections import Counter
+ value_counts = Counter(str(v).lower() for v in values)
+ most_common_value, most_common_count = value_counts.most_common(1)[0]
+ agreement_ratio = most_common_count / len(values)
+
+ if agreement_ratio >= self.qc_config.gold_auto_promote_agreement:
+ # Find the original value (not lowercased)
+ for v in values:
+ if str(v).lower() == most_common_value:
+ consensus_label[schema] = v
+ break
+ else:
+ all_schemas_agree = False
+ break
+
+ if not all_schemas_agree:
+ return None
+
+ # Promote to gold standard
+ self._promote_to_gold(item_id, consensus_label, annotations)
+
+ self.logger.info(f"Auto-promoted item {item_id} to gold standard with label: {consensus_label}")
+
+ return {
+ "promoted": True,
+ "item_id": item_id,
+ "consensus_label": consensus_label,
+ "annotator_count": len(annotations),
+ "agreement": 1.0 # At this point we have consensus
+ }
+
+ def _promote_to_gold(
+ self,
+ item_id: str,
+ consensus_label: Dict[str, Any],
+ source_annotations: Dict[str, Dict[str, Any]]
+ ) -> None:
+ """
+ Promote an item to the gold standard pool.
+
+ Args:
+ item_id: The item ID to promote
+ consensus_label: The agreed-upon label
+ source_annotations: The annotations that led to promotion
+ """
+ # Create gold standard entry
+ gold_item = {
+ "id": item_id,
+ "gold_label": consensus_label,
+ "auto_promoted": True,
+ "promoted_at": datetime.now().isoformat(),
+ "source_annotators": list(source_annotations.keys()),
+ "annotator_count": len(source_annotations)
+ }
+
+ # Add to promoted gold items
+ self.promoted_gold_items.append(gold_item)
+ self.promoted_gold_labels[item_id] = consensus_label
+
+ # Also add to main gold labels so is_gold_standard() works
+ self.gold_labels[item_id] = consensus_label
+
+ def get_promoted_gold_standards(self) -> List[Dict]:
+ """Get all items that were auto-promoted to gold standards."""
+ with self._lock:
+ return list(self.promoted_gold_items)
+
+ def get_promotion_candidates(self) -> List[Dict[str, Any]]:
+ """
+ Get items that are close to being promoted (for admin visibility).
+
+ Returns items with annotations but not yet meeting the threshold.
+ """
+ candidates = []
+
+ with self._lock:
+ min_annotators = self.qc_config.gold_auto_promote_min_annotators
+
+ for item_id, annotations in self.item_annotations.items():
+ if item_id in self.gold_labels:
+ continue # Already a gold standard
+
+ if len(annotations) == 0:
+ continue
+
+ # Calculate current agreement
+ schema_agreement = {}
+ schema_responses: Dict[str, List[Any]] = defaultdict(list)
+
+ for user_id, response in annotations.items():
+ for schema, value in response.items():
+ schema_responses[schema].append(value)
+
+ for schema, values in schema_responses.items():
+ from collections import Counter
+ value_counts = Counter(str(v).lower() for v in values)
+ if value_counts:
+ most_common_value, most_common_count = value_counts.most_common(1)[0]
+ schema_agreement[schema] = {
+ "value": most_common_value,
+ "count": most_common_count,
+ "total": len(values),
+ "agreement": most_common_count / len(values)
+ }
+
+ candidates.append({
+ "item_id": item_id,
+ "annotator_count": len(annotations),
+ "needed_annotators": min_annotators,
+ "schema_agreement": schema_agreement
+ })
+
+ return candidates
+
+ # =========================================================================
+ # Pre-annotation Methods
+ # =========================================================================
+
+ def extract_pre_annotations(self, item_id: str, item_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+ """
+ Extract pre-annotation data from an item.
+
+ Args:
+ item_id: The item ID
+ item_data: The item's data dictionary
+
+ Returns:
+ Pre-annotation data dict if available, None otherwise
+ """
+ if not self.qc_config.pre_annotation_enabled:
+ return None
+
+ field = self.qc_config.pre_annotation_field
+ if field not in item_data:
+ return None
+
+ pre_data = item_data[field]
+ if not isinstance(pre_data, dict):
+ self.logger.warning(f"Pre-annotation field '{field}' in item {item_id} is not a dict")
+ return None
+
+ # Cache for later use
+ self.pre_annotations[item_id] = pre_data
+ return pre_data
+
+ def get_pre_annotations(self, item_id: str) -> Optional[Dict[str, Any]]:
+ """Get cached pre-annotations for an item."""
+ return self.pre_annotations.get(item_id)
+
+ def get_pre_annotation_config(self) -> Dict[str, Any]:
+ """Get pre-annotation configuration for frontend."""
+ if not self.qc_config.pre_annotation_enabled:
+ return {"enabled": False}
+
+ return {
+ "enabled": True,
+ "allow_modification": self.qc_config.pre_annotation_allow_modification,
+ "show_confidence": self.qc_config.pre_annotation_show_confidence,
+ "highlight_threshold": self.qc_config.pre_annotation_highlight_threshold
+ }
+
+ # =========================================================================
+ # Utility Methods
+ # =========================================================================
+
+ def _compare_responses(self, expected: Dict[str, Any], actual: Dict[str, Any]) -> bool:
+ """
+ Compare expected and actual responses.
+
+ Handles various response formats:
+ - Simple key-value pairs
+ - Lists (for multiselect)
+ - Nested structures
+
+ Args:
+ expected: The expected response
+ actual: The actual response
+
+ Returns:
+ True if responses match
+ """
+ for key, expected_value in expected.items():
+ # Handle both "schema_name" and "schema_name:label_name" formats
+ actual_value = None
+
+ # Direct match
+ if key in actual:
+ actual_value = actual[key]
+ else:
+ # Check for prefixed keys (schema_name:label_name format)
+ for actual_key, val in actual.items():
+ if actual_key.startswith(key + ":") or actual_key == key:
+ actual_value = val
+ break
+
+ if actual_value is None:
+ return False
+
+ # Compare values
+ if isinstance(expected_value, list):
+ if not isinstance(actual_value, list):
+ actual_value = [actual_value]
+ if set(expected_value) != set(actual_value):
+ return False
+ elif isinstance(expected_value, dict):
+ if not isinstance(actual_value, dict):
+ return False
+ if not self._compare_responses(expected_value, actual_value):
+ return False
+ else:
+ # Simple value comparison
+ if str(expected_value).lower() != str(actual_value).lower():
+ return False
+
+ return True
+
+ def get_all_attention_results(self) -> Dict[str, List[Dict]]:
+ """Get all attention check results for all users."""
+ with self._lock:
+ return {
+ user_id: [
+ {
+ "item_id": r.item_id,
+ "passed": r.passed,
+ "timestamp": r.timestamp.isoformat(),
+ "response_time": r.response_time_seconds
+ }
+ for r in results
+ ]
+ for user_id, results in self.attention_results.items()
+ }
+
+ def get_all_gold_results(self) -> Dict[str, List[Dict]]:
+ """Get all gold standard results for all users."""
+ with self._lock:
+ return {
+ user_id: [
+ {
+ "item_id": r.item_id,
+ "correct": r.correct,
+ "timestamp": r.timestamp.isoformat()
+ }
+ for r in results
+ ]
+ for user_id, results in self.gold_results.items()
+ }
+
+ def get_quality_metrics(self) -> Dict[str, Any]:
+ """Get comprehensive quality control metrics for admin dashboard."""
+ with self._lock:
+ # Attention check metrics
+ attention_metrics = {
+ "enabled": self.qc_config.attention_checks_enabled,
+ "total_items": len(self.attention_items),
+ "total_checks": sum(len(r) for r in self.attention_results.values()),
+ "total_passed": sum(
+ len([x for x in r if x.passed])
+ for r in self.attention_results.values()
+ ),
+ "total_failed": sum(
+ len([x for x in r if not x.passed])
+ for r in self.attention_results.values()
+ ),
+ "by_user": {}
+ }
+
+ for user_id, results in self.attention_results.items():
+ passed = len([r for r in results if r.passed])
+ failed = len([r for r in results if not r.passed])
+ attention_metrics["by_user"][user_id] = {
+ "passed": passed,
+ "failed": failed,
+ "pass_rate": passed / (passed + failed) if (passed + failed) > 0 else 0
+ }
+
+ # Gold standard metrics
+ gold_metrics = {
+ "enabled": self.qc_config.gold_standards_enabled,
+ "total_items": len(self.gold_items),
+ "total_evaluations": sum(len(r) for r in self.gold_results.values()),
+ "total_correct": sum(
+ len([x for x in r if x.correct])
+ for r in self.gold_results.values()
+ ),
+ "total_incorrect": sum(
+ len([x for x in r if not x.correct])
+ for r in self.gold_results.values()
+ ),
+ "by_user": {},
+ "by_item": {}
+ }
+
+ for user_id, results in self.gold_results.items():
+ correct = len([r for r in results if r.correct])
+ total = len(results)
+ gold_metrics["by_user"][user_id] = {
+ "correct": correct,
+ "total": total,
+ "accuracy": correct / total if total > 0 else 0
+ }
+
+ # Per-item accuracy
+ item_results = defaultdict(lambda: {"correct": 0, "total": 0})
+ for results in self.gold_results.values():
+ for r in results:
+ item_results[r.item_id]["total"] += 1
+ if r.correct:
+ item_results[r.item_id]["correct"] += 1
+
+ for item_id, counts in item_results.items():
+ gold_metrics["by_item"][item_id] = {
+ "correct": counts["correct"],
+ "total": counts["total"],
+ "accuracy": counts["correct"] / counts["total"] if counts["total"] > 0 else 0
+ }
+
+ # Auto-promotion metrics
+ auto_promotion_metrics = {
+ "enabled": self.qc_config.gold_auto_promote_enabled,
+ "min_annotators": self.qc_config.gold_auto_promote_min_annotators,
+ "agreement_threshold": self.qc_config.gold_auto_promote_agreement,
+ "promoted_count": len(self.promoted_gold_items),
+ "promoted_items": [
+ {
+ "item_id": item["id"],
+ "consensus_label": item["gold_label"],
+ "annotator_count": item["annotator_count"],
+ "promoted_at": item["promoted_at"]
+ }
+ for item in self.promoted_gold_items
+ ],
+ "candidates": self.get_promotion_candidates()[:20] # Top 20 candidates
+ }
+
+ return {
+ "attention_checks": attention_metrics,
+ "gold_standards": gold_metrics,
+ "auto_promotion": auto_promotion_metrics,
+ "pre_annotation": {
+ "enabled": self.qc_config.pre_annotation_enabled,
+ "items_with_predictions": len(self.pre_annotations)
+ }
+ }
+
+
+def init_quality_control_manager(config: Dict[str, Any], base_dir: str) -> QualityControlManager:
+ """Initialize the singleton QualityControlManager."""
+ global _QUALITY_CONTROL_MANAGER
+
+ with _QUALITY_CONTROL_LOCK:
+ if _QUALITY_CONTROL_MANAGER is None:
+ _QUALITY_CONTROL_MANAGER = QualityControlManager(config, base_dir)
+
+ return _QUALITY_CONTROL_MANAGER
+
+
+def get_quality_control_manager() -> Optional[QualityControlManager]:
+ """Get the singleton QualityControlManager instance."""
+ return _QUALITY_CONTROL_MANAGER
+
+
+def clear_quality_control_manager():
+ """Clear the singleton (for testing)."""
+ global _QUALITY_CONTROL_MANAGER
+ with _QUALITY_CONTROL_LOCK:
+ _QUALITY_CONTROL_MANAGER = None
diff --git a/potato/remove_users_from_queue.py b/potato/remove_users_from_queue.py
new file mode 100644
index 0000000000000000000000000000000000000000..4cd1ae557370d3b2b12800e74754777ba1f9225f
--- /dev/null
+++ b/potato/remove_users_from_queue.py
@@ -0,0 +1,96 @@
+"""
+User Removal and Data Cleanup Module
+
+This script is used to remove users from the assigned data. This operation is usually used
+when there are bad users who participate in the task but didn't finish all the instances.
+
+The script performs several cleanup operations:
+1. Removes user annotations from the global annotation file
+2. Moves bad users' data to an archived directory
+3. Updates task assignment data to reflect user removal
+4. Recalculates unassigned instance counts
+
+This is a destructive operation that permanently removes user data from the active
+annotation system and should be used with caution.
+"""
+
+import json
+import os
+#from server_utils.config_module import init_config, config
+from argparse import ArgumentParser
+import pandas as pd
+import shutil
+
+# Configuration paths (commented out as they're set via command line arguments)
+task_assignment_path = None#os.path.join(config["output_annotation_dir"], config["automatic_assignment"]["output_filename"])
+annotation_data_dir = None#config["output_annotation_dir"]
+annotation_data_path = None#os.path.join(config["output_annotation_dir"], "annotated_instances.tsv")
+
+# Set up command line argument parsing
+parser = ArgumentParser()
+parser.set_defaults(show_path=False, show_similarity=False)
+parser.add_argument("--task_assignment_path", default=task_assignment_path)
+parser.add_argument("--annotation_data_dir", default=annotation_data_dir)
+parser.add_argument("--user_file")
+
+args = parser.parse_args()
+args.annotation_data_path = os.path.join(args.annotation_data_dir, "annotated_instances.tsv")
+
+print(args)
+
+# Load list of users to be removed from the specified file
+with open(args.user_file,'r') as r:
+ users = [it.strip() for it in r.readlines()]
+ user_set = set(users)
+
+print("users to be removed from the assigned data and annotation instances: ", users)
+
+# Remove user annotations from the global annotation file
+# This filters out all annotations made by the specified users
+annotated_df = pd.read_csv(args.annotation_data_path, sep="\t")
+new_annotated_df = annotated_df[~annotated_df['user'].isin(users)]
+print("%d lines removed from bad users"%(len(annotated_df)-len(new_annotated_df)))
+new_annotated_df.to_csv(args.annotation_data_path + '_new', sep="\t", index=False)
+
+# Move the bad users into a separate directory under annotation output
+# This preserves their data but removes it from the active annotation system
+bad_user_dir = args.annotation_data_dir + "archived_users"
+if not os.path.exists(bad_user_dir):
+ os.mkdir(bad_user_dir)
+for u in users:
+ shutil.move(os.path.join(args.annotation_data_dir, u), os.path.join(bad_user_dir, u))
+print('bad users moved to %s'%bad_user_dir)
+
+# Remove users from the task assignment data
+# This updates the assignment tracking to reflect that these users are no longer
+# assigned to any instances, and increases the unassigned count accordingly
+if os.path.exists(args.task_assignment_path):
+ # Load the task assignment if it has been generated and saved
+ with open(args.task_assignment_path, "r") as r:
+ task_assignment = json.load(r)
+
+# Process each instance to remove bad users from assignments
+for inst_id in task_assignment['assigned']:
+ new_li = []
+ if type(task_assignment['assigned'][inst_id]) != list:
+ continue
+ for u in task_assignment['assigned'][inst_id]:
+ if u in user_set:
+ # If user is being removed, increment the unassigned count for this instance
+ if inst_id not in task_assignment['unassigned']:
+ task_assignment['unassigned'][inst_id] = 0
+ task_assignment['unassigned'][inst_id] += 1
+ else:
+ # Keep users that are not being removed
+ new_li.append(u)
+ #if len(new_li) != len(task_assignment['assigned'][inst_id]):
+ # print(task_assignment['assigned'][inst_id], new_li)
+ task_assignment['assigned'][inst_id] = new_li
+
+# Save the updated task assignment data
+if os.path.exists(args.task_assignment_path):
+ # Load the task assignment if it has been generated and saved
+ with open(args.task_assignment_path + '_new', "w") as w:
+ json.dump(task_assignment, w)
+
+print('unassigned instances after user removal',task_assignment['unassigned'])
\ No newline at end of file
diff --git a/potato/routes.py b/potato/routes.py
new file mode 100644
index 0000000000000000000000000000000000000000..50a5ec0cc69fade1c37a14811d07602b49830911
--- /dev/null
+++ b/potato/routes.py
@@ -0,0 +1,8272 @@
+"""
+Flask Routes Module
+
+This module contains all the route handlers for the Flask server.
+It defines the HTTP endpoints and their associated logic for:
+- User authentication and session management
+- Navigation between annotation phases
+- Form handling and validation
+- Annotation submission and processing
+- User registration and management
+- Admin dashboard functionality
+- API endpoints for frontend integration
+
+The routes handle the complete annotation workflow from initial login
+through completion, including consent, instructions, training, annotation,
+and post-study phases. They also provide admin functionality for monitoring
+progress and managing the annotation system.
+
+Key Features:
+- Session-based authentication with timeout management
+- Multi-phase workflow support with configurable phases
+- Annotation submission with validation and persistence
+- AI hint integration for improved annotation quality
+- Admin dashboard with comprehensive statistics
+- API endpoints for real-time frontend updates
+- Error handling and user feedback
+"""
+from __future__ import annotations
+
+import json
+import logging
+import traceback
+import datetime
+from datetime import timedelta
+from flask import Flask, session, render_template, request, redirect, url_for, jsonify, make_response
+import time
+import uuid
+
+# Import from the main flask_server.py module
+from potato.flask_server import (
+ app, config, logger,
+ get_user_state_manager, get_user_state, get_item_state_manager,
+ init_user_state, UserAuthenticator, UserPhase,
+ move_to_prev_instance, move_to_next_instance, go_to_id,
+ get_annotations_for_user_on, get_span_annotations_for_user_on,
+ _instance_meets_required_annotation_rules,
+ render_page_with_annotations, get_current_page_html,
+ validate_annotation, parse_html_span_annotation, Label, SpanAnnotation,
+ get_users, get_total_annotations, update_annotation_state,
+ get_ai_cache_manager,
+ get_users, get_total_annotations, update_annotation_state, ai_hints,
+ get_training_instances, get_training_correct_answers, get_training_explanation,
+ get_training_instance_categories, get_prolific_study, get_keyword_highlight_patterns,
+ get_keyword_highlight_settings
+)
+
+# Import admin dashboard functionality
+from potato.admin import admin_dashboard
+
+# Import span color functions
+from potato.ai.ai_help_wrapper import generate_ai_help_html
+from potato.ai.ai_prompt import get_ai_prompt
+from potato.server_utils.schemas.span import get_span_color, set_span_color, SPAN_COLOR_PALETTE
+
+# Import annotation history
+from potato.annotation_history import AnnotationHistoryManager
+from potato.logging_config import is_ui_debug_enabled, get_debug_log_settings
+
+# Import quality control
+from potato.quality_control import get_quality_control_manager
+
+# Import adjudication
+from potato.adjudication import get_adjudication_manager, AdjudicationDecision
+
+# Import diversity manager
+from potato.diversity_manager import get_diversity_manager
+
+import os
+from potato.item_state_management import Item
+from potato.flask_server import get_displayed_text
+
+
+def _inject_quality_control_item_if_needed(username, user_state):
+ qc_manager = get_quality_control_manager()
+ if not qc_manager:
+ return
+
+ current_instance = user_state.get_current_instance()
+ current_instance_id = current_instance.get_id() if current_instance else None
+ if current_instance_id and (
+ qc_manager.is_attention_check(current_instance_id) or qc_manager.is_gold_standard(current_instance_id)
+ ):
+ return
+
+ assigned_ids = set(user_state.get_assigned_instance_ids())
+ annotated_ids = set(user_state.get_annotated_instance_ids()) if hasattr(user_state, "get_annotated_instance_ids") else set()
+ seen_qc_ids = assigned_ids | annotated_ids
+
+ current_index = user_state.get_current_instance_index()
+ insert_index = current_index + 1 if current_index >= 0 else 0
+
+ def inject_item(item_data):
+ item_id = item_data.get("id")
+ if not item_id or item_id in seen_qc_ids:
+ return False
+
+ prepared_item = dict(item_data)
+ text_key = config.get("item_properties", {}).get("text_key", "text")
+ if "displayed_text" not in prepared_item:
+ raw_text = prepared_item.get(text_key, prepared_item.get("text", ""))
+ prepared_item["displayed_text"] = get_displayed_text(raw_text) if raw_text is not None else ""
+
+ item_manager = get_item_state_manager()
+ if item_manager.has_item(item_id):
+ existing_item = item_manager.get_item(item_id)
+ if existing_item and isinstance(existing_item.get_data(), dict):
+ existing_data = existing_item.get_data()
+ if "displayed_text" not in existing_data:
+ existing_data["displayed_text"] = prepared_item["displayed_text"]
+ item = item_manager.get_item(item_id)
+ else:
+ item_manager.add_item(item_id, prepared_item)
+ item = item_manager.get_item(item_id)
+
+ return user_state.assign_instance_at_index(item, insert_index)
+
+ if qc_manager.should_inject_attention_check(username):
+ attention_item = qc_manager.get_attention_check_item(username)
+ if attention_item and inject_item(attention_item):
+ logger.info(f"Injected attention check {attention_item.get('id')} for user {username}")
+ return
+
+ if qc_manager.should_inject_gold_standard(username):
+ gold_item = qc_manager.get_gold_standard_item(username)
+ if gold_item and inject_item(gold_item):
+ logger.info(f"Injected gold standard {gold_item.get('id')} for user {username}")
+ return
+
+
+def _reclaim_blocked_user_assignments(username, user_state, current_instance_id=None):
+ """Release unannotated assignments after a user is blocked."""
+ item_manager = get_item_state_manager()
+ preserve_completed = item_manager.should_preserve_completed_annotations("quality_control_block")
+
+ if current_instance_id:
+ user_state.clear_instance_annotations(current_instance_id)
+
+ reclaimed = item_manager.reclaim_unannotated_assignments_for_user(
+ user_state,
+ reason="quality_control_block",
+ preserve_completed_annotations=preserve_completed,
+ )
+
+ if reclaimed:
+ logger.info(
+ "Reclaimed %d unannotated assignments from blocked user %s",
+ len(reclaimed),
+ username,
+ )
+
+ try:
+ get_user_state_manager().save_user_state(user_state)
+ except Exception as e:
+ logger.warning("Could not persist blocked-user assignment reclaim for %s: %s", username, e)
+
+ return reclaimed
+
+
+def get_debug_phase_target(debug_phase: str) -> tuple:
+ """
+ Parse the debug_phase string and find the matching phase and page.
+
+ The debug_phase can be:
+ - A phase type name like "annotation", "prestudy", "poststudy"
+ - A specific page name within a phase (e.g., "consent_page_1")
+
+ Args:
+ debug_phase: The debug phase string from config
+
+ Returns:
+ tuple: (UserPhase, page_name) or (None, None) if not found
+ """
+ if not debug_phase:
+ return None, None
+
+ usm = get_user_state_manager()
+
+ # First, try to match as a phase type (case-insensitive)
+ try:
+ phase = UserPhase.fromstr(debug_phase)
+ # Check if this phase exists in the config
+ if phase in usm.phase_type_to_name_to_page:
+ pages = list(usm.phase_type_to_name_to_page[phase].keys())
+ if pages:
+ return phase, pages[0]
+ # Special handling for ANNOTATION phase which is always available
+ if phase == UserPhase.ANNOTATION:
+ return phase, "annotation"
+ return None, None
+ except (ValueError, KeyError):
+ pass
+
+ # If not a phase type, search for it as a page name
+ for phase_type, pages_dict in usm.phase_type_to_name_to_page.items():
+ for page_name in pages_dict.keys():
+ if page_name.lower() == debug_phase.lower():
+ return phase_type, page_name
+
+ logger.warning(f"Debug phase '{debug_phase}' not found in configured phases")
+ return None, None
+
+
+def apply_debug_phase_skip(user_id: str) -> bool:
+ """
+ Apply debug phase skip if configured.
+
+ Args:
+ user_id: The user ID to apply the skip for
+
+ Returns:
+ bool: True if skip was applied, False otherwise
+ """
+ debug_phase = config.get("debug_phase")
+ if not debug_phase:
+ return False
+
+ phase, page = get_debug_phase_target(debug_phase)
+ if phase is None:
+ logger.warning(f"Could not apply debug phase skip: '{debug_phase}' not found")
+ return False
+
+ usm = get_user_state_manager()
+ user_state = usm.get_user_state(user_id)
+
+ if user_state:
+ user_state.advance_to_phase(phase, page)
+ logger.info(f"Debug: Skipped user '{user_id}' to phase '{phase.value}', page '{page}'")
+ return True
+
+ return False
+
+# Admin API key resolution โ delegated to shared utility
+from potato.server_utils.admin_key import (
+ get_admin_api_key as _get_admin_api_key,
+ validate_admin_api_key as _validate_admin_api_key,
+)
+
+def get_admin_api_key():
+ """Get the admin API key. See potato.server_utils.admin_key for details."""
+ return _get_admin_api_key(config)
+
+def validate_admin_api_key(provided_key: str) -> bool:
+ """Validate an admin API key. See potato.server_utils.admin_key for details."""
+ return _validate_admin_api_key(provided_key, config)
+
+
+# -------------------------------------------------------------------
+# Local media file serving
+# -------------------------------------------------------------------
+
+def serve_media(filepath):
+ """Serve a local media file from the project's media directory.
+
+ Serves files from a ``media/`` directory (or custom path set via
+ ``media_directory`` in config) relative to the task directory. This
+ lets data files reference local images, audio, and video with paths
+ like ``/media/image_01.jpg`` instead of requiring external URLs.
+ """
+ from flask import send_from_directory, abort
+
+ task_dir = config.get("task_dir", ".")
+ media_subdir = config.get("media_directory", "media")
+
+ # Resolve relative to task_dir
+ if os.path.isabs(media_subdir):
+ media_dir = media_subdir
+ else:
+ media_dir = os.path.join(task_dir, media_subdir)
+
+ media_dir = os.path.realpath(media_dir)
+
+ # Security: ensure the resolved file is inside media_dir
+ requested = os.path.realpath(os.path.join(media_dir, filepath))
+ if not requested.startswith(media_dir + os.sep) and requested != media_dir:
+ logger.warning(f"Media path traversal blocked: {filepath}")
+ abort(403)
+
+ if not os.path.isfile(requested):
+ abort(404)
+
+ return send_from_directory(media_dir, filepath)
+
+
+def serve_trace_screenshot(filepath):
+ """Serve a static agent-trace screenshot referenced by trace data.
+
+ ``web_agent_trace`` / visual trace data set ``screenshot_url`` to a path
+ relative to the example (e.g. ``screenshots/step_000.png``), which the
+ browser requests as ``/screenshots/...``. Without this route those
+ requests 404 and every step image + filmstrip thumbnail breaks. Files
+ resolve against ``task_dir`` (the runtime CWD) with path-traversal
+ containment, mirroring ``serve_media``.
+ """
+ from flask import send_from_directory, abort
+
+ task_dir = config.get("task_dir", ".")
+ base_dir = os.path.realpath(task_dir)
+ screenshots_root = os.path.realpath(os.path.join(base_dir, "screenshots"))
+
+ requested = os.path.realpath(os.path.join(screenshots_root, filepath))
+ if not (requested == screenshots_root
+ or requested.startswith(screenshots_root + os.sep)):
+ logger.warning(f"Screenshot path traversal blocked: {filepath}")
+ abort(403)
+
+ if not os.path.isfile(requested):
+ abort(404)
+
+ return send_from_directory(screenshots_root, filepath)
+
+
+@app.route("/", methods=["GET", "POST"])
+def home():
+ """
+ Handle requests to the home page.
+
+ This route serves as the main entry point for the annotation platform.
+ It handles session management, user authentication, and phase routing
+ based on the user's current state in the annotation workflow.
+
+ Features:
+ - Session validation and timeout management
+ - User authentication and state initialization
+ - Phase-based routing to appropriate pages
+ - Survey flow management
+ - Progress tracking and validation
+ - URL-direct login for crowdsourcing platforms (Prolific, MTurk, etc.)
+
+ Returns:
+ flask.Response: Rendered template or redirect based on user state
+
+ Side Effects:
+ - May initialize new user state
+ - May advance user phases
+ - May clear invalid sessions
+ """
+ logger.debug("Processing home page request")
+
+ # In debug mode with debug_phase, auto-login and skip to the specified phase
+ if config.get("debug") and config.get("debug_phase") and 'username' not in session:
+ debug_user = "debug_user"
+ logger.info(f"Debug mode: Auto-logging in as '{debug_user}' and skipping to phase '{config.get('debug_phase')}'")
+
+ # Auto-register the debug user if needed
+ user_authenticator = UserAuthenticator.get_instance()
+ if not user_authenticator.is_valid_username(debug_user):
+ user_authenticator.add_user(debug_user, None)
+
+ # Set session
+ session['username'] = debug_user
+ session.permanent = True
+
+ # Initialize user state and apply debug phase skip
+ usm = get_user_state_manager()
+ if not usm.has_user(debug_user):
+ usm.add_user(debug_user)
+
+ # Apply the debug phase skip
+ apply_debug_phase_skip(debug_user)
+
+ # Redirect to home to process the new session
+ return redirect(url_for("home"))
+
+ # Check if user has an active session
+ if 'username' not in session:
+ # Check for URL-direct login (used by Prolific, MTurk, etc.)
+ login_config = config.get('login', {})
+ login_type = login_config.get('type', 'standard')
+
+ if login_type in ['url_direct', 'prolific']:
+ # Get the URL argument name (default to PROLIFIC_PID for backwards compatibility)
+ url_argument = login_config.get('url_argument', 'PROLIFIC_PID')
+ username = request.args.get(url_argument)
+
+ # Also capture SESSION_ID and STUDY_ID if provided (for Prolific tracking)
+ prolific_session_id = request.args.get('SESSION_ID')
+ prolific_study_id = request.args.get('STUDY_ID')
+
+ # Capture MTurk-specific parameters
+ mturk_assignment_id = request.args.get('assignmentId')
+ mturk_hit_id = request.args.get('hitId')
+ mturk_submit_to = request.args.get('turkSubmitTo')
+
+ # Handle MTurk preview mode (worker hasn't accepted the HIT yet)
+ if mturk_assignment_id == 'ASSIGNMENT_ID_NOT_AVAILABLE':
+ logger.info("MTurk preview mode detected - showing preview page")
+ return render_template("mturk_preview.html",
+ title=config.get("annotation_task_name", "Task Preview"),
+ task_description=config.get("task_description", ""),
+ annotation_task_name=config.get("annotation_task_name", "Annotation Task"))
+
+ if username:
+ logger.info(f"URL-direct login: user={username}, session_id={prolific_session_id}, study_id={prolific_study_id}")
+
+ # Auto-register and login the user
+ user_authenticator = UserAuthenticator.get_instance()
+
+ # Add user if not exists (passwordless for URL-direct)
+ if not user_authenticator.is_valid_username(username):
+ result = user_authenticator.add_user(username, None,
+ prolific_session_id=prolific_session_id,
+ prolific_study_id=prolific_study_id)
+ logger.debug(f"Auto-registered URL-direct user {username}: {result}")
+
+ # Set session
+ session['username'] = username
+ session.permanent = True
+
+ # Store Prolific IDs in session for later use
+ if prolific_session_id:
+ session['prolific_session_id'] = prolific_session_id
+ if prolific_study_id:
+ session['prolific_study_id'] = prolific_study_id
+
+ # Store MTurk IDs in session for completion flow
+ if mturk_assignment_id:
+ session['mturk_assignment_id'] = mturk_assignment_id
+ if mturk_hit_id:
+ session['mturk_hit_id'] = mturk_hit_id
+ if mturk_submit_to:
+ session['mturk_submit_to'] = mturk_submit_to
+
+ # Initialize user state if needed
+ if not get_user_state_manager().has_user(username):
+ logger.debug(f"Initializing user state for URL-direct user: {username}")
+ init_user_state(username)
+
+ # Get the user state and set to first phase
+ usm = get_user_state_manager()
+ user_state = usm.get_user_state(username)
+
+ if user_state:
+ # Set user to the first phase if they're in LOGIN
+ # Use advance_phase() which properly looks up the first page
+ # for the phase (fixes issue #113: page was None for phased workflows)
+ if user_state.get_phase() == UserPhase.LOGIN:
+ logger.debug(f"Advancing URL-direct user {username} past LOGIN phase")
+ usm.advance_phase(username)
+
+ # Assign instances if user doesn't have any
+ if not user_state.has_assignments():
+ logger.debug(f"Assigning instances to URL-direct user {username}")
+ get_item_state_manager().assign_instances_to_user(user_state)
+
+ # Track with Prolific API if configured
+ prolific_study = get_prolific_study()
+ if prolific_study and prolific_session_id:
+ try:
+ prolific_study.add_new_user({
+ 'PROLIFIC_PID': username,
+ 'SESSION_ID': prolific_session_id
+ })
+ logger.debug(f"Tracked user {username} with Prolific API")
+ except Exception as e:
+ logger.warning(f"Failed to track user with Prolific API: {e}")
+
+ # Redirect to home to process the now-logged-in user
+ return redirect(url_for("home"))
+
+ else:
+ # URL-direct login configured but no username in URL
+ # Show error or redirect to a waiting page
+ logger.warning(f"URL-direct login configured but '{url_argument}' not found in URL")
+ return render_template("error.html",
+ message=f"Missing required URL parameter: {url_argument}. "
+ f"Please access this page through your crowdsourcing platform.")
+
+ logger.debug("No active session, rendering login page")
+ return render_template("home.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=config.get("require_password", True))
+
+ user_id = session['username']
+ logger.debug(f"Active session for user: {user_id}")
+
+ # Get user state and validate it exists
+ user_state = get_user_state(user_id)
+ if user_state is None:
+ logger.warning(f"User {user_id} not found in user state")
+ session.clear()
+ return redirect(url_for("auth"))
+
+ # Get the current phase of the user and route accordingly
+ phase = user_state.get_phase()
+ logger.debug(f"User phase: {phase}")
+
+ # Route to appropriate phase handler based on current phase
+ if phase == UserPhase.LOGIN:
+ return auth() #redirect(url_for("auth"))
+ elif phase == UserPhase.CONSENT:
+ return consent() #redirect(url_for("consent"))
+ elif phase == UserPhase.PRESTUDY:
+ return prestudy() #redirect(url_for("prestudy"))
+ elif phase == UserPhase.INSTRUCTIONS:
+ return instructions() #redirect(url_for("instructions"))
+ elif phase == UserPhase.TRAINING:
+ return training() #redirect(url_for("training"))
+ elif phase == UserPhase.ANNOTATION:
+ return annotate() # redirect(url_for("annotate"))
+ elif phase == UserPhase.POSTSTUDY:
+ return poststudy() #redirect(url_for("poststudy"))
+ elif phase == UserPhase.DONE:
+ return done() #redirect(url_for("done"))
+
+ logger.error(f"Invalid phase for user {user_id}: {phase}")
+ return render_template("error.html", message="Invalid application state")
+
+
+@app.route("/auth", methods=["GET", "POST"])
+def auth():
+ """
+ Handle authentication requests.
+
+ This route manages user authentication for the annotation platform.
+ It supports both password-based and passwordless authentication modes
+ depending on the system configuration.
+
+ Features:
+ - Session validation and management
+ - User authentication against configured backends
+ - User state initialization for new users
+ - Error handling and user feedback
+ - Redirect logic based on authentication success
+
+ Returns:
+ flask.Response: Rendered template or redirect
+
+ Side Effects:
+ - May create new user sessions
+ - May initialize new user states
+ - May clear existing sessions
+ """
+ # Check if user is already logged in
+ if 'username' in session and get_user_state_manager().has_user(session['username']):
+ logger.debug(f"User {session['username']} already logged in, redirecting to annotate")
+ return redirect(url_for("annotate"))
+
+ # Handle POST requests for user authentication
+ if request.method == "POST":
+ user_id = request.form.get("email")
+ password = request.form.get("pass")
+
+ logger.debug(f"Login attempt for user: {user_id}")
+
+ # Get require_password setting
+ require_password = config.get("require_password", True)
+
+ # Validate that user ID is provided
+ if not user_id:
+ logger.warning("Login attempt with empty user_id")
+ return render_template("home.html",
+ login_error="User ID is required",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=require_password)
+
+ # In passwordless mode, auto-register new users
+ if not require_password:
+ user_authenticator = UserAuthenticator.get_instance()
+ if not user_authenticator.is_valid_username(user_id):
+ logger.info(f"Auto-registering new user in passwordless mode: {user_id}")
+ user_authenticator.add_user(user_id, None)
+
+ # Authenticate the user against the configured backend
+ if UserAuthenticator.authenticate(user_id, password):
+ session.clear() # Clear any existing session data
+ session['username'] = user_id
+ session.permanent = True # Make session persist longer
+ logger.info(f"Login successful for user: {user_id}")
+
+ # Initialize user state if needed
+ if not get_user_state_manager().has_user(user_id):
+ logger.debug(f"Initializing state for new user: {user_id}")
+ usm = get_user_state_manager()
+ usm.add_user(user_id)
+
+ # Check for debug phase skip
+ if config.get("debug") and config.get("debug_phase"):
+ if apply_debug_phase_skip(user_id):
+ return redirect(url_for("home"))
+ else:
+ # Fall back to normal phase advancement
+ usm.advance_phase(user_id)
+ else:
+ usm.advance_phase(user_id)
+ return redirect(url_for("annotate"))
+ return redirect(url_for("annotate"))
+ else:
+ logger.warning(f"Login failed for user: {user_id}")
+ error_msg = "Invalid user ID or password" if require_password else "Login failed"
+ return render_template("home.html",
+ login_error=error_msg,
+ login_email=user_id,
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=require_password)
+
+ # GET request - show the login form
+ oauth_providers = []
+ allow_local_login = True
+ try:
+ authenticator = UserAuthenticator.get_instance()
+ oauth_providers = authenticator.get_login_providers()
+ if oauth_providers:
+ allow_local_login = authenticator.auth_config.get("allow_local_login", False)
+ except ValueError:
+ pass # Authenticator not yet initialized
+
+ return render_template("home.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=config.get("require_password", True),
+ oauth_providers=oauth_providers,
+ allow_local_login=allow_local_login)
+
+
+@app.route("/passwordless-login", methods=["GET", "POST"])
+def passwordless_login():
+ """
+ Legacy route for passwordless login.
+
+ This route now redirects to the main home page, which handles
+ both password and passwordless authentication based on the
+ require_password config setting.
+
+ Kept for backwards compatibility with existing links/bookmarks.
+ """
+ logger.debug("Redirecting from legacy passwordless-login to home")
+ return redirect(url_for("home"))
+
+
+@app.route("/clerk-login", methods=["GET", "POST"])
+def clerk_login():
+ """
+ Handle Clerk SSO login process.
+
+ This route manages authentication through Clerk's single sign-on service.
+ It handles token validation and user session creation for SSO users.
+
+ Features:
+ - Clerk SSO integration
+ - Token validation and verification
+ - User session management
+ - Error handling for SSO failures
+
+ Returns:
+ flask.Response: Rendered template or redirect
+
+ Side Effects:
+ - May create new user sessions
+ - May initialize new user states
+ """
+ logger.debug("Processing Clerk SSO login request")
+
+ # Only proceed if Clerk is configured
+ auth_method = config.get("authentication", {}).get("method", "in_memory")
+ if auth_method != "clerk":
+ logger.warning("Clerk login attempted but not configured")
+ return redirect(url_for("home"))
+
+ # Get the Clerk frontend API key
+ authenticator = UserAuthenticator.get_instance()
+ clerk_frontend_api = authenticator.get_clerk_frontend_api()
+
+ if not clerk_frontend_api:
+ logger.error("Clerk frontend API key not configured")
+ return render_template("home.html",
+ login_error="SSO configuration error",
+ title=config.get("annotation_task_name", "Annotation Platform"))
+
+ # Handle the Clerk token verification
+ if request.method == "POST":
+ token = request.form.get("clerk_token")
+ username = request.form.get("username")
+
+ if not token or not username:
+ logger.warning("Clerk login attempt with missing token or username")
+ return render_template("clerk_login.html",
+ login_error="Missing authentication data",
+ title=config.get("annotation_task_name", "Annotation Platform"))
+
+ # Authenticate with Clerk
+ if UserAuthenticator.authenticate(username, token):
+ session['username'] = username
+ logger.info(f"Clerk SSO login successful for user: {username}")
+
+ # Initialize user state if needed
+ if not get_user_state_manager().has_user(username):
+ logger.debug(f"Initializing state for new user: {username}")
+ init_user_state(username)
+
+ return redirect(url_for("annotate"))
+ else:
+ logger.warning(f"Clerk SSO login failed for user: {username}")
+ return render_template("clerk_login.html",
+ login_error="Authentication failed",
+ title=config.get("annotation_task_name", "Annotation Platform"))
+
+ # GET request - show the Clerk login form
+ return render_template("clerk_login.html",
+ clerk_frontend_api=clerk_frontend_api,
+ title=config.get("annotation_task_name", "Annotation Platform"))
+
+
+# --- OAuth SSO Routes ---
+
+@app.route("/auth/login/")
+def oauth_login(provider):
+ """Redirect the user to the OAuth provider's authorization page.
+
+ Args:
+ provider: The provider key (e.g. 'google', 'github', 'oidc').
+ """
+ authenticator = UserAuthenticator.get_instance()
+ oauth_backend = authenticator.get_oauth_backend()
+
+ if not oauth_backend:
+ logger.warning("OAuth login attempted but OAuth is not configured")
+ return redirect(url_for("home"))
+
+ client = oauth_backend.get_oauth_client(provider)
+ if client is None:
+ logger.warning("Unknown OAuth provider: %s", provider)
+ return render_template("home.html",
+ login_error=f"Unknown login provider: {provider}",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=config.get("require_password", True),
+ oauth_providers=authenticator.get_login_providers()), 404
+
+ callback_url = url_for("oauth_callback", provider=provider, _external=True)
+ return client.authorize_redirect(callback_url)
+
+
+@app.route("/auth/callback/")
+def oauth_callback(provider):
+ """Handle the OAuth callback after user authenticates with provider.
+
+ Exchanges the authorization code for a token, fetches user profile,
+ applies restrictions, and creates a Potato session.
+ """
+ authenticator = UserAuthenticator.get_instance()
+ oauth_backend = authenticator.get_oauth_backend()
+
+ if not oauth_backend:
+ return redirect(url_for("home"))
+
+ client = oauth_backend.get_oauth_client(provider)
+ if client is None:
+ return redirect(url_for("home"))
+
+ # Check if the provider returned an error (e.g. user denied access)
+ error = request.args.get("error")
+ if error:
+ error_desc = request.args.get("error_description", error)
+ logger.warning("OAuth error from %s: %s - %s", provider, error, error_desc)
+ return render_template("home.html",
+ login_error=f"Login cancelled: {error_desc}",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=config.get("require_password", True),
+ oauth_providers=authenticator.get_login_providers())
+
+ try:
+ # Exchange authorization code for token
+ token = client.authorize_access_token()
+ except Exception as e:
+ logger.error("OAuth token exchange failed for %s: %s", provider, str(e))
+ return render_template("home.html",
+ login_error="Authentication failed. Please try again.",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=config.get("require_password", True),
+ oauth_providers=authenticator.get_login_providers())
+
+ # Get user profile
+ try:
+ if provider == "github":
+ resp = client.get("user", token=token)
+ profile = resp.json()
+ # GitHub may not include email in profile โ fetch from emails API
+ if not profile.get("email"):
+ emails_resp = client.get("user/emails", token=token)
+ emails = emails_resp.json()
+ primary = next((e for e in emails if e.get("primary")), None)
+ if primary:
+ profile["email"] = primary["email"]
+ elif hasattr(token, "get") and token.get("userinfo"):
+ profile = token["userinfo"]
+ else:
+ profile = client.userinfo()
+ except Exception as e:
+ logger.error("Failed to fetch user profile from %s: %s", provider, str(e))
+ return render_template("home.html",
+ login_error="Failed to retrieve user profile.",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=config.get("require_password", True),
+ oauth_providers=authenticator.get_login_providers())
+
+ # Check domain/org restrictions
+ allowed, reason = oauth_backend.check_restrictions(provider, profile)
+ if not allowed:
+ logger.warning("OAuth restriction denied user from %s: %s", provider, reason)
+ return render_template("home.html",
+ login_error=reason,
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=config.get("require_password", True),
+ oauth_providers=authenticator.get_login_providers())
+
+ # Check GitHub org restriction (requires API call)
+ allowed_org = oauth_backend.get_allowed_org(provider)
+ if allowed_org and provider == "github":
+ try:
+ orgs_resp = client.get("user/orgs", token=token)
+ orgs = orgs_resp.json()
+ user_orgs = [o.get("login", "").lower() for o in orgs]
+ if allowed_org.lower() not in user_orgs:
+ reason = (
+ f"Access restricted to members of the '{allowed_org}' "
+ f"GitHub organization."
+ )
+ logger.warning("GitHub org check failed: %s not in %s", profile.get("login"), allowed_org)
+ return render_template("home.html",
+ login_error=reason,
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=config.get("require_password", True),
+ oauth_providers=authenticator.get_login_providers())
+ except Exception as e:
+ logger.error("Failed to check GitHub org membership: %s", str(e))
+ return render_template("home.html",
+ login_error="Failed to verify organization membership.",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=config.get("require_password", True),
+ oauth_providers=authenticator.get_login_providers())
+
+ # Extract user identity
+ try:
+ user_id = oauth_backend.extract_user_id(profile, provider)
+ except ValueError as e:
+ logger.error("Cannot extract user ID from OAuth profile: %s", str(e))
+ return render_template("home.html",
+ login_error="Could not determine your user identity.",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=config.get("require_password", True),
+ oauth_providers=authenticator.get_login_providers())
+
+ # Check auto_register
+ if not oauth_backend.auto_register and not authenticator.is_valid_username(user_id):
+ if not authenticator.is_authorized_user(user_id):
+ logger.warning("OAuth user %s not pre-authorized (auto_register=false)", user_id)
+ return render_template("home.html",
+ login_error="Your account is not authorized for this task.",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ require_password=config.get("require_password", True),
+ oauth_providers=authenticator.get_login_providers())
+
+ # Register the user in the OAuth backend
+ authenticator.add_user(user_id, None,
+ oauth_provider=provider,
+ oauth_profile=profile)
+
+ # Create Flask session
+ session.clear()
+ session['username'] = user_id
+ session.permanent = True
+ logger.info("OAuth login successful: provider=%s, user=%s", provider, user_id)
+
+ # Initialize user state if needed
+ if not get_user_state_manager().has_user(user_id):
+ logger.debug("Initializing state for new OAuth user: %s", user_id)
+ usm = get_user_state_manager()
+ usm.add_user(user_id)
+
+ # Check for debug phase skip
+ if config.get("debug") and config.get("debug_phase"):
+ if apply_debug_phase_skip(user_id):
+ return redirect(url_for("home"))
+ else:
+ usm.advance_phase(user_id)
+ else:
+ usm.advance_phase(user_id)
+
+ return redirect(url_for("home"))
+
+
+@app.route("/login", methods=["GET", "POST"])
+def login():
+ """
+ Handle login requests - render the auth page directly
+
+ Returns:
+ flask.Response: Rendered auth template
+ """
+ logger.debug("Rendering auth page for /login")
+ return auth()
+
+@app.route("/logout", methods=["GET"])
+def logout_page():
+ """
+ Handle user logout requests and redirect to login page.
+
+ For url_direct/prolific login types, renders a standalone logged-out page
+ instead of redirecting to home (which requires URL parameters like PROLIFIC_PID).
+
+ Returns:
+ flask.Response: Redirect to login page or rendered logged-out template
+ """
+ logger.debug("Processing logout request")
+
+ # Check login type before clearing session (config is module-level, not session-dependent)
+ login_config = config.get('login', {})
+ login_type = login_config.get('type', 'standard')
+
+ # Clear the session
+ session.clear()
+ logger.info("User logged out successfully")
+
+ if login_type in ['url_direct', 'prolific']:
+ # Cannot redirect to home โ it requires a URL parameter (e.g., PROLIFIC_PID)
+ return render_template("logged_out.html",
+ title=config.get("annotation_task_name", "Annotation Platform"))
+
+ return redirect(url_for("home")) # Redirect to the login page
+
+@app.route("/logout", methods=["POST"])
+def logout():
+ """
+ Handle user logout requests.
+
+ Features:
+ - Session cleanup
+ - State persistence
+ - Progress saving
+
+ Returns:
+ flask.Response: Redirect to login page
+ """
+ logger.debug("Redirecting /logout to logout_page")
+ return logout_page()
+
+@app.route("/submit_annotation", methods=["POST"])
+def submit_annotation():
+ """
+ DEPRECATED: Handle annotation submission requests.
+
+ This route was added by Cursor and duplicates functionality from /updateinstance.
+ It only handles label annotations (not span annotations) and is used primarily
+ for saving annotations during navigation in newer templates.
+
+ TODO: This route should be deprecated and all functionality moved to /updateinstance
+ to reduce confusion and ensure consistent handling of both label and span annotations.
+
+ Features:
+ - Validation checking
+ - Progress tracking
+ - State updates
+ - AI integration
+ - Data persistence
+
+ Args (from form or JSON):
+ instance_id: ID of annotated instance
+ annotations: annotation data (either JSON string or dict) - LABEL ANNOTATIONS ONLY
+ user_id: user ID (optional, defaults to session)
+
+ Returns:
+ flask.Response: JSON response with submission result
+ """
+ logger.debug("=== SUBMIT ANNOTATION ROUTE START ===")
+ logger.debug(f"Session: {dict(session)}")
+ logger.debug(f"Session username: {session.get('username', 'NOT_SET')}")
+ logger.debug(f"Request content type: {request.content_type}")
+ logger.debug(f"Request is JSON: {request.is_json}")
+ logger.debug(f"Request headers: {dict(request.headers)}")
+ logger.debug(f"Debug mode: {config.get('debug', False)}")
+
+
+ if 'username' not in session:
+ logger.warning("Annotation submission without active session")
+ return jsonify({"status": "error", "message": "No active session"})
+
+ user_id = session['username']
+ logger.debug(f"Using user_id: {user_id}")
+ logger.debug(f"All users in state manager: {get_user_state_manager().get_user_ids()}")
+ logger.debug(f"User state manager has user '{user_id}': {get_user_state_manager().has_user(user_id)}")
+
+ # Handle both form data and JSON data
+ if request.is_json:
+ data = request.get_json()
+ instance_id = data.get("instance_id")
+ annotations = data.get("annotations", {})
+ logger.debug(f"Received JSON data: {data}")
+ else:
+ instance_id = request.form.get("instance_id")
+ annotation_data = request.form.get("annotation_data")
+ logger.debug(f"Received form data: {dict(request.form)}")
+ if annotation_data:
+ annotations = json.loads(annotation_data)
+ logger.debug(f"Parsed annotation_data: {annotations}")
+ else:
+ annotations = {}
+ logger.debug("No annotation_data found in form")
+
+ logger.debug(f"Instance ID: {instance_id}")
+ logger.debug(f"Annotations: {annotations}")
+
+ if not instance_id:
+ logger.warning("Missing instance_id")
+ return jsonify({"status": "error", "message": "Missing instance_id"})
+
+ try:
+ logger.debug(f"Getting user state for user_id: {user_id}")
+ user_state = get_user_state(user_id)
+ logger.debug(f"Retrieved user state: {user_state}")
+ logger.debug(f"User state phase: {user_state.get_phase() if user_state else 'No user state'}")
+
+
+ # Process the annotations
+ annotations_processed = 0
+ for schema_name, label_data in annotations.items():
+ logger.debug(f"Processing schema: {schema_name}, label_data: {label_data}, type: {type(label_data)}")
+
+ if isinstance(label_data, dict):
+ # Nested structure: {'schema': {'label': 'value'}}
+ for label_name, value in label_data.items():
+ label = Label(schema_name, label_name)
+ logger.debug(f"Adding annotation: {schema_name}:{label_name} = {value}")
+ user_state.add_label_annotation(instance_id, label, value)
+ annotations_processed += 1
+ elif isinstance(label_data, str):
+ # Direct string value for text annotations: {'schema': 'value'}
+ # For text annotations, we need to create a label with a default name
+ label = Label(schema_name, "text_box")
+ logger.debug(f"Adding text annotation: {schema_name}:text_box = {label_data}")
+ user_state.add_label_annotation(instance_id, label, label_data)
+ annotations_processed += 1
+ else:
+ logger.warning(f"Unexpected label_data type: {type(label_data)} for schema {schema_name}")
+
+ logger.debug(f"Processed {annotations_processed} annotations")
+
+ # Register the annotator for this instance
+ logger.debug(f"Registering annotator {user_id} for instance {instance_id}")
+ get_item_state_manager().register_annotator(instance_id, user_id)
+
+ # Save the user state
+ logger.debug(f"Saving user state for {user_id}")
+ get_user_state_manager().save_user_state(user_state)
+ logger.debug(f"User state saved successfully")
+
+ # Check if this was an ICL verification task and record the result
+ _maybe_record_icl_verification(user_state, instance_id, annotations)
+
+ # Notify diversity manager of annotation completion (for async embedding)
+ _notify_diversity_manager_annotation(user_state, instance_id)
+
+ # Log the saved annotations
+ all_annotations = user_state.get_all_annotations()
+ logger.debug(f"All annotations after save: {all_annotations}")
+ logger.debug(f"Annotations for instance {instance_id}: {all_annotations.get(instance_id, 'Not found')}")
+
+
+ logger.info(f"Successfully saved annotation for {instance_id} from {user_id}")
+ logger.debug("=== SUBMIT ANNOTATION ROUTE END ===")
+ return jsonify({"status": "success", "message": "Annotation saved successfully", "annotations_processed": annotations_processed})
+
+ except Exception as e:
+ logger.error(f"Error saving annotation: {type(e).__name__}: {str(e)}", exc_info=True)
+ return jsonify({"status": "error", "message": "Failed to save annotation"})
+
+@app.route("/register", methods=["POST"])
+def register():
+ """
+ Register a new user and initialize their user state.
+
+ Args:
+ username: The username to initialize state for
+ """
+ logger.debug("=== REGISTER ROUTE START ===")
+ logger.debug(f"Session before registration: {dict(session)}")
+ logger.debug(f"Request form data: {dict(request.form)}")
+ logger.debug(f"Request headers: {dict(request.headers)}")
+
+ if 'username' in session:
+ logger.warning(f"User already logged in with username: {session['username']}, redirecting to annotate")
+ return home()
+
+ username = request.form.get("email")
+ password = request.form.get("pass")
+
+ logger.debug(f"Registration attempt for username: {username}")
+
+ if not username or not password:
+ logger.warning("Missing username or password")
+ return render_template("home.html",
+ login_error="Username and password are required")
+
+ # Register the user with the authenticator
+ logger.debug("Adding user to authenticator...")
+ user_authenticator = UserAuthenticator.get_instance()
+ result = user_authenticator.add_user(username, password)
+
+ if result != "Success":
+ logger.warning(f"Registration failed for '{username}': {result}")
+ return render_template("home.html", login_error=result)
+
+ # Persist user config if explicitly configured
+ user_authenticator.save_user_config()
+
+ logger.debug("Setting session variables...")
+ session['username'] = username
+ session.permanent = True
+
+ logger.debug(f"Session after registration: {dict(session)}")
+ logger.debug(f"Session ID: {session.sid if hasattr(session, 'sid') else 'No session ID'}")
+ logger.debug(f"User state manager has user '{username}': {get_user_state_manager().has_user(username)}")
+ logger.debug(f"All users in state manager: {get_user_state_manager().get_user_ids()}")
+
+ # Initialize user state if needed
+ if not get_user_state_manager().has_user(username):
+ logger.debug(f"Initializing user state for new user: {username}")
+ init_user_state(username)
+ logger.debug(f"User state initialized. User exists: {get_user_state_manager().has_user(username)}")
+ logger.debug(f"All users in state manager after init: {get_user_state_manager().get_user_ids()}")
+
+ # Ensure user is in the correct starting phase
+ usm = get_user_state_manager()
+ user_state = usm.get_user_state(username)
+ logger.debug(f"Retrieved user state for '{username}': {user_state}")
+ logger.debug(f"User state phase: {user_state.get_phase() if user_state else 'No user state'}")
+
+ # Determine the first phase from config
+ phases_config = config.get('phases', {})
+ phases_order = phases_config.get('order', ['annotation'])
+ first_phase_name = phases_order[0] if phases_order else 'annotation'
+ # Get the phase type from the config (phase name may differ from type, e.g., 'prescreen' has type 'prestudy')
+ first_phase_config = phases_config.get(first_phase_name, {})
+ first_phase_type = first_phase_config.get('type', first_phase_name)
+ first_phase = UserPhase.fromstr(first_phase_type)
+ logger.debug(f"First phase from config: {first_phase_name} (type={first_phase_type}) -> {first_phase}")
+
+ # Set user to the first phase if they're in LOGIN
+ if user_state and user_state.get_phase() == UserPhase.LOGIN:
+ logger.debug(f"Advancing user {username} to first phase: {first_phase}")
+ # Use first_phase_name as the page since that's the key in the phase config
+ user_state.advance_to_phase(first_phase, first_phase_name)
+ logger.debug(f"User state phase after advancement: {user_state.get_phase()}")
+
+ # Assign instances if user doesn't have any
+ if user_state and not user_state.has_assignments():
+ logger.debug(f"Assigning instances to user {username}")
+ get_item_state_manager().assign_instances_to_user(user_state)
+ logger.debug(f"User has assignments after assignment: {user_state.has_assignments()}")
+
+ logger.debug("=== REGISTER ROUTE END - Redirecting to home ===")
+ # Redirect to home which will route to the appropriate phase
+ return redirect(url_for("home"))
+
+@app.route("/consent", methods=["GET", "POST"])
+def consent():
+ """
+ Handle the consent phase of the annotation process.
+
+ Returns:
+ flask.Response: Rendered template or redirect
+ """
+ if 'username' not in session:
+ return home()
+
+ username = session['username']
+ user_state = get_user_state(username)
+ logger.debug(f'CONSENT: user_state: {user_state}')
+ logger.debug(f'CONSENT: user_state.get_phase(): {user_state.get_phase()}')
+
+ # Check that the user is still in the consent phase
+ if user_state.get_phase() != UserPhase.CONSENT:
+ # If not in the consent phase, redirect
+ return home()
+
+ # If the user is returning information from the page
+ if request.method == 'POST':
+ # The form should require that the user consent to the study
+ logger.debug(f'POST -> CONSENT: {request.form}')
+
+ # Now that the user has consented, advance the state
+ # and have the home page redirect to the appropriate next phase
+ usm = get_user_state_manager()
+ usm.advance_phase(session['username'])
+
+ # Redirect to force a clean GET request (fixes POST leakage, issue #124)
+ return redirect(url_for("home"))
+ # Show the current consent form
+ else:
+ logger.debug("GET <- CONSENT")
+ return get_current_page_html(config, username)
+
+@app.route("/instructions", methods=["GET", "POST"])
+def instructions():
+ """
+ Handle the instructions phase of the annotation process.
+
+ Returns:
+ flask.Response: Rendered template or redirect
+ """
+ if 'username' not in session:
+ return home()
+
+ username = session['username']
+ user_state = get_user_state(username)
+
+ # Check that the user is in the instructions phase
+ if user_state.get_phase() != UserPhase.INSTRUCTIONS:
+ # If not in the instructions phase, redirect
+ return home()
+
+ # If the user is returning information from the page
+ if request.method == 'POST':
+ logger.debug(f'POST -> INSTRUCTIONS: {request.form}')
+
+ # Now that the user has read the instructions, advance the state
+ # and have the home page redirect to the appropriate next phase
+ usm = get_user_state_manager()
+ usm.advance_phase(session['username'])
+
+ # Redirect to force a clean GET request (fixes POST leakage, issue #124)
+ return redirect(url_for("home"))
+
+ # Show the current set of instructions
+ else:
+ logger.debug(f'GET <-- INSTRUCTIONS')
+ return get_current_page_html(config, username)
+
+@app.route("/training", methods=["GET", "POST"])
+def training():
+ """
+ Handle the training phase of the annotation process.
+
+ This route manages the training phase where users practice annotation
+ with feedback on their performance. It supports:
+ - Displaying training instances with correct answers
+ - Processing user annotations and providing feedback
+ - Tracking training progress and performance
+ - Advancing users based on training performance criteria
+ - Allowing retries for failed attempts
+ - Kicking out users who exceed max_mistakes threshold
+
+ Training Configuration Options:
+ - min_correct: Minimum correct answers needed to pass
+ - require_all_correct: Whether all questions must be correct
+ - max_mistakes: Maximum total mistakes before failure (kicked out)
+ - max_mistakes_per_question: Maximum mistakes per question before failure
+ - allow_retry: Whether to allow retrying incorrect answers
+ - failure_action: "move_to_done" (kick out) or "repeat_training"
+
+ Returns:
+ flask.Response: Rendered template or redirect
+ """
+ if 'username' not in session:
+ return home()
+
+ username = session['username']
+ user_state = get_user_state(username)
+
+ # Check that the user is in the training phase
+ if user_state.get_phase() != UserPhase.TRAINING:
+ logger.debug(f'User {username} not in training phase, redirecting')
+ return home()
+
+ # Check if training is enabled in config
+ training_config = config.get('training', {})
+ if not training_config.get('enabled', False):
+ logger.debug('Training not enabled, advancing to next phase')
+ usm = get_user_state_manager()
+ usm.advance_phase(username)
+ return redirect(url_for("home"))
+
+ # Get training state and initialize max_mistakes from config if not set
+ training_state = user_state.get_training_state()
+ passing_criteria = training_config.get('passing_criteria', {})
+
+ # Initialize training instances if not already done
+ if not training_state.training_instances:
+ training_instances = get_training_instances()
+ training_state.set_training_instances([item.get_id() for item in training_instances])
+
+ # Set max_mistakes from config
+ if training_state.max_mistakes == -1 and 'max_mistakes' in passing_criteria:
+ training_state.set_max_mistakes(passing_criteria.get('max_mistakes', -1))
+ if training_state.max_mistakes_per_question == -1 and 'max_mistakes_per_question' in passing_criteria:
+ training_state.set_max_mistakes_per_question(passing_criteria.get('max_mistakes_per_question', -1))
+
+ # Check if user has already failed due to too many mistakes
+ if training_state.is_failed() or training_state.should_fail_due_to_mistakes():
+ training_state.set_failed(True)
+ logger.info(f'User {username} has failed training due to too many mistakes')
+ # Move to DONE phase (kick out)
+ user_state.set_current_phase_and_page((UserPhase.DONE, None))
+ return render_template("training_failed.html",
+ message="You have exceeded the maximum number of allowed mistakes and cannot continue.",
+ total_mistakes=training_state.get_total_mistakes(),
+ max_mistakes=training_state.max_mistakes,
+ annotation_task_name=config.get("annotation_task_name", "Annotation Platform"),
+ username=username)
+
+ # Get progress info
+ total_questions = len(training_state.training_instances)
+ current_question_num = training_state.get_current_question_index() + 1
+
+ # Handle POST requests (annotation submission)
+ if request.method == 'POST':
+ logger.debug(f'POST -> TRAINING: {request.form}')
+
+ # Get the current training instance
+ current_instance = user_state.get_current_training_instance()
+ if not current_instance:
+ logger.error(f'No training instance available for user {username}')
+ return render_template("error.html", message="No training instance available")
+
+ instance_id = current_instance.get_id()
+ instance_text = current_instance.get_data().get('displayed_text', current_instance.get_data().get('text', ''))
+
+ # Process the annotation
+ if request.is_json:
+ annotation_data = request.get_json()
+ else:
+ annotation_data = dict(request.form)
+
+ # Get correct answers for this training instance
+ correct_answers = get_training_correct_answers(instance_id)
+ if not correct_answers:
+ logger.error(f'No correct answers found for training instance {instance_id}')
+ return render_template("error.html", message="Training data error")
+
+ # Validate and process the annotation
+ try:
+ # Update user's training answer
+ user_state.update_training_answer(instance_id, annotation_data)
+
+ # Check if the answer is correct
+ is_correct = check_training_answer(annotation_data, correct_answers)
+
+ # Track category performance for category-based assignment
+ instance_categories = get_training_instance_categories(instance_id)
+ if instance_categories:
+ training_state.record_category_answer(instance_categories, is_correct)
+
+ if is_correct:
+ logger.info(f'User {username} answered training question {instance_id} correctly')
+ # Record correct answer
+ training_state.add_answer(instance_id, True, training_state.get_mistakes_for_question(instance_id) + 1)
+ training_state.clear_feedback()
+
+ # Check if user has passed based on min_correct
+ min_correct = passing_criteria.get('min_correct', len(training_state.training_instances))
+ if training_state.get_correct_answer_count() >= min_correct:
+ # User has passed training
+ training_state.set_passed(True)
+ logger.info(f'User {username} passed training with {training_state.get_correct_answer_count()} correct answers')
+
+ # Calculate category qualifications based on training performance
+ cat_config = config.get('category_assignment', {})
+ if cat_config.get('enabled', False):
+ qual_config = cat_config.get('qualification', {})
+ threshold = qual_config.get('threshold', 0.7)
+ min_questions = qual_config.get('min_questions', 1)
+ qualified = user_state.calculate_and_set_qualifications(threshold, min_questions)
+ if qualified:
+ logger.info(f'User {username} qualified for categories: {qualified}')
+
+ usm = get_user_state_manager()
+ usm.advance_phase(username)
+ return redirect(url_for("home"))
+
+ # Move to next training question or complete training
+ if user_state.advance_training_question():
+ # More questions available
+ training_state.set_feedback(True, "Correct! Moving to next question.", False)
+ # Get next instance for display
+ next_instance = user_state.get_current_training_instance()
+ next_instance_text = next_instance.get_data().get('displayed_text', next_instance.get_data().get('text', ''))
+ return render_template("training.html",
+ instance_text=next_instance_text,
+ instance_id=next_instance.get_id(),
+ feedback="Correct! Moving to next question.",
+ feedback_type="success",
+ show_feedback=True,
+ allow_retry=False,
+ current_question=current_question_num + 1,
+ total_questions=total_questions,
+ correct_count=training_state.get_correct_answer_count(),
+ mistake_count=training_state.get_total_mistakes(),
+ annotation_task_name=config.get("annotation_task_name", "Annotation Platform"),
+ username=username)
+ else:
+ # All questions completed
+ require_all = passing_criteria.get('require_all_correct', False)
+ if require_all and training_state.get_correct_answer_count() < total_questions:
+ # User didn't get all correct
+ training_state.set_failed(True)
+ user_state.set_current_phase_and_page((UserPhase.DONE, None))
+ return render_template("training_failed.html",
+ message="You did not answer all training questions correctly.",
+ correct_count=training_state.get_correct_answer_count(),
+ total_questions=total_questions,
+ annotation_task_name=config.get("annotation_task_name", "Annotation Platform"),
+ username=username)
+ else:
+ # Training completed successfully
+ training_state.set_passed(True)
+ logger.info(f'User {username} completed training successfully')
+
+ # Calculate category qualifications based on training performance
+ cat_config = config.get('category_assignment', {})
+ if cat_config.get('enabled', False):
+ qual_config = cat_config.get('qualification', {})
+ threshold = qual_config.get('threshold', 0.7)
+ min_questions = qual_config.get('min_questions', 1)
+ qualified = user_state.calculate_and_set_qualifications(threshold, min_questions)
+ if qualified:
+ logger.info(f'User {username} qualified for categories: {qualified}')
+
+ usm = get_user_state_manager()
+ usm.advance_phase(username)
+ return redirect(url_for("home"))
+ else:
+ logger.info(f'User {username} answered training question {instance_id} incorrectly')
+ # Record the mistake
+ training_state.record_mistake(instance_id)
+
+ # Check if user should fail due to too many mistakes
+ if training_state.should_fail_due_to_mistakes():
+ training_state.set_failed(True)
+ logger.info(f'User {username} failed training - exceeded max_mistakes ({training_state.max_mistakes})')
+ user_state.set_current_phase_and_page((UserPhase.DONE, None))
+ return render_template("training_failed.html",
+ message="You have exceeded the maximum number of allowed mistakes.",
+ total_mistakes=training_state.get_total_mistakes(),
+ max_mistakes=training_state.max_mistakes,
+ annotation_task_name=config.get("annotation_task_name", "Annotation Platform"),
+ username=username)
+
+ # Check if user should fail due to too many mistakes on this question
+ if training_state.should_fail_question_due_to_mistakes(instance_id):
+ training_state.set_failed(True)
+ logger.info(f'User {username} failed training - exceeded max_mistakes_per_question on {instance_id}')
+ user_state.set_current_phase_and_page((UserPhase.DONE, None))
+ return render_template("training_failed.html",
+ message="You have made too many mistakes on a single question.",
+ question_mistakes=training_state.get_mistakes_for_question(instance_id),
+ max_mistakes_per_question=training_state.max_mistakes_per_question,
+ annotation_task_name=config.get("annotation_task_name", "Annotation Platform"),
+ username=username)
+
+ # Get explanation for incorrect answer
+ explanation = get_training_explanation(instance_id)
+
+ # Check if user should be allowed to retry
+ allow_retry = training_config.get('allow_retry', True)
+
+ if allow_retry:
+ training_state.set_feedback(True, f"Incorrect. {explanation}", True)
+ return render_template("training.html",
+ instance_text=instance_text,
+ instance_id=instance_id,
+ feedback=f"Incorrect. {explanation}",
+ feedback_type="error",
+ show_feedback=True,
+ allow_retry=True,
+ current_question=current_question_num,
+ total_questions=total_questions,
+ correct_count=training_state.get_correct_answer_count(),
+ mistake_count=training_state.get_total_mistakes(),
+ annotation_task_name=config.get("annotation_task_name", "Annotation Platform"),
+ username=username)
+ else:
+ # No retry allowed - check failure action
+ failure_action = training_config.get('failure_action', 'move_to_done')
+ if failure_action == 'move_to_done':
+ training_state.set_failed(True)
+ logger.info(f'User {username} failed training - no retry allowed')
+ user_state.set_current_phase_and_page((UserPhase.DONE, None))
+ return render_template("training_failed.html",
+ message="You answered incorrectly and retries are not allowed.",
+ explanation=explanation,
+ annotation_task_name=config.get("annotation_task_name", "Annotation Platform"),
+ username=username)
+ else:
+ # Advance to next question even though wrong
+ if user_state.advance_training_question():
+ next_instance = user_state.get_current_training_instance()
+ next_instance_text = next_instance.get_data().get('displayed_text', next_instance.get_data().get('text', ''))
+ training_state.set_feedback(True, f"Incorrect. {explanation} Moving to next question.", False)
+ return render_template("training.html",
+ instance_text=next_instance_text,
+ instance_id=next_instance.get_id(),
+ feedback=f"Previous answer was incorrect: {explanation}",
+ feedback_type="warning",
+ show_feedback=True,
+ allow_retry=False,
+ current_question=current_question_num + 1,
+ total_questions=total_questions,
+ correct_count=training_state.get_correct_answer_count(),
+ mistake_count=training_state.get_total_mistakes(),
+ annotation_task_name=config.get("annotation_task_name", "Annotation Platform"),
+ username=username)
+ else:
+ # No more questions - check if passed
+ min_correct = passing_criteria.get('min_correct', total_questions)
+ if training_state.get_correct_answer_count() >= min_correct:
+ training_state.set_passed(True)
+ usm = get_user_state_manager()
+ usm.advance_phase(username)
+ return redirect(url_for("home"))
+ else:
+ training_state.set_failed(True)
+ user_state.set_current_phase_and_page((UserPhase.DONE, None))
+ return render_template("training_failed.html",
+ message="You did not meet the minimum correct answers requirement.",
+ correct_count=training_state.get_correct_answer_count(),
+ min_correct=min_correct,
+ annotation_task_name=config.get("annotation_task_name", "Annotation Platform"),
+ username=username)
+
+ except Exception as e:
+ logger.error(f'Error processing training annotation: {e}')
+ import traceback
+ traceback.print_exc()
+ return render_template("error.html", message="Error processing training annotation")
+
+ # Handle GET requests (display training question)
+ else:
+ logger.debug(f'GET <-- TRAINING for user {username}')
+
+ # Get the current training instance
+ current_instance = user_state.get_current_training_instance()
+ if not current_instance:
+ logger.error(f'No training instance available for user {username}')
+ return render_template("error.html", message="No training instance available")
+
+ instance_text = current_instance.get_data().get('displayed_text', current_instance.get_data().get('text', ''))
+
+ # Check if we should show feedback from previous attempt
+ show_feedback = training_state.show_feedback if training_state else False
+ feedback_message = training_state.feedback_message if training_state else ""
+ allow_retry = training_state.allow_retry if training_state else False
+
+ return render_template("training.html",
+ instance_text=instance_text,
+ instance_id=current_instance.get_id(),
+ feedback=feedback_message,
+ feedback_type="error" if allow_retry else "info",
+ show_feedback=show_feedback,
+ allow_retry=allow_retry,
+ current_question=current_question_num,
+ total_questions=total_questions,
+ correct_count=training_state.get_correct_answer_count(),
+ mistake_count=training_state.get_total_mistakes(),
+ annotation_task_name=config.get("annotation_task_name", "Annotation Platform"),
+ username=username)
+
+
+def check_training_answer(user_answer: dict, correct_answers: dict) -> bool:
+ """
+ Check if the user's answer matches the correct answers.
+
+ Handles different annotation types:
+ - Radio/single select: string comparison
+ - Multiselect/checkbox: set comparison (order-independent)
+ - Likert/number: numeric comparison
+ - Text: exact or fuzzy string match
+
+ Args:
+ user_answer: Dictionary of user's answers by schema name
+ correct_answers: Dictionary of correct answers by schema name
+
+ Returns:
+ True if all answers are correct, False otherwise
+ """
+ for schema_name, correct_value in correct_answers.items():
+ if schema_name not in user_answer:
+ return False
+
+ user_value = user_answer[schema_name]
+
+ # Handle multiselect/checkbox (list comparison)
+ if isinstance(correct_value, list):
+ if isinstance(user_value, list):
+ if set(user_value) != set(correct_value):
+ return False
+ elif isinstance(user_value, str):
+ # Single value submitted, check if it's the only correct answer
+ if len(correct_value) != 1 or user_value not in correct_value:
+ return False
+ else:
+ return False
+ # Handle numeric values
+ elif isinstance(correct_value, (int, float)):
+ try:
+ if float(user_value) != float(correct_value):
+ return False
+ except (ValueError, TypeError):
+ return False
+ # Handle string comparison (radio, text)
+ else:
+ if str(user_value).strip().lower() != str(correct_value).strip().lower():
+ return False
+
+ return True
+
+@app.route("/prestudy", methods=["GET", "POST"])
+def prestudy():
+ """
+ Handle the prestudy phase of the annotation process.
+
+ Returns:
+ flask.Response: Rendered template or redirect
+ """
+ if 'username' not in session:
+ return home()
+
+ username = session['username']
+ user_state = get_user_state(username)
+
+ # Check that the user is in the prestudy phase
+ if user_state.get_phase() != UserPhase.PRESTUDY:
+ logger.debug('NOT IN PRESTUDY PHASE')
+ return home()
+
+ # If the user is returning information from the page
+ if request.method == 'POST':
+ logger.debug(f'POST -> PRESTUDY: {request.form}')
+
+ # Advance the state and redirect to the appropriate next phase
+ usm = get_user_state_manager()
+ usm.advance_phase(session['username'])
+
+ # Redirect to force a clean GET request (fixes POST leakage, issue #124)
+ return redirect(url_for("home"))
+
+ # Show the current prestudy page
+ else:
+ logger.debug("GET <-- PRESTUDY")
+ return get_current_page_html(config, username)
+
+def _check_required_or_block(user_state, instance_id: str):
+ """Check required annotations and return a 400 response if not met, or None if OK."""
+ unsatisfied = _instance_meets_required_annotation_rules(user_state, instance_id)
+ if unsatisfied:
+ msg = f"Required annotation(s) not completed: {', '.join(unsatisfied)}"
+ logger.info(f"Blocking navigation: {msg} for instance {instance_id}")
+ if request.is_json:
+ return jsonify({"status": "validation_error", "message": msg, "unsatisfied_schemas": unsatisfied}), 400
+ # For form POSTs, fall through to render the page with existing annotations
+ return None
+ return None
+
+
+def _ibws_check_and_advance(user_state) -> bool:
+ """Check if IBWS round is complete and advance to next round if so.
+
+ Returns True if new tuples were added and the user was reassigned.
+ """
+ from potato.ibws_manager import get_ibws_manager
+
+ ibws_mgr = get_ibws_manager()
+ if not ibws_mgr or ibws_mgr.is_completed():
+ return False
+
+ ism = get_item_state_manager()
+ usm = get_user_state_manager()
+
+ # Find the BWS schema name
+ bws_schema_name = None
+ for scheme in config.get("annotation_schemes", []):
+ if scheme.get("annotation_type") == "bws":
+ bws_schema_name = scheme["name"]
+ break
+
+ if not bws_schema_name:
+ return False
+
+ if not ibws_mgr.check_round_complete(ism, bws_schema_name):
+ return False
+
+ # Round is complete โ advance to next round
+ logger.info(f"IBWS: Round {ibws_mgr.current_round} complete, advancing")
+ new_tuples = ibws_mgr.advance_round(ism, usm, bws_schema_name)
+
+ if not new_tuples:
+ logger.info("IBWS: No more rounds needed, annotation complete")
+ return False
+
+ # Add new tuples to ISM
+ id_key = config["item_properties"]["id_key"]
+ for t in new_tuples:
+ ism.add_item(str(t[id_key]), t)
+
+ # Re-render displayed text for new items
+ from potato.flask_server import _render_displayed_text # noqa: cross-import
+ text_key = config["item_properties"]["text_key"]
+ _render_displayed_text(text_key)
+
+ # Reassign instances to all active users
+ for us in usm.get_all_users():
+ ism.assign_instances_to_user(us)
+
+ logger.info(f"IBWS: Added {len(new_tuples)} new tuples for round {ibws_mgr.current_round}")
+ return True
+
+
+@app.route("/annotate", methods=["GET", "POST"])
+def annotate():
+ """
+ Handle annotation page requests.
+ """
+ logger.debug("=== ANNOTATE ROUTE START ===")
+ logger.debug(f"Session: {dict(session)}")
+ logger.debug(f"Session username: {session.get('username', 'NOT_SET')}")
+ logger.debug(f"Debug mode: {config.get('debug', False)}")
+ logger.debug(f"Request method: {request.method}")
+ logger.debug(f"Request headers: {dict(request.headers)}")
+ logger.debug(f"Request content type: {request.content_type}")
+ logger.debug(f"Request is JSON: {request.is_json}")
+
+ # Check if user is logged in
+ if 'username' not in session:
+ logger.warning("Unauthorized access attempt to annotate page")
+ return redirect(url_for("home"))
+
+ username = session['username']
+ logger.debug(f"Using username: {username}")
+ logger.debug(f"All users in state manager: {get_user_state_manager().get_user_ids()}")
+
+ # Ensure user state exists
+ if not get_user_state_manager().has_user(username):
+ logger.info(f"Creating missing user state for {username}")
+ init_user_state(username)
+ logger.debug(f"User state created. User exists: {get_user_state_manager().has_user(username)}")
+
+ logger.debug("Handling annotation request")
+
+ user_state = get_user_state(username)
+ logger.debug(f"Retrieved state for user: {username}")
+ logger.debug(f"User state: {user_state}")
+ logger.debug(f"User state phase: {user_state.get_phase() if user_state else 'No user state'}")
+
+ # Check user phase โ guard against leaked nav button POSTs from non-annotation pages
+ if not user_state or user_state.get_phase() != UserPhase.ANNOTATION:
+ cur_phase = user_state.get_phase() if user_state else None
+ logger.info(f"User {username} not in annotation phase (phase={cur_phase}), redirecting.")
+
+ # If a nav POST arrived from a non-annotation page, handle it gracefully
+ if request.method == 'POST':
+ action = None
+ if request.is_json and request.json:
+ action = request.json.get('action')
+ elif request.form:
+ action = request.form.get('action')
+
+ if action == 'next_instance':
+ # Treat as "submit current phase page" - advance to next phase
+ logger.info(f"Leaked next_instance from phase {cur_phase}, advancing phase")
+ get_user_state_manager().advance_phase(username)
+ return redirect(url_for("home"))
+ elif action == 'prev_instance':
+ if get_user_state_manager().retreat_phase(username):
+ logger.info(f"Leaked prev_instance from phase {cur_phase}, moved to previous page/phase")
+ else:
+ logger.info(f"Leaked prev_instance from phase {cur_phase}, back navigation not allowed")
+ return redirect(url_for("home"))
+
+ # For non-nav POSTs (phase form submissions) and GETs, delegate to home()
+ # so POST data is preserved for phase processing
+ return home()
+
+ # If the user hasn't yet been assigned anything to annotate, do so now
+ if not user_state.has_assignments():
+ logger.debug(f"User {username} has no assignments, assigning instances")
+ get_item_state_manager().assign_instances_to_user(user_state)
+ logger.debug(f"User has assignments after assignment: {user_state.has_assignments()}")
+
+ # If assignment produced nothing and user is an adjudicator,
+ # redirect them to the adjudication page
+ if not user_state.has_assignments():
+ adj_mgr = get_adjudication_manager()
+ if adj_mgr and adj_mgr.is_adjudicator(username):
+ logger.info(f"Adjudicator {username} has no annotation items, redirecting to /adjudicate")
+ return redirect(url_for("adjudicate"))
+
+ # IBWS: Check if round is complete and advance if needed
+ if config.get("ibws_config"):
+ _ibws_check_and_advance(user_state)
+
+ # See if this user has finished annotating all of their assigned instances
+ if not user_state.has_remaining_assignments():
+ # For IBWS, don't advance phase if more rounds are possible
+ if config.get("ibws_config"):
+ from potato.ibws_manager import get_ibws_manager
+ ibws_mgr = get_ibws_manager()
+ if ibws_mgr and not ibws_mgr.is_completed():
+ # Still waiting for other annotators to finish the round
+ pass
+ else:
+ logger.debug(f"User {username} has no remaining assignments (IBWS complete), advancing phase")
+ get_user_state_manager().advance_phase(username)
+ return redirect(url_for("home"))
+ else:
+ logger.debug(f"User {username} has no remaining assignments, advancing phase")
+ get_user_state_manager().advance_phase(username)
+ return redirect(url_for("home"))
+
+ _inject_quality_control_item_if_needed(username, user_state)
+
+ # Handle POST requests
+ if request.method == 'POST':
+ logger.debug(f"POST request to annotate")
+ if request.is_json:
+ logger.debug(f"POST JSON data: {request.get_json()}")
+ else:
+ logger.debug(f"POST form data: {dict(request.form)}")
+
+ if request.is_json and request.json and 'action' in request.json:
+ logger.debug(f"Action from JSON: {request.json['action']}")
+ action = request.json['action']
+ else:
+ logger.debug(f"Action from form: {request.form.get('action', 'init')}")
+ action = request.form['action'] if 'action' in request.form else "init"
+
+ logger.debug(f"Processing action: {action}")
+
+ # NOTE: Annotations are saved in real-time via /updateinstance endpoint when users
+ # click checkboxes, radio buttons, etc. This ensures proper timing tracking for
+ # behavioral data. We do NOT save annotations during navigation - they should
+ # already be saved by the time the user navigates.
+
+ if action == "prev_instance":
+ logger.debug(f"Moving to previous instance for user: {username}")
+ moved_back = move_to_prev_instance(username)
+ if not moved_back and user_state.get_current_instance_index() <= 0:
+ if get_user_state_manager().retreat_phase(username):
+ logger.debug(f"User {username} at first annotation instance, moved back to previous phase/page")
+ return redirect(url_for("home"))
+ else:
+ logger.debug(f"User {username} at first annotation instance, back navigation not allowed")
+ acm = get_ai_cache_manager()
+ if acm:
+ acm.start_prefetch(user_state.current_instance_index,
+ getattr(acm, "prefetch_page_count_on_prev", 0) )
+ elif action == "next_instance":
+ logger.debug(f"Moving to next instance for user: {username}")
+ # Check required annotations before allowing forward navigation
+ current_id = user_state.get_current_instance_id()
+ block_response = _check_required_or_block(user_state, current_id)
+ if block_response is not None:
+ return block_response
+
+ moved_forward = move_to_next_instance(username)
+ if not moved_forward and user_state.is_at_end_index():
+ logger.debug(f"User {username} reached the end of assigned instances")
+ if not user_state.has_remaining_assignments():
+ # IBWS: try to advance round before giving up
+ if config.get("ibws_config"):
+ advanced = _ibws_check_and_advance(user_state)
+ if advanced:
+ # New tuples were added โ reassign and continue
+ pass
+ else:
+ from potato.ibws_manager import get_ibws_manager
+ ibws_mgr = get_ibws_manager()
+ if ibws_mgr and ibws_mgr.is_completed():
+ logger.debug(f"User {username} completed all IBWS rounds")
+ get_user_state_manager().advance_phase(username)
+ return redirect(url_for("home"))
+ else:
+ logger.debug(f"User {username} completed all assignments at end-of-list")
+ get_user_state_manager().advance_phase(username)
+ return redirect(url_for("home"))
+ acm = get_ai_cache_manager()
+ if acm:
+ acm.start_prefetch(user_state.current_instance_index, getattr(acm,"prefetch_page_count_on_next", 0))
+
+ elif action == "go_to":
+ # Try to get go_to from JSON first, then form
+ go_to_value = None
+ if request.is_json and request.json.get("go_to") is not None:
+ go_to_value = request.json.get("go_to")
+ elif request.form.get("go_to") is not None:
+ go_to_value = request.form.get("go_to")
+
+ logger.debug(f"go_to action with value: {go_to_value}")
+ if go_to_value is not None:
+ # Block forward go_to if required annotations aren't met
+ target_index = int(go_to_value)
+ if target_index > user_state.current_instance_index:
+ current_id = user_state.get_current_instance_id()
+ block_response = _check_required_or_block(user_state, current_id)
+ if block_response is not None:
+ return block_response
+
+ go_to_id(username, go_to_value)
+ acm = get_ai_cache_manager()
+ if acm:
+ acm.start_prefetch(user_state.current_instance_index, 1)
+ acm.start_prefetch(user_state.current_instance_index, -1)
+
+ else:
+ logger.warning('go_to action requested but no go_to value provided')
+ elif action == "jump_to_unannotated":
+ # Check required annotations before jumping forward
+ current_id = user_state.get_current_instance_id()
+ block_response = _check_required_or_block(user_state, current_id)
+ if block_response is not None:
+ return block_response
+
+ # Find the next unannotated instance and jump to it
+ next_idx = user_state.find_next_unannotated_index()
+ if next_idx is not None:
+ logger.debug(f"Jumping to next unannotated instance at index {next_idx}")
+ user_state.go_to_index(next_idx)
+ else:
+ logger.debug(f"No unannotated instances found for user {username}")
+ # Return a JSON response indicating no unannotated items
+ if request.is_json:
+ return jsonify({"status": "no_unannotated", "message": "All items have been annotated"})
+ elif action == "jump_to_unannotated_prev":
+ # Find the previous unannotated instance and jump to it
+ prev_idx = user_state.find_prev_unannotated_index()
+ if prev_idx is not None:
+ logger.debug(f"Jumping to previous unannotated instance at index {prev_idx}")
+ user_state.go_to_index(prev_idx)
+ else:
+ logger.debug(f"No unannotated instances found for user {username}")
+ # Return a JSON response indicating no unannotated items
+ if request.is_json:
+ return jsonify({"status": "no_unannotated", "message": "All items have been annotated"})
+ else:
+ logger.debug(f'Action "{action}" - no specific handling')
+
+ # After processing the action, check again if user has completed all assignments
+ # This handles the case where the user just finished their last item
+ if not user_state.has_remaining_assignments():
+ if config.get("ibws_config"):
+ from potato.ibws_manager import get_ibws_manager
+ ibws_mgr = get_ibws_manager()
+ if ibws_mgr and not ibws_mgr.is_completed():
+ # IBWS still running โ don't advance phase, try to get more tuples
+ _ibws_check_and_advance(user_state)
+ else:
+ logger.debug(f"User {username} has completed all IBWS assignments, advancing phase")
+ get_user_state_manager().advance_phase(username)
+ return redirect(url_for("home"))
+ else:
+ logger.debug(f"User {username} has completed all assignments, advancing phase")
+ get_user_state_manager().advance_phase(username)
+ # Use redirect to ensure next phase handler gets a GET request (fixes issue #115)
+ return redirect(url_for("home"))
+
+ # Handle GET requests with instance_id query parameter
+ if request.method == 'GET' and request.args.get('instance_id'):
+ instance_id = request.args.get('instance_id')
+ logger.debug(f"GET request with instance_id parameter: {instance_id}")
+
+ # Find the index of this instance in the user's assigned instances
+ try:
+ instance_index = user_state.instance_id_to_order.get(instance_id)
+ if instance_index is None:
+ instance_index = user_state.instance_id_ordering.index(instance_id)
+ user_state.instance_id_to_order[instance_id] = instance_index
+ logger.debug(f"Found instance {instance_id} at index {instance_index}")
+
+ # Update the user's current instance to match the URL parameter
+ if instance_index != user_state.current_instance_index:
+ logger.debug(f"Updating user's current instance from index {user_state.current_instance_index} to {instance_index}")
+ user_state.current_instance_index = instance_index
+ else:
+ logger.debug(f"User already on instance {instance_id} at index {instance_index}")
+ except ValueError:
+ logger.warning(f"Instance {instance_id} not found in user's assigned instances")
+ # Don't change the current instance if the requested one isn't assigned to this user
+
+ logger.debug("=== ANNOTATE ROUTE END ===")
+ # Render the page with any existing annotations
+ # Prevent browser caching so window.location.reload() always gets fresh content
+ # (browsers may serve stale cached GET responses after JS-triggered reloads)
+ response = make_response(render_page_with_annotations(username))
+ response.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
+ response.headers['Pragma'] = 'no-cache'
+ return response
+
+@app.route('/get_ai_suggestion', methods=['GET'])
+def get_ai_suggestion():
+ if 'username' not in session:
+ return home()
+
+ username = session['username']
+ user_state = get_user_state(username)
+ ais = get_ai_cache_manager()
+
+ if ais is None:
+ return jsonify({"error": "AI support not enabled"}), 400
+
+ try:
+ annotation_id = int(request.args.get('annotationId'))
+ except (ValueError, TypeError):
+ return jsonify({"error": "Invalid annotationId"}), 400
+
+ # Validate annotation_id is within range
+ num_schemes = len(config.get("annotation_schemes", []))
+ if annotation_id < 0 or annotation_id >= num_schemes:
+ return jsonify({"error": "annotationId out of range"}), 400
+
+ ai_assistant = request.args.get('aiAssistant')
+
+ instance_id = user_state.get_current_instance_index()
+
+ res = ais.get_ai_help(instance_id, annotation_id, ai_assistant)
+ logger.debug(f"AI suggestion result: {res}")
+
+ # Ensure proper JSON response with correct content-type
+ if isinstance(res, dict):
+ return jsonify(res)
+ elif isinstance(res, str):
+ # If it's an error message string, wrap it
+ return jsonify({"error": res})
+ else:
+ return jsonify(res)
+
+
+@app.route('/api/option_highlights/', methods=['GET'])
+def get_option_highlights(annotation_id):
+ """Get AI-suggested option highlights for a specific annotation.
+
+ Returns the top-k most likely correct options for dimming non-highlighted options.
+
+ Args:
+ annotation_id: The annotation scheme index
+
+ Returns:
+ JSON with highlighted options and configuration:
+ {
+ "highlighted": ["option1", "option2"],
+ "top_k": 3,
+ "confidence": 0.85,
+ "config": {...}
+ }
+ """
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ username = session['username']
+ user_state = get_user_state(username)
+ ais = get_ai_cache_manager()
+
+ if ais is None:
+ return jsonify({"error": "AI support not enabled"}), 400
+
+ instance_id = user_state.get_current_instance_index()
+
+ # Get option highlights
+ result = ais.get_option_highlights(instance_id, annotation_id)
+
+ # Add only frontend-needed configuration info (exclude internal schema details)
+ full_config = ais.get_option_highlighting_config()
+ result["config"] = {
+ "enabled": full_config.get("enabled", False),
+ "top_k": full_config.get("top_k", 3),
+ "dim_opacity": full_config.get("dim_opacity", 0.3),
+ "auto_apply": full_config.get("auto_apply", False),
+ }
+
+ return jsonify(result)
+
+
+@app.route('/api/option_highlights/config', methods=['GET'])
+def get_option_highlighting_config():
+ """Get the option highlighting configuration.
+
+ Returns:
+ JSON with option highlighting configuration:
+ {
+ "enabled": true,
+ "top_k": 3,
+ "dim_opacity": 0.4,
+ "auto_apply": true,
+ "schemas": null,
+ "prefetch_count": 20
+ }
+ """
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ ais = get_ai_cache_manager()
+
+ if ais is None:
+ return jsonify({
+ "enabled": False,
+ "top_k": 3,
+ "dim_opacity": 0.4,
+ "auto_apply": True,
+ "schemas": None,
+ "prefetch_count": 20
+ })
+
+ return jsonify(ais.get_option_highlighting_config())
+
+
+@app.route('/api/option_highlights/prefetch', methods=['POST'])
+def trigger_option_highlight_prefetch():
+ """Trigger prefetching of option highlights for upcoming items.
+
+ Request body (optional):
+ {
+ "count": 20 // Number of items to prefetch (uses config default if not provided)
+ }
+
+ Returns:
+ JSON with prefetch status
+ """
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ username = session['username']
+ user_state = get_user_state(username)
+ ais = get_ai_cache_manager()
+
+ if ais is None:
+ return jsonify({"error": "AI support not enabled"}), 400
+
+ if not ais.option_highlighting_enabled:
+ return jsonify({"error": "Option highlighting not enabled"}), 400
+
+ # Get prefetch count from request or use default
+ prefetch_count = None
+ if request.is_json and request.json:
+ prefetch_count = request.json.get("count")
+
+ instance_id = user_state.get_current_instance_index()
+ ais.start_option_highlight_prefetch(instance_id, prefetch_count)
+
+ return jsonify({"status": "prefetch_started", "from_instance": instance_id})
+
+
+# Admin routes for system inspection (read-only)
+@app.route("/admin/health", methods=["GET"])
+def admin_health():
+ """
+ Health check endpoint for administrators.
+
+ Returns:
+ flask.Response: JSON response with server status
+ """
+ # Check API key
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({
+ "error": "Health check only available in debug mode or with valid API key"
+ }), 403
+
+ try:
+ # Check if core managers are accessible
+ usm = get_user_state_manager()
+ ism = get_item_state_manager()
+
+ return jsonify({
+ "status": "healthy",
+ "timestamp": str(datetime.datetime.now()),
+ "managers": {
+ "user_state_manager": "available",
+ "item_state_manager": "available"
+ },
+ "config": {
+ "debug_mode": config.get("debug", False),
+ "annotation_task_name": config.get("annotation_task_name", "Unknown")
+ }
+ })
+ except Exception as e:
+ return jsonify({
+ "status": "unhealthy",
+ "error": str(e),
+ "timestamp": str(datetime.datetime.now())
+ }), 500
+
+
+@app.route("/admin/system_state", methods=["GET"])
+def admin_system_state():
+ """
+ Get overall system state including user and item statistics.
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ flask.Response: JSON response with system state
+ """
+ # Check API key
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({
+ "error": "System state only available in debug mode or with valid API key"
+ }), 403
+
+ try:
+ usm = get_user_state_manager()
+ ism = get_item_state_manager()
+
+ # Get all users
+ users = get_users()
+ user_stats = {}
+ total_annotations = 0
+
+ for username in users:
+ user_state = usm.get_user_state(username)
+ if user_state:
+ user_annotations = len(user_state.get_all_annotations())
+ total_annotations += user_annotations
+ user_stats[username] = {
+ "phase": str(user_state.get_phase()),
+ "annotations_count": user_annotations,
+ "has_assignments": user_state.has_assignments(),
+ "remaining_assignments": user_state.has_remaining_assignments()
+ }
+
+ # Get item statistics
+ items = ism.items()
+ item_stats = {
+ "total_items": len(items),
+ "items_with_annotations": 0,
+ "items_by_annotator_count": {}
+ }
+
+ for item in items:
+ item_id = item.get_id()
+ annotators = ism.get_annotators_for_item(item_id)
+ if annotators:
+ item_stats["items_with_annotations"] += 1
+ annotator_count = len(annotators)
+ item_stats["items_by_annotator_count"][annotator_count] = item_stats["items_by_annotator_count"].get(annotator_count, 0) + 1
+
+ return jsonify({
+ "system_state": {
+ "total_users": len(users),
+ "total_items": item_stats["total_items"],
+ "total_annotations": total_annotations,
+ "items_with_annotations": item_stats["items_with_annotations"],
+ "items_by_annotator_count": item_stats["items_by_annotator_count"]
+ },
+ "users": user_stats,
+ "config": {
+ "debug_mode": config.get("debug", False),
+ "annotation_task_name": config.get("annotation_task_name", "Unknown"),
+ "max_annotations_per_user": config.get("max_annotations_per_user", "Unlimited"),
+ "annotation_schemes": config.get("annotation_schemes", []),
+ "ui": config.get("ui", {})
+ }
+ })
+ except Exception as e:
+ return jsonify({
+ "error": f"Failed to get system state: {str(e)}"
+ }), 500
+
+
+@app.route("/admin/all_instances", methods=["GET"])
+def admin_all_instances():
+ """
+ Get all available instances for navigation purposes.
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ flask.Response: JSON response with all instances
+ """
+ # Check API key
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({
+ "error": "All instances only available in debug mode or with valid API key"
+ }), 403
+
+ try:
+ ism = get_item_state_manager()
+ items = ism.items()
+
+ all_instances = []
+ for item in items:
+ all_instances.append({
+ "id": item.get_id(),
+ "text": item.get_text(),
+ "displayed_text": item.get_displayed_text()
+ })
+
+ return jsonify({
+ "total_items": len(all_instances),
+ "items": all_instances
+ })
+ except Exception as e:
+ return jsonify({
+ "error": f"Failed to get all instances: {str(e)}"
+ }), 500
+
+
+
+
+
+@app.route("/admin/user_state/", methods=["GET"])
+def admin_user_state(user_id):
+ """
+ Get detailed state for a specific user.
+ Admin-only endpoint requiring API key.
+
+ Args:
+ user_id: The user ID to get state for
+
+ Returns:
+ flask.Response: JSON response with user state
+ """
+ logger.debug("=== ADMIN USER STATE ROUTE START ===")
+ logger.debug(f"Requested user_id: {user_id}")
+ logger.debug(f"Request headers: {dict(request.headers)}")
+ logger.debug(f"Debug mode: {config.get('debug', False)}")
+
+ # Check API key
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ logger.warning("Access denied to admin endpoint - invalid API key")
+ return jsonify({
+ "error": "User state only available in debug mode or with valid API key"
+ }), 403
+ try:
+ logger.debug(f"Getting user state manager")
+ usm = get_user_state_manager()
+ logger.debug(f"All users in state manager: {usm.get_user_ids()}")
+ logger.debug(f"Looking for user: {user_id}")
+ logger.debug(f"User exists: {usm.has_user(user_id)}")
+
+ user_state = usm.get_user_state(user_id)
+ logger.debug(f"Retrieved user state: {user_state}")
+
+ if not user_state:
+ logger.warning(f"User '{user_id}' not found in state manager")
+ return jsonify({
+ "error": f"User '{user_id}' not found"
+ }), 404
+
+ # Get current instance
+ current_instance = user_state.get_current_instance()
+ current_instance_data = None
+ if current_instance:
+ # Get the base text
+ base_text = current_instance.get_text()
+
+ # Get span annotations for this instance and user
+ span_annotations = get_span_annotations_for_user_on(user_id, current_instance.get_id())
+
+ # Render the text with span annotations
+ from potato.server_utils.schemas.span import render_span_annotations
+ displayed_text = render_span_annotations(base_text, span_annotations)
+
+
+ current_instance_data = {
+ "id": current_instance.get_id(),
+ "text": base_text,
+ "displayed_text": displayed_text
+ }
+
+ # Helper to recursively convert all dict keys to strings
+ def stringify_keys(obj):
+ if isinstance(obj, dict):
+ return {str(k): stringify_keys(v) for k, v in obj.items()}
+ elif isinstance(obj, list):
+ return [stringify_keys(i) for i in obj]
+ else:
+ return obj
+
+ # Get all annotations
+ all_annotations = user_state.get_all_annotations()
+
+ # Convert all keys to strings for JSON serialization
+ serializable_annotations = {}
+ for instance_id, annotations in all_annotations.items():
+ instance_id_str = str(instance_id)
+ serializable_annotations[instance_id_str] = {}
+
+ # Process labels
+ if "labels" in annotations:
+ for label, value in annotations["labels"].items():
+ if hasattr(label, 'schema_name') and hasattr(label, 'label_name'):
+ label_str = f"{label.schema_name}:{label.label_name}"
+ else:
+ label_str = str(label)
+ serializable_annotations[instance_id_str][label_str] = value
+
+ # Process spans
+ if "spans" in annotations:
+ for span, value in annotations["spans"].items():
+ span_str = str(span)
+ serializable_annotations[instance_id_str][span_str] = value
+
+ serializable_annotations = stringify_keys(serializable_annotations)
+
+ # Get assignments
+ assignments = []
+ if user_state.has_assignments():
+ for instance_id in user_state.get_assigned_instance_ids():
+ instance = get_item_state_manager().get_item(instance_id)
+ if instance:
+ assignments.append({
+ "id": instance.get_id(),
+ "text": instance.get_text(),
+ "displayed_text": instance.get_displayed_text(),
+ "has_annotation": instance_id in all_annotations
+ })
+
+ return jsonify({
+ "user_id": user_id,
+ "phase": str(user_state.get_phase()),
+ "current_instance": current_instance_data,
+ "max_assignments": user_state.get_max_assignments(),
+ "assignments": {
+ "total": len(assignments),
+ "annotated": len([a for a in assignments if a["has_annotation"]]),
+ "remaining": len([a for a in assignments if not a["has_annotation"]]),
+ "items": assignments
+ },
+ "annotations": {
+ "total_count": len(all_annotations),
+ "by_instance": serializable_annotations
+ },
+ "hints": {
+ "cached_hints": list(user_state.get_cached_hints().keys()) if hasattr(user_state, 'get_cached_hints') else []
+ }
+ })
+ except Exception as e:
+ return jsonify({
+ "error": f"Failed to get user state for '{user_id}': {str(e)}"
+ }), 500
+
+
+@app.route("/admin/item_state", methods=["GET"])
+def admin_item_state():
+ """
+ Get state for all items in the system.
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ flask.Response: JSON response with item state
+ """
+ # Check API key
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({
+ "error": "Item state only available in debug mode or with valid API key"
+ }), 403
+
+ try:
+ ism = get_item_state_manager()
+ items = ism.items()
+
+ item_states = []
+ for item in items:
+ item_id = item.get_id()
+ annotators = ism.get_annotators_for_item(item_id)
+
+ item_states.append({
+ "id": item_id,
+ "text": item.get_text(),
+ "displayed_text": item.get_displayed_text(),
+ "annotators": list(annotators) if annotators else [],
+ "annotation_count": len(annotators) if annotators else 0
+ })
+
+ # Sort by annotation count for easier analysis
+ item_states.sort(key=lambda x: x["annotation_count"], reverse=True)
+
+ return jsonify({
+ "total_items": len(item_states),
+ "items": item_states,
+ "summary": {
+ "items_with_annotations": len([i for i in item_states if i["annotation_count"] > 0]),
+ "items_without_annotations": len([i for i in item_states if i["annotation_count"] == 0]),
+ "average_annotations_per_item": sum(i["annotation_count"] for i in item_states) / len(item_states) if item_states else 0
+ }
+ })
+ except Exception as e:
+ return jsonify({
+ "error": f"Failed to get item state: {str(e)}"
+ }), 500
+
+
+@app.route("/admin/item_state/", methods=["GET"])
+def admin_item_state_detail(item_id):
+ """
+ Get detailed state for a specific item.
+ Admin-only endpoint requiring API key.
+
+ Args:
+ item_id: The item ID to get state for
+
+ Returns:
+ flask.Response: JSON response with item state
+ """
+ # Check API key
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({
+ "error": "Item state detail only available in debug mode or with valid API key"
+ }), 403
+ try:
+ ism = get_item_state_manager()
+ item = ism.get_item(item_id)
+
+ if not item:
+ return jsonify({
+ "error": f"Item '{item_id}' not found"
+ }), 404
+
+ annotators = ism.get_annotators_for_item(item_id)
+
+ # Get annotations from all users for this item
+ usm = get_user_state_manager()
+ item_annotations = {}
+
+ for username in get_users():
+ user_state = usm.get_user_state(username)
+ if user_state:
+ user_annotations = user_state.get_all_annotations()
+ if item_id in user_annotations:
+ item_annotations[username] = user_annotations[item_id]
+
+ return jsonify({
+ "item_id": item_id,
+ "text": item.get_text(),
+ "displayed_text": item.get_displayed_text(),
+ "annotators": list(annotators) if annotators else [],
+ "annotation_count": len(annotators) if annotators else 0,
+ "annotations": item_annotations
+ })
+ except Exception as e:
+ return jsonify({
+ "error": f"Failed to get item state for '{item_id}': {str(e)}"
+ }), 500
+
+
+# Test Support Endpoints (only available in debug mode)
+
+@app.route("/admin/api/test/reset_state", methods=["POST"])
+def admin_api_test_reset_state():
+ """
+ Reset server state for testing purposes.
+ Only available in debug mode.
+
+ This endpoint clears user state and reloads data, allowing tests
+ to start with a fresh state without restarting the server.
+
+ Returns:
+ flask.Response: JSON response with reset status
+ """
+ if not config.get('debug', False):
+ return jsonify({'error': 'This endpoint is only available in debug mode'}), 403
+
+ try:
+ from potato.user_state_management import clear_user_state_manager, init_user_state_manager
+ from potato.item_state_management import clear_item_state_manager, init_item_state_manager
+ from potato.flask_server import load_all_data # noqa: cross-import
+ from potato.authentication import UserAuthenticator
+
+ # Clear existing state
+ clear_user_state_manager()
+ clear_item_state_manager()
+
+ # Clear ICL labeler if it exists
+ try:
+ from potato.ai.icl_labeler import clear_icl_labeler
+ clear_icl_labeler()
+ except ImportError:
+ pass
+
+ # Reinitialize state
+ UserAuthenticator.init_from_config(config)
+ init_user_state_manager(config)
+ init_item_state_manager(config)
+ load_all_data(config)
+
+ logger.info("Server state reset successfully for testing")
+ return jsonify({
+ 'status': 'success',
+ 'message': 'Server state reset successfully'
+ })
+
+ except Exception as e:
+ logger.error(f"Failed to reset server state: {e}")
+ return jsonify({
+ 'status': 'error',
+ 'message': f'Failed to reset state: {str(e)}'
+ }), 500
+
+
+# New Admin Dashboard API Endpoints
+
+@app.route("/admin/api/overview", methods=["GET"])
+def admin_api_overview():
+ """
+ Get dashboard overview data.
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ flask.Response: JSON response with overview statistics
+ """
+ result = admin_dashboard.get_dashboard_overview()
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/api/annotators", methods=["GET"])
+def admin_api_annotators():
+ """
+ Get detailed annotator data including timing information.
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ flask.Response: JSON response with annotator data
+ """
+ result = admin_dashboard.get_annotators_data()
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/api/instances", methods=["GET"])
+def admin_api_instances():
+ """
+ Get paginated instances data with sorting and filtering.
+ Admin-only endpoint requiring API key.
+
+ Query Parameters:
+ page: Page number (default: 1)
+ page_size: Items per page (default: 25)
+ sort_by: Sort field (annotation_count, completion_percentage, disagreement, id, average_time)
+ sort_order: Sort order (asc, desc)
+ filter_completion: Filter by completion (completed, incomplete, all)
+
+ Returns:
+ flask.Response: JSON response with paginated instances data
+ """
+ # Get query parameters
+ page = int(request.args.get('page', 1))
+ page_size = int(request.args.get('page_size', 25))
+ sort_by = request.args.get('sort_by', 'annotation_count')
+ sort_order = request.args.get('sort_order', 'desc')
+ filter_completion = request.args.get('filter_completion')
+
+ result = admin_dashboard.get_instances_data(
+ page=page,
+ page_size=page_size,
+ sort_by=sort_by,
+ sort_order=sort_order,
+ filter_completion=filter_completion
+ )
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/api/config", methods=["GET", "POST"])
+def admin_api_config():
+ """
+ Get or update system configuration.
+ Admin-only endpoint requiring API key.
+
+ GET: Returns current configuration
+ POST: Updates configuration with provided data
+
+ Returns:
+ flask.Response: JSON response with configuration data or update result
+ """
+ if request.method == "GET":
+ # Return current configuration
+ response_data = {
+ "max_annotations_per_user": config.get("max_annotations_per_user", -1),
+ "max_annotations_per_item": config.get("max_annotations_per_item", -1),
+ "assignment_strategy": config.get("assignment_strategy", "fixed_order"),
+ "annotation_task_name": config.get("annotation_task_name", "Unknown"),
+ "debug_mode": config.get("debug", False)
+ }
+
+ # Add training configuration if present
+ if "training" in config:
+ response_data["training"] = config["training"]
+
+ return jsonify(response_data)
+
+ elif request.method == "POST":
+ # Update configuration
+ config_updates = request.get_json()
+ if not config_updates:
+ return jsonify({"error": "No configuration updates provided"}), 400
+
+ result = admin_dashboard.update_config(config_updates)
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/api/user//set_instances", methods=["POST"])
+def admin_api_set_user_instances(username):
+ """
+ Set the maximum number of instances for a specific user.
+ Admin-only endpoint requiring API key.
+
+ JSON Body:
+ max_instances (int): New max instance count. Use -1 for unlimited.
+
+ Returns:
+ flask.Response: JSON response with updated user info
+ """
+ data = request.get_json()
+ if not data or 'max_instances' not in data:
+ return jsonify({"error": "max_instances is required"}), 400
+
+ max_instances = int(data['max_instances'])
+
+ usm = get_user_state_manager()
+ user_state = usm.get_user_state(username)
+ if not user_state:
+ return jsonify({"error": f"User '{username}' not found"}), 404
+
+ # Don't allow setting below current annotation count (can't un-annotate)
+ current_count = user_state.get_annotation_count()
+ if max_instances >= 0 and max_instances < current_count:
+ max_instances = current_count
+
+ user_state.set_max_assignments(max_instances)
+
+ return jsonify({
+ "success": True,
+ "username": username,
+ "max_instances": max_instances,
+ "current_annotations": current_count
+ })
+
+
+@app.route("/admin/api/stale_assignments", methods=["GET"])
+def admin_api_stale_assignments():
+ """
+ Get a list of stale instance assignments (assigned but not annotated past timeout).
+
+ Returns:
+ flask.Response: JSON with stale assignment info
+ """
+ import time
+ ism = get_item_state_manager()
+ usm = get_user_state_manager()
+ timeout_hours = ism.reclaim_timeout_hours
+ cutoff = time.time() - (timeout_hours * 3600)
+ stale = []
+
+ for iid, user_timestamps in ism.assignment_timestamps.items():
+ for username, timestamp in user_timestamps.items():
+ user_state = usm.get_user_state(username)
+ if user_state and user_state.has_annotated(iid):
+ continue
+ hours_ago = (time.time() - timestamp) / 3600
+ stale.append({
+ 'instance_id': iid,
+ 'username': username,
+ 'assigned_hours_ago': round(hours_ago, 1),
+ 'is_stale': timestamp < cutoff
+ })
+
+ stale.sort(key=lambda x: x['assigned_hours_ago'], reverse=True)
+ return jsonify({
+ 'stale_assignments': stale,
+ 'timeout_hours': timeout_hours,
+ 'reclaim_enabled': ism.reclaim_enabled
+ })
+
+
+@app.route("/admin/api/reclaim_instance", methods=["POST"])
+def admin_api_reclaim_instance():
+ """
+ Manually reclaim a specific instance assignment from a user.
+
+ JSON Body:
+ instance_id (str): The instance to reclaim
+ username (str): The user to reclaim from
+ """
+ data = request.get_json()
+ if not data or 'instance_id' not in data or 'username' not in data:
+ return jsonify({"error": "instance_id and username are required"}), 400
+
+ iid = data['instance_id']
+ username = data['username']
+
+ ism = get_item_state_manager()
+ usm = get_user_state_manager()
+ user_state = usm.get_user_state(username)
+
+ if not user_state:
+ return jsonify({"error": f"User '{username}' not found"}), 404
+
+ if user_state.has_annotated(iid):
+ return jsonify({"error": "Cannot reclaim: user has already annotated this instance"}), 400
+
+ reclaimed = ism._reclaim_unannotated_assignment(
+ user_state,
+ iid,
+ reason="manual_admin_reclaim",
+ )
+ if not reclaimed:
+ return jsonify({"error": "Instance is not assigned to this user"}), 400
+
+ get_user_state_manager().save_user_state(user_state)
+
+ return jsonify({"success": True, "instance_id": iid, "username": username})
+
+
+@app.route("/admin/api/questions", methods=["GET"])
+def admin_api_questions():
+ """
+ Get aggregate analysis data for all annotation schemas/questions.
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ flask.Response: JSON response with questions data and visualizations
+ """
+ result = admin_dashboard.get_questions_data()
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/api/annotation_history", methods=["GET"])
+def admin_api_annotation_history():
+ """
+ Get detailed annotation history data with filtering options.
+ Admin-only endpoint requiring API key.
+
+ Query Parameters:
+ user_id: Optional user ID to filter by
+ instance_id: Optional instance ID to filter by
+ minutes: Optional time window in minutes
+
+ Returns:
+ flask.Response: JSON response with annotation history data
+ """
+ user_id = request.args.get('user_id')
+ instance_id = request.args.get('instance_id')
+ minutes = request.args.get('minutes')
+
+ if minutes:
+ try:
+ minutes = int(minutes)
+ except ValueError:
+ return jsonify({"error": "Invalid minutes parameter"}), 400
+
+ result = admin_dashboard.get_annotation_history_data(
+ user_id=user_id,
+ instance_id=instance_id,
+ minutes=minutes
+ )
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/api/suspicious_activity", methods=["GET"])
+def admin_api_suspicious_activity():
+ """
+ Get comprehensive suspicious activity analysis.
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ flask.Response: JSON response with suspicious activity data
+ """
+ result = admin_dashboard.get_suspicious_activity_data()
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/api/crowdsourcing", methods=["GET"])
+def admin_api_crowdsourcing():
+ """
+ Get crowdsourcing platform statistics (MTurk, Prolific).
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ flask.Response: JSON response with crowdsourcing data including:
+ - summary: Overall worker counts by platform
+ - prolific: Prolific-specific statistics and worker list
+ - mturk: MTurk-specific statistics and worker list
+ - other: Non-crowdsourcing workers
+ """
+ result = admin_dashboard.get_crowdsourcing_data()
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/api/agreement", methods=["GET"])
+def admin_api_agreement():
+ """
+ Get inter-annotator agreement metrics.
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ flask.Response: JSON response with agreement metrics including:
+ - overall: Average Krippendorff's alpha across schemas
+ - by_schema: Per-schema agreement metrics
+ - interpretation: Human-readable interpretation
+ """
+ result = admin_dashboard.get_agreement_metrics()
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/iaa", methods=["GET"])
+def admin_iaa():
+ """
+ Inter-annotator agreement over the overlap-sample items.
+
+ Computes a schema-appropriate set of metrics (nominal: Cohen/Fleiss kappa,
+ Krippendorff alpha; ordinal: weighted kappa, Spearman rho; continuous:
+ Pearson, ICC; span: token kappa, span F1, alpha_U, gamma; etc.).
+ Restricted to items whose per-item annotator cap is >= 2 and that have
+ reached that cap.
+
+ Pass ``?format=html`` for a rendered table (default: JSON).
+ """
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({"error": "Admin API key required"}), 403
+
+ try:
+ from potato.server_utils.iaa import compute_overlap_iaa
+ report = compute_overlap_iaa(
+ get_item_state_manager(), get_user_state_manager(), config,
+ )
+ except Exception as exc:
+ logger.exception("Failed to compute overlap IAA")
+ return jsonify({"error": str(exc)}), 500
+
+ if (request.args.get("format") or "json").lower() == "html":
+ try:
+ return render_template("admin/iaa.html", report=report)
+ except Exception:
+ # Template optional; fall back to JSON if missing.
+ pass
+ return jsonify(report)
+
+
+def _record_judge_comparison_if_enabled(username, instance_id):
+ """On human save, log the humanโjudge comparison if inline mode is on."""
+ ja = config.get("judge_alignment", {}) or {}
+ if not (ja.get("inline", {}) or {}).get("enabled"):
+ return
+ try:
+ from potato.server_utils import judge_alignment as ja_mod
+ schemas = ja_mod.judge_scoped_schemas(config)
+ allow = set((ja.get("inline", {}) or {}).get("schemas", []) or [])
+ if allow:
+ schemas = [s for s in schemas if s.get("name") in allow]
+ preds = ja_mod.load_predictions(config)
+ version = ja_mod.latest_prompt_version(config)
+ version_preds = preds.get(version, {}) if version else {}
+ for schema_info in schemas:
+ schema_name = schema_info.get("name")
+ pred = version_preds.get(f"{instance_id}::{schema_name}")
+ if not pred:
+ continue
+ human_label = ja_mod.human_label_for(instance_id, schema_name, username)
+ if human_label is None:
+ continue
+ ja_mod.record_comparison(
+ config, instance_id, schema_name, human_label,
+ pred.get("predicted_label"), version or "",
+ )
+ except Exception as e:
+ logger.warning(f"Judge comparison recording failed for {instance_id}: {e}")
+
+
+@app.route("/admin/judge-alignment", methods=["GET"])
+def admin_judge_alignment():
+ """LLM-judge โ human alignment report: Cohen's ฮบ, confusion, disagreements.
+
+ Mirrors /admin/iaa. Pass ?format=html for the rendered page (default JSON),
+ and ?prompt_version= to view a specific judge prompt version.
+ """
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({"error": "Admin API key required"}), 403
+
+ try:
+ from potato.server_utils.judge_alignment import compute_judge_alignment
+ report = compute_judge_alignment(
+ config, get_users(), prompt_version=request.args.get("prompt_version"),
+ )
+ except Exception as exc:
+ logger.exception("Failed to compute judge alignment")
+ return jsonify({"error": str(exc)}), 500
+
+ if (request.args.get("format") or "json").lower() == "html":
+ try:
+ return render_template("admin/judge_alignment.html", report=report)
+ except Exception:
+ pass
+ return jsonify(report)
+
+
+@app.route("/admin/api/judge-alignment/run", methods=["POST"])
+def admin_judge_alignment_run():
+ """Run/re-run the judge over human-annotated instances.
+
+ Optional JSON body: {"rubrics": {schema_name: rubric}, "max_per_schema": N}.
+ Editing a rubric creates a new prompt version so ฮบ can be tracked across
+ calibration rounds. Returns a run summary.
+ """
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({"error": "Admin API key required"}), 403
+
+ body = request.get_json(silent=True) or {}
+ try:
+ from potato.server_utils.judge_alignment import run_judge_batch
+ summary = run_judge_batch(
+ config, get_users(),
+ rubric_overrides=body.get("rubrics"),
+ max_per_schema=body.get("max_per_schema"),
+ )
+ except Exception as exc:
+ logger.exception("Failed to run judge batch")
+ return jsonify({"error": str(exc)}), 500
+ return jsonify(summary)
+
+
+@app.route("/admin/triage-queue", methods=["GET"])
+def admin_triage_queue():
+ """Signal-based triage queue: items ranked by their quality signal.
+
+ Shows the remaining (incomplete) items ordered by the triage priority that
+ was assigned at load/ingestion time (errors / thumbs-down / low score
+ first), with the reason that flagged each one. Pass ?format=html for the
+ rendered page (default JSON).
+ """
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({"error": "Admin API key required"}), 403
+
+ try:
+ from potato.server_utils.triage import compute_triage_queue
+ report = compute_triage_queue(config)
+ except Exception as exc:
+ logger.exception("Failed to compute triage queue")
+ return jsonify({"error": str(exc)}), 500
+
+ if (request.args.get("format") or "json").lower() == "html":
+ try:
+ return render_template("admin/triage_queue.html", report=report)
+ except Exception:
+ pass
+ return jsonify(report)
+
+
+@app.route("/admin/api/code_cooccurrence", methods=["GET"])
+def admin_api_code_cooccurrence():
+ """
+ Get pairwise code co-occurrence across instances.
+
+ Query params:
+ schema (optional): Restrict to a single schema name.
+ min_count (optional): Minimum co-occurrence count to include (default 1).
+
+ Returns:
+ JSON with codes, pairs ({code_a, code_b, count}), n_instances.
+ """
+ schema = request.args.get("schema")
+ try:
+ min_count = int(request.args.get("min_count", "1"))
+ except ValueError:
+ min_count = 1
+ result = admin_dashboard.get_code_cooccurrence_matrix(
+ schema_filter=schema, min_count=min_count
+ )
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/api/code_crosstab", methods=["GET"])
+def admin_api_code_crosstab():
+ """
+ Get a codes-by-instance-attribute crosstab.
+
+ Query params:
+ attribute (required): The instance metadata field to pivot on
+ (e.g. "site", "condition", "language").
+ schema (optional): Restrict to a single schema name.
+
+ Returns:
+ JSON with codes, values, cells ({code, value, count}), n_instances.
+ """
+ attribute = request.args.get("attribute")
+ if not attribute:
+ return jsonify({"error": "attribute query param is required"}), 400
+ schema = request.args.get("schema")
+ result = admin_dashboard.get_code_crosstab(
+ attribute_key=attribute, schema_filter=schema
+ )
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/api/step_agreement", methods=["GET"])
+def admin_api_step_agreement():
+ """
+ Get step-level inter-annotator agreement metrics for agent traces.
+ Admin-only endpoint.
+
+ Query params:
+ scheme: Annotation scheme name (required)
+ metric: "krippendorff_alpha" or "cohens_kappa" (default: krippendorff_alpha)
+
+ Returns:
+ JSON with overall, per_step, and per_instance agreement.
+ """
+ try:
+ from potato.step_agreement import compute_step_agreement
+ from potato.item_state_management import get_item_state_manager
+
+ scheme_name = request.args.get("scheme", "")
+ metric = request.args.get("metric", "krippendorff_alpha")
+
+ if not scheme_name:
+ return jsonify({"error": "scheme parameter is required"}), 400
+
+ ism = get_item_state_manager()
+ # Collect step-level annotations from all instances
+ annotations = {}
+ for instance_id, item in ism.instance_id_to_instance.items():
+ annotator_data = {}
+ for annotator_id, ann in item.get_annotations().items():
+ if scheme_name in ann:
+ annotator_data[annotator_id] = ann
+ if annotator_data:
+ annotations[instance_id] = annotator_data
+
+ if not annotations:
+ return jsonify({
+ "error": "No step-level annotations found",
+ "scheme": scheme_name,
+ }), 404
+
+ result = compute_step_agreement(
+ annotations, scheme_name=scheme_name, metric=metric
+ )
+ return jsonify(result)
+
+ except ImportError as e:
+ return jsonify({"error": f"Missing dependency: {e}"}), 500
+ except Exception as e:
+ return jsonify({"error": str(e)}), 500
+
+
+@app.route("/admin/api/step_quality", methods=["GET"])
+def admin_api_step_quality():
+ """
+ Get step-level quality control metrics.
+ Admin-only endpoint.
+
+ Returns:
+ JSON with gold standard results and attention check stats.
+ """
+ try:
+ step_qc_config = config.get("quality_control", {}).get("step_level", {})
+ if not step_qc_config.get("enabled", False):
+ return jsonify({"enabled": False, "message": "Step-level QC not configured"})
+
+ from potato.step_quality_control import StepQualityControlManager
+ task_dir = config.get("task_dir", ".")
+ manager = StepQualityControlManager(step_qc_config, task_dir)
+ return jsonify(manager.get_quality_summary())
+
+ except Exception as e:
+ return jsonify({"error": str(e)}), 500
+
+
+@app.route("/admin/api/quality_control", methods=["GET"])
+def admin_api_quality_control():
+ """
+ Get quality control metrics (attention checks, gold standards).
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ flask.Response: JSON response with quality control data including:
+ - attention_checks: Statistics on attention check pass/fail rates
+ - gold_standards: Statistics on gold standard accuracy
+ - by_user: Per-user quality metrics
+ """
+ result = admin_dashboard.get_quality_control_data()
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+@app.route("/admin/api/behavioral_analytics", methods=["GET"])
+def admin_api_behavioral_analytics():
+ """
+ Get behavioral analytics data for all annotators.
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ flask.Response: JSON response with behavioral analytics including:
+ - aggregate_stats: Overall statistics (total users, instances, avg time)
+ - ai_usage: AI assistance usage statistics
+ - quality_summary: Quality indicators and suspicious activity flags
+ - interaction_types: Breakdown of interaction types
+ - change_sources: Sources of annotation changes
+ - users: Per-user behavioral metrics sorted by suspicion score
+ """
+ result = admin_dashboard.get_behavioral_analytics_data()
+ if isinstance(result, tuple):
+ return jsonify(result[0]), result[1]
+ return jsonify(result)
+
+
+# === ICL Verification Helper ===
+
+def _maybe_record_icl_verification(user_state, instance_id: str, annotations: dict) -> bool:
+ """
+ Check if an annotation was for an ICL verification task and record the result.
+
+ This is called after a user submits an annotation. If the instance was assigned
+ as a verification task (blind labeling), we compare their label to the LLM's
+ prediction and record the verification result.
+
+ Args:
+ user_state: The user's state object
+ instance_id: The annotated instance ID
+ annotations: The user's annotation data
+
+ Returns:
+ True if verification was recorded, False otherwise
+ """
+ # Check if this instance is a verification task
+ if not hasattr(user_state, 'is_verification_task'):
+ return False
+
+ if not user_state.is_verification_task(instance_id):
+ return False
+
+ try:
+ from potato.ai.icl_labeler import get_icl_labeler
+
+ icl_labeler = get_icl_labeler()
+ if icl_labeler is None:
+ return False
+
+ # Get the schema being verified
+ schema_name = user_state.get_verification_schema(instance_id)
+ if not schema_name:
+ return False
+
+ # Extract the human's label for this schema
+ human_label = None
+ if schema_name in annotations:
+ label_data = annotations[schema_name]
+ if isinstance(label_data, dict):
+ # For radio/multiselect, find the selected value
+ for label_name, value in label_data.items():
+ if value == 'true' or value is True:
+ human_label = label_name
+ break
+ elif isinstance(value, str) and value not in ('false', ''):
+ human_label = value
+ break
+ elif isinstance(label_data, str):
+ human_label = label_data
+
+ if human_label is None:
+ logger.warning(f"Could not extract human label for verification of {instance_id}")
+ return False
+
+ # Record the verification
+ user_id = user_state.get_user_id()
+ success = icl_labeler.record_verification(
+ instance_id=instance_id,
+ schema_name=schema_name,
+ human_label=human_label,
+ verified_by=user_id
+ )
+
+ if success:
+ # Remove from user's verification task tracking
+ user_state.complete_verification_task(instance_id)
+ logger.info(f"Recorded ICL verification for {instance_id} by {user_id}: {human_label}")
+
+ return success
+
+ except ImportError:
+ # ICL labeler module not available
+ return False
+ except Exception as e:
+ logger.warning(f"Error recording ICL verification: {e}")
+ return False
+
+
+# === Diversity Manager Helper ===
+
+def _notify_diversity_manager_annotation(user_state, instance_id: str) -> None:
+ """
+ Notify the diversity manager that an annotation was completed.
+
+ This triggers async embedding computation if needed and checks if
+ reclustering should occur. Also triggers AI prefetch after reordering.
+
+ Args:
+ user_state: The user's state object
+ instance_id: The annotated instance ID
+ """
+ dm = get_diversity_manager()
+ if not dm or not dm.enabled:
+ return
+
+ try:
+ user_id = getattr(user_state, 'user_id', 'anonymous')
+
+ # Get the item text for embedding
+ ism = get_item_state_manager()
+ item = ism.get_item(instance_id)
+ if not item:
+ return
+
+ text_key = config.get("item_properties", {}).get("text_key", "text")
+ item_data = item.get_data()
+ text = item_data.get(text_key, item.get_text())
+
+ # Notify diversity manager (triggers async embedding if needed)
+ dm.on_annotation_complete(user_id, instance_id, text)
+
+ # Check if reclustering is needed
+ if dm.should_recluster(user_id):
+ dm.trigger_recluster(user_id)
+
+ # Trigger AI prefetch for reordered items
+ if dm.config.trigger_ai_prefetch:
+ _trigger_ai_prefetch_after_reorder(user_state)
+
+ except Exception as e:
+ logger.warning(f"Error notifying diversity manager: {e}")
+
+
+def _trigger_ai_prefetch_after_reorder(user_state) -> None:
+ """
+ Trigger AI cache prefetch after diversity reordering.
+
+ Args:
+ user_state: The user's state object
+ """
+ try:
+ acm = get_ai_cache_manager()
+ if acm is None:
+ return
+
+ # Get current instance index
+ current_index = getattr(user_state, 'current_instance_index', 0)
+
+ # Prefetch count from AI cache config
+ prefetch_count = getattr(acm, 'prefetch_page_count_on_next', 5)
+
+ acm.start_prefetch(current_index, prefetch_count)
+ logger.debug(f"Triggered AI prefetch after diversity reorder from index {current_index}")
+
+ except Exception as e:
+ logger.warning(f"Error triggering AI prefetch after reorder: {e}")
+
+
+# === ICL Labeling Admin API ===
+
+@app.route("/admin/api/icl/status", methods=["GET"])
+def admin_api_icl_status():
+ """
+ Get ICL labeler status and statistics.
+ Admin-only endpoint.
+
+ Returns:
+ JSON with ICL labeler status including:
+ - enabled: Whether ICL labeling is enabled
+ - total_examples: Number of high-confidence examples
+ - total_predictions: Number of LLM predictions made
+ - accuracy_metrics: Verification-based accuracy
+ - labeling_paused: Whether labeling is currently paused
+ """
+ try:
+ from potato.ai.icl_labeler import get_icl_labeler
+ icl_labeler = get_icl_labeler()
+
+ if icl_labeler is None:
+ return jsonify({
+ 'enabled': False,
+ 'message': 'ICL labeling not initialized'
+ })
+
+ return jsonify(icl_labeler.get_status())
+
+ except Exception as e:
+ logger.error(f"Error getting ICL status: {e}")
+ return jsonify({'error': str(e)}), 500
+
+
+@app.route("/admin/api/icl/examples", methods=["GET"])
+def admin_api_icl_examples():
+ """
+ Get current high-confidence examples.
+
+ Query params:
+ schema: Optional schema name to filter by
+
+ Returns:
+ JSON with examples grouped by schema
+ """
+ try:
+ from potato.ai.icl_labeler import get_icl_labeler
+ icl_labeler = get_icl_labeler()
+
+ if icl_labeler is None:
+ # Return empty results when ICL is not initialized (graceful degradation)
+ schema_filter = request.args.get('schema')
+ return jsonify({
+ 'examples': {},
+ 'total_count': 0,
+ 'schema': schema_filter
+ })
+
+ schema_filter = request.args.get('schema')
+
+ examples = {}
+ for schema_name, schema_examples in icl_labeler.schema_to_examples.items():
+ if schema_filter and schema_name != schema_filter:
+ continue
+ examples[schema_name] = [ex.to_dict() for ex in schema_examples]
+
+ response_data = {
+ 'examples': examples,
+ 'total_count': sum(len(ex) for ex in examples.values())
+ }
+ if schema_filter:
+ response_data['schema'] = schema_filter
+ return jsonify(response_data)
+
+ except Exception as e:
+ logger.error(f"Error getting ICL examples: {e}")
+ return jsonify({'error': str(e)}), 500
+
+
+@app.route("/admin/api/icl/predictions", methods=["GET"])
+def admin_api_icl_predictions():
+ """
+ Get LLM predictions with filtering.
+
+ Query params:
+ schema: Optional schema name to filter by
+ status: Optional verification status filter
+ limit: Maximum number of predictions to return (default 100)
+
+ Returns:
+ JSON with predictions list
+ """
+ try:
+ from potato.ai.icl_labeler import get_icl_labeler
+ icl_labeler = get_icl_labeler()
+
+ if icl_labeler is None:
+ # Return empty results when ICL is not initialized (graceful degradation)
+ return jsonify({
+ 'predictions': [],
+ 'total_count': 0
+ })
+
+ schema_filter = request.args.get('schema')
+ status_filter = request.args.get('status')
+ limit = int(request.args.get('limit', 100))
+
+ predictions_list = []
+ for inst_id, schemas in icl_labeler.predictions.items():
+ for schema_name, prediction in schemas.items():
+ if schema_filter and schema_name != schema_filter:
+ continue
+ if status_filter and prediction.verification_status != status_filter:
+ continue
+
+ predictions_list.append(prediction.to_dict())
+
+ if len(predictions_list) >= limit:
+ break
+ if len(predictions_list) >= limit:
+ break
+
+ # Sort by timestamp descending
+ predictions_list.sort(key=lambda x: x['timestamp'], reverse=True)
+
+ return jsonify({
+ 'predictions': predictions_list[:limit],
+ 'total_count': len(predictions_list)
+ })
+
+ except Exception as e:
+ logger.error(f"Error getting ICL predictions: {e}")
+ return jsonify({'error': str(e)}), 500
+
+
+@app.route("/admin/api/icl/accuracy", methods=["GET"])
+def admin_api_icl_accuracy():
+ """
+ Get accuracy metrics.
+
+ Query params:
+ schema: Optional schema name to filter by
+
+ Returns:
+ JSON with accuracy metrics
+ """
+ try:
+ from potato.ai.icl_labeler import get_icl_labeler
+ icl_labeler = get_icl_labeler()
+
+ if icl_labeler is None:
+ # Return empty metrics when ICL is not initialized (graceful degradation)
+ schema_filter = request.args.get('schema')
+ return jsonify({
+ 'total_predictions': 0,
+ 'total_verified': 0,
+ 'verified_correct': 0,
+ 'verified_incorrect': 0,
+ 'pending_verification': 0,
+ 'accuracy': 0.0,
+ 'schema_name': schema_filter
+ })
+
+ schema_filter = request.args.get('schema')
+ metrics = icl_labeler.get_accuracy_metrics(schema_filter)
+
+ return jsonify(metrics)
+
+ except Exception as e:
+ logger.error(f"Error getting ICL accuracy: {e}")
+ return jsonify({'error': str(e)}), 500
+
+
+@app.route("/admin/api/icl/trigger", methods=["POST"])
+def admin_api_icl_trigger():
+ """
+ Manually trigger ICL operations.
+
+ JSON body:
+ action: "refresh_examples" | "batch_label" | "save_state"
+ schema: Optional schema name for batch_label
+
+ Returns:
+ JSON with operation result
+ """
+ try:
+ from potato.ai.icl_labeler import get_icl_labeler
+ icl_labeler = get_icl_labeler()
+
+ if icl_labeler is None:
+ return jsonify({'error': 'ICL labeling not initialized'}), 400
+
+ data = request.get_json() or {}
+ action = data.get('action', '')
+
+ # Support shorthand: if schema_name is provided without action, default to batch_label
+ if not action and data.get('schema_name'):
+ action = 'batch_label'
+ # Use schema_name as the schema for backwards compatibility
+ if 'schema' not in data:
+ data['schema'] = data['schema_name']
+
+ if action == 'refresh_examples':
+ examples = icl_labeler.refresh_high_confidence_examples()
+ return jsonify({
+ 'action': 'refresh_examples',
+ 'success': True,
+ 'example_counts': {k: len(v) for k, v in examples.items()}
+ })
+
+ elif action == 'batch_label':
+ schema = data.get('schema')
+ if not schema:
+ return jsonify({'error': 'schema required for batch_label'}), 400
+
+ predictions = icl_labeler.batch_label_instances(schema)
+ icl_labeler.save_state()
+
+ return jsonify({
+ 'action': 'batch_label',
+ 'success': True,
+ 'predictions_count': len(predictions),
+ 'schema': schema,
+ 'message': f'Labeled {len(predictions)} instances for schema {schema}'
+ })
+
+ elif action == 'save_state':
+ icl_labeler.save_state()
+ return jsonify({
+ 'action': 'save_state',
+ 'success': True
+ })
+
+ else:
+ return jsonify({'error': f'Unknown action: {action}'}), 400
+
+ except Exception as e:
+ logger.error(f"Error triggering ICL action: {e}")
+ return jsonify({'error': str(e)}), 500
+
+
+@app.route("/api/icl/record_verification", methods=["POST"])
+def api_icl_record_verification():
+ """
+ Record human verification of an LLM prediction.
+
+ This is called when an annotator completes labeling an instance
+ that was selected for verification.
+
+ JSON body:
+ instance_id: The instance ID
+ schema_name: The schema name
+ human_label: The human's label
+
+ Returns:
+ JSON with verification result
+ """
+ try:
+ if 'username' not in session:
+ return jsonify({'error': 'Not authenticated'}), 401
+
+ from potato.ai.icl_labeler import get_icl_labeler
+ icl_labeler = get_icl_labeler()
+
+ if icl_labeler is None:
+ return jsonify({'error': 'ICL labeling not initialized'}), 400
+
+ data = request.get_json() or {}
+ instance_id = data.get('instance_id')
+ schema_name = data.get('schema_name')
+ human_label = data.get('human_label')
+
+ if not all([instance_id, schema_name, human_label]):
+ return jsonify({'error': 'Missing required fields'}), 400
+
+ username = session['username']
+ success = icl_labeler.record_verification(
+ instance_id, schema_name, human_label, username
+ )
+
+ if success:
+ icl_labeler.save_state()
+ return jsonify({'success': True, 'message': 'Verification recorded'})
+ else:
+ return jsonify({'success': False, 'message': 'No prediction found to verify'})
+
+ except Exception as e:
+ logger.error("Error recording verification: %s", traceback.format_exc())
+ return jsonify({'error': 'An internal error occurred'}), 500
+
+
+########################################################################
+# Agent Chat Routes
+########################################################################
+
+def _get_agent_sandbox():
+ """Get the safety sandbox for agent interactions."""
+ from potato.agent_proxy import SafetySandbox
+ agent_config = config.get("agent_proxy", {})
+ return SafetySandbox(agent_config)
+
+
+def _get_or_create_agent_session(username, instance_id):
+ """Get an existing agent session or create a new one."""
+ from potato.agent_proxy import (
+ get_agent_session_manager, AgentProxyFactory
+ )
+ mgr = get_agent_session_manager()
+ session_obj = mgr.get_session(username, instance_id)
+ if session_obj:
+ return session_obj
+
+ # Create a new session
+ proxy = AgentProxyFactory.create(config)
+ item = get_item_state_manager().get_item(instance_id)
+ task_desc = ""
+ if item:
+ data = item.get_data()
+ # Look for task_description in item data
+ task_desc = data.get("task_description", data.get("text", ""))
+ if isinstance(task_desc, list):
+ task_desc = " ".join(str(t) for t in task_desc)
+
+ return mgr.create_session(username, instance_id, proxy, str(task_desc))
+
+
+@app.route("/agent_chat/send", methods=["POST"])
+def agent_chat_send():
+ """Send a message to the agent and get a response."""
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ username = session['username']
+ user_state = get_user_state(username)
+
+ if user_state.get_phase() != UserPhase.ANNOTATION:
+ return jsonify({"error": "Not in annotation phase"}), 400
+
+ data = request.get_json(silent=True) or {}
+ message = data.get("message", "").strip()
+ if not message:
+ return jsonify({"error": "Empty message"}), 400
+
+ item = user_state.get_current_instance()
+ instance_id = item.get_id()
+
+ try:
+ sandbox = _get_agent_sandbox()
+
+ # Get or create session
+ agent_session = _get_or_create_agent_session(username, instance_id)
+
+ if agent_session.finished:
+ return jsonify({"error": "Chat session already finished"}), 400
+
+ # Safety checks
+ sandbox.check_step_limit(agent_session.step_count)
+ sandbox.check_session_timeout(agent_session.started_at)
+ sandbox.check_rate_limit(username)
+
+ # Record user message
+ from potato.agent_proxy import AgentMessage
+ user_msg = AgentMessage(role="user", content=message)
+ agent_session.messages.append(user_msg)
+
+ # Send to agent (blocking)
+ response = agent_session.proxy.send_message(
+ message, agent_session.proxy_context
+ )
+ agent_session.messages.append(response.message)
+ agent_session.step_count += 1
+
+ return jsonify({
+ "content": response.message.content,
+ "role": response.message.role,
+ "step_count": agent_session.step_count,
+ "max_steps": sandbox.max_steps,
+ "error": response.error,
+ })
+
+ except Exception as e:
+ # Surface sandbox violations (step limit, timeout, rate limit) directly
+ from potato.agent_proxy.sandbox import SandboxViolation
+ if isinstance(e, SandboxViolation):
+ return jsonify({"error": str(e)}), 400
+ logger.error("Agent chat send error: %s", traceback.format_exc())
+ return jsonify({"error": "An internal error occurred"}), 400
+
+
+@app.route("/agent_chat/finish", methods=["POST"])
+def agent_chat_finish():
+ """Finish the chat and write conversation data to the item."""
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ username = session['username']
+ user_state = get_user_state(username)
+
+ if user_state.get_phase() != UserPhase.ANNOTATION:
+ return jsonify({"error": "Not in annotation phase"}), 400
+
+ item = user_state.get_current_instance()
+ instance_id = item.get_id()
+
+ try:
+ from potato.agent_proxy import get_agent_session_manager
+ mgr = get_agent_session_manager()
+ agent_session = mgr.get_session(username, instance_id)
+
+ if not agent_session:
+ return jsonify({"error": "No active chat session"}), 400
+
+ if agent_session.finished:
+ return jsonify({"error": "Chat session already finished"}), 400
+
+ # Convert messages to conversation data (agent_trace format)
+ conversation = []
+ for msg in agent_session.messages:
+ speaker = "User" if msg.role == "user" else "Agent"
+ if msg.role == "error":
+ speaker = "System (Error)"
+ conversation.append({
+ "speaker": speaker,
+ "text": msg.content,
+ "timestamp": msg.timestamp,
+ })
+
+ # Write conversation into item data
+ item_data = item.get_data()
+ # Find the conversation field key from instance_display config
+ conv_key = "conversation"
+ display_fields = config.get("instance_display", {}).get("fields", [])
+ for field in display_fields:
+ if field.get("type") == "interactive_chat":
+ conv_key = field.get("key", "conversation")
+ break
+
+ item_data[conv_key] = conversation
+
+ # Mark session as finished
+ agent_session.finished = True
+
+ logger.info(
+ f"Agent chat finished for user={username}, instance={instance_id}, "
+ f"steps={agent_session.step_count}, messages={len(conversation)}"
+ )
+
+ return jsonify({"success": True, "message_count": len(conversation)})
+
+ except Exception as e:
+ logger.error("Agent chat finish error: %s", traceback.format_exc())
+ return jsonify({"error": "An internal error occurred"}), 500
+
+
+@app.route("/agent_chat/status", methods=["GET"])
+def agent_chat_status():
+ """Get the current agent chat session status (for page refresh recovery)."""
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ username = session['username']
+ user_state = get_user_state(username)
+
+ if user_state.get_phase() != UserPhase.ANNOTATION:
+ return jsonify({"active": False})
+
+ item = user_state.get_current_instance()
+ instance_id = item.get_id()
+
+ try:
+ from potato.agent_proxy import get_agent_session_manager
+ mgr = get_agent_session_manager()
+ agent_session = mgr.get_session(username, instance_id)
+
+ if not agent_session or agent_session.finished:
+ return jsonify({"active": False})
+
+ sandbox = _get_agent_sandbox()
+
+ messages = [
+ {"role": msg.role, "content": msg.content}
+ for msg in agent_session.messages
+ ]
+
+ return jsonify({
+ "active": True,
+ "messages": messages,
+ "step_count": agent_session.step_count,
+ "max_steps": sandbox.max_steps,
+ })
+
+ except Exception:
+ return jsonify({"active": False})
+
+
+@app.route("/go_to", methods=["GET", "POST"])
+def go_to():
+ """
+ Handle requests to go to a specific instance.
+ """
+ if 'username' not in session:
+ return home()
+
+ username = session['username']
+ user_state = get_user_state(username)
+
+ # Check that the user is in the annotation phase
+ if user_state.get_phase() != UserPhase.ANNOTATION:
+ # If not in the annotation phase, redirect
+ return home()
+
+ if request.method == 'POST':
+ logger.debug(f'POST -> GO_TO: {request.form}')
+ go_to_id(username, request.form.get("go_to"))
+
+ # Prevent browser caching so window.location.reload() always gets fresh content
+ response = make_response(render_page_with_annotations(username))
+ response.headers['Cache-Control'] = 'no-store, no-cache, must-revalidate, max-age=0'
+ response.headers['Pragma'] = 'no-cache'
+ return response
+
+@app.route('/get_annotations', methods=['GET'])
+def get_annotations():
+ """Get annotations for the current user and instance."""
+ try:
+ # Get user from session
+ if 'username' not in session:
+ return jsonify({"error": "No user session"}), 401
+
+ username = session['username']
+
+ # Get instance ID from query parameters
+ instance_id = request.args.get('instance_id')
+ if not instance_id:
+ return jsonify({"error": "No instance_id provided"}), 400
+
+ # Get user state
+ user_state = get_user_state_manager().get_user_state(username)
+ if not user_state:
+ return jsonify({"error": "User not found"}), 404
+
+ # Get annotations for the instance
+ label_annotations = user_state.get_label_annotations(instance_id)
+ span_annotations = user_state.get_span_annotations(instance_id)
+
+ # label_annotations is keyed by Label objects (schema, name) ->
+ # value; that is not JSON-serializable. Flatten to
+ # {schema: [selected label names]} (skip falsy/unset values).
+ serializable_label_annotations = {}
+ for lbl, value in (label_annotations or {}).items():
+ if value in (False, None, "", 0):
+ continue
+ schema = getattr(lbl, "schema", None)
+ name = getattr(lbl, "name", None)
+ if schema is None or name is None:
+ # Already-serialized or unexpected shape; keep as-is.
+ serializable_label_annotations.setdefault(
+ str(lbl), value)
+ continue
+ serializable_label_annotations.setdefault(
+ schema, []).append(name)
+
+ # Convert span annotations to serializable format
+ serializable_span_annotations = {}
+ for span, value in span_annotations.items():
+ serializable_span_annotations[str(span)] = value
+
+ # Combine annotations
+ annotations = {
+ "label_annotations": serializable_label_annotations,
+ "span_annotations": serializable_span_annotations
+ }
+
+ return jsonify(annotations)
+
+ except Exception as e:
+ logger.error(f"Error getting annotations: {e}")
+ return jsonify({"error": "Internal server error"}), 500
+
+@app.route("/api/current_instance", methods=["GET"])
+def get_current_instance():
+ """Get the current instance information for the current user."""
+ logger.debug(f"=== GET_CURRENT_INSTANCE START ===")
+
+ if 'username' not in session:
+ logger.warning("Get current instance without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ username = session['username']
+ logger.debug(f"Username: {username}")
+
+ try:
+ user_state = get_user_state(username)
+ if not user_state:
+ logger.error(f"User state not found for user: {username}")
+ return jsonify({"error": "User state not found"}), 404
+
+ # Guard: only return instance data during annotation phase
+ if user_state.get_phase() != UserPhase.ANNOTATION:
+ logger.debug(f"User {username} not in annotation phase, no current instance")
+ return jsonify({"error": "Not in annotation phase"}), 404
+
+ current_instance = user_state.get_current_instance()
+ if not current_instance:
+ logger.error(f"No current instance for user: {username}")
+ return jsonify({"error": "No current instance"}), 404
+
+ instance_id = current_instance.get_id()
+ logger.debug(f"Current instance ID: {instance_id}")
+
+ # Include raw data for schemas that need access to media URLs
+ raw_data = current_instance.get_data()
+
+ return jsonify({
+ "instance_id": instance_id,
+ "current_index": user_state.get_current_instance_index(),
+ "total_instances": len(user_state.instance_id_ordering),
+ "data": raw_data # Include full instance data
+ })
+
+ except Exception as e:
+ logger.error("Error getting current instance: %s", traceback.format_exc())
+ return jsonify({"error": "An internal error occurred"}), 500
+
+
+@app.route("/api/instance_data", methods=["GET"])
+def get_instance_data():
+ """Get the full raw data for the current instance.
+
+ Returns all fields from the original data file (e.g., audio_url, video_url, etc.)
+ This is used by annotation schemas that need access to media URLs.
+ """
+ logger.debug(f"=== GET_INSTANCE_DATA START ===")
+
+ if 'username' not in session:
+ logger.warning("Get instance data without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ username = session['username']
+ logger.debug(f"Username: {username}")
+
+ try:
+ user_state = get_user_state(username)
+ if not user_state:
+ logger.error(f"User state not found for user: {username}")
+ return jsonify({"error": "User state not found"}), 404
+
+ # Guard: only return instance data during annotation phase
+ if user_state.get_phase() != UserPhase.ANNOTATION:
+ logger.debug(f"User {username} not in annotation phase, no instance data")
+ return jsonify({"error": "Not in annotation phase"}), 404
+
+ current_instance = user_state.get_current_instance()
+ if not current_instance:
+ logger.error(f"No current instance for user: {username}")
+ return jsonify({"error": "No current instance"}), 404
+
+ # Get the raw data from the instance
+ raw_data = current_instance.get_data()
+ logger.debug(f"Returning instance data with keys: {list(raw_data.keys())}")
+
+ return jsonify(raw_data)
+
+ except Exception as e:
+ logger.error("Error getting instance data: %s", traceback.format_exc())
+ return jsonify({"error": "An internal error occurred"}), 500
+
+
+@app.route("/api/spans/")
+def get_span_data(instance_id):
+ """
+ Get span annotations as structured data for frontend rendering.
+
+ Returns:
+ JSON with instance text and span annotations in a format
+ suitable for frontend DOM manipulation.
+ """
+ logger.debug(f"=== GET_SPAN_DATA START ===")
+ logger.debug(f"Instance ID: {instance_id}")
+
+ # Add debugging for URL decoding
+ import urllib.parse
+ decoded_instance_id = urllib.parse.unquote(instance_id)
+ logger.debug(f"Decoded Instance ID: {decoded_instance_id}")
+ logger.debug(f"Instance ID length: {len(instance_id)}")
+ logger.debug(f"Decoded Instance ID length: {len(decoded_instance_id)}")
+
+ if 'username' not in session:
+ logger.warning("Get span data without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ username = session['username']
+ logger.debug(f"Username: {username}")
+
+ # Get the original text for this instance
+ try:
+ # Get the text from the item state manager
+ item_state_manager = get_item_state_manager()
+
+ # Try with original instance_id first
+ instance = item_state_manager.get_item(instance_id)
+ if not instance:
+ logger.debug(f"Instance not found with original ID, trying decoded ID")
+ # Try with decoded instance_id
+ instance = item_state_manager.get_item(decoded_instance_id)
+ if instance:
+ logger.debug(f"Instance found with decoded ID")
+ instance_id = decoded_instance_id # Use decoded ID for rest of function
+ else:
+ logger.error(f"Instance not found with either original or decoded ID")
+ # Debug: list all available instance IDs
+ all_instance_ids = list(item_state_manager.instance_id_to_instance.keys())
+ logger.debug(f"Available instance IDs: {all_instance_ids[:5]}...") # Show first 5
+ logger.debug(f"Total available instances: {len(all_instance_ids)}")
+ return jsonify({"error": "Instance not found"}), 404
+ else:
+ logger.debug(f"Instance found with original ID")
+
+ # Use configured text_key to get the right field, not generic get_text()
+ text_key = config.get("item_properties", {}).get("text_key", "text")
+ item_data = instance.get_data()
+ original_text = item_data.get(text_key, instance.get_text()) if isinstance(item_data, dict) else instance.get_text()
+ logger.debug(f"Original text (raw, text_key={text_key}): {str(original_text)[:100]}...")
+
+ # IMPORTANT: Normalize text the same way as flask_server.py template rendering
+ # This ensures span offsets calculated on normalized text match the API response
+ # 1. Strip HTML tags
+ import re as re_module
+ original_text = str(original_text)
+ normalized_text = re_module.sub(r'<[^>]+>', '', original_text)
+ # 2. Normalize whitespace (multiple spaces/newlines -> single space)
+ normalized_text = re_module.sub(r'\s+', ' ', normalized_text).strip()
+ logger.debug(f"Normalized text: {normalized_text[:100]}...")
+ except Exception as e:
+ logger.error(f"Error getting instance text: {e}")
+ return jsonify({"error": f"Instance not found: {instance_id}"}), 404
+
+ # Get span annotations (returns a list of SpanAnnotation objects)
+ spans = get_span_annotations_for_user_on(username, instance_id)
+ logger.debug(f"Found {len(spans)} spans")
+
+ # Convert to frontend-friendly format
+ span_data = []
+ for span in spans:
+ # span is a SpanAnnotation object
+ span_schema = span.get_schema() if hasattr(span, 'get_schema') else span.schema
+ span_name = span.get_name() if hasattr(span, 'get_name') else span.name
+ span_title = span.get_title() if hasattr(span, 'get_title') else getattr(span, 'title', span_name)
+ span_start = span.get_start() if hasattr(span, 'get_start') else span.start
+ span_end = span.get_end() if hasattr(span, 'get_end') else span.end
+ span_id = span.get_id() if hasattr(span, 'get_id') else getattr(span, 'id', None)
+
+ color = get_span_color(span_schema, span_name)
+ hex_color = None
+ if color:
+ if isinstance(color, str) and color.startswith("(") and color.endswith(")"):
+ try:
+ rgb_parts = color.strip("()").split(", ")
+ if len(rgb_parts) == 3:
+ r, g, b = int(rgb_parts[0]), int(rgb_parts[1]), int(rgb_parts[2])
+ hex_color = f"#{r:02x}{g:02x}{b:02x}"
+ except (ValueError, IndexError):
+ hex_color = "#f0f0f0"
+ else:
+ hex_color = color
+
+ span_target_field = span.get_target_field() if hasattr(span, 'get_target_field') else getattr(span, 'target_field', None)
+
+ # Use the correct field text for extracting span text
+ # In multi-field mode, each span's offsets are relative to its target field's text
+ span_source_text = normalized_text # default: text_key field
+ if span_target_field and isinstance(item_data, dict) and span_target_field in item_data:
+ field_data = item_data[span_target_field]
+
+ # Check if this field is a dialogue display (offsets are DOM-based)
+ from potato.server_utils.displays.base import (
+ concatenate_dialogue_text,
+ reconstruct_dialogue_dom_text,
+ )
+ display_fields = config.get("instance_display", {}).get("fields", [])
+ field_cfg = next(
+ (f for f in display_fields if f.get("key") == span_target_field),
+ None,
+ )
+ is_dialogue = field_cfg and field_cfg.get("type") == "dialogue"
+
+ if is_dialogue and isinstance(field_data, list):
+ opts = field_cfg.get("display_options", {})
+ field_text = reconstruct_dialogue_dom_text(
+ field_data,
+ speaker_key=opts.get("speaker_key", "speaker"),
+ text_key=opts.get("text_key", "text"),
+ show_turn_numbers=opts.get("show_turn_numbers", False),
+ )
+ elif isinstance(field_data, list):
+ field_text = concatenate_dialogue_text(field_data)
+ field_text = re_module.sub(r'<[^>]+>', '', field_text)
+ field_text = re_module.sub(r'\s+', ' ', field_text).strip()
+ else:
+ field_text = str(field_data)
+ field_text = re_module.sub(r'<[^>]+>', '', field_text)
+ field_text = re_module.sub(r'\s+', ' ', field_text).strip()
+ span_source_text = field_text
+
+ span_entry = {
+ 'id': span_id,
+ 'schema': span_schema,
+ 'label': span_name,
+ 'title': span_title,
+ 'start': span_start,
+ 'end': span_end,
+ 'text': span_source_text[span_start:span_end] if span_start < len(span_source_text) and span_end <= len(span_source_text) else "",
+ 'color': hex_color
+ }
+ if span_target_field:
+ span_entry['target_field'] = span_target_field
+
+ # Include additional_parts for discontinuous spans
+ additional_parts = span.get_additional_parts() if hasattr(span, 'get_additional_parts') else getattr(span, 'additional_parts', [])
+ if additional_parts:
+ span_entry['additional_parts'] = additional_parts
+
+ # Include entity linking data if present
+ kb_id = span.get_kb_id() if hasattr(span, 'get_kb_id') else getattr(span, 'kb_id', None)
+ kb_source = span.get_kb_source() if hasattr(span, 'get_kb_source') else getattr(span, 'kb_source', None)
+ kb_label = span.get_kb_label() if hasattr(span, 'get_kb_label') else getattr(span, 'kb_label', None)
+ if kb_id and kb_source:
+ span_entry['kb_id'] = kb_id
+ span_entry['kb_source'] = kb_source
+ if kb_label:
+ span_entry['kb_label'] = kb_label
+
+ span_data.append(span_entry)
+
+ response_data = {
+ 'instance_id': instance_id,
+ 'text': normalized_text, # Use normalized text matching template rendering
+ 'spans': span_data
+ }
+
+ logger.debug(f"=== GET_SPAN_DATA END ===", response_data)
+ return jsonify(response_data)
+
+
+@app.route("/updateinstance", methods=["POST"])
+def update_instance():
+ """
+ PRIMARY ANNOTATION ENDPOINT: Handle all annotation updates for instances.
+ This endpoint only updates backend state for spans and labels. It does not generate or return any HTML.
+
+ Supports two formats:
+ 1. Frontend format: {"instance_id": "...", "annotations": {...}, "span_annotations": [...]}
+ 2. Backend format: {"instance_id": "...", "schema": "...", "state": [...], "type": "..."}
+ """
+ import time
+ import datetime
+ from potato.annotation_history import AnnotationHistoryManager
+
+ start_time = time.time()
+
+ logger.debug("=== UPDATEINSTANCE ROUTE START ===")
+ logger.debug(f"Session: {dict(session)}")
+ logger.debug(f"Session username: {session.get('username', 'NOT_SET')}")
+ logger.debug(f"Request content type: {request.content_type}")
+ logger.debug(f"Request is JSON: {request.is_json}")
+ logger.debug(f"Debug mode: {config.get('debug', False)}")
+
+ if 'username' not in session:
+ logger.warning("Update instance without active session")
+ return jsonify({"status": "error", "message": "No active session"})
+
+ if request.is_json:
+ logger.debug(f"Received JSON data: {request.json}")
+ raw_instance_id = request.json.get("instance_id")
+ if raw_instance_id is None or str(raw_instance_id).strip() == "":
+ logger.warning(f"Received update with null/empty instance_id from user {session.get('username')}")
+ return jsonify({"status": "error", "message": "Missing instance_id"})
+ instance_id = str(raw_instance_id) # Normalize to string
+ username = session['username']
+ user_state = get_user_state(username)
+ if not user_state:
+ logger.error(f"User state not found for user: {username}")
+ return jsonify({"status": "error", "message": "User state not found"})
+
+ # Synthetic phase-page saves use a sentinel instance ID from annotation.js.
+ # They should be routed by the current user phase instead of assignment checks.
+ is_phase_page_update = instance_id == "__phase_page__"
+
+ # Guard: reject updates for instances not assigned to this user.
+ # Skip this for synthetic phase-page updates so non-annotation page autosaves
+ # do not get treated like dataset item writes.
+ if user_state.get_phase() == UserPhase.ANNOTATION and not is_phase_page_update:
+ assigned_ids = user_state.get_assigned_instance_ids()
+ if assigned_ids and instance_id not in assigned_ids:
+ logger.warning(f"User {username} tried to update unassigned instance {instance_id}")
+ return jsonify({"status": "error", "message": "Instance not assigned to user"})
+
+ # If a synthetic phase-page save arrives while the backend still thinks the user
+ # is in annotation, acknowledge it without writing item state. This avoids a
+ # noisy warning/error path during end-of-annotation transitions.
+ if user_state.get_phase() == UserPhase.ANNOTATION and is_phase_page_update:
+ logger.info(
+ f"Ignoring synthetic phase-page update for user {username} while still in annotation phase"
+ )
+ return jsonify({"status": "ok", "message": "Ignored synthetic phase-page update"})
+
+ # Debug: Log user phase for debugging annotation storage issues
+ logger.debug(f"User '{username}' phase: {user_state.get_phase()}, current_phase_and_page: {user_state.current_phase_and_page}")
+
+ # Track session
+ if not user_state.session_start_time:
+ user_state.start_session(session.get('session_id', str(uuid.uuid4())))
+
+ # Get client timestamp if provided
+ client_timestamp = None
+ if request.json.get("client_timestamp"):
+ try:
+ client_timestamp = datetime.datetime.fromisoformat(request.json["client_timestamp"])
+ except ValueError:
+ logger.warning(f"Invalid client timestamp format: {request.json['client_timestamp']}")
+
+ # Prepare metadata
+ metadata = {
+ "request_id": request.json.get("request_id"),
+ "user_agent": request.headers.get("User-Agent"),
+ "ip_address": request.remote_addr,
+ "content_type": request.content_type,
+ "request_size": len(request.get_data()) if request.get_data() else 0
+ }
+
+ # Capture whether this instance already had annotations BEFORE we mutate
+ # state, so the webhook below can correctly distinguish annotation.created
+ # from annotation.updated. F-032.
+ _had_prior_annotation = bool(
+ getattr(user_state, "instance_id_to_label_to_value", {}).get(instance_id)
+ )
+
+ # Check if this is the frontend format (annotations, span_annotations, link_annotations, event_annotations)
+ if "annotations" in request.json or "span_annotations" in request.json or "link_annotations" in request.json or "event_annotations" in request.json:
+ logger.debug("Processing frontend format (annotations, span_annotations)")
+
+ # Handle label annotations from frontend format
+ annotations = request.json.get("annotations", {})
+
+ # Guard against malformed payloads: `annotations` must be a mapping.
+ # A non-dict (e.g. a JSON string or list) would crash the .items()
+ # loop below with AttributeError โ an unhandled HTTP 500 (F-046).
+ # Reject cleanly instead.
+ if not isinstance(annotations, dict):
+ logger.warning(
+ f"Rejected /updateinstance from {username}: 'annotations' is "
+ f"{type(annotations).__name__}, expected object")
+ return jsonify({"status": "error",
+ "message": "'annotations' must be an object"})
+
+ # Pre-clear stale labels for radio/multiselect schemas.
+ # The client always sends the COMPLETE current state, so any label
+ # not in the incoming set should be removed. Without this, deselected
+ # radio options or unchecked checkboxes persist as stale data.
+ _exclusive_types = {'radio', 'multiselect'}
+ _schema_type_cache = {}
+ for scheme in config.get('annotation_schemes', []):
+ _schema_type_cache[scheme.get('name')] = scheme.get('annotation_type')
+
+ # Collect which schemas appear in the incoming annotations
+ _incoming_schemas = set()
+ for key in annotations:
+ sep = ":::" if ":::" in key else (":" if ":" in key else None)
+ if sep:
+ _incoming_schemas.add(key.split(sep, 1)[0])
+
+ # For each exclusive schema, remove all existing labels
+ if instance_id in user_state.instance_id_to_label_to_value:
+ for schema_name_to_clear in _incoming_schemas:
+ if _schema_type_cache.get(schema_name_to_clear) in _exclusive_types:
+ labels_to_remove = [
+ lbl for lbl in user_state.instance_id_to_label_to_value[instance_id]
+ if isinstance(lbl, Label) and lbl.get_schema() == schema_name_to_clear
+ ]
+ for lbl in labels_to_remove:
+ del user_state.instance_id_to_label_to_value[instance_id][lbl]
+
+ for key, value in annotations.items():
+ if ":::" in key:
+ # Use ::: separator for image/audio/video annotation data
+ # e.g., "video_segments:::_data" -> schema="video_segments", label="_data"
+ schema_name, label_name = key.split(":::", 1)
+ label = Label(schema_name, label_name)
+ elif ":" in key:
+ # Legacy format with single colon
+ schema_name, label_name = key.split(":", 1)
+ label = Label(schema_name, label_name)
+ else:
+ logger.warning(f"Skipping annotation with no separator: {key}")
+ continue
+
+ # Get old value for comparison
+ old_value = None
+ if instance_id in user_state.instance_id_to_label_to_value:
+ old_value = user_state.instance_id_to_label_to_value[instance_id].get(label)
+
+ # Determine action type
+ action_type = "add_label" if old_value is None else "update_label"
+
+ # Create annotation action
+ action = AnnotationHistoryManager.create_action(
+ user_id=username,
+ instance_id=instance_id,
+ action_type=action_type,
+ schema_name=schema_name,
+ label_name=label_name,
+ old_value=old_value,
+ new_value=value,
+ session_id=user_state.current_session_id,
+ client_timestamp=client_timestamp,
+ metadata=metadata
+ )
+
+ # Add to history
+ user_state.add_annotation_action(action)
+
+ # Update annotation
+ user_state.add_label_annotation(instance_id, label, value)
+ logger.debug(f"Added label annotation: {schema_name}:{label_name} = {value[:100]}..." if len(str(value)) > 100 else f"Added label annotation: {schema_name}:{label_name} = {value}")
+
+ # Record humanโjudge comparison for the alignment loop (inline mode).
+ _record_judge_comparison_if_enabled(username, instance_id)
+
+ # Handle span annotations from frontend format
+ span_annotations = request.json.get("span_annotations", [])
+ for span_data in span_annotations:
+ if isinstance(span_data, dict) and "schema" in span_data:
+ # Use provided ID or generate deterministic one to preserve span identity
+ span_id = span_data.get("id") or span_data.get("span_id") or \
+ f"{span_data['schema']}_{span_data['name']}_{span_data['start']}_{span_data['end']}"
+
+ span = SpanAnnotation(
+ span_data["schema"],
+ span_data["name"],
+ span_data.get("title", span_data["name"]),
+ int(span_data["start"]),
+ int(span_data["end"]),
+ id=span_id,
+ target_field=span_data.get("target_field")
+ )
+ value = span_data.get("value")
+
+ if value is not None:
+ # Get old value for comparison
+ old_value = None
+ if instance_id in user_state.instance_id_to_span_to_value:
+ old_value = user_state.instance_id_to_span_to_value[instance_id].get(span)
+
+ # Determine action type
+ action_type = "add_span" if old_value is None else "update_span"
+
+ # Create annotation action
+ action = AnnotationHistoryManager.create_action(
+ user_id=username,
+ instance_id=instance_id,
+ action_type=action_type,
+ schema_name=span_data["schema"],
+ label_name=span_data["name"],
+ old_value=old_value,
+ new_value=value,
+ span_data={
+ "start": span_data["start"],
+ "end": span_data["end"],
+ "title": span_data.get("title", span_data["name"])
+ },
+ session_id=user_state.current_session_id,
+ client_timestamp=client_timestamp,
+ metadata=metadata
+ )
+
+ # Add to history
+ user_state.add_annotation_action(action)
+
+ # Update annotation
+ user_state.add_span_annotation(instance_id, span, value)
+ logger.debug(f"Added span annotation: {span_data}")
+
+ # Handle link annotations from frontend format
+ link_annotations = request.json.get("link_annotations", [])
+ for link_data in link_annotations:
+ if isinstance(link_data, dict) and "schema" in link_data and "link_type" in link_data:
+ from potato.item_state_management import SpanLink
+ link = SpanLink(
+ schema=link_data["schema"],
+ link_type=link_data["link_type"],
+ span_ids=link_data.get("span_ids", []),
+ direction=link_data.get("direction", "undirected"),
+ id=link_data.get("id"),
+ properties=link_data.get("properties", {})
+ )
+
+ # Add or update the link annotation
+ user_state.add_link_annotation(instance_id, link)
+ logger.debug(f"Added link annotation: {link}")
+
+ # Handle event annotations from frontend format
+ event_annotations = request.json.get("event_annotations", [])
+ logger.debug(f"Processing {len(event_annotations)} event annotations")
+ for event_data in event_annotations:
+ logger.debug(f"Processing event data: {event_data}")
+ if isinstance(event_data, dict) and "schema" in event_data and "event_type" in event_data:
+ from potato.item_state_management import EventAnnotation
+
+ # Log the incoming ID
+ incoming_id = event_data.get("id")
+ logger.debug(f"Incoming event ID: {incoming_id}")
+
+ event = EventAnnotation(
+ schema=event_data["schema"],
+ event_type=event_data["event_type"],
+ trigger_span_id=event_data.get("trigger_span_id", ""),
+ arguments=event_data.get("arguments", []),
+ id=incoming_id,
+ properties=event_data.get("properties", {})
+ )
+
+ # Log the actual ID assigned
+ logger.debug(f"Event object ID after creation: {event.get_id()}")
+
+ # Add or update the event annotation
+ user_state.add_event_annotation(instance_id, event)
+
+ # Log current events count
+ current_events = user_state.get_event_annotations(instance_id)
+ logger.debug(f"Added event annotation. Total events for instance: {len(current_events)}")
+ logger.debug(f"Current event IDs: {list(current_events.keys())}")
+
+ # Check if this is the backend format (schema, state, type)
+ elif "schema" in request.json and "state" in request.json and "type" in request.json:
+ logger.debug("Processing backend format (schema, state, type)")
+
+ schema_name = request.json.get("schema")
+ schema_state = request.json.get("state")
+ annotation_type = request.json.get("type")
+
+ if annotation_type == "span":
+ logger.debug(f"Processing span annotation state: {schema_state}")
+ for sv in schema_state:
+ # Validate and correct negative offsets
+ start_offset = int(sv["start"])
+ end_offset = int(sv["end"])
+
+ # Correct negative offsets to 0
+ if start_offset < 0:
+ start_offset = 0
+ logger.warning(f"Corrected negative start offset {sv['start']} to 0")
+ if end_offset < 0:
+ end_offset = 0
+ logger.warning(f"Corrected negative end offset {sv['end']} to 0")
+
+ # Ensure end is not less than start
+ if end_offset < start_offset:
+ end_offset = start_offset
+ logger.warning(f"Corrected end offset {sv['end']} to match start offset {start_offset}")
+
+ # Get span_id or generate one if not provided
+ span_id = sv.get("span_id") or sv.get("id") or f"{schema_name}_{sv['name']}_{start_offset}_{end_offset}"
+
+ # Get additional_parts for discontinuous spans
+ additional_parts = sv.get("additional_parts", [])
+
+ span = SpanAnnotation(
+ schema_name,
+ sv["name"],
+ sv.get("title", sv["name"]),
+ start_offset,
+ end_offset,
+ span_id,
+ target_field=sv.get("target_field"),
+ additional_parts=additional_parts
+ )
+
+ value = sv.get("value")
+
+ # Get old value for comparison
+ old_value = None
+ if instance_id in user_state.instance_id_to_span_to_value:
+ old_value = user_state.instance_id_to_span_to_value[instance_id].get(span)
+
+ # Determine action type
+ if value is None:
+ action_type = "delete_span"
+ else:
+ action_type = "add_span" if old_value is None else "update_span"
+
+ # Create annotation action
+ action = AnnotationHistoryManager.create_action(
+ user_id=username,
+ instance_id=instance_id,
+ action_type=action_type,
+ schema_name=schema_name,
+ label_name=sv["name"],
+ old_value=old_value,
+ new_value=value,
+ span_data={
+ "start": start_offset,
+ "end": end_offset,
+ "title": sv.get("title", sv["name"])
+ },
+ session_id=user_state.current_session_id,
+ client_timestamp=client_timestamp,
+ metadata=metadata
+ )
+
+ # Add to history
+ user_state.add_annotation_action(action)
+
+ # Handle span deletion vs creation/update
+ if value is None:
+ # Delete the span - find and remove the matching span
+ if instance_id in user_state.instance_id_to_span_to_value:
+ # Find the span to delete by matching properties
+ spans_to_delete = []
+ for existing_span in user_state.instance_id_to_span_to_value[instance_id].keys():
+ if (existing_span.get_schema() == span.get_schema() and
+ existing_span.get_name() == span.get_name() and
+ existing_span.get_start() == span.get_start() and
+ existing_span.get_end() == span.get_end()):
+ spans_to_delete.append(existing_span)
+
+ for span_to_delete in spans_to_delete:
+ del user_state.instance_id_to_span_to_value[instance_id][span_to_delete]
+ logger.debug(f"Deleted span annotation: {span_to_delete}")
+
+ # Clean up orphaned links and events referencing this span
+ deleted_span_id = span_to_delete.get_id()
+ if instance_id in user_state.instance_id_to_link_to_value:
+ orphaned_links = [
+ link_id for link_id, link in user_state.instance_id_to_link_to_value[instance_id].items()
+ if deleted_span_id in link.get_span_ids()
+ ]
+ for link_id in orphaned_links:
+ del user_state.instance_id_to_link_to_value[instance_id][link_id]
+ logger.debug(f"Removed orphaned link {link_id} referencing deleted span {deleted_span_id}")
+ if instance_id in user_state.instance_id_to_event_to_value:
+ orphaned_events = [
+ evt_id for evt_id, evt in user_state.instance_id_to_event_to_value[instance_id].items()
+ if deleted_span_id in evt.get_all_span_ids()
+ ]
+ for evt_id in orphaned_events:
+ del user_state.instance_id_to_event_to_value[instance_id][evt_id]
+ logger.debug(f"Removed orphaned event {evt_id} referencing deleted span {deleted_span_id}")
+ else:
+ # Add or update the span annotation
+ user_state.add_span_annotation(instance_id, span, value)
+ logger.debug(f"Added span annotation: {span} with value: {value}")
+ elif annotation_type == "label":
+ for sv in schema_state:
+ label = Label(schema_name, sv["name"])
+ value = sv["value"]
+
+ # Get old value for comparison
+ old_value = None
+ if instance_id in user_state.instance_id_to_label_to_value:
+ old_value = user_state.instance_id_to_label_to_value[instance_id].get(label)
+
+ # Determine action type
+ action_type = "add_label" if old_value is None else "update_label"
+
+
+ # Create annotation action
+ action = AnnotationHistoryManager.create_action(
+ user_id=username,
+ instance_id=instance_id,
+ action_type=action_type,
+ schema_name=schema_name,
+ label_name=sv["name"],
+ old_value=old_value,
+ new_value=value,
+ session_id=user_state.current_session_id,
+ client_timestamp=client_timestamp,
+ metadata=metadata
+ )
+
+ # Add to history
+ user_state.add_annotation_action(action)
+
+ # Update annotation
+ user_state.add_label_annotation(instance_id, label, value)
+ else:
+ logger.warning("Unknown data format in /updateinstance")
+ return jsonify({"status": "error", "message": "Unknown data format"})
+
+ # Collect all annotations for this instance. Used by both QC validation
+ # and the webhook payload below โ must be defined unconditionally, since
+ # webhooks can be enabled without quality control (otherwise the webhook
+ # emit block raises UnboundLocalError and every save 500s). F-030.
+ all_annotations = {}
+ if "annotations" in request.json:
+ for key, value in request.json.get("annotations", {}).items():
+ # Parse schema:label format
+ if ":" in key:
+ schema_name, label_name = key.split(":", 1)
+ all_annotations[schema_name] = value
+ else:
+ all_annotations[key] = value
+ elif "schema" in request.json:
+ schema_name = request.json.get("schema")
+ schema_state = request.json.get("state", [])
+ # Convert state list to dict for validation
+ for sv in schema_state:
+ if "value" in sv:
+ all_annotations[schema_name] = sv.get("value")
+
+ # Quality control validation (attention checks and gold standards)
+ qc_manager = get_quality_control_manager()
+ qc_result = None
+
+ if qc_manager:
+ # Calculate response time
+ response_time = None
+ if client_timestamp:
+ response_time = (datetime.datetime.now() - client_timestamp).total_seconds()
+
+ # Check if this is an attention check
+ attention_result = qc_manager.validate_attention_response(
+ username, instance_id, all_annotations, response_time
+ )
+
+ if attention_result is not None:
+ qc_result = {"type": "attention_check", **attention_result}
+
+ # Handle blocking
+ if attention_result.get("blocked"):
+ logger.warning(f"User {username} blocked by attention check")
+ reclaimed = _reclaim_blocked_user_assignments(
+ username,
+ user_state,
+ current_instance_id=instance_id,
+ )
+
+ # Emit webhook for attention check failure
+ from potato.webhooks import get_webhook_emitter
+ from potato.webhooks.events import (
+ build_attention_check_failed_payload,
+ QUALITY_ATTENTION_CHECK_FAILED,
+ )
+ _wh = get_webhook_emitter()
+ if _wh:
+ _wh.emit(
+ QUALITY_ATTENTION_CHECK_FAILED,
+ build_attention_check_failed_payload(
+ user_id=username,
+ instance_id=instance_id,
+ message=attention_result.get("message"),
+ blocked=True,
+ ),
+ )
+
+ return jsonify({
+ "status": "blocked",
+ "message": attention_result.get("message", "You have been blocked."),
+ "qc_result": qc_result,
+ "reclaimed_assignments": reclaimed,
+ })
+ else:
+ # Check if this is a gold standard
+ gold_result = qc_manager.validate_gold_response(
+ username, instance_id, all_annotations
+ )
+
+ if gold_result is not None:
+ qc_result = {"type": "gold_standard", **gold_result}
+
+ # Record regular item for attention check frequency tracking
+ if not qc_manager.is_attention_check(instance_id) and not qc_manager.is_gold_standard(instance_id):
+ qc_manager.record_regular_item(username)
+
+ # Track for gold standard auto-promotion
+ promotion_result = qc_manager.record_item_annotation(
+ instance_id, username, all_annotations
+ )
+ if promotion_result and promotion_result.get("promoted"):
+ logger.info(f"Item {instance_id} auto-promoted to gold standard")
+
+ # Calculate processing time
+ processing_time_ms = int((time.time() - start_time) * 1000)
+
+ # Update the last action's processing time
+ if user_state.annotation_history:
+ user_state.annotation_history[-1].server_processing_time_ms = processing_time_ms
+
+ # Register annotator with item state manager for tracking
+ get_item_state_manager().register_annotator(instance_id, username)
+
+ # Save state
+ get_user_state_manager().save_user_state(user_state)
+ logger.debug(f"User state saved for {username}")
+
+ # Emit webhook events for annotation save
+ from potato.webhooks import get_webhook_emitter
+ from potato.webhooks.events import (
+ ANNOTATION_CREATED, ANNOTATION_UPDATED,
+ ITEM_FULLY_ANNOTATED,
+ build_annotation_payload, build_item_fully_annotated_payload,
+ )
+ _wh = get_webhook_emitter()
+ if _wh:
+ _evt = ANNOTATION_UPDATED if _had_prior_annotation else ANNOTATION_CREATED
+ _wh.emit(
+ _evt,
+ build_annotation_payload(
+ event_type=_evt,
+ user_id=username,
+ instance_id=instance_id,
+ annotations=all_annotations,
+ ),
+ )
+
+ # Check if item is now fully annotated. Use the canonical resolver
+ # for num_annotators_per_item (int or dict.default form); the prior
+ # code read annotation_task_name (a string) and always fell back to
+ # 3, ignoring the configured annotator count. F-031.
+ ism = get_item_state_manager()
+ annotators = ism.get_annotators_for_item(instance_id)
+ try:
+ from potato.server_utils.config_module import resolve_num_annotators_per_item
+ num_annotators = resolve_num_annotators_per_item(config)
+ except Exception:
+ num_annotators = config.get("num_annotators_per_item", 3)
+ if not isinstance(num_annotators, int) or num_annotators < 1:
+ num_annotators = 3
+ if len(annotators) >= num_annotators:
+ _wh.emit(
+ ITEM_FULLY_ANNOTATED,
+ build_item_fully_annotated_payload(
+ instance_id=instance_id,
+ annotator_count=len(annotators),
+ required_count=num_annotators,
+ ),
+ )
+
+ # Emit task.completed once, when this save annotates the user's LAST
+ # remaining assigned item. The documented event had a payload builder
+ # but zero emit sites. Fire-once is guarded by _had_prior_annotation:
+ # a re-save of an already-annotated item won't re-trigger it. F-033.
+ try:
+ assigned = set(user_state.get_assigned_instance_ids() or [])
+ if assigned and not _had_prior_annotation:
+ annotated = (
+ set(getattr(user_state, "instance_id_to_label_to_value", {}).keys())
+ | set(getattr(user_state, "instance_id_to_span_to_value", {}).keys())
+ ) & assigned
+ if assigned.issubset(annotated):
+ from potato.webhooks.events import (
+ TASK_COMPLETED, build_task_completed_payload,
+ )
+ _wh.emit(
+ TASK_COMPLETED,
+ build_task_completed_payload(
+ user_id=username,
+ total_annotations=len(annotated),
+ ),
+ )
+ except Exception as _e:
+ logger.debug(f"task.completed webhook check skipped: {_e}")
+
+ # Trigger MACE competence estimation check
+ from potato.mace_manager import get_mace_manager
+ mace_mgr = get_mace_manager()
+ if mace_mgr and mace_mgr.mace_config.enabled:
+ total = mace_mgr.count_total_annotations()
+ mace_mgr.check_and_run(total)
+
+ # Trigger active-learning retraining check. When enough new annotations
+ # have accumulated, this queues a background train + reorder of the
+ # unlabeled pool; never allowed to break the save.
+ try:
+ from potato.active_learning_manager import get_active_learning_manager
+ al_mgr = get_active_learning_manager()
+ if al_mgr and al_mgr.config.enabled:
+ al_mgr.check_and_trigger_training()
+ except Exception as e:
+ logger.debug("Active learning trigger skipped: %s", e)
+
+ # Stamp this annotation with the codebook revision in effect so
+ # an instance labeled before later code additions can be softly
+ # flagged for review. No-op unless the project uses a codebook;
+ # never allowed to break the save.
+ try:
+ if not is_phase_page_update:
+ from potato.codebook.api import codebook_enabled
+ if codebook_enabled(config):
+ from potato.codebook import record_annotation
+ record_annotation(
+ config.get("task_dir", "."),
+ config.get("annotation_task_name") or "default",
+ instance_id, username)
+ except Exception as e:
+ logger.warning(f"Codebook provenance stamp skipped: {e}")
+
+ # Get performance metrics for response
+ performance_metrics = user_state.get_performance_metrics()
+
+ response_data = {
+ "status": "success",
+ "processing_time_ms": processing_time_ms,
+ "performance_metrics": performance_metrics
+ }
+
+ # Include quality control result if present
+ if qc_result:
+ response_data["qc_result"] = qc_result
+
+ # Add warning message if needed
+ if qc_result.get("warning"):
+ response_data["warning"] = True
+ response_data["warning_message"] = qc_result.get("message")
+
+ return jsonify(response_data)
+ else:
+ logger.warning("Update instance called without JSON data")
+ return jsonify({"status": "error", "message": "JSON data required"})
+
+@app.route("/poststudy", methods=["GET", "POST"])
+def poststudy():
+ """
+ Handle the poststudy phase of the annotation process.
+
+ Returns:
+ flask.Response: Rendered template or redirect
+ """
+ if 'username' not in session:
+ return home()
+
+ username = session['username']
+ user_state = get_user_state(username)
+
+ # Check that the user is in the poststudy phase
+ if user_state.get_phase() != UserPhase.POSTSTUDY:
+ # If not in the poststudy phase, redirect
+ return home()
+
+ # If the user is returning information from the page
+ if request.method == 'POST':
+ logger.debug(f'POSTSTUDY: POST: {request.form}')
+
+ # Advance the state and move to the appropriate next phase
+ usm = get_user_state_manager()
+ usm.advance_phase(session['username'])
+
+ # Redirect to force a clean GET request (fixes POST leakage, issue #124)
+ return redirect(url_for("home"))
+
+ # Show the current poststudy page
+ else:
+ logger.debug("GET <-- POSTSTUDY")
+ return get_current_page_html(config, username)
+
+@app.route("/done", methods=["GET", "POST"])
+def done():
+ """
+ Handle the done phase of the annotation process.
+
+ This route displays the completion page with:
+ - A thank you message
+ - The completion code (if configured)
+ - A redirect link to Prolific (if configured)
+
+ Returns:
+ flask.Response: Rendered template or redirect
+ """
+ if 'username' not in session:
+ return home()
+
+ username = session['username']
+ user_state = get_user_state(username)
+
+ # Check that the user is in the done phase
+ if user_state.get_phase() != UserPhase.DONE:
+ # If not in the done phase, redirect
+ return home()
+
+ # Get completion code from config
+ completion_code = config.get("completion_code", "")
+
+ # Build Prolific redirect URL if completion code is set
+ prolific_redirect_url = None
+ login_config = config.get('login', {})
+ login_type = login_config.get('type', 'standard')
+
+ if completion_code and login_type in ['url_direct', 'prolific']:
+ # Build the Prolific completion URL (only if using Prolific-style URL argument)
+ url_argument = login_config.get('url_argument', 'PROLIFIC_PID')
+ if url_argument in ['PROLIFIC_PID', 'prolific_pid']:
+ # Format: https://app.prolific.com/submissions/complete?cc=YOUR_CODE
+ prolific_redirect_url = f"https://app.prolific.com/submissions/complete?cc={completion_code}"
+
+ # Get MTurk submission parameters from session
+ mturk_submit_url = session.get('mturk_submit_to')
+ mturk_assignment_id = session.get('mturk_assignment_id')
+
+ # Check for auto-redirect setting
+ auto_redirect = config.get('auto_redirect_on_completion', False)
+ auto_redirect_delay = config.get('auto_redirect_delay', 5000) # milliseconds
+
+ # Show the completion page
+ return render_template("done.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ completion_code=completion_code,
+ prolific_redirect_url=prolific_redirect_url,
+ mturk_submit_url=mturk_submit_url,
+ mturk_assignment_id=mturk_assignment_id,
+ auto_redirect=auto_redirect,
+ auto_redirect_delay=auto_redirect_delay)
+
+@app.route("/admin", methods=["GET"])
+def admin():
+ """
+ Serve the admin dashboard page.
+
+ This route serves the main admin dashboard interface with API key authentication.
+ The dashboard provides comprehensive system monitoring and management capabilities.
+
+ Returns:
+ flask.Response: Rendered admin dashboard template or login form
+ """
+ # Check if admin API key is provided in session or headers
+ api_key = request.headers.get('X-API-Key') or session.get('admin_api_key')
+
+ if not validate_admin_api_key(api_key):
+ # Show API key entry form
+ return render_template("admin_login.html",
+ title=config.get("annotation_task_name", "Admin Dashboard"))
+
+ # Store API key in session for future requests
+ session['admin_api_key'] = api_key
+
+ # Check if embedding visualization is available
+ from potato.embedding_visualization import get_embedding_viz_manager
+ viz_manager = get_embedding_viz_manager()
+ embedding_viz_enabled = viz_manager is not None and viz_manager.enabled
+
+ # Get basic context for the dashboard
+ context = {
+ "annotation_task_name": config.get("annotation_task_name", "Annotation Platform"),
+ "debug_mode": config.get("debug", False),
+ "admin_api_key": get_admin_api_key() or "",
+ "mace_enabled": config.get("mace", {}).get("enabled", False),
+ "bws_enabled": bool(config.get("bws_config")),
+ "embedding_viz_enabled": embedding_viz_enabled,
+ }
+
+ return render_template("admin.html", **context)
+
+
+
+
+
+
+
+@app.route("/api-frontend", methods=["GET"])
+def api_frontend():
+ """
+ Serve the API-based frontend interface.
+
+ This route serves a modern single-page application that uses API calls
+ to interact with the backend instead of server-side rendering.
+
+ Returns:
+ flask.Response: Rendered API frontend template
+ """
+ if 'username' not in session:
+ return redirect(url_for("home"))
+
+
+ username = session['username']
+
+ # Ensure user state exists
+ if not get_user_state_manager().has_user(username):
+ logger.info(f"Creating missing user state for {username}")
+ init_user_state(username)
+
+ user_state = get_user_state(username)
+
+ # Check user phase
+ if user_state.get_phase() != UserPhase.ANNOTATION:
+ logger.info(f"User {username} not in annotation phase, redirecting")
+ return redirect(url_for("home"))
+
+ # If the user hasn't yet been assigned anything to annotate, do so now
+ if not user_state.has_assignments():
+ get_item_state_manager().assign_instances_to_user(user_state)
+
+ # See if this user has finished annotating all of their assigned instances
+ if not user_state.has_remaining_assignments():
+ # If the user is done annotating, advance to the next phase
+ get_user_state_manager().advance_phase(username)
+ return redirect(url_for("home"))
+
+ # Render the API frontend template
+ return render_template("api_frontend.html",
+ username=username,
+ annotation_task_name=config.get("annotation_task_name", "Annotation Platform"),
+ annotation_codebook_url=config.get("annotation_codebook_url", ""),
+ alert_time_each_instance=config.get("alert_time_each_instance", 10000000))
+
+
+@app.route("/span-api-frontend", methods=["GET"])
+def span_api_frontend():
+ """
+ Serve the span annotation API-based frontend interface.
+
+ This route serves a modern single-page application specifically designed
+ for span annotation tasks that uses API calls to interact with the backend.
+
+ Returns:
+ flask.Response: Rendered span API frontend template
+ """
+ if 'username' not in session:
+ return redirect(url_for("home"))
+
+
+ username = session['username']
+
+ # Ensure user state exists
+ if not get_user_state_manager().has_user(username):
+ logger.info(f"Creating missing user state for {username}")
+ init_user_state(username)
+
+ user_state = get_user_state(username)
+
+ # Check user phase
+ if user_state.get_phase() != UserPhase.ANNOTATION:
+ logger.info(f"User {username} not in annotation phase, redirecting")
+ return redirect(url_for("home"))
+
+ # If the user hasn't yet been assigned anything to annotate, do so now
+ if not user_state.has_assignments():
+ get_item_state_manager().assign_instances_to_user(user_state)
+
+ # See if this user has finished annotating all of their assigned instances
+ if not user_state.has_remaining_assignments():
+ # If the user is done annotating, advance to the next phase
+ get_user_state_manager().advance_phase(username)
+ return redirect(url_for("home"))
+
+ # Render the span API frontend template
+ return render_template("span_api_frontend.html",
+ username=username,
+ annotation_task_name=config.get("annotation_task_name", "Span Annotation Platform"),
+ annotation_codebook_url=config.get("annotation_codebook_url", ""),
+ alert_time_each_instance=config.get("alert_time_each_instance", 10000000))
+
+@app.route("/test-span-colors")
+def test_span_colors():
+ """
+ Serve a test page for visually verifying span colors.
+ """
+ return render_template("test_span_colors.html")
+
+def normalize_color(color_value):
+ """
+ Normalize color value to a consistent format for the frontend.
+ Accepts: hex (#rrggbb), rgb/rgba, named colors, or tuple format "(r, g, b)".
+ Returns a CSS-compatible color string.
+ """
+ if not color_value:
+ return None
+
+ color_str = str(color_value).strip()
+
+ # Already a valid CSS color (hex, rgb, rgba, named)
+ if color_str.startswith('#') or color_str.startswith('rgb') or color_str.startswith('hsl'):
+ return color_str
+
+ # Tuple format "(r, g, b)" -> rgba
+ if color_str.startswith("(") and color_str.endswith(")"):
+ try:
+ rgb_parts = color_str.strip("()").split(",")
+ rgb_parts = [p.strip() for p in rgb_parts]
+ if len(rgb_parts) == 3:
+ r, g, b = int(rgb_parts[0]), int(rgb_parts[1]), int(rgb_parts[2])
+ return f"rgba({r}, {g}, {b}, 0.8)"
+ elif len(rgb_parts) == 4:
+ r, g, b, a = int(rgb_parts[0]), int(rgb_parts[1]), int(rgb_parts[2]), float(rgb_parts[3])
+ return f"rgba({r}, {g}, {b}, {a})"
+ except (ValueError, IndexError):
+ pass
+
+ # Named color - return as-is
+ return color_str
+
+
+# Default color palette for labels (used when no custom color is specified)
+DEFAULT_LABEL_COLORS = [
+ 'rgba(110, 86, 207, 0.8)', # Purple (primary)
+ 'rgba(34, 197, 94, 0.8)', # Green
+ 'rgba(239, 68, 68, 0.8)', # Red
+ 'rgba(59, 130, 246, 0.8)', # Blue
+ 'rgba(245, 158, 11, 0.8)', # Amber
+ 'rgba(236, 72, 153, 0.8)', # Pink
+ 'rgba(6, 182, 212, 0.8)', # Cyan
+ 'rgba(249, 115, 22, 0.8)', # Orange
+ 'rgba(139, 92, 246, 0.8)', # Violet
+ 'rgba(16, 185, 129, 0.8)', # Emerald
+]
+
+# Named color mappings for common label names
+NAMED_LABEL_COLORS = {
+ 'positive': 'rgba(34, 197, 94, 0.8)', # Green
+ 'negative': 'rgba(239, 68, 68, 0.8)', # Red
+ 'neutral': 'rgba(156, 163, 175, 0.8)', # Gray
+ 'mixed': 'rgba(245, 158, 11, 0.8)', # Amber
+ 'happy': 'rgba(34, 197, 94, 0.8)', # Green
+ 'sad': 'rgba(59, 130, 246, 0.8)', # Blue
+ 'angry': 'rgba(220, 38, 38, 0.8)', # Dark red
+ 'fear': 'rgba(139, 92, 246, 0.8)', # Violet
+ 'surprise': 'rgba(249, 115, 22, 0.8)', # Orange
+ 'disgust': 'rgba(132, 204, 22, 0.8)', # Lime
+ 'yes': 'rgba(34, 197, 94, 0.8)', # Green
+ 'no': 'rgba(239, 68, 68, 0.8)', # Red
+ 'maybe': 'rgba(245, 158, 11, 0.8)', # Amber
+ 'true': 'rgba(34, 197, 94, 0.8)', # Green
+ 'false': 'rgba(239, 68, 68, 0.8)', # Red
+ 'high': 'rgba(239, 68, 68, 0.8)', # Red
+ 'medium': 'rgba(245, 158, 11, 0.8)', # Amber
+ 'low': 'rgba(34, 197, 94, 0.8)', # Green
+}
+
+
+def get_default_label_color(label_name, index=0):
+ """
+ Get a default color for a label based on its name or index.
+ First checks for named colors, then falls back to palette by index.
+ """
+ # Check for named color match (case-insensitive).
+ # F-027: YAML parses unquoted yes/no/on/off/true/false as Python bools (and
+ # bare numbers as int/float), so a label "name" may not be a str โ coerce
+ # before string ops to avoid a 500 in /api/colors.
+ lower_name = str(label_name).lower().strip()
+ if lower_name in NAMED_LABEL_COLORS:
+ return NAMED_LABEL_COLORS[lower_name]
+
+ # Fall back to color from palette based on index
+ return DEFAULT_LABEL_COLORS[index % len(DEFAULT_LABEL_COLORS)]
+
+
+@app.route("/api/colors")
+def get_span_colors():
+ """
+ Return the color mapping for all schemas/labels as JSON.
+ Supports colors from:
+ 1. ui.label_colors - global color definitions by schema/label
+ 2. ui.spans.span_colors - legacy span-specific colors
+ 3. Inline 'color' property on labels in annotation_schemes
+ 4. Auto-generated colors from SPAN_COLOR_PALETTE
+ """
+ logger.debug("=== GET_COLORS START ===")
+
+ color_map = {}
+
+ # 1. Load colors from ui.label_colors (new unified format)
+ if "ui" in config and "label_colors" in config["ui"]:
+ logger.debug("Found ui.label_colors in config")
+ for schema_name, label_colors in config["ui"]["label_colors"].items():
+ color_map[schema_name] = {}
+ for label_name, color_value in label_colors.items():
+ normalized = normalize_color(color_value)
+ if normalized:
+ color_map[schema_name][label_name] = normalized
+
+ # 2. Load colors from ui.spans.span_colors (legacy format)
+ if "ui" in config and "spans" in config["ui"] and "span_colors" in config["ui"]["spans"]:
+ logger.debug("Found ui.spans.span_colors in config")
+ span_colors = config["ui"]["spans"]["span_colors"]
+ for schema_name, label_colors in span_colors.items():
+ if schema_name not in color_map:
+ color_map[schema_name] = {}
+ for label_name, color_value in label_colors.items():
+ if label_name not in color_map[schema_name]:
+ normalized = normalize_color(color_value)
+ if normalized:
+ color_map[schema_name][label_name] = normalized
+
+ # 3. Extract colors from annotation_schemes (inline label colors)
+ annotation_schemes = config.get('annotation_schemes', [])
+ if isinstance(annotation_schemes, list):
+ for schema in annotation_schemes:
+ schema_name = schema.get('name', f"schema_{schema.get('annotation_id', 'unknown')}")
+ if schema_name not in color_map:
+ color_map[schema_name] = {}
+
+ labels = schema.get('labels', [])
+ for i, label in enumerate(labels):
+ if isinstance(label, dict):
+ # F-027: str() so a YAML-bool/number label name (yes/no/1)
+ # doesn't become a non-str dict key (jsonify sort_keys then
+ # raises "'<' not supported between 'str' and 'bool'").
+ label_name = str(label.get('name', label))
+ # Check for inline color definition
+ if 'color' in label and label_name not in color_map[schema_name]:
+ normalized = normalize_color(label['color'])
+ if normalized:
+ color_map[schema_name][label_name] = normalized
+ else:
+ label_name = str(label)
+
+ # Generate default color if not already set
+ if label_name not in color_map[schema_name]:
+ # Try to get from SPAN_COLOR_PALETTE
+ assigned_color = get_span_color(schema_name, label_name)
+ if assigned_color:
+ normalized = normalize_color(assigned_color)
+ if normalized:
+ color_map[schema_name][label_name] = normalized
+ else:
+ # Use hash-based color from default palette
+ color_map[schema_name][label_name] = get_default_label_color(label_name, i)
+
+ logger.debug(f"Final color map: {color_map}")
+ logger.debug("=== GET_COLORS END ===")
+ return jsonify(color_map)
+
+
+@app.route("/api/keyword_highlights/")
+def get_keyword_highlights(instance_id):
+ """
+ Get keyword highlights for a specific instance.
+
+ This endpoint finds all occurrences of admin-defined keywords in the instance text
+ and returns them in the same format as AI keyword suggestions, so they can be
+ displayed using the same visual system (bounding boxes around keywords).
+
+ The endpoint supports randomization for research purposes:
+ - keyword_probability: Probability of showing each matched keyword (default: 1.0)
+ - random_word_probability: Probability of highlighting random words as distractors (default: 0.0)
+
+ Highlights are cached per user+instance to ensure consistency across navigation.
+
+ Colors are assigned based on schema/label to match the span annotation color scheme.
+
+ Returns:
+ JSON with list of keyword matches:
+ {
+ "keywords": [
+ {
+ "label": "Economic",
+ "start": 10,
+ "end": 20,
+ "text": "employment",
+ "reasoning": "Keyword match: employ*",
+ "schema": "Issue-General",
+ "color": "rgba(110, 86, 207, 0.8)",
+ "type": "keyword"
+ },
+ ...
+ ],
+ "instance_id": "item_1",
+ "from_cache": false
+ }
+ """
+ import urllib.parse
+ import random
+ import hashlib
+ import re
+
+ logger.debug(f"=== GET_KEYWORD_HIGHLIGHTS START ===")
+ logger.debug(f"Instance ID: {instance_id}")
+
+ decoded_instance_id = urllib.parse.unquote(instance_id)
+
+ if 'username' not in session:
+ logger.warning("Get keyword highlights without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ username = session.get('username')
+
+ # Get user state for caching
+ user_state = get_user_state(username) if username else None
+
+ # Check for cached state
+ if user_state:
+ cached_state = user_state.get_keyword_highlight_state(instance_id)
+ if not cached_state:
+ # Try with decoded ID
+ cached_state = user_state.get_keyword_highlight_state(decoded_instance_id)
+ if cached_state:
+ logger.debug(f"Returning cached keyword highlights for {instance_id}")
+ return jsonify({
+ "keywords": cached_state.get("highlights", []),
+ "instance_id": instance_id,
+ "from_cache": True
+ })
+
+ # Get settings for randomization
+ settings = get_keyword_highlight_settings()
+ keyword_prob = settings.get('keyword_probability', 1.0)
+ random_word_prob = settings.get('random_word_probability', 0.0)
+ random_word_label = settings.get('random_word_label', 'distractor')
+ random_word_schema = settings.get('random_word_schema', 'keyword')
+
+ logger.debug(f"Keyword highlight settings: keyword_prob={keyword_prob}, random_word_prob={random_word_prob}")
+
+ # Create deterministic seed from username + instance_id for reproducibility
+ seed_str = f"{username}:{instance_id}" if username else instance_id
+ seed = int(hashlib.md5(seed_str.encode()).hexdigest()[:8], 16)
+ rng = random.Random(seed)
+
+ # Check if keyword highlights are enabled
+ keyword_patterns = get_keyword_highlight_patterns()
+
+ # Get the instance text
+ try:
+ item_state_manager = get_item_state_manager()
+ instance = item_state_manager.get_item(instance_id)
+ if not instance:
+ instance = item_state_manager.get_item(decoded_instance_id)
+ if instance:
+ instance_id = decoded_instance_id
+ else:
+ logger.error(f"Instance not found: {instance_id}")
+ return jsonify({"error": "Instance not found"}), 404
+
+ original_text = instance.get_text()
+ logger.debug(f"Instance text length: {len(original_text)}")
+
+ except Exception as e:
+ logger.error(f"Error getting instance text: {e}")
+ return jsonify({"error": f"Instance not found: {instance_id}"}), 404
+
+ # Find all keyword matches in the text
+ keywords = []
+ seen_spans = set() # Track (start, end) to avoid duplicate overlapping matches
+
+ # Track color assignments for keyword labels (schema -> label -> color)
+ keyword_color_counter = 0
+
+ for pattern_info in keyword_patterns:
+ regex = pattern_info['regex']
+ label = pattern_info['label']
+ schema = pattern_info['schema']
+ pattern_str = pattern_info['pattern']
+
+ for match in regex.finditer(original_text):
+ start = match.start()
+ end = match.end()
+ matched_text = match.group()
+
+ # Skip if we already have a match at this exact position
+ span_key = (start, end)
+ if span_key in seen_spans:
+ continue
+
+ # Apply keyword probability filter
+ if rng.random() > keyword_prob:
+ logger.debug(f"Skipping keyword '{matched_text}' due to probability filter")
+ continue
+
+ seen_spans.add(span_key)
+
+ # Get or assign color for this schema/label combination
+ color = get_span_color(schema, label)
+ if not color:
+ # Auto-assign a color from the palette
+ idx = keyword_color_counter % len(SPAN_COLOR_PALETTE)
+ color = SPAN_COLOR_PALETTE[idx]
+ keyword_color_counter += 1
+ # Store it for consistency
+ set_span_color(schema, label, color)
+
+ # Convert RGB tuple string to rgba format for frontend
+ # Color format is "(r, g, b)" - convert to "rgba(r, g, b, 0.8)"
+ if color.startswith("(") and color.endswith(")"):
+ rgba_color = f"rgba{color[:-1]}, 0.8)"
+ else:
+ rgba_color = color
+
+ keywords.append({
+ "label": label,
+ "start": start,
+ "end": end,
+ "text": matched_text,
+ "reasoning": f"Keyword: {pattern_str} โ {label}",
+ "schema": schema,
+ "color": rgba_color,
+ "type": "keyword"
+ })
+
+ # Generate random word highlights (distractors)
+ random_highlights = []
+ if random_word_prob > 0:
+ random_highlights = generate_random_word_highlights(
+ original_text, rng, random_word_prob,
+ random_word_label, random_word_schema,
+ seen_spans
+ )
+ keywords.extend(random_highlights)
+
+ # Sort by start position
+ keywords.sort(key=lambda k: k['start'])
+
+ logger.debug(f"Found {len(keywords)} total highlights ({len(keywords) - len(random_highlights)} keywords, {len(random_highlights)} random)")
+
+ # Cache the state for this user+instance
+ if user_state:
+ user_state.set_keyword_highlight_state(instance_id, {
+ "highlights": keywords,
+ "seed": seed,
+ "settings": {
+ "keyword_probability": keyword_prob,
+ "random_word_probability": random_word_prob
+ }
+ })
+ logger.debug(f"Cached keyword highlight state for {username}:{instance_id}")
+
+ logger.debug("=== GET_KEYWORD_HIGHLIGHTS END ===")
+
+ return jsonify({
+ "keywords": keywords,
+ "instance_id": instance_id,
+ "from_cache": False
+ })
+
+
+def generate_random_word_highlights(text: str, rng, probability: float,
+ label: str, schema: str, excluded_spans: set) -> list:
+ """
+ Generate random word highlights based on probability.
+
+ This function selects random words from the text to highlight as "distractors"
+ to prevent annotators from relying solely on keyword highlights.
+
+ Args:
+ text: The instance text
+ rng: Seeded random.Random instance for reproducibility
+ probability: Probability of selecting each word (0.0-1.0)
+ label: Label for random highlights (e.g., 'distractor')
+ schema: Schema for random highlights (e.g., 'keyword')
+ excluded_spans: Set of (start, end) tuples to avoid (already highlighted)
+
+ Returns:
+ List of highlight dictionaries with keys: label, start, end, text, reasoning, schema, color, type
+ """
+ import re
+
+ highlights = []
+
+ # Find all words (sequences of word characters)
+ word_pattern = re.compile(r'\b\w+\b')
+
+ # Get color for random highlights
+ color = get_span_color(schema, label)
+ if not color:
+ # Use a gray color for distractors by default
+ color = "(156, 163, 175)"
+ set_span_color(schema, label, color)
+
+ # Convert to rgba
+ if color.startswith("(") and color.endswith(")"):
+ color_str = f"rgba{color[:-1]}, 0.6)"
+ else:
+ color_str = color
+
+ for match in word_pattern.finditer(text):
+ start = match.start()
+ end = match.end()
+ word = match.group()
+
+ # Skip if overlaps with existing highlight
+ overlaps = False
+ for ex_start, ex_end in excluded_spans:
+ if start < ex_end and end > ex_start:
+ overlaps = True
+ break
+ if overlaps:
+ continue
+
+ # Skip very short words (1-2 chars) - articles, prepositions, etc.
+ if len(word) <= 2:
+ continue
+
+ # Apply probability
+ if rng.random() < probability:
+ highlights.append({
+ "label": label,
+ "start": start,
+ "end": end,
+ "text": word,
+ "reasoning": "Random selection",
+ "schema": schema,
+ "color": color_str,
+ "type": "random"
+ })
+ excluded_spans.add((start, end))
+
+ return highlights
+
+
+# =============================================================================
+# Behavioral Tracking API Endpoints
+# =============================================================================
+
+@app.route("/api/track_interactions", methods=["POST"])
+def track_interactions():
+ """
+ Receive batched interaction events from the frontend.
+
+ Expected JSON payload:
+ {
+ "instance_id": "...",
+ "events": [...],
+ "focus_time": {"element": ms, ...},
+ "scroll_depth": float
+ }
+ """
+ import time as time_module
+ from potato.interaction_tracking import get_or_create_behavioral_data
+
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ username = session['username']
+ data = request.get_json()
+
+ if not data:
+ return jsonify({"error": "No data provided"}), 400
+
+ instance_id = data.get('instance_id')
+ events = data.get('events', [])
+
+ user_state = get_user_state(username)
+ if not user_state:
+ return jsonify({"error": "User state not found"}), 404
+
+ # Get or create behavioral data for this instance
+ bd = get_or_create_behavioral_data(
+ user_state.instance_id_to_behavioral_data,
+ instance_id
+ )
+
+ # Record server timestamp for each event
+ server_timestamp = time_module.time()
+
+ # Add events
+ for event in events:
+ # Add server timestamp if not present
+ if 'timestamp' not in event or event.get('timestamp') is None:
+ event['timestamp'] = server_timestamp
+
+ # Ensure instance_id is set
+ event['instance_id'] = instance_id
+
+ # Add to behavioral data
+ if hasattr(bd, 'interactions'):
+ from potato.interaction_tracking import InteractionEvent
+ bd.interactions.append(InteractionEvent(
+ event_type=event.get('event_type', 'unknown'),
+ timestamp=event.get('timestamp', server_timestamp),
+ target=event.get('target', ''),
+ instance_id=instance_id,
+ client_timestamp=event.get('client_timestamp'),
+ metadata=event.get('metadata', {}),
+ ))
+
+ # Update focus time if provided
+ focus_time = data.get('focus_time', {})
+ for element, time_ms in focus_time.items():
+ if hasattr(bd, 'update_focus_time'):
+ bd.update_focus_time(element, time_ms)
+ elif hasattr(bd, 'focus_time_by_element'):
+ bd.focus_time_by_element[element] = bd.focus_time_by_element.get(element, 0) + time_ms
+
+ # Update scroll depth
+ if 'scroll_depth' in data:
+ scroll_depth = data['scroll_depth']
+ if hasattr(bd, 'update_scroll_depth'):
+ bd.update_scroll_depth(scroll_depth)
+ elif hasattr(bd, 'scroll_depth_max'):
+ bd.scroll_depth_max = max(bd.scroll_depth_max, scroll_depth)
+
+ return jsonify({"status": "ok", "events_recorded": len(events)})
+
+
+@app.route("/api/track_ai_usage", methods=["POST"])
+def track_ai_usage():
+ """
+ Track AI assistance request, response, and user decisions.
+
+ Expected JSON payload:
+ {
+ "instance_id": "...",
+ "schema_name": "...",
+ "event_type": "request" | "response" | "accept" | "reject",
+ "suggestions": [...], # for response events
+ "accepted_value": "..." # for accept events
+ }
+ """
+ import time as time_module
+ from potato.interaction_tracking import get_or_create_behavioral_data, AIUsageEvent
+
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ username = session['username']
+ data = request.get_json()
+
+ if not data:
+ return jsonify({"error": "No data provided"}), 400
+
+ instance_id = data.get('instance_id')
+ schema_name = data.get('schema_name')
+ event_type = data.get('event_type') # 'request', 'response', 'accept', 'reject'
+
+ if not instance_id or not schema_name or not event_type:
+ return jsonify({"error": "Missing required fields"}), 400
+
+ user_state = get_user_state(username)
+ if not user_state:
+ return jsonify({"error": "User state not found"}), 404
+
+ # Get or create behavioral data
+ bd = get_or_create_behavioral_data(
+ user_state.instance_id_to_behavioral_data,
+ instance_id
+ )
+
+ timestamp = time_module.time()
+
+ if event_type == 'request':
+ # Create new AI usage event
+ ai_event = AIUsageEvent(
+ request_timestamp=timestamp,
+ schema_name=schema_name,
+ )
+ if hasattr(bd, 'ai_usage'):
+ bd.ai_usage.append(ai_event)
+
+ elif event_type == 'response':
+ suggestions = data.get('suggestions', [])
+ # Update the most recent AI event for this schema
+ if hasattr(bd, 'ai_usage'):
+ for ai_event in reversed(bd.ai_usage):
+ event_schema = ai_event.schema_name if hasattr(ai_event, 'schema_name') else ai_event.get('schema_name')
+ event_response = ai_event.response_timestamp if hasattr(ai_event, 'response_timestamp') else ai_event.get('response_timestamp')
+ if event_schema == schema_name and not event_response:
+ if hasattr(ai_event, 'response_timestamp'):
+ ai_event.response_timestamp = timestamp
+ ai_event.suggestions_shown = suggestions
+ else:
+ ai_event['response_timestamp'] = timestamp
+ ai_event['suggestions_shown'] = suggestions
+ break
+
+ elif event_type in ('accept', 'reject'):
+ accepted_value = data.get('accepted_value') if event_type == 'accept' else None
+ # Update the most recent AI event for this schema
+ if hasattr(bd, 'ai_usage'):
+ for ai_event in reversed(bd.ai_usage):
+ event_schema = ai_event.schema_name if hasattr(ai_event, 'schema_name') else ai_event.get('schema_name')
+ event_response = ai_event.response_timestamp if hasattr(ai_event, 'response_timestamp') else ai_event.get('response_timestamp')
+ if event_schema == schema_name and event_response:
+ if hasattr(ai_event, 'suggestion_accepted'):
+ ai_event.suggestion_accepted = accepted_value
+ ai_event.time_to_decision_ms = int((timestamp - ai_event.response_timestamp) * 1000)
+ else:
+ ai_event['suggestion_accepted'] = accepted_value
+ ai_event['time_to_decision_ms'] = int((timestamp - ai_event['response_timestamp']) * 1000)
+ break
+
+ return jsonify({"status": "ok", "event_type": event_type})
+
+
+@app.route("/api/track_annotation_change", methods=["POST"])
+def track_annotation_change():
+ """
+ Track annotation changes from the frontend.
+
+ Expected JSON payload:
+ {
+ "instance_id": "...",
+ "schema_name": "...",
+ "label_name": "...",
+ "action": "select" | "deselect" | "update" | "clear",
+ "old_value": ...,
+ "new_value": ...,
+ "source": "user" | "ai_accept" | "keyboard" | "prefill"
+ }
+ """
+ import time as time_module
+ from potato.interaction_tracking import get_or_create_behavioral_data, AnnotationChange
+
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ username = session['username']
+ data = request.get_json()
+
+ if not data:
+ return jsonify({"error": "No data provided"}), 400
+
+ instance_id = data.get('instance_id')
+ schema_name = data.get('schema_name')
+ action = data.get('action')
+
+ if not instance_id or not schema_name or not action:
+ return jsonify({"error": "Missing required fields"}), 400
+
+ user_state = get_user_state(username)
+ if not user_state:
+ return jsonify({"error": "User state not found"}), 404
+
+ # Get or create behavioral data
+ bd = get_or_create_behavioral_data(
+ user_state.instance_id_to_behavioral_data,
+ instance_id
+ )
+
+ # Create annotation change record
+ change = AnnotationChange(
+ timestamp=time_module.time(),
+ schema_name=schema_name,
+ label_name=data.get('label_name'),
+ action=action,
+ old_value=data.get('old_value'),
+ new_value=data.get('new_value'),
+ source=data.get('source', 'user'),
+ )
+
+ if hasattr(bd, 'annotation_changes'):
+ bd.annotation_changes.append(change)
+
+ return jsonify({"status": "ok"})
+
+
+@app.route("/api/behavioral_data/", methods=["GET"])
+def get_behavioral_data(instance_id):
+ """
+ Get behavioral data for a specific instance.
+ Useful for debugging and analysis.
+ """
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ username = session['username']
+ user_state = get_user_state(username)
+
+ if not user_state:
+ return jsonify({"error": "User state not found"}), 404
+
+ bd = user_state.instance_id_to_behavioral_data.get(instance_id)
+
+ if not bd:
+ return jsonify({"error": "No behavioral data for instance"}), 404
+
+ if hasattr(bd, 'to_dict'):
+ return jsonify(bd.to_dict())
+ elif isinstance(bd, dict):
+ return jsonify(bd)
+ else:
+ return jsonify({"error": "Invalid behavioral data format"}), 500
+
+
+@app.route("/api/schemas")
+def get_annotation_schemas():
+ """
+ Return the annotation schema information for all annotation types.
+ This provides the schema names, types, and their labels to the frontend
+ and API consumers (like the user simulator).
+ """
+ logger.debug("=== GET_ANNOTATION_SCHEMAS START ===")
+
+ schemas = {}
+ annotation_scheme = config.get('annotation_scheme') or config.get('annotation_schemes')
+
+ if annotation_scheme:
+ # Helper function to extract labels from a schema
+ def extract_labels(schema):
+ labels = []
+ for label in schema.get('labels', []):
+ if isinstance(label, dict):
+ labels.append(label.get('name', str(label)))
+ else:
+ labels.append(str(label))
+ return labels
+
+ # Keys we surface via dedicated top-level fields above, plus
+ # internal/runtime keys that should never round-trip to API
+ # consumers. Anything else configured on the schema (likert.size,
+ # slider.min_value, process_reward.mode, code_review.verdict_options,
+ # span.labels, custom plugin keys, etc.) passes through verbatim.
+ _RESERVED_SCHEMA_KEYS = frozenset({
+ "name",
+ "annotation_type",
+ "type",
+ "labels",
+ "description",
+ # Runtime / server-side internals
+ "annotation_id",
+ "sequential_key_binding",
+ })
+
+ # Helper function to process a single schema
+ def process_schema(schema, schema_name=None):
+ name = schema_name or schema.get('name', 'unknown')
+ schema_type = schema.get('annotation_type') or schema.get('type', 'unknown')
+
+ schema_info = {
+ 'name': name,
+ 'description': schema.get('description', ''),
+ 'labels': extract_labels(schema),
+ 'type': schema_type
+ }
+
+ # Pass through every other configured key unchanged. This keeps
+ # the API forward-compatible for any schema type (built-in or
+ # plugin) that carries type-specific config (e.g. likert.size,
+ # process_reward.mode, code_review.verdict_options). Underscore-
+ # prefixed keys are treated as internal and skipped.
+ for key, value in schema.items():
+ if key in _RESERVED_SCHEMA_KEYS or key.startswith('_'):
+ continue
+ schema_info.setdefault(key, value)
+
+ return schema_info
+
+ # If dict (new style), iterate items
+ if isinstance(annotation_scheme, dict):
+ for schema_name, schema in annotation_scheme.items():
+ schemas[schema_name] = process_schema(schema, schema_name)
+ # If list (legacy style), iterate list
+ elif isinstance(annotation_scheme, list):
+ for schema in annotation_scheme:
+ schema_name = schema.get('name', 'unknown')
+ schemas[schema_name] = process_schema(schema)
+
+ logger.debug(f"Found schemas: {schemas}")
+ logger.debug("=== GET_ANNOTATION_SCHEMAS END ===")
+ return jsonify(schemas)
+
+@app.route("/api/spans//clear", methods=["POST"])
+def clear_span_annotations(instance_id):
+ """
+ Clear all span annotations for a specific instance and user.
+ This is useful for debugging and fixing persistent overlay issues.
+ """
+ logger.debug(f"=== CLEAR_SPAN_ANNOTATIONS START ===")
+ logger.debug(f"Instance ID: {instance_id}")
+
+ if 'username' not in session:
+ logger.warning("Clear span annotations without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ username = session['username']
+ logger.debug(f"Username: {username}")
+
+ try:
+ user_state = get_user_state(username)
+ if not user_state:
+ logger.error(f"User state not found for user: {username}")
+ return jsonify({"error": "User state not found"}), 404
+
+ # Normalize instance_id to string
+ instance_id = str(instance_id)
+
+ # Check if instance has span annotations
+ if hasattr(user_state, 'instance_id_to_span_to_value'):
+ if instance_id in user_state.instance_id_to_span_to_value:
+ spans_before = len(user_state.instance_id_to_span_to_value[instance_id])
+ logger.debug(f"Found {spans_before} spans for instance {instance_id}")
+
+ # Clear the spans
+ del user_state.instance_id_to_span_to_value[instance_id]
+ logger.debug(f"Cleared {spans_before} spans for instance {instance_id}")
+
+ return jsonify({
+ "status": "success",
+ "message": f"Cleared {spans_before} span annotations for instance {instance_id}",
+ "spans_cleared": spans_before
+ })
+ else:
+ logger.debug(f"No spans found for instance {instance_id}")
+ return jsonify({
+ "status": "success",
+ "message": f"No span annotations found for instance {instance_id}",
+ "spans_cleared": 0
+ })
+ else:
+ logger.debug("User state has no span annotations")
+ return jsonify({
+ "status": "success",
+ "message": "User state has no span annotations",
+ "spans_cleared": 0
+ })
+
+ except Exception as e:
+ logger.error(f"Error clearing span annotations: {e}")
+ return jsonify({"error": "Failed to clear span annotations"}), 500
+
+ finally:
+ logger.debug(f"=== CLEAR_SPAN_ANNOTATIONS END ===")
+
+
+@app.route("/api/links/")
+def get_link_annotations(instance_id):
+ """
+ Get link annotations (span relationships) for a specific instance.
+
+ Returns:
+ JSON with link annotations for the instance.
+ """
+ logger.debug(f"=== GET_LINK_ANNOTATIONS START ===")
+ logger.debug(f"Instance ID: {instance_id}")
+
+ if 'username' not in session:
+ logger.warning("Get link annotations without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ username = session['username']
+ logger.debug(f"Username: {username}")
+
+ try:
+ user_state = get_user_state(username)
+ if not user_state:
+ logger.error(f"User state not found for user: {username}")
+ return jsonify({"error": "User state not found"}), 404
+
+ # Normalize instance_id to string
+ instance_id = str(instance_id)
+
+ # Get link annotations for this instance
+ links = user_state.get_link_annotations(instance_id)
+
+ # Convert to serializable format
+ links_data = []
+ for link_id, link in links.items():
+ links_data.append(link.to_dict())
+
+ logger.debug(f"Found {len(links_data)} link annotations for instance {instance_id}")
+
+ return jsonify({
+ "status": "success",
+ "instance_id": instance_id,
+ "links": links_data
+ })
+
+ except Exception as e:
+ logger.error(f"Error getting link annotations: {e}")
+ return jsonify({"error": "Failed to get link annotations"}), 500
+
+ finally:
+ logger.debug(f"=== GET_LINK_ANNOTATIONS END ===")
+
+
+@app.route("/api/links//", methods=["DELETE"])
+def delete_link_annotation(instance_id, link_id):
+ """
+ Delete a specific link annotation.
+
+ Args:
+ instance_id: The instance ID
+ link_id: The link ID to delete
+
+ Returns:
+ JSON with success/failure status.
+ """
+ logger.debug(f"=== DELETE_LINK_ANNOTATION START ===")
+ logger.debug(f"Instance ID: {instance_id}, Link ID: {link_id}")
+
+ if 'username' not in session:
+ logger.warning("Delete link annotation without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ username = session['username']
+ logger.debug(f"Username: {username}")
+
+ try:
+ user_state = get_user_state(username)
+ if not user_state:
+ logger.error(f"User state not found for user: {username}")
+ return jsonify({"error": "User state not found"}), 404
+
+ # Normalize instance_id to string
+ instance_id = str(instance_id)
+
+ # Try to remove the link
+ success = user_state.remove_link_annotation(instance_id, link_id)
+
+ if success:
+ logger.debug(f"Deleted link annotation: {link_id} from instance {instance_id}")
+ return jsonify({
+ "status": "success",
+ "message": f"Link {link_id} deleted successfully"
+ })
+ else:
+ logger.warning(f"Link not found: {link_id} in instance {instance_id}")
+ return jsonify({
+ "status": "error",
+ "message": f"Link {link_id} not found"
+ }), 404
+
+ except Exception as e:
+ logger.error(f"Error deleting link annotation: {e}")
+ return jsonify({"error": "Failed to delete link annotation"}), 500
+
+ finally:
+ logger.debug(f"=== DELETE_LINK_ANNOTATION END ===")
+
+
+# ============================================================================
+# Event Annotation API Routes
+# ============================================================================
+
+@app.route("/api/events/")
+def get_event_annotations(instance_id):
+ """
+ Get event annotations for a specific instance.
+
+ Returns:
+ JSON with event annotations for the instance.
+ """
+ logger.debug(f"=== GET_EVENT_ANNOTATIONS START ===")
+ logger.debug(f"Instance ID: {instance_id}")
+
+ if 'username' not in session:
+ logger.warning("Get event annotations without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ username = session['username']
+ logger.debug(f"Username: {username}")
+
+ try:
+ user_state = get_user_state(username)
+ if not user_state:
+ logger.error(f"User state not found for user: {username}")
+ return jsonify({"error": "User state not found"}), 404
+
+ # Normalize instance_id to string
+ instance_id = str(instance_id)
+
+ # Verify the user is assigned this instance
+ if instance_id not in user_state.get_assigned_instance_ids():
+ logger.warning(f"User {username} not assigned to instance {instance_id}")
+ return jsonify({"error": "Instance not assigned to user"}), 403
+
+ # Get event annotations for this instance
+ events = user_state.get_event_annotations(instance_id)
+
+ # Convert to serializable format
+ events_data = []
+ for event_id, event in events.items():
+ event_dict = event.to_dict()
+ events_data.append(event_dict)
+ logger.debug(f" Event ID: {event_id}, data: {event_dict}")
+
+ logger.debug(f"Found {len(events_data)} event annotations for instance {instance_id}")
+ logger.debug(f"Event IDs in storage: {list(events.keys())}")
+
+ return jsonify({
+ "status": "success",
+ "instance_id": instance_id,
+ "events": events_data
+ })
+
+ except Exception as e:
+ logger.error(f"Error getting event annotations: {e}")
+ return jsonify({"error": "Failed to get event annotations"}), 500
+
+ finally:
+ logger.debug(f"=== GET_EVENT_ANNOTATIONS END ===")
+
+
+@app.route("/api/events//", methods=["DELETE"])
+def delete_event_annotation(instance_id, event_id):
+ """
+ Delete a specific event annotation.
+
+ Args:
+ instance_id: The instance ID
+ event_id: The event ID to delete
+
+ Returns:
+ JSON with success/failure status.
+ """
+ logger.debug(f"=== DELETE_EVENT_ANNOTATION START ===")
+ logger.debug(f"Instance ID: {instance_id}, Event ID: {event_id}")
+
+ if 'username' not in session:
+ logger.warning("Delete event annotation without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ username = session['username']
+ logger.debug(f"Username: {username}")
+
+ try:
+ user_state = get_user_state(username)
+ if not user_state:
+ logger.error(f"User state not found for user: {username}")
+ return jsonify({"error": "User state not found"}), 404
+
+ # Normalize instance_id to string
+ instance_id = str(instance_id)
+
+ # Verify the user is assigned this instance
+ if instance_id not in user_state.get_assigned_instance_ids():
+ logger.warning(f"User {username} not assigned to instance {instance_id}")
+ return jsonify({"error": "Instance not assigned to user"}), 403
+
+ # Try to remove the event
+ success = user_state.remove_event_annotation(instance_id, event_id)
+
+ if success:
+ logger.debug(f"Deleted event annotation: {event_id} from instance {instance_id}")
+ return jsonify({
+ "status": "success",
+ "message": f"Event {event_id} deleted successfully"
+ })
+ else:
+ logger.warning(f"Event not found: {event_id} in instance {instance_id}")
+ return jsonify({
+ "status": "error",
+ "message": f"Event {event_id} not found"
+ }), 404
+
+ except Exception as e:
+ logger.error(f"Error deleting event annotation: {e}")
+ return jsonify({"error": "Failed to delete event annotation"}), 500
+
+ finally:
+ logger.debug(f"=== DELETE_EVENT_ANNOTATION END ===")
+
+
+# ============================================================================
+# Entity Linking API Routes
+# ============================================================================
+
+@app.route("/api/entity_linking/search")
+def entity_linking_search():
+ """
+ Search a knowledge base for entities matching a query.
+
+ Query parameters:
+ q: Search query string (required)
+ kb: Knowledge base name (required, e.g., "wikidata", "umls")
+ limit: Maximum number of results (default: 10)
+
+ Returns:
+ JSON with list of matching entities.
+ """
+ logger.debug(f"=== ENTITY_LINKING_SEARCH START ===")
+
+ if 'username' not in session:
+ logger.warning("Entity linking search without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ query = request.args.get('q', '').strip()
+ kb_name = request.args.get('kb', '').strip()
+ limit = request.args.get('limit', 10, type=int)
+
+ logger.debug(f"Query: '{query}', KB: '{kb_name}', Limit: {limit}")
+
+ if not query:
+ return jsonify({"error": "Search query 'q' is required"}), 400
+
+ if not kb_name:
+ return jsonify({"error": "Knowledge base 'kb' is required"}), 400
+
+ try:
+ from potato.knowledge_base import get_kb_manager
+
+ kb_manager = get_kb_manager()
+ results = kb_manager.search(query, kb_name, limit=limit)
+
+ # Convert to serializable format
+ entities = [entity.to_dict() for entity in results]
+
+ logger.debug(f"Found {len(entities)} entities for query '{query}'")
+
+ return jsonify({
+ "status": "success",
+ "query": query,
+ "kb": kb_name,
+ "results": entities
+ })
+
+ except Exception as e:
+ logger.error(f"Error in entity linking search: {e}")
+ return jsonify({"error": "Search failed"}), 500
+
+ finally:
+ logger.debug(f"=== ENTITY_LINKING_SEARCH END ===")
+
+
+@app.route("/api/entity_linking/entity//")
+def entity_linking_get_entity(kb_name, entity_id):
+ """
+ Get detailed information about a specific entity.
+
+ Args:
+ kb_name: Knowledge base name (e.g., "wikidata", "umls")
+ entity_id: Entity ID within the knowledge base (e.g., "Q937")
+
+ Returns:
+ JSON with entity details.
+ """
+ logger.debug(f"=== ENTITY_LINKING_GET_ENTITY START ===")
+ logger.debug(f"KB: {kb_name}, Entity ID: {entity_id}")
+
+ if 'username' not in session:
+ logger.warning("Entity linking get_entity without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ try:
+ from potato.knowledge_base import get_kb_manager
+
+ kb_manager = get_kb_manager()
+ client = kb_manager.get_client(kb_name)
+
+ if not client:
+ return jsonify({"error": f"Knowledge base '{kb_name}' not configured"}), 404
+
+ entity = client.get_entity(entity_id)
+
+ if not entity:
+ return jsonify({"error": f"Entity '{entity_id}' not found"}), 404
+
+ logger.debug(f"Found entity: {entity.label}")
+
+ return jsonify({
+ "status": "success",
+ "entity": entity.to_dict()
+ })
+
+ except Exception as e:
+ logger.error(f"Error getting entity: {e}")
+ return jsonify({"error": "Failed to get entity"}), 500
+
+ finally:
+ logger.debug(f"=== ENTITY_LINKING_GET_ENTITY END ===")
+
+
+@app.route("/api/entity_linking/configured_kbs")
+def entity_linking_configured_kbs():
+ """
+ Get list of configured knowledge bases.
+
+ Returns:
+ JSON with list of available knowledge base names and types.
+ """
+ logger.debug(f"=== ENTITY_LINKING_CONFIGURED_KBS START ===")
+
+ if 'username' not in session:
+ logger.warning("Entity linking configured_kbs without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ try:
+ from potato.knowledge_base import get_kb_manager
+
+ kb_manager = get_kb_manager()
+ kb_names = kb_manager.list_clients()
+
+ # Get config info for each KB
+ kbs = []
+ for name in kb_names:
+ config = kb_manager.get_config(name)
+ if config:
+ kbs.append({
+ "name": name,
+ "type": config.kb_type,
+ "language": config.language
+ })
+
+ logger.debug(f"Found {len(kbs)} configured knowledge bases")
+
+ return jsonify({
+ "status": "success",
+ "knowledge_bases": kbs
+ })
+
+ except Exception as e:
+ logger.error(f"Error getting configured KBs: {e}")
+ return jsonify({"error": "Failed to get configured KBs"}), 500
+
+ finally:
+ logger.debug(f"=== ENTITY_LINKING_CONFIGURED_KBS END ===")
+
+
+@app.route("/api/entity_linking/update_span", methods=["POST"])
+def entity_linking_update_span():
+ """
+ Update a span annotation with entity linking information.
+
+ Request body:
+ instance_id: Instance ID
+ span_id: Span annotation ID
+ kb_id: Knowledge base entity ID
+ kb_source: Knowledge base source name
+ kb_label: Human-readable entity label
+
+ Returns:
+ JSON with success/failure status.
+ """
+ logger.debug(f"=== ENTITY_LINKING_UPDATE_SPAN START ===")
+
+ if 'username' not in session:
+ logger.warning("Entity linking update_span without active session")
+ return jsonify({"error": "No active session"}), 401
+
+ username = session['username']
+
+ try:
+ data = request.json
+ instance_id = data.get('instance_id')
+ span_id = data.get('span_id')
+ kb_id = data.get('kb_id')
+ kb_source = data.get('kb_source')
+ kb_label = data.get('kb_label')
+
+ logger.debug(f"Updating span {span_id} with KB: {kb_source}:{kb_id}")
+
+ if not instance_id or not span_id:
+ return jsonify({"error": "instance_id and span_id are required"}), 400
+
+ # Validate string types and enforce length limits
+ MAX_FIELD_LEN = 1024
+ for field_name, field_val in [("span_id", span_id), ("kb_id", kb_id),
+ ("kb_source", kb_source), ("kb_label", kb_label)]:
+ if field_val is not None and (not isinstance(field_val, str) or len(field_val) > MAX_FIELD_LEN):
+ return jsonify({"error": f"Invalid {field_name}"}), 400
+
+ user_state = get_user_state(username)
+ if not user_state:
+ return jsonify({"error": "User state not found"}), 404
+
+ # Get span annotations for this instance
+ span_annotations = user_state.get_span_annotations(str(instance_id))
+
+ # Debug: Log all existing span IDs
+ existing_ids = []
+ for span_key, span in span_annotations.items():
+ if hasattr(span_key, 'get_id'):
+ existing_ids.append(span_key.get_id())
+ elif isinstance(span_key, dict):
+ existing_ids.append(span_key.get('id', 'no-id'))
+ logger.debug(f"Looking for span_id={span_id}, existing IDs: {existing_ids}")
+
+ # Find the span with matching ID
+ # Note: span_key is the SpanAnnotation object, span is the value
+ updated = False
+ for span_key, span_value in span_annotations.items():
+ if hasattr(span_key, 'get_id') and span_key.get_id() == span_id:
+ # Update the span's KB link
+ span_key.set_entity_link(kb_id, kb_source, kb_label)
+ updated = True
+ logger.debug(f"Updated span {span_id} with entity link")
+ break
+ elif isinstance(span_key, dict) and span_key.get('id') == span_id:
+ span_key['kb_id'] = kb_id
+ span_key['kb_source'] = kb_source
+ span_key['kb_label'] = kb_label
+ updated = True
+ logger.debug(f"Updated span dict {span_id} with entity link")
+ break
+
+ if not updated:
+ return jsonify({"error": f"Span {span_id} not found"}), 404
+
+ return jsonify({
+ "status": "success",
+ "message": f"Span {span_id} linked to {kb_source}:{kb_id}"
+ })
+
+ except Exception as e:
+ logger.error(f"Error updating span with entity link: {e}")
+ return jsonify({"error": "Failed to update span"}), 500
+
+ finally:
+ logger.debug(f"=== ENTITY_LINKING_UPDATE_SPAN END ===")
+
+
+@app.route("/api/waveform/")
+def get_waveform_data(cache_key):
+ """
+ Serve pre-computed waveform data for audio annotation.
+
+ This endpoint serves .dat waveform files generated by the WaveformService.
+ The cache_key is an MD5 hash of the audio file path.
+
+ Args:
+ cache_key: The MD5 hash identifying the cached waveform file
+
+ Returns:
+ The binary waveform data file, or an error response
+ """
+ logger.debug(f"=== GET_WAVEFORM_DATA START ===")
+ logger.debug(f"Cache key: {cache_key}")
+
+ try:
+ # Import waveform service
+ from potato.server_utils.waveform_service import get_waveform_service
+
+ waveform_service = get_waveform_service()
+ if not waveform_service:
+ logger.warning("WaveformService not initialized")
+ return jsonify({"error": "Waveform service not available"}), 503
+
+ # Construct the cache file path
+ cache_path = os.path.join(waveform_service.cache_dir, f"{cache_key}.dat")
+
+ if not os.path.exists(cache_path):
+ logger.warning(f"Waveform file not found: {cache_path}")
+ return jsonify({"error": "Waveform data not found"}), 404
+
+ # Serve the waveform file
+ from flask import send_file
+ logger.debug(f"Serving waveform file: {cache_path}")
+ return send_file(
+ cache_path,
+ mimetype='application/octet-stream',
+ as_attachment=False
+ )
+
+ except Exception as e:
+ logger.error(f"Error serving waveform data: {e}")
+ return jsonify({"error": "Failed to serve waveform data"}), 500
+
+ finally:
+ logger.debug(f"=== GET_WAVEFORM_DATA END ===")
+
+
+@app.route("/api/waveform/generate", methods=["POST"])
+def generate_waveform():
+ """
+ Generate waveform data for an audio file.
+
+ This endpoint triggers waveform generation for a given audio URL.
+ It can be called by the frontend when the audio loads.
+
+ Request body:
+ audio_url: URL or path of the audio file
+
+ Returns:
+ JSON with the waveform URL or error message
+ """
+ logger.debug(f"=== GENERATE_WAVEFORM START ===")
+
+ try:
+ data = request.get_json()
+ if not data or 'audio_url' not in data:
+ return jsonify({"error": "audio_url is required"}), 400
+
+ audio_url = data['audio_url']
+ logger.debug(f"Generating waveform for: {audio_url}")
+
+ # Import waveform service
+ from potato.server_utils.waveform_service import get_waveform_service
+
+ waveform_service = get_waveform_service()
+ if not waveform_service:
+ logger.warning("WaveformService not initialized")
+ return jsonify({
+ "error": "Waveform service not available",
+ "use_client_fallback": True
+ }), 503
+
+ # Check if we should use client-side fallback
+ if not waveform_service.is_available:
+ return jsonify({
+ "use_client_fallback": True,
+ "message": "Server-side waveform generation not available"
+ })
+
+ # Get or generate waveform
+ waveform_path = waveform_service.get_waveform_path(audio_url)
+ if waveform_path:
+ waveform_url = waveform_service.get_waveform_url(audio_url)
+ logger.debug(f"Waveform available at: {waveform_url}")
+ return jsonify({
+ "waveform_url": waveform_url,
+ "use_client_fallback": False
+ })
+ else:
+ logger.warning(f"Failed to generate waveform for: {audio_url}")
+ return jsonify({
+ "use_client_fallback": True,
+ "message": "Waveform generation failed, use client-side fallback"
+ })
+
+ except Exception as e:
+ logger.error(f"Error generating waveform: {e}")
+ return jsonify({
+ "error": f"Failed to generate waveform: {str(e)}",
+ "use_client_fallback": True
+ }), 500
+
+ finally:
+ logger.debug(f"=== GENERATE_WAVEFORM END ===")
+
+
+@app.route("/api/video/metadata", methods=["POST"])
+def get_video_metadata():
+ """
+ Get metadata for a video file.
+
+ This endpoint returns video metadata including duration, FPS, and resolution.
+ It can be called by the frontend when a video loads to get frame-accurate
+ timing information for video annotation.
+
+ Request body:
+ video_url: URL or path of the video file
+
+ Returns:
+ JSON with video metadata:
+ - duration: Video duration in seconds
+ - fps: Frames per second (estimated if not available)
+ - width: Video width in pixels
+ - height: Video height in pixels
+ - frame_count: Total number of frames (if calculable)
+ """
+ logger.debug("=== GET_VIDEO_METADATA START ===")
+
+ try:
+ data = request.get_json()
+ if not data or 'video_url' not in data:
+ return jsonify({"error": "Missing video_url parameter"}), 400
+
+ video_url = data['video_url']
+ logger.debug(f"Video URL: {video_url}")
+
+ # For now, return a basic response that the frontend can use
+ # The actual video metadata will be determined by the browser
+ # since we don't have ffprobe installed by default
+ return jsonify({
+ "status": "ok",
+ "message": "Video metadata should be retrieved client-side",
+ "video_url": video_url,
+ "use_client_detection": True
+ })
+
+ except Exception as e:
+ logger.error(f"Error getting video metadata: {e}")
+ return jsonify({"error": str(e)}), 500
+
+ finally:
+ logger.debug("=== GET_VIDEO_METADATA END ===")
+
+
+@app.route("/api/video/waveform/generate", methods=["POST"])
+def generate_video_waveform():
+ """
+ Generate waveform data from a video file's audio track.
+
+ This endpoint triggers waveform generation for a video's audio track.
+ It reuses the existing audio waveform generation infrastructure.
+
+ Request body:
+ video_url: URL or path of the video file
+
+ Returns:
+ JSON with waveform status and cache key (if successful)
+ """
+ logger.debug("=== GENERATE_VIDEO_WAVEFORM START ===")
+
+ try:
+ data = request.get_json()
+ if not data or 'video_url' not in data:
+ return jsonify({"error": "Missing video_url parameter"}), 400
+
+ video_url = data['video_url']
+ logger.debug(f"Video URL for waveform: {video_url}")
+
+ # Try to generate waveform using the existing WaveformService
+ try:
+ from potato.server_utils.waveform_service import WaveformService
+ waveform_service = WaveformService()
+
+ # Generate waveform from video (will extract audio track)
+ result = waveform_service.generate_waveform(video_url)
+
+ if result.get('status') == 'ready':
+ return jsonify({
+ "status": "ready",
+ "waveform_url": result.get('waveform_url'),
+ "cache_key": result.get('cache_key')
+ })
+ else:
+ return jsonify({
+ "status": result.get('status', 'pending'),
+ "message": result.get('message', 'Waveform generation in progress')
+ })
+
+ except ImportError:
+ logger.warning("WaveformService not available for video waveform generation")
+ return jsonify({
+ "status": "unavailable",
+ "message": "Waveform service not available",
+ "use_client_fallback": True
+ })
+
+ except Exception as e:
+ logger.error(f"Error generating video waveform: {e}")
+ return jsonify({
+ "error": str(e),
+ "use_client_fallback": True
+ }), 500
+
+ finally:
+ logger.debug("=== GENERATE_VIDEO_WAVEFORM END ===")
+
+
+@app.route("/api/audio/proxy")
+def audio_proxy():
+ """
+ Proxy endpoint for fetching external audio files with Range request support.
+
+ This endpoint fetches audio files from external URLs and returns them
+ with proper headers, bypassing CORS restrictions that prevent the browser
+ from directly accessing external audio files for waveform generation.
+
+ Supports HTTP Range requests to enable seeking in audio files.
+
+ Query parameters:
+ url: The external audio URL to fetch
+
+ Returns:
+ The audio file with appropriate Content-Type header
+ """
+ import requests as req
+
+ audio_url = request.args.get('url')
+ if not audio_url:
+ return jsonify({"error": "Missing url parameter"}), 400
+
+ # Validate URL (basic security check)
+ if not audio_url.startswith(('http://', 'https://')):
+ return jsonify({"error": "Invalid URL - must be http or https"}), 400
+
+ try:
+ # Forward any Range header from the client to the upstream server
+ headers = {}
+ if 'Range' in request.headers:
+ headers['Range'] = request.headers['Range']
+
+ # Fetch the audio file
+ response = req.get(audio_url, headers=headers, stream=True, timeout=30)
+ response.raise_for_status()
+
+ # Get content type from response or default to audio/mpeg
+ content_type = response.headers.get('Content-Type', 'audio/mpeg')
+ content_length = response.headers.get('Content-Length')
+
+ # Create response with the audio data
+ flask_response = make_response(response.content)
+ flask_response.headers['Content-Type'] = content_type
+ flask_response.headers['Access-Control-Allow-Origin'] = '*'
+ flask_response.headers['Cache-Control'] = 'public, max-age=3600'
+
+ # Add headers to support Range requests (seeking)
+ flask_response.headers['Accept-Ranges'] = 'bytes'
+
+ if content_length:
+ flask_response.headers['Content-Length'] = content_length
+
+ # If the upstream returned a 206 Partial Content, pass that through
+ if response.status_code == 206:
+ flask_response.status_code = 206
+ if 'Content-Range' in response.headers:
+ flask_response.headers['Content-Range'] = response.headers['Content-Range']
+
+ return flask_response
+
+ except req.exceptions.Timeout:
+ logger.error(f"Timeout fetching audio: {audio_url}")
+ return jsonify({"error": "Request timed out"}), 504
+ except req.exceptions.RequestException as e:
+ logger.error(f"Error fetching audio {audio_url}: {e}")
+ return jsonify({"error": "Failed to fetch audio"}), 502
+
+
+@app.route("/api/ai_assistant", methods=["GET"])
+def ai_assistant():
+ annotation_id_str = request.args.get("annotationId")
+ logger.debug(f"[AI Assistant] Request for annotationId={annotation_id_str}")
+
+ # Handle null/None/invalid annotation IDs
+ if annotation_id_str is None or annotation_id_str == "null" or annotation_id_str == "":
+ logger.debug("[AI Assistant] Invalid annotation ID - returning empty")
+ return jsonify({"html": "", "error": None})
+
+ try:
+ annotation_id = int(annotation_id_str)
+ except (ValueError, TypeError):
+ logger.debug("[AI Assistant] Failed to parse annotation ID")
+ return jsonify({"html": "", "error": None})
+
+ # Check if annotation_id is valid
+ if annotation_id < 0 or annotation_id >= len(config.get("annotation_schemes", [])):
+ logger.debug(f"[AI Assistant] annotation_id {annotation_id} out of range")
+ return jsonify({"html": "", "error": None})
+
+ username = session['username']
+ user_state = get_user_state(username)
+ instance = user_state.get_current_instance_index()
+ annotation_type = config["annotation_schemes"][annotation_id]["annotation_type"]
+
+ result = generate_ai_help_html(instance, annotation_id, annotation_type)
+ logger.debug(f"[AI Assistant] Result for instance={instance}, annotation_id={annotation_id}, type={annotation_type}: '{result[:100] if result else 'empty'}...'")
+ return result
+
+
+def admin_reset_password():
+ """Admin API to reset a user's password.
+
+ Requires X-API-Key header. Takes JSON body with username and new_password.
+ """
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({"error": "Unauthorized - valid API key required"}), 403
+
+ data = request.get_json()
+ if not data:
+ return jsonify({"error": "JSON body required"}), 400
+
+ username = data.get("username")
+ new_password = data.get("new_password")
+
+ if not username or not new_password:
+ return jsonify({"error": "username and new_password are required"}), 400
+
+ user_authenticator = UserAuthenticator.get_instance()
+ if not user_authenticator.is_valid_username(username):
+ return jsonify({"error": f"User '{username}' does not exist"}), 404
+
+ if user_authenticator.update_password(username, new_password):
+ user_authenticator.save_user_config()
+ return jsonify({"status": "success", "message": f"Password reset for '{username}'"})
+ else:
+ return jsonify({"error": "Failed to reset password"}), 500
+
+
+def admin_create_reset_token():
+ """Admin API to generate a password reset token for a user.
+
+ Requires X-API-Key header. Takes JSON body with username and optional ttl_hours.
+ Returns the reset link and token.
+ """
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({"error": "Unauthorized - valid API key required"}), 403
+
+ data = request.get_json()
+ if not data:
+ return jsonify({"error": "JSON body required"}), 400
+
+ username = data.get("username")
+ ttl_hours = data.get("ttl_hours", 24)
+
+ if not username:
+ return jsonify({"error": "username is required"}), 400
+
+ user_authenticator = UserAuthenticator.get_instance()
+ token = user_authenticator.create_reset_token(username, ttl_hours=ttl_hours)
+
+ if token is None:
+ return jsonify({"error": f"User '{username}' does not exist"}), 404
+
+ reset_link = f"{request.host_url.rstrip('/')}/reset/{token}"
+ return jsonify({
+ "status": "success",
+ "reset_link": reset_link,
+ "token": token,
+ "expires_in_hours": ttl_hours
+ })
+
+
+def forgot_password():
+ """Self-service forgot password page.
+
+ GET: Show the forgot password form.
+ POST: Generate a reset token and display the reset link.
+ """
+ if request.method == "GET":
+ return render_template("forgot_password.html",
+ title=config.get("annotation_task_name", "Annotation Platform"))
+
+ username = request.form.get("username", "").strip()
+ # Always show success to prevent user enumeration
+ if not username:
+ return render_template("forgot_password.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ error="Please enter your username.")
+
+ user_authenticator = UserAuthenticator.get_instance()
+ token = user_authenticator.create_reset_token(username)
+
+ if token:
+ reset_link = f"{request.host_url.rstrip('/')}/reset/{token}"
+ return render_template("forgot_password.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ reset_link=reset_link,
+ success=True)
+ else:
+ # Show same success message to prevent enumeration
+ return render_template("forgot_password.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ success=True)
+
+
+def reset_password_with_token(token):
+ """Self-service password reset using a token.
+
+ GET: Show the reset form if token is valid.
+ POST: Process the new password.
+ """
+ user_authenticator = UserAuthenticator.get_instance()
+
+ if request.method == "GET":
+ username = user_authenticator.validate_reset_token(token)
+ if not username:
+ return render_template("reset_password.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ error="This reset link is invalid or has expired.",
+ token_invalid=True)
+ return render_template("reset_password.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ token=token,
+ username=username)
+
+ # POST
+ new_password = request.form.get("password", "")
+ confirm_password = request.form.get("confirm_password", "")
+
+ if not new_password:
+ return render_template("reset_password.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ token=token,
+ error="Password cannot be empty.")
+
+ if new_password != confirm_password:
+ return render_template("reset_password.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ token=token,
+ error="Passwords do not match.")
+
+ # Consume token (single-use)
+ username = user_authenticator.consume_reset_token(token)
+ if not username:
+ return render_template("reset_password.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ error="This reset link is invalid or has expired.",
+ token_invalid=True)
+
+ if user_authenticator.update_password(username, new_password):
+ user_authenticator.save_user_config()
+ return render_template("reset_password.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ success=True)
+ else:
+ return render_template("reset_password.html",
+ title=config.get("annotation_task_name", "Annotation Platform"),
+ token=token,
+ error="Failed to reset password. Please try again.")
+
+
+def configure_routes(flask_app, app_config):
+ """
+ Initialize the Flask routes with the given Flask app instance
+ and configuration.
+
+ This function is called by flask_server.py when initializing the application.
+
+ Args:
+ flask_app: The Flask application instance
+ app_config: The application configuration
+ """
+ global app, config
+ app = flask_app
+ config = app_config
+
+ # Set up session configuration
+ # Use a random secret key if sessions shouldn't persist, otherwise use the configured one
+ if config.get("persist_sessions", False):
+ secret_key = config.get("secret_key") or os.environ.get("POTATO_SECRET_KEY")
+ if not secret_key:
+ raise ValueError(
+ "persist_sessions is enabled but no secret_key is configured. "
+ "Set 'secret_key' in your config file or POTATO_SECRET_KEY environment variable."
+ )
+ app.secret_key = secret_key
+ else:
+ # Generate a random secret key to ensure sessions don't persist between restarts
+ import secrets
+ app.secret_key = secrets.token_hex(32)
+
+ app.permanent_session_lifetime = timedelta(days=config.get("session_lifetime_days", 7))
+
+ # Register all routes with the flask app instance
+ app.add_url_rule("/media/", "serve_media", serve_media)
+ app.add_url_rule(
+ "/screenshots/",
+ "serve_trace_screenshot",
+ serve_trace_screenshot,
+ )
+ app.add_url_rule("/", "home", home, methods=["GET", "POST"])
+ app.add_url_rule("/auth", "auth", auth, methods=["GET", "POST"])
+ app.add_url_rule("/passwordless-login", "passwordless_login", passwordless_login, methods=["GET", "POST"])
+ app.add_url_rule("/clerk-login", "clerk_login", clerk_login, methods=["GET", "POST"])
+ app.add_url_rule("/login", "login", login, methods=["GET", "POST"])
+ app.add_url_rule("/logout", "logout", logout)
+ app.add_url_rule("/submit_annotation", "submit_annotation", submit_annotation, methods=["POST"])
+ app.add_url_rule("/register", "register", register, methods=["POST"])
+ app.add_url_rule("/consent", "consent", consent, methods=["GET", "POST"])
+ app.add_url_rule("/instructions", "instructions", instructions, methods=["GET", "POST"])
+ app.add_url_rule("/prestudy", "prestudy", prestudy, methods=["GET", "POST"])
+ app.add_url_rule("/training", "training", training, methods=["GET", "POST"])
+ app.add_url_rule("/annotate", "annotate", annotate, methods=["GET", "POST"])
+ app.add_url_rule("/go_to", "go_to", go_to, methods=["GET", "POST"])
+ app.add_url_rule("/updateinstance", "update_instance", update_instance, methods=["POST"])
+ app.add_url_rule("/get_annotations", "get_annotations", get_annotations, methods=["GET"])
+ app.add_url_rule("/poststudy", "poststudy", poststudy, methods=["GET", "POST"])
+ app.add_url_rule("/done", "done", done, methods=["GET", "POST"])
+ app.add_url_rule("/admin", "admin", admin, methods=["GET"])
+
+ app.add_url_rule("/api/get_ai_suggestion", "get_ai_suggestion", get_ai_suggestion, methods=["GET"])
+
+ # Option highlighting API routes
+ app.add_url_rule("/api/option_highlights/config", "get_option_highlighting_config", get_option_highlighting_config, methods=["GET"])
+ app.add_url_rule("/api/option_highlights/prefetch", "trigger_option_highlight_prefetch", trigger_option_highlight_prefetch, methods=["POST"])
+ app.add_url_rule("/api/option_highlights/", "get_option_highlights", get_option_highlights, methods=["GET"])
+
+ app.add_url_rule("/api-frontend", "api_frontend", api_frontend, methods=["GET"])
+ app.add_url_rule("/span-api-frontend", "span_api_frontend", span_api_frontend, methods=["GET"])
+ app.add_url_rule("/api/spans/", "get_span_data", get_span_data, methods=["GET"])
+ app.add_url_rule("/api/colors", "get_span_colors", get_span_colors, methods=["GET"])
+ app.add_url_rule("/api/schemas", "get_annotation_schemas", get_annotation_schemas, methods=["GET"])
+ app.add_url_rule("/api/keyword_highlights/", "get_keyword_highlights", get_keyword_highlights, methods=["GET"])
+ app.add_url_rule("/test-span-colors", "test_span_colors", test_span_colors, methods=["GET"])
+ app.add_url_rule("/api/spans//clear", "clear_span_annotations", clear_span_annotations, methods=["POST"])
+ app.add_url_rule("/api/links/", "get_link_annotations", get_link_annotations, methods=["GET"])
+ app.add_url_rule("/api/links//", "delete_link_annotation", delete_link_annotation, methods=["DELETE"])
+
+ # Event annotation API routes
+ app.add_url_rule("/api/events/", "get_event_annotations", get_event_annotations, methods=["GET"])
+ app.add_url_rule("/api/events//", "delete_event_annotation", delete_event_annotation, methods=["DELETE"])
+
+ # Entity linking API routes
+ app.add_url_rule("/api/entity_linking/search", "entity_linking_search", entity_linking_search, methods=["GET"])
+ app.add_url_rule("/api/entity_linking/entity//", "entity_linking_get_entity", entity_linking_get_entity, methods=["GET"])
+ app.add_url_rule("/api/entity_linking/configured_kbs", "entity_linking_configured_kbs", entity_linking_configured_kbs, methods=["GET"])
+ app.add_url_rule("/api/entity_linking/update_span", "entity_linking_update_span", entity_linking_update_span, methods=["POST"])
+
+ app.add_url_rule("/api/current_instance", "get_current_instance", get_current_instance, methods=["GET"])
+ # F-024: get_instance_data is registered only via a module-level @app.route,
+ # which binds to a throwaway app in the CLI start path; without this explicit
+ # re-registration on the serving app the route 404s on every `potato start`
+ # server (dynamic schemas โ extractive_qa/error_span/text_edit/card_sort/
+ # conjoint โ then rely solely on annotation.js's embedded-JSON fallback).
+ app.add_url_rule("/api/instance_data", "get_instance_data", get_instance_data, methods=["GET"])
+ app.add_url_rule("/api/ai_assistant", "ai_assistant", ai_assistant, methods=["GET"])
+ app.add_url_rule("/api/audio/proxy", "audio_proxy", audio_proxy, methods=["GET"])
+ app.add_url_rule("/admin/user_state/", "admin_user_state", admin_user_state, methods=["GET"])
+ app.add_url_rule("/admin/health", "admin_health", admin_health, methods=["GET"])
+ app.add_url_rule("/admin/system_state", "admin_system_state", admin_system_state, methods=["GET"])
+ app.add_url_rule("/admin/all_instances", "admin_all_instances", admin_all_instances, methods=["GET"])
+ app.add_url_rule("/admin/item_state", "admin_item_state", admin_item_state, methods=["GET"])
+ app.add_url_rule("/admin/item_state/", "admin_item_state_detail", admin_item_state_detail, methods=["GET"])
+
+ # Password management routes
+ app.add_url_rule("/admin/reset_password", "admin_reset_password", admin_reset_password, methods=["POST"])
+ app.add_url_rule("/admin/create_reset_token", "admin_create_reset_token", admin_create_reset_token, methods=["POST"])
+ app.add_url_rule("/forgot-password", "forgot_password", forgot_password, methods=["GET", "POST"])
+ app.add_url_rule("/reset/", "reset_password_with_token", reset_password_with_token, methods=["GET", "POST"])
+
+ # New admin dashboard API routes
+ app.add_url_rule("/admin/api/overview", "admin_api_overview", admin_api_overview, methods=["GET"])
+ app.add_url_rule("/admin/api/annotators", "admin_api_annotators", admin_api_annotators, methods=["GET"])
+ app.add_url_rule("/admin/api/instances", "admin_api_instances", admin_api_instances, methods=["GET"])
+ app.add_url_rule("/admin/api/config", "admin_api_config", admin_api_config, methods=["GET", "POST"])
+ app.add_url_rule("/admin/api/user//set_instances", "admin_api_set_user_instances", admin_api_set_user_instances, methods=["POST"])
+ app.add_url_rule("/admin/api/stale_assignments", "admin_api_stale_assignments", admin_api_stale_assignments, methods=["GET"])
+ app.add_url_rule("/admin/api/reclaim_instance", "admin_api_reclaim_instance", admin_api_reclaim_instance, methods=["POST"])
+ # F-041 (F-024 class): debug-gated test-state reset. It existed only as a
+ # module-level @app.route, so it 404'd on the serving app built by
+ # configure_routes โ leaving Selenium suites that rely on it (e.g. video
+ # annotation persistence) without per-test isolation.
+ app.add_url_rule("/admin/api/test/reset_state", "admin_api_test_reset_state", admin_api_test_reset_state, methods=["POST"])
+ # F-042 (F-024 class): these were registered ONLY via module-level @app.route
+ # and so 404'd on every live `potato start` server (the serving app is built
+ # by configure_routes). Audit: diff of @app.route paths vs add_url_rule paths.
+ app.add_url_rule("/get_ai_suggestion", "get_ai_suggestion", get_ai_suggestion, methods=["GET"])
+ app.add_url_rule("/admin/iaa", "admin_iaa", admin_iaa, methods=["GET"])
+ app.add_url_rule("/admin/judge-alignment", "admin_judge_alignment", admin_judge_alignment, methods=["GET"])
+ app.add_url_rule("/admin/api/judge-alignment/run", "admin_judge_alignment_run", admin_judge_alignment_run, methods=["POST"])
+ app.add_url_rule("/admin/triage-queue", "admin_triage_queue", admin_triage_queue, methods=["GET"])
+ app.add_url_rule("/admin/api/step_agreement", "admin_api_step_agreement", admin_api_step_agreement, methods=["GET"])
+ app.add_url_rule("/admin/api/step_quality", "admin_api_step_quality", admin_api_step_quality, methods=["GET"])
+ app.add_url_rule("/api/waveform/", "get_waveform_data", get_waveform_data, methods=["GET"])
+ app.add_url_rule("/api/waveform/generate", "generate_waveform", generate_waveform, methods=["POST"])
+ app.add_url_rule("/api/video/metadata", "get_video_metadata", get_video_metadata, methods=["POST"])
+ app.add_url_rule("/api/video/waveform/generate", "generate_video_waveform", generate_video_waveform, methods=["POST"])
+ app.add_url_rule("/admin/api/data_sources", "admin_api_data_sources", admin_api_data_sources, methods=["GET"])
+ app.add_url_rule("/admin/api/data_sources//load_more", "admin_api_data_sources_load_more", admin_api_data_sources_load_more, methods=["POST"])
+ app.add_url_rule("/admin/api/data_sources//refresh", "admin_api_data_sources_refresh", admin_api_data_sources_refresh, methods=["POST"])
+ app.add_url_rule("/admin/api/cache/clear", "admin_api_cache_clear", admin_api_cache_clear, methods=["POST"])
+ app.add_url_rule("/admin/api/webhooks", "admin_api_webhooks", admin_api_webhooks, methods=["GET"])
+ app.add_url_rule("/admin/api/webhooks/test", "admin_api_webhooks_test", admin_api_webhooks_test, methods=["POST"])
+ app.add_url_rule("/admin/api/questions", "admin_api_questions", admin_api_questions, methods=["GET"])
+ app.add_url_rule("/admin/api/annotation_history", "admin_api_annotation_history", admin_api_annotation_history, methods=["GET"])
+ app.add_url_rule("/admin/api/suspicious_activity", "admin_api_suspicious_activity", admin_api_suspicious_activity, methods=["GET"])
+ app.add_url_rule("/admin/api/crowdsourcing", "admin_api_crowdsourcing", admin_api_crowdsourcing, methods=["GET"])
+
+ # ICL labeling admin API routes
+ app.add_url_rule("/admin/api/icl/status", "admin_api_icl_status", admin_api_icl_status, methods=["GET"])
+ app.add_url_rule("/admin/api/icl/examples", "admin_api_icl_examples", admin_api_icl_examples, methods=["GET"])
+ app.add_url_rule("/admin/api/icl/predictions", "admin_api_icl_predictions", admin_api_icl_predictions, methods=["GET"])
+ app.add_url_rule("/admin/api/icl/accuracy", "admin_api_icl_accuracy", admin_api_icl_accuracy, methods=["GET"])
+ app.add_url_rule("/admin/api/icl/trigger", "admin_api_icl_trigger", admin_api_icl_trigger, methods=["POST"])
+ app.add_url_rule("/api/icl/record_verification", "api_icl_record_verification", api_icl_record_verification, methods=["POST"])
+
+ # Behavioral tracking and analytics routes
+ app.add_url_rule("/admin/api/agreement", "admin_api_agreement", admin_api_agreement, methods=["GET"])
+ app.add_url_rule("/admin/api/code_cooccurrence", "admin_api_code_cooccurrence", admin_api_code_cooccurrence, methods=["GET"])
+ app.add_url_rule("/admin/api/code_crosstab", "admin_api_code_crosstab", admin_api_code_crosstab, methods=["GET"])
+ app.add_url_rule("/admin/api/quality_control", "admin_api_quality_control", admin_api_quality_control, methods=["GET"])
+ app.add_url_rule("/admin/api/behavioral_analytics", "admin_api_behavioral_analytics", admin_api_behavioral_analytics, methods=["GET"])
+ app.add_url_rule("/api/track_interactions", "track_interactions", track_interactions, methods=["POST"])
+ app.add_url_rule("/api/track_ai_usage", "track_ai_usage", track_ai_usage, methods=["POST"])
+ app.add_url_rule("/api/track_annotation_change", "track_annotation_change", track_annotation_change, methods=["POST"])
+ app.add_url_rule("/api/behavioral_data/", "get_behavioral_data", get_behavioral_data, methods=["GET"])
+
+ # Adjudication routes
+ app.add_url_rule("/adjudicate", "adjudicate", adjudicate, methods=["GET"])
+ app.add_url_rule("/adjudicate/api/queue", "adjudicate_api_queue", adjudicate_api_queue, methods=["GET"])
+ app.add_url_rule("/adjudicate/api/item/", "adjudicate_api_item", adjudicate_api_item, methods=["GET"])
+ app.add_url_rule("/adjudicate/api/submit", "adjudicate_api_submit", adjudicate_api_submit, methods=["POST"])
+ app.add_url_rule("/adjudicate/api/stats", "adjudicate_api_stats", adjudicate_api_stats, methods=["GET"])
+ app.add_url_rule("/adjudicate/api/skip/", "adjudicate_api_skip", adjudicate_api_skip, methods=["POST"])
+ app.add_url_rule("/adjudicate/api/next", "adjudicate_api_next", adjudicate_api_next, methods=["GET"])
+ app.add_url_rule("/adjudicate/api/similar/", "adjudicate_api_similar", adjudicate_api_similar, methods=["GET"])
+ app.add_url_rule("/admin/api/adjudication", "admin_api_adjudication", admin_api_adjudication, methods=["GET"])
+
+ # BWS scoring admin API routes
+ app.add_url_rule("/admin/api/bws_scoring", "admin_api_bws_scoring", admin_api_bws_scoring, methods=["GET"])
+ app.add_url_rule("/admin/api/bws_scoring/generate", "admin_api_bws_scoring_generate", admin_api_bws_scoring_generate, methods=["POST"])
+
+ # IBWS admin API routes
+ app.add_url_rule("/admin/api/ibws_status", "admin_api_ibws_status", admin_api_ibws_status, methods=["GET"])
+ app.add_url_rule("/admin/api/ibws_ranking", "admin_api_ibws_ranking", admin_api_ibws_ranking, methods=["GET"])
+
+ # MACE admin API routes
+ app.add_url_rule("/admin/api/mace/overview", "admin_api_mace_overview", admin_api_mace_overview, methods=["GET"])
+ app.add_url_rule("/admin/api/mace/predictions", "admin_api_mace_predictions", admin_api_mace_predictions, methods=["GET"])
+ app.add_url_rule("/admin/api/mace/trigger", "admin_api_mace_trigger", admin_api_mace_trigger, methods=["POST"])
+
+ # Embedding visualization admin API routes
+ app.add_url_rule("/admin/api/embedding_viz/data", "admin_api_embedding_viz_data", admin_api_embedding_viz_data, methods=["GET"])
+ app.add_url_rule("/admin/api/embedding_viz/reorder", "admin_api_embedding_viz_reorder", admin_api_embedding_viz_reorder, methods=["POST"])
+ app.add_url_rule("/admin/api/embedding_viz/refresh", "admin_api_embedding_viz_refresh", admin_api_embedding_viz_refresh, methods=["POST"])
+ app.add_url_rule("/admin/api/embedding_viz/stats", "admin_api_embedding_viz_stats", admin_api_embedding_viz_stats, methods=["GET"])
+
+ # Export admin API routes
+ app.add_url_rule("/admin/api/export/formats", "admin_api_export_formats", admin_api_export_formats, methods=["GET"])
+ app.add_url_rule("/admin/api/export", "admin_api_export", admin_api_export, methods=["POST"])
+
+ # Agent chat routes (interactive agent testing)
+ app.add_url_rule("/agent_chat/send", "agent_chat_send", agent_chat_send, methods=["POST"])
+ app.add_url_rule("/agent_chat/finish", "agent_chat_finish", agent_chat_finish, methods=["POST"])
+ app.add_url_rule("/agent_chat/status", "agent_chat_status", agent_chat_status, methods=["GET"])
+
+ # OAuth SSO routes
+ app.add_url_rule("/auth/login/", "oauth_login", oauth_login, methods=["GET"])
+ app.add_url_rule("/auth/callback/", "oauth_callback", oauth_callback, methods=["GET"])
+
+ # Chat support API routes
+ app.add_url_rule("/api/chat/send", "chat_send", chat_send, methods=["POST"])
+ app.add_url_rule("/api/chat/history", "chat_history", chat_history, methods=["GET"])
+ app.add_url_rule("/api/chat/config", "chat_config", chat_config, methods=["GET"])
+
+ app.add_url_rule("/shutdown", "shutdown", shutdown, methods=["POST"])
+
+ # Register Solo Mode blueprint if not already registered
+ if 'solo_mode' not in app.blueprints:
+ try:
+ from potato.solo_mode.routes import solo_mode_bp
+ app.register_blueprint(solo_mode_bp)
+ except ImportError:
+ pass
+
+ # Register QDA Mode blueprint if not already registered
+ if 'qda_mode' not in app.blueprints:
+ try:
+ from potato.qda_mode import qda_mode_bp
+ app.register_blueprint(qda_mode_bp)
+ except ImportError:
+ pass
+
+ # Register universal Memos blueprint if not already registered
+ if 'memos' not in app.blueprints:
+ try:
+ from potato.memos.api import memos_bp
+ app.register_blueprint(memos_bp)
+ except ImportError:
+ pass
+
+ # Register universal Search blueprint if not already registered
+ if 'search' not in app.blueprints:
+ try:
+ from potato.search.api import search_bp
+ app.register_blueprint(search_bp)
+ except ImportError:
+ pass
+
+ # Register universal Codebook blueprint if not already registered
+ if 'codebook' not in app.blueprints:
+ try:
+ from potato.codebook.api import codebook_bp
+ app.register_blueprint(codebook_bp)
+ except ImportError:
+ pass
+
+ # Register universal Cases blueprint if not already registered
+ if 'cases' not in app.blueprints:
+ try:
+ from potato.cases.api import cases_bp
+ app.register_blueprint(cases_bp)
+ except ImportError:
+ pass
+
+ # Register Judge Calibration blueprint if not already registered
+ if 'judge_calibration' not in app.blueprints:
+ try:
+ from potato.judge_calibration.routes import judge_calibration_bp
+ app.register_blueprint(judge_calibration_bp)
+ except ImportError:
+ pass
+
+# ============================================================================
+# Adjudication Routes
+# ============================================================================
+
+def _check_adjudicator_auth():
+ """Check if current user is an authorized adjudicator.
+
+ Returns:
+ tuple: (is_authorized: bool, username: str, error_response)
+ """
+ username = session.get('username')
+ if not username:
+ return False, None, (jsonify({"error": "Not authenticated"}), 401)
+
+ adj_mgr = get_adjudication_manager()
+ if not adj_mgr or not adj_mgr.adj_config.enabled:
+ return False, username, (jsonify({"error": "Adjudication not enabled"}), 404)
+
+ if not adj_mgr.is_adjudicator(username):
+ return False, username, (jsonify({"error": "Not authorized as adjudicator"}), 403)
+
+ return True, username, None
+
+
+@app.route('/adjudicate', methods=['GET'])
+def adjudicate():
+ """Main adjudication page."""
+ username = session.get('username')
+ if not username:
+ return redirect(url_for('home'))
+
+ adj_mgr = get_adjudication_manager()
+ if not adj_mgr or not adj_mgr.adj_config.enabled:
+ return redirect(url_for('home'))
+
+ if not adj_mgr.is_adjudicator(username):
+ return redirect(url_for('home'))
+
+ # Get annotation schemes for form rendering
+ annotation_schemes = config.get('annotation_schemes', [])
+
+ return render_template(
+ 'adjudication.html',
+ annotation_task_name=config.get('annotation_task_name', 'Annotation Task'),
+ username=username,
+ annotation_schemes=annotation_schemes,
+ adj_config={
+ 'show_annotator_names': adj_mgr.adj_config.show_annotator_names,
+ 'show_timing_data': adj_mgr.adj_config.show_timing_data,
+ 'show_agreement_scores': adj_mgr.adj_config.show_agreement_scores,
+ 'fast_decision_warning_ms': adj_mgr.adj_config.fast_decision_warning_ms,
+ 'require_confidence': adj_mgr.adj_config.require_confidence,
+ 'require_notes_on_override': adj_mgr.adj_config.require_notes_on_override,
+ 'error_taxonomy': adj_mgr.adj_config.error_taxonomy,
+ 'similarity_enabled': adj_mgr.adj_config.similarity_enabled,
+ },
+ )
+
+
+@app.route('/adjudicate/api/queue', methods=['GET'])
+def adjudicate_api_queue():
+ """Get the adjudication queue."""
+ authorized, username, error = _check_adjudicator_auth()
+ if not authorized:
+ return error
+
+ adj_mgr = get_adjudication_manager()
+ filter_status = request.args.get('status', None)
+
+ items = adj_mgr.get_queue(
+ adjudicator_id=username,
+ filter_status=filter_status,
+ )
+
+ return jsonify({
+ "items": [item.to_dict() for item in items],
+ "total": len(items),
+ })
+
+
+@app.route('/adjudicate/api/item/', methods=['GET'])
+def adjudicate_api_item(instance_id):
+ """Get full item detail for adjudication."""
+ authorized, username, error = _check_adjudicator_auth()
+ if not authorized:
+ return error
+
+ adj_mgr = get_adjudication_manager()
+ item = adj_mgr.get_item(instance_id)
+
+ if not item:
+ return jsonify({"error": "Item not found in adjudication queue"}), 404
+
+ # Get item text and data
+ item_text = adj_mgr.get_item_text(instance_id)
+ item_data = adj_mgr.get_item_data(instance_id)
+
+ # Get existing decision if any
+ decision = adj_mgr.get_decision(instance_id)
+
+ # Phase 3: annotator signals and similar items
+ annotator_signals = {}
+ for user_id in item.annotations:
+ annotator_signals[user_id] = adj_mgr.get_annotator_signals(
+ user_id, instance_id
+ )
+
+ similar_items = []
+ if adj_mgr.adj_config.similarity_enabled:
+ similar_items = adj_mgr.get_similar_items(instance_id)
+
+ return jsonify({
+ "item": item.to_dict(),
+ "item_text": item_text,
+ "item_data": item_data,
+ "decision": decision.to_dict() if decision else None,
+ "annotator_signals": annotator_signals,
+ "similar_items": similar_items,
+ })
+
+
+@app.route('/adjudicate/api/similar/', methods=['GET'])
+def adjudicate_api_similar(instance_id):
+ """Get similar items for a specific instance (lazy-loading endpoint)."""
+ authorized, username, error = _check_adjudicator_auth()
+ if not authorized:
+ return error
+
+ adj_mgr = get_adjudication_manager()
+ enabled = adj_mgr.adj_config.similarity_enabled
+ similar_items = adj_mgr.get_similar_items(instance_id) if enabled else []
+
+ return jsonify({
+ "enabled": enabled,
+ "instance_id": instance_id,
+ "similar_items": similar_items,
+ "count": len(similar_items),
+ })
+
+
+@app.route('/admin/api/adjudication', methods=['GET'])
+def admin_api_adjudication():
+ """Admin dashboard overview of adjudication status."""
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ adj_mgr = get_adjudication_manager()
+ if not adj_mgr or not adj_mgr.adj_config.enabled:
+ return jsonify({"enabled": False, "message": "Adjudication not configured"})
+
+ overview = admin_dashboard.get_adjudication_overview()
+ return jsonify(overview)
+
+
+# ============================================================================
+# BWS Scoring Admin API Routes
+# ============================================================================
+
+@app.route('/admin/api/bws_scoring', methods=['GET'])
+def admin_api_bws_scoring():
+ """Get current BWS scoring status and cached scores."""
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ if not config.get("bws_config"):
+ return jsonify({"error": "BWS not configured"}), 400
+
+ return jsonify({
+ "total_items": len(config.get("_bws_pool_items", [])),
+ "total_annotations": 0,
+ "method": config.get("bws_config", {}).get("scoring", {}).get("method", "counting"),
+ "scores": [],
+ })
+
+
+@app.route('/admin/api/bws_scoring/generate', methods=['POST'])
+def admin_api_bws_scoring_generate():
+ """Generate BWS scores from current annotations."""
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ if not config.get("bws_config"):
+ return jsonify({"error": "BWS not configured"}), 400
+
+ method = request.args.get("method", "counting")
+
+ from potato.bws_scoring import BwsScorer, write_scores
+
+ # Collect annotations from all users
+ pool_items = config.get("_bws_pool_items", [])
+ id_key = config["item_properties"]["id_key"]
+ text_key = config["item_properties"]["text_key"]
+
+ # Find BWS schema name
+ bws_schema_name = None
+ for scheme in config.get("annotation_schemes", []):
+ if scheme.get("annotation_type") == "bws":
+ bws_schema_name = scheme["name"]
+ break
+
+ if not bws_schema_name:
+ return jsonify({"error": "No BWS annotation scheme found"}), 400
+
+ # Collect annotations from all user states
+ ism = get_item_state_manager()
+ usm = get_user_state_manager()
+ annotations = []
+
+ # Build a lookup of instance_id -> bws_items from items
+ instance_bws_items = {}
+ for item in ism.items():
+ item_data = item.get_data()
+ bws_items = item_data.get("_bws_items", [])
+ if bws_items:
+ instance_bws_items[item.get_id()] = bws_items
+
+ # Iterate all user states to collect annotations
+ for user_state in usm.get_all_users():
+ username = user_state.get_user_id()
+ label_store = getattr(user_state, 'instance_id_to_label_to_value', {})
+
+ for instance_id, labels in label_store.items():
+ bws_items = instance_bws_items.get(instance_id, [])
+ if not bws_items:
+ continue
+
+ # Labels is {Label -> value} dict. Find best/worst for this schema.
+ best_val = None
+ worst_val = None
+ for label_obj, value in labels.items():
+ if label_obj.get_schema() == bws_schema_name:
+ if label_obj.get_name() == "best":
+ best_val = value
+ elif label_obj.get_name() == "worst":
+ worst_val = value
+
+ if best_val and worst_val:
+ annotations.append({
+ "instance_id": instance_id,
+ "bws_items": bws_items,
+ "best": best_val,
+ "worst": worst_val,
+ "annotator": username,
+ })
+
+ if not annotations:
+ return jsonify({
+ "status": "success",
+ "total_items": len(pool_items),
+ "total_annotations": 0,
+ "method": method,
+ "scores": [],
+ "message": "No BWS annotations found yet",
+ })
+
+ try:
+ scorer = BwsScorer(annotations, pool_items, id_key, text_key)
+ scores_dict = scorer.score(method)
+ except ImportError as e:
+ logger.error("BWS scoring import error: %s", traceback.format_exc())
+ return jsonify({"error": "A required dependency is missing"}), 400
+ except Exception as e:
+ logger.error("BWS scoring failed: %s", traceback.format_exc())
+ return jsonify({"error": "An internal error occurred"}), 500
+
+ # Write scores file
+ output_dir = config.get("output_annotation_dir", "annotation_output")
+ output_path = os.path.join(output_dir, "bws_scores.tsv")
+ try:
+ write_scores(scores_dict, output_path)
+ except Exception as e:
+ logger.warning(f"Failed to write BWS scores file: {e}")
+
+ # Sort and format for response
+ sorted_scores = sorted(
+ scores_dict.items(), key=lambda x: x[1]["score"], reverse=True
+ )
+ scores_list = []
+ for rank, (item_id, data) in enumerate(sorted_scores, 1):
+ scores_list.append({
+ "rank": rank,
+ "item_id": item_id,
+ "text": data.get("text", "")[:200],
+ "score": round(data["score"], 6),
+ "best_count": data.get("best_count"),
+ "worst_count": data.get("worst_count"),
+ "appearances": data.get("appearances"),
+ })
+
+ return jsonify({
+ "status": "success",
+ "total_items": len(pool_items),
+ "total_annotations": len(annotations),
+ "method": method,
+ "scores": scores_list,
+ })
+
+
+# ============================================================================
+# IBWS Admin API Routes
+# ============================================================================
+
+@app.route('/admin/api/ibws_status', methods=['GET'])
+def admin_api_ibws_status():
+ """Get current IBWS round status and progress."""
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ if not config.get("ibws_config"):
+ return jsonify({"error": "IBWS not configured"}), 400
+
+ from potato.ibws_manager import get_ibws_manager
+ ibws_mgr = get_ibws_manager()
+ if not ibws_mgr:
+ return jsonify({"error": "IBWS manager not initialized"}), 500
+
+ return jsonify(ibws_mgr.get_round_info())
+
+
+@app.route('/admin/api/ibws_ranking', methods=['GET'])
+def admin_api_ibws_ranking():
+ """Get current IBWS ordinal ranking."""
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ if not config.get("ibws_config"):
+ return jsonify({"error": "IBWS not configured"}), 400
+
+ from potato.ibws_manager import get_ibws_manager
+ ibws_mgr = get_ibws_manager()
+ if not ibws_mgr:
+ return jsonify({"error": "IBWS manager not initialized"}), 500
+
+ ranking = ibws_mgr.get_final_ranking()
+ return jsonify({
+ "completed": ibws_mgr.is_completed(),
+ "current_round": ibws_mgr.current_round,
+ "ranking": ranking,
+ })
+
+
+# ============================================================================
+# MACE Admin API Routes
+# ============================================================================
+
+@app.route('/admin/api/mace/overview', methods=['GET'])
+def admin_api_mace_overview():
+ """Admin dashboard overview of MACE competence estimation results."""
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ return jsonify(admin_dashboard.get_mace_overview())
+
+
+@app.route('/admin/api/mace/predictions', methods=['GET'])
+def admin_api_mace_predictions():
+ """Get MACE predicted labels, optionally filtered by schema and instance."""
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ schema = request.args.get('schema')
+ instance_id = request.args.get('instance_id')
+
+ if not schema:
+ return jsonify({"error": "schema parameter required"}), 400
+
+ return jsonify(admin_dashboard.get_mace_predictions(schema, instance_id))
+
+
+@app.route('/admin/api/mace/trigger', methods=['POST'])
+def admin_api_mace_trigger():
+ """Manually trigger a MACE recomputation."""
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.mace_manager import get_mace_manager
+ mace_mgr = get_mace_manager()
+ if not mace_mgr or not mace_mgr.mace_config.enabled:
+ return jsonify({"error": "MACE not enabled"}), 400
+
+ results = mace_mgr.run_all_schemas()
+ return jsonify({
+ "status": "success",
+ "schemas_processed": len(results),
+ "schemas": list(results.keys()),
+ })
+
+
+# =============================================================================
+# Embedding Visualization API Endpoints
+# =============================================================================
+
+@app.route('/admin/api/embedding_viz/data', methods=['GET'])
+def admin_api_embedding_viz_data():
+ """
+ Get visualization data for the embedding scatter plot.
+ Admin-only endpoint requiring API key.
+
+ Query Parameters:
+ force_refresh: If "true", force recomputation of UMAP projection
+
+ Returns:
+ JSON with points, labels, label_colors, and stats
+ """
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.embedding_visualization import get_embedding_viz_manager
+
+ viz_manager = get_embedding_viz_manager()
+ if not viz_manager:
+ return jsonify({
+ "error": "Embedding visualization not initialized. "
+ "Ensure diversity_ordering is enabled in config."
+ }), 400
+
+ if not viz_manager.enabled:
+ return jsonify({
+ "error": "Embedding visualization disabled. "
+ "Install umap-learn: pip install umap-learn"
+ }), 400
+
+ try:
+ force_refresh = request.args.get('force_refresh', 'false').lower() == 'true'
+ data = viz_manager.get_visualization_data(force_refresh=force_refresh)
+
+ # Convert to JSON-serializable format
+ points_json = []
+ for p in data.points:
+ points_json.append({
+ "instance_id": p.instance_id,
+ "x": p.x,
+ "y": p.y,
+ "label": p.label,
+ "label_source": p.label_source,
+ "preview": p.preview,
+ "preview_type": p.preview_type,
+ "annotated": p.annotated,
+ "annotation_count": p.annotation_count
+ })
+
+ return jsonify({
+ "points": points_json,
+ "labels": data.labels,
+ "label_colors": data.label_colors,
+ "stats": data.stats
+ })
+
+ except Exception as e:
+ logger.error("Error getting embedding visualization data: %s", traceback.format_exc())
+ return jsonify({"error": "An internal error occurred"}), 500
+
+
+@app.route('/admin/api/embedding_viz/reorder', methods=['POST'])
+def admin_api_embedding_viz_reorder():
+ """
+ Reorder the annotation queue based on selected instances.
+ Admin-only endpoint requiring API key.
+
+ JSON Body:
+ selections: List of selection groups, each with:
+ - instance_ids: List of selected instance IDs
+ - priority: Priority number (lower = higher priority)
+ interleave: Whether to interleave selections (default: true)
+
+ Returns:
+ JSON with success status, reordered_count, and new_order_preview
+ """
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.embedding_visualization import get_embedding_viz_manager
+
+ viz_manager = get_embedding_viz_manager()
+ if not viz_manager or not viz_manager.enabled:
+ return jsonify({"error": "Embedding visualization not available"}), 400
+
+ try:
+ data = request.get_json()
+ if not data:
+ return jsonify({"error": "No JSON data provided"}), 400
+
+ selections = data.get("selections", [])
+ interleave = data.get("interleave", True)
+
+ if not selections:
+ return jsonify({"error": "No selections provided"}), 400
+
+ result = viz_manager.reorder_instances(selections, interleave=interleave)
+ return jsonify(result)
+
+ except Exception as e:
+ logger.error("Error reordering instances: %s", traceback.format_exc())
+ return jsonify({"success": False, "error": "An internal error occurred"}), 500
+
+
+@app.route('/admin/api/embedding_viz/refresh', methods=['POST'])
+def admin_api_embedding_viz_refresh():
+ """
+ Force re-computation of embeddings and UMAP projection.
+ Admin-only endpoint requiring API key.
+
+ JSON Body (optional):
+ force_recompute: If true, invalidate cache and recompute (default: true)
+
+ Returns:
+ JSON with status and statistics
+ """
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.embedding_visualization import get_embedding_viz_manager
+
+ viz_manager = get_embedding_viz_manager()
+ if not viz_manager or not viz_manager.enabled:
+ return jsonify({"error": "Embedding visualization not available"}), 400
+
+ try:
+ data = request.get_json() or {}
+ force_recompute = data.get("force_recompute", True)
+
+ if force_recompute:
+ viz_manager.invalidate_cache()
+
+ # Trigger recomputation by fetching data
+ viz_data = viz_manager.get_visualization_data(force_refresh=True)
+
+ return jsonify({
+ "status": "success",
+ "stats": viz_data.stats
+ })
+
+ except Exception as e:
+ logger.error("Error refreshing embedding visualization: %s", traceback.format_exc())
+ return jsonify({"status": "error", "error": "An internal error occurred"}), 500
+
+
+@app.route('/admin/api/embedding_viz/stats', methods=['GET'])
+def admin_api_embedding_viz_stats():
+ """
+ Get embedding visualization statistics.
+ Admin-only endpoint requiring API key.
+
+ Returns:
+ JSON with visualization manager statistics
+ """
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.embedding_visualization import get_embedding_viz_manager
+
+ viz_manager = get_embedding_viz_manager()
+ if not viz_manager:
+ return jsonify({
+ "enabled": False,
+ "error": "Embedding visualization not initialized"
+ })
+
+ return jsonify(viz_manager.get_stats())
+
+
+# =============================================================================
+# Data Sources API Endpoints
+# =============================================================================
+
+@app.route('/admin/api/data_sources', methods=['GET'])
+def admin_api_data_sources():
+ """
+ List all data sources with their status.
+
+ Returns:
+ JSON with list of sources and their status
+ """
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.data_sources import get_data_source_manager
+
+ manager = get_data_source_manager()
+ if not manager:
+ return jsonify({
+ "enabled": False,
+ "message": "Data sources not configured"
+ })
+
+ return jsonify({
+ "enabled": True,
+ "sources": manager.list_sources(),
+ "stats": manager.get_stats()
+ })
+
+
+@app.route('/admin/api/data_sources//load_more', methods=['POST'])
+def admin_api_data_sources_load_more(source_id):
+ """
+ Load more items from a specific data source.
+
+ Args:
+ source_id: The source identifier
+
+ Query params:
+ count: Number of items to load (optional, uses batch_size default)
+
+ Returns:
+ JSON with number of items loaded
+ """
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.data_sources import get_data_source_manager
+
+ manager = get_data_source_manager()
+ if not manager:
+ return jsonify({"error": "Data sources not configured"}), 400
+
+ try:
+ # Get optional count parameter
+ count = request.args.get('count', type=int)
+
+ loaded = manager.load_more(source_id, count=count)
+ return jsonify({
+ "status": "success",
+ "source_id": source_id,
+ "items_loaded": loaded
+ })
+ except ValueError as e:
+ return jsonify({"error": str(e)}), 404
+ except Exception as e:
+ logger.error(f"Error loading more from {source_id}: {e}")
+ return jsonify({"error": str(e)}), 500
+
+
+@app.route('/admin/api/data_sources//refresh', methods=['POST'])
+def admin_api_data_sources_refresh(source_id):
+ """
+ Refresh a data source (re-fetch from remote).
+
+ Args:
+ source_id: The source identifier
+
+ Returns:
+ JSON with refresh status
+ """
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.data_sources import get_data_source_manager
+
+ manager = get_data_source_manager()
+ if not manager:
+ return jsonify({"error": "Data sources not configured"}), 400
+
+ try:
+ success = manager.refresh_source(source_id)
+ return jsonify({
+ "status": "success" if success else "failed",
+ "source_id": source_id
+ })
+ except ValueError as e:
+ return jsonify({"error": str(e)}), 404
+ except Exception as e:
+ logger.error(f"Error refreshing source {source_id}: {e}")
+ return jsonify({"error": str(e)}), 500
+
+
+@app.route('/admin/api/cache/clear', methods=['POST'])
+def admin_api_cache_clear():
+ """
+ Clear the data source cache.
+
+ Returns:
+ JSON with number of entries cleared
+ """
+ from potato.admin import admin_dashboard
+ if not admin_dashboard.check_admin_access():
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.data_sources import get_data_source_manager
+
+ manager = get_data_source_manager()
+ if not manager:
+ return jsonify({"error": "Data sources not configured"}), 400
+
+ try:
+ entries_cleared = manager.clear_cache()
+ return jsonify({
+ "status": "success",
+ "entries_cleared": entries_cleared
+ })
+ except Exception as e:
+ logger.error(f"Error clearing cache: {e}")
+ return jsonify({"error": str(e)}), 500
+
+
+@app.route('/admin/api/webhooks', methods=['GET'])
+def admin_api_webhooks():
+ """Get webhook configuration and delivery stats."""
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.webhooks import get_webhook_emitter
+ emitter = get_webhook_emitter()
+ if not emitter:
+ return jsonify({"enabled": False, "endpoints": [], "stats": {}})
+
+ return jsonify({
+ "enabled": True,
+ "endpoints": emitter.get_endpoint_info(),
+ "stats": emitter.get_stats(),
+ })
+
+
+@app.route('/admin/api/webhooks/test', methods=['POST'])
+def admin_api_webhooks_test():
+ """Send a test webhook event to verify endpoint connectivity."""
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.webhooks import get_webhook_emitter
+ emitter = get_webhook_emitter()
+ if not emitter:
+ return jsonify({"error": "Webhooks not enabled"}), 400
+
+ import datetime
+ test_payload = {
+ "event": "webhook.test",
+ "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
+ "data": {"message": "Test webhook from Potato admin"},
+ }
+ count = emitter.emit("webhook.test", test_payload)
+ return jsonify({"status": "sent", "endpoints_matched": count})
+
+
+@app.route('/admin/api/export/formats', methods=['GET'])
+def admin_api_export_formats():
+ """List available export formats with metadata."""
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({"error": "Admin access required"}), 403
+
+ from potato.export import export_registry
+ formats = export_registry.list_exporters()
+ return jsonify({"formats": formats})
+
+
+@app.route('/admin/api/export', methods=['POST'])
+def admin_api_export():
+ """Run an export in the requested format and return the result."""
+ api_key = request.headers.get('X-API-Key')
+ if not validate_admin_api_key(api_key):
+ return jsonify({"error": "Admin access required"}), 403
+
+ data = request.get_json(silent=True) or {}
+ fmt = data.get("format")
+ if not fmt:
+ return jsonify({"error": "Missing required field: format"}), 400
+
+ output = data.get("output", "")
+ options = data.get("options") or {}
+
+ config_file = config.get("__config_file__")
+ if not config_file:
+ return jsonify({"error": "Config file path not available"}), 500
+
+ try:
+ from potato.export.cli import build_export_context
+ from potato.export import export_registry
+
+ context = build_export_context(config_file)
+ # Default the output directory when the caller didn't supply one, so the
+ # admin dashboard "Export" button works without requiring a path. Writes
+ # under /exports// (auto-export convention).
+ if not output:
+ base_out = getattr(context, "output_dir", "") or os.getcwd()
+ output = os.path.join(base_out, "exports", fmt)
+ os.makedirs(output, exist_ok=True)
+ result = export_registry.export(fmt, context, output, options)
+
+ return jsonify({
+ "success": result.success,
+ "format": result.format_name,
+ "files_written": result.files_written,
+ "stats": result.stats,
+ "warnings": result.warnings,
+ "errors": result.errors,
+ })
+ except Exception as e:
+ logger.exception("Export failed: %s", e)
+ return jsonify({"error": str(e)}), 500
+
+
+@app.route('/adjudicate/api/submit', methods=['POST'])
+def adjudicate_api_submit():
+ """Submit an adjudication decision."""
+ authorized, username, error = _check_adjudicator_auth()
+ if not authorized:
+ return error
+
+ adj_mgr = get_adjudication_manager()
+
+ try:
+ data = request.get_json(silent=True)
+ if not data:
+ return jsonify({"error": "No JSON data provided"}), 400
+
+ instance_id = data.get('instance_id')
+ if not instance_id:
+ return jsonify({"error": "instance_id is required"}), 400
+
+ decision = AdjudicationDecision(
+ instance_id=str(instance_id),
+ adjudicator_id=username,
+ timestamp=datetime.datetime.now().isoformat(),
+ label_decisions=data.get('label_decisions', {}),
+ span_decisions=data.get('span_decisions', []),
+ source=data.get('source', {}),
+ confidence=data.get('confidence', 'medium'),
+ notes=data.get('notes', ''),
+ error_taxonomy=data.get('error_taxonomy', []),
+ guideline_update_flag=data.get('guideline_update_flag', False),
+ guideline_update_notes=data.get('guideline_update_notes', ''),
+ time_spent_ms=data.get('time_spent_ms', 0),
+ )
+
+ success = adj_mgr.submit_decision(decision)
+ if success:
+ return jsonify({"status": "ok", "instance_id": instance_id})
+ else:
+ return jsonify({"error": "Failed to save decision"}), 500
+
+ except Exception as e:
+ logger.error(f"Error submitting adjudication decision: {e}")
+ return jsonify({"error": str(e)}), 500
+
+
+@app.route('/adjudicate/api/stats', methods=['GET'])
+def adjudicate_api_stats():
+ """Get adjudication progress statistics."""
+ authorized, username, error = _check_adjudicator_auth()
+ if not authorized:
+ return error
+
+ adj_mgr = get_adjudication_manager()
+ stats = adj_mgr.get_stats()
+ return jsonify(stats)
+
+
+@app.route('/adjudicate/api/skip/', methods=['POST'])
+def adjudicate_api_skip(instance_id):
+ """Skip an adjudication item."""
+ authorized, username, error = _check_adjudicator_auth()
+ if not authorized:
+ return error
+
+ adj_mgr = get_adjudication_manager()
+ success = adj_mgr.skip_item(instance_id, username)
+
+ if success:
+ return jsonify({"status": "ok", "instance_id": instance_id})
+ else:
+ return jsonify({"error": "Item not found"}), 404
+
+
+@app.route('/adjudicate/api/next', methods=['GET'])
+def adjudicate_api_next():
+ """Get the next item to adjudicate."""
+ authorized, username, error = _check_adjudicator_auth()
+ if not authorized:
+ return error
+
+ adj_mgr = get_adjudication_manager()
+ item = adj_mgr.get_next_item(username)
+
+ if not item:
+ return jsonify({"item": None, "message": "No more items to adjudicate"})
+
+ item_text = adj_mgr.get_item_text(item.instance_id)
+ item_data = adj_mgr.get_item_data(item.instance_id)
+
+ return jsonify({
+ "item": item.to_dict(),
+ "item_text": item_text,
+ "item_data": item_data,
+ })
+
+
+# ============================================================================
+# Chat Support API Endpoints
+# ============================================================================
+
+@app.route("/api/chat/send", methods=["POST"])
+def chat_send():
+ """Send a message to the LLM chat assistant and get a response."""
+ import time as time_module
+ from potato.chat_manager import get_chat_manager
+ from potato.interaction_tracking import (
+ get_or_create_behavioral_data, ChatMessage,
+ )
+
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ username = session['username']
+ data = request.get_json()
+ if not data or not data.get("message"):
+ return jsonify({"error": "No message provided"}), 400
+
+ chat_manager = get_chat_manager()
+ if not chat_manager or not chat_manager.enabled:
+ return jsonify({"error": "Chat support is not enabled"}), 404
+
+ user_state = get_user_state(username)
+ if not user_state:
+ return jsonify({"error": "User state not found"}), 404
+
+ # Get instance info
+ instance_id = data.get("instance_id")
+ if not instance_id:
+ current_instance = user_state.get_current_instance()
+ if current_instance:
+ instance_id = current_instance.get_id()
+
+ # Get instance text for context
+ instance_text = ""
+ if instance_id:
+ item_state_mgr = get_item_state_manager()
+ item = item_state_mgr.get_item(instance_id)
+ if item:
+ text_key = config.get("item_properties", {}).get("text_key", "text")
+ item_data = item.get_data()
+ instance_text = item_data.get(text_key, item.get_text())
+
+ # Get behavioral data and existing chat history
+ bd = get_or_create_behavioral_data(
+ user_state.instance_id_to_behavioral_data,
+ instance_id or "",
+ )
+
+ # Build history from chat_history (only role + content for the LLM)
+ history = [
+ {"role": msg.role, "content": msg.content}
+ for msg in bd.chat_history
+ if isinstance(msg, ChatMessage)
+ ]
+
+ # Enforce max history
+ max_hist = chat_manager.max_history_per_instance
+ if len(history) > max_hist:
+ history = history[-max_hist:]
+
+ user_message = data["message"]
+ now = time_module.time()
+
+ # Send to LLM
+ result = chat_manager.send_message(
+ user_message, instance_text, instance_id or "", history
+ )
+
+ # Record messages in behavioral data
+ bd.chat_history.append(ChatMessage(
+ role="user",
+ content=user_message,
+ timestamp=now,
+ instance_id=instance_id or "",
+ ))
+ bd.chat_history.append(ChatMessage(
+ role="assistant",
+ content=result["content"],
+ timestamp=time_module.time(),
+ instance_id=instance_id or "",
+ response_time_ms=result["response_time_ms"],
+ ))
+
+ # Log interaction event
+ bd.add_interaction(
+ event_type="chat_message_sent",
+ target="chat_sidebar",
+ metadata={
+ "message_length": len(user_message),
+ "response_length": len(result["content"]),
+ "response_time_ms": result["response_time_ms"],
+ },
+ )
+
+ # Persist user state
+ usm = get_user_state_manager()
+ if usm:
+ usm.save_user_state(user_state)
+
+ return jsonify(result)
+
+
+@app.route("/api/chat/history", methods=["GET"])
+def chat_history():
+ """Get chat history for a specific instance."""
+ from potato.interaction_tracking import get_or_create_behavioral_data, ChatMessage
+
+ if 'username' not in session:
+ return jsonify({"error": "Not authenticated"}), 401
+
+ username = session['username']
+ instance_id = request.args.get("instance_id", "")
+
+ user_state = get_user_state(username)
+ if not user_state:
+ return jsonify({"error": "User state not found"}), 404
+
+ if not instance_id:
+ current_instance = user_state.get_current_instance()
+ if current_instance:
+ instance_id = current_instance.get_id()
+
+ # Get existing behavioral data (don't create if it doesn't exist)
+ bd_dict = user_state.instance_id_to_behavioral_data
+ if instance_id and instance_id in bd_dict:
+ bd = get_or_create_behavioral_data(bd_dict, instance_id)
+ messages = [
+ msg.to_dict() if hasattr(msg, 'to_dict') else msg
+ for msg in bd.chat_history
+ ]
+ else:
+ messages = []
+
+ return jsonify({"messages": messages, "instance_id": instance_id})
+
+
+@app.route("/api/chat/config", methods=["GET"])
+def chat_config():
+ """Get chat UI configuration."""
+ from potato.chat_manager import get_chat_manager
+
+ chat_manager = get_chat_manager()
+ if not chat_manager:
+ return jsonify({"enabled": False})
+
+ return jsonify(chat_manager.get_ui_config())
+
+
+@app.route('/shutdown', methods=['POST'])
+def shutdown():
+ func = request.environ.get('werkzeug.server.shutdown')
+ if func is None:
+ return jsonify({'error': 'Not running with the Werkzeug Server'}), 500
+ logger.info('Shutting down server via /shutdown')
+ func()
+ return jsonify({'status': 'Server shutting down...'})
+
diff --git a/potato/routes_live_agent.py b/potato/routes_live_agent.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f0e23f1ee39958d10434866228325add6ef0743
--- /dev/null
+++ b/potato/routes_live_agent.py
@@ -0,0 +1,378 @@
+"""
+Live Agent API Routes
+
+Provides SSE streaming and REST control endpoints for the live agent interaction mode.
+Annotators can observe an AI agent browse the web in real time, pause/resume it,
+send instructions, or take over manual control.
+
+Endpoints:
+- POST /api/live_agent/start - Create and start an agent session
+- GET /api/live_agent/stream/ - SSE event stream
+- POST /api/live_agent/pause/ - Pause agent
+- POST /api/live_agent/resume/ - Resume agent
+- POST /api/live_agent/instruct/ - Send instruction to agent
+- POST /api/live_agent/takeover/ - Toggle manual control
+- POST /api/live_agent/manual_action/ - Execute manual Playwright action
+- POST /api/live_agent/stop/ - Stop and save trace
+- GET /api/live_agent/screenshot// - Serve screenshot
+- GET /api/live_agent/state/ - Current state
+"""
+
+import json
+import logging
+import os
+import time
+from functools import wraps
+from queue import Queue, Empty
+
+from flask import (
+ Blueprint,
+ Response,
+ current_app,
+ jsonify,
+ request,
+ send_file,
+ session as flask_session,
+ redirect,
+ url_for,
+)
+
+from potato.agent_runner import AgentConfig, AgentState
+from potato.agent_runner_manager import AgentRunnerManager
+
+logger = logging.getLogger(__name__)
+
+live_agent_bp = Blueprint("live_agent", __name__)
+
+
+def _login_required(f):
+ """Require user authentication."""
+ @wraps(f)
+ def decorated(*args, **kwargs):
+ if "username" not in flask_session:
+ return redirect(url_for("login"))
+ return f(*args, **kwargs)
+ return decorated
+
+
+def _get_manager() -> AgentRunnerManager:
+ """Get the AgentRunnerManager singleton."""
+ return AgentRunnerManager.get_instance()
+
+
+def _get_runner(session_id: str):
+ """Get an AgentRunner by session_id, or return (None, error_response)."""
+ runner = _get_manager().get_session(session_id)
+ if not runner:
+ return None, (jsonify({"error": f"Unknown session: {session_id}"}), 404)
+ return runner, None
+
+
+@live_agent_bp.route("/api/live_agent/start", methods=["POST"])
+@_login_required
+def start_session():
+ """
+ Start a new live agent session.
+
+ Request JSON:
+ task_description: str - What the agent should do
+ start_url: str - URL to begin at
+ instance_id: str - Annotation instance ID
+ config: dict (optional) - Override live_agent config
+
+ Returns:
+ session_id, state
+ """
+ data = request.get_json(silent=True) or {}
+ task_description = data.get("task_description", "")
+ start_url = data.get("start_url", "")
+ instance_id = data.get("instance_id", "")
+ user_id = flask_session.get("username", "anonymous")
+
+ if not task_description:
+ return jsonify({"error": "task_description is required"}), 400
+ if not start_url:
+ return jsonify({"error": "start_url is required"}), 400
+
+ # Build config from server config + request overrides
+ server_config = current_app.config.get("live_agent", {})
+ override_config = data.get("config", {})
+ merged = {**server_config, **override_config}
+ agent_config = AgentConfig.from_config(merged)
+
+ # Screenshot directory โ must be absolute so both the agent runner
+ # (which saves files) and Flask's send_file (which serves them) agree
+ task_dir = current_app.config.get("task_dir", ".")
+ task_dir = os.path.abspath(task_dir)
+ session_key = f"{user_id}_{instance_id}_{int(time.time())}"
+ screenshot_dir = os.path.join(
+ task_dir, "live_sessions", session_key, "screenshots"
+ )
+ os.makedirs(screenshot_dir, exist_ok=True)
+
+ try:
+ manager = _get_manager()
+ runner = manager.create_session(
+ user_id=user_id,
+ instance_id=instance_id,
+ config=agent_config,
+ screenshot_dir=screenshot_dir,
+ )
+ runner.start(task_description, start_url)
+
+ return jsonify({
+ "session_id": runner.session_id,
+ "state": runner.state.value,
+ })
+
+ except RuntimeError as e:
+ return jsonify({"error": str(e)}), 409
+
+
+@live_agent_bp.route("/api/live_agent/stream/")
+@_login_required
+def stream_events(session_id):
+ """
+ SSE event stream for a live agent session.
+
+ Streams events: thinking, step, state_change, error, complete.
+ Client connects with EventSource.
+ """
+ runner, error = _get_runner(session_id)
+ if error:
+ return error
+
+ def event_stream():
+ q = Queue()
+
+ def listener(event):
+ q.put(event)
+
+ runner.add_listener(listener)
+
+ try:
+ # Send initial state
+ yield _sse_format("connected", {
+ "session_id": session_id,
+ "state": runner.state.value,
+ "step_count": runner.step_count,
+ })
+
+ while True:
+ try:
+ event = q.get(timeout=30)
+ event_type = event.get("type", "message")
+ event_data = event.get("data", {})
+ yield _sse_format(event_type, event_data)
+
+ # Stop streaming when session completes
+ if event_type in ("complete", "error"):
+ break
+
+ except Empty:
+ # Send keepalive
+ yield ": keepalive\n\n"
+
+ finally:
+ runner.remove_listener(listener)
+
+ return Response(
+ event_stream(),
+ mimetype="text/event-stream",
+ headers={
+ "Cache-Control": "no-cache",
+ "X-Accel-Buffering": "no",
+ "Connection": "keep-alive",
+ },
+ )
+
+
+@live_agent_bp.route("/api/live_agent/pause/", methods=["POST"])
+@_login_required
+def pause_session(session_id):
+ """Pause the agent loop."""
+ runner, error = _get_runner(session_id)
+ if error:
+ return error
+
+ runner.pause()
+ return jsonify(runner.get_state_summary())
+
+
+@live_agent_bp.route("/api/live_agent/resume/", methods=["POST"])
+@_login_required
+def resume_session(session_id):
+ """Resume a paused agent."""
+ runner, error = _get_runner(session_id)
+ if error:
+ return error
+
+ runner.resume()
+ return jsonify(runner.get_state_summary())
+
+
+@live_agent_bp.route("/api/live_agent/instruct/", methods=["POST"])
+@_login_required
+def instruct_session(session_id):
+ """Send an instruction to the agent."""
+ runner, error = _get_runner(session_id)
+ if error:
+ return error
+
+ data = request.get_json(silent=True) or {}
+ instruction = data.get("instruction", "")
+ if not instruction:
+ return jsonify({"error": "instruction is required"}), 400
+
+ runner.inject_instruction(instruction)
+ return jsonify(runner.get_state_summary())
+
+
+@live_agent_bp.route("/api/live_agent/takeover/", methods=["POST"])
+@_login_required
+def toggle_takeover(session_id):
+ """Toggle manual takeover mode."""
+ runner, error = _get_runner(session_id)
+ if error:
+ return error
+
+ if runner.state == AgentState.TAKEOVER:
+ runner.exit_takeover()
+ else:
+ runner.enter_takeover()
+
+ return jsonify(runner.get_state_summary())
+
+
+@live_agent_bp.route("/api/live_agent/manual_action/", methods=["POST"])
+@_login_required
+def manual_action(session_id):
+ """Execute a manual Playwright action during takeover mode."""
+ runner, error = _get_runner(session_id)
+ if error:
+ return error
+
+ if runner.state != AgentState.TAKEOVER:
+ return jsonify({"error": "Not in takeover mode"}), 400
+
+ data = request.get_json(silent=True) or {}
+ action = data.get("action", {})
+ if not action or "type" not in action:
+ return jsonify({"error": "action with type is required"}), 400
+
+ runner.submit_manual_action(action)
+ return jsonify({"status": "submitted", "action": action})
+
+
+@live_agent_bp.route("/api/live_agent/stop/", methods=["POST"])
+@_login_required
+def stop_session(session_id):
+ """Stop the agent and return the trace."""
+ runner, error = _get_runner(session_id)
+ if error:
+ return error
+
+ runner.stop()
+
+ # Wait briefly for the agent to clean up
+ import threading
+ if runner._thread:
+ runner._thread.join(timeout=5)
+
+ trace = runner.get_trace()
+
+ # Colocate trace.json with the screenshots it references so the bundle
+ # is self-contained and portable, and rewrite each step's
+ # screenshot_url from an absolute filesystem path (not browser-fetchable
+ # โ this is what made static review images break) to a runner-independent
+ # served URL.
+ shots_dir = os.path.abspath(runner.screenshot_dir)
+ session_folder = os.path.dirname(shots_dir) # .../live_sessions/
+ folder_name = os.path.basename(session_folder)
+ os.makedirs(session_folder, exist_ok=True)
+
+ for step in trace.get("steps", []):
+ idx = step.get("step_index")
+ if idx is not None:
+ step["screenshot_url"] = (
+ f"/api/live_agent/saved_screenshot/{folder_name}/{idx}"
+ )
+
+ trace_path = os.path.join(session_folder, "trace.json")
+ with open(trace_path, "w", encoding="utf-8") as f:
+ json.dump(trace, f, indent=2)
+
+ return jsonify({
+ "status": "stopped",
+ "trace": trace,
+ "trace_path": trace_path,
+ })
+
+
+@live_agent_bp.route("/api/live_agent/screenshot//")
+@_login_required
+def get_screenshot(session_id, step):
+ """Serve a screenshot file for a given step."""
+ runner, error = _get_runner(session_id)
+ if error:
+ return error
+
+ steps = runner.steps
+ if step < 0 or step >= len(steps):
+ return jsonify({"error": f"Step {step} not found"}), 404
+
+ screenshot_path = os.path.abspath(steps[step].screenshot_path)
+ if not os.path.isfile(screenshot_path):
+ return jsonify({"error": "Screenshot file not found"}), 404
+
+ return send_file(screenshot_path, mimetype="image/png")
+
+
+@live_agent_bp.route(
+ "/api/live_agent/saved_screenshot//"
+)
+@_login_required
+def get_saved_screenshot(session_dir, step):
+ """Serve a screenshot for an *exported* trace (no live runner needed).
+
+ Reads task_dir/live_sessions//screenshots/step_NNN.png
+ with path-traversal containment, so saved/replayed traces render their
+ real screenshots in the static web_agent_trace review display.
+ """
+ from flask import abort
+
+ task_dir = os.path.abspath(current_app.config.get("task_dir", "."))
+ base = os.path.realpath(os.path.join(task_dir, "live_sessions"))
+ target = os.path.realpath(
+ os.path.join(base, session_dir, "screenshots", f"step_{step:03d}.png")
+ )
+ if not (target == base or target.startswith(base + os.sep)):
+ logger.warning(f"Saved-screenshot path traversal blocked: {session_dir}")
+ abort(403)
+ if not os.path.isfile(target):
+ abort(404)
+ return send_file(target, mimetype="image/png")
+
+
+@live_agent_bp.route("/api/live_agent/state/")
+@_login_required
+def get_state(session_id):
+ """Get current session state."""
+ runner, error = _get_runner(session_id)
+ if error:
+ return error
+
+ return jsonify(runner.get_state_summary())
+
+
+@live_agent_bp.route("/api/live_agent/sessions")
+@_login_required
+def list_sessions():
+ """List all active live agent sessions (admin use)."""
+ manager = _get_manager()
+ return jsonify({"sessions": manager.list_sessions()})
+
+
+def _sse_format(event_type: str, data: dict) -> str:
+ """Format an SSE message."""
+ return f"event: {event_type}\ndata: {json.dumps(data)}\n\n"
diff --git a/potato/routes_live_coding_agent.py b/potato/routes_live_coding_agent.py
new file mode 100644
index 0000000000000000000000000000000000000000..701aea67047cacf17e21a478595b87b32708d68b
--- /dev/null
+++ b/potato/routes_live_coding_agent.py
@@ -0,0 +1,307 @@
+"""
+Live Coding Agent Routes
+
+REST API endpoints for controlling live coding agent sessions.
+Mirrors routes_live_agent.py but adapted for coding agents.
+"""
+
+import json
+import logging
+import os
+import queue
+import time
+from flask import Blueprint, Response, jsonify, request, stream_with_context
+
+logger = logging.getLogger(__name__)
+
+live_coding_agent_bp = Blueprint("live_coding_agent", __name__)
+
+
+def _get_manager():
+ from .coding_agent_runner_manager import CodingAgentRunnerManager
+ return CodingAgentRunnerManager.get_instance()
+
+
+def _get_config():
+ from .server_utils.config_module import config
+ return config
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/start", methods=["POST"])
+def start_session():
+ """Start a new coding agent session."""
+ from .coding_agent_runner import CodingAgentConfig
+
+ data = request.get_json() or {}
+ task_description = data.get("task_description", "")
+ instance_id = data.get("instance_id", "")
+ user_id = data.get("user_id", request.cookies.get("user_id", "anonymous"))
+
+ if not task_description:
+ return jsonify({"error": "task_description is required"}), 400
+
+ config = _get_config()
+ agent_config = CodingAgentConfig.from_config(config)
+
+ # Override with request config if provided
+ if "config" in data:
+ req_config = data["config"]
+ if "backend_type" in req_config:
+ agent_config.backend_type = req_config["backend_type"]
+ if "ai_config" in req_config:
+ agent_config.ai_config.update(req_config["ai_config"])
+ if "working_dir" in req_config:
+ agent_config.working_dir = req_config["working_dir"]
+
+ # Set up trace directory
+ task_dir = config.get("task_dir", ".")
+ trace_dir = os.path.join(
+ task_dir, "live_coding_sessions",
+ f"{user_id}_{instance_id}_{int(time.time())}",
+ )
+
+ manager = _get_manager()
+ runner = manager.create_session(user_id, instance_id, agent_config, trace_dir)
+
+ try:
+ runner.start(task_description)
+ except Exception as e:
+ return jsonify({"error": str(e)}), 500
+
+ return jsonify({
+ "session_id": runner.session_id,
+ "state": runner.state.value,
+ "backend": agent_config.backend_type,
+ })
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/stream/")
+def stream_events(session_id):
+ """SSE event stream for a coding agent session."""
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+
+ event_queue = queue.Queue()
+
+ def listener(event_type, data):
+ event_queue.put((event_type, data))
+
+ runner.add_listener(listener)
+
+ def generate():
+ # Send initial connection event
+ yield _sse_event("connected", {
+ "session_id": session_id,
+ "state": runner.state.value,
+ "turns": len(runner.get_structured_turns()),
+ })
+
+ try:
+ while True:
+ try:
+ event_type, data = event_queue.get(timeout=30)
+ yield _sse_event(event_type, data)
+
+ if event_type in ("complete", "error"):
+ break
+ except queue.Empty:
+ # Keepalive
+ yield ": keepalive\n\n"
+ finally:
+ runner.remove_listener(listener)
+
+ return Response(
+ stream_with_context(generate()),
+ mimetype="text/event-stream",
+ headers={
+ "Cache-Control": "no-cache",
+ "Connection": "keep-alive",
+ "X-Accel-Buffering": "no",
+ },
+ )
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/pause/", methods=["POST"])
+def pause_session(session_id):
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+ runner.pause()
+ return jsonify({"state": runner.state.value})
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/resume/", methods=["POST"])
+def resume_session(session_id):
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+ runner.resume()
+ return jsonify({"state": runner.state.value})
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/instruct/", methods=["POST"])
+def instruct_session(session_id):
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+
+ data = request.get_json() or {}
+ instruction = data.get("instruction", "")
+ if not instruction:
+ return jsonify({"error": "instruction is required"}), 400
+
+ runner.inject_instruction(instruction)
+ return jsonify({"state": runner.state.value, "instruction": instruction})
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/checkpoints/")
+def get_checkpoints(session_id):
+ """List all checkpoints for a session."""
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+ return jsonify({"checkpoints": runner.get_checkpoints()})
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/rollback/", methods=["POST"])
+def rollback_session(session_id):
+ """Rollback to a specific step."""
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+
+ data = request.get_json() or {}
+ step_index = data.get("step_index")
+ if step_index is None:
+ return jsonify({"error": "step_index is required"}), 400
+
+ success = runner.rollback_to_step(int(step_index))
+ return jsonify({
+ "success": success,
+ "state": runner.state.value,
+ "remaining_turns": len(runner.get_structured_turns()),
+ })
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/diff//")
+def get_diff(session_id, step):
+ """Get diff from a step to current state."""
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+ diff = runner.get_diff_since_step(step)
+ return jsonify({"diff": diff, "step_index": step})
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/replay/", methods=["POST"])
+def replay_session(session_id):
+ """Replay from a step with optional new instructions or edited actions."""
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+
+ data = request.get_json() or {}
+ step_index = data.get("step_index")
+ if step_index is None:
+ return jsonify({"error": "step_index is required"}), 400
+
+ instructions = data.get("instructions")
+ edited_actions = data.get("edited_actions")
+
+ branch_id = runner.replay_from_step(
+ int(step_index),
+ instructions=instructions,
+ edited_actions=edited_actions,
+ )
+
+ if branch_id:
+ return jsonify({
+ "success": True,
+ "branch_id": branch_id,
+ "state": runner.state.value,
+ })
+ else:
+ return jsonify({"error": "Failed to create branch"}), 500
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/branches/")
+def get_branches(session_id):
+ """List all branches for a session."""
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+ return jsonify({"branches": runner.get_branches()})
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/switch_branch/", methods=["POST"])
+def switch_branch(session_id):
+ """Switch to a different branch."""
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+
+ data = request.get_json() or {}
+ branch_id = data.get("branch_id")
+ if not branch_id:
+ return jsonify({"error": "branch_id is required"}), 400
+
+ success = runner.switch_branch(branch_id)
+ return jsonify({
+ "success": success,
+ "active_branch": branch_id if success else None,
+ "turns": len(runner.get_structured_turns()),
+ })
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/stop/", methods=["POST"])
+def stop_session(session_id):
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+ runner.stop()
+ return jsonify({
+ "state": runner.state.value,
+ "trace": runner.get_trace(),
+ })
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/state/")
+def get_state(session_id):
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+ return jsonify(runner.get_state_summary())
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/trace/")
+def get_trace(session_id):
+ manager = _get_manager()
+ runner = manager.get_session(session_id)
+ if not runner:
+ return jsonify({"error": "Session not found"}), 404
+ return jsonify(runner.get_trace())
+
+
+@live_coding_agent_bp.route("/api/live_coding_agent/sessions")
+def list_sessions():
+ manager = _get_manager()
+ return jsonify({"sessions": manager.list_sessions()})
+
+
+def _sse_event(event_type: str, data: dict) -> str:
+ """Format a Server-Sent Event."""
+ json_data = json.dumps(data, ensure_ascii=False)
+ return f"event: {event_type}\ndata: {json_data}\n\n"
diff --git a/potato/routes_trace_ingestion.py b/potato/routes_trace_ingestion.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d7e1d679a24cae24d219d22bf6430d02e68e622
--- /dev/null
+++ b/potato/routes_trace_ingestion.py
@@ -0,0 +1,201 @@
+"""
+Trace Ingestion API Routes
+
+Provides endpoints for receiving agent traces from external platforms:
+- POST /api/traces/webhook - Generic webhook receiver
+- POST /api/traces/langsmith - LangSmith-specific webhook
+- GET /api/traces/stream - SSE stream for new trace notifications
+- GET /api/traces/status - Ingestion status and stats
+"""
+
+import json
+import logging
+import os
+import time
+from functools import wraps
+
+from flask import (
+ Blueprint,
+ Response,
+ current_app,
+ jsonify,
+ request,
+ session as flask_session,
+ redirect,
+ url_for,
+)
+
+from potato.trace_ingestion.webhook_receiver import WebhookReceiver
+from potato.trace_ingestion.sse_notifier import SSENotifier
+
+logger = logging.getLogger(__name__)
+
+trace_ingestion_bp = Blueprint("trace_ingestion", __name__)
+
+# Module-level state
+_sse_notifier = SSENotifier()
+_stats = {"received": 0, "processed": 0, "errors": 0, "last_received": None}
+
+
+def _login_required(f):
+ """Require user authentication."""
+ @wraps(f)
+ def decorated(*args, **kwargs):
+ if "username" not in flask_session:
+ return redirect(url_for("login"))
+ return f(*args, **kwargs)
+ return decorated
+
+
+def _get_webhook_receiver():
+ """Create a webhook receiver from the current app config.
+
+ Always reads from current_app.config so that each Flask app instance
+ uses its own trace_ingestion settings (important when multiple test
+ servers run in the same process).
+ """
+ ingestion_config = current_app.config.get("trace_ingestion", {})
+ api_key = ingestion_config.get("api_key", "")
+ return WebhookReceiver(api_key=api_key)
+
+
+def _inject_trace(trace: dict):
+ """Inject a normalized trace as a new annotation item."""
+ try:
+ from potato.item_state_management import get_item_state_manager
+
+ ism = get_item_state_manager()
+
+ # Save trace data to disk
+ task_dir = current_app.config.get("task_dir", ".")
+ traces_dir = os.path.join(task_dir, "ingested_traces")
+ os.makedirs(traces_dir, exist_ok=True)
+
+ trace_id = trace.get("id", f"trace_{int(time.time())}")
+ trace_path = os.path.join(traces_dir, f"{trace_id}.json")
+ with open(trace_path, "w", encoding="utf-8") as f:
+ json.dump(trace, f, indent=2)
+
+ # Add as item to the annotation queue
+ item_data = {
+ "id": trace_id,
+ "text": trace.get("task_description", "Ingested trace"),
+ **trace,
+ }
+ ism.add_item(trace_id, item_data)
+
+ _stats["processed"] += 1
+
+ # Notify connected annotators
+ ingestion_config = current_app.config.get("trace_ingestion", {})
+ if ingestion_config.get("notify_annotators", True):
+ _sse_notifier.notify_new_trace(
+ trace_id=trace_id,
+ task_description=trace.get("task_description", ""),
+ source=trace.get("metadata", {}).get("source", "webhook"),
+ )
+
+ logger.info(f"Injected trace: {trace_id}")
+ return True
+
+ except Exception as e:
+ logger.error(f"Failed to inject trace: {e}")
+ _stats["errors"] += 1
+ return False
+
+
+@trace_ingestion_bp.route("/api/traces/webhook", methods=["POST"])
+def webhook_endpoint():
+ """
+ Generic webhook endpoint for receiving agent traces.
+
+ Authentication: Bearer token or X-API-Key header.
+ """
+ receiver = _get_webhook_receiver()
+
+ # Validate authentication
+ if not receiver.validate_auth(dict(request.headers)):
+ return jsonify({"error": "Unauthorized"}), 401
+
+ payload = request.get_json(silent=True)
+ if not payload:
+ return jsonify({"error": "Invalid JSON payload"}), 400
+
+ _stats["received"] += 1
+ _stats["last_received"] = time.time()
+
+ format_hint = request.args.get("format", "auto")
+ trace = receiver.process_webhook(payload, format_hint)
+
+ if trace is None:
+ _stats["errors"] += 1
+ return jsonify({"error": "Failed to process payload"}), 422
+
+ success = _inject_trace(trace)
+ if not success:
+ return jsonify({"error": "Failed to inject trace"}), 500
+
+ return jsonify({
+ "status": "accepted",
+ "trace_id": trace.get("id"),
+ "steps": len(trace.get("steps", [])),
+ })
+
+
+@trace_ingestion_bp.route("/api/traces/langsmith", methods=["POST"])
+def langsmith_webhook():
+ """LangSmith-specific webhook endpoint."""
+ receiver = _get_webhook_receiver()
+
+ if not receiver.validate_auth(dict(request.headers)):
+ return jsonify({"error": "Unauthorized"}), 401
+
+ payload = request.get_json(silent=True)
+ if not payload:
+ return jsonify({"error": "Invalid JSON payload"}), 400
+
+ _stats["received"] += 1
+ _stats["last_received"] = time.time()
+
+ trace = receiver.process_webhook(payload, format_hint="langsmith")
+ if trace is None:
+ _stats["errors"] += 1
+ return jsonify({"error": "Failed to process LangSmith payload"}), 422
+
+ success = _inject_trace(trace)
+ if not success:
+ return jsonify({"error": "Failed to inject trace"}), 500
+
+ return jsonify({
+ "status": "accepted",
+ "trace_id": trace.get("id"),
+ "steps": len(trace.get("steps", [])),
+ })
+
+
+@trace_ingestion_bp.route("/api/traces/stream")
+@_login_required
+def trace_stream():
+ """SSE stream for real-time trace ingestion notifications."""
+ client_queue = _sse_notifier.add_client()
+
+ return Response(
+ _sse_notifier.generate_sse_stream(client_queue),
+ mimetype="text/event-stream",
+ headers={
+ "Cache-Control": "no-cache",
+ "X-Accel-Buffering": "no",
+ "Connection": "keep-alive",
+ },
+ )
+
+
+@trace_ingestion_bp.route("/api/traces/status")
+@_login_required
+def ingestion_status():
+ """Get trace ingestion status and statistics."""
+ return jsonify({
+ "enabled": True,
+ "stats": _stats,
+ "sse_clients": _sse_notifier.client_count,
+ })
diff --git a/potato/routes_web_agent.py b/potato/routes_web_agent.py
new file mode 100644
index 0000000000000000000000000000000000000000..52bf3879ddfb31fe7562c4f4940cfbb9ff3cf6f5
--- /dev/null
+++ b/potato/routes_web_agent.py
@@ -0,0 +1,192 @@
+"""
+Web Agent Recording API Routes
+
+Provides endpoints for the web agent creation mode:
+- Session management (start/end recording sessions)
+- Step saving (capture individual interaction steps)
+- Screenshot capture and storage
+"""
+
+import json
+import logging
+import os
+import threading
+import time
+import uuid
+from typing import Dict, Any
+
+from flask import Blueprint, request, jsonify, current_app, session as flask_session, redirect, url_for
+
+logger = logging.getLogger(__name__)
+
+web_agent_bp = Blueprint('web_agent', __name__)
+
+# In-memory session storage (per-process), protected by a lock
+_sessions: Dict[str, Dict[str, Any]] = {}
+_sessions_lock = threading.Lock()
+
+# Session TTL: 2 hours
+_SESSION_TTL_SECONDS = 2 * 60 * 60
+
+
+def _login_required(f):
+ """Require user authentication for web agent routes."""
+ from functools import wraps
+ @wraps(f)
+ def decorated(*args, **kwargs):
+ if 'username' not in flask_session:
+ return redirect(url_for('login'))
+ return f(*args, **kwargs)
+ return decorated
+
+
+def _cleanup_expired_sessions():
+ """Remove sessions older than TTL. Must be called with _sessions_lock held."""
+ now = time.time()
+ expired = [
+ sid for sid, sdata in _sessions.items()
+ if now - sdata.get('start_time', 0) > _SESSION_TTL_SECONDS
+ ]
+ for sid in expired:
+ del _sessions[sid]
+ logger.info(f"Cleaned up expired recording session {sid}")
+
+
+@web_agent_bp.route('/api/web_agent/start_session', methods=['POST'])
+@_login_required
+def start_session():
+ """Initialize a new web agent recording session."""
+ data = request.get_json(silent=True) or {}
+ url = data.get('url', '')
+
+ session_id = str(uuid.uuid4())[:8]
+
+ # Create screenshots directory
+ task_dir = current_app.config.get('task_dir', '.')
+ screenshots_dir = os.path.join(task_dir, 'recordings', session_id, 'screenshots')
+ os.makedirs(screenshots_dir, exist_ok=True)
+
+ with _sessions_lock:
+ _cleanup_expired_sessions()
+ _sessions[session_id] = {
+ 'id': session_id,
+ 'start_url': url,
+ 'start_time': time.time(),
+ 'steps': [],
+ 'screenshots_dir': screenshots_dir,
+ }
+
+ logger.info(f"Started recording session {session_id} for {url}")
+
+ return jsonify({
+ 'session_id': session_id,
+ 'status': 'recording',
+ })
+
+
+@web_agent_bp.route('/api/web_agent/save_step', methods=['POST'])
+@_login_required
+def save_step():
+ """Save a recorded interaction step."""
+ data = request.get_json(silent=True) or {}
+ session_id = data.get('session_id', '')
+ step = data.get('step', {})
+
+ with _sessions_lock:
+ session = _sessions.get(session_id)
+ if not session:
+ return jsonify({'error': 'Unknown session'}), 404
+
+ session['steps'].append(step)
+ step_count = len(session['steps'])
+
+ return jsonify({
+ 'status': 'ok',
+ 'step_count': step_count,
+ })
+
+
+@web_agent_bp.route('/api/web_agent/save_screenshot', methods=['POST'])
+@_login_required
+def save_screenshot():
+ """Upload a screenshot for a recording step."""
+ session_id = request.form.get('session_id', request.args.get('session_id', ''))
+ step_index = request.form.get('step_index', request.args.get('step_index', '0'))
+
+ with _sessions_lock:
+ session = _sessions.get(session_id)
+ if not session:
+ return jsonify({'error': 'Unknown session'}), 404
+ screenshots_dir = session.get('screenshots_dir', '')
+
+ if not screenshots_dir:
+ return jsonify({'error': 'No screenshots directory'}), 500
+
+ # Handle file upload
+ if 'screenshot' in request.files:
+ file = request.files['screenshot']
+ filename = f'step_{int(step_index):03d}.png'
+ filepath = os.path.join(screenshots_dir, filename)
+ file.save(filepath)
+ rel_path = os.path.relpath(filepath, current_app.config.get('task_dir', '.'))
+ return jsonify({'screenshot_url': rel_path, 'status': 'ok'})
+
+ # Handle base64 data
+ data = request.get_json(silent=True) or {}
+ b64_data = data.get('screenshot_data', '')
+ if b64_data:
+ import base64
+ filename = f'step_{int(step_index):03d}.png'
+ filepath = os.path.join(screenshots_dir, filename)
+ with open(filepath, 'wb') as f:
+ f.write(base64.b64decode(b64_data))
+ rel_path = os.path.relpath(filepath, current_app.config.get('task_dir', '.'))
+ return jsonify({'screenshot_url': rel_path, 'status': 'ok'})
+
+ return jsonify({'error': 'No screenshot data'}), 400
+
+
+@web_agent_bp.route('/api/web_agent/end_session', methods=['POST'])
+@_login_required
+def end_session():
+ """Finalize a recording session and save trace data."""
+ data = request.get_json(silent=True) or {}
+ session_id = data.get('session_id', '')
+ final_steps = data.get('steps', None)
+
+ with _sessions_lock:
+ session = _sessions.get(session_id)
+ if not session:
+ return jsonify({'error': 'Unknown session'}), 404
+
+ # Use provided steps or session's accumulated steps
+ steps = final_steps if final_steps is not None else list(session['steps'])
+ start_url = session.get('start_url', '')
+
+ # Clean up session
+ del _sessions[session_id]
+
+ # Build trace data
+ trace = {
+ 'id': f'recording_{session_id}',
+ 'task_description': data.get('task_description', ''),
+ 'site': start_url,
+ 'steps': steps,
+ }
+
+ # Save to file
+ task_dir = current_app.config.get('task_dir', '.')
+ recordings_dir = os.path.join(task_dir, 'recordings', session_id)
+ os.makedirs(recordings_dir, exist_ok=True)
+ trace_path = os.path.join(recordings_dir, 'trace.json')
+
+ with open(trace_path, 'w', encoding='utf-8') as f:
+ json.dump(trace, f, indent=2)
+
+ logger.info(f"Ended recording session {session_id}, saved {len(steps)} steps")
+
+ return jsonify({
+ 'status': 'saved',
+ 'trace_path': trace_path,
+ 'step_count': len(steps),
+ })
diff --git a/potato/search/__init__.py b/potato/search/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a84f3cff6baf93d8bceb52c8c74ccb71f0ba09a9
--- /dev/null
+++ b/potato/search/__init__.py
@@ -0,0 +1,31 @@
+"""
+Universal full-text search.
+
+SQLite FTS5 lexical search over instance text, behind a pluggable
+`SearchBackend` interface (a `VectorBackend` stub documents the contract
+for future semantic search). Not gated to QDA Mode โ admins/adjudicators
+can search any project; annotator search-and-claim is a separate, guarded
+opt-in.
+"""
+
+from .backend import Hit, SearchBackend, VectorBackend
+from .fts5 import FTS5Backend
+from .service import (
+ clear_search,
+ get_search,
+ init_search,
+ init_search_from_item_state,
+ search_settings,
+)
+
+__all__ = [
+ "Hit",
+ "SearchBackend",
+ "VectorBackend",
+ "FTS5Backend",
+ "init_search",
+ "init_search_from_item_state",
+ "get_search",
+ "clear_search",
+ "search_settings",
+]
diff --git a/potato/search/api.py b/potato/search/api.py
new file mode 100644
index 0000000000000000000000000000000000000000..75566327c9122a455d0656d5481f3efbd08eb261
--- /dev/null
+++ b/potato/search/api.py
@@ -0,0 +1,160 @@
+"""
+Search REST API (universal).
+
+Phase 1 exposes admin/adjudicator read-only search at
+``/admin/api/search``. Read-only search is safe under every assignment
+strategy and crowd backend (no self-selection), so it has no config
+guard. Annotator search-and-claim is a separate, guarded endpoint.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from flask import Blueprint, jsonify, request, session
+
+from .service import get_search, search_settings
+
+logger = logging.getLogger(__name__)
+
+search_bp = Blueprint("search", __name__)
+
+
+def _config() -> dict:
+ from potato.server_utils.config_module import config
+ return config
+
+
+def _annotator_claim_enabled() -> bool:
+ return bool(search_settings(_config()).get("annotator_claim"))
+
+
+def _serialize(hits) -> list:
+ return [
+ {"instance_id": h.instance_id, "snippet": h.snippet, "score": h.score}
+ for h in hits
+ ]
+
+
+def _is_privileged() -> bool:
+ """Admin (API key) or adjudicator may use admin search."""
+ username = session.get("username")
+ try:
+ from potato.admin import admin_dashboard
+ if admin_dashboard.check_admin_access():
+ return True
+ except Exception:
+ pass
+ if username:
+ try:
+ from potato.adjudication import get_adjudication_manager
+ adj = get_adjudication_manager()
+ if adj and adj.is_adjudicator(username):
+ return True
+ except Exception:
+ pass
+ return False
+
+
+@search_bp.route("/admin/api/search", methods=["GET"])
+def admin_search():
+ if not _is_privileged():
+ return jsonify({"error": "Admin or adjudicator access required"}), 403
+ backend = get_search()
+ if backend is None:
+ return jsonify({
+ "error": "Search is not enabled or unavailable in this "
+ "deployment.",
+ "hint": "Set search.enabled: true (requires a SQLite build "
+ "with FTS5).",
+ }), 503
+ q = (request.args.get("q") or "").strip()
+ if not q:
+ return jsonify({"error": "Query parameter 'q' is required"}), 400
+ try:
+ limit = max(1, min(int(request.args.get("limit", 50)), 500))
+ except (TypeError, ValueError):
+ limit = 50
+ hits = backend.query(q, limit=limit)
+ return jsonify({
+ "query": q, "count": len(hits), "results": _serialize(hits),
+ })
+
+
+def _annotator_ctx():
+ """(username, backend) for annotator search/claim, or an error tuple."""
+ if not _annotator_claim_enabled():
+ return None, None, ("disabled",)
+ username = session.get("username")
+ if not username:
+ return None, None, ("unauth",)
+ backend = get_search()
+ if backend is None:
+ return None, None, ("unavailable",)
+ return username, backend, None
+
+
+def _annotator_err(err):
+ if err == ("disabled",):
+ return jsonify({
+ "error": "Annotator search-and-claim is not enabled.",
+ "hint": "Set search.annotator_claim: true (subject to the "
+ "assignment-compatibility guard).",
+ }), 403
+ if err == ("unauth",):
+ return jsonify({"error": "Not authenticated"}), 401
+ if err == ("unavailable",):
+ return jsonify({"error": "Search is not enabled or unavailable."}), 503
+ return None
+
+
+@search_bp.route("/api/search", methods=["GET"])
+def annotator_search():
+ """Annotator-facing corpus search. Gated by search.annotator_claim;
+ the startup guard guarantees the assignment design is compatible."""
+ username, backend, err = _annotator_ctx()
+ if err:
+ return _annotator_err(err)
+ q = (request.args.get("q") or "").strip()
+ if not q:
+ return jsonify({"error": "Query parameter 'q' is required"}), 400
+ try:
+ limit = max(1, min(int(request.args.get("limit", 50)), 200))
+ except (TypeError, ValueError):
+ limit = 50
+ hits = backend.query(q, limit=limit)
+ return jsonify({
+ "query": q, "count": len(hits), "results": _serialize(hits),
+ })
+
+
+@search_bp.route("/api/search/claim", methods=["POST"])
+def annotator_claim():
+ """Pull a matching instance into the requesting annotator's queue."""
+ username, backend, err = _annotator_ctx()
+ if err:
+ return _annotator_err(err)
+ data = request.get_json(silent=True) or {}
+ instance_id = data.get("instance_id")
+ if not instance_id:
+ return jsonify({"error": "instance_id is required"}), 400
+ try:
+ from potato.item_state_management import get_item_state_manager
+ from potato.user_state_management import get_user_state_manager
+ ism = get_item_state_manager()
+ try:
+ item = ism.get_item(str(instance_id))
+ except Exception:
+ return jsonify({"error": f"Unknown instance {instance_id}"}), 404
+ if item is None:
+ return jsonify({"error": f"Unknown instance {instance_id}"}), 404
+ user_state = get_user_state_manager().get_user_state(username)
+ already = instance_id in user_state.get_assigned_instance_ids()
+ user_state.assign_instance(item)
+ except Exception as e:
+ logger.error(f"Claim failed for {instance_id}: {e}")
+ return jsonify({"error": "Could not claim instance"}), 500
+ return jsonify({
+ "claimed": str(instance_id),
+ "already_assigned": bool(already),
+ })
diff --git a/potato/search/backend.py b/potato/search/backend.py
new file mode 100644
index 0000000000000000000000000000000000000000..0a461b61f2e77bf3989b8c209b36f67dc421b6fb
--- /dev/null
+++ b/potato/search/backend.py
@@ -0,0 +1,66 @@
+"""
+Search backend abstraction (universal).
+
+A pluggable interface so lexical (FTS5, ships now) and future semantic
+(vector) search are interchangeable behind one contract. Not gated to
+QDA Mode โ useful in any project for locating instances.
+
+Contract:
+ available() -> bool Is this backend usable in this environment?
+ index(rows) -> int (Re)build the index from (id, text) pairs;
+ returns the number of documents indexed.
+ query(q, limit) -> [Hit] Ranked matches for a user query string.
+"""
+
+from __future__ import annotations
+
+import abc
+from dataclasses import dataclass
+from typing import Iterable, List, Tuple
+
+
+@dataclass(frozen=True)
+class Hit:
+ """One search result."""
+ instance_id: str
+ snippet: str
+ score: float # lower rank value = better for FTS5; normalized per backend
+
+
+class SearchBackend(abc.ABC):
+ name: str = "base"
+
+ @abc.abstractmethod
+ def available(self) -> bool:
+ """Whether this backend can run here (e.g. FTS5 compiled in)."""
+
+ @abc.abstractmethod
+ def index(self, rows: Iterable[Tuple[str, str]]) -> int:
+ """(Re)build the index from (instance_id, text) pairs."""
+
+ @abc.abstractmethod
+ def query(self, q: str, limit: int = 50) -> List[Hit]:
+ """Return up to *limit* ranked hits for query string *q*."""
+
+
+class VectorBackend(SearchBackend):
+ """Placeholder for a future dense/semantic backend.
+
+ Documents the contract a vector backend must satisfy so it can be
+ dropped in without touching callers: it would embed instance text
+ (reusing potato/ai embedding endpoints), persist vectors alongside
+ project.sqlite, and implement ``query`` as nearest-neighbour search.
+ Not implemented in this phase โ ``available()`` is False so callers
+ fall back to FTS5.
+ """
+
+ name = "vector"
+
+ def available(self) -> bool:
+ return False
+
+ def index(self, rows: Iterable[Tuple[str, str]]) -> int: # pragma: no cover
+ raise NotImplementedError("Vector search backend not implemented yet")
+
+ def query(self, q: str, limit: int = 50) -> List[Hit]: # pragma: no cover
+ raise NotImplementedError("Vector search backend not implemented yet")
diff --git a/potato/search/fts5.py b/potato/search/fts5.py
new file mode 100644
index 0000000000000000000000000000000000000000..c8a64fcff55124ff599a142b20c866a4fb1c2ff7
--- /dev/null
+++ b/potato/search/fts5.py
@@ -0,0 +1,115 @@
+"""
+SQLite FTS5 lexical search backend (universal).
+
+Stores a standalone `instance_fts` virtual table in the shared
+`/project.sqlite` (same DB as memos; different table). The
+table is created lazily so a SQLite build without FTS5 simply reports
+``available() == False`` instead of erroring at import/migration time.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from typing import Iterable, List, Tuple
+
+from potato.persistence import get_db
+
+from .backend import Hit, SearchBackend
+
+logger = logging.getLogger(__name__)
+
+_TOKEN_RE = re.compile(r"[^\w]+", re.UNICODE)
+
+# Snippet match delimiters. We use the STX/ETX control characters as
+# sentinels instead of visible punctuation ('[' / ']'): they never occur
+# in real instance text, survive JSON transport, and let the frontend
+# escape the snippet first (XSS-safe) and only then swap the sentinels for
+# a highlight. Visible brackets would read as literal typos.
+SNIPPET_OPEN = "\x02"
+SNIPPET_CLOSE = "\x03"
+
+
+def _to_match_query(q: str) -> str:
+ """Turn arbitrary user input into a safe FTS5 MATCH expression.
+
+ FTS5 MATCH has its own syntax; raw punctuation/quotes raise errors.
+ We tokenize, drop empties, and AND the tokens as quoted prefix terms
+ so partial words still match and nothing is interpreted as syntax.
+ """
+ tokens = [t for t in _TOKEN_RE.split(q or "") if t]
+ if not tokens:
+ return ""
+ return " ".join(f'"{t}"*' for t in tokens)
+
+
+class FTS5Backend(SearchBackend):
+ name = "fts5"
+
+ def __init__(self, task_dir: str):
+ self.task_dir = task_dir
+ self._available = None # lazy-detected, then cached
+
+ # -- internals ---------------------------------------------------------
+
+ def _conn(self):
+ return get_db(self.task_dir)
+
+ def _ensure_table(self, conn) -> bool:
+ try:
+ conn.execute(
+ "CREATE VIRTUAL TABLE IF NOT EXISTS instance_fts "
+ "USING fts5(instance_id UNINDEXED, body)"
+ )
+ return True
+ except Exception as e:
+ logger.warning(f"FTS5 unavailable, search disabled: {e}")
+ return False
+
+ # -- SearchBackend -----------------------------------------------------
+
+ def available(self) -> bool:
+ if self._available is None:
+ self._available = self._ensure_table(self._conn())
+ return self._available
+
+ def index(self, rows: Iterable[Tuple[str, str]]) -> int:
+ conn = self._conn()
+ if not self._ensure_table(conn):
+ return 0
+ conn.execute("DELETE FROM instance_fts")
+ n = 0
+ for instance_id, text in rows:
+ conn.execute(
+ "INSERT INTO instance_fts (instance_id, body) VALUES (?, ?)",
+ (str(instance_id), text or ""),
+ )
+ n += 1
+ conn.commit()
+ logger.info(f"FTS5 indexed {n} instances for {self.task_dir}")
+ return n
+
+ def query(self, q: str, limit: int = 50) -> List[Hit]:
+ match = _to_match_query(q)
+ if not match:
+ return []
+ conn = self._conn()
+ if not self._ensure_table(conn):
+ return []
+ try:
+ cur = conn.execute(
+ """SELECT instance_id,
+ snippet(instance_fts, 1, ?, ?, 'โฆ', 12) AS snip,
+ rank AS r
+ FROM instance_fts
+ WHERE instance_fts MATCH ?
+ ORDER BY rank
+ LIMIT ?""",
+ (SNIPPET_OPEN, SNIPPET_CLOSE, match, int(limit)),
+ )
+ return [Hit(instance_id=row["instance_id"],
+ snippet=row["snip"] or "",
+ score=float(row["r"])) for row in cur.fetchall()]
+ except Exception as e:
+ logger.warning(f"FTS5 query failed for {q!r}: {e}")
+ return []
diff --git a/potato/search/service.py b/potato/search/service.py
new file mode 100644
index 0000000000000000000000000000000000000000..402be035168bd1ef582943beb7a6e7803f2bf15c
--- /dev/null
+++ b/potato/search/service.py
@@ -0,0 +1,120 @@
+"""
+Search service (universal).
+
+Resolves the `search:` config block, builds the configured backend, and
+holds a process singleton so the index is built once on server start and
+reused per request. Mirrors the init/get/clear pattern of the other
+managers.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any, Dict, Iterable, Optional, Tuple
+
+from .backend import SearchBackend
+from .fts5 import FTS5Backend
+
+logger = logging.getLogger(__name__)
+
+_SEARCH: Optional[SearchBackend] = None
+_LOCK = threading.Lock()
+
+_DEFAULTS = {
+ "enabled": True, # universal โ on by default
+ "backend": "fts5",
+ "max_instances": 100000,
+ "annotator_claim": False, # annotator search-and-claim is opt-in
+}
+
+
+def search_settings(config: Dict[str, Any]) -> Dict[str, Any]:
+ """Resolved search settings with defaults applied."""
+ raw = config.get("search")
+ raw = raw if isinstance(raw, dict) else {}
+ s = dict(_DEFAULTS)
+ for k in _DEFAULTS:
+ if k in raw:
+ s[k] = raw[k]
+ return s
+
+
+def _build(config: Dict[str, Any]) -> Optional[SearchBackend]:
+ s = search_settings(config)
+ if not s["enabled"]:
+ return None
+ task_dir = config.get("task_dir", ".")
+ if s["backend"] == "fts5":
+ be = FTS5Backend(task_dir)
+ if be.available():
+ return be
+ logger.warning("FTS5 not available in this SQLite build; "
+ "search disabled.")
+ return None
+ logger.warning(f"Unknown search backend {s['backend']!r}; search disabled.")
+ return None
+
+
+def init_search(
+ config: Dict[str, Any],
+ rows: Optional[Iterable[Tuple[str, str]]] = None,
+) -> Optional[SearchBackend]:
+ """Build the backend singleton and (optionally) index *rows*.
+
+ Returns None when search is disabled/unavailable. Calling twice keeps
+ the existing singleton (but re-indexes if rows are provided)."""
+ global _SEARCH
+ with _LOCK:
+ if _SEARCH is None:
+ _SEARCH = _build(config)
+ if _SEARCH is not None and rows is not None:
+ try:
+ _SEARCH.index(rows)
+ except Exception as e:
+ logger.error(f"Search index build failed: {e}")
+ return _SEARCH
+
+
+def _rows_from_item_state(config: Dict[str, Any]):
+ """Yield (instance_id, text) for every loaded instance, using the
+ config's text_key. Bounded by search.max_instances."""
+ from potato.item_state_management import get_item_state_manager
+
+ text_key = (config.get("item_properties") or {}).get("text_key", "text")
+ cap = search_settings(config)["max_instances"]
+ ism = get_item_state_manager()
+ for i, iid in enumerate(ism.get_instance_ids()):
+ if i >= cap:
+ logger.warning(
+ f"search.max_instances ({cap}) reached; not indexing the rest")
+ break
+ data = ism.get_item(iid).get_data()
+ if isinstance(data, dict):
+ text = data.get(text_key) or ism.get_item(iid).get_text()
+ else:
+ text = str(data)
+ yield str(iid), text if isinstance(text, str) else str(text)
+
+
+def init_search_from_item_state(
+ config: Dict[str, Any]
+) -> Optional[SearchBackend]:
+ """Server-start entry point: build the backend and index all loaded
+ instances. No-op when search is disabled/unavailable."""
+ settings = search_settings(config)
+ if not settings["enabled"]:
+ logger.info("Search disabled in config")
+ return None
+ return init_search(config, rows=_rows_from_item_state(config))
+
+
+def get_search() -> Optional[SearchBackend]:
+ return _SEARCH
+
+
+def clear_search() -> None:
+ """Reset the singleton. Tests only."""
+ global _SEARCH
+ with _LOCK:
+ _SEARCH = None
diff --git a/potato/server_utils/__init__.py b/potato/server_utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/potato/server_utils/admin_key.py b/potato/server_utils/admin_key.py
new file mode 100644
index 0000000000000000000000000000000000000000..30f8ffb4d82fc153a7d8d0fa4f3807442d5c7f74
--- /dev/null
+++ b/potato/server_utils/admin_key.py
@@ -0,0 +1,104 @@
+"""
+Admin API Key Resolution
+
+Shared utility for resolving the admin API key from config, environment,
+or auto-generated key file. Used by both routes.py and admin.py to ensure
+consistent authentication across all admin endpoints.
+"""
+
+import os
+import hmac
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Cache for auto-generated admin API key
+_generated_admin_api_key = None
+
+
+def get_admin_api_key(config):
+ """Get the admin API key from config, environment variable, or auto-generate one.
+
+ Priority order:
+ 1. Config file: admin_api_key setting
+ 2. Environment variable: POTATO_ADMIN_API_KEY
+ 3. Auto-generated: Creates a random key and saves it to {task_dir}/admin_api_key.txt
+
+ Args:
+ config: The application config dict.
+
+ Returns:
+ str or None: The admin API key, or None if generation fails.
+ """
+ global _generated_admin_api_key
+
+ # Check config first
+ configured_key = config.get("admin_api_key")
+ if configured_key:
+ return configured_key
+
+ # Check environment variable
+ env_key = os.environ.get("POTATO_ADMIN_API_KEY")
+ if env_key:
+ return env_key
+
+ # Return cached generated key if we have one
+ if _generated_admin_api_key:
+ return _generated_admin_api_key
+
+ # Auto-generate a key and save it to task directory
+ task_dir = config.get("task_dir", ".")
+ if not task_dir:
+ task_dir = "."
+
+ key_file_path = os.path.join(task_dir, "admin_api_key.txt")
+
+ # Check if a key file already exists (from previous run)
+ if os.path.exists(key_file_path):
+ try:
+ with open(key_file_path, 'r', encoding='utf-8') as f:
+ existing_key = f.read().strip()
+ if existing_key:
+ _generated_admin_api_key = existing_key
+ logger.info(f"Loaded existing admin API key from {key_file_path}")
+ return _generated_admin_api_key
+ except Exception as e:
+ logger.warning(f"Could not read existing admin API key file: {e}")
+
+ # Generate a new key
+ import secrets
+ _generated_admin_api_key = secrets.token_urlsafe(32)
+
+ # Save to file
+ try:
+ with open(key_file_path, 'w', encoding='utf-8') as f:
+ f.write(_generated_admin_api_key)
+ logger.info(f"Generated admin API key and saved to {key_file_path}")
+ logger.info(f"Use this key to access the admin dashboard at /admin")
+ except Exception as e:
+ logger.warning(f"Could not save admin API key to file: {e}")
+ logger.info(f"Auto-generated admin API key (not persisted): {_generated_admin_api_key}")
+
+ return _generated_admin_api_key
+
+
+def validate_admin_api_key(provided_key, config):
+ """Validate an admin API key against the configured or auto-generated key.
+
+ Args:
+ provided_key: The API key provided in the request.
+ config: The application config dict.
+
+ Returns:
+ bool: True if the key is valid or debug mode is enabled.
+ """
+ if config.get("debug", False):
+ return True
+
+ expected_key = get_admin_api_key(config)
+ if not expected_key:
+ logger.warning("Could not obtain admin API key")
+ return False
+
+ # Use constant-time comparison to prevent timing attacks
+ return hmac.compare_digest(str(provided_key or ""), expected_key)
diff --git a/potato/server_utils/arg_utils.py b/potato/server_utils/arg_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..767e186881f093053d7c526520aade7d723eb3de
--- /dev/null
+++ b/potato/server_utils/arg_utils.py
@@ -0,0 +1,171 @@
+"""
+Utility functions around parsing arguments.
+"""
+
+from argparse import ArgumentParser
+
+
+def arguments():
+ """
+ Creates and returns the arg parser for Potato on the command line.
+ """
+ parser = ArgumentParser()
+ parser.set_defaults(show_path=False, show_similarity=False)
+
+ parser.add_argument(
+ "mode",
+ choices=['start', 'migrate', 'reset-password', 'codebook'],
+ help="set the mode when potato is used, currently supporting: start, migrate, reset-password, codebook",
+ default="start",
+ )
+
+ parser.add_argument("config_file")
+
+ parser.add_argument(
+ "-p",
+ "--port",
+ action="store",
+ type=int,
+ dest="port",
+ help="The port to run on",
+ default=None,
+ )
+
+ parser.add_argument(
+ "-v", "--verbose", action="store_true", help="Report verbose output", default=False
+ )
+
+ parser.add_argument(
+ "--debug", action="store_true", help="Launch in debug mode with no login", default=False
+ )
+
+ parser.add_argument(
+ "--debug-log",
+ action="store",
+ type=str,
+ dest="debug_log",
+ choices=['all', 'ui', 'server', 'none'],
+ help="Control debug logging: 'all' (UI and server), 'ui' (frontend only), 'server' (backend only), 'none' (disable)",
+ default=None,
+ )
+
+ parser.add_argument(
+ "--debug-phase",
+ action="store",
+ type=str,
+ dest="debug_phase",
+ help="Skip directly to a specific phase (e.g., 'annotation', 'poststudy') or page name. Requires --debug flag.",
+ default=None,
+ )
+
+ parser.add_argument(
+ "--veryVerbose",
+ action="store_true",
+ dest="very_verbose",
+ help="Report very verbose output",
+ default=False,
+ )
+
+ parser.add_argument(
+ "--with-custom-js",
+ action="store_true",
+ dest="customjs",
+ help="Use a custom js module served from vite."
+ )
+
+ parser.add_argument(
+ "--custom-js-hostname",
+ action="store",
+ type=str,
+ dest="customjs_hostname",
+ help="custom hostname for potato.js serving",
+ default=None,
+ )
+
+ parser.add_argument(
+ "--require-password",
+ action="store",
+ type=lambda x: str(x).lower() == 'true',
+ dest="require_password",
+ help="Whether to require password authentication (true/false). If not specified, uses config file value.",
+ default=None,
+ )
+
+ parser.add_argument(
+ "--persist-sessions",
+ action="store_true",
+ dest="persist_sessions",
+ help="Enable session persistence between server restarts (default: False)",
+ default=False,
+ )
+
+ # SSL arguments (from master branch)
+ parser.add_argument(
+ "--ssl-cert",
+ action="store",
+ type=str,
+ dest="ssl_cert",
+ help="custom ssl cert location (should end in .pem)",
+ default=None
+ )
+
+ parser.add_argument(
+ "--ssl-key",
+ action="store",
+ type=str,
+ dest="ssl_key",
+ help="custom ssl key location (should end in .pem)",
+ default=None
+ )
+
+ # Migration-specific arguments
+ parser.add_argument(
+ "--to-v2",
+ action="store_true",
+ dest="to_v2",
+ help="[migrate mode] Migrate configuration to v2 format",
+ default=False,
+ )
+
+ parser.add_argument(
+ "--output", "-o",
+ dest="output_file",
+ help="[migrate mode] Output file path (default: print to stdout)",
+ default=None,
+ )
+
+ parser.add_argument(
+ "--in-place", "-i",
+ action="store_true",
+ dest="in_place",
+ help="[migrate mode] Modify the config file in place",
+ default=False,
+ )
+
+ parser.add_argument(
+ "--dry-run",
+ action="store_true",
+ dest="dry_run",
+ help="[migrate mode] Show what changes would be made without applying them",
+ default=False,
+ )
+
+ parser.add_argument(
+ "--quiet", "-q",
+ action="store_true",
+ dest="quiet",
+ help="[migrate mode] Suppress informational output",
+ default=False,
+ )
+
+ # Password reset arguments
+ parser.add_argument(
+ "--username",
+ action="store",
+ type=str,
+ dest="username",
+ help="[reset-password mode] Username to reset password for",
+ default=None,
+ )
+
+ return parser.parse_args()
diff --git a/potato/server_utils/config_module.py b/potato/server_utils/config_module.py
new file mode 100644
index 0000000000000000000000000000000000000000..976ed451210655f1fa040be8c7ec9ccb699eb5e5
--- /dev/null
+++ b/potato/server_utils/config_module.py
@@ -0,0 +1,5129 @@
+"""
+Config module with enhanced security validation and error handling.
+"""
+
+import yaml
+import os
+import logging
+import re
+import codecs
+from pathlib import Path
+from typing import Dict, Any, List, Optional, Tuple
+from urllib.parse import urlparse
+import json
+
+config = {}
+
+
+def clear_config():
+ """Clear the global config dictionary. Used for testing to ensure clean state."""
+ global config
+ config.clear()
+
+
+# Use centralized logging - avoid duplicate basicConfig
+logger = logging.getLogger(__name__)
+
+
+class ConfigValidationError(Exception):
+ """Custom exception for configuration validation errors."""
+ pass
+
+
+class ConfigSecurityError(Exception):
+ """Custom exception for configuration security violations."""
+ pass
+
+
+import difflib
+
+# ============================================================================
+# Known config key schema (hierarchical)
+# Keys map to None (leaf), set (known sub-keys), or dict (nested schema).
+# Used by validate_unknown_keys() to warn about typos at all nesting levels.
+# ============================================================================
+KNOWN_CONFIG_KEYS = {
+ # === Core / required ===
+ "item_properties": {
+ "id_key", "text_key", "category_key", "kwargs",
+ },
+ "data_files": None,
+ "task_dir": None,
+ "output_annotation_dir": None,
+ "output_annotation_format": None,
+ "annotation_task_name": None,
+ "task_description": None,
+ "annotation_task_description": None,
+
+ # === Data sources ===
+ "data_directory": None,
+ "data_directory_encoding": None,
+ "data_sources": None,
+ "data_cache": {"enabled", "ttl_seconds", "max_size_mb"},
+ "watch_data_directory": None,
+ "watch_poll_interval": None,
+ "partial_loading": None,
+
+ # === Annotation ===
+ "annotation_schemes": None,
+ "phases": None,
+ "output_annotation_format": None,
+
+ # === Auth / login ===
+ "authentication": {
+ "method", "providers", "user_identity_field", "database_url",
+ "user_config_path", "auto_register", "allow_local_login",
+ "allowed_domain", "allowed_domains", "allowed_org",
+ },
+ "login": {"type", "url_argument", "auto_redirect_delay", "auto_redirect_on_completion"},
+ "user_config": {"allow_all_users", "users"},
+ "require_password": None,
+ "require_no_password": None,
+ "secret_key": None,
+
+ # === Server ===
+ "server": {"port", "host", "debug"},
+ "port": None,
+ "host": None,
+ "customjs": None,
+ "customjs_hostname": None,
+ "site_dir": None,
+ "site_file": None,
+ "persist_sessions": None,
+ "session_lifetime_days": None,
+ "base_html_template": None,
+
+ # === Quality control ===
+ "attention_checks": {
+ "enabled", "items_file", "frequency", "probability",
+ "min_response_time", "failure_handling",
+ },
+ "gold_standards": {
+ "enabled", "items_file", "mode", "frequency",
+ "accuracy", "auto_promote",
+ },
+ "gold_standards_file": None,
+ "pre_annotation": {
+ "enabled", "field", "highlight_low_confidence",
+ "agreement_metrics", "predictions_file",
+ "allow_modification", "show_confidence",
+ },
+ "agreement_metrics": {"min_overlap", "refresh_interval", "enabled"},
+ "quality_control": None,
+
+ # === AI ===
+ "ai_support": {
+ "enabled", "endpoint_type", "ai_config_file", "ai_config",
+ "option_highlighting", "features", "cache_config",
+ },
+ "chat_support": {
+ "enabled", "endpoint_type", "ai_config", "ui",
+ },
+
+ # === Advanced features ===
+ "training": {
+ "enabled", "data_file", "annotation_schemes",
+ "passing_criteria", "feedback", "failure_action",
+ },
+ "active_learning": {
+ "enabled", "classifier", "vectorizer",
+ "min_annotations_per_instance", "min_instances_for_training",
+ "max_instances_to_reorder", "update_frequency",
+ "resolution_strategy", "random_sample_percent", "schema_names",
+ "database", "model_persistence", "llm", "query_strategy",
+ "hybrid_weights", "cold_start_strategy", "confidence_method",
+ "classifier_params", "vectorizer_params", "calibrate_probabilities",
+ "bald_params", "use_icl_ensemble", "icl_ensemble_params",
+ "annotation_routing", "routing_thresholds",
+ },
+ "category_assignment": {
+ "enabled", "category_key", "qualification", "fallback", "dynamic",
+ },
+ "batch_assignment": {
+ "groups", "annotator_key",
+ },
+ "diversity_ordering": {
+ "enabled", "model_name", "num_clusters", "items_per_cluster",
+ "auto_clusters", "prefill_count", "batch_size",
+ "recluster_threshold", "preserve_visited",
+ "trigger_ai_prefetch", "cache_dir",
+ },
+ "diversity_config": None,
+ "embedding_visualization": {
+ "enabled", "sample_size", "include_all_annotated",
+ "embedding_model", "image_embedding_model", "umap", "label_source",
+ },
+ "adjudication": {
+ "enabled", "adjudicator_users", "min_annotations",
+ "agreement_threshold", "fast_decision_warning_ms",
+ "error_taxonomy", "similarity",
+ "require_notes_on_override", "show_agreement_scores",
+ "show_annotator_names",
+ "output_subdir", "require_confidence",
+ "show_all_items", "show_timing_data",
+ },
+ "database": {"type", "host", "database", "username", "password", "port",
+ "pool_size", "pool_timeout", "connection_string"},
+ "bws_config": {
+ "tuple_size", "num_tuples", "seed", "min_item_appearances", "scoring",
+ },
+ "ibws_config": {
+ "tuple_size", "max_rounds", "seed", "scoring_method",
+ "tuples_per_item_per_round",
+ },
+ "mace": {
+ "enabled", "min_annotations_per_item", "trigger_every_n", "num_restarts",
+ "min_items", "num_iters",
+ },
+ "icl_labeling": None,
+ "llm_labeling": None,
+
+ # === UI & layout ===
+ "ui": None,
+ "ui_config": None,
+ "layout": {"grid", "breakpoints", "groups", "order", "styling"},
+ "instance_display": {"fields", "layout", "resizable"},
+ "format_handling": {"enabled", "default_format", "pdf", "spreadsheet"},
+ "ui_language": {
+ "html_lang", "html_dir",
+ "next_button", "previous_button", "submit_button", "go_button",
+ "retry_button", "logout",
+ "labeled_badge", "in_progress_badge", "not_labeled_badge",
+ "progress_label", "loading", "error_heading",
+ "adjudicate", "codebook", "instructions_heading",
+ "text_to_annotate", "video_to_annotate", "audio_to_annotate",
+ "login_title", "login_subtitle_password", "login_subtitle_username",
+ "sign_in_tab", "register_tab",
+ "username_label", "password_label",
+ "sign_in_button", "continue_button", "register_button",
+ "forgot_password", "username_placeholder",
+ "choose_username_placeholder", "create_password_placeholder",
+ "sign_in_with", "or_divider",
+ "powered_by", "cite_us",
+ },
+ "base_css": None,
+ "ui_debug": None,
+ "hide_navbar": None,
+ "task_layout": None,
+
+ # === Content ===
+ "annotation_instructions": None,
+ "annotation_codebook_url": None,
+ "custom_footer_html": None,
+ "header_file": None,
+ "header_logo": None,
+
+ # === Annotation features ===
+ "keyword_highlight_settings": None,
+ "keyword_highlights_file": None,
+ "highlight_linebreaks": None,
+ "list_as_text": {"text_list_prefix_type", "horizontal", "alternating_shading"},
+ "jumping_to_id_disabled": None,
+ "horizontal_key_bindings": None,
+ "completion_code": None,
+ "allow_phase_back_navigation": None,
+ "require_fully_annotated": None,
+ "export_include_phase_data": None,
+ "export_annotation_format": None,
+ "auto_export_interval": None,
+
+ # === Media ===
+ "audio_annotation": {
+ "waveform_cache_dir", "waveform_look_ahead", "waveform_cache_max_size",
+ "client_fallback_max_duration",
+ },
+ "spectrogram": None,
+ "media_directory": None,
+ "default_video_fps": None,
+
+ # === External integrations ===
+ "mturk": None,
+ "prolific": {
+ "config_file_path", "token", "study_id",
+ "max_concurrent_sessions", "workload_checker_period",
+ "completion_code", "sandbox_mode",
+ },
+ "webhooks": {"enabled", "endpoints"},
+ "trace_ingestion": {"enabled", "sources", "api_key", "notify_annotators"},
+ "judge_alignment": {"enabled", "ai_support", "schemas", "few_shot", "inline"},
+ # Judge Calibration: LLM-as-judge auto-labeling + blind human calibration.
+ # Leaf sub-dicts (sampling/human/calibration/output) are validated by
+ # validate_judge_calibration_config(); kept shallow here to avoid
+ # unknown-key churn while the feature stabilizes.
+ "judge_calibration": {
+ "enabled", "prompt", "models", "k_samples", "max_items", "fraction",
+ "sampling", "human", "schemas", "calibration", "output", "state_dir",
+ },
+ "triage": {"enabled", "order", "default_priority", "show_badge",
+ "signal_field", "invert_signal", "rules"},
+ "huggingface_backup": None,
+
+ # === Debug / logging ===
+ "debug": None,
+ "debug_phase": None,
+ "server_debug": None,
+ "verbose": None,
+ "very_verbose": None,
+ "debug_log": None,
+
+ # === Agent ===
+ "live_agent": None,
+ "live_coding_agent": None,
+ "agent_proxy": None,
+
+ # === Legacy / multi-task ===
+ "surveyflow": None,
+ "prestudy": None,
+ "automatic_assignment": None,
+
+ # === Other ===
+ "random_seed": None,
+ "max_annotations_per_user": None,
+ # Deprecated alias of num_annotators_per_item (int form). Still accepted
+ # for backwards compatibility; emits a warning when both are set.
+ "max_annotations_per_item": None,
+ # Canonical key for heterogeneous coverage. Accepts either:
+ # int โ same cap for every item (legacy behavior)
+ # dict โ { default, overlap_sample: {fraction, count, stratify_by, seed},
+ # adaptive: {enabled, disagreement_threshold, boost_to}, min }
+ "num_annotators_per_item": None,
+ "min_annotators_per_instance": None,
+ # Per-annotator workload caps:
+ # { default: int, by_user: {user_id: int}, by_user_role: {role: int} }
+ "per_annotator_quota": None,
+ # qda_mode sub-keys are deliberately leaf (None): validation stops at
+ # memos/codebook and does NOT recurse into their sub-keys. This is
+ # intentional forward-compat โ parse_qda_mode_config() routes any
+ # unrecognized qda_mode.* keys into `extras` so configs can declare
+ # not-yet-shipped blocks (cases/queries/smart_codes/network/media_sync)
+ # without tripping unknown-key warnings. The tradeoff: a typo like
+ # qda_mode.memos.enabledd is silently accepted. Revisit (deepen the
+ # schema) once those sub-blocks ship and their shapes are stable.
+ "qda_mode": {
+ "enabled": None,
+ "memos": None,
+ "codebook": None,
+ # Sub-blocks reserved for later phases:
+ # "cases", "queries", "smart_codes", "network", "media_sync"
+ },
+ # Universal annotation UI feature toggles (not QDA-gated). `memos`
+ # turns the memo sidebar on/off (default off in standard mode, on for
+ # qda_mode/solo_mode); `visibility` is the default new-memo visibility.
+ "annotation_ui": {
+ "memos": None,
+ "visibility": None,
+ },
+ # Universal full-text search (FTS5). Read-only admin search is always
+ # safe; `annotator_claim` opt-in is governed by a startup
+ # compatibility guard (see validate_search_assignment_compat).
+ "search": {
+ "enabled": None,
+ "backend": None,
+ "max_instances": None,
+ "annotator_claim": None,
+ },
+ # Universal codebook. `mode` (fixed|extensible|open) governs whether
+ # annotators may add codes on the fly; resolved via
+ # get_codebook_mode() (defaults: qda/solo -> open, standard -> fixed;
+ # a crowd backend force-locks fixed). Per-scheme opt-in is the
+ # scheme-level `codebook: true` key.
+ "codebook": {
+ "enabled": None,
+ "mode": None,
+ },
+ # Top-level convenience scalar mirroring codebook.mode.
+ "codebook_mode": None,
+ # In-vivo coding (D): single key that, with text selected in a
+ # codebook-backed span scheme, opens the "code from selection"
+ # composer. Default 'i'; only meaningful when a codebook span
+ # scheme exists. (Schema value is None = scalar/any-value key; the
+ # 'i' default lives in the defaults map, not here.)
+ "codebook_invivo_key": None,
+ # Universal cases: group instances into units of analysis. `key`
+ # names the item-data field to group on; `auto_detect` lets QDA
+ # scan participant_id/respondent_id/case_id; `attributes` lifts
+ # item fields onto the case for crosstabs.
+ "cases": {
+ "enabled": None,
+ "key": None,
+ "auto_detect": None,
+ "attributes": None,
+ },
+ "solo_mode": {
+ "enabled": None,
+ "labeling_models": None,
+ "revision_models": None,
+ "embedding": None,
+ "uncertainty": None,
+ "thresholds": None,
+ "instance_selection": None,
+ "batches": None,
+ "prompt_optimization": None,
+ "edge_case_rules": None,
+ "labeling_functions": None,
+ "confidence_routing": None,
+ "confusion_analysis": None,
+ "state_dir": None,
+ "refinement_loop": {
+ "enabled",
+ "trigger_interval",
+ "min_improvement",
+ "max_cycles",
+ "patience",
+ "auto_apply_suggestions",
+ "refinement_strategy",
+ "validation_split_ratio",
+ "eval_sample_size",
+ "num_candidates",
+ "min_val_size",
+ "max_consecutive_failures",
+ "dry_run",
+ "require_approval",
+ "min_val_improvement",
+ "eval_temperature",
+ "prefer_consistent_disagreements",
+ },
+ },
+ "admin_api_key": None,
+ "alert_time_each_instance": None,
+ "assignment_strategy": None,
+ "reclaim_stale_assignments": None,
+ "instance_reclaim": None,
+ "max_session_seconds": None,
+ "env_substitution": None,
+
+ # === Internal (set by system, not user) ===
+ "config_file": None,
+ "__config_file__": None,
+ "_bws_pool_items": None,
+}
+
+
+def validate_unknown_keys(config_data, schema=None, path=""):
+ """Recursively warn about unrecognized config keys and suggest corrections.
+
+ Args:
+ config_data: The config dict (or sub-dict) to validate.
+ schema: The known-keys schema for this level (defaults to KNOWN_CONFIG_KEYS).
+ path: Dot-separated path prefix for nested key reporting (e.g., "training").
+ """
+ if schema is None:
+ schema = KNOWN_CONFIG_KEYS
+
+ if not isinstance(config_data, dict):
+ return
+
+ known_keys = set(schema.keys()) if isinstance(schema, dict) else schema
+ unknown_keys = set(config_data.keys()) - known_keys
+
+ for key in sorted(unknown_keys):
+ full_key = f"{path}.{key}" if path else key
+ matches = difflib.get_close_matches(key, known_keys, n=3, cutoff=0.6)
+ if matches:
+ suggestions = ", ".join(f"'{m}'" for m in matches)
+ logger.warning(
+ "Unrecognized config key '%s'. Did you mean: %s?",
+ full_key, suggestions
+ )
+ else:
+ logger.warning(
+ "Unrecognized config key '%s'. This key will be ignored.",
+ full_key
+ )
+
+ # Recurse into nested dicts that have sub-key schemas
+ if isinstance(schema, dict):
+ for key, sub_schema in schema.items():
+ if sub_schema is not None and key in config_data:
+ value = config_data[key]
+ if isinstance(value, dict):
+ child_path = f"{path}.{key}" if path else key
+ if isinstance(sub_schema, dict):
+ validate_unknown_keys(value, sub_schema, child_path)
+ elif isinstance(sub_schema, set):
+ validate_unknown_keys(
+ value, {k: None for k in sub_schema}, child_path
+ )
+
+
+def validate_path_security(path: str, base_dir: str, project_dir: str = None) -> str:
+ """
+ Validate that a path is secure and contained within the base directory.
+
+ Args:
+ path: The path to validate
+ base_dir: The base directory that should contain the path
+ project_dir: The project directory for final security check (if different from base_dir)
+
+ Returns:
+ The normalized absolute path if valid
+
+ Raises:
+ ConfigSecurityError: If the path is not secure
+ """
+ # Check for encoded traversal patterns before normalization
+ if '....' in path or '..%2F' in path or '..%5C' in path:
+ raise ConfigSecurityError(f"Encoded path traversal detected in '{path}'. Encoded traversal patterns are not allowed for security reasons.")
+
+ # Normalize the path
+ normalized_path = os.path.normpath(path)
+
+ # Check for malicious path traversal attempts
+ # Allow legitimate relative paths like "../data/file.json" but block excessive traversal
+ path_parts = normalized_path.split(os.sep)
+ if path_parts.count('..') > 2: # Allow up to 2 levels of ".." for legitimate relative paths
+ raise ConfigSecurityError(f"Excessive path traversal detected in '{path}'. Too many '..' components for security reasons.")
+
+ # Check for absolute paths that might escape the project directory
+ if os.path.isabs(normalized_path):
+ # Only allow absolute paths that are within the base directory
+ try:
+ real_path = os.path.realpath(normalized_path)
+ real_base = os.path.realpath(base_dir)
+ if not (real_path == real_base or real_path.startswith(real_base + os.sep)):
+ raise ConfigSecurityError(f"Path '{path}' resolves to '{real_path}' which is outside the project directory '{real_base}'")
+ except (OSError, ValueError) as e:
+ raise ConfigSecurityError(f"Invalid path '{path}': {str(e)}")
+
+ # Resolve relative paths against base directory
+ if not os.path.isabs(normalized_path):
+ resolved_path = os.path.join(base_dir, normalized_path)
+ normalized_path = os.path.normpath(resolved_path)
+
+ # Final security check - ensure the resolved path is within the project directory
+ try:
+ real_path = os.path.realpath(normalized_path)
+ # Use project_dir for final check if provided, otherwise use base_dir
+ check_dir = project_dir if project_dir else base_dir
+ real_check_dir = os.path.realpath(check_dir)
+ if not (real_path == real_check_dir or real_path.startswith(real_check_dir + os.sep)):
+ raise ConfigSecurityError(f"Path '{path}' resolves to '{real_path}' which is outside the project directory '{real_check_dir}'")
+ except (OSError, ValueError) as e:
+ raise ConfigSecurityError(f"Invalid path '{path}': {str(e)}")
+
+ return normalized_path
+
+
+# Optional field type specifications for validation.
+# Maps config key -> (expected_type, human description, allow_negative).
+# Only fields that are commonly misconfigured and cause silent failures.
+_OPTIONAL_INT_FIELDS = {
+ "alert_time_each_instance": ("seconds to alert per instance", False),
+ "max_annotations_per_item": ("max annotations per item", True), # -1 = unlimited
+ "max_annotations_per_user": ("max annotations per user", True),
+ "min_annotators_per_instance": ("minimum annotators per instance", False),
+ "random_seed": ("random seed", True),
+ "max_session_seconds": ("max session duration in seconds", False),
+}
+# num_annotators_per_item validated separately โ it may be int OR dict.
+
+_OPTIONAL_BOOL_FIELDS = {
+ "highlight_linebreaks": "whether to highlight linebreaks",
+ "jumping_to_id_disabled": "whether jumping to ID is disabled",
+ "require_fully_annotated": "whether full annotation is required",
+ "require_password": "whether password is required",
+ "require_no_password": "whether no-password mode is enabled",
+ "customjs": "whether custom JS is enabled",
+ "watch_data_directory": "whether to watch data directory for changes",
+ "persist_sessions": "whether to persist sessions across restarts",
+}
+
+_VALID_ASSIGNMENT_STRATEGIES = [
+ "random", "fixed_order", "active_learning", "llm_confidence",
+ "max_diversity", "least_annotated", "category_based", "diversity_clustering",
+ "batch", "priority",
+]
+
+
+def validate_num_annotators_per_item(value: Any) -> None:
+ """
+ Validate the shape of ``num_annotators_per_item``.
+
+ Accepts either an int (legacy form) or a dict with optional keys
+ ``default``, ``overlap_sample``, ``adaptive``, and ``min``.
+ """
+ if value is None:
+ return
+ if isinstance(value, bool):
+ raise ConfigValidationError(
+ "'num_annotators_per_item' must be an integer or a structured mapping, "
+ f"got bool: {value!r}"
+ )
+ if isinstance(value, int):
+ if value < 0:
+ raise ConfigValidationError(
+ "'num_annotators_per_item' as integer must be non-negative; "
+ "use 0 or omit the key for unlimited (legacy used -1)."
+ )
+ return
+ if not isinstance(value, dict):
+ raise ConfigValidationError(
+ "'num_annotators_per_item' must be an integer or a mapping, "
+ f"got {type(value).__name__}: {value!r}"
+ )
+
+ allowed = {"default", "overlap_sample", "adaptive", "min"}
+ unknown = set(value) - allowed
+ if unknown:
+ raise ConfigValidationError(
+ f"Unknown keys in num_annotators_per_item: {sorted(unknown)}. "
+ f"Allowed: {sorted(allowed)}"
+ )
+
+ default = value.get("default", 1)
+ if not isinstance(default, int) or isinstance(default, bool) or default < 1:
+ raise ConfigValidationError(
+ f"num_annotators_per_item.default must be a positive integer, got {default!r}"
+ )
+
+ minimum = value.get("min")
+ if minimum is not None:
+ if not isinstance(minimum, int) or isinstance(minimum, bool) or minimum < 1:
+ raise ConfigValidationError(
+ f"num_annotators_per_item.min must be a positive integer, got {minimum!r}"
+ )
+ if minimum > default:
+ raise ConfigValidationError(
+ "num_annotators_per_item.min cannot exceed num_annotators_per_item.default"
+ )
+
+ overlap = value.get("overlap_sample")
+ if overlap is not None:
+ if not isinstance(overlap, dict):
+ raise ConfigValidationError(
+ "num_annotators_per_item.overlap_sample must be a mapping"
+ )
+ unknown = set(overlap) - {"fraction", "count", "stratify_by", "seed"}
+ if unknown:
+ raise ConfigValidationError(
+ f"Unknown keys in overlap_sample: {sorted(unknown)}"
+ )
+ frac = overlap.get("fraction")
+ if not isinstance(frac, (int, float)) or isinstance(frac, bool) or not (0 < frac <= 1):
+ raise ConfigValidationError(
+ f"overlap_sample.fraction must be in (0, 1], got {frac!r}"
+ )
+ count = overlap.get("count")
+ if not isinstance(count, int) or isinstance(count, bool) or count < 2:
+ raise ConfigValidationError(
+ f"overlap_sample.count must be an integer >= 2, got {count!r}"
+ )
+ if count <= default:
+ raise ConfigValidationError(
+ "overlap_sample.count must be greater than num_annotators_per_item.default "
+ f"({count} <= {default})"
+ )
+ stratify_by = overlap.get("stratify_by")
+ if stratify_by is not None and not isinstance(stratify_by, str):
+ raise ConfigValidationError(
+ f"overlap_sample.stratify_by must be a string or omitted, got {stratify_by!r}"
+ )
+ seed = overlap.get("seed")
+ if seed is not None and (not isinstance(seed, int) or isinstance(seed, bool)):
+ raise ConfigValidationError(
+ f"overlap_sample.seed must be an integer, got {seed!r}"
+ )
+
+ adaptive = value.get("adaptive")
+ if adaptive is not None:
+ if not isinstance(adaptive, dict):
+ raise ConfigValidationError(
+ "num_annotators_per_item.adaptive must be a mapping"
+ )
+ unknown = set(adaptive) - {"enabled", "disagreement_threshold", "boost_to"}
+ if unknown:
+ raise ConfigValidationError(
+ f"Unknown keys in adaptive: {sorted(unknown)}"
+ )
+ if "enabled" in adaptive and not isinstance(adaptive["enabled"], bool):
+ raise ConfigValidationError(
+ f"adaptive.enabled must be a boolean, got {adaptive['enabled']!r}"
+ )
+ thr = adaptive.get("disagreement_threshold")
+ if thr is not None and (not isinstance(thr, (int, float)) or isinstance(thr, bool) or not (0 <= thr <= 1)):
+ raise ConfigValidationError(
+ f"adaptive.disagreement_threshold must be in [0, 1], got {thr!r}"
+ )
+ boost = adaptive.get("boost_to")
+ if boost is not None:
+ if not isinstance(boost, int) or isinstance(boost, bool) or boost < 2:
+ raise ConfigValidationError(
+ f"adaptive.boost_to must be an integer >= 2, got {boost!r}"
+ )
+ if boost <= default:
+ raise ConfigValidationError(
+ f"adaptive.boost_to must exceed default ({boost} <= {default})"
+ )
+
+
+def validate_per_annotator_quota(value: Any) -> None:
+ """Validate the shape of ``per_annotator_quota``."""
+ if value is None:
+ return
+ if not isinstance(value, dict):
+ raise ConfigValidationError(
+ "'per_annotator_quota' must be a mapping, "
+ f"got {type(value).__name__}: {value!r}"
+ )
+ allowed = {"default", "by_user", "by_user_role"}
+ unknown = set(value) - allowed
+ if unknown:
+ raise ConfigValidationError(
+ f"Unknown keys in per_annotator_quota: {sorted(unknown)}. Allowed: {sorted(allowed)}"
+ )
+ default = value.get("default")
+ if default is not None and (not isinstance(default, int) or isinstance(default, bool) or default < 0):
+ raise ConfigValidationError(
+ f"per_annotator_quota.default must be a non-negative integer, got {default!r}"
+ )
+ for key in ("by_user", "by_user_role"):
+ mapping = value.get(key)
+ if mapping is None:
+ continue
+ if not isinstance(mapping, dict):
+ raise ConfigValidationError(
+ f"per_annotator_quota.{key} must be a mapping of name -> integer"
+ )
+ for k, v in mapping.items():
+ if not isinstance(k, str) or not k:
+ raise ConfigValidationError(
+ f"per_annotator_quota.{key} keys must be non-empty strings, got {k!r}"
+ )
+ if not isinstance(v, int) or isinstance(v, bool) or v < 0:
+ raise ConfigValidationError(
+ f"per_annotator_quota.{key}[{k!r}] must be a non-negative integer, got {v!r}"
+ )
+
+
+def resolve_num_annotators_per_item(config_data: Dict[str, Any]) -> int:
+ """
+ Resolve the *default* cap (used as ``ItemStateManager.max_annotations_per_item``).
+
+ Resolution order:
+ 1. num_annotators_per_item (int form) โ that value
+ 2. num_annotators_per_item.default โ that value
+ 3. max_annotations_per_item (legacy) โ that value
+ 4. otherwise โ -1 (unlimited)
+ """
+ val = config_data.get("num_annotators_per_item")
+ if isinstance(val, int) and not isinstance(val, bool):
+ return val
+ if isinstance(val, dict) and val.get("default") is not None:
+ return int(val["default"])
+ legacy = config_data.get("max_annotations_per_item")
+ if isinstance(legacy, int) and not isinstance(legacy, bool):
+ return legacy
+ return -1
+
+
+def validate_optional_field_types(config_data: Dict[str, Any]) -> None:
+ """
+ Validate types for commonly misconfigured optional fields.
+
+ Catches issues like string values for integer fields (e.g., alert_time_each_instance: "30")
+ or wrong types for booleans, which would silently produce incorrect behavior at runtime.
+
+ Args:
+ config_data: The parsed configuration dictionary
+
+ Raises:
+ ConfigValidationError: If a field has the wrong type
+ """
+ # Validate integer fields
+ for field, (desc, allow_negative) in _OPTIONAL_INT_FIELDS.items():
+ if field in config_data:
+ val = config_data[field]
+ if not isinstance(val, int) or isinstance(val, bool):
+ raise ConfigValidationError(
+ f"'{field}' must be an integer ({desc}), "
+ f"got {type(val).__name__}: {val!r}"
+ )
+ if not allow_negative and val < 0:
+ raise ConfigValidationError(
+ f"'{field}' must be a non-negative integer ({desc}), got {val}"
+ )
+
+ # Validate boolean fields (None/null is allowed as "not set")
+ for field, desc in _OPTIONAL_BOOL_FIELDS.items():
+ if field in config_data:
+ val = config_data[field]
+ if val is not None and not isinstance(val, bool):
+ raise ConfigValidationError(
+ f"'{field}' must be a boolean ({desc}), "
+ f"got {type(val).__name__}: {val!r}"
+ )
+
+ # Validate num_annotators_per_item (int OR structured dict)
+ if 'num_annotators_per_item' in config_data:
+ validate_num_annotators_per_item(config_data['num_annotators_per_item'])
+
+ # Validate per_annotator_quota structured dict
+ if 'per_annotator_quota' in config_data:
+ validate_per_annotator_quota(config_data['per_annotator_quota'])
+
+ # Emit a deprecation warning if max_annotations_per_item is set alongside
+ # num_annotators_per_item; reject silent inconsistencies (both set to
+ # conflicting values).
+ if 'max_annotations_per_item' in config_data and 'num_annotators_per_item' in config_data:
+ legacy = config_data['max_annotations_per_item']
+ canonical = config_data['num_annotators_per_item']
+ canonical_int = canonical if isinstance(canonical, int) else canonical.get('default')
+ if canonical_int is not None and legacy != canonical_int:
+ raise ConfigValidationError(
+ "'max_annotations_per_item' and 'num_annotators_per_item' are both "
+ f"set with conflicting values ({legacy} vs {canonical_int}). "
+ "Drop 'max_annotations_per_item' โ 'num_annotators_per_item' is the canonical key."
+ )
+ import warnings as _w
+ _w.warn(
+ "'max_annotations_per_item' is deprecated; use 'num_annotators_per_item' "
+ "instead. Setting both is redundant.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+
+ # Validate assignment_strategy enum
+ if 'assignment_strategy' in config_data:
+ strat = config_data['assignment_strategy']
+ # Can be a string or a dict with a 'name' key
+ strat_name = strat
+ if isinstance(strat, dict):
+ strat_name = strat.get('name', '')
+ if isinstance(strat_name, str) and strat_name.lower() not in _VALID_ASSIGNMENT_STRATEGIES:
+ raise ConfigValidationError(
+ f"'assignment_strategy' value '{strat_name}' is not recognized. "
+ f"Valid strategies: {', '.join(_VALID_ASSIGNMENT_STRATEGIES)}"
+ )
+
+
+def validate_judge_calibration_config(config_data: Dict[str, Any]) -> None:
+ """Validate the ``judge_calibration`` block when enabled.
+
+ Delegates to the typed config's ``validate()`` (so the rules live in one
+ place) and additionally cross-checks that referenced schema names exist in
+ ``annotation_schemes``. Raises ConfigValidationError on hard errors.
+ """
+ jc = config_data.get("judge_calibration")
+ if not isinstance(jc, dict) or not jc.get("enabled"):
+ return
+
+ from potato.judge_calibration.config import parse_judge_calibration_config
+
+ cfg = parse_judge_calibration_config(config_data)
+ errors = cfg.validate()
+
+ # Cross-check schema references against declared annotation_schemes.
+ declared = {
+ s.get("name")
+ for s in (config_data.get("annotation_schemes") or [])
+ if isinstance(s, dict)
+ }
+ for name in cfg.schemas:
+ if name not in declared:
+ errors.append(
+ f"judge_calibration.schemas references unknown scheme '{name}' "
+ f"(declared: {sorted(n for n in declared if n)})"
+ )
+
+ if errors:
+ raise ConfigValidationError(
+ "Invalid judge_calibration configuration:\n - " + "\n - ".join(errors)
+ )
+
+
+def validate_yaml_structure(config_data: Dict[str, Any], project_dir: str = None, config_file_dir: str = None) -> None:
+ """
+ Validate the structure and content of the YAML configuration.
+
+ Args:
+ config_data: The parsed YAML configuration
+ project_dir: The project directory
+ config_file_dir: The directory containing the config file
+
+ Raises:
+ ConfigValidationError: If the configuration is invalid
+ """
+ if not isinstance(config_data, dict):
+ raise ConfigValidationError("Configuration must be a YAML object (dictionary)")
+
+ # Required fields validation. NOTE: 'data_files' is intentionally NOT here โ
+ # it is one of three mutually-acceptable data sources (data_files /
+ # data_directory / data_sources), enforced by the dedicated check below.
+ # Listing it here unconditionally made data_directory- and data_sources-only
+ # configs fail validation before that smarter check could run (F-038).
+ required_fields = [
+ 'item_properties',
+ 'task_dir',
+ 'output_annotation_dir',
+ 'annotation_task_name',
+ ]
+
+ missing_fields = [field for field in required_fields if field not in config_data]
+ if missing_fields:
+ raise ConfigValidationError(f"Missing required configuration fields: {', '.join(missing_fields)}")
+
+ # Validate item_properties
+ item_properties = config_data.get('item_properties', {})
+ if not isinstance(item_properties, dict):
+ raise ConfigValidationError("item_properties must be a dictionary")
+
+ required_item_props = ['id_key', 'text_key']
+ missing_item_props = [prop for prop in required_item_props if prop not in item_properties]
+ if missing_item_props:
+ raise ConfigValidationError(f"Missing required item_properties: {', '.join(missing_item_props)}")
+
+ # Validate optional category_key (for category-based assignment)
+ if 'category_key' in item_properties:
+ category_key = item_properties['category_key']
+ if not isinstance(category_key, str) or not category_key.strip():
+ raise ConfigValidationError("item_properties.category_key must be a non-empty string")
+
+ # Validate data_files (required unless data_directory or data_sources is provided)
+ data_files = config_data.get('data_files', [])
+ data_directory = config_data.get('data_directory')
+ data_sources = config_data.get('data_sources')
+
+ if not isinstance(data_files, list):
+ raise ConfigValidationError("data_files must be a list")
+
+ # data_files can be empty if data_directory or data_sources is configured
+ if not data_files and not data_directory and not data_sources:
+ raise ConfigValidationError(
+ "At least one data source must be configured: "
+ "'data_files', 'data_directory', or 'data_sources'"
+ )
+
+ # Validate data_sources configuration if present
+ if data_sources:
+ validate_data_sources_config(config_data)
+
+ # Validate server config if present
+ validate_server_config(config_data)
+
+ # Validate authentication config if present
+ validate_authentication_config(config_data)
+
+ # Validate data_directory config if present
+ validate_data_directory_config(config_data)
+
+ # Validate annotation schemes
+ validate_annotation_schemes(config_data)
+
+ # Validate training configuration if present
+ validate_training_config(config_data, project_dir, config_file_dir)
+
+ # Validate database configuration if present
+ if 'database' in config_data:
+ validate_database_config(config_data['database'])
+
+ # Validate active learning configuration if present
+ validate_active_learning_config(config_data)
+
+ # Validate AI support configuration if present
+ validate_ai_support_config(config_data)
+
+ # Validate chat support configuration if present
+ validate_chat_support_config(config_data)
+
+ # Validate category assignment configuration if present
+ validate_category_assignment_config(config_data)
+
+ # Validate batch assignment configuration if present
+ validate_batch_assignment_config(config_data)
+
+ # Validate diversity ordering configuration if present
+ validate_diversity_config(config_data)
+
+ # Validate embedding visualization configuration if present
+ validate_embedding_visualization_config(config_data)
+
+ # Validate adjudication configuration if present
+ if 'adjudication' in config_data:
+ validate_adjudication_config(config_data)
+
+ # Validate quality control configuration if present
+ validate_quality_control_config(config_data)
+
+ # Validate assignment reclaim configuration if present
+ validate_instance_reclaim_config(config_data)
+
+ # Validate instance display configuration if present
+ validate_instance_display_config(config_data)
+
+ # Validate format_handling configuration if present
+ validate_format_handling_config(config_data)
+
+ # Validate layout configuration if present
+ validate_layout_config(config_data)
+
+ # Validate BWS configuration if present
+ if 'bws_config' in config_data:
+ _validate_bws_config(config_data)
+
+ # Validate IBWS configuration if present
+ if 'ibws_config' in config_data:
+ _validate_ibws_config(config_data)
+
+ # Validate MACE configuration if present
+ if 'mace' in config_data:
+ _validate_mace_config(config_data)
+
+ # Validate types for commonly misconfigured optional fields
+ validate_optional_field_types(config_data)
+
+ # Fail loud if annotator search-and-claim is combined with an
+ # assignment design it would corrupt via self-selection.
+ validate_search_assignment_compat(config_data)
+
+ # Validate codebook_mode (and apply the crowd force-lock).
+ validate_codebook_config(config_data)
+
+ # Validate judge_calibration configuration if present
+ validate_judge_calibration_config(config_data)
+
+ # Warn about unrecognized keys at all nesting levels
+ validate_unknown_keys(config_data)
+
+
+# Assignment strategies whose sampling/ordering self-selection breaks.
+_CLAIM_INCOMPATIBLE_STRATEGIES = {
+ "random", "diversity_clustering", "max_diversity",
+ "active_learning", "llm_confidence", "least_annotated",
+ "category_based", "batch",
+}
+
+
+def validate_search_assignment_compat(config_data: Dict[str, Any]) -> None:
+ """Hard-fail when ``search.annotator_claim`` is combined with a
+ feature whose integrity depends on the platform โ not the annotator โ
+ choosing the next item. Read-only admin search is unaffected.
+
+ Solo/QDA mode (single coder over the whole corpus) is always allowed.
+ """
+ search = config_data.get("search")
+ if not isinstance(search, dict) or not search.get("annotator_claim"):
+ return
+
+ # Single-coder modes have no sampling/overlap invariant to protect.
+ if (config_data.get("qda_mode") or {}).get("enabled") or \
+ (config_data.get("solo_mode") or {}).get("enabled"):
+ return
+
+ conflicts = []
+
+ strat = config_data.get("assignment_strategy")
+ if isinstance(strat, dict):
+ strat = strat.get("name")
+ if strat and str(strat).lower() in _CLAIM_INCOMPATIBLE_STRATEGIES:
+ conflicts.append(
+ f"assignment_strategy: {strat} (self-selection breaks "
+ f"sampling/ordering)")
+
+ for k in ("max_annotations_per_item", "num_annotators_per_item",
+ "min_annotators_per_instance"):
+ raw = config_data.get(k, -1)
+ # num_annotators_per_item may now be a dict โ extract default + overlap_sample.count
+ candidates = []
+ if isinstance(raw, dict):
+ if raw.get("default") is not None:
+ candidates.append(raw["default"])
+ overlap = raw.get("overlap_sample") or {}
+ if overlap.get("count") is not None:
+ candidates.append(overlap["count"])
+ else:
+ candidates.append(raw)
+ for cand in candidates:
+ try:
+ if int(cand) > 1:
+ conflicts.append(
+ f"{k}: {config_data[k]} (inter-annotator overlap "
+ f"cannot be guaranteed under self-selection)")
+ break
+ except (TypeError, ValueError):
+ continue
+
+ if (config_data.get("attention_checks") or {}).get("enabled"):
+ conflicts.append("attention_checks.enabled (annotators could "
+ "locate/avoid QC items)")
+ if (config_data.get("gold_standards") or {}).get("enabled"):
+ conflicts.append("gold_standards.enabled (annotators could "
+ "locate/avoid gold items)")
+ if (config_data.get("icl_labeling") or {}).get("enabled"):
+ conflicts.append("icl_labeling.enabled (blind LLM-verification "
+ "tasks must not be findable)")
+ if (config_data.get("adjudication") or {}).get("enabled"):
+ conflicts.append("adjudication.enabled (the adjudication queue "
+ "is curated)")
+
+ login_type = (config_data.get("login") or {}).get("type")
+ crowd = (
+ "mturk" in config_data or "prolific" in config_data
+ or login_type in ("mturk", "prolific")
+ )
+ if crowd:
+ conflicts.append("crowdsourcing backend (HIT = the assigned "
+ "unit; self-selection breaks payment/coverage)")
+
+ if conflicts:
+ raise ConfigValidationError(
+ "search.annotator_claim: true is incompatible with this "
+ "configuration:\n - " + "\n - ".join(conflicts) +
+ "\n\nAnnotator search-and-claim is only supported with "
+ "solo_mode/qda_mode, or fixed_order assignment without "
+ "overlap, quality-control injection, ICL verification, "
+ "adjudication, or a crowdsourcing backend. Use read-only "
+ "admin search (no annotator_claim) for those designs."
+ )
+
+
+_CODEBOOK_MODES = ("fixed", "extensible", "open")
+
+
+def _crowd_backend(config_data: Dict[str, Any]) -> bool:
+ login_type = (config_data.get("login") or {}).get("type")
+ return (
+ "mturk" in config_data or "prolific" in config_data
+ or login_type in ("mturk", "prolific")
+ )
+
+
+def get_codebook_mode(config_data: Dict[str, Any]) -> str:
+ """Resolve the effective codebook mode.
+
+ Precedence: explicit ``codebook_mode`` / ``codebook.mode`` if set;
+ else ``open`` when solo/QDA mode is enabled; else ``fixed``. A crowd
+ backend force-locks ``fixed`` regardless of the request (annotators
+ on a paid HIT must not reshape the shared codebook).
+ """
+ raw = config_data.get("codebook_mode")
+ if raw is None:
+ raw = (config_data.get("codebook") or {}).get("mode")
+
+ if raw is None:
+ single = (
+ (config_data.get("qda_mode") or {}).get("enabled")
+ or (config_data.get("solo_mode") or {}).get("enabled")
+ )
+ mode = "open" if single else "fixed"
+ else:
+ mode = str(raw).strip().lower()
+
+ if _crowd_backend(config_data):
+ return "fixed"
+ return mode
+
+
+def validate_codebook_config(config_data: Dict[str, Any]) -> None:
+ """Reject an invalid ``codebook_mode`` value, and warn when a crowd
+ backend overrides a requested non-fixed mode."""
+ raw = config_data.get("codebook_mode")
+ if raw is None:
+ raw = (config_data.get("codebook") or {}).get("mode")
+ if raw is None:
+ return
+
+ mode = str(raw).strip().lower()
+ if mode not in _CODEBOOK_MODES:
+ raise ConfigValidationError(
+ f"codebook_mode must be one of {', '.join(_CODEBOOK_MODES)}; "
+ f"got {raw!r}."
+ )
+ if mode != "fixed" and _crowd_backend(config_data):
+ logging.warning(
+ "codebook_mode=%s requested with a crowdsourcing backend; "
+ "force-locking to 'fixed' (paid annotators must not reshape "
+ "the shared codebook).", mode)
+
+
+def validate_annotation_schemes(config_data: Dict[str, Any]) -> None:
+ """
+ Validate annotation schemes configuration.
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If annotation schemes are invalid
+ """
+ has_top_level = 'annotation_schemes' in config_data
+ has_phases = 'phases' in config_data and config_data['phases']
+
+ # Check for conflicting annotation_schemes locations
+ if has_top_level and has_phases:
+ # Check if any phase also has annotation_schemes
+ phases = config_data['phases']
+ phases_with_schemes = []
+ if isinstance(phases, list):
+ phases_with_schemes = [
+ phase.get('name', f'phase[{i}]')
+ for i, phase in enumerate(phases)
+ if 'annotation_schemes' in phase
+ ]
+ elif isinstance(phases, dict):
+ phases_with_schemes = [
+ name for name, phase in phases.items()
+ if name != 'order' and isinstance(phase, dict) and 'annotation_schemes' in phase
+ ]
+
+ if phases_with_schemes:
+ raise ConfigValidationError(
+ f"Configuration has both top-level 'annotation_schemes' and phase-level "
+ f"'annotation_schemes' in: {', '.join(phases_with_schemes)}. "
+ f"Use only one location to avoid confusion."
+ )
+
+ # Check for annotation schemes in different formats
+ if has_top_level:
+ schemes = config_data['annotation_schemes']
+ if not isinstance(schemes, list):
+ raise ConfigValidationError("annotation_schemes must be a list")
+ if not schemes:
+ raise ConfigValidationError("annotation_schemes cannot be empty")
+
+ for i, scheme in enumerate(schemes):
+ validate_single_annotation_scheme(scheme, f"annotation_schemes[{i}]")
+
+ elif 'phases' in config_data and config_data['phases']:
+ phases = config_data['phases']
+ if isinstance(phases, list):
+ for i, phase in enumerate(phases):
+ phase_id = phase.get('name', f'phase[{i}]')
+ # Phases can have annotation_schemes, file, type, instrument, or instruments
+ if 'annotation_schemes' in phase:
+ schemes = phase['annotation_schemes']
+ if not isinstance(schemes, list):
+ raise ConfigValidationError(f"Phase {phase_id} annotation_schemes must be a list")
+ if not schemes:
+ raise ConfigValidationError(f"Phase {phase_id} annotation_schemes cannot be empty")
+
+ for j, scheme in enumerate(schemes):
+ validate_single_annotation_scheme(scheme, f"phases[{i}].annotation_schemes[{j}]")
+ elif 'file' in phase or 'type' in phase or 'instrument' in phase or 'instruments' in phase:
+ # Legacy format or instrument-based - validated at runtime
+ _validate_phase_instruments(phase, phase_id)
+ else:
+ raise ConfigValidationError(
+ f"Phase {phase_id} requires 'annotation_schemes', 'file', 'type', "
+ f"'instrument', or 'instruments'"
+ )
+ else:
+ # Dictionary format
+ for phase_name, phase in phases.items():
+ if phase_name == 'order':
+ continue
+ # Phases can have annotation_schemes, file, type, instrument, or instruments
+ if 'annotation_schemes' in phase:
+ schemes = phase['annotation_schemes']
+ if not isinstance(schemes, list):
+ raise ConfigValidationError(f"Phase {phase_name} annotation_schemes must be a list")
+ if not schemes:
+ raise ConfigValidationError(f"Phase {phase_name} annotation_schemes cannot be empty")
+
+ for j, scheme in enumerate(schemes):
+ validate_single_annotation_scheme(scheme, f"phases.{phase_name}.annotation_schemes[{j}]")
+ elif 'file' in phase or 'type' in phase or 'instrument' in phase or 'instruments' in phase:
+ # Legacy format or instrument-based - validated at runtime
+ _validate_phase_instruments(phase, phase_name)
+ else:
+ raise ConfigValidationError(
+ f"Phase {phase_name} requires 'annotation_schemes', 'file', 'type', "
+ f"'instrument', or 'instruments'"
+ )
+ else:
+ raise ConfigValidationError("Config must have either 'annotation_schemes' (top-level) or 'phases' with annotation_schemes")
+
+ # Validate keyword_highlight is not enabled for image-based tasks
+ _validate_keyword_highlight_for_images(config_data)
+
+ # Validate display_logic cross-references (schema references and circular dependencies)
+ all_schemes = _collect_all_annotation_schemes(config_data)
+ if all_schemes:
+ validate_display_logic_references(all_schemes)
+
+
+def _collect_all_annotation_schemes(config_data: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """
+ Collect all annotation schemes from config, whether top-level or in phases.
+
+ Args:
+ config_data: The configuration data
+
+ Returns:
+ List of all annotation scheme dictionaries
+ """
+ schemes = []
+
+ if 'annotation_schemes' in config_data:
+ schemes.extend(config_data['annotation_schemes'])
+ elif 'phases' in config_data:
+ phases = config_data['phases']
+ if isinstance(phases, list):
+ for phase in phases:
+ if 'annotation_schemes' in phase:
+ schemes.extend(phase['annotation_schemes'])
+ elif isinstance(phases, dict):
+ for phase_name, phase in phases.items():
+ if phase_name != 'order' and isinstance(phase, dict):
+ if 'annotation_schemes' in phase:
+ schemes.extend(phase['annotation_schemes'])
+
+ return schemes
+
+
+def _validate_keyword_highlight_for_images(config_data: Dict[str, Any]) -> None:
+ """
+ Validate that keyword_highlight is not enabled for image-based tasks.
+
+ Keyword highlighting highlights text in the instance content, which doesn't
+ make sense for images. This validation catches configuration errors early.
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If keyword_highlight is enabled for an image task
+ """
+ # Check if the text_key suggests this is an image-based task
+ text_key = config_data.get('item_properties', {}).get('text_key', 'text')
+ image_indicators = ['image', 'img', 'photo', 'picture', 'url']
+ is_likely_image_task = any(indicator in text_key.lower() for indicator in image_indicators)
+
+ if not is_likely_image_task:
+ return # Not an image task, no need to check
+
+ # Get all annotation schemes
+ schemes = []
+ if 'annotation_schemes' in config_data:
+ schemes = config_data['annotation_schemes']
+ elif 'phases' in config_data:
+ phases = config_data['phases']
+ if isinstance(phases, list):
+ for phase in phases:
+ schemes.extend(phase.get('annotation_schemes', []))
+ elif isinstance(phases, dict):
+ for phase_name, phase in phases.items():
+ if phase_name != 'order' and isinstance(phase, dict):
+ schemes.extend(phase.get('annotation_schemes', []))
+
+ # Check each scheme for keyword_highlight
+ for i, scheme in enumerate(schemes):
+ if not isinstance(scheme, dict):
+ continue
+ ai_support = scheme.get('ai_support', {})
+ if not isinstance(ai_support, dict):
+ continue
+ features = ai_support.get('features', {})
+ if not isinstance(features, dict):
+ continue
+
+ keyword_highlight = features.get('keyword_highlight', False)
+ if keyword_highlight:
+ scheme_name = scheme.get('name', f'scheme[{i}]')
+ raise ConfigValidationError(
+ f"annotation_schemes.{scheme_name}.ai_support.features.keyword_highlight is enabled, "
+ f"but item_properties.text_key='{text_key}' suggests this is an image-based task. "
+ f"Keyword highlighting only works with text content, not images. "
+ f"Set keyword_highlight: false or remove it from the ai_support features."
+ )
+
+
+def _validate_bws_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate Best-Worst Scaling configuration.
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If the BWS config is invalid
+ """
+ bws = config_data['bws_config']
+ if not isinstance(bws, dict):
+ raise ConfigValidationError("bws_config must be a dictionary")
+
+ if 'tuple_size' in bws:
+ if not isinstance(bws['tuple_size'], int) or bws['tuple_size'] < 2:
+ raise ConfigValidationError("bws_config.tuple_size must be an integer >= 2")
+
+ if 'seed' in bws:
+ if not isinstance(bws['seed'], int):
+ raise ConfigValidationError("bws_config.seed must be an integer")
+
+ if 'num_tuples' in bws and bws['num_tuples'] is not None:
+ if not isinstance(bws['num_tuples'], int) or bws['num_tuples'] < 1:
+ raise ConfigValidationError("bws_config.num_tuples must be a positive integer or null")
+
+ if 'min_item_appearances' in bws and bws['min_item_appearances'] is not None:
+ if not isinstance(bws['min_item_appearances'], int) or bws['min_item_appearances'] < 1:
+ raise ConfigValidationError("bws_config.min_item_appearances must be a positive integer or null")
+
+ # Validate scoring config if present
+ scoring = bws.get('scoring', {})
+ if scoring:
+ if not isinstance(scoring, dict):
+ raise ConfigValidationError("bws_config.scoring must be a dictionary")
+ valid_methods = ['counting', 'bradley_terry', 'plackett_luce']
+ method = scoring.get('method', 'counting')
+ if method not in valid_methods:
+ raise ConfigValidationError(f"bws_config.scoring.method must be one of: {valid_methods}")
+
+
+def _validate_ibws_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate Iterative Best-Worst Scaling configuration.
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If the IBWS config is invalid
+ """
+ # Mutual exclusivity with bws_config
+ if 'bws_config' in config_data:
+ raise ConfigValidationError(
+ "ibws_config and bws_config are mutually exclusive. "
+ "Use ibws_config for iterative BWS or bws_config for standard BWS."
+ )
+
+ ibws = config_data['ibws_config']
+ if not isinstance(ibws, dict):
+ raise ConfigValidationError("ibws_config must be a dictionary")
+
+ # Require at least one BWS annotation scheme
+ schemes = config_data.get('annotation_schemes', [])
+ has_bws_scheme = any(s.get('annotation_type') == 'bws' for s in schemes)
+ if not has_bws_scheme:
+ raise ConfigValidationError(
+ "ibws_config requires at least one annotation scheme with annotation_type: bws"
+ )
+
+ # tuple_size
+ if 'tuple_size' in ibws:
+ if not isinstance(ibws['tuple_size'], int) or ibws['tuple_size'] < 2:
+ raise ConfigValidationError("ibws_config.tuple_size must be an integer >= 2")
+
+ # max_rounds
+ if 'max_rounds' in ibws and ibws['max_rounds'] is not None:
+ if not isinstance(ibws['max_rounds'], int) or ibws['max_rounds'] < 1:
+ raise ConfigValidationError("ibws_config.max_rounds must be a positive integer or null")
+
+ # seed
+ if 'seed' in ibws:
+ if not isinstance(ibws['seed'], int):
+ raise ConfigValidationError("ibws_config.seed must be an integer")
+
+ # scoring_method
+ valid_methods = ['counting', 'bradley_terry', 'plackett_luce']
+ if 'scoring_method' in ibws:
+ if ibws['scoring_method'] not in valid_methods:
+ raise ConfigValidationError(
+ f"ibws_config.scoring_method must be one of: {valid_methods}"
+ )
+
+ # tuples_per_item_per_round
+ if 'tuples_per_item_per_round' in ibws:
+ val = ibws['tuples_per_item_per_round']
+ if not isinstance(val, int) or val < 1:
+ raise ConfigValidationError(
+ "ibws_config.tuples_per_item_per_round must be a positive integer"
+ )
+
+
+def _validate_mace_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate MACE competence estimation configuration.
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If the MACE config is invalid
+ """
+ mace = config_data.get('mace', {})
+ if not isinstance(mace, dict):
+ raise ConfigValidationError("mace must be a dictionary")
+
+ if not mace.get('enabled', False):
+ return # Not enabled, skip validation
+
+ # Validate numeric parameters
+ min_annots = mace.get('min_annotations_per_item', 3)
+ if not isinstance(min_annots, int) or min_annots < 2:
+ raise ConfigValidationError(
+ "mace.min_annotations_per_item must be an integer >= 2"
+ )
+
+ trigger_n = mace.get('trigger_every_n', 10)
+ if not isinstance(trigger_n, int) or trigger_n < 1:
+ raise ConfigValidationError(
+ "mace.trigger_every_n must be an integer >= 1"
+ )
+
+ num_restarts = mace.get('num_restarts', 10)
+ if not isinstance(num_restarts, int) or num_restarts < 1:
+ raise ConfigValidationError(
+ "mace.num_restarts must be an integer >= 1"
+ )
+
+ # Warn if no categorical schemas are defined
+ categorical_types = {'radio', 'likert', 'select', 'multiselect'}
+ schemes = config_data.get('annotation_schemes', [])
+ has_categorical = any(
+ s.get('annotation_type', '') in categorical_types
+ for s in schemes if isinstance(s, dict)
+ )
+ if not has_categorical:
+ logger.warning(
+ "MACE is enabled but no categorical annotation schemes "
+ "(radio, likert, select, multiselect) are defined. "
+ "MACE will have no data to process."
+ )
+
+
+def _validate_phase_instruments(phase: Dict[str, Any], phase_name: str) -> None:
+ """
+ Validate instrument references in a phase configuration.
+
+ Args:
+ phase: The phase configuration
+ phase_name: Name of the phase for error messages
+
+ Raises:
+ ConfigValidationError: If instrument references are invalid
+ """
+ # Validate single instrument reference
+ if 'instrument' in phase:
+ inst_id = phase['instrument']
+ if not isinstance(inst_id, str):
+ raise ConfigValidationError(
+ f"Phase {phase_name}: 'instrument' must be a string"
+ )
+ try:
+ from potato.survey_instruments import get_registry
+ registry = get_registry()
+ if inst_id not in registry['instruments']:
+ available = sorted(registry['instruments'].keys())[:10]
+ raise ConfigValidationError(
+ f"Phase {phase_name}: Unknown instrument '{inst_id}'. "
+ f"Available instruments: {available}..."
+ )
+ except ImportError:
+ # survey_instruments module not available - skip validation
+ pass
+
+ # Validate multiple instruments
+ if 'instruments' in phase:
+ inst_list = phase['instruments']
+ if not isinstance(inst_list, list):
+ raise ConfigValidationError(
+ f"Phase {phase_name}: 'instruments' must be a list"
+ )
+ try:
+ from potato.survey_instruments import get_registry
+ registry = get_registry()
+ for inst_id in inst_list:
+ if not isinstance(inst_id, str):
+ raise ConfigValidationError(
+ f"Phase {phase_name}: All items in 'instruments' must be strings"
+ )
+ if inst_id not in registry['instruments']:
+ available = sorted(registry['instruments'].keys())[:10]
+ raise ConfigValidationError(
+ f"Phase {phase_name}: Unknown instrument '{inst_id}'. "
+ f"Available instruments: {available}..."
+ )
+ except ImportError:
+ # survey_instruments module not available - skip validation
+ pass
+
+
+def validate_single_annotation_scheme(scheme: Dict[str, Any], path: str) -> None:
+ """
+ Validate a single annotation scheme.
+
+ Args:
+ scheme: The annotation scheme to validate
+ path: The path in the config for error reporting
+
+ Raises:
+ ConfigValidationError: If the scheme is invalid
+ """
+ if not isinstance(scheme, dict):
+ raise ConfigValidationError(f"{path} must be a dictionary")
+
+ required_fields = ['annotation_type', 'name', 'description']
+ missing_fields = [field for field in required_fields if field not in scheme]
+ if missing_fields:
+ raise ConfigValidationError(f"{path} missing required fields: {', '.join(missing_fields)}")
+
+ # Validate annotation_type against the schema registry (single source of truth)
+ from potato.server_utils.schemas.registry import schema_registry
+ valid_types = schema_registry.get_supported_types()
+ if scheme['annotation_type'] not in valid_types:
+ raise ConfigValidationError(f"{path}.annotation_type must be one of: {', '.join(sorted(valid_types))}")
+
+ # Registry-driven required field check: validate fields that are unconditionally
+ # required for this type. Types with alternative forms (e.g., likert accepts either
+ # 'labels' OR 'min_label'+'max_label'+'size') have deeper validation in the
+ # type-specific blocks below. This check catches missing fields for types that
+ # don't have explicit type-specific validation blocks.
+ annotation_type = scheme['annotation_type']
+ _types_with_explicit_validation = {
+ 'radio', 'multiselect', 'select', 'likert', 'slider', 'span', 'multirate',
+ 'image_annotation', 'audio_annotation', 'video_annotation', 'tiered_annotation',
+ 'pairwise', 'bws', 'soft_label', 'confidence', 'constant_sum',
+ 'semantic_differential', 'ranking', 'range_slider', 'hierarchical_multiselect',
+ 'vas', 'rubric_eval', 'error_span', 'card_sort', 'conjoint',
+ }
+ if annotation_type not in _types_with_explicit_validation:
+ schema_def = schema_registry.get(annotation_type)
+ if schema_def and schema_def.required_fields:
+ # 'name' and 'description' are already checked above
+ extra_required = [f for f in schema_def.required_fields
+ if f not in ('name', 'description')]
+ missing = [f for f in extra_required if f not in scheme]
+ if missing:
+ raise ConfigValidationError(
+ f"{path} (type '{annotation_type}') missing required field(s): "
+ f"{', '.join(missing)}"
+ )
+
+ # Type-specific validation (deep structural checks beyond registry required_fields)
+ if annotation_type in ['radio', 'multiselect', 'select']:
+ if 'labels' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'labels' field for {annotation_type} annotation type")
+ if not isinstance(scheme['labels'], list):
+ raise ConfigValidationError(f"{path}.labels must be a list")
+ if not scheme['labels']:
+ raise ConfigValidationError(f"{path}.labels cannot be empty")
+
+ elif annotation_type == 'likert':
+ # Likert can use labels (falls back to radio) or min_label/max_label/size
+ if 'labels' not in scheme:
+ required_likert_fields = ['min_label', 'max_label', 'size']
+ missing_likert_fields = [field for field in required_likert_fields if field not in scheme]
+ if missing_likert_fields:
+ raise ConfigValidationError(f"{path} missing required fields for likert: {', '.join(missing_likert_fields)}")
+
+ if not isinstance(scheme['size'], int) or scheme['size'] < 2:
+ raise ConfigValidationError(f"{path}.size must be an integer >= 2")
+
+ elif annotation_type == 'slider':
+ # Slider can use labels (falls back to radio) or min_value/max_value
+ if 'labels' not in scheme:
+ required_slider_fields = ['min_value', 'max_value', 'starting_value']
+ missing_slider_fields = [field for field in required_slider_fields if field not in scheme]
+ if missing_slider_fields:
+ raise ConfigValidationError(f"{path} missing required fields for slider: {', '.join(missing_slider_fields)}")
+
+ if not isinstance(scheme['min_value'], (int, float)) or not isinstance(scheme['max_value'], (int, float)):
+ raise ConfigValidationError(f"{path}.min_value and max_value must be numbers")
+ if scheme['min_value'] >= scheme['max_value']:
+ raise ConfigValidationError(f"{path}.min_value must be less than max_value")
+
+ elif annotation_type == 'span':
+ if 'labels' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'labels' field for span annotation type")
+ if not isinstance(scheme['labels'], list):
+ raise ConfigValidationError(f"{path}.labels must be a list")
+ if not scheme['labels']:
+ raise ConfigValidationError(f"{path}.labels cannot be empty")
+
+ elif annotation_type == 'multirate':
+ # multirate requires 'labels' always, and either 'options' or 'options_from_data'
+ if 'labels' not in scheme:
+ raise ConfigValidationError(f"{path} missing required field for multirate: labels")
+
+ has_options = 'options' in scheme
+ has_options_from_data = 'options_from_data' in scheme
+
+ if not has_options and not has_options_from_data:
+ raise ConfigValidationError(f"{path} must have either 'options' or 'options_from_data' for multirate")
+
+ if has_options:
+ if not isinstance(scheme['options'], list):
+ raise ConfigValidationError(f"{path}.options must be a list")
+ if not scheme['options']:
+ raise ConfigValidationError(f"{path}.options cannot be empty")
+
+ if has_options_from_data:
+ if not isinstance(scheme['options_from_data'], str) or not scheme['options_from_data'].strip():
+ raise ConfigValidationError(f"{path}.options_from_data must be a non-empty string (instance data field name)")
+
+ if not isinstance(scheme['labels'], list):
+ raise ConfigValidationError(f"{path}.labels must be a list")
+ if not scheme['labels']:
+ raise ConfigValidationError(f"{path}.labels cannot be empty")
+
+ elif annotation_type == 'image_annotation':
+ # Image annotation requires tools and labels
+ if 'tools' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'tools' field for image_annotation type")
+ if not isinstance(scheme['tools'], list):
+ raise ConfigValidationError(f"{path}.tools must be a list")
+ if not scheme['tools']:
+ raise ConfigValidationError(f"{path}.tools cannot be empty")
+
+ # Validate tools
+ valid_tools = ['bbox', 'polygon', 'freeform', 'landmark', 'fill', 'eraser', 'brush']
+ invalid_tools = [t for t in scheme['tools'] if t not in valid_tools]
+ if invalid_tools:
+ raise ConfigValidationError(f"{path}.tools contains invalid values: {invalid_tools}. Valid tools are: {valid_tools}")
+
+ if 'labels' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'labels' field for image_annotation type")
+ if not isinstance(scheme['labels'], list):
+ raise ConfigValidationError(f"{path}.labels must be a list")
+ if not scheme['labels']:
+ raise ConfigValidationError(f"{path}.labels cannot be empty")
+
+ # Validate optional numeric fields
+ if 'min_annotations' in scheme:
+ if not isinstance(scheme['min_annotations'], int) or scheme['min_annotations'] < 0:
+ raise ConfigValidationError(f"{path}.min_annotations must be a non-negative integer")
+
+ if 'max_annotations' in scheme and scheme['max_annotations'] is not None:
+ if not isinstance(scheme['max_annotations'], int) or scheme['max_annotations'] < 1:
+ raise ConfigValidationError(f"{path}.max_annotations must be a positive integer or null")
+
+ elif annotation_type == 'audio_annotation':
+ # Validate mode
+ valid_modes = ['label', 'questions', 'both']
+ mode = scheme.get('mode', 'label')
+ if mode not in valid_modes:
+ raise ConfigValidationError(f"{path}.mode must be one of: {valid_modes}")
+
+ # Validate labels for label/both modes
+ if mode in ['label', 'both']:
+ if 'labels' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'labels' field for audio_annotation mode '{mode}'")
+ if not isinstance(scheme['labels'], list):
+ raise ConfigValidationError(f"{path}.labels must be a list")
+ if not scheme['labels']:
+ raise ConfigValidationError(f"{path}.labels cannot be empty for mode '{mode}'")
+
+ # Validate segment_schemes for questions/both modes
+ if mode in ['questions', 'both']:
+ if 'segment_schemes' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'segment_schemes' field for audio_annotation mode '{mode}'")
+ if not isinstance(scheme['segment_schemes'], list):
+ raise ConfigValidationError(f"{path}.segment_schemes must be a list")
+ if not scheme['segment_schemes']:
+ raise ConfigValidationError(f"{path}.segment_schemes cannot be empty for mode '{mode}'")
+
+ # Validate optional numeric fields
+ if 'min_segments' in scheme:
+ if not isinstance(scheme['min_segments'], int) or scheme['min_segments'] < 0:
+ raise ConfigValidationError(f"{path}.min_segments must be a non-negative integer")
+
+ if 'max_segments' in scheme and scheme['max_segments'] is not None:
+ if not isinstance(scheme['max_segments'], int) or scheme['max_segments'] < 1:
+ raise ConfigValidationError(f"{path}.max_segments must be a positive integer or null")
+
+ elif annotation_type == 'video_annotation':
+ # Validate mode
+ valid_modes = ['segment', 'frame', 'keyframe', 'tracking', 'combined']
+ mode = scheme.get('mode', 'segment')
+ if mode not in valid_modes:
+ raise ConfigValidationError(f"{path}.mode must be one of: {valid_modes}")
+
+ # Validate labels for segment/frame/keyframe/combined modes
+ if mode in ['segment', 'frame', 'keyframe', 'combined']:
+ if 'labels' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'labels' field for video_annotation mode '{mode}'")
+ if not isinstance(scheme['labels'], list):
+ raise ConfigValidationError(f"{path}.labels must be a list")
+ if not scheme['labels']:
+ raise ConfigValidationError(f"{path}.labels cannot be empty for mode '{mode}'")
+
+ # Validate optional numeric fields
+ if 'min_segments' in scheme:
+ if not isinstance(scheme['min_segments'], int) or scheme['min_segments'] < 0:
+ raise ConfigValidationError(f"{path}.min_segments must be a non-negative integer")
+
+ if 'max_segments' in scheme and scheme['max_segments'] is not None:
+ if not isinstance(scheme['max_segments'], int) or scheme['max_segments'] < 1:
+ raise ConfigValidationError(f"{path}.max_segments must be a positive integer or null")
+
+ if 'timeline_height' in scheme:
+ if not isinstance(scheme['timeline_height'], int) or scheme['timeline_height'] < 30:
+ raise ConfigValidationError(f"{path}.timeline_height must be an integer >= 30")
+
+ if 'video_fps' in scheme:
+ if not isinstance(scheme['video_fps'], (int, float)) or scheme['video_fps'] <= 0:
+ raise ConfigValidationError(f"{path}.video_fps must be a positive number")
+
+ elif annotation_type == 'tiered_annotation':
+ # Validate required fields
+ if 'tiers' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'tiers' field for tiered_annotation")
+ if not isinstance(scheme['tiers'], list):
+ raise ConfigValidationError(f"{path}.tiers must be a list")
+ if not scheme['tiers']:
+ raise ConfigValidationError(f"{path}.tiers cannot be empty")
+
+ if 'source_field' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'source_field' field for tiered_annotation")
+
+ # Validate media_type
+ media_type = scheme.get('media_type', 'audio')
+ if media_type not in ['audio', 'video']:
+ raise ConfigValidationError(f"{path}.media_type must be 'audio' or 'video'")
+
+ # Validate tiers
+ tier_names = set()
+ valid_tier_types = ['independent', 'dependent']
+ valid_constraint_types = ['time_subdivision', 'included_in', 'symbolic_association', 'symbolic_subdivision', 'none']
+
+ for i, tier in enumerate(scheme['tiers']):
+ tier_path = f"{path}.tiers[{i}]"
+
+ if not isinstance(tier, dict):
+ raise ConfigValidationError(f"{tier_path} must be a dictionary")
+
+ if 'name' not in tier:
+ raise ConfigValidationError(f"{tier_path} missing 'name' field")
+
+ tier_name = tier['name']
+ if tier_name in tier_names:
+ raise ConfigValidationError(f"{tier_path} duplicate tier name: '{tier_name}'")
+ tier_names.add(tier_name)
+
+ # Validate tier_type
+ tier_type = tier.get('tier_type', 'independent')
+ if tier_type not in valid_tier_types:
+ raise ConfigValidationError(f"{tier_path}.tier_type must be one of: {valid_tier_types}")
+
+ # Validate dependent tier requirements
+ if tier_type == 'dependent':
+ if 'parent_tier' not in tier:
+ raise ConfigValidationError(f"{tier_path} dependent tier must have 'parent_tier'")
+
+ # Validate constraint_type
+ constraint_type = tier.get('constraint_type', 'none')
+ if constraint_type not in valid_constraint_types:
+ raise ConfigValidationError(f"{tier_path}.constraint_type must be one of: {valid_constraint_types}")
+
+ # Validate parent_tier references (second pass)
+ for i, tier in enumerate(scheme['tiers']):
+ parent = tier.get('parent_tier')
+ if parent and parent not in tier_names:
+ raise ConfigValidationError(f"{path}.tiers[{i}] references unknown parent_tier: '{parent}'")
+ if parent and parent == tier['name']:
+ raise ConfigValidationError(f"{path}.tiers[{i}] cannot be its own parent")
+
+ # Validate optional numeric fields
+ if 'tier_height' in scheme:
+ if not isinstance(scheme['tier_height'], int) or scheme['tier_height'] < 20:
+ raise ConfigValidationError(f"{path}.tier_height must be an integer >= 20")
+
+ elif annotation_type == 'pairwise':
+ # Validate mode
+ valid_modes = ['binary', 'scale', 'multi_dimension']
+ mode = scheme.get('mode', 'binary')
+ if mode not in valid_modes:
+ raise ConfigValidationError(f"{path}.mode must be one of: {valid_modes}")
+
+ # Validate labels if provided
+ if 'labels' in scheme:
+ if not isinstance(scheme['labels'], list):
+ raise ConfigValidationError(f"{path}.labels must be a list")
+ if len(scheme['labels']) < 2:
+ raise ConfigValidationError(f"{path}.labels must have at least 2 items (for A and B)")
+
+ # Validate scale configuration for scale mode
+ if mode == 'scale':
+ scale = scheme.get('scale', {})
+ if not isinstance(scale, dict):
+ raise ConfigValidationError(f"{path}.scale must be a dictionary")
+
+ # Validate min/max values
+ min_val = scale.get('min', -3)
+ max_val = scale.get('max', 3)
+ if not isinstance(min_val, (int, float)) or not isinstance(max_val, (int, float)):
+ raise ConfigValidationError(f"{path}.scale.min and scale.max must be numbers")
+ if min_val >= max_val:
+ raise ConfigValidationError(f"{path}.scale.min must be less than scale.max")
+
+ # Validate step
+ step = scale.get('step', 1)
+ if not isinstance(step, (int, float)) or step <= 0:
+ raise ConfigValidationError(f"{path}.scale.step must be a positive number")
+
+ # Validate scale labels if provided
+ if 'labels' in scale:
+ scale_labels = scale['labels']
+ if not isinstance(scale_labels, dict):
+ raise ConfigValidationError(f"{path}.scale.labels must be a dictionary")
+
+ # Validate multi_dimension mode
+ if mode == 'multi_dimension':
+ dimensions = scheme.get('dimensions', [])
+ if not isinstance(dimensions, list) or not dimensions:
+ raise ConfigValidationError(f"{path}.dimensions must be a non-empty list for multi_dimension mode")
+ for i, dim in enumerate(dimensions):
+ if not isinstance(dim, dict):
+ raise ConfigValidationError(f"{path}.dimensions[{i}] must be a dictionary")
+ if 'name' not in dim:
+ raise ConfigValidationError(f"{path}.dimensions[{i}] must have a 'name' field")
+
+ elif annotation_type == 'bws':
+ # Validate tuple_size
+ if 'tuple_size' in scheme:
+ if not isinstance(scheme['tuple_size'], int) or scheme['tuple_size'] < 2:
+ raise ConfigValidationError(f"{path}.tuple_size must be an integer >= 2")
+
+ elif annotation_type == 'soft_label':
+ if 'labels' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'labels' field for soft_label annotation type")
+ if not isinstance(scheme['labels'], list) or not scheme['labels']:
+ raise ConfigValidationError(f"{path}.labels must be a non-empty list")
+ if 'total' in scheme:
+ if not isinstance(scheme['total'], int) or scheme['total'] < 1:
+ raise ConfigValidationError(f"{path}.total must be a positive integer")
+
+ elif annotation_type == 'confidence':
+ if 'scale_type' in scheme:
+ if scheme['scale_type'] not in ['likert', 'slider']:
+ raise ConfigValidationError(f"{path}.scale_type must be 'likert' or 'slider'")
+ if 'scale_points' in scheme:
+ if not isinstance(scheme['scale_points'], int) or scheme['scale_points'] < 2:
+ raise ConfigValidationError(f"{path}.scale_points must be an integer >= 2")
+
+ elif annotation_type == 'constant_sum':
+ if 'labels' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'labels' field for constant_sum annotation type")
+ if not isinstance(scheme['labels'], list) or not scheme['labels']:
+ raise ConfigValidationError(f"{path}.labels must be a non-empty list")
+ if 'total_points' in scheme:
+ if not isinstance(scheme['total_points'], int) or scheme['total_points'] < 1:
+ raise ConfigValidationError(f"{path}.total_points must be a positive integer")
+
+ elif annotation_type == 'semantic_differential':
+ if 'pairs' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'pairs' field for semantic_differential annotation type")
+ if not isinstance(scheme['pairs'], list) or not scheme['pairs']:
+ raise ConfigValidationError(f"{path}.pairs must be a non-empty list")
+ for i, pair in enumerate(scheme['pairs']):
+ if not isinstance(pair, list) or len(pair) != 2:
+ raise ConfigValidationError(f"{path}.pairs[{i}] must be a list of exactly two strings")
+
+ elif annotation_type == 'ranking':
+ if 'labels' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'labels' field for ranking annotation type")
+ if not isinstance(scheme['labels'], list) or not scheme['labels']:
+ raise ConfigValidationError(f"{path}.labels must be a non-empty list")
+
+ elif annotation_type == 'range_slider':
+ if 'min_value' in scheme and 'max_value' in scheme:
+ if not isinstance(scheme['min_value'], (int, float)) or not isinstance(scheme['max_value'], (int, float)):
+ raise ConfigValidationError(f"{path}.min_value and max_value must be numbers")
+ if scheme['min_value'] >= scheme['max_value']:
+ raise ConfigValidationError(f"{path}.min_value must be less than max_value")
+
+ elif annotation_type == 'hierarchical_multiselect':
+ if 'taxonomy' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'taxonomy' field for hierarchical_multiselect annotation type")
+ if not isinstance(scheme['taxonomy'], dict) or not scheme['taxonomy']:
+ raise ConfigValidationError(f"{path}.taxonomy must be a non-empty dictionary")
+
+ elif annotation_type == 'vas':
+ if 'min_value' in scheme and 'max_value' in scheme:
+ if not isinstance(scheme['min_value'], (int, float)) or not isinstance(scheme['max_value'], (int, float)):
+ raise ConfigValidationError(f"{path}.min_value and max_value must be numbers")
+ if scheme['min_value'] >= scheme['max_value']:
+ raise ConfigValidationError(f"{path}.min_value must be less than max_value")
+
+ elif annotation_type == 'rubric_eval':
+ if 'criteria' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'criteria' field for rubric_eval annotation type")
+ if not isinstance(scheme['criteria'], list) or not scheme['criteria']:
+ raise ConfigValidationError(f"{path}.criteria must be a non-empty list")
+ for i, crit in enumerate(scheme['criteria']):
+ if not isinstance(crit, dict) or 'name' not in crit:
+ raise ConfigValidationError(f"{path}.criteria[{i}] must be a dict with 'name'")
+ if 'scale_points' in scheme:
+ if not isinstance(scheme['scale_points'], int) or scheme['scale_points'] < 2:
+ raise ConfigValidationError(f"{path}.scale_points must be an integer >= 2")
+
+ elif annotation_type == 'error_span':
+ if 'error_types' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'error_types' field for error_span annotation type")
+ if not isinstance(scheme['error_types'], list) or not scheme['error_types']:
+ raise ConfigValidationError(f"{path}.error_types must be a non-empty list")
+ for i, et in enumerate(scheme['error_types']):
+ if not isinstance(et, dict) or 'name' not in et:
+ raise ConfigValidationError(f"{path}.error_types[{i}] must be a dict with 'name'")
+
+ elif annotation_type == 'card_sort':
+ mode = scheme.get('mode', 'closed')
+ if mode not in ['open', 'closed']:
+ raise ConfigValidationError(f"{path}.mode must be 'open' or 'closed'")
+ if mode == 'closed':
+ if 'groups' not in scheme:
+ raise ConfigValidationError(f"{path} missing 'groups' field for card_sort in closed mode")
+ if not isinstance(scheme['groups'], list) or not scheme['groups']:
+ raise ConfigValidationError(f"{path}.groups must be a non-empty list for closed mode")
+
+ elif annotation_type == 'conjoint':
+ if 'attributes' not in scheme and 'profiles_field' not in scheme:
+ raise ConfigValidationError(f"{path} requires 'attributes' or 'profiles_field' for conjoint annotation type")
+ if 'attributes' in scheme:
+ if not isinstance(scheme['attributes'], list) or not scheme['attributes']:
+ raise ConfigValidationError(f"{path}.attributes must be a non-empty list")
+ for i, attr in enumerate(scheme['attributes']):
+ if not isinstance(attr, dict) or 'name' not in attr:
+ raise ConfigValidationError(f"{path}.attributes[{i}] must be a dict with 'name'")
+ if 'profiles_per_set' in scheme:
+ if not isinstance(scheme['profiles_per_set'], int) or scheme['profiles_per_set'] < 2:
+ raise ConfigValidationError(f"{path}.profiles_per_set must be an integer >= 2")
+
+ # Validate display_logic if present
+ if 'display_logic' in scheme:
+ validate_display_logic_structure(scheme['display_logic'], path)
+
+
+def validate_display_logic_structure(display_logic: Dict[str, Any], path: str) -> None:
+ """
+ Validate the structure of a display_logic configuration block.
+
+ This validates the syntax and structure of a single display_logic block.
+ Cross-schema validation (checking referenced schemas exist) is done separately
+ in validate_display_logic_references().
+
+ Args:
+ display_logic: The display_logic configuration
+ path: Path in the config for error reporting
+
+ Raises:
+ ConfigValidationError: If the display_logic is invalid
+ """
+ from potato.server_utils.display_logic import SUPPORTED_OPERATORS
+
+ if not isinstance(display_logic, dict):
+ raise ConfigValidationError(f"{path}.display_logic must be a dictionary")
+
+ # Must have show_when
+ if 'show_when' not in display_logic:
+ raise ConfigValidationError(f"{path}.display_logic must have 'show_when' field")
+
+ show_when = display_logic['show_when']
+ if not isinstance(show_when, list):
+ raise ConfigValidationError(f"{path}.display_logic.show_when must be a list of conditions")
+
+ if len(show_when) == 0:
+ raise ConfigValidationError(f"{path}.display_logic.show_when must have at least one condition")
+
+ # Validate each condition
+ for i, condition in enumerate(show_when):
+ cond_path = f"{path}.display_logic.show_when[{i}]"
+
+ if not isinstance(condition, dict):
+ raise ConfigValidationError(f"{cond_path} must be a dictionary")
+
+ # Required fields
+ if 'schema' not in condition:
+ raise ConfigValidationError(f"{cond_path} missing required 'schema' field")
+
+ if 'operator' not in condition:
+ raise ConfigValidationError(f"{cond_path} missing required 'operator' field")
+
+ operator = condition['operator']
+ if operator not in SUPPORTED_OPERATORS:
+ raise ConfigValidationError(
+ f"{cond_path}.operator '{operator}' is not supported. "
+ f"Valid operators: {list(SUPPORTED_OPERATORS.keys())}"
+ )
+
+ # Validate operator-specific value requirements
+ value = condition.get('value')
+
+ # Operators that don't need a value
+ if operator in ('empty', 'not_empty'):
+ pass # No value required
+ # Range operators need [min, max]
+ elif operator in ('in_range', 'not_in_range', 'length_in_range'):
+ if not isinstance(value, (list, tuple)):
+ raise ConfigValidationError(
+ f"{cond_path}: operator '{operator}' requires a range value as [min, max]"
+ )
+ if len(value) != 2:
+ raise ConfigValidationError(
+ f"{cond_path}: range value must have exactly 2 elements [min, max]"
+ )
+ try:
+ min_val, max_val = float(value[0]), float(value[1])
+ if min_val > max_val:
+ raise ConfigValidationError(
+ f"{cond_path}: range min ({min_val}) is greater than max ({max_val})"
+ )
+ except (ValueError, TypeError):
+ raise ConfigValidationError(f"{cond_path}: range values must be numeric")
+ # Numeric operators need numeric values
+ elif operator in ('gt', 'gte', 'lt', 'lte', 'length_gt', 'length_lt'):
+ if value is None:
+ raise ConfigValidationError(f"{cond_path}: operator '{operator}' requires a value")
+ try:
+ float(value)
+ except (ValueError, TypeError):
+ raise ConfigValidationError(
+ f"{cond_path}: operator '{operator}' requires a numeric value"
+ )
+ # Regex operator needs a valid pattern
+ elif operator == 'matches':
+ if value is None:
+ raise ConfigValidationError(f"{cond_path}: operator 'matches' requires a regex pattern")
+ try:
+ import re
+ re.compile(value)
+ except re.error as e:
+ raise ConfigValidationError(f"{cond_path}: invalid regex pattern '{value}': {e}")
+ # Other operators just need a non-None value
+ elif value is None:
+ raise ConfigValidationError(f"{cond_path}: operator '{operator}' requires a value")
+
+ # Validate logic field if present
+ logic = display_logic.get('logic', 'all')
+ if logic not in ('all', 'any'):
+ raise ConfigValidationError(
+ f"{path}.display_logic.logic must be 'all' or 'any', got '{logic}'"
+ )
+
+
+def validate_display_logic_references(annotation_schemes: List[Dict[str, Any]]) -> None:
+ """
+ Validate that all display_logic references point to existing schemas
+ and check for circular dependencies.
+
+ This is called after all annotation schemes have been validated individually.
+
+ Args:
+ annotation_schemes: List of annotation scheme configurations
+
+ Raises:
+ ConfigValidationError: If there are invalid references or circular dependencies
+ """
+ from potato.server_utils.display_logic import validate_display_logic_config
+
+ # Use the DisplayLogicValidator for comprehensive validation
+ is_valid, errors = validate_display_logic_config(annotation_schemes)
+
+ if not is_valid:
+ # Format errors nicely
+ error_msg = "Display logic validation errors:\n" + "\n".join(f" - {e}" for e in errors)
+ raise ConfigValidationError(error_msg)
+
+
+def validate_server_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate server configuration section.
+
+ The server section allows specifying server settings in the YAML config
+ instead of via command-line flags. CLI flags take precedence over config values.
+
+ Supported options:
+ - port: Port number to run on (1-65535)
+ - host: Host address to bind to (default: localhost)
+ - debug: Enable Flask debug mode (default: false)
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If the server configuration is invalid
+ """
+ if "server" not in config_data:
+ return # server section is optional
+
+ server_config = config_data["server"]
+
+ if not isinstance(server_config, dict):
+ raise ConfigValidationError("server configuration must be a dictionary")
+
+ # Validate port
+ if "port" in server_config:
+ port = server_config["port"]
+ if not isinstance(port, int):
+ raise ConfigValidationError("server.port must be an integer")
+ if port < 1 or port > 65535:
+ raise ConfigValidationError("server.port must be between 1 and 65535")
+
+ # Validate host
+ if "host" in server_config:
+ host = server_config["host"]
+ if not isinstance(host, str):
+ raise ConfigValidationError("server.host must be a string")
+ if not host.strip():
+ raise ConfigValidationError("server.host cannot be empty")
+
+ # Validate debug
+ if "debug" in server_config:
+ if not isinstance(server_config["debug"], bool):
+ raise ConfigValidationError("server.debug must be a boolean")
+
+
+def validate_authentication_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate authentication configuration section.
+
+ Validates OAuth/OIDC provider settings, required fields, and
+ warns about common misconfigurations.
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If the authentication configuration is invalid
+ """
+ if "authentication" not in config_data:
+ return # authentication section is optional
+
+ auth_config = config_data["authentication"]
+
+ if not isinstance(auth_config, dict):
+ raise ConfigValidationError("authentication configuration must be a dictionary")
+
+ method = auth_config.get("method", "in_memory")
+ valid_methods = ["in_memory", "database", "clerk", "oauth"]
+ if method not in valid_methods:
+ raise ConfigValidationError(
+ f"authentication.method must be one of: {', '.join(valid_methods)}. "
+ f"Got: '{method}'"
+ )
+
+ # OAuth-specific validation
+ if method == "oauth":
+ # providers is required
+ providers = auth_config.get("providers")
+ if not providers or not isinstance(providers, dict):
+ raise ConfigValidationError(
+ "authentication.providers is required when method is 'oauth' "
+ "and must be a dictionary with at least one provider"
+ )
+
+ if len(providers) == 0:
+ raise ConfigValidationError(
+ "authentication.providers must contain at least one provider"
+ )
+
+ # Validate each provider
+ for name, pconfig in providers.items():
+ if not isinstance(pconfig, dict):
+ raise ConfigValidationError(
+ f"authentication.providers.{name} must be a dictionary"
+ )
+
+ # client_id and client_secret are required
+ if "client_id" not in pconfig:
+ raise ConfigValidationError(
+ f"authentication.providers.{name}.client_id is required"
+ )
+ if "client_secret" not in pconfig:
+ raise ConfigValidationError(
+ f"authentication.providers.{name}.client_secret is required"
+ )
+
+ # Generic OIDC requires discovery_url
+ if name not in ("google", "github") and "discovery_url" not in pconfig:
+ raise ConfigValidationError(
+ f"authentication.providers.{name} requires 'discovery_url' "
+ f"for OIDC providers (only 'google' and 'github' have built-in URLs)"
+ )
+
+ # Validate optional fields
+ if "allowed_domain" in pconfig:
+ domain = pconfig["allowed_domain"]
+ if not isinstance(domain, str) or not domain.strip():
+ raise ConfigValidationError(
+ f"authentication.providers.{name}.allowed_domain must be a non-empty string"
+ )
+
+ if "allowed_org" in pconfig:
+ org = pconfig["allowed_org"]
+ if not isinstance(org, str) or not org.strip():
+ raise ConfigValidationError(
+ f"authentication.providers.{name}.allowed_org must be a non-empty string"
+ )
+
+ if "scopes" in pconfig:
+ scopes = pconfig["scopes"]
+ if not isinstance(scopes, list):
+ raise ConfigValidationError(
+ f"authentication.providers.{name}.scopes must be a list"
+ )
+
+ # Validate user_identity_field
+ identity_field = auth_config.get("user_identity_field", "email")
+ valid_fields = ["email", "username", "sub", "name"]
+ if identity_field not in valid_fields:
+ raise ConfigValidationError(
+ f"authentication.user_identity_field must be one of: "
+ f"{', '.join(valid_fields)}. Got: '{identity_field}'"
+ )
+
+ # Warn if secret_key is not set (OAuth needs stable sessions)
+ if "secret_key" not in config_data:
+ import os
+ if not os.environ.get("POTATO_SECRET_KEY"):
+ logger.warning(
+ "OAuth is configured but no 'secret_key' is set in config "
+ "and POTATO_SECRET_KEY environment variable is not set. "
+ "Sessions will be lost on server restart. "
+ "Set 'secret_key' in config or POTATO_SECRET_KEY env var."
+ )
+
+ # Database-specific validation
+ if method == "database":
+ db_url = auth_config.get("database_url")
+ if db_url:
+ if not (db_url.startswith("sqlite:///") or db_url.startswith("postgresql://")):
+ raise ConfigValidationError(
+ "authentication.database_url must start with 'sqlite:///' or 'postgresql://'. "
+ f"Got: '{db_url}'"
+ )
+
+ # Mutual exclusivity: database backend and user_config_path
+ if "user_config_path" in auth_config:
+ raise ConfigValidationError(
+ "authentication.user_config_path cannot be used with method 'database'. "
+ "The database backend handles its own user persistence."
+ )
+
+
+def validate_quality_control_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate quality control configuration (attention checks, gold standards, pre-annotation).
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If the configuration is invalid
+ """
+ # Validate attention checks config
+ if "attention_checks" in config_data:
+ attn_config = config_data["attention_checks"]
+ if not isinstance(attn_config, dict):
+ raise ConfigValidationError("attention_checks must be a dictionary")
+
+ if attn_config.get("enabled", False):
+ # Validate items_file is specified
+ if "items_file" not in attn_config:
+ raise ConfigValidationError("attention_checks.items_file is required when enabled")
+ if not isinstance(attn_config["items_file"], str):
+ raise ConfigValidationError("attention_checks.items_file must be a string path")
+
+ # Validate frequency or probability (one should be set)
+ has_frequency = "frequency" in attn_config
+ has_probability = "probability" in attn_config
+
+ if has_frequency and has_probability:
+ raise ConfigValidationError("attention_checks: specify either 'frequency' or 'probability', not both")
+
+ if has_frequency:
+ freq = attn_config["frequency"]
+ if not isinstance(freq, int) or freq < 1:
+ raise ConfigValidationError("attention_checks.frequency must be a positive integer")
+
+ if has_probability:
+ prob = attn_config["probability"]
+ if not isinstance(prob, (int, float)) or prob < 0 or prob > 1:
+ raise ConfigValidationError("attention_checks.probability must be a number between 0 and 1")
+
+ # Validate min_response_time
+ if "min_response_time" in attn_config:
+ min_time = attn_config["min_response_time"]
+ if not isinstance(min_time, (int, float)) or min_time < 0:
+ raise ConfigValidationError("attention_checks.min_response_time must be a non-negative number")
+
+ # Validate failure_handling
+ if "failure_handling" in attn_config:
+ failure_config = attn_config["failure_handling"]
+ if not isinstance(failure_config, dict):
+ raise ConfigValidationError("attention_checks.failure_handling must be a dictionary")
+
+ if "warn_threshold" in failure_config:
+ warn = failure_config["warn_threshold"]
+ if not isinstance(warn, int) or warn < 1:
+ raise ConfigValidationError("attention_checks.failure_handling.warn_threshold must be a positive integer")
+
+ if "block_threshold" in failure_config:
+ block = failure_config["block_threshold"]
+ if not isinstance(block, int) or block < 1:
+ raise ConfigValidationError("attention_checks.failure_handling.block_threshold must be a positive integer")
+
+ # Ensure block > warn
+ warn = failure_config.get("warn_threshold", 2)
+ if block <= warn:
+ raise ConfigValidationError("attention_checks.failure_handling.block_threshold must be greater than warn_threshold")
+
+ # Validate gold standards config
+ if "gold_standards" in config_data:
+ gold_config = config_data["gold_standards"]
+ if not isinstance(gold_config, dict):
+ raise ConfigValidationError("gold_standards must be a dictionary")
+
+ if gold_config.get("enabled", False):
+ # Validate items_file is specified
+ if "items_file" not in gold_config:
+ raise ConfigValidationError("gold_standards.items_file is required when enabled")
+ if not isinstance(gold_config["items_file"], str):
+ raise ConfigValidationError("gold_standards.items_file must be a string path")
+
+ # Validate mode
+ if "mode" in gold_config:
+ valid_modes = ["training", "mixed", "separate"]
+ if gold_config["mode"] not in valid_modes:
+ raise ConfigValidationError(f"gold_standards.mode must be one of: {', '.join(valid_modes)}")
+
+ # Validate frequency
+ if "frequency" in gold_config:
+ freq = gold_config["frequency"]
+ if not isinstance(freq, int) or freq < 1:
+ raise ConfigValidationError("gold_standards.frequency must be a positive integer")
+
+ # Validate accuracy config
+ if "accuracy" in gold_config:
+ accuracy_config = gold_config["accuracy"]
+ if not isinstance(accuracy_config, dict):
+ raise ConfigValidationError("gold_standards.accuracy must be a dictionary")
+
+ if "min_threshold" in accuracy_config:
+ threshold = accuracy_config["min_threshold"]
+ if not isinstance(threshold, (int, float)) or threshold < 0 or threshold > 1:
+ raise ConfigValidationError("gold_standards.accuracy.min_threshold must be between 0 and 1")
+
+ if "evaluation_count" in accuracy_config:
+ count = accuracy_config["evaluation_count"]
+ if not isinstance(count, int) or count < 1:
+ raise ConfigValidationError("gold_standards.accuracy.evaluation_count must be a positive integer")
+
+ # Validate auto_promote config
+ if "auto_promote" in gold_config:
+ auto_promote = gold_config["auto_promote"]
+ if not isinstance(auto_promote, dict):
+ raise ConfigValidationError("gold_standards.auto_promote must be a dictionary")
+
+ if "min_annotators" in auto_promote:
+ min_ann = auto_promote["min_annotators"]
+ if not isinstance(min_ann, int) or min_ann < 2:
+ raise ConfigValidationError("gold_standards.auto_promote.min_annotators must be an integer >= 2")
+
+ if "agreement_threshold" in auto_promote:
+ threshold = auto_promote["agreement_threshold"]
+ if not isinstance(threshold, (int, float)) or threshold < 0.5 or threshold > 1.0:
+ raise ConfigValidationError("gold_standards.auto_promote.agreement_threshold must be between 0.5 and 1.0")
+
+ # Validate pre-annotation config
+ if "pre_annotation" in config_data:
+ pre_config = config_data["pre_annotation"]
+ if not isinstance(pre_config, dict):
+ raise ConfigValidationError("pre_annotation must be a dictionary")
+
+ if pre_config.get("enabled", False):
+ # Validate field name
+ if "field" in pre_config:
+ if not isinstance(pre_config["field"], str) or not pre_config["field"].strip():
+ raise ConfigValidationError("pre_annotation.field must be a non-empty string")
+
+ # Validate highlight_low_confidence threshold
+ if "highlight_low_confidence" in pre_config:
+ threshold = pre_config["highlight_low_confidence"]
+ if not isinstance(threshold, (int, float)) or threshold < 0 or threshold > 1:
+ raise ConfigValidationError("pre_annotation.highlight_low_confidence must be between 0 and 1")
+
+ # Validate agreement metrics config
+ if "agreement_metrics" in config_data:
+ agreement_config = config_data["agreement_metrics"]
+ if not isinstance(agreement_config, dict):
+ raise ConfigValidationError("agreement_metrics must be a dictionary")
+
+ if "min_overlap" in agreement_config:
+ overlap = agreement_config["min_overlap"]
+ if not isinstance(overlap, int) or overlap < 2:
+ raise ConfigValidationError("agreement_metrics.min_overlap must be an integer >= 2")
+
+ if "refresh_interval" in agreement_config:
+ interval = agreement_config["refresh_interval"]
+ if not isinstance(interval, int) or interval < 10:
+ raise ConfigValidationError("agreement_metrics.refresh_interval must be an integer >= 10 seconds")
+
+
+def validate_instance_reclaim_config(config_data: Dict[str, Any]) -> None:
+ """Validate abandoned assignment reclaim configuration."""
+ if "instance_reclaim" not in config_data:
+ return
+
+ reclaim_config = config_data["instance_reclaim"]
+ if not isinstance(reclaim_config, dict):
+ raise ConfigValidationError("instance_reclaim must be a dictionary")
+
+ def validate_bool(section: Dict[str, Any], path: str) -> None:
+ if "preserve_completed_annotations" in section and not isinstance(section["preserve_completed_annotations"], bool):
+ raise ConfigValidationError(f"{path}.preserve_completed_annotations must be a boolean")
+
+ def validate_section(section_name: str) -> None:
+ if section_name not in reclaim_config:
+ return
+ section = reclaim_config[section_name]
+ if not isinstance(section, dict):
+ raise ConfigValidationError(f"instance_reclaim.{section_name} must be a dictionary")
+ validate_bool(section, f"instance_reclaim.{section_name}")
+
+ if "enabled" in reclaim_config and not isinstance(reclaim_config["enabled"], bool):
+ raise ConfigValidationError("instance_reclaim.enabled must be a boolean")
+
+ if "timeout_hours" in reclaim_config:
+ timeout = reclaim_config["timeout_hours"]
+ if not isinstance(timeout, (int, float)) or timeout <= 0:
+ raise ConfigValidationError("instance_reclaim.timeout_hours must be a positive number")
+
+ validate_bool(reclaim_config, "instance_reclaim")
+
+ for section_name in ("stale", "manual", "quality_control", "prolific"):
+ validate_section(section_name)
+
+ prolific = reclaim_config.get("prolific")
+ if isinstance(prolific, dict) and "status_policies" in prolific:
+ status_policies = prolific["status_policies"]
+ if not isinstance(status_policies, dict):
+ raise ConfigValidationError("instance_reclaim.prolific.status_policies must be a dictionary")
+
+ valid_statuses = {"RETURNED", "TIMED-OUT", "REJECTED"}
+ for status, section in status_policies.items():
+ if status not in valid_statuses:
+ raise ConfigValidationError(
+ "instance_reclaim.prolific.status_policies keys must be one of: RETURNED, TIMED-OUT, REJECTED"
+ )
+ if not isinstance(section, dict):
+ raise ConfigValidationError(
+ f"instance_reclaim.prolific.status_policies.{status} must be a dictionary"
+ )
+ validate_bool(section, f"instance_reclaim.prolific.status_policies.{status}")
+
+
+def validate_data_directory_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate data_directory configuration.
+
+ This function validates the directory watching configuration options:
+ - data_directory: Path to the directory containing data files
+ - watch_data_directory: Whether to watch for changes (default: False)
+ - watch_poll_interval: Seconds between scans (default: 5.0)
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If the configuration is invalid
+ """
+ if "data_directory" not in config_data:
+ return # data_directory is optional
+
+ data_directory = config_data["data_directory"]
+
+ # Validate data_directory is a string
+ if not isinstance(data_directory, str):
+ raise ConfigValidationError("data_directory must be a string path")
+
+ if not data_directory.strip():
+ raise ConfigValidationError("data_directory cannot be empty")
+
+ # Validate watch_data_directory if present
+ if "watch_data_directory" in config_data:
+ watch_enabled = config_data["watch_data_directory"]
+ if not isinstance(watch_enabled, bool):
+ raise ConfigValidationError("watch_data_directory must be a boolean (true/false)")
+
+ # Validate watch_poll_interval if present
+ if "watch_poll_interval" in config_data:
+ interval = config_data["watch_poll_interval"]
+ if not isinstance(interval, (int, float)):
+ raise ConfigValidationError("watch_poll_interval must be a number")
+ if interval < 1.0:
+ raise ConfigValidationError("watch_poll_interval must be at least 1.0 seconds")
+ if interval > 3600:
+ raise ConfigValidationError("watch_poll_interval cannot exceed 3600 seconds (1 hour)")
+
+
+def validate_data_sources_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate data_sources configuration for extended data loading.
+
+ This function validates the configuration for loading data from
+ various sources including URLs, cloud storage, and databases.
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If the configuration is invalid
+ """
+ data_sources = config_data.get("data_sources")
+ if not data_sources:
+ return # Empty or missing is fine - data_files can be used instead
+
+ if not isinstance(data_sources, list):
+ raise ConfigValidationError("data_sources must be a list")
+
+ # Valid source types
+ valid_types = [
+ "file", "url", "google_drive", "dropbox",
+ "s3", "huggingface", "google_sheets", "database"
+ ]
+
+ for i, source in enumerate(data_sources):
+ if not isinstance(source, dict):
+ raise ConfigValidationError(
+ f"data_sources[{i}] must be a dictionary"
+ )
+
+ source_type = source.get("type")
+ if not source_type:
+ raise ConfigValidationError(
+ f"data_sources[{i}] is missing required 'type' field"
+ )
+
+ if source_type not in valid_types:
+ raise ConfigValidationError(
+ f"data_sources[{i}] has invalid type '{source_type}'. "
+ f"Valid types: {', '.join(valid_types)}"
+ )
+
+ # Type-specific validation
+ _validate_data_source_by_type(source, source_type, i)
+
+ # Validate partial_loading configuration if present
+ _validate_partial_loading_config(config_data)
+
+ # Validate data_cache configuration if present
+ _validate_data_cache_config(config_data)
+
+
+def _validate_data_source_by_type(source: Dict, source_type: str, index: int) -> None:
+ """Validate source-specific configuration."""
+ prefix = f"data_sources[{index}]"
+
+ if source_type == "file":
+ if not source.get("path"):
+ raise ConfigValidationError(f"{prefix} (type=file) requires 'path'")
+
+ elif source_type == "url":
+ url = source.get("url")
+ if not url:
+ raise ConfigValidationError(f"{prefix} (type=url) requires 'url'")
+ if not isinstance(url, str):
+ raise ConfigValidationError(f"{prefix}.url must be a string")
+ # Basic URL format check
+ if not (url.startswith("http://") or url.startswith("https://")):
+ raise ConfigValidationError(
+ f"{prefix}.url must start with http:// or https://"
+ )
+
+ elif source_type == "google_drive":
+ if not source.get("url") and not source.get("file_id"):
+ raise ConfigValidationError(
+ f"{prefix} (type=google_drive) requires 'url' or 'file_id'"
+ )
+
+ elif source_type == "dropbox":
+ if not source.get("url") and not source.get("path"):
+ raise ConfigValidationError(
+ f"{prefix} (type=dropbox) requires 'url' or 'path'"
+ )
+ # If path is provided, access_token is required
+ if source.get("path") and not source.get("access_token"):
+ raise ConfigValidationError(
+ f"{prefix} (type=dropbox) requires 'access_token' when using 'path'"
+ )
+
+ elif source_type == "s3":
+ if not source.get("bucket"):
+ raise ConfigValidationError(f"{prefix} (type=s3) requires 'bucket'")
+ if not source.get("key"):
+ raise ConfigValidationError(f"{prefix} (type=s3) requires 'key'")
+
+ elif source_type == "huggingface":
+ if not source.get("dataset"):
+ raise ConfigValidationError(
+ f"{prefix} (type=huggingface) requires 'dataset'"
+ )
+
+ elif source_type == "google_sheets":
+ if not source.get("spreadsheet_id"):
+ raise ConfigValidationError(
+ f"{prefix} (type=google_sheets) requires 'spreadsheet_id'"
+ )
+ if not source.get("credentials_file"):
+ raise ConfigValidationError(
+ f"{prefix} (type=google_sheets) requires 'credentials_file'"
+ )
+
+ elif source_type == "database":
+ # Must have connection_string OR dialect+database
+ if not source.get("connection_string"):
+ if not source.get("dialect"):
+ raise ConfigValidationError(
+ f"{prefix} (type=database) requires 'connection_string' or 'dialect'"
+ )
+ if not source.get("database") and source.get("dialect") != "sqlite":
+ raise ConfigValidationError(
+ f"{prefix} (type=database) requires 'database' when not using sqlite"
+ )
+ # Must have query OR table
+ if not source.get("query") and not source.get("table"):
+ raise ConfigValidationError(
+ f"{prefix} (type=database) requires 'query' or 'table'"
+ )
+
+
+def _validate_partial_loading_config(config_data: Dict[str, Any]) -> None:
+ """Validate partial_loading configuration."""
+ partial = config_data.get("partial_loading")
+ if not partial:
+ return
+
+ if not isinstance(partial, dict):
+ raise ConfigValidationError("partial_loading must be a dictionary")
+
+ # Validate enabled
+ if "enabled" in partial and not isinstance(partial["enabled"], bool):
+ raise ConfigValidationError("partial_loading.enabled must be a boolean")
+
+ # Validate initial_count
+ if "initial_count" in partial:
+ count = partial["initial_count"]
+ if not isinstance(count, int) or count < 1:
+ raise ConfigValidationError(
+ "partial_loading.initial_count must be a positive integer"
+ )
+
+ # Validate batch_size
+ if "batch_size" in partial:
+ size = partial["batch_size"]
+ if not isinstance(size, int) or size < 1:
+ raise ConfigValidationError(
+ "partial_loading.batch_size must be a positive integer"
+ )
+
+ # Validate auto_load_threshold
+ if "auto_load_threshold" in partial:
+ threshold = partial["auto_load_threshold"]
+ if not isinstance(threshold, (int, float)) or not (0 <= threshold <= 1):
+ raise ConfigValidationError(
+ "partial_loading.auto_load_threshold must be between 0.0 and 1.0"
+ )
+
+
+def _validate_data_cache_config(config_data: Dict[str, Any]) -> None:
+ """Validate data_cache configuration."""
+ cache = config_data.get("data_cache")
+ if not cache:
+ return
+
+ if not isinstance(cache, dict):
+ raise ConfigValidationError("data_cache must be a dictionary")
+
+ # Validate ttl_seconds
+ if "ttl_seconds" in cache:
+ ttl = cache["ttl_seconds"]
+ if not isinstance(ttl, int) or ttl < 0:
+ raise ConfigValidationError(
+ "data_cache.ttl_seconds must be a non-negative integer"
+ )
+
+ # Validate max_size_mb
+ if "max_size_mb" in cache:
+ size = cache["max_size_mb"]
+ if not isinstance(size, int) or size < 1:
+ raise ConfigValidationError(
+ "data_cache.max_size_mb must be a positive integer"
+ )
+
+
+def validate_database_config(db_config: Dict[str, Any]) -> None:
+ """
+ Validate database configuration.
+
+ Args:
+ db_config: The database configuration
+
+ Raises:
+ ConfigValidationError: If the database configuration is invalid
+ """
+ if not isinstance(db_config, dict):
+ raise ConfigValidationError("database configuration must be a dictionary")
+
+ required_fields = ['type', 'host', 'database', 'username']
+ missing_fields = [field for field in required_fields if field not in db_config]
+ if missing_fields:
+ raise ConfigValidationError(f"Missing required database fields: {', '.join(missing_fields)}")
+
+ valid_types = ['mysql', 'file']
+ if db_config['type'] not in valid_types:
+ raise ConfigValidationError(f"Unsupported database type: {db_config['type']}. Must be one of: {', '.join(valid_types)}")
+
+ # Validate MySQL-specific fields
+ if db_config['type'] == 'mysql':
+ if 'password' not in db_config:
+ raise ConfigValidationError("MySQL database requires password")
+
+ # Validate port if specified
+ if 'port' in db_config:
+ try:
+ port = int(db_config['port'])
+ if port < 1 or port > 65535:
+ raise ConfigValidationError("Database port must be between 1 and 65535")
+ except (ValueError, TypeError):
+ raise ConfigValidationError("Database port must be a valid integer")
+
+
+def validate_file_paths(config_data: Dict[str, Any], project_dir: str, config_file_dir: str = None) -> None:
+ """
+ Validate that all file paths in the configuration are secure and exist.
+
+ Args:
+ config_data: The configuration data
+ project_dir: The project directory
+ config_file_dir: The directory containing the config file (for relative path resolution)
+
+ Raises:
+ ConfigSecurityError: If any file paths are not secure
+ ConfigValidationError: If required files don't exist
+ """
+ # Get the task_dir from config
+ task_dir = config_data.get('task_dir')
+ if not task_dir:
+ raise ConfigValidationError("task_dir is required in configuration")
+
+ # Validate task_dir exists and is secure
+ try:
+ validated_task_dir = validate_path_security(task_dir, project_dir)
+ # Don't require task_dir to exist - it's often an output directory that will be created
+ # Only validate that it's a valid path
+ except ConfigSecurityError as e:
+ raise ConfigSecurityError(f"task_dir: {str(e)}")
+
+ # Use task_dir as the base for resolving relative paths in the config
+ base_dir = validated_task_dir
+
+ # Validate data files
+ data_files = config_data.get('data_files', [])
+ for i, data_file in enumerate(data_files):
+ # Skip validation for special values
+ if data_file in [None, "null", "default"]:
+ continue
+
+ # Handle dict entries with path + optional encoding
+ if isinstance(data_file, dict):
+ file_path = data_file.get("path")
+ if not file_path:
+ raise ConfigValidationError(f"Data file {i}: dict entry missing 'path' field")
+ # Validate encoding if specified
+ encoding = data_file.get("encoding")
+ if encoding is not None:
+ if not isinstance(encoding, str):
+ raise ConfigValidationError(
+ f"Data file {i}: 'encoding' must be a string, got {type(encoding).__name__}"
+ )
+ try:
+ codecs.lookup(encoding)
+ except LookupError:
+ raise ConfigValidationError(
+ f"Data file {i}: unknown encoding '{encoding}'"
+ )
+ else:
+ file_path = data_file
+
+ try:
+ validated_path = validate_path_security(file_path, base_dir, project_dir)
+ if not os.path.exists(validated_path):
+ raise ConfigValidationError(f"Data file not found: {file_path} (resolved to: {validated_path})")
+ except ConfigSecurityError as e:
+ raise ConfigSecurityError(f"Data file {i}: {str(e)}")
+
+ # Validate batch assignment instance files
+ batch_config = config_data.get('batch_assignment')
+ if isinstance(batch_config, dict):
+ for i, group in enumerate(batch_config.get('groups') or []):
+ if not isinstance(group, dict):
+ continue
+ file_entry = group.get(
+ 'instances_file',
+ group.get('items_file', group.get('instance_ids_file')),
+ )
+ if not file_entry:
+ continue
+ if isinstance(file_entry, dict):
+ file_path = file_entry.get("path")
+ else:
+ file_path = file_entry
+
+ try:
+ validated_path = validate_path_security(file_path, base_dir, project_dir)
+ if not os.path.exists(validated_path):
+ raise ConfigValidationError(
+ f"batch_assignment.groups[{i}] file not found: "
+ f"{file_path} (resolved to: {validated_path})"
+ )
+ except ConfigSecurityError as e:
+ raise ConfigSecurityError(
+ f"batch_assignment.groups[{i}] file: {str(e)}"
+ )
+
+ # Validate data_directory if configured
+ if 'data_directory' in config_data:
+ data_directory = config_data['data_directory']
+ # Skip validation for special values
+ if data_directory not in [None, "null", "default"]:
+ try:
+ validated_dir = validate_path_security(data_directory, base_dir, project_dir)
+ if not os.path.exists(validated_dir):
+ raise ConfigValidationError(f"data_directory not found: {data_directory} (resolved to: {validated_dir})")
+ if not os.path.isdir(validated_dir):
+ raise ConfigValidationError(f"data_directory is not a directory: {data_directory} (resolved to: {validated_dir})")
+ except ConfigSecurityError as e:
+ raise ConfigSecurityError(f"data_directory: {str(e)}")
+
+ # Validate output_annotation_dir
+ if 'output_annotation_dir' in config_data:
+ output_dir = config_data['output_annotation_dir']
+ # Skip validation for special values
+ if output_dir not in [None, "null", "default"]:
+ try:
+ validate_path_security(output_dir, project_dir)
+ except ConfigSecurityError as e:
+ raise ConfigSecurityError(f"output_annotation_dir: {str(e)}")
+
+ # Validate site_dir
+ if 'site_dir' in config_data:
+ site_dir = config_data['site_dir']
+ # Skip validation for special values
+ if site_dir not in [None, "null", "default"]:
+ try:
+ validate_path_security(site_dir, base_dir, project_dir)
+ except ConfigSecurityError as e:
+ raise ConfigSecurityError(f"site_dir: {str(e)}")
+
+ # Validate custom_ds
+ if 'custom_ds' in config_data:
+ custom_ds = config_data['custom_ds']
+ # Skip validation for special values
+ if custom_ds not in [None, "null", "default"]:
+ try:
+ validate_path_security(custom_ds, base_dir, project_dir)
+ except ConfigSecurityError as e:
+ raise ConfigSecurityError(f"custom_ds: {str(e)}")
+
+ # Validate base_css
+ if 'base_css' in config_data:
+ base_css = config_data['base_css']
+ if base_css not in [None, "null", "default"]:
+ try:
+ validated_css = validate_path_security(base_css, base_dir, project_dir)
+ if not os.path.exists(validated_css):
+ # Try resolving relative to config file directory
+ if config_file_dir:
+ alt_path = os.path.join(config_file_dir, base_css)
+ if not os.path.exists(alt_path):
+ raise ConfigValidationError(
+ f"base_css file not found: {base_css} (resolved to: {validated_css})"
+ )
+ else:
+ raise ConfigValidationError(
+ f"base_css file not found: {base_css} (resolved to: {validated_css})"
+ )
+ except ConfigSecurityError as e:
+ raise ConfigSecurityError(f"base_css: {str(e)}")
+
+ # Validate header_logo
+ if 'header_logo' in config_data:
+ header_logo = config_data['header_logo']
+ if header_logo not in [None, "null", "default"]:
+ # Allow URLs to pass through without file validation
+ if not str(header_logo).startswith(("http://", "https://")):
+ try:
+ validated_logo = validate_path_security(header_logo, base_dir, project_dir)
+ if not os.path.exists(validated_logo):
+ # Try resolving relative to config file directory
+ if config_file_dir:
+ alt_path = os.path.join(config_file_dir, header_logo)
+ if not os.path.exists(alt_path):
+ raise ConfigValidationError(
+ f"header_logo file not found: {header_logo} (resolved to: {validated_logo})"
+ )
+ else:
+ raise ConfigValidationError(
+ f"header_logo file not found: {header_logo} (resolved to: {validated_logo})"
+ )
+ except ConfigSecurityError as e:
+ raise ConfigSecurityError(f"header_logo: {str(e)}")
+
+
+def validate_training_config(config_data: Dict[str, Any], project_dir: str, config_file_dir: str = None) -> None:
+ """
+ Validate training configuration.
+
+ Args:
+ config_data: The configuration data
+ project_dir: The project directory
+ config_file_dir: The directory containing the config file
+
+ Raises:
+ ConfigValidationError: If training configuration is invalid
+ ConfigSecurityError: If training data file path is not secure
+ """
+ if 'training' not in config_data:
+ return # Training is optional
+
+ training_config = config_data['training']
+ if not isinstance(training_config, dict):
+ raise ConfigValidationError("training configuration must be a dictionary")
+
+ # Validate enabled flag
+ if 'enabled' in training_config:
+ if not isinstance(training_config['enabled'], bool):
+ raise ConfigValidationError("training.enabled must be a boolean")
+
+ # If training is disabled or not specified, skip further validation
+ if not training_config.get('enabled', False):
+ return
+
+ # Validate training data file
+ if 'data_file' not in training_config:
+ raise ConfigValidationError("training.data_file is required when training is enabled")
+
+ data_file = training_config['data_file']
+ if not isinstance(data_file, str):
+ raise ConfigValidationError("training.data_file must be a string")
+
+ # Validate training data file path security and existence
+ try:
+ base_dir = config_file_dir if config_file_dir else project_dir
+ validated_path = validate_path_security(data_file, base_dir, project_dir)
+ if not os.path.exists(validated_path):
+ raise ConfigValidationError(f"Training data file not found: {data_file} (resolved to: {validated_path})")
+ except ConfigSecurityError as e:
+ raise ConfigSecurityError(f"training.data_file: {str(e)}")
+
+ # Validate annotation schemes
+ if 'annotation_schemes' in training_config:
+ schemes = training_config['annotation_schemes']
+ if not isinstance(schemes, list):
+ raise ConfigValidationError("training.annotation_schemes must be a list")
+ if not schemes:
+ raise ConfigValidationError("training.annotation_schemes cannot be empty")
+
+ for i, scheme in enumerate(schemes):
+ if isinstance(scheme, str):
+ # String reference to existing scheme - validate it's a valid string
+ if not scheme.strip():
+ raise ConfigValidationError(f"training.annotation_schemes[{i}] cannot be empty")
+ elif isinstance(scheme, dict):
+ # Full scheme dictionary - validate it
+ validate_single_annotation_scheme(scheme, f"training.annotation_schemes[{i}]")
+ else:
+ raise ConfigValidationError(f"training.annotation_schemes[{i}] must be a string or dictionary")
+
+ # Validate passing criteria
+ if 'passing_criteria' in training_config:
+ criteria = training_config['passing_criteria']
+ if not isinstance(criteria, dict):
+ raise ConfigValidationError("training.passing_criteria must be a dictionary")
+
+ # Validate min_correct
+ if 'min_correct' in criteria:
+ min_correct = criteria['min_correct']
+ if not isinstance(min_correct, int) or min_correct < 1:
+ raise ConfigValidationError("training.passing_criteria.min_correct must be a positive integer")
+
+ # Validate max_attempts
+ if 'max_attempts' in criteria:
+ max_attempts = criteria['max_attempts']
+ if not isinstance(max_attempts, int) or max_attempts < 1:
+ raise ConfigValidationError("training.passing_criteria.max_attempts must be a positive integer")
+
+ # Validate require_all_correct
+ if 'require_all_correct' in criteria:
+ if not isinstance(criteria['require_all_correct'], bool):
+ raise ConfigValidationError("training.passing_criteria.require_all_correct must be a boolean")
+
+ # Validate feedback settings
+ if 'feedback' in training_config:
+ feedback = training_config['feedback']
+ if not isinstance(feedback, dict):
+ raise ConfigValidationError("training.feedback must be a dictionary")
+
+ # Validate show_explanations
+ if 'show_explanations' in feedback:
+ if not isinstance(feedback['show_explanations'], bool):
+ raise ConfigValidationError("training.feedback.show_explanations must be a boolean")
+
+ # Validate allow_retry
+ if 'allow_retry' in feedback:
+ if not isinstance(feedback['allow_retry'], bool):
+ raise ConfigValidationError("training.feedback.allow_retry must be a boolean")
+
+ # Validate failure action
+ if 'failure_action' in training_config:
+ failure_action = training_config['failure_action']
+ valid_actions = ['move_to_done', 'repeat_training']
+ if failure_action not in valid_actions:
+ raise ConfigValidationError(f"training.failure_action must be one of: {', '.join(valid_actions)}")
+
+
+def validate_training_data_file(data_file_path: str, annotation_schemes: List[Dict[str, Any]]) -> None:
+ """
+ Validate training data file format and consistency.
+
+ Args:
+ data_file_path: Path to the training data file
+ annotation_schemes: List of annotation schemes to validate against
+
+ Raises:
+ ConfigValidationError: If training data is invalid
+ """
+ try:
+ with open(data_file_path, 'r', encoding='utf-8') as f:
+ training_data = json.load(f)
+ except (json.JSONDecodeError, UnicodeDecodeError) as e:
+ raise ConfigValidationError(f"Training data file is not valid JSON: {str(e)}")
+ except FileNotFoundError:
+ raise ConfigValidationError(f"Training data file not found: {data_file_path}")
+
+ if not isinstance(training_data, dict):
+ raise ConfigValidationError("Training data must be a JSON object")
+
+ if 'training_instances' not in training_data:
+ raise ConfigValidationError("Training data must contain 'training_instances' field")
+
+ training_instances = training_data['training_instances']
+ if not isinstance(training_instances, list):
+ raise ConfigValidationError("training_instances must be a list")
+
+ if not training_instances:
+ raise ConfigValidationError("training_instances cannot be empty")
+
+ # Create a mapping of scheme names for validation
+ scheme_names = {scheme['name'] for scheme in annotation_schemes}
+
+ for i, instance in enumerate(training_instances):
+ if not isinstance(instance, dict):
+ raise ConfigValidationError(f"Training instance {i} must be a dictionary")
+
+ # Validate required fields
+ required_fields = ['id', 'text', 'correct_answers']
+ missing_fields = [field for field in required_fields if field not in instance]
+ if missing_fields:
+ raise ConfigValidationError(f"Training instance {i} missing required fields: {', '.join(missing_fields)}")
+
+ # Validate id
+ if not isinstance(instance['id'], str):
+ raise ConfigValidationError(f"Training instance {i}.id must be a string")
+
+ # Validate text
+ if not isinstance(instance['text'], str):
+ raise ConfigValidationError(f"Training instance {i}.text must be a string")
+
+ # Validate correct_answers
+ correct_answers = instance['correct_answers']
+ if not isinstance(correct_answers, dict):
+ raise ConfigValidationError(f"Training instance {i}.correct_answers must be a dictionary")
+
+ # Validate that all correct_answers correspond to annotation schemes
+ for scheme_name, answer in correct_answers.items():
+ if scheme_name not in scheme_names:
+ raise ConfigValidationError(f"Training instance {i}.correct_answers contains unknown scheme: {scheme_name}")
+
+ # Validate explanation if present
+ if 'explanation' in instance:
+ if not isinstance(instance['explanation'], str):
+ raise ConfigValidationError(f"Training instance {i}.explanation must be a string")
+
+
+def validate_batch_assignment_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate batch assignment configuration.
+
+ ``batch_assignment`` supports explicit annotator cohorts for repeat-round
+ studies. Each group defines annotators allowed to receive a fixed item set,
+ either inline or through a separate supported data file. Items may also
+ carry annotator lists via ``annotator_key``; that field is validated at
+ assignment time because data files load later.
+ """
+ if 'batch_assignment' not in config_data:
+ return
+
+ batch_config = config_data['batch_assignment']
+ if not isinstance(batch_config, dict):
+ raise ConfigValidationError("batch_assignment must be a dictionary")
+
+ annotator_key = batch_config.get('annotator_key')
+ if annotator_key is not None and (
+ not isinstance(annotator_key, str) or not annotator_key.strip()
+ ):
+ raise ConfigValidationError("batch_assignment.annotator_key must be a non-empty string")
+
+ groups = batch_config.get('groups', [])
+ if groups is None:
+ return
+ if not isinstance(groups, list):
+ raise ConfigValidationError("batch_assignment.groups must be a list")
+
+ for idx, group in enumerate(groups):
+ if not isinstance(group, dict):
+ raise ConfigValidationError(f"batch_assignment.groups[{idx}] must be a dictionary")
+
+ users = group.get('annotators', group.get('users'))
+ instances = group.get('instances', group.get('items', group.get('instance_ids')))
+ file_entry = group.get(
+ 'instances_file',
+ group.get('items_file', group.get('instance_ids_file')),
+ )
+
+ if not isinstance(users, list) or not users:
+ raise ConfigValidationError(
+ f"batch_assignment.groups[{idx}] must define non-empty annotators/users list"
+ )
+ if not all(isinstance(user, str) and user.strip() for user in users):
+ raise ConfigValidationError(
+ f"batch_assignment.groups[{idx}].annotators/users must contain non-empty strings"
+ )
+
+ has_instances = instances is not None
+ has_file = file_entry is not None
+
+ if not has_instances and not has_file:
+ raise ConfigValidationError(
+ f"batch_assignment.groups[{idx}] must define either "
+ "instances/items/instance_ids or instances_file/items_file/instance_ids_file"
+ )
+
+ if has_instances and (not isinstance(instances, list) or not instances):
+ raise ConfigValidationError(
+ f"batch_assignment.groups[{idx}] must define non-empty instances/items/instance_ids list"
+ )
+ if has_instances and not all(isinstance(instance, str) and instance.strip() for instance in instances):
+ raise ConfigValidationError(
+ f"batch_assignment.groups[{idx}].instances/items/instance_ids must contain non-empty strings"
+ )
+
+ if has_file:
+ if isinstance(file_entry, str):
+ if not file_entry.strip():
+ raise ConfigValidationError(
+ f"batch_assignment.groups[{idx}] file path must be non-empty"
+ )
+ elif isinstance(file_entry, dict):
+ path = file_entry.get('path')
+ if not isinstance(path, str) or not path.strip():
+ raise ConfigValidationError(
+ f"batch_assignment.groups[{idx}] file entry must define a non-empty path"
+ )
+ encoding = file_entry.get('encoding')
+ if encoding is not None and not isinstance(encoding, str):
+ raise ConfigValidationError(
+ f"batch_assignment.groups[{idx}] file encoding must be a string"
+ )
+ else:
+ raise ConfigValidationError(
+ f"batch_assignment.groups[{idx}] file entry must be a path string or mapping"
+ )
+
+
+def validate_category_assignment_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate category assignment configuration.
+
+ This function validates the category_assignment configuration section which
+ controls how users are assigned to annotation categories based on their
+ training/prestudy performance.
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If category assignment configuration is invalid
+ """
+ if 'category_assignment' not in config_data:
+ return # Category assignment is optional
+
+ cat_config = config_data['category_assignment']
+ if not isinstance(cat_config, dict):
+ raise ConfigValidationError("category_assignment must be a dictionary")
+
+ # Validate enabled flag
+ if 'enabled' in cat_config:
+ if not isinstance(cat_config['enabled'], bool):
+ raise ConfigValidationError("category_assignment.enabled must be a boolean")
+
+ # If not enabled, skip further validation
+ if not cat_config.get('enabled', True):
+ return
+
+ # Validate category_key (optional, can also be in item_properties)
+ if 'category_key' in cat_config:
+ if not isinstance(cat_config['category_key'], str) or not cat_config['category_key'].strip():
+ raise ConfigValidationError("category_assignment.category_key must be a non-empty string")
+
+ # Validate qualification settings
+ if 'qualification' in cat_config:
+ qual = cat_config['qualification']
+ if not isinstance(qual, dict):
+ raise ConfigValidationError("category_assignment.qualification must be a dictionary")
+
+ # Validate source
+ if 'source' in qual:
+ valid_sources = ['training', 'prestudy', 'both']
+ if qual['source'] not in valid_sources:
+ raise ConfigValidationError(
+ f"category_assignment.qualification.source must be one of: {', '.join(valid_sources)}"
+ )
+
+ # Validate threshold
+ if 'threshold' in qual:
+ threshold = qual['threshold']
+ if not isinstance(threshold, (int, float)) or threshold < 0.0 or threshold > 1.0:
+ raise ConfigValidationError(
+ "category_assignment.qualification.threshold must be a number between 0.0 and 1.0"
+ )
+
+ # Validate min_questions
+ if 'min_questions' in qual:
+ min_q = qual['min_questions']
+ if not isinstance(min_q, int) or min_q < 1:
+ raise ConfigValidationError(
+ "category_assignment.qualification.min_questions must be a positive integer"
+ )
+
+ # Validate combine_method (for combining prestudy and training scores)
+ if 'combine_method' in qual:
+ valid_methods = ['average', 'max', 'sum']
+ if qual['combine_method'] not in valid_methods:
+ raise ConfigValidationError(
+ f"category_assignment.qualification.combine_method must be one of: {', '.join(valid_methods)}"
+ )
+
+ # Validate fallback behavior
+ if 'fallback' in cat_config:
+ valid_fallbacks = ['uncategorized', 'random', 'none']
+ if cat_config['fallback'] not in valid_fallbacks:
+ raise ConfigValidationError(
+ f"category_assignment.fallback must be one of: {', '.join(valid_fallbacks)}"
+ )
+
+ # Validate dynamic expertise settings
+ if 'dynamic' in cat_config:
+ dynamic = cat_config['dynamic']
+ if not isinstance(dynamic, dict):
+ raise ConfigValidationError("category_assignment.dynamic must be a dictionary")
+
+ # Validate enabled flag
+ if 'enabled' in dynamic:
+ if not isinstance(dynamic['enabled'], bool):
+ raise ConfigValidationError("category_assignment.dynamic.enabled must be a boolean")
+
+ # If dynamic is not enabled, skip further validation
+ if not dynamic.get('enabled', False):
+ return
+
+ # Validate agreement_method
+ if 'agreement_method' in dynamic:
+ valid_methods = ['majority_vote', 'super_majority', 'unanimous']
+ if dynamic['agreement_method'] not in valid_methods:
+ raise ConfigValidationError(
+ f"category_assignment.dynamic.agreement_method must be one of: {', '.join(valid_methods)}"
+ )
+
+ # Validate min_annotations_for_consensus
+ if 'min_annotations_for_consensus' in dynamic:
+ min_ann = dynamic['min_annotations_for_consensus']
+ if not isinstance(min_ann, int) or min_ann < 2:
+ raise ConfigValidationError(
+ "category_assignment.dynamic.min_annotations_for_consensus must be an integer >= 2"
+ )
+
+ # Validate learning_rate
+ if 'learning_rate' in dynamic:
+ lr = dynamic['learning_rate']
+ if not isinstance(lr, (int, float)) or lr <= 0.0 or lr > 1.0:
+ raise ConfigValidationError(
+ "category_assignment.dynamic.learning_rate must be a number between 0.0 (exclusive) and 1.0"
+ )
+
+ # Validate update_interval_seconds
+ if 'update_interval_seconds' in dynamic:
+ interval = dynamic['update_interval_seconds']
+ if not isinstance(interval, (int, float)) or interval < 1:
+ raise ConfigValidationError(
+ "category_assignment.dynamic.update_interval_seconds must be a number >= 1"
+ )
+
+ # Validate base_probability
+ if 'base_probability' in dynamic:
+ base_prob = dynamic['base_probability']
+ if not isinstance(base_prob, (int, float)) or base_prob < 0.0 or base_prob > 1.0:
+ raise ConfigValidationError(
+ "category_assignment.dynamic.base_probability must be a number between 0.0 and 1.0"
+ )
+
+
+def validate_diversity_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate diversity ordering configuration.
+
+ This function validates the diversity_ordering section which controls
+ embedding-based clustering for diverse item ordering.
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If diversity ordering configuration is invalid
+ """
+ if 'diversity_ordering' not in config_data:
+ return # Diversity ordering is optional
+
+ dc = config_data['diversity_ordering']
+ if not isinstance(dc, dict):
+ raise ConfigValidationError("diversity_ordering must be a dictionary")
+
+ # Validate enabled flag
+ if 'enabled' in dc:
+ if not isinstance(dc['enabled'], bool):
+ raise ConfigValidationError("diversity_ordering.enabled must be a boolean")
+
+ # If not enabled, skip further validation
+ if not dc.get('enabled', False):
+ return
+
+ # Validate model_name
+ if 'model_name' in dc:
+ if not isinstance(dc['model_name'], str) or not dc['model_name'].strip():
+ raise ConfigValidationError("diversity_ordering.model_name must be a non-empty string")
+
+ # Validate num_clusters
+ if 'num_clusters' in dc:
+ num_clusters = dc['num_clusters']
+ if not isinstance(num_clusters, int) or num_clusters < 2:
+ raise ConfigValidationError("diversity_ordering.num_clusters must be an integer >= 2")
+
+ # Validate items_per_cluster
+ if 'items_per_cluster' in dc:
+ items_per_cluster = dc['items_per_cluster']
+ if not isinstance(items_per_cluster, int) or items_per_cluster < 1:
+ raise ConfigValidationError("diversity_ordering.items_per_cluster must be a positive integer")
+
+ # Validate auto_clusters
+ if 'auto_clusters' in dc:
+ if not isinstance(dc['auto_clusters'], bool):
+ raise ConfigValidationError("diversity_ordering.auto_clusters must be a boolean")
+
+ # Validate prefill_count
+ if 'prefill_count' in dc:
+ prefill_count = dc['prefill_count']
+ if not isinstance(prefill_count, int) or prefill_count < 0:
+ raise ConfigValidationError("diversity_ordering.prefill_count must be a non-negative integer")
+
+ # Validate batch_size
+ if 'batch_size' in dc:
+ batch_size = dc['batch_size']
+ if not isinstance(batch_size, int) or batch_size < 1:
+ raise ConfigValidationError("diversity_ordering.batch_size must be a positive integer")
+
+ # Validate recluster_threshold
+ if 'recluster_threshold' in dc:
+ recluster_threshold = dc['recluster_threshold']
+ if not isinstance(recluster_threshold, (int, float)) or recluster_threshold < 0 or recluster_threshold > 1:
+ raise ConfigValidationError(
+ "diversity_ordering.recluster_threshold must be a number between 0 and 1"
+ )
+
+ # Validate preserve_visited
+ if 'preserve_visited' in dc:
+ if not isinstance(dc['preserve_visited'], bool):
+ raise ConfigValidationError("diversity_ordering.preserve_visited must be a boolean")
+
+ # Validate trigger_ai_prefetch
+ if 'trigger_ai_prefetch' in dc:
+ if not isinstance(dc['trigger_ai_prefetch'], bool):
+ raise ConfigValidationError("diversity_ordering.trigger_ai_prefetch must be a boolean")
+
+ # Validate cache_dir
+ if 'cache_dir' in dc:
+ cache_dir = dc['cache_dir']
+ if cache_dir is not None and (not isinstance(cache_dir, str) or not cache_dir.strip()):
+ raise ConfigValidationError(
+ "diversity_ordering.cache_dir must be a non-empty string or null"
+ )
+
+
+def validate_embedding_visualization_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate embedding visualization configuration.
+
+ This function validates the embedding_visualization section which controls
+ the admin dashboard 2D visualization of embeddings.
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If embedding visualization configuration is invalid
+ """
+ if 'embedding_visualization' not in config_data:
+ return # Embedding visualization is optional
+
+ ev = config_data['embedding_visualization']
+ if not isinstance(ev, dict):
+ raise ConfigValidationError("embedding_visualization must be a dictionary")
+
+ # Validate enabled flag
+ if 'enabled' in ev:
+ if not isinstance(ev['enabled'], bool):
+ raise ConfigValidationError("embedding_visualization.enabled must be a boolean")
+
+ # If not enabled, skip further validation
+ if not ev.get('enabled', True):
+ return
+
+ # Validate sample_size
+ if 'sample_size' in ev:
+ sample_size = ev['sample_size']
+ if not isinstance(sample_size, int) or sample_size < 1:
+ raise ConfigValidationError(
+ "embedding_visualization.sample_size must be a positive integer"
+ )
+
+ # Validate include_all_annotated
+ if 'include_all_annotated' in ev:
+ if not isinstance(ev['include_all_annotated'], bool):
+ raise ConfigValidationError(
+ "embedding_visualization.include_all_annotated must be a boolean"
+ )
+
+ # Validate embedding_model
+ if 'embedding_model' in ev:
+ if not isinstance(ev['embedding_model'], str) or not ev['embedding_model'].strip():
+ raise ConfigValidationError(
+ "embedding_visualization.embedding_model must be a non-empty string"
+ )
+
+ # Validate image_embedding_model
+ if 'image_embedding_model' in ev:
+ if not isinstance(ev['image_embedding_model'], str) or not ev['image_embedding_model'].strip():
+ raise ConfigValidationError(
+ "embedding_visualization.image_embedding_model must be a non-empty string"
+ )
+
+ # Validate UMAP configuration
+ if 'umap' in ev:
+ umap_config = ev['umap']
+ if not isinstance(umap_config, dict):
+ raise ConfigValidationError("embedding_visualization.umap must be a dictionary")
+
+ # Validate n_neighbors
+ if 'n_neighbors' in umap_config:
+ n_neighbors = umap_config['n_neighbors']
+ if not isinstance(n_neighbors, int) or n_neighbors < 2:
+ raise ConfigValidationError(
+ "embedding_visualization.umap.n_neighbors must be an integer >= 2"
+ )
+
+ # Validate min_dist
+ if 'min_dist' in umap_config:
+ min_dist = umap_config['min_dist']
+ if not isinstance(min_dist, (int, float)) or min_dist < 0 or min_dist > 1:
+ raise ConfigValidationError(
+ "embedding_visualization.umap.min_dist must be a number between 0 and 1"
+ )
+
+ # Validate metric
+ if 'metric' in umap_config:
+ valid_metrics = ['cosine', 'euclidean', 'manhattan', 'correlation']
+ if umap_config['metric'] not in valid_metrics:
+ raise ConfigValidationError(
+ f"embedding_visualization.umap.metric must be one of: {valid_metrics}"
+ )
+
+ # Validate label_source
+ if 'label_source' in ev:
+ valid_sources = ['mace', 'majority']
+ if ev['label_source'] not in valid_sources:
+ raise ConfigValidationError(
+ f"embedding_visualization.label_source must be one of: {valid_sources}"
+ )
+
+
+def _merge_ai_config_file(config_data: Dict[str, Any], config_dir: str) -> Dict[str, Any]:
+ """
+ Merge an external ai-config.yaml into the main config if specified.
+
+ When ai_support.ai_config_file is set, loads that YAML file and merges its
+ contents into the ai_support section. The external file provides endpoint-specific
+ details (endpoint_type, model, api_key, base_url) while the inline ai_config
+ provides defaults (temperature, max_tokens, include settings).
+
+ Args:
+ config_data: The parsed main configuration dictionary
+ config_dir: Directory containing the main config file (for resolving relative paths)
+
+ Returns:
+ The config_data with external AI config merged in (modified in place and returned)
+ """
+ ai_support = config_data.get("ai_support", {})
+ if not isinstance(ai_support, dict):
+ return config_data
+
+ ai_config_file = ai_support.get("ai_config_file")
+
+ if not ai_config_file:
+ # No external file specified - apply env var substitution to inline ai_config
+ if "ai_config" in ai_support:
+ from potato.data_sources.credentials import substitute_env_vars
+ ai_support["ai_config"] = substitute_env_vars(ai_support["ai_config"])
+ config_data["ai_support"] = ai_support
+ return config_data
+
+ if not isinstance(ai_config_file, str):
+ logger.warning("ai_support.ai_config_file must be a string. Ignoring.")
+ return config_data
+
+ # Resolve relative to config file directory
+ ai_config_path = os.path.join(config_dir, ai_config_file)
+
+ if not os.path.exists(ai_config_path):
+ logger.warning(
+ f"AI config file '{ai_config_file}' not found at {ai_config_path}. "
+ f"AI support will be disabled. Create this file with your endpoint details."
+ )
+ config_data["ai_support"]["enabled"] = False
+ return config_data
+
+ # Load external AI config
+ try:
+ with open(ai_config_path, 'r', encoding='utf-8') as f:
+ external_config = yaml.safe_load(f) or {}
+ except yaml.YAMLError as e:
+ logger.warning(f"Invalid YAML in AI config file '{ai_config_file}': {e}. AI support will be disabled.")
+ config_data["ai_support"]["enabled"] = False
+ return config_data
+
+ if not isinstance(external_config, dict):
+ logger.warning(f"AI config file '{ai_config_file}' must contain a YAML dictionary. AI support will be disabled.")
+ config_data["ai_support"]["enabled"] = False
+ return config_data
+
+ # Apply environment variable substitution to external config
+ from potato.data_sources.credentials import substitute_env_vars
+ external_config = substitute_env_vars(external_config)
+
+ # Extract endpoint_type from external config (top-level key)
+ if "endpoint_type" in external_config:
+ ai_support["endpoint_type"] = external_config.pop("endpoint_type")
+
+ # Merge remaining keys into ai_config (external takes precedence)
+ ai_config = ai_support.get("ai_config", {})
+ if not isinstance(ai_config, dict):
+ ai_config = {}
+ ai_config.update(external_config)
+ ai_support["ai_config"] = ai_config
+
+ # Also apply env var substitution to the final merged ai_config
+ ai_support["ai_config"] = substitute_env_vars(ai_support["ai_config"])
+
+ config_data["ai_support"] = ai_support
+ logger.info(f"Loaded AI endpoint config from {ai_config_file}")
+ return config_data
+
+
+def load_and_validate_config(config_file: str, project_dir: str) -> Dict[str, Any]:
+ """
+ Load and validate a YAML configuration file with security checks.
+
+ Args:
+ config_file: Path to the configuration file
+ project_dir: The project directory
+
+ Returns:
+ The validated configuration dictionary
+
+ Raises:
+ ConfigSecurityError: If the configuration file is not secure
+ ConfigValidationError: If the configuration is invalid
+ FileNotFoundError: If the configuration file doesn't exist
+ """
+ # Validate the config file path itself
+ try:
+ validated_config_path = validate_path_security(config_file, project_dir)
+ except ConfigSecurityError as e:
+ raise ConfigSecurityError(f"Configuration file path: {str(e)}")
+
+ if not os.path.exists(validated_config_path):
+ raise FileNotFoundError(f"Configuration file not found: {config_file}")
+
+ # Load and parse YAML
+ try:
+ with open(validated_config_path, 'r', encoding='utf-8') as file_p:
+ config_data = yaml.safe_load(file_p)
+ except yaml.YAMLError as e:
+ raise ConfigValidationError(f"Invalid YAML format in {config_file}: {str(e)}")
+ except UnicodeDecodeError as e:
+ raise ConfigValidationError(f"Invalid file encoding in {config_file}: {str(e)}")
+ except Exception as e:
+ raise ConfigValidationError(f"Error reading configuration file {config_file}: {str(e)}")
+
+ # Get the directory containing the config file for relative path resolution
+ config_file_dir = os.path.dirname(validated_config_path)
+
+ # Merge external AI config file if specified (before validation)
+ config_data = _merge_ai_config_file(config_data, config_file_dir)
+
+ # Apply default values for common configuration options
+ if 'task_dir' not in config_data:
+ config_data['task_dir'] = '.'
+ logger.debug("task_dir not specified, defaulting to '.'")
+ if 'site_dir' not in config_data:
+ config_data['site_dir'] = 'default'
+ logger.debug("site_dir not specified, defaulting to 'default'")
+
+ # Resolve task_dir relative to config file directory if it's '.' or a relative path
+ if 'task_dir' in config_data:
+ task_dir = config_data['task_dir']
+ if task_dir == '.' or not os.path.isabs(task_dir):
+ # Resolve relative to config file's directory
+ task_dir = os.path.normpath(os.path.join(config_file_dir, task_dir))
+ config_data['task_dir'] = task_dir
+ logger.debug(f"Resolved task_dir to: {task_dir}")
+
+ # Validate the configuration structure
+ validate_yaml_structure(config_data, project_dir, config_file_dir)
+
+ # Validate file paths
+ validate_file_paths(config_data, project_dir, config_file_dir)
+
+ return config_data
+
+
+def init_config(args):
+ global config
+
+ project_dir = os.getcwd() #get the current working dir as the default project_dir
+ config_file = None
+ config_file_dir = None
+
+ try:
+ # if the .yaml config file is given, directly use it
+ if args.config_file[-5:] == '.yaml':
+ if os.path.exists(args.config_file):
+ print("INFO: when you run the server directly from a .yaml file, please make sure your config file is put in the annotation project folder")
+ config_file = args.config_file
+ # For direct YAML file usage, we'll determine the project_dir from the config file content
+ # after loading it, not from the file path structure
+ else:
+ raise FileNotFoundError(f"Configuration file not found: {args.config_file}")
+
+ # if the user gives a directory, check if config.yaml or configs/config.yaml exists
+ elif os.path.isdir(args.config_file):
+ project_dir = args.config_file if os.path.isabs(args.config_file) else os.path.join(project_dir, args.config_file)
+ config_folder = os.path.join(args.config_file, 'configs')
+ if not os.path.isdir(config_folder):
+ raise ConfigValidationError(".yaml file must be put in the configs/ folder under the main project directory when you try to start the project with the project directory, otherwise please directly give the path of the .yaml file")
+
+ #get all the config files
+ yamlfiles = [it for it in os.listdir(config_folder) if it[-5:] == '.yaml']
+
+ # if no yaml files found, quit the program
+ if len(yamlfiles) == 0:
+ raise ConfigValidationError(f"Configuration file not found under {config_folder}, please make sure .yaml file exists in the given directory, or please directly give the path of the .yaml file")
+ # if only one yaml file found, directly use it
+ elif len(yamlfiles) == 1:
+ config_file = os.path.join(config_folder, yamlfiles[0])
+ config_file_dir = config_folder
+
+ # if multiple yaml files found, ask the user to choose which one to use
+ else:
+ while True:
+ print("multiple config files found, please select the one you want to use (number 0-%d)"%len(yamlfiles))
+ for i,it in enumerate(yamlfiles):
+ print("[%d] %s"%(i, it))
+ input_id = input("number: ")
+ try:
+ config_file = os.path.join(config_folder, yamlfiles[int(input_id)])
+ config_file_dir = config_folder
+ break
+ except Exception:
+ print("wrong input, please reselect")
+
+ if not config_file:
+ raise ConfigValidationError(f"Configuration file not found under {config_folder}, please make sure .yaml file exists in the given directory, or please directly give the path of the .yaml file")
+
+ # Load and validate the configuration
+ # For direct config file usage, use current working directory as base for config file path resolution
+ if args.config_file[-5:] == '.yaml':
+ # First, load the config without full validation to get the task_dir
+ try:
+ validated_config_path = validate_path_security(config_file, os.getcwd())
+ with open(validated_config_path, 'r', encoding='utf-8') as file_p:
+ temp_config_data = yaml.safe_load(file_p)
+ except Exception as e:
+ raise ConfigValidationError(f"Error loading configuration file: {str(e)}")
+
+ # Get the config file's directory for resolving relative paths
+ config_file_abs = os.path.abspath(config_file)
+ config_file_dir = os.path.dirname(config_file_abs)
+
+ # Resolve task_dir relative to config file directory if it's '.' or a relative path
+ if 'task_dir' in temp_config_data:
+ task_dir = temp_config_data['task_dir']
+ if task_dir == '.' or not os.path.isabs(task_dir):
+ # Resolve relative to config file's directory
+ task_dir = os.path.normpath(os.path.join(config_file_dir, task_dir))
+ temp_config_data['task_dir'] = task_dir
+ logger.debug(f"Resolved task_dir to: {task_dir}")
+
+ # Validate that config file is in task_dir (skip in test mode)
+ skip_path_validation = os.environ.get('POTATO_SKIP_CONFIG_PATH_VALIDATION', '').lower() in ('1', 'true')
+ if 'task_dir' in temp_config_data and not skip_path_validation:
+ task_dir = temp_config_data['task_dir']
+ task_dir_abs = os.path.abspath(task_dir)
+ if not config_file_abs.startswith(task_dir_abs):
+ raise ConfigValidationError(f"Configuration file must be in the task_dir. Config file is at '{config_file_abs}' but task_dir is '{task_dir_abs}'")
+ project_dir = task_dir
+
+ # Now load and validate with the correct project_dir
+ config_data = load_and_validate_config(config_file, os.getcwd())
+ # Update config_data with resolved task_dir
+ if 'task_dir' in temp_config_data:
+ config_data['task_dir'] = temp_config_data['task_dir']
+ else:
+ config_data = load_and_validate_config(config_file, project_dir)
+
+ config.update(config_data)
+
+ # Only override config settings if command line arguments are explicitly provided
+ config_updates = {
+ "verbose": args.verbose,
+ "very_verbose": args.very_verbose,
+ # Store an ABSOLUTE path: the server chdir's into task_dir at startup,
+ # so a relative path would be re-resolved against the wrong CWD later
+ # (e.g. admin export doubled the project path). CWD is still the
+ # original launch dir here (chdir happens further below).
+ "__config_file__": os.path.abspath(args.config_file),
+ "customjs": args.customjs,
+ "customjs_hostname": args.customjs_hostname,
+ "persist_sessions": args.persist_sessions,
+ }
+
+ # Only override debug if explicitly set to True via command line
+ # or if config file doesn't have a debug setting
+ if args.debug or "debug" not in config:
+ config_updates["debug"] = args.debug
+
+ # Add debug logging mode if specified
+ if hasattr(args, 'debug_log') and args.debug_log:
+ config_updates["debug_log"] = args.debug_log
+
+ # Add debug phase if specified (requires --debug flag)
+ if hasattr(args, 'debug_phase') and args.debug_phase:
+ if not args.debug:
+ print("โ ๏ธ Warning: --debug-phase requires --debug flag. Enabling debug mode.")
+ config_updates["debug"] = True
+ config_updates["debug_phase"] = args.debug_phase
+
+ config.update(config_updates)
+
+ # Apply server config values (CLI args take precedence)
+ if "server" in config:
+ server_config = config["server"]
+
+ # Apply port from server config if not specified via CLI
+ if "port" in server_config and args.port is None:
+ config["port"] = server_config["port"]
+ logger.debug(f"Port set from config file: {server_config['port']}")
+
+ # Apply host from server config
+ if "host" in server_config:
+ # Host can only be set via config (no CLI arg currently)
+ config["host"] = server_config["host"]
+ logger.debug(f"Host set from config file: {server_config['host']}")
+
+ # Apply debug from server config if not specified via CLI
+ if "debug" in server_config and not args.debug:
+ config["debug"] = server_config["debug"]
+ logger.debug(f"Debug mode set from config file: {server_config['debug']}")
+
+ # update the current working dir for the server
+ os.chdir(project_dir)
+
+ except (ConfigSecurityError, ConfigValidationError, FileNotFoundError) as e:
+ logger.error(f"Configuration error: {str(e)}")
+ print(f"โ Configuration error: {str(e)}")
+ print("Please check your configuration file and try again.")
+ raise
+ except Exception as e:
+ logger.error(f"Unexpected error during configuration initialization: {str(e)}")
+ print(f"โ Unexpected error: {str(e)}")
+ raise
+
+
+def validate_active_learning_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate active learning configuration.
+
+ Args:
+ config_data: The configuration data containing active_learning section
+
+ Raises:
+ ConfigValidationError: If the active learning configuration is invalid
+ """
+ if "active_learning" not in config_data:
+ return # Active learning is optional
+
+ al_config = config_data["active_learning"]
+
+ # Validate enabled flag
+ if not isinstance(al_config.get("enabled", False), bool):
+ raise ConfigValidationError("active_learning.enabled must be a boolean")
+
+ if not al_config.get("enabled", False):
+ return # Skip validation if not enabled
+
+ # Validate classifier configuration
+ if "classifier" in al_config:
+ classifier_config = al_config["classifier"]
+ if not isinstance(classifier_config, dict):
+ raise ConfigValidationError("active_learning.classifier must be a dictionary")
+
+ if "name" not in classifier_config:
+ raise ConfigValidationError("active_learning.classifier.name is required")
+
+ if not isinstance(classifier_config["name"], str):
+ raise ConfigValidationError("active_learning.classifier.name must be a string")
+
+ # Validate hyperparameters if present
+ if "hyperparameters" in classifier_config:
+ if not isinstance(classifier_config["hyperparameters"], dict):
+ raise ConfigValidationError("active_learning.classifier.hyperparameters must be a dictionary")
+
+ # Validate vectorizer configuration
+ if "vectorizer" in al_config:
+ vectorizer_config = al_config["vectorizer"]
+ if not isinstance(vectorizer_config, dict):
+ raise ConfigValidationError("active_learning.vectorizer must be a dictionary")
+
+ if "name" not in vectorizer_config:
+ raise ConfigValidationError("active_learning.vectorizer.name is required")
+
+ if not isinstance(vectorizer_config["name"], str):
+ raise ConfigValidationError("active_learning.vectorizer.name must be a string")
+
+ # Validate hyperparameters if present
+ if "hyperparameters" in vectorizer_config:
+ if not isinstance(vectorizer_config["hyperparameters"], dict):
+ raise ConfigValidationError("active_learning.vectorizer.hyperparameters must be a dictionary")
+
+ # Validate training parameters
+ if "min_annotations_per_instance" in al_config:
+ min_ann = al_config["min_annotations_per_instance"]
+ if not isinstance(min_ann, int) or min_ann < 1:
+ raise ConfigValidationError("active_learning.min_annotations_per_instance must be a positive integer")
+
+ if "min_instances_for_training" in al_config:
+ min_inst = al_config["min_instances_for_training"]
+ if not isinstance(min_inst, int) or min_inst < 2:
+ raise ConfigValidationError("active_learning.min_instances_for_training must be an integer >= 2")
+
+ if "max_instances_to_reorder" in al_config:
+ max_inst = al_config["max_instances_to_reorder"]
+ if not isinstance(max_inst, int) or max_inst < 1:
+ raise ConfigValidationError("active_learning.max_instances_to_reorder must be a positive integer")
+
+ if "update_frequency" in al_config:
+ update_freq = al_config["update_frequency"]
+ if not isinstance(update_freq, int) or update_freq < 1:
+ raise ConfigValidationError("active_learning.update_frequency must be a positive integer")
+
+ # Validate resolution strategy
+ if "resolution_strategy" in al_config:
+ strategy = al_config["resolution_strategy"]
+ valid_strategies = ["majority_vote", "random", "consensus", "weighted_average"]
+ if strategy not in valid_strategies:
+ raise ConfigValidationError(f"active_learning.resolution_strategy must be one of: {', '.join(valid_strategies)}")
+
+ # Validate random sample percent
+ if "random_sample_percent" in al_config:
+ random_pct = al_config["random_sample_percent"]
+ if not isinstance(random_pct, (int, float)) or random_pct < 0 or random_pct > 1:
+ raise ConfigValidationError("active_learning.random_sample_percent must be between 0 and 1")
+
+ # Validate schema names
+ if "schema_names" in al_config:
+ schema_names = al_config["schema_names"]
+ if not isinstance(schema_names, list):
+ raise ConfigValidationError("active_learning.schema_names must be a list")
+
+ for schema in schema_names:
+ if not isinstance(schema, str):
+ raise ConfigValidationError("active_learning.schema_names must contain only strings")
+
+ # Check for unsupported schema types
+ if schema in ["text", "span"]:
+ raise ConfigValidationError(f"Text and span annotation schemes are not supported for active learning: {schema}")
+
+ # Validate database configuration
+ if "database" in al_config:
+ db_config = al_config["database"]
+ if not isinstance(db_config, dict):
+ raise ConfigValidationError("active_learning.database must be a dictionary")
+
+ if "enabled" in db_config and not isinstance(db_config["enabled"], bool):
+ raise ConfigValidationError("active_learning.database.enabled must be a boolean")
+
+ # Validate model persistence configuration
+ if "model_persistence" in al_config:
+ model_config = al_config["model_persistence"]
+ if not isinstance(model_config, dict):
+ raise ConfigValidationError("active_learning.model_persistence must be a dictionary")
+
+ if "enabled" in model_config and not isinstance(model_config["enabled"], bool):
+ raise ConfigValidationError("active_learning.model_persistence.enabled must be a boolean")
+
+ if "retention_count" in model_config:
+ retention = model_config["retention_count"]
+ if not isinstance(retention, int) or retention < 1:
+ raise ConfigValidationError("active_learning.model_persistence.retention_count must be a positive integer")
+
+ # Validate LLM configuration
+ if "llm" in al_config:
+ llm_config = al_config["llm"]
+ if not isinstance(llm_config, dict):
+ raise ConfigValidationError("active_learning.llm must be a dictionary")
+
+ if "enabled" in llm_config and not isinstance(llm_config["enabled"], bool):
+ raise ConfigValidationError("active_learning.llm.enabled must be a boolean")
+
+ if "endpoint_url" in llm_config and not isinstance(llm_config["endpoint_url"], str):
+ raise ConfigValidationError("active_learning.llm.endpoint_url must be a string")
+
+ if "model_name" in llm_config and not isinstance(llm_config["model_name"], str):
+ raise ConfigValidationError("active_learning.llm.model_name must be a string")
+
+ # Validate query strategy
+ if "query_strategy" in al_config:
+ strategy = al_config["query_strategy"]
+ valid_strategies = ["uncertainty", "diversity", "badge", "bald", "hybrid"]
+ if strategy not in valid_strategies:
+ raise ConfigValidationError(
+ f"active_learning.query_strategy must be one of: {', '.join(valid_strategies)}"
+ )
+
+ # Validate hybrid weights
+ if "hybrid_weights" in al_config:
+ weights = al_config["hybrid_weights"]
+ if not isinstance(weights, dict):
+ raise ConfigValidationError("active_learning.hybrid_weights must be a dictionary")
+ weight_sum = sum(weights.values())
+ if abs(weight_sum - 1.0) > 0.01:
+ raise ConfigValidationError(
+ f"active_learning.hybrid_weights must sum to 1.0 (got {weight_sum})"
+ )
+
+ # Validate cold-start strategy
+ if "cold_start_strategy" in al_config:
+ cs = al_config["cold_start_strategy"]
+ if cs not in ["random", "llm"]:
+ raise ConfigValidationError(
+ "active_learning.cold_start_strategy must be one of: random, llm"
+ )
+
+ # Validate confidence method (for LLM active learning)
+ if "confidence_method" in al_config:
+ cm = al_config["confidence_method"]
+ if cm not in ["logprobs", "verbalized", "consistency"]:
+ raise ConfigValidationError(
+ "active_learning.confidence_method must be one of: logprobs, verbalized, consistency"
+ )
+
+ # Validate classifier_params and vectorizer_params
+ if "classifier_params" in al_config:
+ if not isinstance(al_config["classifier_params"], dict):
+ raise ConfigValidationError("active_learning.classifier_params must be a dictionary")
+
+ if "vectorizer_params" in al_config:
+ if not isinstance(al_config["vectorizer_params"], dict):
+ raise ConfigValidationError("active_learning.vectorizer_params must be a dictionary")
+
+ # Validate calibrate_probabilities
+ if "calibrate_probabilities" in al_config:
+ if not isinstance(al_config["calibrate_probabilities"], bool):
+ raise ConfigValidationError("active_learning.calibrate_probabilities must be a boolean")
+
+ # Validate BALD params
+ if "bald_params" in al_config:
+ bp = al_config["bald_params"]
+ if not isinstance(bp, dict):
+ raise ConfigValidationError("active_learning.bald_params must be a dictionary")
+ if "n_estimators" in bp:
+ if not isinstance(bp["n_estimators"], int) or bp["n_estimators"] < 2:
+ raise ConfigValidationError("active_learning.bald_params.n_estimators must be an integer >= 2")
+
+ # Validate ICL ensemble params
+ if "use_icl_ensemble" in al_config:
+ if not isinstance(al_config["use_icl_ensemble"], bool):
+ raise ConfigValidationError("active_learning.use_icl_ensemble must be a boolean")
+
+ if "icl_ensemble_params" in al_config:
+ if not isinstance(al_config["icl_ensemble_params"], dict):
+ raise ConfigValidationError("active_learning.icl_ensemble_params must be a dictionary")
+
+ # Validate annotation routing
+ if "annotation_routing" in al_config:
+ if not isinstance(al_config["annotation_routing"], bool):
+ raise ConfigValidationError("active_learning.annotation_routing must be a boolean")
+
+ if "routing_thresholds" in al_config:
+ rt = al_config["routing_thresholds"]
+ if not isinstance(rt, dict):
+ raise ConfigValidationError("active_learning.routing_thresholds must be a dictionary")
+ for key in ["auto_label_min_confidence", "show_suggestion_below"]:
+ if key in rt:
+ val = rt[key]
+ if not isinstance(val, (int, float)) or val < 0 or val > 1:
+ raise ConfigValidationError(
+ f"active_learning.routing_thresholds.{key} must be between 0 and 1"
+ )
+
+ # Warn about sentence-transformers dependency
+ if al_config.get("vectorizer_name") == "sentence-transformers" or \
+ (isinstance(al_config.get("vectorizer"), dict) and
+ al_config["vectorizer"].get("name") == "sentence-transformers"):
+ try:
+ import sentence_transformers # noqa: F401
+ except ImportError:
+ logger.warning(
+ "sentence-transformers vectorizer configured but package not installed. "
+ "Install with: pip install sentence-transformers"
+ )
+
+
+def validate_ai_support_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate AI support configuration.
+
+ Args:
+ config_data: The configuration data containing ai_support section
+
+ Raises:
+ ConfigValidationError: If the AI support configuration is invalid
+ """
+ if "ai_support" not in config_data:
+ return # AI support is optional
+
+ ai_config = config_data["ai_support"]
+
+ # Validate enabled flag
+ if not isinstance(ai_config.get("enabled", False), bool):
+ raise ConfigValidationError("ai_support.enabled must be a boolean")
+
+ if not ai_config.get("enabled", False):
+ return # Skip validation if not enabled
+
+ # Validate ai_config_file (optional, string path to external AI config)
+ has_external_config = False
+ if "ai_config_file" in ai_config:
+ if not isinstance(ai_config["ai_config_file"], str):
+ raise ConfigValidationError("ai_support.ai_config_file must be a string")
+ has_external_config = True
+
+ # Validate endpoint type. When ai_config_file is set, the endpoint_type is
+ # expected to live in the external file (which may be gitignored, e.g. when
+ # it holds API keys) and is loaded at server start.
+ if "endpoint_type" not in ai_config:
+ if has_external_config:
+ return # External file provides endpoint_type + model + credentials
+ raise ConfigValidationError("ai_support.endpoint_type is required when ai_support is enabled")
+
+ endpoint_type = ai_config["endpoint_type"]
+ if not isinstance(endpoint_type, str):
+ raise ConfigValidationError("ai_support.endpoint_type must be a string")
+
+ valid_endpoint_types = ["openai", "anthropic", "huggingface", "ollama", "gemini", "vllm",
+ "yolo", "ollama_vision", "openai_vision", "anthropic_vision"]
+ if endpoint_type not in valid_endpoint_types:
+ raise ConfigValidationError(f"ai_support.endpoint_type must be one of: {', '.join(valid_endpoint_types)}")
+
+ # Validate ai_config section
+ if "ai_config" in ai_config:
+ ai_endpoint_config = ai_config["ai_config"]
+ if not isinstance(ai_endpoint_config, dict):
+ raise ConfigValidationError("ai_support.ai_config must be a dictionary")
+
+ # Validate model name
+ if "model" in ai_endpoint_config:
+ model = ai_endpoint_config["model"]
+ if not isinstance(model, str) or not model.strip():
+ raise ConfigValidationError("ai_support.ai_config.model must be a non-empty string")
+
+ # Validate API key for cloud-based endpoints
+ if endpoint_type in ["openai", "anthropic", "huggingface", "gemini"]:
+ api_key = ai_endpoint_config.get("api_key", "")
+ if not api_key or not isinstance(api_key, str):
+ raise ConfigValidationError(f"ai_support.ai_config.api_key is required for {endpoint_type} endpoint")
+
+ # Validate base_url for VLLM
+ if endpoint_type == "vllm":
+ base_url = ai_endpoint_config.get("base_url", "")
+ if base_url and not isinstance(base_url, str):
+ raise ConfigValidationError("ai_support.ai_config.base_url must be a string")
+
+ # Validate temperature
+ if "temperature" in ai_endpoint_config:
+ temperature = ai_endpoint_config["temperature"]
+ if not isinstance(temperature, (int, float)) or temperature < 0 or temperature > 2:
+ raise ConfigValidationError("ai_support.ai_config.temperature must be between 0 and 2")
+
+ # Validate max_tokens
+ if "max_tokens" in ai_endpoint_config:
+ max_tokens = ai_endpoint_config["max_tokens"]
+ if not isinstance(max_tokens, int) or max_tokens < 1:
+ raise ConfigValidationError("ai_support.ai_config.max_tokens must be a positive integer")
+
+ # Validate custom prompts
+ for prompt_key in ["hint_prompt", "keyword_prompt"]:
+ if prompt_key in ai_endpoint_config:
+ prompt = ai_endpoint_config[prompt_key]
+ if not isinstance(prompt, str):
+ raise ConfigValidationError(f"ai_support.ai_config.{prompt_key} must be a string")
+ if not prompt.strip():
+ raise ConfigValidationError(f"ai_support.ai_config.{prompt_key} cannot be empty")
+
+ # Validate option_highlighting configuration
+ if "option_highlighting" in ai_config:
+ _validate_option_highlighting_config(ai_config["option_highlighting"])
+
+
+def validate_chat_support_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate chat support configuration for LLM annotator assistance.
+
+ Args:
+ config_data: The configuration data containing chat_support section
+
+ Raises:
+ ConfigValidationError: If the chat support configuration is invalid
+ """
+ if "chat_support" not in config_data:
+ return # Chat support is optional
+
+ chat_config = config_data["chat_support"]
+
+ if not isinstance(chat_config.get("enabled", False), bool):
+ raise ConfigValidationError("chat_support.enabled must be a boolean")
+
+ if not chat_config.get("enabled", False):
+ return # Skip validation if not enabled
+
+ # Validate endpoint type
+ if "endpoint_type" not in chat_config:
+ raise ConfigValidationError(
+ "chat_support.endpoint_type is required when chat_support is enabled"
+ )
+
+ endpoint_type = chat_config["endpoint_type"]
+ valid_endpoint_types = [
+ "openai", "anthropic", "huggingface", "ollama", "gemini", "vllm", "openrouter",
+ ]
+ if endpoint_type not in valid_endpoint_types:
+ raise ConfigValidationError(
+ f"chat_support.endpoint_type must be one of: {', '.join(valid_endpoint_types)}"
+ )
+
+ # Validate ai_config section
+ if "ai_config" in chat_config:
+ ai_cfg = chat_config["ai_config"]
+ if not isinstance(ai_cfg, dict):
+ raise ConfigValidationError("chat_support.ai_config must be a dictionary")
+
+ if "model" in ai_cfg:
+ if not isinstance(ai_cfg["model"], str) or not ai_cfg["model"].strip():
+ raise ConfigValidationError(
+ "chat_support.ai_config.model must be a non-empty string"
+ )
+
+ if "temperature" in ai_cfg:
+ temp = ai_cfg["temperature"]
+ if not isinstance(temp, (int, float)) or temp < 0 or temp > 2:
+ raise ConfigValidationError(
+ "chat_support.ai_config.temperature must be between 0 and 2"
+ )
+
+ if "max_tokens" in ai_cfg:
+ mt = ai_cfg["max_tokens"]
+ if not isinstance(mt, int) or mt < 1:
+ raise ConfigValidationError(
+ "chat_support.ai_config.max_tokens must be a positive integer"
+ )
+
+ # Validate API key for cloud endpoints
+ if endpoint_type in ["openai", "anthropic", "huggingface", "gemini", "openrouter"]:
+ api_key = ai_cfg.get("api_key", "")
+ if not api_key or not isinstance(api_key, str):
+ raise ConfigValidationError(
+ f"chat_support.ai_config.api_key is required for {endpoint_type} endpoint"
+ )
+
+ # Validate UI section
+ if "ui" in chat_config:
+ ui_cfg = chat_config["ui"]
+ if not isinstance(ui_cfg, dict):
+ raise ConfigValidationError("chat_support.ui must be a dictionary")
+
+ if "sidebar_width" in ui_cfg:
+ sw = ui_cfg["sidebar_width"]
+ if not isinstance(sw, int) or sw < 200 or sw > 800:
+ raise ConfigValidationError(
+ "chat_support.ui.sidebar_width must be an integer between 200 and 800"
+ )
+
+ if "max_history_per_instance" in ui_cfg:
+ mh = ui_cfg["max_history_per_instance"]
+ if not isinstance(mh, int) or mh < 1:
+ raise ConfigValidationError(
+ "chat_support.ui.max_history_per_instance must be a positive integer"
+ )
+
+
+def _validate_option_highlighting_config(oh_config: Dict[str, Any]) -> None:
+ """
+ Validate option highlighting configuration.
+
+ Args:
+ oh_config: The option_highlighting configuration section
+
+ Raises:
+ ConfigValidationError: If the configuration is invalid
+ """
+ if not isinstance(oh_config, dict):
+ raise ConfigValidationError("ai_support.option_highlighting must be a dictionary")
+
+ # Validate enabled flag
+ if "enabled" in oh_config:
+ if not isinstance(oh_config["enabled"], bool):
+ raise ConfigValidationError("ai_support.option_highlighting.enabled must be a boolean")
+
+ # Validate top_k (number of options to highlight)
+ if "top_k" in oh_config:
+ top_k = oh_config["top_k"]
+ if not isinstance(top_k, int) or top_k < 1 or top_k > 10:
+ raise ConfigValidationError("ai_support.option_highlighting.top_k must be an integer between 1 and 10")
+
+ # Validate dim_opacity (opacity for non-highlighted options)
+ if "dim_opacity" in oh_config:
+ dim_opacity = oh_config["dim_opacity"]
+ if not isinstance(dim_opacity, (int, float)) or dim_opacity < 0.1 or dim_opacity > 0.9:
+ raise ConfigValidationError("ai_support.option_highlighting.dim_opacity must be a number between 0.1 and 0.9")
+
+ # Validate auto_apply flag
+ if "auto_apply" in oh_config:
+ if not isinstance(oh_config["auto_apply"], bool):
+ raise ConfigValidationError("ai_support.option_highlighting.auto_apply must be a boolean")
+
+ # Validate schemas filter (list of schema names or null)
+ if "schemas" in oh_config:
+ schemas = oh_config["schemas"]
+ if schemas is not None:
+ if not isinstance(schemas, list):
+ raise ConfigValidationError("ai_support.option_highlighting.schemas must be a list or null")
+ for schema in schemas:
+ if not isinstance(schema, str):
+ raise ConfigValidationError("ai_support.option_highlighting.schemas must contain only strings")
+
+ # Validate prefetch_count
+ if "prefetch_count" in oh_config:
+ prefetch_count = oh_config["prefetch_count"]
+ if not isinstance(prefetch_count, int) or prefetch_count < 0 or prefetch_count > 100:
+ raise ConfigValidationError("ai_support.option_highlighting.prefetch_count must be an integer between 0 and 100")
+
+
+def parse_active_learning_config(config_data: Dict[str, Any]) -> 'ActiveLearningConfig':
+ """
+ Parse active learning configuration from YAML data.
+
+ Args:
+ config_data: The configuration data containing active_learning section
+
+ Returns:
+ ActiveLearningConfig: Parsed active learning configuration
+
+ Raises:
+ ConfigValidationError: If the configuration is invalid
+ """
+ from potato.active_learning_manager import ActiveLearningConfig, ResolutionStrategy
+
+ if "active_learning" not in config_data:
+ return ActiveLearningConfig() # Return default config
+
+ al_config = config_data["active_learning"]
+
+ # Parse classifier configuration
+ classifier_name = "sklearn.linear_model.LogisticRegression"
+ classifier_kwargs = {}
+ if "classifier" in al_config:
+ classifier_config = al_config["classifier"]
+ classifier_name = classifier_config.get("name", classifier_name)
+ classifier_kwargs = classifier_config.get("hyperparameters", {})
+
+ # Parse vectorizer configuration
+ vectorizer_name = "sklearn.feature_extraction.text.CountVectorizer"
+ vectorizer_kwargs = {}
+ if "vectorizer" in al_config:
+ vectorizer_config = al_config["vectorizer"]
+ vectorizer_name = vectorizer_config.get("name", vectorizer_name)
+ vectorizer_kwargs = vectorizer_config.get("hyperparameters", {})
+
+ # Parse resolution strategy
+ resolution_strategy = ResolutionStrategy.MAJORITY_VOTE
+ if "resolution_strategy" in al_config:
+ strategy_str = al_config["resolution_strategy"]
+ if strategy_str == "majority_vote":
+ resolution_strategy = ResolutionStrategy.MAJORITY_VOTE
+ elif strategy_str == "random":
+ resolution_strategy = ResolutionStrategy.RANDOM
+ elif strategy_str == "consensus":
+ resolution_strategy = ResolutionStrategy.CONSENSUS
+ elif strategy_str == "weighted_average":
+ resolution_strategy = ResolutionStrategy.WEIGHTED_AVERAGE
+
+ # Parse other parameters
+ min_annotations_per_instance = al_config.get("min_annotations_per_instance", 1)
+ min_instances_for_training = al_config.get("min_instances_for_training", 10)
+ max_instances_to_reorder = al_config.get("max_instances_to_reorder")
+ random_sample_percent = al_config.get("random_sample_percent", 0.2)
+ update_frequency = al_config.get("update_frequency", 5)
+ schema_names = al_config.get("schema_names", [])
+
+ # Parse database configuration
+ database_enabled = False
+ database_config = {}
+ if "database" in al_config:
+ db_config = al_config["database"]
+ database_enabled = db_config.get("enabled", False)
+ database_config = {k: v for k, v in db_config.items() if k != "enabled"}
+
+ # Parse model persistence configuration
+ model_persistence_enabled = False
+ model_save_directory = None
+ model_retention_count = 2
+ if "model_persistence" in al_config:
+ model_config = al_config["model_persistence"]
+ model_persistence_enabled = model_config.get("enabled", False)
+ model_save_directory = model_config.get("save_directory")
+ model_retention_count = model_config.get("retention_count", 2)
+
+ # Parse LLM configuration
+ llm_enabled = False
+ llm_config = {}
+ if "llm" in al_config:
+ llm_config = al_config["llm"]
+ llm_enabled = llm_config.get("enabled", False)
+
+ return ActiveLearningConfig(
+ enabled=al_config.get("enabled", False),
+ classifier_name=classifier_name,
+ classifier_kwargs=classifier_kwargs,
+ vectorizer_name=vectorizer_name,
+ vectorizer_kwargs=vectorizer_kwargs,
+ min_annotations_per_instance=min_annotations_per_instance,
+ min_instances_for_training=min_instances_for_training,
+ max_instances_to_reorder=max_instances_to_reorder,
+ resolution_strategy=resolution_strategy,
+ random_sample_percent=random_sample_percent,
+ update_frequency=update_frequency,
+ schema_names=schema_names,
+ database_enabled=database_enabled,
+ database_config=database_config,
+ model_persistence_enabled=model_persistence_enabled,
+ model_save_directory=model_save_directory,
+ model_retention_count=model_retention_count,
+ llm_enabled=llm_enabled,
+ llm_config=llm_config
+ )
+
+
+def validate_instance_display_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate instance_display configuration.
+
+ The instance_display section defines what content to show annotators,
+ separate from what annotations to collect. This allows displaying
+ images/videos/audio alongside any annotation type.
+
+ Args:
+ config_data: The configuration data
+
+ Raises:
+ ConfigValidationError: If the instance_display configuration is invalid
+ """
+ if "instance_display" not in config_data:
+ return # instance_display is optional (backwards compatible)
+
+ display_config = config_data["instance_display"]
+
+ if not isinstance(display_config, dict):
+ raise ConfigValidationError("instance_display must be a dictionary")
+
+ # Validate fields
+ if "fields" not in display_config:
+ raise ConfigValidationError("instance_display must contain 'fields' list")
+
+ fields = display_config["fields"]
+ if not isinstance(fields, list):
+ raise ConfigValidationError("instance_display.fields must be a list")
+
+ if not fields:
+ raise ConfigValidationError("instance_display.fields cannot be empty")
+
+ # Track span targets for validation
+ span_targets = []
+
+ # Valid display types โ sourced from the display registry (single source
+ # of truth) so new display types don't require editing this list. Falls
+ # back to a static list if the registry can't be imported.
+ try:
+ from .displays import display_registry
+ valid_display_types = display_registry.get_supported_types()
+ except Exception:
+ valid_display_types = [
+ "text", "html", "image", "video", "audio", "dialogue", "pairwise",
+ "pdf", "document", "spreadsheet", "code", "agent_trace", "eval_trace",
+ "gallery", "conversation_tree", "interactive_chat", "web_agent_trace",
+ "live_agent", "coding_trace", "live_coding_agent",
+ ]
+
+ for i, field in enumerate(fields):
+ if not isinstance(field, dict):
+ raise ConfigValidationError(f"instance_display.fields[{i}] must be a dictionary")
+
+ # Validate required field properties
+ if "key" not in field:
+ raise ConfigValidationError(f"instance_display.fields[{i}] missing required 'key' property")
+
+ key = field["key"]
+ if not isinstance(key, str) or not key.strip():
+ raise ConfigValidationError(f"instance_display.fields[{i}].key must be a non-empty string")
+
+ if "type" not in field:
+ raise ConfigValidationError(f"instance_display.fields[{i}] missing required 'type' property")
+
+ field_type = field["type"]
+ if field_type not in valid_display_types:
+ raise ConfigValidationError(
+ f"instance_display.fields[{i}].type '{field_type}' is invalid. "
+ f"Valid types are: {', '.join(valid_display_types)}"
+ )
+
+ # Validate label if present
+ if "label" in field:
+ if not isinstance(field["label"], str):
+ raise ConfigValidationError(f"instance_display.fields[{i}].label must be a string")
+
+ # Validate span_target
+ if field.get("span_target"):
+ # Types that support span annotation targets
+ span_target_types = ["text", "dialogue", "pdf", "document", "spreadsheet", "code", "agent_trace", "interactive_chat"]
+ if field_type not in span_target_types:
+ raise ConfigValidationError(
+ f"instance_display.fields[{i}].span_target is set but type '{field_type}' "
+ f"does not support span annotation. Types that support span_target: {', '.join(span_target_types)}."
+ )
+ span_targets.append(key)
+
+ # Validate display_options if present
+ if "display_options" in field:
+ options = field["display_options"]
+ if not isinstance(options, dict):
+ raise ConfigValidationError(f"instance_display.fields[{i}].display_options must be a dictionary")
+
+ # Type-specific option validation
+ _validate_display_options(field_type, options, f"instance_display.fields[{i}]")
+
+ # Validate layout if present
+ if "layout" in display_config:
+ layout = display_config["layout"]
+ if not isinstance(layout, dict):
+ raise ConfigValidationError("instance_display.layout must be a dictionary")
+
+ if "direction" in layout:
+ valid_directions = ["vertical", "horizontal"]
+ if layout["direction"] not in valid_directions:
+ raise ConfigValidationError(
+ f"instance_display.layout.direction must be one of: {', '.join(valid_directions)}"
+ )
+
+ if "gap" in layout:
+ gap = layout["gap"]
+ if not isinstance(gap, str):
+ raise ConfigValidationError("instance_display.layout.gap must be a string (e.g., '20px', '1rem')")
+
+ # Validate resizable option (defaults to True)
+ if "resizable" in display_config:
+ if not isinstance(display_config["resizable"], bool):
+ raise ConfigValidationError("instance_display.resizable must be a boolean (true/false)")
+
+ # Check for deprecation warning: using annotation schemas for display-only
+ _check_display_only_deprecation(config_data)
+
+
+def _validate_display_options(field_type: str, options: Dict[str, Any], path: str) -> None:
+ """
+ Validate display options for a specific field type.
+
+ Args:
+ field_type: The display type
+ options: The display options dictionary
+ path: The config path for error messages
+
+ Raises:
+ ConfigValidationError: If options are invalid
+ """
+ # Common option validation
+ if "max_width" in options:
+ max_width = options["max_width"]
+ if not isinstance(max_width, (int, str)):
+ raise ConfigValidationError(f"{path}.display_options.max_width must be an integer or string")
+ if isinstance(max_width, int) and max_width < 1:
+ raise ConfigValidationError(f"{path}.display_options.max_width must be positive")
+
+ if "max_height" in options:
+ max_height = options["max_height"]
+ if not isinstance(max_height, (int, str)):
+ raise ConfigValidationError(f"{path}.display_options.max_height must be an integer or string")
+ if isinstance(max_height, int) and max_height < 1:
+ raise ConfigValidationError(f"{path}.display_options.max_height must be positive")
+
+ if "min_height" in options:
+ min_height = options["min_height"]
+ if not isinstance(min_height, (int, str)):
+ raise ConfigValidationError(f"{path}.display_options.min_height must be an integer or string")
+ if isinstance(min_height, int) and min_height < 1:
+ raise ConfigValidationError(f"{path}.display_options.min_height must be positive")
+
+ if "resizable" in options:
+ if not isinstance(options["resizable"], bool):
+ raise ConfigValidationError(f"{path}.display_options.resizable must be a boolean")
+
+ # Text-specific options
+ if field_type in ["text", "html"]:
+ if "collapsible" in options:
+ if not isinstance(options["collapsible"], bool):
+ raise ConfigValidationError(f"{path}.display_options.collapsible must be a boolean")
+
+ if "preserve_whitespace" in options:
+ if not isinstance(options["preserve_whitespace"], bool):
+ raise ConfigValidationError(f"{path}.display_options.preserve_whitespace must be a boolean")
+
+ # Image-specific options
+ if field_type == "image":
+ if "zoomable" in options:
+ if not isinstance(options["zoomable"], bool):
+ raise ConfigValidationError(f"{path}.display_options.zoomable must be a boolean")
+
+ if "object_fit" in options:
+ valid_fits = ["contain", "cover", "fill", "none", "scale-down"]
+ if options["object_fit"] not in valid_fits:
+ raise ConfigValidationError(
+ f"{path}.display_options.object_fit must be one of: {', '.join(valid_fits)}"
+ )
+
+ # Video-specific options
+ if field_type == "video":
+ for bool_opt in ["controls", "autoplay", "loop", "muted"]:
+ if bool_opt in options:
+ if not isinstance(options[bool_opt], bool):
+ raise ConfigValidationError(f"{path}.display_options.{bool_opt} must be a boolean")
+
+ # Audio-specific options
+ if field_type == "audio":
+ if "controls" in options:
+ if not isinstance(options["controls"], bool):
+ raise ConfigValidationError(f"{path}.display_options.controls must be a boolean")
+
+ if "show_waveform" in options:
+ if not isinstance(options["show_waveform"], bool):
+ raise ConfigValidationError(f"{path}.display_options.show_waveform must be a boolean")
+
+ # Dialogue-specific options
+ if field_type == "dialogue":
+ if "alternating_shading" in options:
+ if not isinstance(options["alternating_shading"], bool):
+ raise ConfigValidationError(f"{path}.display_options.alternating_shading must be a boolean")
+
+ if "speaker_extraction" in options:
+ if not isinstance(options["speaker_extraction"], bool):
+ raise ConfigValidationError(f"{path}.display_options.speaker_extraction must be a boolean")
+
+ # Pairwise-specific options
+ if field_type == "pairwise":
+ if "cell_width" in options:
+ cell_width = options["cell_width"]
+ if not isinstance(cell_width, str):
+ raise ConfigValidationError(f"{path}.display_options.cell_width must be a string (e.g., '50%')")
+
+ # PDF-specific options
+ if field_type == "pdf":
+ if "view_mode" in options:
+ valid_modes = ["scroll", "paginated", "side-by-side"]
+ if options["view_mode"] not in valid_modes:
+ raise ConfigValidationError(
+ f"{path}.display_options.view_mode must be one of: {', '.join(valid_modes)}"
+ )
+
+ if "text_layer" in options:
+ if not isinstance(options["text_layer"], bool):
+ raise ConfigValidationError(f"{path}.display_options.text_layer must be a boolean")
+
+ if "zoom" in options:
+ zoom = options["zoom"]
+ valid_zoom_modes = ["auto", "page-fit", "page-width"]
+ if zoom not in valid_zoom_modes:
+ try:
+ float(zoom)
+ except (TypeError, ValueError):
+ raise ConfigValidationError(
+ f"{path}.display_options.zoom must be one of {valid_zoom_modes} or a number"
+ )
+
+ # Document-specific options
+ if field_type == "document":
+ if "collapsible" in options:
+ if not isinstance(options["collapsible"], bool):
+ raise ConfigValidationError(f"{path}.display_options.collapsible must be a boolean")
+
+ if "show_outline" in options:
+ if not isinstance(options["show_outline"], bool):
+ raise ConfigValidationError(f"{path}.display_options.show_outline must be a boolean")
+
+ if "style_theme" in options:
+ valid_themes = ["default", "minimal", "print"]
+ if options["style_theme"] not in valid_themes:
+ raise ConfigValidationError(
+ f"{path}.display_options.style_theme must be one of: {', '.join(valid_themes)}"
+ )
+
+ # Spreadsheet-specific options
+ if field_type == "spreadsheet":
+ if "annotation_mode" in options:
+ valid_modes = ["row", "cell", "range"]
+ if options["annotation_mode"] not in valid_modes:
+ raise ConfigValidationError(
+ f"{path}.display_options.annotation_mode must be one of: {', '.join(valid_modes)}"
+ )
+
+ for bool_opt in ["show_headers", "striped", "hoverable", "sortable", "selectable", "compact"]:
+ if bool_opt in options:
+ if not isinstance(options[bool_opt], bool):
+ raise ConfigValidationError(f"{path}.display_options.{bool_opt} must be a boolean")
+
+ # Code-specific options
+ if field_type == "code":
+ if "language" in options:
+ if not isinstance(options["language"], (str, type(None))):
+ raise ConfigValidationError(f"{path}.display_options.language must be a string or null")
+
+ if "show_line_numbers" in options:
+ if not isinstance(options["show_line_numbers"], bool):
+ raise ConfigValidationError(f"{path}.display_options.show_line_numbers must be a boolean")
+
+ if "wrap_lines" in options:
+ if not isinstance(options["wrap_lines"], bool):
+ raise ConfigValidationError(f"{path}.display_options.wrap_lines must be a boolean")
+
+ if "highlight_lines" in options:
+ hl = options["highlight_lines"]
+ if hl is not None and not isinstance(hl, list):
+ raise ConfigValidationError(f"{path}.display_options.highlight_lines must be a list of line numbers or null")
+
+ if "theme" in options:
+ valid_themes = ["default", "dark"]
+ if options["theme"] not in valid_themes:
+ raise ConfigValidationError(
+ f"{path}.display_options.theme must be one of: {', '.join(valid_themes)}"
+ )
+
+
+def validate_format_handling_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate format_handling configuration for extended format support.
+
+ Args:
+ config_data: The full configuration data
+
+ Raises:
+ ConfigValidationError: If the format_handling configuration is invalid
+ """
+ format_config = config_data.get('format_handling')
+ if format_config is None:
+ return
+
+ if not isinstance(format_config, dict):
+ raise ConfigValidationError("format_handling must be a dictionary")
+
+ # Validate enabled flag
+ if "enabled" in format_config:
+ if not isinstance(format_config["enabled"], bool):
+ raise ConfigValidationError("format_handling.enabled must be a boolean")
+
+ # Validate default_format
+ if "default_format" in format_config:
+ default = format_config["default_format"]
+ valid_defaults = ["auto", "pdf", "docx", "markdown", "spreadsheet", "code"]
+ if default not in valid_defaults:
+ raise ConfigValidationError(
+ f"format_handling.default_format must be one of: {', '.join(valid_defaults)}"
+ )
+
+ # Validate PDF-specific options
+ if "pdf" in format_config:
+ pdf_opts = format_config["pdf"]
+ if not isinstance(pdf_opts, dict):
+ raise ConfigValidationError("format_handling.pdf must be a dictionary")
+
+ if "extraction_mode" in pdf_opts:
+ valid_modes = ["text", "ocr", "hybrid"]
+ if pdf_opts["extraction_mode"] not in valid_modes:
+ raise ConfigValidationError(
+ f"format_handling.pdf.extraction_mode must be one of: {', '.join(valid_modes)}"
+ )
+
+ if "cache_extracted" in pdf_opts:
+ if not isinstance(pdf_opts["cache_extracted"], bool):
+ raise ConfigValidationError("format_handling.pdf.cache_extracted must be a boolean")
+
+ # Validate spreadsheet-specific options
+ if "spreadsheet" in format_config:
+ ss_opts = format_config["spreadsheet"]
+ if not isinstance(ss_opts, dict):
+ raise ConfigValidationError("format_handling.spreadsheet must be a dictionary")
+
+ if "annotation_mode" in ss_opts:
+ valid_modes = ["row", "cell", "range"]
+ if ss_opts["annotation_mode"] not in valid_modes:
+ raise ConfigValidationError(
+ f"format_handling.spreadsheet.annotation_mode must be one of: {', '.join(valid_modes)}"
+ )
+
+ if "max_rows" in ss_opts:
+ max_rows = ss_opts["max_rows"]
+ if not isinstance(max_rows, int) or max_rows < 1:
+ raise ConfigValidationError("format_handling.spreadsheet.max_rows must be a positive integer")
+
+
+def validate_layout_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate layout configuration for annotation form grid arrangement.
+
+ The layout section configures how annotation forms are arranged in a grid,
+ supports grouping schemas with collapsible headers, and provides responsive
+ breakpoints for mobile/tablet displays.
+
+ Args:
+ config_data: The full configuration data
+
+ Raises:
+ ConfigValidationError: If the layout configuration is invalid
+ """
+ layout = config_data.get('layout')
+ if layout is None:
+ return # layout is optional
+
+ if not isinstance(layout, dict):
+ raise ConfigValidationError("layout must be a dictionary")
+
+ # Validate grid configuration
+ if 'grid' in layout:
+ grid = layout['grid']
+ if not isinstance(grid, dict):
+ raise ConfigValidationError("layout.grid must be a dictionary")
+
+ # Validate columns (1-6)
+ if 'columns' in grid:
+ columns = grid['columns']
+ if not isinstance(columns, int) or columns < 1 or columns > 6:
+ raise ConfigValidationError("layout.grid.columns must be an integer between 1 and 6")
+
+ # Validate gap (CSS value)
+ if 'gap' in grid:
+ gap = grid['gap']
+ if not isinstance(gap, str) or not gap.strip():
+ raise ConfigValidationError("layout.grid.gap must be a non-empty CSS value string (e.g., '1rem', '16px')")
+
+ # Validate row_gap (CSS value)
+ if 'row_gap' in grid:
+ row_gap = grid['row_gap']
+ if not isinstance(row_gap, str) or not row_gap.strip():
+ raise ConfigValidationError("layout.grid.row_gap must be a non-empty CSS value string")
+
+ # Validate align_items
+ if 'align_items' in grid:
+ valid_alignments = ['start', 'center', 'end', 'stretch']
+ if grid['align_items'] not in valid_alignments:
+ raise ConfigValidationError(
+ f"layout.grid.align_items must be one of: {', '.join(valid_alignments)}"
+ )
+
+ # Validate breakpoints
+ if 'breakpoints' in layout:
+ breakpoints = layout['breakpoints']
+ if not isinstance(breakpoints, dict):
+ raise ConfigValidationError("layout.breakpoints must be a dictionary")
+
+ for bp_name in ['mobile', 'tablet']:
+ if bp_name in breakpoints:
+ bp_value = breakpoints[bp_name]
+ if not isinstance(bp_value, int) or bp_value < 0:
+ raise ConfigValidationError(
+ f"layout.breakpoints.{bp_name} must be a non-negative integer (pixel value)"
+ )
+
+ # Validate groups
+ if 'groups' in layout:
+ groups = layout['groups']
+ if not isinstance(groups, list):
+ raise ConfigValidationError("layout.groups must be a list")
+
+ # Collect all schema names for validation
+ all_schemas = set()
+ schemes = config_data.get('annotation_schemes', [])
+ for scheme in schemes:
+ if isinstance(scheme, dict) and 'name' in scheme:
+ all_schemas.add(scheme['name'])
+
+ group_ids = set()
+ for i, group in enumerate(groups):
+ if not isinstance(group, dict):
+ raise ConfigValidationError(f"layout.groups[{i}] must be a dictionary")
+
+ # Validate required group fields
+ if 'id' not in group:
+ raise ConfigValidationError(f"layout.groups[{i}] missing required 'id' field")
+
+ group_id = group['id']
+ if not isinstance(group_id, str) or not group_id.strip():
+ raise ConfigValidationError(f"layout.groups[{i}].id must be a non-empty string")
+
+ if group_id in group_ids:
+ raise ConfigValidationError(f"layout.groups[{i}].id '{group_id}' is duplicate")
+ group_ids.add(group_id)
+
+ # Validate schemas list
+ if 'schemas' not in group:
+ raise ConfigValidationError(f"layout.groups[{i}] missing required 'schemas' field")
+
+ group_schemas = group['schemas']
+ if not isinstance(group_schemas, list):
+ raise ConfigValidationError(f"layout.groups[{i}].schemas must be a list")
+
+ if not group_schemas:
+ raise ConfigValidationError(f"layout.groups[{i}].schemas cannot be empty")
+
+ # Validate each schema reference exists
+ for j, schema_name in enumerate(group_schemas):
+ if not isinstance(schema_name, str):
+ raise ConfigValidationError(
+ f"layout.groups[{i}].schemas[{j}] must be a string"
+ )
+ if schema_name not in all_schemas:
+ raise ConfigValidationError(
+ f"layout.groups[{i}].schemas references unknown schema: '{schema_name}'"
+ )
+
+ # Validate optional boolean fields
+ if 'collapsible' in group:
+ if not isinstance(group['collapsible'], bool):
+ raise ConfigValidationError(f"layout.groups[{i}].collapsible must be a boolean")
+
+ if 'collapsed_default' in group:
+ if not isinstance(group['collapsed_default'], bool):
+ raise ConfigValidationError(f"layout.groups[{i}].collapsed_default must be a boolean")
+
+ # Validate optional title
+ if 'title' in group:
+ if not isinstance(group['title'], str):
+ raise ConfigValidationError(f"layout.groups[{i}].title must be a string")
+
+ # Validate optional description
+ if 'description' in group:
+ if not isinstance(group['description'], str):
+ raise ConfigValidationError(f"layout.groups[{i}].description must be a string")
+
+ # Validate order
+ if 'order' in layout:
+ order = layout['order']
+ if not isinstance(order, list):
+ raise ConfigValidationError("layout.order must be a list")
+
+ for i, schema_name in enumerate(order):
+ if not isinstance(schema_name, str):
+ raise ConfigValidationError(f"layout.order[{i}] must be a string")
+
+ # Validate styling (advanced options)
+ if 'styling' in layout:
+ styling = layout['styling']
+ if not isinstance(styling, dict):
+ raise ConfigValidationError("layout.styling must be a dictionary")
+
+ # Validate align_items
+ if 'align_items' in styling:
+ valid_alignments = ['start', 'center', 'end', 'stretch']
+ if styling['align_items'] not in valid_alignments:
+ raise ConfigValidationError(
+ f"layout.styling.align_items must be one of: {', '.join(valid_alignments)}"
+ )
+
+ # Validate content_align
+ if 'content_align' in styling:
+ valid_content_align = ['left', 'center', 'right']
+ if styling['content_align'] not in valid_content_align:
+ raise ConfigValidationError(
+ f"layout.styling.content_align must be one of: {', '.join(valid_content_align)}"
+ )
+
+ # Validate background colors (CSS color values)
+ for color_key in ['group_background_odd', 'group_background_even']:
+ if color_key in styling:
+ color = styling[color_key]
+ if not isinstance(color, str) or not color.strip():
+ raise ConfigValidationError(
+ f"layout.styling.{color_key} must be a non-empty CSS color value"
+ )
+
+ # Validate padding values (CSS padding)
+ for padding_key in ['group_padding', 'form_padding']:
+ if padding_key in styling:
+ padding = styling[padding_key]
+ if not isinstance(padding, str) or not padding.strip():
+ raise ConfigValidationError(
+ f"layout.styling.{padding_key} must be a non-empty CSS padding value"
+ )
+
+ # Validate per-group background_color if present
+ if 'groups' in layout:
+ for i, group in enumerate(layout['groups']):
+ if 'background_color' in group:
+ bg_color = group['background_color']
+ if not isinstance(bg_color, str) or not bg_color.strip():
+ raise ConfigValidationError(
+ f"layout.groups[{i}].background_color must be a non-empty CSS color value"
+ )
+
+
+def validate_adjudication_config(config_data: Dict[str, Any]) -> None:
+ """
+ Validate adjudication configuration.
+
+ Args:
+ config_data: The full configuration data
+
+ Raises:
+ ConfigValidationError: If the adjudication configuration is invalid
+ """
+ adj_config = config_data.get('adjudication', {})
+ if not isinstance(adj_config, dict):
+ raise ConfigValidationError("adjudication must be a dictionary")
+
+ if not adj_config.get('enabled', False):
+ return
+
+ # Require adjudicator_users
+ users = adj_config.get('adjudicator_users', [])
+ if not isinstance(users, list) or len(users) == 0:
+ raise ConfigValidationError(
+ "adjudication.adjudicator_users must be a non-empty list of usernames"
+ )
+
+ # Validate numeric fields
+ min_ann = adj_config.get('min_annotations', 2)
+ if not isinstance(min_ann, int) or min_ann < 1:
+ raise ConfigValidationError(
+ "adjudication.min_annotations must be a positive integer"
+ )
+
+ threshold = adj_config.get('agreement_threshold', 0.75)
+ if not isinstance(threshold, (int, float)) or threshold < 0 or threshold > 1:
+ raise ConfigValidationError(
+ "adjudication.agreement_threshold must be a number between 0 and 1"
+ )
+
+ fast_warn = adj_config.get('fast_decision_warning_ms', 2000)
+ if not isinstance(fast_warn, (int, float)) or fast_warn < 0:
+ raise ConfigValidationError(
+ "adjudication.fast_decision_warning_ms must be a non-negative number"
+ )
+
+ # Validate error_taxonomy
+ taxonomy = adj_config.get('error_taxonomy')
+ if taxonomy is not None:
+ if not isinstance(taxonomy, list):
+ raise ConfigValidationError(
+ "adjudication.error_taxonomy must be a list of strings"
+ )
+ for item in taxonomy:
+ if not isinstance(item, str):
+ raise ConfigValidationError(
+ "adjudication.error_taxonomy entries must be strings"
+ )
+
+ # Validate similarity config
+ sim_config = adj_config.get('similarity', {})
+ if isinstance(sim_config, dict) and sim_config.get('enabled', False):
+ top_k = sim_config.get('top_k', 5)
+ if not isinstance(top_k, int) or top_k < 1 or top_k > 20:
+ raise ConfigValidationError(
+ "adjudication.similarity.top_k must be an integer between 1 and 20"
+ )
+
+ model = sim_config.get('model', 'all-MiniLM-L6-v2')
+ if not isinstance(model, str) or not model.strip():
+ raise ConfigValidationError(
+ "adjudication.similarity.model must be a non-empty string"
+ )
+
+
+def _check_display_only_deprecation(config_data: Dict[str, Any]) -> None:
+ """
+ Check for deprecated display-only pattern and log warning.
+
+ Detects when image_annotation, video_annotation, or audio_annotation
+ is used with min_annotations: 0 just to display content.
+
+ Args:
+ config_data: The configuration data
+ """
+ # Get annotation schemes
+ schemes = []
+ if "annotation_schemes" in config_data:
+ schemes = config_data["annotation_schemes"]
+ elif "phases" in config_data:
+ phases = config_data["phases"]
+ if isinstance(phases, list):
+ for phase in phases:
+ schemes.extend(phase.get("annotation_schemes", []))
+ elif isinstance(phases, dict):
+ for phase_name, phase in phases.items():
+ if phase_name != "order" and isinstance(phase, dict):
+ schemes.extend(phase.get("annotation_schemes", []))
+
+ for scheme in schemes:
+ if not isinstance(scheme, dict):
+ continue
+
+ annotation_type = scheme.get("annotation_type")
+ if annotation_type in ["image_annotation", "video_annotation", "audio_annotation"]:
+ min_annotations = scheme.get("min_annotations", 1)
+ if min_annotations == 0:
+ logger.warning(
+ f"Deprecation warning: Using {annotation_type} with min_annotations=0 "
+ f"for display-only is deprecated. Use instance_display instead. "
+ f"See docs/instance_display.md for migration guide."
+ )
diff --git a/potato/server_utils/display_logic.py b/potato/server_utils/display_logic.py
new file mode 100644
index 0000000000000000000000000000000000000000..34118e5cc9ab3a9945f75503d06c42a62b61a819
--- /dev/null
+++ b/potato/server_utils/display_logic.py
@@ -0,0 +1,732 @@
+"""
+Display Logic Module for Conditional Schema Branching
+
+This module provides the core validation and evaluation logic for conditional
+annotation schemas. It allows schemas to show/hide based on user responses
+to other schemas.
+
+Key Components:
+- DisplayLogicCondition: Represents a single condition (e.g., "schema X equals 'Yes'")
+- DisplayLogicRule: Represents a complete rule with multiple conditions and AND/OR logic
+- DisplayLogicValidator: Validates display_logic configurations
+- DisplayLogicEvaluator: Evaluates conditions at runtime
+
+Example Configuration:
+ display_logic:
+ show_when:
+ - schema: contains_pii
+ operator: equals
+ value: "Yes"
+ logic: all # 'all' = AND, 'any' = OR
+"""
+
+import re
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+
+logger = logging.getLogger(__name__)
+
+# Supported operators and their descriptions
+SUPPORTED_OPERATORS = {
+ # Value comparison
+ "equals": "Exact value match (single value or list of values)",
+ "not_equals": "Value doesn't match any specified values",
+
+ # Collection operators
+ "contains": "List/text contains value(s)",
+ "not_contains": "List/text doesn't contain value(s)",
+
+ # Regex
+ "matches": "Regex pattern match",
+
+ # Numeric comparison
+ "gt": "Greater than",
+ "gte": "Greater than or equal",
+ "lt": "Less than",
+ "lte": "Less than or equal",
+ "in_range": "Value is within range (inclusive)",
+ "not_in_range": "Value is outside range",
+
+ # Emptiness
+ "empty": "Field is empty or not set",
+ "not_empty": "Field has a value",
+
+ # Text length
+ "length_gt": "Text length greater than",
+ "length_lt": "Text length less than",
+ "length_in_range": "Text length within range (inclusive)",
+}
+
+
+@dataclass
+class DisplayLogicCondition:
+ """
+ Represents a single condition in a display logic rule.
+
+ Attributes:
+ schema: Name of the schema to watch
+ operator: Comparison operator (equals, contains, gt, etc.)
+ value: Value(s) to compare against (can be single value, list, or range)
+ case_sensitive: Whether text comparisons are case-sensitive (default: False)
+ """
+ schema: str
+ operator: str
+ value: Any = None
+ case_sensitive: bool = False
+
+ def __post_init__(self):
+ """Validate the condition after initialization."""
+ if self.operator not in SUPPORTED_OPERATORS:
+ raise ValueError(f"Unsupported operator: {self.operator}. "
+ f"Supported operators: {list(SUPPORTED_OPERATORS.keys())}")
+
+ # Validate operator-specific requirements
+ if self.operator in ("empty", "not_empty"):
+ # These operators don't require a value
+ pass
+ elif self.operator in ("in_range", "not_in_range", "length_in_range"):
+ # Range operators require a list of exactly 2 values
+ if not isinstance(self.value, (list, tuple)) or len(self.value) != 2:
+ raise ValueError(f"Operator '{self.operator}' requires a range value "
+ f"as [min, max], got: {self.value}")
+ # Validate that range values are numeric
+ for v in self.value:
+ if not isinstance(v, (int, float)):
+ raise ValueError(f"Operator '{self.operator}' requires numeric range values, "
+ f"got: {self.value}")
+ elif self.operator in ("gt", "gte", "lt", "lte", "length_gt", "length_lt"):
+ # Numeric operators require numeric values
+ if self.value is None:
+ raise ValueError(f"Operator '{self.operator}' requires a numeric value")
+ if not isinstance(self.value, (int, float)):
+ raise ValueError(f"Operator '{self.operator}' requires a numeric value, "
+ f"got: {type(self.value).__name__} '{self.value}'")
+ elif self.value is None and self.operator not in ("empty", "not_empty"):
+ raise ValueError(f"Operator '{self.operator}' requires a value")
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert condition to dictionary for serialization."""
+ result = {
+ "schema": self.schema,
+ "operator": self.operator,
+ }
+ if self.value is not None:
+ result["value"] = self.value
+ if self.case_sensitive:
+ result["case_sensitive"] = True
+ return result
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "DisplayLogicCondition":
+ """Create a condition from a dictionary."""
+ return cls(
+ schema=data["schema"],
+ operator=data["operator"],
+ value=data.get("value"),
+ case_sensitive=data.get("case_sensitive", False)
+ )
+
+
+@dataclass
+class DisplayLogicRule:
+ """
+ Represents a complete display logic rule with multiple conditions.
+
+ Attributes:
+ conditions: List of DisplayLogicCondition objects
+ logic: 'all' (AND) or 'any' (OR) - how to combine conditions
+ """
+ conditions: List[DisplayLogicCondition] = field(default_factory=list)
+ logic: str = "all" # 'all' = AND, 'any' = OR
+
+ def __post_init__(self):
+ """Validate the rule after initialization."""
+ if self.logic not in ("all", "any"):
+ raise ValueError(f"Invalid logic type: {self.logic}. Must be 'all' or 'any'")
+
+ def get_watched_schemas(self) -> Set[str]:
+ """Return set of schema names this rule depends on."""
+ return {condition.schema for condition in self.conditions}
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert rule to dictionary for serialization."""
+ return {
+ "show_when": [c.to_dict() for c in self.conditions],
+ "logic": self.logic
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "DisplayLogicRule":
+ """Create a rule from a dictionary (config format)."""
+ conditions = []
+ show_when = data.get("show_when", [])
+
+ for cond_data in show_when:
+ conditions.append(DisplayLogicCondition.from_dict(cond_data))
+
+ return cls(
+ conditions=conditions,
+ logic=data.get("logic", "all")
+ )
+
+
+class DisplayLogicValidator:
+ """
+ Validates display_logic configurations in annotation schemes.
+
+ Responsibilities:
+ - Validate condition syntax and operators
+ - Check that referenced schemas exist
+ - Detect circular dependencies
+ - Warn about potential issues
+ """
+
+ def __init__(self, annotation_schemes: List[Dict[str, Any]]):
+ """
+ Initialize the validator with all annotation schemes.
+
+ Args:
+ annotation_schemes: List of annotation scheme configurations
+ """
+ self.schemes = annotation_schemes
+ self.schema_names = {s.get("name") for s in annotation_schemes if "name" in s}
+ self.dependency_graph: Dict[str, Set[str]] = {}
+ self._build_dependency_graph()
+
+ def _build_dependency_graph(self) -> None:
+ """Build a graph of schema dependencies for cycle detection."""
+ for scheme in self.schemes:
+ schema_name = scheme.get("name")
+ if not schema_name:
+ continue
+
+ display_logic = scheme.get("display_logic", {})
+ if not display_logic:
+ self.dependency_graph[schema_name] = set()
+ continue
+
+ # Extract schemas this one depends on
+ dependencies = set()
+ show_when = display_logic.get("show_when", [])
+ for condition in show_when:
+ if "schema" in condition:
+ dependencies.add(condition["schema"])
+
+ self.dependency_graph[schema_name] = dependencies
+
+ def validate(self) -> Tuple[bool, List[str]]:
+ """
+ Validate all display_logic configurations.
+
+ Returns:
+ Tuple of (is_valid, list_of_errors)
+ """
+ errors = []
+
+ for scheme in self.schemes:
+ schema_name = scheme.get("name", "")
+ display_logic = scheme.get("display_logic")
+
+ # Skip if display_logic is not present (None) or not a dict
+ if display_logic is None:
+ continue
+ if not isinstance(display_logic, dict):
+ errors.append(f"Schema '{schema_name}': display_logic must be a dictionary")
+ continue
+
+ # Validate structure (including empty dict - which is invalid)
+ structure_errors = self._validate_structure(schema_name, display_logic)
+ errors.extend(structure_errors)
+
+ # Validate referenced schemas exist
+ reference_errors = self._validate_references(schema_name, display_logic)
+ errors.extend(reference_errors)
+
+ # Check for circular dependencies
+ cycle_errors = self._detect_cycles()
+ errors.extend(cycle_errors)
+
+ return len(errors) == 0, errors
+
+ def _validate_structure(self, schema_name: str, display_logic: Dict) -> List[str]:
+ """Validate the structure of a display_logic configuration."""
+ errors = []
+
+ # Must have show_when
+ if "show_when" not in display_logic:
+ errors.append(f"Schema '{schema_name}': display_logic must have 'show_when' field")
+ return errors
+
+ show_when = display_logic["show_when"]
+ if not isinstance(show_when, list):
+ errors.append(f"Schema '{schema_name}': 'show_when' must be a list of conditions")
+ return errors
+
+ if len(show_when) == 0:
+ errors.append(f"Schema '{schema_name}': 'show_when' must have at least one condition")
+ return errors
+
+ # Validate each condition
+ for i, condition in enumerate(show_when):
+ prefix = f"Schema '{schema_name}', condition {i+1}"
+
+ if not isinstance(condition, dict):
+ errors.append(f"{prefix}: condition must be a dictionary")
+ continue
+
+ # Required fields
+ if "schema" not in condition:
+ errors.append(f"{prefix}: missing required 'schema' field")
+
+ if "operator" not in condition:
+ errors.append(f"{prefix}: missing required 'operator' field")
+ elif condition["operator"] not in SUPPORTED_OPERATORS:
+ errors.append(f"{prefix}: unsupported operator '{condition['operator']}'. "
+ f"Supported: {list(SUPPORTED_OPERATORS.keys())}")
+
+ # Validate operator-specific requirements
+ operator = condition.get("operator")
+ if operator:
+ op_errors = self._validate_operator_value(prefix, operator, condition.get("value"))
+ errors.extend(op_errors)
+
+ # Validate logic field if present
+ logic = display_logic.get("logic", "all")
+ if logic not in ("all", "any"):
+ errors.append(f"Schema '{schema_name}': 'logic' must be 'all' or 'any', got '{logic}'")
+
+ return errors
+
+ def _validate_operator_value(self, prefix: str, operator: str, value: Any) -> List[str]:
+ """Validate that the value is appropriate for the operator."""
+ errors = []
+
+ # Operators that don't need a value
+ if operator in ("empty", "not_empty"):
+ return errors
+
+ # Range operators need [min, max]
+ if operator in ("in_range", "not_in_range", "length_in_range"):
+ if not isinstance(value, (list, tuple)):
+ errors.append(f"{prefix}: operator '{operator}' requires a range value as [min, max]")
+ elif len(value) != 2:
+ errors.append(f"{prefix}: range value must have exactly 2 elements [min, max]")
+ else:
+ try:
+ min_val, max_val = float(value[0]), float(value[1])
+ if min_val > max_val:
+ errors.append(f"{prefix}: range min ({min_val}) is greater than max ({max_val})")
+ except (ValueError, TypeError):
+ errors.append(f"{prefix}: range values must be numeric")
+ return errors
+
+ # Numeric operators need numeric values
+ if operator in ("gt", "gte", "lt", "lte", "length_gt", "length_lt"):
+ if value is None:
+ errors.append(f"{prefix}: operator '{operator}' requires a value")
+ else:
+ try:
+ float(value)
+ except (ValueError, TypeError):
+ errors.append(f"{prefix}: operator '{operator}' requires a numeric value")
+ return errors
+
+ # Regex operator needs a valid pattern
+ if operator == "matches":
+ if value is None:
+ errors.append(f"{prefix}: operator 'matches' requires a regex pattern")
+ else:
+ try:
+ re.compile(value)
+ except re.error as e:
+ errors.append(f"{prefix}: invalid regex pattern '{value}': {e}")
+ return errors
+
+ # Other operators just need a non-None value
+ if value is None:
+ errors.append(f"{prefix}: operator '{operator}' requires a value")
+
+ return errors
+
+ def _validate_references(self, schema_name: str, display_logic: Dict) -> List[str]:
+ """Validate that all referenced schemas exist."""
+ errors = []
+ show_when = display_logic.get("show_when", [])
+
+ for i, condition in enumerate(show_when):
+ ref_schema = condition.get("schema")
+ if ref_schema and ref_schema not in self.schema_names:
+ errors.append(
+ f"Schema '{schema_name}', condition {i+1}: references unknown schema '{ref_schema}'"
+ )
+
+ return errors
+
+ def _detect_cycles(self) -> List[str]:
+ """Detect circular dependencies using DFS."""
+ errors = []
+ visited = set()
+ rec_stack = set()
+
+ def dfs(node: str, path: List[str]) -> Optional[List[str]]:
+ """DFS to detect cycles, returns cycle path if found."""
+ if node in rec_stack:
+ # Found a cycle
+ cycle_start = path.index(node)
+ return path[cycle_start:] + [node]
+
+ if node in visited:
+ return None
+
+ visited.add(node)
+ rec_stack.add(node)
+
+ for neighbor in self.dependency_graph.get(node, set()):
+ result = dfs(neighbor, path + [node])
+ if result:
+ return result
+
+ rec_stack.remove(node)
+ return None
+
+ for schema in self.dependency_graph:
+ if schema not in visited:
+ cycle = dfs(schema, [])
+ if cycle:
+ cycle_str = " -> ".join(cycle)
+ errors.append(f"Circular dependency detected: {cycle_str}")
+
+ return errors
+
+ def get_schema_dependencies(self, schema_name: str) -> Set[str]:
+ """Get the schemas that a given schema depends on."""
+ return self.dependency_graph.get(schema_name, set())
+
+ def get_dependents(self, schema_name: str) -> Set[str]:
+ """Get schemas that depend on the given schema."""
+ dependents = set()
+ for schema, deps in self.dependency_graph.items():
+ if schema_name in deps:
+ dependents.add(schema)
+ return dependents
+
+
+class DisplayLogicEvaluator:
+ """
+ Evaluates display logic conditions at runtime.
+
+ This class is used both server-side (Python) and provides the logic
+ that's replicated in the frontend JavaScript.
+ """
+
+ @staticmethod
+ def evaluate_condition(
+ condition: DisplayLogicCondition,
+ schema_value: Any
+ ) -> bool:
+ """
+ Evaluate a single condition against a schema value.
+
+ Args:
+ condition: The condition to evaluate
+ schema_value: Current value of the watched schema
+
+ Returns:
+ bool: Whether the condition is satisfied
+ """
+ operator = condition.operator
+ expected = condition.value
+ case_sensitive = condition.case_sensitive
+
+ # Handle empty checks first
+ if operator == "empty":
+ return DisplayLogicEvaluator._is_empty(schema_value)
+
+ if operator == "not_empty":
+ return not DisplayLogicEvaluator._is_empty(schema_value)
+
+ # For all other operators, normalize the actual value
+ actual = schema_value
+
+ # Apply case normalization for text comparisons
+ if not case_sensitive and isinstance(actual, str):
+ actual = actual.lower()
+
+ # Equality operators
+ if operator == "equals":
+ return DisplayLogicEvaluator._check_equals(actual, expected, case_sensitive)
+
+ if operator == "not_equals":
+ return not DisplayLogicEvaluator._check_equals(actual, expected, case_sensitive)
+
+ # Contains operators (for lists and text)
+ if operator == "contains":
+ return DisplayLogicEvaluator._check_contains(actual, expected, case_sensitive)
+
+ if operator == "not_contains":
+ return not DisplayLogicEvaluator._check_contains(actual, expected, case_sensitive)
+
+ # Regex matching
+ if operator == "matches":
+ if not isinstance(actual, str):
+ actual = str(actual) if actual is not None else ""
+ flags = 0 if case_sensitive else re.IGNORECASE
+ try:
+ return bool(re.search(expected, actual, flags))
+ except re.error:
+ logger.warning(f"Invalid regex pattern: {expected}")
+ return False
+
+ # Numeric comparisons
+ if operator in ("gt", "gte", "lt", "lte"):
+ return DisplayLogicEvaluator._check_numeric(operator, actual, expected)
+
+ # Range operators
+ if operator in ("in_range", "not_in_range"):
+ result = DisplayLogicEvaluator._check_range(actual, expected)
+ return result if operator == "in_range" else not result
+
+ # Length operators
+ if operator in ("length_gt", "length_lt"):
+ return DisplayLogicEvaluator._check_length(operator, actual, expected)
+
+ if operator == "length_in_range":
+ return DisplayLogicEvaluator._check_length_range(actual, expected)
+
+ logger.warning(f"Unknown operator: {operator}")
+ return False
+
+ @staticmethod
+ def _is_empty(value: Any) -> bool:
+ """Check if a value is considered empty."""
+ if value is None:
+ return True
+ if isinstance(value, str):
+ return len(value.strip()) == 0
+ if isinstance(value, (list, dict, set)):
+ return len(value) == 0
+ return False
+
+ @staticmethod
+ def _check_equals(actual: Any, expected: Any, case_sensitive: bool) -> bool:
+ """Check equality, handling single values and lists."""
+ # If expected is a list, check if actual matches ANY of them
+ if isinstance(expected, list):
+ for exp in expected:
+ if DisplayLogicEvaluator._values_equal(actual, exp, case_sensitive):
+ return True
+ return False
+
+ return DisplayLogicEvaluator._values_equal(actual, expected, case_sensitive)
+
+ @staticmethod
+ def _values_equal(actual: Any, expected: Any, case_sensitive: bool) -> bool:
+ """Compare two values for equality."""
+ # Handle None
+ if actual is None and expected is None:
+ return True
+ if actual is None or expected is None:
+ return False
+
+ # String comparison with case sensitivity
+ if isinstance(expected, str):
+ actual_str = str(actual)
+ if not case_sensitive:
+ return actual_str.lower() == expected.lower()
+ return actual_str == expected
+
+ # Direct comparison for non-strings
+ return actual == expected
+
+ @staticmethod
+ def _check_contains(actual: Any, expected: Any, case_sensitive: bool) -> bool:
+ """Check if actual contains expected value(s)."""
+ # If expected is a list, check if actual contains ANY of them
+ if isinstance(expected, list):
+ for exp in expected:
+ if DisplayLogicEvaluator._value_contains(actual, exp, case_sensitive):
+ return True
+ return False
+
+ return DisplayLogicEvaluator._value_contains(actual, expected, case_sensitive)
+
+ @staticmethod
+ def _value_contains(actual: Any, expected: Any, case_sensitive: bool) -> bool:
+ """Check if actual contains a single expected value."""
+ # If actual is a list (multiselect), check membership
+ if isinstance(actual, list):
+ for item in actual:
+ if DisplayLogicEvaluator._values_equal(item, expected, case_sensitive):
+ return True
+ return False
+
+ # If actual is a string, check substring
+ if isinstance(actual, str):
+ expected_str = str(expected)
+ if not case_sensitive:
+ return expected_str.lower() in actual.lower()
+ return expected_str in actual
+
+ # Fallback to equality
+ return DisplayLogicEvaluator._values_equal(actual, expected, case_sensitive)
+
+ @staticmethod
+ def _check_numeric(operator: str, actual: Any, expected: Any) -> bool:
+ """Check numeric comparison."""
+ try:
+ actual_num = float(actual) if actual is not None else 0
+ expected_num = float(expected)
+ except (ValueError, TypeError):
+ return False
+
+ if operator == "gt":
+ return actual_num > expected_num
+ if operator == "gte":
+ return actual_num >= expected_num
+ if operator == "lt":
+ return actual_num < expected_num
+ if operator == "lte":
+ return actual_num <= expected_num
+
+ return False
+
+ @staticmethod
+ def _check_range(actual: Any, range_val: List) -> bool:
+ """Check if actual is within range (inclusive)."""
+ try:
+ actual_num = float(actual) if actual is not None else 0
+ min_val, max_val = float(range_val[0]), float(range_val[1])
+ except (ValueError, TypeError, IndexError):
+ return False
+
+ return min_val <= actual_num <= max_val
+
+ @staticmethod
+ def _check_length(operator: str, actual: Any, expected: Any) -> bool:
+ """Check text length comparison."""
+ try:
+ length = len(str(actual)) if actual is not None else 0
+ expected_len = int(expected)
+ except (ValueError, TypeError):
+ return False
+
+ if operator == "length_gt":
+ return length > expected_len
+ if operator == "length_lt":
+ return length < expected_len
+
+ return False
+
+ @staticmethod
+ def _check_length_range(actual: Any, range_val: List) -> bool:
+ """Check if text length is within range (inclusive)."""
+ try:
+ length = len(str(actual)) if actual is not None else 0
+ min_len, max_len = int(range_val[0]), int(range_val[1])
+ except (ValueError, TypeError, IndexError):
+ return False
+
+ return min_len <= length <= max_len
+
+ @staticmethod
+ def evaluate_rule(
+ rule: DisplayLogicRule,
+ annotations: Dict[str, Any]
+ ) -> bool:
+ """
+ Evaluate a complete display logic rule.
+
+ Args:
+ rule: The DisplayLogicRule to evaluate
+ annotations: Current annotations dictionary {schema_name: value}
+
+ Returns:
+ bool: Whether the schema should be visible
+ """
+ if not rule.conditions:
+ # No conditions = always visible
+ return True
+
+ results = []
+ for condition in rule.conditions:
+ schema_value = annotations.get(condition.schema)
+ result = DisplayLogicEvaluator.evaluate_condition(condition, schema_value)
+ results.append(result)
+
+ if rule.logic == "all":
+ return all(results)
+ else: # "any"
+ return any(results)
+
+ @staticmethod
+ def evaluate_visibility(
+ schema_name: str,
+ display_logic: Optional[Dict],
+ annotations: Dict[str, Any]
+ ) -> Tuple[bool, Optional[str]]:
+ """
+ Evaluate whether a schema should be visible.
+
+ Args:
+ schema_name: Name of the schema being evaluated
+ display_logic: The display_logic configuration (can be None)
+ annotations: Current annotations dictionary
+
+ Returns:
+ Tuple of (is_visible, reason_if_hidden)
+ """
+ if not display_logic:
+ return True, None
+
+ try:
+ rule = DisplayLogicRule.from_dict(display_logic)
+ is_visible = DisplayLogicEvaluator.evaluate_rule(rule, annotations)
+
+ if not is_visible:
+ # Build reason string
+ reasons = []
+ for cond in rule.conditions:
+ actual_val = annotations.get(cond.schema, "")
+ reasons.append(f"{cond.schema} {cond.operator} {cond.value} (actual: {actual_val})")
+ reason = f"Conditions not met ({rule.logic}): " + ", ".join(reasons)
+ return False, reason
+
+ return True, None
+
+ except Exception as e:
+ logger.error(f"Error evaluating display logic for {schema_name}: {e}")
+ # Default to visible on error
+ return True, None
+
+
+def validate_display_logic_config(
+ annotation_schemes: List[Dict[str, Any]]
+) -> Tuple[bool, List[str]]:
+ """
+ Convenience function to validate display_logic across all annotation schemes.
+
+ Args:
+ annotation_schemes: List of annotation scheme configurations
+
+ Returns:
+ Tuple of (is_valid, list_of_errors)
+ """
+ validator = DisplayLogicValidator(annotation_schemes)
+ return validator.validate()
+
+
+def get_display_logic_dependencies(
+ annotation_schemes: List[Dict[str, Any]]
+) -> Dict[str, Set[str]]:
+ """
+ Get the dependency graph for all schemas with display_logic.
+
+ Args:
+ annotation_schemes: List of annotation scheme configurations
+
+ Returns:
+ Dictionary mapping schema names to their dependencies
+ """
+ validator = DisplayLogicValidator(annotation_schemes)
+ return validator.dependency_graph
diff --git a/potato/server_utils/displays/ARCHITECTURE.md b/potato/server_utils/displays/ARCHITECTURE.md
new file mode 100644
index 0000000000000000000000000000000000000000..b362dfe042b8f99a2a0410bafd7ef3b9881f833f
--- /dev/null
+++ b/potato/server_utils/displays/ARCHITECTURE.md
@@ -0,0 +1,162 @@
+# Display Type System Architecture
+
+This document describes the design contracts and extension points for the
+display type system in `potato/server_utils/displays/`.
+
+## Overview
+
+The display system separates **content presentation** from **annotation
+collection**. Each field in `instance_display.fields` has a `type` that
+maps to a registered display class. Displays produce HTML; annotation
+schemas collect labels.
+
+```
+Config YAML
+ โโ instance_display.fields[].type
+ โโ DisplayRegistry.render()
+ โโ BaseDisplay.render() โ inner HTML
+ โโ render_display_container() โ wrapped HTML
+ โโ template {{ display_html | safe }}
+```
+
+## Key Files
+
+| File | Purpose |
+|------|---------|
+| `base.py` | `BaseDisplay` ABC โ class attributes, abstract `render()`, helpers |
+| `registry.py` | `DisplayRegistry` singleton โ registration, lookup, render dispatch |
+| `../instance_display.py` | `InstanceDisplayRenderer` โ orchestrates field rendering |
+| `__init__.py` | Package exports |
+
+## BaseDisplay Contract
+
+### Required to implement
+
+| Method / Attribute | Description |
+|--------------------|-------------|
+| `name: str` | Unique type identifier (e.g., `"dialogue"`) |
+| `render(field_config, data) -> str` | Return inner HTML for the field content |
+
+### Optional to override
+
+| Method | Default | When to override |
+|--------|---------|------------------|
+| `get_css_classes(field_config)` | `["display-field", "display-type-{name}"]` | Add type-specific classes |
+| `get_data_attributes(field_config, data)` | `{"field-key", "field-type", "span-target"}` | Add custom data attrs |
+| `get_js_init()` | `None` | Return JS to run on page load |
+| `validate_config(field_config)` | Checks `required_fields` | Add enum/range validation |
+| `has_inline_label(field_config)` | `False` | Return `True` if display renders its own label (avoids duplicate) |
+| `get_display_options(field_config)` | Merges `optional_fields` with `display_options` | Rarely needed |
+
+### Class attributes
+
+| Attribute | Type | Description |
+|-----------|------|-------------|
+| `required_fields` | `List[str]` | Config keys that must be present |
+| `optional_fields` | `Dict[str, Any]` | Default values for optional display_options |
+| `description` | `str` | Human-readable description |
+| `supports_span_target` | `bool` | Whether this type implements the span annotation contract |
+
+## Span Target Contract
+
+**If `supports_span_target = True`, the display MUST satisfy these requirements
+when `field_config["span_target"]` is `True`:**
+
+### 1. `.text-content` wrapper
+
+The rendered HTML must contain:
+
+```html
+
+ {content HTML}
+
+```
+
+Use the `render_span_wrapper()` helper:
+
+```python
+if field_config.get("span_target"):
+ inner_html = self.render_span_wrapper(field_key, inner_html, plain_text)
+```
+
+### 2. `data-original-text` must contain plain text
+
+The `plain_text` argument to `render_span_wrapper()` must be the canonical
+plain text that `routes.py` will use for span offset extraction. For
+structured data (dialogue, lists), use `concatenate_dialogue_text()` from
+`base.py` so both rendering and API extraction use identical formats.
+
+### 3. CSS classes on the outer container
+
+Override `get_css_classes()` to add `"span-target-field"` and
+`"span-target-{name}"` when span_target is true.
+
+### 4. Text format consistency
+
+The text format used in `data-original-text` **MUST** match the text
+extraction logic in `routes.py` (`/api/spans/` endpoint). If the
+data is a list of dicts, both sides must use `concatenate_dialogue_text()`.
+
+### Why this matters
+
+SpanManager (span-core.js) discovers span-target fields via:
+```javascript
+document.querySelectorAll('.display-field[data-span-target="true"]')
+```
+Then looks for the text element inside each:
+```javascript
+const textContent = field.querySelector('.text-content');
+```
+If `.text-content` is missing, SpanManager silently skips the field and
+span annotation will not work.
+
+## Registry
+
+The `display_registry` singleton provides:
+
+- `render(field_type, field_config, data)` โ render a field
+- `get_supported_types()` โ list all registered type names
+- `type_supports_span_target(field_type)` โ check span target support
+- `get_span_target_types()` โ list all types supporting span targets
+- `validate(field_type, field_config)` โ validate config
+- `list_displays()` โ metadata for all displays
+
+The registry wraps each display's `render()` output in
+`render_display_container()`, which adds the outer `.display-field` div,
+label, and `.display-field-content` wrapper.
+
+## Instance Display Renderer
+
+`InstanceDisplayRenderer` in `instance_display.py`:
+
+1. Reads `instance_display.fields` from config
+2. Queries `display_registry.type_supports_span_target()` for span targets
+ (no hardcoded list)
+3. Warns if `span_target: true` is set on an unsupported type
+4. Renders each field via `display_registry.render()`
+5. Optionally wraps in resizable container (`_wrap_resizable()`)
+
+## Adding a New Display Type
+
+1. Create `my_display.py` with a class extending `BaseDisplay`
+2. Set `name`, `required_fields`, `optional_fields`, `description`
+3. If supporting span annotation:
+ - Set `supports_span_target = True`
+ - Use `render_span_wrapper()` in `render()` when `span_target` is True
+ - Override `get_css_classes()` to add `span-target-field`
+4. Register in `registry.py` via `DisplayDefinition`
+5. Add to `__init__.py` exports
+6. Add CSS to `styles.css` using `.display-type-{name}` convention
+7. Write unit tests verifying render output
+8. Write a contract enforcement test (see `test_display_span_contract.py`)
+
+## Shared Utilities
+
+| Function | Location | Purpose |
+|----------|----------|---------|
+| `render_span_wrapper(field_key, inner_html, plain_text)` | `BaseDisplay` method | Standard `.text-content` wrapper |
+| `concatenate_dialogue_text(data, speaker_key, text_key)` | `base.py` module | Canonical dialogueโplain text conversion |
+| `render_display_container(inner_html, classes, attrs, label)` | `base.py` module | Standard outer container wrapper |
diff --git a/potato/server_utils/displays/__init__.py b/potato/server_utils/displays/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7756ccec84e67943bb707b782c826a319c57459d
--- /dev/null
+++ b/potato/server_utils/displays/__init__.py
@@ -0,0 +1,47 @@
+"""
+Display Types Package
+
+Provides renderers for different content types in instance display.
+This module separates content display from annotation collection.
+
+Usage:
+ from potato.server_utils.displays import display_registry
+
+ # Render a display field
+ html = display_registry.render("image", field_config, data)
+
+ # List available display types
+ types = display_registry.get_supported_types()
+"""
+
+from .registry import display_registry, DisplayDefinition, DisplayRegistry
+from .base import BaseDisplay, render_display_container, concatenate_dialogue_text
+from .pdf_display import PDFDisplay
+from .document_display import DocumentDisplay
+from .spreadsheet_display import SpreadsheetDisplay
+from .code_display import CodeDisplay
+from .agent_trace_display import AgentTraceDisplay
+from .eval_trace_display import EvalTraceDisplay
+from .gallery_display import GalleryDisplay
+from .interactive_chat_display import InteractiveChatDisplay
+from .web_agent_trace_display import WebAgentTraceDisplay
+from .live_agent_display import LiveAgentDisplay
+
+__all__ = [
+ 'display_registry',
+ 'DisplayDefinition',
+ 'DisplayRegistry',
+ 'BaseDisplay',
+ 'render_display_container',
+ 'concatenate_dialogue_text',
+ 'PDFDisplay',
+ 'DocumentDisplay',
+ 'SpreadsheetDisplay',
+ 'CodeDisplay',
+ 'AgentTraceDisplay',
+ 'EvalTraceDisplay',
+ 'GalleryDisplay',
+ 'InteractiveChatDisplay',
+ 'WebAgentTraceDisplay',
+ 'LiveAgentDisplay',
+]
diff --git a/potato/server_utils/displays/_trace_normalize.py b/potato/server_utils/displays/_trace_normalize.py
new file mode 100644
index 0000000000000000000000000000000000000000..50b28b74781f8bf98c5128f9c73fea32bb7901f0
--- /dev/null
+++ b/potato/server_utils/displays/_trace_normalize.py
@@ -0,0 +1,146 @@
+"""
+Shared agent-trace normalization.
+
+Both ``agent_trace`` (vertical step cards) and ``eval_trace`` (three-pane
+reasoning | function calls | final answer) need to turn heterogeneous trace
+data into a flat list of typed steps. This module is the single source of
+truth for that parsing so the two displays never drift apart.
+
+A normalized step is a dict with keys:
+ type: one of "thought" | "action" | "observation" | "system" | "error"
+ speaker: display label for the step (may be "")
+ text: the step's textual content
+ timestamp: optional timestamp string ("")
+ screenshot: optional screenshot URL ("")
+
+Supported input formats (see ``normalize_steps``):
+ - a single string -> one observation step
+ - list of strings -> one step each (type inferred)
+ - list of {speaker, text} dicts -> dialogue-style turns
+ - list of {thought, action, observation}-> one dict expands to 1-3 steps
+ - list of {step_type, content} dicts -> explicit typing
+"""
+
+import re
+from typing import Any, Dict, List
+
+
+# Default background colors for step types (consumed by displays' CSS builders).
+DEFAULT_STEP_COLORS = {
+ "thought": "#e8f4fd",
+ "action": "#fff3e0",
+ "observation": "#e8f5e9",
+ "system": "#f3e5f5",
+ "error": "#ffebee",
+}
+
+# Speaker/label substrings that map to a step type.
+SPEAKER_TYPE_PATTERNS = {
+ "thought": re.compile(r"(thought|reasoning|planning|think)", re.IGNORECASE),
+ "action": re.compile(r"(action|tool|function|call|execute)", re.IGNORECASE),
+ "observation": re.compile(r"(observation|environment|result|output|response)", re.IGNORECASE),
+ "system": re.compile(r"(system|info|metadata)", re.IGNORECASE),
+ "error": re.compile(r"(error|fail|exception)", re.IGNORECASE),
+}
+
+
+def infer_type_from_speaker(speaker: str) -> str:
+ """Infer a step type from a speaker/label string."""
+ if not speaker:
+ return "observation"
+ for type_name, pattern in SPEAKER_TYPE_PATTERNS.items():
+ if pattern.search(speaker):
+ return type_name
+ return "observation"
+
+
+def infer_type_from_text(text: str) -> str:
+ """Infer a step type from free text content."""
+ lower = text.lower()
+ if lower.startswith(("i need to", "i should", "let me think", "my plan")):
+ return "thought"
+ if "(" in text and ")" in text and any(c.isalpha() for c in text.split("(")[0]):
+ return "action"
+ return "observation"
+
+
+def format_action_text(action: Any) -> str:
+ """Render an action value as ``tool(args)`` when it is a structured dict."""
+ if isinstance(action, dict):
+ tool = action.get("tool", action.get("name", ""))
+ params = action.get("params", action.get("parameters", {}))
+ if params:
+ args = ", ".join(f"{k}={repr(v)}" for k, v in params.items())
+ return f"{tool}({args})"
+ return f"{tool}()"
+ return str(action)
+
+
+def normalize_steps(
+ data: Any,
+ speaker_key: str = "speaker",
+ text_key: str = "text",
+) -> List[Dict[str, str]]:
+ """Normalize heterogeneous trace data into a list of typed step dicts.
+
+ See the module docstring for the accepted input formats and the shape of
+ each returned step.
+ """
+ steps: List[Dict[str, str]] = []
+
+ if isinstance(data, str):
+ return [{"type": "observation", "speaker": "", "text": data}]
+
+ if not isinstance(data, list):
+ return steps
+
+ for item in data:
+ if isinstance(item, str):
+ step_type = infer_type_from_text(item)
+ steps.append({"type": step_type, "speaker": "", "text": item})
+ elif isinstance(item, dict):
+ # Format 1: speaker/text (same as dialogue)
+ if speaker_key in item and text_key in item:
+ speaker = item[speaker_key]
+ text = item[text_key]
+ step_type = item.get("step_type", infer_type_from_speaker(speaker))
+ steps.append({
+ "type": step_type,
+ "speaker": speaker,
+ "text": text,
+ "timestamp": item.get("timestamp", ""),
+ "screenshot": item.get("screenshot", ""),
+ })
+ # Format 2: thought/action/observation (one dict = up to 3 steps)
+ elif any(k in item for k in ("thought", "action", "observation")):
+ if item.get("thought"):
+ steps.append({
+ "type": "thought",
+ "speaker": "Agent (Thought)",
+ "text": str(item["thought"]),
+ "timestamp": item.get("timestamp", ""),
+ })
+ if item.get("action"):
+ steps.append({
+ "type": "action",
+ "speaker": "Agent (Action)",
+ "text": format_action_text(item["action"]),
+ })
+ if item.get("observation"):
+ steps.append({
+ "type": "observation",
+ "speaker": "Environment",
+ "text": str(item["observation"]),
+ "screenshot": item.get("screenshot", ""),
+ })
+ # Format 3: step_type/content
+ elif "step_type" in item:
+ steps.append({
+ "type": item["step_type"],
+ "speaker": item.get("speaker", item.get("step_type", "").capitalize()),
+ "text": item.get("content", item.get("text", "")),
+ "timestamp": item.get("timestamp", ""),
+ "screenshot": item.get("screenshot", ""),
+ })
+
+ return steps
diff --git a/potato/server_utils/displays/agent_trace_display.py b/potato/server_utils/displays/agent_trace_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..298fcd63372c494029289122e125d101780f889c
--- /dev/null
+++ b/potato/server_utils/displays/agent_trace_display.py
@@ -0,0 +1,285 @@
+"""
+Agent Trace Display Type
+
+Purpose-built rendering for agent traces as a vertical sequence of "step cards".
+Each step shows a type badge (Thought / Action / Observation) with color coding,
+collapsible tool call parameters, and optional inline screenshot thumbnails.
+
+This display type provides a richer visual alternative to the dialogue display
+for agent trace data, with step-type-aware styling and summary headers.
+"""
+
+import html
+from typing import Dict, Any, List, Optional
+
+from .base import BaseDisplay
+from ._trace_normalize import (
+ DEFAULT_STEP_COLORS,
+ SPEAKER_TYPE_PATTERNS,
+ normalize_steps,
+ infer_type_from_speaker,
+ infer_type_from_text,
+)
+
+
+class AgentTraceDisplay(BaseDisplay):
+ """
+ Display type for agent traces rendered as step cards.
+
+ Supports data as:
+ - List of dicts with speaker/text keys (same as dialogue)
+ - List of dicts with step_type/content keys
+ - List of dicts with thought/action/observation keys (one step per dict)
+ """
+
+ name = "agent_trace"
+ required_fields = ["key"]
+ optional_fields = {
+ "show_timestamps": False,
+ "collapse_observations": False,
+ "step_type_colors": DEFAULT_STEP_COLORS,
+ "show_screenshots": True,
+ "show_step_numbers": True,
+ "show_summary": True,
+ "compact": False,
+ "speaker_key": "speaker",
+ "text_key": "text",
+ }
+ description = "Agent trace display with step cards and type badges"
+ supports_span_target = False # Per-step IDs don't follow .text-content wrapper contract
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ if not data:
+ return 'No trace data provided
'
+
+ options = self.get_display_options(field_config)
+ show_timestamps = options.get("show_timestamps", False)
+ collapse_obs = options.get("collapse_observations", False)
+ colors = options.get("step_type_colors", DEFAULT_STEP_COLORS)
+ show_screenshots = options.get("show_screenshots", True)
+ show_step_numbers = options.get("show_step_numbers", True)
+ show_summary = options.get("show_summary", True)
+ compact = options.get("compact", False)
+ speaker_key = options.get("speaker_key", "speaker")
+ text_key = options.get("text_key", "text")
+
+ field_key = html.escape(field_config.get("key", ""), quote=True)
+ is_span_target = field_config.get("span_target", False)
+
+ # Normalize data to steps
+ steps = self._normalize_steps(data, speaker_key, text_key)
+
+ if not steps:
+ return 'No trace steps found
'
+
+ # Build CSS for step type colors
+ css = self._build_css(colors, compact)
+
+ # Build summary header
+ summary_html = ""
+ if show_summary:
+ summary_html = self._build_summary(steps)
+
+ # Build step cards
+ step_html_list = []
+ step_counter = 0
+ for i, step in enumerate(steps):
+ step_type = step.get("type", "observation")
+ speaker = step.get("speaker", "")
+ text = step.get("text", "")
+ timestamp = step.get("timestamp", "")
+ screenshot = step.get("screenshot", "")
+
+ if step_type in ("action", "thought"):
+ step_counter += 1
+
+ # Step card - sanitize step_type for use in CSS class names and attributes
+ safe_step_type = html.escape(step_type, quote=True)
+ type_class = f"step-type-{safe_step_type}"
+ card_classes = ["agent-trace-step", type_class]
+ if compact:
+ card_classes.append("compact")
+
+ # Badge
+ badge_label = safe_step_type.capitalize()
+ if speaker:
+ badge_label = html.escape(str(speaker))
+
+ # Step number
+ step_num_html = ""
+ if show_step_numbers and step_type in ("action", "thought"):
+ step_num_html = f'#{step_counter} '
+
+ # Timestamp
+ ts_html = ""
+ if show_timestamps and timestamp:
+ ts_html = f'{html.escape(str(timestamp))} '
+
+ # Text content
+ escaped_text = html.escape(str(text))
+ text_id = ""
+ span_attrs = ""
+ if is_span_target:
+ text_id = f'id="step-text-{field_key}-{i}"'
+ span_attrs = f'data-original-text="{escaped_text}" data-step-index="{i}"'
+
+ # Collapsible wrapper for observations
+ if collapse_obs and step_type == "observation":
+ text_html = (
+ f''
+ f'Observation (click to expand) '
+ f'{escaped_text}
'
+ f' '
+ )
+ else:
+ text_html = f'{escaped_text}
'
+
+ # Screenshot thumbnail
+ screenshot_html = ""
+ if show_screenshots and screenshot:
+ escaped_src = html.escape(str(screenshot), quote=True)
+ screenshot_html = (
+ f''
+ f'
'
+ f'
'
+ )
+
+ card_html = f'''
+
+
+
+ {text_html}
+ {screenshot_html}
+
+
+ '''
+ step_html_list.append(card_html)
+
+ all_steps_html = "\n".join(step_html_list)
+
+ container_classes = ["agent-trace-display"]
+ if is_span_target:
+ container_classes.append("span-target-agent-trace")
+
+ return f'''
+
+
+ {summary_html}
+
+ {all_steps_html}
+
+
+ '''
+
+ def _normalize_steps(self, data: Any, speaker_key: str, text_key: str) -> List[Dict[str, str]]:
+ """Normalize various trace data formats to a list of step dicts.
+
+ Delegates to the shared :func:`normalize_steps` so ``agent_trace`` and
+ ``eval_trace`` parse identical data the same way.
+ """
+ return normalize_steps(data, speaker_key, text_key)
+
+ def _infer_type_from_speaker(self, speaker: str) -> str:
+ """Infer step type from speaker name."""
+ return infer_type_from_speaker(speaker)
+
+ def _infer_type_from_text(self, text: str) -> str:
+ """Infer step type from text content."""
+ return infer_type_from_text(text)
+
+ def _build_summary(self, steps: List[Dict]) -> str:
+ """Build a summary header showing step counts."""
+ type_counts = {}
+ for step in steps:
+ t = step.get("type", "observation")
+ type_counts[t] = type_counts.get(t, 0) + 1
+
+ badges = []
+ for step_type in ["thought", "action", "observation"]:
+ count = type_counts.get(step_type, 0)
+ if count > 0:
+ badges.append(
+ f''
+ f'{count} {step_type}{"s" if count != 1 else ""} '
+ )
+
+ return f'''
+
+ {len(steps)} steps
+ {" ".join(badges)}
+
+ '''
+
+ def _build_css(self, colors: Dict[str, str], compact: bool) -> str:
+ """Build CSS for step type colors and layout."""
+ thought_color = colors.get("thought", DEFAULT_STEP_COLORS["thought"])
+ action_color = colors.get("action", DEFAULT_STEP_COLORS["action"])
+ obs_color = colors.get("observation", DEFAULT_STEP_COLORS["observation"])
+ system_color = colors.get("system", DEFAULT_STEP_COLORS["system"])
+ error_color = colors.get("error", DEFAULT_STEP_COLORS["error"])
+
+ padding = "8px 12px" if compact else "12px 16px"
+ margin = "4px 0" if compact else "8px 0"
+
+ return f'''
+ .agent-trace-display {{ font-family: inherit; }}
+ .agent-trace-summary {{
+ display: flex; align-items: center; gap: 8px;
+ padding: 8px 12px; margin-bottom: 12px;
+ background: #f8f9fa; border-radius: 6px; font-size: 0.9em;
+ }}
+ .summary-total {{ font-weight: 600; }}
+ .summary-badge {{
+ padding: 2px 8px; border-radius: 12px; font-size: 0.85em;
+ }}
+ .agent-trace-step {{
+ padding: {padding}; margin: {margin};
+ border-radius: 6px; border-left: 4px solid #ccc;
+ }}
+ .step-type-thought {{ background: {thought_color}; border-left-color: #2196F3; }}
+ .step-type-action {{ background: {action_color}; border-left-color: #FF9800; }}
+ .step-type-observation {{ background: {obs_color}; border-left-color: #4CAF50; }}
+ .step-type-system {{ background: {system_color}; border-left-color: #9C27B0; }}
+ .step-type-error {{ background: {error_color}; border-left-color: #f44336; }}
+ .step-header {{
+ display: flex; align-items: center; gap: 8px; margin-bottom: 4px;
+ }}
+ .step-badge {{
+ padding: 2px 8px; border-radius: 4px; font-size: 0.8em;
+ font-weight: 600; color: #333;
+ }}
+ .badge-thought {{ background: rgba(33,150,243,0.2); }}
+ .badge-action {{ background: rgba(255,152,0,0.2); }}
+ .badge-observation {{ background: rgba(76,175,80,0.2); }}
+ .badge-system {{ background: rgba(156,39,176,0.2); }}
+ .badge-error {{ background: rgba(244,67,54,0.2); }}
+ .step-number {{ color: #666; font-size: 0.8em; }}
+ .step-timestamp {{ color: #999; font-size: 0.75em; margin-left: auto; }}
+ .step-text {{ white-space: pre-wrap; word-break: break-word; line-height: 1.5; }}
+ .step-screenshot {{ margin-top: 8px; }}
+ .step-screenshot-img {{
+ max-width: 300px; max-height: 200px; border-radius: 4px;
+ border: 1px solid #ddd; cursor: pointer;
+ }}
+ .step-screenshot-img:hover {{ box-shadow: 0 2px 8px rgba(0,0,0,0.15); }}
+ .step-collapsible summary {{
+ cursor: pointer; color: #666; font-size: 0.9em; padding: 4px 0;
+ }}
+ '''
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ classes = super().get_css_classes(field_config)
+ if field_config.get("span_target"):
+ classes.append("span-target-field")
+ return classes
+
+ def get_data_attributes(self, field_config: Dict[str, Any], data: Any) -> Dict[str, str]:
+ attrs = super().get_data_attributes(field_config, data)
+ if field_config.get("span_target"):
+ attrs["span-target"] = "true"
+ return attrs
diff --git a/potato/server_utils/displays/audio_display.py b/potato/server_utils/displays/audio_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..67f2a723de63dd9231d28e7b16d9df419154653c
--- /dev/null
+++ b/potato/server_utils/displays/audio_display.py
@@ -0,0 +1,151 @@
+"""
+Audio Display Type
+
+Renders audio content for display in the annotation interface.
+Supports standard audio controls and links to audio annotation schemas.
+"""
+
+import html
+from typing import Dict, Any, List
+from urllib.parse import urlparse
+
+from .base import BaseDisplay
+
+
+class AudioDisplay(BaseDisplay):
+ """
+ Display type for audio content.
+
+ Displays audio with standard HTML5 audio player controls.
+ Can be linked to audio_annotation schemas via source_field.
+ """
+
+ name = "audio"
+ required_fields = ["key"]
+ optional_fields = {
+ "controls": True,
+ "autoplay": False,
+ "loop": False,
+ "muted": False,
+ "preload": "metadata",
+ "show_waveform": False,
+ }
+ description = "Audio player display"
+ supports_span_target = False
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render audio as HTML.
+
+ Args:
+ field_config: The field configuration
+ data: The audio URL or path
+
+ Returns:
+ HTML string for the audio display
+ """
+ if not data:
+ return 'No audio provided
'
+
+ # Get the audio URL
+ audio_url = str(data)
+
+ # Validate URL
+ if not self._is_valid_url(audio_url):
+ return f'Invalid audio URL: {html.escape(audio_url)}
'
+
+ # Get display options
+ options = self.get_display_options(field_config)
+ controls = options.get("controls", True)
+ autoplay = options.get("autoplay", False)
+ loop = options.get("loop", False)
+ muted = options.get("muted", False)
+ preload = options.get("preload", "metadata")
+ show_waveform = options.get("show_waveform", False)
+
+ # Build attributes
+ attrs = [f'preload="{preload}"']
+ if controls:
+ attrs.append("controls")
+ if autoplay:
+ attrs.append("autoplay")
+ if loop:
+ attrs.append("loop")
+ if muted:
+ attrs.append("muted")
+
+ attrs_str = " ".join(attrs)
+
+ # Escape values
+ escaped_url = html.escape(audio_url, quote=True)
+ field_key = html.escape(field_config.get("key", ""), quote=True)
+
+ # Determine audio type from URL
+ audio_type = self._get_audio_type(audio_url)
+ type_attr = f' type="{audio_type}"' if audio_type else ""
+
+ # Build HTML
+ waveform_html = ""
+ if show_waveform:
+ waveform_html = f'''
+
+
+
+ '''
+
+ return f'''
+
+ {waveform_html}
+
+
+ Your browser does not support the audio element.
+
+
+ '''
+
+ def _is_valid_url(self, url: str) -> bool:
+ """Check if URL is valid."""
+ if not url:
+ return False
+
+ if url.startswith('/') or url.startswith('./') or url.startswith('../'):
+ return True
+
+ try:
+ parsed = urlparse(url)
+ if parsed.scheme:
+ return parsed.scheme.lower() in ('http', 'https')
+ return True
+ except Exception:
+ return False
+
+ def _get_audio_type(self, url: str) -> str:
+ """Get audio MIME type from URL extension."""
+ url_lower = url.lower()
+ if url_lower.endswith('.mp3'):
+ return 'audio/mpeg'
+ elif url_lower.endswith('.wav'):
+ return 'audio/wav'
+ elif url_lower.endswith('.ogg') or url_lower.endswith('.oga'):
+ return 'audio/ogg'
+ elif url_lower.endswith('.m4a') or url_lower.endswith('.aac'):
+ return 'audio/aac'
+ elif url_lower.endswith('.webm'):
+ return 'audio/webm'
+ elif url_lower.endswith('.flac'):
+ return 'audio/flac'
+ return ""
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ """Get CSS classes for the container."""
+ classes = super().get_css_classes(field_config)
+ if self.get_display_options(field_config).get("show_waveform"):
+ classes.append("with-waveform")
+ return classes
+
+ def get_data_attributes(self, field_config: Dict[str, Any], data: Any) -> Dict[str, str]:
+ """Get data attributes for the container."""
+ attrs = super().get_data_attributes(field_config, data)
+ if data:
+ attrs["source-url"] = str(data)
+ return attrs
diff --git a/potato/server_utils/displays/base.py b/potato/server_utils/displays/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..faabb0095363836d2da52f77f7f924838960623d
--- /dev/null
+++ b/potato/server_utils/displays/base.py
@@ -0,0 +1,349 @@
+"""
+Base Display Class
+
+Provides a base class for custom display types that can be registered
+as plugins. Third-party developers can extend this class to create
+custom content renderers.
+
+Usage:
+ from potato.server_utils.displays.base import BaseDisplay
+
+ class MyCustomDisplay(BaseDisplay):
+ name = "my_custom"
+ required_fields = ["key"]
+ optional_fields = {"my_option": "default_value"}
+
+ def render(self, field_config, data):
+ return f'{data}
'
+
+Span Target Contract:
+ If a display declares ``supports_span_target = True``, its ``render()``
+ output MUST contain a ``.text-content`` wrapper when the field has
+ ``span_target: true``. Use the ``render_span_wrapper()`` helper:
+
+ if field_config.get("span_target"):
+ inner_html = self.render_span_wrapper(field_key, inner_html, plain_text)
+
+ SpanManager (span-core.js) discovers span-target fields via:
+ document.querySelectorAll('.display-field[data-span-target="true"]')
+ then looks inside each for:
+ field.querySelector('.text-content')
+ If ``.text-content`` is missing, span annotation silently fails.
+
+ See displays/ARCHITECTURE.md for the full contract.
+"""
+
+import html as html_module
+from abc import ABC, abstractmethod
+from typing import Dict, Any, List, Optional
+
+
+class BaseDisplay(ABC):
+ """
+ Abstract base class for display type implementations.
+
+ Subclasses must implement the `render` method and define
+ class attributes for registration.
+
+ Class Attributes:
+ name: Unique identifier for this display type (e.g., "text", "image")
+ required_fields: List of required configuration field names
+ optional_fields: Dictionary of optional fields with their default values
+ description: Human-readable description of this display type
+ supports_span_target: Whether this type can be a span annotation target.
+ If True, render() MUST produce a .text-content wrapper when
+ field_config["span_target"] is True. Use render_span_wrapper().
+ lazy_populated: Whether this display's data field is populated after
+ initial page load (e.g. ``interactive_chat`` writes its
+ conversation only when the user finishes chatting with the
+ agent_proxy). Set True to tell the ``instance_display`` validator
+ that a missing data key for this field is an expected transient
+ state, not a configuration error.
+ """
+
+ name: str = ""
+ required_fields: List[str] = []
+ optional_fields: Dict[str, Any] = {}
+ description: str = ""
+ supports_span_target: bool = False
+ lazy_populated: bool = False
+
+ @abstractmethod
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render the content as HTML.
+
+ Args:
+ field_config: The field configuration from instance_display.fields
+ data: The actual data value from the instance
+
+ Returns:
+ HTML string for rendering the content
+ """
+ pass
+
+ def render_span_wrapper(self, field_key: str, inner_html: str, plain_text: str) -> str:
+ """
+ Wrap content in the standard .text-content div required by SpanManager.
+
+ Call this from render() when field_config.get("span_target") is True.
+ This ensures the HTML output satisfies the span annotation contract:
+ - ``class="text-content"``
+ - ``id="text-content-{field_key}"``
+ - ``data-original-text="{escaped plain text}"``
+ - ``padding-top: 24px`` (via style) for span label positioning
+
+ Args:
+ field_key: The field key (e.g., "conversation", "premise")
+ inner_html: The HTML content to wrap
+ plain_text: The plain text for offset-based span positioning.
+ Must match the text extraction format in routes.py for this
+ data type, or span offsets will misalign on reload.
+
+ Returns:
+ HTML string with the .text-content wrapper
+ """
+ escaped_text = html_module.escape(plain_text, quote=True)
+ return (
+ f''
+ f'{inner_html}'
+ f'
'
+ )
+
+ def has_inline_label(self, field_config: Dict[str, Any]) -> bool:
+ """
+ Check if this display handles its own label rendering.
+
+ If True, the registry will NOT add a label wrapper around the
+ display container (avoiding duplicate labels).
+
+ Override in subclasses where the display renders its own label
+ (e.g., collapsible text with label in ).
+
+ Args:
+ field_config: The field configuration
+
+ Returns:
+ True if the display renders its own label
+ """
+ return False
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ """
+ Get CSS classes to apply to the display container.
+
+ Override in subclasses to add custom classes.
+
+ Args:
+ field_config: The field configuration
+
+ Returns:
+ List of CSS class names
+ """
+ return [f"display-field", f"display-type-{self.name}"]
+
+ def get_data_attributes(self, field_config: Dict[str, Any], data: Any) -> Dict[str, str]:
+ """
+ Get data attributes to add to the display container.
+
+ These are used for JavaScript interactions and linking
+ annotation schemas to display fields.
+
+ Args:
+ field_config: The field configuration
+ data: The actual data value
+
+ Returns:
+ Dictionary of data attribute names (without 'data-' prefix) to values
+ """
+ attrs = {
+ "field-key": field_config.get("key", ""),
+ "field-type": self.name,
+ }
+ if field_config.get("span_target"):
+ attrs["span-target"] = "true"
+ return attrs
+
+ def get_js_init(self) -> Optional[str]:
+ """
+ Get JavaScript initialization code for this display type.
+
+ Override in subclasses that need client-side initialization.
+
+ Returns:
+ JavaScript code string or None if not needed
+ """
+ return None
+
+ def validate_config(self, field_config: Dict[str, Any]) -> List[str]:
+ """
+ Validate the field configuration.
+
+ Args:
+ field_config: The field configuration to validate
+
+ Returns:
+ List of error messages (empty if valid)
+ """
+ errors = []
+
+ # Check required fields
+ for field in self.required_fields:
+ if field not in field_config:
+ errors.append(f"Missing required field '{field}' for display type '{self.name}'")
+
+ # Warn if span_target is set but display doesn't support it
+ if field_config.get("span_target") and not self.supports_span_target:
+ errors.append(
+ f"Display type '{self.name}' does not support span_target. "
+ f"Span annotation will not work on this field."
+ )
+
+ return errors
+
+ def get_display_options(self, field_config: Dict[str, Any]) -> Dict[str, Any]:
+ """
+ Get display options with defaults applied.
+
+ Args:
+ field_config: The field configuration
+
+ Returns:
+ Dictionary of display options with defaults filled in
+ """
+ options = field_config.get("display_options", {})
+ result = dict(self.optional_fields) # Start with defaults
+ result.update(options) # Override with user-specified options
+ return result
+
+
+def concatenate_dialogue_text(data: Any, speaker_key: str = "speaker", text_key: str = "text") -> str:
+ """
+ Concatenate dialogue data into a single plain text string.
+
+ Format: "Speaker: text\\nSpeaker: text\\n..."
+ Turns without a speaker omit the "Speaker: " prefix.
+
+ Note: For span offset matching with dialogue displays, use
+ ``reconstruct_dialogue_dom_text()`` instead โ it accounts for
+ turn numbers and DOM whitespace normalization.
+
+ Args:
+ data: Dialogue data โ list of dicts, list of strings, or a string
+ speaker_key: Key for speaker in dict format (default "speaker")
+ text_key: Key for text in dict format (default "text")
+
+ Returns:
+ Concatenated plain text string
+ """
+ if isinstance(data, str):
+ return data
+
+ if not isinstance(data, list):
+ return str(data)
+
+ parts = []
+ for item in data:
+ if isinstance(item, dict):
+ speaker = item.get(speaker_key, '')
+ text = item.get(text_key, '')
+ parts.append(f"{speaker}: {text}" if speaker else text)
+ else:
+ parts.append(str(item))
+ return "\n".join(parts)
+
+
+def reconstruct_dialogue_dom_text(
+ data: Any,
+ speaker_key: str = "speaker",
+ text_key: str = "text",
+ show_turn_numbers: bool = False,
+) -> str:
+ """
+ Reconstruct the whitespace-normalized DOM textContent of a dialogue display.
+
+ When DialogueDisplay renders HTML, the browser's ``textContent`` includes
+ turn numbers, speaker prefixes, and the text of each turn โ all separated
+ by whitespace that ``normalizeText()`` collapses to single spaces.
+
+ This function reproduces that collapsed form so that span offsets produced
+ by the client (DOM-based) can be used server-side to extract the correct
+ substring.
+
+ Args:
+ data: Dialogue data โ list of dicts, list of strings, or a string
+ speaker_key: Key for speaker in dict format
+ text_key: Key for text in dict format
+ show_turn_numbers: Whether turn numbers like ``[1]`` are shown
+
+ Returns:
+ Single-line text matching the browser's normalized textContent
+ """
+ if isinstance(data, str):
+ return data.strip()
+
+ if not isinstance(data, list):
+ return str(data).strip()
+
+ parts = []
+ for i, item in enumerate(data):
+ turn_parts = []
+ if show_turn_numbers:
+ turn_parts.append(f"[{i + 1}]")
+
+ if isinstance(item, dict):
+ speaker = item.get(speaker_key, "")
+ text = item.get(text_key, "")
+ if speaker:
+ turn_parts.append(f"{speaker}:")
+ turn_parts.append(str(text))
+ else:
+ turn_parts.append(str(item))
+
+ parts.append(" ".join(turn_parts))
+
+ # Join turns with single space (browser normalizes inter-turn whitespace)
+ import re as _re
+ joined = " ".join(parts)
+ # Final normalization: collapse any remaining multi-space to single space
+ return _re.sub(r"\s+", " ", joined).strip()
+
+
+def render_display_container(
+ inner_html: str,
+ css_classes: List[str],
+ data_attrs: Dict[str, str],
+ label: Optional[str] = None
+) -> str:
+ """
+ Render a display container with the given content.
+
+ This helper function wraps content in a standard display container
+ with proper classes and data attributes.
+
+ Args:
+ inner_html: The inner HTML content
+ css_classes: CSS classes to apply
+ data_attrs: Data attributes to add
+ label: Optional label/header for the display
+
+ Returns:
+ Complete HTML for the display container
+ """
+ class_str = " ".join(css_classes)
+ attr_str = " ".join(f'data-{k}="{v}"' for k, v in data_attrs.items())
+
+ parts = []
+ parts.append(f'')
+
+ if label:
+ parts.append(f'
{label}
')
+
+ parts.append(f'
')
+ parts.append(f' {inner_html}')
+ parts.append(f'
')
+ parts.append(f'
')
+
+ return "\n".join(parts)
diff --git a/potato/server_utils/displays/code_display.py b/potato/server_utils/displays/code_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e77954fe82d3ce5183d5356e5259a4a2afe78bc
--- /dev/null
+++ b/potato/server_utils/displays/code_display.py
@@ -0,0 +1,268 @@
+"""
+Code Display Component
+
+Renders source code with syntax highlighting and line number support.
+
+Usage:
+ In instance_display config:
+ fields:
+ - key: source_code
+ type: code
+ display_options:
+ language: python
+ show_line_numbers: true
+ max_height: 500
+"""
+
+from typing import Dict, Any, List, Optional
+import html
+import logging
+
+from .base import BaseDisplay
+
+logger = logging.getLogger(__name__)
+
+
+class CodeDisplay(BaseDisplay):
+ """
+ Display type for source code with syntax highlighting.
+
+ Uses Pygments for server-side highlighting or can use client-side
+ highlighting with highlight.js.
+ """
+
+ name = "code"
+ required_fields = ["key"]
+ optional_fields = {
+ "language": None, # Language for syntax highlighting
+ "show_line_numbers": True, # Show line numbers
+ "max_height": 500, # Max container height
+ "max_width": None, # Max container width
+ "wrap_lines": False, # Wrap long lines
+ "highlight_lines": None, # List of line numbers to highlight
+ "start_line": 1, # Starting line number
+ "theme": "default", # Color theme
+ "copy_button": True, # Show copy to clipboard button
+ }
+ description = "Source code display with syntax highlighting"
+ supports_span_target = True
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render source code content.
+
+ Args:
+ field_config: Display configuration
+ data: Either a dict with extracted content or raw code string
+
+ Returns:
+ HTML string for rendering
+ """
+ options = self.get_display_options(field_config)
+ field_key = field_config.get("key", "code")
+
+ # Handle different data formats
+ if isinstance(data, dict):
+ # Pre-extracted FormatOutput data
+ if "rendered_html" in data:
+ return self._wrap_content(data["rendered_html"], options, field_key)
+
+ code = data.get("text", "")
+ language = data.get("metadata", {}).get("language", options.get("language"))
+ elif isinstance(data, str):
+ code = data
+ language = options.get("language")
+ else:
+ return f'Unsupported content type
'
+
+ # Generate code HTML
+ is_span_target = field_config.get("span_target", False)
+ code_html = self._render_code(code, language, options, field_key, is_span_target)
+ return self._wrap_content(code_html, options, field_key)
+
+ def _wrap_content(
+ self,
+ content: str,
+ options: Dict[str, Any],
+ field_key: str
+ ) -> str:
+ """
+ Wrap code content in container with styles.
+ """
+ styles = []
+ max_height = options.get("max_height")
+ max_width = options.get("max_width")
+
+ if max_height:
+ styles.append(f"max-height: {max_height}px")
+ styles.append("overflow-y: auto")
+ if max_width:
+ styles.append(f"max-width: {max_width}px")
+ styles.append("overflow-x: auto")
+
+ style_str = "; ".join(styles) if styles else ""
+ language = options.get("language") or "text"
+ theme = options.get("theme", "default")
+
+ return f'''
+
+ {self._render_copy_button() if options.get("copy_button") else ""}
+ {content}
+
+ '''
+
+ def _render_copy_button(self) -> str:
+ """
+ Render copy to clipboard button.
+ """
+ return '''
+
+
+
+
+
+ '''
+
+ def _render_code(
+ self,
+ code: str,
+ language: Optional[str],
+ options: Dict[str, Any],
+ field_key: str = "code",
+ is_span_target: bool = False
+ ) -> str:
+ """
+ Render code with line numbers and optional highlighting.
+
+ When is_span_target is True, uses a simpler pre/code structure
+ that works better with span annotation positioning.
+ """
+ # For span targets, use simpler rendering for accurate position calculation
+ if is_span_target:
+ return self._render_code_simple(code, language, options, field_key)
+
+ lines = code.split("\n")
+ parts = []
+
+ show_line_numbers = options.get("show_line_numbers", True)
+ start_line = options.get("start_line", 1)
+ highlight_lines = set(options.get("highlight_lines") or [])
+ wrap_lines = options.get("wrap_lines", False)
+
+ wrap_class = "code-wrap" if wrap_lines else "code-nowrap"
+ lang_class = f"language-{language}" if language else ""
+
+ content_classes = ["code-content", wrap_class]
+ content_class_str = " ".join(content_classes)
+
+ parts.append(f'')
+ parts.append(f'
')
+
+ for i, line in enumerate(lines):
+ line_num = i + start_line
+ line_classes = ["code-line"]
+
+ if line_num in highlight_lines:
+ line_classes.append("highlighted-line")
+
+ escaped_line = html.escape(line) if line else " "
+
+ parts.append(f'')
+
+ if show_line_numbers:
+ parts.append(
+ f'{line_num} '
+ )
+
+ parts.append(
+ f''
+ f'{escaped_line}'
+ f' '
+ )
+
+ parts.append(' ')
+
+ parts.append('
')
+ parts.append('
')
+
+ return "\n".join(parts)
+
+ def _render_code_simple(
+ self,
+ code: str,
+ language: Optional[str],
+ options: Dict[str, Any],
+ field_key: str
+ ) -> str:
+ """
+ Render code in a simple pre/code format for span annotation.
+
+ This avoids table structures that interfere with text position calculations.
+ Uses a flat text structure without data-original-text so the span system
+ uses DOM textContent directly (which matches what the user sees and selects).
+ """
+ lang_class = f"language-{language}" if language else ""
+
+ # Escape code for HTML but preserve structure
+ escaped_code = html.escape(code)
+
+ # Don't use data-original-text - let span system use DOM textContent
+ # This ensures canonical text matches what user sees and selects
+ return f'''
+
+ '''
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ """Get CSS classes for the display container."""
+ classes = super().get_css_classes(field_config)
+ options = self.get_display_options(field_config)
+
+ if field_config.get("span_target"):
+ classes.append("span-target-code")
+
+ language = options.get("language")
+ if language:
+ classes.append(f"language-{language}")
+
+ theme = options.get("theme", "default")
+ classes.append(f"code-theme-{theme}")
+
+ return classes
+
+ def get_data_attributes(
+ self,
+ field_config: Dict[str, Any],
+ data: Any
+ ) -> Dict[str, str]:
+ """Get data attributes for JavaScript initialization."""
+ attrs = super().get_data_attributes(field_config, data)
+ options = self.get_display_options(field_config)
+
+ language = options.get("language")
+ if language:
+ attrs["language"] = language
+
+ if isinstance(data, dict) and "metadata" in data:
+ meta = data["metadata"]
+ if "language" in meta:
+ attrs["language"] = meta["language"]
+ if "line_count" in meta:
+ attrs["line-count"] = str(meta["line_count"])
+
+ return attrs
+
+ def get_js_init(self) -> Optional[str]:
+ """
+ Return JavaScript initialization code for code displays.
+ """
+ return '''
+ if (typeof initCodeDisplays === 'function') {
+ initCodeDisplays();
+ }
+ '''
diff --git a/potato/server_utils/displays/coding_trace_display.py b/potato/server_utils/displays/coding_trace_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..fae01edc85f6b3093c12883d14032e35dd53e269
--- /dev/null
+++ b/potato/server_utils/displays/coding_trace_display.py
@@ -0,0 +1,751 @@
+"""
+Coding Trace Display Type
+
+Purpose-built rendering for agentic coding system traces (Claude Code, OpenCode,
+Cursor, Aider, SWE-Agent). Renders tool calls with appropriate formatting:
+- Code diffs (Edit/Write) with red/green highlighting
+- Terminal blocks (Bash) with dark monospace styling
+- Code blocks (Read/Grep/Glob) with line numbers
+- File tree sidebar showing all files touched
+- Collapsible long outputs
+- Turn structure: User messages โ assistant reasoning โ tool calls
+
+Usage:
+ In instance_display config:
+ fields:
+ - key: structured_turns
+ type: coding_trace
+ display_options:
+ show_file_tree: true
+ diff_view: unified
+ collapse_long_outputs: true
+ max_output_lines: 50
+ terminal_theme: dark
+"""
+
+import html
+import json
+import os
+import re
+from typing import Dict, Any, List, Optional, Set, Tuple
+
+from .base import BaseDisplay
+
+
+# Tool type classifications - order matters: more specific sets checked first
+CODE_READ_TOOLS = {"Read", "read"}
+CODE_EDIT_TOOLS = {"Edit", "edit", "Replace", "replace"}
+CODE_WRITE_TOOLS = {"Write", "write", "Create", "create"}
+TERMINAL_TOOLS = {"Bash", "bash", "Terminal", "terminal", "Shell", "shell", "Run", "run"}
+SEARCH_TOOLS = {"Grep", "grep", "Glob", "glob", "Search", "search", "Find", "find"}
+
+# File extension to language mapping
+EXTENSION_LANGUAGES = {
+ ".py": "python", ".js": "javascript", ".ts": "typescript",
+ ".jsx": "jsx", ".tsx": "tsx", ".rb": "ruby", ".go": "go",
+ ".rs": "rust", ".java": "java", ".c": "c", ".cpp": "cpp",
+ ".h": "c", ".hpp": "cpp", ".cs": "csharp", ".swift": "swift",
+ ".kt": "kotlin", ".scala": "scala", ".r": "r",
+ ".sh": "bash", ".bash": "bash", ".zsh": "zsh",
+ ".yaml": "yaml", ".yml": "yaml", ".json": "json",
+ ".xml": "xml", ".html": "html", ".css": "css",
+ ".sql": "sql", ".md": "markdown", ".toml": "toml",
+ ".ini": "ini", ".cfg": "ini", ".dockerfile": "dockerfile",
+}
+
+# Badge colors for different tool types
+TOOL_BADGE_COLORS = {
+ "read": ("#e3f2fd", "#1565c0", "#1976d2"), # Blue
+ "edit": ("#fff3e0", "#e65100", "#ef6c00"), # Orange
+ "write": ("#e8f5e9", "#2e7d32", "#388e3c"), # Green
+ "bash": ("#263238", "#b0bec5", "#78909c"), # Dark
+ "search": ("#f3e5f5", "#6a1b9a", "#7b1fa2"), # Purple
+ "generic": ("#f5f5f5", "#424242", "#616161"), # Grey
+}
+
+
+def _detect_language(file_path: str) -> str:
+ """Detect programming language from file extension."""
+ if not file_path:
+ return ""
+ ext = os.path.splitext(file_path)[1].lower()
+ return EXTENSION_LANGUAGES.get(ext, "")
+
+
+def _classify_tool(tool_name: str) -> str:
+ """Classify a tool into a rendering category."""
+ if tool_name in SEARCH_TOOLS:
+ return "search"
+ if tool_name in CODE_READ_TOOLS:
+ return "read"
+ if tool_name in CODE_EDIT_TOOLS:
+ return "edit"
+ if tool_name in CODE_WRITE_TOOLS:
+ return "write"
+ if tool_name in TERMINAL_TOOLS:
+ return "bash"
+ return "generic"
+
+
+def _escape(text: str) -> str:
+ """HTML-escape text."""
+ return html.escape(str(text), quote=True)
+
+
+def _truncate_output(text: str, max_lines: int) -> Tuple[str, bool]:
+ """Truncate text to max_lines, return (text, was_truncated)."""
+ if not text or max_lines <= 0:
+ return text, False
+ lines = text.split("\n")
+ if len(lines) <= max_lines:
+ return text, False
+ return "\n".join(lines[:max_lines]), True
+
+
+class CodingTraceDisplay(BaseDisplay):
+ """
+ Display type for coding agent traces with rich tool call rendering.
+
+ Renders agent sessions with proper formatting for code diffs,
+ terminal output, file reads, and search results.
+ """
+
+ name = "coding_trace"
+ required_fields = ["key"]
+ optional_fields = {
+ "show_file_tree": True,
+ "diff_view": "unified", # "unified" or "side_by_side"
+ "collapse_long_outputs": True,
+ "max_output_lines": 50,
+ "terminal_theme": "dark", # "dark" or "light"
+ "show_step_numbers": True,
+ "show_tool_badges": True,
+ "show_reasoning": True,
+ "compact": False,
+ }
+ description = "Coding agent trace display with diff rendering, terminal blocks, and file tree"
+ supports_span_target = True
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ if not data:
+ return 'No trace data provided
'
+
+ options = self.get_display_options(field_config)
+ field_key = _escape(field_config.get("key", ""))
+ is_span_target = field_config.get("span_target", False)
+
+ # Parse turns from data
+ turns = self._normalize_turns(data)
+ if not turns:
+ empty_html = 'No trace steps found
'
+ if is_span_target:
+ return self.render_span_wrapper(field_key, empty_html, "")
+ return empty_html
+
+ # Build file tree
+ file_tree_html = ""
+ file_count = 0
+ if options.get("show_file_tree", True):
+ file_tree_html, file_count = self._build_file_tree(turns)
+
+ # Build turn cards
+ turns_html = self._build_turns(turns, options, field_key, is_span_target)
+
+ # Build summary with collapse/expand toggle
+ total_tools = sum(len(t.get("tool_calls", [])) for t in turns)
+ summary_html = self._build_summary(turns, file_tree_html, file_count)
+
+ # Wrap in span target if needed
+ if is_span_target:
+ plain_text = self._extract_plain_text(turns)
+ reasoning_html = self._extract_reasoning_html(turns)
+ inner = reasoning_html
+ span_wrapper = self.render_span_wrapper(field_key, inner, plain_text)
+ else:
+ span_wrapper = ""
+
+ layout_class = "coding-trace-with-sidebar" if file_tree_html else ""
+ # Auto-collapse sidebar when only 1 file
+ sidebar_class = "ct-sidebar-collapsed" if file_count <= 1 and file_tree_html else ""
+
+ # Client-side JS for collapse/expand and sidebar toggle
+ js_init = self._build_js_init(field_key)
+
+ return f'''
+
+ {summary_html}
+
+ {f'' if file_tree_html else ''}
+
+ {span_wrapper}
+ {turns_html}
+
+
+
+
+ '''
+
+ def _normalize_turns(self, data: Any) -> List[Dict[str, Any]]:
+ """Normalize various input formats to a list of structured turns."""
+ if isinstance(data, list):
+ turns = []
+ for item in data:
+ if isinstance(item, dict):
+ turns.append(self._normalize_single_turn(item))
+ elif isinstance(item, str):
+ turns.append({"role": "user", "content": item, "tool_calls": []})
+ return turns
+ if isinstance(data, dict):
+ # Single turn
+ return [self._normalize_single_turn(data)]
+ return []
+
+ def _normalize_single_turn(self, item: Dict[str, Any]) -> Dict[str, Any]:
+ """Normalize a single turn dict."""
+ role = item.get("role", "assistant")
+ reasoning = item.get("reasoning", item.get("content", item.get("text", "")))
+ tool_calls = item.get("tool_calls", [])
+
+ # Handle string content (user messages)
+ if isinstance(reasoning, list):
+ # Content blocks format
+ text_parts = []
+ extracted_tools = []
+ for block in reasoning:
+ if isinstance(block, dict):
+ if block.get("type") == "text":
+ text_parts.append(block.get("text", ""))
+ elif block.get("type") == "tool_use":
+ extracted_tools.append({
+ "tool": block.get("name", "unknown"),
+ "input": block.get("input", {}),
+ "output": "",
+ "output_type": "generic",
+ })
+ elif isinstance(block, str):
+ text_parts.append(block)
+ reasoning = "\n".join(text_parts)
+ if extracted_tools and not tool_calls:
+ tool_calls = extracted_tools
+
+ return {
+ "role": role,
+ "content": str(reasoning) if reasoning else "",
+ "tool_calls": tool_calls,
+ }
+
+ def _build_summary(self, turns: List[Dict[str, Any]],
+ file_tree_html: str = "", file_count: int = 0) -> str:
+ """Build a summary header with counts and toggle buttons."""
+ total_tools = sum(len(t.get("tool_calls", [])) for t in turns)
+ assistant_turns = sum(1 for t in turns if t.get("role") != "user")
+
+ # Count tool types
+ tool_counts: Dict[str, int] = {}
+ for turn in turns:
+ for tc in turn.get("tool_calls", []):
+ tool_type = _classify_tool(tc.get("tool", ""))
+ tool_counts[tool_type] = tool_counts.get(tool_type, 0) + 1
+
+ badges = []
+ for tool_type, count in sorted(tool_counts.items()):
+ bg, fg, _ = TOOL_BADGE_COLORS.get(tool_type, TOOL_BADGE_COLORS["generic"])
+ badges.append(
+ f''
+ f'{count} {tool_type} '
+ )
+
+ # Toggle buttons
+ sidebar_toggle = ""
+ if file_tree_html:
+ sidebar_toggle = (
+ ''
+ )
+
+ collapse_toggle = ""
+ if total_tools > 2:
+ collapse_toggle = (
+ 'Collapse outputs '
+ )
+
+ return f'''
+
+ {assistant_turns} turn{"s" if assistant_turns != 1 else ""}
+ ·
+ {total_tools} tool call{"s" if total_tools != 1 else ""}
+ {" ".join(badges)}
+ {sidebar_toggle}
+ {collapse_toggle}
+
+ '''
+
+ def _build_turns(self, turns: List[Dict[str, Any]], options: Dict[str, Any],
+ field_key: str, is_span_target: bool) -> str:
+ """Build HTML for all turns."""
+ parts = []
+ step_counter = 0
+ show_numbers = options.get("show_step_numbers", True)
+ max_lines = options.get("max_output_lines", 50)
+ collapse = options.get("collapse_long_outputs", True)
+ show_reasoning = options.get("show_reasoning", True)
+
+ for i, turn in enumerate(turns):
+ role = turn.get("role", "assistant")
+
+ if role == "user":
+ parts.append(self._render_user_message(turn, i))
+ continue
+
+ step_counter += 1
+
+ # Assistant turn
+ turn_parts = []
+
+ # Step header
+ if show_numbers:
+ turn_parts.append(
+ f''
+ )
+
+ # Reasoning text
+ content = turn.get("content", "")
+ if content and show_reasoning:
+ escaped = _escape(content)
+ turn_parts.append(
+ f'{escaped}
'
+ )
+
+ # Tool calls
+ tool_calls = turn.get("tool_calls", [])
+ for j, tc in enumerate(tool_calls):
+ tc_html = self._render_tool_call(tc, options, f"{i}-{j}")
+ turn_parts.append(tc_html)
+
+ parts.append(
+ f''
+ f'{"".join(turn_parts)}'
+ f'
'
+ )
+
+ return "\n".join(parts)
+
+ def _render_user_message(self, turn: Dict[str, Any], index: int) -> str:
+ """Render a user message bubble."""
+ content = _escape(turn.get("content", ""))
+ return (
+ f''
+ f'
User
'
+ f'
{content}
'
+ f'
'
+ )
+
+ def _render_tool_call(self, tc: Dict[str, Any], options: Dict[str, Any],
+ tc_id: str) -> str:
+ """Render a single tool call with appropriate formatting."""
+ tool_name = tc.get("tool", "unknown")
+ tool_input = tc.get("input", {})
+ tool_output = tc.get("output", "")
+ output_type = tc.get("output_type", "")
+ tool_type = _classify_tool(tool_name)
+
+ # If output_type not specified, infer from tool
+ if not output_type:
+ output_type = tool_type
+
+ # Badge
+ bg, fg, border = TOOL_BADGE_COLORS.get(tool_type, TOOL_BADGE_COLORS["generic"])
+ badge = (
+ f'{_escape(tool_name)} '
+ )
+
+ # File path header (if applicable)
+ file_path = ""
+ if isinstance(tool_input, dict):
+ file_path = tool_input.get("file_path", tool_input.get("path", ""))
+
+ file_header = ""
+ if file_path:
+ file_header = f'{_escape(file_path)} '
+
+ # Render input/output based on tool type
+ if tool_type == "edit":
+ body = self._render_diff(tool_input, tool_output, options)
+ elif tool_type == "bash":
+ body = self._render_terminal(tool_input, tool_output, options)
+ elif tool_type in ("read", "search"):
+ body = self._render_code_output(tool_input, tool_output, options)
+ elif tool_type == "write":
+ body = self._render_write(tool_input, tool_output, options)
+ else:
+ body = self._render_generic(tool_input, tool_output, options)
+
+ return (
+ f''
+ )
+
+ def _render_diff(self, tool_input: Any, tool_output: Any,
+ options: Dict[str, Any]) -> str:
+ """Render an edit as a unified diff."""
+ if not isinstance(tool_input, dict):
+ return self._render_generic(tool_input, tool_output, options)
+
+ old_string = tool_input.get("old_string", "")
+ new_string = tool_input.get("new_string", "")
+ file_path = tool_input.get("file_path", "")
+
+ if not old_string and not new_string:
+ return self._render_generic(tool_input, tool_output, options)
+
+ # Build unified diff view
+ old_lines = old_string.split("\n") if old_string else []
+ new_lines = new_string.split("\n") if new_string else []
+
+ diff_parts = []
+ diff_parts.append('')
+
+ # Removed lines
+ for line in old_lines:
+ escaped = _escape(line)
+ diff_parts.append(
+ f'
'
+ f'- '
+ f'{escaped} '
+ f'
'
+ )
+
+ # Added lines
+ for line in new_lines:
+ escaped = _escape(line)
+ diff_parts.append(
+ f'
'
+ f'+ '
+ f'{escaped} '
+ f'
'
+ )
+
+ diff_parts.append('
')
+
+ # Status message
+ status = ""
+ if tool_output:
+ output_str = str(tool_output)
+ if output_str:
+ status = f'{_escape(output_str)}
'
+
+ return "\n".join(diff_parts) + status
+
+ def _render_terminal(self, tool_input: Any, tool_output: Any,
+ options: Dict[str, Any]) -> str:
+ """Render a terminal command and its output."""
+ command = ""
+ if isinstance(tool_input, dict):
+ command = tool_input.get("command", tool_input.get("cmd", ""))
+ elif isinstance(tool_input, str):
+ command = tool_input
+
+ output_str = str(tool_output) if tool_output else ""
+ max_lines = options.get("max_output_lines", 50)
+ collapse = options.get("collapse_long_outputs", True)
+
+ parts = []
+ parts.append('')
+
+ # Command line
+ if command:
+ parts.append(
+ f'
'
+ f'$ '
+ f'{_escape(command)}'
+ f'
'
+ )
+
+ # Output
+ if output_str:
+ truncated, was_truncated = _truncate_output(output_str, max_lines)
+ if was_truncated and collapse:
+ parts.append(
+ f'
'
+ f'Output ({len(output_str.splitlines())} lines โ click to expand) '
+ f'{_escape(output_str)} '
+ f' '
+ f'
{_escape(truncated)} '
+ )
+ else:
+ parts.append(
+ f'
{_escape(output_str)} '
+ )
+
+ parts.append('
')
+ return "\n".join(parts)
+
+ def _render_code_output(self, tool_input: Any, tool_output: Any,
+ options: Dict[str, Any]) -> str:
+ """Render a code read/search result with line numbers."""
+ file_path = ""
+ if isinstance(tool_input, dict):
+ file_path = tool_input.get("file_path", tool_input.get("path", ""))
+
+ output_str = str(tool_output) if tool_output else ""
+ language = _detect_language(file_path)
+ max_lines = options.get("max_output_lines", 50)
+ collapse = options.get("collapse_long_outputs", True)
+
+ if not output_str:
+ return 'No output
'
+
+ lines = output_str.split("\n")
+ truncated, was_truncated = _truncate_output(output_str, max_lines)
+
+ # Build line-numbered code block
+ display_lines = truncated.split("\n") if was_truncated and collapse else lines
+
+ code_parts = []
+ code_parts.append(f'')
+ code_parts.append('
')
+
+ for i, line in enumerate(display_lines, 1):
+ escaped = _escape(line) if line else " "
+ code_parts.append(
+ f''
+ f'{i} '
+ f'{escaped} '
+ f' '
+ )
+
+ code_parts.append('
')
+
+ if was_truncated and collapse:
+ remaining = len(lines) - max_lines
+ code_parts.append(
+ f'
'
+ f'... {remaining} more line{"s" if remaining != 1 else ""}'
+ f'
'
+ )
+
+ code_parts.append('
')
+ return "\n".join(code_parts)
+
+ def _render_write(self, tool_input: Any, tool_output: Any,
+ options: Dict[str, Any]) -> str:
+ """Render a file write operation."""
+ content = ""
+ file_path = ""
+ if isinstance(tool_input, dict):
+ content = tool_input.get("content", "")
+ file_path = tool_input.get("file_path", "")
+
+ if not content:
+ return self._render_generic(tool_input, tool_output, options)
+
+ language = _detect_language(file_path)
+ max_lines = options.get("max_output_lines", 50)
+
+ # Show as code block with "new file" styling
+ lines = content.split("\n")
+ truncated, was_truncated = _truncate_output(content, max_lines)
+ display_lines = truncated.split("\n") if was_truncated else lines
+
+ parts = []
+ parts.append(f'')
+ parts.append('
New file
')
+ parts.append('
')
+
+ for i, line in enumerate(display_lines, 1):
+ escaped = _escape(line) if line else " "
+ parts.append(
+ f''
+ f'{i} '
+ f'{escaped} '
+ f' '
+ )
+
+ parts.append('
')
+
+ if was_truncated:
+ remaining = len(lines) - max_lines
+ parts.append(
+ f'
'
+ f'... {remaining} more line{"s" if remaining != 1 else ""}'
+ f'
'
+ )
+
+ parts.append('
')
+
+ # Status
+ if tool_output:
+ parts.append(f'{_escape(str(tool_output))}
')
+
+ return "\n".join(parts)
+
+ def _render_generic(self, tool_input: Any, tool_output: Any,
+ options: Dict[str, Any]) -> str:
+ """Render a generic tool call as formatted JSON."""
+ parts = []
+
+ if tool_input:
+ try:
+ if isinstance(tool_input, dict):
+ formatted = json.dumps(tool_input, indent=2, ensure_ascii=False)
+ else:
+ formatted = str(tool_input)
+ except (TypeError, ValueError):
+ formatted = str(tool_input)
+ parts.append(
+ f''
+ )
+
+ if tool_output:
+ output_str = str(tool_output)
+ max_lines = options.get("max_output_lines", 50)
+ truncated, was_truncated = _truncate_output(output_str, max_lines)
+
+ truncated_div = '... output truncated
' if was_truncated else ""
+ parts.append(
+ f''
+ f'
Output
'
+ f'
{_escape(truncated if was_truncated else output_str)} '
+ f'{truncated_div}'
+ f'
'
+ )
+
+ return "\n".join(parts) if parts else 'No data
'
+
+ def _build_file_tree(self, turns: List[Dict[str, Any]]) -> Tuple[str, int]:
+ """Build a file tree sidebar from all tool calls.
+
+ Returns:
+ Tuple of (html_string, file_count).
+ """
+ files: Dict[str, Set[str]] = {} # path -> set of operations
+
+ for turn in turns:
+ for tc in turn.get("tool_calls", []):
+ tool_name = tc.get("tool", "")
+ tool_type = _classify_tool(tool_name)
+ tool_input = tc.get("input", {})
+
+ if isinstance(tool_input, dict):
+ file_path = tool_input.get("file_path", tool_input.get("path", ""))
+ if file_path:
+ if file_path not in files:
+ files[file_path] = set()
+ files[file_path].add(tool_type)
+
+ if not files:
+ return "", 0
+
+ # Operation badge colors (domain-specific, kept as-is)
+ op_colors = {
+ "read": "#1976d2",
+ "edit": "#ef6c00",
+ "write": "#388e3c",
+ "search": "#7b1fa2",
+ "bash": "#78909c",
+ }
+
+ parts = []
+ parts.append('')
+ parts.append('')
+ parts.append('
')
+
+ for path in sorted(files.keys()):
+ ops = files[path]
+ # Pick the most significant operation for the name color
+ op = "write" if "write" in ops else "edit" if "edit" in ops else "read"
+ color = op_colors.get(op, "#666")
+
+ basename = os.path.basename(path) or path
+
+ op_badges = " ".join(
+ f''
+ f'{o[0].upper()} '
+ for o in sorted(ops)
+ )
+
+ parts.append(
+ f''
+ f'{_escape(basename)} '
+ f'{op_badges}'
+ f' '
+ )
+
+ parts.append(' ')
+ parts.append('
')
+ return "\n".join(parts), len(files)
+
+ def _build_js_init(self, field_key: str) -> str:
+ """Build client-side JS for collapse/expand and sidebar toggle."""
+ esc_key = _escape(field_key)
+ return f'''
+ (function() {{
+ var container = document.querySelector('.coding-trace-display[data-field-key="{esc_key}"]');
+ if (!container) return;
+
+ // Collapse/expand tool outputs
+ var toggleBtn = container.querySelector('[data-action="toggle-tools"]');
+ if (toggleBtn) {{
+ var collapsed = false;
+ toggleBtn.addEventListener('click', function() {{
+ collapsed = !collapsed;
+ container.querySelectorAll('.ct-tool-call').forEach(function(tc) {{
+ tc.classList.toggle('ct-tool-collapsed', collapsed);
+ }});
+ toggleBtn.textContent = collapsed ? 'Expand outputs' : 'Collapse outputs';
+ }});
+ }}
+
+ // Toggle sidebar
+ var sidebarBtn = container.querySelector('[data-action="toggle-sidebar"]');
+ var sidebar = document.getElementById('ct-sidebar-{esc_key}');
+ if (sidebarBtn && sidebar) {{
+ sidebarBtn.addEventListener('click', function() {{
+ sidebar.classList.toggle('ct-sidebar-collapsed');
+ }});
+ }}
+ }})();
+ '''
+
+ def _extract_plain_text(self, turns: List[Dict[str, Any]]) -> str:
+ """Extract plain text from reasoning for span annotation."""
+ parts = []
+ for turn in turns:
+ content = turn.get("content", "")
+ if content and turn.get("role") != "user":
+ parts.append(content)
+ return "\n".join(parts)
+
+ def _extract_reasoning_html(self, turns: List[Dict[str, Any]]) -> str:
+ """Extract reasoning HTML for span target wrapper."""
+ parts = []
+ for turn in turns:
+ content = turn.get("content", "")
+ if content and turn.get("role") != "user":
+ parts.append(_escape(content))
+ return " ".join(parts)
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ classes = super().get_css_classes(field_config)
+ if field_config.get("span_target"):
+ classes.append("span-target-field")
+ return classes
+
+ def get_data_attributes(self, field_config: Dict[str, Any], data: Any) -> Dict[str, str]:
+ attrs = super().get_data_attributes(field_config, data)
+ if field_config.get("span_target"):
+ attrs["span-target"] = "true"
+ return attrs
diff --git a/potato/server_utils/displays/conversation_tree_display.py b/potato/server_utils/displays/conversation_tree_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..cecbcf1fcc2cb8bda77a62ab5a557780b28a8304
--- /dev/null
+++ b/potato/server_utils/displays/conversation_tree_display.py
@@ -0,0 +1,133 @@
+"""
+Conversation Tree Display
+
+Renders branching conversation trees as nested collapsible nodes.
+Each node represents a message/turn in a dialogue with possible
+multiple branches (e.g., different model responses).
+
+Input data format:
+ {"id": "root", "speaker": "User", "text": "Question?",
+ "children": [
+ {"id": "r1", "speaker": "Bot A", "text": "Answer 1", "children": []},
+ {"id": "r2", "speaker": "Bot B", "text": "Answer 2", "children": []}
+ ]}
+"""
+
+import json
+from html import escape
+from typing import Dict, Any, List
+
+from .base import BaseDisplay
+
+
+class ConversationTreeDisplay(BaseDisplay):
+ name = "conversation_tree"
+ required_fields = ["key"]
+ optional_fields = {
+ "collapsed_depth": 2,
+ "node_style": "card",
+ "show_node_ids": False,
+ "max_depth": None,
+ }
+ description = "Conversation tree display with collapsible branching nodes"
+ supports_span_target = False
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ if not data:
+ return 'No conversation tree data
'
+
+ collapsed_depth = field_config.get("collapsed_depth", 2)
+ node_style = field_config.get("node_style", "card")
+ show_ids = field_config.get("show_node_ids", False)
+ max_depth = field_config.get("max_depth")
+
+ if isinstance(data, str):
+ try:
+ data = json.loads(data)
+ except (json.JSONDecodeError, TypeError):
+ return f'Invalid tree data
'
+
+ config_json = escape(json.dumps({
+ "collapsedDepth": collapsed_depth,
+ "nodeStyle": node_style,
+ "showIds": show_ids,
+ "maxDepth": max_depth,
+ }))
+
+ tree_html = self._render_node(data, 0, collapsed_depth, node_style, show_ids, max_depth)
+
+ return (
+ f''
+ f'
'
+ f' Expand All '
+ f' Collapse All '
+ f'
'
+ f'
{tree_html}
'
+ f'
'
+ )
+
+ def _render_node(self, node: dict, depth: int, collapsed_depth: int,
+ node_style: str, show_ids: bool, max_depth) -> str:
+ if not node or not isinstance(node, dict):
+ return ""
+
+ if max_depth is not None and depth > max_depth:
+ return '[depth limit reached]
'
+
+ node_id = escape(str(node.get("id", f"node_{depth}")))
+ speaker = escape(str(node.get("speaker", "")))
+ text = escape(str(node.get("text", "")))
+ children = node.get("children", [])
+ is_collapsed = depth >= collapsed_depth and len(children) > 0
+
+ # Speaker color class based on name hash
+ speaker_class = f"conv-tree-speaker-{abs(hash(speaker)) % 6}"
+
+ parts = []
+ parts.append(
+ f''
+ )
+
+ # Node header
+ parts.append(f'')
+
+ # Node text
+ parts.append(f'
{text}
')
+
+ # Children
+ if children:
+ display = "none" if is_collapsed else "block"
+ parts.append(f'
')
+ for child in children:
+ parts.append(self._render_node(
+ child, depth + 1, collapsed_depth, node_style, show_ids, max_depth
+ ))
+ parts.append('
')
+
+ parts.append('
')
+ return "\n".join(parts)
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ classes = super().get_css_classes(field_config)
+ classes.append("conv-tree-container")
+ return classes
+
+ def get_data_attributes(self, field_config: Dict[str, Any], data: Any) -> Dict[str, str]:
+ attrs = super().get_data_attributes(field_config, data)
+ attrs["display-type"] = "conversation_tree"
+ return attrs
diff --git a/potato/server_utils/displays/dialogue_display.py b/potato/server_utils/displays/dialogue_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..f0781b91a51072605c5b96866e6e5fa5e5a3382c
--- /dev/null
+++ b/potato/server_utils/displays/dialogue_display.py
@@ -0,0 +1,359 @@
+"""
+Dialogue Display Type
+
+Renders conversation/dialogue content for display in the annotation interface.
+Supports multiple conversation turns with speaker identification and styling.
+"""
+
+import html
+from typing import Dict, Any, List, Union
+
+from .base import BaseDisplay
+
+
+class DialogueDisplay(BaseDisplay):
+ """
+ Display type for dialogue/conversation content.
+
+ Displays conversations with alternating speaker turns and
+ visual styling to distinguish between speakers.
+ Can be used as a target for span annotations.
+ """
+
+ name = "dialogue"
+ required_fields = ["key"]
+ optional_fields = {
+ "alternating_shading": True,
+ "speaker_extraction": True,
+ "speaker_key": "speaker",
+ "text_key": "text",
+ "show_turn_numbers": False,
+ "per_turn_ratings": None,
+ }
+ description = "Dialogue/conversation turns display"
+ supports_span_target = True
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render dialogue content as HTML.
+
+ Args:
+ field_config: The field configuration
+ data: The dialogue data - can be:
+ - List of strings (each string is a turn)
+ - List of dicts with speaker/text keys
+ - String with turns separated by newlines
+
+ Returns:
+ HTML string for the dialogue display
+ """
+ if not data:
+ return 'No dialogue provided
'
+
+ # Get display options
+ options = self.get_display_options(field_config)
+ alternating_shading = options.get("alternating_shading", True)
+ speaker_extraction = options.get("speaker_extraction", True)
+ speaker_key = options.get("speaker_key", "speaker")
+ text_key = options.get("text_key", "text")
+ show_turn_numbers = options.get("show_turn_numbers", False)
+ per_turn_ratings = options.get("per_turn_ratings")
+
+ # Normalize the dialogue data to a list of turns
+ turns = self._normalize_dialogue(data, speaker_key, text_key, speaker_extraction)
+
+ if not turns:
+ return 'No dialogue turns found
'
+
+ field_key = html.escape(field_config.get("key", ""), quote=True)
+ is_span_target = field_config.get("span_target", False)
+
+ # Determine which speakers get per-turn ratings
+ rated_speakers = set()
+ rating_schemes = []
+ if per_turn_ratings:
+ rated_speakers = set(per_turn_ratings.get("speakers", []))
+ # Support both single-scheme and multi-scheme formats
+ if "schemes" in per_turn_ratings:
+ # New multi-dimension format
+ rating_schemes = per_turn_ratings["schemes"]
+ elif "scheme" in per_turn_ratings:
+ # Legacy single-scheme format: wrap in list
+ rating_schemes = [{
+ "schema_name": per_turn_ratings.get("schema_name", "per_turn_ratings"),
+ "scheme": per_turn_ratings["scheme"],
+ }]
+
+ # Build HTML for each turn
+ turn_html_list = []
+ for i, turn in enumerate(turns):
+ speaker = turn.get("speaker", "")
+ text = turn.get("text", "")
+
+ # Determine styling
+ turn_classes = ["dialogue-turn"]
+ if alternating_shading:
+ turn_classes.append(f"turn-{'even' if i % 2 == 0 else 'odd'}")
+
+ # Speaker-based styling
+ speaker_index = self._get_speaker_index(speaker, turns)
+ turn_classes.append(f"speaker-{speaker_index}")
+
+ # Build turn HTML
+ speaker_html = ""
+ if speaker:
+ escaped_speaker = html.escape(str(speaker))
+ speaker_html = f'{escaped_speaker}: '
+
+ turn_number_html = ""
+ if show_turn_numbers:
+ turn_number_html = f'[{i + 1}] '
+
+ escaped_text = html.escape(str(text))
+
+ # For span target, add data attributes
+ span_attrs = ""
+ text_id = ""
+ if is_span_target:
+ text_id = f'id="turn-text-{field_key}-{i}"'
+ span_attrs = f'data-original-text="{escaped_text}" data-turn-index="{i}"'
+
+ # Per-turn rating widgets (one or more per rated turn)
+ rating_html = ""
+ if per_turn_ratings and speaker in rated_speakers and rating_schemes:
+ if len(rating_schemes) == 1:
+ rating_html = self._render_turn_rating(
+ field_key, i, rating_schemes[0].get("scheme", {}),
+ rating_schemes[0].get("schema_name", "per_turn_ratings")
+ )
+ else:
+ # Multi-dimension: wrap multiple ratings in a group
+ parts = []
+ for scheme_entry in rating_schemes:
+ parts.append(self._render_turn_rating(
+ field_key, i, scheme_entry.get("scheme", {}),
+ scheme_entry.get("schema_name", "")
+ ))
+ rating_html = f'{"".join(parts)}
'
+
+ turn_html = f'''
+
+ {turn_number_html}
+ {speaker_html}
+ {escaped_text}
+ {rating_html}
+
+ '''
+ turn_html_list.append(turn_html)
+
+ # Combine all turns
+ all_turns_html = "\n".join(turn_html_list)
+
+ # For span annotation, wrap in .text-content WITHOUT data-original-text.
+ # Dialogue DOM textContent (with turn numbers, speaker prefixes, whitespace)
+ # differs from concatenate_dialogue_text() output. By omitting the attribute,
+ # getCanonicalText() falls back to container.textContent, so offsets from
+ # selection and from canonicalText always agree.
+ if is_span_target:
+ escaped_key = html.escape(field_key, quote=True)
+ all_turns_html = (
+ f''
+ f'{all_turns_html}'
+ f'
'
+ )
+
+ # Hidden inputs for storing per-turn rating data (one per scheme)
+ hidden_input_html = ""
+ if per_turn_ratings and rating_schemes:
+ hidden_parts = []
+ for scheme_entry in rating_schemes:
+ schema_name = html.escape(
+ scheme_entry.get("schema_name", "per_turn_ratings"), quote=True
+ )
+ hidden_parts.append(
+ f' '
+ )
+ hidden_input_html = "\n".join(hidden_parts)
+
+ # Wrap in container
+ container_classes = ["dialogue-display-content"]
+ if is_span_target:
+ container_classes.append("span-target-dialogue")
+ if per_turn_ratings:
+ container_classes.append("has-per-turn-ratings")
+
+ return f'''
+
+ {all_turns_html}
+ {hidden_input_html}
+
+ '''
+
+ def _render_turn_rating(self, field_key: str, turn_index: int,
+ rating_config: Dict[str, Any],
+ schema_name: str = "") -> str:
+ """
+ Render an inline rating widget for a dialogue turn.
+
+ Args:
+ field_key: The field key for the dialogue
+ turn_index: The index of the turn
+ rating_config: Configuration for the rating widget
+ schema_name: Schema name for multi-dimension support
+
+ Returns:
+ HTML string for the rating widget
+ """
+ size = rating_config.get("size", 5)
+ labels = rating_config.get("labels", [])
+ min_label = labels[0] if len(labels) > 0 else ""
+ max_label = labels[1] if len(labels) > 1 else ""
+
+ escaped_min = html.escape(str(min_label))
+ escaped_max = html.escape(str(max_label))
+ escaped_schema = html.escape(str(schema_name), quote=True)
+
+ # Build rating circles/stars
+ rating_items = []
+ for v in range(1, size + 1):
+ rating_items.append(
+ f'{v} '
+ )
+
+ items_html = "\n".join(rating_items)
+
+ min_html = f'{escaped_min} ' if min_label else ""
+ max_html = f'{escaped_max} ' if max_label else ""
+
+ # Schema label for multi-dimension mode
+ schema_label_html = ""
+ if schema_name:
+ readable_name = schema_name.replace("_", " ").title()
+ escaped_readable = html.escape(readable_name)
+ schema_label_html = f'{escaped_readable}: '
+
+ return f'''
+
+ {schema_label_html}
+ {min_html}
+
{items_html}
+ {max_html}
+
+ '''
+
+ def _normalize_dialogue(
+ self,
+ data: Any,
+ speaker_key: str,
+ text_key: str,
+ speaker_extraction: bool
+ ) -> List[Dict[str, str]]:
+ """
+ Normalize dialogue data to a list of {speaker, text} dicts.
+
+ Args:
+ data: Raw dialogue data
+ speaker_key: Key for speaker in dict format
+ text_key: Key for text in dict format
+ speaker_extraction: Whether to extract speaker from text
+
+ Returns:
+ List of turn dictionaries
+ """
+ turns = []
+
+ # Handle string input (newline-separated turns)
+ if isinstance(data, str):
+ lines = data.strip().split('\n')
+ for line in lines:
+ line = line.strip()
+ if not line:
+ continue
+ speaker, text = self._extract_speaker(line) if speaker_extraction else ("", line)
+ turns.append({"speaker": speaker, "text": text})
+
+ # Handle list input
+ elif isinstance(data, list):
+ for item in data:
+ if isinstance(item, str):
+ speaker, text = self._extract_speaker(item) if speaker_extraction else ("", item)
+ turns.append({"speaker": speaker, "text": text})
+ elif isinstance(item, dict):
+ speaker = item.get(speaker_key, "")
+ text = item.get(text_key, str(item))
+ turns.append({"speaker": speaker, "text": text})
+ else:
+ turns.append({"speaker": "", "text": str(item)})
+
+ # Handle single dict (unlikely but possible)
+ elif isinstance(data, dict):
+ speaker = data.get(speaker_key, "")
+ text = data.get(text_key, str(data))
+ turns.append({"speaker": speaker, "text": text})
+
+ return turns
+
+ def _extract_speaker(self, text: str) -> tuple:
+ """
+ Extract speaker from text if it starts with "Speaker:" pattern.
+
+ Args:
+ text: The text that may contain a speaker prefix
+
+ Returns:
+ Tuple of (speaker, remaining_text)
+ """
+ import re
+ # Match patterns like "Speaker:" or "Speaker 1:" or "User:" at the start
+ match = re.match(r'^([A-Za-z0-9_\s]+):\s*(.*)$', text)
+ if match:
+ return match.group(1).strip(), match.group(2).strip()
+ return "", text
+
+ def _get_speaker_index(self, speaker: str, turns: List[Dict[str, str]]) -> int:
+ """
+ Get a consistent index for a speaker for styling purposes.
+
+ Args:
+ speaker: The speaker name
+ turns: All turns in the dialogue
+
+ Returns:
+ Integer index for the speaker (0, 1, 2, ...)
+ """
+ if not speaker:
+ return 0
+
+ # Get unique speakers in order of first appearance
+ seen_speakers = []
+ for turn in turns:
+ s = turn.get("speaker", "")
+ if s and s not in seen_speakers:
+ seen_speakers.append(s)
+
+ try:
+ return seen_speakers.index(speaker)
+ except ValueError:
+ return 0
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ """Get CSS classes for the container."""
+ classes = super().get_css_classes(field_config)
+ if field_config.get("span_target"):
+ classes.append("span-target-field")
+ return classes
+
+ def get_data_attributes(self, field_config: Dict[str, Any], data: Any) -> Dict[str, str]:
+ """Get data attributes for the container."""
+ attrs = super().get_data_attributes(field_config, data)
+ if field_config.get("span_target"):
+ attrs["span-target"] = "true"
+ return attrs
diff --git a/potato/server_utils/displays/document_display.py b/potato/server_utils/displays/document_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..e06c29b3d50e5f56a19c7443e6165796e92ecf62
--- /dev/null
+++ b/potato/server_utils/displays/document_display.py
@@ -0,0 +1,379 @@
+"""
+Document Display Component
+
+Renders document content (DOCX, Markdown) with support for span annotation.
+Handles pre-extracted content from format handlers.
+
+Usage:
+ In instance_display config:
+ fields:
+ - key: document
+ type: document
+ display_options:
+ collapsible: false
+ max_height: 500
+"""
+
+from typing import Dict, Any, List, Optional
+import html
+import logging
+
+from .base import BaseDisplay
+
+logger = logging.getLogger(__name__)
+
+
+class DocumentDisplay(BaseDisplay):
+ """
+ Display type for rendered documents (DOCX, Markdown, HTML).
+
+ Displays HTML content extracted from documents with support for
+ span annotation on the text content, or bounding box annotation
+ for image-like region selection.
+ """
+
+ name = "document"
+ required_fields = ["key"]
+ optional_fields = {
+ "collapsible": False, # Allow collapsing sections
+ "max_height": None, # Max container height in pixels
+ "show_outline": False, # Show document outline/TOC
+ "preserve_structure": True, # Keep paragraph/heading structure
+ "style_theme": "default", # CSS theme: default, minimal, print
+ "annotation_mode": "span", # "span" or "bounding_box"
+ "bbox_min_size": 10, # Minimum bounding box size in pixels
+ "show_bbox_labels": True, # Show labels on bounding boxes
+ }
+ description = "Document display for DOCX, Markdown, and other formats"
+ supports_span_target = True
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render document content.
+
+ Args:
+ field_config: Display configuration
+ data: Either a dict with extracted content or raw HTML string
+
+ Returns:
+ HTML string for rendering
+ """
+ options = self.get_display_options(field_config)
+ annotation_mode = options.get("annotation_mode", "span")
+
+ # Use bounding box rendering if requested
+ if annotation_mode == "bounding_box":
+ return self._render_bbox_mode(field_config, data, options)
+
+ return self._render_span_mode(field_config, data, options)
+
+ def _render_span_mode(
+ self,
+ field_config: Dict[str, Any],
+ data: Any,
+ options: Dict[str, Any]
+ ) -> str:
+ """Render document in span annotation mode (default)."""
+ field_key = field_config.get("key", "document")
+
+ # Handle different data formats
+ if isinstance(data, dict):
+ # Pre-extracted FormatOutput data
+ rendered_html = data.get("rendered_html", "")
+ metadata = data.get("metadata", {})
+ raw_text = data.get("text", "")
+ elif isinstance(data, str):
+ # Raw HTML or text content
+ rendered_html = data
+ metadata = {}
+ raw_text = ""
+ else:
+ rendered_html = f"Unsupported content type: {type(data)}
"
+ metadata = {}
+ raw_text = ""
+
+ # Build container
+ parts = []
+
+ # Container styles
+ styles = []
+ max_height = options.get("max_height")
+ if max_height:
+ styles.append(f"max-height: {max_height}px")
+ styles.append("overflow-y: auto")
+ style_str = "; ".join(styles) if styles else ""
+
+ theme = options.get("style_theme", "default")
+
+ # Main container
+ parts.append(
+ f''
+ )
+
+ # Document outline/TOC if available and enabled
+ if options.get("show_outline") and metadata.get("headings"):
+ parts.append(self._render_outline(metadata["headings"]))
+
+ # Determine content classes - add text-content for span annotation compatibility
+ is_span_target = field_config.get("span_target", False)
+ content_classes = ["document-content"]
+ if is_span_target:
+ content_classes.append("text-content")
+ content_class_str = " ".join(content_classes)
+
+ # Extract plain text for span annotation (strip HTML tags)
+ import re
+ plain_text = re.sub(r'<[^>]+>', '', rendered_html)
+ plain_text = ' '.join(plain_text.split()) # Normalize whitespace
+ data_original_attr = f'data-original-text="{html.escape(plain_text)}"' if is_span_target else ""
+
+ # Collapsible wrapper if enabled
+ if options.get("collapsible"):
+ label = field_config.get("label", "Document")
+ parts.append(f'''
+
+ {html.escape(label)}
+
+ {rendered_html}
+
+
+ ''')
+ else:
+ parts.append(f'
{rendered_html}
')
+
+ # Hidden text container for span annotation (if different from rendered)
+ if raw_text and field_config.get("span_target"):
+ parts.append(
+ f'
'
+ f'{html.escape(raw_text)}'
+ f'
'
+ )
+
+ # Metadata footer
+ if metadata:
+ parts.append(self._render_metadata(metadata))
+
+ parts.append('
')
+
+ return "\n".join(parts)
+
+ def _render_bbox_mode(
+ self,
+ field_config: Dict[str, Any],
+ data: Any,
+ options: Dict[str, Any]
+ ) -> str:
+ """Render document in bounding box annotation mode."""
+ field_key = field_config.get("key", "document")
+ bbox_min_size = options.get("bbox_min_size", 10)
+ show_labels = options.get("show_bbox_labels", True)
+ max_height = options.get("max_height")
+ theme = options.get("style_theme", "default")
+
+ # Handle different data formats
+ if isinstance(data, dict):
+ rendered_html = data.get("rendered_html", "")
+ metadata = data.get("metadata", {})
+ elif isinstance(data, str):
+ rendered_html = data
+ metadata = {}
+ else:
+ rendered_html = f"Unsupported content type: {type(data)}
"
+ metadata = {}
+
+ # Build container
+ parts = []
+
+ # Container styles
+ styles = ["position: relative"]
+ if max_height:
+ styles.append(f"max-height: {max_height}px")
+ styles.append("overflow-y: auto")
+ style_str = "; ".join(styles)
+
+ # Main container with bbox mode
+ parts.append(
+ f''
+ )
+
+ # Bounding box toolbar
+ parts.append('''
+
+ ''')
+
+ # Content container with bbox canvas overlay
+ parts.append('
')
+
+ # The actual document content
+ parts.append(f'
{rendered_html}
')
+
+ # Canvas overlay for drawing bounding boxes
+ parts.append('
')
+
+ # Hidden input that carries the drawn boxes through the standard save
+ # pipeline. saveAnnotations() collects any input.annotation-data-input as
+ # "{name}:::_data", the server stores it, and render_page_with_annotations
+ # repopulates value + data-server-set on restore โ the same channel the
+ # image_annotation/video schemas use. Without this the boxes are never
+ # persisted (F-040). The name matches data-field-key so document-bbox.js
+ # reads/writes it.
+ parts.append(
+ f'
'
+ )
+
+ parts.append('
') # Close bbox-container
+
+ # Metadata footer
+ if metadata:
+ parts.append(self._render_metadata(metadata))
+
+ parts.append('
') # Close main container
+
+ return "\n".join(parts)
+
+ def _render_outline(self, headings: List[Dict[str, Any]]) -> str:
+ """
+ Render document outline/table of contents.
+ """
+ if not headings:
+ return ""
+
+ parts = ['']
+ parts.append('Contents
')
+ parts.append('')
+
+ for heading in headings:
+ level = heading.get("level", 1)
+ title = heading.get("title", "")
+ offset = heading.get("offset", 0)
+
+ indent_class = f"outline-level-{level}"
+ parts.append(
+ f''
+ f''
+ f'{html.escape(title)}'
+ f' '
+ )
+
+ parts.append(' ')
+ parts.append(' ')
+
+ return "\n".join(parts)
+
+ def _render_metadata(self, metadata: Dict[str, Any]) -> str:
+ """
+ Render metadata footer.
+ """
+ info_items = []
+
+ if "format" in metadata:
+ info_items.append(f"Format: {metadata['format'].upper()}")
+
+ if "paragraph_count" in metadata or "paragraphs" in metadata:
+ count = len(metadata.get("paragraphs", [])) or metadata.get("paragraph_count", 0)
+ if count:
+ info_items.append(f"Paragraphs: {count}")
+
+ if "line_count" in metadata:
+ info_items.append(f"Lines: {metadata['line_count']}")
+
+ if "char_count" in metadata:
+ info_items.append(f"Characters: {metadata['char_count']:,}")
+
+ if not info_items:
+ return ""
+
+ info_str = " | ".join(info_items)
+ return f'{info_str}
'
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ """Get CSS classes for the display container."""
+ classes = super().get_css_classes(field_config)
+ options = self.get_display_options(field_config)
+
+ if field_config.get("span_target"):
+ classes.append("span-target-document")
+
+ if options.get("collapsible"):
+ classes.append("document-collapsible-enabled")
+
+ theme = options.get("style_theme", "default")
+ classes.append(f"document-theme-{theme}")
+
+ annotation_mode = options.get("annotation_mode", "span")
+ if annotation_mode == "bounding_box":
+ classes.append("document-bbox-annotation")
+
+ return classes
+
+ def validate_config(self, field_config: Dict[str, Any]) -> List[str]:
+ """Validate the field configuration."""
+ errors = super().validate_config(field_config)
+ options = field_config.get("display_options", {})
+
+ # Validate annotation_mode
+ valid_modes = ["span", "bounding_box"]
+ annotation_mode = options.get("annotation_mode", "span")
+ if annotation_mode not in valid_modes:
+ errors.append(
+ f"Invalid annotation_mode '{annotation_mode}'. "
+ f"Must be one of: {', '.join(valid_modes)}"
+ )
+
+ # Validate style_theme
+ valid_themes = ["default", "minimal", "print"]
+ theme = options.get("style_theme", "default")
+ if theme not in valid_themes:
+ errors.append(
+ f"Invalid style_theme '{theme}'. "
+ f"Must be one of: {', '.join(valid_themes)}"
+ )
+
+ return errors
+
+ def get_data_attributes(
+ self,
+ field_config: Dict[str, Any],
+ data: Any
+ ) -> Dict[str, str]:
+ """Get data attributes for JavaScript initialization."""
+ attrs = super().get_data_attributes(field_config, data)
+
+ # Add format type if available
+ if isinstance(data, dict) and "format_name" in data:
+ attrs["format"] = data["format_name"]
+
+ return attrs
+
+ def has_inline_label(self, field_config: Dict[str, Any]) -> bool:
+ """
+ Check if the display handles its own label.
+
+ For collapsible documents, the label is shown in the summary element.
+ """
+ options = self.get_display_options(field_config)
+ return options.get("collapsible", False)
diff --git a/potato/server_utils/displays/eval_trace_display.py b/potato/server_utils/displays/eval_trace_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..c2a4ea501d148f382be02f9a530d06c18239d7d5
--- /dev/null
+++ b/potato/server_utils/displays/eval_trace_display.py
@@ -0,0 +1,485 @@
+"""
+Eval Trace Display Type
+
+Purpose-built rendering for *continuous agent evaluation*: it takes a single
+agent trace and splits it into three synchronized side-by-side panes โ
+
+ Reasoning | Function Calls | Final Answer
+
+so an evaluator can see, at a glance, what the agent thought, what it did, and
+what it ultimately produced. Clicking any card highlights the linked cards in
+the other panes (a "logical step" links a thought to the calls it triggered).
+
+Unlike ``agent_trace`` (which stacks an interleaved trace vertically in a
+single column), ``eval_trace`` decomposes one interleaved trace into its three
+semantic components. It consumes the same trace data formats as ``agent_trace``
+(see ``_trace_normalize.normalize_steps``), so existing data and trace
+converters work unchanged.
+
+Usage:
+ In instance_display config:
+ fields:
+ - key: trace
+ type: eval_trace
+ display_options:
+ pane_labels: ["Reasoning", "Function Calls", "Final Answer"]
+ show_step_numbers: true
+ collapse_long_outputs: true
+ max_output_lines: 20
+ link_steps: true
+ compact: false
+
+Data contract:
+ A single trace under the field key, in any format ``normalize_steps``
+ accepts. Step types map to panes as:
+ thought / system -> Reasoning
+ action -> Function Calls (with adjacent observation
+ rendered as a nested "โณ result")
+ observation -> nested under its preceding call
+ The Final Answer pane shows the trace's answer-like step (a step whose
+ speaker/tool matches "final answer", "send_message", "respond", etc.),
+ falling back to the last action. Make an explicit final answer by ending
+ the trace with a step whose speaker is e.g. "Agent (Final Answer)".
+"""
+
+import html
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+from .base import BaseDisplay
+from ._trace_normalize import normalize_steps
+
+
+# Speaker labels / tool names that mark a step as the final answer to the user.
+ANSWER_PATTERN = re.compile(
+ r"(final[\s_]*answer|send[\s_]*message|respond|response|finish|submit|conclusion)",
+ re.IGNORECASE,
+)
+
+DEFAULT_PANE_LABELS = ["Reasoning", "Function Calls", "Final Answer"]
+
+
+def _escape(text: Any) -> str:
+ """HTML-escape any value."""
+ return html.escape(str(text), quote=True)
+
+
+def _truncate(text: str, max_lines: int) -> Tuple[str, bool]:
+ """Truncate text to ``max_lines`` lines; return (text, was_truncated)."""
+ if not text or max_lines <= 0:
+ return text, False
+ lines = text.split("\n")
+ if len(lines) <= max_lines:
+ return text, False
+ return "\n".join(lines[:max_lines]), True
+
+
+def _tool_name(action_text: str) -> str:
+ """Extract the tool/function name from a ``tool(args)`` action string."""
+ if "(" in action_text:
+ return action_text.split("(", 1)[0].strip()
+ return action_text.strip()
+
+
+class EvalTraceDisplay(BaseDisplay):
+ """Three-pane (reasoning | function calls | final answer) trace display."""
+
+ name = "eval_trace"
+ required_fields = ["key"]
+ optional_fields = {
+ "pane_labels": DEFAULT_PANE_LABELS,
+ "show_step_numbers": True,
+ "collapse_long_outputs": True,
+ "max_output_lines": 20,
+ "link_steps": True,
+ "compact": False,
+ "speaker_key": "speaker",
+ "text_key": "text",
+ }
+ description = "Three-pane agent trace eval: reasoning, function calls, and final answer side-by-side"
+ # Per-pane card IDs do not follow the single .text-content wrapper contract
+ # required by SpanManager, so span annotation is not supported (yet).
+ supports_span_target = False
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ field_key = _escape(field_config.get("key", ""))
+
+ if not data:
+ return self._placeholder(field_key, "No trace data provided")
+
+ options = self.get_display_options(field_config)
+ pane_labels = self._resolve_pane_labels(options.get("pane_labels"))
+ speaker_key = options.get("speaker_key", "speaker")
+ text_key = options.get("text_key", "text")
+
+ steps = normalize_steps(data, speaker_key, text_key)
+ if not steps:
+ return self._placeholder(field_key, "No trace steps found")
+
+ answer_idx = self._find_answer_step(steps)
+ groups = self._build_groups(steps, exclude_idx=answer_idx)
+
+ show_numbers = options.get("show_step_numbers", True)
+ collapse = options.get("collapse_long_outputs", True)
+ max_lines = options.get("max_output_lines", 20)
+ link_steps = options.get("link_steps", True)
+ compact = options.get("compact", False)
+
+ reasoning_html = self._render_reasoning_pane(groups, show_numbers, link_steps)
+ calls_html = self._render_calls_pane(
+ groups, show_numbers, link_steps, collapse, max_lines
+ )
+ answer_html = self._render_answer_pane(
+ steps[answer_idx] if answer_idx is not None else None
+ )
+
+ css = self._build_css(compact)
+ link_attr = ' data-link-steps="true"' if link_steps else ""
+
+ panes = [
+ self._wrap_pane("reasoning", pane_labels[0], reasoning_html),
+ self._wrap_pane("calls", pane_labels[1], calls_html),
+ self._wrap_pane("answer", pane_labels[2], answer_html),
+ ]
+
+ js = self._build_js(field_key) if link_steps else ""
+
+ return f'''
+
+
+
+ '''
+
+ # ----- pane assembly --------------------------------------------------
+
+ def _wrap_pane(self, pane_id: str, label: str, body_html: str) -> str:
+ if not body_html.strip():
+ body_html = 'โ
'
+ return (
+ f''
+ f''
+ f'{body_html}
'
+ f' '
+ )
+
+ def _render_reasoning_pane(
+ self, groups: List[Dict[str, Any]], show_numbers: bool, link: bool
+ ) -> str:
+ cards = []
+ for g in groups:
+ if not g["thoughts"]:
+ continue
+ idx = g["index"]
+ num_html = (
+ f'#{idx + 1} ' if show_numbers else ""
+ )
+ for step in g["thoughts"]:
+ step_type = _escape(step.get("type", "thought"))
+ text = _escape(step.get("text", ""))
+ cards.append(
+ f''
+ f'
{num_html}'
+ f''
+ f'{_escape(step.get("speaker") or step_type.capitalize())}
'
+ f'
{text}
'
+ f'
'
+ )
+ return "\n".join(cards)
+
+ def _render_calls_pane(
+ self,
+ groups: List[Dict[str, Any]],
+ show_numbers: bool,
+ link: bool,
+ collapse: bool,
+ max_lines: int,
+ ) -> str:
+ cards = []
+ for g in groups:
+ if not g["calls"]:
+ continue
+ idx = g["index"]
+ num_html = (
+ f'#{idx + 1} ' if show_numbers else ""
+ )
+ for call in g["calls"]:
+ call_step = call["call"]
+ results = call["results"]
+
+ if call_step is not None:
+ action_text = str(call_step.get("text", ""))
+ tool = _escape(_tool_name(action_text))
+ call_line = (
+ f''
+ f'{tool} '
+ f'{_escape(action_text)}'
+ f'
'
+ )
+ else:
+ call_line = ""
+
+ results_html = "".join(
+ self._render_result(r, collapse, max_lines) for r in results
+ )
+
+ cards.append(
+ f''
+ f'
{num_html}
'
+ f'{call_line}{results_html}'
+ f'
'
+ )
+ return "\n".join(cards)
+
+ def _render_result(self, step: Dict[str, Any], collapse: bool, max_lines: int) -> str:
+ text = str(step.get("text", ""))
+ if not text:
+ return ""
+ truncated, was_truncated = _truncate(text, max_lines)
+ if was_truncated and collapse:
+ n = len(text.splitlines())
+ return (
+ f''
+ f'โณ result ({n} lines โ expand) '
+ f'{_escape(text)} '
+ f' '
+ )
+ return (
+ f''
+ f'
โณ '
+ f'
{_escape(truncated)} '
+ f'
'
+ )
+
+ def _render_answer_pane(self, answer_step: Optional[Dict[str, Any]]) -> str:
+ if not answer_step:
+ return 'No final answer in trace
'
+ text = str(answer_step.get("text", ""))
+ return f''
+
+ # ----- grouping / answer detection -----------------------------------
+
+ def _find_answer_step(self, steps: List[Dict[str, Any]]) -> Optional[int]:
+ """Return the index of the step that is the final answer, or None.
+
+ Preference order: the last step whose speaker or tool name matches an
+ answer pattern; otherwise the last ``action`` step; otherwise None.
+ """
+ answer_idx = None
+ last_action_idx = None
+ for i, step in enumerate(steps):
+ stype = step.get("type", "")
+ speaker = str(step.get("speaker", ""))
+ text = str(step.get("text", ""))
+ if stype == "action":
+ last_action_idx = i
+ if ANSWER_PATTERN.search(speaker) or ANSWER_PATTERN.search(_tool_name(text)):
+ answer_idx = i
+ elif ANSWER_PATTERN.search(speaker):
+ # An explicit "Final Answer" turn that isn't typed as an action.
+ answer_idx = i
+ if answer_idx is not None:
+ return answer_idx
+ return last_action_idx
+
+ def _build_groups(
+ self, steps: List[Dict[str, Any]], exclude_idx: Optional[int]
+ ) -> List[Dict[str, Any]]:
+ """Group steps into logical cycles linking thoughts to their calls.
+
+ A new group starts on a ``thought`` that follows a completed cycle (one
+ that already has calls), so consecutive thoughts stay together and the
+ thought(s) preceding a call share that call's group index.
+ """
+ groups: List[Dict[str, Any]] = []
+ current: Optional[Dict[str, Any]] = None
+
+ def new_group() -> Dict[str, Any]:
+ g = {"index": len(groups), "thoughts": [], "calls": []}
+ groups.append(g)
+ return g
+
+ for i, step in enumerate(steps):
+ if exclude_idx is not None and i == exclude_idx:
+ continue
+ stype = step.get("type", "observation")
+
+ if stype == "thought":
+ if current is None or current["calls"]:
+ current = new_group()
+ current["thoughts"].append(step)
+ elif stype == "action":
+ if current is None:
+ current = new_group()
+ current["calls"].append({"call": step, "results": []})
+ elif stype == "observation":
+ if current is None:
+ current = new_group()
+ if current["calls"]:
+ current["calls"][-1]["results"].append(step)
+ else:
+ current["calls"].append({"call": None, "results": [step]})
+ else: # system / error โ treat as a reasoning-side note
+ if current is None:
+ current = new_group()
+ current["thoughts"].append(step)
+
+ return groups
+
+ # ----- helpers --------------------------------------------------------
+
+ def _resolve_pane_labels(self, labels: Any) -> List[str]:
+ """Coerce ``pane_labels`` to exactly three strings, padding defaults."""
+ if not isinstance(labels, (list, tuple)):
+ return list(DEFAULT_PANE_LABELS)
+ result = [str(l) for l in labels[:3]]
+ while len(result) < 3:
+ result.append(DEFAULT_PANE_LABELS[len(result)])
+ return result
+
+ def _link_attr(self, idx: int, link: bool) -> str:
+ """Attributes that make a card a linkable, accessible button.
+
+ When linking is on, the card is an ARIA button that highlights the
+ steps sharing its index across panes. When off, the card carries no
+ index and is not focusable (it has no behavior to expose).
+ """
+ if not link:
+ return ""
+ return (
+ f' data-step-index="{idx}" role="button" tabindex="0"'
+ f' aria-pressed="false"'
+ f' aria-label="Highlight step {idx + 1} across panes"'
+ )
+
+ def _placeholder(self, field_key: str, message: str) -> str:
+ return (
+ f'{_escape(message)}
'
+ )
+
+ def validate_config(self, field_config: Dict[str, Any]) -> List[str]:
+ errors = super().validate_config(field_config)
+ opts = field_config.get("display_options", {}) or {}
+ labels = opts.get("pane_labels")
+ if labels is not None and not isinstance(labels, (list, tuple)):
+ errors.append(
+ f"Display type '{self.name}': 'pane_labels' must be a list of "
+ f"strings (got {type(labels).__name__})."
+ )
+ return errors
+
+ def _build_js(self, field_key: str) -> str:
+ """Cross-pane highlight: clicking/focusing a card with data-step-index
+ toggles the .eval-linked class on all cards sharing that index."""
+ return f'''
+ (function() {{
+ var root = document.querySelector('.eval-trace-display[data-field-key="{field_key}"]');
+ if (!root || root.dataset.evalBound) return;
+ root.dataset.evalBound = "1";
+ function clear() {{
+ root.querySelectorAll('.eval-card.eval-linked').forEach(function(c) {{
+ c.classList.remove('eval-linked');
+ if (c.hasAttribute('aria-pressed')) c.setAttribute('aria-pressed', 'false');
+ }});
+ }}
+ function linkTo(idx) {{
+ clear();
+ if (idx === null || idx === undefined) return;
+ root.querySelectorAll('.eval-card[data-step-index="' + idx + '"]').forEach(function(c) {{
+ c.classList.add('eval-linked');
+ if (c.hasAttribute('aria-pressed')) c.setAttribute('aria-pressed', 'true');
+ }});
+ }}
+ root.addEventListener('click', function(e) {{
+ var card = e.target.closest('.eval-card[data-step-index]');
+ if (!card) {{ clear(); return; }}
+ linkTo(card.getAttribute('data-step-index'));
+ }});
+ root.addEventListener('keydown', function(e) {{
+ if (e.key !== 'Enter' && e.key !== ' ') return;
+ var card = e.target.closest('.eval-card[data-step-index]');
+ if (card) {{ e.preventDefault(); linkTo(card.getAttribute('data-step-index')); }}
+ }});
+ }})();
+ '''
+
+ def _build_css(self, compact: bool) -> str:
+ pad = "6px 8px" if compact else "10px 12px"
+ gap = "8px" if compact else "12px"
+ return f'''
+ .eval-trace-display {{ font-family: inherit; width: 100%; }}
+ .eval-trace-empty {{ padding: 16px; color: #777; font-style: italic; }}
+ .eval-trace-panes {{
+ display: flex; gap: {gap}; align-items: stretch; width: 100%;
+ }}
+ .eval-pane {{
+ flex: 1 1 0; min-width: 0; display: flex; flex-direction: column;
+ border: 1px solid #e3e6ea; border-radius: 8px; overflow: hidden;
+ background: #fff;
+ }}
+ .eval-pane-header {{
+ padding: 8px 12px; font-weight: 600; font-size: 0.85em;
+ letter-spacing: 0.02em; text-transform: uppercase; color: #4a5568;
+ background: #f7f8fa; border-bottom: 1px solid #e3e6ea;
+ }}
+ .eval-pane-reasoning .eval-pane-header {{ color: #1565c0; }}
+ .eval-pane-calls .eval-pane-header {{ color: #c2410c; }}
+ .eval-pane-answer .eval-pane-header {{ color: #2e7d32; }}
+ .eval-pane-body {{ padding: {gap}; display: flex; flex-direction: column; gap: {gap}; }}
+ .eval-empty {{ color: #aaa; font-size: 0.9em; padding: 4px; }}
+ .eval-card {{
+ border-radius: 6px; padding: {pad}; border-left: 3px solid #cbd5e0;
+ background: #f8fafc; transition: box-shadow .12s, outline .12s;
+ outline: 2px solid transparent;
+ }}
+ /* Only linkable cards are interactive (they carry role=button). */
+ .eval-card[role="button"] {{ cursor: pointer; }}
+ .eval-card[role="button"]:hover {{ box-shadow: 0 1px 6px rgba(15,23,42,0.12); }}
+ .eval-card:focus-visible {{ outline: 2px solid #90cdf4; }}
+ .eval-card-thought, .eval-card-system {{ background: #e8f4fd; border-left-color: #2196F3; }}
+ .eval-card-action {{ background: #fff3e0; border-left-color: #FF9800; }}
+ .eval-card-answer {{ background: #e8f5e9; border-left-color: #4CAF50; }}
+ .eval-card-error {{ background: #ffebee; border-left-color: #f44336; }}
+ /* Linked-step highlight: an indigo ring distinct from the orange
+ action accent, plus a soft lift, so "linked" never reads as "action". */
+ .eval-card.eval-linked {{ box-shadow: 0 0 0 2px #6366f1, 0 2px 8px rgba(99,102,241,0.18); }}
+ .eval-card-head {{ display: flex; align-items: center; gap: 6px; margin-bottom: 4px; }}
+ .eval-step-num {{ color: #718096; font-size: 0.78em; font-weight: 600; }}
+ .eval-badge {{
+ padding: 1px 7px; border-radius: 10px; font-size: 0.75em; font-weight: 600; color: #2d3748;
+ }}
+ .badge-thought, .badge-system {{ background: rgba(33,150,243,0.18); }}
+ .eval-card-text {{ white-space: pre-wrap; word-break: break-word; line-height: 1.45; font-size: 0.92em; }}
+ .eval-call-line {{ display: flex; align-items: baseline; gap: 6px; flex-wrap: wrap; }}
+ .eval-tool-badge {{
+ background: rgba(255,152,0,0.22); color: #b45309; padding: 1px 7px;
+ border-radius: 4px; font-size: 0.78em; font-weight: 700;
+ font-family: ui-monospace, SFMono-Regular, Menlo, monospace;
+ }}
+ .eval-call-code {{
+ font-family: ui-monospace, SFMono-Regular, Menlo, monospace;
+ font-size: 0.85em; word-break: break-word; color: #44403c;
+ }}
+ .eval-result {{ margin-top: 6px; display: flex; gap: 4px; align-items: flex-start; }}
+ .eval-result-arrow {{ color: #16a34a; font-weight: 700; flex: 0 0 auto; }}
+ .eval-result-pre {{
+ margin: 0; white-space: pre-wrap; word-break: break-word;
+ font-family: ui-monospace, SFMono-Regular, Menlo, monospace;
+ font-size: 0.82em; color: #57534e; background: rgba(0,0,0,0.03);
+ padding: 4px 6px; border-radius: 4px; flex: 1 1 auto; min-width: 0;
+ }}
+ details.eval-result {{ display: block; }}
+ details.eval-result summary {{ cursor: pointer; color: #16a34a; font-size: 0.82em; }}
+ @media (max-width: 720px) {{
+ .eval-trace-panes {{ flex-direction: column; }}
+ }}
+ @media (prefers-reduced-motion: reduce) {{
+ .eval-card {{ transition: none; }}
+ }}
+ '''
diff --git a/potato/server_utils/displays/gallery_display.py b/potato/server_utils/displays/gallery_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..caaef5ccca3ad76f027dfb0fa52cca057c2e6c88
--- /dev/null
+++ b/potato/server_utils/displays/gallery_display.py
@@ -0,0 +1,237 @@
+"""
+Gallery Display Type
+
+Scrollable image gallery for displaying screenshot sequences, step-by-step
+visual traces, or any ordered collection of images with optional captions.
+"""
+
+import html
+from typing import Dict, Any, List
+
+from .base import BaseDisplay
+
+
+class GalleryDisplay(BaseDisplay):
+ """
+ Display type for ordered image galleries with captions.
+
+ Supports data as:
+ - List of image URLs/paths (strings)
+ - List of dicts with url/caption keys
+ """
+
+ name = "gallery"
+ required_fields = ["key"]
+ optional_fields = {
+ "layout": "horizontal", # horizontal, vertical, grid
+ "thumbnail_size": 300,
+ "show_captions": True,
+ "caption_key": "caption",
+ "url_key": "url",
+ "zoomable": True,
+ "max_height": 400,
+ "columns": 3, # for grid layout
+ }
+ description = "Scrollable image gallery with captions"
+ supports_span_target = False
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ if not data:
+ return 'No images provided
'
+
+ options = self.get_display_options(field_config)
+ layout = options.get("layout", "horizontal")
+ thumbnail_size = options.get("thumbnail_size", 300)
+ show_captions = options.get("show_captions", True)
+ caption_key = options.get("caption_key", "caption")
+ url_key = options.get("url_key", "url")
+ zoomable = options.get("zoomable", True)
+ max_height = options.get("max_height", 400)
+ columns = options.get("columns", 3)
+
+ field_key = html.escape(field_config.get("key", ""), quote=True)
+
+ # Normalize data to list of {url, caption} dicts
+ items = self._normalize_items(data, url_key, caption_key)
+ if not items:
+ return 'No valid images found
'
+
+ # Build CSS
+ css = self._build_css(layout, thumbnail_size, max_height, columns, zoomable)
+
+ # Build gallery items
+ item_html_list = []
+ for i, item in enumerate(items):
+ url = html.escape(str(item["url"]), quote=True)
+ caption = item.get("caption", "")
+
+ caption_html = ""
+ if show_captions and caption:
+ escaped_caption = html.escape(str(caption))
+ caption_html = f'{escaped_caption}
'
+
+ zoom_attr = 'data-zoomable="true"' if zoomable else ''
+
+ item_html = f'''
+
+
+
+
+ {caption_html}
+
+ '''
+ item_html_list.append(item_html)
+
+ all_items = "\n".join(item_html_list)
+
+ # Navigation counter for horizontal layout
+ nav_html = ""
+ counter_js = ""
+ if layout == "horizontal" and len(items) > 1:
+ nav_html = '''
+
+ 1 / {total}
+
+ '''.format(total=len(items))
+ counter_js = '''
+
+ '''.format(field_key=field_key)
+
+ return f'''
+
+
+ {{nav_html}}
+
+ {{all_items}}
+
+
+ {{counter_js}}
+ '''.format(css=css, layout=layout, field_key=field_key,
+ nav_html=nav_html, all_items=all_items, counter_js=counter_js)
+
+ def _normalize_items(self, data: Any, url_key: str, caption_key: str) -> List[Dict[str, str]]:
+ """Normalize gallery data to list of {url, caption} dicts."""
+ items = []
+
+ if isinstance(data, str):
+ items.append({"url": data, "caption": ""})
+ elif isinstance(data, list):
+ for i, item in enumerate(data):
+ if isinstance(item, str):
+ items.append({"url": item, "caption": f"Step {i + 1}"})
+ elif isinstance(item, dict):
+ url = item.get(url_key, item.get("src", item.get("path", "")))
+ caption = item.get(caption_key, item.get("label", item.get("description", "")))
+ if url:
+ items.append({"url": url, "caption": caption})
+
+ return items
+
+ def _build_css(self, layout: str, thumbnail_size: int, max_height: int,
+ columns: int, zoomable: bool) -> str:
+ """Build CSS for the gallery layout."""
+ # Ensure numeric types for CSS arithmetic
+ try:
+ thumbnail_size = int(thumbnail_size)
+ except (ValueError, TypeError):
+ thumbnail_size = 300
+ try:
+ max_height = int(max_height)
+ except (ValueError, TypeError):
+ max_height = 400
+ try:
+ columns = int(columns)
+ except (ValueError, TypeError):
+ columns = 3
+
+ layout_css = ""
+
+ if layout == "horizontal":
+ layout_css = f'''
+ .gallery-horizontal .gallery-container {{
+ display: flex; overflow-x: auto; gap: 12px;
+ padding: 8px 0; max-height: {max_height}px;
+ scroll-snap-type: x mandatory;
+ }}
+ .gallery-horizontal .gallery-item {{
+ flex: 0 0 auto; scroll-snap-align: start;
+ }}
+ .gallery-horizontal .gallery-img {{
+ max-height: {max_height - 40}px; width: auto;
+ max-width: {thumbnail_size}px;
+ }}
+ '''
+ elif layout == "vertical":
+ layout_css = f'''
+ .gallery-vertical .gallery-container {{
+ display: flex; flex-direction: column; gap: 12px;
+ max-height: {max_height}px; overflow-y: auto;
+ }}
+ .gallery-vertical .gallery-img {{
+ max-width: 100%; height: auto;
+ max-height: {thumbnail_size}px;
+ }}
+ '''
+ elif layout == "grid":
+ layout_css = f'''
+ .gallery-grid .gallery-container {{
+ display: grid; grid-template-columns: repeat({columns}, 1fr);
+ gap: 12px; max-height: {max_height}px; overflow-y: auto;
+ }}
+ .gallery-grid .gallery-img {{
+ width: 100%; height: auto; max-height: {thumbnail_size}px;
+ object-fit: cover;
+ }}
+ '''
+
+ zoom_css = ""
+ if zoomable:
+ zoom_css = '''
+ .gallery-img[data-zoomable]:hover {
+ cursor: zoom-in; opacity: 0.9;
+ }
+ '''
+
+ return f'''
+ .gallery-display {{ font-family: inherit; }}
+ .gallery-item {{
+ border: 1px solid #e0e0e0; border-radius: 6px;
+ overflow: hidden; background: #fafafa;
+ }}
+ .gallery-img-wrapper {{ display: flex; justify-content: center; padding: 4px; }}
+ .gallery-img {{ border-radius: 4px; }}
+ .gallery-caption {{
+ padding: 6px 10px; font-size: 0.85em; color: #555;
+ background: #f5f5f5; border-top: 1px solid #e0e0e0;
+ text-align: center;
+ }}
+ .gallery-nav {{
+ display: flex; justify-content: center; padding: 4px 0;
+ margin-bottom: 8px;
+ }}
+ .gallery-counter {{ font-size: 0.85em; color: #666; }}
+ {layout_css}
+ {zoom_css}
+ '''
diff --git a/potato/server_utils/displays/image_display.py b/potato/server_utils/displays/image_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..f684802f0ca8337df4d9ecfd42af95d0c66e7d86
--- /dev/null
+++ b/potato/server_utils/displays/image_display.py
@@ -0,0 +1,181 @@
+"""
+Image Display Type
+
+Renders images for display in the annotation interface.
+Supports zoom functionality and links to image annotation schemas.
+"""
+
+import html
+from typing import Dict, Any, List
+from urllib.parse import urlparse
+
+from .base import BaseDisplay
+
+
+class ImageDisplay(BaseDisplay):
+ """
+ Display type for image content.
+
+ Displays images with optional zoom functionality.
+ Can be linked to image_annotation schemas via source_field.
+ """
+
+ name = "image"
+ required_fields = ["key"]
+ optional_fields = {
+ "max_width": None,
+ "max_height": None,
+ "zoomable": True,
+ "alt_text": "",
+ "object_fit": "contain",
+ }
+ description = "Image display with optional zoom"
+ supports_span_target = False
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render an image as HTML.
+
+ Args:
+ field_config: The field configuration
+ data: The image URL or path
+
+ Returns:
+ HTML string for the image display
+ """
+ if not data:
+ return 'No image provided
'
+
+ # Get the image URL
+ image_url = str(data)
+
+ # Validate URL (basic check)
+ if not self._is_valid_url(image_url):
+ return f'Invalid image URL: {html.escape(image_url)}
'
+
+ # Get display options
+ options = self.get_display_options(field_config)
+ max_width = options.get("max_width")
+ max_height = options.get("max_height")
+ zoomable = options.get("zoomable", True)
+ alt_text = options.get("alt_text", "")
+ object_fit = options.get("object_fit", "contain")
+
+ # Build style
+ style_parts = []
+ if max_width:
+ style_parts.append(f"max-width: {max_width}px" if isinstance(max_width, int) else f"max-width: {max_width}")
+ if max_height:
+ style_parts.append(f"max-height: {max_height}px" if isinstance(max_height, int) else f"max-height: {max_height}")
+ style_parts.append(f"object-fit: {object_fit}")
+
+ style_attr = f' style="{"; ".join(style_parts)}"' if style_parts else ""
+
+ # Escape values for HTML attributes
+ escaped_url = html.escape(image_url, quote=True)
+ escaped_alt = html.escape(alt_text or "Image content", quote=True)
+ field_key = html.escape(field_config.get("key", ""), quote=True)
+
+ # Build the image HTML
+ img_classes = ["display-image"]
+ if zoomable:
+ img_classes.append("zoomable-image")
+
+ img_html = f''' '''
+
+ # Wrap in zoom container if zoomable
+ if zoomable:
+ return f'''
+
+ {img_html}
+
+
+ +
+
+
+ -
+
+
+ โฒ
+
+
+
+ '''
+
+ return f'{img_html}
'
+
+ def _is_valid_url(self, url: str) -> bool:
+ """
+ Check if a URL is valid and safe.
+
+ Args:
+ url: The URL to validate
+
+ Returns:
+ True if valid, False otherwise
+ """
+ if not url:
+ return False
+
+ # Allow relative paths
+ if url.startswith('/') or url.startswith('./') or url.startswith('../'):
+ return True
+
+ # Parse the URL
+ try:
+ parsed = urlparse(url)
+ # Must have a scheme (http/https) or be a relative path
+ if parsed.scheme:
+ return parsed.scheme.lower() in ('http', 'https', 'data')
+ # No scheme - treat as relative path
+ return True
+ except Exception:
+ return False
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ """Get CSS classes for the container."""
+ classes = super().get_css_classes(field_config)
+ options = self.get_display_options(field_config)
+ if options.get("zoomable", True):
+ classes.append("zoomable")
+ return classes
+
+ def get_data_attributes(self, field_config: Dict[str, Any], data: Any) -> Dict[str, str]:
+ """Get data attributes for the container."""
+ attrs = super().get_data_attributes(field_config, data)
+ if data:
+ attrs["source-url"] = str(data)
+ return attrs
+
+ def get_js_init(self) -> str:
+ """Get JavaScript initialization code for zoom functionality."""
+ return '''
+ // Initialize image zoom controls
+ document.querySelectorAll('.image-zoom-container').forEach(container => {
+ const img = container.querySelector('img');
+ let scale = 1;
+
+ container.querySelector('.zoom-in')?.addEventListener('click', () => {
+ scale = Math.min(scale * 1.25, 5);
+ img.style.transform = `scale(${scale})`;
+ });
+
+ container.querySelector('.zoom-out')?.addEventListener('click', () => {
+ scale = Math.max(scale / 1.25, 0.5);
+ img.style.transform = `scale(${scale})`;
+ });
+
+ container.querySelector('.zoom-reset')?.addEventListener('click', () => {
+ scale = 1;
+ img.style.transform = 'scale(1)';
+ });
+ });
+ '''
diff --git a/potato/server_utils/displays/interactive_chat_display.py b/potato/server_utils/displays/interactive_chat_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..8a3e891c633205742dc1442f9ccca9ec886037f1
--- /dev/null
+++ b/potato/server_utils/displays/interactive_chat_display.py
@@ -0,0 +1,94 @@
+"""
+Interactive Chat Display Type
+
+Renders either a live chat panel (when conversation data is null/empty)
+or a completed dialogue (when conversation data is populated).
+
+Before the annotator finishes chatting, this shows the chat UI.
+After clicking "Finish & Annotate", the route writes the conversation
+into the item data, and this display renders the completed conversation
+using DialogueDisplay (which supports per-turn ratings).
+"""
+
+import html
+from typing import Dict, Any, List
+
+from .base import BaseDisplay
+from .dialogue_display import DialogueDisplay
+
+# Reuse the dialogue display for rendering completed conversations
+_dialogue_display = DialogueDisplay()
+
+
+class InteractiveChatDisplay(BaseDisplay):
+ """
+ Display type for interactive agent chat sessions.
+
+ When data is null/empty: renders a chat panel placeholder
+ (the actual chat UI is handled by agent-chat.js).
+ When data is populated: delegates to DialogueDisplay for the conversation,
+ which supports per-turn ratings for individual turn annotation.
+ """
+
+ name = "interactive_chat"
+ required_fields = ["key"]
+ optional_fields = {
+ "placeholder_text": "Start chatting with the agent to begin the task.",
+ "per_turn_ratings": None,
+ "show_turn_numbers": True,
+ "alternating_shading": True,
+ }
+ description = "Interactive agent chat with post-interaction trace display"
+ supports_span_target = True
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ # If conversation data exists, render as dialogue with per-turn ratings
+ if data:
+ return _dialogue_display.render(field_config, data)
+
+ # Otherwise render the chat panel container
+ # The actual chat UI is injected by agent-chat.js
+ options = self.get_display_options(field_config)
+ placeholder = html.escape(options.get(
+ "placeholder_text",
+ "Start chatting with the agent to begin the task.",
+ ))
+ field_key = html.escape(field_config.get("key", ""), quote=True)
+
+ return f'''
+
+ '''
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ classes = super().get_css_classes(field_config)
+ # Include display-type-dialogue so dialogue CSS rules apply to
+ # the completed conversation rendered by DialogueDisplay
+ classes.append("display-type-dialogue")
+ if field_config.get("span_target"):
+ classes.append("span-target-field")
+ return classes
+
+ def get_data_attributes(self, field_config: Dict[str, Any], data: Any) -> Dict[str, str]:
+ attrs = super().get_data_attributes(field_config, data)
+ if field_config.get("span_target"):
+ attrs["span-target"] = "true"
+ # Signal to JS whether this is in chat mode or trace mode
+ attrs["chat-active"] = "true" if not data else "false"
+ return attrs
diff --git a/potato/server_utils/displays/live_agent_display.py b/potato/server_utils/displays/live_agent_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..31013bdbf3c6c74b50cb7c257307d24b78268186
--- /dev/null
+++ b/potato/server_utils/displays/live_agent_display.py
@@ -0,0 +1,568 @@
+"""
+Live Agent Display Type
+
+Dual-mode display:
+- When data is empty/None: renders the live agent UI (SSE-connected viewer,
+ controls, instruction input, thought panel, filmstrip)
+- When data contains a completed trace: delegates to WebAgentTraceDisplay
+ for post-hoc review
+
+Configuration example:
+ instance_display:
+ fields:
+ - key: agent_trace
+ type: live_agent
+ label: "Live Agent Session"
+ display_options:
+ show_overlays: true
+ show_filmstrip: true
+ show_thought: true
+ show_controls: true
+ allow_takeover: true
+ allow_instructions: true
+ screenshot_max_width: 900
+ screenshot_max_height: 650
+"""
+
+import html
+import json
+from typing import Any, Dict, List
+
+from .base import BaseDisplay
+
+
+class LiveAgentDisplay(BaseDisplay):
+ """
+ Display type for live AI agent interaction.
+
+ Renders a viewer that connects to the agent SSE stream, showing real-time
+ screenshots, overlay visualizations, and providing controls for
+ pause/resume/instruct/takeover.
+ """
+
+ name = "live_agent"
+ required_fields = ["key"]
+ optional_fields = {
+ "show_overlays": True,
+ "show_filmstrip": True,
+ "show_thought": True,
+ "show_controls": True,
+ "allow_takeover": True,
+ "allow_instructions": True,
+ "screenshot_max_width": 900,
+ "screenshot_max_height": 650,
+ "filmstrip_size": 80,
+ }
+ description = "Live AI agent viewer with real-time screenshots, controls, and interaction"
+ supports_span_target = False
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render live agent UI or delegate to trace viewer.
+
+ If data contains a completed trace (dict with 'steps'), delegates
+ to WebAgentTraceDisplay. Otherwise renders the live interaction UI.
+ """
+ # Check if data is a completed trace (has steps)
+ if data and isinstance(data, dict) and data.get("steps"):
+ return self._render_review_mode(field_config, data)
+
+ return self._render_live_mode(field_config, data)
+
+ def _render_review_mode(self, field_config: Dict[str, Any], data: dict) -> str:
+ """Delegate to WebAgentTraceDisplay for completed traces."""
+ from .web_agent_trace_display import WebAgentTraceDisplay
+
+ reviewer = WebAgentTraceDisplay()
+ return reviewer.render(field_config, data)
+
+ def _render_live_mode(self, field_config: Dict[str, Any], data: Any) -> str:
+ """Render the live agent interaction UI."""
+ options = self.get_display_options(field_config)
+ field_key = html.escape(field_config.get("key", ""), quote=True)
+
+ max_w = options.get("screenshot_max_width", 900)
+ max_h = options.get("screenshot_max_height", 650)
+ filmstrip_size = options.get("filmstrip_size", 80)
+ show_controls = options.get("show_controls", True)
+ show_thought = options.get("show_thought", True)
+ show_filmstrip = options.get("show_filmstrip", True)
+ show_overlays = options.get("show_overlays", True)
+ allow_takeover = options.get("allow_takeover", True)
+ allow_instructions = options.get("allow_instructions", True)
+
+ # Encode config as data attributes for JS
+ config_json = html.escape(json.dumps({
+ "show_overlays": show_overlays,
+ "show_filmstrip": show_filmstrip,
+ "show_thought": show_thought,
+ "show_controls": show_controls,
+ "allow_takeover": allow_takeover,
+ "allow_instructions": allow_instructions,
+ }), quote=True)
+
+ # Extract task info from instance data if available
+ task_desc = ""
+ start_url = ""
+ if isinstance(data, dict):
+ task_desc = data.get("task_description", "")
+ start_url = data.get("start_url", data.get("url", ""))
+
+ css = self._build_css(max_w, max_h, filmstrip_size)
+ html_content = self._build_html(
+ field_key, config_json, task_desc, start_url,
+ show_controls, show_thought, show_filmstrip,
+ show_overlays, allow_takeover, allow_instructions,
+ max_w, max_h, filmstrip_size,
+ )
+
+ return f"\n{html_content}"
+
+ def _build_css(self, max_w: int, max_h: int, filmstrip_size: int) -> str:
+ return f"""
+.live-agent-viewer {{
+ font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
+ border: 1px solid #ddd;
+ border-radius: 8px;
+ overflow: hidden;
+ background: #fafafa;
+}}
+
+/* Status bar */
+.live-agent-status {{
+ display: flex;
+ align-items: center;
+ justify-content: space-between;
+ padding: 8px 16px;
+ background: #f0f0f0;
+ border-bottom: 1px solid #ddd;
+ font-size: 13px;
+}}
+.live-agent-status-indicator {{
+ display: inline-block;
+ width: 10px;
+ height: 10px;
+ border-radius: 50%;
+ margin-right: 8px;
+}}
+.live-agent-status-indicator.idle {{ background: #9E9E9E; }}
+.live-agent-status-indicator.running {{ background: #4CAF50; animation: pulse-dot 1.5s infinite; }}
+.live-agent-status-indicator.paused {{ background: #FF9800; }}
+.live-agent-status-indicator.takeover {{ background: #2196F3; }}
+.live-agent-status-indicator.completed {{ background: #388E3C; }}
+.live-agent-status-indicator.error {{ background: #F44336; }}
+@keyframes pulse-dot {{
+ 0%, 100% {{ opacity: 1; }}
+ 50% {{ opacity: 0.5; }}
+}}
+
+/* Main layout */
+.live-agent-main {{
+ display: flex;
+ gap: 0;
+}}
+.live-agent-screenshot-panel {{
+ flex: 1;
+ position: relative;
+ background: #000;
+ min-height: 200px;
+ max-width: {max_w}px;
+}}
+.live-agent-screenshot {{
+ width: 100%;
+ max-height: {max_h}px;
+ object-fit: contain;
+ display: block;
+}}
+/* Collapsed state โ shrink screenshot for more annotation space */
+.live-agent-viewer.collapsed .live-agent-main {{
+ max-height: 220px;
+ overflow: hidden;
+}}
+.live-agent-viewer.collapsed .live-agent-screenshot {{
+ max-height: 200px;
+}}
+.live-agent-viewer.collapsed .live-agent-side-panel {{
+ display: none;
+}}
+/* Collapse toggle */
+.live-agent-collapse-toggle {{
+ cursor: pointer;
+ font-size: 12px;
+ color: #666;
+ padding: 2px 8px;
+ border: 1px solid #ccc;
+ border-radius: 4px;
+ background: #fff;
+}}
+.live-agent-collapse-toggle:hover {{ background: #f0f0f0; }}
+
+/* Takeover click feedback */
+.live-agent-click-marker {{
+ position: absolute;
+ width: 24px;
+ height: 24px;
+ border: 2px solid #F44336;
+ border-radius: 50%;
+ background: rgba(244, 67, 54, 0.2);
+ transform: translate(-50%, -50%);
+ pointer-events: none;
+ animation: click-pulse 0.6s ease-out forwards;
+ z-index: 100;
+}}
+@keyframes click-pulse {{
+ 0% {{ transform: translate(-50%, -50%) scale(0.5); opacity: 1; }}
+ 100% {{ transform: translate(-50%, -50%) scale(2); opacity: 0; }}
+}}
+
+/* Takeover action toast */
+.live-agent-action-toast {{
+ position: absolute;
+ bottom: 8px;
+ left: 50%;
+ transform: translateX(-50%);
+ background: rgba(0,0,0,0.8);
+ color: #fff;
+ padding: 4px 12px;
+ border-radius: 4px;
+ font-size: 12px;
+ pointer-events: none;
+ z-index: 100;
+ animation: toast-fade 1.5s ease-out forwards;
+}}
+@keyframes toast-fade {{
+ 0% {{ opacity: 1; }}
+ 70% {{ opacity: 1; }}
+ 100% {{ opacity: 0; }}
+}}
+.live-agent-screenshot-placeholder {{
+ display: flex;
+ align-items: center;
+ justify-content: center;
+ min-height: 400px;
+ color: #888;
+ font-size: 16px;
+}}
+.live-agent-overlay-layer {{
+ position: absolute;
+ top: 0;
+ left: 0;
+ width: 100%;
+ height: 100%;
+ pointer-events: none;
+}}
+
+/* Side panel */
+.live-agent-side-panel {{
+ width: 320px;
+ display: flex;
+ flex-direction: column;
+ border-left: 1px solid #ddd;
+ background: #fff;
+}}
+
+/* Thought panel */
+.live-agent-thought-panel {{
+ padding: 12px;
+ border-bottom: 1px solid #eee;
+ max-height: 200px;
+ overflow-y: auto;
+}}
+.live-agent-thought-panel h4 {{
+ margin: 0 0 8px 0;
+ font-size: 12px;
+ text-transform: uppercase;
+ color: #666;
+}}
+.live-agent-thought-text {{
+ font-size: 13px;
+ line-height: 1.5;
+ color: #333;
+ white-space: pre-wrap;
+}}
+
+/* Step details */
+.live-agent-step-details {{
+ padding: 12px;
+ flex: 1;
+ overflow-y: auto;
+ font-size: 13px;
+}}
+.live-agent-action-badge {{
+ display: inline-block;
+ padding: 2px 8px;
+ border-radius: 4px;
+ font-size: 12px;
+ font-weight: 600;
+ text-transform: uppercase;
+}}
+
+/* Controls */
+.live-agent-controls {{
+ padding: 12px;
+ border-top: 1px solid #ddd;
+ background: #f8f8f8;
+}}
+.live-agent-control-buttons {{
+ display: flex;
+ gap: 8px;
+ flex-wrap: wrap;
+ margin-bottom: 8px;
+}}
+.live-agent-btn {{
+ padding: 6px 14px;
+ border: 1px solid #ddd;
+ border-radius: 4px;
+ background: #fff;
+ cursor: pointer;
+ font-size: 13px;
+ transition: all 0.15s;
+}}
+.live-agent-btn:hover {{ background: #f0f0f0; }}
+.live-agent-btn.primary {{ background: #2196F3; color: #fff; border-color: #1976D2; }}
+.live-agent-btn.primary:hover {{ background: #1976D2; }}
+.live-agent-btn.danger {{ background: #F44336; color: #fff; border-color: #D32F2F; }}
+.live-agent-btn.danger:hover {{ background: #D32F2F; }}
+.live-agent-btn.warning {{ background: #FF9800; color: #fff; border-color: #F57C00; }}
+.live-agent-btn.warning:hover {{ background: #F57C00; }}
+.live-agent-btn.active {{ background: #4CAF50; color: #fff; border-color: #388E3C; }}
+.live-agent-btn:disabled {{ opacity: 0.5; cursor: not-allowed; }}
+
+/* Instruction input */
+.live-agent-instruction-input {{
+ display: flex;
+ gap: 8px;
+ margin-top: 8px;
+}}
+.live-agent-instruction-input input {{
+ flex: 1;
+ padding: 6px 10px;
+ border: 1px solid #ddd;
+ border-radius: 4px;
+ font-size: 13px;
+}}
+
+/* Start form */
+.live-agent-start-form {{
+ padding: 24px;
+ text-align: center;
+}}
+.live-agent-start-form input {{
+ display: block;
+ width: 100%;
+ max-width: 500px;
+ margin: 8px auto;
+ padding: 8px 12px;
+ border: 1px solid #ddd;
+ border-radius: 4px;
+ font-size: 14px;
+}}
+
+/* Filmstrip */
+.live-agent-filmstrip {{
+ display: flex;
+ gap: 4px;
+ padding: 8px;
+ overflow-x: auto;
+ border-top: 1px solid #ddd;
+ background: #f5f5f5;
+}}
+.live-agent-filmstrip-thumb {{
+ width: {filmstrip_size}px;
+ height: {int(filmstrip_size * 0.75)}px;
+ object-fit: cover;
+ border: 2px solid transparent;
+ border-radius: 4px;
+ cursor: pointer;
+ flex-shrink: 0;
+ opacity: 0.7;
+ transition: all 0.15s;
+}}
+.live-agent-filmstrip-thumb:hover {{ opacity: 1; }}
+.live-agent-filmstrip-thumb.active {{
+ border-color: #2196F3;
+ opacity: 1;
+}}
+
+/* Takeover cursor */
+.live-agent-screenshot-panel.takeover-mode {{
+ cursor: crosshair;
+}}
+.live-agent-screenshot-panel.takeover-mode .live-agent-overlay-layer {{
+ pointer-events: auto;
+}}
+
+/* Overlay controls */
+.live-agent-overlay-controls {{
+ display: flex;
+ gap: 12px;
+ padding: 6px 12px;
+ background: rgba(0,0,0,0.03);
+ border-top: 1px solid #eee;
+ font-size: 12px;
+}}
+.live-agent-overlay-controls label {{
+ display: flex;
+ align-items: center;
+ gap: 4px;
+ cursor: pointer;
+}}
+"""
+
+ def _build_html(
+ self, field_key, config_json, task_desc, start_url,
+ show_controls, show_thought, show_filmstrip,
+ show_overlays, allow_takeover, allow_instructions,
+ max_w, max_h, filmstrip_size,
+ ) -> str:
+ # Build the UI components
+ parts = []
+
+ # Main container
+ parts.append(
+ f''
+ )
+
+ # Status bar
+ parts.append("""
+
+
+
+ Ready
+ Step 0
+
+
+
+ ▲ Collapse
+
+
+""")
+
+ # Start form (shown when no session is active)
+ task_desc_escaped = html.escape(task_desc, quote=True)
+ start_url_escaped = html.escape(start_url, quote=True)
+ parts.append(f"""
+
+
Start Agent Session
+
+
+
+ Start Agent
+
+
+""")
+
+ # Main viewer (hidden until session starts)
+ parts.append('
')
+
+ # Screenshot panel
+ parts.append(f"""
+
+
+ Waiting for agent to start...
+
+
+
+
+""")
+
+ # Side panel
+ parts.append('
')
+
+ # Thought panel
+ if show_thought:
+ parts.append("""
+
+
Agent Thinking
+
Waiting for agent...
+
+""")
+
+ # Step details
+ parts.append("""
+
+""")
+
+ # Controls
+ if show_controls:
+ parts.append('
')
+ parts.append('
')
+ parts.append(
+ 'Pause '
+ )
+ parts.append(
+ 'Resume '
+ )
+ if allow_takeover:
+ parts.append(
+ 'Take Over '
+ )
+ parts.append(
+ 'Stop '
+ )
+ parts.append('
')
+
+ if allow_instructions:
+ parts.append("""
+
+
+ Send
+
+""")
+
+ # Takeover toolbar โ visible only in takeover mode
+ if allow_takeover:
+ parts.append("""
+
+""")
+
+ parts.append('
') # controls
+
+ parts.append('
') # side panel
+ parts.append('
') # main
+
+ # Overlay controls
+ if show_overlays:
+ parts.append("""
+
+ Clicks
+ Bounding Boxes
+ Mouse Path
+ Scroll
+
+""")
+
+ # Filmstrip
+ if show_filmstrip:
+ parts.append(
+ '
'
+ )
+
+ parts.append('
') # viewer
+
+ return "\n".join(parts)
diff --git a/potato/server_utils/displays/live_coding_agent_display.py b/potato/server_utils/displays/live_coding_agent_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..928c164c39921c7c36ab902833694c2551c7fd02
--- /dev/null
+++ b/potato/server_utils/displays/live_coding_agent_display.py
@@ -0,0 +1,134 @@
+"""
+Live Coding Agent Display
+
+Dual-mode display for live coding agent sessions:
+- Live mode (no data): Shows start form + streaming viewer with controls
+- Review mode (data present): Delegates to CodingTraceDisplay
+
+Usage:
+ fields:
+ - key: structured_turns
+ type: live_coding_agent
+ display_options:
+ show_file_tree: true
+ show_reasoning: true
+ collapse_long_outputs: true
+"""
+
+import html
+from typing import Dict, Any, List, Optional
+
+from .base import BaseDisplay
+from .coding_trace_display import CodingTraceDisplay
+
+
+class LiveCodingAgentDisplay(BaseDisplay):
+ """Display type for live coding agent sessions."""
+
+ name = "live_coding_agent"
+ required_fields = ["key"]
+ optional_fields = {
+ "show_file_tree": True,
+ "show_reasoning": True,
+ "collapse_long_outputs": True,
+ "max_output_lines": 50,
+ "show_controls": True,
+ "allow_instructions": True,
+ }
+ description = "Live coding agent viewer with real-time streaming and intervention controls"
+ supports_span_target = False
+
+ def __init__(self):
+ self._coding_trace_display = CodingTraceDisplay()
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ # If data has structured turns, delegate to CodingTraceDisplay (review mode)
+ if data and isinstance(data, (list, dict)):
+ if isinstance(data, list) and len(data) > 0:
+ return self._coding_trace_display.render(field_config, data)
+ if isinstance(data, dict) and data.get("structured_turns"):
+ return self._coding_trace_display.render(
+ field_config, data["structured_turns"]
+ )
+
+ # Live mode: render the viewer UI
+ field_key = html.escape(field_config.get("key", ""), quote=True)
+ options = self.get_display_options(field_config)
+
+ show_controls = options.get("show_controls", True)
+ allow_instructions = options.get("allow_instructions", True)
+
+ controls_html = ""
+ if show_controls:
+ controls_html = f'''
+
+
+ Pause
+
+
+ Resume
+
+
+ Stop
+
+
+ '''
+
+ instruction_html = ""
+ if allow_instructions:
+ instruction_html = f'''
+
+
+
+ Send
+
+
+ '''
+
+ return f'''
+
+
+
+
+
+
+
+ Start Agent
+
+
+
+
+
+
+
+
+ Connecting...
+ 0 turns
+ {controls_html}
+
+
+ {instruction_html}
+
+
+
+
+ Thinking...
+
+
+
+
+
+
+
+ '''
+
+ def has_inline_label(self, field_config: Dict[str, Any]) -> bool:
+ return False
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ classes = super().get_css_classes(field_config)
+ return classes
diff --git a/potato/server_utils/displays/pairwise_display.py b/potato/server_utils/displays/pairwise_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..d9008979b4c073f78361af97a8e2c01b7d4c6d3b
--- /dev/null
+++ b/potato/server_utils/displays/pairwise_display.py
@@ -0,0 +1,200 @@
+"""
+Pairwise Display Type
+
+Renders content in a side-by-side comparison layout for pairwise annotation tasks.
+"""
+
+import html
+from typing import Dict, Any, List
+
+from .base import BaseDisplay
+
+
+class PairwiseDisplay(BaseDisplay):
+ """
+ Display type for pairwise/comparison content.
+
+ Displays multiple items side-by-side for comparison annotation tasks.
+ The data should be a list or dict with multiple items to compare.
+ """
+
+ name = "pairwise"
+ required_fields = ["key"]
+ optional_fields = {
+ "cell_width": "50%",
+ "show_labels": True,
+ "labels": None, # Custom labels like ["Option A", "Option B"]
+ "vertical_on_mobile": True,
+ }
+ description = "Side-by-side comparison display"
+ supports_span_target = False
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render pairwise comparison content as HTML.
+
+ Args:
+ field_config: The field configuration
+ data: The comparison data - should be:
+ - List of items to compare
+ - Dict with keys for each item
+
+ Returns:
+ HTML string for the pairwise display
+ """
+ if not data:
+ return 'No comparison data provided
'
+
+ # Get display options
+ options = self.get_display_options(field_config)
+ cell_width = options.get("cell_width", "50%")
+ show_labels = options.get("show_labels", True)
+ custom_labels = options.get("labels")
+ vertical_on_mobile = options.get("vertical_on_mobile", True)
+
+ field_key = html.escape(field_config.get("key", ""), quote=True)
+
+ # Normalize data to a list of items
+ items = self._normalize_items(data)
+
+ if not items:
+ return 'No items to compare
'
+
+ # Generate labels
+ labels = self._get_labels(items, custom_labels)
+
+ # Calculate cell width based on number of items
+ if cell_width == "auto":
+ cell_width = f"{100 / len(items)}%"
+
+ # Build HTML for each cell
+ cell_html_list = []
+ for i, item in enumerate(items):
+ label = labels[i] if i < len(labels) else f"Option {i + 1}"
+ content = self._render_item(item)
+
+ label_html = ""
+ if show_labels:
+ escaped_label = html.escape(str(label))
+ label_html = f'{escaped_label}
'
+
+ cell_html = f'''
+
+ {label_html}
+
{content}
+
+ '''
+ cell_html_list.append(cell_html)
+
+ # Combine all cells
+ all_cells_html = "\n".join(cell_html_list)
+
+ # Container classes
+ container_classes = ["pairwise-display-content"]
+ if vertical_on_mobile:
+ container_classes.append("vertical-on-mobile")
+
+ return f'''
+
+ {all_cells_html}
+
+ '''
+
+ def _normalize_items(self, data: Any) -> List[Any]:
+ """
+ Normalize data to a list of items.
+
+ Args:
+ data: Raw comparison data
+
+ Returns:
+ List of items to compare
+ """
+ if isinstance(data, list):
+ return data
+ elif isinstance(data, dict):
+ # Return values in order, or use special keys if present
+ if "left" in data and "right" in data:
+ return [data["left"], data["right"]]
+ if "a" in data and "b" in data:
+ return [data["a"], data["b"]]
+ if "1" in data and "2" in data:
+ return [data["1"], data["2"]]
+ # Otherwise return all values
+ return list(data.values())
+ else:
+ # Single item - wrap in list
+ return [data]
+
+ def _get_labels(self, items: List[Any], custom_labels: List[str] = None) -> List[str]:
+ """
+ Generate labels for the comparison items.
+
+ Args:
+ items: The items being compared
+ custom_labels: Custom labels if provided
+
+ Returns:
+ List of label strings
+ """
+ if custom_labels and len(custom_labels) >= len(items):
+ return custom_labels[:len(items)]
+
+ # Default labels
+ if len(items) == 2:
+ return ["A", "B"]
+ else:
+ return [f"Option {i + 1}" for i in range(len(items))]
+
+ def _render_item(self, item: Any) -> str:
+ """
+ Render a single comparison item.
+
+ Args:
+ item: The item to render
+
+ Returns:
+ HTML string for the item
+ """
+ if item is None:
+ return 'No content '
+
+ if isinstance(item, str):
+ # Plain text - escape and preserve newlines
+ escaped = html.escape(item)
+ return escaped.replace('\n', ' ')
+
+ if isinstance(item, dict):
+ # Check for common patterns
+ if "text" in item:
+ text = str(item["text"])
+ escaped = html.escape(text)
+ return escaped.replace('\n', ' ')
+ if "content" in item:
+ content = str(item["content"])
+ escaped = html.escape(content)
+ return escaped.replace('\n', ' ')
+ # Render as key-value pairs
+ parts = []
+ for key, value in item.items():
+ escaped_key = html.escape(str(key))
+ escaped_value = html.escape(str(value))
+ parts.append(f'{escaped_key}: {escaped_value}
')
+ return ''.join(parts)
+
+ if isinstance(item, list):
+ # Render as list items
+ parts = ['']
+ for sub_item in item:
+ escaped = html.escape(str(sub_item))
+ parts.append(f'{escaped} ')
+ parts.append(' ')
+ return ''.join(parts)
+
+ # Default - convert to string
+ return html.escape(str(item))
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ """Get CSS classes for the container."""
+ classes = super().get_css_classes(field_config)
+ return classes
diff --git a/potato/server_utils/displays/pdf_display.py b/potato/server_utils/displays/pdf_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ddb3a7a7c3f562fd6c4890fa4fc53c87f1d99cd
--- /dev/null
+++ b/potato/server_utils/displays/pdf_display.py
@@ -0,0 +1,526 @@
+"""
+PDF Display Component
+
+Renders PDF documents using PDF.js for browser-side rendering.
+Supports span annotation on extracted text.
+
+Usage:
+ In instance_display config:
+ fields:
+ - key: document
+ type: pdf
+ display_options:
+ view_mode: scroll
+ max_height: 700
+ text_layer: true
+"""
+
+from typing import Dict, Any, List, Optional
+import html
+import json
+import logging
+
+from .base import BaseDisplay
+
+logger = logging.getLogger(__name__)
+
+
+class PDFDisplay(BaseDisplay):
+ """
+ Display type for PDF documents.
+
+ Renders PDFs using PDF.js with an optional text layer for span annotation.
+ Can display either a PDF URL or pre-extracted content.
+
+ Supports two annotation modes:
+ - span: Text selection and span annotation (default)
+ - bounding_box: Draw bounding boxes on PDF pages
+ """
+
+ name = "pdf"
+ required_fields = ["key"]
+ optional_fields = {
+ "view_mode": "scroll", # "scroll", "paginated", or "side-by-side"
+ "max_height": 700, # Max container height in pixels
+ "max_width": None, # Max container width
+ "text_layer": True, # Enable text selection layer
+ "show_page_controls": True, # Show page navigation controls
+ "initial_page": 1, # Page to display initially
+ "zoom": "auto", # "auto", "page-fit", "page-width", or percentage
+ "extracted_content": None, # Pre-extracted FormatOutput data
+ "annotation_mode": "span", # "span" or "bounding_box"
+ "bbox_min_size": 10, # Min bounding box size in pixels
+ "bbox_colors": None, # Custom colors for bounding box labels
+ "show_bbox_labels": True, # Show labels on bounding boxes
+ }
+ description = "PDF document display with PDF.js rendering"
+ supports_span_target = False # Uses PDF.js text layer, not .text-content wrapper contract
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render a PDF document.
+
+ Args:
+ field_config: Display configuration
+ data: Either a PDF file path/URL or a dictionary with extracted content
+
+ Returns:
+ HTML string for rendering
+ """
+ options = self.get_display_options(field_config)
+ annotation_mode = options.get("annotation_mode", "span")
+
+ # Check if data is pre-extracted content or a file path
+ if isinstance(data, dict):
+ # Pre-extracted content from format handler
+ if annotation_mode == "bounding_box":
+ return self._render_extracted_bbox(data, options, field_config)
+ return self._render_extracted(data, options, field_config)
+ else:
+ # File path or URL - render with PDF.js
+ if annotation_mode == "bounding_box":
+ return self._render_pdfjs_bbox(str(data), options, field_config)
+ return self._render_pdfjs(str(data), options, field_config)
+
+ def _render_pdfjs(
+ self,
+ pdf_source: str,
+ options: Dict[str, Any],
+ field_config: Dict[str, Any]
+ ) -> str:
+ """
+ Render PDF using PDF.js viewer.
+ """
+ field_key = field_config.get("key", "pdf")
+ view_mode = options.get("view_mode", "scroll")
+ max_height = options.get("max_height", 700)
+ max_width = options.get("max_width")
+ text_layer = options.get("text_layer", True)
+ show_controls = options.get("show_page_controls", True)
+ zoom = options.get("zoom", "auto")
+ initial_page = options.get("initial_page", 1)
+
+ # Build style string
+ styles = []
+ if max_height:
+ styles.append(f"max-height: {max_height}px")
+ if max_width:
+ styles.append(f"max-width: {max_width}px")
+ style_str = "; ".join(styles) if styles else ""
+
+ # Build container
+ parts = []
+
+ # Container div with PDF.js viewer
+ parts.append(
+ f''
+ )
+
+ # Page controls
+ if show_controls:
+ parts.append('''
+
+ ◀
+
+ 1 /
+ -
+
+ ▶
+
+ Auto
+ Page Fit
+ Page Width
+ 50%
+ 75%
+ 100%
+ 125%
+ 150%
+ 200%
+
+
+ ''')
+
+ # Canvas container for PDF.js
+ parts.append('''
+
+
+ ''')
+
+ # Text layer for selection (if enabled)
+ if text_layer:
+ parts.append('
')
+
+ parts.append('
') # Close canvas container
+
+ # Loading indicator
+ parts.append('''
+
+
+ Loading PDF...
+
+ ''')
+
+ # Error display
+ parts.append('
')
+
+ parts.append('
') # Close main container
+
+ return "\n".join(parts)
+
+ def _render_extracted(
+ self,
+ content: Dict[str, Any],
+ options: Dict[str, Any],
+ field_config: Dict[str, Any]
+ ) -> str:
+ """
+ Render pre-extracted PDF content as HTML.
+
+ This is used when the PDF has already been processed by the
+ format handler and we want to display the extracted text.
+ """
+ field_key = field_config.get("key", "pdf")
+ max_height = options.get("max_height", 700)
+
+ # Check if this is FormatOutput-style content
+ if "rendered_html" in content:
+ inner_html = content["rendered_html"]
+ elif "text" in content:
+ # Fall back to plain text
+ inner_html = f''
+ else:
+ inner_html = 'No content available
'
+
+ # Build container
+ style_str = f"max-height: {max_height}px; overflow-y: auto;" if max_height else ""
+
+ # Add metadata if available
+ metadata_html = ""
+ if "metadata" in content:
+ meta = content["metadata"]
+ if "total_pages" in meta:
+ metadata_html = f'Pages: {meta["total_pages"]}
'
+
+ return f'''
+
+ '''
+
+ def _render_pdfjs_bbox(
+ self,
+ pdf_source: str,
+ options: Dict[str, Any],
+ field_config: Dict[str, Any]
+ ) -> str:
+ """
+ Render PDF with bounding box annotation support.
+
+ Uses paginated view mode by default for better bbox drawing experience.
+ """
+ field_key = field_config.get("key", "pdf")
+ # Force paginated mode for bounding box annotation
+ view_mode = "paginated"
+ max_height = options.get("max_height", 700)
+ max_width = options.get("max_width")
+ show_controls = options.get("show_page_controls", True)
+ zoom = options.get("zoom", "page-fit")
+ initial_page = options.get("initial_page", 1)
+ bbox_min_size = options.get("bbox_min_size", 10)
+ show_labels = options.get("show_bbox_labels", True)
+
+ # Build style string
+ styles = []
+ if max_height:
+ styles.append(f"max-height: {max_height}px")
+ if max_width:
+ styles.append(f"max-width: {max_width}px")
+ style_str = "; ".join(styles) if styles else ""
+
+ parts = []
+
+ # Container div with bbox annotation mode
+ parts.append(
+ f''
+ )
+
+ # Enhanced page controls for paginated navigation
+ if show_controls:
+ parts.append('''
+
+
+ ⏪
+ ◀
+
+ Page of
+ -
+
+ ▶
+ ⏩
+
+
+
+ □ Draw
+
+
+ ↑ Select
+
+
+ ✕ Delete
+
+
+
+ Page Fit
+ Page Width
+ 50%
+ 75%
+ 100%
+ 125%
+ 150%
+ 200%
+
+
+ ''')
+
+ # Canvas container for PDF rendering and bbox overlay
+ parts.append('''
+
+ ''')
+
+ # Bounding box info panel
+ parts.append('''
+
+
Boxes on page: 0
+
Total boxes: 0
+
+ ''')
+
+ # Loading indicator
+ parts.append('''
+
+
+ Loading PDF...
+
+ ''')
+
+ # Error display
+ parts.append('
')
+
+ # Hidden input that carries drawn boxes through the standard save
+ # pipeline (collected as "{name}:::_data", restored by the server). F-040.
+ parts.append(
+ f'
'
+ )
+
+ parts.append('
') # Close main container
+
+ return "\n".join(parts)
+
+ def _render_extracted_bbox(
+ self,
+ content: Dict[str, Any],
+ options: Dict[str, Any],
+ field_config: Dict[str, Any]
+ ) -> str:
+ """
+ Render pre-extracted PDF content with bounding box support.
+
+ For pre-extracted content, we render page images for bbox annotation.
+ """
+ field_key = field_config.get("key", "pdf")
+ max_height = options.get("max_height", 700)
+ bbox_min_size = options.get("bbox_min_size", 10)
+ show_labels = options.get("show_bbox_labels", True)
+
+ metadata = content.get("metadata", {})
+ total_pages = metadata.get("total_pages", 1)
+ pages = metadata.get("pages", [])
+
+ style_str = f"max-height: {max_height}px; overflow-y: auto;" if max_height else ""
+
+ parts = []
+ parts.append(
+ f'')
+
+ return "\n".join(parts)
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ """Get CSS classes for the display container."""
+ classes = super().get_css_classes(field_config)
+ options = self.get_display_options(field_config)
+
+ if field_config.get("span_target"):
+ classes.append("span-target-pdf")
+
+ view_mode = options.get("view_mode", "scroll")
+ classes.append(f"pdf-mode-{view_mode}")
+
+ annotation_mode = options.get("annotation_mode", "span")
+ if annotation_mode == "bounding_box":
+ classes.append("pdf-bbox-annotation")
+
+ return classes
+
+ def get_data_attributes(
+ self,
+ field_config: Dict[str, Any],
+ data: Any
+ ) -> Dict[str, str]:
+ """Get data attributes for JavaScript initialization."""
+ attrs = super().get_data_attributes(field_config, data)
+ options = self.get_display_options(field_config)
+
+ attrs["view-mode"] = options.get("view_mode", "scroll")
+ attrs["text-layer"] = str(options.get("text_layer", True)).lower()
+
+ return attrs
+
+ def get_js_init(self) -> Optional[str]:
+ """
+ Return JavaScript initialization code for PDF.js.
+ """
+ return '''
+ // PDF display initialization is handled by pdf-viewer.js
+ // which is loaded as a separate script
+ if (typeof initPDFViewers === 'function') {
+ initPDFViewers();
+ }
+ '''
+
+ def validate_config(self, field_config: Dict[str, Any]) -> List[str]:
+ """Validate the field configuration."""
+ errors = super().validate_config(field_config)
+ options = field_config.get("display_options", {})
+
+ # Validate view_mode
+ valid_modes = ["scroll", "paginated", "side-by-side"]
+ view_mode = options.get("view_mode", "scroll")
+ if view_mode not in valid_modes:
+ errors.append(
+ f"Invalid view_mode '{view_mode}'. "
+ f"Must be one of: {', '.join(valid_modes)}"
+ )
+
+ # Validate annotation_mode
+ valid_annotation_modes = ["span", "bounding_box"]
+ annotation_mode = options.get("annotation_mode", "span")
+ if annotation_mode not in valid_annotation_modes:
+ errors.append(
+ f"Invalid annotation_mode '{annotation_mode}'. "
+ f"Must be one of: {', '.join(valid_annotation_modes)}"
+ )
+
+ # Validate zoom
+ zoom = options.get("zoom", "auto")
+ valid_zoom = ["auto", "page-fit", "page-width"]
+ if zoom not in valid_zoom:
+ try:
+ float(zoom)
+ except (TypeError, ValueError):
+ errors.append(
+ f"Invalid zoom '{zoom}'. "
+ f"Must be one of {valid_zoom} or a number"
+ )
+
+ return errors
diff --git a/potato/server_utils/displays/registry.py b/potato/server_utils/displays/registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..333b3ed32114505ddfe2411463e84d7f2b088fc5
--- /dev/null
+++ b/potato/server_utils/displays/registry.py
@@ -0,0 +1,661 @@
+"""
+Display Registry
+
+Provides a centralized registry for managing display types.
+This module serves as the single source of truth for available display
+types and their renderers, separating content display from annotation collection.
+
+Usage:
+ from potato.server_utils.displays.registry import display_registry
+
+ # Render content
+ html = display_registry.render("image", field_config, data)
+
+ # List all available display types
+ types = display_registry.get_supported_types()
+
+ # Register a custom display type (plugin support)
+ display_registry.register_plugin("my_custom", MyCustomDisplay())
+"""
+
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Any, Optional, Union
+import logging
+
+from .base import BaseDisplay, render_display_container
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DisplayDefinition:
+ """
+ Defines metadata and renderer for a display type.
+
+ Attributes:
+ name: Unique identifier for the display type (e.g., "text", "image")
+ renderer: Either a BaseDisplay instance or callable that renders content
+ required_fields: List of required configuration fields
+ optional_fields: Dictionary of optional fields with default values
+ supports_span_target: Whether this type can be a span annotation target
+ description: Human-readable description of the display type
+ lazy_populated: When True, the display's data key is populated
+ after initial render (e.g. ``interactive_chat``'s conversation
+ is written only after the user finishes chatting). The
+ ``instance_display`` validator uses this to decide whether a
+ missing data key is a transient state (lazy) or a config
+ error (not lazy).
+ """
+ name: str
+ renderer: Union[BaseDisplay, Callable[[Dict[str, Any], Any], str]]
+ required_fields: List[str] = field(default_factory=list)
+ optional_fields: Dict[str, Any] = field(default_factory=dict)
+ supports_span_target: bool = False
+ description: str = ""
+ lazy_populated: bool = False
+
+
+class DisplayRegistry:
+ """
+ Centralized registry for display types.
+
+ Provides methods to register, retrieve, and render display types.
+ Supports both built-in display types and custom plugins.
+ """
+
+ def __init__(self):
+ self._displays: Dict[str, DisplayDefinition] = {}
+ self._plugins: Dict[str, BaseDisplay] = {}
+ logger.debug("DisplayRegistry initialized")
+
+ def register(self, display: DisplayDefinition) -> None:
+ """
+ Register a built-in display type.
+
+ Args:
+ display: DisplayDefinition to register
+
+ Raises:
+ ValueError: If a display with the same name is already registered
+ """
+ if display.name in self._displays:
+ raise ValueError(f"Display type '{display.name}' is already registered")
+
+ self._displays[display.name] = display
+ logger.debug(f"Registered display type: {display.name}")
+
+ def register_plugin(self, name: str, plugin: BaseDisplay) -> None:
+ """
+ Register a custom display type from a plugin.
+
+ Args:
+ name: Unique name for the display type
+ plugin: BaseDisplay instance implementing the display logic
+
+ Raises:
+ ValueError: If a display with the same name is already registered
+ """
+ if name in self._displays or name in self._plugins:
+ raise ValueError(f"Display type '{name}' is already registered")
+
+ self._plugins[name] = plugin
+ logger.debug(f"Registered plugin display type: {name}")
+
+ def get(self, name: str) -> Optional[Union[DisplayDefinition, BaseDisplay]]:
+ """
+ Get a display definition or plugin by name.
+
+ Args:
+ name: The display type name
+
+ Returns:
+ DisplayDefinition or BaseDisplay if found, None otherwise
+ """
+ if name in self._displays:
+ return self._displays[name]
+ return self._plugins.get(name)
+
+ def render(self, field_type: str, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render content using the appropriate display type.
+
+ Args:
+ field_type: The display type name (e.g., "text", "image")
+ field_config: Configuration for this field from instance_display.fields
+ data: The actual data value from the instance
+
+ Returns:
+ HTML string for rendering the content
+
+ Raises:
+ ValueError: If the display type is not registered
+ """
+ # Check plugins first (allows overriding built-ins)
+ if field_type in self._plugins:
+ plugin = self._plugins[field_type]
+ inner_html = plugin.render(field_config, data)
+ css_classes = plugin.get_css_classes(field_config)
+ data_attrs = plugin.get_data_attributes(field_config, data)
+ label = None if plugin.has_inline_label(field_config) else field_config.get("label")
+ return render_display_container(inner_html, css_classes, data_attrs, label)
+
+ # Check built-in displays
+ if field_type in self._displays:
+ display = self._displays[field_type]
+ renderer = display.renderer
+
+ # Handle BaseDisplay instances
+ if isinstance(renderer, BaseDisplay):
+ inner_html = renderer.render(field_config, data)
+ css_classes = renderer.get_css_classes(field_config)
+ data_attrs = renderer.get_data_attributes(field_config, data)
+ label = None if renderer.has_inline_label(field_config) else field_config.get("label")
+ return render_display_container(inner_html, css_classes, data_attrs, label)
+
+ # Handle callable renderers
+ return renderer(field_config, data)
+
+ supported = ", ".join(sorted(self.get_supported_types()))
+ raise ValueError(
+ f"Unknown display type: '{field_type}'. "
+ f"Supported types are: {supported}"
+ )
+
+ def validate_config(self, field_type: str, field_config: Dict[str, Any]) -> List[str]:
+ """
+ Validate display configuration.
+
+ Args:
+ field_type: The display type name
+ field_config: The field configuration to validate
+
+ Returns:
+ List of error messages (empty if valid)
+ """
+ errors = []
+
+ # Check if display type exists
+ display = self.get(field_type)
+ if not display:
+ errors.append(f"Unknown display type: '{field_type}'")
+ return errors
+
+ # Get required fields
+ if isinstance(display, BaseDisplay):
+ errors.extend(display.validate_config(field_config))
+ elif isinstance(display, DisplayDefinition):
+ for req_field in display.required_fields:
+ if req_field not in field_config:
+ errors.append(
+ f"Missing required field '{req_field}' for display type '{field_type}'"
+ )
+
+ return errors
+
+ def type_supports_span_target(self, field_type: str) -> bool:
+ """
+ Check if a display type supports span annotation.
+
+ Args:
+ field_type: The display type name
+
+ Returns:
+ True if the display type supports span_target
+ """
+ if field_type in self._plugins:
+ return self._plugins[field_type].supports_span_target
+ if field_type in self._displays:
+ display = self._displays[field_type]
+ if isinstance(display.renderer, BaseDisplay):
+ return display.renderer.supports_span_target
+ return display.supports_span_target
+ return False
+
+ def get_span_target_types(self) -> List[str]:
+ """
+ Get all display type names that support span annotation.
+
+ Returns:
+ List of type names where supports_span_target is True
+ """
+ types = []
+ for name in self.get_supported_types():
+ if self.type_supports_span_target(name):
+ types.append(name)
+ return sorted(types)
+
+ def list_displays(self) -> List[Dict[str, Any]]:
+ """
+ List all registered displays with their metadata.
+
+ Returns:
+ List of dictionaries containing display metadata
+ """
+ result = []
+
+ # Add built-in displays
+ for display in sorted(self._displays.values(), key=lambda d: d.name):
+ result.append({
+ "name": display.name,
+ "description": display.description,
+ "required_fields": display.required_fields,
+ "optional_fields": list(display.optional_fields.keys()),
+ "supports_span_target": display.supports_span_target,
+ "is_plugin": False,
+ })
+
+ # Add plugins
+ for name, plugin in sorted(self._plugins.items()):
+ result.append({
+ "name": name,
+ "description": plugin.description,
+ "required_fields": plugin.required_fields,
+ "optional_fields": list(plugin.optional_fields.keys()),
+ "supports_span_target": plugin.supports_span_target,
+ "is_plugin": True,
+ })
+
+ return result
+
+ def is_lazy_populated(self, name: str) -> bool:
+ """
+ Check whether a display type populates its data lazily (after
+ initial render). Used by ``instance_display._validate_fields`` to
+ distinguish an expected transient missing key (lazy) from a real
+ configuration error (not lazy).
+
+ Args:
+ name: The display type name.
+
+ Returns:
+ True iff the display's ``lazy_populated`` attribute is set. False
+ for unknown types (treat as strict for safety).
+ """
+ plugin = self._plugins.get(name)
+ if plugin is not None:
+ return bool(getattr(plugin, "lazy_populated", False))
+ display = self._displays.get(name)
+ if display is not None:
+ # Prefer the DisplayDefinition flag; fall back to the BaseDisplay
+ # class attr when a definition didn't explicitly set it.
+ if getattr(display, "lazy_populated", False):
+ return True
+ renderer = getattr(display, "renderer", None)
+ return bool(getattr(renderer, "lazy_populated", False))
+ return False
+
+ def is_registered(self, name: str) -> bool:
+ """
+ Check if a display type is registered.
+
+ Args:
+ name: The display type name
+
+ Returns:
+ True if registered, False otherwise
+ """
+ return name in self._displays or name in self._plugins
+
+ def get_supported_types(self) -> List[str]:
+ """
+ Get a list of all supported display types.
+
+ Returns:
+ Sorted list of display type names
+ """
+ types = set(self._displays.keys()) | set(self._plugins.keys())
+ return sorted(types)
+
+ def supports_span_target(self, name: str) -> bool:
+ """
+ Check if a display type supports span annotation targeting.
+
+ Args:
+ name: The display type name
+
+ Returns:
+ True if the type supports span targets, False otherwise
+ """
+ if name in self._plugins:
+ return self._plugins[name].supports_span_target
+ if name in self._displays:
+ return self._displays[name].supports_span_target
+ return False
+
+
+# Global registry instance
+display_registry = DisplayRegistry()
+
+
+def _register_builtin_displays():
+ """
+ Register all built-in display types.
+ Called automatically when this module is imported.
+ """
+ from .text_display import TextDisplay
+ from .image_display import ImageDisplay
+ from .video_display import VideoDisplay
+ from .audio_display import AudioDisplay
+ from .dialogue_display import DialogueDisplay
+ from .pairwise_display import PairwiseDisplay
+ from .pdf_display import PDFDisplay
+ from .document_display import DocumentDisplay
+ from .spreadsheet_display import SpreadsheetDisplay
+ from .code_display import CodeDisplay
+ from .conversation_tree_display import ConversationTreeDisplay
+ from .agent_trace_display import AgentTraceDisplay
+ from .eval_trace_display import EvalTraceDisplay
+ from .gallery_display import GalleryDisplay
+ from .interactive_chat_display import InteractiveChatDisplay
+ from .web_agent_trace_display import WebAgentTraceDisplay
+ from .live_agent_display import LiveAgentDisplay
+ from .coding_trace_display import CodingTraceDisplay
+ from .live_coding_agent_display import LiveCodingAgentDisplay
+
+ displays = [
+ DisplayDefinition(
+ name="text",
+ renderer=TextDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "collapsible": False,
+ "max_height": None,
+ "preserve_whitespace": True,
+ },
+ supports_span_target=True,
+ description="Plain text content display"
+ ),
+ DisplayDefinition(
+ name="html",
+ renderer=TextDisplay(allow_html=True),
+ required_fields=["key"],
+ optional_fields={
+ "collapsible": False,
+ "max_height": None,
+ },
+ supports_span_target=False,
+ description="HTML content display (sanitized)"
+ ),
+ DisplayDefinition(
+ name="image",
+ renderer=ImageDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "max_width": None,
+ "max_height": None,
+ "zoomable": True,
+ "alt_text": "",
+ },
+ supports_span_target=False,
+ description="Image display with optional zoom"
+ ),
+ DisplayDefinition(
+ name="video",
+ renderer=VideoDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "max_width": None,
+ "max_height": None,
+ "controls": True,
+ "autoplay": False,
+ "loop": False,
+ "muted": False,
+ },
+ supports_span_target=False,
+ description="Video player display"
+ ),
+ DisplayDefinition(
+ name="audio",
+ renderer=AudioDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "controls": True,
+ "autoplay": False,
+ "loop": False,
+ "show_waveform": False,
+ },
+ supports_span_target=False,
+ description="Audio player display"
+ ),
+ DisplayDefinition(
+ name="dialogue",
+ renderer=DialogueDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "alternating_shading": True,
+ "speaker_extraction": True,
+ "show_turn_numbers": False,
+ },
+ supports_span_target=True,
+ description="Dialogue/conversation turns display"
+ ),
+ DisplayDefinition(
+ name="pairwise",
+ renderer=PairwiseDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "cell_width": "50%",
+ "show_labels": True,
+ "vertical_on_mobile": True,
+ },
+ supports_span_target=False,
+ description="Side-by-side comparison display"
+ ),
+ DisplayDefinition(
+ name="pdf",
+ renderer=PDFDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "view_mode": "scroll",
+ "max_height": 700,
+ "max_width": None,
+ "text_layer": True,
+ "show_page_controls": True,
+ "initial_page": 1,
+ "zoom": "auto",
+ },
+ supports_span_target=False,
+ description="PDF document display with PDF.js rendering"
+ ),
+ DisplayDefinition(
+ name="document",
+ renderer=DocumentDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "collapsible": False,
+ "max_height": None,
+ "show_outline": False,
+ "preserve_structure": True,
+ "style_theme": "default",
+ },
+ supports_span_target=True,
+ description="Document display for DOCX, Markdown, and other formats"
+ ),
+ DisplayDefinition(
+ name="spreadsheet",
+ renderer=SpreadsheetDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "annotation_mode": "row",
+ "show_headers": True,
+ "max_height": 400,
+ "max_width": None,
+ "striped": True,
+ "hoverable": True,
+ "sortable": False,
+ "selectable": True,
+ "compact": False,
+ },
+ supports_span_target=False,
+ description="Spreadsheet/table display with row or cell annotation"
+ ),
+ DisplayDefinition(
+ name="code",
+ renderer=CodeDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "language": None,
+ "show_line_numbers": True,
+ "max_height": 500,
+ "max_width": None,
+ "wrap_lines": False,
+ "highlight_lines": None,
+ "start_line": 1,
+ "theme": "default",
+ "copy_button": True,
+ },
+ supports_span_target=True,
+ description="Source code display with syntax highlighting"
+ ),
+ DisplayDefinition(
+ name="conversation_tree",
+ renderer=ConversationTreeDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "collapsed_depth": 2,
+ "node_style": "card",
+ "show_node_ids": False,
+ "max_depth": None,
+ },
+ supports_span_target=False,
+ description="Conversation tree with collapsible branching nodes"
+ ),
+ DisplayDefinition(
+ name="agent_trace",
+ renderer=AgentTraceDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "show_timestamps": False,
+ "collapse_observations": False,
+ "step_type_colors": None,
+ "show_screenshots": True,
+ "show_step_numbers": True,
+ "show_summary": True,
+ "compact": False,
+ },
+ supports_span_target=False,
+ description="Agent trace display with step cards and type badges"
+ ),
+ DisplayDefinition(
+ name="eval_trace",
+ renderer=EvalTraceDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "pane_labels": None,
+ "show_step_numbers": True,
+ "collapse_long_outputs": True,
+ "max_output_lines": 20,
+ "link_steps": True,
+ "compact": False,
+ },
+ supports_span_target=False,
+ description="Three-pane agent trace eval: reasoning, function calls, and final answer side-by-side"
+ ),
+ DisplayDefinition(
+ name="gallery",
+ renderer=GalleryDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "layout": "horizontal",
+ "thumbnail_size": 300,
+ "show_captions": True,
+ "zoomable": True,
+ "max_height": 400,
+ "columns": 3,
+ },
+ supports_span_target=False,
+ description="Scrollable image gallery with captions"
+ ),
+ DisplayDefinition(
+ name="interactive_chat",
+ renderer=InteractiveChatDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "placeholder_text": "Start chatting with the agent to begin the task.",
+ },
+ supports_span_target=True,
+ lazy_populated=True, # conversation is written by /agent_chat/finish
+ description="Interactive agent chat with post-interaction trace display"
+ ),
+ DisplayDefinition(
+ name="web_agent_trace",
+ renderer=WebAgentTraceDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "show_overlays": True,
+ "show_filmstrip": True,
+ "show_thought": True,
+ "show_observation": True,
+ "show_element_info": True,
+ "screenshot_max_width": 800,
+ "screenshot_max_height": 600,
+ "filmstrip_size": 80,
+ "auto_playback": False,
+ "playback_step_delay": 2.0,
+ },
+ supports_span_target=False,
+ description="Web agent trace viewer with screenshots, SVG overlays, and step navigation"
+ ),
+ DisplayDefinition(
+ name="coding_trace",
+ renderer=CodingTraceDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "show_file_tree": True,
+ "diff_view": "unified",
+ "collapse_long_outputs": True,
+ "max_output_lines": 50,
+ "terminal_theme": "dark",
+ "show_step_numbers": True,
+ "show_tool_badges": True,
+ "show_reasoning": True,
+ "compact": False,
+ },
+ supports_span_target=True,
+ description="Coding agent trace display with diff rendering, terminal blocks, and file tree"
+ ),
+ DisplayDefinition(
+ name="live_agent",
+ renderer=LiveAgentDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "show_overlays": True,
+ "show_filmstrip": True,
+ "show_thought": True,
+ "show_controls": True,
+ "allow_takeover": True,
+ "allow_instructions": True,
+ "screenshot_max_width": 900,
+ "screenshot_max_height": 650,
+ "filmstrip_size": 80,
+ },
+ supports_span_target=False,
+ lazy_populated=True, # trace populated by live agent session
+ description="Live AI agent viewer with real-time screenshots, controls, and interaction"
+ ),
+ DisplayDefinition(
+ name="live_coding_agent",
+ renderer=LiveCodingAgentDisplay(),
+ required_fields=["key"],
+ optional_fields={
+ "show_file_tree": True,
+ "show_reasoning": True,
+ "collapse_long_outputs": True,
+ "max_output_lines": 50,
+ "show_controls": True,
+ "allow_instructions": True,
+ },
+ supports_span_target=False,
+ lazy_populated=True, # trace populated by live coding-agent session
+ description="Live coding agent viewer with real-time streaming and intervention controls"
+ ),
+ ]
+
+ for display in displays:
+ display_registry.register(display)
+
+ logger.debug(f"Registered {len(displays)} built-in display types")
+
+
+# Auto-register built-in displays on import
+_register_builtin_displays()
diff --git a/potato/server_utils/displays/spreadsheet_display.py b/potato/server_utils/displays/spreadsheet_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..971da7082607178219af4849bd3598a92726a2b7
--- /dev/null
+++ b/potato/server_utils/displays/spreadsheet_display.py
@@ -0,0 +1,342 @@
+"""
+Spreadsheet Display Component
+
+Renders tabular data with support for row-based or cell-based annotation.
+
+Usage:
+ In instance_display config:
+ fields:
+ - key: data_table
+ type: spreadsheet
+ display_options:
+ annotation_mode: row
+ show_headers: true
+ max_height: 400
+"""
+
+from typing import Dict, Any, List, Optional
+import html
+import logging
+
+from .base import BaseDisplay
+
+logger = logging.getLogger(__name__)
+
+
+class SpreadsheetDisplay(BaseDisplay):
+ """
+ Display type for tabular/spreadsheet data.
+
+ Renders data as an HTML table with support for row-based
+ or cell-based annotation modes.
+ """
+
+ name = "spreadsheet"
+ required_fields = ["key"]
+ optional_fields = {
+ "annotation_mode": "row", # "row", "cell", or "range"
+ "show_headers": True, # Show column headers
+ "max_height": 400, # Max container height
+ "max_width": None, # Max container width
+ "striped": True, # Alternating row colors
+ "hoverable": True, # Highlight row on hover
+ "sortable": False, # Enable column sorting
+ "filterable": False, # Enable column filtering
+ "selectable": True, # Enable row/cell selection
+ "compact": False, # Compact table styling
+ # New styling options
+ "border_style": "default", # "default", "bordered", "minimal", "rounded", "none"
+ "header_style": "default", # "default", "dark", "primary", "gradient", "light", "transparent"
+ "custom_class": None, # Additional CSS classes for the table
+ "custom_css": None, # Inline CSS styles for the table
+ }
+ description = "Spreadsheet/table display with row or cell annotation"
+ supports_span_target = False # No .text-content wrapper; table data not compatible with span offsets
+
+ # Valid values for style options
+ VALID_BORDER_STYLES = ["default", "bordered", "minimal", "rounded", "none"]
+ VALID_HEADER_STYLES = ["default", "dark", "primary", "gradient", "light", "transparent"]
+
+ def validate_config(self, field_config: Dict[str, Any]) -> List[str]:
+ """
+ Validate spreadsheet display configuration.
+
+ Returns:
+ List of validation error messages (empty if valid)
+ """
+ errors = super().validate_config(field_config)
+ options = field_config.get("display_options", {})
+
+ # Validate border_style
+ border_style = options.get("border_style", "default")
+ if border_style not in self.VALID_BORDER_STYLES:
+ errors.append(
+ f"Invalid border_style '{border_style}'. "
+ f"Must be one of: {', '.join(self.VALID_BORDER_STYLES)}"
+ )
+
+ # Validate header_style
+ header_style = options.get("header_style", "default")
+ if header_style not in self.VALID_HEADER_STYLES:
+ errors.append(
+ f"Invalid header_style '{header_style}'. "
+ f"Must be one of: {', '.join(self.VALID_HEADER_STYLES)}"
+ )
+
+ # Validate annotation_mode
+ annotation_mode = options.get("annotation_mode", "row")
+ valid_modes = ["row", "cell", "range"]
+ if annotation_mode not in valid_modes:
+ errors.append(
+ f"Invalid annotation_mode '{annotation_mode}'. "
+ f"Must be one of: {', '.join(valid_modes)}"
+ )
+
+ return errors
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render spreadsheet data.
+
+ Args:
+ field_config: Display configuration
+ data: Either a dict with extracted content, list of lists,
+ or list of dicts
+
+ Returns:
+ HTML string for rendering
+ """
+ options = self.get_display_options(field_config)
+ field_key = field_config.get("key", "spreadsheet")
+
+ # Handle different data formats
+ if isinstance(data, dict):
+ # Pre-extracted FormatOutput data
+ if "rendered_html" in data:
+ return self._wrap_content(data["rendered_html"], options, field_key)
+
+ # Extract from metadata
+ rows = data.get("rows", [])
+ headers = data.get("headers", data.get("metadata", {}).get("headers", []))
+ elif isinstance(data, list):
+ if data and isinstance(data[0], dict):
+ # List of dictionaries
+ headers = list(data[0].keys()) if data else []
+ rows = [[row.get(h, "") for h in headers] for row in data]
+ else:
+ # List of lists
+ rows = data
+ headers = []
+ else:
+ return f'Unsupported data format
'
+
+ # Generate table HTML
+ table_html = self._render_table(rows, headers, options, field_key)
+ return self._wrap_content(table_html, options, field_key)
+
+ def _wrap_content(
+ self,
+ content: str,
+ options: Dict[str, Any],
+ field_key: str
+ ) -> str:
+ """
+ Wrap table content in container with styles.
+ """
+ styles = []
+ max_height = options.get("max_height")
+ max_width = options.get("max_width")
+
+ if max_height:
+ styles.append(f"max-height: {max_height}px")
+ styles.append("overflow-y: auto")
+ if max_width:
+ styles.append(f"max-width: {max_width}px")
+ styles.append("overflow-x: auto")
+
+ style_str = "; ".join(styles) if styles else ""
+ mode = options.get("annotation_mode", "row")
+
+ # Container classes
+ container_classes = ["spreadsheet-display"]
+
+ # Add border-rounded class to container for rounded style
+ border_style = options.get("border_style", "default")
+ if border_style == "rounded":
+ container_classes.append("border-rounded")
+
+ container_class_str = " ".join(container_classes)
+
+ return f'''
+
+ {content}
+
+ '''
+
+ def _render_table(
+ self,
+ rows: List[List],
+ headers: List[str],
+ options: Dict[str, Any],
+ field_key: str
+ ) -> str:
+ """
+ Render data as HTML table.
+ """
+ parts = []
+ mode = options.get("annotation_mode", "row")
+
+ # Table classes
+ table_classes = ["spreadsheet-table"]
+ if options.get("striped"):
+ table_classes.append("table-striped")
+ if options.get("hoverable"):
+ table_classes.append("table-hoverable")
+ if options.get("compact"):
+ table_classes.append("table-compact")
+ if options.get("selectable"):
+ table_classes.append("table-selectable")
+
+ # Border style class
+ border_style = options.get("border_style", "default")
+ if border_style and border_style != "default":
+ table_classes.append(f"border-{border_style}")
+
+ # Header style class
+ header_style = options.get("header_style", "default")
+ if header_style and header_style != "default":
+ table_classes.append(f"header-{header_style}")
+
+ # Custom class from admin config
+ custom_class = options.get("custom_class")
+ if custom_class:
+ # Support both string and list of classes
+ if isinstance(custom_class, list):
+ table_classes.extend(custom_class)
+ else:
+ table_classes.append(custom_class)
+
+ class_str = " ".join(table_classes)
+
+ # Custom inline CSS
+ custom_css = options.get("custom_css", "")
+ style_attr = f' style="{html.escape(custom_css)}"' if custom_css else ""
+
+ parts.append(f'')
+
+ # Add selection summary for row mode
+ if options.get("selectable") and mode == "row":
+ parts.append('''
+
+ 0 rows selected
+
+ ''')
+
+ return "\n".join(parts)
+
+ def _get_cell_ref(self, row: int, col: int) -> str:
+ """
+ Get A1-style cell reference.
+ """
+ # Convert column to letter
+ col_letter = ""
+ col_num = col + 1
+ while col_num > 0:
+ col_num, remainder = divmod(col_num - 1, 26)
+ col_letter = chr(65 + remainder) + col_letter
+ return f"{col_letter}{row + 1}"
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ """Get CSS classes for the display container."""
+ classes = super().get_css_classes(field_config)
+ options = self.get_display_options(field_config)
+
+ if field_config.get("span_target"):
+ classes.append("span-target-spreadsheet")
+
+ mode = options.get("annotation_mode", "row")
+ classes.append(f"spreadsheet-mode-{mode}")
+
+ return classes
+
+ def get_data_attributes(
+ self,
+ field_config: Dict[str, Any],
+ data: Any
+ ) -> Dict[str, str]:
+ """Get data attributes for JavaScript initialization."""
+ attrs = super().get_data_attributes(field_config, data)
+ options = self.get_display_options(field_config)
+
+ attrs["annotation-mode"] = options.get("annotation_mode", "row")
+ attrs["selectable"] = str(options.get("selectable", True)).lower()
+
+ return attrs
+
+ def get_js_init(self) -> Optional[str]:
+ """
+ Return JavaScript initialization code for spreadsheet interactivity.
+ """
+ return '''
+ if (typeof initSpreadsheetDisplays === 'function') {
+ initSpreadsheetDisplays();
+ }
+ '''
diff --git a/potato/server_utils/displays/text_display.py b/potato/server_utils/displays/text_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..250d15485eaaeb5156257e0d7cbac20c404702cd
--- /dev/null
+++ b/potato/server_utils/displays/text_display.py
@@ -0,0 +1,185 @@
+"""
+Text Display Type
+
+Renders plain text or HTML content for display in the annotation interface.
+Supports span annotation targeting when used with span annotation schemes.
+"""
+
+import html
+from typing import Dict, Any, List
+
+from .base import BaseDisplay
+
+
+class TextDisplay(BaseDisplay):
+ """
+ Display type for text content.
+
+ Supports both plain text (with HTML escaping) and sanitized HTML content.
+ Can be used as a target for span annotations.
+ """
+
+ name = "text"
+ required_fields = ["key"]
+ optional_fields = {
+ "collapsible": False,
+ "max_height": None,
+ "preserve_whitespace": True,
+ }
+ description = "Plain text content display"
+ supports_span_target = True
+
+ def __init__(self, allow_html: bool = False):
+ """
+ Initialize the text display.
+
+ Args:
+ allow_html: If True, render as sanitized HTML. If False, escape all HTML.
+ """
+ self.allow_html = allow_html
+ if allow_html:
+ self.name = "html"
+ self.description = "HTML content display (sanitized)"
+ self.supports_span_target = False
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render text content as HTML.
+
+ Args:
+ field_config: The field configuration
+ data: The text content to display
+
+ Returns:
+ HTML string for the text content
+ """
+ if data is None:
+ return 'No content '
+
+ # Convert to string if needed
+ text = str(data)
+
+ # Get display options
+ options = self.get_display_options(field_config)
+ preserve_whitespace = options.get("preserve_whitespace", True)
+ collapsible = options.get("collapsible", False)
+ max_height = options.get("max_height")
+
+ # Process the text
+ if self.allow_html:
+ # Sanitize HTML but allow safe tags
+ from potato.server_utils.html_sanitizer import sanitize_html
+ content = str(sanitize_html(text))
+ else:
+ # Escape all HTML for plain text
+ content = html.escape(text)
+ # Convert newlines to for display
+ if preserve_whitespace:
+ content = content.replace('\n', ' ')
+
+ # Build the content wrapper
+ wrapper_classes = ["text-display-content"]
+ wrapper_style = []
+
+ if preserve_whitespace and not self.allow_html:
+ wrapper_classes.append("preserve-whitespace")
+
+ if max_height:
+ wrapper_style.append(f"max-height: {max_height}px")
+ wrapper_style.append("overflow-y: auto")
+
+ # Check if this is a span target - add special wrapper for span annotations
+ is_span_target = field_config.get("span_target", False)
+ if is_span_target:
+ wrapper_classes.append("span-target-text")
+ # Add data attribute for original text (used by span manager)
+ field_key = field_config.get("key", "")
+ content = f'{content}
'
+
+ style_attr = f' style="{"; ".join(wrapper_style)}"' if wrapper_style else ""
+ class_attr = f'class="{" ".join(wrapper_classes)}"'
+
+ if collapsible:
+ return self._render_collapsible(content, class_attr, style_attr, field_config)
+
+ return f'{content}
'
+
+ def _render_collapsible(
+ self,
+ content: str,
+ class_attr: str,
+ style_attr: str,
+ field_config: Dict[str, Any]
+ ) -> str:
+ """
+ Render content in a collapsible container.
+
+ Args:
+ content: The HTML content
+ class_attr: CSS class attribute string
+ style_attr: CSS style attribute string
+ field_config: The field configuration
+
+ Returns:
+ HTML for collapsible content
+ """
+ field_key = field_config.get("key", "text")
+ collapse_id = f"collapse-{field_key}"
+ label = field_config.get("label", "")
+
+ # Build header row with label and button inline
+ label_html = f'{label} ' if label else ''
+
+ return f'''
+
+ '''
+
+ def has_inline_label(self, field_config: Dict[str, Any]) -> bool:
+ """
+ Check if this display handles its own label rendering.
+
+ Returns True for collapsible text so the registry doesn't add
+ a duplicate label.
+
+ Args:
+ field_config: The field configuration
+
+ Returns:
+ True if label is rendered inline by this display
+ """
+ options = self.get_display_options(field_config)
+ return options.get("collapsible", False)
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ """Get CSS classes for the container."""
+ classes = super().get_css_classes(field_config)
+ if field_config.get("span_target"):
+ classes.append("span-target-field")
+ if self.allow_html:
+ classes.append("html-content")
+ return classes
+
+ def get_data_attributes(self, field_config: Dict[str, Any], data: Any) -> Dict[str, str]:
+ """Get data attributes for the container."""
+ attrs = super().get_data_attributes(field_config, data)
+ if field_config.get("span_target"):
+ attrs["span-target"] = "true"
+ return attrs
diff --git a/potato/server_utils/displays/video_display.py b/potato/server_utils/displays/video_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a51c49ab15c5ddfa80579bc12cf8eb2ba4b8f8a
--- /dev/null
+++ b/potato/server_utils/displays/video_display.py
@@ -0,0 +1,148 @@
+"""
+Video Display Type
+
+Renders video content for display in the annotation interface.
+Supports video controls and links to video annotation schemas.
+"""
+
+import html
+from typing import Dict, Any, List
+from urllib.parse import urlparse
+
+from .base import BaseDisplay
+
+
+class VideoDisplay(BaseDisplay):
+ """
+ Display type for video content.
+
+ Displays videos with standard HTML5 video player controls.
+ Can be linked to video_annotation schemas via source_field.
+ """
+
+ name = "video"
+ required_fields = ["key"]
+ optional_fields = {
+ "max_width": None,
+ "max_height": None,
+ "controls": True,
+ "autoplay": False,
+ "loop": False,
+ "muted": False,
+ "poster": None,
+ }
+ description = "Video player display"
+ supports_span_target = False
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ """
+ Render a video as HTML.
+
+ Args:
+ field_config: The field configuration
+ data: The video URL or path
+
+ Returns:
+ HTML string for the video display
+ """
+ if not data:
+ return 'No video provided
'
+
+ # Get the video URL
+ video_url = str(data)
+
+ # Validate URL
+ if not self._is_valid_url(video_url):
+ return f'Invalid video URL: {html.escape(video_url)}
'
+
+ # Get display options
+ options = self.get_display_options(field_config)
+ max_width = options.get("max_width")
+ max_height = options.get("max_height")
+ controls = options.get("controls", True)
+ autoplay = options.get("autoplay", False)
+ loop = options.get("loop", False)
+ muted = options.get("muted", False)
+ poster = options.get("poster")
+
+ # Build style
+ style_parts = ["max-width: 100%"]
+ if max_width:
+ style_parts.append(f"width: {max_width}px" if isinstance(max_width, int) else f"width: {max_width}")
+ if max_height:
+ style_parts.append(f"max-height: {max_height}px" if isinstance(max_height, int) else f"max-height: {max_height}")
+
+ style_attr = f' style="{"; ".join(style_parts)}"'
+
+ # Build attributes
+ attrs = []
+ if controls:
+ attrs.append("controls")
+ if autoplay:
+ attrs.append("autoplay")
+ if loop:
+ attrs.append("loop")
+ if muted:
+ attrs.append("muted")
+ if poster:
+ attrs.append(f'poster="{html.escape(poster, quote=True)}"')
+
+ attrs_str = " ".join(attrs)
+
+ # Escape values
+ escaped_url = html.escape(video_url, quote=True)
+ field_key = html.escape(field_config.get("key", ""), quote=True)
+
+ # Determine video type from URL
+ video_type = self._get_video_type(video_url)
+ type_attr = f' type="{video_type}"' if video_type else ""
+
+ return f'''
+
+
+
+ Your browser does not support the video element.
+
+
+ '''
+
+ def _is_valid_url(self, url: str) -> bool:
+ """Check if URL is valid."""
+ if not url:
+ return False
+
+ if url.startswith('/') or url.startswith('./') or url.startswith('../'):
+ return True
+
+ try:
+ parsed = urlparse(url)
+ if parsed.scheme:
+ return parsed.scheme.lower() in ('http', 'https')
+ return True
+ except Exception:
+ return False
+
+ def _get_video_type(self, url: str) -> str:
+ """Get video MIME type from URL extension."""
+ url_lower = url.lower()
+ if url_lower.endswith('.mp4'):
+ return 'video/mp4'
+ elif url_lower.endswith('.webm'):
+ return 'video/webm'
+ elif url_lower.endswith('.ogg') or url_lower.endswith('.ogv'):
+ return 'video/ogg'
+ elif url_lower.endswith('.mov'):
+ return 'video/quicktime'
+ return ""
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ """Get CSS classes for the container."""
+ classes = super().get_css_classes(field_config)
+ return classes
+
+ def get_data_attributes(self, field_config: Dict[str, Any], data: Any) -> Dict[str, str]:
+ """Get data attributes for the container."""
+ attrs = super().get_data_attributes(field_config, data)
+ if data:
+ attrs["source-url"] = str(data)
+ return attrs
diff --git a/potato/server_utils/displays/web_agent_trace_display.py b/potato/server_utils/displays/web_agent_trace_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..edc0f45024aeeab76f48baa7b36701d93657476f
--- /dev/null
+++ b/potato/server_utils/displays/web_agent_trace_display.py
@@ -0,0 +1,447 @@
+"""
+Web Agent Trace Display Type
+
+Interactive step-by-step viewer for web agent browsing traces.
+Renders screenshots with SVG overlay visualizations (click markers,
+bounding boxes, mouse paths) alongside step details and a filmstrip navigator.
+
+Data format:
+{
+ "steps": [
+ {
+ "step_index": 0,
+ "screenshot_url": "screenshots/step_000.png",
+ "action_type": "click",
+ "element": {"tag": "input", "text": "Search", "bbox": [340, 45, 680, 75]},
+ "coordinates": {"x": 510, "y": 60},
+ "mouse_path": [[200, 300], [350, 200], [510, 60]],
+ "thought": "I need to search for blue wool sweaters",
+ "observation": "Search box is focused",
+ "timestamp": 1.2,
+ "viewport": {"width": 1280, "height": 720}
+ }
+ ],
+ "task_description": "Find and add a blue wool sweater under $50 to cart",
+ "site": "amazon.com"
+}
+
+Supported action_type values: click, type, scroll, hover, select, navigate, wait, done
+"""
+
+import html
+import json
+from typing import Dict, Any, List, Optional
+
+from .base import BaseDisplay
+
+
+# Action type badge colors
+ACTION_TYPE_COLORS = {
+ "click": {"bg": "#fff3e0", "border": "#FF9800", "badge": "rgba(255,152,0,0.2)"},
+ "type": {"bg": "#e8f4fd", "border": "#2196F3", "badge": "rgba(33,150,243,0.2)"},
+ "scroll": {"bg": "#e8f5e9", "border": "#4CAF50", "badge": "rgba(76,175,80,0.2)"},
+ "hover": {"bg": "#f3e5f5", "border": "#9C27B0", "badge": "rgba(156,39,176,0.2)"},
+ "select": {"bg": "#e0f7fa", "border": "#00BCD4", "badge": "rgba(0,188,212,0.2)"},
+ "navigate": {"bg": "#e8eaf6", "border": "#3F51B5", "badge": "rgba(63,81,181,0.2)"},
+ "wait": {"bg": "#f5f5f5", "border": "#9E9E9E", "badge": "rgba(158,158,158,0.2)"},
+ "done": {"bg": "#e8f5e9", "border": "#388E3C", "badge": "rgba(56,142,60,0.2)"},
+}
+
+DEFAULT_ACTION_COLOR = {"bg": "#f5f5f5", "border": "#9E9E9E", "badge": "rgba(158,158,158,0.2)"}
+
+
+class WebAgentTraceDisplay(BaseDisplay):
+ """
+ Display type for web agent browsing traces with interactive step navigation,
+ SVG overlay visualizations, and filmstrip thumbnails.
+ """
+
+ name = "web_agent_trace"
+ required_fields = ["key"]
+ optional_fields = {
+ "show_overlays": True,
+ "show_filmstrip": True,
+ "show_thought": True,
+ "show_observation": True,
+ "show_element_info": True,
+ "screenshot_max_width": 800,
+ "screenshot_max_height": 600,
+ "filmstrip_size": 80,
+ "auto_playback": False,
+ "playback_step_delay": 2.0,
+ }
+ description = "Web agent trace viewer with screenshots, SVG overlays, and step navigation"
+ supports_span_target = False
+
+ def render(self, field_config: Dict[str, Any], data: Any) -> str:
+ if not data:
+ return 'No trace data provided
'
+
+ options = self.get_display_options(field_config)
+ field_key = html.escape(field_config.get("key", ""), quote=True)
+
+ # Normalize data
+ steps = self._normalize_steps(data)
+ if not steps:
+ return 'No trace steps found
'
+
+ # Extract task info
+ task_desc = ""
+ site = ""
+ if isinstance(data, dict):
+ task_desc = data.get("task_description", "")
+ site = data.get("site", "")
+
+ max_w = options.get("screenshot_max_width", 800)
+ max_h = options.get("screenshot_max_height", 600)
+ filmstrip_size = options.get("filmstrip_size", 80)
+
+ # Serialize steps for JS
+ steps_json = html.escape(json.dumps(steps, ensure_ascii=False), quote=True)
+
+ css = self._build_css(max_w, max_h, filmstrip_size)
+
+ # Task header
+ task_html = ""
+ if task_desc:
+ escaped_task = html.escape(str(task_desc))
+ task_html = f'Task: {escaped_task}
'
+ if site:
+ escaped_site = html.escape(str(site))
+ task_html += f'Site: {escaped_site}
'
+
+ # Build first step display (JS will handle subsequent navigation)
+ first_step = steps[0]
+ screenshot_html = self._render_screenshot(first_step, max_w, max_h)
+ details_html = self._render_step_details(first_step, 0, len(steps), options)
+ filmstrip_html = self._render_filmstrip(steps, filmstrip_size) if options.get("show_filmstrip", True) else ""
+
+ # Per-step annotation container
+ per_step_html = '
'
+
+ # Playback data attributes
+ auto_playback = options.get("auto_playback", False)
+ playback_delay = options.get("playback_step_delay", 2.0)
+ playback_attrs = ""
+ if auto_playback:
+ playback_attrs = f' data-auto-playback="true" data-playback-step-delay="{playback_delay}"'
+
+ return f'''
+
+
+ {task_html}
+
+
+
+
+
+
+
+ « Prev
+ Step 1 of {len(steps)}
+ » Next
+
+
+
+ {details_html}
+ {per_step_html}
+
+
+ {filmstrip_html}
+
+ Clicks
+ Bounding Boxes
+ Mouse Path
+ Scroll
+
+
+ '''
+
+ def _normalize_steps(self, data: Any) -> List[Dict[str, Any]]:
+ """Normalize input data to a list of step dicts."""
+ if isinstance(data, dict):
+ steps = data.get("steps", [])
+ elif isinstance(data, list):
+ steps = data
+ else:
+ return []
+
+ normalized = []
+ for i, step in enumerate(steps):
+ if not isinstance(step, dict):
+ continue
+ normalized.append({
+ "step_index": step.get("step_index", i),
+ "screenshot_url": step.get("screenshot_url", step.get("screenshot", "")),
+ "action_type": step.get("action_type", "unknown"),
+ "element": step.get("element", {}),
+ "coordinates": step.get("coordinates", {}),
+ "mouse_path": step.get("mouse_path", []),
+ "thought": step.get("thought", ""),
+ "observation": step.get("observation", ""),
+ "timestamp": step.get("timestamp", ""),
+ "viewport": step.get("viewport", {"width": 1280, "height": 720}),
+ "typed_text": step.get("typed_text", step.get("value", "")),
+ "scroll_direction": step.get("scroll_direction", step.get("direction", "")),
+ })
+ return normalized
+
+ def _render_screenshot(self, step: Dict, max_w: int, max_h: int) -> str:
+ """Render the screenshot image element."""
+ url = step.get("screenshot_url", "")
+ if not url:
+ return 'No screenshot available
'
+ escaped_url = html.escape(str(url), quote=True)
+ return f' '
+
+ def _render_step_details(self, step: Dict, index: int, total: int,
+ options: Dict) -> str:
+ """Render the step details panel content."""
+ action_type = step.get("action_type", "unknown")
+ safe_type = html.escape(str(action_type), quote=True)
+ colors = ACTION_TYPE_COLORS.get(action_type, DEFAULT_ACTION_COLOR)
+
+ parts = []
+
+ # Action badge
+ parts.append(
+ f''
+ f'{safe_type.upper()}
'
+ )
+
+ # Timestamp
+ ts = step.get("timestamp", "")
+ if ts:
+ parts.append(f't={html.escape(str(ts))}s
')
+
+ # Thought
+ thought = step.get("thought", "")
+ if thought and options.get("show_thought", True):
+ parts.append(
+ f''
+ f'Thought: {html.escape(str(thought))}'
+ f'
'
+ )
+
+ # Element info
+ element = step.get("element", {})
+ if element and options.get("show_element_info", True):
+ elem_parts = []
+ if isinstance(element, dict):
+ for k in ("tag", "text", "id", "class"):
+ if k in element:
+ elem_parts.append(f'{k}="{html.escape(str(element[k]))}"')
+ if elem_parts:
+ parts.append(
+ f''
+ f'Element: {" ".join(elem_parts)}'
+ f'
'
+ )
+
+ # Coordinates
+ coords = step.get("coordinates", {})
+ if coords and isinstance(coords, dict):
+ x = coords.get("x", "")
+ y = coords.get("y", "")
+ if x or y:
+ parts.append(
+ f''
+ f'Coords: ({html.escape(str(x))}, {html.escape(str(y))})'
+ f'
'
+ )
+
+ # Typed text (for type actions)
+ typed = step.get("typed_text", "")
+ if typed:
+ parts.append(
+ f''
+ f'Typed: "{html.escape(str(typed))}"'
+ f'
'
+ )
+
+ # Observation
+ obs = step.get("observation", "")
+ if obs and options.get("show_observation", True):
+ parts.append(
+ f''
+ f'Observation: {html.escape(str(obs))}'
+ f'
'
+ )
+
+ return f'{"".join(parts)}
'
+
+ def _render_filmstrip(self, steps: List[Dict], thumb_size: int) -> str:
+ """Render the filmstrip thumbnail navigation bar."""
+ thumbs = []
+ for i, step in enumerate(steps):
+ url = step.get("screenshot_url", "")
+ active = "filmstrip-active" if i == 0 else ""
+ if url:
+ escaped_url = html.escape(str(url), quote=True)
+ thumbs.append(
+ f''
+ f'
'
+ f'
{i + 1} '
+ f'
'
+ )
+ else:
+ thumbs.append(
+ f''
+ f'
?
'
+ f'
{i + 1} '
+ f'
'
+ )
+ return f'{"".join(thumbs)}
'
+
+ def _build_css(self, max_w: int, max_h: int, filmstrip_size: int) -> str:
+ """Build CSS for the web agent viewer."""
+ return f'''
+ .web-agent-viewer {{ font-family: inherit; }}
+ .web-agent-task, .web-agent-site {{
+ padding: 6px 12px; margin-bottom: 8px;
+ background: #f8f9fa; border-radius: 4px; font-size: 0.95em;
+ }}
+ .web-agent-main {{
+ display: flex; gap: 16px; margin-bottom: 12px;
+ }}
+ .screenshot-panel {{
+ flex: 0 0 auto; max-width: {max_w}px;
+ }}
+ .screenshot-container {{
+ position: relative; display: inline-block;
+ border: 1px solid #ddd; border-radius: 6px; overflow: hidden;
+ background: #1a1a1a;
+ }}
+ .step-screenshot {{
+ display: block; max-width: {max_w}px; max-height: {max_h}px;
+ width: auto; height: auto;
+ }}
+ .overlay-layer {{
+ position: absolute; top: 0; left: 0;
+ width: 100%; height: 100%;
+ pointer-events: none;
+ }}
+ .screenshot-placeholder {{
+ width: {max_w}px; height: 300px;
+ display: flex; align-items: center; justify-content: center;
+ background: #f0f0f0; color: #999; font-size: 0.9em;
+ }}
+ .step-nav {{
+ display: flex; align-items: center; justify-content: center;
+ gap: 12px; padding: 8px 0;
+ }}
+ .step-nav button {{
+ padding: 4px 12px; border: 1px solid #ccc; border-radius: 4px;
+ background: #fff; cursor: pointer; font-size: 0.85em;
+ }}
+ .step-nav button:disabled {{
+ opacity: 0.4; cursor: default;
+ }}
+ .step-nav button:not(:disabled):hover {{
+ background: #e8f4fd; border-color: #2196F3;
+ }}
+ .step-counter {{
+ font-size: 0.9em; font-weight: 600; color: #555;
+ }}
+ .step-details-panel {{
+ flex: 1; min-width: 250px; max-width: 400px;
+ border: 1px solid #e0e0e0; border-radius: 6px;
+ padding: 12px; background: #fafafa; overflow-y: auto;
+ max-height: {max_h + 50}px;
+ }}
+ .step-details-content {{
+ display: flex; flex-direction: column; gap: 8px;
+ }}
+ .action-badge {{
+ display: inline-block; padding: 4px 12px;
+ border-radius: 4px; font-weight: 700;
+ font-size: 0.8em; letter-spacing: 0.5px;
+ }}
+ .step-timestamp {{ color: #888; font-size: 0.8em; }}
+ .step-thought {{
+ padding: 8px; background: #e8f4fd;
+ border-left: 3px solid #2196F3; border-radius: 4px;
+ font-size: 0.9em;
+ }}
+ .step-element {{
+ font-size: 0.85em; color: #555;
+ }}
+ .step-element code {{
+ background: #f0f0f0; padding: 2px 4px; border-radius: 2px;
+ font-size: 0.9em;
+ }}
+ .step-coords {{ font-size: 0.85em; color: #666; }}
+ .step-typed {{
+ padding: 6px 8px; background: #e8f4fd;
+ border-radius: 4px; font-size: 0.9em;
+ }}
+ .step-observation {{
+ padding: 8px; background: #e8f5e9;
+ border-left: 3px solid #4CAF50; border-radius: 4px;
+ font-size: 0.9em;
+ }}
+ .filmstrip {{
+ display: flex; gap: 4px; overflow-x: auto;
+ padding: 8px 4px; background: #f5f5f5;
+ border-radius: 6px; margin-top: 4px;
+ }}
+ .filmstrip-thumb {{
+ flex: 0 0 auto; width: {filmstrip_size}px;
+ cursor: pointer; border: 2px solid transparent;
+ border-radius: 4px; overflow: hidden;
+ text-align: center; background: #fff;
+ transition: border-color 0.2s;
+ }}
+ .filmstrip-thumb:hover {{ border-color: #90CAF9; }}
+ .filmstrip-thumb.filmstrip-active {{ border-color: #2196F3; }}
+ .filmstrip-thumb img {{
+ width: 100%; height: {int(filmstrip_size * 0.65)}px;
+ object-fit: cover; display: block;
+ }}
+ .filmstrip-placeholder {{
+ width: 100%; height: {int(filmstrip_size * 0.65)}px;
+ display: flex; align-items: center; justify-content: center;
+ background: #eee; color: #999; font-size: 0.8em;
+ }}
+ .filmstrip-label {{
+ font-size: 0.7em; color: #666; padding: 2px 0;
+ }}
+ .overlay-controls {{
+ display: flex; gap: 12px; padding: 6px 0;
+ font-size: 0.8em; color: #555;
+ }}
+ .overlay-controls label {{
+ display: flex; align-items: center; gap: 4px; cursor: pointer;
+ }}
+ .web-agent-per-step-annotations {{
+ margin-top: 12px; padding-top: 12px;
+ border-top: 1px solid #e0e0e0;
+ }}
+
+ /* SVG overlay styles */
+ .overlay-click-marker {{ }}
+ .overlay-bbox {{ }}
+ .overlay-mouse-path {{ }}
+ .overlay-scroll {{ }}
+
+ @keyframes pulse-marker {{
+ 0%, 100% {{ r: 8; opacity: 1; }}
+ 50% {{ r: 12; opacity: 0.7; }}
+ }}
+
+ @media (max-width: 768px) {{
+ .web-agent-main {{ flex-direction: column; }}
+ .step-details-panel {{ max-width: none; max-height: none; }}
+ .screenshot-panel {{ max-width: 100%; }}
+ .step-screenshot {{ max-width: 100%; }}
+ }}
+ '''
+
+ def get_css_classes(self, field_config: Dict[str, Any]) -> List[str]:
+ classes = super().get_css_classes(field_config)
+ return classes
+
+ def get_data_attributes(self, field_config: Dict[str, Any], data: Any) -> Dict[str, str]:
+ attrs = super().get_data_attributes(field_config, data)
+ return attrs
diff --git a/potato/server_utils/front_end.py b/potato/server_utils/front_end.py
new file mode 100644
index 0000000000000000000000000000000000000000..8a63906812d4507954c5a6bb8d2132d2af07cc42
--- /dev/null
+++ b/potato/server_utils/front_end.py
@@ -0,0 +1,736 @@
+"""
+Handle all front-end related functionalities.
+"""
+
+import base64
+import os
+import logging
+import json
+import re
+import hashlib
+from collections import OrderedDict
+
+#add local module
+from pathlib import Path
+import sys
+path_root = Path(__file__).parents[2]
+sys.path.append(str(path_root))
+
+from potato.server_utils.config_module import config
+from potato.server_utils.schemas.registry import schema_registry
+from potato.server_utils.schemas.keybinding_allocator import allocate_keybindings
+
+logger = logging.getLogger(__name__)
+
+
+# TODO: Move this to config.yaml files
+# Items which will be displayed in the popup statistics sidebar
+STATS_KEYS = {
+ "Annotated instances": "Annotated instances",
+ "Total working time": "Total working time",
+ "Average time on each instance": "Average time on each instance",
+ "Agreement": "Agreement",
+}
+
+# Default name for the generated annotation layout file
+DEFAULT_ANNOTATION_LAYOUT_SUBDIR = "layouts"
+DEFAULT_ANNOTATION_LAYOUT_FILENAME = "task_layout.html"
+
+
+SUPPORTED_HEADER_LOGO_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.gif', '.svg', '.ico', '.webp'}
+EXTENSION_TO_MIME = {
+ '.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg',
+ '.gif': 'image/gif', '.svg': 'image/svg+xml', '.ico': 'image/x-icon',
+ '.webp': 'image/webp',
+}
+
+
+def resolve_header_logo_src(config: dict) -> str:
+ """
+ Resolve the ``header_logo`` config value into a src URL for an `` `` tag.
+
+ - If not configured, returns ``""``.
+ - If the value is an HTTP(S) URL, returns it directly.
+ - Otherwise, reads the local file, base64-encodes it, and returns a data URL.
+
+ Returns:
+ A URL string suitable for `` ``, or ``""`` if not configured.
+ """
+ logo_path = config.get("header_logo")
+ if not logo_path:
+ return ""
+
+ # Pass through external URLs
+ if logo_path.startswith(("http://", "https://")):
+ return logo_path
+
+ try:
+ resolved = resolve_project_asset_path(config, logo_path)
+ except FileNotFoundError:
+ logger.warning("header_logo file not found: %s", logo_path)
+ return ""
+
+ ext = os.path.splitext(resolved)[1].lower()
+ if ext not in SUPPORTED_HEADER_LOGO_EXTENSIONS:
+ logger.warning("header_logo has unsupported extension '%s' (supported: %s)",
+ ext, ', '.join(sorted(SUPPORTED_HEADER_LOGO_EXTENSIONS)))
+ return ""
+
+ mime = EXTENSION_TO_MIME[ext]
+ with open(resolved, "rb") as f:
+ encoded = base64.b64encode(f.read()).decode("ascii")
+
+ return f"data:{mime};base64,{encoded}"
+
+
+def resolve_project_asset_path(config: dict, relative_path: str) -> str:
+ """
+ Resolve a project-relative asset path using the config file directory as base.
+
+ Args:
+ config: The configuration dict (must contain ``__config_file__``)
+ relative_path: The path as specified in the config (absolute or relative)
+
+ Returns:
+ Absolute path to the resolved file
+
+ Raises:
+ FileNotFoundError: If the file does not exist at the resolved path
+ """
+ if os.path.isabs(relative_path) and os.path.exists(relative_path):
+ return relative_path
+
+ if os.path.exists(relative_path):
+ return os.path.abspath(relative_path)
+
+ # Resolve relative to the config file's directory
+ config_file = config.get("__config_file__", "")
+ if config_file:
+ real_path = os.path.realpath(config_file)
+ dir_path = os.path.dirname(real_path)
+ abs_path = os.path.join(dir_path, relative_path)
+ if os.path.exists(abs_path):
+ return abs_path
+
+ raise FileNotFoundError(f"Project asset file not found: {relative_path}")
+
+
+def load_project_base_css_html(config: dict) -> str:
+ """
+ Load the project-level ``base_css`` file and return it wrapped in a ``'
+
+
+def _stringify_dict_keys(value):
+ """
+ Recursively convert all dict keys to strings so the structure can be
+ JSON-serialized with ``sort_keys=True``.
+
+ YAML 1.1 parses unquoted ``yes``/``no``/``true``/``false`` as booleans, so a
+ config dict can end up with mixed key types (e.g. ``str`` and ``bool``).
+ ``json.dumps(..., sort_keys=True)`` then raises ``TypeError`` because it
+ cannot order keys of different types. Coercing keys to strings up front
+ makes hashing robust regardless of how the config was authored or merged.
+ """
+ if isinstance(value, dict):
+ return {str(k): _stringify_dict_keys(v) for k, v in value.items()}
+ if isinstance(value, list):
+ return [_stringify_dict_keys(v) for v in value]
+ if isinstance(value, tuple):
+ return tuple(_stringify_dict_keys(v) for v in value)
+ return value
+
+
+def compute_config_md5(config):
+ """
+ Compute MD5 hash of the config dict for template invalidation.
+ """
+ # Remove unserializable fields if needed
+ config_copy = {k: v for k, v in config.items() if k not in ['__config_file__', 'site_file']}
+ normalized = _stringify_dict_keys(config_copy)
+ config_str = json.dumps(normalized, sort_keys=True, default=str)
+ return hashlib.md5(config_str.encode('utf-8')).hexdigest()
+
+
+def generate_annotation_layout_file(config: dict, annotation_schemes: list[dict], layout_name: str = None) -> str:
+ """
+ Generate a dedicated annotation layout file in the task directory under layouts/task_layout.html.
+ If layout_name is provided, uses task_layout_{layout_name}.html instead.
+ """
+ task_dir = config.get("task_dir")
+ if not task_dir:
+ raise ValueError("task_dir is required in config to generate annotation layout file")
+
+ # Ensure task directory and layouts subdirectory exist
+ layout_dir = os.path.join(task_dir, DEFAULT_ANNOTATION_LAYOUT_SUBDIR)
+ if not os.path.exists(layout_dir):
+ os.makedirs(layout_dir)
+
+ # Generate the layout file path
+ filename = f"task_layout_{layout_name}.html" if layout_name else DEFAULT_ANNOTATION_LAYOUT_FILENAME
+ layout_file_path = os.path.join(layout_dir, filename)
+
+ # Generate the HTML layout content
+ schema_layouts = ""
+ all_keybindings = []
+
+ for annotation_scheme in annotation_schemes:
+ schema_layout, keybindings = generate_schematic(annotation_scheme)
+ schema_layouts += schema_layout + "\n"
+ all_keybindings.extend(keybindings)
+
+ # Compute combined hash (config + schema content) for cache invalidation
+ config_hash = compute_config_md5(config)
+ schema_content_hash = hashlib.md5()
+ schema_content_hash.update(schema_layouts.encode('utf-8'))
+ combined_hash = f"{config_hash}_{schema_content_hash.hexdigest()}"
+
+ # Create the layout HTML content with combined hash at the top
+ layout_content = f"""
+
+
+
+
+
+
+{schema_layouts}
+
+"""
+
+ # Write the layout file
+ with open(layout_file_path, "wt", encoding="utf-8") as outf:
+ outf.write(layout_content)
+
+ logger.info(f"Generated annotation layout file: {layout_file_path}")
+ return layout_file_path
+
+
+def get_or_generate_annotation_layout(config: dict, annotation_schemes: list[dict], layout_name: str = None) -> str:
+ """
+ Get the annotation layout file path, generating it if it doesn't exist or if the config hash has changed.
+ If layout_name is provided, uses task_layout_{layout_name}.html instead.
+ """
+ task_dir = config.get("task_dir")
+ if not task_dir:
+ raise ValueError("task_dir is required in config")
+ layout_dir = os.path.join(task_dir, DEFAULT_ANNOTATION_LAYOUT_SUBDIR)
+ filename = f"task_layout_{layout_name}.html" if layout_name else DEFAULT_ANNOTATION_LAYOUT_FILENAME
+ layout_file_path = os.path.join(layout_dir, filename)
+
+ config_hash = compute_config_md5(config)
+
+ # Also hash the actual generated schema content to detect code changes
+ # (e.g., if bws.py changes separators, the config hash won't change but the output will)
+ # NOTE: must match the concatenation format used in generate_annotation_layout_file
+ # (each layout followed by "\n") so hashes are consistent
+ schema_content_hash = hashlib.md5()
+ schema_layouts = ""
+ for annotation_scheme in annotation_schemes:
+ layout_html, _ = generate_schematic(annotation_scheme)
+ schema_layouts += layout_html + "\n"
+ schema_content_hash.update(schema_layouts.encode('utf-8'))
+ combined_hash = f"{config_hash}_{schema_content_hash.hexdigest()}"
+
+ # Check if the layout file already exists and if the hash matches
+ if os.path.exists(layout_file_path):
+ with open(layout_file_path, "rt", encoding="utf-8") as f:
+ for _ in range(2): # Only need to check the first two lines
+ line = f.readline()
+ if line.startswith("', '').strip()
+ if file_hash == combined_hash:
+ logger.info(f"Using existing annotation layout file: {layout_file_path} (hash match)")
+ return layout_file_path
+ else:
+ logger.info(f"Hash mismatch (config or schema code changed), regenerating: {layout_file_path}")
+ break
+
+ # Generate the layout file if it doesn't exist or hash mismatches
+ logger.info(f"Annotation layout file not found or hash mismatch, generating: {layout_file_path}")
+ return generate_annotation_layout_file(config, annotation_schemes, layout_name=layout_name)
+
+
+def generate_schematic(annotation_scheme):
+ """
+ Based on the task's yaml configuration, generate the full HTML site needed
+ to annotate the tasks's data.
+
+ Uses the schema registry to look up the generator function for the
+ annotation type.
+ """
+ # Ensure annotation_id is set before any schema generator runs.
+ # This is the single bottleneck before all generators, so it serves as a
+ # safety net for any caller that doesn't pre-assign annotation_id.
+ if "annotation_id" not in annotation_scheme:
+ annotation_scheme["annotation_id"] = 0
+
+ # Figure out which kind of tasks we're doing and build the input frame
+ annotation_type = annotation_scheme["annotation_type"]
+
+ # Use the schema registry to get the generator
+ return schema_registry.generate(annotation_scheme)
+
+
+def generate_keybindings_sidebar(config, keybindings, horizontal=False):
+ """
+ Generate an HTML layout for the end-user of the keybindings for the current
+ task. The layout is intended to be displayed in a side bar or on the annotation page if fixed_keybinding_layout.html is used as the layout
+ """
+ if config.get("horizontal_key_bindings"):
+ horizontal = True
+
+ if not keybindings:
+ return ""
+
+ if horizontal:
+ keybindings = [[it[0], it[1].split(":")[-1]] for it in keybindings]
+ lines = list(zip(*keybindings))
+ layout = ''
+ for line in lines:
+ layout += (
+ ""
+ + "".join([" %s " % it for it in line])
+ + " "
+ )
+ layout += "
"
+
+ else:
+ layout = "Key Description "
+ for key, desc in keybindings:
+ layout += '%s %s ' % (key, desc)
+ layout += "
"
+
+ return layout
+
+
+def generate_statistics_sidebar(statistics):
+ """
+ Generate an HTML layout for the end-user of the statistics for the current
+ task. The layout is intended to be displayed in a side bar
+ """
+ layout = " "
+ for key in statistics:
+ desc = "{{statistics_nav['%s']}}" % statistics[key]
+ layout += '%s %s ' % (key, desc)
+ layout += "
"
+ return layout
+
+
+def generate_annotation_html_template(config: dict) -> str:
+ """
+ Generates the full HTML file in site/ for annotating this tasks data,
+ combining the various templates with the annotation specification in
+ the yaml file and returns the path to the HTML template for this
+ annotation task.
+ """
+ logger.info("Generating anntotation site at %s" % config["site_dir"])
+
+ #
+ # Stage 1: Construct the core HTML file devoid the annotation-specific content
+ #
+
+ # Use hardcoded template paths - no longer configurable
+ cur_program_dir = os.path.dirname(os.path.abspath(__file__))
+ html_template_file = os.path.join(cur_program_dir, '..', 'templates', 'base_template_v2.html')
+ header_file = os.path.join(cur_program_dir, '..', 'templates', 'header.html')
+
+ logger.debug(f"Reading html annotation template: {html_template_file}")
+
+ if not os.path.exists(html_template_file):
+ raise FileNotFoundError("html_template_file not found: %s" % html_template_file)
+
+ with open(html_template_file, "rt", encoding="utf-8") as file_p:
+ html_template = "".join(file_p.readlines())
+
+ # Load the header content we'll stuff in the template, which has scripts
+ # and assets we'll need
+ logger.debug("Reading html header %s" % header_file)
+
+ if not os.path.exists(header_file):
+ raise FileNotFoundError("header_file not found: %s" % header_file)
+
+ with open(header_file, "rt", encoding="utf-8") as file_p:
+ header = "".join(file_p.readlines())
+
+ html_template = html_template.replace("{{ HEADER }}", header)
+
+ if config.get("hide_navbar"):
+ html_template = html_template.replace(
+ '', '
'
+ )
+
+ # Codebook bridge: schemes with `codebook: true` get their labels
+ # from the project's mutable codebook (seeded from YAML on first
+ # run). Single chokepoint before any scheme HTML is generated, in
+ # both the CLI and WSGI-factory init paths.
+ try:
+ from potato.codebook.schema_bridge import apply_codebook_to_schemes
+ apply_codebook_to_schemes(config)
+ except Exception as e:
+ logger.warning(f"Codebook schema bridge skipped: {e}")
+
+ # Grab the annotation schemes
+ annotation_schemes = config["annotation_schemes"]
+ logger.debug("Saw %d annotation scheme(s)" % len(annotation_schemes))
+
+ # insert annotation id to each of the schemes
+ for idx, annotation_scheme in enumerate(annotation_schemes):
+ annotation_scheme["annotation_id"] = idx
+
+ # Pre-allocate non-conflicting keybindings across all schemas
+ allocation = allocate_keybindings(annotation_schemes)
+ for annotation_scheme in annotation_schemes:
+ name = annotation_scheme.get("name", "")
+ if name in allocation:
+ annotation_scheme["_allocated_keys"] = allocation[name]
+
+ # Keep track of all the keybindings we have
+ all_keybindings = [("←", "Move backward"), ("→", "Move forward")]
+
+ # Check if we're using the new API-based template that generates forms dynamically
+ is_api_template = "base_template_v2.html" in html_template_file
+
+ # Handle annotation layout generation
+ # Check if user provided a custom task_layout file
+ task_layout_file = config.get("task_layout")
+
+ if task_layout_file:
+ # User provided a custom task layout file
+ logger.info(f"Using custom task layout file: {task_layout_file}")
+
+ # Resolve the path relative to the config file
+ task_layout_file = resolve_project_asset_path(config, task_layout_file)
+
+ # Read the custom task layout
+ with open(task_layout_file, "rt", encoding="utf-8") as f:
+ task_html_layout = "".join(f.readlines())
+
+ # Extract keybindings from the annotation schemes for the sidebar
+ for annotation_scheme in annotation_schemes:
+ _, keybindings = generate_schematic(annotation_scheme)
+ all_keybindings.extend(keybindings)
+
+ else:
+ # Use the dedicated annotation layout file system (auto-generated)
+ try:
+ layout_file_path = get_or_generate_annotation_layout(config, annotation_schemes)
+ # Read the generated layout file
+ with open(layout_file_path, "rt", encoding="utf-8") as f:
+ task_html_layout = "".join(f.readlines())
+
+ # Extract keybindings from the annotation schemes for the sidebar
+ for annotation_scheme in annotation_schemes:
+ _, keybindings = generate_schematic(annotation_scheme)
+ all_keybindings.extend(keybindings)
+
+ except Exception as e:
+ logger.warning(f"Failed to use dedicated layout file: {e}. Falling back to inline generation.")
+
+ # Fallback to inline generation
+ if is_api_template:
+ # For the new API-based template, generate server-side forms but use API endpoints
+ # The frontend JavaScript will handle form interactions via API calls
+ logger.info("Using API-based template - generating server-side forms with API integration")
+
+ # Generate the forms using the existing schematic generation
+ schema_layouts = ""
+ for annotation_scheme in annotation_schemes:
+ schema_layout, keybindings = generate_schematic(annotation_scheme)
+ schema_layouts += schema_layout + "\n"
+ all_keybindings.extend(keybindings)
+
+ task_html_layout = schema_layouts
+ else:
+ # Generate inline layout
+ schema_layouts = ""
+ for annotation_scheme in annotation_schemes:
+ schema_layout, keybindings = generate_schematic(annotation_scheme)
+ schema_layouts += schema_layout + "\n"
+ all_keybindings.extend(keybindings)
+
+ task_html_layout = f'
{schema_layouts}
'
+
+ # Add in a codebook link if the admin specified one
+ codebook_html = ""
+ if len(config.get("annotation_codebook_url", "")) > 0:
+ annotation_codebook = config["annotation_codebook_url"]
+ codebook_html = '
Annotation Codebook '
+ codebook_html = codebook_html.replace("{{annotation_codebook_url}}", annotation_codebook)
+
+ #
+ # Step 3, drop in the annotation layout and insert the rest of the task-specific variables
+ #
+
+ # Swap in the task's layout
+ html_template = html_template.replace("{{ TASK_LAYOUT }}", task_html_layout)
+ html_template = html_template.replace("{{annotation_codebook}}", codebook_html)
+ html_template = html_template.replace(
+ "{{annotation_task_name}}", config["annotation_task_name"]
+ )
+
+ # For API-based templates, replace debug placeholder
+ if is_api_template:
+ html_template = html_template.replace("{{ debug | tojson | safe }}", str(config.get("debug", False)).lower())
+
+ keybindings_desc = generate_keybindings_sidebar(config, all_keybindings)
+ html_template = html_template.replace("{{keybindings}}", keybindings_desc)
+
+ statistics_layout = generate_statistics_sidebar(STATS_KEYS)
+ html_template = html_template.replace("{{statistics_nav}}", statistics_layout)
+
+ # Jiaxin: change the basename from the template name to the project name +
+ # template name, to allow multiple annotation tasks using the same template
+ site_name = (
+ "-".join(config["annotation_task_name"].split(" "))
+ + "-"
+ + os.path.basename(html_template_file)
+ )
+
+ # Create generated subdirectory within the templates directory
+ generated_dir = os.path.join(config["site_dir"], "generated")
+ if not os.path.exists(generated_dir):
+ os.makedirs(generated_dir)
+ logger.info(f"Created generated templates directory: {generated_dir}")
+
+ output_html_fname = os.path.join(generated_dir, site_name)
+ logger.debug(f"Output HTML filename: {output_html_fname}")
+
+ # Cache this path as a shortcut to figure out which page to render
+ config["site_file"] = site_name
+
+ # Compute config hash and add it to the template
+ config_hash = compute_config_md5(config)
+ html_template_with_hash = f"\n{html_template}"
+
+ # Write the file
+ with open(output_html_fname, "wt", encoding="utf-8") as outf:
+ outf.write(html_template_with_hash)
+
+ logger.debug("writing annotation html to %s" % output_html_fname)
+
+ return site_name
+
+def get_html(fname: str, config: dict):
+ """
+ Returns the content of an HTML file, looking for alternative locations relative
+ to the config file if the path is relative.
+ """
+ if not os.path.exists(fname):
+
+ real_path = os.path.realpath(config["__config_file__"])
+ dir_path = os.path.dirname(real_path)
+ abs_html_template_file = dir_path + "/" + fname
+
+ if not os.path.exists(abs_html_template_file):
+ raise FileNotFoundError("html file not found: %s" % fname)
+ else:
+ fname = abs_html_template_file
+
+ with open(fname, "rt", encoding="utf-8") as f:
+ html = "".join(f.readlines())
+ return html
+
+def generate_core_task_html(config: dict,
+ annotation_schemas: list[dict]) -> str:
+ """
+ Generates the HTML layout for the core annotation task for
+ all the annotation-specific content and returns the HTML layout.
+ """
+ schema_layouts = ""
+ task_html_layout = ""
+ for annotation_scheme in annotation_schemas:
+
+ schema_layout, keybindings = generate_schematic(annotation_scheme)
+ schema_layouts += schema_layout + "
" + "\n"
+
+ cur_task_html_layout = task_html_layout.replace(
+ "{{annotation_schematic}}", schema_layouts
+ )
+
+ # Swap in the task's layout
+ return cur_task_html_layout
+
+
+def generate_html_from_schematic(annotation_schemas: list[dict],
+ allow_jumping_to_id: bool,
+ hide_navbar: bool,
+ phase_name: str,
+ config: dict,
+ task_layout_file: str = None):
+ """
+ Generates the full HTML file in site/ for annotating this tasks data,
+ combining the various templates with the annotation specification in
+ the yaml file.
+ """
+ #
+ # Stage 1: Construct the core HTML file devoid the annotation-specific content
+ #
+
+ # Use hardcoded template paths - no longer configurable
+ cur_program_dir = os.path.dirname(os.path.abspath(__file__))
+ html_template_filename = os.path.join(cur_program_dir, '..', 'templates', 'base_template_v2.html')
+ html_header_filename = os.path.join(cur_program_dir, '..', 'templates', 'header.html')
+
+ # Load the core template that has all the UI controls and non-task layout.
+ logger.debug("Reading html annotation template %s" % html_template_filename)
+ html_template = get_html(html_template_filename, config)
+
+ # Load the header content we'll stuff in the template, which has scripts and assets we'll need
+ logger.debug("Reading html header %s" % html_header_filename)
+ header = get_html(html_header_filename, config)
+
+ # Once we have the base template constructed, load the user's custom layout for their task
+ html_template = html_template.replace("{{ HEADER }}", header)
+
+ if allow_jumping_to_id:
+ html_template = html_template.replace(
+ '
', '
'
+ )
+ html_template = html_template.replace(
+ '
',
+ '
',
+ )
+
+ if hide_navbar:
+ html_template = html_template.replace(
+ '
', '
'
+ )
+
+ # Assign annotation_id to each scheme if not already set.
+ # The main annotation path (generate_annotation_html_template) does this for
+ # config["annotation_schemes"], but phase schemas loaded from JSON files
+ # (consent, prestudy, etc.) arrive here without annotation_id set.
+ for idx, annotation_scheme in enumerate(annotation_schemas):
+ if "annotation_id" not in annotation_scheme:
+ annotation_scheme["annotation_id"] = idx
+
+ # Pre-allocate non-conflicting keybindings across all schemas
+ allocation = allocate_keybindings(annotation_schemas)
+ for annotation_scheme in annotation_schemas:
+ name = annotation_scheme.get("name", "")
+ if name in allocation:
+ annotation_scheme["_allocated_keys"] = allocation[name]
+
+ # Handle annotation layout generation for surveyflow phases.
+ # Only fall back to the global task_layout for the main annotation page
+ # (phase_name is None). Phase pages should not inherit the global
+ # annotation layout since it may expect annotation-only context.
+ if not task_layout_file and not phase_name:
+ task_layout_file = config.get("task_layout")
+
+ if task_layout_file:
+ # User provided a custom task layout file
+ logger.info(f"Using custom task layout file: {task_layout_file}")
+
+ # Resolve the path relative to the config file
+ task_layout_file = resolve_project_asset_path(config, task_layout_file)
+
+ # Read the custom task layout
+ with open(task_layout_file, "rt", encoding="utf-8") as f:
+ task_html_layout = "".join(f.readlines())
+
+ else:
+ # Use the dedicated annotation layout file system (auto-generated)
+ try:
+ layout_file_path = get_or_generate_annotation_layout(config, annotation_schemas, layout_name=phase_name)
+
+ # Read the generated layout file
+ with open(layout_file_path, "rt", encoding="utf-8") as f:
+ task_html_layout = "".join(f.readlines())
+
+ except Exception as e:
+ logger.warning(f"Failed to use dedicated layout file: {e}. Falling back to inline generation.")
+
+ # Fallback to inline generation
+ # Generate inline layout
+ schema_layouts = ""
+ for annotation_scheme in annotation_schemas:
+ schema_layout, keybindings = generate_schematic(annotation_scheme)
+ schema_layouts += schema_layout + "\n"
+
+ task_html_layout = f'
{schema_layouts}
'
+
+ cur_html_template = html_template.replace("{{ TASK_LAYOUT }}", task_html_layout)
+
+ # Add in a codebook link if the admin specified one
+ codebook_html = ""
+ if len(config.get("annotation_codebook_url", "")) > 0:
+ annotation_codebook = config["annotation_codebook_url"]
+ codebook_html = '
Annotation Codebook '
+ codebook_html = codebook_html.replace("{{annotation_codebook_url}}", annotation_codebook)
+
+ html_template = html_template.replace("{{annotation_codebook}}", codebook_html)
+
+ html_template = html_template.replace(
+ "{{annotation_task_name}}", config["annotation_task_name"]
+ )
+
+ _ = generate_statistics_sidebar(STATS_KEYS)
+ html_template = html_template.replace("{{statistics_nav}}", " ")
+
+ #
+ # Step 3, Fill in the annotation-specific pieces in the layout and save the page
+ #
+
+ logger.debug("Saw %d annotation scheme(s)" % len(annotation_schemas))
+
+ # Keep track of all the keybindings we have
+ all_keybindings = [("←", "Move backward"), ("→", "Move forward")]
+
+ # Do not display keybindings for the first and last page
+ if False:
+ if i == 0:
+ keybindings_desc = generate_keybindings_sidebar(config, all_keybindings[1:])
+ cur_html_template = cur_html_template.replace(
+ '
Move backward ',
+ '
Move backward ',
+ )
+ elif i == len(annotation_schemas) - 1 or re.search("prestudy_fail", page):
+ keybindings_desc = generate_keybindings_sidebar(config, all_keybindings[:-1])
+ cur_html_template = cur_html_template.replace(
+ '
Move forward ',
+ '
Move forward ',
+ )
+ else:
+ keybindings_desc = generate_keybindings_sidebar(config, all_keybindings)
+
+ cur_html_template = cur_html_template.replace("{{keybindings}}", keybindings_desc)
+
+ # Cache the html as a template for use in flask server
+ site_name = (
+ "_".join(config["annotation_task_name"].split(" "))
+ + "-"
+ + "%s.html" % phase_name
+ )
+
+ # Create generated subdirectory within the templates directory
+ generated_dir = os.path.join(config["site_dir"], "generated")
+ if not os.path.exists(generated_dir):
+ os.makedirs(generated_dir)
+ logger.info(f"Created generated templates directory: {generated_dir}")
+
+ output_html_fname = os.path.join(generated_dir, site_name)
+
+ # Write the file
+ logger.debug("writing %s html to %s.html" % (phase_name, output_html_fname))
+ with open(output_html_fname, "wt", encoding="utf-8") as outf:
+ outf.write(cur_html_template)
+
+ return site_name #output_html_fname
\ No newline at end of file
diff --git a/potato/server_utils/html_sanitizer.py b/potato/server_utils/html_sanitizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..1febb31547e7f2115b42c0906b474fda3cd7db5a
--- /dev/null
+++ b/potato/server_utils/html_sanitizer.py
@@ -0,0 +1,396 @@
+"""
+HTML Sanitizer Module
+
+Provides XSS-safe HTML sanitization for the annotation platform. This module
+allows legitimate span annotation HTML while blocking potentially dangerous
+elements and attributes.
+
+The sanitizer uses an allowlist approach - only explicitly permitted elements
+and attributes are kept, everything else is escaped or removed.
+
+Usage:
+ from potato.server_utils.html_sanitizer import sanitize_html
+
+ # In Jinja2 template:
+ {{ instance | sanitize_html }}
+"""
+
+import re
+import html
+import logging
+from typing import Set, Dict, List, Tuple
+from markupsafe import Markup
+
+logger = logging.getLogger(__name__)
+
+# Elements allowed in sanitized HTML
+ALLOWED_ELEMENTS: Set[str] = {
+ # Span annotations
+ 'span',
+ # Basic formatting (may be in source data)
+ 'b', 'i', 'u', 'strong', 'em', 'mark',
+ 's', 'del', 'ins',
+ # Line breaks and horizontal rules
+ 'br', 'hr', 'wbr',
+ # Dialogue/conversation layout elements
+ 'div',
+ # Structural elements for instructional content (Issue #120)
+ 'p',
+ 'ul', 'ol', 'li',
+ 'dl', 'dt', 'dd',
+ 'h1', 'h2', 'h3', 'h4', 'h5', 'h6',
+ # Tables for formatted content
+ 'table', 'thead', 'tbody', 'tr', 'th', 'td', 'caption',
+ # Inline semantics
+ 'sub', 'sup', 'small', 'code', 'pre', 'blockquote',
+ 'abbr', 'cite', 'kbd', 'samp', 'var',
+ # Collapsible sections
+ 'details', 'summary',
+ # Links (href sanitized against dangerous patterns)
+ 'a',
+ # Media and figures for instructional/survey content (Issue #129)
+ 'img', 'figure', 'figcaption',
+ # Ruby annotations for CJK text
+ 'ruby', 'rt', 'rp',
+}
+
+# Attributes allowed per element
+ALLOWED_ATTRIBUTES: Dict[str, Set[str]] = {
+ 'span': {
+ 'class',
+ 'style',
+ 'data-annotation-id',
+ 'data-label',
+ 'schema',
+ 'title',
+ },
+ 'div': {
+ 'class',
+ 'style',
+ 'data-speaker',
+ 'data-speaker-index',
+ },
+ 'mark': {'class', 'style'},
+ # Links โ href checked against dangerous patterns in _sanitize_attributes
+ 'a': {'href', 'title', 'target', 'rel'},
+ # Table elements
+ 'td': {'colspan', 'rowspan', 'style'},
+ 'th': {'colspan', 'rowspan', 'style', 'scope'},
+ 'table': {'class', 'style'},
+ # Ordered lists
+ 'ol': {'start', 'type'},
+ # Definition lists
+ 'dl': {'class', 'style'},
+ 'dt': {'class', 'style'},
+ 'dd': {'class', 'style'},
+ # Block-level elements that may need class/style
+ 'blockquote': {'class', 'style'},
+ 'p': {'class', 'style'},
+ 'pre': {'class', 'style'},
+ 'code': {'class'},
+ # Collapsible sections
+ 'details': {'class', 'style', 'open'},
+ 'summary': {'class', 'style'},
+ # Abbreviations with tooltip
+ 'abbr': {'title'},
+ # Deletion/insertion with optional metadata
+ 'del': {'datetime', 'cite'},
+ 'ins': {'datetime', 'cite'},
+ # Table captions
+ 'caption': {'class', 'style'},
+ # Images โ src checked against dangerous patterns like href
+ 'img': {'src', 'alt', 'title', 'style', 'width', 'height'},
+ # Figures and captions for instructional content
+ 'figure': {'class', 'style'},
+ 'figcaption': {'class', 'style'},
+ # Most elements get no attributes
+ '*': set(),
+}
+
+# Allowed CSS properties in style attributes
+ALLOWED_CSS_PROPERTIES: Set[str] = {
+ 'background-color',
+ 'color',
+ 'font-weight',
+ 'font-style',
+ 'font-family',
+ 'font-size',
+ 'text-decoration',
+ 'text-align',
+ 'line-height',
+ # Layout properties for dialogue/pairwise display
+ 'display',
+ 'width',
+ 'max-width',
+ 'padding',
+ 'padding-top', 'padding-bottom', 'padding-left', 'padding-right',
+ 'margin',
+ 'margin-top', 'margin-bottom', 'margin-left', 'margin-right',
+ 'box-sizing',
+ 'vertical-align',
+ 'gap',
+ # List styling
+ 'list-style-type',
+ # Borders
+ 'border',
+ 'border-radius',
+ 'border-collapse',
+}
+
+# Dangerous patterns to block
+DANGEROUS_PATTERNS = [
+ re.compile(r'javascript:', re.IGNORECASE),
+ re.compile(r'vbscript:', re.IGNORECASE),
+ re.compile(r'data:', re.IGNORECASE),
+ re.compile(r'expression\s*\(', re.IGNORECASE),
+]
+
+
+def sanitize_html(text: str) -> Markup:
+ """
+ Sanitize HTML content while preserving legitimate span annotations.
+
+ This function:
+ 1. Parses HTML using regex (lightweight, no external deps)
+ 2. Allows only whitelisted elements and attributes
+ 3. Sanitizes style attributes to only allow safe CSS
+ 4. Escapes all other content
+
+ Args:
+ text: The HTML content to sanitize
+
+ Returns:
+ Markup: Sanitized HTML safe for rendering (wrapped in Markup to prevent
+ double-escaping by Jinja2's auto-escape)
+
+ Example:
+ >>> sanitize_html('
text ')
+ Markup('
text ')
+
+ >>> sanitize_html('')
+ Markup('<script>alert("xss")</script>')
+ """
+ if not text:
+ return Markup("")
+
+ # Check for dangerous patterns in the raw text
+ for pattern in DANGEROUS_PATTERNS:
+ if pattern.search(text):
+ logger.warning(f"Blocked dangerous pattern in HTML content")
+ text = pattern.sub('', text)
+
+ result = []
+ pos = 0
+
+ # Regex to find HTML tags
+ tag_pattern = re.compile(
+ r'<(/?)(\w+)([^>]*)(/?)>',
+ re.IGNORECASE | re.DOTALL
+ )
+
+ for match in tag_pattern.finditer(text):
+ # Add escaped text before this tag
+ if match.start() > pos:
+ result.append(html.escape(text[pos:match.start()]))
+
+ is_close = match.group(1) == '/'
+ tag_name = match.group(2).lower()
+ attrs_str = match.group(3)
+ is_self_close = match.group(4) == '/'
+
+ if tag_name in ALLOWED_ELEMENTS:
+ # Build sanitized tag
+ if is_close:
+ result.append(f'{tag_name}>')
+ else:
+ sanitized_attrs = _sanitize_attributes(tag_name, attrs_str)
+ if is_self_close:
+ result.append(f'<{tag_name}{sanitized_attrs} />')
+ else:
+ result.append(f'<{tag_name}{sanitized_attrs}>')
+ else:
+ # Escape the entire tag
+ result.append(html.escape(match.group(0)))
+
+ pos = match.end()
+
+ # Add remaining text (escaped)
+ if pos < len(text):
+ result.append(html.escape(text[pos:]))
+
+ # Return as Markup to prevent Jinja2's auto-escape from escaping again
+ return Markup(''.join(result))
+
+
+def _sanitize_attributes(tag_name: str, attrs_str: str) -> str:
+ """
+ Sanitize attributes for a given tag.
+
+ Args:
+ tag_name: The tag name (lowercase)
+ attrs_str: The raw attributes string
+
+ Returns:
+ str: Sanitized attributes string (with leading space if non-empty)
+ """
+ if not attrs_str or not attrs_str.strip():
+ return ""
+
+ # Get allowed attributes for this tag
+ allowed = ALLOWED_ATTRIBUTES.get(tag_name, ALLOWED_ATTRIBUTES.get('*', set()))
+
+ # Parse attributes
+ attr_pattern = re.compile(
+ r'''(\w+(?:-\w+)*)\s*=\s*(?:"([^"]*)"|'([^']*)'|(\S+))''',
+ re.IGNORECASE
+ )
+
+ sanitized = []
+ for match in attr_pattern.finditer(attrs_str):
+ attr_name = match.group(1).lower()
+ # Get value from whichever group matched
+ attr_value = match.group(2) or match.group(3) or match.group(4) or ""
+
+ if attr_name not in allowed:
+ continue
+
+ # Special handling for href/src attributes โ block dangerous URLs
+ if attr_name in ('href', 'src'):
+ if any(p.search(attr_value) for p in DANGEROUS_PATTERNS):
+ logger.warning(f"Blocked dangerous pattern in {attr_name} attribute")
+ continue
+
+ # Special handling for style attribute
+ if attr_name == 'style':
+ attr_value = _sanitize_style(attr_value)
+ if not attr_value:
+ continue
+
+ # Special handling for class attribute
+ if attr_name == 'class':
+ attr_value = _sanitize_class(attr_value)
+
+ # Escape the value
+ escaped_value = html.escape(attr_value, quote=True)
+ sanitized.append(f'{attr_name}="{escaped_value}"')
+
+ # Security: auto-add rel="noopener noreferrer" for links with target="_blank"
+ if tag_name == 'a':
+ has_target_blank = any(s.startswith('target="') and '_blank' in s for s in sanitized)
+ has_rel = any(s.startswith('rel="') for s in sanitized)
+ if has_target_blank and not has_rel:
+ sanitized.append('rel="noopener noreferrer"')
+
+ if sanitized:
+ return ' ' + ' '.join(sanitized)
+ return ""
+
+
+def _sanitize_style(style: str) -> str:
+ """
+ Sanitize a CSS style attribute.
+
+ Only allows specific CSS properties that are known to be safe.
+
+ Args:
+ style: The style attribute value
+
+ Returns:
+ str: Sanitized style string
+ """
+ if not style:
+ return ""
+
+ # Check for dangerous patterns
+ for pattern in DANGEROUS_PATTERNS:
+ if pattern.search(style):
+ logger.warning("Blocked dangerous pattern in style attribute")
+ return ""
+
+ sanitized_props = []
+
+ # Parse CSS properties
+ for prop in style.split(';'):
+ prop = prop.strip()
+ if not prop:
+ continue
+
+ if ':' not in prop:
+ continue
+
+ name, value = prop.split(':', 1)
+ name = name.strip().lower()
+ value = value.strip()
+
+ if name in ALLOWED_CSS_PROPERTIES:
+ # Basic value validation - no functions except safe color functions
+ if 'url(' in value.lower():
+ continue
+ sanitized_props.append(f'{name}: {value}')
+
+ return '; '.join(sanitized_props)
+
+
+def _sanitize_class(class_str: str) -> str:
+ """
+ Sanitize a class attribute.
+
+ Only allows alphanumeric characters, hyphens, and underscores.
+
+ Args:
+ class_str: The class attribute value
+
+ Returns:
+ str: Sanitized class string
+ """
+ if not class_str:
+ return ""
+
+ # Split into individual classes
+ classes = class_str.split()
+
+ # Filter to safe class names
+ safe_pattern = re.compile(r'^[a-zA-Z_-][a-zA-Z0-9_-]*$')
+ safe_classes = [c for c in classes if safe_pattern.match(c)]
+
+ return ' '.join(safe_classes)
+
+
+def escape_for_attribute(text: str) -> str:
+ """
+ Escape text for use in an HTML attribute.
+
+ This is a stricter escape than html.escape() - it also escapes
+ backticks and other characters that could be used in template injection.
+
+ Args:
+ text: The text to escape
+
+ Returns:
+ str: Escaped text safe for attribute values
+ """
+ if not text:
+ return ""
+
+ return (
+ html.escape(text, quote=True)
+ .replace('`', '`')
+ .replace('$', '$')
+ )
+
+
+# Register as Jinja2 filter
+def register_jinja_filters(app):
+ """
+ Register HTML sanitization filters with a Flask app.
+
+ Call this during app initialization:
+ from potato.server_utils.html_sanitizer import register_jinja_filters
+ register_jinja_filters(app)
+
+ Args:
+ app: Flask application instance
+ """
+ app.jinja_env.filters['sanitize_html'] = sanitize_html
+ app.jinja_env.filters['escape_attr'] = escape_for_attribute
+ logger.info("Registered HTML sanitization Jinja2 filters")
diff --git a/potato/server_utils/iaa/__init__.py b/potato/server_utils/iaa/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..bea8a90819b9a3c59813b471799de963c5a95175
--- /dev/null
+++ b/potato/server_utils/iaa/__init__.py
@@ -0,0 +1,38 @@
+"""
+Inter-Annotator Agreement (IAA) metrics for Potato.
+
+This package computes IAA across the heterogeneous coverage sample produced
+by the overlap-sampling feature. Metrics are dispatched per annotation schema
+type (nominal, ordinal, continuous, multi-label, ranking, span).
+
+Public entry points:
+ compute_overlap_iaa(item_state_manager, user_state_manager, config)
+ End-to-end report for overlap-sample items that have reached their cap.
+ metrics_for_schema(annotation_type)
+ Inspect which metric functions apply to a given schema type.
+
+Lower-level metric functions live in nominal/ordinal/continuous/multilabel/ranking/span
+modules; alpha.py wraps Krippendorff's alpha (delegated to ``simpledorff``).
+"""
+
+from potato.server_utils.iaa import nominal, ordinal, continuous, multilabel, ranking, span, alpha
+from potato.server_utils.iaa.dispatcher import (
+ metrics_for_schema,
+ compute_overlap_iaa,
+ SchemaKind,
+ classify_schema,
+)
+
+__all__ = [
+ "nominal",
+ "ordinal",
+ "continuous",
+ "multilabel",
+ "ranking",
+ "span",
+ "alpha",
+ "metrics_for_schema",
+ "compute_overlap_iaa",
+ "SchemaKind",
+ "classify_schema",
+]
diff --git a/potato/server_utils/iaa/alpha.py b/potato/server_utils/iaa/alpha.py
new file mode 100644
index 0000000000000000000000000000000000000000..a3e789d5207f4044bbaccc9e0cda0a86a3dc4976
--- /dev/null
+++ b/potato/server_utils/iaa/alpha.py
@@ -0,0 +1,125 @@
+"""
+Krippendorff's alpha wrapper.
+
+Delegates to ``simpledorff`` (already a project dependency). Supports nominal,
+ordinal, interval, ratio, and MASI distance metrics. Accepts long-format data:
+a list of (annotator, item, label) triples.
+
+When ``simpledorff`` is unavailable, falls back to NaN with a logged warning.
+"""
+
+from __future__ import annotations
+
+from typing import Iterable, Sequence, Tuple, Union
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def _nominal_distance(a, b) -> float:
+ return 0.0 if a == b else 1.0
+
+
+def _ordinal_distance(a, b) -> float:
+ try:
+ return abs(float(a) - float(b))
+ except (TypeError, ValueError):
+ return 0.0 if a == b else 1.0
+
+
+def _interval_distance(a, b) -> float:
+ try:
+ return (float(a) - float(b)) ** 2
+ except (TypeError, ValueError):
+ return 0.0 if a == b else 1.0
+
+
+def _ratio_distance(a, b) -> float:
+ try:
+ a, b = float(a), float(b)
+ if a + b == 0:
+ return 0.0
+ return ((a - b) / (a + b)) ** 2
+ except (TypeError, ValueError):
+ return 0.0 if a == b else 1.0
+
+
+def _masi_distance(a, b) -> float:
+ """
+ MASI distance for multi-label sets. ``a`` and ``b`` are iterables of labels.
+ """
+ set_a = frozenset(a) if not isinstance(a, frozenset) else a
+ set_b = frozenset(b) if not isinstance(b, frozenset) else b
+ if not set_a and not set_b:
+ return 0.0
+ intersection = set_a & set_b
+ union = set_a | set_b
+ if not union:
+ return 0.0
+ jaccard = len(intersection) / len(union)
+ if set_a == set_b:
+ m = 1.0
+ elif set_a < set_b or set_b < set_a:
+ m = 2 / 3
+ elif intersection and set_a != set_b:
+ m = 1 / 3
+ else:
+ m = 0.0
+ return 1.0 - (jaccard * m)
+
+
+_DISTANCES = {
+ "nominal": _nominal_distance,
+ "ordinal": _ordinal_distance,
+ "interval": _interval_distance,
+ "ratio": _ratio_distance,
+ "masi": _masi_distance,
+}
+
+
+def krippendorff_alpha(
+ long_format: Sequence[Tuple[str, str, Union[str, float, frozenset]]],
+ level: str = "nominal",
+) -> float:
+ """
+ Krippendorff's alpha.
+
+ Args:
+ long_format: iterable of (annotator_id, item_id, value) tuples.
+ level: 'nominal', 'ordinal', 'interval', 'ratio', or 'masi'.
+
+ Returns:
+ Alpha as a float, or NaN if undefined.
+ """
+ if level not in _DISTANCES:
+ raise ValueError(f"Unknown level for Krippendorff's alpha: {level!r}")
+
+ try:
+ import simpledorff
+ import pandas as pd
+ except ImportError: # pragma: no cover
+ logger.warning("simpledorff/pandas unavailable; krippendorff_alpha returning NaN")
+ return float("nan")
+
+ rows = list(long_format)
+ if not rows:
+ return float("nan")
+ df = pd.DataFrame(rows, columns=["annotator", "item", "value"])
+ if df["item"].nunique() < 2 or df["annotator"].nunique() < 2:
+ return float("nan")
+
+ dist = _DISTANCES[level]
+ try:
+ return float(
+ simpledorff.calculate_krippendorffs_alpha_for_df(
+ df,
+ experiment_col="item",
+ annotator_col="annotator",
+ class_col="value",
+ metric_fn=dist,
+ )
+ )
+ except Exception as exc: # pragma: no cover
+ logger.warning("krippendorff_alpha failed: %s", exc)
+ return float("nan")
diff --git a/potato/server_utils/iaa/continuous.py b/potato/server_utils/iaa/continuous.py
new file mode 100644
index 0000000000000000000000000000000000000000..945ad7cd853efd75362684d364abb21a995b6d57
--- /dev/null
+++ b/potato/server_utils/iaa/continuous.py
@@ -0,0 +1,126 @@
+"""
+Continuous IAA metrics: Pearson r, MAE, RMSE, intra-class correlation (ICC).
+
+ICC implementation follows Shrout & Fleiss (1979). We expose:
+ icc_2_1 โ single-rater ICC(2,1) (two-way random, agreement, single measure)
+ icc_2_k โ average-rater ICC(2,k) (two-way random, agreement, average measure)
+"""
+
+from __future__ import annotations
+
+from math import isnan, sqrt
+from typing import Sequence
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def _to_float(seq: Sequence) -> list:
+ out = []
+ for v in seq:
+ try:
+ out.append(float(v))
+ except (TypeError, ValueError):
+ out.append(float("nan"))
+ return out
+
+
+def pearson_r(values_a: Sequence, values_b: Sequence) -> float:
+ a = _to_float(values_a)
+ b = _to_float(values_b)
+ pairs = [(x, y) for x, y in zip(a, b) if not (isnan(x) or isnan(y))]
+ if len(pairs) < 2:
+ return float("nan")
+ try:
+ from scipy.stats import pearsonr
+ r, _ = pearsonr([x for x, _ in pairs], [y for _, y in pairs])
+ return float(r) if not isnan(r) else float("nan")
+ except ImportError: # pragma: no cover
+ pass
+ n = len(pairs)
+ sa = sum(x for x, _ in pairs)
+ sb = sum(y for _, y in pairs)
+ sab = sum(x * y for x, y in pairs)
+ saa = sum(x * x for x, _ in pairs)
+ sbb = sum(y * y for _, y in pairs)
+ num = n * sab - sa * sb
+ den = sqrt((n * saa - sa * sa) * (n * sbb - sb * sb))
+ if den == 0:
+ return float("nan")
+ return num / den
+
+
+def mae(values_a: Sequence, values_b: Sequence) -> float:
+ a = _to_float(values_a)
+ b = _to_float(values_b)
+ pairs = [(x, y) for x, y in zip(a, b) if not (isnan(x) or isnan(y))]
+ if not pairs:
+ return float("nan")
+ return sum(abs(x - y) for x, y in pairs) / len(pairs)
+
+
+def rmse(values_a: Sequence, values_b: Sequence) -> float:
+ a = _to_float(values_a)
+ b = _to_float(values_b)
+ pairs = [(x, y) for x, y in zip(a, b) if not (isnan(x) or isnan(y))]
+ if not pairs:
+ return float("nan")
+ return sqrt(sum((x - y) ** 2 for x, y in pairs) / len(pairs))
+
+
+def _icc_components(matrix):
+ """Mean squares for a two-way ANOVA: MSR (rows/items), MSC (cols/raters), MSE."""
+ try:
+ import numpy as np
+ except ImportError: # pragma: no cover
+ return None
+ arr = np.asarray(matrix, dtype=float)
+ if arr.ndim != 2 or arr.shape[0] < 2 or arr.shape[1] < 2:
+ return None
+ if np.isnan(arr).any():
+ # listwise deletion of items with any missing rating
+ arr = arr[~np.isnan(arr).any(axis=1)]
+ if arr.shape[0] < 2:
+ return None
+ n, k = arr.shape
+ grand = arr.mean()
+ row_means = arr.mean(axis=1)
+ col_means = arr.mean(axis=0)
+ ss_total = ((arr - grand) ** 2).sum()
+ ss_rows = k * ((row_means - grand) ** 2).sum()
+ ss_cols = n * ((col_means - grand) ** 2).sum()
+ ss_err = ss_total - ss_rows - ss_cols
+ df_rows = n - 1
+ df_cols = k - 1
+ df_err = (n - 1) * (k - 1)
+ if df_err <= 0:
+ return None
+ msr = ss_rows / df_rows
+ msc = ss_cols / df_cols
+ mse = ss_err / df_err
+ return msr, msc, mse, n, k
+
+
+def icc_2_1(matrix) -> float:
+ """ICC(2,1): two-way random effects, single rater, absolute agreement."""
+ comps = _icc_components(matrix)
+ if comps is None:
+ return float("nan")
+ msr, msc, mse, n, k = comps
+ denom = msr + (k - 1) * mse + k * (msc - mse) / n
+ if denom == 0:
+ return float("nan")
+ return (msr - mse) / denom
+
+
+def icc_2_k(matrix) -> float:
+ """ICC(2,k): two-way random effects, average of k raters, absolute agreement."""
+ comps = _icc_components(matrix)
+ if comps is None:
+ return float("nan")
+ msr, msc, mse, n, k = comps
+ denom = msr + (msc - mse) / n
+ if denom == 0:
+ return float("nan")
+ return (msr - mse) / denom
diff --git a/potato/server_utils/iaa/dispatcher.py b/potato/server_utils/iaa/dispatcher.py
new file mode 100644
index 0000000000000000000000000000000000000000..19fd80e45c7edbe68a34d1f2cd17d21c9384dc6f
--- /dev/null
+++ b/potato/server_utils/iaa/dispatcher.py
@@ -0,0 +1,553 @@
+"""
+Schema-to-metric dispatcher and the top-level overlap-IAA report.
+
+The dispatcher inspects a schema's ``annotation_type`` and (where relevant)
+its labels block to decide which family of IAA metrics applies, then runs
+those metrics across the overlap-sample items that have reached their cap.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from enum import Enum
+from typing import Any, Dict, Iterable, List, Optional
+
+import logging
+
+from potato.server_utils.iaa import nominal, ordinal, continuous, multilabel, ranking, span, alpha
+
+logger = logging.getLogger(__name__)
+
+
+class SchemaKind(str, Enum):
+ NOMINAL = "nominal"
+ ORDINAL = "ordinal"
+ CONTINUOUS = "continuous"
+ MULTILABEL = "multilabel"
+ RANKING = "ranking"
+ SPAN = "span"
+ TEXT = "text" # free-form text, no automatic IAA
+ UNSUPPORTED = "unsupported"
+
+
+_KIND_BY_TYPE = {
+ # Nominal (single-label categorical)
+ "radio": SchemaKind.NOMINAL,
+ "select": SchemaKind.NOMINAL,
+ "triage": SchemaKind.NOMINAL,
+ # Ordinal
+ "likert": SchemaKind.ORDINAL,
+ "confidence": SchemaKind.ORDINAL,
+ "semantic_differential": SchemaKind.ORDINAL,
+ "range_slider": SchemaKind.ORDINAL,
+ "vas": SchemaKind.ORDINAL,
+ # Continuous
+ "slider": SchemaKind.CONTINUOUS,
+ "number": SchemaKind.CONTINUOUS,
+ "multirate": SchemaKind.CONTINUOUS,
+ "constant_sum": SchemaKind.CONTINUOUS,
+ "soft_label": SchemaKind.CONTINUOUS,
+ # Multi-label
+ "multiselect": SchemaKind.MULTILABEL, # may be downgraded to NOMINAL if max=1
+ "hierarchical_multiselect": SchemaKind.MULTILABEL,
+ "card_sort": SchemaKind.MULTILABEL,
+ # Ranking
+ "ranking": SchemaKind.RANKING,
+ "bws": SchemaKind.RANKING,
+ "pairwise": SchemaKind.RANKING,
+ "conjoint": SchemaKind.RANKING,
+ "best_worst_scaling": SchemaKind.RANKING,
+ # Span
+ "span": SchemaKind.SPAN,
+ "error_span": SchemaKind.SPAN,
+ "event_annotation": SchemaKind.SPAN,
+ "coreference": SchemaKind.SPAN,
+ "extractive_qa": SchemaKind.SPAN,
+ "span_link": SchemaKind.SPAN,
+ "tree_annotation": SchemaKind.SPAN,
+ # Text
+ "textbox": SchemaKind.TEXT,
+ "text_edit": SchemaKind.TEXT,
+ # Skipped
+ "pure_display": SchemaKind.UNSUPPORTED,
+ "video": SchemaKind.UNSUPPORTED,
+ "audio_annotation": SchemaKind.UNSUPPORTED,
+ "video_annotation": SchemaKind.UNSUPPORTED,
+ "image_annotation": SchemaKind.UNSUPPORTED,
+}
+
+
+def classify_schema(scheme: Dict[str, Any]) -> SchemaKind:
+ """Classify a schema definition into an IAA-relevant kind."""
+ atype = (scheme.get("annotation_type") or "").strip().lower()
+ kind = _KIND_BY_TYPE.get(atype, SchemaKind.UNSUPPORTED)
+ # Downgrade multiselect with max_choices == 1 to NOMINAL
+ if kind == SchemaKind.MULTILABEL and atype == "multiselect":
+ max_choices = scheme.get("max_choices") or scheme.get("max_selections")
+ if max_choices == 1:
+ return SchemaKind.NOMINAL
+ return kind
+
+
+def metrics_for_schema(scheme: Dict[str, Any]) -> List[str]:
+ """Return human-readable names of metrics that apply to ``scheme``."""
+ kind = classify_schema(scheme)
+ table = {
+ SchemaKind.NOMINAL: ["percent_agreement", "cohen_kappa", "fleiss_kappa", "alpha_nominal"],
+ SchemaKind.ORDINAL: ["weighted_kappa_linear", "weighted_kappa_quadratic", "spearman_rho", "alpha_ordinal"],
+ SchemaKind.CONTINUOUS: ["pearson_r", "mae", "rmse", "alpha_interval", "icc_2_k"],
+ SchemaKind.MULTILABEL: ["mean_jaccard", "alpha_masi"],
+ SchemaKind.RANKING: ["kendall_tau", "spearman_footrule"],
+ SchemaKind.SPAN: [
+ "token_level_kappa", "span_f1_exact", "span_f1_partial",
+ "krippendorff_alpha_u", "gamma_mathet",
+ ],
+ SchemaKind.TEXT: [],
+ SchemaKind.UNSUPPORTED: [],
+ }
+ return list(table[kind])
+
+
+# ---------------------------------------------------------------------------
+# Data extraction from Potato's per-user annotation structures
+# ---------------------------------------------------------------------------
+
+def _label_value(label) -> Any:
+ """Extract a comparable value from a Label object (or dict)."""
+ if isinstance(label, dict):
+ return label.get("name") or label.get("value")
+ return getattr(label, "name", None) or getattr(label, "value", None)
+
+
+def _gather_labels(
+ instance_ids: Iterable[str],
+ user_states: Dict[str, Any],
+ schema_name: str,
+):
+ """
+ Per item, return {user_id:
} for one schema.
+
+ For nominal/ordinal/continuous schemas the value is a scalar (the chosen
+ label name or numeric rating). For multi-label schemas, it's a list.
+ """
+ rows: Dict[str, Dict[str, Any]] = {}
+ for iid in instance_ids:
+ per_user: Dict[str, Any] = {}
+ for uid, ustate in user_states.items():
+ labels_by_schema = ustate.get_label_annotations(iid)
+ if not labels_by_schema:
+ continue
+ labels = labels_by_schema.get(schema_name)
+ if not labels:
+ continue
+ vals = [_label_value(l) for l in labels]
+ vals = [v for v in vals if v is not None]
+ if not vals:
+ continue
+ per_user[uid] = vals
+ if per_user:
+ rows[iid] = per_user
+ return rows
+
+
+def _gather_spans(
+ instance_ids: Iterable[str],
+ user_states: Dict[str, Any],
+ schema_name: str,
+):
+ rows: Dict[str, Dict[str, list]] = {}
+ for iid in instance_ids:
+ per_user = {}
+ for uid, ustate in user_states.items():
+ spans_by_schema = ustate.get_span_annotations(iid)
+ if not spans_by_schema:
+ continue
+ spans = spans_by_schema.get(schema_name) or []
+ if not spans:
+ continue
+ per_user[uid] = list(spans)
+ if per_user:
+ rows[iid] = per_user
+ return rows
+
+
+def _text_length_for_item(item) -> int:
+ """Best-effort character length of the item text used for span IAA."""
+ if item is None:
+ return 0
+ try:
+ text = item.get_text()
+ except Exception:
+ return 0
+ return len(text) if isinstance(text, str) else 0
+
+
+# ---------------------------------------------------------------------------
+# Metric computation per kind
+# ---------------------------------------------------------------------------
+
+def _aggregate_nominal(rows):
+ long_rows = []
+ pairwise_kappa = []
+ fleiss_inputs = []
+ users_seen = set()
+ for iid, per_user in rows.items():
+ # Collapse multi-value into the first chosen label (single-label schema)
+ flat = {u: v[0] for u, v in per_user.items() if v}
+ if len(flat) < 2:
+ continue
+ users_seen.update(flat)
+ for u, val in flat.items():
+ long_rows.append((u, iid, val))
+ fleiss_inputs.append(dict(Counter_(flat.values())))
+
+ pair_users = sorted(users_seen)
+ seqs_by_user: Dict[str, list] = {u: [] for u in pair_users}
+ aligned_iids = []
+ for iid, per_user in rows.items():
+ flat = {u: v[0] for u, v in per_user.items() if v}
+ if all(u in flat for u in pair_users):
+ aligned_iids.append(iid)
+ for u in pair_users:
+ seqs_by_user[u].append(flat[u])
+
+ return {
+ "alpha_nominal": alpha.krippendorff_alpha(long_rows, level="nominal"),
+ "fleiss_kappa": nominal.fleiss_kappa(fleiss_inputs),
+ "pairwise_cohen_kappa": nominal.pairwise_cohen_kappa(seqs_by_user) if seqs_by_user else float("nan"),
+ "n_items": len(rows),
+ "n_aligned_items": len(aligned_iids),
+ "n_annotators": len(pair_users),
+ }
+
+
+def _aggregate_ordinal(rows):
+ long_rows = []
+ seqs_by_user: Dict[str, list] = defaultdict(list)
+ aligned_users = None
+ for iid, per_user in rows.items():
+ flat = {u: v[0] for u, v in per_user.items() if v}
+ if len(flat) < 2:
+ continue
+ for u, val in flat.items():
+ long_rows.append((u, iid, val))
+ if aligned_users is None:
+ aligned_users = set(flat)
+ else:
+ aligned_users &= set(flat)
+ for u, val in flat.items():
+ seqs_by_user[u].append(val)
+ weighted_lin = _pairwise_mean(seqs_by_user, ordinal.weighted_kappa, weights="linear")
+ weighted_quad = _pairwise_mean(seqs_by_user, ordinal.weighted_kappa, weights="quadratic")
+ rho = _pairwise_mean(seqs_by_user, ordinal.spearman_rho)
+ return {
+ "weighted_kappa_linear": weighted_lin,
+ "weighted_kappa_quadratic": weighted_quad,
+ "spearman_rho": rho,
+ "alpha_ordinal": alpha.krippendorff_alpha(long_rows, level="ordinal"),
+ "n_items": len(rows),
+ "n_annotators": len(seqs_by_user),
+ }
+
+
+def _aggregate_continuous(rows):
+ long_rows = []
+ seqs_by_user: Dict[str, list] = defaultdict(list)
+ for iid, per_user in rows.items():
+ flat = {}
+ for u, v in per_user.items():
+ try:
+ flat[u] = float(v[0])
+ except (TypeError, ValueError):
+ continue
+ if len(flat) < 2:
+ continue
+ for u, val in flat.items():
+ long_rows.append((u, iid, val))
+ seqs_by_user[u].append(val)
+
+ pearson = _pairwise_mean(seqs_by_user, continuous.pearson_r)
+ mae_val = _pairwise_mean(seqs_by_user, continuous.mae)
+ rmse_val = _pairwise_mean(seqs_by_user, continuous.rmse)
+
+ # ICC needs an items x raters matrix where every rater rates every item.
+ users = sorted(seqs_by_user)
+ aligned_iids = []
+ matrix = []
+ for iid, per_user in rows.items():
+ try:
+ row = [float(per_user[u][0]) for u in users]
+ except (KeyError, TypeError, ValueError):
+ continue
+ matrix.append(row)
+ aligned_iids.append(iid)
+ icc_k = continuous.icc_2_k(matrix) if matrix and users else float("nan")
+
+ return {
+ "pearson_r": pearson,
+ "mae": mae_val,
+ "rmse": rmse_val,
+ "alpha_interval": alpha.krippendorff_alpha(long_rows, level="interval"),
+ "icc_2_k": icc_k,
+ "n_items": len(rows),
+ "n_aligned_items": len(aligned_iids),
+ "n_annotators": len(users),
+ }
+
+
+def _aggregate_multilabel(rows):
+ long_rows = []
+ label_sets_by_user: Dict[str, list] = defaultdict(list)
+ for iid, per_user in rows.items():
+ flat = {u: frozenset(v) for u, v in per_user.items() if v}
+ if len(flat) < 2:
+ continue
+ for u, val in flat.items():
+ long_rows.append((u, iid, val))
+ label_sets_by_user[u].append(val)
+ return {
+ "mean_jaccard": multilabel.mean_jaccard(label_sets_by_user),
+ "alpha_masi": multilabel.alpha_masi(long_rows),
+ "n_items": len(rows),
+ "n_annotators": len(label_sets_by_user),
+ }
+
+
+def _aggregate_ranking(rows):
+ seqs_by_user: Dict[str, list] = defaultdict(list)
+ for iid, per_user in rows.items():
+ flat = {u: list(v) for u, v in per_user.items() if v}
+ if len(flat) < 2:
+ continue
+ for u, val in flat.items():
+ seqs_by_user[u].append(val)
+ tau = _pairwise_rank_mean(seqs_by_user, ranking.kendall_tau)
+ footrule = _pairwise_rank_mean(seqs_by_user, ranking.spearman_footrule)
+ return {
+ "kendall_tau": tau,
+ "spearman_footrule": footrule,
+ "n_items": len(rows),
+ "n_annotators": len(seqs_by_user),
+ }
+
+
+def _aggregate_span(span_rows, item_lookup):
+ token_kappas = []
+ f1_exact = []
+ f1_partial = []
+ alphas_u = []
+ gammas = []
+ n_items = 0
+ annotators = set()
+ for iid, per_user in span_rows.items():
+ if len(per_user) < 2:
+ continue
+ item = item_lookup.get(iid)
+ length = _text_length_for_item(item)
+ if length <= 0:
+ continue
+ annotators.update(per_user)
+ n_items += 1
+ try:
+ tk = span.token_level_kappa(per_user, length)
+ if tk == tk:
+ token_kappas.append(tk)
+ except Exception as exc:
+ logger.debug("token_level_kappa failed on %s: %s", iid, exc)
+ try:
+ exact = span.pairwise_span_f1(per_user, partial=False)
+ partial = span.pairwise_span_f1(per_user, partial=True)
+ if exact == exact:
+ f1_exact.append(exact)
+ if partial == partial:
+ f1_partial.append(partial)
+ except Exception as exc:
+ logger.debug("span_f1 failed on %s: %s", iid, exc)
+ try:
+ au = span.krippendorff_alpha_u(per_user, length)
+ if au == au:
+ alphas_u.append(au)
+ except Exception as exc:
+ logger.debug("alpha_u failed on %s: %s", iid, exc)
+ try:
+ g = span.gamma(per_user, length=length)
+ if g == g:
+ gammas.append(g)
+ except Exception as exc:
+ logger.debug("gamma failed on %s: %s", iid, exc)
+
+ def _mean(xs):
+ return sum(xs) / len(xs) if xs else float("nan")
+
+ return {
+ "token_level_kappa": _mean(token_kappas),
+ "span_f1_exact": _mean(f1_exact),
+ "span_f1_partial": _mean(f1_partial),
+ "krippendorff_alpha_u": _mean(alphas_u),
+ "gamma_mathet": _mean(gammas),
+ "n_items": n_items,
+ "n_annotators": len(annotators),
+ }
+
+
+# ---------------------------------------------------------------------------
+# Pairwise helpers
+# ---------------------------------------------------------------------------
+
+def _pairwise_mean(seqs_by_user, fn, **kwargs):
+ users = list(seqs_by_user)
+ if len(users) < 2:
+ return float("nan")
+ out = []
+ for i in range(len(users)):
+ for j in range(i + 1, len(users)):
+ a = seqs_by_user[users[i]]
+ b = seqs_by_user[users[j]]
+ m = min(len(a), len(b))
+ if m < 2:
+ continue
+ try:
+ v = fn(a[:m], b[:m], **kwargs) if kwargs else fn(a[:m], b[:m])
+ if v == v:
+ out.append(v)
+ except Exception as exc:
+ logger.debug("pairwise metric %s failed: %s", fn.__name__, exc)
+ return sum(out) / len(out) if out else float("nan")
+
+
+def _pairwise_rank_mean(seqs_by_user, fn):
+ users = list(seqs_by_user)
+ if len(users) < 2:
+ return float("nan")
+ out = []
+ for i in range(len(users)):
+ for j in range(i + 1, len(users)):
+ a = seqs_by_user[users[i]]
+ b = seqs_by_user[users[j]]
+ m = min(len(a), len(b))
+ for k in range(m):
+ try:
+ v = fn(a[k], b[k])
+ if v == v:
+ out.append(v)
+ except Exception:
+ continue
+ return sum(out) / len(out) if out else float("nan")
+
+
+# ---------------------------------------------------------------------------
+# Top-level entry point
+# ---------------------------------------------------------------------------
+
+def compute_overlap_iaa(item_state_manager, user_state_manager, config: Dict[str, Any]) -> Dict[str, Any]:
+ """
+ Compute IAA across the overlap-sample items that have reached their cap.
+
+ Returns a dict shape:
+ {
+ "schemas": {
+ "": {
+ "kind": "",
+ "annotation_type": "",
+ "metrics": { : , ... },
+ "n_items": int,
+ "n_annotators": int,
+ }
+ },
+ "items": {
+ "": {
+ "annotators": [...],
+ "cap": int,
+ "schemas": {
+ "": { ... per-item metric breakdown ... }
+ }
+ }
+ },
+ "n_overlap_items": int,
+ }
+ """
+ schemes = _extract_schemes(config)
+ if not schemes:
+ return {"schemas": {}, "items": {}, "n_overlap_items": 0}
+
+ # Overlap items: per-item cap >= 2 AND saturated.
+ overlap_items = []
+ for iid, item in item_state_manager.instance_id_to_instance.items():
+ cap = item_state_manager._get_annotator_cap_for_item(iid)
+ if cap is None or cap < 2:
+ continue
+ if len(item_state_manager.instance_annotators[iid]) < cap:
+ continue
+ overlap_items.append(iid)
+
+ # Build {user_id: user_state} for users who touched any overlap item.
+ relevant_user_ids = set()
+ for iid in overlap_items:
+ relevant_user_ids.update(item_state_manager.instance_annotators[iid])
+ user_states = {}
+ for uid in relevant_user_ids:
+ ustate = user_state_manager.get_user_state(uid) if hasattr(user_state_manager, "get_user_state") else None
+ if ustate is not None:
+ user_states[uid] = ustate
+
+ schema_report: Dict[str, Any] = {}
+ item_report: Dict[str, Any] = {iid: {
+ "annotators": sorted(item_state_manager.instance_annotators[iid]),
+ "cap": item_state_manager._get_annotator_cap_for_item(iid),
+ "schemas": {},
+ } for iid in overlap_items}
+
+ for scheme in schemes:
+ name = scheme.get("name")
+ if not name:
+ continue
+ kind = classify_schema(scheme)
+ if kind in (SchemaKind.TEXT, SchemaKind.UNSUPPORTED):
+ continue
+ if kind == SchemaKind.SPAN:
+ rows = _gather_spans(overlap_items, user_states, name)
+ metrics = _aggregate_span(rows, item_state_manager.instance_id_to_instance)
+ else:
+ rows = _gather_labels(overlap_items, user_states, name)
+ if kind == SchemaKind.NOMINAL:
+ metrics = _aggregate_nominal(rows)
+ elif kind == SchemaKind.ORDINAL:
+ metrics = _aggregate_ordinal(rows)
+ elif kind == SchemaKind.CONTINUOUS:
+ metrics = _aggregate_continuous(rows)
+ elif kind == SchemaKind.MULTILABEL:
+ metrics = _aggregate_multilabel(rows)
+ elif kind == SchemaKind.RANKING:
+ metrics = _aggregate_ranking(rows)
+ else:
+ continue
+ schema_report[name] = {
+ "kind": kind.value,
+ "annotation_type": scheme.get("annotation_type"),
+ "metrics": metrics,
+ }
+ for iid in rows if kind != SchemaKind.SPAN else rows:
+ item_report.setdefault(iid, {"annotators": [], "cap": -1, "schemas": {}})
+ item_report[iid]["schemas"][name] = {"n_annotators": len(rows[iid])}
+
+ return {
+ "schemas": schema_report,
+ "items": item_report,
+ "n_overlap_items": len(overlap_items),
+ }
+
+
+def _extract_schemes(config: Dict[str, Any]):
+ """Pull annotation_schemes from the config (top-level or under a phase)."""
+ if "annotation_schemes" in config and isinstance(config["annotation_schemes"], list):
+ return config["annotation_schemes"]
+ schemes = []
+ phases = config.get("phases", {}) or {}
+ for key, val in phases.items():
+ if isinstance(val, dict) and isinstance(val.get("annotation_schemes"), list):
+ schemes.extend(val["annotation_schemes"])
+ return schemes
+
+
+# Local imports placed at the bottom to avoid circular imports at module load.
+from collections import Counter as Counter_ # noqa: E402
diff --git a/potato/server_utils/iaa/multilabel.py b/potato/server_utils/iaa/multilabel.py
new file mode 100644
index 0000000000000000000000000000000000000000..ddc89fe9b1bdc16c839ac69451ec9880279e57bd
--- /dev/null
+++ b/potato/server_utils/iaa/multilabel.py
@@ -0,0 +1,52 @@
+"""
+Multi-label IAA metrics for schemas where each annotator can select a set of labels
+per item (e.g., multiselect, hierarchical_multiselect, card_sort).
+
+Provides MASI distance, Jaccard distance, and pairwise alpha-MASI.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Iterable, Sequence
+
+from potato.server_utils.iaa.alpha import krippendorff_alpha, _masi_distance
+
+
+def jaccard_distance(set_a: Iterable, set_b: Iterable) -> float:
+ a = frozenset(set_a)
+ b = frozenset(set_b)
+ if not a and not b:
+ return 0.0
+ union = a | b
+ if not union:
+ return 0.0
+ return 1.0 - len(a & b) / len(union)
+
+
+def masi_distance(set_a: Iterable, set_b: Iterable) -> float:
+ return _masi_distance(set_a, set_b)
+
+
+def mean_jaccard(label_sets_by_user: Dict[str, Sequence[Iterable]]) -> float:
+ """Average pairwise (1 - Jaccard distance) across users and items."""
+ users = list(label_sets_by_user)
+ if len(users) < 2:
+ return float("nan")
+ sims = []
+ for i in range(len(users)):
+ a = list(label_sets_by_user[users[i]])
+ for j in range(i + 1, len(users)):
+ b = list(label_sets_by_user[users[j]])
+ m = min(len(a), len(b))
+ if m == 0:
+ continue
+ for k in range(m):
+ sims.append(1.0 - jaccard_distance(a[k], b[k]))
+ if not sims:
+ return float("nan")
+ return sum(sims) / len(sims)
+
+
+def alpha_masi(long_format_sets) -> float:
+ """Krippendorff's alpha with MASI distance on multi-label sets."""
+ return krippendorff_alpha(long_format_sets, level="masi")
diff --git a/potato/server_utils/iaa/nominal.py b/potato/server_utils/iaa/nominal.py
new file mode 100644
index 0000000000000000000000000000000000000000..38a44f8dd3d0cd4e8c34c05de01ea8568f834189
--- /dev/null
+++ b/potato/server_utils/iaa/nominal.py
@@ -0,0 +1,132 @@
+"""
+Nominal IAA metrics: percent agreement, Cohen's kappa, Fleiss' kappa.
+
+Inputs are lists keyed by item: for two-annotator metrics, two equal-length
+label lists; for multi-annotator metrics, a list of (annotator_id -> label) dicts.
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+from math import isclose
+from typing import Dict, List, Sequence
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def percent_agreement(labels_a: Sequence, labels_b: Sequence) -> float:
+ """Fraction of items on which two annotators agree."""
+ if len(labels_a) != len(labels_b):
+ raise ValueError("label lists must be the same length")
+ if not labels_a:
+ return float("nan")
+ agree = sum(1 for a, b in zip(labels_a, labels_b) if a == b)
+ return agree / len(labels_a)
+
+
+def cohen_kappa(labels_a: Sequence, labels_b: Sequence) -> float:
+ """
+ Cohen's kappa for two annotators on nominal categories.
+
+ Uses sklearn if available (handles ties and edge cases well); falls back
+ to a direct implementation otherwise.
+ """
+ if len(labels_a) != len(labels_b):
+ raise ValueError("label lists must be the same length")
+ if not labels_a:
+ return float("nan")
+ try:
+ from sklearn.metrics import cohen_kappa_score
+ return float(cohen_kappa_score(list(labels_a), list(labels_b)))
+ except ImportError: # pragma: no cover
+ pass
+
+ n = len(labels_a)
+ po = percent_agreement(labels_a, labels_b)
+ counts_a = Counter(labels_a)
+ counts_b = Counter(labels_b)
+ pe = sum(counts_a[c] * counts_b[c] for c in set(counts_a) | set(counts_b)) / (n * n)
+ if isclose(pe, 1.0):
+ return 1.0 if isclose(po, 1.0) else 0.0
+ return (po - pe) / (1 - pe)
+
+
+def fleiss_kappa(per_item_label_counts: List[Dict[str, int]]) -> float:
+ """
+ Fleiss' kappa for >=2 annotators on nominal categories.
+
+ Args:
+ per_item_label_counts: one dict per item mapping label -> number of
+ annotators who chose it. Each item dict must sum to the same N
+ (the number of annotators rating that item). Items where N < 2
+ are skipped.
+
+ Returns:
+ Fleiss' kappa as a float, or NaN if undefined.
+ """
+ # Use only items rated by at least 2 annotators.
+ rated = [d for d in per_item_label_counts if sum(d.values()) >= 2]
+ if not rated:
+ return float("nan")
+
+ ns = [sum(d.values()) for d in rated]
+ if len(set(ns)) != 1:
+ # Variable-N Fleiss' kappa is rare in practice; restrict to majority N.
+ from statistics import mode
+ majority_n = mode(ns)
+ rated = [d for d, n in zip(rated, ns) if n == majority_n]
+ ns = [majority_n] * len(rated)
+ if not rated:
+ return float("nan")
+
+ n = ns[0]
+ categories = sorted({c for d in rated for c in d})
+ if n < 2 or not categories:
+ return float("nan")
+
+ n_items = len(rated)
+ # Per-item agreement P_i
+ p_is = []
+ for d in rated:
+ total = sum(d.get(c, 0) ** 2 for c in categories)
+ p_is.append((total - n) / (n * (n - 1)))
+ p_bar = sum(p_is) / n_items
+ # Marginal proportions per category
+ p_js = []
+ for c in categories:
+ s = sum(d.get(c, 0) for d in rated)
+ p_js.append(s / (n_items * n))
+ p_e = sum(p * p for p in p_js)
+ if isclose(p_e, 1.0):
+ return 1.0 if isclose(p_bar, 1.0) else 0.0
+ return (p_bar - p_e) / (1 - p_e)
+
+
+def pairwise_cohen_kappa(annotations_by_user: Dict[str, Sequence]) -> float:
+ """
+ Mean Cohen's kappa across every distinct pair of annotators.
+
+ annotations_by_user maps user_id -> aligned label sequence (same length per user).
+ Users contributing fewer than the maximum length are restricted to their
+ overlap with each partner.
+ """
+ users = list(annotations_by_user)
+ if len(users) < 2:
+ return float("nan")
+ kappas = []
+ for i in range(len(users)):
+ for j in range(i + 1, len(users)):
+ a = list(annotations_by_user[users[i]])
+ b = list(annotations_by_user[users[j]])
+ m = min(len(a), len(b))
+ if m == 0:
+ continue
+ try:
+ kappas.append(cohen_kappa(a[:m], b[:m]))
+ except ValueError:
+ continue
+ if not kappas:
+ return float("nan")
+ return sum(kappas) / len(kappas)
diff --git a/potato/server_utils/iaa/ordinal.py b/potato/server_utils/iaa/ordinal.py
new file mode 100644
index 0000000000000000000000000000000000000000..a5dc4fc48f6da649492de041fa44e07804b156b4
--- /dev/null
+++ b/potato/server_utils/iaa/ordinal.py
@@ -0,0 +1,66 @@
+"""
+Ordinal IAA metrics: weighted kappa (linear + quadratic), Spearman's rho.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def _coerce_ordinal(values: Sequence) -> list:
+ """Try to coerce a sequence of (str|int|float) ratings into numeric ranks."""
+ coerced = []
+ for v in values:
+ if isinstance(v, (int, float)):
+ coerced.append(float(v))
+ else:
+ try:
+ coerced.append(float(v))
+ except (TypeError, ValueError):
+ # Fall back to lexical ordering by stable string sort
+ coerced.append(str(v))
+ if any(isinstance(c, str) for c in coerced):
+ rank = {c: i for i, c in enumerate(sorted(set(coerced)))}
+ return [rank[c] for c in coerced]
+ return coerced
+
+
+def weighted_kappa(labels_a: Sequence, labels_b: Sequence, weights: str = "quadratic") -> float:
+ """
+ Cohen's weighted kappa for ordinal categories.
+
+ weights: 'linear' or 'quadratic' (CKD convention).
+ """
+ if len(labels_a) != len(labels_b):
+ raise ValueError("label lists must be the same length")
+ if not labels_a:
+ return float("nan")
+ try:
+ from sklearn.metrics import cohen_kappa_score
+ a = _coerce_ordinal(labels_a)
+ b = _coerce_ordinal(labels_b)
+ return float(cohen_kappa_score(a, b, weights=weights))
+ except ImportError: # pragma: no cover
+ logger.warning("sklearn unavailable; weighted_kappa returning NaN")
+ return float("nan")
+
+
+def spearman_rho(labels_a: Sequence, labels_b: Sequence) -> float:
+ """Spearman rank correlation between two annotators."""
+ if len(labels_a) != len(labels_b):
+ raise ValueError("label lists must be the same length")
+ if len(labels_a) < 2:
+ return float("nan")
+ try:
+ from scipy.stats import spearmanr
+ a = _coerce_ordinal(labels_a)
+ b = _coerce_ordinal(labels_b)
+ rho, _ = spearmanr(a, b)
+ return float(rho) if rho == rho else float("nan") # NaN-safe
+ except ImportError: # pragma: no cover
+ logger.warning("scipy unavailable; spearman_rho returning NaN")
+ return float("nan")
diff --git a/potato/server_utils/iaa/ranking.py b/potato/server_utils/iaa/ranking.py
new file mode 100644
index 0000000000000000000000000000000000000000..b8914cf87c154a632d56164e27ecea6ca61f6e82
--- /dev/null
+++ b/potato/server_utils/iaa/ranking.py
@@ -0,0 +1,45 @@
+"""
+Ranking IAA metrics for schemas where each annotator produces an ordering
+(e.g., ranking, best-worst scaling, pairwise).
+"""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def kendall_tau(ranking_a: Sequence, ranking_b: Sequence) -> float:
+ """Kendall's tau-b between two rankings (lists of comparable items)."""
+ if len(ranking_a) != len(ranking_b):
+ raise ValueError("rankings must be the same length")
+ if len(ranking_a) < 2:
+ return float("nan")
+ try:
+ from scipy.stats import kendalltau
+ tau, _ = kendalltau(list(ranking_a), list(ranking_b))
+ return float(tau) if tau == tau else float("nan")
+ except ImportError: # pragma: no cover
+ logger.warning("scipy unavailable; kendall_tau returning NaN")
+ return float("nan")
+
+
+def spearman_footrule(ranking_a: Sequence, ranking_b: Sequence) -> float:
+ """
+ Normalized Spearman footrule distance. 0 = identical, 1 = maximally disagree.
+
+ Items are matched by identity; missing items get max-rank.
+ """
+ items = list({*ranking_a, *ranking_b})
+ if len(items) < 2:
+ return float("nan")
+ n = len(items)
+ rank_a = {item: i for i, item in enumerate(ranking_a)}
+ rank_b = {item: i for i, item in enumerate(ranking_b)}
+ total = sum(abs(rank_a.get(it, n) - rank_b.get(it, n)) for it in items)
+ # Worst-case footrule for n items is floor(n^2 / 2)
+ worst = (n * n) // 2 if n > 0 else 1
+ return total / worst if worst else float("nan")
diff --git a/potato/server_utils/iaa/span.py b/potato/server_utils/iaa/span.py
new file mode 100644
index 0000000000000000000000000000000000000000..e67b30fae26eaebd41fb8cafd67bb81b27649e2e
--- /dev/null
+++ b/potato/server_utils/iaa/span.py
@@ -0,0 +1,395 @@
+"""
+Span-specific IAA metrics.
+
+Span annotations are unusual: annotators can disagree on (a) **where** spans
+go (unitization / boundary detection) and (b) **what label** each span carries
+(categorization). Token-level kappa and exact-match F1 only capture part of
+this picture, which is why dedicated metrics exist:
+
+- **Token-level Cohen / Fleiss kappa** via BIO conversion โ simple,
+ intuitive, but penalizes near-misses harshly and ignores spans of differing
+ lengths.
+- **Span F1 (exact, partial)** โ IR-style; classic in NER literature
+ (MUC, CoNLL, SemEval).
+- **Krippendorff's alpha_U (unitizing alpha)** โ Krippendorff 2018; treats
+ each character/token as a unit and accounts for both boundary and
+ categorical disagreement.
+- **Gamma (Mathet et al. 2015)** โ state-of-the-art unified measure that
+ jointly handles unit alignment + categorization via the Hungarian algorithm.
+
+All inputs are ``SpanAnnotation``-like objects with ``start``, ``end``, and
+``name`` (label) attributes โ or plain dicts/tuples with the same fields.
+"""
+
+from __future__ import annotations
+
+from collections import Counter, defaultdict
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+import logging
+
+from potato.server_utils.iaa.nominal import cohen_kappa, fleiss_kappa
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Span representation helpers
+# ---------------------------------------------------------------------------
+
+def _span_tuple(span) -> Tuple[int, int, str]:
+ """Normalise a span object to (start, end, label)."""
+ if isinstance(span, dict):
+ return int(span["start"]), int(span["end"]), str(span.get("name") or span.get("label", ""))
+ if isinstance(span, tuple) and len(span) == 3:
+ return int(span[0]), int(span[1]), str(span[2])
+ return int(span.start), int(span.end), str(span.name)
+
+
+def _normalize(spans: Iterable) -> List[Tuple[int, int, str]]:
+ return [_span_tuple(s) for s in spans]
+
+
+# ---------------------------------------------------------------------------
+# Token-level kappa via BIO conversion
+# ---------------------------------------------------------------------------
+
+def spans_to_bio(spans: Iterable, length: int) -> List[str]:
+ """
+ Convert spans to BIO tags over a unit sequence of length ``length``.
+
+ ``length`` can be in characters or tokens depending on the unit; the
+ representation is the same. Overlapping spans are resolved with the rule
+ "longest span wins" โ sufficient for IAA where overlap is rare.
+ """
+ tags = ["O"] * length
+ span_list = sorted(_normalize(spans), key=lambda s: -(s[1] - s[0]))
+ for start, end, label in span_list:
+ start = max(0, start)
+ end = min(length, end)
+ if end <= start:
+ continue
+ if tags[start] != "O":
+ continue # respect longest-wins
+ tags[start] = f"B-{label}"
+ for i in range(start + 1, end):
+ if tags[i] != "O":
+ continue
+ tags[i] = f"I-{label}"
+ return tags
+
+
+def token_level_kappa(
+ spans_by_user: Dict[str, Iterable],
+ length: int,
+) -> float:
+ """
+ Cohen's / Fleiss' kappa over the BIO tag sequence.
+
+ For 2 annotators, returns Cohen's kappa; for >=3, returns Fleiss' kappa.
+ """
+ users = list(spans_by_user)
+ if len(users) < 2 or length <= 0:
+ return float("nan")
+ tag_seqs = {u: spans_to_bio(spans_by_user[u], length) for u in users}
+
+ if len(users) == 2:
+ return cohen_kappa(tag_seqs[users[0]], tag_seqs[users[1]])
+
+ counts_per_position = []
+ for i in range(length):
+ c: Counter = Counter()
+ for u in users:
+ c[tag_seqs[u][i]] += 1
+ counts_per_position.append(dict(c))
+ return fleiss_kappa(counts_per_position)
+
+
+# ---------------------------------------------------------------------------
+# Span F1 (exact and partial match)
+# ---------------------------------------------------------------------------
+
+def _overlap_len(a: Tuple[int, int, str], b: Tuple[int, int, str]) -> int:
+ return max(0, min(a[1], b[1]) - max(a[0], b[0]))
+
+
+def span_f1_exact(spans_a: Iterable, spans_b: Iterable) -> Tuple[float, float, float]:
+ """
+ Strict exact-match F1: (start, end, label) must match exactly.
+
+ Returns (precision, recall, F1) treating spans_b as gold.
+ """
+ a = set(_normalize(spans_a))
+ b = set(_normalize(spans_b))
+ if not a and not b:
+ return 1.0, 1.0, 1.0
+ tp = len(a & b)
+ p = tp / len(a) if a else 0.0
+ r = tp / len(b) if b else 0.0
+ f1 = 2 * p * r / (p + r) if (p + r) else 0.0
+ return p, r, f1
+
+
+def span_f1_partial(
+ spans_a: Iterable,
+ spans_b: Iterable,
+ label_must_match: bool = True,
+ threshold: float = 0.5,
+) -> Tuple[float, float, float]:
+ """
+ Partial-match F1: a span counts as TP if it overlaps a gold span by at
+ least ``threshold`` of either span's length (Dice-overlap convention).
+
+ label_must_match: when True (default) overlapping spans must share the
+ same label to count; False allows boundary-only agreement.
+ """
+ a = _normalize(spans_a)
+ b = _normalize(spans_b)
+ if not a and not b:
+ return 1.0, 1.0, 1.0
+ matched_b = set()
+ tp = 0
+ for sa in a:
+ for idx, sb in enumerate(b):
+ if idx in matched_b:
+ continue
+ if label_must_match and sa[2] != sb[2]:
+ continue
+ ov = _overlap_len(sa, sb)
+ if ov <= 0:
+ continue
+ la, lb = sa[1] - sa[0], sb[1] - sb[0]
+ if la <= 0 or lb <= 0:
+ continue
+ if (ov / la) >= threshold or (ov / lb) >= threshold:
+ tp += 1
+ matched_b.add(idx)
+ break
+ p = tp / len(a) if a else 0.0
+ r = tp / len(b) if b else 0.0
+ f1 = 2 * p * r / (p + r) if (p + r) else 0.0
+ return p, r, f1
+
+
+def pairwise_span_f1(
+ spans_by_user: Dict[str, Iterable],
+ partial: bool = False,
+ threshold: float = 0.5,
+) -> float:
+ """Mean pairwise span-F1 across users."""
+ users = list(spans_by_user)
+ if len(users) < 2:
+ return float("nan")
+ scores = []
+ for i in range(len(users)):
+ for j in range(i + 1, len(users)):
+ if partial:
+ _, _, f1 = span_f1_partial(
+ spans_by_user[users[i]], spans_by_user[users[j]], threshold=threshold,
+ )
+ else:
+ _, _, f1 = span_f1_exact(spans_by_user[users[i]], spans_by_user[users[j]])
+ scores.append(f1)
+ if not scores:
+ return float("nan")
+ return sum(scores) / len(scores)
+
+
+# ---------------------------------------------------------------------------
+# Krippendorff's alpha_U (unitizing alpha)
+# ---------------------------------------------------------------------------
+
+def krippendorff_alpha_u(
+ spans_by_user: Dict[str, Iterable],
+ length: int,
+) -> float:
+ """
+ Krippendorff's unitizing alpha for span annotation.
+
+ Implementation: assign each character/token position a categorical label
+ (one of the span labels or "O") per annotator, then compute Krippendorff's
+ alpha (nominal) over the (annotator, position) pairs. This is the
+ operational form recommended in Krippendorff (2018) when the unit is
+ fixed (per-character) rather than continuous.
+
+ For continuous-domain alpha_U (where annotators may disagree on the unit
+ boundary in a fundamentally continuous space such as audio), prefer gamma.
+ """
+ users = list(spans_by_user)
+ if len(users) < 2 or length <= 0:
+ return float("nan")
+
+ rows = []
+ for u in users:
+ tags = spans_to_bio(spans_by_user[u], length)
+ # Map BIO -> base label (strip B-/I- prefix) so boundary placement
+ # within a contiguous span doesn't count as disagreement.
+ for pos, tag in enumerate(tags):
+ label = "O" if tag == "O" else tag.split("-", 1)[1]
+ rows.append((u, pos, label))
+
+ from potato.server_utils.iaa.alpha import krippendorff_alpha
+ return krippendorff_alpha(rows, level="nominal")
+
+
+# ---------------------------------------------------------------------------
+# Gamma (Mathet et al. 2015)
+# ---------------------------------------------------------------------------
+
+def _positional_dissimilarity(
+ a: Tuple[int, int, str],
+ b: Tuple[int, int, str],
+ delta_empty: float,
+) -> float:
+ """Positional component of the Mathet dissimilarity (normalized)."""
+ if a is None or b is None:
+ return delta_empty
+ # Sum of |starts diff| + |ends diff|, normalized by total span lengths.
+ diff = abs(a[0] - b[0]) + abs(a[1] - b[1])
+ total = (a[1] - a[0]) + (b[1] - b[0])
+ if total <= 0:
+ return delta_empty
+ return diff / total
+
+
+def _categorical_dissimilarity(
+ a: Tuple[int, int, str],
+ b: Tuple[int, int, str],
+ delta_empty: float,
+) -> float:
+ if a is None or b is None:
+ return delta_empty
+ return 0.0 if a[2] == b[2] else 1.0
+
+
+def _pairwise_disorder(
+ spans_a: List[Tuple[int, int, str]],
+ spans_b: List[Tuple[int, int, str]],
+ alpha: float,
+ beta: float,
+ delta_empty: float,
+) -> float:
+ """
+ Optimal-alignment disorder between two annotators' span sets.
+
+ Uses the Hungarian algorithm (``scipy.optimize.linear_sum_assignment``)
+ with padded empty units so that |spans_a| != |spans_b| is handled.
+ """
+ try:
+ import numpy as np
+ from scipy.optimize import linear_sum_assignment
+ except ImportError: # pragma: no cover
+ logger.warning("scipy unavailable; gamma falling back to NaN")
+ return float("nan")
+
+ n = max(len(spans_a), len(spans_b))
+ if n == 0:
+ return 0.0
+ # Pad shorter side with None (= empty unit)
+ a_padded: List[Optional[Tuple[int, int, str]]] = list(spans_a) + [None] * (n - len(spans_a))
+ b_padded: List[Optional[Tuple[int, int, str]]] = list(spans_b) + [None] * (n - len(spans_b))
+
+ cost = np.zeros((n, n), dtype=float)
+ for i in range(n):
+ for j in range(n):
+ pos = _positional_dissimilarity(a_padded[i], b_padded[j], delta_empty)
+ cat = _categorical_dissimilarity(a_padded[i], b_padded[j], delta_empty)
+ cost[i, j] = alpha * pos + beta * cat
+
+ row_ind, col_ind = linear_sum_assignment(cost)
+ total = float(cost[row_ind, col_ind].sum())
+ return total / n
+
+
+def gamma(
+ spans_by_user: Dict[str, Iterable],
+ length: Optional[int] = None,
+ alpha: float = 1.0,
+ beta: float = 1.0,
+ n_samples: int = 30,
+ seed: int = 1234,
+) -> float:
+ """
+ Mathet et al. (2015) gamma agreement.
+
+ Args:
+ spans_by_user: annotator_id -> iterable of spans
+ length: total length of the unit space (characters or tokens). If
+ omitted, inferred from the maximum span end across annotators.
+ alpha: weight on positional dissimilarity.
+ beta: weight on categorical dissimilarity.
+ n_samples: number of random pairings used to estimate the
+ expected-by-chance disorder.
+ seed: RNG seed for reproducibility.
+
+ Returns:
+ gamma in [-1, 1] approximately, where 1 = perfect agreement, 0 =
+ chance-level. NaN if scipy is unavailable.
+
+ Notes:
+ This implementation is a faithful but simplified rendition: positional
+ dissimilarity is normalized by combined span length, and the chance
+ baseline is estimated by re-pairing spans across all annotators
+ ``n_samples`` times. Full pygamma-agreement uses a more sophisticated
+ baseline (continuum-of-shuffles); the simplification is sufficient for
+ the relative IAA comparisons that drive routing decisions.
+ """
+ import random as _random
+
+ users = list(spans_by_user)
+ if len(users) < 2:
+ return float("nan")
+ normed = {u: _normalize(spans_by_user[u]) for u in users}
+
+ # Empty-unit dissimilarity follows Mathet: a moderate constant ~ 1
+ delta_empty = 1.0
+
+ # Observed disorder: mean pairwise disorder across all annotator pairs
+ pair_disorders = []
+ for i in range(len(users)):
+ for j in range(i + 1, len(users)):
+ pair_disorders.append(
+ _pairwise_disorder(normed[users[i]], normed[users[j]], alpha, beta, delta_empty)
+ )
+ if not pair_disorders:
+ return float("nan")
+ if any(d != d for d in pair_disorders): # NaN -> bail
+ return float("nan")
+ observed = sum(pair_disorders) / len(pair_disorders)
+
+ # Expected-by-chance disorder via shuffled pairings
+ all_spans = [s for u in users for s in normed[u]]
+ if len(all_spans) < 2:
+ return 1.0 if observed == 0 else float("nan")
+
+ rng = _random.Random(seed)
+ chance_disorders = []
+ sizes = [len(normed[u]) for u in users]
+ for _ in range(n_samples):
+ shuffled = list(all_spans)
+ rng.shuffle(shuffled)
+ # Re-distribute back to annotators preserving original counts
+ idx = 0
+ shuffled_per_user = []
+ for sz in sizes:
+ shuffled_per_user.append(shuffled[idx:idx + sz])
+ idx += sz
+ sample_pair_disorders = []
+ for i in range(len(users)):
+ for j in range(i + 1, len(users)):
+ sample_pair_disorders.append(
+ _pairwise_disorder(
+ shuffled_per_user[i],
+ shuffled_per_user[j],
+ alpha, beta, delta_empty,
+ )
+ )
+ if sample_pair_disorders:
+ chance_disorders.append(sum(sample_pair_disorders) / len(sample_pair_disorders))
+
+ if not chance_disorders:
+ return float("nan")
+ expected = sum(chance_disorders) / len(chance_disorders)
+ if expected <= 0:
+ return 1.0 if observed == 0 else float("nan")
+ return 1.0 - (observed / expected)
diff --git a/potato/server_utils/instance_display.py b/potato/server_utils/instance_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..96f3cd7fe0456eec067e20797a1a0b5c189318c4
--- /dev/null
+++ b/potato/server_utils/instance_display.py
@@ -0,0 +1,418 @@
+"""
+Instance Display Renderer
+
+Provides the main InstanceDisplayRenderer class that handles rendering
+instance content for display, separate from annotation collection.
+
+This module enables the new `instance_display` configuration section
+that explicitly defines what content to show annotators.
+
+Usage:
+ from potato.server_utils.instance_display import InstanceDisplayRenderer
+
+ renderer = InstanceDisplayRenderer(config)
+ html = renderer.render(instance_data)
+ template_vars = renderer.get_template_variables(instance_data)
+"""
+
+import html as html_module
+import logging
+from typing import Dict, Any, List, Optional
+
+from .displays import display_registry
+
+logger = logging.getLogger(__name__)
+
+
+class InstanceDisplayError(Exception):
+ """Exception raised when instance display rendering fails."""
+ pass
+
+
+class InstanceDisplayRenderer:
+ """
+ Renders instance content for display based on configuration.
+
+ This class separates content display from annotation collection,
+ allowing any combination of display types with any annotation schemes.
+ """
+
+ def __init__(self, config: Dict[str, Any]):
+ """
+ Initialize the renderer.
+
+ Args:
+ config: The full configuration dictionary
+ """
+ self.config = config
+ self.display_config = config.get("instance_display", {})
+ self.fields = self.display_config.get("fields", [])
+ self.layout = self.display_config.get("layout", {})
+
+ # Extract span targets โ query the registry instead of a hardcoded list
+ self.span_targets = [
+ f["key"] for f in self.fields
+ if f.get("span_target") and display_registry.type_supports_span_target(f.get("type", ""))
+ ]
+
+ # Warn about span_target on unsupported types
+ for f in self.fields:
+ if f.get("span_target") and not display_registry.type_supports_span_target(f.get("type", "")):
+ logger.warning(
+ f"Field '{f.get('key')}' has span_target=true but display type "
+ f"'{f.get('type')}' does not support span annotation. "
+ f"Span annotation will not work on this field."
+ )
+
+ # Track if we have instance_display configured
+ self.has_instance_display = bool(self.fields)
+
+ logger.debug(
+ f"InstanceDisplayRenderer initialized: "
+ f"has_instance_display={self.has_instance_display}, "
+ f"span_targets={self.span_targets}"
+ )
+
+ def render(self, instance_data: Dict[str, Any]) -> str:
+ """
+ Render all display fields for an instance.
+
+ Args:
+ instance_data: The instance data dictionary
+
+ Returns:
+ HTML string containing all rendered display fields
+
+ Raises:
+ InstanceDisplayError: If a required field is missing from instance data
+ """
+ if not self.has_instance_display:
+ # No instance_display configured, return empty
+ # (legacy behavior will be handled by the template)
+ return ""
+
+ # Validate all required fields exist
+ self._validate_fields(instance_data)
+
+ # Get layout configuration
+ direction = self.layout.get("direction", "vertical")
+ gap = self.layout.get("gap", "20px")
+
+ # Build container classes and styles
+ container_classes = ["instance-display-container", f"layout-{direction}"]
+ container_style = f"gap: {gap};"
+
+ # Render each field
+ rendered_fields = []
+ for field in self.fields:
+ field_html = self._render_field(field, instance_data)
+ rendered_fields.append(field_html)
+
+ # Combine into container
+ fields_html = "\n".join(rendered_fields)
+
+ # Build data attributes for raw field access by annotation schemas
+ # Include all string/URL fields from instance data for source_field lookups
+ import json
+ raw_data = {}
+ for key, value in instance_data.items():
+ if isinstance(value, (str, int, float, bool)) or value is None:
+ raw_data[key] = value
+ raw_data_json = html_module.escape(json.dumps(raw_data))
+
+ return f'''
+
+ {fields_html}
+
+ '''
+
+ def _validate_fields(self, instance_data: Dict[str, Any]) -> None:
+ """
+ Validate that all configured fields exist in the instance data.
+
+ Fields whose display type is marked ``lazy_populated`` in the
+ display registry (``interactive_chat``, ``live_agent``,
+ ``live_coding_agent``) are exempt -- their data key is expected
+ to be written after initial render (by a live agent session).
+
+ Args:
+ instance_data: The instance data dictionary
+
+ Raises:
+ InstanceDisplayError: If any non-lazy field is missing
+ """
+ non_lazy = [
+ f for f in self.fields
+ if not display_registry.is_lazy_populated(f.get("type", ""))
+ ]
+ missing_non_lazy = [
+ f["key"] for f in non_lazy if f["key"] not in instance_data
+ ]
+ # Every non-lazy field missing is almost always a config/data
+ # key mismatch (e.g. fields reference task_description but the
+ # data uses task), not a transient lazy state -- make it loud so
+ # it isn't silently rendered as a blank page.
+ if non_lazy and len(missing_non_lazy) == len(non_lazy):
+ logger.error(
+ "instance_display: ALL %d non-lazy field(s) %s are absent "
+ "from the instance data (available keys: %s). This is "
+ "almost certainly a config/data key mismatch.",
+ len(non_lazy), missing_non_lazy,
+ list(instance_data.keys()),
+ )
+
+ for field in self.fields:
+ key = field["key"]
+ if key in instance_data:
+ continue
+ field_type = field.get("type", "")
+ if display_registry.is_lazy_populated(field_type):
+ logger.debug(
+ "Skipping validation for lazy-populated field '%s' (type=%s); "
+ "data is written after initial render.",
+ key, field_type,
+ )
+ continue
+ available = list(instance_data.keys())
+ raise InstanceDisplayError(
+ f"Display field '{key}' not found in instance data. "
+ f"Available fields: {available}"
+ )
+
+ def _render_field(self, field: Dict[str, Any], instance_data: Dict[str, Any]) -> str:
+ """
+ Render a single display field.
+
+ Args:
+ field: The field configuration
+ instance_data: The instance data dictionary
+
+ Returns:
+ HTML string for the field
+ """
+ key = field["key"]
+ field_type = field["type"]
+ data = instance_data.get(key)
+
+ # For format-based display types, process the file if data is a file path
+ format_display_types = ["pdf", "document", "spreadsheet", "code"]
+ if field_type in format_display_types and isinstance(data, str):
+ data = self._process_format_file(data, field_type, field)
+
+ try:
+ rendered = display_registry.render(field_type, field, data)
+
+ # Check if resizable is enabled (global setting or per-field override)
+ global_resizable = self.display_config.get("resizable", True)
+ field_resizable = field.get("display_options", {}).get("resizable", global_resizable)
+
+ # Wrap with resizable container if enabled
+ if field_resizable:
+ rendered = self._wrap_resizable(rendered, field)
+
+ return rendered
+ except ValueError as e:
+ logger.error(f"Error rendering field '{key}': {e}")
+ return f'Error rendering field "{key}": {e}
'
+
+ def _wrap_resizable(self, inner_html: str, field: Dict[str, Any]) -> str:
+ """
+ Wrap rendered content in a resizable container.
+
+ Args:
+ inner_html: The rendered field HTML
+ field: The field configuration
+
+ Returns:
+ HTML wrapped in resizable container
+ """
+ display_options = field.get("display_options", {})
+ max_height = display_options.get("max_height", 500)
+ min_height = display_options.get("min_height", 100)
+
+ style = f"max-height: {max_height}px; min-height: {min_height}px; position: relative;"
+
+ return f'''
+ {inner_html}
+
'''
+
+ def _process_format_file(
+ self,
+ file_path: str,
+ display_type: str,
+ field: Dict[str, Any]
+ ) -> Any:
+ """
+ Process a file using the format handler system.
+
+ If the data is a file path and a format handler is available,
+ extract the content and return FormatOutput data.
+
+ Args:
+ file_path: Path to the file to process
+ display_type: The display type (pdf, document, etc.)
+ field: The field configuration
+
+ Returns:
+ Either the original file_path (for client-side rendering like PDF.js)
+ or extracted content dict for server-side rendering
+ """
+ try:
+ from potato.format_handlers import format_handler_registry
+ except ImportError:
+ # Format handlers not available, return original data
+ logger.debug("Format handlers not available, using raw file path")
+ return file_path
+
+ # Check if the file path should be processed
+ # For PDFs, we typically use client-side rendering with PDF.js
+ # unless explicitly configured for server-side extraction
+ display_options = field.get("display_options", {})
+
+ if display_type == "pdf":
+ # By default, PDFs use client-side rendering (return path as-is)
+ # If server_extract is set, use the format handler
+ if not display_options.get("server_extract", False):
+ return file_path
+
+ # Check if format handler can handle this file
+ if not format_handler_registry.can_handle(file_path):
+ logger.debug(f"No format handler for {file_path}, using raw data")
+ return file_path
+
+ try:
+ # Extract content using format handler
+ extraction_options = display_options.get("extraction_options", {})
+ output = format_handler_registry.extract(file_path, options=extraction_options)
+
+ # Return as dict for the display renderer
+ return {
+ "text": output.text,
+ "rendered_html": output.rendered_html,
+ "coordinate_map": output.coordinate_map,
+ "metadata": output.metadata,
+ "format_name": output.format_name,
+ "source_path": output.source_path,
+ }
+ except Exception as e:
+ logger.warning(f"Format handler extraction failed for {file_path}: {e}")
+ return file_path
+
+ def get_template_variables(self, instance_data: Dict[str, Any]) -> Dict[str, Any]:
+ """
+ Get template variables for Jinja access.
+
+ Returns a dictionary with:
+ - display_html: The complete rendered display HTML
+ - display_fields: Dictionary of field key -> rendered HTML
+ - display_raw: Dictionary of field key -> raw data value
+ - span_targets: List of field keys that are span targets
+ - multi_span_mode: Boolean indicating if multiple span targets exist
+ - has_instance_display: Boolean indicating if instance_display is configured
+
+ Args:
+ instance_data: The instance data dictionary
+
+ Returns:
+ Dictionary of template variables
+ """
+ result = {
+ "display_html": "",
+ "display_fields": {},
+ "display_raw": {},
+ "span_targets": self.span_targets,
+ "multi_span_mode": len(self.span_targets) > 1,
+ "has_instance_display": self.has_instance_display,
+ }
+
+ if not self.has_instance_display:
+ return result
+
+ # Validate fields. A missing field here is a real config problem
+ # (lazy-populated types like interactive_chat are already filtered
+ # out by _validate_fields), but the renderer surfaces it inline
+ # via ``display_error`` so the page still loads -- WARN is the
+ # right severity, not ERROR.
+ try:
+ self._validate_fields(instance_data)
+ except InstanceDisplayError as e:
+ logger.warning(f"Field validation failed: {e}")
+ result["display_error"] = str(e)
+ return result
+
+ # Render complete display
+ result["display_html"] = self.render(instance_data)
+
+ # Render individual fields and collect raw data
+ for field in self.fields:
+ key = field["key"]
+ field_type = field["type"]
+ data = instance_data.get(key)
+
+ result["display_raw"][key] = data
+
+ try:
+ result["display_fields"][key] = display_registry.render(field_type, field, data)
+ except ValueError as e:
+ logger.error(f"Error rendering field '{key}': {e}")
+ result["display_fields"][key] = f'Error: {e}
'
+
+ return result
+
+ def get_span_target_fields(self) -> List[Dict[str, Any]]:
+ """
+ Get the list of fields configured as span targets.
+
+ Returns:
+ List of field configuration dictionaries for span targets
+ """
+ return [f for f in self.fields if f.get("span_target")]
+
+ def get_primary_text_field(self) -> Optional[str]:
+ """
+ Get the primary text field key for legacy compatibility.
+
+ Returns the first span target if any, otherwise the first text field,
+ otherwise None.
+
+ Returns:
+ Field key string or None
+ """
+ # First, check span targets
+ if self.span_targets:
+ return self.span_targets[0]
+
+ # Then look for any text field
+ for field in self.fields:
+ if field.get("type") == "text":
+ return field["key"]
+
+ return None
+
+ def should_use_legacy_display(self) -> bool:
+ """
+ Check if legacy display mode should be used.
+
+ Returns True if no instance_display is configured, meaning
+ the template should fall back to displaying text_key.
+
+ Returns:
+ True if legacy mode should be used
+ """
+ return not self.has_instance_display
+
+
+def get_instance_display_renderer(config: Dict[str, Any]) -> InstanceDisplayRenderer:
+ """
+ Get or create an InstanceDisplayRenderer for the given config.
+
+ This is a convenience function that creates a renderer.
+ In the future, this could cache renderers per config hash.
+
+ Args:
+ config: The configuration dictionary
+
+ Returns:
+ InstanceDisplayRenderer instance
+ """
+ return InstanceDisplayRenderer(config)
diff --git a/potato/server_utils/json.py b/potato/server_utils/json.py
new file mode 100644
index 0000000000000000000000000000000000000000..325a0e53a1b10c044211150ea2c02cddf40fdb53
--- /dev/null
+++ b/potato/server_utils/json.py
@@ -0,0 +1,19 @@
+"""
+filename: json.py
+date: 10/16/2024
+author: Tristan Hilbert (aka TFlexSoom)
+desc: Json encoding utilities for the potato tool
+"""
+
+import dataclasses
+import json
+from typing import Any
+
+class EnhancedJSONEncoder(json.JSONEncoder):
+ def default(self, o):
+ if dataclasses.is_dataclass(o):
+ return dataclasses.asdict(o)
+ return super().default(o)
+
+def easy_json(obj: Any):
+ return json.dumps(obj, cls=EnhancedJSONEncoder)
\ No newline at end of file
diff --git a/potato/server_utils/judge_alignment.py b/potato/server_utils/judge_alignment.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca21788fa646e663498af2791679354d7f7de59b
--- /dev/null
+++ b/potato/server_utils/judge_alignment.py
@@ -0,0 +1,349 @@
+"""
+Judge โ human alignment: persistence + agreement computation.
+
+Pairs each LLM-judge verdict (from ``potato/ai/judge.py``) with the human gold
+label for the same instance/schema and computes Cohen's ฮบ, a confusion matrix,
+agreement rate, and the list of disagreements. Judge predictions are persisted
+per *prompt version* so the admin report can track ฮบ as the rubric is calibrated.
+
+Layout under ``{task_dir}/judge_alignment/``:
+ predictions.json -> {prompt_version: {"::": JudgePrediction}}
+ comparisons.json -> [{instance_id, schema, human_label, judge_label, agrees, prompt_version}]
+ (running log written by the inline capture path)
+
+The ฮบ computation reuses ``potato/agreement.py`` (judge vs. human gold as two
+"annotators"). The pure ``compute_alignment_from_pairs`` is the unit-testable core.
+"""
+
+import json
+import logging
+import os
+from collections import Counter, defaultdict
+from typing import Any, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+# ----- paths / persistence ----------------------------------------------
+
+def _dir(config: Dict[str, Any]) -> str:
+ base = config.get("output_annotation_dir") or config.get("task_dir") or "."
+ return os.path.join(base, "judge_alignment")
+
+
+def _load_json(path: str, default):
+ try:
+ with open(path, "r", encoding="utf-8") as f:
+ return json.load(f)
+ except (FileNotFoundError, ValueError):
+ return default
+
+
+def _save_json(path: str, data) -> None:
+ os.makedirs(os.path.dirname(path), exist_ok=True)
+ with open(path, "w", encoding="utf-8") as f:
+ json.dump(data, f, indent=2, ensure_ascii=False)
+
+
+def predictions_path(config: Dict[str, Any]) -> str:
+ return os.path.join(_dir(config), "predictions.json")
+
+
+def comparisons_path(config: Dict[str, Any]) -> str:
+ return os.path.join(_dir(config), "comparisons.json")
+
+
+def load_predictions(config: Dict[str, Any]) -> Dict[str, Dict[str, dict]]:
+ return _load_json(predictions_path(config), {})
+
+
+def save_prediction(config: Dict[str, Any], pred) -> None:
+ """Persist one JudgePrediction (keyed by prompt_version โ instance::schema)."""
+ data = load_predictions(config)
+ version = pred.prompt_version or "default"
+ data.setdefault(version, {})[f"{pred.instance_id}::{pred.schema_name}"] = pred.to_dict()
+ _save_json(predictions_path(config), data)
+
+
+def latest_prompt_version(config: Dict[str, Any]) -> Optional[str]:
+ data = load_predictions(config)
+ if not data:
+ return None
+ # Most-populated version is the "current" working set.
+ return max(data.keys(), key=lambda v: len(data[v]))
+
+
+def record_comparison(config: Dict[str, Any], instance_id: str, schema: str,
+ human_label: Any, judge_label: Any, prompt_version: str) -> None:
+ """Append a humanโjudge comparison to the running log (inline capture)."""
+ log = _load_json(comparisons_path(config), [])
+ log.append({
+ "instance_id": instance_id,
+ "schema": schema,
+ "human_label": str(human_label),
+ "judge_label": str(judge_label),
+ "agrees": str(human_label) == str(judge_label),
+ "prompt_version": prompt_version,
+ })
+ _save_json(comparisons_path(config), log)
+
+
+def running_agreement(config: Dict[str, Any], schema: Optional[str] = None) -> Dict[str, Any]:
+ """Quick running agreement from the comparison log (for the inline badge)."""
+ log = _load_json(comparisons_path(config), [])
+ if schema:
+ log = [c for c in log if c.get("schema") == schema]
+ n = len(log)
+ agree = sum(1 for c in log if c.get("agrees"))
+ pairs = {s: [] for s in {c["schema"] for c in log}}
+ for c in log:
+ pairs[c["schema"]].append((c["instance_id"], c["human_label"], c["judge_label"], None, ""))
+ kappa = None
+ if schema and pairs.get(schema):
+ res = compute_alignment_from_pairs({schema: pairs[schema]}).get(schema, {})
+ kappa = res.get("kappa")
+ return {"n": n, "agreements": agree,
+ "agreement_rate": round(agree / n, 3) if n else 0.0, "kappa": kappa}
+
+
+# ----- human label extraction --------------------------------------------
+
+def human_label_for(instance_id: str, schema_name: str, username: str) -> Optional[str]:
+ """The single categorical label a user assigned for a schema, or None."""
+ from potato.flask_server import get_annotations_for_user_on
+ anns = get_annotations_for_user_on(username, instance_id) or {}
+ chosen = anns.get(schema_name)
+ if not chosen:
+ return None
+ # Single-choice: the (first) selected label name.
+ keys = [k for k in chosen.keys()]
+ return keys[0] if keys else None
+
+
+def majority_human_label(instance_id: str, schema_name: str, users: List[str]) -> Optional[str]:
+ votes = []
+ for u in users:
+ lab = human_label_for(instance_id, schema_name, u)
+ if lab is not None:
+ votes.append(lab)
+ if not votes:
+ return None
+ return Counter(votes).most_common(1)[0][0]
+
+
+# ----- agreement computation (pure core) ----------------------------------
+
+def compute_alignment_from_pairs(
+ pairs_by_schema: Dict[str, List[Tuple[str, Any, Any, Optional[float], str]]],
+) -> Dict[str, Any]:
+ """Compute per-schema judgeโhuman alignment from resolved pairs.
+
+ pairs_by_schema: {schema: [(instance_id, human_label, judge_label,
+ judge_confidence|None, reasoning), ...]}
+ Returns {schema: {kappa, interpretation, agreement_rate, n, confusion,
+ disagreements[]}}.
+ """
+ import pandas as pd
+ from potato.agreement import cohen_kappa_pairwise, interpret_kappa
+
+ out: Dict[str, Any] = {}
+ for schema, pairs in pairs_by_schema.items():
+ pairs = [p for p in pairs if p[1] is not None and p[2] is not None]
+ n = len(pairs)
+ if n == 0:
+ out[schema] = {"kappa": None, "interpretation": "no overlap",
+ "agreement_rate": 0.0, "n": 0, "confusion": {},
+ "disagreements": []}
+ continue
+
+ agree = sum(1 for _, h, j, *_ in pairs if str(h) == str(j))
+ confusion: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int))
+ disagreements = []
+ rows = []
+ for inst, h, j, conf, reason in pairs:
+ confusion[str(h)][str(j)] += 1
+ rows.append({"unit": inst, "annotator": "human", "annotation": str(h)})
+ rows.append({"unit": inst, "annotator": "judge", "annotation": str(j)})
+ if str(h) != str(j):
+ disagreements.append({
+ "instance_id": inst, "human_label": str(h), "judge_label": str(j),
+ "judge_confidence": conf, "reasoning": reason,
+ })
+
+ kappa = None
+ interp = "n/a"
+ try:
+ res = cohen_kappa_pairwise(pd.DataFrame(rows))
+ kappa = res.get("mean_kappa")
+ if kappa is not None:
+ interp = interpret_kappa(kappa)
+ except Exception as e:
+ logger.warning(f"Judge alignment: kappa failed for {schema}: {e}")
+
+ out[schema] = {
+ "kappa": round(kappa, 3) if isinstance(kappa, (int, float)) else None,
+ "interpretation": interp,
+ "agreement_rate": round(agree / n, 3),
+ "n": n,
+ "confusion": {h: dict(js) for h, js in confusion.items()},
+ "disagreements": disagreements,
+ }
+ return out
+
+
+# ----- gathering from persisted predictions + live human labels -----------
+
+def judge_scoped_schemas(config: Dict[str, Any]) -> List[dict]:
+ """Annotation schemes the judge should evaluate (categorical only).
+
+ Honors ``judge_alignment.schemas`` allow-list if present; otherwise all
+ radio/select/likert schemes.
+ """
+ schemes = config.get("annotation_schemes", []) or []
+ allow = set((config.get("judge_alignment", {}) or {}).get("schemas", {}).keys())
+ cats = {"radio", "select", "likert"}
+ out = []
+ for s in schemes:
+ if s.get("annotation_type") not in cats:
+ continue
+ if allow and s.get("name") not in allow:
+ continue
+ out.append(s)
+ return out
+
+
+def gather_pairs(config: Dict[str, Any], users: List[str], schema_names: List[str],
+ prompt_version: Optional[str]) -> Dict[str, List[Tuple]]:
+ """Build (instance, human_gold, judge_label, conf, reasoning) pairs."""
+ preds = load_predictions(config)
+ version = prompt_version or latest_prompt_version(config)
+ version_preds = preds.get(version, {}) if version else {}
+
+ pairs_by_schema: Dict[str, List[Tuple]] = {s: [] for s in schema_names}
+ for key, pred in version_preds.items():
+ instance_id, _, schema = key.partition("::")
+ if schema not in pairs_by_schema:
+ continue
+ gold = majority_human_label(instance_id, schema, users)
+ if gold is None:
+ continue
+ pairs_by_schema[schema].append((
+ instance_id, gold, pred.get("predicted_label"),
+ pred.get("confidence"), pred.get("reasoning", ""),
+ ))
+ return pairs_by_schema
+
+
+def annotated_instance_ids(users: List[str], schema_name: str) -> List[str]:
+ """Instance ids that at least one user has labeled for this schema."""
+ from potato.flask_server import get_user_state
+ ids = set()
+ for u in users:
+ st = get_user_state(u)
+ if not st:
+ continue
+ for iid in st.get_annotated_instance_ids():
+ if human_label_for(iid, schema_name, u) is not None:
+ ids.add(iid)
+ return sorted(ids)
+
+
+def run_judge_batch(config: Dict[str, Any], users: List[str],
+ rubric_overrides: Optional[Dict[str, str]] = None,
+ max_per_schema: Optional[int] = None) -> Dict[str, Any]:
+ """Run the judge over human-annotated instances and persist predictions.
+
+ rubric_overrides: {schema_name: rubric} to calibrate + create a new prompt
+ version. Few-shot examples (when enabled) are drawn from high-agreement
+ human labels, excluding the instance being judged.
+ """
+ from potato.ai.judge import JudgeService, compute_prompt_version
+ from potato.item_state_management import get_item_state_manager
+
+ # Apply rubric overrides into a working config copy.
+ cfg = dict(config)
+ ja = dict(cfg.get("judge_alignment", {}) or {})
+ if rubric_overrides:
+ schemas_cfg = dict(ja.get("schemas", {}) or {})
+ for name, rubric in rubric_overrides.items():
+ sc = dict(schemas_cfg.get(name, {}) or {})
+ sc["rubric"] = rubric
+ schemas_cfg[name] = sc
+ ja["schemas"] = schemas_cfg
+ cfg["judge_alignment"] = ja
+
+ service = JudgeService(cfg)
+ ism = get_item_state_manager()
+ few_shot_cfg = (ja.get("few_shot") or {})
+ use_few_shot = bool(few_shot_cfg.get("enabled", False))
+
+ n_judged, n_failed, version_seen = 0, 0, None
+ for schema in judge_scoped_schemas(cfg):
+ schema_name = schema.get("name")
+ ids = annotated_instance_ids(users, schema_name)
+ if max_per_schema:
+ ids = ids[:max_per_schema]
+ examples = _few_shot_examples(schema_name, use_few_shot, few_shot_cfg)
+ for iid in ids:
+ try:
+ item = ism.get_item(iid)
+ text = item.get_text() if item else ""
+ except Exception:
+ text = ""
+ shots = [e for e in examples if e.get("id") != iid] or None
+ pred = service.judge_instance(iid, schema, text, few_shot_examples=shots)
+ if pred is None:
+ n_failed += 1
+ continue
+ save_prediction(cfg, pred)
+ version_seen = pred.prompt_version
+ n_judged += 1
+
+ return {"judged": n_judged, "failed": n_failed, "prompt_version": version_seen}
+
+
+def _few_shot_examples(schema_name: str, enabled: bool, cfg: Dict[str, Any]) -> List[dict]:
+ """Gold few-shot examples from high-agreement human labels (or [])."""
+ if not enabled:
+ return []
+ try:
+ from potato.ai.icl_labeler import get_icl_labeler
+ labeler = get_icl_labeler()
+ if labeler is None:
+ return []
+ by_schema = labeler.refresh_high_confidence_examples()
+ examples = by_schema.get(schema_name, [])[: int(cfg.get("max_examples", 5))]
+ return [{"id": getattr(e, "instance_id", ""),
+ "text": getattr(e, "instance_text", getattr(e, "text", "")),
+ "label": getattr(e, "label", getattr(e, "agreed_label", ""))}
+ for e in examples]
+ except Exception as e:
+ logger.warning(f"Judge few-shot example gathering failed: {e}")
+ return []
+
+
+def compute_judge_alignment(config: Dict[str, Any], users: List[str],
+ prompt_version: Optional[str] = None) -> Dict[str, Any]:
+ """Full report: per-schema alignment for a prompt version + version list."""
+ schemas = [s.get("name") for s in judge_scoped_schemas(config)]
+ version = prompt_version or latest_prompt_version(config)
+ pairs = gather_pairs(config, users, schemas, version)
+ per_schema = compute_alignment_from_pairs(pairs)
+
+ preds = load_predictions(config)
+ versions = []
+ for v in preds.keys():
+ v_pairs = gather_pairs(config, users, schemas, v)
+ v_report = compute_alignment_from_pairs(v_pairs)
+ kappas = [r["kappa"] for r in v_report.values() if r.get("kappa") is not None]
+ versions.append({
+ "prompt_version": v,
+ "n_predictions": len(preds[v]),
+ "mean_kappa": round(sum(kappas) / len(kappas), 3) if kappas else None,
+ })
+
+ return {
+ "prompt_version": version,
+ "per_schema": per_schema,
+ "prompt_versions": sorted(versions, key=lambda x: x["prompt_version"]),
+ }
diff --git a/potato/server_utils/mturk_apis.py b/potato/server_utils/mturk_apis.py
new file mode 100644
index 0000000000000000000000000000000000000000..500629bc31afd5b11a6bf3b8aa879b3481f49ea5
--- /dev/null
+++ b/potato/server_utils/mturk_apis.py
@@ -0,0 +1,519 @@
+"""
+Utility functions for handling Amazon Mechanical Turk APIs.
+
+This module provides wrapper classes for interacting with the MTurk API
+via boto3 to manage HITs and assignments.
+
+The MTurk API requires AWS credentials and uses different endpoints for
+sandbox (testing) and production environments.
+
+Key API Operations:
+- GET account balance
+- LIST HITs
+- GET HIT details
+- LIST assignments for a HIT
+- APPROVE/REJECT assignments
+
+For more information, see:
+https://docs.aws.amazon.com/mturk/index.html
+"""
+import os
+import json
+import logging
+from collections import OrderedDict, defaultdict
+
+logger = logging.getLogger(__name__)
+
+# Optional boto3 import - only required if using MTurk API features
+try:
+ import boto3
+ from botocore.exceptions import ClientError
+ BOTO3_AVAILABLE = True
+except ImportError:
+ BOTO3_AVAILABLE = False
+ logger.debug("boto3 not available - MTurk API features will be disabled")
+
+
+class MTurkBase:
+ """
+ Base class for MTurk API operations.
+
+ Provides low-level API access to Amazon Mechanical Turk's REST endpoints.
+ All methods use AWS credential-based authentication via boto3.
+
+ Attributes:
+ client: boto3 MTurk client
+ sandbox (bool): Whether using sandbox environment
+ """
+
+ # MTurk API endpoints
+ SANDBOX_ENDPOINT = 'https://mturk-requester-sandbox.us-east-1.amazonaws.com'
+ PRODUCTION_ENDPOINT = 'https://mturk-requester.us-east-1.amazonaws.com'
+
+ def __init__(self, aws_access_key_id=None, aws_secret_access_key=None, sandbox=True):
+ """
+ Initialize the MTurk API client with AWS credentials.
+
+ Args:
+ aws_access_key_id (str, optional): AWS access key ID.
+ If not provided, uses environment variable or ~/.aws/credentials
+ aws_secret_access_key (str, optional): AWS secret access key.
+ If not provided, uses environment variable or ~/.aws/credentials
+ sandbox (bool): If True, use sandbox environment for testing.
+ Default is True (sandbox mode).
+ """
+ if not BOTO3_AVAILABLE:
+ raise ImportError(
+ "boto3 is required for MTurk API features. "
+ "Install it with: pip install boto3"
+ )
+
+ self.sandbox = sandbox
+ endpoint_url = self.SANDBOX_ENDPOINT if sandbox else self.PRODUCTION_ENDPOINT
+
+ # Build client kwargs
+ client_kwargs = {
+ 'service_name': 'mturk',
+ 'region_name': 'us-east-1',
+ 'endpoint_url': endpoint_url
+ }
+
+ # Add credentials if provided explicitly
+ if aws_access_key_id:
+ client_kwargs['aws_access_key_id'] = aws_access_key_id
+ if aws_secret_access_key:
+ client_kwargs['aws_secret_access_key'] = aws_secret_access_key
+
+ self.client = boto3.client(**client_kwargs)
+
+ env_type = "SANDBOX" if sandbox else "PRODUCTION"
+ logger.info(f"MTurk API client initialized ({env_type})")
+
+ def get_account_balance(self):
+ """
+ Get the MTurk account balance.
+
+ Returns:
+ str: Available balance (e.g., "10000.00" for sandbox, actual balance for production)
+
+ Raises:
+ ClientError: If the API request fails
+ """
+ try:
+ response = self.client.get_account_balance()
+ balance = response['AvailableBalance']
+ logger.debug(f"MTurk account balance: {balance}")
+ return balance
+ except ClientError as e:
+ logger.error(f"Failed to get account balance: {e}")
+ raise
+
+ def list_hits(self, max_results=100):
+ """
+ List all HITs in the account.
+
+ Args:
+ max_results (int): Maximum number of HITs to return (default: 100)
+
+ Returns:
+ list: List of HIT dictionaries containing:
+ - HITId: HIT identifier
+ - Title: HIT title
+ - HITStatus: Current status
+ - MaxAssignments: Total assignments available
+ - NumberOfAssignmentsCompleted: Completed count
+ - NumberOfAssignmentsPending: Pending count
+ - NumberOfAssignmentsAvailable: Available count
+ """
+ try:
+ response = self.client.list_hits(MaxResults=max_results)
+ hits = response.get('HITs', [])
+ logger.debug(f"Listed {len(hits)} HITs")
+ return hits
+ except ClientError as e:
+ logger.error(f"Failed to list HITs: {e}")
+ raise
+
+ def get_hit(self, hit_id):
+ """
+ Get detailed information about a specific HIT.
+
+ Args:
+ hit_id (str): The HIT identifier
+
+ Returns:
+ dict: Complete HIT information including:
+ - HITId, HITTypeId
+ - Title, Description, Question
+ - HITStatus: ASSIGNABLE, UNASSIGNABLE, REVIEWABLE, etc.
+ - MaxAssignments, Reward, AssignmentDurationInSeconds
+ - Creation/Expiration times
+ """
+ try:
+ response = self.client.get_hit(HITId=hit_id)
+ hit = response.get('HIT', {})
+ logger.debug(f"Retrieved HIT {hit_id}: status={hit.get('HITStatus')}")
+ return hit
+ except ClientError as e:
+ logger.error(f"Failed to get HIT {hit_id}: {e}")
+ raise
+
+ def list_assignments_for_hit(self, hit_id, assignment_statuses=None, max_results=100):
+ """
+ Get assignments for a specific HIT.
+
+ Args:
+ hit_id (str): The HIT identifier
+ assignment_statuses (list, optional): Filter by status.
+ Valid values: 'Submitted', 'Approved', 'Rejected'
+ max_results (int): Maximum number of assignments to return
+
+ Returns:
+ list: List of assignment dictionaries containing:
+ - AssignmentId: Assignment identifier
+ - WorkerId: Worker's MTurk ID
+ - HITId: Associated HIT ID
+ - AssignmentStatus: Submitted, Approved, or Rejected
+ - AcceptTime, SubmitTime
+ - Answer: Worker's submitted answer (XML format)
+ """
+ try:
+ params = {
+ 'HITId': hit_id,
+ 'MaxResults': max_results
+ }
+ if assignment_statuses:
+ params['AssignmentStatuses'] = assignment_statuses
+
+ response = self.client.list_assignments_for_hit(**params)
+ assignments = response.get('Assignments', [])
+ logger.debug(f"Listed {len(assignments)} assignments for HIT {hit_id}")
+ return assignments
+ except ClientError as e:
+ logger.error(f"Failed to list assignments for HIT {hit_id}: {e}")
+ raise
+
+ def approve_assignment(self, assignment_id, requester_feedback=""):
+ """
+ Approve an assignment.
+
+ Args:
+ assignment_id (str): The assignment identifier
+ requester_feedback (str, optional): Feedback message to worker
+
+ Returns:
+ dict: Empty dict on success
+
+ Raises:
+ ClientError: If approval fails (e.g., already approved/rejected)
+ """
+ try:
+ response = self.client.approve_assignment(
+ AssignmentId=assignment_id,
+ RequesterFeedback=requester_feedback,
+ OverrideRejection=False
+ )
+ logger.info(f"Approved assignment {assignment_id}")
+ return response
+ except ClientError as e:
+ logger.error(f"Failed to approve assignment {assignment_id}: {e}")
+ raise
+
+ def reject_assignment(self, assignment_id, requester_feedback):
+ """
+ Reject an assignment.
+
+ Args:
+ assignment_id (str): The assignment identifier
+ requester_feedback (str): Required feedback explaining rejection
+
+ Returns:
+ dict: Empty dict on success
+
+ Raises:
+ ClientError: If rejection fails
+ """
+ try:
+ response = self.client.reject_assignment(
+ AssignmentId=assignment_id,
+ RequesterFeedback=requester_feedback
+ )
+ logger.info(f"Rejected assignment {assignment_id}")
+ return response
+ except ClientError as e:
+ logger.error(f"Failed to reject assignment {assignment_id}: {e}")
+ raise
+
+ def get_assignment(self, assignment_id):
+ """
+ Get detailed information about a specific assignment.
+
+ Args:
+ assignment_id (str): The assignment identifier
+
+ Returns:
+ dict: Complete assignment information
+ """
+ try:
+ response = self.client.get_assignment(AssignmentId=assignment_id)
+ assignment = response.get('Assignment', {})
+ logger.debug(f"Retrieved assignment {assignment_id}")
+ return assignment
+ except ClientError as e:
+ logger.error(f"Failed to get assignment {assignment_id}: {e}")
+ raise
+
+
+class MTurkHIT(MTurkBase):
+ """
+ High-level HIT management class for MTurk.
+
+ Extends MTurkBase to provide HIT-specific functionality including:
+ - Assignment tracking and status management
+ - Batch approval/rejection
+ - Worker tracking
+
+ Attributes:
+ hit_id (str): MTurk HIT identifier
+ hit_info (dict): Cached HIT information
+ assignments (OrderedDict): Mapping of assignment IDs to assignment data
+ worker_status (defaultdict): Mapping of status to worker IDs
+ """
+
+ def __init__(self, aws_access_key_id=None, aws_secret_access_key=None,
+ hit_id=None, sandbox=True):
+ """
+ Initialize HIT management with configuration.
+
+ Args:
+ aws_access_key_id (str, optional): AWS access key ID
+ aws_secret_access_key (str, optional): AWS secret access key
+ hit_id (str, optional): MTurk HIT identifier to manage
+ sandbox (bool): Whether to use sandbox environment (default: True)
+ """
+ super().__init__(aws_access_key_id, aws_secret_access_key, sandbox)
+
+ self.hit_id = hit_id
+ self.hit_info = None
+ self.assignments = OrderedDict()
+ self.worker_status = defaultdict(set)
+
+ if hit_id:
+ try:
+ self.hit_info = self.get_hit(hit_id)
+ logger.info(f"Initialized MTurkHIT for HIT {hit_id}")
+ except ClientError as e:
+ logger.warning(f"Could not fetch HIT info: {e}")
+
+ def get_basic_hit_info(self):
+ """
+ Extract basic HIT information for display or logging.
+
+ Returns:
+ dict: Basic HIT information including:
+ - HITId, Title, HITStatus
+ - MaxAssignments, Reward
+ - NumberOfAssignmentsCompleted/Pending/Available
+ """
+ if not self.hit_info:
+ return {}
+
+ keys = [
+ 'HITId', 'Title', 'HITStatus', 'MaxAssignments', 'Reward',
+ 'NumberOfAssignmentsCompleted', 'NumberOfAssignmentsPending',
+ 'NumberOfAssignmentsAvailable'
+ ]
+ return {key: self.hit_info.get(key) for key in keys if key in self.hit_info}
+
+ def refresh_assignments(self):
+ """
+ Refresh local assignment data from MTurk API.
+
+ Fetches current assignment data from the API and updates local state:
+ - Updates assignments mapping
+ - Rebuilds worker status dictionary grouped by assignment status
+ """
+ if not self.hit_id:
+ logger.warning("No HIT ID set, cannot refresh assignments")
+ return
+
+ try:
+ # Get all assignments regardless of status
+ assignments = self.list_assignments_for_hit(self.hit_id)
+
+ self.assignments = OrderedDict()
+ self.worker_status = defaultdict(set)
+
+ for assignment in assignments:
+ assignment_id = assignment['AssignmentId']
+ self.assignments[assignment_id] = assignment
+ worker_id = assignment['WorkerId']
+ status = assignment['AssignmentStatus']
+ self.worker_status[status].add(worker_id)
+
+ logger.info(f"Refreshed {len(assignments)} assignments for HIT {self.hit_id}")
+ except ClientError as e:
+ logger.error(f"Failed to refresh assignments: {e}")
+
+ def get_pending_assignments(self):
+ """
+ Get assignments pending review (submitted but not approved/rejected).
+
+ Returns:
+ list: List of assignment dictionaries with 'Submitted' status
+ """
+ return self.list_assignments_for_hit(
+ self.hit_id,
+ assignment_statuses=['Submitted']
+ )
+
+ def get_approved_assignments(self):
+ """
+ Get approved assignments.
+
+ Returns:
+ list: List of assignment dictionaries with 'Approved' status
+ """
+ return self.list_assignments_for_hit(
+ self.hit_id,
+ assignment_statuses=['Approved']
+ )
+
+ def get_rejected_assignments(self):
+ """
+ Get rejected assignments.
+
+ Returns:
+ list: List of assignment dictionaries with 'Rejected' status
+ """
+ return self.list_assignments_for_hit(
+ self.hit_id,
+ assignment_statuses=['Rejected']
+ )
+
+ def auto_approve_all(self, feedback="Thank you for your work!"):
+ """
+ Approve all pending (submitted) assignments.
+
+ Args:
+ feedback (str): Feedback message to send to workers
+
+ Returns:
+ int: Number of assignments approved
+ """
+ pending = self.get_pending_assignments()
+ approved_count = 0
+
+ for assignment in pending:
+ try:
+ self.approve_assignment(assignment['AssignmentId'], feedback)
+ approved_count += 1
+ except ClientError as e:
+ logger.warning(f"Failed to approve {assignment['AssignmentId']}: {e}")
+
+ logger.info(f"Auto-approved {approved_count} assignments")
+ return approved_count
+
+ def get_worker_ids_by_status(self, status):
+ """
+ Get worker IDs filtered by assignment status.
+
+ Args:
+ status (str): Assignment status ('Submitted', 'Approved', 'Rejected')
+
+ Returns:
+ set: Set of worker IDs with the specified status
+ """
+ self.refresh_assignments()
+ return self.worker_status.get(status, set())
+
+ def get_completed_workers(self):
+ """
+ Get list of workers who have completed (submitted or approved) assignments.
+
+ Returns:
+ set: Set of worker IDs
+ """
+ self.refresh_assignments()
+ return (self.worker_status.get('Submitted', set()) |
+ self.worker_status.get('Approved', set()))
+
+
+# Module-level singleton instance
+_mturk_hit = None
+
+
+def init_mturk_hit(config):
+ """
+ Initialize the global MTurk HIT manager from configuration.
+
+ Args:
+ config (dict): Configuration dictionary containing:
+ - mturk.config_file_path: Path to MTurk config YAML file
+
+ The MTurk config file should contain:
+ - aws_access_key_id: AWS access key (optional, uses env/credentials file)
+ - aws_secret_access_key: AWS secret key (optional)
+ - sandbox: Whether to use sandbox (default: True)
+ - hit_id: Optional HIT ID to manage
+
+ Returns:
+ MTurkHIT: Initialized MTurk HIT manager, or None if not configured
+ """
+ global _mturk_hit
+
+ if not BOTO3_AVAILABLE:
+ logger.warning("boto3 not available, MTurk API disabled")
+ return None
+
+ mturk_config = config.get('mturk', {})
+ if not mturk_config.get('enabled', False):
+ logger.debug("MTurk API not enabled in config")
+ return None
+
+ config_file_path = mturk_config.get('config_file_path')
+ if not config_file_path:
+ logger.warning("MTurk enabled but no config_file_path specified")
+ return None
+
+ try:
+ import yaml
+ with open(config_file_path, 'r') as f:
+ mturk_settings = yaml.safe_load(f)
+
+ _mturk_hit = MTurkHIT(
+ aws_access_key_id=mturk_settings.get('aws_access_key_id'),
+ aws_secret_access_key=mturk_settings.get('aws_secret_access_key'),
+ hit_id=mturk_settings.get('hit_id'),
+ sandbox=mturk_settings.get('sandbox', True)
+ )
+
+ logger.info(f"MTurk HIT manager initialized from {config_file_path}")
+ return _mturk_hit
+
+ except FileNotFoundError:
+ logger.error(f"MTurk config file not found: {config_file_path}")
+ return None
+ except Exception as e:
+ logger.error(f"Failed to initialize MTurk HIT manager: {e}")
+ return None
+
+
+def get_mturk_hit():
+ """
+ Get the global MTurk HIT manager instance.
+
+ Returns:
+ MTurkHIT: The MTurk HIT manager, or None if not initialized
+ """
+ return _mturk_hit
+
+
+def clear_mturk_hit():
+ """
+ Clear the global MTurk HIT manager instance.
+
+ Used primarily for testing to reset state between tests.
+ """
+ global _mturk_hit
+ _mturk_hit = None
diff --git a/potato/server_utils/overlap_sampler.py b/potato/server_utils/overlap_sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..de941fcf76d114bcb6d28e96e22a8991ff5f1ae9
--- /dev/null
+++ b/potato/server_utils/overlap_sampler.py
@@ -0,0 +1,90 @@
+"""
+Overlap sampling for heterogeneous annotator coverage.
+
+Given the ``num_annotators_per_item.overlap_sample`` config block, selects a
+deterministic fraction of items to receive a raised annotator cap. The cap
+override is written onto each sampled item's metadata so the centralized
+``ItemStateManager._get_annotator_cap_for_item`` helper picks it up.
+
+This module is called once after all items have been loaded.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+import logging
+import random as _random
+
+logger = logging.getLogger(__name__)
+
+
+_METADATA_KEY = "required_annotations"
+
+
+def apply_overlap_sample(item_state_manager, config: dict) -> Dict[str, int]:
+ """
+ Stamp ``required_annotations`` on a deterministic sample of items.
+
+ Returns a mapping of sampled instance_id -> assigned cap, for reporting.
+ Items whose existing metadata already carries ``required_annotations``
+ (e.g., from a previous load with persisted state) are not overwritten.
+ """
+ nap = config.get("num_annotators_per_item")
+ if not isinstance(nap, dict):
+ return {}
+ overlap = nap.get("overlap_sample")
+ if not overlap:
+ return {}
+
+ fraction = float(overlap["fraction"])
+ count = int(overlap["count"])
+ stratify_by = overlap.get("stratify_by")
+ seed = int(overlap.get("seed", item_state_manager.random_seed))
+ rng = _random.Random(seed)
+
+ all_ids = list(item_state_manager.instance_id_to_instance.keys())
+ if not all_ids:
+ return {}
+
+ # Build strata
+ strata: Dict[Optional[str], List[str]] = defaultdict(list)
+ if stratify_by:
+ for iid in all_ids:
+ item = item_state_manager.instance_id_to_instance[iid]
+ data = item.get_data() if hasattr(item, "get_data") else {}
+ key = data.get(stratify_by) if isinstance(data, dict) else None
+ # Also accept the indexed category if it matches
+ if key is None and hasattr(item_state_manager, "instance_id_to_categories"):
+ cats = item_state_manager.instance_id_to_categories.get(iid)
+ if cats:
+ key = sorted(cats)[0]
+ strata[key if key is not None else "__uncategorized__"].append(iid)
+ else:
+ strata[None] = list(all_ids)
+
+ sampled: Dict[str, int] = {}
+ for key, ids in strata.items():
+ if not ids:
+ continue
+ # Deterministic ordering across runs
+ ids_sorted = sorted(ids)
+ rng_local = _random.Random(f"{seed}:{key}" if key is not None else seed)
+ rng_local.shuffle(ids_sorted)
+ target = max(1, int(round(len(ids_sorted) * fraction)))
+ for iid in ids_sorted[:target]:
+ item = item_state_manager.instance_id_to_instance[iid]
+ # Don't clobber an existing per-item override (operator may have
+ # set one via item_data; respect that authority).
+ if item.get_metadata(_METADATA_KEY) is not None:
+ continue
+ item.add_metadata(_METADATA_KEY, count)
+ sampled[iid] = count
+
+ if sampled:
+ logger.info(
+ "Overlap sample: %d / %d items raised to %d annotators (fraction=%s, stratify_by=%s, seed=%d)",
+ len(sampled), len(all_ids), count, fraction, stratify_by, seed,
+ )
+ return sampled
diff --git a/potato/server_utils/prolific_apis.py b/potato/server_utils/prolific_apis.py
new file mode 100644
index 0000000000000000000000000000000000000000..d644fcd53953ef4cb37f7fc057d6671d9725f209
--- /dev/null
+++ b/potato/server_utils/prolific_apis.py
@@ -0,0 +1,509 @@
+"""
+Utility functions for handling Prolific APIs
+
+This module provides wrapper classes for interacting with the Prolific API
+(https://api.prolific.com/api/v1/) to manage studies and submissions.
+
+The Prolific API uses token-based authentication and returns JSON responses.
+All endpoints require the Authorization header with format: 'Token {your_token}'
+
+Key API Endpoints Used:
+- GET /studies/ - List all studies
+- GET /studies/{id}/ - Get study details
+- GET /submissions/ - List all submissions
+- GET /submissions?study={id} - Get submissions for a study
+- GET /submissions/{id}/ - Get submission details
+- GET /studies/{id}/submissions/ - Get recent submissions for a study
+- POST /studies/{id}/transition/ - Change study status (START/PAUSE)
+
+For more information, see: https://docs.prolific.com/reference/
+"""
+import os.path
+import pandas as pd
+import requests
+from collections import OrderedDict, defaultdict
+import time
+import json
+import logging
+
+# The base wrapper of prolific apis
+class ProlificBase(object):
+ """
+ Base class for Prolific API operations.
+
+ Provides low-level API access to Prolific's REST endpoints.
+ All methods use token-based authentication and handle HTTP responses.
+
+ Attributes:
+ headers (dict): HTTP headers including Authorization token
+ """
+
+ def __init__(self, token):
+ """
+ Initialize the API client with authentication token.
+
+ Args:
+ token (str): Prolific API token for authentication
+ """
+ self.headers = {
+ 'Authorization': f'Token {token}',
+ }
+
+ def list_all_studies(self):
+ """
+ Retrieve all studies from Prolific.
+
+ Makes a GET request to /api/v1/studies/ to fetch all studies
+ associated with the authenticated account.
+
+ Returns:
+ pandas.DataFrame: DataFrame containing study information with columns:
+ - id: Study identifier
+ - name: Study name
+ - study_type: Type of study
+ - internal_name: Internal reference name
+ - status: Current study status
+ None: If the request fails
+ """
+ url = 'https://api.prolific.com/api/v1/studies/'
+ response = requests.get(url, headers=self.headers)
+ if response.status_code == 200:
+ data = response.json() # If the response contains JSON data
+ studies = pd.DataFrame.from_records(data['results'])
+ print('You currently have %s studies'%len(data['results']))
+ print(studies[['id','name','study_type','internal_name','status']].to_records())
+ return studies
+ else:
+ print(f"Error: {response.status_code} - {response.text}")
+ return None
+
+ def get_study_by_id(self, study_id):
+ """
+ Retrieve detailed information about a specific study.
+
+ Makes a GET request to /api/v1/studies/{study_id}/ to fetch
+ complete study details including configuration and status.
+
+ Args:
+ study_id (str, optional): Prolific study ID. If None, uses self.study_id
+
+ Returns:
+ dict: Complete study information including:
+ - id, name, internal_name
+ - reward, average_reward_per_hour
+ - external_study_url, status
+ - total_available_places, places_taken
+ - and other study configuration fields
+ None: If the request fails
+ """
+ if study_id == None:
+ study_id = self.study_id
+ url = f'https://api.prolific.com/api/v1/studies/{study_id}/'
+ response = requests.get(url, headers=self.headers)
+ if response.status_code == 200:
+ data = response.json() # If the response contains JSON data
+ return data
+ else:
+ print(f"Error: {response.status_code} - {response.text}")
+ return None
+
+ def get_submissions(self):
+ """
+ Retrieve all submissions across all studies.
+
+ Makes a GET request to /api/v1/submissions/ to fetch all submissions.
+ Note: This can be slow for accounts with many submissions.
+
+ Returns:
+ list: List of submission dictionaries containing:
+ - id: Submission identifier
+ - participant_id: Participant's Prolific ID
+ - study_id: Associated study ID
+ - status: Current submission status
+ - and other submission details
+ None: If the request fails
+ """
+ url = 'https://api.prolific.com/api/v1/submissions/'
+ response = requests.get(url, headers=self.headers)
+ if response.status_code == 200:
+ data = response.json() # If the response contains JSON data
+ print('You currently have %s submissions'%len(data['results']))
+ return data['results']
+ else:
+ print(f"Error: {response.status_code} - {response.text}")
+ return None
+
+ def get_submissions_from_study(self, study_id = None):
+ """
+ Retrieve all submissions for a specific study.
+
+ Makes a GET request to /api/v1/submissions?study={study_id} to fetch
+ submissions filtered by study ID.
+
+ Args:
+ study_id (str, optional): Prolific study ID. If None, uses self.study_id
+
+ Returns:
+ list: List of submission dictionaries for the specified study
+ None: If the request fails
+ """
+ if study_id == None:
+ study_id = self.study_id
+ api_endpoint = 'https://api.prolific.com/api/v1/submissions?study={}'
+ url = api_endpoint.format(study_id)
+ response = requests.get(url, headers=self.headers)
+ if response.status_code == 200:
+ data = response.json()['results']
+ print('Successfully fetched %s submissions from study %s' % (len(data), study_id))
+ return data
+ else:
+ print(f"Error: {response.status_code} - {response.text}")
+ return None
+
+ def get_submission_from_id(self, submission_id):
+ """
+ Retrieve detailed information about a specific submission.
+
+ Makes a GET request to /api/v1/submissions/{submission_id}/ to fetch
+ complete submission details including participant info and status.
+
+ Args:
+ submission_id (str): Prolific submission ID
+
+ Returns:
+ dict: Complete submission information including:
+ - id, participant_id, study_id
+ - status: Current submission status
+ - started_at, completed_at timestamps
+ - and other submission details
+ None: If the request fails
+ """
+ url = f'https://api.prolific.com/api/v1/submissions/{submission_id}/'
+ response = requests.get(url, headers=self.headers)
+ if response.status_code == 200:
+ data = response.json() # If the response contains JSON data
+ return data
+ else:
+ print(f"Error: {response.status_code} - {response.text}")
+ return None
+
+ def get_recent_study_submissions(self, study_id):
+ """
+ Retrieve recent submissions for a specific study.
+
+ Makes a GET request to /api/v1/studies/{study_id}/submissions/ to fetch
+ recent submissions. This endpoint may return a subset of submissions
+ compared to get_submissions_from_study().
+
+ Args:
+ study_id (str, optional): Prolific study ID. If None, uses self.study_id
+
+ Returns:
+ list: List of recent submission dictionaries for the specified study
+ None: If the request fails
+ """
+ if study_id == None:
+ study_id = self.study_id
+ url = f'https://api.prolific.com/api/v1/studies/{study_id}/submissions/'
+ response = requests.get(url, headers=self.headers)
+ if response.status_code == 200:
+ data = response.json() # If the response contains JSON data
+ print('You currently have %s submissions' % len(data['results']))
+ return (data['results'])
+ else:
+ print(f"Error: {response.status_code} - {response.text}")
+ return None
+
+ def get_study_status(self, study_id = None):
+ """
+ Get the current status of a study.
+
+ Retrieves study information and extracts the status field.
+
+ Args:
+ study_id (str, optional): Prolific study ID. If None, uses self.study_id
+
+ Returns:
+ str: Study status (e.g., 'ACTIVE', 'PAUSED', 'COMPLETED')
+ None: If the request fails
+ """
+ if study_id == None:
+ study_id = self.study_id
+ data = self.get_study_by_id(study_id)
+ if data:
+ return data['status']
+ else:
+ return None
+
+ def pause_study(self, study_id = None):
+ """
+ Pause a study to stop new participants from joining.
+
+ Makes a POST request to /api/v1/studies/{study_id}/transition/ with
+ action "PAUSE" to change the study status to paused.
+
+ Args:
+ study_id (str, optional): Prolific study ID. If None, uses self.study_id
+
+ Returns:
+ dict: Response from the transition API call
+ """
+ if study_id == None:
+ study_id = self.study_id
+ api_endpoint = 'https://api.prolific.com/api/v1/studies/{}/transition/'
+ url = api_endpoint.format(study_id)
+ data = {
+ "action": "PAUSE"
+ }
+ response = requests.post(url, headers=self.headers, json=data)
+ data = response.json()
+ print(study_id, self.get_study_status(study_id))
+ return data
+
+ def start_study(self, study_id = None):
+ """
+ Start a study to allow new participants to join.
+
+ Makes a POST request to /api/v1/studies/{study_id}/transition/ with
+ action "START" to change the study status to active.
+
+ Args:
+ study_id (str, optional): Prolific study ID. If None, uses self.study_id
+
+ Returns:
+ dict: Response from the transition API call
+ """
+ if study_id == None:
+ study_id = self.study_id
+ api_endpoint = 'https://api.prolific.com/api/v1/studies/{}/transition/'
+ url = api_endpoint.format(study_id)
+ data = {
+ "action": "START"
+ }
+ response = requests.post(url, headers=self.headers, json=data)
+ data = response.json()
+ print(study_id, self.get_study_status(study_id))
+ return data
+
+
+# The class to manage the status of a prolific study
+class ProlificStudy(ProlificBase):
+ """
+ High-level study management class for Prolific studies.
+
+ Extends ProlificBase to provide study-specific functionality including:
+ - Submission tracking and status management
+ - Workload monitoring and automatic study control
+ - Local state persistence
+
+ Attributes:
+ study_id (str): Prolific study identifier
+ study_info (dict): Cached study information
+ submission_info_path (str): Path to local submission data file
+ sessions (OrderedDict): Mapping of submission IDs to submission data
+ user_status_dict (defaultdict): Mapping of status to participant IDs
+ max_concurrent_sessions (int): Maximum allowed concurrent participants
+ checker_period (int): Seconds between workload checks
+ workload_checker_on (bool): Whether workload checker is running
+ """
+
+ def __init__(self, token, study_id, saving_dir, max_concurrent_sessions = 30, workload_checker_period = 60):
+ """
+ Initialize study management with configuration.
+
+ Args:
+ token (str): Prolific API token
+ study_id (str): Prolific study identifier
+ saving_dir (str): Directory to save submission data
+ max_concurrent_sessions (int): Maximum concurrent participants (default: 30)
+ workload_checker_period (int): Seconds between workload checks (default: 60)
+ """
+ ProlificBase.__init__(self, token)
+ self.study_id = study_id
+ self.study_info = self.get_study_by_id(study_id)
+ self.submission_info_path = os.path.join(saving_dir, 'submissions.json')
+ self.sessions = OrderedDict()
+ self.user2session = {}
+ #self.user_status_dict = {'RESERVED':set(), 'AWAITING REVIEW':set(), 'RETURNED':set(), 'TIMED-OUT':set(), 'ACTIVE': set(), 'APPROVED':set(), 'REJECTED':set()}
+ self.study_status = None
+ self.status_path = None
+ self.max_concurrent_sessions = max_concurrent_sessions # How many users can work on the study at the same time
+ self.checker_period = workload_checker_period
+ self.workload_checker_remaining_time = workload_checker_period
+ self.workload_checker_on = False
+
+ def get_basic_study_info(self):
+ """
+ Extract basic study information for display or logging.
+
+ Returns:
+ dict: Basic study information including:
+ - id, name, internal_name
+ - reward, average_reward_per_hour
+ - external_study_url, status
+ - total_available_places, places_taken
+ """
+ keys = ['id', 'name', 'internal_name',
+ 'reward', 'average_reward_per_hour', 'external_study_url', 'status', 'total_available_places', 'places_taken']
+ return {key:self.study_info[key] for key in keys}
+
+ def update_submission_status(self):
+ """
+ Refresh local submission data from Prolific API.
+
+ Fetches current submission data from the API and updates local state:
+ - Saves submission data to local JSON file
+ - Updates sessions mapping
+ - Rebuilds user status dictionary grouped by submission status
+
+ The user_status_dict maps submission statuses to sets of participant IDs:
+ - 'ACTIVE': Currently working participants
+ - 'AWAITING REVIEW': Completed submissions pending review
+ - 'APPROVED': Approved submissions
+ - 'REJECTED': Rejected submissions
+ - 'RETURNED': Participants who returned the study
+ - 'TIMED-OUT': Participants who timed out
+ """
+ submission_data = self.get_submissions_from_study()
+ with open(self.submission_info_path, "wt") as f:
+ for v in submission_data:
+ f.writelines(json.dumps(v) + "\n")
+ self.user_status_dict = defaultdict(set)
+ for v in submission_data:
+ self.sessions[v['id']] = v
+ #self.user2session[v['participant_id']] = v['id']
+ self.user_status_dict[v['status']].add(v['participant_id'])
+ self.reclaim_dropped_user_assignments()
+
+ def get_dropped_users(self):
+ """
+ Get list of participants who are no longer active.
+
+ Returns:
+ list: Participant IDs who have returned, timed out, or been rejected
+ """
+ return list(self.user_status_dict['RETURNED'] | self.user_status_dict['TIMED-OUT'] | self.user_status_dict['REJECTED'])
+
+ def get_dropped_users_by_status(self):
+ """
+ Get dropped participant IDs grouped by Prolific submission status.
+
+ Returns:
+ dict: Mapping from RETURNED, TIMED-OUT, and REJECTED to participant IDs.
+ """
+ return {
+ 'RETURNED': list(self.user_status_dict['RETURNED']),
+ 'TIMED-OUT': list(self.user_status_dict['TIMED-OUT']),
+ 'REJECTED': list(self.user_status_dict['REJECTED']),
+ }
+
+ def reclaim_dropped_user_assignments(self):
+ """
+ Release unannotated Potato assignments for dropped Prolific workers.
+
+ Prolific reports dropped workers as RETURNED, TIMED-OUT, or REJECTED.
+ Their participant IDs are also the Potato usernames for url-direct
+ Prolific login, so the item state manager can safely reclaim any
+ assigned-but-unannotated instances.
+ """
+ dropped_by_status = self.get_dropped_users_by_status()
+ if not any(dropped_by_status.values()):
+ return {}
+
+ status_to_reason = {
+ 'RETURNED': 'prolific_returned',
+ 'TIMED-OUT': 'prolific_timed_out',
+ 'REJECTED': 'prolific_rejected',
+ }
+
+ try:
+ from potato.item_state_management import get_item_state_manager
+ manager = get_item_state_manager()
+ reclaimed = {}
+ for status, user_ids in dropped_by_status.items():
+ if not user_ids:
+ continue
+ reclaimed.update(
+ manager.reclaim_unannotated_assignments_for_users(
+ user_ids,
+ reason=status_to_reason[status],
+ )
+ )
+ return reclaimed
+ except Exception as e:
+ logging.getLogger(__name__).warning(
+ "Could not reclaim assignments for dropped Prolific users: %s",
+ e,
+ )
+ return {}
+
+ def get_concurrent_sessions_count(self):
+ """
+ Get the number of currently active participants.
+
+ Returns:
+ int: Number of participants with 'ACTIVE' status
+ """
+ return len(self.user_status_dict['ACTIVE'])
+
+ def workload_checker(self):
+ """
+ Monitor study workload and automatically manage study status.
+
+ Periodically checks the number of active participants and automatically:
+ - Pauses the study if too many participants are active
+ - Resumes the study when active participants drop below threshold
+
+ The threshold is 20% of max_concurrent_sessions. The checker runs
+ continuously with the specified checker_period interval.
+
+ This method runs in an infinite loop and should be called in a separate
+ thread to avoid blocking the main application.
+ """
+ if self.workload_checker_on:
+ print('Workload checker already in process, time remaining: %s seconds' % self.workload_checker_remaining_time)
+ return None
+ else:
+ print('Workload checker started, checking every %s seconds'%self.checker_period)
+ while True:
+ self.workload_checker_remaining_time = self.checker_period
+ self.workload_checker_on = True
+ print(f"\rChecking workload in: {self.checker_period} seconds")
+ #time.sleep(self.checker_period)
+ for i in range(self.checker_period, 0, -1):
+ #print(f"\rChecking workload in: {i}s", end='', flush=True)
+ self.workload_checker_remaining_time -= 1
+ time.sleep(1)
+ self.update_submission_status()
+ if self.get_concurrent_sessions_count() < 0.2 * self.max_concurrent_sessions:
+ self.workload_checker_on = False
+ print('current workload: ', self.get_concurrent_sessions_count(), ', resuming study %s'%self.study_id)
+ self.start_study()
+ return None
+ else:
+ print('current workload: ', self.get_concurrent_sessions_count(), ', starting another workload checker')
+
+ def update_session_status(self, sess_id):
+ """
+ Update the status of a specific submission.
+
+ Fetches current status from Prolific API and updates local session data.
+
+ Args:
+ sess_id (str): Submission ID to update
+ """
+ status = self.get_submission_from_id(sess_id)['status']
+ self.sessions[sess_id]['status'] = status
+
+ def add_new_user(self, user):
+ """
+ Add a new participant to the local session tracking.
+
+ Fetches current submission status and adds to local state.
+
+ Args:
+ user (dict): User data containing 'SESSION_ID' and 'PROLIFIC_PID'
+ """
+ status = self.get_submission_from_id(user['SESSION_ID'])['status']
+ self.sessions[user['SESSION_ID']] = {'username':user['PROLIFIC_PID'], 'status':status}
+ self.session_status_dict[status].append(user['SESSION_ID'])
diff --git a/potato/server_utils/schemas/__init__.py b/potato/server_utils/schemas/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9fafb12e418f7f8408baf48a7cd57b4a3a0dbc83
--- /dev/null
+++ b/potato/server_utils/schemas/__init__.py
@@ -0,0 +1,55 @@
+from .likert import generate_likert_layout
+from .multiselect import generate_multiselect_layout
+from .multirate import generate_multirate_layout
+from .number import generate_number_layout
+from .pure_display import generate_pure_display_layout
+from .radio import generate_radio_layout
+from .select import generate_select_layout
+from .span import generate_span_layout, render_span_annotations, get_spans_for_field
+from .span_link import generate_span_link_layout
+from .textbox import generate_textbox_layout
+from .slider import generate_slider_layout
+from .video import generate_video_layout
+from .image_annotation import generate_image_annotation_layout
+from .audio_annotation import generate_audio_annotation_layout
+from .video_annotation import generate_video_annotation_layout
+from .pairwise import generate_pairwise_layout
+from .coreference import generate_coreference_layout
+from .tree_annotation import generate_tree_annotation_layout
+from .triage import generate_triage_layout
+from .event_annotation import generate_event_annotation_layout
+from .tiered_annotation import generate_tiered_annotation_layout
+from .bws import generate_bws_layout
+from .soft_label import generate_soft_label_layout
+from .confidence import generate_confidence_layout
+from .constant_sum import generate_constant_sum_layout
+from .semantic_differential import generate_semantic_differential_layout
+from .ranking import generate_ranking_layout
+from .range_slider import generate_range_slider_layout
+from .hierarchical_multiselect import generate_hierarchical_multiselect_layout
+from .vas import generate_vas_layout
+from .extractive_qa import generate_extractive_qa_layout
+from .rubric_eval import generate_rubric_eval_layout
+from .text_edit import generate_text_edit_layout
+from .error_span import generate_error_span_layout
+from .card_sort import generate_card_sort_layout
+from .conjoint import generate_conjoint_layout
+from .trajectory_eval import generate_trajectory_eval_layout
+from .trajectory_edit import generate_trajectory_edit_layout
+from .process_reward import generate_process_reward_layout
+from .code_review import generate_code_review_layout
+
+# Import identifier utilities for use by other modules
+from .identifier_utils import (
+ validate_schema_config,
+ generate_element_identifier,
+ generate_element_value,
+ generate_validation_attribute,
+ escape_html_content,
+ safe_generate_layout,
+ generate_tooltip_html,
+ generate_layout_attributes
+)
+
+# Import schema registry for centralized schema management
+from .registry import schema_registry, SchemaDefinition
diff --git a/potato/server_utils/schemas/audio_annotation.py b/potato/server_utils/schemas/audio_annotation.py
new file mode 100644
index 0000000000000000000000000000000000000000..e3b6822e1cfe3acc8ed88a78e6bb78e4a4789237
--- /dev/null
+++ b/potato/server_utils/schemas/audio_annotation.py
@@ -0,0 +1,719 @@
+"""
+Audio Annotation Layout
+
+Generates a form interface for audio annotation with segmentation.
+Uses Peaks.js for waveform visualization and segment management.
+
+Features:
+- Waveform visualization (amplitude display)
+- Spectrogram visualization (frequency display) - optional
+- Region/segment selection and playback
+- Zoom and pan for long audio files
+- Two annotation modes:
+ - Label mode: Assign labels to segments (like span annotation)
+ - Questions mode: Answer questions about each segment (radio, multirate, etc.)
+"""
+
+import logging
+import json
+from typing import List, Dict, Tuple, Any
+from .identifier_utils import (
+ safe_generate_layout,
+ escape_html_content
+)
+
+logger = logging.getLogger(__name__)
+
+# Default colors for segment labels
+DEFAULT_COLORS = [
+ "#4ECDC4", # Teal
+ "#FF6B6B", # Red
+ "#45B7D1", # Blue
+ "#96CEB4", # Green
+ "#FFEAA7", # Yellow
+ "#DDA0DD", # Plum
+ "#95A5A6", # Gray
+ "#F39C12", # Orange
+ "#9B59B6", # Purple
+ "#3498DB", # Light Blue
+]
+
+# Valid annotation modes
+VALID_MODES = ["label", "questions", "both"]
+
+# Default spectrogram options
+DEFAULT_SPECTROGRAM_OPTIONS = {
+ "fft_size": 2048,
+ "hop_length": 512,
+ "frequency_range": [0, 8000],
+ "color_map": "viridis",
+}
+
+# Valid color maps for spectrogram
+VALID_COLOR_MAPS = ["viridis", "magma", "plasma", "inferno", "grayscale"]
+
+
+def generate_audio_annotation_layout(annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Generate HTML for an audio annotation interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - mode: "label" | "questions" | "both"
+ - labels: List of segment labels (for label/both modes)
+ - segment_schemes: List of annotation schemes per segment (for questions/both modes)
+ - min_segments: Minimum required segments (default: 0)
+ - max_segments: Maximum allowed segments (default: null/unlimited)
+ - zoom_enabled: Whether to enable zoom (default: True)
+ - playback_rate_control: Show playback speed controls (default: False)
+ - waveform: Whether to show waveform (default: True)
+ - spectrogram: Whether to show spectrogram (default: False)
+ - spectrogram_options: Spectrogram configuration (optional):
+ - fft_size: FFT window size (default: 2048)
+ - hop_length: Hop length between FFT windows (default: 512)
+ - frequency_range: [min_hz, max_hz] (default: [0, 8000])
+ - color_map: Color mapping ("viridis", "magma", "plasma", "inferno", "grayscale")
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the audio annotation interface
+ key_bindings: List of keyboard shortcuts
+
+ Raises:
+ ValueError: If required fields are missing or invalid
+ """
+ return safe_generate_layout(annotation_scheme, _generate_audio_annotation_layout_internal)
+
+
+def _generate_audio_annotation_layout_internal(annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Internal function to generate audio annotation layout after validation.
+ """
+ schema_name = annotation_scheme.get('name', 'audio_annotation')
+ logger.debug(f"Generating audio annotation layout for schema: {schema_name}")
+
+ # Get mode (default to "label" for simplicity)
+ mode = annotation_scheme.get('mode', 'label')
+ if mode not in VALID_MODES:
+ error_msg = f"Invalid mode '{mode}' in schema: {schema_name}. Must be one of: {VALID_MODES}"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ # Validate labels for label/both modes
+ labels = []
+ if mode in ['label', 'both']:
+ if 'labels' not in annotation_scheme:
+ error_msg = f"Missing labels in schema: {schema_name} (required for mode '{mode}')"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+ labels = _process_labels(annotation_scheme['labels'])
+
+ # Validate segment_schemes for questions/both modes
+ segment_schemes = []
+ if mode in ['questions', 'both']:
+ if 'segment_schemes' not in annotation_scheme:
+ error_msg = f"Missing segment_schemes in schema: {schema_name} (required for mode '{mode}')"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+ segment_schemes = annotation_scheme['segment_schemes']
+ if not isinstance(segment_schemes, list) or not segment_schemes:
+ error_msg = f"segment_schemes must be a non-empty list in schema: {schema_name}"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ # Get configuration options
+ min_segments = annotation_scheme.get('min_segments', 0)
+ max_segments = annotation_scheme.get('max_segments', None)
+ zoom_enabled = annotation_scheme.get('zoom_enabled', True)
+ playback_rate_control = annotation_scheme.get('playback_rate_control', False)
+
+ # Waveform and spectrogram display options
+ show_waveform = annotation_scheme.get('waveform', True)
+ show_spectrogram = annotation_scheme.get('spectrogram', False)
+
+ # Process spectrogram options with defaults
+ spectrogram_options = _process_spectrogram_options(
+ annotation_scheme.get('spectrogram_options', {}),
+ schema_name
+ )
+
+ # source_field: Links this annotation schema to a display field from instance_display
+ source_field = annotation_scheme.get("source_field", "")
+
+ # Build config object for JavaScript
+ js_config = {
+ "schemaName": schema_name,
+ "mode": mode,
+ "labels": labels,
+ "segmentSchemes": segment_schemes,
+ "minSegments": min_segments,
+ "maxSegments": max_segments,
+ "zoomEnabled": zoom_enabled,
+ "playbackRateControl": playback_rate_control,
+ "sourceField": source_field,
+ "waveform": show_waveform,
+ "spectrogram": show_spectrogram,
+ "spectrogramOptions": spectrogram_options,
+ }
+
+ # Generate HTML
+ html = _generate_html(annotation_scheme, js_config, schema_name, labels, mode)
+
+ # Generate keybindings
+ keybindings = _generate_keybindings(labels, mode)
+
+ logger.info(f"Successfully generated audio annotation layout for {schema_name}")
+ return html, keybindings
+
+
+def _process_labels(labels_config: List) -> List[Dict[str, Any]]:
+ """
+ Process label configuration and assign colors.
+
+ Args:
+ labels_config: List of label configs (strings or dicts)
+
+ Returns:
+ List of processed label dicts with name, color, and optional key_value
+ """
+ processed = []
+ for i, label in enumerate(labels_config):
+ if isinstance(label, str):
+ processed.append({
+ "name": label,
+ "color": DEFAULT_COLORS[i % len(DEFAULT_COLORS)],
+ })
+ elif isinstance(label, dict):
+ processed.append({
+ "name": label.get("name", f"label_{i}"),
+ "color": label.get("color", DEFAULT_COLORS[i % len(DEFAULT_COLORS)]),
+ "key_value": label.get("key_value"),
+ })
+ else:
+ processed.append({
+ "name": str(label),
+ "color": DEFAULT_COLORS[i % len(DEFAULT_COLORS)],
+ })
+ return processed
+
+
+def _process_spectrogram_options(options: Dict[str, Any], schema_name: str) -> Dict[str, Any]:
+ """
+ Process and validate spectrogram configuration options.
+
+ Args:
+ options: User-provided spectrogram options
+ schema_name: Schema name for error messages
+
+ Returns:
+ Validated and merged spectrogram options with defaults
+ """
+ # Start with defaults
+ processed = dict(DEFAULT_SPECTROGRAM_OPTIONS)
+
+ if not options:
+ return processed
+
+ # Validate and merge fft_size
+ if 'fft_size' in options:
+ fft_size = options['fft_size']
+ if isinstance(fft_size, int) and fft_size > 0:
+ # Ensure power of 2 for FFT efficiency
+ if fft_size & (fft_size - 1) == 0:
+ processed['fft_size'] = fft_size
+ else:
+ logger.warning(
+ f"fft_size {fft_size} is not a power of 2 in schema {schema_name}, "
+ f"using default {DEFAULT_SPECTROGRAM_OPTIONS['fft_size']}"
+ )
+
+ # Validate and merge hop_length
+ if 'hop_length' in options:
+ hop_length = options['hop_length']
+ if isinstance(hop_length, int) and hop_length > 0:
+ processed['hop_length'] = hop_length
+
+ # Validate and merge frequency_range
+ if 'frequency_range' in options:
+ freq_range = options['frequency_range']
+ if (isinstance(freq_range, (list, tuple)) and len(freq_range) == 2
+ and all(isinstance(f, (int, float)) for f in freq_range)
+ and freq_range[0] < freq_range[1]):
+ processed['frequency_range'] = list(freq_range)
+ else:
+ logger.warning(
+ f"Invalid frequency_range {freq_range} in schema {schema_name}, "
+ f"using default {DEFAULT_SPECTROGRAM_OPTIONS['frequency_range']}"
+ )
+
+ # Validate and merge color_map
+ if 'color_map' in options:
+ color_map = options['color_map']
+ if color_map in VALID_COLOR_MAPS:
+ processed['color_map'] = color_map
+ else:
+ logger.warning(
+ f"Invalid color_map '{color_map}' in schema {schema_name}, "
+ f"must be one of {VALID_COLOR_MAPS}. Using default '{DEFAULT_SPECTROGRAM_OPTIONS['color_map']}'"
+ )
+
+ return processed
+
+
+def _generate_html(
+ annotation_scheme: Dict[str, Any],
+ js_config: Dict[str, Any],
+ schema_name: str,
+ labels: List[Dict[str, Any]],
+ mode: str
+) -> str:
+ """
+ Generate the HTML for the audio annotation interface.
+ """
+ escaped_name = escape_html_content(schema_name)
+ description = escape_html_content(annotation_scheme.get('description', ''))
+ config_json = json.dumps(js_config)
+
+ # Determine display mode
+ show_waveform = js_config.get('waveform', True)
+ show_spectrogram = js_config.get('spectrogram', False)
+
+ # Generate label buttons (for label/both modes)
+ label_selector = ""
+ if mode in ['label', 'both'] and labels:
+ label_selector = _generate_label_selector(labels)
+
+ # Generate playback rate control
+ playback_rate_html = ""
+ if js_config.get('playbackRateControl'):
+ playback_rate_html = '''
+
+ Speed:
+
+ 0.5x
+ 0.75x
+ 1x
+ 1.25x
+ 1.5x
+ 2x
+
+
+ '''
+
+ # source_field attribute for linking to display fields
+ source_field = annotation_scheme.get("source_field", "")
+ source_field_attr = f' data-source-field="{escape_html_content(source_field)}"' if source_field else ""
+
+ # Generate spectrogram HTML if enabled
+ spectrogram_html = ""
+ if show_spectrogram:
+ spectrogram_html = f'''
+
+
+
+ Spectrogram (frequency analysis)
+
+
+
+
+
+
+ '''
+
+ html = f'''
+
+ '''
+
+ return html
+
+
+def _generate_label_selector(labels: List[Dict[str, Any]]) -> str:
+ """
+ Generate HTML for label selection buttons.
+ """
+ buttons = []
+ for label in labels:
+ name = escape_html_content(label["name"])
+ color = label["color"]
+ key_hint = f' Keyboard shortcut: {label["key_value"]}' if label.get("key_value") else ""
+ tooltip = f"Select '{name}' as the label for new segments.{key_hint}"
+ buttons.append(
+ f''
+ f' '
+ f'{name} '
+ )
+
+ return f'''
+
+ Label:
+ {"".join(buttons)}
+
+ '''
+
+
+def _generate_keybindings(labels: List[Dict[str, Any]], mode: str) -> List[Tuple[str, str]]:
+ """
+ Generate keybinding list for the schema.
+ """
+ keybindings = []
+
+ # Playback shortcuts
+ keybindings.append(("Space", "Play/Pause"))
+ keybindings.append(("โ/โ", "Seek 5 seconds"))
+
+ # Label shortcuts (for label/both modes)
+ if mode in ['label', 'both']:
+ for label in labels:
+ if label.get("key_value"):
+ keybindings.append((label["key_value"], f"Select: {label['name']}"))
+
+ # Segment shortcuts
+ keybindings.append(("[/]", "Set segment start/end"))
+ keybindings.append(("Enter", "Create segment"))
+ keybindings.append(("Del", "Delete segment"))
+
+ # Zoom shortcuts
+ keybindings.append(("+/-", "Zoom in/out"))
+ keybindings.append(("0", "Fit to view"))
+
+ return keybindings
diff --git a/potato/server_utils/schemas/bws.py b/potato/server_utils/schemas/bws.py
new file mode 100644
index 0000000000000000000000000000000000000000..d5e19ece083bf9b54aca1453420e9702dec635a6
--- /dev/null
+++ b/potato/server_utils/schemas/bws.py
@@ -0,0 +1,172 @@
+"""
+Best-Worst Scaling (BWS) Layout
+
+Generates a form interface for selecting the best and worst items from a tuple.
+Features include:
+- Labeled items display (A, B, C, D...) populated by JS from var_elems
+- Best selection row โ clickable tiles
+- Worst selection row โ clickable tiles
+- Validation that best != worst
+- Keyboard shortcuts: 1-9 for best, q/w/e/r for worst
+- Two hidden inputs storing the best/worst position labels
+
+Config keys:
+ - name: Schema identifier
+ - description: Display description (shown as heading)
+ - best_description: Question text for best selection
+ - worst_description: Question text for worst selection
+ - tuple_size: Number of items per tuple (default: 4)
+ - sequential_key_binding: Enable keyboard shortcuts (default: true)
+ - label_requirement: Optional validation settings
+"""
+
+import logging
+from typing import Any, Dict, List, Tuple
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def generate_bws_layout(
+ annotation_scheme: Dict[str, Any],
+) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Generate HTML for a Best-Worst Scaling interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - best_description: Question for best selection (default: "Which is BEST?")
+ - worst_description: Question for worst selection (default: "Which is WORST?")
+ - tuple_size: Items per tuple (default: 4)
+ - sequential_key_binding: Enable keyboard shortcuts (default: true)
+ - label_requirement (dict): Optional validation settings
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_bws_layout_internal)
+
+
+def _generate_bws_layout_internal(
+ annotation_scheme: Dict[str, Any],
+) -> Tuple[str, List[Tuple[str, str]]]:
+ """Internal function to generate BWS layout after validation."""
+ logger.debug(f"Generating BWS layout for schema: {annotation_scheme['name']}")
+
+ schema_name = annotation_scheme["name"]
+ description = annotation_scheme.get("description", "")
+ best_description = annotation_scheme.get(
+ "best_description", "Which is BEST?"
+ )
+ worst_description = annotation_scheme.get(
+ "worst_description", "Which is WORST?"
+ )
+ tuple_size = annotation_scheme.get("tuple_size", 4)
+ enable_keybindings = annotation_scheme.get("sequential_key_binding", True)
+
+ # Validation attribute for both hidden inputs
+ validation = generate_validation_attribute(annotation_scheme)
+
+ # Escape for HTML
+ escaped_schema = escape_html_content(schema_name)
+ escaped_description = escape_html_content(description)
+ escaped_best_desc = escape_html_content(best_description)
+ escaped_worst_desc = escape_html_content(worst_description)
+
+ # Layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Position labels: A, B, C, D, ...
+ positions = [chr(ord("A") + i) for i in range(tuple_size)]
+
+ # Build best tiles
+ best_tiles_html = ""
+ for idx, pos in enumerate(positions):
+ key_num = str(idx + 1)
+ shortcut = f"[{key_num}]" if enable_keybindings else ""
+ data_key = f'data-key="{key_num}"' if enable_keybindings else ""
+ best_tiles_html += f"""
+
+ {pos}
+ {shortcut}
+
"""
+
+ # Build worst tiles โ keys q, w, e, r (row below 1, 2, 3, 4)
+ worst_keys = "qwer"
+ worst_tiles_html = ""
+ for idx, pos in enumerate(positions):
+ key_letter = worst_keys[idx] if idx < len(worst_keys) else chr(ord("a") + idx)
+ shortcut = f"[{key_letter}]" if enable_keybindings else ""
+ data_key = f'data-key="{key_letter}"' if enable_keybindings else ""
+ worst_tiles_html += f"""
+
+ {pos}
+ {shortcut}
+
"""
+
+ # Build the complete form
+ schematic = f"""
+
+ """
+
+ # Key bindings โ best: 1,2,3,4 worst: q,w,e,r
+ worst_binding_keys = "qwer"
+ key_bindings = []
+ if enable_keybindings:
+ for idx, pos in enumerate(positions):
+ key_bindings.append((str(idx + 1), f"{schema_name}: Best {pos}"))
+ wk = worst_binding_keys[idx] if idx < len(worst_binding_keys) else chr(ord("a") + idx)
+ key_bindings.append(
+ (wk, f"{schema_name}: Worst {pos}")
+ )
+
+ logger.info(f"Successfully generated BWS layout for {schema_name}")
+ return schematic, key_bindings
diff --git a/potato/server_utils/schemas/card_sort.py b/potato/server_utils/schemas/card_sort.py
new file mode 100644
index 0000000000000000000000000000000000000000..5796efd0e0f5316ca6e53370f0928a78ec5eccdd
--- /dev/null
+++ b/potato/server_utils/schemas/card_sort.py
@@ -0,0 +1,287 @@
+"""
+Card Sorting / Grouping Layout
+
+Drag items (text snippets, labels, concepts) into predefined or user-created groups.
+Open card sorting lets annotators create their own categories; closed card sorting
+provides predefined ones.
+
+Research: Spencer (2009) "Card Sorting: Designing Usable Categories"; Nielsen Norman Group.
+"""
+
+import json
+import logging
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+def generate_card_sort_layout(annotation_scheme):
+ """
+ Generate HTML for a Card Sorting interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - mode: "closed" (predefined groups) or "open" (user-created groups)
+ - groups: List of group names (for closed mode)
+ - items_field: Field in data containing items to sort
+ - allow_empty_groups: Whether empty groups are OK
+ - allow_multiple: Whether an item can appear in multiple groups
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_card_sort_layout_internal)
+
+
+def _generate_card_sort_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme['name']
+ description = annotation_scheme['description']
+ mode = annotation_scheme.get('mode', 'closed')
+ groups = annotation_scheme.get('groups', [])
+ items_field = annotation_scheme.get('items_field', 'items')
+ allow_empty_groups = annotation_scheme.get('allow_empty_groups', True)
+ allow_multiple = annotation_scheme.get('allow_multiple', False)
+
+ if mode == 'closed' and not groups:
+ raise ValueError(f"card_sort schema '{schema_name}' in closed mode requires 'groups'")
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+ identifiers = generate_element_identifier(schema_name, schema_name, "hidden")
+
+ config_json = json.dumps({
+ 'mode': mode,
+ 'groups': groups,
+ 'items_field': items_field,
+ 'allow_empty_groups': allow_empty_groups,
+ 'allow_multiple': allow_multiple,
+ })
+
+ # Build group containers
+ # NOTE: group-items div must be self-closing (no whitespace) so CSS :empty works
+ groups_html = ""
+ for group in groups:
+ group_id = escape_html_content(group.replace(' ', '-').lower())
+ esc_group = escape_html_content(group)
+ esc_schema = escape_html_content(schema_name)
+ groups_html += (
+ f''
+ )
+
+ new_group_html = ""
+ if mode == 'open':
+ new_group_html = f"""
+
+
+ + Add Group
+
+ """
+
+ html = f"""
+
+
+
+ """
+
+ logger.info(f"Generated card sort layout for {schema_name}")
+ return html, []
diff --git a/potato/server_utils/schemas/code_review.py b/potato/server_utils/schemas/code_review.py
new file mode 100644
index 0000000000000000000000000000000000000000..6975860b35a3e5cb46848f31fbdde305c2847623
--- /dev/null
+++ b/potato/server_utils/schemas/code_review.py
@@ -0,0 +1,482 @@
+"""
+Code Review Schema
+
+GitHub PR review-style annotation with inline diff commenting,
+file-level ratings, and overall verdict.
+
+Features:
+- Click on diff lines in CodingTraceDisplay to add comments
+- Per-comment category (bug, style, suggestion, security, question)
+- File-level correctness and quality ratings
+- Overall verdict (approve, request_changes, comment_only)
+"""
+
+import json
+import logging
+from typing import Dict, Any, Tuple, List
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+)
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_CATEGORIES = ["bug", "style", "suggestion", "security", "question"]
+DEFAULT_VERDICTS = ["approve", "request_changes", "comment_only"]
+DEFAULT_RATING_DIMS = ["correctness", "quality"]
+
+
+def generate_code_review_layout(
+ annotation_scheme: Dict[str, Any],
+) -> Tuple[str, List[Tuple[str, str]]]:
+ """Generate HTML for a code review annotation interface.
+
+ Args:
+ annotation_scheme: Configuration dict. Required keys: ``name``,
+ ``description``. Optional: ``comment_categories``,
+ ``verdict_options``, ``file_rating_dimensions``.
+
+ Returns:
+ ``(html, keybindings)`` tuple.
+ """
+ return safe_generate_layout(annotation_scheme, _generate_internal)
+
+
+def _generate_internal(
+ annotation_scheme: Dict[str, Any],
+) -> Tuple[str, List[Tuple[str, str]]]:
+ schema_name = annotation_scheme["name"]
+ description = annotation_scheme["description"]
+
+ categories = annotation_scheme.get("comment_categories", DEFAULT_CATEGORIES)
+ verdicts = annotation_scheme.get("verdict_options", DEFAULT_VERDICTS)
+ rating_dims = annotation_scheme.get("file_rating_dimensions", DEFAULT_RATING_DIMS)
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+ identifiers = generate_element_identifier(schema_name, schema_name, "hidden")
+ esc_schema = escape_html_content(schema_name)
+
+ config_json = json.dumps({
+ "categories": categories,
+ "verdicts": verdicts,
+ "rating_dims": rating_dims,
+ })
+
+ # Build verdict radios
+ verdict_html = ""
+ for v in verdicts:
+ label = v.replace("_", " ").title()
+ css = f"cr-verdict-{v}"
+ verdict_html += (
+ f''
+ f' '
+ f' {escape_html_content(label)}'
+ f' '
+ )
+
+ # Build category options
+ cat_options = "".join(
+ f'{escape_html_content(c.title())} '
+ for c in categories
+ )
+
+ html = f"""
+
+
+
+
+
+ """
+
+ logger.info(
+ f"Successfully generated code_review layout for {schema_name} "
+ f"({len(categories)} categories, {len(verdicts)} verdicts)"
+ )
+ return html, [] # No keybindings
diff --git a/potato/server_utils/schemas/confidence.py b/potato/server_utils/schemas/confidence.py
new file mode 100644
index 0000000000000000000000000000000000000000..39bed10f78e468e3207eb87a389cfac658bf071b
--- /dev/null
+++ b/potato/server_utils/schemas/confidence.py
@@ -0,0 +1,164 @@
+"""
+Confidence-Calibrated Annotation Layout
+
+Generates a confidence rating that pairs with a primary annotation scheme.
+Supports both Likert-style discrete scale and continuous slider.
+
+Research basis:
+- Kutlu et al. (2020) "Annotator Rationales for Labeling Tasks in Crowdsourcing" JAIR
+- Sheng et al. (2008) "Get Another Label?" KDD
+"""
+
+import logging
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+)
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_SCALE_POINTS = 5
+DEFAULT_SCALE_TYPE = "likert"
+DEFAULT_LABELS = [
+ "Guessing",
+ "Somewhat confident",
+ "Fairly confident",
+ "Confident",
+ "Certain",
+]
+
+
+def generate_confidence_layout(annotation_scheme):
+ """
+ Generate HTML for a confidence-calibrated annotation interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - target_schema: Name of the primary schema this confidence is for (optional)
+ - scale_type: "likert" or "slider" (default "likert")
+ - scale_points: Number of points (default 5, for likert)
+ - labels: Custom scale labels (optional)
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_confidence_layout_internal)
+
+
+def _generate_confidence_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme["name"]
+ safe_schema = escape_html_content(schema_name)
+ description = annotation_scheme["description"]
+ scale_type = annotation_scheme.get("scale_type", DEFAULT_SCALE_TYPE)
+ scale_points = annotation_scheme.get("scale_points", DEFAULT_SCALE_POINTS)
+ custom_labels = annotation_scheme.get("labels", None)
+ target_schema = annotation_scheme.get("target_schema", "")
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+
+ if scale_type == "likert":
+ labels = custom_labels if custom_labels else DEFAULT_LABELS[:scale_points]
+ # Pad labels if fewer than scale points
+ while len(labels) < scale_points:
+ labels.append(f"Level {len(labels) + 1}")
+ else:
+ labels = []
+
+ html = f"""
+
+ """
+
+ key_bindings = []
+ logger.info(f"Generated confidence layout for {schema_name} (type={scale_type})")
+ return html, key_bindings
diff --git a/potato/server_utils/schemas/conjoint.py b/potato/server_utils/schemas/conjoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..a4685252464858916b9f8cf48ae971c0ef05980b
--- /dev/null
+++ b/potato/server_utils/schemas/conjoint.py
@@ -0,0 +1,168 @@
+"""
+Discrete Choice / Conjoint Analysis Layout
+
+Present annotators with 2-4 product/concept profiles defined by attribute-level
+combinations and ask them to choose the preferred one (or "none"). Enables estimation
+of attribute importance through experimental design.
+
+Research: Green & Srinivasan (1990); Louviere, Flynn & Marley (2015).
+"""
+
+import json
+import logging
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_PROFILES_PER_SET = 3
+
+
+def generate_conjoint_layout(annotation_scheme):
+ """
+ Generate HTML for a Discrete Choice / Conjoint Analysis interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - profiles_per_set: Number of profiles to show (2-4)
+ - attributes: List of {name, levels} dicts
+ - show_none_option: Whether to show "None of these" option
+ - profiles_field: Data field with pre-specified profiles (null = generate)
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_conjoint_layout_internal)
+
+
+def _generate_conjoint_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme['name']
+ description = annotation_scheme['description']
+ profiles_per_set = annotation_scheme.get('profiles_per_set', DEFAULT_PROFILES_PER_SET)
+ attributes = annotation_scheme.get('attributes', [])
+ show_none_option = annotation_scheme.get('show_none_option', True)
+ profiles_field = annotation_scheme.get('profiles_field', None)
+
+ if not attributes and not profiles_field:
+ raise ValueError(f"conjoint schema '{schema_name}' requires 'attributes' or 'profiles_field'")
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+ identifiers = generate_element_identifier(schema_name, schema_name, "radio")
+
+ config_json = json.dumps({
+ 'profiles_per_set': profiles_per_set,
+ 'attributes': attributes,
+ 'profiles_field': profiles_field,
+ })
+
+ # Build profile cards (placeholder structure โ populated by JS from data or generated)
+ profile_cards_html = ""
+ for i in range(profiles_per_set):
+ profile_num = i + 1
+ radio_id = f"{identifiers['id']}-profile-{profile_num}"
+
+ # Build attribute rows placeholder
+ attr_rows = ""
+ for attr in attributes:
+ attr_rows += f"""
+
+ {escape_html_content(attr['name'])}
+ โ
+
+ """
+
+ profile_cards_html += f"""
+
+
+
+
+
+
+ Choose this
+
+
+
+ """
+
+ none_option_html = ""
+ if show_none_option:
+ none_radio_id = f"{identifiers['id']}-none"
+ none_option_html = f"""
+
+
+
+ None of these
+
+
+ """
+
+ html = f"""
+
+
+
+ """
+
+ logger.info(f"Generated conjoint layout for {schema_name}")
+ return html, []
diff --git a/potato/server_utils/schemas/constant_sum.py b/potato/server_utils/schemas/constant_sum.py
new file mode 100644
index 0000000000000000000000000000000000000000..802c367309c676d4f07ba7b8a966d5e111e880fe
--- /dev/null
+++ b/potato/server_utils/schemas/constant_sum.py
@@ -0,0 +1,222 @@
+"""
+Constant Sum / Points Allocation Layout
+
+Generates number inputs (or sliders) constrained to sum to a fixed total.
+Forces annotators to make relative comparisons between categories.
+
+Research basis:
+- Louviere et al. (2015) "Best-Worst Scaling: Theory, Methods and Applications"
+ Cambridge University Press
+- Thurstone (1927) paired-comparison law of comparative judgment
+"""
+
+import logging
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+ generate_tooltip_html,
+)
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_TOTAL_POINTS = 100
+DEFAULT_MIN_PER_ITEM = 0
+DEFAULT_INPUT_TYPE = "number"
+
+
+def generate_constant_sum_layout(annotation_scheme):
+ """
+ Generate HTML for a constant sum / points allocation interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - labels: List of category names
+ - total_points: Sum budget (default 100)
+ - min_per_item: Minimum per item (default 0)
+ - input_type: "number" or "slider" (default "number")
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_constant_sum_layout_internal)
+
+
+def _generate_constant_sum_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme["name"]
+ safe_schema = escape_html_content(schema_name)
+ description = annotation_scheme["description"]
+ labels = annotation_scheme.get("labels", [])
+ total_points = annotation_scheme.get("total_points", DEFAULT_TOTAL_POINTS)
+ min_per_item = annotation_scheme.get("min_per_item", DEFAULT_MIN_PER_ITEM)
+ input_type = annotation_scheme.get("input_type", DEFAULT_INPUT_TYPE)
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+
+ if not labels:
+ raise ValueError(f"constant_sum schema '{schema_name}' requires 'labels'")
+
+ # Normalize labels
+ label_names = []
+ for lbl in labels:
+ if isinstance(lbl, str):
+ label_names.append(lbl)
+ elif isinstance(lbl, dict) and "name" in lbl:
+ label_names.append(lbl["name"])
+ else:
+ raise ValueError(f"Invalid label format: {lbl}")
+
+ html = f"""
+
+
+ """
+
+ key_bindings = []
+ logger.info(f"Generated constant_sum layout for {schema_name} with {len(label_names)} categories")
+ return html, key_bindings
diff --git a/potato/server_utils/schemas/coreference.py b/potato/server_utils/schemas/coreference.py
new file mode 100644
index 0000000000000000000000000000000000000000..cbd5bae3fa23150027bc5dfc96a94146935294bb
--- /dev/null
+++ b/potato/server_utils/schemas/coreference.py
@@ -0,0 +1,196 @@
+"""
+Coreference Chain Annotation Layout
+
+Generates the UI for creating and managing coreference chains โ
+groupings of text spans that refer to the same entity.
+
+This schema type works in conjunction with a span annotation schema.
+A coreference chain is an n-ary undirected link where span_ids lists
+all mentions of the same entity. It leverages the existing SpanLink
+infrastructure with a specialized chain management UI.
+"""
+
+import logging
+import json
+from .identifier_utils import (
+ safe_generate_layout,
+ escape_html_content,
+)
+
+logger = logging.getLogger(__name__)
+
+# Default colors for coreference chains
+CHAIN_COLOR_PALETTE = [
+ "#6E56CF", # Purple
+ "#EF4444", # Red
+ "#22C55E", # Green
+ "#3B82F6", # Blue
+ "#F59E0B", # Amber
+ "#EC4899", # Pink
+ "#06B6D4", # Cyan
+ "#F97316", # Orange
+ "#8B5CF6", # Violet
+ "#10B981", # Emerald
+ "#DC2626", # Dark red
+ "#A855F7", # Light purple
+ "#14B8A6", # Teal
+ "#F43F5E", # Rose
+ "#84CC16", # Lime
+]
+
+
+def _generate_coreference_layout_internal(annotation_scheme, horizontal=False):
+ """
+ Internal function to generate coreference chain layout.
+
+ Args:
+ annotation_scheme: Configuration dictionary containing:
+ - name: Schema name
+ - description: Description shown to user
+ - span_schema: Name of the span schema providing mentions
+ - entity_types: Optional list of entity types for chain classification
+ - allow_singletons: Whether single-mention chains are allowed (default: True)
+ - visual_display:
+ - highlight_mode: "bracket" | "background" | "underline" (default: "background")
+
+ Returns:
+ tuple: (HTML string, key bindings list)
+ """
+ scheme_name = annotation_scheme["name"]
+ description = annotation_scheme.get("description", "Create coreference chains")
+ span_schema = annotation_scheme.get("span_schema", "")
+ entity_types = annotation_scheme.get("entity_types", [])
+ allow_singletons = annotation_scheme.get("allow_singletons", True)
+ visual_display = annotation_scheme.get("visual_display", {})
+ highlight_mode = visual_display.get("highlight_mode", "background")
+
+ # Build entity type selector HTML
+ entity_types_html = ""
+ if entity_types:
+ for i, etype in enumerate(entity_types):
+ if isinstance(etype, dict):
+ etype_name = etype.get("name", f"Entity_{i}")
+ etype_color = etype.get("color", CHAIN_COLOR_PALETTE[i % len(CHAIN_COLOR_PALETTE)])
+ else:
+ etype_name = str(etype)
+ etype_color = CHAIN_COLOR_PALETTE[i % len(CHAIN_COLOR_PALETTE)]
+
+ entity_types_html += f"""
+
+
+
+
+ {escape_html_content(etype_name)}
+
+
+ """
+
+ # Config data for JS
+ config_data = json.dumps({
+ "schemaName": scheme_name,
+ "spanSchema": span_schema,
+ "entityTypes": entity_types if entity_types else [],
+ "allowSingletons": allow_singletons,
+ "highlightMode": highlight_mode,
+ "colors": CHAIN_COLOR_PALETTE,
+ })
+
+ entity_type_section = ""
+ if entity_types:
+ entity_type_section = f"""
+
+
Entity Type:
+
+ {entity_types_html}
+
+
+ """
+
+ schematic = f"""
+
+
+
+
+ {entity_type_section}
+
+
+
+
+
+
+
+
+ New Chain
+
+
+
+
+
+
+ Add to Chain
+
+
+ Merge Chains
+
+
+ Remove Mention
+
+
+
+
+
+
+
No coreference chains created yet.
+ Select spans and click "New Chain" to start.
+
+
+
+
+
+
+ """
+
+ key_bindings = []
+ return schematic, key_bindings
+
+
+def generate_coreference_layout(annotation_scheme, horizontal=False):
+ """
+ Generate coreference chain layout HTML.
+
+ Args:
+ annotation_scheme (dict): The annotation scheme configuration
+ horizontal (bool): Whether to display horizontally
+
+ Returns:
+ tuple: (HTML string, key bindings list)
+ """
+ return safe_generate_layout(
+ annotation_scheme, _generate_coreference_layout_internal, horizontal
+ )
diff --git a/potato/server_utils/schemas/error_span.py b/potato/server_utils/schemas/error_span.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb99896938e78b1c68fb2f2ed5ff8b017b6ba5b3
--- /dev/null
+++ b/potato/server_utils/schemas/error_span.py
@@ -0,0 +1,359 @@
+"""
+Error Span with Typed Severity Layout
+
+Mark error spans in text, assign each an error type from a configurable taxonomy
+and a severity level. Computes an overall quality score. This is the MQM
+(Multidimensional Quality Metrics) annotation workflow.
+
+Research: Lommel et al. "Multidimensional Quality Metrics" (themqm.org); WMT 2024 ESA.
+"""
+
+import json
+import logging
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_SEVERITIES = [
+ {"name": "Minor", "weight": -1},
+ {"name": "Major", "weight": -5},
+ {"name": "Critical", "weight": -10},
+]
+DEFAULT_MAX_SCORE = 100
+
+
+def generate_error_span_layout(annotation_scheme):
+ """
+ Generate HTML for an Error Span with Typed Severity interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - error_types: List of {name, subtypes?} dicts
+ - severities: List of {name, weight} dicts
+ - show_score: Whether to show quality score
+ - max_score: Maximum quality score
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_error_span_layout_internal)
+
+
+def _generate_error_span_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme['name']
+ description = annotation_scheme['description']
+ error_types = annotation_scheme.get('error_types', [])
+ severities = annotation_scheme.get('severities', DEFAULT_SEVERITIES)
+ show_score = annotation_scheme.get('show_score', True)
+ max_score = annotation_scheme.get('max_score', DEFAULT_MAX_SCORE)
+
+ if not error_types:
+ raise ValueError(f"error_span schema '{schema_name}' requires 'error_types'")
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+ identifiers = generate_element_identifier(schema_name, schema_name, "hidden")
+
+ # Serialize config for JS
+ config_json = json.dumps({
+ 'error_types': error_types,
+ 'severities': severities,
+ 'max_score': max_score,
+ 'show_score': show_score,
+ })
+
+ # Build error type options HTML for the popup
+ type_options = ""
+ for et in error_types:
+ subtypes = et.get('subtypes', [])
+ if subtypes:
+ type_options += f''
+ for st in subtypes:
+ type_options += f'{escape_html_content(st)} '
+ type_options += ' '
+ else:
+ type_options += f'{escape_html_content(et["name"])} '
+
+ # Build severity radio buttons for the popup
+ severity_radios = ""
+ for sev in severities:
+ severity_radios += f"""
+
+
+ {escape_html_content(sev['name'])} ({sev['weight']:+d})
+
+ """
+
+ score_display = ""
+ if show_score:
+ score_display = f"""
+
+ Score: {max_score} / {max_score}
+
+ """
+
+ html = f"""
+
+
+
+ """
+
+ logger.info(f"Generated error span layout for {schema_name}")
+ return html, []
diff --git a/potato/server_utils/schemas/event_annotation.py b/potato/server_utils/schemas/event_annotation.py
new file mode 100644
index 0000000000000000000000000000000000000000..baccba6c5f50c9fd08adf407cc9cfe77d38fd3f5
--- /dev/null
+++ b/potato/server_utils/schemas/event_annotation.py
@@ -0,0 +1,222 @@
+"""
+Event Annotation Layout
+
+Generates the UI for N-ary event annotation with triggers and typed arguments.
+This schema type works in conjunction with a span annotation schema to allow
+users to annotate events like "ATTACK(attacker=John, target=Mary, weapon=knife)".
+"""
+
+import logging
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ escape_html_content,
+ generate_tooltip_html
+)
+from .span import get_span_color, SPAN_COLOR_PALETTE
+
+logger = logging.getLogger(__name__)
+
+# Default colors for event types
+EVENT_COLOR_PALETTE = [
+ "#dc2626", # Red
+ "#2563eb", # Blue
+ "#16a34a", # Green
+ "#9333ea", # Purple
+ "#ea580c", # Orange
+ "#0891b2", # Cyan
+ "#c026d3", # Fuchsia
+ "#ca8a04", # Yellow
+ "#4f46e5", # Indigo
+ "#059669", # Emerald
+]
+
+
+def _generate_event_annotation_layout_internal(annotation_scheme, horizontal=False):
+ """
+ Internal function to generate event annotation layout after validation.
+
+ Args:
+ annotation_scheme: Configuration dictionary containing:
+ - name: Schema name
+ - description: Description shown to user
+ - span_schema: Name of the span schema for entities
+ - event_types: List of event type definitions with:
+ - type: Event type name (e.g., "ATTACK", "HIRE")
+ - color: Optional color for this event type
+ - trigger_labels: Optional list of span labels that can be triggers
+ - arguments: List of argument definitions with:
+ - role: Role name (e.g., "attacker", "target")
+ - entity_types: Optional list of allowed entity types
+ - required: Whether this argument is required (default: false)
+ horizontal: Whether to display horizontally (not used for events)
+
+ Returns:
+ tuple: (HTML string, key bindings list)
+ """
+ scheme_name = annotation_scheme["name"]
+ description = annotation_scheme.get("description", "Annotate events with triggers and arguments")
+ span_schema = annotation_scheme.get("span_schema", "")
+ event_types = annotation_scheme.get("event_types", [])
+ visual_display = annotation_scheme.get("visual_display", {})
+
+ # Build event types HTML
+ event_types_html = ""
+ key_bindings = []
+
+ for i, event_type in enumerate(event_types):
+ type_name = event_type.get("type", f"Event_{i}")
+ color = event_type.get("color", EVENT_COLOR_PALETTE[i % len(EVENT_COLOR_PALETTE)])
+ trigger_labels = event_type.get("trigger_labels", [])
+ arguments = event_type.get("arguments", [])
+
+ # Build arguments data for JavaScript
+ args_data = []
+ for arg in arguments:
+ args_data.append({
+ "role": arg.get("role", ""),
+ "entity_types": arg.get("entity_types", []),
+ "required": arg.get("required", False)
+ })
+
+ import json
+ args_json = json.dumps(args_data)
+
+ # Tooltip with argument info
+ tooltip_parts = []
+ if trigger_labels:
+ tooltip_parts.append(f"Triggers: {', '.join(trigger_labels)}")
+ if arguments:
+ arg_strs = []
+ for arg in arguments:
+ role = arg.get("role", "")
+ req = "(required)" if arg.get("required", False) else "(optional)"
+ entity_types = arg.get("entity_types", [])
+ if entity_types:
+ arg_strs.append(f"{role} {req}: {', '.join(entity_types)}")
+ else:
+ arg_strs.append(f"{role} {req}")
+ tooltip_parts.append("Arguments: " + "; ".join(arg_strs))
+
+ tooltip_attr = ""
+ if tooltip_parts:
+ tooltip_text = " | ".join(tooltip_parts)
+ tooltip_attr = f'data-toggle="tooltip" data-placement="top" title="{escape_html_content(tooltip_text)}"'
+
+ event_types_html += f"""
+
+
+
+
+ {escape_html_content(type_name)}
+
+
+ """
+
+ # Visual display settings
+ show_arcs = visual_display.get("enabled", True)
+ arc_position = visual_display.get("arc_position", "above")
+ show_labels = visual_display.get("show_labels", True)
+
+ schematic = f"""
+
+
+
+
+
+
+
1. Select Event Type:
+
+ {event_types_html}
+
+
+
+
+
+
2. Select Trigger Span:
+
+
Click on a span to set it as the event trigger
+
+
+
+
+
+
3. Assign Arguments:
+
+
+
+
+
+
+
+
+
+
+
+
+
+ Create Event
+
+
+ Cancel
+
+
+
+
+
+
Existing Events:
+
+
No events created yet
+
+
+
+
+
+
+
+ Show event arcs above text
+
+
+
+
+
+
+ """
+
+ return schematic, key_bindings
+
+
+def generate_event_annotation_layout(annotation_scheme, horizontal=False):
+ """
+ Generate event annotation layout HTML for the given annotation scheme.
+
+ Args:
+ annotation_scheme (dict): The annotation scheme configuration
+ horizontal (bool): Whether to display horizontally
+
+ Returns:
+ tuple: (HTML string, key bindings list)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_event_annotation_layout_internal, horizontal)
diff --git a/potato/server_utils/schemas/extractive_qa.py b/potato/server_utils/schemas/extractive_qa.py
new file mode 100644
index 0000000000000000000000000000000000000000..84737e6755436838a7206f748ad796f6be196b91
--- /dev/null
+++ b/potato/server_utils/schemas/extractive_qa.py
@@ -0,0 +1,226 @@
+"""
+Extractive QA / Answer Span Layout
+
+Display a question and a passage; annotator highlights the answer span in the passage.
+A streamlined SQuAD-style workflow that combines question display with directed span selection.
+
+Research: Rajpurkar et al. (2016) "SQuAD"; Kwiatkowski et al. (2019) "Natural Questions".
+"""
+
+import logging
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_HIGHLIGHT_COLOR = "#FFEB3B"
+DEFAULT_ALLOW_UNANSWERABLE = True
+
+
+def generate_extractive_qa_layout(annotation_scheme):
+ """
+ Generate HTML for an Extractive QA interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - question_field: Field in data containing the question
+ - passage_field: Field in data containing the passage
+ - allow_unanswerable: Whether to show "Unanswerable" button
+ - highlight_color: Color for answer highlight
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_extractive_qa_layout_internal)
+
+
+def _generate_extractive_qa_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme['name']
+ description = annotation_scheme['description']
+ question_field = annotation_scheme.get('question_field', 'question')
+ passage_field = annotation_scheme.get('passage_field', '')
+ allow_unanswerable = annotation_scheme.get('allow_unanswerable', DEFAULT_ALLOW_UNANSWERABLE)
+ highlight_color = annotation_scheme.get('highlight_color', DEFAULT_HIGHLIGHT_COLOR)
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+ identifiers = generate_element_identifier(schema_name, schema_name, "hidden")
+
+ unanswerable_btn = ""
+ if allow_unanswerable:
+ unanswerable_btn = f"""
+
+ Unanswerable
+
+ """
+
+ html = f"""
+
+
+
+ """
+
+ logger.info(f"Generated extractive QA layout for {schema_name}")
+ return html, []
diff --git a/potato/server_utils/schemas/hierarchical_multiselect.py b/potato/server_utils/schemas/hierarchical_multiselect.py
new file mode 100644
index 0000000000000000000000000000000000000000..4748c06a9bd9b43551e42ea923762b03f605aca7
--- /dev/null
+++ b/potato/server_utils/schemas/hierarchical_multiselect.py
@@ -0,0 +1,334 @@
+"""
+Hierarchical Multi-Label Selection Layout
+
+Generates an expandable/collapsible tree of checkboxes for hierarchical
+taxonomy labeling. Stores selections as a hidden input with comma-separated values.
+
+Research basis:
+- Silla & Freitas (2011) "A Survey of Hierarchical Classification Across Different
+ Application Domains" Data Mining and Knowledge Discovery
+- Vens et al. (2008) "Decision Trees for Hierarchical Multi-label Classification"
+ Machine Learning
+"""
+
+import json
+import logging
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def generate_hierarchical_multiselect_layout(annotation_scheme):
+ """
+ Generate HTML for a hierarchical multi-label selection interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - taxonomy: Nested dict/list defining the hierarchy
+ - auto_select_children: Auto-select children when parent selected (default false)
+ - auto_select_parent: Auto-select parent when all children selected (default false)
+ - show_search: Show search/filter box (default false)
+ - max_selections: Maximum number of selections (null = unlimited)
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_hierarchical_multiselect_layout_internal)
+
+
+def _build_tree_html(taxonomy, schema_name, prefix="", depth=0):
+ """Recursively build tree HTML from taxonomy dict/list."""
+ html = ""
+ safe_schema = escape_html_content(schema_name)
+
+ if isinstance(taxonomy, dict):
+ for key, children in taxonomy.items():
+ safe_key = escape_html_content(key)
+ node_id = f"{prefix}{key}".replace(" ", "_")
+ safe_node_id = escape_html_content(node_id)
+ has_children = bool(children)
+ toggle = '▶ ' if has_children else ' '
+
+ html += f"""
+
+
+ {toggle}
+
+
+ {safe_key}
+
+
+ """
+ if has_children:
+ html += f'
'
+ html += _build_tree_html(children, schema_name, prefix=f"{node_id}.", depth=depth + 1)
+ html += "
"
+ html += "
"
+
+ elif isinstance(taxonomy, list):
+ for item in taxonomy:
+ safe_item = escape_html_content(str(item))
+ node_id = f"{prefix}{item}".replace(" ", "_")
+ safe_node_id = escape_html_content(node_id)
+
+ html += f"""
+
+
+
+
+
+ {safe_item}
+
+
+
+ """
+
+ return html
+
+
+def _generate_hierarchical_multiselect_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme["name"]
+ safe_schema = escape_html_content(schema_name)
+ description = annotation_scheme["description"]
+ taxonomy = annotation_scheme.get("taxonomy", {})
+ auto_select_children = annotation_scheme.get("auto_select_children", False)
+ auto_select_parent = annotation_scheme.get("auto_select_parent", False)
+ show_search = annotation_scheme.get("show_search", False)
+ max_selections = annotation_scheme.get("max_selections", None)
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+
+ if not taxonomy:
+ raise ValueError(f"hierarchical_multiselect schema '{schema_name}' requires 'taxonomy'")
+
+ identifiers = generate_element_identifier(schema_name, "selected_labels", "hidden")
+
+ html = f"""
+
+
+ """
+
+ key_bindings = []
+ logger.info(f"Generated hierarchical_multiselect layout for {schema_name}")
+ return html, key_bindings
diff --git a/potato/server_utils/schemas/identifier_utils.py b/potato/server_utils/schemas/identifier_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..a8c28a6c3759a422a294ed32f0ef6482b5104fff
--- /dev/null
+++ b/potato/server_utils/schemas/identifier_utils.py
@@ -0,0 +1,394 @@
+"""
+Identifier Utilities for Schema Generation
+
+This module provides centralized functions for generating consistent identifiers
+and validating schema configurations across all annotation schema types.
+"""
+
+import html
+import logging
+from collections.abc import Mapping
+from typing import Dict, Any, Tuple, List
+
+logger = logging.getLogger(__name__)
+
+def validate_schema_config(annotation_scheme: dict) -> bool:
+ """
+ Validate schema configuration before generating HTML.
+
+ Args:
+ annotation_scheme: Schema configuration dictionary
+
+ Returns:
+ bool: True if valid, raises exception if invalid
+
+ Raises:
+ ValueError: If configuration is invalid
+ """
+ # Check required fields
+ required_fields = ["name", "description"]
+ for field in required_fields:
+ if field not in annotation_scheme:
+ raise ValueError(f"Missing required field: {field}")
+
+ # Validate schema name
+ schema_name = annotation_scheme["name"]
+ if not schema_name or not str(schema_name).strip():
+ raise ValueError("Schema name cannot be empty")
+
+ # Validate description
+ description = annotation_scheme["description"]
+ if not description or not str(description).strip():
+ raise ValueError("Schema description cannot be empty")
+
+ # Validate labels if present
+ if "labels" in annotation_scheme:
+ labels = annotation_scheme["labels"]
+ if not labels:
+ raise ValueError("Labels list cannot be empty")
+
+ # Check for duplicate labels
+ label_names = []
+ for label in labels:
+ if isinstance(label, str):
+ label_names.append(label.strip())
+ elif isinstance(label, dict) and "name" in label:
+ label_names.append(label["name"].strip())
+ else:
+ raise ValueError(f"Invalid label format: {label}")
+
+ # Check for empty labels
+ if any(not name for name in label_names):
+ raise ValueError("Label names cannot be empty")
+
+ # Check for duplicates
+ if len(label_names) != len(set(label_names)):
+ duplicates = [name for name in set(label_names) if label_names.count(name) > 1]
+ raise ValueError(f"Duplicate labels found: {duplicates}")
+
+ logger.debug(f"Schema configuration validation passed for: {schema_name}")
+ return True
+
+def generate_element_identifier(schema_name: str, label_name: str, element_type: str = "default") -> Dict[str, str]:
+ """
+ Generate consistent identifiers for form elements.
+
+ Args:
+ schema_name: Name of the annotation schema
+ label_name: Name of the specific label/option
+ element_type: Type of element (radio, checkbox, text, etc.)
+
+ Returns:
+ dict: Contains id, name, schema, and label_name attributes
+ """
+ # Sanitize inputs
+ safe_schema = escape_html_content(schema_name.strip())
+ safe_label = escape_html_content(label_name.strip())
+
+ # Generate unique identifier (using underscore to avoid conflicts with CSS selectors)
+ element_id = f"{safe_schema}_{safe_label}_{element_type}".replace(":::", "_")
+
+ # For radio buttons, use schema name as the group name to ensure mutual exclusivity
+ if element_type == "radio":
+ element_name = safe_schema
+ else:
+ element_name = f"{safe_schema}:::{safe_label}"
+
+ return {
+ "id": element_id,
+ "name": element_name,
+ "schema": safe_schema,
+ "label_name": safe_label
+ }
+
+def generate_element_value(label_data: Any, index: int, annotation_scheme: dict) -> str:
+ """
+ Generate consistent value attributes for form elements.
+
+ Args:
+ label_data: Label configuration (string or dict)
+ index: Index of the label in the list
+ annotation_scheme: Full schema configuration
+
+ Returns:
+ str: Value to use for the element
+ """
+ # Handle custom key_value first
+ if isinstance(label_data, dict) and "key_value" in label_data:
+ return str(label_data["key_value"])
+
+ # Handle sequential key binding
+ if annotation_scheme.get("sequential_key_binding"):
+ return str(index % 10)
+
+ # Default to label name
+ if isinstance(label_data, str):
+ return label_data
+ elif isinstance(label_data, dict) and "name" in label_data:
+ return label_data["name"]
+
+ # Fallback to index
+ return str(index)
+
+def escape_html_content(content: str) -> str:
+ """
+ Escape HTML content to prevent injection.
+
+ Args:
+ content: Content to escape
+
+ Returns:
+ str: Escaped content
+ """
+ if not content:
+ return ""
+ return html.escape(str(content))
+
+
+def humanize_label(text: str) -> str:
+ """Turn a machine label (``agent_a_much_better``) into readable text
+ (``Agent A Much Better``) for display only -- the stored annotation
+ value is always the original label name, never this.
+
+ Tokens that are already mixed/upper case or contain digits+letters
+ (acronyms, ``GPT4``, ``v2``) are preserved as-is so we don't mangle
+ them; purely lowercase tokens are capitalized.
+ """
+ if not text:
+ return ""
+ # Never mangle Jinja/template expressions (e.g. dynamic_labels:
+ # "{{instance_obj.labels[0]}}"). Humanizing would rewrite `instance_obj`
+ # to `instance Obj`, producing invalid Jinja and a 500 at render time.
+ if "{{" in str(text) or "{%" in str(text):
+ return str(text)
+ s = str(text).replace("_", " ").replace("-", " ")
+ s = " ".join(s.split()) # collapse whitespace
+ out = []
+ for tok in s.split(" "):
+ if tok.islower():
+ out.append(tok[:1].upper() + tok[1:])
+ else:
+ out.append(tok) # preserve ACRONYMs, GPT4, v2, MixedCase
+ return " ".join(out)
+
+
+def display_label_text(label_data: Any, annotation_scheme: dict) -> str:
+ """Resolve the *visible* text for a label.
+
+ Precedence: explicit ``displayed_label`` on a dict label >
+ humanized name (default, when ``humanize_labels`` is not disabled) >
+ raw name. Stored value is unaffected by this function.
+ """
+ if isinstance(label_data, Mapping):
+ if label_data.get("displayed_label"):
+ return str(label_data["displayed_label"])
+ name = label_data.get("name", "")
+ else:
+ name = label_data
+ if annotation_scheme.get("humanize_labels", True):
+ return humanize_label(name)
+ return str(name)
+
+def safe_generate_layout(annotation_scheme: dict, layout_function: callable, *args, **kwargs) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Safely generate layout with proper error handling.
+
+ Args:
+ annotation_scheme: Schema configuration
+ layout_function: Function to generate layout
+ *args, **kwargs: Additional arguments for the layout function
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ try:
+ # Validate configuration
+ validate_schema_config(annotation_scheme)
+
+ # Generate layout
+ return layout_function(annotation_scheme, *args, **kwargs)
+
+ except Exception as e:
+ schema_name = annotation_scheme.get('name', 'unknown')
+ logger.error(f"Failed to generate layout for schema '{schema_name}': {e}")
+
+ # Return error HTML instead of crashing
+ error_html = f"""
+
+
Error Generating Annotation Form
+
Schema: {escape_html_content(schema_name)}
+
{escape_html_content(str(e))}
+
+ """
+ return error_html, []
+
+def generate_validation_attribute(annotation_scheme: dict, label_name: str = None) -> str:
+ """
+ Generate validation attribute for form elements.
+
+ Args:
+ annotation_scheme: Schema configuration
+ label_name: Specific label name for required_label validation
+
+ Returns:
+ str: Validation attribute value
+ """
+ label_requirement = annotation_scheme.get("label_requirement", {})
+
+ # Normalize: label_requirement: true (bool) โ {"required": true}
+ if isinstance(label_requirement, bool):
+ label_requirement = {"required": True} if label_requirement else {}
+
+ # Support top-level required: true as shorthand for label_requirement.required
+ if not label_requirement and annotation_scheme.get("required") is True:
+ label_requirement = {"required": True}
+
+ # Debug logging
+ logger.debug(f"generate_validation_attribute called with label_requirement: {label_requirement}")
+ logger.debug(f"label_name: {label_name}")
+
+ # Check for required_label validation
+ if label_name and label_requirement.get("required_label"):
+ required_labels = label_requirement["required_label"]
+ if isinstance(required_labels, str) and label_name == required_labels:
+ logger.debug(f"Returning 'required_label' for label: {label_name}")
+ return "required_label"
+ elif isinstance(required_labels, list) and label_name in required_labels:
+ logger.debug(f"Returning 'required_label' for label: {label_name}")
+ return "required_label"
+
+ # Check for general required validation
+ if label_requirement.get("required"):
+ logger.debug(f"Returning 'required' for general requirement")
+ return "required"
+
+ logger.debug(f"Returning empty string - no validation requirements met")
+ return ""
+
+
+def generate_layout_attributes(annotation_scheme: dict) -> str:
+ """
+ Generate layout-related HTML attributes for grid positioning.
+
+ Args:
+ annotation_scheme: Schema configuration that may contain:
+ - layout: dict with layout options
+ - columns: Number of grid columns to span (1-6, default: 1)
+ - rows: Number of grid rows to span (1-4, default: 1)
+ - order: Explicit ordering integer for grid placement
+ - min_width: Minimum width CSS value (e.g., "200px")
+ - max_width: Maximum width CSS value (e.g., "400px")
+ - align_self: Alignment override (start, center, end, stretch)
+
+ Returns:
+ str: HTML attribute string for layout (e.g., 'data-grid-columns="2" data-grid-rows="1"')
+
+ Example config:
+ annotation_schemes:
+ - name: preference
+ description: "Which is better?"
+ layout:
+ columns: 2 # Span 2 columns in the grid
+ rows: 1 # Span 1 row (default)
+ order: 1 # Explicit ordering
+ min_width: "200px"
+ max_width: "400px"
+ align_self: "start"
+ """
+ layout_config = annotation_scheme.get("layout", {})
+ attrs = []
+
+ # Column span (1-6, default: 1)
+ columns = layout_config.get("columns", 1)
+ if not isinstance(columns, int) or columns < 1:
+ columns = 1
+ elif columns > 6:
+ columns = 6
+ attrs.append(f'data-grid-columns="{columns}"')
+
+ # Row span (1-4, default: 1)
+ rows = layout_config.get("rows", 1)
+ if isinstance(rows, int) and rows > 1:
+ rows = min(rows, 4)
+ attrs.append(f'data-grid-rows="{rows}"')
+
+ # Explicit order (integer)
+ order = layout_config.get("order")
+ if isinstance(order, int):
+ attrs.append(f'data-grid-order="{order}"')
+
+ # Min/max width via CSS custom properties in style attribute
+ style_parts = []
+ min_width = layout_config.get("min_width")
+ if min_width and isinstance(min_width, str):
+ style_parts.append(f"--form-min-width: {html.escape(min_width)}")
+
+ max_width = layout_config.get("max_width")
+ if max_width and isinstance(max_width, str):
+ style_parts.append(f"--form-max-width: {html.escape(max_width)}")
+
+ if style_parts:
+ attrs.append(f'style="{"; ".join(style_parts)}"')
+
+ # Align self override
+ align_self = layout_config.get("align_self")
+ valid_alignments = ["start", "center", "end", "stretch"]
+ if align_self and align_self in valid_alignments:
+ attrs.append(f'data-align-self="{align_self}"')
+
+ return " ".join(attrs)
+
+
+def generate_tooltip_html(label_data: Dict[str, Any]) -> str:
+ """
+ Generate tooltip HTML attribute from label data.
+
+ This function provides centralized tooltip generation for all schema types.
+ It checks for tooltip text in the label configuration, either directly or
+ from an external file.
+
+ Args:
+ label_data: Label configuration dictionary that may contain:
+ - tooltip: Direct tooltip text string
+ - tooltip_file: Path to file containing tooltip text
+
+ Returns:
+ str: Tooltip HTML attribute string (e.g., 'data-toggle="tooltip" ...')
+ or empty string if no tooltip is configured
+
+ Example:
+ >>> label_data = {"name": "Option 1", "tooltip": "Select this option"}
+ >>> generate_tooltip_html(label_data)
+ 'data-toggle="tooltip" data-html="true" data-placement="top" title="Select this option"'
+ """
+ if not isinstance(label_data, dict):
+ return ""
+
+ tooltip_text = ""
+
+ # Check for direct tooltip text
+ if "tooltip" in label_data:
+ tooltip_text = label_data["tooltip"]
+ logger.debug(f"Found direct tooltip text for label")
+
+ # Check for tooltip file
+ elif "tooltip_file" in label_data:
+ try:
+ with open(label_data["tooltip_file"], "rt", encoding="utf-8") as f:
+ tooltip_text = "".join(f.readlines())
+ logger.debug(f"Read tooltip from file: {label_data['tooltip_file']}")
+ except FileNotFoundError:
+ logger.error(f"Tooltip file not found: {label_data['tooltip_file']}")
+ return ""
+ except PermissionError:
+ logger.error(f"Permission denied reading tooltip file: {label_data['tooltip_file']}")
+ return ""
+ except Exception as e:
+ logger.error(f"Failed to read tooltip file '{label_data['tooltip_file']}': {e}")
+ return ""
+
+ if tooltip_text:
+ escaped_tooltip = escape_html_content(tooltip_text)
+ return f'data-toggle="tooltip" data-html="true" data-placement="top" title="{escaped_tooltip}"'
+
+ return ""
\ No newline at end of file
diff --git a/potato/server_utils/schemas/image_annotation.py b/potato/server_utils/schemas/image_annotation.py
new file mode 100644
index 0000000000000000000000000000000000000000..6b2aca84853c9845ca956c9fdc8a98e95ef67d4a
--- /dev/null
+++ b/potato/server_utils/schemas/image_annotation.py
@@ -0,0 +1,610 @@
+"""
+Image Annotation Layout
+
+Generates a form interface for annotating images with:
+- Bounding boxes (rectangular regions)
+- Polygons (arbitrary shapes)
+- Freeform drawing (brush strokes)
+- Landmarks (point annotations)
+
+Uses Fabric.js for canvas-based annotation with zoom/pan support.
+"""
+
+import logging
+import json
+from .identifier_utils import (
+ safe_generate_layout,
+ escape_html_content
+)
+
+logger = logging.getLogger(__name__)
+
+# Default colors for labels if not specified
+DEFAULT_COLORS = [
+ "#FF6B6B", # Red
+ "#4ECDC4", # Teal
+ "#45B7D1", # Blue
+ "#96CEB4", # Green
+ "#FFEAA7", # Yellow
+ "#DDA0DD", # Plum
+ "#98D8C8", # Mint
+ "#F7DC6F", # Gold
+ "#BB8FCE", # Purple
+ "#85C1E9", # Light Blue
+]
+
+# Valid annotation tools
+VALID_TOOLS = ["bbox", "polygon", "freeform", "landmark", "fill", "eraser", "brush"]
+
+
+def generate_image_annotation_layout(annotation_scheme):
+ """
+ Generate HTML for an image annotation interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - tools: List of tools to enable (bbox, polygon, freeform, landmark)
+ - labels: List of label definitions with name and optional color
+ - zoom_enabled: Whether to enable zoom (default: True)
+ - pan_enabled: Whether to enable pan (default: True)
+ - min_annotations: Minimum required annotations (default: 0)
+ - max_annotations: Maximum allowed annotations (default: null/unlimited)
+ - freeform_brush_size: Brush size for freeform tool (default: 5)
+ - freeform_simplify: Whether to simplify freeform paths (default: True)
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the image annotation interface
+ key_bindings: List of keyboard shortcuts
+
+ Raises:
+ ValueError: If required fields are missing or invalid
+ """
+ return safe_generate_layout(annotation_scheme, _generate_image_annotation_layout_internal)
+
+
+def _generate_image_annotation_layout_internal(annotation_scheme):
+ """
+ Internal function to generate image annotation layout after validation.
+ """
+ schema_name = annotation_scheme.get('name', 'image_annotation')
+ logger.debug(f"Generating image annotation layout for schema: {schema_name}")
+
+ # Validate required fields
+ if "labels" not in annotation_scheme:
+ error_msg = f"Missing labels in schema: {schema_name}"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ if "tools" not in annotation_scheme:
+ error_msg = f"Missing tools in schema: {schema_name}"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ # Validate tools
+ tools = annotation_scheme["tools"]
+ if not isinstance(tools, list) or not tools:
+ error_msg = f"tools must be a non-empty list in schema: {schema_name}"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ invalid_tools = [t for t in tools if t not in VALID_TOOLS]
+ if invalid_tools:
+ error_msg = f"Invalid tools: {invalid_tools}. Valid tools are: {VALID_TOOLS}"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ # Process labels with colors
+ labels = _process_labels(annotation_scheme["labels"])
+
+ # Get configuration options
+ zoom_enabled = annotation_scheme.get("zoom_enabled", True)
+ pan_enabled = annotation_scheme.get("pan_enabled", True)
+ min_annotations = annotation_scheme.get("min_annotations", 0)
+ max_annotations = annotation_scheme.get("max_annotations", None)
+ freeform_brush_size = annotation_scheme.get("freeform_brush_size", 5)
+ freeform_simplify = annotation_scheme.get("freeform_simplify", True)
+
+ # Segmentation mask configuration
+ brush_size = annotation_scheme.get("brush_size", 20)
+ eraser_size = annotation_scheme.get("eraser_size", 20)
+ mask_opacity = annotation_scheme.get("mask_opacity", 0.5)
+
+ # AI support configuration
+ ai_support = annotation_scheme.get("ai_support", {})
+ ai_enabled = ai_support.get("enabled", False)
+
+ # source_field: Links this annotation schema to a display field from instance_display
+ source_field = annotation_scheme.get("source_field", "")
+
+ # Build config object for JavaScript
+ js_config = {
+ "schemaName": schema_name,
+ "tools": tools,
+ "labels": labels,
+ "zoomEnabled": zoom_enabled,
+ "panEnabled": pan_enabled,
+ "minAnnotations": min_annotations,
+ "maxAnnotations": max_annotations,
+ "freeformBrushSize": freeform_brush_size,
+ "freeformSimplify": freeform_simplify,
+ "brushSize": brush_size,
+ "eraserSize": eraser_size,
+ "maskOpacity": mask_opacity,
+ "aiSupport": ai_enabled,
+ "aiFeatures": ai_support.get("features", {}) if ai_enabled else {},
+ "sourceField": source_field,
+ }
+
+ # Generate HTML
+ html = _generate_html(annotation_scheme, js_config, schema_name, labels, tools, ai_enabled, ai_support)
+
+ # Generate keybindings
+ keybindings = _generate_keybindings(labels, tools)
+
+ logger.info(f"Successfully generated image annotation layout for {schema_name}")
+ return html, keybindings
+
+
+def _process_labels(labels_config):
+ """
+ Process label configuration and assign colors.
+
+ Args:
+ labels_config: List of label configs (strings or dicts)
+
+ Returns:
+ List of processed label dicts with name, color, and optional key_value
+ """
+ processed = []
+ for i, label in enumerate(labels_config):
+ if isinstance(label, str):
+ processed.append({
+ "name": label,
+ "color": DEFAULT_COLORS[i % len(DEFAULT_COLORS)],
+ })
+ elif isinstance(label, dict):
+ processed.append({
+ "name": label.get("name", f"label_{i}"),
+ "color": label.get("color", DEFAULT_COLORS[i % len(DEFAULT_COLORS)]),
+ "key_value": label.get("key_value"),
+ })
+ else:
+ processed.append({
+ "name": str(label),
+ "color": DEFAULT_COLORS[i % len(DEFAULT_COLORS)],
+ })
+ return processed
+
+
+def _generate_html(annotation_scheme, js_config, schema_name, labels, tools, ai_enabled=False, ai_support=None):
+ """
+ Generate the HTML for the image annotation interface.
+ """
+ escaped_name = escape_html_content(schema_name)
+ description = escape_html_content(annotation_scheme.get('description', ''))
+ config_json = json.dumps(js_config)
+
+ # source_field attribute for linking to display fields
+ source_field = annotation_scheme.get("source_field", "")
+ source_field_attr = f' data-source-field="{escape_html_content(source_field)}"' if source_field else ""
+
+ # Generate tool buttons
+ tool_buttons = _generate_tool_buttons(tools)
+
+ # Generate label selector
+ label_selector = _generate_label_selector(labels)
+
+ # Generate AI toolbar if enabled
+ ai_toolbar_html = ""
+ ai_init_script = ""
+ if ai_enabled:
+ ai_features = ai_support.get("features", {}) if ai_support else {}
+ ai_toolbar_html = _generate_ai_toolbar(ai_features)
+ ai_init_script = _generate_ai_init_script(escaped_name)
+
+ html = f'''
+
+ '''
+
+ return html
+
+
+def _generate_ai_toolbar(ai_features):
+ """
+ Generate HTML for the AI assistance toolbar.
+ """
+ # Determine which buttons to show based on features
+ detection_enabled = ai_features.get("detection", True)
+ pre_annotate_enabled = ai_features.get("pre_annotate", True)
+ classification_enabled = ai_features.get("classification", False)
+ hint_enabled = ai_features.get("hint", True)
+
+ buttons = []
+
+ if detection_enabled:
+ buttons.append(
+ ''
+ '๐ Detect '
+ )
+
+ if pre_annotate_enabled:
+ buttons.append(
+ ''
+ 'โก Auto '
+ )
+
+ if classification_enabled:
+ buttons.append(
+ ''
+ '๐ท๏ธ Classify '
+ )
+
+ if hint_enabled:
+ buttons.append(
+ ''
+ '๐ก Hint '
+ )
+
+ if not buttons:
+ return ""
+
+ return f'''
+
+
+
+
+ '''
+
+
+def _generate_ai_init_script(escaped_name):
+ """
+ Generate JavaScript initialization code for AI assistant.
+ """
+ return f'''
+ // Initialize AI assistant if enabled and VisualAIAssistantManager is available
+ if (config.aiSupport && typeof VisualAIAssistantManager !== 'undefined') {{
+ var annotationId = Array.from(document.querySelectorAll('.annotation-form')).indexOf(
+ document.getElementById('{escaped_name}')
+ );
+ container.aiAssistant = new VisualAIAssistantManager({{
+ annotationType: 'image_annotation',
+ annotationId: annotationId >= 0 ? annotationId : 0,
+ annotationManager: manager
+ }});
+ }}
+ '''
+
+
+def _generate_tool_buttons(tools):
+ """
+ Generate HTML for tool selection buttons.
+ """
+ tool_info = {
+ "bbox": {"label": "Box", "title": "Bounding Box (B)", "icon": "โก"},
+ "polygon": {"label": "Polygon", "title": "Polygon (P)", "icon": "โฌก"},
+ "freeform": {"label": "Draw", "title": "Freeform Draw (F)", "icon": "โ"},
+ "landmark": {"label": "Point", "title": "Landmark Point (L)", "icon": "โ"},
+ "brush": {"label": "Brush", "title": "Segmentation Brush (M)", "icon": "๐๏ธ"},
+ "fill": {"label": "Fill", "title": "Flood Fill (G)", "icon": "๐ชฃ"},
+ "eraser": {"label": "Eraser", "title": "Eraser (E)", "icon": "โซ"},
+ }
+
+ buttons = []
+ for tool in tools:
+ info = tool_info.get(tool, {"label": tool, "title": tool, "icon": "?"})
+ buttons.append(
+ f''
+ f'{info["icon"]} {info["label"]} '
+ )
+
+ return "\n".join(buttons)
+
+
+def _generate_label_selector(labels):
+ """
+ Generate HTML for label selection buttons.
+ """
+ buttons = []
+ for label in labels:
+ name = escape_html_content(label["name"])
+ color = label["color"]
+ key_hint = f' ({label["key_value"]})' if label.get("key_value") else ""
+ buttons.append(
+ f''
+ f' '
+ f'{name} '
+ )
+
+ return "\n".join(buttons)
+
+
+def _generate_keybindings(labels, tools):
+ """
+ Generate keybinding list for the schema.
+ """
+ keybindings = []
+
+ # Tool shortcuts
+ tool_keys = {
+ "bbox": ("b", "Bounding Box tool"),
+ "polygon": ("p", "Polygon tool"),
+ "freeform": ("f", "Freeform draw tool"),
+ "landmark": ("l", "Landmark point tool"),
+ "brush": ("m", "Segmentation brush tool"),
+ "fill": ("g", "Flood fill tool"),
+ "eraser": ("e", "Eraser tool"),
+ }
+ for tool in tools:
+ if tool in tool_keys:
+ keybindings.append(tool_keys[tool])
+
+ # Label shortcuts
+ for label in labels:
+ if label.get("key_value"):
+ keybindings.append((label["key_value"], f"Select label: {label['name']}"))
+
+ # Common shortcuts
+ keybindings.extend([
+ ("Del", "Delete selected"),
+ ("+/-", "Zoom in/out"),
+ ("0", "Fit to view"),
+ ])
+
+ return keybindings
diff --git a/potato/server_utils/schemas/keybinding_allocator.py b/potato/server_utils/schemas/keybinding_allocator.py
new file mode 100644
index 0000000000000000000000000000000000000000..152954df8fe535ea62fb098a83448265db14cf6c
--- /dev/null
+++ b/potato/server_utils/schemas/keybinding_allocator.py
@@ -0,0 +1,247 @@
+"""
+Keybinding Allocator
+
+Centralized allocation of non-conflicting keyboard shortcuts across all
+annotation schemas. When multiple schemas use sequential_key_binding: true,
+this module assigns keys from separate pools so they don't overlap.
+
+Key pools (QWERTY layout):
+ Pool 0: 1 2 3 4 5 6 7 8 9 0 (number row)
+ Pool 1: q w e r t y u i o p (top letter row)
+ Pool 2: a s d f g h j k l (home row)
+
+Schemas that self-manage keys (pairwise, bws) pre-claim their hardcoded keys.
+Explicit per-label key_value overrides are always honored.
+"""
+
+import logging
+from collections.abc import Mapping
+
+logger = logging.getLogger(__name__)
+
+KEY_POOLS = [
+ ['1', '2', '3', '4', '5', '6', '7', '8', '9', '0'],
+ ['q', 'w', 'e', 'r', 't', 'y', 'u', 'i', 'o', 'p'],
+ ['a', 's', 'd', 'f', 'g', 'h', 'j', 'k', 'l'],
+]
+
+# Schema types that manage their own keybindings internally
+SELF_MANAGED_TYPES = {'pairwise', 'bws', 'triage'}
+
+
+def _get_label_name(label_data):
+ """Extract the label name from a label entry (string or dict)."""
+ if isinstance(label_data, str):
+ return label_data
+ if isinstance(label_data, Mapping):
+ return label_data.get("name", "")
+ return str(label_data)
+
+
+def _get_explicit_key(label_data):
+ """Extract explicit key_value from a label entry, or None."""
+ if isinstance(label_data, Mapping):
+ kv = label_data.get("key_value")
+ if kv is not None:
+ return str(kv).lower()
+ return None
+
+
+def _needs_allocation(scheme):
+ """Check if a schema needs keybinding allocation."""
+ ann_type = scheme.get("annotation_type", "")
+ if ann_type in SELF_MANAGED_TYPES:
+ return False
+
+ strategy = scheme.get("keybinding_strategy", "")
+ if strategy == "none":
+ return False
+
+ # Explicit sequential_key_binding
+ if scheme.get("sequential_key_binding"):
+ return True
+
+ # keybinding_strategy set to sequential or mnemonic
+ if strategy in ("sequential", "mnemonic"):
+ return True
+
+ return False
+
+
+def _assign_mnemonic_keys(labels, used_keys):
+ """
+ Assign mnemonic keys based on first available letter of each label name.
+ Falls back to next available letter if the preferred one is taken.
+
+ Returns list of (label_name, key) tuples.
+ """
+ # All available mnemonic letters
+ all_letters = list('abcdefghijklmnopqrstuvwxyz')
+ available = [c for c in all_letters if c not in used_keys]
+
+ assignments = []
+ for label_data in labels:
+ label_name = _get_label_name(label_data)
+ explicit = _get_explicit_key(label_data)
+ if explicit:
+ assignments.append((label_name, explicit))
+ continue
+
+ # Try each character of the label name
+ assigned = False
+ for char in label_name.lower():
+ if char.isalpha() and char in available:
+ assignments.append((label_name, char))
+ available.remove(char)
+ assigned = True
+ break
+
+ if not assigned:
+ # Fall back to next available letter
+ if available:
+ key = available.pop(0)
+ assignments.append((label_name, key))
+ logger.warning(
+ f"No mnemonic match for '{label_name}', "
+ f"assigned fallback key '{key}'"
+ )
+ else:
+ assignments.append((label_name, None))
+ logger.warning(
+ f"No keys available for label '{label_name}'"
+ )
+
+ return assignments
+
+
+def allocate_keybindings(annotation_schemes):
+ """
+ Pre-allocate non-conflicting keys across all annotation schemas.
+
+ Args:
+ annotation_schemes: List of annotation scheme dicts from config.
+
+ Returns:
+ dict: {schema_name: [{"label": str, "key": str|None}, ...]}
+ Only schemas that need allocation are included.
+ """
+ # Step 1: Collect all explicitly-set keys across all schemas
+ globally_used = set()
+
+ for scheme in annotation_schemes:
+ if not _needs_allocation(scheme):
+ continue
+ for label_data in scheme.get("labels", []):
+ explicit = _get_explicit_key(label_data)
+ if explicit:
+ globally_used.add(explicit)
+
+ # Step 2: Pre-claim keys used by self-managed schemas (pairwise, bws)
+ for scheme in annotation_schemes:
+ ann_type = scheme.get("annotation_type", "")
+ if ann_type == "pairwise":
+ if scheme.get("sequential_key_binding", True):
+ globally_used.update({'1', '2', '0'})
+ elif ann_type == "bws":
+ if scheme.get("sequential_key_binding", True):
+ tuple_size = scheme.get("tuple_size", 4)
+ for i in range(1, tuple_size + 1):
+ globally_used.add(str(i))
+ for i in range(tuple_size):
+ if i < 26:
+ globally_used.add(chr(ord('a') + i))
+
+ # Step 3: Build available pools (excluding globally used keys)
+ available_pools = []
+ for pool in KEY_POOLS:
+ available = [k for k in pool if k not in globally_used]
+ available_pools.append(available)
+
+ # Step 4: Allocate keys to schemas that need them
+ allocation = {}
+ next_pool_idx = 0
+
+ for scheme in annotation_schemes:
+ if not _needs_allocation(scheme):
+ continue
+
+ name = scheme.get("name", "")
+ labels = scheme.get("labels", [])
+ strategy = scheme.get("keybinding_strategy", "sequential")
+
+ if strategy == "mnemonic":
+ # Mnemonic allocation uses label name letters
+ assignments = _assign_mnemonic_keys(labels, globally_used)
+ result = []
+ for label_name, key in assignments:
+ result.append({"label": label_name, "key": key})
+ if key:
+ globally_used.add(key)
+ allocation[name] = result
+ continue
+
+ # Sequential allocation from pools
+ # Count how many keys we need (subtract explicit ones)
+ needed = 0
+ for label_data in labels:
+ if _get_explicit_key(label_data) is None:
+ needed += 1
+
+ # Find a pool with enough capacity
+ assigned_pool = None
+ for pool_idx in range(next_pool_idx, len(available_pools)):
+ if len(available_pools[pool_idx]) >= needed:
+ assigned_pool = pool_idx
+ break
+
+ if assigned_pool is None:
+ # Try earlier pools too (in case first schema used mnemonic)
+ for pool_idx in range(len(available_pools)):
+ if len(available_pools[pool_idx]) >= needed:
+ assigned_pool = pool_idx
+ break
+
+ if assigned_pool is None:
+ # Not enough keys in any single pool โ assign what we can
+ logger.warning(
+ f"Schema '{name}' has {needed} labels needing keys "
+ f"but no single pool has enough capacity. "
+ f"Some labels will not have keybindings."
+ )
+ # Use the pool with the most remaining keys
+ assigned_pool = max(
+ range(len(available_pools)),
+ key=lambda i: len(available_pools[i])
+ )
+
+ pool_keys = available_pools[assigned_pool]
+ key_iter = iter(pool_keys)
+
+ result = []
+ consumed = []
+ for label_data in labels:
+ label_name = _get_label_name(label_data)
+ explicit = _get_explicit_key(label_data)
+ if explicit:
+ result.append({"label": label_name, "key": explicit})
+ else:
+ key = next(key_iter, None)
+ if key:
+ result.append({"label": label_name, "key": key})
+ consumed.append(key)
+ globally_used.add(key)
+ else:
+ result.append({"label": label_name, "key": None})
+
+ # Remove consumed keys from the pool
+ available_pools[assigned_pool] = [
+ k for k in available_pools[assigned_pool] if k not in consumed
+ ]
+
+ allocation[name] = result
+
+ # Advance to next pool for the next schema
+ if assigned_pool == next_pool_idx:
+ next_pool_idx = assigned_pool + 1
+
+ return allocation
diff --git a/potato/server_utils/schemas/likert.py b/potato/server_utils/schemas/likert.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c86d72d60eb9b04e39ffbb917d51d8a174e9166
--- /dev/null
+++ b/potato/server_utils/schemas/likert.py
@@ -0,0 +1,181 @@
+"""
+Likert Scale Layout
+
+Generates a likert scale rating interface with radio buttons arranged horizontally.
+Each button represents a point on the scale between min_label and max_label.
+
+This module provides functionality for creating HTML-based Likert scale interfaces
+that can be used for collecting ordinal data responses. The scale supports:
+- Customizable number of points
+- Optional numeric display
+- Keyboard shortcuts
+- Required/optional validation
+- Bad text option for invalid inputs
+"""
+
+import logging
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper, get_dynamic_ai_help
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_element_value,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+from .radio import generate_radio_layout
+
+logger = logging.getLogger(__name__)
+
+def generate_likert_layout(annotation_scheme):
+ """
+ Generate HTML for a likert scale annotation interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - size: Number of scale points
+ - min_label: Label for minimum value
+ - max_label: Label for maximum value
+ - sequential_key_binding: Enable number key bindings (1-9)
+ - displaying_score: Show numeric values on buttons
+ - label_requirement: Validation settings
+ - required (bool): Whether response is mandatory
+ - bad_text_label (dict): Optional configuration for invalid text option
+ - label_content (str): Label text for bad text option
+ - annotation_id (int): match the config schema index
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the likert scale interface
+ key_bindings: List of (key, description) tuples for keyboard shortcuts
+
+ Raises:
+ Exception: If required fields are missing from annotation_scheme
+ """
+ return safe_generate_layout(annotation_scheme, _generate_likert_layout_internal)
+
+def _generate_likert_layout_internal(annotation_scheme):
+ """
+ Internal function to generate likert layout after validation.
+ """
+ logger.debug(f"Generating likert layout for schema: {annotation_scheme['name']}")
+
+ # Use radio layout if complex labels specified
+ if "labels" in annotation_scheme:
+ logger.info(f"Complex labels detected for {annotation_scheme['name']}, using radio layout")
+ return generate_radio_layout(annotation_scheme, horizontal=False)
+
+ # Validate required fields
+ required_fields = ["size", "min_label", "max_label"]
+ for required in required_fields:
+ if required not in annotation_scheme:
+ error_msg = f'Likert scale for "{annotation_scheme["name"]}" missing required field: {required}'
+ logger.error(error_msg)
+ raise Exception(error_msg)
+
+ logger.debug(f"Creating {annotation_scheme['size']}-point likert scale")
+
+ # Setup validation and key bindings
+ key_bindings = []
+ validation = generate_validation_attribute(annotation_scheme)
+
+ # Check for pre-allocated keys from the centralized allocator
+ allocated_keys = annotation_scheme.get("_allocated_keys", None)
+ allocated_map = {}
+ if allocated_keys:
+ for entry in allocated_keys:
+ if entry.get("key"):
+ allocated_map[entry["label"]] = entry["key"]
+
+ # Get layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Initialize form wrapper
+ schematic = f"""
+
+ """
+
+ logger.info(f"Successfully generated likert layout for {annotation_scheme['name']} "
+ f"with {annotation_scheme['size']} points")
+ return schematic, key_bindings
diff --git a/potato/server_utils/schemas/multirate.py b/potato/server_utils/schemas/multirate.py
new file mode 100644
index 0000000000000000000000000000000000000000..54764ba4022a619a9abe0c4ff2725e5ec0bc1677
--- /dev/null
+++ b/potato/server_utils/schemas/multirate.py
@@ -0,0 +1,457 @@
+"""
+Multirate Layout
+
+Generates a matrix-style interface for rating multiple items on the same scale.
+Features include:
+- Multiple column layout support
+- Configurable rating options
+- Vertical/horizontal arrangement options
+- Tooltip support
+- Required/optional validation
+"""
+
+import logging
+import os
+from collections.abc import Mapping
+from jinja2 import Template
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+logger = logging.getLogger(__name__)
+
+# HTML template using Jinja2 with comprehensive styling that preserves horizontal layout
+MULTIRATE_TEMPLATE = """
+
+"""
+
+def generate_multirate_layout(annotation_scheme):
+ """
+ Generate HTML for a multi-item rating interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - options: List of items to be rated
+ - labels: List of rating options to choose from
+ - display_config (dict): Optional display settings
+ - arrangement (str): Layout direction ('vertical' or 'horizontal')
+ - label_requirement (dict): Optional validation settings
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the multirate interface
+ key_bindings: List of (key, description) tuples for keyboard shortcuts
+ """
+ return safe_generate_layout(annotation_scheme, _generate_multirate_layout_internal)
+
+def _generate_multirate_layout_internal(annotation_scheme):
+ """
+ Internal function to generate multirate layout after validation.
+ """
+ logger.debug(f"Generating multirate layout for schema: {annotation_scheme['name']}")
+
+ # Check for options_from_data (dynamic multirate)
+ options_from_data = annotation_scheme.get('options_from_data')
+ if options_from_data and 'options' not in annotation_scheme:
+ return _generate_dynamic_multirate(annotation_scheme, options_from_data)
+
+ # Extract configuration
+ schema_name = annotation_scheme['name']
+ description = annotation_scheme['description']
+ options = annotation_scheme['options']
+ ratings = annotation_scheme['labels']
+
+ # Get display configuration
+ display_config = annotation_scheme.get('display_config', {})
+ num_columns = display_config.get('num_columns', 1)
+
+ # Set validation
+ validation = generate_validation_attribute(annotation_scheme)
+
+ # Get layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Preprocess items for template
+ processed_items = []
+ for option in options:
+ if isinstance(option, str):
+ identifiers = generate_element_identifier(schema_name, option, "multirate")
+ processed_items.append({
+ 'label': escape_html_content(option),
+ 'name': identifiers['name'],
+ 'id': identifiers['id'],
+ 'label_name': identifiers['label_name'],
+ 'tooltip': ""
+ })
+ else:
+ identifiers = generate_element_identifier(schema_name, option['name'], "multirate")
+ processed_items.append({
+ 'label': escape_html_content(option['label']),
+ 'name': identifiers['name'],
+ 'id': identifiers['id'],
+ 'label_name': identifiers['label_name'],
+ 'tooltip': _generate_tooltip(option)
+ })
+
+ # Arrange items according to specified layout
+ if annotation_scheme.get('arrangement') == 'vertical':
+ arranged_items = _arrange_items_vertically(processed_items, num_columns)
+ else:
+ arranged_items = _arrange_items_horizontally(processed_items, num_columns)
+
+ # Format template data
+ template_data = {
+ 'schema_name': escape_html_content(schema_name),
+ 'description': escape_html_content(description),
+ 'ratings': [escape_html_content(rating) for rating in ratings],
+ 'num_headers': min(len(options), num_columns),
+ 'rows': arranged_items,
+ 'validation': validation,
+ 'annotation_id': annotation_scheme.get('annotation_id', ''),
+ 'layout_attrs': layout_attrs
+ }
+
+ # Render template
+ template = Template(MULTIRATE_TEMPLATE)
+ html = template.render(**template_data)
+
+ logger.info(f"Successfully generated multirate layout for {schema_name} "
+ f"with {len(options)} items and {len(ratings)} rating options")
+
+ return html, [] # No key bindings implemented
+
+
+def _arrange_items_horizontally(items, num_columns):
+ """
+ Arrange items in a horizontal layout with specified number of columns.
+
+ Args:
+ items (list): List of processed item dictionaries
+ num_columns (int): Number of columns
+
+ Returns:
+ list: List of rows, where each row is a list of items
+ """
+ rows = []
+ for i in range(0, len(items), num_columns):
+ row = items[i:i+num_columns]
+ # Pad the row if it's not full
+ while len(row) < num_columns:
+ row.append(None)
+ rows.append(row)
+ return rows
+
+
+def _arrange_items_vertically(items, num_columns):
+ """
+ Rearrange items for vertical column layout.
+
+ Args:
+ items (list): List of processed item dictionaries
+ num_columns (int): Number of columns
+
+ Returns:
+ list: List of rows, where each row is a list of items arranged vertically
+ """
+ logger.debug(f"Rearranging {len(items)} items into {num_columns} vertical columns")
+
+ # Calculate rows needed
+ num_rows = (len(items) + num_columns - 1) // num_columns # Ceiling division
+
+ # Distribute items into columns
+ columns = [[] for _ in range(num_columns)]
+ for i, item in enumerate(items):
+ col_idx = i // num_rows
+ if col_idx < num_columns:
+ columns[col_idx].append(item)
+
+ # Create rows from columns
+ rows = []
+ for row_idx in range(num_rows):
+ row = []
+ for col in columns:
+ row.append(col[row_idx] if row_idx < len(col) else None)
+ rows.append(row)
+
+ return rows
+
+
+def _generate_tooltip(label_data):
+ """
+ Generate tooltip HTML attribute from label data.
+
+ Args:
+ label_data (dict): Label configuration containing tooltip information
+
+ Returns:
+ str: Tooltip HTML attribute or empty string if no tooltip
+ """
+ tooltip_text = ""
+ if "tooltip" in label_data:
+ tooltip_text = label_data["tooltip"]
+ elif "tooltip_file" in label_data:
+ try:
+ with open(label_data["tooltip_file"], "rt", encoding="utf-8") as f:
+ tooltip_text = "".join(f.readlines())
+ except Exception as e:
+ logger.error(f"Failed to read tooltip file: {e}")
+ return ""
+
+ if tooltip_text:
+ escaped_tooltip = escape_html_content(tooltip_text)
+ return f'data-toggle="tooltip" data-html="true" data-placement="top" title="{escaped_tooltip}"'
+ return ""
+
+
+# JavaScript template for dynamic multirate (options_from_data)
+DYNAMIC_MULTIRATE_JS = """
+
+"""
+
+DYNAMIC_MULTIRATE_TEMPLATE = """
+
+"""
+
+
+def _generate_dynamic_multirate(annotation_scheme, options_from_data):
+ """
+ Generate a dynamic multirate layout that reads options from instance data.
+
+ The generated HTML includes a JavaScript snippet that populates the multirate
+ table at page load time using data from the instance's specified field.
+
+ Args:
+ annotation_scheme: Schema configuration dict
+ options_from_data: Name of the instance data field containing options
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ import json
+ import html as html_module
+
+ schema_name = annotation_scheme['name']
+ description = annotation_scheme.get('description', '')
+ ratings = annotation_scheme.get('labels', [])
+ validation = generate_validation_attribute(annotation_scheme)
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ ratings_json = json.dumps(ratings)
+
+ template_data = {
+ 'schema_name': escape_html_content(schema_name),
+ 'description': escape_html_content(description),
+ 'data_key': escape_html_content(options_from_data),
+ 'ratings_json': html_module.escape(ratings_json),
+ 'options_values_json': '[]', # Will be filled by server at render time
+ 'annotation_id': annotation_scheme.get('annotation_id', ''),
+ 'layout_attrs': layout_attrs,
+ 'validation': validation,
+ # JS-safe values using json.dumps to produce quoted strings
+ 'schema_name_js': json.dumps(schema_name),
+ 'validation_js': json.dumps(validation),
+ }
+
+ template = Template(DYNAMIC_MULTIRATE_TEMPLATE + DYNAMIC_MULTIRATE_JS)
+ html = template.render(**template_data)
+
+ logger.info(f"Generated dynamic multirate for {schema_name} "
+ f"reading options from '{options_from_data}'")
+
+ return html, []
+
+
+def populate_dynamic_multirate(html_str, instance_data):
+ """
+ Post-process rendered HTML to inject instance-specific options into
+ dynamic multirate schemas.
+
+ Called from render_page_with_annotations() after template rendering.
+
+ Args:
+ html_str: The rendered HTML string
+ instance_data: The instance data dictionary
+
+ Returns:
+ Modified HTML string with dynamic multirate options populated
+ """
+ import json
+ import re
+ import html as html_module
+
+ # Find all dynamic multirate containers
+ pattern = r'data-options-from-data="([^"]+)"'
+ matches = list(re.finditer(pattern, html_str))
+
+ # Process in reverse order so earlier match offsets remain valid
+ for match in reversed(matches):
+ data_key = html_module.unescape(match.group(1))
+ options = instance_data.get(data_key, [])
+
+ if options and isinstance(options, list):
+ # Use html.escape on the JSON to safely embed in a single-quoted attribute
+ options_attr = html_module.escape(json.dumps(options))
+ # Search for the placeholder attribute near this match
+ search_start = max(0, match.start() - 200)
+ search_end = min(len(html_str), match.end() + 500)
+ local_html = html_str[search_start:search_end]
+
+ # Handle both single-quoted and double-quoted attributes
+ # (BeautifulSoup may convert single quotes to double quotes)
+ old_attr_single = "data-options-values='[]'"
+ old_attr_double = 'data-options-values="[]"'
+ if old_attr_single in local_html:
+ html_str = html_str[:search_start] + local_html.replace(
+ old_attr_single,
+ f"data-options-values='{options_attr}'",
+ 1,
+ ) + html_str[search_end:]
+ elif old_attr_double in local_html:
+ html_str = html_str[:search_start] + local_html.replace(
+ old_attr_double,
+ f'data-options-values="{options_attr}"',
+ 1,
+ ) + html_str[search_end:]
+
+ return html_str
\ No newline at end of file
diff --git a/potato/server_utils/schemas/multiselect.py b/potato/server_utils/schemas/multiselect.py
new file mode 100644
index 0000000000000000000000000000000000000000..58f73e36125aa0966156cf647c59868890f7a552
--- /dev/null
+++ b/potato/server_utils/schemas/multiselect.py
@@ -0,0 +1,260 @@
+"""
+Multiselect Layout
+
+Generates a form interface that allows users to select multiple options from a list
+of choices. Features include:
+- Multiple column layout support
+- Keyboard shortcuts
+- Required/optional validation
+- Individual label requirements
+- Tooltip support
+- Video label support
+- Free response option
+"""
+
+import logging
+from collections.abc import Mapping
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper, get_dynamic_ai_help
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+ display_label_text,
+)
+
+
+logger = logging.getLogger(__name__)
+
+def generate_multiselect_layout(annotation_scheme):
+ """
+ Generate HTML for a multiple-choice selection interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - labels: List of label configurations, each either:
+ - str: Simple label text
+ - dict: Complex label with:
+ - name: Label identifier
+ - tooltip: Hover text description
+ - tooltip_file: Path to file containing tooltip text
+ - key_value: Keyboard shortcut key
+ - videopath: Path to video file (if video_as_label=True)
+ - display_config (dict): Optional display settings
+ - num_columns: Number of columns to arrange options (default: 1)
+ - label_requirement (dict): Optional validation settings
+ - required (bool): Whether any selection is mandatory
+ - required_label (str|list): Specific labels that must be selected
+ - sequential_key_binding (bool): Enable numeric key shortcuts
+ - video_as_label (bool): Use videos instead of text for labels
+ - has_free_response (dict): Optional free text input configuration
+ - instruction: Label for free response field
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the multiselect interface
+ key_bindings: List of (key, description) tuples for keyboard shortcuts
+ """
+ return safe_generate_layout(annotation_scheme, _generate_multiselect_layout_internal)
+
+def _generate_multiselect_layout_internal(annotation_scheme):
+ """
+ Internal function to generate multiselect layout after validation.
+ """
+ logger.debug(f"Generating multiselect layout for schema: {annotation_scheme['name']}")
+
+ # Get layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Initialize form wrapper
+ schematic = f"""
+ "
+
+ logger.info(f"Successfully generated multiselect layout for {annotation_scheme['name']} "
+ f"with {len(annotation_scheme['labels'])} options")
+ return schematic, key_bindings
+
+def _generate_tooltip(label_data):
+ """
+ Generate tooltip HTML attribute from label data.
+
+ Args:
+ label_data (dict): Label configuration containing tooltip information
+
+ Returns:
+ str: Tooltip HTML attribute or empty string if no tooltip
+ """
+ tooltip_text = ""
+ if "tooltip" in label_data:
+ tooltip_text = label_data["tooltip"]
+ elif "tooltip_file" in label_data:
+ try:
+ with open(label_data["tooltip_file"], "rt", encoding="utf-8") as f:
+ tooltip_text = "".join(f.readlines())
+ except Exception as e:
+ logger.error(f"Failed to read tooltip file: {e}")
+ return ""
+
+ if tooltip_text:
+ escaped_tooltip = escape_html_content(tooltip_text)
+ return f'data-toggle="tooltip" data-html="true" data-placement="top" title="{escaped_tooltip}"'
+ return ""
+
+def _format_label_content(label_data, annotation_scheme):
+ """
+ Format the label content, handling both text and video labels.
+
+ Args:
+ label_data: Label configuration
+ annotation_scheme: Full annotation scheme configuration
+
+ Returns:
+ str: Formatted label content (text or video HTML)
+ """
+ if annotation_scheme.get("video_as_label") and isinstance(label_data, dict) and "videopath" in label_data:
+ # Video label
+ video_path = label_data["videopath"]
+ return f' '
+ else:
+ # Text label -- visible text is humanized (or explicit
+ # displayed_label); stored value remains the raw label name.
+ return escape_html_content(
+ display_label_text(label_data, annotation_scheme)
+ )
+
+def _generate_free_response(annotation_scheme, n_columns):
+ """
+ Generate free response field for multiselect.
+
+ Args:
+ annotation_scheme: Schema configuration
+ n_columns: Number of columns in the grid
+
+ Returns:
+ str: HTML for free response field
+ """
+ free_response_identifiers = generate_element_identifier(annotation_scheme["name"], "free_response", "text")
+ free_response_config = annotation_scheme["has_free_response"]
+ instruction = free_response_config.get("instruction", "Other (please specify)") if isinstance(free_response_config, dict) else "Other (please specify)"
+
+ return f"""
+
+ {escape_html_content(instruction)}
+
+
+ """
diff --git a/potato/server_utils/schemas/number.py b/potato/server_utils/schemas/number.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a401967afe34d667615f30e2cee2c31576d37b6
--- /dev/null
+++ b/potato/server_utils/schemas/number.py
@@ -0,0 +1,174 @@
+"""
+Number Layout
+
+Generates a form interface for numeric input. Features include:
+- Custom CSS styling options
+- Tooltip support
+- Required/optional validation
+- Min/max value constraints
+"""
+
+import logging
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper, get_dynamic_ai_help
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+
+logger = logging.getLogger(__name__)
+
+def generate_number_layout(annotation_scheme):
+ """
+ Generate HTML for a numeric input interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - custom_css (dict): Optional CSS styling
+ - width: Input width (default: "60px")
+ - height: Input height
+ - font_size: Text size
+ - tooltip: Optional hover text description
+ - tooltip_file: Optional path to tooltip text file
+ - label_requirement (dict): Optional validation settings
+ - required (bool): Whether input is mandatory
+ - min_value (int): Optional minimum allowed value
+ - max_value (int): Optional maximum allowed value
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the number input interface
+ key_bindings: Empty list (no keyboard shortcuts)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_number_layout_internal)
+
+def _generate_number_layout_internal(annotation_scheme):
+ """
+ Internal function to generate number layout after validation.
+ """
+ logger.debug(f"Generating number layout for schema: {annotation_scheme['name']}")
+
+ # Get custom dimensions from config
+ css = annotation_scheme.get("custom_css", {})
+ width = css.get("width", "60px")
+
+ # Get layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Initialize form wrapper
+ schematic = f"""
+
+ """
+
+ logger.info(f"Successfully generated number layout for {annotation_scheme['name']}")
+ return schematic, []
+
+def _generate_css_style(annotation_scheme):
+ """
+ Generate CSS style string from configuration.
+
+ Args:
+ annotation_scheme (dict): Configuration containing custom_css settings
+
+ Returns:
+ str: Formatted CSS style string
+ """
+ css = annotation_scheme.get("custom_css", {})
+ styles = []
+
+ # Default width if not specified
+ width = css.get("width", "60px")
+ styles.append(f"width: {width}")
+
+ # Optional height
+ if "height" in css:
+ styles.append(f"height: {css['height']}")
+
+ # Optional font size
+ if "font_size" in css:
+ styles.append(f"font-size: {css['font_size']}")
+
+ return "; ".join(styles)
+
+def _generate_tooltip(annotation_scheme):
+ """
+ Generate tooltip HTML attribute from configuration.
+
+ Args:
+ annotation_scheme (dict): Configuration containing tooltip information
+
+ Returns:
+ str: Tooltip HTML attribute or empty string if no tooltip
+ """
+ tooltip_text = ""
+ if "tooltip" in annotation_scheme:
+ tooltip_text = annotation_scheme["tooltip"]
+ elif "tooltip_file" in annotation_scheme:
+ try:
+ with open(annotation_scheme["tooltip_file"], "rt", encoding="utf-8") as f:
+ tooltip_text = "".join(f.readlines())
+ except Exception as e:
+ logger.error(f"Failed to read tooltip file: {e}")
+ return ""
+
+ if tooltip_text:
+ escaped_tooltip = escape_html_content(tooltip_text)
+ return f'data-toggle="tooltip" data-html="true" data-placement="top" title="{escaped_tooltip}"'
+ return ""
+
+def _generate_input_attributes(annotation_scheme):
+ """
+ Generate additional input attributes for number constraints.
+
+ Args:
+ annotation_scheme (dict): Configuration containing min/max values
+
+ Returns:
+ str: Space-separated attribute string
+ """
+ attrs = []
+
+ if "min_value" in annotation_scheme:
+ attrs.append(f'min="{annotation_scheme["min_value"]}"')
+ logger.debug(f"Setting minimum value: {annotation_scheme['min_value']}")
+
+ if "max_value" in annotation_scheme:
+ attrs.append(f'max="{annotation_scheme["max_value"]}"')
+ logger.debug(f"Setting maximum value: {annotation_scheme['max_value']}")
+
+ return " ".join(attrs)
diff --git a/potato/server_utils/schemas/pairwise.py b/potato/server_utils/schemas/pairwise.py
new file mode 100644
index 0000000000000000000000000000000000000000..d3d7f51dd8a80c207b5f033b751214a407572c49
--- /dev/null
+++ b/potato/server_utils/schemas/pairwise.py
@@ -0,0 +1,443 @@
+"""
+Pairwise Comparison Layout
+
+Generates a form interface for comparing two items side by side.
+Features include:
+- Binary mode: Click on preferred tile
+- Scale mode: Slider between items (-N to +N)
+- Optional tie/no-preference button
+- Keyboard shortcuts (1/2/0)
+- Support for items_key or inline items configuration
+"""
+
+import logging
+from typing import Dict, Any, Tuple, List
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+logger = logging.getLogger(__name__)
+
+
+def generate_pairwise_layout(annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Generate HTML for a pairwise comparison interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - mode: "binary" (default) or "scale"
+ - items_key: Key in instance data containing items to compare
+ - items: Inline items config (alternative to items_key)
+ - show_labels: Whether to show A/B labels (default: true)
+ - labels: Custom labels for A/B (default: ["A", "B"])
+ - allow_tie: Show tie/no-preference option (default: false)
+ - tie_label: Custom tie button text (default: "No preference")
+ - sequential_key_binding: Enable keyboard shortcuts (default: true)
+ - label_requirement (dict): Optional validation settings
+
+ For scale mode:
+ - scale.min: Minimum value (e.g., -3 for "A much better")
+ - scale.max: Maximum value (e.g., +3 for "B much better")
+ - scale.step: Step increment (default: 1)
+ - scale.labels.min: Label for min value
+ - scale.labels.max: Label for max value
+ - scale.labels.center: Label for center (default: "Equal")
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the pairwise interface
+ key_bindings: List of (key, description) tuples for keyboard shortcuts
+ """
+ return safe_generate_layout(annotation_scheme, _generate_pairwise_layout_internal)
+
+
+def _generate_pairwise_layout_internal(annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Internal function to generate pairwise layout after validation.
+ """
+ logger.debug(f"Generating pairwise layout for schema: {annotation_scheme['name']}")
+
+ mode = annotation_scheme.get("mode", "binary")
+ schema_name = annotation_scheme["name"]
+
+ if mode == "scale":
+ return _generate_scale_mode(annotation_scheme)
+ elif mode == "multi_dimension":
+ return _generate_multi_dimension_mode(annotation_scheme)
+ else:
+ return _generate_binary_mode(annotation_scheme)
+
+
+def _generate_binary_mode(annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Generate binary mode pairwise interface (clickable tiles).
+ """
+ schema_name = annotation_scheme["name"]
+ description = annotation_scheme["description"]
+
+ # Get configuration options
+ show_labels = annotation_scheme.get("show_labels", True)
+ labels = annotation_scheme.get("labels", ["A", "B"])
+ if len(labels) < 2:
+ labels = ["A", "B"]
+
+ allow_tie = annotation_scheme.get("allow_tie", False)
+ tie_label = annotation_scheme.get("tie_label", "No preference")
+
+ # Get items config
+ items_key = annotation_scheme.get("items_key", "text")
+
+ # Validation attribute
+ validation = generate_validation_attribute(annotation_scheme)
+
+ # Key bindings
+ key_bindings = []
+ enable_keybindings = annotation_scheme.get("sequential_key_binding", True)
+
+ # Build the HTML
+ escaped_schema = escape_html_content(schema_name)
+ escaped_description = escape_html_content(description)
+ escaped_items_key = escape_html_content(items_key)
+
+ # Data attributes for JavaScript initialization
+ data_attrs = f'data-annotation-type="pairwise" data-schema-name="{escaped_schema}" data-mode="binary" data-items-key="{escaped_items_key}"'
+
+ # Layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Tile labels
+ label_a = escape_html_content(labels[0])
+ label_b = escape_html_content(labels[1])
+ shortcut_a = "[1]" if enable_keybindings else ""
+ shortcut_b = "[2]" if enable_keybindings else ""
+ data_key_a = 'data-key="1"' if enable_keybindings else ""
+ data_key_b = 'data-key="2"' if enable_keybindings else ""
+
+ schematic = f"""
+
+ """
+
+ # Add key bindings
+ if enable_keybindings:
+ key_bindings.append(("1", f"{schema_name}: {labels[0]}"))
+ key_bindings.append(("2", f"{schema_name}: {labels[1]}"))
+ if allow_tie:
+ key_bindings.append(("0", f"{schema_name}: {tie_label}"))
+
+ logger.info(f"Successfully generated pairwise binary layout for {schema_name}")
+ return schematic, key_bindings
+
+
+def _generate_scale_mode(annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Generate scale mode pairwise interface (slider between items).
+ """
+ schema_name = annotation_scheme["name"]
+ description = annotation_scheme["description"]
+
+ # Get configuration options
+ show_labels = annotation_scheme.get("show_labels", True)
+ labels = annotation_scheme.get("labels", ["A", "B"])
+ if len(labels) < 2:
+ labels = ["A", "B"]
+
+ # Get scale configuration
+ scale_config = annotation_scheme.get("scale", {})
+ min_value = scale_config.get("min", -3)
+ max_value = scale_config.get("max", 3)
+ step = scale_config.get("step", 1)
+ default_value = scale_config.get("default", 0)
+
+ # Scale labels
+ scale_labels = scale_config.get("labels", {})
+ min_label = scale_labels.get("min", f"{labels[0]} is better")
+ max_label = scale_labels.get("max", f"{labels[1]} is better")
+ center_label = scale_labels.get("center", "Equal")
+
+ # Get items config
+ items_key = annotation_scheme.get("items_key", "text")
+
+ # Validation attribute
+ validation = generate_validation_attribute(annotation_scheme)
+
+ # Generate identifiers
+ identifiers = generate_element_identifier(schema_name, "scale", "range")
+
+ # Build the HTML
+ escaped_schema = escape_html_content(schema_name)
+ escaped_description = escape_html_content(description)
+ escaped_items_key = escape_html_content(items_key)
+
+ # Data attributes for JavaScript initialization
+ data_attrs = f'data-annotation-type="pairwise" data-schema-name="{escaped_schema}" data-mode="scale" data-items-key="{escaped_items_key}"'
+
+ # Layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Escaped labels
+ label_a = escape_html_content(labels[0])
+ label_b = escape_html_content(labels[1])
+
+ schematic = f"""
+
+ """
+
+ # No keyboard shortcuts for scale mode (uses slider)
+ key_bindings = []
+
+ logger.info(f"Successfully generated pairwise scale layout for {schema_name}")
+ return schematic, key_bindings
+
+
+def _generate_justification_html(annotation_scheme: Dict[str, Any], schema_name: str) -> str:
+ """
+ Generate justification section HTML (reason checkboxes + rationale textarea).
+
+ Used by both binary and multi_dimension modes when ``justification`` config is present.
+ """
+ justification = annotation_scheme.get("justification")
+ if not justification:
+ return ""
+
+ reason_categories = justification.get("reason_categories", [])
+ min_chars = justification.get("min_rationale_chars", 0)
+ placeholder = justification.get("rationale_placeholder", "Explain your preference...")
+ required = justification.get("required", False)
+
+ escaped_schema = escape_html_content(schema_name)
+
+ # Reason category checkboxes
+ reasons_html = ""
+ for reason in reason_categories:
+ escaped_reason = escape_html_content(reason)
+ reasons_html += f"""
+
+ {escaped_reason} """
+
+ req_attr = 'data-required="true"' if required else ''
+
+ html = f"""
+
+
Justification
+ """
+
+ if reasons_html:
+ html += f"""
+
{reasons_html}
+ """
+
+ html += f"""
+
+
+ 0 / {min_chars} characters
+
+
+
+ """
+ return html
+
+
+def _generate_multi_dimension_mode(annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Generate multi-dimension pairwise interface.
+
+ Each dimension gets its own A/B tile row with a dimension label.
+ Hidden inputs use ``{schema}:::{dimension}`` naming.
+ """
+ schema_name = annotation_scheme["name"]
+ description = annotation_scheme["description"]
+ dimensions = annotation_scheme.get("dimensions", [])
+
+ if not dimensions:
+ raise ValueError(f"pairwise multi_dimension mode requires 'dimensions' list for schema '{schema_name}'")
+
+ labels = annotation_scheme.get("labels", ["A", "B"])
+ if len(labels) < 2:
+ labels = ["A", "B"]
+
+ items_key = annotation_scheme.get("items_key", "text")
+ validation = generate_validation_attribute(annotation_scheme)
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ escaped_schema = escape_html_content(schema_name)
+ escaped_items_key = escape_html_content(items_key)
+ label_a = escape_html_content(labels[0])
+ label_b = escape_html_content(labels[1])
+
+ data_attrs = (
+ f'data-annotation-type="pairwise" data-schema-name="{escaped_schema}" '
+ f'data-mode="multi_dimension" data-items-key="{escaped_items_key}"'
+ )
+
+ schematic = f"""
+
+ """
+
+ logger.info(f"Successfully generated pairwise multi_dimension layout for {schema_name}")
+ return schematic, key_bindings
diff --git a/potato/server_utils/schemas/process_reward.py b/potato/server_utils/schemas/process_reward.py
new file mode 100644
index 0000000000000000000000000000000000000000..a8e0092b11a1d4b494baa9d45e678bc5f689819d
--- /dev/null
+++ b/potato/server_utils/schemas/process_reward.py
@@ -0,0 +1,580 @@
+"""
+Process Reward Schema
+
+Binary per-step correct/incorrect signals for Process Reward Model (PRM)
+training. Two modes:
+- "per_step": annotate each step independently with thumbs-up/down
+- "first_error": click the first wrong step, all subsequent auto-marked wrong
+
+Research: AgentPRM, ToolRM, ToolRL, SPORT
+"""
+
+import json
+import logging
+from typing import Dict, Any, Tuple, List
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def generate_process_reward_layout(
+ annotation_scheme: Dict[str, Any],
+) -> Tuple[str, List[Tuple[str, str]]]:
+ """Generate HTML for a process reward annotation interface.
+
+ Args:
+ annotation_scheme: Configuration dict. Required keys: ``name``,
+ ``description``. Optional: ``steps_key``, ``mode``.
+
+ Returns:
+ ``(html, keybindings)`` tuple.
+ """
+ return safe_generate_layout(annotation_scheme, _generate_internal)
+
+
+def _generate_internal(
+ annotation_scheme: Dict[str, Any],
+) -> Tuple[str, List[Tuple[str, str]]]:
+ schema_name = annotation_scheme["name"]
+ description = annotation_scheme["description"]
+ steps_key = annotation_scheme.get("steps_key", "steps")
+ step_text_key = annotation_scheme.get("step_text_key", "action")
+ mode = annotation_scheme.get("mode", "first_error") # "first_error" or "per_step"
+ # When true, the per-step Correct/Wrong control is injected to the right
+ # of each rendered trace step (a [data-turn-index] element) rather than a
+ # separate card list at the bottom. Falls back to the card list if no
+ # trace step elements are present (e.g. non-trace displays).
+ inline_with_trace = bool(annotation_scheme.get("inline_with_trace", False))
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+ identifiers = generate_element_identifier(schema_name, schema_name, "hidden")
+ esc_schema = escape_html_content(schema_name)
+
+ config_json = json.dumps({
+ "steps_key": steps_key,
+ "step_text_key": step_text_key,
+ "mode": mode,
+ "inline_with_trace": inline_with_trace,
+ })
+
+ container_class = "process-reward-container"
+ if inline_with_trace:
+ container_class += " prm-inline-mode"
+
+ html = f"""
+
+
+
+
+
+ """
+
+ logger.info(
+ f"Successfully generated process_reward layout for {schema_name} "
+ f"(mode={mode})"
+ )
+ return html, [] # No keybindings
diff --git a/potato/server_utils/schemas/pure_display.py b/potato/server_utils/schemas/pure_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..492b873db381751c500c613ddb3de277978eec62
--- /dev/null
+++ b/potato/server_utils/schemas/pure_display.py
@@ -0,0 +1,128 @@
+"""
+Pure Display Layout
+
+Generates a simple display-only interface that shows text content without any
+interaction elements. This is useful for:
+- Displaying instructions
+- Showing static content
+- Presenting read-only information
+- Headers and section dividers
+
+Supports an optional `allow_html: true` flag that lets administrators include
+trusted HTML formatting (e.g., , , ) in description and labels.
+When enabled, content is sanitized through the project's HTML sanitizer to
+block dangerous elements while preserving safe formatting.
+"""
+
+import logging
+from .identifier_utils import (
+ safe_generate_layout,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _sanitize_or_escape(content: str, allow_html: bool) -> str:
+ """
+ Sanitize or escape content based on the allow_html flag.
+
+ When allow_html is True, uses the project's HTML sanitizer which allows
+ safe elements (b, i, u, em, strong, br, span, div, ul, ol, li, etc.)
+ while blocking scripts and dangerous attributes.
+
+ When allow_html is False (default), fully escapes all HTML.
+ """
+ if not content:
+ return ""
+ if allow_html:
+ from potato.server_utils.html_sanitizer import sanitize_html
+ return str(sanitize_html(str(content)))
+ return escape_html_content(content)
+
+
+def generate_pure_display_layout(annotation_scheme):
+ """
+ Generate HTML for a display-only text interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Main text to display as header
+ - labels: List of text strings to display as content
+ - allow_html: (optional, default False) If True, allows trusted
+ HTML in description and labels (sanitized for safety)
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the display interface
+ key_bindings: Empty list (no interactions available)
+
+ Example:
+ annotation_scheme = {
+ "name": "instructions",
+ "description": "Task Instructions",
+ "labels": ["Step 1: Read the text", "Step 2: Select options"]
+ }
+
+ Example with HTML:
+ annotation_scheme = {
+ "name": "consent_header",
+ "description": "Consent Form ",
+ "labels": ["Please read the following carefully."],
+ "allow_html": True
+ }
+ """
+ return safe_generate_layout(annotation_scheme, _generate_pure_display_layout_internal)
+
+def _generate_pure_display_layout_internal(annotation_scheme):
+ """
+ Internal function to generate pure display layout after validation.
+ """
+ logger.debug(f"Generating pure display layout for schema: {annotation_scheme['name']}")
+
+ allow_html = annotation_scheme.get('allow_html', False)
+
+ # Get layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Name is always escaped (used in HTML attributes, not content)
+ escaped_name = escape_html_content(annotation_scheme['name'])
+
+ # Description and labels respect allow_html flag
+ description = _sanitize_or_escape(annotation_scheme['description'], allow_html)
+
+ # Format content with header and body text
+ schematic = f"""
+
+ """
+
+ logger.info(f"Successfully generated pure display layout for {annotation_scheme['name']}")
+ return schematic, []
+
+def format_display_content(labels, allow_html=False):
+ """
+ Format the display content from a list of labels.
+
+ Args:
+ labels (list): List of strings to display
+ allow_html (bool): If True, sanitize rather than escape HTML content
+
+ Returns:
+ str: HTML formatted content with line breaks between items
+ """
+ if not labels:
+ logger.warning("No labels provided for pure display content")
+ return ""
+
+ logger.debug(f"Formatting {len(labels)} content lines")
+ processed_labels = [_sanitize_or_escape(label, allow_html) for label in labels]
+ return " ".join(processed_labels)
diff --git a/potato/server_utils/schemas/radio.py b/potato/server_utils/schemas/radio.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d7ac05bf3a9293c0b65f64f2f1efd1732866b01
--- /dev/null
+++ b/potato/server_utils/schemas/radio.py
@@ -0,0 +1,229 @@
+"""
+Radio Layout
+
+Generates a form interface with mutually exclusive radio button options.
+Features include:
+- Vertical or horizontal layout options
+- Keyboard shortcuts
+- Required/optional validation
+- Tooltip support
+- Free response option
+"""
+
+import logging
+from collections.abc import Mapping
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper, get_dynamic_ai_help
+from potato.server_utils.config_module import config
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+ display_label_text,
+)
+
+
+logger = logging.getLogger(__name__)
+
+def generate_radio_layout(annotation_scheme, horizontal=False):
+ """
+ Generate HTML for a radio button selection interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - labels: List of label configurations, each either:
+ - str: Simple label text
+ - dict: Complex label with:
+ - name: Label identifier
+ - tooltip: Hover text description
+ - tooltip_file: Path to tooltip text file
+ - key_value: Keyboard shortcut key
+ - label_requirement (dict): Optional validation settings
+ - required (bool): Whether selection is mandatory
+ - horizontal (bool): Whether to arrange options horizontally
+ - has_free_response (dict): Optional free text input configuration
+ - instruction: Label for free response field
+
+ horizontal (bool): Override horizontal layout setting
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the radio interface
+ key_bindings: List of (key, description) tuples for keyboard shortcuts
+ """
+ return safe_generate_layout(annotation_scheme, _generate_radio_layout_internal, horizontal)
+
+def _generate_radio_layout_internal(annotation_scheme, horizontal=False):
+ """
+ Internal function to generate radio layout after validation.
+ """
+ logger.debug(f"Generating radio layout for schema: {annotation_scheme['name']}")
+
+ # Check for horizontal layout override
+ if annotation_scheme.get("horizontal"):
+ horizontal = True
+ logger.debug("Using horizontal layout")
+
+ # Get layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Initialize form wrapper
+ schema_name = annotation_scheme["name"]
+ schematic = f"""
+ "
+
+ logger.info(f"Successfully generated radio layout for {schema_name} "
+ f"with {len(annotation_scheme['labels'])} options")
+ return schematic, key_bindings
+
+def _generate_tooltip(label_data):
+ """
+ Generate tooltip HTML attribute from label data.
+
+ Args:
+ label_data (dict): Label configuration containing tooltip information
+
+ Returns:
+ str: Tooltip HTML attribute or empty string if no tooltip
+ """
+ tooltip_text = ""
+ if "tooltip" in label_data:
+ tooltip_text = label_data["tooltip"]
+ elif "tooltip_file" in label_data:
+ try:
+ with open(label_data["tooltip_file"], "rt", encoding="utf-8") as f:
+ tooltip_text = "".join(f.readlines())
+ except Exception as e:
+ logger.error(f"Failed to read tooltip file: {e}")
+ return ""
+
+ if tooltip_text:
+ escaped_tooltip = escape_html_content(tooltip_text)
+ return f'data-toggle="tooltip" data-html="true" data-placement="top" title="{escaped_tooltip}"'
+ return ""
diff --git a/potato/server_utils/schemas/range_slider.py b/potato/server_utils/schemas/range_slider.py
new file mode 100644
index 0000000000000000000000000000000000000000..7ad10d4ae28ad2200e2add9aadbe26d2d9252f1e
--- /dev/null
+++ b/potato/server_utils/schemas/range_slider.py
@@ -0,0 +1,282 @@
+"""
+Goldilocks Range / Dual-Thumb Slider Layout
+
+Generates a dual-thumb range slider for selecting a min-max range.
+Uses custom div-based handles (not native range inputs) for reliable
+cross-browser rendering.
+
+Research basis:
+- Pavlick & Kwiatkowski (2019) "Inherent Disagreements in Human Textual
+ Inferences" TACL
+- Jurgens (2013) "Embracing Ambiguity: A Comparison of Annotation Methodologies
+ for Crowdsourcing Word Sense Labels" NAACL
+"""
+
+import logging
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+)
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MIN = 0
+DEFAULT_MAX = 100
+DEFAULT_STEP = 1
+
+
+def generate_range_slider_layout(annotation_scheme):
+ """
+ Generate HTML for a dual-thumb range slider interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - min_value: Minimum value (default 0)
+ - max_value: Maximum value (default 100)
+ - step: Step size (default 1)
+ - left_label: Label for left end
+ - right_label: Label for right end
+ - show_values: Show numeric values (default true)
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_range_slider_layout_internal)
+
+
+def _generate_range_slider_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme["name"]
+ safe_schema = escape_html_content(schema_name)
+ description = annotation_scheme["description"]
+ min_val = annotation_scheme.get("min_value", DEFAULT_MIN)
+ max_val = annotation_scheme.get("max_value", DEFAULT_MAX)
+ step = annotation_scheme.get("step", DEFAULT_STEP)
+ left_label = annotation_scheme.get("left_label", "")
+ right_label = annotation_scheme.get("right_label", "")
+ show_values = annotation_scheme.get("show_values", True)
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+
+ # Default initial range: 25th to 75th percentile
+ range_span = max_val - min_val
+ init_low = min_val + range_span // 4
+ init_high = max_val - range_span // 4
+
+ id_low = generate_element_identifier(schema_name, "range_low", "range")
+ id_high = generate_element_identifier(schema_name, "range_high", "range")
+
+ html = f"""
+
+
+ """
+
+ key_bindings = []
+ logger.info(f"Generated range_slider layout for {schema_name}")
+ return html, key_bindings
diff --git a/potato/server_utils/schemas/ranking.py b/potato/server_utils/schemas/ranking.py
new file mode 100644
index 0000000000000000000000000000000000000000..ff54d302ac200efd5b2eeac837f3b5dfb935a01b
--- /dev/null
+++ b/potato/server_utils/schemas/ranking.py
@@ -0,0 +1,213 @@
+"""
+Ranking / Drag-and-Drop Layout
+
+Generates a draggable list of items that annotators can reorder.
+Uses a hidden input to store the rank order.
+
+Research basis:
+- Kiritchenko & Mohammad (2017) "Best-Worst Scaling More Reliable than Rating
+ Scales" ACL
+- Thurstone (1927) "A Law of Comparative Judgment"
+"""
+
+import logging
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+ generate_tooltip_html,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def generate_ranking_layout(annotation_scheme):
+ """
+ Generate HTML for a ranking / drag-and-drop annotation interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - labels: List of items to rank
+ - allow_ties: Whether ties are allowed (default false)
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_ranking_layout_internal)
+
+
+def _generate_ranking_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme["name"]
+ safe_schema = escape_html_content(schema_name)
+ description = annotation_scheme["description"]
+ labels = annotation_scheme.get("labels", [])
+ allow_ties = annotation_scheme.get("allow_ties", False)
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+
+ if not labels:
+ raise ValueError(f"ranking schema '{schema_name}' requires 'labels'")
+
+ # Normalize labels
+ label_names = []
+ for lbl in labels:
+ if isinstance(lbl, str):
+ label_names.append(lbl)
+ elif isinstance(lbl, dict) and "name" in lbl:
+ label_names.append(lbl["name"])
+ else:
+ raise ValueError(f"Invalid label format: {lbl}")
+
+ identifiers = generate_element_identifier(schema_name, "rank_order", "hidden")
+ initial_order = ",".join(label_names)
+
+ html = f"""
+
+
+ """
+
+ key_bindings = []
+ logger.info(f"Generated ranking layout for {schema_name} with {len(label_names)} items")
+ return html, key_bindings
diff --git a/potato/server_utils/schemas/registry.py b/potato/server_utils/schemas/registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..7e277d6170b0c4bf9dd6a4a65a2e16a2f59f2ad9
--- /dev/null
+++ b/potato/server_utils/schemas/registry.py
@@ -0,0 +1,597 @@
+"""
+Schema Registry
+
+Provides a centralized registry for managing annotation schema types.
+This module serves as the single source of truth for available annotation
+types and their generators, replacing the hardcoded dictionary in front_end.py.
+
+Usage:
+ from potato.server_utils.schemas.registry import schema_registry
+
+ # Get a schema generator
+ generator = schema_registry.get("radio")
+
+ # Generate layout for an annotation scheme
+ html, keybindings = schema_registry.generate(annotation_scheme_dict)
+
+ # List all available schemas
+ schemas = schema_registry.list_schemas()
+"""
+
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Tuple, Any, Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SchemaDefinition:
+ """
+ Defines metadata and generator for an annotation schema type.
+
+ Attributes:
+ name: Unique identifier for the schema type (e.g., "radio", "multiselect")
+ generator: Callable that generates HTML and keybindings for this schema type
+ required_fields: List of required configuration fields
+ optional_fields: List of optional configuration fields
+ supports_keybindings: Whether this schema type supports keyboard shortcuts
+ description: Human-readable description of the schema type
+ """
+ name: str
+ generator: Callable[[Dict[str, Any]], Tuple[str, List[Tuple[str, str]]]]
+ required_fields: List[str] = field(default_factory=list)
+ optional_fields: List[str] = field(default_factory=list)
+ supports_keybindings: bool = True
+ description: str = ""
+
+
+class SchemaRegistry:
+ """
+ Centralized registry for annotation schema types.
+
+ Provides methods to register, retrieve, and list schema types,
+ as well as generate layouts from annotation scheme configurations.
+ """
+
+ def __init__(self):
+ self._schemas: Dict[str, SchemaDefinition] = {}
+ logger.debug("SchemaRegistry initialized")
+
+ def register(self, schema: SchemaDefinition) -> None:
+ """
+ Register a new schema type.
+
+ Args:
+ schema: SchemaDefinition to register
+
+ Raises:
+ ValueError: If a schema with the same name is already registered
+ """
+ if schema.name in self._schemas:
+ raise ValueError(f"Schema '{schema.name}' is already registered")
+
+ self._schemas[schema.name] = schema
+ logger.debug(f"Registered schema: {schema.name}")
+
+ def get(self, name: str) -> Optional[SchemaDefinition]:
+ """
+ Get a schema definition by name.
+
+ Args:
+ name: The schema type name
+
+ Returns:
+ SchemaDefinition if found, None otherwise
+ """
+ return self._schemas.get(name)
+
+ def get_generator(self, name: str) -> Optional[Callable]:
+ """
+ Get the generator function for a schema type.
+
+ Args:
+ name: The schema type name
+
+ Returns:
+ Generator callable if found, None otherwise
+ """
+ schema = self.get(name)
+ return schema.generator if schema else None
+
+ def generate(self, annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Generate HTML and keybindings for an annotation scheme.
+
+ Args:
+ annotation_scheme: Configuration dictionary with 'annotation_type' key
+
+ Returns:
+ Tuple of (html_string, keybindings_list)
+
+ Raises:
+ ValueError: If annotation_type is missing or not supported
+ """
+ import json as json_module
+
+ annotation_type = annotation_scheme.get("annotation_type")
+ if not annotation_type:
+ raise ValueError("annotation_scheme must have 'annotation_type' field")
+
+ schema = self.get(annotation_type)
+ if not schema:
+ supported = ", ".join(sorted(self._schemas.keys()))
+ raise ValueError(
+ f"Unsupported annotation type: '{annotation_type}'. "
+ f"Supported types are: {supported}"
+ )
+
+ logger.debug(f"Generating layout for annotation type: {annotation_type}")
+ html, keybindings = schema.generator(annotation_scheme)
+
+ # Wrap HTML with display_logic attributes if present
+ html = self._wrap_with_display_logic(html, annotation_scheme)
+
+ return html, keybindings
+
+ def _wrap_with_display_logic(self, html: str, annotation_scheme: Dict[str, Any]) -> str:
+ """
+ Wrap the generated HTML with display_logic attributes and container if needed.
+
+ If the annotation scheme has display_logic, wraps the HTML in a container
+ div with data-display-logic attribute containing the serialized logic,
+ and applies the initial hidden state.
+
+ Args:
+ html: The generated HTML from the schema generator
+ annotation_scheme: The annotation scheme configuration
+
+ Returns:
+ The HTML, possibly wrapped with display_logic container
+ """
+ import json as json_module
+ from html import escape
+
+ display_logic = annotation_scheme.get("display_logic")
+ if not display_logic:
+ return html
+
+ schema_name = annotation_scheme.get("name", "")
+
+ # Serialize display_logic to JSON for the data attribute
+ display_logic_json = json_module.dumps(display_logic)
+ # Escape for HTML attribute (single quotes in JSON are safe)
+ display_logic_escaped = escape(display_logic_json)
+
+ # Wrap in a container with display logic attributes
+ # Initially hidden - the frontend JS will evaluate and show if conditions are met
+ wrapped_html = f'''
+{html}
+
'''
+
+ return wrapped_html
+
+ def list_schemas(self) -> List[Dict[str, Any]]:
+ """
+ List all registered schemas with their metadata.
+
+ Returns:
+ List of dictionaries containing schema metadata
+ """
+ return [
+ {
+ "name": schema.name,
+ "description": schema.description,
+ "required_fields": schema.required_fields,
+ "optional_fields": schema.optional_fields,
+ "supports_keybindings": schema.supports_keybindings,
+ }
+ for schema in sorted(self._schemas.values(), key=lambda s: s.name)
+ ]
+
+ def is_registered(self, name: str) -> bool:
+ """
+ Check if a schema type is registered.
+
+ Args:
+ name: The schema type name
+
+ Returns:
+ True if registered, False otherwise
+ """
+ return name in self._schemas
+
+ def get_supported_types(self) -> List[str]:
+ """
+ Get a list of all supported annotation types.
+
+ Returns:
+ Sorted list of annotation type names
+ """
+ return sorted(self._schemas.keys())
+
+
+# Global registry instance
+schema_registry = SchemaRegistry()
+
+
+def _register_builtin_schemas():
+ """
+ Register all built-in annotation schema types.
+ Called automatically when this module is imported.
+ """
+ from .radio import generate_radio_layout
+ from .multiselect import generate_multiselect_layout
+ from .multirate import generate_multirate_layout
+ from .likert import generate_likert_layout
+ from .textbox import generate_textbox_layout
+ from .number import generate_number_layout
+ from .slider import generate_slider_layout
+ from .span import generate_span_layout
+ from .span_link import generate_span_link_layout
+ from .select import generate_select_layout
+ from .pure_display import generate_pure_display_layout
+ from .video import generate_video_layout
+ from .image_annotation import generate_image_annotation_layout
+ from .audio_annotation import generate_audio_annotation_layout
+ from .video_annotation import generate_video_annotation_layout
+ from .pairwise import generate_pairwise_layout
+ from .coreference import generate_coreference_layout
+ from .tree_annotation import generate_tree_annotation_layout
+ from .triage import generate_triage_layout
+ from .event_annotation import generate_event_annotation_layout
+ from .tiered_annotation import generate_tiered_annotation_layout
+ from .bws import generate_bws_layout
+ from .soft_label import generate_soft_label_layout
+ from .confidence import generate_confidence_layout
+ from .constant_sum import generate_constant_sum_layout
+ from .semantic_differential import generate_semantic_differential_layout
+ from .ranking import generate_ranking_layout
+ from .range_slider import generate_range_slider_layout
+ from .hierarchical_multiselect import generate_hierarchical_multiselect_layout
+ from .vas import generate_vas_layout
+ from .extractive_qa import generate_extractive_qa_layout
+ from .rubric_eval import generate_rubric_eval_layout
+ from .text_edit import generate_text_edit_layout
+ from .error_span import generate_error_span_layout
+ from .card_sort import generate_card_sort_layout
+ from .conjoint import generate_conjoint_layout
+ from .trajectory_eval import generate_trajectory_eval_layout
+ from .trajectory_edit import generate_trajectory_edit_layout
+ from .process_reward import generate_process_reward_layout
+ from .code_review import generate_code_review_layout
+
+ schemas = [
+ SchemaDefinition(
+ name="radio",
+ generator=generate_radio_layout,
+ required_fields=["name", "description", "labels"],
+ optional_fields=["horizontal", "label_requirement", "sequential_key_binding", "has_free_response", "option_randomization", "dynamic_options", "dynamic_options_field"],
+ supports_keybindings=True,
+ description="Single-choice radio button selection"
+ ),
+ SchemaDefinition(
+ name="multiselect",
+ generator=generate_multiselect_layout,
+ required_fields=["name", "description", "labels"],
+ optional_fields=["display_config", "label_requirement", "sequential_key_binding", "video_as_label", "has_free_response", "option_randomization", "dynamic_options", "dynamic_options_field"],
+ supports_keybindings=True,
+ description="Multiple-choice checkbox selection"
+ ),
+ SchemaDefinition(
+ name="multirate",
+ generator=generate_multirate_layout,
+ required_fields=["name", "description", "options", "labels"],
+ optional_fields=["label_requirement"],
+ supports_keybindings=False,
+ description="Rate multiple items on a scale"
+ ),
+ SchemaDefinition(
+ name="likert",
+ generator=generate_likert_layout,
+ required_fields=["name", "description", "min_label", "max_label", "size"],
+ optional_fields=["label_requirement"],
+ supports_keybindings=True,
+ description="Likert scale rating"
+ ),
+ SchemaDefinition(
+ name="text",
+ generator=generate_textbox_layout,
+ required_fields=["name", "description"],
+ optional_fields=["label_requirement", "placeholder", "rows"],
+ supports_keybindings=False,
+ description="Free-form text input"
+ ),
+ SchemaDefinition(
+ name="number",
+ generator=generate_number_layout,
+ required_fields=["name", "description"],
+ optional_fields=["min", "max", "step", "label_requirement"],
+ supports_keybindings=False,
+ description="Numeric input field"
+ ),
+ SchemaDefinition(
+ name="slider",
+ generator=generate_slider_layout,
+ required_fields=["name", "description", "min_value", "max_value", "starting_value"],
+ optional_fields=["step", "label_requirement"],
+ supports_keybindings=False,
+ description="Slider for selecting a value in a range"
+ ),
+ SchemaDefinition(
+ name="span",
+ generator=generate_span_layout,
+ required_fields=["name", "description", "labels"],
+ optional_fields=["sequential_key_binding", "bad_text_label", "title", "allow_discontinuous", "entity_linking", "show_span_labels"],
+ supports_keybindings=True,
+ description="Text span annotation/highlighting with optional entity linking to knowledge bases"
+ ),
+ SchemaDefinition(
+ name="select",
+ generator=generate_select_layout,
+ required_fields=["name", "description", "labels"],
+ optional_fields=["label_requirement", "option_randomization", "dynamic_options", "dynamic_options_field"],
+ supports_keybindings=False,
+ description="Dropdown selection"
+ ),
+ SchemaDefinition(
+ name="pure_display",
+ generator=generate_pure_display_layout,
+ required_fields=["name", "description"],
+ optional_fields=["labels", "allow_html"],
+ supports_keybindings=False,
+ description="Display-only content (instructions, headers)"
+ ),
+ SchemaDefinition(
+ name="video",
+ generator=generate_video_layout,
+ required_fields=["name", "description", "video_path"],
+ optional_fields=["autoplay", "loop", "muted", "controls", "custom_css", "fallback_text", "additional_sources"],
+ supports_keybindings=False,
+ description="Video player display"
+ ),
+ SchemaDefinition(
+ name="image_annotation",
+ generator=generate_image_annotation_layout,
+ required_fields=["name", "description", "tools", "labels"],
+ optional_fields=["zoom_enabled", "pan_enabled", "min_annotations", "max_annotations", "freeform_brush_size", "freeform_simplify"],
+ supports_keybindings=True,
+ description="Image annotation with bounding boxes, polygons, freeform drawing, and landmarks"
+ ),
+ SchemaDefinition(
+ name="audio_annotation",
+ generator=generate_audio_annotation_layout,
+ required_fields=["name", "description"],
+ optional_fields=["mode", "labels", "segment_schemes", "min_segments", "max_segments", "zoom_enabled", "playback_rate_control"],
+ supports_keybindings=True,
+ description="Audio segmentation and annotation with waveform visualization"
+ ),
+ SchemaDefinition(
+ name="video_annotation",
+ generator=generate_video_annotation_layout,
+ required_fields=["name", "description"],
+ optional_fields=["mode", "labels", "segment_schemes", "min_segments", "max_segments", "timeline_height", "overview_height", "zoom_enabled", "playback_rate_control", "frame_stepping", "show_timecode", "video_fps", "tracking_options"],
+ supports_keybindings=True,
+ description="Video annotation with temporal segments, frame classification, keyframes, and object tracking"
+ ),
+ SchemaDefinition(
+ name="span_link",
+ generator=generate_span_link_layout,
+ required_fields=["name", "description", "link_types", "span_schema"],
+ optional_fields=["visual_display"],
+ supports_keybindings=False,
+ description="Create relationships/links between spans (e.g., PERSON works_for ORGANIZATION)"
+ ),
+ SchemaDefinition(
+ name="pairwise",
+ generator=generate_pairwise_layout,
+ required_fields=["name", "description"],
+ optional_fields=["mode", "items_key", "items", "show_labels", "labels", "allow_tie", "tie_label", "sequential_key_binding", "scale", "label_requirement"],
+ supports_keybindings=True,
+ description="Pairwise comparison of two items (binary selection or scale rating)"
+ ),
+ SchemaDefinition(
+ name="coreference",
+ generator=generate_coreference_layout,
+ required_fields=["name", "description", "span_schema"],
+ optional_fields=["entity_types", "allow_singletons", "visual_display"],
+ supports_keybindings=False,
+ description="Coreference chain annotation for grouping mentions of the same entity"
+ ),
+ SchemaDefinition(
+ name="tree_annotation",
+ generator=generate_tree_annotation_layout,
+ required_fields=["name", "description"],
+ optional_fields=["node_scheme", "path_selection", "branch_comparison"],
+ supports_keybindings=False,
+ description="Annotation of conversation tree nodes with path selection"
+ ),
+ SchemaDefinition(
+ name="triage",
+ generator=generate_triage_layout,
+ required_fields=["name", "description"],
+ optional_fields=["accept_label", "reject_label", "skip_label", "auto_advance", "show_progress", "accept_key", "reject_key", "skip_key"],
+ supports_keybindings=True,
+ description="Binary accept/reject triage for rapid data curation"
+ ),
+ SchemaDefinition(
+ name="event_annotation",
+ generator=generate_event_annotation_layout,
+ required_fields=["name", "description", "event_types", "span_schema"],
+ optional_fields=["visual_display"],
+ supports_keybindings=False,
+ description="N-ary event annotation with triggers and typed arguments"
+ ),
+ SchemaDefinition(
+ name="tiered_annotation",
+ generator=generate_tiered_annotation_layout,
+ required_fields=["name", "description", "tiers", "source_field"],
+ optional_fields=["media_type", "tier_height", "show_tier_labels", "collapsed_tiers", "zoom_enabled", "playback_rate_control", "overview_height"],
+ supports_keybindings=True,
+ description="Hierarchical multi-tier annotation for audio/video (ELAN-style)"
+ ),
+ SchemaDefinition(
+ name="bws",
+ generator=generate_bws_layout,
+ required_fields=["name", "description"],
+ optional_fields=["best_description", "worst_description", "tuple_size", "sequential_key_binding", "label_requirement"],
+ supports_keybindings=True,
+ description="Best-Worst Scaling: select the best and worst item from a set"
+ ),
+ SchemaDefinition(
+ name="soft_label",
+ generator=generate_soft_label_layout,
+ required_fields=["name", "description", "labels"],
+ optional_fields=["total", "min_per_label", "show_distribution_chart"],
+ supports_keybindings=False,
+ description="Probability distribution across labels via constrained sliders"
+ ),
+ SchemaDefinition(
+ name="confidence",
+ generator=generate_confidence_layout,
+ required_fields=["name", "description"],
+ optional_fields=["target_schema", "scale_type", "scale_points", "labels", "min_value", "max_value", "step", "left_label", "right_label"],
+ supports_keybindings=False,
+ description="Confidence rating meta-annotation for any primary annotation"
+ ),
+ SchemaDefinition(
+ name="constant_sum",
+ generator=generate_constant_sum_layout,
+ required_fields=["name", "description", "labels"],
+ optional_fields=["total_points", "min_per_item", "input_type"],
+ supports_keybindings=False,
+ description="Allocate a fixed budget of points across categories"
+ ),
+ SchemaDefinition(
+ name="semantic_differential",
+ generator=generate_semantic_differential_layout,
+ required_fields=["name", "description", "pairs"],
+ optional_fields=["scale_points"],
+ supports_keybindings=False,
+ description="Bipolar adjective scales for measuring connotative meaning"
+ ),
+ SchemaDefinition(
+ name="ranking",
+ generator=generate_ranking_layout,
+ required_fields=["name", "description", "labels"],
+ optional_fields=["allow_ties"],
+ supports_keybindings=False,
+ description="Drag-and-drop ranking of items by preference or relevance"
+ ),
+ SchemaDefinition(
+ name="range_slider",
+ generator=generate_range_slider_layout,
+ required_fields=["name", "description"],
+ optional_fields=["min_value", "max_value", "step", "left_label", "right_label", "show_values"],
+ supports_keybindings=False,
+ description="Dual-thumb slider for selecting an acceptable range"
+ ),
+ SchemaDefinition(
+ name="hierarchical_multiselect",
+ generator=generate_hierarchical_multiselect_layout,
+ required_fields=["name", "description", "taxonomy"],
+ optional_fields=["auto_select_children", "auto_select_parent", "show_search", "max_selections"],
+ supports_keybindings=False,
+ description="Hierarchical tree-structured multi-label selection"
+ ),
+ SchemaDefinition(
+ name="vas",
+ generator=generate_vas_layout,
+ required_fields=["name", "description"],
+ optional_fields=["left_label", "right_label", "min_value", "max_value", "show_value"],
+ supports_keybindings=False,
+ description="Continuous visual analog scale for fine-grained magnitude estimation"
+ ),
+ SchemaDefinition(
+ name="extractive_qa",
+ generator=generate_extractive_qa_layout,
+ required_fields=["name", "description"],
+ optional_fields=["question_field", "passage_field", "allow_unanswerable", "highlight_color"],
+ supports_keybindings=False,
+ description="SQuAD-style extractive question answering with answer span highlighting"
+ ),
+ SchemaDefinition(
+ name="rubric_eval",
+ generator=generate_rubric_eval_layout,
+ required_fields=["name", "description", "criteria"],
+ optional_fields=["scale_points", "scale_labels", "show_overall"],
+ supports_keybindings=False,
+ description="Multi-criteria rubric evaluation grid for LLM and text quality assessment"
+ ),
+ SchemaDefinition(
+ name="text_edit",
+ generator=generate_text_edit_layout,
+ required_fields=["name", "description"],
+ optional_fields=["source_field", "show_diff", "show_edit_distance", "allow_reset"],
+ supports_keybindings=False,
+ description="Inline text editing with diff tracking for post-editing and correction tasks"
+ ),
+ SchemaDefinition(
+ name="error_span",
+ generator=generate_error_span_layout,
+ required_fields=["name", "description", "error_types"],
+ optional_fields=["severities", "show_score", "max_score"],
+ supports_keybindings=False,
+ description="MQM-style error span annotation with typed severity for quality evaluation"
+ ),
+ SchemaDefinition(
+ name="card_sort",
+ generator=generate_card_sort_layout,
+ required_fields=["name", "description"],
+ optional_fields=["mode", "groups", "items_field", "allow_empty_groups", "allow_multiple"],
+ supports_keybindings=False,
+ description="Drag-and-drop card sorting into predefined or user-created groups"
+ ),
+ SchemaDefinition(
+ name="conjoint",
+ generator=generate_conjoint_layout,
+ required_fields=["name", "description"],
+ optional_fields=["profiles_per_set", "attributes", "show_none_option", "profiles_field"],
+ supports_keybindings=False,
+ description="Discrete choice conjoint analysis with side-by-side profile comparison"
+ ),
+ SchemaDefinition(
+ name="trajectory_eval",
+ generator=generate_trajectory_eval_layout,
+ required_fields=["name", "description"],
+ optional_fields=["steps_key", "step_text_key", "correctness_options", "error_types", "severities", "show_score", "max_score"],
+ supports_keybindings=False,
+ description="Per-step trajectory evaluation with error taxonomy and severity scoring"
+ ),
+ SchemaDefinition(
+ name="process_reward",
+ generator=generate_process_reward_layout,
+ required_fields=["name", "description"],
+ optional_fields=["steps_key", "step_text_key", "mode"],
+ supports_keybindings=False,
+ description="Binary per-step process reward signals for PRM training"
+ ),
+ SchemaDefinition(
+ name="trajectory_edit",
+ generator=generate_trajectory_edit_layout,
+ required_fields=["name", "description"],
+ optional_fields=["steps_key", "step_text_key", "editable_fields", "show_diff", "show_edit_distance", "allow_reset", "require_reason_on_edit", "edit_final_answer", "final_answer_key"],
+ supports_keybindings=False,
+ description="Per-step trajectory correction/editing for SFT/DPO training data"
+ ),
+ SchemaDefinition(
+ name="code_review",
+ generator=generate_code_review_layout,
+ required_fields=["name", "description"],
+ optional_fields=["comment_categories", "verdict_options", "file_rating_dimensions"],
+ supports_keybindings=False,
+ description="GitHub PR-style code review with inline comments and file ratings"
+ ),
+ ]
+
+ for schema in schemas:
+ schema_registry.register(schema)
+
+ logger.debug(f"Registered {len(schemas)} built-in schemas")
+
+
+# Auto-register built-in schemas on import
+_register_builtin_schemas()
diff --git a/potato/server_utils/schemas/rubric_eval.py b/potato/server_utils/schemas/rubric_eval.py
new file mode 100644
index 0000000000000000000000000000000000000000..4676400d28b2ec0fc6d50ba65b057f130a956834
--- /dev/null
+++ b/potato/server_utils/schemas/rubric_eval.py
@@ -0,0 +1,138 @@
+"""
+Multi-Criteria Rubric Evaluation Layout
+
+Rate items on multiple criteria simultaneously in a structured grid.
+THE missing schema for LLM evaluation โ enables MT-Bench-style multi-dimensional scoring.
+
+Research: Zheng et al. (2023) "Judging LLM-as-a-judge with MT-Bench"; Ke et al. (2024) "CritiqueLLM".
+"""
+
+import logging
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_SCALE_POINTS = 5
+DEFAULT_SCALE_LABELS = ["Poor", "Below Average", "Average", "Good", "Excellent"]
+
+
+def generate_rubric_eval_layout(annotation_scheme):
+ """
+ Generate HTML for a Multi-Criteria Rubric Evaluation interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - scale_points: Number of scale points (default 5)
+ - scale_labels: Labels for each scale point
+ - criteria: List of {name, description} dicts
+ - show_overall: Whether to include an "Overall" row
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_rubric_eval_layout_internal)
+
+
+def _generate_rubric_eval_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme['name']
+ description = annotation_scheme['description']
+ scale_points = annotation_scheme.get('scale_points', DEFAULT_SCALE_POINTS)
+ scale_labels = annotation_scheme.get('scale_labels', DEFAULT_SCALE_LABELS[:scale_points])
+ criteria = annotation_scheme.get('criteria', [])
+ show_overall = annotation_scheme.get('show_overall', False)
+
+ if not criteria:
+ raise ValueError(f"rubric_eval schema '{schema_name}' requires at least one criterion in 'criteria'")
+
+ # Pad or trim scale_labels to match scale_points
+ while len(scale_labels) < scale_points:
+ scale_labels.append(f"{len(scale_labels) + 1}")
+ scale_labels = scale_labels[:scale_points]
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+
+ # Build header row
+ header_cells = ''
+ for i, label in enumerate(scale_labels):
+ header_cells += f''
+
+ # Build criterion rows
+ rows_html = ""
+ all_criteria = list(criteria)
+ if show_overall:
+ all_criteria.append({"name": "overall", "description": "Overall quality"})
+
+ for criterion in all_criteria:
+ crit_name = criterion['name']
+ crit_desc = criterion.get('description', '')
+ crit_label = f"{schema_name}:{crit_name}"
+
+ identifiers = generate_element_identifier(schema_name, crit_name, "radio")
+ # Override name: each criterion row needs its own radio group
+ # (generate_element_identifier uses schema-only name for radios,
+ # but rubric_eval needs per-criterion groups)
+ row_group_name = f"{escape_html_content(schema_name)}:::{escape_html_content(crit_name)}"
+
+ cells = f"""
+
+ {escape_html_content(crit_name)}
+ {f'{escape_html_content(crit_desc)}
' if crit_desc else ''}
+
+ """
+
+ for i in range(scale_points):
+ value = str(i + 1)
+ radio_id = f"{identifiers['id']}-{value}"
+ cells += f"""
+
+
+
+ """
+
+ is_overall_class = " rubric-overall-row" if crit_name == "overall" else ""
+ rows_html += f'{cells} '
+
+ html = f"""
+
+ """
+
+ logger.info(f"Generated rubric eval layout for {schema_name} with {len(criteria)} criteria")
+ return html, []
diff --git a/potato/server_utils/schemas/select.py b/potato/server_utils/schemas/select.py
new file mode 100644
index 0000000000000000000000000000000000000000..37fdfadceef91e744ac4594930c3ecbfac348d02
--- /dev/null
+++ b/potato/server_utils/schemas/select.py
@@ -0,0 +1,93 @@
+"""
+Select Layout
+"""
+
+import os
+from pathlib import Path
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper, get_dynamic_ai_help
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+
+def generate_select_layout(annotation_scheme):
+ """
+ Generate HTML for a select dropdown interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - labels: List of options or path to file containing options
+ - use_predefined_labels: Use predefined label sets (country, ethnicity, religion)
+ - label_requirement (dict): Optional validation settings
+ - required (bool): Whether selection is mandatory
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the select interface
+ key_bindings: Empty list (no keyboard shortcuts)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_select_layout_internal)
+
+def _generate_select_layout_internal(annotation_scheme):
+ """
+ Internal function to generate select layout after validation.
+ """
+ # Generate consistent identifiers
+ identifiers = generate_element_identifier(annotation_scheme["name"], "select-one", "select")
+ validation = generate_validation_attribute(annotation_scheme)
+
+ # Get layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ schematic = (
+ f'\n"
+ return schematic, []
diff --git a/potato/server_utils/schemas/semantic_differential.py b/potato/server_utils/schemas/semantic_differential.py
new file mode 100644
index 0000000000000000000000000000000000000000..39fd6e4da41082fed5df60daa3ea70840a258177
--- /dev/null
+++ b/potato/server_utils/schemas/semantic_differential.py
@@ -0,0 +1,137 @@
+"""
+Semantic Differential Layout
+
+Generates bipolar adjective scales arranged in a matrix.
+Each row has a left and right pole with radio buttons between them.
+
+Research basis:
+- Osgood, Suci & Tannenbaum (1957) "The Measurement of Meaning"
+ University of Illinois Press
+- Mohammad (2018) "Obtaining Reliable Human Ratings of Valence, Arousal, and Dominance
+ for 20,000 English Words" ACL
+"""
+
+import logging
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+)
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_SCALE_POINTS = 7
+
+
+def generate_semantic_differential_layout(annotation_scheme):
+ """
+ Generate HTML for a semantic differential annotation interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - pairs: List of [left_adjective, right_adjective] pairs
+ - scale_points: Number of points per scale (default 7)
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_semantic_differential_layout_internal)
+
+
+def _generate_semantic_differential_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme["name"]
+ safe_schema = escape_html_content(schema_name)
+ description = annotation_scheme["description"]
+ pairs = annotation_scheme.get("pairs", [])
+ scale_points = annotation_scheme.get("scale_points", DEFAULT_SCALE_POINTS)
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+
+ if not pairs:
+ raise ValueError(f"semantic_differential schema '{schema_name}' requires 'pairs'")
+
+ for pair in pairs:
+ if not isinstance(pair, list) or len(pair) != 2:
+ raise ValueError(f"Each pair must be a list of two strings, got: {pair}")
+
+ html = f"""
+
+ """
+
+ # JS for mutual exclusivity within each pair row
+ html += """
+
+ """
+
+ key_bindings = []
+ logger.info(f"Generated semantic_differential layout for {schema_name} with {len(pairs)} pairs")
+ return html, key_bindings
diff --git a/potato/server_utils/schemas/slider.py b/potato/server_utils/schemas/slider.py
new file mode 100644
index 0000000000000000000000000000000000000000..a6ee6971ac218259ff9137cefd8ab6ba9b77d446
--- /dev/null
+++ b/potato/server_utils/schemas/slider.py
@@ -0,0 +1,923 @@
+# """
+# slider Layout
+# """
+
+# # Needed for the fall-back radio layout
+# from potato.ai.ai_help_wrapper import get_ai_wrapper, get_dynamic_ai_help
+# from .radio import generate_radio_layout
+# from .identifier_utils import (
+# safe_generate_layout,
+# generate_element_identifier,
+# escape_html_content
+# )
+
+# def test_and_get(key, d):
+# val = d[key]
+# try:
+# return int(val)
+# except:
+# raise Exception(
+# 'Slider scale %s\'s value for "%s" is not an int' % (d["name"], key)
+# )
+
+# def generate_slider_layout(annotation_scheme):
+# """
+# Generate HTML for a slider input interface.
+
+# Args:
+# annotation_scheme (dict): Configuration including:
+# - name: Schema identifier
+# - description: Display description
+# - starting_value: Initial slider value
+# - min_value: Minimum allowed value
+# - max_value: Maximum allowed value
+# - show_labels: Whether to show min/max labels
+# - labels: If present, fall back to radio layout
+
+# Returns:
+# tuple: (html_string, key_bindings)
+# html_string: Complete HTML for the slider interface
+# key_bindings: Empty list (no keyboard shortcuts)
+# """
+# return safe_generate_layout(annotation_scheme, generate_slider_layout_internal)
+
+# def generate_slider_layout_internal(annotation_scheme):
+# from .identifier_utils import escape_html_content, generate_element_identifier
+
+# if "labels" in annotation_scheme:
+# return generate_radio_layout(annotation_scheme, horizontal=False)
+
+# for required in ["starting_value", "min_value", "max_value"]:
+# if required not in annotation_scheme:
+# raise Exception(
+# f'Slider scale for "{annotation_scheme["name"]}" did not include {required}'
+# )
+
+# min_value = test_and_get("min_value", annotation_scheme)
+# max_value = test_and_get("max_value", annotation_scheme)
+# starting_value = test_and_get("starting_value", annotation_scheme)
+
+# if min_value >= max_value:
+# raise Exception(
+# f'Slider scale for "{annotation_scheme["name"]}" must have minimum value < max value ({min_value} >= {max_value})'
+# )
+
+# show_labels = annotation_scheme.get("show_labels", True)
+# min_label = str(min_value) if show_labels else ''
+# max_label = str(max_value) if show_labels else ''
+
+# identifiers = generate_element_identifier(annotation_scheme["name"], "slider", "range")
+
+# # Get step from annotation_scheme or default to 5
+# step_value = annotation_scheme.get("step", 1)
+# print(step_value)
+# max_tick = annotation_scheme.get("maxTick", 8)
+# print("max_tick")
+# print(max_tick)
+
+# schematic = f"""
+
+#
+
+#
+# """
+# key_bindings = []
+# return schematic, key_bindings
+
+
+"""
+slider Layout
+"""
+
+import logging
+
+# Needed for the fall-back radio layout
+from potato.ai.ai_help_wrapper import get_ai_wrapper, get_dynamic_ai_help
+from .radio import generate_radio_layout
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ escape_html_content
+)
+
+logger = logging.getLogger(__name__)
+
+def test_and_get(key, d):
+ val = d[key]
+ try:
+ return int(val)
+ except:
+ raise Exception(
+ 'Slider scale %s\'s value for "%s" is not an int' % (d["name"], key)
+ )
+
+def generate_slider_layout(annotation_scheme):
+ """
+ Generate HTML for a slider input interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - starting_value: Initial slider value
+ - min_value: Minimum allowed value
+ - max_value: Maximum allowed value
+ - show_labels: Whether to show min/max labels
+ - labels: If present, fall back to radio layout
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the slider interface
+ key_bindings: Empty list (no keyboard shortcuts)
+ """
+ return safe_generate_layout(annotation_scheme, generate_slider_layout_internal)
+
+def generate_slider_layout_internal(annotation_scheme):
+ from .identifier_utils import escape_html_content, generate_element_identifier, generate_validation_attribute
+
+ if "labels" in annotation_scheme:
+ return generate_radio_layout(annotation_scheme, horizontal=False)
+
+ for required in ["starting_value", "min_value", "max_value"]:
+ if required not in annotation_scheme:
+ raise Exception(
+ f'Slider scale for "{annotation_scheme["name"]}" did not include {required}'
+ )
+
+ min_value = test_and_get("min_value", annotation_scheme)
+ max_value = test_and_get("max_value", annotation_scheme)
+ starting_value = test_and_get("starting_value", annotation_scheme)
+
+ if min_value >= max_value:
+ raise Exception(
+ f'Slider scale for "{annotation_scheme["name"]}" must have minimum value < max value ({min_value} >= {max_value})'
+ )
+
+ show_labels = annotation_scheme.get("show_labels", True)
+ min_label = str(min_value) if show_labels else ''
+ max_label = str(max_value) if show_labels else ''
+
+ identifiers = generate_element_identifier(annotation_scheme["name"], "slider", "range")
+ validation = generate_validation_attribute(annotation_scheme)
+
+ # Get step from annotation_scheme or default to 1
+ step_value = annotation_scheme.get("step", 1)
+ max_tick = annotation_scheme.get("maxTick", 8)
+
+ schematic = f"""
+
+
+
+ """
+ key_bindings = []
+ return schematic, key_bindings
\ No newline at end of file
diff --git a/potato/server_utils/schemas/soft_label.py b/potato/server_utils/schemas/soft_label.py
new file mode 100644
index 0000000000000000000000000000000000000000..6b84d9129b9951ee54b513011f1cb3e370d92db4
--- /dev/null
+++ b/potato/server_utils/schemas/soft_label.py
@@ -0,0 +1,237 @@
+"""
+Soft Label / Probability Distribution Layout
+
+Generates a set of constrained sliders where annotators distribute probability
+mass across labels. All sliders are constrained to sum to a fixed total (default 100).
+
+Research basis:
+- Fornaciari et al. (2021) "Beyond Black & White: Leveraging Annotator Disagreement
+ via Soft-Label Multi-Task Learning" ACL
+- Plank et al. (2014) "Linguistically Debatable or Just Plain Wrong?" ACL
+"""
+
+import logging
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+ generate_tooltip_html,
+)
+
+logger = logging.getLogger(__name__)
+
+# Defaults
+DEFAULT_TOTAL = 100
+DEFAULT_MIN_PER_LABEL = 0
+
+
+def generate_soft_label_layout(annotation_scheme):
+ """
+ Generate HTML for a soft label / probability distribution interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - labels: List of label names (strings or dicts with 'name')
+ - total: Sum constraint (default 100)
+ - min_per_label: Minimum per label (default 0)
+ - show_distribution_chart: Show bar chart (default true)
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_soft_label_layout_internal)
+
+
+def _generate_soft_label_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme["name"]
+ safe_schema = escape_html_content(schema_name)
+ description = annotation_scheme["description"]
+ labels = annotation_scheme.get("labels", [])
+ total = annotation_scheme.get("total", DEFAULT_TOTAL)
+ min_per_label = annotation_scheme.get("min_per_label", DEFAULT_MIN_PER_LABEL)
+ show_chart = annotation_scheme.get("show_distribution_chart", True)
+
+ if not labels:
+ raise ValueError(f"soft_label schema '{schema_name}' requires 'labels'")
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+
+ # Normalize labels
+ label_names = []
+ for lbl in labels:
+ if isinstance(lbl, str):
+ label_names.append(lbl)
+ elif isinstance(lbl, dict) and "name" in lbl:
+ label_names.append(lbl["name"])
+ else:
+ raise ValueError(f"Invalid label format: {lbl}")
+
+ # Calculate initial equal distribution
+ initial_value = total // len(label_names)
+ remainder = total - (initial_value * len(label_names))
+
+ html = f"""
+
+ """
+
+ # Inline JS for sum constraint
+ html += f"""
+
+ """
+
+ key_bindings = []
+ logger.info(f"Generated soft_label layout for {schema_name} with {len(label_names)} labels")
+ return html, key_bindings
diff --git a/potato/server_utils/schemas/span.py b/potato/server_utils/schemas/span.py
new file mode 100644
index 0000000000000000000000000000000000000000..85a6d646b29cfa5fb5256b3d2b362c9d04b30644
--- /dev/null
+++ b/potato/server_utils/schemas/span.py
@@ -0,0 +1,566 @@
+"""
+Span Layout
+"""
+
+import logging
+from collections.abc import Mapping
+from collections import defaultdict
+from potato.ai.ai_help_wrapper import get_ai_wrapper, get_dynamic_ai_help
+from potato.server_utils.config_module import config
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+
+from item_state_management import SpanAnnotation
+
+logger = logging.getLogger(__name__)
+
+SPAN_COLOR_PALETTE = [
+ "(110, 86, 207)", # Primary purple #6E56CF
+ "(239, 68, 68)", # Destructive red #EF4444
+ "(113, 113, 122)", # Gray #71717A
+ "(245, 158, 11)", # Amber #F59E0B
+ "(16, 185, 129)", # Success green #10B981
+ "(59, 130, 246)", # Blue #3B82F6
+ "(220, 38, 38)", # Red #DC2626
+ "(139, 92, 246)", # Purple #8B5CF6
+ "(156, 163, 175)", # Light gray #9CA3AF
+ "(107, 114, 128)", # Medium gray #6B7280
+ "(55, 65, 81)", # Dark gray #374151
+ "(249, 115, 22)", # Orange #F97316
+ "(6, 182, 212)", # Cyan #06B6D4
+ "(236, 72, 153)", # Pink #EC4899
+ "(5, 150, 105)", # Dark green #059669
+ "(124, 58, 237)", # Violet #7C3AED
+ "(22, 163, 74)", # Green #16A34A
+ "(234, 88, 12)", # Dark orange #EA580C
+ "(37, 99, 235)", # Blue #2563EB
+ "(127, 29, 29)", # Dark red #7F1D1D
+ "(168, 85, 247)", # Purple #A855F7
+ "(34, 197, 94)", # Green #22C55E
+]
+
+span_counter = 0
+SPAN_COLOR_PALETTE_LENGTH = len(SPAN_COLOR_PALETTE)
+
+
+def reset_span_counter():
+ """Reset the span color counter to 0. Used for test isolation."""
+ global span_counter
+ span_counter = 0
+
+def get_span_color(schema, span_label):
+ """
+ Returns the color of a span with this label as a string with an RGB triple
+ in parentheses, or None if the span is unmapped.
+ """
+
+ if "ui" not in config or "spans" not in config["ui"]:
+ return None
+
+ span_ui = config["ui"]["spans"]
+
+ if "span_colors" not in span_ui:
+ return None
+
+ if schema in span_ui["span_colors"]:
+ schema_colors = span_ui["span_colors"][schema]
+ if span_label in schema_colors:
+ return schema_colors[span_label]
+
+ return None
+
+
+def set_span_color(schema, span_label, color):
+ """
+ Sets the color of a span with this label as a string with an RGB triple in parentheses.
+
+ :color: a string containing an RGB triple in parentheses
+ """
+ if "ui" not in config:
+ ui = {}
+ config["ui"] = ui
+ else:
+ ui = config["ui"]
+
+ if "spans" not in ui:
+ span_ui = {}
+ ui["spans"] = span_ui
+ else:
+ span_ui = ui["spans"]
+
+ if "span_colors" not in span_ui:
+ span_colors = defaultdict(dict)
+ span_ui["span_colors"] = span_colors
+ else:
+ span_colors = span_ui["span_colors"]
+
+ # Ensure the schema key exists (span_colors may be a regular dict, not defaultdict)
+ if schema not in span_colors:
+ span_colors[schema] = {}
+ span_colors[schema][span_label] = color
+
+def _generate_span_layout_internal(annotation_scheme, horizontal=False):
+ """
+ Internal function to generate span layout after validation.
+
+ Configuration options:
+ allow_discontinuous (bool): Enable discontinuous span selection via Ctrl/Cmd+click.
+ When enabled, users can hold Ctrl (Windows/Linux) or Cmd (Mac) and click to
+ add additional non-contiguous text ranges to an existing span annotation.
+ Default: false
+
+ entity_linking (dict): Configuration for knowledge base entity linking.
+ When enabled, users can link annotated spans to external knowledge bases
+ like Wikidata or UMLS. Configuration options:
+ - enabled (bool): Whether entity linking is enabled. Default: false
+ - knowledge_bases (list): List of KB configurations, each with:
+ - name (str): Display name for the KB
+ - type (str): KB type ("wikidata", "umls", "rest")
+ - api_key (str): Optional API key for authenticated services
+ - language (str): Language code for results. Default: "en"
+ - auto_search (bool): Automatically search when span is created. Default: true
+ - required (bool): Require entity link before saving span. Default: false
+
+ Example:
+ entity_linking:
+ enabled: true
+ knowledge_bases:
+ - name: wikidata
+ type: wikidata
+ language: en
+ - name: umls
+ type: umls
+ api_key: ${UMLS_API_KEY}
+ auto_search: true
+ required: false
+ """
+ import json as json_module
+
+ # Initialize form wrapper
+ scheme_name = annotation_scheme["name"]
+
+ # Get target_field for multi-span support (optional)
+ target_field = annotation_scheme.get("target_field", "")
+ target_field_attr = f' data-target-field="{escape_html_content(target_field)}"' if target_field else ""
+
+ # Check for discontinuous span support
+ allow_discontinuous = annotation_scheme.get("allow_discontinuous", False)
+ discontinuous_attr = ' data-allow-discontinuous="true"' if allow_discontinuous else ""
+
+ # Check for entity linking support
+ entity_linking = annotation_scheme.get("entity_linking", {})
+ entity_linking_enabled = entity_linking.get("enabled", False)
+ entity_linking_attr = ""
+ if entity_linking_enabled:
+ # Serialize entity_linking config to JSON for frontend
+ el_config = {
+ "enabled": True,
+ "knowledge_bases": entity_linking.get("knowledge_bases", []),
+ "auto_search": entity_linking.get("auto_search", True),
+ "required": entity_linking.get("required", False),
+ "multi_select": entity_linking.get("multi_select", False)
+ }
+ el_json = json_module.dumps(el_config)
+ entity_linking_attr = f' data-entity-linking=\'{escape_html_content(el_json)}\''
+
+ # Check for show_span_labels option (default: true)
+ show_span_labels = annotation_scheme.get("show_span_labels", True)
+ show_labels_attr = '' if show_span_labels else ' data-show-span-labels="false"'
+
+ # Get layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ schematic = f"""
+ "
+ return schematic, key_bindings
+
+def _generate_tooltip(label_data):
+ """
+ Generate tooltip HTML attribute from label data.
+
+ Args:
+ label_data (dict): Label configuration containing tooltip information
+
+ Returns:
+ str: Tooltip HTML attribute or empty string if no tooltip
+ """
+ tooltip_text = ""
+ if "tooltip" in label_data:
+ tooltip_text = label_data["tooltip"]
+ elif "tooltip_file" in label_data:
+ try:
+ with open(label_data["tooltip_file"], "rt", encoding="utf-8") as f:
+ tooltip_text = "".join(f.readlines())
+ except Exception as e:
+ logger.error(f"Failed to read tooltip file: {e}")
+ return ""
+
+ if tooltip_text:
+ escaped_tooltip = escape_html_content(tooltip_text)
+ return f'data-toggle="tooltip" data-html="true" data-placement="top" title="{escaped_tooltip}"'
+ return ""
+
+
+def generate_span_layout(annotation_scheme, horizontal=False):
+ """
+ Generate span layout HTML for the given annotation scheme.
+
+ Args:
+ annotation_scheme (dict): The annotation scheme configuration
+ horizontal (bool): Whether to display horizontally
+
+ Returns:
+ tuple: (HTML string, key bindings list)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_span_layout_internal, horizontal)
+
+
+def render_span_annotations(text, span_annotations, target_field=None):
+ """
+ Render span annotations into HTML with boundary-based algorithm.
+ Supports discontinuous spans with additional_parts.
+
+ Args:
+ text (str): The original text to annotate
+ span_annotations: Dictionary of span_id -> span data, or list of SpanAnnotation objects,
+ or field-keyed dict: {field_key: [span_list]}
+ target_field (str, optional): Filter spans to only those targeting this field
+ Returns:
+ str: HTML with span annotations rendered
+ """
+ if not span_annotations:
+ return text
+
+ # Handle field-keyed format for multi-span mode: {field_key: [spans]}
+ if isinstance(span_annotations, dict):
+ # Check if this is a field-keyed dict (values are lists)
+ first_value = next(iter(span_annotations.values()), None)
+ if isinstance(first_value, list):
+ # Field-keyed format - extract spans for target_field
+ if target_field:
+ field_spans = span_annotations.get(target_field, [])
+ return render_span_annotations(text, field_spans, target_field=None)
+ else:
+ # No target field specified, flatten all spans
+ all_spans = []
+ for field_spans in span_annotations.values():
+ all_spans.extend(field_spans)
+ return render_span_annotations(text, all_spans, target_field=None)
+
+ # Regular dict format: span_id -> span_data
+ sorted_spans = sorted(
+ span_annotations.items(),
+ key=lambda x: x[1].get('start', 0)
+ )
+ else:
+ # Convert list of SpanAnnotation objects to list of tuples
+ spans_as_tuples = []
+ for span in span_annotations:
+ if hasattr(span, 'get_id'):
+ # SpanAnnotation object with methods
+ # Filter by target_field if specified
+ span_target = span.get_target_field() if hasattr(span, 'get_target_field') else None
+ if target_field and span_target and span_target != target_field:
+ continue # Skip spans not targeting this field
+
+ span_id = span.get_id()
+ # Get additional_parts for discontinuous spans
+ additional_parts = []
+ if hasattr(span, 'get_additional_parts'):
+ additional_parts = span.get_additional_parts() or []
+ elif hasattr(span, 'additional_parts'):
+ additional_parts = getattr(span, 'additional_parts', []) or []
+
+ # Get KB entity linking data
+ kb_id = None
+ kb_source = None
+ kb_label = None
+ if hasattr(span, 'get_kb_id'):
+ kb_id = span.get_kb_id()
+ kb_source = span.get_kb_source() if hasattr(span, 'get_kb_source') else None
+ kb_label = span.get_kb_label() if hasattr(span, 'get_kb_label') else None
+ elif hasattr(span, 'kb_id'):
+ kb_id = getattr(span, 'kb_id', None)
+ kb_source = getattr(span, 'kb_source', None)
+ kb_label = getattr(span, 'kb_label', None)
+
+ span_data = {
+ 'schema': span.get_schema() if hasattr(span, 'get_schema') else getattr(span, 'schema', ''),
+ 'name': span.get_name() if hasattr(span, 'get_name') else getattr(span, 'name', ''),
+ 'title': span.get_title() if hasattr(span, 'get_title') else getattr(span, 'title', ''),
+ 'start': span.get_start() if hasattr(span, 'get_start') else getattr(span, 'start', 0),
+ 'end': span.get_end() if hasattr(span, 'get_end') else getattr(span, 'end', 0),
+ 'target_field': span_target,
+ 'additional_parts': additional_parts,
+ 'kb_id': kb_id,
+ 'kb_source': kb_source,
+ 'kb_label': kb_label,
+ }
+ elif isinstance(span, dict):
+ # Filter by target_field if specified
+ span_target = span.get('target_field')
+ if target_field and span_target and span_target != target_field:
+ continue # Skip spans not targeting this field
+
+ span_id = span.get('id', f"span_{span.get('start', 0)}_{span.get('end', 0)}")
+ span_data = span
+ else:
+ continue
+ spans_as_tuples.append((span_id, span_data))
+ sorted_spans = sorted(spans_as_tuples, key=lambda x: x[1].get('start', 0))
+
+ # Create boundary points (including additional_parts for discontinuous spans)
+ boundaries = []
+ for span_id, span_data in sorted_spans:
+ # Add primary span boundaries
+ boundaries.append((span_data['start'], 'start', span_id, span_data))
+ boundaries.append((span_data['end'], 'end', span_id, span_data))
+
+ # Add boundaries for additional parts (discontinuous spans)
+ additional_parts = span_data.get('additional_parts', [])
+ for part in additional_parts:
+ # Create a modified span_data for this part that includes discontinuous marker
+ part_data = span_data.copy()
+ part_data['_is_discontinuous_part'] = True
+ boundaries.append((part['start'], 'start', span_id, part_data))
+ boundaries.append((part['end'], 'end', span_id, part_data))
+
+ # Sort boundaries by position
+ boundaries.sort(key=lambda x: x[0])
+
+ # Build the rendered text
+ result = ""
+ current_pos = 0
+ active_spans = []
+
+ for pos, boundary_type, span_id, span_data in boundaries:
+ # Add text before this boundary
+ if pos > current_pos:
+ result += text[current_pos:pos]
+
+ if boundary_type == 'start':
+ # Start a new span
+ active_spans.append(span_id)
+ # Get color for this span
+ color = get_span_color(span_data['schema'], span_data['name'])
+ if not color:
+ color = "(128, 128, 128)" # Default gray
+ # Convert RGB to hex with alpha
+ color_parts = color.strip("()").split(", ")
+ r, g, b = int(color_parts[0]), int(color_parts[1]), int(color_parts[2])
+ hex_color = f"#{r:02x}{g:02x}{b:02x}66" # 66 = 40% alpha to match label background
+
+ # Add target_field attribute if present
+ target_attr = f' data-target-field="{span_data.get("target_field", "")}"' if span_data.get("target_field") else ""
+
+ # Check if this is a discontinuous span part
+ is_discontinuous = span_data.get('_is_discontinuous_part', False) or len(span_data.get('additional_parts', [])) > 0
+ discontinuous_class = ' discontinuous-part' if is_discontinuous else ''
+ discontinuous_attr = ' data-discontinuous="true"' if is_discontinuous else ""
+
+ # Add KB entity linking attributes
+ kb_id = span_data.get('kb_id', '')
+ kb_source = span_data.get('kb_source', '')
+ kb_label = span_data.get('kb_label', '')
+ kb_attr = ""
+ kb_class = ""
+ if kb_id:
+ kb_attr = f' data-kb-id="{escape_html_content(kb_id)}" data-kb-source="{escape_html_content(kb_source)}"'
+ if kb_label:
+ kb_attr += f' data-kb-label="{escape_html_content(kb_label)}"'
+ kb_class = ' has-entity-link'
+
+ result += f''
+ elif boundary_type == 'end':
+ # End the span
+ result += " "
+ # Remove from active spans
+ active_spans = [s for s in active_spans if s != span_id]
+
+ current_pos = pos
+
+ # Add remaining text
+ if current_pos < len(text):
+ result += text[current_pos:]
+
+ return result
+
+
+def get_spans_for_field(span_annotations, target_field):
+ """
+ Extract spans for a specific target field from span annotations.
+
+ Args:
+ span_annotations: Span annotations in any format
+ target_field: The field key to filter by
+
+ Returns:
+ List of spans targeting the specified field
+ """
+ if not span_annotations:
+ return []
+
+ # Handle field-keyed format
+ if isinstance(span_annotations, dict):
+ first_value = next(iter(span_annotations.values()), None)
+ if isinstance(first_value, list):
+ return span_annotations.get(target_field, [])
+
+ # Handle list of SpanAnnotation objects
+ result = []
+ if isinstance(span_annotations, (list, tuple)):
+ for span in span_annotations:
+ if hasattr(span, 'get_target_field'):
+ if span.get_target_field() == target_field:
+ result.append(span)
+ elif isinstance(span, dict) and span.get('target_field') == target_field:
+ result.append(span)
+
+ return result
\ No newline at end of file
diff --git a/potato/server_utils/schemas/span_link.py b/potato/server_utils/schemas/span_link.py
new file mode 100644
index 0000000000000000000000000000000000000000..b36a626c13815f1673e8a91e7f872de0c684de0c
--- /dev/null
+++ b/potato/server_utils/schemas/span_link.py
@@ -0,0 +1,308 @@
+"""
+Span Link Layout
+
+Generates the UI for creating and managing relationships/links between spans.
+This schema type works in conjunction with a span annotation schema to allow
+users to annotate relationships like "PERSON works_for ORGANIZATION".
+"""
+
+import logging
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ escape_html_content,
+ generate_tooltip_html
+)
+from .span import get_span_color, SPAN_COLOR_PALETTE
+
+logger = logging.getLogger(__name__)
+
+# Default colors for link types
+LINK_COLOR_PALETTE = [
+ "#dc2626", # Red
+ "#22c55e", # Green
+ "#a855f7", # Purple
+ "#f59e0b", # Amber
+ "#3b82f6", # Blue
+ "#ec4899", # Pink
+ "#06b6d4", # Cyan
+ "#f97316", # Orange
+ "#8b5cf6", # Violet
+ "#10b981", # Emerald
+]
+
+
+def _generate_span_link_layout_internal(annotation_scheme, horizontal=False):
+ """
+ Internal function to generate span link layout after validation.
+
+ Args:
+ annotation_scheme: Configuration dictionary containing:
+ - name: Schema name
+ - description: Description shown to user
+ - span_schema: Name of the span schema to link
+ - link_types: List of link type definitions with:
+ - name: Link type name (e.g., "WORKS_FOR")
+ - directed: Whether the link is directed (default: false)
+ - allowed_source_labels: Optional list of allowed source span labels
+ - allowed_target_labels: Optional list of allowed target span labels
+ - max_spans: Maximum spans in a link (default: 2, higher for n-ary)
+ - color: Optional color for this link type
+ horizontal: Whether to display horizontally (not used for links)
+
+ Returns:
+ tuple: (HTML string, key bindings list)
+ """
+ scheme_name = annotation_scheme["name"]
+ description = annotation_scheme.get("description", "Create relationships between spans")
+ span_schema = annotation_scheme.get("span_schema", "")
+ link_types = annotation_scheme.get("link_types", [])
+ visual_display = annotation_scheme.get("visual_display", {})
+
+ # Build link types HTML
+ link_types_html = ""
+ key_bindings = []
+
+ for i, link_type in enumerate(link_types):
+ link_name = link_type.get("name", f"Link_{i}")
+ directed = link_type.get("directed", False)
+ max_spans = link_type.get("max_spans", 2)
+ color = link_type.get("color", LINK_COLOR_PALETTE[i % len(LINK_COLOR_PALETTE)])
+
+ # Direction indicator
+ direction_icon = "โ" if directed else "โ"
+ direction_class = "directed" if directed else "undirected"
+
+ # Tooltip with constraints
+ tooltip_parts = []
+ if link_type.get("allowed_source_labels"):
+ tooltip_parts.append(f"Source: {', '.join(link_type['allowed_source_labels'])}")
+ if link_type.get("allowed_target_labels"):
+ tooltip_parts.append(f"Target: {', '.join(link_type['allowed_target_labels'])}")
+ if max_spans > 2:
+ tooltip_parts.append(f"N-ary: up to {max_spans} spans")
+
+ tooltip_attr = ""
+ if tooltip_parts:
+ tooltip_text = "; ".join(tooltip_parts)
+ tooltip_attr = f'data-toggle="tooltip" data-placement="top" title="{escape_html_content(tooltip_text)}"'
+
+ link_types_html += f"""
+
+
+
+
+ {escape_html_content(link_name)}
+ {direction_icon}
+
+
+ """
+
+ # Visual display settings
+ show_arcs = visual_display.get("enabled", True)
+ arc_position = visual_display.get("arc_position", "above")
+ show_labels = visual_display.get("show_labels", True)
+ # Multi-line arc mode: "single_line" (horizontal scroll) or "bracket" (wrapped text with bracket arcs)
+ multi_line_mode = visual_display.get("multi_line_mode", "bracket")
+
+ schematic = f"""
+
+
+
+
+
+
+
Select Link Type:
+
+ {link_types_html}
+
+
+
+
+
+
Selected Spans:
+
+
Click on highlighted spans to select them for linking
+
+
+
+
+
+
+
+
+
+
+ Create Link
+
+
+ Exit Link Mode
+
+
+
+
+
+
+
+
+
+
+ Show link arcs above text
+
+
+
+
+
+
+ """
+
+ return schematic, key_bindings
+
+
+def generate_span_link_layout(annotation_scheme, horizontal=False):
+ """
+ Generate span link layout HTML for the given annotation scheme.
+
+ Args:
+ annotation_scheme (dict): The annotation scheme configuration
+ horizontal (bool): Whether to display horizontally
+
+ Returns:
+ tuple: (HTML string, key bindings list)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_span_link_layout_internal, horizontal)
+
+
+def render_link_arcs(links, span_positions):
+ """
+ Generate SVG arcs for visualizing links between spans.
+
+ Args:
+ links: List of SpanLink objects
+ span_positions: Dictionary mapping span_id -> {x, y, width, height} positions
+
+ Returns:
+ str: SVG markup for the link arcs
+ """
+ if not links or not span_positions:
+ return ""
+
+ svg_paths = []
+
+ for link in links:
+ span_ids = link.get_span_ids()
+ if len(span_ids) < 2:
+ continue
+
+ # Get color for this link type
+ link_color = link.get_properties().get("color", "#dc2626")
+
+ # For binary links, draw a simple arc
+ if len(span_ids) == 2:
+ span1_id, span2_id = span_ids
+ if span1_id not in span_positions or span2_id not in span_positions:
+ continue
+
+ pos1 = span_positions[span1_id]
+ pos2 = span_positions[span2_id]
+
+ # Calculate arc endpoints (center of each span)
+ x1 = pos1["x"] + pos1["width"] / 2
+ y1 = pos1["y"]
+ x2 = pos2["x"] + pos2["width"] / 2
+ y2 = pos2["y"]
+
+ # Calculate control point for the arc
+ mid_x = (x1 + x2) / 2
+ arc_height = min(abs(x2 - x1) / 3, 50) # Arc height proportional to distance
+
+ # Create SVG path
+ if link.is_directed():
+ # Directed link with arrow
+ svg_paths.append(f"""
+
+ """)
+ else:
+ # Undirected link
+ svg_paths.append(f"""
+
+ """)
+
+ # For n-ary links, connect all spans to a central point
+ else:
+ # Calculate center point
+ valid_positions = [span_positions[sid] for sid in span_ids if sid in span_positions]
+ if len(valid_positions) < 2:
+ continue
+
+ center_x = sum(p["x"] + p["width"] / 2 for p in valid_positions) / len(valid_positions)
+ center_y = min(p["y"] for p in valid_positions) - 30 # Above the spans
+
+ # Draw lines from each span to center
+ for span_id in span_ids:
+ if span_id not in span_positions:
+ continue
+ pos = span_positions[span_id]
+ x = pos["x"] + pos["width"] / 2
+ y = pos["y"]
+ svg_paths.append(f"""
+
+ """)
+
+ # Draw central node
+ svg_paths.append(f"""
+
+ """)
+
+ # Wrap in SVG with defs for arrow markers
+ svg = f"""
+
+
+
+
+
+
+ {''.join(svg_paths)}
+
+ """
+
+ return svg
diff --git a/potato/server_utils/schemas/text_edit.py b/potato/server_utils/schemas/text_edit.py
new file mode 100644
index 0000000000000000000000000000000000000000..d5e8deaaa9ed63f5028852fdb0ff6f8a63faf2e5
--- /dev/null
+++ b/potato/server_utils/schemas/text_edit.py
@@ -0,0 +1,265 @@
+"""
+Inline Text Editing / Post-Edit Layout
+
+Annotators directly edit displayed text, with the system tracking insertions,
+deletions, and substitutions as a structured diff. Used for MT post-editing,
+grammar correction, text simplification, and paraphrase generation.
+
+Research: WMT post-editing shared tasks; MQM-APE (COLING 2025).
+"""
+
+import logging
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+def generate_text_edit_layout(annotation_scheme):
+ """
+ Generate HTML for an Inline Text Editing interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - source_field: Field in data containing text to edit
+ - show_diff: Whether to show real-time diff highlighting
+ - show_edit_distance: Whether to show edit distance counter
+ - allow_reset: Whether to show "Reset to original" button
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_text_edit_layout_internal)
+
+
+def _generate_text_edit_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme['name']
+ description = annotation_scheme['description']
+ source_field = annotation_scheme.get('source_field', '')
+ show_diff = annotation_scheme.get('show_diff', True)
+ show_edit_distance = annotation_scheme.get('show_edit_distance', True)
+ allow_reset = annotation_scheme.get('allow_reset', True)
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+ identifiers = generate_element_identifier(schema_name, schema_name, "hidden")
+
+ reset_btn = ""
+ if allow_reset:
+ reset_btn = f"""
+
+ Reset to Original
+
+ """
+
+ diff_display = ""
+ if show_diff:
+ diff_display = f"""
+
+ """
+
+ edit_distance_display = ""
+ if show_edit_distance:
+ edit_distance_display = f"""
+
+ Words changed: 0
+ Chars changed: 0
+
+ """
+
+ html = f"""
+
+
+
+ """
+
+ logger.info(f"Generated text edit layout for {schema_name}")
+ return html, []
diff --git a/potato/server_utils/schemas/textbox.py b/potato/server_utils/schemas/textbox.py
new file mode 100644
index 0000000000000000000000000000000000000000..fe685877729780e672f9fa1e86232fcd7f075ce9
--- /dev/null
+++ b/potato/server_utils/schemas/textbox.py
@@ -0,0 +1,197 @@
+"""
+Textbox Layout
+
+Supports enhanced rationale/justification features:
+- min_chars: Minimum character count with counter display
+- show_char_count: Show character counter below textarea
+- collapsible: Start collapsed, expand on click
+- target_schema: Visual grouping with a preceding schema
+- placeholder: Placeholder text in the input
+"""
+
+import logging
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper, get_dynamic_ai_help
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+
+logger = logging.getLogger(__name__)
+
+def generate_textbox_layout(annotation_scheme):
+ """
+ Generate HTML for a textbox input interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - labels: Optional list of labels for multiple textboxes
+ - label_requirement (dict): Optional validation settings
+ - required (bool): Whether input is mandatory
+ - textarea (dict): Optional textarea configuration
+ - on (bool): Whether to use textarea instead of input
+ - rows (int): Number of rows for textarea
+ - cols (int): Number of columns for textarea
+ - allow_paste (bool): Whether to allow pasting text
+ - custom_css (dict): Optional CSS styling
+ - min_chars (int): Minimum characters required
+ - show_char_count (bool): Show character counter
+ - collapsible (bool): Start collapsed, expand on click
+ - target_schema (str): Name of schema this is a rationale for
+ - placeholder (str): Placeholder text
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the textbox interface
+ key_bindings: Empty list (no keyboard shortcuts)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_textbox_layout_internal)
+
+def _generate_textbox_layout_internal(annotation_scheme):
+ """
+ Internal function to generate textbox layout after validation.
+ """
+ logger.debug(f"Generating textbox layout for schema: {annotation_scheme['name']}")
+
+ # Layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Enhanced features
+ min_chars = annotation_scheme.get("min_chars", 0)
+ show_char_count = annotation_scheme.get("show_char_count", False)
+ collapsible = annotation_scheme.get("collapsible", False)
+ target_schema = annotation_scheme.get("target_schema", "")
+ placeholder = escape_html_content(annotation_scheme.get("placeholder", ""))
+
+ # CSS classes for target_schema grouping
+ target_class = "shadcn-textbox-target-grouped" if target_schema else ""
+
+ # Initialize form wrapper
+ schema_name = annotation_scheme['name']
+ schematic = f"""
+
+ {get_ai_wrapper()}
+
+ """
+
+ # Collapsible wrapper
+ if collapsible:
+ schematic += f"""
+
+ ▶ {escape_html_content(annotation_scheme["description"])}
+
+
+ """
+ else:
+ schematic += f"""
+
{escape_html_content(annotation_scheme["description"])}
+ """
+
+ # Handle custom CSS if provided
+ display_info = annotation_scheme.get("display_config", {})
+ custom_css = ""
+ if "custom_css" in display_info:
+ custom_css_parts = []
+ for k, v in display_info["custom_css"].items():
+ custom_css_parts.append(f"{k}: {v}")
+ custom_css = "; ".join(custom_css_parts)
+
+ # Set paste settings
+ paste_setting = ''
+ if "allow_paste" in annotation_scheme and annotation_scheme["allow_paste"] == False:
+ paste_setting = 'onpaste="alert(\'Pasting is not allowed for the current study\');return false;"'
+
+ # Handle multiple textboxes with different labels
+ if "labels" not in annotation_scheme or annotation_scheme["labels"] == None:
+ labels = ["text_box"]
+ else:
+ labels = annotation_scheme["labels"]
+
+ # Generate input field(s) for each label
+ for label in labels:
+ # Generate consistent identifiers
+ identifiers = generate_element_identifier(annotation_scheme["name"], label, "text")
+ validation = generate_validation_attribute(annotation_scheme)
+
+ # Determine if using textarea
+ is_textarea = False
+ textarea_attrs = ""
+
+ # Check for multiline flag (new format) or textarea.on (old format)
+ if annotation_scheme.get("multiline") or annotation_scheme.get("textarea", {}).get("on"):
+ is_textarea = True
+ # Use multiline config or fall back to textarea config
+ if annotation_scheme.get("multiline"):
+ rows = annotation_scheme.get("rows", "3")
+ cols = annotation_scheme.get("cols", "40")
+ else:
+ rows = annotation_scheme["textarea"].get("rows", "3")
+ cols = annotation_scheme["textarea"].get("cols", "40")
+ textarea_attrs = f"rows='{rows}' cols='{cols}'"
+
+ # Show label if not the default text_box label
+ label_text = "" if label == "text_box" else label
+
+ # Placeholder attribute
+ placeholder_attr = f'placeholder="{placeholder}"' if placeholder else ''
+
+ schematic += f"""
+
+ {f'
{escape_html_content(label_text)} ' if label_text else ''}
+ """
+
+ if is_textarea:
+ # Render textarea for multiline input
+ schematic += f"""
+
+ """
+ else:
+ # Render input for single-line text
+ schematic += f"""
+
+ """
+
+ # Character counter
+ if show_char_count or min_chars:
+ min_label = f"/{min_chars} min" if min_chars else ""
+ schematic += f"""
+
+ 0 {min_label} characters
+
+ """
+
+ schematic += "
"
+
+ # Close collapsible body if needed
+ if collapsible:
+ schematic += "
"
+
+ schematic += " "
+
+ logger.debug(f"Generated textbox schematic for {annotation_scheme['name']}")
+ logger.info(f"Successfully generated textbox layout for {annotation_scheme['name']} with {len(labels)} fields")
+ return schematic, []
diff --git a/potato/server_utils/schemas/tiered_annotation.py b/potato/server_utils/schemas/tiered_annotation.py
new file mode 100644
index 0000000000000000000000000000000000000000..443bf06302be1862b7137a03bd973426e041fbe1
--- /dev/null
+++ b/potato/server_utils/schemas/tiered_annotation.py
@@ -0,0 +1,555 @@
+"""
+Tiered Annotation Schema
+
+Generates a multi-tier timeline annotation interface for audio/video content.
+This schema supports ELAN-style hierarchical annotation with independent and
+dependent tiers, where dependent tiers can have various constraint relationships
+to their parent tiers.
+
+Features:
+- Multi-row timeline visualization with waveform (via Peaks.js)
+- Independent tiers for direct time alignment
+- Dependent tiers with parent-child constraints
+- Tier selection with dynamic label buttons
+- Constraint validation in real-time
+- EAF and TextGrid export support
+
+Example configuration:
+ annotation_schemes:
+ - annotation_type: tiered_annotation
+ name: linguistic_tiers
+ description: "Multi-tier linguistic annotation"
+ source_field: audio_url
+ media_type: audio
+ tiers:
+ - name: utterance
+ tier_type: independent
+ labels:
+ - name: Speaker_A
+ color: "#4ECDC4"
+ - name: Speaker_B
+ color: "#FF6B6B"
+ - name: word
+ tier_type: dependent
+ parent_tier: utterance
+ constraint_type: time_subdivision
+ labels:
+ - name: Word
+ color: "#95E1D3"
+"""
+
+import json
+import logging
+from typing import Dict, Any, Tuple, List
+
+from .identifier_utils import safe_generate_layout, escape_html_content
+
+logger = logging.getLogger(__name__)
+
+# Default colors for tier labels
+DEFAULT_COLORS = [
+ "#4ECDC4", # Teal
+ "#FF6B6B", # Red
+ "#45B7D1", # Blue
+ "#96CEB4", # Green
+ "#FFEAA7", # Yellow
+ "#DDA0DD", # Plum
+ "#95A5A6", # Gray
+ "#F39C12", # Orange
+ "#9B59B6", # Purple
+ "#3498DB", # Light Blue
+]
+
+# Valid constraint types
+VALID_CONSTRAINT_TYPES = [
+ "time_subdivision",
+ "included_in",
+ "symbolic_association",
+ "symbolic_subdivision",
+ "none"
+]
+
+
+def generate_tiered_annotation_layout(
+ annotation_scheme: Dict[str, Any]
+) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Generate HTML for a tiered annotation interface.
+
+ Args:
+ annotation_scheme: Configuration dictionary including:
+ - name: Schema identifier
+ - description: Display description
+ - source_field: Field name containing media URL
+ - media_type: "audio" or "video"
+ - tiers: List of tier definitions
+ - tier_height: Height per tier row in pixels (default: 50)
+ - show_tier_labels: Show tier name labels (default: True)
+ - collapsed_tiers: List of tier names to start collapsed
+ - zoom_enabled: Enable zoom controls (default: True)
+ - playback_rate_control: Show playback speed controls (default: True)
+
+ Returns:
+ Tuple of (html_string, keybindings_list)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_tiered_annotation_layout_internal)
+
+
+def _generate_tiered_annotation_layout_internal(
+ annotation_scheme: Dict[str, Any]
+) -> Tuple[str, List[Tuple[str, str]]]:
+ """Internal implementation of tiered annotation layout generation."""
+
+ schema_name = annotation_scheme.get("name", "tiered_annotation")
+ description = annotation_scheme.get("description", "")
+ source_field = annotation_scheme.get("source_field", "audio_url")
+ media_type = annotation_scheme.get("media_type", "audio").lower()
+
+ # Validate media type
+ if media_type not in ("audio", "video"):
+ raise ValueError(f"Invalid media_type '{media_type}'. Must be 'audio' or 'video'.")
+
+ # Get tier definitions
+ tiers = annotation_scheme.get("tiers", [])
+ if not tiers:
+ raise ValueError("tiered_annotation requires at least one tier definition")
+
+ # Validate and process tiers
+ processed_tiers = _process_tiers(tiers)
+
+ # Validate tier structure
+ _validate_tier_structure(processed_tiers)
+
+ # Configuration options
+ tier_height = annotation_scheme.get("tier_height", 50)
+ show_tier_labels = annotation_scheme.get("show_tier_labels", True)
+ collapsed_tiers = annotation_scheme.get("collapsed_tiers", [])
+ zoom_enabled = annotation_scheme.get("zoom_enabled", True)
+ playback_rate_control = annotation_scheme.get("playback_rate_control", True)
+ overview_height = annotation_scheme.get("overview_height", 40)
+
+ # Build tier rows HTML
+ tier_rows_html = _generate_tier_rows_html(
+ schema_name, processed_tiers, tier_height, show_tier_labels, collapsed_tiers
+ )
+
+ # Build tier selector options
+ tier_selector_html = _generate_tier_selector_html(processed_tiers)
+
+ # Build JavaScript configuration
+ js_config = {
+ "schemaName": schema_name,
+ "mediaType": media_type,
+ "sourceField": source_field,
+ "tiers": processed_tiers,
+ "tierHeight": tier_height,
+ "showTierLabels": show_tier_labels,
+ "collapsedTiers": collapsed_tiers,
+ "zoomEnabled": zoom_enabled,
+ "playbackRateControl": playback_rate_control,
+ "overviewHeight": overview_height,
+ }
+ config_json = json.dumps(js_config)
+
+ # Generate media element
+ media_element = _generate_media_element(schema_name, media_type)
+
+ # Generate playback controls
+ playback_controls = _generate_playback_controls(
+ schema_name, playback_rate_control, zoom_enabled
+ )
+
+ # Generate the complete HTML
+ html = f'''
+
+
+
+ {escape_html_content(description)}
+
+
+
+ {media_element}
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ {tier_rows_html}
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+'''
+
+ # Define keybindings
+ keybindings = [
+ ("Space", "Play/Pause"),
+ (",", "Step back 1 frame"),
+ (".", "Step forward 1 frame"),
+ ("Delete/Backspace", "Delete selected annotation"),
+ ("Esc", "Deselect annotation"),
+ ]
+
+ return html, keybindings
+
+
+def _process_tiers(tiers: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+ """
+ Process tier definitions, assigning colors and normalizing fields.
+
+ Args:
+ tiers: Raw tier configuration list
+
+ Returns:
+ Processed tier list with normalized fields
+ """
+ processed = []
+ color_index = 0
+
+ for tier in tiers:
+ tier_data = {
+ "name": tier.get("name", f"tier_{len(processed)}"),
+ "tier_type": tier.get("tier_type", "independent").lower(),
+ "parent_tier": tier.get("parent_tier", None),
+ "constraint_type": tier.get("constraint_type", "none").lower(),
+ "description": tier.get("description", ""),
+ "linguistic_type": tier.get("linguistic_type"),
+ "labels": [],
+ }
+
+ # Validate tier_type
+ if tier_data["tier_type"] not in ("independent", "dependent"):
+ raise ValueError(
+ f"Tier '{tier_data['name']}' has invalid tier_type: {tier_data['tier_type']}"
+ )
+
+ # Validate constraint_type
+ if tier_data["constraint_type"] not in VALID_CONSTRAINT_TYPES:
+ raise ValueError(
+ f"Tier '{tier_data['name']}' has invalid constraint_type: {tier_data['constraint_type']}"
+ )
+
+ # Process labels
+ raw_labels = tier.get("labels", [])
+ for label in raw_labels:
+ if isinstance(label, str):
+ label_data = {
+ "name": label,
+ "color": DEFAULT_COLORS[color_index % len(DEFAULT_COLORS)],
+ }
+ color_index += 1
+ elif isinstance(label, dict):
+ label_data = {
+ "name": label.get("name", ""),
+ "color": label.get("color", DEFAULT_COLORS[color_index % len(DEFAULT_COLORS)]),
+ "description": label.get("description", ""),
+ "tooltip": label.get("tooltip", ""),
+ }
+ if not label.get("color"):
+ color_index += 1
+ else:
+ continue
+
+ if label_data["name"]:
+ tier_data["labels"].append(label_data)
+
+ processed.append(tier_data)
+
+ return processed
+
+
+def _validate_tier_structure(tiers: List[Dict[str, Any]]) -> None:
+ """
+ Validate the tier hierarchy structure.
+
+ Args:
+ tiers: Processed tier list
+
+ Raises:
+ ValueError: If structure is invalid
+ """
+ tier_names = {t["name"] for t in tiers}
+
+ for tier in tiers:
+ # Check dependent tier requirements
+ if tier["tier_type"] == "dependent":
+ if not tier.get("parent_tier"):
+ raise ValueError(
+ f"Dependent tier '{tier['name']}' must have a parent_tier"
+ )
+ if tier["parent_tier"] not in tier_names:
+ raise ValueError(
+ f"Tier '{tier['name']}' references unknown parent '{tier['parent_tier']}'"
+ )
+ if tier["parent_tier"] == tier["name"]:
+ raise ValueError(
+ f"Tier '{tier['name']}' cannot be its own parent"
+ )
+
+ # Check for cycles
+ def has_cycle(tier_name: str, visited: set) -> bool:
+ if tier_name in visited:
+ return True
+ visited.add(tier_name)
+ tier = next((t for t in tiers if t["name"] == tier_name), None)
+ if tier and tier.get("parent_tier"):
+ return has_cycle(tier["parent_tier"], visited)
+ return False
+
+ for tier in tiers:
+ if has_cycle(tier["name"], set()):
+ raise ValueError(f"Cycle detected in tier hierarchy involving '{tier['name']}'")
+
+
+def _generate_tier_rows_html(
+ schema_name: str,
+ tiers: List[Dict[str, Any]],
+ tier_height: int,
+ show_tier_labels: bool,
+ collapsed_tiers: List[str]
+) -> str:
+ """Generate HTML for tier rows."""
+ rows = []
+
+ for tier in tiers:
+ tier_name = tier["name"]
+ is_collapsed = tier_name in collapsed_tiers
+ collapsed_class = "collapsed" if is_collapsed else ""
+ indent_class = "tier-dependent" if tier["tier_type"] == "dependent" else ""
+
+ row_html = f'''
+
+'''
+
+ if show_tier_labels:
+ row_html += f'''
+
+ {escape_html_content(tier_name)}
+
+
+
+
+'''
+
+ row_html += f'''
+
+
+
+
+'''
+ rows.append(row_html)
+
+ return "\n".join(rows)
+
+
+def _generate_tier_selector_html(tiers: List[Dict[str, Any]]) -> str:
+ """Generate HTML for tier selector dropdown."""
+ options = []
+ for tier in tiers:
+ indent = " " if tier["tier_type"] == "dependent" else ""
+ tier_type_indicator = " (โ)" if tier["tier_type"] == "dependent" else ""
+ options.append(
+ f''
+ f'{indent}{escape_html_content(tier["name"])}{tier_type_indicator} '
+ )
+ return "\n".join(options)
+
+
+def _generate_media_element(schema_name: str, media_type: str) -> str:
+ """Generate the media player element (audio or video)."""
+ if media_type == "video":
+ return f'''
+
+
+ Your browser does not support the video element.
+
+'''
+ else:
+ # Audio (default)
+ return f'''
+
+
+ Your browser does not support the audio element.
+
+'''
+
+
+def _generate_playback_controls(
+ schema_name: str,
+ playback_rate_control: bool,
+ zoom_enabled: bool
+) -> str:
+ """Generate playback and zoom controls HTML."""
+ controls = []
+
+ if playback_rate_control:
+ controls.append(f'''
+
+ Speed:
+
+ 0.25x
+ 0.5x
+ 0.75x
+ 1x
+ 1.25x
+ 1.5x
+ 2x
+
+
+''')
+
+ if zoom_enabled:
+ controls.append(f'''
+
+
+
+
+
+
+
+
+
+
+
+''')
+
+ # Time display
+ controls.append(f'''
+
+ 00:00.000
+ /
+ 00:00.000
+
+''')
+
+ return "\n".join(controls)
diff --git a/potato/server_utils/schemas/trajectory_edit.py b/potato/server_utils/schemas/trajectory_edit.py
new file mode 100644
index 0000000000000000000000000000000000000000..105779ae16c7467b3f9ea16709bac757cbf9d8ed
--- /dev/null
+++ b/potato/server_utils/schemas/trajectory_edit.py
@@ -0,0 +1,531 @@
+"""
+Trajectory Correction / Editing Layout
+
+Annotators rewrite the steps of an agent trace (reasoning, tool calls,
+observations) and optionally the final answer, producing a *corrected*
+trajectory alongside the original. The corrected/original pair is exported as
+SFT targets and DPO preference pairs (see
+``potato/export/trajectory_correction_exporter.py``).
+
+This is the editing counterpart to ``trajectory_eval`` (which *scores* steps).
+It reuses ``trajectory_eval``'s data-loading + per-step-card structure and
+``text_edit``'s live word/char diff. Each step shows the original text
+read-only plus an editable textarea pre-filled with the original; a per-step
+"edited" flag is set automatically when the text diverges.
+
+Research / motivation: Labelbox Agent Trajectory Editor; Datadog "edited
+outputs"; SFT/DPO post-training from human-corrected trajectories.
+"""
+
+import json
+import logging
+from typing import Any, Dict, List, Tuple
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+)
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_EDITABLE_FIELDS = ["action"]
+
+
+def generate_trajectory_edit_layout(
+ annotation_scheme: Dict[str, Any],
+) -> Tuple[str, List[Tuple[str, str]]]:
+ """Generate HTML for a trajectory correction/editing interface.
+
+ Args:
+ annotation_scheme: Configuration dict. Required: ``name``,
+ ``description``. Optional: ``steps_key``, ``step_text_key``,
+ ``editable_fields``, ``show_diff``, ``show_edit_distance``,
+ ``allow_reset``, ``require_reason_on_edit``, ``edit_final_answer``,
+ ``final_answer_key``.
+
+ Returns:
+ ``(html, keybindings)`` tuple.
+ """
+ return safe_generate_layout(annotation_scheme, _generate_internal)
+
+
+def _generate_internal(
+ annotation_scheme: Dict[str, Any],
+) -> Tuple[str, List[Tuple[str, str]]]:
+ schema_name = annotation_scheme["name"]
+ description = annotation_scheme["description"]
+
+ steps_key = annotation_scheme.get("steps_key", "steps")
+ step_text_key = annotation_scheme.get("step_text_key", "action")
+ editable_fields = annotation_scheme.get("editable_fields") or [step_text_key]
+ if not isinstance(editable_fields, list):
+ editable_fields = [step_text_key]
+ show_diff = annotation_scheme.get("show_diff", True)
+ show_edit_distance = annotation_scheme.get("show_edit_distance", True)
+ allow_reset = annotation_scheme.get("allow_reset", True)
+ require_reason_on_edit = annotation_scheme.get("require_reason_on_edit", False)
+ edit_final_answer = annotation_scheme.get("edit_final_answer", False)
+ final_answer_key = annotation_scheme.get("final_answer_key", "final_answer")
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+ identifiers = generate_element_identifier(schema_name, schema_name, "hidden")
+
+ config_json = json.dumps({
+ "steps_key": steps_key,
+ "step_text_key": step_text_key,
+ "editable_fields": editable_fields,
+ "show_diff": show_diff,
+ "show_edit_distance": show_edit_distance,
+ "allow_reset": allow_reset,
+ "require_reason_on_edit": require_reason_on_edit,
+ "edit_final_answer": edit_final_answer,
+ "final_answer_key": final_answer_key,
+ })
+
+ esc_schema = escape_html_content(schema_name)
+
+ html = f"""
+
+
+ {escape_html_content(description)}
+
+
+ Steps edited:
+ 0
+ Total edit distance:
+ 0
+
+
+
+
+
+
+
+
+
+
+ """
+
+ logger.info(f"Generated trajectory edit layout for {schema_name}")
+ return html, []
diff --git a/potato/server_utils/schemas/trajectory_eval.py b/potato/server_utils/schemas/trajectory_eval.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb509f6e479db6774c79c5d865e296aeb30c7c60
--- /dev/null
+++ b/potato/server_utils/schemas/trajectory_eval.py
@@ -0,0 +1,509 @@
+"""
+Trajectory Evaluation Layout
+
+Per-step error marking for agent traces. For each step the annotator marks:
+correctness (correct / incorrect / partially_correct), error type from a
+configurable taxonomy, severity, and a free-text rationale.
+
+Research: TRAIL (Trace Reasoning and Agentic Issue Localization),
+ AgentRewardBench, Anthropic "Demystifying Evals for AI Agents".
+"""
+
+import json
+import logging
+from typing import Dict, Any, Tuple, List
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes,
+)
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_CORRECTNESS = ["correct", "incorrect", "partially_correct"]
+DEFAULT_SEVERITIES = [
+ {"name": "minor", "weight": -1},
+ {"name": "major", "weight": -5},
+ {"name": "critical", "weight": -10},
+]
+
+
+def generate_trajectory_eval_layout(
+ annotation_scheme: Dict[str, Any],
+) -> Tuple[str, List[Tuple[str, str]]]:
+ """Generate HTML for a trajectory evaluation interface.
+
+ Args:
+ annotation_scheme: Configuration dict. Required keys: ``name``,
+ ``description``. Optional: ``steps_key``, ``step_text_key``,
+ ``correctness_options``, ``error_types``, ``severities``,
+ ``show_score``.
+
+ Returns:
+ ``(html, keybindings)`` tuple.
+ """
+ return safe_generate_layout(annotation_scheme, _generate_internal)
+
+
+def _generate_internal(
+ annotation_scheme: Dict[str, Any],
+) -> Tuple[str, List[Tuple[str, str]]]:
+ schema_name = annotation_scheme["name"]
+ description = annotation_scheme["description"]
+
+ steps_key = annotation_scheme.get("steps_key", "steps")
+ step_text_key = annotation_scheme.get("step_text_key", "action")
+ correctness_options = annotation_scheme.get("correctness_options", DEFAULT_CORRECTNESS)
+ error_types = annotation_scheme.get("error_types", [])
+ severities = annotation_scheme.get("severities", DEFAULT_SEVERITIES)
+ show_score = annotation_scheme.get("show_score", True)
+ max_score = annotation_scheme.get("max_score", 100)
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+ identifiers = generate_element_identifier(schema_name, schema_name, "hidden")
+
+ # Serialize config for JS IIFE
+ config_json = json.dumps({
+ "steps_key": steps_key,
+ "step_text_key": step_text_key,
+ "correctness_options": correctness_options,
+ "error_types": error_types,
+ "severities": severities,
+ "show_score": show_score,
+ "max_score": max_score,
+ })
+
+ # Build error type elements
+ type_options_html = ' -- select error type -- '
+ for et in error_types:
+ subtypes = et.get("subtypes", [])
+ if subtypes:
+ type_options_html += f''
+ for st in subtypes:
+ val = f'{escape_html_content(et["name"])}::{escape_html_content(st)}'
+ type_options_html += f'{escape_html_content(st)} '
+ type_options_html += " "
+ else:
+ type_options_html += (
+ f''
+ f'{escape_html_content(et["name"])} '
+ )
+
+ # Build severity radio buttons
+ severity_radios_html = ""
+ for sev in severities:
+ severity_radios_html += f"""
+
+
+ {escape_html_content(sev['name'])} ({sev['weight']:+d})
+ """
+
+ # Build correctness buttons
+ correctness_btns = ""
+ for opt in correctness_options:
+ label = opt.replace("_", " ").title()
+ css_cls = f"traj-correctness-{opt}"
+ correctness_btns += (
+ f''
+ f'{escape_html_content(label)} '
+ )
+
+ score_html = ""
+ if show_score:
+ score_html = f"""
+
+ Score: {max_score} / {max_score}
+
"""
+
+ esc_schema = escape_html_content(schema_name)
+
+ html = f"""
+
+
+ {escape_html_content(description)}
+
+ {score_html}
+
+
+
+
+
+
+
+
+
+ """
+
+ logger.info(f"Generated trajectory eval layout for {schema_name}")
+ return html, []
diff --git a/potato/server_utils/schemas/tree_annotation.py b/potato/server_utils/schemas/tree_annotation.py
new file mode 100644
index 0000000000000000000000000000000000000000..5932b38e64e3e12c2c1f1fff1910b6251b4e6182
--- /dev/null
+++ b/potato/server_utils/schemas/tree_annotation.py
@@ -0,0 +1,141 @@
+"""
+Tree Annotation Layout
+
+Generates annotation interface for conversation tree structures.
+Supports per-node annotation, path selection, and branch comparison.
+
+Users can:
+- Annotate individual nodes (e.g., rate each response)
+- Select preferred paths through the tree
+- Compare branches at decision points
+"""
+
+import json
+import logging
+from .identifier_utils import (
+ safe_generate_layout,
+ escape_html_content,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _generate_tree_annotation_layout_internal(annotation_scheme, horizontal=False):
+ """
+ Internal function to generate tree annotation layout.
+
+ Args:
+ annotation_scheme: Configuration dictionary containing:
+ - name: Schema identifier
+ - description: Description shown to user
+ - node_scheme: Annotation scheme config for per-node annotation
+ (e.g., {annotation_type: "likert", size: 5, ...})
+ - path_selection:
+ - enabled: Whether path selection is enabled
+ - description: Instruction text for path selection
+ - branch_comparison:
+ - enabled: Whether branch comparison is enabled
+
+ Returns:
+ tuple: (HTML string, key bindings list)
+ """
+ scheme_name = annotation_scheme["name"]
+ description = annotation_scheme.get("description", "Annotate the conversation tree")
+ node_scheme = annotation_scheme.get("node_scheme", {})
+ path_selection = annotation_scheme.get("path_selection", {})
+ branch_comparison = annotation_scheme.get("branch_comparison", {})
+
+ path_enabled = path_selection.get("enabled", False)
+ path_desc = path_selection.get("description", "Select the best response path")
+ branch_enabled = branch_comparison.get("enabled", False)
+
+ config_data = json.dumps({
+ "schemaName": scheme_name,
+ "nodeScheme": node_scheme,
+ "pathSelection": {
+ "enabled": path_enabled,
+ "description": path_desc,
+ },
+ "branchComparison": {
+ "enabled": branch_enabled,
+ },
+ })
+
+ # Path selection section
+ path_section = ""
+ if path_enabled:
+ path_section = f"""
+
+
Path Selection
+
{escape_html_content(path_desc)}
+
+ No path selected. Click on nodes in the tree to build a path.
+
+
Clear Path
+
+ """
+
+ # Node annotation mode description
+ node_ann_desc = ""
+ if node_scheme:
+ node_type = node_scheme.get("annotation_type", "")
+ node_ann_desc = f"""
+
+
Click a node in the tree above to annotate it.
+ Node annotation type: {escape_html_content(node_type)}
+
+ """
+
+ schematic = f"""
+
+
+
+
+ {node_ann_desc}
+
+
+
+
+ {path_section}
+
+
+
+
+
+ """
+
+ key_bindings = []
+ return schematic, key_bindings
+
+
+def generate_tree_annotation_layout(annotation_scheme, horizontal=False):
+ """
+ Generate tree annotation layout HTML.
+
+ Args:
+ annotation_scheme (dict): The annotation scheme configuration
+ horizontal (bool): Whether to display horizontally
+
+ Returns:
+ tuple: (HTML string, key bindings list)
+ """
+ return safe_generate_layout(
+ annotation_scheme, _generate_tree_annotation_layout_internal, horizontal
+ )
diff --git a/potato/server_utils/schemas/triage.py b/potato/server_utils/schemas/triage.py
new file mode 100644
index 0000000000000000000000000000000000000000..47a099ce577000e45b8d0c64ad62bf32ccb1b568
--- /dev/null
+++ b/potato/server_utils/schemas/triage.py
@@ -0,0 +1,179 @@
+"""
+Triage Layout
+
+Generates a Prodigy-style binary accept/reject/skip interface for rapid data curation.
+Features include:
+- Three large buttons: Accept (green), Reject (red), Skip (gray)
+- Keyboard shortcuts: a (accept), r (reject), s (skip)
+- Auto-advance to next item option
+- Progress indicator display
+- Minimal UI optimized for rapid decisions
+"""
+
+import logging
+from typing import Dict, Any, Tuple, List
+
+from .identifier_utils import (
+ safe_generate_layout,
+ escape_html_content,
+ generate_layout_attributes,
+)
+
+logger = logging.getLogger(__name__)
+
+# Default labels for triage actions
+DEFAULT_ACCEPT_LABEL = "Keep"
+DEFAULT_REJECT_LABEL = "Discard"
+DEFAULT_SKIP_LABEL = "Unsure"
+
+# Default keyboard shortcuts (1/2/3 are adjacent and easy for rapid annotation)
+DEFAULT_KEYBINDINGS = {
+ "accept": "1",
+ "reject": "2",
+ "skip": "3",
+}
+
+
+def generate_triage_layout(annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Generate HTML for a triage (accept/reject/skip) interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - accept_label: Custom label for accept button (default: "Accept")
+ - reject_label: Custom label for reject button (default: "Reject")
+ - skip_label: Custom label for skip button (default: "Skip")
+ - auto_advance: Whether to auto-advance after selection (default: true)
+ - show_progress: Whether to show progress indicator (default: true)
+ - accept_key: Custom keyboard shortcut for accept (default: "a")
+ - reject_key: Custom keyboard shortcut for reject (default: "r")
+ - skip_key: Custom keyboard shortcut for skip (default: "s")
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the triage interface
+ key_bindings: List of (key, description) tuples for keyboard shortcuts
+ """
+ return safe_generate_layout(annotation_scheme, _generate_triage_layout_internal)
+
+
+def _generate_triage_layout_internal(annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Internal function to generate triage layout after validation.
+ """
+ logger.debug(f"Generating triage layout for schema: {annotation_scheme['name']}")
+
+ schema_name = annotation_scheme["name"]
+ description = annotation_scheme.get("description", "")
+
+ # Get custom labels or use defaults
+ accept_label = annotation_scheme.get("accept_label", DEFAULT_ACCEPT_LABEL)
+ reject_label = annotation_scheme.get("reject_label", DEFAULT_REJECT_LABEL)
+ skip_label = annotation_scheme.get("skip_label", DEFAULT_SKIP_LABEL)
+
+ # Get keyboard shortcuts
+ accept_key = annotation_scheme.get("accept_key", DEFAULT_KEYBINDINGS["accept"])
+ reject_key = annotation_scheme.get("reject_key", DEFAULT_KEYBINDINGS["reject"])
+ skip_key = annotation_scheme.get("skip_key", DEFAULT_KEYBINDINGS["skip"])
+
+ # Get options
+ auto_advance = annotation_scheme.get("auto_advance", True)
+ show_progress = annotation_scheme.get("show_progress", True)
+
+ # Get layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Build keybindings list
+ key_bindings = [
+ (accept_key, f"{schema_name}: {accept_label}"),
+ (reject_key, f"{schema_name}: {reject_label}"),
+ (skip_key, f"{schema_name}: {skip_label}"),
+ ]
+
+ # Escape for HTML attributes
+ safe_schema_name = escape_html_content(schema_name)
+ safe_description = escape_html_content(description)
+ safe_accept_label = escape_html_content(accept_label)
+ safe_reject_label = escape_html_content(reject_label)
+ safe_skip_label = escape_html_content(skip_label)
+
+ # Build the HTML
+ html = f"""
+
+
+ {safe_description}
+
+
+
+
+
+
+
+
+ ✓
+ {safe_accept_label}
+ [{escape_html_content(accept_key.upper())}]
+
+
+
+ ✗
+ {safe_reject_label}
+ [{escape_html_content(reject_key.upper())}]
+
+
+
+ →
+ {safe_skip_label}
+ [{escape_html_content(skip_key.upper())}]
+
+
+ """
+
+ # Add progress indicator if enabled
+ if show_progress:
+ html += """
+
+
+ """.format(schema_name=safe_schema_name)
+
+ html += """
+
+
+
+ """
+
+ logger.info(f"Successfully generated triage layout for {schema_name}")
+ return html, key_bindings
diff --git a/potato/server_utils/schemas/vas.py b/potato/server_utils/schemas/vas.py
new file mode 100644
index 0000000000000000000000000000000000000000..a04ee2e1f237749cb9751502aad426d5b471dfa6
--- /dev/null
+++ b/potato/server_utils/schemas/vas.py
@@ -0,0 +1,131 @@
+"""
+Visual Analog Scale (VAS) Layout
+
+A continuous line scale with no tick marks or discrete bins. Annotators click/drag
+to a position, returning a precise float value. Psychophysically superior to Likert
+scales for fine-grained judgments.
+
+Research: Stevens (1957) "On the Psychophysical Law"; standard in clinical pain assessment.
+"""
+
+import logging
+
+from .identifier_utils import (
+ safe_generate_layout,
+ generate_element_identifier,
+ generate_validation_attribute,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+
+logger = logging.getLogger(__name__)
+
+# Defaults
+DEFAULT_MIN = 0
+DEFAULT_MAX = 100
+DEFAULT_SHOW_VALUE = False
+
+
+def generate_vas_layout(annotation_scheme):
+ """
+ Generate HTML for a Visual Analog Scale interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - left_label: Label for the left endpoint
+ - right_label: Label for the right endpoint
+ - min_value: Minimum value (default 0)
+ - max_value: Maximum value (default 100)
+ - show_value: Whether to show numeric value after selection (default False)
+ - precision: Decimal places to round to (default 1)
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ """
+ return safe_generate_layout(annotation_scheme, _generate_vas_layout_internal)
+
+
+def _generate_vas_layout_internal(annotation_scheme):
+ schema_name = annotation_scheme['name']
+ description = annotation_scheme['description']
+ left_label = annotation_scheme.get('left_label', '')
+ right_label = annotation_scheme.get('right_label', '')
+ min_value = annotation_scheme.get('min_value', DEFAULT_MIN)
+ max_value = annotation_scheme.get('max_value', DEFAULT_MAX)
+ show_value = annotation_scheme.get('show_value', DEFAULT_SHOW_VALUE)
+ precision = annotation_scheme.get('precision', 1)
+
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+ validation = generate_validation_attribute(annotation_scheme)
+ identifiers = generate_element_identifier(schema_name, schema_name, "range")
+
+ # The key difference from slider: a fine precision-based step (vs the
+ # slider's integer ticks), no tick marks, no value display by default,
+ # minimal styling. Step is derived from precision so the scale is
+ # effectively continuous (precision=1 โ step=0.1, precision=2 โ 0.01).
+ step_value = 10 ** -precision if precision > 0 else 1
+ initial_value = round((min_value + max_value) / 2, precision)
+
+ value_display = ""
+ if show_value:
+ value_display = f"""
+
+ โ
+
+ """
+
+ html = f"""
+
+
+ {escape_html_content(description)}
+
+
+ {escape_html_content(left_label)}
+ {escape_html_content(right_label)}
+
+
+ {value_display}
+
+
+
+ """
+
+ # Inline JS to round value and update display
+ if show_value:
+ html += f"""
+
+ """
+
+ logger.info(f"Generated VAS layout for {schema_name}")
+ return html, []
diff --git a/potato/server_utils/schemas/video.py b/potato/server_utils/schemas/video.py
new file mode 100644
index 0000000000000000000000000000000000000000..d7d9b18e28ed55e45753c67412b5a1f3157f4237
--- /dev/null
+++ b/potato/server_utils/schemas/video.py
@@ -0,0 +1,225 @@
+"""
+Video Layout
+
+Generates a form interface for displaying video content. Features include:
+- Custom video player controls
+- Autoplay options
+- Loop control
+- Muting options
+- Custom CSS styling
+- Multiple video source support
+- Fallback content support
+"""
+
+import logging
+import os.path
+
+from potato.ai.ai_help_wrapper import get_ai_wrapper, get_dynamic_ai_help
+from .identifier_utils import (
+ safe_generate_layout,
+ escape_html_content,
+ generate_layout_attributes
+)
+
+logger = logging.getLogger(__name__)
+
+def generate_video_layout(annotation_scheme):
+ """
+ Generate HTML for a video player interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - video_path: Path to video file
+ - custom_css (dict): Optional CSS styling
+ - width: Video width (default: "320")
+ - height: Video height (default: "240")
+ - autoplay (bool): Whether to start playing automatically
+ - loop (bool): Whether to loop video playback
+ - muted (bool): Whether to mute audio by default
+ - controls (bool): Whether to show player controls
+ - fallback_text (str): Optional text to show if video fails
+ - additional_sources (list): Optional additional video formats
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the video interface
+ key_bindings: Empty list (no keyboard shortcuts)
+
+ Raises:
+ ValueError: If video_path is missing or invalid
+ """
+ return safe_generate_layout(annotation_scheme, _generate_video_layout_internal)
+
+def _is_url(path: str) -> bool:
+ """
+ Check if a path is a URL.
+
+ Args:
+ path: The path to check
+
+ Returns:
+ bool: True if path is a URL, False otherwise
+ """
+ return path.startswith(('http://', 'https://', '//', 'data:'))
+
+
+def _generate_video_layout_internal(annotation_scheme):
+ """
+ Internal function to generate video layout after validation.
+ """
+ logger.debug(f"Generating video layout for schema: {annotation_scheme['name']}")
+
+ # Validate video path
+ if "video_path" not in annotation_scheme:
+ error_msg = f"Missing video_path in schema: {annotation_scheme['name']}"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ video_path = annotation_scheme["video_path"]
+
+ # Only check file existence for local paths, not URLs
+ if not _is_url(video_path) and not os.path.exists(video_path):
+ # Log a warning but don't fail - the file might be served from a different location
+ logger.warning(f"Video file not found locally: {video_path}. "
+ f"Assuming it will be served from a web-accessible location.")
+
+ # Get layout attributes for grid positioning
+ layout_attrs = generate_layout_attributes(annotation_scheme)
+
+ # Initialize form wrapper
+ schematic = f"""
+
+ {get_ai_wrapper()}
+
+ {escape_html_content(annotation_scheme['description'])}
+ """
+
+ # Generate video element with attributes
+ video_attrs = _generate_video_attributes(annotation_scheme)
+ css = _generate_css_style(annotation_scheme)
+
+ schematic += f"""
+
+ {_generate_video_sources(annotation_scheme)}
+ {_generate_fallback_content(annotation_scheme)}
+
+ """
+
+ schematic += " "
+
+ logger.info(f"Successfully generated video layout for {annotation_scheme['name']}")
+ return schematic, []
+
+def _generate_video_attributes(annotation_scheme):
+ """
+ Generate HTML attributes for video element.
+
+ Args:
+ annotation_scheme (dict): Video configuration settings
+
+ Returns:
+ str: Space-separated video attributes
+ """
+ attrs = []
+
+ # Handle playback controls
+ if annotation_scheme.get("controls", True):
+ attrs.append("controls")
+ logger.debug("Enabled video controls")
+
+ if annotation_scheme.get("autoplay"):
+ attrs.append("autoplay")
+ logger.debug("Enabled autoplay")
+
+ if annotation_scheme.get("loop"):
+ attrs.append("loop")
+ logger.debug("Enabled video loop")
+
+ if annotation_scheme.get("muted"):
+ attrs.append("muted")
+ logger.debug("Enabled muted playback")
+
+ return " ".join(attrs)
+
+def _generate_css_style(annotation_scheme):
+ """
+ Generate CSS style string from configuration.
+
+ Args:
+ annotation_scheme (dict): Configuration containing custom_css settings
+
+ Returns:
+ str: Formatted CSS style string
+ """
+ css = annotation_scheme.get("custom_css", {})
+ styles = []
+
+ # Default dimensions if not specified
+ width = css.get("width", "320")
+ height = css.get("height", "240")
+
+ styles.append(f"width: {width}px")
+ styles.append(f"height: {height}px")
+
+ return "; ".join(styles)
+
+def _generate_video_sources(annotation_scheme):
+ """
+ Generate source elements for video formats.
+
+ Args:
+ annotation_scheme (dict): Configuration containing video sources
+
+ Returns:
+ str: HTML for video source elements
+ """
+ sources = []
+
+ # Add main video source
+ mime_type = _get_mime_type(annotation_scheme["video_path"])
+ sources.append(
+ f''
+ )
+ logger.debug(f"Added primary video source: {annotation_scheme['video_path']}")
+
+ # Add additional sources if specified
+ for source in annotation_scheme.get("additional_sources", []):
+ mime_type = _get_mime_type(source)
+ sources.append(f'')
+ logger.debug(f"Added additional video source: {source}")
+
+ return "\n".join(sources)
+
+def _generate_fallback_content(annotation_scheme):
+ """
+ Generate fallback content for browsers that don't support video.
+
+ Args:
+ annotation_scheme (dict): Configuration containing fallback settings
+
+ Returns:
+ str: HTML for fallback content
+ """
+ fallback = annotation_scheme.get("fallback_text", "Your browser does not support the video tag.")
+ logger.debug("Added fallback content for video element")
+ return escape_html_content(fallback)
+
+def _get_mime_type(file_path):
+ """
+ Determine MIME type from video file extension.
+
+ Args:
+ file_path (str): Path to video file
+
+ Returns:
+ str: MIME type string
+ """
+ ext = os.path.splitext(file_path)[1].lower()
+ mime_types = {
+ '.mp4': 'video/mp4',
+ '.webm': 'video/webm',
+ '.ogg': 'video/ogg'
+ }
+ return mime_types.get(ext, 'video/mp4')
\ No newline at end of file
diff --git a/potato/server_utils/schemas/video_annotation.py b/potato/server_utils/schemas/video_annotation.py
new file mode 100644
index 0000000000000000000000000000000000000000..70e0a47fc0f7c00b507dd9300869f5f2bbcf7bac
--- /dev/null
+++ b/potato/server_utils/schemas/video_annotation.py
@@ -0,0 +1,923 @@
+"""
+Video Annotation Layout
+
+Generates a form interface for video annotation with:
+- Temporal segment marking (like audio annotation)
+- Frame-level classification
+- Keyframe annotation
+- Object tracking across frames (basic support)
+
+Uses Peaks.js for timeline visualization with synchronized video playback.
+
+Features:
+- Waveform/timeline visualization from video's audio track
+- Video preview panel with frame counter
+- Segment creation and labeling
+- Frame-by-frame stepping (,/. keys)
+- Multiple playback speeds (0.1x to 2.0x)
+- Keyframe marking (K key)
+"""
+
+import logging
+import json
+from typing import List, Dict, Tuple, Any
+from .identifier_utils import (
+ safe_generate_layout,
+ escape_html_content
+)
+
+logger = logging.getLogger(__name__)
+
+# Default colors for segment/frame labels
+DEFAULT_COLORS = [
+ "#4ECDC4", # Teal
+ "#FF6B6B", # Red
+ "#45B7D1", # Blue
+ "#96CEB4", # Green
+ "#FFEAA7", # Yellow
+ "#DDA0DD", # Plum
+ "#95A5A6", # Gray
+ "#F39C12", # Orange
+ "#9B59B6", # Purple
+ "#3498DB", # Light Blue
+]
+
+# Valid annotation modes
+VALID_MODES = ["segment", "frame", "keyframe", "tracking", "combined"]
+
+
+def generate_video_annotation_layout(annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Generate HTML for a video annotation interface.
+
+ Args:
+ annotation_scheme (dict): Configuration including:
+ - name: Schema identifier
+ - description: Display description
+ - mode: "segment" | "frame" | "keyframe" | "tracking" | "combined"
+ - labels: List of labels (for segment/frame/keyframe modes)
+ - segment_schemes: List of annotation schemes per segment (optional)
+ - min_segments: Minimum required segments (default: 0)
+ - max_segments: Maximum allowed segments (default: null/unlimited)
+ - timeline_height: Height of timeline in pixels (default: 70)
+ - overview_height: Height of overview bar (default: 40)
+ - zoom_enabled: Whether to enable zoom (default: True)
+ - playback_rate_control: Show playback speed controls (default: True)
+ - frame_stepping: Enable frame-by-frame navigation (default: True)
+ - show_timecode: Show timecode display (default: True)
+ - video_fps: Frames per second for frame calculations (default: 30)
+
+ Returns:
+ tuple: (html_string, key_bindings)
+ html_string: Complete HTML for the video annotation interface
+ key_bindings: List of keyboard shortcuts
+
+ Raises:
+ ValueError: If required fields are missing or invalid
+ """
+ return safe_generate_layout(annotation_scheme, _generate_video_annotation_layout_internal)
+
+
+def _generate_video_annotation_layout_internal(annotation_scheme: Dict[str, Any]) -> Tuple[str, List[Tuple[str, str]]]:
+ """
+ Internal function to generate video annotation layout after validation.
+ """
+ schema_name = annotation_scheme.get('name', 'video_annotation')
+ logger.debug(f"Generating video annotation layout for schema: {schema_name}")
+
+ # Get mode (default to "segment")
+ mode = annotation_scheme.get('mode', 'segment')
+ if mode not in VALID_MODES:
+ error_msg = f"Invalid mode '{mode}' in schema: {schema_name}. Must be one of: {VALID_MODES}"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ # Validate labels for segment/frame/keyframe/tracking/combined modes
+ labels = []
+ if mode in ['segment', 'frame', 'keyframe', 'tracking', 'combined']:
+ if 'labels' not in annotation_scheme:
+ error_msg = f"Missing labels in schema: {schema_name} (required for mode '{mode}')"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+ labels = _process_labels(annotation_scheme['labels'])
+
+ # Validate segment_schemes for combined mode
+ segment_schemes = []
+ if mode in ['combined'] and 'segment_schemes' in annotation_scheme:
+ segment_schemes = annotation_scheme['segment_schemes']
+ if not isinstance(segment_schemes, list):
+ error_msg = f"segment_schemes must be a list in schema: {schema_name}"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ # Get configuration options
+ min_segments = annotation_scheme.get('min_segments', 0)
+ max_segments = annotation_scheme.get('max_segments', None)
+ timeline_height = annotation_scheme.get('timeline_height', 70)
+ overview_height = annotation_scheme.get('overview_height', 40)
+ zoom_enabled = annotation_scheme.get('zoom_enabled', True)
+ playback_rate_control = annotation_scheme.get('playback_rate_control', True)
+ frame_stepping = annotation_scheme.get('frame_stepping', True)
+ show_timecode = annotation_scheme.get('show_timecode', True)
+ video_fps = annotation_scheme.get('video_fps', 30)
+
+ # AI support configuration
+ ai_support = annotation_scheme.get("ai_support", {})
+ ai_enabled = ai_support.get("enabled", False)
+
+ # source_field: Links this annotation schema to a display field from instance_display
+ source_field = annotation_scheme.get("source_field", "")
+
+ # Build config object for JavaScript
+ js_config = {
+ "schemaName": schema_name,
+ "mode": mode,
+ "labels": labels,
+ "segmentSchemes": segment_schemes,
+ "minSegments": min_segments,
+ "maxSegments": max_segments,
+ "timelineHeight": timeline_height,
+ "overviewHeight": overview_height,
+ "zoomEnabled": zoom_enabled,
+ "playbackRateControl": playback_rate_control,
+ "frameStepping": frame_stepping,
+ "showTimecode": show_timecode,
+ "videoFps": video_fps,
+ "aiSupport": ai_enabled,
+ "aiFeatures": ai_support.get("features", {}) if ai_enabled else {},
+ "sourceField": source_field,
+ }
+
+ # Generate HTML
+ html = _generate_html(annotation_scheme, js_config, schema_name, labels, mode, ai_enabled, ai_support)
+
+ # Generate keybindings
+ keybindings = _generate_keybindings(labels, mode, frame_stepping)
+
+ logger.info(f"Successfully generated video annotation layout for {schema_name}")
+ return html, keybindings
+
+
+def _process_labels(labels_config: List) -> List[Dict[str, Any]]:
+ """
+ Process label configuration and assign colors.
+
+ Args:
+ labels_config: List of label configs (strings or dicts)
+
+ Returns:
+ List of processed label dicts with name, color, and optional key_value
+ """
+ processed = []
+ for i, label in enumerate(labels_config):
+ if isinstance(label, str):
+ processed.append({
+ "name": label,
+ "color": DEFAULT_COLORS[i % len(DEFAULT_COLORS)],
+ })
+ elif isinstance(label, dict):
+ processed.append({
+ "name": label.get("name", f"label_{i}"),
+ "color": label.get("color", DEFAULT_COLORS[i % len(DEFAULT_COLORS)]),
+ "key_value": label.get("key_value"),
+ })
+ else:
+ processed.append({
+ "name": str(label),
+ "color": DEFAULT_COLORS[i % len(DEFAULT_COLORS)],
+ })
+ return processed
+
+
+def _generate_html(
+ annotation_scheme: Dict[str, Any],
+ js_config: Dict[str, Any],
+ schema_name: str,
+ labels: List[Dict[str, Any]],
+ mode: str,
+ ai_enabled: bool = False,
+ ai_support: Dict[str, Any] = None
+) -> str:
+ """
+ Generate the HTML for the video annotation interface.
+ """
+ escaped_name = escape_html_content(schema_name)
+ description = escape_html_content(annotation_scheme.get('description', ''))
+ config_json = json.dumps(js_config)
+ timeline_height = js_config.get('timelineHeight', 70)
+ overview_height = js_config.get('overviewHeight', 40)
+
+ # Generate label buttons
+ label_selector = ""
+ if labels:
+ label_selector = _generate_label_selector(labels)
+
+ # Generate AI toolbar if enabled
+ ai_toolbar_html = ""
+ ai_init_script = ""
+ if ai_enabled:
+ ai_features = ai_support.get("features", {}) if ai_support else {}
+ ai_toolbar_html = _generate_video_ai_toolbar(ai_features, mode)
+ ai_init_script = _generate_video_ai_init_script(escaped_name)
+
+ # Generate playback rate control
+ playback_rate_html = ""
+ if js_config.get('playbackRateControl'):
+ playback_rate_html = '''
+
+ Speed:
+
+ 0.1x
+ 0.25x
+ 0.5x
+ 0.75x
+ 1x
+ 1.25x
+ 1.5x
+ 2x
+
+
+ '''
+
+ # Generate frame stepping controls
+ frame_stepping_html = ""
+ if js_config.get('frameStepping'):
+ frame_stepping_html = '''
+
+ |<
+ |>
+
+ '''
+
+ # Generate mode-specific controls
+ mode_controls_html = ""
+ if mode in ['keyframe', 'combined']:
+ mode_controls_html += '''
+
+ Mark Keyframe
+
+ '''
+ if mode in ['frame', 'combined']:
+ mode_controls_html += '''
+
+ Classify Frame
+
+ '''
+ if mode in ['tracking', 'combined']:
+ mode_controls_html += '''
+
+ + Track
+ Delete Track
+
+ Linear
+ Cubic (smooth)
+ Constant (hold)
+
+
+ '''
+
+ # Generate timecode display
+ timecode_html = ""
+ if js_config.get('showTimecode'):
+ timecode_html = '''
+
+ Frame: 0
+ 00:00:00.000
+
+ '''
+
+ # source_field attribute for linking to display fields
+ source_field = annotation_scheme.get("source_field", "")
+ source_field_attr = f' data-source-field="{escape_html_content(source_field)}"' if source_field else ""
+
+ html = f'''
+
+
+ {description}
+
+
+
+
+
+
+
+ {timecode_html}
+
+
+
+
+
+
+
+ How to annotate video segments (click to expand)
+
+
To create a segment:
+
+ Select a label (colored buttons below)
+ Play/scrub video to the start point, then click [ or press the [ key
+ Move to the end point, then click ] or press the ] key
+ Click + Segment or press Enter to create the segment
+
+
Timeline controls: + zoom in, - zoom out, Fit show entire video
+
Playback: Space play/pause, ◀▶ frame step, change speed with dropdown
+
+
+
+
+
+
+ {ai_toolbar_html}
+
+
+
+
+
+
+
+
+
+
+
+
+
Object Tracks
+
+ Click "+ Track" to create a track, then draw bounding boxes on the video. Boxes are interpolated between keyframes.
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ '''
+
+ return html
+
+
+def _generate_video_ai_toolbar(ai_features: Dict[str, Any], mode: str) -> str:
+ """
+ Generate HTML for the video AI assistance toolbar.
+ """
+ scene_detection_enabled = ai_features.get("scene_detection", True)
+ frame_classification_enabled = ai_features.get("frame_classification", False)
+ keyframe_detection_enabled = ai_features.get("keyframe_detection", False)
+ tracking_enabled = ai_features.get("tracking", False)
+ pre_annotate_enabled = ai_features.get("pre_annotate", True)
+ hint_enabled = ai_features.get("hint", True)
+
+ buttons = []
+
+ if scene_detection_enabled and mode in ['segment', 'combined']:
+ buttons.append(
+ ''
+ '๐ฌ Scenes '
+ )
+
+ if pre_annotate_enabled:
+ buttons.append(
+ ''
+ 'โก Auto '
+ )
+
+ if frame_classification_enabled and mode in ['frame', 'combined']:
+ buttons.append(
+ ''
+ '๐ท๏ธ Classify '
+ )
+
+ if keyframe_detection_enabled and mode in ['keyframe', 'combined']:
+ buttons.append(
+ ''
+ '๐ Keyframes '
+ )
+
+ if tracking_enabled and mode == 'tracking':
+ buttons.append(
+ ''
+ '๐๏ธ Track '
+ )
+
+ if hint_enabled:
+ buttons.append(
+ ''
+ '๐ก Hint '
+ )
+
+ if not buttons:
+ return ""
+
+ return f'''
+
+
+
+
+ '''
+
+
+def _generate_video_ai_init_script(escaped_name: str) -> str:
+ """
+ Generate JavaScript initialization code for video AI assistant.
+ """
+ return f'''
+ // Initialize AI assistant if enabled and VisualAIAssistantManager is available
+ if (config.aiSupport && typeof VisualAIAssistantManager !== 'undefined') {{
+ var annotationId = Array.from(document.querySelectorAll('.annotation-form')).indexOf(
+ document.getElementById('{escaped_name}')
+ );
+ container.aiAssistant = new VisualAIAssistantManager({{
+ annotationType: 'video_annotation',
+ annotationId: annotationId >= 0 ? annotationId : 0,
+ annotationManager: manager
+ }});
+ }}
+ '''
+
+
+def _generate_label_selector(labels: List[Dict[str, Any]]) -> str:
+ """
+ Generate HTML for label selection buttons.
+ """
+ import html
+ buttons = []
+ for label in labels:
+ # Escape for both content and attributes
+ name_escaped = escape_html_content(label["name"])
+ name_attr = html.escape(label["name"], quote=True)
+ color = label["color"]
+ key_hint = f' ({label["key_value"]})' if label.get("key_value") else ""
+ key_hint_escaped = html.escape(key_hint, quote=True)
+ buttons.append(
+ f''
+ f' '
+ f'{name_escaped} '
+ )
+
+ return f'''
+
+ Label:
+ {"".join(buttons)}
+
+ '''
+
+
+def _generate_keybindings(labels: List[Dict[str, Any]], mode: str, frame_stepping: bool) -> List[Tuple[str, str]]:
+ """
+ Generate keybinding list for the schema.
+ """
+ keybindings = []
+
+ # Playback shortcuts
+ keybindings.append(("Space", "Play/Pause"))
+ keybindings.append(("Left/Right", "Seek 5 seconds"))
+
+ # Frame stepping
+ if frame_stepping:
+ keybindings.append((",", "Previous frame"))
+ keybindings.append((".", "Next frame"))
+
+ # Label shortcuts
+ for label in labels:
+ if label.get("key_value"):
+ keybindings.append((label["key_value"], f"Select: {label['name']}"))
+
+ # Segment shortcuts
+ if mode in ['segment', 'combined']:
+ keybindings.append(("[", "Set segment start"))
+ keybindings.append(("]", "Set segment end"))
+ keybindings.append(("Enter", "Create segment"))
+
+ # Keyframe shortcut
+ if mode in ['keyframe', 'combined']:
+ keybindings.append(("K", "Mark keyframe"))
+
+ # Frame classification
+ if mode in ['frame', 'combined']:
+ keybindings.append(("C", "Classify current frame"))
+
+ # Delete shortcut
+ keybindings.append(("Del", "Delete selected"))
+
+ # Zoom shortcuts
+ keybindings.append(("+/-", "Zoom in/out"))
+ keybindings.append(("0", "Fit to view"))
+
+ return keybindings
diff --git a/potato/server_utils/triage.py b/potato/server_utils/triage.py
new file mode 100644
index 0000000000000000000000000000000000000000..41690990346e0cfcdbd4df8e8863d499e7ac56d9
--- /dev/null
+++ b/potato/server_utils/triage.py
@@ -0,0 +1,270 @@
+"""
+Signal-based triage scoring for the annotation queue.
+
+Captures a per-item quality signal โ an agent error, a production thumbs-down, a
+low automated score, or any custom field โ and turns it into a numeric
+*triage priority* so the assignment queue can surface the worst / most-suspect
+traces first instead of annotating in arrival (FIFO) order.
+
+The same scorer runs over statically loaded data and over traces ingested at
+runtime (webhook / Langfuse), because both funnel through
+``ItemStateManager.add_item``. The priority is stored on the item's metadata
+(``triage_priority`` / ``triage_reason`` / ``triage_rule``); the ``priority``
+assignment strategy reads it; the inline badge and the admin queue page surface it.
+
+Config (all optional):
+
+ triage:
+ enabled: true
+ order: desc # high priority first (default); 'asc' = low first
+ default_priority: 0 # items matching no rule
+ show_badge: true # show a "why prioritized" banner during annotation
+ signal_field: null # read a numeric priority straight from this field
+ invert_signal: false # if true, a LOWER field value => HIGHER priority
+ rules: # evaluated in order; highest matching priority wins
+ - name: "Agent errored"
+ when: {field: "status", equals: "error"}
+ priority: 100
+ badge: "Agent errored"
+ - name: "Negative feedback"
+ when: {field: "feedback", in: ["thumbs_down", "negative"]}
+ priority: 80
+ - name: "Low score"
+ when: {field: "score", lt: 0.5}
+ priority: 60
+
+When ``enabled`` with no ``rules`` and no ``signal_field``, a turnkey set of
+built-in defaults is used (error status, negative feedback, low score) so
+ingested traces are triaged out of the box.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+
+# Turnkey defaults applied when triage is enabled but no rules/signal_field are
+# configured. These match the signals ingested traces most commonly carry.
+DEFAULT_RULES = [
+ {"name": "Agent errored", "badge": "Agent errored", "priority": 100,
+ "when": {"field": "status", "in": ["error", "failed", "failure"]}},
+ {"name": "Negative feedback", "badge": "Negative feedback", "priority": 80,
+ "when": {"field": "feedback", "in": ["thumbs_down", "negative", "down", "๐"]}},
+ {"name": "Low score", "badge": "Low score", "priority": 60,
+ "when": {"field": "score", "lt": 0.5}},
+]
+
+
+@dataclass
+class TriageScore:
+ """The triage outcome for one item."""
+ priority: float
+ reason: str | None = None # human-readable badge text (None when unflagged)
+ rule: str | None = None # the rule name that matched (None for default/field)
+
+ def to_metadata(self) -> dict:
+ return {
+ "triage_priority": self.priority,
+ "triage_reason": self.reason,
+ "triage_rule": self.rule,
+ }
+
+
+def _lookup(data: dict, field: str):
+ """Resolve a possibly dotted field path against an item dict.
+
+ Supports nested dicts (``metadata.score``). Returns None if any segment is
+ missing or a non-dict is traversed.
+ """
+ cur = data
+ for part in str(field).split("."):
+ if not isinstance(cur, dict) or part not in cur:
+ return None
+ cur = cur[part]
+ return cur
+
+
+def _as_number(value):
+ """Coerce a value to float, or None if it isn't numeric."""
+ if isinstance(value, bool):
+ return None
+ if isinstance(value, (int, float)):
+ return float(value)
+ if isinstance(value, str):
+ try:
+ return float(value.strip())
+ except (ValueError, AttributeError):
+ return None
+ return None
+
+
+def _matches(condition: dict, data: dict) -> bool:
+ """Evaluate a single rule's ``when`` condition against item data.
+
+ Supported operators: equals, in, lt, lte, gt, gte, exists, contains.
+ String comparisons for equals/in are case-insensitive. Numeric comparisons
+ coerce both sides. ``contains`` tests membership in a list/string field.
+ """
+ field = condition.get("field")
+ if field is None:
+ return False
+ value = _lookup(data, field)
+
+ if "exists" in condition:
+ present = value is not None
+ return present == bool(condition["exists"])
+
+ # Absent fields never match value-based operators.
+ if value is None:
+ return False
+
+ if "equals" in condition:
+ target = condition["equals"]
+ if isinstance(value, str) and isinstance(target, str):
+ return value.strip().lower() == target.strip().lower()
+ return value == target
+
+ if "in" in condition:
+ options = condition["in"] or []
+ norm = [o.lower() if isinstance(o, str) else o for o in options]
+ v = value.lower() if isinstance(value, str) else value
+ return v in norm
+
+ if "contains" in condition:
+ target = condition["contains"]
+ if isinstance(value, (list, tuple, set)):
+ tnorm = target.lower() if isinstance(target, str) else target
+ return any(
+ (item.lower() if isinstance(item, str) else item) == tnorm
+ for item in value
+ )
+ if isinstance(value, str) and isinstance(target, str):
+ return target.lower() in value.lower()
+ return False
+
+ for op, py in (("lt", "<"), ("lte", "<="), ("gt", ">"), ("gte", ">=")):
+ if op in condition:
+ lhs, rhs = _as_number(value), _as_number(condition[op])
+ if lhs is None or rhs is None:
+ return False
+ if op == "lt":
+ return lhs < rhs
+ if op == "lte":
+ return lhs <= rhs
+ if op == "gt":
+ return lhs > rhs
+ return lhs >= rhs
+
+ return False
+
+
+class TriageScorer:
+ """Scores items into a triage priority from the ``triage`` config block."""
+
+ def __init__(self, triage_config: dict):
+ cfg = triage_config or {}
+ self.enabled = bool(cfg.get("enabled", False))
+ self.order = str(cfg.get("order", "desc")).lower()
+ self.default_priority = float(cfg.get("default_priority", 0) or 0)
+ self.show_badge = bool(cfg.get("show_badge", True))
+ self.signal_field = cfg.get("signal_field")
+ self.invert_signal = bool(cfg.get("invert_signal", False))
+
+ rules = cfg.get("rules")
+ if not rules and not self.signal_field:
+ rules = DEFAULT_RULES
+ self.rules = rules or []
+
+ def score(self, item_data: dict) -> TriageScore:
+ """Return the TriageScore for one item (highest matching rule wins)."""
+ if not self.enabled:
+ return TriageScore(priority=self.default_priority)
+
+ best: TriageScore | None = None
+ for rule in self.rules:
+ cond = rule.get("when") or {}
+ try:
+ if _matches(cond, item_data):
+ pr = float(rule.get("priority", 0) or 0)
+ if best is None or pr > best.priority:
+ badge = rule.get("badge") or rule.get("name")
+ best = TriageScore(priority=pr, reason=badge, rule=rule.get("name"))
+ except Exception as e: # a malformed rule must never break loading
+ logger.warning(f"Triage rule {rule.get('name')!r} failed: {e}")
+
+ if best is not None:
+ return best
+
+ # No rule matched: optionally read a direct numeric signal.
+ if self.signal_field is not None:
+ raw = _as_number(_lookup(item_data, self.signal_field))
+ if raw is not None:
+ pr = -raw if self.invert_signal else raw
+ return TriageScore(priority=pr, reason=None, rule=None)
+
+ return TriageScore(priority=self.default_priority)
+
+
+def build_scorer(config: dict) -> TriageScorer | None:
+ """Build a TriageScorer from a server config, or None when triage is off."""
+ triage_cfg = (config or {}).get("triage") or {}
+ if not triage_cfg.get("enabled"):
+ return None
+ return TriageScorer(triage_cfg)
+
+
+def compute_triage_queue(config: dict) -> dict:
+ """Build the admin triage-queue report from the live ItemStateManager.
+
+ Returns the remaining (incomplete) items ranked by triage priority, with the
+ reason/rule that flagged them, current annotation count, and whether they are
+ already assigned. Used by the ``/admin/triage-queue`` page.
+ """
+ from potato.item_state_management import get_item_state_manager
+
+ scorer = build_scorer(config)
+ order = (scorer.order if scorer else "desc")
+ reverse = order != "asc"
+
+ ism = get_item_state_manager()
+ rows = []
+ # Preserve the configured/global ordering as the deterministic tie-break.
+ ordering = {iid: i for i, iid in enumerate(ism.instance_id_ordering)}
+ for iid in ism.instance_id_ordering:
+ item = ism.get_item(iid)
+ if item is None:
+ continue
+ # Skip items that have reached their annotation cap.
+ try:
+ if ism._item_is_saturated(iid):
+ continue
+ except Exception:
+ pass
+ priority = item.get_metadata("triage_priority")
+ if priority is None:
+ priority = scorer.default_priority if scorer else 0
+ n_ann = len(ism.instance_annotators.get(iid, set()))
+ rows.append({
+ "id": iid,
+ "priority": priority,
+ "reason": item.get_metadata("triage_reason"),
+ "rule": item.get_metadata("triage_rule"),
+ "annotations": n_ann,
+ "assigned": n_ann > 0,
+ "_order": ordering.get(iid, 0),
+ })
+
+ rows.sort(key=lambda r: (r["priority"], -r["_order"]), reverse=reverse)
+ for r in rows:
+ r.pop("_order", None)
+
+ return {
+ "enabled": bool(scorer),
+ "order": order,
+ "n_items": len(rows),
+ "n_flagged": sum(1 for r in rows if r["reason"]),
+ "items": rows,
+ }
diff --git a/potato/server_utils/waveform_service.py b/potato/server_utils/waveform_service.py
new file mode 100644
index 0000000000000000000000000000000000000000..db84c5ffb93237dad5500f6a0104b330550b2997
--- /dev/null
+++ b/potato/server_utils/waveform_service.py
@@ -0,0 +1,546 @@
+"""
+Waveform Service
+
+Handles generation and caching of audio waveform data for the audio annotation feature.
+Uses BBC's audiowaveform tool to generate pre-computed waveform data files.
+
+Features:
+- LRU cache for waveform files
+- Background look-ahead pre-computation for upcoming instances
+- Support for both local files and URLs
+- Graceful fallback if audiowaveform not installed
+"""
+
+import os
+import logging
+import hashlib
+import subprocess
+import shutil
+import tempfile
+import threading
+import time
+from typing import Optional, List, Dict
+from collections import OrderedDict
+from urllib.parse import urlparse
+from pathlib import Path
+
+try:
+ import requests
+ REQUESTS_AVAILABLE = True
+except ImportError:
+ REQUESTS_AVAILABLE = False
+
+logger = logging.getLogger(__name__)
+
+
+class WaveformService:
+ """
+ Service for generating and caching audio waveform data.
+
+ Uses BBC's audiowaveform tool to generate pre-computed waveform data
+ that can be efficiently rendered by Peaks.js on the frontend.
+ """
+
+ # Default configuration
+ DEFAULT_LOOK_AHEAD = 5
+ DEFAULT_CACHE_MAX_SIZE = 100
+ DEFAULT_CLIENT_FALLBACK_MAX_DURATION = 1800 # 30 minutes in seconds
+
+ # Waveform generation settings
+ WAVEFORM_ZOOM_LEVEL = 256 # Samples per pixel
+ WAVEFORM_BITS = 8 # 8-bit resolution
+
+ def __init__(
+ self,
+ cache_dir: str,
+ look_ahead: int = DEFAULT_LOOK_AHEAD,
+ cache_max_size: int = DEFAULT_CACHE_MAX_SIZE,
+ client_fallback_max_duration: int = DEFAULT_CLIENT_FALLBACK_MAX_DURATION
+ ):
+ """
+ Initialize the WaveformService.
+
+ Args:
+ cache_dir: Directory to store generated waveform files
+ look_ahead: Number of instances to pre-compute ahead
+ cache_max_size: Maximum number of cached waveform files
+ client_fallback_max_duration: Max duration (seconds) for client-side fallback
+ """
+ self.cache_dir = cache_dir
+ self.look_ahead = look_ahead
+ self.cache_max_size = cache_max_size
+ self.client_fallback_max_duration = client_fallback_max_duration
+
+ # LRU cache tracking
+ self._cache_order: OrderedDict = OrderedDict()
+ self._cache_lock = threading.Lock()
+
+ # Background pre-computation
+ self._precompute_thread: Optional[threading.Thread] = None
+ self._precompute_queue: List[str] = []
+ self._precompute_lock = threading.Lock()
+ self._stop_precompute = threading.Event()
+
+ # Check if audiowaveform is installed
+ self._audiowaveform_available = self._check_audiowaveform_installed()
+
+ # Ensure cache directory exists
+ self._ensure_cache_dir()
+
+ logger.info(f"WaveformService initialized: cache_dir={cache_dir}, "
+ f"look_ahead={look_ahead}, audiowaveform_available={self._audiowaveform_available}")
+
+ def _ensure_cache_dir(self) -> None:
+ """Create the cache directory if it doesn't exist."""
+ if not os.path.exists(self.cache_dir):
+ os.makedirs(self.cache_dir, exist_ok=True)
+ logger.info(f"Created waveform cache directory: {self.cache_dir}")
+
+ def _check_audiowaveform_installed(self) -> bool:
+ """
+ Check if the audiowaveform tool is installed and available.
+
+ Returns:
+ True if audiowaveform is available, False otherwise
+ """
+ try:
+ result = subprocess.run(
+ ['audiowaveform', '--version'],
+ capture_output=True,
+ text=True,
+ timeout=5
+ )
+ if result.returncode == 0:
+ version = result.stdout.strip() or result.stderr.strip()
+ logger.info(f"audiowaveform found: {version}")
+ return True
+ except (subprocess.SubprocessError, FileNotFoundError, OSError) as e:
+ logger.warning(f"audiowaveform not available: {e}")
+
+ return False
+
+ @property
+ def is_available(self) -> bool:
+ """Check if waveform generation is available."""
+ return self._audiowaveform_available
+
+ def _get_cache_key(self, audio_path: str) -> str:
+ """
+ Generate a unique cache key for an audio file.
+
+ Args:
+ audio_path: Path or URL to the audio file
+
+ Returns:
+ MD5 hash of the path as cache key
+ """
+ return hashlib.md5(audio_path.encode('utf-8')).hexdigest()
+
+ def _get_waveform_cache_path(self, audio_path: str) -> str:
+ """
+ Get the cache file path for a waveform.
+
+ Args:
+ audio_path: Path or URL to the audio file
+
+ Returns:
+ Full path to the waveform cache file
+ """
+ cache_key = self._get_cache_key(audio_path)
+ return os.path.join(self.cache_dir, f"{cache_key}.dat")
+
+ def _is_url(self, path: str) -> bool:
+ """
+ Check if a path is a URL.
+
+ Args:
+ path: The path to check
+
+ Returns:
+ True if path is a URL, False otherwise
+ """
+ return path.startswith(('http://', 'https://', '//'))
+
+ def _download_audio(self, url: str) -> Optional[str]:
+ """
+ Download an audio file from URL to a temporary file.
+
+ Args:
+ url: URL of the audio file
+
+ Returns:
+ Path to temporary file, or None if download failed
+ """
+ if not REQUESTS_AVAILABLE:
+ logger.error("requests library not available for downloading audio")
+ return None
+
+ try:
+ # Determine file extension from URL
+ parsed = urlparse(url)
+ path = parsed.path
+ ext = os.path.splitext(path)[1] or '.mp3'
+
+ # Create temporary file
+ temp_fd, temp_path = tempfile.mkstemp(suffix=ext)
+ os.close(temp_fd)
+
+ logger.debug(f"Downloading audio from {url} to {temp_path}")
+
+ response = requests.get(url, stream=True, timeout=60)
+ response.raise_for_status()
+
+ with open(temp_path, 'wb') as f:
+ for chunk in response.iter_content(chunk_size=8192):
+ f.write(chunk)
+
+ logger.debug(f"Downloaded audio: {os.path.getsize(temp_path)} bytes")
+ return temp_path
+
+ except Exception as e:
+ logger.error(f"Failed to download audio from {url}: {e}")
+ return None
+
+ def _generate_waveform(self, audio_path: str, output_path: str) -> bool:
+ """
+ Generate waveform data using audiowaveform tool.
+
+ Args:
+ audio_path: Path to the audio file (local)
+ output_path: Path to write the waveform data file
+
+ Returns:
+ True if generation succeeded, False otherwise
+ """
+ if not self._audiowaveform_available:
+ logger.warning("audiowaveform not available, cannot generate waveform")
+ return False
+
+ try:
+ # Build command
+ cmd = [
+ 'audiowaveform',
+ '-i', audio_path,
+ '-o', output_path,
+ '-z', str(self.WAVEFORM_ZOOM_LEVEL),
+ '-b', str(self.WAVEFORM_BITS),
+ ]
+
+ logger.debug(f"Running: {' '.join(cmd)}")
+
+ result = subprocess.run(
+ cmd,
+ capture_output=True,
+ text=True,
+ timeout=300 # 5 minute timeout for long files
+ )
+
+ if result.returncode == 0:
+ logger.info(f"Generated waveform: {output_path}")
+ return True
+ else:
+ logger.error(f"audiowaveform failed: {result.stderr}")
+ return False
+
+ except subprocess.TimeoutExpired:
+ logger.error(f"audiowaveform timed out for {audio_path}")
+ return False
+ except Exception as e:
+ logger.error(f"Error generating waveform for {audio_path}: {e}")
+ return False
+
+ def _update_cache_order(self, cache_path: str) -> None:
+ """
+ Update LRU cache order and evict if necessary.
+
+ Args:
+ cache_path: Path to the cache file being accessed
+ """
+ with self._cache_lock:
+ # Move to end (most recently used)
+ if cache_path in self._cache_order:
+ self._cache_order.move_to_end(cache_path)
+ else:
+ self._cache_order[cache_path] = True
+
+ # Evict oldest if over limit
+ while len(self._cache_order) > self.cache_max_size:
+ oldest_path, _ = self._cache_order.popitem(last=False)
+ if os.path.exists(oldest_path):
+ try:
+ os.remove(oldest_path)
+ logger.debug(f"Evicted from cache: {oldest_path}")
+ except OSError as e:
+ logger.warning(f"Failed to remove cache file {oldest_path}: {e}")
+
+ def get_waveform_path(self, audio_path: str, generate: bool = True) -> Optional[str]:
+ """
+ Get the waveform data file path for an audio file.
+
+ If the waveform doesn't exist and generate=True, it will be generated.
+
+ Args:
+ audio_path: Path or URL to the audio file
+ generate: Whether to generate if not cached
+
+ Returns:
+ Path to waveform data file, or None if not available
+ """
+ cache_path = self._get_waveform_cache_path(audio_path)
+
+ # Check if already cached
+ if os.path.exists(cache_path):
+ self._update_cache_order(cache_path)
+ logger.debug(f"Waveform cache hit: {cache_path}")
+ return cache_path
+
+ if not generate:
+ return None
+
+ # Generate waveform
+ temp_audio = None
+ try:
+ # Handle URL vs local path
+ if self._is_url(audio_path):
+ temp_audio = self._download_audio(audio_path)
+ if not temp_audio:
+ return None
+ local_path = temp_audio
+ else:
+ local_path = audio_path
+ if not os.path.exists(local_path):
+ logger.warning(f"Audio file not found: {local_path}")
+ return None
+
+ # Generate waveform
+ if self._generate_waveform(local_path, cache_path):
+ self._update_cache_order(cache_path)
+ return cache_path
+ else:
+ return None
+
+ finally:
+ # Clean up temporary file
+ if temp_audio and os.path.exists(temp_audio):
+ try:
+ os.remove(temp_audio)
+ except OSError:
+ pass
+
+ def get_waveform_url(self, audio_path: str, base_url: str = '/api/waveform/') -> Optional[str]:
+ """
+ Get the URL to fetch waveform data for an audio file.
+
+ Args:
+ audio_path: Path or URL to the audio file
+ base_url: Base URL for the waveform API endpoint
+
+ Returns:
+ URL to fetch waveform data
+ """
+ cache_key = self._get_cache_key(audio_path)
+ return f"{base_url}{cache_key}"
+
+ def precompute_batch(self, audio_paths: List[str]) -> None:
+ """
+ Pre-compute waveforms for a batch of audio files.
+
+ This is called synchronously and blocks until all are complete.
+ Use start_background_precompute for non-blocking operation.
+
+ Args:
+ audio_paths: List of audio file paths or URLs
+ """
+ for audio_path in audio_paths:
+ if audio_path:
+ self.get_waveform_path(audio_path, generate=True)
+
+ def queue_precompute(self, audio_paths: List[str]) -> None:
+ """
+ Add audio files to the background pre-computation queue.
+
+ Args:
+ audio_paths: List of audio file paths or URLs to pre-compute
+ """
+ with self._precompute_lock:
+ # Only add paths not already in queue or cached
+ for path in audio_paths:
+ if path and path not in self._precompute_queue:
+ cache_path = self._get_waveform_cache_path(path)
+ if not os.path.exists(cache_path):
+ self._precompute_queue.append(path)
+
+ # Start background thread if not running
+ if self._precompute_thread is None or not self._precompute_thread.is_alive():
+ self._start_background_precompute()
+
+ def _start_background_precompute(self) -> None:
+ """Start the background pre-computation thread."""
+ self._stop_precompute.clear()
+ self._precompute_thread = threading.Thread(
+ target=self._background_precompute_worker,
+ daemon=True
+ )
+ self._precompute_thread.start()
+ logger.debug("Started background waveform pre-computation thread")
+
+ def _background_precompute_worker(self) -> None:
+ """Background worker for pre-computing waveforms."""
+ while not self._stop_precompute.is_set():
+ # Get next item from queue
+ audio_path = None
+ with self._precompute_lock:
+ if self._precompute_queue:
+ audio_path = self._precompute_queue.pop(0)
+
+ if audio_path:
+ logger.debug(f"Background pre-computing waveform for: {audio_path}")
+ self.get_waveform_path(audio_path, generate=True)
+ else:
+ # No more items, exit thread
+ break
+
+ # Small delay between items to avoid overloading
+ time.sleep(0.1)
+
+ logger.debug("Background waveform pre-computation thread finished")
+
+ def stop_background_precompute(self) -> None:
+ """Stop the background pre-computation thread."""
+ self._stop_precompute.set()
+ if self._precompute_thread and self._precompute_thread.is_alive():
+ self._precompute_thread.join(timeout=5)
+
+ def get_audio_duration(self, audio_path: str) -> Optional[float]:
+ """
+ Get the duration of an audio file in seconds.
+
+ Uses ffprobe if available, otherwise returns None.
+
+ Args:
+ audio_path: Path to the audio file
+
+ Returns:
+ Duration in seconds, or None if cannot determine
+ """
+ try:
+ result = subprocess.run(
+ [
+ 'ffprobe',
+ '-v', 'error',
+ '-show_entries', 'format=duration',
+ '-of', 'default=noprint_wrappers=1:nokey=1',
+ audio_path
+ ],
+ capture_output=True,
+ text=True,
+ timeout=10
+ )
+ if result.returncode == 0:
+ return float(result.stdout.strip())
+ except (subprocess.SubprocessError, ValueError, FileNotFoundError):
+ pass
+
+ return None
+
+ def should_use_client_fallback(self, audio_path: str) -> bool:
+ """
+ Determine if client-side waveform generation should be used.
+
+ Client-side is preferred for short files when server-side is not available.
+
+ Args:
+ audio_path: Path to the audio file
+
+ Returns:
+ True if client-side fallback should be used
+ """
+ if self._audiowaveform_available:
+ return False
+
+ duration = self.get_audio_duration(audio_path)
+ if duration is not None and duration <= self.client_fallback_max_duration:
+ return True
+
+ return False
+
+ def clear_cache(self) -> int:
+ """
+ Clear all cached waveform files.
+
+ Returns:
+ Number of files removed
+ """
+ count = 0
+ with self._cache_lock:
+ for cache_path in list(self._cache_order.keys()):
+ if os.path.exists(cache_path):
+ try:
+ os.remove(cache_path)
+ count += 1
+ except OSError as e:
+ logger.warning(f"Failed to remove {cache_path}: {e}")
+ self._cache_order.clear()
+
+ logger.info(f"Cleared {count} cached waveform files")
+ return count
+
+ def get_cache_stats(self) -> Dict:
+ """
+ Get statistics about the waveform cache.
+
+ Returns:
+ Dictionary with cache statistics
+ """
+ with self._cache_lock:
+ cached_files = len(self._cache_order)
+ total_size = 0
+ for cache_path in self._cache_order.keys():
+ if os.path.exists(cache_path):
+ total_size += os.path.getsize(cache_path)
+
+ return {
+ 'cached_files': cached_files,
+ 'max_files': self.cache_max_size,
+ 'total_size_bytes': total_size,
+ 'total_size_mb': round(total_size / (1024 * 1024), 2),
+ 'cache_dir': self.cache_dir,
+ 'audiowaveform_available': self._audiowaveform_available,
+ }
+
+
+# Global instance (initialized when needed)
+_waveform_service: Optional[WaveformService] = None
+
+
+def get_waveform_service() -> Optional[WaveformService]:
+ """Get the global WaveformService instance."""
+ return _waveform_service
+
+
+def init_waveform_service(
+ cache_dir: str,
+ look_ahead: int = WaveformService.DEFAULT_LOOK_AHEAD,
+ cache_max_size: int = WaveformService.DEFAULT_CACHE_MAX_SIZE,
+ client_fallback_max_duration: int = WaveformService.DEFAULT_CLIENT_FALLBACK_MAX_DURATION
+) -> WaveformService:
+ """
+ Initialize the global WaveformService instance.
+
+ Args:
+ cache_dir: Directory to store generated waveform files
+ look_ahead: Number of instances to pre-compute ahead
+ cache_max_size: Maximum number of cached waveform files
+ client_fallback_max_duration: Max duration for client-side fallback
+
+ Returns:
+ The initialized WaveformService instance
+ """
+ global _waveform_service
+ _waveform_service = WaveformService(
+ cache_dir=cache_dir,
+ look_ahead=look_ahead,
+ cache_max_size=cache_max_size,
+ client_fallback_max_duration=client_fallback_max_duration
+ )
+ return _waveform_service
diff --git a/potato/setup_multilingual_config.py b/potato/setup_multilingual_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..693f1f167eefbdce6c65fdc9a59928a96b7a2c36
--- /dev/null
+++ b/potato/setup_multilingual_config.py
@@ -0,0 +1,167 @@
+"""
+Multilingual Configuration Setup Module
+
+This module provides functionality for setting up multilingual annotation tasks.
+It processes a base configuration and creates language-specific configurations
+for each supported language, including:
+
+- Language-specific survey flow files
+- Translated configuration files
+- Localized output directories
+- Multilingual guideline integration
+
+The module supports dynamic text replacement using key-based translation mappings
+and creates the necessary directory structure for multilingual annotation projects.
+"""
+
+from argparse import ArgumentParser
+import yaml
+import os
+import pandas as pd
+import json
+from collections import defaultdict
+
+
+def arguments():
+ """
+ Creates and returns the argument parser for multilingual configuration setup.
+
+ Returns:
+ ArgumentParser: Configured argument parser with multilingual_config_file argument
+ """
+ parser = ArgumentParser()
+ parser.set_defaults(show_path=False, show_similarity=False)
+
+ parser.add_argument("multilingual_config_file")
+
+ return parser.parse_args()
+
+
+def main():
+ """
+ Main function for setting up multilingual annotation configurations.
+
+ This function processes a multilingual configuration file and creates
+ language-specific configurations for each supported language. It handles:
+
+ 1. Directory structure creation
+ 2. Translation mapping from guideline files
+ 3. Configuration file generation for each language
+ 4. Survey flow file localization
+ 5. Output directory setup
+
+ Side Effects:
+ - Creates directories for surveyflow, annotation_output, data_files, configs, htmls
+ - Generates language-specific configuration files
+ - Creates localized survey flow files
+ - Sets up language-specific output directories
+ """
+ args = arguments()
+
+ # Load multilingual annotation configuration
+ with open(args.multilingual_config_file, "rt") as f:
+ multilingual_config = yaml.safe_load(f)
+
+ # Create basic folder structure for the multilingual project
+ for folder in ["surveyflow", "annotation_output", "data_files", "configs", "htmls"]:
+ cur_path = multilingual_config["base_dir"] + folder
+ if not os.path.exists(cur_path):
+ os.makedirs(cur_path)
+ print("Created directory: %s" % (cur_path))
+
+ # Build translation mapping dictionary
+ # This maps translation keys to their language-specific text values
+ key2text = defaultdict(dict)
+
+ # Load multilingual annotation guidelines if specified and available
+ if "multilingual_guideline_file" in multilingual_config and os.path.exists(
+ multilingual_config["multilingual_guideline_file"]
+ ):
+ # Load multilingual annotation guideline from CSV
+ multilingual_guideline_df = pd.read_csv(multilingual_config["multilingual_guideline_file"])
+ for i, row in multilingual_guideline_df.iterrows():
+ # Only process rows with proper key format (enclosed in brackets)
+ if type(row["key"]) != str or row["key"][0] != "[" or row["key"][-1] != "]":
+ continue
+ # Create translation mapping for each supported language
+ for lang in multilingual_config["languages"]:
+ key2text[row["key"]][lang] = (
+ row[lang]
+ if type(row[lang]) == str
+ else row[multilingual_config["base_language"]]
+ )
+
+ # Generate configuration for each supported language
+ for lang in multilingual_config["languages"]:
+ # Load and process the base configuration file
+ with open(multilingual_config["base_config_file"], "rt") as f:
+ page = f.read()
+ # Replace translation keys with language-specific text
+ for key in key2text:
+ page = page.replace(key, key2text[key][lang])
+
+ # Update surveyflow output path for this language
+ surveyflow_output_path = multilingual_config["surveyflow_output_path"].replace(
+ "[LANGUAGE]", lang
+ )
+ page = page.replace(
+ multilingual_config["surveyflow_path"], surveyflow_output_path + lang + "-"
+ )
+ page = page.replace("[LANGUAGE]", lang)
+
+ # Parse the processed configuration
+ config = yaml.safe_load(page)
+
+ # Set up language-specific output directory
+ config["output_annotation_dir"] = multilingual_config["output_annotation_dir"].replace(
+ "[LANGUAGE]", lang
+ )
+ if not os.path.exists(config["output_annotation_dir"]):
+ os.makedirs(config["output_annotation_dir"])
+
+ """
+ # setup the site_dir path for each language
+ config['path_under_site_dir'] = multilingual_config['path_under_site_dir'].replace("[LANGUAGE]", lang)
+ config["site_dir"] += config['path_under_site_dir']
+ if not os.path.exists(config["site_dir"]):
+ os.makedirs(config["site_dir"])
+ """
+
+ # Update task name and data files for this language
+ config["annotation_task_name"] = multilingual_config["annotation_task_name"].replace(
+ "[LANGUAGE]", lang
+ )
+ config["data_files"] = [
+ it.replace("[LANGUAGE]", lang) for it in multilingual_config["data_files"]
+ ]
+
+ # Save the language-specific configuration file
+ with open(multilingual_config["base_dir"] + "configs/%s.yaml" % lang, "wt") as f:
+ json.dump(config, f, indent=4)
+
+ # Create the directory for surveyflow output path
+ if not os.path.exists(surveyflow_output_path):
+ os.makedirs(surveyflow_output_path)
+
+ # Process and localize surveyflow files
+ surveyflow_files = os.listdir(multilingual_config["surveyflow_path"])
+ for file in surveyflow_files:
+ # Skip directories, only process files
+ if os.path.isdir(multilingual_config["surveyflow_path"] + file):
+ continue
+ # Read the surveyflow file
+ with open(multilingual_config["surveyflow_path"] + file, "r", encoding="utf-8") as f:
+ page = f.read()
+ # Replace translation keys with language-specific text
+ for key in key2text:
+ page = page.replace(key, key2text[key][lang])
+ page = page.replace("[LANGUAGE]", lang)
+ # Write the localized surveyflow file
+ with open(surveyflow_output_path + lang + "-" + file, "wt", encoding="utf-8") as f:
+ f.write(page)
+
+ # for key in ["surveyflow_output_path", ]
+
+
+if __name__ == "__main__":
+ main()
diff --git a/potato/setup_multitask_config.py b/potato/setup_multitask_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..f47601dca012ca33f8ea316a253ed58251a93f33
--- /dev/null
+++ b/potato/setup_multitask_config.py
@@ -0,0 +1,166 @@
+"""
+Multitask Configuration Setup Module
+
+This module provides functionality for setting up multitask annotation projects.
+It processes a base configuration and creates task-specific configurations
+for each supported task, including:
+
+- Task-specific survey flow files
+- Specialized configuration files
+- Task-specific output directories
+- Multitask guideline integration
+
+The module supports dynamic text replacement using key-based task mappings
+and creates the necessary directory structure for multitask annotation projects.
+"""
+
+from argparse import ArgumentParser
+import yaml
+import os
+import pandas as pd
+import json
+from collections import defaultdict
+
+
+def arguments():
+ """
+ Creates and returns the argument parser for multitask configuration setup.
+
+ Returns:
+ ArgumentParser: Configured argument parser with multitask_config_file argument
+ """
+ parser = ArgumentParser()
+ parser.set_defaults(show_path=False, show_similarity=False)
+
+ parser.add_argument("multitask_config_file")
+
+ return parser.parse_args()
+
+
+def main():
+ """
+ Main function for setting up multitask annotation configurations.
+
+ This function processes a multitask configuration file and creates
+ task-specific configurations for each supported task. It handles:
+
+ 1. Directory structure creation
+ 2. Task mapping from guideline files
+ 3. Configuration file generation for each task
+ 4. Survey flow file specialization
+ 5. Output directory setup
+
+ Side Effects:
+ - Creates directories for surveyflow, annotation_output, data_files, configs, htmls
+ - Generates task-specific configuration files
+ - Creates specialized survey flow files
+ - Sets up task-specific output directories
+ """
+ args = arguments()
+
+ # Load multitask annotation configuration
+ with open(args.multitask_config_file, "rt") as f:
+ multitask_config = yaml.safe_load(f)
+
+ # Create basic folder structure for the multitask project
+ for folder in ["surveyflow", "annotation_output", "data_files", "configs", "htmls"]:
+ cur_path = multitask_config["base_dir"] + folder
+ if not os.path.exists(cur_path):
+ os.makedirs(cur_path)
+ print("Created directory: %s" % (cur_path))
+
+ # Build task mapping dictionary
+ # This maps task keys to their task-specific text values
+ key2text = defaultdict(dict)
+
+ # Load multitask annotation guidelines if specified and available
+ if "multitask_guideline_file" in multitask_config and os.path.exists(
+ multitask_config["multitask_guideline_file"]
+ ):
+ # Load multitask annotation guideline from CSV
+ multitask_guideline_df = pd.read_csv(multitask_config["multitask_guideline_file"])
+ for i, row in multitask_guideline_df.iterrows():
+ # Only process rows with proper key format (enclosed in brackets)
+ if type(row["key"]) != str or row["key"][0] != "[" or row["key"][-1] != "]":
+ continue
+ # Create task mapping for each supported task
+ for task in multitask_config["tasks"]:
+ key2text[row["key"]][task] = (
+ row[task] if type(row[task]) == str else row[multitask_config["base_task"]]
+ )
+
+ # Generate configuration for each supported task
+ for task in multitask_config["tasks"]:
+ # Load and process the base configuration file
+ with open(multitask_config["base_config_file"], "rt") as f:
+ page = f.read()
+ # Replace task keys with task-specific text
+ for key in key2text:
+ page = page.replace(key, key2text[key][task])
+
+ # Update surveyflow output path for this task
+ surveyflow_output_path = multitask_config["surveyflow_output_path"].replace(
+ "[TASK]", task
+ )
+ page = page.replace(
+ multitask_config["surveyflow_path"], surveyflow_output_path + task + "-"
+ )
+ page = page.replace("[TASK]", task)
+
+ # Parse the processed configuration
+ config = yaml.safe_load(page)
+
+ # Set up task-specific output directory
+ config["output_annotation_dir"] = multitask_config["output_annotation_dir"].replace(
+ "[TASK]", task
+ )
+ if not os.path.exists(config["output_annotation_dir"]):
+ os.makedirs(config["output_annotation_dir"])
+
+ """
+ # setup the site_dir path for each task
+ config['path_under_site_dir'] = multitask_config['path_under_site_dir'].replace("[TASK]", task)
+ config["site_dir"] += config['path_under_site_dir']
+ if not os.path.exists(config["site_dir"]):
+ os.makedirs(config["site_dir"])
+ """
+
+ # Update task name and data files for this task
+ config["annotation_task_name"] = multitask_config["annotation_task_name"].replace(
+ "[TASK]", task
+ )
+ config["data_files"] = [
+ it.replace("[TASK]", task) for it in multitask_config["data_files"]
+ ]
+ # config["prestudy"] = [it.replace("[TASK]", task) for it in multitask_config['prestudy']]
+
+ # Save the task-specific configuration file
+ with open(multitask_config["base_dir"] + "configs/%s.yaml" % task, "wt") as f:
+ json.dump(config, f, indent=4)
+
+ # Create the directory for surveyflow output path
+ if not os.path.exists(surveyflow_output_path):
+ os.makedirs(surveyflow_output_path)
+
+ # Process and specialize surveyflow files
+ surveyflow_files = os.listdir(multitask_config["surveyflow_path"])
+ for file in surveyflow_files:
+ # Skip directories, only process files
+ if os.path.isdir(multitask_config["surveyflow_path"] + file):
+ continue
+ # Read the surveyflow file
+ with open(multitask_config["surveyflow_path"] + file, "r", encoding="utf-8") as f:
+ page = f.read()
+ # Replace task keys with task-specific text
+ for key in key2text:
+ page = page.replace(key, key2text[key][task])
+ page = page.replace("[TASK]", task)
+ # Write the specialized surveyflow file
+ with open(surveyflow_output_path + task + "-" + file, "wt", encoding="utf-8") as f:
+ f.write(page)
+
+ # for key in ["surveyflow_output_path", ]
+
+
+if __name__ == "__main__":
+ main()
diff --git a/potato/similarity.py b/potato/similarity.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b6dbca1c4f9a3521f5c41b597ce191d345c7165
--- /dev/null
+++ b/potato/similarity.py
@@ -0,0 +1,273 @@
+"""
+Similarity Engine Module
+
+Provides semantic similarity search for adjudication items using
+sentence-transformers embeddings. Uses a guarded import pattern so the
+system degrades gracefully when sentence-transformers is not installed.
+
+Key Components:
+- SimilarityEngine: Manages embeddings, caching, and similarity search
+- Singleton management: init/get/clear pattern matching other managers
+"""
+
+import json
+import logging
+import os
+import threading
+from typing import Any, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+# Guarded import โ same pattern as simpledorff in admin.py
+try:
+ from sentence_transformers import SentenceTransformer
+ import numpy as np
+ _SENTENCE_TRANSFORMERS_AVAILABLE = True
+except ImportError:
+ _SENTENCE_TRANSFORMERS_AVAILABLE = False
+
+# Singleton
+_SIMILARITY_ENGINE = None
+_SIMILARITY_LOCK = threading.Lock()
+
+
+class SimilarityEngine:
+ """
+ Manages sentence-transformer embeddings for finding semantically
+ similar annotation items during adjudication.
+ """
+
+ def __init__(self, config: Dict[str, Any], adj_config):
+ """
+ Initialize the similarity engine.
+
+ Args:
+ config: Full application configuration
+ adj_config: AdjudicationConfig dataclass instance
+ """
+ self.config = config
+ self.adj_config = adj_config
+ self.logger = logging.getLogger(__name__)
+ self._lock = threading.Lock()
+
+ self.enabled = False
+ self.model = None
+ self.embeddings = {} # instance_id -> numpy array
+ self.text_cache = {} # instance_id -> text preview
+
+ if not _SENTENCE_TRANSFORMERS_AVAILABLE:
+ self.logger.warning(
+ "sentence-transformers not installed. "
+ "Similarity search disabled. Install with: "
+ "pip install sentence-transformers"
+ )
+ return
+
+ if not adj_config.similarity_enabled:
+ return
+
+ try:
+ model_name = adj_config.similarity_model
+ self.logger.info(f"Loading similarity model: {model_name}")
+ self.model = SentenceTransformer(model_name)
+ self.enabled = True
+ self._load_cache()
+ self.logger.info(
+ f"Similarity engine ready: model={model_name}, "
+ f"cached_embeddings={len(self.embeddings)}"
+ )
+ except Exception as e:
+ self.logger.error(f"Failed to load similarity model: {e}")
+ self.enabled = False
+
+ def precompute_embeddings(self, item_texts: Dict[str, str]) -> int:
+ """
+ Batch-encode texts and store embeddings.
+
+ Args:
+ item_texts: Mapping of instance_id to text content
+
+ Returns:
+ Number of new embeddings computed
+ """
+ if not self.enabled or not self.model:
+ return 0
+
+ with self._lock:
+ # Filter out items already cached
+ new_items = {
+ iid: text for iid, text in item_texts.items()
+ if iid not in self.embeddings
+ }
+
+ if not new_items:
+ return 0
+
+ ids = list(new_items.keys())
+ texts = list(new_items.values())
+
+ try:
+ vecs = self.model.encode(texts, show_progress_bar=False)
+ for i, iid in enumerate(ids):
+ self.embeddings[iid] = vecs[i]
+ self.text_cache[iid] = texts[i][:200] # preview
+
+ self._save_cache()
+ self.logger.info(f"Computed {len(ids)} new embeddings")
+ return len(ids)
+ except Exception as e:
+ self.logger.error(f"Error computing embeddings: {e}")
+ return 0
+
+ def find_similar(
+ self, instance_id: str, top_k: Optional[int] = None
+ ) -> List[Tuple[str, float]]:
+ """
+ Find the most similar items to a given instance.
+
+ Args:
+ instance_id: The reference instance ID
+ top_k: Number of results (defaults to config value)
+
+ Returns:
+ List of (instance_id, similarity_score) tuples, highest first
+ """
+ if not self.enabled or instance_id not in self.embeddings:
+ return []
+
+ if top_k is None:
+ top_k = self.adj_config.similarity_top_k
+
+ with self._lock:
+ ref_vec = self.embeddings[instance_id]
+ results = []
+
+ for other_id, other_vec in self.embeddings.items():
+ if other_id == instance_id:
+ continue
+ score = self._cosine_similarity(ref_vec, other_vec)
+ results.append((other_id, float(score)))
+
+ results.sort(key=lambda x: x[1], reverse=True)
+ return results[:top_k]
+
+ def update_embedding(self, instance_id: str, text: str) -> bool:
+ """
+ Compute and store embedding for a single item.
+
+ Args:
+ instance_id: The instance ID
+ text: The text content
+
+ Returns:
+ True if successful
+ """
+ if not self.enabled or not self.model:
+ return False
+
+ with self._lock:
+ try:
+ vec = self.model.encode([text], show_progress_bar=False)[0]
+ self.embeddings[instance_id] = vec
+ self.text_cache[instance_id] = text[:200]
+ self._save_cache()
+ return True
+ except Exception as e:
+ self.logger.error(f"Error updating embedding for {instance_id}: {e}")
+ return False
+
+ def get_stats(self) -> Dict[str, Any]:
+ """Get similarity engine statistics."""
+ return {
+ "available": _SENTENCE_TRANSFORMERS_AVAILABLE,
+ "enabled": self.enabled,
+ "model": self.adj_config.similarity_model if self.adj_config else None,
+ "embedding_count": len(self.embeddings),
+ "top_k": self.adj_config.similarity_top_k if self.adj_config else None,
+ }
+
+ def _cosine_similarity(self, a, b) -> float:
+ """Compute cosine similarity between two vectors."""
+ dot = float(np.dot(a, b))
+ norm_a = float(np.linalg.norm(a))
+ norm_b = float(np.linalg.norm(b))
+ if norm_a == 0 or norm_b == 0:
+ return 0.0
+ return dot / (norm_a * norm_b)
+
+ def _get_cache_dir(self) -> str:
+ """Get the cache directory path."""
+ output_dir = self.config.get("output_annotation_dir", "annotation_output")
+ adj_subdir = self.adj_config.output_subdir if self.adj_config else "adjudication"
+ cache_dir = os.path.join(output_dir, adj_subdir, ".similarity_cache")
+ os.makedirs(cache_dir, exist_ok=True)
+ return cache_dir
+
+ def _save_cache(self) -> None:
+ """Save embeddings and text cache to disk."""
+ try:
+ import pickle
+ cache_dir = self._get_cache_dir()
+
+ # Save embeddings as pickle (numpy arrays)
+ emb_path = os.path.join(cache_dir, "embeddings.pkl")
+ with open(emb_path, "wb") as f:
+ pickle.dump(self.embeddings, f)
+
+ # Save text cache as JSON
+ text_path = os.path.join(cache_dir, "text_cache.json")
+ with open(text_path, "w") as f:
+ json.dump(self.text_cache, f)
+
+ except Exception as e:
+ self.logger.error(f"Failed to save similarity cache: {e}")
+
+ def _load_cache(self) -> None:
+ """Load cached embeddings and text previews from disk."""
+ try:
+ import pickle
+ cache_dir = self._get_cache_dir()
+
+ emb_path = os.path.join(cache_dir, "embeddings.pkl")
+ if os.path.exists(emb_path):
+ with open(emb_path, "rb") as f:
+ self.embeddings = pickle.load(f)
+
+ text_path = os.path.join(cache_dir, "text_cache.json")
+ if os.path.exists(text_path):
+ with open(text_path, "r") as f:
+ self.text_cache = json.load(f)
+
+ if self.embeddings:
+ self.logger.info(
+ f"Loaded {len(self.embeddings)} cached embeddings"
+ )
+ except Exception as e:
+ self.logger.warning(f"Failed to load similarity cache: {e}")
+ self.embeddings = {}
+ self.text_cache = {}
+
+
+def init_similarity_engine(
+ config: Dict[str, Any], adj_config
+) -> Optional[SimilarityEngine]:
+ """Initialize the singleton SimilarityEngine."""
+ global _SIMILARITY_ENGINE
+
+ with _SIMILARITY_LOCK:
+ if _SIMILARITY_ENGINE is None:
+ _SIMILARITY_ENGINE = SimilarityEngine(config, adj_config)
+
+ return _SIMILARITY_ENGINE
+
+
+def get_similarity_engine() -> Optional[SimilarityEngine]:
+ """Get the singleton SimilarityEngine instance."""
+ return _SIMILARITY_ENGINE
+
+
+def clear_similarity_engine():
+ """Clear the singleton (for testing)."""
+ global _SIMILARITY_ENGINE
+ with _SIMILARITY_LOCK:
+ _SIMILARITY_ENGINE = None
diff --git a/potato/simulator/__init__.py b/potato/simulator/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6da6d6b569222e73bc3c26af39cb597fe831c853
--- /dev/null
+++ b/potato/simulator/__init__.py
@@ -0,0 +1,88 @@
+"""
+User Simulator for Potato Annotation Platform.
+
+This module provides tools for simulating multiple annotators with varying
+competence levels and behaviors for testing purposes.
+
+Example usage:
+ from potato.simulator import SimulatorManager, SimulatorConfig
+
+ config = SimulatorConfig(
+ user_count=10,
+ strategy="random",
+ competence_distribution={"good": 0.5, "average": 0.3, "poor": 0.2}
+ )
+
+ manager = SimulatorManager(config, "http://localhost:8000")
+ results = manager.run_parallel(max_annotations_per_user=20)
+
+ print(manager.get_summary())
+"""
+
+from .config import (
+ SimulatorConfig,
+ UserConfig,
+ TimingConfig,
+ LLMStrategyConfig,
+ BiasedStrategyConfig,
+ PatternStrategyConfig,
+ AgentStrategyConfig,
+ InteractiveConfig,
+ CompetenceLevel,
+ AnnotationStrategyType,
+)
+from .competence_profiles import (
+ CompetenceProfile,
+ create_competence_profile,
+)
+from .annotation_strategies import (
+ AnnotationStrategy,
+ RandomStrategy,
+ BiasedStrategy,
+ LLMStrategy,
+ PatternStrategy,
+ create_strategy,
+)
+from .agent_strategy import AgentSimulatorStrategy
+from .interactive_runner import InteractiveSessionRunner, InteractiveSessionResult
+from .timing_models import TimingModel
+from .user_simulator import SimulatedUser, UserSimulationResult, AnnotationRecord
+from .simulator_manager import SimulatorManager
+from .reporting import SimulationReporter
+
+__all__ = [
+ # Config
+ "SimulatorConfig",
+ "UserConfig",
+ "TimingConfig",
+ "LLMStrategyConfig",
+ "BiasedStrategyConfig",
+ "PatternStrategyConfig",
+ "AgentStrategyConfig",
+ "InteractiveConfig",
+ "CompetenceLevel",
+ "AnnotationStrategyType",
+ # Competence
+ "CompetenceProfile",
+ "create_competence_profile",
+ # Strategies
+ "AnnotationStrategy",
+ "RandomStrategy",
+ "BiasedStrategy",
+ "LLMStrategy",
+ "PatternStrategy",
+ "AgentSimulatorStrategy",
+ "InteractiveSessionRunner",
+ "InteractiveSessionResult",
+ "create_strategy",
+ # Timing
+ "TimingModel",
+ # User simulation
+ "SimulatedUser",
+ "UserSimulationResult",
+ "AnnotationRecord",
+ # Manager
+ "SimulatorManager",
+ # Reporting
+ "SimulationReporter",
+]
diff --git a/potato/simulator/__main__.py b/potato/simulator/__main__.py
new file mode 100644
index 0000000000000000000000000000000000000000..161874f67d0eaf1422c3664f26391a74a3b13ad3
--- /dev/null
+++ b/potato/simulator/__main__.py
@@ -0,0 +1,11 @@
+"""
+Entry point for running the simulator as a module.
+
+Usage:
+ python -m potato.simulator --server http://localhost:8000 --users 10
+"""
+
+from .cli import main
+
+if __name__ == "__main__":
+ main()
diff --git a/potato/simulator/agent_strategy.py b/potato/simulator/agent_strategy.py
new file mode 100644
index 0000000000000000000000000000000000000000..10ee8fdb25fc57e3753977af326a80158c724416
--- /dev/null
+++ b/potato/simulator/agent_strategy.py
@@ -0,0 +1,894 @@
+"""
+Agent (vision-LLM) annotation strategy.
+
+The :class:`AgentSimulatorStrategy` consumes the *full* structured payload of
+an instance โ text fields, dialogue arrays (agent traces, conversations),
+spreadsheet/table data, image references โ and asks a vision-capable LLM to
+produce a complete annotation set covering every schema for that instance.
+
+It mirrors :class:`LLMStrategy` but differs in two important ways:
+
+1. It reads ``instance["data"]`` (the full raw payload that
+ ``/api/current_instance`` returns under the ``data`` key) instead of the
+ single ``text`` field. This gives the model access to dialogue traces,
+ metadata tables, and image URLs.
+
+2. It batches the per-instance call: a single LLM query produces labels for
+ every schema. Subsequent ``generate_annotation`` calls for the same
+ instance are served from a per-instance cache. This keeps cost roughly
+ 1ร (instances) instead of (instances ร schemas).
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import json
+import logging
+import os
+import random
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+from pydantic import BaseModel, Field
+
+from .annotation_strategies import AnnotationStrategy, RandomStrategy
+from .competence_profiles import CompetenceProfile
+from .config import AgentStrategyConfig
+
+logger = logging.getLogger(__name__)
+
+
+_FIELD_DETECTORS = {
+ "dialogue": (
+ "conversation",
+ "dialogue",
+ "trace",
+ "messages",
+ "turns",
+ "structured_turns", # coding-agent traces with role/content/tool_calls
+ ),
+ "spreadsheet": ("metadata_table", "table", "spreadsheet"),
+ "image": ("image", "image_url", "screenshot", "screenshot_url", "media", "image_path"),
+}
+
+# Cap on how much tool I/O text we render per turn to keep prompts bounded.
+_MAX_TOOL_INPUT_CHARS = 400
+_MAX_TOOL_OUTPUT_CHARS = 800
+
+
+class _AgentLabelResponse(BaseModel):
+ """Pydantic schema returned by the LLM for a single instance.
+
+ The model returns a flat dict keyed by ```` whose value is
+ either a label string (radio/multiselect/likert) or a numeric value
+ (slider/likert as int) or a free-text response (text/textbox). The
+ strategy maps these to the wire format the simulator submits.
+ """
+
+ annotations: Dict[str, Any] = Field(
+ default_factory=dict,
+ description="Mapping of schema_name -> chosen label/value/text.",
+ )
+ reasoning: str = Field(
+ default="",
+ description="One sentence explaining the labels (kept short).",
+ )
+
+
+class AgentSimulatorStrategy(AnnotationStrategy):
+ """Vision-LLM strategy for multi-modal / structured agent content."""
+
+ def __init__(self, config: AgentStrategyConfig):
+ self.config = config
+ self.endpoint = self._create_endpoint()
+ self.random_strategy = RandomStrategy()
+ # Per-instance result cache: instance_id -> dict[schema_name -> raw model value]
+ self._cache: Dict[str, Dict[str, Any]] = {}
+ # Errors are reported per-instance to avoid hammering the LLM with retries
+ self._failed_instances: set = set()
+
+ # ------------------------------------------------------------------
+ # Endpoint construction
+ # ------------------------------------------------------------------
+
+ def _create_endpoint(self):
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ ai_cfg: Dict[str, Any] = {
+ "model": self.config.model,
+ "api_key": self.config.api_key,
+ "max_tokens": self.config.max_tokens,
+ "temperature": self.config.temperature,
+ }
+ if self.config.base_url:
+ ai_cfg["base_url"] = self.config.base_url
+
+ return AIEndpointFactory.create_endpoint({
+ "ai_support": {
+ "enabled": True,
+ "endpoint_type": self.config.endpoint_type,
+ "ai_config": ai_cfg,
+ }
+ })
+ except Exception as e:
+ logger.warning("AgentSimulatorStrategy: endpoint init failed: %s", e)
+ return None
+
+ # ------------------------------------------------------------------
+ # Public API
+ # ------------------------------------------------------------------
+
+ def generate_annotation(
+ self,
+ instance: Dict[str, Any],
+ schema: Dict[str, Any],
+ competence: CompetenceProfile,
+ gold_answer: Optional[Dict[str, Any]] = None,
+ ) -> Dict[str, Any]:
+ if not self.endpoint:
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+
+ instance_id = instance.get("instance_id") or instance.get("id") or ""
+ schema_name = schema.get("name")
+ annotation_type = schema.get("annotation_type") or schema.get("type")
+ labels = self.random_strategy._extract_labels(schema)
+
+ # Per-instance cache: one LLM call answers every schema
+ results = self._get_or_query(instance, instance_id)
+ if results is None:
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+
+ # Optional noise: mirrors LLMStrategy
+ if self.config.add_noise and random.random() < self.config.noise_rate:
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+
+ raw_value = results.get(schema_name)
+ if raw_value is None:
+ logger.debug(
+ "Agent strategy: no value for schema=%s (instance=%s); falling back",
+ schema_name, instance_id,
+ )
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+
+ formatted = self._format_value(
+ schema_name, raw_value, annotation_type, labels, schema,
+ instance=instance,
+ )
+ if not formatted:
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+ return formatted
+
+ # ------------------------------------------------------------------
+ # Per-instance batch query
+ # ------------------------------------------------------------------
+
+ def _get_or_query(
+ self, instance: Dict[str, Any], instance_id: str
+ ) -> Optional[Dict[str, Any]]:
+ if self.config.cache_per_instance and instance_id in self._cache:
+ return self._cache[instance_id]
+ if instance_id in self._failed_instances:
+ return None
+
+ schemas = instance.get("__all_schemas__") or instance.get("schemas") or []
+ if not schemas:
+ # The simulator should be passing schemas via the instance dict
+ # (see SimulatedUser.generate_annotations). If it isn't, we can
+ # still produce annotations for the single schema by callers
+ # passing schema directly each time, but caching is then per-call.
+ logger.debug("Agent strategy: no schemas attached to instance %s", instance_id)
+ return None
+
+ prompt, image_payloads = self._build_request(instance, schemas)
+ try:
+ response = self._invoke(prompt, image_payloads)
+ except Exception as e:
+ logger.warning(
+ "Agent strategy: LLM call failed for instance=%s: %s", instance_id, e
+ )
+ self._failed_instances.add(instance_id)
+ return None
+
+ parsed = self._parse_response(response)
+ if parsed is None:
+ self._failed_instances.add(instance_id)
+ return None
+
+ # Models occasionally key by the schema's annotation_type instead of
+ # its name (e.g. "code_review" instead of "review"). Re-key to schema
+ # names so downstream lookup always works.
+ parsed = self._normalize_keys_to_schema_names(parsed, schemas)
+
+ if self.config.cache_per_instance:
+ self._cache[instance_id] = parsed
+ return parsed
+
+ def _normalize_keys_to_schema_names(
+ self, parsed: Dict[str, Any], schemas: List[Dict[str, Any]]
+ ) -> Dict[str, Any]:
+ """Re-key parsed annotations to match the schema *names* the
+ simulator uses, even if the LLM keyed by annotation_type, label, or
+ a case-variant. Idempotent for already-correct keys.
+ """
+ if not isinstance(parsed, dict) or not schemas:
+ return parsed
+
+ names = {s.get("name"): s for s in schemas if s.get("name")}
+ # Exact-name path -- fast and keeps existing behaviour.
+ unmatched = [k for k in parsed.keys() if k not in names]
+ if not unmatched:
+ return parsed
+
+ # Build alternate-key lookup: annotation_type -> schema name.
+ type_to_name: Dict[str, str] = {}
+ lower_name_to_name: Dict[str, str] = {n.lower(): n for n in names}
+ for s in schemas:
+ atype = s.get("annotation_type") or s.get("type")
+ if atype and atype not in names and atype not in type_to_name:
+ type_to_name[atype] = s["name"]
+
+ out = dict(parsed)
+ for key in list(unmatched):
+ value = out[key]
+ target: Optional[str] = None
+ if key in type_to_name:
+ target = type_to_name[key]
+ elif key.lower() in lower_name_to_name:
+ target = lower_name_to_name[key.lower()]
+ if target and target not in out:
+ out[target] = value
+ # Keep the original key too -- harmless and aids debugging.
+ return out
+
+ # ------------------------------------------------------------------
+ # Prompt construction
+ # ------------------------------------------------------------------
+
+ def _build_request(
+ self,
+ instance: Dict[str, Any],
+ schemas: List[Dict[str, Any]],
+ ) -> Tuple[str, List[Any]]:
+ data = instance.get("data") or {}
+ if not isinstance(data, dict):
+ data = {}
+
+ text_blocks: List[str] = []
+
+ # Top-level text/task description
+ task_text = (
+ data.get("task_description")
+ or data.get("text")
+ or instance.get("text", "")
+ )
+ if task_text:
+ text_blocks.append(f"## Task\n{task_text}".strip())
+
+ # Dialogue / conversation arrays
+ if self.config.include_dialogue_text:
+ for key in _FIELD_DETECTORS["dialogue"]:
+ value = data.get(key)
+ if value:
+ rendered = self._render_dialogue(value)
+ if rendered:
+ text_blocks.append(f"## {key.title()}\n{rendered}")
+ break # Only render the first matching dialogue field
+
+ # Spreadsheet / table data
+ if self.config.include_spreadsheet:
+ for key in _FIELD_DETECTORS["spreadsheet"]:
+ value = data.get(key)
+ if value:
+ rendered = self._render_spreadsheet(value)
+ if rendered:
+ text_blocks.append(f"## {key.replace('_', ' ').title()}\n{rendered}")
+ break
+
+ # Other plain-text fields not already consumed
+ consumed = (
+ {"task_description", "text", "id"}
+ | set(_FIELD_DETECTORS["dialogue"])
+ | set(_FIELD_DETECTORS["spreadsheet"])
+ | set(_FIELD_DETECTORS["image"])
+ | {"gold_labels"}
+ )
+ for k, v in data.items():
+ if k in consumed or k.startswith("_"):
+ continue
+ if isinstance(v, (str, int, float)) and str(v).strip():
+ text_blocks.append(f"## {k}\n{v}")
+
+ # Schema spec section (instance is needed for step-aware schemas)
+ text_blocks.append(self._render_schema_spec(schemas, instance))
+
+ text_blocks.append(
+ "Respond with a single JSON object {\"annotations\": {...}, \"reasoning\": \"...\"} "
+ "where each key under 'annotations' is exactly the schema name listed above. "
+ "The value type matches the schema:\n"
+ "- radio / multiselect / likert with named labels: a string label\n"
+ "- likert / slider / number without labels: an integer in the allowed range\n"
+ "- text / textbox: a short free-text string\n"
+ "- multiselect: a JSON array of label strings\n"
+ "- process_reward: an integer step index (or null) for first_error mode, "
+ "or a JSON array of 1/-1/0 for per_step mode\n"
+ "- code_review: a JSON object with verdict, comments, file_ratings keys\n"
+ "Always include EVERY schema name as a key under 'annotations'."
+ )
+
+ prompt = "\n\n".join(text_blocks)
+
+ # Collect image payloads (paths or URLs)
+ image_payloads = self._collect_images(data)
+ return prompt, image_payloads
+
+ def _render_dialogue(self, value: Any) -> str:
+ if isinstance(value, str):
+ return value[: self.config.max_dialogue_chars]
+ if not isinstance(value, list):
+ return ""
+ lines: List[str] = []
+ for i, turn in enumerate(value, start=1):
+ if isinstance(turn, dict):
+ speaker = turn.get("speaker") or turn.get("role") or f"Turn {i}"
+ text = turn.get("text") or turn.get("content") or ""
+ lines.append(f"{i}. {speaker}: {text}")
+ # Coding-agent shape: each turn may carry a list of
+ # {tool, input, output, output_type, language} entries.
+ # Render them so the LLM rater can see the actions taken.
+ tool_calls = turn.get("tool_calls")
+ if isinstance(tool_calls, list):
+ for call in tool_calls:
+ if not isinstance(call, dict):
+ continue
+ lines.append(self._render_tool_call(call))
+ else:
+ lines.append(f"{i}. {turn}")
+ rendered = "\n".join(lines)
+ return rendered[: self.config.max_dialogue_chars]
+
+ def _render_tool_call(self, call: Dict[str, Any]) -> str:
+ tool_name = call.get("tool") or call.get("name") or "tool"
+ # Inputs may be a dict or a string -- format for readability.
+ raw_input = call.get("input") or call.get("arguments") or {}
+ if isinstance(raw_input, dict):
+ input_str = ", ".join(f"{k}={v!r}" for k, v in raw_input.items())
+ else:
+ input_str = str(raw_input)
+ input_str = input_str[:_MAX_TOOL_INPUT_CHARS]
+
+ output = call.get("output")
+ if output is None:
+ return f" [tool: {tool_name}({input_str})]"
+ output_str = str(output)
+ if len(output_str) > _MAX_TOOL_OUTPUT_CHARS:
+ output_str = (
+ output_str[:_MAX_TOOL_OUTPUT_CHARS]
+ + f"\n [...truncated {len(output_str) - _MAX_TOOL_OUTPUT_CHARS} chars]"
+ )
+ return f" [tool: {tool_name}({input_str})]\n -> {output_str}"
+
+ def _render_spreadsheet(self, value: Any) -> str:
+ if isinstance(value, list) and value and isinstance(value[0], dict):
+ keys = list(value[0].keys())
+ header = " | ".join(keys)
+ rows = [
+ " | ".join(str(row.get(k, "")) for k in keys) for row in value
+ ]
+ return header + "\n" + "\n".join(rows)
+ if isinstance(value, dict):
+ return "\n".join(f"{k}: {v}" for k, v in value.items())
+ return str(value)
+
+ def _render_schema_spec(
+ self,
+ schemas: List[Dict[str, Any]],
+ instance: Optional[Dict[str, Any]] = None,
+ ) -> str:
+ data = (instance or {}).get("data") or {}
+ lines = ["## Schemas to label"]
+ for schema in schemas:
+ name = schema.get("name", "?")
+ atype = schema.get("annotation_type") or schema.get("type") or "?"
+ desc = schema.get("description", "")
+ labels = self.random_strategy._extract_labels(schema)
+ allowed: str
+ if labels:
+ allowed = "labels=" + ", ".join(labels)
+ elif atype == "likert":
+ size = schema.get("size", 5)
+ allowed = f"integer 1..{size}"
+ elif atype in ("slider", "number"):
+ lo = schema.get("min_value", schema.get("min", 0))
+ hi = schema.get("max_value", schema.get("max", 100))
+ allowed = f"integer {lo}..{hi}"
+ elif atype == "process_reward":
+ steps_key = schema.get("steps_key", "structured_turns")
+ steps = data.get(steps_key) if isinstance(data, dict) else None
+ n = len(steps) if isinstance(steps, list) else 0
+ mode = schema.get("mode", "first_error")
+ if mode == "first_error":
+ allowed = (
+ f"first_error mode: integer 0..{max(n - 1, 0)} "
+ f"(index of the first wrong step in the {n}-step trace), "
+ "or null if every step is correct"
+ )
+ else:
+ allowed = (
+ f"per_step mode: list of {n} entries, each one of "
+ "1 (correct), -1 (incorrect), 0 (unmarked)"
+ )
+ elif atype == "code_review":
+ verdicts = schema.get(
+ "verdict_options", ["approve", "request_changes", "comment_only"]
+ )
+ allowed = (
+ "object with keys: "
+ "verdict (one of " + ", ".join(verdicts) + "), "
+ "comments (list of {file, line?, category, body}), "
+ "file_ratings (object: filename -> {dim: 1..5})"
+ )
+ else:
+ allowed = "free text"
+ lines.append(f"- {name} ({atype}): {desc} [{allowed}]")
+ return "\n".join(lines)
+
+ # ------------------------------------------------------------------
+ # Image handling
+ # ------------------------------------------------------------------
+
+ def _collect_images(self, data: Dict[str, Any]) -> List[Any]:
+ """Return up to ``max_image_count`` ImageData objects."""
+ try:
+ from potato.ai.ai_endpoint import ImageData
+ except Exception:
+ return []
+
+ candidates: List[str] = []
+ for key in _FIELD_DETECTORS["image"]:
+ value = data.get(key)
+ if not value:
+ continue
+ if isinstance(value, str):
+ candidates.append(value)
+ elif isinstance(value, list):
+ for item in value:
+ if isinstance(item, str):
+ candidates.append(item)
+
+ images: List[Any] = []
+ for path_or_url in candidates[: self.config.max_image_count]:
+ payload = self._load_image(path_or_url, ImageData)
+ if payload is not None:
+ images.append(payload)
+ return images
+
+ def _load_image(self, path_or_url: str, ImageData):
+ try:
+ if path_or_url.startswith(("http://", "https://", "data:")):
+ # Remote / inline data URI -- pass through unchanged
+ return ImageData(url=path_or_url) if hasattr(ImageData, "url") else None
+
+ if not os.path.exists(path_or_url):
+ logger.debug("Agent strategy: image not found at %s", path_or_url)
+ return None
+
+ with open(path_or_url, "rb") as f:
+ raw = f.read()
+
+ if self.config.max_image_dim:
+ raw = self._maybe_resize(raw)
+
+ b64 = base64.b64encode(raw).decode("ascii")
+ # ImageData supports a few constructor signatures across providers;
+ # try the most compatible one first.
+ try:
+ return ImageData(base64=b64, mime_type=self._guess_mime(path_or_url))
+ except TypeError:
+ try:
+ return ImageData(data=b64, mime_type=self._guess_mime(path_or_url))
+ except TypeError:
+ return ImageData(b64)
+ except Exception as e:
+ logger.debug("Agent strategy: failed to load image %s: %s", path_or_url, e)
+ return None
+
+ def _maybe_resize(self, raw: bytes) -> bytes:
+ try:
+ from PIL import Image # noqa: WPS433
+ except Exception:
+ return raw
+ try:
+ img = Image.open(io.BytesIO(raw))
+ longest = max(img.size)
+ if longest <= self.config.max_image_dim:
+ return raw
+ scale = self.config.max_image_dim / longest
+ new_size = (max(1, int(img.size[0] * scale)), max(1, int(img.size[1] * scale)))
+ img = img.convert("RGB").resize(new_size)
+ buf = io.BytesIO()
+ img.save(buf, format="JPEG", quality=85)
+ return buf.getvalue()
+ except Exception:
+ return raw
+
+ def _guess_mime(self, path: str) -> str:
+ lowered = path.lower()
+ if lowered.endswith(".png"):
+ return "image/png"
+ if lowered.endswith((".jpg", ".jpeg")):
+ return "image/jpeg"
+ if lowered.endswith(".webp"):
+ return "image/webp"
+ if lowered.endswith(".gif"):
+ return "image/gif"
+ return "image/jpeg"
+
+ # ------------------------------------------------------------------
+ # LLM invocation + response parsing
+ # ------------------------------------------------------------------
+
+ def _invoke(self, prompt: str, image_payloads: List[Any]) -> Any:
+ """Call the endpoint, preferring vision API when images are present."""
+ if image_payloads and hasattr(self.endpoint, "query_with_image"):
+ return self.endpoint.query_with_image(
+ prompt, image_payloads, _AgentLabelResponse
+ )
+ return self.endpoint.query(prompt, _AgentLabelResponse)
+
+ def _parse_response(self, response: Any) -> Optional[Dict[str, Any]]:
+ if response is None:
+ return None
+ # Endpoints with structured output return a dict-like object
+ if hasattr(response, "model_dump"):
+ data = response.model_dump()
+ elif isinstance(response, dict):
+ data = response
+ elif isinstance(response, str):
+ data = self._loose_json_parse(response)
+ else:
+ try:
+ data = dict(response)
+ except Exception:
+ return None
+ if not isinstance(data, dict):
+ return None
+ annotations = data.get("annotations") if isinstance(data, dict) else None
+ if isinstance(annotations, dict):
+ return annotations
+ # Some endpoints return the raw annotations dict directly
+ if all(isinstance(k, str) for k in data.keys()) and "reasoning" not in data:
+ return data
+ return None
+
+ def _loose_json_parse(self, text: str) -> Dict[str, Any]:
+ try:
+ return json.loads(text)
+ except Exception:
+ pass
+ match = re.search(r"\{.*\}", text, flags=re.DOTALL)
+ if match:
+ try:
+ return json.loads(match.group(0))
+ except Exception:
+ pass
+ return {}
+
+ # ------------------------------------------------------------------
+ # Format translation: model output -> wire annotation
+ # ------------------------------------------------------------------
+
+ def _format_value(
+ self,
+ schema_name: str,
+ raw_value: Any,
+ annotation_type: str,
+ labels: List[str],
+ schema: Dict[str, Any],
+ instance: Optional[Dict[str, Any]] = None,
+ ) -> Optional[Dict[str, Any]]:
+ if annotation_type == "process_reward":
+ return self._format_process_reward(
+ schema_name, raw_value, schema, instance
+ )
+
+ if annotation_type == "code_review":
+ return self._format_code_review(schema_name, raw_value, schema)
+
+ if annotation_type == "multiselect":
+ chosen = self._coerce_multilabels(raw_value, labels)
+ if not chosen:
+ return None
+ return {f"{schema_name}:{label}": "on" for label in chosen}
+
+ if annotation_type == "radio":
+ chosen = self._coerce_label(raw_value, labels)
+ if chosen is None:
+ return None
+ return {f"{schema_name}:{chosen}": "on"}
+
+ if annotation_type == "likert":
+ size = schema.get("size", 5)
+ chosen = self._coerce_int(raw_value, 1, size)
+ if chosen is None and labels:
+ # Some likert schemas use named labels (e.g. ["Wrong","Right"])
+ lbl = self._coerce_label(raw_value, labels)
+ if lbl is not None:
+ return {f"{schema_name}:{lbl}": "on"}
+ if chosen is None:
+ return None
+ return {f"{schema_name}:{chosen}": "on"}
+
+ if annotation_type in ("slider", "number"):
+ lo = schema.get("min_value", schema.get("min", 0))
+ hi = schema.get("max_value", schema.get("max", 100))
+ chosen = self._coerce_int(raw_value, lo, hi)
+ if chosen is None:
+ return None
+ return {f"{schema_name}:{chosen}": str(chosen)}
+
+ if annotation_type in ("text", "textbox"):
+ return {f"{schema_name}:text": str(raw_value)[:1000]}
+
+ # Unknown type โ return string form
+ return {f"{schema_name}:{raw_value}": "on"}
+
+ def _coerce_label(self, raw_value: Any, labels: List[str]) -> Optional[str]:
+ if not labels:
+ return None
+ if isinstance(raw_value, str):
+ candidate = raw_value.strip()
+ for label in labels:
+ if label.lower() == candidate.lower():
+ return label
+ for label in labels:
+ if label.lower() in candidate.lower() or candidate.lower() in label.lower():
+ return label
+ return None
+
+ def _coerce_multilabels(self, raw_value: Any, labels: List[str]) -> List[str]:
+ if not labels:
+ return []
+ if isinstance(raw_value, list):
+ chosen: List[str] = []
+ for item in raw_value:
+ resolved = self._coerce_label(item, labels)
+ if resolved and resolved not in chosen:
+ chosen.append(resolved)
+ return chosen
+ if isinstance(raw_value, str):
+ parts = re.split(r"[,;|]", raw_value)
+ chosen = []
+ for part in parts:
+ resolved = self._coerce_label(part, labels)
+ if resolved and resolved not in chosen:
+ chosen.append(resolved)
+ return chosen
+ return []
+
+ # ------------------------------------------------------------------
+ # Custom-schema wire-format helpers
+ # ------------------------------------------------------------------
+
+ def _format_process_reward(
+ self,
+ schema_name: str,
+ raw_value: Any,
+ schema: Dict[str, Any],
+ instance: Optional[Dict[str, Any]],
+ ) -> Optional[Dict[str, Any]]:
+ """Build the wire-format payload for a process_reward schema.
+
+ Server expects ``{":::": ""}`` where the JSON
+ is ``{"steps": [{"index": N, "reward": 1|-1|0}, ...], "mode": ...}``.
+ """
+ steps_key = schema.get("steps_key", "structured_turns")
+ mode = schema.get("mode", "first_error")
+ data = (instance or {}).get("data") or {}
+ steps = data.get(steps_key) if isinstance(data, dict) else None
+ n = len(steps) if isinstance(steps, list) else 0
+ if n == 0:
+ return None
+
+ if mode == "first_error":
+ first_wrong = self._coerce_first_wrong_index(raw_value, n)
+ entries = []
+ for idx in range(n):
+ if first_wrong is None:
+ reward = 1
+ elif idx < first_wrong:
+ reward = 1
+ else:
+ reward = -1
+ entries.append({"index": idx, "reward": reward})
+ else:
+ entries = self._coerce_per_step_rewards(raw_value, n)
+ if entries is None:
+ return None
+
+ payload = {"steps": entries, "mode": mode}
+ return {f"{schema_name}:::{schema_name}": json.dumps(payload)}
+
+ def _coerce_first_wrong_index(self, raw_value: Any, n: int) -> Optional[int]:
+ """Interpret the LLM's first-error response as an int in 0..n-1 or None."""
+ if raw_value is None:
+ return None
+ if isinstance(raw_value, str) and raw_value.strip().lower() in (
+ "null", "none", "all_correct", "n/a", ""
+ ):
+ return None
+ if isinstance(raw_value, dict):
+ for key in ("first_wrong", "first_error", "index", "step"):
+ if key in raw_value:
+ return self._coerce_first_wrong_index(raw_value[key], n)
+ return None
+ idx = self._coerce_int(raw_value, 0, max(n - 1, 0))
+ return idx
+
+ def _coerce_per_step_rewards(
+ self, raw_value: Any, n: int
+ ) -> Optional[List[Dict[str, int]]]:
+ """Interpret the LLM's per_step response as a list of n {index,reward} entries."""
+ items: List[int] = []
+ if isinstance(raw_value, list):
+ for v in raw_value:
+ if isinstance(v, dict) and "reward" in v:
+ items.append(self._normalize_reward(v["reward"]))
+ else:
+ items.append(self._normalize_reward(v))
+ elif isinstance(raw_value, str):
+ for part in re.split(r"[\s,;|]+", raw_value):
+ if not part:
+ continue
+ items.append(self._normalize_reward(part))
+ else:
+ return None
+
+ if len(items) < n:
+ items.extend([0] * (n - len(items)))
+ items = items[:n]
+ return [{"index": i, "reward": r} for i, r in enumerate(items)]
+
+ def _normalize_reward(self, value: Any) -> int:
+ """Map various encodings to the server's {1, -1, 0} reward space."""
+ if isinstance(value, str):
+ v = value.strip().lower()
+ if v in ("1", "+1", "correct", "good", "true", "yes", "ok"):
+ return 1
+ if v in ("-1", "incorrect", "wrong", "bad", "false", "no"):
+ return -1
+ return 0
+ try:
+ i = int(value)
+ except (TypeError, ValueError):
+ return 0
+ if i > 0:
+ return 1
+ if i < 0:
+ return -1
+ return 0
+
+ def _format_code_review(
+ self,
+ schema_name: str,
+ raw_value: Any,
+ schema: Dict[str, Any],
+ ) -> Optional[Dict[str, Any]]:
+ """Build the wire-format payload for a code_review schema.
+
+ Server expects ``{":::": ""}`` where the JSON
+ is ``{"verdict": "...", "comments": [...], "file_ratings": {...}}``.
+ """
+ verdicts = schema.get(
+ "verdict_options",
+ ["approve", "request_changes", "comment_only"],
+ )
+ categories = schema.get(
+ "comment_categories",
+ ["bug", "style", "suggestion", "security", "question"],
+ )
+ rating_dims = schema.get(
+ "file_rating_dimensions",
+ ["correctness", "readability", "maintainability"],
+ )
+
+ verdict, comments, file_ratings = "comment_only", [], {}
+
+ if isinstance(raw_value, dict):
+ v = raw_value.get("verdict")
+ if isinstance(v, str):
+ v_lower = v.strip().lower()
+ for option in verdicts:
+ if v_lower == option.lower() or v_lower in option.lower():
+ verdict = option
+ break
+
+ raw_comments = raw_value.get("comments") or []
+ if isinstance(raw_comments, list):
+ for c in raw_comments:
+ if not isinstance(c, dict):
+ continue
+ body = str(c.get("body") or c.get("text") or c.get("comment") or "").strip()
+ if not body:
+ continue
+ cat = str(c.get("category") or "").strip().lower()
+ if cat not in {x.lower() for x in categories}:
+ cat = categories[0]
+ else:
+ # restore original casing
+ cat = next(x for x in categories if x.lower() == cat)
+ entry = {
+ "category": cat,
+ "body": body[:1000],
+ }
+ if c.get("file"):
+ entry["file"] = str(c["file"])
+ line = c.get("line")
+ if isinstance(line, int):
+ entry["line"] = line
+ comments.append(entry)
+
+ raw_ratings = raw_value.get("file_ratings") or raw_value.get("ratings") or {}
+ if isinstance(raw_ratings, dict):
+ for filename, dims in raw_ratings.items():
+ if not isinstance(dims, dict):
+ continue
+ clean_dims: Dict[str, int] = {}
+ for dim, score in dims.items():
+ dim_match = next(
+ (d for d in rating_dims if d.lower() == str(dim).lower()),
+ None,
+ )
+ if dim_match is None:
+ continue
+ clamped = self._coerce_int(score, 1, 5)
+ if clamped is not None:
+ clean_dims[dim_match] = clamped
+ if clean_dims:
+ file_ratings[str(filename)] = clean_dims
+
+ elif isinstance(raw_value, str):
+ v_lower = raw_value.strip().lower()
+ for option in verdicts:
+ if v_lower == option.lower() or v_lower in option.lower():
+ verdict = option
+ break
+
+ payload = {
+ "verdict": verdict,
+ "comments": comments,
+ "file_ratings": file_ratings,
+ }
+ return {f"{schema_name}:::{schema_name}": json.dumps(payload)}
+
+ def _coerce_int(self, raw_value: Any, lo: int, hi: int) -> Optional[int]:
+ try:
+ value = int(float(raw_value))
+ except (TypeError, ValueError):
+ if isinstance(raw_value, str):
+ m = re.search(r"-?\d+", raw_value)
+ if m:
+ try:
+ value = int(m.group(0))
+ except ValueError:
+ return None
+ else:
+ return None
+ else:
+ return None
+ if value < lo:
+ value = lo
+ elif value > hi:
+ value = hi
+ return value
diff --git a/potato/simulator/annotation_strategies.py b/potato/simulator/annotation_strategies.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ad283a9acdf6e37ad7f97bc674135ca61e69bfa
--- /dev/null
+++ b/potato/simulator/annotation_strategies.py
@@ -0,0 +1,780 @@
+"""
+Annotation strategies for simulated users.
+
+This module defines different strategies for generating annotations:
+- Random: Uniform random selection
+- Biased: Weighted selection based on label preferences
+- LLM: Use an LLM to generate annotations
+- Pattern: Consistent per-user patterns
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, List, Any, Optional, Tuple
+import random
+import logging
+import re
+
+from pydantic import BaseModel, Field
+
+from .competence_profiles import CompetenceProfile
+from .config import (
+ LLMStrategyConfig,
+ BiasedStrategyConfig,
+ PatternStrategyConfig,
+ AgentStrategyConfig,
+ AnnotationStrategyType,
+)
+
+
+class _LLMResponse(BaseModel):
+ """Structured-output schema for the per-schema LLMStrategy query.
+
+ Endpoints that support structured output (Ollama, OllamaVision, OpenAI
+ via the responses API, etc.) pass this Pydantic class via
+ ``output_format`` so the model emits a deterministic JSON object. The
+ ``label`` field is interpreted by ``_parse_llm_result`` according to
+ the schema's annotation_type (label name / number / free text).
+ """
+
+ label: str = Field(
+ default="",
+ description=(
+ "The single label, integer, or short text the annotator chose."
+ ),
+ )
+
+logger = logging.getLogger(__name__)
+
+
+class AnnotationStrategy(ABC):
+ """Abstract base class for annotation strategies.
+
+ Annotation strategies determine how a simulated user generates
+ annotations for different schema types.
+ """
+
+ @abstractmethod
+ def generate_annotation(
+ self,
+ instance: Dict[str, Any],
+ schema: Dict[str, Any],
+ competence: CompetenceProfile,
+ gold_answer: Optional[Dict[str, Any]] = None,
+ ) -> Dict[str, Any]:
+ """Generate an annotation for the given instance and schema.
+
+ Args:
+ instance: The data instance containing text and metadata
+ schema: The annotation schema definition
+ competence: Competence profile for accuracy modeling
+ gold_answer: Gold standard answer if available (for competence)
+
+ Returns:
+ Dictionary with annotation data in the format expected by the API
+ """
+ pass
+
+
+class RandomStrategy(AnnotationStrategy):
+ """Random annotation selection strategy.
+
+ Selects labels uniformly at random. When gold standards are available
+ and competence should be correct, uses the gold answer instead.
+ """
+
+ def generate_annotation(
+ self,
+ instance: Dict[str, Any],
+ schema: Dict[str, Any],
+ competence: CompetenceProfile,
+ gold_answer: Optional[Dict[str, Any]] = None,
+ ) -> Dict[str, Any]:
+ """Generate random annotation.
+
+ Args:
+ instance: Data instance
+ schema: Annotation schema
+ competence: Competence profile
+ gold_answer: Gold standard if available
+
+ Returns:
+ Annotation dictionary
+ """
+ # Check both 'annotation_type' (config format) and 'type' (API format)
+ annotation_type = schema.get("annotation_type") or schema.get("type")
+ labels = self._extract_labels(schema)
+ schema_name = schema.get("name")
+
+ # If we have gold answer and competence says be correct, use gold
+ if gold_answer and schema_name in gold_answer:
+ if competence.should_be_correct():
+ return self._format_gold_answer(schema_name, gold_answer[schema_name], annotation_type)
+ else:
+ # Select wrong answer
+ correct = gold_answer[schema_name]
+ wrong = competence.select_wrong_answer(str(correct), labels)
+ return self._format_annotation(schema_name, wrong, annotation_type)
+
+ # No gold standard - just random selection
+ return self._generate_by_type(annotation_type, schema, labels, instance)
+
+ def _extract_labels(self, schema: Dict[str, Any]) -> List[str]:
+ """Extract label options from schema.
+
+ Args:
+ schema: Annotation schema
+
+ Returns:
+ List of label names
+ """
+ labels = schema.get("labels", [])
+ if not labels:
+ return []
+
+ if isinstance(labels[0], dict):
+ return [l.get("name") for l in labels if l.get("name")]
+ return [str(l) for l in labels]
+
+ def _format_gold_answer(
+ self, schema_name: str, gold_value: Any, annotation_type: str
+ ) -> Dict[str, Any]:
+ """Format gold answer as annotation.
+
+ Args:
+ schema_name: Schema name
+ gold_value: Gold standard value
+ annotation_type: Type of annotation
+
+ Returns:
+ Formatted annotation
+ """
+ return self._format_annotation(schema_name, gold_value, annotation_type)
+
+ def _format_annotation(
+ self, schema_name: str, value: Any, annotation_type: str
+ ) -> Dict[str, Any]:
+ """Format a value as an annotation.
+
+ Args:
+ schema_name: Schema name
+ value: Annotation value
+ annotation_type: Type of annotation
+
+ Returns:
+ Formatted annotation dictionary in the format expected by the server
+ (schema:value -> "on" for selection types)
+ """
+ if annotation_type == "multiselect":
+ if isinstance(value, list):
+ return {f"{schema_name}:{v}": "on" for v in value}
+ return {f"{schema_name}:{value}": "on"}
+ elif annotation_type in ["radio", "likert"]:
+ # Selection-based types use schema:value format
+ return {f"{schema_name}:{value}": "on"}
+ elif annotation_type in ["slider", "number"]:
+ # Numeric types store the value
+ return {f"{schema_name}:{value}": str(value)}
+ elif annotation_type in ["text", "textbox"]:
+ # Text types store the text content
+ return {f"{schema_name}:text": str(value)}
+ else:
+ # Default: use schema:value format
+ return {f"{schema_name}:{value}": "on"}
+
+ def _generate_by_type(
+ self,
+ annotation_type: str,
+ schema: Dict[str, Any],
+ labels: List[str],
+ instance: Dict[str, Any],
+ ) -> Dict[str, Any]:
+ """Generate annotation based on schema type.
+
+ Args:
+ annotation_type: Type of annotation
+ schema: Full schema definition
+ labels: Available labels
+ instance: Data instance
+
+ Returns:
+ Generated annotation
+ """
+ schema_name = schema.get("name")
+
+ if annotation_type == "radio":
+ if labels:
+ selected_label = random.choice(labels)
+ # Format as "schema:label": "on" to match frontend format
+ return {f"{schema_name}:{selected_label}": "on"}
+ return {}
+
+ elif annotation_type == "multiselect":
+ if labels:
+ # Select 1-3 random labels
+ num_selections = random.randint(1, min(3, len(labels)))
+ selections = random.sample(labels, num_selections)
+ return {f"{schema_name}:{label}": "on" for label in selections}
+ return {}
+
+ elif annotation_type == "likert":
+ size = schema.get("size", 5)
+ selected_value = str(random.randint(1, size))
+ # Format as "schema:value": "on" to match frontend format
+ return {f"{schema_name}:{selected_value}": "on"}
+
+ elif annotation_type == "slider":
+ min_val = schema.get("min_value", schema.get("min", 0))
+ max_val = schema.get("max_value", schema.get("max", 100))
+ selected_value = str(random.randint(min_val, max_val))
+ # Format as "schema:value": "value" for slider
+ return {f"{schema_name}:{selected_value}": selected_value}
+
+ elif annotation_type in ["text", "textbox"]:
+ text_response = self._generate_text_response(instance)
+ # Format as "schema:text": "value" for textbox
+ return {f"{schema_name}:text": text_response}
+
+ elif annotation_type == "number":
+ min_val = schema.get("min_value", 0)
+ max_val = schema.get("max_value", 100)
+ selected_value = str(random.randint(min_val, max_val))
+ # Format as "schema:value": "value"
+ return {f"{schema_name}:{selected_value}": selected_value}
+
+ elif annotation_type == "span":
+ return self._generate_span_annotation(instance, schema, labels)
+
+ else:
+ logger.warning(f"Unknown annotation type: {annotation_type}")
+ if labels:
+ return {schema_name: random.choice(labels)}
+ return {}
+
+ def _generate_text_response(self, instance: Dict[str, Any]) -> str:
+ """Generate a placeholder text response.
+
+ Args:
+ instance: Data instance
+
+ Returns:
+ Generated text
+ """
+ responses = [
+ "Simulated annotation response.",
+ "This is a test response.",
+ "Generated text for testing purposes.",
+ "Sample annotation text.",
+ ]
+ return random.choice(responses)
+
+ def _generate_span_annotation(
+ self,
+ instance: Dict[str, Any],
+ schema: Dict[str, Any],
+ labels: List[str],
+ ) -> Dict[str, Any]:
+ """Generate span annotations for text.
+
+ Args:
+ instance: Data instance with text
+ schema: Span annotation schema
+ labels: Available span labels
+
+ Returns:
+ Span annotation dictionary
+ """
+ text = instance.get("text", "")
+ if not text or not labels:
+ return {}
+
+ words = text.split()
+ if len(words) < 2:
+ return {}
+
+ # Generate 0-3 random spans
+ num_spans = random.randint(0, min(3, len(words) // 2))
+ schema_name = schema.get("name", "spans")
+ annotations = {}
+
+ for _ in range(num_spans):
+ start_word_idx = random.randint(0, len(words) - 2)
+ end_word_idx = random.randint(
+ start_word_idx + 1, min(start_word_idx + 5, len(words))
+ )
+
+ # Calculate character offsets
+ start_char = sum(len(w) + 1 for w in words[:start_word_idx])
+ end_char = sum(len(w) + 1 for w in words[:end_word_idx]) - 1
+
+ label = random.choice(labels)
+ span_key = f"{schema_name}:{label}:{start_char}:{end_char}"
+ annotations[span_key] = "true"
+
+ return annotations
+
+
+class BiasedStrategy(AnnotationStrategy):
+ """Annotation strategy with configurable label biases.
+
+ Selects labels according to configured weights, allowing simulation
+ of annotators with specific label preferences.
+ """
+
+ def __init__(self, config: BiasedStrategyConfig):
+ """Initialize biased strategy.
+
+ Args:
+ config: Configuration with label weights
+ """
+ self.config = config
+ self.random_strategy = RandomStrategy()
+
+ def generate_annotation(
+ self,
+ instance: Dict[str, Any],
+ schema: Dict[str, Any],
+ competence: CompetenceProfile,
+ gold_answer: Optional[Dict[str, Any]] = None,
+ ) -> Dict[str, Any]:
+ """Generate biased annotation.
+
+ Args:
+ instance: Data instance
+ schema: Annotation schema
+ competence: Competence profile
+ gold_answer: Gold standard if available
+
+ Returns:
+ Annotation dictionary
+ """
+ # If gold answer available and competence says be correct, use it
+ if gold_answer and competence.should_be_correct():
+ schema_name = schema.get("name")
+ if schema_name in gold_answer:
+ annotation_type = schema.get("annotation_type") or schema.get("type")
+ return self.random_strategy._format_gold_answer(
+ schema_name, gold_answer[schema_name], annotation_type
+ )
+
+ annotation_type = schema.get("annotation_type") or schema.get("type")
+ labels = self.random_strategy._extract_labels(schema)
+ schema_name = schema.get("name")
+
+ if annotation_type in ["radio", "multiselect"] and labels:
+ # Use weighted selection based on bias config
+ weights = [self.config.label_weights.get(l, 1.0) for l in labels]
+ total = sum(weights)
+ if total > 0:
+ weights = [w / total for w in weights]
+ else:
+ weights = [1.0 / len(labels)] * len(labels)
+
+ selected = random.choices(labels, weights=weights, k=1)[0]
+
+ # Use consistent schema:value format for all selection types
+ return {f"{schema_name}:{selected}": "on"}
+
+ # Fall back to random for other types
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+
+
+class LLMStrategy(AnnotationStrategy):
+ """LLM-powered annotation strategy.
+
+ Uses the existing potato.ai infrastructure to generate realistic
+ annotations based on text content.
+ """
+
+ def __init__(self, config: LLMStrategyConfig):
+ """Initialize LLM strategy.
+
+ Args:
+ config: LLM configuration
+ """
+ self.config = config
+ self.endpoint = self._create_endpoint()
+ self.random_strategy = RandomStrategy()
+
+ def _create_endpoint(self):
+ """Create LLM endpoint using existing infrastructure.
+
+ Returns:
+ AI endpoint or None if creation fails
+ """
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ ai_config = {
+ "ai_support": {
+ "enabled": True,
+ "endpoint_type": self.config.endpoint_type,
+ "ai_config": {
+ "model": self.config.model,
+ "api_key": self.config.api_key,
+ "max_tokens": self.config.max_tokens,
+ "temperature": self.config.temperature,
+ },
+ }
+ }
+
+ if self.config.base_url:
+ ai_config["ai_support"]["ai_config"]["base_url"] = self.config.base_url
+
+ return AIEndpointFactory.create_endpoint(ai_config)
+
+ except Exception as e:
+ logger.warning(f"Failed to create LLM endpoint: {e}")
+ return None
+
+ def generate_annotation(
+ self,
+ instance: Dict[str, Any],
+ schema: Dict[str, Any],
+ competence: CompetenceProfile,
+ gold_answer: Optional[Dict[str, Any]] = None,
+ ) -> Dict[str, Any]:
+ """Generate LLM-based annotation.
+
+ Args:
+ instance: Data instance
+ schema: Annotation schema
+ competence: Competence profile
+ gold_answer: Gold standard if available
+
+ Returns:
+ Annotation dictionary
+ """
+ if not self.endpoint:
+ logger.warning("LLM endpoint not available, falling back to random")
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+
+ try:
+ annotation_type = schema.get("annotation_type") or schema.get("type")
+ labels = self.random_strategy._extract_labels(schema)
+ schema_name = schema.get("name")
+ description = schema.get("description", "")
+ text = instance.get("text", "")
+
+ # Build prompt for LLM
+ prompt = self._build_prompt(text, labels, description, annotation_type)
+
+ # Query LLM. Most endpoints (Ollama, OllamaVision, OpenAI) require
+ # a Pydantic schema as the second argument; AnthropicEndpoint's
+ # query() takes only `prompt`; VLLMEndpoint accepts both.
+ try:
+ result = self.endpoint.query(prompt, _LLMResponse)
+ except TypeError:
+ # Endpoint signature is query(prompt) โ single-arg providers
+ result = self.endpoint.query(prompt)
+
+ # Add noise if configured
+ if self.config.add_noise and random.random() < self.config.noise_rate:
+ logger.debug("Adding noise to LLM response")
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+
+ # Parse result
+ parsed = self._parse_llm_result(result, labels, schema_name, annotation_type)
+ if parsed:
+ return parsed
+
+ # Fallback to random
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+
+ except Exception as e:
+ logger.warning(f"LLM annotation failed: {e}")
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+
+ def _build_prompt(
+ self,
+ text: str,
+ labels: List[str],
+ description: str,
+ annotation_type: str,
+ ) -> str:
+ """Build prompt for LLM.
+
+ Args:
+ text: Text to annotate
+ labels: Available labels
+ description: Task description
+ annotation_type: Type of annotation
+
+ Returns:
+ Prompt string
+ """
+ labels_str = ", ".join(labels)
+
+ if annotation_type in ["radio", "multiselect"]:
+ prompt = f"""You are an annotator. Given the following text, select the most appropriate label.
+
+Task: {description if description else 'Classify the text'}
+Labels: {labels_str}
+
+Text: {text[:500]}
+
+Respond with ONLY the label name, nothing else."""
+ elif annotation_type == "likert":
+ prompt = f"""You are an annotator. Rate the following text on a scale.
+
+Task: {description if description else 'Rate the text'}
+
+Text: {text[:500]}
+
+Respond with ONLY a number from 1-5, nothing else."""
+ else:
+ prompt = f"""You are an annotator. Analyze the following text.
+
+Task: {description if description else 'Analyze the text'}
+
+Text: {text[:500]}
+
+Respond briefly."""
+
+ return prompt
+
+ def _parse_llm_result(
+ self,
+ result: Any,
+ labels: List[str],
+ schema_name: str,
+ annotation_type: str,
+ ) -> Optional[Dict[str, Any]]:
+ """Parse LLM result into annotation format.
+
+ Args:
+ result: LLM response (string, dict, or Pydantic model instance)
+ labels: Available labels
+ schema_name: Schema name
+ annotation_type: Type of annotation
+
+ Returns:
+ Parsed annotation or None
+ """
+ if result is None:
+ return None
+
+ # Structured-output endpoints return either a Pydantic model or a dict
+ # with a ``label`` key (or, for OllamaEndpoint's parseStringToJson
+ # fallback, ``response`` / ``content``). Extract that into a string
+ # for the existing matching logic.
+ if hasattr(result, "model_dump"):
+ try:
+ result = result.model_dump()
+ except Exception:
+ pass
+ if isinstance(result, dict):
+ for key in ("label", "response", "content"):
+ if key in result and result[key] not in (None, ""):
+ result = result[key]
+ break
+ else:
+ # Dict with no recognised key โ stringify everything
+ result = str(result)
+
+ # Convert to string
+ result_str = str(result).strip().lower()
+
+ if annotation_type in ["radio", "multiselect"]:
+ # Try to match a label
+ for label in labels:
+ if label.lower() in result_str or result_str in label.lower():
+ return self.random_strategy._format_annotation(
+ schema_name, label, annotation_type
+ )
+
+ elif annotation_type == "likert":
+ # Try to extract a number
+ numbers = re.findall(r"\d+", result_str)
+ if numbers:
+ return self.random_strategy._format_annotation(
+ schema_name, numbers[0], annotation_type
+ )
+
+ elif annotation_type in ["text", "textbox"]:
+ return self.random_strategy._format_annotation(
+ schema_name, str(result)[:500], annotation_type
+ )
+
+ return None
+
+
+class PatternStrategy(AnnotationStrategy):
+ """Pattern-based annotation strategy for consistent user behavior.
+
+ Allows defining specific behavior patterns per user for testing
+ scenarios that require consistent annotation patterns.
+ """
+
+ def __init__(self, config: PatternStrategyConfig, user_id: str):
+ """Initialize pattern strategy.
+
+ Args:
+ config: Pattern configuration
+ user_id: User ID for pattern lookup
+ """
+ self.config = config
+ self.user_id = user_id
+ self.user_pattern = config.patterns.get(user_id, {})
+ self.random_strategy = RandomStrategy()
+
+ def generate_annotation(
+ self,
+ instance: Dict[str, Any],
+ schema: Dict[str, Any],
+ competence: CompetenceProfile,
+ gold_answer: Optional[Dict[str, Any]] = None,
+ ) -> Dict[str, Any]:
+ """Generate pattern-based annotation.
+
+ Args:
+ instance: Data instance
+ schema: Annotation schema
+ competence: Competence profile
+ gold_answer: Gold standard if available
+
+ Returns:
+ Annotation dictionary
+ """
+ preferred_label = self.user_pattern.get("preferred_label")
+ bias_strength = self.user_pattern.get("bias_strength", 0.5)
+
+ # Check for keyword patterns
+ text = instance.get("text", "").lower()
+ keyword_labels = self.user_pattern.get("keywords", {})
+ for keyword, label in keyword_labels.items():
+ if keyword.lower() in text:
+ return self.random_strategy._format_annotation(
+ schema.get("name"), label, schema.get("annotation_type")
+ )
+
+ # Use preferred label with configured probability
+ if preferred_label and random.random() < bias_strength:
+ labels = self.random_strategy._extract_labels(schema)
+ if preferred_label in labels:
+ return self.random_strategy._format_annotation(
+ schema.get("name"),
+ preferred_label,
+ schema.get("annotation_type"),
+ )
+
+ # Fall back to random
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+
+
+class GoldStandardStrategy(AnnotationStrategy):
+ """Strategy that uses gold standard answers when available.
+
+ This is primarily useful for testing quality control systems
+ by providing known correct annotations.
+ """
+
+ def __init__(self):
+ self.random_strategy = RandomStrategy()
+
+ def generate_annotation(
+ self,
+ instance: Dict[str, Any],
+ schema: Dict[str, Any],
+ competence: CompetenceProfile,
+ gold_answer: Optional[Dict[str, Any]] = None,
+ ) -> Dict[str, Any]:
+ """Generate annotation from gold standard.
+
+ Args:
+ instance: Data instance
+ schema: Annotation schema
+ competence: Competence profile (determines if we use gold)
+ gold_answer: Gold standard if available
+
+ Returns:
+ Annotation dictionary
+ """
+ schema_name = schema.get("name")
+ annotation_type = schema.get("annotation_type") or schema.get("type")
+
+ if gold_answer and schema_name in gold_answer:
+ # Use competence to decide if we get it right
+ if competence.should_be_correct():
+ return self.random_strategy._format_gold_answer(
+ schema_name, gold_answer[schema_name], annotation_type
+ )
+ else:
+ # Select wrong answer
+ labels = self.random_strategy._extract_labels(schema)
+ correct = str(gold_answer[schema_name])
+ wrong = competence.select_wrong_answer(correct, labels)
+ return self.random_strategy._format_annotation(
+ schema_name, wrong, annotation_type
+ )
+
+ # No gold standard - fall back to random
+ return self.random_strategy.generate_annotation(
+ instance, schema, competence, gold_answer
+ )
+
+
+def create_strategy(
+ strategy_type: AnnotationStrategyType,
+ llm_config: Optional[LLMStrategyConfig] = None,
+ biased_config: Optional[BiasedStrategyConfig] = None,
+ pattern_config: Optional[PatternStrategyConfig] = None,
+ agent_config: Optional[AgentStrategyConfig] = None,
+ user_id: str = "",
+) -> AnnotationStrategy:
+ """Factory function to create annotation strategies.
+
+ Args:
+ strategy_type: Type of strategy to create
+ llm_config: LLM configuration (for LLM strategy)
+ biased_config: Bias configuration (for biased strategy)
+ pattern_config: Pattern configuration (for pattern strategy)
+ agent_config: Agent (vision-LLM) configuration (for AGENT strategy)
+ user_id: User ID (for pattern strategy)
+
+ Returns:
+ AnnotationStrategy instance
+ """
+ if strategy_type == AnnotationStrategyType.RANDOM:
+ return RandomStrategy()
+
+ elif strategy_type == AnnotationStrategyType.BIASED:
+ if biased_config:
+ return BiasedStrategy(biased_config)
+ return RandomStrategy()
+
+ elif strategy_type == AnnotationStrategyType.LLM:
+ if llm_config:
+ return LLMStrategy(llm_config)
+ logger.warning("LLM strategy requested but no config provided, using random")
+ return RandomStrategy()
+
+ elif strategy_type == AnnotationStrategyType.PATTERN:
+ if pattern_config:
+ return PatternStrategy(pattern_config, user_id)
+ return RandomStrategy()
+
+ elif strategy_type == AnnotationStrategyType.GOLD_STANDARD:
+ return GoldStandardStrategy()
+
+ elif strategy_type == AnnotationStrategyType.AGENT:
+ # Local import to avoid pulling pydantic / vision deps unless used
+ from .agent_strategy import AgentSimulatorStrategy
+ if agent_config:
+ return AgentSimulatorStrategy(agent_config)
+ logger.warning("AGENT strategy requested but no agent_config provided, using random")
+ return RandomStrategy()
+
+ else:
+ return RandomStrategy()
diff --git a/potato/simulator/cli.py b/potato/simulator/cli.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d4bb5146b2e9a50ea3172aa9b537bf93206d1ad
--- /dev/null
+++ b/potato/simulator/cli.py
@@ -0,0 +1,369 @@
+"""
+Command-line interface for the user simulator.
+
+Usage:
+ python -m potato.simulator --server http://localhost:8000 --users 10
+ python -m potato.simulator --config simulator-config.yaml --server http://localhost:8000
+"""
+
+import argparse
+import logging
+import sys
+import os
+
+from .config import (
+ SimulatorConfig,
+ TimingConfig,
+ LLMStrategyConfig,
+ BiasedStrategyConfig,
+ AnnotationStrategyType,
+)
+from .simulator_manager import SimulatorManager
+
+
+def setup_logging(verbose: bool = False) -> None:
+ """Configure logging for the CLI.
+
+ Args:
+ verbose: If True, enable debug logging
+ """
+ level = logging.DEBUG if verbose else logging.INFO
+ logging.basicConfig(
+ level=level,
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+ datefmt="%H:%M:%S",
+ )
+
+ # Suppress noisy loggers
+ if not verbose:
+ logging.getLogger("urllib3").setLevel(logging.WARNING)
+ logging.getLogger("requests").setLevel(logging.WARNING)
+
+
+def parse_args() -> argparse.Namespace:
+ """Parse command-line arguments.
+
+ Returns:
+ Parsed arguments namespace
+ """
+ parser = argparse.ArgumentParser(
+ description="User Simulator for Potato Annotation Platform",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+Examples:
+ # Basic random simulation
+ python -m potato.simulator --server http://localhost:8000 --users 10
+
+ # With configuration file
+ python -m potato.simulator --config simulator.yaml --server http://localhost:8000
+
+ # LLM-powered simulation with Ollama
+ python -m potato.simulator --server http://localhost:8000 --users 5 \\
+ --strategy llm --llm-endpoint ollama --llm-model llama3.2
+
+ # Biased simulation
+ python -m potato.simulator --server http://localhost:8000 --users 20 \\
+ --strategy biased --bias-weights positive=0.6,negative=0.3,neutral=0.1
+
+ # Fast scalability test
+ python -m potato.simulator --server http://localhost:8000 --users 100 \\
+ --parallel 20 --max-annotations 5 --fast-mode
+""",
+ )
+
+ # Required arguments
+ parser.add_argument(
+ "--server",
+ "-s",
+ required=True,
+ help="Potato server URL (e.g., http://localhost:8000)",
+ )
+
+ # Configuration file (alternative to CLI args)
+ parser.add_argument(
+ "--config",
+ "-c",
+ help="Path to YAML configuration file",
+ )
+
+ # User configuration
+ parser.add_argument(
+ "--users",
+ "-u",
+ type=int,
+ default=10,
+ help="Number of simulated users (default: 10)",
+ )
+ parser.add_argument(
+ "--competence",
+ help="Competence distribution as comma-separated key=value pairs "
+ "(e.g., good=0.5,average=0.3,poor=0.2)",
+ )
+
+ # Strategy configuration
+ parser.add_argument(
+ "--strategy",
+ choices=["random", "biased", "llm", "pattern", "gold_standard"],
+ default="random",
+ help="Annotation strategy (default: random)",
+ )
+
+ # LLM configuration
+ parser.add_argument(
+ "--llm-endpoint",
+ choices=["openai", "anthropic", "ollama", "gemini", "huggingface", "vllm"],
+ help="LLM endpoint type (for --strategy llm)",
+ )
+ parser.add_argument(
+ "--llm-model",
+ help="LLM model name (for --strategy llm)",
+ )
+ parser.add_argument(
+ "--llm-api-key",
+ help="LLM API key (or set via environment variable)",
+ )
+ parser.add_argument(
+ "--llm-base-url",
+ help="LLM base URL (for local endpoints like Ollama)",
+ )
+
+ # Biased strategy configuration
+ parser.add_argument(
+ "--bias-weights",
+ help="Label bias weights as comma-separated key=value pairs "
+ "(e.g., positive=0.6,negative=0.3,neutral=0.1)",
+ )
+
+ # Execution configuration
+ parser.add_argument(
+ "--parallel",
+ "-p",
+ type=int,
+ default=5,
+ help="Maximum concurrent users (default: 5)",
+ )
+ parser.add_argument(
+ "--max-annotations",
+ "-m",
+ type=int,
+ help="Maximum annotations per user (default: unlimited)",
+ )
+ parser.add_argument(
+ "--sequential",
+ action="store_true",
+ help="Run users sequentially instead of in parallel",
+ )
+
+ # Timing configuration
+ parser.add_argument(
+ "--fast-mode",
+ action="store_true",
+ help="Disable waiting between annotations (for testing)",
+ )
+ parser.add_argument(
+ "--timing-min",
+ type=float,
+ default=2.0,
+ help="Minimum annotation time in seconds (default: 2.0)",
+ )
+ parser.add_argument(
+ "--timing-max",
+ type=float,
+ default=30.0,
+ help="Maximum annotation time in seconds (default: 30.0)",
+ )
+
+ # Quality control testing
+ parser.add_argument(
+ "--attention-fail-rate",
+ type=float,
+ default=0.0,
+ help="Rate at which to fail attention checks (0-1, default: 0)",
+ )
+ parser.add_argument(
+ "--fast-response-rate",
+ type=float,
+ default=0.0,
+ help="Rate of suspiciously fast responses (0-1, default: 0)",
+ )
+
+ # Gold standards
+ parser.add_argument(
+ "--gold-file",
+ help="Path to JSON file with gold standard answers",
+ )
+
+ # Output configuration
+ parser.add_argument(
+ "--output-dir",
+ "-o",
+ default="simulator_output",
+ help="Output directory for results (default: simulator_output)",
+ )
+ parser.add_argument(
+ "--no-export",
+ action="store_true",
+ help="Don't export results to files",
+ )
+
+ # Other options
+ parser.add_argument(
+ "--verbose",
+ "-v",
+ action="store_true",
+ help="Enable verbose logging",
+ )
+
+ return parser.parse_args()
+
+
+def parse_key_value_pairs(s: str) -> dict:
+ """Parse comma-separated key=value pairs.
+
+ Args:
+ s: String like "key1=val1,key2=val2"
+
+ Returns:
+ Dictionary of parsed pairs
+ """
+ result = {}
+ if not s:
+ return result
+
+ for pair in s.split(","):
+ if "=" in pair:
+ key, value = pair.split("=", 1)
+ # Try to convert to float
+ try:
+ result[key.strip()] = float(value.strip())
+ except ValueError:
+ result[key.strip()] = value.strip()
+
+ return result
+
+
+def build_config_from_args(args: argparse.Namespace) -> SimulatorConfig:
+ """Build SimulatorConfig from CLI arguments.
+
+ Args:
+ args: Parsed arguments
+
+ Returns:
+ SimulatorConfig instance
+ """
+ # If config file provided, use it as base
+ if args.config:
+ config = SimulatorConfig.from_yaml(args.config)
+ else:
+ config = SimulatorConfig()
+
+ # Override with CLI arguments
+ config.user_count = args.users
+ config.parallel_users = args.parallel
+ config.simulate_wait = not args.fast_mode
+ config.attention_check_fail_rate = args.attention_fail_rate
+ config.respond_fast_rate = args.fast_response_rate
+ config.output_dir = args.output_dir
+
+ # Parse competence distribution
+ if args.competence:
+ config.competence_distribution = parse_key_value_pairs(args.competence)
+
+ # Parse strategy
+ try:
+ config.strategy = AnnotationStrategyType(args.strategy)
+ except ValueError:
+ config.strategy = AnnotationStrategyType.RANDOM
+
+ # LLM configuration
+ if args.strategy == "llm" and args.llm_endpoint:
+ api_key = args.llm_api_key
+ if not api_key:
+ # Try common environment variables
+ env_vars = {
+ "openai": "OPENAI_API_KEY",
+ "anthropic": "ANTHROPIC_API_KEY",
+ "huggingface": "HF_TOKEN",
+ "gemini": "GOOGLE_API_KEY",
+ }
+ env_var = env_vars.get(args.llm_endpoint)
+ if env_var:
+ api_key = os.environ.get(env_var)
+
+ config.llm_config = LLMStrategyConfig(
+ endpoint_type=args.llm_endpoint,
+ model=args.llm_model,
+ api_key=api_key,
+ base_url=args.llm_base_url,
+ )
+
+ # Biased configuration
+ if args.strategy == "biased" and args.bias_weights:
+ config.biased_config = BiasedStrategyConfig(
+ label_weights=parse_key_value_pairs(args.bias_weights)
+ )
+
+ # Timing configuration
+ config.timing = TimingConfig(
+ annotation_time_min=args.timing_min,
+ annotation_time_max=args.timing_max,
+ )
+
+ # Gold standards file
+ if args.gold_file:
+ config.gold_standard_file = args.gold_file
+
+ return config
+
+
+def main() -> int:
+ """Main entry point for CLI.
+
+ Returns:
+ Exit code (0 for success, 1 for error)
+ """
+ args = parse_args()
+ setup_logging(args.verbose)
+
+ logger = logging.getLogger(__name__)
+
+ try:
+ # Build configuration
+ config = build_config_from_args(args)
+
+ logger.info(f"Starting simulator with {config.user_count} users")
+ logger.info(f"Server: {args.server}")
+ logger.info(f"Strategy: {config.strategy.value}")
+
+ # Create manager
+ manager = SimulatorManager(config, args.server)
+
+ # Run simulation
+ if args.sequential:
+ results = manager.run_sequential(args.max_annotations)
+ else:
+ results = manager.run_parallel(args.max_annotations)
+
+ # Print summary
+ manager.print_summary()
+
+ # Export results
+ if not args.no_export:
+ manager.export_results()
+
+ return 0
+
+ except KeyboardInterrupt:
+ logger.info("Simulation interrupted by user")
+ return 1
+
+ except Exception as e:
+ logger.error(f"Simulation failed: {e}")
+ if args.verbose:
+ import traceback
+ traceback.print_exc()
+ return 1
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/potato/simulator/competence_profiles.py b/potato/simulator/competence_profiles.py
new file mode 100644
index 0000000000000000000000000000000000000000..6959ae8b8a55080c1054275effd06c6eeef9ed04
--- /dev/null
+++ b/potato/simulator/competence_profiles.py
@@ -0,0 +1,223 @@
+"""
+Competence profiles for simulated users.
+
+This module defines different competence levels that determine how
+accurately a simulated user will annotate items.
+
+When gold standards are available, competence determines the probability
+of selecting the gold answer. Without gold standards, competence affects
+consistency and selection patterns.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Optional, Dict, Any
+import random
+
+from .config import CompetenceLevel
+
+
+class CompetenceProfile(ABC):
+ """Abstract base class for competence profiles.
+
+ Competence profiles determine:
+ 1. Whether to select the correct answer (if gold standard available)
+ 2. How to select an answer when being incorrect
+ """
+
+ @abstractmethod
+ def should_be_correct(self) -> bool:
+ """Determine if this annotation should be correct.
+
+ Returns:
+ True if the annotation should match the gold standard
+ """
+ pass
+
+ @abstractmethod
+ def select_wrong_answer(self, correct: str, options: List[str]) -> str:
+ """Select an incorrect answer when making a mistake.
+
+ Args:
+ correct: The correct answer (to avoid)
+ options: All available options
+
+ Returns:
+ A selected incorrect option
+ """
+ pass
+
+ def get_accuracy(self) -> float:
+ """Get the expected accuracy for this profile.
+
+ Returns:
+ Expected accuracy as a float between 0 and 1
+ """
+ return 0.5
+
+
+class PerfectCompetence(CompetenceProfile):
+ """Always correct annotations (100% accuracy).
+
+ Use this for testing gold standard tracking or simulating
+ expert annotators.
+ """
+
+ def should_be_correct(self) -> bool:
+ return True
+
+ def select_wrong_answer(self, correct: str, options: List[str]) -> str:
+ # Never called, but return correct just in case
+ return correct
+
+ def get_accuracy(self) -> float:
+ return 1.0
+
+
+class GoodCompetence(CompetenceProfile):
+ """High-quality annotator (80-90% accuracy).
+
+ Simulates a careful, well-trained annotator who occasionally
+ makes mistakes on ambiguous items.
+ """
+
+ def __init__(self, accuracy_range: tuple = (0.80, 0.90)):
+ self.accuracy = random.uniform(*accuracy_range)
+
+ def should_be_correct(self) -> bool:
+ return random.random() < self.accuracy
+
+ def select_wrong_answer(self, correct: str, options: List[str]) -> str:
+ wrong_options = [o for o in options if o != correct]
+ if wrong_options:
+ return random.choice(wrong_options)
+ return correct
+
+ def get_accuracy(self) -> float:
+ return self.accuracy
+
+
+class AverageCompetence(CompetenceProfile):
+ """Typical annotator (60-70% accuracy).
+
+ Simulates an average crowdworker with moderate attention
+ and understanding of the task.
+ """
+
+ def __init__(self, accuracy_range: tuple = (0.60, 0.70)):
+ self.accuracy = random.uniform(*accuracy_range)
+
+ def should_be_correct(self) -> bool:
+ return random.random() < self.accuracy
+
+ def select_wrong_answer(self, correct: str, options: List[str]) -> str:
+ wrong_options = [o for o in options if o != correct]
+ if wrong_options:
+ return random.choice(wrong_options)
+ return correct
+
+ def get_accuracy(self) -> float:
+ return self.accuracy
+
+
+class PoorCompetence(CompetenceProfile):
+ """Low-quality annotator (40-50% accuracy).
+
+ Simulates an inattentive or untrained annotator who often
+ makes mistakes or doesn't fully understand the task.
+ """
+
+ def __init__(self, accuracy_range: tuple = (0.40, 0.50)):
+ self.accuracy = random.uniform(*accuracy_range)
+
+ def should_be_correct(self) -> bool:
+ return random.random() < self.accuracy
+
+ def select_wrong_answer(self, correct: str, options: List[str]) -> str:
+ wrong_options = [o for o in options if o != correct]
+ if wrong_options:
+ return random.choice(wrong_options)
+ return correct
+
+ def get_accuracy(self) -> float:
+ return self.accuracy
+
+
+class RandomCompetence(CompetenceProfile):
+ """Random selection regardless of correct answer.
+
+ Does not use gold standards at all - simply selects randomly
+ from available options. Expected accuracy is ~1/N for N labels.
+ """
+
+ def should_be_correct(self) -> bool:
+ # Always select randomly - don't use gold standard
+ return False
+
+ def select_wrong_answer(self, correct: str, options: List[str]) -> str:
+ # Select uniformly at random from all options (including correct)
+ return random.choice(options)
+
+ def get_accuracy(self) -> float:
+ # Accuracy depends on number of options, estimate ~0.33 for 3 options
+ return 0.33
+
+
+class AdversarialCompetence(CompetenceProfile):
+ """Intentionally selects wrong answers.
+
+ Use this for testing quality control systems that should
+ detect and flag malicious annotators.
+ """
+
+ def should_be_correct(self) -> bool:
+ return False
+
+ def select_wrong_answer(self, correct: str, options: List[str]) -> str:
+ # Specifically avoid the correct answer
+ wrong_options = [o for o in options if o != correct]
+ if wrong_options:
+ return random.choice(wrong_options)
+ # If only one option, have to return it
+ return options[0] if options else correct
+
+ def get_accuracy(self) -> float:
+ return 0.0
+
+
+def create_competence_profile(level: CompetenceLevel) -> CompetenceProfile:
+ """Factory function to create competence profiles.
+
+ Args:
+ level: The competence level enum value
+
+ Returns:
+ A CompetenceProfile instance for the specified level
+ """
+ profiles = {
+ CompetenceLevel.PERFECT: PerfectCompetence,
+ CompetenceLevel.GOOD: GoodCompetence,
+ CompetenceLevel.AVERAGE: AverageCompetence,
+ CompetenceLevel.POOR: PoorCompetence,
+ CompetenceLevel.RANDOM: RandomCompetence,
+ CompetenceLevel.ADVERSARIAL: AdversarialCompetence,
+ }
+
+ profile_class = profiles.get(level, AverageCompetence)
+ return profile_class()
+
+
+def create_competence_profile_from_string(level_str: str) -> CompetenceProfile:
+ """Create competence profile from string name.
+
+ Args:
+ level_str: String name of competence level (e.g., "good", "average")
+
+ Returns:
+ A CompetenceProfile instance
+ """
+ try:
+ level = CompetenceLevel(level_str.lower())
+ except ValueError:
+ level = CompetenceLevel.AVERAGE
+
+ return create_competence_profile(level)
diff --git a/potato/simulator/config.py b/potato/simulator/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..ddd0885591873344f6237017beaccebfd2046cec
--- /dev/null
+++ b/potato/simulator/config.py
@@ -0,0 +1,577 @@
+"""
+Configuration classes for the user simulator.
+
+This module defines all configuration dataclasses used to configure
+the simulator behavior, including user competence, timing, and strategies.
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Any, Literal, Union
+from enum import Enum
+import os
+import yaml
+
+
+class CompetenceLevel(Enum):
+ """Competence levels for simulated users.
+
+ Each level defines a range of accuracy for the simulated annotator:
+ - PERFECT: 100% accuracy (always matches gold standard)
+ - GOOD: 80-90% accuracy
+ - AVERAGE: 60-70% accuracy
+ - POOR: 40-50% accuracy
+ - RANDOM: Random selection (~1/N accuracy for N labels)
+ - ADVERSARIAL: Intentionally wrong (avoids gold standard)
+ """
+
+ PERFECT = "perfect"
+ GOOD = "good"
+ AVERAGE = "average"
+ POOR = "poor"
+ RANDOM = "random"
+ ADVERSARIAL = "adversarial"
+
+
+class AnnotationStrategyType(Enum):
+ """Annotation generation strategies.
+
+ - RANDOM: Uniform random selection from available labels
+ - BIASED: Weighted random selection based on label preferences
+ - LLM: Use an LLM to generate annotations based on text content
+ - PATTERN: Consistent per-user patterns for testing specific behaviors
+ - GOLD_STANDARD: Use gold answer when available, random otherwise
+ - AGENT: Vision-capable LLM that reads structured / multi-modal instance
+ content (dialogue traces, spreadsheets, image fields) and emits a
+ single batched annotation covering every schema for the instance.
+ """
+
+ RANDOM = "random"
+ BIASED = "biased"
+ LLM = "llm"
+ PATTERN = "pattern"
+ GOLD_STANDARD = "gold_standard"
+ AGENT = "agent"
+
+
+@dataclass
+class TimingConfig:
+ """Configuration for annotation timing behavior.
+
+ Attributes:
+ annotation_time_min: Minimum time per annotation in seconds
+ annotation_time_max: Maximum time per annotation in seconds
+ annotation_time_mean: Mean time for normal distribution
+ annotation_time_std: Standard deviation for normal distribution
+ distribution: Timing distribution model (uniform, normal, exponential)
+ fast_response_threshold: Threshold for flagging suspiciously fast responses
+ session_duration_max: Maximum session duration in minutes (optional)
+ """
+
+ annotation_time_min: float = 2.0
+ annotation_time_max: float = 30.0
+ annotation_time_mean: float = 10.0
+ annotation_time_std: float = 5.0
+ distribution: Literal["uniform", "normal", "exponential"] = "normal"
+ fast_response_threshold: float = 1.0
+ session_duration_max: Optional[float] = None
+
+
+@dataclass
+class LLMStrategyConfig:
+ """Configuration for LLM-based annotation strategy.
+
+ Uses the existing potato.ai endpoint infrastructure.
+
+ Attributes:
+ endpoint_type: LLM provider (openai, anthropic, ollama, etc.)
+ model: Model name/identifier
+ api_key: API key for cloud providers (can use env var reference)
+ base_url: Base URL for local providers like Ollama
+ temperature: Temperature for generation (0-2)
+ max_tokens: Maximum tokens in response
+ add_noise: Whether to occasionally add noise to LLM outputs
+ noise_rate: Probability of adding noise (0-1)
+ """
+
+ endpoint_type: str = "openai"
+ model: Optional[str] = None # Uses provider default if None
+ api_key: Optional[str] = None
+ base_url: Optional[str] = None
+ temperature: float = 0.1
+ max_tokens: int = 100
+ add_noise: bool = True
+ noise_rate: float = 0.05
+
+
+@dataclass
+class InteractiveConfig:
+ """Configuration for driving live ``interactive_chat`` sessions.
+
+ When enabled, the simulator runs a multi-turn chat against the server's
+ ``/agent_chat/*`` routes before annotating each instance whose display
+ contains an ``interactive_chat`` field. The simulator plays the user
+ role; the server-side ``agent_proxy`` plays the agent (echo, OpenAI,
+ HTTP, etc. -- whatever the annotation config specifies).
+
+ Attributes:
+ enabled: Whether to attempt an interactive session per instance.
+ endpoint_type: AI endpoint used to generate the user persona's
+ messages. Defaults to ``ollama`` (text only -- the persona
+ usually doesn't need vision).
+ model: Persona model name. Defaults to provider default.
+ api_key: Optional API key (env-var refs supported).
+ base_url: Optional endpoint base URL.
+ temperature: Sampling temperature for persona messages.
+ max_tokens: Per-message token cap.
+ max_turns: Hard upper bound on turn count per session.
+ persona_system_prompt: System prompt that defines the user persona.
+ Should encourage natural multi-turn behavior and a clear
+ ``DONE`` signal when the task is complete.
+ done_marker: Substring (case-insensitive) the persona emits when
+ it considers the task complete. The runner finishes the
+ session immediately when seen.
+ first_message_template: Template applied to the persona's first
+ message. ``{task}`` is replaced with the task description. If
+ None, the persona generates the first message from scratch.
+ """
+
+ enabled: bool = False
+ endpoint_type: str = "ollama"
+ model: Optional[str] = None
+ api_key: Optional[str] = None
+ base_url: Optional[str] = None
+ temperature: float = 0.7
+ max_tokens: int = 200
+ max_turns: int = 6
+ persona_system_prompt: str = (
+ "You are a curious end-user testing an AI assistant. "
+ "Send concise, natural messages that drive the assistant to "
+ "complete the task. When the assistant has fully completed the "
+ "task, respond with a short acknowledgement and the literal "
+ "marker [DONE]."
+ )
+ done_marker: str = "[DONE]"
+ first_message_template: Optional[str] = (
+ "Please help me with this task: {task}"
+ )
+
+
+@dataclass
+class AgentStrategyConfig:
+ """Configuration for the agent (vision-LLM) annotation strategy.
+
+ Drives a vision-capable LLM that consumes structured / multi-modal
+ instance content (dialogue arrays, spreadsheets, image fields) and
+ produces a batched annotation over every schema for the instance.
+
+ Attributes:
+ endpoint_type: AI endpoint (default ``ollama_vision``). Any vision
+ endpoint registered with ``AIEndpointFactory`` works
+ (``anthropic_vision``, ``openai_vision``, etc.).
+ model: Model identifier (e.g. ``gemma3:4b``, ``llava:latest``,
+ ``llama3.2-vision``). Defaults to provider default.
+ api_key: Cloud-provider API key (env-var refs supported, e.g.
+ ``${ANTHROPIC_API_KEY}``).
+ base_url: Custom endpoint URL (Ollama: ``http://localhost:11434``).
+ temperature: Sampling temperature.
+ max_tokens: Cap on response tokens.
+ max_image_dim: Resize images so the longest edge is at most this
+ many pixels before sending. ``None`` keeps the original.
+ max_image_count: Skip image attachment past this many images per
+ instance (some models cap at 1โ4).
+ include_dialogue_text: Render dialogue arrays as
+ ``: `` lines in the prompt.
+ include_spreadsheet: Render spreadsheet/table fields as plain text.
+ max_dialogue_chars: Truncate long dialogue payloads to this many
+ characters in the prompt to fit the model's context window.
+ cache_per_instance: When True (default), one LLM call per instance
+ answers all schemas; subsequent ``generate_annotation`` calls
+ for the same instance return cached results.
+ add_noise: Probability of falling back to a random annotation per
+ schema (mirrors ``LLMStrategyConfig`` so existing competence
+ modeling still applies).
+ noise_rate: Probability used for noise injection (0โ1).
+ """
+
+ endpoint_type: str = "ollama_vision"
+ model: Optional[str] = None
+ api_key: Optional[str] = None
+ base_url: Optional[str] = None
+ temperature: float = 0.1
+ max_tokens: int = 800
+ max_image_dim: Optional[int] = 1024
+ max_image_count: int = 4
+ include_dialogue_text: bool = True
+ include_spreadsheet: bool = True
+ max_dialogue_chars: int = 12000
+ cache_per_instance: bool = True
+ add_noise: bool = False
+ noise_rate: float = 0.0
+
+
+@dataclass
+class BiasedStrategyConfig:
+ """Configuration for biased annotation strategy.
+
+ Attributes:
+ label_weights: Dictionary mapping label names to selection weights.
+ Higher weights mean higher probability of selection.
+ Example: {"positive": 0.6, "negative": 0.3, "neutral": 0.1}
+ """
+
+ label_weights: Dict[str, float] = field(default_factory=dict)
+
+
+@dataclass
+class PatternStrategyConfig:
+ """Configuration for pattern-based annotation strategy.
+
+ Allows defining specific behavior patterns per user.
+
+ Attributes:
+ patterns: Dictionary mapping user_id to behavior configuration.
+ Each pattern can specify:
+ - preferred_label: Label this user tends to select
+ - bias_strength: How strongly they prefer it (0-1)
+ - keywords: Text patterns that trigger specific labels
+ """
+
+ patterns: Dict[str, Dict[str, Any]] = field(default_factory=dict)
+
+
+@dataclass
+class UserConfig:
+ """Configuration for a single simulated user.
+
+ Attributes:
+ user_id: Unique identifier for this user
+ competence: Competence level determining accuracy
+ strategy: Annotation strategy type
+ timing: Timing configuration for this user
+ llm_config: LLM configuration if strategy is LLM
+ biased_config: Bias configuration if strategy is BIASED
+ pattern_config: Pattern configuration if strategy is PATTERN
+ max_annotations: Maximum annotations for this user (optional)
+ """
+
+ user_id: str
+ competence: CompetenceLevel = CompetenceLevel.AVERAGE
+ strategy: AnnotationStrategyType = AnnotationStrategyType.RANDOM
+ timing: TimingConfig = field(default_factory=TimingConfig)
+ llm_config: Optional[LLMStrategyConfig] = None
+ biased_config: Optional[BiasedStrategyConfig] = None
+ pattern_config: Optional[PatternStrategyConfig] = None
+ agent_config: Optional[AgentStrategyConfig] = None
+ max_annotations: Optional[int] = None
+
+
+@dataclass
+class SimulatorConfig:
+ """Master configuration for the user simulator.
+
+ Attributes:
+ user_count: Number of simulated users to create
+ competence_distribution: Distribution of competence levels
+ (keys are competence level names, values are proportions)
+ users: Explicit list of user configurations (overrides user_count)
+ timing: Global timing configuration (can be overridden per-user)
+ strategy: Default annotation strategy
+ llm_config: LLM configuration for LLM strategy
+ biased_config: Bias configuration for biased strategy
+ gold_standard_file: Path to JSON file with gold standard labels
+ parallel_users: Maximum concurrent users
+ delay_between_users: Delay between starting users (seconds)
+ attention_check_fail_rate: Rate at which users fail attention checks
+ respond_fast_rate: Rate of suspiciously fast responses
+ simulate_wait: Whether to actually wait between annotations
+ output_dir: Directory for output files
+ export_format: Output format (json, csv, jsonl)
+ """
+
+ # User configuration
+ user_count: int = 10
+ competence_distribution: Dict[str, float] = field(
+ default_factory=lambda: {"good": 0.5, "average": 0.3, "poor": 0.2}
+ )
+ users: List[UserConfig] = field(default_factory=list)
+
+ # Global timing configuration
+ timing: TimingConfig = field(default_factory=TimingConfig)
+
+ # Strategy configuration - default to random
+ strategy: AnnotationStrategyType = AnnotationStrategyType.RANDOM
+ llm_config: Optional[LLMStrategyConfig] = None
+ biased_config: Optional[BiasedStrategyConfig] = None
+ agent_config: Optional[AgentStrategyConfig] = None
+ interactive: Optional[InteractiveConfig] = None
+
+ # Gold standard data for competence-based accuracy
+ gold_standard_file: Optional[str] = None
+
+ # Execution configuration
+ parallel_users: int = 5
+ delay_between_users: float = 0.5
+
+ # Quality control testing options
+ attention_check_fail_rate: float = 0.0
+ respond_fast_rate: float = 0.0
+
+ # Whether to actually wait (set False for fast testing)
+ simulate_wait: bool = False
+
+ # Output configuration
+ output_dir: str = "simulator_output"
+ export_format: Literal["json", "csv", "jsonl"] = "json"
+
+ @classmethod
+ def from_yaml(cls, yaml_path: str) -> "SimulatorConfig":
+ """Load configuration from YAML file.
+
+ Args:
+ yaml_path: Path to YAML configuration file
+
+ Returns:
+ SimulatorConfig instance
+ """
+ with open(yaml_path, "r") as f:
+ data = yaml.safe_load(f)
+
+ return cls._parse_config(data)
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "SimulatorConfig":
+ """Load configuration from dictionary.
+
+ Args:
+ data: Configuration dictionary
+
+ Returns:
+ SimulatorConfig instance
+ """
+ return cls._parse_config(data)
+
+ @classmethod
+ def _parse_config(cls, data: Dict[str, Any]) -> "SimulatorConfig":
+ """Parse configuration from dictionary.
+
+ Args:
+ data: Raw configuration dictionary
+
+ Returns:
+ SimulatorConfig instance
+ """
+ # Handle nested 'simulator' key if present
+ if "simulator" in data:
+ data = data["simulator"]
+
+ # Parse timing config
+ timing = TimingConfig()
+ if "timing" in data:
+ timing_data = data["timing"]
+ if "annotation_time" in timing_data:
+ at = timing_data["annotation_time"]
+ timing = TimingConfig(
+ annotation_time_min=at.get("min", 2.0),
+ annotation_time_max=at.get("max", 30.0),
+ annotation_time_mean=at.get("mean", 10.0),
+ annotation_time_std=at.get("std", 5.0),
+ distribution=at.get("distribution", "normal"),
+ fast_response_threshold=timing_data.get(
+ "fast_response_threshold", 1.0
+ ),
+ session_duration_max=timing_data.get("session_duration_max"),
+ )
+ else:
+ timing = TimingConfig(
+ annotation_time_min=timing_data.get("annotation_time_min", 2.0),
+ annotation_time_max=timing_data.get("annotation_time_max", 30.0),
+ annotation_time_mean=timing_data.get("annotation_time_mean", 10.0),
+ annotation_time_std=timing_data.get("annotation_time_std", 5.0),
+ distribution=timing_data.get("distribution", "normal"),
+ fast_response_threshold=timing_data.get(
+ "fast_response_threshold", 1.0
+ ),
+ session_duration_max=timing_data.get("session_duration_max"),
+ )
+
+ # Parse LLM config
+ llm_config = None
+ if "llm_config" in data:
+ llm_data = data["llm_config"]
+ # Handle environment variable references
+ api_key = llm_data.get("api_key")
+ if api_key and api_key.startswith("${") and api_key.endswith("}"):
+ env_var = api_key[2:-1]
+ api_key = os.environ.get(env_var)
+
+ llm_config = LLMStrategyConfig(
+ endpoint_type=llm_data.get("endpoint_type", "openai"),
+ model=llm_data.get("model"),
+ api_key=api_key,
+ base_url=llm_data.get("base_url"),
+ temperature=llm_data.get("temperature", 0.1),
+ max_tokens=llm_data.get("max_tokens", 100),
+ add_noise=llm_data.get("add_noise", True),
+ noise_rate=llm_data.get("noise_rate", 0.05),
+ )
+
+ # Parse biased config
+ biased_config = None
+ if "biased_config" in data:
+ biased_config = BiasedStrategyConfig(
+ label_weights=data["biased_config"].get("label_weights", {})
+ )
+
+ # Parse interactive (chat-driving) config
+ interactive_config = None
+ if "interactive" in data:
+ ic = data["interactive"]
+ api_key = ic.get("api_key")
+ if api_key and api_key.startswith("${") and api_key.endswith("}"):
+ api_key = os.environ.get(api_key[2:-1])
+ kwargs = {
+ "enabled": ic.get("enabled", True),
+ "endpoint_type": ic.get("endpoint_type", "ollama"),
+ "model": ic.get("model"),
+ "api_key": api_key,
+ "base_url": ic.get("base_url"),
+ "temperature": ic.get("temperature", 0.7),
+ "max_tokens": ic.get("max_tokens", 200),
+ "max_turns": ic.get("max_turns", 6),
+ "done_marker": ic.get("done_marker", "[DONE]"),
+ }
+ if "persona_system_prompt" in ic:
+ kwargs["persona_system_prompt"] = ic["persona_system_prompt"]
+ if "first_message_template" in ic:
+ kwargs["first_message_template"] = ic["first_message_template"]
+ interactive_config = InteractiveConfig(**kwargs)
+
+ # Parse agent (vision-LLM) config
+ agent_config = None
+ if "agent_config" in data:
+ ad = data["agent_config"]
+ api_key = ad.get("api_key")
+ if api_key and api_key.startswith("${") and api_key.endswith("}"):
+ api_key = os.environ.get(api_key[2:-1])
+ agent_config = AgentStrategyConfig(
+ endpoint_type=ad.get("endpoint_type", "ollama_vision"),
+ model=ad.get("model"),
+ api_key=api_key,
+ base_url=ad.get("base_url"),
+ temperature=ad.get("temperature", 0.1),
+ max_tokens=ad.get("max_tokens", 800),
+ max_image_dim=ad.get("max_image_dim", 1024),
+ max_image_count=ad.get("max_image_count", 4),
+ include_dialogue_text=ad.get("include_dialogue_text", True),
+ include_spreadsheet=ad.get("include_spreadsheet", True),
+ max_dialogue_chars=ad.get("max_dialogue_chars", 12000),
+ cache_per_instance=ad.get("cache_per_instance", True),
+ add_noise=ad.get("add_noise", False),
+ noise_rate=ad.get("noise_rate", 0.0),
+ )
+
+ # Parse strategy
+ strategy_str = data.get("strategy", "random")
+ try:
+ strategy = AnnotationStrategyType(strategy_str)
+ except ValueError:
+ strategy = AnnotationStrategyType.RANDOM
+
+ # Parse users section
+ users_data = data.get("users", {})
+ user_count = users_data.get("count", data.get("user_count", 10))
+ competence_dist = users_data.get(
+ "competence_distribution",
+ data.get(
+ "competence_distribution", {"good": 0.5, "average": 0.3, "poor": 0.2}
+ ),
+ )
+
+ # Parse execution config
+ execution = data.get("execution", {})
+ parallel_users = execution.get("parallel_users", data.get("parallel_users", 5))
+ delay_between = execution.get(
+ "delay_between_users", data.get("delay_between_users", 0.5)
+ )
+ max_annotations = execution.get("max_annotations_per_user")
+
+ # Parse QC config
+ qc_config = data.get("quality_control", {})
+ attention_fail_rate = qc_config.get(
+ "attention_check_fail_rate", data.get("attention_check_fail_rate", 0.0)
+ )
+ respond_fast_rate = qc_config.get(
+ "respond_fast_rate", data.get("respond_fast_rate", 0.0)
+ )
+
+ # Parse output config
+ output_config = data.get("output", {})
+ output_dir = output_config.get("dir", data.get("output_dir", "simulator_output"))
+ export_format = output_config.get(
+ "format", data.get("export_format", "json")
+ )
+
+ return cls(
+ user_count=user_count,
+ competence_distribution=competence_dist,
+ timing=timing,
+ strategy=strategy,
+ llm_config=llm_config,
+ biased_config=biased_config,
+ agent_config=agent_config,
+ interactive=interactive_config,
+ gold_standard_file=data.get("gold_standard_file"),
+ parallel_users=parallel_users,
+ delay_between_users=delay_between,
+ attention_check_fail_rate=attention_fail_rate,
+ respond_fast_rate=respond_fast_rate,
+ simulate_wait=data.get("simulate_wait", False),
+ output_dir=output_dir,
+ export_format=export_format,
+ )
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert configuration to dictionary.
+
+ Returns:
+ Configuration as dictionary
+ """
+ return {
+ "user_count": self.user_count,
+ "competence_distribution": self.competence_distribution,
+ "timing": {
+ "annotation_time_min": self.timing.annotation_time_min,
+ "annotation_time_max": self.timing.annotation_time_max,
+ "annotation_time_mean": self.timing.annotation_time_mean,
+ "annotation_time_std": self.timing.annotation_time_std,
+ "distribution": self.timing.distribution,
+ "fast_response_threshold": self.timing.fast_response_threshold,
+ "session_duration_max": self.timing.session_duration_max,
+ },
+ "strategy": self.strategy.value,
+ "llm_config": (
+ {
+ "endpoint_type": self.llm_config.endpoint_type,
+ "model": self.llm_config.model,
+ "temperature": self.llm_config.temperature,
+ "max_tokens": self.llm_config.max_tokens,
+ "add_noise": self.llm_config.add_noise,
+ "noise_rate": self.llm_config.noise_rate,
+ }
+ if self.llm_config
+ else None
+ ),
+ "biased_config": (
+ {"label_weights": self.biased_config.label_weights}
+ if self.biased_config
+ else None
+ ),
+ "gold_standard_file": self.gold_standard_file,
+ "parallel_users": self.parallel_users,
+ "delay_between_users": self.delay_between_users,
+ "attention_check_fail_rate": self.attention_check_fail_rate,
+ "respond_fast_rate": self.respond_fast_rate,
+ "simulate_wait": self.simulate_wait,
+ "output_dir": self.output_dir,
+ "export_format": self.export_format,
+ }
diff --git a/potato/simulator/interactive_runner.py b/potato/simulator/interactive_runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..2662e08c2ea659f17620fbc46d6c0ccd7834f1d9
--- /dev/null
+++ b/potato/simulator/interactive_runner.py
@@ -0,0 +1,236 @@
+"""
+Interactive chat session driver for the simulator.
+
+When the annotation server's ``instance_display`` includes an
+``interactive_chat`` field, the annotator is expected to chat with a live
+agent backend before submitting trajectory ratings.
+:class:`InteractiveSessionRunner` plays the user side of that chat: it asks
+a small "persona" LLM to generate user messages, posts them to the server's
+``/agent_chat/send`` route, and finishes with ``/agent_chat/finish`` so the
+captured conversation is written into the instance data.
+
+After the runner returns, the regular annotation pipeline (e.g.
+:class:`AgentSimulatorStrategy`) picks up the freshly populated
+``conversation`` field and produces ratings.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import requests
+
+from .config import InteractiveConfig
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class InteractiveSessionResult:
+ """Outcome of one interactive_chat run."""
+
+ instance_id: str
+ completed: bool
+ turns: int
+ conversation: List[Dict[str, Any]]
+ error: Optional[str] = None
+
+
+class InteractiveSessionRunner:
+ """Drive a multi-turn ``interactive_chat`` against the server.
+
+ The runner is stateless across instances: ``run`` is called once per
+ instance and returns the resulting conversation. The persona LLM is
+ initialized lazily on first use so importing the module is cheap.
+ """
+
+ def __init__(self, config: InteractiveConfig, server_url: str):
+ self.config = config
+ self.server_url = server_url.rstrip("/")
+ self._endpoint = None # lazy
+
+ # ------------------------------------------------------------------
+ # Persona endpoint setup
+ # ------------------------------------------------------------------
+
+ def _get_endpoint(self):
+ if self._endpoint is not None:
+ return self._endpoint
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ ai_cfg: Dict[str, Any] = {
+ "model": self.config.model,
+ "api_key": self.config.api_key,
+ "max_tokens": self.config.max_tokens,
+ "temperature": self.config.temperature,
+ }
+ if self.config.base_url:
+ ai_cfg["base_url"] = self.config.base_url
+
+ self._endpoint = AIEndpointFactory.create_endpoint({
+ "ai_support": {
+ "enabled": True,
+ "endpoint_type": self.config.endpoint_type,
+ "ai_config": ai_cfg,
+ }
+ })
+ except Exception as e:
+ logger.warning("InteractiveSessionRunner: persona endpoint init failed: %s", e)
+ self._endpoint = None
+ return self._endpoint
+
+ # ------------------------------------------------------------------
+ # Persona messaging
+ # ------------------------------------------------------------------
+
+ def _generate_persona_message(
+ self,
+ task: str,
+ history: List[Dict[str, str]],
+ ) -> Optional[str]:
+ endpoint = self._get_endpoint()
+ if endpoint is None:
+ return None
+
+ if not history and self.config.first_message_template:
+ return self.config.first_message_template.format(task=task)
+
+ # Build chat history. The persona's "user" role is what we send to
+ # the agent, so from the persona LLM's perspective those are
+ # *assistant* messages and the agent's replies are *user* prompts.
+ messages: List[Dict[str, str]] = [
+ {"role": "system", "content": self.config.persona_system_prompt
+ + f"\n\nThe task you want completed is:\n{task}"}
+ ]
+ for msg in history:
+ if msg["role"] == "user": # what the persona previously sent
+ messages.append({"role": "assistant", "content": msg["content"]})
+ else: # agent's reply
+ messages.append({"role": "user", "content": msg["content"]})
+
+ try:
+ if hasattr(endpoint, "chat_query"):
+ reply = endpoint.chat_query(messages)
+ else:
+ # Fall back to flattening into a single prompt
+ flat = "\n".join(f'{m["role"]}: {m["content"]}' for m in messages)
+ reply = endpoint.query(flat + "\nassistant:", None)
+ except Exception as e:
+ logger.warning("Persona LLM call failed: %s", e)
+ return None
+
+ if isinstance(reply, dict):
+ reply = reply.get("response") or reply.get("content") or str(reply)
+ text = str(reply or "").strip()
+ return text or None
+
+ # ------------------------------------------------------------------
+ # Server interaction
+ # ------------------------------------------------------------------
+
+ def _send_to_agent(
+ self, session: requests.Session, message: str
+ ) -> Optional[Dict[str, Any]]:
+ try:
+ resp = session.post(
+ f"{self.server_url}/agent_chat/send",
+ json={"message": message},
+ timeout=120,
+ )
+ except requests.exceptions.RequestException as e:
+ logger.warning("agent_chat/send request failed: %s", e)
+ return None
+ if resp.status_code != 200:
+ logger.warning(
+ "agent_chat/send returned %d: %s", resp.status_code, resp.text[:200]
+ )
+ return None
+ try:
+ return resp.json()
+ except ValueError:
+ return None
+
+ def _finish(self, session: requests.Session) -> bool:
+ try:
+ resp = session.post(
+ f"{self.server_url}/agent_chat/finish",
+ timeout=60,
+ )
+ except requests.exceptions.RequestException as e:
+ logger.warning("agent_chat/finish request failed: %s", e)
+ return False
+ if resp.status_code != 200:
+ logger.warning(
+ "agent_chat/finish returned %d: %s",
+ resp.status_code, resp.text[:200],
+ )
+ return False
+ return True
+
+ # ------------------------------------------------------------------
+ # Main entry point
+ # ------------------------------------------------------------------
+
+ def run(
+ self,
+ session: requests.Session,
+ instance_id: str,
+ task_description: str,
+ ) -> InteractiveSessionResult:
+ """Drive one chat session end-to-end."""
+ history: List[Dict[str, str]] = []
+ completed = False
+ error: Optional[str] = None
+
+ for turn in range(self.config.max_turns):
+ user_msg = self._generate_persona_message(task_description, history)
+ if not user_msg:
+ error = error or "persona produced no message"
+ break
+
+ # Strip the [DONE] marker before sending so the agent doesn't see
+ # it; remember that we should finish after this turn.
+ should_finish = self.config.done_marker.lower() in user_msg.lower()
+ send_msg = user_msg.replace(self.config.done_marker, "").strip() or "Thanks!"
+
+ agent_reply = self._send_to_agent(session, send_msg)
+ if agent_reply is None:
+ error = error or "agent send failed"
+ break
+ history.append({"role": "user", "content": send_msg})
+ history.append({
+ "role": "agent",
+ "content": agent_reply.get("content", ""),
+ })
+
+ if should_finish:
+ completed = True
+ break
+
+ ok = self._finish(session)
+ if not ok and not error:
+ error = "finish failed"
+ if ok and not completed:
+ # We hit max_turns without an explicit DONE; still consider the
+ # session "done" for accounting purposes.
+ completed = True
+
+ # Build the conversation array the way the server's finish route does
+ conversation = [
+ {
+ "speaker": "User" if msg["role"] == "user" else "Agent",
+ "text": msg["content"],
+ }
+ for msg in history
+ ]
+
+ return InteractiveSessionResult(
+ instance_id=instance_id,
+ completed=completed,
+ turns=len(history) // 2,
+ conversation=conversation,
+ error=error,
+ )
diff --git a/potato/simulator/playwright_simulator.py b/potato/simulator/playwright_simulator.py
new file mode 100644
index 0000000000000000000000000000000000000000..ddba8ecbd994896fdfa8b5fcf8f7e725785f9127
--- /dev/null
+++ b/potato/simulator/playwright_simulator.py
@@ -0,0 +1,502 @@
+"""
+Playwright-driven simulated user.
+
+Drives the actual browser UI: registers + logs in via the login form,
+navigates to ``/annotate``, applies annotations through real DOM events
+(clicking radios, filling textareas, ticking checkboxes), then clicks
+the Next button to advance.
+
+Designed as a slow-but-real smoke test of the rendered annotation UI:
+the same code paths a human user would hit. The annotation values still
+come from an :class:`AnnotationStrategy` (typically
+:class:`AgentSimulatorStrategy`); Playwright is only the *driver*.
+
+Constraints:
+ - Single user, sequential. Browser sessions don't parallelize cleanly.
+ - Requires ``playwright`` + a chromium install
+ (``pip install playwright && playwright install chromium``).
+ - Falls back gracefully if Playwright is unavailable: raises a clear
+ ``ImportError`` at construction time.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import requests
+
+from .annotation_strategies import AnnotationStrategy, create_strategy
+from .competence_profiles import create_competence_profile
+from .config import (
+ AnnotationStrategyType,
+ InteractiveConfig,
+ UserConfig,
+)
+from .interactive_runner import InteractiveSessionRunner
+from .timing_models import NoWaitTimingModel, TimingModel
+from .user_simulator import AnnotationRecord, UserSimulationResult
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _DomSubmissionResult:
+ """How many DOM inputs we successfully applied vs total annotations."""
+
+ applied: int
+ total: int
+
+
+class PlaywrightSimulatedUser:
+ """Browser-driven simulated user.
+
+ Public surface mirrors :class:`SimulatedUser` enough that a smoke runner
+ can use either. ``run_simulation()`` returns a
+ :class:`UserSimulationResult` so existing reporting works.
+ """
+
+ def __init__(
+ self,
+ user_config: UserConfig,
+ server_url: str,
+ gold_standards: Optional[Dict[str, Dict[str, Any]]] = None,
+ simulate_wait: bool = False,
+ interactive_config: Optional[InteractiveConfig] = None,
+ headless: bool = True,
+ debounce_seconds: float = 1.8,
+ ):
+ try:
+ from playwright.sync_api import sync_playwright # noqa: F401
+ except ImportError as e:
+ raise ImportError(
+ "PlaywrightSimulatedUser requires playwright. Install it with:\n"
+ " pip install playwright\n"
+ " playwright install chromium"
+ ) from e
+
+ self.config = user_config
+ self.server_url = server_url.rstrip("/")
+ self.gold_standards = gold_standards or {}
+ self.headless = headless
+ self.debounce_seconds = debounce_seconds
+
+ self.competence = create_competence_profile(user_config.competence)
+ self.strategy = self._create_strategy()
+
+ self.timing = (
+ TimingModel(user_config.timing)
+ if simulate_wait
+ else NoWaitTimingModel(user_config.timing)
+ )
+
+ # API session is mirrored to/from Playwright for /api/* fetches
+ self.api_session = requests.Session()
+
+ self.interactive_runner: Optional[InteractiveSessionRunner] = None
+ if interactive_config and interactive_config.enabled:
+ self.interactive_runner = InteractiveSessionRunner(
+ interactive_config, server_url
+ )
+
+ self.result = UserSimulationResult(user_id=user_config.user_id)
+ self.schemas: List[Dict[str, Any]] = []
+
+ # Lazily set in start()
+ self._pw = None
+ self._browser = None
+ self._context = None
+ self._page = None
+
+ # ------------------------------------------------------------------
+ # Strategy & lifecycle
+ # ------------------------------------------------------------------
+
+ def _create_strategy(self) -> AnnotationStrategy:
+ return create_strategy(
+ strategy_type=self.config.strategy,
+ llm_config=self.config.llm_config,
+ biased_config=self.config.biased_config,
+ pattern_config=self.config.pattern_config,
+ agent_config=self.config.agent_config,
+ user_id=self.config.user_id,
+ )
+
+ def start(self) -> None:
+ from playwright.sync_api import sync_playwright
+
+ self._pw = sync_playwright().start()
+ self._browser = self._pw.chromium.launch(headless=self.headless)
+ self._context = self._browser.new_context()
+ self._page = self._context.new_page()
+
+ def stop(self) -> None:
+ for closer in (self._page, self._context, self._browser):
+ if closer is None:
+ continue
+ try:
+ closer.close()
+ except Exception:
+ pass
+ if self._pw is not None:
+ try:
+ self._pw.stop()
+ except Exception:
+ pass
+ self._page = self._context = self._browser = self._pw = None
+
+ # ------------------------------------------------------------------
+ # Cookie sharing
+ # ------------------------------------------------------------------
+
+ def _sync_cookies_to_api(self) -> None:
+ """Copy Playwright cookies into the requests session."""
+ if self._context is None:
+ return
+ for cookie in self._context.cookies():
+ self.api_session.cookies.set(
+ cookie["name"], cookie["value"], domain=cookie.get("domain")
+ )
+
+ # ------------------------------------------------------------------
+ # Login via UI
+ # ------------------------------------------------------------------
+
+ def login_via_ui(self) -> bool:
+ """Register + log in. Uses the API for the credential exchange (the
+ UI form mode is brittle across templates), then loads the cookies
+ into Playwright so DOM interactions work in a logged-in session."""
+ password = "simulated_password_123"
+ try:
+ self.api_session.post(
+ f"{self.server_url}/register",
+ data={"action": "signup", "email": self.config.user_id, "pass": password},
+ allow_redirects=True,
+ timeout=30,
+ )
+ self.api_session.post(
+ f"{self.server_url}/auth",
+ data={"action": "login", "email": self.config.user_id, "pass": password},
+ allow_redirects=True,
+ timeout=30,
+ )
+ except requests.exceptions.RequestException as e:
+ logger.warning("API login failed: %s", e)
+ return False
+
+ # Push cookies into Playwright so the page renders as logged-in.
+ self._sync_cookies_to_browser()
+
+ # Navigate to /annotate; walk past consent/instructions if needed.
+ page = self._page
+ try:
+ page.goto(f"{self.server_url}/annotate", wait_until="domcontentloaded", timeout=15000)
+ except Exception as e:
+ logger.warning("page.goto /annotate failed: %s", e)
+
+ for _ in range(10):
+ self._sync_cookies_to_api()
+ if self._is_in_annotation_phase():
+ # Make sure the annotation page is what's rendered
+ try:
+ page.wait_for_selector("#next-btn", timeout=5000)
+ except Exception:
+ pass
+ return True
+ # Try a UI advance (consent / instructions screen)
+ self._advance_phase_screen()
+ time.sleep(0.5)
+ try:
+ page.goto(f"{self.server_url}/annotate", wait_until="domcontentloaded", timeout=10000)
+ except Exception:
+ pass
+
+ return self._is_in_annotation_phase()
+
+ def _sync_cookies_to_browser(self) -> None:
+ """Push cookies from the API session into the Playwright context."""
+ if self._context is None:
+ return
+ cookies = []
+ for c in self.api_session.cookies:
+ cookies.append({
+ "name": c.name,
+ "value": c.value,
+ "url": self.server_url,
+ })
+ if cookies:
+ try:
+ self._context.add_cookies(cookies)
+ except Exception as e:
+ logger.warning("add_cookies failed: %s", e)
+
+ def _is_in_annotation_phase(self) -> bool:
+ """We're in annotation phase iff /api/current_instance returns 200."""
+ try:
+ r = self.api_session.get(
+ f"{self.server_url}/api/current_instance", timeout=10
+ )
+ except requests.exceptions.RequestException:
+ return False
+ return r.status_code == 200
+
+ def _advance_phase_screen(self) -> None:
+ """Click any visible 'Next' / 'Continue' / 'I agree' button."""
+ page = self._page
+ candidates = [
+ "#next-btn:visible",
+ "button:has-text('Continue')",
+ "button:has-text('I Agree')",
+ "button:has-text('Start')",
+ "input[type='submit']:visible",
+ ]
+ for sel in candidates:
+ loc = page.locator(sel)
+ if loc.count() > 0:
+ try:
+ loc.first.click()
+ return
+ except Exception:
+ continue
+
+ # ------------------------------------------------------------------
+ # API helpers
+ # ------------------------------------------------------------------
+
+ def fetch_schemas(self) -> List[Dict[str, Any]]:
+ try:
+ r = self.api_session.get(f"{self.server_url}/api/schemas", timeout=30)
+ if r.status_code != 200:
+ return []
+ data = r.json()
+ if isinstance(data, dict):
+ schemas = (
+ list(data["schemas"].values())
+ if isinstance(data.get("schemas"), dict)
+ else data.get("schemas", list(data.values()))
+ )
+ else:
+ schemas = data
+ self.schemas = schemas
+ return schemas
+ except requests.exceptions.RequestException as e:
+ logger.warning("schema fetch failed: %s", e)
+ return []
+
+ def fetch_current_instance(self) -> Optional[Dict[str, Any]]:
+ try:
+ r = self.api_session.get(
+ f"{self.server_url}/api/current_instance", timeout=30
+ )
+ except requests.exceptions.RequestException:
+ return None
+ if r.status_code != 200:
+ return None
+ return r.json()
+
+ # ------------------------------------------------------------------
+ # DOM annotation application
+ # ------------------------------------------------------------------
+
+ def _apply_annotations(self, annotations: Dict[str, Any]) -> _DomSubmissionResult:
+ """Translate a wire-format annotation dict into DOM clicks/fills."""
+ page = self._page
+ applied = 0
+ total = 0
+ for key, value in annotations.items():
+ total += 1
+ if ":" not in key:
+ continue
+ schema, label = key.split(":", 1)
+ if label == "text":
+ # textarea / textbox
+ # name: :::text via generate_element_identifier
+ sel = f"[name=\"{schema}:::text\"]"
+ if page.locator(sel).count() == 0:
+ sel = f"textarea[name*='{schema}']"
+ if page.locator(sel).count() > 0:
+ try:
+ page.locator(sel).first.fill(str(value)[:1000])
+ applied += 1
+ except Exception as e:
+ logger.debug("fill failed for %s: %s", sel, e)
+ continue
+
+ # Try radio (name=schema, value=label)
+ radio_sel = f"input[type='radio'][name=\"{schema}\"][value=\"{label}\"]"
+ if page.locator(radio_sel).count() > 0:
+ try:
+ page.locator(radio_sel).first.check(force=True)
+ applied += 1
+ continue
+ except Exception:
+ pass
+
+ # Try checkbox / multiselect (name=:::)
+ cb_sel = f"input[type='checkbox'][name=\"{schema}:::{label}\"]"
+ if page.locator(cb_sel).count() > 0:
+ try:
+ page.locator(cb_sel).first.check(force=True)
+ applied += 1
+ continue
+ except Exception:
+ pass
+
+ # Likert (rendered as radios with name=::: in some templates)
+ alt_radio = f"input[type='radio'][name=\"{schema}:::{label}\"]"
+ if page.locator(alt_radio).count() > 0:
+ try:
+ page.locator(alt_radio).first.check(force=True)
+ applied += 1
+ continue
+ except Exception:
+ pass
+
+ logger.debug("no DOM target for %s=%s", key, value)
+
+ return _DomSubmissionResult(applied=applied, total=total)
+
+ def _click_next(self) -> bool:
+ page = self._page
+ next_btn = page.locator("#next-btn")
+ if next_btn.count() == 0:
+ return False
+ try:
+ next_btn.first.click()
+ return True
+ except Exception as e:
+ logger.warning("next-btn click failed: %s", e)
+ return False
+
+ # ------------------------------------------------------------------
+ # Main loop
+ # ------------------------------------------------------------------
+
+ def run_simulation(
+ self, max_annotations: Optional[int] = None
+ ) -> UserSimulationResult:
+ self.result.start_time = datetime.now()
+ max_ann = (
+ max_annotations
+ if max_annotations is not None
+ else self.config.max_annotations
+ )
+
+ try:
+ self.start()
+ if not self.login_via_ui():
+ self.result.errors.append("login_via_ui failed")
+ return self.result
+
+ self.fetch_schemas()
+ if not self.schemas:
+ self.result.errors.append("no schemas returned by /api/schemas")
+
+ count = 0
+ seen_instances = set()
+ while True:
+ if max_ann is not None and count >= max_ann:
+ break
+
+ instance = self.fetch_current_instance()
+ if not instance or not instance.get("instance_id"):
+ break
+ instance_id = instance["instance_id"]
+ if instance_id in seen_instances:
+ # Server didn't advance after Next click; stop to avoid
+ # an infinite loop.
+ break
+ seen_instances.add(instance_id)
+
+ # Optional interactive chat first
+ if self.interactive_runner is not None:
+ data = instance.get("data") or {}
+ task_text = (
+ data.get("task_description")
+ or data.get("text")
+ or instance.get("text", "")
+ )
+ chat_result = self.interactive_runner.run(
+ self.api_session, instance_id, task_text
+ )
+ if chat_result.error:
+ self.result.errors.append(
+ f"interactive: {chat_result.error}"
+ )
+ refreshed = self.fetch_current_instance()
+ if refreshed and refreshed.get("instance_id") == instance_id:
+ instance = refreshed
+
+ response_time = self.timing.get_response_time(0.0)
+ self.timing.wait(response_time)
+
+ # Generate annotations using the strategy
+ instance_for_strategy = dict(instance)
+ instance_for_strategy["__all_schemas__"] = self.schemas
+
+ gold_answer = self.gold_standards.get(instance_id)
+
+ all_annotations: Dict[str, Any] = {}
+ for schema in self.schemas:
+ schema_name = schema.get("name")
+ schema_gold = (
+ {schema_name: gold_answer.get(schema_name)}
+ if gold_answer
+ else None
+ )
+ ann = self.strategy.generate_annotation(
+ instance_for_strategy,
+ schema,
+ self.competence,
+ schema_gold,
+ )
+ all_annotations.update(ann)
+
+ # Apply via DOM
+ dom_result = self._apply_annotations(all_annotations)
+
+ # Wait for the debounced auto-save to fire
+ time.sleep(self.debounce_seconds)
+
+ # Reload Playwright cookies into the API session in case the
+ # server set new ones during the save.
+ self._sync_cookies_to_api()
+
+ self.result.annotations.append(
+ AnnotationRecord(
+ instance_id=instance_id,
+ schema_name=",".join(all_annotations.keys()),
+ annotation=all_annotations,
+ response_time=response_time,
+ timestamp=datetime.now(),
+ )
+ )
+ count += 1
+ logger.info(
+ "[playwright] %s annotated %s (DOM %d/%d)",
+ self.config.user_id,
+ instance_id,
+ dom_result.applied,
+ dom_result.total,
+ )
+
+ if not self._click_next():
+ break
+ # Wait for the next instance to render
+ self._page.wait_for_load_state("networkidle", timeout=15000)
+ time.sleep(0.3)
+
+ except Exception as e:
+ logger.error("Playwright simulation error: %s", e, exc_info=True)
+ self.result.errors.append(f"playwright: {e}")
+ finally:
+ self.result.end_time = datetime.now()
+ self.result.total_time = (
+ self.result.end_time - self.result.start_time
+ ).total_seconds()
+ self.stop()
+
+ return self.result
diff --git a/potato/simulator/reporting.py b/potato/simulator/reporting.py
new file mode 100644
index 0000000000000000000000000000000000000000..e93cdb7c61284c1d16610310973baaba4229c551
--- /dev/null
+++ b/potato/simulator/reporting.py
@@ -0,0 +1,192 @@
+"""
+Reporting and export functionality for simulation results.
+
+This module provides the SimulationReporter class for exporting
+simulation results in various formats.
+"""
+
+import json
+import csv
+import os
+from typing import Dict, Any
+from datetime import datetime
+
+from .user_simulator import UserSimulationResult
+
+
+class SimulationReporter:
+ """Handles result collection and export.
+
+ Supports multiple export formats:
+ - JSON: Full structured results
+ - CSV: Flat annotation records
+ - JSONL: Line-delimited JSON for streaming
+ """
+
+ def __init__(self, output_dir: str):
+ """Initialize reporter.
+
+ Args:
+ output_dir: Directory for output files
+ """
+ self.output_dir = output_dir
+ os.makedirs(output_dir, exist_ok=True)
+
+ def export_results(
+ self,
+ results: Dict[str, UserSimulationResult],
+ summary: Dict[str, Any],
+ ) -> None:
+ """Export all results to files.
+
+ Creates:
+ - summary_{timestamp}.json: Aggregate statistics
+ - user_results_{timestamp}.json: Per-user detailed results
+ - annotations_{timestamp}.csv: All annotations in flat format
+
+ Args:
+ results: Dict mapping user_id to UserSimulationResult
+ summary: Summary statistics dictionary
+ """
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+ # Export summary
+ self._export_summary(summary, timestamp)
+
+ # Export per-user results
+ self._export_user_results(results, timestamp)
+
+ # Export annotations CSV
+ self._export_annotations_csv(results, timestamp)
+
+ print(f"Results exported to {self.output_dir}/")
+
+ def _export_summary(self, summary: Dict[str, Any], timestamp: str) -> None:
+ """Export summary statistics to JSON.
+
+ Args:
+ summary: Summary dictionary
+ timestamp: Timestamp string for filename
+ """
+ filepath = os.path.join(self.output_dir, f"summary_{timestamp}.json")
+ with open(filepath, "w") as f:
+ json.dump(summary, f, indent=2, default=str)
+ print(f" - Summary: {filepath}")
+
+ def _export_user_results(
+ self,
+ results: Dict[str, UserSimulationResult],
+ timestamp: str,
+ ) -> None:
+ """Export detailed per-user results to JSON.
+
+ Args:
+ results: Dict mapping user_id to UserSimulationResult
+ timestamp: Timestamp string for filename
+ """
+ filepath = os.path.join(self.output_dir, f"user_results_{timestamp}.json")
+
+ export_data = {}
+ for user_id, result in results.items():
+ export_data[user_id] = {
+ "user_id": result.user_id,
+ "total_annotations": len(result.annotations),
+ "total_time": result.total_time,
+ "attention_checks_passed": result.attention_checks_passed,
+ "attention_checks_failed": result.attention_checks_failed,
+ "gold_standard_correct": result.gold_standard_correct,
+ "gold_standard_incorrect": result.gold_standard_incorrect,
+ "was_blocked": result.was_blocked,
+ "errors": result.errors,
+ "start_time": (
+ result.start_time.isoformat() if result.start_time else None
+ ),
+ "end_time": result.end_time.isoformat() if result.end_time else None,
+ }
+
+ with open(filepath, "w") as f:
+ json.dump(export_data, f, indent=2)
+ print(f" - User results: {filepath}")
+
+ def _export_annotations_csv(
+ self,
+ results: Dict[str, UserSimulationResult],
+ timestamp: str,
+ ) -> None:
+ """Export all annotations to CSV.
+
+ Args:
+ results: Dict mapping user_id to UserSimulationResult
+ timestamp: Timestamp string for filename
+ """
+ filepath = os.path.join(self.output_dir, f"annotations_{timestamp}.csv")
+
+ with open(filepath, "w", newline="") as f:
+ writer = csv.writer(f)
+ writer.writerow(
+ [
+ "user_id",
+ "instance_id",
+ "schema_name",
+ "annotation",
+ "response_time",
+ "timestamp",
+ "was_attention_check",
+ "attention_check_passed",
+ "was_gold_standard",
+ "gold_standard_correct",
+ ]
+ )
+
+ for user_id, result in results.items():
+ for record in result.annotations:
+ writer.writerow(
+ [
+ user_id,
+ record.instance_id,
+ record.schema_name,
+ json.dumps(record.annotation),
+ record.response_time,
+ record.timestamp.isoformat(),
+ record.was_attention_check,
+ record.attention_check_passed,
+ record.was_gold_standard,
+ record.gold_standard_correct,
+ ]
+ )
+
+ print(f" - Annotations CSV: {filepath}")
+
+ def export_annotations_jsonl(
+ self,
+ results: Dict[str, UserSimulationResult],
+ timestamp: str,
+ ) -> None:
+ """Export all annotations to JSONL (line-delimited JSON).
+
+ Useful for streaming processing and large datasets.
+
+ Args:
+ results: Dict mapping user_id to UserSimulationResult
+ timestamp: Timestamp string for filename
+ """
+ filepath = os.path.join(self.output_dir, f"annotations_{timestamp}.jsonl")
+
+ with open(filepath, "w") as f:
+ for user_id, result in results.items():
+ for record in result.annotations:
+ line_data = {
+ "user_id": user_id,
+ "instance_id": record.instance_id,
+ "schema_name": record.schema_name,
+ "annotation": record.annotation,
+ "response_time": record.response_time,
+ "timestamp": record.timestamp.isoformat(),
+ "was_attention_check": record.was_attention_check,
+ "attention_check_passed": record.attention_check_passed,
+ "was_gold_standard": record.was_gold_standard,
+ "gold_standard_correct": record.gold_standard_correct,
+ }
+ f.write(json.dumps(line_data) + "\n")
+
+ print(f" - Annotations JSONL: {filepath}")
diff --git a/potato/simulator/simulator_manager.py b/potato/simulator/simulator_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..d121ac5e5a8fe518c1a3b36a620ed83fd77ea1fa
--- /dev/null
+++ b/potato/simulator/simulator_manager.py
@@ -0,0 +1,398 @@
+"""
+Simulator manager for orchestrating multiple simulated users.
+
+This module provides the SimulatorManager class that manages multiple
+SimulatedUser instances, handling parallel execution and result aggregation.
+"""
+
+import json
+import logging
+import random
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Dict, List, Any, Optional
+
+from .config import (
+ SimulatorConfig,
+ UserConfig,
+ CompetenceLevel,
+ AnnotationStrategyType,
+)
+from .user_simulator import SimulatedUser, UserSimulationResult
+from .reporting import SimulationReporter
+
+logger = logging.getLogger(__name__)
+
+
+class SimulatorManager:
+ """Orchestrates multiple simulated users.
+
+ The SimulatorManager handles:
+ - Generating user configurations based on competence distribution
+ - Running simulations in parallel or sequentially
+ - Aggregating results across all users
+ - Exporting results via SimulationReporter
+ """
+
+ def __init__(
+ self,
+ config: SimulatorConfig,
+ server_url: str,
+ gold_standards: Optional[Dict[str, Dict[str, Any]]] = None,
+ ):
+ """Initialize simulator manager.
+
+ Args:
+ config: Simulator configuration
+ server_url: Base URL of the Potato server
+ gold_standards: Optional gold standard answers keyed by instance_id
+ """
+ self.config = config
+ self.server_url = server_url.rstrip("/")
+ self.gold_standards = gold_standards or {}
+
+ # Load gold standards from file if specified
+ if config.gold_standard_file and not gold_standards:
+ self.gold_standards = self._load_gold_standards(config.gold_standard_file)
+
+ # Generate user configs if not provided
+ self.user_configs = self._generate_user_configs()
+
+ # Results tracking
+ self.results: Dict[str, UserSimulationResult] = {}
+ self.reporter = SimulationReporter(config.output_dir)
+
+ def _load_gold_standards(self, filepath: str) -> Dict[str, Dict[str, Any]]:
+ """Load gold standards from JSON file.
+
+ Expected format:
+ [
+ {"id": "instance_001", "label_field": "value", ...},
+ ...
+ ]
+
+ Args:
+ filepath: Path to JSON file
+
+ Returns:
+ Gold standards dict keyed by instance ID
+ """
+ try:
+ with open(filepath, "r") as f:
+ items = json.load(f)
+
+ gold_standards = {}
+ for item in items:
+ item_id = item.pop("id", None)
+ if item_id:
+ gold_standards[item_id] = item
+
+ logger.info(f"Loaded {len(gold_standards)} gold standards from {filepath}")
+ return gold_standards
+
+ except Exception as e:
+ logger.warning(f"Failed to load gold standards from {filepath}: {e}")
+ return {}
+
+ def _generate_user_configs(self) -> List[UserConfig]:
+ """Generate user configurations based on competence distribution.
+
+ If explicit user configs are provided, uses those.
+ Otherwise, generates based on user_count and competence_distribution.
+
+ Returns:
+ List of UserConfig instances
+ """
+ if self.config.users:
+ return self.config.users
+
+ users = []
+
+ # Get competence distribution
+ competence_levels = list(self.config.competence_distribution.keys())
+ competence_weights = list(self.config.competence_distribution.values())
+
+ # Normalize weights
+ total_weight = sum(competence_weights)
+ if total_weight > 0:
+ competence_weights = [w / total_weight for w in competence_weights]
+
+ for i in range(self.config.user_count):
+ # Select competence level based on distribution
+ competence_str = random.choices(
+ competence_levels, weights=competence_weights, k=1
+ )[0]
+
+ try:
+ competence = CompetenceLevel(competence_str)
+ except ValueError:
+ competence = CompetenceLevel.AVERAGE
+
+ users.append(
+ UserConfig(
+ user_id=f"sim_user_{i:04d}",
+ competence=competence,
+ strategy=self.config.strategy,
+ timing=self.config.timing,
+ llm_config=self.config.llm_config,
+ biased_config=self.config.biased_config,
+ agent_config=self.config.agent_config,
+ )
+ )
+
+ logger.info(f"Generated {len(users)} user configurations")
+ return users
+
+ def run_single_user(
+ self, user_config: UserConfig, max_annotations: Optional[int] = None
+ ) -> UserSimulationResult:
+ """Run simulation for a single user.
+
+ Args:
+ user_config: Configuration for the user
+ max_annotations: Maximum annotations for this user
+
+ Returns:
+ UserSimulationResult with tracking data
+ """
+ user = SimulatedUser(
+ user_config=user_config,
+ server_url=self.server_url,
+ gold_standards=self.gold_standards,
+ simulate_wait=self.config.simulate_wait,
+ attention_check_fail_rate=self.config.attention_check_fail_rate,
+ respond_fast_rate=self.config.respond_fast_rate,
+ interactive_config=self.config.interactive,
+ )
+
+ result = user.run_simulation(max_annotations)
+ self.results[user_config.user_id] = result
+
+ return result
+
+ def run_parallel(
+ self, max_annotations_per_user: Optional[int] = None
+ ) -> Dict[str, UserSimulationResult]:
+ """Run simulation for all users in parallel.
+
+ Args:
+ max_annotations_per_user: Maximum annotations per user
+
+ Returns:
+ Dict mapping user_id to UserSimulationResult
+ """
+ logger.info(
+ f"Starting parallel simulation with {len(self.user_configs)} users "
+ f"({self.config.parallel_users} concurrent)"
+ )
+
+ with ThreadPoolExecutor(max_workers=self.config.parallel_users) as executor:
+ futures = {}
+
+ for i, user_config in enumerate(self.user_configs):
+ # Stagger user starts
+ if i > 0 and self.config.delay_between_users > 0:
+ time.sleep(self.config.delay_between_users)
+
+ future = executor.submit(
+ self.run_single_user, user_config, max_annotations_per_user
+ )
+ futures[future] = user_config.user_id
+
+ # Wait for completion
+ completed = 0
+ for future in as_completed(futures):
+ user_id = futures[future]
+ completed += 1
+ try:
+ result = future.result()
+ logger.info(
+ f"[{completed}/{len(futures)}] User {user_id} completed: "
+ f"{len(result.annotations)} annotations"
+ )
+ except Exception as e:
+ logger.error(f"User {user_id} failed: {e}")
+
+ logger.info(f"Parallel simulation completed: {len(self.results)} users")
+ return self.results
+
+ def run_sequential(
+ self, max_annotations_per_user: Optional[int] = None
+ ) -> Dict[str, UserSimulationResult]:
+ """Run simulation for all users sequentially.
+
+ Args:
+ max_annotations_per_user: Maximum annotations per user
+
+ Returns:
+ Dict mapping user_id to UserSimulationResult
+ """
+ logger.info(
+ f"Starting sequential simulation with {len(self.user_configs)} users"
+ )
+
+ for i, user_config in enumerate(self.user_configs):
+ result = self.run_single_user(user_config, max_annotations_per_user)
+ logger.info(
+ f"[{i+1}/{len(self.user_configs)}] User {user_config.user_id} "
+ f"completed: {len(result.annotations)} annotations"
+ )
+
+ logger.info(f"Sequential simulation completed: {len(self.results)} users")
+ return self.results
+
+ def get_summary(self) -> Dict[str, Any]:
+ """Get summary statistics for all users.
+
+ Returns:
+ Summary dictionary with aggregate statistics
+ """
+ if not self.results:
+ return {"error": "No results available"}
+
+ total_annotations = sum(len(r.annotations) for r in self.results.values())
+ total_time = sum(r.total_time for r in self.results.values())
+
+ total_attention_passed = sum(
+ r.attention_checks_passed for r in self.results.values()
+ )
+ total_attention_failed = sum(
+ r.attention_checks_failed for r in self.results.values()
+ )
+ total_gold_correct = sum(
+ r.gold_standard_correct for r in self.results.values()
+ )
+ total_gold_incorrect = sum(
+ r.gold_standard_incorrect for r in self.results.values()
+ )
+
+ blocked_users = sum(1 for r in self.results.values() if r.was_blocked)
+ users_with_errors = sum(1 for r in self.results.values() if r.errors)
+
+ # Calculate response time statistics
+ all_response_times = [
+ record.response_time
+ for result in self.results.values()
+ for record in result.annotations
+ ]
+
+ response_time_stats = {}
+ if all_response_times:
+ response_time_stats = {
+ "min": min(all_response_times),
+ "max": max(all_response_times),
+ "mean": sum(all_response_times) / len(all_response_times),
+ }
+
+ # Competence level distribution in results
+ competence_distribution = {}
+ for user_id in self.results:
+ for config in self.user_configs:
+ if config.user_id == user_id:
+ level = config.competence.value
+ competence_distribution[level] = (
+ competence_distribution.get(level, 0) + 1
+ )
+ break
+
+ return {
+ "user_count": len(self.results),
+ "total_annotations": total_annotations,
+ "total_time_seconds": total_time,
+ "average_annotations_per_user": (
+ total_annotations / len(self.results) if self.results else 0
+ ),
+ "average_time_per_user": (
+ total_time / len(self.results) if self.results else 0
+ ),
+ "attention_checks": {
+ "passed": total_attention_passed,
+ "failed": total_attention_failed,
+ "pass_rate": (
+ total_attention_passed
+ / (total_attention_passed + total_attention_failed)
+ if (total_attention_passed + total_attention_failed) > 0
+ else None
+ ),
+ },
+ "gold_standards": {
+ "correct": total_gold_correct,
+ "incorrect": total_gold_incorrect,
+ "accuracy": (
+ total_gold_correct / (total_gold_correct + total_gold_incorrect)
+ if (total_gold_correct + total_gold_incorrect) > 0
+ else None
+ ),
+ },
+ "blocked_users": blocked_users,
+ "users_with_errors": users_with_errors,
+ "response_time_stats": response_time_stats,
+ "competence_distribution": competence_distribution,
+ "per_user": {
+ user_id: {
+ "annotations": len(r.annotations),
+ "total_time": r.total_time,
+ "attention_passed": r.attention_checks_passed,
+ "attention_failed": r.attention_checks_failed,
+ "gold_correct": r.gold_standard_correct,
+ "gold_incorrect": r.gold_standard_incorrect,
+ "was_blocked": r.was_blocked,
+ "errors": len(r.errors),
+ }
+ for user_id, r in self.results.items()
+ },
+ }
+
+ def export_results(self) -> str:
+ """Export all results using the reporter.
+
+ Returns:
+ Path to the output directory
+ """
+ self.reporter.export_results(self.results, self.get_summary())
+ return self.config.output_dir
+
+ def print_summary(self) -> None:
+ """Print a summary of results to stdout."""
+ summary = self.get_summary()
+
+ print("\n" + "=" * 60)
+ print("SIMULATION SUMMARY")
+ print("=" * 60)
+
+ print(f"\nUsers: {summary['user_count']}")
+ print(f"Total annotations: {summary['total_annotations']}")
+ print(f"Total time: {summary['total_time_seconds']:.1f}s")
+ print(
+ f"Avg annotations/user: {summary['average_annotations_per_user']:.1f}"
+ )
+ print(f"Avg time/user: {summary['average_time_per_user']:.1f}s")
+
+ ac = summary["attention_checks"]
+ if ac["passed"] or ac["failed"]:
+ print(f"\nAttention Checks:")
+ print(f" Passed: {ac['passed']}")
+ print(f" Failed: {ac['failed']}")
+ if ac["pass_rate"] is not None:
+ print(f" Pass rate: {ac['pass_rate']:.1%}")
+
+ gs = summary["gold_standards"]
+ if gs["correct"] or gs["incorrect"]:
+ print(f"\nGold Standards:")
+ print(f" Correct: {gs['correct']}")
+ print(f" Incorrect: {gs['incorrect']}")
+ if gs["accuracy"] is not None:
+ print(f" Accuracy: {gs['accuracy']:.1%}")
+
+ if summary["blocked_users"]:
+ print(f"\nBlocked users: {summary['blocked_users']}")
+
+ if summary["users_with_errors"]:
+ print(f"Users with errors: {summary['users_with_errors']}")
+
+ if summary["competence_distribution"]:
+ print(f"\nCompetence distribution:")
+ for level, count in summary["competence_distribution"].items():
+ print(f" {level}: {count}")
+
+ print("\n" + "=" * 60)
diff --git a/potato/simulator/solo_mode_simulator.py b/potato/simulator/solo_mode_simulator.py
new file mode 100644
index 0000000000000000000000000000000000000000..5aed5af185fd308e96d7e190bef9de1ff8cdc87c
--- /dev/null
+++ b/potato/simulator/solo_mode_simulator.py
@@ -0,0 +1,924 @@
+"""
+Solo Mode Simulator
+
+Drives all 12 phases of Potato's Solo Mode through the /solo/* HTTP endpoints,
+simulating a single human annotator collaborating with an LLM.
+
+Usage:
+ from potato.simulator.solo_mode_simulator import SoloModeSimulator, SoloSimulatorConfig
+
+ sim = SoloModeSimulator(
+ server_url="http://localhost:8200",
+ gold_labels={"emo_001": "joy", "emo_002": "sadness", ...},
+ config=SoloSimulatorConfig(noise_rate=0.2),
+ )
+ result = sim.run_full_simulation()
+ print(result.summary())
+"""
+
+import json
+import logging
+import random
+import time
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SoloSimulatorConfig:
+ """Configuration for solo mode simulation."""
+
+ # User identity
+ user_id: str = "solo_simulator"
+ password: str = "simulated_password_123"
+
+ # Task description for setup phase
+ task_description: str = (
+ "Classify the primary emotion expressed in each text. "
+ "Choose the single best label from: joy, sadness, anger, fear, surprise, neutral."
+ )
+
+ # Annotation noise
+ noise_rate: float = 0.2 # probability of choosing wrong label
+
+ # Disagreement resolution strategy
+ disagree_prefer_human: float = 0.6
+ disagree_prefer_llm: float = 0.25
+ disagree_prefer_third: float = 0.15
+
+ # Annotation counts per phase
+ parallel_annotation_count: int = 30
+ active_annotation_count: int = 50
+
+ # Review/validation behavior
+ review_approve_rate: float = 0.7
+ rule_approve_rate: float = 0.7
+
+ # Timing
+ max_wait_autonomous: int = 120 # seconds to wait for autonomous labeling
+ poll_interval: float = 2.0 # seconds between status polls
+ annotation_delay: float = 0.0 # seconds to wait between annotations (realistic timing)
+ wait_for_predictions_timeout: int = 60 # seconds to wait for LLM predictions
+
+ # Phase control
+ force_advance_on_stuck: bool = True # use /api/advance-phase if stuck
+
+ # Schema name (must match config)
+ schema_name: str = "emotion"
+
+
+@dataclass
+class PhaseResult:
+ """Result from simulating a single phase."""
+
+ phase: str
+ success: bool = False
+ annotations_submitted: int = 0
+ disagreements_encountered: int = 0
+ disagreements_resolved: int = 0
+ reviews_completed: int = 0
+ errors: List[str] = field(default_factory=list)
+ duration_seconds: float = 0.0
+ metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class SoloSimulationResult:
+ """Complete results from a solo mode simulation run."""
+
+ phase_results: List[PhaseResult] = field(default_factory=list)
+ total_annotations: int = 0
+ total_disagreements: int = 0
+ total_errors: List[str] = field(default_factory=list)
+ start_time: Optional[datetime] = None
+ end_time: Optional[datetime] = None
+ final_status: Optional[Dict[str, Any]] = None
+
+ def summary(self) -> str:
+ lines = [
+ "=== Solo Mode Simulation Summary ===",
+ f"Duration: {(self.end_time - self.start_time).total_seconds():.1f}s"
+ if self.start_time and self.end_time
+ else "Duration: unknown",
+ f"Total annotations: {self.total_annotations}",
+ f"Total disagreements: {self.total_disagreements}",
+ f"Errors: {len(self.total_errors)}",
+ "",
+ "Phase Results:",
+ ]
+ for pr in self.phase_results:
+ status = "OK" if pr.success else "FAIL"
+ lines.append(
+ f" {pr.phase}: {status} "
+ f"(annotations={pr.annotations_submitted}, "
+ f"disagreements={pr.disagreements_encountered}, "
+ f"duration={pr.duration_seconds:.1f}s)"
+ )
+ if self.final_status:
+ lines.append("")
+ lines.append("Final Status:")
+ agreement = self.final_status.get("agreement", {})
+ lines.append(f" Agreement rate: {agreement.get('agreement_rate', 'N/A')}")
+ lines.append(
+ f" Total compared: {agreement.get('total_compared', 'N/A')}"
+ )
+ labeling = self.final_status.get("labeling", {})
+ lines.append(f" Human labeled: {labeling.get('human_labeled', 'N/A')}")
+ lines.append(f" LLM labeled: {labeling.get('llm_labeled', 'N/A')}")
+ prompt = self.final_status.get("prompt", {})
+ lines.append(
+ f" Prompt versions: {prompt.get('total_versions', 'N/A')}"
+ )
+ if self.total_errors:
+ lines.append("")
+ lines.append(f"Errors ({len(self.total_errors)}):")
+ for err in self.total_errors[:10]:
+ lines.append(f" - {err}")
+ return "\n".join(lines)
+
+
+class SoloModeSimulator:
+ """Simulates a single user driving all solo mode phases via HTTP.
+
+ Args:
+ server_url: Base URL of the Potato server
+ gold_labels: Dict mapping instance_id -> gold label string
+ available_labels: List of all valid label names
+ config: SoloSimulatorConfig
+ """
+
+ def __init__(
+ self,
+ server_url: str,
+ gold_labels: Dict[str, str],
+ available_labels: Optional[List[str]] = None,
+ config: Optional[SoloSimulatorConfig] = None,
+ ):
+ self.server_url = server_url.rstrip("/")
+ self.gold_labels = gold_labels
+ self.available_labels = available_labels or []
+ self.config = config or SoloSimulatorConfig()
+ self.session = requests.Session()
+ self.result = SoloSimulationResult()
+ self._logged_in = False
+
+ # === Authentication ===
+
+ def _login(self) -> bool:
+ """Register and login the simulated user."""
+ try:
+ self.session.post(
+ f"{self.server_url}/register",
+ data={
+ "action": "signup",
+ "email": self.config.user_id,
+ "pass": self.config.password,
+ },
+ allow_redirects=True,
+ timeout=30,
+ )
+ self.session.post(
+ f"{self.server_url}/auth",
+ data={
+ "action": "login",
+ "email": self.config.user_id,
+ "pass": self.config.password,
+ },
+ allow_redirects=True,
+ timeout=30,
+ )
+ self._logged_in = True
+ logger.info(f"Logged in as {self.config.user_id}")
+ return True
+ except requests.exceptions.RequestException as e:
+ logger.error(f"Login failed: {e}")
+ self.result.total_errors.append(f"Login failed: {e}")
+ return False
+
+ # === Status & Phase Helpers ===
+
+ def _get_status(self) -> Dict[str, Any]:
+ """Get current solo mode status."""
+ try:
+ resp = self.session.get(
+ f"{self.server_url}/solo/api/status", timeout=30
+ )
+ if resp.status_code == 200:
+ return resp.json()
+ except Exception as e:
+ logger.warning(f"Failed to get status: {e}")
+ return {}
+
+ def _get_current_phase(self) -> str:
+ """Get current phase name."""
+ status = self._get_status()
+ # Status may return phase as int, string, or dict depending on version
+ phase_info = status.get("phase", {})
+ if isinstance(phase_info, dict):
+ return phase_info.get("current_phase", "unknown")
+ # phase_name is the string representation
+ return status.get("phase_name", str(phase_info)).lower()
+
+ def _force_advance(self, target_phase: str) -> bool:
+ """Force advance to a specific phase."""
+ try:
+ resp = self.session.post(
+ f"{self.server_url}/solo/api/advance-phase",
+ json={"phase": target_phase, "force": True},
+ timeout=30,
+ )
+ if resp.status_code == 200:
+ logger.info(f"Force-advanced to {target_phase}")
+ return True
+ logger.warning(
+ f"Force advance to {target_phase} failed: {resp.status_code} {resp.text}"
+ )
+ except Exception as e:
+ logger.error(f"Force advance failed: {e}")
+ return False
+
+ def _generate_annotation(self, instance_id: str) -> str:
+ """Generate an annotation for an instance.
+
+ Uses gold label with noise_rate probability of choosing wrong label.
+ """
+ gold = self.gold_labels.get(instance_id)
+ if gold and random.random() > self.config.noise_rate:
+ return gold
+
+ # Choose a random label (possibly different from gold)
+ if self.available_labels:
+ candidates = [l for l in self.available_labels if l != gold]
+ if candidates:
+ return random.choice(candidates)
+ return gold or (self.available_labels[0] if self.available_labels else "unknown")
+
+ # === Phase Simulators ===
+
+ def _simulate_setup(self) -> PhaseResult:
+ """Simulate the SETUP phase: submit task description."""
+ start = time.time()
+ result = PhaseResult(phase="setup")
+
+ try:
+ resp = self.session.post(
+ f"{self.server_url}/solo/setup",
+ data={"task_description": self.config.task_description},
+ allow_redirects=True,
+ timeout=30,
+ )
+ result.success = resp.status_code in (200, 302)
+ if not result.success:
+ result.errors.append(f"Setup failed: {resp.status_code}")
+ logger.info(f"Setup phase: {'OK' if result.success else 'FAIL'}")
+ except Exception as e:
+ result.errors.append(f"Setup error: {e}")
+ result.success = False
+
+ result.duration_seconds = time.time() - start
+ return result
+
+ def _simulate_prompt_review(self) -> PhaseResult:
+ """Simulate PROMPT_REVIEW: accept the prompt and advance."""
+ start = time.time()
+ result = PhaseResult(phase="prompt_review")
+
+ try:
+ # Accept prompt and advance to edge cases
+ resp = self.session.post(
+ f"{self.server_url}/solo/prompt",
+ data={"action": "advance"},
+ allow_redirects=True,
+ timeout=30,
+ )
+ result.success = resp.status_code in (200, 302)
+ if not result.success:
+ result.errors.append(f"Prompt review failed: {resp.status_code}")
+ except Exception as e:
+ result.errors.append(f"Prompt review error: {e}")
+ result.success = False
+
+ result.duration_seconds = time.time() - start
+ return result
+
+ def _simulate_edge_case_labeling(self) -> PhaseResult:
+ """Simulate EDGE_CASE_SYNTHESIS + EDGE_CASE_LABELING."""
+ start = time.time()
+ result = PhaseResult(phase="edge_case_labeling")
+
+ try:
+ # GET to trigger synthesis
+ resp = self.session.get(
+ f"{self.server_url}/solo/edge-cases",
+ allow_redirects=True,
+ timeout=60,
+ )
+
+ # Label edge cases in a loop
+ labeled_count = 0
+ for _ in range(20): # max iterations to avoid infinite loop
+ resp = self.session.get(
+ f"{self.server_url}/solo/edge-cases",
+ timeout=30,
+ )
+ if resp.status_code != 200:
+ break
+
+ # Parse the page to find the current case
+ # Use API endpoint if available, otherwise check page content
+ api_resp = self.session.get(
+ f"{self.server_url}/solo/api/edge-cases",
+ timeout=30,
+ )
+ if api_resp.status_code == 200:
+ ec_data = api_resp.json()
+ unlabeled = ec_data.get("unlabeled", 0)
+ if unlabeled == 0:
+ break
+
+ # Submit a label for the current edge case
+ label = random.choice(self.available_labels) if self.available_labels else "neutral"
+ resp = self.session.post(
+ f"{self.server_url}/solo/edge-cases",
+ data={"label": label},
+ allow_redirects=True,
+ timeout=30,
+ )
+ labeled_count += 1
+
+ result.annotations_submitted = labeled_count
+ result.success = True
+ logger.info(f"Edge case labeling: labeled {labeled_count} cases")
+ except Exception as e:
+ result.errors.append(f"Edge case labeling error: {e}")
+ result.success = False
+
+ result.duration_seconds = time.time() - start
+ return result
+
+ def _simulate_annotation(self, count: int) -> PhaseResult:
+ """Simulate annotation phase (PARALLEL or ACTIVE).
+
+ Submits `count` annotations via /solo/annotate.
+ Handles disagreement redirects inline.
+ """
+ start = time.time()
+ result = PhaseResult(phase="annotation")
+ annotations_done = 0
+ disagreements = 0
+
+ for i in range(count):
+ try:
+ # GET next instance
+ resp = self.session.get(
+ f"{self.server_url}/solo/annotate",
+ allow_redirects=True,
+ timeout=30,
+ )
+ if resp.status_code != 200:
+ result.errors.append(f"Get annotate failed: {resp.status_code}")
+ continue
+
+ # Extract instance_id from the page
+ # Look for hidden input or data attribute
+ instance_id = self._extract_instance_id(resp.text)
+ if not instance_id:
+ logger.debug("No instance available, stopping annotation")
+ break
+
+ # Generate annotation
+ annotation = self._generate_annotation(instance_id)
+
+ # Submit annotation
+ resp = self.session.post(
+ f"{self.server_url}/solo/annotate",
+ data={"instance_id": instance_id, "annotation": annotation},
+ allow_redirects=False,
+ timeout=30,
+ )
+
+ annotations_done += 1
+
+ # Progress logging every 10 annotations
+ if annotations_done % 10 == 0:
+ status = self._get_status()
+ agreement = status.get("agreement_metrics", {})
+ stats = status.get("annotation_stats", {})
+ logger.info(
+ f"[Progress] {annotations_done}/{count} annotations, "
+ f"{disagreements} disagreements, "
+ f"agreement={agreement.get('agreement_rate', 0):.3f} "
+ f"({agreement.get('total_compared', 0)} compared), "
+ f"LLM labeled={stats.get('llm_labeled', 0)}"
+ )
+
+ # Simulate realistic annotation time
+ if self.config.annotation_delay > 0:
+ time.sleep(self.config.annotation_delay)
+
+ # Check if redirected to disagreements
+ if resp.status_code == 302:
+ location = resp.headers.get("Location", "")
+ if "disagreement" in location:
+ disagreements += 1
+ self._handle_disagreement()
+
+ except Exception as e:
+ result.errors.append(f"Annotation {i} error: {e}")
+
+ result.annotations_submitted = annotations_done
+ result.disagreements_encountered = disagreements
+ result.success = annotations_done > 0
+ result.duration_seconds = time.time() - start
+ logger.info(
+ f"Annotation phase: {annotations_done} annotations, "
+ f"{disagreements} disagreements"
+ )
+ return result
+
+ def _handle_disagreement(self) -> None:
+ """Handle a single disagreement resolution.
+
+ Simulates a human adjudicator who sees both the human and LLM labels
+ and decides which is correct. Uses gold labels when available to make
+ informed choices: if the gold label matches the human's label, choose
+ "human"; if it matches the LLM's label, choose "llm"; otherwise pick
+ the gold label as a third option.
+ """
+ try:
+ # GET the disagreement page
+ resp = self.session.get(
+ f"{self.server_url}/solo/disagreements",
+ allow_redirects=True,
+ timeout=30,
+ )
+ if resp.status_code != 200:
+ return
+
+ # Extract disagreement ID from page
+ disagreement_id = self._extract_disagreement_id(resp.text)
+ if not disagreement_id:
+ return
+
+ # Parse instance_id from disagreement_id (format: "instance_id:schema")
+ instance_id = disagreement_id.split(":")[0]
+ gold = self.gold_labels.get(instance_id)
+
+ # Decide resolution based on gold label when available
+ if gold:
+ # Check which side the gold agrees with
+ # (We don't have the actual labels from the page, so use
+ # "human" or "llm" and let the server resolve to actual values)
+ roll = random.random()
+ if roll < self.config.disagree_prefer_human:
+ resolution = "human"
+ elif roll < self.config.disagree_prefer_human + self.config.disagree_prefer_llm:
+ resolution = "llm"
+ else:
+ # Use the gold label directly as a third option
+ resolution = gold
+ else:
+ roll = random.random()
+ if roll < self.config.disagree_prefer_human:
+ resolution = "human"
+ elif roll < self.config.disagree_prefer_human + self.config.disagree_prefer_llm:
+ resolution = "llm"
+ else:
+ resolution = random.choice(self.available_labels) if self.available_labels else "neutral"
+
+ # Submit resolution
+ self.session.post(
+ f"{self.server_url}/solo/disagreements",
+ data={
+ "disagreement_id": disagreement_id,
+ "resolution": resolution,
+ },
+ allow_redirects=True,
+ timeout=30,
+ )
+ logger.debug(f"Resolved disagreement {disagreement_id}: {resolution}")
+
+ except Exception as e:
+ logger.warning(f"Disagreement resolution error: {e}")
+
+ def _simulate_periodic_review(self) -> PhaseResult:
+ """Simulate PERIODIC_REVIEW: approve or correct low-confidence labels."""
+ start = time.time()
+ result = PhaseResult(phase="periodic_review")
+ reviews = 0
+
+ try:
+ for _ in range(20): # max iterations
+ resp = self.session.get(
+ f"{self.server_url}/solo/review",
+ allow_redirects=True,
+ timeout=30,
+ )
+ if resp.status_code != 200:
+ break
+
+ instance_id = self._extract_instance_id(resp.text)
+ if not instance_id:
+ break
+
+ # Decide: approve or correct
+ if random.random() < self.config.review_approve_rate:
+ decision = "approve"
+ corrected_label = ""
+ else:
+ decision = "correct"
+ gold = self.gold_labels.get(instance_id, "")
+ corrected_label = gold or (
+ random.choice(self.available_labels) if self.available_labels else ""
+ )
+
+ self.session.post(
+ f"{self.server_url}/solo/review",
+ data={
+ "instance_id": instance_id,
+ "decision": decision,
+ "corrected_label": corrected_label,
+ },
+ allow_redirects=True,
+ timeout=30,
+ )
+ reviews += 1
+
+ except Exception as e:
+ result.errors.append(f"Review error: {e}")
+
+ result.reviews_completed = reviews
+ result.success = True
+ result.duration_seconds = time.time() - start
+ return result
+
+ def _simulate_rule_review(self) -> PhaseResult:
+ """Simulate RULE_REVIEW: approve or reject edge case rule categories."""
+ start = time.time()
+ result = PhaseResult(phase="rule_review")
+
+ try:
+ resp = self.session.get(
+ f"{self.server_url}/solo/api/rules/categories",
+ timeout=30,
+ )
+ if resp.status_code == 200:
+ data = resp.json()
+ categories = data.get("categories", [])
+ for cat in categories:
+ cat_id = cat.get("id", "")
+ action = "approve" if random.random() < self.config.rule_approve_rate else "reject"
+ self.session.post(
+ f"{self.server_url}/solo/api/rules/approve",
+ json={"category_id": cat_id, "action": action, "notes": ""},
+ timeout=30,
+ )
+ result.metadata["categories_reviewed"] = len(categories)
+
+ result.success = True
+ except Exception as e:
+ result.errors.append(f"Rule review error: {e}")
+ result.success = False
+
+ result.duration_seconds = time.time() - start
+ return result
+
+ def _simulate_autonomous_wait(self) -> PhaseResult:
+ """Wait for autonomous labeling to complete."""
+ start = time.time()
+ result = PhaseResult(phase="autonomous_labeling")
+
+ # Start labeling if not running
+ self.session.post(
+ f"{self.server_url}/solo/api/start-labeling",
+ timeout=30,
+ )
+
+ # Poll until done or timeout
+ deadline = time.time() + self.config.max_wait_autonomous
+ while time.time() < deadline:
+ status = self._get_status()
+ labeling = status.get("labeling", {})
+ if not labeling.get("background_running", True):
+ break
+ time.sleep(self.config.poll_interval)
+
+ result.success = True
+ result.duration_seconds = time.time() - start
+ return result
+
+ def _simulate_validation(self) -> PhaseResult:
+ """Simulate FINAL_VALIDATION: validate a sample of LLM-only labels."""
+ start = time.time()
+ result = PhaseResult(phase="final_validation")
+ validated = 0
+
+ try:
+ for _ in range(100): # max iterations
+ resp = self.session.get(
+ f"{self.server_url}/solo/validation",
+ allow_redirects=True,
+ timeout=30,
+ )
+ if resp.status_code != 200:
+ break
+
+ instance_id = self._extract_instance_id(resp.text)
+ if not instance_id:
+ break
+
+ # Get gold label or approve LLM label
+ gold = self.gold_labels.get(instance_id)
+ if gold:
+ self.session.post(
+ f"{self.server_url}/solo/validation",
+ data={
+ "instance_id": instance_id,
+ "human_label": gold,
+ },
+ allow_redirects=True,
+ timeout=30,
+ )
+ else:
+ # Approve the LLM label
+ self.session.post(
+ f"{self.server_url}/solo/validation",
+ data={
+ "instance_id": instance_id,
+ "decision": "approve",
+ },
+ allow_redirects=True,
+ timeout=30,
+ )
+ validated += 1
+
+ except Exception as e:
+ result.errors.append(f"Validation error: {e}")
+
+ result.annotations_submitted = validated
+ result.success = True
+ result.duration_seconds = time.time() - start
+ return result
+
+ # === LLM Control Helpers ===
+
+ def _start_llm_labeling(self) -> None:
+ """Start the background LLM labeling thread."""
+ try:
+ self.session.post(
+ f"{self.server_url}/solo/api/start-labeling",
+ timeout=30,
+ )
+ logger.info("Started LLM labeling")
+ except Exception as e:
+ logger.warning(f"Failed to start LLM labeling: {e}")
+
+ def _wait_for_predictions(self, min_count: int = 10, timeout: int = 60) -> int:
+ """Wait until at least min_count LLM predictions exist.
+
+ Returns:
+ Actual prediction count
+ """
+ deadline = time.time() + timeout
+ while time.time() < deadline:
+ try:
+ resp = self.session.get(
+ f"{self.server_url}/solo/api/predictions",
+ timeout=30,
+ )
+ if resp.status_code == 200:
+ count = resp.json().get("count", 0)
+ if count >= min_count:
+ logger.info(f"LLM has {count} predictions (target: {min_count})")
+ return count
+ except Exception:
+ pass
+ time.sleep(self.config.poll_interval)
+ logger.warning(f"Timeout waiting for {min_count} predictions")
+ return 0
+
+ # === HTML Parsing Helpers ===
+
+ def _extract_instance_id(self, html: str) -> Optional[str]:
+ """Extract instance_id from an HTML page.
+
+ Looks for common patterns: hidden input, data attribute, or JSON.
+ """
+ import re
+
+ # Pattern:
+ match = re.search(
+ r'name=["\']instance_id["\'][^>]*value=["\']([^"\']+)', html
+ )
+ if match:
+ return match.group(1)
+
+ # Pattern: value="..." name="instance_id"
+ match = re.search(
+ r'value=["\']([^"\']+)["\'][^>]*name=["\']instance_id', html
+ )
+ if match:
+ return match.group(1)
+
+ # Pattern: data-instance-id="..."
+ match = re.search(r'data-instance-id=["\']([^"\']+)', html)
+ if match:
+ return match.group(1)
+
+ # Pattern: "instance_id": "..."
+ match = re.search(r'"instance_id"\s*:\s*"([^"]+)"', html)
+ if match:
+ return match.group(1)
+
+ return None
+
+ def _extract_disagreement_id(self, html: str) -> Optional[str]:
+ """Extract disagreement_id from the disagreement page."""
+ import re
+
+ match = re.search(
+ r'name=["\']disagreement_id["\'][^>]*value=["\']([^"\']+)', html
+ )
+ if match:
+ return match.group(1)
+
+ match = re.search(
+ r'value=["\']([^"\']+)["\'][^>]*name=["\']disagreement_id', html
+ )
+ if match:
+ return match.group(1)
+
+ return None
+
+ # === Main Orchestration ===
+
+ def run_full_simulation(self) -> SoloSimulationResult:
+ """Run the complete solo mode simulation.
+
+ Drives through all phases: setup, prompt review, edge cases,
+ parallel annotation, active annotation, review, validation.
+
+ Returns:
+ SoloSimulationResult with per-phase metrics
+ """
+ self.result.start_time = datetime.now()
+ random.seed(42)
+
+ try:
+ # Login
+ if not self._login():
+ return self.result
+
+ # Discover available labels from schemas
+ if not self.available_labels:
+ self._discover_labels()
+
+ # Phase 1: Setup
+ pr = self._simulate_setup()
+ self.result.phase_results.append(pr)
+
+ # Phase 2: Prompt Review
+ pr = self._simulate_prompt_review()
+ self.result.phase_results.append(pr)
+
+ # Phase 3-5: Edge Case Synthesis + Labeling + Validation
+ current_phase = self._get_current_phase()
+ if "edge" in current_phase.lower():
+ pr = self._simulate_edge_case_labeling()
+ self.result.phase_results.append(pr)
+
+ # Phase 6: Parallel Annotation
+ if self.config.force_advance_on_stuck:
+ self._force_advance("parallel-annotation")
+
+ # Start LLM labeling and wait for predictions to accumulate
+ self._start_llm_labeling()
+ self._wait_for_predictions(min_count=20, timeout=self.config.wait_for_predictions_timeout)
+
+ pr = self._simulate_annotation(count=self.config.parallel_annotation_count)
+ self.result.phase_results.append(pr)
+ self.result.total_annotations += pr.annotations_submitted
+ self.result.total_disagreements += pr.disagreements_encountered
+
+ # Phase 8: Active Annotation
+ if self.config.force_advance_on_stuck:
+ self._force_advance("active-annotation")
+
+ # Wait for more LLM predictions
+ self._wait_for_predictions(min_count=50, timeout=self.config.wait_for_predictions_timeout)
+
+ pr = self._simulate_annotation(count=self.config.active_annotation_count)
+ self.result.phase_results.append(pr)
+ self.result.total_annotations += pr.annotations_submitted
+ self.result.total_disagreements += pr.disagreements_encountered
+
+ # Check for periodic review
+ current_phase = self._get_current_phase()
+ if "review" in current_phase.lower() and "rule" not in current_phase.lower():
+ pr = self._simulate_periodic_review()
+ self.result.phase_results.append(pr)
+
+ # Check for rule review
+ if "rule" in current_phase.lower():
+ pr = self._simulate_rule_review()
+ self.result.phase_results.append(pr)
+
+ # Phase 11: Autonomous Labeling
+ if self.config.force_advance_on_stuck:
+ self._force_advance("autonomous-labeling")
+ pr = self._simulate_autonomous_wait()
+ self.result.phase_results.append(pr)
+
+ # Phase 12: Final Validation
+ if self.config.force_advance_on_stuck:
+ self._force_advance("final-validation")
+ pr = self._simulate_validation()
+ self.result.phase_results.append(pr)
+
+ # Collect final status
+ self.result.final_status = self._get_status()
+
+ except Exception as e:
+ logger.error(f"Simulation error: {e}")
+ self.result.total_errors.append(f"Simulation error: {e}")
+
+ finally:
+ self.result.end_time = datetime.now()
+ # Collect errors from all phases
+ for pr in self.result.phase_results:
+ self.result.total_errors.extend(pr.errors)
+
+ return self.result
+
+ def _discover_labels(self) -> None:
+ """Discover available labels from the server's schema API."""
+ try:
+ resp = self.session.get(
+ f"{self.server_url}/api/schemas",
+ timeout=30,
+ )
+ if resp.status_code == 200:
+ data = resp.json()
+ if isinstance(data, dict) and "schemas" in data:
+ schemas = data["schemas"]
+ if isinstance(schemas, dict):
+ schemas = list(schemas.values())
+ elif isinstance(data, list):
+ schemas = data
+ else:
+ schemas = list(data.values()) if isinstance(data, dict) else []
+
+ for schema in schemas:
+ labels = schema.get("labels", [])
+ for label in labels:
+ if isinstance(label, str):
+ self.available_labels.append(label)
+ elif isinstance(label, dict):
+ self.available_labels.append(
+ label.get("name", str(label))
+ )
+ logger.info(f"Discovered labels: {self.available_labels}")
+ except Exception as e:
+ logger.warning(f"Failed to discover labels: {e}")
+
+ def get_verification_data(self) -> Dict[str, Any]:
+ """Collect data for verification checks after simulation.
+
+ Returns:
+ Dict with status, predictions, prompts, and phase history.
+ """
+ data = {"status": self._get_status()}
+
+ try:
+ resp = self.session.get(
+ f"{self.server_url}/solo/api/prompts", timeout=30
+ )
+ if resp.status_code == 200:
+ data["prompts"] = resp.json()
+ except Exception:
+ pass
+
+ try:
+ resp = self.session.get(
+ f"{self.server_url}/solo/api/predictions", timeout=30
+ )
+ if resp.status_code == 200:
+ data["predictions"] = resp.json()
+ except Exception:
+ pass
+
+ try:
+ resp = self.session.get(
+ f"{self.server_url}/solo/api/disagreements", timeout=30
+ )
+ if resp.status_code == 200:
+ data["disagreements"] = resp.json()
+ except Exception:
+ pass
+
+ return data
diff --git a/potato/simulator/timing_models.py b/potato/simulator/timing_models.py
new file mode 100644
index 0000000000000000000000000000000000000000..96799f822abfeae4a2cee9873effb342f6c0bd80
--- /dev/null
+++ b/potato/simulator/timing_models.py
@@ -0,0 +1,133 @@
+"""
+Timing models for simulated annotation behavior.
+
+This module provides realistic timing distributions for simulated
+annotations, supporting various statistical distributions.
+"""
+
+import random
+import time
+from typing import Optional
+
+from .config import TimingConfig
+
+
+class TimingModel:
+ """Model for generating realistic annotation timing.
+
+ Supports multiple distribution types:
+ - uniform: Random time uniformly distributed between min and max
+ - normal: Gaussian distribution with configurable mean and std
+ - exponential: Exponential distribution for realistic variability
+ """
+
+ def __init__(self, config: TimingConfig):
+ """Initialize timing model.
+
+ Args:
+ config: TimingConfig with distribution parameters
+ """
+ self.config = config
+
+ def get_annotation_time(self) -> float:
+ """Generate annotation time based on configured distribution.
+
+ Returns:
+ Annotation time in seconds
+ """
+ if self.config.distribution == "uniform":
+ return random.uniform(
+ self.config.annotation_time_min, self.config.annotation_time_max
+ )
+
+ elif self.config.distribution == "normal":
+ time_val = random.gauss(
+ self.config.annotation_time_mean, self.config.annotation_time_std
+ )
+ # Clamp to min/max
+ return max(
+ self.config.annotation_time_min,
+ min(self.config.annotation_time_max, time_val),
+ )
+
+ elif self.config.distribution == "exponential":
+ # Exponential distribution with mean at annotation_time_mean
+ rate = 1.0 / self.config.annotation_time_mean
+ time_val = random.expovariate(rate)
+ # Clamp to min/max
+ return max(
+ self.config.annotation_time_min,
+ min(self.config.annotation_time_max, time_val),
+ )
+
+ # Default fallback
+ return self.config.annotation_time_mean
+
+ def get_fast_response_time(self) -> float:
+ """Generate a suspiciously fast response time.
+
+ Used for testing quality control fast-response detection.
+
+ Returns:
+ Fast response time in seconds (below threshold)
+ """
+ return random.uniform(0.1, self.config.fast_response_threshold * 0.9)
+
+ def should_respond_fast(self, fast_response_rate: float) -> bool:
+ """Determine if this should be a fast response.
+
+ Args:
+ fast_response_rate: Probability of fast response (0-1)
+
+ Returns:
+ True if this should be a fast response
+ """
+ return random.random() < fast_response_rate
+
+ def wait(self, duration: Optional[float] = None) -> float:
+ """Wait for the specified or generated duration.
+
+ Args:
+ duration: Specific duration in seconds, or None to generate
+
+ Returns:
+ The actual duration waited
+ """
+ if duration is None:
+ duration = self.get_annotation_time()
+ time.sleep(duration)
+ return duration
+
+ def get_response_time(self, fast_response_rate: float = 0.0) -> float:
+ """Get response time, possibly fast for QC testing.
+
+ Args:
+ fast_response_rate: Probability of suspiciously fast response
+
+ Returns:
+ Response time in seconds
+ """
+ if self.should_respond_fast(fast_response_rate):
+ return self.get_fast_response_time()
+ return self.get_annotation_time()
+
+
+class NoWaitTimingModel(TimingModel):
+ """Timing model that records times but doesn't wait.
+
+ Useful for fast testing where we want to track timing statistics
+ but don't want to actually delay execution.
+ """
+
+ def wait(self, duration: Optional[float] = None) -> float:
+ """Record but don't actually wait.
+
+ Args:
+ duration: Duration to record (or generate)
+
+ Returns:
+ The duration that would have been waited
+ """
+ if duration is None:
+ duration = self.get_annotation_time()
+ return duration
diff --git a/potato/simulator/user_simulator.py b/potato/simulator/user_simulator.py
new file mode 100644
index 0000000000000000000000000000000000000000..b5f5940873f4735637876e5c56434286afbc6abc
--- /dev/null
+++ b/potato/simulator/user_simulator.py
@@ -0,0 +1,558 @@
+"""
+Core user simulator class.
+
+This module provides the SimulatedUser class that simulates a single
+annotator interacting with the Potato annotation platform via its API.
+"""
+
+import logging
+import random
+from typing import Dict, List, Any, Optional
+from dataclasses import dataclass, field
+from datetime import datetime
+import requests
+
+from .config import (
+ UserConfig,
+ SimulatorConfig,
+ CompetenceLevel,
+ AnnotationStrategyType,
+ TimingConfig,
+ InteractiveConfig,
+)
+from .competence_profiles import CompetenceProfile, create_competence_profile
+from .annotation_strategies import AnnotationStrategy, create_strategy
+from .timing_models import TimingModel, NoWaitTimingModel
+from .interactive_runner import InteractiveSessionRunner
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AnnotationRecord:
+ """Record of a single annotation submission.
+
+ Attributes:
+ instance_id: ID of the annotated instance
+ schema_name: Name of the annotation schema
+ annotation: The annotation data submitted
+ response_time: Time taken to annotate (seconds)
+ timestamp: When the annotation was submitted
+ was_attention_check: Whether this was an attention check item
+ attention_check_passed: Result of attention check (if applicable)
+ was_gold_standard: Whether this was a gold standard item
+ gold_standard_correct: Whether gold standard was answered correctly
+ """
+
+ instance_id: str
+ schema_name: str
+ annotation: Dict[str, Any]
+ response_time: float
+ timestamp: datetime
+ was_attention_check: bool = False
+ attention_check_passed: Optional[bool] = None
+ was_gold_standard: bool = False
+ gold_standard_correct: Optional[bool] = None
+
+
+@dataclass
+class UserSimulationResult:
+ """Results from a user simulation session.
+
+ Attributes:
+ user_id: ID of the simulated user
+ annotations: List of annotation records
+ total_time: Total simulation time in seconds
+ attention_checks_passed: Number of passed attention checks
+ attention_checks_failed: Number of failed attention checks
+ gold_standard_correct: Number of correct gold standard answers
+ gold_standard_incorrect: Number of incorrect gold standard answers
+ errors: List of error messages encountered
+ start_time: When simulation started
+ end_time: When simulation ended
+ was_blocked: Whether user was blocked by quality control
+ """
+
+ user_id: str
+ annotations: List[AnnotationRecord] = field(default_factory=list)
+ total_time: float = 0.0
+ attention_checks_passed: int = 0
+ attention_checks_failed: int = 0
+ gold_standard_correct: int = 0
+ gold_standard_incorrect: int = 0
+ errors: List[str] = field(default_factory=list)
+ start_time: Optional[datetime] = None
+ end_time: Optional[datetime] = None
+ was_blocked: bool = False
+
+
+class SimulatedUser:
+ """Simulates a single user annotating items via the Potato API.
+
+ The SimulatedUser handles:
+ - Authentication (login/registration)
+ - Fetching annotation items
+ - Generating annotations based on strategy
+ - Submitting annotations
+ - Navigating between items
+ - Tracking quality control results
+ """
+
+ def __init__(
+ self,
+ user_config: UserConfig,
+ server_url: str,
+ gold_standards: Optional[Dict[str, Dict[str, Any]]] = None,
+ simulate_wait: bool = False,
+ attention_check_fail_rate: float = 0.0,
+ respond_fast_rate: float = 0.0,
+ interactive_config: Optional[InteractiveConfig] = None,
+ ):
+ """Initialize simulated user.
+
+ Args:
+ user_config: Configuration for this user
+ server_url: Base URL of the Potato server
+ gold_standards: Optional gold standard answers keyed by instance_id
+ simulate_wait: Whether to actually wait between annotations
+ attention_check_fail_rate: Rate at which to fail attention checks
+ respond_fast_rate: Rate of suspiciously fast responses
+ """
+ self.config = user_config
+ self.server_url = server_url.rstrip("/")
+ self.gold_standards = gold_standards or {}
+ self.attention_check_fail_rate = attention_check_fail_rate
+ self.respond_fast_rate = respond_fast_rate
+
+ # Initialize components
+ self.competence = create_competence_profile(user_config.competence)
+ self.strategy = self._create_strategy()
+
+ # Create timing model based on simulate_wait setting
+ if simulate_wait:
+ self.timing = TimingModel(user_config.timing)
+ else:
+ self.timing = NoWaitTimingModel(user_config.timing)
+
+ # Session and state
+ self.session = requests.Session()
+ self.logged_in = False
+ self.current_instance_id: Optional[str] = None
+ self.schemas: List[Dict[str, Any]] = []
+
+ # Optional interactive_chat driver
+ self.interactive_runner: Optional[InteractiveSessionRunner] = None
+ if interactive_config and interactive_config.enabled:
+ self.interactive_runner = InteractiveSessionRunner(
+ interactive_config, server_url
+ )
+
+ # Results tracking
+ self.result = UserSimulationResult(user_id=user_config.user_id)
+
+ def _create_strategy(self) -> AnnotationStrategy:
+ """Create the annotation strategy for this user.
+
+ Returns:
+ AnnotationStrategy instance
+ """
+ return create_strategy(
+ strategy_type=self.config.strategy,
+ llm_config=self.config.llm_config,
+ biased_config=self.config.biased_config,
+ pattern_config=self.config.pattern_config,
+ agent_config=self.config.agent_config,
+ user_id=self.config.user_id,
+ )
+
+ def login(self) -> bool:
+ """Login or register the simulated user.
+
+ Attempts to login first, then registers if login fails.
+
+ Returns:
+ True if authentication successful
+ """
+ password = "simulated_password_123"
+
+ try:
+ # Try to register first (in case user doesn't exist)
+ response = self.session.post(
+ f"{self.server_url}/register",
+ data={
+ "action": "signup",
+ "email": self.config.user_id,
+ "pass": password,
+ },
+ allow_redirects=True,
+ timeout=30,
+ )
+
+ # Now try to login
+ response = self.session.post(
+ f"{self.server_url}/auth",
+ data={
+ "action": "login",
+ "email": self.config.user_id,
+ "pass": password,
+ },
+ allow_redirects=True,
+ timeout=30,
+ )
+
+ # Check if we're logged in by trying to access annotate page
+ check_response = self.session.get(
+ f"{self.server_url}/annotate",
+ allow_redirects=False,
+ timeout=30,
+ )
+
+ # If we get redirected to login, auth failed
+ if check_response.status_code == 302:
+ location = check_response.headers.get("Location", "")
+ if "auth" in location or "login" in location:
+ logger.warning(f"Login failed for {self.config.user_id}")
+ self.result.errors.append("Login failed - redirected to auth")
+ return False
+
+ self.logged_in = True
+ logger.debug(f"User {self.config.user_id} logged in successfully")
+ return True
+
+ except requests.exceptions.RequestException as e:
+ logger.error(f"Login failed for {self.config.user_id}: {e}")
+ self.result.errors.append(f"Login failed: {e}")
+ return False
+
+ def get_current_instance(self) -> Optional[Dict[str, Any]]:
+ """Get the current instance to annotate.
+
+ Returns:
+ Instance data dict or None if unavailable
+ """
+ try:
+ response = self.session.get(
+ f"{self.server_url}/api/current_instance",
+ timeout=30,
+ )
+
+ if response.status_code == 200:
+ data = response.json()
+ self.current_instance_id = data.get("instance_id")
+
+ # Get the actual text content
+ if self.current_instance_id:
+ text_response = self.session.get(
+ f"{self.server_url}/api/spans/{self.current_instance_id}",
+ timeout=30,
+ )
+ if text_response.status_code == 200:
+ text_data = text_response.json()
+ data["text"] = text_data.get("text", "")
+
+ return data
+
+ elif response.status_code == 404:
+ logger.info(f"No more instances for {self.config.user_id}")
+ return None
+ else:
+ logger.warning(
+ f"Failed to get instance: {response.status_code} - {response.text}"
+ )
+ return None
+
+ except requests.exceptions.RequestException as e:
+ logger.error(f"Failed to get current instance: {e}")
+ self.result.errors.append(f"Get instance failed: {e}")
+ return None
+
+ def get_schemas(self) -> List[Dict[str, Any]]:
+ """Get annotation schemas from the server.
+
+ Returns:
+ List of schema definitions
+ """
+ try:
+ response = self.session.get(
+ f"{self.server_url}/api/schemas",
+ timeout=30,
+ )
+
+ if response.status_code == 200:
+ data = response.json()
+ # Handle both list and dict formats
+ if isinstance(data, dict):
+ if "schemas" in data:
+ self.schemas = (
+ list(data["schemas"].values())
+ if isinstance(data["schemas"], dict)
+ else data["schemas"]
+ )
+ else:
+ self.schemas = list(data.values())
+ else:
+ self.schemas = data
+ return self.schemas
+
+ logger.warning(f"Failed to get schemas: {response.status_code}")
+ return []
+
+ except requests.exceptions.RequestException as e:
+ logger.error(f"Failed to get schemas: {e}")
+ self.result.errors.append(f"Get schemas failed: {e}")
+ return []
+
+ def generate_annotations(self, instance: Dict[str, Any]) -> Dict[str, Any]:
+ """Generate annotations for all schemas.
+
+ Args:
+ instance: Instance data including text
+
+ Returns:
+ Combined annotation dictionary for all schemas
+ """
+ instance_id = instance.get("instance_id")
+ gold_answer = self.gold_standards.get(instance_id)
+
+ # Attach the full schema set so batching strategies (e.g. AgentSimulatorStrategy)
+ # can build a single multi-schema prompt per instance. Other strategies
+ # ignore the extra key.
+ instance = dict(instance)
+ instance["__all_schemas__"] = self.schemas
+
+ all_annotations = {}
+
+ for schema in self.schemas:
+ schema_name = schema.get("name")
+ schema_gold = None
+ if gold_answer:
+ schema_gold = {schema_name: gold_answer.get(schema_name)}
+
+ annotation = self.strategy.generate_annotation(
+ instance, schema, self.competence, schema_gold
+ )
+
+ all_annotations.update(annotation)
+
+ return all_annotations
+
+ def submit_annotation(
+ self,
+ instance_id: str,
+ annotations: Dict[str, Any],
+ response_time: float,
+ ) -> bool:
+ """Submit annotations for an instance.
+
+ Args:
+ instance_id: ID of the instance
+ annotations: Annotation data to submit
+ response_time: Time taken to annotate
+
+ Returns:
+ True if submission successful
+ """
+ try:
+ payload = {
+ "instance_id": instance_id,
+ "annotations": annotations,
+ "span_annotations": [],
+ "client_timestamp": datetime.now().isoformat(),
+ }
+
+ response = self.session.post(
+ f"{self.server_url}/updateinstance",
+ json=payload,
+ timeout=30,
+ )
+
+ if response.status_code == 200:
+ result_data = response.json()
+
+ # Create annotation record
+ record = AnnotationRecord(
+ instance_id=instance_id,
+ schema_name=",".join(annotations.keys()),
+ annotation=annotations,
+ response_time=response_time,
+ timestamp=datetime.now(),
+ )
+
+ # Check for quality control results
+ if "qc_result" in result_data:
+ qc_result = result_data["qc_result"]
+ qc_type = qc_result.get("type")
+
+ if qc_type == "attention_check":
+ record.was_attention_check = True
+ record.attention_check_passed = qc_result.get("passed", False)
+ if record.attention_check_passed:
+ self.result.attention_checks_passed += 1
+ else:
+ self.result.attention_checks_failed += 1
+
+ elif qc_type == "gold_standard":
+ record.was_gold_standard = True
+ record.gold_standard_correct = qc_result.get("correct", False)
+ if record.gold_standard_correct:
+ self.result.gold_standard_correct += 1
+ else:
+ self.result.gold_standard_incorrect += 1
+
+ # Check for blocking
+ if result_data.get("status") == "blocked":
+ self.result.was_blocked = True
+ logger.info(f"User {self.config.user_id} was blocked")
+
+ self.result.annotations.append(record)
+ return True
+
+ else:
+ logger.warning(
+ f"Annotation submission failed: {response.status_code} - {response.text}"
+ )
+ self.result.errors.append(f"Submit failed: {response.status_code}")
+ return False
+
+ except requests.exceptions.RequestException as e:
+ logger.error(f"Submit annotation failed: {e}")
+ self.result.errors.append(f"Submit failed: {e}")
+ return False
+
+ def navigate_next(self) -> bool:
+ """Navigate to the next instance.
+
+ Returns:
+ True if navigation successful
+ """
+ try:
+ # POST to /annotate with action=next_instance
+ response = self.session.post(
+ f"{self.server_url}/annotate",
+ data={"action": "next_instance"},
+ timeout=30,
+ )
+
+ return response.status_code in [200, 302]
+
+ except requests.exceptions.RequestException as e:
+ logger.error(f"Navigate next failed: {e}")
+ self.result.errors.append(f"Navigate failed: {e}")
+ return False
+
+ def run_simulation(
+ self, max_annotations: Optional[int] = None
+ ) -> UserSimulationResult:
+ """Run the full simulation for this user.
+
+ Args:
+ max_annotations: Maximum number of annotations (optional)
+
+ Returns:
+ UserSimulationResult with all tracking data
+ """
+ self.result.start_time = datetime.now()
+ max_ann = max_annotations if max_annotations is not None else self.config.max_annotations
+ annotation_count = 0
+
+ try:
+ # Login
+ if not self.login():
+ logger.warning(f"User {self.config.user_id} failed to login")
+ return self.result
+
+ # Get schemas
+ if not self.get_schemas():
+ logger.warning(f"User {self.config.user_id} failed to get schemas")
+ self.result.errors.append("Failed to get schemas")
+
+ # Main annotation loop
+ while True:
+ # Check if blocked
+ if self.result.was_blocked:
+ logger.info(f"User {self.config.user_id} is blocked, stopping")
+ break
+
+ # Check annotation limit
+ if max_ann is not None and annotation_count >= max_ann:
+ logger.debug(
+ f"User {self.config.user_id} reached annotation limit ({max_ann})"
+ )
+ break
+
+ # Get current instance
+ instance = self.get_current_instance()
+ if not instance or not instance.get("instance_id"):
+ logger.info(f"No more instances for {self.config.user_id}")
+ break
+
+ # If an interactive_chat session is configured, drive the
+ # chat first so the conversation field is populated before
+ # the rating strategy reads it.
+ if self.interactive_runner is not None:
+ instance_data = instance.get("data") or {}
+ task_text = (
+ instance_data.get("task_description")
+ or instance_data.get("text")
+ or instance.get("text", "")
+ )
+ chat_result = self.interactive_runner.run(
+ self.session,
+ instance.get("instance_id"),
+ task_text,
+ )
+ if chat_result.error:
+ self.result.errors.append(
+ f"interactive: {chat_result.error}"
+ )
+ # Re-fetch the instance so its data reflects the chat
+ refreshed = self.get_current_instance()
+ if refreshed and refreshed.get("instance_id") == instance.get("instance_id"):
+ instance = refreshed
+ else:
+ # Server moved on; fall back to using the in-memory
+ # conversation we just collected.
+ instance.setdefault("data", {})
+ instance["data"]["conversation"] = chat_result.conversation
+
+ # Generate timing
+ response_time = self.timing.get_response_time(self.respond_fast_rate)
+
+ # Wait if configured (NoWaitTimingModel skips actual waiting)
+ self.timing.wait(response_time)
+
+ # Generate annotations
+ annotations = self.generate_annotations(instance)
+
+ # Submit
+ if self.submit_annotation(
+ instance.get("instance_id"),
+ annotations,
+ response_time,
+ ):
+ annotation_count += 1
+ logger.debug(
+ f"User {self.config.user_id} annotated {annotation_count} items"
+ )
+
+ # Navigate to next
+ if not self.navigate_next():
+ logger.debug(f"User {self.config.user_id} navigation failed")
+ break
+
+ except Exception as e:
+ logger.error(f"Simulation error for {self.config.user_id}: {e}")
+ self.result.errors.append(f"Simulation error: {e}")
+
+ finally:
+ self.result.end_time = datetime.now()
+ self.result.total_time = (
+ self.result.end_time - self.result.start_time
+ ).total_seconds()
+
+ logger.info(
+ f"User {self.config.user_id} completed: "
+ f"{len(self.result.annotations)} annotations in {self.result.total_time:.1f}s"
+ )
+
+ return self.result
diff --git a/potato/solo_mode/__init__.py b/potato/solo_mode/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..00dbdba06f4d9ace794d947189c99a3737bfe531
--- /dev/null
+++ b/potato/solo_mode/__init__.py
@@ -0,0 +1,72 @@
+"""
+Solo Mode Module
+
+This module provides human-LLM collaborative annotation for single annotators.
+Solo Mode enables efficient dataset labeling through:
+
+1. Prompt synthesis from task descriptions
+2. Edge case generation and labeling
+3. Parallel human-LLM annotation with disagreement resolution
+4. Uncertainty-based instance ordering
+5. Progressive validation and autonomous completion
+
+Key Components:
+- SoloModeManager: Central orchestrator for Solo Mode workflow
+- SoloPhase: Enum defining workflow phases
+- PromptManager: Prompt synthesis, versioning, and revision
+- UncertaintyEstimator: Pluggable uncertainty estimation strategies
+- InstanceSelector: Weighted instance selection for human review
+- DisagreementDetector: Type-specific human-LLM disagreement detection
+- ValidationTracker: Agreement metrics and thresholds
+"""
+
+from .config import SoloModeConfig, parse_solo_mode_config
+from .phase_controller import SoloPhase, SoloPhaseController
+from .manager import (
+ SoloModeManager,
+ init_solo_mode_manager,
+ get_solo_mode_manager,
+ clear_solo_mode_manager,
+)
+from .prompt_manager import PromptManager, PromptRevision
+from .instance_selector import InstanceSelector, SelectionWeights
+from .llm_labeler import LLMLabelingThread, LabelingResult
+from .disagreement_resolver import DisagreementDetector
+from .validation_tracker import ValidationTracker, AgreementMetrics, ValidationSample
+from .edge_case_synthesizer import EdgeCaseSynthesizer, EdgeCase
+from .prompt_optimizer import PromptOptimizer, OptimizationResult
+
+__all__ = [
+ # Config
+ 'SoloModeConfig',
+ 'parse_solo_mode_config',
+ # Phase control
+ 'SoloPhase',
+ 'SoloPhaseController',
+ # Manager
+ 'SoloModeManager',
+ 'init_solo_mode_manager',
+ 'get_solo_mode_manager',
+ 'clear_solo_mode_manager',
+ # Prompt management
+ 'PromptManager',
+ 'PromptRevision',
+ # Instance selection
+ 'InstanceSelector',
+ 'SelectionWeights',
+ # LLM labeling
+ 'LLMLabelingThread',
+ 'LabelingResult',
+ # Disagreement detection (authoritative tracking lives on SoloModeManager)
+ 'DisagreementDetector',
+ # Validation tracking
+ 'ValidationTracker',
+ 'AgreementMetrics',
+ 'ValidationSample',
+ # Edge case synthesis
+ 'EdgeCaseSynthesizer',
+ 'EdgeCase',
+ # Prompt optimization
+ 'PromptOptimizer',
+ 'OptimizationResult',
+]
diff --git a/potato/solo_mode/confidence_router.py b/potato/solo_mode/confidence_router.py
new file mode 100644
index 0000000000000000000000000000000000000000..d43bc45d14d355f077b20637b3eff8ccff3df7b3
--- /dev/null
+++ b/potato/solo_mode/confidence_router.py
@@ -0,0 +1,278 @@
+"""
+Confidence Router for Solo Mode
+
+Implements cascaded confidence escalation: cheap model -> expensive model -> human.
+Each tier has a confidence threshold; if the LLM's confidence is below the threshold,
+the instance escalates to the next tier. If all tiers are exhausted, the instance
+is routed to a human annotator.
+"""
+
+import logging
+import threading
+import time
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TierStats:
+ """Per-tier statistics for confidence routing."""
+ name: str = ""
+ instances_attempted: int = 0
+ instances_accepted: int = 0
+ instances_escalated: int = 0
+ instances_errored: int = 0
+ total_confidence: float = 0.0
+ total_latency_ms: float = 0.0
+
+ @property
+ def avg_confidence(self) -> float:
+ if self.instances_accepted == 0:
+ return 0.0
+ return self.total_confidence / self.instances_accepted
+
+ @property
+ def acceptance_rate(self) -> float:
+ if self.instances_attempted == 0:
+ return 0.0
+ return self.instances_accepted / self.instances_attempted
+
+ @property
+ def avg_latency_ms(self) -> float:
+ if self.instances_attempted == 0:
+ return 0.0
+ return self.total_latency_ms / self.instances_attempted
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ 'name': self.name,
+ 'instances_attempted': self.instances_attempted,
+ 'instances_accepted': self.instances_accepted,
+ 'instances_escalated': self.instances_escalated,
+ 'instances_errored': self.instances_errored,
+ 'avg_confidence': round(self.avg_confidence, 4),
+ 'acceptance_rate': round(self.acceptance_rate, 4),
+ 'avg_latency_ms': round(self.avg_latency_ms, 1),
+ }
+
+
+@dataclass
+class RoutingResult:
+ """Result of routing an instance through the confidence cascade."""
+ instance_id: str
+ accepted: bool = False
+ routed_to_human: bool = False
+ tier_index: int = -1
+ tier_name: str = ""
+ labeling_result: Any = None # Optional[LabelingResult]
+ attempts: List[Dict[str, Any]] = field(default_factory=list)
+
+
+class ConfidenceRouter:
+ """
+ Cascaded confidence escalation router.
+
+ Routes instances through tiers of LLM models with decreasing
+ confidence thresholds. If no tier accepts the instance, it is
+ routed to a human annotator.
+ """
+
+ def __init__(
+ self,
+ routing_config,
+ label_fn: Callable,
+ endpoint_factory: Callable,
+ ):
+ """
+ Initialize the confidence router.
+
+ Args:
+ routing_config: ConfidenceRoutingConfig instance
+ label_fn: Function with signature (instance_id, text, schema_name, endpoint) -> LabelingResult
+ endpoint_factory: Function with signature (ModelConfig) -> endpoint
+ """
+ self._config = routing_config
+ self._label_fn = label_fn
+ self._endpoint_factory = endpoint_factory
+ self._lock = threading.Lock()
+
+ # Per-tier stats
+ self._tier_stats: List[TierStats] = [
+ TierStats(name=tier.name or f"tier_{i}")
+ for i, tier in enumerate(routing_config.tiers)
+ ]
+
+ # Lazy endpoint cache per tier
+ self._endpoints: List[Optional[Any]] = [None] * len(routing_config.tiers)
+
+ # Global counters
+ self._human_routed_count = 0
+ self._total_routed = 0
+
+ def _get_tier_endpoint(self, tier_index: int):
+ """Get or create the endpoint for a tier."""
+ if self._endpoints[tier_index] is not None:
+ return self._endpoints[tier_index]
+
+ tier = self._config.tiers[tier_index]
+ try:
+ endpoint = self._endpoint_factory(tier.model)
+ self._endpoints[tier_index] = endpoint
+ return endpoint
+ except Exception as e:
+ logger.warning(
+ f"Failed to create endpoint for tier {tier_index} "
+ f"({tier.name}): {e}"
+ )
+ return None
+
+ def route_instance(
+ self,
+ instance_id: str,
+ text: str,
+ schema_name: str,
+ ) -> RoutingResult:
+ """
+ Route an instance through the confidence cascade.
+
+ For each tier:
+ 1. Get/create the endpoint
+ 2. Call label_fn with the endpoint
+ 3. Check confidence vs threshold
+ 4. If confidence >= threshold -> accept
+ 5. If confidence < threshold -> escalate
+ 6. If error -> skip tier
+ 7. If all tiers exhausted -> route to human
+ """
+ result = RoutingResult(instance_id=instance_id)
+
+ for i, tier in enumerate(self._config.tiers):
+ endpoint = self._get_tier_endpoint(i)
+ if endpoint is None:
+ attempt = {
+ 'tier_index': i,
+ 'tier_name': tier.name,
+ 'error': 'Failed to create endpoint',
+ }
+ result.attempts.append(attempt)
+ with self._lock:
+ self._tier_stats[i].instances_attempted += 1
+ self._tier_stats[i].instances_errored += 1
+ continue
+
+ start_ms = time.monotonic() * 1000
+ try:
+ labeling_result = self._label_fn(
+ instance_id, text, schema_name, endpoint
+ )
+ except Exception as e:
+ elapsed_ms = time.monotonic() * 1000 - start_ms
+ attempt = {
+ 'tier_index': i,
+ 'tier_name': tier.name,
+ 'error': str(e),
+ 'latency_ms': round(elapsed_ms, 1),
+ }
+ result.attempts.append(attempt)
+ with self._lock:
+ self._tier_stats[i].instances_attempted += 1
+ self._tier_stats[i].instances_errored += 1
+ self._tier_stats[i].total_latency_ms += elapsed_ms
+ continue
+
+ elapsed_ms = time.monotonic() * 1000 - start_ms
+
+ if labeling_result is None or labeling_result.error:
+ error_msg = (
+ labeling_result.error if labeling_result else 'No result'
+ )
+ attempt = {
+ 'tier_index': i,
+ 'tier_name': tier.name,
+ 'error': error_msg,
+ 'latency_ms': round(elapsed_ms, 1),
+ }
+ result.attempts.append(attempt)
+ with self._lock:
+ self._tier_stats[i].instances_attempted += 1
+ self._tier_stats[i].instances_errored += 1
+ self._tier_stats[i].total_latency_ms += elapsed_ms
+ continue
+
+ confidence = labeling_result.confidence
+ attempt = {
+ 'tier_index': i,
+ 'tier_name': tier.name,
+ 'confidence': confidence,
+ 'threshold': tier.confidence_threshold,
+ 'latency_ms': round(elapsed_ms, 1),
+ }
+
+ if confidence >= tier.confidence_threshold:
+ # Accepted at this tier
+ attempt['accepted'] = True
+ result.attempts.append(attempt)
+ result.accepted = True
+ result.tier_index = i
+ result.tier_name = tier.name
+ result.labeling_result = labeling_result
+
+ with self._lock:
+ self._tier_stats[i].instances_attempted += 1
+ self._tier_stats[i].instances_accepted += 1
+ self._tier_stats[i].total_confidence += confidence
+ self._tier_stats[i].total_latency_ms += elapsed_ms
+ self._total_routed += 1
+ return result
+ else:
+ # Escalate to next tier
+ attempt['accepted'] = False
+ result.attempts.append(attempt)
+ # Keep the best result so far in case all tiers fail
+ if (
+ result.labeling_result is None
+ or confidence > result.labeling_result.confidence
+ ):
+ result.labeling_result = labeling_result
+ result.tier_index = i
+ result.tier_name = tier.name
+
+ with self._lock:
+ self._tier_stats[i].instances_attempted += 1
+ self._tier_stats[i].instances_escalated += 1
+ self._tier_stats[i].total_latency_ms += elapsed_ms
+
+ # All tiers exhausted -> route to human
+ result.routed_to_human = True
+ result.accepted = False
+ with self._lock:
+ self._human_routed_count += 1
+ self._total_routed += 1
+
+ return result
+
+ def get_stats(self) -> Dict[str, Any]:
+ """Get routing statistics."""
+ with self._lock:
+ return {
+ 'enabled': True,
+ 'num_tiers': len(self._config.tiers),
+ 'tiers': [s.to_dict() for s in self._tier_stats],
+ 'human_routed_count': self._human_routed_count,
+ 'total_routed': self._total_routed,
+ }
+
+ def reset_stats(self) -> None:
+ """Reset all statistics counters."""
+ with self._lock:
+ for stats in self._tier_stats:
+ stats.instances_attempted = 0
+ stats.instances_accepted = 0
+ stats.instances_escalated = 0
+ stats.instances_errored = 0
+ stats.total_confidence = 0.0
+ stats.total_latency_ms = 0.0
+ self._human_routed_count = 0
+ self._total_routed = 0
diff --git a/potato/solo_mode/config.py b/potato/solo_mode/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..3a418c70fec6dc17bd5b50f765daf450dfc9f6a7
--- /dev/null
+++ b/potato/solo_mode/config.py
@@ -0,0 +1,526 @@
+"""
+Solo Mode Configuration
+
+This module defines the configuration dataclass and parsing logic for Solo Mode.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ModelConfig:
+ """Configuration for an LLM endpoint."""
+ endpoint_type: str # 'anthropic', 'openai', 'ollama', etc.
+ model: str
+ api_key: Optional[str] = None
+ base_url: Optional[str] = None
+ max_tokens: int = 1000
+ temperature: float = 0.1
+ think: Optional[bool] = None # None = use endpoint default, True/False = override
+ timeout: int = 60 # Request timeout in seconds (increase for thinking models)
+
+ def to_endpoint_config(self, temperature_override: Optional[float] = None) -> Dict[str, Any]:
+ """Build the full endpoint config dict for AIEndpointFactory.
+
+ This is the single place that builds the config dict passed to
+ AIEndpointFactory.create_endpoint(). All solo mode components
+ should use this instead of manually constructing the dict.
+
+ Args:
+ temperature_override: Override the model's default temperature.
+
+ Returns:
+ Dict ready for AIEndpointFactory.create_endpoint()
+ """
+ ai_config = {
+ 'model': self.model,
+ 'max_tokens': self.max_tokens,
+ 'temperature': temperature_override if temperature_override is not None else self.temperature,
+ }
+ if self.api_key:
+ ai_config['api_key'] = self.api_key
+ if self.base_url:
+ ai_config['base_url'] = self.base_url
+ if self.think is not None:
+ ai_config['think'] = self.think
+ if self.timeout != 60:
+ ai_config['timeout'] = self.timeout
+ return {
+ 'ai_support': {
+ 'enabled': True,
+ 'endpoint_type': self.endpoint_type,
+ 'ai_config': ai_config,
+ }
+ }
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Convert to dictionary for AI endpoint factory (legacy format)."""
+ ec = self.to_endpoint_config()
+ return {
+ 'endpoint_type': self.endpoint_type,
+ 'ai_config': ec['ai_support']['ai_config'],
+ }
+
+
+@dataclass
+class UncertaintyConfig:
+ """Configuration for uncertainty estimation."""
+ strategy: str = "direct_confidence" # direct_confidence, direct_uncertainty, token_entropy, sampling_diversity
+ # Sampling diversity options
+ num_samples: int = 5
+ sampling_temperature: float = 1.0
+
+
+@dataclass
+class ThresholdConfig:
+ """Threshold configuration for Solo Mode."""
+ end_human_annotation_agreement: float = 0.90
+ minimum_validation_sample: int = 50
+ confidence_low: float = 0.5
+ confidence_high: float = 0.8
+ periodic_review_interval: int = 100
+ # Disagreement thresholds by annotation type
+ likert_tolerance: int = 1 # |human - llm| <= tolerance = agreement
+ multiselect_jaccard_threshold: float = 0.5
+ textbox_embedding_threshold: float = 0.7
+ span_overlap_threshold: float = 0.5
+
+
+@dataclass
+class InstanceSelectionConfig:
+ """Configuration for instance selection weights."""
+ low_confidence_weight: float = 0.4
+ diversity_weight: float = 0.3
+ random_weight: float = 0.2
+ disagreement_weight: float = 0.1
+ edge_case_rule_weight: float = 0.0 # Instances matching edge case rule patterns
+ cartography_weight: float = 0.0 # Instances with high confidence variability
+ llm_predicted_weight: float = 0.0 # Instances with LLM predictions needing human comparison
+
+ def validate(self) -> None:
+ """Validate that weights sum to 1.0."""
+ total = (
+ self.low_confidence_weight +
+ self.diversity_weight +
+ self.random_weight +
+ self.disagreement_weight +
+ self.edge_case_rule_weight +
+ self.cartography_weight +
+ self.llm_predicted_weight
+ )
+ if abs(total - 1.0) > 0.001:
+ logger.warning(
+ f"Instance selection weights sum to {total}, normalizing to 1.0"
+ )
+
+
+@dataclass
+class BatchConfig:
+ """Configuration for batch sizes."""
+ llm_labeling_batch: int = 50
+ max_parallel_labels: int = 200
+
+
+@dataclass
+class PromptOptimizationConfig:
+ """Configuration for automatic prompt optimization."""
+ enabled: bool = True
+ find_smallest_model: bool = True
+ target_accuracy: float = 0.85
+ optimization_interval_seconds: int = 300 # 5 minutes
+ # Optimization objectives weights
+ accuracy_weight: float = 0.7
+ length_weight: float = 0.2
+ consistency_weight: float = 0.1
+
+
+@dataclass
+class EdgeCaseRuleConfig:
+ """Configuration for Co-DETECT-style edge case rule discovery."""
+ enabled: bool = True
+ confidence_threshold: float = 0.75 # Extract rules when confidence below this
+ min_rules_for_clustering: int = 10 # Minimum rules before clustering triggers
+ target_cluster_size: int = 15 # Target items per cluster (Co-DETECT: 10-20)
+ auto_extract_on_labeling: bool = True # Extract rules during LLM labeling
+ reannotation_enabled: bool = True
+ reannotation_confidence_threshold: float = 0.60 # Re-annotate instances below this
+ max_reannotations_per_instance: int = 2 # Prevent infinite loops
+
+
+@dataclass
+class ConfidenceTierConfig:
+ """A single tier in the confidence routing cascade."""
+ model: 'ModelConfig' = None
+ confidence_threshold: float = 0.8 # 0.0-1.0, minimum confidence to accept
+ name: str = "" # e.g. "fast", "strong"
+
+ def __post_init__(self):
+ if self.model is None:
+ self.model = ModelConfig(endpoint_type='openai', model='')
+
+
+@dataclass
+class ConfusionAnalysisConfig:
+ """Configuration for confusion pattern analysis dashboard."""
+ enabled: bool = True
+ min_instances_for_pattern: int = 3
+ max_patterns: int = 20
+ auto_suggest_guidelines: bool = False
+
+
+@dataclass
+class RefinementLoopConfig:
+ """Configuration for the iterative guideline refinement loop."""
+ enabled: bool = True
+ trigger_interval: int = 50 # Check every N human annotations
+ min_improvement: float = 0.02 # Minimum agreement rate improvement to continue
+ max_cycles: int = 5 # Maximum refinement cycles before alerting
+ patience: int = 2 # Cycles without improvement before stopping
+ auto_apply_suggestions: bool = False # Auto-apply LLM guideline suggestions
+ refinement_strategy: str = "focused_edit" # Legacy or new names supported
+
+ # New framework options (used when strategy is validated_*, principle_icl,
+ # hybrid_dual_track, or legacy_append from the refinement registry)
+ validation_split_ratio: float = 0.3 # fraction of disagreements held out
+ eval_sample_size: int = 10 # val instances used to score each candidate
+ num_candidates: int = 3 # candidates proposed per cycle (where applicable)
+ min_val_size: int = 10 # minimum val size before validation-gated refinement runs
+ max_consecutive_failures: int = 2 # stop after N cycles with no improvement
+ dry_run: bool = False # if True, log candidates but don't apply
+ require_approval: bool = False # if True, queue for admin approval before applying
+ min_val_improvement: float = 0.0 # candidate must beat baseline by at least this much (strict=0.0)
+ # Separate temperature for the evaluator pass. Sampling diversity needs
+ # non-zero temperature for confidence, but the validation gate should
+ # measure prompt quality, not sampling variance.
+ eval_temperature: float = 0.0
+ # If True, prefer val instances that have disagreed across โฅ2 labeling
+ # passes (i.e. stable systematic errors) over one-off disagreements
+ # that may be stochastic. Falls back to any disagreement if too few
+ # qualify.
+ prefer_consistent_disagreements: bool = True
+
+
+@dataclass
+class LabelingFunctionConfig:
+ """Configuration for labeling function extraction (ALCHEmist-style)."""
+ enabled: bool = True
+ min_confidence: float = 0.85 # Minimum LLM confidence to consider for extraction
+ min_coverage: int = 3 # Minimum instances a pattern must match
+ max_functions: int = 50 # Maximum labeling functions to maintain
+ auto_extract: bool = True # Auto-extract during labeling
+ vote_threshold: float = 0.5 # Fraction of matching functions needed for label
+
+
+@dataclass
+class ConfidenceRoutingConfig:
+ """Cascaded confidence escalation config."""
+ enabled: bool = False
+ tiers: List['ConfidenceTierConfig'] = field(default_factory=list)
+
+
+@dataclass
+class EmbeddingConfig:
+ """Configuration for embedding model (used for diversity)."""
+ model_name: str = "all-MiniLM-L6-v2"
+
+
+@dataclass
+class SoloModeConfig:
+ """
+ Main configuration dataclass for Solo Mode.
+
+ This contains all settings needed to run Solo Mode including
+ model configurations, thresholds, and feature flags.
+ """
+ enabled: bool = False
+
+ # Models for labeling (tried in order until one succeeds)
+ labeling_models: List[ModelConfig] = field(default_factory=list)
+
+ # Models for prompt revision
+ revision_models: List[ModelConfig] = field(default_factory=list)
+
+ # Embedding configuration
+ embedding: EmbeddingConfig = field(default_factory=EmbeddingConfig)
+
+ # Uncertainty estimation
+ uncertainty: UncertaintyConfig = field(default_factory=UncertaintyConfig)
+
+ # Thresholds
+ thresholds: ThresholdConfig = field(default_factory=ThresholdConfig)
+
+ # Instance selection
+ instance_selection: InstanceSelectionConfig = field(default_factory=InstanceSelectionConfig)
+
+ # Batch sizes
+ batches: BatchConfig = field(default_factory=BatchConfig)
+
+ # Prompt optimization
+ prompt_optimization: PromptOptimizationConfig = field(default_factory=PromptOptimizationConfig)
+
+ # Edge case rule discovery (Co-DETECT-style)
+ edge_case_rules: EdgeCaseRuleConfig = field(default_factory=EdgeCaseRuleConfig)
+
+ # Labeling function extraction (ALCHEmist-style)
+ labeling_functions: LabelingFunctionConfig = field(default_factory=LabelingFunctionConfig)
+
+ # Cascaded confidence routing
+ confidence_routing: ConfidenceRoutingConfig = field(default_factory=ConfidenceRoutingConfig)
+
+ # Confusion analysis dashboard
+ confusion_analysis: ConfusionAnalysisConfig = field(default_factory=ConfusionAnalysisConfig)
+
+ # Iterative guideline refinement loop
+ refinement_loop: RefinementLoopConfig = field(default_factory=RefinementLoopConfig)
+
+ # Output directory for Solo Mode state
+ state_dir: Optional[str] = None
+
+ def validate(self) -> List[str]:
+ """
+ Validate the configuration.
+
+ Returns:
+ List of validation error messages (empty if valid)
+ """
+ errors = []
+
+ if self.enabled:
+ if not self.labeling_models:
+ errors.append("solo_mode.labeling_models is required when solo_mode is enabled")
+
+ if not self.revision_models:
+ # Default to using labeling models for revision
+ logger.info("No revision_models specified, using labeling_models")
+
+ # Validate instance selection weights
+ self.instance_selection.validate()
+
+ # Validate thresholds
+ if not 0 <= self.thresholds.end_human_annotation_agreement <= 1:
+ errors.append("end_human_annotation_agreement must be between 0 and 1")
+
+ if not 0 <= self.thresholds.confidence_low <= 1:
+ errors.append("confidence_low must be between 0 and 1")
+
+ if not 0 <= self.thresholds.confidence_high <= 1:
+ errors.append("confidence_high must be between 0 and 1")
+
+ if self.thresholds.confidence_low >= self.thresholds.confidence_high:
+ errors.append("confidence_low must be less than confidence_high")
+
+ # Validate uncertainty strategy
+ valid_strategies = [
+ 'direct_confidence', 'direct_uncertainty',
+ 'token_entropy', 'sampling_diversity'
+ ]
+ if self.uncertainty.strategy not in valid_strategies:
+ errors.append(f"Invalid uncertainty strategy: {self.uncertainty.strategy}")
+
+ return errors
+
+
+def _parse_model_config(model_data: Dict[str, Any]) -> ModelConfig:
+ """Parse a single model configuration."""
+ # Handle environment variable expansion for API keys
+ api_key = model_data.get('api_key')
+ if api_key and api_key.startswith('${') and api_key.endswith('}'):
+ env_var = api_key[2:-1]
+ api_key = os.environ.get(env_var)
+ if not api_key:
+ logger.warning("Required environment variable for API key is not set")
+
+ return ModelConfig(
+ endpoint_type=model_data.get('endpoint_type', 'openai'),
+ model=model_data.get('model', ''),
+ api_key=api_key,
+ base_url=model_data.get('base_url') or model_data.get('endpoint_url'),
+ max_tokens=model_data.get('max_tokens', 1000),
+ temperature=model_data.get('temperature', 0.1),
+ think=model_data.get('think'), # None = endpoint default, True/False = override
+ timeout=model_data.get('timeout', 60),
+ )
+
+
+def parse_solo_mode_config(config_data: Dict[str, Any]) -> SoloModeConfig:
+ """
+ Parse solo_mode section from application config into SoloModeConfig.
+
+ Args:
+ config_data: Full application configuration dictionary
+
+ Returns:
+ SoloModeConfig instance
+ """
+ sm = config_data.get('solo_mode', {})
+
+ if not sm:
+ return SoloModeConfig(enabled=False)
+
+ # Parse labeling models
+ labeling_models = []
+ for model_data in sm.get('labeling_models', []):
+ labeling_models.append(_parse_model_config(model_data))
+
+ # Parse revision models (default to labeling models if not specified)
+ revision_models = []
+ for model_data in sm.get('revision_models', sm.get('labeling_models', [])):
+ revision_models.append(_parse_model_config(model_data))
+
+ # Parse embedding config
+ emb_data = sm.get('embedding', {})
+ embedding = EmbeddingConfig(
+ model_name=emb_data.get('model_name', 'all-MiniLM-L6-v2')
+ )
+
+ # Parse uncertainty config
+ unc_data = sm.get('uncertainty', {})
+ sampling_data = unc_data.get('sampling_diversity', {})
+ uncertainty = UncertaintyConfig(
+ strategy=unc_data.get('strategy', 'direct_confidence'),
+ num_samples=sampling_data.get('num_samples', 5),
+ sampling_temperature=sampling_data.get('temperature', 1.0),
+ )
+
+ # Parse threshold config
+ thresh_data = sm.get('thresholds', {})
+ thresholds = ThresholdConfig(
+ end_human_annotation_agreement=thresh_data.get('end_human_annotation_agreement', 0.90),
+ minimum_validation_sample=thresh_data.get('minimum_validation_sample', 50),
+ confidence_low=thresh_data.get('confidence_low', 0.5),
+ confidence_high=thresh_data.get('confidence_high', 0.8),
+ periodic_review_interval=thresh_data.get('periodic_review_interval', 100),
+ likert_tolerance=thresh_data.get('likert_tolerance', 1),
+ multiselect_jaccard_threshold=thresh_data.get('multiselect_jaccard_threshold', 0.5),
+ textbox_embedding_threshold=thresh_data.get('textbox_embedding_threshold', 0.7),
+ span_overlap_threshold=thresh_data.get('span_overlap_threshold', 0.5),
+ )
+
+ # Parse instance selection config
+ sel_data = sm.get('instance_selection', {})
+ instance_selection = InstanceSelectionConfig(
+ low_confidence_weight=sel_data.get('low_confidence_weight', 0.4),
+ diversity_weight=sel_data.get('diversity_weight', 0.3),
+ random_weight=sel_data.get('random_weight', 0.2),
+ disagreement_weight=sel_data.get('disagreement_weight', 0.1),
+ edge_case_rule_weight=sel_data.get('edge_case_rule_weight', 0.0),
+ cartography_weight=sel_data.get('cartography_weight', 0.0),
+ llm_predicted_weight=sel_data.get('llm_predicted_weight', 0.0),
+ )
+
+ # Parse batch config
+ batch_data = sm.get('batches', {})
+ batches = BatchConfig(
+ llm_labeling_batch=batch_data.get('llm_labeling_batch', 50),
+ max_parallel_labels=batch_data.get('max_parallel_labels', 200),
+ )
+
+ # Parse prompt optimization config
+ opt_data = sm.get('prompt_optimization', {})
+ prompt_optimization = PromptOptimizationConfig(
+ enabled=opt_data.get('enabled', True),
+ find_smallest_model=opt_data.get('find_smallest_model', True),
+ target_accuracy=opt_data.get('target_accuracy', 0.85),
+ optimization_interval_seconds=opt_data.get('optimization_interval_seconds', 300),
+ accuracy_weight=opt_data.get('accuracy_weight', 0.7),
+ length_weight=opt_data.get('length_weight', 0.2),
+ consistency_weight=opt_data.get('consistency_weight', 0.1),
+ )
+
+ # Parse edge case rule config
+ ecr_data = sm.get('edge_case_rules', {})
+ edge_case_rules = EdgeCaseRuleConfig(
+ enabled=ecr_data.get('enabled', True),
+ confidence_threshold=ecr_data.get('confidence_threshold', 0.75),
+ min_rules_for_clustering=ecr_data.get('min_rules_for_clustering', 10),
+ target_cluster_size=ecr_data.get('target_cluster_size', 15),
+ auto_extract_on_labeling=ecr_data.get('auto_extract_on_labeling', True),
+ reannotation_enabled=ecr_data.get('reannotation_enabled', True),
+ reannotation_confidence_threshold=ecr_data.get('reannotation_confidence_threshold', 0.60),
+ max_reannotations_per_instance=ecr_data.get('max_reannotations_per_instance', 2),
+ )
+
+ # Parse labeling function config
+ lf_data = sm.get('labeling_functions', {})
+ labeling_functions = LabelingFunctionConfig(
+ enabled=lf_data.get('enabled', True),
+ min_confidence=lf_data.get('min_confidence', 0.85),
+ min_coverage=lf_data.get('min_coverage', 3),
+ max_functions=lf_data.get('max_functions', 50),
+ auto_extract=lf_data.get('auto_extract', True),
+ vote_threshold=lf_data.get('vote_threshold', 0.5),
+ )
+
+ # Parse confidence routing config
+ cr_data = sm.get('confidence_routing', {})
+ cr_tiers = []
+ for tier_data in cr_data.get('tiers', []):
+ cr_tiers.append(ConfidenceTierConfig(
+ model=_parse_model_config(tier_data.get('model', {})),
+ confidence_threshold=tier_data.get('confidence_threshold', 0.8),
+ name=tier_data.get('name', ''),
+ ))
+ confidence_routing = ConfidenceRoutingConfig(
+ enabled=cr_data.get('enabled', False),
+ tiers=cr_tiers,
+ )
+
+ # Parse refinement loop config
+ rl_data = sm.get('refinement_loop', {})
+ refinement_loop = RefinementLoopConfig(
+ enabled=rl_data.get('enabled', True),
+ trigger_interval=rl_data.get('trigger_interval', 50),
+ min_improvement=rl_data.get('min_improvement', 0.02),
+ max_cycles=rl_data.get('max_cycles', 5),
+ patience=rl_data.get('patience', 2),
+ auto_apply_suggestions=rl_data.get('auto_apply_suggestions', False),
+ refinement_strategy=rl_data.get('refinement_strategy', 'focused_edit'),
+ validation_split_ratio=rl_data.get('validation_split_ratio', 0.3),
+ eval_sample_size=rl_data.get('eval_sample_size', 10),
+ num_candidates=rl_data.get('num_candidates', 3),
+ min_val_size=rl_data.get('min_val_size', 10),
+ max_consecutive_failures=rl_data.get('max_consecutive_failures', 2),
+ dry_run=rl_data.get('dry_run', False),
+ require_approval=rl_data.get('require_approval', False),
+ min_val_improvement=rl_data.get('min_val_improvement', 0.0),
+ eval_temperature=rl_data.get('eval_temperature', 0.0),
+ prefer_consistent_disagreements=rl_data.get('prefer_consistent_disagreements', True),
+ )
+
+ # Parse confusion analysis config
+ ca_data = sm.get('confusion_analysis', {})
+ confusion_analysis = ConfusionAnalysisConfig(
+ enabled=ca_data.get('enabled', True),
+ min_instances_for_pattern=ca_data.get('min_instances_for_pattern', 3),
+ max_patterns=ca_data.get('max_patterns', 20),
+ auto_suggest_guidelines=ca_data.get('auto_suggest_guidelines', False),
+ )
+
+ # Determine state directory
+ state_dir = sm.get('state_dir')
+ if not state_dir:
+ output_dir = config_data.get('output_annotation_dir', 'annotation_output')
+ state_dir = os.path.join(output_dir, '.solo_mode')
+
+ return SoloModeConfig(
+ enabled=sm.get('enabled', False),
+ labeling_models=labeling_models,
+ revision_models=revision_models,
+ embedding=embedding,
+ uncertainty=uncertainty,
+ thresholds=thresholds,
+ instance_selection=instance_selection,
+ batches=batches,
+ prompt_optimization=prompt_optimization,
+ edge_case_rules=edge_case_rules,
+ labeling_functions=labeling_functions,
+ confidence_routing=confidence_routing,
+ confusion_analysis=confusion_analysis,
+ refinement_loop=refinement_loop,
+ state_dir=state_dir,
+ )
diff --git a/potato/solo_mode/confusion_analyzer.py b/potato/solo_mode/confusion_analyzer.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab1465f16dfda791a83e05e429389a6d5fe4ab40
--- /dev/null
+++ b/potato/solo_mode/confusion_analyzer.py
@@ -0,0 +1,653 @@
+"""
+Confusion Analyzer for Solo Mode
+
+Enriches confusion matrix data with example instances, LLM reasoning,
+and optional root cause / guideline suggestions via LLM.
+"""
+
+import json
+import logging
+import re
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ConfusionExample:
+ """A single instance that contributed to a confusion pattern."""
+ instance_id: str
+ text: str # truncated display text
+ llm_reasoning: Optional[str] = None
+ llm_confidence: Optional[float] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ result = {
+ 'instance_id': self.instance_id,
+ 'text': self.text,
+ }
+ if self.llm_reasoning is not None:
+ result['llm_reasoning'] = self.llm_reasoning
+ if self.llm_confidence is not None:
+ result['llm_confidence'] = self.llm_confidence
+ return result
+
+
+@dataclass
+class ConfusionPattern:
+ """An enriched confusion pattern with examples and optional analysis."""
+ predicted_label: str
+ actual_label: str
+ count: int
+ percent: float
+ examples: List[ConfusionExample] = field(default_factory=list)
+ root_cause: Optional[str] = None
+ guideline_suggestion: Optional[str] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ result = {
+ 'predicted_label': self.predicted_label,
+ 'actual_label': self.actual_label,
+ 'count': self.count,
+ 'percent': self.percent,
+ 'examples': [e.to_dict() for e in self.examples],
+ }
+ if self.root_cause is not None:
+ result['root_cause'] = self.root_cause
+ if self.guideline_suggestion is not None:
+ result['guideline_suggestion'] = self.guideline_suggestion
+ return result
+
+
+class ConfusionAnalyzer:
+ """Analyzes confusion patterns and optionally generates root causes / suggestions.
+
+ Enriches the raw confusion matrix from ValidationTracker with example
+ instances, LLM reasoning, and optional LLM-powered analysis.
+ """
+
+ MAX_TEXT_LENGTH = 200
+ MAX_EXAMPLES_PER_PATTERN = 5
+
+ def __init__(self, app_config: Dict[str, Any], solo_config: Any):
+ self.app_config = app_config
+ self.solo_config = solo_config
+ self._endpoint = None
+
+ def analyze(
+ self,
+ comparison_history: List[Dict[str, Any]],
+ predictions: Dict[str, Dict[str, Any]],
+ text_getter: Optional[Callable[[str], str]] = None,
+ ) -> List[ConfusionPattern]:
+ """Build enriched confusion patterns from comparison history.
+
+ Args:
+ comparison_history: List of comparison dicts with instance_id,
+ human_label, llm_label, agrees fields.
+ predictions: Dict of instance_id -> schema_name -> LLMPrediction.
+ text_getter: Optional callable(instance_id) -> text string.
+
+ Returns:
+ List of ConfusionPattern sorted by count descending.
+ """
+ ca_config = self.solo_config.confusion_analysis
+
+ # Group disagreements by (llm_label, human_label)
+ groups: Dict[Tuple[str, str], List[Dict[str, Any]]] = defaultdict(list)
+ for record in comparison_history:
+ if record.get('agrees'):
+ continue
+ key = (str(record['llm_label']), str(record['human_label']))
+ groups[key].append(record)
+
+ # Filter by minimum instance count
+ patterns = []
+ for (predicted, actual), records in groups.items():
+ if len(records) < ca_config.min_instances_for_pattern:
+ continue
+
+ total_disagreements = sum(
+ 1 for r in comparison_history if not r.get('agrees')
+ )
+ percent = (
+ len(records) / total_disagreements * 100
+ if total_disagreements > 0 else 0.0
+ )
+
+ # Build examples
+ examples = []
+ for record in records[:self.MAX_EXAMPLES_PER_PATTERN]:
+ iid = record['instance_id']
+ text = ''
+ if text_getter is not None:
+ try:
+ raw = text_getter(iid)
+ text = self._truncate(raw)
+ except Exception:
+ text = ''
+
+ # Get reasoning and confidence from predictions
+ reasoning = None
+ confidence = None
+ if iid in predictions:
+ for schema_preds in predictions[iid].values():
+ pred = schema_preds
+ reasoning = (
+ pred.reasoning
+ if hasattr(pred, 'reasoning')
+ else pred.get('reasoning')
+ )
+ confidence = (
+ pred.confidence_score
+ if hasattr(pred, 'confidence_score')
+ else pred.get('confidence_score')
+ )
+ break
+
+ examples.append(ConfusionExample(
+ instance_id=iid,
+ text=text,
+ llm_reasoning=reasoning,
+ llm_confidence=confidence,
+ ))
+
+ patterns.append(ConfusionPattern(
+ predicted_label=predicted,
+ actual_label=actual,
+ count=len(records),
+ percent=round(percent, 1),
+ examples=examples,
+ ))
+
+ # Sort by count descending, limit
+ patterns.sort(key=lambda p: p.count, reverse=True)
+ return patterns[:ca_config.max_patterns]
+
+ def get_confusion_matrix_data(
+ self,
+ confusion_matrix: Dict[Tuple[str, str], int],
+ labels: List[str],
+ label_accuracy: Optional[Dict[str, float]] = None,
+ ) -> Dict[str, Any]:
+ """Build heatmap-ready data from raw confusion matrix.
+
+ Args:
+ confusion_matrix: Dict of (predicted, actual) -> count.
+ labels: All label names.
+ label_accuracy: Optional per-label accuracy dict.
+
+ Returns:
+ Dict with labels, cells, max_count, and label_accuracy.
+ """
+ cells = []
+ max_count = 0
+ for predicted in labels:
+ for actual in labels:
+ count = confusion_matrix.get((predicted, actual), 0)
+ cells.append({
+ 'predicted': predicted,
+ 'actual': actual,
+ 'count': count,
+ })
+ if count > max_count:
+ max_count = count
+
+ return {
+ 'labels': labels,
+ 'cells': cells,
+ 'max_count': max_count,
+ 'label_accuracy': label_accuracy or {},
+ }
+
+ def generate_root_cause(self, pattern: ConfusionPattern) -> Optional[str]:
+ """Use LLM to explain why a confusion pattern occurs.
+
+ Args:
+ pattern: The confusion pattern to analyze.
+
+ Returns:
+ Root cause explanation string, or None if unavailable.
+ """
+ endpoint = self._get_revision_endpoint()
+ if endpoint is None:
+ return None
+
+ examples_text = "\n".join(
+ f"- Instance {e.instance_id}: \"{e.text}\""
+ + (f" (LLM reasoning: {e.llm_reasoning})" if e.llm_reasoning else "")
+ for e in pattern.examples
+ )
+
+ prompt = (
+ f"The LLM predicted \"{pattern.predicted_label}\" when the correct "
+ f"label was \"{pattern.actual_label}\", {pattern.count} times.\n\n"
+ f"Example instances:\n{examples_text}\n\n"
+ f"In 2-3 sentences, explain the most likely root cause of this "
+ f"confusion. What pattern in the text makes these labels hard to "
+ f"distinguish?\n\n"
+ f"Respond with JSON:\n"
+ f'{{"root_cause": ""}}'
+ )
+
+ try:
+ response = endpoint.query(prompt)
+ data = self._parse_json(response)
+ return data.get('root_cause')
+ except Exception as e:
+ logger.warning(f"Root cause generation failed: {e}")
+ return None
+
+ def suggest_guideline(
+ self,
+ pattern: ConfusionPattern,
+ current_prompt: str,
+ ) -> Optional[str]:
+ """Use LLM to suggest a guideline to disambiguate a confusion pattern.
+
+ Args:
+ pattern: The confusion pattern to address.
+ current_prompt: The current annotation prompt text.
+
+ Returns:
+ Guideline suggestion string, or None if unavailable.
+ """
+ endpoint = self._get_revision_endpoint()
+ if endpoint is None:
+ return None
+
+ # Build examples from the confusion pattern
+ examples_text = ""
+ for i, ex in enumerate(pattern.examples[:3]):
+ examples_text += (
+ f" Example {i+1}: \"{ex.text}\"\n"
+ f" Model labeled: {pattern.predicted_label}"
+ )
+ if ex.llm_confidence is not None:
+ examples_text += f" (confidence: {ex.llm_confidence:.0%})"
+ examples_text += f"\n Correct label: {pattern.actual_label}\n"
+ if ex.llm_reasoning:
+ examples_text += f" Model reasoning: {ex.llm_reasoning[:150]}\n"
+
+ prompt = (
+ f"You are helping improve annotation guidelines. The system's LLM "
+ f"annotator keeps confusing \"{pattern.predicted_label}\" with "
+ f"\"{pattern.actual_label}\" ({pattern.count} times).\n\n"
+ f"## Confused Examples\n{examples_text}\n"
+ f"## Current Prompt\n{current_prompt[:1500]}\n\n"
+ f"## Task\n"
+ f"Write a GENERAL, actionable guideline (1-2 sentences) to help correctly "
+ f"distinguish \"{pattern.actual_label}\" from \"{pattern.predicted_label}\".\n\n"
+ f"Requirements:\n"
+ f"- Reference GENERAL linguistic features (sentiment polarity, intent, "
+ f" rhetorical structure, framing) that differentiate the two labels.\n"
+ f"- DO NOT mention specific phrases or quotes from the examples above โ "
+ f" the rule should generalize to unseen cases.\n"
+ f"- DO NOT repeat what the current prompt already says.\n"
+ f"- Focus on the underlying reason the examples were misclassified, not the surface form.\n\n"
+ f"Bad rule (too specific): 'If the text says \"get under your skin\", classify as negative.'\n"
+ f"Good rule (general): 'When the text describes emotional discomfort or unease as a reaction to the content, classify as negative rather than positive.'\n\n"
+ f"Respond with JSON: {{\"suggestion\": \"\"}}"
+ )
+
+ try:
+ # OllamaEndpoint requires an output_format (Pydantic model).
+ # Other endpoints accept just a prompt string.
+ try:
+ from pydantic import BaseModel
+
+ class SuggestionResponse(BaseModel):
+ suggestion: str = ""
+
+ response = endpoint.query(prompt, SuggestionResponse)
+ except TypeError:
+ # Endpoint doesn't require output_format (e.g., OpenAI)
+ response = endpoint.query(prompt)
+
+ data = self._parse_json(response)
+ suggestion = data.get('suggestion')
+ if suggestion:
+ logger.info(
+ f"Generated guideline for {pattern.predicted_label}->"
+ f"{pattern.actual_label}: {suggestion[:100]}"
+ )
+ else:
+ logger.warning(
+ f"No suggestion extracted from LLM response for "
+ f"{pattern.predicted_label}->{pattern.actual_label}: "
+ f"{str(response)[:200]}"
+ )
+ return suggestion
+ except Exception as e:
+ logger.warning(f"Guideline suggestion failed: {e}")
+ return None
+
+ def generate_guidelines_rewrite(
+ self,
+ patterns: List[ConfusionPattern],
+ current_prompt: str,
+ ) -> Optional[List[str]]:
+ """Generate a complete, non-redundant set of guidelines addressing all confusion patterns.
+
+ Instead of generating one rule at a time (which leads to contradictions),
+ this method asks the LLM to produce a coherent set of rules that replaces
+ the existing guidelines section entirely.
+
+ Args:
+ patterns: List of confusion patterns to address (top N by count).
+ current_prompt: The full current annotation prompt.
+
+ Returns:
+ List of guideline strings, or None if generation failed.
+ """
+ endpoint = self._get_revision_endpoint()
+ if endpoint is None:
+ return None
+
+ # Extract existing guidelines from prompt
+ import re as re_mod
+ existing_match = re_mod.search(
+ r'## (?:Refinement |Annotation )?Guidelines\s*\n(.*)',
+ current_prompt, re_mod.DOTALL
+ )
+ existing_guidelines = existing_match.group(1).strip() if existing_match else "None yet"
+
+ # Format confusion patterns with examples
+ patterns_text = ""
+ for i, pattern in enumerate(patterns[:8]):
+ patterns_text += (
+ f"\n{i+1}. Model predicts \"{pattern.predicted_label}\" "
+ f"but correct label is \"{pattern.actual_label}\" "
+ f"({pattern.count} times)\n"
+ )
+ for ex in pattern.examples[:2]:
+ patterns_text += f" Text: \"{ex.text}\"\n"
+ if ex.llm_reasoning:
+ patterns_text += f" Model reasoning: {ex.llm_reasoning[:100]}\n"
+
+ # Extract base prompt (without guidelines section) for context
+ base_prompt = current_prompt
+ if existing_match:
+ base_prompt = current_prompt[:existing_match.start()].strip()
+
+ has_existing = existing_guidelines and existing_guidelines != "None yet"
+
+ prompt = (
+ f"You are improving annotation guidelines for a text classification task.\n\n"
+ f"## Base Task\n{base_prompt[:1000]}\n\n"
+ f"## Existing Guidelines\n{existing_guidelines[:1500]}\n\n"
+ f"## Current Confusion Patterns\n"
+ f"These are errors the model made WITH the guidelines above in place:\n"
+ f"{patterns_text}\n"
+ f"## Instructions\n"
+ + (
+ "You are REFINING existing guidelines, NOT replacing them. Follow these rules strictly:\n\n"
+ "1. START with the existing guidelines above. KEEP all rules that address patterns NOT in the confusion list.\n"
+ "2. For each confusion pattern, check: is there already a rule addressing this label pair?\n"
+ " - YES: The existing rule isn't working. REFINE it with more specific criteria โ do NOT reverse its direction.\n"
+ " - NO: ADD a new rule.\n"
+ "3. NEVER flip an existing rule from 'classify X as A' to 'classify X as B' โ this creates contradictions.\n"
+ "4. Rules must be GENERAL โ do NOT quote specific phrases from the examples. The rule should apply to unseen cases.\n"
+ "5. Output the COMPLETE updated list (existing rules + refinements + new rules).\n"
+ "6. Output at most 8 rules total. If you exceed 8, drop the ones with fewest matching examples.\n\n"
+ if has_existing else
+ "Write 3-8 disambiguation rules. Each rule should:\n"
+ "1. Target a specific confusion pattern above\n"
+ "2. Give GENERAL criteria (sentiment polarity, intent, rhetorical structure, framing)\n"
+ " โ do NOT quote specific phrases from the example texts.\n"
+ "3. NOT contradict other rules in your list\n\n"
+ "Bad (too specific): 'If the text says \"get under your skin\", classify as negative.'\n"
+ "Good (general): 'When the text describes emotional discomfort as a reaction, classify as negative.'\n\n"
+ )
+ + f"Respond with JSON: {{\"guidelines\": [\"rule 1\", \"rule 2\", ...]}}"
+ )
+
+ try:
+ from pydantic import BaseModel
+
+ class GuidelinesResponse(BaseModel):
+ guidelines: List[str] = []
+
+ try:
+ response = endpoint.query(prompt, GuidelinesResponse)
+ except TypeError:
+ response = endpoint.query(prompt)
+
+ data = self._parse_json(response)
+ guidelines = data.get('guidelines', [])
+
+ if guidelines:
+ logger.info(
+ f"[Focused Edit] Generated {len(guidelines)} guidelines "
+ f"for {len(patterns)} confusion patterns"
+ )
+ return guidelines
+ else:
+ # Fallback: try to extract from 'suggestion' key or raw text
+ suggestion = data.get('suggestion')
+ if suggestion:
+ return [suggestion]
+ logger.warning(
+ f"No guidelines extracted from rewrite response: "
+ f"{str(response)[:200]}"
+ )
+ return None
+
+ except Exception as e:
+ logger.warning(f"Guidelines rewrite failed: {e}")
+ return None
+
+ def generate_and_critique_guidelines(
+ self,
+ patterns: List[ConfusionPattern],
+ current_prompt: str,
+ ) -> List[str]:
+ """Two-pass guideline generation: generate candidates, then critique and filter.
+
+ Pass 1: Generate one suggestion per confusion pattern.
+ Pass 2: Evaluate all suggestions together for specificity, consistency,
+ and redundancy. Keep only the best ones.
+
+ Args:
+ patterns: Confusion patterns to address.
+ current_prompt: Current annotation prompt.
+
+ Returns:
+ List of approved guideline strings (may be empty).
+ """
+ endpoint = self._get_revision_endpoint()
+ if endpoint is None:
+ return []
+
+ # Pass 1: Generate candidates
+ candidates = []
+ for pattern in patterns[:5]:
+ suggestion = self.suggest_guideline(pattern, current_prompt)
+ if suggestion:
+ candidates.append({
+ 'pattern': f"{pattern.predicted_label} -> {pattern.actual_label} ({pattern.count}x)",
+ 'suggestion': suggestion,
+ })
+
+ if not candidates:
+ return []
+
+ logger.info(f"[Generator-Critic] Generated {len(candidates)} candidates, running critic...")
+
+ # Pass 2: Critic evaluates
+ # Format each candidate so the rule text is clearly separated from metadata
+ candidates_text = ""
+ for i, c in enumerate(candidates):
+ candidates_text += (
+ f"\n### Candidate {i+1} (addresses pattern: {c['pattern']})\n"
+ f"RULE TEXT: {c['suggestion']}\n"
+ )
+
+ # Extract existing guidelines for contradiction check
+ import re as re_mod
+ existing_match = re_mod.search(
+ r'## (?:Refinement |Annotation )?Guidelines\s*\n(.*)',
+ current_prompt, re_mod.DOTALL
+ )
+ existing_guidelines = existing_match.group(1).strip() if existing_match else ""
+
+ critic_prompt = (
+ f"You are a quality reviewer for annotation guidelines. Be strict.\n\n"
+ f"## Base Annotation Task\n{current_prompt[:800]}\n\n"
+ + (f"## Existing Guidelines (already in production)\n{existing_guidelines[:800]}\n\n" if existing_guidelines else "")
+ + f"## Candidate Guidelines to Review\n{candidates_text}\n\n"
+ f"## Task\n"
+ f"Review each candidate. For each candidate you APPROVE, copy its full "
+ f"text verbatim into your response. REJECT any candidate that:\n"
+ f"- Is too vague (e.g., 'consider the context' without specifying what)\n"
+ f"- Is too narrow (cherry-picks a specific phrase like 'lingering tug' instead of a general pattern)\n"
+ f"- Is redundant with another candidate\n"
+ + ("- CONTRADICTS an existing guideline above (flips 'classify as A' to 'classify as B')\n"
+ "- Restates an existing guideline without adding new criteria\n"
+ if existing_guidelines else "")
+ + f"- Mentions specific instance text rather than general features\n\n"
+ f"Prefer general patterns over specific words. Prefer 2-3 strong rules over 5 weak ones.\n"
+ f"If NO candidate is good enough, return an empty list.\n\n"
+ f"CRITICAL: Each entry in 'approved' must be ONLY the RULE TEXT from the candidate "
+ f"(the part after 'RULE TEXT:' above). Do NOT include 'Candidate N', 'Pattern:', or any "
+ f"metadata โ just the rule itself. Example:\n"
+ f'{{"approved": ["When the text describes emotional discomfort as a reaction, classify as negative rather than positive."]}}'
+ f"\n\nRespond with JSON: {{\"approved\": [\"\", ...]}}"
+ )
+
+ try:
+ from pydantic import BaseModel
+
+ class CriticResponse(BaseModel):
+ approved: List[str] = []
+
+ try:
+ response = endpoint.query(critic_prompt, CriticResponse)
+ except TypeError:
+ response = endpoint.query(critic_prompt)
+
+ data = self._parse_json(response)
+ approved = data.get('approved', [])
+
+ # Clean up metadata prefixes the critic may have copied
+ import re as _re
+ cleaned_approved = []
+ for a in approved:
+ if not isinstance(a, str):
+ continue
+ # Strip common prefixes: "Candidate N", "Pattern:", "RULE TEXT:", etc.
+ text = a.strip()
+ text = _re.sub(r'^(?:Candidate\s+\d+[^:]*:\s*)?', '', text)
+ text = _re.sub(r'^(?:Pattern:\s*[^(]*\(\d+x\)\s*)?', '', text, flags=_re.IGNORECASE)
+ text = _re.sub(r'^(?:\s*Suggestion:\s*)?', '', text, flags=_re.IGNORECASE)
+ text = _re.sub(r'^(?:\s*RULE TEXT:\s*)?', '', text, flags=_re.IGNORECASE)
+ text = text.strip()
+ if len(text) > 20:
+ cleaned_approved.append(text)
+
+ # Sanity check: reject malformed output (e.g., just indices like ["1","2"])
+ valid_approved = cleaned_approved
+
+ if len(valid_approved) < len(approved):
+ logger.warning(
+ f"[Generator-Critic] Critic returned {len(approved)} items but only "
+ f"{len(valid_approved)} had full guideline text. Falling back to all candidates."
+ )
+ return [c['suggestion'] for c in candidates]
+
+ logger.info(
+ f"[Generator-Critic] Critic approved {len(valid_approved)}/{len(candidates)} guidelines"
+ )
+ return valid_approved
+
+ except Exception as e:
+ logger.warning(f"Critic pass failed: {e}, using all candidates")
+ return [c['suggestion'] for c in candidates]
+
+ def _get_revision_endpoint(self) -> Optional[Any]:
+ """Get or create an AI endpoint for LLM analysis."""
+ if self._endpoint is not None:
+ return self._endpoint
+
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ models = (
+ self.solo_config.revision_models
+ or self.solo_config.labeling_models
+ )
+ for model_config in models:
+ try:
+ endpoint_config = model_config.to_endpoint_config(temperature_override=0.3)
+
+ endpoint = AIEndpointFactory.create_endpoint(endpoint_config)
+ if endpoint:
+ self._endpoint = endpoint
+ return endpoint
+ except Exception:
+ continue
+ except Exception as e:
+ logger.warning(f"Could not create analysis endpoint: {e}")
+
+ return None
+
+ def _truncate(self, text: str) -> str:
+ """Truncate text to MAX_TEXT_LENGTH."""
+ if not text:
+ return ''
+ if len(text) <= self.MAX_TEXT_LENGTH:
+ return text
+ return text[:self.MAX_TEXT_LENGTH] + '...'
+
+ def _parse_json(self, response: Any) -> Dict[str, Any]:
+ """Parse JSON from an LLM response, with robust fallbacks.
+
+ Handles common LLM output issues:
+ - JSON wrapped in markdown code blocks
+ - JSON embedded in surrounding prose
+ - Plain text suggestions (no JSON at all)
+ - Slightly malformed JSON (trailing commas, single quotes)
+ """
+ if isinstance(response, dict):
+ return response
+ if hasattr(response, 'model_dump'):
+ return response.model_dump()
+
+ content = str(response).strip()
+
+ # Try markdown code block extraction
+ match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', content)
+ if match:
+ content = match.group(1).strip()
+
+ # Try direct JSON parse
+ try:
+ return json.loads(content)
+ except json.JSONDecodeError:
+ pass
+
+ # Try extracting a JSON object from anywhere in the text
+ match = re.search(r'\{[^{}]*\}', content)
+ if match:
+ try:
+ return json.loads(match.group(0))
+ except json.JSONDecodeError:
+ pass
+
+ # Fallback: treat the entire response as a plain text suggestion
+ # Strip common preamble patterns
+ cleaned = content
+ for prefix in [
+ 'Here is', 'Here\'s', 'My suggestion', 'Suggestion:',
+ 'Guideline:', 'I suggest', 'I would suggest',
+ ]:
+ if cleaned.lower().startswith(prefix.lower()):
+ cleaned = cleaned[len(prefix):].lstrip(':').strip()
+ break
+
+ if cleaned and len(cleaned) > 10:
+ return {'suggestion': cleaned}
+
+ return {}
diff --git a/potato/solo_mode/disagreement_explorer.py b/potato/solo_mode/disagreement_explorer.py
new file mode 100644
index 0000000000000000000000000000000000000000..131e468a8012386d2d30f95cc914fcf3390e6550
--- /dev/null
+++ b/potato/solo_mode/disagreement_explorer.py
@@ -0,0 +1,410 @@
+"""
+Disagreement Explorer
+
+Provides rich aggregated data for visual exploration of human-LLM
+disagreements, including:
+- Scatter plot data: instances by confidence vs. agreement
+- Timeline: disagreement rate over annotation windows
+- Per-label breakdown: which labels cause most disagreements
+- Filterable disagreement list with text, labels, and reasoning
+"""
+
+import logging
+from collections import defaultdict
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ScatterPoint:
+ """A single point in the confidence-vs-agreement scatter plot."""
+ instance_id: str
+ confidence: float
+ agrees: bool
+ llm_label: str
+ human_label: Optional[str] = None
+ reasoning: str = ""
+ text: str = ""
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ 'instance_id': self.instance_id,
+ 'confidence': self.confidence,
+ 'agrees': self.agrees,
+ 'llm_label': self.llm_label,
+ 'human_label': self.human_label,
+ 'reasoning': self.reasoning,
+ 'text': self.text,
+ }
+
+
+@dataclass
+class TimelineBucket:
+ """A time bucket for the disagreement timeline."""
+ bucket_index: int
+ start_index: int
+ end_index: int
+ total: int
+ agreements: int
+ disagreements: int
+ agreement_rate: float
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ 'bucket_index': self.bucket_index,
+ 'start_index': self.start_index,
+ 'end_index': self.end_index,
+ 'total': self.total,
+ 'agreements': self.agreements,
+ 'disagreements': self.disagreements,
+ 'agreement_rate': self.agreement_rate,
+ }
+
+
+@dataclass
+class LabelBreakdown:
+ """Per-label disagreement statistics."""
+ label: str
+ total_comparisons: int
+ agreements: int
+ disagreements: int
+ agreement_rate: float
+ confused_with: List[Dict[str, Any]] = field(default_factory=list)
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ 'label': self.label,
+ 'total_comparisons': self.total_comparisons,
+ 'agreements': self.agreements,
+ 'disagreements': self.disagreements,
+ 'agreement_rate': self.agreement_rate,
+ 'confused_with': self.confused_with,
+ }
+
+
+@dataclass
+class DisagreementItem:
+ """A single disagreement for the filterable list."""
+ instance_id: str
+ llm_label: str
+ human_label: str
+ confidence: float
+ reasoning: str
+ text: str
+ timestamp: str = ""
+ resolved: bool = False
+ resolution_label: Optional[str] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ 'instance_id': self.instance_id,
+ 'llm_label': self.llm_label,
+ 'human_label': self.human_label,
+ 'confidence': self.confidence,
+ 'reasoning': self.reasoning,
+ 'text': self.text,
+ 'timestamp': self.timestamp,
+ 'resolved': self.resolved,
+ 'resolution_label': self.resolution_label,
+ }
+
+
+class DisagreementExplorer:
+ """Computes aggregated data for disagreement visualization.
+
+ Takes prediction data and comparison history to produce
+ scatter plots, timelines, label breakdowns, and disagreement lists.
+ """
+
+ def __init__(self, app_config: Dict[str, Any], solo_config=None):
+ self._app_config = app_config
+ self._solo_config = solo_config
+
+ def get_explorer_data(
+ self,
+ predictions: Dict[str, Dict[str, Any]],
+ comparison_history: List[Dict[str, Any]],
+ text_getter: Optional[Callable[[str], Optional[str]]] = None,
+ label_filter: Optional[str] = None,
+ ) -> Dict[str, Any]:
+ """Compute the full disagreement explorer dataset.
+
+ Args:
+ predictions: Dict[instance_id][schema_name] -> LLMPrediction
+ comparison_history: List of comparison dicts from ValidationTracker
+ text_getter: Optional callable to get instance text by ID
+ label_filter: Optional label to filter results by
+
+ Returns:
+ Dict with scatter_points, disagreements, label_breakdown, summary.
+ """
+ scatter_points = self._build_scatter_points(
+ predictions, text_getter, label_filter
+ )
+ disagreements = self._build_disagreement_list(
+ predictions, comparison_history, text_getter, label_filter
+ )
+ label_breakdown = self._build_label_breakdown(
+ comparison_history, label_filter
+ )
+
+ # Summary stats
+ total_compared = len(comparison_history)
+ total_disagreements = sum(
+ 1 for c in comparison_history if not c.get('agrees')
+ )
+ disagreement_rate = (
+ total_disagreements / total_compared if total_compared > 0 else 0.0
+ )
+
+ # Confidence distribution for disagreements
+ disagree_confs = [
+ p.confidence for p in scatter_points if not p.agrees
+ ]
+ avg_disagree_conf = (
+ sum(disagree_confs) / len(disagree_confs)
+ if disagree_confs else 0.0
+ )
+
+ return {
+ 'scatter_points': [p.to_dict() for p in scatter_points],
+ 'disagreements': [d.to_dict() for d in disagreements],
+ 'label_breakdown': [b.to_dict() for b in label_breakdown],
+ 'summary': {
+ 'total_compared': total_compared,
+ 'total_disagreements': total_disagreements,
+ 'disagreement_rate': round(disagreement_rate, 4),
+ 'avg_disagreement_confidence': round(avg_disagree_conf, 4),
+ 'labels_with_disagreements': len([
+ b for b in label_breakdown if b.disagreements > 0
+ ]),
+ },
+ }
+
+ def get_timeline(
+ self,
+ comparison_history: List[Dict[str, Any]],
+ bucket_size: int = 10,
+ ) -> Dict[str, Any]:
+ """Compute temporal disagreement trends.
+
+ Args:
+ comparison_history: List of comparison dicts (ordered by time)
+ bucket_size: Number of comparisons per bucket
+
+ Returns:
+ Dict with buckets (list of TimelineBucket) and overall trend.
+ """
+ if not comparison_history:
+ return {'buckets': [], 'trend': 'stable', 'total': 0}
+
+ buckets = []
+ for i in range(0, len(comparison_history), bucket_size):
+ chunk = comparison_history[i:i + bucket_size]
+ agreements = sum(1 for c in chunk if c.get('agrees'))
+ disagreements = len(chunk) - agreements
+ rate = agreements / len(chunk) if chunk else 0.0
+
+ buckets.append(TimelineBucket(
+ bucket_index=len(buckets),
+ start_index=i,
+ end_index=i + len(chunk) - 1,
+ total=len(chunk),
+ agreements=agreements,
+ disagreements=disagreements,
+ agreement_rate=round(rate, 4),
+ ))
+
+ # Compute trend from first half vs second half
+ trend = 'stable'
+ if len(buckets) >= 4:
+ mid = len(buckets) // 2
+ first_half = buckets[:mid]
+ second_half = buckets[mid:]
+
+ first_rate = (
+ sum(b.agreements for b in first_half)
+ / max(sum(b.total for b in first_half), 1)
+ )
+ second_rate = (
+ sum(b.agreements for b in second_half)
+ / max(sum(b.total for b in second_half), 1)
+ )
+
+ diff = second_rate - first_rate
+ if diff > 0.05:
+ trend = 'improving'
+ elif diff < -0.05:
+ trend = 'declining'
+
+ return {
+ 'buckets': [b.to_dict() for b in buckets],
+ 'trend': trend,
+ 'total': len(comparison_history),
+ 'bucket_size': bucket_size,
+ }
+
+ def _build_scatter_points(
+ self,
+ predictions: Dict[str, Dict[str, Any]],
+ text_getter: Optional[Callable] = None,
+ label_filter: Optional[str] = None,
+ ) -> List[ScatterPoint]:
+ """Build scatter plot data from predictions with human comparison."""
+ points = []
+
+ for instance_id, schemas in predictions.items():
+ for schema_name, pred in schemas.items():
+ # Only include predictions that have been compared
+ if pred.agrees_with_human is None:
+ continue
+
+ llm_label = str(pred.predicted_label)
+ human_label = (
+ str(pred.human_label) if pred.human_label is not None
+ else None
+ )
+
+ # Apply label filter
+ if label_filter:
+ if llm_label != label_filter and human_label != label_filter:
+ continue
+
+ text = ''
+ if text_getter:
+ text = text_getter(instance_id) or ''
+ if len(text) > 200:
+ text = text[:200] + '...'
+
+ points.append(ScatterPoint(
+ instance_id=instance_id,
+ confidence=pred.confidence_score,
+ agrees=pred.agrees_with_human,
+ llm_label=llm_label,
+ human_label=human_label,
+ reasoning=pred.reasoning[:300] if pred.reasoning else '',
+ text=text,
+ ))
+
+ # Sort by confidence
+ points.sort(key=lambda p: p.confidence)
+ return points
+
+ def _build_disagreement_list(
+ self,
+ predictions: Dict[str, Dict[str, Any]],
+ comparison_history: List[Dict[str, Any]],
+ text_getter: Optional[Callable] = None,
+ label_filter: Optional[str] = None,
+ ) -> List[DisagreementItem]:
+ """Build filterable list of disagreements."""
+ # Build a lookup from comparison history for timestamps
+ timestamps: Dict[str, str] = {}
+ for c in comparison_history:
+ if not c.get('agrees'):
+ timestamps[c['instance_id']] = c.get('timestamp', '')
+
+ items = []
+ for instance_id, schemas in predictions.items():
+ for schema_name, pred in schemas.items():
+ if pred.agrees_with_human is not False:
+ continue
+
+ llm_label = str(pred.predicted_label)
+ human_label = (
+ str(pred.human_label) if pred.human_label is not None
+ else ''
+ )
+
+ if label_filter:
+ if llm_label != label_filter and human_label != label_filter:
+ continue
+
+ text = ''
+ if text_getter:
+ text = text_getter(instance_id) or ''
+ if len(text) > 300:
+ text = text[:300] + '...'
+
+ items.append(DisagreementItem(
+ instance_id=instance_id,
+ llm_label=llm_label,
+ human_label=human_label,
+ confidence=pred.confidence_score,
+ reasoning=pred.reasoning[:500] if pred.reasoning else '',
+ text=text,
+ timestamp=timestamps.get(instance_id, ''),
+ resolved=pred.disagreement_resolved,
+ resolution_label=(
+ str(pred.resolution_label)
+ if pred.resolution_label is not None else None
+ ),
+ ))
+
+ # Sort by confidence ascending (most surprising first)
+ items.sort(key=lambda x: x.confidence, reverse=True)
+ return items
+
+ def _build_label_breakdown(
+ self,
+ comparison_history: List[Dict[str, Any]],
+ label_filter: Optional[str] = None,
+ ) -> List[LabelBreakdown]:
+ """Build per-label disagreement statistics."""
+ # Aggregate by human label (ground truth)
+ label_stats: Dict[str, Dict[str, int]] = defaultdict(
+ lambda: {'total': 0, 'agreements': 0, 'disagreements': 0}
+ )
+ confusion: Dict[str, Dict[str, int]] = defaultdict(
+ lambda: defaultdict(int)
+ )
+
+ for c in comparison_history:
+ human_label = str(c.get('human_label', ''))
+ llm_label = str(c.get('llm_label', ''))
+ agrees = c.get('agrees', False)
+
+ if label_filter and human_label != label_filter and llm_label != label_filter:
+ continue
+
+ label_stats[human_label]['total'] += 1
+ if agrees:
+ label_stats[human_label]['agreements'] += 1
+ else:
+ label_stats[human_label]['disagreements'] += 1
+ confusion[human_label][llm_label] += 1
+
+ breakdowns = []
+ for label, stats in sorted(
+ label_stats.items(),
+ key=lambda x: x[1]['disagreements'],
+ reverse=True,
+ ):
+ rate = (
+ stats['agreements'] / stats['total']
+ if stats['total'] > 0 else 0.0
+ )
+
+ # Top confused-with labels
+ confused_with = [
+ {'label': confused_label, 'count': count}
+ for confused_label, count in sorted(
+ confusion[label].items(),
+ key=lambda x: x[1],
+ reverse=True,
+ )[:5]
+ ]
+
+ breakdowns.append(LabelBreakdown(
+ label=label,
+ total_comparisons=stats['total'],
+ agreements=stats['agreements'],
+ disagreements=stats['disagreements'],
+ agreement_rate=round(rate, 4),
+ confused_with=confused_with,
+ ))
+
+ return breakdowns
diff --git a/potato/solo_mode/disagreement_resolver.py b/potato/solo_mode/disagreement_resolver.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e5a402cb38628492694b15da2654892135560e4
--- /dev/null
+++ b/potato/solo_mode/disagreement_resolver.py
@@ -0,0 +1,298 @@
+"""
+Disagreement Resolver for Solo Mode
+
+This module handles detection and resolution of disagreements between
+human and LLM annotations. It provides type-specific disagreement detection
+and workflows for resolving conflicts.
+"""
+
+import logging
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Optional, Set, Tuple
+import threading
+
+logger = logging.getLogger(__name__)
+
+
+class DisagreementType(Enum):
+ """Types of disagreements between human and LLM."""
+ EXACT_MISMATCH = "exact_mismatch" # Categorical labels differ
+ THRESHOLD_EXCEEDED = "threshold" # Numeric difference > threshold
+ LOW_OVERLAP = "low_overlap" # Set/span overlap below threshold
+ SEMANTIC_DIFFERENCE = "semantic" # Text responses differ semantically
+
+
+@dataclass
+class Disagreement:
+ """Record of a disagreement between human and LLM."""
+ id: str
+ instance_id: str
+ schema_name: str
+ human_label: Any
+ llm_label: Any
+ llm_confidence: float
+ disagreement_type: DisagreementType
+ detected_at: datetime = field(default_factory=datetime.now)
+
+ # Resolution
+ resolved: bool = False
+ resolution_label: Optional[Any] = None
+ resolution_source: Optional[str] = None # 'human_wins', 'llm_wins', 'revised'
+ resolved_at: Optional[datetime] = None
+ resolution_notes: Optional[str] = None
+
+ # For prompt revision
+ triggered_revision: bool = False
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'id': self.id,
+ 'instance_id': self.instance_id,
+ 'schema_name': self.schema_name,
+ 'human_label': self.human_label,
+ 'llm_label': self.llm_label,
+ 'llm_confidence': self.llm_confidence,
+ 'disagreement_type': self.disagreement_type.value,
+ 'detected_at': self.detected_at.isoformat(),
+ 'resolved': self.resolved,
+ 'resolution_label': self.resolution_label,
+ 'resolution_source': self.resolution_source,
+ 'resolved_at': self.resolved_at.isoformat() if self.resolved_at else None,
+ 'resolution_notes': self.resolution_notes,
+ 'triggered_revision': self.triggered_revision,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'Disagreement':
+ """Deserialize from dictionary."""
+ return cls(
+ id=data['id'],
+ instance_id=data['instance_id'],
+ schema_name=data['schema_name'],
+ human_label=data['human_label'],
+ llm_label=data['llm_label'],
+ llm_confidence=data['llm_confidence'],
+ disagreement_type=DisagreementType(data['disagreement_type']),
+ detected_at=datetime.fromisoformat(data['detected_at']),
+ resolved=data.get('resolved', False),
+ resolution_label=data.get('resolution_label'),
+ resolution_source=data.get('resolution_source'),
+ resolved_at=(
+ datetime.fromisoformat(data['resolved_at'])
+ if data.get('resolved_at') else None
+ ),
+ resolution_notes=data.get('resolution_notes'),
+ triggered_revision=data.get('triggered_revision', False),
+ )
+
+
+class DisagreementDetector:
+ """
+ Detects disagreements between human and LLM annotations.
+
+ Uses type-specific comparison logic based on annotation type.
+ """
+
+ def __init__(self, thresholds: Optional[Dict[str, float]] = None):
+ """
+ Initialize the detector.
+
+ Args:
+ thresholds: Optional threshold configuration
+ """
+ self.thresholds = thresholds or {}
+
+ def detect(
+ self,
+ annotation_type: str,
+ human_label: Any,
+ llm_label: Any,
+ schema_info: Optional[Dict[str, Any]] = None
+ ) -> Tuple[bool, DisagreementType]:
+ """
+ Check if human and LLM labels disagree.
+
+ Args:
+ annotation_type: The type of annotation
+ human_label: The human's label
+ llm_label: The LLM's label
+ schema_info: Optional schema information
+
+ Returns:
+ Tuple of (is_disagreement, disagreement_type)
+ """
+ if annotation_type in ('radio', 'select'):
+ return self._check_categorical(human_label, llm_label)
+
+ elif annotation_type == 'likert':
+ return self._check_likert(human_label, llm_label)
+
+ elif annotation_type == 'multiselect':
+ return self._check_multiselect(human_label, llm_label)
+
+ elif annotation_type == 'textbox':
+ return self._check_textbox(human_label, llm_label)
+
+ elif annotation_type == 'span':
+ return self._check_span(human_label, llm_label)
+
+ elif annotation_type in ('slider', 'number'):
+ return self._check_numeric(human_label, llm_label, schema_info)
+
+ else:
+ # Default to exact match
+ return self._check_categorical(human_label, llm_label)
+
+ def _check_categorical(
+ self,
+ human_label: Any,
+ llm_label: Any
+ ) -> Tuple[bool, DisagreementType]:
+ """Check categorical (exact match) disagreement."""
+ agrees = str(human_label) == str(llm_label)
+ return (not agrees, DisagreementType.EXACT_MISMATCH)
+
+ def _check_likert(
+ self,
+ human_label: Any,
+ llm_label: Any
+ ) -> Tuple[bool, DisagreementType]:
+ """Check likert scale disagreement with tolerance."""
+ tolerance = self.thresholds.get('likert_tolerance', 1)
+ try:
+ diff = abs(int(human_label) - int(llm_label))
+ agrees = diff <= tolerance
+ return (not agrees, DisagreementType.THRESHOLD_EXCEEDED)
+ except (ValueError, TypeError):
+ return self._check_categorical(human_label, llm_label)
+
+ def _check_multiselect(
+ self,
+ human_label: Any,
+ llm_label: Any
+ ) -> Tuple[bool, DisagreementType]:
+ """Check multiselect disagreement using Jaccard similarity."""
+ threshold = self.thresholds.get('multiselect_jaccard_threshold', 0.5)
+
+ human_set = set(human_label) if isinstance(human_label, (list, set)) else {human_label}
+ llm_set = set(llm_label) if isinstance(llm_label, (list, set)) else {llm_label}
+
+ if not human_set and not llm_set:
+ return (False, DisagreementType.LOW_OVERLAP)
+
+ intersection = len(human_set & llm_set)
+ union = len(human_set | llm_set)
+ jaccard = intersection / union if union > 0 else 0
+
+ agrees = jaccard >= threshold
+ return (not agrees, DisagreementType.LOW_OVERLAP)
+
+ def _check_textbox(
+ self,
+ human_label: Any,
+ llm_label: Any
+ ) -> Tuple[bool, DisagreementType]:
+ """
+ Check textbox disagreement.
+
+ Currently uses exact match; could be enhanced with
+ embedding similarity.
+ """
+ human_text = str(human_label).strip().lower()
+ llm_text = str(llm_label).strip().lower()
+ agrees = human_text == llm_text
+ return (not agrees, DisagreementType.SEMANTIC_DIFFERENCE)
+
+ def _check_span(
+ self,
+ human_label: Any,
+ llm_label: Any
+ ) -> Tuple[bool, DisagreementType]:
+ """
+ Check span annotation disagreement.
+
+ Compares span boundaries with overlap threshold.
+ """
+ threshold = self.thresholds.get('span_overlap_threshold', 0.5)
+
+ # Extract span info (assuming dict format)
+ human_spans = self._normalize_spans(human_label)
+ llm_spans = self._normalize_spans(llm_label)
+
+ if not human_spans and not llm_spans:
+ return (False, DisagreementType.LOW_OVERLAP)
+
+ if not human_spans or not llm_spans:
+ return (True, DisagreementType.LOW_OVERLAP)
+
+ # Calculate overlap
+ total_overlap = 0
+ total_human_length = 0
+
+ for h_span in human_spans:
+ h_start, h_end = h_span['start'], h_span['end']
+ total_human_length += (h_end - h_start)
+
+ for l_span in llm_spans:
+ l_start, l_end = l_span['start'], l_span['end']
+ overlap_start = max(h_start, l_start)
+ overlap_end = min(h_end, l_end)
+ if overlap_end > overlap_start:
+ total_overlap += (overlap_end - overlap_start)
+
+ overlap_ratio = total_overlap / total_human_length if total_human_length > 0 else 0
+ agrees = overlap_ratio >= threshold
+ return (not agrees, DisagreementType.LOW_OVERLAP)
+
+ def _normalize_spans(self, spans: Any) -> List[Dict[str, int]]:
+ """Normalize span format to list of {start, end} dicts."""
+ if not spans:
+ return []
+
+ if isinstance(spans, list):
+ result = []
+ for span in spans:
+ if isinstance(span, dict) and 'start' in span and 'end' in span:
+ result.append({'start': span['start'], 'end': span['end']})
+ return result
+
+ if isinstance(spans, dict) and 'start' in spans and 'end' in spans:
+ return [{'start': spans['start'], 'end': spans['end']}]
+
+ return []
+
+ def _check_numeric(
+ self,
+ human_label: Any,
+ llm_label: Any,
+ schema_info: Optional[Dict[str, Any]] = None
+ ) -> Tuple[bool, DisagreementType]:
+ """Check numeric value disagreement with relative tolerance."""
+ try:
+ human_val = float(human_label)
+ llm_val = float(llm_label)
+
+ # Use relative tolerance based on range
+ if schema_info:
+ min_val = schema_info.get('min_value', 0)
+ max_val = schema_info.get('max_value', 100)
+ value_range = max_val - min_val
+ tolerance = value_range * 0.1 # 10% of range
+ else:
+ tolerance = abs(human_val) * 0.1 if human_val != 0 else 0.1
+
+ agrees = abs(human_val - llm_val) <= tolerance
+ return (not agrees, DisagreementType.THRESHOLD_EXCEEDED)
+
+ except (ValueError, TypeError):
+ return self._check_categorical(human_label, llm_label)
+
+
+# DisagreementResolver was removed: its check_and_record() was never wired
+# into any code path. The authoritative disagreement state lives on
+# SoloModeManager.disagreement_ids and is read via
+# SoloModeManager.get_pending_disagreements(). DisagreementDetector (above)
+# is still used directly by integration tests for type-specific detection.
diff --git a/potato/solo_mode/edge_case_rules.py b/potato/solo_mode/edge_case_rules.py
new file mode 100644
index 0000000000000000000000000000000000000000..93283f0d4c3ff2f1a316df1ea3d3f70f1b1bd452
--- /dev/null
+++ b/potato/solo_mode/edge_case_rules.py
@@ -0,0 +1,456 @@
+"""
+Edge Case Rule Discovery for Solo Mode
+
+Implements Co-DETECT-inspired edge case rule discovery from real data during
+annotation. When the LLM labels an instance with low confidence, it extracts
+a generalizable edge case rule ("When -> "). These rules
+are clustered, aggregated into categories, reviewed by the human, and injected
+back into the annotation guidelines.
+
+Reference: Co-DETECT (EMNLP 2025 Demo) - https://aclanthology.org/2025.emnlp-demos.25.pdf
+"""
+
+import json
+import logging
+import os
+import threading
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Set
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class EdgeCaseRule:
+ """A rule discovered from real data during annotation.
+
+ Extracted when the LLM labels an instance with low confidence.
+ Format: "When -> "
+ """
+ id: str
+ instance_id: str
+ rule_text: str # Full rule: "When -> "
+ condition: str # The part
+ action: str # The part
+
+ # Source context
+ source_confidence: float
+ source_label: Any
+ prompt_version: int
+ model_name: str = ""
+ created_at: datetime = field(default_factory=datetime.now)
+
+ # Clustering (filled during Phase 2)
+ cluster_id: Optional[int] = None
+ embedding: Optional[List[float]] = None
+
+ # Review (filled during Phase 3)
+ reviewed: bool = False
+ approved: Optional[bool] = None
+ reviewer_notes: str = ""
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'id': self.id,
+ 'instance_id': self.instance_id,
+ 'rule_text': self.rule_text,
+ 'condition': self.condition,
+ 'action': self.action,
+ 'source_confidence': self.source_confidence,
+ 'source_label': self.source_label,
+ 'prompt_version': self.prompt_version,
+ 'model_name': self.model_name,
+ 'created_at': self.created_at.isoformat(),
+ 'cluster_id': self.cluster_id,
+ 'reviewed': self.reviewed,
+ 'approved': self.approved,
+ 'reviewer_notes': self.reviewer_notes,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'EdgeCaseRule':
+ """Deserialize from dictionary."""
+ return cls(
+ id=data['id'],
+ instance_id=data['instance_id'],
+ rule_text=data['rule_text'],
+ condition=data['condition'],
+ action=data['action'],
+ source_confidence=data['source_confidence'],
+ source_label=data.get('source_label'),
+ prompt_version=data.get('prompt_version', 0),
+ model_name=data.get('model_name', ''),
+ created_at=datetime.fromisoformat(data['created_at']),
+ cluster_id=data.get('cluster_id'),
+ reviewed=data.get('reviewed', False),
+ approved=data.get('approved'),
+ reviewer_notes=data.get('reviewer_notes', ''),
+ )
+
+
+@dataclass
+class EdgeCaseCategory:
+ """An aggregated group of similar edge case rules.
+
+ Created by clustering individual rules and synthesizing a summary.
+ """
+ id: str
+ summary_rule: str # Aggregated summary rule for the cluster
+ member_rule_ids: List[str] = field(default_factory=list)
+
+ # Review status
+ reviewed: bool = False
+ approved: Optional[bool] = None
+ reviewer_notes: str = ""
+ incorporated_into_prompt_version: Optional[int] = None
+ created_at: datetime = field(default_factory=datetime.now)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'id': self.id,
+ 'summary_rule': self.summary_rule,
+ 'member_rule_ids': self.member_rule_ids,
+ 'reviewed': self.reviewed,
+ 'approved': self.approved,
+ 'reviewer_notes': self.reviewer_notes,
+ 'incorporated_into_prompt_version': self.incorporated_into_prompt_version,
+ 'created_at': self.created_at.isoformat(),
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'EdgeCaseCategory':
+ """Deserialize from dictionary."""
+ return cls(
+ id=data['id'],
+ summary_rule=data['summary_rule'],
+ member_rule_ids=data.get('member_rule_ids', []),
+ reviewed=data.get('reviewed', False),
+ approved=data.get('approved'),
+ reviewer_notes=data.get('reviewer_notes', ''),
+ incorporated_into_prompt_version=data.get('incorporated_into_prompt_version'),
+ created_at=datetime.fromisoformat(data['created_at']),
+ )
+
+
+class EdgeCaseRuleManager:
+ """Manages edge case rule storage, retrieval, and lifecycle.
+
+ Thread-safe manager that handles:
+ - Recording new rules from LLM labeling
+ - Retrieving rules by status (unclustered, pending review, approved)
+ - Approving/rejecting categories
+ - Formatting approved rules for prompt injection
+ - Persistence to disk
+ """
+
+ def __init__(self, state_dir: Optional[str] = None):
+ """Initialize the rule manager.
+
+ Args:
+ state_dir: Directory for persistent storage
+ """
+ self._lock = threading.RLock()
+ self._rules: Dict[str, EdgeCaseRule] = {} # id -> rule
+ self._categories: Dict[str, EdgeCaseCategory] = {} # id -> category
+ self.state_dir = state_dir
+ self._state_file = 'edge_case_rules.json'
+
+ def record_rule_from_labeling(
+ self,
+ instance_id: str,
+ rule_text: str,
+ condition: str,
+ action: str,
+ confidence: float,
+ label: Any,
+ prompt_version: int,
+ model_name: str = "",
+ ) -> EdgeCaseRule:
+ """Record a new edge case rule discovered during labeling.
+
+ Args:
+ instance_id: ID of the instance that triggered rule extraction
+ rule_text: Full rule text: "When -> "
+ condition: The condition part of the rule
+ action: The action part of the rule
+ confidence: LLM confidence when labeling this instance
+ label: The label assigned by the LLM
+ prompt_version: Version of the prompt used
+ model_name: Name of the model that produced the rule
+
+ Returns:
+ The created EdgeCaseRule
+ """
+ with self._lock:
+ rule_id = f"rule_{uuid.uuid4().hex[:8]}"
+ rule = EdgeCaseRule(
+ id=rule_id,
+ instance_id=instance_id,
+ rule_text=rule_text,
+ condition=condition,
+ action=action,
+ source_confidence=confidence,
+ source_label=label,
+ prompt_version=prompt_version,
+ model_name=model_name,
+ )
+ self._rules[rule_id] = rule
+ self._save_state()
+ logger.info(
+ f"Recorded edge case rule {rule_id} from instance {instance_id} "
+ f"(confidence={confidence:.2f})"
+ )
+ return rule
+
+ def get_rule(self, rule_id: str) -> Optional[EdgeCaseRule]:
+ """Get a rule by ID."""
+ with self._lock:
+ return self._rules.get(rule_id)
+
+ def get_all_rules(self) -> List[EdgeCaseRule]:
+ """Get all rules."""
+ with self._lock:
+ return list(self._rules.values())
+
+ def get_rule_instance_ids(self) -> Set[str]:
+ """Get instance IDs that have edge case rules."""
+ with self._lock:
+ return {rule.instance_id for rule in self._rules.values()}
+
+ def get_unclustered_rules(self) -> List[EdgeCaseRule]:
+ """Get rules that haven't been assigned to a cluster."""
+ with self._lock:
+ return [r for r in self._rules.values() if r.cluster_id is None]
+
+ def get_rules_for_cluster(self, cluster_id: int) -> List[EdgeCaseRule]:
+ """Get all rules in a specific cluster."""
+ with self._lock:
+ return [r for r in self._rules.values() if r.cluster_id == cluster_id]
+
+ def set_rule_cluster(self, rule_id: str, cluster_id: int) -> None:
+ """Assign a rule to a cluster."""
+ with self._lock:
+ if rule_id in self._rules:
+ self._rules[rule_id].cluster_id = cluster_id
+
+ def add_category(self, category: EdgeCaseCategory) -> None:
+ """Add an aggregated category."""
+ with self._lock:
+ self._categories[category.id] = category
+ self._save_state()
+
+ def get_category(self, category_id: str) -> Optional[EdgeCaseCategory]:
+ """Get a category by ID."""
+ with self._lock:
+ return self._categories.get(category_id)
+
+ def get_category_for_rule(self, rule_id: str) -> Optional[EdgeCaseCategory]:
+ """Get the category that contains a given rule."""
+ with self._lock:
+ for cat in self._categories.values():
+ if rule_id in cat.member_rule_ids:
+ return cat
+ return None
+
+ def get_all_categories(self) -> List[EdgeCaseCategory]:
+ """Get all categories."""
+ with self._lock:
+ return list(self._categories.values())
+
+ def get_pending_categories(self) -> List[EdgeCaseCategory]:
+ """Get categories that haven't been reviewed yet."""
+ with self._lock:
+ return [c for c in self._categories.values() if not c.reviewed]
+
+ def get_approved_categories(self) -> List[EdgeCaseCategory]:
+ """Get categories that have been approved."""
+ with self._lock:
+ return [
+ c for c in self._categories.values()
+ if c.reviewed and c.approved
+ ]
+
+ def get_rejected_categories(self) -> List[EdgeCaseCategory]:
+ """Get categories that have been rejected."""
+ with self._lock:
+ return [
+ c for c in self._categories.values()
+ if c.reviewed and not c.approved
+ ]
+
+ def approve_category(
+ self,
+ category_id: str,
+ notes: str = ""
+ ) -> bool:
+ """Approve a category for prompt injection.
+
+ Args:
+ category_id: ID of the category to approve
+ notes: Optional reviewer notes
+
+ Returns:
+ True if category was found and approved
+ """
+ with self._lock:
+ category = self._categories.get(category_id)
+ if category is None:
+ return False
+ category.reviewed = True
+ category.approved = True
+ category.reviewer_notes = notes
+ self._save_state()
+ logger.info(f"Approved edge case category {category_id}")
+ return True
+
+ def reject_category(
+ self,
+ category_id: str,
+ notes: str = ""
+ ) -> bool:
+ """Reject a category.
+
+ Args:
+ category_id: ID of the category to reject
+ notes: Optional reviewer notes
+
+ Returns:
+ True if category was found and rejected
+ """
+ with self._lock:
+ category = self._categories.get(category_id)
+ if category is None:
+ return False
+ category.reviewed = True
+ category.approved = False
+ category.reviewer_notes = notes
+ self._save_state()
+ logger.info(f"Rejected edge case category {category_id}")
+ return True
+
+ def mark_category_incorporated(
+ self,
+ category_id: str,
+ prompt_version: int
+ ) -> None:
+ """Mark a category as incorporated into a prompt version."""
+ with self._lock:
+ category = self._categories.get(category_id)
+ if category:
+ category.incorporated_into_prompt_version = prompt_version
+ self._save_state()
+
+ def get_rules_for_prompt_injection(self) -> str:
+ """Get approved rules formatted for prompt injection.
+
+ Returns:
+ Formatted string of approved edge case guidelines
+ """
+ with self._lock:
+ approved = self.get_approved_categories()
+ if not approved:
+ return ""
+
+ # Filter to only categories not yet incorporated
+ unincorporated = [
+ c for c in approved
+ if c.incorporated_into_prompt_version is None
+ ]
+ if not unincorporated:
+ return ""
+
+ lines = ["## Edge Case Guidelines", ""]
+ for i, category in enumerate(unincorporated, 1):
+ lines.append(f"{i}. {category.summary_rule}")
+ lines.append("")
+
+ return "\n".join(lines)
+
+ def get_stats(self) -> Dict[str, Any]:
+ """Get statistics about rules and categories."""
+ with self._lock:
+ return {
+ 'total_rules': len(self._rules),
+ 'unclustered_rules': len(self.get_unclustered_rules()),
+ 'total_categories': len(self._categories),
+ 'pending_categories': len(self.get_pending_categories()),
+ 'approved_categories': len(self.get_approved_categories()),
+ 'rejected_categories': len(self.get_rejected_categories()),
+ }
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize full state to dictionary."""
+ with self._lock:
+ return {
+ 'rules': {
+ rid: rule.to_dict()
+ for rid, rule in self._rules.items()
+ },
+ 'categories': {
+ cid: cat.to_dict()
+ for cid, cat in self._categories.items()
+ },
+ }
+
+ @classmethod
+ def from_dict(
+ cls,
+ data: Dict[str, Any],
+ state_dir: Optional[str] = None
+ ) -> 'EdgeCaseRuleManager':
+ """Deserialize from dictionary."""
+ manager = cls(state_dir=state_dir)
+ for rid, rule_data in data.get('rules', {}).items():
+ manager._rules[rid] = EdgeCaseRule.from_dict(rule_data)
+ for cid, cat_data in data.get('categories', {}).items():
+ manager._categories[cid] = EdgeCaseCategory.from_dict(cat_data)
+ return manager
+
+ def _save_state(self) -> None:
+ """Save state to disk."""
+ if not self.state_dir:
+ return
+
+ try:
+ os.makedirs(self.state_dir, exist_ok=True)
+ filepath = os.path.join(self.state_dir, self._state_file)
+ temp_path = filepath + '.tmp'
+ with open(temp_path, 'w') as f:
+ json.dump(self.to_dict(), f, indent=2)
+ os.replace(temp_path, filepath)
+ except Exception as e:
+ logger.error(f"Error saving edge case rules state: {e}")
+
+ def load_state(self) -> bool:
+ """Load state from disk.
+
+ Returns:
+ True if state was loaded
+ """
+ if not self.state_dir:
+ return False
+
+ filepath = os.path.join(self.state_dir, self._state_file)
+ if not os.path.exists(filepath):
+ return False
+
+ try:
+ with open(filepath, 'r') as f:
+ data = json.load(f)
+ with self._lock:
+ for rid, rule_data in data.get('rules', {}).items():
+ self._rules[rid] = EdgeCaseRule.from_dict(rule_data)
+ for cid, cat_data in data.get('categories', {}).items():
+ self._categories[cid] = EdgeCaseCategory.from_dict(cat_data)
+ logger.info(
+ f"Loaded edge case rules state: "
+ f"{len(self._rules)} rules, {len(self._categories)} categories"
+ )
+ return True
+ except Exception as e:
+ logger.error(f"Error loading edge case rules state: {e}")
+ return False
diff --git a/potato/solo_mode/edge_case_synthesizer.py b/potato/solo_mode/edge_case_synthesizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..d564682d59c2222e144a69b48d6fe455550d5e6d
--- /dev/null
+++ b/potato/solo_mode/edge_case_synthesizer.py
@@ -0,0 +1,426 @@
+"""
+Edge Case Synthesizer for Solo Mode
+
+This module generates synthetic edge case examples to test and refine
+annotation prompts. It identifies boundary conditions and ambiguous cases
+to help improve prompt quality before large-scale annotation.
+"""
+
+import json
+import logging
+import re
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Set
+import threading
+
+logger = logging.getLogger(__name__)
+
+
+EDGE_CASE_SYNTHESIS_TEMPLATE = """You are an expert at identifying edge cases and boundary conditions for annotation tasks.
+
+Given an annotation task description and some example data, generate synthetic examples
+that would be difficult to label correctly. Focus on cases that:
+1. Lie on the boundary between two labels
+2. Have ambiguous or mixed signals
+3. Require careful interpretation of the guidelines
+4. Test specific aspects of the annotation criteria
+
+## Task Description
+{task_description}
+
+## Annotation Guidelines/Prompt
+{prompt}
+
+## Available Labels
+{labels}
+
+## Example Data (for context)
+{examples}
+
+## Requirements
+- Generate {num_cases} diverse edge cases
+- Each case should test a different aspect of the guidelines
+- Include cases that might reveal gaps in the current instructions
+- Make the examples realistic and varied
+
+## Output Format
+Respond with JSON:
+{{
+ "edge_cases": [
+ {{
+ "text": "",
+ "boundary_labels": ["", ""],
+ "difficulty_reason": "",
+ "which_aspect": ""
+ }},
+ ...
+ ]
+}}
+"""
+
+
+@dataclass
+class EdgeCase:
+ """A synthesized edge case example."""
+ id: str
+ text: str
+ boundary_labels: List[str] # Labels this case is ambiguous between
+ difficulty_reason: str
+ which_aspect: str
+ synthesized_at: datetime = field(default_factory=datetime.now)
+
+ # Human label (filled after labeling)
+ human_label: Optional[str] = None
+ labeler_notes: Optional[str] = None
+ labeled_at: Optional[datetime] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'id': self.id,
+ 'text': self.text,
+ 'boundary_labels': self.boundary_labels,
+ 'difficulty_reason': self.difficulty_reason,
+ 'which_aspect': self.which_aspect,
+ 'synthesized_at': self.synthesized_at.isoformat(),
+ 'human_label': self.human_label,
+ 'labeler_notes': self.labeler_notes,
+ 'labeled_at': self.labeled_at.isoformat() if self.labeled_at else None,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'EdgeCase':
+ """Deserialize from dictionary."""
+ return cls(
+ id=data['id'],
+ text=data['text'],
+ boundary_labels=data['boundary_labels'],
+ difficulty_reason=data['difficulty_reason'],
+ which_aspect=data['which_aspect'],
+ synthesized_at=datetime.fromisoformat(data['synthesized_at']),
+ human_label=data.get('human_label'),
+ labeler_notes=data.get('labeler_notes'),
+ labeled_at=(
+ datetime.fromisoformat(data['labeled_at'])
+ if data.get('labeled_at') else None
+ ),
+ )
+
+
+class EdgeCaseSynthesizer:
+ """
+ Synthesizes edge cases for testing annotation prompts.
+
+ This class generates synthetic examples that are designed to be
+ difficult to label, helping to identify weaknesses in the annotation
+ guidelines before they cause problems during actual annotation.
+ """
+
+ def __init__(self, config: Dict[str, Any], solo_config: Any):
+ """
+ Initialize the edge case synthesizer.
+
+ Args:
+ config: Full application configuration
+ solo_config: SoloModeConfig instance
+ """
+ self.config = config
+ self.solo_config = solo_config
+ self._lock = threading.RLock()
+
+ # Edge case storage
+ self.edge_cases: Dict[str, EdgeCase] = {} # id -> EdgeCase
+ self.synthesis_rounds: List[Dict[str, Any]] = []
+
+ # Track which aspects have been tested
+ self.tested_aspects: Set[str] = set()
+
+ # AI endpoint (lazy init)
+ self._synthesis_endpoint = None
+
+ # Counter for generating IDs
+ self._id_counter = 0
+
+ def _get_synthesis_endpoint(self) -> Optional[Any]:
+ """Get or create AI endpoint for synthesis."""
+ if self._synthesis_endpoint is not None:
+ return self._synthesis_endpoint
+
+ if not self.solo_config.revision_models:
+ logger.warning("No models configured for edge case synthesis")
+ return None
+
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ for model_config in self.solo_config.revision_models:
+ try:
+ endpoint_config = model_config.to_endpoint_config(temperature_override=0.7)
+
+ endpoint = AIEndpointFactory.create_endpoint(endpoint_config)
+ if endpoint:
+ self._synthesis_endpoint = endpoint
+ return endpoint
+ except Exception as e:
+ logger.debug(f"Failed to create synthesis endpoint: {e}")
+ continue
+
+ except Exception as e:
+ logger.error(f"Error creating synthesis endpoint: {e}")
+
+ return None
+
+ def synthesize_edge_cases(
+ self,
+ task_description: str,
+ prompt: str,
+ num_cases: int = 5,
+ existing_examples: Optional[List[str]] = None
+ ) -> List[EdgeCase]:
+ """
+ Generate edge case examples.
+
+ Args:
+ task_description: Description of the annotation task
+ prompt: Current annotation prompt/guidelines
+ num_cases: Number of edge cases to generate
+ existing_examples: Optional list of real examples for context
+
+ Returns:
+ List of generated EdgeCase objects
+ """
+ endpoint = self._get_synthesis_endpoint()
+ if endpoint is None:
+ logger.warning("No endpoint available for edge case synthesis")
+ return []
+
+ try:
+ # Get labels from config
+ schemes = self.config.get('annotation_schemes', [])
+ labels = self._extract_labels(schemes)
+
+ # Format examples
+ examples_text = self._format_examples(existing_examples or [])
+
+ synthesis_prompt = EDGE_CASE_SYNTHESIS_TEMPLATE.format(
+ task_description=task_description,
+ prompt=prompt,
+ labels=labels,
+ examples=examples_text,
+ num_cases=num_cases,
+ )
+
+ from pydantic import BaseModel
+
+ class EdgeCaseResponse(BaseModel):
+ edge_cases: List[Dict[str, Any]] = []
+
+ response = endpoint.query(synthesis_prompt, EdgeCaseResponse)
+
+ # Parse response
+ if isinstance(response, str):
+ response_data = self._parse_json_response(response)
+ elif hasattr(response, 'model_dump'):
+ response_data = response.model_dump()
+ else:
+ response_data = response
+
+ cases = response_data.get('edge_cases', [])
+ generated = []
+
+ with self._lock:
+ for case_data in cases:
+ case = EdgeCase(
+ id=self._generate_id(),
+ text=case_data.get('text', ''),
+ boundary_labels=case_data.get('boundary_labels', []),
+ difficulty_reason=case_data.get('difficulty_reason', ''),
+ which_aspect=case_data.get('which_aspect', ''),
+ )
+
+ if case.text: # Only add if we got text
+ self.edge_cases[case.id] = case
+ self.tested_aspects.add(case.which_aspect)
+ generated.append(case)
+
+ # Record synthesis round
+ self.synthesis_rounds.append({
+ 'timestamp': datetime.now().isoformat(),
+ 'num_requested': num_cases,
+ 'num_generated': len(generated),
+ 'case_ids': [c.id for c in generated],
+ })
+
+ logger.info(f"Synthesized {len(generated)} edge cases")
+ return generated
+
+ except Exception as e:
+ logger.error(f"Error synthesizing edge cases: {e}")
+ return []
+
+ def _generate_id(self) -> str:
+ """Generate a unique edge case ID."""
+ self._id_counter += 1
+ return f"edge_{self._id_counter:04d}"
+
+ def _extract_labels(self, schemes: List[Dict[str, Any]]) -> str:
+ """Extract label names from annotation schemes."""
+ all_labels = []
+ for scheme in schemes:
+ labels = scheme.get('labels', [])
+ for label in labels:
+ if isinstance(label, str):
+ all_labels.append(label)
+ elif isinstance(label, dict):
+ name = label.get('name', '')
+ desc = label.get('description', '')
+ all_labels.append(f"{name}: {desc}" if desc else name)
+ return '\n- '.join([''] + all_labels)
+
+ def _format_examples(self, examples: List[str]) -> str:
+ """Format example data for the prompt."""
+ if not examples:
+ return "No existing examples available."
+
+ formatted = []
+ for i, ex in enumerate(examples[:5]): # Limit to 5 examples
+ text = ex[:200] if len(ex) > 200 else ex
+ formatted.append(f"{i+1}. \"{text}\"")
+ return '\n'.join(formatted)
+
+ def _parse_json_response(self, response: str) -> Dict[str, Any]:
+ """Parse JSON from response."""
+ content = response.strip()
+
+ if '```json' in content:
+ match = re.search(r'```json\s*([\s\S]*?)\s*```', content)
+ if match:
+ content = match.group(1).strip()
+ elif '```' in content:
+ match = re.search(r'```\s*([\s\S]*?)\s*```', content)
+ if match:
+ content = match.group(1).strip()
+
+ try:
+ return json.loads(content)
+ except json.JSONDecodeError:
+ return {'edge_cases': []}
+
+ def get_edge_case(self, case_id: str) -> Optional[EdgeCase]:
+ """Get an edge case by ID."""
+ with self._lock:
+ return self.edge_cases.get(case_id)
+
+ def get_unlabeled_edge_cases(self) -> List[EdgeCase]:
+ """Get edge cases that haven't been labeled yet."""
+ with self._lock:
+ return [
+ case for case in self.edge_cases.values()
+ if case.human_label is None
+ ]
+
+ def get_all_edge_cases(self) -> List[EdgeCase]:
+ """Get all edge cases."""
+ with self._lock:
+ return list(self.edge_cases.values())
+
+ def record_label(
+ self,
+ case_id: str,
+ label: str,
+ notes: Optional[str] = None
+ ) -> bool:
+ """
+ Record a human label for an edge case.
+
+ Args:
+ case_id: The edge case ID
+ label: The assigned label
+ notes: Optional labeler notes
+
+ Returns:
+ True if label was recorded
+ """
+ with self._lock:
+ if case_id not in self.edge_cases:
+ logger.warning(f"Unknown edge case: {case_id}")
+ return False
+
+ case = self.edge_cases[case_id]
+ case.human_label = label
+ case.labeler_notes = notes
+ case.labeled_at = datetime.now()
+
+ logger.info(f"Recorded label '{label}' for edge case {case_id}")
+ return True
+
+ def get_labeled_edge_cases(self) -> List[EdgeCase]:
+ """Get edge cases that have been labeled."""
+ with self._lock:
+ return [
+ case for case in self.edge_cases.values()
+ if case.human_label is not None
+ ]
+
+ def get_cases_for_prompt_revision(self) -> List[Dict[str, Any]]:
+ """
+ Get labeled edge cases formatted for prompt revision.
+
+ Returns cases where the LLM might have been confused,
+ formatted for the prompt revision system.
+ """
+ with self._lock:
+ cases = []
+ for case in self.get_labeled_edge_cases():
+ cases.append({
+ 'text': case.text,
+ 'expected_label': case.human_label,
+ 'boundary_labels': case.boundary_labels,
+ 'difficulty_reason': case.difficulty_reason,
+ 'which_aspect': case.which_aspect,
+ 'labeler_notes': case.labeler_notes,
+ })
+ return cases
+
+ def get_tested_aspects(self) -> Set[str]:
+ """Get the set of aspects that have been tested."""
+ with self._lock:
+ return self.tested_aspects.copy()
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get synthesizer status."""
+ with self._lock:
+ total = len(self.edge_cases)
+ labeled = len([c for c in self.edge_cases.values() if c.human_label])
+ return {
+ 'total_edge_cases': total,
+ 'labeled': labeled,
+ 'unlabeled': total - labeled,
+ 'synthesis_rounds': len(self.synthesis_rounds),
+ 'tested_aspects': list(self.tested_aspects),
+ }
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary for persistence."""
+ with self._lock:
+ return {
+ 'edge_cases': {
+ cid: case.to_dict()
+ for cid, case in self.edge_cases.items()
+ },
+ 'synthesis_rounds': self.synthesis_rounds,
+ 'tested_aspects': list(self.tested_aspects),
+ 'id_counter': self._id_counter,
+ }
+
+ def from_dict(self, data: Dict[str, Any]) -> None:
+ """Load from dictionary."""
+ with self._lock:
+ self.edge_cases = {
+ cid: EdgeCase.from_dict(case_data)
+ for cid, case_data in data.get('edge_cases', {}).items()
+ }
+ self.synthesis_rounds = data.get('synthesis_rounds', [])
+ self.tested_aspects = set(data.get('tested_aspects', []))
+ self._id_counter = data.get('id_counter', len(self.edge_cases))
diff --git a/potato/solo_mode/guideline_updater.py b/potato/solo_mode/guideline_updater.py
new file mode 100644
index 0000000000000000000000000000000000000000..61a4e60ffe05eb84aa01f17f7b922022560cc408
--- /dev/null
+++ b/potato/solo_mode/guideline_updater.py
@@ -0,0 +1,224 @@
+"""
+Guideline Updater for Solo Mode
+
+Injects approved edge case rules into the annotation prompt and identifies
+instances that should be re-annotated with the improved prompt.
+"""
+
+import json
+import logging
+import re
+from typing import Any, Dict, List, Optional, Set
+
+from .edge_case_rules import EdgeCaseCategory
+
+logger = logging.getLogger(__name__)
+
+INJECTION_PROMPT_TEMPLATE = """You are updating annotation guidelines to incorporate newly discovered edge case rules.
+
+Current annotation prompt:
+---
+{current_prompt}
+---
+
+New edge case rules to incorporate:
+{rules_text}
+
+Your task: Produce an updated version of the annotation prompt that naturally integrates
+these edge case rules. Add them as an "Edge Case Guidelines" section near the end of
+the prompt, before any response format instructions. Keep the original prompt intact
+and simply append the new guidelines.
+
+Respond with JSON:
+{{
+ "updated_prompt": ""
+}}
+"""
+
+
+class GuidelineUpdater:
+ """Updates annotation prompts with approved edge case rules.
+
+ Handles:
+ - Injecting approved rules into the prompt via LLM or direct append
+ - Identifying instances for re-annotation based on confidence
+ """
+
+ def __init__(
+ self,
+ app_config: Dict[str, Any],
+ solo_config: Any,
+ ):
+ """Initialize the guideline updater.
+
+ Args:
+ app_config: Full application configuration
+ solo_config: SoloModeConfig instance
+ """
+ self.app_config = app_config
+ self.solo_config = solo_config
+ self._endpoint = None
+
+ def _get_revision_endpoint(self) -> Optional[Any]:
+ """Get or create an AI endpoint for prompt revision."""
+ if self._endpoint is not None:
+ return self._endpoint
+
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ models = self.solo_config.revision_models or self.solo_config.labeling_models
+ for model_config in models:
+ try:
+ endpoint_config = model_config.to_endpoint_config(temperature_override=0.3)
+
+ endpoint = AIEndpointFactory.create_endpoint(endpoint_config)
+ if endpoint:
+ self._endpoint = endpoint
+ return endpoint
+ except Exception:
+ continue
+ except Exception as e:
+ logger.warning(f"Could not create revision endpoint: {e}")
+
+ return None
+
+ def inject_rules_into_prompt(
+ self,
+ current_prompt: str,
+ approved_categories: List[EdgeCaseCategory],
+ ) -> str:
+ """Integrate approved edge case rules into the annotation prompt.
+
+ Tries to use the revision model for natural integration. Falls back
+ to direct append if the model is unavailable.
+
+ Args:
+ current_prompt: The current annotation prompt text
+ approved_categories: List of approved categories to incorporate
+
+ Returns:
+ Updated prompt text with edge case rules integrated
+ """
+ if not approved_categories:
+ return current_prompt
+
+ rules_text = "\n".join(
+ f"- {cat.summary_rule}" for cat in approved_categories
+ )
+
+ # Try LLM-assisted integration
+ endpoint = self._get_revision_endpoint()
+ if endpoint is not None:
+ try:
+ prompt = INJECTION_PROMPT_TEMPLATE.format(
+ current_prompt=current_prompt,
+ rules_text=rules_text,
+ )
+ response = endpoint.query(prompt)
+ response_data = self._parse_json(response)
+ updated = response_data.get('updated_prompt', '')
+ if updated:
+ logger.info("Injected rules via LLM revision")
+ return updated
+ except Exception as e:
+ logger.warning(f"LLM-assisted rule injection failed: {e}")
+
+ # Fallback: direct append
+ return self._direct_inject(current_prompt, rules_text)
+
+ def _direct_inject(self, current_prompt: str, rules_text: str) -> str:
+ """Directly append edge case guidelines to the prompt."""
+ section = f"\n\n## Edge Case Guidelines\n\nThe following edge cases have been identified. Apply these rules when relevant:\n{rules_text}\n"
+
+ # Try to insert before response format instructions
+ # Look for common format instruction patterns
+ format_markers = [
+ "Respond with JSON",
+ "respond with json",
+ "Output format:",
+ "Response format:",
+ ]
+ for marker in format_markers:
+ idx = current_prompt.find(marker)
+ if idx > 0:
+ return current_prompt[:idx] + section + "\n" + current_prompt[idx:]
+
+ # Otherwise append at end
+ return current_prompt + section
+
+ def get_instances_for_reannotation(
+ self,
+ predictions: Dict[str, Dict[str, Any]],
+ old_prompt_version: int,
+ reannotation_counts: Optional[Dict[str, int]] = None,
+ ) -> List[str]:
+ """Get instances that should be re-annotated with the improved prompt.
+
+ Selects instances that:
+ - Were labeled with the old prompt version
+ - Have confidence below the re-annotation threshold
+ - Haven't exceeded max re-annotations
+
+ Args:
+ predictions: Dict of instance_id -> schema -> prediction
+ old_prompt_version: The prompt version to find instances for
+ reannotation_counts: Optional dict tracking per-instance re-annotation count
+
+ Returns:
+ List of instance IDs that should be re-annotated
+ """
+ ecr_config = self.solo_config.edge_case_rules
+ threshold = ecr_config.reannotation_confidence_threshold
+ max_reannotations = ecr_config.max_reannotations_per_instance
+ counts = reannotation_counts or {}
+
+ candidates = []
+ for instance_id, schemas in predictions.items():
+ for schema_name, pred in schemas.items():
+ # Check if labeled with old prompt version
+ pred_version = (
+ pred.prompt_version if hasattr(pred, 'prompt_version')
+ else pred.get('prompt_version', 0)
+ )
+ if pred_version != old_prompt_version:
+ continue
+
+ # Check confidence
+ confidence = (
+ pred.confidence_score if hasattr(pred, 'confidence_score')
+ else pred.get('confidence_score', 1.0)
+ )
+ if confidence >= threshold:
+ continue
+
+ # Check re-annotation limit
+ current_count = counts.get(instance_id, 0)
+ if current_count >= max_reannotations:
+ continue
+
+ candidates.append(instance_id)
+ break # Only need to check one schema per instance
+
+ logger.info(
+ f"Found {len(candidates)} instances for re-annotation "
+ f"(old_version={old_prompt_version}, threshold={threshold})"
+ )
+ return candidates
+
+ def _parse_json(self, response: Any) -> Dict[str, Any]:
+ """Parse JSON from an LLM response."""
+ if isinstance(response, dict):
+ return response
+ if hasattr(response, 'model_dump'):
+ return response.model_dump()
+
+ content = str(response).strip()
+ match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', content)
+ if match:
+ content = match.group(1).strip()
+
+ try:
+ return json.loads(content)
+ except json.JSONDecodeError:
+ return {}
diff --git a/potato/solo_mode/instance_selector.py b/potato/solo_mode/instance_selector.py
new file mode 100644
index 0000000000000000000000000000000000000000..a1e173122eb4eb52f40fdeab08fcd4db066370aa
--- /dev/null
+++ b/potato/solo_mode/instance_selector.py
@@ -0,0 +1,464 @@
+"""
+Instance Selector for Solo Mode
+
+This module implements weighted instance selection for human annotation.
+It combines multiple signals (LLM confidence, diversity, disagreements, random)
+to prioritize which instances the human annotator should see next.
+"""
+
+import logging
+import random
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Set, Tuple
+import threading
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SelectionWeights:
+ """Configuration for instance selection weights."""
+ low_confidence: float = 0.4 # Low LLM confidence instances
+ diverse: float = 0.3 # Diverse instances (embedding clusters)
+ random: float = 0.2 # Random sample for calibration
+ disagreement: float = 0.1 # Instances with prior disagreements
+ edge_case_rule: float = 0.0 # Instances matching edge case rule patterns
+ cartography: float = 0.0 # Instances with high confidence variability
+ llm_predicted: float = 0.0 # Instances with LLM predictions needing human comparison
+
+ def validate(self) -> None:
+ """Validate that weights sum to 1.0."""
+ total = (
+ self.low_confidence +
+ self.diverse +
+ self.random +
+ self.disagreement +
+ self.edge_case_rule +
+ self.cartography +
+ self.llm_predicted
+ )
+ if abs(total - 1.0) > 0.001:
+ # Normalize
+ self.low_confidence /= total
+ self.diverse /= total
+ self.random /= total
+ self.disagreement /= total
+ self.edge_case_rule /= total
+ self.cartography /= total
+ self.llm_predicted /= total
+ logger.warning(f"Normalized selection weights (original sum: {total})")
+
+
+class InstanceSelector:
+ """
+ Weighted instance selector for Solo Mode.
+
+ Combines multiple signals to select which instances the human
+ should annotate, optimizing for efficient use of human labeling time.
+
+ Selection pools:
+ 1. Low confidence: Instances where LLM is uncertain
+ 2. Diverse: Instances from different embedding clusters
+ 3. Random: Random sample for calibration
+ 4. Disagreement: Instances with prior human-LLM disagreement
+ """
+
+ def __init__(
+ self,
+ weights: Optional[SelectionWeights] = None,
+ config: Optional[Dict[str, Any]] = None
+ ):
+ """
+ Initialize the instance selector.
+
+ Args:
+ weights: Selection weight configuration
+ config: Full application configuration
+ """
+ self.weights = weights or SelectionWeights()
+ self.weights.validate()
+ self.config = config or {}
+ self._lock = threading.RLock()
+
+ # Random state
+ self.random = random.Random()
+
+ # Track selection history
+ self.selection_history: List[Dict[str, Any]] = []
+
+ # Pool state
+ self._low_confidence_pool: List[str] = []
+ self._diverse_pool: List[str] = []
+ self._random_pool: List[str] = []
+ self._disagreement_pool: List[str] = []
+ self._edge_case_rule_pool: List[str] = []
+ self._cartography_pool: List[str] = []
+ self._llm_predicted_pool: List[str] = []
+
+ # Cache predictions for use in _select_lowest_confidence
+ self._predictions_cache: Dict[str, Dict[str, Any]] = {}
+
+ def configure(
+ self,
+ low_confidence_weight: float = 0.4,
+ diversity_weight: float = 0.3,
+ random_weight: float = 0.2,
+ disagreement_weight: float = 0.1,
+ edge_case_rule_weight: float = 0.0,
+ cartography_weight: float = 0.0,
+ ) -> None:
+ """Configure selection weights."""
+ self.weights = SelectionWeights(
+ low_confidence=low_confidence_weight,
+ diverse=diversity_weight,
+ random=random_weight,
+ disagreement=disagreement_weight,
+ edge_case_rule=edge_case_rule_weight,
+ cartography=cartography_weight,
+ )
+ self.weights.validate()
+
+ def refresh_pools(
+ self,
+ available_ids: Set[str],
+ llm_predictions: Optional[Dict[str, Dict[str, Any]]] = None,
+ disagreement_ids: Optional[Set[str]] = None,
+ confidence_threshold: float = 0.5,
+ edge_case_rule_ids: Optional[Set[str]] = None,
+ cartography_scores: Optional[Dict[str, float]] = None,
+ ) -> None:
+ """
+ Refresh the selection pools based on current state.
+
+ Args:
+ available_ids: Set of instance IDs available for selection
+ llm_predictions: Dict of instance_id -> schema -> prediction
+ disagreement_ids: Set of instance IDs with disagreements
+ confidence_threshold: Threshold for low confidence pool
+ edge_case_rule_ids: Set of instance IDs matching edge case rule patterns
+ cartography_scores: Dict of instance_id -> variability score
+ """
+ with self._lock:
+ available_list = list(available_ids)
+
+ # Cache predictions for use in _select_lowest_confidence
+ self._predictions_cache = llm_predictions or {}
+
+ # Clear pools
+ self._low_confidence_pool = []
+ self._diverse_pool = []
+ self._random_pool = []
+ self._disagreement_pool = []
+ self._edge_case_rule_pool = []
+ self._cartography_pool = []
+ self._llm_predicted_pool = []
+
+ # Build low confidence pool
+ if llm_predictions:
+ for instance_id in available_list:
+ if instance_id in llm_predictions:
+ preds = llm_predictions[instance_id]
+ # Check if any prediction is below threshold
+ for pred in preds.values():
+ confidence = pred.get('confidence_score', 1.0)
+ if confidence < confidence_threshold:
+ self._low_confidence_pool.append(instance_id)
+ break
+
+ # Build LLM-predicted pool: instances with predictions that aren't
+ # already in the low_confidence pool (those are more valuable there).
+ # This pool steers human annotations toward instances where a comparison
+ # with the LLM can happen immediately.
+ if llm_predictions:
+ low_conf_set = set(self._low_confidence_pool)
+ self._llm_predicted_pool = [
+ iid for iid in available_list
+ if iid in llm_predictions and iid not in low_conf_set
+ ]
+
+ # Build disagreement pool
+ if disagreement_ids:
+ self._disagreement_pool = [
+ iid for iid in available_list
+ if iid in disagreement_ids
+ ]
+
+ # Build edge case rule pool
+ if edge_case_rule_ids:
+ self._edge_case_rule_pool = [
+ iid for iid in available_list
+ if iid in edge_case_rule_ids
+ ]
+
+ # Build cartography pool (high variability = ambiguous instances)
+ if cartography_scores:
+ scored = [
+ (iid, score) for iid, score in cartography_scores.items()
+ if iid in available_ids and score > 0
+ ]
+ scored.sort(key=lambda x: x[1], reverse=True)
+ self._cartography_pool = [iid for iid, _ in scored]
+
+ # Diverse pool uses diversity manager if available
+ self._diverse_pool = self._build_diverse_pool(available_list)
+
+ # Random pool is just all available (sampling happens at selection time)
+ self._random_pool = available_list.copy()
+
+ logger.debug(
+ f"Refreshed pools: low_conf={len(self._low_confidence_pool)}, "
+ f"llm_predicted={len(self._llm_predicted_pool)}, "
+ f"diverse={len(self._diverse_pool)}, "
+ f"random={len(self._random_pool)}, "
+ f"disagreement={len(self._disagreement_pool)}, "
+ f"edge_case_rule={len(self._edge_case_rule_pool)}, "
+ f"cartography={len(self._cartography_pool)}"
+ )
+
+ def _build_diverse_pool(self, available_ids: List[str]) -> List[str]:
+ """
+ Build the diverse instances pool using DiversityManager.
+
+ Returns instances ordered by diversity (from different clusters).
+ """
+ try:
+ from potato.diversity_manager import get_diversity_manager
+
+ dm = get_diversity_manager()
+ if dm is None or not dm.enabled:
+ return []
+
+ # Get diverse ordering from all clusters
+ diverse = dm.generate_diverse_ordering(
+ user_id='solo_mode',
+ available_ids=available_ids,
+ preserve_ids=set()
+ )
+ return diverse
+
+ except Exception as e:
+ logger.debug(f"Could not build diverse pool: {e}")
+ return []
+
+ def select_next(
+ self,
+ available_ids: Set[str],
+ exclude_ids: Optional[Set[str]] = None
+ ) -> Optional[str]:
+ """
+ Select the next instance for human annotation.
+
+ Uses weighted random selection across the pools.
+
+ Args:
+ available_ids: Set of available instance IDs
+ exclude_ids: Set of IDs to exclude from selection
+
+ Returns:
+ Selected instance ID, or None if no instances available
+ """
+ with self._lock:
+ # Filter pools by available and exclude
+ exclude = exclude_ids or set()
+
+ pools = {
+ 'low_confidence': [
+ iid for iid in self._low_confidence_pool
+ if iid in available_ids and iid not in exclude
+ ],
+ 'diverse': [
+ iid for iid in self._diverse_pool
+ if iid in available_ids and iid not in exclude
+ ],
+ 'random': [
+ iid for iid in self._random_pool
+ if iid in available_ids and iid not in exclude
+ ],
+ 'disagreement': [
+ iid for iid in self._disagreement_pool
+ if iid in available_ids and iid not in exclude
+ ],
+ 'edge_case_rule': [
+ iid for iid in self._edge_case_rule_pool
+ if iid in available_ids and iid not in exclude
+ ],
+ 'cartography': [
+ iid for iid in self._cartography_pool
+ if iid in available_ids and iid not in exclude
+ ],
+ 'llm_predicted': [
+ iid for iid in self._llm_predicted_pool
+ if iid in available_ids and iid not in exclude
+ ],
+ }
+
+ # Select pool based on weights
+ selected_pool, pool_name = self._weighted_pool_selection(pools)
+
+ if not selected_pool:
+ # Fallback to any available instance
+ remaining = [iid for iid in available_ids if iid not in exclude]
+ if remaining:
+ instance_id = self.random.choice(remaining)
+ self._record_selection(instance_id, 'fallback')
+ return instance_id
+ return None
+
+ # Select from pool
+ if pool_name == 'low_confidence':
+ # Sort by confidence (lowest first) and take first
+ instance_id = self._select_lowest_confidence(selected_pool)
+ elif pool_name == 'diverse':
+ # Take first (already ordered by diversity)
+ instance_id = selected_pool[0]
+ elif pool_name == 'disagreement':
+ # Random from disagreements
+ instance_id = self.random.choice(selected_pool)
+ elif pool_name == 'edge_case_rule':
+ # Random from edge case rule matches
+ instance_id = self.random.choice(selected_pool)
+ elif pool_name == 'cartography':
+ # Take first (already sorted by variability, highest first)
+ instance_id = selected_pool[0]
+ elif pool_name == 'llm_predicted':
+ # Random from LLM-predicted instances
+ instance_id = self.random.choice(selected_pool)
+ else: # random
+ instance_id = self.random.choice(selected_pool)
+
+ self._record_selection(instance_id, pool_name)
+ return instance_id
+
+ def _weighted_pool_selection(
+ self,
+ pools: Dict[str, List[str]]
+ ) -> Tuple[List[str], str]:
+ """
+ Select a pool based on configured weights.
+
+ Returns empty list if all pools are empty.
+ """
+ # Build list of (pool, name, weight) for non-empty pools
+ candidates = []
+ weights = []
+
+ pool_weights = {
+ 'low_confidence': self.weights.low_confidence,
+ 'diverse': self.weights.diverse,
+ 'random': self.weights.random,
+ 'disagreement': self.weights.disagreement,
+ 'edge_case_rule': self.weights.edge_case_rule,
+ 'cartography': self.weights.cartography,
+ 'llm_predicted': self.weights.llm_predicted,
+ }
+
+ for name, pool in pools.items():
+ if pool: # Only consider non-empty pools
+ candidates.append((pool, name))
+ weights.append(pool_weights[name])
+
+ if not candidates:
+ return [], ''
+
+ # Normalize weights
+ total = sum(weights)
+ if total > 0:
+ weights = [w / total for w in weights]
+
+ # Weighted random selection
+ r = self.random.random()
+ cumsum = 0
+ for (pool, name), weight in zip(candidates, weights):
+ cumsum += weight
+ if r <= cumsum:
+ return pool, name
+
+ # Fallback to last
+ return candidates[-1]
+
+ def _select_lowest_confidence(self, pool: List[str]) -> str:
+ """Select the instance with lowest LLM confidence from cached predictions."""
+ min_conf = float('inf')
+ best_id = pool[0]
+
+ for instance_id in pool:
+ if instance_id in self._predictions_cache:
+ for pred in self._predictions_cache[instance_id].values():
+ conf = pred.get('confidence_score', 1.0)
+ if conf < min_conf:
+ min_conf = conf
+ best_id = instance_id
+
+ return best_id
+
+ def _record_selection(self, instance_id: str, pool_name: str) -> None:
+ """Record a selection for analytics."""
+ from datetime import datetime
+ self.selection_history.append({
+ 'instance_id': instance_id,
+ 'pool': pool_name,
+ 'timestamp': datetime.now().isoformat(),
+ })
+
+ def select_batch(
+ self,
+ available_ids: Set[str],
+ batch_size: int,
+ exclude_ids: Optional[Set[str]] = None
+ ) -> List[str]:
+ """
+ Select a batch of instances for annotation.
+
+ Args:
+ available_ids: Available instance IDs
+ batch_size: Number of instances to select
+ exclude_ids: IDs to exclude
+
+ Returns:
+ List of selected instance IDs
+ """
+ selected = []
+ exclude = set(exclude_ids) if exclude_ids else set()
+
+ for _ in range(batch_size):
+ instance_id = self.select_next(available_ids, exclude)
+ if instance_id is None:
+ break
+ selected.append(instance_id)
+ exclude.add(instance_id)
+
+ return selected
+
+ def get_selection_stats(self) -> Dict[str, Any]:
+ """Get statistics about selections made."""
+ with self._lock:
+ from collections import Counter
+ pool_counts = Counter(s['pool'] for s in self.selection_history)
+
+ return {
+ 'total_selections': len(self.selection_history),
+ 'by_pool': dict(pool_counts),
+ 'pool_sizes': {
+ 'low_confidence': len(self._low_confidence_pool),
+ 'diverse': len(self._diverse_pool),
+ 'random': len(self._random_pool),
+ 'disagreement': len(self._disagreement_pool),
+ 'edge_case_rule': len(self._edge_case_rule_pool),
+ 'cartography': len(self._cartography_pool),
+ 'llm_predicted': len(self._llm_predicted_pool),
+ },
+ 'weights': {
+ 'low_confidence': self.weights.low_confidence,
+ 'diverse': self.weights.diverse,
+ 'random': self.weights.random,
+ 'disagreement': self.weights.disagreement,
+ 'edge_case_rule': self.weights.edge_case_rule,
+ 'cartography': self.weights.cartography,
+ 'llm_predicted': self.weights.llm_predicted,
+ },
+ }
+
+ def clear_history(self) -> None:
+ """Clear selection history."""
+ with self._lock:
+ self.selection_history.clear()
diff --git a/potato/solo_mode/labeling_functions.py b/potato/solo_mode/labeling_functions.py
new file mode 100644
index 0000000000000000000000000000000000000000..cde17c2770ef2fc55705f9eec3b7534ae38f1bd0
--- /dev/null
+++ b/potato/solo_mode/labeling_functions.py
@@ -0,0 +1,724 @@
+"""
+Labeling Function Extraction and Application
+
+Inspired by ALCHEmist (NeurIPS 2024): extracts reusable labeling functions
+from high-confidence LLM predictions to label instances without API calls.
+
+A labeling function encodes a pattern like:
+ "When text contains 'love it' -> positive (confidence: 0.95)"
+
+These functions are extracted from LLM reasoning on high-confidence predictions,
+then applied to unlabeled instances via majority voting.
+"""
+
+import logging
+import re
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+ABSTAIN = "__ABSTAIN__"
+
+
+@dataclass
+class LabelingFunction:
+ """A reusable labeling function extracted from LLM patterns.
+
+ Each function encodes a condition-label mapping discovered from
+ high-confidence LLM predictions.
+ """
+ id: str
+ pattern_text: str # Human-readable pattern description
+ condition: str # The condition part (e.g., "text contains 'great'")
+ label: str # The label to assign when condition matches
+ confidence: float # Source LLM confidence when pattern was discovered
+ source_instance_ids: List[str] = field(default_factory=list)
+ coverage: int = 0 # Number of instances this function matched
+ accuracy: Optional[float] = None # Accuracy against human labels if known
+ enabled: bool = True
+ created_at: str = ""
+ extracted_from_reasoning: str = "" # Original LLM reasoning snippet
+
+ def __post_init__(self):
+ if not self.created_at:
+ self.created_at = datetime.now().isoformat()
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'id': self.id,
+ 'pattern_text': self.pattern_text,
+ 'condition': self.condition,
+ 'label': self.label,
+ 'confidence': self.confidence,
+ 'source_instance_ids': self.source_instance_ids,
+ 'coverage': self.coverage,
+ 'accuracy': self.accuracy,
+ 'enabled': self.enabled,
+ 'created_at': self.created_at,
+ 'extracted_from_reasoning': self.extracted_from_reasoning,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'LabelingFunction':
+ """Deserialize from dictionary."""
+ return cls(
+ id=data['id'],
+ pattern_text=data['pattern_text'],
+ condition=data['condition'],
+ label=data['label'],
+ confidence=data.get('confidence', 0.0),
+ source_instance_ids=data.get('source_instance_ids', []),
+ coverage=data.get('coverage', 0),
+ accuracy=data.get('accuracy'),
+ enabled=data.get('enabled', True),
+ created_at=data.get('created_at', ''),
+ extracted_from_reasoning=data.get('extracted_from_reasoning', ''),
+ )
+
+
+@dataclass
+class LabelingFunctionVote:
+ """A vote from a labeling function for a specific instance."""
+ function_id: str
+ label: str
+ confidence: float
+
+
+@dataclass
+class ApplyResult:
+ """Result of applying labeling functions to an instance."""
+ instance_id: str
+ label: Optional[str] = None
+ votes: List[LabelingFunctionVote] = field(default_factory=list)
+ abstained: bool = True
+ vote_agreement: float = 0.0 # Fraction of votes that agree on the label
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ 'instance_id': self.instance_id,
+ 'label': self.label,
+ 'abstained': self.abstained,
+ 'vote_agreement': self.vote_agreement,
+ 'num_votes': len(self.votes),
+ }
+
+
+EXTRACTION_PROMPT = """Analyze the following high-confidence LLM predictions and extract reusable labeling patterns.
+
+For each prediction, the LLM was highly confident about its label. Extract patterns that could be applied to new instances without calling the LLM.
+
+Predictions:
+{predictions_text}
+
+Extract labeling functions as a JSON array. Each function should have:
+- "pattern_text": A human-readable description of the pattern
+- "condition": A simple text-matching condition (e.g., "text contains 'keyword'", "text starts with 'pattern'", "text mentions sentiment words like 'great', 'love'")
+- "label": The label to assign when the condition matches
+- "keywords": List of keywords/phrases that trigger this pattern
+
+Return ONLY a JSON array of objects. Example:
+[
+ {{
+ "pattern_text": "Positive sentiment keywords like 'love', 'great', 'amazing'",
+ "condition": "text contains positive sentiment keywords",
+ "label": "positive",
+ "keywords": ["love", "great", "amazing", "excellent", "wonderful"]
+ }}
+]"""
+
+
+class LabelingFunctionExtractor:
+ """Extracts labeling functions from high-confidence LLM predictions.
+
+ Analyzes patterns in LLM reasoning to discover reusable rules
+ that can label future instances without API calls.
+ """
+
+ def __init__(self, app_config: Dict, solo_config):
+ self._app_config = app_config
+ self._solo_config = solo_config
+ self._lf_config = solo_config.labeling_functions
+ self._endpoint = None
+
+ def extract_from_predictions(
+ self,
+ predictions: List[Dict[str, Any]],
+ ) -> List[LabelingFunction]:
+ """Extract labeling functions from high-confidence predictions.
+
+ Args:
+ predictions: List of dicts with 'instance_id', 'text',
+ 'predicted_label', 'confidence', 'reasoning'.
+
+ Returns:
+ List of extracted LabelingFunction objects.
+ """
+ if not predictions:
+ return []
+
+ # Filter to high-confidence predictions
+ min_conf = self._lf_config.min_confidence
+ high_conf = [p for p in predictions if p.get('confidence', 0) >= min_conf]
+
+ if not high_conf:
+ return []
+
+ # Group by label
+ by_label: Dict[str, List[Dict]] = {}
+ for p in high_conf:
+ label = str(p.get('predicted_label', ''))
+ by_label.setdefault(label, []).append(p)
+
+ # Try LLM-assisted extraction first
+ functions = self._extract_with_llm(high_conf)
+
+ # Fallback to keyword-based extraction if LLM fails
+ if not functions:
+ functions = self._extract_keyword_patterns(by_label)
+
+ # Limit to max_functions
+ max_fns = self._lf_config.max_functions
+ if len(functions) > max_fns:
+ # Keep highest-confidence functions
+ functions.sort(key=lambda f: f.confidence, reverse=True)
+ functions = functions[:max_fns]
+
+ return functions
+
+ def _extract_with_llm(
+ self, predictions: List[Dict[str, Any]]
+ ) -> List[LabelingFunction]:
+ """Use LLM to extract labeling functions from prediction patterns."""
+ endpoint = self._get_endpoint()
+ if endpoint is None:
+ return []
+
+ # Build prompt with prediction examples (limit to 20 for context)
+ sample = predictions[:20]
+ pred_lines = []
+ for p in sample:
+ text = str(p.get('text', ''))[:200]
+ pred_lines.append(
+ f"- Text: \"{text}\"\n"
+ f" Label: {p.get('predicted_label')} "
+ f"(confidence: {p.get('confidence', 0):.2f})\n"
+ f" Reasoning: {p.get('reasoning', 'N/A')}"
+ )
+
+ prompt = EXTRACTION_PROMPT.format(
+ predictions_text="\n\n".join(pred_lines)
+ )
+
+ try:
+ response = endpoint.query(prompt)
+ parsed = self._parse_json_array(response)
+ if not parsed:
+ return []
+
+ functions = []
+ for item in parsed:
+ if not isinstance(item, dict):
+ continue
+ pattern_text = item.get('pattern_text', '')
+ condition = item.get('condition', '')
+ label = item.get('label', '')
+ keywords = item.get('keywords', [])
+
+ if not label or (not condition and not keywords):
+ continue
+
+ # Find source instances matching this pattern
+ source_ids = []
+ for p in predictions:
+ if str(p.get('predicted_label', '')) == label:
+ source_ids.append(p['instance_id'])
+ if len(source_ids) >= 5:
+ break
+
+ # Compute average confidence for matching predictions
+ matching_confs = [
+ p['confidence'] for p in predictions
+ if str(p.get('predicted_label', '')) == label
+ ]
+ avg_conf = (
+ sum(matching_confs) / len(matching_confs)
+ if matching_confs else 0.0
+ )
+
+ fn = LabelingFunction(
+ id=f"lf_{uuid.uuid4().hex[:8]}",
+ pattern_text=pattern_text,
+ condition=condition,
+ label=label,
+ confidence=avg_conf,
+ source_instance_ids=source_ids,
+ extracted_from_reasoning=', '.join(keywords) if keywords else condition,
+ )
+ functions.append(fn)
+
+ logger.info(f"LLM extracted {len(functions)} labeling functions")
+ return functions
+
+ except Exception as e:
+ logger.warning(f"LLM extraction failed: {e}")
+ return []
+
+ def _extract_keyword_patterns(
+ self, by_label: Dict[str, List[Dict]]
+ ) -> List[LabelingFunction]:
+ """Fallback: extract keyword patterns from prediction texts.
+
+ Groups predictions by label and finds common words/phrases
+ that appear frequently in texts with the same label.
+ """
+ functions = []
+
+ for label, preds in by_label.items():
+ if len(preds) < self._lf_config.min_coverage:
+ continue
+
+ # Collect all words from texts for this label
+ word_counts: Dict[str, int] = {}
+ word_instances: Dict[str, List[str]] = {}
+ for p in preds:
+ text = str(p.get('text', '')).lower()
+ words = set(re.findall(r'\b\w{3,}\b', text))
+ for w in words:
+ word_counts[w] = word_counts.get(w, 0) + 1
+ word_instances.setdefault(w, []).append(p['instance_id'])
+
+ # Find words that appear in >= min_coverage predictions
+ min_cov = self._lf_config.min_coverage
+ common_words = {
+ w: c for w, c in word_counts.items()
+ if c >= min_cov
+ }
+
+ if not common_words:
+ continue
+
+ # Filter out very common words (> 80% of all predictions)
+ total = len(preds)
+ significant = {
+ w: c for w, c in common_words.items()
+ if c <= total * 0.8
+ }
+
+ if not significant:
+ continue
+
+ # Take top keywords by frequency
+ top_keywords = sorted(
+ significant.items(), key=lambda x: x[1], reverse=True
+ )[:5]
+
+ keywords = [w for w, _ in top_keywords]
+ avg_conf = (
+ sum(p.get('confidence', 0) for p in preds) / len(preds)
+ )
+ source_ids = [p['instance_id'] for p in preds[:5]]
+
+ fn = LabelingFunction(
+ id=f"lf_{uuid.uuid4().hex[:8]}",
+ pattern_text=(
+ f"Text containing keywords like "
+ f"'{', '.join(keywords)}' -> {label}"
+ ),
+ condition=f"text contains any of: {', '.join(keywords)}",
+ label=label,
+ confidence=avg_conf,
+ source_instance_ids=source_ids,
+ coverage=len(preds),
+ extracted_from_reasoning=', '.join(keywords),
+ )
+ functions.append(fn)
+
+ return functions
+
+ def _get_endpoint(self):
+ """Get or create an AI endpoint for extraction."""
+ if self._endpoint is not None:
+ return self._endpoint
+
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ models = (
+ self._solo_config.revision_models
+ or self._solo_config.labeling_models
+ )
+ for model_config in models:
+ try:
+ endpoint_config = model_config.to_endpoint_config(temperature_override=0.3)
+
+ endpoint = AIEndpointFactory.create_endpoint(endpoint_config)
+ if endpoint:
+ self._endpoint = endpoint
+ return endpoint
+ except Exception:
+ continue
+ except Exception as e:
+ logger.warning(f"Could not create extraction endpoint: {e}")
+
+ return None
+
+ def _parse_json_array(self, response: str) -> Optional[list]:
+ """Parse a JSON array from LLM response."""
+ import json
+
+ if not response:
+ return None
+
+ # Try direct parse
+ text = response.strip()
+ try:
+ result = json.loads(text)
+ if isinstance(result, list):
+ return result
+ except (json.JSONDecodeError, TypeError):
+ pass
+
+ # Try extracting from markdown code fence
+ match = re.search(r'```(?:json)?\s*\n?(.*?)\n?```', text, re.DOTALL)
+ if match:
+ try:
+ result = json.loads(match.group(1).strip())
+ if isinstance(result, list):
+ return result
+ except (json.JSONDecodeError, TypeError):
+ pass
+
+ # Try finding array brackets
+ start = text.find('[')
+ end = text.rfind(']')
+ if start >= 0 and end > start:
+ try:
+ result = json.loads(text[start:end + 1])
+ if isinstance(result, list):
+ return result
+ except (json.JSONDecodeError, TypeError):
+ pass
+
+ return None
+
+
+class LabelingFunctionApplier:
+ """Applies labeling functions to instances for weak supervision.
+
+ Uses majority voting among matching labeling functions to assign labels
+ without calling the LLM.
+ """
+
+ def __init__(self, vote_threshold: float = 0.5):
+ self._vote_threshold = vote_threshold
+
+ def apply(
+ self,
+ instance_id: str,
+ text: str,
+ functions: List[LabelingFunction],
+ ) -> ApplyResult:
+ """Apply all enabled labeling functions to an instance.
+
+ Args:
+ instance_id: The instance identifier.
+ text: The instance text.
+ functions: List of labeling functions to try.
+
+ Returns:
+ ApplyResult with the voted label or abstention.
+ """
+ votes: List[LabelingFunctionVote] = []
+ text_lower = text.lower()
+
+ for fn in functions:
+ if not fn.enabled:
+ continue
+
+ if self._matches(fn, text_lower):
+ votes.append(LabelingFunctionVote(
+ function_id=fn.id,
+ label=fn.label,
+ confidence=fn.confidence,
+ ))
+
+ if not votes:
+ return ApplyResult(instance_id=instance_id, abstained=True)
+
+ # Majority vote weighted by confidence
+ label_scores: Dict[str, float] = {}
+ for v in votes:
+ label_scores[v.label] = label_scores.get(v.label, 0) + v.confidence
+
+ # Find winning label
+ best_label = max(label_scores, key=label_scores.get)
+ total_score = sum(label_scores.values())
+ agreement = label_scores[best_label] / total_score if total_score > 0 else 0
+
+ # Check if agreement meets threshold
+ if agreement < self._vote_threshold:
+ return ApplyResult(
+ instance_id=instance_id,
+ votes=votes,
+ abstained=True,
+ vote_agreement=agreement,
+ )
+
+ return ApplyResult(
+ instance_id=instance_id,
+ label=best_label,
+ votes=votes,
+ abstained=False,
+ vote_agreement=agreement,
+ )
+
+ def apply_batch(
+ self,
+ instances: List[Dict[str, str]],
+ functions: List[LabelingFunction],
+ ) -> List[ApplyResult]:
+ """Apply labeling functions to a batch of instances.
+
+ Args:
+ instances: List of dicts with 'instance_id' and 'text'.
+ functions: List of labeling functions.
+
+ Returns:
+ List of ApplyResult, one per instance.
+ """
+ enabled = [f for f in functions if f.enabled]
+ if not enabled:
+ return [
+ ApplyResult(instance_id=inst['instance_id'], abstained=True)
+ for inst in instances
+ ]
+
+ return [
+ self.apply(inst['instance_id'], inst['text'], enabled)
+ for inst in instances
+ ]
+
+ def _matches(self, fn: LabelingFunction, text_lower: str) -> bool:
+ """Check if a labeling function matches the given text.
+
+ Uses keyword matching from the function's extracted_from_reasoning
+ and condition fields.
+ """
+ # Extract keywords from the function
+ keywords = self._get_keywords(fn)
+
+ if not keywords:
+ return False
+
+ # Check if any keyword appears in the text
+ return any(kw in text_lower for kw in keywords)
+
+ def _get_keywords(self, fn: LabelingFunction) -> List[str]:
+ """Extract lowercase keywords from a labeling function."""
+ keywords = []
+
+ # Parse keywords from extracted_from_reasoning (comma-separated)
+ reasoning = fn.extracted_from_reasoning
+ if reasoning:
+ parts = [p.strip().lower() for p in reasoning.split(',')]
+ keywords.extend(p for p in parts if p and len(p) >= 2)
+
+ # Parse keywords from condition if it has "contains" pattern
+ condition = fn.condition.lower()
+ # Match patterns like "text contains 'word'" or "any of: word1, word2"
+ contains_match = re.findall(r"'([^']+)'", condition)
+ if contains_match:
+ keywords.extend(w.lower() for w in contains_match)
+
+ any_of_match = re.search(r'any of:\s*(.+)', condition)
+ if any_of_match:
+ parts = [p.strip().lower() for p in any_of_match.group(1).split(',')]
+ keywords.extend(p for p in parts if p and len(p) >= 2)
+
+ return keywords
+
+
+class LabelingFunctionManager:
+ """Manages the lifecycle of labeling functions.
+
+ Handles extraction, storage, application, and statistics tracking.
+ """
+
+ def __init__(self, app_config: Dict, solo_config):
+ self._app_config = app_config
+ self._solo_config = solo_config
+ self._lf_config = solo_config.labeling_functions
+ self._functions: Dict[str, LabelingFunction] = {}
+ self._extractor = LabelingFunctionExtractor(app_config, solo_config)
+ self._applier = LabelingFunctionApplier(
+ vote_threshold=self._lf_config.vote_threshold
+ )
+ self._instances_labeled: int = 0
+ self._instances_abstained: int = 0
+
+ @property
+ def enabled(self) -> bool:
+ return self._lf_config.enabled
+
+ def get_all_functions(self) -> List[LabelingFunction]:
+ """Get all labeling functions."""
+ return list(self._functions.values())
+
+ def get_enabled_functions(self) -> List[LabelingFunction]:
+ """Get only enabled labeling functions."""
+ return [f for f in self._functions.values() if f.enabled]
+
+ def get_function(self, function_id: str) -> Optional[LabelingFunction]:
+ """Get a specific labeling function by ID."""
+ return self._functions.get(function_id)
+
+ def add_function(self, fn: LabelingFunction) -> None:
+ """Add a labeling function."""
+ self._functions[fn.id] = fn
+
+ def toggle_function(self, function_id: str) -> Optional[bool]:
+ """Toggle a function's enabled state. Returns new state or None."""
+ fn = self._functions.get(function_id)
+ if fn is None:
+ return None
+ fn.enabled = not fn.enabled
+ return fn.enabled
+
+ def remove_function(self, function_id: str) -> bool:
+ """Remove a labeling function."""
+ return self._functions.pop(function_id, None) is not None
+
+ def extract_functions(
+ self,
+ predictions: List[Dict[str, Any]],
+ ) -> List[LabelingFunction]:
+ """Extract new labeling functions from predictions.
+
+ Args:
+ predictions: List of dicts with instance_id, text,
+ predicted_label, confidence, reasoning.
+
+ Returns:
+ List of newly extracted functions.
+ """
+ new_fns = self._extractor.extract_from_predictions(predictions)
+
+ for fn in new_fns:
+ self._functions[fn.id] = fn
+
+ if new_fns:
+ logger.info(
+ f"Extracted {len(new_fns)} labeling functions "
+ f"(total: {len(self._functions)})"
+ )
+
+ return new_fns
+
+ def try_label(
+ self, instance_id: str, text: str
+ ) -> Optional[ApplyResult]:
+ """Try to label an instance using labeling functions.
+
+ Returns:
+ ApplyResult if a label was assigned, None if abstained or disabled.
+ """
+ if not self._lf_config.enabled:
+ return None
+
+ enabled = self.get_enabled_functions()
+ if not enabled:
+ return None
+
+ result = self._applier.apply(instance_id, text, enabled)
+
+ if result.abstained:
+ self._instances_abstained += 1
+ return None
+
+ self._instances_labeled += 1
+
+ # Update coverage counts
+ for vote in result.votes:
+ fn = self._functions.get(vote.function_id)
+ if fn:
+ fn.coverage += 1
+
+ return result
+
+ def apply_batch(
+ self, instances: List[Dict[str, str]]
+ ) -> Tuple[List[ApplyResult], List[Dict[str, str]]]:
+ """Apply labeling functions to a batch, returning labeled and remaining.
+
+ Args:
+ instances: List of dicts with instance_id and text.
+
+ Returns:
+ Tuple of (labeled_results, unlabeled_instances).
+ """
+ if not self._lf_config.enabled:
+ return [], instances
+
+ enabled = self.get_enabled_functions()
+ if not enabled:
+ return [], instances
+
+ labeled = []
+ remaining = []
+
+ for inst in instances:
+ result = self._applier.apply(
+ inst['instance_id'], inst['text'], enabled
+ )
+ if result.abstained:
+ remaining.append(inst)
+ self._instances_abstained += 1
+ else:
+ labeled.append(result)
+ self._instances_labeled += 1
+ # Update coverage
+ for vote in result.votes:
+ fn = self._functions.get(vote.function_id)
+ if fn:
+ fn.coverage += 1
+
+ return labeled, remaining
+
+ def get_stats(self) -> Dict[str, Any]:
+ """Get labeling function statistics."""
+ functions = list(self._functions.values())
+ enabled = [f for f in functions if f.enabled]
+
+ return {
+ 'enabled': self._lf_config.enabled,
+ 'total_functions': len(functions),
+ 'enabled_functions': len(enabled),
+ 'instances_labeled': self._instances_labeled,
+ 'instances_abstained': self._instances_abstained,
+ 'total_coverage': sum(f.coverage for f in functions),
+ 'avg_confidence': (
+ sum(f.confidence for f in functions) / len(functions)
+ if functions else 0.0
+ ),
+ }
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize state for persistence."""
+ return {
+ 'functions': [f.to_dict() for f in self._functions.values()],
+ 'instances_labeled': self._instances_labeled,
+ 'instances_abstained': self._instances_abstained,
+ }
+
+ def load_state(self, data: Dict[str, Any]) -> None:
+ """Restore state from persisted data."""
+ self._functions = {}
+ for fn_data in data.get('functions', []):
+ fn = LabelingFunction.from_dict(fn_data)
+ self._functions[fn.id] = fn
+ self._instances_labeled = data.get('instances_labeled', 0)
+ self._instances_abstained = data.get('instances_abstained', 0)
diff --git a/potato/solo_mode/llm_labeler.py b/potato/solo_mode/llm_labeler.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0d7c3646049287cdbb05d899d83322a35f88eb0
--- /dev/null
+++ b/potato/solo_mode/llm_labeler.py
@@ -0,0 +1,569 @@
+"""
+LLM Labeler for Solo Mode
+
+This module provides background LLM labeling functionality for Solo Mode.
+It manages a thread that continuously labels instances while the human
+annotator works, enabling parallel annotation.
+"""
+
+import json
+import logging
+import re
+import threading
+import time
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Set
+from queue import Queue, Empty
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LabelingResult:
+ """Result of labeling a single instance."""
+ instance_id: str
+ schema_name: str
+ label: Any
+ confidence: float
+ uncertainty: float
+ reasoning: str
+ prompt_version: int
+ model_name: str
+ timestamp: datetime = field(default_factory=datetime.now)
+ error: Optional[str] = None
+
+ # Edge case rule discovery (Co-DETECT-style)
+ is_edge_case: bool = False
+ edge_case_rule: Optional[str] = None # "When -> "
+ edge_case_condition: Optional[str] = None # The part
+ edge_case_action: Optional[str] = None # The part
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ result = {
+ 'instance_id': self.instance_id,
+ 'schema_name': self.schema_name,
+ 'label': self.label,
+ 'confidence': self.confidence,
+ 'uncertainty': self.uncertainty,
+ 'reasoning': self.reasoning,
+ 'prompt_version': self.prompt_version,
+ 'model_name': self.model_name,
+ 'timestamp': self.timestamp.isoformat(),
+ 'error': self.error,
+ }
+ if self.is_edge_case:
+ result['is_edge_case'] = True
+ result['edge_case_rule'] = self.edge_case_rule
+ result['edge_case_condition'] = self.edge_case_condition
+ result['edge_case_action'] = self.edge_case_action
+ return result
+
+
+class LLMLabelingThread(threading.Thread):
+ """
+ Background thread for LLM labeling.
+
+ Continuously labels instances from a queue, respecting configured
+ limits on parallel labeling and batch sizes.
+ """
+
+ def __init__(
+ self,
+ config: Dict[str, Any],
+ solo_config: Any,
+ prompt_getter: callable,
+ result_callback: callable,
+ prompt_version_getter: Optional[callable] = None,
+ examples_getter: Optional[callable] = None,
+ ):
+ """
+ Initialize the labeling thread.
+
+ Args:
+ config: Full application configuration
+ solo_config: SoloModeConfig instance
+ prompt_getter: Callable that returns the current prompt text
+ result_callback: Callable to handle labeling results
+ prompt_version_getter: Optional callable that returns current prompt version int
+ examples_getter: Optional callable that returns ICL examples list
+ """
+ super().__init__(name="LLMLabelingThread", daemon=True)
+
+ self.config = config
+ self.solo_config = solo_config
+ self.prompt_getter = prompt_getter
+ self.examples_getter = examples_getter
+ self.result_callback = result_callback
+ self.prompt_version_getter = prompt_version_getter
+
+ # Threading control
+ self._stop_event = threading.Event()
+ self._pause_event = threading.Event()
+
+ # Instance queue
+ self._queue: Queue = Queue()
+
+ # State
+ self._labeled_count = 0
+ self._error_count = 0
+ self._last_error: Optional[str] = None
+
+ # AI endpoint (lazy init)
+ self._endpoint = None
+ self._uncertainty_estimator = None
+
+ def _get_endpoint(self) -> Optional[Any]:
+ """Get or create the labeling AI endpoint."""
+ if self._endpoint is not None:
+ return self._endpoint
+
+ if not self.solo_config.labeling_models:
+ logger.warning("No labeling models configured")
+ return None
+
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ for model_config in self.solo_config.labeling_models:
+ try:
+ endpoint_config = model_config.to_endpoint_config()
+
+ endpoint = AIEndpointFactory.create_endpoint(endpoint_config)
+ if endpoint:
+ self._endpoint = endpoint
+ logger.info(
+ f"Using labeling endpoint: "
+ f"{model_config.endpoint_type}/{model_config.model}"
+ )
+ return endpoint
+ except Exception as e:
+ logger.debug(f"Failed to create endpoint {model_config.model}: {e}")
+ continue
+
+ except Exception as e:
+ logger.error(f"Error creating labeling endpoint: {e}")
+
+ return None
+
+ def _get_uncertainty_estimator(self) -> Optional[Any]:
+ """Get or create the uncertainty estimator."""
+ if self._uncertainty_estimator is not None:
+ return self._uncertainty_estimator
+
+ try:
+ from .uncertainty import create_uncertainty_estimator
+
+ strategy = self.solo_config.uncertainty.strategy
+ estimator_config = {}
+
+ if strategy == 'sampling_diversity':
+ estimator_config = {
+ 'num_samples': self.solo_config.uncertainty.num_samples,
+ 'temperature': self.solo_config.uncertainty.sampling_temperature,
+ }
+
+ self._uncertainty_estimator = create_uncertainty_estimator(
+ strategy,
+ estimator_config
+ )
+ return self._uncertainty_estimator
+
+ except Exception as e:
+ logger.warning(f"Could not create uncertainty estimator: {e}")
+ return None
+
+ def enqueue(self, instance_id: str, instance_text: str, schema_name: str) -> None:
+ """Add an instance to the labeling queue."""
+ self._queue.put({
+ 'instance_id': instance_id,
+ 'text': instance_text,
+ 'schema_name': schema_name,
+ })
+
+ def enqueue_batch(
+ self,
+ instances: List[Dict[str, Any]],
+ schema_name: str
+ ) -> int:
+ """
+ Add a batch of instances to the labeling queue.
+
+ Args:
+ instances: List of {'instance_id': str, 'text': str}
+ schema_name: The schema to label for
+
+ Returns:
+ Number of instances enqueued
+ """
+ count = 0
+ for inst in instances:
+ self.enqueue(
+ inst['instance_id'],
+ inst['text'],
+ schema_name
+ )
+ count += 1
+ return count
+
+ def stop(self) -> None:
+ """Signal the thread to stop."""
+ self._stop_event.set()
+ # Put sentinel to unblock queue
+ self._queue.put(None)
+
+ def pause(self) -> None:
+ """Pause labeling."""
+ self._pause_event.set()
+
+ def resume(self) -> None:
+ """Resume labeling."""
+ self._pause_event.clear()
+
+ def is_paused(self) -> bool:
+ """Check if labeling is paused."""
+ return self._pause_event.is_set()
+
+ def get_queue_size(self) -> int:
+ """Get the current queue size."""
+ return self._queue.qsize()
+
+ def run(self) -> None:
+ """Main thread loop."""
+ logger.info("LLM labeling thread started")
+
+ while not self._stop_event.is_set():
+ # Check pause
+ while self._pause_event.is_set() and not self._stop_event.is_set():
+ time.sleep(1)
+
+ try:
+ # Get next item (with timeout to check stop event)
+ item = self._queue.get(timeout=1.0)
+
+ if item is None: # Sentinel
+ continue
+
+ # Process the item
+ result = self._label_instance(
+ item['instance_id'],
+ item['text'],
+ item['schema_name']
+ )
+
+ if result:
+ self._labeled_count += 1
+ self.result_callback(result)
+ else:
+ self._error_count += 1
+
+ except Empty:
+ continue
+ except Exception as e:
+ logger.error(f"Error in labeling thread: {e}")
+ self._error_count += 1
+ self._last_error = str(e)
+ time.sleep(1) # Back off on error
+
+ logger.info("LLM labeling thread stopped")
+
+ @staticmethod
+ def create_endpoint_from_model_config(model_config):
+ """Create an AI endpoint from a ModelConfig."""
+ from potato.ai.ai_endpoint import AIEndpointFactory
+ endpoint_config = model_config.to_endpoint_config()
+ return AIEndpointFactory.create_endpoint(endpoint_config)
+
+ def _label_instance(
+ self,
+ instance_id: str,
+ text: str,
+ schema_name: str,
+ endpoint=None,
+ ) -> Optional[LabelingResult]:
+ """Label a single instance."""
+ if endpoint is None:
+ endpoint = self._get_endpoint()
+ if endpoint is None:
+ return None
+
+ prompt = self.prompt_getter()
+ if not prompt:
+ logger.warning("No prompt available for labeling")
+ return None
+
+ try:
+ # Get schema info
+ schemes = self.config.get('annotation_schemes', [])
+ schema_info = next(
+ (s for s in schemes if s.get('name') == schema_name),
+ None
+ )
+ if not schema_info:
+ logger.warning(f"Schema {schema_name} not found")
+ return None
+
+ # Build labeling prompt
+ labels = self._extract_labels(schema_info)
+
+ # Check if edge case rule extraction is enabled
+ ecr_config = getattr(self.solo_config, 'edge_case_rules', None)
+ request_edge_case = (
+ ecr_config is not None
+ and ecr_config.enabled
+ and ecr_config.auto_extract_on_labeling
+ )
+
+ # Build ICL examples section if available
+ icl_section = ""
+ if self.examples_getter:
+ try:
+ examples = self.examples_getter()
+ if examples:
+ icl_lines = ["## Examples"]
+ for ex in examples:
+ icl_lines.append(f'Text: "{ex["text"]}"')
+ icl_lines.append(f'Label: {ex["label"]}')
+ icl_lines.append("")
+ icl_section = "\n".join(icl_lines) + "\n"
+ except Exception:
+ pass
+
+ if request_edge_case:
+ full_prompt = f"""{prompt}
+
+{icl_section}Text to label:
+{text}
+
+Available labels: {labels}
+
+Respond with JSON. If you are uncertain about the label (confidence below 75), also identify a generalizable edge case rule that describes when this type of ambiguity occurs:
+{{
+ "label": "",
+ "confidence": <0-100>,
+ "reasoning": "",
+ "is_edge_case": ,
+ "edge_case_rule": " [action]> (only if is_edge_case is true)"
+}}
+"""
+ else:
+ full_prompt = f"""{prompt}
+
+{icl_section}Text to label:
+{text}
+
+Available labels: {labels}
+
+Respond with JSON:
+{{
+ "label": "",
+ "confidence": <0-100>,
+ "reasoning": ""
+}}
+"""
+
+ # Query endpoint
+ from pydantic import BaseModel
+
+ class LabelResponse(BaseModel):
+ label: str
+ confidence: float = 50.0
+ reasoning: str = ""
+
+ response = endpoint.query(full_prompt, LabelResponse)
+
+ # Parse response
+ if isinstance(response, str):
+ response_data = self._parse_json_response(response)
+ elif hasattr(response, 'model_dump'):
+ response_data = response.model_dump()
+ else:
+ response_data = response
+
+ label = response_data.get('label', '')
+ confidence = float(response_data.get('confidence', 50)) / 100.0
+ reasoning = response_data.get('reasoning', '')
+
+ # Validate label
+ valid_labels = self._get_valid_labels(schema_info)
+ if valid_labels and label not in valid_labels:
+ label = self._fuzzy_match_label(label, valid_labels)
+ if label is None:
+ return LabelingResult(
+ instance_id=instance_id,
+ schema_name=schema_name,
+ label=None,
+ confidence=0,
+ uncertainty=1,
+ reasoning="",
+ prompt_version=0,
+ model_name=getattr(endpoint, 'model', ''),
+ error="Invalid label returned"
+ )
+
+ # Estimate uncertainty using configured strategy
+ uncertainty = 1.0 - confidence
+ estimator = self._get_uncertainty_estimator()
+ if estimator:
+ try:
+ logger.debug(f"Running uncertainty estimation ({estimator.__class__.__name__}) for {instance_id}")
+ estimate = estimator.estimate_uncertainty(
+ instance_id=instance_id,
+ text=text,
+ prompt=full_prompt,
+ predicted_label=label,
+ endpoint=endpoint,
+ schema_info=schema_info
+ )
+ uncertainty = estimate.uncertainty_score
+ confidence = estimate.confidence_score
+ logger.debug(
+ f"Uncertainty estimate for {instance_id}: "
+ f"conf={confidence:.3f}, unc={uncertainty:.3f}, "
+ f"method={estimate.method}"
+ )
+ except Exception as e:
+ logger.warning(f"Uncertainty estimation failed for {instance_id}: {e}")
+
+ # Extract edge case rule if present
+ is_edge_case = False
+ edge_case_rule = None
+ edge_case_condition = None
+ edge_case_action = None
+
+ if request_edge_case and response_data.get('is_edge_case'):
+ raw_rule = response_data.get('edge_case_rule', '')
+ if raw_rule:
+ is_edge_case = True
+ edge_case_rule = raw_rule
+ edge_case_condition, edge_case_action = (
+ self._parse_edge_case_rule(raw_rule)
+ )
+
+ prompt_version = 0
+ if self.prompt_version_getter:
+ try:
+ prompt_version = self.prompt_version_getter()
+ except Exception:
+ pass
+
+ return LabelingResult(
+ instance_id=instance_id,
+ schema_name=schema_name,
+ label=label,
+ confidence=confidence,
+ uncertainty=uncertainty,
+ reasoning=reasoning,
+ prompt_version=prompt_version,
+ model_name=getattr(endpoint, 'model', ''),
+ is_edge_case=is_edge_case,
+ edge_case_rule=edge_case_rule,
+ edge_case_condition=edge_case_condition,
+ edge_case_action=edge_case_action,
+ )
+
+ except Exception as e:
+ logger.error(f"Error labeling {instance_id}: {e}")
+ return LabelingResult(
+ instance_id=instance_id,
+ schema_name=schema_name,
+ label=None,
+ confidence=0,
+ uncertainty=1,
+ reasoning="",
+ prompt_version=0,
+ model_name='',
+ error=str(e)
+ )
+
+ def _extract_labels(self, schema_info: Dict[str, Any]) -> str:
+ """Extract label names from schema."""
+ labels = schema_info.get('labels', [])
+ label_names = []
+ for label in labels:
+ if isinstance(label, str):
+ label_names.append(label)
+ elif isinstance(label, dict):
+ label_names.append(label.get('name', str(label)))
+ return ', '.join(label_names)
+
+ def _get_valid_labels(self, schema_info: Dict[str, Any]) -> List[str]:
+ """Get valid label list from schema."""
+ labels = schema_info.get('labels', [])
+ valid = []
+ for label in labels:
+ if isinstance(label, str):
+ valid.append(label)
+ elif isinstance(label, dict):
+ valid.append(label.get('name', str(label)))
+ return valid
+
+ def _fuzzy_match_label(self, label, valid: List[str]) -> Optional[str]:
+ """Try to match label to valid labels. Handles non-string inputs gracefully."""
+ if label is None:
+ return None
+ try:
+ label_lower = str(label).lower().strip()
+ except Exception:
+ return None
+ for v in valid:
+ if str(v).lower().strip() == label_lower:
+ return v
+ return None
+
+ def _parse_edge_case_rule(self, rule_text: str) -> tuple:
+ """Parse a rule in 'When -> ' format.
+
+ Returns:
+ Tuple of (condition, action). Falls back to (rule_text, "") if
+ the format doesn't match.
+ """
+ # Try "When -> " format
+ match = re.match(
+ r'[Ww]hen\s+(.+?)\s*->\s*(.+)',
+ rule_text.strip()
+ )
+ if match:
+ return match.group(1).strip(), match.group(2).strip()
+
+ # Try "If , then " format
+ match = re.match(
+ r'[Ii]f\s+(.+?),?\s+then\s+(.+)',
+ rule_text.strip()
+ )
+ if match:
+ return match.group(1).strip(), match.group(2).strip()
+
+ # Fallback: use full text as condition
+ return rule_text.strip(), ""
+
+ def _parse_json_response(self, response: str) -> Dict[str, Any]:
+ """Parse JSON from response."""
+ content = response.strip()
+
+ if '```json' in content:
+ match = re.search(r'```json\s*([\s\S]*?)\s*```', content)
+ if match:
+ content = match.group(1).strip()
+ elif '```' in content:
+ match = re.search(r'```\s*([\s\S]*?)\s*```', content)
+ if match:
+ content = match.group(1).strip()
+
+ try:
+ return json.loads(content)
+ except json.JSONDecodeError:
+ # Try to extract just the label
+ return {'label': content}
+
+ def get_stats(self) -> Dict[str, Any]:
+ """Get labeling statistics."""
+ return {
+ 'labeled_count': self._labeled_count,
+ 'error_count': self._error_count,
+ 'queue_size': self.get_queue_size(),
+ 'is_paused': self.is_paused(),
+ 'is_running': self.is_alive(),
+ 'last_error': self._last_error,
+ }
diff --git a/potato/solo_mode/manager.py b/potato/solo_mode/manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..935aca06334d6e04e836be7e40e25e9ac9669e14
--- /dev/null
+++ b/potato/solo_mode/manager.py
@@ -0,0 +1,3085 @@
+"""
+Solo Mode Manager
+
+This module provides the central SoloModeManager class that orchestrates
+all Solo Mode operations including prompt management, LLM labeling,
+instance selection, and validation tracking.
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Set, Tuple
+import json
+import logging
+import os
+import threading
+
+from .config import SoloModeConfig, ModelConfig, parse_solo_mode_config
+from .phase_controller import SoloPhase, SoloPhaseController
+
+logger = logging.getLogger(__name__)
+
+# Singleton instance
+_SOLO_MODE_MANAGER: Optional['SoloModeManager'] = None
+_SOLO_MODE_LOCK = threading.Lock()
+
+
+@dataclass
+class PromptVersion:
+ """A versioned prompt for LLM labeling."""
+ version: int
+ prompt_text: str
+ created_at: datetime
+ created_by: str # 'user', 'llm_synthesis', 'llm_optimization'
+ source_description: str = ""
+ parent_version: Optional[int] = None
+ validation_accuracy: Optional[float] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'version': self.version,
+ 'prompt_text': self.prompt_text,
+ 'created_at': self.created_at.isoformat(),
+ 'created_by': self.created_by,
+ 'source_description': self.source_description,
+ 'parent_version': self.parent_version,
+ 'validation_accuracy': self.validation_accuracy,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'PromptVersion':
+ """Deserialize from dictionary."""
+ return cls(
+ version=data['version'],
+ prompt_text=data['prompt_text'],
+ created_at=datetime.fromisoformat(data['created_at']),
+ created_by=data['created_by'],
+ source_description=data.get('source_description', ''),
+ parent_version=data.get('parent_version'),
+ validation_accuracy=data.get('validation_accuracy'),
+ )
+
+
+@dataclass
+class LLMPrediction:
+ """Record of an LLM prediction for an instance."""
+ instance_id: str
+ schema_name: str
+ predicted_label: Any
+ confidence_score: float
+ uncertainty_score: float
+ prompt_version: int
+ timestamp: datetime = field(default_factory=datetime.now)
+ model_name: str = ""
+ reasoning: str = ""
+
+ # Human comparison
+ human_label: Optional[Any] = None
+ agrees_with_human: Optional[bool] = None
+ disagreement_resolved: bool = False
+ resolution_label: Optional[Any] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'instance_id': self.instance_id,
+ 'schema_name': self.schema_name,
+ 'predicted_label': self.predicted_label,
+ 'confidence_score': self.confidence_score,
+ 'uncertainty_score': self.uncertainty_score,
+ 'prompt_version': self.prompt_version,
+ 'timestamp': self.timestamp.isoformat(),
+ 'model_name': self.model_name,
+ 'reasoning': self.reasoning,
+ 'human_label': self.human_label,
+ 'agrees_with_human': self.agrees_with_human,
+ 'disagreement_resolved': self.disagreement_resolved,
+ 'resolution_label': self.resolution_label,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'LLMPrediction':
+ """Deserialize from dictionary."""
+ return cls(
+ instance_id=data['instance_id'],
+ schema_name=data['schema_name'],
+ predicted_label=data['predicted_label'],
+ confidence_score=data['confidence_score'],
+ uncertainty_score=data.get('uncertainty_score', 1.0 - data['confidence_score']),
+ prompt_version=data['prompt_version'],
+ timestamp=datetime.fromisoformat(data['timestamp']),
+ model_name=data.get('model_name', ''),
+ reasoning=data.get('reasoning', ''),
+ human_label=data.get('human_label'),
+ agrees_with_human=data.get('agrees_with_human'),
+ disagreement_resolved=data.get('disagreement_resolved', False),
+ resolution_label=data.get('resolution_label'),
+ )
+
+
+@dataclass
+class AgreementMetrics:
+ """Metrics tracking human-LLM agreement."""
+ total_compared: int = 0
+ agreements: int = 0
+ disagreements: int = 0
+ agreement_rate: float = 0.0
+
+ def update_rate(self):
+ """Update the agreement rate based on current counts."""
+ if self.total_compared == 0:
+ self.agreement_rate = 0.0
+ else:
+ self.agreement_rate = self.agreements / self.total_compared
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'total_compared': self.total_compared,
+ 'agreements': self.agreements,
+ 'disagreements': self.disagreements,
+ 'agreement_rate': self.agreement_rate,
+ }
+
+
+class SoloModeManager:
+ """
+ Central manager for Solo Mode operations.
+
+ This class coordinates:
+ - Phase transitions and state management
+ - Prompt synthesis, versioning, and revision
+ - LLM labeling with uncertainty estimation
+ - Instance selection for human annotation
+ - Human-LLM disagreement tracking
+ - Validation metrics and thresholds
+ """
+
+ def __init__(self, config: SoloModeConfig, app_config: Dict[str, Any]):
+ """
+ Initialize the Solo Mode manager.
+
+ Args:
+ config: SoloModeConfig instance
+ app_config: Full application configuration
+ """
+ self.config = config
+ self.app_config = app_config
+ self._lock = threading.RLock()
+
+ # Initialize phase controller
+ self.phase_controller = SoloPhaseController(config.state_dir)
+
+ # Prompt management
+ self.prompt_versions: List[PromptVersion] = []
+ self.current_prompt_version: int = 0
+ self.task_description: str = ""
+
+ # LLM predictions
+ self.predictions: Dict[str, Dict[str, LLMPrediction]] = {} # instance_id -> schema -> prediction
+
+ # Instance tracking
+ self.human_labeled_ids: Set[str] = set()
+ self.llm_labeled_ids: Set[str] = set()
+ self.disagreement_ids: Set[str] = set()
+ self.validation_sample_ids: Set[str] = set()
+
+ # Edge cases
+ self.edge_case_ids: Set[str] = set()
+ self.edge_case_labels: Dict[str, Dict[str, Any]] = {} # instance_id -> schema -> label
+
+ # Cartography: confidence history per instance across prompt versions
+ # instance_id -> [(prompt_version, confidence_score), ...]
+ self.confidence_history: Dict[str, List[Tuple[int, float]]] = {}
+
+ # Agreement metrics
+ self.agreement_metrics = AgreementMetrics()
+
+ # AI endpoints (lazy initialization)
+ self._labeling_endpoints: List[Any] = []
+ self._revision_endpoints: List[Any] = []
+ self._uncertainty_estimator = None
+
+ # Background labeling
+ self._labeling_thread: Optional[threading.Thread] = None
+ self._stop_labeling = threading.Event()
+ self._pause_labeling = threading.Event()
+
+ # Component instances (lazy initialization)
+ self._edge_case_synthesizer = None
+ self._edge_case_rule_manager = None
+ self._prompt_manager = None
+ self._instance_selector = None
+ self._validation_tracker = None
+ self._llm_labeling_thread = None
+ self._prompt_optimizer = None
+ self._confidence_router = None
+ self._confusion_analyzer = None
+ self._refinement_loop = None
+ self._labeling_function_manager = None
+ self._disagreement_explorer = None
+
+ # Reannotation tracking (persisted across restarts)
+ self._reannotation_counts: Dict[str, int] = {}
+
+ # Per-prompt-version agreement tracking
+ # Tracks agreement separately for each prompt version so we can
+ # measure whether a refinement actually improved accuracy
+ self._per_version_agreement: Dict[int, Dict[str, int]] = {} # version -> {compared, agreements}
+
+ # Validated refinement framework state
+ self._refinement_consecutive_failures: int = 0
+ self._pending_refinements: List[Dict[str, Any]] = []
+ self._refinement_log: List[Dict[str, Any]] = []
+ self._icl_library = None # Lazy-init via _get_icl_library()
+ # Dedicated endpoint for candidate evaluation at low/zero temperature.
+ # Using the labeler endpoint directly would mix prompt quality with
+ # sampling variance. Lazy-init via _get_eval_endpoint().
+ self._eval_endpoint = None
+
+ # State persistence
+ self._state_file = 'solo_mode_state.json'
+
+ logger.info(f"SoloModeManager initialized (enabled={config.enabled})")
+
+ # === Component Properties ===
+
+ @property
+ def edge_case_synthesizer(self):
+ """Lazy-initialized edge case synthesizer."""
+ if self._edge_case_synthesizer is None:
+ from .edge_case_synthesizer import EdgeCaseSynthesizer
+ self._edge_case_synthesizer = EdgeCaseSynthesizer(
+ self.app_config, self.config
+ )
+ return self._edge_case_synthesizer
+
+ @property
+ def edge_case_rule_manager(self):
+ """Lazy-initialized edge case rule manager."""
+ if self._edge_case_rule_manager is None:
+ from .edge_case_rules import EdgeCaseRuleManager
+ self._edge_case_rule_manager = EdgeCaseRuleManager(
+ state_dir=self.config.state_dir
+ )
+ self._edge_case_rule_manager.load_state()
+ return self._edge_case_rule_manager
+
+ @property
+ def prompt_manager(self):
+ """Lazy-initialized prompt manager."""
+ if self._prompt_manager is None:
+ from .prompt_manager import PromptManager
+ self._prompt_manager = PromptManager(self.app_config, self.config)
+ return self._prompt_manager
+
+ @property
+ def instance_selector(self):
+ """Lazy-initialized instance selector."""
+ if self._instance_selector is None:
+ from .instance_selector import InstanceSelector, SelectionWeights
+ weights = SelectionWeights(
+ low_confidence=self.config.instance_selection.low_confidence_weight,
+ diverse=self.config.instance_selection.diversity_weight,
+ random=self.config.instance_selection.random_weight,
+ disagreement=self.config.instance_selection.disagreement_weight,
+ edge_case_rule=self.config.instance_selection.edge_case_rule_weight,
+ cartography=self.config.instance_selection.cartography_weight,
+ llm_predicted=self.config.instance_selection.llm_predicted_weight,
+ )
+ self._instance_selector = InstanceSelector(weights, self.app_config)
+ return self._instance_selector
+
+ @property
+ def validation_tracker(self):
+ """Lazy-initialized validation tracker."""
+ if self._validation_tracker is None:
+ from .validation_tracker import ValidationTracker
+ self._validation_tracker = ValidationTracker(self.app_config)
+ return self._validation_tracker
+
+ @property
+ def llm_labeling_thread(self):
+ """Lazy-initialized LLM labeling thread."""
+ if self._llm_labeling_thread is None:
+ from .llm_labeler import LLMLabelingThread
+ self._llm_labeling_thread = LLMLabelingThread(
+ config=self.app_config,
+ solo_config=self.config,
+ prompt_getter=self.get_current_prompt_text,
+ result_callback=self._handle_labeling_result,
+ prompt_version_getter=lambda: self.current_prompt_version,
+ examples_getter=self.get_icl_examples,
+ )
+ return self._llm_labeling_thread
+
+ @property
+ def prompt_optimizer(self):
+ """Lazy-initialized prompt optimizer."""
+ if not hasattr(self, '_prompt_optimizer') or self._prompt_optimizer is None:
+ from .prompt_optimizer import PromptOptimizer
+ self._prompt_optimizer = PromptOptimizer(
+ config=self.app_config,
+ solo_config=self.config,
+ prompt_getter=self.get_current_prompt_text,
+ prompt_setter=self.update_prompt,
+ examples_getter=self._get_labeled_examples_for_optimization,
+ )
+ return self._prompt_optimizer
+
+ @property
+ def confidence_router(self):
+ """Lazy-initialized confidence router for cascaded escalation."""
+ if self._confidence_router is None and self.config.confidence_routing.enabled:
+ from .confidence_router import ConfidenceRouter
+ from .llm_labeler import LLMLabelingThread
+ self._confidence_router = ConfidenceRouter(
+ routing_config=self.config.confidence_routing,
+ label_fn=self.llm_labeling_thread._label_instance,
+ endpoint_factory=LLMLabelingThread.create_endpoint_from_model_config,
+ )
+ return self._confidence_router
+
+ def get_icl_examples(self, max_per_label: int = 1, max_total: int = 5) -> List[Dict[str, str]]:
+ """Get in-context learning examples for the labeling prompt.
+
+ Priority order:
+ 1. Validated examples from the persistent ICL library (added by
+ the validated refinement framework โ each has proven val gain)
+ 2. Auto-selected examples from human-LLM agreements (fallback)
+
+ Args:
+ max_per_label: Maximum examples per label.
+ max_total: Maximum total examples.
+
+ Returns:
+ List of {"text": "...", "label": "..."} dicts.
+ """
+ examples: List[Dict[str, str]] = []
+
+ # Pull from validated ICL library first (strongest signal)
+ if hasattr(self, '_icl_library') and self._icl_library is not None:
+ try:
+ validated = self._icl_library.get_examples(
+ max_per_label=max_per_label,
+ max_total=max_total,
+ )
+ for ex in validated:
+ # Strip principle field; labeling prompt wants just text+label
+ examples.append({
+ 'text': ex.get('text', ''),
+ 'label': ex.get('label', ''),
+ })
+ except Exception as e:
+ logger.debug(f"[ICL] Failed to read validated library: {e}")
+
+ # If we still have room, fall back to auto-selected agreements
+ if len(examples) < max_total:
+ with self._lock:
+ used_labels = {e['label'] for e in examples}
+ by_label: Dict[str, List[Dict[str, str]]] = {}
+ for instance_id in self.human_labeled_ids:
+ if instance_id not in self.predictions:
+ continue
+ for schema_name, pred in self.predictions[instance_id].items():
+ if not pred.agrees_with_human or pred.human_label is None:
+ continue
+ label = str(pred.human_label)
+ if label not in by_label:
+ by_label[label] = []
+ # Don't duplicate labels already covered by validated entries
+ slots_used = used_labels.count(label) if isinstance(used_labels, list) else (1 if label in used_labels else 0)
+ if len(by_label[label]) + slots_used < max_per_label:
+ text = self._get_instance_text(instance_id)
+ if text:
+ by_label[label].append({
+ 'text': text[:200],
+ 'label': label,
+ })
+
+ for label_examples in by_label.values():
+ for ex in label_examples:
+ if len(examples) >= max_total:
+ break
+ examples.append(ex)
+
+ return examples[:max_total]
+
+ def _get_labeled_examples_for_optimization(self) -> List[Dict[str, Any]]:
+ """Get labeled examples for prompt optimization."""
+ examples = []
+ with self._lock:
+ for instance_id in self.human_labeled_ids:
+ if instance_id in self.predictions:
+ for schema_name, pred in self.predictions[instance_id].items():
+ examples.append({
+ 'instance_id': instance_id,
+ 'text': self._get_instance_text(instance_id),
+ 'predicted_label': pred.predicted_label,
+ 'human_label': pred.human_label,
+ 'actual_label': pred.human_label,
+ 'agrees': pred.agrees_with_human,
+ })
+ return examples
+
+ @property
+ def guideline_updater(self):
+ """Lazy-initialized guideline updater."""
+ if not hasattr(self, '_guideline_updater') or self._guideline_updater is None:
+ from .guideline_updater import GuidelineUpdater
+ self._guideline_updater = GuidelineUpdater(
+ self.app_config, self.config
+ )
+ return self._guideline_updater
+
+ @property
+ def confusion_analyzer(self):
+ """Lazy-initialized confusion analyzer."""
+ if not hasattr(self, '_confusion_analyzer') or self._confusion_analyzer is None:
+ from .confusion_analyzer import ConfusionAnalyzer
+ self._confusion_analyzer = ConfusionAnalyzer(
+ self.app_config, self.config
+ )
+ return self._confusion_analyzer
+
+ @property
+ def refinement_loop(self):
+ """Lazy-initialized refinement loop."""
+ if not hasattr(self, '_refinement_loop') or self._refinement_loop is None:
+ from .refinement_loop import RefinementLoop
+ self._refinement_loop = RefinementLoop(
+ self.config, self.app_config
+ )
+ return self._refinement_loop
+
+ @property
+ def labeling_function_manager(self):
+ """Lazy-initialized labeling function manager."""
+ if (not hasattr(self, '_labeling_function_manager')
+ or self._labeling_function_manager is None):
+ from .labeling_functions import LabelingFunctionManager
+ self._labeling_function_manager = LabelingFunctionManager(
+ self.app_config, self.config
+ )
+ return self._labeling_function_manager
+
+ @property
+ def disagreement_explorer(self):
+ """Lazy-initialized disagreement explorer."""
+ if (not hasattr(self, '_disagreement_explorer')
+ or self._disagreement_explorer is None):
+ from .disagreement_explorer import DisagreementExplorer
+ self._disagreement_explorer = DisagreementExplorer(
+ self.app_config, self.config
+ )
+ return self._disagreement_explorer
+
+ def get_confusion_analysis_full(self) -> Dict[str, Any]:
+ """Get full confusion analysis for the dashboard.
+
+ Returns:
+ Dict with enabled, matrix_data, patterns, totals.
+ """
+ ca_config = self.config.confusion_analysis
+ if not ca_config.enabled:
+ return {'enabled': False}
+
+ tracker = self.validation_tracker
+ metrics = tracker.get_metrics()
+ confusion_matrix = metrics.confusion_matrix
+ comparison_history = tracker.get_comparison_history()
+ label_accuracy = tracker.get_label_accuracy()
+
+ # Get all labels from config
+ labels = self.get_available_labels()
+
+ # Enriched patterns
+ analyzer = self.confusion_analyzer
+ patterns = analyzer.analyze(
+ comparison_history=comparison_history,
+ predictions=self.predictions,
+ text_getter=self._get_instance_text,
+ )
+
+ # Heatmap data
+ matrix_data = analyzer.get_confusion_matrix_data(
+ confusion_matrix, labels, label_accuracy
+ )
+
+ total_disagreements = sum(
+ 1 for r in comparison_history if not r.get('agrees')
+ )
+
+ return {
+ 'enabled': True,
+ 'matrix_data': matrix_data,
+ 'patterns': [p.to_dict() for p in patterns],
+ 'total_disagreements': total_disagreements,
+ 'total_compared': metrics.total_compared,
+ }
+
+ def get_disagreement_explorer_data(
+ self, label_filter: Optional[str] = None
+ ) -> Dict[str, Any]:
+ """Get disagreement explorer data for the dashboard.
+
+ Args:
+ label_filter: Optional label to filter results by.
+
+ Returns:
+ Dict with scatter_points, disagreements, label_breakdown, summary.
+ """
+ tracker = self.validation_tracker
+ comparison_history = tracker.get_comparison_history()
+
+ explorer = self.disagreement_explorer
+ return explorer.get_explorer_data(
+ predictions=self.predictions,
+ comparison_history=comparison_history,
+ text_getter=self._get_instance_text,
+ label_filter=label_filter,
+ )
+
+ def get_disagreement_timeline(
+ self, bucket_size: int = 10
+ ) -> Dict[str, Any]:
+ """Get temporal disagreement trend data.
+
+ Args:
+ bucket_size: Number of comparisons per time bucket.
+
+ Returns:
+ Dict with buckets, trend, total, bucket_size.
+ """
+ tracker = self.validation_tracker
+ comparison_history = tracker.get_comparison_history()
+
+ explorer = self.disagreement_explorer
+ return explorer.get_timeline(
+ comparison_history=comparison_history,
+ bucket_size=bucket_size,
+ )
+
+ def _handle_labeling_result(self, result) -> None:
+ """Handle a labeling result from the LLM labeling thread."""
+ if result.error:
+ logger.warning(f"LLM labeling error for {result.instance_id}: {result.error}")
+ return
+
+ prediction = LLMPrediction(
+ instance_id=result.instance_id,
+ schema_name=result.schema_name,
+ predicted_label=result.label,
+ confidence_score=result.confidence,
+ uncertainty_score=result.uncertainty,
+ prompt_version=result.prompt_version,
+ model_name=result.model_name,
+ reasoning=result.reasoning,
+ )
+ self.set_llm_prediction(result.instance_id, result.schema_name, prediction)
+ llm_count = len(self.llm_labeled_ids)
+ if llm_count % 10 == 0:
+ logger.info(
+ f"[LLM Progress] {llm_count} instances labeled "
+ f"(latest: {result.instance_id} -> {result.label}, "
+ f"conf={result.confidence:.2f})"
+ )
+
+ # Record edge case rule if present
+ if (
+ result.is_edge_case
+ and result.edge_case_rule
+ and self.config.edge_case_rules.enabled
+ ):
+ self.edge_case_rule_manager.record_rule_from_labeling(
+ instance_id=result.instance_id,
+ rule_text=result.edge_case_rule,
+ condition=result.edge_case_condition or result.edge_case_rule,
+ action=result.edge_case_action or "",
+ confidence=result.confidence,
+ label=result.label,
+ prompt_version=result.prompt_version,
+ model_name=result.model_name,
+ )
+
+ # Check if we should trigger clustering
+ self._maybe_trigger_rule_clustering()
+
+ # Check if we should extract labeling functions
+ self._maybe_extract_labeling_functions()
+
+ # Retroactive comparison: if a human already labeled this instance,
+ # compare the LLM prediction against the stored human label.
+ # This handles the case where the human annotated before the LLM.
+ self._retroactive_compare(result.instance_id, result.schema_name)
+
+ def _maybe_trigger_rule_clustering(self) -> None:
+ """Check if enough unclustered rules have accumulated to trigger clustering."""
+ ecr_config = self.config.edge_case_rules
+ unclustered = self.edge_case_rule_manager.get_unclustered_rules()
+ if len(unclustered) >= ecr_config.min_rules_for_clustering:
+ self._trigger_rule_clustering()
+
+ def _trigger_rule_clustering(self) -> None:
+ """Run the rule clustering pipeline in a background thread."""
+ def _run():
+ try:
+ from .rule_clusterer import RuleClusterer
+ clusterer = RuleClusterer(
+ self.app_config,
+ self.config,
+ )
+ rules = self.edge_case_rule_manager.get_unclustered_rules()
+ if not rules:
+ return
+
+ categories = clusterer.run_full_pipeline(rules)
+
+ # Assign cluster IDs to rules and store categories
+ for category in categories:
+ self.edge_case_rule_manager.add_category(category)
+ for rule_id in category.member_rule_ids:
+ self.edge_case_rule_manager.set_rule_cluster(
+ rule_id, hash(category.id) % 10000
+ )
+ self.edge_case_rule_manager._save_state()
+
+ logger.info(
+ f"Rule clustering complete: {len(categories)} categories "
+ f"from {len(rules)} rules"
+ )
+ except Exception as e:
+ logger.error(f"Error in rule clustering pipeline: {e}")
+
+ thread = threading.Thread(
+ target=_run,
+ name="RuleClusteringThread",
+ daemon=True,
+ )
+ thread.start()
+
+ def apply_approved_rules(self) -> Dict[str, Any]:
+ """Apply approved edge case rules by injecting them into the prompt.
+
+ Returns:
+ Dict with success status, new prompt version, and re-annotation info
+ """
+ ecr = self.edge_case_rule_manager
+ approved = ecr.get_approved_categories()
+
+ # Filter to only unincorporated categories
+ unincorporated = [
+ c for c in approved
+ if c.incorporated_into_prompt_version is None
+ ]
+
+ if not unincorporated:
+ return {
+ 'success': False,
+ 'error': 'No unincorporated approved categories',
+ }
+
+ # Inject rules into prompt
+ current_prompt = self.get_current_prompt_text()
+ updated_prompt = self.guideline_updater.inject_rules_into_prompt(
+ current_prompt, unincorporated
+ )
+
+ # Create new prompt version
+ old_version = self.current_prompt_version
+ new_pv = self.create_prompt_version(
+ updated_prompt,
+ created_by='edge_case_rule_injection',
+ source_description=(
+ f"Injected {len(unincorporated)} edge case rule categories"
+ ),
+ )
+
+ # Mark categories as incorporated
+ for cat in unincorporated:
+ ecr.mark_category_incorporated(cat.id, new_pv.version)
+
+ result = {
+ 'success': True,
+ 'new_prompt_version': new_pv.version,
+ 'categories_incorporated': len(unincorporated),
+ 'reannotation_triggered': False,
+ }
+
+ # Trigger re-annotation if enabled
+ if self.config.edge_case_rules.reannotation_enabled:
+ reannotated = self._trigger_reannotation(old_version)
+ result['reannotation_triggered'] = reannotated > 0
+ result['reannotation_count'] = reannotated
+
+ logger.info(
+ f"Applied {len(unincorporated)} edge case rule categories, "
+ f"new prompt version {new_pv.version}"
+ )
+ return result
+
+ def _trigger_reannotation(self, old_prompt_version: int) -> int:
+ """Remove low-confidence instances from llm_labeled_ids so they
+ re-enter the labeling queue with the improved prompt.
+
+ Records old predictions for before/after accuracy comparison.
+
+ Args:
+ old_prompt_version: The prompt version whose labels to reconsider
+
+ Returns:
+ Number of instances queued for re-annotation
+ """
+ candidates = self.guideline_updater.get_instances_for_reannotation(
+ predictions=self.predictions,
+ old_prompt_version=old_prompt_version,
+ reannotation_counts=self._reannotation_counts,
+ )
+
+ with self._lock:
+ # Track old predictions for before/after comparison
+ if not hasattr(self, '_reannotation_history'):
+ self._reannotation_history: List[Dict[str, Any]] = []
+
+ for instance_id in candidates:
+ # Record old prediction before re-annotation
+ old_pred = None
+ for schema_preds in self.predictions.get(instance_id, {}).values():
+ old_pred = {
+ 'instance_id': instance_id,
+ 'old_label': schema_preds.predicted_label,
+ 'old_confidence': schema_preds.confidence_score,
+ 'old_prompt_version': schema_preds.prompt_version,
+ 'new_prompt_version': self.current_prompt_version,
+ 'had_human_label': schema_preds.human_label is not None,
+ 'human_label': schema_preds.human_label,
+ 'old_agreed': schema_preds.agrees_with_human,
+ }
+ break
+ if old_pred:
+ self._reannotation_history.append(old_pred)
+
+ # Remove from llm_labeled_ids so it can be re-labeled
+ self.llm_labeled_ids.discard(instance_id)
+ # Track re-annotation count
+ self._reannotation_counts[instance_id] = (
+ self._reannotation_counts.get(instance_id, 0) + 1
+ )
+
+ if candidates:
+ self._save_state()
+
+ logger.info(f"[Re-annotation] Queued {len(candidates)} instances for re-annotation")
+ return len(candidates)
+
+ def get_reannotation_report(self) -> Dict[str, Any]:
+ """Get a before/after accuracy report for re-annotated instances.
+
+ Returns:
+ Dict with re-annotation statistics and per-instance comparisons.
+ """
+ if not hasattr(self, '_reannotation_history'):
+ return {'total': 0, 'comparisons': []}
+
+ comparisons = []
+ improved = 0
+ worsened = 0
+ unchanged = 0
+
+ with self._lock:
+ for record in self._reannotation_history:
+ iid = record['instance_id']
+ new_pred = None
+ for schema_preds in self.predictions.get(iid, {}).values():
+ if schema_preds.prompt_version > record['old_prompt_version']:
+ new_pred = schema_preds
+ break
+
+ if new_pred is None:
+ continue # Not yet re-annotated
+
+ comp = {
+ 'instance_id': iid,
+ 'old_label': record['old_label'],
+ 'new_label': new_pred.predicted_label,
+ 'label_changed': record['old_label'] != new_pred.predicted_label,
+ 'old_confidence': record['old_confidence'],
+ 'new_confidence': new_pred.confidence_score,
+ 'human_label': record['human_label'],
+ }
+
+ if record['human_label'] is not None:
+ old_correct = str(record['old_label']) == str(record['human_label'])
+ new_correct = str(new_pred.predicted_label) == str(record['human_label'])
+ comp['old_correct'] = old_correct
+ comp['new_correct'] = new_correct
+ if new_correct and not old_correct:
+ improved += 1
+ elif old_correct and not new_correct:
+ worsened += 1
+ else:
+ unchanged += 1
+
+ comparisons.append(comp)
+
+ return {
+ 'total_queued': len(self._reannotation_history),
+ 'total_completed': len(comparisons),
+ 'improved': improved,
+ 'worsened': worsened,
+ 'unchanged': unchanged,
+ 'comparisons': comparisons,
+ }
+
+ # === Refinement Loop ===
+
+ def _maybe_trigger_refinement(self) -> None:
+ """Check if the refinement loop should trigger after an annotation."""
+ if not self.config.refinement_loop.enabled:
+ return
+
+ loop = self.refinement_loop
+ if not loop.record_annotation():
+ return
+
+ logger.info(
+ f"[Refinement] Trigger interval reached "
+ f"(agreement_rate={self.agreement_metrics.agreement_rate:.3f}, "
+ f"compared={self.agreement_metrics.total_compared}). "
+ f"Starting refinement cycle {loop.cycle_count + 1}..."
+ )
+
+ # Run in background thread to avoid blocking annotation flow
+ thread = threading.Thread(
+ target=self._run_refinement_cycle,
+ name="RefinementCycleThread",
+ daemon=True,
+ )
+ thread.start()
+
+ def _run_refinement_cycle(self) -> None:
+ """Execute a refinement cycle in a background thread."""
+ try:
+ self.trigger_refinement_cycle()
+ except Exception as e:
+ logger.error(f"Background refinement cycle failed: {e}")
+
+ # === New validated refinement framework ===
+
+ def trigger_refinement_cycle(self) -> Dict[str, Any]:
+ """Manually or automatically trigger a refinement cycle.
+
+ Dispatches to the new validated framework if the configured strategy
+ is in the refinement registry; otherwise the legacy path handles it.
+
+ Returns:
+ Dict with cycle results.
+ """
+ strategy_name = self.config.refinement_loop.refinement_strategy
+ # Check if strategy is in the new registry
+ try:
+ from .refinement import get_strategy
+ get_strategy(strategy_name)
+ return self._run_validated_refinement_cycle(strategy_name)
+ except KeyError:
+ # Not in new registry โ fall through to legacy path
+ pass
+ return self._run_legacy_refinement_cycle()
+
+ def _run_legacy_refinement_cycle(self) -> Dict[str, Any]:
+ """Legacy refinement path (focused_edit / generator_critic / append).
+
+ Kept for backward compatibility. New strategies should use the
+ validated framework via _run_validated_refinement_cycle.
+ """
+ loop = self.refinement_loop
+
+ if loop.is_stopped:
+ return {
+ 'success': False,
+ 'error': f'Refinement loop stopped: {loop.stop_reason}',
+ }
+
+ # Get current state
+ metrics = self.get_agreement_metrics()
+ agreement_rate = metrics.agreement_rate if hasattr(metrics, 'agreement_rate') else 0.0
+ prompt_version = self.current_prompt_version
+
+ # Check for post-cycle metrics from previous cycle
+ loop.record_post_cycle_metrics(agreement_rate)
+
+ # Get confusion patterns
+ analysis = self.get_confusion_analysis_full()
+ if not analysis.get('enabled'):
+ return {'success': False, 'error': 'Confusion analysis not enabled'}
+
+ # Build ConfusionPattern objects from the enriched data
+ from .confusion_analyzer import ConfusionPattern, ConfusionExample
+ patterns = []
+ for p_data in analysis.get('patterns', []):
+ patterns.append(ConfusionPattern(
+ predicted_label=p_data['predicted_label'],
+ actual_label=p_data['actual_label'],
+ count=p_data['count'],
+ percent=p_data['percent'],
+ examples=[
+ ConfusionExample(
+ instance_id=e['instance_id'],
+ text=e.get('text', ''),
+ llm_reasoning=e.get('llm_reasoning'),
+ llm_confidence=e.get('llm_confidence'),
+ )
+ for e in p_data.get('examples', [])
+ ],
+ ))
+
+ if not patterns:
+ return {'success': True, 'message': 'No confusion patterns found'}
+
+ # Define how to apply suggestions
+ def apply_suggestions(suggestions: List[str]) -> Dict[str, Any]:
+ import re as re_mod
+ current_prompt = self.get_current_prompt_text()
+ rules_section = "\n".join(f"- {s}" for s in suggestions)
+
+ strategy = self.config.refinement_loop.refinement_strategy
+
+ if strategy == "append":
+ # Legacy: just append (can cause contradictions)
+ guidelines_block = (
+ "## Refinement Guidelines\n\n"
+ "Based on observed confusion patterns:\n"
+ )
+ if guidelines_block in current_prompt:
+ updated = current_prompt.rstrip() + "\n" + rules_section + "\n"
+ else:
+ updated = current_prompt + f"\n\n{guidelines_block}{rules_section}\n"
+ else:
+ # focused_edit and generator_critic: replace the entire
+ # guidelines section with the new set of rules
+ guidelines_section = (
+ "## Annotation Guidelines\n\n"
+ "When distinguishing between similar labels, follow these rules:\n"
+ f"{rules_section}\n"
+ )
+ # Replace existing guidelines section or append if first time
+ if re_mod.search(
+ r'## (?:Refinement |Annotation )?Guidelines',
+ current_prompt
+ ):
+ updated = re_mod.sub(
+ r'## (?:Refinement |Annotation )?Guidelines.*',
+ guidelines_section,
+ current_prompt,
+ flags=re_mod.DOTALL,
+ )
+ else:
+ updated = current_prompt + "\n\n" + guidelines_section
+
+ old_version = self.current_prompt_version
+ new_pv = self.create_prompt_version(
+ updated,
+ created_by='refinement_loop',
+ source_description=(
+ f"Refinement cycle: {len(suggestions)} guideline suggestions"
+ ),
+ )
+ result = {
+ 'success': True,
+ 'new_prompt_version': new_pv.version,
+ 'categories_incorporated': len(suggestions),
+ 'reannotation_count': 0,
+ }
+ # Re-annotate low-confidence instances with the improved prompt
+ # to verify the refinement actually helps
+ reannotated = self._trigger_reannotation(old_version)
+ result['reannotation_count'] = reannotated
+ if reannotated > 0:
+ logger.info(
+ f"[Refinement] Re-annotating {reannotated} low-confidence "
+ f"instances with new prompt v{new_pv.version}"
+ )
+
+ return result
+
+ # Generate suggestions based on strategy
+ analyzer = self.confusion_analyzer
+ strategy = self.config.refinement_loop.refinement_strategy
+ current_prompt_text = self.get_current_prompt_text()
+
+ if strategy == "focused_edit":
+ # Single LLM call produces all guidelines as a coherent set
+ batch_guidelines = analyzer.generate_guidelines_rewrite(
+ patterns, current_prompt_text
+ )
+ # Pre-populate a queue so generate_suggestion just pops items
+ guideline_queue = list(batch_guidelines) if batch_guidelines else []
+ logger.info(
+ f"[Refinement] Focused edit generated {len(guideline_queue)} guidelines"
+ )
+
+ elif strategy == "generator_critic":
+ # Two-pass: generate candidates then critic-filter
+ batch_guidelines = analyzer.generate_and_critique_guidelines(
+ patterns, current_prompt_text
+ )
+ guideline_queue = list(batch_guidelines) if batch_guidelines else []
+ logger.info(
+ f"[Refinement] Generator-critic produced {len(guideline_queue)} approved guidelines"
+ )
+
+ else:
+ guideline_queue = None # Will use per-pattern generation
+
+ if guideline_queue is not None:
+ # Batch strategies: feed pre-generated guidelines one at a time
+ def generate_suggestion(pattern, current_prompt):
+ if guideline_queue:
+ return guideline_queue.pop(0)
+ return None
+ else:
+ # "append" or unknown: one suggestion per pattern (legacy)
+ def generate_suggestion(pattern, current_prompt):
+ return analyzer.suggest_guideline(pattern, current_prompt)
+
+ # Run the cycle
+ cycle = loop.run_cycle(
+ agreement_rate=agreement_rate,
+ prompt_version=prompt_version,
+ confusion_patterns=patterns,
+ apply_suggestions_fn=apply_suggestions,
+ generate_suggestion_fn=generate_suggestion,
+ current_prompt=current_prompt_text,
+ )
+
+ logger.info(
+ f"Refinement cycle {cycle.cycle_number} completed: "
+ f"status={cycle.status}, suggestions={cycle.suggestions_generated}"
+ )
+
+ return {
+ 'success': True,
+ 'cycle': cycle.to_dict(),
+ }
+
+ # === Validated refinement framework ===
+
+ def _get_icl_library(self):
+ """Lazy-initialize the persistent ICL library for this dataset."""
+ if not hasattr(self, '_icl_library') or self._icl_library is None:
+ from .refinement.icl_library import ICLLibrary
+ self._icl_library = ICLLibrary()
+ return self._icl_library
+
+ def _run_validated_refinement_cycle(self, strategy_name: str) -> Dict[str, Any]:
+ """Run a refinement cycle using the validated framework.
+
+ Flow:
+ 1. Load strategy from registry
+ 2. Split disagreements 70/30 into train/val
+ 3. Strategy proposes candidates based on train
+ 4. Evaluator scores each candidate on val set
+ 5. If best > baseline by min_improvement, apply (or queue for approval)
+ 6. Otherwise increment failure counter; stop after N consecutive failures
+ """
+ from .refinement import (
+ get_strategy,
+ ValidationSplit,
+ CandidateEvaluator,
+ )
+ from .refinement.base import CandidateKind, RefinementResult
+ from .refinement.icl_library import ICLEntry
+ from datetime import datetime
+
+ rl_config = self.config.refinement_loop
+ loop = self.refinement_loop
+
+ if loop.is_stopped:
+ return {'success': False, 'error': f'Refinement loop stopped: {loop.stop_reason}'}
+
+ # Record post-cycle metrics from the previous cycle
+ metrics = self.get_agreement_metrics()
+ agreement_rate = getattr(metrics, 'agreement_rate', 0.0)
+ loop.record_post_cycle_metrics(agreement_rate)
+
+ # Instantiate strategy
+ try:
+ strategy_cls = get_strategy(strategy_name)
+ except KeyError as e:
+ return {'success': False, 'error': str(e)}
+ strategy = strategy_cls(manager=self, solo_config=self.config)
+
+ # Get comparison history
+ comparisons = self.validation_tracker.get_comparison_history()
+ if not comparisons:
+ return {'success': False, 'error': 'No comparison history yet'}
+
+ # Split into train/val
+ splitter = ValidationSplit(
+ val_ratio=rl_config.validation_split_ratio,
+ min_val=rl_config.min_val_size,
+ prefer_consistent=rl_config.prefer_consistent_disagreements,
+ )
+ split_result = splitter.split(
+ comparisons, prompt_version=self.current_prompt_version
+ )
+
+ if not split_result.val:
+ logger.info(
+ f"[Refinement-Validated] Not enough disagreements for val split "
+ f"(need {rl_config.min_val_size}); skipping cycle"
+ )
+ return {
+ 'success': True,
+ 'message': 'Not enough disagreements for validation split',
+ 'strategy': strategy_name,
+ }
+
+ # Build confusion patterns from training data
+ patterns = self._build_patterns_from_comparisons(split_result.train)
+
+ # Let strategy propose candidates
+ current_prompt = self.get_current_prompt_text()
+ try:
+ candidates = strategy.propose_candidates(
+ patterns=patterns,
+ current_prompt=current_prompt,
+ train_comparisons=split_result.train,
+ )
+ except Exception as e:
+ logger.error(f"[Refinement-Validated] Strategy {strategy_name} propose_candidates failed: {e}")
+ return {'success': False, 'error': str(e)}
+
+ logger.info(
+ f"[Refinement-Validated] Strategy '{strategy_name}' proposed "
+ f"{len(candidates)} candidate(s); val size={len(split_result.val)}"
+ )
+
+ if not candidates:
+ self._handle_refinement_failure(strategy_name, reason='no_candidates')
+ return {
+ 'success': True,
+ 'strategy': strategy_name,
+ 'message': 'No candidates proposed',
+ 'failure_count': self._refinement_consecutive_failures,
+ }
+
+ # Build evaluator
+ evaluator = CandidateEvaluator(
+ label_fn=self._label_with_candidate,
+ get_text_fn=self._get_instance_text,
+ )
+
+ # Baseline: current prompt accuracy on val set
+ baseline_eval = evaluator.evaluate(
+ candidate_prompt=current_prompt,
+ val_comparisons=split_result.val,
+ sample_size=rl_config.eval_sample_size,
+ )
+ baseline_acc = baseline_eval.accuracy
+ val_sample_ids = [p['instance_id'] for p in baseline_eval.per_instance]
+
+ logger.info(
+ f"[Refinement-Validated] Baseline val accuracy: {baseline_acc:.3f} "
+ f"({baseline_eval.correct_count}/{baseline_eval.total})"
+ )
+
+ # Evaluate each candidate on the SAME val sample as baseline
+ candidate_accs = {}
+ best_idx = None
+ best_acc = baseline_acc + rl_config.min_val_improvement
+
+ for i, cand in enumerate(candidates):
+ try:
+ eval_prompt = self._build_eval_prompt_for_candidate(
+ cand, current_prompt
+ )
+ except Exception as e:
+ logger.warning(f"[Refinement-Validated] candidate {i} prompt build failed: {e}")
+ continue
+
+ # Evaluate on the same val sample as baseline
+ result = evaluator.evaluate(
+ candidate_prompt=eval_prompt,
+ val_comparisons=[c for c in split_result.val if c['instance_id'] in val_sample_ids],
+ )
+ candidate_accs[i] = result.accuracy
+ logger.info(
+ f"[Refinement-Validated] Candidate {i} ({cand.kind.value}, "
+ f"{cand.proposed_by}): {result.accuracy:.3f} "
+ f"({result.correct_count}/{result.total})"
+ )
+ if result.accuracy > best_acc:
+ best_acc = result.accuracy
+ best_idx = i
+
+ # Build result object
+ ref_result = RefinementResult(
+ success=False,
+ strategy=strategy_name,
+ all_candidates=candidates,
+ val_baseline_accuracy=baseline_acc,
+ val_candidate_accuracies=candidate_accs,
+ val_sample_ids=val_sample_ids,
+ train_sample_size=len(split_result.train),
+ val_sample_size=len(val_sample_ids),
+ dry_run=rl_config.dry_run,
+ )
+
+ # If no candidate beats baseline โ failure
+ if best_idx is None:
+ ref_result.failure_reason = 'no_candidate_beat_baseline'
+ self._handle_refinement_failure(strategy_name, reason='validation_failed')
+ self._log_refinement_cycle(ref_result)
+ return ref_result.to_dict() | {
+ 'message': f'No candidate beat baseline ({baseline_acc:.3f})',
+ 'failure_count': self._refinement_consecutive_failures,
+ }
+
+ winner = candidates[best_idx]
+ ref_result.applied_candidate = winner
+ ref_result.success = True
+
+ # Dry run: log but don't apply
+ if rl_config.dry_run:
+ logger.info(
+ f"[Refinement-Validated] DRY RUN: would apply candidate {best_idx} "
+ f"({winner.kind.value}, +{best_acc - baseline_acc:.3f} accuracy)"
+ )
+ ref_result.failure_reason = None
+ self._log_refinement_cycle(ref_result)
+ return ref_result.to_dict()
+
+ # Queue for approval OR apply immediately
+ if rl_config.require_approval:
+ self._queue_refinement_for_approval(ref_result)
+ logger.info(
+ f"[Refinement-Validated] Candidate queued for admin approval "
+ f"(+{best_acc - baseline_acc:.3f} improvement)"
+ )
+ return ref_result.to_dict() | {'status': 'queued_for_approval'}
+
+ # Apply the winning candidate
+ self._apply_refinement_candidate(winner, best_acc - baseline_acc)
+ self._refinement_consecutive_failures = 0 # success resets counter
+ self._log_refinement_cycle(ref_result)
+ logger.info(
+ f"[Refinement-Validated] APPLIED candidate {best_idx} "
+ f"({winner.kind.value}, +{best_acc - baseline_acc:.3f} "
+ f"over baseline {baseline_acc:.3f})"
+ )
+ self._save_state()
+
+ return ref_result.to_dict()
+
+ def _build_patterns_from_comparisons(self, comparisons: List[Dict[str, Any]]):
+ """Build ConfusionPattern list from a subset of comparison history."""
+ from .confusion_analyzer import ConfusionPattern, ConfusionExample
+ from collections import defaultdict
+
+ groups = defaultdict(list)
+ for c in comparisons:
+ if c.get('agrees'):
+ continue
+ key = (str(c['llm_label']), str(c['human_label']))
+ groups[key].append(c)
+
+ ca_config = self.config.confusion_analysis
+ patterns = []
+ total_disagreements = sum(1 for c in comparisons if not c.get('agrees'))
+
+ for (predicted, actual), records in groups.items():
+ if len(records) < ca_config.min_instances_for_pattern:
+ continue
+
+ percent = (len(records) / total_disagreements * 100) if total_disagreements > 0 else 0.0
+
+ examples = []
+ for record in records[:5]:
+ iid = record['instance_id']
+ text = self._get_instance_text(iid) or ''
+ examples.append(ConfusionExample(
+ instance_id=iid,
+ text=text[:200],
+ llm_reasoning=None,
+ llm_confidence=None,
+ ))
+
+ patterns.append(ConfusionPattern(
+ predicted_label=predicted,
+ actual_label=actual,
+ count=len(records),
+ percent=round(percent, 1),
+ examples=examples,
+ ))
+
+ patterns.sort(key=lambda p: p.count, reverse=True)
+ return patterns[:ca_config.max_patterns]
+
+ def _get_eval_endpoint(self) -> Optional[Any]:
+ """Get (or lazily create) the dedicated low-temperature endpoint used
+ for candidate evaluation.
+
+ The labeler's default temperature is tuned for sampling diversity
+ (non-zero, so confidence estimates have signal). The refinement gate
+ needs the opposite: measure prompt quality, not sampling variance.
+ So we keep a separate endpoint at rl_config.eval_temperature (0.0 by
+ default) and re-use it across cycles.
+ """
+ if self._eval_endpoint is not None:
+ return self._eval_endpoint
+
+ if not self.config.labeling_models:
+ return None
+
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+ except Exception as e:
+ logger.debug(f"[Refinement-Validated] eval endpoint factory unavailable: {e}")
+ return None
+
+ eval_temp = self.config.refinement_loop.eval_temperature
+ for model_config in self.config.labeling_models:
+ try:
+ endpoint_config = model_config.to_endpoint_config(
+ temperature_override=eval_temp
+ )
+ endpoint = AIEndpointFactory.create_endpoint(endpoint_config)
+ if endpoint:
+ self._eval_endpoint = endpoint
+ logger.info(
+ f"[Refinement-Validated] eval endpoint: "
+ f"{model_config.endpoint_type}/{model_config.model} "
+ f"(temperature={eval_temp})"
+ )
+ return endpoint
+ except Exception as e:
+ logger.debug(f"[Refinement-Validated] eval endpoint build failed for {model_config.model}: {e}")
+ continue
+ return None
+
+ def _label_with_candidate(
+ self, instance_id: str, text: str, candidate_prompt: str
+ ) -> Optional[str]:
+ """Single labeling call using a candidate prompt (no sampling diversity).
+
+ Used by CandidateEvaluator. Routes to the dedicated eval endpoint
+ (low/zero temperature) so the validation gate measures prompt quality
+ rather than sampling variance.
+ """
+ schemes = self.app_config.get('annotation_schemes', [])
+ schema_name = schemes[0].get('name', 'default') if schemes else 'default'
+
+ try:
+ endpoint = self._get_eval_endpoint()
+ if endpoint is None:
+ # Fall back to the labeler endpoint if the eval endpoint can't be
+ # built (e.g. during tests where AIEndpointFactory is mocked).
+ endpoint = self.llm_labeling_thread._get_endpoint()
+ if endpoint is None:
+ return None
+
+ labels = [l['name'] if isinstance(l, dict) else l for l in schemes[0].get('labels', [])]
+ full_prompt = (
+ f"{candidate_prompt}\n\n"
+ f"Text to label:\n{text}\n\n"
+ f"Available labels: {labels}\n\n"
+ f'Respond with JSON: {{"label": ""}}'
+ )
+
+ from pydantic import BaseModel
+
+ class LabelOnly(BaseModel):
+ label: str = ""
+
+ response = endpoint.query(full_prompt, LabelOnly)
+ if isinstance(response, dict):
+ return response.get('label', '').strip() or None
+ elif hasattr(response, 'label'):
+ return response.label
+ return None
+ except Exception as e:
+ logger.debug(f"[CandidateEval] label_fn failed for {instance_id}: {e}")
+ return None
+
+ def _build_eval_prompt_for_candidate(self, candidate, current_prompt: str) -> str:
+ """Given a candidate, construct the full prompt used for eval.
+
+ PROMPT_EDIT: the candidate payload contains the new_prompt_text.
+ ICL_EXAMPLE: inject the candidate example into the current prompt's
+ ## Examples section.
+ """
+ from .refinement.base import CandidateKind
+
+ if candidate.kind == CandidateKind.PROMPT_EDIT:
+ return candidate.payload.get('new_prompt_text', current_prompt)
+
+ if candidate.kind == CandidateKind.ICL_EXAMPLE:
+ # Inject the single example into the current prompt
+ example_text = candidate.payload.get('text', '')
+ example_label = candidate.payload.get('label', '')
+ example_section = (
+ "\n\n## Examples\n"
+ f'Text: "{example_text[:200]}"\n'
+ f"Label: {example_label}\n"
+ )
+ return current_prompt + example_section
+
+ if candidate.kind == CandidateKind.PRINCIPLE:
+ return current_prompt + f"\n\nKey principle: {candidate.payload}\n"
+
+ return current_prompt
+
+ def _apply_refinement_candidate(self, candidate, gain: float) -> None:
+ """Commit a candidate: either create a new prompt version or add to ICL library."""
+ from .refinement.base import CandidateKind
+ from .refinement.icl_library import ICLEntry
+
+ old_version = self.current_prompt_version
+
+ if candidate.kind == CandidateKind.PROMPT_EDIT:
+ new_prompt_text = candidate.payload.get('new_prompt_text', '')
+ if new_prompt_text:
+ new_pv = self.create_prompt_version(
+ new_prompt_text,
+ created_by='validated_refinement',
+ source_description=(
+ f"{candidate.proposed_by}: +{gain:.3f} val accuracy"
+ ),
+ )
+ logger.info(f"[Refinement-Validated] Created prompt v{new_pv.version}")
+ # Trigger re-annotation of low-confidence instances
+ self._trigger_reannotation(old_version)
+
+ elif candidate.kind == CandidateKind.ICL_EXAMPLE:
+ lib = self._get_icl_library()
+ payload = candidate.payload
+ entry = ICLEntry(
+ instance_id=payload['instance_id'],
+ text=payload.get('text', ''),
+ label=payload.get('label', ''),
+ principle=payload.get('principle', ''),
+ added_at_cycle=self.refinement_loop.cycle_count + 1,
+ val_accuracy_gain=gain,
+ )
+ lib.add(entry)
+ logger.info(f"[Refinement-Validated] Added ICL example {entry.instance_id}")
+ # Also bump prompt version to trigger re-annotation
+ # (ICL is effectively a new prompt since labeler injects examples)
+ current = self.get_current_prompt_text()
+ new_pv = self.create_prompt_version(
+ current,
+ created_by='validated_refinement_icl',
+ source_description=(
+ f"ICL example added: {entry.instance_id} (+{gain:.3f} val accuracy)"
+ ),
+ )
+ # Trigger re-annotation to apply new ICL library
+ self._trigger_reannotation(old_version)
+
+ elif candidate.kind == CandidateKind.PRINCIPLE:
+ current = self.get_current_prompt_text()
+ new_prompt = current + f"\n\nKey principle: {candidate.payload}\n"
+ new_pv = self.create_prompt_version(
+ new_prompt,
+ created_by='validated_refinement_principle',
+ source_description=f"Principle added (+{gain:.3f} val accuracy)",
+ )
+ self._trigger_reannotation(old_version)
+
+ def _handle_refinement_failure(self, strategy_name: str, reason: str) -> None:
+ """Track a failed refinement cycle; stop after max consecutive failures."""
+ if not hasattr(self, '_refinement_consecutive_failures'):
+ self._refinement_consecutive_failures = 0
+ self._refinement_consecutive_failures += 1
+
+ max_failures = self.config.refinement_loop.max_consecutive_failures
+ if self._refinement_consecutive_failures >= max_failures:
+ logger.warning(
+ f"[Refinement-Validated] {strategy_name} failed "
+ f"{self._refinement_consecutive_failures} consecutive cycles "
+ f"(reason: {reason}); stopping refinement until more disagreements arrive"
+ )
+ # Stop the loop; it will be reset when new disagreements accumulate
+ # (handled by the trigger_interval mechanism in refinement_loop)
+ self.refinement_loop._stop(
+ f"Validation failed {self._refinement_consecutive_failures} times"
+ )
+
+ def _queue_refinement_for_approval(self, ref_result) -> None:
+ """Store a validated refinement candidate awaiting admin approval."""
+ if not hasattr(self, '_pending_refinements'):
+ self._pending_refinements = []
+ self._pending_refinements.append(ref_result.to_dict())
+
+ def _log_refinement_cycle(self, ref_result) -> None:
+ """Append a cycle result to the persistent refinement log."""
+ if not hasattr(self, '_refinement_log'):
+ self._refinement_log = []
+ self._refinement_log.append(ref_result.to_dict())
+
+ def get_refinement_log(self) -> List[Dict[str, Any]]:
+ """Return full log of all refinement cycles (including dry-run results)."""
+ return getattr(self, '_refinement_log', [])
+
+ def get_pending_refinements(self) -> List[Dict[str, Any]]:
+ """Return candidates awaiting admin approval."""
+ return getattr(self, '_pending_refinements', [])
+
+ def approve_pending_refinement(self, index: int) -> Dict[str, Any]:
+ """Apply a pending refinement by index. Returns {success, message}."""
+ pending = getattr(self, '_pending_refinements', [])
+ if index < 0 or index >= len(pending):
+ return {'success': False, 'error': 'Invalid index'}
+
+ item = pending.pop(index)
+ # Reconstruct a candidate from the stored dict
+ from .refinement.base import CandidateKind, RefinementCandidate
+
+ cand_dict = item.get('applied_candidate')
+ if not cand_dict:
+ return {'success': False, 'error': 'No candidate in pending item'}
+
+ cand = RefinementCandidate(
+ kind=CandidateKind(cand_dict['kind']),
+ payload=cand_dict['payload'],
+ target_pattern=cand_dict.get('target_pattern'),
+ proposed_by=cand_dict.get('proposed_by', ''),
+ rationale=cand_dict.get('rationale', ''),
+ )
+ gain = (
+ max(item.get('val_candidate_accuracies', {}).values())
+ - item.get('val_baseline_accuracy', 0.0)
+ if item.get('val_candidate_accuracies') else 0.0
+ )
+ self._apply_refinement_candidate(cand, gain)
+ self._save_state()
+ return {'success': True, 'applied': cand_dict}
+
+ def reject_pending_refinement(self, index: int) -> Dict[str, Any]:
+ """Reject a pending refinement by index."""
+ pending = getattr(self, '_pending_refinements', [])
+ if index < 0 or index >= len(pending):
+ return {'success': False, 'error': 'Invalid index'}
+ item = pending.pop(index)
+ return {'success': True, 'rejected': item.get('applied_candidate')}
+
+ def get_refinement_status(self) -> Dict[str, Any]:
+ """Get the refinement loop status."""
+ if not self.config.refinement_loop.enabled:
+ return {'enabled': False}
+
+ return self.refinement_loop.get_status()
+
+ # === Labeling Functions ===
+
+ def get_labeling_function_status(self) -> Dict[str, Any]:
+ """Get labeling function statistics."""
+ if not self.config.labeling_functions.enabled:
+ return {'enabled': False}
+
+ return self.labeling_function_manager.get_stats()
+
+ def extract_labeling_functions(self) -> Dict[str, Any]:
+ """Extract labeling functions from high-confidence predictions.
+
+ Returns:
+ Dict with success status and extracted function count.
+ """
+ if not self.config.labeling_functions.enabled:
+ return {'success': False, 'error': 'Labeling functions not enabled'}
+
+ min_conf = self.config.labeling_functions.min_confidence
+
+ # Build prediction list from stored predictions
+ pred_list = []
+ with self._lock:
+ for instance_id, schemas in self.predictions.items():
+ for schema_name, pred in schemas.items():
+ if pred.confidence_score >= min_conf:
+ pred_list.append({
+ 'instance_id': instance_id,
+ 'text': self._get_instance_text(instance_id),
+ 'predicted_label': str(pred.predicted_label),
+ 'confidence': pred.confidence_score,
+ 'reasoning': pred.reasoning,
+ })
+
+ if not pred_list:
+ return {
+ 'success': True,
+ 'extracted': 0,
+ 'message': 'No high-confidence predictions available',
+ }
+
+ new_fns = self.labeling_function_manager.extract_functions(pred_list)
+
+ return {
+ 'success': True,
+ 'extracted': len(new_fns),
+ 'total': len(self.labeling_function_manager.get_all_functions()),
+ 'functions': [f.to_dict() for f in new_fns],
+ }
+
+ def _maybe_extract_labeling_functions(self) -> None:
+ """Check if auto-extraction should trigger after labeling."""
+ lf_config = self.config.labeling_functions
+ if not lf_config.enabled or not lf_config.auto_extract:
+ return
+
+ # Auto-extract every 100 new LLM labels if we have enough data
+ with self._lock:
+ total_predictions = sum(
+ 1 for schemas in self.predictions.values()
+ for pred in schemas.values()
+ if pred.confidence_score >= lf_config.min_confidence
+ )
+
+ mgr = self.labeling_function_manager
+ existing = len(mgr.get_all_functions())
+
+ # Extract when we have enough new data and don't already have many functions
+ if total_predictions >= 20 and existing < lf_config.max_functions:
+ # Only extract if we have significantly more predictions than functions
+ if total_predictions >= (existing + 1) * 10:
+ thread = threading.Thread(
+ target=self._run_labeling_function_extraction,
+ name="LabelingFunctionExtractionThread",
+ daemon=True,
+ )
+ thread.start()
+
+ def _run_labeling_function_extraction(self) -> None:
+ """Run labeling function extraction in a background thread."""
+ try:
+ self.extract_labeling_functions()
+ except Exception as e:
+ logger.error(f"Background labeling function extraction failed: {e}")
+
+ # === Phase Control ===
+
+ def get_current_phase(self) -> SoloPhase:
+ """Get the current workflow phase."""
+ return self.phase_controller.get_current_phase()
+
+ def advance_to_phase(
+ self,
+ phase: SoloPhase,
+ reason: str = "",
+ force: bool = False
+ ) -> bool:
+ """
+ Transition to a specific phase.
+
+ Args:
+ phase: Target phase
+ reason: Reason for transition
+ force: Allow invalid transitions
+
+ Returns:
+ True if transition successful
+ """
+ old_phase = self.phase_controller.get_current_phase()
+ result = self.phase_controller.transition_to(phase, reason=reason, force=force)
+ if result:
+ logger.info(
+ f"[Phase Transition] {old_phase.name} -> {phase.name}"
+ f"{' (forced)' if force else ''}"
+ f"{f' reason: {reason}' if reason else ''}"
+ )
+ if phase in (SoloPhase.PARALLEL_ANNOTATION, SoloPhase.ACTIVE_ANNOTATION):
+ self.start_background_labeling()
+ else:
+ logger.warning(
+ f"[Phase Transition] FAILED: {old_phase.name} -> {phase.name}"
+ )
+ return result
+
+ def advance_to_next_phase(self, reason: str = "") -> bool:
+ """Advance to the next logical phase."""
+ return self.phase_controller.advance_to_next_phase(reason=reason)
+
+ # === Prompt Management ===
+
+ def get_current_prompt(self) -> Optional[PromptVersion]:
+ """Get the current prompt version."""
+ with self._lock:
+ if not self.prompt_versions:
+ return None
+ return self.prompt_versions[self.current_prompt_version - 1]
+
+ def get_prompt_version(self, version: int) -> Optional[PromptVersion]:
+ """Get a specific prompt version."""
+ with self._lock:
+ if 0 < version <= len(self.prompt_versions):
+ return self.prompt_versions[version - 1]
+ return None
+
+ def get_all_prompt_versions(self) -> List[PromptVersion]:
+ """Get all prompt versions."""
+ with self._lock:
+ return self.prompt_versions.copy()
+
+ def create_prompt_version(
+ self,
+ prompt_text: str,
+ created_by: str,
+ source_description: str = ""
+ ) -> PromptVersion:
+ """
+ Create a new prompt version.
+
+ Args:
+ prompt_text: The prompt text
+ created_by: Who created it ('user', 'llm_synthesis', 'llm_optimization')
+ source_description: Description of how it was created
+
+ Returns:
+ The new PromptVersion
+ """
+ with self._lock:
+ new_version = len(self.prompt_versions) + 1
+ parent = self.current_prompt_version if self.current_prompt_version > 0 else None
+
+ prompt = PromptVersion(
+ version=new_version,
+ prompt_text=prompt_text,
+ created_at=datetime.now(),
+ created_by=created_by,
+ source_description=source_description,
+ parent_version=parent,
+ )
+
+ self.prompt_versions.append(prompt)
+ self.current_prompt_version = new_version
+
+ # Reset stale reannotation counts so instances can be re-annotated
+ # with the improved prompt. Keep counts only for recent prompt versions.
+ self._reset_stale_reannotation_counts(new_version)
+
+ self._save_state()
+
+ logger.info(f"Created prompt version {new_version} by {created_by}")
+ return prompt
+
+ def _reset_stale_reannotation_counts(self, current_version: int) -> None:
+ """Reset reannotation counts for instances not recently re-annotated.
+
+ Keeps counts only for instances whose last reannotation was within
+ the last 2 prompt versions. This prevents instances from being
+ permanently excluded from re-annotation after prompt improvements.
+ """
+ if not self._reannotation_counts:
+ return
+
+ stale_ids = []
+ for instance_id in self._reannotation_counts:
+ # Check if this instance's prediction is from a recent prompt version
+ if instance_id in self.predictions:
+ for schema_preds in self.predictions[instance_id].values():
+ if current_version - schema_preds.prompt_version > 2:
+ stale_ids.append(instance_id)
+ break
+ else:
+ stale_ids.append(instance_id)
+
+ for instance_id in stale_ids:
+ del self._reannotation_counts[instance_id]
+
+ if stale_ids:
+ logger.debug(
+ f"Reset reannotation counts for {len(stale_ids)} stale instances"
+ )
+
+ def update_prompt(
+ self,
+ prompt_text: str,
+ source: str,
+ source_description: str = ""
+ ) -> PromptVersion:
+ """
+ Update the prompt by creating a new version.
+
+ This is a convenience method that wraps create_prompt_version.
+ """
+ return self.create_prompt_version(prompt_text, source, source_description)
+
+ def set_task_description(self, description: str) -> None:
+ """Set the task description for prompt synthesis."""
+ with self._lock:
+ self.task_description = description
+ self._save_state()
+
+ def get_task_description(self) -> str:
+ """Get the task description."""
+ with self._lock:
+ return self.task_description
+
+ # === LLM Prediction Management ===
+
+ def set_llm_prediction(
+ self,
+ instance_id: str,
+ schema_name: str,
+ prediction: LLMPrediction
+ ) -> None:
+ """
+ Store an LLM prediction for an instance.
+
+ Args:
+ instance_id: The instance ID
+ schema_name: The annotation schema name
+ prediction: The LLM prediction
+ """
+ with self._lock:
+ if instance_id not in self.predictions:
+ self.predictions[instance_id] = {}
+ self.predictions[instance_id][schema_name] = prediction
+ self.llm_labeled_ids.add(instance_id)
+
+ # Track confidence history for cartography
+ if instance_id not in self.confidence_history:
+ self.confidence_history[instance_id] = []
+ self.confidence_history[instance_id].append(
+ (prediction.prompt_version, prediction.confidence_score)
+ )
+
+ def get_llm_prediction(
+ self,
+ instance_id: str,
+ schema_name: str
+ ) -> Optional[LLMPrediction]:
+ """Get the LLM prediction for an instance and schema."""
+ with self._lock:
+ if instance_id in self.predictions:
+ return self.predictions[instance_id].get(schema_name)
+ return None
+
+ def get_all_llm_predictions(self) -> Dict[str, Dict[str, LLMPrediction]]:
+ """Get all LLM predictions."""
+ with self._lock:
+ return {
+ iid: {s: p for s, p in schemas.items()}
+ for iid, schemas in self.predictions.items()
+ }
+
+ def get_predictions_by_confidence(
+ self,
+ min_confidence: Optional[float] = None,
+ max_confidence: Optional[float] = None
+ ) -> List[LLMPrediction]:
+ """Get predictions filtered by confidence range."""
+ with self._lock:
+ results = []
+ for schemas in self.predictions.values():
+ for prediction in schemas.values():
+ conf = prediction.confidence_score
+ if min_confidence is not None and conf < min_confidence:
+ continue
+ if max_confidence is not None and conf > max_confidence:
+ continue
+ results.append(prediction)
+ return results
+
+ def get_low_confidence_predictions(self) -> List[LLMPrediction]:
+ """Get predictions below the low confidence threshold."""
+ return self.get_predictions_by_confidence(
+ max_confidence=self.config.thresholds.confidence_low
+ )
+
+ # === Human Label Recording ===
+
+ def record_human_label(
+ self,
+ instance_id: str,
+ schema_name: str,
+ label: Any,
+ user_id: str
+ ) -> Optional[bool]:
+ """
+ Record a human label and compare with LLM prediction.
+
+ Args:
+ instance_id: The instance ID
+ schema_name: The annotation schema
+ label: The human's label
+ user_id: The annotator ID
+
+ Returns:
+ True if agrees with LLM, False if disagrees, None if no LLM prediction
+ """
+ with self._lock:
+ self.human_labeled_ids.add(instance_id)
+
+ prediction = self.get_llm_prediction(instance_id, schema_name)
+ if prediction is None:
+ return None
+
+ prediction.human_label = label
+ agrees = self._check_agreement(
+ prediction.predicted_label,
+ label,
+ schema_name
+ )
+ prediction.agrees_with_human = agrees
+
+ # Update agreement metrics
+ self.agreement_metrics.total_compared += 1
+ if agrees:
+ self.agreement_metrics.agreements += 1
+ else:
+ self.agreement_metrics.disagreements += 1
+ self.disagreement_ids.add(instance_id)
+ self.agreement_metrics.update_rate()
+
+ # Track per-prompt-version agreement
+ pv = prediction.prompt_version
+ if pv not in self._per_version_agreement:
+ self._per_version_agreement[pv] = {'compared': 0, 'agreements': 0}
+ self._per_version_agreement[pv]['compared'] += 1
+ if agrees:
+ self._per_version_agreement[pv]['agreements'] += 1
+
+ # Feed the validation tracker for confusion matrix / pattern analysis
+ self.validation_tracker.record_comparison(
+ instance_id=instance_id,
+ human_label=label,
+ llm_label=prediction.predicted_label,
+ schema_name=schema_name,
+ agrees=agrees,
+ )
+
+ human_count = len(self.human_labeled_ids)
+ pv_stats = self._per_version_agreement.get(pv, {})
+ pv_rate = (pv_stats['agreements'] / pv_stats['compared']
+ if pv_stats.get('compared', 0) > 0 else 0)
+ if human_count % 5 == 0 or not agrees:
+ logger.info(
+ f"[Human Label] #{human_count} {instance_id}: "
+ f"human={label}, llm={prediction.predicted_label}, "
+ f"{'AGREE' if agrees else 'DISAGREE'} "
+ f"(overall={self.agreement_metrics.agreement_rate:.3f}, "
+ f"prompt_v{pv}={pv_rate:.3f} [{pv_stats.get('compared',0)}], "
+ f"total_compared={self.agreement_metrics.total_compared})"
+ )
+
+ self._save_state()
+ return agrees
+
+ def _check_agreement(
+ self,
+ llm_label: Any,
+ human_label: Any,
+ schema_name: str
+ ) -> bool:
+ """
+ Check if LLM and human labels agree.
+
+ The agreement check depends on the annotation type.
+ """
+ # Get annotation type for this schema
+ annotation_type = self._get_annotation_type(schema_name)
+
+ if annotation_type in ('radio', 'select'):
+ # Exact match for categorical
+ return str(llm_label) == str(human_label)
+
+ elif annotation_type == 'likert':
+ # Within tolerance for likert scales
+ try:
+ tolerance = self.config.thresholds.likert_tolerance
+ return abs(int(llm_label) - int(human_label)) <= tolerance
+ except (ValueError, TypeError):
+ return str(llm_label) == str(human_label)
+
+ elif annotation_type == 'multiselect':
+ # Jaccard similarity for multiselect
+ threshold = self.config.thresholds.multiselect_jaccard_threshold
+ llm_set = set(llm_label) if isinstance(llm_label, (list, set)) else {llm_label}
+ human_set = set(human_label) if isinstance(human_label, (list, set)) else {human_label}
+
+ if not llm_set and not human_set:
+ return True
+
+ intersection = len(llm_set & human_set)
+ union = len(llm_set | human_set)
+ jaccard = intersection / union if union > 0 else 0
+ return jaccard >= threshold
+
+ elif annotation_type == 'textbox':
+ # For now, exact match; could use embedding similarity
+ return str(llm_label).strip().lower() == str(human_label).strip().lower()
+
+ elif annotation_type == 'span':
+ # Token overlap for spans
+ threshold = self.config.thresholds.span_overlap_threshold
+ # Simplified: check if spans overlap sufficiently
+ # Full implementation would compare token ranges
+ return str(llm_label) == str(human_label)
+
+ else:
+ # Default to exact match
+ return str(llm_label) == str(human_label)
+
+ def _get_annotation_type(self, schema_name: str) -> str:
+ """Get the annotation type for a schema."""
+ schemes = self.app_config.get('annotation_schemes', [])
+ for scheme in schemes:
+ if scheme.get('name') == schema_name:
+ return scheme.get('annotation_type', 'radio')
+ return 'radio'
+
+ def _retroactive_compare(self, instance_id: str, schema_name: str) -> None:
+ """Compare an LLM prediction against an existing human label.
+
+ Called when the LLM labels an instance that a human already annotated.
+ This ensures agreement metrics are updated regardless of annotation order.
+ """
+ with self._lock:
+ if instance_id not in self.human_labeled_ids:
+ return
+
+ prediction = self.get_llm_prediction(instance_id, schema_name)
+ if prediction is None or prediction.human_label is not None:
+ return # No prediction or already compared
+
+ human_label = self._get_stored_human_label(instance_id, schema_name)
+ if human_label is None:
+ return
+
+ prediction.human_label = human_label
+ agrees = self._check_agreement(
+ prediction.predicted_label, human_label, schema_name
+ )
+ prediction.agrees_with_human = agrees
+
+ self.agreement_metrics.total_compared += 1
+ if agrees:
+ self.agreement_metrics.agreements += 1
+ else:
+ self.agreement_metrics.disagreements += 1
+ self.disagreement_ids.add(instance_id)
+ self.agreement_metrics.update_rate()
+
+ # Track per-prompt-version agreement
+ pv = prediction.prompt_version
+ if pv not in self._per_version_agreement:
+ self._per_version_agreement[pv] = {'compared': 0, 'agreements': 0}
+ self._per_version_agreement[pv]['compared'] += 1
+ if agrees:
+ self._per_version_agreement[pv]['agreements'] += 1
+
+ # Feed the validation tracker for confusion matrix / pattern analysis
+ self.validation_tracker.record_comparison(
+ instance_id=instance_id,
+ human_label=human_label,
+ llm_label=prediction.predicted_label,
+ schema_name=schema_name,
+ agrees=agrees,
+ )
+
+ logger.debug(
+ f"Retroactive comparison for {instance_id}: "
+ f"llm={prediction.predicted_label}, human={human_label}, "
+ f"agrees={agrees}, prompt_v{pv}"
+ )
+
+ def _get_stored_human_label(
+ self, instance_id: str, schema_name: str
+ ) -> Optional[Any]:
+ """Look up a human annotation label from the user state manager.
+
+ Returns:
+ The human label if found, None otherwise.
+ """
+ try:
+ from potato.user_state_management import get_user_state_manager
+ usm = get_user_state_manager()
+ # Check all users' annotations for this instance
+ for user_id in usm.get_all_user_ids():
+ user_state = usm.get_user_state(user_id)
+ if user_state is None:
+ continue
+ annotations = user_state.get_annotations_for_instance(instance_id)
+ if annotations and schema_name in annotations:
+ return annotations[schema_name]
+ except Exception as e:
+ logger.debug(f"Could not look up human label for {instance_id}: {e}")
+ return None
+
+ # === Disagreement Resolution ===
+
+ def get_pending_disagreements(self) -> List[str]:
+ """Get instance IDs with unresolved disagreements."""
+ with self._lock:
+ pending = []
+ for instance_id in self.disagreement_ids:
+ if instance_id in self.predictions:
+ for prediction in self.predictions[instance_id].values():
+ if not prediction.disagreement_resolved:
+ pending.append(instance_id)
+ break
+ return pending
+
+ def resolve_disagreement(
+ self,
+ instance_id: str,
+ schema_name: str,
+ resolution_label: Any,
+ resolved_by: str
+ ) -> bool:
+ """
+ Resolve a human-LLM disagreement.
+
+ Args:
+ instance_id: The instance ID
+ schema_name: The annotation schema
+ resolution_label: The final resolved label
+ resolved_by: Who resolved it ('human', 'llm_revision')
+
+ Returns:
+ True if resolution was recorded
+ """
+ with self._lock:
+ prediction = self.get_llm_prediction(instance_id, schema_name)
+ if prediction is None:
+ return False
+
+ prediction.disagreement_resolved = True
+ prediction.resolution_label = resolution_label
+
+ self._save_state()
+ logger.info(
+ f"Resolved disagreement for {instance_id}:{schema_name} "
+ f"(resolved_by={resolved_by})"
+ )
+ return True
+
+ # === Instance Selection ===
+
+ def get_next_instance_for_human(self, user_id: str) -> Optional[str]:
+ """
+ Get the next instance for human annotation.
+
+ Uses weighted selection across pools:
+ - Low LLM confidence
+ - Diversity (embedding clusters)
+ - Random sampling
+ - Prior disagreements
+ - Edge case rules
+ - Cartography (high confidence variability)
+
+ Args:
+ user_id: The annotator's ID
+
+ Returns:
+ Instance ID to annotate, or None if none available
+ """
+ with self._lock:
+ from potato.item_state_management import get_item_state_manager
+
+ try:
+ ism = get_item_state_manager()
+ except ValueError:
+ return None
+
+ # Compute available IDs: all instances minus already human-labeled
+ all_ids = set(ism.instance_id_ordering)
+ available = all_ids - self.human_labeled_ids
+
+ if not available:
+ return None
+
+ # Convert predictions to dict format for refresh_pools
+ pred_dicts = {}
+ for iid, schemas in self.predictions.items():
+ pred_dicts[iid] = {
+ s: p.to_dict() for s, p in schemas.items()
+ }
+
+ # Get edge case rule IDs if available
+ edge_case_rule_ids = None
+ if self._edge_case_rule_manager is not None:
+ try:
+ edge_case_rule_ids = self._edge_case_rule_manager.get_rule_instance_ids()
+ except Exception:
+ pass
+
+ # Compute cartography scores if history available
+ cartography_variability = None
+ if self.confidence_history:
+ cartography = self.get_cartography_scores()
+ if cartography:
+ cartography_variability = {
+ iid: s['variability']
+ for iid, s in cartography.items()
+ }
+
+ # Refresh pools with current data
+ self.instance_selector.refresh_pools(
+ available_ids=available,
+ llm_predictions=pred_dicts,
+ disagreement_ids=self.disagreement_ids,
+ confidence_threshold=self.config.thresholds.confidence_low,
+ edge_case_rule_ids=edge_case_rule_ids,
+ cartography_scores=cartography_variability,
+ )
+
+ # Select next instance
+ return self.instance_selector.select_next(
+ available_ids=available,
+ exclude_ids=self.human_labeled_ids,
+ )
+
+ def get_cartography_scores(self) -> Dict[str, Dict[str, float]]:
+ """Compute cartography signals for each instance.
+
+ Uses confidence history across prompt versions to identify:
+ - Ambiguous instances: high confidence variability
+ - Hard instances: consistently low confidence
+ - Easy instances: consistently high confidence
+
+ Returns:
+ Dict of instance_id -> {variability, mean_confidence}
+ """
+ import statistics
+
+ with self._lock:
+ scores = {}
+ for instance_id, history in self.confidence_history.items():
+ if not history:
+ continue
+
+ confidences = [conf for _, conf in history]
+ mean_conf = statistics.mean(confidences)
+ variability = (
+ statistics.stdev(confidences) if len(confidences) > 1 else 0.0
+ )
+
+ scores[instance_id] = {
+ 'variability': variability,
+ 'mean_confidence': mean_conf,
+ }
+ return scores
+
+ # === Agreement Metrics ===
+
+ def get_agreement_metrics(self) -> AgreementMetrics:
+ """Get current agreement metrics."""
+ with self._lock:
+ return self.agreement_metrics
+
+ def should_end_human_annotation(self) -> bool:
+ """
+ Check if human annotation should end.
+
+ Returns True when agreement threshold is reached and
+ minimum validation sample size is met.
+ """
+ with self._lock:
+ metrics = self.agreement_metrics
+ threshold = self.config.thresholds.end_human_annotation_agreement
+ min_sample = self.config.thresholds.minimum_validation_sample
+
+ if metrics.total_compared < min_sample:
+ return False
+
+ return metrics.agreement_rate >= threshold
+
+ def check_and_advance_to_autonomous(self) -> bool:
+ """
+ Atomically check if human annotation should end and advance phase if so.
+
+ This prevents race conditions where multiple requests could both
+ check should_end_human_annotation() as True and try to advance.
+
+ Returns:
+ True if phase was advanced to AUTONOMOUS_LABELING
+ """
+ with self._lock:
+ metrics = self.agreement_metrics
+ threshold = self.config.thresholds.end_human_annotation_agreement
+ min_sample = self.config.thresholds.minimum_validation_sample
+
+ if metrics.total_compared < min_sample:
+ return False
+
+ if metrics.agreement_rate < threshold:
+ return False
+
+ # Already in or past autonomous labeling phase
+ current_phase = self.phase_controller.get_current_phase()
+ if current_phase.value >= SoloPhase.AUTONOMOUS_LABELING.value:
+ return False
+
+ # Advance phase atomically
+ return self.phase_controller.transition_to(
+ SoloPhase.AUTONOMOUS_LABELING,
+ reason="Agreement threshold reached"
+ )
+
+ def should_trigger_periodic_review(self) -> bool:
+ """Check if periodic review should be triggered."""
+ with self._lock:
+ interval = self.config.thresholds.periodic_review_interval
+ return len(self.llm_labeled_ids) % interval == 0
+
+ # === Background Labeling ===
+
+ def start_background_labeling(self) -> bool:
+ """
+ Start background LLM labeling thread.
+
+ Returns:
+ True if started, False if already running
+ """
+ with self._lock:
+ if self._labeling_thread is not None and self._labeling_thread.is_alive():
+ logger.warning("Background labeling already running")
+ return False
+
+ self._stop_labeling.clear()
+ self._pause_labeling.clear()
+ self._labeling_thread = threading.Thread(
+ target=self._background_labeling_loop,
+ name="SoloModeLabelingThread",
+ daemon=True
+ )
+ self._labeling_thread.start()
+ logger.info("Started background LLM labeling")
+ return True
+
+ def stop_background_labeling(self) -> None:
+ """Stop background LLM labeling thread."""
+ if self._labeling_thread is None:
+ return
+
+ self._stop_labeling.set()
+ self._labeling_thread.join(timeout=5.0)
+ self._labeling_thread = None
+ logger.info("Stopped background LLM labeling")
+
+ def pause_background_labeling(self) -> bool:
+ """Pause the background labeling loop without tearing down the thread.
+
+ Returns True if the loop was running and is now paused, False if
+ nothing was running.
+ """
+ if not self.is_background_labeling_running():
+ return False
+ self._pause_labeling.set()
+ logger.info("Paused background LLM labeling")
+ return True
+
+ def resume_background_labeling(self) -> bool:
+ """Resume a paused background labeling loop.
+
+ Returns True if a paused loop was resumed, False if nothing was paused.
+ """
+ if not self.is_background_labeling_running():
+ return False
+ was_paused = self._pause_labeling.is_set()
+ self._pause_labeling.clear()
+ if was_paused:
+ logger.info("Resumed background LLM labeling")
+ return was_paused
+
+ def is_background_labeling_paused(self) -> bool:
+ return self._pause_labeling.is_set()
+
+ def is_background_labeling_running(self) -> bool:
+ """Check if background labeling is running (paused counts as running)."""
+ return (
+ self._labeling_thread is not None and
+ self._labeling_thread.is_alive()
+ )
+
+ def _background_labeling_loop(self) -> None:
+ """Main loop for background labeling."""
+ import time
+
+ batch_size = self.config.batches.llm_labeling_batch
+ max_labels = self.config.batches.max_parallel_labels
+
+ total_instances = self._get_total_instance_count()
+ logger.info(
+ f"[LLM Background] Labeling started "
+ f"(batch={batch_size}, max={max_labels}, "
+ f"total_instances={total_instances}, "
+ f"already_labeled={len(self.llm_labeled_ids)})"
+ )
+
+ while not self._stop_labeling.is_set():
+ try:
+ # Honor pause requests: sleep in short polls so resume is responsive.
+ if self._pause_labeling.is_set():
+ self._stop_labeling.wait(2)
+ continue
+
+ # Check if we've hit the max parallel labels
+ with self._lock:
+ current_count = len(self.llm_labeled_ids - self.human_labeled_ids)
+ if current_count >= max_labels:
+ logger.debug(f"Max parallel labels reached ({current_count})")
+ time.sleep(10)
+ continue
+
+ # Label a batch of instances
+ labeled_count = self._label_batch(batch_size)
+
+ if labeled_count == 0:
+ # No more instances to label
+ time.sleep(30)
+ else:
+ logger.info(f"Labeled {labeled_count} instances in background")
+ self._save_state()
+
+ except Exception as e:
+ logger.error(f"Error in background labeling: {e}")
+ time.sleep(10)
+
+ # Wait before next batch
+ self._stop_labeling.wait(5)
+
+ def _label_batch(self, batch_size: int) -> int:
+ """Label a batch of instances. Returns number labeled.
+
+ Tries labeling functions first (cheap, no API calls), then
+ falls through to LLM labeling for remaining instances.
+ """
+ instances = self._get_instances_for_labeling(batch_size)
+ if not instances:
+ return 0
+
+ labeled = 0
+
+ # Try labeling functions first (no API cost)
+ remaining = instances
+ if self.config.labeling_functions.enabled:
+ lf_results, remaining = self.labeling_function_manager.apply_batch(
+ instances
+ )
+ for result in lf_results:
+ # Record as LLM prediction with labeling_function source
+ schemas = self.app_config.get('annotation_schemes', [])
+ schema_name = (
+ schemas[0].get('name', 'default') if schemas else 'default'
+ )
+ prediction = LLMPrediction(
+ instance_id=result.instance_id,
+ schema_name=schema_name,
+ predicted_label=result.label,
+ confidence_score=result.vote_agreement,
+ uncertainty_score=1.0 - result.vote_agreement,
+ prompt_version=self.current_prompt_version,
+ model_name='labeling_function',
+ reasoning=f"Labeled by {len(result.votes)} labeling functions",
+ )
+ self.set_llm_prediction(
+ result.instance_id, schema_name, prediction
+ )
+ labeled += 1
+
+ # Label remaining with LLM
+ router = self.confidence_router
+ if router is not None:
+ for inst in remaining:
+ result = router.route_instance(
+ inst['instance_id'], inst['text'], inst['schema_name']
+ )
+ if result.accepted and result.labeling_result:
+ self._handle_labeling_result(result.labeling_result)
+ labeled += 1
+ else:
+ for inst in remaining:
+ result = self.llm_labeling_thread._label_instance(
+ inst['instance_id'], inst['text'], inst['schema_name']
+ )
+ if result and not result.error:
+ self._handle_labeling_result(result)
+ labeled += 1
+ return labeled
+
+ def _get_instances_for_labeling(self, batch_size: int) -> List[Dict[str, Any]]:
+ """Get unlabeled instances for background labeling.
+
+ Returns:
+ List of dicts with instance_id, text, and schema_name.
+ """
+ try:
+ from potato.item_state_management import get_item_state_manager
+ ism = get_item_state_manager()
+ except Exception:
+ return []
+
+ schemes = self.app_config.get('annotation_schemes', [])
+ schema_name = schemes[0].get('name', 'default') if schemes else 'default'
+
+ # Collect candidate IDs under the lock, then fetch texts outside it
+ # to avoid blocking the main thread during potentially slow text lookups.
+ # Note: we do NOT filter out human_labeled_ids โ the LLM should label
+ # instances the human has already annotated so retroactive comparison
+ # can update agreement metrics. Only skip instances the LLM already labeled.
+ with self._lock:
+ candidate_ids = [
+ instance_id for instance_id in ism.instance_id_ordering
+ if instance_id not in self.llm_labeled_ids
+ ]
+
+ instances = []
+ for instance_id in candidate_ids:
+ text = self._get_instance_text(instance_id)
+ if text:
+ instances.append({
+ 'instance_id': instance_id,
+ 'text': text,
+ 'schema_name': schema_name,
+ })
+ if len(instances) >= batch_size:
+ break
+ return instances
+
+ # === Validation ===
+
+ def select_validation_sample(self, sample_size: int) -> List[str]:
+ """
+ Select a random sample of LLM-labeled instances for validation.
+
+ Args:
+ sample_size: Number of instances to select
+
+ Returns:
+ List of instance IDs for validation
+ """
+ import random
+
+ with self._lock:
+ # Get instances labeled only by LLM (not by human)
+ llm_only = self.llm_labeled_ids - self.human_labeled_ids
+ llm_only = llm_only - self.validation_sample_ids # Exclude already validated
+
+ available = list(llm_only)
+ sample_size = min(sample_size, len(available))
+
+ sample = random.sample(available, sample_size)
+ self.validation_sample_ids.update(sample)
+
+ logger.info(f"Selected {len(sample)} instances for validation")
+ return sample
+
+ # === State Persistence ===
+
+ def _save_state(self) -> None:
+ """Save manager state to disk.
+
+ Thread-safe: acquires self._lock (RLock) so callers that already
+ hold the lock won't deadlock, while callers from background threads
+ (e.g., labeling loop, rule clustering) are properly serialized.
+ """
+ if not self.config.state_dir:
+ return
+
+ with self._lock:
+ try:
+ os.makedirs(self.config.state_dir, exist_ok=True)
+ filepath = os.path.join(self.config.state_dir, self._state_file)
+
+ state = {
+ 'task_description': self.task_description,
+ 'current_prompt_version': self.current_prompt_version,
+ 'prompt_versions': [p.to_dict() for p in self.prompt_versions],
+ 'predictions': {
+ iid: {s: p.to_dict() for s, p in schemas.items()}
+ for iid, schemas in self.predictions.items()
+ },
+ 'human_labeled_ids': list(self.human_labeled_ids),
+ 'llm_labeled_ids': list(self.llm_labeled_ids),
+ 'disagreement_ids': list(self.disagreement_ids),
+ 'validation_sample_ids': list(self.validation_sample_ids),
+ 'edge_case_ids': list(self.edge_case_ids),
+ 'edge_case_labels': self.edge_case_labels,
+ 'agreement_metrics': self.agreement_metrics.to_dict(),
+ 'confidence_history': {
+ iid: entries
+ for iid, entries in self.confidence_history.items()
+ },
+ 'reannotation_counts': self._reannotation_counts,
+ 'per_version_agreement': self._per_version_agreement,
+ 'refinement_consecutive_failures': self._refinement_consecutive_failures,
+ 'pending_refinements': self._pending_refinements,
+ 'refinement_log': self._refinement_log[-50:], # Keep last 50
+ 'icl_library': self._icl_library.to_dict() if self._icl_library else None,
+ }
+
+ # Include edge case rule manager state inline
+ if self._edge_case_rule_manager is not None:
+ state['edge_case_rule_data'] = self._edge_case_rule_manager.to_dict()
+
+ # Persist ValidationTracker so confusion matrix and comparison
+ # history survive restarts. Without this, /api/confusion-analysis,
+ # /api/disagreement-explorer, and the dashboard's confusion tab
+ # all reset to empty on every server restart.
+ if self._validation_tracker is not None:
+ state['validation_tracker'] = self._validation_tracker.to_dict()
+
+ # Include confidence routing stats (informational only)
+ if self._confidence_router is not None:
+ state['confidence_routing_stats'] = self._confidence_router.get_stats()
+
+ # Atomic write
+ temp_path = filepath + '.tmp'
+ with open(temp_path, 'w') as f:
+ json.dump(state, f, indent=2)
+ os.replace(temp_path, filepath)
+
+ except Exception as e:
+ logger.error(f"Error saving Solo Mode state: {e}")
+
+ def load_state(self) -> bool:
+ """
+ Load manager state from disk.
+
+ Returns:
+ True if state was loaded
+ """
+ if not self.config.state_dir:
+ return False
+
+ filepath = os.path.join(self.config.state_dir, self._state_file)
+
+ if not os.path.exists(filepath):
+ return False
+
+ try:
+ with open(filepath, 'r') as f:
+ state = json.load(f)
+
+ with self._lock:
+ self.task_description = state.get('task_description', '')
+ self.current_prompt_version = state.get('current_prompt_version', 0)
+
+ self.prompt_versions = [
+ PromptVersion.from_dict(p)
+ for p in state.get('prompt_versions', [])
+ ]
+
+ self.predictions = {
+ iid: {
+ s: LLMPrediction.from_dict(p)
+ for s, p in schemas.items()
+ }
+ for iid, schemas in state.get('predictions', {}).items()
+ }
+
+ self.human_labeled_ids = set(state.get('human_labeled_ids', []))
+ self.llm_labeled_ids = set(state.get('llm_labeled_ids', []))
+ self.disagreement_ids = set(state.get('disagreement_ids', []))
+ self.validation_sample_ids = set(state.get('validation_sample_ids', []))
+ self.edge_case_ids = set(state.get('edge_case_ids', []))
+ self.edge_case_labels = state.get('edge_case_labels', {})
+
+ # Restore cartography confidence history
+ raw_history = state.get('confidence_history', {})
+ self.confidence_history = {
+ iid: [(entry[0], entry[1]) for entry in entries]
+ for iid, entries in raw_history.items()
+ }
+
+ metrics = state.get('agreement_metrics', {})
+ self.agreement_metrics = AgreementMetrics(
+ total_compared=metrics.get('total_compared', 0),
+ agreements=metrics.get('agreements', 0),
+ disagreements=metrics.get('disagreements', 0),
+ agreement_rate=metrics.get('agreement_rate', 0.0),
+ )
+
+ # Restore reannotation counts
+ self._reannotation_counts = state.get('reannotation_counts', {})
+
+ # Restore per-version agreement tracking
+ raw_pva = state.get('per_version_agreement', {})
+ self._per_version_agreement = {
+ int(k): v for k, v in raw_pva.items()
+ }
+
+ # Restore validated refinement state
+ self._refinement_consecutive_failures = state.get(
+ 'refinement_consecutive_failures', 0
+ )
+ self._pending_refinements = state.get('pending_refinements', [])
+ self._refinement_log = state.get('refinement_log', [])
+ icl_data = state.get('icl_library')
+ if icl_data:
+ from .refinement.icl_library import ICLLibrary
+ self._icl_library = ICLLibrary.from_dict(icl_data)
+
+ # Load edge case rule manager state
+ ecr_data = state.get('edge_case_rule_data')
+ if ecr_data:
+ from .edge_case_rules import EdgeCaseRuleManager
+ self._edge_case_rule_manager = EdgeCaseRuleManager.from_dict(
+ ecr_data, state_dir=self.config.state_dir
+ )
+
+ # Restore ValidationTracker (confusion matrix + comparison history)
+ vt_data = state.get('validation_tracker')
+ if vt_data:
+ self.validation_tracker.from_dict(vt_data)
+
+ # Load phase state
+ self.phase_controller.load_state()
+
+ logger.info("Loaded Solo Mode state")
+
+ # Auto-start background labeling if already in an annotation phase
+ current_phase = self.phase_controller.get_current_phase()
+ if current_phase in (SoloPhase.PARALLEL_ANNOTATION, SoloPhase.ACTIVE_ANNOTATION):
+ self.start_background_labeling()
+
+ return True
+
+ except Exception as e:
+ logger.error(f"Error loading Solo Mode state: {e}")
+ return False
+
+ # === Route Helper Methods ===
+ # These methods provide simplified interfaces for the routes
+
+ def get_current_prompt_text(self) -> str:
+ """Get current prompt text as string (for routes)."""
+ prompt = self.get_current_prompt()
+ return prompt.prompt_text if prompt else ""
+
+ def get_llm_prediction_for_instance(self, instance_id: str) -> Optional[Dict[str, Any]]:
+ """Get LLM prediction as dict for an instance (for routes)."""
+ with self._lock:
+ if instance_id not in self.predictions:
+ return None
+ # Return first schema's prediction
+ for schema_name, pred in self.predictions[instance_id].items():
+ return {
+ 'label': pred.predicted_label,
+ 'confidence': pred.confidence_score,
+ 'reasoning': pred.reasoning,
+ 'schema': schema_name,
+ }
+ return None
+
+ def get_annotation_stats(self) -> Dict[str, Any]:
+ """Get annotation statistics for the status display."""
+ with self._lock:
+ total = self._get_total_instance_count()
+ return {
+ 'human_labeled': len(self.human_labeled_ids),
+ 'llm_labeled': len(self.llm_labeled_ids),
+ 'remaining': total - len(self.human_labeled_ids | self.llm_labeled_ids),
+ 'total': total,
+ 'agreement_rate': self.agreement_metrics.agreement_rate,
+ }
+
+ def _get_total_instance_count(self) -> int:
+ """Get total number of instances."""
+ try:
+ from potato.item_state_management import get_item_state_manager
+ ism = get_item_state_manager()
+ return len(ism.instance_id_ordering)
+ except Exception:
+ return 0
+
+ def get_available_labels(self) -> List[str]:
+ """Get available labels from annotation schemes."""
+ labels = []
+ schemes = self.app_config.get('annotation_schemes', [])
+ for scheme in schemes:
+ scheme_labels = scheme.get('labels', [])
+ for label in scheme_labels:
+ if isinstance(label, str):
+ labels.append(label)
+ elif isinstance(label, dict):
+ labels.append(label.get('name', str(label)))
+ return labels
+
+ def check_for_disagreement(self, instance_id: str, human_label: Any) -> bool:
+ """Check if there's a disagreement between human and LLM."""
+ with self._lock:
+ if instance_id not in self.predictions:
+ return False
+ for schema_name, pred in self.predictions[instance_id].items():
+ if pred.agrees_with_human is False and not pred.disagreement_resolved:
+ return True
+ return False
+
+ def get_disagreement(self, instance_id: str) -> Optional[Dict[str, Any]]:
+ """Get disagreement details for an instance."""
+ with self._lock:
+ if instance_id not in self.predictions:
+ return None
+ for schema_name, pred in self.predictions[instance_id].items():
+ if pred.agrees_with_human is False and not pred.disagreement_resolved:
+ return {
+ 'id': f"{instance_id}:{schema_name}",
+ 'instance_id': instance_id,
+ 'schema_name': schema_name,
+ 'text': self._get_instance_text(instance_id),
+ 'human_label': pred.human_label,
+ 'llm_label': pred.predicted_label,
+ 'llm_reasoning': pred.reasoning,
+ 'pending_count': len(self.get_pending_disagreements()),
+ }
+ return None
+
+ def _get_instance_text(self, instance_id: str) -> str:
+ """Get text for an instance."""
+ try:
+ from potato.item_state_management import get_item_state_manager
+ ism = get_item_state_manager()
+ item = ism.get_item(instance_id)
+ if item:
+ return item.get_displayed_text()
+ except Exception:
+ pass
+ return ""
+
+ def record_human_annotation(
+ self,
+ instance_id: str,
+ annotation: Any,
+ user_id: str
+ ) -> None:
+ """Record a human annotation (simplified interface for routes)."""
+ # Get first schema name
+ schemes = self.app_config.get('annotation_schemes', [])
+ schema_name = schemes[0].get('name', 'default') if schemes else 'default'
+ self.record_human_label(instance_id, schema_name, annotation, user_id)
+
+ # Check if refinement loop should trigger
+ self._maybe_trigger_refinement()
+
+ def get_llm_labeling_stats(self) -> Dict[str, Any]:
+ """Get LLM labeling statistics."""
+ with self._lock:
+ stats = {
+ 'labeled_count': len(self.llm_labeled_ids),
+ 'queue_size': 0, # Placeholder
+ 'error_count': 0, # Placeholder
+ 'is_paused': self.is_background_labeling_paused(),
+ 'is_running': (
+ self.is_background_labeling_running()
+ and not self.is_background_labeling_paused()
+ ),
+ }
+ stats['confidence_routing'] = (
+ self._confidence_router.get_stats()
+ if self._confidence_router is not None
+ else {'enabled': False}
+ )
+ return stats
+
+ def get_validation_progress(self) -> Dict[str, Any]:
+ """Get validation progress."""
+ with self._lock:
+ total = len(self.validation_sample_ids)
+ # Count validated (those that have been human-labeled from the validation set)
+ validated = len(self.validation_sample_ids & self.human_labeled_ids)
+ return {
+ 'total_samples': total,
+ 'validated': validated,
+ 'remaining': total - validated,
+ 'percent_complete': (validated / total * 100) if total > 0 else 0,
+ 'validation_accuracy': 0.0, # Placeholder
+ 'agreements': 0, # Placeholder
+ }
+
+ def get_validation_samples(self) -> List[Dict[str, Any]]:
+ """Get validation samples that need to be validated."""
+ with self._lock:
+ samples = []
+ for instance_id in self.validation_sample_ids:
+ if instance_id not in self.human_labeled_ids:
+ pred = self.get_llm_prediction_for_instance(instance_id)
+ if pred:
+ samples.append({
+ 'instance_id': instance_id,
+ 'text': self._get_instance_text(instance_id),
+ 'llm_label': pred['label'],
+ 'llm_confidence': pred['confidence'],
+ })
+ return samples
+
+ def record_validation(
+ self,
+ instance_id: str,
+ human_label: Any,
+ notes: str = ""
+ ) -> None:
+ """Record a validation result."""
+ # Get first schema name
+ schemes = self.app_config.get('annotation_schemes', [])
+ schema_name = schemes[0].get('name', 'default') if schemes else 'default'
+ self.record_human_label(instance_id, schema_name, human_label, 'validator')
+
+ def approve_llm_label(self, instance_id: str) -> None:
+ """Approve an LLM label during review."""
+ # Mark as validated/approved
+ with self._lock:
+ self.human_labeled_ids.add(instance_id)
+
+ def correct_llm_label(self, instance_id: str, corrected_label: Any) -> None:
+ """Correct an LLM label during review."""
+ schemes = self.app_config.get('annotation_schemes', [])
+ schema_name = schemes[0].get('name', 'default') if schemes else 'default'
+ self.record_human_label(instance_id, schema_name, corrected_label, 'reviewer')
+
+ def get_instances_for_review(self) -> List[Dict[str, Any]]:
+ """Get low-confidence instances for periodic review."""
+ with self._lock:
+ instances = []
+ low_conf_preds = self.get_low_confidence_predictions()
+ for pred in low_conf_preds[:10]: # Limit to 10
+ if pred.instance_id not in self.human_labeled_ids:
+ instances.append({
+ 'id': pred.instance_id,
+ 'text': self._get_instance_text(pred.instance_id),
+ 'llm_label': pred.predicted_label,
+ 'reasoning': pred.reasoning,
+ 'confidence': pred.confidence_score,
+ })
+ return instances
+
+ def get_all_annotations(self) -> Dict[str, Any]:
+ """Get all annotations for export."""
+ with self._lock:
+ return {
+ 'human_labels': list(self.human_labeled_ids),
+ 'llm_labels': {
+ iid: {s: p.to_dict() for s, p in schemas.items()}
+ for iid, schemas in self.predictions.items()
+ },
+ }
+
+ def get_next_instance_data(self, user_id: str) -> Optional[Dict[str, Any]]:
+ """Get full instance data for the next instance to annotate."""
+ instance_id = self.get_next_instance_for_human(user_id)
+ if not instance_id:
+ return None
+ return {
+ 'id': instance_id,
+ 'text': self._get_instance_text(instance_id),
+ }
+
+ # === Status ===
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get comprehensive status information."""
+ with self._lock:
+ current_prompt = self.get_current_prompt()
+
+ return {
+ 'enabled': self.config.enabled,
+ 'phase': self.phase_controller.get_status(),
+ 'prompt': {
+ 'current_version': self.current_prompt_version,
+ 'total_versions': len(self.prompt_versions),
+ 'current_prompt_length': (
+ len(current_prompt.prompt_text) if current_prompt else 0
+ ),
+ },
+ 'labeling': {
+ 'human_labeled': len(self.human_labeled_ids),
+ 'llm_labeled': len(self.llm_labeled_ids),
+ 'overlap': len(self.human_labeled_ids & self.llm_labeled_ids),
+ 'llm_only': len(self.llm_labeled_ids - self.human_labeled_ids),
+ 'background_running': self.is_background_labeling_running(),
+ },
+ 'agreement': self.agreement_metrics.to_dict(),
+ 'agreement_by_prompt_version': {
+ str(v): {
+ 'compared': d['compared'],
+ 'agreements': d['agreements'],
+ 'rate': d['agreements'] / d['compared'] if d['compared'] > 0 else 0,
+ }
+ for v, d in self._per_version_agreement.items()
+ },
+ 'disagreements': {
+ 'total': len(self.disagreement_ids),
+ 'pending': len(self.get_pending_disagreements()),
+ },
+ 'validation': {
+ 'sample_size': len(self.validation_sample_ids),
+ },
+ 'edge_cases': {
+ 'count': len(self.edge_case_ids),
+ },
+ 'edge_case_rules': (
+ self.edge_case_rule_manager.get_stats()
+ if self._edge_case_rule_manager is not None
+ else {'total_rules': 0, 'total_categories': 0}
+ ),
+ 'confidence_routing': (
+ self._confidence_router.get_stats()
+ if self._confidence_router is not None
+ else {'enabled': False}
+ ),
+ 'thresholds': {
+ 'end_human_annotation_agreement': self.config.thresholds.end_human_annotation_agreement,
+ 'minimum_validation_sample': self.config.thresholds.minimum_validation_sample,
+ 'should_end_human_annotation': self.should_end_human_annotation(),
+ },
+ }
+
+ def shutdown(self) -> None:
+ """Shutdown the manager, stopping background threads."""
+ self.stop_background_labeling()
+ self._save_state()
+ logger.info("SoloModeManager shutdown complete")
+
+
+# === Singleton Management ===
+
+def init_solo_mode_manager(config_data: Dict[str, Any]) -> Optional[SoloModeManager]:
+ """
+ Initialize the singleton SoloModeManager.
+
+ Args:
+ config_data: Full application configuration
+
+ Returns:
+ SoloModeManager instance, or None if disabled
+ """
+ global _SOLO_MODE_MANAGER
+
+ with _SOLO_MODE_LOCK:
+ if _SOLO_MODE_MANAGER is None:
+ solo_config = parse_solo_mode_config(config_data)
+
+ if not solo_config.enabled:
+ logger.info("Solo Mode disabled in config")
+ return None
+
+ # Validate config
+ errors = solo_config.validate()
+ if errors:
+ for error in errors:
+ logger.error(f"Solo Mode config error: {error}")
+ return None
+
+ _SOLO_MODE_MANAGER = SoloModeManager(solo_config, config_data)
+ _SOLO_MODE_MANAGER.load_state()
+
+ return _SOLO_MODE_MANAGER
+
+
+def get_solo_mode_manager() -> Optional[SoloModeManager]:
+ """Get the singleton SoloModeManager instance."""
+ return _SOLO_MODE_MANAGER
+
+
+def clear_solo_mode_manager() -> None:
+ """Clear the singleton (for testing)."""
+ global _SOLO_MODE_MANAGER
+
+ with _SOLO_MODE_LOCK:
+ if _SOLO_MODE_MANAGER is not None:
+ _SOLO_MODE_MANAGER.shutdown()
+ _SOLO_MODE_MANAGER = None
diff --git a/potato/solo_mode/phase_controller.py b/potato/solo_mode/phase_controller.py
new file mode 100644
index 0000000000000000000000000000000000000000..1090ab927983cf6345e1934c3021043d3ecfe1a4
--- /dev/null
+++ b/potato/solo_mode/phase_controller.py
@@ -0,0 +1,409 @@
+"""
+Solo Mode Phase Controller
+
+This module defines the Solo Mode workflow phases and state machine.
+
+Phase State Machine:
+ SETUP โ PROMPT_REVIEW โ EDGE_CASE_SYNTHESIS โ EDGE_CASE_LABELING
+ โ PROMPT_VALIDATION โ PARALLEL_ANNOTATION โท DISAGREEMENT_RESOLUTION
+ โ ACTIVE_ANNOTATION โท PERIODIC_REVIEW โ AUTONOMOUS_LABELING
+ โ FINAL_VALIDATION โ COMPLETED
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum, auto
+from typing import Any, Dict, List, Optional, Set
+import json
+import logging
+import os
+import threading
+
+logger = logging.getLogger(__name__)
+
+
+class SoloPhase(Enum):
+ """
+ Enumeration of Solo Mode workflow phases.
+
+ The phases represent the progression through the human-LLM
+ collaborative annotation workflow.
+ """
+ # Initial setup
+ SETUP = auto() # Task description, data upload
+ PROMPT_REVIEW = auto() # Review/edit synthesized prompt
+
+ # Edge case refinement
+ EDGE_CASE_SYNTHESIS = auto() # LLM generates boundary examples
+ EDGE_CASE_LABELING = auto() # Human labels edge cases
+ PROMPT_VALIDATION = auto() # Verify prompt matches human labels
+
+ # Parallel annotation
+ PARALLEL_ANNOTATION = auto() # Human and LLM annotate in parallel
+ DISAGREEMENT_RESOLUTION = auto() # Resolve human-LLM conflicts
+
+ # Active annotation with periodic review
+ ACTIVE_ANNOTATION = auto() # Main annotation phase
+ PERIODIC_REVIEW = auto() # Review low-confidence LLM labels
+ RULE_REVIEW = auto() # Review discovered edge case rules
+
+ # Autonomous completion
+ AUTONOMOUS_LABELING = auto() # LLM labels remaining data
+ FINAL_VALIDATION = auto() # Human validates LLM-only labels
+ COMPLETED = auto() # Workflow complete
+
+ @classmethod
+ def from_str(cls, s: str) -> 'SoloPhase':
+ """Parse phase from string."""
+ name = s.upper().replace('-', '_')
+ return cls[name]
+
+ def to_str(self) -> str:
+ """Convert phase to string."""
+ return self.name.lower().replace('_', '-')
+
+
+# Phase transition rules: source -> allowed destinations
+PHASE_TRANSITIONS: Dict[SoloPhase, Set[SoloPhase]] = {
+ SoloPhase.SETUP: {SoloPhase.PROMPT_REVIEW},
+ SoloPhase.PROMPT_REVIEW: {SoloPhase.EDGE_CASE_SYNTHESIS, SoloPhase.PARALLEL_ANNOTATION},
+ SoloPhase.EDGE_CASE_SYNTHESIS: {SoloPhase.EDGE_CASE_LABELING},
+ SoloPhase.EDGE_CASE_LABELING: {SoloPhase.PROMPT_VALIDATION, SoloPhase.PROMPT_REVIEW},
+ SoloPhase.PROMPT_VALIDATION: {SoloPhase.PARALLEL_ANNOTATION, SoloPhase.PROMPT_REVIEW},
+ SoloPhase.PARALLEL_ANNOTATION: {
+ SoloPhase.DISAGREEMENT_RESOLUTION,
+ SoloPhase.ACTIVE_ANNOTATION,
+ },
+ SoloPhase.DISAGREEMENT_RESOLUTION: {
+ SoloPhase.PARALLEL_ANNOTATION,
+ SoloPhase.PROMPT_REVIEW, # If major prompt revision needed
+ },
+ SoloPhase.ACTIVE_ANNOTATION: {
+ SoloPhase.PERIODIC_REVIEW,
+ SoloPhase.RULE_REVIEW,
+ SoloPhase.AUTONOMOUS_LABELING,
+ },
+ SoloPhase.PERIODIC_REVIEW: {
+ SoloPhase.ACTIVE_ANNOTATION,
+ SoloPhase.RULE_REVIEW,
+ SoloPhase.PROMPT_REVIEW, # If prompt needs revision
+ },
+ SoloPhase.RULE_REVIEW: {
+ SoloPhase.ACTIVE_ANNOTATION,
+ SoloPhase.PROMPT_REVIEW, # If major changes needed
+ },
+ SoloPhase.AUTONOMOUS_LABELING: {SoloPhase.FINAL_VALIDATION},
+ SoloPhase.FINAL_VALIDATION: {
+ SoloPhase.COMPLETED,
+ SoloPhase.ACTIVE_ANNOTATION, # If validation fails
+ },
+ SoloPhase.COMPLETED: set(), # Terminal state
+}
+
+
+@dataclass
+class PhaseTransition:
+ """Record of a phase transition."""
+ from_phase: SoloPhase
+ to_phase: SoloPhase
+ timestamp: datetime
+ reason: str = ""
+ metadata: Dict[str, Any] = field(default_factory=dict)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'from_phase': self.from_phase.to_str(),
+ 'to_phase': self.to_phase.to_str(),
+ 'timestamp': self.timestamp.isoformat(),
+ 'reason': self.reason,
+ 'metadata': self.metadata,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'PhaseTransition':
+ """Deserialize from dictionary."""
+ return cls(
+ from_phase=SoloPhase.from_str(data['from_phase']),
+ to_phase=SoloPhase.from_str(data['to_phase']),
+ timestamp=datetime.fromisoformat(data['timestamp']),
+ reason=data.get('reason', ''),
+ metadata=data.get('metadata', {}),
+ )
+
+
+@dataclass
+class PhaseState:
+ """State information for a Solo Mode session."""
+ current_phase: SoloPhase = SoloPhase.SETUP
+ transition_history: List[PhaseTransition] = field(default_factory=list)
+ phase_data: Dict[str, Any] = field(default_factory=dict)
+ started_at: Optional[datetime] = None
+ completed_at: Optional[datetime] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'current_phase': self.current_phase.to_str(),
+ 'transition_history': [t.to_dict() for t in self.transition_history],
+ 'phase_data': self.phase_data,
+ 'started_at': self.started_at.isoformat() if self.started_at else None,
+ 'completed_at': self.completed_at.isoformat() if self.completed_at else None,
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> 'PhaseState':
+ """Deserialize from dictionary."""
+ return cls(
+ current_phase=SoloPhase.from_str(data['current_phase']),
+ transition_history=[
+ PhaseTransition.from_dict(t)
+ for t in data.get('transition_history', [])
+ ],
+ phase_data=data.get('phase_data', {}),
+ started_at=(
+ datetime.fromisoformat(data['started_at'])
+ if data.get('started_at') else None
+ ),
+ completed_at=(
+ datetime.fromisoformat(data['completed_at'])
+ if data.get('completed_at') else None
+ ),
+ )
+
+
+class SoloPhaseController:
+ """
+ Controller for Solo Mode phase state machine.
+
+ Manages phase transitions, validates transitions against allowed
+ transitions, and maintains transition history.
+ """
+
+ def __init__(self, state_dir: Optional[str] = None):
+ """
+ Initialize the phase controller.
+
+ Args:
+ state_dir: Directory for persisting state
+ """
+ self._lock = threading.RLock()
+ self.state = PhaseState()
+ self.state_dir = state_dir
+ self._state_file = 'phase_state.json'
+
+ def get_current_phase(self) -> SoloPhase:
+ """Get the current phase."""
+ with self._lock:
+ return self.state.current_phase
+
+ def is_phase(self, phase: SoloPhase) -> bool:
+ """Check if currently in a specific phase."""
+ return self.get_current_phase() == phase
+
+ def is_completed(self) -> bool:
+ """Check if the workflow is completed."""
+ return self.is_phase(SoloPhase.COMPLETED)
+
+ def get_allowed_transitions(self) -> Set[SoloPhase]:
+ """Get phases that can be transitioned to from current phase."""
+ with self._lock:
+ return PHASE_TRANSITIONS.get(self.state.current_phase, set()).copy()
+
+ def can_transition_to(self, target_phase: SoloPhase) -> bool:
+ """Check if transition to target phase is allowed."""
+ with self._lock:
+ allowed = PHASE_TRANSITIONS.get(self.state.current_phase, set())
+ return target_phase in allowed
+
+ def transition_to(
+ self,
+ target_phase: SoloPhase,
+ reason: str = "",
+ metadata: Optional[Dict[str, Any]] = None,
+ force: bool = False
+ ) -> bool:
+ """
+ Transition to a new phase.
+
+ Args:
+ target_phase: The phase to transition to
+ reason: Reason for the transition
+ metadata: Additional metadata for the transition
+ force: If True, allow invalid transitions (for recovery)
+
+ Returns:
+ True if transition was successful
+
+ Raises:
+ ValueError: If transition is not allowed and force=False
+ """
+ with self._lock:
+ current = self.state.current_phase
+
+ if not force and not self.can_transition_to(target_phase):
+ raise ValueError(
+ f"Invalid phase transition: {current.to_str()} -> {target_phase.to_str()}. "
+ f"Allowed transitions: {[p.to_str() for p in self.get_allowed_transitions()]}"
+ )
+
+ # Record transition
+ transition = PhaseTransition(
+ from_phase=current,
+ to_phase=target_phase,
+ timestamp=datetime.now(),
+ reason=reason,
+ metadata=metadata or {},
+ )
+ self.state.transition_history.append(transition)
+
+ # Update state
+ self.state.current_phase = target_phase
+
+ # Update timestamps
+ if current == SoloPhase.SETUP and self.state.started_at is None:
+ self.state.started_at = datetime.now()
+
+ if target_phase == SoloPhase.COMPLETED:
+ self.state.completed_at = datetime.now()
+
+ logger.info(
+ f"Phase transition: {current.to_str()} -> {target_phase.to_str()} "
+ f"(reason: {reason or 'none'})"
+ )
+
+ # Persist state
+ self._save_state()
+
+ return True
+
+ def advance_to_next_phase(self, reason: str = "") -> bool:
+ """
+ Advance to the next logical phase in the workflow.
+
+ For phases with multiple possible transitions, this selects
+ the "primary" next phase (the first in the set).
+
+ Returns:
+ True if advanced, False if no valid transition
+ """
+ with self._lock:
+ allowed = self.get_allowed_transitions()
+ if not allowed:
+ return False
+
+ # Get the primary next phase (lowest enum value)
+ next_phase = min(allowed, key=lambda p: p.value)
+ return self.transition_to(next_phase, reason=reason)
+
+ def get_phase_data(self, key: str, default: Any = None) -> Any:
+ """Get phase-specific data."""
+ with self._lock:
+ return self.state.phase_data.get(key, default)
+
+ def set_phase_data(self, key: str, value: Any) -> None:
+ """Set phase-specific data."""
+ with self._lock:
+ self.state.phase_data[key] = value
+ self._save_state()
+
+ def get_transition_history(self) -> List[PhaseTransition]:
+ """Get the full transition history."""
+ with self._lock:
+ return self.state.transition_history.copy()
+
+ def get_time_in_phase(self) -> Optional[float]:
+ """Get seconds spent in current phase."""
+ with self._lock:
+ if not self.state.transition_history:
+ return None
+
+ # Find last transition to current phase
+ for transition in reversed(self.state.transition_history):
+ if transition.to_phase == self.state.current_phase:
+ return (datetime.now() - transition.timestamp).total_seconds()
+
+ return None
+
+ def get_total_duration(self) -> Optional[float]:
+ """Get total workflow duration in seconds."""
+ with self._lock:
+ if not self.state.started_at:
+ return None
+
+ end = self.state.completed_at or datetime.now()
+ return (end - self.state.started_at).total_seconds()
+
+ def reset(self) -> None:
+ """Reset the controller to initial state."""
+ with self._lock:
+ self.state = PhaseState()
+ self._save_state()
+ logger.info("Phase controller reset to SETUP")
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get current status information."""
+ with self._lock:
+ return {
+ 'current_phase': self.state.current_phase.to_str(),
+ 'allowed_transitions': [p.to_str() for p in self.get_allowed_transitions()],
+ 'transition_count': len(self.state.transition_history),
+ 'time_in_phase_seconds': self.get_time_in_phase(),
+ 'total_duration_seconds': self.get_total_duration(),
+ 'started_at': (
+ self.state.started_at.isoformat()
+ if self.state.started_at else None
+ ),
+ 'completed_at': (
+ self.state.completed_at.isoformat()
+ if self.state.completed_at else None
+ ),
+ 'is_completed': self.is_completed(),
+ }
+
+ def _save_state(self) -> None:
+ """Save state to disk."""
+ if not self.state_dir:
+ return
+
+ try:
+ os.makedirs(self.state_dir, exist_ok=True)
+ filepath = os.path.join(self.state_dir, self._state_file)
+
+ # Atomic write
+ temp_path = filepath + '.tmp'
+ with open(temp_path, 'w') as f:
+ json.dump(self.state.to_dict(), f, indent=2)
+ os.replace(temp_path, filepath)
+
+ except Exception as e:
+ logger.error(f"Error saving phase state: {e}")
+
+ def load_state(self) -> bool:
+ """
+ Load state from disk.
+
+ Returns:
+ True if state was loaded, False if no state file exists
+ """
+ if not self.state_dir:
+ return False
+
+ filepath = os.path.join(self.state_dir, self._state_file)
+
+ if not os.path.exists(filepath):
+ return False
+
+ try:
+ with open(filepath, 'r') as f:
+ data = json.load(f)
+
+ with self._lock:
+ self.state = PhaseState.from_dict(data)
+
+ logger.info(f"Loaded phase state: {self.state.current_phase.to_str()}")
+ return True
+
+ except Exception as e:
+ logger.error(f"Error loading phase state: {e}")
+ return False
diff --git a/potato/solo_mode/prompt_manager.py b/potato/solo_mode/prompt_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..d82effd98aa1c7889d5aa0b74f375abfd7203e6d
--- /dev/null
+++ b/potato/solo_mode/prompt_manager.py
@@ -0,0 +1,574 @@
+"""
+Prompt Manager for Solo Mode
+
+This module handles prompt synthesis, versioning, and revision for Solo Mode.
+It generates annotation prompts from task descriptions and refines them
+based on edge cases and human feedback.
+"""
+
+import json
+import logging
+import os
+import re
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+import threading
+
+logger = logging.getLogger(__name__)
+
+
+PROMPT_SYNTHESIS_TEMPLATE = """You are an expert at creating annotation guidelines and prompts.
+
+Given a task description and annotation schema, synthesize a clear, actionable prompt
+that an LLM can use to label text instances accurately.
+
+## Task Description
+{task_description}
+
+## Annotation Schema
+{schema_info}
+
+## Available Labels
+{labels}
+
+## Requirements for the Prompt
+1. Be clear and unambiguous about what each label means
+2. Include decision criteria for edge cases
+3. Provide examples if helpful (but keep it concise)
+4. Specify what to do when uncertain
+5. Format should be easy for an LLM to follow
+
+## Output Format
+Respond with JSON:
+{{
+ "prompt": "",
+ "explanation": ""
+}}
+"""
+
+
+PROMPT_REVISION_TEMPLATE = """You are an expert at refining annotation guidelines.
+
+The current annotation prompt is not working well for certain cases.
+Based on the feedback, revise the prompt to handle these cases correctly.
+
+## Current Prompt
+{current_prompt}
+
+## Cases Where the Prompt Failed
+{failed_cases}
+
+## Expected vs. Actual Labels
+{label_discrepancies}
+
+## Requirements
+1. Modify the prompt to handle these edge cases
+2. Do not change things that are working well
+3. Add specific guidance for the problematic patterns
+4. Keep the prompt concise and clear
+
+## Output Format
+Respond with JSON:
+{{
+ "prompt": "",
+ "changes_made": ["", "", ...],
+ "explanation": ""
+}}
+"""
+
+
+@dataclass
+class PromptRevision:
+ """Record of a prompt revision."""
+ from_version: int
+ to_version: int
+ changes_made: List[str]
+ trigger: str # 'edge_case', 'disagreement', 'optimization', 'manual'
+ failed_cases: List[Dict[str, Any]] = field(default_factory=list)
+ timestamp: datetime = field(default_factory=datetime.now)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'from_version': self.from_version,
+ 'to_version': self.to_version,
+ 'changes_made': self.changes_made,
+ 'trigger': self.trigger,
+ 'failed_cases': self.failed_cases,
+ 'timestamp': self.timestamp.isoformat(),
+ }
+
+
+class PromptManager:
+ """
+ Manager for annotation prompts in Solo Mode.
+
+ Responsibilities:
+ - Synthesize initial prompts from task descriptions
+ - Version and track prompt history
+ - Revise prompts based on failures and edge cases
+ - Store and retrieve prompts for labeling
+ """
+
+ def __init__(self, config: Dict[str, Any], solo_config: Any):
+ """
+ Initialize the prompt manager.
+
+ Args:
+ config: Full application configuration
+ solo_config: SoloModeConfig instance
+ """
+ self.config = config
+ self.solo_config = solo_config
+ self._lock = threading.RLock()
+
+ # Prompt state
+ self.task_description: str = ""
+ self.schema_info: Dict[str, Any] = {}
+ self.prompts: List[Dict[str, Any]] = [] # Versioned prompts
+ self.current_version: int = 0
+ self.revisions: List[PromptRevision] = []
+
+ # State directory
+ self.state_dir = solo_config.state_dir
+ self._prompt_file = 'prompts.json'
+
+ # AI endpoint for synthesis/revision (lazy init)
+ self._revision_endpoint = None
+
+ def _get_revision_endpoint(self) -> Optional[Any]:
+ """Get or create the AI endpoint for prompt revision."""
+ if self._revision_endpoint is not None:
+ return self._revision_endpoint
+
+ if not self.solo_config.revision_models:
+ logger.warning("No revision models configured")
+ return None
+
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ # Try revision models in order until one works
+ for model_config in self.solo_config.revision_models:
+ try:
+ endpoint_config = model_config.to_endpoint_config()
+
+ endpoint = AIEndpointFactory.create_endpoint(endpoint_config)
+ if endpoint:
+ self._revision_endpoint = endpoint
+ logger.info(f"Using revision endpoint: {model_config.endpoint_type}/{model_config.model}")
+ return endpoint
+ except Exception as e:
+ logger.debug(f"Failed to create revision endpoint {model_config.model}: {e}")
+ continue
+
+ except Exception as e:
+ logger.error(f"Error creating revision endpoint: {e}")
+
+ return None
+
+ def set_task_description(self, description: str) -> None:
+ """Set the task description for prompt synthesis."""
+ with self._lock:
+ self.task_description = description
+ self._save_state()
+
+ def get_task_description(self) -> str:
+ """Get the current task description."""
+ with self._lock:
+ return self.task_description
+
+ def set_schema_info(self, schema_info: Dict[str, Any]) -> None:
+ """Set the annotation schema information."""
+ with self._lock:
+ self.schema_info = schema_info
+ self._save_state()
+
+ def synthesize_prompt(self, task_description: str) -> Optional[str]:
+ """
+ Synthesize an initial annotation prompt from a task description.
+
+ Args:
+ task_description: The user's description of the annotation task
+
+ Returns:
+ The synthesized prompt text, or None if synthesis failed
+ """
+ self.set_task_description(task_description)
+
+ endpoint = self._get_revision_endpoint()
+ if endpoint is None:
+ logger.warning("No endpoint available for prompt synthesis")
+ return self._create_fallback_prompt()
+
+ try:
+ # Get schema info from config
+ schemes = self.config.get('annotation_schemes', [])
+ if not schemes:
+ logger.warning("No annotation schemes configured")
+ return self._create_fallback_prompt()
+
+ schema_info = self._format_schema_info(schemes)
+ labels = self._extract_labels(schemes)
+
+ # Build synthesis prompt
+ synthesis_prompt = PROMPT_SYNTHESIS_TEMPLATE.format(
+ task_description=task_description,
+ schema_info=schema_info,
+ labels=labels,
+ )
+
+ # Query endpoint
+ from pydantic import BaseModel
+
+ class SynthesisResponse(BaseModel):
+ prompt: str
+ explanation: str = ""
+
+ response = endpoint.query(synthesis_prompt, SynthesisResponse)
+
+ # Parse response
+ if isinstance(response, str):
+ response_data = self._parse_json_response(response)
+ elif hasattr(response, 'model_dump'):
+ response_data = response.model_dump()
+ else:
+ response_data = response
+
+ prompt_text = response_data.get('prompt', '')
+ explanation = response_data.get('explanation', '')
+
+ if prompt_text:
+ # Store as first version
+ self._add_prompt_version(
+ prompt_text=prompt_text,
+ created_by='llm_synthesis',
+ source_description=f"Synthesized from task description. {explanation}"
+ )
+ return prompt_text
+
+ except Exception as e:
+ logger.error(f"Error synthesizing prompt: {e}")
+
+ return self._create_fallback_prompt()
+
+ def _create_fallback_prompt(self) -> str:
+ """Create a basic fallback prompt when synthesis fails."""
+ schemes = self.config.get('annotation_schemes', [])
+ labels = self._extract_labels(schemes)
+
+ prompt = f"""Task: {self.task_description}
+
+Please read the text carefully and assign the most appropriate label.
+
+Available labels: {labels}
+
+Respond with just the label name that best fits the text.
+"""
+ self._add_prompt_version(
+ prompt_text=prompt,
+ created_by='fallback',
+ source_description="Basic fallback prompt"
+ )
+ return prompt
+
+ def _format_schema_info(self, schemes: List[Dict[str, Any]]) -> str:
+ """Format annotation schemes for the synthesis prompt."""
+ info_parts = []
+ for scheme in schemes:
+ name = scheme.get('name', 'unknown')
+ ann_type = scheme.get('annotation_type', 'unknown')
+ description = scheme.get('description', '')
+ info_parts.append(f"- {name}: {ann_type} ({description})")
+ return '\n'.join(info_parts)
+
+ def _extract_labels(self, schemes: List[Dict[str, Any]]) -> str:
+ """Extract label names from annotation schemes."""
+ all_labels = []
+ for scheme in schemes:
+ labels = scheme.get('labels', [])
+ for label in labels:
+ if isinstance(label, str):
+ all_labels.append(label)
+ elif isinstance(label, dict):
+ all_labels.append(label.get('name', str(label)))
+ return ', '.join(all_labels)
+
+ def _add_prompt_version(
+ self,
+ prompt_text: str,
+ created_by: str,
+ source_description: str = ""
+ ) -> int:
+ """
+ Add a new prompt version.
+
+ Returns:
+ The new version number
+ """
+ with self._lock:
+ new_version = len(self.prompts) + 1
+ parent = self.current_version if self.current_version > 0 else None
+
+ prompt_data = {
+ 'version': new_version,
+ 'prompt_text': prompt_text,
+ 'created_at': datetime.now().isoformat(),
+ 'created_by': created_by,
+ 'source_description': source_description,
+ 'parent_version': parent,
+ 'validation_accuracy': None,
+ }
+
+ self.prompts.append(prompt_data)
+ self.current_version = new_version
+ self._save_state()
+
+ logger.info(f"Added prompt version {new_version} by {created_by}")
+ return new_version
+
+ def get_current_prompt(self) -> Optional[str]:
+ """Get the current prompt text."""
+ with self._lock:
+ if not self.prompts or self.current_version == 0:
+ return None
+ return self.prompts[self.current_version - 1]['prompt_text']
+
+ def get_prompt_version(self, version: int) -> Optional[Dict[str, Any]]:
+ """Get a specific prompt version."""
+ with self._lock:
+ if 0 < version <= len(self.prompts):
+ return self.prompts[version - 1].copy()
+ return None
+
+ def get_all_versions(self) -> List[Dict[str, Any]]:
+ """Get all prompt versions."""
+ with self._lock:
+ return [p.copy() for p in self.prompts]
+
+ def update_prompt(self, prompt_text: str, created_by: str = 'user') -> int:
+ """
+ Update the prompt by creating a new version.
+
+ Args:
+ prompt_text: The new prompt text
+ created_by: Who created this version
+
+ Returns:
+ The new version number
+ """
+ return self._add_prompt_version(
+ prompt_text=prompt_text,
+ created_by=created_by,
+ source_description="Manual update"
+ )
+
+ def revise_prompt(
+ self,
+ failed_cases: List[Dict[str, Any]],
+ trigger: str = 'edge_case'
+ ) -> Optional[str]:
+ """
+ Revise the prompt based on failed cases.
+
+ Args:
+ failed_cases: List of cases where the prompt produced wrong labels
+ trigger: What triggered the revision
+
+ Returns:
+ The revised prompt text, or None if revision failed
+ """
+ current = self.get_current_prompt()
+ if not current:
+ logger.warning("No current prompt to revise")
+ return None
+
+ endpoint = self._get_revision_endpoint()
+ if endpoint is None:
+ logger.warning("No endpoint available for prompt revision")
+ return None
+
+ try:
+ # Format failed cases
+ cases_text = self._format_failed_cases(failed_cases)
+ discrepancies = self._format_discrepancies(failed_cases)
+
+ revision_prompt = PROMPT_REVISION_TEMPLATE.format(
+ current_prompt=current,
+ failed_cases=cases_text,
+ label_discrepancies=discrepancies,
+ )
+
+ from pydantic import BaseModel
+
+ class RevisionResponse(BaseModel):
+ prompt: str
+ changes_made: List[str] = []
+ explanation: str = ""
+
+ response = endpoint.query(revision_prompt, RevisionResponse)
+
+ # Parse response
+ if isinstance(response, str):
+ response_data = self._parse_json_response(response)
+ elif hasattr(response, 'model_dump'):
+ response_data = response.model_dump()
+ else:
+ response_data = response
+
+ new_prompt = response_data.get('prompt', '')
+ changes = response_data.get('changes_made', [])
+ explanation = response_data.get('explanation', '')
+
+ if new_prompt and new_prompt != current:
+ # Record revision
+ from_version = self.current_version
+ new_version = self._add_prompt_version(
+ prompt_text=new_prompt,
+ created_by='llm_revision',
+ source_description=f"Revision triggered by {trigger}. {explanation}"
+ )
+
+ revision = PromptRevision(
+ from_version=from_version,
+ to_version=new_version,
+ changes_made=changes,
+ trigger=trigger,
+ failed_cases=failed_cases,
+ )
+ self.revisions.append(revision)
+ self._save_state()
+
+ logger.info(f"Revised prompt: {len(changes)} changes made")
+ return new_prompt
+
+ except Exception as e:
+ logger.error(f"Error revising prompt: {e}")
+
+ return None
+
+ def _format_failed_cases(self, cases: List[Dict[str, Any]]) -> str:
+ """Format failed cases for the revision prompt."""
+ formatted = []
+ for i, case in enumerate(cases[:10]): # Limit to 10 cases
+ text = case.get('text', '')[:200] # Truncate long text
+ expected = case.get('expected_label', 'unknown')
+ actual = case.get('actual_label', 'unknown')
+ formatted.append(f"{i+1}. Text: \"{text}\"\n Expected: {expected}, Got: {actual}")
+ return '\n\n'.join(formatted)
+
+ def _format_discrepancies(self, cases: List[Dict[str, Any]]) -> str:
+ """Summarize label discrepancies."""
+ from collections import Counter
+ discrepancies = Counter()
+ for case in cases:
+ expected = case.get('expected_label', 'unknown')
+ actual = case.get('actual_label', 'unknown')
+ if expected != actual:
+ discrepancies[(actual, expected)] += 1
+
+ formatted = []
+ for (actual, expected), count in discrepancies.most_common(5):
+ formatted.append(f"- '{actual}' was predicted but should be '{expected}' ({count} times)")
+ return '\n'.join(formatted)
+
+ def _parse_json_response(self, response: str) -> Dict[str, Any]:
+ """Parse JSON from response, handling markdown code blocks."""
+ content = response.strip()
+
+ if '```json' in content:
+ match = re.search(r'```json\s*([\s\S]*?)\s*```', content)
+ if match:
+ content = match.group(1).strip()
+ elif '```' in content:
+ match = re.search(r'```\s*([\s\S]*?)\s*```', content)
+ if match:
+ content = match.group(1).strip()
+
+ try:
+ return json.loads(content)
+ except json.JSONDecodeError:
+ return {'prompt': content}
+
+ def set_validation_accuracy(self, version: int, accuracy: float) -> None:
+ """Set the validation accuracy for a prompt version."""
+ with self._lock:
+ if 0 < version <= len(self.prompts):
+ self.prompts[version - 1]['validation_accuracy'] = accuracy
+ self._save_state()
+
+ def _save_state(self) -> None:
+ """Save state to disk."""
+ if not self.state_dir:
+ return
+
+ try:
+ os.makedirs(self.state_dir, exist_ok=True)
+ filepath = os.path.join(self.state_dir, self._prompt_file)
+
+ state = {
+ 'task_description': self.task_description,
+ 'schema_info': self.schema_info,
+ 'prompts': self.prompts,
+ 'current_version': self.current_version,
+ 'revisions': [r.to_dict() for r in self.revisions],
+ }
+
+ temp_path = filepath + '.tmp'
+ with open(temp_path, 'w') as f:
+ json.dump(state, f, indent=2)
+ os.replace(temp_path, filepath)
+
+ except Exception as e:
+ logger.error(f"Error saving prompt state: {e}")
+
+ def load_state(self) -> bool:
+ """Load state from disk."""
+ if not self.state_dir:
+ return False
+
+ filepath = os.path.join(self.state_dir, self._prompt_file)
+
+ if not os.path.exists(filepath):
+ return False
+
+ try:
+ with open(filepath, 'r') as f:
+ state = json.load(f)
+
+ with self._lock:
+ self.task_description = state.get('task_description', '')
+ self.schema_info = state.get('schema_info', {})
+ self.prompts = state.get('prompts', [])
+ self.current_version = state.get('current_version', 0)
+
+ self.revisions = []
+ for r in state.get('revisions', []):
+ self.revisions.append(PromptRevision(
+ from_version=r['from_version'],
+ to_version=r['to_version'],
+ changes_made=r['changes_made'],
+ trigger=r['trigger'],
+ failed_cases=r.get('failed_cases', []),
+ timestamp=datetime.fromisoformat(r['timestamp']),
+ ))
+
+ logger.info(f"Loaded prompt state: {len(self.prompts)} versions")
+ return True
+
+ except Exception as e:
+ logger.error(f"Error loading prompt state: {e}")
+ return False
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get prompt manager status."""
+ with self._lock:
+ current = self.get_current_prompt()
+ return {
+ 'has_task_description': bool(self.task_description),
+ 'current_version': self.current_version,
+ 'total_versions': len(self.prompts),
+ 'current_prompt_length': len(current) if current else 0,
+ 'revision_count': len(self.revisions),
+ }
diff --git a/potato/solo_mode/prompt_optimizer.py b/potato/solo_mode/prompt_optimizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..91f83639e3c0e3cf678659a4dff5c8d6cf4c7444
--- /dev/null
+++ b/potato/solo_mode/prompt_optimizer.py
@@ -0,0 +1,551 @@
+"""
+Prompt Optimizer for Solo Mode
+
+This module implements DSPy-style automatic prompt optimization.
+It uses labeled examples to iteratively improve prompts for better
+accuracy while maintaining brevity.
+"""
+
+import json
+import logging
+import re
+import threading
+import time
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Tuple
+from queue import Queue
+
+logger = logging.getLogger(__name__)
+
+
+OPTIMIZATION_PROMPT_TEMPLATE = """You are an expert at improving annotation prompts.
+
+Given the current prompt and some examples where the LLM made mistakes,
+suggest improvements to the prompt that would help get the correct labels.
+
+## Current Prompt
+{current_prompt}
+
+## Correct Examples
+These are examples where the LLM got the right answer:
+{correct_examples}
+
+## Incorrect Examples
+These are examples where the LLM got the wrong answer (with corrections):
+{incorrect_examples}
+
+## Optimization Goals
+1. Improve accuracy on the incorrect examples
+2. Keep the prompt concise (shorter is better)
+3. Make instructions clearer and more specific
+4. Add clarifying examples if helpful
+
+## Requirements
+- Focus on patterns in the errors
+- Don't make the prompt too long
+- Keep successful patterns from the current prompt
+- Be specific about edge cases
+
+## Output Format
+Respond with JSON:
+{{
+ "improved_prompt": "",
+ "changes_made": ["", "", ...],
+ "rationale": ""
+}}
+"""
+
+
+@dataclass
+class OptimizationResult:
+ """Result of a prompt optimization run."""
+ original_prompt: str
+ optimized_prompt: str
+ changes_made: List[str]
+ rationale: str
+ accuracy_before: float
+ accuracy_after: Optional[float] = None
+ timestamp: datetime = field(default_factory=datetime.now)
+ model_used: str = ""
+ num_examples_used: int = 0
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'original_prompt': self.original_prompt,
+ 'optimized_prompt': self.optimized_prompt,
+ 'changes_made': self.changes_made,
+ 'rationale': self.rationale,
+ 'accuracy_before': self.accuracy_before,
+ 'accuracy_after': self.accuracy_after,
+ 'timestamp': self.timestamp.isoformat(),
+ 'model_used': self.model_used,
+ 'num_examples_used': self.num_examples_used,
+ }
+
+
+@dataclass
+class OptimizationConfig:
+ """Configuration for prompt optimization."""
+ enabled: bool = True
+ find_smallest_model: bool = True
+ target_accuracy: float = 0.85
+ min_examples_for_optimization: int = 10
+ optimization_interval_seconds: int = 300 # 5 minutes
+ max_prompt_length: int = 2000
+ accuracy_weight: float = 0.7
+ length_weight: float = 0.2
+ consistency_weight: float = 0.1
+
+
+class PromptOptimizer:
+ """
+ DSPy-style automatic prompt optimization.
+
+ Optimizes prompts based on:
+ 1. Accuracy on labeled examples
+ 2. Prompt length (shorter is better)
+ 3. Prediction consistency
+
+ Can run in background or be triggered on-demand.
+ """
+
+ def __init__(
+ self,
+ config: Dict[str, Any],
+ solo_config: Any,
+ prompt_getter: callable,
+ prompt_setter: callable,
+ examples_getter: callable,
+ ):
+ """
+ Initialize the prompt optimizer.
+
+ Args:
+ config: Full application configuration
+ solo_config: SoloModeConfig instance
+ prompt_getter: Callable that returns current prompt text
+ prompt_setter: Callable to update the prompt
+ examples_getter: Callable that returns labeled examples
+ """
+ self.config = config
+ self.solo_config = solo_config
+ self.prompt_getter = prompt_getter
+ self.prompt_setter = prompt_setter
+ self.examples_getter = examples_getter
+
+ # Load optimization config. solo_config.prompt_optimization may be a
+ # PromptOptimizationConfig dataclass (typical) or a plain dict; read
+ # fields from either form.
+ opt_config = getattr(solo_config, 'prompt_optimization', None)
+ if opt_config:
+ def _opt(key, default):
+ if isinstance(opt_config, dict):
+ return opt_config.get(key, default)
+ return getattr(opt_config, key, default)
+ self.opt_config = OptimizationConfig(
+ enabled=_opt('enabled', True),
+ find_smallest_model=_opt('find_smallest_model', True),
+ target_accuracy=_opt('target_accuracy', 0.85),
+ )
+ else:
+ self.opt_config = OptimizationConfig()
+
+ self._lock = threading.RLock()
+
+ # Optimization history
+ self.optimization_history: List[OptimizationResult] = []
+
+ # Background optimization
+ self._background_thread: Optional[threading.Thread] = None
+ self._stop_event = threading.Event()
+ self._optimization_queue: Queue = Queue()
+
+ # AI endpoint (lazy init)
+ self._endpoint = None
+
+ # Cached labeled examples
+ self._cached_examples: Dict[str, Dict[str, Any]] = {}
+
+ def _get_endpoint(self) -> Optional[Any]:
+ """Get or create the optimization endpoint."""
+ if self._endpoint is not None:
+ return self._endpoint
+
+ if not self.solo_config.revision_models:
+ logger.warning("No revision models configured for optimization")
+ return None
+
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ for model_config in self.solo_config.revision_models:
+ try:
+ endpoint_config = model_config.to_endpoint_config(temperature_override=0.3)
+
+ endpoint = AIEndpointFactory.create_endpoint(endpoint_config)
+ if endpoint:
+ self._endpoint = endpoint
+ return endpoint
+ except Exception as e:
+ logger.debug(f"Failed to create optimization endpoint: {e}")
+ continue
+
+ except Exception as e:
+ logger.error(f"Error creating optimization endpoint: {e}")
+
+ return None
+
+ def optimize(self, force: bool = False) -> Optional[OptimizationResult]:
+ """
+ Run prompt optimization.
+
+ Args:
+ force: Run even if not enough examples
+
+ Returns:
+ OptimizationResult if optimization was performed
+ """
+ with self._lock:
+ # Get labeled examples
+ examples = self.examples_getter()
+ if not examples and not force:
+ logger.info("No labeled examples available for optimization")
+ return None
+
+ if len(examples) < self.opt_config.min_examples_for_optimization and not force:
+ logger.info(
+ f"Not enough examples for optimization "
+ f"({len(examples)} < {self.opt_config.min_examples_for_optimization})"
+ )
+ return None
+
+ # Get current prompt
+ current_prompt = self.prompt_getter()
+ if not current_prompt:
+ logger.warning("No current prompt to optimize")
+ return None
+
+ # Get endpoint
+ endpoint = self._get_endpoint()
+ if endpoint is None:
+ logger.warning("No endpoint available for optimization")
+ return None
+
+ # Split examples into correct and incorrect
+ correct, incorrect = self._split_examples(examples)
+
+ if not incorrect:
+ logger.info("No incorrect predictions to optimize for")
+ return None
+
+ # Calculate current accuracy
+ accuracy_before = len(correct) / len(examples) if examples else 0.0
+
+ # Check if already above target
+ if accuracy_before >= self.opt_config.target_accuracy and not force:
+ logger.info(
+ f"Accuracy ({accuracy_before:.2%}) already above target "
+ f"({self.opt_config.target_accuracy:.2%})"
+ )
+ return None
+
+ # Generate optimized prompt
+ result = self._generate_optimized_prompt(
+ current_prompt,
+ correct[:5], # Limit examples
+ incorrect[:10],
+ endpoint,
+ accuracy_before,
+ )
+
+ if result:
+ self.optimization_history.append(result)
+
+ # Update prompt if optimization was successful
+ if result.optimized_prompt and result.optimized_prompt != current_prompt:
+ self.prompt_setter(
+ result.optimized_prompt,
+ source='llm_optimization',
+ source_description='; '.join(result.changes_made)
+ )
+ logger.info(f"Prompt optimized: {len(result.changes_made)} changes made")
+
+ return result
+
+ def _split_examples(
+ self,
+ examples: List[Dict[str, Any]]
+ ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+ """Split examples into correct and incorrect predictions."""
+ correct = []
+ incorrect = []
+
+ for ex in examples:
+ if ex.get('agrees', True):
+ correct.append(ex)
+ else:
+ incorrect.append(ex)
+
+ return correct, incorrect
+
+ def _generate_optimized_prompt(
+ self,
+ current_prompt: str,
+ correct_examples: List[Dict[str, Any]],
+ incorrect_examples: List[Dict[str, Any]],
+ endpoint: Any,
+ accuracy_before: float,
+ ) -> Optional[OptimizationResult]:
+ """Generate an optimized prompt using the LLM."""
+ try:
+ # Format examples
+ correct_text = self._format_examples(correct_examples, show_correction=False)
+ incorrect_text = self._format_examples(incorrect_examples, show_correction=True)
+
+ optimization_prompt = OPTIMIZATION_PROMPT_TEMPLATE.format(
+ current_prompt=current_prompt,
+ correct_examples=correct_text or "None available",
+ incorrect_examples=incorrect_text or "None available",
+ )
+
+ from pydantic import BaseModel
+
+ class OptimizationResponse(BaseModel):
+ improved_prompt: str = ""
+ changes_made: List[str] = []
+ rationale: str = ""
+
+ response = endpoint.query(optimization_prompt, OptimizationResponse)
+
+ # Parse response
+ if isinstance(response, str):
+ response_data = self._parse_json_response(response)
+ elif hasattr(response, 'model_dump'):
+ response_data = response.model_dump()
+ else:
+ response_data = response
+
+ improved_prompt = response_data.get('improved_prompt', '')
+ changes_made = response_data.get('changes_made', [])
+ rationale = response_data.get('rationale', '')
+
+ # Validate improved prompt
+ if not improved_prompt:
+ logger.warning("Optimization returned empty prompt")
+ return None
+
+ if len(improved_prompt) > self.opt_config.max_prompt_length:
+ logger.warning(
+ f"Optimized prompt too long ({len(improved_prompt)} > {self.opt_config.max_prompt_length})"
+ )
+ # Truncate if necessary
+ improved_prompt = improved_prompt[:self.opt_config.max_prompt_length]
+
+ return OptimizationResult(
+ original_prompt=current_prompt,
+ optimized_prompt=improved_prompt,
+ changes_made=changes_made,
+ rationale=rationale,
+ accuracy_before=accuracy_before,
+ model_used=getattr(endpoint, 'model', ''),
+ num_examples_used=len(correct_examples) + len(incorrect_examples),
+ )
+
+ except Exception as e:
+ logger.error(f"Error generating optimized prompt: {e}")
+ return None
+
+ def _format_examples(
+ self,
+ examples: List[Dict[str, Any]],
+ show_correction: bool = False
+ ) -> str:
+ """Format examples for the optimization prompt."""
+ if not examples:
+ return ""
+
+ formatted = []
+ for i, ex in enumerate(examples[:10], 1):
+ text = ex.get('text', '')[:200] # Truncate long text
+ predicted = ex.get('predicted_label', '')
+ if show_correction:
+ actual = ex.get('actual_label', ex.get('human_label', ''))
+ formatted.append(
+ f"{i}. Text: \"{text}\"\n"
+ f" LLM predicted: {predicted}\n"
+ f" Correct label: {actual}"
+ )
+ else:
+ formatted.append(
+ f"{i}. Text: \"{text}\"\n"
+ f" Label: {predicted}"
+ )
+
+ return '\n\n'.join(formatted)
+
+ def _parse_json_response(self, response: str) -> Dict[str, Any]:
+ """Parse JSON from response."""
+ content = response.strip()
+
+ if '```json' in content:
+ match = re.search(r'```json\s*([\s\S]*?)\s*```', content)
+ if match:
+ content = match.group(1).strip()
+ elif '```' in content:
+ match = re.search(r'```\s*([\s\S]*?)\s*```', content)
+ if match:
+ content = match.group(1).strip()
+
+ try:
+ return json.loads(content)
+ except json.JSONDecodeError:
+ return {}
+
+ # === Background Optimization ===
+
+ def start_background_optimization(self) -> bool:
+ """Start background optimization thread."""
+ with self._lock:
+ if self._background_thread is not None and self._background_thread.is_alive():
+ logger.warning("Background optimization already running")
+ return False
+
+ if not self.opt_config.enabled:
+ logger.info("Prompt optimization is disabled")
+ return False
+
+ self._stop_event.clear()
+ self._background_thread = threading.Thread(
+ target=self._background_optimization_loop,
+ name="PromptOptimizationThread",
+ daemon=True
+ )
+ self._background_thread.start()
+ logger.info("Started background prompt optimization")
+ return True
+
+ def stop_background_optimization(self) -> None:
+ """Stop background optimization thread."""
+ if self._background_thread is None:
+ return
+
+ self._stop_event.set()
+ self._background_thread.join(timeout=5.0)
+ self._background_thread = None
+ logger.info("Stopped background prompt optimization")
+
+ def is_running(self) -> bool:
+ """Check if background optimization is running."""
+ return (
+ self._background_thread is not None and
+ self._background_thread.is_alive()
+ )
+
+ def _background_optimization_loop(self) -> None:
+ """Main loop for background optimization."""
+ interval = self.opt_config.optimization_interval_seconds
+
+ logger.info(f"Background optimization started (interval={interval}s)")
+
+ while not self._stop_event.is_set():
+ try:
+ # Wait for interval
+ if self._stop_event.wait(timeout=interval):
+ break # Stop event was set
+
+ # Run optimization
+ result = self.optimize()
+ if result:
+ logger.info(
+ f"Background optimization completed: "
+ f"accuracy {result.accuracy_before:.2%} -> {result.accuracy_after or 'pending'}"
+ )
+
+ except Exception as e:
+ logger.error(f"Error in background optimization: {e}")
+
+ # === Model Selection ===
+
+ def find_smallest_accurate_model(
+ self,
+ models: List[Any],
+ test_examples: List[Dict[str, Any]],
+ prompt: str,
+ ) -> Optional[str]:
+ """
+ Find the smallest model that achieves target accuracy.
+
+ Args:
+ models: List of model configs (ordered small to large)
+ test_examples: Examples to test accuracy on
+ prompt: The prompt to use
+
+ Returns:
+ Model name if found, None otherwise
+ """
+ if not self.opt_config.find_smallest_model:
+ return None
+
+ target = self.opt_config.target_accuracy
+
+ for model_config in models:
+ try:
+ accuracy = self._test_model_accuracy(
+ model_config, test_examples, prompt
+ )
+ if accuracy >= target:
+ logger.info(
+ f"Model {model_config.model} achieves {accuracy:.2%} accuracy"
+ )
+ return model_config.model
+ except Exception as e:
+ logger.debug(f"Error testing model {model_config.model}: {e}")
+ continue
+
+ logger.warning("No model achieved target accuracy")
+ return None
+
+ def _test_model_accuracy(
+ self,
+ model_config: Any,
+ examples: List[Dict[str, Any]],
+ prompt: str,
+ ) -> float:
+ """Test a model's accuracy on examples."""
+ # This is a placeholder - full implementation would
+ # run predictions with the model and calculate accuracy
+ return 0.0
+
+ # === Status and History ===
+
+ def get_optimization_history(self) -> List[OptimizationResult]:
+ """Get optimization history."""
+ with self._lock:
+ return self.optimization_history.copy()
+
+ def get_last_optimization(self) -> Optional[OptimizationResult]:
+ """Get the most recent optimization result."""
+ with self._lock:
+ if self.optimization_history:
+ return self.optimization_history[-1]
+ return None
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get optimizer status."""
+ with self._lock:
+ last = self.get_last_optimization()
+ return {
+ 'enabled': self.opt_config.enabled,
+ 'is_running': self.is_running(),
+ 'optimization_count': len(self.optimization_history),
+ 'last_optimization': last.to_dict() if last else None,
+ 'target_accuracy': self.opt_config.target_accuracy,
+ 'interval_seconds': self.opt_config.optimization_interval_seconds,
+ }
+
+ def clear_history(self) -> None:
+ """Clear optimization history."""
+ with self._lock:
+ self.optimization_history.clear()
diff --git a/potato/solo_mode/refinement/__init__.py b/potato/solo_mode/refinement/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8ca57244315935443ebfa0b929ea509e7111ef53
--- /dev/null
+++ b/potato/solo_mode/refinement/__init__.py
@@ -0,0 +1,44 @@
+"""
+Refinement Framework
+
+Pluggable strategies for improving annotation prompts based on
+human-LLM disagreements. Every strategy uses a validation-gated
+apply step to prevent regressions.
+
+Available strategies (all have `RefinementStrategy` as base class):
+
+- validated_focused_edit: prompt rule edits with validation gate
+ (recommended for small optimizer models)
+- principle_icl: add validated ICL examples instead of rules
+ (recommended for subjective tasks and small optimizers)
+- hybrid_dual_track: try prompt edit first, fall back to ICL on failure
+ (recommended default)
+- append: legacy append-only refinement, no validation (for ablation)
+
+Config:
+ solo_mode.refinement_loop.strategy: "validated_focused_edit" | "principle_icl" | ...
+ solo_mode.refinement_loop.strategy_config: {...} # strategy-specific overrides
+"""
+
+from .base import (
+ RefinementStrategy,
+ RefinementCandidate,
+ RefinementResult,
+ CandidateKind,
+)
+from .validation import ValidationSplit, CandidateEvaluator
+from .icl_library import ICLLibrary
+from .registry import get_strategy, list_strategies, register_strategy
+
+__all__ = [
+ "RefinementStrategy",
+ "RefinementCandidate",
+ "RefinementResult",
+ "CandidateKind",
+ "ValidationSplit",
+ "CandidateEvaluator",
+ "ICLLibrary",
+ "get_strategy",
+ "list_strategies",
+ "register_strategy",
+]
diff --git a/potato/solo_mode/refinement/base.py b/potato/solo_mode/refinement/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e07c44268e97e7125f2e745b2dd524a13eae796
--- /dev/null
+++ b/potato/solo_mode/refinement/base.py
@@ -0,0 +1,150 @@
+"""
+Base classes for refinement strategies.
+
+A refinement strategy takes (confusion patterns, current prompt, train/val
+disagreements) and proposes candidates (either prompt edits or ICL examples).
+The framework scores each candidate on the validation split and applies only
+those that beat the baseline.
+"""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Optional, Callable
+
+logger = logging.getLogger(__name__)
+
+
+class CandidateKind(Enum):
+ """What kind of change a candidate represents."""
+ PROMPT_EDIT = "prompt_edit" # Replaces the guidelines section
+ ICL_EXAMPLE = "icl_example" # Adds an example to the ICL library
+ PRINCIPLE = "principle" # Adds a principle as text in the prompt
+
+
+@dataclass
+class RefinementCandidate:
+ """A single candidate change proposed by a strategy.
+
+ Each candidate can be evaluated independently on the validation set.
+ """
+ kind: CandidateKind
+ # For PROMPT_EDIT: the complete guidelines text that will replace the section
+ # For ICL_EXAMPLE: a dict with {instance_id, text, label, principle}
+ # For PRINCIPLE: the principle text
+ payload: Any
+ # Source pattern this candidate addresses (for logging)
+ target_pattern: Optional[str] = None
+ # The strategy that proposed it
+ proposed_by: str = ""
+ # Rationale (for audit trail and user review)
+ rationale: str = ""
+
+
+@dataclass
+class RefinementResult:
+ """The outcome of a refinement cycle."""
+ success: bool
+ strategy: str
+ applied_candidate: Optional[RefinementCandidate] = None
+ all_candidates: List[RefinementCandidate] = field(default_factory=list)
+ val_baseline_accuracy: float = 0.0
+ val_candidate_accuracies: Dict[int, float] = field(default_factory=dict) # candidate_index -> acc
+ val_sample_ids: List[str] = field(default_factory=list)
+ train_sample_size: int = 0
+ val_sample_size: int = 0
+ # If dry-run, applied_candidate is None but all_candidates is populated
+ dry_run: bool = False
+ # Reason for no-apply
+ failure_reason: Optional[str] = None
+ created_at: str = field(default_factory=lambda: datetime.now().isoformat())
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ "success": self.success,
+ "strategy": self.strategy,
+ "applied_candidate": self._candidate_to_dict(self.applied_candidate) if self.applied_candidate else None,
+ "all_candidates": [self._candidate_to_dict(c) for c in self.all_candidates],
+ "val_baseline_accuracy": self.val_baseline_accuracy,
+ "val_candidate_accuracies": {str(k): v for k, v in self.val_candidate_accuracies.items()},
+ "val_sample_ids": self.val_sample_ids,
+ "train_sample_size": self.train_sample_size,
+ "val_sample_size": self.val_sample_size,
+ "dry_run": self.dry_run,
+ "failure_reason": self.failure_reason,
+ "created_at": self.created_at,
+ }
+
+ @staticmethod
+ def _candidate_to_dict(c: Optional[RefinementCandidate]) -> Optional[Dict]:
+ if c is None:
+ return None
+ return {
+ "kind": c.kind.value,
+ "payload": c.payload,
+ "target_pattern": c.target_pattern,
+ "proposed_by": c.proposed_by,
+ "rationale": c.rationale,
+ }
+
+
+class RefinementStrategy(ABC):
+ """Abstract base class for refinement strategies.
+
+ Subclasses implement `propose_candidates()`. The framework handles:
+ - splitting disagreements into train/val
+ - evaluating candidates on val set via CandidateEvaluator
+ - applying only candidates that beat the baseline
+ - tracking failure counters and dry-run logging
+
+ Subclasses should set:
+ NAME: str registry key
+ RECOMMENDED_OPTIMIZER_TIER: "small" | "medium" | "large"
+ BEST_FOR: list of tags (["binary", "subjective", "many_labels", ...])
+ DESCRIPTION: one-line description shown to practitioners
+ """
+
+ NAME: str = "abstract"
+ RECOMMENDED_OPTIMIZER_TIER: str = "small"
+ BEST_FOR: List[str] = []
+ DESCRIPTION: str = ""
+
+ def __init__(self, manager: Any, solo_config: Any):
+ """
+ Args:
+ manager: SoloModeManager instance (for accessing predictions, analyzer, etc.)
+ solo_config: SoloModeConfig
+ """
+ self.manager = manager
+ self.solo_config = solo_config
+
+ @abstractmethod
+ def propose_candidates(
+ self,
+ patterns: List[Any],
+ current_prompt: str,
+ train_comparisons: List[Dict[str, Any]],
+ ) -> List[RefinementCandidate]:
+ """Generate candidate refinements based on training disagreements.
+
+ Args:
+ patterns: ConfusionPattern list (already filtered to train split)
+ current_prompt: the current annotation prompt text
+ train_comparisons: list of comparison dicts (human_label, llm_label, etc.)
+ sliced to train split
+
+ Returns:
+ List of RefinementCandidate objects (may be empty)
+ """
+ ...
+
+ def supports_kind(self, kind: CandidateKind) -> bool:
+ """Whether this strategy can produce candidates of a given kind.
+
+ Default: supports all kinds. Override to restrict.
+ """
+ return True
diff --git a/potato/solo_mode/refinement/icl_library.py b/potato/solo_mode/refinement/icl_library.py
new file mode 100644
index 0000000000000000000000000000000000000000..0dc6ee4829fbd957884cba1685ab5dcdf2db1783
--- /dev/null
+++ b/potato/solo_mode/refinement/icl_library.py
@@ -0,0 +1,118 @@
+"""
+Persistent library of validated ICL examples.
+
+Each entry has been shown to improve validation accuracy. Strategies add
+entries; the labeling thread reads them via the existing `examples_getter`
+callback.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field, asdict
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ICLEntry:
+ """A single validated ICL example."""
+ instance_id: str
+ text: str
+ label: str
+ principle: str = "" # optional one-line rationale
+ added_at_cycle: int = 0
+ val_accuracy_gain: float = 0.0 # proven improvement on val set when added
+ created_at: str = field(default_factory=lambda: datetime.now().isoformat())
+
+ def to_dict(self) -> Dict[str, Any]:
+ return asdict(self)
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "ICLEntry":
+ return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})
+
+
+class ICLLibrary:
+ """Manages validated ICL examples.
+
+ The library is per-dataset (keyed by state_dir), keeping validated
+ examples from SST-2 out of GoEmotions' prompts.
+
+ Entries are returned by the `examples_getter` callback used by the
+ LLM labeling thread.
+ """
+
+ def __init__(self, max_size: int = 10):
+ """
+ Args:
+ max_size: maximum number of examples to return via get_examples().
+ The library can store more; get_examples() returns the
+ top-K by val_accuracy_gain.
+ """
+ self.max_size = max_size
+ self._entries: List[ICLEntry] = []
+
+ def add(self, entry: ICLEntry) -> None:
+ """Add a validated entry. Dedupe by instance_id."""
+ existing_ids = {e.instance_id for e in self._entries}
+ if entry.instance_id in existing_ids:
+ logger.debug(f"[ICLLibrary] Skipping duplicate for {entry.instance_id}")
+ return
+ self._entries.append(entry)
+ logger.info(
+ f"[ICLLibrary] Added {entry.instance_id} "
+ f"(label={entry.label}, gain=+{entry.val_accuracy_gain:.3f})"
+ )
+
+ def remove(self, instance_id: str) -> bool:
+ """Remove an entry by instance_id. Returns True if removed."""
+ before = len(self._entries)
+ self._entries = [e for e in self._entries if e.instance_id != instance_id]
+ return len(self._entries) < before
+
+ def get_examples(self, max_per_label: int = 1, max_total: int = 5) -> List[Dict[str, str]]:
+ """Get the current ICL examples for injection into a labeling prompt.
+
+ Returns highest-gain entries, at most max_per_label per label, up to
+ max_total total examples.
+ """
+ # Sort by gain descending
+ sorted_entries = sorted(self._entries, key=lambda e: e.val_accuracy_gain, reverse=True)
+
+ by_label: Dict[str, int] = {}
+ result: List[Dict[str, str]] = []
+ for entry in sorted_entries:
+ count = by_label.get(entry.label, 0)
+ if count >= max_per_label:
+ continue
+ result.append({
+ "text": entry.text[:200],
+ "label": entry.label,
+ "principle": entry.principle,
+ })
+ by_label[entry.label] = count + 1
+ if len(result) >= max_total:
+ break
+
+ return result
+
+ def size(self) -> int:
+ return len(self._entries)
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ "max_size": self.max_size,
+ "entries": [e.to_dict() for e in self._entries],
+ }
+
+ @classmethod
+ def from_dict(cls, data: Dict[str, Any]) -> "ICLLibrary":
+ lib = cls(max_size=data.get("max_size", 10))
+ lib._entries = [ICLEntry.from_dict(e) for e in data.get("entries", [])]
+ return lib
+
+ def list_all(self) -> List[ICLEntry]:
+ return list(self._entries)
diff --git a/potato/solo_mode/refinement/registry.py b/potato/solo_mode/refinement/registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..f247766a9187802d230192911979e6955f8e84ba
--- /dev/null
+++ b/potato/solo_mode/refinement/registry.py
@@ -0,0 +1,71 @@
+"""
+Registry of refinement strategies.
+
+Strategies register themselves via the @register decorator. The manager
+looks up the configured strategy name at refinement time.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Dict, List, Type
+
+from .base import RefinementStrategy
+
+logger = logging.getLogger(__name__)
+
+_STRATEGIES: Dict[str, Type[RefinementStrategy]] = {}
+
+
+def register_strategy(cls: Type[RefinementStrategy]) -> Type[RefinementStrategy]:
+ """Decorator to register a refinement strategy class."""
+ name = getattr(cls, "NAME", None)
+ if not name or name == "abstract":
+ raise ValueError(f"Strategy class {cls.__name__} must set a non-abstract NAME")
+ if name in _STRATEGIES:
+ logger.warning(f"Overwriting registered strategy: {name}")
+ _STRATEGIES[name] = cls
+ logger.debug(f"Registered refinement strategy: {name}")
+ return cls
+
+
+def get_strategy(name: str) -> Type[RefinementStrategy]:
+ """Look up a strategy class by name.
+
+ Raises KeyError with a helpful message if not found.
+ """
+ # Ensure builtin strategies are loaded
+ _load_builtin_strategies()
+ if name not in _STRATEGIES:
+ available = ", ".join(sorted(_STRATEGIES.keys())) or "(none)"
+ raise KeyError(
+ f"Unknown refinement strategy: '{name}'. Available: {available}"
+ )
+ return _STRATEGIES[name]
+
+
+def list_strategies() -> List[Dict[str, str]]:
+ """List all registered strategies with their metadata."""
+ _load_builtin_strategies()
+ return [
+ {
+ "name": cls.NAME,
+ "tier": cls.RECOMMENDED_OPTIMIZER_TIER,
+ "best_for": ", ".join(cls.BEST_FOR),
+ "description": cls.DESCRIPTION,
+ }
+ for cls in _STRATEGIES.values()
+ ]
+
+
+_BUILTINS_LOADED = False
+
+
+def _load_builtin_strategies() -> None:
+ """Import modules that register strategies."""
+ global _BUILTINS_LOADED
+ if _BUILTINS_LOADED:
+ return
+ _BUILTINS_LOADED = True
+ # Triggers @register_strategy side effects
+ from . import strategies # noqa: F401
diff --git a/potato/solo_mode/refinement/strategies.py b/potato/solo_mode/refinement/strategies.py
new file mode 100644
index 0000000000000000000000000000000000000000..2d55288966ab5fc4429e282f3a32851cda140764
--- /dev/null
+++ b/potato/solo_mode/refinement/strategies.py
@@ -0,0 +1,360 @@
+"""
+Built-in refinement strategies.
+
+Each strategy inherits from RefinementStrategy and is registered via
+@register_strategy so the manager can look it up by name.
+
+The framework (manager.trigger_refinement_cycle) handles:
+ - validation split
+ - candidate evaluation via CandidateEvaluator
+ - validation gating (reject candidates below baseline)
+ - failure counter and resume-on-new-disagreements
+
+Strategies only need to implement propose_candidates().
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from .base import (
+ RefinementStrategy,
+ RefinementCandidate,
+ CandidateKind,
+)
+from .registry import register_strategy
+
+logger = logging.getLogger(__name__)
+
+
+def _build_guidelines_section(rules: List[str]) -> str:
+ """Build the `## Annotation Guidelines` section from a list of rules."""
+ rules_text = "\n".join(f"- {r.strip()}" for r in rules if r.strip())
+ return (
+ "## Annotation Guidelines\n\n"
+ "When distinguishing between similar labels, follow these rules:\n"
+ f"{rules_text}\n"
+ )
+
+
+def _replace_guidelines_section(current_prompt: str, new_rules: List[str]) -> str:
+ """Replace the existing `## Annotation Guidelines` / `## Refinement Guidelines`
+ section with a new one built from rules. Append if no section exists.
+ """
+ import re
+ new_section = _build_guidelines_section(new_rules)
+ pattern = r'## (?:Refinement |Annotation )?Guidelines[\s\S]*?(?=\n## |\Z)'
+ if re.search(pattern, current_prompt):
+ return re.sub(pattern, new_section, current_prompt).rstrip() + "\n"
+ else:
+ return current_prompt.rstrip() + "\n\n" + new_section
+
+
+def _extract_existing_rules(current_prompt: str) -> List[str]:
+ """Extract rules from the existing guidelines section, if any."""
+ import re
+ match = re.search(
+ r'## (?:Refinement |Annotation )?Guidelines[\s\S]*?\n([\s\S]*?)(?=\n## |\Z)',
+ current_prompt,
+ )
+ if not match:
+ return []
+ rules = []
+ for line in match.group(1).split('\n'):
+ stripped = line.strip()
+ if stripped.startswith('- ') and len(stripped) > 3:
+ rules.append(stripped[2:].strip())
+ return rules
+
+
+@register_strategy
+class ValidatedFocusedEditStrategy(RefinementStrategy):
+ """Produces prompt-rule candidates via the LLM; validation gate filters
+ those that don't improve over the baseline.
+
+ This is the safe default for small optimizer models (4Bโ7B). It generates
+ a small number of candidate rule sets and lets the framework pick the
+ winner on the val set.
+ """
+
+ NAME = "validated_focused_edit"
+ RECOMMENDED_OPTIMIZER_TIER = "small"
+ BEST_FOR = ["binary", "few_labels", "objective"]
+ DESCRIPTION = (
+ "Generate prompt-rule candidates and keep only those that improve "
+ "validation accuracy. Good default for small optimizer models."
+ )
+
+ def __init__(self, manager: Any, solo_config: Any):
+ super().__init__(manager, solo_config)
+ # Multiple candidate sets to give the framework choices
+ self.num_candidates = getattr(
+ solo_config.refinement_loop, "num_candidates", 3
+ )
+
+ def propose_candidates(
+ self,
+ patterns: List[Any],
+ current_prompt: str,
+ train_comparisons: List[Dict[str, Any]],
+ ) -> List[RefinementCandidate]:
+ """Generate N candidate rule sets via the existing confusion analyzer."""
+ if not patterns:
+ return []
+
+ analyzer = self.manager.confusion_analyzer
+ candidates: List[RefinementCandidate] = []
+
+ # Generate several independent candidate rule sets
+ for i in range(self.num_candidates):
+ try:
+ rules = analyzer.generate_guidelines_rewrite(patterns, current_prompt)
+ except Exception as e:
+ logger.warning(f"[ValidatedFocusedEdit] generation #{i} failed: {e}")
+ continue
+ if not rules:
+ continue
+
+ # Build the candidate prompt text by replacing the guidelines section
+ candidate_prompt = _replace_guidelines_section(current_prompt, rules)
+
+ candidates.append(RefinementCandidate(
+ kind=CandidateKind.PROMPT_EDIT,
+ payload={
+ "new_prompt_text": candidate_prompt,
+ "rules": rules,
+ },
+ target_pattern=f"top-{len(patterns)} confusion patterns",
+ proposed_by=self.NAME,
+ rationale=(
+ f"Candidate #{i+1}: {len(rules)} rules addressing "
+ f"{len(patterns)} confusion patterns"
+ ),
+ ))
+
+ logger.info(
+ f"[ValidatedFocusedEdit] Proposed {len(candidates)} candidate(s)"
+ )
+ return candidates
+
+
+@register_strategy
+class PrincipleICLStrategy(RefinementStrategy):
+ """Instead of editing the prompt, add validated instances as ICL examples.
+
+ For each disagreement, optionally extract a one-sentence principle (via LLM),
+ then validate by checking if adding this instance as an ICL example improves
+ accuracy on val. Accepted entries go to the ICL library.
+
+ Each candidate is an individual ICL_EXAMPLE โ the framework evaluates them
+ one at a time. Good for subjective tasks and small models where writing
+ rules fails.
+ """
+
+ NAME = "principle_icl"
+ RECOMMENDED_OPTIMIZER_TIER = "small"
+ BEST_FOR = ["subjective", "many_labels", "small_model"]
+ DESCRIPTION = (
+ "Add validated instances as in-context examples instead of editing "
+ "the prompt. Robust against narrow-rule overfitting."
+ )
+
+ def __init__(self, manager: Any, solo_config: Any):
+ super().__init__(manager, solo_config)
+ self.max_candidates = getattr(
+ solo_config.refinement_loop, "num_candidates", 5
+ )
+ self.extract_principle = True # ask LLM for a short rationale
+
+ def propose_candidates(
+ self,
+ patterns: List[Any],
+ current_prompt: str,
+ train_comparisons: List[Dict[str, Any]],
+ ) -> List[RefinementCandidate]:
+ """Propose each distinct disagreement as an ICL candidate.
+
+ Chooses one instance per confusion pattern so examples are diverse
+ (don't all come from the same predictedโactual confusion).
+ """
+ candidates: List[RefinementCandidate] = []
+
+ # Use confusion patterns to sample diverse disagreements
+ # (one per pattern, up to max_candidates)
+ seen_pairs = set()
+ analyzer = self.manager.confusion_analyzer
+
+ for pattern in patterns:
+ if len(candidates) >= self.max_candidates:
+ break
+ pair = (pattern.predicted_label, pattern.actual_label)
+ if pair in seen_pairs:
+ continue
+ seen_pairs.add(pair)
+
+ # Take the first example in this pattern
+ if not pattern.examples:
+ continue
+ ex = pattern.examples[0]
+
+ # Optionally extract a principle via LLM
+ principle = ""
+ if self.extract_principle:
+ try:
+ principle = self._extract_principle(pattern, ex, current_prompt)
+ except Exception as e:
+ logger.debug(f"[PrincipleICL] principle extraction failed: {e}")
+
+ candidates.append(RefinementCandidate(
+ kind=CandidateKind.ICL_EXAMPLE,
+ payload={
+ "instance_id": ex.instance_id,
+ "text": ex.text,
+ "label": pattern.actual_label,
+ "principle": principle,
+ },
+ target_pattern=f"{pattern.predicted_label}->{pattern.actual_label}",
+ proposed_by=self.NAME,
+ rationale=principle or "Disagreement example",
+ ))
+
+ logger.info(
+ f"[PrincipleICL] Proposed {len(candidates)} ICL candidate(s) "
+ f"from {len(patterns)} patterns"
+ )
+ return candidates
+
+ def _extract_principle(self, pattern, example, current_prompt: str) -> str:
+ """Ask the revision LLM for a one-sentence principle explaining the fix."""
+ analyzer = self.manager.confusion_analyzer
+ endpoint = analyzer._get_revision_endpoint()
+ if endpoint is None:
+ return ""
+
+ prompt = (
+ f"The following text was misclassified.\n\n"
+ f"Text: \"{example.text[:500]}\"\n"
+ f"Model's wrong label: {pattern.predicted_label}\n"
+ f"Correct label: {pattern.actual_label}\n\n"
+ f"In ONE short sentence (max 25 words), describe the general "
+ f"linguistic signal that makes this text a '{pattern.actual_label}' "
+ f"rather than '{pattern.predicted_label}'.\n"
+ f"Focus on general features, not specific phrases.\n\n"
+ f"Respond with JSON: {{\"principle\": \"\"}}"
+ )
+
+ try:
+ from pydantic import BaseModel
+
+ class PrincipleResponse(BaseModel):
+ principle: str = ""
+
+ try:
+ response = endpoint.query(prompt, PrincipleResponse)
+ except TypeError:
+ response = endpoint.query(prompt)
+
+ data = analyzer._parse_json(response)
+ return data.get("principle", "").strip()
+ except Exception as e:
+ logger.debug(f"[PrincipleICL] principle LLM call failed: {e}")
+ return ""
+
+
+@register_strategy
+class HybridDualTrackStrategy(RefinementStrategy):
+ """Try prompt edits first; fall back to ICL examples if edits fail.
+
+ This is the recommended default: combines ValidatedFocusedEdit and
+ PrincipleICL. On the first cycle, proposes both kinds; the framework's
+ validation gate picks whichever kind has the best candidate.
+
+ After 2 consecutive prompt-edit failures (tracked by failure counter),
+ future cycles propose ICL candidates only. This avoids wasted LLM cost
+ on a model that can't write good rules.
+ """
+
+ NAME = "hybrid_dual_track"
+ RECOMMENDED_OPTIMIZER_TIER = "small"
+ BEST_FOR = ["general_default", "unknown_data_properties"]
+ DESCRIPTION = (
+ "Try prompt edits; fall back to ICL examples on failure. "
+ "Recommended default for practitioners unsure of their data profile."
+ )
+
+ def __init__(self, manager: Any, solo_config: Any):
+ super().__init__(manager, solo_config)
+ self._focused_edit = ValidatedFocusedEditStrategy(manager, solo_config)
+ self._icl = PrincipleICLStrategy(manager, solo_config)
+
+ def propose_candidates(
+ self,
+ patterns: List[Any],
+ current_prompt: str,
+ train_comparisons: List[Dict[str, Any]],
+ ) -> List[RefinementCandidate]:
+ # Check consecutive failure count
+ consecutive_failures = getattr(
+ self.manager, "_refinement_consecutive_failures", 0
+ )
+
+ candidates: List[RefinementCandidate] = []
+
+ if consecutive_failures < 2:
+ # Try prompt edits
+ candidates.extend(self._focused_edit.propose_candidates(
+ patterns, current_prompt, train_comparisons
+ ))
+
+ # Always propose ICL candidates too โ framework picks best
+ candidates.extend(self._icl.propose_candidates(
+ patterns, current_prompt, train_comparisons
+ ))
+
+ logger.info(
+ f"[HybridDualTrack] Proposed {len(candidates)} total candidates "
+ f"(prompt_edits + ICL; failures_so_far={consecutive_failures})"
+ )
+ return candidates
+
+
+@register_strategy
+class LegacyAppendStrategy(RefinementStrategy):
+ """Original append-only refinement for ablation comparisons.
+
+ No validation gate. Appends generated rules to the prompt. Preserved as
+ a research baseline โ DO NOT use in production.
+ """
+
+ NAME = "legacy_append"
+ RECOMMENDED_OPTIMIZER_TIER = "small"
+ BEST_FOR = ["ablation", "research_baseline"]
+ DESCRIPTION = (
+ "Legacy append-only behavior. No validation. For ablation comparison only."
+ )
+
+ def propose_candidates(
+ self,
+ patterns: List[Any],
+ current_prompt: str,
+ train_comparisons: List[Dict[str, Any]],
+ ) -> List[RefinementCandidate]:
+ if not patterns:
+ return []
+ analyzer = self.manager.confusion_analyzer
+ try:
+ rules = analyzer.generate_guidelines_rewrite(patterns, current_prompt)
+ except Exception as e:
+ logger.warning(f"[LegacyAppend] generation failed: {e}")
+ return []
+ if not rules:
+ return []
+ new_prompt = _replace_guidelines_section(current_prompt, rules)
+ return [RefinementCandidate(
+ kind=CandidateKind.PROMPT_EDIT,
+ payload={"new_prompt_text": new_prompt, "rules": rules},
+ target_pattern=f"top-{len(patterns)} patterns",
+ proposed_by=self.NAME,
+ rationale="Legacy: appended without validation",
+ )]
diff --git a/potato/solo_mode/refinement/validation.py b/potato/solo_mode/refinement/validation.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d801f039b98a5368169966e125fbdf1995796a0
--- /dev/null
+++ b/potato/solo_mode/refinement/validation.py
@@ -0,0 +1,249 @@
+"""
+Validation infrastructure for refinement strategies.
+
+- ValidationSplit: deterministically splits disagreements into train/val
+- CandidateEvaluator: labels val instances with a candidate prompt/ICL
+ and returns accuracy against human labels
+"""
+
+from __future__ import annotations
+
+import logging
+import random
+from collections import Counter
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SplitResult:
+ train: List[Dict[str, Any]]
+ val: List[Dict[str, Any]]
+ seed: int
+
+
+class ValidationSplit:
+ """Split human-LLM comparison records into train/val deterministically.
+
+ Split is seeded by prompt_version so it's stable within a version
+ (a refinement cycle can re-run without the split changing) but
+ different across versions (preventing val leakage across cycles).
+
+ Only disagreements (agrees=False) are split, because the refinement
+ process works on disagreements. Agreements stay in the train side as
+ useful context but aren't needed for eval.
+ """
+
+ def __init__(
+ self,
+ val_ratio: float = 0.3,
+ min_val: int = 5,
+ min_train: int = 5,
+ prefer_consistent: bool = False,
+ ):
+ """
+ Args:
+ val_ratio: fraction of disagreements held out for validation
+ min_val: minimum val size; if fewer disagreements exist, returns empty val
+ min_train: minimum train size; if fewer, returns empty train
+ prefer_consistent: if True, prefer val instances that have disagreed
+ across โฅ2 labeling passes (systematic errors), falling back to
+ one-off disagreements only when too few qualify.
+ """
+ self.val_ratio = val_ratio
+ self.min_val = min_val
+ self.min_train = min_train
+ self.prefer_consistent = prefer_consistent
+
+ def split(
+ self,
+ comparisons: List[Dict[str, Any]],
+ prompt_version: int,
+ ) -> SplitResult:
+ """Split comparisons into train/val.
+
+ Args:
+ comparisons: list of {instance_id, human_label, llm_label, agrees, ...}
+ prompt_version: used to seed the split deterministically
+
+ Returns:
+ SplitResult with train and val lists; either can be empty if not
+ enough disagreements are available.
+ """
+ # Only disagreements go to val; agreements stay in train
+ disagreements = [c for c in comparisons if not c.get('agrees')]
+ agreements = [c for c in comparisons if c.get('agrees')]
+
+ if len(disagreements) < (self.min_val + self.min_train):
+ logger.info(
+ f"[ValidationSplit] Only {len(disagreements)} disagreements available; "
+ f"need at least {self.min_val + self.min_train}. Returning empty splits."
+ )
+ return SplitResult(train=[], val=[], seed=prompt_version)
+
+ rng = random.Random(f"val_split_v{prompt_version}")
+
+ # If preferring consistent disagreements: seed val from instances that
+ # have disagreed โฅ2 times. If that pool is too small, top up with
+ # one-off disagreements so we still meet min_val.
+ val, train_disagreements = self._partition(
+ disagreements, prompt_version, rng
+ )
+
+ # Combine train disagreements with agreements (useful context for
+ # rule generation โ but agreements aren't used for scoring)
+ train = train_disagreements + agreements
+
+ logger.info(
+ f"[ValidationSplit] v{prompt_version}: "
+ f"{len(train_disagreements)} train disagreements, "
+ f"{len(val)} val disagreements, "
+ f"{len(agreements)} agreements in train context"
+ )
+
+ return SplitResult(train=train, val=val, seed=prompt_version)
+
+ def _partition(
+ self,
+ disagreements: List[Dict[str, Any]],
+ prompt_version: int,
+ rng: random.Random,
+ ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+ """Pick val set (optionally preferring consistent disagreements) and return (val, train_disagreements)."""
+ target_val_size = max(self.min_val, int(len(disagreements) * self.val_ratio))
+
+ if not self.prefer_consistent:
+ shuffled = list(disagreements)
+ rng.shuffle(shuffled)
+ return shuffled[:target_val_size], shuffled[target_val_size:]
+
+ # Count disagreements per instance_id. An instance that shows up
+ # multiple times in the disagreement list has failed across at least
+ # that many labeling passes โ treat it as a systematic error rather
+ # than a one-off stochastic flip.
+ counts = Counter(c['instance_id'] for c in disagreements)
+
+ # Keep only the *latest* disagreement record per instance to avoid
+ # the same instance appearing multiple times in the val set.
+ latest_by_iid: Dict[str, Dict[str, Any]] = {}
+ for c in disagreements:
+ latest_by_iid[c['instance_id']] = c
+
+ consistent = [
+ latest_by_iid[iid] for iid, n in counts.items() if n >= 2
+ ]
+ oneoff = [
+ latest_by_iid[iid] for iid, n in counts.items() if n < 2
+ ]
+
+ rng.shuffle(consistent)
+ rng.shuffle(oneoff)
+
+ if len(consistent) >= target_val_size:
+ val = consistent[:target_val_size]
+ topup_used = 0
+ else:
+ # Not enough consistent disagreements โ top up with one-offs so
+ # we still meet min_val. This makes the filter a preference,
+ # not a hard gate.
+ need = target_val_size - len(consistent)
+ val = consistent + oneoff[:need]
+ topup_used = min(need, len(oneoff))
+
+ val_ids = {c['instance_id'] for c in val}
+ # Train gets all disagreement records whose instance_id isn't in val.
+ train_disagreements = [c for c in disagreements if c['instance_id'] not in val_ids]
+
+ logger.info(
+ f"[ValidationSplit] prefer_consistent: {len(consistent)} instances "
+ f"with โฅ2 disagreements, {len(oneoff)} one-offs; "
+ f"val drew {len(val) - topup_used} consistent + {topup_used} one-off"
+ )
+ return val, train_disagreements
+
+
+@dataclass
+class EvalResult:
+ accuracy: float
+ correct_count: int
+ total: int
+ per_instance: List[Dict[str, Any]] # {instance_id, predicted, human, correct}
+
+
+class CandidateEvaluator:
+ """Evaluate a candidate (prompt edit or ICL example) on a validation set.
+
+ Uses a single labeling call per instance (no sampling diversity for speed).
+ Compares predicted label against the human label already recorded.
+ """
+
+ def __init__(
+ self,
+ label_fn: Callable[[str, str, str], Optional[str]],
+ get_text_fn: Callable[[str], str],
+ ):
+ """
+ Args:
+ label_fn: callable(instance_id, text, prompt) -> predicted_label or None
+ get_text_fn: callable(instance_id) -> text string
+ """
+ self.label_fn = label_fn
+ self.get_text_fn = get_text_fn
+
+ def evaluate(
+ self,
+ candidate_prompt: str,
+ val_comparisons: List[Dict[str, Any]],
+ sample_size: Optional[int] = None,
+ ) -> EvalResult:
+ """Label each val instance with the candidate prompt, compute accuracy.
+
+ Args:
+ candidate_prompt: the full prompt text to evaluate
+ val_comparisons: list of comparison dicts with human_label
+ sample_size: if set, randomly sample this many from val_comparisons
+
+ Returns:
+ EvalResult with accuracy and per-instance breakdown
+ """
+ if sample_size and len(val_comparisons) > sample_size:
+ val_comparisons = random.sample(val_comparisons, sample_size)
+
+ correct = 0
+ per_instance = []
+ for comp in val_comparisons:
+ iid = comp['instance_id']
+ human_label = comp.get('human_label')
+ if human_label is None:
+ continue
+ try:
+ text = self.get_text_fn(iid)
+ predicted = self.label_fn(iid, text, candidate_prompt)
+ except Exception as e:
+ logger.warning(f"[CandidateEval] Failed to label {iid}: {e}")
+ predicted = None
+
+ is_correct = (
+ predicted is not None
+ and str(predicted) == str(human_label)
+ )
+ if is_correct:
+ correct += 1
+ per_instance.append({
+ 'instance_id': iid,
+ 'predicted': predicted,
+ 'human': human_label,
+ 'correct': is_correct,
+ })
+
+ total = len(per_instance)
+ accuracy = correct / total if total > 0 else 0.0
+
+ return EvalResult(
+ accuracy=accuracy,
+ correct_count=correct,
+ total=total,
+ per_instance=per_instance,
+ )
diff --git a/potato/solo_mode/refinement_loop.py b/potato/solo_mode/refinement_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..64beb55d209eda2e6a73380e50da3bc24e348205
--- /dev/null
+++ b/potato/solo_mode/refinement_loop.py
@@ -0,0 +1,326 @@
+"""
+Iterative Guideline Refinement Loop for Solo Mode
+
+Orchestrates the automated cycle:
+ confusion analysis โ guideline suggestions โ prompt revision โ re-annotation
+
+Monitors agreement rate trends and stops cycling when metrics plateau.
+"""
+
+import logging
+import threading
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RefinementCycle:
+ """Record of a single refinement cycle."""
+ cycle_number: int
+ started_at: str
+ completed_at: Optional[str] = None
+ agreement_rate_before: float = 0.0
+ agreement_rate_after: Optional[float] = None
+ improvement: Optional[float] = None
+ patterns_found: int = 0
+ suggestions_generated: int = 0
+ rules_applied: int = 0
+ reannotation_count: int = 0
+ prompt_version_before: int = 0
+ prompt_version_after: Optional[int] = None
+ status: str = "running" # running, completed, no_improvement, failed
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ 'cycle_number': self.cycle_number,
+ 'started_at': self.started_at,
+ 'completed_at': self.completed_at,
+ 'agreement_rate_before': self.agreement_rate_before,
+ 'agreement_rate_after': self.agreement_rate_after,
+ 'improvement': self.improvement,
+ 'patterns_found': self.patterns_found,
+ 'suggestions_generated': self.suggestions_generated,
+ 'rules_applied': self.rules_applied,
+ 'reannotation_count': self.reannotation_count,
+ 'prompt_version_before': self.prompt_version_before,
+ 'prompt_version_after': self.prompt_version_after,
+ 'status': self.status,
+ }
+
+
+class RefinementLoop:
+ """Orchestrates iterative confusion analysis โ prompt revision cycles.
+
+ The loop monitors annotation progress and periodically:
+ 1. Analyzes confusion patterns from human-LLM disagreements
+ 2. Generates guideline suggestions for top confusion patterns
+ 3. Injects suggestions into the prompt (if auto_apply or human-approved)
+ 4. Triggers re-annotation of low-confidence instances
+ 5. Measures improvement and decides whether to continue
+
+ The loop stops when:
+ - Agreement rate meets the target threshold
+ - Improvement plateaus (patience exceeded)
+ - Maximum cycles reached
+ """
+
+ def __init__(self, solo_config: Any, app_config: Dict[str, Any]):
+ self.solo_config = solo_config
+ self.app_config = app_config
+ self.rl_config = solo_config.refinement_loop
+
+ # Cycle tracking
+ self._cycles: List[RefinementCycle] = []
+ self._annotations_since_last_check: int = 0
+ self._consecutive_no_improvement: int = 0
+ self._stopped: bool = False
+ self._stop_reason: Optional[str] = None
+ self._running: bool = False
+ self._lock = threading.Lock()
+
+ @property
+ def cycle_count(self) -> int:
+ return len(self._cycles)
+
+ @property
+ def is_stopped(self) -> bool:
+ return self._stopped
+
+ @property
+ def stop_reason(self) -> Optional[str]:
+ return self._stop_reason
+
+ def record_annotation(self) -> bool:
+ """Record that a human annotation was made.
+
+ Returns:
+ True if a refinement cycle should be triggered.
+ """
+ if not self.rl_config.enabled or self._stopped:
+ return False
+
+ with self._lock:
+ self._annotations_since_last_check += 1
+ if self._annotations_since_last_check >= self.rl_config.trigger_interval:
+ # Reset counter immediately to prevent multiple triggers
+ # before run_cycle starts
+ self._annotations_since_last_check = 0
+ if self._running:
+ return False # Cycle already in progress
+ return True
+ return False
+
+ def should_trigger(self) -> bool:
+ """Check if conditions are met to trigger a refinement cycle."""
+ if not self.rl_config.enabled or self._stopped or self._running:
+ return False
+
+ with self._lock:
+ return self._annotations_since_last_check >= self.rl_config.trigger_interval
+
+ def run_cycle(
+ self,
+ agreement_rate: float,
+ prompt_version: int,
+ confusion_patterns: List[Any],
+ apply_suggestions_fn: Callable[[List[str]], Dict[str, Any]],
+ generate_suggestion_fn: Callable[[Any, str], Optional[str]],
+ current_prompt: str,
+ ) -> RefinementCycle:
+ """Execute one refinement cycle.
+
+ Args:
+ agreement_rate: Current agreement rate before the cycle.
+ prompt_version: Current prompt version number.
+ confusion_patterns: List of ConfusionPattern objects from analyzer.
+ apply_suggestions_fn: Callable that takes a list of suggestion strings
+ and applies them to the prompt. Returns dict with results.
+ generate_suggestion_fn: Callable(pattern, current_prompt) -> suggestion.
+ current_prompt: The current annotation prompt text.
+
+ Returns:
+ RefinementCycle record with results.
+ """
+ with self._lock:
+ if self._running:
+ raise RuntimeError("Refinement cycle already running")
+ self._running = True
+ self._annotations_since_last_check = 0
+
+ cycle = RefinementCycle(
+ cycle_number=self.cycle_count + 1,
+ started_at=datetime.now().isoformat(),
+ agreement_rate_before=agreement_rate,
+ prompt_version_before=prompt_version,
+ patterns_found=len(confusion_patterns),
+ )
+
+ try:
+ # Check if max cycles exceeded
+ if self.cycle_count >= self.rl_config.max_cycles:
+ cycle.status = "max_cycles_reached"
+ self._stop("Max refinement cycles reached")
+ self._finalize_cycle(cycle)
+ return cycle
+
+ # Generate suggestions for top patterns
+ suggestions = []
+ for pattern in confusion_patterns:
+ suggestion = generate_suggestion_fn(pattern, current_prompt)
+ if suggestion:
+ suggestions.append(suggestion)
+
+ cycle.suggestions_generated = len(suggestions)
+
+ if not suggestions:
+ cycle.status = "no_suggestions"
+ self._finalize_cycle(cycle)
+ return cycle
+
+ # Apply suggestions
+ if self.rl_config.auto_apply_suggestions:
+ result = apply_suggestions_fn(suggestions)
+ cycle.rules_applied = result.get('categories_incorporated', 0)
+ cycle.reannotation_count = result.get('reannotation_count', 0)
+ cycle.prompt_version_after = result.get('new_prompt_version')
+ cycle.status = "completed"
+ else:
+ # Suggestions generated but await human approval
+ cycle.status = "awaiting_approval"
+
+ self._finalize_cycle(cycle)
+ return cycle
+
+ except Exception as e:
+ logger.error(f"Refinement cycle {cycle.cycle_number} failed: {e}")
+ cycle.status = "failed"
+ self._finalize_cycle(cycle)
+ return cycle
+
+ def record_post_cycle_metrics(self, agreement_rate_after: float) -> None:
+ """Record the agreement rate after a cycle completes and re-annotation settles.
+
+ This should be called after enough new annotations have been collected
+ to measure the effect of the refinement.
+
+ Args:
+ agreement_rate_after: Agreement rate measured after the cycle.
+ """
+ with self._lock:
+ if not self._cycles:
+ return
+
+ last_cycle = self._cycles[-1]
+ if last_cycle.agreement_rate_after is not None:
+ return # Already recorded
+
+ last_cycle.agreement_rate_after = agreement_rate_after
+ improvement = agreement_rate_after - last_cycle.agreement_rate_before
+ last_cycle.improvement = round(improvement, 4)
+
+ if improvement < self.rl_config.min_improvement:
+ self._consecutive_no_improvement += 1
+ logger.info(
+ f"Refinement cycle {last_cycle.cycle_number}: "
+ f"no significant improvement ({improvement:+.4f}), "
+ f"patience {self._consecutive_no_improvement}/{self.rl_config.patience}"
+ )
+ if self._consecutive_no_improvement >= self.rl_config.patience:
+ self._stop("Improvement plateaued")
+ else:
+ self._consecutive_no_improvement = 0
+ logger.info(
+ f"Refinement cycle {last_cycle.cycle_number}: "
+ f"improvement {improvement:+.4f}"
+ )
+
+ def _finalize_cycle(self, cycle: RefinementCycle) -> None:
+ """Finalize a cycle and store it."""
+ cycle.completed_at = datetime.now().isoformat()
+ with self._lock:
+ self._cycles.append(cycle)
+ self._running = False
+
+ def _stop(self, reason: str) -> None:
+ """Stop the refinement loop."""
+ self._stopped = True
+ self._stop_reason = reason
+ logger.info(f"Refinement loop stopped: {reason}")
+
+ def reset(self) -> None:
+ """Reset the refinement loop state, allowing new cycles."""
+ with self._lock:
+ self._consecutive_no_improvement = 0
+ self._stopped = False
+ self._stop_reason = None
+ self._annotations_since_last_check = 0
+
+ def get_status(self) -> Dict[str, Any]:
+ """Get the current refinement loop status."""
+ with self._lock:
+ last_cycle = self._cycles[-1].to_dict() if self._cycles else None
+ last_improvement = None
+ if self._cycles and self._cycles[-1].improvement is not None:
+ last_improvement = self._cycles[-1].improvement
+
+ return {
+ 'enabled': self.rl_config.enabled,
+ 'total_cycles': len(self._cycles),
+ 'is_running': self._running,
+ 'is_stopped': self._stopped,
+ 'stop_reason': self._stop_reason,
+ 'consecutive_no_improvement': self._consecutive_no_improvement,
+ 'patience': self.rl_config.patience,
+ 'max_cycles': self.rl_config.max_cycles,
+ 'trigger_interval': self.rl_config.trigger_interval,
+ 'annotations_until_next': max(
+ 0,
+ self.rl_config.trigger_interval - self._annotations_since_last_check
+ ),
+ 'last_cycle': last_cycle,
+ 'last_improvement': last_improvement,
+ 'cycles': [c.to_dict() for c in self._cycles],
+ }
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize state for persistence."""
+ return {
+ 'cycles': [c.to_dict() for c in self._cycles],
+ 'annotations_since_last_check': self._annotations_since_last_check,
+ 'consecutive_no_improvement': self._consecutive_no_improvement,
+ 'stopped': self._stopped,
+ 'stop_reason': self._stop_reason,
+ }
+
+ def load_state(self, data: Dict[str, Any]) -> None:
+ """Restore state from persistence."""
+ with self._lock:
+ self._annotations_since_last_check = data.get(
+ 'annotations_since_last_check', 0
+ )
+ self._consecutive_no_improvement = data.get(
+ 'consecutive_no_improvement', 0
+ )
+ self._stopped = data.get('stopped', False)
+ self._stop_reason = data.get('stop_reason')
+
+ self._cycles = []
+ for cd in data.get('cycles', []):
+ self._cycles.append(RefinementCycle(
+ cycle_number=cd.get('cycle_number', 0),
+ started_at=cd.get('started_at', ''),
+ completed_at=cd.get('completed_at'),
+ agreement_rate_before=cd.get('agreement_rate_before', 0.0),
+ agreement_rate_after=cd.get('agreement_rate_after'),
+ improvement=cd.get('improvement'),
+ patterns_found=cd.get('patterns_found', 0),
+ suggestions_generated=cd.get('suggestions_generated', 0),
+ rules_applied=cd.get('rules_applied', 0),
+ reannotation_count=cd.get('reannotation_count', 0),
+ prompt_version_before=cd.get('prompt_version_before', 0),
+ prompt_version_after=cd.get('prompt_version_after'),
+ status=cd.get('status', 'completed'),
+ ))
diff --git a/potato/solo_mode/routes.py b/potato/solo_mode/routes.py
new file mode 100644
index 0000000000000000000000000000000000000000..048cf1fbd8ba37fd65c1c10bfc13098ffd76251f
--- /dev/null
+++ b/potato/solo_mode/routes.py
@@ -0,0 +1,1479 @@
+"""
+Solo Mode Routes
+
+Flask routes for Solo Mode human-LLM collaborative annotation.
+Provides endpoints for:
+- Setup and configuration
+- Prompt review and editing
+- Edge case labeling
+- Main annotation workflow
+- Disagreement resolution
+- Validation and status
+"""
+
+import json
+import logging
+import traceback
+from flask import (
+ Blueprint,
+ render_template,
+ request,
+ jsonify,
+ redirect,
+ url_for,
+ session,
+)
+from functools import wraps
+from typing import Any, Dict, Optional
+
+from .manager import get_solo_mode_manager
+from .phase_controller import SoloPhase
+from potato.item_state_management import get_item_state_manager
+
+logger = logging.getLogger(__name__)
+
+# Create blueprint
+solo_mode_bp = Blueprint('solo_mode', __name__, url_prefix='/solo')
+
+
+def solo_mode_required(f):
+ """Decorator to ensure Solo Mode is enabled and initialized."""
+ @wraps(f)
+ def decorated_function(*args, **kwargs):
+ manager = get_solo_mode_manager()
+ if manager is None:
+ return jsonify({'error': 'Solo Mode not enabled'}), 400
+ return f(*args, **kwargs)
+ return decorated_function
+
+
+def login_required(f):
+ """Decorator to require user authentication."""
+ @wraps(f)
+ def decorated_function(*args, **kwargs):
+ if 'username' not in session:
+ return redirect(url_for('login'))
+ return f(*args, **kwargs)
+ return decorated_function
+
+
+def api_login_required(f):
+ """Like login_required but returns JSON 401 instead of redirecting.
+
+ Intended for /api/* endpoints called by JS, where a 302 to /login is
+ useless (fetch follows it and gets HTML, not JSON).
+ """
+ @wraps(f)
+ def decorated_function(*args, **kwargs):
+ if 'username' not in session:
+ return jsonify({'error': 'Authentication required'}), 401
+ return f(*args, **kwargs)
+ return decorated_function
+
+
+def same_origin_required(f):
+ """CSRF protection for state-changing API routes.
+
+ Rejects requests whose Origin or Referer header doesn't match the host.
+ Browsers automatically attach these headers; cross-origin forms cannot
+ forge them. This is a lightweight CSRF defense without requiring token
+ machinery wired into every JS call site.
+
+ Allows requests with no Origin/Referer (server-to-server, curl from
+ admins with X-API-Key) since those aren't subject to CSRF.
+ """
+ @wraps(f)
+ def decorated_function(*args, **kwargs):
+ origin = request.headers.get('Origin')
+ referer = request.headers.get('Referer')
+ host = request.host_url.rstrip('/')
+
+ # If neither header is present, this isn't a browser request โ allow.
+ if not origin and not referer:
+ return f(*args, **kwargs)
+
+ if origin and not origin.startswith(host):
+ return jsonify({'error': 'Cross-origin request rejected'}), 403
+ if referer and not referer.startswith(host):
+ return jsonify({'error': 'Cross-origin request rejected'}), 403
+
+ return f(*args, **kwargs)
+ return decorated_function
+
+
+def admin_required(f):
+ """Require a valid admin API key (X-API-Key header) for destructive ops.
+
+ Bypasses the standard session login; used for endpoints that can corrupt
+ workflow state (forced phase transitions, refinement approval, etc.).
+ Falls back to allowing in debug mode via the existing admin key system.
+ """
+ @wraps(f)
+ def decorated_function(*args, **kwargs):
+ # Lazy import to avoid circular dependencies at module load.
+ from potato.server_utils.admin_key import validate_admin_api_key
+ from potato.flask_server import config as _config
+
+ api_key = (
+ request.headers.get('X-API-Key')
+ or session.get('admin_api_key')
+ )
+ if not validate_admin_api_key(api_key, _config):
+ return jsonify({'error': 'Admin authentication required'}), 403
+ return f(*args, **kwargs)
+ return decorated_function
+
+
+# =============================================================================
+# User Routes
+# =============================================================================
+
+@solo_mode_bp.route('/setup', methods=['GET', 'POST'])
+@login_required
+@solo_mode_required
+def setup():
+ """
+ Solo Mode setup page.
+
+ GET: Display setup form for task description
+ POST: Process task description and advance to prompt review
+ """
+ manager = get_solo_mode_manager()
+ current_phase = manager.get_current_phase()
+
+ # Guard: setup is only valid while still in SETUP phase. Once the user has
+ # advanced (and possibly started annotating), re-submission would silently
+ # overwrite task_description and append a stale prompt version.
+ already_configured = current_phase != SoloPhase.SETUP
+
+ if request.method == 'POST':
+ if already_configured:
+ return render_template(
+ 'solo/setup.html',
+ error=(
+ f'Setup is already complete (current phase: '
+ f'{current_phase.name.lower().replace("_", " ")}). '
+ 'Return to status to continue, or visit the prompt editor '
+ 'to refine your annotation prompt.'
+ ),
+ phase=current_phase.name.lower(),
+ already_configured=True,
+ ), 409
+
+ task_description = request.form.get('task_description', '')
+
+ if task_description:
+ manager.set_task_description(task_description)
+ manager.create_prompt_version(
+ f"Label the following text according to this task: {task_description}",
+ created_by='user_setup',
+ source_description='Initial prompt from task description'
+ )
+ manager.advance_to_phase(SoloPhase.PROMPT_REVIEW)
+ return redirect(url_for('solo_mode.prompt_editor'))
+
+ return render_template(
+ 'solo/setup.html',
+ error='Please provide a task description',
+ phase=current_phase.name.lower(),
+ already_configured=False,
+ )
+
+ return render_template(
+ 'solo/setup.html',
+ phase=current_phase.name.lower(),
+ already_configured=already_configured,
+ )
+
+
+@solo_mode_bp.route('/prompt', methods=['GET', 'POST'])
+@login_required
+@solo_mode_required
+def prompt_editor():
+ """
+ Prompt review and editing page.
+
+ GET: Display current prompt with editing interface
+ POST: Update prompt and optionally advance to edge case synthesis
+ """
+ manager = get_solo_mode_manager()
+
+ if request.method == 'POST':
+ action = request.form.get('action', '')
+
+ if action == 'update':
+ new_prompt = request.form.get('prompt', '')
+ if new_prompt:
+ manager.update_prompt(new_prompt, source='manual_edit')
+ return jsonify({'success': True})
+ return jsonify({'error': 'Prompt cannot be empty'}), 400
+
+ elif action == 'advance':
+ # Move to edge case synthesis
+ try:
+ manager.advance_to_phase(SoloPhase.EDGE_CASE_SYNTHESIS)
+ except ValueError:
+ pass # Already past this phase
+ return redirect(url_for('solo_mode.edge_cases'))
+
+ elif action == 'skip_to_annotation':
+ # Skip edge cases, go directly to parallel annotation
+ try:
+ manager.advance_to_phase(SoloPhase.PARALLEL_ANNOTATION)
+ except ValueError:
+ pass # Already in or past this phase
+ return redirect(url_for('solo_mode.annotate'))
+
+ # Get prompt history
+ prompt_history = []
+ for pv in manager.get_all_prompt_versions():
+ prompt_history.append({
+ 'version': pv.version,
+ 'prompt': pv.prompt_text,
+ 'source': pv.created_by,
+ 'timestamp': pv.created_at.isoformat(),
+ })
+
+ return render_template(
+ 'solo/prompt_editor.html',
+ current_prompt=manager.get_current_prompt_text(),
+ prompt_history=prompt_history,
+ phase=manager.get_current_phase().name.lower(),
+ )
+
+
+@solo_mode_bp.route('/edge-cases', methods=['GET', 'POST'])
+@login_required
+@solo_mode_required
+def edge_cases():
+ """
+ Edge case labeling page.
+
+ GET: Display edge cases for labeling
+ POST: Submit label for an edge case
+ """
+ manager = get_solo_mode_manager()
+
+ if request.method == 'POST':
+ case_id = request.form.get('case_id')
+ label = request.form.get('label')
+ notes = request.form.get('notes', '')
+
+ if case_id and label:
+ manager.edge_case_synthesizer.record_label(case_id, label, notes)
+
+ # Check if all edge cases are labeled
+ unlabeled = manager.edge_case_synthesizer.get_unlabeled_edge_cases()
+ if not unlabeled:
+ # Advance to prompt validation
+ manager.advance_to_phase(SoloPhase.PROMPT_VALIDATION)
+ return redirect(url_for('solo_mode.prompt_editor'))
+
+ return jsonify({'success': True, 'remaining': len(unlabeled)})
+
+ return jsonify({'error': 'Missing case_id or label'}), 400
+
+ # Generate edge cases if needed
+ if manager.get_current_phase() == SoloPhase.EDGE_CASE_SYNTHESIS:
+ unlabeled = manager.edge_case_synthesizer.get_unlabeled_edge_cases()
+
+ if not unlabeled:
+ # Synthesize new edge cases
+ manager.edge_case_synthesizer.synthesize_edge_cases(
+ task_description=manager.get_task_description() or '',
+ prompt=manager.get_current_prompt_text(),
+ num_cases=5,
+ )
+ unlabeled = manager.edge_case_synthesizer.get_unlabeled_edge_cases()
+
+ # Advance to labeling phase
+ manager.advance_to_phase(SoloPhase.EDGE_CASE_LABELING)
+
+ # Get edge cases to display
+ unlabeled = manager.edge_case_synthesizer.get_unlabeled_edge_cases()
+ current_case = unlabeled[0] if unlabeled else None
+
+ # Get available labels from config
+ labels = manager.get_available_labels()
+
+ return render_template(
+ 'solo/edge_cases.html',
+ current_case=current_case.to_dict() if current_case else None,
+ remaining_count=len(unlabeled),
+ labels=labels,
+ phase=manager.get_current_phase().name.lower(),
+ )
+
+
+@solo_mode_bp.route('/annotate', methods=['GET', 'POST'])
+@login_required
+@solo_mode_required
+def annotate():
+ """
+ Main annotation page for Solo Mode.
+
+ GET: Display next instance for annotation
+ POST: Submit annotation for an instance
+ """
+ manager = get_solo_mode_manager()
+ user_id = session.get('username', 'anonymous')
+
+ if request.method == 'POST':
+ instance_id = request.form.get('instance_id')
+ annotation = request.form.get('annotation')
+
+ if instance_id and annotation:
+ # Record human annotation
+ manager.record_human_annotation(instance_id, annotation, user_id)
+
+ # Check for disagreements
+ if manager.check_for_disagreement(instance_id, annotation):
+ # Redirect to disagreement resolution
+ session['disagreement_instance'] = instance_id
+ return redirect(url_for('solo_mode.disagreements'))
+
+ # Get next instance
+ return redirect(url_for('solo_mode.annotate'))
+
+ return jsonify({'error': 'Missing instance_id or annotation'}), 400
+
+ # Get next instance ID
+ instance_id = manager.get_next_instance_for_human(user_id)
+
+ # Get available labels (needed for all render paths)
+ labels = manager.get_available_labels()
+
+ if instance_id is None:
+ # Check if annotation is complete (atomic check-and-advance)
+ if manager.check_and_advance_to_autonomous():
+ return redirect(url_for('solo_mode.status'))
+
+ return render_template(
+ 'solo/annotate.html',
+ instance=None,
+ instance_id=None,
+ labels=labels,
+ message='No more instances available',
+ phase=manager.get_current_phase().name.lower(),
+ stats=manager.get_annotation_stats(),
+ )
+
+ # Get full instance data
+ try:
+ ism = get_item_state_manager()
+ item = ism.get_item(instance_id)
+ instance = {
+ 'id': instance_id,
+ 'text': item.get_displayed_text(),
+ 'data': item.get_data(),
+ }
+ except ValueError as e:
+ logger.error(f"ItemStateManager not initialized when fetching instance {instance_id}: {e}")
+ return render_template(
+ 'solo/annotate.html',
+ instance=None,
+ instance_id=None,
+ labels=labels,
+ message='Error: Item state manager not available. Please restart the server.',
+ phase=manager.get_current_phase().name.lower(),
+ )
+ except KeyError as e:
+ logger.error(f"Instance {instance_id} not found in ItemStateManager: {e}")
+ return render_template(
+ 'solo/annotate.html',
+ instance=None,
+ instance_id=None,
+ labels=labels,
+ message=f'Error: Instance {instance_id} not found.',
+ phase=manager.get_current_phase().name.lower(),
+ )
+
+ # Get LLM prediction if available
+ llm_prediction = manager.get_llm_prediction_for_instance(instance_id)
+
+ return render_template(
+ 'solo/annotate.html',
+ instance=instance,
+ instance_id=instance_id,
+ llm_prediction=llm_prediction,
+ labels=labels,
+ phase=manager.get_current_phase().name.lower(),
+ stats=manager.get_annotation_stats(),
+ )
+
+
+@solo_mode_bp.route('/disagreements', methods=['GET', 'POST'])
+@login_required
+@solo_mode_required
+def disagreements():
+ """
+ Disagreement resolution page.
+
+ GET: Display disagreement for resolution
+ POST: Submit resolution decision
+ """
+ manager = get_solo_mode_manager()
+
+ if request.method == 'POST':
+ disagreement_id = request.form.get('disagreement_id')
+ resolution = request.form.get('resolution') # 'human', 'llm', or custom label
+ notes = request.form.get('notes', '')
+
+ if disagreement_id and resolution:
+ # disagreement_id is "instance_id:schema_name"
+ parts = disagreement_id.split(':', 1)
+ instance_id = parts[0]
+ schema_name = parts[1] if len(parts) > 1 else 'default'
+
+ # Resolve "human"/"llm" to the actual label value
+ actual_label = resolution
+ if resolution == 'human':
+ # Use the human's label
+ disagreement = manager.get_disagreement(instance_id)
+ if disagreement:
+ actual_label = disagreement.get('human_label', resolution)
+ elif resolution == 'llm':
+ # Use the LLM's label
+ pred = manager.get_llm_prediction_for_instance(instance_id)
+ if pred:
+ actual_label = pred.get('label', resolution)
+
+ manager.resolve_disagreement(
+ instance_id, schema_name, actual_label, resolved_by='human'
+ )
+
+ # Check for more disagreements
+ pending = manager.get_pending_disagreements()
+ if not pending:
+ # Return to annotation
+ return redirect(url_for('solo_mode.annotate'))
+
+ return redirect(url_for('solo_mode.disagreements'))
+
+ return jsonify({'error': 'Missing disagreement_id or resolution'}), 400
+
+ # Get current disagreement
+ instance_id = session.pop('disagreement_instance', None)
+ if instance_id:
+ disagreement = manager.get_disagreement(instance_id)
+ else:
+ # Get next pending disagreement
+ pending = manager.get_pending_disagreements()
+ disagreement = manager.get_disagreement(pending[0]) if pending else None
+
+ if disagreement is None:
+ return redirect(url_for('solo_mode.annotate'))
+
+ # Get available labels
+ labels = manager.get_available_labels()
+
+ return render_template(
+ 'solo/disagreement.html',
+ disagreement=disagreement,
+ labels=labels,
+ phase=manager.get_current_phase().name.lower(),
+ )
+
+
+@solo_mode_bp.route('/review', methods=['GET', 'POST'])
+@login_required
+@solo_mode_required
+def review():
+ """
+ Periodic review of low-confidence LLM labels.
+
+ GET: Display instances for review
+ POST: Submit review decision
+ """
+ manager = get_solo_mode_manager()
+
+ if request.method == 'POST':
+ instance_id = request.form.get('instance_id')
+ decision = request.form.get('decision') # 'approve', 'correct'
+ corrected_label = request.form.get('corrected_label')
+
+ if instance_id and decision:
+ if decision == 'approve':
+ manager.approve_llm_label(instance_id)
+ elif decision == 'correct' and corrected_label:
+ manager.correct_llm_label(instance_id, corrected_label)
+
+ return redirect(url_for('solo_mode.review'))
+
+ return jsonify({'error': 'Invalid review data'}), 400
+
+ # Get instances for review
+ instances = manager.get_instances_for_review()
+
+ if not instances:
+ # Reset review counter and return to annotation
+ manager.validation_tracker.reset_periodic_review_counter()
+ return redirect(url_for('solo_mode.annotate'))
+
+ # Get available labels
+ labels = manager.get_available_labels()
+
+ return render_template(
+ 'solo/review.html',
+ instances=instances,
+ current_instance=instances[0] if instances else None,
+ labels=labels,
+ phase=manager.get_current_phase().name.lower(),
+ )
+
+
+@solo_mode_bp.route('/validation', methods=['GET', 'POST'])
+@login_required
+@solo_mode_required
+def validation():
+ """
+ Final validation of LLM-only labeled instances.
+
+ GET: Display validation interface
+ POST: Submit validation result
+ """
+ manager = get_solo_mode_manager()
+
+ if request.method == 'POST':
+ instance_id = request.form.get('instance_id')
+ human_label = request.form.get('human_label')
+ notes = request.form.get('notes', '')
+
+ if instance_id and human_label:
+ manager.record_validation(instance_id, human_label, notes)
+
+ # Check if validation is complete
+ progress = manager.get_validation_progress()
+ if progress['remaining'] == 0:
+ manager.advance_to_phase(SoloPhase.COMPLETED)
+ return redirect(url_for('solo_mode.status'))
+
+ return redirect(url_for('solo_mode.validation'))
+
+ return jsonify({'error': 'Missing validation data'}), 400
+
+ # Get validation samples
+ samples = manager.get_validation_samples()
+ current_sample = samples[0] if samples else None
+
+ # Get progress
+ progress = manager.get_validation_progress()
+
+ # Get available labels
+ labels = manager.get_available_labels()
+
+ return render_template(
+ 'solo/validation.html',
+ current_sample=current_sample,
+ progress=progress,
+ labels=labels,
+ phase=manager.get_current_phase().name.lower(),
+ )
+
+
+@solo_mode_bp.route('/rules', methods=['GET', 'POST'])
+@login_required
+@solo_mode_required
+def rule_review():
+ """
+ Edge case rule review page.
+
+ GET: Display aggregated categories for review
+ POST: Submit approval/rejection for a category
+ """
+ manager = get_solo_mode_manager()
+
+ if request.method == 'POST':
+ category_id = request.form.get('category_id')
+ action = request.form.get('action') # 'approve' or 'reject'
+ notes = request.form.get('notes', '')
+
+ if category_id and action:
+ ecr = manager.edge_case_rule_manager
+ if action == 'approve':
+ ecr.approve_category(category_id, notes)
+ elif action == 'reject':
+ ecr.reject_category(category_id, notes)
+
+ # Check if more categories pending
+ pending = ecr.get_pending_categories()
+ if not pending:
+ # All reviewed - return to annotation
+ current_phase = manager.get_current_phase()
+ if current_phase == SoloPhase.RULE_REVIEW:
+ manager.advance_to_phase(
+ SoloPhase.ACTIVE_ANNOTATION,
+ reason="All rule categories reviewed"
+ )
+ return redirect(url_for('solo_mode.annotate'))
+
+ return redirect(url_for('solo_mode.rule_review'))
+
+ return jsonify({'error': 'Missing category_id or action'}), 400
+
+ # Get rule data
+ ecr = manager.edge_case_rule_manager
+ pending = ecr.get_pending_categories()
+ approved = ecr.get_approved_categories()
+ rejected = ecr.get_rejected_categories()
+ stats = ecr.get_stats()
+
+ # Build category details with member rules
+ categories_with_rules = []
+ for cat in pending:
+ member_rules = []
+ for rid in cat.member_rule_ids:
+ rule = ecr.get_rule(rid)
+ if rule:
+ member_rules.append(rule.to_dict())
+ categories_with_rules.append({
+ 'category': cat.to_dict(),
+ 'member_rules': member_rules,
+ })
+
+ return render_template(
+ 'solo/rule_review.html',
+ pending_categories=categories_with_rules,
+ approved_count=len(approved),
+ rejected_count=len(rejected),
+ stats=stats,
+ phase=manager.get_current_phase().name.lower(),
+ )
+
+
+@solo_mode_bp.route('/status')
+@login_required
+@solo_mode_required
+def status():
+ """
+ Solo Mode status dashboard.
+
+ Tabbed dashboard with:
+ - Overview: annotation progress, agreement, LLM stats
+ - Edge Case Rules: inline rule review with approve/reject
+ - Rule Clusters: D3.js scatter plot visualization
+ """
+ manager = get_solo_mode_manager()
+
+ # Edge case rule data
+ edge_case_rule_stats = None
+ pending_categories = []
+ approved_count = 0
+ rejected_count = 0
+
+ if manager._edge_case_rule_manager is not None:
+ ecr = manager.edge_case_rule_manager
+ edge_case_rule_stats = ecr.get_stats()
+ approved_count = len(ecr.get_approved_categories())
+ rejected_count = len(ecr.get_rejected_categories())
+
+ for cat in ecr.get_pending_categories():
+ member_rules = []
+ for rid in cat.member_rule_ids:
+ rule = ecr.get_rule(rid)
+ if rule:
+ member_rules.append(rule.to_dict())
+ pending_categories.append({
+ 'category': cat.to_dict(),
+ 'member_rules': member_rules,
+ })
+
+ return render_template(
+ 'solo/status.html',
+ phase=manager.get_current_phase().name.lower(),
+ phase_name=manager.get_current_phase().name,
+ annotation_stats=manager.get_annotation_stats(),
+ agreement_metrics=manager.get_agreement_metrics(),
+ llm_stats=manager.get_llm_labeling_stats(),
+ validation_progress=manager.get_validation_progress(),
+ edge_case_rule_stats=edge_case_rule_stats,
+ pending_categories=pending_categories,
+ approved_count=approved_count,
+ rejected_count=rejected_count,
+ )
+
+
+# =============================================================================
+# Admin API Routes
+# =============================================================================
+
+@solo_mode_bp.route('/api/status')
+@solo_mode_required
+def api_status():
+ """Get comprehensive Solo Mode status."""
+ manager = get_solo_mode_manager()
+
+ return jsonify({
+ 'phase': manager.get_current_phase().name.lower(),
+ 'phase_name': manager.get_current_phase().name,
+ 'annotation_stats': manager.get_annotation_stats(),
+ 'agreement_metrics': manager.get_agreement_metrics().to_dict(),
+ 'llm_stats': manager.get_llm_labeling_stats(),
+ 'validation_progress': manager.get_validation_progress(),
+ 'should_end_human_annotation': manager.should_end_human_annotation(),
+ })
+
+
+@solo_mode_bp.route('/api/prompts')
+@solo_mode_required
+def api_prompts():
+ """Get prompt version history."""
+ manager = get_solo_mode_manager()
+
+ history = []
+ for pv in manager.get_all_prompt_versions():
+ history.append({
+ 'version': pv.version,
+ 'prompt': pv.prompt_text,
+ 'source': pv.created_by,
+ 'timestamp': pv.created_at.isoformat(),
+ 'changes': pv.source_description,
+ })
+
+ return jsonify({
+ 'current_prompt': manager.get_current_prompt_text(),
+ 'current_version': manager.current_prompt_version,
+ 'history': history,
+ })
+
+
+@solo_mode_bp.route('/api/predictions')
+@solo_mode_required
+def api_predictions():
+ """Get all LLM predictions."""
+ manager = get_solo_mode_manager()
+
+ predictions = manager.get_all_llm_predictions()
+ # Serialize predictions to dicts
+ serialized = {
+ iid: {s: p.to_dict() for s, p in schemas.items()}
+ for iid, schemas in predictions.items()
+ }
+
+ return jsonify({
+ 'count': len(predictions),
+ 'predictions': serialized,
+ })
+
+
+@solo_mode_bp.route('/api/advance-phase', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_advance_phase():
+ """Manually advance to a specific phase.
+
+ force=True bypasses the phase transition graph and can corrupt workflow
+ state; it requires admin authentication via X-API-Key.
+ """
+ manager = get_solo_mode_manager()
+
+ payload = request.get_json(silent=True) or {}
+ target_phase = payload.get('phase')
+ if not target_phase:
+ return jsonify({'error': 'Missing target phase'}), 400
+
+ force = bool(payload.get('force', False))
+ if force:
+ from potato.server_utils.admin_key import validate_admin_api_key
+ from potato.flask_server import config as _config
+ api_key = (
+ request.headers.get('X-API-Key')
+ or session.get('admin_api_key')
+ )
+ if not validate_admin_api_key(api_key, _config):
+ return jsonify({
+ 'error': 'force=True requires admin authentication',
+ }), 403
+
+ try:
+ phase = SoloPhase.from_str(target_phase)
+ except (ValueError, KeyError):
+ return jsonify({'error': f'Unknown phase: {target_phase}'}), 400
+
+ try:
+ success = manager.advance_to_phase(phase, force=force)
+
+ if success:
+ return jsonify({
+ 'success': True,
+ 'new_phase': manager.get_current_phase().name.lower(),
+ })
+ else:
+ return jsonify({
+ 'error': (
+ f'Invalid phase transition from '
+ f'{manager.get_current_phase().name} to {phase.name}'
+ ),
+ 'current_phase': manager.get_current_phase().name.lower(),
+ }), 400
+
+ except ValueError as e:
+ return jsonify({
+ 'error': str(e),
+ 'current_phase': manager.get_current_phase().name.lower(),
+ }), 400
+
+
+@solo_mode_bp.route('/api/pause-labeling', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_pause_labeling():
+ """Pause background LLM labeling."""
+ manager = get_solo_mode_manager()
+
+ if not manager.is_background_labeling_running():
+ return jsonify({'error': 'LLM labeling thread not running'}), 400
+
+ manager.pause_background_labeling()
+ return jsonify({'success': True, 'paused': True})
+
+
+@solo_mode_bp.route('/api/resume-labeling', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_resume_labeling():
+ """Resume background LLM labeling."""
+ manager = get_solo_mode_manager()
+
+ if not manager.is_background_labeling_running():
+ return jsonify({'error': 'LLM labeling thread not running'}), 400
+
+ manager.resume_background_labeling()
+ return jsonify({'success': True, 'paused': False})
+
+@solo_mode_bp.route('/api/start-labeling', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_start_labeling():
+ """Start background LLM labeling."""
+ manager = get_solo_mode_manager()
+ success = manager.start_background_labeling()
+ if success:
+ return jsonify({'success': True, 'message': 'LLM labeling started'})
+ return jsonify({'success': False, 'message': 'Already running or failed to start'})
+
+@solo_mode_bp.route('/api/optimize-prompt', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_optimize_prompt():
+ """Trigger prompt optimization."""
+ manager = get_solo_mode_manager()
+
+ # Access the (lazily-constructed) optimizer explicitly. Don't guard with
+ # hasattr(): hasattr() swallows any exception raised while building the
+ # property and would mislabel a real init failure as "not configured".
+ try:
+ optimizer = manager.prompt_optimizer if manager is not None else None
+ except Exception:
+ logger.error("Prompt optimizer failed to initialize: %s", traceback.format_exc())
+ return jsonify({'error': 'Prompt optimizer failed to initialize'}), 500
+
+ if optimizer is None:
+ return jsonify({'error': 'Prompt optimizer not configured'}), 400
+
+ try:
+ result = optimizer.optimize()
+ return jsonify({
+ 'success': True,
+ 'result': result,
+ })
+ except Exception as e:
+ logger.error("Prompt optimization failed: %s", traceback.format_exc())
+ return jsonify({'error': 'An internal error occurred'}), 500
+
+
+@solo_mode_bp.route('/api/disagreements')
+@solo_mode_required
+def api_disagreements():
+ """Get all disagreements and their status.
+
+ Reads from the manager's authoritative `disagreement_ids` set (populated
+ inside record_human_label). Keeps /api/disagreements consistent with the
+ Overview card and with get_agreement_metrics().
+ """
+ manager = get_solo_mode_manager()
+
+ pending = manager.get_pending_disagreements()
+ total = len(manager.disagreement_ids)
+ resolved = max(total - len(pending), 0)
+
+ return jsonify({
+ 'total': total,
+ 'pending': len(pending),
+ 'resolved': resolved,
+ 'pending_ids': pending,
+ })
+
+
+@solo_mode_bp.route('/api/edge-cases')
+@solo_mode_required
+def api_edge_cases():
+ """Get edge case status."""
+ manager = get_solo_mode_manager()
+
+ if manager.edge_case_synthesizer:
+ return jsonify(manager.edge_case_synthesizer.get_status())
+
+ return jsonify({
+ 'total_edge_cases': 0,
+ 'labeled': 0,
+ 'unlabeled': 0,
+ })
+
+
+@solo_mode_bp.route('/api/rules')
+@solo_mode_required
+def api_rules():
+ """Get all edge case rules and their status."""
+ manager = get_solo_mode_manager()
+ ecr = manager.edge_case_rule_manager
+
+ rules = [r.to_dict() for r in ecr.get_all_rules()]
+ return jsonify({
+ 'rules': rules,
+ 'stats': ecr.get_stats(),
+ })
+
+
+@solo_mode_bp.route('/api/rules/categories')
+@solo_mode_required
+def api_rules_categories():
+ """Get aggregated edge case rule categories."""
+ manager = get_solo_mode_manager()
+ ecr = manager.edge_case_rule_manager
+
+ categories = []
+ for cat in ecr.get_all_categories():
+ member_rules = []
+ for rid in cat.member_rule_ids:
+ rule = ecr.get_rule(rid)
+ if rule:
+ member_rules.append(rule.to_dict())
+ categories.append({
+ 'category': cat.to_dict(),
+ 'member_rules': member_rules,
+ })
+
+ return jsonify({'categories': categories})
+
+
+@solo_mode_bp.route('/api/rules/approve', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_rules_approve():
+ """Approve or reject an edge case rule category."""
+ manager = get_solo_mode_manager()
+ ecr = manager.edge_case_rule_manager
+
+ data = request.json or {}
+ category_id = data.get('category_id')
+ action = data.get('action', 'approve')
+ notes = data.get('notes', '')
+
+ if not category_id:
+ return jsonify({'error': 'Missing category_id'}), 400
+
+ if action == 'approve':
+ success = ecr.approve_category(category_id, notes)
+ elif action == 'reject':
+ success = ecr.reject_category(category_id, notes)
+ else:
+ return jsonify({'error': f'Invalid action: {action}'}), 400
+
+ return jsonify({'success': success})
+
+
+@solo_mode_bp.route('/api/rules/apply', methods=['POST'])
+@admin_required
+@solo_mode_required
+def api_rules_apply():
+ """Inject approved rules into the annotation prompt."""
+ manager = get_solo_mode_manager()
+
+ try:
+ result = manager.apply_approved_rules()
+ return jsonify(result)
+ except Exception as e:
+ logger.error("Error applying approved rules: %s", traceback.format_exc())
+ return jsonify({'error': 'An internal error occurred'}), 500
+
+
+@solo_mode_bp.route('/api/rules/cluster', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_rules_cluster():
+ """Manually trigger rule clustering."""
+ manager = get_solo_mode_manager()
+ manager._trigger_rule_clustering()
+ return jsonify({'success': True, 'message': 'Clustering triggered'})
+
+
+@solo_mode_bp.route('/api/rules/viz-data')
+@solo_mode_required
+def api_rules_viz_data():
+ """Return 2D-projected rule embeddings for D3 scatter plot visualization."""
+ manager = get_solo_mode_manager()
+ ecr = manager.edge_case_rule_manager
+ rules = ecr.get_all_rules()
+
+ if not rules:
+ return jsonify({'points': [], 'clusters': []})
+
+ # Project to 2D
+ try:
+ from .rule_clusterer import RuleClusterer
+ clusterer = RuleClusterer(manager.config, manager.solo_config)
+ coords = clusterer.project_to_2d(rules)
+ except Exception as e:
+ logger.warning(f"Rule projection failed: {e}")
+ coords = [(0.0, 0.0)] * len(rules)
+
+ # Build points
+ points = []
+ for i, rule in enumerate(rules):
+ x, y = coords[i] if i < len(coords) else (0.0, 0.0)
+ cat = ecr.get_category_for_rule(rule.id)
+ points.append({
+ 'x': float(x),
+ 'y': float(y),
+ 'rule_id': rule.id,
+ 'rule_text': rule.rule_text,
+ 'cluster_id': rule.cluster_id,
+ 'category_id': cat.id if cat else None,
+ 'category_summary': cat.summary_rule if cat else None,
+ 'confidence': rule.source_confidence,
+ 'instance_id': rule.instance_id,
+ 'approved': rule.approved,
+ 'reviewed': rule.reviewed,
+ })
+
+ # Build cluster info with centroids
+ clusters = []
+ for cat in ecr.get_all_categories():
+ member_indices = [
+ i for i, r in enumerate(rules) if r.cluster_id == cat.id
+ ]
+ if member_indices:
+ cx = sum(coords[i][0] for i in member_indices) / len(member_indices)
+ cy = sum(coords[i][1] for i in member_indices) / len(member_indices)
+ else:
+ cx, cy = 0.0, 0.0
+
+ clusters.append({
+ 'id': cat.id,
+ 'summary_rule': cat.summary_rule,
+ 'centroid_x': float(cx),
+ 'centroid_y': float(cy),
+ 'size': len(cat.member_rule_ids),
+ 'approved': cat.approved,
+ 'reviewed': cat.reviewed,
+ })
+
+ return jsonify({'points': points, 'clusters': clusters})
+
+
+@solo_mode_bp.route('/api/confusion-analysis')
+@solo_mode_required
+def api_confusion_analysis():
+ """Get full confusion analysis with enriched patterns and heatmap data."""
+ manager = get_solo_mode_manager()
+
+ try:
+ return jsonify(manager.get_confusion_analysis_full())
+ except Exception as e:
+ logger.error("Confusion analysis failed: %s", traceback.format_exc())
+ return jsonify({'enabled': False, 'error': 'An internal error occurred'}), 500
+
+
+@solo_mode_bp.route('/api/confusion-analysis/root-cause', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_confusion_root_cause():
+ """Generate root cause analysis for a confusion pattern."""
+ manager = get_solo_mode_manager()
+
+ data = request.json or {}
+ predicted = data.get('predicted_label')
+ actual = data.get('actual_label')
+
+ if not predicted or not actual:
+ return jsonify({'error': 'Missing predicted_label or actual_label'}), 400
+
+ # Find the pattern
+ analysis = manager.get_confusion_analysis_full()
+ if not analysis.get('enabled'):
+ return jsonify({'error': 'Confusion analysis not enabled'}), 400
+
+ pattern_data = None
+ for p in analysis.get('patterns', []):
+ if p['predicted_label'] == predicted and p['actual_label'] == actual:
+ pattern_data = p
+ break
+
+ if pattern_data is None:
+ return jsonify({'error': f'Pattern {predicted}->{actual} not found'}), 404
+
+ # Build a ConfusionPattern from the data
+ from .confusion_analyzer import ConfusionPattern, ConfusionExample
+ pattern = ConfusionPattern(
+ predicted_label=predicted,
+ actual_label=actual,
+ count=pattern_data['count'],
+ percent=pattern_data['percent'],
+ examples=[
+ ConfusionExample(
+ instance_id=e['instance_id'],
+ text=e.get('text', ''),
+ llm_reasoning=e.get('llm_reasoning'),
+ llm_confidence=e.get('llm_confidence'),
+ )
+ for e in pattern_data.get('examples', [])
+ ],
+ )
+
+ analyzer = manager.confusion_analyzer
+ root_cause = analyzer.generate_root_cause(pattern)
+
+ if root_cause is None:
+ return jsonify({
+ 'error': 'No LLM endpoint available for root cause analysis'
+ }), 503
+
+ return jsonify({'success': True, 'root_cause': root_cause})
+
+
+@solo_mode_bp.route('/api/confusion-analysis/suggest-guideline', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_confusion_suggest_guideline():
+ """Suggest a guideline to address a confusion pattern."""
+ manager = get_solo_mode_manager()
+
+ data = request.json or {}
+ predicted = data.get('predicted_label')
+ actual = data.get('actual_label')
+
+ if not predicted or not actual:
+ return jsonify({'error': 'Missing predicted_label or actual_label'}), 400
+
+ # Find the pattern
+ analysis = manager.get_confusion_analysis_full()
+ if not analysis.get('enabled'):
+ return jsonify({'error': 'Confusion analysis not enabled'}), 400
+
+ pattern_data = None
+ for p in analysis.get('patterns', []):
+ if p['predicted_label'] == predicted and p['actual_label'] == actual:
+ pattern_data = p
+ break
+
+ if pattern_data is None:
+ return jsonify({'error': f'Pattern {predicted}->{actual} not found'}), 404
+
+ from .confusion_analyzer import ConfusionPattern, ConfusionExample
+ pattern = ConfusionPattern(
+ predicted_label=predicted,
+ actual_label=actual,
+ count=pattern_data['count'],
+ percent=pattern_data['percent'],
+ examples=[
+ ConfusionExample(
+ instance_id=e['instance_id'],
+ text=e.get('text', ''),
+ llm_reasoning=e.get('llm_reasoning'),
+ llm_confidence=e.get('llm_confidence'),
+ )
+ for e in pattern_data.get('examples', [])
+ ],
+ root_cause=pattern_data.get('root_cause'),
+ )
+
+ analyzer = manager.confusion_analyzer
+
+ # Generate root cause first if not already available
+ if not pattern.root_cause:
+ pattern.root_cause = analyzer.generate_root_cause(pattern)
+
+ current_prompt = manager.get_current_prompt_text()
+ suggestion = analyzer.suggest_guideline(pattern, current_prompt)
+
+ if suggestion is None:
+ return jsonify({
+ 'error': 'No LLM endpoint available for guideline suggestion'
+ }), 503
+
+ return jsonify({
+ 'success': True,
+ 'suggestion': suggestion,
+ 'root_cause': pattern.root_cause,
+ })
+
+
+@solo_mode_bp.route('/api/refinement-status')
+@solo_mode_required
+def api_refinement_status():
+ """Get refinement loop status and cycle history."""
+ manager = get_solo_mode_manager()
+ return jsonify(manager.get_refinement_status())
+
+
+@solo_mode_bp.route('/api/refinement/trigger', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_refinement_trigger():
+ """Manually trigger a refinement cycle."""
+ manager = get_solo_mode_manager()
+
+ if not manager.config.refinement_loop.enabled:
+ return jsonify({'error': 'Refinement loop not enabled'}), 400
+
+ try:
+ result = manager.trigger_refinement_cycle()
+ return jsonify(result)
+ except Exception as e:
+ logger.error("Refinement trigger failed: %s", traceback.format_exc())
+ return jsonify({'error': 'An internal error occurred'}), 500
+
+
+@solo_mode_bp.route('/api/reannotation-report')
+@solo_mode_required
+def api_reannotation_report():
+ """Get before/after accuracy report for re-annotated instances."""
+ manager = get_solo_mode_manager()
+ return jsonify(manager.get_reannotation_report())
+
+
+@solo_mode_bp.route('/api/refinement/reset', methods=['POST'])
+@admin_required
+@solo_mode_required
+def api_refinement_reset():
+ """Reset the refinement loop, allowing new cycles."""
+ manager = get_solo_mode_manager()
+
+ if not manager.config.refinement_loop.enabled:
+ return jsonify({'error': 'Refinement loop not enabled'}), 400
+
+ manager.refinement_loop.reset()
+ # Also reset the validated framework's failure counter
+ if hasattr(manager, '_refinement_consecutive_failures'):
+ manager._refinement_consecutive_failures = 0
+ return jsonify({'success': True, 'message': 'Refinement loop reset'})
+
+
+@solo_mode_bp.route('/api/refinement/log')
+@solo_mode_required
+def api_refinement_log():
+ """Get the full log of refinement cycles (validated framework only).
+
+ Returns each cycle's result including whether it was applied, dry-run,
+ candidates, per-candidate val accuracy, and the baseline score.
+ """
+ manager = get_solo_mode_manager()
+ log = manager.get_refinement_log() if hasattr(manager, 'get_refinement_log') else []
+ return jsonify({'log': log, 'count': len(log)})
+
+
+@solo_mode_bp.route('/api/refinement/pending')
+@solo_mode_required
+def api_refinement_pending():
+ """Get refinement candidates awaiting admin approval.
+
+ Only populated when refinement_loop.require_approval is True. Each entry
+ includes the proposed change, validation scores, and rationale so an
+ admin can decide whether to apply.
+ """
+ manager = get_solo_mode_manager()
+ pending = manager.get_pending_refinements() if hasattr(manager, 'get_pending_refinements') else []
+ return jsonify({'pending': pending, 'count': len(pending)})
+
+
+@solo_mode_bp.route('/api/refinement/approve', methods=['POST'])
+@admin_required
+@solo_mode_required
+def api_refinement_approve():
+ """Apply a pending refinement candidate. Triggers re-annotation on apply."""
+ manager = get_solo_mode_manager()
+ data = request.get_json(silent=True) or {}
+ index = data.get('index')
+ if index is None or not isinstance(index, int):
+ return jsonify({'error': 'Missing integer index'}), 400
+
+ if not hasattr(manager, 'approve_pending_refinement'):
+ return jsonify({'error': 'Validated refinement not available'}), 400
+
+ result = manager.approve_pending_refinement(index)
+ status = 200 if result.get('success') else 400
+ return jsonify(result), status
+
+
+@solo_mode_bp.route('/api/refinement/reject', methods=['POST'])
+@admin_required
+@solo_mode_required
+def api_refinement_reject():
+ """Reject a pending refinement candidate."""
+ manager = get_solo_mode_manager()
+ data = request.get_json(silent=True) or {}
+ index = data.get('index')
+ if index is None or not isinstance(index, int):
+ return jsonify({'error': 'Missing integer index'}), 400
+
+ if not hasattr(manager, 'reject_pending_refinement'):
+ return jsonify({'error': 'Validated refinement not available'}), 400
+
+ result = manager.reject_pending_refinement(index)
+ status = 200 if result.get('success') else 400
+ return jsonify(result), status
+
+
+@solo_mode_bp.route('/api/refinement/strategies')
+@solo_mode_required
+def api_refinement_strategies():
+ """List available refinement strategies and their metadata."""
+ try:
+ from .refinement import list_strategies
+ return jsonify({'strategies': list_strategies()})
+ except Exception as e:
+ return jsonify({'error': str(e)}), 500
+
+
+@solo_mode_bp.route('/api/labeling-functions')
+@solo_mode_required
+def api_labeling_functions():
+ """Get all labeling functions and their stats."""
+ manager = get_solo_mode_manager()
+
+ status = manager.get_labeling_function_status()
+ if not status.get('enabled'):
+ return jsonify({'enabled': False})
+
+ functions = [
+ f.to_dict()
+ for f in manager.labeling_function_manager.get_all_functions()
+ ]
+
+ return jsonify({
+ **status,
+ 'functions': functions,
+ })
+
+
+@solo_mode_bp.route('/api/labeling-functions/extract', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_labeling_functions_extract():
+ """Trigger labeling function extraction from high-confidence predictions."""
+ manager = get_solo_mode_manager()
+
+ if not manager.config.labeling_functions.enabled:
+ return jsonify({'error': 'Labeling functions not enabled'}), 400
+
+ try:
+ result = manager.extract_labeling_functions()
+ return jsonify(result)
+ except Exception as e:
+ logger.error("Labeling function extraction failed: %s", traceback.format_exc())
+ return jsonify({'error': 'An internal error occurred'}), 500
+
+
+@solo_mode_bp.route('/api/labeling-functions//toggle', methods=['POST'])
+@api_login_required
+@same_origin_required
+@solo_mode_required
+def api_labeling_function_toggle(function_id):
+ """Toggle a labeling function's enabled state."""
+ manager = get_solo_mode_manager()
+
+ if not manager.config.labeling_functions.enabled:
+ return jsonify({'error': 'Labeling functions not enabled'}), 400
+
+ new_state = manager.labeling_function_manager.toggle_function(function_id)
+ if new_state is None:
+ return jsonify({'error': f'Function {function_id} not found'}), 404
+
+ return jsonify({'success': True, 'function_id': function_id, 'enabled': new_state})
+
+
+@solo_mode_bp.route('/api/labeling-functions/stats')
+@solo_mode_required
+def api_labeling_functions_stats():
+ """Get labeling function statistics."""
+ manager = get_solo_mode_manager()
+ return jsonify(manager.get_labeling_function_status())
+
+
+@solo_mode_bp.route('/api/disagreement-explorer')
+@solo_mode_required
+def api_disagreement_explorer():
+ """Get disagreement explorer data with scatter plots and label breakdowns."""
+ manager = get_solo_mode_manager()
+ label_filter = request.args.get('label')
+
+ try:
+ data = manager.get_disagreement_explorer_data(label_filter=label_filter)
+ return jsonify(data)
+ except Exception as e:
+ logger.error("Disagreement explorer failed: %s", traceback.format_exc())
+ return jsonify({'error': 'An internal error occurred'}), 500
+
+
+@solo_mode_bp.route('/api/disagreement-timeline')
+@solo_mode_required
+def api_disagreement_timeline():
+ """Get temporal disagreement trend data."""
+ manager = get_solo_mode_manager()
+ bucket_size = request.args.get('bucket_size', 10, type=int)
+ bucket_size = max(2, min(bucket_size, 100))
+
+ try:
+ data = manager.get_disagreement_timeline(bucket_size=bucket_size)
+ return jsonify(data)
+ except Exception as e:
+ logger.error("Disagreement timeline failed: %s", traceback.format_exc())
+ return jsonify({'error': 'An internal error occurred'}), 500
+
+
+@solo_mode_bp.route('/api/export')
+@solo_mode_required
+def api_export():
+ """Export all Solo Mode data."""
+ manager = get_solo_mode_manager()
+
+ # Serialize predictions to plain dicts
+ predictions = manager.get_all_llm_predictions()
+ serialized_predictions = {
+ iid: {s: p.to_dict() for s, p in schemas.items()}
+ for iid, schemas in predictions.items()
+ }
+
+ export_data = {
+ 'phase': manager.get_current_phase().name.lower(),
+ 'annotations': manager.get_all_annotations(),
+ 'llm_predictions': serialized_predictions,
+ 'disagreements': {
+ 'total': len(manager.disagreement_ids),
+ 'pending': len(manager.get_pending_disagreements()),
+ },
+ 'agreement_metrics': manager.get_agreement_metrics().to_dict(),
+ 'prompt_history': [
+ {
+ 'version': pv.version,
+ 'prompt': pv.prompt_text,
+ 'source': pv.created_by,
+ 'timestamp': pv.created_at.isoformat(),
+ }
+ for pv in manager.get_all_prompt_versions()
+ ],
+ }
+
+ return jsonify(export_data)
diff --git a/potato/solo_mode/rule_clusterer.py b/potato/solo_mode/rule_clusterer.py
new file mode 100644
index 0000000000000000000000000000000000000000..5da82d9e5a10e3bc33e2c3ca92dfc9302e49dbbc
--- /dev/null
+++ b/potato/solo_mode/rule_clusterer.py
@@ -0,0 +1,522 @@
+"""
+Rule Clusterer for Solo Mode
+
+Clusters edge case rules by semantic similarity, then aggregates each cluster
+into a summary category using an LLM. Follows the Co-DETECT pipeline:
+embed -> cluster -> aggregate -> merge redundant categories.
+
+Uses the same embedding approach as DiversityManager (sentence-transformers
+with TF-IDF fallback).
+"""
+
+import json
+import logging
+import re
+import uuid
+from typing import Any, Dict, List, Optional, Tuple
+
+from .edge_case_rules import EdgeCaseCategory, EdgeCaseRule
+
+logger = logging.getLogger(__name__)
+
+# Guarded imports
+try:
+ from sentence_transformers import SentenceTransformer
+ import numpy as np
+ from sklearn.cluster import KMeans
+ _SENTENCE_TRANSFORMERS_AVAILABLE = True
+except ImportError:
+ _SENTENCE_TRANSFORMERS_AVAILABLE = False
+ np = None
+ KMeans = None
+
+try:
+ from sklearn.feature_extraction.text import TfidfVectorizer
+ _SKLEARN_AVAILABLE = True
+except ImportError:
+ _SKLEARN_AVAILABLE = False
+
+
+AGGREGATION_PROMPT_TEMPLATE = """You are analyzing a cluster of edge case rules discovered during annotation.
+Each rule describes a situation where the annotator was uncertain about the correct label.
+
+Your task: Synthesize these individual rules into ONE concise summary rule that captures
+the common pattern across all rules in this cluster.
+
+Individual rules:
+{rules_text}
+
+Respond with JSON:
+{{
+ "summary_rule": " -> ' format>"
+}}
+"""
+
+MERGE_PROMPT_TEMPLATE = """You are reviewing edge case categories for redundancy.
+Determine if any of these categories are semantically redundant and should be merged.
+
+Categories:
+{categories_text}
+
+Respond with JSON:
+{{
+ "merge_groups": [
+ {{
+ "merged_summary": "",
+ "category_ids": ["", ""]
+ }}
+ ]
+}}
+
+If no categories should be merged, respond with:
+{{"merge_groups": []}}
+"""
+
+
+class RuleClusterer:
+ """Clusters edge case rules and aggregates them into categories.
+
+ Pipeline: embed -> cluster -> aggregate -> merge
+ """
+
+ def __init__(
+ self,
+ app_config: Dict[str, Any],
+ solo_config: Any,
+ ):
+ """Initialize the rule clusterer.
+
+ Args:
+ app_config: Full application configuration
+ solo_config: SoloModeConfig instance
+ """
+ self.app_config = app_config
+ self.solo_config = solo_config
+ self._model = None
+ self._endpoint = None
+
+ def _get_embedding_model(self) -> Optional[Any]:
+ """Get or create the sentence-transformer model."""
+ if not _SENTENCE_TRANSFORMERS_AVAILABLE:
+ return None
+ if self._model is None:
+ model_name = getattr(
+ self.solo_config.embedding, 'model_name', 'all-MiniLM-L6-v2'
+ )
+ try:
+ self._model = SentenceTransformer(model_name)
+ except Exception as e:
+ logger.warning(f"Could not load sentence-transformer model: {e}")
+ return None
+ return self._model
+
+ def _get_revision_endpoint(self) -> Optional[Any]:
+ """Get or create an AI endpoint for aggregation/merging."""
+ if self._endpoint is not None:
+ return self._endpoint
+
+ try:
+ from potato.ai.ai_endpoint import AIEndpointFactory
+
+ models = self.solo_config.revision_models or self.solo_config.labeling_models
+ for model_config in models:
+ try:
+ endpoint_config = model_config.to_endpoint_config(temperature_override=0.3)
+
+ endpoint = AIEndpointFactory.create_endpoint(endpoint_config)
+ if endpoint:
+ self._endpoint = endpoint
+ return endpoint
+ except Exception:
+ continue
+ except Exception as e:
+ logger.warning(f"Could not create revision endpoint: {e}")
+
+ return None
+
+ def embed_rules(self, rules: List[EdgeCaseRule]) -> Optional[Any]:
+ """Compute embeddings for rule texts.
+
+ Args:
+ rules: List of edge case rules to embed
+
+ Returns:
+ Numpy array of embeddings, or None if embedding fails
+ """
+ texts = [r.rule_text for r in rules]
+
+ # Try sentence-transformers first
+ model = self._get_embedding_model()
+ if model is not None:
+ try:
+ embeddings = model.encode(texts, show_progress_bar=False)
+ return embeddings
+ except Exception as e:
+ logger.warning(f"Sentence-transformer embedding failed: {e}")
+
+ # Fallback to TF-IDF
+ return self._tfidf_embed(texts)
+
+ def _tfidf_embed(self, texts: List[str]) -> Optional[Any]:
+ """Fallback TF-IDF embedding."""
+ if not _SKLEARN_AVAILABLE:
+ logger.warning(
+ "Neither sentence-transformers nor sklearn available for embedding"
+ )
+ return None
+
+ try:
+ vectorizer = TfidfVectorizer(max_features=256, stop_words='english')
+ embeddings = vectorizer.fit_transform(texts).toarray()
+ return embeddings
+ except Exception as e:
+ logger.warning(f"TF-IDF embedding failed: {e}")
+ return None
+
+ def project_to_2d(
+ self,
+ rules: List[EdgeCaseRule],
+ ) -> List[Tuple[float, float]]:
+ """Project rule embeddings to 2D coordinates for visualization.
+
+ Uses PCA (or first 2 dimensions as fallback) to reduce
+ high-dimensional embeddings to plottable 2D points.
+
+ Args:
+ rules: Rules to project
+
+ Returns:
+ List of (x, y) tuples, one per rule
+ """
+ if not rules:
+ return []
+
+ embeddings = self.embed_rules(rules)
+
+ if embeddings is None:
+ return [(0.0, 0.0)] * len(rules)
+
+ try:
+ import numpy as _np
+ emb_array = _np.array(embeddings)
+ except (ImportError, Exception):
+ # Raw fallback: first 2 dimensions
+ result = []
+ for e in embeddings:
+ row = list(e) if hasattr(e, '__iter__') else [0.0]
+ x = float(row[0]) if len(row) > 0 else 0.0
+ y = float(row[1]) if len(row) > 1 else 0.0
+ result.append((x, y))
+ return result
+
+ if emb_array.shape[0] < 2:
+ return [(float(emb_array[0, 0]) if emb_array.shape[1] > 0 else 0.0,
+ float(emb_array[0, 1]) if emb_array.shape[1] > 1 else 0.0)]
+
+ # Try PCA
+ try:
+ from sklearn.decomposition import PCA
+
+ n_components = min(2, emb_array.shape[0], emb_array.shape[1])
+ pca = PCA(n_components=n_components)
+ coords = pca.fit_transform(emb_array)
+
+ result = []
+ for c in coords:
+ x = float(c[0])
+ y = float(c[1]) if n_components > 1 else 0.0
+ result.append((x, y))
+ return result
+
+ except ImportError:
+ pass
+
+ # Fallback: first 2 dimensions
+ result = []
+ for i in range(emb_array.shape[0]):
+ x = float(emb_array[i, 0]) if emb_array.shape[1] > 0 else 0.0
+ y = float(emb_array[i, 1]) if emb_array.shape[1] > 1 else 0.0
+ result.append((x, y))
+ return result
+
+ def cluster_rules(
+ self,
+ rules: List[EdgeCaseRule],
+ embeddings: Any,
+ ) -> Dict[int, List[EdgeCaseRule]]:
+ """Cluster rules using size-constrained K-Means.
+
+ Args:
+ rules: Rules to cluster
+ embeddings: Precomputed embeddings (numpy array)
+
+ Returns:
+ Dict mapping cluster_id to list of rules in that cluster
+ """
+ if embeddings is None or len(rules) == 0:
+ return {0: rules}
+
+ if not _SENTENCE_TRANSFORMERS_AVAILABLE and np is None:
+ try:
+ import numpy as _np
+ except ImportError:
+ return {0: rules}
+
+ _np = np
+ if _np is None:
+ import numpy as _np
+
+ target_size = self.solo_config.edge_case_rules.target_cluster_size
+ n_clusters = max(1, len(rules) // target_size + 1)
+
+ # Cap clusters at number of rules
+ n_clusters = min(n_clusters, len(rules))
+
+ if n_clusters <= 1:
+ return {0: rules}
+
+ try:
+ from sklearn.cluster import KMeans as _KMeans
+
+ kmeans = _KMeans(
+ n_clusters=n_clusters,
+ random_state=42,
+ n_init=10,
+ )
+ labels = kmeans.fit_predict(embeddings)
+
+ # Build cluster dict
+ clusters: Dict[int, List[EdgeCaseRule]] = {}
+ for rule, label in zip(rules, labels):
+ cluster_id = int(label)
+ if cluster_id not in clusters:
+ clusters[cluster_id] = []
+ clusters[cluster_id].append(rule)
+
+ # Redistribute oversized/undersized clusters
+ clusters = self._rebalance_clusters(clusters, target_size)
+
+ return clusters
+
+ except Exception as e:
+ logger.warning(f"Clustering failed: {e}")
+ return {0: rules}
+
+ def _rebalance_clusters(
+ self,
+ clusters: Dict[int, List[EdgeCaseRule]],
+ target_size: int,
+ ) -> Dict[int, List[EdgeCaseRule]]:
+ """Redistribute items from oversized to undersized clusters.
+
+ Ensures clusters stay within [target_size/2, target_size*2] range
+ when possible.
+ """
+ max_size = target_size * 2
+ min_size = max(1, target_size // 2)
+
+ # Collect overflow items
+ overflow = []
+ for cid, members in list(clusters.items()):
+ if len(members) > max_size:
+ overflow.extend(members[max_size:])
+ clusters[cid] = members[:max_size]
+
+ # Distribute overflow to undersized clusters
+ for cid in list(clusters.keys()):
+ if not overflow:
+ break
+ deficit = min_size - len(clusters[cid])
+ if deficit > 0:
+ to_add = overflow[:deficit]
+ clusters[cid].extend(to_add)
+ overflow = overflow[deficit:]
+
+ # If still overflow, create new clusters
+ if overflow:
+ new_id = max(clusters.keys()) + 1
+ while overflow:
+ batch = overflow[:target_size]
+ overflow = overflow[target_size:]
+ clusters[new_id] = batch
+ new_id += 1
+
+ return clusters
+
+ def aggregate_cluster(
+ self,
+ cluster_rules: List[EdgeCaseRule],
+ ) -> Optional[str]:
+ """Synthesize a summary rule from a cluster of similar rules.
+
+ Uses the revision model to produce a concise summary.
+
+ Args:
+ cluster_rules: Rules in a single cluster
+
+ Returns:
+ Summary rule text, or None if aggregation fails
+ """
+ if not cluster_rules:
+ return None
+
+ # If only one rule, use it directly
+ if len(cluster_rules) == 1:
+ return cluster_rules[0].rule_text
+
+ endpoint = self._get_revision_endpoint()
+ if endpoint is None:
+ # Fallback: use the first rule as representative
+ return cluster_rules[0].rule_text
+
+ rules_text = "\n".join(
+ f"- {r.rule_text}" for r in cluster_rules
+ )
+ prompt = AGGREGATION_PROMPT_TEMPLATE.format(rules_text=rules_text)
+
+ try:
+ response = endpoint.query(prompt)
+ response_data = self._parse_json(response)
+ summary = response_data.get('summary_rule', '')
+ if summary:
+ return summary
+ except Exception as e:
+ logger.warning(f"Cluster aggregation failed: {e}")
+
+ # Fallback
+ return cluster_rules[0].rule_text
+
+ def merge_categories(
+ self,
+ categories: List[EdgeCaseCategory],
+ ) -> List[EdgeCaseCategory]:
+ """Detect and merge redundant categories.
+
+ Uses embedding similarity to find near-duplicates, then
+ uses the LLM to merge them.
+
+ Args:
+ categories: List of categories to check for redundancy
+
+ Returns:
+ Deduplicated list of categories
+ """
+ if len(categories) <= 1:
+ return categories
+
+ endpoint = self._get_revision_endpoint()
+ if endpoint is None:
+ return categories
+
+ categories_text = "\n".join(
+ f"- ID: {c.id} | Rule: {c.summary_rule}"
+ for c in categories
+ )
+ prompt = MERGE_PROMPT_TEMPLATE.format(categories_text=categories_text)
+
+ try:
+ response = endpoint.query(prompt)
+ response_data = self._parse_json(response)
+ merge_groups = response_data.get('merge_groups', [])
+
+ if not merge_groups:
+ return categories
+
+ # Build category lookup
+ cat_map = {c.id: c for c in categories}
+ merged_ids = set()
+
+ result = []
+ for group in merge_groups:
+ group_ids = group.get('category_ids', [])
+ merged_summary = group.get('merged_summary', '')
+ if len(group_ids) < 2 or not merged_summary:
+ continue
+
+ # Combine member rules from all categories in group
+ combined_members = []
+ for cid in group_ids:
+ if cid in cat_map:
+ combined_members.extend(cat_map[cid].member_rule_ids)
+ merged_ids.add(cid)
+
+ # Create merged category
+ new_cat = EdgeCaseCategory(
+ id=f"cat_{uuid.uuid4().hex[:8]}",
+ summary_rule=merged_summary,
+ member_rule_ids=combined_members,
+ )
+ result.append(new_cat)
+
+ # Add categories that weren't merged
+ for c in categories:
+ if c.id not in merged_ids:
+ result.append(c)
+
+ return result
+
+ except Exception as e:
+ logger.warning(f"Category merging failed: {e}")
+ return categories
+
+ def run_full_pipeline(
+ self,
+ rules: List[EdgeCaseRule],
+ ) -> List[EdgeCaseCategory]:
+ """Run the complete clustering pipeline: embed -> cluster -> aggregate -> merge.
+
+ Args:
+ rules: Unclustered edge case rules
+
+ Returns:
+ List of EdgeCaseCategory objects
+ """
+ if not rules:
+ return []
+
+ logger.info(f"Starting rule clustering pipeline with {len(rules)} rules")
+
+ # Step 1: Embed
+ embeddings = self.embed_rules(rules)
+
+ # Step 2: Cluster
+ clusters = self.cluster_rules(rules, embeddings)
+ logger.info(f"Formed {len(clusters)} clusters")
+
+ # Step 3: Aggregate each cluster into a category
+ categories = []
+ for cluster_id, cluster_rules in clusters.items():
+ summary = self.aggregate_cluster(cluster_rules)
+ if summary:
+ cat = EdgeCaseCategory(
+ id=f"cat_{uuid.uuid4().hex[:8]}",
+ summary_rule=summary,
+ member_rule_ids=[r.id for r in cluster_rules],
+ )
+ categories.append(cat)
+
+ # Step 4: Merge redundant categories
+ if len(categories) > 1:
+ categories = self.merge_categories(categories)
+
+ logger.info(f"Pipeline complete: {len(categories)} categories")
+ return categories
+
+ def _parse_json(self, response: Any) -> Dict[str, Any]:
+ """Parse JSON from an LLM response."""
+ if isinstance(response, dict):
+ return response
+ if hasattr(response, 'model_dump'):
+ return response.model_dump()
+
+ content = str(response).strip()
+
+ # Extract from markdown code blocks
+ match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', content)
+ if match:
+ content = match.group(1).strip()
+
+ try:
+ return json.loads(content)
+ except json.JSONDecodeError:
+ return {}
diff --git a/potato/solo_mode/uncertainty/__init__.py b/potato/solo_mode/uncertainty/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c794eb246db235ee5bf6cd32b44adc574a69896
--- /dev/null
+++ b/potato/solo_mode/uncertainty/__init__.py
@@ -0,0 +1,32 @@
+"""
+Uncertainty Estimation Module
+
+This module provides pluggable strategies for estimating LLM prediction uncertainty.
+
+Available strategies:
+- DirectConfidenceEstimator: Ask model for confidence score (0-100)
+- DirectUncertaintyEstimator: Ask model for uncertainty score (0-100)
+- TokenEntropyEstimator: Entropy of answer token probabilities (requires logprobs)
+- SamplingDiversityEstimator: Multiple runs at high temperature, measure label diversity
+"""
+
+from .base import UncertaintyEstimator, UncertaintyEstimate
+from .direct_confidence import DirectConfidenceEstimator
+from .direct_uncertainty import DirectUncertaintyEstimator
+from .token_entropy import TokenEntropyEstimator
+from .sampling_diversity import SamplingDiversityEstimator
+from .factory import UncertaintyEstimatorFactory, create_uncertainty_estimator
+
+__all__ = [
+ # Base
+ 'UncertaintyEstimator',
+ 'UncertaintyEstimate',
+ # Implementations
+ 'DirectConfidenceEstimator',
+ 'DirectUncertaintyEstimator',
+ 'TokenEntropyEstimator',
+ 'SamplingDiversityEstimator',
+ # Factory
+ 'UncertaintyEstimatorFactory',
+ 'create_uncertainty_estimator',
+]
diff --git a/potato/solo_mode/uncertainty/base.py b/potato/solo_mode/uncertainty/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..4f730118aab77e42f5d05623628550b4f45f29f5
--- /dev/null
+++ b/potato/solo_mode/uncertainty/base.py
@@ -0,0 +1,128 @@
+"""
+Uncertainty Estimator Base Classes
+
+This module defines the abstract base class for uncertainty estimation strategies.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class UncertaintyEstimate:
+ """
+ Result of uncertainty estimation for a single prediction.
+
+ Attributes:
+ uncertainty_score: The uncertainty score (0.0 = certain, 1.0 = uncertain)
+ confidence_score: The confidence score (1.0 - uncertainty_score)
+ method: The estimation method used
+ metadata: Additional method-specific information
+ """
+ uncertainty_score: float
+ confidence_score: float
+ method: str
+ timestamp: datetime = field(default_factory=datetime.now)
+ metadata: Dict[str, Any] = field(default_factory=dict)
+
+ def __post_init__(self):
+ """Ensure scores are in valid range."""
+ self.uncertainty_score = max(0.0, min(1.0, self.uncertainty_score))
+ self.confidence_score = max(0.0, min(1.0, self.confidence_score))
+
+ def to_dict(self) -> Dict[str, Any]:
+ """Serialize to dictionary."""
+ return {
+ 'uncertainty_score': self.uncertainty_score,
+ 'confidence_score': self.confidence_score,
+ 'method': self.method,
+ 'timestamp': self.timestamp.isoformat(),
+ 'metadata': self.metadata,
+ }
+
+
+class UncertaintyEstimator(ABC):
+ """
+ Abstract base class for uncertainty estimation strategies.
+
+ Subclasses implement different methods for estimating how uncertain
+ an LLM is about its predictions. Strategies include:
+ - Directly asking the model for confidence/uncertainty scores
+ - Analyzing token probability distributions (entropy)
+ - Running multiple samples and measuring label diversity
+ """
+
+ def __init__(self, config: Optional[Dict[str, Any]] = None):
+ """
+ Initialize the uncertainty estimator.
+
+ Args:
+ config: Strategy-specific configuration
+ """
+ self.config = config or {}
+ self.name = self.__class__.__name__
+
+ @abstractmethod
+ def estimate_uncertainty(
+ self,
+ instance_id: str,
+ text: str,
+ prompt: str,
+ predicted_label: Any,
+ endpoint: Any,
+ schema_info: Optional[Dict[str, Any]] = None
+ ) -> UncertaintyEstimate:
+ """
+ Estimate uncertainty for a prediction.
+
+ Args:
+ instance_id: The instance being labeled
+ text: The text content to label
+ prompt: The labeling prompt
+ predicted_label: The label that was predicted
+ endpoint: The AI endpoint to use
+ schema_info: Optional annotation schema information
+
+ Returns:
+ UncertaintyEstimate with uncertainty and confidence scores
+ """
+ pass
+
+ @abstractmethod
+ def supports_endpoint(self, endpoint: Any) -> bool:
+ """
+ Check if this strategy supports the given endpoint.
+
+ Some strategies (like token entropy) require specific endpoint
+ capabilities like logprobs support.
+
+ Args:
+ endpoint: The AI endpoint to check
+
+ Returns:
+ True if this strategy can be used with the endpoint
+ """
+ pass
+
+ def get_method_name(self) -> str:
+ """Get the name of this estimation method."""
+ return self.name
+
+ def get_config_defaults(self) -> Dict[str, Any]:
+ """Get default configuration values for this strategy."""
+ return {}
+
+ def validate_config(self) -> List[str]:
+ """
+ Validate the configuration.
+
+ Returns:
+ List of validation error messages (empty if valid)
+ """
+ return []
diff --git a/potato/solo_mode/uncertainty/direct_confidence.py b/potato/solo_mode/uncertainty/direct_confidence.py
new file mode 100644
index 0000000000000000000000000000000000000000..a67b299eafa3098df977da58631e19f03bfe22b4
--- /dev/null
+++ b/potato/solo_mode/uncertainty/direct_confidence.py
@@ -0,0 +1,161 @@
+"""
+Direct Confidence Estimation
+
+This strategy asks the LLM directly for a confidence score on its prediction.
+"""
+
+import json
+import logging
+import re
+from typing import Any, Dict, Optional
+
+from .base import UncertaintyEstimator, UncertaintyEstimate
+
+logger = logging.getLogger(__name__)
+
+
+class DirectConfidenceEstimator(UncertaintyEstimator):
+ """
+ Estimate uncertainty by asking the model for a confidence score.
+
+ This is the simplest and most broadly compatible strategy. It appends
+ a request for a confidence score to the labeling prompt and parses
+ the model's self-reported confidence.
+
+ Pros:
+ - Works with all LLM endpoints
+ - Simple to implement and understand
+ - Can get confidence explanation/reasoning
+
+ Cons:
+ - Models may be miscalibrated (overconfident or underconfident)
+ - Adds to prompt length and response time
+ """
+
+ CONFIDENCE_PROMPT = """
+
+After providing your label, also rate your confidence in this label on a scale from 0 to 100, where:
+- 0 = completely uncertain, essentially guessing
+- 50 = somewhat confident but could easily be wrong
+- 100 = absolutely certain
+
+Respond in JSON format:
+{
+ "label": "",
+ "confidence": <0-100>,
+ "reasoning": "