Add public trace example

Author: Viicos

docs/index.md

@@ -134,13 +134,10 @@ async def add_customer_name(ctx: RunContext[SupportDependencies]) -> str:
 
 
 @support_agent.tool  # (6)!
-async def customer_balance(
-    ctx: RunContext[SupportDependencies], include_pending: bool
-) -> float:
+async def customer_balance(ctx: RunContext[SupportDependencies]) -> float:
     """Returns the customer's current account balance."""  # (7)!
     return await ctx.deps.db.customer_balance(
         id=ctx.deps.customer_id,
-        include_pending=include_pending,
     )
 
 
@@ -196,7 +193,7 @@ import logfire
 
 logfire.configure()  # (1)!
 logfire.instrument_pydantic_ai()  # (2)!
-logfire.instrument_asyncpg()  # (3)!
+logfire.instrument_sqlite3()  # (3)!
 
 ...
 
@@ -213,11 +210,14 @@ support_agent = Agent(
 
 1. Configure the Logfire SDK, this will fail if project is not set up.
 2. This will instrument all Pydantic AI agents used from here on out. If you want to instrument only a specific agent, you can pass the [`instrument=True` keyword argument][pydantic_ai.Agent.__init__] to the agent.
-3. In our demo, `DatabaseConn` uses [`asyncpg`]() to connect to a PostgreSQL database, so [`logfire.instrument_asyncpg()`](https://magicstack.github.io/asyncpg/current/) is used to log the database queries.
+3. In our demo, `DatabaseConn` uses [`sqlite3`][] to connect to a PostgreSQL database, so [`logfire.instrument_sqlite3()`](https://logfire.pydantic.dev/docs/integrations/databases/sqlite3/)
+   is used to log the database queries.
 
 That's enough to get the following view of your agent in action:
 
-{{ video('9078b98c4f75d01f912a0368bbbdb97a', 25, 55) }}
+/// public-trace | https://logfire-eu.pydantic.dev/public-trace/a2957caa-b7b7-4883-a529-777742649004?spanId=31aade41ab896144
+    title: 'Logfire instrumentation for the bank agent'
+///
 
 See [Monitoring and Performance](logfire.md) to learn more.
 

examples/pydantic_ai_examples/bank_support.py

@@ -5,36 +5,37 @@
     uv run -m pydantic_ai_examples.bank_support
 """
 
+import sqlite3
 from dataclasses import dataclass
 
 from pydantic import BaseModel
 
 from pydantic_ai import Agent, RunContext
 
 
+@dataclass
 class DatabaseConn:
-    """This is a fake database for example purposes.
-
-    In reality, you'd be connecting to an external database
-    (e.g. PostgreSQL) to get information about customers.
-    """
-
-    @classmethod
-    async def customer_name(cls, *, id: int) -> str | None:
-        if id == 123:
-            return 'John'
-
-    @classmethod
-    async def customer_balance(cls, *, id: int, include_pending: bool) -> float:
-        if id == 123:
-            if include_pending:
-                return 123.45
-            else:
-                return 100.00
+    """A wrapper over the SQLite connection."""
+
+    sqlite_conn: sqlite3.Connection
+
+    async def customer_name(self, *, id: int) -> str | None:
+        res = cur.execute('SELECT name FROM customers WHERE id=?', (id,))
+        row = res.fetchone()
+        if row:
+            return row[0]
+        return None
+
+    async def customer_balance(self, *, id: int) -> float:
+        res = cur.execute('SELECT balance FROM customers WHERE id=?', (id,))
+        row = res.fetchone()
+        if row:
+            return row[0]
         else:
             raise ValueError('Customer not found')
 
 
+
 @dataclass
 class SupportDependencies:
     customer_id: int
@@ -69,27 +70,33 @@ async def add_customer_name(ctx: RunContext[SupportDependencies]) -> str:
 
 
 @support_agent.tool
-async def customer_balance(
-    ctx: RunContext[SupportDependencies], include_pending: bool
-) -> str:
+async def customer_balance(ctx: RunContext[SupportDependencies]) -> str:
     """Returns the customer's current account balance."""
     balance = await ctx.deps.db.customer_balance(
         id=ctx.deps.customer_id,
-        include_pending=include_pending,
     )
     return f'${balance:.2f}'
 
 
 if __name__ == '__main__':
-    deps = SupportDependencies(customer_id=123, db=DatabaseConn())
-    result = support_agent.run_sync('What is my balance?', deps=deps)
-    print(result.output)
-    """
-    support_advice='Hello John, your current account balance, including pending transactions, is $123.45.' block_card=False risk=1
-    """
-
-    result = support_agent.run_sync('I just lost my card!', deps=deps)
-    print(result.output)
-    """
-    support_advice="I'm sorry to hear that, John. We are temporarily blocking your card to prevent unauthorized transactions." block_card=True risk=8
-    """
+    with sqlite3.connect(':memory:') as con:
+        cur = con.cursor()
+        cur.execute('CREATE TABLE customers(id, name, balance)')
+        cur.execute("""
+            INSERT INTO customers VALUES
+                (123, 'John', 123.45)
+        """)
+        con.commit()
+
+        deps = SupportDependencies(customer_id=123, db=DatabaseConn(sqlite_conn=con))
+        result = support_agent.run_sync('What is my balance?', deps=deps)
+        print(result.output)
+        """
+        support_advice='Hello John, your current account balance, including pending transactions, is $123.45.' block_card=False risk=1
+        """
+
+        result = support_agent.run_sync('I just lost my card!', deps=deps)
+        print(result.output)
+        """
+        support_advice="I'm sorry to hear that, John. We are temporarily blocking your card to prevent unauthorized transactions." block_card=True risk=8
+        """