refactor(component): refactor base components and file store implementations

- Remove abstract methods from base component start/close
- Update BaseJob to remove name parameter and simplify initialization
- Change file modification time field from mtime_ms to modified_time in seconds
- Add type checking imports and improve typing annotations
- Implement LocalFileStore with JSONL persistence for file chunks
- Add MdFileParser with markdown and frontmatter support
- Simplify HttpClient call method with proper kwargs handling
- Remove unused ReMe class methods and create backup version
- Update StreamJob to use step_components instead of steps attribute
This commit is contained in:
jinli.yl 2026-04-23 17:44:25 +08:00
parent 42a3343cb5
commit ca5970c6b0
17 changed files with 201 additions and 292 deletions

View file

@ -37,11 +37,9 @@ class BaseComponent(ABC):
self._is_started: bool = False
self._lock: asyncio.Lock = asyncio.Lock()
@abstractmethod
async def _start(self) -> None:
"""Start the component."""
@abstractmethod
async def _close(self) -> None:
"""Close the component."""

View file

@ -22,5 +22,5 @@ class BaseClient(BaseComponent):
"""Close the client."""
@abstractmethod
async def __call__(self, action: str, **kwargs) -> dict:
async def __call__(self) -> dict:
"""Invoke an action with the given configuration."""

View file

@ -24,7 +24,7 @@ class HttpClient(BaseClient):
):
super().__init__(**kwargs)
if host and port:
if host is not None and port is not None:
pass
elif service_info := os.environ.get(REME_SERVICE_INFO):
try:
@ -49,7 +49,9 @@ class HttpClient(BaseClient):
timeout=self.timeout,
)
async def __call__(self, **_kwargs) -> dict:
async def __call__(self) -> dict:
if self.client is None:
await self._start()
response = await self.client.post(f"/{self.action}", json=self.kwargs)
response.raise_for_status()
return response.json()

View file

@ -42,7 +42,7 @@ class DefaultFileParser(BaseFileParser):
file_meta = FileMetadata(
hash=file_hash,
mtime_ms=stat.st_mtime * 1000,
modified_time=stat.st_mtime,
size=stat.st_size,
path=str(file_path.absolute()),
content=content,

View file

@ -1,6 +1,7 @@
"""Markdown file parser."""
import asyncio
import os
from pathlib import Path
import frontmatter
@ -28,13 +29,14 @@ class MdFileParser(BaseFileParser):
raw = file_path.read_text(encoding=self.encoding)
post = frontmatter.loads(raw)
stat = file_path.stat()
os.stat()
return stat, dict(post.metadata), post.content
stat, metadata, content = await asyncio.to_thread(_read_and_parse)
file_meta = FileMetadata(
hash=hash_text(content),
mtime_ms=stat.st_mtime * 1000,
modified_time=stat.st_mtime,
size=stat.st_size,
path=str(file_path.absolute()),
content=content,

View file

@ -129,7 +129,7 @@ class ChromaFileStore(BaseFileStore):
if file_meta.path:
self._metadata_cache[file_meta.path] = FileMetadata(
hash=file_meta.hash,
mtime_ms=file_meta.mtime_ms,
modified_time=file_meta.modified_time,
size=file_meta.size,
path=file_meta.path,
chunk_count=file_meta.chunk_count,

View file

@ -120,7 +120,7 @@ class LocalFileStore(BaseFileStore):
if file_meta.path:
self._files[file_meta.path] = FileMetadata(
hash=file_meta.hash,
mtime_ms=file_meta.mtime_ms,
modified_time=file_meta.modified_time,
size=file_meta.size,
path=file_meta.path,
chunk_count=file_meta.chunk_count,

View file

@ -168,7 +168,7 @@ class SqliteFileStore(BaseFileStore):
metadata = json.loads(meta_str) if meta_str else {}
result[path] = FileMetadata(
hash=hash_val,
mtime_ms=mtime,
modified_time=mtime,
size=size,
path=path,
metadata=metadata,
@ -195,7 +195,7 @@ class SqliteFileStore(BaseFileStore):
(
file_meta.path,
file_meta.hash,
file_meta.mtime_ms,
file_meta.modified_time,
file_meta.size,
json.dumps(file_meta.metadata, ensure_ascii=False) if file_meta.metadata else None,
len(chunks),
@ -317,7 +317,7 @@ class SqliteFileStore(BaseFileStore):
metadata = json.loads(meta_str) if meta_str else {}
return FileMetadata(
hash=hash_val,
mtime_ms=mtime,
modified_time=mtime,
size=size,
path=path,
chunk_count=chunk_count,

View file

@ -1,4 +1,5 @@
"""Base job component for sequential step execution."""
from typing import TYPE_CHECKING
from ..base_component import BaseComponent
from ..component_registry import R
@ -6,75 +7,49 @@ from ..runtime_context import RuntimeContext
from ...enumeration import ComponentEnum
from ...schema import Response, ComponentConfig
if TYPE_CHECKING:
from ..base_step import BaseStep
@R.register("base")
class BaseJob(BaseComponent):
"""Base job that executes a sequence of steps.
A job orchestrates multiple steps in sequence, passing a runtime context
through each step. Steps are configured via ComponentConfig and instantiated
lazily when the job starts.
"""
"""Job that executes steps sequentially."""
component_type = ComponentEnum.JOB
def __init__(
self,
name: str = "",
description: str = "",
parameters: dict | None = None,
steps: list[ComponentConfig] | None = None,
**kwargs,
):
"""Initialize the job.
Args:
name: Job name identifier.
description: Human-readable description.
parameters: Default parameters passed to steps.
steps: List of step configurations to execute.
**kwargs: Additional arguments passed to BaseComponent.
"""
super().__init__(**kwargs)
self.name: str = name
self.description: str = description
self.parameters: dict = parameters or {}
self.step_configs: list[ComponentConfig] = steps or []
self.steps: list = []
self.description = description
self.parameters = parameters or {}
self.step_configs = steps or []
self.step_components: list["BaseStep"] = []
async def _start(self) -> None:
"""Instantiate all configured steps."""
assert self.app_context is not None, "app_context must be provided"
for step_config in self.step_configs:
if not step_config.backend:
raise ValueError(f"{step_config.backend} backend is not specified.")
backend_cls = R.get(ComponentEnum.STEP, step_config.backend)
if not backend_cls:
raise ValueError(f"{step_config.backend} is not registered.")
step = backend_cls(
language=self.app_context.app_config.language,
**step_config.model_dump(exclude={"backend"}),
)
self.steps.append(step)
for config in self.step_configs:
if not config.backend:
raise ValueError(f"Step is missing the required 'backend' field")
step_cls = R.get(ComponentEnum.STEP, config.backend)
if not step_cls:
raise ValueError(
f"Step references an unregistered backend '{config.backend}' "
f"of type '{ComponentEnum.STEP}'",
)
params = config.model_dump()
params["app_context"] = self.app_context
self.step_components.append(step_cls(**params))
async def _close(self) -> None:
"""Clear all instantiated steps."""
self.steps.clear()
self.step_components.clear()
async def __call__(self, **kwargs) -> Response:
"""Execute all steps sequentially.
Args:
**kwargs: Parameters passed to the runtime context.
Returns:
The final response from the runtime context.
"""
context = RuntimeContext(**kwargs)
for step in self.steps:
await step(context)
for step_component in self.step_components:
await step_component(context)
return context.response

View file

@ -1,7 +1,5 @@
"""Streaming job for real-time output delivery."""
import asyncio
from .base_job import BaseJob
from ..component_registry import R
from ..runtime_context import RuntimeContext
@ -10,28 +8,13 @@ from ...enumeration import ChunkEnum
@R.register("stream")
class StreamJob(BaseJob):
"""Job that streams execution results in real-time.
"""Job that streams results to a queue in real-time."""
Unlike BaseJob which returns a final response, StreamJob pushes
intermediate results to a queue as they are produced, allowing
clients to receive updates incrementally.
"""
async def __call__(self, **kwargs) -> asyncio.Queue:
"""Execute all steps with streaming enabled.
Args:
**kwargs: Parameters passed to the runtime context.
Returns:
An asyncio.Queue containing streamed chunks.
"""
async def __call__(self, **kwargs):
context = RuntimeContext(stream=True, **kwargs)
try:
for step in self.steps:
for step in self.step_components:
await step(context)
except Exception as e:
await context.add_stream_string(str(e), ChunkEnum.ERROR)
await context.add_stream_done()
return context.stream_queue

View file

@ -12,39 +12,28 @@ if TYPE_CHECKING:
class BaseService(BaseComponent):
"""Abstract base class for services that expose jobs.
Services provide different ways to invoke jobs (HTTP, CLI, MCP, etc.).
Subclasses must implement add_job to register jobs with the service.
"""
"""Abstract base class for services that expose jobs (HTTP, MCP, etc.)."""
component_type = ComponentEnum.SERVICE
def __init__(self, **kwargs):
"""Initialize the service."""
super().__init__(**kwargs)
self.service = None
async def _start(self) -> None:
async def _close(self) -> None:
"""Default empty implementation for sync services."""
@abstractmethod
def build_service(self, app: "Application") -> None:
"""Build the service."""
...
@abstractmethod
def add_job(self, job: BaseJob) -> None:
"""Register a job with the service."""
...
@abstractmethod
def start_service(self, app: "Application") -> None:
"""Start the service."""
...
def add_jobs(self, app: "Application") -> None:
"""Register all jobs from the application context."""
for name, job in app.context.jobs.values():
for name, job in app.context.jobs.items():
try:
self.add_job(job)
self.logger.info(f"Successfully Added job {name}")
@ -52,11 +41,6 @@ class BaseService(BaseComponent):
self.logger.error(f"Failed to add job {name}: {e}")
def run_app(self, app: "Application") -> None:
"""Register all jobs from the application and start the service.
Args:
app: The application containing jobs to register.
"""
self.build_service(app)
self.add_jobs(app)
self.start_service(app)

View file

@ -25,34 +25,14 @@ if TYPE_CHECKING:
@R.register("http")
class HttpService(BaseService):
"""HTTP service that exposes jobs as REST endpoints.
Regular jobs return JSON responses, while StreamJobs return
server-sent events (SSE) for real-time streaming.
"""
# Removed the circular import - using string annotation instead
# from ...application import Application
"""HTTP service: jobs -> JSON endpoints, StreamJobs -> SSE endpoints."""
def __init__(self, host: str = REME_DEFAULT_HOST, port: int = REME_DEFAULT_PORT, **kwargs):
"""Initialize the HTTP service.
Args:
host: Bind address for the server.
port: Port number for the server.
**kwargs: Additional arguments passed to uvicorn.
"""
super().__init__(**kwargs)
self.host: str = host
self.port: int = port
def _add_job(self, job: BaseJob) -> None:
"""Register a regular job as a POST endpoint.
Args:
job: The job to register.
"""
async def execute_endpoint(request: Request) -> Response:
return await job(**request.model_dump(exclude_none=True))
@ -63,12 +43,6 @@ class HttpService(BaseService):
)(execute_endpoint)
def _add_stream_job(self, job: StreamJob) -> None:
"""Register a stream job as an SSE endpoint.
Args:
job: The stream job to register.
"""
async def execute_stream_endpoint(request: Request) -> StreamingResponse:
stream_queue = asyncio.Queue()
task = asyncio.create_task(
@ -90,41 +64,23 @@ class HttpService(BaseService):
self.service.post(f"/{job.name}")(execute_stream_endpoint)
def add_job(self, job: BaseJob) -> None:
"""Register a job with the HTTP service.
StreamJobs are registered as SSE endpoints, regular jobs as JSON endpoints.
Args:
job: The job to register.
"""
if isinstance(job, StreamJob):
self._add_stream_job(job)
else:
self._add_job(job)
def build_service(self, app: "Application") -> None:
"""Build the FastAPI application with CORS middleware.
Args:
app: The application instance.
"""
@asynccontextmanager
async def lifespan(_: FastAPI):
await app.start()
service_info = json.dumps(
{
"host": self.host,
"port": self.port,
},
)
service_info = json.dumps({"host": self.host, "port": self.port})
os.environ[REME_SERVICE_INFO] = service_info
self.logger.info(f"ReMe Service started: {REME_SERVICE_INFO}={service_info}")
yield
await app.close()
self.service = FastAPI(title=app.config.app_name, lifespan=lifespan)
self.service.add_middleware(
CORSMiddleware, # type: ignore[arg-type]
allow_origins=["*"],
@ -135,9 +91,4 @@ class HttpService(BaseService):
self.service.post("/health")(lambda: {"status": "healthy"})
def start_service(self, app: "Application") -> None:
"""Start the HTTP server.
Args:
app: The application instance.
"""
uvicorn.run(self.service, host=self.host, port=self.port, **self.kwargs)

View file

@ -1,130 +1,17 @@
"""ReMe CLI application entry point."""
import sys
from pathlib import Path
from agentscope.formatter import FormatterBase
from agentscope.message import Msg
from agentscope.model import ChatModelBase
from agentscope.token import HuggingFaceTokenCounter, TokenCounterBase
from agentscope.tool import Toolkit
from .application import Application
from .component import R, RuntimeContext
from .component import R
from .config import parse_args
from .enumeration import ComponentEnum
from .file_based.summarizer import Summarizer
from .utils import run_coro_safely
class ReMe(Application):
"""ReMe memory management application."""
async def summarize(
self,
messages: list[Msg],
as_llm: str | ChatModelBase = "default",
as_llm_formatter: str | FormatterBase = "default",
as_token_counter: str | TokenCounterBase | HuggingFaceTokenCounter = "default",
toolkit: Toolkit | None = None,
language: str = "zh",
max_input_length: float = 128 * 1024,
compact_ratio: float = 0.7,
timezone: str | None = None,
add_thinking_block: bool = True,
) -> str:
"""Summarize and compact memory messages.
Args:
messages: List of AgentScope messages to summarize.
as_llm: LLM model name or instance.
as_llm_formatter: Formatter name or instance.
as_token_counter: Token counter name or instance.
toolkit: Optional toolkit for the summarizer agent.
language: Language for prompts (zh or en).
max_input_length: Maximum input token length.
compact_ratio: Ratio of max_input_length to use as compact threshold.
timezone: Optional timezone for date formatting.
add_thinking_block: Whether to include thinking blocks.
Returns:
Summarized memory string.
"""
working_dir = Path(self.config.working_dir).absolute()
memory_dir = working_dir / "memory"
memory_compact_threshold = int(max_input_length * compact_ratio)
# Resolve token counter - use provided instance or create default
token_counter_instance = None
if isinstance(as_token_counter, HuggingFaceTokenCounter):
token_counter_instance = as_token_counter
else:
token_counter_instance = HuggingFaceTokenCounter()
summarizer = Summarizer(
working_dir=str(working_dir),
memory_dir=str(memory_dir),
memory_compact_threshold=memory_compact_threshold,
toolkit=toolkit,
timezone=timezone,
add_thinking_block=add_thinking_block,
as_token_counter=token_counter_instance,
language=language,
as_llm=as_llm if isinstance(as_llm, str) else "default",
as_llm_formatter=as_llm_formatter if isinstance(as_llm_formatter, str) else "default",
)
context = RuntimeContext(
messages=messages,
application_context=self.context,
)
result = await summarizer(context=context)
return result or ""
async def memory_search(self, query: str, max_results: int = 5, min_score: float = 0.1) -> str:
"""Search memory for relevant entries."""
from .file_based.memory_search import MemorySearch
try:
search_step = MemorySearch()
self.logger.info(f"Running memory search with {query} {max_results} {min_score}")
return await search_step(query=query, max_results=max_results, min_score=min_score)
except Exception as e:
return str(e)
async def dream(
self,
as_llm: str | ChatModelBase = "default",
as_llm_formatter: str | FormatterBase = "default",
as_token_counter: str | TokenCounterBase = "default",
toolkit: Toolkit | None = None,
language: str = "zh",
timezone: str | None = None,
) -> str:
"""Process and consolidate memories in background."""
return ""
async def proactive(
self,
as_llm: str | ChatModelBase = "default",
as_llm_formatter: str | FormatterBase = "default",
as_token_counter: str | TokenCounterBase = "default",
toolkit: Toolkit | None = None,
language: str = "zh",
timezone: str | None = None,
) -> str:
"""Generate proactive memory insights."""
return ""
class ReMeLight(ReMe):
"""ReMe memory management application."""
def __init__(self, **kwargs) -> None:
super().__init__(**kwargs)
self.context.app_config.service.backend = "http"
def main():
"""Entry point for ReMe CLI."""

144
reme2/reme_backup.py Normal file
View file

@ -0,0 +1,144 @@
"""ReMe CLI application entry point."""
import sys
from pathlib import Path
from agentscope.formatter import FormatterBase
from agentscope.message import Msg
from agentscope.model import ChatModelBase
from agentscope.token import HuggingFaceTokenCounter, TokenCounterBase
from agentscope.tool import Toolkit
from .application import Application
from .component import R, RuntimeContext
from .config import parse_args
from .enumeration import ComponentEnum
from .file_based.summarizer import Summarizer
from .utils import run_coro_safely
class ReMe(Application):
"""ReMe memory management application."""
async def summarize(
self,
messages: list[Msg],
as_llm: str | ChatModelBase = "default",
as_llm_formatter: str | FormatterBase = "default",
as_token_counter: str | TokenCounterBase | HuggingFaceTokenCounter = "default",
toolkit: Toolkit | None = None,
language: str = "zh",
max_input_length: float = 128 * 1024,
compact_ratio: float = 0.7,
timezone: str | None = None,
add_thinking_block: bool = True,
) -> str:
"""Summarize and compact memory messages.
Args:
messages: List of AgentScope messages to summarize.
as_llm: LLM model name or instance.
as_llm_formatter: Formatter name or instance.
as_token_counter: Token counter name or instance.
toolkit: Optional toolkit for the summarizer agent.
language: Language for prompts (zh or en).
max_input_length: Maximum input token length.
compact_ratio: Ratio of max_input_length to use as compact threshold.
timezone: Optional timezone for date formatting.
add_thinking_block: Whether to include thinking blocks.
Returns:
Summarized memory string.
"""
working_dir = Path(self.config.working_dir).absolute()
memory_dir = working_dir / "memory"
memory_compact_threshold = int(max_input_length * compact_ratio)
# Resolve token counter - use provided instance or create default
token_counter_instance = None
if isinstance(as_token_counter, HuggingFaceTokenCounter):
token_counter_instance = as_token_counter
else:
token_counter_instance = HuggingFaceTokenCounter()
summarizer = Summarizer(
working_dir=str(working_dir),
memory_dir=str(memory_dir),
memory_compact_threshold=memory_compact_threshold,
toolkit=toolkit,
timezone=timezone,
add_thinking_block=add_thinking_block,
as_token_counter=token_counter_instance,
language=language,
as_llm=as_llm if isinstance(as_llm, str) else "default",
as_llm_formatter=as_llm_formatter if isinstance(as_llm_formatter, str) else "default",
)
context = RuntimeContext(
messages=messages,
application_context=self.context,
)
result = await summarizer(context=context)
return result or ""
async def memory_search(self, query: str, max_results: int = 5, min_score: float = 0.1) -> str:
"""Search memory for relevant entries."""
from .file_based.memory_search import MemorySearch
try:
search_step = MemorySearch()
self.logger.info(f"Running memory search with {query} {max_results} {min_score}")
return await search_step(query=query, max_results=max_results, min_score=min_score)
except Exception as e:
return str(e)
async def dream(
self,
as_llm: str | ChatModelBase = "default",
as_llm_formatter: str | FormatterBase = "default",
as_token_counter: str | TokenCounterBase = "default",
toolkit: Toolkit | None = None,
language: str = "zh",
timezone: str | None = None,
) -> str:
"""Process and consolidate memories in background."""
return ""
async def proactive(
self,
as_llm: str | ChatModelBase = "default",
as_llm_formatter: str | FormatterBase = "default",
as_token_counter: str | TokenCounterBase = "default",
toolkit: Toolkit | None = None,
language: str = "zh",
timezone: str | None = None,
) -> str:
"""Generate proactive memory insights."""
return ""
class ReMeLight(ReMe):
"""ReMe memory management application."""
def __init__(self, **kwargs) -> None:
super().__init__(**kwargs)
self.context.app_config.service.backend = "http"
def main():
"""Entry point for ReMe CLI."""
action, config = parse_args(sys.argv[1:])
if action == "start":
reme = ReMe(**config)
reme.run_app()
else:
backend: str = config.pop("backend", "http")
client_cls = R.get(ComponentEnum.CLIENT, backend)
client = client_cls(action=action, **config)
run_coro_safely(client())
if __name__ == "__main__":
main()

View file

@ -1,29 +1,12 @@
"""Base node schema module.
This module defines the BaseNode model, which serves as the foundational
data structure for nodes in the knowledge graph or document processing pipeline.
"""
from uuid import uuid4
from pydantic import BaseModel, Field
class BaseNode(BaseModel):
"""Base node model for graph and document structures.
"""Base node model for graph and document structures."""
This model represents a single node in the knowledge graph or
a chunk in the document processing pipeline. It contains text content,
optional embeddings, and associated metadata.
Attributes:
id: Unique identifier for the node, auto-generated if not provided.
text: Text content of the node.
embedding: Optional vector embedding of the text content.
metadata: Additional metadata associated with the node.
"""
id: str = Field(default_factory=lambda: uuid4().hex, description="Unique node identifier")
text: str = Field(default="", description="Text content of the node")
embedding: list[float] | None = Field(default=None, description="Vector embedding of text")
metadata: dict = Field(default_factory=dict, description="Additional metadata")
id: str = Field(default_factory=lambda: uuid4().hex)
text: str = Field(default="")
embedding: list[float] | None = Field(default=None)
metadata: dict = Field(default_factory=dict)

View file

@ -4,7 +4,7 @@ from .base_node import BaseNode
class FileChunk(BaseNode):
"""文件内容分块,包含位置和评分元数据。"""
"""File content chunk with positional and scoring metadata."""
path: str = Field(...)
start_line: int = Field(...)

View file

@ -16,7 +16,7 @@ class FileMetadata(BaseModel):
Attributes:
hash: Hash of the file content for change detection.
mtime_ms: Last modification time in milliseconds since epoch.
modified_time: Last modification time in seconds since epoch.
size: File size in bytes.
path: Relative path to the file within the workspace.
content: Parsed content from the file (optional, memory-intensive).
@ -25,7 +25,7 @@ class FileMetadata(BaseModel):
"""
hash: str = Field(..., description="Hash of file content for change detection")
mtime_ms: float = Field(..., description="Last modification time in milliseconds")
modified_time: float = Field(..., description="Last modification time in seconds")
size: int = Field(..., description="File size in bytes")
path: str | None = Field(default=None, description="Relative path within workspace")
content: str | None = Field(default=None, description="Parsed content (optional)")