Add Anthropic Claude provider

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,
+\       },
 \   },
 \}
 

autoload/neural/source/anthropic.vim

@@ -0,0 +1,10 @@
+" Author: w0rp <[email protected]>
+" 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

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.
 

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