From 3dc413638bc22ae68e3b0157b55cbcb8a2262954 Mon Sep 17 00:00:00 2001
From: RockChinQ <1010553892@qq.com>
Date: Tue, 9 Jul 2024 23:37:34 +0800
Subject: [PATCH] =?UTF-8?q?feat(ollama):=20=E9=85=8D=E7=BD=AE=E6=96=87?=
 =?UTF-8?q?=E4=BB=B6=E8=BF=81=E7=A7=BB?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../m010_ollama_requester_config.py           | 23 +++++++++++++++++++
 pkg/core/stages/migrate.py                    |  1 +
 templates/provider.json                       |  5 ++++
 3 files changed, 29 insertions(+)
 create mode 100644 pkg/config/migrations/m010_ollama_requester_config.py

diff --git a/pkg/config/migrations/m010_ollama_requester_config.py b/pkg/config/migrations/m010_ollama_requester_config.py
new file mode 100644
index 00000000..56e49663
--- /dev/null
+++ b/pkg/config/migrations/m010_ollama_requester_config.py
@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from .. import migration
+
+
+@migration.migration_class("ollama-requester-config", 10)
+class MsgTruncatorConfigMigration(migration.Migration):
+    """迁移"""
+
+    async def need_migrate(self) -> bool:
+        """判断当前环境是否需要运行此迁移"""
+        return 'ollama-chat' not in self.ap.provider_cfg.data['requester']
+
+    async def run(self):
+        """执行迁移"""
+        
+        self.ap.provider_cfg.data['requester']['ollama-chat'] = {
+            "base-url": "http://127.0.0.1:11434",
+            "args": {},
+            "timeout": 600
+        }
+
+        await self.ap.provider_cfg.dump_config()
diff --git a/pkg/core/stages/migrate.py b/pkg/core/stages/migrate.py
index 2ad1e974..862d90af 100644
--- a/pkg/core/stages/migrate.py
+++ b/pkg/core/stages/migrate.py
@@ -6,6 +6,7 @@ from .. import stage, app
 from ...config import migration
 from ...config.migrations import m001_sensitive_word_migration, m002_openai_config_migration, m003_anthropic_requester_cfg_completion, m004_moonshot_cfg_completion
 from ...config.migrations import m005_deepseek_cfg_completion, m006_vision_config, m007_qcg_center_url, m008_ad_fixwin_config_migrate, m009_msg_truncator_cfg
+from ...config.migrations import m010_ollama_requester_config
 
 
 @stage.stage_class("MigrationStage")
diff --git a/templates/provider.json b/templates/provider.json
index 309fb827..32878fe9 100644
--- a/templates/provider.json
+++ b/templates/provider.json
@@ -37,6 +37,11 @@
             "base-url": "https://api.deepseek.com",
             "args": {},
             "timeout": 120
+        },
+        "ollama-chat": {
+            "base-url": "http://127.0.0.1:11434",
+            "args": {},
+            "timeout": 600
         }
     },
     "model": "gpt-3.5-turbo",