A client for IBM Watson X AI that provides similar interface to OllamaClient for seamless integration with the RAG system.
| 6 | |
| 7 | |
| 8 | class WatsonXClient: |
| 9 | """ |
| 10 | A client for IBM Watson X AI that provides similar interface to OllamaClient |
| 11 | for seamless integration with the RAG system. |
| 12 | """ |
| 13 | def __init__( |
| 14 | self, |
| 15 | api_key: str, |
| 16 | project_id: str, |
| 17 | url: str = "https://us-south.ml.cloud.ibm.com", |
| 18 | ): |
| 19 | """ |
| 20 | Initialize the Watson X client. |
| 21 | |
| 22 | Args: |
| 23 | api_key: IBM Cloud API key for authentication |
| 24 | project_id: Watson X project ID |
| 25 | url: Watson X service URL (default: us-south region) |
| 26 | """ |
| 27 | self.api_key = api_key |
| 28 | self.project_id = project_id |
| 29 | self.url = url |
| 30 | |
| 31 | try: |
| 32 | from ibm_watsonx_ai import APIClient |
| 33 | from ibm_watsonx_ai import Credentials |
| 34 | from ibm_watsonx_ai.foundation_models import ModelInference |
| 35 | from ibm_watsonx_ai.foundation_models.schema import TextGenParameters |
| 36 | except ImportError: |
| 37 | raise ImportError( |
| 38 | "ibm-watsonx-ai package is required. " |
| 39 | "Install it with: pip install ibm-watsonx-ai" |
| 40 | ) |
| 41 | |
| 42 | self._APIClient = APIClient |
| 43 | self._Credentials = Credentials |
| 44 | self._ModelInference = ModelInference |
| 45 | self._TextGenParameters = TextGenParameters |
| 46 | |
| 47 | self.credentials = self._Credentials( |
| 48 | api_key=self.api_key, |
| 49 | url=self.url |
| 50 | ) |
| 51 | |
| 52 | self.client = self._APIClient(self.credentials) |
| 53 | self.client.set.default_project(self.project_id) |
| 54 | |
| 55 | def _image_to_base64(self, image: Image.Image) -> str: |
| 56 | """Converts a Pillow Image to a base64 string.""" |
| 57 | buffered = BytesIO() |
| 58 | image.save(buffered, format="PNG") |
| 59 | return base64.b64encode(buffered.getvalue()).decode('utf-8') |
| 60 | |
| 61 | def generate_embedding(self, model: str, text: str) -> List[float]: |
| 62 | """ |
| 63 | Generate embeddings using Watson X embedding models. |
| 64 | Note: This requires using Watson X embedding models through the embeddings API. |
| 65 | """ |
no outgoing calls
no test coverage detected