polymath-agent 0.4.0__py3-none-any.whl

polymath/__init__.py

@@ -0,0 +1,2 @@
+"""polymath — multi-model AI orchestrator"""
+__version__ = "0.1.0"

polymath/adapters/__init__.py

@@ -0,0 +1,7 @@
+from polymath.adapters.base import BaseAdapter, Message
+from polymath.adapters.claude import ClaudeAdapter
+from polymath.adapters.gemini import GeminiAdapter
+from polymath.adapters.openai_adapter import OpenAIAdapter
+from polymath.adapters.ollama import OllamaAdapter
+
+__all__ = ["BaseAdapter", "Message", "ClaudeAdapter", "GeminiAdapter", "OpenAIAdapter", "OllamaAdapter"]

polymath/adapters/base.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, AsyncIterator, Dict, List, Optional
+
+
+#: Output ceilings. Current models write up to 128K tokens, and 4096 was
+#: cutting plans, reviews and file rewrites off mid-sentence. The
+#: non-streaming figure stays well inside the SDK's HTTP timeout; streaming
+#: has no such constraint, so it gets room.
+DEFAULT_MAX_TOKENS = 16_000
+DEFAULT_STREAM_MAX_TOKENS = 64_000
+
+
+def truncation_notice(limit: int) -> str:
+    """Appended to a reply the provider cut off at the output limit.
+
+    Without it a truncated answer is indistinguishable from a bad one: it
+    simply stops, and every layer downstream treats it as complete.
+    """
+    return (
+        f"\n\n[truncated: the model hit its output limit of {limit:,} tokens. "
+        "Ask it to continue, or raise max_tokens.]"
+    )
+
+
+class AuthExpiredError(Exception):
+    """Raised when an adapter receives an auth/401 error from the provider."""
+    def __init__(self, provider: str, relogin_cmd: str) -> None:
+        self.provider = provider
+        self.relogin_cmd = relogin_cmd
+        super().__init__(f"{provider} auth expired — run: {relogin_cmd}")
+
+
+def is_auth_failure(
+    exc: BaseException,
+    auth_types: tuple[type, ...] = (),
+) -> bool:
+    """Whether a provider exception means "your credentials are no longer good".
+
+    Every SDK raises a distinct class for this, so the class is the answer
+    whenever we can resolve it. Matching on the rendered message is the
+    fallback for the case where the SDK is too old to expose the class, and
+    it is deliberately narrow: an earlier version tested `"401" in text`,
+    which reported a rate limit whose body happened to mention 401 as an
+    expired login and sent the user off to log in again.
+    """
+    if auth_types and isinstance(exc, auth_types):
+        return True
+    status = getattr(exc, "status_code", None)
+    if status is None:
+        response = getattr(exc, "response", None)
+        status = getattr(response, "status_code", None)
+    if status in (401, 403):
+        return True
+    text = str(exc).lower()
+    return any(
+        marker in text
+        for marker in (
+            "authentication_error", "authentication error",
+            "authentication failed", "unauthorized",
+            "invalid api key", "invalid_api_key",
+            "unauthenticated", "invalid_grant", "expired token",
+            "token has expired", "credentials",
+        )
+    )
+
+
+def is_rate_limit(exc: BaseException) -> bool:
+    """Whether a provider exception means "slow down or you are out of quota".
+
+    Status first, then the SDK class name, then unambiguous phrases. Matching
+    on `"429" in str(exc)` alone classifies any error whose body happens to
+    contain those digits as a rate limit.
+    """
+    status = getattr(exc, "status_code", None)
+    if status is None:
+        response = getattr(exc, "response", None)
+        status = getattr(response, "status_code", None)
+    if status == 429:
+        return True
+    if type(exc).__name__ in ("RateLimitError", "ResourceExhausted",
+                              "TooManyRequests"):
+        return True
+    text = str(exc).lower()
+    return any(
+        marker in text
+        for marker in ("rate limit", "rate_limit", "insufficient_quota",
+                       "quota exceeded", "too many requests")
+    )
+
+
+def resolve_exception_types(module_name: str, *names: str) -> tuple[type, ...]:
+    """Look up exception classes on an SDK without importing it at module
+    scope — the SDKs are optional extras, so a missing one must not break
+    the import. Names the SDK does not define are skipped."""
+    import importlib
+    try:
+        module = importlib.import_module(module_name)
+    except Exception:
+        return ()
+    found = []
+    for name in names:
+        cls = getattr(module, name, None)
+        if isinstance(cls, type) and issubclass(cls, BaseException):
+            found.append(cls)
+    return tuple(found)
+
+
+@dataclass
+class Message:
+    role: str      # "user" | "assistant" | "system" | "tool"
+    content: str
+    model: str = ""
+    step: str = ""
+    tool_calls: List[ToolCall] = field(default_factory=list)
+    tool_call_id: Optional[str] = None
+
+
+@dataclass
+class ToolCall:
+    id: str
+    name: str
+    arguments: Dict[str, Any]
+
+
+class BaseAdapter(ABC):
+    """Every model adapter must implement this interface."""
+
+    provider: str = ""
+
+    @abstractmethod
+    async def complete(
+        self,
+        messages: list[Message],
+        model_id: str,
+        system: str = "",
+        tools: list[Any] | None = None,
+        temperature: float = 0.7,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ) -> str | Message:
+        """Send messages and return the full response text or Message with tool calls."""
+        ...
+
+    @abstractmethod
+    async def stream(
+        self,
+        messages: list[Message],
+        model_id: str,
+        system: str = "",
+        tools: list[Any] | None = None,
+        temperature: float = 0.7,
+        max_tokens: int = DEFAULT_STREAM_MAX_TOKENS,
+    ) -> AsyncIterator[str | ToolCall]:
+        """Stream response tokens or tool call chunks."""
+        ...
+
+    def count_tokens(self, text: str) -> int:
+        """Local token estimate. Never a network call, never exact.
+
+        Used where an async context is not available. Prefer
+        count_tokens_async, which asks the provider where the provider can
+        answer.
+        """
+        return len(text) // 4
+
+    async def count_tokens_async(self, text: str, model_id: str = "") -> int:
+        """Token count, exact where the provider offers a counting endpoint.
+
+        Defaults to the local estimate. Adapters that can do better override
+        this; they must fall back to the estimate rather than raise, since a
+        count is only ever used to budget context.
+        """
+        return self.count_tokens(text)

polymath/adapters/claude.py

@@ -0,0 +1,280 @@
+from __future__ import annotations
+
+import hashlib
+from typing import Any, AsyncIterator
+
+from polymath.adapters.base import (
+    AuthExpiredError, BaseAdapter, DEFAULT_MAX_TOKENS,
+    DEFAULT_STREAM_MAX_TOKENS, Message, ToolCall,
+    is_auth_failure, resolve_exception_types, truncation_notice,
+)
+from polymath.config import get_model
+
+
+def _auth_types() -> tuple[type, ...]:
+    return resolve_exception_types(
+        "anthropic", "AuthenticationError", "PermissionDeniedError",
+    )
+
+
+#: Rough characters-per-token for English prose. Only used when the API
+#: cannot be reached; count_tokens_async is the real answer.
+_CLAUDE_CHARS_PER_TOKEN = 4
+
+#: Model used for counting when the caller did not name one. Counting is
+#: model-specific, so this is a fallback, not a default worth relying on.
+_DEFAULT_COUNTING_MODEL = "claude-sonnet-5"
+
+#: (model, sha256 of text) -> token count. Cleared wholesale when full;
+#: counts are cheap to recompute and this is only a budget figure.
+_TOKEN_COUNT_CACHE: dict[tuple[str, str], int] = {}
+_TOKEN_COUNT_CACHE_MAX = 512
+
+
+def _accepts_temperature(model_id: str) -> bool:
+    """Whether to send a sampling temperature to this model.
+
+    Claude removed temperature/top_p/top_k with the 4.7 generation; sending
+    one returns a 400. An id we do not recognise is assumed to be a newer
+    model and gets no temperature — omitting it is valid on every Claude
+    model ever released, so unknown-means-omit is the safe default.
+    """
+    model = get_model(model_id)
+    return bool(model and model.supports_temperature)
+
+
+class ClaudeAdapter(BaseAdapter):
+    provider = "claude"
+
+    def __init__(self, api_key: str) -> None:
+        import anthropic
+        self._client = anthropic.AsyncAnthropic(api_key=api_key)
+        self._last_input_tokens: int = 0
+        self._last_output_tokens: int = 0
+
+    def _convert_tools(self, tools: list[Any] | None) -> list[dict] | None:
+        if not tools:
+            return None
+        claude_tools = []
+        for t in tools:
+            schema = t.to_json_schema()
+            claude_tools.append({
+                "name": schema["name"],
+                "description": schema["description"],
+                "input_schema": schema["parameters"]
+            })
+        return claude_tools
+
+    async def complete(
+        self,
+        messages: list[Message],
+        model_id: str,
+        system: str = "",
+        tools: list[Any] | None = None,
+        temperature: float = 0.7,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ) -> str | Message:
+        claude_msgs = []
+        for m in messages:
+            if m.role == "system":
+                continue
+            
+            content = []
+            if m.content:
+                content.append({"type": "text", "text": m.content})
+            
+            for tc in m.tool_calls:
+                content.append({
+                    "type": "tool_use",
+                    "id": tc.id,
+                    "name": tc.name,
+                    "input": tc.arguments
+                })
+            
+            if m.role == "tool":
+                # Special handling for tool results in Claude
+                claude_msgs.append({
+                    "role": "user",
+                    "content": [{
+                        "type": "tool_result",
+                        "tool_use_id": m.tool_call_id,
+                        "content": m.content
+                    }]
+                })
+                continue
+
+            claude_msgs.append({"role": m.role, "content": content})
+
+        kwargs = dict(
+            model=model_id,
+            max_tokens=max_tokens,
+            messages=claude_msgs,
+        )
+        if _accepts_temperature(model_id):
+            kwargs["temperature"] = temperature
+        if system:
+            kwargs["system"] = system
+        if tools:
+            kwargs["tools"] = self._convert_tools(tools)
+
+        try:
+            resp = await self._client.messages.create(**kwargs)
+        except Exception as e:
+            if is_auth_failure(e, _auth_types()):
+                raise AuthExpiredError("claude", "claude") from e
+            raise
+
+        if hasattr(resp, "usage") and resp.usage:
+            self._last_input_tokens = getattr(resp.usage, "input_tokens", 0)
+            self._last_output_tokens = getattr(resp.usage, "output_tokens", 0)
+
+        # Collect all tool_use blocks (regardless of stop_reason)
+        tool_calls = []
+        content_text = ""
+        for block in resp.content:
+            if block.type == "text":
+                content_text += block.text
+            elif block.type == "tool_use":
+                tool_calls.append(ToolCall(
+                    id=block.id,
+                    name=block.name,
+                    arguments=block.input,
+                ))
+
+        if getattr(resp, "stop_reason", None) == "max_tokens":
+            content_text += truncation_notice(max_tokens)
+
+        if tool_calls:
+            return Message(
+                role="assistant",
+                content=content_text,
+                tool_calls=tool_calls,
+            )
+
+        return content_text
+
+    async def stream(
+        self,
+        messages: list[Message],
+        model_id: str,
+        system: str = "",
+        tools: list[Any] | None = None,
+        temperature: float = 0.7,
+        max_tokens: int = DEFAULT_STREAM_MAX_TOKENS,
+    ) -> AsyncIterator[str | ToolCall]:
+        # Fall back to complete() when tools are present since streaming + tools is complex.
+        if tools:
+            result = await self.complete(
+                messages=messages,
+                model_id=model_id,
+                system=system,
+                tools=tools,
+                temperature=temperature,
+                max_tokens=max_tokens,
+            )
+            if isinstance(result, Message):
+                if result.content:
+                    yield result.content
+                for tc in result.tool_calls:
+                    yield tc
+            else:
+                yield result
+            return
+
+        # Build properly formatted messages for streaming (no tools)
+        claude_msgs = []
+        for m in messages:
+            if m.role == "system":
+                continue
+            if m.role == "tool":
+                claude_msgs.append({
+                    "role": "user",
+                    "content": [{
+                        "type": "tool_result",
+                        "tool_use_id": m.tool_call_id,
+                        "content": m.content,
+                    }]
+                })
+                continue
+            content = []
+            if m.content:
+                content.append({"type": "text", "text": m.content})
+            for tc in m.tool_calls:
+                content.append({
+                    "type": "tool_use",
+                    "id": tc.id,
+                    "name": tc.name,
+                    "input": tc.arguments,
+                })
+            claude_msgs.append({"role": m.role, "content": content})
+
+        kwargs = dict(
+            model=model_id,
+            max_tokens=max_tokens,
+            messages=claude_msgs,
+        )
+        if _accepts_temperature(model_id):
+            kwargs["temperature"] = temperature
+        if system:
+            kwargs["system"] = system
+
+        try:
+            async with self._client.messages.stream(**kwargs) as stream:
+                async for text in stream.text_stream:
+                    yield text
+                # The answer is already in the user's hands by this point.
+                # Reporting truncation is a nicety, so it must never be able
+                # to fail the run that just succeeded — an SDK without
+                # get_final_message would otherwise turn a complete answer
+                # into an exception.
+                try:
+                    final = await stream.get_final_message()
+                    truncated = getattr(final, "stop_reason", None) == "max_tokens"
+                except Exception:
+                    truncated = False
+                if truncated:
+                    yield truncation_notice(max_tokens)
+        except Exception as e:
+            if is_auth_failure(e, _auth_types()):
+                raise AuthExpiredError("claude", "claude") from e
+            raise
+
+    def count_tokens(self, text: str) -> int:
+        """Local estimate for Claude.
+
+        This used to load tiktoken's cl100k_base, which is OpenAI's encoding
+        and not Claude's, so the number was wrong by an unknown margin in an
+        unknown direction. There is no local Claude tokenizer; a character
+        ratio is the honest local answer, and count_tokens_async gets the
+        real one from the API.
+        """
+        return max(1, len(text) // _CLAUDE_CHARS_PER_TOKEN) if text else 0
+
+    async def count_tokens_async(self, text: str, model_id: str = "") -> int:
+        """Exact count from Anthropic's own counting endpoint.
+
+        Cached per (model, text): the context budget is recomputed on every
+        turn over a prompt that is mostly unchanged, and this would otherwise
+        be one network round trip each time.
+        """
+        if not text:
+            return 0
+        model = model_id or _DEFAULT_COUNTING_MODEL
+        key = (model, hashlib.sha256(text.encode("utf-8")).hexdigest())
+        cached = _TOKEN_COUNT_CACHE.get(key)
+        if cached is not None:
+            return cached
+        try:
+            resp = await self._client.messages.count_tokens(
+                model=model,
+                messages=[{"role": "user", "content": text}],
+            )
+            count = int(resp.input_tokens)
+        except Exception:
+            # Offline, unkeyed, or an unknown model id. A budget figure is
+            # not worth failing a run over.
+            return self.count_tokens(text)
+        if len(_TOKEN_COUNT_CACHE) >= _TOKEN_COUNT_CACHE_MAX:
+            _TOKEN_COUNT_CACHE.clear()
+        _TOKEN_COUNT_CACHE[key] = count
+        return count

polymath/adapters/gemini.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+import json
+from typing import Any, AsyncIterator
+
+from polymath.adapters.base import (
+    AuthExpiredError, BaseAdapter, DEFAULT_MAX_TOKENS,
+    DEFAULT_STREAM_MAX_TOKENS, Message, ToolCall,
+    is_auth_failure, resolve_exception_types,
+)
+
+
+def _auth_types() -> tuple[type, ...]:
+    return resolve_exception_types(
+        "google.genai.errors", "ClientError", "UnauthenticatedError",
+    )
+
+
+class GeminiAdapter(BaseAdapter):
+    provider = "gemini"
+
+    def __init__(self, api_key: str) -> None:
+        from google import genai
+        from google.genai import types
+        # OAuth tokens (from gemini CLI) start with "ya29." — use http_options
+        # to pass them as a bearer token instead of an API key.
+        if api_key.startswith("ya29."):
+            self._client = genai.Client(
+                http_options={"headers": {"Authorization": f"Bearer {api_key}"}},
+            )
+        else:
+            self._client = genai.Client(api_key=api_key)
+        self._types = types
+        self._last_input_tokens: int = 0
+        self._last_output_tokens: int = 0
+
+    def _build_contents(self, messages: list[Message]) -> list:
+        """Convert Message list to Gemini contents format."""
+        contents = []
+        for m in messages:
+            if m.role == "system":
+                continue
+            role = "user" if m.role in ("user", "tool") else "model"
+            parts = []
+
+            if m.role == "tool":
+                # Tool result — send as function_response
+                parts.append(self._types.Part(
+                    function_response=self._types.FunctionResponse(
+                        name=m.model,
+                        response={"result": m.content},
+                    )
+                ))
+            elif m.tool_calls:
+                # Assistant message with tool calls
+                if m.content:
+                    parts.append(self._types.Part(text=m.content))
+                for tc in m.tool_calls:
+                    parts.append(self._types.Part(
+                        function_call=self._types.FunctionCall(
+                            name=tc.name,
+                            args=tc.arguments,
+                        )
+                    ))
+            else:
+                parts.append(self._types.Part(text=m.content or ""))
+
+            contents.append(self._types.Content(role=role, parts=parts))
+        return contents
+
+    def _build_tools(self, tools: list[Any] | None):
+        if not tools:
+            return None
+        declarations = []
+        for t in tools:
+            schema = t.parameters.copy() if t.parameters else {}
+            declarations.append(self._types.FunctionDeclaration(
+                name=t.name,
+                description=t.description,
+                parameters=schema,
+            ))
+        return [self._types.Tool(function_declarations=declarations)]
+
+    def _parse_response(self, resp) -> str | Message:
+        try:
+            usage = resp.usage_metadata
+            if usage:
+                self._last_input_tokens = getattr(usage, "prompt_token_count", 0) or 0
+                self._last_output_tokens = getattr(usage, "candidates_token_count", 0) or 0
+        except Exception:
+            pass
+
+        try:
+            tool_calls = []
+            content_text = ""
+            for part in resp.candidates[0].content.parts:
+                if part.function_call and part.function_call.name:
+                    fc = part.function_call
+                    tool_calls.append(ToolCall(
+                        id=fc.name,
+                        name=fc.name,
+                        arguments=dict(fc.args) if fc.args else {},
+                    ))
+                elif part.text:
+                    content_text += part.text
+            if tool_calls:
+                return Message(role="assistant", content=content_text, tool_calls=tool_calls)
+            return content_text
+        except Exception:
+            pass
+
+        try:
+            return resp.text
+        except Exception:
+            return ""
+
+    async def complete(
+        self,
+        messages: list[Message],
+        model_id: str,
+        system: str = "",
+        tools: list[Any] | None = None,
+        temperature: float = 0.7,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ) -> str | Message:
+        contents = self._build_contents(messages)
+        config = self._types.GenerateContentConfig(
+            system_instruction=system or None,
+            temperature=temperature,
+            max_output_tokens=max_tokens,
+            tools=self._build_tools(tools),
+        )
+        try:
+            resp = await self._client.aio.models.generate_content(
+                model=model_id,
+                contents=contents,
+                config=config,
+            )
+        except Exception as e:
+            if is_auth_failure(e):
+                raise AuthExpiredError("gemini", "gemini auth login") from e
+            raise
+        return self._parse_response(resp)
+
+    async def stream(
+        self,
+        messages: list[Message],
+        model_id: str,
+        system: str = "",
+        tools: list[Any] | None = None,
+        temperature: float = 0.7,
+        max_tokens: int = DEFAULT_STREAM_MAX_TOKENS,
+    ) -> AsyncIterator[str | ToolCall]:
+        # If tools requested, fall back to complete() — streaming + tools is complex
+        if tools:
+            result = await self.complete(messages, model_id, system, tools, temperature, max_tokens)
+            if isinstance(result, Message):
+                if result.content:
+                    yield result.content
+                for tc in result.tool_calls:
+                    yield tc
+            else:
+                yield result
+            return
+
+        contents = self._build_contents(messages)
+        config = self._types.GenerateContentConfig(
+            system_instruction=system or None,
+            temperature=temperature,
+            max_output_tokens=max_tokens,
+        )
+        try:
+            async for chunk in await self._client.aio.models.generate_content_stream(
+                model=model_id,
+                contents=contents,
+                config=config,
+            ):
+                try:
+                    if chunk.text:
+                        yield chunk.text
+                except Exception:
+                    pass
+        except Exception as e:
+            if is_auth_failure(e):
+                raise AuthExpiredError("gemini", "gemini auth login") from e
+            raise