From 62c4cf6b43d95ad0ff353030c54d3fe042711254 Mon Sep 17 00:00:00 2001
From: Bastian Fredriksson <bastian.fredriksson@keyfactor.com>
Date: Sat, 5 Oct 2024 09:00:23 +0200
Subject: [PATCH] feat: add provider for github models

---
 .devcontainer/requirements-dev.txt |  1 +
 README.md                          | 19 +++++++++++++++++--
 pyproject.toml                     |  3 ++-
 src/fish_ai/engine.py              | 15 +++++++++++++++
 4 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/.devcontainer/requirements-dev.txt b/.devcontainer/requirements-dev.txt
index 2895268..eb8979d 100644
--- a/.devcontainer/requirements-dev.txt
+++ b/.devcontainer/requirements-dev.txt
@@ -11,3 +11,4 @@ iterfzf
 hugchat
 mistralai
 binaryornot
+azure-ai-inference
diff --git a/README.md b/README.md
index 2faa86a..9366460 100644
--- a/README.md
+++ b/README.md
@@ -10,8 +10,8 @@ Python and git is installed.
 Originally based on [Tom Dörr's `fish.codex` repository](https://github.com/tom-doerr/codex.fish),
 but with some additional functionality.
 
-It can be hooked up to OpenAI, Azure OpenAI, Google, Hugging Face, Mistral or a
-self-hosted LLM behind any OpenAI-compatible API.
+It can be hooked up to OpenAI, Azure OpenAI, Google, Hugging Face, Mistral,
+GitHub or a self-hosted LLM behind any OpenAI-compatible API.
 
 If you like it, please add a ⭐. If you don't like it, create a PR. 😆
 
@@ -118,6 +118,21 @@ provider = mistral
 api_key = <your API key>
 ```
 
+If you use [GitHub Models](https://github.com/marketplace/models):
+
+```ini
+[fish-ai]
+configuration = github
+
+[github]
+provider = github
+api_key = <GitHub PAT>
+model = Meta-Llama-3.1-70B-Instruct
+```
+
+You can create a personal access token (PAT) [here](https://github.com/settings/tokens).
+The PAT does not require any permissions.
+
 ### Install `fish-ai`
 
 Install the plugin. You can install it using [`fisher`](https://github.com/jorgebucaran/fisher).
diff --git a/pyproject.toml b/pyproject.toml
index a7880b5..ce4748d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fish_ai"
-version = "0.9.10"
+version = "0.10.0"
 authors = [{ name = "Bastian Fredriksson", email = "realiserad@gmail.com" }]
 description = "Provides core functionality for fish-ai, an AI plugin for the fish shell."
 readme = "README.md"
@@ -22,6 +22,7 @@ dependencies = [
   "hugchat==0.4.11",
   "mistralai==1.0.2",
   "binaryornot==0.4.4",
+  "azure-ai-inference==1.0.0b4",
 ]
 
 [project.urls]
diff --git a/src/fish_ai/engine.py b/src/fish_ai/engine.py
index 1bb56d4..4130405 100644
--- a/src/fish_ai/engine.py
+++ b/src/fish_ai/engine.py
@@ -19,6 +19,8 @@
 from mistralai import Mistral
 from fish_ai.redact import redact
 import itertools
+from azure.ai.inference import ChatCompletionsClient
+from azure.core.credentials import AzureKeyCredential
 
 config = ConfigParser()
 config.read(path.expanduser('~/.config/fish-ai.ini'))
@@ -249,6 +251,19 @@ def get_response(messages):
             temperature=float(get_config('temperature') or '0.2'),
         )
         response = completions.choices[0].message.content.strip(' `')
+    elif get_config('provider') == 'github':
+        client = ChatCompletionsClient(
+            endpoint='https://models.inference.ai.azure.com',
+            credential=AzureKeyCredential(get_config('api_key')),
+        )
+
+        completions = client.complete(
+            messages=messages,
+            model=get_config('model') or 'Meta-Llama-3.1-70B-Instruct',
+            temperature=float(get_config('temperature') or '0.2'),
+            max_tokens=1024
+        )
+        response = completions.choices[0].message.content.strip(' `')
     else:
         completions = get_openai_client().chat.completions.create(
             model=get_config('model') or 'gpt-4o',