data_studio/sources.py

from __future__ import annotations

import os
from typing import Any

from bson import ObjectId
from pymongo.database import Database

from data_studio.utils import utc_now_iso


def default_source_configs() -> list[dict[str, Any]]:
    now = utc_now_iso()
    sources = [
        {
            "dataset_key": "sangraha_verified_telugu",
            "name": "Sangraha Verified Telugu",
            "description": "Pure Telugu document review queue from Sangraha verified Telugu.",
            "queue_name": "pure_telugu",
            "sample_type": "document",
            "hf_dataset": "ai4bharat/sangraha",
            "hf_config": None,
            "hf_split": "train",
            "hf_data_dir": "verified/tel",
            "field_mapping": {
                "source_record_id": "doc_id",
                "text": "text",
                "doc_type": "type",
                "title": None,
            },
            "filters": {"min_chars": 30},
            "enabled": True,
            "created_at": now,
            "updated_at": now,
        },
        {
            "dataset_key": "samanantar_telugu_english",
            "name": "Samanantar Telugu-English",
            "description": "Telugu-English translation review queue from AI4Bharat Samanantar.",
            "queue_name": "translation_tel_eng",
            "sample_type": "translation_pair",
            "hf_dataset": "ai4bharat/samanantar",
            "hf_config": "te",
            "hf_split": "train",
            "hf_data_dir": None,
            "field_mapping": {
                "source_record_id": "idx",
                "source_text": "src",
                "target_text": "tgt",
                "source_lang": "en",
                "target_lang": "te",
            },
            "filters": {"min_chars": 10},
            "enabled": True,
            "created_at": now,
            "updated_at": now,
        },
        {
            "dataset_key": "dakshina_telugu_transliteration",
            "name": "Dakshina Telugu Transliteration",
            "description": "Telugu native-script to Latin transliteration review queue.",
            "queue_name": "transliteration",
            "sample_type": "transliteration_pair",
            "hf_dataset": "ramgopal-reddy/Telugu_to_English_Transliteration_Dakshina",
            "hf_config": None,
            "hf_split": "train",
            "hf_data_dir": None,
            "field_mapping": {
                "source_record_id": "unique_identifier",
                "native_text": "native_word",
                "latin_text": "english_word",
                "source_name": "source",
            },
            "filters": {"min_chars": 1},
            "enabled": True,
            "created_at": now,
            "updated_at": now,
        },
    ]

    custom_repo = os.getenv("CUSTOM_PURE_TELUGU_HF_DATASET", "").strip()
    if custom_repo:
        custom_name = (
            os.getenv("CUSTOM_PURE_TELUGU_SOURCE_NAME", "Rachana Combined Telugu Content").strip()
            or "Rachana Combined Telugu Content"
        )
        sources.append(
            {
                "dataset_key": "rachana_combined_telugu_content",
                "name": custom_name,
                "description": "Custom pure Telugu document review queue from combined cleaned content.",
                "queue_name": "pure_telugu",
                "sample_type": "document",
                "hf_dataset": custom_repo,
                "hf_config": None,
                "hf_split": "train",
                "hf_data_dir": None,
                "field_mapping": {
                    "source_record_id": None,
                    "text": "Content",
                    "doc_type": None,
                    "title": None,
                    "url": None,
                },
                "filters": {"min_chars": 30},
                "enabled": True,
                "created_at": now,
                "updated_at": now,
            }
        )

    return sources


def seed_default_sources(db: Database) -> int:
    inserted = 0
    for source in default_source_configs():
        result = db["source_configs"].update_one(
            {"dataset_key": source["dataset_key"]},
            {
                "$set": {
                    "name": source["name"],
                    "description": source["description"],
                    "queue_name": source["queue_name"],
                    "sample_type": source["sample_type"],
                    "hf_dataset": source["hf_dataset"],
                    "hf_config": source["hf_config"],
                    "hf_split": source["hf_split"],
                    "hf_data_dir": source["hf_data_dir"],
                    "field_mapping": source["field_mapping"],
                    "filters": source["filters"],
                    "enabled": source["enabled"],
                    "updated_at": utc_now_iso(),
                },
                "$setOnInsert": {"created_at": source["created_at"]},
            },
            upsert=True,
        )
        if result.upserted_id is not None:
            inserted += 1
    db["source_configs"].delete_many({"dataset_key": "wikipedia_telugu"})
    return inserted


def _serialize_source(source: dict[str, Any]) -> dict[str, Any]:
    serialized = dict(source)
    if isinstance(serialized.get("_id"), ObjectId):
        serialized["_id"] = str(serialized["_id"])
    return serialized


def list_sources(db: Database) -> list[dict[str, Any]]:
    return [_serialize_source(source) for source in db["source_configs"].find().sort("dataset_key", 1)]


def get_source(db: Database, dataset_key: str) -> dict[str, Any] | None:
    return db["source_configs"].find_one({"dataset_key": dataset_key})