From 2bdc3468d1727d3d6288fe20ad4e3e1e93bc2576 Mon Sep 17 00:00:00 2001
From: canyuan <canyuan0801@outlook.com>
Date: Mon, 8 Jul 2024 21:08:07 +0800
Subject: [PATCH] add ollama cmd

---
 pkg/command/cmdmgr.py           |   2 +-
 pkg/command/operators/ollama.py | 112 ++++++++++++++++++++++++++++++++
 2 files changed, 113 insertions(+), 1 deletion(-)
 create mode 100644 pkg/command/operators/ollama.py

diff --git a/pkg/command/cmdmgr.py b/pkg/command/cmdmgr.py
index 93ed8f8..1d7b92f 100644
--- a/pkg/command/cmdmgr.py
+++ b/pkg/command/cmdmgr.py
@@ -8,7 +8,7 @@ from . import entities, operator, errors
 from ..config import manager as cfg_mgr
 
 # 引入所有算子以便注册
-from .operators import func, plugin, default, reset, list as list_cmd, last, next, delc, resend, prompt, cmd, help, version, update
+from .operators import func, plugin, default, reset, list as list_cmd, last, next, delc, resend, prompt, cmd, help, version, update, ollama
 
 
 class CommandManager:
diff --git a/pkg/command/operators/ollama.py b/pkg/command/operators/ollama.py
new file mode 100644
index 0000000..3932b17
--- /dev/null
+++ b/pkg/command/operators/ollama.py
@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import json
+import typing
+
+import ollama
+from .. import operator, entities
+
+
+@operator.operator_class(
+    name="ollama_list",
+    help="ollama模型列表",
+    usage="!ollama_list"
+)
+class OllamaListOperator(operator.CommandOperator):
+    async def execute(
+            self, context: entities.ExecuteContext
+    ) -> typing.AsyncGenerator[entities.CommandReturn, None]:
+        content: str = '模型列表:\n'
+        model_list: list = ollama.list().get('models', [])
+        for model in model_list:
+            content += f"name: {model['name']}\n"
+            content += f"modified_at: {model['modified_at']}\n"
+            content += f"size: {bytes_to_mb(model['size'])}mb\n\n"
+        yield entities.CommandReturn(text=f"{content}")
+
+
+def bytes_to_mb(num_bytes):
+    mb: float = num_bytes / 1024 / 1024
+    return format(mb, '.2f')
+
+
+@operator.operator_class(
+    name="ollama_show",
+    help="ollama模型详情",
+    usage="!ollama_show <模型名>"
+)
+class OllamaShowOperator(operator.CommandOperator):
+    async def execute(
+            self, context: entities.ExecuteContext
+    ) -> typing.AsyncGenerator[entities.CommandReturn, None]:
+        content: str = '模型详情:\n'
+        try:
+            show: dict = ollama.show(model=context.crt_params[0])
+            model_info: dict = show.get('model_info', {})
+            ignore_show: str = 'too long to show...'
+
+            for key in ['license', 'modelfile']:
+                show[key] = ignore_show
+
+            for key in ['tokenizer.chat_template.rag', 'tokenizer.chat_template.tool_use']:
+                model_info[key] = ignore_show
+
+            content += json.dumps(show, indent=4)
+        except ollama.ResponseError as e:
+            content = f"{e.error}"
+
+        yield entities.CommandReturn(text=content)
+
+
+@operator.operator_class(
+    name="ollama_pull",
+    help="ollama模型拉取",
+    usage="!ollama_pull <模型名>"
+)
+class OllamaPullOperator(operator.CommandOperator):
+    async def execute(
+            self, context: entities.ExecuteContext
+    ) -> typing.AsyncGenerator[entities.CommandReturn, None]:
+        model_list: list = ollama.list().get('models', [])
+        if context.crt_params[0] in [model['name'] for model in model_list]:
+            yield entities.CommandReturn(text="模型已存在")
+            return
+
+        on_progress: bool = False
+        progress_count: int = 0
+        try:
+            for resp in ollama.pull(model=context.crt_params[0], stream=True):
+                total: typing.Any = resp.get('total')
+                if not on_progress:
+                    if total is not None:
+                        on_progress = True
+                    yield entities.CommandReturn(text=resp.get('status'))
+                else:
+                    if total is None:
+                        on_progress = False
+
+                    completed: typing.Any = resp.get('completed')
+                    if isinstance(completed, int) and isinstance(total, int):
+                        percentage_completed = (completed / total) * 100
+                        if percentage_completed > progress_count:
+                            progress_count += 10
+                            yield entities.CommandReturn(
+                                text=f"下载进度: {completed}/{total} = {percentage_completed:.2f}%")
+        except ollama.ResponseError as e:
+            yield entities.CommandReturn(text=f"拉取失败: {e.error}")
+
+
+@operator.operator_class(
+    name="ollama_del",
+    help="ollama模型删除",
+    usage="!ollama_del <模型名>"
+)
+class OllamaDelOperator(operator.CommandOperator):
+    async def execute(
+            self, context: entities.ExecuteContext
+    ) -> typing.AsyncGenerator[entities.CommandReturn, None]:
+        try:
+            ret: str = ollama.delete(model=context.crt_params[0])['status']
+        except ollama.ResponseError as e:
+            ret = f"{e.error}"
+        yield entities.CommandReturn(text=ret)