diff --git a/autoload/neural/config.vim b/autoload/neural/config.vim index 933ee4f..caefd97 100644 --- a/autoload/neural/config.vim +++ b/autoload/neural/config.vim @@ -41,6 +41,13 @@ let s:defaults = { \ 'temperature': 0.2, \ 'top_p': 1, \ }, +\ 'anthropic': { +\ 'api_key': '', +\ 'max_tokens': 1024, +\ 'model': 'claude-2', +\ 'temperature': 0.2, +\ 'top_p': 1, +\ }, \ }, \} diff --git a/autoload/neural/source/anthropic.vim b/autoload/neural/source/anthropic.vim new file mode 100644 index 0000000..be1f6d8 --- /dev/null +++ b/autoload/neural/source/anthropic.vim @@ -0,0 +1,10 @@ +" Author: w0rp +" Description: A script describing how to use Anthropic Claude with Neural + +function! neural#source#anthropic#Get() abort + return { + \ 'name': 'anthropic', + \ 'script_language': 'python', + \ 'script': neural#GetScriptDir() . '/anthropic.py', + \} +endfunction diff --git a/doc/neural.txt b/doc/neural.txt index 7da310e..c1a2cc8 100644 --- a/doc/neural.txt +++ b/doc/neural.txt @@ -12,7 +12,8 @@ CONTENTS *neural-contents* 4.2 Neural Buffer ............................. |neural-buffer| 4.3 OpenAI .................................... |neural-openai| 4.4 ChatGPT ................................... |neural-chatgpt| - 4.5 Highlights ................................ |neural-highlights| + 4.5 Anthropic Claude .......................... |neural-anthropic| + 4.6 Highlights ................................ |neural-highlights| 5. API ........................................ |neural-api| 6. Environment Variables ...................... |neural-env| 6.1 Linux + KDE ............................. |neural-env-kde| @@ -33,6 +34,7 @@ generate text, code, and much more. Neural supports the following tools. 1. OpenAI - https://beta.openai.com/signup/ +2. Anthropic Claude - https://www.anthropic.com/ To select the tool that Neural will use, set |g:neural.selected| to the appropriate value. OpenAI is the default data source. @@ -160,6 +162,7 @@ g:neural.selected *g:neural.selected* 1. `'openai'` - OpenAI 2. `'chatgpt'` - ChatGPT + 3. `'anthropic'` - Anthropic Claude g:neural.set_default_keybinds *g:neural.set_default_keybinds* @@ -419,7 +422,53 @@ g:neural.source.chatgpt.top_p *g:neural.source.chatgpt.top_p* ------------------------------------------------------------------------------- -4.5 Highlights *neural-highlights* +4.5 Anthropic Claude *neural-anthropic* + +Options for configuring Anthropic Claude are listed below. + + +g:neural.source.anthropic.api_key *g:neural.source.anthropic.api_key* + *vim.g.neural.source.anthropic.api_key* + Type: |String| + Default: `''` + + The Anthropic API key. See: https://www.anthropic.com/ + + +g:neural.source.anthropic.max_tokens *g:neural.source.anthropic.max_tokens* + *vim.g.neural.source.anthropic.max_tokens* + Type: |Number| + Default: `1024` + + The maximum number of tokens to generate in the output. + + +g:neural.source.anthropic.model *g:neural.source.anthropic.model* + *vim.g.neural.source.anthropic.model* + Type: |String| + Default: `'claude-2'` + + The model to use for Anthropic Claude. + + +g:neural.source.anthropic.temperature *g:neural.source.anthropic.temperature* + *vim.g.neural.source.anthropic.temperature* + Type: |Number| or |Float| + Default: `0.2` + + The Anthropic sampling temperature. + + +g:neural.source.anthropic.top_p *g:neural.source.anthropic.top_p* + *vim.g.neural.source.anthropic.top_p* + Type: |Number| or |Float| + Default: `1` + + The Anthropic nucleus sampling between `0` and `1`. + + +------------------------------------------------------------------------------- +4.6 Highlights *neural-highlights* The following highlights can be configured to change |neural|'s colors. diff --git a/neural_providers/anthropic.py b/neural_providers/anthropic.py new file mode 100644 index 0000000..ce3304f --- /dev/null +++ b/neural_providers/anthropic.py @@ -0,0 +1,159 @@ +""" +A Neural datasource for loading generated text via Anthropic Claude. +""" +import json +import platform +import ssl +import sys +import urllib.error +import urllib.request +from typing import Any, Dict, List, Optional, Union + +API_ENDPOINT = 'https://api.anthropic.com/v1/complete' + +ANTHROPIC_DATA_HEADER = 'data: ' +ANTHROPIC_DONE = '[DONE]' + + +class Config: + """The sanitised configuration.""" + + def __init__( + self, + api_key: str, + model: str, + temperature: float, + top_p: float, + max_tokens: int, + ) -> None: + self.api_key = api_key + self.model = model + self.temperature = temperature + self.top_p = top_p + self.max_tokens = max_tokens + + +def get_claude_completion( + config: Config, + prompt: Union[str, List[Dict[str, str]]], +) -> None: + headers = { + "Content-Type": "application/json", + "x-api-key": config.api_key, + "anthropic-version": "2023-06-01", + } + data = { + "model": config.model, + "prompt": ( + prompt + if isinstance(prompt, str) + else ''.join([msg.get("content", "") for msg in prompt]) + ), + "temperature": config.temperature, + "top_p": config.top_p, + "max_tokens_to_sample": config.max_tokens, + "stream": True, + } + + req = urllib.request.Request( + API_ENDPOINT, + data=json.dumps(data).encode("utf-8"), + headers=headers, + method="POST", + ) + role: Optional[str] = None + + context = ( + ssl._create_unverified_context() # type: ignore + if platform.system() == "Darwin" else + None + ) + + with urllib.request.urlopen(req, context=context) as response: + while True: + line_bytes = response.readline() + + if not line_bytes: + break + + line = line_bytes.decode("utf-8", errors="replace") + line_data = ( + line[len(ANTHROPIC_DATA_HEADER):-1] + if line.startswith(ANTHROPIC_DATA_HEADER) else None + ) + + if line_data and line_data != ANTHROPIC_DONE: + chunk = json.loads(line_data) + + if "completion" in chunk: + print(chunk["completion"], end="", flush=True) + + print() + + +def load_config(raw_config: Dict[str, Any]) -> Config: + if not isinstance(raw_config, dict): # type: ignore + raise ValueError("anthropic config is not a dictionary") + + api_key = raw_config.get('api_key') + if not isinstance(api_key, str) or not api_key: # type: ignore + raise ValueError("anthropic.api_key is not defined") + + model = raw_config.get('model') + if not isinstance(model, str) or not model: + raise ValueError("anthropic.model is not defined") + + temperature = raw_config.get('temperature', 0.2) + if not isinstance(temperature, (int, float)): + raise ValueError("anthropic.temperature is invalid") + + top_p = raw_config.get('top_p', 1) + if not isinstance(top_p, (int, float)): + raise ValueError("anthropic.top_p is invalid") + + max_tokens = raw_config.get('max_tokens', 1024) + if not isinstance(max_tokens, int): + raise ValueError("anthropic.max_tokens is invalid") + + return Config( + api_key=api_key, + model=model, + temperature=temperature, + top_p=top_p, + max_tokens=max_tokens, + ) + + +def get_error_message(error: urllib.error.HTTPError) -> str: + message = error.read().decode('utf-8', errors='ignore') + + try: + message = json.loads(message)['error']['message'] + except Exception: + pass + + return message + + +def main() -> None: + input_data = json.loads(sys.stdin.readline()) + + try: + config = load_config(input_data["config"]) + except ValueError as err: + sys.exit(str(err)) + + try: + get_claude_completion(config, input_data["prompt"]) + except urllib.error.HTTPError as error: + if error.code in (400, 401): + message = get_error_message(error) + sys.exit('Neural error: Anthropic request failure: ' + message) + elif error.code == 429: + sys.exit('Neural error: Anthropic request limit reached!') + else: + raise + + +if __name__ == "__main__": # pragma: no cover + main() # pragma: no cover diff --git a/test/python/test_anthropic.py b/test/python/test_anthropic.py new file mode 100644 index 0000000..6308909 --- /dev/null +++ b/test/python/test_anthropic.py @@ -0,0 +1,164 @@ +import json +import urllib.error +import urllib.request +from io import BytesIO +from typing import Any, Dict, Optional, cast +from unittest import mock + +import pytest +import sys + +from neural_providers import anthropic + + +def get_valid_config() -> Dict[str, Any]: + return { + "api_key": ".", + "model": "foo", + "temperature": 1, + "top_p": 1, + "max_tokens": 1, + } + + +def test_load_config_errors() -> None: + with pytest.raises(ValueError) as exc: + anthropic.load_config(cast(Any, 0)) + + assert str(exc.value) == "anthropic config is not a dictionary" + + config: Dict[str, Any] = {} + + for modification, expected_error in [ + ({}, "anthropic.api_key is not defined"), + ({"api_key": ""}, "anthropic.api_key is not defined"), + ({"api_key": "."}, "anthropic.model is not defined"), + ({"model": ""}, "anthropic.model is not defined"), + ({"model": "x", "temperature": "x"}, "anthropic.temperature is invalid"), + ({"temperature": 1, "top_p": "x"}, "anthropic.top_p is invalid"), + ({"top_p": 1, "max_tokens": "x"}, "anthropic.max_tokens is invalid"), + ]: + config.update(modification) + + with pytest.raises(ValueError) as exc: + anthropic.load_config(config) + + assert str(exc.value) == expected_error, config + + +def test_main_function_rate_other_error() -> None: + with mock.patch.object(sys.stdin, 'readline') as readline_mock, \ + mock.patch.object(anthropic, 'get_claude_completion') as compl_mock: + + compl_mock.side_effect = urllib.error.HTTPError( + url='', + msg='', + hdrs=mock.Mock(), + fp=None, + code=500, + ) + readline_mock.return_value = json.dumps({ + "config": get_valid_config(), + "prompt": "hello there", + }) + + with pytest.raises(urllib.error.HTTPError): + anthropic.main() + + +def test_print_anthropic_results() -> None: + result_data = ( + b'data: {"completion":"Hi"}\n' # noqa + b'\n' + b'data: {"completion":"!"}\n' # noqa + b'\n' + b'data: [DONE]\n' + b'\n' + ) + + with mock.patch.object(sys.stdin, 'readline') as readline_mock, \ + mock.patch.object(urllib.request, 'urlopen') as urlopen_mock, \ + mock.patch('builtins.print') as print_mock: + + urlopen_mock.return_value.__enter__.return_value = BytesIO(result_data) + + readline_mock.return_value = json.dumps({ + "config": get_valid_config(), + "prompt": "hello there", + }) + anthropic.main() + + assert print_mock.call_args_list == [ + mock.call('Hi', end='', flush=True), + mock.call('!', end='', flush=True), + mock.call(), + ] + + +def test_main_function_bad_config() -> None: + with mock.patch.object(sys.stdin, 'readline') as readline_mock, \ + mock.patch.object(anthropic, 'load_config') as load_config_mock: + + load_config_mock.side_effect = ValueError("expect this") + readline_mock.return_value = json.dumps({"config": {}}) + + with pytest.raises(SystemExit) as exc: + anthropic.main() + + assert str(exc.value) == 'expect this' + + +@pytest.mark.parametrize( + 'code, error_text, expected_message', + ( + pytest.param( + 429, + None, + 'Anthropic request limit reached!', + id="request_limit", + ), + pytest.param( + 400, + '{]', + 'Anthropic request failure: {]', + id="error_with_mangled_json", + ), + pytest.param( + 400, + json.dumps({'error': {}}), + 'Anthropic request failure: {"error": {}}', + id="error_with_missing_message_key", + ), + pytest.param( + 401, + json.dumps({'error': {'message': 'Bad authentication error'}}), + 'Anthropic request failure: Bad authentication error', + id="unauthorised_failure", + ), + ) +) +def test_api_error( + code: int, + error_text: Optional[str], + expected_message: str, +) -> None: + with mock.patch.object(sys.stdin, 'readline') as readline_mock, \ + mock.patch.object(anthropic, 'get_claude_completion') as compl_mock: + + compl_mock.side_effect = urllib.error.HTTPError( + url='', + msg='', + hdrs=mock.Mock(), + fp=BytesIO(error_text.encode('utf-8')) if error_text else None, + code=code, + ) + + readline_mock.return_value = json.dumps({ + "config": get_valid_config(), + "prompt": "hello there", + }) + + with pytest.raises(SystemExit) as exc: + anthropic.main() + + assert str(exc.value) == f'Neural error: {expected_message}' diff --git a/test/vim/test_anthropic.vader b/test/vim/test_anthropic.vader new file mode 100644 index 0000000..dfb40ad --- /dev/null +++ b/test/vim/test_anthropic.vader @@ -0,0 +1,8 @@ +Execute(The Anthropic source Dictionary should be correct): + AssertEqual + \ { + \ 'name': 'anthropic', + \ 'script_language': 'python', + \ 'script': neural#GetScriptDir() . '/anthropic.py' + \ }, + \ neural#source#anthropic#Get() diff --git a/test/vim/test_config.vader b/test/vim/test_config.vader index b5242f0..8629125 100644 --- a/test/vim/test_config.vader +++ b/test/vim/test_config.vader @@ -66,6 +66,19 @@ Execute(The default chatgpt settings should be correct): \ }, \ get(g:neural.source, 'chatgpt') +Execute(The default anthropic settings should be correct): + call neural#config#Load() + + AssertEqual + \ { + \ 'api_key': '', + \ 'max_tokens': 1024, + \ 'model': 'claude-2', + \ 'temperature': 0.2, + \ 'top_p': 1, + \ }, + \ get(g:neural.source, 'anthropic') + Execute(The default neural buffer settings should be correct): call neural#config#Load() " call filter(g:neural.buffer, {key -> key =~ 'completion'})