Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion samples/adk-sql-agent/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import google.auth
import google.auth.transport.requests
import grpc
Expand Down Expand Up @@ -97,6 +96,7 @@ def setup_opentelemetry() -> None:

# Load instrumentors
SQLite3Instrumentor().instrument()
# ADK uses Vertex AI and Google Gen AI SDKs.
VertexAIInstrumentor().instrument()
GoogleGenAiSdkInstrumentor().instrument()

Expand Down
5 changes: 0 additions & 5 deletions samples/adk-sql-agent/sql_agent/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,6 @@

import sqlite3

from opentelemetry import trace

# from utils import ask_prompt, console, print_markdown, render_messages


SYSTEM_PROMPT = f"""\
You are a helpful AI assistant with a mastery of database design and querying. You have access
to an ephemeral sqlite3 database that you can query and modify through some tools. Help answer
Expand Down
2 changes: 2 additions & 0 deletions samples/adk-sql-agent/sql_agent/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ class SqlRunResult(TypedDict):
rows: NotRequired[list[tuple[str, ...]]]
"""The rows returned by the SQL query"""

# [START opentelemetry_adk_agent_span]
@tracer.start_as_current_span("create_database")
def create_database_tool(tool_context: ToolContext) -> dict[str, Any]:
"""Creates a temporary file in the /tmp directory to hold an ephemeral
Expand All @@ -46,6 +47,7 @@ def create_database_tool(tool_context: ToolContext) -> dict[str, Any]:
tool_context.state[SESSION_DB_KEY] = path
return {"resp": "Created an ephemeral database"}
return {"resp": f"Skipping database creation, {tool_context.state[SESSION_DB_KEY]} already exists"}
# [END opentelemetry_adk_agent_span]

@tracer.start_as_current_span("run_sql")
def run_sql_tool(sql_query: str, tool_context: ToolContext) -> dict[str, Any]:
Expand Down