51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
from __future__ import annotations
|
|
import anthropic
|
|
|
|
|
|
INPUT_COST_PER_1M = 3.0 # claude-sonnet-4-6 $/1M tokens
|
|
OUTPUT_COST_PER_1M = 15.0
|
|
|
|
|
|
class ClaudeClient:
|
|
MODEL = "claude-sonnet-4-6"
|
|
|
|
def __init__(self, api_key: str):
|
|
self._client = anthropic.Anthropic(api_key=api_key)
|
|
|
|
def ask(
|
|
self,
|
|
prompt: str,
|
|
system: str = "",
|
|
max_tokens: int = 1024,
|
|
) -> tuple[str, dict]:
|
|
messages = [{"role": "user", "content": prompt}]
|
|
kwargs = {"model": self.MODEL, "max_tokens": max_tokens, "messages": messages}
|
|
if system:
|
|
kwargs["system"] = system
|
|
|
|
response = self._client.messages.create(**kwargs)
|
|
text = response.content[0].text if response.content else ""
|
|
usage = {
|
|
"input_tokens": response.usage.input_tokens,
|
|
"output_tokens": response.usage.output_tokens,
|
|
"cost_usd": self._estimate_cost(response.usage.input_tokens, response.usage.output_tokens),
|
|
}
|
|
return text, usage
|
|
|
|
def stream_ask(self, prompt: str, system: str = "", max_tokens: int = 1024):
|
|
messages = [{"role": "user", "content": prompt}]
|
|
kwargs = {"model": self.MODEL, "max_tokens": max_tokens, "messages": messages}
|
|
if system:
|
|
kwargs["system"] = system
|
|
|
|
with self._client.messages.stream(**kwargs) as stream:
|
|
for text in stream.text_stream:
|
|
yield text
|
|
|
|
def _estimate_cost(self, input_tokens: int, output_tokens: int) -> float:
|
|
return (input_tokens / 1_000_000 * INPUT_COST_PER_1M) + (output_tokens / 1_000_000 * OUTPUT_COST_PER_1M)
|
|
|
|
@staticmethod
|
|
def estimate_tokens(text: str) -> int:
|
|
return max(1, len(text) // 4)
|