ibm-watsonx-orchestrate-evaluation-framework 1.0.3__py3-none-any.whl → 1.1.8b0__py3-none-any.whl

wxo_agentic_evaluation/service_provider/ollama_provider.py

@@ -1,40 +1,410 @@
-import requests
 import json
-from wxo_agentic_evaluation.service_provider.provider import Provider
-from typing import List
+import logging
 import os
+import time
+import uuid
+from typing import Any, Dict, Iterator, List, Optional, Sequence
+
+import requests
+
+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
+        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,
     ):
-        self.url = OLLAMA_URL + "/api/generate"
-        self.model_id = model_id
-        super().__init__()
+        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")
 
-    def query(self, sentence: str) -> str:
-        payload = {"model": self.model_id, "prompt": sentence}
-        resp = requests.post(self.url, json=payload, stream=True)
+        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/portkey_provider.py

@@ -0,0 +1,229 @@
+import logging
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+from wxo_agentic_evaluation.service_provider.provider import (
+    ChatResult,
+    Provider,
+)
+
+
+logger = logging.getLogger(__name__)
+try:
+    from portkey_ai.api_resources.types.chat_complete_type import (
+        ChatCompletions,
+    )
+except Exception as e:
+    logger.warning(e)
+
+
+def _extract_text_from_response(resp: Any) -> str:
+    """Extract assistant text from common Portkey response shapes.
+
+    The exact return type from the client may vary. Try several common
+    patterns and fall back to str(resp).
+    """
+    try:
+        # Common pattern like OpenAI: choices[0].message.content (string or list)
+        if isinstance(resp, dict):
+            choices = resp.get("choices")
+            if choices and len(choices) > 0:
+                choice = choices[0]
+                msg = choice.get("message") or choice.get("delta") or {}
+                if isinstance(msg, dict):
+                    content = msg.get("content")
+                    # content might be a string or a list of content blocks
+                    if isinstance(content, str):
+                        return content
+                    if isinstance(content, list) and content:
+                        # content blocks might be {"type":"text","text":...}
+                        first = content[0]
+                        if isinstance(first, dict) and "text" in first:
+                            return first.get("text", "")
+                        return str(first)
+
+                # fallback: some clients return choices[0].text
+                if "text" in choice:
+                    return choice.get("text") or ""
+
+        # If not a dict, try objects with attributes
+        if hasattr(resp, "choices"):
+            choices = getattr(resp, "choices")
+            if choices:
+                c0 = choices[0]
+                if hasattr(c0, "message") and getattr(c0, "message"):
+                    m = getattr(c0, "message")
+                    if isinstance(m, dict):
+                        return _extract_text_from_response({"choices": [m]})
+                    # message may be an object; try to get content attr
+                    if hasattr(m, "content"):
+                        return getattr(m, "content")
+
+    except Exception:
+        # parsing should never raise to caller; fall through to str(resp)
+        pass
+
+    # Last resort
+    try:
+        return str(resp)
+    except Exception:
+        return ""
+
+
+class PortkeyProvider(Provider):
+    """Provider that delegates to the Portkey AI client"""
+
+    def __init__(
+        self,
+        provider: str,
+        api_key: Optional[str] = None,
+        model_id: Optional[str] = None,
+        embedding_model: Optional[str] = None,
+        base_url: Optional[str] = None,
+        timeout: int = 60,
+        system_prompt: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+
+        self.provider = provider
+        self.api_key = api_key
+        self.model_id = model_id
+        self.embedding_model = embedding_model
+        self.base_url = base_url
+        self.timeout = timeout * 1000  # convert to ms
+        self.system_prompt = system_prompt
+
+        # Lazy import - avoid hard dependency at import time
+        self._client = None
+        if self.api_key is not None:
+            try:
+
+                from portkey_ai import Portkey  # type: ignore
+
+                client_kwargs = {
+                    "provider": self.provider,
+                    "Authorization": self.api_key,
+                }
+                if self.base_url:
+                    client_kwargs["base_url"] = base_url
+                if self.timeout:
+                    client_kwargs["request_timeout"] = self.timeout
+                # Add any remaining kwargs
+                client_kwargs.update(kwargs)
+                # construct client
+                self._client = Portkey(**client_kwargs)
+            except Exception as e:  # ImportError or runtime errors
+                # Do not fail hard on import; surface when used
+                logger.debug("portkey_ai import/initialization failed: %s", e)
+                self._client = None
+
+    def _require_client(self) -> None:
+        if self._client is None:
+            raise ImportError(
+                "portkey_ai client is not available. Install 'portkey_ai' and provide a valid api_key."
+            )
+
+    def old_query(self, sentence: str, extract_text: bool = False) -> str:
+        return self.new_query(sentence, extract_text)
+
+    def new_query(self, sentence: str, extract_text: bool = False) -> str:
+        """Send a single user message and return assistant text."""
+        self._require_client()
+
+        messages = []
+        if self.system_prompt:
+            messages.append({"role": "system", "content": self.system_prompt})
+        messages.append({"role": "user", "content": sentence})
+
+        resp = self._client.chat.completions.create(
+            messages=messages, model=self.model_id
+        )
+
+        if extract_text:
+            return _extract_text_from_response(resp)
+
+        return resp
+
+    def chat(
+        self,
+        messages: Sequence[Dict[str, str]],
+        params: Optional[Dict[str, Any]] = None,
+        return_chat_completions: bool = True,
+    ) -> Union[ChatResult, ChatCompletions]:
+        self._require_client()
+
+        # build messages for Portkey: pass them mostly through
+        port_messages = []
+        for m in messages:
+            # Portkey expects simple role/content pairs for chat
+            role = m.get("role")
+            content = m.get("content", "")
+            port_messages.append({"role": role, "content": content})
+
+        kwargs = {}
+        if params:
+            kwargs.update(params)
+
+        try:
+            resp = self._client.chat.completions.create(
+                messages=port_messages, model=self.model_id, **kwargs
+            )
+        except TypeError:
+            # fallback if client signature differs
+            resp = self._client.chat.completions.create(
+                messages=port_messages, model=self.model_id
+            )
+
+        if return_chat_completions:
+            return resp
+
+        # try to extract text, usage and finish reason
+        text = _extract_text_from_response(resp)
+
+        usage = None
+        finish_reason = None
+        if isinstance(resp, dict):
+            usage = resp.get("usage")
+            try:
+                finish_reason = resp.get("choices", [])[0].get("finish_reason")
+            except Exception:
+                finish_reason = None
+
+        return ChatResult(
+            text=text, usage=usage, finish_reason=finish_reason, raw=resp
+        )
+
+    def encode(self, sentences: List[str]) -> List[list]:
+        if self.embedding_model is None:
+            raise Exception(
+                "embedding model id must be specified for text encoding"
+            )
+
+        self._require_client()
+
+        try:
+            resp = self._client.embeddings.create(
+                inputs=sentences, model=self.embedding_model
+            )
+        except TypeError:
+            resp = self._client.embeddings.create(
+                inputs=sentences, model=self.embedding_model
+            )
+
+        # Try common shapes: {'data': [{'embedding': [...]}, ...]} or {'results': ...}
+        if isinstance(resp, dict):
+            if "data" in resp:
+                return [d.get("embedding") for d in resp.get("data", [])]
+            if "results" in resp:
+                return [r.get("embedding") for r in resp.get("results", [])]
+
+        # If the client returns list directly
+        if isinstance(resp, list):
+            # expect list of embeddings
+            return resp
+
+        # Unknown shape -> try to coerce
+        try:
+            return [list(e) for e in resp]
+        except Exception:
+            raise ValueError("Unexpected response from embeddings request")