wehub-resource-sync
992 lines
40 KiB
Python
992 lines
40 KiB
Python
import json
|
|
import os
|
|
import logging
|
|
import time
|
|
import re
|
|
import urllib.parse
|
|
try:
|
|
from dotenv import load_dotenv
|
|
env_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env")
|
|
load_dotenv(env_path, override=True)
|
|
except ImportError:
|
|
pass
|
|
import google.auth
|
|
from google.auth.transport.requests import AuthorizedSession
|
|
|
|
def get_endpoint_host(location: str) -> str:
|
|
"""Returns the host domain based on the geographic location."""
|
|
if location == "global":
|
|
return "discoveryengine.googleapis.com"
|
|
return f"{location}-discoveryengine.googleapis.com"
|
|
|
|
def get_session():
|
|
"""Returns an authorized Google API requests session."""
|
|
credentials, project = google.auth.default(
|
|
scopes=["https://www.googleapis.com/auth/cloud-platform"]
|
|
)
|
|
return AuthorizedSession(credentials)
|
|
|
|
def list_notebooks(project_number: str, location: str = "global") -> list:
|
|
"""Lists recently viewed notebooks in the specified project."""
|
|
session = get_session()
|
|
host = get_endpoint_host(location)
|
|
url = f"https://{host}/v1alpha/projects/{project_number}/locations/{location}/notebooks:listRecentlyViewed"
|
|
|
|
resp = session.get(url)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
return data.get("notebooks", [])
|
|
|
|
def list_sources_and_types(notebook_id: str, project_number: str, location: str = "global") -> list:
|
|
"""Lists sources and maps their types for a given notebook."""
|
|
session = get_session()
|
|
host = get_endpoint_host(location)
|
|
base_url = f"https://{host}/v1alpha/projects/{project_number}/locations/{location}/notebooks/{notebook_id}"
|
|
|
|
get_resp = session.get(base_url)
|
|
get_resp.raise_for_status()
|
|
nb_data = get_resp.json()
|
|
sources = nb_data.get("sources", [])
|
|
|
|
results = []
|
|
for src in sources:
|
|
src_id = src.get("sourceId", {}).get("id")
|
|
src_title = src.get("title")
|
|
|
|
src_url = f"{base_url}/sources/{src_id}"
|
|
src_resp = session.get(src_url)
|
|
src_resp.raise_for_status()
|
|
src_data = src_resp.json()
|
|
|
|
metadata = src_data.get("metadata", {})
|
|
|
|
source_type = "copied text"
|
|
source_location = "N/A"
|
|
|
|
if "webpageMetadata" in metadata:
|
|
source_type = "website"
|
|
source_location = metadata["webpageMetadata"].get("webpageUrl")
|
|
elif "googleDocsMetadata" in metadata:
|
|
source_type = "google docs"
|
|
doc_id = metadata["googleDocsMetadata"].get("documentId")
|
|
source_location = f"https://docs.google.com/document/d/{doc_id}/edit"
|
|
|
|
results.append({
|
|
"title": src_title,
|
|
"id": src_id,
|
|
"type": source_type,
|
|
"location": source_location,
|
|
"raw_data": src_data
|
|
})
|
|
|
|
return results
|
|
|
|
def create_notebook(target_project_number: str, target_location: str, title: str) -> dict:
|
|
"""Creates a new empty notebook in the target project."""
|
|
logging.info(f"DEBUG: create_notebook called with title='{title}', project='{target_project_number}', location='{target_location}'")
|
|
session = get_session()
|
|
host = get_endpoint_host(target_location)
|
|
url = f"https://{host}/v1alpha/projects/{target_project_number}/locations/{target_location}/notebooks"
|
|
|
|
resp = session.post(url, json={"title": title})
|
|
resp.raise_for_status()
|
|
return resp.json()
|
|
|
|
def add_source_to_notebook(target_project_number: str, target_location: str, notebook_id: str, source_content: dict) -> dict:
|
|
"""Adds a single source (userContent payload) to the specified notebook."""
|
|
logging.info(f"DEBUG: add_source_to_notebook called with notebook_id='{notebook_id}', project='{target_project_number}'")
|
|
session = get_session()
|
|
host = get_endpoint_host(target_location)
|
|
url = f"https://{host}/v1alpha/projects/{target_project_number}/locations/{target_location}/notebooks/{notebook_id}/sources:batchCreate"
|
|
|
|
# Make a copy to avoid mutating inputs
|
|
content_obj = json.loads(json.dumps(source_content))
|
|
if "sourceName" in content_obj:
|
|
logging.info(f"DEBUG: Removing invalid 'sourceName' from payload: {content_obj['sourceName']}")
|
|
content_obj.pop("sourceName")
|
|
|
|
logging.info(f"DEBUG: Sending request to {url} with payload: {json.dumps(content_obj)}")
|
|
resp = session.post(url, json={"userContents": [content_obj]}, timeout=60)
|
|
logging.info(f"DEBUG: Response status: {resp.status_code}")
|
|
resp.raise_for_status()
|
|
return resp.json()
|
|
|
|
def extract_notebook_source_payload(src_raw_data: dict) -> dict:
|
|
"""Intelligently extracts the userContent payload from a raw source data dict, reconstructs it if needed, and removes read-only fields."""
|
|
user_content = src_raw_data.get("userContent")
|
|
if user_content:
|
|
payload = json.loads(json.dumps(user_content))
|
|
else:
|
|
payload = {}
|
|
metadata = src_raw_data.get("metadata", {})
|
|
if "webpageMetadata" in metadata:
|
|
payload["webContent"] = {
|
|
"url": metadata["webpageMetadata"].get("webpageUrl")
|
|
}
|
|
elif "googleDocsMetadata" in metadata:
|
|
payload["googleDriveContent"] = {
|
|
"documentId": metadata["googleDocsMetadata"].get("documentId"),
|
|
"mimeType": "application/vnd.google-apps.document"
|
|
}
|
|
elif "textContent" in src_raw_data:
|
|
payload["textContent"] = src_raw_data["textContent"]
|
|
elif "textContent" in metadata:
|
|
payload["textContent"] = metadata["textContent"]
|
|
|
|
if "sourceName" in payload:
|
|
payload.pop("sourceName")
|
|
return payload
|
|
|
|
def migrate_notebook_pipeline(
|
|
notebook_id_or_title: str,
|
|
target_project_number: str,
|
|
target_location: str,
|
|
source_project_number: str,
|
|
source_location: str = "global",
|
|
backup_bucket: str = ""
|
|
) -> dict:
|
|
"""Migrates an entire notebook and all its sources deterministically in a python loop."""
|
|
logging.info(f"Starting deterministic notebook migration for '{notebook_id_or_title}'")
|
|
|
|
# 1. List source notebooks
|
|
notebooks = list_notebooks(source_project_number, source_location)
|
|
source_notebook = None
|
|
for nb in notebooks:
|
|
nb_id = nb.get("name", "").split("/")[-1]
|
|
nb_title = nb.get("title", "")
|
|
if nb_id == notebook_id_or_title or nb_title == notebook_id_or_title:
|
|
source_notebook = nb
|
|
break
|
|
|
|
if not source_notebook:
|
|
raise ValueError(f"Source notebook '{notebook_id_or_title}' not found.")
|
|
|
|
source_nb_id = source_notebook.get("name", "").split("/")[-1]
|
|
source_nb_title = source_notebook.get("title", "Untitled Notebook")
|
|
|
|
# 2. Get all sources
|
|
sources = list_sources_and_types(source_nb_id, source_project_number, source_location)
|
|
logging.info(f"Found {len(sources)} sources to migrate for notebook '{source_nb_title}'")
|
|
|
|
# 2.5 Upload to GCS (Backup) implicitly
|
|
bucket_name = backup_bucket if backup_bucket else os.environ.get("GCS_BUCKET_NAME")
|
|
if bucket_name:
|
|
try:
|
|
timestamp = int(time.time())
|
|
object_name = f"exports/{source_nb_title}_{timestamp}.json"
|
|
exported_sources = []
|
|
for src in sources:
|
|
payload = extract_notebook_source_payload(src.get("raw_data", {}))
|
|
exported_sources.append({
|
|
"title": src["title"],
|
|
"userContent": payload
|
|
})
|
|
notebook_def = {
|
|
"title": source_nb_title,
|
|
"sources": exported_sources
|
|
}
|
|
session = get_session()
|
|
encoded_obj = urllib.parse.quote(object_name, safe="")
|
|
upload_url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_obj}"
|
|
up_resp = session.post(
|
|
upload_url,
|
|
data=json.dumps(notebook_def, indent=2),
|
|
headers={"Content-Type": "application/json"}
|
|
)
|
|
up_resp.raise_for_status()
|
|
logging.info(f"Implicitly backed up notebook '{source_nb_title}' to gs://{bucket_name}/{object_name}")
|
|
except Exception as e:
|
|
logging.error(f"Failed to implicitly backup notebook to GCS: {e}")
|
|
|
|
# 3. Create target notebook
|
|
target_nb = create_notebook(target_project_number, target_location, source_nb_title)
|
|
target_nb_id = target_nb.get("name", "").split("/")[-1]
|
|
logging.info(f"Created target notebook '{source_nb_title}' with ID: {target_nb_id}")
|
|
|
|
# 4. For each source, map and add
|
|
migrated_sources = []
|
|
failed_sources = []
|
|
|
|
for src in sources:
|
|
title = src.get("title", "Untitled Source")
|
|
raw_data = src.get("raw_data", {})
|
|
payload = extract_notebook_source_payload(raw_data)
|
|
|
|
logging.info(f"Adding source '{title}' to target notebook '{target_nb_id}'")
|
|
try:
|
|
add_resp = add_source_to_notebook(target_project_number, target_location, target_nb_id, payload)
|
|
migrated_sources.append({
|
|
"title": title,
|
|
"status": "success",
|
|
"response": add_resp
|
|
})
|
|
except Exception as e:
|
|
logging.error(f"Failed to add source '{title}': {e}")
|
|
failed_sources.append({
|
|
"title": title,
|
|
"status": "failed",
|
|
"error": str(e)
|
|
})
|
|
|
|
return {
|
|
"success": len(failed_sources) == 0,
|
|
"source_notebook_id": source_nb_id,
|
|
"source_notebook_title": source_nb_title,
|
|
"target_notebook_id": target_nb_id,
|
|
"migrated_sources_count": len(migrated_sources),
|
|
"failed_sources_count": len(failed_sources),
|
|
"migrated": migrated_sources,
|
|
"failed": failed_sources
|
|
}
|
|
|
|
def list_employee_agents(
|
|
project_id: str,
|
|
location: str = "global",
|
|
engine_id: str = "enterprise-search-17416389_1741638989378",
|
|
basic: bool = False
|
|
) -> list:
|
|
"""Lists employee-made low-code agents."""
|
|
session = get_session()
|
|
host = get_endpoint_host(location)
|
|
base_url = f"https://{host}/v1alpha"
|
|
parent = f"projects/{project_id}/locations/{location}/collections/default_collection/engines/{engine_id}/assistants/default_assistant"
|
|
url = f"{base_url}/{parent}/agents"
|
|
|
|
logging.info(f"DEBUG: list_employee_agents called for {parent}")
|
|
resp = session.get(url)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
agents = data.get("agents", [])
|
|
|
|
employee_agents = []
|
|
for agent in agents:
|
|
displayName = (agent.get("displayName") or "").replace("\r\n", " ").replace("\n", " ")
|
|
name = agent.get("name")
|
|
description = (agent.get("description") or "").replace("\r\n", " ").replace("\n", " ")
|
|
|
|
root_instructions = "No instructions found."
|
|
root_tools = []
|
|
root_knowledge = []
|
|
sub_agents = []
|
|
|
|
if "lowCodeAgentDefinition" in agent:
|
|
definition = agent.get("lowCodeAgentDefinition", {})
|
|
nodes = definition.get("nodes", [])
|
|
root_id = definition.get("rootAgentId")
|
|
|
|
for node in nodes:
|
|
node_id = node.get("id")
|
|
llm_node = node.get("llmAgentNode", {})
|
|
node_instruction = llm_node.get("instruction", "No instructions found.")
|
|
|
|
# Extract tools for this node
|
|
node_tools = []
|
|
for t in llm_node.get("selectedTools", {}).get("tool", []):
|
|
node_tools.append(t.get("name", "Unknown Tool"))
|
|
for spec in llm_node.get("dataStoreSpecs", {}).get("specs", []):
|
|
ds = spec.get("dataStore", "")
|
|
if ds:
|
|
ds_name = ds.split("/")[-1]
|
|
node_tools.append(f"DataStore: {ds_name}")
|
|
node_knowledge = []
|
|
for k_field in ["groundingSources", "userContents", "files", "knowledge"]:
|
|
for item in llm_node.get(k_field, []):
|
|
title = item.get("displayName") or item.get("googleDriveContent", {}).get("documentId")
|
|
if title and title not in node_knowledge:
|
|
node_knowledge.append(title)
|
|
for match in re.findall(r'\[([^\]]+)\]\(https://drive\.google\.com/[^\)]+\)', node_instruction):
|
|
if match and match not in node_knowledge:
|
|
node_knowledge.append(match)
|
|
|
|
if node_id == root_id:
|
|
root_instructions = node_instruction
|
|
for nt in node_tools:
|
|
if nt not in root_tools:
|
|
root_tools.append(nt)
|
|
for nk in node_knowledge:
|
|
if nk not in root_knowledge:
|
|
root_knowledge.append(nk)
|
|
else:
|
|
sub_agents.append({
|
|
"displayName": node.get("displayName", "Sub-Agent"),
|
|
"description": llm_node.get("description", ""),
|
|
"model": llm_node.get("model", "Unknown Model"),
|
|
"instructions": node_instruction,
|
|
"tools": node_tools,
|
|
"knowledge": node_knowledge
|
|
})
|
|
elif "skillAgentDefinition" in agent:
|
|
definition = agent.get("skillAgentDefinition", {})
|
|
root_instructions = definition.get("instruction", "No instructions found.")
|
|
|
|
if basic:
|
|
employee_agents.append({
|
|
"displayName": displayName,
|
|
"name": name,
|
|
"description": description,
|
|
"connectors_and_tools": root_tools,
|
|
})
|
|
else:
|
|
employee_agents.append({
|
|
"displayName": displayName,
|
|
"name": name,
|
|
"description": description,
|
|
"instructions": root_instructions,
|
|
"connectors_and_tools": root_tools,
|
|
"knowledge": root_knowledge,
|
|
"sub_agents": sub_agents,
|
|
"raw_agent": agent
|
|
})
|
|
return employee_agents
|
|
|
|
def extract_connectors(agent_def: dict) -> list:
|
|
"""Extracts unique connector and datastore names from an agent definition."""
|
|
connectors = set()
|
|
nodes = agent_def.get("nodes", [])
|
|
for node in nodes:
|
|
llm_node = node.get("llmAgentNode", {})
|
|
|
|
selected_tools = llm_node.get("selectedTools", {})
|
|
tools_list = selected_tools.get("tool", [])
|
|
for t in tools_list:
|
|
name = t.get("name")
|
|
if name:
|
|
connectors.add(name)
|
|
|
|
for spec in llm_node.get("dataStoreSpecs", {}).get("specs", []):
|
|
ds = spec.get("dataStore", "")
|
|
if ds:
|
|
ds_name = ds.split("/")[-1]
|
|
connectors.add(ds_name)
|
|
|
|
return list(connectors)
|
|
|
|
def list_target_connectors(session, base_url, target_parent) -> list:
|
|
"""Lists connectors used by agents in the target environment."""
|
|
target_url = f"{base_url}/{target_parent}/agents"
|
|
try:
|
|
logging.info(f"Fetching target agents from {target_url}")
|
|
resp = session.get(target_url)
|
|
resp.raise_for_status()
|
|
agents = resp.json().get("agents", [])
|
|
|
|
target_connectors = set()
|
|
for agent in agents:
|
|
if "lowCodeAgentDefinition" in agent:
|
|
definition = agent.get("lowCodeAgentDefinition", {})
|
|
connectors = extract_connectors(definition)
|
|
for c in connectors:
|
|
target_connectors.add(c)
|
|
return list(target_connectors)
|
|
except Exception as e:
|
|
logging.error(f"Failed to list target connectors: {e}")
|
|
return []
|
|
|
|
def extract_agent_datastores(
|
|
source_agent_name: str,
|
|
source_project_id: str,
|
|
source_location: str = "global",
|
|
source_engine_id: str = "enterprise-search-17416389_1741638989378"
|
|
) -> dict:
|
|
"""Extracts and lists the names of datastores used by an agent and its subagents."""
|
|
session = get_session()
|
|
host = get_endpoint_host(source_location)
|
|
base_url = f"https://{host}/v1alpha"
|
|
source_parent = f"projects/{source_project_id}/locations/{source_location}/collections/default_collection/engines/{source_engine_id}/assistants/default_assistant"
|
|
source_url = f"{base_url}/{source_parent}/agents"
|
|
|
|
resp = session.get(source_url)
|
|
resp.raise_for_status()
|
|
agents = resp.json().get("agents", [])
|
|
|
|
agent_to_check = None
|
|
for agent in agents:
|
|
if agent.get("displayName") == source_agent_name and "lowCodeAgentDefinition" in agent:
|
|
agent_to_check = agent
|
|
break
|
|
|
|
if not agent_to_check:
|
|
raise ValueError(f"Source agent '{source_agent_name}' not found or is not a low-code agent.")
|
|
|
|
definition = agent_to_check.get("lowCodeAgentDefinition", {})
|
|
nodes = definition.get("nodes", [])
|
|
|
|
datastore_report = {}
|
|
for node in nodes:
|
|
node_name = node.get("displayName", "Unknown Node")
|
|
llm_node = node.get("llmAgentNode", {})
|
|
datastores = []
|
|
for spec in llm_node.get("dataStoreSpecs", {}).get("specs", []):
|
|
ds = spec.get("dataStore", "")
|
|
if ds:
|
|
ds_name = ds.split("/")[-1]
|
|
datastores.append(ds_name)
|
|
if datastores:
|
|
datastore_report[node_name] = datastores
|
|
|
|
return {
|
|
"agent_name": source_agent_name,
|
|
"datastores": datastore_report
|
|
}
|
|
|
|
def clean_connector_name(name: str) -> str:
|
|
"""Strips common suffixes like May29, dates, punctuation for robust fuzzy matching."""
|
|
clean = re.sub(r'(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\d*', '', name, flags=re.IGNORECASE)
|
|
return re.sub(r'[^a-zA-Z0-9]', '', clean).lower()
|
|
|
|
def lookup_and_map_connectors(
|
|
source_agent_name: str,
|
|
target_project_id: str,
|
|
target_location: str,
|
|
target_engine_id: str,
|
|
source_project_id: str,
|
|
source_location: str = "global",
|
|
source_engine_id: str = "enterprise-search-17416389_1741638989378"
|
|
) -> dict:
|
|
"""Intelligently matches source agent connectors against target environment connectors."""
|
|
session = get_session()
|
|
source_host = get_endpoint_host(source_location)
|
|
source_base_url = f"https://{source_host}/v1alpha"
|
|
source_parent = f"projects/{source_project_id}/locations/{source_location}/collections/default_collection/engines/{source_engine_id}/assistants/default_assistant"
|
|
source_url = f"{source_base_url}/{source_parent}/agents"
|
|
|
|
resp = session.get(source_url)
|
|
if resp.status_code == 200:
|
|
agents = resp.json().get("agents", [])
|
|
else:
|
|
agents = []
|
|
|
|
src_connectors = set()
|
|
for ag in agents:
|
|
if source_agent_name.upper() == "ALL" or ag.get("displayName", "").lower() == source_agent_name.lower():
|
|
for c in extract_connectors(ag.get("lowCodeAgentDefinition", {})):
|
|
src_connectors.add(c)
|
|
src_connectors = list(src_connectors)
|
|
|
|
target_host = get_endpoint_host(target_location)
|
|
target_base_url = f"https://{target_host}/v1alpha"
|
|
target_parent = f"projects/{target_project_id}/locations/{target_location}/collections/default_collection/engines/{target_engine_id}/assistants/default_assistant"
|
|
target_connectors = list_target_connectors(session, target_base_url, target_parent)
|
|
try:
|
|
source_ds_objs = list_datastores(source_project_id, source_location)
|
|
except Exception:
|
|
source_ds_objs = []
|
|
try:
|
|
target_ds_objs = list_datastores(target_project_id, target_location)
|
|
except Exception:
|
|
target_ds_objs = []
|
|
|
|
canonical_tools = ["googleSearch", "urlContext", "geGmail", "snowflakeMcp"]
|
|
verified_target_ids = [ds["id"] for ds in target_ds_objs]
|
|
all_targets = canonical_tools + verified_target_ids
|
|
|
|
mapping = {}
|
|
missing = []
|
|
for sc in src_connectors:
|
|
sc_clean = clean_connector_name(sc)
|
|
matched = None
|
|
for s_ds in source_ds_objs:
|
|
if s_ds["id"] == sc or s_ds["displayName"].lower() == sc.lower():
|
|
for t_ds in target_ds_objs:
|
|
if t_ds["displayName"].lower() == s_ds["displayName"].lower():
|
|
if s_ds["displayName"].lower() == "mcp_data":
|
|
s_pref = s_ds["id"].split("_")[0].split("-")[0].lower()
|
|
t_pref = t_ds["id"].split("_")[0].split("-")[0].lower()
|
|
if s_pref != t_pref:
|
|
continue
|
|
matched = t_ds["id"]
|
|
break
|
|
break
|
|
|
|
if not matched:
|
|
for tc in all_targets:
|
|
if clean_connector_name(tc) == sc_clean:
|
|
matched = tc
|
|
break
|
|
|
|
if matched:
|
|
mapping[sc] = matched
|
|
else:
|
|
mapping[sc] = "❌ Missing (Not Ingested in Target)"
|
|
missing.append(sc)
|
|
|
|
return {
|
|
"source_agent": source_agent_name,
|
|
"source_connectors": src_connectors,
|
|
"target_available": all_targets,
|
|
"proposed_mapping": mapping,
|
|
"missing_connectors": missing
|
|
}
|
|
|
|
def migrate_employee_agent(
|
|
source_agent_name: str,
|
|
target_project_id: str,
|
|
target_location: str,
|
|
target_engine_id: str,
|
|
source_project_id: str,
|
|
source_location: str = "global",
|
|
source_engine_id: str = "enterprise-search-17416389_1741638989378",
|
|
force: bool = False,
|
|
backup_bucket: str = "",
|
|
connector_mapping: str = ""
|
|
) -> dict:
|
|
"""Migrates an employee-made low-code agent to a target environment."""
|
|
session = get_session()
|
|
|
|
target_host = get_endpoint_host(target_location)
|
|
target_base_url = f"https://{target_host}/v1alpha"
|
|
target_parent = f"projects/{target_project_id}/locations/{target_location}/collections/default_collection/engines/{target_engine_id}/assistants/default_assistant"
|
|
target_url = f"{target_base_url}/{target_parent}/agents"
|
|
|
|
# 1. Fetch the source agent definition
|
|
source_host = get_endpoint_host(source_location)
|
|
source_base_url = f"https://{source_host}/v1alpha"
|
|
source_parent = f"projects/{source_project_id}/locations/{source_location}/collections/default_collection/engines/{source_engine_id}/assistants/default_assistant"
|
|
source_url = f"{source_base_url}/{source_parent}/agents"
|
|
|
|
logging.info(f"Fetching source agent from {source_url}")
|
|
resp = session.get(source_url)
|
|
resp.raise_for_status()
|
|
agents = resp.json().get("agents", [])
|
|
|
|
agent_to_migrate = None
|
|
for agent in agents:
|
|
if agent.get("displayName") == source_agent_name and "lowCodeAgentDefinition" in agent:
|
|
agent_to_migrate = agent
|
|
break
|
|
|
|
if not agent_to_migrate:
|
|
raise ValueError(f"Source agent '{source_agent_name}' not found or is not a low-code agent.")
|
|
|
|
# 1.4 Validation and Reporting
|
|
src_connectors = extract_connectors(agent_to_migrate.get("lowCodeAgentDefinition", {}))
|
|
target_connectors = list_target_connectors(session, target_base_url, target_parent)
|
|
|
|
missing_connectors = [c for c in src_connectors if c not in target_connectors]
|
|
|
|
report = {
|
|
"source_connectors": src_connectors,
|
|
"target_connectors_used": target_connectors,
|
|
"missing_connectors": missing_connectors
|
|
}
|
|
|
|
logging.info(f"Connector Validation Report: {json.dumps(report)}")
|
|
|
|
if missing_connectors and not force:
|
|
return {
|
|
"success": False,
|
|
"warning": "Missing dependencies in target environment.",
|
|
"report": report,
|
|
"message": f"The following dependencies are missing in the target environment: {missing_connectors}."
|
|
}
|
|
|
|
# 1.5 Upload to GCS (Backup) implicitly
|
|
bucket_name = backup_bucket if backup_bucket else os.environ.get("GCS_BUCKET_NAME")
|
|
if bucket_name:
|
|
try:
|
|
timestamp = int(time.time())
|
|
object_name = f"exports/{source_agent_name}_{timestamp}.json"
|
|
encoded_obj = urllib.parse.quote(object_name, safe="")
|
|
upload_url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_obj}"
|
|
up_resp = session.post(
|
|
upload_url,
|
|
data=json.dumps(agent_to_migrate, indent=2),
|
|
headers={"Content-Type": "application/json"}
|
|
)
|
|
up_resp.raise_for_status()
|
|
logging.info(f"Implicitly backed up agent '{source_agent_name}' to gs://{bucket_name}/{object_name}")
|
|
except Exception as e:
|
|
logging.error(f"Failed to implicitly backup agent to GCS: {e}")
|
|
|
|
# 2. Create the payload for the target environment
|
|
definition = agent_to_migrate.get("lowCodeAgentDefinition", {})
|
|
if "session" in definition:
|
|
del definition["session"]
|
|
|
|
mapping_dict = {}
|
|
if connector_mapping:
|
|
try:
|
|
mapping_dict = json.loads(connector_mapping)
|
|
except Exception:
|
|
for pair in connector_mapping.split(","):
|
|
if ":" in pair:
|
|
k, v = pair.split(":", 1)
|
|
mapping_dict[k.strip()] = v.strip()
|
|
|
|
# Apply tool mapping and ensure googleSearch
|
|
all_nodes = definition.get("nodes", []) + definition.get("deployedNodes", [])
|
|
for node in all_nodes:
|
|
llm_node = node.get("llmAgentNode", {})
|
|
if "selectedTools" not in llm_node:
|
|
llm_node["selectedTools"] = {"tool": []}
|
|
selected_tools = llm_node["selectedTools"]
|
|
tools_list = selected_tools.get("tool", [])
|
|
|
|
new_tools = []
|
|
has_google_search = False
|
|
for t in tools_list:
|
|
t_name = t.get("name", "")
|
|
if t_name in mapping_dict:
|
|
t_name = mapping_dict[t_name]
|
|
elif clean_connector_name(t_name) == clean_connector_name("snowflakeMcp"):
|
|
t_name = mapping_dict.get("Snowflake Mcp May29", "snowflakeMcp")
|
|
|
|
if t_name == "googleSearch":
|
|
has_google_search = True
|
|
if t_name:
|
|
new_tools.append({"name": t_name})
|
|
|
|
if not has_google_search:
|
|
new_tools.append({"name": "googleSearch"})
|
|
selected_tools["tool"] = new_tools
|
|
|
|
new_specs = []
|
|
for spec in llm_node.get("dataStoreSpecs", {}).get("specs", []):
|
|
ds = spec.get("dataStore", "")
|
|
if ds:
|
|
ds_id = ds.split("/")[-1]
|
|
if ds_id in mapping_dict:
|
|
ds_id = mapping_dict[ds_id]
|
|
new_ds = f"projects/{target_project_id}/locations/{target_location}/collections/default_collection/dataStores/{ds_id}"
|
|
new_specs.append({"dataStore": new_ds})
|
|
if new_specs:
|
|
llm_node["dataStoreSpecs"] = {"specs": new_specs}
|
|
if any("snowflake" in s.get("dataStore", "").lower() for s in new_specs):
|
|
sf_conn = mapping_dict.get("Snowflake Mcp May29", "custom_mcp")
|
|
if not any(t.get("name") == sf_conn for t in new_tools):
|
|
new_tools.append({"name": sf_conn})
|
|
selected_tools["tool"] = new_tools
|
|
if any("drive" in s.get("dataStore", "").lower() for s in new_specs):
|
|
dr_conn = mapping_dict.get("ge-drive-all", "Drive")
|
|
if not any(t.get("name") == dr_conn for t in new_tools):
|
|
new_tools.append({"name": dr_conn})
|
|
selected_tools["tool"] = new_tools
|
|
|
|
payload_str = json.dumps({
|
|
"displayName": agent_to_migrate.get("displayName"),
|
|
"description": agent_to_migrate.get("description", ""),
|
|
"lowCodeAgentDefinition": definition
|
|
})
|
|
payload_str = payload_str.replace(f"projects/{source_project_id}", f"projects/{target_project_id}")
|
|
payload_str = payload_str.replace(source_engine_id, target_engine_id)
|
|
payload = json.loads(payload_str)
|
|
|
|
logging.info(f"Creating new agent at target {target_url}")
|
|
create_resp = session.post(target_url, json=payload)
|
|
create_resp.raise_for_status()
|
|
|
|
created_agent = create_resp.json()
|
|
tgt_extracted = extract_connectors(created_agent.get("lowCodeAgentDefinition", {}))
|
|
message = f"Successfully migrated agent '{source_agent_name}' to target environment.\nConnectors in Source Agent: {src_connectors}\nConnectors in Target Agent: {tgt_extracted}"
|
|
truly_missing = [c for c in missing_connectors if c not in mapping_dict]
|
|
if truly_missing:
|
|
message += f"\nWARNING: Missing connectors were ignored: {truly_missing}"
|
|
|
|
return {
|
|
"success": True,
|
|
"message": message,
|
|
"target_agent": created_agent
|
|
}
|
|
|
|
|
|
|
|
def export_agent_to_gcs(
|
|
source_agent_name: str,
|
|
object_name: str,
|
|
bucket_name: str = "",
|
|
source_project_id: str = "",
|
|
source_location: str = "global",
|
|
source_engine_id: str = "enterprise-search-17416389_1741638989378"
|
|
) -> dict:
|
|
"""Exports an employee-made low-code agent definition to a GCS bucket."""
|
|
if not bucket_name:
|
|
bucket_name = os.environ.get("GCS_BUCKET_NAME")
|
|
if not bucket_name:
|
|
raise ValueError("bucket_name not provided and GCS_BUCKET_NAME not found in environment.")
|
|
|
|
session = get_session()
|
|
host = get_endpoint_host(source_location)
|
|
base_url = f"https://{host}/v1alpha"
|
|
source_parent = f"projects/{source_project_id}/locations/{source_location}/collections/default_collection/engines/{source_engine_id}/assistants/default_assistant"
|
|
source_url = f"{base_url}/{source_parent}/agents"
|
|
|
|
resp = session.get(source_url)
|
|
resp.raise_for_status()
|
|
agents = resp.json().get("agents", [])
|
|
|
|
agent_to_export = None
|
|
for agent in agents:
|
|
if agent.get("displayName") == source_agent_name and "lowCodeAgentDefinition" in agent:
|
|
agent_to_export = agent
|
|
break
|
|
|
|
if not agent_to_export:
|
|
raise ValueError(f"Source agent '{source_agent_name}' not found or is not a low-code agent.")
|
|
|
|
encoded_obj = urllib.parse.quote(object_name, safe="")
|
|
upload_url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_obj}"
|
|
up_resp = session.post(
|
|
upload_url,
|
|
data=json.dumps(agent_to_export, indent=2),
|
|
headers={"Content-Type": "application/json"}
|
|
)
|
|
up_resp.raise_for_status()
|
|
|
|
return {
|
|
"success": True,
|
|
"message": f"Successfully exported agent '{source_agent_name}' to gs://{bucket_name}/{object_name}"
|
|
}
|
|
|
|
def import_agent_from_gcs(
|
|
object_name: str,
|
|
target_project_id: str,
|
|
target_location: str,
|
|
target_engine_id: str,
|
|
bucket_name: str = "",
|
|
connector_mapping: str = ""
|
|
) -> dict:
|
|
"""Imports an agent definition from GCS and creates it in a target environment."""
|
|
if not bucket_name:
|
|
bucket_name = os.environ.get("GCS_BUCKET_NAME")
|
|
if not bucket_name:
|
|
raise ValueError("bucket_name not provided and GCS_BUCKET_NAME not found in environment.")
|
|
|
|
session = get_session()
|
|
host = get_endpoint_host(target_location)
|
|
base_url = f"https://{host}/v1alpha"
|
|
|
|
encoded_object = urllib.parse.quote(object_name, safe="")
|
|
gcs_url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{encoded_object}?alt=media"
|
|
dl_resp = session.get(gcs_url)
|
|
dl_resp.raise_for_status()
|
|
agent_data = dl_resp.json()
|
|
|
|
mapping_dict = {}
|
|
if connector_mapping:
|
|
try:
|
|
mapping_dict = json.loads(connector_mapping)
|
|
except Exception:
|
|
for pair in connector_mapping.split(","):
|
|
if ":" in pair:
|
|
k, v = pair.split(":", 1)
|
|
mapping_dict[k.strip()] = v.strip()
|
|
|
|
definition = agent_data.get("lowCodeAgentDefinition", {})
|
|
if "session" in definition:
|
|
del definition["session"]
|
|
|
|
nodes = definition.get("nodes", [])
|
|
for node in nodes:
|
|
llm_node = node.get("llmAgentNode", {})
|
|
if "selectedTools" not in llm_node:
|
|
llm_node["selectedTools"] = {"tool": []}
|
|
selected_tools = llm_node["selectedTools"]
|
|
if "tool" not in selected_tools:
|
|
selected_tools["tool"] = []
|
|
|
|
tools_list = selected_tools["tool"]
|
|
has_google_search = False
|
|
for t in tools_list:
|
|
if t.get("name") == "googleSearch":
|
|
has_google_search = True
|
|
break
|
|
if not has_google_search:
|
|
tools_list.append({"name": "googleSearch"})
|
|
|
|
new_specs = llm_node.get("dataStoreSpecs", {}).get("specs", [])
|
|
if any("snowflake" in s.get("dataStore", "").lower() for s in new_specs):
|
|
sf_conn = mapping_dict.get("Snowflake Mcp May29", "custom_mcp")
|
|
if not any(t.get("name") == sf_conn for t in tools_list):
|
|
tools_list.append({"name": sf_conn})
|
|
if any("drive" in s.get("dataStore", "").lower() for s in new_specs):
|
|
dr_conn = mapping_dict.get("ge-drive-all", "Drive")
|
|
if not any(t.get("name") == dr_conn for t in tools_list):
|
|
tools_list.append({"name": dr_conn})
|
|
|
|
if "selectedTools" in llm_node:
|
|
for t in llm_node["selectedTools"].get("tool", []):
|
|
tool_name = t.get("name")
|
|
if tool_name in mapping_dict:
|
|
t["name"] = mapping_dict[tool_name]
|
|
|
|
if "dataStoreSpecs" in llm_node:
|
|
for s in llm_node["dataStoreSpecs"].get("specs", []):
|
|
ds = s.get("dataStore", "")
|
|
for src_id, tgt_id in mapping_dict.items():
|
|
if src_id in ds:
|
|
s["dataStore"] = ds.replace(src_id, tgt_id)
|
|
ds = s["dataStore"]
|
|
|
|
source_name = agent_data.get("name", "")
|
|
match = re.search(r"projects/([^/]+)/locations/([^/]+)/collections/default_collection/engines/([^/]+)", source_name)
|
|
|
|
payload_str = json.dumps({
|
|
"displayName": agent_data.get("displayName"),
|
|
"description": agent_data.get("description", ""),
|
|
"lowCodeAgentDefinition": definition
|
|
})
|
|
|
|
if match:
|
|
source_project_id = match.group(1)
|
|
source_engine_id = match.group(3)
|
|
payload_str = payload_str.replace(f"projects/{source_project_id}", f"projects/{target_project_id}")
|
|
payload_str = payload_str.replace(source_engine_id, target_engine_id)
|
|
|
|
payload = json.loads(payload_str)
|
|
target_parent = f"projects/{target_project_id}/locations/{target_location}/collections/default_collection/engines/{target_engine_id}/assistants/default_assistant"
|
|
target_url = f"{base_url}/{target_parent}/agents"
|
|
|
|
logging.info(f"Creating new agent from GCS at target {target_url}")
|
|
create_resp = session.post(target_url, json=payload)
|
|
create_resp.raise_for_status()
|
|
|
|
return {
|
|
"success": True,
|
|
"message": f"Successfully imported agent from gs://{bucket_name}/{object_name} to target environment.",
|
|
"target_agent": create_resp.json()
|
|
}
|
|
|
|
def list_datastores(project_id: str, location: str = "global", collection: str = "default_collection") -> list:
|
|
"""Lists all available datastores and their IDs in a target project/location."""
|
|
session = get_session()
|
|
host = get_endpoint_host(location)
|
|
base_url = f"https://{host}/v1alpha"
|
|
url = f"{base_url}/projects/{project_id}/locations/{location}/collections/{collection}/dataStores"
|
|
|
|
logging.info(f"Fetching datastores from {url}")
|
|
resp = session.get(url)
|
|
resp.raise_for_status()
|
|
|
|
datastores = resp.json().get("dataStores", [])
|
|
results = []
|
|
for ds in datastores:
|
|
ds_name = ds.get("name", "")
|
|
ds_id = ds_name.split("/")[-1] if ds_name else "Unknown ID"
|
|
results.append({
|
|
"id": ds_id,
|
|
"displayName": ds.get("displayName", ""),
|
|
"name": ds_name
|
|
})
|
|
return results
|
|
|
|
def export_notebook_to_gcs(
|
|
notebook_id_or_title: str,
|
|
object_name: str,
|
|
bucket_name: str = "",
|
|
source_project_number: str = "",
|
|
source_location: str = "global"
|
|
) -> dict:
|
|
"""Exports a notebook definition (including metadata and sources) to GCS."""
|
|
if not bucket_name:
|
|
bucket_name = os.environ.get("GCS_BUCKET_NAME")
|
|
if not bucket_name:
|
|
raise ValueError("bucket_name not provided and GCS_BUCKET_NAME not found in environment.")
|
|
|
|
notebooks = list_notebooks(source_project_number, source_location)
|
|
source_notebook = None
|
|
for nb in notebooks:
|
|
nb_id = nb.get("name", "").split("/")[-1]
|
|
nb_title = nb.get("title", "")
|
|
if nb_id == notebook_id_or_title or nb_title == notebook_id_or_title:
|
|
source_notebook = nb
|
|
break
|
|
|
|
if not source_notebook:
|
|
raise ValueError(f"Source notebook '{notebook_id_or_title}' not found.")
|
|
|
|
source_nb_id = source_notebook.get("name", "").split("/")[-1]
|
|
source_nb_title = source_notebook.get("title", "Untitled Notebook")
|
|
|
|
sources = list_sources_and_types(source_nb_id, source_project_number, source_location)
|
|
|
|
exported_sources = []
|
|
for src in sources:
|
|
title = src.get("title", "Untitled Source")
|
|
raw_data = src.get("raw_data", {})
|
|
payload = extract_notebook_source_payload(raw_data)
|
|
|
|
exported_sources.append({
|
|
"title": title,
|
|
"userContent": payload
|
|
})
|
|
|
|
notebook_def = {
|
|
"title": source_nb_title,
|
|
"sources": exported_sources
|
|
}
|
|
|
|
session = get_session()
|
|
encoded_obj = urllib.parse.quote(object_name, safe="")
|
|
upload_url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_obj}"
|
|
up_resp = session.post(
|
|
upload_url,
|
|
data=json.dumps(notebook_def, indent=2),
|
|
headers={"Content-Type": "application/json"}
|
|
)
|
|
up_resp.raise_for_status()
|
|
|
|
return {
|
|
"success": True,
|
|
"message": f"Successfully exported notebook '{source_nb_title}' to gs://{bucket_name}/{object_name}",
|
|
"notebook_def": notebook_def
|
|
}
|
|
|
|
def import_notebook_from_gcs(
|
|
object_name: str,
|
|
target_project_number: str,
|
|
target_location: str,
|
|
bucket_name: str = ""
|
|
) -> dict:
|
|
"""Imports a notebook definition from GCS and creates it in the target project."""
|
|
if not bucket_name:
|
|
bucket_name = os.environ.get("GCS_BUCKET_NAME")
|
|
if not bucket_name:
|
|
raise ValueError("bucket_name not provided and GCS_BUCKET_NAME not found in environment.")
|
|
|
|
session = get_session()
|
|
encoded_object = urllib.parse.quote(object_name, safe="")
|
|
gcs_url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{encoded_object}?alt=media"
|
|
dl_resp = session.get(gcs_url)
|
|
dl_resp.raise_for_status()
|
|
nb_data = dl_resp.json()
|
|
|
|
title = nb_data.get("title", "Untitled Notebook")
|
|
sources = nb_data.get("sources", [])
|
|
|
|
target_nb = create_notebook(target_project_number, target_location, title)
|
|
target_nb_id = target_nb.get("name", "").split("/")[-1]
|
|
logging.info(f"Created target notebook '{title}' from GCS with ID: {target_nb_id}")
|
|
|
|
migrated_sources = []
|
|
failed_sources = []
|
|
|
|
for src in sources:
|
|
src_title = src.get("title", "Untitled Source")
|
|
payload = src.get("userContent", {})
|
|
|
|
logging.info(f"Adding source '{src_title}' to target notebook '{target_nb_id}'")
|
|
try:
|
|
add_resp = add_source_to_notebook(target_project_number, target_location, target_nb_id, payload)
|
|
migrated_sources.append({
|
|
"title": src_title,
|
|
"status": "success",
|
|
"response": add_resp
|
|
})
|
|
except Exception as e:
|
|
logging.error(f"Failed to add source '{src_title}': {e}")
|
|
failed_sources.append({
|
|
"title": src_title,
|
|
"status": "failed",
|
|
"error": str(e)
|
|
})
|
|
|
|
return {
|
|
"success": len(failed_sources) == 0,
|
|
"target_notebook_id": target_nb_id,
|
|
"target_notebook_title": title,
|
|
"migrated_sources_count": len(migrated_sources),
|
|
"failed_sources_count": len(failed_sources),
|
|
"migrated": migrated_sources,
|
|
"failed": failed_sources
|
|
}
|