ibm-watsonx-orchestrate-evaluation-framework 1.1.1__py3-none-any.whl → 1.1.2__py3-none-any.whl

wxo_agentic_evaluation/service_provider/model_proxy_provider.py

@@ -1,8 +1,7 @@
 import os
 import requests
 import time
-
-from typing import List
+from typing import List, Tuple
 from threading import Lock
 
 from wxo_agentic_evaluation.service_provider.provider import Provider
@@ -12,6 +11,22 @@ AUTH_ENDPOINT_AWS = "https://iam.platform.saas.ibm.com/siusermgr/api/1.0/apikeys
 AUTH_ENDPOINT_IBM_CLOUD = "https://iam.cloud.ibm.com/identity/token"
 DEFAULT_PARAM = {"min_new_tokens": 1, "decoding_method": "greedy", "max_new_tokens": 400}
 
+def _infer_cpd_auth_url(instance_url: str) -> str:
+    inst = (instance_url or "").rstrip("/")
+    if not inst:
+        return "/icp4d-api/v1/authorize"
+    if "/orchestrate" in inst:
+        base = inst.split("/orchestrate", 1)[0].rstrip("/")
+        return base + "/icp4d-api/v1/authorize"
+    return inst + "/icp4d-api/v1/authorize"
+
+
+def _normalize_cpd_auth_url(url: str) -> str:
+    u = (url or "").rstrip("/")
+    if u.endswith("/icp4d-api"):
+        return u + "/v1/authorize"
+    return url
+
 
 class ModelProxyProvider(Provider):
     def __init__(
@@ -26,20 +41,43 @@ class ModelProxyProvider(Provider):
         super().__init__()
 
         instance_url = os.environ.get("WO_INSTANCE", instance_url)
-        api_key = os.environ.get("WO_API_KEY", api_key)
-        if not instance_url or not api_key:
-            raise RuntimeError("instance url and WO apikey must be specified to use WO model proxy")
+        if not instance_url:
+            raise RuntimeError("instance url must be specified to use WO model proxy")
 
         self.timeout = timeout
-        self.model_id = model_id
-
+        self.model_id = os.environ.get("MODEL_OVERRIDE",model_id)
         self.embedding_model_id = embedding_model_id
 
-        self.api_key = api_key
+        self.api_key = os.environ.get("WO_API_KEY", api_key)
+        self.username = os.environ.get("WO_USERNAME", None)
+        self.password = os.environ.get("WO_PASSWORD", None)
+        self.auth_type = os.environ.get("WO_AUTH_TYPE", "").lower()  # explicit override if set, otherwise inferred- match ADK values
+        explicit_auth_url = os.environ.get("AUTHORIZATION_URL", None)
+
         self.is_ibm_cloud = is_ibm_cloud_url(instance_url)
-        self.auth_url = AUTH_ENDPOINT_IBM_CLOUD if self.is_ibm_cloud else AUTH_ENDPOINT_AWS
-    
-        self.instance_url = instance_url
+        self.instance_url = instance_url.rstrip("/")
+
+        self.auth_mode, self.auth_url = self._resolve_auth_mode_and_url(explicit_auth_url=explicit_auth_url)
+        self._wo_ssl_verify = os.environ.get("WO_SSL_VERIFY", "true").lower() != "false"
+        env_space_id = os.environ.get("WATSONX_SPACE_ID", None)
+        if self.auth_mode == "cpd":
+            if not env_space_id or not env_space_id.strip():
+                raise RuntimeError("CPD mode requires WATSONX_SPACE_ID environment variable to be set")
+            self.space_id = env_space_id.strip()
+        else:
+            self.space_id = (env_space_id.strip() if env_space_id and env_space_id.strip() else "1")
+
+        if self.auth_mode == "cpd":
+            if "/orchestrate" in self.instance_url:
+                self.instance_url = self.instance_url.split("/orchestrate", 1)[0].rstrip("/")
+            if not self.username:
+                raise RuntimeError("CPD auth requires WO_USERNAME to be set")
+            if not (self.password or self.api_key):
+                raise RuntimeError("CPD auth requires either WO_PASSWORD or WO_API_KEY to be set (with WO_USERNAME)")
+        else:
+            if not self.api_key:
+                raise RuntimeError("WO_API_KEY must be specified for SaaS or IBM IAM auth")
+
         self.url = self.instance_url + "/ml/v1/text/generation?version=2024-05-01"
         self.embedding_url = self.instance_url + "/ml/v1/text/embeddings"
 
@@ -47,20 +85,85 @@ class ModelProxyProvider(Provider):
         self.token, self.refresh_time = self.get_token()
         self.params = params if params else DEFAULT_PARAM
 
-    def get_token(self):
+    def _resolve_auth_mode_and_url(
+        self,
+        explicit_auth_url: str | None
+    ) -> Tuple[str, str]:
+        """
+        Returns (auth_mode, auth_url)
+        - auth_mode: "cpd" | "ibm_iam" | "saas"
+        """
+        if explicit_auth_url:
+            if "/icp4d-api" in explicit_auth_url:
+                return "cpd", _normalize_cpd_auth_url(explicit_auth_url)
+            if self.auth_type == "ibm_iam":
+                return "ibm_iam", explicit_auth_url
+            elif self.auth_type == "saas":
+                return "saas", explicit_auth_url
+            else:
+                mode = "ibm_iam" if self.is_ibm_cloud else "saas"
+                return mode, explicit_auth_url
+
+        if self.auth_type == "cpd":
+            inferred_cpd_url = _infer_cpd_auth_url(self.instance_url)
+            return "cpd", inferred_cpd_url
+        if self.auth_type == "ibm_iam":
+            return "ibm_iam", AUTH_ENDPOINT_IBM_CLOUD
+        if self.auth_type == "saas":
+            return "saas", AUTH_ENDPOINT_AWS
+
+        if "/orchestrate" in self.instance_url:
+            inferred_cpd_url = _infer_cpd_auth_url(self.instance_url)
+            return "cpd", inferred_cpd_url
+
         if self.is_ibm_cloud:
-            payload = {"grant_type": "urn:ibm:params:oauth:grant-type:apikey", "apikey": self.api_key}
-            resp = requests.post(self.auth_url, data=payload)
-            token_key = "access_token"
+            return "ibm_iam", AUTH_ENDPOINT_IBM_CLOUD
         else:
-            payload = {"apikey": self.api_key}
-            resp = requests.post(self.auth_url, json=payload)
-            token_key = "token"
+            return "saas", AUTH_ENDPOINT_AWS
+
+    def get_token(self):
+        headers = {}
+        post_args = {}
+        timeout = 10
+        exchange_url = self.auth_url
+
+        if self.auth_mode == "ibm_iam":
+            headers = {"Accept": "application/json", "Content-Type": "application/x-www-form-urlencoded"}
+            form_data = {
+                "grant_type": "urn:ibm:params:oauth:grant-type:apikey",
+                "apikey": self.api_key
+            }
+            post_args = {"data": form_data}
+            resp = requests.post(exchange_url, headers=headers, timeout=timeout, verify=self._wo_ssl_verify, **post_args)
+        elif self.auth_mode == "cpd":
+            headers = {"Accept": "application/json", "Content-Type": "application/json"}
+            body = {"username": self.username}
+            if self.password:
+                body["password"] = self.password
+            else:
+                body["api_key"] = self.api_key
+            timeout = self.timeout
+            resp = requests.post(exchange_url, headers=headers, json=body, timeout=timeout, verify=self._wo_ssl_verify)
+        else:
+            headers = {"Accept": "application/json", "Content-Type": "application/json"}
+            post_args = {"json": {"apikey": self.api_key}}
+            resp = requests.post(exchange_url, headers=headers, timeout=timeout, verify=self._wo_ssl_verify, **post_args)
+
         if resp.status_code == 200:
             json_obj = resp.json()
-            token = json_obj[token_key]
-            expires_in = json_obj["expires_in"]
-            refresh_time = time.time() + int(0.8*expires_in)
+            token = json_obj.get("access_token") or json_obj.get("token")
+            if not token:
+                raise RuntimeError(f"No token field found in response: {json_obj!r}")
+
+            expires_in = json_obj.get("expires_in")
+            try:
+                expires_in = int(expires_in) if expires_in is not None else None
+            except Exception:
+                expires_in = None
+            if not expires_in or expires_in <= 0:
+                expires_in = int(os.environ.get("TOKEN_DEFAULT_EXPIRES_IN", 1))
+
+            refresh_time = time.time() + int(0.8 * expires_in)
             return token, refresh_time
 
         resp.raise_for_status()
@@ -80,9 +183,9 @@ class ModelProxyProvider(Provider):
 
         self.refresh_token_if_expires()
         headers = self.get_header()
-        payload = {"inputs": sentences, "model_id": self.embedding_model_id, "space_id": "1"}
+        payload = {"inputs": sentences, "model_id": self.embedding_model_id, "space_id": self.space_id}
                    #"timeout": self.timeout}
-        resp = requests.post(self.embedding_url, json=payload, headers=headers)
+        resp = requests.post(self.embedding_url, json=payload, headers=headers, verify=self._wo_ssl_verify)
 
         if resp.status_code == 200:
             json_obj = resp.json()
@@ -95,9 +198,9 @@ class ModelProxyProvider(Provider):
             raise Exception("model id must be specified for text generation")
         self.refresh_token_if_expires()
         headers = self.get_header()
-        payload = {"input": sentence, "model_id": self.model_id, "space_id": "1",
+        payload = {"input": sentence, "model_id": self.model_id, "space_id": self.space_id,
                    "timeout": self.timeout, "parameters": self.params}
-        resp = requests.post(self.url, json=payload, headers=headers)
+        resp = requests.post(self.url, json=payload, headers=headers, verify=self._wo_ssl_verify)
         if resp.status_code == 200:
             return resp.json()["results"][0]["generated_text"]
 
@@ -106,4 +209,4 @@ class ModelProxyProvider(Provider):
 
 if __name__ == "__main__":
     provider = ModelProxyProvider(model_id="meta-llama/llama-3-3-70b-instruct", embedding_model_id="ibm/slate-30m-english-rtrvr")
-    print(provider.query("ok"))
+    print(provider.query("ok"))

wxo_agentic_evaluation/service_provider/ollama_provider.py

@@ -1,17 +1,16 @@
-import requests
 import json
-from wxo_agentic_evaluation.service_provider.provider import Provider
-from typing import List
 import os
+from typing import List
+
+import requests
+
+from wxo_agentic_evaluation.service_provider.provider import Provider
 
 OLLAMA_URL = os.environ.get("OLLAMA_HOST", "http://localhost:11434")
 
 
 class OllamaProvider(Provider):
-    def __init__(
-        self,
-        model_id=None
-    ):
+    def __init__(self, model_id=None):
         self.url = OLLAMA_URL + "/api/generate"
         self.model_id = model_id
         super().__init__()
@@ -20,14 +19,14 @@ class OllamaProvider(Provider):
         payload = {"model": self.model_id, "prompt": sentence}
         resp = requests.post(self.url, json=payload, stream=True)
         final_text = ""
-        data = b''
+        data = b""
         for chunk in resp:
             data += chunk
-            if data.endswith(b'\n'):
+            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''
+                data = b""
 
         return final_text
 
@@ -37,4 +36,4 @@ class OllamaProvider(Provider):
 
 if __name__ == "__main__":
     provider = OllamaProvider(model_id="llama3.1:8b")
-    print(provider.query("ok"))
+    print(provider.query("ok"))

wxo_agentic_evaluation/service_provider/provider.py

@@ -16,4 +16,3 @@ class Provider(ABC):
     @abstractmethod
     def encode(self, sentences: List[str]) -> List[list]:
         pass
-

wxo_agentic_evaluation/service_provider/referenceless_provider_wrapper.py

@@ -1,11 +1,15 @@
-import requests
-from typing import List, Mapping, Union, Optional, Any
 from abc import ABC, abstractmethod
+from typing import Any, List, Mapping, Optional, Union
 
+import requests
 import rich
 
-from wxo_agentic_evaluation.service_provider.model_proxy_provider import ModelProxyProvider
-from wxo_agentic_evaluation.service_provider.watsonx_provider import WatsonXProvider
+from wxo_agentic_evaluation.service_provider.model_proxy_provider import (
+    ModelProxyProvider,
+)
+from wxo_agentic_evaluation.service_provider.watsonx_provider import (
+    WatsonXProvider,
+)
 
 
 class LLMResponse:
@@ -14,7 +18,9 @@ class LLMResponse:
     Response object that can contain both content and tool calls
     """
 
-    def __init__(self, content: str, tool_calls: Optional[List[Mapping[str, Any]]] = None):
+    def __init__(
+        self, content: str, tool_calls: Optional[List[Mapping[str, Any]]] = None
+    ):
         self.content = content
         self.tool_calls = tool_calls or []
 
@@ -26,25 +32,26 @@ class LLMResponse:
         """Return a string representation of the LLMResponse object."""
         return f"LLMResponse(content='{self.content}', tool_calls={self.tool_calls})"
 
+
 class LLMKitWrapper(ABC):
-    """ In the future this wrapper won't be neccesary.
+    """In the future this wrapper won't be neccesary.
     Right now the referenceless code requires a `generate()` function for the metrics client.
     In refactor, rewrite referenceless code so this wrapper is not needed.
     """
+
     @abstractmethod
     def chat():
         pass
 
     def generate(
-            self,
-            prompt: Union[str, List[Mapping[str, str]]],
-            *,
-            schema,
-            retries: int = 3,
-            generation_args: Optional[Any] = None,
-            **kwargs: Any
-        ):
-
+        self,
+        prompt: Union[str, List[Mapping[str, str]]],
+        *,
+        schema,
+        retries: int = 3,
+        generation_args: Optional[Any] = None,
+        **kwargs: Any,
+    ):
         """
         In future, implement validation of response like in llmevalkit
         """
@@ -55,7 +62,9 @@ class LLMKitWrapper(ABC):
                 response = self._parse_llm_response(raw_response)
                 return response
             except Exception as e:
-                rich.print(f"[b][r] Generation failed with error '{str(e)}' during `quick-eval` ... Attempt ({attempt} / {retries}))")
+                rich.print(
+                    f"[b][r] Generation failed with error '{str(e)}' during `quick-eval` ... Attempt ({attempt} / {retries}))"
+                )
 
     def _parse_llm_response(self, raw: Any) -> Union[str, LLMResponse]:
         """
@@ -82,10 +91,12 @@ class LLMKitWrapper(ABC):
                                 "id": tool_call.get("id"),
                                 "type": tool_call.get("type", "function"),
                                 "function": {
-                                    "name": tool_call.get("function", {}).get("name"),
-                                    "arguments": tool_call.get("function", {}).get(
-                                        "arguments"
+                                    "name": tool_call.get("function", {}).get(
+                                        "name"
                                     ),
+                                    "arguments": tool_call.get(
+                                        "function", {}
+                                    ).get("arguments"),
                                 },
                             }
                             tool_calls.append(tool_call_dict)
@@ -101,6 +112,7 @@ class LLMKitWrapper(ABC):
 
         return content
 
+
 class ModelProxyProviderLLMKitWrapper(ModelProxyProvider, LLMKitWrapper):
     def chat(self, sentence: List[str]):
         if self.model_id is None:
@@ -113,7 +125,7 @@ class ModelProxyProviderLLMKitWrapper(ModelProxyProvider, LLMKitWrapper):
             "messages": sentence,
             "parameters": self.params,
             "space_id": "1",
-            "timeout": self.timeout
+            "timeout": self.timeout,
         }
         resp = requests.post(url=chat_url, headers=headers, json=data)
         if resp.status_code == 200:
@@ -121,6 +133,7 @@ class ModelProxyProviderLLMKitWrapper(ModelProxyProvider, LLMKitWrapper):
         else:
             resp.raise_for_status()
 
+
 class WatsonXLLMKitWrapper(WatsonXProvider, LLMKitWrapper):
     def chat(self, sentence: list):
         chat_url = f"{self.api_endpoint}/ml/v1/text/chat?version=2023-05-02"
@@ -129,7 +142,7 @@ class WatsonXLLMKitWrapper(WatsonXProvider, LLMKitWrapper):
             "model_id": self.model_id,
             "messages": sentence,
             "parameters": self.params,
-            "space_id": self.space_id
+            "space_id": self.space_id,
         }
         resp = requests.post(url=chat_url, headers=headers, json=data)
         if resp.status_code == 200:

wxo_agentic_evaluation/service_provider/watsonx_provider.py

@@ -1,11 +1,13 @@
-import os
-import requests
+import dataclasses
 import json
+import os
+import time
+from threading import Lock
 from types import MappingProxyType
 from typing import List, Mapping, Union
-import dataclasses
-from threading import Lock
-import time
+
+import requests
+
 from wxo_agentic_evaluation.service_provider.provider import Provider
 
 ACCESS_URL = "https://iam.cloud.ibm.com/identity/token"
@@ -36,7 +38,9 @@ class WatsonXProvider(Provider):
         super().__init__()
         self.url = url
         if (embedding_model_id is None) and (model_id is None):
-            raise Exception("either model_id or embedding_model_id must be specified")
+            raise Exception(
+                "either model_id or embedding_model_id must be specified"
+            )
         self.model_id = model_id
         api_key = os.environ.get("WATSONX_APIKEY", api_key)
         if not api_key:
@@ -56,7 +60,7 @@ class WatsonXProvider(Provider):
         self.lock = Lock()
 
         self.params = params if params else DEFAULT_PARAM
-        
+
         if isinstance(self.params, MappingProxyType):
             self.params = dict(self.params)
         if dataclasses.is_dataclass(self.params):
@@ -68,7 +72,10 @@ class WatsonXProvider(Provider):
 
     def _get_access_token(self):
         response = requests.post(
-            self.url, headers=ACCESS_HEADER, data=self.access_data, timeout=self.timeout
+            self.url,
+            headers=ACCESS_HEADER,
+            data=self.access_data,
+            timeout=self.timeout,
         )
         if response.status_code == 200:
             token_data = json.loads(response.text)
@@ -84,16 +91,24 @@ class WatsonXProvider(Provider):
         )
 
     def prepare_header(self):
-        headers = {"Authorization": f"Bearer {self.access_token}",
-                  "Content-Type": "application/json"}
+        headers = {
+            "Authorization": f"Bearer {self.access_token}",
+            "Content-Type": "application/json",
+        }
         return headers
 
     def _query(self, sentence: str):
         headers = self.prepare_header()
 
-        data = {"model_id": self.model_id, "input": sentence,
-                "parameters": self.params, "space_id": self.space_id}
-        generation_url = f"{self.api_endpoint}/ml/v1/text/generation?version=2023-05-02"
+        data = {
+            "model_id": self.model_id,
+            "input": sentence,
+            "parameters": self.params,
+            "space_id": self.space_id,
+        }
+        generation_url = (
+            f"{self.api_endpoint}/ml/v1/text/generation?version=2023-05-02"
+        )
         resp = requests.post(url=generation_url, headers=headers, json=data)
         if resp.status_code == 200:
             return resp.json()["results"][0]
@@ -105,20 +120,25 @@ class WatsonXProvider(Provider):
         if not self.access_token or time.time() > self.refresh_time:
             with self.lock:
                 if not self.access_token or time.time() > self.refresh_time:
-                    self.access_token, self.refresh_time = self._get_access_token()
+                    (
+                        self.access_token,
+                        self.refresh_time,
+                    ) = self._get_access_token()
 
     def query(self, sentence: Union[str, Mapping[str, str]]) -> str:
         if self.model_id is None:
             raise Exception("model id must be specified for text generation")
         try:
             response = self._query(sentence)
-            if (generated_text := response.get("generated_text")):
+            if generated_text := response.get("generated_text"):
                 return generated_text
-            elif (message := response.get("message")):
+            elif message := response.get("message"):
                 return message
             else:
-                raise ValueError(f"Unexpected response from WatsonX: {response}")
-            
+                raise ValueError(
+                    f"Unexpected response from WatsonX: {response}"
+                )
+
         except Exception as e:
             with self.lock:
                 if "authentication_token_expired" in str(e):
@@ -130,12 +150,18 @@ class WatsonXProvider(Provider):
 
     def encode(self, sentences: List[str]) -> List[list]:
         if self.embedding_model_id is None:
-            raise Exception("embedding model id must be specified for text encoding")
+            raise Exception(
+                "embedding model id must be specified for text encoding"
+            )
 
         headers = self.prepare_header()
         url = f"{self.api_endpoint}/ml/v1/text/embeddings?version=2023-10-25"
 
-        data = {"inputs": sentences, "model_id": self.model_id, "space_id": self.space_id}
+        data = {
+            "inputs": sentences,
+            "model_id": self.model_id,
+            "space_id": self.space_id,
+        }
         resp = requests.post(url=url, headers=headers, json=data)
         if resp.status_code == 200:
             return [entry["embedding"] for entry in resp.json()["results"]]
@@ -144,7 +170,9 @@ class WatsonXProvider(Provider):
 
 
 if __name__ == "__main__":
-    provider = WatsonXProvider(model_id="meta-llama/llama-3-2-90b-vision-instruct")
+    provider = WatsonXProvider(
+        model_id="meta-llama/llama-3-2-90b-vision-instruct"
+    )
 
     prompt = """
 <|begin_of_text|><|start_header_id|>system<|end_header_id|>
@@ -176,4 +204,4 @@ Usernwaters did not take anytime off during the period<|eot_id|>
 <|eot_id|><|start_header_id|>user<|end_header_id|>
 """
 
-    print(provider.query(prompt))
+    print(provider.query(prompt))