From ff2956a41f5e6e9db16cf41598f1ca4bf1777eb5 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:35:21 +0800 Subject: [PATCH] refactor: simplify core, Studio, and unit test scaffolding (#658) * refactor: share backend search, file responses, and vector digests * refactor(studio): remove unused scaffolding and share workspace updates * test: simplify async tests and share deterministic store fixtures --- .github/workflows/_build-docs.yml | 2 +- .github/workflows/ci-reme-studio.yml | 2 +- example.env | 1 - .../file_store/faiss_local_file_store.py | 23 - .../components/file_store/local_file_store.py | 18 + .../file_store/zvec_local_file_store.py | 21 - reme/steps/file_io/_base.py | 13 + reme/steps/file_io/daily_write.py | 11 +- reme/steps/file_io/edit.py | 11 +- reme/steps/file_io/list.py | 12 +- reme/steps/file_io/load.py | 10 +- reme/steps/file_io/read_image.py | 11 +- reme/steps/file_io/save.py | 10 +- reme/steps/file_io/write.py | 11 +- reme/steps/index/_dedup.py | 29 +- reme/steps/index/_single_search.py | 68 + reme/steps/index/_tool_context.py | 23 + reme/steps/index/bm25_search.py | 65 +- reme/steps/index/draft.py | 11 +- reme/steps/index/search.py | 12 +- reme/steps/index/vector_search.py | 65 +- reme/utils/tushare.py | 17 - reme_studio/README.md | 6 +- reme_studio/README_ZH.md | 6 +- reme_studio/app/chatgpt-auth.ts | 90 - reme_studio/app/globals.css | 21 - reme_studio/app/open-file.ts | 21 +- reme_studio/app/store.ts | 331 +- reme_studio/app/workspace-chat.tsx | 296 ++ reme_studio/app/workspace-tabs.tsx | 136 + reme_studio/app/workspace-theme.tsx | 51 + reme_studio/app/workspace.tsx | 514 +-- reme_studio/cloudflare-env.d.ts | 5 - reme_studio/db/index.ts | 13 - reme_studio/db/schema.ts | 4 - reme_studio/drizzle.config.ts | 7 - reme_studio/drizzle/meta/_journal.json | 5 - .../examples/d1/app/api/notes/route.ts | 63 - reme_studio/examples/d1/db/schema.ts | 11 - reme_studio/package-lock.json | 1183 +------ reme_studio/package.json | 12 +- reme_studio/public/file.svg | 1 - reme_studio/public/globe.svg | 1 - reme_studio/public/window.svg | 1 - .../{ => distribution}/demo-build.test.mjs | 2 +- .../{ => distribution}/rendered-html.test.mjs | 4 +- .../{ => distribution}/static-build.test.mjs | 10 +- reme_studio/tooling/sites-vite-plugin.ts | 8 +- reme_studio/worker/index.ts | 1 - tests/unit/file_store_test_support.py | 79 + tests/unit/test_background_steps.py | 101 +- tests/unit/test_base_component.py | 28 - tests/unit/test_common_steps.py | 29 - tests/unit/test_component_registry.py | 17 - tests/unit/test_crud_steps.py | 124 +- tests/unit/test_daily_steps.py | 63 +- tests/unit/test_default_file_chunker.py | 33 - tests/unit/test_faiss_index_maintenance.py | 443 +-- tests/unit/test_file_catalog.py | 29 +- tests/unit/test_file_graph.py | 34 +- tests/unit/test_file_store_consistency.py | 3109 ++++++++--------- tests/unit/test_frontmatter_steps.py | 18 +- tests/unit/test_injected_job_kwargs.py | 4 - tests/unit/test_jsonl_file_chunker.py | 45 - tests/unit/test_keyword_index.py | 1543 ++++---- tests/unit/test_link_expansion.py | 37 +- tests/unit/test_local_embedding_store.py | 1508 ++++---- tests/unit/test_markdown_file_chunker.py | 68 +- tests/unit/test_path_permission.py | 22 +- tests/unit/test_prompt_handler.py | 36 - tests/unit/test_read_image_steps.py | 27 +- tests/unit/test_read_with_neighbors.py | 22 +- tests/unit/test_runtime_context.py | 22 - tests/unit/test_tokenizer.py | 23 - tests/unit/test_tushare.py | 21 - tests/unit/test_wikilink_utils.py | 54 +- tests/unit/test_write_metadata_lock.py | 23 +- tests/unit/test_zvec_file_store.py | 708 ++-- 78 files changed, 4171 insertions(+), 7348 deletions(-) create mode 100644 reme/steps/file_io/_base.py create mode 100644 reme/steps/index/_single_search.py create mode 100644 reme/steps/index/_tool_context.py delete mode 100644 reme/utils/tushare.py delete mode 100644 reme_studio/app/chatgpt-auth.ts create mode 100644 reme_studio/app/workspace-chat.tsx create mode 100644 reme_studio/app/workspace-tabs.tsx create mode 100644 reme_studio/app/workspace-theme.tsx delete mode 100644 reme_studio/cloudflare-env.d.ts delete mode 100644 reme_studio/db/index.ts delete mode 100644 reme_studio/db/schema.ts delete mode 100644 reme_studio/drizzle.config.ts delete mode 100644 reme_studio/drizzle/meta/_journal.json delete mode 100644 reme_studio/examples/d1/app/api/notes/route.ts delete mode 100644 reme_studio/examples/d1/db/schema.ts delete mode 100644 reme_studio/public/file.svg delete mode 100644 reme_studio/public/globe.svg delete mode 100644 reme_studio/public/window.svg rename reme_studio/tests/{ => distribution}/demo-build.test.mjs (93%) rename reme_studio/tests/{ => distribution}/rendered-html.test.mjs (88%) rename reme_studio/tests/{ => distribution}/static-build.test.mjs (74%) create mode 100644 tests/unit/file_store_test_support.py delete mode 100644 tests/unit/test_tushare.py diff --git a/.github/workflows/_build-docs.yml b/.github/workflows/_build-docs.yml index 1b3eb8f5..d2b8a577 100644 --- a/.github/workflows/_build-docs.yml +++ b/.github/workflows/_build-docs.yml @@ -67,7 +67,7 @@ jobs: - name: Verify browser demo bundle if: inputs.run_tests working-directory: reme_studio - run: node --test tests/demo-build.test.mjs + run: node --test tests/distribution/demo-build.test.mjs - name: Configure Pages if: inputs.upload_pages_artifact diff --git a/.github/workflows/ci-reme-studio.yml b/.github/workflows/ci-reme-studio.yml index fa8f5d5f..eb9ccf3f 100644 --- a/.github/workflows/ci-reme-studio.yml +++ b/.github/workflows/ci-reme-studio.yml @@ -66,7 +66,7 @@ jobs: - name: Verify npm package run: | - npm pack --pack-destination "${RUNNER_TEMP}" + npm pack --ignore-scripts --pack-destination "${RUNNER_TEMP}" tar -tzf "${RUNNER_TEMP}"/agentscope-ai-reme_studio-*.tgz | grep '^package/dist-static/index.html$' - name: Set up Python diff --git a/example.env b/example.env index 37c3aa76..b8e9e3d9 100644 --- a/example.env +++ b/example.env @@ -14,4 +14,3 @@ LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 # Optional data-source mirrors. Unset variables use the official services. # HF_MIRROR_URL=https://hf-mirror.com # ARXIV_MIRROR_URL=https://export.arxiv.org -# TUSHARE_MIRROR_URL=http://112.124.63.173:4000/tushare diff --git a/reme/components/file_store/faiss_local_file_store.py b/reme/components/file_store/faiss_local_file_store.py index 40dfd654..c78ae41d 100644 --- a/reme/components/file_store/faiss_local_file_store.py +++ b/reme/components/file_store/faiss_local_file_store.py @@ -1,7 +1,6 @@ """FAISS-backed file store: chunk JSONL stays authoritative; FAISS HNSW replaces the linear vector scan.""" import asyncio -import hashlib import json from collections.abc import Callable from contextlib import suppress @@ -389,28 +388,6 @@ class FaissLocalFileStore(LocalFileStore): # -- persistence ------------------------------------------------------ - def _chunks_embedding_digest(self) -> str: - """Order-independent digest of the live (chunk_id, embedding) set. - - Hashes float16 canonical bytes (the chunk JSONL serialization dtype) so - the value is identical whether an embedding sits fresh in memory - (possibly float32 from the provider; assignment bypasses the pydantic - validator) or has round-tripped through the JSONL. Written into the - idmap sidecar at dump time and recomputed from ``self.file_chunks`` at - load time: a mismatch means the sidecar vectors belong to a different - chunk generation than the authoritative JSONL, which the live-id set - check alone cannot detect for same-ID in-place updates. - """ - digest = hashlib.sha256() - eligible = sorted( - cid for cid, chunk in self.file_chunks.items() if self._embedding_dim_matches(chunk.embedding) - ) - for cid in eligible: - digest.update(cid.encode("utf-8")) - digest.update(b"\x00") - digest.update(np.asarray(self.file_chunks[cid].embedding, dtype=np.float16).tobytes()) - return digest.hexdigest() - async def load(self) -> None: """Load chunks via the parent, then attach FAISS state (sidecar or rebuild).""" await super().load() diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index dcbe6177..ee36674b 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -3,6 +3,7 @@ import asyncio import base64 import datetime +import hashlib import heapq import json import time @@ -107,6 +108,23 @@ class LocalFileStore(BaseFileStore): return 0 return max(0, int(getattr(self.embedding_store, "dimensions", 0))) + def _chunks_embedding_digest(self) -> str: + """Order-independent digest of the live (chunk_id, embedding) set. + + Float16 bytes match the authoritative chunk JSONL serialization even + before a fresh provider vector has round-tripped through persistence. + Vector backends use this digest to reject stale index sidecars. + """ + digest = hashlib.sha256() + eligible = sorted( + cid for cid, chunk in self.file_chunks.items() if self._embedding_dim_matches(chunk.embedding) + ) + for cid in eligible: + digest.update(cid.encode("utf-8")) + digest.update(b"\x00") + digest.update(np.asarray(self.file_chunks[cid].embedding, dtype=np.float16).tobytes()) + return digest.hexdigest() + # -- lifecycle ------------------------------------------------------------ async def _start(self) -> None: diff --git a/reme/components/file_store/zvec_local_file_store.py b/reme/components/file_store/zvec_local_file_store.py index 84db2dbd..8d81609d 100644 --- a/reme/components/file_store/zvec_local_file_store.py +++ b/reme/components/file_store/zvec_local_file_store.py @@ -1,6 +1,5 @@ """Zvec-backed file store: chunk JSONL stays authoritative; a zvec collection replaces the linear vector scan.""" -import hashlib import json import shutil import time @@ -227,26 +226,6 @@ class ZvecLocalFileStore(LocalFileStore): # -- persistence ------------------------------------------------------ - def _chunks_embedding_digest(self) -> str: - """Order-independent digest of the live (chunk_id, embedding) set. - - Hashes float16 canonical bytes (the chunk JSONL serialization dtype) so - the value is identical whether an embedding sits fresh in memory or has - round-tripped through the JSONL. Written into the sidecar at dump time - and recomputed from ``self.file_chunks`` at load time: a mismatch means - the collection belongs to a different chunk generation than the - authoritative JSONL. - """ - digest = hashlib.sha256() - eligible = sorted( - cid for cid, chunk in self.file_chunks.items() if self._embedding_dim_matches(chunk.embedding) - ) - for cid in eligible: - digest.update(cid.encode("utf-8")) - digest.update(b"\x00") - digest.update(np.asarray(self.file_chunks[cid].embedding, dtype=np.float16).tobytes()) - return digest.hexdigest() - async def load(self) -> None: """Load chunks via the parent, then attach the zvec collection (open or rebuild).""" await super().load() diff --git a/reme/steps/file_io/_base.py b/reme/steps/file_io/_base.py new file mode 100644 index 00000000..7c42bc81 --- /dev/null +++ b/reme/steps/file_io/_base.py @@ -0,0 +1,13 @@ +"""Common response handling for file operations.""" + +from ..base_step import BaseStep + + +class _FileIOStep(BaseStep): + """Preserve the file-operation failure envelope without changing its payload.""" + + def _fail(self, message: str, **meta) -> None: + assert self.context is not None + self.context.response.success = False + self.context.response.answer = f"Error: {message}" + self.context.response.metadata.update(meta) diff --git a/reme/steps/file_io/daily_write.py b/reme/steps/file_io/daily_write.py index 31ec050b..addff538 100644 --- a/reme/steps/file_io/daily_write.py +++ b/reme/steps/file_io/daily_write.py @@ -1,8 +1,8 @@ """Write a daily markdown note by dispatching the generic ``write`` step.""" +from ._base import _FileIOStep from ._daily_index import parse_daily_date, refresh_day_index, validate_session_id from ._path import validate_filename_component -from ..base_step import BaseStep from ..index import normalize_posix_path from ...components import R from ...steps.evolve import now @@ -11,16 +11,9 @@ _RESERVED_METADATA_KEYS = {"name", "description", "session_id", "source_conversa @R.register("daily_write_step") -class DailyWriteStep(BaseStep): +class DailyWriteStep(_FileIOStep): """Write ``daily//.md`` with conversation frontmatter.""" - def _fail(self, message: str, **meta) -> None: - assert self.context is not None - self.context.response.success = False - self.context.response.answer = f"Error: {message}" - if meta: - self.context.response.metadata.update(meta) - def _session_dir(self) -> str: return normalize_posix_path(str(self.config_value("session_dir"))) diff --git a/reme/steps/file_io/edit.py b/reme/steps/file_io/edit.py index 2e412d54..b9a457d5 100644 --- a/reme/steps/file_io/edit.py +++ b/reme/steps/file_io/edit.py @@ -3,14 +3,14 @@ import frontmatter import yaml +from ._base import _FileIOStep from ._file_io import get_path_lock, read_file_safe, write_file_safe from ._path import _check_path_permission, NON_MD_WARNING, gate_md, resolve_path -from ..base_step import BaseStep from ...components import R @R.register("edit_step") -class EditStep(BaseStep): +class EditStep(_FileIOStep): """Replace every occurrence of ``old`` with ``new`` inside the file body. The YAML front matter block (if any) is parsed out, kept verbatim and @@ -24,13 +24,6 @@ class EditStep(BaseStep): read-modify-write cycle against the same file (multi-worker / multi- process safety is out of scope).""" - def _fail(self, message: str, **meta) -> None: - assert self.context is not None - self.context.response.success = False - self.context.response.answer = f"Error: {message}" - if meta: - self.context.response.metadata.update(meta) - async def execute(self): # pylint: disable=too-many-return-statements assert self.context is not None raw = str(self.context.get("path") or "") diff --git a/reme/steps/file_io/list.py b/reme/steps/file_io/list.py index df217e81..7f62185e 100644 --- a/reme/steps/file_io/list.py +++ b/reme/steps/file_io/list.py @@ -19,8 +19,8 @@ from pathlib import Path from stat import S_ISREG from typing import Iterable +from ._base import _FileIOStep from ._path import resolve_path -from ..base_step import BaseStep from ...components import R # Default cap on returned items so huge workspaces don't blow up the response. @@ -28,17 +28,9 @@ DEFAULT_LIMIT = 100 @R.register("list_step") -class ListStep(BaseStep): +class ListStep(_FileIOStep): """Enumerate files under a directory in the workspace.""" - def _fail(self, message: str, **meta) -> None: - """Set a failed response (matches the read/edit/... fail envelope).""" - assert self.context is not None - self.context.response.success = False - self.context.response.answer = f"Error: {message}" - if meta: - self.context.response.metadata.update(meta) - def _collect_params(self) -> tuple[str, bool, int, str, frozenset[str]]: """Read ``path`` / ``recursive`` / ``limit`` from context; coerce permissively.""" assert self.context is not None diff --git a/reme/steps/file_io/load.py b/reme/steps/file_io/load.py index 94bc23fc..4403c72e 100644 --- a/reme/steps/file_io/load.py +++ b/reme/steps/file_io/load.py @@ -2,24 +2,18 @@ from datetime import datetime +from ._base import _FileIOStep from ._file_io import read_file_safe from ._path import _check_path_permission, resolve_path -from ..base_step import BaseStep from ...components import R DEFAULT_EDITOR_MAX_BYTES = 5 * 1024 * 1024 @R.register("load_step") -class LoadStep(BaseStep): +class LoadStep(_FileIOStep): """Return complete text without the truncation used by agent-facing reads.""" - def _fail(self, message: str, **metadata) -> None: - assert self.context is not None - self.context.response.success = False - self.context.response.answer = f"Error: {message}" - self.context.response.metadata.update(metadata) - async def execute(self): assert self.context is not None raw_path = str(self.context.get("path") or "") diff --git a/reme/steps/file_io/read_image.py b/reme/steps/file_io/read_image.py index 67342e0c..84666eec 100644 --- a/reme/steps/file_io/read_image.py +++ b/reme/steps/file_io/read_image.py @@ -5,14 +5,14 @@ import base64 import aiofiles import aiofiles.os +from ._base import _FileIOStep from ._path import NON_IMAGE_WARNING, gate_image, resolve_path -from ..base_step import BaseStep from ...components import R from ...constants import DEFAULT_MAX_IMAGE_BYTES @R.register("read_image_step") -class ReadImageStep(BaseStep): +class ReadImageStep(_FileIOStep): """Read an image file under ``workspace_dir`` and return base64 in ``answer``. Step-level attributes (``kwargs``, configured in yaml under ``steps:`` — @@ -32,13 +32,6 @@ class ReadImageStep(BaseStep): base64; ``metadata.non_image_warning=True`` flags it. """ - def _fail(self, message: str, **meta) -> None: - assert self.context is not None - self.context.response.success = False - self.context.response.answer = f"Error: {message}" - if meta: - self.context.response.metadata.update(meta) - async def execute(self): # pylint: disable=too-many-return-statements assert self.context is not None raw = str(self.context.get("path") or "") diff --git a/reme/steps/file_io/save.py b/reme/steps/file_io/save.py index 0b77cf7a..9ec9845f 100644 --- a/reme/steps/file_io/save.py +++ b/reme/steps/file_io/save.py @@ -2,14 +2,14 @@ from datetime import datetime +from ._base import _FileIOStep from ._file_io import detect_file_encoding, get_path_lock, write_file_safe from ._path import _check_path_permission, resolve_path -from ..base_step import BaseStep from ...components import R @R.register("save_step") -class SaveStep(BaseStep): +class SaveStep(_FileIOStep): """Persist editor content without interpreting or rebuilding frontmatter. ``expected_mtime`` is the ISO timestamp returned by ``stat``. When supplied, @@ -17,12 +17,6 @@ class SaveStep(BaseStep): filesystem authoritative when ReMe and another editor are used together. """ - def _fail(self, message: str, **metadata) -> None: - assert self.context is not None - self.context.response.success = False - self.context.response.answer = f"Error: {message}" - self.context.response.metadata.update(metadata) - async def execute(self): assert self.context is not None raw_path = str(self.context.get("path") or "") diff --git a/reme/steps/file_io/write.py b/reme/steps/file_io/write.py index 44d1425d..ff5fd081 100644 --- a/reme/steps/file_io/write.py +++ b/reme/steps/file_io/write.py @@ -2,14 +2,14 @@ import frontmatter +from ._base import _FileIOStep from ._file_io import detect_file_encoding, get_path_lock, write_file_safe from ._path import _check_path_permission, NON_MD_WARNING, gate_md, resolve_path -from ..base_step import BaseStep from ...components import R @R.register("write_step") -class WriteStep(BaseStep): +class WriteStep(_FileIOStep): """Write (create or overwrite) a markdown file. Frontmatter accepts two reserved string fields (``name`` / ``description``) @@ -26,13 +26,6 @@ class WriteStep(BaseStep): scope). """ - def _fail(self, message: str, **meta) -> None: - assert self.context is not None - self.context.response.success = False - self.context.response.answer = f"Error: {message}" - if meta: - self.context.response.metadata.update(meta) - async def execute(self): # pylint: disable=too-many-return-statements assert self.context is not None raw = str(self.context.get("path") or "") diff --git a/reme/steps/index/_dedup.py b/reme/steps/index/_dedup.py index 96929e7b..230aebfb 100644 --- a/reme/steps/index/_dedup.py +++ b/reme/steps/index/_dedup.py @@ -12,15 +12,13 @@ When ``app_context`` is absent the same structure is mirrored under """ import datetime -from typing import TYPE_CHECKING, Any, Callable, Final +from typing import Callable, Final +from ._tool_context import _ToolContextMixin from ...schema import FileChunk -if TYPE_CHECKING: - from ...components import ApplicationContext - -class _ToolContextDedupMixin: +class _ToolContextDedupMixin(_ToolContextMixin): """Mixin providing tool_context-scoped chunk dedup with TTL. Must be mixed into a ``BaseStep`` subclass (e.g. ``SearchStep``); it @@ -29,7 +27,6 @@ class _ToolContextDedupMixin: set by the host step's ``__init__`` (default 24h). """ - TOOL_CONTEXTS_KEY: Final[str] = "tool_contexts" SEARCH_SEEN_RANGES_KEY: Final[str] = "search_seen_chunk_ranges" def __init_subclass__(cls, **kwargs): @@ -53,25 +50,7 @@ class _ToolContextDedupMixin: ) return super().__new__(cls, *args, **kwargs) - if TYPE_CHECKING: - # Declared by BaseStep; repeated here so static analysis resolves - # attribute access on the mixin without inheriting BaseStep. - app_context: "ApplicationContext | None" - kwargs: dict[str, Any] - seen_ttl_hours: float - - def _tool_context_store(self, tool_context_id: str) -> dict: - """Return the mutable state bucket for a tool context. - - The bucket is created lazily on first access and lives at - ``metadata["tool_contexts"][tool_context_id]`` (or the same path under - ``kwargs`` when no ``app_context`` is available, e.g. in unit tests). - """ - if self.app_context is not None: - contexts = self.app_context.metadata.setdefault(self.TOOL_CONTEXTS_KEY, {}) - else: - contexts = self.kwargs.setdefault(self.TOOL_CONTEXTS_KEY, {}) - return contexts.setdefault(tool_context_id, {}) + seen_ttl_hours: float @staticmethod def _now_ts() -> float: diff --git a/reme/steps/index/_single_search.py b/reme/steps/index/_single_search.py new file mode 100644 index 00000000..06973f02 --- /dev/null +++ b/reme/steps/index/_single_search.py @@ -0,0 +1,68 @@ +"""Shared execution for single-backend search steps.""" + +from abc import abstractmethod +from typing import Final + +from ._dedup import _ToolContextDedupMixin +from ._source_format import ALL_RETURNED_MESSAGE, NO_RESULTS_MESSAGE, join_chunk_entries +from ._source_format import merge_session_chunk_intervals, render_chunk_entries +from ..base_step import BaseStep +from ...schema import FileChunk + +_MAX_CANDIDATES: Final = 200 +_CANDIDATE_MULTIPLIER: Final = 10 + + +class _SingleSearchStep(_ToolContextDedupMixin, BaseStep): + """Retrieve candidates, then apply the common single-backend search contract.""" + + def __init__(self, *args, seen_ttl_hours: float = 24, include_source: bool = True, **kwargs): + super().__init__(*args, **kwargs) + self.seen_ttl_hours = seen_ttl_hours + self.include_source = include_source + + async def execute(self): + assert self.context is not None + raw_query = self.context.get("query") + query: str = "" if raw_query is None else str(raw_query).strip() + limit: int = int(self.context.get("limit") or 5) + min_score: float = float(self.context.get("min_score") or 0.0) + tool_context_id: str = (self.context.get("tool_context_id", "") or "").strip() + + if not query: + self.context.response.success = False + self.context.response.answer = "Error: query cannot be empty" + return self.context.response + assert limit > 0, f"limit must be positive, got {limit}" + + candidates = min(_MAX_CANDIDATES, max(1, limit * _CANDIDATE_MULTIPLIER)) + results = await self._retrieve(query, candidates) + self.logger.info(f"[{self.name}] query={query!r} candidates={candidates} hits={len(results)}") + + if min_score > 0.0: + results = [chunk for chunk in results if chunk.score >= min_score] + + pre_dedup_count = 0 + if tool_context_id: + pre_dedup_count = len(results) + results, _ = self._dedupe_tool_context(results, tool_context_id, limit) + else: + results = results[:limit] + + session_dir = self.config_value("session_dir") + entries = render_chunk_entries( + merge_session_chunk_intervals(results, session_dir), + session_dir, + include_source=self.include_source, + ) + self.context.response.answer = join_chunk_entries(entries) + if not results: + self.context.response.answer = ALL_RETURNED_MESSAGE if pre_dedup_count > 0 else NO_RESULTS_MESSAGE + self.context.response.metadata["results"] = [ + c.model_dump(exclude_none=True, exclude={"embedding"}) for c in results + ] + return self.context.response + + @abstractmethod + async def _retrieve(self, query: str, candidates: int) -> list[FileChunk]: + """Fetch ranked candidates from the concrete backend.""" diff --git a/reme/steps/index/_tool_context.py b/reme/steps/index/_tool_context.py new file mode 100644 index 00000000..4b14a6d5 --- /dev/null +++ b/reme/steps/index/_tool_context.py @@ -0,0 +1,23 @@ +"""Shared access to application-lifetime tool context state.""" + +from typing import TYPE_CHECKING, Any, Final + +if TYPE_CHECKING: + from ...components import ApplicationContext + + +class _ToolContextMixin: + """Share storage access while leaving each step's state format independent.""" + + TOOL_CONTEXTS_KEY: Final[str] = "tool_contexts" + + if TYPE_CHECKING: + app_context: "ApplicationContext | None" + kwargs: dict[str, Any] + + def _tool_context_store(self, tool_context_id: str) -> dict: + if self.app_context is not None: + contexts = self.app_context.metadata.setdefault(self.TOOL_CONTEXTS_KEY, {}) + else: + contexts = self.kwargs.setdefault(self.TOOL_CONTEXTS_KEY, {}) + return contexts.setdefault(tool_context_id, {}) diff --git a/reme/steps/index/bm25_search.py b/reme/steps/index/bm25_search.py index 99ca8a4d..59bd2460 100644 --- a/reme/steps/index/bm25_search.py +++ b/reme/steps/index/bm25_search.py @@ -1,64 +1,13 @@ -"""``bm25_search_step`` — plain BM25 keyword search with tool_context dedup.""" +"""Plain BM25 keyword search with tool-context deduplication.""" -from typing import Final - -from ._dedup import _ToolContextDedupMixin -from ._source_format import ALL_RETURNED_MESSAGE, NO_RESULTS_MESSAGE, join_chunk_entries -from ._source_format import merge_session_chunk_intervals, render_chunk_entries -from ..base_step import BaseStep +from ._single_search import _SingleSearchStep from ...components import R - -_MAX_CANDIDATES: Final = 200 -_CANDIDATE_MULTIPLIER: Final = 10 +from ...schema import FileChunk @R.register("bm25_search_step") -class Bm25SearchStep(_ToolContextDedupMixin, BaseStep): - """BM25-only search: retrieve, filter by min_score, dedup by tool_context, truncate.""" +class Bm25SearchStep(_SingleSearchStep): + """Search the BM25 keyword backend using the shared result contract.""" - def __init__(self, *args, seen_ttl_hours: float = 24, include_source: bool = True, **kwargs): - super().__init__(*args, **kwargs) - self.seen_ttl_hours = seen_ttl_hours - self.include_source = include_source - - async def execute(self): - assert self.context is not None - raw_query = self.context.get("query") - query: str = "" if raw_query is None else str(raw_query).strip() - limit: int = int(self.context.get("limit") or 5) - min_score: float = float(self.context.get("min_score") or 0.0) - tool_context_id: str = (self.context.get("tool_context_id", "") or "").strip() - - if not query: - self.context.response.success = False - self.context.response.answer = "Error: query cannot be empty" - return self.context.response - assert limit > 0, f"limit must be positive, got {limit}" - - candidates = min(_MAX_CANDIDATES, max(1, limit * _CANDIDATE_MULTIPLIER)) - results = await self.file_store.keyword_search(query, candidates, {}) - self.logger.info(f"[{self.name}] query={query!r} candidates={candidates} hits={len(results)}") - - if min_score > 0.0: - results = [chunk for chunk in results if chunk.score >= min_score] - - pre_dedup_count = 0 - if tool_context_id: - pre_dedup_count = len(results) - results, _ = self._dedupe_tool_context(results, tool_context_id, limit) - else: - results = results[:limit] - - session_dir = self.config_value("session_dir") - entries = render_chunk_entries( - merge_session_chunk_intervals(results, session_dir), - session_dir, - include_source=self.include_source, - ) - self.context.response.answer = join_chunk_entries(entries) - if not results: - self.context.response.answer = ALL_RETURNED_MESSAGE if pre_dedup_count > 0 else NO_RESULTS_MESSAGE - self.context.response.metadata["results"] = [ - c.model_dump(exclude_none=True, exclude={"embedding"}) for c in results - ] - return self.context.response + async def _retrieve(self, query: str, candidates: int) -> list[FileChunk]: + return await self.file_store.keyword_search(query, candidates, {}) diff --git a/reme/steps/index/draft.py b/reme/steps/index/draft.py index a33d3db2..5684e329 100644 --- a/reme/steps/index/draft.py +++ b/reme/steps/index/draft.py @@ -2,24 +2,17 @@ from typing import Final +from ._tool_context import _ToolContextMixin from ..base_step import BaseStep from ...components import R @R.register("add_draft_step") -class AddDraftStep(BaseStep): +class AddDraftStep(_ToolContextMixin, BaseStep): """Append one draft text item to the current tool context.""" - TOOL_CONTEXTS_KEY: Final[str] = "tool_contexts" DRAFTS_KEY: Final[str] = "drafts" - def _tool_context_store(self, tool_context_id: str) -> dict: - if self.app_context is not None: - contexts = self.app_context.metadata.setdefault(self.TOOL_CONTEXTS_KEY, {}) - else: - contexts = self.kwargs.setdefault(self.TOOL_CONTEXTS_KEY, {}) - return contexts.setdefault(tool_context_id, {}) - async def execute(self): assert self.context is not None text = self.context.get("text", "") diff --git a/reme/steps/index/search.py b/reme/steps/index/search.py index 928587f5..77ff4a4c 100644 --- a/reme/steps/index/search.py +++ b/reme/steps/index/search.py @@ -5,6 +5,7 @@ import datetime import os from typing import Final +from ._tool_context import _ToolContextMixin from ..base_step import BaseStep from ..file_io import extract_daily_date from ...components import R @@ -35,10 +36,9 @@ def _filter_values(value: object) -> set: @R.register("search_step") -class SearchStep(BaseStep): +class SearchStep(_ToolContextMixin, BaseStep): """Hybrid search: run vector + keyword in parallel, fuse via RRF, filter, truncate.""" - TOOL_CONTEXTS_KEY: Final[str] = "tool_contexts" SEARCH_SEEN_KEY: Final[str] = "search_seen_chunk_ids" def __init__( @@ -98,14 +98,6 @@ class SearchStep(BaseStep): def _now_ts() -> float: return datetime.datetime.now().timestamp() - def _tool_context_store(self, tool_context_id: str) -> dict: - """Return the mutable state bucket for a tool context.""" - if self.app_context is not None: - contexts = self.app_context.metadata.setdefault(self.TOOL_CONTEXTS_KEY, {}) - else: - contexts = self.kwargs.setdefault(self.TOOL_CONTEXTS_KEY, {}) - return contexts.setdefault(tool_context_id, {}) - def _dedupe_tool_context( self, chunks: list[FileChunk], diff --git a/reme/steps/index/vector_search.py b/reme/steps/index/vector_search.py index 5de5aa94..858db2bb 100644 --- a/reme/steps/index/vector_search.py +++ b/reme/steps/index/vector_search.py @@ -1,64 +1,13 @@ -"""``vector_search_step`` — plain vector search with tool_context dedup.""" +"""Plain vector search with tool-context deduplication.""" -from typing import Final - -from ._dedup import _ToolContextDedupMixin -from ._source_format import ALL_RETURNED_MESSAGE, NO_RESULTS_MESSAGE, join_chunk_entries -from ._source_format import merge_session_chunk_intervals, render_chunk_entries -from ..base_step import BaseStep +from ._single_search import _SingleSearchStep from ...components import R - -_MAX_CANDIDATES: Final = 200 -_CANDIDATE_MULTIPLIER: Final = 10 +from ...schema import FileChunk @R.register("vector_search_step") -class VectorSearchStep(_ToolContextDedupMixin, BaseStep): - """Vector-only search: retrieve, filter by min_score, dedup by tool_context, truncate.""" +class VectorSearchStep(_SingleSearchStep): + """Search the vector backend using the shared result contract.""" - def __init__(self, *args, seen_ttl_hours: float = 24, include_source: bool = True, **kwargs): - super().__init__(*args, **kwargs) - self.seen_ttl_hours = seen_ttl_hours - self.include_source = include_source - - async def execute(self): - assert self.context is not None - raw_query = self.context.get("query") - query: str = "" if raw_query is None else str(raw_query).strip() - limit: int = int(self.context.get("limit") or 5) - min_score: float = float(self.context.get("min_score") or 0.0) - tool_context_id: str = (self.context.get("tool_context_id", "") or "").strip() - - if not query: - self.context.response.success = False - self.context.response.answer = "Error: query cannot be empty" - return self.context.response - assert limit > 0, f"limit must be positive, got {limit}" - - candidates = min(_MAX_CANDIDATES, max(1, limit * _CANDIDATE_MULTIPLIER)) - results = await self.file_store.vector_search(query, candidates, {}) - self.logger.info(f"[{self.name}] query={query!r} candidates={candidates} hits={len(results)}") - - if min_score > 0.0: - results = [chunk for chunk in results if chunk.score >= min_score] - - pre_dedup_count = 0 - if tool_context_id: - pre_dedup_count = len(results) - results, _ = self._dedupe_tool_context(results, tool_context_id, limit) - else: - results = results[:limit] - - session_dir = self.config_value("session_dir") - entries = render_chunk_entries( - merge_session_chunk_intervals(results, session_dir), - session_dir, - include_source=self.include_source, - ) - self.context.response.answer = join_chunk_entries(entries) - if not results: - self.context.response.answer = ALL_RETURNED_MESSAGE if pre_dedup_count > 0 else NO_RESULTS_MESSAGE - self.context.response.metadata["results"] = [ - c.model_dump(exclude_none=True, exclude={"embedding"}) for c in results - ] - return self.context.response + async def _retrieve(self, query: str, candidates: int) -> list[FileChunk]: + return await self.file_store.vector_search(query, candidates, {}) diff --git a/reme/utils/tushare.py b/reme/utils/tushare.py deleted file mode 100644 index 68218151..00000000 --- a/reme/utils/tushare.py +++ /dev/null @@ -1,17 +0,0 @@ -"""TuShare client construction with an optional mirror endpoint.""" - -import os - - -def create_tushare_api(token: str): - """Create a TuShare DataApi, using TUSHARE_MIRROR_URL when configured.""" - try: - import tushare as ts - except ImportError as exc: # pragma: no cover - optional core dependency. - raise RuntimeError("tushare is required for market-data research") from exc - - api = ts.pro_api(token) - api._DataApi__timeout = 600 # pylint: disable=protected-access - if mirror_url := os.getenv("TUSHARE_MIRROR_URL", "").strip(): - api._DataApi__http_url = mirror_url.rstrip("/") # pylint: disable=protected-access - return api diff --git a/reme_studio/README.md b/reme_studio/README.md index 8667cd91..2cbeea61 100644 --- a/reme_studio/README.md +++ b/reme_studio/README.md @@ -402,8 +402,10 @@ Run proportionate checks from `reme_studio/`: ```bash npm run format:check npm run lint -npm run build -npm run build:static +npm run test:unit +# Build and verify the Sites, static, and demo distributions: +npm run test:build +# Run both suites: npm test ``` diff --git a/reme_studio/README_ZH.md b/reme_studio/README_ZH.md index f3a11d59..3351d3a9 100644 --- a/reme_studio/README_ZH.md +++ b/reme_studio/README_ZH.md @@ -357,8 +357,10 @@ npm run build:demo ```bash npm run format:check npm run lint -npm run build -npm run build:static +npm run test:unit +# 构建并验证 Sites、静态和演示分发： +npm run test:build +# 运行以上全部检查： npm test ``` diff --git a/reme_studio/app/chatgpt-auth.ts b/reme_studio/app/chatgpt-auth.ts deleted file mode 100644 index a0ae2ed9..00000000 --- a/reme_studio/app/chatgpt-auth.ts +++ /dev/null @@ -1,90 +0,0 @@ -import { headers } from "next/headers"; -import { redirect } from "next/navigation"; - -export type ChatGPTUser = { - userId: string; - displayName: string; - email: string; - fullName: string | null; -}; - -const USER_ID_HEADER = "oai-authenticated-user-id"; -const USER_EMAIL_HEADER = "oai-authenticated-user-email"; -const USER_FULL_NAME_HEADER = "oai-authenticated-user-full-name"; -const USER_FULL_NAME_ENCODING_HEADER = - "oai-authenticated-user-full-name-encoding"; -const PERCENT_ENCODED_UTF8 = "percent-encoded-utf-8"; -const SIGN_IN_PATH = "/signin-with-chatgpt"; -const SIGN_OUT_PATH = "/signout-with-chatgpt"; -const CALLBACK_PATH = "/callback"; - -export async function getChatGPTUser(): Promise { - const requestHeaders = await headers(); - const userId = requestHeaders.get(USER_ID_HEADER); - const email = requestHeaders.get(USER_EMAIL_HEADER); - if (!userId || !email) return null; - - const encodedFullName = requestHeaders.get(USER_FULL_NAME_HEADER); - const fullName = - encodedFullName && - requestHeaders.get(USER_FULL_NAME_ENCODING_HEADER) === PERCENT_ENCODED_UTF8 - ? safeDecodeURIComponent(encodedFullName) - : null; - - return { - userId, - displayName: fullName ?? email, - email, - fullName, - }; -} - -export async function requireChatGPTUser( - returnTo: string, -): Promise { - const user = await getChatGPTUser(); - if (user) return user; - - redirect(chatGPTSignInPath(returnTo)); -} - -export function chatGPTSignInPath(returnTo: string): string { - const safeReturnTo = safeRelativeReturnPath(returnTo); - return `${SIGN_IN_PATH}?return_to=${encodeURIComponent(safeReturnTo)}`; -} - -export function chatGPTSignOutPath(returnTo = "/"): string { - const safeReturnTo = safeRelativeReturnPath(returnTo); - return `${SIGN_OUT_PATH}?return_to=${encodeURIComponent(safeReturnTo)}`; -} - -function safeRelativeReturnPath(value: string): string { - if (!value.startsWith("/") || value.startsWith("//")) return "/"; - - let url: URL; - try { - url = new URL(value, "https://app.local"); - } catch { - return "/"; - } - if (url.origin !== "https://app.local") return "/"; - if (isReservedAuthPath(url.pathname)) return "/"; - - return `${url.pathname}${url.search}${url.hash}`; -} - -function isReservedAuthPath(pathname: string): boolean { - return ( - pathname === SIGN_IN_PATH || - pathname === SIGN_OUT_PATH || - pathname === CALLBACK_PATH - ); -} - -function safeDecodeURIComponent(value: string): string | null { - try { - return decodeURIComponent(value); - } catch { - return null; - } -} diff --git a/reme_studio/app/globals.css b/reme_studio/app/globals.css index 2fc41164..2621b818 100644 --- a/reme_studio/app/globals.css +++ b/reme_studio/app/globals.css @@ -1610,27 +1610,6 @@ textarea:focus-visible { margin: 8px 0 10px; font-size: 12px; } -.component-row { - min-height: 48px; - display: flex; - align-items: center; - justify-content: space-between; - gap: 16px; - border-top: 1px solid var(--line); -} -.component-row > span { - display: flex; - flex-direction: column; - font-size: 12px; -} -.component-row small { - margin-top: 2px; - color: var(--muted); - font-size: 10px; -} -.component-row strong { - font-size: 12px; -} .index-card { padding: 22px; } diff --git a/reme_studio/app/open-file.ts b/reme_studio/app/open-file.ts index 41b965c0..9fe34dab 100644 --- a/reme_studio/app/open-file.ts +++ b/reme_studio/app/open-file.ts @@ -9,13 +9,24 @@ export async function openWorkspaceFile(path: string) { ); const id = store.openMarkdown(path); if (existing) return; + await loadWorkspaceFile(id, path); +} + +/** Load both newly opened and restored tabs, respecting the caller's lifetime. */ +export async function loadWorkspaceFile( + id: string, + path: string, + shouldApply: () => boolean = () => true, +) { + const store = useWorkspaceStore.getState(); try { const file = await readWorkspaceFile(path); - store.hydrateMarkdown(id, file.content, file.stat.mtime); + if (shouldApply()) store.hydrateMarkdown(id, file.content, file.stat.mtime); } catch (error) { - store.failMarkdown( - id, - error instanceof Error ? error.message : "Failed to read file", - ); + if (shouldApply()) + store.failMarkdown( + id, + error instanceof Error ? error.message : "Failed to read file", + ); } } diff --git a/reme_studio/app/store.ts b/reme_studio/app/store.ts index 5f2fb64f..98ec968d 100644 --- a/reme_studio/app/store.ts +++ b/reme_studio/app/store.ts @@ -54,6 +54,19 @@ interface WorkspaceState { finishChat: (tabId: string, messageId: string, error?: string) => void; } +function updateChatMessage( + tab: Extract, + messageId: string, + update: (message: ChatMessage) => ChatMessage, +) { + return { + ...tab, + messages: tab.messages.map((message) => + message.id === messageId ? update(message) : message, + ), + }; +} + const fileTitle = (path: string) => path.split("/").pop() || path; // Preserve the original key so existing tabs, chats, and unsaved drafts survive the Studio rename. const WORKSPACE_STORAGE_KEY = IS_DEMO @@ -62,189 +75,153 @@ const WORKSPACE_STORAGE_KEY = IS_DEMO export const useWorkspaceStore = create()( persist( - (set, get) => ({ - tabs: [], - activeTabId: undefined, - openMarkdown: (path) => { - const existing = get().tabs.find( - (tab) => tab.type === "markdown" && tab.path === path, - ); - if (existing) { - set({ activeTabId: existing.id }); - return existing.id; - } - const id = `file:${path}`; + (set, get) => { + const updateTab = ( + id: string, + type: Type, + update: (tab: Extract) => WorkspaceTab, + ) => set((state) => ({ - tabs: [ - ...state.tabs, - { - id, - type: "markdown", - title: fileTitle(path), - path, - content: "", - savedContent: "", - loading: true, - }, - ], - activeTabId: id, + tabs: state.tabs.map((tab) => + tab.id === id && tab.type === type + ? update(tab as Extract) + : tab, + ), })); - return id; - }, - openAgent: () => { - const id = `agent:${crypto.randomUUID()}`; - set((state) => ({ - tabs: [...state.tabs, { id, type: "agent", title: "", messages: [] }], - activeTabId: id, - })); - return id; - }, - openGraph: (root) => { - const id = `graph:${root}`; - const existing = get().tabs.find( - (tab) => tab.id === id && tab.type === "graph", - ); - if (existing) { - set({ activeTabId: id }); + return { + tabs: [], + activeTabId: undefined, + openMarkdown: (path) => { + const existing = get().tabs.find( + (tab) => tab.type === "markdown" && tab.path === path, + ); + if (existing) { + set({ activeTabId: existing.id }); + return existing.id; + } + const id = `file:${path}`; + set((state) => ({ + tabs: [ + ...state.tabs, + { + id, + type: "markdown", + title: fileTitle(path), + path, + content: "", + savedContent: "", + loading: true, + }, + ], + activeTabId: id, + })); return id; - } - set((state) => ({ - tabs: [ - ...state.tabs, - { id, type: "graph", title: `${root} graph`, root }, - ], - activeTabId: id, - })); - return id; - }, - closeTab: (id, discardUnsaved = false) => - set((state) => { - if ( - !discardUnsaved && - unsavedTabsClosedBy(state.tabs, id, false).length - ) - return state; - const index = state.tabs.findIndex((tab) => tab.id === id); - const tabs = state.tabs.filter((tab) => tab.id !== id); - const activeTabId = - state.activeTabId === id - ? tabs[Math.max(0, index - 1)]?.id - : state.activeTabId; - return { tabs, activeTabId }; - }), - closeOtherTabs: (id, discardUnsaved = false) => - set((state) => { - if ( - !discardUnsaved && - unsavedTabsClosedBy(state.tabs, id, true).length - ) - return state; - const tab = state.tabs.find((item) => item.id === id); - return tab ? { tabs: [tab], activeTabId: id } : state; - }), - setActiveTab: (activeTabId) => set({ activeTabId }), - hydrateMarkdown: (id, content, mtime) => - set((state) => ({ - tabs: state.tabs.map((tab) => - tab.id === id && tab.type === "markdown" - ? { - ...tab, - content, - savedContent: content, - mtime, - loading: false, - error: undefined, - } - : tab, - ), - })), - failMarkdown: (id, error) => - set((state) => ({ - tabs: state.tabs.map((tab) => - tab.id === id && tab.type === "markdown" - ? { ...tab, loading: false, error } - : tab, - ), - })), - updateMarkdown: (id, content) => - set((state) => ({ - tabs: state.tabs.map((tab) => - tab.id === id && tab.type === "markdown" - ? { ...tab, content } - : tab, - ), - })), - markSaved: (id, content, mtime) => - set((state) => ({ - tabs: state.tabs.map((tab) => - tab.id === id && tab.type === "markdown" - ? markMarkdownContentSaved(tab, content, mtime) - : tab, - ), - })), - addChatTurn: (tabId, user, assistant) => - set((state) => ({ - tabs: state.tabs.map((tab) => - tab.id === tabId && tab.type === "agent" - ? { - ...tab, - title: tab.messages.length - ? tab.title - : user.content.slice(0, 18), - messages: [...tab.messages, user, assistant], - streaming: true, - } - : tab, - ), - })), - applyChatChunk: (tabId, messageId, chunk) => - set((state) => ({ - tabs: state.tabs.map((tab) => { - if (tab.id !== tabId || tab.type !== "agent") return tab; - const messages = tab.messages.map((message) => - message.id === messageId - ? applyStreamChunk(message, chunk) - : message, - ); - return { - ...tab, - sessionId: chunk.session_id || tab.sessionId, - messages, - }; + }, + openAgent: () => { + const id = `agent:${crypto.randomUUID()}`; + set((state) => ({ + tabs: [ + ...state.tabs, + { id, type: "agent", title: "", messages: [] }, + ], + activeTabId: id, + })); + return id; + }, + openGraph: (root) => { + const id = `graph:${root}`; + const existing = get().tabs.find( + (tab) => tab.id === id && tab.type === "graph", + ); + if (existing) { + set({ activeTabId: id }); + return id; + } + set((state) => ({ + tabs: [ + ...state.tabs, + { id, type: "graph", title: `${root} graph`, root }, + ], + activeTabId: id, + })); + return id; + }, + closeTab: (id, discardUnsaved = false) => + set((state) => { + if ( + !discardUnsaved && + unsavedTabsClosedBy(state.tabs, id, false).length + ) + return state; + const index = state.tabs.findIndex((tab) => tab.id === id); + const tabs = state.tabs.filter((tab) => tab.id !== id); + const activeTabId = + state.activeTabId === id + ? tabs[Math.max(0, index - 1)]?.id + : state.activeTabId; + return { tabs, activeTabId }; }), - })), - toggleChatBlock: (tabId, messageId, blockId, expanded) => - set((state) => ({ - tabs: state.tabs.map((tab) => - tab.id === tabId && tab.type === "agent" - ? { - ...tab, - messages: tab.messages.map((message) => - message.id === messageId - ? toggleChatBlock(message, blockId, expanded) - : message, - ), - } - : tab, + closeOtherTabs: (id, discardUnsaved = false) => + set((state) => { + if ( + !discardUnsaved && + unsavedTabsClosedBy(state.tabs, id, true).length + ) + return state; + const tab = state.tabs.find((item) => item.id === id); + return tab ? { tabs: [tab], activeTabId: id } : state; + }), + setActiveTab: (activeTabId) => set({ activeTabId }), + hydrateMarkdown: (id, content, mtime) => + updateTab(id, "markdown", (tab) => ({ + ...tab, + content, + savedContent: content, + mtime, + loading: false, + error: undefined, + })), + failMarkdown: (id, error) => + updateTab(id, "markdown", (tab) => ({ + ...tab, + loading: false, + error, + })), + updateMarkdown: (id, content) => + updateTab(id, "markdown", (tab) => ({ ...tab, content })), + markSaved: (id, content, mtime) => + updateTab(id, "markdown", (tab) => + markMarkdownContentSaved(tab, content, mtime), ), - })), - finishChat: (tabId, messageId, error) => - set((state) => ({ - tabs: state.tabs.map((tab) => - tab.id === tabId && tab.type === "agent" - ? { - ...tab, - streaming: false, - messages: tab.messages.map((message) => - message.id === messageId - ? finishChatMessage(message, error) - : message, - ), - } - : tab, + addChatTurn: (tabId, user, assistant) => + updateTab(tabId, "agent", (tab) => ({ + ...tab, + title: tab.messages.length ? tab.title : user.content.slice(0, 18), + messages: [...tab.messages, user, assistant], + streaming: true, + })), + applyChatChunk: (tabId, messageId, chunk) => + updateTab(tabId, "agent", (tab) => ({ + ...updateChatMessage(tab, messageId, (message) => + applyStreamChunk(message, chunk), + ), + sessionId: chunk.session_id || tab.sessionId, + })), + toggleChatBlock: (tabId, messageId, blockId, expanded) => + updateTab(tabId, "agent", (tab) => + updateChatMessage(tab, messageId, (message) => + toggleChatBlock(message, blockId, expanded), + ), ), - })), - }), + finishChat: (tabId, messageId, error) => + updateTab(tabId, "agent", (tab) => ({ + ...updateChatMessage(tab, messageId, (message) => + finishChatMessage(message, error), + ), + streaming: false, + })), + }; + }, { name: WORKSPACE_STORAGE_KEY, version: 1, diff --git a/reme_studio/app/workspace-chat.tsx b/reme_studio/app/workspace-chat.tsx new file mode 100644 index 00000000..c1b8d349 --- /dev/null +++ b/reme_studio/app/workspace-chat.tsx @@ -0,0 +1,296 @@ +"use client"; + +import { useEffect, useRef, useState } from "react"; +import { Check, CircleAlert, LoaderCircle, Send } from "lucide-react"; +import WorkspaceMarkdown from "./files-workspace/WorkspaceMarkdown"; +import { streamChat } from "./api"; +import { chatStreamError, formatStreamPayloads } from "./chat-stream"; +import { useI18n } from "./i18n"; +import { useWorkspaceStore } from "./store"; +import { IS_DEMO } from "./studio-mode"; +import type { + ChatBlock, + ChatMessage, + DetailBlock, + WorkspaceTab, +} from "./types"; +import { + appendWorkspaceFileReference, + WORKSPACE_FILE_DRAG_TYPE, +} from "./workspace-drag"; +import { scenarios } from "../demo/scenarios"; + +function MarkdownView({ content }: { content: string }) { + return ( +
+ +
+ ); +} + +function ChunkStatus({ + active, + error = false, +}: { + active: boolean; + error?: boolean; +}) { + if (error) return ; + return active ? ( + + ) : ( + + ); +} + +function detailLabel( + block: DetailBlock, + t: ReturnType["t"], +): string { + if (block.type === "think") return t("thinking"); + if (block.type === "data") + return block.mediaType + ? `${t("dataChunk")} · ${block.mediaType}` + : t("dataChunk"); + if (block.type === "approval") return t("approvalChunk"); + if (block.type === "usage") return t("usageChunk"); + return t("unknownChunk", { type: block.sourceType }); +} + +function ChatBlockView({ + block, + onToggle, +}: { + block: ChatBlock; + onToggle: (blockId: string, expanded: boolean) => void; +}) { + const { t } = useI18n(); + if (block.type === "content") return ; + if (block.type === "error") + return ( +
+ + {block.text} +
+ ); + if (block.type === "usage") return null; + + const active = + block.status === "streaming" || + block.status === "calling" || + block.status === "running"; + const failed = block.status === "error"; + if (block.type === "tool") { + const call = formatStreamPayloads(block.callPayloads); + const result = formatStreamPayloads(block.resultPayloads); + return ( +
{ + if (event.currentTarget.open !== block.expanded) + onToggle(block.id, event.currentTarget.open); + }} + > + + + {block.name} + {active ? t("streaming") : t("completed")} + +
+ {call && ( +
+ {t("toolCall")} + 
{call}
+
+ )} + {result && ( +
+ {t("toolResult")} + 
{result}
+
+ )} +
+
+ ); + } + + const detail = formatStreamPayloads(block.payloads); + return ( +
{ + if (event.currentTarget.open !== block.expanded) + onToggle(block.id, event.currentTarget.open); + }} + > + + + {detailLabel(block, t)} + {active ? t("streaming") : t("completed")} + + {detail && ( +
+ 
{detail}
+
+ )} +
+ ); +} + +export default function Chat({ + tab, +}: { + tab: Extract; +}) { + const { language, t } = useI18n(); + const [input, setInput] = useState(""); + const [fileDragOver, setFileDragOver] = useState(false); + const endRef = useRef(null); + const controller = useRef(null); + const { addChatTurn, applyChatChunk, toggleChatBlock, finishChat } = + useWorkspaceStore(); + + useEffect(() => { + endRef.current?.scrollIntoView({ behavior: "smooth" }); + }, [tab.messages]); + useEffect(() => () => controller.current?.abort(), []); + + const suggestions = ( +
+ {(IS_DEMO + ? scenarios[language].map((scenario) => scenario.question) + : [t("promptRecent"), t("promptTasks"), t("promptIdeas")] + ).map((question) => ( + + ))} +
+ ); + + const send = async () => { + const query = input.trim(); + if (!query || tab.streaming) return; + setInput(""); + const user: ChatMessage = { + id: crypto.randomUUID(), + role: "user", + content: query, + }; + const assistant: ChatMessage = { + id: crypto.randomUUID(), + role: "assistant", + content: "", + blocks: [], + }; + addChatTurn(tab.id, user, assistant); + const requestController = new AbortController(); + controller.current = requestController; + const error = await chatStreamError( + () => + streamChat(query, tab.sessionId, requestController.signal, (chunk) => + applyChatChunk(tab.id, assistant.id, chunk), + ), + requestController.signal, + t("chatFailed"), + ); + finishChat(tab.id, assistant.id, error); + if (controller.current === requestController) controller.current = null; + }; + + return ( +
+
+ {!tab.messages.length && ( +
+ + )} + {tab.messages.map((message) => ( +
+
+ {message.role === "user" ? ( + t("you") + ) : ( + + )} +
+
+ {message.role === "user" &&

{message.content}

} + {message.role === "assistant" && + message.blocks?.map((block) => ( + + toggleChatBlock(tab.id, message.id, blockId, expanded) + } + /> + ))} + {message.role === "assistant" && + tab.streaming && + !message.blocks?.some((block) => block.type !== "usage") && ( + + )} +
+
+ ))} + {IS_DEMO && tab.messages.length > 0 && suggestions} +
+
+
{ + if (!event.dataTransfer.types.includes(WORKSPACE_FILE_DRAG_TYPE)) + return; + event.preventDefault(); + event.dataTransfer.dropEffect = "copy"; + setFileDragOver(true); + }} + onDragLeave={(event) => { + if (!event.currentTarget.contains(event.relatedTarget as Node | null)) + setFileDragOver(false); + }} + onDrop={(event) => { + const path = event.dataTransfer.getData(WORKSPACE_FILE_DRAG_TYPE); + if (!path) return; + event.preventDefault(); + setFileDragOver(false); + setInput((current) => appendWorkspaceFileReference(current, path)); + }} + > +
+