diff --git a/samples/adk-sql-agent/main.py b/samples/adk-sql-agent/main.py index cae2acde..c3563f81 100644 --- a/samples/adk-sql-agent/main.py +++ b/samples/adk-sql-agent/main.py @@ -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 @@ -97,19 +96,20 @@ def setup_opentelemetry() -> None: # Load instrumentors SQLite3Instrumentor().instrument() + # ADK uses Vertex AI and Google Gen AI SDKs. VertexAIInstrumentor().instrument() GoogleGenAiSdkInstrumentor().instrument() # [END opentelemetry_adk_otel_setup] - +# [START opentelemetry_adk_launch_web_interface] def main() -> None: # Make sure to set: # OTEL_PYTHON_LOGGING_AUTO_INSTRUMENTATION_ENABLED=true # OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT=true # in order to full prompts and responses and logs messages. - # For this sample, these can be set by loading the `main.env` file. + # For this sample, these can be set by loading the `opentelemetry.env` file. setup_opentelemetry() # Call the function to get the FastAPI app instance. @@ -125,7 +125,9 @@ def main() -> None: web=SERVE_WEB_INTERFACE, ) + # Lauch the web interface on port 8080. uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8080))) +# [END opentelemetry_adk_launch_web_interface] if __name__ == "__main__": diff --git a/samples/adk-sql-agent/sql_agent/agent.py b/samples/adk-sql-agent/sql_agent/agent.py index 64ee513a..a3f584f7 100644 --- a/samples/adk-sql-agent/sql_agent/agent.py +++ b/samples/adk-sql-agent/sql_agent/agent.py @@ -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 diff --git a/samples/adk-sql-agent/sql_agent/tools.py b/samples/adk-sql-agent/sql_agent/tools.py index 912b659c..97410ffb 100644 --- a/samples/adk-sql-agent/sql_agent/tools.py +++ b/samples/adk-sql-agent/sql_agent/tools.py @@ -33,7 +33,6 @@ class SqlRunResult(TypedDict): rows: NotRequired[list[tuple[str, ...]]] """The rows returned by the SQL query""" -@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 sqlite3 database if a database is not found for the current session. @@ -47,7 +46,6 @@ def create_database_tool(tool_context: ToolContext) -> dict[str, Any]: return {"resp": "Created an ephemeral database"} return {"resp": f"Skipping database creation, {tool_context.state[SESSION_DB_KEY]} already exists"} -@tracer.start_as_current_span("run_sql") def run_sql_tool(sql_query: str, tool_context: ToolContext) -> dict[str, Any]: """Runs a SQLite query. The SQL query can be DDL or DML. Returns the rows if it's a SELECT query.""" current_session_db_path = tool_context.state.get(SESSION_DB_KEY)