mirror of
https://github.com/dataforcanada/d4c-service-geo-assistant.git
synced 2026-06-13 14:31:01 +02:00
Wei Ji
e3373026d6
* 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
51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
"""Tests for the satellite image summarization tool."""
|
|
|
|
import base64
|
|
import uuid
|
|
|
|
import pytest
|
|
import requests
|
|
from langchain_core.tools.base import ToolCall
|
|
|
|
from geo_assistant.agent.state import GeoAssistantState
|
|
from geo_assistant.tools.summarize import summarize_sat_img
|
|
|
|
# Sample test data
|
|
TEST_IMAGE_URL = "https://petapixel.com/assets/uploads/2022/08/French-Officials-Use-Satellite-Photos-and-AI-to-Spot-Unregistered-Pools-1536x806.jpg"
|
|
|
|
|
|
@pytest.mark.xfail
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"img_url,summary",
|
|
[
|
|
(TEST_IMAGE_URL, "building"),
|
|
],
|
|
)
|
|
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",
|
|
}
|
|
resp = requests.get(img_url, headers=headers)
|
|
resp.raise_for_status()
|
|
img_base64 = base64.b64encode(resp.content).decode("utf-8")
|
|
command = await summarize_sat_img.ainvoke(
|
|
ToolCall(
|
|
name="summarize_sat_img",
|
|
type="tool_call",
|
|
args={
|
|
"state": GeoAssistantState(naip_img_bytes=img_base64, messages=[]),
|
|
"tool_call_id": str(uuid.uuid4()),
|
|
},
|
|
id=str(uuid.uuid4()),
|
|
),
|
|
)
|
|
|
|
print(command.update.get("messages"))
|
|
assert summary in command.update.get("messages")[-1].content
|