ibm-watsonx-orchestrate-evaluation-framework 1.1.6__py3-none-any.whl → 1.1.7__py3-none-any.whl

wxo_agentic_evaluation/service_provider/ollama_provider.py

@@ -1,39 +1,410 @@
 import json
+import logging
 import os
-from typing import List
+import time
+import uuid
+from typing import Any, Dict, Iterator, List, Optional, Sequence
 
 import requests
 
-from wxo_agentic_evaluation.service_provider.provider import Provider
+from wxo_agentic_evaluation.service_provider.provider import (
+    ChatResult,
+    Provider,
+)
+
+logger = logging.getLogger(__name__)
 
 OLLAMA_URL = os.environ.get("OLLAMA_HOST", "http://localhost:11434")
 
 
+def _truncate(value: Any, max_len: int = 1000) -> str:
+    if value is None:
+        return ""
+    s = str(value)
+    return (
+        s
+        if len(s) <= max_len
+        else s[:max_len] + f"... [truncated {len(s) - max_len} chars]"
+    )
+
+
+def _translate_params_to_ollama_options(
+    params: Optional[Dict[str, Any]]
+) -> Dict[str, Any]:
+    """
+    Map generic params to Ollama 'options' field.
+    Ollama options docs: https://github.com/ollama/ollama/blob/main/docs/modelfile.md#parameters
+    """
+    p = params or {}
+    out: Dict[str, Any] = {}
+
+    for key in ("temperature", "top_p", "top_k", "stop", "seed"):
+        if key in p:
+            out[key] = p[key]
+
+    if "max_new_tokens" in p:
+        out["num_predict"] = p["max_new_tokens"]
+    elif "max_tokens" in p:
+        out["num_predict"] = p["max_tokens"]
+
+    if "repeat_penalty" in p:
+        out["repeat_penalty"] = p["repeat_penalty"]
+    if "repeat_last_n" in p:
+        out["repeat_last_n"] = p["repeat_last_n"]
+
+    return out
+
+
 class OllamaProvider(Provider):
-    def __init__(self, model_id=None):
-        self.url = OLLAMA_URL + "/api/generate"
-        self.model_id = model_id
-        super().__init__()
-
-    def query(self, sentence: str) -> str:
-        payload = {"model": self.model_id, "prompt": sentence}
-        resp = requests.post(self.url, json=payload, stream=True)
+    def __init__(
+        self,
+        model_id: Optional[str] = None,
+        params: Optional[Dict[str, Any]] = None,
+        timeout: int = 300,
+        use_legacy_query: Optional[bool] = None,
+        system_prompt: Optional[str] = None,
+        token: Optional[str] = None,
+        instance_url: Optional[str] = None,
+    ):
+        super().__init__(use_legacy_query=use_legacy_query)
+        self.generate_url = (
+            OLLAMA_URL.rstrip("/") + "/api/generate"
+        )  # legacy text generation
+        self.chat_url = OLLAMA_URL.rstrip("/") + "/api/chat"  # chat endpoint
+        self.model_id = os.environ.get("MODEL_OVERRIDE", model_id)
+        logger.info("[d b]Using inference model %s", self.model_id)
+        self.params = params or {}
+        self.timeout = timeout
+        self.system_prompt = system_prompt
+
+    def old_query(self, sentence: str) -> str:
+        # Legacy /api/generate
+        if not self.model_id:
+            raise ValueError("model_id must be specified for Ollama generation")
+
+        options = _translate_params_to_ollama_options(self.params)
+        payload: Dict[str, Any] = {
+            "model": self.model_id,
+            "prompt": sentence,
+            "stream": True,
+        }
+        if options:
+            payload["options"] = options
+
+        request_id = str(uuid.uuid4())
+        t0 = time.time()
+
+        logger.debug(
+            "[d][b]Sending Ollama generate request | request_id=%s url=%s model=%s params=%s input_preview=%s",
+            request_id,
+            self.generate_url,
+            self.model_id,
+            json.dumps(options, sort_keys=True, ensure_ascii=False),
+            _truncate(sentence, 200),
+        )
+
+        resp = None
         final_text = ""
-        data = b""
-        for chunk in resp:
-            data += chunk
-            if data.endswith(b"\n"):
-                json_obj = json.loads(data)
-                if not json_obj["done"] and json_obj["response"]:
-                    final_text += json_obj["response"]
-                data = b""
+        usage: Dict[str, Any] = {}
+
+        try:
+            resp = requests.post(
+                self.generate_url,
+                json=payload,
+                stream=True,
+                timeout=self.timeout,
+            )
+
+            if resp.status_code != 200:
+                resp_text_preview = _truncate(getattr(resp, "text", ""), 2000)
+                duration_ms = int((time.time() - t0) * 1000)
+                logger.error(
+                    "[d b red]Ollama generate request failed (non-200) | request_id=%s status_code=%s duration_ms=%s response_text_preview=%s",
+                    request_id,
+                    resp.status_code,
+                    duration_ms,
+                    resp_text_preview,
+                )
+                resp.raise_for_status()
+
+            for line in resp.iter_lines(decode_unicode=True):
+                if not line:
+                    continue
+                try:
+                    obj = json.loads(line)
+                except Exception:
+                    logger.warning(
+                        "Skipping unparsable line from Ollama generate | request_id=%s line_preview=%s",
+                        request_id,
+                        _truncate(line, 500),
+                    )
+                    continue
+
+                if not obj.get("done"):
+                    chunk = obj.get("response", "")
+                    if chunk:
+                        final_text += chunk
+                else:
+                    # Final metrics frame
+                    usage = {
+                        "prompt_eval_count": obj.get("prompt_eval_count"),
+                        "eval_count": obj.get("eval_count"),
+                        "prompt_eval_duration_ns": obj.get(
+                            "prompt_eval_duration"
+                        ),
+                        "eval_duration_ns": obj.get("eval_duration"),
+                        "total_duration_ns": obj.get("total_duration"),
+                        "load_duration_ns": obj.get("load_duration"),
+                    }
+
+            duration_ms = int((time.time() - t0) * 1000)
+            logger.debug(
+                "[d][b]Ollama generate response received | request_id=%s status_code=%s duration_ms=%s usage=%s output_preview=%s",
+                request_id,
+                resp.status_code,
+                duration_ms,
+                json.dumps(usage, sort_keys=True, ensure_ascii=False),
+                _truncate(final_text, 2000),
+            )
+
+            return final_text
+
+        except Exception:
+            duration_ms = int((time.time() - t0) * 1000)
+            status_code = getattr(resp, "status_code", None)
+            resp_text_preview = None
+            try:
+                if resp is not None and not getattr(resp, "raw", None):
+                    resp_text_preview = _truncate(
+                        getattr(resp, "text", None), 2000
+                    )
+            except Exception:
+                pass
+
+            logger.exception(
+                "Ollama generate request encountered an error | request_id=%s status_code=%s duration_ms=%s response_text_preview=%s",
+                request_id,
+                status_code,
+                duration_ms,
+                resp_text_preview,
+            )
+            raise
+
+    def new_query(self, sentence: str) -> str:
+        """
+        /api/chat
+        Returns assistant message content.
+        """
+        if not self.model_id:
+            raise ValueError("model_id must be specified for Ollama chat")
+
+        options = _translate_params_to_ollama_options(self.params)
 
-        return final_text
+        messages: List[Dict[str, str]] = []
+        if self.system_prompt:
+            messages.append({"role": "system", "content": self.system_prompt})
+        messages.append({"role": "user", "content": sentence})
+
+        payload: Dict[str, Any] = {
+            "model": self.model_id,
+            "messages": messages,
+            "stream": False,
+        }
+        if options:
+            payload["options"] = options
+
+        request_id = str(uuid.uuid4())
+        t0 = time.time()
+
+        logger.debug(
+            "[d][b]Sending Ollama chat request (non-streaming) | request_id=%s url=%s model=%s params=%s input_preview=%s",
+            request_id,
+            self.chat_url,
+            self.model_id,
+            json.dumps(options, sort_keys=True, ensure_ascii=False),
+            _truncate(sentence, 200),
+        )
+
+        resp = None
+        try:
+            resp = requests.post(
+                self.chat_url, json=payload, timeout=self.timeout
+            )
+            duration_ms = int((time.time() - t0) * 1000)
+            resp.raise_for_status()
+            data = resp.json()
+
+            # Non-streaming chat response: { "message": {"role": "assistant", "content": "..."} , "done": true, ... }
+            message = data.get("message") or {}
+            content = message.get("content", "") or ""
+            finish_reason = data.get("finish_reason")
+            usage = {
+                "prompt_eval_count": data.get("prompt_eval_count"),
+                "eval_count": data.get("eval_count"),
+                "prompt_eval_duration_ns": data.get("prompt_eval_duration"),
+                "eval_duration_ns": data.get("eval_duration"),
+                "total_duration_ns": data.get("total_duration"),
+                "load_duration_ns": data.get("load_duration"),
+            }
+
+            logger.debug(
+                "[d][b]Ollama chat response received | request_id=%s status_code=%s duration_ms=%s finish_reason=%s usage=%s output_preview=%s",
+                request_id,
+                resp.status_code,
+                duration_ms,
+                finish_reason,
+                json.dumps(usage, sort_keys=True, ensure_ascii=False),
+                _truncate(content, 2000),
+            )
+
+            return content
+
+        except Exception:
+            duration_ms = int((time.time() - t0) * 1000)
+            status_code = getattr(resp, "status_code", None)
+            resp_text_preview = (
+                _truncate(getattr(resp, "text", None), 2000)
+                if resp is not None
+                else None
+            )
+
+            logger.exception(
+                "Ollama chat request encountered an error | request_id=%s status_code=%s duration_ms=%s response_text_preview=%s",
+                request_id,
+                status_code,
+                duration_ms,
+                resp_text_preview,
+            )
+            raise
+
+    def chat(
+        self,
+        messages: Sequence[Dict[str, str]],
+        params: Optional[Dict[str, Any]] = None,
+    ) -> ChatResult:
+        """
+        Non-streaming chat via /api/chat.
+        """
+        if not self.model_id:
+            raise ValueError("model_id must be specified for Ollama chat")
+
+        merged_params = dict(self.params or {})
+        if params:
+            merged_params.update(params)
+        options = _translate_params_to_ollama_options(merged_params)
+
+        payload: Dict[str, Any] = {
+            "model": self.model_id,
+            "messages": list(messages),
+            "stream": False,
+        }
+        if options:
+            payload["options"] = options
+
+        last_user = next(
+            (
+                m.get("content", "")
+                for m in reversed(messages)
+                if m.get("role") == "user"
+            ),
+            "",
+        )
+        request_id = str(uuid.uuid4())
+        t0 = time.time()
+
+        logger.debug(
+            "[d][b]Sending Ollama chat request (non-streaming, multi-message) | request_id=%s url=%s model=%s params=%s input_preview=%s",
+            request_id,
+            self.chat_url,
+            self.model_id,
+            json.dumps(options, sort_keys=True, ensure_ascii=False),
+            _truncate(last_user, 200),
+        )
+
+        resp = None
+        try:
+            resp = requests.post(
+                self.chat_url, json=payload, timeout=self.timeout
+            )
+            duration_ms = int((time.time() - t0) * 1000)
+            resp.raise_for_status()
+            data = resp.json()
+
+            message = data.get("message") or {}
+            content = message.get("content", "") or ""
+            finish_reason = data.get("finish_reason")
+            usage = {
+                "prompt_eval_count": data.get("prompt_eval_count"),
+                "eval_count": data.get("eval_count"),
+                "prompt_eval_duration_ns": data.get("prompt_eval_duration"),
+                "eval_duration_ns": data.get("eval_duration"),
+                "total_duration_ns": data.get("total_duration"),
+                "load_duration_ns": data.get("load_duration"),
+            }
+
+            logger.debug(
+                "[d][b]Ollama chat response received (non-streaming, multi-message) | request_id=%s status_code=%s duration_ms=%s finish_reason=%s usage=%s output_preview=%s",
+                request_id,
+                resp.status_code,
+                duration_ms,
+                finish_reason,
+                json.dumps(usage, sort_keys=True, ensure_ascii=False),
+                _truncate(content, 2000),
+            )
+
+            return ChatResult(
+                text=content, usage=usage, finish_reason=finish_reason, raw=data
+            )
+
+        except Exception:
+            duration_ms = int((time.time() - t0) * 1000)
+            status_code = getattr(resp, "status_code", None)
+            resp_text_preview = (
+                _truncate(getattr(resp, "text", None), 2000)
+                if resp is not None
+                else None
+            )
+
+            logger.exception(
+                "Ollama chat request (non-streaming, multi-message) encountered an error | request_id=%s status_code=%s duration_ms=%s response_text_preview=%s",
+                request_id,
+                status_code,
+                duration_ms,
+                resp_text_preview,
+            )
+            raise
 
     def encode(self, sentences: List[str]) -> List[list]:
-        pass
+        raise NotImplementedError(
+            "encode is not implemented for OllamaProvider"
+        )
 
 
 if __name__ == "__main__":
-    provider = OllamaProvider(model_id="llama3.1:8b")
-    print(provider.query("ok"))
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s %(levelname)s %(name)s %(message)s",
+    )
+
+    provider = OllamaProvider(model_id="llama3.1:8b", use_legacy_query=False)
+
+    print("new_query:", provider.query("Say hello in one sentence."))
+
+    # chat API
+    messages = [
+        {"role": "system", "content": "You are concise."},
+        {"role": "user", "content": "List three fruits."},
+    ]
+    result = provider.chat(messages)
+    print("chat:", result.text)
+
+    # Streaming chat
+    print("stream_chat:")
+    assembled = []
+    for chunk in provider.stream_chat(
+        [{"role": "user", "content": "Stream a short sentence."}]
+    ):
+        if chunk.get("delta"):
+            assembled.append(chunk["delta"])
+        if chunk.get("is_final"):
+            print("".join(assembled))

wxo_agentic_evaluation/service_provider/provider.py

@@ -1,18 +1,138 @@
-from abc import ABC, abstractmethod
-from typing import List
+from __future__ import annotations
 
+import logging
+import os
+from abc import ABC, ABCMeta, abstractmethod
+from dataclasses import dataclass
+from threading import Lock
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
-class Provider(ABC):
-    def __init__(self):
-        pass
+from wxo_agentic_evaluation.type import ProviderInstancesCacheKey
+
+
+class SingletonProviderMeta(type):
+
+    _provider_instances: Dict[str, "Provider"] = {}
+    _instantiation_lock = Lock()
+
+    def __call__(cls, *args, **kwargs):
+
+        key_str: str = str(cls._get_key(cls.__name__, args, kwargs))
+
+        if key_str not in cls._provider_instances:
+            with cls._instantiation_lock:
+                if key_str not in cls._provider_instances:
+                    cls._provider_instances[key_str] = super().__call__(
+                        *args, **kwargs
+                    )
+
+        return cls._provider_instances[key_str]
+
+    @staticmethod
+    def _get_key(
+        provider: str, args: Tuple[Any, ...], kwargs: Dict[str, Any]
+    ) -> ProviderInstancesCacheKey:
+
+        args_str = str(args) if args else "noargs"
+        kwargs_str = str(sorted(kwargs.items())) if kwargs else "nokwargs"
+
+        return ProviderInstancesCacheKey(
+            provider=provider,
+            hashed_args=args_str,
+            hashed_kwargs=kwargs_str,
+        )
+
+
+class SingletonProviderABCMeta(ABCMeta, SingletonProviderMeta):
+    pass
+
+
+@dataclass
+class ChatResult:
+    text: str
+    usage: Optional[Dict[str, Any]] = None
+    finish_reason: Optional[str] = None
+    raw: Optional[Any] = None
+
+
+class Provider(ABC, metaclass=SingletonProviderABCMeta):
+    def __init__(
+        self,
+        use_legacy_query: Optional[bool] = None,
+        logger: Optional[logging.Logger] = None,
+    ) -> None:
+        self.logger = logger or logging.getLogger(self.__class__.__name__)
+
+        env_use_legacy = os.environ.get("USE_LEGACY_QUERY")
+        if env_use_legacy is not None:
+            self.use_legacy_query: bool = env_use_legacy.strip().lower() in (
+                "1",
+                "true",
+                "yes",
+                "on",
+            )
+        else:
+            self.use_legacy_query = (
+                bool(use_legacy_query) if use_legacy_query is not None else True
+            )
+        if self.use_legacy_query:
+            self.logger.debug("[d][b]Using legacy /text/generation queries")
+        else:
+            self.logger.debug("[d][b]Using new /chat/completions queries")
 
     @abstractmethod
-    def query(self, sentence: str) -> str:
-        pass
+    def old_query(self, sentence: str) -> str:
+        raise NotImplementedError
 
-    def batch_query(self, sentences: List[str]) -> List[str]:
-        return [self.query(sentence) for sentence in sentences]
+    @abstractmethod
+    def new_query(self, sentence: str) -> str:
+        raise NotImplementedError
 
     @abstractmethod
     def encode(self, sentences: List[str]) -> List[list]:
-        pass
+        raise NotImplementedError
+
+    def query(self, sentence: str) -> str:
+        if self.use_legacy_query:
+            return self.old_query(sentence)
+        return self.new_query(sentence)
+
+    def chat(
+        self,
+        messages: Sequence[Dict[str, str]],
+        params: Optional[Dict[str, Any]] = None,
+    ) -> ChatResult:
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not implement chat()."
+        )
+
+    def batch_query(
+        self,
+        sentences: List[str],
+        max_workers: Optional[int] = None,
+    ) -> List[str]:
+        if not sentences:
+            return []
+
+        if not max_workers or max_workers <= 1:
+            return [self.query(sentence) for sentence in sentences]
+
+        from concurrent.futures import ThreadPoolExecutor, as_completed
+
+        results: List[Optional[str]] = [None] * len(sentences)
+        with ThreadPoolExecutor(max_workers=max_workers) as pool:
+            future_to_idx = {
+                pool.submit(self.query, s): i for i, s in enumerate(sentences)
+            }
+            for fut in as_completed(future_to_idx):
+                idx = future_to_idx[fut]
+                results[idx] = fut.result()
+
+        return [r if r is not None else "" for r in results]
+
+    def set_routing(self, use_legacy_query: Optional[bool] = None) -> None:
+        if use_legacy_query is not None:
+            self.use_legacy_query = bool(use_legacy_query)
+
+    def close(self) -> None:
+        return

wxo_agentic_evaluation/service_provider/referenceless_provider_wrapper.py

@@ -1,9 +1,14 @@
+import uuid
 from abc import ABC, abstractmethod
 from typing import Any, List, Mapping, Optional, Union
 
 import requests
 import rich
 
+from wxo_agentic_evaluation.service_provider.gateway_provider import (
+    GatewayProvider,
+    _translate_params_to_chat,
+)
 from wxo_agentic_evaluation.service_provider.model_proxy_provider import (
     ModelProxyProvider,
 )
@@ -149,3 +154,50 @@ class WatsonXLLMKitWrapper(WatsonXProvider, LLMKitWrapper):
             return resp.json()
         else:
             resp.raise_for_status()
+
+
+class GatewayProviderLLMKitWrapper(GatewayProvider, LLMKitWrapper):
+    def chat(self, sentence: Union[str, List[Mapping[str, str]]]):
+        if isinstance(sentence, str):
+            messages = []
+            if self.system_prompt:
+                messages.append(
+                    {"role": "system", "content": self.system_prompt}
+                )
+            messages.append({"role": "user", "content": sentence})
+        else:
+            messages = sentence
+
+        if self.model_id is None:
+            raise Exception("model id must be specified for text generation")
+
+        self.refresh_token_if_expires()
+
+        merged_params = dict(self.params or {})
+        chat_params = _translate_params_to_chat(merged_params)
+        chat_params.pop("stream", None)
+
+        override_params = dict(merged_params)
+        override_params["model"] = self.model_id
+
+        payload = {
+            "model": self._payload_model_str(self.model_id),
+            "messages": list(messages),
+            **chat_params,
+        }
+
+        request_id = str(uuid.uuid4())
+        headers = self._headers(request_id, override_params)
+
+        resp = requests.post(
+            self.chat_url,
+            json=payload,
+            headers=headers,
+            verify=self._wo_ssl_verify,
+            timeout=self.timeout,
+        )
+
+        if resp.status_code == 200:
+            return resp.json()
+        else:
+            resp.raise_for_status()