Uprev pydantic AI (#6)

* uprev pydantic-ai

* slugify

* add typechecking

* add ci for pre-commit

Author: samuelcolvin

.github/workflows/ci.yml

@@ -0,0 +1,28 @@
+name: CI
+on:
+  push:
+    branches:
+      - main
+    tags:
+      - "**"
+  pull_request: {}
+
+env:
+  COLUMNS: 150
+  UV_PYTHON: 3.13
+  UV_FROZEN: "1"
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
+      - run: uv sync --all-packages
+      - uses: actions/cache@v4
+        with:
+          path: ~/.cache/pre-commit
+          key: pre-commit|${{ env.UV_PYTHON }}|${{ hashFiles('.pre-commit-config.yaml') }}
+      - run: uvx pre-commit run --color=always --all-files --verbose

.pre-commit-config.yaml

@@ -19,15 +19,15 @@ repos:
     hooks:
       - id: format
         name: Format Python
-        entry: uv
-        args: [run, ruff, format]
+        entry: make
+        args: [format]
         language: system
         types: [python]
         pass_filenames: false
-      - id: lint
-        name: lint Python
-        entry: uv
-        args: [run, ruff, check, --fix]
+      - id: typecheck
+        name: Typecheck Python
+        entry: make
+        args: [typecheck]
         language: system
         types: [python]
         pass_filenames: false

Makefile

@@ -0,0 +1,42 @@
+.DEFAULT_GOAL := all
+
+.PHONY: .uv
+.uv: ## Check that uv is installed
+	@uv --version || echo 'Please install uv: https://docs.astral.sh/uv/getting-started/installation/'
+
+.PHONY: .pre-commit
+.pre-commit: ## Check that pre-commit is installed
+	@pre-commit -V || echo 'Please install pre-commit: https://pre-commit.com/'
+
+.PHONY: install
+install: .uv .pre-commit ## Install the package, dependencies, and pre-commit for local development
+	uv sync --frozen --all-packages --all-extras
+	pre-commit install --install-hooks
+
+.PHONY: format
+format: ## Format the code
+	uv run ruff format
+	uv run ruff check --fix --fix-only
+
+.PHONY: lint
+lint: ## Lint the code
+	uv run ruff format --check
+	uv run ruff check
+
+.PHONY: typecheck
+typecheck:
+	uv run basedpyright
+
+
+.PHONY: help
+help: ## Show this help (usage: make help)
+	@echo "Usage: make [recipe]"
+	@echo "Recipes:"
+	@awk '/^[a-zA-Z0-9_-]+:.*?##/ { \
+		helpMessage = match($$0, /## (.*)/); \
+		if (helpMessage) { \
+			recipe = $$1; \
+			sub(/:/, "", recipe); \
+			printf "  \033[36m%-20s\033[0m %s\n", recipe, substr($$0, RSTART + 3, RLENGTH); \
+		} \
+	}' $(MAKEFILE_LIST)

distributed-frontend-example/main.py

@@ -3,7 +3,6 @@
 from pathlib import Path
 from uuid import uuid4
 
-import httpx
 import logfire
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
@@ -80,13 +79,14 @@ async def generate_image(prompt: str) -> GenerateResponse:
 # Proxy to Logfire for client traces from the browser
 @app.post('/client-traces')
 async def client_traces(request: Request):
-    async with httpx.AsyncClient() as client:
-        response = await client.request(
-            method=request.method,
-            url=f'{logfire_base_url}v1/traces',
-            headers=dict(Authorization=logfire_token),
-            json=await request.json(),
-        )
+    assert logfire_token is not None, 'Logfire token is not set'
+    response = await http_client.request(
+        method=request.method,
+        url=f'{logfire_base_url}v1/traces',
+        headers={'Authorization': logfire_token},
+        json=await request.json(),
+    )
+    response.raise_for_status()
 
     return {
         'status_code': response.status_code,

pai-mcp-sampling/client.py

@@ -9,11 +9,12 @@
 logfire.instrument_mcp()
 
 server = MCPServerStdio(command='python', args=[str(Path(__file__).parent / 'generate_svg.py')])
-agent = Agent('openai:gpt-4.1-mini', mcp_servers=[server])
+agent = Agent('openai:gpt-4.1-mini', toolsets=[server])
+agent.set_mcp_sampling_model()
 
 
 async def main():
-    async with agent.run_mcp_servers():
+    async with agent:
         result = await agent.run('Create an image of a robot in a punk style, it should be pink.')
     print(result.output)
 

pai-mcp-sampling/generate_svg.py

@@ -21,7 +21,8 @@ async def image_generator(ctx: Context[ServerSessionT, LifespanContextT, Request
     # run the agent, using MCPSamplingModel to proxy the LLM call through the client.
     svg_result = await svg_agent.run(f'{subject=} {style=}', model=MCPSamplingModel(ctx.session))
 
-    path = Path(f'{subject}_{style}.svg')
+    path = Path(f'{slugify(subject)}_{slugify(style)}.svg')
+    logfire.info(f'writing file to {path}')
     # remove triple backticks if the svg was returned within markdown
     if m := re.search(r'^```\w*$(.+?)```$', svg_result.output, re.S | re.M):
         path.write_text(m.group(1))
@@ -30,6 +31,10 @@ async def image_generator(ctx: Context[ServerSessionT, LifespanContextT, Request
     return f'See {path}'
 
 
+def slugify(text: str) -> str:
+    return re.sub(r'\W+', '-', text.lower())
+
+
 if __name__ == '__main__':
     # run the server via stdio
     app.run()

pyproject.toml

@@ -8,12 +8,17 @@ dependencies = [
     "asyncpg>=0.30.0",
     "devtools>=0.12.2",
     "fastapi>=0.115.14",
-    "logfire[asyncpg,fastapi,httpx]==4.3.6",
-    "pydantic-ai>=0.3.6",
+    "logfire[asyncpg,fastapi,httpx]>=4.10",
+    "mcp>=1.15.0",
+    "pydantic-ai>=1",
 ]
 
 [dependency-groups]
-dev = ["ruff>=0.12.2", "asyncpg-stubs>=0.30.2"]
+dev = [
+    "ruff>=0.12.2",
+    "asyncpg-stubs>=0.30.2",
+    "basedpyright>=1.31.6",
+]
 
 [tool.ruff]
 line-length = 120
@@ -32,3 +37,13 @@ convention = "google"
 # don't format python in docstrings, pytest-examples takes care of it
 docstring-code-format = false
 quote-style = "single"
+
+[tool.pyright]
+pythonVersion = "3.12"
+typeCheckingMode = "strict"
+reportUnnecessaryTypeIgnoreComment = true
+reportIgnoreCommentWithoutRule = true
+reportUnusedParameter = true
+reportIncompatibleUnannotatedOverride = true
+reportImplicitAbstractClass = true
+venv = ".venv"