""" title: Multi-Model Manifold Pipe author: Your Name version: 0.3.8 license: MIT """ import json from pydantic import BaseModel, Field from typing import Union, Iterator, List from google.oauth2 import service_account from google.auth.transport.requests import AuthorizedSession # Set DEBUG to True to enable detailed logging DEBUG = True # Replace this string with your actual service account JSON content SERVICE_ACCOUNT_JSON = r""" { "type": "", "project_id": "", "private_key_id": "", "private_key": "-----BEGIN PRIVATE KEY-----\xxxxx\n-----END PRIVATE KEY-----\n", "client_email": "", "client_id": "", "auth_uri": "https://accounts.google.com/o/oauth2/auth", "token_uri": "https://oauth2.googleapis.com/token", "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", "client_x509_cert_url": "", "universe_domain": "googleapis.com" } """ class Pipe: class Valves(BaseModel): PROJECT_ID: str = Field( default="", # Set this to your Project ID description="Your Google Cloud Project ID.", ) LOCATION: str = Field( default="us-central1", description="Default region for Vertex AI models." ) def __init__(self): self.id = ( "vertex_ai" # Set this to your function's ID in OpenWebUI ) self.type = "manifold" self.name = "Vertex AI Models" self.valves = self.Valves() # Load credentials from the hardcoded JSON credentials_info = json.loads(SERVICE_ACCOUNT_JSON) self.credentials = service_account.Credentials.from_service_account_info( credentials_info, scopes=["https://www.googleapis.com/auth/cloud-platform"] ) self.authed_session = AuthorizedSession(self.credentials) def pipes(self) -> List[dict]: # Return a list of models return [ {"id": "gemini-1.5-pro-002", "name": "Gemini 1.5 Pro"}, {"id": "claude-3-5-sonnet@20240620", "name": "Claude 3.5 Sonnet V1"}, {"id": "claude-3-haiku@20240307", "name": "Claude 3 Haiku"}, ] def pipe(self, body: dict) -> Union[str, Iterator[str]]: project_id = self.valves.PROJECT_ID default_location = self.valves.LOCATION model_id = body.get("model") if not model_id: return "Error: No model specified." if not model_id: return "Error: No model specified." if DEBUG: print(f"Original model_id: {model_id}") # Remove the function ID prefix if present function_id_with_dot = f"{self.id}." if model_id.startswith(function_id_with_dot): model_id = model_id[len(function_id_with_dot) :] if DEBUG: print( f"Removed prefix '{function_id_with_dot}' from model_id, new model_id: {model_id}" ) else: if DEBUG: print("No prefix to remove from model_id.") if DEBUG: print(f"Using model_id: {model_id}") try: # Determine which model is selected and call the appropriate method if model_id.startswith("gemini"): return self.call_gemini_model( body, project_id, default_location, model_id ) elif model_id.startswith("claude"): # Use 'us-east5' for Claude models return self.call_claude_model(body, project_id, "us-east5", model_id) elif model_id.startswith("llama"): return self.call_llama_model( body, project_id, default_location, model_id ) else: return f"Error: Unknown model {model_id}" except Exception as e: if DEBUG: print(f"Error in pipe method: {e}") return f"Error: {e}" def call_gemini_model(self, body, project_id, location, model_id): # Use the working endpoint for Gemini models url = f"https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project_id}/locations/{location}/endpoints/openapi/chat/completions" # Prepare the prompt messages messages = body["messages"] formatted_messages = self.format_messages(messages) # Prepare request data data = { "model": f"google/{model_id}", "messages": formatted_messages, "temperature": body.get("temperature", 0.7), "max_tokens": body.get("max_tokens", 512), "top_p": body.get("top_p", 0.8), "top_k": body.get("top_k", 40), # Include other parameters if needed } if DEBUG: print("Request URL:", url) print("Request Data:", json.dumps(data, indent=2)) # Make the POST request response = self.authed_session.post(url, json=data) # Check the response if response.status_code == 200: response_data = response.json() if DEBUG: print("Response Data:", json.dumps(response_data, indent=2)) if "choices" in response_data and len(response_data["choices"]) > 0: # Extract the generated text generated_text = response_data["choices"][0]["message"]["content"] return generated_text else: if DEBUG: print("No choices found in the response.") print("Response Data:", json.dumps(response_data, indent=2)) return "No response generated by the model." else: if DEBUG: print(f"Error: {response.status_code}") print("Response Headers:", response.headers) print("Response Content:", response.text) return f"Error: {response.status_code} {response.text}" def call_claude_model(self, body, project_id, location, model_id): # Endpoint URL for Claude models url = f"https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/publishers/anthropic/models/{model_id}:predict" # Prepare the prompt messages messages = body["messages"] formatted_messages = self.format_claude_messages(messages) # Prepare request data data = { "instances": [ { "messages": formatted_messages, "anthropic_version": "vertex-2023-10-16", } ], "parameters": { "temperature": body.get("temperature", 0.7), "max_output_tokens": body.get("max_tokens", 512), "top_p": body.get("top_p", 0.8), "top_k": body.get("top_k", 40), }, } if DEBUG: print("Request URL:", url) print("Request Data:", json.dumps(data, indent=2)) # Make the POST request response = self.authed_session.post(url, json=data) # Check the response if response.status_code == 200: response_data = response.json() if DEBUG: print("Response Data:", json.dumps(response_data, indent=2)) # Extract the generated text if "predictions" in response_data and len(response_data["predictions"]) > 0: generated_text = response_data["predictions"][0].get("content", "") return generated_text else: if DEBUG: print("No predictions found in the response.") print("Response Data:", json.dumps(response_data, indent=2)) return "No response generated by the model." else: if DEBUG: print(f"Error: {response.status_code}") print("Response Headers:", response.headers) print("Response Content:", response.text) return f"Error: {response.status_code} {response.text}" def call_llama_model(self, body, project_id, location, model_id): # Placeholder for Llama model implementation return "Llama model is not implemented yet." def format_messages(self, messages): formatted_messages = [] for msg in messages: role = msg.get("role") content = msg.get("content") if role == "user": formatted_messages.append({"role": "user", "content": content}) elif role == "assistant": formatted_messages.append({"role": "assistant", "content": content}) elif role == "system": formatted_messages.append({"role": "system", "content": content}) return formatted_messages def format_claude_messages(self, messages): # For Claude models, messages might need to be formatted differently formatted_messages = [] for msg in messages: role = msg.get("role") content = msg.get("content") formatted_messages.append({"role": role, "content": content}) return formatted_messages