Created
October 24, 2024 16:30
-
-
Save koverholt/f6a7725168c707871a19189834692275 to your computer and use it in GitHub Desktop.
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
| PROJECT_ID = "[your-project-id]" # @param {type:"string"} | |
| LOCATION = "us-central1" # @param {type:"string"} | |
| import vertexai | |
| vertexai.init(project=PROJECT_ID, location=LOCATION) | |
| from vertexai.generative_models import ( | |
| FunctionDeclaration, | |
| GenerationConfig, | |
| GenerativeModel, | |
| Tool, | |
| ToolConfig, | |
| ) | |
| list_tables_func = FunctionDeclaration( | |
| name="list_tables", | |
| description="List tables in a dataset that will help answer the user's question", | |
| parameters={ | |
| "type": "object", | |
| "properties": { | |
| "dataset_id": { | |
| "type": "string", | |
| "description": "Dataset ID to fetch tables from.", | |
| } | |
| }, | |
| "required": [ | |
| "dataset_id", | |
| ], | |
| }, | |
| ) | |
| tool = Tool( | |
| function_declarations=[ | |
| list_tables_func, | |
| ], | |
| ) | |
| model = GenerativeModel( | |
| "gemini-1.5-pro", | |
| generation_config=GenerationConfig(temperature=0), | |
| tools=[tool], | |
| ) | |
| prompt = "What's the application of the Navier-Stokes equation?" | |
| response = model.generate_content( | |
| prompt, | |
| tools=[tool], | |
| tool_config=ToolConfig( | |
| function_calling_config=ToolConfig.FunctionCallingConfig( | |
| mode=ToolConfig.FunctionCallingConfig.Mode.AUTO, # AUTO is the default model behavior, just adding it here to be explicit | |
| ) | |
| ) | |
| ) | |
| print(response.text) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment