kplam
51dc8dc4f6
- FastAPI + Qdrant + Redis + Ollama 技术栈 - L1→L2→L3→chunk 四层分层检索（dense + sparse RRF 融合） - 文档三级总结与 2.5 级回退 - query 解析路由与分类 - /admin 管理页面
63 lines
1.8 KiB
Python
63 lines
1.8 KiB
Python
"""Ollama 客户端
|
||
|
||
通过 Ollama HTTP API 调用本地模型进行文本生成。
|
||
"""
|
||
|
||
import httpx
|
||
import structlog
|
||
|
||
from app.config import settings
|
||
|
||
logger = structlog.get_logger()
|
||
|
||
|
||
class OllamaClient:
|
||
"""Ollama HTTP API 客户端"""
|
||
|
||
def __init__(
|
||
self,
|
||
base_url: str | None = None,
|
||
model: str | None = None,
|
||
timeout: float = 120.0,
|
||
) -> None:
|
||
self.base_url = (base_url or settings.ollama_base_url).rstrip("/")
|
||
self.model = model or settings.ollama_model
|
||
self.timeout = timeout
|
||
|
||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||
"""调用 Ollama 生成文本
|
||
|
||
Args:
|
||
prompt: 输入提示词
|
||
json_mode: 为 True 时启用 Ollama 原生 JSON 约束输出（payload 加 format: json）
|
||
|
||
Returns:
|
||
生成的文本内容
|
||
"""
|
||
url = f"{self.base_url}/api/generate"
|
||
payload = {
|
||
"model": self.model,
|
||
"prompt": prompt,
|
||
"stream": False,
|
||
}
|
||
if json_mode:
|
||
payload["format"] = "json"
|
||
|
||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||
resp = await client.post(url, json=payload)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
|
||
result = data.get("response", "")
|
||
logger.debug("Ollama 生成完成", model=self.model, output_length=len(result))
|
||
return result
|
||
|
||
async def is_available(self) -> bool:
|
||
"""检查 Ollama 服务是否可用"""
|
||
try:
|
||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||
resp = await client.get(f"{self.base_url}/api/tags")
|
||
return resp.status_code == 200
|
||
except httpx.HTTPError:
|
||
return False
|