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

wxo_agentic_evaluation/service_provider/model_proxy_provider.py

@@ -1,13 +1,21 @@
+import json
+import logging
 import os
 import time
+import uuid
 from threading import Lock
-from typing import List, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 import requests
 
-from wxo_agentic_evaluation.service_provider.provider import Provider
+from wxo_agentic_evaluation.service_provider.provider import (
+    ChatResult,
+    Provider,
+)
 from wxo_agentic_evaluation.utils.utils import is_ibm_cloud_url
 
+logger = logging.getLogger(__name__)
+
 AUTH_ENDPOINT_AWS = (
     "https://iam.platform.saas.ibm.com/siusermgr/api/1.0/apikeys/token"
 )
@@ -15,7 +23,6 @@ AUTH_ENDPOINT_IBM_CLOUD = "https://iam.cloud.ibm.com/identity/token"
 DEFAULT_PARAM = {
     "min_new_tokens": 1,
     "decoding_method": "greedy",
-    "max_new_tokens": 400,
 }
 
 
@@ -36,17 +43,62 @@ def _normalize_cpd_auth_url(url: str) -> str:
     return url
 
 
+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_chat(params: Dict[str, Any] | None) -> Dict[str, Any]:
+    # Translate legacy generation params to chat.completions params.
+    p = params or {}
+    out: Dict[str, Any] = {}
+
+    passthrough = {
+        "temperature",
+        "top_p",
+        "n",
+        "stream",
+        "stop",
+        "presence_penalty",
+        "frequency_penalty",
+        "logit_bias",
+        "user",
+        "max_tokens",
+        "seed",
+        "response_format",
+    }
+    for k in passthrough:
+        if k in p:
+            out[k] = p[k]
+
+    if "max_new_tokens" in p and "max_tokens" not in out:
+        out["max_tokens"] = p["max_new_tokens"]
+
+    return out
+
+
 class ModelProxyProvider(Provider):
     def __init__(
         self,
-        model_id=None,
-        api_key=None,
-        instance_url=None,
-        timeout=300,
-        embedding_model_id=None,
-        params=None,
+        model_id: Optional[str] = None,
+        api_key: Optional[str] = None,
+        instance_url: Optional[str] = None,
+        timeout: int = 300,
+        embedding_model_id: Optional[str] = None,
+        params: Optional[Dict[str, Any]] = None,
+        use_legacy_query: Optional[
+            bool
+        ] = None,  # Provider routes query() to old/new based on this
+        system_prompt: Optional[str] = None,
+        token: Optional[str] = None,
     ):
-        super().__init__()
+        super().__init__(use_legacy_query=use_legacy_query)
 
         instance_url = os.environ.get("WO_INSTANCE", instance_url)
         if not instance_url:
@@ -56,6 +108,7 @@ class ModelProxyProvider(Provider):
 
         self.timeout = timeout
         self.model_id = os.environ.get("MODEL_OVERRIDE", model_id)
+        logger.info("[d b]Using inference model %s", self.model_id)
         self.embedding_model_id = embedding_model_id
 
         self.api_key = os.environ.get("WO_API_KEY", api_key)
@@ -106,14 +159,17 @@ class ModelProxyProvider(Provider):
                     "WO_API_KEY must be specified for SaaS or IBM IAM auth"
                 )
 
+        # Endpoints
         self.url = (
             self.instance_url + "/ml/v1/text/generation?version=2024-05-01"
-        )
+        )  # legacy
+        self.chat_url = self.instance_url + "/ml/v1/chat/completions"  # chat
         self.embedding_url = self.instance_url + "/ml/v1/text/embeddings"
 
         self.lock = Lock()
         self.token, self.refresh_time = self.get_token()
         self.params = params if params else DEFAULT_PARAM
+        self.system_prompt = system_prompt
 
     def _resolve_auth_mode_and_url(
         self, explicit_auth_url: str | None
@@ -254,6 +310,7 @@ class ModelProxyProvider(Provider):
             json=payload,
             headers=headers,
             verify=self._wo_ssl_verify,
+            timeout=self.timeout,
         )
 
         if resp.status_code == 200:
@@ -262,9 +319,11 @@ class ModelProxyProvider(Provider):
 
         resp.raise_for_status()
 
-    def query(self, sentence: str) -> str:
+    def old_query(self, sentence: str) -> str:
+        # Legacy /ml/v1/text/generation
         if self.model_id is None:
             raise Exception("model id must be specified for text generation")
+
         self.refresh_token_if_expires()
         headers = self.get_header()
         payload = {
@@ -274,18 +333,357 @@ class ModelProxyProvider(Provider):
             "timeout": self.timeout,
             "parameters": self.params,
         }
-        resp = requests.post(
-            self.url, json=payload, headers=headers, verify=self._wo_ssl_verify
+
+        request_id = str(uuid.uuid4())
+        start_time = time.time()
+
+        # Input logging
+        logger.debug(
+            "[d][b]Sending text.generation request | request_id=%s url=%s model=%s space_id=%s params=%s input_preview=%s",
+            request_id,
+            self.url,
+            self.model_id,
+            self.space_id,
+            json.dumps(
+                payload.get("parameters", {}),
+                sort_keys=True,
+                ensure_ascii=False,
+            ),
+            _truncate(sentence, 200),
         )
-        if resp.status_code == 200:
-            return resp.json()["results"][0]["generated_text"]
 
-        resp.raise_for_status()
+        resp = None
+        try:
+            resp = requests.post(
+                self.url,
+                json=payload,
+                headers=headers,
+                verify=self._wo_ssl_verify,
+                timeout=self.timeout,
+            )
+
+            duration_ms = int((time.time() - start_time) * 1000)
+            resp.raise_for_status()
+            data = resp.json()
+
+            result = (
+                data["results"][0]
+                if "results" in data and data["results"]
+                else data
+            )
+            output_text = (
+                (
+                    result.get("generated_text")
+                    if isinstance(result, dict)
+                    else None
+                )
+                or (result.get("message") if isinstance(result, dict) else None)
+                or ""
+            )
+
+            # Usage (best-effort)
+            usage = data.get("usage") or {}
+            if not usage and isinstance(result, dict):
+                in_tok = result.get("input_token_count")
+                out_tok = result.get("generated_token_count") or result.get(
+                    "output_token_count"
+                )
+                if in_tok is not None or out_tok is not None:
+                    usage = {
+                        "prompt_tokens": in_tok,
+                        "completion_tokens": out_tok,
+                        "total_tokens": (in_tok or 0) + (out_tok or 0),
+                    }
+
+            api_request_id = resp.headers.get(
+                "x-request-id"
+            ) or resp.headers.get("request-id")
+
+            # Output logging
+            logger.debug(
+                "[d][b]text.generation response received | request_id=%s status_code=%s duration_ms=%s usage=%s output_preview=%s api_request_id=%s",
+                request_id,
+                resp.status_code,
+                duration_ms,
+                json.dumps(usage, sort_keys=True, ensure_ascii=False),
+                _truncate(output_text, 2000),
+                api_request_id,
+            )
+
+            if output_text:
+                return output_text
+            else:
+                raise ValueError(
+                    f"Unexpected response from legacy endpoint: {data}"
+                )
+
+        except Exception as e:
+            duration_ms = int((time.time() - start_time) * 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(
+                "text.generation request failed | request_id=%s status_code=%s duration_ms=%s response_text_preview=%s",
+                request_id,
+                status_code,
+                duration_ms,
+                resp_text_preview,
+            )
+
+            with self.lock:
+                if (
+                    "authentication_token_expired" in str(e)
+                    or status_code == 401
+                ):
+                    self.token, self.refresh_time = self.get_token()
+            raise
+
+    def new_query(self, sentence: str) -> str:
+        """
+        New /ml/v1/chat/completions
+        Returns assistant message content of the first choice.
+        """
+        if self.model_id is None:
+            raise Exception("model id must be specified for text generation")
+
+        self.refresh_token_if_expires()
+        headers = self.get_header()
+
+        messages: List[Dict[str, Any]] = []
+        if getattr(self, "system_prompt", None):
+            messages.append({"role": "system", "content": self.system_prompt})
+        messages.append(
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": sentence}],
+            }
+        )
+
+        chat_params = _translate_params_to_chat(self.params)
+        if isinstance(self.params, dict) and "time_limit" in self.params:
+            chat_params["time_limit"] = self.params["time_limit"]
+
+        payload: Dict[str, Any] = {
+            "model_id": self.model_id,
+            "space_id": self.space_id,
+            "messages": messages,
+            **chat_params,
+        }
+
+        url = f"{self.instance_url}/ml/v1/text/chat?version=2024-10-08"
+        request_id = str(uuid.uuid4())
+        start_time = time.time()
+
+        logger.debug(
+            "[d][b]Sending chat.completions request | request_id=%s url=%s model=%s space_id=%s params=%s input_preview=%s",
+            request_id,
+            url,
+            self.model_id,
+            self.space_id,
+            json.dumps(chat_params, sort_keys=True, ensure_ascii=False),
+            _truncate(sentence, 200),
+        )
+
+        resp = None
+        try:
+            resp = requests.post(
+                url,
+                json=payload,
+                headers=headers,
+                verify=self._wo_ssl_verify,
+                timeout=self.timeout,
+            )
+            duration_ms = int((time.time() - start_time) * 1000)
+            resp.raise_for_status()
+            data = resp.json()
+
+            choice = data["choices"][0]
+            content = choice["message"]["content"]
+            finish_reason = choice.get("finish_reason")
+            usage = data.get("usage", {})
+            api_request_id = resp.headers.get(
+                "x-request-id"
+            ) or resp.headers.get("request-id")
+
+            logger.debug(
+                "[d][b]chat.completions response received | request_id=%s status_code=%s duration_ms=%s finish_reason=%s usage=%s output_preview=%s api_request_id=%s",
+                request_id,
+                resp.status_code,
+                duration_ms,
+                finish_reason,
+                json.dumps(usage, sort_keys=True, ensure_ascii=False),
+                _truncate(content, 2000),
+                api_request_id,
+            )
+
+            return content
+
+        except Exception as e:
+            duration_ms = int((time.time() - start_time) * 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(
+                "chat.completions request failed | request_id=%s status_code=%s duration_ms=%s response_text_preview=%s",
+                request_id,
+                status_code,
+                duration_ms,
+                resp_text_preview,
+            )
+
+            with self.lock:
+                if (
+                    "authentication_token_expired" in str(e)
+                    or status_code == 401
+                ):
+                    self.token, self.refresh_time = self.get_token()
+            raise
+
+    def chat(
+        self,
+        messages: Sequence[Dict[str, str]],
+        params: Optional[Dict[str, Any]] = None,
+    ) -> ChatResult:
+        # Non-streaming chat using /ml/v1/chat/completions.
+        if self.model_id is None:
+            raise Exception("model id must be specified for chat")
+
+        self.refresh_token_if_expires()
+        headers = self.get_header()
+
+        # Convert messages to watsonx format: user content is typed list
+        wx_messages: List[Dict[str, Any]] = []
+        for m in messages:
+            role = m.get("role")
+            content = m.get("content", "")
+            if role == "user" and isinstance(content, str):
+                wx_messages.append(
+                    {
+                        "role": "user",
+                        "content": [{"type": "text", "text": content}],
+                    }
+                )
+            else:
+                wx_messages.append({"role": role, "content": content})
+
+        merged_params = dict(self.params or {})
+        if params:
+            merged_params.update(params)
+        chat_params = _translate_params_to_chat(merged_params)
+        chat_params.pop("stream", None)  # force non-streaming
+        if "time_limit" in merged_params:
+            chat_params["time_limit"] = merged_params["time_limit"]
+
+        payload: Dict[str, Any] = {
+            "model_id": self.model_id,
+            "space_id": self.space_id,
+            "messages": wx_messages,
+            **chat_params,
+        }
+
+        url = f"{self.instance_url}/ml/v1/text/chat?version=2024-10-08"
+
+        last_user = next(
+            (
+                m.get("content", "")
+                for m in reversed(messages)
+                if m.get("role") == "user"
+            ),
+            "",
+        )
+        request_id = str(uuid.uuid4())
+        start_time = time.time()
+
+        logger.debug(
+            "[d][b]Sending chat.completions request (non-streaming) | request_id=%s url=%s model=%s space_id=%s params=%s input_preview=%s",
+            request_id,
+            url,
+            self.model_id,
+            self.space_id,
+            json.dumps(chat_params, sort_keys=True, ensure_ascii=False),
+            _truncate(last_user, 200),
+        )
+
+        resp = None
+        try:
+            resp = requests.post(
+                url,
+                json=payload,
+                headers=headers,
+                verify=self._wo_ssl_verify,
+                timeout=self.timeout,
+            )
+            duration_ms = int((time.time() - start_time) * 1000)
+            resp.raise_for_status()
+            data = resp.json()
+
+            choice = data["choices"][0]
+            content = choice["message"]["content"]
+            finish_reason = choice.get("finish_reason")
+            usage = data.get("usage", {})
+            api_request_id = resp.headers.get(
+                "x-request-id"
+            ) or resp.headers.get("request-id")
+
+            logger.debug(
+                "[d][b]chat.completions response received (non-streaming) | request_id=%s status_code=%s duration_ms=%s finish_reason=%s usage=%s output_preview=%s api_request_id=%s",
+                request_id,
+                resp.status_code,
+                duration_ms,
+                finish_reason,
+                json.dumps(usage, sort_keys=True, ensure_ascii=False),
+                _truncate(content, 2000),
+                api_request_id,
+            )
+
+            return ChatResult(
+                text=content, usage=usage, finish_reason=finish_reason, raw=data
+            )
+
+        except Exception as e:
+            duration_ms = int((time.time() - start_time) * 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(
+                "chat.completions request failed (non-streaming) | request_id=%s status_code=%s duration_ms=%s response_text_preview=%s",
+                request_id,
+                status_code,
+                duration_ms,
+                resp_text_preview,
+            )
+            with self.lock:
+                if (
+                    "authentication_token_expired" in str(e)
+                    or status_code == 401
+                ):
+                    self.token, self.refresh_time = self.get_token()
+            raise
 
 
 if __name__ == "__main__":
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s %(levelname)s %(name)s %(message)s",
+    )
+
     provider = ModelProxyProvider(
         model_id="meta-llama/llama-3-3-70b-instruct",
         embedding_model_id="ibm/slate-30m-english-rtrvr",
+        use_legacy_query=False,
+        system_prompt="",
     )
+    # Base class will route .query() to new_query() by default (unless USE_LEGACY_QUERY=true)
     print(provider.query("ok"))