1+ import uuid
2+ from typing import Any
3+
4+ import psycopg
5+ from ag_ui .core import ActivitySnapshotEvent
6+ from pydantic_ai import Agent , ToolReturn
7+
8+ from ravnar .agents import PydanticAiAgentWrapper
9+
10+
11+ class PostgresDatabase :
12+ def __init__ (
13+ self ,
14+ * ,
15+ host : str ,
16+ port : int = 5432 ,
17+ name : str ,
18+ user : str ,
19+ password : str ,
20+ ) -> None :
21+ self ._host = host
22+ self ._port = port
23+ self ._name = name
24+ self ._user = user
25+ self ._password = password
26+
27+ async def _connect (self ) -> psycopg .AsyncConnection [Any ]:
28+ return await psycopg .AsyncConnection .connect (
29+ dbname = self ._name ,
30+ user = self ._user ,
31+ password = self ._password ,
32+ host = self ._host ,
33+ port = self ._port ,
34+ )
35+
36+ async def execute (self , query : str ) -> list [tuple ]:
37+ async with await self ._connect () as conn , conn .cursor () as cur :
38+ await cur .execute (query )
39+ return list (await cur .fetchall ())
40+
41+
42+ def create_agent (
43+ agent : Agent ,
44+ database : PostgresDatabase ,
45+ ) -> PydanticAiAgentWrapper :
46+ @agent .system_prompt
47+ def _system_prompt () -> str :
48+ return """
49+ You are an agent with the sole task of answering the user's
50+ questions related to the Austin Permits Database. You have a suite
51+ of tools available to fetch the db schema, make queries, and create
52+ visualizations from the data. The database is mounted as readonly,
53+ so don't generate any SQL that would modify the database. Your
54+ query will fail if it attempts to modify the db. Plan your actions
55+ accordingly by ensuring that you understand the db schema before
56+ generating SQL queries.
57+ """
58+
59+ @agent .tool_plain
60+ async def get_db_schema () -> list [tuple ]:
61+ """Get the schema for the Austin Permits database.
62+
63+ Use this tool to understand the db structure before issuing
64+ any queries to the db.
65+
66+ """
67+ query = """
68+ SELECT
69+ t.table_name,
70+ c.column_name,
71+ c.data_type,
72+ c.is_nullable,
73+ c.ordinal_position
74+ FROM information_schema.tables t
75+ JOIN information_schema.columns c ON t.table_name = c.table_name
76+ WHERE t.table_schema = 'public' AND t.table_type = 'BASE TABLE'
77+ ORDER BY t.table_name, c.ordinal_position;
78+ """
79+ return await database .execute (query )
80+
81+ @agent .tool_plain
82+ async def execute_query (query : str ) -> list [tuple ]:
83+ """Execute a query against the Austin Permits database.
84+
85+ The query should not attempt to modify the db in any way.
86+
87+ The db is mounted as readonly, so queries that attempt to modify it will fail.
88+
89+ This tool takes a single argument, which is the SQL to execute
90+ against the db. The db is Postgres, so use Postgres-compatible
91+ SQL syntax.
92+
93+ """
94+ return await database .execute (query )
95+
96+ @agent .tool_plain
97+ async def create_chart (option : dict ) -> ToolReturn :
98+ """Create a chart from the data retrieved by a query to the db.
99+
100+ This tool takes a single argument, which is a Python dictionary
101+ that follows the format of an Apache ECharts configuration object.
102+
103+ The object must be serializable to JSON, so it cannot include
104+ JS function callbacks. Only use the features of the Apache ECharts
105+ config that can be serialized as plain JSON data.
106+
107+ """
108+ return ToolReturn (
109+ return_value = "Chart Created" ,
110+ metadata = [
111+ ActivitySnapshotEvent (
112+ message_id = str (uuid .uuid4 ()), activity_type = "application/json+echart" , content = option
113+ )
114+ ],
115+ )
116+
117+ @agent .tool_plain
118+ async def create_map (data : dict ) -> ToolReturn :
119+ """Create a map with markers from the data retrieved by a query to the db.
120+
121+ This tool takes a single argument, which is a Python dictionary of
122+ the following form:
123+ { 'center': [number, number], 'features': GeoJSONFeatureCollection }
124+
125+ Use the 'center' key to define the center of the map in
126+ [latitude, longitude] floating point numbers.
127+
128+ Use the 'features' key to define an additional GeoJSON feature
129+ collection, such as markers with popup metadata, to add to the map.
130+ This dictionary must be a valid GeoJSON feature collection.
131+
132+ To add a popup, create a 'popup' key in the properties of the feature.
133+ Use an HTML string to define the contents of the popup. Any other key
134+ in the properties of a feature besides 'popup' will be ignored.
135+
136+ If a link to the permit is available, include it in any popup.
137+
138+ """
139+ return ToolReturn (
140+ return_value = "Map Created" ,
141+ metadata = [
142+ ActivitySnapshotEvent (
143+ message_id = str (uuid .uuid4 ()), activity_type = "application/json+leaflet" , content = data
144+ )
145+ ],
146+ )
147+
148+ return PydanticAiAgentWrapper (agent )
0 commit comments