mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
Merge branch 'main' into dev_0227
This commit is contained in:
commit
bb0d0ee248
182 changed files with 1525 additions and 2185 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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/*
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,9 +0,0 @@
|
|||
"""A simple chatbot."""
|
||||
|
||||
from . import chat
|
||||
from . import memory
|
||||
|
||||
__all__ = [
|
||||
"chat",
|
||||
"memory",
|
||||
]
|
||||
|
|
@ -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)
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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
|
||||
|
|
@ -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."
|
||||
}}
|
||||
```
|
||||
|
|
@ -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,
|
||||
}
|
||||
|
|
@ -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 they’re 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"}}
|
||||
|
||||
|
|
@ -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
|
||||
|
|
@ -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.
|
||||
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
64
reme/config/service.yaml
Normal 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
|
||||
|
||||
|
|
@ -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:
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
23
reme/core/file_store/__init__.py
Normal file
23
reme/core/file_store/__init__.py
Normal 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)
|
||||
|
|
@ -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."""
|
||||
|
|
@ -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,
|
||||
|
|
@ -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}'")
|
||||
|
|
@ -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()
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -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()
|
||||
|
|
@ -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__(
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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] = {}
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
45
reme/core/tools/__init__.py
Normal file
45
reme/core/tools/__init__.py
Normal 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)
|
||||
|
|
@ -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):
|
||||
|
|
@ -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):
|
||||
|
|
@ -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)}"
|
||||
|
|
@ -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:
|
||||
|
|
@ -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):
|
||||
|
|
@ -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):
|
||||
|
|
@ -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:
|
||||
|
|
@ -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:
|
||||
|
|
@ -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:
|
||||
|
|
@ -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
|
||||
|
|
@ -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:
|
||||
|
|
@ -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):
|
||||
|
|
@ -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):
|
||||
|
|
@ -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):
|
||||
|
|
@ -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):
|
||||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
18
reme/extension/__init__.py
Normal file
18
reme/extension/__init__.py
Normal 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)
|
||||
|
|
@ -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
Loading…
Add table
Reference in a new issue