runapprentice 0.1.0__py3-none-any.whl

runapprentice/langchain.py

@@ -0,0 +1,403 @@
+"""LangChain integration: zero-code-change capture via callbacks.
+
+Usage:
+    from runapprentice import Apprentice
+    from runapprentice.langchain import ApprenticeCallback
+
+    client = Apprentice(base_url=..., api_key=...)
+    model = init_chat_model("gpt-5.5", callbacks=[ApprenticeCallback("content_gap", client)])
+
+Every LLM call then logs input/output/latency to the task's dataset (raw tier).
+
+FAIL-OPEN CONTRACT: nothing in this module may ever raise into the customer's
+call path. All capture work is wrapped; worst case is a missing trace.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Sequence
+from typing import Any, Callable
+from uuid import UUID
+
+from runapprentice.client import Apprentice
+from runapprentice.models import TraceRecord
+
+try:
+    from langchain_core.callbacks import BaseCallbackHandler
+except ImportError:  # pragma: no cover - exercised only without the optional extra
+    BaseCallbackHandler = object  # type: ignore[assignment, misc]
+
+logger = logging.getLogger("runapprentice.langchain")
+
+RedactFn = Callable[[str], str]
+
+_SUPPORTED_FORMAT = "f-string"
+_ROLE_BY_CLASS = {
+    "SystemMessagePromptTemplate": "system",
+    "HumanMessagePromptTemplate": "human",
+    "AIMessagePromptTemplate": "ai",
+}
+
+
+class ApprenticeCallback(BaseCallbackHandler):
+    """Captures chat-model calls for one task. Streaming/async safe: LangChain
+    invokes `on_llm_end` after the final chunk for streamed runs."""
+
+    # never let a callback error abort the customer's chain
+    raise_error = False
+
+    def __init__(
+        self,
+        task: str,
+        client: Apprentice,
+        redact: RedactFn | None = None,
+    ) -> None:
+        _require_langchain()
+        self.task = task
+        self._client = client
+        self._redact = redact
+        self._starts: dict[UUID, tuple[str, float, str | None]] = {}
+        self._parents: dict[UUID, UUID | None] = {}
+        self._questions: dict[UUID, str] = {}
+        self._contexts: dict[UUID, str] = {}
+
+    # -- chain / retriever lifecycle -------------------------------------------
+
+    def on_chain_start(
+        self,
+        serialized: dict[str, Any],
+        inputs: dict[str, Any] | Any,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            self._parents[run_id] = parent_run_id
+            question = _question_text(inputs)
+            if question is not None:
+                self._questions[run_id] = question
+        except Exception as exc:  # noqa: BLE001 — fail-open contract
+            logger.debug("chain capture start failed: %s", exc)
+
+    def on_chain_end(
+        self,
+        outputs: dict[str, Any] | Any,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            self._questions.pop(run_id, None)
+            self._contexts.pop(run_id, None)
+            self._parents.pop(run_id, None)
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("chain capture end failed: %s", exc)
+
+    def on_retriever_start(
+        self,
+        serialized: dict[str, Any],
+        query: str,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            self._parents[run_id] = parent_run_id
+            if parent_run_id is None:
+                self._questions[run_id] = query
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("retriever capture start failed: %s", exc)
+
+    def on_retriever_end(
+        self,
+        documents: Sequence[Any],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            context = _documents_context(documents)
+            if context is None:
+                return
+            # Store context on the lineage ROOT, not the immediate parent: in the
+            # canonical LCEL shape ({"context": retriever | format, "question": ...}
+            # | prompt | llm) the retriever sits under a RunnableParallel sibling, so
+            # only the shared root run is also an ancestor of the model call. The
+            # model's on_llm_end walks its lineage to that root to pair context with
+            # question. ponytail: one retriever per simple RAG chain; use explicit
+            # capture() when custom formatting or multiple retrievers matter.
+            root = _lineage(run_id, self._parents)[-1]
+            self._contexts[root] = context
+            # The retriever run gets no on_chain_end, so drop its own transient
+            # entries here — otherwise they leak for the life of the process.
+            self._parents.pop(run_id, None)
+            if root != run_id:
+                self._questions.pop(run_id, None)
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("retriever capture end failed: %s", exc)
+
+    # -- chat model lifecycle ---------------------------------------------------
+
+    def on_chat_model_start(
+        self,
+        serialized: dict[str, Any],
+        messages: list[list[Any]],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            self._parents[run_id] = parent_run_id
+            text = "\n".join(f"{m.type}: {_message_text(m)}" for batch in messages for m in batch)
+            model_name = _model_name(serialized, kwargs)
+            self._starts[run_id] = (text, time.monotonic(), model_name)
+        except Exception as exc:  # noqa: BLE001 — fail-open contract
+            logger.debug("capture start failed: %s", exc)
+
+    def on_llm_start(
+        self,
+        serialized: dict[str, Any],
+        prompts: list[str],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            self._parents[run_id] = parent_run_id
+            self._starts[run_id] = (
+                "\n".join(prompts),
+                time.monotonic(),
+                _model_name(serialized, kwargs),
+            )
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("capture start failed: %s", exc)
+
+    def on_llm_end(self, response: Any, *, run_id: UUID, **kwargs: Any) -> None:
+        try:
+            start = self._starts.pop(run_id, None)
+            input_text, t0, model_name = start if start else ("", None, None)
+            output_text = "\n".join(gen.text for batch in response.generations for gen in batch)
+            usage = (response.llm_output or {}).get("token_usage", {}) or {}
+            if self._redact is not None:
+                input_text = self._redact(input_text)
+                output_text = self._redact(output_text)
+            inputs = self._structured_inputs(run_id)
+            if inputs is not None and self._redact is not None:
+                inputs = {key: self._redact(value) for key, value in inputs.items()}
+            record = TraceRecord(
+                task=self.task,
+                input=input_text,
+                inputs=inputs,
+                output=output_text,
+                model=model_name,
+                latency_ms=(time.monotonic() - t0) * 1000 if t0 is not None else None,
+                prompt_tokens=usage.get("prompt_tokens"),
+                completion_tokens=usage.get("completion_tokens"),
+                metadata={"structured_capture": inputs is not None},
+            )
+            self._client.post_trace_failopen(record)
+            self._parents.pop(run_id, None)
+        except Exception as exc:  # noqa: BLE001 — fail-open contract
+            logger.debug("capture end failed: %s", exc)
+
+    def on_llm_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any) -> None:
+        try:
+            start = self._starts.pop(run_id, None)
+            input_text, t0, model_name = start if start else ("", None, None)
+            if self._redact is not None:
+                input_text = self._redact(input_text)
+            self._client.post_trace_failopen(
+                TraceRecord(
+                    task=self.task,
+                    input=input_text,
+                    output="",
+                    model=model_name,
+                    latency_ms=(time.monotonic() - t0) * 1000 if t0 is not None else None,
+                    error=str(error)[:2000],
+                )
+            )
+            self._parents.pop(run_id, None)
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("capture error-path failed: %s", exc)
+
+    def _structured_inputs(self, run_id: UUID) -> dict[str, str] | None:
+        question: str | None = None
+        context: str | None = None
+        for current in _lineage(run_id, self._parents):
+            question = question or self._questions.get(current)
+            context = context or self._contexts.get(current)
+            if question is not None and context is not None:
+                return {"question": question, "context": context}
+        return None
+
+
+def _message_text(m: Any) -> str:
+    """Message text across LangChain versions: `.text` is a property in current
+    releases (was a method before; deprecated). Prefer the property, fall back
+    to the callable form, then to `.content`."""
+    text_attr = getattr(type(m), "text", None)
+    if isinstance(text_attr, property):
+        text = text_attr.__get__(m, type(m))
+        if text is not None:
+            return str(text)
+    text = getattr(m, "text", None)
+    if callable(text):  # pre-property LangChain
+        try:
+            return str(text())
+        except Exception:
+            return str(getattr(m, "content", ""))
+    if text is not None:
+        return str(text)
+    return str(getattr(m, "content", ""))
+
+
+def _question_text(inputs: Any) -> str | None:
+    if isinstance(inputs, str):
+        return inputs
+    if not isinstance(inputs, dict):
+        return None
+    for key in ("question", "input", "query"):
+        value = inputs.get(key)
+        if isinstance(value, str):
+            return value
+    strings = [value for value in inputs.values() if isinstance(value, str)]
+    return strings[0] if len(strings) == 1 else None
+
+
+def _documents_context(documents: Sequence[Any]) -> str | None:
+    parts = [str(content) for doc in documents if (content := getattr(doc, "page_content", None))]
+    return "\n\n".join(parts) if parts else None
+
+
+def _lineage(run_id: UUID, parents: dict[UUID, UUID | None]) -> list[UUID]:
+    lineage = [run_id]
+    current = run_id
+    seen = {run_id}
+    while (parent := parents.get(current)) is not None and parent not in seen:
+        lineage.append(parent)
+        seen.add(parent)
+        current = parent
+    return lineage
+
+
+def _model_name(serialized: dict[str, Any], kwargs: dict[str, Any]) -> str | None:
+    invocation = kwargs.get("invocation_params") or {}
+    name = invocation.get("model") or invocation.get("model_name")
+    if name:
+        return str(name)
+    ids = serialized.get("id")
+    return str(ids[-1]) if isinstance(ids, list) and ids else None
+
+
+def serialize_prompt(template: Any) -> dict[str, Any]:
+    """Serialize LangChain prompt templates to Apprentice Contract 2.
+
+    Plain `PromptTemplate` becomes one human message because it has no chat role.
+    """
+    ChatPromptTemplate, PromptTemplate, MessagesPlaceholder = _prompt_classes()
+    if isinstance(template, ChatPromptTemplate):
+        messages = [_serialize_message(m, MessagesPlaceholder) for m in template.messages]
+        formats = {_message_format(m, MessagesPlaceholder) for m in template.messages}
+        template_format = formats.pop() if formats else _SUPPORTED_FORMAT
+        if formats:
+            raise ValueError("mixed template formats are not supported; use one f-string format")
+        _reject_format(template_format, messages)
+        return {
+            "format": template_format,
+            "messages": messages,
+            "input_variables": list(template.input_variables),
+            "partial_variables": dict(getattr(template, "partial_variables", {}) or {}),
+        }
+    if isinstance(template, PromptTemplate):
+        template_format = str(getattr(template, "template_format", _SUPPORTED_FORMAT))
+        messages = [{"role": "human", "template": str(template.template)}]
+        _reject_format(template_format, messages)
+        return {
+            "format": template_format,
+            "messages": messages,
+            "input_variables": list(template.input_variables),
+            "partial_variables": dict(getattr(template, "partial_variables", {}) or {}),
+        }
+    raise ValueError("unsupported LangChain prompt type; use ChatPromptTemplate or PromptTemplate")
+
+
+def artifact_to_langchain(artifact: dict[str, Any]) -> Any:
+    ChatPromptTemplate, _, _ = _prompt_classes()
+    template_format = str(artifact.get("format", _SUPPORTED_FORMAT))
+    messages = artifact.get("messages")
+    if not isinstance(messages, list):
+        raise ValueError("template artifact must include messages list")
+    _reject_format(template_format, messages)
+    lc_messages = []
+    for message in messages:
+        if not isinstance(message, dict):
+            raise ValueError("template artifact messages must be objects")
+        role = str(message.get("role", ""))
+        if role not in {"system", "human", "ai"}:
+            raise ValueError(f"unsupported message role {role!r}; use system, human, or ai")
+        lc_messages.append((role, str(message.get("template", ""))))
+    return ChatPromptTemplate.from_messages(lc_messages, template_format=template_format)
+
+
+def _serialize_message(message: Any, MessagesPlaceholder: type[Any]) -> dict[str, str]:
+    if isinstance(message, MessagesPlaceholder):
+        raise ValueError(
+            "MessagesPlaceholder is not supported yet; upload flattened scalar inputs instead"
+        )
+    prompt = getattr(message, "prompt", None)
+    template = getattr(prompt, "template", None)
+    if template is None:
+        raise ValueError("only message prompt templates are supported; use tuple/string templates")
+    role = _ROLE_BY_CLASS.get(type(message).__name__)
+    if role is None:
+        raise ValueError("unsupported chat message role; use system, human, or ai message templates")
+    return {"role": role, "template": str(template)}
+
+
+def _message_format(message: Any, MessagesPlaceholder: type[Any]) -> str:
+    if isinstance(message, MessagesPlaceholder):
+        return _SUPPORTED_FORMAT
+    prompt = getattr(message, "prompt", None)
+    return str(getattr(prompt, "template_format", _SUPPORTED_FORMAT))
+
+
+def _reject_format(template_format: str, messages: list[dict[str, Any]]) -> None:
+    if any(_has_template_loop(template_format, str(m.get("template", ""))) for m in messages):
+        raise ValueError(
+            "jinja2/mustache loops are not supported yet; use flat f-string variables"
+        )
+    if template_format != _SUPPORTED_FORMAT:
+        raise ValueError(
+            f"template_format={template_format!r} is not supported yet; "
+            "use f-string templates for Apprentice dynamic prompts"
+        )
+
+
+def _has_template_loop(template_format: str, template: str) -> bool:
+    if template_format == "jinja2":
+        return "{% for " in template or "{% endfor" in template
+    if template_format == "mustache":
+        return "{{#" in template or "{{/" in template
+    return False
+
+
+def _prompt_classes() -> tuple[type[Any], type[Any], type[Any]]:
+    try:
+        from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate
+    except ImportError as exc:
+        raise ImportError(
+            "LangChain support requires the optional extra: "
+            "`uv add 'runapprentice[langchain]'` or install `langchain-core>=0.3`."
+        ) from exc
+    return ChatPromptTemplate, PromptTemplate, MessagesPlaceholder
+
+
+def _require_langchain() -> None:
+    _prompt_classes()

runapprentice/models.py

@@ -0,0 +1,62 @@
+"""Pydantic models for the Apprentice SDK. Mirrors the control-plane API contract."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class DatasetRow(BaseModel):
+    input: str
+    output: str
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class DatasetStatus(BaseModel):
+    task: str
+    gold: int = 0
+    silver: int = 0
+    raw: int = 0
+    ready_for_optimization: bool = False
+
+
+class TraceRecord(BaseModel):
+    task: str
+    input: str
+    inputs: dict[str, str] | None = None
+    output: str
+    model: str | None = None
+    latency_ms: float | None = None
+    prompt_tokens: int | None = None
+    completion_tokens: int | None = None
+    error: str | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class OptimizationReport(BaseModel):
+    task: str
+    baseline_score: float | None = None
+    optimized_score: float | None = None
+    examples_used: int | None = None
+    optimized_prompt: str | None = None
+    optimized_template: dict[str, Any] | None = None
+    detail: dict[str, Any] = Field(default_factory=dict)
+
+
+class JobInfo(BaseModel):
+    job_id: str
+    task: str
+    kind: str
+    status: str  # queued | running | succeeded | failed
+    created_at: datetime | None = None
+    error: str | None = None
+    report: OptimizationReport | None = None
+
+
+class PromptVersion(BaseModel):
+    task: str
+    version: int
+    text: str
+    score: float | None = None

runapprentice/openai.py

@@ -0,0 +1,126 @@
+"""OpenAI SDK integration: zero-config capture by wrapping the client.
+
+Usage:
+    from openai import OpenAI            # or AsyncOpenAI
+    from runapprentice import Apprentice
+    from runapprentice.openai import wrap
+
+    client = wrap(OpenAI(), Apprentice(api_key="..."), task="ticket-triage")
+    # every client.chat.completions.create(...) call is captured, unchanged.
+
+Works with both the sync `OpenAI` and async `AsyncOpenAI` clients (the wrapper
+matches whichever `create` it finds).
+
+This captures traces (it does NOT proxy or route your calls — your request goes
+straight to OpenAI as before). Routing/replacement is a later, opt-in stage.
+
+FAIL-OPEN CONTRACT: capture never raises into the customer's call path, and never
+blocks it (delivery is buffered off-thread). The real OpenAI call happens first and
+its result is always returned; worst case is a missing trace. Real OpenAI errors
+propagate unchanged.
+
+STREAMING: when `stream=True`, the response is an iterator that must not be consumed
+by the capture layer (that would break the caller's stream). Such calls pass through
+uncaptured — breaking streaming is never worth a trace.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+import time
+from typing import Any, Callable
+
+from runapprentice.client import Apprentice
+from runapprentice.models import TraceRecord
+
+logger = logging.getLogger("runapprentice.openai")
+
+RedactFn = Callable[[str], str]
+
+
+def wrap(openai_client: Any, client: Apprentice, task: str, redact: RedactFn | None = None) -> Any:
+    """Wrap an OpenAI client so chat.completions.create calls are captured.
+
+    Returns the same client (mutated in place) for convenient one-line use.
+    Idempotent: wrapping an already-wrapped client is a no-op. Supports both sync
+    (`OpenAI`) and async (`AsyncOpenAI`) clients.
+    """
+    completions = openai_client.chat.completions
+    if getattr(completions, "_apprentice_wrapped", False):
+        return openai_client
+    original = completions.create
+
+    def _capture(response: Any, kwargs: dict[str, Any], args: tuple[Any, ...], t0: float) -> None:
+        try:
+            input_text = _messages_to_text(kwargs.get("messages") or (args[0] if args else []))
+            output_text = _response_to_text(response)
+            if redact is not None:
+                input_text = redact(input_text)
+                output_text = redact(output_text)
+            usage = getattr(response, "usage", None)
+            client.post_trace_failopen(
+                TraceRecord(
+                    task=task,
+                    input=input_text,
+                    output=output_text,
+                    model=getattr(response, "model", None) or kwargs.get("model"),
+                    latency_ms=(time.monotonic() - t0) * 1000,
+                    prompt_tokens=getattr(usage, "prompt_tokens", None),
+                    completion_tokens=getattr(usage, "completion_tokens", None),
+                )
+            )
+        except Exception as exc:  # noqa: BLE001 — fail-open contract
+            logger.debug("capture failed: %s", exc)
+
+    if inspect.iscoroutinefunction(original):
+
+        @functools.wraps(original)
+        async def wrapped(*args: Any, **kwargs: Any) -> Any:
+            t0 = time.monotonic()
+            response = await original(*args, **kwargs)  # real call — never intercepted
+            if not kwargs.get("stream"):
+                _capture(response, kwargs, args, t0)
+            return response
+
+    else:
+
+        @functools.wraps(original)
+        def wrapped(*args: Any, **kwargs: Any) -> Any:
+            t0 = time.monotonic()
+            response = original(*args, **kwargs)  # real call — never intercepted
+            if not kwargs.get("stream"):
+                _capture(response, kwargs, args, t0)
+            return response
+
+    wrapped._apprentice_wrapped = True  # type: ignore[attr-defined]
+    completions.create = wrapped
+    completions._apprentice_wrapped = True
+    return openai_client
+
+
+def _messages_to_text(messages: Any) -> str:
+    if isinstance(messages, str):
+        return messages
+    parts = []
+    for m in messages or []:
+        role = m.get("role", "") if isinstance(m, dict) else getattr(m, "role", "")
+        content = m.get("content", "") if isinstance(m, dict) else getattr(m, "content", "")
+        if isinstance(content, list):  # multimodal: keep the text parts
+            content = "\n".join(
+                p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text"
+            )
+        parts.append(f"{role}: {content}".strip(": ").strip() if role else str(content))
+    return "\n".join(parts)
+
+
+def _response_to_text(response: Any) -> str:
+    choices = getattr(response, "choices", None) or []
+    out = []
+    for c in choices:
+        msg = getattr(c, "message", None)
+        text = getattr(msg, "content", None) if msg is not None else getattr(c, "text", None)
+        if text:
+            out.append(text)
+    return "\n".join(out)