Adding asset tracking telemetry testing (#5)

* Update gitlab.py * Update gitlab.py
GoogleCloudPlatform · Apr 19, 2024 · 8877556 · 8877556
1 parent f0364ea
commit 8877556
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/outer-loop-cli/src/devai/commands/gitlab.py b/outer-loop-cli/src/devai/commands/gitlab.py
@@ -18,8 +18,12 @@
 from langchain_google_vertexai import ChatVertexAI
 from langchain_community.agent_toolkits.gitlab.toolkit import GitLabToolkit
 from langchain_community.utilities.gitlab import GitLabAPIWrapper
+from google.cloud.aiplatform import telemetry
+
+USER_AGENT = 'cloud-solutions/genai-outer-loop-devai'
 
-llm = ChatVertexAI(model_name="gemini-pro",
+with telemetry.tool_context_manager(USER_AGENT):
+  llm = ChatVertexAI(model_name="gemini-pro",
                    convert_system_message_to_human=True,
                    project=os.environ["PROJECT_ID"],
                    location=os.environ["LOCATION"])