Pinscope open-source core
Agentic schematic validation: datasheet extraction via Claude Console Skills, netlist/BOM design graph, per-IC direct datasheet review with page citations, capacitor derating, Next.js report UI. Extracted from the Pinscope cloud codebase. Auth and billing live in the private gateway repo behind stable seams (billing_hook.py, adapter files listed in CLAUDE.md).
This commit is contained in:
Siddharth Kothari
@@ -0,0 +1,36 @@
|
||||
"""Provider-agnostic LLM client layer.
|
||||
|
||||
All Claude API calls in the backend route through this package via the
|
||||
``LLMProvider`` interface. The default provider is Anthropic; per-stage
|
||||
overrides via ``Settings.provider_*`` env vars route specific stages to
|
||||
other providers (currently Anthropic + Gemini).
|
||||
"""
|
||||
|
||||
from backend.services.llm.factory import call_with_fallback, get_provider
|
||||
from backend.services.llm.types import (
|
||||
Completion,
|
||||
ContentBlock,
|
||||
Message,
|
||||
PdfBlock,
|
||||
TextBlock,
|
||||
ToolCall,
|
||||
ToolChoice,
|
||||
ToolResultBlock,
|
||||
ToolSchema,
|
||||
Usage,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Completion",
|
||||
"ContentBlock",
|
||||
"Message",
|
||||
"PdfBlock",
|
||||
"TextBlock",
|
||||
"ToolCall",
|
||||
"ToolChoice",
|
||||
"ToolResultBlock",
|
||||
"ToolSchema",
|
||||
"Usage",
|
||||
"call_with_fallback",
|
||||
"get_provider",
|
||||
]
|
||||
@@ -0,0 +1,363 @@
|
||||
"""Anthropic provider — wraps AsyncAnthropic + Console Skills.
|
||||
|
||||
Translates the unified ``Message`` / ``Completion`` shapes into Anthropic's
|
||||
native message-block format and back. Caching is per-block via
|
||||
``cache_control: ephemeral``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import re
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import anthropic
|
||||
|
||||
from backend.config import settings
|
||||
from backend.services.llm.base import LLMProvider, LLMSession
|
||||
from backend.services.llm.types import (
|
||||
Completion,
|
||||
ContentBlock,
|
||||
Message,
|
||||
PdfBlock,
|
||||
TextBlock,
|
||||
ToolCall,
|
||||
ToolChoice,
|
||||
ToolResultBlock,
|
||||
ToolSchema,
|
||||
Usage,
|
||||
)
|
||||
|
||||
|
||||
_SKILL_MAX_TURNS = 10
|
||||
|
||||
# Sampling params were removed on newer Claude models (Sonnet 5, Opus 4.7+,
|
||||
# Fable/Mythos 5) — sending `temperature` returns 400 "`temperature` is
|
||||
# deprecated for this model". Allowlist the families that still accept it so
|
||||
# unknown/future models fail safe (omit → default sampling) instead of
|
||||
# 400-ing every call in the session.
|
||||
_TEMPERATURE_OK = re.compile(r"^claude-(3-|opus-4-[0-6]|sonnet-4-|haiku-)")
|
||||
|
||||
|
||||
def _model_accepts_temperature(model: str) -> bool:
|
||||
return bool(_TEMPERATURE_OK.match(model))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Translation helpers — unified types ↔ Anthropic dicts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _encode_pdf_block(path: Path | str, *, cache: bool) -> dict:
|
||||
data = base64.standard_b64encode(Path(path).read_bytes()).decode()
|
||||
block: dict = {
|
||||
"type": "document",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": data},
|
||||
}
|
||||
if cache:
|
||||
block["cache_control"] = {"type": "ephemeral"}
|
||||
return block
|
||||
|
||||
|
||||
def _to_anthropic_block(b: ContentBlock) -> dict:
|
||||
if isinstance(b, TextBlock):
|
||||
d: dict = {"type": "text", "text": b.text}
|
||||
if b.cacheable:
|
||||
d["cache_control"] = {"type": "ephemeral"}
|
||||
return d
|
||||
if isinstance(b, PdfBlock):
|
||||
return _encode_pdf_block(b.path, cache=b.cacheable)
|
||||
if isinstance(b, ToolCall):
|
||||
return {"type": "tool_use", "id": b.id, "name": b.name, "input": b.input}
|
||||
if isinstance(b, ToolResultBlock):
|
||||
return {
|
||||
"type": "tool_result",
|
||||
"tool_use_id": b.tool_use_id,
|
||||
"content": b.content,
|
||||
}
|
||||
raise TypeError(f"Unknown ContentBlock: {type(b).__name__}")
|
||||
|
||||
|
||||
def _to_anthropic_message(m: Message) -> dict:
|
||||
return {"role": m.role, "content": [_to_anthropic_block(b) for b in m.content]}
|
||||
|
||||
|
||||
# Anthropic allows at most 4 cache_control breakpoints per request. The system
|
||||
# prompt always consumes one (see AnthropicSession.complete), leaving 3 for
|
||||
# message content. A multi-turn review attaches a cacheable PDF for each
|
||||
# get_datasheet_excerpt fetch (validation_tools.py), so a hub IC that verifies
|
||||
# two interface excerpts produced 5 breakpoints — system + initial PDF + initial
|
||||
# context + 2 excerpts — and the API rejected the request with
|
||||
# "A maximum of 4 blocks with cache_control may be provided. Found 5."
|
||||
#
|
||||
# Cap the message-block breakpoints in the translated request, keeping the most
|
||||
# valuable ones: the first cacheable block (the full-datasheet anchor — a stable,
|
||||
# guaranteed cache hit every turn) plus the two most recent (incremental caching
|
||||
# of the growing tail). Any caller-set cache_control beyond that is dropped.
|
||||
_MAX_MESSAGE_CACHE_BREAKPOINTS = 3
|
||||
|
||||
|
||||
def _enforce_cache_breakpoint_limit(messages: list[dict]) -> None:
|
||||
"""Strip excess cache_control markers from message blocks in place so that
|
||||
system(1) + message breakpoints never exceed Anthropic's per-request limit."""
|
||||
marked: list[dict] = []
|
||||
for m in messages:
|
||||
content = m.get("content")
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
for block in content:
|
||||
if isinstance(block, dict) and "cache_control" in block:
|
||||
marked.append(block)
|
||||
if len(marked) <= _MAX_MESSAGE_CACHE_BREAKPOINTS:
|
||||
return
|
||||
keep = {id(marked[0]), id(marked[-1]), id(marked[-2])}
|
||||
for block in marked:
|
||||
if id(block) not in keep:
|
||||
block.pop("cache_control", None)
|
||||
|
||||
|
||||
def _to_anthropic_tool(t: ToolSchema) -> dict:
|
||||
return {"name": t.name, "description": t.description, "input_schema": t.input_schema}
|
||||
|
||||
|
||||
def _to_anthropic_tool_choice(c: ToolChoice) -> dict:
|
||||
if c == "auto":
|
||||
return {"type": "auto"}
|
||||
if c == "none":
|
||||
return {"type": "none"}
|
||||
if isinstance(c, dict) and "name" in c:
|
||||
return {"type": "tool", "name": c["name"]}
|
||||
raise ValueError(f"Invalid tool_choice: {c!r}")
|
||||
|
||||
|
||||
def _from_anthropic_response(resp) -> Completion:
|
||||
"""Parse an Anthropic message response into a unified Completion."""
|
||||
text_parts: list[str] = []
|
||||
tool_calls: list[ToolCall] = []
|
||||
raw_blocks: list[ContentBlock] = []
|
||||
|
||||
for block in resp.content:
|
||||
btype = getattr(block, "type", None)
|
||||
if btype == "text":
|
||||
text_parts.append(block.text)
|
||||
raw_blocks.append(TextBlock(text=block.text))
|
||||
elif btype == "tool_use":
|
||||
tc = ToolCall(id=block.id, name=block.name, input=dict(block.input))
|
||||
tool_calls.append(tc)
|
||||
raw_blocks.append(tc)
|
||||
# Other block types (server tool calls etc.) are pass-through ignored
|
||||
|
||||
usage = Usage(
|
||||
input_tokens=resp.usage.input_tokens,
|
||||
output_tokens=resp.usage.output_tokens,
|
||||
cache_creation_tokens=getattr(resp.usage, "cache_creation_input_tokens", 0) or 0,
|
||||
cache_read_tokens=getattr(resp.usage, "cache_read_input_tokens", 0) or 0,
|
||||
)
|
||||
|
||||
return Completion(
|
||||
text="".join(text_parts),
|
||||
tool_calls=tool_calls,
|
||||
usage=usage,
|
||||
stop_reason=resp.stop_reason or "unknown",
|
||||
raw_assistant_blocks=raw_blocks,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AnthropicSession(LLMSession):
|
||||
provider_name = "anthropic"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client: anthropic.AsyncAnthropic,
|
||||
model: str,
|
||||
system: str,
|
||||
max_tokens: int,
|
||||
temperature: float | None = None,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self.model = model
|
||||
self._system = system
|
||||
self._max_tokens = max_tokens
|
||||
self._temperature = temperature
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
*,
|
||||
messages: list[Message],
|
||||
tools: list[ToolSchema] | None = None,
|
||||
tool_choice: ToolChoice = "auto",
|
||||
) -> Completion:
|
||||
kwargs: dict = {
|
||||
"model": self.model,
|
||||
"max_tokens": self._max_tokens,
|
||||
"system": [{
|
||||
"type": "text",
|
||||
"text": self._system,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}],
|
||||
"messages": [_to_anthropic_message(m) for m in messages],
|
||||
}
|
||||
_enforce_cache_breakpoint_limit(kwargs["messages"])
|
||||
if self._temperature is not None and _model_accepts_temperature(self.model):
|
||||
kwargs["temperature"] = self._temperature
|
||||
if tools:
|
||||
kwargs["tools"] = [_to_anthropic_tool(t) for t in tools]
|
||||
kwargs["tool_choice"] = _to_anthropic_tool_choice(tool_choice)
|
||||
|
||||
# Streaming, not create(): SDK 0.83+ raises ValueError pre-flight on
|
||||
# `messages.create` whenever max_tokens crosses ~21k for Sonnet
|
||||
# (the "may take longer than 10 minutes" guard). Review uses 32k
|
||||
# max_tokens for Gemini thinking headroom; streaming bypasses that
|
||||
# client-side timeout cap. get_final_message() returns the same
|
||||
# shape as create(), so _from_anthropic_response is reused as-is.
|
||||
async with self._client.messages.stream(**kwargs) as stream:
|
||||
resp = await stream.get_final_message()
|
||||
return _from_anthropic_response(resp)
|
||||
|
||||
async def close(self) -> None:
|
||||
# Anthropic ephemeral cache cleans up on its own (5-min TTL).
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Provider
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AnthropicProvider(LLMProvider):
|
||||
name = "anthropic"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._client = anthropic.AsyncAnthropic(api_key=settings.anthropic_api_key)
|
||||
|
||||
async def create_session(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
system: str,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float | None = None,
|
||||
) -> LLMSession:
|
||||
return AnthropicSession(
|
||||
client=self._client,
|
||||
model=model,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
async def run_skill(
|
||||
self,
|
||||
*,
|
||||
skill_name: str,
|
||||
model: str,
|
||||
system: str,
|
||||
user_text: str,
|
||||
pdf_path: str | None,
|
||||
output_tool: ToolSchema,
|
||||
) -> tuple[dict, Completion]:
|
||||
"""Anthropic Console Skills — multi-turn skill execution with the
|
||||
``skills-2025-10-02`` + ``code-execution-2025-08-25`` betas.
|
||||
|
||||
Skill mounts in a per-call container; the model reads ``SKILL.md``,
|
||||
runs ``validate.py`` server-side via code_execution, and voluntarily
|
||||
calls ``output_tool`` once it has well-formed data.
|
||||
"""
|
||||
skill_id, version = settings.get_skill(skill_name)
|
||||
|
||||
# Build initial user content
|
||||
user_content: list[dict] = []
|
||||
if pdf_path:
|
||||
user_content.append(_encode_pdf_block(pdf_path, cache=True))
|
||||
user_content.append({"type": "text", "text": user_text})
|
||||
|
||||
messages: list[dict] = [{"role": "user", "content": user_content}]
|
||||
container: dict = {
|
||||
"skills": [{
|
||||
"type": "custom",
|
||||
"skill_id": skill_id,
|
||||
"version": version,
|
||||
}],
|
||||
}
|
||||
|
||||
total_input = 0
|
||||
total_output = 0
|
||||
total_cache_creation = 0
|
||||
total_cache_read = 0
|
||||
t0 = time.monotonic()
|
||||
last_resp = None
|
||||
|
||||
for turn in range(_SKILL_MAX_TURNS):
|
||||
resp = await self._client.beta.messages.create(
|
||||
model=model,
|
||||
max_tokens=16384,
|
||||
system=[{
|
||||
"type": "text",
|
||||
"text": system,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}],
|
||||
tools=[
|
||||
{"type": "code_execution_20250825", "name": "code_execution"},
|
||||
_to_anthropic_tool(output_tool),
|
||||
],
|
||||
container=container,
|
||||
messages=messages,
|
||||
betas=["skills-2025-10-02", "code-execution-2025-08-25"],
|
||||
)
|
||||
last_resp = resp
|
||||
|
||||
total_input += resp.usage.input_tokens
|
||||
total_output += resp.usage.output_tokens
|
||||
total_cache_creation += getattr(resp.usage, "cache_creation_input_tokens", 0) or 0
|
||||
total_cache_read += getattr(resp.usage, "cache_read_input_tokens", 0) or 0
|
||||
|
||||
# Reuse container for subsequent turns
|
||||
if hasattr(resp, "container") and resp.container:
|
||||
container = {"id": resp.container.id}
|
||||
|
||||
for block in resp.content:
|
||||
if (
|
||||
getattr(block, "type", None) == "tool_use"
|
||||
and block.name == output_tool.name
|
||||
):
|
||||
completion = Completion(
|
||||
text="",
|
||||
tool_calls=[ToolCall(id=block.id, name=block.name, input=dict(block.input))],
|
||||
usage=Usage(
|
||||
input_tokens=total_input,
|
||||
output_tokens=total_output,
|
||||
cache_creation_tokens=total_cache_creation,
|
||||
cache_read_tokens=total_cache_read,
|
||||
),
|
||||
stop_reason=resp.stop_reason or "unknown",
|
||||
)
|
||||
# Stash turns count via attribute for callers that need it
|
||||
completion.turns = turn + 1 # type: ignore[attr-defined]
|
||||
completion.duration_ms = int((time.monotonic() - t0) * 1000) # type: ignore[attr-defined]
|
||||
return dict(block.input), completion
|
||||
|
||||
messages.append({"role": "assistant", "content": resp.content})
|
||||
|
||||
if resp.stop_reason == "pause_turn":
|
||||
continue
|
||||
|
||||
if resp.stop_reason == "end_turn":
|
||||
messages.append({
|
||||
"role": "user",
|
||||
"content": f"Please call {output_tool.name} with the extracted data.",
|
||||
})
|
||||
continue
|
||||
|
||||
# tool_use from code_execution — let the loop continue
|
||||
continue
|
||||
|
||||
raise RuntimeError(
|
||||
f"Skill {skill_name!r} did not produce {output_tool.name} "
|
||||
f"in {_SKILL_MAX_TURNS} turns"
|
||||
)
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Abstract LLMProvider + LLMSession interfaces."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Protocol
|
||||
|
||||
from backend.services.llm.types import (
|
||||
Completion,
|
||||
Message,
|
||||
ToolChoice,
|
||||
ToolSchema,
|
||||
)
|
||||
|
||||
|
||||
class LLMSession(ABC):
|
||||
"""A multi-turn conversation with provider-specific cache lifecycle.
|
||||
|
||||
Lifecycle::
|
||||
|
||||
session = await provider.create_session(model=..., system=...)
|
||||
try:
|
||||
messages = [Message("user", [
|
||||
PdfBlock(path, cacheable=True),
|
||||
TextBlock(context, cacheable=True),
|
||||
])]
|
||||
for turn in range(N):
|
||||
completion = await session.complete(
|
||||
messages=messages, tools=..., tool_choice=...,
|
||||
)
|
||||
# process tool_calls, append to messages, repeat
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
Caching: blocks with ``cacheable=True`` participate in provider caching.
|
||||
Anthropic stamps ``cache_control: ephemeral`` on each cacheable block on
|
||||
every call. Gemini collects all cacheable blocks (plus the system prompt)
|
||||
on the first ``complete()`` call into a ``CachedContent`` object and
|
||||
references it on subsequent calls. The system prompt is always cached.
|
||||
"""
|
||||
|
||||
provider_name: str
|
||||
"""Provider identifier ("anthropic", "gemini") — used for api_logs."""
|
||||
model: str
|
||||
|
||||
@abstractmethod
|
||||
async def complete(
|
||||
self,
|
||||
*,
|
||||
messages: list[Message],
|
||||
tools: list[ToolSchema] | None = None,
|
||||
tool_choice: ToolChoice = "auto",
|
||||
) -> Completion:
|
||||
"""Run one inference turn."""
|
||||
|
||||
@abstractmethod
|
||||
async def close(self) -> None:
|
||||
"""Release any provider-side resources (e.g. delete a cache object).
|
||||
Safe to call multiple times."""
|
||||
|
||||
|
||||
class LLMProvider(Protocol):
|
||||
"""Top-level provider interface."""
|
||||
|
||||
name: str
|
||||
|
||||
async def create_session(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
system: str,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float | None = None,
|
||||
) -> LLMSession:
|
||||
"""Construct a session. ``system`` is always cached by the session.
|
||||
|
||||
``temperature`` — if not None, applied to every ``complete()`` call on
|
||||
this session. ``None`` means use the provider's default. Set to 0.0
|
||||
for deterministic-as-possible behavior in agentic loops where the same
|
||||
inputs should produce the same outputs."""
|
||||
...
|
||||
|
||||
async def run_skill(
|
||||
self,
|
||||
*,
|
||||
skill_name: str,
|
||||
model: str,
|
||||
system: str,
|
||||
user_text: str,
|
||||
pdf_path: str | None,
|
||||
output_tool: ToolSchema,
|
||||
) -> tuple[dict, "Completion"]:
|
||||
"""Execute a managed Skill and return (forced-tool input, Completion).
|
||||
|
||||
Anthropic uses Console Skills (skill_id + container + code_execution
|
||||
beta). Gemini raises ``NotImplementedError`` — there is no
|
||||
Gemini-managed-Skill equivalent today; if you want a Gemini path for
|
||||
skill-style extraction, inline the SKILL.md content as ``system`` and
|
||||
run validation locally."""
|
||||
...
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Provider factory + per-stage routing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
from typing import Awaitable, Callable, TypeVar
|
||||
|
||||
from backend.config import settings
|
||||
from backend.services.llm.base import LLMProvider
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@lru_cache(maxsize=4)
|
||||
def get_provider_by_name(name: str) -> LLMProvider:
|
||||
"""Return a singleton provider instance for ``name`` ("anthropic" |
|
||||
"gemini"). Used by :func:`get_provider` and :func:`call_with_fallback`."""
|
||||
if name == "anthropic":
|
||||
from backend.services.llm.anthropic_provider import AnthropicProvider
|
||||
return AnthropicProvider()
|
||||
if name == "gemini":
|
||||
from backend.services.llm.gemini_provider import GeminiProvider
|
||||
return GeminiProvider()
|
||||
raise ValueError(f"Unknown LLM provider: {name!r}")
|
||||
|
||||
|
||||
# Backwards-compatible alias
|
||||
_get_provider_by_name = get_provider_by_name
|
||||
|
||||
|
||||
def get_provider(stage: str) -> LLMProvider:
|
||||
"""Return the provider configured for ``stage``.
|
||||
|
||||
Falls back to ``settings.provider_default`` if no per-stage override.
|
||||
Providers are cached per-name, so repeated calls return the same
|
||||
instance (and share the underlying SDK client)."""
|
||||
name = settings.provider_for_stage(stage)
|
||||
return get_provider_by_name(name)
|
||||
|
||||
|
||||
async def call_with_fallback(
|
||||
stage: str,
|
||||
body: Callable[[LLMProvider, str], Awaitable[T]],
|
||||
) -> T:
|
||||
"""Run ``body(provider, model)`` for ``stage``; on any exception,
|
||||
retry once with the fallback provider/model if one is configured via
|
||||
``FALLBACK_PROVIDER_<STAGE>`` / ``FALLBACK_MODEL_<STAGE>``.
|
||||
|
||||
The fallback runs ``body`` from scratch — any tokens spent in the
|
||||
primary attempt are lost (and not logged). ``asyncio.CancelledError``
|
||||
is always re-raised so cancellation still works.
|
||||
"""
|
||||
primary_provider = get_provider(stage)
|
||||
primary_model = settings.model_for_stage(stage)
|
||||
try:
|
||||
return await body(primary_provider, primary_model)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
fb = settings.fallback_for_stage(stage)
|
||||
if fb is None:
|
||||
raise
|
||||
log.warning(
|
||||
"[%s] primary %s/%s failed (%s) — falling back to %s/%s",
|
||||
stage, primary_provider.name, primary_model,
|
||||
exc, fb[0], fb[1],
|
||||
)
|
||||
fallback_provider = get_provider_by_name(fb[0])
|
||||
return await body(fallback_provider, fb[1])
|
||||
@@ -0,0 +1,379 @@
|
||||
"""Gemini provider — wraps google-genai async client.
|
||||
|
||||
Translates the unified ``Message`` / ``Completion`` shapes into Gemini's
|
||||
native ``Content`` / ``Part`` format. Caching uses ``CachedContent``: on the
|
||||
first ``complete()`` call, cacheable blocks (system + any block flagged
|
||||
``cacheable=True`` in the first user message) are uploaded as a
|
||||
``CachedContent`` with TTL=30min; subsequent calls reference the cache by
|
||||
name. On ``close()`` the cache is deleted. If creation fails (e.g.
|
||||
sub-threshold token count), the session falls back to inline content with no
|
||||
caching for the remainder of the conversation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from google import genai
|
||||
from google.genai import types as gtypes
|
||||
|
||||
from backend.config import settings
|
||||
from backend.services.llm.base import LLMProvider, LLMSession
|
||||
from backend.services.llm.types import (
|
||||
Completion,
|
||||
ContentBlock,
|
||||
Message,
|
||||
PdfBlock,
|
||||
TextBlock,
|
||||
ToolCall,
|
||||
ToolChoice,
|
||||
ToolResultBlock,
|
||||
ToolSchema,
|
||||
Usage,
|
||||
)
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
_CACHE_TTL = "1800s" # 30 min — covers our longest agent loop with margin
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Translation helpers — unified types ↔ Gemini Parts/Contents
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _block_to_part(b: ContentBlock) -> gtypes.Part:
|
||||
if isinstance(b, TextBlock):
|
||||
return gtypes.Part(
|
||||
text=b.text,
|
||||
thought_signature=b.thought_signature,
|
||||
)
|
||||
if isinstance(b, PdfBlock):
|
||||
return gtypes.Part(
|
||||
inline_data=gtypes.Blob(
|
||||
mime_type="application/pdf",
|
||||
data=Path(b.path).read_bytes(),
|
||||
),
|
||||
)
|
||||
if isinstance(b, ToolCall):
|
||||
return gtypes.Part(
|
||||
function_call=gtypes.FunctionCall(
|
||||
id=b.id or None,
|
||||
name=b.name,
|
||||
args=b.input,
|
||||
),
|
||||
thought_signature=b.thought_signature,
|
||||
)
|
||||
if isinstance(b, ToolResultBlock):
|
||||
return gtypes.Part(
|
||||
function_response=gtypes.FunctionResponse(
|
||||
id=b.tool_use_id or None,
|
||||
name=b.name,
|
||||
# FunctionResponse.response is a dict — wrap string content
|
||||
response={"result": b.content},
|
||||
),
|
||||
)
|
||||
raise TypeError(f"Unknown ContentBlock: {type(b).__name__}")
|
||||
|
||||
|
||||
def _message_to_content(m: Message) -> gtypes.Content:
|
||||
# Gemini uses "user" and "model" (not "assistant")
|
||||
role = "model" if m.role == "assistant" else "user"
|
||||
return gtypes.Content(
|
||||
role=role,
|
||||
parts=[_block_to_part(b) for b in m.content],
|
||||
)
|
||||
|
||||
|
||||
def _tool_to_function_declaration(t: ToolSchema) -> gtypes.FunctionDeclaration:
|
||||
return gtypes.FunctionDeclaration(
|
||||
name=t.name,
|
||||
description=t.description,
|
||||
parameters_json_schema=t.input_schema,
|
||||
)
|
||||
|
||||
|
||||
def _tools_to_gemini(tools: list[ToolSchema]) -> list[gtypes.Tool]:
|
||||
return [
|
||||
gtypes.Tool(
|
||||
function_declarations=[_tool_to_function_declaration(t) for t in tools],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _tool_choice_to_config(c: ToolChoice) -> gtypes.ToolConfig:
|
||||
if c == "auto":
|
||||
return gtypes.ToolConfig(
|
||||
function_calling_config=gtypes.FunctionCallingConfig(mode="AUTO"),
|
||||
)
|
||||
if c == "none":
|
||||
return gtypes.ToolConfig(
|
||||
function_calling_config=gtypes.FunctionCallingConfig(mode="NONE"),
|
||||
)
|
||||
if isinstance(c, dict) and "name" in c:
|
||||
return gtypes.ToolConfig(
|
||||
function_calling_config=gtypes.FunctionCallingConfig(
|
||||
mode="ANY",
|
||||
allowed_function_names=[c["name"]],
|
||||
),
|
||||
)
|
||||
raise ValueError(f"Invalid tool_choice: {c!r}")
|
||||
|
||||
|
||||
def _from_gemini_response(resp: Any) -> Completion:
|
||||
"""Parse a Gemini GenerateContentResponse into a unified Completion."""
|
||||
text_parts: list[str] = []
|
||||
tool_calls: list[ToolCall] = []
|
||||
raw_blocks: list[ContentBlock] = []
|
||||
stop_reason = "unknown"
|
||||
|
||||
candidates = getattr(resp, "candidates", None) or []
|
||||
if candidates:
|
||||
cand = candidates[0]
|
||||
finish = getattr(cand, "finish_reason", None)
|
||||
if finish:
|
||||
stop_reason = str(finish).lower().split(".")[-1]
|
||||
content = getattr(cand, "content", None)
|
||||
if content and content.parts:
|
||||
for part in content.parts:
|
||||
# Preserve thought_signature (Gemini 3 thinking-mode) for
|
||||
# exact replay on subsequent turns; missing signatures cause
|
||||
# 400 INVALID_ARGUMENT on the next call.
|
||||
sig = getattr(part, "thought_signature", None)
|
||||
if getattr(part, "text", None):
|
||||
text_parts.append(part.text)
|
||||
raw_blocks.append(TextBlock(
|
||||
text=part.text, thought_signature=sig,
|
||||
))
|
||||
elif getattr(part, "function_call", None):
|
||||
fc = part.function_call
|
||||
tc = ToolCall(
|
||||
id=fc.id or f"{fc.name}_{len(tool_calls)}",
|
||||
name=fc.name,
|
||||
input=dict(fc.args or {}),
|
||||
thought_signature=sig,
|
||||
)
|
||||
tool_calls.append(tc)
|
||||
raw_blocks.append(tc)
|
||||
|
||||
usage_md = getattr(resp, "usage_metadata", None)
|
||||
if usage_md is not None:
|
||||
prompt_tokens = usage_md.prompt_token_count or 0
|
||||
cached_tokens = usage_md.cached_content_token_count or 0
|
||||
# Gemini reports prompt_token_count as the TOTAL prompt tokens —
|
||||
# cached tokens are billed at the cache-read rate, the rest at the
|
||||
# input rate. Subtract so they don't double-count.
|
||||
non_cached = max(0, prompt_tokens - cached_tokens)
|
||||
# Thinking-mode models (2.5 Pro, 3 series) report reasoning tokens
|
||||
# in thoughts_token_count, billed at the output rate. Fold into
|
||||
# output_tokens so cost accounting matches Gemini's actual bill.
|
||||
thoughts_tokens = getattr(usage_md, "thoughts_token_count", 0) or 0
|
||||
usage = Usage(
|
||||
input_tokens=non_cached,
|
||||
output_tokens=(usage_md.candidates_token_count or 0) + thoughts_tokens,
|
||||
cache_creation_tokens=0, # Gemini doesn't expose this separately
|
||||
cache_read_tokens=cached_tokens,
|
||||
)
|
||||
else:
|
||||
usage = Usage()
|
||||
|
||||
return Completion(
|
||||
text="".join(text_parts),
|
||||
tool_calls=tool_calls,
|
||||
usage=usage,
|
||||
stop_reason=stop_reason,
|
||||
raw_assistant_blocks=raw_blocks,
|
||||
)
|
||||
|
||||
|
||||
def _is_first_user_message_fully_cacheable(messages: list[Message]) -> bool:
|
||||
"""We cache only when EVERY block in the very first user message is
|
||||
flagged cacheable. This matches our actual usage (validation + power
|
||||
tree both pass entirely cacheable initial messages) and avoids brittle
|
||||
partial-cache scenarios."""
|
||||
if not messages:
|
||||
return False
|
||||
first = messages[0]
|
||||
if first.role != "user" or not first.content:
|
||||
return False
|
||||
return all(
|
||||
isinstance(b, (TextBlock, PdfBlock)) and b.cacheable
|
||||
for b in first.content
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class GeminiSession(LLMSession):
|
||||
provider_name = "gemini"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client: genai.Client,
|
||||
model: str,
|
||||
system: str,
|
||||
max_tokens: int,
|
||||
temperature: float | None = None,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self.model = model
|
||||
self._system = system
|
||||
self._max_tokens = max_tokens
|
||||
self._temperature = temperature
|
||||
self._cache_name: str | None = None
|
||||
self._cache_attempted = False
|
||||
|
||||
async def _try_create_cache(self, first_msg: Message) -> str | None:
|
||||
"""Attempt to create a CachedContent from system + first user message.
|
||||
Returns the cache name on success, None on failure."""
|
||||
try:
|
||||
parts = [_block_to_part(b) for b in first_msg.content]
|
||||
cache = await self._client.aio.caches.create(
|
||||
model=self.model,
|
||||
config=gtypes.CreateCachedContentConfig(
|
||||
system_instruction=self._system,
|
||||
contents=[gtypes.Content(role="user", parts=parts)],
|
||||
ttl=_CACHE_TTL,
|
||||
),
|
||||
)
|
||||
log.info(
|
||||
"Gemini cache created (%s, model=%s, ttl=%s)",
|
||||
cache.name, self.model, _CACHE_TTL,
|
||||
)
|
||||
return cache.name
|
||||
except Exception as exc:
|
||||
log.info(
|
||||
"Gemini cache creation skipped (%s) — falling back to inline",
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
*,
|
||||
messages: list[Message],
|
||||
tools: list[ToolSchema] | None = None,
|
||||
tool_choice: ToolChoice = "auto",
|
||||
) -> Completion:
|
||||
if not messages:
|
||||
raise ValueError("Gemini complete() requires at least one message")
|
||||
|
||||
# First call: decide whether to cache
|
||||
if not self._cache_attempted:
|
||||
self._cache_attempted = True
|
||||
if _is_first_user_message_fully_cacheable(messages):
|
||||
self._cache_name = await self._try_create_cache(messages[0])
|
||||
|
||||
# Build per-call contents
|
||||
if self._cache_name:
|
||||
# Skip the cached first message — its contents are in the cache
|
||||
contents = [_message_to_content(m) for m in messages[1:]]
|
||||
else:
|
||||
contents = [_message_to_content(m) for m in messages]
|
||||
|
||||
# Build config
|
||||
config_kwargs: dict[str, Any] = {
|
||||
"max_output_tokens": self._max_tokens,
|
||||
}
|
||||
if self._temperature is not None:
|
||||
config_kwargs["temperature"] = self._temperature
|
||||
if self._cache_name:
|
||||
config_kwargs["cached_content"] = self._cache_name
|
||||
else:
|
||||
config_kwargs["system_instruction"] = self._system
|
||||
if tools:
|
||||
config_kwargs["tools"] = _tools_to_gemini(tools)
|
||||
config_kwargs["tool_config"] = _tool_choice_to_config(tool_choice)
|
||||
|
||||
config = gtypes.GenerateContentConfig(**config_kwargs)
|
||||
|
||||
# When using cached_content, Gemini still requires non-empty contents.
|
||||
# If the cached path leaves us with no per-call contents (only happens
|
||||
# on the very first turn with a cached initial message), seed with a
|
||||
# minimal continuation prompt.
|
||||
if self._cache_name and not contents:
|
||||
contents = [gtypes.Content(role="user", parts=[gtypes.Part(text="Continue.")])]
|
||||
|
||||
try:
|
||||
resp = await self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=contents,
|
||||
config=config,
|
||||
)
|
||||
except Exception as exc:
|
||||
# Cache may have expired mid-loop — drop it and retry inline once
|
||||
if self._cache_name and "cache" in str(exc).lower():
|
||||
log.warning("Gemini cache failed (%s) — retrying inline", exc)
|
||||
self._cache_name = None
|
||||
return await self.complete(
|
||||
messages=messages, tools=tools, tool_choice=tool_choice,
|
||||
)
|
||||
raise
|
||||
|
||||
return _from_gemini_response(resp)
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._cache_name:
|
||||
try:
|
||||
await self._client.aio.caches.delete(name=self._cache_name)
|
||||
except Exception as exc:
|
||||
log.warning("Gemini cache delete failed (%s): %s", self._cache_name, exc)
|
||||
finally:
|
||||
self._cache_name = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Provider
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class GeminiProvider(LLMProvider):
|
||||
name = "gemini"
|
||||
|
||||
def __init__(self) -> None:
|
||||
api_key = settings.gemini_api_key
|
||||
if not api_key:
|
||||
raise RuntimeError(
|
||||
"GEMINI_API_KEY is not set. Either set it in .env or route "
|
||||
"this stage to Anthropic via PROVIDER_<STAGE>=anthropic."
|
||||
)
|
||||
self._client = genai.Client(api_key=api_key)
|
||||
|
||||
async def create_session(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
system: str,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float | None = None,
|
||||
) -> LLMSession:
|
||||
return GeminiSession(
|
||||
client=self._client,
|
||||
model=model,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
async def run_skill(
|
||||
self,
|
||||
*,
|
||||
skill_name: str,
|
||||
model: str,
|
||||
system: str,
|
||||
user_text: str,
|
||||
pdf_path: str | None,
|
||||
output_tool: ToolSchema,
|
||||
) -> tuple[dict, Completion]:
|
||||
raise NotImplementedError(
|
||||
f"GeminiProvider.run_skill() not implemented (skill={skill_name!r}). "
|
||||
f"Anthropic Console Skills have no Gemini equivalent. To migrate "
|
||||
f"this skill to Gemini, inline its SKILL.md as the system prompt "
|
||||
f"and run validate.py locally."
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Per-provider pricing tables and cost computation.
|
||||
|
||||
Replaces the flat ``PRICING`` dict that used to live in
|
||||
``backend/services/api_logs.py``. Indexed by (provider, model).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
# Per-million-token USD rates. Source-of-truth links:
|
||||
# Anthropic: https://docs.anthropic.com/en/docs/about-claude/pricing
|
||||
# Google: https://ai.google.dev/pricing
|
||||
# Last updated: 2026-07-01
|
||||
PRICING: dict[str, dict[str, dict[str, float]]] = {
|
||||
"anthropic": {
|
||||
"claude-opus-4-6": {"input": 5.00, "output": 25.00},
|
||||
"claude-opus-4-5": {"input": 5.00, "output": 25.00},
|
||||
"claude-opus-4-1": {"input": 15.00, "output": 75.00},
|
||||
"claude-opus-4": {"input": 15.00, "output": 75.00},
|
||||
"claude-sonnet-4-6": {"input": 3.00, "output": 15.00},
|
||||
# Sonnet 5 standard rate (== Sonnet 4.6). Introductory pricing of
|
||||
# $2/$10 runs through 2026-08-31; intentionally NOT tracked here —
|
||||
# chosen set-and-forget so no dated bump is needed on 2026-09-01.
|
||||
# (New tokenizer emits ~30% more tokens, so per-run cost still rises.)
|
||||
"claude-sonnet-5": {"input": 3.00, "output": 15.00},
|
||||
"claude-sonnet-4-5": {"input": 3.00, "output": 15.00},
|
||||
"claude-sonnet-4": {"input": 3.00, "output": 15.00},
|
||||
"claude-haiku-4-5-20251001": {"input": 1.00, "output": 5.00},
|
||||
"claude-haiku-4-5": {"input": 1.00, "output": 5.00},
|
||||
"claude-haiku-3-5": {"input": 0.80, "output": 4.00},
|
||||
"default": {"input": 3.00, "output": 15.00},
|
||||
},
|
||||
"gemini": {
|
||||
# Gemini 3 Flash pricing (per 1M tokens). Preview alias mirrors GA.
|
||||
"gemini-3-flash-preview": {"input": 0.30, "output": 2.50},
|
||||
"gemini-3-flash": {"input": 0.30, "output": 2.50},
|
||||
"gemini-flash-latest": {"input": 0.30, "output": 2.50},
|
||||
"gemini-2.5-flash": {"input": 0.30, "output": 2.50},
|
||||
"gemini-2.5-pro": {"input": 1.25, "output": 10.00},
|
||||
# Gemini 3.1 Pro Preview — standard tier, prompts ≤200k tokens.
|
||||
# Above 200k Google charges $4.00/$18.00; we don't yet split by
|
||||
# prompt size, so we use the smaller-tier rate. Almost every
|
||||
# pipeline call here is well under 200k.
|
||||
"gemini-3.1-pro-preview": {"input": 2.00, "output": 12.00},
|
||||
"gemini-3-pro-preview": {"input": 2.00, "output": 12.00},
|
||||
"default": {"input": 0.30, "output": 2.50},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# Per-provider cache token multipliers, applied on top of the input rate.
|
||||
# create: cost when a cache is *written* (Anthropic charges 1.25× input;
|
||||
# Gemini charges 1.0× input — caching writes are billed as a
|
||||
# normal input pass)
|
||||
# read: cost when a cached prefix is *reused* (much cheaper)
|
||||
CACHE_RATES: dict[str, dict[str, float]] = {
|
||||
"anthropic": {"create": 1.25, "read": 0.10},
|
||||
"gemini": {"create": 1.00, "read": 0.25},
|
||||
}
|
||||
|
||||
|
||||
def cost_for_entry(entry: dict) -> float:
|
||||
"""USD cost for an api_logs entry. Reads ``provider`` (default
|
||||
``anthropic`` for legacy entries) and ``model`` to pick rates."""
|
||||
provider = entry.get("provider") or "anthropic"
|
||||
table = PRICING.get(provider) or PRICING["anthropic"]
|
||||
rates = table.get(entry.get("model", ""), table["default"])
|
||||
cache_rates = CACHE_RATES.get(provider, CACHE_RATES["anthropic"])
|
||||
input_rate = rates["input"]
|
||||
output_rate = rates["output"]
|
||||
return (
|
||||
entry.get("input_tokens", 0) * input_rate
|
||||
+ entry.get("cache_creation_input_tokens", 0) * input_rate * cache_rates["create"]
|
||||
+ entry.get("cache_read_input_tokens", 0) * input_rate * cache_rates["read"]
|
||||
+ entry.get("output_tokens", 0) * output_rate
|
||||
) / 1_000_000
|
||||
|
||||
|
||||
def total_cost(entries: list[dict]) -> float:
|
||||
"""Sum USD across entries."""
|
||||
return round(sum(cost_for_entry(e) for e in entries), 6)
|
||||
@@ -0,0 +1,116 @@
|
||||
"""Provider-agnostic message and completion types.
|
||||
|
||||
These dataclasses are the lingua franca between calling code and providers.
|
||||
Each provider implementation translates these into its native shape on the
|
||||
way out and back into these on the way in.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Content blocks — what goes inside a Message
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextBlock:
|
||||
text: str
|
||||
cacheable: bool = False
|
||||
# Gemini 3 / thinking-mode: opaque bytes the model returns alongside text
|
||||
# parts that came from internal reasoning. Must be replayed verbatim when
|
||||
# this turn is fed back into the conversation, or the next call 400s.
|
||||
# Anthropic: always None.
|
||||
thought_signature: bytes | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PdfBlock:
|
||||
"""Inline PDF document. Provider encodes as base64 (Anthropic) or
|
||||
inline_data (Gemini) and applies caching policy if cacheable=True."""
|
||||
path: Path
|
||||
cacheable: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""Assistant turn: model called a tool."""
|
||||
id: str
|
||||
name: str
|
||||
input: dict[str, Any]
|
||||
# Same purpose as TextBlock.thought_signature — Gemini 3 attaches one
|
||||
# to every function_call part when thinking is on. Round-trip required.
|
||||
thought_signature: bytes | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolResultBlock:
|
||||
"""User turn: result fed back from a tool the model invoked previously."""
|
||||
tool_use_id: str
|
||||
name: str
|
||||
content: str
|
||||
|
||||
|
||||
ContentBlock = TextBlock | PdfBlock | ToolCall | ToolResultBlock
|
||||
|
||||
|
||||
@dataclass
|
||||
class Message:
|
||||
role: Literal["user", "assistant"]
|
||||
content: list[ContentBlock]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tools
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolSchema:
|
||||
"""JSON-schema tool definition. Both providers accept the same shape."""
|
||||
name: str
|
||||
description: str
|
||||
input_schema: dict[str, Any]
|
||||
|
||||
|
||||
# Tool choice: "auto" (model picks), "none" (no tools), or a forced name
|
||||
ToolChoice = Literal["auto", "none"] | dict # {"name": "save_xyz"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Completion / usage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class Usage:
|
||||
"""Token usage normalised across providers.
|
||||
|
||||
Anthropic exposes cache_creation_input_tokens (write) and
|
||||
cache_read_input_tokens (hit). Gemini only exposes a cache hit count
|
||||
(cached_content_token_count) — its cache writes don't bill as input.
|
||||
|
||||
For Gemini, ``cache_creation_tokens`` is always 0; ``cache_read_tokens``
|
||||
holds the cached hit count when a cache was used.
|
||||
"""
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_creation_tokens: int = 0
|
||||
cache_read_tokens: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class Completion:
|
||||
"""Result of a single provider.complete() / session.complete() call."""
|
||||
text: str # any text block(s) concatenated
|
||||
tool_calls: list[ToolCall]
|
||||
usage: Usage
|
||||
stop_reason: str
|
||||
raw_assistant_blocks: list[ContentBlock] = field(default_factory=list)
|
||||
"""The full assistant message, in our normalised content-block form, so
|
||||
callers can append it back to the conversation history when continuing
|
||||
the loop."""
|
||||
Reference in New Issue
Block a user