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)