From e3373026d6626077bff67c047088d5813dae6447 Mon Sep 17 00:00:00 2001 From: Wei Ji <23487320+weiji14@users.noreply.github.com> Date: Tue, 9 Dec 2025 21:22:01 +1300 Subject: [PATCH] Enable pydocstyle (D) ruff rule and add more docs (#21) * Use pydocstyle (D) rule with google convention Add a ruff rule to catch missing documentation. Using google convention so that undocumented-param (D417) rule is enabled to catch missing params, xref https://docs.astral.sh/ruff/rules/undocumented-param. Extended to include D213 (instead of D212) and D410 rules too. * Fix D100 Missing docstring in public module * Fix D101 Missing docstring in public class * Fix D103 Missing docstring in public function Also ignore rule D205 to allow first sentence of docstring to wrap to multiple lines. * Fix D417 Missing argument description in the docstring * Update indent in pyproject.toml file --- pyproject.toml | 14 ++++++++++++++ src/geo_assistant/agent/graph.py | 3 +++ src/geo_assistant/agent/llms.py | 2 ++ src/geo_assistant/agent/state.py | 4 ++++ src/geo_assistant/api/app.py | 8 ++++++-- src/geo_assistant/api/schemas/chat.py | 6 ++++++ src/geo_assistant/frontend/app.py | 2 ++ src/geo_assistant/tools/__init__.py | 2 ++ src/geo_assistant/tools/buffer.py | 10 +++++++++- src/geo_assistant/tools/naip.py | 6 ++++-- src/geo_assistant/tools/overture.py | 25 ++++++++++++++++++++----- src/geo_assistant/tools/summarize.py | 15 +++++++++++---- tests/test_api.py | 5 ++++- tests/tools/test_buffer.py | 4 ++++ tests/tools/test_naip.py | 4 ++-- tests/tools/test_overture.py | 7 +++++++ tests/tools/test_summarize.py | 4 ++++ 17 files changed, 104 insertions(+), 17 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b360213..a2a0cc6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,8 +59,22 @@ asyncio_mode = "auto" [tool.ruff.lint] select = [ "COM", # flake8-commas + "D", # pydocstyle "F", # pyflakes "I", # isort "RUF", # ruff-specific "UP", # pyupgrade ] +extend-select = [ + "D213", # Summary lines should be positioned on the second physical line of the docstring. + "D410", # A blank line after section headings. +] +ignore = [ + "D205", # 1 blank line required between summary line and description + "D212", # Multi-line docstring summary should start at the first line +] + +[tool.ruff.lint.pydocstyle] +# See https://docs.astral.sh/ruff/faq/#does-ruff-support-numpy-or-google-style-docstrings +# for the enabled/disabled rules for the "google" convention. +convention = "google" diff --git a/src/geo_assistant/agent/graph.py b/src/geo_assistant/agent/graph.py index 4ef8de8..ab4687d 100644 --- a/src/geo_assistant/agent/graph.py +++ b/src/geo_assistant/agent/graph.py @@ -1,3 +1,5 @@ +"""Create agent graph that calls tools.""" + import datetime from langchain.agents import create_agent @@ -31,6 +33,7 @@ The current date and time is {now}. async def create_graph(): + """Create langchain agent graph with a list of tools.""" checkpointer = InMemorySaver() graph = create_agent( model=llm, diff --git a/src/geo_assistant/agent/llms.py b/src/geo_assistant/agent/llms.py index 7fc3cfb..71f56ac 100644 --- a/src/geo_assistant/agent/llms.py +++ b/src/geo_assistant/agent/llms.py @@ -1,3 +1,5 @@ +"""Ollama chat model.""" + import os from dotenv import load_dotenv diff --git a/src/geo_assistant/agent/state.py b/src/geo_assistant/agent/state.py index b727944..7d3c340 100644 --- a/src/geo_assistant/agent/state.py +++ b/src/geo_assistant/agent/state.py @@ -1,3 +1,5 @@ +"""State schema for the geo-assistant agent.""" + from typing import NotRequired from geojson_pydantic import Feature, FeatureCollection @@ -6,6 +8,8 @@ from pydantic import Field class GeoAssistantState(AgentState): + """Schema for the geo-assistant agent's state.""" + place: NotRequired[Feature | None] = None search_area: NotRequired[Feature | None] = None places_within_buffer: NotRequired[FeatureCollection | None] = None diff --git a/src/geo_assistant/api/app.py b/src/geo_assistant/api/app.py index f08a191..f37677b 100644 --- a/src/geo_assistant/api/app.py +++ b/src/geo_assistant/api/app.py @@ -1,3 +1,5 @@ +"""Chat app API endpoint.""" + import logging from collections.abc import AsyncGenerator from contextlib import aclosing, asynccontextmanager @@ -21,12 +23,12 @@ UI_SET_FIELDS_WHITELIST = ["point", "messages"] @asynccontextmanager -async def lifespan(app: FastAPI): +async def _lifespan(app: FastAPI): app.state.chatbot = await create_graph() yield -app = FastAPI(title="Geo Assistant", lifespan=lifespan) +app = FastAPI(title="Geo Assistant", lifespan=_lifespan) app.add_middleware( @@ -44,6 +46,7 @@ async def stream_chat( chatbot: Any, request: Request, ) -> AsyncGenerator[bytes]: + """Agent chat stream.""" config: dict[str, Any] = { "configurable": { "thread_id": str(thread_id), @@ -101,6 +104,7 @@ async def stream_chat( @app.post("/chat") async def chat(request: ChatRequestBody, http_request: Request) -> StreamingResponse: + """HTTP POST endpoint at /chat.""" generator = stream_chat( ui_state_update=request.agent_state_input, thread_id=request.thread_id, diff --git a/src/geo_assistant/api/schemas/chat.py b/src/geo_assistant/api/schemas/chat.py index 56927b6..1d1e16d 100644 --- a/src/geo_assistant/api/schemas/chat.py +++ b/src/geo_assistant/api/schemas/chat.py @@ -1,13 +1,19 @@ +"""Chat API schemas.""" + from pydantic import BaseModel from geo_assistant.agent.state import GeoAssistantState class ChatRequestBody(BaseModel): + """Schema for the request to the Chat API.""" + thread_id: str agent_state_input: GeoAssistantState class ChatResponse(BaseModel): + """Schema for the response from the Chat API.""" + thread_id: str state: GeoAssistantState diff --git a/src/geo_assistant/frontend/app.py b/src/geo_assistant/frontend/app.py index 5f188cc..f8a56e3 100644 --- a/src/geo_assistant/frontend/app.py +++ b/src/geo_assistant/frontend/app.py @@ -1,3 +1,5 @@ +"""Chat app frontend.""" + import base64 import json import os diff --git a/src/geo_assistant/tools/__init__.py b/src/geo_assistant/tools/__init__.py index c2233c1..b7cc7d0 100644 --- a/src/geo_assistant/tools/__init__.py +++ b/src/geo_assistant/tools/__init__.py @@ -1,3 +1,5 @@ +"""List of tools available to the agent.""" + from geo_assistant.tools.buffer import get_search_area from geo_assistant.tools.naip import fetch_naip_img from geo_assistant.tools.overture import get_place, get_places_within_buffer diff --git a/src/geo_assistant/tools/buffer.py b/src/geo_assistant/tools/buffer.py index 8331624..37633b0 100644 --- a/src/geo_assistant/tools/buffer.py +++ b/src/geo_assistant/tools/buffer.py @@ -1,3 +1,5 @@ +"""Tool to create a buffer polygon around a geometry feature.""" + from typing import Annotated import geopandas as gpd @@ -17,8 +19,14 @@ async def get_search_area( state: Annotated[GeoAssistantState, InjectedState], tool_call_id: Annotated[str, InjectedToolCallId] = "", ) -> Command: - """Get a search area buffer in km around the place defined in the agent state.""" + """ + Get a search area buffer in km around the place defined in the agent state. + Args: + buffer_size_km: Radius of the buffer in kilometres. + state: Pass in 'place' as state into this agent. + tool_call_id: Optional ID for tracking the tool call. + """ place_feature = state.get("place") if not place_feature: diff --git a/src/geo_assistant/tools/naip.py b/src/geo_assistant/tools/naip.py index 63969d4..e10f597 100644 --- a/src/geo_assistant/tools/naip.py +++ b/src/geo_assistant/tools/naip.py @@ -1,4 +1,5 @@ -# tools/naip_mpc_tools.py +"""Tool to query Planetary Computer STAC API for NAIP imagery.""" + import base64 from concurrent.futures import ThreadPoolExecutor from io import BytesIO @@ -39,7 +40,8 @@ async def fetch_naip_img( Args: start_date: Start date (YYYY-MM-DD). end_date: End date (YYYY-MM-DD). - + state: Pass in search_area as state into this agent. + tool_call_id: Optional ID for tracking the tool call """ if not state["search_area"]: return Command( diff --git a/src/geo_assistant/tools/overture.py b/src/geo_assistant/tools/overture.py index a77a673..a2b5ba0 100644 --- a/src/geo_assistant/tools/overture.py +++ b/src/geo_assistant/tools/overture.py @@ -1,3 +1,5 @@ +"""Tool to find closest matching Overture place based on user input.""" + import json import os from typing import Annotated @@ -21,7 +23,8 @@ load_dotenv() def create_database_connection(): - """Create and configure a DuckDB connection with necessary extensions. + """ + Create and configure a DuckDB connection with necessary extensions. Args: database_path: Path to the DuckDB database file @@ -43,8 +46,14 @@ async def get_place( place_name: str, tool_call_id: Annotated[str, InjectedToolCallId] = "", ) -> Command: - """Get place location from Overture Maps based on user input place name.""" + """ + Get place location from Overture Maps based on user input place name. + Args: + place_name: An address or location given as a human-readable string. + tool_call_id: Optional ID for tracking the tool call. + + """ db_connection = create_database_connection() source = os.getenv("OVERTURE_SOURCE", "local") if source == "s3": @@ -162,10 +171,16 @@ async def get_places_within_buffer( state: Annotated[GeoAssistantState, InjectedState], tool_call_id: Annotated[str, InjectedToolCallId], ) -> Command: - """Get places from Overture Maps within user specified area and user specified Overture place type. - - Accepts: restaurant(s), cafe(s), coffee shop(s), bar(s), pub(s) - case insensitive.""" + """ + Get places from Overture Maps within user specified area and user specified Overture + place type. + Args: + place: Overture place type. Accepts: restaurant(s), cafe(s), coffee shop(s), + bar(s), pub(s) - case insensitive. + state: Pass in 'search_area' as state into this agent. + tool_call_id: Optional ID for tracking the tool call. + """ # Normalize the place type place = normalize_place_type(place) diff --git a/src/geo_assistant/tools/summarize.py b/src/geo_assistant/tools/summarize.py index ad689cb..65b4e1c 100644 --- a/src/geo_assistant/tools/summarize.py +++ b/src/geo_assistant/tools/summarize.py @@ -17,7 +17,7 @@ dotenv.load_dotenv() class SatImgSummary(dspy.Signature): - "Describe things you see in the satellite image." + """Describe things you see in the satellite image.""" img: dspy.Image = dspy.InputField(desc="A satellite image") answer: str = dspy.OutputField(desc="Description of the image") @@ -33,7 +33,8 @@ class SatImgSummaryAgent(dspy.Module): temperature: float = 0.5, max_tokens: int = 4_096, ) -> None: - """Initialize the satellite image summary agent. + """ + Initialize the satellite image summary agent. Args: model: The Ollama model to use for summarization @@ -53,7 +54,8 @@ class SatImgSummaryAgent(dspy.Module): self.summarizer = dspy.Predict(SatImgSummary) def forward(self, img_url: str) -> dspy.Prediction: - """Generate a summary for the given image URL. + """ + Generate a summary for the given image URL. Args: img_url: URL of the image to summarize @@ -73,7 +75,12 @@ async def summarize_sat_img( state: Annotated[GeoAssistantState, InjectedState], tool_call_id: Annotated[str | None, InjectedToolCallId] = None, ) -> Command: - """Summarize the contents of a satellite image using an LLM. + """ + Summarize the contents of a satellite image using an LLM. + + Args: + state: Pass in 'naip_img_bytes' as state into this agent. + tool_call_id: Optional ID for tracking the tool call. Returns: Command containing the image summary and metadata diff --git a/tests/test_api.py b/tests/test_api.py index 796ed96..9088a44 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,3 +1,5 @@ +"""Tests for chat API endpoint.""" + from uuid import uuid4 import pytest @@ -10,7 +12,7 @@ from geo_assistant.api.app import app @pytest_asyncio.fixture async def initialized_app(): - """Initialize the app's chatbot before testing""" + """Initialize the app's chatbot before testing.""" # Manually initialize the chatbot as the lifespan would app.state.chatbot = await create_graph() yield app @@ -21,6 +23,7 @@ async def initialized_app(): @pytest.mark.xfail async def test_call_api(initialized_app): + """Test calling the API at the /chat HTTP POST endpoint.""" async with AsyncClient( transport=ASGITransport(app=initialized_app), base_url="http://test", diff --git a/tests/tools/test_buffer.py b/tests/tools/test_buffer.py index 44ab50d..ddac46a 100644 --- a/tests/tools/test_buffer.py +++ b/tests/tools/test_buffer.py @@ -1,3 +1,5 @@ +"""Tests for buffer tool.""" + from geojson_pydantic import Feature, Point from langchain_core.tools.base import ToolCall from pytest import fixture @@ -8,6 +10,7 @@ from geo_assistant.tools.buffer import get_search_area @fixture def geo_assistant_fixture(): + """Fixture with a GeoJSON point feature in a GeoAssistantState.""" place_geojson = Feature( type="Feature", geometry=Point(type="Point", coordinates=[-9.1393, 38.7223]), @@ -22,6 +25,7 @@ def geo_assistant_fixture(): async def test_get_search_area(geo_assistant_fixture): + """Ensure that `get_search_area` tool returns a buffer Polygon.""" # Call the underlying function directly to test the logic # This bypasses the injection framework which is better suited for integration tests command = await get_search_area.ainvoke( diff --git a/tests/tools/test_naip.py b/tests/tools/test_naip.py index 269769a..c3f6d55 100644 --- a/tests/tools/test_naip.py +++ b/tests/tools/test_naip.py @@ -1,3 +1,5 @@ +"""Tests for NAIP tool.""" + from types import NoneType import pytest @@ -19,7 +21,6 @@ async def test_fetch_naip(): - Internet access (to reach Planetary Computer STAC + blobs) - Planetary Computer / NAIP service to be up """ - # Union Market coordinates from GeoNames: 38.90789, -76.99831 # N 38°54'28" W 76°59'54" # We'll use a small neighborhood AOI around that point. @@ -60,7 +61,6 @@ async def test_fetch_naip_too_large(): - Internet access (to reach Planetary Computer STAC + blobs) - Planetary Computer / NAIP service to be up """ - # Union Market coordinates from GeoNames: 38.90789, -76.99831 # N 38°54'28" W 76°59'54" # We'll use a small neighborhood AOI around that point. diff --git a/tests/tools/test_overture.py b/tests/tools/test_overture.py index 6fae726..b821f58 100644 --- a/tests/tools/test_overture.py +++ b/tests/tools/test_overture.py @@ -1,3 +1,5 @@ +"""Tests for Overture tool.""" + import os import geopandas as gpd @@ -57,6 +59,7 @@ def geo_assistant_with_buffer_fixture(): async def test_get_place(): + """Ensure that `get_place` tool returns an Overture place given a place_name.""" command = await get_place.ainvoke( ToolCall( name="get_place", @@ -69,6 +72,10 @@ async def test_get_place(): async def test_get_places_within_buffer(geo_assistant_with_buffer_fixture): + """ + Ensure that `get_places_within_buffer` tool returns multiple Overture places that + fit match the category 'cafe' within a specific buffer area around a location. + """ command = await get_places_within_buffer.ainvoke( ToolCall( name="get_places_within_buffer", diff --git a/tests/tools/test_summarize.py b/tests/tools/test_summarize.py index 532ab4e..b4c2a7d 100644 --- a/tests/tools/test_summarize.py +++ b/tests/tools/test_summarize.py @@ -23,6 +23,10 @@ TEST_IMAGE_URL = "https://petapixel.com/assets/uploads/2022/08/French-Officials- ], ) async def test_summarize_sat_img(img_url, summary): + """ + Ensure that the `summarize_sat_img` tool can describe a satellite image in JPEG + format. + """ # Load the image from the supplied URL and encode it in base64 headers = { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",