first commit
This commit is contained in:
@@ -0,0 +1,211 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
import mimetypes
|
||||
import re
|
||||
from typing import Any, Dict, Iterable, Optional
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
CODE_FILE_PATTERN = re.compile(r"<\|code_file\|>([^<]*?)</\|code_file(?:\|)?>")
|
||||
OUTPUT_FILE_PATTERN = re.compile(r"<\|output_file\|>([^<]*?)</\|output_file(?:\|)?>")
|
||||
WORKSPACE_DIRNAME = "workspace"
|
||||
|
||||
MODEL_CONTENT_TYPES = {
|
||||
".stl": "model/stl",
|
||||
".obj": "text/plain; charset=utf-8",
|
||||
".ply": "application/octet-stream",
|
||||
".glb": "model/gltf-binary",
|
||||
".gltf": "model/gltf+json",
|
||||
".step": "model/step",
|
||||
".stp": "model/step",
|
||||
}
|
||||
|
||||
|
||||
def content_to_text(content: Any) -> str:
|
||||
if content is None:
|
||||
return ""
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
fragments: list[str] = []
|
||||
for item in content:
|
||||
if isinstance(item, dict):
|
||||
item_type = item.get("type")
|
||||
if item_type == "text":
|
||||
fragments.append(str(item.get("text", "")))
|
||||
elif item_type == "image_url":
|
||||
fragments.append("[Image]")
|
||||
else:
|
||||
fragments.append(str(item))
|
||||
else:
|
||||
fragments.append(str(item))
|
||||
return "\n".join(fragment for fragment in fragments if fragment).strip()
|
||||
return str(content)
|
||||
|
||||
|
||||
def resolve_project_path(
|
||||
path_text: str, project_root: Path = PROJECT_ROOT
|
||||
) -> Optional[Path]:
|
||||
candidate_text = path_text.strip()
|
||||
if not candidate_text:
|
||||
return None
|
||||
|
||||
candidate = Path(candidate_text).expanduser()
|
||||
if not candidate.is_absolute():
|
||||
candidate = project_root / candidate
|
||||
|
||||
try:
|
||||
resolved_project_root = project_root.resolve(strict=False)
|
||||
resolved = candidate.resolve(strict=False)
|
||||
resolved.relative_to(resolved_project_root)
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
|
||||
return resolved
|
||||
|
||||
|
||||
def resolve_artifact_path(
|
||||
path_text: str,
|
||||
project_root: Path = PROJECT_ROOT,
|
||||
reference_path: Optional[Path] = None,
|
||||
) -> Optional[Path]:
|
||||
candidate_text = path_text.strip()
|
||||
if not candidate_text:
|
||||
return None
|
||||
|
||||
candidate_paths: list[Path] = []
|
||||
direct_path = resolve_project_path(candidate_text, project_root=project_root)
|
||||
if direct_path is not None:
|
||||
candidate_paths.append(direct_path)
|
||||
|
||||
raw_candidate = Path(candidate_text).expanduser()
|
||||
if raw_candidate.is_absolute():
|
||||
return (
|
||||
direct_path if direct_path is not None and direct_path.is_file() else None
|
||||
)
|
||||
|
||||
normalized_relative = candidate_text.replace("\\", "/")
|
||||
while normalized_relative.startswith("./"):
|
||||
normalized_relative = normalized_relative[2:]
|
||||
|
||||
relative_fragments = [candidate_text]
|
||||
if normalized_relative and normalized_relative != candidate_text:
|
||||
relative_fragments.append(normalized_relative)
|
||||
|
||||
search_roots: list[Path] = []
|
||||
if reference_path is not None:
|
||||
search_roots.append(reference_path.parent)
|
||||
|
||||
workspace_root = project_root / WORKSPACE_DIRNAME
|
||||
if workspace_root.exists():
|
||||
search_roots.append(workspace_root)
|
||||
|
||||
try:
|
||||
seen: set[Path] = set(candidate_paths)
|
||||
for root in search_roots:
|
||||
for fragment in relative_fragments:
|
||||
resolved = resolve_project_path(
|
||||
str(root / fragment), project_root=project_root
|
||||
)
|
||||
if resolved is not None and resolved not in seen:
|
||||
candidate_paths.append(resolved)
|
||||
seen.add(resolved)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
for candidate_path in candidate_paths:
|
||||
if candidate_path.is_file():
|
||||
return candidate_path
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def remember_recent_path(paths: list[Path], next_path: Path) -> None:
|
||||
try:
|
||||
paths.remove(next_path)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
paths.append(next_path)
|
||||
|
||||
|
||||
def paths_by_recency(paths: list[Path]) -> list[Path]:
|
||||
return list(reversed(paths))
|
||||
|
||||
|
||||
def guess_content_type(path: Path) -> str:
|
||||
suffix = path.suffix.lower()
|
||||
explicit = MODEL_CONTENT_TYPES.get(suffix)
|
||||
if explicit:
|
||||
return explicit
|
||||
|
||||
guessed, _ = mimetypes.guess_type(path.name)
|
||||
return guessed or "application/octet-stream"
|
||||
|
||||
|
||||
def latest_stl_for_code_file(code_path: Path) -> Optional[Path]:
|
||||
if code_path.name != "model.py" or not code_path.exists():
|
||||
return None
|
||||
|
||||
stl_files = [path for path in code_path.parent.glob("*.stl") if path.is_file()]
|
||||
if not stl_files:
|
||||
return None
|
||||
|
||||
return max(stl_files, key=lambda item: item.stat().st_mtime)
|
||||
|
||||
|
||||
def extract_latest_artifacts(
|
||||
messages: Iterable[Dict[str, Any]], project_root: Path = PROJECT_ROOT
|
||||
) -> Dict[str, Any]:
|
||||
latest_code_path: Optional[Path] = None
|
||||
latest_model_path: Optional[Path] = None
|
||||
code_paths: list[Path] = []
|
||||
model_paths: list[Path] = []
|
||||
output_file_paths: list[Path] = []
|
||||
|
||||
for message in messages:
|
||||
if message.get("role") != "assistant":
|
||||
continue
|
||||
|
||||
text = content_to_text(message.get("content"))
|
||||
|
||||
for match in CODE_FILE_PATTERN.findall(text):
|
||||
resolved = resolve_artifact_path(match, project_root=project_root)
|
||||
if resolved is not None:
|
||||
latest_code_path = resolved
|
||||
remember_recent_path(code_paths, resolved)
|
||||
|
||||
for match in OUTPUT_FILE_PATTERN.findall(text):
|
||||
resolved = resolve_artifact_path(
|
||||
match,
|
||||
project_root=project_root,
|
||||
reference_path=latest_code_path,
|
||||
)
|
||||
if resolved is not None:
|
||||
remember_recent_path(output_file_paths, resolved)
|
||||
latest_model_path = resolved
|
||||
remember_recent_path(model_paths, resolved)
|
||||
|
||||
if latest_code_path is not None and latest_model_path is None:
|
||||
latest_model_path = latest_stl_for_code_file(latest_code_path)
|
||||
if latest_model_path is not None:
|
||||
remember_recent_path(model_paths, latest_model_path)
|
||||
|
||||
return {
|
||||
"code_path": latest_code_path,
|
||||
"code_paths": paths_by_recency(code_paths),
|
||||
"model_path": latest_model_path,
|
||||
"model_paths": paths_by_recency(model_paths),
|
||||
"output_paths": paths_by_recency(output_file_paths),
|
||||
}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"PROJECT_ROOT",
|
||||
"content_to_text",
|
||||
"extract_latest_artifacts",
|
||||
"guess_content_type",
|
||||
"resolve_artifact_path",
|
||||
"resolve_project_path",
|
||||
]
|
||||
Reference in New Issue
Block a user