diff --git a/reme/memory/vector_based/personal/personal_retriever.py b/reme/memory/vector_based/personal/personal_retriever.py index b8803717..2b0c094e 100644 --- a/reme/memory/vector_based/personal/personal_retriever.py +++ b/reme/memory/vector_based/personal/personal_retriever.py @@ -106,5 +106,5 @@ class PersonalRetriever(BaseMemoryAgent): ], ) - result["retrieved_nodes"] = self.retrieved_nodes + result["retrieved_nodes"] = self.retri eved_nodes return result diff --git a/reme/memory/vector_based/personal/personal_retriever.yaml b/reme/memory/vector_based/personal/personal_retriever.yaml index 99df1374..e277df92 100644 --- a/reme/memory/vector_based/personal/personal_retriever.yaml +++ b/reme/memory/vector_based/personal/personal_retriever.yaml @@ -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 diff --git a/reme/memory/vector_based/personal/personal_summarizer.yaml b/reme/memory/vector_based/personal/personal_summarizer.yaml index 83befa96..bb0eb1b7 100644 --- a/reme/memory/vector_based/personal/personal_summarizer.yaml +++ b/reme/memory/vector_based/personal/personal_summarizer.yaml @@ -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 [] : {context} diff --git a/reme2/application.py b/reme2/application.py index 3672be46..c9f52e39 100644 --- a/reme2/application.py +++ b/reme2/application.py @@ -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 diff --git a/reme2/component/application_context.py b/reme2/component/application_context.py index 4b372353..4c6c5587 100644 --- a/reme2/component/application_context.py +++ b/reme2/component/application_context.py @@ -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] = {} diff --git a/reme2/component/as_llm/__init__.py b/reme2/component/as_llm/__init__.py index d3f8625c..18ffdd3d 100644 --- a/reme2/component/as_llm/__init__.py +++ b/reme2/component/as_llm/__init__.py @@ -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", ] diff --git a/reme2/component/as_llm_formatter/__init__.py b/reme2/component/as_llm_formatter/__init__.py index 0dc3894d..3e8eb5b0 100644 --- a/reme2/component/as_llm_formatter/__init__.py +++ b/reme2/component/as_llm_formatter/__init__.py @@ -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", ] diff --git a/reme2/component/as_token_counter/__init__.py b/reme2/component/as_token_counter/__init__.py index 1bb4389e..8c102add 100644 --- a/reme2/component/as_token_counter/__init__.py +++ b/reme2/component/as_token_counter/__init__.py @@ -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) diff --git a/reme2/component/as_token_counter/estimate_token_counter.py b/reme2/component/as_token_counter/estimate_token_counter.py index 790dca10..eca3f446 100644 --- a/reme2/component/as_token_counter/estimate_token_counter.py +++ b/reme2/component/as_token_counter/estimate_token_counter.py @@ -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: diff --git a/reme2/component/base_component.py b/reme2/component/base_component.py index 61518fc2..0900f8da 100644 --- a/reme2/component/base_component.py +++ b/reme2/component/base_component.py @@ -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: diff --git a/reme2/component/base_step.py b/reme2/component/base_step.py index 7892334c..e4ec87ac 100644 --- a/reme2/component/base_step.py +++ b/reme2/component/base_step.py @@ -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) diff --git a/reme2/component/client/base_client.py b/reme2/component/client/base_client.py index d5da961b..d92b7d50 100644 --- a/reme2/component/client/base_client.py +++ b/reme2/component/client/base_client.py @@ -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.""" diff --git a/reme2/component/client/http_client.py b/reme2/component/client/http_client.py index c93870eb..d01b81a0 100644 --- a/reme2/component/client/http_client.py +++ b/reme2/component/client/http_client.py @@ -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( diff --git a/reme2/component/embedding/base_embedding_model.py b/reme2/component/embedding/base_embedding_model.py index 69526b40..bbc4ed3e 100644 --- a/reme2/component/embedding/base_embedding_model.py +++ b/reme2/component/embedding/base_embedding_model.py @@ -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() diff --git a/reme2/component/embedding/openai_embedding_model.py b/reme2/component/embedding/openai_embedding_model.py index faa97706..e2405b41 100644 --- a/reme2/component/embedding/openai_embedding_model.py +++ b/reme2/component/embedding/openai_embedding_model.py @@ -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) diff --git a/reme2/component/file_parser/base_file_parser.py b/reme2/component/file_parser/base_file_parser.py index a7e9d793..fe310a29 100644 --- a/reme2/component/file_parser/base_file_parser.py +++ b/reme2/component/file_parser/base_file_parser.py @@ -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): diff --git a/reme2/component/file_store/base_file_store.py b/reme2/component/file_store/base_file_store.py index 19f78e01..9167426d 100644 --- a/reme2/component/file_store/base_file_store.py +++ b/reme2/component/file_store/base_file_store.py @@ -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] diff --git a/reme2/component/file_store/local_file_store.py b/reme2/component/file_store/local_file_store.py index 98b785b8..fb0df8c6 100644 --- a/reme2/component/file_store/local_file_store.py +++ b/reme2/component/file_store/local_file_store.py @@ -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.""" diff --git a/reme2/component/file_watcher/base_file_watcher.py b/reme2/component/file_watcher/base_file_watcher.py index e7d8541c..512c65f4 100644 --- a/reme2/component/file_watcher/base_file_watcher.py +++ b/reme2/component/file_watcher/base_file_watcher.py @@ -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: diff --git a/reme2/component/job/base_job.py b/reme2/component/job/base_job.py index 8cf2776f..88c41c58 100644 --- a/reme2/component/job/base_job.py +++ b/reme2/component/job/base_job.py @@ -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) diff --git a/reme2/component/service/base_service.py b/reme2/component/service/base_service.py index cb33069e..c4820e2d 100644 --- a/reme2/component/service/base_service.py +++ b/reme2/component/service/base_service.py @@ -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.""" diff --git a/reme2/component/service/http_service.py b/reme2/component/service/http_service.py index 3f3477bd..d264f87c 100644 --- a/reme2/component/service/http_service.py +++ b/reme2/component/service/http_service.py @@ -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): diff --git a/reme2/config/paw.yaml b/reme2/config/paw.yaml index 0b9c1d07..1ef43102 100644 --- a/reme2/config/paw.yaml +++ b/reme2/config/paw.yaml @@ -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