|
| 1 | +from collections.abc import AsyncGenerator |
| 2 | + |
| 3 | +from anthropic import AsyncAnthropic |
| 4 | + |
| 5 | +from astrbot.core.provider.entities import LLMResponse |
| 6 | + |
| 7 | +from ..register import register_provider_adapter |
| 8 | +from .anthropic_source import ProviderAnthropic |
| 9 | + |
| 10 | +_OAUTH_DEFAULT_HEADERS = { |
| 11 | + "anthropic-beta": "claude-code-20250219,oauth-2025-04-20,context-1m-2025-08-07", |
| 12 | + "user-agent": "claude-cli/1.0.0 (external, cli)", |
| 13 | + "x-app": "cli", |
| 14 | + "anthropic-dangerous-direct-browser-access": "true", |
| 15 | +} |
| 16 | + |
| 17 | +_CLAUDE_CODE_SYSTEM_PREFIX = ( |
| 18 | + "You are Claude Code, Anthropic's official CLI for Claude.\n\n" |
| 19 | +) |
| 20 | + |
| 21 | +# 支持 1M 上下文窗口的模型前缀（需配合 context-1m beta header）。 |
| 22 | +# 新增 4.6+ 模型时需同步更新此列表。 |
| 23 | +_1M_CONTEXT_MODEL_PREFIXES = ( |
| 24 | + "claude-opus-4-6", |
| 25 | + "claude-sonnet-4-6", |
| 26 | +) |
| 27 | + |
| 28 | + |
| 29 | +@register_provider_adapter( |
| 30 | + "anthropic_oauth", |
| 31 | + "Anthropic Claude Code OAuth provider adapter", |
| 32 | +) |
| 33 | +class ProviderAnthropicOAuth(ProviderAnthropic): |
| 34 | + def __init__( |
| 35 | + self, |
| 36 | + provider_config: dict, |
| 37 | + provider_settings: dict, |
| 38 | + ) -> None: |
| 39 | + # 禁用父类的 API key 客户端初始化，避免重复构造客户端 |
| 40 | + super().__init__(provider_config, provider_settings, use_api_key=False) |
| 41 | + |
| 42 | + # 手动解析 key 列表（父类跳过了 _init_api_key） |
| 43 | + self.api_keys: list = self.get_keys() |
| 44 | + self.chosen_api_key: str = self.api_keys[0] if self.api_keys else "" |
| 45 | + |
| 46 | + # 使用 auth_token（OAuth Bearer 认证）构建客户端 |
| 47 | + self.client = AsyncAnthropic( |
| 48 | + auth_token=self.chosen_api_key, |
| 49 | + timeout=self.timeout, |
| 50 | + base_url=self.base_url, |
| 51 | + default_headers=_OAUTH_DEFAULT_HEADERS, |
| 52 | + http_client=self._create_http_client(provider_config), |
| 53 | + ) |
| 54 | + |
| 55 | + def set_model(self, model_name: str) -> None: |
| 56 | + super().set_model(model_name) |
| 57 | + if any(model_name.startswith(p) for p in _1M_CONTEXT_MODEL_PREFIXES): |
| 58 | + if self.provider_config.get("max_context_tokens", 0) <= 0: |
| 59 | + self.provider_config["max_context_tokens"] = 1_000_000 |
| 60 | + |
| 61 | + def get_model_metadata_overrides(self, model_ids: list[str]) -> dict[str, dict]: |
| 62 | + overrides = {} |
| 63 | + for mid in model_ids: |
| 64 | + if any(mid.startswith(p) for p in _1M_CONTEXT_MODEL_PREFIXES): |
| 65 | + overrides[mid] = {"limit": {"context": 1_000_000}} |
| 66 | + return overrides |
| 67 | + |
| 68 | + def set_key(self, key: str) -> None: |
| 69 | + self.chosen_api_key = key |
| 70 | + # 切换 key 时需要重建客户端以使用新的 auth_token |
| 71 | + self.client = AsyncAnthropic( |
| 72 | + auth_token=key, |
| 73 | + timeout=self.timeout, |
| 74 | + base_url=self.base_url, |
| 75 | + default_headers=_OAUTH_DEFAULT_HEADERS, |
| 76 | + http_client=self._create_http_client(self.provider_config), |
| 77 | + ) |
| 78 | + |
| 79 | + async def get_models(self) -> list[str]: |
| 80 | + return await super().get_models() |
| 81 | + |
| 82 | + async def test(self, timeout: float = 45.0) -> None: |
| 83 | + await super().test(timeout) |
| 84 | + |
| 85 | + async def text_chat( |
| 86 | + self, |
| 87 | + prompt=None, |
| 88 | + session_id=None, |
| 89 | + image_urls=None, |
| 90 | + func_tool=None, |
| 91 | + contexts=None, |
| 92 | + system_prompt=None, |
| 93 | + tool_calls_result=None, |
| 94 | + model=None, |
| 95 | + extra_user_content_parts=None, |
| 96 | + **kwargs, |
| 97 | + ) -> LLMResponse: |
| 98 | + system_prompt = _CLAUDE_CODE_SYSTEM_PREFIX + (system_prompt or "") |
| 99 | + |
| 100 | + return await super().text_chat( |
| 101 | + prompt=prompt, |
| 102 | + session_id=session_id, |
| 103 | + image_urls=image_urls, |
| 104 | + func_tool=func_tool, |
| 105 | + contexts=contexts, |
| 106 | + system_prompt=system_prompt, |
| 107 | + tool_calls_result=tool_calls_result, |
| 108 | + model=model, |
| 109 | + extra_user_content_parts=extra_user_content_parts, |
| 110 | + **kwargs, |
| 111 | + ) |
| 112 | + |
| 113 | + async def text_chat_stream( |
| 114 | + self, |
| 115 | + prompt=None, |
| 116 | + session_id=None, |
| 117 | + image_urls=None, |
| 118 | + func_tool=None, |
| 119 | + contexts=None, |
| 120 | + system_prompt=None, |
| 121 | + tool_calls_result=None, |
| 122 | + model=None, |
| 123 | + extra_user_content_parts=None, |
| 124 | + **kwargs, |
| 125 | + ) -> AsyncGenerator[LLMResponse, None]: |
| 126 | + system_prompt = _CLAUDE_CODE_SYSTEM_PREFIX + (system_prompt or "") |
| 127 | + |
| 128 | + async for llm_response in super().text_chat_stream( |
| 129 | + prompt=prompt, |
| 130 | + session_id=session_id, |
| 131 | + image_urls=image_urls, |
| 132 | + func_tool=func_tool, |
| 133 | + contexts=contexts, |
| 134 | + system_prompt=system_prompt, |
| 135 | + tool_calls_result=tool_calls_result, |
| 136 | + model=model, |
| 137 | + extra_user_content_parts=extra_user_content_parts, |
| 138 | + **kwargs, |
| 139 | + ): |
| 140 | + yield llm_response |
0 commit comments