Repository: madox2/vim-ai

Issue:
Additional Providers (AWS Bedrock) 
Hello there -

Use Case:

I would like to see AWS Bedrock provider support capabilities in the plugin.  I would like to be able to use it without using local proxies or creating infrastructure AWS side for supporting OpenAI api format.

Why:
I work at an enterprise and can't use OpenAI for code assist. I can however use Amazon bedrock. While I can put infrastructure in place to get bedrock to respond to OpenAI requests, it would require maintance and additional costs that the company wouldn't be willing to take on for the small group of vim users.

Request/Question:

I'm wondering if you would be open to PR(s) that would enable support for Amazon Bedrock. I've been playing around with a hacked up version of this plugin (not working yet) but I believe it's possible to configure/add more providers. It would require some refactoring to make the integration clean/scalable. Is this something you'd be open to?

Hints from issue comments:
@madox2: I am not planning to add built-in support for other providers. I don't want this repo/plugin end up maintaining many integrations, there are already projects which focus solely on that (like LiteLLM).

However I am fine with opening vim-ai for an extension. Here is a rough idea how it could work:

```vim
Plug 'madox2/vim-ai'
Plug 'maxres-ch/vim-ai-provider-bedrock'

let g:vim_ai_chat = {
\  "provider": "bedrock",
\  "options": {
\    ...
\  },
\}
```

That would require a good abstraction and a handful of refactoring as you already mentioned.@maxres-ch: Thanks so much for the fast response.

>  I don't want this repo/plugin end up maintaining many integrations, there are already projects which focus solely on that (like LiteLLM).

Yea totally understand and 100% respect this. 

> However I am fine with opening vim-ai for an extension. Here is a rough idea how it could work:

That seems like a good approach, to me. Leaves you with less debt. 

> That would require a good abstraction and a handful of refactoring as you already mentioned.

I will look into this as an avenue. Would you have a preferred way for me to propose the refactor? I'd hate to write a bunch of code and you not liking it. This smells like potentially using some classes, I've been writing so much Golang lately I'll have to get back to python dev mode. I generally avoid classes unless needed, do you have a gut reaction to that?@madox2: Cool, I think that it will be a really nice feature since other people were also requesting integration to other providers. 

Using a python class sounds like a clean way how to define the interface to the extension, so no problem with that. If you want, create a (draft) PR so I can review early or maybe even participate.

Thank you for being involved in this!@grumps: Hi @madox2  (from my personal account). I've [drafted up a PR for provider integration](https://github.com/madox2/vim-ai/pull/143), it's not complete and won't work for OpenAI. It does work for my bedrock work. Maybe this would be better to complete in the PR.

I created a basic class with typing for an "api" as follows: 

```python
from collections.abc import Sequence, Mapping
from typing import TypedDict, Protocol


class AIProvider(Protocol):

  def __init__(self, config: Mapping[str, str|float|dict]) -> None:
      pass

  def request(self, messages: Sequence[Message]) -> List[str]:
      pass

class Message(TypedDict):
    role: str
    content: str
```

I propose changing the main on `chat.py` and `complete.py` to be something like below.  I do think passing the whole config object to class. 

```python
        provider_class = load_provider(config['provider'])
        provider = provider_class(config)
```

Here's the current way of loading the class:

```python
def load_provider(provider):
    provider_name, provider_module = provider["name"].split(".")
    if provider_name != "openai":

        provider_path = os.path.join(f"{plugin_root}",
                                     "..",
                                     f"vim-ai-{provider_name}",
                                     "py",
                                     f"{provider_module}.py")
    else:
        return openai_request

    vim.command(f"py3file {provider_path}")
    try:
        provider_class = globals()[provider["class"]]
    except KeyError as error:
        printDebug("[load-provider] provider: {}", error)
        raise KeyError(error.message, "provider not found")
    return provider_class
```

O do think this could be improved by using importlib, something like this:

```python

import importlib.util
def import_class(module_path, module_name, class_name):
    spec = importlib.util.spec_from_file_location(module_name, module_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return getattr(module, class_name)

```


Here's the suggested change to the config (probably in the default config):
```vim
\
  "provider": {
\    "name": "openai.complete", 
\    "class": "OpenAIChat"
\    "options: { ... } # additional options for the provider could be picked up here (or possibly another variable/map specific for the  provider?
\ }
```

Here's the python code for a provider I created:

```python
import boto3
import logging

import json

from botocore.exceptions import ClientError

logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(name)-12s %(levelname)-8s %(message)s',
                    datefmt='%m-%d %H:%M',
                    filename='/tmp/vim-ai-bedrock.log',
                    filemode='w')
logger = logging.getLogger(__name__)




# Create a Bedrock Runtime client in the AWS Region of your choice.
boto3.setup_default_session(profile_name='coder')
client = boto3.client('bedrock-runtime',
                      region_name='us-west-2')

# Set the model ID, e.g., Claude 3 Haiku.
model_id = "us.meta.llama3-2-3b-instruct-v1:0"
# model_id = "us.anthropic.claude-3-5-sonnet-20241022-v2:0"


test_messages = [
  {'role': 'user',
   'content': 'write a rust program that writes hello world to a file'},
  {'role': 'user',
   'content': 'write a rust program that writes hello world to a file'},
  {'role': 'user',
   'content': 'write a rust program that writes hello world to a file'}
]

class BedrockClient:

    def __init__(self, config):
        self.config = config

    def map_msgs(self, messages):
        user_msgs, system_msgs = [], []
        for msg in messages:
            # bedrock doesn't use include, we switch the role to user
            if msg['role'] == 'include':
                msg['role'] = 'user'
            elif msg['role'] == 'system':
                system_msgs.append({
                    'text': msg['content'].strip().strip('\n')
                })
                continue
            # build the dict into the structure that's supported by bedrock
            user_msgs.append({
              'role': msg['role'],
                'content': [
                    {
                      'text': msg['content'],
                    },
                ]
            })
        return user_msgs, system_msgs

    def request(self, messages):
        modelId='us.meta.llama3-3-70b-instruct-v1:0'
        user_messages, system_messages = self.map_msgs(messages)

        logger.info(user_messages)
        # this needs to come from config, not inlined
        inferenceConfig={
            'maxTokens': 3000,
            'temperature': 0.1,
            'topP': 0.4,
            'stopSequences': [
                'string',
            ]
        }

        try:
            # Invoke the model with the request.
            response = client.converse(modelId=modelId,
                                       messages=user_messages,
                                       inferenceConfig=inferenceConfig,
                                       system=system_messages)
        except (ClientError, Exception) as e:
            print(f"ERROR: Can't invoke '{modelId}'. Reason: {e}")
            exit(1)
        # Decode the response body.

        response_text= []
        content = response.get('output', {}).get('message', {}).get('content', [{}])
        for txt in content:
            response_text.append(txt['text'])
        return response_text
 ```
@madox2: Hi @grumps , the idea is good! I have couple of things in mind we need to figure out, I will continue the discussion directly on the code/PR.@madox2: Support for custom AI providers added

Solution patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 9fa10c9..213798f 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -1,5 +1,11 @@
 name: Tests
-on: [push]
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    types: [opened, synchronize, reopened]
+  workflow_dispatch:
 jobs:
   build:
     runs-on: ubuntu-latest
diff --git a/py/chat.py b/py/chat.py
index 6861b51..810fc46 100644
--- a/py/chat.py
+++ b/py/chat.py
@@ -63,8 +63,10 @@ def run_ai_chat(context):
 
             print('Answering...')
             vim.command("redraw")
+            provider_class = load_provider(config['provider'])
+            provider = provider_class(config)
+            text_chunks = provider.request(messages)
 
-            text_chunks = make_chat_text_chunks(messages, options)
             render_text_chunks(text_chunks)
 
             vim.command("normal! a\n\n>>> user\n\n")
diff --git a/py/complete.py b/py/complete.py
index 8078dea..765d6e0 100644
--- a/py/complete.py
+++ b/py/complete.py
@@ -42,7 +42,10 @@ def run_ai_completition(context):
         if prompt:
             print('Completing...')
             vim.command("redraw")
-            text_chunks = engines[engine](prompt)
+            provider_class = load_provider(config['provider'])
+            provider = provider_class(config)
+            messages = parse_chat_messages(f">>> user\n\n{prompt}".strip())
+            text_chunks = provider.request(messages)
             render_text_chunks(text_chunks)
             clear_echo_message()
     except BaseException as error:
diff --git a/py/types.py b/py/types.py
new file mode 100644
index 0000000..4eb0ccc
--- /dev/null
+++ b/py/types.py
@@ -0,0 +1,32 @@
+from collections.abc import Sequence, Mapping
+from typing import TypedDict, Literal, Union, List
+
+
+class AIProvider(Protocol):
+
+  def __init__(self, config: Mapping[str, str]) -> None:
+      pass
+
+  def request(self, messages: Sequence[Message]) -> Generator[str]:
+      pass
+
+class Message(TypedDict):
+    role: str
+    content: str
+    type: str
+
+
+class TextContent(TypedDict):
+    type: Literal['text']
+    text: str
+
+class ImageUrlContent(TypedDict):
+    type: Literal['image_url']
+    image_url: dict[str, str]  # {'url': str}
+
+MessageContent = Union[TextContent, ImageUrlContent]
+
+class Message(TypedDict):
+    role: Literal['system', 'user', 'assistant']
+    content: List[MessageContent]
+
diff --git a/py/utils.py b/py/utils.py
index 8cd204f..5dbc2d6 100644
--- a/py/utils.py
+++ b/py/utils.py
@@ -142,6 +142,7 @@ def make_image_message(path):
 
 def make_text_file_message(path):
     try:
+        path = os.path.relpath(path)
         with open(path, 'r') as file:
             file_content = file.read().strip()
             return { 'type': 'text', 'text': f'==> {path} <==\n' + file_content.strip() }
@@ -359,3 +360,23 @@ def save_b64_to_file(path, b64_data):
     f = open(path, "wb")
     f.write(base64.b64decode(b64_data))
     f.close()
+
+def load_provider(provider):
+    provider_name, provider_module = provider["name"].split(".")
+    if provider_name != "openai":
+        provider_path = os.path.join(f"{plugin_root}",
+                                     "..",
+                                     f"vim-ai-{provider_name}",
+                                     "py",
+                                     f"{provider_module}.py")
+    else:
+        return openai_request
+    vim.command(f"py3file {provider_path}")
+    try:
+        provider_class = globals()[provider["class"]]
+    except KeyError as error:
+        printDebug("[load-provider] provider: {}", error)
+        raise KeyError(error.message, "provider not found")
+    return provider_class
+
+


Test patch:
diff --git a/tests/chat_test.py b/tests/chat_test.py
index e29368c..f781c56 100644
--- a/tests/chat_test.py
+++ b/tests/chat_test.py
@@ -2,6 +2,7 @@ from utils import parse_chat_messages
 import os
 
 curr_dir = os.path.dirname(__file__)
+relative_dir = 'tests'
 
 def strip_text(txt):
     txt = txt.strip()
@@ -165,7 +166,7 @@ def test_parse_include_single_file_message():
                 },
                 {
                     'type': 'text',
-                    'text': f'==> {curr_dir}/resources/test1.include.txt <==\nhello world',
+                    'text': f'==> {relative_dir}/resources/test1.include.txt <==\nhello world',
                 },
             ],
         },
@@ -212,11 +213,11 @@ def test_parse_include_multiple_files_message():
                 },
                 {
                     'type': 'text',
-                    'text': f'==> {curr_dir}/resources/test1.include.txt <==\nhello world',
+                    'text': f'==> {relative_dir}/resources/test1.include.txt <==\nhello world',
                 },
                 {
                     'type': 'text',
-                    'text': f'==> {curr_dir}/resources/test2.include.txt <==\nvim is awesome',
+                    'text': f'==> {relative_dir}/resources/test2.include.txt <==\nvim is awesome',
                 },
             ],
         },
@@ -243,11 +244,11 @@ def test_parse_include_glob_files_message():
                 },
                 {
                     'type': 'text',
-                    'text': f'==> {curr_dir}/resources/test1.include.txt <==\nhello world',
+                    'text': f'==> {relative_dir}/resources/test1.include.txt <==\nhello world',
                 },
                 {
                     'type': 'text',
-                    'text': f'==> {curr_dir}/resources/test2.include.txt <==\nvim is awesome',
+                    'text': f'==> {relative_dir}/resources/test2.include.txt <==\nvim is awesome',
                 },
             ],
         },
@@ -297,7 +298,7 @@ def test_parse_include_image_with_files_message():
             'content': [
                 {
                     'type': 'text',
-                    'text': f'==> {curr_dir}/resources/test1.include.txt <==\nhello world',
+                    'text': f'==> {relative_dir}/resources/test1.include.txt <==\nhello world',
                 },
                 {
                     'type': 'image_url',
@@ -307,7 +308,7 @@ def test_parse_include_image_with_files_message():
                 },
                 {
                     'type': 'text',
-                    'text': f'==> {curr_dir}/resources/test2.include.txt <==\nvim is awesome',
+                    'text': f'==> {relative_dir}/resources/test2.include.txt <==\nvim is awesome',
                 },
             ],
         },
@@ -327,11 +328,11 @@ def test_parse_include_unsupported_binary_file():
             'content': [
                 {
                     'type': 'text',
-                    'text': f'==> {curr_dir}/resources/binary_file.bin <==\nBinary file, cannot display',
+                    'text': f'==> {relative_dir}/resources/binary_file.bin <==\nBinary file, cannot display',
                 },
                 {
                     'type': 'text',
-                    'text': f'==> {curr_dir}/resources/test1.include.txt <==\nhello world',
+                    'text': f'==> {relative_dir}/resources/test1.include.txt <==\nhello world',
                 },
             ],
         },