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,215 @@
+"""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 an optional dependency (``example_plugins[llm]``). It 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.
+
+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
+
+from osprey.worker.lib.config import Config
+from osprey.worker.lib.llm.base import (
+    BaseLLMProvider,
+    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. "
+                "Install it with the optional extra, e.g. `uv pip install 'example_plugins[llm]'`."
+            ) 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] = {
+            'model': model or self._default_model,
+            'max_tokens': max_tokens or self._default_max_tokens,
+            'messages': self._to_anthropic_messages(messages),
+        }
+
+        system_text = self._collect_system_text(system, messages)
+        if system_text is not None:
+            request['system'] = system_text
+
+        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 ------------------------------------------------
+
+    @staticmethod
+    def _collect_system_text(system: Optional[str], messages: Sequence[LLMMessage]) -> Optional[str]:
+        parts: List[str] = []
+        if system:
+            parts.append(system)
+        # Anthropic carries the system prompt as a top-level field, not a message,
+        # so fold any role='system' messages into it.
+        for message in messages:
+            if message.role == 'system' and message.content:
+                parts.append(message.content)
+        if not parts:
+            return None
+        return '\n\n'.join(parts)
+
+    @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:
+                cache_control: Dict[str, Any] = {'type': 'ephemeral'}
+                if message.cache_control.ttl is not None:
+                    cache_control['ttl'] = message.cache_control.ttl
+                blocks[-1]['cache_control'] = 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':
+                tool_calls.append(
+                    ToolCall(
+                        id=getattr(block, 'id', ''),
+                        name=getattr(block, 'name', ''),
+                        arguments=dict(getattr(block, 'input', {}) or {}),
+                    )
+                )
+
+        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,
+        )