from __future__ import annotations import asyncio import base64 import json import secrets from collections.abc import AsyncIterator from typing import Any import httpx from app.models.contracts import ChatMessage from app.services.engine_service import build_revision, load_engine from app.services.library import CdslLibrary from app.services.sse import event from app.services.storage import WorkspaceStore from app.settings import ProviderConfig, ProviderModel, Settings TOOL_SCHEMAS: list[dict[str, Any]] = [ { "type": "function", "function": { "name": "search_cdsl_library", "description": "Search the official local CDSL library for similar geometry and feature sequences.", "parameters": { "type": "object", "properties": {"query": {"type": "string"}, "limit": {"type": "integer", "minimum": 1, "maximum": 8}}, "required": ["query"], }, }, }, { "type": "function", "function": { "name": "read_cdsl_reference", "description": "Read one official CDSL sample by part_id. Use this before creating geometry based on a reference.", "parameters": {"type": "object", "properties": {"part_id": {"type": "string"}}, "required": ["part_id"]}, }, }, { "type": "function", "function": { "name": "read_current_cdsl", "description": "Read the current task's latest CDSL before making a natural-language revision.", "parameters": {"type": "object", "properties": {}, "additionalProperties": False}, }, }, { "type": "function", "function": { "name": "generate_cdsl_model", "description": "Validate and execute a complete parameterized CDSL model. Use only for explicit CAD generation or revision.", "parameters": { "type": "object", "properties": { "cdsl": {"type": "object", "description": "Complete cad.cdsl.llm.v1 JSON object."}, "summary": {"type": "string"}, "assumptions": {"type": "array", "items": {"type": "string"}}, }, "required": ["cdsl", "summary"], }, }, }, ] def text_from_message(message: ChatMessage) -> str: return "\n".join(part.text or "" for part in message.parts if part.type == "text").strip() def messages_for_model(messages: list[ChatMessage]) -> list[dict[str, Any]]: result: list[dict[str, Any]] = [] for message in messages[-20:]: text = text_from_message(message) if text: result.append({"role": message.role, "content": text}) return result def system_prompt(settings: Settings) -> str: skill_path = settings.engine_root.parent.parent / "agent" / "skills" / "cad-engine" / "SKILL.md" skill = skill_path.read_text(encoding="utf-8") if skill_path.is_file() else "" readme_path = settings.engine_root / "README.md" engine_readme = readme_path.read_text(encoding="utf-8") if readme_path.is_file() else "" supported_profiles = ", ".join(sorted(load_engine(settings).SHAPE_GENERATORS)) return f"""You are the CDSL CAD Agent for CDSL CAD Studio. You generate parameterized CDSL, never raw CAD source code. For new CAD requests: 1. Search the local official CDSL library. 2. Read at least one relevant reference when a match exists. 3. Call generate_cdsl_model only when the request is sufficiently specified. 4. Never claim success unless the tool returns a successful CDSL-only STEP and GLB artifact. For a revision, call read_current_cdsl first and preserve unrelated features. Do not output compiler_context, unknown_shape, complex_arc_shape, entities, contour_edges_mm, or contour_regions_mm. Use only self-contained named profiles supported by the engine. Ask a concise clarification question when essential dimensions or intent are missing. Ordinary explanations must not create CAD. Local skill: {skill} Local engine guide: {engine_readme} Supported named profile types: {supported_profiles} """ class AgentService: def __init__(self, settings: Settings, store: WorkspaceStore, library: CdslLibrary) -> None: self.settings = settings self.store = store self.library = library async def stream( self, messages: list[ChatMessage], conversation_id: str | None, selected_task_id: str | None, provider_id: str | None = None, model_id: str | None = None, ) -> AsyncIterator[bytes]: latest_user = next((message for message in reversed(messages) if message.role == "user"), None) if latest_user is None: yield event("cad_error", {"stage": "request", "message": "A user message is required."}) yield event("done", {}) return user_text = text_from_message(latest_user) conversation = self.store.ensure_conversation(conversation_id, selected_task_id) self.store.append_conversation_message(conversation["conversation_id"], latest_user.model_dump(), selected_task_id) task_id = selected_task_id or conversation.get("current_task_id") or "" assistant_parts: list[dict[str, Any]] = [] assistant_id = f"assistant_{secrets.token_hex(8)}" successful_result: dict[str, Any] | None = None error_payload: dict[str, Any] | None = None try: provider, model = self.settings.resolve_model(provider_id, model_id) except ValueError as error: provider = None model = None configuration_error = str(error) else: configuration_error = "" if not self.settings.llm_configured or provider is None or model is None: message = "Agent 尚未配置模型。请设置 CDSL_LLM_BASE_URL、CDSL_LLM_API_KEY 和 CDSL_LLM_MODEL。" if configuration_error: message = configuration_error error_payload = {"stage": "configuration", "message": message} assistant_parts.append({"type": "data-cad-error", "data": error_payload}) yield event("cad_error", error_payload) self._persist_assistant(conversation["conversation_id"], assistant_id, assistant_parts, task_id) yield event("done", {}) return try: attachment_message = self._attachment_message(conversation, model) except ValueError as error: error_payload = {"stage": "attachment", "message": str(error)} assistant_parts.append({"type": "data-cad-error", "data": error_payload}) yield event("cad_error", error_payload) self._persist_assistant(conversation["conversation_id"], assistant_id, assistant_parts, task_id) yield event("done", {}) return yield event("progress", {"step": "analyze_request", "label": "分析需求", "status": "running", "message": "正在整理当前会话和 CAD 需求。"}) references: list[str] = [] model_messages: list[dict[str, Any]] = [{"role": "system", "content": system_prompt(self.settings)}] model_messages.extend(messages_for_model(messages)) if attachment_message: model_messages.append({"role": "user", "content": attachment_message}) tools = TOOL_SCHEMAS try: for iteration in range(8): response = await self._complete(model_messages, tools, provider, model) choice = response["choices"][0]["message"] tool_calls = choice.get("tool_calls") or [] content = str(choice.get("content") or "") if content: assistant_parts.append({"type": "text", "text": content}) for chunk in self._chunks(content): yield event("text_delta", {"text": chunk}) if not tool_calls: break model_messages.append(choice) for call in tool_calls: name = str(call.get("function", {}).get("name") or "") arguments = json.loads(call.get("function", {}).get("arguments") or "{}") yield event("progress", { "step": name, "label": self._tool_label(name), "status": "running", "message": "Agent 正在调用本地 CAD 工具。", }) result, generated = await self._run_tool(name, arguments, task_id, user_text, references) if generated: task_id = generated["task_id"] model_messages.append({ "role": "tool", "tool_call_id": call.get("id", ""), "content": json.dumps(result, ensure_ascii=False), }) yield event("progress", { "step": name, "label": self._tool_label(name), "status": "success" if result.get("ok", True) else "error", "message": result.get("message") or result.get("summary") or "", }) if generated: result_payload = { "taskId": generated["task_id"], "revisionId": generated["revision_id"], "cdslPath": generated["cdsl_path"], "stepPath": generated["step_path"], "glbPath": generated["glb_path"], "reportPath": generated["report_path"], "parametersPath": generated.get("parameters_path"), "selectorPath": generated.get("selector_path"), "edgesPath": generated.get("edges_path"), "summary": generated["summary"], "referenceIds": generated["reference_ids"], "engine": generated["engine"], } successful_result = result_payload assistant_parts.append({"type": "data-cad-result", "data": result_payload}) yield event("cad_result", result_payload) if iteration == 7: error_payload = {"stage": "agent", "message": "Agent tool loop reached its safety limit."} assistant_parts.append({"type": "data-cad-error", "data": error_payload}) yield event("cad_error", error_payload) except Exception as error: error_payload = {"stage": "agent", "message": str(error)} assistant_parts.append({"type": "data-cad-error", "data": error_payload}) yield event("cad_error", error_payload) if task_id: self.store.ensure_conversation(conversation["conversation_id"], task_id) if not assistant_parts: assistant_parts.append({ "type": "text", "text": "我暂时没有生成可执行的 CAD 结果。请补充尺寸、形状或修改目标。", }) if successful_result and not any(part.get("type") == "text" for part in assistant_parts): assistant_parts.insert(0, {"type": "text", "text": f"已生成：{successful_result['summary']}。"}) self._persist_assistant(conversation["conversation_id"], assistant_id, assistant_parts, task_id) yield event("progress", {"step": "agent_stream", "label": "调用模型和工具", "status": "success", "message": "Agent 请求已完成。"}) yield event("done", {}) def _persist_assistant( self, conversation_id: str, assistant_id: str, parts: list[dict[str, Any]], task_id: str, ) -> None: self.store.append_conversation_message( conversation_id, { "id": assistant_id or f"assistant_{conversation_id}_{len(parts)}", "role": "assistant", "parts": parts, }, task_id or None, ) async def _complete( self, messages: list[dict[str, Any]], tools: list[dict[str, Any]], provider: ProviderConfig, model: ProviderModel, ) -> dict[str, Any]: url = f"{provider.base_url}/chat/completions" headers = {"Authorization": f"Bearer {provider.api_key}", "Content-Type": "application/json"} payload = {"model": model.id, "messages": messages, "tools": tools, "tool_choice": "auto", "temperature": 0.1} async with httpx.AsyncClient(timeout=self.settings.llm_timeout_s) as client: response = await client.post(url, headers=headers, json=payload) if response.status_code >= 400: raise RuntimeError(f"LLM request failed ({response.status_code}): {response.text[:800]}") return response.json() async def _run_tool(self, name: str, arguments: dict[str, Any], task_id: str, request: str, references: list[str]) -> tuple[dict[str, Any], dict[str, Any] | None]: if name == "search_cdsl_library": results = self.library.search(str(arguments.get("query") or request), int(arguments.get("limit") or 5)) return {"ok": True, "results": results}, None if name == "read_cdsl_reference": part_id = str(arguments.get("part_id") or "") sample = self.library.read_sample(part_id) if part_id not in references: references.append(part_id) return {"ok": True, "part_id": part_id, "cdsl": sample}, None if name == "read_current_cdsl": if not task_id: return {"ok": False, "message": "No current task exists. This is a new model request."}, None path = self.store.current_cdsl_path(task_id) if path is None: return {"ok": False, "message": "The current task has no successful CDSL revision."}, None return {"ok": True, "task_id": task_id, "cdsl": json.loads(path.read_text(encoding="utf-8"))}, None if name == "generate_cdsl_model": cdsl = arguments.get("cdsl") if isinstance(cdsl, str): cdsl = json.loads(cdsl) if not isinstance(cdsl, dict): raise ValueError("generate_cdsl_model requires a CDSL JSON object") summary = str(arguments.get("summary") or "CDSL CAD model") yieldable = await asyncio.to_thread( build_revision, settings=self.settings, store=self.store, task_id=task_id or None, request=request, cdsl=cdsl, reference_ids=list(references), summary=summary, ) return {"ok": True, "summary": summary, "task_id": yieldable["task_id"], "revision_id": yieldable["revision_id"]}, yieldable raise ValueError(f"Unknown agent tool: {name}") @staticmethod def _chunks(text: str) -> list[str]: return [text[index:index + 96] for index in range(0, len(text), 96)] @staticmethod def _tool_label(name: str) -> str: return { "search_cdsl_library": "检索 CDSL 模型库", "read_cdsl_reference": "读取 CDSL 参考模型", "read_current_cdsl": "读取当前 CDSL", "generate_cdsl_model": "生成 CDSL CAD 模型", }.get(name, "调用 CAD 工具") def _attachment_message(self, conversation: dict[str, Any], model: ProviderModel) -> list[dict[str, Any]] | str: attachments = conversation.get("attachments") or [] if not attachments: return "" content: list[dict[str, Any]] = [{"type": "text", "text": "The following local attachments are part of the CAD request."}] for attachment in attachments: if not isinstance(attachment, dict): continue kind = str(attachment.get("kind") or "") task_id = str(attachment.get("task_id") or "") relative = str(attachment.get("path") or "") if not task_id or not relative: continue path = self.store.artifact_path(task_id, relative) if kind == "image": if not model.vision: raise ValueError("The selected model does not support images. Choose a vision-capable OpenAI or Kimi model.") mime = str(attachment.get("mime") or "image/png") encoded = base64.b64encode(path.read_bytes()).decode("ascii") content.append({"type": "image_url", "image_url": {"url": f"data:{mime};base64,{encoded}"}}) elif kind == "document": extracted = str(attachment.get("extracted_path") or "") if extracted: text_path = self.store.artifact_path(task_id, extracted) text = text_path.read_text(encoding="utf-8")[:30_000] content.append({"type": "text", "text": f"Document {attachment.get('name')}:\n{text}"}) return content