openai-gabriel 1.0.1__py3-none-any.whl

gabriel/tasks/extract.py

@@ -0,0 +1,455 @@
+from __future__ import annotations
+
+import hashlib
+import asyncio
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from collections import defaultdict
+from typing import Any, DefaultDict, Dict, List, Optional, Set, Tuple
+
+import pandas as pd
+
+from ..core.prompt_template import PromptTemplate, resolve_template
+from ..utils.openai_utils import get_all_responses
+from ..utils import (
+    safest_json,
+    load_image_inputs,
+    load_audio_inputs,
+    load_pdf_inputs,
+    warn_if_modality_mismatch,
+)
+from ..utils.logging import announce_prompt_rendering
+from ._attribute_utils import load_persisted_attributes
+
+
+@dataclass
+class ExtractConfig:
+    attributes: Dict[str, str]
+    save_dir: str = "extraction"
+    file_name: str = "extraction.csv"
+    model: str = "gpt-5-mini"
+    n_parallels: int = 650
+    n_runs: int = 1
+    use_dummy: bool = False
+    max_timeout: Optional[float] = None
+    additional_instructions: Optional[str] = None
+    modality: str = "entity"
+    n_attributes_per_run: int = 8
+    reasoning_effort: Optional[str] = None
+    reasoning_summary: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        if self.additional_instructions is not None:
+            cleaned = str(self.additional_instructions).strip()
+            self.additional_instructions = cleaned or None
+
+
+class Extract:
+    """Extract attributes from passages using an LLM."""
+
+    def __init__(
+        self,
+        cfg: ExtractConfig,
+        template: Optional[PromptTemplate] = None,
+        template_path: Optional[str] = None,
+    ) -> None:
+        expanded = Path(os.path.expandvars(os.path.expanduser(cfg.save_dir)))
+        expanded.mkdir(parents=True, exist_ok=True)
+        cfg.save_dir = str(expanded)
+        self.cfg = cfg
+        self.template = resolve_template(
+            template=template,
+            template_path=template_path,
+            reference_filename="extraction_prompt.jinja2",
+        )
+
+    async def _parse(
+        self, raw: Any, attrs: List[str]
+    ) -> List[Tuple[Optional[str], Dict[str, str]]]:
+        obj = await safest_json(raw)
+        attr_names = list(attrs)
+
+        def _default_attr_map() -> Dict[str, str]:
+            return {attr: "unknown" for attr in attr_names}
+
+        def _clean_name(name: Any) -> Optional[str]:
+            if isinstance(name, str):
+                cleaned = name.strip()
+                return cleaned or None
+            if name is None:
+                return None
+            text = str(name).strip()
+            return text or None
+
+        def _build_entry(
+            entity_name: Optional[str], payload: Optional[Dict[str, Any]]
+        ) -> Tuple[Optional[str], Dict[str, str]]:
+            values = _default_attr_map()
+            if isinstance(payload, dict):
+                for attr in attr_names:
+                    val = payload.get(attr)
+                    values[attr] = str(val) if val is not None else "unknown"
+            return (_clean_name(entity_name), values)
+
+        entries: List[Tuple[Optional[str], Dict[str, str]]] = []
+
+        if isinstance(obj, dict):
+            attr_keys = set(attr_names)
+            nested_candidates = [
+                (key, val)
+                for key, val in obj.items()
+                if isinstance(val, dict)
+                and (not attr_keys or key not in attr_keys)
+            ]
+            if nested_candidates:
+                for name, payload in nested_candidates:
+                    entries.append(_build_entry(name, payload))
+                if entries:
+                    return entries
+            entries.append(_build_entry(None, obj))
+            return entries
+
+        if isinstance(obj, list):
+            for item in obj:
+                if not isinstance(item, dict):
+                    continue
+                attr_payload: Optional[Dict[str, Any]] = None
+                if isinstance(item.get("attributes"), dict):
+                    attr_payload = item.get("attributes")  # type: ignore[assignment]
+                else:
+                    attr_payload = item
+                name = item.get("entity_name") or item.get("entity") or item.get("name")
+                entries.append(_build_entry(name, attr_payload))
+            if entries:
+                return entries
+
+        return [
+            _build_entry(None, None),
+        ]
+
+    async def run(
+        self,
+        df: pd.DataFrame,
+        column_name: str,
+        *,
+        reset_files: bool = False,
+        types: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> pd.DataFrame:
+        df_proc = df.reset_index(drop=True).copy()
+        input_columns = list(df_proc.columns)
+        base_name = os.path.splitext(self.cfg.file_name)[0]
+        self.cfg.attributes = load_persisted_attributes(
+            save_dir=self.cfg.save_dir,
+            incoming=self.cfg.attributes,
+            reset_files=reset_files,
+            task_name="Extract",
+            item_name="attributes",
+            legacy_filename=f"{base_name}_attrs.json",
+        )
+        values = df_proc[column_name].tolist()
+        texts = [str(v) for v in values]
+
+        warn_if_modality_mismatch(values, self.cfg.modality, column_name=column_name)
+
+        base_ids: List[str] = []
+        id_to_rows: DefaultDict[str, List[int]] = defaultdict(list)
+        id_to_val: Dict[str, Any] = {}
+        prompt_texts: Dict[str, str] = {}
+        row_ids: List[str] = []
+
+        for row, (passage, orig) in enumerate(zip(texts, values)):
+            sha8 = hashlib.sha1(passage.encode()).hexdigest()[:8]
+            row_ids.append(sha8)
+            id_to_rows[sha8].append(row)
+            if len(id_to_rows[sha8]) > 1:
+                continue
+            id_to_val[sha8] = orig
+            prompt_texts[sha8] = passage if self.cfg.modality in {"text", "entity", "web"} else ""
+            base_ids.append(sha8)
+
+        df_proc["_gid"] = row_ids
+
+        if not base_ids:
+            out_path = os.path.join(self.cfg.save_dir, f"{base_name}_cleaned.csv")
+            result = df_proc.drop(columns=["_gid"])
+            result["entity_name"] = pd.NA
+            for attr in self.cfg.attributes.keys():
+                result[attr] = pd.NA
+            result.to_csv(out_path, index=False)
+            return result
+
+        attr_items = list(self.cfg.attributes.items())
+        attr_count = len(attr_items)
+        if attr_count > self.cfg.n_attributes_per_run:
+            batches = (
+                attr_count + self.cfg.n_attributes_per_run - 1
+            ) // self.cfg.n_attributes_per_run
+            print(
+                f"[Extract] {attr_count} attributes provided. n_attributes_per_run={self.cfg.n_attributes_per_run}. "
+                f"Splitting into {batches} prompt batches. Increase n_attributes_per_run if you want all attributes "
+                "to be processed in the same prompt."
+            )
+        attr_batches: List[Dict[str, str]] = [
+            dict(attr_items[i : i + self.cfg.n_attributes_per_run])
+            for i in range(0, len(attr_items), self.cfg.n_attributes_per_run)
+        ]
+
+        prompts: List[str] = []
+        ids: List[str] = []
+        announce_prompt_rendering("Extract", len(base_ids) * len(attr_batches))
+        for batch_idx, batch_attrs in enumerate(attr_batches):
+            for ident in base_ids:
+                prompts.append(
+                    self.template.render(
+                        text=prompt_texts[ident],
+                        attributes=batch_attrs,
+                        additional_instructions=self.cfg.additional_instructions or "",
+                        modality=self.cfg.modality,
+                    )
+                )
+                ids.append(f"{ident}_batch{batch_idx}")
+
+        prompt_images: Optional[Dict[str, List[str]]] = None
+        prompt_audio: Optional[Dict[str, List[Dict[str, str]]]] = None
+        prompt_pdfs: Optional[Dict[str, List[Dict[str, str]]]] = None
+        if self.cfg.modality == "image":
+            tmp: Dict[str, List[str]] = {}
+            for ident, rows in id_to_rows.items():
+                imgs = load_image_inputs(values[rows[0]])
+                if imgs:
+                    for batch_idx in range(len(attr_batches)):
+                        tmp[f"{ident}_batch{batch_idx}"] = imgs
+            prompt_images = tmp or None
+        elif self.cfg.modality == "audio":
+            tmp_a: Dict[str, List[Dict[str, str]]] = {}
+            for ident, rows in id_to_rows.items():
+                auds = load_audio_inputs(values[rows[0]])
+                if auds:
+                    for batch_idx in range(len(attr_batches)):
+                        tmp_a[f"{ident}_batch{batch_idx}"] = auds
+            prompt_audio = tmp_a or None
+        elif self.cfg.modality == "pdf":
+            tmp_p: Dict[str, List[Dict[str, str]]] = {}
+            for ident, rows in id_to_rows.items():
+                pdfs = load_pdf_inputs(values[rows[0]])
+                if pdfs:
+                    for batch_idx in range(len(attr_batches)):
+                        tmp_p[f"{ident}_batch{batch_idx}"] = pdfs
+            prompt_pdfs = tmp_p or None
+
+        csv_path = os.path.join(self.cfg.save_dir, f"{base_name}_raw_responses.csv")
+
+        kwargs.setdefault("web_search", self.cfg.modality == "web")
+
+        if not isinstance(self.cfg.n_runs, int) or self.cfg.n_runs < 1:
+            raise ValueError("n_runs must be an integer >= 1")
+
+        existing_ids: Set[str] = set()
+        if not reset_files and os.path.exists(csv_path):
+            try:
+                existing_df = pd.read_csv(csv_path, usecols=["Identifier"])
+                existing_ids = set(existing_df["Identifier"].astype(str))
+            except Exception:
+                existing_ids = set()
+
+        run_identifier_lists: List[List[str]] = []
+        for run_idx in range(1, self.cfg.n_runs + 1):
+            run_ids: List[str] = []
+            for ident in ids:
+                if run_idx == 1:
+                    legacy_ident = f"{ident}_run1"
+                    run_ids.append(legacy_ident if legacy_ident in existing_ids else ident)
+                else:
+                    run_ids.append(f"{ident}_run{run_idx}")
+            run_identifier_lists.append(run_ids)
+
+        prompts_all: List[str] = []
+        ids_all: List[str] = []
+        for run_ids in run_identifier_lists:
+            prompts_all.extend(prompts)
+            ids_all.extend(run_ids)
+
+        prompt_images_all: Optional[Dict[str, List[str]]] = None
+        if prompt_images:
+            prompt_images_all = {}
+            for run_ids in run_identifier_lists:
+                for base_ident, run_ident in zip(ids, run_ids):
+                    imgs = prompt_images.get(base_ident)
+                    if imgs:
+                        prompt_images_all[run_ident] = imgs
+        prompt_audio_all: Optional[Dict[str, List[Dict[str, str]]]] = None
+        if prompt_audio:
+            prompt_audio_all = {}
+            for run_ids in run_identifier_lists:
+                for base_ident, run_ident in zip(ids, run_ids):
+                    auds = prompt_audio.get(base_ident)
+                    if auds:
+                        prompt_audio_all[run_ident] = auds
+        prompt_pdfs_all: Optional[Dict[str, List[Dict[str, str]]]] = None
+        if prompt_pdfs:
+            prompt_pdfs_all = {}
+            for run_ids in run_identifier_lists:
+                for base_ident, run_ident in zip(ids, run_ids):
+                    pdfs = prompt_pdfs.get(base_ident)
+                    if pdfs:
+                        prompt_pdfs_all[run_ident] = pdfs
+
+        df_resp_all = await get_all_responses(
+            prompts=prompts_all,
+            identifiers=ids_all,
+            prompt_images=prompt_images_all,
+            prompt_audio=prompt_audio_all,
+            prompt_pdfs=prompt_pdfs_all,
+            n_parallels=self.cfg.n_parallels,
+            model=self.cfg.model,
+            save_path=csv_path,
+            use_dummy=self.cfg.use_dummy,
+            max_timeout=self.cfg.max_timeout,
+            json_mode=self.cfg.modality != "audio",
+            reset_files=reset_files,
+            reasoning_effort=self.cfg.reasoning_effort,
+            reasoning_summary=self.cfg.reasoning_summary,
+            **kwargs,
+        )
+        if not isinstance(df_resp_all, pd.DataFrame):
+            raise RuntimeError("get_all_responses returned no DataFrame")
+
+        df_resps = []
+        for run_idx, run_ids in enumerate(run_identifier_lists, start=1):
+            suffix = f"_run{run_idx}"
+            sub = df_resp_all[df_resp_all.Identifier.isin(run_ids)].copy()
+            sub.Identifier = sub.Identifier.str.replace(suffix + "$", "", regex=True)
+            df_resps.append(sub)
+
+        full_records: List[Dict[str, Any]] = []
+        base_attrs = list(self.cfg.attributes.keys())
+        for run_idx, df_resp in enumerate(df_resps, start=1):
+            id_to_entity_vals: Dict[str, Dict[Optional[str], Dict[str, str]]] = {
+                ident: {} for ident in base_ids
+            }
+            for ident_batch, raw in zip(df_resp.Identifier, df_resp.Response):
+                if "_batch" not in ident_batch:
+                    continue
+                base_ident, batch_part = ident_batch.rsplit("_batch", 1)
+                batch_idx = int(batch_part)
+                attrs = list(attr_batches[batch_idx].keys())
+                parsed_entities = await self._parse(raw, attrs)
+                entity_store = id_to_entity_vals.setdefault(base_ident, {})
+                for entity_name, entity_attrs in parsed_entities:
+                    key = entity_name if entity_name is not None else None
+                    if key not in entity_store:
+                        entity_store[key] = {attr: "unknown" for attr in base_attrs}
+                    for attr in attrs:
+                        entity_store[key][attr] = entity_attrs.get(attr, "unknown")
+            for ident in base_ids:
+                entity_map = id_to_entity_vals.get(ident) or {
+                    None: {attr: "unknown" for attr in base_attrs}
+                }
+                for entity_name, attr_values in entity_map.items():
+                    rec: Dict[str, Any] = {
+                        "id": ident,
+                        "entity_name": entity_name,
+                        "text": id_to_val[ident],
+                        "run": run_idx,
+                    }
+                    rec.update({attr: attr_values.get(attr, "unknown") for attr in base_attrs})
+                    full_records.append(rec)
+
+        full_df = pd.DataFrame(full_records).set_index(["id", "entity_name", "run"])
+        if self.cfg.n_runs > 1:
+            disagg_path = os.path.join(
+                self.cfg.save_dir, f"{base_name}_full_disaggregated.csv"
+            )
+            full_df.to_csv(disagg_path, index_label=["id", "entity_name", "run"])
+
+        def _pick_first(s: pd.Series) -> str:
+            for val in s.dropna():
+                if str(val).strip().lower() != "unknown":
+                    return str(val)
+            return "unknown"
+
+        entity_index = full_df.index.droplevel("run").unique()
+        if base_attrs:
+            agg_series = {
+                attr: full_df[attr]
+                .groupby(level=["id", "entity_name"], sort=False)
+                .apply(_pick_first)
+                for attr in base_attrs
+            }
+            agg_df = pd.DataFrame(agg_series)
+        else:
+            agg_df = pd.DataFrame(index=entity_index)
+        agg_df = agg_df.reindex(entity_index)
+
+        unknown_counts = {attr: (agg_df[attr] == "unknown").sum() for attr in base_attrs}
+
+        out_path = os.path.join(self.cfg.save_dir, f"{base_name}_cleaned.csv")
+        agg_reset = agg_df.reset_index()
+        result = df_proc.merge(agg_reset, left_on="_gid", right_on="id", how="left")
+        drop_cols = [col for col in ("_gid", "id") if col in result.columns]
+        if drop_cols:
+            result = result.drop(columns=drop_cols)
+
+        original_cols = [col for col in df_proc.columns if col != "_gid"]
+        final_order: List[str] = []
+        for col in original_cols:
+            if col in result.columns:
+                final_order.append(col)
+        if "entity_name" in result.columns and "entity_name" not in final_order:
+            final_order.append("entity_name")
+        for attr in base_attrs:
+            if attr in result.columns:
+                final_order.append(attr)
+        remaining = [col for col in result.columns if col not in final_order]
+        if remaining:
+            final_order.extend(remaining)
+        result = result[final_order]
+
+        result.to_csv(out_path, index=False)
+
+        result = result.replace("unknown", pd.NA)
+
+        duplicate_rows = len(result) - len(df_proc)
+
+        if types:
+            coerced = result.copy()
+            fail_logs: Dict[str, int] = {}
+            for col, typ in types.items():
+                if col not in coerced:
+                    continue
+                orig = coerced[col]
+                non_null = orig.notna()
+                target = str(typ).lower()
+                if target in {"datetime", "date"}:
+                    conv = pd.to_datetime(orig, errors="coerce")
+                else:
+                    conv = pd.to_numeric(orig, errors="coerce")
+                    if target in {"int", "int64"}:
+                        conv = conv.round().astype("Int64")
+                coerced[col] = conv
+                fail_logs[col] = int((non_null & conv.isna()).sum())
+            coerced_path = os.path.join(self.cfg.save_dir, f"{base_name}_cleaned_coerced.csv")
+            coerced.to_csv(coerced_path, index=False)
+            for col, n_fail in fail_logs.items():
+                print(f"[Extract] Failed to coerce {n_fail} values in column '{col}'.")
+            result = coerced
+
+        total = len(agg_df)
+        print("\n=== Extraction coverage ===")
+        for attr in base_attrs:
+            known = total - unknown_counts[attr]
+            print(f"{attr:<55s}: {known:5d} extracted, {unknown_counts[attr]:5d} unknown")
+        print("============================\n")
+
+        if duplicate_rows > 0:
+            subset_hint = ", ".join(f"'{col}'" for col in input_columns)
+            print(
+                "[Extract] Multiple entity names were returned for at least one input.\n"
+                f"          {duplicate_rows} additional row(s) were created so the cleaned file now has {len(result)} rows versus {len(df_proc)} inputs.\n"
+                "          If you only need one row per original input, deduplicate on your source columns\n"
+                f"          (e.g. `result = result.drop_duplicates(subset=[{subset_hint}], keep='first')`).\n"
+            )
+
+        return result

gabriel/tasks/filter.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+import os
+import random
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Set
+
+import pandas as pd
+
+from ..core.prompt_template import PromptTemplate, resolve_template
+from ..utils.openai_utils import get_all_responses
+from ..utils.logging import announce_prompt_rendering
+from ..utils import safest_json
+
+
+@dataclass
+class FilterConfig:
+    """Configuration for :class:`Filter`."""
+
+    condition: str
+    save_dir: str
+    file_name: str = "filter_responses.csv"
+    model: str = "gpt-5-nano"
+    n_parallels: int = 650
+    entities_per_call: int = 150
+    shuffle: bool = True
+    random_seed: int = 42
+    n_runs: int = 1
+    threshold: float = 0.5
+    additional_instructions: Optional[str] = None
+    use_dummy: bool = False
+    max_timeout: Optional[float] = None
+    fix_json_with_llm: bool = False
+    json_fix_timeout: Optional[float] = 60.0
+
+    def __post_init__(self) -> None:
+        if self.additional_instructions is not None:
+            cleaned = str(self.additional_instructions).strip()
+            self.additional_instructions = cleaned or None
+
+
+class Filter:
+    """Filter entities in a DataFrame column based on a condition using an LLM."""
+
+    def __init__(
+        self,
+        cfg: FilterConfig,
+        template: Optional[PromptTemplate] = None,
+        template_path: Optional[str] = None,
+    ) -> None:
+        expanded = Path(os.path.expandvars(os.path.expanduser(cfg.save_dir)))
+        expanded.mkdir(parents=True, exist_ok=True)
+        cfg.save_dir = str(expanded)
+        self.cfg = cfg
+        self.template = resolve_template(
+            template=template,
+            template_path=template_path,
+            reference_filename="filter_prompt.jinja2",
+        )
+
+    # ------------------------------------------------------------------
+    async def run(
+        self,
+        df: pd.DataFrame,
+        column_name: str,
+        *,
+        reset_files: bool = False,
+        **kwargs: Any,
+    ) -> pd.DataFrame:
+        df_proc = df.reset_index(drop=True).copy()
+        raw_entities = [str(x) for x in df_proc[column_name].dropna()]
+
+        # unique while preserving order
+        seen: Set[str] = set()
+        entities: List[str] = []
+        for ent in raw_entities:
+            if ent not in seen:
+                seen.add(ent)
+                entities.append(ent)
+
+        prompts: List[str] = []
+        identifiers: List[str] = []
+        total_chunks = 0
+        for run in range(self.cfg.n_runs):
+            ents = list(entities)
+            if self.cfg.shuffle:
+                rnd = random.Random(self.cfg.random_seed + run)
+                rnd.shuffle(ents)
+            chunks = [
+                ents[i : i + self.cfg.entities_per_call]
+                for i in range(0, len(ents), self.cfg.entities_per_call)
+            ]
+            total_chunks += len(chunks)
+            for idx, chunk in enumerate(chunks):
+                prompts.append(
+                    self.template.render(
+                        condition=self.cfg.condition,
+                        entities=chunk,
+                        additional_instructions=self.cfg.additional_instructions,
+                    )
+                )
+                identifiers.append(f"filter_{run:03d}_{idx:05d}")
+
+        announce_prompt_rendering("Filter", total_chunks)
+
+        save_path = os.path.join(self.cfg.save_dir, self.cfg.file_name)
+        if prompts:
+            resp_df = await get_all_responses(
+                prompts=prompts,
+                identifiers=identifiers,
+                n_parallels=self.cfg.n_parallels,
+                model=self.cfg.model,
+                save_path=save_path,
+                use_dummy=self.cfg.use_dummy,
+                max_timeout=self.cfg.max_timeout,
+                json_mode=True,
+                reset_files=reset_files,
+                **kwargs,
+            )
+        else:
+            resp_df = pd.DataFrame(columns=["Identifier", "Response"])
+
+        resp_map: Dict[str, Any] = dict(
+            zip(resp_df.get("Identifier", []), resp_df.get("Response", []))
+        )
+
+        meets_by_run: List[Set[str]] = [set() for _ in range(self.cfg.n_runs)]
+        for ident, raw in resp_map.items():
+            parts = ident.split("_")
+            if len(parts) < 3:
+                continue
+            try:
+                run_idx = int(parts[1])
+            except ValueError:
+                continue
+            parsed = await safest_json(
+                raw,
+                model=self.cfg.model if self.cfg.fix_json_with_llm else None,
+                use_llm_fallback=self.cfg.fix_json_with_llm,
+                llm_timeout=self.cfg.json_fix_timeout,
+            )
+            ent_list: Optional[List[str]] = None
+            if isinstance(parsed, dict):
+                val = parsed.get("entities meeting condition") or parsed.get(
+                    "entities_meeting_condition"
+                )
+                if isinstance(val, list):
+                    ent_list = [str(v) for v in val if isinstance(v, str)]
+            elif isinstance(parsed, list):
+                ent_list = [str(v) for v in parsed if isinstance(v, str)]
+            if ent_list and 0 <= run_idx < self.cfg.n_runs:
+                for ent in ent_list:
+                    meets_by_run[run_idx].add(ent.strip())
+
+        run_cols: List[str] = []
+        for run_idx in range(self.cfg.n_runs):
+            meets_norm = {m.lower() for m in meets_by_run[run_idx]}
+            col = f"meets_condition_run_{run_idx + 1}"
+            run_cols.append(col)
+            df_proc[col] = [
+                str(v).lower() in meets_norm if not pd.isna(v) else False
+                for v in df_proc[column_name]
+            ]
+
+        df_proc["meets_condition"] = (
+            df_proc[run_cols].sum(axis=1) / self.cfg.n_runs >= self.cfg.threshold
+        )
+        return df_proc