Add pluggable LLM API provider hook with tool calling

docs/DEVELOPMENT.md

@@ -125,6 +125,52 @@ def register_ast_validators() -> None:
     # Register AST validators
 ```
 
+### Available hooks
+
+Implement any subset of these in your plugin's `register_plugins.py`:
+
+| Hook | Returns | Notes |
+| --- | --- | --- |
+| `register_udfs` | `Sequence[Type[UDFBase]]` | Custom user-defined functions. |
+| `register_output_sinks` | `Sequence[BaseOutputSink]` | Where execution results go. |
+| `register_ast_validators` | `Sequence[Type[BaseValidator]]` | Extra SML validators. |
+| `register_action_proto_deserializer` | `ActionProtoDeserializer \| None` | Custom action proto → JSON. |
+| `register_input_stream` | `BaseInputStream` | Single-provider (`firstresult`). |
+| `register_execution_result_store` | `ExecutionResultStore` | Single-provider (`firstresult`). |
+| `register_labels_service_or_provider` | `LabelsServiceBase \| LabelsProvider` | Single-provider (`firstresult`). |
+| `register_llm_provider` | `BaseLLMProvider` | Single-provider (`firstresult`). LLM API access for AI-assisted features. |
+
+### LLM provider hook
+
+`register_llm_provider` lets a plugin supply the LLM API client used by AI-assisted
+features (e.g. natural-language query building). The interface lives in
+`osprey.worker.lib.llm` and is vendor-neutral and **tool-calling aware**: you pass
+`ToolDefinition`s in, the model may return `ToolCall`s, and you feed `ToolResult`s
+back on the next `chat()` call.
+
+```python
+from osprey.worker.lib.config import Config
+from osprey.worker.lib.llm.base import BaseLLMProvider
+
+@hookimpl_osprey
+def register_llm_provider(config: Config) -> BaseLLMProvider:
+    return MyLLMProvider(config)
+```
+
+Only one provider may be registered (`firstresult=True`). Retrieve it with
+`bootstrap_llm_provider(config)` from `osprey.worker.adaptor.plugin_manager`, which
+returns `None` when no plugin registers one — so callers should null-check.
+
+A direct Anthropic implementation is provided as a reference in
+`example_plugins/src/llm/anthropic_provider.py`, including the request/response and
+`tool_use` translation. It imports the `anthropic` SDK lazily; the SDK is **not**
+declared as a workspace dependency (it conflicts with the pinned
+`typing-extensions`), so install it manually to actually run the provider:
+
+```bash
+uv pip install anthropic
+```
+
 ## Rules
 
 Rules are written in SML, some examples are provided in `example_rules/` with YAML config, the rules are mounted to the worker processes when the containers start via environment variables. ex:

example_plugins/pyproject.toml

@@ -7,6 +7,15 @@ dependencies = [
     "pluggy==1.5.0"
 ]
 
+# NOTE: The example Anthropic LLM provider (llm/anthropic_provider.py) needs the
+# `anthropic` SDK, but it is intentionally NOT declared here. The workspace pins
+# `typing-extensions==4.6.3`, while every modern `anthropic` requires a newer
+# typing-extensions, so declaring it (even as an optional extra) makes the uv
+# workspace lock unsatisfiable. The provider imports `anthropic` lazily and raises
+# a clear error if it is missing; install it manually to use the provider, e.g.:
+#   uv pip install anthropic
+# (You may also need to relax the typing-extensions pin in the root workspace.)
+
 [tool.setuptools]
 package-dir = {"" = "src"}
 

example_plugins/src/llm/__init__.py

@@ -0,0 +1,5 @@
+"""Example LLM provider plugins for Osprey.
+
+See :mod:`llm.anthropic_provider` for a direct Anthropic Messages API implementation
+of :class:`osprey.worker.lib.llm.base.BaseLLMProvider`.
+"""

example_plugins/src/llm/anthropic_provider.py

@@ -0,0 +1,249 @@
+"""Example LLM provider backed directly by the Anthropic Messages API.
+
+This demonstrates implementing :class:`osprey.worker.lib.llm.base.BaseLLMProvider`,
+including tool calling: it translates the vendor-neutral ``LLMMessage`` /
+``ToolDefinition`` types into Anthropic's request format, and maps the response
+(including ``tool_use`` blocks) back into ``LLMResponse`` / ``ToolCall``.
+
+The ``anthropic`` SDK is imported lazily so the base example package, and Osprey's
+CI, do not require the SDK, an API key, or network access unless this provider is
+actually used. It is intentionally not a declared dependency (it conflicts with the
+pinned ``typing-extensions``; see ``example_plugins/pyproject.toml``), so install it
+manually to run the provider: ``uv pip install anthropic``.
+
+Configuration (via Osprey ``Config`` or environment):
+
+- API key: ``LLM_ANTHROPIC_API_KEY`` config key, else the ``ANTHROPIC_API_KEY``
+  environment variable (read by the SDK itself if neither is set explicitly).
+- Default model: ``LLM_ANTHROPIC_MODEL`` config key
+  (default: ``claude-3-5-sonnet-latest``).
+- Default max tokens: ``LLM_ANTHROPIC_MAX_TOKENS`` config key (default: ``1024``).
+"""
+
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union
+
+from osprey.worker.lib.config import Config
+from osprey.worker.lib.llm.base import (
+    BaseLLMProvider,
+    CacheControl,
+    LLMMessage,
+    LLMResponse,
+    LLMUsage,
+    ToolCall,
+    ToolDefinition,
+)
+
+if TYPE_CHECKING:
+    from anthropic import Anthropic
+
+DEFAULT_MODEL = 'claude-3-5-sonnet-latest'
+DEFAULT_MAX_TOKENS = 1024
+
+
+class AnthropicLLMProvider(BaseLLMProvider):
+    """A :class:`BaseLLMProvider` that calls the Anthropic Messages API directly."""
+
+    def __init__(self, config: Config, client: Optional['Anthropic'] = None) -> None:
+        self._config = config
+        self._default_model = config.get_str('LLM_ANTHROPIC_MODEL', DEFAULT_MODEL)
+        self._default_max_tokens = config.get_int('LLM_ANTHROPIC_MAX_TOKENS', DEFAULT_MAX_TOKENS)
+        # Allow injecting a client (used in tests); otherwise build lazily on first use.
+        self._client = client
+
+    def _get_client(self) -> 'Anthropic':
+        if self._client is not None:
+            return self._client
+
+        try:
+            import anthropic
+        except ImportError as exc:  # pragma: no cover - exercised only without the optional dep
+            raise RuntimeError(
+                "The 'anthropic' package is required to use AnthropicLLMProvider. "
+                'It is not a declared workspace dependency (it conflicts with the pinned '
+                'typing-extensions), so install it manually, e.g. `uv pip install anthropic`.'
+            ) from exc
+
+        api_key = self._config.get_optional_str('LLM_ANTHROPIC_API_KEY') or os.environ.get('ANTHROPIC_API_KEY')
+        # If api_key is None the SDK still reads ANTHROPIC_API_KEY from the environment itself.
+        self._client = anthropic.Anthropic(api_key=api_key) if api_key else anthropic.Anthropic()
+        return self._client
+
+    def chat(
+        self,
+        *,
+        messages: Sequence[LLMMessage],
+        system: Optional[str] = None,
+        tools: Optional[Sequence[ToolDefinition]] = None,
+        model: Optional[str] = None,
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        **params: Any,
+    ) -> LLMResponse:
+        request: Dict[str, Any] = {
+            # Use `is not None` rather than `or` so an explicit max_tokens=0 or
+            # model='' is passed through (and rejected by the API) instead of
+            # silently falling back to the default.
+            'model': model if model is not None else self._default_model,
+            'max_tokens': max_tokens if max_tokens is not None else self._default_max_tokens,
+            'messages': self._to_anthropic_messages(messages),
+        }
+
+        system_value = self._build_system(system, messages)
+        if system_value is not None:
+            request['system'] = system_value
+
+        if tools:
+            request['tools'] = [self._to_anthropic_tool(tool) for tool in tools]
+
+        if temperature is not None:
+            request['temperature'] = temperature
+
+        # Provider-specific passthrough (e.g. top_p, stop_sequences, tool_choice).
+        request.update(params)
+
+        response = self._get_client().messages.create(**request)
+        return self._from_anthropic_response(response)
+
+    # --- request translation ------------------------------------------------
+
+    @classmethod
+    def _build_system(
+        cls, system: Optional[str], messages: Sequence[LLMMessage]
+    ) -> Union[str, List[Dict[str, Any]], None]:
+        # Anthropic carries the system prompt as a top-level field, not a message,
+        # so the `system` argument and any role='system' messages are folded here.
+        # Each part keeps its own optional cache_control breakpoint.
+        parts: List[Tuple[str, Optional[CacheControl]]] = []
+        if system:
+            parts.append((system, None))
+        for message in messages:
+            if message.role == 'system' and message.content:
+                parts.append((message.content, message.cache_control))
+
+        if not parts:
+            return None
+
+        # When nothing needs a cache breakpoint, the simple string form suffices.
+        if all(cache_control is None for _, cache_control in parts):
+            return '\n\n'.join(text for text, _ in parts)
+
+        # Otherwise emit the structured block form so per-part cache_control is
+        # preserved (a string `system` cannot carry cache breakpoints).
+        blocks: List[Dict[str, Any]] = []
+        for text, cache_control in parts:
+            block: Dict[str, Any] = {'type': 'text', 'text': text}
+            if cache_control is not None:
+                block['cache_control'] = cls._cache_control_dict(cache_control)
+            blocks.append(block)
+        return blocks
+
+    @staticmethod
+    def _cache_control_dict(cache_control: CacheControl) -> Dict[str, Any]:
+        result: Dict[str, Any] = {'type': 'ephemeral'}
+        if cache_control.ttl is not None:
+            # A non-default ttl (e.g. '1h') requires the Anthropic extended-cache-ttl
+            # beta header on the client, or the API rejects the request; the default
+            # 5m ephemeral cache needs no header.
+            result['ttl'] = cache_control.ttl
+        return result
+
+    @staticmethod
+    def _to_anthropic_tool(tool: ToolDefinition) -> Dict[str, Any]:
+        return {
+            'name': tool.name,
+            'description': tool.description,
+            'input_schema': tool.input_schema,
+        }
+
+    @classmethod
+    def _to_anthropic_messages(cls, messages: Sequence[LLMMessage]) -> List[Dict[str, Any]]:
+        out: List[Dict[str, Any]] = []
+        for message in messages:
+            # System messages are handled separately via the top-level `system` field.
+            if message.role == 'system':
+                continue
+
+            blocks = cls._message_content_blocks(message)
+            if not blocks:
+                continue
+
+            if message.cache_control is not None:
+                blocks[-1]['cache_control'] = cls._cache_control_dict(message.cache_control)
+
+            # Tool results are surfaced to Anthropic as a user-role message.
+            role = 'user' if message.role == 'tool' else message.role
+            out.append({'role': role, 'content': blocks})
+        return out
+
+    @staticmethod
+    def _message_content_blocks(message: LLMMessage) -> List[Dict[str, Any]]:
+        blocks: List[Dict[str, Any]] = []
+
+        if message.content:
+            blocks.append({'type': 'text', 'text': message.content})
+
+        for tool_call in message.tool_calls:
+            blocks.append(
+                {
+                    'type': 'tool_use',
+                    'id': tool_call.id,
+                    'name': tool_call.name,
+                    'input': tool_call.arguments,
+                }
+            )
+
+        for tool_result in message.tool_results:
+            blocks.append(
+                {
+                    'type': 'tool_result',
+                    'tool_use_id': tool_result.tool_call_id,
+                    'content': tool_result.content,
+                    'is_error': tool_result.is_error,
+                }
+            )
+
+        return blocks
+
+    # --- response translation ------------------------------------------------
+
+    @staticmethod
+    def _from_anthropic_response(response: Any) -> LLMResponse:
+        text_parts: List[str] = []
+        tool_calls: List[ToolCall] = []
+
+        for block in getattr(response, 'content', None) or []:
+            block_type = getattr(block, 'type', None)
+            if block_type == 'text':
+                text_parts.append(getattr(block, 'text', '') or '')
+            elif block_type == 'tool_use':
+                raw_input = getattr(block, 'input', None)
+                tool_calls.append(
+                    ToolCall(
+                        id=getattr(block, 'id', ''),
+                        name=getattr(block, 'name', ''),
+                        # Degrade gracefully (like the other fields) if the SDK ever
+                        # hands back a non-dict input rather than crashing.
+                        arguments=dict(raw_input) if isinstance(raw_input, dict) else {},
+                    )
+                )
+
+        usage: Optional[LLMUsage] = None
+        raw_usage = getattr(response, 'usage', None)
+        if raw_usage is not None:
+            usage = LLMUsage(
+                input_tokens=getattr(raw_usage, 'input_tokens', 0) or 0,
+                output_tokens=getattr(raw_usage, 'output_tokens', 0) or 0,
+                cache_read_tokens=getattr(raw_usage, 'cache_read_input_tokens', 0) or 0,
+                cache_write_tokens=getattr(raw_usage, 'cache_creation_input_tokens', 0) or 0,
+            )
+
+        return LLMResponse(
+            text=''.join(text_parts),
+            tool_calls=tool_calls,
+            stop_reason=getattr(response, 'stop_reason', None),
+            usage=usage,
+            raw=response,
+        )