mirror of
https://github.com/langbot-app/LangBot.git
synced 2025-11-25 03:15:06 +08:00
122 lines
4.5 KiB
Python
122 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import typing
|
|
import traceback
|
|
|
|
import ollama
|
|
from .. import operator, entities, errors
|
|
|
|
|
|
@operator.operator_class(
|
|
name="ollama",
|
|
help="ollama平台操作",
|
|
usage="!ollama\n!ollama show <模型名>\n!ollama pull <模型名>\n!ollama del <模型名>"
|
|
)
|
|
class OllamaOperator(operator.CommandOperator):
|
|
async def execute(
|
|
self, context: entities.ExecuteContext
|
|
) -> typing.AsyncGenerator[entities.CommandReturn, None]:
|
|
try:
|
|
content: str = '模型列表:\n'
|
|
model_list: list = ollama.list().get('models', [])
|
|
for model in model_list:
|
|
content += f"名称: {model['name']}\n"
|
|
content += f"修改时间: {model['modified_at']}\n"
|
|
content += f"大小: {bytes_to_mb(model['size'])}MB\n\n"
|
|
yield entities.CommandReturn(text=f"{content.strip()}")
|
|
except ollama.ResponseError as e:
|
|
yield entities.CommandReturn(error=errors.CommandError(f"无法获取模型列表，请确认 Ollama 服务正常"))
|
|
|
|
|
|
def bytes_to_mb(num_bytes):
|
|
mb: float = num_bytes / 1024 / 1024
|
|
return format(mb, '.2f')
|
|
|
|
|
|
@operator.operator_class(
|
|
name="show",
|
|
help="ollama模型详情",
|
|
privilege=2,
|
|
parent_class=OllamaOperator
|
|
)
|
|
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)
|
|
yield entities.CommandReturn(text=content.strip())
|
|
except ollama.ResponseError as e:
|
|
yield entities.CommandReturn(error=errors.CommandError(f"无法获取模型详情，请确认 Ollama 服务正常"))
|
|
|
|
@operator.operator_class(
|
|
name="pull",
|
|
help="ollama模型拉取",
|
|
privilege=2,
|
|
parent_class=OllamaOperator
|
|
)
|
|
class OllamaPullOperator(operator.CommandOperator):
|
|
async def execute(
|
|
self, context: entities.ExecuteContext
|
|
) -> typing.AsyncGenerator[entities.CommandReturn, None]:
|
|
try:
|
|
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
|
|
except ollama.ResponseError as e:
|
|
yield entities.CommandReturn(error=errors.CommandError(f"无法获取模型列表，请确认 Ollama 服务正常"))
|
|
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="del",
|
|
help="ollama模型删除",
|
|
privilege=2,
|
|
parent_class=OllamaOperator
|
|
)
|
|
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)
|