Local Ollama intent extractor with schema output

# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Sarang Datye
"""Local Ollama intent extraction with deterministic validation."""

from __future__ import annotations

from dataclasses import dataclass
import json
import os
import socket
import time
from typing import Any
import urllib.error
import urllib.request

from deterministic_capabilities import NormalizedRequest, validate_request


MODEL_NAME = os.environ.get("OLLAMA_MODEL", "qwen3:8b")
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434/api/chat")
REQUEST_TIMEOUT_SECONDS = 30

SYSTEM_PROMPT = """You are an intent and parameter extractor.
Do not answer the user's request or execute any requested operation.
Return one JSON object only, with exactly: intent, confidence, parameters.
Use only: calculator, email_validator, create_support_ticket, general_explanation, unsupported.
Use unsupported with empty parameters when required values cannot be safely extracted.
Never invent required parameters.
"""


def request_variant(
    intent: str, parameter_properties: dict[str, Any], required: list[str]
) -> dict[str, Any]:
    return {
        "type": "object",
        "additionalProperties": False,
        "required": ["intent", "confidence", "parameters"],
        "properties": {
            "intent": {"const": intent},
            "confidence": {"type": "number", "minimum": 0, "maximum": 1},
            "parameters": {
                "type": "object",
                "additionalProperties": False,
                "required": required,
                "properties": parameter_properties,
            },
        },
    }


INTENT_RESPONSE_SCHEMA = {
    "oneOf": [
        request_variant(
            "calculator",
            {
                "operation": {"type": "string", "enum": ["add", "subtract", "multiply", "divide"]},
                "left": {"type": "number"},
                "right": {"type": "number"},
            },
            ["operation", "left", "right"],
        ),
        request_variant("email_validator", {"email": {"type": "string", "minLength": 1}}, ["email"]),
        request_variant(
            "create_support_ticket",
            {
                "application": {"type": "string", "minLength": 1},
                "description": {"type": "string", "minLength": 1},
                "priority": {"type": "string", "minLength": 1},
            },
            ["application", "description", "priority"],
        ),
        request_variant("general_explanation", {"prompt": {"type": "string"}}, ["prompt"]),
        request_variant("unsupported", {}, []),
    ]
}


class OllamaClientError(RuntimeError):
    """Base error for local Ollama intent extraction."""


class OllamaResponseError(OllamaClientError):
    """Raised when Ollama returns an unusable response."""


class OllamaJSONError(OllamaClientError):
    """Raised when model content is not exactly one JSON object."""


@dataclass(frozen=True)
class InvocationMetadata:
    model: str
    elapsed_ms: float
    prompt_tokens: int | None
    completion_tokens: int | None


@dataclass(frozen=True)
class IntentInvocation:
    request: dict[str, Any]
    metadata: InvocationMetadata


def build_payload(user_input: str) -> dict[str, Any]:
    return {
        "model": MODEL_NAME,
        "stream": False,
        "think": False,
        "format": INTENT_RESPONSE_SCHEMA,
        "options": {"temperature": 0},
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_input},
        ],
    }


def post_chat(payload: dict[str, Any]) -> tuple[dict[str, Any], float]:
    request = urllib.request.Request(
        OLLAMA_URL,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    started = time.perf_counter()
    try:
        with urllib.request.urlopen(request, timeout=REQUEST_TIMEOUT_SECONDS) as response:
            envelope = json.load(response)
    except urllib.error.HTTPError as error:
        raise OllamaClientError(f"Ollama returned HTTP {error.code}: {error.reason}") from error
    except (TimeoutError, socket.timeout) as error:
        raise OllamaClientError(f"Ollama timed out after {REQUEST_TIMEOUT_SECONDS} seconds.") from error
    except urllib.error.URLError as error:
        raise OllamaClientError(f"Cannot reach local Ollama: {error.reason}") from error
    except json.JSONDecodeError as error:
        raise OllamaResponseError("Ollama response envelope is not valid JSON.") from error
    elapsed_ms = (time.perf_counter() - started) * 1000
    if not isinstance(envelope, dict):
        raise OllamaResponseError("Ollama response envelope must be an object.")
    return envelope, elapsed_ms


def token_count(envelope: dict[str, Any], field: str) -> int | None:
    value = envelope.get(field)
    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
        return None
    return value


def request_dict(normalized: NormalizedRequest) -> dict[str, Any]:
    return {
        "intent": normalized.intent,
        "confidence": normalized.confidence,
        "parameters": normalized.parameters,
    }


def extract_intent_detailed(user_input: str) -> IntentInvocation:
    if not isinstance(user_input, str):
        raise TypeError("user_input must be a string.")
    if not user_input.strip():
        raise ValueError("user_input must be a nonempty string.")

    envelope, elapsed_ms = post_chat(build_payload(user_input))
    message = envelope.get("message")
    if not isinstance(message, dict) or not isinstance(message.get("content"), str):
        raise OllamaResponseError("Ollama response is missing message.content.")
    try:
        parsed = json.loads(message["content"])
    except json.JSONDecodeError as error:
        raise OllamaJSONError("Ollama model content must be one unfenced JSON object.") from error
    if not isinstance(parsed, dict):
        raise OllamaJSONError("Ollama model content must be one JSON object.")

    normalized = validate_request(parsed)
    return IntentInvocation(
        request=request_dict(normalized),
        metadata=InvocationMetadata(
            model=envelope.get("model") if isinstance(envelope.get("model"), str) else MODEL_NAME,
            elapsed_ms=elapsed_ms,
            prompt_tokens=token_count(envelope, "prompt_eval_count"),
            completion_tokens=token_count(envelope, "eval_count"),
        ),
    )


def extract_intent(user_input: str) -> dict[str, Any]:
    return extract_intent_detailed(user_input).request

python3 -m py_compile deterministic_capabilities.py ollama_intent_extractor.py