Phase 1: initial codes

stockmind/ai/claude_client.py

@@ -0,0 +1,50 @@
+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)