135 lines
6.7 KiB
Python
135 lines
6.7 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import io
|
|
import json
|
|
from pathlib import Path
|
|
from tempfile import TemporaryDirectory
|
|
|
|
from PIL import Image
|
|
|
|
from app.services.image_observation import (
|
|
merge_image_observations,
|
|
normalize_image_observation,
|
|
normalize_sketch_candidates,
|
|
render_image_observation_context,
|
|
)
|
|
from app.services.image_processing import cv_hints, image_metadata
|
|
from app.services.agent_service import tools_for_model
|
|
from app.services.agent_service import AgentService
|
|
from app.services.attachments import attachment_record
|
|
from app.services.library import CdslLibrary
|
|
from app.services.storage import WorkspaceStore
|
|
from app.settings import ProviderConfig, ProviderModel, Settings
|
|
from app.models.contracts import ChatMessage, MessagePart
|
|
|
|
|
|
def test_observation_keeps_multiview_profiles_and_measurement_sources() -> None:
|
|
result = normalize_image_observation({
|
|
"part_type": "bent bracket",
|
|
"visible_features": ["plate", "irregular opening"],
|
|
"uncertain_features": ["inner bend radius"],
|
|
"views": [{"attachment_id": "upload_a", "view_role": "front", "confidence": 0.8}],
|
|
"profiles": [{
|
|
"id": "opening_01",
|
|
"role": "cutout",
|
|
"closed": True,
|
|
"source_images": ["upload_a"],
|
|
"segments": [
|
|
{"type": "line", "start": [0, 0], "end": [10, 0]},
|
|
{"type": "arc", "start": [10, 0], "end": [10, 4], "center": [8, 2], "radius_mm": 2},
|
|
{"type": "polyline", "points": [[10, 4], [5, 8], [0, 4]]},
|
|
],
|
|
}],
|
|
"measurements": [{"name": "plate_thickness", "value_mm": 1.2, "source": "user"}],
|
|
}, attachment_ids=["upload_a"])
|
|
|
|
assert result["schema_version"] == "cad.image-observation.v2"
|
|
assert result["profiles"][0]["segments"][1]["type"] == "arc"
|
|
assert result["measurements"][0]["source"] == "user"
|
|
assert "opening_01" in render_image_observation_context(result)
|
|
|
|
|
|
def test_sketch_merge_does_not_replace_user_measurement() -> None:
|
|
survey = normalize_image_observation({
|
|
"part_type": "bracket",
|
|
"visible_features": ["plate"],
|
|
"uncertain_features": [],
|
|
"views": [{"attachment_id": "upload_a"}],
|
|
"measurements": [{"name": "thickness", "value_mm": 1.2, "source": "user"}],
|
|
}, attachment_ids=["upload_a"])
|
|
sketches = normalize_sketch_candidates({
|
|
"profiles": [],
|
|
"measurements": [{"name": "thickness", "value_mm": 1.6, "source": "image"}],
|
|
"uncertainties": ["bend radius"],
|
|
}, attachment_ids=["upload_a"])
|
|
merged = merge_image_observations(survey, sketches)
|
|
|
|
assert merged["measurements"][0]["value_mm"] == 1.2
|
|
assert merged["uncertainties"] == ["bend radius"]
|
|
|
|
|
|
def test_image_metadata_and_cv_degrade_without_required_cv_support() -> None:
|
|
output = io.BytesIO()
|
|
Image.new("RGB", (32, 16), "white").save(output, format="PNG")
|
|
metadata = image_metadata(output.getvalue())
|
|
assert metadata["width"] == 32
|
|
assert metadata["height"] == 16
|
|
assert "available" in cv_hints(output.getvalue())
|
|
|
|
|
|
def test_image_tool_stages_expose_only_the_required_tool() -> None:
|
|
model = ProviderModel("vision", vision=True)
|
|
assert [tool["function"]["name"] for tool in tools_for_model(model, image_stage="survey")] == ["analyze_image_reference"]
|
|
assert [tool["function"]["name"] for tool in tools_for_model(model, image_stage="sketch", include_image_analysis=False, include_image_sketches=True)] == ["extract_image_sketch_candidates"]
|
|
|
|
|
|
def test_agent_runs_survey_then_sketch_stage_before_normal_tools() -> None:
|
|
class TwoStageAgent(AgentService):
|
|
def __init__(self, *args: object, **kwargs: object) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
self.required: list[str | None] = []
|
|
self.responses = [
|
|
{"choices": [{"message": {"role": "assistant", "content": "", "tool_calls": [{
|
|
"id": "survey", "type": "function", "function": {"name": "analyze_image_reference", "arguments": json.dumps({
|
|
"part_type": "bracket", "visible_features": ["plate"], "uncertain_features": [],
|
|
"views": [{"attachment_id": "upload_a", "view_role": "front"}], "profiles": [],
|
|
"measurements": [], "uncertainties": [],
|
|
})},
|
|
}]}}]},
|
|
{"choices": [{"message": {"role": "assistant", "content": "", "tool_calls": [{
|
|
"id": "sketch", "type": "function", "function": {"name": "extract_image_sketch_candidates", "arguments": json.dumps({
|
|
"profiles": [{"id": "opening", "role": "cutout", "closed": True, "segments": [{"type": "line", "start": [0, 0], "end": [2, 0]}]}],
|
|
"measurements": [], "uncertainties": [],
|
|
})},
|
|
}]}}]},
|
|
{"choices": [{"message": {"role": "assistant", "content": "继续建模。", "tool_calls": []}}]},
|
|
]
|
|
|
|
async def _complete(self, *args: object, **kwargs: object) -> dict[str, object]:
|
|
self.required.append(kwargs.get("required_tool_name") if "required_tool_name" in kwargs else args[4] if len(args) > 4 else None)
|
|
return self.responses.pop(0)
|
|
|
|
with TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
backend_root = Path(__file__).resolve().parents[1]
|
|
provider = ProviderConfig("test", "Test", "https://example.invalid/v1", "key", (ProviderModel("vision", vision=True),))
|
|
settings = Settings(root / "tasks", root / "conversations", backend_root / "cdsl_library", backend_root / "engine" / "cdsl_engine", "", "", "", 5, "test", (provider,))
|
|
store = WorkspaceStore(settings)
|
|
conversation_id = "conv_000000000001"
|
|
store.ensure_conversation(conversation_id)
|
|
relative, _ = store.write_conversation_upload(conversation_id, "part.png", b"png")
|
|
store.add_conversation_attachment(conversation_id, attachment_record(conversation_id, "part.png", "image/png", relative, b"png", "image"))
|
|
agent = TwoStageAgent(settings, store, CdslLibrary(settings))
|
|
message = ChatMessage(id="user", role="user", parts=[MessagePart(type="text", text="根据图片继续建模")])
|
|
|
|
async def collect() -> list[dict[str, object]]:
|
|
events = []
|
|
async for chunk in agent.stream([message], conversation_id, None):
|
|
events.append(json.loads(chunk.decode("utf-8").split("data: ", 1)[1]))
|
|
return events
|
|
|
|
events = asyncio.run(collect())
|
|
assert agent.required[:2] == ["analyze_image_reference", "extract_image_sketch_candidates"]
|
|
assert any(event.get("observationStage") == "complete" for event in events)
|