14
14
15
15
16
16
import logging
17
- from typing import Any , NotRequired , TypedDict
18
- from opentelemetry import trace
19
17
import sqlite3
20
- from google . adk . tools import ToolContext
18
+ from typing import Any , NotRequired , TypedDict
21
19
22
20
from google .genai import Client
23
21
from google .genai import types as genai_types
24
22
23
+ from opentelemetry import trace
24
+ from opentelemetry .instrumentation .sqlite3 import SQLite3Instrumentor
25
+
26
+
25
27
tracer = trace .get_tracer (__name__ )
26
28
logger = logging .getLogger (__name__ )
27
29
@@ -35,11 +37,13 @@ class SqlRunResult(TypedDict):
35
37
36
38
37
39
def create_run_sql_tool (dbpath : str ):
40
+ instrumentor = SQLite3Instrumentor ()
41
+
38
42
@tracer .start_as_current_span ("run_sql" )
39
43
def run_sql (sql_query : str ) -> dict [str , Any ]:
40
44
"""Runs a SQLite query. The SQL query can be DDL or DML. Returns the rows if it's a SELECT query."""
41
45
42
- with sqlite3 .connect (dbpath ) as db :
46
+ with instrumentor . instrument_connection ( sqlite3 .connect (dbpath ) ) as db :
43
47
try :
44
48
cursor = db .cursor ()
45
49
cursor .execute (sql_query )
@@ -81,17 +85,25 @@ def run_sql(sql_query: str) -> dict[str, Any]:
81
85
# Tool that calls gemini with an image
82
86
@tracer .start_as_current_span ("describe_uri_tool" )
83
87
def describe_uri_tool (uri : str , mime_type : str ) -> dict [str , Any ]:
84
- """describe_uri_tool is a tool that takes a URI and mime type for the URI (you can guess
85
- the type your best) and returns a description or error
88
+ """describe_uri_tool is a tool that takes a URI (http(s) or gs:// for GCS) and mime type
89
+ for the URI (you can guess the type your best) and returns a description or error.
86
90
"""
87
91
client = Client ()
88
- res = client .models .generate_content (
89
- model = "gemini-2.0-flash-lite" ,
90
- contents = [
91
- genai_types .Part .from_text (
92
- text = "Describe the following attached file:"
93
- ),
94
- genai_types .Part .from_uri (file_uri = uri , mime_type = mime_type ),
95
- ],
96
- )
97
- return res .model_dump ()
92
+
93
+ try :
94
+ res = client .models .generate_content (
95
+ model = "gemini-2.0-flash-lite" ,
96
+ contents = [
97
+ genai_types .Part .from_text (
98
+ text = "Describe the following attached file:"
99
+ ),
100
+ genai_types .Part .from_uri (file_uri = uri , mime_type = mime_type ),
101
+ ],
102
+ )
103
+ except Exception as e :
104
+ return {"error" : str (e )}
105
+
106
+ if not res .candidates :
107
+ return {"error" : "No candidates returned" }
108
+
109
+ return res .candidates [0 ].model_dump (mode = "json" )
0 commit comments