lambdagent/agentexample/agent67/core/ollama_lam.py

"""
agent67.core.ollama_lam — 基于 Ollama 本地模型的 Lam 实现

无需 API Key，本地运行，完全离线可用。

Lambda 语义不变:
    OllamaLam("name", "prompt") ≡ λ_D . F_{ollama,D}
    调用 = β-规约 = ollama.chat 解码

用法:
    brain = OllamaLam("assistant", "你是一个助手", model="qwen2.5:7b")
    result = brain("帮我查看文件", ctx)
"""
from __future__ import annotations

import json
import time
import urllib.request
import urllib.error
from typing import Any, Callable, Optional

import sys
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from lambdagent.core import Term, Context


OLLAMA_BASE_URL = "http://localhost:11434"


def ollama_available(base_url=OLLAMA_BASE_URL):
    """检测 Ollama 服务是否在运行。"""
    try:
        req = urllib.request.Request("{}/api/tags".format(base_url))
        with urllib.request.urlopen(req, timeout=3) as resp:
            data = json.loads(resp.read())
            return [m["name"] for m in data.get("models", [])]
    except Exception:
        return []


class OllamaLam(Term):
    """
    基于 Ollama 本地模型的 Lambda 抽象。

    通过 Ollama HTTP API (localhost:11434) 调用本地模型，
    支持 Qwen2.5、GLM-4、Llama3 等所有 Ollama 支持的模型。

    Lambda: OllamaLam("name", "prompt", model="qwen2.5:7b") = λx. ollama(prompt, x)
    """

    def __init__(
        self,
        name: str,
        prompt: str,
        model: str = "qwen2.5:7b",
        temperature: float = 0.3,
        max_tokens: int = 4096,
        base_url: str = OLLAMA_BASE_URL,
        output_parser: Callable[[str], Any] | None = None,
    ):
        super().__init__(name)
        self.prompt = prompt
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.base_url = base_url
        self.output_parser = output_parser or (lambda x: x)

    def apply(self, input: Any, ctx: Context | None = None) -> Any:
        """β-规约: (λ_D x) → ollama.chat(prompt, x)"""
        ctx = ctx or Context()
        t0 = time.time()

        raw = self._call_ollama(str(input))
        duration = (time.time() - t0) * 1000
        result = self.output_parser(raw)

        ctx.log(self._name, self._trace_id, input, result, duration,
                "ollama/{}".format(self.model))
        return result

    def _call_ollama(self, input_text: str) -> str:
        """通过 Ollama HTTP API 执行 β-规约。"""
        url = "{}/api/chat".format(self.base_url)

        body = json.dumps({
            "model": self.model,
            "messages": [
                {"role": "system", "content": self.prompt},
                {"role": "user", "content": input_text},
            ],
            "options": {
                "temperature": self.temperature,
                "num_predict": self.max_tokens,
            },
            "stream": False,
        }).encode("utf-8")

        req = urllib.request.Request(
            url, data=body,
            headers={"Content-Type": "application/json"},
        )

        try:
            with urllib.request.urlopen(req, timeout=120) as resp:
                data = json.loads(resp.read())
                message = data.get("message", {})
                return message.get("content", "").strip() or "[无输出]"
        except urllib.error.URLError as e:
            return "[Ollama 错误] 连接失败: {}。请确保 ollama serve 在运行".format(e)
        except Exception as e:
            return "[Ollama 错误] {}".format(e)

    def __rshift__(self, other):
        """支持 >> 组合"""
        from lambdagent.primitives import Compose
        return Compose(self, other)