schematize 0.1.0__py3-none-any.whl

schematize/__init__.py

@@ -0,0 +1,32 @@
+"""schematize — agentic extraction-schema generation."""
+
+__version__ = "0.1.0"
+
+from schematize.agents.agent_state import AgentState, agent_state_to_json
+from schematize.agents.schema_generator import SchemaGenerator, SchemaGeneratorPrompts
+from schematize.eval.evaluator import SchemaEvaluator
+from schematize.retrieval.base import DocumentRetriever
+from schematize.schema.model import (
+    DynamicModelFactory,
+    FieldDef,
+    FieldType,
+    NamedFieldDef,
+    SchemaFields,
+)
+from schematize.utils.load import load_prompts
+
+__all__ = [
+    "__version__",
+    "SchemaGenerator",
+    "SchemaGeneratorPrompts",
+    "SchemaEvaluator",
+    "DocumentRetriever",
+    "SchemaFields",
+    "FieldDef",
+    "NamedFieldDef",
+    "FieldType",
+    "DynamicModelFactory",
+    "AgentState",
+    "agent_state_to_json",
+    "load_prompts",
+]

schematize/_scripts/evaluate_annotator_agreement.py

@@ -0,0 +1,100 @@
+import asyncio
+import json
+import logging
+from itertools import permutations
+from pathlib import Path
+from typing import Any
+
+import hydra
+import yaml
+from langchain_openai import ChatOpenAI
+from omegaconf import DictConfig
+from tqdm.asyncio import tqdm as atqdm
+
+from schematize.eval.evaluator import AnnotatorAgreementEvaluator
+from schematize.settings import PROMPTS_PATH
+from schematize.utils.langchain import setup_langchain_llm_cache
+
+_CONFIG_PATH = str(Path(__file__).parent / "../../../config")
+
+
+@hydra.main(version_base=None, config_path=_CONFIG_PATH, config_name="evaluate_agreement")
+def main(cfg: DictConfig) -> None:
+    logging.getLogger("httpx").setLevel(logging.WARNING)
+
+    if cfg.llm_cache:
+        setup_langchain_llm_cache()
+
+    llm = ChatOpenAI(
+        model=cfg.model_name,
+        base_url=cfg.api_url,
+        api_key=cfg.api_key,
+        temperature=cfg.llm.temperature,
+        max_tokens=cfg.llm.max_tokens,
+        use_responses_api=False,
+        **cfg.llm.model_kwargs,
+    )
+
+    with open(PROMPTS_PATH / "eval" / "annotator_agreement.yaml", "r") as f:
+        prompts = yaml.safe_load(f)
+
+    evaluator = AnnotatorAgreementEvaluator(llm, prompts["annotator_agreement_prompt"])
+
+    case_dir = Path(cfg.case_dir)
+    if not case_dir.exists():
+        raise ValueError(f"Case directory not found at {case_dir}")
+
+    expert_files = sorted(case_dir.glob("expert_*.yaml"))
+    assert len(expert_files) >= 2, f"Need at least 2 expert files in {case_dir}"
+
+    results = asyncio.run(_evaluate_all_pairs(evaluator, case_dir, expert_files, cfg))
+
+    output_path = Path(cfg.output) / "agreement.json"
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, "w") as f:
+        json.dump(results, f, indent=2)
+    print(f"Annotator agreement results saved to {output_path}")
+
+
+async def _evaluate_all_pairs(
+    evaluator: AnnotatorAgreementEvaluator,
+    case_dir: Path,
+    expert_files: list[Path],
+    cfg: DictConfig,
+) -> dict[str, Any]:
+    semaphore = asyncio.Semaphore(cfg.get("max_concurrent", 10))
+
+    async def eval_pair(target_file: Path, reference_file: Path) -> dict:
+        target_questions = _load_expert_questions(target_file)
+        reference_questions = _load_expert_questions(reference_file)
+        async with semaphore:
+            evaluation = await evaluator.aevaluate_agreement(target_questions, reference_questions)
+        return {
+            "target": target_file.stem,
+            "reference": reference_file.stem,
+            "total_questions": evaluation.total_questions,
+            "covered_questions": evaluation.covered_questions,
+            "high_confidence": evaluation.high_confidence_coverage,
+            "medium_confidence": evaluation.medium_confidence_coverage,
+            "low_confidence": evaluation.low_confidence_coverage,
+        }
+
+    pairs = await atqdm.gather(
+        *[eval_pair(target, reference) for target, reference in permutations(expert_files, 2)],
+        desc="Expert pairs",
+    )
+
+    return {
+        "case_dir": str(case_dir),
+        "num_experts": len(expert_files),
+        "pairs": pairs,
+    }
+
+
+def _load_expert_questions(expert_path: Path) -> list[str]:
+    with open(expert_path, "r") as f:
+        return yaml.safe_load(f).get("questions", [])
+
+
+if __name__ == "__main__":
+    main()

schematize/_scripts/evaluate_schema.py

@@ -0,0 +1,169 @@
+import asyncio
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+import hydra
+import yaml
+from langchain_openai import ChatOpenAI
+from omegaconf import DictConfig
+from tqdm.asyncio import tqdm as atqdm
+
+from schematize.eval.evaluator import SchemaEvaluator
+from schematize.settings import PROMPTS_PATH
+from schematize.utils.langchain import setup_langchain_llm_cache
+
+_CONFIG_PATH = str(Path(__file__).parent / "../../../config")
+
+
+@hydra.main(version_base=None, config_path=_CONFIG_PATH, config_name="evaluate")
+def main(cfg: DictConfig) -> None:
+    logging.getLogger("httpx").setLevel(logging.WARNING)
+
+    if cfg.llm_cache:
+        setup_langchain_llm_cache()
+
+    llm = ChatOpenAI(
+        model=cfg.model_name,
+        base_url=cfg.api_url,
+        api_key=cfg.api_key,
+        temperature=cfg.llm.temperature,
+        max_tokens=cfg.llm.max_tokens,
+        use_responses_api=False,
+        **cfg.llm.model_kwargs,
+    )
+
+    with open(PROMPTS_PATH / "eval" / "schema_evaluator.yaml", "r") as f:
+        prompts = yaml.safe_load(f)
+
+    evaluator = SchemaEvaluator(llm, prompts["schema_evaluator_prompt"])
+
+    state_paths = _get_state_paths(Path(cfg.state_dir))
+    case_dir = Path(cfg.case_dir)
+
+    if not case_dir.exists():
+        raise ValueError(f"Case directory not found at {case_dir}")
+
+    expert_files = sorted(case_dir.glob("expert_*.yaml"))
+    assert len(expert_files) > 0, f"No expert files found in {case_dir}"
+
+    asyncio.run(_run_all(evaluator, state_paths, case_dir, expert_files, cfg))
+
+
+async def _run_all(
+    evaluator: SchemaEvaluator,
+    state_paths: list[Path],
+    case_dir: Path,
+    expert_files: list[Path],
+    cfg: DictConfig,
+) -> None:
+    semaphore = asyncio.Semaphore(cfg.get("max_concurrent", 10))
+    await atqdm.gather(
+        *[
+            _evaluate_and_save(evaluator, p, case_dir, expert_files, cfg, semaphore)
+            for p in state_paths
+        ],
+        desc="States",
+    )
+
+
+async def _evaluate_and_save(
+    evaluator: SchemaEvaluator,
+    state_path: Path,
+    case_dir: Path,
+    expert_files: list[Path],
+    cfg: DictConfig,
+    semaphore: asyncio.Semaphore,
+) -> None:
+    state = _load_state(state_path)
+    if cfg.final_only:
+        schema_history = [state["current_schema"]] if state.get("current_schema") else []
+    else:
+        schema_history = state.get("schema_history", [])
+        if state.get("current_schema"):
+            schema_history.append(state["current_schema"])
+
+    results = await _evaluate(evaluator, state_path, case_dir, schema_history, expert_files, semaphore)
+
+    output_path = (
+        Path(cfg.output) / state_path.relative_to(cfg.state_dir).parent / "evaluation.json"
+    )
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, "w") as f:
+        json.dump(results, f, indent=2)
+    print(f"Evaluation results saved to {output_path}")
+
+
+async def _evaluate(
+    evaluator: SchemaEvaluator,
+    state_path: Path,
+    case_dir: Path,
+    schema_history: list[dict],
+    expert_files: list[Path],
+    semaphore: asyncio.Semaphore,
+) -> dict[str, Any]:
+    schema_results: dict[int, dict] = {
+        i: {"schema_index": i, "num_fields": len(schema), "experts": []}
+        for i, schema in enumerate(schema_history)
+    }
+
+    async def eval_one(schema_idx: int, schema: dict, expert_file: Path) -> tuple:
+        questions = _load_expert_questions(expert_file)
+        async with semaphore:
+            evaluation = await evaluator.aevaluate_schema(schema, questions)
+        return schema_idx, expert_file.stem, evaluation
+
+    tasks = [
+        eval_one(schema_idx, schema, expert_file)
+        for schema_idx, schema in enumerate(schema_history)
+        for expert_file in expert_files
+    ]
+
+    for coro in asyncio.as_completed(tasks):
+        schema_idx, expert_stem, evaluation = await coro
+        schema_results[schema_idx]["experts"].append(
+            {
+                "expert": expert_stem,
+                "total_questions": evaluation.total_questions,
+                "covered_questions": evaluation.covered_questions,
+                "high_confidence": evaluation.high_confidence_coverage,
+                "medium_confidence": evaluation.medium_confidence_coverage,
+                "low_confidence": evaluation.low_confidence_coverage,
+            }
+        )
+
+    return {
+        "state_path": str(state_path),
+        "case_dir": str(case_dir),
+        "num_schemas": len(schema_history),
+        "num_experts": len(expert_files),
+        "evaluations": [schema_results[i] for i in range(len(schema_history))],
+    }
+
+
+def _load_state(state_path: Path) -> dict:
+    with open(state_path, "r") as f:
+        return json.load(f)
+
+
+def _load_expert_questions(expert_path: Path) -> list[str]:
+    with open(expert_path, "r") as f:
+        return yaml.safe_load(f).get("questions", [])
+
+
+def _get_state_paths(state_dir: Path) -> list[Path]:
+    direct = state_dir / "state.json"
+    if direct.exists():
+        return [direct]
+    result = [
+        p for p in state_dir.glob("*/state.json")
+        if not (p.parent / "evaluation.json").exists()
+    ]
+    if not result:
+        raise ValueError(f"No unevaluated state files found in {state_dir}")
+    return result
+
+
+if __name__ == "__main__":
+    main()

schematize/_scripts/schema_generator.py

@@ -0,0 +1,116 @@
+import os
+from pathlib import Path
+from typing import Annotated, Optional
+
+import typer
+from dotenv import load_dotenv
+from langchain_openai import ChatOpenAI
+
+from schematize.agents.agent_state import agent_state_to_json
+from schematize.agents.schema_generator import SchemaGenerator, SchemaGeneratorPrompts
+from schematize.utils.load import load_prompts
+
+app = typer.Typer()
+
+
+@app.command()
+def main(
+    model: Annotated[str, typer.Option()] = "gpt-4o-mini",
+    language: Annotated[str, typer.Option(help="pl or en")] = "en",
+    system_type: Annotated[str, typer.Option(help="tax or law")] = "tax",
+    temperature: Annotated[float, typer.Option()] = 0.2,
+    max_tokens: Annotated[int, typer.Option()] = 32_000,
+    reasoning_effort: Annotated[Optional[str], typer.Option(help="low | medium | high")] = None,
+    output: Annotated[Path, typer.Option()] = Path("state.json"),
+    api_url: Annotated[Optional[str], typer.Option()] = None,
+    api_key: Annotated[Optional[str], typer.Option()] = None,
+    retriever_type: Annotated[str, typer.Option(help="weaviate | mmlw | huggingface")] = "weaviate",
+    dataset: Annotated[Optional[str], typer.Option(help="HuggingFace dataset identifier")] = None,
+    text_column: Annotated[str, typer.Option(help="Dataset column to embed and search")] = "text",
+    max_documents: Annotated[Optional[int], typer.Option(help="Max rows to index (None = all)")] = None,
+    index_path: Annotated[Optional[str], typer.Option(help="Directory to cache the FAISS index")] = None,
+    collection_name: Annotated[Optional[str], typer.Option(help="Weaviate collection name")] = None,
+    target_vector: Annotated[Optional[str], typer.Option(help="Weaviate named vector")] = None,
+    wv_filter: Annotated[Optional[list[str]], typer.Option(help="Weaviate filters as key=value")] = None,
+    max_refinement_rounds: Annotated[int, typer.Option()] = 3,
+    min_refinement_rounds: Annotated[int, typer.Option()] = 2,
+    max_data_refinement_rounds: Annotated[int, typer.Option()] = 3,
+    min_data_refinement_rounds: Annotated[int, typer.Option()] = 2,
+    data_assessment_top_k: Annotated[int, typer.Option()] = 50,
+    data_assessment_num_examples: Annotated[int, typer.Option()] = 3,
+    data_assessment_random_seed: Annotated[int, typer.Option()] = 17,
+    data_assessment_document_max_chars: Annotated[int, typer.Option()] = 32_000,
+) -> None:
+    load_dotenv(".env")
+
+    if language not in ("pl", "en"):
+        raise typer.BadParameter(f"language must be 'pl' or 'en', got '{language}'")
+    if system_type not in ("tax", "law"):
+        raise typer.BadParameter(f"system_type must be 'tax' or 'law', got '{system_type}'")
+
+    prompts = load_prompts(language, system_type)
+
+    if retriever_type == "weaviate":
+        try:
+            from schematize.retrieval.weaviate import WeaviateRetriever
+        except ImportError:
+            raise typer.Exit("Weaviate retriever requires: pip install schematize[weaviate]")
+        if not collection_name:
+            raise typer.BadParameter("--collection-name is required for weaviate retriever")
+        filters = dict(f.split("=", 1) for f in (wv_filter or []))
+        retriever = WeaviateRetriever(
+            collection_name=collection_name,
+            target_vector=target_vector,
+            filters=filters,
+        )
+    else:
+        try:
+            from schematize.retrieval.huggingface import (
+                HuggingFaceRetriever,
+                MMLWRobertaV2Retriever,
+            )
+        except ImportError:
+            raise typer.Exit("HuggingFace retriever requires: pip install schematize[huggingface]")
+        if not dataset:
+            raise typer.BadParameter("--dataset is required for huggingface/mmlw retriever")
+        kwargs = dict(dataset_name=dataset, text_column=text_column, max_documents=max_documents, index_path=index_path)
+        retriever = MMLWRobertaV2Retriever(**kwargs) if retriever_type == "mmlw" else HuggingFaceRetriever(**kwargs)
+
+    llm = ChatOpenAI(
+        model=model,
+        base_url=api_url or os.getenv("API_URL"),
+        api_key=api_key or os.getenv("API_KEY"),
+        temperature=temperature,
+        max_tokens=max_tokens,
+        use_responses_api=False,
+        reasoning_effort=reasoning_effort,
+    )
+
+    schema_system = SchemaGenerator(
+        llm,
+        retriever,
+        SchemaGeneratorPrompts(**prompts),
+        max_refinement_rounds=max_refinement_rounds,
+        min_refinement_rounds=min_refinement_rounds,
+        max_data_refinement_rounds=max_data_refinement_rounds,
+        min_data_refinement_rounds=min_data_refinement_rounds,
+        data_assessment_top_k=data_assessment_top_k,
+        data_assessment_num_examples=data_assessment_num_examples,
+        data_assessment_random_seed=data_assessment_random_seed,
+        data_assessment_document_max_chars=data_assessment_document_max_chars,
+        recursion_limit=100,
+    )
+
+    typer.echo("Describe the schema to generate:")
+    input_text = input()
+
+    final_state = schema_system.stream_graph_updates(input_text)
+
+    with output.open("w") as f:
+        f.write(agent_state_to_json(final_state))
+
+    typer.echo(f"State saved to {output}")
+
+
+if __name__ == "__main__":
+    app()

schematize/_scripts/schema_generator_ablation.py

@@ -0,0 +1,151 @@
+import logging
+import sys
+from pathlib import Path
+from typing import Any
+
+import hydra
+import yaml
+from dotenv import load_dotenv
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_openai import ChatOpenAI
+from loguru import logger
+from omegaconf import DictConfig
+
+from schematize.agents.agent_state import AgentState, agent_state_to_json
+from schematize.agents.schema_generator import SchemaGenerator, SchemaGeneratorPrompts
+from schematize.utils.langchain import setup_langchain_llm_cache
+from schematize.utils.load import load_prompts
+
+_CONFIG_PATH = str(Path(__file__).parent / "../../../config")
+
+load_dotenv()
+
+logger.remove()
+logger.add(sys.stderr, level="INFO")
+logging.getLogger("httpx").setLevel(logging.WARNING)
+
+
+@hydra.main(version_base=None, config_path=_CONFIG_PATH, config_name="ablation")
+def main(cfg: DictConfig) -> None:
+    assert cfg.case.language in ["pl", "en"], "Invalid language"
+    assert cfg.case.system_type in ["law", "tax"], "Invalid system type"
+
+    if cfg.llm_cache:
+        setup_langchain_llm_cache()
+
+    prompts = load_prompts(cfg.case.language, cfg.case.system_type)
+    cases = _load_cases(Path(cfg.cases_path))
+
+    case_name = cfg.case.name
+    assert case_name in cases, f"Case '{case_name}' not found. Available: {list(cases.keys())}"
+    case_data = cases[case_name]
+
+    retriever = _build_retriever(cfg.retriever)
+
+    llm = ChatOpenAI(
+        model=cfg.model_name,
+        base_url=cfg.api_url,
+        api_key=cfg.api_key,
+        temperature=cfg.llm.temperature,
+        max_tokens=cfg.llm.max_tokens,
+        use_responses_api=False,
+        **cfg.llm.model_kwargs,
+    )
+
+    sg_cfg = cfg.schema_generator
+    schema_system = SchemaGenerator(
+        llm,
+        retriever,
+        SchemaGeneratorPrompts(**prompts),
+        max_refinement_rounds=sg_cfg.max_refinement_rounds,
+        min_refinement_rounds=sg_cfg.min_refinement_rounds,
+        max_data_refinement_rounds=sg_cfg.max_data_refinement_rounds,
+        min_data_refinement_rounds=sg_cfg.min_data_refinement_rounds,
+        data_assessment_top_k=sg_cfg.data_assessment_top_k,
+        data_assessment_num_examples=sg_cfg.data_assessment_num_examples,
+        data_assessment_random_seed=sg_cfg.data_assessment_random_seed,
+        data_assessment_document_max_chars=sg_cfg.data_assessment_document_max_chars,
+        skip_problem_definition=sg_cfg.skip_problem_definition,
+        skip_refinement=sg_cfg.skip_refinement,
+        skip_data_grounded=sg_cfg.skip_data_grounded,
+        recursion_limit=100,
+    )
+
+    if not sg_cfg.skip_problem_definition:
+        if "problem_help" in case_data:
+            schema_system.problem_definer_helper = _mock_problem_helper(case_data["problem_help"])
+        if "user_feedback" in case_data:
+            schema_system.human_feedback = _mock_user_feedback(case_data["user_feedback"])
+    if "human_message" in case_data:
+        schema_system.human_message = _mock_human_message(case_data["human_message"])
+
+    schema_system.graph = schema_system.build_graph()
+
+    user_input = case_data.get("user_input", "Generate a schema for legal documents")
+    final_state = schema_system.stream_graph_updates(user_input)
+
+    output_path = Path(cfg.output) / "state.json"
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with output_path.open("w") as f:
+        f.write(agent_state_to_json(final_state))
+    logger.info("State saved to {}", output_path)
+
+    schema_path = Path(cfg.output) / "schema.json"
+    schema_path.write_text(final_state.get("current_schema").model_dump_json(indent=2))
+    logger.info("Schema saved to {}", schema_path)
+
+
+def _build_retriever(r_cfg):
+    if r_cfg.type == "weaviate":
+        try:
+            from schematize.retrieval.weaviate import WeaviateRetriever
+        except ImportError:
+            raise SystemExit("Weaviate retriever requires: pip install schematize[weaviate]")
+        filters = dict(r_cfg.wv_filters) if r_cfg.wv_filters else {}
+        return WeaviateRetriever(
+            collection_name=r_cfg.collection_name,
+            target_vector=r_cfg.target_vector or None,
+            filters=filters,
+        )
+    try:
+        from schematize.retrieval.huggingface import HuggingFaceRetriever, MMLWRobertaV2Retriever
+    except ImportError:
+        raise SystemExit("HuggingFace retriever requires: pip install schematize[huggingface]")
+    kwargs = dict(
+        dataset_name=r_cfg.dataset_name,
+        text_column=r_cfg.text_column,
+        split=r_cfg.split,
+        max_documents=r_cfg.max_documents,
+        index_path=r_cfg.index_path,
+    )
+    return MMLWRobertaV2Retriever(**kwargs) if r_cfg.type == "mmlw" else HuggingFaceRetriever(**kwargs)
+
+
+def _load_cases(cases_path: Path) -> dict[str, dict[str, str]]:
+    assert cases_path.exists(), f"Cases directory does not exist: {cases_path}"
+    return {p.stem: yaml.safe_load(p.read_text()) for p in cases_path.glob("*.yaml")}
+
+
+def _mock_problem_helper(text: str):
+    def _fn(state: AgentState) -> dict[str, Any]:
+        return {"messages": [AIMessage(content=text)], "problem_help": text}
+    return _fn
+
+
+def _mock_user_feedback(text: str):
+    def _fn(state: AgentState) -> dict[str, Any]:
+        return {"user_feedback": text, "messages": [HumanMessage(content=text)]}
+    return _fn
+
+
+def _mock_human_message(text: str):
+    def _fn(state: AgentState) -> dict[str, Any]:
+        return {
+            "messages": [HumanMessage(content=text)],
+            "final_messages": [HumanMessage(content=text)],
+        }
+    return _fn
+
+
+if __name__ == "__main__":
+    main()