docent-python 0.1.20a0__py3-none-any.whl → 0.1.22a0__py3-none-any.whl

docent/_llm_util/providers/openrouter.py

@@ -0,0 +1,375 @@
+"""OpenRouter provider implementation using aiohttp library."""
+
+import json
+import os
+from typing import Any, Literal
+
+import aiohttp
+import backoff
+from backoff.types import Details
+
+from docent._llm_util.data_models.exceptions import (
+    CompletionTooLongException,
+    ContextWindowException,
+    NoResponseException,
+    RateLimitException,
+)
+from docent._llm_util.data_models.llm_output import (
+    AsyncSingleLLMOutputStreamingCallback,
+    LLMCompletion,
+    LLMOutput,
+    UsageMetrics,
+)
+from docent._log_util import get_logger
+from docent.data_models.chat import ChatMessage, Content, ToolCall, ToolInfo
+
+logger = get_logger(__name__)
+
+OPENROUTER_API_BASE = "https://openrouter.ai/api/v1"
+
+
+class OpenRouterClient:
+    """Async client for OpenRouter API using aiohttp."""
+
+    def __init__(self, api_key: str | None = None):
+        self.api_key = api_key or os.getenv("OPENROUTER_API_KEY")
+        self.base_url = OPENROUTER_API_BASE
+
+    def _get_headers(self) -> dict[str, str]:
+        return {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+    async def chat_completions_create(
+        self,
+        model: str,
+        messages: list[dict[str, Any]],
+        tools: list[dict[str, Any]] | None = None,
+        tool_choice: str | None = None,
+        max_tokens: int = 32,
+        temperature: float = 1.0,
+        timeout: float = 30.0,
+    ) -> dict[str, Any]:
+        """Make an async chat completion request."""
+        url = f"{self.base_url}/chat/completions"
+
+        payload: dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+        }
+
+        if tools:
+            payload["tools"] = tools
+        if tool_choice:
+            payload["tool_choice"] = tool_choice
+
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                url,
+                json=payload,
+                headers=self._get_headers(),
+                timeout=aiohttp.ClientTimeout(total=timeout),
+            ) as response:
+                if response.status != 200:
+                    try:
+                        error_data: dict[str, Any] = await response.json()
+                        error_msg: Any = error_data.get("error", {}).get(
+                            "message", await response.text()
+                        )
+                    except Exception:
+                        error_msg = await response.text()
+                    if response.status == 429:
+                        raise RateLimitException(f"OpenRouter rate limit: {error_msg}")
+                    elif response.status == 400 and "context" in str(error_msg).lower():
+                        raise ContextWindowException()
+                    else:
+                        raise Exception(f"OpenRouter API error ({response.status}): {error_msg}")
+
+                return await response.json()
+
+
+def get_openrouter_client_async(api_key: str | None = None) -> OpenRouterClient:
+    return OpenRouterClient(api_key=api_key)
+
+
+def _print_backoff_message(e: Details):
+    logger.warning(
+        f"OpenRouter backing off for {e['wait']:.2f}s due to {e['exception'].__class__.__name__}"  # type: ignore
+    )
+
+
+def _is_retryable_error(e: BaseException) -> bool:
+    if isinstance(e, RateLimitException):
+        return True
+    if isinstance(e, ContextWindowException):
+        return False
+    if isinstance(e, aiohttp.ClientError):
+        return True
+    return False
+
+
+def _parse_message_content(
+    content: str | list[Content],
+) -> str | list[dict[str, str]]:
+    if isinstance(content, str):
+        return content
+    else:
+        result: list[dict[str, str]] = []
+        for sub_content in content:
+            if sub_content.type == "text":
+                result.append({"type": "text", "text": sub_content.text})
+            else:
+                raise ValueError(f"Unsupported content type: {sub_content.type}")
+        return result
+
+
+def parse_chat_messages(messages: list[ChatMessage]) -> list[dict[str, Any]]:
+    """Convert ChatMessage list to OpenRouter format."""
+    result: list[dict[str, Any]] = []
+
+    for message in messages:
+        if message.role == "user":
+            result.append(
+                {
+                    "role": "user",
+                    "content": _parse_message_content(message.content),
+                }
+            )
+        elif message.role == "assistant":
+            msg: dict[str, Any] = {
+                "role": "assistant",
+                "content": _parse_message_content(message.content),
+            }
+            if message.tool_calls:
+                msg["tool_calls"] = [
+                    {
+                        "id": tc.id,
+                        "type": "function",
+                        "function": {
+                            "name": tc.function,
+                            "arguments": json.dumps(tc.arguments),
+                        },
+                    }
+                    for tc in message.tool_calls
+                ]
+            result.append(msg)
+        elif message.role == "tool":
+            result.append(
+                {
+                    "role": "tool",
+                    "content": _parse_message_content(message.content),
+                    "tool_call_id": str(message.tool_call_id),
+                }
+            )
+        elif message.role == "system":
+            result.append(
+                {
+                    "role": "system",
+                    "content": _parse_message_content(message.content),
+                }
+            )
+
+    return result
+
+
+def parse_tools(tools: list[ToolInfo]) -> list[dict[str, Any]]:
+    """Convert ToolInfo objects to OpenRouter format."""
+    result: list[dict[str, Any]] = []
+
+    for tool in tools:
+        result.append(
+            {
+                "type": "function",
+                "function": {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "parameters": tool.parameters.model_dump(exclude_none=True),
+                },
+            }
+        )
+
+    return result
+
+
+def _parse_openrouter_tool_call(tc: dict[str, Any]) -> ToolCall:
+    """Parse tool call from OpenRouter response."""
+    if tc.get("type") != "function":
+        return ToolCall(
+            id=tc.get("id", "unknown"),
+            function="unknown",
+            arguments={},
+            parse_error=f"Unsupported tool call type: {tc.get('type')}",
+            type=None,
+        )
+
+    function_data = tc.get("function", {})
+    arguments: dict[str, Any] = {}
+    try:
+        arguments = json.loads(function_data.get("arguments", "{}"))
+        parse_error = None
+    except Exception as e:
+        arguments = {"__parse_error_raw_args": function_data.get("arguments", "")}
+        parse_error = f"Couldn't parse tool call arguments as JSON: {e}"
+
+    return ToolCall(
+        id=tc.get("id", "unknown"),
+        function=function_data.get("name", "unknown"),
+        arguments=arguments,
+        parse_error=parse_error,
+        type="function",
+    )
+
+
+def parse_openrouter_completion(response: dict[str, Any], model: str) -> LLMOutput:
+    """Parse OpenRouter completion response."""
+    choices = response.get("choices", [])
+    if not choices:
+        return LLMOutput(
+            model=model,
+            completions=[],
+            errors=[NoResponseException()],
+        )
+
+    usage_data = response.get("usage", {})
+    usage = UsageMetrics(
+        input=usage_data.get("prompt_tokens", 0),
+        output=usage_data.get("completion_tokens", 0),
+    )
+
+    completions: list[LLMCompletion] = []
+    for choice in choices:
+        message = choice.get("message", {})
+        tool_calls_data = message.get("tool_calls")
+
+        completions.append(
+            LLMCompletion(
+                text=message.get("content"),
+                finish_reason=choice.get("finish_reason"),
+                tool_calls=(
+                    [_parse_openrouter_tool_call(tc) for tc in tool_calls_data]
+                    if tool_calls_data
+                    else None
+                ),
+            )
+        )
+
+    return LLMOutput(
+        model=response.get("model", model),
+        completions=completions,
+        usage=usage,
+    )
+
+
+@backoff.on_exception(
+    backoff.expo,
+    exception=(Exception,),
+    giveup=lambda e: not _is_retryable_error(e),
+    max_tries=5,
+    factor=3.0,
+    on_backoff=_print_backoff_message,
+)
+async def get_openrouter_chat_completion_async(
+    client: OpenRouterClient,
+    messages: list[ChatMessage],
+    model_name: str,
+    tools: list[ToolInfo] | None = None,
+    tool_choice: Literal["auto", "required"] | None = None,
+    max_new_tokens: int = 32,
+    temperature: float = 1.0,
+    reasoning_effort: Literal["low", "medium", "high"] | None = None,
+    logprobs: bool = False,
+    top_logprobs: int | None = None,
+    timeout: float = 30.0,
+) -> LLMOutput:
+    """Get completion from OpenRouter."""
+    if logprobs or top_logprobs is not None:
+        raise NotImplementedError(
+            "We have not implemented logprobs or top_logprobs for OpenRouter yet."
+        )
+
+    if reasoning_effort is not None:
+        logger.warning("OpenRouter does not support reasoning_effort parameter, ignoring.")
+
+    input_messages = parse_chat_messages(messages)
+    input_tools = parse_tools(tools) if tools else None
+
+    response = await client.chat_completions_create(
+        model=model_name,
+        messages=input_messages,
+        tools=input_tools,
+        tool_choice=tool_choice,
+        max_tokens=max_new_tokens,
+        temperature=temperature,
+        timeout=timeout,
+    )
+
+    output = parse_openrouter_completion(response, model_name)
+
+    if output.first and output.first.finish_reason == "length" and output.first.no_text:
+        raise CompletionTooLongException(
+            "Completion empty due to truncation. Consider increasing max_new_tokens."
+        )
+
+    return output
+
+
+@backoff.on_exception(
+    backoff.expo,
+    exception=(Exception,),
+    giveup=lambda e: not _is_retryable_error(e),
+    max_tries=5,
+    factor=3.0,
+    on_backoff=_print_backoff_message,
+)
+async def get_openrouter_chat_completion_streaming_async(
+    client: OpenRouterClient,
+    streaming_callback: AsyncSingleLLMOutputStreamingCallback | None,
+    messages: list[ChatMessage],
+    model_name: str,
+    tools: list[ToolInfo] | None = None,
+    tool_choice: Literal["auto", "required"] | None = None,
+    max_new_tokens: int = 32,
+    temperature: float = 1.0,
+    reasoning_effort: Literal["low", "medium", "high"] | None = None,
+    logprobs: bool = False,
+    top_logprobs: int | None = None,
+    timeout: float = 30.0,
+) -> LLMOutput:
+    """Get streaming completion from OpenRouter (falls back to non-streaming)."""
+    logger.warning("Streaming not yet implemented for OpenRouter, using non-streaming.")
+
+    return await get_openrouter_chat_completion_async(
+        client=client,
+        messages=messages,
+        model_name=model_name,
+        tools=tools,
+        tool_choice=tool_choice,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        reasoning_effort=reasoning_effort,
+        logprobs=logprobs,
+        top_logprobs=top_logprobs,
+        timeout=timeout,
+    )
+
+
+async def is_openrouter_api_key_valid(api_key: str) -> bool:
+    """Test whether an OpenRouter API key is valid."""
+    client = OpenRouterClient(api_key=api_key)
+
+    try:
+        # Make a minimal request to test the key
+        await client.chat_completions_create(
+            model="openai/gpt-5-nano",
+            messages=[{"role": "user", "content": "hi"}],
+            max_tokens=1,
+            timeout=10.0,
+        )
+        return True
+    except Exception as e:
+        if "authentication" in str(e).lower() or "authorization" in str(e).lower():
+            return False
+        raise

docent/_llm_util/providers/preference_types.py

@@ -0,0 +1,104 @@
+"""Provides preferences of which LLM models to use for different Docent functions."""
+
+from functools import cached_property
+from typing import Literal
+
+from pydantic import BaseModel
+
+from docent._llm_util.model_registry import get_context_window
+from docent._log_util import get_logger
+
+logger = get_logger(__name__)
+
+
+class ModelOption(BaseModel):
+    """Configuration for a specific model from a provider. Not to be confused with ModelInfo.
+
+    Attributes:
+        provider: The name of the LLM provider (e.g., "openai", "anthropic").
+        model_name: The specific model to use from the provider.
+        reasoning_effort: Optional indication of computational effort to use.
+    """
+
+    provider: str
+    model_name: str
+    reasoning_effort: Literal["low", "medium", "high"] | None = None
+
+
+class ModelOptionWithContext(BaseModel):
+    """Enhanced model option that includes context window information for frontend use.
+    Not to be confused with ModelInfo or ModelOption.
+
+    Attributes:
+        provider: The name of the LLM provider (e.g., "openai", "anthropic").
+        model_name: The specific model to use from the provider.
+        reasoning_effort: Optional indication of computational effort to use.
+        context_window: The context window size in tokens.
+        uses_byok: Whether this model would use the user's own API key.
+    """
+
+    provider: str
+    model_name: str
+    reasoning_effort: Literal["low", "medium", "high"] | None = None
+    context_window: int
+    uses_byok: bool
+
+    @classmethod
+    def from_model_option(
+        cls, model_option: ModelOption, uses_byok: bool = False
+    ) -> "ModelOptionWithContext":
+        """Create a ModelOptionWithContext from a ModelOption.
+
+        Args:
+            model_option: The base model option
+            uses_byok: Whether this model requires bring-your-own-key
+
+        Returns:
+            ModelOptionWithContext with context window looked up from global mapping
+        """
+        context_window = get_context_window(model_option.model_name)
+
+        return cls(
+            provider=model_option.provider,
+            model_name=model_option.model_name,
+            reasoning_effort=model_option.reasoning_effort,
+            context_window=context_window,
+            uses_byok=uses_byok,
+        )
+
+
+def merge_models_with_byok(
+    defaults: list[ModelOption],
+    byok: list[ModelOption],
+    api_keys: dict[str, str] | None,
+) -> list[ModelOptionWithContext]:
+    user_keys = api_keys or {}
+
+    merged: list[ModelOption] = list(defaults)
+    if user_keys:
+        merged.extend([m for m in byok if m.provider in user_keys])
+
+    return [ModelOptionWithContext.from_model_option(m, m.provider in user_keys) for m in merged]
+
+
+class PublicProviderPreferences(BaseModel):
+    @cached_property
+    def default_judge_models(self) -> list[ModelOption]:
+        """Judge models that any user can access without providing their own API key"""
+
+        return [
+            ModelOption(provider="openai", model_name="gpt-5", reasoning_effort="medium"),
+            ModelOption(provider="openai", model_name="gpt-5", reasoning_effort="low"),
+            ModelOption(provider="openai", model_name="gpt-5", reasoning_effort="high"),
+            ModelOption(provider="openai", model_name="gpt-5-mini", reasoning_effort="low"),
+            ModelOption(provider="openai", model_name="gpt-5-mini", reasoning_effort="medium"),
+            ModelOption(provider="openai", model_name="gpt-5-mini", reasoning_effort="high"),
+            ModelOption(
+                provider="anthropic",
+                model_name="claude-sonnet-4-20250514",
+                reasoning_effort="medium",
+            ),
+        ]
+
+
+PUBLIC_PROVIDER_PREFERENCES = PublicProviderPreferences()

docent/_llm_util/providers/provider_registry.py

@@ -0,0 +1,164 @@
+"""Registry for LLM providers with their configurations."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Literal, Protocol, TypedDict
+
+from docent._llm_util.data_models.llm_output import (
+    AsyncSingleLLMOutputStreamingCallback,
+    LLMOutput,
+)
+from docent._llm_util.providers import anthropic, google, openai, openrouter
+from docent._llm_util.providers.anthropic import (
+    get_anthropic_chat_completion_async,
+    get_anthropic_chat_completion_streaming_async,
+)
+from docent._llm_util.providers.google import (
+    get_google_chat_completion_async,
+    get_google_chat_completion_streaming_async,
+)
+from docent._llm_util.providers.openai import (
+    get_openai_chat_completion_async,
+    get_openai_chat_completion_streaming_async,
+)
+from docent._llm_util.providers.openrouter import (
+    get_openrouter_chat_completion_async,
+    get_openrouter_chat_completion_streaming_async,
+)
+from docent.data_models.chat import ChatMessage, ToolInfo
+
+
+class SingleOutputGetter(Protocol):
+    """Protocol for getting non-streaming output from an LLM.
+
+    Defines the interface for async functions that retrieve a single
+    non-streaming response from an LLM provider.
+    """
+
+    async def __call__(
+        self,
+        client: Any,
+        messages: list[ChatMessage],
+        model_name: str,
+        *,
+        tools: list[ToolInfo] | None,
+        tool_choice: Literal["auto", "required"] | None,
+        max_new_tokens: int,
+        temperature: float,
+        reasoning_effort: Literal["low", "medium", "high"] | None,
+        logprobs: bool,
+        top_logprobs: int | None,
+        timeout: float,
+    ) -> LLMOutput:
+        """Get a single completion from an LLM.
+
+        Args:
+            client: The provider-specific client instance.
+            messages: The list of messages in the conversation.
+            model_name: The name of the model to use.
+            tools: Optional list of tools available to the model.
+            tool_choice: Optional specification for tool usage.
+            max_new_tokens: Maximum number of tokens to generate.
+            temperature: Controls randomness in output generation.
+            reasoning_effort: Optional control for model reasoning depth.
+            logprobs: Whether to return log probabilities.
+            top_logprobs: Number of most likely tokens to return probabilities for.
+            timeout: Maximum time to wait for a response in seconds.
+
+        Returns:
+            LLMOutput: The model's response.
+        """
+        ...
+
+
+class SingleStreamingOutputGetter(Protocol):
+    """Protocol for getting streaming output from an LLM.
+
+    Defines the interface for async functions that retrieve streaming
+    responses from an LLM provider.
+    """
+
+    async def __call__(
+        self,
+        client: Any,
+        streaming_callback: AsyncSingleLLMOutputStreamingCallback | None,
+        messages: list[ChatMessage],
+        model_name: str,
+        *,
+        tools: list[ToolInfo] | None,
+        tool_choice: Literal["auto", "required"] | None,
+        max_new_tokens: int,
+        temperature: float,
+        reasoning_effort: Literal["low", "medium", "high"] | None,
+        logprobs: bool,
+        top_logprobs: int | None,
+        timeout: float,
+    ) -> LLMOutput:
+        """Get a streaming completion from an LLM.
+
+        Args:
+            client: The provider-specific client instance.
+            streaming_callback: Optional callback for processing streaming chunks.
+            messages: The list of messages in the conversation.
+            model_name: The name of the model to use.
+            tools: Optional list of tools available to the model.
+            tool_choice: Optional specification for tool usage.
+            max_new_tokens: Maximum number of tokens to generate.
+            temperature: Controls randomness in output generation.
+            reasoning_effort: Optional control for model reasoning depth.
+            logprobs: Whether to return log probabilities.
+            top_logprobs: Number of most likely tokens to return probabilities for.
+            timeout: Maximum time to wait for a response in seconds.
+
+        Returns:
+            LLMOutput: The complete model response after streaming finishes.
+        """
+        ...
+
+
+class ProviderConfig(TypedDict):
+    """Configuration for an LLM provider.
+
+    Contains the necessary functions to create clients and interact with
+    a specific LLM provider.
+
+    Attributes:
+        async_client_getter: Function to get an async client for the provider.
+        single_output_getter: Function to get a non-streaming completion.
+        single_streaming_output_getter: Function to get a streaming completion.
+    """
+
+    async_client_getter: Callable[[str | None], Any]
+    single_output_getter: SingleOutputGetter
+    single_streaming_output_getter: SingleStreamingOutputGetter
+
+
+# Registry of supported LLM providers with their respective configurations
+PROVIDERS: dict[str, ProviderConfig] = {
+    "anthropic": ProviderConfig(
+        async_client_getter=anthropic.get_anthropic_client_async,
+        single_output_getter=get_anthropic_chat_completion_async,
+        single_streaming_output_getter=get_anthropic_chat_completion_streaming_async,
+    ),
+    "google": ProviderConfig(
+        async_client_getter=google.get_google_client_async,
+        single_output_getter=get_google_chat_completion_async,
+        single_streaming_output_getter=get_google_chat_completion_streaming_async,
+    ),
+    "openai": ProviderConfig(
+        async_client_getter=openai.get_openai_client_async,
+        single_output_getter=get_openai_chat_completion_async,
+        single_streaming_output_getter=get_openai_chat_completion_streaming_async,
+    ),
+    "azure_openai": ProviderConfig(
+        async_client_getter=openai.get_azure_openai_client_async,
+        single_output_getter=get_openai_chat_completion_async,
+        single_streaming_output_getter=get_openai_chat_completion_streaming_async,
+    ),
+    "openrouter": ProviderConfig(
+        async_client_getter=openrouter.get_openrouter_client_async,
+        single_output_getter=get_openrouter_chat_completion_async,
+        single_streaming_output_getter=get_openrouter_chat_completion_streaming_async,
+    ),
+}
+"""Registry of supported LLM providers with their respective configurations."""

docent/data_models/transcript.py

@@ -42,6 +42,8 @@ Important notes:
 - Citations are self-contained. Do NOT label them as citation or evidence. Just insert the citation by itself at the appropriate place in the text.
 - Citations must come immediately after the part of a claim that they support. This may be in the middle of a sentence.
 - Each pair of brackets must contain only one citation. To cite multiple blocks, use multiple pairs of brackets, like [T0B0] [T0B1].
+- Outside of citations, do not refer to transcript numbers or block numbers.
+- Outside of citations, avoid quoting or paraphrasing the transcript.
 """
 
 BLOCK_CITE_INSTRUCTION = """Each transcript and each block has a unique index. Cite the relevant indices in brackets when relevant, like [T<idx>B<idx>]. Use multiple tags to cite multiple blocks, like [T<idx1>B<idx1>][T<idx2>B<idx2>]. Remember to cite specific blocks and NOT action units."""

docent/judges/__init__.py

@@ -0,0 +1,21 @@
+from docent.judges.impl import BaseJudge, MajorityVotingJudge, MultiReflectionJudge
+from docent.judges.types import (
+    JudgeResult,
+    JudgeResultCompletionCallback,
+    JudgeResultWithCitations,
+    ResultType,
+    Rubric,
+)
+
+__all__ = [
+    # Judges
+    "MajorityVotingJudge",
+    "MultiReflectionJudge",
+    "BaseJudge",
+    # Types
+    "Rubric",
+    "JudgeResult",
+    "JudgeResultWithCitations",
+    "JudgeResultCompletionCallback",
+    "ResultType",
+]