scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-09-17 23:51:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(components): add Anthropic LLM support and enhance component architecture

Browse source 
- Integrate AnthropicChatModel with new AnthropicAsLLM component
- Add component formatters for OpenAI and Anthropic chat models
- Implement token counter component with estimated token counting
- Create base client component for ReMe service communication
- Refactor BaseComponent to remove app_context parameter from _start/_close
- Update embedding model base class to remove retry logic and use npz cache
- Add job component for sequential step execution with BaseJob
- Implement step component base class for LLM workflow execution
- Enhance application context with proper type annotations
- Update component initialization to pass app_context automatically
- Remove asyncio dependency from embedding model cache operations
This commit is contained in:
jinli.yl 2026-04-23 14:18:57 +08:00
parent ec66fcf2ed
commit fa249a3e38
23 changed files with 216 additions and 352 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
reme/memory/vector_based/personal/personal_retriever.py
View file

@ -106,5 +106,5 @@ class PersonalRetriever(BaseMemoryAgent):
],
)
result["retrieved_nodes"] = self.retrieved_nodes
result["retrieved_nodes"] = self.retri eved_nodes
return result

3
reme/memory/vector_based/personal/personal_retriever.yaml
View file

@ -46,7 +46,8 @@ user_message: |
- Can read multiple histories at once by passing multiple history_ids
- Be selective: choose only the top 1-3 most promising histories
- Use this to understand the full conversation surrounding a memory
{
}
## Response Guidelines
- Base your answer EXCLUSIVELY on user profile, retrieved memories, and history data
- Never infer, assume, or hallucinate information

2
reme/memory/vector_based/personal/personal_summarizer.yaml
View file

@ -1,6 +1,6 @@
user_message_s1: |
You are a Memory Agent responsible for managing {memory_type} memories about {memory_target}.
{soul}
## Latest Conversation
Format: round<index> [<timestamp>] <role/name>: <content>
{context}

34
reme2/application.py
View file

@ -4,8 +4,8 @@ import asyncio
from pathlib import Path
from typing import AsyncGenerator
from .enumeration import ComponentEnum
from .component import BaseComponent, ApplicationContext
from .enumeration import ComponentEnum
from .schema import Response, StreamChunk
from .utils import execute_stream_task, print_logo, get_logger
@ -43,7 +43,9 @@ class Application(BaseComponent):
f"Service references an unregistered backend '{service_config.backend}' "
f"of type '{ComponentEnum.SERVICE}'",
)
self.context.service = service_cls(**service_config.model_dump(exclude={"backend"}))
params = service_config.model_dump()
params["app_context"] = self.context
self.context.service = service_cls(**params)
# Initialize all components grouped by type and name
for component_type, component_configs in self.config.components.items():
@ -57,7 +59,10 @@ class Application(BaseComponent):
f"Component '{name}' references an unregistered backend '{config.backend}' "
f"of type '{component_type}'",
)
self.context.components[component_type][name] = backend_cls(**config.model_dump(exclude={"backend"}))
params = config.model_dump()
params.setdefault("name", name)
params["app_context"] = self.context
self.context.components[component_type][name] = backend_cls(**params)
# Initialize all jobs
for job_config in self.config.jobs:
@ -70,25 +75,28 @@ class Application(BaseComponent):
f"Job '{job_config.name}' references an unregistered backend '{job_config.backend}' "
f"of type '{ComponentEnum.JOB}'",
)
self.context.jobs[job_config.name] = job_cls(**job_config.model_dump(exclude={"backend"}))
params = job_config.model_dump()
params.setdefault("name", job_config.name)
params["app_context"] = self.context
self.context.jobs[job_config.name] = job_cls(**params)
@property
def config(self):
"""Get application configuration."""
return self.context.app_config
async def _start(self, app_context=None) -> None:
async def _start(self) -> None:
"""Start the application."""
for components in self.context.components.values():
for component in components.values():
try:
await component.start(self.context)
await component.start()
except Exception as e:
self.logger.exception(f"Failed to start component {component.__class__.__name__}: {e}")
for name, job in self.context.jobs.items():
try:
await job.start(self.context)
await job.start()
except Exception as e:
self.logger.exception(f"Failed to start job '{name}': {e}")
@ -112,7 +120,7 @@ class Application(BaseComponent):
if name not in self.context.jobs:
raise KeyError(f"Job '{name}' not found")
job = self.context.jobs[name]
return await job(app_context=self.context, **kwargs)
return await job(**kwargs)
async def run_stream_job(self, name: str, **kwargs) -> AsyncGenerator[StreamChunk, None]:
"""Execute a streaming job and yield chunks."""
@ -120,12 +128,12 @@ class Application(BaseComponent):
raise KeyError(f"Job '{name}' not found")
job = self.context.jobs[name]
stream_queue = asyncio.Queue()
task = asyncio.create_task(job(stream_queue=stream_queue, app_context=self.context, **kwargs))
task = asyncio.create_task(job(stream_queue=stream_queue, **kwargs))
async for chunk in execute_stream_task(
stream_queue=stream_queue,
task=task,
task_name=name,
output_format="chunk",
stream_queue=stream_queue,
task=task,
task_name=name,
output_format="chunk",
):
assert isinstance(chunk, StreamChunk)
yield chunk

5
reme2/component/application_context.py
View file

@ -26,8 +26,9 @@ class ApplicationContext:
self.app_config: ApplicationConfig = ApplicationConfig(**kwargs)
from .base_component import BaseComponent
from .job.base_job import BaseJob
from .job import BaseJob
from .service import BaseService
self.service = None
self.service: BaseService | None = None
self.components: dict[ComponentEnum, dict[str, BaseComponent]] = {}
self.jobs: dict[str, BaseJob] = {}

36
reme2/component/as_llm/__init__.py
View file

@ -1,8 +1,6 @@
"""AgentScope LLM model wrappers."""
import asyncio
from agentscope.model import OpenAIChatModel, ChatModelBase
from agentscope.model import OpenAIChatModel, ChatModelBase, AnthropicChatModel
from ..base_component import BaseComponent
from ..component_registry import R
@ -22,8 +20,8 @@ class BaseAsLLM(BaseComponent):
super().__init__(**kwargs)
self.model: ChatModelBase | None = None
async def _start(self, app_context=None) -> None:
"""Initialize the AgentScope model. Override in subclasses."""
async def _start(self) -> None:
"""Initialize the model."""
async def _close(self) -> None:
"""Release model resources."""
@ -34,24 +32,34 @@ class BaseAsLLM(BaseComponent):
class OpenAIAsLLM(BaseAsLLM):
"""OpenAI chat model wrapper."""
async def _start(self, app_context=None) -> None:
async def _start(self) -> None:
"""Initialize the OpenAI chat model."""
self.model = OpenAIChatModel(**self.kwargs)
async def _close(self) -> None:
"""Close the HTTP client and release resources."""
if self.model is not None:
client = getattr(self.model, "client", None)
if client is not None and hasattr(client, "close"):
close_method = client.close
if asyncio.iscoroutinefunction(close_method):
await close_method()
else:
close_method()
self.model = None
assert isinstance(self.model, OpenAIChatModel)
await self.model.client.close()
@R.register("anthropic")
class AnthropicAsLLM(BaseAsLLM):
"""Anthropic chat model wrapper."""
async def _start(self) -> None:
"""Initialize the Anthropic chat model."""
self.model = AnthropicChatModel(**self.kwargs)
async def _close(self) -> None:
"""Close the HTTP client and release resources."""
if self.model is not None:
assert isinstance(self.model, AnthropicChatModel)
await self.model.client.close()
__all__ = [
"BaseAsLLM",
"OpenAIAsLLM",
"AnthropicAsLLM",
]

16
reme2/component/as_llm_formatter/__init__.py
View file

@ -1,6 +1,6 @@
"""Module for AgentScope LLM formatter components."""
from agentscope.formatter import FormatterBase
from agentscope.formatter import FormatterBase, AnthropicChatFormatter
from .reme_openai_chat_formatter import ReMeOpenAIChatFormatter
from ..base_component import BaseComponent
@ -17,8 +17,7 @@ class BaseAsLLMFormatter(BaseComponent):
super().__init__(**kwargs)
self.formatter: FormatterBase | None = None
async def _start(self, app_context=None) -> None:
"""Initialize the formatter instance."""
async def _start(self) -> None:
async def _close(self) -> None:
self.formatter = None
@ -28,11 +27,20 @@ class BaseAsLLMFormatter(BaseComponent):
class AsOpenAIChatFormatter(BaseAsLLMFormatter):
"""Wrapper for OpenAI chat completion formatter."""
async def _start(self, app_context=None) -> None:
async def _start(self) -> None:
self.formatter = ReMeOpenAIChatFormatter(**self.kwargs)
@R.register("anthropic")
class AsAnthropicChatFormatter(BaseAsLLMFormatter):
"""Wrapper for Anthropic chat completion formatter."""
async def _start(self) -> None:
self.formatter = AnthropicChatFormatter(**self.kwargs)
__all__ = [
"BaseAsLLMFormatter",
"AsOpenAIChatFormatter",
"AsAnthropicChatFormatter",
]

5
reme2/component/as_token_counter/__init__.py
View file

@ -21,8 +21,7 @@ class BaseAsTokenCounter(BaseComponent):
super().__init__(**kwargs)
self.token_counter: TokenCounterBase | None = None
async def _start(self, app_context=None) -> None:
"""Initialize the token counter. Override in subclasses."""
async def _start(self) -> None:
async def _close(self) -> None:
"""Release token counter resources."""
@ -33,7 +32,7 @@ class BaseAsTokenCounter(BaseComponent):
class EstimatedAsTokenCounter(BaseAsTokenCounter):
"""Estimated token counter using character-based estimation."""
async def _start(self, app_context=None) -> None:
async def _start(self) -> None:
"""Initialize the estimated token counter."""
self.token_counter = EstimatedTokenCounter(**self.kwargs)

2
reme2/component/as_token_counter/estimate_token_counter.py
View file

@ -20,8 +20,6 @@ class EstimatedTokenCounter(TokenCounterBase):
Args:
estimate_divisor: The divisor for character-to-token estimation.
Default 4 assumes roughly 4 characters per token.
Use 2-3 for Chinese/Japanese text, 4-5 for English.
encoding: The character encoding to use for byte length calculation.
"""
if estimate_divisor <= 0:

181
reme2/component/base_component.py
View file

@ -1,132 +1,54 @@
"""Base class for components."""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
from ..enumeration import ComponentEnum
from ..utils.logger_utils import get_logger
if TYPE_CHECKING:
from .application_context import ApplicationContext
class BaseComponent(ABC):
"""Base class for all application components.
"""Async lifecycle base class with context manager support.
Provides an asynchronous lifecycle with start/close operations and
async context manager support. State tracking prevents duplicate
start or close calls.
Subclasses must implement ``_start`` and ``_close`` to define their
specific initialization and teardown logic.
Examples:
Direct usage::
comp = MyComponent()
await comp.start()
# ... use component ...
await comp.close()
Context manager usage::
async with MyComponent() as comp:
# ... use component ...
Attributes:
component_type: The type identifier for this component, used during
registry lookup. Defaults to ``ComponentEnum.BASE``.
_is_started: Internal flag indicating whether the component has been
started and not yet closed.
Subclasses must implement ``_start`` and ``_close``.
"""
from .application_context import ApplicationContext
component_type = ComponentEnum.BASE
def __init__(self, **kwargs) -> None:
"""Initialize a component instance.
Sets up the component's internal state, binds a structured logger
with the component's class name, and stores any additional keyword
arguments for downstream use by subclasses.
Args:
**kwargs: Arbitrary keyword arguments forwarded to the component
subclass. Typically provided by the registry when the
component is instantiated from configuration.
"""
def __init__(
self,
name: str | None = None,
backend: str | None = None,
app_context: "ApplicationContext | None" = None,
**kwargs,
) -> None:
self.name: str = name or self.__class__.__name__
self.backend: str | None = backend
self.app_context: "ApplicationContext | None" = app_context
self.kwargs: dict = dict(kwargs)
self.logger = get_logger()
if hasattr(self.logger, "bind"):
self.logger = self.logger.bind(component=self.__class__.__name__)
self.logger = self.logger.bind(component=self.name)
self._is_started: bool = False
@abstractmethod
async def _start(self, app_context: ApplicationContext | None = None) -> None:
"""Perform the actual initialization logic for this component.
Subclasses must implement this method to set up resources such as
connections, caches, or background tasks. This method is called
internally by ``start()`` after verifying the component is not
already started.
Args:
app_context: The shared application context that provides access
to other initialized components and the application service.
May be ``None`` if the component does not require cross-component
references.
Raises:
ValueError: If required configuration or dependencies are missing
or invalid.
Exception: Any exception raised during resource acquisition will
propagate to the caller of ``start()``.
"""
async def _start(self) -> None: ...
@abstractmethod
async def _close(self) -> None:
"""Perform the actual teardown logic for this component.
async def _close(self) -> None: ...
Subclasses must implement this method to release resources such as
closing connections, flushing buffers, or cancelling background tasks.
This method is called internally by ``close()`` after verifying the
component is in a started state.
Raises:
ValueError: If the component is in an unexpected state during shutdown.
Exception: Any exception raised during resource cleanup will
propagate to the caller of ``close()``.
"""
async def start(self, app_context: ApplicationContext | None = None) -> None:
"""Start the component and transition it to an active state.
This is the public entry point for component initialization. It guards
against duplicate starts by returning immediately if the component is
already running, then delegates to ``_start`` for the subclass-specific
setup.
Args:
app_context: The shared application context to pass to ``_start``.
Raises:
ValueError: If the component configuration is invalid or required
dependencies are unavailable (raised by the subclass ``_start``).
"""
async def start(self) -> None:
"""Start the component. No-op if already started."""
if self._is_started:
return
await self._start(app_context)
await self._start()
self._is_started = True
async def close(self) -> None:
"""Close the component and release its resources.
This is the public entry point for component teardown. It guards
against redundant closes by returning immediately if the component
has not been started or is already closed, then delegates to ``_close``
for the subclass-specific cleanup.
Raises:
ValueError: If the component is in an inconsistent state that
prevents safe shutdown (raised by the subclass ``_close``).
"""
"""Close the component. No-op if not started."""
if not self._is_started:
return
try:
@ -134,50 +56,18 @@ class BaseComponent(ABC):
finally:
self._is_started = False
async def restart(self, app_context: ApplicationContext | None = None) -> None:
"""Restart the component by closing and then starting it again.
This method safely tears down the component if it is currently running
and reinitialized it. If the component is not started, it will simply
be started.
If either the close or start operation fails, the exception propagates
immediately and the component will be left in a non-started state.
Args:
app_context: The shared application context to pass during startup.
Raises:
ValueError: If the component cannot be cleanly shut down or
reinitialized due to invalid state or configuration.
"""
async def restart(self) -> None:
"""Close then start."""
await self.close()
await self.start(app_context)
await self.start()
@property
def is_started(self) -> bool:
"""Return whether the component is currently in a started state.
Returns:
``True`` if ``start()`` has been called and ``close()`` has not
been called since; ``False`` otherwise.
"""
return self._is_started
async def __call__(self, **kwargs):
"""Call the component instance as a function."""
async def __call__(self, **kwargs): ...
async def __aenter__(self) -> "BaseComponent":
"""Enter the async context manager by starting the component.
Returns:
The component instance, allowing it to be bound in an ``async with``
statement.
Raises:
ValueError: If the component fails to start due to invalid
configuration or missing dependencies.
"""
await self.start()
return self
@ -187,23 +77,6 @@ class BaseComponent(ABC):
exc_val: BaseException | None,
exc_tb,
) -> bool:
"""Exit the async context manager by closing the component.
Any exception raised within the context block is not suppressed.
If both the context block and ``close()`` raise exceptions, the
original exception from the context block is preserved and the
close exception is attached as its ``__cause__`` to maintain
the full error chain.
Args:
exc_type: The exception type if an exception was raised in the
context block, otherwise ``None``.
exc_val: The exception value if an exception was raised.
exc_tb: The traceback if an exception was raised.
Returns:
``False`` to indicate that exceptions should not be suppressed.
"""
if self._is_started:
if exc_val is not None:
try:

2
reme2/component/base_step.py
View file

@ -49,7 +49,7 @@ class BaseStep(BaseComponent):
self.output_mapping = output_mapping
self.context: RuntimeContext | None = None
async def _start(self, app_context=None) -> None:
async def _start(self) -> None:
"""Apply input mapping before execution."""
if self.input_mapping and self.context:
self.context.apply_mapping(self.input_mapping)

3
reme2/component/client/base_client.py
View file

@ -15,8 +15,7 @@ class BaseClient(BaseComponent):
super().__init__(**kwargs)
self.client = None
async def _start(self, app_context=None) -> None:
"""Initialize the client."""
async def _start(self) -> None:
async def _close(self) -> None:
"""Close the client."""

2
reme2/component/client/http_client.py
View file

@ -41,7 +41,7 @@ class HttpClient(BaseClient):
self.base_url = f"http://{host}:{port}"
self.timeout = timeout
async def _start(self, app_context=None) -> None:
async def _start(self) -> None:
"""Initialize the HTTP client."""
if self.client is None:
self.client = httpx.AsyncClient(

226
reme2/component/embedding/base_embedding_model.py
View file

@ -1,9 +1,10 @@
"""Base embedding model with caching, batching, and retry support."""
"""Base embedding model with caching and batching support."""
import asyncio
import hashlib
import json
import os
import time
import numpy as np
from abc import abstractmethod
from collections import OrderedDict
from pathlib import Path
@ -14,33 +15,30 @@ from ...schema import BaseNode
class BaseEmbeddingModel(BaseComponent):
"""Abstract base class for embedding models with LRU cache and retry logic.
"""Abstract base class for embedding models with LRU cache.
Provides:
- LRU in-memory cache with disk persistence (JSONL)
- LRU in-memory cache with disk persistence (npz)
- Automatic text truncation to max_input_length
- Retry logic with exponential backoff
- Batch embedding support
"""
component_type = ComponentEnum.EMBEDDING_MODEL
def __init__(
self,
api_key: str | None = None,
base_url: str | None = None,
model_name: str = "",
dimensions: int = 1024,
use_dimensions: bool = False,
max_batch_size: int = 10,
max_retries: int = 3,
raise_exception: bool = True,
max_input_length: int = 8192,
cache_dir: str | Path = ".reme",
max_cache_size: int = 2000,
enable_cache: bool = True,
encoding: str = "utf-8",
**kwargs,
self,
api_key: str | None = None,
base_url: str | None = None,
model_name: str = "",
dimensions: int = 1024,
pass_dimensions: bool = False,
max_batch_size: int = 10,
max_input_length: int = 8192,
cache_dir: str | Path = ".reme",
max_cache_size: int = 2000,
enable_cache: bool = True,
encoding: str = "utf-8",
**kwargs,
):
"""Initialize embedding model configuration.
@ -49,10 +47,8 @@ class BaseEmbeddingModel(BaseComponent):
base_url: Base URL for the embedding service.
model_name: Name of the embedding model.
dimensions: Vector dimensions.
use_dimensions: Whether to pass dimensions parameter to API.
pass_dimensions: Whether to pass dimensions parameter to API.
max_batch_size: Maximum batch size for embedding requests.
max_retries: Maximum retry attempts on failure.
raise_exception: Whether to raise exceptions on failure.
max_input_length: Maximum input text length.
cache_dir: Directory for cache storage.
max_cache_size: Maximum LRU cache size.
@ -60,14 +56,12 @@ class BaseEmbeddingModel(BaseComponent):
encoding: Text encoding for cache file operations.
"""
super().__init__(**kwargs)
self.api_key: str | None = api_key
self.base_url: str | None = base_url
self.api_key: str = api_key or os.environ.get("EMBEDDING_API_KEY", "")
self.base_url: str = base_url or os.environ.get("EMBEDDING_BASE_URL", "")
self.model_name = model_name
self.dimensions = dimensions
self.use_dimensions = use_dimensions
self.pass_dimensions = pass_dimensions
self.max_batch_size = max_batch_size
self.max_retries = max_retries
self.raise_exception = raise_exception
self.max_input_length = max_input_length
self.cache_dir = cache_dir
self.max_cache_size = max_cache_size
@ -79,6 +73,17 @@ class BaseEmbeddingModel(BaseComponent):
self._cache_misses = 0
self.cache_path: Path = Path(self.cache_dir)
async def _start(self) -> None:
"""Load cache on start."""
assert self.app_context is not None, "app_context must be provided"
working_path = Path(self.app_context.app_config.working_dir)
working_path.embedding_cache = self.cache_path
self._load_cache()
async def _close(self) -> None:
"""Save cache on close."""
self._save_cache()
def _truncate_text(self, text: str) -> str:
"""Truncate text to max_input_length."""
return text[: self.max_input_length] if len(text) > self.max_input_length else text
@ -106,11 +111,11 @@ class BaseEmbeddingModel(BaseComponent):
return hashlib.sha256(cache_string.encode(self.encoding)).hexdigest()
def _get_cache_file_path(self) -> Path:
"""Return path to the cache JSONL file."""
return self.cache_path / "embedding_cache.jsonl"
"""Return path to the cache npz file."""
return self.cache_path / "embedding_cache.npz"
def _load_cache(self) -> None:
"""Load embedding cache from disk (JSONL format)."""
"""Load embedding cache from disk (npz format)."""
if not self.enable_cache:
return
@ -122,38 +127,25 @@ class BaseEmbeddingModel(BaseComponent):
try:
load_start = time.time()
with open(cache_file, "r", encoding=self.encoding) as f:
lines = f.readlines()
data = np.load(cache_file)
keys = data["keys"]
embeddings = data["embeddings"]
loaded_count = 0
for line in reversed(lines):
line = line.strip()
if not line:
for key, emb in zip(keys, embeddings):
key_str = str(key)
emb_list = emb.tolist()
if len(emb_list) != self.dimensions:
self.logger.warning(
f"Cache dimension mismatch for {key_str}: "
f"expected {self.dimensions}, got {len(emb_list)}",
)
continue
try:
data = json.loads(line)
except json.JSONDecodeError as e:
self.logger.warning(f"Failed to parse cache line: {e}")
continue
if not data:
continue
cache_key, embedding = next(iter(data.items()))
if cache_key and embedding and isinstance(embedding, list):
if cache_key in self._embedding_cache:
continue
if len(embedding) != self.dimensions:
self.logger.warning(
f"Cache dimension mismatch for {cache_key}: "
f"expected {self.dimensions}, got {len(embedding)}",
)
continue
if len(self._embedding_cache) >= self.max_cache_size:
self.logger.info(f"Cache limit reached ({self.max_cache_size}), loaded {loaded_count}")
break
self._embedding_cache[cache_key] = embedding
loaded_count += 1
if len(self._embedding_cache) >= self.max_cache_size:
self.logger.info(f"Cache limit reached ({self.max_cache_size}), loaded {loaded_count}")
break
self._embedding_cache[key_str] = emb_list
loaded_count += 1
self.logger.info(f"Loaded {loaded_count} embeddings from {cache_file} in {time.time() - load_start:.2f}s")
except Exception as e:
@ -164,19 +156,27 @@ class BaseEmbeddingModel(BaseComponent):
self.logger.error(f"Failed to delete cache file: {del_e}")
def _save_cache(self) -> None:
"""Save embedding cache to disk (JSONL format)."""
"""Save embedding cache to disk (npz format)."""
if not self.enable_cache or not self._embedding_cache:
return
cache_file = self._get_cache_file_path()
try:
with open(cache_file, "w", encoding=self.encoding) as f:
for cache_key, embedding in self._embedding_cache.items():
if len(embedding) != self.dimensions:
self.logger.warning(f"Cache dimension mismatch for {cache_key}")
continue
f.write(json.dumps({cache_key: embedding}, ensure_ascii=False) + "\n")
self.logger.info(f"Saved {len(self._embedding_cache)} embeddings to {cache_file}")
keys = []
embeddings = []
for cache_key, embedding in self._embedding_cache.items():
if len(embedding) != self.dimensions:
self.logger.warning(f"Cache dimension mismatch for {cache_key}")
continue
keys.append(cache_key)
embeddings.append(embedding)
np.savez(
cache_file,
keys=np.array(keys, dtype=str),
embeddings=np.array(embeddings, dtype=np.float32),
)
self.logger.info(f"Saved {len(keys)} embeddings to {cache_file}")
except Exception as e:
self.logger.error(f"Failed to save cache to {cache_file}: {e}")
@ -242,34 +242,23 @@ class BaseEmbeddingModel(BaseComponent):
"""Fetch embeddings for a batch of texts. Override in subclasses."""
async def get_embedding(self, input_text: str, **kwargs) -> list[float]:
"""Get embedding for a single text with cache and retry."""
"""Get embedding for a single text with cache."""
truncated_text = self._truncate_text(input_text)
cached = self._get_from_cache(truncated_text)
if cached is not None:
return cached
for retry in range(self.max_retries):
try:
result = await self._get_embeddings([truncated_text], **kwargs)
if result and len(result) == 1:
embedding = self._validate_and_adjust_embedding(result[0])
self._put_to_cache(truncated_text, embedding)
return embedding
self.logger.warning(
f"Model {self.model_name} returned {len(result) if result else 0} results, expected 1",
)
if retry == self.max_retries - 1:
if self.raise_exception:
raise RuntimeError("Embedding API returned empty result")
return []
await asyncio.sleep(retry + 1)
except Exception as e:
self.logger.error(f"Model {self.model_name} failed: {e}")
if retry == self.max_retries - 1:
if self.raise_exception:
raise
return []
await asyncio.sleep(retry + 1)
try:
result = await self._get_embeddings([truncated_text], **kwargs)
if result and len(result) == 1:
embedding = self._validate_and_adjust_embedding(result[0])
self._put_to_cache(truncated_text, embedding)
return embedding
self.logger.warning(
f"Model {self.model_name} returned {len(result) if result else 0} results, expected 1",
)
except Exception as e:
self.logger.error(f"Model {self.model_name} failed: {e}")
return []
async def get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float]]:
@ -288,40 +277,29 @@ class BaseEmbeddingModel(BaseComponent):
if texts_to_compute:
uncached_texts = [text for _, text in texts_to_compute]
for i in range(0, len(uncached_texts), self.max_batch_size):
batch_texts = uncached_texts[i : i + self.max_batch_size]
batch_indices = [idx for idx, _ in texts_to_compute[i : i + self.max_batch_size]]
batch_texts = uncached_texts[i: i + self.max_batch_size]
batch_indices = [idx for idx, _ in texts_to_compute[i: i + self.max_batch_size]]
for retry in range(self.max_retries):
try:
batch_embeddings = await self._get_embeddings(batch_texts, **kwargs)
if batch_embeddings and len(batch_embeddings) == len(batch_texts):
for orig_idx, text, embedding in zip(batch_indices, batch_texts, batch_embeddings):
adjusted = self._validate_and_adjust_embedding(embedding)
results[orig_idx] = adjusted
self._put_to_cache(text, adjusted)
break
try:
batch_embeddings = await self._get_embeddings(batch_texts, **kwargs)
if batch_embeddings and len(batch_embeddings) == len(batch_texts):
for orig_idx, text, embedding in zip(batch_indices, batch_texts, batch_embeddings):
adjusted = self._validate_and_adjust_embedding(embedding)
results[orig_idx] = adjusted
self._put_to_cache(text, adjusted)
else:
self.logger.warning(
f"Batch returned {len(batch_embeddings) if batch_embeddings else 0} "
f"results for {len(batch_texts)} inputs",
)
if retry == self.max_retries - 1:
if self.raise_exception:
raise RuntimeError(f"Batch embedding failed after {self.max_retries} retries")
for orig_idx in batch_indices:
if results[orig_idx] is None:
results[orig_idx] = []
else:
await asyncio.sleep(retry + 1)
except Exception as e:
self.logger.error(f"Model {self.model_name} batch failed: {e}")
if retry == self.max_retries - 1:
if self.raise_exception:
raise
for orig_idx in batch_indices:
if results[orig_idx] is None:
results[orig_idx] = []
else:
await asyncio.sleep(retry + 1)
for orig_idx in batch_indices:
if results[orig_idx] is None:
results[orig_idx] = []
except Exception as e:
self.logger.error(f"Model {self.model_name} batch failed: {e}")
for orig_idx in batch_indices:
if results[orig_idx] is None:
results[orig_idx] = []
return [r if r is not None else [] for r in results]
@ -337,10 +315,4 @@ class BaseEmbeddingModel(BaseComponent):
self.logger.warning(f"Mismatch: {len(embeddings)} vectors for {len(nodes)} nodes, skipping assignment")
return nodes
async def _start(self, app_context=None) -> None:
"""Load cache on start."""
self._load_cache()
async def _close(self) -> None:
"""Save cache on close."""
self._save_cache()

6
reme2/component/embedding/openai_embedding_model.py
View file

@ -15,10 +15,10 @@ class OpenAIEmbeddingModel(BaseEmbeddingModel):
super().__init__(**kwargs)
self._client: AsyncOpenAI | None = None
async def _start(self, app_context=None) -> None:
async def _start(self) -> None:
"""Initialize the AsyncOpenAI client."""
self._client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url, **self.kwargs)
await super()._start(app_context)
await super()._start()
async def _close(self) -> None:
"""Close the AsyncOpenAI client."""
@ -37,7 +37,7 @@ class OpenAIEmbeddingModel(BaseEmbeddingModel):
"input": input_text,
**kwargs,
}
if self.use_dimensions:
if self.pass_dimensions:
create_kwargs["dimensions"] = self.dimensions
completion = await self._client.embeddings.create(**create_kwargs)

2
reme2/component/file_parser/base_file_parser.py
View file

@ -23,7 +23,7 @@ class BaseFileParser(BaseComponent):
self.chunk_tokens = chunk_tokens
self.chunk_overlap = chunk_overlap
async def _start(self, app_context=None):
async def _start(self):
pass
async def _close(self):

6
reme2/component/file_store/base_file_store.py
View file

@ -44,12 +44,12 @@ class BaseFileStore(BaseComponent):
if not self.vector_enabled and not self.fts_enabled:
raise ValueError("At least one of embedding_model or fts_enabled must be set.")
async def _start(self, app_context=None):
async def _start(self):
"""Resolve embedding model from app_context."""
if not self._embedding_model_name:
return
assert app_context is not None, "app_context must be provided"
models = app_context.components.get(ComponentEnum.EMBEDDING_MODEL, {})
assert self.app_context is not None, "app_context must be provided"
models = self.app_context.components.get(ComponentEnum.EMBEDDING_MODEL, {})
if self._embedding_model_name not in models:
raise ValueError(f"Embedding model '{self._embedding_model_name}' not found.")
model = models[self._embedding_model_name]

4
reme2/component/file_store/local_file_store.py
View file

@ -87,7 +87,7 @@ class LocalFileStore(BaseFileStore):
# -- Lifecycle ----------------------------------------------------------
async def _start(self, app_context=None) -> None:
async def _start(self) -> None:
"""Load persisted data into memory."""
await self._load_metadata()
await self._load_chunks()
@ -95,7 +95,7 @@ class LocalFileStore(BaseFileStore):
f"LocalFileStore '{self.store_name}' ready: "
f"{len(self._chunks)} chunks, metadata at {self._metadata_file}",
)
await super()._start(app_context)
await super()._start()
async def _close(self) -> None:
"""Flush state to disk and clear memory."""

8
reme2/component/file_watcher/base_file_watcher.py
View file

@ -54,12 +54,12 @@ class BaseFileWatcher(BaseComponent):
self._stop_event = asyncio.Event()
self._watch_task: asyncio.Task | None = None
async def _start(self, app_context=None):
async def _start(self):
"""Resolve file_store and start watching task."""
if self._file_store_name:
assert app_context is not None, "app_context must be provided"
assert self.app_context is not None, "app_context must be provided"
stores = app_context.components.get(ComponentEnum.FILE_STORE, {})
stores = self.app_context.components.get(ComponentEnum.FILE_STORE, {})
if self._file_store_name not in stores:
raise ValueError(f"File store '{self._file_store_name}' not found.")
store = stores[self._file_store_name]
@ -67,7 +67,7 @@ class BaseFileWatcher(BaseComponent):
raise TypeError(f"Expected BaseFileStore, got {type(store).__name__}")
self.file_store = store
parsers = app_context.components.get(ComponentEnum.FILE_PARSER, {})
parsers = self.app_context.components.get(ComponentEnum.FILE_PARSER, {})
for parser in parsers.values():
if isinstance(parser, BaseFileParser):
for suffix in parser.suffixes:

14
reme2/component/job/base_job.py
View file

@ -43,15 +43,9 @@ class BaseJob(BaseComponent):
self.step_configs: list[ComponentConfig] = steps or []
self.steps: list = []
async def _start(self, app_context=None) -> None:
"""Instantiate all configured steps.
Args:
app_context: Application context for dependency injection.
Raises:
ValueError: If a step backend is not specified or not registered.
"""
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.")
@ -61,7 +55,7 @@ class BaseJob(BaseComponent):
raise ValueError(f"{step_config.backend} is not registered.")
step = backend_cls(
language=app_context.app_config.language,
language=self.app_context.app_config.language,
**step_config.model_dump(exclude={"backend"}),
)
self.steps.append(step)

3
reme2/component/service/base_service.py
View file

@ -25,8 +25,7 @@ class BaseService(BaseComponent):
super().__init__(**kwargs)
self.service = None
async def _start(self, app_context=None) -> None:
"""Default empty implementation for sync services."""
async def _start(self) -> None:
async def _close(self) -> None:
"""Default empty implementation for sync services."""

4
reme2/component/service/http_service.py
View file

@ -5,6 +5,7 @@ import json
import os
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING
import uvicorn
from fastapi import FastAPI
@ -18,6 +19,9 @@ from ...constants import REME_DEFAULT_HOST, REME_DEFAULT_PORT, REME_SERVICE_INFO
from ...schema import Request, Response
from ...utils import execute_stream_task
if TYPE_CHECKING:
from ...application import Application
@R.register("http")
class HttpService(BaseService):

2
reme2/config/paw.yaml
View file

@ -36,7 +36,7 @@ components:
backend: openai
model_name: text-embedding-v3
dimensions: 1024
use_dimensions: false
pass_dimensions: false
enable_cache: true
max_batch_size: 10
max_cache_size: 2000