Merge branch 'main' into dev_0227

This commit is contained in:
caozouying.czy 2026-02-27 14:06:11 +08:00
commit bb0d0ee248
182 changed files with 1525 additions and 2185 deletions

View file

@ -320,7 +320,6 @@ class MemoryProcessor:
return_dict=True,
enable_time_filter=True,
enable_thinking_params=self.enable_thinking_params,
llm_config_name="qwen-plus-t",
)
duration_ms = (time.time() - start) * 1000
@ -355,7 +354,6 @@ class MemoryProcessor:
return_dict=True,
enable_time_filter=True,
enable_thinking_params=self.enable_thinking_params,
llm_config_name="qwen-plus-t",
)
# Extract memories from response
@ -569,15 +567,6 @@ class HaluMemEvaluator:
default_llm_config={
"model_name": self.config.reme_model_name,
},
llms={
"qwen-plus-t": {
"backend": "openai",
"model_name": "qwen-plus",
"extra_body": {
"enable_thinking": True,
},
},
},
)
# Load evaluation prompts into ReMe's prompt handler

View file

@ -113,5 +113,11 @@ remecli = "reme.reme_cli:main"
[tool.pytest.ini_options]
asyncio_default_fixture_loop_scope = "function"
# Script-style tests that need to be run with `python test_*.py`
testpaths = ["tests"]
python_files = ["test_*.py"]
python_functions = ["test_*"]
# Exclude script-style tests that require manual execution
addopts = "--ignore=tests/test_embedding.py --ignore=tests/test_embedding_cache.py --ignore=tests/test_embedding_sync.py --ignore=tests/test_file_store.py"
# python -m build && twine upload dist/*

View file

@ -1,27 +1,27 @@
"""ReMe"""
from . import agent
from . import config
from . import core
from . import tool
from . import extension
from . import memory
from . import workflow
from .reme import ReMe
from .reme_cli import ReMeCli
from .reme_fs import ReMeFs
from .reme_fb import ReMeFb
__version__ = "0.3.0.0b5"
__all__ = [
"agent",
"config",
"core",
"tool",
"extension",
"memory",
"workflow",
"ReMe",
"ReMeCli",
"ReMeFs",
"ReMeFb",
]
__version__ = "0.3.0.0b4"
"""
conda create -n fl_test2 python=3.10
conda activate fl_test2

View file

@ -1,9 +0,0 @@
"""A simple chatbot."""
from . import chat
from . import memory
__all__ = [
"chat",
"memory",
]

View file

@ -1,16 +0,0 @@
"""chat agent"""
from .fs_cli import FsCli
from .simple_chat import SimpleChat
from .stream_chat import StreamChat
from ...core import R
__all__ = [
"FsCli",
"StreamChat",
"SimpleChat",
]
R.ops.register(FsCli)
R.ops.register(SimpleChat)
R.ops.register(StreamChat)

View file

@ -1,11 +0,0 @@
"""File system agents for memory management."""
from .fs_compactor import FsCompactor
from .fs_context_checker import FsContextChecker
from .fs_summarizer import FsSummarizer
__all__ = [
"FsSummarizer",
"FsCompactor",
"FsContextChecker",
]

View file

@ -1,92 +0,0 @@
"""Personal memory retriever agent for retrieving personal memories through vector search."""
from ..base_memory_agent import BaseMemoryAgent
from ....core.enumeration import Role, MemoryType
from ....core.op import BaseTool
from ....core.schema import Message
from ....core.utils import format_messages
class PersonalHalumemRetriever(BaseMemoryAgent):
"""Retrieve personal memories through vector search and history reading."""
memory_type: MemoryType = MemoryType.PERSONAL
async def build_messages(self) -> list[Message]:
if self.context.get("query"):
context = self.context.query
elif self.context.get("messages"):
context = self.description + "\n" + format_messages(self.context.messages)
else:
raise ValueError("input must have either `query` or `messages`")
read_all_profiles_tool: BaseTool | None = self.pop_tool("read_all_profiles")
if read_all_profiles_tool is not None:
all_profiles = await read_all_profiles_tool.call(
memory_target=self.memory_target,
service_context=self.service_context,
)
else:
all_profiles = ""
return [
Message(
role=Role.SYSTEM,
content=self.prompt_format(
prompt_name="system_prompt",
memory_type=self.memory_type.value,
memory_target=self.memory_target,
user_profile=all_profiles,
context=context.strip(),
),
),
Message(
role=Role.USER,
content=self.prompt_format(
prompt_name="user_message",
memory_type=self.memory_type.value,
memory_target=self.memory_target,
user_profile=all_profiles,
context=context.strip(),
),
),
]
async def _acting_step(
self,
assistant_message: Message,
tools: list[BaseTool],
step: int,
stage: str = "",
**kwargs,
) -> tuple[list[BaseTool], list[Message]]:
"""Execute tool calls with memory context."""
return await super()._acting_step(
assistant_message,
tools,
step,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
retrieved_nodes=self.retrieved_nodes,
**kwargs,
)
async def execute(self):
result = await super().execute()
answer = result["answer"]
if "MEMORY_NOT_FOUND" in answer:
result["answer"] = "\n".join(
[
n.format(
include_memory_id=False,
include_when_to_use=False,
include_content=True,
include_message_time=False,
ref_memory_id_key="",
)
for n in self.retrieved_nodes
],
)
result["retrieved_nodes"] = self.retrieved_nodes
return result

View file

@ -1,91 +0,0 @@
system_prompt: |
# Role Definition:
You are a Memory Retrieval Agent specialized in retrieving {memory_type} memories about {memory_target}.
## Multi-Phase Retrieval Strategy
Follow these phases sequentially to gather comprehensive information:
### Tool Rules
**Tool**: `retrieve_memory` (without time constraints)
**Objective**: Cast a wide net to find potentially relevant memories
**Approach**:
- Execute 3-5 diverse search queries using different formulations:
* Original question verbatim
* Rephrased variations (different wording, synonyms)
* Entity-focused queries (extract and search specific names, places, events)
* Keyword-based searches (core concepts, topics)
* Related context queries (broader themes)
- Review all results before proceeding to next phase
**Tool**: `retrieve_memory` (with time filter)
**When to use**: Only if the user question contains temporal references
**Time Filter Format**:
- Single date: `20200101`
- Date range: `20200101,20200102` (inclusive: 20200101 ≤ time ≤ 20200102)
- Before date: `0,20200102` (up to and including 20200102)
- After date: `20200101,99999999` (from 20200101 onwards)
**Approach**:
- Identify temporal constraints from the user question
- Refine Phase 1 queries with appropriate time filters
- Try multiple time ranges if initial searches yield no results
**Tool**: `read_history`
**When to use**: After exhausting retrieval attempts OR when specific conversation context is needed
**Approach**:
- Extract `history_id` from retrieved memory references
- Prioritize histories that are most relevant or recent
- Read multiple histories if necessary for complete context
- Use this to understand the full conversation surrounding a memory
user_message: |
## User Profile
{user_profile}
## User Question
{context}
# Core Objective:
Before responding to the user, you must strictly follow the **[Memory Retrieve -> Original Source Tracing -> Broad Search Fallback]** retrieval strategy. It is strictly forbidden to directly opt for an indiscriminate search of massive historical original texts.
## Retrieval Strategy & Workflow (Strictly Enforced Chain of Thought)
### Phase 1: Intent Decomposition and Primary Retrieval (Summary First)
1. **Analyze Intent**: Analyze the user's current Query, decomposing it into 1-3 core search intents.2. **Summary Priority**: First, retrieve from **high-level memories**.
- **Action**: Call `vector_retrieve_memory` using at least two different `query`.
- **Filter**: (Optional) Set metadata filter {{"timestamp": "YYYY-MM-DD"}}
- **Goal**: Obtain refined conclusions such as entity attributes, task status, user preferences, or environmental information.
### Phase 2: Memory Evaluation and Deep Tracing (Drill Down)
Check the retrieval results of Phase 1:
- **Case A (Sufficient Information)**: If the summarized memory contains all the details needed for the answer, proceed directly to Phase 4 for the response.
- **Case B (Vague/Complex Information)**: If summarized memory exists (e.g., 'discussed project architecture') but lacks specific details (e.g., 'specific parameter configuration'), use clues from the summary to trace the original text.
- **Action**: Call `read_history` using ref_memory_id from the retrieved memory.
- **Goal**: Obtain the specific conversation context at that time.
### Phase 3: Fallback Retrieval and Strategy Adjustment (Fallback & Expand)
If no valid information is found in both Phase 1 and Phase 2 (result is empty or similarity is too low): Rewrite the Query based on the context (remove non-keywords, synonym substitution), and search again.
### Phase 4: Result Compilation and Response
- Combine the retrieved content (summary or original text) with the current conversation context.
- If all retrieved results are irrelevant, **it is strictly forbidden to fabricate memories**; directly inform the user that no relevant information was found.
## Output Format
Before the final reply, ensure at least 3 tool calls for retrieval, and then output your answer in ten words.
- Base your answer EXCLUSIVELY on retrieved memories, user profile, and history data
- Never infer, assume, or hallucinate information
- Always cite sources with timestamps: `[timestamp] Memory content`
- Present conflicting information transparently with respective timestamps
- Exhaust all search strategies before concluding information doesn't exist
Before the final reply, ensure at least 3 tool calls for retrieval, and then output the most relevant JSON retrieval summary, followed by your answer:
```json
{{
"retrieved_memories": [
{{"type": "profile", "timestamp":"...", "content": "..."}},
{{"type": "personal", "timestamp":"...", "content": "..."}},
{{"type": "history", "timestamp":"...", "content": "..."}},
....
],
"summary": "Fill in your summarized answer here."
}}
```

View file

@ -1,140 +0,0 @@
"""Personal memory summarizer agent for two-phase personal memory processing."""
from loguru import logger
from ..base_memory_agent import BaseMemoryAgent
from ....core.enumeration import Role, MemoryType
from ....core.op import BaseTool
from ....core.schema import Message
class PersonalHalumemSummarizer(BaseMemoryAgent):
"""Two-phase personal memory processor: retrieve/add memories then update profile."""
memory_type: MemoryType = MemoryType.PERSONAL
async def _build_s1_messages(self) -> list[Message]:
return [
Message(
role=Role.SYSTEM,
content=self.prompt_format(
prompt_name="system_prompt_s1",
context=self.context.history_node.content,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
),
),
Message(
role=Role.USER,
# content=self.get_prompt("user_message_s1"),
content=self.prompt_format(
prompt_name="user_message_s1",
context=self.context.history_node.content,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
),
),
]
async def _build_s2_messages(self, user_profile: str) -> list[Message]:
return [
Message(
role=Role.SYSTEM,
content=self.prompt_format(
prompt_name="system_prompt_s2",
context=self.context.history_node.content,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
user_profile=user_profile,
),
),
Message(
role=Role.USER,
content=self.prompt_format(
prompt_name="user_message_s2",
context=self.context.history_node.content,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
user_profile=user_profile,
),
),
]
async def _acting_step(
self,
assistant_message: Message,
tools: list[BaseTool],
step: int,
stage: str = "",
**kwargs,
) -> tuple[list[BaseTool], list[Message]]:
"""Execute tool calls with memory context."""
return await super()._acting_step(
assistant_message,
tools,
step,
stage=stage,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
history_node=self.history_node,
author=self.author,
retrieved_nodes=self.retrieved_nodes,
**kwargs,
)
async def execute(self):
memory_tools = []
profile_tools = []
for i, tool in enumerate(self.tools):
tool_name = tool.tool_call.name
if "_memory" in tool_name:
memory_tools.append(tool)
elif "_profile" in tool_name:
profile_tools.append(tool)
else:
raise ValueError(f"[{self.__class__.__name__}] unknown tool_name={tool_name}")
logger.info(f"[{self.__class__.__name__}] tool_call[{i}]={tool.tool_call.simple_input_dump(as_dict=False)}")
stage = "s1-memory"
messages_s1 = await self._build_s1_messages()
for i, message in enumerate(messages_s1):
role = message.name or message.role
logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}")
tools_s1, messages_s1, success_s1 = await self.react(messages_s1, memory_tools, stage=stage)
if profile_tools:
read_all_profiles_tool: BaseTool | None = self.pop_tool("read_all_profiles")
if read_all_profiles_tool is not None:
all_profiles = await read_all_profiles_tool.call(
memory_target=self.memory_target,
service_context=self.service_context,
)
else:
all_profiles = ""
stage = "s2-profile"
messages_s2 = await self._build_s2_messages(user_profile=all_profiles)
for i, message in enumerate(messages_s2):
role = message.name or message.role
logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}")
tools_s2, messages_s2, success_s2 = await self.react(messages_s2, profile_tools, stage=stage)
else:
tools_s2, messages_s2, success_s2 = [], [], True
answer = (messages_s1[-1].content if success_s1 else "") + (messages_s2[-1].content if success_s2 else "")
success = success_s1 and success_s2
messages = messages_s1 + messages_s2
tools = tools_s1 + tools_s2
memory_nodes = []
for tool in tools:
if tool.memory_nodes:
memory_nodes.extend(tool.memory_nodes)
return {
"answer": answer,
"success": success,
"messages": messages,
"tools": tools,
"memory_nodes": memory_nodes,
}

View file

@ -1,96 +0,0 @@
system_prompt_s1: |
You are a Memory Agent responsible for managing {memory_type} memories about {memory_target}.
## Tool Rules
1. `add_and_retrieve_similar_memory`: Create a memory in the vector store.
- Use this tool to add memories, and it will return the relevant content related to the added memories.
- Use actual names from the conversation (e.g., "Bob likes apples") instead of generic references (e.g., "user likes apples")
- The tool will retrieve similar historical memories via vector search to help you consolidate in Step 2
2. `update_memory`: Update memories in the vector store.
**What to Delete** (via `memory_ids_to_delete`):
- Duplicate memories with identical or highly similar content
- Memories that should be merged into a single consolidated entry
**What to Add** (via `memories_to_add` with message_time and memory_content):
- For each topic with changes: add ONE consolidated memory that merges related information
- New distinct memories that don't overlap with existing ones
- Updated memories that capture the latest state while preserving temporal evolution
user_message_s1: |
## Latest Conversation
Format: round<index> [<timestamp>] <role/name>: <content>
{context}
## Task
### Step 1: Create Memory
- At this step, you can call the tool multiple times to store memories, or you can call it once to store multiple memories.
### Step 2: Update Memory Store
- Update the vector store using `update_memory` to keep it well-organized and consolidated.
## Storage Scope (Biographical & Behavioral ONLY)
- Personal Biography: Significant milestones, past experiences, and life events.
- Behavioral Patterns: How the agent reacts, specific actions taken, and recurring habits.
- **EXCLUSION**: DO NOT record objective world facts, general knowledge, or user-specific health/states.
Extraction & Formatting Rules
- Fact Filtering: Only extract information that builds the biography of **{memory_target}**.
- Subject Splitting: If a conversation mentions multiple subject (e.g., the User's childhood and their Father's career), create separate memory entries for each subject.
- Atomic Content: Each entry should focus on one specific event or trait. Keep descriptions concise to ensure efficient retrieval.
system_prompt_s2: |
You are a Profile Agent responsible for managing profiles about {memory_target}.
## Tool Rules
1. Update Profile with `update_profile`
Synchronize profile with new information from the conversation:
- `profile_ids_to_delete`: Remove conflicting, or redundant entries (array of profile IDs).
- `profiles_to_add`:
- `conversation_time`: Time of conversation (format: `YYYY-MM-DD HH:MM:SS`, e.g., `2024-01-15 14:30:00`)
- `profile_content`: Complete, self-contained profile description with full context
Update user profile using `update_profile` based on the conversation and current profile.
user_message_s2: |
You are a memory agent managing **{memory_type}** memories about **{memory_target}**.
## Latest Conversation:
{context}
Message format: `round<index> [<timestamp>] <role/name>: <content>` (timestamp: YYYY-MM-DD HH:MM:SS).
**CRITICAL**: Extract ONLY explicitly stated information. DO NOT infer, assume, or fabricate.
## Current User Profile:
{user_profile}
## Task
### Step 1: ADD Profile
- Add new, relevant, and up-to-date information to the user profile using `update_profile` (via `profiles_to_add`).
### Step 2: DELETE Profile
- Delete outdated, redundant, or resolved states from the user profile using `update_profile` (via `profile_ids_to_delete`).
## Storage Scope (Current States ONLY)
- **EXCLUSION PRINCIPLE**: DO NOT record any user *actions*, *requests*, *queries*, or *interactions with the system* (e.g., "asked for code", "solved a puzzle", "requested translation"). These are interaction logs, not user states.
- Identity: Geography, job title, work content, income.
- Background: Education, family, relationships, hobbies, interests, and other personal preferences.
- Temporary States: Physical health (e.g., "Has a cold"), emotional mood, stress levels, and specific prohibitions (e.g., "Cannot drink alcohol due to medication").
## Profile Management Rules
- Subject Splitting (CRITICAL): If the conversation mentions multiple subjects (e.g., the User's job and their Spouse's health), you MUST create separate profile entries for each unique subject.
- Conflict Resolution: Use profile_ids_to_delete to remove outdated, redundant, or resolved states (e.g., if a user is "Recovered," delete the "Illness" entry).
- Each profile entry MUST describe a **persistent or temporary state of the user themselves** (e.g., who they are, what they like, what theyre dealing with), NOT an event they participated in or a request they made.
## Profile Format
- **ONLY record what the user EXPLICITLY STATES about themselves as a state or preference.**
- The key in the record represents the category of memory, and the value should record the specific content. For example
{{"message_time": "YYYY-MM-DD HH:MM:SS", "profile_key": "the category of memory", "profile_value": "content" }}
- When there is no information conflict or outdated information, you don't need to delete any of the memory. If there is no information that meets the requirements, it is also acceptable not to add it.
## Forbidden Case
1.There is no need to record user behavior: {{ "profile_key": "workouts", "profile_content": "confident in new running shoes' suitability for chosen route; they have significantly improved morning jogs"}}
2. There is no need to record the users' plans or requirements: {{ "profile_key": "plans.vacation", "profile_content": "planning to go to a nearby city for a week and ask for a job change"}}

View file

@ -1,86 +0,0 @@
"""Personal memory retriever agent for retrieving personal memories through vector search."""
from ..base_memory_agent import BaseMemoryAgent
from ....core.enumeration import Role, MemoryType
from ....core.op import BaseTool
from ....core.schema import Message
from ....core.utils import format_messages
class PersonalRetriever(BaseMemoryAgent):
"""Retrieve personal memories through vector search and history reading."""
memory_type: MemoryType = MemoryType.PERSONAL
async def build_messages(self) -> list[Message]:
if self.context.get("query"):
context = self.context.query
elif self.context.get("messages"):
context = self.description + "\n" + format_messages(self.context.messages)
else:
raise ValueError("input must have either `query` or `messages`")
read_all_profiles_tool: BaseTool | None = self.pop_tool("read_all_profiles")
if read_all_profiles_tool is not None:
all_profiles = await read_all_profiles_tool.call(
memory_target=self.memory_target,
service_context=self.service_context,
)
else:
all_profiles = ""
return [
Message(
role=Role.SYSTEM,
content=self.prompt_format(
prompt_name="system_prompt",
memory_type=self.memory_type.value,
memory_target=self.memory_target,
user_profile=all_profiles,
context=context.strip(),
),
),
Message(
role=Role.USER,
content=self.get_prompt("user_message"),
),
]
async def _acting_step(
self,
assistant_message: Message,
tools: list[BaseTool],
step: int,
stage: str = "",
**kwargs,
) -> tuple[list[BaseTool], list[Message]]:
"""Execute tool calls with memory context."""
return await super()._acting_step(
assistant_message,
tools,
step,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
retrieved_nodes=self.retrieved_nodes,
**kwargs,
)
async def execute(self):
result = await super().execute()
answer = result["answer"]
if "MEMORY_NOT_FOUND" in answer:
result["answer"] = "\n".join(
[
n.format(
include_memory_id=False,
include_when_to_use=False,
include_content=True,
include_message_time=False,
ref_memory_id_key="",
)
for n in self.retrieved_nodes
],
)
result["retrieved_nodes"] = self.retrieved_nodes
return result

View file

@ -1,66 +0,0 @@
system_prompt: |
You are a Memory Retrieval Agent specialized in retrieving {memory_type} memories about {memory_target}.
## User Profile
{user_profile}
## User Question
{context}
## Multi-Phase Retrieval Strategy
Follow these phases sequentially to gather comprehensive information:
### Phase 1: Semantic Search (No Time Filter)
**Tool**: `retrieve_memory` (without time constraints)
**Objective**: Cast a wide net to find potentially relevant memories
**Approach**:
- Execute 3-5 diverse search queries using different formulations:
* Original question verbatim
* Rephrased variations (different wording, synonyms)
* Entity-focused queries (extract and search specific names, places, events)
* Keyword-based searches (core concepts, topics)
* Related context queries (broader themes)
- Review all results before proceeding to next phase
### Phase 2: Temporal Search (Optional)
**Tool**: `retrieve_memory` (with time filter)
**When to use**: Only if the user question contains temporal references
**Time Filter Format**:
- Single date: `20200101`
- Date range: `20200101,20200102` (inclusive: 20200101 ≤ time ≤ 20200102)
- Before date: `0,20200102` (up to and including 20200102)
- After date: `20200101,99999999` (from 20200101 onwards)
**Approach**:
- Identify temporal constraints from the user question
- Refine Phase 1 queries with appropriate time filters
- Try multiple time ranges if initial searches yield no results
### Phase 3: Deep Dive into History
**Tool**: `read_history`
**When to use**: After exhausting retrieval attempts OR when specific conversation context is needed
**Approach**:
- Extract `history_id` from retrieved memory references
- Prioritize histories that are most relevant or recent
- Read multiple histories if necessary for complete context
- Use this to understand the full conversation surrounding a memory
## Response Guidelines
**Critical Rules**:
- Base your answer EXCLUSIVELY on retrieved memories, user profile, and history data
- Never infer, assume, or hallucinate information
- Always cite sources with timestamps: `[timestamp] Memory content`
- Present conflicting information transparently with respective timestamps
- Exhaust all search strategies before concluding information doesn't exist
**Output Format**:
- When information is found:
<MEMORY_FOUND>
[timestamp] All relevant memory/profile/history content
- When no information is found after thorough search (5+ queries across phases):
<MEMORY_NOT_FOUND>
No relevant information found after exhaustive search using multiple query strategies and retrieval phases.
user_message: |
Retrieve relevant memories following the multi-phase strategy outlined above.

View file

@ -1,65 +0,0 @@
system_prompt_s1: |
You are a Memory Agent responsible for managing {memory_type} memories about {memory_target}.
## Latest Conversation
Format: round<index> [<timestamp>] <role/name>: <content>
{context}
## Task
### Step 1: Create Memory Draft
Create a memory draft in `add_draft_and_retrieve_similar_memory` based on the latest conversation.
- Use actual names from the conversation (e.g., "Bob likes apples") instead of generic references (e.g., "user likes apples")
- Always record memories with real names
- The tool will retrieve similar historical memories via vector search to help you consolidate in Step 2
### Step 2: Update Memory Store
Update the vector store using `update_memory` to keep it well-organized and consolidated:
**What to Delete** (via `memory_ids_to_delete`):
- Duplicate memories with identical or highly similar content
- Memories that should be merged into a single consolidated entry
**What to Add** (via `memories_to_add` with message_time and memory_content):
- For each topic with changes: add ONE consolidated memory that merges related information
- New distinct memories that don't overlap with existing ones
- Updated memories that capture the latest state while preserving temporal evolution
## Requirements
- Extract only what's explicitly stated—no inferences, assumptions, or fabrications
- Preserve temporal evolution: capture how things change over time within the same topic
- Maintain organization: group related memories by topic and eliminate all redundancy
user_message_s1: |
Complete the task by following Step 1 and Step 2 in order
system_prompt_s2: |
You are a Profile Agent responsible for managing profiles about {memory_target}.
## Latest Conversation
Format: round<index> [<timestamp>] <role/name>: <content>
{context}
## Task
### Step 1: Create Profile Draft
Create a profile draft in `add_draft_and_read_all_profiles` based on the latest conversation.
- The tool will return all existing profiles to help you maintain the profile store in Step 2
### Step 2: Update Profile Store
Update the profile store using `update_profile` to keep it well-organized and consolidated:
**What to Delete** (via `profile_ids_to_delete`):
- Duplicate profiles with identical keys or values
- Conflicting profiles that contradict the new information
- Profiles that should be merged into a single consolidated entry
**What to Add** (via `profiles_to_add` with message_time, profile_key, and profile_value):
- For each profile key with changes: add ONE consolidated profile that merges related information
- New distinct profiles that don't overlap with existing ones
- Updated profiles that capture the latest state
## Requirements
- Extract only what's explicitly stated—no inferences, assumptions, or fabrications
- Maintain organization: group related profiles by key and eliminate all redundancy
user_message_s2: |
Complete the task by following Step 1 and Step 2 in order

View file

@ -1,123 +0,0 @@
"""Personal memory summarizer agent for two-phase personal memory processing."""
from loguru import logger
from ..base_memory_agent import BaseMemoryAgent
from ....core.enumeration import Role, MemoryType
from ....core.op import BaseTool
from ....core.schema import Message
class PersonalV1Summarizer(BaseMemoryAgent):
"""Two-phase personal memory processor: retrieve/add memories then update profile."""
memory_type: MemoryType = MemoryType.PERSONAL
async def _build_s1_messages(self) -> list[Message]:
return [
Message(
role=Role.USER,
content=self.prompt_format(
prompt_name="user_message_s1",
context=self.context.history_node.content,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
),
),
]
async def _build_s2_messages(self, profiles: str) -> list[Message]:
return [
Message(
role=Role.USER,
content=self.prompt_format(
prompt_name="user_message_s2",
profiles=profiles,
context=self.context.history_node.content,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
),
),
]
async def _acting_step(
self,
assistant_message: Message,
tools: list[BaseTool],
step: int,
stage: str = "",
**kwargs,
) -> tuple[list[BaseTool], list[Message]]:
"""Execute tool calls with memory context."""
return await super()._acting_step(
assistant_message,
tools,
step,
stage=stage,
memory_type=self.memory_type.value,
memory_target=self.memory_target,
history_node=self.history_node,
author=self.author,
retrieved_nodes=self.retrieved_nodes,
**kwargs,
)
async def execute(self):
memory_tools = []
profile_tools = []
read_all_profiles_tool: BaseTool | None = None
for i, tool in enumerate(self.tools):
tool_name = tool.tool_call.name
if tool_name == "read_all_profiles":
read_all_profiles_tool = tool
elif "_memory" in tool_name:
memory_tools.append(tool)
elif "_profile" in tool_name:
profile_tools.append(tool)
else:
raise ValueError(f"[{self.__class__.__name__}] unknown tool_name={tool_name}")
logger.info(f"[{self.__class__.__name__}] tool_call[{i}]={tool.tool_call.simple_input_dump(as_dict=False)}")
stage = "s1-memory"
messages_s1 = await self._build_s1_messages()
for i, message in enumerate(messages_s1):
role = message.name or message.role
logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}")
tools_s1, messages_s1, success_s1 = await self.react(messages_s1, memory_tools, stage=stage)
if read_all_profiles_tool is not None:
profiles = await read_all_profiles_tool.call(
memory_target=self.memory_target,
service_context=self.service_context,
)
else:
profiles = ""
if profile_tools:
stage = "s2-profile"
messages_s2 = await self._build_s2_messages(profiles)
for i, message in enumerate(messages_s2):
role = message.name or message.role
logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}")
tools_s2, messages_s2, success_s2 = await self.react(messages_s2, profile_tools, stage=stage)
else:
tools_s2, messages_s2, success_s2 = [], [], True
answer = (messages_s1[-1].content if success_s1 and messages_s1 else "") + (
messages_s2[-1].content if success_s2 and messages_s2 else ""
)
success = success_s1 and success_s2
messages = messages_s1 + messages_s2
tools = tools_s1 + tools_s2
memory_nodes = []
for tool in tools:
if tool.memory_nodes:
memory_nodes.extend(tool.memory_nodes)
return {
"answer": answer,
"success": success,
"messages": messages,
"tools": tools,
"memory_nodes": memory_nodes,
}

View file

@ -11,7 +11,7 @@ metadata:
llms:
default:
backend: openai
# model_name: qwen3-235b-a22b-thinking-2507
# model_name: qwen3-235b-a22b-thinking-2507
model_name: qwen3.5-plus
request_interval: 1
@ -20,12 +20,12 @@ embedding_models:
backend: openai
model_name: text-embedding-v4
dimensions: 1024
enable_cache: true
memory_stores:
file_stores:
default:
backend: chroma
# backend: local
db_name: reme.db
store_name: reme
embedding_model: default
fts_enabled: true
@ -34,7 +34,7 @@ memory_stores:
file_watchers:
default:
backend: full
memory_store: default
file_store: default
watch_paths: [ ".reme", ".reme/memory" ]
suffix_filters: [ ".md" ]
recursive: false

View file

@ -1,26 +1,23 @@
backend: cmd
working_dir: .reme
llms:
default:
backend: openai
# model_name: qwen3-30b-a3b-instruct-2507
# model_name: qwen3-30b-a3b-thinking-2507
model_name: qwen3-235b-a22b-thinking-2507
model_name: qwen3.5-plus
request_interval: 1
# temperature: 0.0001
embedding_models:
default:
backend: openai
model_name: text-embedding-v4
dimensions: 1024
enable_cache: true
memory_stores:
file_stores:
default:
# backend: sqlite
backend: chroma
# backend: local
db_name: reme.db
store_name: reme
embedding_model: default
fts_enabled: true
@ -29,7 +26,7 @@ memory_stores:
file_watchers:
default:
backend: full
memory_store: default
file_store: default
watch_paths: [ ".reme", ".reme/memory" ]
suffix_filters: [ ".md" ]
recursive: false

64
reme/config/service.yaml Normal file
View file

@ -0,0 +1,64 @@
backend: http
working_dir: .reme
thread_pool_max_workers: 64
mcp:
transport: sse
host: "0.0.0.0"
port: 8001
http:
host: "0.0.0.0"
port: 8002
timeout_keep_alive: 600
limit_concurrency: 64
flows:
test:
flow_content: TestOp()
description: "test"
# curl -X POST http://localhost:8002/simple_chat \
# -H "Content-Type: application/json" \
# -d '{
# "query": "hello"
# }'
simple_chat:
flow_content: SimpleChat()
description: "test"
stream_chat:
flow_content: StreamChat()
description: "test"
stream: true
llms:
default:
backend: openai
model_name: qwen3.5-plus
request_interval: 1
# temperature: 0.0001
embedding_models:
default:
backend: openai
model_name: text-embedding-v4
dimensions: 1024
enable_cache: false
vector_stores:
default:
backend: chroma
# backend: local
collection_name: reme
embedding_model: default
token_counters:
default:
backend: base
hf:
backend: hf
model_name: Qwen/Qwen3-Coder-30B-A3B-Instruct
use_mirror: true

View file

@ -1,21 +1,5 @@
backend: http
thread_pool_max_workers: 64
mcp:
transport: sse
host: "0.0.0.0"
port: 8001
http:
host: "0.0.0.0"
port: 8002
timeout_keep_alive: 600
limit_concurrency: 64
flows:
test:
flow_content: TestOp()
description: "test"
backend: cmd
working_dir: .reme
retrieve_task_memory:
flow_content: BuildQuery() >> MemoryRetrieval() >> RerankMemory() >> RewriteMemory()
@ -162,8 +146,7 @@ flows:
llms:
default:
backend: openai
model_name: qwen3-30b-a3b-instruct-2507
# model_name: qwen3-30b-a3b-thinking-2507
model_name: qwen3.5-plus
request_interval: 1
# temperature: 0.0001
@ -184,13 +167,14 @@ embedding_models:
backend: openai
model_name: text-embedding-v4
dimensions: 1024
enable_cache: false
vector_stores:
default:
backend: chroma
# backend: local
embedding_model: default
collection_name: reme
embedding_model: default
token_counters:
default:

View file

@ -1,12 +1,11 @@
"""Core"""
from . import context
from . import embedding
from . import enumeration
from . import file_store
from . import file_watcher
from . import flow
from . import llm
from . import memory_store
from . import op
from . import schema
from . import service
@ -14,22 +13,33 @@ from . import token_counter
from . import utils
from . import vector_store
from .application import Application
from .context import R
from .base_dict import BaseDict
from .prompt_handler import PromptHandler
from .registry_factory import R, Registry, RegistryFactory
from .runtime_context import RuntimeContext
from .service_context import ServiceContext
__all__ = [
"context",
# Submodules
"embedding",
"enumeration",
"file_watcher",
"flow",
"llm",
"memory_store",
"file_store",
"op",
"schema",
"service",
"token_counter",
"utils",
"vector_store",
# Classes
"Application",
"BaseDict",
"PromptHandler",
"R",
"Registry",
"RegistryFactory",
"RuntimeContext",
"ServiceContext",
]

View file

@ -6,15 +6,17 @@ from pathlib import Path
from loguru import logger
from .context import PromptHandler, ServiceContext, R
from .embedding import BaseEmbeddingModel
from .file_store import BaseFileStore
from .file_watcher import BaseFileWatcher
from .flow import BaseFlow
from .llm import BaseLLM
from .memory_store import BaseMemoryStore
from .prompt_handler import PromptHandler
from .registry_factory import R
from .schema import Response, ServiceConfig
from .service_context import ServiceContext
from .token_counter import BaseTokenCounter
from .utils import execute_stream_task, PydanticConfigParser, init_logger, print_logo, MCPClient
from .utils import execute_stream_task, PydanticConfigParser, init_logger, MCPClient, print_logo
from .vector_store import BaseVectorStore
@ -36,7 +38,7 @@ class Application:
default_llm_config: dict | None = None,
default_embedding_model_config: dict | None = None,
default_vector_store_config: dict | None = None,
default_memory_store_config: dict | None = None,
default_file_store_config: dict | None = None,
default_token_counter_config: dict | None = None,
default_file_watcher_config: dict | None = None,
**kwargs,
@ -56,12 +58,17 @@ class Application:
default_llm_config=default_llm_config,
default_embedding_model_config=default_embedding_model_config,
default_vector_store_config=default_vector_store_config,
default_memory_store_config=default_memory_store_config,
default_file_store_config=default_file_store_config,
default_token_counter_config=default_token_counter_config,
default_file_watcher_config=default_file_watcher_config,
**kwargs,
)
self.prompt_handler = PromptHandler(language=self.service_config.language)
# NOTE: flows are initialized here to start service!
self.init_flows()
self._started: bool = False
@classmethod
@ -71,40 +78,8 @@ class Application:
await instance.start()
return instance
@property
def service_config(self) -> ServiceConfig:
"""Get the service configuration."""
return self.service_context.service_config
async def start(self):
"""Start the service context by initializing all configured components."""
if self._started:
logger.warning("Application has already started.")
return self
init_logger(log_to_console=self.service_config.log_to_console)
logger.info(f"Init ReMe with config: {self.service_config.model_dump_json()}")
working_path = Path(self.service_config.working_dir)
working_path.mkdir(parents=True, exist_ok=True)
if self.service_config.enable_logo:
print_logo(service_config=self.service_config)
if self.service_config.ray_max_workers > 1:
import ray
if not ray.is_initialized():
ray.init(num_cpus=self.service_config.ray_max_workers)
if (
self.service_context.thread_pool is None
or self.service_context.thread_pool._shutdown # pylint: disable=protected-access
):
self.service_context.thread_pool = ThreadPoolExecutor(
max_workers=self.service_config.thread_pool_max_workers,
)
def init_flows(self):
"""Initialize flows."""
expression_flow_cls = None
for name, flow_cls in R.flows.items():
if not self._filter_flows(name):
@ -129,12 +104,56 @@ class Application:
else:
logger.info("No expression flow found, please check your configuration.")
def _filter_flows(self, name: str) -> bool:
"""Filter flows based on enabled_flows and disabled_flows configuration."""
if self.service_config.enabled_flows:
return name in self.service_config.enabled_flows
elif self.service_config.disabled_flows:
return name not in self.service_config.disabled_flows
else:
return True
@property
def service_config(self) -> ServiceConfig:
"""Get the service configuration."""
return self.service_context.service_config
async def start(self):
"""Start the service context by initializing all configured components."""
if self._started:
logger.warning("Application has already started.")
return self
init_logger(log_to_console=self.service_config.log_to_console)
logger.info(f"Init ReMe with config: {self.service_config.model_dump_json()}")
working_path = Path(self.service_config.working_dir)
working_path.mkdir(parents=True, exist_ok=True)
if self.service_config.ray_max_workers > 1:
import ray
if not ray.is_initialized():
ray.init(num_cpus=self.service_config.ray_max_workers)
if (
self.service_context.thread_pool is None
or self.service_context.thread_pool._shutdown # pylint: disable=protected-access
):
self.service_context.thread_pool = ThreadPoolExecutor(
max_workers=self.service_config.thread_pool_max_workers,
)
if self.service_context.service_config.enable_logo:
print_logo(service_config=self.service_config)
for name, config in self.service_config.llms.items():
if config.backend not in R.llms:
logger.warning(f"LLM backend {config.backend} is not supported.")
else:
config_dict = config.model_dump(exclude={"backend"})
self.service_context.llms[name] = R.llms[config.backend](**config_dict)
await self.service_context.llms[name].start()
for name, config in self.service_config.embedding_models.items():
if config.backend not in R.embedding_models:
@ -143,6 +162,7 @@ class Application:
config_dict = config.model_dump(exclude={"backend"})
config_dict["cache_dir"] = working_path / "embedding_cache"
self.service_context.embedding_models[name] = R.embedding_models[config.backend](**config_dict)
await self.service_context.embedding_models[name].start()
for name, config in self.service_config.token_counters.items():
if config.backend not in R.token_counters:
@ -159,33 +179,32 @@ class Application:
config_dict.update(
{
"embedding_model": self.service_context.embedding_models[config.embedding_model],
"thread_pool": self.service_context.thread_pool,
"db_path": working_path / "vector_store",
},
)
self.service_context.vector_stores[name] = R.vector_stores[config.backend](**config_dict)
await self.service_context.vector_stores[name].create_collection(config.collection_name)
await self.service_context.vector_stores[name].start()
for name, config in self.service_config.memory_stores.items():
if config.backend not in R.memory_stores:
logger.warning(f"Memory store backend {config.backend} is not supported.")
for name, config in self.service_config.file_stores.items():
if config.backend not in R.file_stores:
logger.warning(f"File store backend {config.backend} is not supported.")
else:
config_dict = config.model_dump(exclude={"backend", "embedding_model"})
config_dict.update(
{
"embedding_model": self.service_context.embedding_models[config.embedding_model],
"thread_pool": self.service_context.thread_pool,
"db_path": working_path / "memory_store",
"db_path": working_path / "file_store",
},
)
self.service_context.memory_stores[name] = R.memory_stores[config.backend](**config_dict)
await self.service_context.memory_stores[name].start()
self.service_context.file_stores[name] = R.file_stores[config.backend](**config_dict)
await self.service_context.file_stores[name].start()
for name, config in self.service_config.file_watchers.items():
if config.backend not in R.file_watchers:
logger.warning(f"File watcher backend {config.backend} is not supported.")
else:
config_dict = config.model_dump(exclude={"backend", "memory_store"})
config_dict["memory_store"] = self.service_context.memory_stores[config.memory_store]
config_dict = config.model_dump(exclude={"backend", "file_store"})
config_dict["file_store"] = self.service_context.file_stores[config.file_store]
self.service_context.file_watchers[name] = R.file_watchers[config.backend](**config_dict)
await self.service_context.file_watchers[name].start()
@ -195,15 +214,6 @@ class Application:
self._started = True
return self
def _filter_flows(self, name: str) -> bool:
"""Filter flows based on enabled_flows and disabled_flows configuration."""
if self.service_config.enabled_flows:
return name in self.service_config.enabled_flows
elif self.service_config.disabled_flows:
return name not in self.service_config.disabled_flows
else:
return True
async def prepare_mcp_servers(self):
"""Prepare and initialize MCP server connections."""
mcp_client = MCPClient(config={"mcpServers": self.service_config.mcp_servers})
@ -228,9 +238,9 @@ class Application:
logger.info(f"Closing vector store: {name}")
await vector_store.close()
for name, memory_store in self.service_context.memory_stores.items():
logger.info(f"Closing memory store: {name}")
await memory_store.close()
for name, file_store in self.service_context.file_stores.items():
logger.info(f"Closing file store: {name}")
await file_store.close()
for name, file_watcher in self.service_context.file_watchers.items():
logger.info(f"Closing file watcher: {name}")
@ -327,13 +337,13 @@ class Application:
return self.service_context.vector_stores.get(name)
@property
def default_memory_store(self) -> BaseMemoryStore:
"""Get the default memory store instance."""
return self.service_context.memory_stores.get("default")
def default_file_store(self) -> BaseFileStore:
"""Get the default file store instance."""
return self.service_context.file_stores.get("default")
def get_memory_store(self, name: str):
"""Get a memory store instance by name."""
return self.service_context.memory_stores.get(name)
def get_file_store(self, name: str):
"""Get a file store instance by name."""
return self.service_context.file_stores.get(name)
@property
def default_file_watcher(self) -> BaseFileWatcher:
@ -358,7 +368,7 @@ class Application:
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)
service = R.services[self.service_config.backend](service_context=self.service_context)
service = R.services[self.service_config.backend](app=self)
service.run()
async def reset_default_collection(self, collection_name: str):

View file

@ -6,7 +6,7 @@ _KT = TypeVar("_KT")
_VT = TypeVar("_VT")
class BaseContext(dict, Generic[_KT, _VT]):
class BaseDict(dict, Generic[_KT, _VT]):
"""A dictionary subclass that enables accessing and modifying keys as attributes."""
def __getattr__(self, name: str) -> _VT:

View file

@ -1,15 +0,0 @@
"""context"""
from .base_context import BaseContext
from .prompt_handler import PromptHandler
from .registry_factory import R
from .runtime_context import RuntimeContext
from .service_context import ServiceContext
__all__ = [
"BaseContext",
"PromptHandler",
"R",
"RuntimeContext",
"ServiceContext",
]

View file

@ -3,7 +3,7 @@
from .base_embedding_model import BaseEmbeddingModel
from .openai_embedding_model import OpenAIEmbeddingModel
from .openai_embedding_model_sync import OpenAIEmbeddingModelSync
from ..context import R
from ..registry_factory import R
__all__ = [
"BaseEmbeddingModel",

View file

@ -14,8 +14,7 @@ from pathlib import Path
from loguru import logger
from ..schema import VectorNode
from ..schema.memory_chunk import MemoryChunk
from ..schema import VectorNode, MemoryChunk
class BaseEmbeddingModel(ABC):
@ -76,9 +75,6 @@ class BaseEmbeddingModel(ABC):
self.cache_path: Path = Path(self.cache_dir)
self.cache_path.mkdir(parents=True, exist_ok=True)
# Load cache from disk if available
self._load_cache()
@property
def api_key(self) -> str | None:
"""Get API key from environment variable."""
@ -531,6 +527,14 @@ class BaseEmbeddingModel(ABC):
logger.warning(f"Mismatch: got {len(embeddings)} vectors for {len(chunks)} chunks")
return chunks
def start_sync(self):
"""Synchronously initialize resources and load cache."""
self._load_cache()
async def start(self):
"""Asynchronously initialize resources and load cache."""
self._load_cache()
def close_sync(self):
"""Synchronously release resources and close connections."""
self._save_cache()

View file

@ -45,6 +45,10 @@ class OpenAIEmbeddingModel(BaseEmbeddingModel):
result_emb[emb.index] = emb.embedding
return result_emb
async def start(self):
"""Initialize the asynchronous OpenAI embedding model and load cache."""
await super().start()
async def close(self):
"""Close the asynchronous OpenAI client and release network resources."""
if self._client is not None:

View file

@ -15,8 +15,8 @@ class RegistryEnum(str, Enum):
# Databases or storage systems for vector search
VECTOR_STORE = "vector_store"
# Databases or storage systems for long-term memory storage
MEMORY_STORE = "memory_store"
# Databases or storage systems for long-term file storage
FILE_STORE = "file_store"
# Atomic operations or functional units
OP = "op"

View file

@ -0,0 +1,23 @@
"""File store module for persistent memory management.
This module provides storage backends for memory chunks and file metadata,
including SQLite-based, ChromaDB-based, and pure-Python local implementations
with vector and full-text search.
"""
from .base_file_store import BaseFileStore
from .chroma_file_store import ChromaFileStore
from .local_file_store import LocalFileStore
from .sqlite_file_store import SqliteFileStore
from ..registry_factory import R
__all__ = [
"BaseFileStore",
"ChromaFileStore",
"LocalFileStore",
"SqliteFileStore",
]
R.file_stores.register("sqlite")(SqliteFileStore)
R.file_stores.register("chroma")(ChromaFileStore)
R.file_stores.register("local")(LocalFileStore)

View file

@ -1,27 +1,22 @@
"""Base storage interface for memory manager."""
"""Base storage interface for file store."""
import asyncio
import re
from abc import ABC, abstractmethod
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from pathlib import Path
from typing import Callable
from ..embedding import BaseEmbeddingModel
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
class BaseMemoryStore(ABC):
"""Abstract base class for memory storage backends."""
class BaseFileStore(ABC):
"""Abstract base class for file storage backends."""
def __init__(
self,
store_name: str,
db_path: str | Path,
thread_pool: ThreadPoolExecutor,
embedding_model: BaseEmbeddingModel,
embedding_model: BaseEmbeddingModel | None = None,
vector_enabled: bool = False,
fts_enabled: bool = True,
**kwargs,
@ -36,11 +31,14 @@ class BaseMemoryStore(ABC):
if not vector_enabled and not fts_enabled:
raise ValueError("At least one of vector_enabled or fts_enabled must be True.")
# Ensure embedding_model is provided when vector search is enabled
if vector_enabled and embedding_model is None:
raise ValueError("embedding_model is required when vector_enabled is True.")
self.store_name: str = store_name
self.db_path: Path = Path(db_path)
self.db_path.mkdir(parents=True, exist_ok=True)
self.thread_pool: ThreadPoolExecutor = thread_pool
self.embedding_model: BaseEmbeddingModel = embedding_model
self.embedding_model: BaseEmbeddingModel | None = embedding_model
self.vector_enabled: bool = vector_enabled
self.fts_enabled: bool = fts_enabled
self.kwargs: dict = kwargs
@ -48,6 +46,8 @@ class BaseMemoryStore(ABC):
@property
def embedding_dim(self) -> int:
"""Get the embedding model's dimensionality."""
if self.embedding_model is None:
return 1024
return self.embedding_model.dimensions
def _get_mock_embedding(self) -> list[float]:
@ -82,11 +82,6 @@ class BaseMemoryStore(ABC):
return chunks
return await self.embedding_model.get_chunk_embeddings(chunks, **kwargs)
async def _run_sync_in_executor(self, sync_func: Callable, *args, **kwargs):
"""Run a synchronous function in the context-defined thread pool executor."""
loop = asyncio.get_running_loop()
return await loop.run_in_executor(self.thread_pool, partial(sync_func, *args, **kwargs)) # noqa
@abstractmethod
async def start(self):
"""Initialize the storage backend."""
@ -115,6 +110,18 @@ class BaseMemoryStore(ABC):
async def get_file_metadata(self, path: str, source: MemorySource) -> FileMetadata | None:
"""Get full file metadata with statistics."""
@abstractmethod
async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None:
"""Update file metadata without affecting chunks.
This is useful for incremental updates where only metadata needs to be updated
(e.g., after adding/removing chunks in delta file watcher).
Args:
file_meta: Updated file metadata (hash, mtime_ms, size, chunk_count)
source: Memory source
"""
@abstractmethod
async def get_file_chunks(self, path: str, source: MemorySource) -> list[MemoryChunk]:
"""Get all chunks for a file."""

View file

@ -1,4 +1,4 @@
"""ChromaDB storage backend for memory index."""
"""ChromaDB storage backend for file store."""
import json
import time
@ -6,7 +6,7 @@ from pathlib import Path
from loguru import logger
from .base_memory_store import BaseMemoryStore
from .base_file_store import BaseFileStore
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
@ -21,10 +21,10 @@ except ImportError:
Settings = None
class ChromaMemoryStore(BaseMemoryStore):
"""ChromaDB memory storage with vector and full-text search.
class ChromaFileStore(BaseFileStore):
"""ChromaDB file storage with vector and full-text search.
Inherits embedding methods from BaseMemoryStore:
Inherits embedding methods from BaseFileStore:
- get_chunk_embedding / get_chunk_embeddings (async)
- get_embedding / get_embeddings (async)
@ -40,7 +40,7 @@ class ChromaMemoryStore(BaseMemoryStore):
):
if not CHROMADB_AVAILABLE:
raise ImportError(
"chromadb package is required for ChromaMemoryStore. Install it with: pip install chromadb",
"chromadb package is required for ChromaFileStore. Install it with: pip install chromadb",
)
super().__init__(**kwargs)
@ -287,6 +287,19 @@ class ChromaMemoryStore(BaseMemoryStore):
return None
return self._metadata_cache[source.value].get(path)
async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None:
"""Update file metadata without affecting chunks."""
if source.value not in self._metadata_cache:
self._metadata_cache[source.value] = {}
self._metadata_cache[source.value][file_meta.path] = FileMetadata(
hash=file_meta.hash,
mtime_ms=file_meta.mtime_ms,
size=file_meta.size,
path=file_meta.path,
chunk_count=file_meta.chunk_count,
)
async def get_file_chunks(
self,
path: str,

View file

@ -1,40 +1,24 @@
"""Pure-Python in-memory storage backend for memory index, with JSON file persistence."""
"""Pure-Python in-memory storage backend for file store, with JSON file persistence."""
import json
import time
from dataclasses import dataclass
from pathlib import Path
import numpy as np
from loguru import logger
from .base_memory_store import BaseMemoryStore
from .base_file_store import BaseFileStore
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
from ..utils.common_utils import cosine_similarity
from ..utils.common_utils import batch_cosine_similarity
@dataclass
class _ChunkRecord:
"""Internal in-memory representation of a stored chunk."""
id: str
path: str
source: str
start_line: int
end_line: int
text: str
hash: str
embedding: list[float] | None
updated_at: int
class LocalMemoryStore(BaseMemoryStore):
"""Pure-Python in-memory memory storage with JSONL file persistence.
class LocalFileStore(BaseFileStore):
"""Pure-Python in-memory file storage with JSONL file persistence.
No external dependencies required. All data lives in Python dicts;
writes are persisted to JSONL files on disk so state survives restarts.
Inherits embedding methods from BaseMemoryStore:
Inherits embedding methods from BaseFileStore:
- get_chunk_embedding / get_chunk_embeddings (async)
- get_embedding / get_embeddings (async)
@ -48,9 +32,9 @@ class LocalMemoryStore(BaseMemoryStore):
super().__init__(**kwargs)
self._started: bool = False
# In-memory indexes
self._chunks: dict[str, _ChunkRecord] = {}
self._chunks: dict[str, MemoryChunk] = {}
self._files: dict[str, dict[str, FileMetadata]] = {} # source -> path -> meta
# Persistence paths (mirror ChromaMemoryStore convention)
# Persistence paths (mirror ChromaFileStore convention)
self._chunks_file: Path = self.db_path / f"{self.store_name}_chunks.jsonl"
self._metadata_file: Path = self.db_path / f"{self.store_name}_file_metadata.json"
@ -69,8 +53,8 @@ class LocalMemoryStore(BaseMemoryStore):
if not line:
continue
rec = json.loads(line)
chunk_id = rec["id"]
self._chunks[chunk_id] = _ChunkRecord(**rec)
chunk = MemoryChunk.model_validate(rec)
self._chunks[chunk.id] = chunk
logger.debug(f"Loaded {len(self._chunks)} chunks from {self._chunks_file}")
except Exception as e:
logger.warning(f"Failed to load chunks from {self._chunks_file}: {e}")
@ -79,18 +63,8 @@ class LocalMemoryStore(BaseMemoryStore):
"""Persist chunks to JSONL file."""
try:
lines = []
for rec in self._chunks.values():
chunk_dict = {
"id": rec.id,
"path": rec.path,
"source": rec.source,
"start_line": rec.start_line,
"end_line": rec.end_line,
"text": rec.text,
"hash": rec.hash,
"embedding": rec.embedding,
"updated_at": rec.updated_at,
}
for chunk in self._chunks.values():
chunk_dict = chunk.model_dump(mode="json")
lines.append(json.dumps(chunk_dict, ensure_ascii=False))
data = "\n".join(lines)
self._chunks_file.write_text(data, encoding="utf-8")
@ -145,7 +119,7 @@ class LocalMemoryStore(BaseMemoryStore):
await self._load_metadata()
await self._load_chunks()
logger.info(
f"LocalMemoryStore '{self.store_name}' ready: "
f"LocalFileStore '{self.store_name}' ready: "
f"{len(self._chunks)} chunks, metadata at {self._metadata_file}",
)
@ -177,19 +151,8 @@ class LocalMemoryStore(BaseMemoryStore):
# Batch generate embeddings (base class returns mock embeddings when vector_enabled=False)
chunks = await self.get_chunk_embeddings(chunks)
now = int(time.time() * 1000)
for chunk in chunks:
self._chunks[chunk.id] = _ChunkRecord(
id=chunk.id,
path=file_meta.path,
source=source.value,
start_line=chunk.start_line,
end_line=chunk.end_line,
text=chunk.text,
hash=chunk.hash,
embedding=chunk.embedding,
updated_at=now,
)
self._chunks[chunk.id] = chunk
if source.value not in self._files:
self._files[source.value] = {}
@ -203,7 +166,7 @@ class LocalMemoryStore(BaseMemoryStore):
async def delete_file(self, path: str, source: MemorySource) -> None:
"""Delete file and all its chunks."""
to_delete = [cid for cid, rec in self._chunks.items() if rec.path == path and rec.source == source.value]
to_delete = [cid for cid, chunk in self._chunks.items() if chunk.path == path and chunk.source == source]
for cid in to_delete:
del self._chunks[cid]
@ -218,10 +181,12 @@ class LocalMemoryStore(BaseMemoryStore):
for cid in chunk_ids:
self._chunks.pop(cid, None)
# Recalculate chunk_count in file metadata
for source_meta in self._files.values():
# Recalculate chunk_count in file metadata (per source)
for source_key, source_meta in self._files.items():
if path in source_meta:
source_meta[path].chunk_count = sum(1 for rec in self._chunks.values() if rec.path == path)
source_meta[path].chunk_count = sum(
1 for chunk in self._chunks.values() if chunk.path == path and chunk.source.value == source_key
)
async def upsert_chunks(
self,
@ -234,19 +199,8 @@ class LocalMemoryStore(BaseMemoryStore):
chunks = await self.get_chunk_embeddings(chunks)
now = int(time.time() * 1000)
for chunk in chunks:
self._chunks[chunk.id] = _ChunkRecord(
id=chunk.id,
path=chunk.path,
source=source.value,
start_line=chunk.start_line,
end_line=chunk.end_line,
text=chunk.text,
hash=chunk.hash,
embedding=chunk.embedding,
updated_at=now,
)
self._chunks[chunk.id] = chunk
# ------------------------------------------------------------------
# Read operations
@ -264,27 +218,28 @@ class LocalMemoryStore(BaseMemoryStore):
"""Get file metadata."""
return self._files.get(source.value, {}).get(path)
async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None:
"""Update file metadata without affecting chunks."""
if source.value not in self._files:
self._files[source.value] = {}
self._files[source.value][file_meta.path] = FileMetadata(
hash=file_meta.hash,
mtime_ms=file_meta.mtime_ms,
size=file_meta.size,
path=file_meta.path,
chunk_count=file_meta.chunk_count,
)
async def get_file_chunks(
self,
path: str,
source: MemorySource,
) -> list[MemoryChunk]:
"""Get all chunks for a file, sorted by start_line."""
records = [rec for rec in self._chunks.values() if rec.path == path and rec.source == source.value]
records.sort(key=lambda r: r.start_line)
return [
MemoryChunk(
id=rec.id,
path=rec.path,
source=MemorySource(rec.source),
start_line=rec.start_line,
end_line=rec.end_line,
text=rec.text,
hash=rec.hash,
embedding=rec.embedding,
)
for rec in records
]
chunks = [chunk for chunk in self._chunks.values() if chunk.path == path and chunk.source == source]
chunks.sort(key=lambda c: c.start_line)
return chunks
# ------------------------------------------------------------------
# Search
@ -304,26 +259,32 @@ class LocalMemoryStore(BaseMemoryStore):
if not query_embedding:
return []
source_values = {s.value for s in sources} if sources else None
results = []
for rec in self._chunks.values():
if source_values and rec.source not in source_values:
continue
if not rec.embedding:
continue
# Collect candidate chunks with embeddings
candidates = [
chunk for chunk in self._chunks.values() if (not sources or chunk.source in sources) and chunk.embedding
]
similarity = cosine_similarity(query_embedding, rec.embedding)
results.append(
MemorySearchResult(
path=rec.path,
start_line=rec.start_line,
end_line=rec.end_line,
score=similarity,
snippet=rec.text,
source=MemorySource(rec.source),
raw_metric=1.0 - similarity, # distance equivalent
),
if not candidates:
return []
# Build embedding matrix and compute similarities in batch
query_array = np.array([query_embedding]) # Shape: (1, emb_size)
chunk_embeddings = np.array([chunk.embedding for chunk in candidates]) # Shape: (n, emb_size)
similarities = batch_cosine_similarity(query_array, chunk_embeddings)[0] # Shape: (n,)
# Build results
results = [
MemorySearchResult(
path=chunk.path,
start_line=chunk.start_line,
end_line=chunk.end_line,
score=float(similarity),
snippet=chunk.text,
source=chunk.source,
raw_metric=1.0 - float(similarity),
)
for chunk, similarity in zip(candidates, similarities)
]
results.sort(key=lambda r: r.score, reverse=True)
return results[:limit]
@ -346,13 +307,12 @@ class LocalMemoryStore(BaseMemoryStore):
words_lower = [w.lower() for w in words]
n_words = len(words)
source_values = {s.value for s in sources} if sources else None
results = []
for rec in self._chunks.values():
if source_values and rec.source not in source_values:
for chunk in self._chunks.values():
if sources and chunk.source not in sources:
continue
text_lower = rec.text.lower()
text_lower = chunk.text.lower()
match_count = sum(1 for w in words_lower if w in text_lower)
if match_count == 0:
continue
@ -364,12 +324,12 @@ class LocalMemoryStore(BaseMemoryStore):
results.append(
MemorySearchResult(
path=rec.path,
start_line=rec.start_line,
end_line=rec.end_line,
path=chunk.path,
start_line=chunk.start_line,
end_line=chunk.end_line,
score=score,
snippet=rec.text,
source=MemorySource(rec.source),
snippet=chunk.text,
source=chunk.source,
),
)
@ -452,18 +412,21 @@ class LocalMemoryStore(BaseMemoryStore):
merged: dict[str, MemorySearchResult] = {}
for result in vector:
result.score = result.score * vector_weight
result.metadata["_weighted_score"] = result.score * vector_weight
merged[result.merge_key] = result
for result in keyword:
key = result.merge_key
if key in merged:
merged[key].score += result.score * text_weight
merged[key].metadata["_weighted_score"] += result.score * text_weight
else:
result.score = result.score * text_weight
result.metadata["_weighted_score"] = result.score * text_weight
merged[key] = result
results = list(merged.values())
for r in results:
r.score = r.metadata.pop("_weighted_score")
results.sort(key=lambda r: r.score, reverse=True)
return results
@ -473,4 +436,4 @@ class LocalMemoryStore(BaseMemoryStore):
self._files.clear()
await self._save_chunks()
await self._save_metadata()
logger.info(f"Cleared all data from LocalMemoryStore '{self.store_name}'")
logger.info(f"Cleared all data from LocalFileStore '{self.store_name}'")

View file

@ -1,4 +1,4 @@
"""SQLite storage backend for memory index."""
"""SQLite storage backend for file store."""
import json
import sqlite3
@ -7,15 +7,15 @@ import time
from loguru import logger
from .base_memory_store import BaseMemoryStore
from .base_file_store import BaseFileStore
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
class SqliteMemoryStore(BaseMemoryStore):
"""SQLite memory storage with vector and full-text search.
class SqliteFileStore(BaseFileStore):
"""SQLite file storage with vector and full-text search.
Inherits embedding methods from BaseMemoryStore:
Inherits embedding methods from BaseFileStore:
- get_chunk_embedding / get_chunk_embeddings (async)
- get_chunk_embedding_sync / get_chunk_embeddings_sync (sync)
- get_embedding / get_embeddings (async)
@ -467,6 +467,24 @@ class SqliteMemoryStore(BaseMemoryStore):
finally:
cursor.close()
async def update_file_metadata(self, file_meta: FileMetadata, source: MemorySource) -> None:
"""Update file metadata without affecting chunks."""
cursor = self.conn.cursor()
try:
cursor.execute(
f"""
INSERT OR REPLACE INTO {self.files_table_name} (path, source, hash, mtime, size)
VALUES (?, ?, ?, ?, ?)
""",
(file_meta.path, source.value, file_meta.hash, file_meta.mtime_ms, file_meta.size),
)
self.conn.commit()
except Exception as e:
logger.error(f"Failed to update file metadata for {file_meta.path}: {e}")
raise
finally:
cursor.close()
async def get_file_chunks(self, path: str, source: MemorySource) -> list[MemoryChunk]:
"""Get all chunks for a file."""
cursor = self.conn.cursor()

View file

@ -7,7 +7,7 @@ and updating memory stores accordingly.
from .base_file_watcher import BaseFileWatcher
from .delta_file_watcher import DeltaFileWatcher
from .full_file_watcher import FullFileWatcher
from ..context import R
from ..registry_factory import R
__all__ = [
"BaseFileWatcher",

View file

@ -13,7 +13,7 @@ from loguru import logger
from watchfiles import awatch, Change
from ..enumeration import MemorySource
from ..memory_store import BaseMemoryStore
from ..file_store import BaseFileStore
class BaseFileWatcher:
@ -32,7 +32,7 @@ class BaseFileWatcher:
debounce: int = 500, # Millisecond debounce
chunk_tokens: int = 400,
chunk_overlap: int = 80,
memory_store: BaseMemoryStore | None = None,
file_store: BaseFileStore | None = None,
callback: Callable[[set[tuple[Change, str]]], None | Coroutine[Any, Any, None]] | None = None,
scan_on_start: bool = False,
**kwargs,
@ -47,7 +47,7 @@ class BaseFileWatcher:
debounce: Debounce time in milliseconds
chunk_tokens: Token size for chunking
chunk_overlap: Overlap size for chunks
memory_store: Memory store instance
file_store: File store instance
callback: Callback function for changes
scan_on_start: If True, scan existing files on start and trigger on_changes with Change.added
**kwargs: Additional keyword arguments
@ -58,7 +58,7 @@ class BaseFileWatcher:
self.debounce: int = debounce
self.chunk_tokens: int = chunk_tokens
self.chunk_overlap: int = chunk_overlap
self.memory_store: BaseMemoryStore = memory_store
self.file_store: BaseFileStore = file_store
self.callback = callback
self.scan_on_start: bool = scan_on_start
self.kwargs: dict = kwargs
@ -140,9 +140,9 @@ class BaseFileWatcher:
else:
logger.info("[SCAN_ON_START] No existing files found matching watch criteria")
files: list[str] = await self.memory_store.list_files(MemorySource.MEMORY)
files: list[str] = await self.file_store.list_files(MemorySource.MEMORY)
for file_path in files:
chunks = await self.memory_store.get_file_chunks(file_path, MemorySource.MEMORY)
chunks = await self.file_store.get_file_chunks(file_path, MemorySource.MEMORY)
logger.info(f"Found existing file: {file_path}, {len(chunks)} chunks")
async def _watch_loop(self):

View file

@ -158,17 +158,17 @@ class DeltaFileWatcher(BaseFileWatcher):
)
if chunks:
chunks = await self.memory_store.get_chunk_embeddings(chunks)
chunks = await self.file_store.get_chunk_embeddings(chunks)
file_meta.chunk_count = len(chunks)
await self.memory_store.upsert_file(file_meta, MemorySource.MEMORY, chunks)
await self.file_store.upsert_file(file_meta, MemorySource.MEMORY, chunks)
logger.info(f"File added: {path} ({len(chunks)} chunks)")
else:
logger.warning(f"No chunks generated for new file {path}")
elif change_type == Change.modified:
# Get existing data
old_chunks = await self.memory_store.get_file_chunks(path, MemorySource.MEMORY)
old_file_meta = await self.memory_store.get_file_metadata(path, MemorySource.MEMORY)
old_chunks = await self.file_store.get_file_chunks(path, MemorySource.MEMORY)
old_file_meta = await self.file_store.get_file_metadata(path, MemorySource.MEMORY)
# Read new file
file_meta = await self._build_file_metadata(path)
@ -187,10 +187,10 @@ class DeltaFileWatcher(BaseFileWatcher):
or []
)
if chunks:
chunks = await self.memory_store.get_chunk_embeddings(chunks)
chunks = await self.file_store.get_chunk_embeddings(chunks)
file_meta.chunk_count = len(chunks)
await self.memory_store.delete_file(path, MemorySource.MEMORY)
await self.memory_store.upsert_file(
await self.file_store.delete_file(path, MemorySource.MEMORY)
await self.file_store.upsert_file(
file_meta,
MemorySource.MEMORY,
chunks,
@ -216,10 +216,10 @@ class DeltaFileWatcher(BaseFileWatcher):
or []
)
if chunks:
chunks = await self.memory_store.get_chunk_embeddings(chunks)
chunks = await self.file_store.get_chunk_embeddings(chunks)
file_meta.chunk_count = len(chunks)
await self.memory_store.delete_file(path, MemorySource.MEMORY)
await self.memory_store.upsert_file(file_meta, MemorySource.MEMORY, chunks)
await self.file_store.delete_file(path, MemorySource.MEMORY)
await self.file_store.upsert_file(file_meta, MemorySource.MEMORY, chunks)
logger.info(f"File modified (full): {path} ({len(chunks)} chunks)")
else:
# Append-only: incremental update
@ -247,16 +247,22 @@ class DeltaFileWatcher(BaseFileWatcher):
f"{chunk.source}:{chunk.path}:{chunk.start_line}:" f"{chunk.end_line}:{chunk.hash}:{idx}",
)
new_chunks = await self.memory_store.get_chunk_embeddings(new_chunks)
new_chunks = await self.file_store.get_chunk_embeddings(new_chunks)
chunks_to_delete = [c.id for c in old_chunks_sorted if c.start_line >= cutoff_line]
# Apply incremental updates
if chunks_to_delete:
await self.memory_store.delete_file_chunks(path, chunks_to_delete)
await self.file_store.delete_file_chunks(path, chunks_to_delete)
if new_chunks:
await self.memory_store.upsert_chunks(new_chunks, MemorySource.MEMORY)
await self.file_store.upsert_chunks(new_chunks, MemorySource.MEMORY)
# Update file metadata to reflect the changes
# Calculate new chunk count: old chunks - deleted + new chunks
new_chunk_count = len(old_chunks) - len(chunks_to_delete) + len(new_chunks)
file_meta.chunk_count = new_chunk_count
await self.file_store.update_file_metadata(file_meta, MemorySource.MEMORY)
logger.info(
f"File modified (incremental): {path} "
@ -265,7 +271,7 @@ class DeltaFileWatcher(BaseFileWatcher):
)
elif change_type == Change.deleted:
await self.memory_store.delete_file(path, MemorySource.MEMORY)
await self.file_store.delete_file(path, MemorySource.MEMORY)
logger.info(f"File deleted: {path}")
else:

View file

@ -58,17 +58,17 @@ class FullFileWatcher(BaseFileWatcher):
or []
)
if chunks:
chunks = await self.memory_store.get_chunk_embeddings(chunks)
chunks = await self.file_store.get_chunk_embeddings(chunks)
file_meta.chunk_count = len(chunks)
await self.memory_store.delete_file(file_meta.path, MemorySource.MEMORY)
await self.file_store.delete_file(file_meta.path, MemorySource.MEMORY)
logger.info(f"delete_file {file_meta.path}")
await self.memory_store.upsert_file(file_meta, MemorySource.MEMORY, chunks)
await self.file_store.upsert_file(file_meta, MemorySource.MEMORY, chunks)
logger.info(f"Upserted {file_meta.chunk_count} chunks for {file_meta.path}")
elif change_type == Change.deleted:
await self.memory_store.delete_file(path, MemorySource.MEMORY)
await self.file_store.delete_file(path, MemorySource.MEMORY)
logger.info(f"Deleted {path}")
else:

View file

@ -3,7 +3,7 @@
from .base_flow import BaseFlow
from .cmd_flow import CmdFlow
from .expression_flow import ExpressionFlow
from ..context import R
from ..registry_factory import R
__all__ = [
"BaseFlow",

View file

@ -7,10 +7,12 @@ from abc import ABC, abstractmethod
from loguru import logger
from ..context import RuntimeContext, ServiceContext, R
from ..enumeration import ChunkEnum
from ..op import BaseOp, SequentialOp, ParallelOp
from ..registry_factory import R
from ..runtime_context import RuntimeContext
from ..schema import Response, ToolCall
from ..service_context import ServiceContext
from ..utils import camel_to_snake, CacheHandler

View file

@ -1,9 +1,9 @@
"""Expression-based flow implementation driven by configuration objects."""
from .base_flow import BaseFlow
from ..context import ServiceContext
from ..op import BaseOp
from ..schema import FlowConfig, ToolCall
from ..service_context import ServiceContext
class ExpressionFlow(BaseFlow):

View file

@ -5,7 +5,7 @@ from .lite_llm import LiteLLM
from .lite_llm_sync import LiteLLMSync
from .openai_llm import OpenAILLM
from .openai_llm_sync import OpenAILLMSync
from ..context import R
from ..registry_factory import R
__all__ = [
"BaseLLM",

View file

@ -10,9 +10,7 @@ from typing import Callable, Generator, AsyncGenerator, Any
from loguru import logger
from ..enumeration import ChunkEnum, Role
from ..schema import Message
from ..schema import StreamChunk
from ..schema import ToolCall
from ..schema import Message, StreamChunk, ToolCall
from ..utils import extract_content
@ -510,8 +508,14 @@ class BaseLLM(ABC):
**kwargs,
)
async def close(self):
"""Release async resources."""
def start_sync(self):
"""Synchronously initialize resources."""
async def start(self):
"""Asynchronously initialize resources."""
def close_sync(self):
"""Release sync resources."""
"""Synchronously release resources and close connections."""
async def close(self):
"""Asynchronously release resources and close connections."""

View file

@ -6,9 +6,7 @@ from loguru import logger
from .base_llm import BaseLLM
from ..enumeration import ChunkEnum
from ..schema import Message
from ..schema import StreamChunk
from ..schema import ToolCall
from ..schema import Message, StreamChunk, ToolCall
class LiteLLM(BaseLLM):

View file

@ -4,9 +4,7 @@ from typing import Generator
from .lite_llm import LiteLLM
from ..enumeration import ChunkEnum
from ..schema import Message
from ..schema import StreamChunk
from ..schema import ToolCall
from ..schema import Message, StreamChunk, ToolCall
class LiteLLMSync(LiteLLM):

View file

@ -7,9 +7,7 @@ from openai import AsyncOpenAI
from .base_llm import BaseLLM
from ..enumeration import ChunkEnum
from ..schema import Message
from ..schema import StreamChunk
from ..schema import ToolCall
from ..schema import Message, StreamChunk, ToolCall
class OpenAILLM(BaseLLM):

View file

@ -6,9 +6,7 @@ from openai import OpenAI
from .openai_llm import OpenAILLM
from ..enumeration import ChunkEnum
from ..schema import Message
from ..schema import StreamChunk
from ..schema import ToolCall
from ..schema import Message, StreamChunk, ToolCall
class OpenAILLMSync(OpenAILLM):

View file

@ -1,23 +0,0 @@
"""Memory store module for persistent memory management.
This module provides storage backends for memory chunks and file metadata,
including SQLite-based, ChromaDB-based, and pure-Python local implementations
with vector and full-text search.
"""
from .base_memory_store import BaseMemoryStore
from .chroma_memory_store import ChromaMemoryStore
from .local_memory_store import LocalMemoryStore
from .sqlite_memory_store import SqliteMemoryStore
from ..context import R
__all__ = [
"BaseMemoryStore",
"ChromaMemoryStore",
"LocalMemoryStore",
"SqliteMemoryStore",
]
R.memory_stores.register("sqlite")(SqliteMemoryStore)
R.memory_stores.register("chroma")(ChromaMemoryStore)
R.memory_stores.register("local")(LocalMemoryStore)

View file

@ -8,7 +8,7 @@ from .base_tool import BaseTool
from .mcp_tool import MCPTool
from .parallel_op import ParallelOp
from .sequential_op import SequentialOp
from ..context import R
from ..registry_factory import R
__all__ = [
"BaseOp",

View file

@ -10,11 +10,14 @@ from typing import Callable, Optional, Any
from loguru import logger
from tqdm import tqdm
from ..context import RuntimeContext, PromptHandler, ServiceContext
from ..embedding import BaseEmbeddingModel
from ..file_store import BaseFileStore
from ..llm import BaseLLM
from ..memory_store import BaseMemoryStore
from ..schema import Response
from ..prompt_handler import PromptHandler
from ..runtime_context import RuntimeContext
from ..schema import Response, ServiceConfig
from ..schema.service_config import OpConfig
from ..service_context import ServiceContext
from ..token_counter import BaseTokenCounter
from ..utils import camel_to_snake, CacheHandler, timer
from ..vector_store import BaseVectorStore
@ -42,7 +45,7 @@ class BaseOp(metaclass=ABCMeta):
llm: str | BaseLLM = "default",
embedding_model: str | BaseEmbeddingModel = "default",
vector_store: str | BaseVectorStore = "default",
memory_store: str | BaseMemoryStore = "default",
file_store: str | BaseFileStore = "default",
token_counter: str | BaseTokenCounter = "default",
enable_cache: bool = False,
cache_path: str = "cache/op",
@ -64,7 +67,7 @@ class BaseOp(metaclass=ABCMeta):
self._llm = llm
self._embedding_model = embedding_model
self._vector_store = vector_store
self._memory_store = memory_store
self._file_store = file_store
self._token_counter = token_counter
self.enable_cache = enable_cache
@ -121,6 +124,11 @@ class BaseOp(metaclass=ABCMeta):
assert self.context, "Service context is not initialized!"
return self.context.service_context
@property
def service_config(self) -> ServiceConfig:
"""Access the service configuration."""
return self.service_context.service_config
@property
def llm(self) -> BaseLLM:
"""Get the LLM instance from ServiceContext."""
@ -143,11 +151,11 @@ class BaseOp(metaclass=ABCMeta):
return self._vector_store
@property
def memory_store(self) -> BaseMemoryStore:
"""Lazily initialize and return the memory store instance."""
if isinstance(self._memory_store, str):
self._memory_store = self.service_context.memory_stores[self._memory_store]
return self._memory_store
def file_store(self) -> BaseFileStore:
"""Lazily initialize and return the file store instance."""
if isinstance(self._file_store, str):
self._file_store = self.service_context.file_stores[self._file_store]
return self._file_store
@property
def token_counter(self) -> BaseTokenCounter:
@ -159,7 +167,7 @@ class BaseOp(metaclass=ABCMeta):
@property
def service_metadata(self) -> dict:
"""Get service configuration metadata."""
return self.service_context.service_config.model_extra
return self.service_context.service_config.metadata
@property
def response(self) -> Response:
@ -167,12 +175,42 @@ class BaseOp(metaclass=ABCMeta):
return self.context.response
def before_execute_sync(self):
"""Prepare context and validate before sync execution."""
"""Prepare context and validate before sync execution.
This method performs the following steps:
1. Apply input mapping to transform context variables
2. Load operator-specific configuration from service config if available
3. Override operator parameters and prompts based on config
"""
self.context.apply_mapping(self.input_mapping)
if self.context.service_context is None:
return
service_config = self.service_context.service_config
if self.name not in service_config.ops:
return
op_config: OpConfig = service_config.ops[self.name]
# Override operator parameters from config
if op_config.params:
for k, v in op_config.params.items():
if hasattr(self, k):
setattr(self, k, v)
logger.info(f"[{self.__class__.__name__}] Set attribute '{k}' = {v}")
else:
self.op_params[k] = v
logger.info(f"[{self.__class__.__name__}] Set op_param '{k}' = {v}")
# Load custom prompt templates from config
if op_config.prompt_dict:
self.prompt.load_prompt_dict(op_config.prompt_dict)
logger.info(f"[{self.__class__.__name__}] Loaded prompt keys={list(op_config.prompt_dict.keys())}")
async def before_execute(self):
"""Prepare context and validate before async execution."""
self.context.apply_mapping(self.input_mapping)
self.before_execute_sync()
def execute_sync(self):
"""Define core sync logic in subclasses."""

View file

@ -8,7 +8,7 @@ from loguru import logger
from tqdm import tqdm
from .base_op import BaseOp
from ..context import BaseContext
from ..base_dict import BaseDict
_RAY_IMPORT_ERROR = None
@ -53,7 +53,7 @@ class BaseRayOp(BaseOp, metaclass=ABCMeta):
# Put large shared objects into the Ray Object Store once
optimized_kwargs = {
k: (ray.put(v) if isinstance(v, (pd.DataFrame, pd.Series, dict, list, BaseContext)) else v)
k: (ray.put(v) if isinstance(v, (pd.DataFrame, pd.Series, dict, list, BaseDict)) else v)
for k, v in kwargs.items()
}

View file

@ -32,7 +32,14 @@ class MCPTool(BaseTool):
self.timeout: float | None = timeout
# Example MCP marketplace: https://bailian.console.aliyun.com/?tab=mcp#/mcp-market
self._client = MCPClient(self.service_context.service_config.mcp_servers)
self._client: MCPClient | None = None
@property
def client(self) -> MCPClient:
"""Lazily initialize and return the MCP client."""
if self._client is None:
self._client = MCPClient(self.service_context.service_config.mcp_servers)
return self._client
def _build_tool_call(self) -> ToolCall:
tool_call_dict = self.service_context.mcp_server_mapping[self.mcp_server]
@ -61,7 +68,7 @@ class MCPTool(BaseTool):
return tool_call
async def execute(self):
tool_result: CallToolResult = await self._client.call_tool(
tool_result: CallToolResult = await self.client.call_tool(
server_name=self.mcp_server,
tool_name=self.tool_name,
arguments=self.input_dict,

View file

@ -8,15 +8,16 @@ from typing import Any, Dict, Optional, Union
import yaml
from loguru import logger
from .base_context import BaseContext
from .base_dict import BaseDict
class PromptHandler(BaseContext):
class PromptHandler(BaseDict):
"""A context-aware handler for loading, retrieving, and formatting prompt templates."""
def __init__(self, language: str = "", **kwargs):
super().__init__(**kwargs)
self.language: str = language.strip()
# Use object.__setattr__ to avoid storing 'language' in the dict
object.__setattr__(self, "language", language.strip())
def load_prompt_by_file(
self,

View file

@ -3,13 +3,13 @@
import inspect
from typing import Callable, TypeVar
from .base_context import BaseContext
from ..utils import singleton
from .base_dict import BaseDict
from .utils import singleton
T = TypeVar("T")
class Registry(BaseContext):
class Registry(BaseDict):
"""A registry container that uses decorators to map and store class references."""
def register(self, name: str | type = "") -> Callable[[type[T]], type[T]] | type[T]:
@ -36,7 +36,7 @@ class RegistryFactory:
self.llms = Registry()
self.embedding_models = Registry()
self.vector_stores = Registry()
self.memory_stores = Registry()
self.file_stores = Registry()
self.ops = Registry()
self.flows = Registry()
self.services = Registry()

View file

@ -2,13 +2,13 @@
import asyncio
from .base_context import BaseContext
from .base_dict import BaseDict
from .enumeration import ChunkEnum
from .schema import Response, StreamChunk
from .service_context import ServiceContext
from ..enumeration import ChunkEnum
from ..schema import Response, StreamChunk
class RuntimeContext(BaseContext):
class RuntimeContext(BaseDict):
"""Context for execution state, response metadata, and stream queues."""
def __init__(

View file

@ -11,10 +11,12 @@ from .response import Response
from .service_config import (
CmdConfig,
EmbeddingModelConfig,
FileWatcherConfig,
FlowConfig,
HttpConfig,
LLMConfig,
MCPConfig,
FileStoreConfig,
ServiceConfig,
TokenCounterConfig,
VectorStoreConfig,
@ -30,6 +32,7 @@ __all__ = [
"ContentBlock",
"EmbeddingModelConfig",
"FileMetadata",
"FileWatcherConfig",
"FlowConfig",
"HttpConfig",
"LLMConfig",
@ -37,6 +40,7 @@ __all__ = [
"MemoryChunk",
"MemoryNode",
"MemorySearchResult",
"FileStoreConfig",
"Message",
"Request",
"Response",

View file

@ -78,6 +78,14 @@ class Message(BaseModel):
return self.content
return [block.simple_dump() for block in self.content]
def get_text_content(self) -> str:
"""Extract plain text content from message, handling both str and list[ContentBlock]."""
if isinstance(self.content, str):
return self.content
return " ".join(
block.content if isinstance(block.content, str) else str(block.content) for block in self.content
)
def simple_dump(
self,
add_name: bool = False,

View file

@ -36,6 +36,15 @@ class CmdConfig(BaseModel):
flow: str = Field(default="")
class OpConfig(BaseModel):
"""Configuration for op settings and parameters."""
model_config = ConfigDict(extra="allow")
prompt_dict: dict[str, str] = Field(default_factory=dict)
params: dict = Field(default_factory=dict)
class FlowConfig(ToolCall):
"""Configuration for workflow execution, caching, and error handling."""
@ -77,8 +86,8 @@ class VectorStoreConfig(BaseModel):
embedding_model: str = Field(default="default")
class MemoryStoreConfig(BaseModel):
"""Configuration for memory database storage and associated embeddings."""
class FileStoreConfig(BaseModel):
"""Configuration for file store database storage and associated embeddings."""
model_config = ConfigDict(extra="allow")
@ -102,7 +111,7 @@ class FileWatcherConfig(BaseModel):
model_config = ConfigDict(extra="allow")
backend: str = Field(default="")
memory_store: str = Field(default="")
file_store: str = Field(default="")
watch_paths: list[str] = Field(default_factory=list)
@ -126,11 +135,12 @@ class ServiceConfig(BaseModel):
mcp: MCPConfig = Field(default_factory=MCPConfig)
http: HttpConfig = Field(default_factory=HttpConfig)
cmd: CmdConfig = Field(default_factory=CmdConfig)
ops: dict[str, OpConfig] = Field(default_factory=dict)
flows: dict[str, FlowConfig] = Field(default_factory=dict)
llms: dict[str, LLMConfig] = Field(default_factory=dict)
embedding_models: dict[str, EmbeddingModelConfig] = Field(default_factory=dict)
vector_stores: dict[str, VectorStoreConfig] = Field(default_factory=dict)
memory_stores: dict[str, MemoryStoreConfig] = Field(default_factory=dict)
file_stores: dict[str, FileStoreConfig] = Field(default_factory=dict)
token_counters: dict[str, TokenCounterConfig] = Field(default_factory=dict)
file_watchers: dict[str, FileWatcherConfig] = Field(default_factory=dict)

View file

@ -4,7 +4,7 @@ from .base_service import BaseService
from .cmd_service import CmdService
from .http_service import HttpService
from .mcp_service import MCPService
from ..context import R
from ..registry_factory import R
__all__ = [
"BaseService",

View file

@ -1,22 +1,26 @@
"""Base service definitions for flow management."""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
from loguru import logger
from pydantic import BaseModel
from ..context import ServiceContext
from ..flow import BaseFlow
from ..schema import ToolCall
from ..utils import create_pydantic_model
if TYPE_CHECKING:
from ..application import Application
class BaseService(ABC):
"""Abstract base class for services that integrate and execute flows."""
def __init__(self, service_context: ServiceContext, **kwargs):
def __init__(self, app: "Application", **kwargs):
"""Initialize the base service."""
self.service_context: ServiceContext = service_context
self.app: "Application" = app
self.service_context = self.app.service_context
self.service_config = self.service_context.service_config
self.kwargs = kwargs

View file

@ -4,7 +4,7 @@ from loguru import logger
from .base_service import BaseService
from ..flow import CmdFlow, BaseFlow
from ..utils.common_utils import run_coro_safely
from ..utils import run_coro_safely
class CmdService(BaseService):
@ -18,10 +18,14 @@ class CmdService(BaseService):
def integrate_flow(self, flow: BaseFlow) -> str | None:
"""Integrate the workflow configuration into the command service."""
self._cmd_flow = CmdFlow(flow=self.service_config.cmd.flow, service_context=self.service_context)
return self._cmd_flow.tool_call.name if self._cmd_flow else None
def run(self):
"""Execute the command flow in either asynchronous or synchronous mode."""
super().run()
if not self._cmd_flow:
logger.warning("No command flow configured, skipping execution")
return
kwargs = self.service_config.cmd.model_extra
if self._cmd_flow.async_mode:

View file

@ -12,7 +12,7 @@ from fastapi.responses import StreamingResponse
from .base_service import BaseService
from ..flow import BaseFlow
from ..schema import Response
from ..utils.common_utils import execute_stream_task
from ..utils import execute_stream_task
class HttpService(BaseService):
@ -24,20 +24,20 @@ class HttpService(BaseService):
@asynccontextmanager
async def lifespan(_: FastAPI):
await self.service_context.start()
await self.app.start()
yield
await self.service_context.close()
await self.app.close()
self.app = FastAPI(title=self.service_config.app_name, lifespan=lifespan)
self.http_service = FastAPI(title=self.service_config.app_name, lifespan=lifespan)
self.app.add_middleware(
self.http_service.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
self.app.get("/health")(lambda: {"status": "healthy"})
self.http_service.get("/health")(lambda: {"status": "healthy"})
def _integrate_flow(self, flow: BaseFlow) -> str:
"""Register a standard flow as a POST endpoint."""
@ -46,7 +46,7 @@ class HttpService(BaseService):
async def execute_endpoint(request: request_model) -> Response:
return await flow.call(**request.model_dump(exclude_none=True))
self.app.post(
self.http_service.post(
path=f"/{tool_call.name}",
response_model=Response,
description=tool_call.description,
@ -72,7 +72,7 @@ class HttpService(BaseService):
return StreamingResponse(generate_stream(), media_type="text/event-stream")
self.app.post(f"/{tool_call.name}")(execute_stream_endpoint)
self.http_service.post(f"/{tool_call.name}")(execute_stream_endpoint)
return tool_call.name
def integrate_flow(self, flow: BaseFlow) -> str | None:
@ -84,7 +84,7 @@ class HttpService(BaseService):
super().run()
cfg = self.service_config.http
uvicorn.run(
self.app,
self.http_service,
host=cfg.host,
port=cfg.port,
timeout_keep_alive=cfg.timeout_keep_alive,

View file

@ -18,11 +18,11 @@ class MCPService(BaseService):
@asynccontextmanager
async def lifespan(_: FastMCP):
await self.service_context.start()
await self.app.start()
yield {}
await self.service_context.close()
await self.app.close()
self.mcp = FastMCP(name=self.service_config.app_name, lifespan=lifespan)
self.mcp_service = FastMCP(name=self.service_config.app_name, lifespan=lifespan)
def integrate_flow(self, flow: BaseFlow) -> str | None:
"""Register a non-streaming flow as an MCP tool."""
@ -37,7 +37,7 @@ class MCPService(BaseService):
response = await flow.call(**request_instance.model_dump(exclude_none=True))
return response.answer
self.mcp.add_tool(
self.mcp_service.add_tool(
FunctionTool(
name=tool_call.name, # noqa
description=tool_call.description, # noqa
@ -50,8 +50,8 @@ class MCPService(BaseService):
def run(self):
"""Run the MCP server with specified transport protocol."""
super().run()
cfg = self.service_config.mcp
cfg = self.service_config.mcp_service
run_args: dict = {"transport": cfg.transport, "show_banner": False, **cfg.model_extra}
if cfg.transport != "stdio":
run_args.update({"host": cfg.host, "port": cfg.port})
self.mcp.run(**run_args)
self.mcp_service.run(**run_args)

View file

@ -6,21 +6,21 @@ from typing import TYPE_CHECKING
from loguru import logger
from .base_context import BaseContext
from ..schema import ServiceConfig
from ..utils import load_env, PydanticConfigParser
from .base_dict import BaseDict
from .schema import ServiceConfig
from .utils import load_env, PydanticConfigParser
if TYPE_CHECKING:
from ..llm import BaseLLM
from ..embedding import BaseEmbeddingModel
from ..vector_store import BaseVectorStore
from ..memory_store import BaseMemoryStore
from ..token_counter import BaseTokenCounter
from ..flow import BaseFlow
from ..file_watcher import BaseFileWatcher
from .llm import BaseLLM
from .embedding import BaseEmbeddingModel
from .vector_store import BaseVectorStore
from .file_store import BaseFileStore
from .token_counter import BaseTokenCounter
from .flow import BaseFlow
from .file_watcher import BaseFileWatcher
class ServiceContext(BaseContext):
class ServiceContext(BaseDict):
"""Service context."""
def __init__(
@ -39,7 +39,7 @@ class ServiceContext(BaseContext):
default_llm_config: dict | None = None,
default_embedding_model_config: dict | None = None,
default_vector_store_config: dict | None = None,
default_memory_store_config: dict | None = None,
default_file_store_config: dict | None = None,
default_token_counter_config: dict | None = None,
default_file_watcher_config: dict | None = None,
**kwargs,
@ -57,7 +57,7 @@ class ServiceContext(BaseContext):
if service_config is None:
parser_class = parser if parser is not None else PydanticConfigParser
parser = parser_class(ServiceConfig)
parser_instance = parser_class(ServiceConfig)
input_args = []
if config_path:
input_args.append(f"config={config_path}")
@ -72,8 +72,8 @@ class ServiceContext(BaseContext):
self._update_section_config(kwargs, "token_counters", **default_token_counter_config)
if default_vector_store_config:
self._update_section_config(kwargs, "vector_stores", **default_vector_store_config)
if default_memory_store_config:
self._update_section_config(kwargs, "memory_stores", **default_memory_store_config)
if default_file_store_config:
self._update_section_config(kwargs, "file_stores", **default_file_store_config)
if default_file_watcher_config:
self._update_section_config(kwargs, "file_watchers", **default_file_watcher_config)
@ -85,7 +85,7 @@ class ServiceContext(BaseContext):
},
)
logger.info(f"update with args: {input_args} kwargs: {kwargs}")
service_config = parser.parse_args(*input_args, **kwargs)
service_config = parser_instance.parse_args(*input_args, **kwargs)
self.service_config: ServiceConfig = service_config
@ -94,7 +94,7 @@ class ServiceContext(BaseContext):
self.embedding_models: dict[str, "BaseEmbeddingModel"] = {}
self.token_counters: dict[str, "BaseTokenCounter"] = {}
self.vector_stores: dict[str, "BaseVectorStore"] = {}
self.memory_stores: dict[str, "BaseMemoryStore"] = {}
self.file_stores: dict[str, "BaseFileStore"] = {}
self.file_watchers: dict[str, "BaseFileWatcher"] = {}
self.flows: dict[str, "BaseFlow"] = {}
self.mcp_server_mapping: dict[str, dict] = {}

View file

@ -3,7 +3,7 @@
from .base_token_counter import BaseTokenCounter
from .hf_token_counter import HFTokenCounter
from .openai_token_counter import OpenAITokenCounter
from ..context import R
from ..registry_factory import R
__all__ = [
"BaseTokenCounter",

View file

@ -39,10 +39,7 @@ class BaseTokenCounter:
# Extract text from messages
segments = []
for msg in messages:
content = msg.content
if isinstance(content, bytes):
content = content.decode("utf-8", errors="ignore")
segments.extend([content, msg.reasoning_content])
segments.extend([msg.get_text_content(), msg.reasoning_content])
# Extract text from tools
if tools:

View file

@ -42,7 +42,7 @@ class OpenAITokenCounter(BaseTokenCounter):
# Every message has <|start|>{role/name}\n{content}<|end|>\n
total_tokens += 3 # Base overhead per message
if msg.content:
total_tokens += len(enc.encode(msg.content))
total_tokens += len(enc.encode(msg.get_text_content()))
if msg.tool_calls:
for tc in msg.tool_calls:

View file

@ -0,0 +1,45 @@
"""tools"""
from .execute_code import ExecuteCode
from .execute_shell import ExecuteShell
# file tools
from .file.base_file_tool import BaseFileTool
from .file.bash_tool import BashTool
from .file.edit_tool import EditTool
from .file.find_tool import FindTool
from .file.grep_tool import GrepTool
from .file.ls_tool import LsTool
from .file.read_tool import ReadTool
from .file.write_tool import WriteTool
# search tools
from .search.dashscope_search import DashscopeSearch
from .search.mock_search import MockSearch
from .search.tavily_search import TavilySearch
from .think_tool import ThinkTool
from ..registry_factory import R
__all__ = [
# base tools
"ThinkTool",
"ExecuteCode",
"ExecuteShell",
# file tools
"BaseFileTool",
"BashTool",
"EditTool",
"FindTool",
"GrepTool",
"LsTool",
"ReadTool",
"WriteTool",
# search tools
"DashscopeSearch",
"TavilySearch",
"MockSearch",
]
for name in __all__:
tool_class = globals()[name]
R.ops.register(tool_class)

View file

@ -4,10 +4,9 @@ This module provides an operation that can execute Python code strings
and return the output or error messages.
"""
from ...core.op import BaseTool
from ...core.schema import ToolCall
from ...core.utils import exec_code, async_exec_code
from ..op import BaseTool
from ..schema import ToolCall
from ..utils import exec_code, async_exec_code
class ExecuteCode(BaseTool):

View file

@ -4,10 +4,9 @@ This module provides an operation that can execute shell commands
asynchronously and return the output, error, and exit code.
"""
from ...core.op import BaseTool
from ...core.schema import ToolCall
from ...core.utils import run_shell_command
from ..op import BaseTool
from ..schema import ToolCall
from ..utils import run_shell_command
class ExecuteShell(BaseTool):

View file

@ -2,11 +2,11 @@
from loguru import logger
from ...core.context import RuntimeContext
from ...core.op import BaseTool
from ...op import BaseTool
from ...runtime_context import RuntimeContext
class BaseFsTool(BaseTool):
class BaseFileTool(BaseTool):
"""Base class for file system tools.
Features:
@ -34,6 +34,7 @@ class BaseFsTool(BaseTool):
response = await self.execute()
response = await self.after_execute(response)
return response
except Exception as e:
# Return error message to LLM instead of raising
error_msg = f"{self.__class__.__name__} failed: {str(e)}"

View file

@ -11,9 +11,9 @@ import platform
import signal
from pathlib import Path
from .base_fs_tool import BaseFsTool
from .base_file_tool import BaseFileTool
from .truncate import DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, truncate_tail
from ...core.schema import ToolCall, TruncationResult
from ...schema import ToolCall, TruncationResult
def get_shell_config() -> tuple[str, list[str]]:
@ -53,7 +53,7 @@ def kill_process_tree(pid: int) -> None:
pass # Best effort
class BashTool(BaseFsTool):
class BashTool(BaseFileTool):
"""Production-grade tool for executing bash commands.
Features:

View file

@ -3,7 +3,7 @@
import os
from pathlib import Path
from .base_fs_tool import BaseFsTool
from .base_file_tool import BaseFileTool
from .edit_diff import (
detect_line_ending,
fuzzy_find_text,
@ -13,10 +13,10 @@ from .edit_diff import (
restore_line_endings,
strip_bom,
)
from ...core.schema import ToolCall
from ...schema import ToolCall
class EditTool(BaseFsTool):
class EditTool(BaseFileTool):
"""Edit a file by replacing exact text."""
def __init__(self, cwd: str | None = None):

View file

@ -3,12 +3,12 @@
import os
from pathlib import Path
from .base_fs_tool import BaseFsTool
from .base_file_tool import BaseFileTool
from .truncate import FIND_MAX_BYTES, FIND_MAX_LINES, format_size, truncate_head
from ...core.schema import ToolCall
from ...schema import ToolCall
class FindTool(BaseFsTool):
class FindTool(BaseFileTool):
"""Search for files by glob pattern, respecting .gitignore."""
def __init__(self, cwd: str | None = None):

View file

@ -13,7 +13,7 @@ import os
import shutil
from pathlib import Path
from .base_fs_tool import BaseFsTool
from .base_file_tool import BaseFileTool
from .truncate import (
DEFAULT_MAX_BYTES,
GREP_MAX_LINE_LENGTH,
@ -21,13 +21,13 @@ from .truncate import (
truncate_head,
truncate_line,
)
from ...core.schema import ToolCall
from ...schema import ToolCall
# Default limits
DEFAULT_LIMIT = 100 # Maximum number of matches
class GrepTool(BaseFsTool):
class GrepTool(BaseFileTool):
"""Tool for searching file contents using ripgrep.
Features:

View file

@ -3,14 +3,14 @@
import os
from pathlib import Path
from .base_fs_tool import BaseFsTool
from .base_file_tool import BaseFileTool
from .truncate import DEFAULT_MAX_BYTES, truncate_head
from ...core.schema import ToolCall
from ...schema import ToolCall
DEFAULT_LIMIT = 500
class LsTool(BaseFsTool):
class LsTool(BaseFileTool):
"""List directory contents with smart truncation.
Features:

View file

@ -9,9 +9,9 @@ Features:
import os
from pathlib import Path
from .base_fs_tool import BaseFsTool
from .base_file_tool import BaseFileTool
from .truncate import DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, format_size, truncate_head
from ...core.schema import ToolCall
from ...schema import ToolCall
# Supported image extensions
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp"}
@ -29,7 +29,7 @@ def is_image_file(path: str) -> bool:
return Path(path).suffix.lower() in IMAGE_EXTENSIONS
class ReadTool(BaseFsTool):
class ReadTool(BaseFileTool):
"""Read file contents with smart truncation.
Features:

View file

@ -2,7 +2,7 @@
from typing import Literal
from ...core.schema import TruncationResult
from ...schema import TruncationResult
# Default limits for output truncation
DEFAULT_MAX_LINES = 1000 # Maximum lines to keep for tail truncation

View file

@ -8,11 +8,11 @@ This module provides a tool for writing content to files with:
import os
from .base_fs_tool import BaseFsTool
from ...core.schema import ToolCall
from .base_file_tool import BaseFileTool
from ...schema import ToolCall
class WriteTool(BaseFsTool):
class WriteTool(BaseFileTool):
"""Tool for writing content to files.
Features:

View file

@ -9,8 +9,8 @@ from typing import Literal
from loguru import logger
from ...core.op import BaseTool
from ...core.schema import ToolCall
from ...op import BaseTool
from ...schema import ToolCall
class DashscopeSearch(BaseTool):

View file

@ -9,10 +9,10 @@ import random
from loguru import logger
from ...core.enumeration import Role
from ...core.op import BaseTool
from ...core.schema import ToolCall, Message
from ...core.utils import extract_content
from ...enumeration import Role
from ...op import BaseTool
from ...schema import ToolCall, Message
from ...utils import extract_content
class MockSearch(BaseTool):

View file

@ -9,8 +9,8 @@ import os
from loguru import logger
from ...core.op import BaseTool
from ...core.schema import ToolCall
from ...op import BaseTool
from ...schema import ToolCall
class TavilySearch(BaseTool):

View file

@ -4,8 +4,8 @@ This module provides a tool that prompts the model for explicit reflection
before taking actions, helping agents reason about their next steps.
"""
from ...core.op import BaseTool
from ...core.schema import ToolCall
from ..op import BaseTool
from ..schema import ToolCall
class ThinkTool(BaseTool):

View file

@ -54,8 +54,6 @@ def load_env(path: str | Path | None = None, enable_log: bool = True) -> None:
_ENV_LOADED = True
return
logger.warning(".env file not found in search path")
def reset_env_flag() -> None:
"""Reset the internal load state flag."""

View file

@ -1,17 +1,21 @@
"""Module providing a decorator to implement the Singleton design pattern."""
import threading
def singleton(cls):
"""A class decorator that ensures only one instance of a class exists."""
# Dictionary to cache the single instance of the class
_instance = {}
_lock = threading.Lock()
def _singleton(*args, **kwargs):
"""Return the existing instance or create a new one if it doesn't exist."""
if cls not in _instance:
# Create and store the instance if it's the first call
_instance[cls] = cls(*args, **kwargs)
with _lock:
if cls not in _instance:
# Create and store the instance if it's the first call
_instance[cls] = cls(*args, **kwargs)
return _instance[cls]
return _singleton

View file

@ -6,7 +6,7 @@ from .es_vector_store import ESVectorStore
from .local_vector_store import LocalVectorStore
from .pgvector_store import PGVectorStore
from .qdrant_vector_store import QdrantVectorStore
from ..context import R
from ..registry_factory import R
__all__ = [
"BaseVectorStore",

View file

@ -1,10 +1,7 @@
"""Base vector store interface for managing vector embeddings and similarity search."""
import asyncio
from abc import ABC, abstractmethod
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from pathlib import Path
from ..embedding import BaseEmbeddingModel
from ..schema import VectorNode
@ -16,21 +13,16 @@ class BaseVectorStore(ABC):
def __init__(
self,
collection_name: str,
db_path: str | Path,
embedding_model: BaseEmbeddingModel,
thread_pool: ThreadPoolExecutor,
**kwargs,
):
"""Initialize the vector store with a collection name and an embedding model."""
self.collection_name: str = collection_name
self.db_path: Path = Path(db_path)
self.embedding_model: BaseEmbeddingModel = embedding_model
self.thread_pool: ThreadPoolExecutor = thread_pool
self.kwargs: dict = kwargs
async def _run_sync_in_executor(self, sync_func: Callable, *args, **kwargs):
"""Run a synchronous function in the context-defined thread pool executor."""
loop = asyncio.get_running_loop()
return await loop.run_in_executor(self.thread_pool, partial(sync_func, *args, **kwargs)) # noqa
async def get_node_embedding(self, node: VectorNode) -> VectorNode:
"""Generate and assign embedding for a single vector node."""
return await self.embedding_model.get_node_embedding(node)
@ -119,5 +111,13 @@ class BaseVectorStore(ABC):
reverse: If True, sort in descending order; if False, sort in ascending order
"""
async def start(self) -> None:
"""Initialize the vector store and ensure the collection exists.
This method should be called after instantiation to perform async initialization.
Subclasses should call super().start() to ensure collection creation.
"""
await self.create_collection(self.collection_name)
async def close(self) -> None:
"""Release resources and close active connections to the vector store."""

View file

@ -1,6 +1,6 @@
"""ChromaDB vector store implementation for the ReMe framework."""
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any
from loguru import logger
@ -26,12 +26,11 @@ class ChromaVectorStore(BaseVectorStore):
def __init__(
self,
collection_name: str,
db_path: str | Path,
embedding_model: BaseEmbeddingModel,
thread_pool: ThreadPoolExecutor,
client: chromadb.ClientAPI | None = None,
host: str | None = None,
port: int | None = None,
path: str | None = None,
api_key: str | None = None,
tenant: str | None = None,
database: str | None = None,
@ -45,13 +44,14 @@ class ChromaVectorStore(BaseVectorStore):
super().__init__(
collection_name=collection_name,
db_path=db_path,
embedding_model=embedding_model,
thread_pool=thread_pool,
**kwargs,
)
self.client: chromadb.ClientAPI
self.collection: chromadb.Collection
self.is_local = client is None and not (api_key and tenant) and not (host and port)
if client:
self.client = client
@ -66,18 +66,9 @@ class ChromaVectorStore(BaseVectorStore):
logger.info(f"Initializing ChromaDB HTTP client at {host}:{port}")
self.client = chromadb.HttpClient(host=host, port=port)
else:
if path is None:
path = "./chroma_vector_store"
logger.info(f"Initializing local ChromaDB at {path}")
self.client = chromadb.PersistentClient(
path=path,
settings=Settings(anonymized_telemetry=False),
)
self.client = None # Will be initialized in start()
self.collection = self.client.get_or_create_collection(
name=collection_name,
metadata={"hnsw:space": "cosine"},
)
self.collection: chromadb.Collection | None = None
@staticmethod
def _parse_results(
@ -213,63 +204,47 @@ class ChromaVectorStore(BaseVectorStore):
async def list_collections(self) -> list[str]:
"""Retrieve a list of all existing collection names."""
def _list():
return [col.name for col in self.client.list_collections()]
return await self._run_sync_in_executor(_list)
return [col.name for col in self.client.list_collections()]
async def create_collection(self, collection_name: str, **kwargs):
"""Create a new collection with specified distance metrics and metadata."""
def _create():
distance_metric = kwargs.get("distance_metric", "cosine")
metadata = kwargs.get("metadata", {})
metadata["hnsw:space"] = distance_metric
return self.client.get_or_create_collection(name=collection_name, metadata=metadata)
new_collection = await self._run_sync_in_executor(_create)
distance_metric = kwargs.get("distance_metric", "cosine")
metadata = kwargs.get("metadata", {})
metadata["hnsw:space"] = distance_metric
new_collection = self.client.get_or_create_collection(name=collection_name, metadata=metadata)
if collection_name == self.collection_name:
self.collection = new_collection
logger.info(f"Created collection `{collection_name}`")
async def delete_collection(self, collection_name: str, **kwargs):
"""Delete a specified collection from the database."""
def _delete():
try:
self.client.delete_collection(name=collection_name)
return True
except Exception as _e:
logger.warning(f"Failed to delete collection {collection_name}: {_e}")
return False
deleted = await self._run_sync_in_executor(_delete)
try:
self.client.delete_collection(name=collection_name)
deleted = True
except Exception as _e:
logger.warning(f"Failed to delete collection {collection_name}: {_e}")
deleted = False
if deleted and collection_name == self.collection_name:
self.collection = None
logger.info(f"Deleted collection {collection_name}")
async def copy_collection(self, collection_name: str, **kwargs):
"""Copy all data from the current collection to a new collection."""
source_data = self.collection.get(include=["documents", "metadatas", "embeddings"])
if not source_data["ids"]:
logger.warning(f"Source collection {self.collection_name} is empty")
return
def _copy():
source_data = self.collection.get(include=["documents", "metadatas", "embeddings"])
if not source_data["ids"]:
logger.warning(f"Source collection {self.collection_name} is empty")
return
target_collection = self.client.get_or_create_collection(
name=collection_name,
metadata={"hnsw:space": "cosine"},
)
target_collection.add(
ids=source_data["ids"],
documents=source_data["documents"],
metadatas=source_data["metadatas"],
embeddings=source_data["embeddings"],
)
await self._run_sync_in_executor(_copy)
target_collection = self.client.get_or_create_collection(
name=collection_name,
metadata={"hnsw:space": "cosine"},
)
target_collection.add(
ids=source_data["ids"],
documents=source_data["documents"],
metadatas=source_data["metadatas"],
embeddings=source_data["embeddings"],
)
logger.info(f"Copied collection {self.collection_name} to {collection_name}")
async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
@ -291,16 +266,14 @@ class ChromaVectorStore(BaseVectorStore):
batch_size = kwargs.get("batch_size", 100)
def _insert_batch(batch_nodes: list[VectorNode]):
for i in range(0, len(nodes_to_insert), batch_size):
batch_nodes = nodes_to_insert[i : i + batch_size]
self.collection.add(
ids=[n.vector_id for n in batch_nodes],
documents=[n.content for n in batch_nodes],
embeddings=[n.vector for n in batch_nodes],
metadatas=[n.metadata for n in batch_nodes],
)
for i in range(0, len(nodes_to_insert), batch_size):
await self._run_sync_in_executor(_insert_batch, nodes_to_insert[i : i + batch_size])
logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}")
async def search(
@ -315,18 +288,15 @@ class ChromaVectorStore(BaseVectorStore):
where_clause = self._generate_where_clause(filters)
include_embeddings = kwargs.get("include_embeddings", False)
def _search():
include: list = ["documents", "metadatas", "distances"]
if include_embeddings:
include.append("embeddings")
return self.collection.query(
query_embeddings=[query_vector],
n_results=limit,
where=where_clause,
include=include,
)
results = await self._run_sync_in_executor(_search)
include: list = ["documents", "metadatas", "distances"]
if include_embeddings:
include.append("embeddings")
results = self.collection.query(
query_embeddings=[query_vector],
n_results=limit,
where=where_clause,
include=include,
)
nodes = self._parse_results(results, include_score=True)
score_threshold = kwargs.get("score_threshold")
@ -341,26 +311,19 @@ class ChromaVectorStore(BaseVectorStore):
if not vector_ids:
return
def _delete():
self.collection.delete(ids=vector_ids)
await self._run_sync_in_executor(_delete)
self.collection.delete(ids=vector_ids)
logger.info(f"Deleted {len(vector_ids)} nodes from {self.collection_name}")
async def delete_all(self, **kwargs):
"""Remove all vectors from the collection."""
def _delete_all():
# Get all IDs in the collection
result = self.collection.get()
if result and result.get("ids"):
ids = result["ids"]
if ids:
self.collection.delete(ids=ids)
return len(ids)
return 0
count = await self._run_sync_in_executor(_delete_all)
# Get all IDs in the collection
result = self.collection.get()
count = 0
if result and result.get("ids"):
ids = result["ids"]
if ids:
self.collection.delete(ids=ids)
count = len(ids)
logger.info(f"Deleted all {count} nodes from {self.collection_name}")
async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
@ -380,15 +343,12 @@ class ChromaVectorStore(BaseVectorStore):
else:
nodes_to_update = nodes
def _update():
self.collection.upsert(
ids=[n.vector_id for n in nodes_to_update],
documents=[n.content for n in nodes_to_update],
embeddings=[n.vector for n in nodes_to_update if n.vector] or None,
metadatas=[n.metadata for n in nodes_to_update],
)
await self._run_sync_in_executor(_update)
self.collection.upsert(
ids=[n.vector_id for n in nodes_to_update],
documents=[n.content for n in nodes_to_update],
embeddings=[n.vector for n in nodes_to_update],
metadatas=[n.metadata for n in nodes_to_update],
)
logger.info(f"Updated {len(nodes_to_update)} nodes in {self.collection_name}")
async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None:
@ -396,10 +356,7 @@ class ChromaVectorStore(BaseVectorStore):
is_single = isinstance(vector_ids, str)
ids = [vector_ids] if is_single else vector_ids
def _get():
return self.collection.get(ids=ids, include=["documents", "metadatas", "embeddings"])
results = await self._run_sync_in_executor(_get)
results = self.collection.get(ids=ids, include=["documents", "metadatas", "embeddings"])
nodes = self._parse_results(results)
return nodes[0] if is_single and nodes else (nodes if not is_single else None)
@ -423,14 +380,11 @@ class ChromaVectorStore(BaseVectorStore):
# If sorting is needed, fetch all records first, then apply limit after sorting
fetch_limit = None if sort_key else limit
def _list():
return self.collection.get(
where=where_clause,
limit=fetch_limit,
include=["documents", "metadatas", "embeddings"],
)
results = await self._run_sync_in_executor(_list)
results = self.collection.get(
where=where_clause,
limit=fetch_limit,
include=["documents", "metadatas", "embeddings"],
)
nodes = self._parse_results(results)
# Apply sorting if sort_key is provided
@ -453,22 +407,38 @@ class ChromaVectorStore(BaseVectorStore):
async def count(self) -> int:
"""Return the total number of vectors in the current collection."""
return await self._run_sync_in_executor(self.collection.count)
return self.collection.count()
async def reset(self):
"""Reset the current collection by clearing all its data."""
logger.warning(f"Resetting collection {self.collection_name}...")
await self.delete_collection(self.collection_name)
def _recreate():
self.collection = self.client.get_or_create_collection(
name=self.collection_name,
metadata={"hnsw:space": "cosine"},
)
await self._run_sync_in_executor(_recreate)
self.collection = self.client.get_or_create_collection(
name=self.collection_name,
metadata={"hnsw:space": "cosine"},
)
logger.info(f"Collection {self.collection_name} has been reset")
async def start(self) -> None:
"""Initialize the ChromaDB collection.
Creates or retrieves the collection with cosine similarity metric.
For local mode, creates the db_path directory if it doesn't exist.
"""
if self.is_local:
self.db_path.mkdir(parents=True, exist_ok=True)
logger.info(f"Initializing local ChromaDB at {self.db_path}")
self.client = chromadb.PersistentClient(
path=str(self.db_path),
settings=Settings(anonymized_telemetry=False),
)
self.collection = self.client.get_or_create_collection(
name=self.collection_name,
metadata={"hnsw:space": "cosine"},
)
logger.info(f"ChromaDB collection {self.collection_name} initialized")
async def close(self):
"""Close the vector store and log the shutdown process."""
logger.info(f"ChromaDB vector store for collection {self.collection_name} closed")

View file

@ -4,7 +4,7 @@ This module provides an Elasticsearch-based vector store that implements the Bas
interface for high-performance dense vector storage and retrieval.
"""
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any
from loguru import logger
@ -30,8 +30,8 @@ class ESVectorStore(BaseVectorStore):
def __init__(
self,
collection_name: str,
db_path: str | Path,
embedding_model: BaseEmbeddingModel,
thread_pool: ThreadPoolExecutor,
hosts: str | list[str] | None = None,
basic_auth: tuple[str, str] | None = None,
cloud_id: str | None = None,
@ -44,8 +44,8 @@ class ESVectorStore(BaseVectorStore):
Args:
collection_name: Name of the Elasticsearch index (converted to lowercase).
db_path: Database path (not used for remote Elasticsearch, kept for API consistency).
embedding_model: Model instance used to generate vector embeddings.
thread_pool: ThreadPoolExecutor for running synchronous operations.
hosts: Connection host(s) for the Elasticsearch cluster.
basic_auth: Credentials for basic authentication.
cloud_id: Deployment ID for Elastic Cloud.
@ -64,8 +64,8 @@ class ESVectorStore(BaseVectorStore):
super().__init__(
collection_name=collection_name,
db_path=db_path,
embedding_model=embedding_model,
thread_pool=thread_pool,
**kwargs,
)
@ -521,6 +521,14 @@ class ESVectorStore(BaseVectorStore):
await self.create_collection(collection_name)
logger.info(f"Collection reset to {collection_name}")
async def start(self) -> None:
"""Initialize the Elasticsearch index.
Creates the index with dense vector mappings if it doesn't exist.
"""
await super().start()
logger.info(f"Elasticsearch index {self.collection_name} initialized")
async def close(self):
"""Terminate the Elasticsearch client session and release resources."""
await self.client.close()

View file

@ -1,85 +1,95 @@
"""Local file system vector store implementation for ReMe."""
import json
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import numpy as np
from loguru import logger
from .base_vector_store import BaseVectorStore
from ..embedding import BaseEmbeddingModel
from ..schema import VectorNode
from ..utils import cosine_similarity
from ..utils import batch_cosine_similarity
class LocalVectorStore(BaseVectorStore):
"""Local file system-based vector store using JSON files and manual cosine similarity."""
"""Local file system-based vector store with in-memory caching.
All operations are performed in memory after start().
Changes are persisted to disk on close().
"""
def __init__(
self,
collection_name: str,
db_path: str | Path,
embedding_model: BaseEmbeddingModel,
thread_pool: ThreadPoolExecutor,
root_path: str = "./local_vector_store",
**kwargs,
):
"""Initialize the local vector store with a root path and collection name."""
"""Initialize the local vector store with a db_path and collection name."""
super().__init__(
collection_name=collection_name,
db_path=db_path,
embedding_model=embedding_model,
thread_pool=thread_pool,
**kwargs,
)
self.root_path = Path(root_path)
self.collection_path = self.root_path / collection_name
self.root_path.mkdir(parents=True, exist_ok=True)
# In-memory cache: vector_id -> VectorNode
self._cache: dict[str, VectorNode] = {}
self._dirty: bool = False # Track if cache has unsaved changes
def _get_collection_path(self, collection_name: str) -> Path:
"""Get the file system path for a specific collection."""
return self.root_path / collection_name
return self.db_path / collection_name
def _get_node_file_path(self, vector_id: str, collection_name: str | None = None) -> Path:
"""Get the JSON file path for a specific vector node."""
col_path = self._get_collection_path(collection_name or self.collection_name)
return col_path / f"{vector_id}.json"
def _save_node(self, node: VectorNode, collection_name: str | None = None):
def _save_node_to_disk(self, node: VectorNode):
"""Save a vector node to a JSON file on disk."""
file_path = self._get_node_file_path(node.vector_id, collection_name)
file_path = self._get_node_file_path(node.vector_id)
file_path.parent.mkdir(parents=True, exist_ok=True)
with open(file_path, "w", encoding="utf-8") as f:
json.dump(node.model_dump(), f, ensure_ascii=False, indent=2)
def _load_node(self, vector_id: str, collection_name: str | None = None) -> VectorNode | None:
"""Load a vector node from a JSON file."""
file_path = self._get_node_file_path(vector_id, collection_name)
if not file_path.exists():
return None
with open(file_path, "r", encoding="utf-8") as f:
data = json.load(f)
return VectorNode(**data)
def _load_all_nodes(self, collection_name: str | None = None) -> list[VectorNode]:
"""Load all vector nodes existing in a collection."""
col_path = self._get_collection_path(collection_name or self.collection_name)
def _load_all_from_disk(self) -> dict[str, VectorNode]:
"""Load all vector nodes from disk into a dictionary."""
col_path = self._get_collection_path(self.collection_name)
if not col_path.exists():
return []
return {}
nodes = []
nodes = {}
for file_path in col_path.glob("*.json"):
try:
with open(file_path, "r", encoding="utf-8") as f:
data = json.load(f)
nodes.append(VectorNode(**data))
node = VectorNode(**data)
nodes[node.vector_id] = node
except Exception as e:
logger.warning(f"Failed to load node from {file_path}: {e}")
return nodes
def _flush_to_disk(self):
"""Persist all cached nodes to disk."""
col_path = self._get_collection_path(self.collection_name)
col_path.mkdir(parents=True, exist_ok=True)
# Remove files that are no longer in cache
existing_files = set(col_path.glob("*.json"))
cached_ids = set(self._cache.keys())
for file_path in existing_files:
vector_id = file_path.stem
if vector_id not in cached_ids:
file_path.unlink()
# Write all cached nodes
for node in self._cache.values():
self._save_node_to_disk(node)
self._dirty = False
logger.info(f"Flushed {len(self._cache)} nodes to disk")
@staticmethod
def _match_filters(node: VectorNode, filters: dict | None) -> bool:
"""Check if a vector node matches the provided metadata filters.
@ -114,11 +124,11 @@ class LocalVectorStore(BaseVectorStore):
return True
async def list_collections(self) -> list[str]:
"""List all collection directories in the root path."""
if not self.root_path.exists():
"""List all collection directories in the db_path."""
if not self.db_path.exists():
return []
return [d.name for d in self.root_path.iterdir() if d.is_dir() and not d.name.startswith(".")]
return [d.name for d in self.db_path.iterdir() if d.is_dir() and not d.name.startswith(".")]
async def create_collection(self, collection_name: str, **kwargs):
"""Create a new collection directory."""
@ -158,7 +168,7 @@ class LocalVectorStore(BaseVectorStore):
logger.info(f"Copied collection {self.collection_name} to {collection_name}")
async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
"""Insert vector nodes into the local store, generating embeddings if necessary."""
"""Insert vector nodes into the cache, generating embeddings if necessary."""
if isinstance(nodes, VectorNode):
nodes = [nodes]
@ -171,8 +181,9 @@ class LocalVectorStore(BaseVectorStore):
nodes_to_insert = nodes
for node in nodes_to_insert:
self._save_node(node)
self._cache[node.vector_id] = node
self._dirty = True
logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}")
async def search(
@ -182,73 +193,69 @@ class LocalVectorStore(BaseVectorStore):
filters: dict | None = None,
**kwargs,
) -> list[VectorNode]:
"""Search for nodes similar to the query using brute-force cosine similarity."""
"""Search for nodes similar to the query using batch cosine similarity."""
query_vector = await self.get_embedding(query)
all_nodes = self._load_all_nodes()
filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)]
scored_nodes = []
for node in filtered_nodes:
if node.vector is None:
logger.warning(f"Node {node.vector_id} has no vector, skipping")
continue
# Filter nodes from cache
filtered_nodes = [node for node in self._cache.values() if self._match_filters(node, filters)]
try:
score = cosine_similarity(query_vector, node.vector)
scored_nodes.append((node, score))
except ValueError as e:
logger.warning(f"Failed to calculate similarity for node {node.vector_id}: {e}")
# Separate nodes with and without vectors
nodes_with_vectors = [node for node in filtered_nodes if node.vector is not None]
if not nodes_with_vectors:
return []
scored_nodes.sort(key=lambda x: x[1], reverse=True)
# Build matrix for batch similarity computation
node_vectors = np.array([node.vector for node in nodes_with_vectors])
query_matrix = np.array([query_vector])
# Compute similarities in batch: shape (1, num_nodes) -> flatten to (num_nodes,)
similarities = batch_cosine_similarity(query_matrix, node_vectors).flatten()
# Apply score threshold if specified
score_threshold = kwargs.get("score_threshold")
# Pair nodes with scores and filter/sort
scored_nodes = list(zip(nodes_with_vectors, similarities))
if score_threshold is not None:
scored_nodes = [(node, score) for node, score in scored_nodes if score >= score_threshold]
scored_nodes.sort(key=lambda x: x[1], reverse=True)
scored_nodes = scored_nodes[:limit]
# Attach scores to metadata
results = []
for node, score in scored_nodes:
node.metadata["score"] = score
node.metadata["score"] = float(score)
results.append(node)
return results
async def delete(self, vector_ids: str | list[str], **kwargs):
"""Delete specific vector nodes by their IDs."""
"""Delete specific vector nodes by their IDs from cache."""
if isinstance(vector_ids, str):
vector_ids = [vector_ids]
deleted_count = 0
for vector_id in vector_ids:
file_path = self._get_node_file_path(vector_id)
if file_path.exists():
file_path.unlink()
if vector_id in self._cache:
del self._cache[vector_id]
deleted_count += 1
else:
logger.warning(f"Node {vector_id} does not exist")
if deleted_count > 0:
self._dirty = True
logger.info(f"Deleted {deleted_count} nodes from {self.collection_name}")
async def delete_all(self, **kwargs):
"""Remove all vectors from the collection."""
col_path = self._get_collection_path(self.collection_name)
if not col_path.exists():
logger.warning(f"Collection {self.collection_name} does not exist")
return
deleted_count = 0
for file_path in col_path.glob("*.json"):
try:
file_path.unlink()
deleted_count += 1
except Exception as e:
logger.warning(f"Failed to delete file {file_path}: {e}")
logger.info(f"Deleted all {deleted_count} nodes from {self.collection_name}")
"""Remove all vectors from the cache."""
count = len(self._cache)
self._cache.clear()
self._dirty = True
logger.info(f"Deleted all {count} nodes from {self.collection_name}")
async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
"""Update existing vector nodes with new data or embeddings."""
"""Update existing vector nodes in the cache."""
if isinstance(nodes, VectorNode):
nodes = [nodes]
@ -262,23 +269,24 @@ class LocalVectorStore(BaseVectorStore):
updated_count = 0
for node in nodes_to_update:
file_path = self._get_node_file_path(node.vector_id)
if file_path.exists():
self._save_node(node)
if node.vector_id in self._cache:
self._cache[node.vector_id] = node
updated_count += 1
else:
logger.warning(f"Node {node.vector_id} does not exist, skipping update")
if updated_count > 0:
self._dirty = True
logger.info(f"Updated {updated_count} nodes in {self.collection_name}")
async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]:
"""Retrieve one or more vector nodes by their unique IDs."""
"""Retrieve one or more vector nodes from cache by their unique IDs."""
is_single = isinstance(vector_ids, str)
ids = [vector_ids] if is_single else vector_ids
results = []
for vector_id in ids:
node = self._load_node(vector_id)
node = self._cache.get(vector_id)
if node:
results.append(node)
else:
@ -291,9 +299,9 @@ class LocalVectorStore(BaseVectorStore):
filters: dict | None = None,
limit: int | None = None,
sort_key: str | None = None,
reverse: bool = False,
reverse: bool = True,
) -> list[VectorNode]:
"""List vector nodes in the collection with optional filtering and limits.
"""List vector nodes from cache with optional filtering and limits.
Args:
filters: Dictionary of filter conditions to match vectors
@ -301,16 +309,14 @@ class LocalVectorStore(BaseVectorStore):
sort_key: Key to sort the results by (e.g., field name in metadata). None for no sorting
reverse: If True, sort in descending order; if False, sort in ascending order
"""
all_nodes = self._load_all_nodes()
filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)]
filtered_nodes = [node for node in self._cache.values() if self._match_filters(node, filters)]
# Apply sorting if sort_key is provided
if sort_key:
# Sort with proper handling of None and missing values
def sort_key_func(node):
value = node.metadata.get(sort_key)
if value is None:
# Return appropriate default based on reverse flag
return float("-inf") if not reverse else float("inf")
return value
@ -321,6 +327,15 @@ class LocalVectorStore(BaseVectorStore):
return filtered_nodes
async def start(self) -> None:
"""Initialize the local vector store and load all nodes into memory."""
await super().start()
self._cache = self._load_all_from_disk()
self._dirty = False
logger.info(f"Local vector store loaded {len(self._cache)} nodes from {self.collection_name}")
async def close(self):
"""Close the vector store (no-op for local file system)."""
"""Persist all cached data to disk and close the vector store."""
if self._dirty:
self._flush_to_disk()
logger.info("Local vector store closed")

View file

@ -2,7 +2,7 @@
import json
import re
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any
from loguru import logger
@ -47,8 +47,8 @@ class PGVectorStore(BaseVectorStore):
def __init__(
self,
collection_name: str,
db_path: str | Path,
embedding_model: BaseEmbeddingModel,
thread_pool: ThreadPoolExecutor,
host: str = "localhost",
port: int = 5432,
database: str = "postgres",
@ -72,8 +72,8 @@ class PGVectorStore(BaseVectorStore):
super().__init__(
collection_name=collection_name,
db_path=db_path,
embedding_model=embedding_model,
thread_pool=thread_pool,
**kwargs,
)
@ -117,12 +117,6 @@ class PGVectorStore(BaseVectorStore):
return self._pool
async def _ensure_collection_exists(self):
"""Check if the collection table exists and create it if missing."""
collections = await self.list_collections()
if self.collection_name not in collections:
await self.create_collection(self.collection_name)
async def list_collections(self) -> list[str]:
"""List all available table names in the current database."""
pool = await self._get_pool()
@ -219,8 +213,6 @@ class PGVectorStore(BaseVectorStore):
async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
"""Insert or upsert vector nodes into the PostgreSQL collection."""
await self._ensure_collection_exists()
if isinstance(nodes, VectorNode):
nodes = [nodes]
@ -337,8 +329,6 @@ class PGVectorStore(BaseVectorStore):
**kwargs,
) -> list[VectorNode]:
"""Perform vector similarity search with optional metadata filtering."""
await self._ensure_collection_exists()
query_vector = await self.get_embedding(query)
vector_str = f"[{','.join(map(str, query_vector))}]"
pool = await self._get_pool()
@ -394,8 +384,6 @@ class PGVectorStore(BaseVectorStore):
async def delete(self, vector_ids: str | list[str], **kwargs):
"""Remove specific vector records from the collection by their IDs."""
await self._ensure_collection_exists()
if isinstance(vector_ids, str):
vector_ids = [vector_ids]
@ -414,8 +402,6 @@ class PGVectorStore(BaseVectorStore):
async def delete_all(self, **kwargs):
"""Remove all vectors from the collection."""
await self._ensure_collection_exists()
pool = await self._get_pool()
async with pool.acquire() as conn:
result = await conn.execute(f"DELETE FROM {self.collection_name}")
@ -424,8 +410,6 @@ class PGVectorStore(BaseVectorStore):
async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
"""Update existing vector nodes with new content, embeddings, or metadata."""
await self._ensure_collection_exists()
if isinstance(nodes, VectorNode):
nodes = [nodes]
@ -474,8 +458,6 @@ class PGVectorStore(BaseVectorStore):
async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None:
"""Retrieve vector nodes by their unique identifiers."""
await self._ensure_collection_exists()
single_result = isinstance(vector_ids, str)
if single_result:
vector_ids = [vector_ids]
@ -531,8 +513,6 @@ class PGVectorStore(BaseVectorStore):
sort_key: Key to sort the results by (e.g., field name in metadata). None for no sorting
reverse: If True, sort in descending order; if False, sort in ascending order
"""
await self._ensure_collection_exists()
pool = await self._get_pool()
filter_clause, filter_params = self._build_filter_clause(filters)
@ -613,6 +593,15 @@ class PGVectorStore(BaseVectorStore):
await self.create_collection(collection_name)
logger.info(f"Collection reset to {collection_name}")
async def start(self) -> None:
"""Initialize the PGVector store.
Creates the connection pool and ensures the collection table exists.
"""
await self._get_pool()
await super().start()
logger.info(f"PGVector collection {self.collection_name} initialized")
async def close(self):
"""Terminate the database connection pool and release associated resources."""
if self._pool is not None:

View file

@ -1,6 +1,6 @@
"""Qdrant vector store implementation for the ReMe project."""
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any
from loguru import logger
@ -42,11 +42,10 @@ class QdrantVectorStore(BaseVectorStore):
def __init__(
self,
collection_name: str,
db_path: str | Path,
embedding_model: BaseEmbeddingModel,
thread_pool: ThreadPoolExecutor,
host: str | None = None,
port: int = 6333,
path: str | None = None,
url: str | None = None,
api_key: str | None = None,
https: bool | None = None,
@ -60,11 +59,10 @@ class QdrantVectorStore(BaseVectorStore):
Args:
collection_name: Name of the collection.
db_path: Local storage path for on-disk/in-memory mode.
embedding_model: Model used for generating vector embeddings.
thread_pool: ThreadPoolExecutor for running synchronous operations.
host: Server host address.
port: HTTP port for the server.
path: Local storage path for on-disk/in-memory mode.
url: Full connection URL.
api_key: Authentication key for Qdrant Cloud.
https: Use secure connection if True.
@ -81,26 +79,23 @@ class QdrantVectorStore(BaseVectorStore):
super().__init__(
collection_name=collection_name,
db_path=db_path,
embedding_model=embedding_model,
thread_pool=thread_pool,
**kwargs,
)
client_kwargs = {k: v for k, v in kwargs.items() if k != "thread_pool"}
self.is_local = host is None and url is None and api_key is None
self.client: AsyncQdrantClient
self.client = AsyncQdrantClient(
host=host,
port=port,
path=path,
url=url,
api_key=api_key,
https=https,
grpc_port=grpc_port,
prefer_grpc=prefer_grpc,
**client_kwargs,
)
# Store connection parameters for deferred initialization in start()
self._host = host
self._port = port
self._url = url
self._api_key = api_key
self._https = https
self._grpc_port = grpc_port
self._prefer_grpc = prefer_grpc
self.is_local = path is not None
distance_map = {
"cosine": Distance.COSINE,
"euclid": Distance.EUCLID,
@ -526,6 +521,29 @@ class QdrantVectorStore(BaseVectorStore):
return results
async def start(self) -> None:
"""Initialize the Qdrant collection.
Creates the collection if it doesn't exist with configured vector parameters.
For local mode, creates the db_path directory if it doesn't exist.
"""
if self.is_local:
self.db_path.mkdir(parents=True, exist_ok=True)
self.client = AsyncQdrantClient(path=str(self.db_path))
else:
self.client = AsyncQdrantClient(
host=self._host,
port=self._port,
url=self._url,
api_key=self._api_key,
https=self._https,
grpc_port=self._grpc_port,
prefer_grpc=self._prefer_grpc,
**self.kwargs,
)
await super().start()
logger.info(f"Qdrant collection {self.collection_name} initialized")
async def close(self):
"""Close the AsyncQdrantClient connection and release resources."""
await self.client.close()

View file

@ -0,0 +1,18 @@
"""Extension operations and tools."""
from .simple_chat import SimpleChat
from .stream_chat import StreamChat
from .test_op import TestOp
from .translate_ts import TranslateTs
from ..core.registry_factory import R
__all__ = [
"SimpleChat",
"StreamChat",
"TestOp",
"TranslateTs",
]
for name in __all__:
op_class = globals()[name]
R.ops.register(op_class)

View file

@ -2,9 +2,9 @@
from loguru import logger
from ...core.enumeration import Role
from ...core.op import BaseTool
from ...core.schema import Message, ToolCall
from ..core.enumeration import Role
from ..core.op import BaseTool
from ..core.schema import Message, ToolCall
class SimpleChat(BaseTool):

Some files were not shown because too many files have changed in this diff Show more