Created
December 31, 2024 09:19
-
-
Save yongyct/ecc0a6c9e53b18450d6f767043987534 to your computer and use it in GitHub Desktop.
Moonshot Custom Connector Example - DataRobotConnector
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import io | |
| import logging | |
| import requests | |
| import traceback | |
| from typing import Any | |
| import datarobot as dr | |
| from datarobot_predict.deployment import predict | |
| import pandas as pd | |
| from moonshot.src.connectors.connector import Connector, perform_retry | |
| from moonshot.src.connectors_endpoints.connector_endpoint_arguments import ( | |
| ConnectorEndpointArguments, | |
| ) | |
| logging.basicConfig(level=logging.DEBUG) | |
| logger = logging.getLogger(__name__) | |
| class DataRobotConnector(Connector): | |
| def __init__(self, ep_arguments: ConnectorEndpointArguments): | |
| # Initialize super class | |
| super().__init__(ep_arguments) | |
| # Set DataRobot | |
| if self.token and self.endpoint: | |
| dr.Client(token=self.token, endpoint=self.endpoint) | |
| self.deployment = dr.Deployment.get(self.params["deployment_id"]) | |
| @Connector.rate_limited | |
| @perform_retry | |
| async def get_response(self, prompt: str) -> str: | |
| """ | |
| Retrieve and return a response. | |
| This method is used to retrieve a response, typically from an object or service represented by | |
| the current instance. | |
| Returns: | |
| str: retrieved response data | |
| """ | |
| connector_prompt = f"{self.pre_prompt}{prompt}{self.post_prompt}" | |
| try: | |
| response = requests.post( | |
| url=f"{self.uri}/deployments/{self.deployment.id}/predictions", | |
| headers={ | |
| "Authorization": f"bearer {self.token}", | |
| "Content-Type": "text/csv", | |
| "Accept": "text/csv", | |
| }, | |
| data=pd.DataFrame([{"promptText": connector_prompt}]).to_csv(index=False), | |
| ) | |
| except Exception as e: | |
| logger.error(traceback.format_exc()) | |
| return await self._process_response(pd.read_csv(io.BytesIO(response.content))) | |
| async def _process_response(self, response: Any) -> str: | |
| """ | |
| Process an HTTP response and extract relevant information as a string. | |
| This function takes an HTTP response object as input and processes it to extract | |
| relevant information as a string. The extracted information may include data | |
| from the response body, headers, or other attributes. | |
| Args: | |
| response (OpenAIObject): An HTTP response object containing the response data. | |
| Returns: | |
| str: A string representing the relevant information extracted from the response. | |
| """ | |
| return response["resultText_PREDICTION"][0] |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment