tracelabel 0.1.0__py3-none-any.whl

tracelabel/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

tracelabel/__main__.py

@@ -0,0 +1,4 @@
+from tracelabel.cli.app import run
+
+if __name__ == "__main__":
+    run()

tracelabel/api/app.py

@@ -0,0 +1,40 @@
+from pathlib import Path
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+
+from tracelabel.config.models import ResolvedTaskConfig
+from tracelabel.db.database import Database
+from tracelabel.errors import NotFoundError, UserError
+
+from .labeling import LabelingService
+from .routes import create_api_router
+from .static import configure_static, default_static_dir
+
+
+def create_app(
+    database: Database,
+    config: ResolvedTaskConfig,
+    queue: list[str],
+    static_dir: Path | None = None,
+) -> FastAPI:
+    service = LabelingService(
+        config,
+        queue,
+        database.traces,
+        database.tasks,
+        database.annotations,
+    )
+    app = FastAPI()
+
+    @app.exception_handler(NotFoundError)
+    async def not_found(_request: Request, error: NotFoundError) -> JSONResponse:
+        return JSONResponse(status_code=404, content={"detail": str(error)})
+
+    @app.exception_handler(UserError)
+    async def user_error(_request: Request, error: UserError) -> JSONResponse:
+        return JSONResponse(status_code=422, content={"detail": str(error)})
+
+    app.include_router(create_api_router(service))
+    configure_static(app, static_dir or default_static_dir())
+    return app

tracelabel/api/labeling.py

@@ -0,0 +1,202 @@
+import sqlite3
+from typing import Any, cast
+
+from tracelabel.config.models import ResolvedTaskConfig
+from tracelabel.config.validation import AnnotationValidator
+from tracelabel.db.annotations import AnnotationRepository
+from tracelabel.db.database import decode_json
+from tracelabel.db.tasks import TaskRepository
+from tracelabel.db.traces import TraceRepository
+from tracelabel.errors import NotFoundError, UserError
+
+from .models import (
+    AnnotationIn,
+    AnnotationOut,
+    DocumentOut,
+    Progress,
+    QueueEntry,
+    SessionInfo,
+    SuggestionOut,
+    TraceDetail,
+    TraceInfo,
+    TurnOut,
+)
+
+
+class LabelingService:
+    def __init__(
+        self,
+        config: ResolvedTaskConfig,
+        queue: list[str],
+        traces: TraceRepository,
+        tasks: TaskRepository,
+        annotations: AnnotationRepository,
+        validator: AnnotationValidator | None = None,
+    ) -> None:
+        self._config = config
+        self._queue = queue
+        self._traces = traces
+        self._tasks = tasks
+        self._annotations = annotations
+        self._validator = validator or AnnotationValidator(config.fields)
+
+    def session(self) -> SessionInfo:
+        return SessionInfo(
+            task=self._config.name,
+            level=self._config.level,
+            fields=self._config.fields,
+            label_roles=self._config.label_roles,
+            annotator=self._config.annotator,
+            schema_hash=self._config.schema_hash,
+            shuffle=self._config.shuffle,
+        )
+
+    def queue(self) -> list[QueueEntry]:
+        counts = self._annotations.target_counts(self._task(), self._config.annotator)
+        entries: list[QueueEntry] = []
+        for position, trace_id in enumerate(self._queue):
+            target_count, labeled_count, skipped_count = counts.get(trace_id, (0, 0, 0))
+            entries.append(
+                QueueEntry(
+                    trace_id=trace_id,
+                    position=position,
+                    n_targets=target_count,
+                    n_labeled=labeled_count,
+                    n_skipped=skipped_count,
+                )
+            )
+        return entries
+
+    def trace_detail(self, trace_id: str) -> TraceDetail:
+        trace = self._traces.get(trace_id)
+        if trace is None:
+            raise NotFoundError(f"unknown trace '{trace_id}'")
+        document = self._document_out(trace)
+        turns = (
+            [] if document else [self._turn_out(row) for row in self._traces.get_turns(trace_id)]
+        )
+        annotations = {
+            str(row["target_id"]): self._annotation_out(row)
+            for row in self._annotations.annotations_for_trace(
+                self._config.name,
+                self._config.annotator,
+                trace_id,
+            )
+        }
+        suggestions = {
+            str(row["target_id"]): self._suggestion_out(row)
+            for row in self._annotations.suggestions_for_trace(self._config.name, trace_id)
+        }
+        return TraceDetail(
+            trace=TraceInfo(
+                id=trace["id"],
+                source=trace["source"],
+                metadata=self._json_object(trace["metadata"]),
+            ),
+            turns=turns,
+            document=document,
+            annotations=annotations,
+            suggestions=suggestions,
+        )
+
+    def put_annotation(self, annotation: AnnotationIn) -> AnnotationOut:
+        self._validate_target(annotation)
+        self._validator.validate(annotation.values, annotation.status)
+        row = self._annotations.upsert_annotation(
+            task=self._config.name,
+            target_type=annotation.target_type,
+            target_id=annotation.target_id,
+            status=annotation.status,
+            values=dict(annotation.values),
+            annotator=self._config.annotator,
+            schema_hash=self._config.schema_hash,
+            prefill_model=annotation.prefill_model,
+        )
+        return self._annotation_out(row)
+
+    def progress(self) -> Progress:
+        counts = self._annotations.target_counts(self._task(), self._config.annotator)
+        return Progress(
+            unit="turns" if self._config.level == "turn" else "traces",
+            total=sum(count[0] for count in counts.values()),
+            labeled=sum(count[1] for count in counts.values()),
+            skipped=sum(count[2] for count in counts.values()),
+        )
+
+    def _task(self) -> sqlite3.Row:
+        task = self._tasks.get(self._config.name)
+        if task is None:
+            raise UserError(f"task '{self._config.name}' is not open")
+        return task
+
+    def _validate_target(self, annotation: AnnotationIn) -> None:
+        if annotation.target_type != self._config.level:
+            raise UserError(
+                f"target_type '{annotation.target_type}' must match task level "
+                f"'{self._config.level}'"
+            )
+        if annotation.target_type == "trace":
+            if self._traces.get(annotation.target_id) is None:
+                raise NotFoundError(f"unknown trace '{annotation.target_id}'")
+            return
+        turn = self._traces.get_turn(annotation.target_id)
+        if turn is None:
+            raise NotFoundError(f"unknown turn '{annotation.target_id}'")
+        if turn["role"] not in self._config.label_roles:
+            raise UserError(
+                f"turn '{annotation.target_id}' has role '{turn['role']}', not labelable"
+            )
+
+    @staticmethod
+    def _document_out(trace: sqlite3.Row) -> DocumentOut | None:
+        content = trace["content"]
+        if content is None:
+            return None
+        return DocumentOut(content=content, content_type=trace["content_type"] or "text")
+
+    def _turn_out(self, row: sqlite3.Row) -> TurnOut:
+        raw_tool_calls = row["tool_calls"]
+        tool_calls = (
+            cast(list[dict[str, Any]], decode_json(raw_tool_calls))
+            if raw_tool_calls is not None
+            else None
+        )
+        return TurnOut(
+            id=row["id"],
+            idx=row["idx"],
+            role=row["role"],
+            content=row["content"],
+            content_type=row["content_type"],
+            tool_calls=tool_calls,
+            tool_call_id=row["tool_call_id"],
+            name=row["name"],
+            labelable=(self._config.level == "turn" and row["role"] in self._config.label_roles),
+            metadata=self._json_object(row["metadata"]),
+        )
+
+    @classmethod
+    def _annotation_out(cls, row: sqlite3.Row) -> AnnotationOut:
+        return AnnotationOut(
+            target_type=row["target_type"],
+            target_id=row["target_id"],
+            status=row["status"],
+            values=cls._json_object(row["values"]),
+            prefill_model=row["prefill_model"],
+            schema_hash=row["schema_hash"],
+            annotator=row["annotator"],
+            created_at=row["created_at"],
+            updated_at=row["updated_at"],
+        )
+
+    @classmethod
+    def _suggestion_out(cls, row: sqlite3.Row) -> SuggestionOut:
+        return SuggestionOut(
+            target_id=row["target_id"],
+            values=cls._json_object(row["values"]),
+            model=row["model"],
+            created_at=row["created_at"],
+        )
+
+    @staticmethod
+    def _json_object(raw: str) -> dict[str, Any]:
+        return cast(dict[str, Any], decode_json(raw))

tracelabel/api/models.py

@@ -0,0 +1,92 @@
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict
+
+from tracelabel.config.models import AnnotationStatus, Level
+from tracelabel.db.annotations import TargetType
+
+
+class SessionInfo(BaseModel):
+    task: str
+    level: Level
+    fields: list[dict[str, Any]]
+    label_roles: list[str]
+    annotator: str
+    schema_hash: str
+    shuffle: bool
+
+
+class QueueEntry(BaseModel):
+    trace_id: str
+    position: int
+    n_targets: int
+    n_labeled: int
+    n_skipped: int
+
+
+class TraceInfo(BaseModel):
+    id: str
+    source: str | None = None
+    metadata: dict[str, Any]
+
+
+class DocumentOut(BaseModel):
+    content: str
+    content_type: Literal["text", "json", "html", "markdown"]
+
+
+class TurnOut(BaseModel):
+    id: str
+    idx: int
+    role: str
+    content: str
+    content_type: Literal["text", "json", "html", "parts"]
+    tool_calls: list[dict[str, Any]] | None = None
+    tool_call_id: str | None = None
+    name: str | None = None
+    labelable: bool
+    metadata: dict[str, Any]
+
+
+class AnnotationIn(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    target_type: TargetType
+    target_id: str
+    status: AnnotationStatus
+    values: dict[str, str | list[str]]
+    prefill_model: str | None = None
+
+
+class AnnotationOut(BaseModel):
+    target_type: TargetType
+    target_id: str
+    status: AnnotationStatus
+    values: dict[str, Any]
+    prefill_model: str | None = None
+    schema_hash: str
+    annotator: str
+    created_at: str
+    updated_at: str
+
+
+class SuggestionOut(BaseModel):
+    target_id: str
+    values: dict[str, Any]
+    model: str
+    created_at: str
+
+
+class TraceDetail(BaseModel):
+    trace: TraceInfo
+    turns: list[TurnOut]
+    document: DocumentOut | None = None
+    annotations: dict[str, AnnotationOut]
+    suggestions: dict[str, SuggestionOut]
+
+
+class Progress(BaseModel):
+    unit: Literal["turns", "traces"]
+    total: int
+    labeled: int
+    skipped: int

tracelabel/api/routes.py

@@ -0,0 +1,37 @@
+from fastapi import APIRouter
+
+from .labeling import LabelingService
+from .models import (
+    AnnotationIn,
+    AnnotationOut,
+    Progress,
+    QueueEntry,
+    SessionInfo,
+    TraceDetail,
+)
+
+
+def create_api_router(service: LabelingService) -> APIRouter:
+    router = APIRouter(prefix="/api")
+
+    @router.get("/session", response_model=SessionInfo)
+    async def get_session() -> SessionInfo:
+        return service.session()
+
+    @router.get("/queue", response_model=list[QueueEntry])
+    async def get_queue() -> list[QueueEntry]:
+        return service.queue()
+
+    @router.get("/traces/{trace_id}", response_model=TraceDetail)
+    async def get_trace(trace_id: str) -> TraceDetail:
+        return service.trace_detail(trace_id)
+
+    @router.put("/annotations", response_model=AnnotationOut)
+    async def put_annotation(annotation: AnnotationIn) -> AnnotationOut:
+        return service.put_annotation(annotation)
+
+    @router.get("/progress", response_model=Progress)
+    async def get_progress() -> Progress:
+        return service.progress()
+
+    return router

tracelabel/api/static.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import FileResponse, Response
+from fastapi.staticfiles import StaticFiles
+from starlette.types import Scope
+
+
+class ImmutableStaticFiles(StaticFiles):
+    async def get_response(self, path: str, scope: Scope) -> Response:
+        response = await super().get_response(path, scope)
+        response.headers["Cache-Control"] = "public, max-age=31536000, immutable"
+        return response
+
+
+def default_static_dir() -> Path:
+    return Path(__file__).parent.parent / "static"
+
+
+def configure_static(app: FastAPI, static_dir: Path) -> None:
+    assets_dir = static_dir / "assets"
+    if assets_dir.is_dir():
+        app.mount("/assets", ImmutableStaticFiles(directory=assets_dir), name="assets")
+
+    async def spa(full_path: str) -> FileResponse:
+        if full_path == "api" or full_path.startswith("api/"):
+            raise HTTPException(status_code=404, detail="not found")
+        index = static_dir / "index.html"
+        if not index.is_file():
+            raise HTTPException(
+                status_code=503,
+                detail="frontend not built — run `npm run build` in frontend/ "
+                "or use the Vite dev server",
+            )
+        return FileResponse(index)
+
+    app.add_api_route("/{full_path:path}", spa, methods=["GET"])

tracelabel/cli/app.py

@@ -0,0 +1,163 @@
+from pathlib import Path
+from typing import Annotated
+
+import click
+import typer
+
+from tracelabel.config.loader import raw_config_for_target
+from tracelabel.config.models import CliArgs
+from tracelabel.config.resolver import ConfigResolver
+from tracelabel.errors import TraceLabelError, UserError
+
+from .commands import (
+    DemoCommand,
+    ExportCommand,
+    ImportCommand,
+    ServeCommand,
+    SuggestCommand,
+    TasksListCommand,
+)
+from .options import (
+    FormatChoice,
+    FromChoice,
+    LevelChoice,
+    OnConflictChoice,
+    StatusChoice,
+)
+
+app = typer.Typer(add_completion=False)
+tasks_app = typer.Typer(add_completion=False)
+app.add_typer(tasks_app, name="tasks")
+
+
+@app.callback()
+def main() -> None:
+    pass
+
+
+def target_path(target: str | None) -> Path:
+    if target is not None:
+        return Path(target)
+    fallback = Path("config.yaml")
+    if fallback.exists():
+        return fallback
+    raise UserError("No data file given (arg or `data:` in YAML)")
+
+
+@app.command()
+def serve(
+    target: Annotated[str | None, typer.Argument()] = None,
+    task: Annotated[str | None, typer.Option("--task")] = None,
+    level: Annotated[LevelChoice | None, typer.Option("--level")] = None,
+    annotator: Annotated[str | None, typer.Option("--annotator")] = None,
+    shuffle: Annotated[bool | None, typer.Option("--shuffle/--no-shuffle")] = None,
+    db: Annotated[Path | None, typer.Option("--db")] = None,
+    port: Annotated[int, typer.Option("--port")] = 8377,
+    no_browser: Annotated[bool, typer.Option("--no-browser")] = False,
+    yes: Annotated[bool, typer.Option("--yes")] = False,
+) -> None:
+    path = target_path(target)
+    cli = CliArgs(
+        task=task,
+        level=level.value if level else None,
+        annotator=annotator,
+        shuffle=shuffle,
+        db=db,
+        yes=yes,
+    )
+    config = ConfigResolver().resolve(raw_config_for_target(path), cli)
+    project_dir = path if path.is_dir() else path.parent
+    ServeCommand().execute(config, project_dir, db, port, no_browser, yes)
+
+
+@app.command(name="import")
+def import_(
+    target: Annotated[str, typer.Argument()],
+    from_: Annotated[FromChoice, typer.Option("--from")] = FromChoice.auto,
+    db: Annotated[Path | None, typer.Option("--db")] = None,
+    on_conflict: Annotated[
+        OnConflictChoice,
+        typer.Option("--on-conflict"),
+    ] = OnConflictChoice.fail,
+    skip_invalid: Annotated[bool, typer.Option("--skip-invalid")] = False,
+    as_documents: Annotated[bool, typer.Option("--as-documents")] = False,
+) -> None:
+    ImportCommand().execute(
+        Path(target),
+        database_path=db,
+        from_=from_.value,
+        on_conflict=on_conflict.value,
+        skip_invalid=skip_invalid,
+        as_documents=as_documents,
+    )
+
+
+@app.command()
+def export(
+    task: Annotated[str | None, typer.Option("--task")] = None,
+    db: Annotated[Path | None, typer.Option("--db")] = None,
+    format: Annotated[FormatChoice, typer.Option("--format")] = FormatChoice.jsonl,
+    joined: Annotated[bool, typer.Option("--joined")] = False,
+    out: Annotated[str | None, typer.Option("--out")] = None,
+    status: Annotated[StatusChoice, typer.Option("--status")] = StatusChoice.all,
+) -> None:
+    ExportCommand().execute(
+        task=task,
+        database_path=db,
+        format=format.value,
+        joined=joined,
+        out=Path(out) if out is not None else None,
+        status=status.value,
+    )
+
+
+@tasks_app.command("list")
+def tasks_list(db: Annotated[Path | None, typer.Option("--db")] = None) -> None:
+    TasksListCommand().execute(db)
+
+
+@app.command()
+def suggest(
+    target: Annotated[str | None, typer.Argument()] = None,
+    task: Annotated[str | None, typer.Option("--task")] = None,
+    db: Annotated[Path | None, typer.Option("--db")] = None,
+    limit: Annotated[int | None, typer.Option("--limit")] = None,
+    overwrite: Annotated[bool, typer.Option("--overwrite")] = False,
+    concurrency: Annotated[int, typer.Option("--concurrency")] = 4,
+) -> None:
+    path = target_path(target)
+    config = ConfigResolver().resolve(
+        raw_config_for_target(path),
+        CliArgs(task=task, db=db),
+    )
+    summary = SuggestCommand().execute(
+        config,
+        path.parent,
+        db,
+        limit=limit,
+        overwrite=overwrite,
+        concurrency=concurrency,
+    )
+    attempted = summary.ok + summary.failed
+    typer.echo(f"suggested {summary.ok}/{attempted} · {summary.failed} failed (re-run to retry)")
+
+
+@app.command()
+def demo(
+    port: Annotated[int, typer.Option("--port")] = 8377,
+    no_browser: Annotated[bool, typer.Option("--no-browser")] = False,
+) -> None:
+    DemoCommand().execute(port, no_browser)
+
+
+def run() -> None:
+    try:
+        app(standalone_mode=False)
+    except TraceLabelError as error:
+        typer.echo(str(error), err=True)
+        raise SystemExit(error.exit_code) from error
+    except (click.exceptions.Abort, KeyboardInterrupt) as error:
+        raise SystemExit(130) from error
+    except click.exceptions.ClickException as error:
+        error.show()
+        raise SystemExit(error.exit_code) from error