chore: add yaml files to the ADK Vertex AI Search datastore

PiperOrigin-RevId: 808895175
This commit is contained in:
Xuan Yang
2025-09-18 23:27:34 -07:00
committed by Copybara-Service
parent 006a406f5b
commit 4d39563ea4
2 changed files with 19 additions and 2 deletions
@@ -72,8 +72,8 @@ def upload_directory_to_gcs(
# into hidden directories.
dirs[:] = [d for d in dirs if not d.startswith(".")]
# Keep only .md and .py files.
files = [f for f in files if f.endswith(".md") or f.endswith(".py")]
# Keep only .md, .py and .yaml files.
files = [f for f in files if f.endswith((".md", ".py", ".yaml"))]
for filename in files:
local_path = os.path.join(root, filename)
@@ -99,6 +99,19 @@ def upload_directory_to_gcs(
bucket.blob(gcs_path).upload_from_string(
html_content, content_type=content_type
)
elif filename.lower().endswith(".yaml"):
# Vertex AI search doesn't recognize yaml,
# convert it to text and use text/plain instead
content_type = "text/plain"
with open(local_path, "r", encoding="utf-8") as f:
yaml_content = f.read()
if not yaml_content:
print(" - Skipped empty file: " + local_path)
continue
gcs_path = gcs_path.removesuffix(".yaml") + ".txt"
bucket.blob(gcs_path).upload_from_string(
yaml_content, content_type=content_type
)
else: # Python files
bucket.blob(gcs_path).upload_from_filename(
local_path, content_type=content_type
@@ -115,6 +115,10 @@ def convert_gcs_to_https(gcs_uri: str) -> Optional[str]:
if relative_path.endswith(".html"):
relative_path = relative_path.removesuffix(".html") + ".md"
# Replace .txt with .yaml
if relative_path.endswith(".txt"):
relative_path = relative_path.removesuffix(".txt") + ".yaml"
# Convert the links for adk-docs
if prefix == "adk-docs" and relative_path.startswith("docs/"):
path_after_docs = relative_path[len("docs/") :]