151 lines
4.3 KiB
Python
151 lines
4.3 KiB
Python
"""
|
|
Ollama LLM API 服務模組
|
|
支援一般請求與串流模式
|
|
"""
|
|
|
|
import requests
|
|
import json
|
|
from typing import Generator, Optional
|
|
from config import Config
|
|
|
|
|
|
class LLMService:
|
|
"""Ollama API 服務封裝"""
|
|
|
|
def __init__(self, api_url: str = None, default_model: str = None):
|
|
self.api_url = api_url or Config.OLLAMA_API_URL
|
|
self.default_model = default_model or Config.OLLAMA_DEFAULT_MODEL
|
|
|
|
def get_available_models(self) -> list:
|
|
"""取得可用模型列表"""
|
|
try:
|
|
response = requests.get(f"{self.api_url}/v1/models", timeout=10)
|
|
response.raise_for_status()
|
|
models = response.json()
|
|
return [m['id'] for m in models.get('data', [])]
|
|
except requests.RequestException as e:
|
|
raise LLMServiceError(f"無法取得模型列表: {e}")
|
|
|
|
def chat(
|
|
self,
|
|
messages: list,
|
|
model: str = None,
|
|
temperature: float = 0.7,
|
|
system_prompt: str = None
|
|
) -> str:
|
|
"""
|
|
發送聊天請求 (非串流)
|
|
|
|
Args:
|
|
messages: 訊息列表 [{"role": "user", "content": "..."}]
|
|
model: 模型名稱
|
|
temperature: 溫度參數 (0-1)
|
|
system_prompt: 系統提示詞
|
|
|
|
Returns:
|
|
AI 回應內容
|
|
"""
|
|
model = model or self.default_model
|
|
|
|
# 加入系統提示詞
|
|
if system_prompt:
|
|
messages = [{"role": "system", "content": system_prompt}] + messages
|
|
|
|
payload = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"stream": False
|
|
}
|
|
|
|
try:
|
|
response = requests.post(
|
|
f"{self.api_url}/v1/chat/completions",
|
|
json=payload,
|
|
timeout=60
|
|
)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
return result['choices'][0]['message']['content']
|
|
except requests.RequestException as e:
|
|
raise LLMServiceError(f"聊天請求失敗: {e}")
|
|
|
|
def chat_stream(
|
|
self,
|
|
messages: list,
|
|
model: str = None,
|
|
temperature: float = 0.7,
|
|
system_prompt: str = None
|
|
) -> Generator[str, None, None]:
|
|
"""
|
|
發送聊天請求 (串流模式)
|
|
|
|
Args:
|
|
messages: 訊息列表
|
|
model: 模型名稱
|
|
temperature: 溫度參數
|
|
system_prompt: 系統提示詞
|
|
|
|
Yields:
|
|
串流回應的每個片段
|
|
"""
|
|
model = model or self.default_model
|
|
|
|
if system_prompt:
|
|
messages = [{"role": "system", "content": system_prompt}] + messages
|
|
|
|
payload = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"stream": True
|
|
}
|
|
|
|
try:
|
|
response = requests.post(
|
|
f"{self.api_url}/v1/chat/completions",
|
|
json=payload,
|
|
stream=True,
|
|
timeout=120
|
|
)
|
|
response.raise_for_status()
|
|
|
|
for line in response.iter_lines():
|
|
if line:
|
|
if line.startswith(b"data: "):
|
|
data_str = line[6:].decode('utf-8')
|
|
if data_str.strip() != "[DONE]":
|
|
try:
|
|
data = json.loads(data_str)
|
|
if 'choices' in data:
|
|
delta = data['choices'][0].get('delta', {})
|
|
if 'content' in delta:
|
|
yield delta['content']
|
|
except json.JSONDecodeError:
|
|
continue
|
|
except requests.RequestException as e:
|
|
raise LLMServiceError(f"串流請求失敗: {e}")
|
|
|
|
def simple_query(self, prompt: str, system_prompt: str = None) -> str:
|
|
"""
|
|
簡單查詢 (單一問題)
|
|
|
|
Args:
|
|
prompt: 使用者問題
|
|
system_prompt: 系統提示詞
|
|
|
|
Returns:
|
|
AI 回應
|
|
"""
|
|
messages = [{"role": "user", "content": prompt}]
|
|
return self.chat(messages, system_prompt=system_prompt)
|
|
|
|
|
|
class LLMServiceError(Exception):
|
|
"""LLM 服務錯誤"""
|
|
pass
|
|
|
|
|
|
# 全域實例
|
|
llm_service = LLMService()
|