Skip to content

Add vertex search tool with integration tests#73

Open
amrit110 wants to merge 5 commits intomainfrom
add_vertex_search_tool
Open

Add vertex search tool with integration tests#73
amrit110 wants to merge 5 commits intomainfrom
add_vertex_search_tool

Conversation

@amrit110
Copy link
Member

@amrit110 amrit110 commented Mar 2, 2026

Summary

Clickup Ticket(s): Link(s) if applicable.

Type of Change

  • 🐛 Bug fix (non-breaking change that fixes an issue)
  • ✨ New feature (non-breaking change that adds functionality)
  • 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • 📝 Documentation update
  • 🔧 Refactoring (no functional changes)
  • ⚡ Performance improvement
  • 🧪 Test improvements
  • 🔒 Security fix

Changes Made

  • Add a new tool that allows for searching over custom knowledge base using vertex AI
  • Add tests including a complete integration test using a dummy database
  • Add a notebook showing how the custom knowledge base can be setup and integrated

Testing

  • Tests pass locally (uv run pytest tests/)
  • Type checking passes (uv run mypy <src_dir>)
  • Linting passes (uv run ruff check src_dir/)
  • Manual testing performed (describe below)

Manual testing details:

Screenshots/Recordings

Related Issues

Deployment Notes

Checklist

  • Code follows the project's style guidelines
  • Self-review of code completed
  • Documentation updated (if applicable)
  • No sensitive information (API keys, credentials) exposed

@amrit110 amrit110 self-assigned this Mar 2, 2026
@amrit110 amrit110 added enhancement New feature or request tests Add or update unit or integration tests labels Mar 2, 2026
Copy link
Collaborator

@lotif lotif left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approved with minor comments.

}

except Exception as e:
logger.error(f"Vertex AI Search failed: {e}")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you use logger.exception the log will include additional information about the exception automatically.

datastore_id = config.vertex_datastore_id
location = config.google_cloud_location

async def vertex_search(query: str) -> dict[str, Any]:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Inner functions are a bit confusing. This would work the same way if this function is made into a module level function with the added benefit of being a regular function that doesn't need a lot of explanations and is easier to debug.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When is the user supposed to run this script? I don't see it on any README or notebook.

gcloud auth application-default login

# Run from the repo root
uv run python scripts/create_test_datastore.py \\
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think a better practice is to run scripts as modules instead: uv run python -m scripts.create_test_datastore.

https://stackoverflow.com/a/22241810/1116627

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request tests Add or update unit or integration tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants