| from __future__ import annotations |
|
|
| import json |
| import os |
| from urllib.error import HTTPError, URLError |
| from urllib.request import Request, urlopen |
|
|
| from prompt_injection_framework.models.base import BaseModelAdapter, ModelRequest |
|
|
|
|
| class OllamaModelAdapter(BaseModelAdapter): |
| provider = "ollama" |
|
|
| def __init__(self, model_name: str, base_url: str | None = None) -> None: |
| super().__init__(model_name=model_name) |
| configured_base = base_url or os.getenv("OLLAMA_BASE_URL") or "http://127.0.0.1:11434" |
| self.base_url = configured_base.rstrip("/") |
| self.api_url = f"{self.base_url}/api/generate" |
|
|
| def _generate_text(self, request: ModelRequest) -> tuple[str, dict[str, object]]: |
| prompt_sections: list[str] = [] |
| if request.context: |
| prompt_sections.append(f"[SYSTEM]\n{request.context}") |
| if request.conversation_history: |
| history_lines = [] |
| for turn in request.conversation_history: |
| history_lines.append(f"[{turn.role.upper()}]\n{turn.content}") |
| prompt_sections.append("\n\n".join(history_lines)) |
|
|
| user_parts = [request.prompt] |
| if request.task_input: |
| user_parts.append(request.task_input) |
| user_content = "\n\n".join(part for part in user_parts if part) |
| prompt_sections.append(f"[USER]\n{user_content}") |
|
|
| payload = { |
| "model": self.model_name, |
| "prompt": "\n\n".join(section for section in prompt_sections if section), |
| "stream": False, |
| "options": { |
| "temperature": request.temperature, |
| "num_predict": request.max_tokens, |
| }, |
| } |
| if request.seed is not None: |
| payload["options"]["seed"] = request.seed |
|
|
| req = Request( |
| self.api_url, |
| data=json.dumps(payload).encode("utf-8"), |
| headers={ |
| "Content-Type": "application/json", |
| "Accept": "application/json", |
| "User-Agent": "llm-prompt-injection-security-eval-framework/0.1", |
| }, |
| method="POST", |
| ) |
|
|
| try: |
| with urlopen(req, timeout=120) as response: |
| raw_response = json.loads(response.read().decode("utf-8")) |
| except HTTPError as exc: |
| detail = exc.read().decode("utf-8", errors="replace") |
| raise RuntimeError(f"Ollama request failed with status {exc.code}: {detail}") from exc |
| except URLError as exc: |
| raise RuntimeError(f"Ollama request failed: {exc.reason}") from exc |
|
|
| text = str(raw_response.get("response", "")) |
| if not text.strip(): |
| raise RuntimeError("Ollama returned an empty response.") |
| return text, raw_response |
|
|
