mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-21 00:22:45 +00:00
up
This commit is contained in:
parent
e5960d0ad0
commit
c8e96b5ae8
76 changed files with 6141 additions and 0 deletions
378
docs4/expert.yaml
Normal file
378
docs4/expert.yaml
Normal file
|
|
@ -0,0 +1,378 @@
|
|||
app_name: reme-expert
|
||||
enable_logo: false
|
||||
log_to_console: true
|
||||
log_to_file: false
|
||||
|
||||
# Expert-exposure profile: every memory_* read + write primitive plus
|
||||
# `sync` (event log) are surfaced as MCP tools. The agent picks whatever
|
||||
# it needs — finest granularity, no opinion enforced beyond the
|
||||
# wikilink-uniqueness gate.
|
||||
#
|
||||
# **No `ingest` tool.** Cold-path R-M-W in expert mode is owned by the
|
||||
# host agent (Claude Code, with the `reme-expert` SKILL loaded), not by
|
||||
# reme2's internal ReActAgent. The reme-expert plugin spawns a
|
||||
# `reme-distiller` subagent at SessionEnd to handle distillation in its
|
||||
# own context window, using the same memory_* primitives.
|
||||
#
|
||||
# Use when you want the agent to manage memory directly with full
|
||||
# control. For an opinionated minimal surface (where reme2's internal
|
||||
# Ingestor owns R-M-W), see ./service.yaml.
|
||||
|
||||
service:
|
||||
backend: mcp
|
||||
transport: stdio
|
||||
sidecar_http: true
|
||||
sidecar_http_host: "127.0.0.1"
|
||||
sidecar_http_port: 8765
|
||||
sidecar_info_path: "./vault/.reme/sidecar.json"
|
||||
|
||||
jobs:
|
||||
# -- Write entry points ------------------------------------------------
|
||||
|
||||
- backend: base
|
||||
name: sync
|
||||
description: |
|
||||
Hot-path event sync: idempotent upsert of an event FOLDER under
|
||||
`events/{date}/{name}/`. The folder contains the index `{name}.md`
|
||||
(Event schema, frontmatter + narrative + Materials footer) plus
|
||||
any raw materials you pass — conversation snippets, tool outputs,
|
||||
data dumps. The watcher indexes everything inside.
|
||||
|
||||
CONTINUITY MODEL: pick a stable `name` per logical thread and
|
||||
call `sync` repeatedly through the task. Each call extends the
|
||||
same folder:
|
||||
* new `content` → appended under a `## Update — {iso}` section
|
||||
* new `materials` → siblings (auto-suffix on filename collision)
|
||||
* `topics` + `tags` merged (union) into frontmatter
|
||||
* Materials footer regenerated to list every artifact in the folder
|
||||
|
||||
First call (folder doesn't exist) → CREATE; subsequent calls with
|
||||
the same `name` while the event is `status: active` → APPEND.
|
||||
If the event is `status: distilled` / `archived`, `sync` REFUSES
|
||||
and returns a `suggested_name` so you start a fresh thread
|
||||
instead of mutating prior cognition.
|
||||
|
||||
Zero LLM cost. Call CONTINUOUSLY through a task as facts land,
|
||||
and especially at PreCompact to dump verbose raw text into
|
||||
`materials` before context truncation loses it. The folder is
|
||||
the unit `ingest` later reads from.
|
||||
parameters:
|
||||
type: object
|
||||
required: [name]
|
||||
properties:
|
||||
name:
|
||||
type: string
|
||||
description: "kebab-case event identifier (folder + index stem). Reuse the same name across calls in one thread to keep extending the same folder."
|
||||
description: { type: string, description: "one-line summary for index frontmatter (set on initial create only)" }
|
||||
content: { type: string, description: "markdown body. Initial create: the body. Subsequent calls: appended as a `## Update — {iso}` section." }
|
||||
topics:
|
||||
type: array
|
||||
items: { type: string }
|
||||
description: "related topic wikilinks ('[[X]]' or '[[topics/X/X]]'). Unioned into frontmatter on append."
|
||||
tags: { type: array, items: { type: string }, description: "free-form tags; unioned on append." }
|
||||
materials:
|
||||
type: array
|
||||
description: "Raw artifacts written as siblings of the index inside the event folder. Filenames must be safe (letters/digits/dot/underscore/dash). On filename collision with an existing artifact, auto-suffixes (foo.txt → foo-2.txt)."
|
||||
items:
|
||||
type: object
|
||||
required: [filename, content]
|
||||
properties:
|
||||
filename: { type: string, description: "e.g. 'raw-prompt.md', 'tool-output.txt', 'snapshot.json'" }
|
||||
content: { type: string }
|
||||
on_date: { type: string, description: "ISO date for events/{date}/ bucket; defaults to today" }
|
||||
origin_session_id: { type: string, description: "set on initial create only" }
|
||||
steps:
|
||||
- backend: sync
|
||||
|
||||
# -- Read tools --------------------------------------------------------
|
||||
|
||||
- backend: base
|
||||
name: memory_search
|
||||
description: "Hybrid (vector + keyword) search over chunks."
|
||||
parameters:
|
||||
type: object
|
||||
required: [query]
|
||||
properties:
|
||||
query: { type: string }
|
||||
max_results: { type: integer, default: 5 }
|
||||
min_score: { type: number, default: 0.1 }
|
||||
paths: { type: array, items: { type: string } }
|
||||
tags: { type: array, items: { type: string } }
|
||||
exclude_paths: { type: array, items: { type: string } }
|
||||
steps:
|
||||
- backend: memory_search
|
||||
|
||||
- backend: base
|
||||
name: memory_graph_search
|
||||
description: |
|
||||
Three-way fusion search: vector + keyword + graph (BFS over wikilinks).
|
||||
Pulls in chunks reachable through linked topics/events that pure
|
||||
relevance search would miss.
|
||||
parameters:
|
||||
type: object
|
||||
properties:
|
||||
query: { type: string }
|
||||
max_results: { type: integer, default: 5 }
|
||||
min_score: { type: number, default: 0.0 }
|
||||
graph_depth: { type: integer, default: 1 }
|
||||
seeds: { type: array, items: { type: string } }
|
||||
paths: { type: array, items: { type: string } }
|
||||
tags: { type: array, items: { type: string } }
|
||||
exclude_paths: { type: array, items: { type: string } }
|
||||
steps:
|
||||
- backend: memory_graph_search
|
||||
|
||||
- backend: base
|
||||
name: memory_get
|
||||
description: "Read frontmatter + body of a single file."
|
||||
parameters:
|
||||
type: object
|
||||
required: [path]
|
||||
properties:
|
||||
path: { type: string }
|
||||
include_chunks: { type: boolean, default: false }
|
||||
steps:
|
||||
- backend: memory_get
|
||||
|
||||
- backend: base
|
||||
name: memory_list
|
||||
description: |
|
||||
List indexed files filtered by frontmatter exact-match, tags, or
|
||||
path prefix. Returns {items: [{path, metadata}], count}.
|
||||
parameters:
|
||||
type: object
|
||||
properties:
|
||||
metadata: { type: object }
|
||||
tags: { type: array, items: { type: string } }
|
||||
path_prefix: { type: string }
|
||||
limit: { type: integer, default: 100 }
|
||||
steps:
|
||||
- backend: memory_list
|
||||
|
||||
- backend: base
|
||||
name: memory_backlinks
|
||||
description: "Files linking TO the given path (with edge predicates)."
|
||||
parameters:
|
||||
type: object
|
||||
required: [path]
|
||||
properties:
|
||||
path: { type: string }
|
||||
steps:
|
||||
- backend: memory_backlinks
|
||||
|
||||
- backend: base
|
||||
name: memory_links
|
||||
description: "Files the given path links to (resolved, with edge predicates)."
|
||||
parameters:
|
||||
type: object
|
||||
required: [path]
|
||||
properties:
|
||||
path: { type: string }
|
||||
steps:
|
||||
- backend: memory_links
|
||||
|
||||
- backend: base
|
||||
name: memory_resolve_wikilink
|
||||
description: |
|
||||
Resolve a `[[target]]` wikilink to a vault path. Stem-form (`X`)
|
||||
consults the file_store's stem index; path-form (`a/b` or `a/b.md`)
|
||||
is anchored at the vault root.
|
||||
parameters:
|
||||
type: object
|
||||
required: [wikilink]
|
||||
properties:
|
||||
wikilink: { type: string }
|
||||
steps:
|
||||
- backend: memory_resolve_wikilink
|
||||
|
||||
- backend: base
|
||||
name: memory_count_tokens
|
||||
description: "Estimate token count for a file body or raw text."
|
||||
parameters:
|
||||
type: object
|
||||
properties:
|
||||
path: { type: string }
|
||||
text: { type: string }
|
||||
steps:
|
||||
- backend: memory_count_tokens
|
||||
|
||||
- backend: base
|
||||
name: memory_lint
|
||||
description: |
|
||||
Read-only projection of the Maintainer's lint findings. Walks the
|
||||
indexed files (optionally restricted to a path prefix) and returns
|
||||
structural issues — broken wikilinks, schema violations, stem
|
||||
collisions. Never mutates. The agent decides what to do with each
|
||||
finding using existing memory_* primitives (rename, update,
|
||||
property_update, archive, delete).
|
||||
parameters:
|
||||
type: object
|
||||
properties:
|
||||
target_prefix:
|
||||
type: string
|
||||
description: "restrict scan to relpaths starting with this prefix (e.g. 'events/2026-05-09/'). Empty string scans the whole vault."
|
||||
steps:
|
||||
- backend: memory_lint
|
||||
|
||||
# -- Write primitives (raw building blocks) ----------------------------
|
||||
|
||||
- backend: base
|
||||
name: memory_create
|
||||
description: |
|
||||
Create a new file (raw primitive — no LLM reasoning). Prefer
|
||||
`ingest` when you want an LLM curator. Two gates run unless
|
||||
`force=true`: (1) path template — must be `topics/{folder}/{name}.md`,
|
||||
`events/{date}/{name}/...`, or `Archive/...`; (2) wikilink uniqueness
|
||||
— `[[stem]]` must resolve to ≤1 path post-create.
|
||||
parameters:
|
||||
type: object
|
||||
required: [path]
|
||||
properties:
|
||||
path: { type: string }
|
||||
metadata: { type: object }
|
||||
content: { type: string }
|
||||
overwrite: { type: boolean, default: false }
|
||||
force:
|
||||
type: boolean
|
||||
default: false
|
||||
description: "bypass BOTH the path-template gate and the wikilink-uniqueness gate. Use only when you intentionally need a non-template path or accept the ambiguity."
|
||||
steps:
|
||||
- backend: memory_create
|
||||
|
||||
- backend: base
|
||||
name: memory_update
|
||||
description: |
|
||||
Edit-style content update: replace `old_string` with `new_string`
|
||||
in the file body. For frontmatter changes use `memory_property_update`.
|
||||
parameters:
|
||||
type: object
|
||||
required: [path, old_string, new_string]
|
||||
properties:
|
||||
path: { type: string }
|
||||
old_string: { type: string }
|
||||
new_string: { type: string }
|
||||
replace_all: { type: boolean, default: false }
|
||||
steps:
|
||||
- backend: memory_update
|
||||
|
||||
- backend: base
|
||||
name: memory_property_update
|
||||
description: |
|
||||
Update a single YAML frontmatter key. value=null deletes the key.
|
||||
When key='status', enforces the active → distilled → archived
|
||||
single-direction state machine (skip / reverse refused). Pass
|
||||
`force=true` to bypass.
|
||||
parameters:
|
||||
type: object
|
||||
required: [path, key]
|
||||
properties:
|
||||
path: { type: string }
|
||||
key: { type: string }
|
||||
value: {}
|
||||
force:
|
||||
type: boolean
|
||||
default: false
|
||||
description: "bypass the status state machine (only meaningful when key='status')."
|
||||
steps:
|
||||
- backend: memory_property_update
|
||||
|
||||
- backend: base
|
||||
name: memory_rename
|
||||
description: |
|
||||
Rename a file and rewrite all incoming `[[wikilink]]` references
|
||||
across the vault. Refuses on destination conflict or stem ambiguity.
|
||||
parameters:
|
||||
type: object
|
||||
required: [old_path, new_path]
|
||||
properties:
|
||||
old_path: { type: string }
|
||||
new_path: { type: string }
|
||||
steps:
|
||||
- backend: memory_rename
|
||||
|
||||
- backend: base
|
||||
name: memory_delete
|
||||
description: "Delete a file."
|
||||
parameters:
|
||||
type: object
|
||||
required: [path]
|
||||
properties:
|
||||
path: { type: string }
|
||||
steps:
|
||||
- backend: memory_delete
|
||||
|
||||
- backend: base
|
||||
name: memory_archive
|
||||
description: |
|
||||
Archive a file: flip `status: archived` then move under
|
||||
`<vault>/<archive_dir>/<original_relative_path>`.
|
||||
parameters:
|
||||
type: object
|
||||
required: [path]
|
||||
properties:
|
||||
path: { type: string }
|
||||
archive_dir: { type: string, default: "Archive" }
|
||||
steps:
|
||||
- backend: memory_archive
|
||||
|
||||
components:
|
||||
# Ingestor LLM (opt-in). Without this the Ingestor degrades to a
|
||||
# direct create from explicit `target_path`; edits/renames/deletes
|
||||
# require the LLM. Uncomment + provide LLM_API_KEY to enable.
|
||||
#
|
||||
# as_llm:
|
||||
# default:
|
||||
# backend: openai
|
||||
# model_name: ${LLM_MODEL_NAME:-gpt-4o-mini}
|
||||
# api_key: ${LLM_API_KEY}
|
||||
# client_kwargs:
|
||||
# base_url: ${LLM_BASE_URL:-https://api.openai.com/v1}
|
||||
# stream: false
|
||||
#
|
||||
# as_llm_formatter:
|
||||
# default:
|
||||
# backend: openai
|
||||
|
||||
as_token_counter:
|
||||
default:
|
||||
backend: estimated
|
||||
|
||||
# Embedding is opt-in: leave embedding_model="" on file_store to run
|
||||
# keyword-only; uncomment + flip to "default" to enable hybrid search.
|
||||
#
|
||||
# embedding_model:
|
||||
# default:
|
||||
# backend: openai
|
||||
# model_name: ${EMBEDDING_MODEL_NAME:-text-embedding-3-small}
|
||||
# dimensions: 1536
|
||||
# pass_dimensions: false
|
||||
# enable_cache: true
|
||||
# max_batch_size: 10
|
||||
# max_cache_size: 2000
|
||||
# max_input_length: 8192
|
||||
|
||||
file_parser:
|
||||
md:
|
||||
backend: md
|
||||
default:
|
||||
backend: text
|
||||
|
||||
file_store:
|
||||
default:
|
||||
backend: local
|
||||
embedding_model: ""
|
||||
store_name: "reme"
|
||||
db_path: "./vault/.reme"
|
||||
working_dir: "./vault"
|
||||
|
||||
file_watcher:
|
||||
default:
|
||||
backend: full
|
||||
file_store: default
|
||||
default_parser: md
|
||||
recursive: true
|
||||
|
||||
# Retriever (`hybrid`) is a Step, not a pre-instantiated component —
|
||||
# the `memory_search` / `memory_graph_search` shells build it on
|
||||
# demand. To tune defaults, pass knobs (`vector_weight`, `graph_*`,
|
||||
# …) on the step config under each job below.
|
||||
51
docs4/obsidian.md
Normal file
51
docs4/obsidian.md
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
# 基础Job
|
||||
@jinli
|
||||
| 分类 | 能力 | 参数 |
|
||||
|--------|---------|-----------------------------------------------------------|
|
||||
| 通用 | help | |
|
||||
| 通用 | start | |
|
||||
| 通用 | restart | |
|
||||
| 通用 | version | |
|
||||
| search | search | query="search term" limit=10 tag="[]" score=0.1 copy=true |
|
||||
|
||||
@sen
|
||||
| tags | stat | 返回特定tag信息 |
|
||||
| tags | list | 返回所有tag列表 |
|
||||
| crud | upload/download | 其他文件 |
|
||||
| file | stat | path |
|
||||
| file | list | path |
|
||||
| property | property:read | |
|
||||
| property | property:update | path="My Note" status=done xx=xxx |
|
||||
| property | property:delete | keys="[xxxx, xxxx]" |
|
||||
| graph | traverse | path="My Note" directtion=forward/backward depth=1 predicat=xxx |
|
||||
|
||||
@wangce
|
||||
| crud | create | path="New Note" content="# Hello" title="xxx" tags="[]" status="" |
|
||||
| crud | read | path="Templates/Recipe.md" |
|
||||
| crud | edit | path="Templates/Recipe.md" old="xxx" new="xxx" |
|
||||
| crud | append | path="My Note" content="New line" |
|
||||
| crud | prepend | path="My Note" content="New line" |
|
||||
| crud | delete | path="My Note
|
||||
| daily:crud | daily:xxx | 与 crud 参数保持一致 |
|
||||
|
||||
|
||||
# 日记类型
|
||||
|
||||
| 类型 | 路径 | 说明 |
|
||||
|-----------|-----------------------------------------------|-----------------------------|
|
||||
| daily | {daily}/xxxx-mm-dd.md + xxxx-mm-dd/{event}.md | 按日期归档的原始信息记录 |
|
||||
| topic | topic/{topic:-personal(agent)}/{xxxx}.md | 按主题聚类的二次加工内容 |
|
||||
| proactive | todo | 基于 daily / topic 思考后主动推送的消息 |
|
||||
|
||||
# 生成Job
|
||||
|
||||
| 任务 | 输入 | 输出 | 触发时机 | 说明 |
|
||||
|-------------------------|---------------|-----------------------------------------------|-----------------------------|------------------------------------------------------|
|
||||
| 日记summary @sen @wangce | msg | {daily}/xxxx-mm-dd.md + xxxx-mm-dd/{event}.md | freq (every_n_turn、compact) | 把 msg 的信息写入 daily 目录 |
|
||||
| 主题dream + 生成链接 @sen | daily/xxx | knowledge/xxx | /dream | 把 daily 目录的内容按主题聚类合并到 topic 目录, 主动在文档中建立 [[link]] 关联 |
|
||||
| 主动proactive @wangce | daily / topic | proactive_query | pre_query | 思考 daily / topic 信息,主动决定推送给用户的消息 |
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
86
docs4/reme_todo.md
Normal file
86
docs4/reme_todo.md
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
ReMe重构
|
||||
1. 根目录
|
||||
a. ✅ vault_root → working_dir @sen
|
||||
2. file_parser
|
||||
a. 抽象基类 parse: @jinli
|
||||
ⅰ. 输入是path:相对路径
|
||||
ⅱ. 输出是FileMetadata & list[FileChunks] & list[FileEdge]
|
||||
b. default parser 兼容老方案 @jinli
|
||||
ⅰ. 带overlap的chunking策略 ,不输出FileEdge
|
||||
c. markdown parser @sen
|
||||
ⅰ. 根据markdown ast做chunk,不需要overlap
|
||||
ⅱ. 增加一个索引的chunk chunk_type @锦鲤 file_chunk_type content/index
|
||||
ⅲ. 增加link的正则解析:predicate:: [[path#anchor]]
|
||||
3. file_store @sen
|
||||
a. 抽象存储:
|
||||
ⅰ. filenode = file + path + st_mtime + metadata + list[FileEdge]
|
||||
ⅱ. graph=dict[str, filenode] 内存+json
|
||||
ⅲ. list[FileChunk] 存db
|
||||
b. 抽象基类
|
||||
ⅰ. graph:fellow dict的操作 update/get/set
|
||||
ⅱ. chunks dict[str, list[chunk]]
|
||||
1. delete_chunks_by_path
|
||||
2. update_chunks_by_path
|
||||
3. list_chunks_by_path
|
||||
4. vector_search/keyword_search
|
||||
ⅲ. 手写一个bm25检索
|
||||
ⅳ. 【核心】检索机制 vector bm25 graph 如何进行融合
|
||||
4. file_watcher @jinli
|
||||
a. 抽象基类
|
||||
ⅰ. on_start:
|
||||
1. file_store 的start 在前,加载graph,file_watcher在后,递归扫描目录
|
||||
a. 通过ms_time对比graph,on_change 进行改动
|
||||
ⅱ. on_change:
|
||||
1. 更新/增加:
|
||||
a. delete_chunks_by_path 更新数据库
|
||||
b. upate_chunks_by_path 更新数据库
|
||||
c. 更新graph
|
||||
2. 删除
|
||||
a. delete_chunks_by_path 更新数据库
|
||||
|
||||
MemorySchema
|
||||
1. markdown文件结构 @sen
|
||||
a. formatter:
|
||||
ⅰ. title
|
||||
ⅱ. desc
|
||||
ⅲ. tags
|
||||
ⅳ.
|
||||
2. memory文件结构目录
|
||||
a. MEMORY.md
|
||||
b. msg/files -> daily/YYYYMMDD/YYYYMMDD.md + xxxx.md
|
||||
ⅰ. YYYYMMDD.md
|
||||
1. xxx -> xxxx.md
|
||||
2. xxx -> xxxd.md
|
||||
ⅱ.
|
||||
c. daily -> topic/topic_l1/topic_l1.md + xxx.md + topic_l2
|
||||
d. proactive
|
||||
|
||||
steps:
|
||||
1. 治理(算法+LLM):
|
||||
a. 节点关联P0:现有的链接做补充,挖掘新的LLM的link
|
||||
ⅰ. /Users/yuli/workspace/ReMe/reme2/component/edge_extractor/llm_edge_extractor.py
|
||||
ⅱ. 移动到steps
|
||||
b. 节点整合/节点拆分/节点归档
|
||||
c. 健康度检查
|
||||
2. retrieve 调用store的检索
|
||||
3. 原子steps:reme edit
|
||||
4. 组合steps:总结:
|
||||
a. - freq (every_n_turn、compact) -> daily_summarizer
|
||||
b. topic (/dream ) -> topic_summarizer(daily_xx -> topic_xx)
|
||||
c. proactive -> proactive_summarizer(personal_xxx -> proactive_query - pre_query
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Qwenpaw
|
||||
270
docs4/service.yaml
Normal file
270
docs4/service.yaml
Normal file
|
|
@ -0,0 +1,270 @@
|
|||
app_name: reme-service
|
||||
enable_logo: false
|
||||
log_to_console: true
|
||||
log_to_file: false
|
||||
|
||||
# Service profile: service-aligned MCP surface — three tools, one per
|
||||
# memory service:
|
||||
# `retrieve` — graph-aware hybrid retrieval (Retriever projection)
|
||||
# `remember` — single write entry point (Ingestor projection)
|
||||
# * mode=log → zero-LLM event-folder upsert
|
||||
# * mode=distill → LLM R-M-W into topic graph
|
||||
# `maintain` — vault hygiene sweep (Maintainer projection)
|
||||
# runs lint + decay; merge/split require an LLM and
|
||||
# are off by default.
|
||||
#
|
||||
# The agent calls `remember(mode=log, name=…, content=…, materials=…)`
|
||||
# continuously through the task — picking a stable `name` per logical
|
||||
# thread so each call extends the same event folder rather than
|
||||
# fragmenting. At task completion (or PreCompact / SessionEnd) it
|
||||
# calls `remember(mode=distill, content=…, related_paths=…)` once to
|
||||
# distill the active events into topic-level cognition. `maintain` is
|
||||
# typically cron-driven, but the agent can invoke it explicitly when
|
||||
# it suspects vault drift.
|
||||
#
|
||||
# For full direct control over every memory_* primitive (raw writes,
|
||||
# fine-grained reads, separate sync/ingest tools, memory_lint) see
|
||||
# ./expert.yaml.
|
||||
|
||||
service:
|
||||
backend: mcp
|
||||
transport: stdio
|
||||
sidecar_http: true
|
||||
sidecar_http_host: "127.0.0.1"
|
||||
sidecar_http_port: 8765
|
||||
sidecar_info_path: "./vault/.reme/sidecar.json"
|
||||
|
||||
jobs:
|
||||
- backend: base
|
||||
name: retrieve
|
||||
description: |
|
||||
Graph-aware hybrid retrieval: vector + keyword + 1-hop wikilink
|
||||
BFS fusion. Use for "what do I know about X" / "did I work on Y" /
|
||||
"what's connected to [[Z]]". Anchor mode: include `[[Target]]` in
|
||||
the query to seed BFS at that file. Topic-rooted mode: pass `seeds`
|
||||
explicitly. Returns chunks ranked by combined relevance + graph
|
||||
proximity, each tagged with `graph_hop`.
|
||||
parameters:
|
||||
type: object
|
||||
properties:
|
||||
query: { type: string }
|
||||
max_results: { type: integer, default: 5 }
|
||||
min_score: { type: number, default: 0.0 }
|
||||
graph_depth:
|
||||
type: integer
|
||||
default: 1
|
||||
description: "BFS hops from seeds. 1 covers immediate neighbors."
|
||||
seeds:
|
||||
type: array
|
||||
items: { type: string }
|
||||
description: "Explicit seed paths (topic-rooted mode)."
|
||||
paths: { type: array, items: { type: string } }
|
||||
tags: { type: array, items: { type: string } }
|
||||
exclude_paths: { type: array, items: { type: string } }
|
||||
steps:
|
||||
- backend: memory_graph_search
|
||||
|
||||
- backend: base
|
||||
name: remember
|
||||
description: |
|
||||
Single write entry point — projects the Ingestor service. Two
|
||||
modes via the `mode` parameter:
|
||||
|
||||
* `mode: log` (zero LLM, hot path) — idempotent upsert of an
|
||||
event FOLDER under `events/{date}/{name}/`. The folder
|
||||
contains the index `{name}.md` (Event schema, frontmatter +
|
||||
narrative + Materials footer) plus any raw materials you
|
||||
pass — conversation snippets, tool outputs, data dumps. The
|
||||
watcher indexes everything inside.
|
||||
|
||||
CONTINUITY MODEL: pick a stable `name` per logical thread
|
||||
and call `remember(mode=log, ...)` repeatedly through the
|
||||
task. Each call extends the same folder:
|
||||
- new `content` → appended under a `## Update — {iso}`
|
||||
section
|
||||
- new `materials` → siblings (auto-suffix on filename
|
||||
collision)
|
||||
- `topics` + `tags` merged (union) into frontmatter
|
||||
- Materials footer regenerated to list every artifact
|
||||
|
||||
First call (folder doesn't exist) → CREATE; subsequent calls
|
||||
with the same `name` while the event is `status: active` →
|
||||
APPEND. If `status: distilled` / `archived`, REFUSES and
|
||||
returns `suggested_name` so you start a fresh thread instead
|
||||
of mutating prior cognition. Call CONTINUOUSLY through a task
|
||||
as facts land, especially at PreCompact to dump verbose raw
|
||||
text into `materials` before context truncation.
|
||||
|
||||
* `mode: distill` (LLM R-M-W loop, cold path) — DEFAULT. Run
|
||||
on EXPLICIT HANDOFF only: task completion / SessionEnd / when
|
||||
you decide the working set is ready. Not a per-turn tool.
|
||||
|
||||
Hand off the working set in two interchangeable forms:
|
||||
- `content` — inline material to distill (hint, summary, or
|
||||
raw text)
|
||||
- `related_paths` — pointers (event folder indexes; the
|
||||
Ingestor follows `## Materials` to read each artifact,
|
||||
individual material files, candidate topics).
|
||||
|
||||
The Ingestor reads the working set + linked topics, decides
|
||||
which existing topics to update / create, and flips each
|
||||
distilled event's status to "distilled". Returns an audit
|
||||
trail.
|
||||
parameters:
|
||||
type: object
|
||||
required: [content]
|
||||
properties:
|
||||
mode:
|
||||
type: string
|
||||
enum: [log, distill]
|
||||
default: distill
|
||||
description: "log = zero-LLM event-folder upsert (requires `name`); distill = LLM R-M-W into topic graph (default)."
|
||||
# mode=log params
|
||||
name:
|
||||
type: string
|
||||
description: "(mode=log) kebab-case event identifier (folder + index stem). Reuse the same name across calls in one thread to keep extending the same folder."
|
||||
description:
|
||||
type: string
|
||||
description: "(mode=log) one-line summary for index frontmatter (set on initial create only)."
|
||||
topics:
|
||||
type: array
|
||||
items: { type: string }
|
||||
description: "(mode=log) related topic wikilinks. Unioned into frontmatter on append."
|
||||
tags:
|
||||
type: array
|
||||
items: { type: string }
|
||||
description: "(mode=log) free-form tags; unioned on append."
|
||||
materials:
|
||||
type: array
|
||||
description: "(mode=log) raw artifacts written as siblings of the index. Filenames must be safe (letters/digits/dot/underscore/dash). Filename collision auto-suffixes (foo.txt → foo-2.txt)."
|
||||
items:
|
||||
type: object
|
||||
required: [filename, content]
|
||||
properties:
|
||||
filename: { type: string, description: "e.g. 'raw-prompt.md', 'tool-output.txt'" }
|
||||
content: { type: string }
|
||||
on_date:
|
||||
type: string
|
||||
description: "(mode=log) ISO date for events/{date}/ bucket; defaults to today."
|
||||
origin_session_id:
|
||||
type: string
|
||||
description: "(mode=log) optional source session identifier (set on initial create only)."
|
||||
# shared / mode=distill params
|
||||
content:
|
||||
type: string
|
||||
description: "Required. mode=log: markdown body for the index (initial create) or appended `## Update — {iso}` section. mode=distill: inline material the Ingestor distills — hint, summary, or raw text."
|
||||
hint:
|
||||
type: string
|
||||
description: "(mode=distill) caller guidance about target / intent."
|
||||
target_path:
|
||||
type: string
|
||||
description: "(mode=distill) optional suggested path; required for the no-LLM degraded path."
|
||||
metadata:
|
||||
type: object
|
||||
description: "(mode=distill) suggested frontmatter for any new topic."
|
||||
related_paths:
|
||||
type: array
|
||||
items: { type: string }
|
||||
description: "(mode=distill) pointers — event folder indexes (Ingestor reads materials from `## Materials`), individual material files, or candidate topics."
|
||||
steps:
|
||||
- backend: ingestor
|
||||
|
||||
- backend: base
|
||||
name: maintain
|
||||
description: |
|
||||
Vault hygiene sweep — projects the Maintainer service. One pass:
|
||||
scan signals → propose ops → resolve conflicts → apply. Returns
|
||||
an audit trail of what ran and what changed.
|
||||
|
||||
Default behavior runs `lint` (broken wikilinks, schema violations,
|
||||
stem collisions — read-only diagnostics) and `decay` (move stale
|
||||
events past their freshness window under `<vault>/Archive/`).
|
||||
Merge / split require an LLM and are off unless explicitly opted
|
||||
into via `ops`.
|
||||
|
||||
Typically cron-driven; the agent invokes it on demand when it
|
||||
suspects vault drift (after a heavy session of edits, after a
|
||||
bulk rename, etc.). Dry-run by default — flip `dry_run=false` to
|
||||
apply changes.
|
||||
parameters:
|
||||
type: object
|
||||
properties:
|
||||
target_prefix:
|
||||
type: string
|
||||
description: "restrict scan to relpaths starting with this prefix (e.g. 'events/2026-05-09/'). Empty string scans the whole vault."
|
||||
dry_run:
|
||||
type: boolean
|
||||
default: true
|
||||
description: "if true (default) returns the plan without mutating; flip to false to actually apply lint+decay (and merge/split if opted in)."
|
||||
ops:
|
||||
type: array
|
||||
items:
|
||||
type: string
|
||||
enum: [lint, decay, merge, split]
|
||||
description: "subset of ops to run. Defaults to ['lint','decay']. Merge/split are LLM-driven and currently scaffolded — enabling them without an LLM is a no-op."
|
||||
decay_days:
|
||||
type: integer
|
||||
description: "freshness window for the decay proposer (default = Maintainer constructor `decay_days`, typically 90)."
|
||||
steps:
|
||||
- backend: maintainer
|
||||
|
||||
components:
|
||||
# Ingestor LLM (opt-in). Without this the Ingestor degrades to a
|
||||
# direct create from explicit `target_path`; edits/renames/deletes
|
||||
# require the LLM. Uncomment + provide LLM_API_KEY to enable.
|
||||
#
|
||||
# as_llm:
|
||||
# default:
|
||||
# backend: openai
|
||||
# model_name: ${LLM_MODEL_NAME:-gpt-4o-mini}
|
||||
# api_key: ${LLM_API_KEY}
|
||||
# client_kwargs:
|
||||
# base_url: ${LLM_BASE_URL:-https://api.openai.com/v1}
|
||||
# stream: false
|
||||
#
|
||||
# as_llm_formatter:
|
||||
# default:
|
||||
# backend: openai
|
||||
|
||||
as_token_counter:
|
||||
default:
|
||||
backend: estimated
|
||||
|
||||
# Embedding is opt-in: leave embedding_model="" on file_store to run
|
||||
# keyword-only; uncomment + flip to "default" to enable hybrid search.
|
||||
#
|
||||
# embedding_model:
|
||||
# default:
|
||||
# backend: openai
|
||||
# model_name: ${EMBEDDING_MODEL_NAME:-text-embedding-3-small}
|
||||
# dimensions: 1536
|
||||
# pass_dimensions: false
|
||||
# enable_cache: true
|
||||
# max_batch_size: 10
|
||||
# max_cache_size: 2000
|
||||
# max_input_length: 8192
|
||||
|
||||
file_parser:
|
||||
md:
|
||||
backend: md
|
||||
default:
|
||||
backend: text
|
||||
|
||||
file_store:
|
||||
default:
|
||||
backend: local
|
||||
embedding_model: ""
|
||||
store_name: "reme"
|
||||
db_path: "./vault/.reme"
|
||||
working_dir: "./vault"
|
||||
|
||||
file_watcher:
|
||||
default:
|
||||
backend: full
|
||||
file_store: default
|
||||
default_parser: md
|
||||
recursive: true
|
||||
|
||||
# Retriever (`hybrid`) is a Step, not a pre-instantiated component —
|
||||
# the `query` shell builds it on demand. Tune defaults by attaching
|
||||
# knobs to the `memory_graph_search` step under the `query` job above.
|
||||
|
|
@ -0,0 +1,24 @@
|
|||
"""ReMe CLI package."""
|
||||
|
||||
from . import config
|
||||
from . import constants
|
||||
from . import enumeration
|
||||
from . import schema
|
||||
from . import steps
|
||||
from . import utils
|
||||
from .application import Application
|
||||
from .components import BaseComponent
|
||||
from .reme import ReMe
|
||||
|
||||
__all__ = [
|
||||
"Application",
|
||||
"BaseComponent",
|
||||
"ReMe",
|
||||
# submodules
|
||||
"config",
|
||||
"constants",
|
||||
"enumeration",
|
||||
"schema",
|
||||
"steps",
|
||||
"utils",
|
||||
]
|
||||
165
reme4/application.py
Normal file
165
reme4/application.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
import asyncio
|
||||
import heapq
|
||||
from pathlib import Path
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from .components import BaseComponent, ApplicationContext
|
||||
from .enumeration import ComponentEnum
|
||||
from .schema import Response, StreamChunk
|
||||
from .utils import execute_stream_task, print_logo, get_logger
|
||||
|
||||
|
||||
class Application(BaseComponent):
|
||||
"""Main application: initializes components, resolves dependencies, runs jobs."""
|
||||
|
||||
def __init__(self, **kwargs) -> None:
|
||||
self.context = ApplicationContext(**kwargs)
|
||||
|
||||
working_path = Path(self.config.working_dir).absolute()
|
||||
working_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if self.config.enable_logo:
|
||||
print_logo(self.config)
|
||||
|
||||
logger = get_logger(
|
||||
log_to_console=self.config.log_to_console,
|
||||
log_to_file=self.config.log_to_file,
|
||||
force_init=True,
|
||||
)
|
||||
logger.info(f"Initializing {self.config.app_name} Application")
|
||||
super().__init__()
|
||||
|
||||
from .components import R
|
||||
|
||||
# Service
|
||||
service_config = self.config.service
|
||||
if not service_config.backend:
|
||||
raise ValueError("Service configuration is missing the required 'backend' field")
|
||||
service_cls = R.get(ComponentEnum.SERVICE, service_config.backend)
|
||||
if not service_cls:
|
||||
raise ValueError(f"Unregistered service backend '{service_config.backend}'")
|
||||
params = service_config.model_dump()
|
||||
params["app_context"] = self.context
|
||||
self.context.service = service_cls(**params)
|
||||
|
||||
# Components
|
||||
for component_type, component_configs in self.config.components.items():
|
||||
self.context.components[component_type] = {}
|
||||
for name, config in component_configs.items():
|
||||
if not config.backend:
|
||||
raise ValueError(f"Component '{name}' is missing the required 'backend' field")
|
||||
backend_cls = R.get(component_type, config.backend)
|
||||
if not backend_cls:
|
||||
raise ValueError(f"Unregistered backend '{config.backend}' for component '{name}'")
|
||||
params = config.model_dump()
|
||||
params.setdefault("name", name)
|
||||
params["app_context"] = self.context
|
||||
self.context.components[component_type][name] = backend_cls(**params)
|
||||
|
||||
# Jobs
|
||||
for job_config in self.config.jobs:
|
||||
if not job_config.backend:
|
||||
raise ValueError(f"Job '{job_config.name}' is missing the required 'backend' field")
|
||||
job_cls = R.get(ComponentEnum.JOB, job_config.backend)
|
||||
if not job_cls:
|
||||
raise ValueError(f"Unregistered backend '{job_config.backend}' for job '{job_config.name}'")
|
||||
params = job_config.model_dump()
|
||||
params["app_context"] = self.context
|
||||
self.context.jobs[job_config.name] = job_cls(**params)
|
||||
|
||||
@property
|
||||
def config(self):
|
||||
return self.context.app_config
|
||||
|
||||
def _topological_order(self) -> list[BaseComponent]:
|
||||
"""Kahn's algorithm. Raises on missing required dep or cycle."""
|
||||
nodes: dict[tuple[ComponentEnum, str], BaseComponent] = {
|
||||
(ctype, name): comp for ctype, group in self.context.components.items() for name, comp in group.items()
|
||||
}
|
||||
|
||||
in_degree: dict[tuple[ComponentEnum, str], int] = dict.fromkeys(nodes, 0)
|
||||
dependents: dict[tuple[ComponentEnum, str], list[tuple[ComponentEnum, str]]] = {k: [] for k in nodes}
|
||||
for key, comp in nodes.items():
|
||||
for dep in comp.dependencies:
|
||||
dep_key = (dep.ctype, dep.name)
|
||||
if dep_key in nodes:
|
||||
dependents[dep_key].append(key)
|
||||
in_degree[key] += 1
|
||||
elif not dep.optional:
|
||||
raise ValueError(
|
||||
f"Component {key[0].value}:{key[1]} depends on {dep.ctype.value}:{dep.name}, not registered",
|
||||
)
|
||||
|
||||
ready = [k for k, d in in_degree.items() if d == 0]
|
||||
heapq.heapify(ready)
|
||||
ordered: list[BaseComponent] = []
|
||||
while ready:
|
||||
key = heapq.heappop(ready)
|
||||
ordered.append(nodes[key])
|
||||
for downstream in dependents[key]:
|
||||
in_degree[downstream] -= 1
|
||||
if in_degree[downstream] == 0:
|
||||
heapq.heappush(ready, downstream)
|
||||
|
||||
if len(ordered) != len(nodes):
|
||||
unresolved = [f"{k[0].value}:{k[1]}" for k, d in in_degree.items() if d > 0]
|
||||
raise ValueError(f"Circular dependency detected among: {unresolved}")
|
||||
return ordered
|
||||
|
||||
async def _start(self) -> None:
|
||||
"""Start components in topological order, then jobs."""
|
||||
start_order = self._topological_order()
|
||||
order_str = " -> ".join(f"{c.component_type.value}:{c.name}" for c in start_order)
|
||||
self.logger.info(f"Component start order: {order_str}")
|
||||
|
||||
for component in start_order:
|
||||
try:
|
||||
await component.start()
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to start {component.component_type.value}:{component.name}: {e}")
|
||||
|
||||
for name, job in self.context.jobs.items():
|
||||
try:
|
||||
await job.start()
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to start job '{name}': {e}")
|
||||
|
||||
async def _close(self) -> None:
|
||||
"""Close all jobs, then components in reverse."""
|
||||
for name, job in self.context.jobs.items():
|
||||
try:
|
||||
await job.close()
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to close job '{name}': {e}")
|
||||
|
||||
for components in self.context.components.values():
|
||||
for component in components.values():
|
||||
try:
|
||||
await component.close()
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to close {component.component_type.value}:{component.name}: {e}")
|
||||
|
||||
async def run_job(self, name: str, /, **kwargs) -> Response:
|
||||
"""Execute a registered job by name."""
|
||||
if name not in self.context.jobs:
|
||||
raise KeyError(f"Job '{name}' not found")
|
||||
return await self.context.jobs[name](**kwargs)
|
||||
|
||||
async def run_stream_job(self, name: str, /, **kwargs) -> AsyncGenerator[StreamChunk, None]:
|
||||
"""Execute a streaming job and yield chunks."""
|
||||
if name not in self.context.jobs:
|
||||
raise KeyError(f"Job '{name}' not found")
|
||||
job = self.context.jobs[name]
|
||||
stream_queue = asyncio.Queue()
|
||||
task = asyncio.create_task(job(stream_queue=stream_queue, **kwargs))
|
||||
async for chunk in execute_stream_task(
|
||||
stream_queue=stream_queue,
|
||||
task=task,
|
||||
task_name=name,
|
||||
output_format="chunk",
|
||||
):
|
||||
assert isinstance(chunk, StreamChunk)
|
||||
yield chunk
|
||||
|
||||
def run_app(self):
|
||||
self.context.service.run_app(app=self)
|
||||
43
reme4/components/__init__.py
Normal file
43
reme4/components/__init__.py
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
"""Components"""
|
||||
|
||||
from . import as_llm
|
||||
from . import as_llm_formatter
|
||||
from . import as_token_counter
|
||||
from . import client
|
||||
from . import embedding
|
||||
from . import file_graph
|
||||
from . import file_parser
|
||||
from . import file_store
|
||||
from . import file_watcher
|
||||
from . import job
|
||||
from . import keyword_index
|
||||
from . import service
|
||||
from . import tokenizer
|
||||
from .application_context import ApplicationContext
|
||||
from .base_component import BaseComponent
|
||||
from .component_registry import ComponentRegistry, R
|
||||
from .prompt_handler import PromptHandler
|
||||
from .runtime_context import RuntimeContext
|
||||
|
||||
__all__ = [
|
||||
"ApplicationContext",
|
||||
"BaseComponent",
|
||||
"ComponentRegistry",
|
||||
"R",
|
||||
"PromptHandler",
|
||||
"RuntimeContext",
|
||||
# base components
|
||||
"as_llm",
|
||||
"as_llm_formatter",
|
||||
"as_token_counter",
|
||||
"client",
|
||||
"embedding",
|
||||
"file_graph",
|
||||
"file_parser",
|
||||
"file_store",
|
||||
"file_watcher",
|
||||
"job",
|
||||
"keyword_index",
|
||||
"service",
|
||||
"tokenizer",
|
||||
]
|
||||
28
reme4/components/application_context.py
Normal file
28
reme4/components/application_context.py
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
"""Application context: shared state container for components, jobs, and service."""
|
||||
|
||||
from ..enumeration import ComponentEnum
|
||||
from ..schema import ApplicationConfig
|
||||
|
||||
|
||||
class ApplicationContext:
|
||||
"""Holds the parsed config and instantiated components, jobs, and service.
|
||||
|
||||
Acts as a passive state container. The actual wiring (resolving backends from
|
||||
the registry and instantiating each component) is performed by Application.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
# Parse and validate raw config kwargs into a typed ApplicationConfig.
|
||||
self.app_config: ApplicationConfig = ApplicationConfig(**kwargs)
|
||||
|
||||
# Local imports to avoid circular dependencies during module init.
|
||||
from .base_component import BaseComponent
|
||||
from .job import BaseJob
|
||||
from .service import BaseService
|
||||
|
||||
# Service endpoint (e.g. HTTP/MCP). Populated by Application.__init__.
|
||||
self.service: BaseService | None = None
|
||||
# Components keyed by type then by user-defined name.
|
||||
self.components: dict[ComponentEnum, dict[str, BaseComponent]] = {}
|
||||
# Jobs keyed by user-defined name.
|
||||
self.jobs: dict[str, BaseJob] = {}
|
||||
53
reme4/components/as_llm/__init__.py
Normal file
53
reme4/components/as_llm/__init__.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
"""AgentScope LLM model wrappers."""
|
||||
|
||||
from agentscope.model import AnthropicChatModel, ChatModelBase, OpenAIChatModel
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ..component_registry import R
|
||||
from ...enumeration import ComponentEnum
|
||||
|
||||
|
||||
class BaseAsLLM(BaseComponent):
|
||||
"""Base wrapper for AgentScope chat models. Builds ``self.model`` in ``_start``."""
|
||||
|
||||
component_type = ComponentEnum.AS_LLM
|
||||
|
||||
def __init__(self, **kwargs) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.model: ChatModelBase | None = None
|
||||
|
||||
async def _close(self) -> None:
|
||||
self.model = None
|
||||
|
||||
|
||||
@R.register("openai")
|
||||
class OpenAIAsLLM(BaseAsLLM):
|
||||
"""OpenAI chat model wrapper."""
|
||||
|
||||
async def _start(self) -> None:
|
||||
self.model = OpenAIChatModel(**self.kwargs)
|
||||
|
||||
async def _close(self) -> None:
|
||||
if self.model is not None:
|
||||
assert isinstance(self.model, OpenAIChatModel)
|
||||
await self.model.client.close()
|
||||
|
||||
|
||||
@R.register("anthropic")
|
||||
class AnthropicAsLLM(BaseAsLLM):
|
||||
"""Anthropic chat model wrapper."""
|
||||
|
||||
async def _start(self) -> None:
|
||||
self.model = AnthropicChatModel(**self.kwargs)
|
||||
|
||||
async def _close(self) -> None:
|
||||
if self.model is not None:
|
||||
assert isinstance(self.model, AnthropicChatModel)
|
||||
await self.model.client.close()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BaseAsLLM",
|
||||
"OpenAIAsLLM",
|
||||
"AnthropicAsLLM",
|
||||
]
|
||||
44
reme4/components/as_llm_formatter/__init__.py
Normal file
44
reme4/components/as_llm_formatter/__init__.py
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
"""AgentScope LLM formatter wrappers."""
|
||||
|
||||
from agentscope.formatter import AnthropicChatFormatter, FormatterBase
|
||||
|
||||
from .reme_openai_chat_formatter import ReMeOpenAIChatFormatter
|
||||
from ..base_component import BaseComponent
|
||||
from ..component_registry import R
|
||||
from ...enumeration import ComponentEnum
|
||||
|
||||
|
||||
class BaseAsLLMFormatter(BaseComponent):
|
||||
"""Base wrapper for AgentScope formatters. Builds ``self.formatter`` in ``_start``."""
|
||||
|
||||
component_type = ComponentEnum.AS_LLM_FORMATTER
|
||||
|
||||
def __init__(self, **kwargs) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.formatter: FormatterBase | None = None
|
||||
|
||||
async def _close(self) -> None:
|
||||
self.formatter = None
|
||||
|
||||
|
||||
@R.register("openai")
|
||||
class AsOpenAIChatFormatter(BaseAsLLMFormatter):
|
||||
"""OpenAI chat formatter wrapper (uses ReMe extensions)."""
|
||||
|
||||
async def _start(self) -> None:
|
||||
self.formatter = ReMeOpenAIChatFormatter(**self.kwargs)
|
||||
|
||||
|
||||
@R.register("anthropic")
|
||||
class AsAnthropicChatFormatter(BaseAsLLMFormatter):
|
||||
"""Anthropic chat formatter wrapper."""
|
||||
|
||||
async def _start(self) -> None:
|
||||
self.formatter = AnthropicChatFormatter(**self.kwargs)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BaseAsLLMFormatter",
|
||||
"AsOpenAIChatFormatter",
|
||||
"AsAnthropicChatFormatter",
|
||||
]
|
||||
139
reme4/components/as_llm_formatter/reme_openai_chat_formatter.py
Normal file
139
reme4/components/as_llm_formatter/reme_openai_chat_formatter.py
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
"""OpenAI chat formatter with ReMe extensions: image promotion and reasoning_content."""
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from agentscope.formatter import OpenAIChatFormatter
|
||||
from agentscope.formatter import _openai_formatter as _of
|
||||
|
||||
_format_openai_image_block = getattr(_of, "_format_openai_image_block")
|
||||
_to_openai_audio_data = getattr(_of, "_to_openai_audio_data")
|
||||
from agentscope.message import ImageBlock, Msg, TextBlock, URLSource
|
||||
|
||||
|
||||
def _format_openai_video_block(video_block: dict) -> dict[str, Any]:
|
||||
"""Convert a video block to OpenAI ``video_url`` content."""
|
||||
source = video_block["source"]
|
||||
if source["type"] == "url":
|
||||
url = source["url"]
|
||||
elif source["type"] == "base64":
|
||||
url = f"data:{source['media_type']};base64,{source['data']}"
|
||||
else:
|
||||
raise ValueError(f"Unsupported video source type: {source['type']}")
|
||||
return {"type": "video_url", "video_url": {"url": url}}
|
||||
|
||||
|
||||
class ReMeOpenAIChatFormatter(OpenAIChatFormatter):
|
||||
"""OpenAIChatFormatter + tool-result image promotion + reasoning_content passthrough."""
|
||||
|
||||
async def _format(self, msgs: list[Msg]) -> list[dict[str, Any]]:
|
||||
"""Format ``Msg`` list into OpenAI chat-completion message dicts."""
|
||||
self.assert_list_of_msgs(msgs)
|
||||
|
||||
messages: list[dict] = []
|
||||
i = 0
|
||||
while i < len(msgs):
|
||||
msg = msgs[i]
|
||||
content_blocks = []
|
||||
tool_calls = []
|
||||
reasoning_content_blocks = []
|
||||
|
||||
for block in msg.get_content_blocks():
|
||||
typ = block.get("type")
|
||||
|
||||
if typ == "text":
|
||||
content_blocks.append({**block})
|
||||
|
||||
elif typ == "thinking":
|
||||
reasoning_content_blocks.append({**block})
|
||||
|
||||
elif typ == "tool_use":
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": block.get("id"),
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": block.get("name"),
|
||||
"arguments": json.dumps(block.get("input", {}), ensure_ascii=False),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
elif typ == "tool_result":
|
||||
textual_output, multimodal_data = self.convert_tool_result_to_string(block["output"])
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": block.get("id"),
|
||||
"content": textual_output,
|
||||
"name": block.get("name"),
|
||||
},
|
||||
)
|
||||
|
||||
# OpenAI tool messages can't carry images; promote to a follow-up user message.
|
||||
promoted_blocks = []
|
||||
for url, multimodal_block in multimodal_data:
|
||||
if multimodal_block["type"] == "image" and self.promote_tool_result_images:
|
||||
promoted_blocks.extend(
|
||||
[
|
||||
TextBlock(type="text", text=f"\n- The image from '{url}': "),
|
||||
ImageBlock(type="image", source=URLSource(type="url", url=url)),
|
||||
],
|
||||
)
|
||||
|
||||
if promoted_blocks:
|
||||
promoted_blocks = [
|
||||
TextBlock(
|
||||
type="text",
|
||||
text="<system-info>The following are the image contents from the tool "
|
||||
f"result of '{block['name']}':",
|
||||
),
|
||||
*promoted_blocks,
|
||||
TextBlock(type="text", text="</system-info>"),
|
||||
]
|
||||
msgs.insert(
|
||||
i + 1,
|
||||
Msg(name="user", content=promoted_blocks, role="user"),
|
||||
)
|
||||
|
||||
elif typ == "image":
|
||||
content_blocks.append(_format_openai_image_block(block))
|
||||
|
||||
elif typ == "audio":
|
||||
# Skip assistant audio — not a valid input modality.
|
||||
if msg.role == "assistant":
|
||||
continue
|
||||
content_blocks.append(
|
||||
{
|
||||
"type": "input_audio",
|
||||
"input_audio": _to_openai_audio_data(block["source"]),
|
||||
},
|
||||
)
|
||||
|
||||
elif typ == "video":
|
||||
# Skip assistant video — not a valid input modality.
|
||||
if msg.role == "assistant":
|
||||
continue
|
||||
content_blocks.append(_format_openai_video_block(block))
|
||||
|
||||
msg_openai = {
|
||||
"role": msg.role,
|
||||
"name": msg.name,
|
||||
"content": content_blocks or None,
|
||||
}
|
||||
|
||||
if tool_calls:
|
||||
msg_openai["tool_calls"] = tool_calls
|
||||
|
||||
# Merge thinking blocks into reasoning_content for compatible models.
|
||||
if reasoning_content_blocks:
|
||||
reasoning_msg = "\n".join(r.get("thinking", "") for r in reasoning_content_blocks)
|
||||
if reasoning_msg:
|
||||
msg_openai["reasoning_content"] = reasoning_msg
|
||||
|
||||
if msg_openai["content"] or msg_openai.get("tool_calls"):
|
||||
messages.append(msg_openai)
|
||||
|
||||
i += 1
|
||||
|
||||
return messages
|
||||
35
reme4/components/as_token_counter/__init__.py
Normal file
35
reme4/components/as_token_counter/__init__.py
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
"""AgentScope token counter wrappers."""
|
||||
|
||||
from agentscope.token import TokenCounterBase
|
||||
|
||||
from .estimate_token_counter import EstimatedTokenCounter
|
||||
from ..base_component import BaseComponent
|
||||
from ..component_registry import R
|
||||
from ...enumeration import ComponentEnum
|
||||
|
||||
|
||||
class BaseAsTokenCounter(BaseComponent):
|
||||
"""Base wrapper for AgentScope token counters. Builds ``self.token_counter`` in ``_start``."""
|
||||
|
||||
component_type = ComponentEnum.AS_TOKEN_COUNTER
|
||||
|
||||
def __init__(self, **kwargs) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.token_counter: TokenCounterBase | None = None
|
||||
|
||||
async def _close(self) -> None:
|
||||
self.token_counter = None
|
||||
|
||||
|
||||
@R.register("estimated")
|
||||
class EstimatedAsTokenCounter(BaseAsTokenCounter):
|
||||
"""Character-based estimated token counter — fast but approximate."""
|
||||
|
||||
async def _start(self) -> None:
|
||||
self.token_counter = EstimatedTokenCounter(**self.kwargs)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BaseAsTokenCounter",
|
||||
"EstimatedAsTokenCounter",
|
||||
]
|
||||
21
reme4/components/as_token_counter/estimate_token_counter.py
Normal file
21
reme4/components/as_token_counter/estimate_token_counter.py
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
"""Character-based token-count estimator."""
|
||||
|
||||
from agentscope.token import TokenCounterBase
|
||||
|
||||
|
||||
class EstimatedTokenCounter(TokenCounterBase):
|
||||
"""Approximate token count as ``encoded_byte_len / divisor``.
|
||||
|
||||
Cheap proxy when exact counts aren't needed; use the model's real
|
||||
tokenizer for accuracy.
|
||||
"""
|
||||
|
||||
def __init__(self, estimate_divisor: float = 4, encoding: str = "utf-8"):
|
||||
if estimate_divisor <= 0:
|
||||
raise ValueError("estimate_divisor must be positive")
|
||||
self.estimate_divisor: float = estimate_divisor
|
||||
self.encoding: str = encoding
|
||||
|
||||
async def count(self, text: str, **kwargs) -> int:
|
||||
"""Estimated token count for ``text``."""
|
||||
return int(len(text.encode(self.encoding)) / self.estimate_divisor + 0.5)
|
||||
168
reme4/components/base_component.py
Normal file
168
reme4/components/base_component.py
Normal file
|
|
@ -0,0 +1,168 @@
|
|||
"""Base class for components."""
|
||||
|
||||
import asyncio
|
||||
from abc import ABC
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, TypeVar, cast
|
||||
|
||||
from .application_context import ApplicationContext
|
||||
from ..enumeration import ComponentEnum
|
||||
from ..utils import get_logger
|
||||
|
||||
T = TypeVar("T", bound="BaseComponent")
|
||||
|
||||
|
||||
class Dependency:
|
||||
"""Declared dependency: bind() return value, instance attribute placeholder, and topological-sort edge."""
|
||||
|
||||
__slots__ = ("ctype", "name", "default_factory", "optional")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ctype: ComponentEnum,
|
||||
name: str,
|
||||
default_factory: Callable[[], Any] | None = None,
|
||||
optional: bool = True,
|
||||
) -> None:
|
||||
self.ctype = ctype
|
||||
self.name = name
|
||||
self.default_factory = default_factory
|
||||
self.optional = optional
|
||||
|
||||
def __repr__(self) -> str:
|
||||
suffix = "?" if self.optional else ""
|
||||
return f"<unresolved {self.ctype.value}:{self.name}{suffix}>"
|
||||
|
||||
def __getattr__(self, item: str) -> Any:
|
||||
# Guard against using the dependency before start() resolves it.
|
||||
raise RuntimeError(
|
||||
f"Dependency {self.ctype.value}:{self.name} accessed before start() (attribute '{item}')",
|
||||
)
|
||||
|
||||
|
||||
class BaseComponent(ABC):
|
||||
"""Async lifecycle base class with bind-based dependency injection."""
|
||||
|
||||
component_type = ComponentEnum.BASE
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str | None = None,
|
||||
backend: str = "",
|
||||
app_context: "ApplicationContext | None" = None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
self.name: str = name or self.__class__.__name__
|
||||
self.backend: str = backend
|
||||
self.app_context: "ApplicationContext | None" = app_context
|
||||
self.kwargs: dict = dict(kwargs)
|
||||
self.logger = get_logger()
|
||||
if hasattr(self.logger, "bind"):
|
||||
self.logger = self.logger.bind(component=self.name)
|
||||
|
||||
self._is_started: bool = False
|
||||
self._lock: asyncio.Lock = asyncio.Lock()
|
||||
# Components created from bind() default_factory in standalone mode (auto-managed lifecycle).
|
||||
self._owned: list["BaseComponent"] = []
|
||||
|
||||
@property
|
||||
def is_started(self) -> bool:
|
||||
return self._is_started
|
||||
|
||||
# ----- Dependency declaration ----------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def bind(
|
||||
name: str | None,
|
||||
base_cls: type[T],
|
||||
*,
|
||||
default_factory: Callable[[], T] | None = None,
|
||||
optional: bool = True,
|
||||
) -> T | None:
|
||||
"""Declare a dependency on another component; resolved at start(). Empty name → None."""
|
||||
if not name:
|
||||
return None
|
||||
ctype = getattr(base_cls, "component_type", None)
|
||||
if not isinstance(ctype, ComponentEnum) or ctype is ComponentEnum.BASE:
|
||||
raise TypeError(f"{base_cls.__name__} must declare a non-BASE ComponentEnum 'component_type'")
|
||||
return cast(T, Dependency(ctype, name, default_factory, optional))
|
||||
|
||||
@property
|
||||
def dependencies(self) -> list[Dependency]:
|
||||
"""All unresolved bindings declared on this instance."""
|
||||
return [v for v in self.__dict__.values() if isinstance(v, Dependency)]
|
||||
|
||||
async def _resolve_bindings(self) -> None:
|
||||
"""Replace Dependency placeholders with real components (or default_factory / None for optional)."""
|
||||
for attr, value in list(self.__dict__.items()):
|
||||
if not isinstance(value, Dependency):
|
||||
continue
|
||||
if self.app_context is None:
|
||||
# Standalone mode: factory or (optional → None) or keep placeholder.
|
||||
if value.default_factory is not None:
|
||||
instance = value.default_factory()
|
||||
setattr(self, attr, instance)
|
||||
if isinstance(instance, BaseComponent):
|
||||
self._owned.append(instance)
|
||||
elif value.optional:
|
||||
setattr(self, attr, None)
|
||||
else:
|
||||
target = self.app_context.components.get(value.ctype, {}).get(value.name)
|
||||
if target is not None:
|
||||
setattr(self, attr, target)
|
||||
elif value.optional:
|
||||
setattr(self, attr, None)
|
||||
else:
|
||||
raise ValueError(f"{value.ctype.value} '{value.name}' not found.")
|
||||
|
||||
# ----- Lookup --------------------------------------------------------
|
||||
|
||||
@property
|
||||
def working_path(self) -> Path:
|
||||
if self.app_context is None:
|
||||
return Path.cwd()
|
||||
return Path(self.app_context.app_config.working_dir)
|
||||
|
||||
# ----- Lifecycle -----------------------------------------------------
|
||||
|
||||
async def _start(self) -> None:
|
||||
"""Subclass hook: start logic."""
|
||||
|
||||
async def _close(self) -> None:
|
||||
"""Subclass hook: close logic."""
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Resolve bindings → start owned fallbacks → _start(). No-op if already started."""
|
||||
async with self._lock:
|
||||
if self._is_started:
|
||||
return
|
||||
await self._resolve_bindings()
|
||||
for owned in self._owned:
|
||||
await owned.start()
|
||||
await self._start()
|
||||
self._is_started = True
|
||||
|
||||
async def close(self) -> None:
|
||||
"""_close() → close owned fallbacks in reverse. No-op if not started."""
|
||||
async with self._lock:
|
||||
if not self._is_started:
|
||||
return
|
||||
await self._close()
|
||||
for owned in reversed(self._owned):
|
||||
await owned.close()
|
||||
self._is_started = False
|
||||
|
||||
async def restart(self) -> None:
|
||||
"""Close then start."""
|
||||
await self.close()
|
||||
await self.start()
|
||||
|
||||
async def __call__(self, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
async def __aenter__(self) -> "BaseComponent":
|
||||
await self.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
|
||||
await self.close()
|
||||
6
reme4/components/client/__init__.py
Normal file
6
reme4/components/client/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
"""Client components."""
|
||||
|
||||
from .base_client import BaseClient
|
||||
from .http_client import HttpClient
|
||||
|
||||
__all__ = ["BaseClient", "HttpClient"]
|
||||
26
reme4/components/client/base_client.py
Normal file
26
reme4/components/client/base_client.py
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
"""Base client abstraction."""
|
||||
|
||||
from abc import abstractmethod
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ...enumeration import ComponentEnum
|
||||
|
||||
|
||||
class BaseClient(BaseComponent):
|
||||
"""Abstract base for clients that communicate with ReMe services."""
|
||||
|
||||
component_type = ComponentEnum.CLIENT
|
||||
|
||||
def __init__(self, **kwargs) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.client = None
|
||||
|
||||
async def _start(self) -> None:
|
||||
"""Initialize the client."""
|
||||
|
||||
async def _close(self) -> None:
|
||||
"""Close the client and release resources."""
|
||||
|
||||
@abstractmethod
|
||||
async def __call__(self) -> dict:
|
||||
"""Execute the configured action and return the response."""
|
||||
62
reme4/components/client/http_client.py
Normal file
62
reme4/components/client/http_client.py
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
"""HTTP client for ReMe services."""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import httpx
|
||||
|
||||
from .base_client import BaseClient
|
||||
from ..component_registry import R
|
||||
from ...constants import REME_SERVICE_INFO, REME_DEFAULT_HOST, REME_DEFAULT_PORT
|
||||
|
||||
|
||||
@R.register("http")
|
||||
class HttpClient(BaseClient):
|
||||
"""HTTP client that communicates with ReMe service via REST API."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
action: str,
|
||||
host: str | None = None,
|
||||
port: int | None = None,
|
||||
timeout: float = 30.0,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Resolve host/port: explicit args > env var > defaults
|
||||
if not (host and port):
|
||||
if service_info := os.environ.get(REME_SERVICE_INFO):
|
||||
try:
|
||||
data = json.loads(service_info)
|
||||
host = data["host"]
|
||||
port = data["port"]
|
||||
except Exception:
|
||||
self.logger.warning(f"Invalid service info: {service_info}")
|
||||
host, port = REME_DEFAULT_HOST, REME_DEFAULT_PORT
|
||||
else:
|
||||
host, port = REME_DEFAULT_HOST, REME_DEFAULT_PORT
|
||||
|
||||
self.action = action
|
||||
self.base_url = f"http://{host}:{port}"
|
||||
self.timeout = timeout
|
||||
|
||||
async def _start(self) -> None:
|
||||
"""Initialize the HTTP client."""
|
||||
if self.client is None:
|
||||
self.client = httpx.AsyncClient(base_url=self.base_url, timeout=self.timeout)
|
||||
|
||||
async def __call__(self) -> dict:
|
||||
"""Send POST request to the configured action endpoint."""
|
||||
if self.client is None:
|
||||
raise RuntimeError("Client not initialized. Call _start() first.")
|
||||
|
||||
response = await self.client.post(f"/{self.action}", json=self.kwargs)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def _close(self) -> None:
|
||||
"""Close the HTTP client."""
|
||||
if self.client is not None:
|
||||
await self.client.aclose()
|
||||
self.client = None
|
||||
77
reme4/components/component_registry.py
Normal file
77
reme4/components/component_registry.py
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
"""Global registry mapping (ComponentEnum, name) -> component class."""
|
||||
|
||||
from typing import Callable, TypeVar, cast
|
||||
|
||||
from .base_component import BaseComponent
|
||||
from ..enumeration import ComponentEnum
|
||||
from ..utils import get_logger
|
||||
|
||||
T = TypeVar("T", bound=BaseComponent)
|
||||
|
||||
|
||||
class ComponentRegistry:
|
||||
"""Two-level registry: component_type -> name -> class.
|
||||
|
||||
Supports both direct calls — ``R.register(MyClass, "name")`` — and
|
||||
decorator usage — ``@R.register("name")``.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._registry: dict[ComponentEnum, dict[str, type[BaseComponent]]] = {}
|
||||
self.logger = get_logger()
|
||||
|
||||
def _do_register(self, cls: type[T], name: str) -> type[T]:
|
||||
"""Insert `cls` under its `component_type` group; warn on overwrite."""
|
||||
component_type = getattr(cls, "component_type", None)
|
||||
if not isinstance(component_type, ComponentEnum):
|
||||
raise TypeError(f"{cls.__name__} must have a ComponentEnum 'component_type' attribute")
|
||||
if not name:
|
||||
raise ValueError("Component name cannot be empty")
|
||||
|
||||
group = self._registry.setdefault(component_type, {})
|
||||
if name in group:
|
||||
self.logger.warning(f"Component '{name}' already registered for {component_type}, overwriting")
|
||||
group[name] = cls
|
||||
return cls
|
||||
|
||||
def register(
|
||||
self,
|
||||
cls_or_name: type[T] | str,
|
||||
name: str | None = None,
|
||||
) -> Callable[[type[T]], type[T]] | type[T]:
|
||||
"""Register a component class directly, or return a decorator that does so."""
|
||||
# Direct mode: first arg is the class itself.
|
||||
if isinstance(cls_or_name, type):
|
||||
return self._do_register(cast(type[T], cls_or_name), name if name is not None else cls_or_name.__name__)
|
||||
|
||||
# Decorator mode: first arg is the registration name.
|
||||
if not isinstance(cls_or_name, str):
|
||||
raise TypeError(f"Expected a class or string, got {type(cls_or_name).__name__}")
|
||||
|
||||
def decorator(decorated_cls: type[T]) -> type[T]:
|
||||
return self._do_register(decorated_cls, cls_or_name)
|
||||
|
||||
return decorator
|
||||
|
||||
def get(self, component_type: ComponentEnum, name: str) -> type[BaseComponent] | None:
|
||||
"""Look up a registered class; return None if not found."""
|
||||
return self._registry.get(component_type, {}).get(name)
|
||||
|
||||
def get_all(self, component_type: ComponentEnum) -> dict[str, type[BaseComponent]]:
|
||||
"""Return a shallow copy of all classes registered under `component_type`."""
|
||||
return dict(self._registry.get(component_type, {}))
|
||||
|
||||
def unregister(self, component_type: ComponentEnum, name: str) -> bool:
|
||||
"""Remove an entry; return True if it existed, False otherwise."""
|
||||
if (group := self._registry.get(component_type)) and name in group:
|
||||
del group[name]
|
||||
return True
|
||||
return False
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Drop every registered entry."""
|
||||
self._registry.clear()
|
||||
|
||||
|
||||
# Process-wide singleton used throughout the codebase.
|
||||
R = ComponentRegistry()
|
||||
6
reme4/components/embedding/__init__.py
Normal file
6
reme4/components/embedding/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
"""Embedding model implementations."""
|
||||
|
||||
from .base_embedding_model import BaseEmbeddingModel
|
||||
from .openai_embedding_model import OpenAIEmbeddingModel
|
||||
|
||||
__all__ = ["BaseEmbeddingModel", "OpenAIEmbeddingModel"]
|
||||
194
reme4/components/embedding/base_embedding_model.py
Normal file
194
reme4/components/embedding/base_embedding_model.py
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
"""Base embedding model with LRU cache and disk persistence."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
from abc import abstractmethod
|
||||
from collections import OrderedDict
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ...enumeration import ComponentEnum
|
||||
from ...schema import EmbNode
|
||||
|
||||
|
||||
class BaseEmbeddingModel(BaseComponent):
|
||||
"""Embedding model with LRU cache and disk persistence."""
|
||||
|
||||
component_type = ComponentEnum.EMBEDDING_MODEL
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
model_name: str = "",
|
||||
dimensions: int = 1024,
|
||||
pass_dimensions: bool = False,
|
||||
max_batch_size: int = 10,
|
||||
max_input_length: int = 8192,
|
||||
max_cache_size: int = 10000,
|
||||
enable_cache: bool = True,
|
||||
max_retries: int = 3,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.api_key = api_key or os.environ.get("EMBEDDING_API_KEY", "")
|
||||
self.base_url = base_url or os.environ.get("EMBEDDING_BASE_URL", "")
|
||||
self.model_name = model_name
|
||||
self.dimensions = dimensions
|
||||
self.pass_dimensions = pass_dimensions
|
||||
self.max_batch_size = max_batch_size
|
||||
self.max_input_length = max_input_length
|
||||
self.max_cache_size = max_cache_size
|
||||
self.enable_cache = enable_cache
|
||||
self.max_retries = max_retries
|
||||
self._embedding_cache: OrderedDict[str, np.ndarray] = OrderedDict()
|
||||
|
||||
@property
|
||||
def cache_path(self) -> Path:
|
||||
"""Disk path for the embedding cache file."""
|
||||
return self.working_path / "embedding_cache" / f"{self.name}.npz"
|
||||
|
||||
async def _start(self) -> None:
|
||||
"""Load cache from disk on startup."""
|
||||
self._embedding_cache.clear()
|
||||
self._load_cache()
|
||||
|
||||
async def _close(self) -> None:
|
||||
"""Persist cache to disk on shutdown."""
|
||||
self._save_cache()
|
||||
|
||||
# -- Public API --
|
||||
|
||||
async def get_embedding(self, input_text: str, **kwargs) -> list[float] | None:
|
||||
"""Get embedding for a single text."""
|
||||
results = await self.get_embeddings([input_text], **kwargs)
|
||||
return results[0] if results else None
|
||||
|
||||
async def get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]:
|
||||
"""Get embeddings for a list of texts, with caching and batching."""
|
||||
truncated = [t[: self.max_input_length] for t in input_text]
|
||||
results: list[list[float] | None] = [None] * len(truncated)
|
||||
to_compute: list[tuple[int, str]] = []
|
||||
|
||||
# Split into cache hits and misses
|
||||
for idx, text in enumerate(truncated):
|
||||
cached = self._get_from_cache(text)
|
||||
if cached is not None:
|
||||
results[idx] = cached.tolist()
|
||||
else:
|
||||
to_compute.append((idx, text))
|
||||
|
||||
# Batch-compute misses with retry
|
||||
if to_compute:
|
||||
for i in range(0, len(to_compute), self.max_batch_size):
|
||||
batch = to_compute[i : i + self.max_batch_size]
|
||||
indices = [idx for idx, _ in batch]
|
||||
texts = [text for _, text in batch]
|
||||
|
||||
embeddings = None
|
||||
for attempt in range(self.max_retries):
|
||||
try:
|
||||
embeddings = await self._get_embeddings(texts, **kwargs)
|
||||
if embeddings and len(embeddings) == len(texts):
|
||||
break
|
||||
except (TimeoutError, ConnectionError, OSError):
|
||||
if attempt < self.max_retries - 1:
|
||||
await asyncio.sleep(2**attempt)
|
||||
except Exception:
|
||||
self.logger.exception("Embedding request failed")
|
||||
break
|
||||
|
||||
if not embeddings or len(embeddings) != len(texts):
|
||||
continue
|
||||
|
||||
# Normalize dimensions and cache
|
||||
for orig_idx, text, emb in zip(indices, texts, embeddings):
|
||||
if emb is None:
|
||||
continue
|
||||
emb_array = np.asarray(emb, dtype=np.float16)
|
||||
if len(emb_array) != self.dimensions:
|
||||
if len(emb_array) < self.dimensions:
|
||||
emb_array = np.pad(emb_array, (0, self.dimensions - len(emb_array)))
|
||||
else:
|
||||
emb_array = emb_array[: self.dimensions]
|
||||
results[orig_idx] = emb_array.tolist()
|
||||
self._put_to_cache(text, emb_array)
|
||||
|
||||
return results
|
||||
|
||||
async def get_node_embeddings(self, nodes: list[EmbNode], **kwargs) -> list[EmbNode]:
|
||||
"""Compute and assign embeddings for EmbNode objects."""
|
||||
embeddings = await self.get_embeddings([n.text for n in nodes], **kwargs)
|
||||
if len(embeddings) == len(nodes):
|
||||
for node, vec in zip(nodes, embeddings):
|
||||
if vec is not None:
|
||||
node.embedding = vec
|
||||
return nodes
|
||||
|
||||
@abstractmethod
|
||||
async def _get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]:
|
||||
"""Get raw embeddings from the underlying provider."""
|
||||
|
||||
# -- Cache Operations --
|
||||
|
||||
def _get_from_cache(self, text: str) -> np.ndarray | None:
|
||||
"""Lookup text in LRU cache, promoting on hit."""
|
||||
if not self.enable_cache:
|
||||
return None
|
||||
key = self._get_cache_key(text)
|
||||
if key not in self._embedding_cache:
|
||||
return None
|
||||
self._embedding_cache.move_to_end(key)
|
||||
return self._embedding_cache[key]
|
||||
|
||||
def _put_to_cache(self, text: str, embedding: np.ndarray) -> None:
|
||||
"""Insert into LRU cache, evicting oldest if full."""
|
||||
if not self.enable_cache or self.max_cache_size <= 0 or len(embedding) != self.dimensions:
|
||||
return
|
||||
key = self._get_cache_key(text)
|
||||
if len(self._embedding_cache) >= self.max_cache_size and key not in self._embedding_cache:
|
||||
self._embedding_cache.popitem(last=False)
|
||||
self._embedding_cache[key] = embedding
|
||||
self._embedding_cache.move_to_end(key)
|
||||
|
||||
def _get_cache_key(self, text: str) -> str:
|
||||
"""Generate cache key from text, model name, and dimensions."""
|
||||
return hashlib.sha256(f"{text}|{self.model_name}|{self.dimensions}".encode()).hexdigest()
|
||||
|
||||
# -- Cache Persistence --
|
||||
|
||||
def _load_cache(self) -> None:
|
||||
"""Load cached embeddings from disk (npz format)."""
|
||||
if not self.enable_cache:
|
||||
return
|
||||
self.cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not self.cache_path.exists():
|
||||
return
|
||||
|
||||
try:
|
||||
data = np.load(self.cache_path)
|
||||
except Exception:
|
||||
self.logger.exception("Failed to load embedding cache, removing")
|
||||
self.cache_path.unlink(missing_ok=True)
|
||||
return
|
||||
|
||||
for key, emb in zip(data["keys"], data["embeddings"]):
|
||||
if len(emb) != self.dimensions:
|
||||
continue
|
||||
if len(self._embedding_cache) >= self.max_cache_size:
|
||||
break
|
||||
self._embedding_cache[str(key)] = emb.astype(np.float16)
|
||||
|
||||
def _save_cache(self) -> None:
|
||||
"""Persist in-memory cache to disk (npz format)."""
|
||||
if not self.enable_cache or not self._embedding_cache:
|
||||
return
|
||||
keys = list(self._embedding_cache.keys())
|
||||
embeddings = np.stack(list(self._embedding_cache.values()))
|
||||
try:
|
||||
np.savez(self.cache_path, keys=np.array(keys, dtype=str), embeddings=embeddings)
|
||||
except Exception:
|
||||
self.logger.exception("Failed to save embedding cache")
|
||||
52
reme4/components/embedding/openai_embedding_model.py
Normal file
52
reme4/components/embedding/openai_embedding_model.py
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
"""OpenAI-compatible async embedding model."""
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from .base_embedding_model import BaseEmbeddingModel
|
||||
from ..component_registry import R
|
||||
|
||||
|
||||
@R.register("openai")
|
||||
class OpenAIEmbeddingModel(BaseEmbeddingModel):
|
||||
"""Embedding model backed by any OpenAI-compatible API."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._client: AsyncOpenAI | None = None
|
||||
|
||||
async def _start(self) -> None:
|
||||
"""Initialize async OpenAI client."""
|
||||
self._client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url, **self.kwargs)
|
||||
await super()._start()
|
||||
|
||||
async def _close(self) -> None:
|
||||
"""Close the async OpenAI client."""
|
||||
if self._client:
|
||||
await self._client.close()
|
||||
self._client = None
|
||||
await super()._close()
|
||||
|
||||
async def _get_embeddings(self, input_text: list[str], **kwargs) -> list[list[float] | None]:
|
||||
"""Call the embeddings API and return results aligned to input order."""
|
||||
if self._client is None:
|
||||
raise RuntimeError("Client not initialized. Call _start() first.")
|
||||
|
||||
create_kwargs: dict = {"model": self.model_name, "input": input_text, **kwargs}
|
||||
if self.pass_dimensions:
|
||||
create_kwargs["dimensions"] = self.dimensions
|
||||
|
||||
completion = await self._client.embeddings.create(**create_kwargs)
|
||||
|
||||
# Map API results back to input order
|
||||
result: list[list[float] | None] = [None] * len(input_text)
|
||||
for emb in completion.data:
|
||||
if 0 <= emb.index < len(input_text):
|
||||
vec = emb.embedding or getattr(emb, "dense_embedding", None)
|
||||
if vec is not None:
|
||||
result[emb.index] = list(vec)
|
||||
else:
|
||||
self.logger.warning(f"Empty embedding at index {emb.index}")
|
||||
else:
|
||||
self.logger.warning(f"Index {emb.index} out of range for input length {len(input_text)}")
|
||||
|
||||
return result
|
||||
7
reme4/components/file_graph/__init__.py
Normal file
7
reme4/components/file_graph/__init__.py
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
"""File graph module."""
|
||||
|
||||
from .base_file_graph import BaseFileGraph
|
||||
from .local_file_graph import LocalFileGraph
|
||||
from .nx_file_graph import NxFileGraph
|
||||
|
||||
__all__ = ["BaseFileGraph", "LocalFileGraph", "NxFileGraph"]
|
||||
50
reme4/components/file_graph/base_file_graph.py
Normal file
50
reme4/components/file_graph/base_file_graph.py
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
from abc import abstractmethod
|
||||
from pathlib import Path
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ...enumeration import ComponentEnum
|
||||
from ...schema import FileLink, FileNode
|
||||
|
||||
|
||||
class BaseFileGraph(BaseComponent):
|
||||
"""Abstract base for file-graph backends."""
|
||||
|
||||
component_type = ComponentEnum.FILE_GRAPH
|
||||
|
||||
def __init__(self, graph_name: str = "default", **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.graph_name: str = graph_name or self.name
|
||||
self.graph_path: Path = self.working_path / self.component_type.value
|
||||
self.graph_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# -- Node CRUD ---------------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def upsert_nodes(self, nodes: list[FileNode]) -> None:
|
||||
"""Insert or update nodes in the graph."""
|
||||
|
||||
@abstractmethod
|
||||
async def delete_nodes(self, paths: list[str]) -> None:
|
||||
"""Delete nodes by path."""
|
||||
|
||||
@abstractmethod
|
||||
async def get_nodes(self, paths: list[str] | None = None) -> list[FileNode]:
|
||||
"""Return nodes by paths; None = all real nodes; [] = []."""
|
||||
|
||||
@abstractmethod
|
||||
async def rebuild_links(self) -> None:
|
||||
"""Rebuild all edges from each node's link payload."""
|
||||
|
||||
@abstractmethod
|
||||
async def clear(self):
|
||||
"""Remove all nodes and edges."""
|
||||
|
||||
# -- Link access -------------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def get_outlinks(self, path: str) -> list[FileLink]:
|
||||
"""Return outgoing links for *path*."""
|
||||
|
||||
@abstractmethod
|
||||
async def get_inlinks(self, path: str) -> list[FileLink]:
|
||||
"""Return incoming links for *path*."""
|
||||
130
reme4/components/file_graph/local_file_graph.py
Normal file
130
reme4/components/file_graph/local_file_graph.py
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
"""Pure-Python file-graph backend (no external deps)."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from .base_file_graph import BaseFileGraph
|
||||
from ..component_registry import R
|
||||
from ...schema import FileLink, FileNode
|
||||
|
||||
|
||||
@R.register("local")
|
||||
class LocalFileGraph(BaseFileGraph):
|
||||
"""Dict-backed file graph; uses FileLink.path for adjacency."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._nodes: dict[str, FileNode] = {}
|
||||
self._inverse: dict[str, set[str]] = {} # target → {sources}
|
||||
self._pending: dict[str, set[str]] = {} # virtual target → {sources}
|
||||
self._graph_file: Path = self.graph_path / f"{self.graph_name}.jsonl"
|
||||
|
||||
# -- Lifecycle ---------------------------------------------------------
|
||||
|
||||
async def _start(self) -> None:
|
||||
await super()._start()
|
||||
self._load()
|
||||
await self.rebuild_links()
|
||||
self.logger.info(
|
||||
f"LocalFileGraph '{self.graph_name}' ready: "
|
||||
f"{len(self._nodes)} nodes, {sum(len(s) for s in self._inverse.values())} edges, "
|
||||
f"{sum(len(s) for s in self._pending.values())} pending",
|
||||
)
|
||||
|
||||
async def _close(self) -> None:
|
||||
self._dump()
|
||||
await super()._close()
|
||||
|
||||
def _load(self) -> None:
|
||||
"""Load nodes from JSONL file into memory."""
|
||||
if not self._graph_file.exists():
|
||||
return
|
||||
with open(self._graph_file, "r", encoding="utf-8") as f:
|
||||
self._nodes.update((n.path, n) for line in f if line.strip() for n in [FileNode.model_validate_json(line)])
|
||||
|
||||
def _dump(self) -> None:
|
||||
"""Persist all nodes to JSONL via atomic rename."""
|
||||
tmp = self._graph_file.with_suffix(".tmp")
|
||||
with open(tmp, "w", encoding="utf-8") as f:
|
||||
f.writelines(f"{n.model_dump_json()}\n" for n in self._nodes.values())
|
||||
tmp.replace(self._graph_file)
|
||||
|
||||
# -- Edge bookkeeping --------------------------------------------------
|
||||
|
||||
def _add_edge(self, src: str, target: str) -> None:
|
||||
"""Register src→target; route to pending if target is virtual."""
|
||||
bucket = self._inverse if target in self._nodes else self._pending
|
||||
bucket.setdefault(target, set()).add(src)
|
||||
|
||||
def _remove_edge(self, src: str, target: str) -> None:
|
||||
"""Remove src→target from both inverse and pending buckets."""
|
||||
for bucket in (self._inverse, self._pending):
|
||||
srcs = bucket.get(target)
|
||||
if srcs is None or src not in srcs:
|
||||
continue
|
||||
srcs.discard(src)
|
||||
if not srcs:
|
||||
del bucket[target]
|
||||
|
||||
# -- Node CRUD ---------------------------------------------------------
|
||||
|
||||
async def upsert_nodes(self, nodes: list[FileNode]) -> None:
|
||||
for node in nodes:
|
||||
path = node.path
|
||||
old = self._nodes.get(path)
|
||||
if old is not None:
|
||||
for link in old.links:
|
||||
if link.path:
|
||||
self._remove_edge(path, link.path)
|
||||
self._nodes[path] = node
|
||||
for link in node.links:
|
||||
if link.path:
|
||||
self._add_edge(path, link.path)
|
||||
# Promote pending edges that now target a real node.
|
||||
promoted = self._pending.pop(path, None)
|
||||
if promoted:
|
||||
self._inverse.setdefault(path, set()).update(promoted)
|
||||
|
||||
async def delete_nodes(self, paths: list[str]) -> None:
|
||||
for path in paths:
|
||||
node = self._nodes.pop(path, None)
|
||||
if node is None:
|
||||
continue
|
||||
for link in node.links:
|
||||
if link.path:
|
||||
self._remove_edge(path, link.path)
|
||||
# Demote inbound edges to pending (sources still reference this path).
|
||||
demoted = self._inverse.pop(path, None)
|
||||
if demoted:
|
||||
self._pending.setdefault(path, set()).update(demoted)
|
||||
|
||||
async def get_nodes(self, paths: list[str] | None = None) -> list[FileNode]:
|
||||
if paths is None:
|
||||
return list(self._nodes.values())
|
||||
return [self._nodes[p] for p in paths if p in self._nodes]
|
||||
|
||||
async def rebuild_links(self) -> None:
|
||||
"""Rebuild inverse/pending indexes from all node link payloads."""
|
||||
self._inverse.clear()
|
||||
self._pending.clear()
|
||||
for src, node in self._nodes.items():
|
||||
for link in node.links:
|
||||
if link.path:
|
||||
self._add_edge(src, link.path)
|
||||
|
||||
async def clear(self):
|
||||
self._nodes.clear()
|
||||
self._inverse.clear()
|
||||
self._pending.clear()
|
||||
|
||||
# -- Link access -------------------------------------------------------
|
||||
|
||||
async def get_outlinks(self, path: str) -> list[FileLink]:
|
||||
node = self._nodes.get(path)
|
||||
if node is None:
|
||||
return []
|
||||
return [lnk for lnk in node.links if lnk.path and lnk.path in self._nodes]
|
||||
|
||||
async def get_inlinks(self, path: str) -> list[FileLink]:
|
||||
if path not in self._nodes:
|
||||
return []
|
||||
return [link for src in self._inverse.get(path, ()) for link in self._nodes[src].links if link.path == path]
|
||||
126
reme4/components/file_graph/nx_file_graph.py
Normal file
126
reme4/components/file_graph/nx_file_graph.py
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
"""Networkx file-graph backend."""
|
||||
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
import networkx as nx
|
||||
except ImportError:
|
||||
nx = None
|
||||
|
||||
from .base_file_graph import BaseFileGraph
|
||||
from ..component_registry import R
|
||||
from ...schema import FileLink, FileNode
|
||||
|
||||
|
||||
@R.register("nx")
|
||||
class NxFileGraph(BaseFileGraph):
|
||||
"""Networkx-backed file graph; uses FileLink.path for adjacency."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
if nx is None:
|
||||
raise ImportError("NxFileGraph requires networkx — pip install networkx")
|
||||
self._graph: nx.MultiDiGraph = nx.MultiDiGraph()
|
||||
self._graph_file: Path = self.graph_path / f"{self.graph_name}.pkl"
|
||||
|
||||
# -- Lifecycle ---------------------------------------------------------
|
||||
|
||||
async def _start(self) -> None:
|
||||
await super()._start()
|
||||
loaded = self._load()
|
||||
if loaded is not None:
|
||||
self._graph = loaded
|
||||
n_real = sum(1 for _, d in self._graph.nodes(data=True) if "node" in d)
|
||||
self.logger.info(
|
||||
f"NxFileGraph '{self.graph_name}' ready: "
|
||||
f"{n_real} nodes, {self._graph.number_of_edges()} edges, "
|
||||
f"{self._graph.number_of_nodes() - n_real} virtual",
|
||||
)
|
||||
|
||||
async def _close(self) -> None:
|
||||
self._dump()
|
||||
await super()._close()
|
||||
|
||||
def _load(self) -> nx.MultiDiGraph | None:
|
||||
"""Load graph from pickle file; return None on failure."""
|
||||
if not self._graph_file.exists():
|
||||
return None
|
||||
try:
|
||||
with open(self._graph_file, "rb") as f:
|
||||
return pickle.load(f)
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to load {self._graph_file}: {e}")
|
||||
return None
|
||||
|
||||
def _dump(self) -> None:
|
||||
"""Persist graph to pickle via atomic rename."""
|
||||
try:
|
||||
tmp = self._graph_file.with_suffix(".tmp")
|
||||
with open(tmp, "wb") as f:
|
||||
pickle.dump(self._graph, f, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
tmp.replace(self._graph_file)
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to write {self._graph_file}: {e}")
|
||||
|
||||
# -- Node CRUD ---------------------------------------------------------
|
||||
|
||||
async def upsert_nodes(self, nodes: list[FileNode]) -> None:
|
||||
for node in nodes:
|
||||
path = node.path
|
||||
if self._graph.has_node(path):
|
||||
# Drop outgoing edges; inbound stay intact.
|
||||
self._graph.remove_edges_from(list(self._graph.out_edges(path, keys=True)))
|
||||
self._graph.add_node(path, node=node) # promotes virtual node if present
|
||||
# Missing targets become attr-less virtual nodes.
|
||||
self._graph.add_edges_from((path, lnk.path, {"link": lnk}) for lnk in node.links if lnk.path)
|
||||
|
||||
async def delete_nodes(self, paths: list[str]) -> None:
|
||||
for path in paths:
|
||||
if not self._graph.has_node(path):
|
||||
continue
|
||||
self._graph.remove_edges_from(list(self._graph.out_edges(path, keys=True)))
|
||||
# Demote to virtual: keep inbound edges, drop node payload.
|
||||
self._graph.nodes[path].pop("node", None)
|
||||
if self._graph.in_degree(path) == 0:
|
||||
self._graph.remove_node(path) # remove orphan virtual node
|
||||
|
||||
async def get_nodes(self, paths: list[str] | None = None) -> list[FileNode]:
|
||||
nodes_view = self._graph.nodes
|
||||
if paths is None:
|
||||
return [d["node"] for _, d in nodes_view(data=True) if "node" in d]
|
||||
return [nodes_view[path]["node"] for path in paths if path in nodes_view and "node" in nodes_view[path]]
|
||||
|
||||
async def rebuild_links(self) -> None:
|
||||
"""Rebuild all edges from real node payloads; drop virtual nodes."""
|
||||
self._graph.remove_edges_from(list(self._graph.edges(keys=True)))
|
||||
virtual = [n for n, d in self._graph.nodes(data=True) if "node" not in d]
|
||||
self._graph.remove_nodes_from(virtual)
|
||||
self._graph.add_edges_from(
|
||||
(path, lnk.path, {"link": lnk})
|
||||
for path, data in self._graph.nodes(data=True)
|
||||
for lnk in data["node"].links
|
||||
if lnk.path
|
||||
)
|
||||
|
||||
async def clear(self):
|
||||
"""Remove all nodes and edges."""
|
||||
self._graph.clear()
|
||||
|
||||
# -- Link access -------------------------------------------------------
|
||||
|
||||
async def get_outlinks(self, path: str) -> list[FileLink]:
|
||||
nodes_view = self._graph.nodes
|
||||
if path not in nodes_view or "node" not in nodes_view[path]:
|
||||
return []
|
||||
return [
|
||||
d["link"]
|
||||
for _, target, d in self._graph.out_edges(path, data=True)
|
||||
if "link" in d and "node" in nodes_view[target]
|
||||
]
|
||||
|
||||
async def get_inlinks(self, path: str) -> list[FileLink]:
|
||||
nodes_view = self._graph.nodes
|
||||
if path not in nodes_view or "node" not in nodes_view[path]:
|
||||
return []
|
||||
return [d["link"] for _, _, d in self._graph.in_edges(path, data=True) if "link" in d]
|
||||
5
reme4/components/file_parser/__init__.py
Normal file
5
reme4/components/file_parser/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
from .bare_file_parser import BareFileParser
|
||||
from .base_file_parser import BaseFileParser
|
||||
from .default_file_parser import DefaultFileParser
|
||||
|
||||
__all__ = ["BareFileParser", "BaseFileParser", "DefaultFileParser"]
|
||||
20
reme4/components/file_parser/bare_file_parser.py
Normal file
20
reme4/components/file_parser/bare_file_parser.py
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
from pathlib import Path
|
||||
|
||||
from .base_file_parser import BaseFileParser
|
||||
from ..component_registry import R
|
||||
from ...schema import FileChunk, FileNode
|
||||
|
||||
|
||||
@R.register("bare")
|
||||
class BareFileParser(BaseFileParser):
|
||||
"""Stat-only parser for attachment/binary files.
|
||||
|
||||
No content read, no chunking, no link extraction. The resulting FileNode
|
||||
has empty links and chunk_ids; front_matter carries mime and size so
|
||||
retrieval can filter by file type without reopening the file.
|
||||
"""
|
||||
|
||||
async def parse(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
file_path = Path(path)
|
||||
stat = file_path.stat()
|
||||
return FileNode(path=self._get_relative_path(path), st_mtime=stat.st_mtime, links=[], chunk_ids=[]), []
|
||||
28
reme4/components/file_parser/base_file_parser.py
Normal file
28
reme4/components/file_parser/base_file_parser.py
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
from abc import abstractmethod
|
||||
from pathlib import Path
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ...enumeration import ComponentEnum
|
||||
from ...schema import FileChunk, FileNode
|
||||
|
||||
|
||||
class BaseFileParser(BaseComponent):
|
||||
"""Abstract base for file parsers. Subclasses implement `parse`."""
|
||||
|
||||
component_type = ComponentEnum.FILE_PARSER
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.working_dir = self.app_context.app_config.working_dir if self.app_context else ""
|
||||
|
||||
def _get_relative_path(self, path: str | Path) -> str:
|
||||
"""Return path relative to working_dir, or absolute path if outside."""
|
||||
file_path = Path(path).absolute()
|
||||
try:
|
||||
return str(file_path.relative_to(Path(self.working_dir).absolute()))
|
||||
except ValueError:
|
||||
return str(file_path)
|
||||
|
||||
@abstractmethod
|
||||
async def parse(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
"""Parse a file into (node, chunks)."""
|
||||
78
reme4/components/file_parser/default_file_parser.py
Normal file
78
reme4/components/file_parser/default_file_parser.py
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
from bisect import bisect_right
|
||||
from pathlib import Path
|
||||
|
||||
import aiofiles
|
||||
import yaml
|
||||
|
||||
from .base_file_parser import BaseFileParser
|
||||
from ..component_registry import R
|
||||
from ...schema import FileChunk, FileNode, FileFrontMatter
|
||||
|
||||
|
||||
@R.register("default")
|
||||
class DefaultFileParser(BaseFileParser):
|
||||
"""Parser that splits files into byte-based overlapping chunks."""
|
||||
|
||||
def __init__(self, encoding: str = "utf-8", chunk_byte_size: int = 10000, overlap_byte_size: int = 100, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.encoding = encoding
|
||||
self.chunk_byte_size = max(100, chunk_byte_size)
|
||||
self.overlap_byte_size = max(4, overlap_byte_size)
|
||||
|
||||
@staticmethod
|
||||
def _parse_front_matter(text: str) -> tuple[FileFrontMatter, str]:
|
||||
"""Parse YAML front matter delimited by ---, return (front_matter, remaining)."""
|
||||
if not text.startswith("---"):
|
||||
return FileFrontMatter(), text
|
||||
end_idx = text.find("\n---", 3)
|
||||
if end_idx == -1:
|
||||
return FileFrontMatter(), text
|
||||
try:
|
||||
data = yaml.safe_load(text[3:end_idx].strip()) or {}
|
||||
front_matter = FileFrontMatter(**(data if isinstance(data, dict) else {}))
|
||||
except yaml.YAMLError:
|
||||
front_matter = FileFrontMatter()
|
||||
return front_matter, text[end_idx + 4 :].lstrip("\n")
|
||||
|
||||
async def parse(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
file_path = Path(path)
|
||||
stat = file_path.stat()
|
||||
rel_path = self._get_relative_path(path)
|
||||
|
||||
async with aiofiles.open(file_path, encoding=self.encoding) as f:
|
||||
text = await f.read()
|
||||
|
||||
if not text:
|
||||
return FileNode(path=rel_path, st_mtime=stat.st_mtime), []
|
||||
|
||||
front_matter, content = self._parse_front_matter(text)
|
||||
if not content:
|
||||
return FileNode(path=rel_path, st_mtime=stat.st_mtime, front_matter=front_matter), []
|
||||
|
||||
chunks = self._chunk_content(content, rel_path)
|
||||
chunk_ids = [c.id for c in chunks]
|
||||
return FileNode(path=rel_path, st_mtime=stat.st_mtime, front_matter=front_matter, chunk_ids=chunk_ids), chunks
|
||||
|
||||
def _chunk_content(self, content: str, rel_path: str) -> list[FileChunk]:
|
||||
"""Split content into overlapping byte-range chunks with line numbers."""
|
||||
content_bytes = content.encode(self.encoding)
|
||||
newline_positions = [i for i, b in enumerate(content_bytes) if b == ord("\n")]
|
||||
chunks: list[FileChunk] = []
|
||||
step = self.chunk_byte_size - self.overlap_byte_size
|
||||
start = 0
|
||||
|
||||
while start < len(content_bytes):
|
||||
end = min(start + self.chunk_byte_size, len(content_bytes))
|
||||
chunk_text = content_bytes[start:end].decode(self.encoding, errors="ignore")
|
||||
start_line = bisect_right(newline_positions, start - 1) + 1
|
||||
end_line = bisect_right(newline_positions, end - 1) + 1
|
||||
if content_bytes[end - 1] == ord("\n"):
|
||||
end_line -= 1
|
||||
chunks.append(
|
||||
FileChunk(path=rel_path, start_line=start_line, end_line=end_line, text=chunk_text).set_hash_id()
|
||||
)
|
||||
if end >= len(content_bytes):
|
||||
break
|
||||
start += step
|
||||
|
||||
return chunks
|
||||
14
reme4/components/file_store/__init__.py
Normal file
14
reme4/components/file_store/__init__.py
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
"""File store module.
|
||||
|
||||
In-memory + JSONL backend for the (file → chunks) graph. Subclass
|
||||
`BaseFileStore` to add other backends; only `LocalFileStore` is
|
||||
shipped today.
|
||||
"""
|
||||
|
||||
from .base_file_store import BaseFileStore
|
||||
from .local_file_store import LocalFileStore
|
||||
|
||||
__all__ = [
|
||||
"BaseFileStore",
|
||||
"LocalFileStore",
|
||||
]
|
||||
66
reme4/components/file_store/base_file_store.py
Normal file
66
reme4/components/file_store/base_file_store.py
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
from abc import abstractmethod
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ..embedding import BaseEmbeddingModel
|
||||
from ..file_graph import BaseFileGraph
|
||||
from ..keyword_index import BaseKeywordIndex
|
||||
from ...enumeration import ComponentEnum
|
||||
from ...schema import FileChunk, FileNode, FileLink
|
||||
|
||||
|
||||
class BaseFileStore(BaseComponent):
|
||||
component_type = ComponentEnum.FILE_STORE
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store_name: str,
|
||||
embedding_model: str = "default",
|
||||
keyword_index: str = "default",
|
||||
file_graph: str = "default",
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.store_name = store_name or self.name
|
||||
if not embedding_model and not keyword_index:
|
||||
raise ValueError("At least one of embedding_model or keyword_index must be set.")
|
||||
|
||||
self.embedding_model = self.bind(embedding_model, BaseEmbeddingModel)
|
||||
self.keyword_index = self.bind(keyword_index, BaseKeywordIndex)
|
||||
self.file_graph = self.bind(file_graph, BaseFileGraph)
|
||||
self.store_path = self.working_path / self.component_type.value / store_name
|
||||
self.store_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
async def upsert_file(
|
||||
self,
|
||||
file: tuple[FileNode, list[FileChunk]] | list[tuple[FileNode, list[FileChunk]]],
|
||||
) -> None:
|
||||
"""Upsert a file and its chunks into the store."""
|
||||
|
||||
async def delete_by_path(self, path: str | list[str]) -> None:
|
||||
"""Delete files by their paths from the store."""
|
||||
|
||||
async def clear(self):
|
||||
"""Clear the store of all files and chunks."""
|
||||
|
||||
@abstractmethod
|
||||
async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
|
||||
"""Perform vector similarity search."""
|
||||
|
||||
@abstractmethod
|
||||
async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
|
||||
"""Perform full-text keyword search."""
|
||||
|
||||
async def rebuild_links(self) -> None:
|
||||
if not self.file_graph:
|
||||
raise RuntimeError("file_graph is required for delete_by_path")
|
||||
return await self.file_graph.rebuild_links()
|
||||
|
||||
async def get_outlinks(self, path: str) -> list[FileLink]:
|
||||
if not self.file_graph:
|
||||
raise RuntimeError("file_graph is required for delete_by_path")
|
||||
return await self.file_graph.get_outlinks(path)
|
||||
|
||||
async def get_inlinks(self, path: str) -> list[FileLink]:
|
||||
if not self.file_graph:
|
||||
raise RuntimeError("file_graph is required for delete_by_path")
|
||||
return await self.file_graph.get_inlinks(path)
|
||||
157
reme4/components/file_store/local_file_store.py
Normal file
157
reme4/components/file_store/local_file_store.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
"""In-memory file store with JSONL persistence on close."""
|
||||
|
||||
import aiofiles
|
||||
import numpy as np
|
||||
|
||||
from .base_file_store import BaseFileStore
|
||||
from ..component_registry import R
|
||||
from ...schema import FileChunk, FileNode
|
||||
from ...utils import batch_cosine_similarity, get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
@R.register("local")
|
||||
class LocalFileStore(BaseFileStore):
|
||||
"""In-memory file store with deferred JSONL persistence."""
|
||||
|
||||
def __init__(self, encoding: str = "utf-8", **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.encoding = encoding
|
||||
self.file_chunks: dict[str, FileChunk] = {}
|
||||
self.chunks_path = self.store_path / "file_chunks.jsonl"
|
||||
|
||||
# Lifecycle
|
||||
|
||||
async def _start(self) -> None:
|
||||
await super()._start()
|
||||
if self.chunks_path.exists():
|
||||
try:
|
||||
async with aiofiles.open(self.chunks_path, encoding=self.encoding) as f:
|
||||
async for line in f:
|
||||
line = line.strip()
|
||||
if line:
|
||||
chunk = FileChunk.model_validate_json(line)
|
||||
self.file_chunks[chunk.id] = chunk
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to load {self.chunks_path}: {e}")
|
||||
self.logger.info(f"LocalFileStore '{self.store_name}' ready: {len(self.file_chunks)} chunks")
|
||||
|
||||
async def _close(self) -> None:
|
||||
try:
|
||||
tmp = self.chunks_path.with_suffix(".tmp")
|
||||
async with aiofiles.open(tmp, "w", encoding=self.encoding) as f:
|
||||
await f.write("\n".join(c.model_dump_json() for c in self.file_chunks.values()))
|
||||
tmp.replace(self.chunks_path)
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to write {self.chunks_path}: {e}")
|
||||
self.file_chunks.clear()
|
||||
await super()._close()
|
||||
|
||||
# Base class interface
|
||||
|
||||
async def upsert_file(
|
||||
self,
|
||||
file: tuple[FileNode, list[FileChunk]] | list[tuple[FileNode, list[FileChunk]]],
|
||||
) -> None:
|
||||
if not self.file_graph:
|
||||
raise RuntimeError("file_graph is required for upsert_file")
|
||||
if isinstance(file, tuple):
|
||||
file = [file]
|
||||
|
||||
old_map = {n.path: n for n in await self.file_graph.get_nodes([node.path for node, _ in file])}
|
||||
|
||||
new_nodes: list[FileNode] = []
|
||||
needs_embed: list[FileChunk] = []
|
||||
keyword_docs: dict[str, str] = {}
|
||||
for node, chunks in file:
|
||||
old_node: FileNode | None = old_map.get(node.path)
|
||||
cached = {}
|
||||
if old_node and self.embedding_model:
|
||||
for cid in old_node.chunk_ids:
|
||||
old = self.file_chunks.pop(cid, None)
|
||||
if old and old.embedding:
|
||||
cached[cid] = old.embedding
|
||||
|
||||
node.chunk_ids = []
|
||||
for c in chunks:
|
||||
if self.embedding_model and not c.embedding:
|
||||
if c.id in cached:
|
||||
c.embedding = cached[c.id]
|
||||
elif c.text:
|
||||
needs_embed.append(c)
|
||||
node.chunk_ids.append(c.id)
|
||||
self.file_chunks[c.id] = c
|
||||
if c.text:
|
||||
keyword_docs[c.id] = c.text
|
||||
new_nodes.append(node)
|
||||
|
||||
await self.file_graph.upsert_nodes(new_nodes)
|
||||
if needs_embed and self.embedding_model:
|
||||
await self.embedding_model.get_node_embeddings(needs_embed)
|
||||
if self.keyword_index and keyword_docs:
|
||||
await self.keyword_index.add_docs(keyword_docs)
|
||||
|
||||
async def delete_by_path(self, path: str | list[str]) -> None:
|
||||
if not self.file_graph:
|
||||
raise RuntimeError("file_graph is required for delete_by_path")
|
||||
if isinstance(path, str):
|
||||
path = [path]
|
||||
nodes = await self.file_graph.get_nodes(path)
|
||||
if not nodes:
|
||||
return
|
||||
deleted_chunk_ids = [cid for n in nodes for cid in n.chunk_ids]
|
||||
for cid in deleted_chunk_ids:
|
||||
self.file_chunks.pop(cid, None)
|
||||
await self.file_graph.delete_nodes([n.path for n in nodes])
|
||||
if self.keyword_index and deleted_chunk_ids:
|
||||
await self.keyword_index.delete_docs(deleted_chunk_ids)
|
||||
|
||||
async def clear(self) -> None:
|
||||
if not self.file_graph:
|
||||
raise RuntimeError("file_graph is required for clear")
|
||||
self.file_chunks.clear()
|
||||
if self.keyword_index:
|
||||
await self.keyword_index.clear()
|
||||
await self.file_graph.clear()
|
||||
|
||||
# Search
|
||||
|
||||
async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
|
||||
if self.embedding_model is None or not query:
|
||||
return []
|
||||
|
||||
query_embedding = await self.embedding_model.get_embedding(query)
|
||||
if not query_embedding:
|
||||
return []
|
||||
|
||||
candidates = [c for c in self.file_chunks.values() if c.embedding is not None]
|
||||
if not candidates:
|
||||
return []
|
||||
|
||||
candidate_embeddings = np.stack([c.embedding for c in candidates])
|
||||
similarities = batch_cosine_similarity(query_embedding.reshape(1, -1), candidate_embeddings)[0]
|
||||
|
||||
results = [
|
||||
c.model_copy(update={"scores": {"vector": float(s), "score": float(s)}})
|
||||
for c, s in zip(candidates, similarities)
|
||||
]
|
||||
results.sort(key=lambda r: r.score, reverse=True)
|
||||
return results[:limit]
|
||||
|
||||
async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
|
||||
if not self.keyword_index:
|
||||
return []
|
||||
|
||||
query = query.strip()
|
||||
if not query:
|
||||
return []
|
||||
|
||||
doc_id_score_dict = await self.keyword_index.retrieve(query, limit=limit)
|
||||
results = []
|
||||
for doc_id, score in doc_id_score_dict.items():
|
||||
chunk = self.file_chunks.get(doc_id)
|
||||
if chunk:
|
||||
results.append(chunk.model_copy(update={"scores": {"keyword": score, "score": score}}))
|
||||
|
||||
return results
|
||||
9
reme4/components/file_watcher/__init__.py
Normal file
9
reme4/components/file_watcher/__init__.py
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
"""File watcher implementations for monitoring file system changes."""
|
||||
|
||||
from .base_file_watcher import BaseFileWatcher
|
||||
from .lite_file_watcher import LiteFileWatcher
|
||||
|
||||
__all__ = [
|
||||
"BaseFileWatcher",
|
||||
"LiteFileWatcher",
|
||||
]
|
||||
115
reme4/components/file_watcher/base_file_watcher.py
Normal file
115
reme4/components/file_watcher/base_file_watcher.py
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
import asyncio
|
||||
from abc import abstractmethod
|
||||
from pathlib import Path
|
||||
|
||||
from watchfiles import Change
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ..file_parser import BaseFileParser
|
||||
from ..file_store import BaseFileStore
|
||||
from ...enumeration import ComponentEnum
|
||||
from ...utils import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class BaseFileWatcher(BaseComponent):
|
||||
"""Abstract base for file watchers. Subclasses implement watch_loop and event handlers."""
|
||||
|
||||
component_type = ComponentEnum.FILE_WATCHER
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
watch_paths: list[str] | str,
|
||||
suffix_filters: list[str] | None = None,
|
||||
recursive: bool = True,
|
||||
force_polling: bool = True,
|
||||
debounce: int = 2000,
|
||||
poll_delay_ms: int = 2000,
|
||||
file_store: str = "default",
|
||||
file_parser: str = "default",
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
watch_paths = [watch_paths] if isinstance(watch_paths, str) else watch_paths
|
||||
base = self.working_path
|
||||
self.watch_paths: list[Path] = [base / x for x in watch_paths if (base / x).exists()]
|
||||
self.suffix_filters: list[str] = suffix_filters or ["md"]
|
||||
self.recursive: bool = recursive
|
||||
self.force_polling: bool = force_polling
|
||||
self.debounce: int = debounce
|
||||
self.poll_delay_ms: int = poll_delay_ms
|
||||
self.file_store = self.bind(file_store, BaseFileStore)
|
||||
self.file_parser = self.bind(file_parser, BaseFileParser)
|
||||
self._stop_event: asyncio.Event = asyncio.Event()
|
||||
self._background_task: asyncio.Task | None = None
|
||||
self._retry_interval: float = 10
|
||||
|
||||
async def _start(self):
|
||||
self._stop_event = asyncio.Event()
|
||||
self._background_task = asyncio.create_task(self._background_run())
|
||||
logger.info(f"Started watching: {self.watch_paths}")
|
||||
|
||||
async def _background_run(self):
|
||||
"""Sync store then enter watch loop."""
|
||||
await self.update_store()
|
||||
await self.watch_loop()
|
||||
|
||||
async def _close(self):
|
||||
self._stop_event.set()
|
||||
if self._background_task:
|
||||
await self._background_task
|
||||
logger.info("Stopped watching")
|
||||
|
||||
def watch_filter(self, _change: Change, path: str) -> bool:
|
||||
"""Return True if the file suffix matches the filter list."""
|
||||
if not self.suffix_filters:
|
||||
return True
|
||||
return any(path.endswith("." + s.strip(".")) for s in self.suffix_filters)
|
||||
|
||||
async def scan_existing_files(self) -> list[Path]:
|
||||
"""Collect all watchable files under watch_paths."""
|
||||
files: list[Path] = []
|
||||
for path in self.watch_paths:
|
||||
if not path.exists():
|
||||
continue
|
||||
if path.is_file():
|
||||
if self.watch_filter(Change.added, str(path)):
|
||||
files.append(path)
|
||||
else:
|
||||
items = path.rglob("*") if self.recursive else path.iterdir()
|
||||
files.extend(p for p in items if p.is_file() and self.watch_filter(Change.added, str(p)))
|
||||
return files
|
||||
|
||||
async def clear_store(self):
|
||||
"""Remove all entries from the file store."""
|
||||
if self.file_store is None:
|
||||
raise ValueError("file_store is not initialized!")
|
||||
await self.file_store.clear()
|
||||
|
||||
async def reset_store(self):
|
||||
"""Clear the store and re-index all existing files."""
|
||||
if self.file_store is None:
|
||||
raise ValueError("file_store is not initialized!")
|
||||
await self.file_store.clear()
|
||||
await self.on_added(await self.scan_existing_files())
|
||||
|
||||
@abstractmethod
|
||||
async def watch_loop(self):
|
||||
"""Watch for file changes and dispatch events."""
|
||||
|
||||
@abstractmethod
|
||||
async def update_store(self):
|
||||
"""Sync the store with the current state of watch_paths."""
|
||||
|
||||
@abstractmethod
|
||||
async def on_added(self, path: Path | list[Path]):
|
||||
"""Handle file added event."""
|
||||
|
||||
@abstractmethod
|
||||
async def on_modified(self, path: Path | list[Path]):
|
||||
"""Handle file modified event."""
|
||||
|
||||
@abstractmethod
|
||||
async def on_deleted(self, path: Path | list[Path]):
|
||||
"""Handle file deleted event."""
|
||||
128
reme4/components/file_watcher/lite_file_watcher.py
Normal file
128
reme4/components/file_watcher/lite_file_watcher.py
Normal file
|
|
@ -0,0 +1,128 @@
|
|||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
from watchfiles import Change, awatch
|
||||
|
||||
from .base_file_watcher import BaseFileWatcher
|
||||
from ..component_registry import R
|
||||
from ...schema import FileChunk, FileNode
|
||||
from ...utils import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
@R.register("lite")
|
||||
class LiteFileWatcher(BaseFileWatcher):
|
||||
"""Polling-based file watcher using watchfiles awatch."""
|
||||
|
||||
async def _interruptible_sleep(self):
|
||||
"""Sleep until stop or timeout, whichever comes first."""
|
||||
try:
|
||||
await asyncio.wait_for(self._stop_event.wait(), timeout=self._retry_interval)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
async def watch_loop(self):
|
||||
if not self.watch_paths:
|
||||
logger.warning("No watch paths specified")
|
||||
return
|
||||
|
||||
while not self._stop_event.is_set():
|
||||
valid_paths = [p for p in self.watch_paths if p.exists()]
|
||||
if not valid_paths:
|
||||
logger.warning(f"No valid paths, retrying in {self._retry_interval}s...")
|
||||
await self._interruptible_sleep()
|
||||
continue
|
||||
|
||||
invalid = set(self.watch_paths) - set(valid_paths)
|
||||
if invalid:
|
||||
logger.warning(f"Skipping invalid paths: {invalid}")
|
||||
|
||||
try:
|
||||
logger.info(f"Watching: {valid_paths}")
|
||||
async for changes in awatch(
|
||||
*valid_paths,
|
||||
watch_filter=self.watch_filter,
|
||||
recursive=self.recursive,
|
||||
force_polling=self.force_polling,
|
||||
debounce=self.debounce,
|
||||
poll_delay_ms=self.poll_delay_ms,
|
||||
stop_event=self._stop_event,
|
||||
):
|
||||
if self._stop_event.is_set():
|
||||
break
|
||||
await self._dispatch_changes(changes)
|
||||
except Exception:
|
||||
logger.exception(f"Watch error, retrying in {self._retry_interval}s...")
|
||||
if not self._stop_event.is_set():
|
||||
await self._interruptible_sleep()
|
||||
|
||||
async def _dispatch_changes(self, changes: set[tuple[Change, str]]):
|
||||
"""Classify raw changes and dispatch to event handlers."""
|
||||
added = [Path(p) for c, p in changes if c == Change.added]
|
||||
modified = [Path(p) for c, p in changes if c == Change.modified]
|
||||
deleted = [Path(p) for c, p in changes if c == Change.deleted]
|
||||
if added:
|
||||
logger.info(f"Detected {len(added)} added file(s)")
|
||||
await self.on_added(added)
|
||||
if modified:
|
||||
logger.info(f"Detected {len(modified)} modified file(s)")
|
||||
await self.on_modified(modified)
|
||||
if deleted:
|
||||
logger.info(f"Detected {len(deleted)} deleted file(s)")
|
||||
await self.on_deleted(deleted)
|
||||
|
||||
async def update_store(self):
|
||||
if self.file_store is None:
|
||||
raise ValueError("file_store is not initialized!")
|
||||
|
||||
# Diff existing files against indexed entries by path and mtime.
|
||||
existing = {str(p): p.stat().st_mtime for p in await self.scan_existing_files()}
|
||||
indexed = {p: n.st_mtime for p, n in self.file_store.file_nodes.items()}
|
||||
existing_keys, indexed_keys = set(existing), set(indexed)
|
||||
|
||||
to_delete = indexed_keys - existing_keys
|
||||
to_add = existing_keys - indexed_keys
|
||||
to_modify = [p for p in existing_keys & indexed_keys if existing[p] != indexed[p]]
|
||||
|
||||
if to_modify:
|
||||
logger.info(f"Updating {len(to_modify)} modified file(s)")
|
||||
await self.on_modified([Path(p) for p in to_modify])
|
||||
if to_delete:
|
||||
logger.info(f"Removing {len(to_delete)} deleted file(s)")
|
||||
await self.on_deleted([Path(p) for p in to_delete])
|
||||
if to_add:
|
||||
logger.info(f"Indexing {len(to_add)} new file(s)")
|
||||
await self.on_added([Path(p) for p in to_add])
|
||||
if not to_modify and not to_delete and not to_add:
|
||||
logger.info("Store is up to date")
|
||||
|
||||
async def _parse_and_upsert(self, paths: list[Path], action: str):
|
||||
"""Parse files and upsert into store. Shared by on_added / on_modified."""
|
||||
if self.file_parser is None or self.file_store is None:
|
||||
raise RuntimeError("file_parser or file_store is not initialized!")
|
||||
|
||||
parsed: list[tuple[FileNode, list[FileChunk]]] = []
|
||||
for p in paths:
|
||||
if p.is_file():
|
||||
logger.info(f"{action} file: {p}")
|
||||
parsed.append(await self.file_parser.parse(p))
|
||||
if parsed:
|
||||
file_paths = [str(p) for p in paths if p.is_file()]
|
||||
await self.file_store.delete_by_path(file_paths)
|
||||
await self.file_store.upsert_file(parsed)
|
||||
|
||||
async def on_added(self, path: Path | list[Path]):
|
||||
paths = [path] if isinstance(path, Path) else path
|
||||
await self._parse_and_upsert(paths, "Adding")
|
||||
|
||||
async def on_modified(self, path: Path | list[Path]):
|
||||
paths = [path] if isinstance(path, Path) else path
|
||||
await self._parse_and_upsert(paths, "Updating")
|
||||
|
||||
async def on_deleted(self, path: Path | list[Path]):
|
||||
if self.file_store is None:
|
||||
raise RuntimeError("file_store is not initialized!")
|
||||
paths = [path] if isinstance(path, Path) else path
|
||||
logger.info(f"Deleting {len(paths)} file(s)")
|
||||
await self.file_store.delete_by_path([str(p) for p in paths])
|
||||
6
reme4/components/job/__init__.py
Normal file
6
reme4/components/job/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
"""Job components for executing workflows."""
|
||||
|
||||
from .base_job import BaseJob
|
||||
from .stream_job import StreamJob
|
||||
|
||||
__all__ = ["BaseJob", "StreamJob"]
|
||||
53
reme4/components/job/base_job.py
Normal file
53
reme4/components/job/base_job.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
"""Base job component for sequential step execution."""
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ..component_registry import R
|
||||
from ..runtime_context import RuntimeContext
|
||||
from ...enumeration import ComponentEnum
|
||||
from ...schema import ComponentConfig, Response
|
||||
|
||||
|
||||
@R.register("base")
|
||||
class BaseJob(BaseComponent):
|
||||
"""Job that executes steps sequentially and returns a Response."""
|
||||
|
||||
component_type = ComponentEnum.JOB
|
||||
|
||||
def __init__(self, description: str, parameters: dict, steps: list[ComponentConfig | dict], **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.description = description
|
||||
self.parameters = parameters or {}
|
||||
self.step_configs = steps or []
|
||||
|
||||
from ...steps import BaseStep
|
||||
|
||||
self.step_components: list[BaseStep] = []
|
||||
|
||||
async def _start(self) -> None:
|
||||
"""Resolve step configs into instantiated step components."""
|
||||
assert self.app_context is not None, "app_context must be provided"
|
||||
for raw in self.step_configs:
|
||||
config = raw if isinstance(raw, ComponentConfig) else ComponentConfig(**raw)
|
||||
if not config.backend:
|
||||
raise ValueError("Step is missing the required 'backend' field")
|
||||
step_cls = R.get(ComponentEnum.STEP, config.backend)
|
||||
if not step_cls:
|
||||
raise ValueError(f"Unregistered backend '{config.backend}' of type '{ComponentEnum.STEP}'")
|
||||
params = config.model_dump()
|
||||
params["app_context"] = self.app_context
|
||||
self.step_components.append(step_cls(**params))
|
||||
|
||||
async def _close(self) -> None:
|
||||
"""Release all step components."""
|
||||
self.step_components.clear()
|
||||
|
||||
async def __call__(self, **kwargs) -> Response:
|
||||
"""Execute all steps in order and return the final response."""
|
||||
context = RuntimeContext(**kwargs)
|
||||
try:
|
||||
for step in self.step_components:
|
||||
await step(context)
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to execute job: {e}")
|
||||
context.response.answer = str(e)
|
||||
return context.response
|
||||
21
reme4/components/job/stream_job.py
Normal file
21
reme4/components/job/stream_job.py
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
"""Streaming job for real-time output delivery."""
|
||||
|
||||
from .base_job import BaseJob
|
||||
from ..component_registry import R
|
||||
from ..runtime_context import RuntimeContext
|
||||
from ...enumeration import ChunkEnum
|
||||
|
||||
|
||||
@R.register("stream")
|
||||
class StreamJob(BaseJob):
|
||||
"""Job that streams chunks to a queue instead of returning a Response."""
|
||||
|
||||
async def __call__(self, **kwargs) -> None:
|
||||
"""Execute steps and stream output; errors are sent as ERROR chunks."""
|
||||
context = RuntimeContext(**kwargs)
|
||||
try:
|
||||
for step in self.step_components:
|
||||
await step(context)
|
||||
except Exception as e:
|
||||
await context.add_stream_string(str(e), ChunkEnum.ERROR)
|
||||
await context.add_stream_done()
|
||||
4
reme4/components/keyword_index/__init__.py
Normal file
4
reme4/components/keyword_index/__init__.py
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
from .base_keyword_index import BaseKeywordIndex
|
||||
from .bm25_index import BM25Index
|
||||
|
||||
__all__ = ["BaseKeywordIndex", "BM25Index"]
|
||||
80
reme4/components/keyword_index/base_keyword_index.py
Normal file
80
reme4/components/keyword_index/base_keyword_index.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
"""Abstract base class for keyword index implementations."""
|
||||
|
||||
from abc import abstractmethod
|
||||
from pathlib import Path
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ..tokenizer import BaseTokenizer
|
||||
from ...enumeration import ComponentEnum
|
||||
|
||||
|
||||
class BaseKeywordIndex(BaseComponent):
|
||||
"""Abstract base class for keyword index implementations."""
|
||||
|
||||
component_type = ComponentEnum.KEYWORD_INDEX
|
||||
|
||||
def __init__(self, tokenizer: str = "default", **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
from ..tokenizer import RegexTokenizer
|
||||
|
||||
self.tokenizer = self.bind(tokenizer, BaseTokenizer, default_factory=RegexTokenizer)
|
||||
self.index_path = self.working_path / self.component_type.value
|
||||
self.index_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
async def _start(self) -> None:
|
||||
"""Load existing index from disk if available."""
|
||||
if self.index_file.exists():
|
||||
await self.load()
|
||||
self.logger.info(f"Loaded index from {self.index_path}")
|
||||
|
||||
async def _close(self) -> None:
|
||||
"""Save index to disk on shutdown."""
|
||||
await self.dump()
|
||||
self.logger.info(f"Saved index to {self.index_path}")
|
||||
|
||||
@property
|
||||
def index_file(self) -> Path:
|
||||
"""Return the pickle file path derived from tokenizer name."""
|
||||
if self.tokenizer is None:
|
||||
raise RuntimeError("Tokenizer not initialized. Call start() first.")
|
||||
name = type(self.tokenizer).__name__.replace("Tokenizer", "").lower()
|
||||
return self.index_path / f"bm25_{name}.pkl"
|
||||
|
||||
def _tokenize(self, text: str) -> list[str]:
|
||||
"""Tokenize a text string into tokens."""
|
||||
if self.tokenizer is None:
|
||||
raise RuntimeError("Tokenizer not initialized. Call start() first.")
|
||||
return self.tokenizer.tokenize([text])[0]
|
||||
|
||||
@abstractmethod
|
||||
async def add_docs(self, docs_dict: dict[str, str]) -> None:
|
||||
"""Index or update documents. Mapping of doc_id to content."""
|
||||
|
||||
@abstractmethod
|
||||
async def delete_docs(self, doc_ids: list[str]) -> None:
|
||||
"""Remove documents by their IDs."""
|
||||
|
||||
@abstractmethod
|
||||
async def retrieve(self, query: str, limit: int = 3) -> dict[str, float]:
|
||||
"""Search documents. Returns {doc_id: score} sorted descending."""
|
||||
|
||||
@abstractmethod
|
||||
async def dump(self) -> None:
|
||||
"""Persist index to disk."""
|
||||
|
||||
@abstractmethod
|
||||
async def load(self) -> None:
|
||||
"""Load index from disk."""
|
||||
|
||||
@abstractmethod
|
||||
async def clear(self) -> None:
|
||||
"""Reset index to empty state."""
|
||||
|
||||
async def reset_index(self, docs_dict: dict[str, str]) -> None:
|
||||
"""Clear index, re-add all documents, and persist."""
|
||||
await self.clear()
|
||||
await self.add_docs(docs_dict)
|
||||
await self.dump()
|
||||
|
||||
async def optimize_index(self) -> None:
|
||||
"""Optimize index for performance. Override in subclass if needed."""
|
||||
194
reme4/components/keyword_index/bm25_index.py
Normal file
194
reme4/components/keyword_index/bm25_index.py
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
"""BM25 search engine with persistent index support.
|
||||
|
||||
Implements Okapi BM25 ranking with an inverted index for efficient
|
||||
document lookup, incremental updates, and pickle-based persistence.
|
||||
"""
|
||||
|
||||
import math
|
||||
import pickle
|
||||
from collections import Counter
|
||||
from typing import TypedDict
|
||||
|
||||
from .base_keyword_index import BaseKeywordIndex
|
||||
|
||||
|
||||
class DocMeta(TypedDict):
|
||||
"""Per-document metadata: token count and unique token ID set."""
|
||||
|
||||
len: int
|
||||
token_ids: set[int]
|
||||
|
||||
|
||||
class BM25Index(BaseKeywordIndex):
|
||||
"""BM25 search engine with file-based persistence.
|
||||
|
||||
Args:
|
||||
k1: Term frequency saturation parameter (default 1.5).
|
||||
b: Document length normalization parameter (default 0.75).
|
||||
"""
|
||||
|
||||
def __init__(self, k1: float = 1.5, b: float = 0.75, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.k1 = k1
|
||||
self.b = b
|
||||
self.vocab: dict[str, int] = {} # token -> token_id
|
||||
self.inverted_index: dict[int, dict[str, int]] = {} # token_id -> {doc_id: tf}
|
||||
self.doc_meta: dict[str, DocMeta] = {} # doc_id -> metadata
|
||||
self.total_len: int = 0
|
||||
self._idf_cache: dict[int, float] = {}
|
||||
|
||||
# -- Properties -----------------------------------------------------------
|
||||
|
||||
@property
|
||||
def n_docs(self) -> int:
|
||||
"""Number of indexed documents."""
|
||||
return len(self.doc_meta)
|
||||
|
||||
@property
|
||||
def avg_len(self) -> float:
|
||||
"""Average document length in tokens."""
|
||||
return self.total_len / self.n_docs if self.n_docs > 0 else 0.0
|
||||
|
||||
# -- Internal helpers -----------------------------------------------------
|
||||
|
||||
def _tokens_to_ids(self, tokens: list[str]) -> list[int]:
|
||||
"""Map tokens to integer IDs, assigning new IDs on first encounter."""
|
||||
ids = []
|
||||
for token in tokens:
|
||||
token = token.strip()
|
||||
if token:
|
||||
ids.append(self.vocab.setdefault(token, len(self.vocab)))
|
||||
return ids
|
||||
|
||||
def _remove_doc(self, doc_id: str) -> None:
|
||||
"""Remove a single document from all internal structures."""
|
||||
if doc_id not in self.doc_meta:
|
||||
return
|
||||
meta = self.doc_meta[doc_id]
|
||||
self.total_len -= meta["len"]
|
||||
for tid in meta["token_ids"]:
|
||||
if tid in self.inverted_index:
|
||||
self.inverted_index[tid].pop(doc_id, None)
|
||||
if not self.inverted_index[tid]:
|
||||
del self.inverted_index[tid]
|
||||
del self.doc_meta[doc_id]
|
||||
|
||||
def _get_idf(self, token_id: int) -> float:
|
||||
"""Compute and cache IDF for a token ID."""
|
||||
if token_id in self._idf_cache:
|
||||
return self._idf_cache[token_id]
|
||||
df = len(self.inverted_index.get(token_id, {}))
|
||||
self._idf_cache[token_id] = math.log(1 + (self.n_docs - df + 0.5) / (df + 0.5)) if df else 0.0
|
||||
return self._idf_cache[token_id]
|
||||
|
||||
# -- Public API -----------------------------------------------------------
|
||||
|
||||
async def add_docs(self, docs_dict: dict[str, str]) -> None:
|
||||
"""Index or update multiple documents. Mapping of doc_id to content."""
|
||||
for doc_id, content in docs_dict.items():
|
||||
if doc_id in self.doc_meta:
|
||||
self._remove_doc(doc_id)
|
||||
tokens = self._tokenize(content)
|
||||
if not tokens:
|
||||
continue
|
||||
token_ids = self._tokens_to_ids(tokens)
|
||||
token_counts = Counter(token_ids)
|
||||
for tid, tf in token_counts.items():
|
||||
self.inverted_index.setdefault(tid, {})[doc_id] = tf
|
||||
self.doc_meta[doc_id] = {"len": len(token_ids), "token_ids": set(token_counts)}
|
||||
self.total_len += len(token_ids)
|
||||
self._idf_cache = {}
|
||||
|
||||
async def delete_docs(self, doc_ids: list[str]) -> None:
|
||||
"""Remove documents by their IDs."""
|
||||
for doc_id in doc_ids:
|
||||
self._remove_doc(doc_id)
|
||||
self._idf_cache = {}
|
||||
|
||||
async def retrieve(self, query: str, limit: int = 3) -> dict[str, float]:
|
||||
"""Search documents. Returns {doc_id: score} sorted descending."""
|
||||
query_ids = [self.vocab[t] for t in self._tokenize(query) if t in self.vocab]
|
||||
if not query_ids or self.n_docs == 0:
|
||||
return {}
|
||||
|
||||
scores: dict[str, float] = {}
|
||||
avg_len = self.avg_len
|
||||
for tid in query_ids:
|
||||
if tid not in self.inverted_index:
|
||||
continue
|
||||
idf = self._get_idf(tid)
|
||||
for doc_id, tf in self.inverted_index[tid].items():
|
||||
doc_len = self.doc_meta[doc_id]["len"]
|
||||
tf_score = tf * (self.k1 + 1) / (tf + self.k1 * (1 - self.b + self.b * doc_len / avg_len))
|
||||
scores[doc_id] = scores.get(doc_id, 0.0) + idf * tf_score
|
||||
|
||||
return dict(sorted(scores.items(), key=lambda x: x[1], reverse=True)[:limit]) if scores else {}
|
||||
|
||||
async def dump(self) -> None:
|
||||
"""Persist index to disk via pickle."""
|
||||
with open(self.index_file, "wb") as f:
|
||||
pickle.dump(
|
||||
{
|
||||
"vocab": self.vocab,
|
||||
"inverted_index": self.inverted_index,
|
||||
"doc_meta": self.doc_meta,
|
||||
"total_len": self.total_len,
|
||||
"k1": self.k1,
|
||||
"b": self.b,
|
||||
},
|
||||
f,
|
||||
)
|
||||
|
||||
async def load(self) -> None:
|
||||
"""Load index from disk. Clears index on failure."""
|
||||
try:
|
||||
with open(self.index_file, "rb") as f:
|
||||
data = pickle.load(f)
|
||||
self.vocab = data["vocab"]
|
||||
self.inverted_index = data["inverted_index"]
|
||||
self.doc_meta = data["doc_meta"]
|
||||
self.total_len = data.get("total_len", 0)
|
||||
self.k1 = data.get("k1", 1.5)
|
||||
self.b = data.get("b", 0.75)
|
||||
self._idf_cache = {}
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to load index: {e}")
|
||||
self.index_file.unlink(missing_ok=True)
|
||||
await self.clear()
|
||||
|
||||
async def clear(self) -> None:
|
||||
"""Reset index to empty state."""
|
||||
self.vocab = {}
|
||||
self.inverted_index = {}
|
||||
self.doc_meta = {}
|
||||
self.total_len = 0
|
||||
self._idf_cache = {}
|
||||
|
||||
async def optimize_index(self) -> None:
|
||||
"""Rebuild vocab to remove unused tokens and compact token IDs."""
|
||||
used_token_ids: set[int] = set()
|
||||
for tid in self.inverted_index:
|
||||
used_token_ids.add(tid)
|
||||
if not used_token_ids:
|
||||
await self.clear()
|
||||
return
|
||||
|
||||
# Build compact ID mapping
|
||||
old_to_new: dict[int, int] = {}
|
||||
new_vocab: dict[str, int] = {}
|
||||
for token, old_tid in self.vocab.items():
|
||||
if old_tid in used_token_ids:
|
||||
new_tid = len(new_vocab)
|
||||
new_vocab[token] = new_tid
|
||||
old_to_new[old_tid] = new_tid
|
||||
|
||||
# Rebuild inverted index and doc_meta with new IDs
|
||||
new_inverted_index: dict[int, dict[str, int]] = {}
|
||||
for old_tid, postings in self.inverted_index.items():
|
||||
new_inverted_index[old_to_new[old_tid]] = postings
|
||||
for meta in self.doc_meta.values():
|
||||
meta["token_ids"] = {old_to_new[t] for t in meta["token_ids"] if t in old_to_new}
|
||||
|
||||
self.vocab = new_vocab
|
||||
self.inverted_index = new_inverted_index
|
||||
self._idf_cache = {}
|
||||
125
reme4/components/prompt_handler.py
Normal file
125
reme4/components/prompt_handler.py
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
"""Prompt template loader and formatter with conditional-line and i18n support."""
|
||||
|
||||
import inspect
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
from string import Formatter
|
||||
|
||||
import yaml
|
||||
|
||||
# Matches a leading flag tag like "[verbose] some text".
|
||||
_FLAG_PATTERN = re.compile(r"^\[(\w+)]")
|
||||
|
||||
|
||||
class PromptHandler:
|
||||
"""Loads prompts from YAML/JSON or class-adjacent files and formats them.
|
||||
|
||||
Templates may carry a language suffix (``key_en``, ``key_zh``); ``get_prompt``
|
||||
falls back to the bare key when no localized variant exists. ``prompt_format``
|
||||
additionally supports per-line flags such as ``[verbose] extra text`` that
|
||||
are kept only when the matching flag kwarg is truthy.
|
||||
"""
|
||||
|
||||
_SUPPORTED_EXTENSIONS = {".yaml", ".yml", ".json"}
|
||||
|
||||
def __init__(self, language: str = "", **kwargs):
|
||||
# Only string entries are treated as prompts; other kwargs are ignored.
|
||||
self.data: dict[str, str] = {k: v for k, v in kwargs.items() if isinstance(v, str)}
|
||||
self.language: str = language.strip()
|
||||
|
||||
def load_prompt_by_file(
|
||||
self,
|
||||
prompt_file_path: str | Path | None = None,
|
||||
overwrite: bool = True,
|
||||
) -> "PromptHandler":
|
||||
"""Load prompts from a YAML or JSON file; silently skip on any error."""
|
||||
if prompt_file_path is None:
|
||||
return self
|
||||
|
||||
path = Path(prompt_file_path)
|
||||
if not path.exists() or path.suffix.lower() not in self._SUPPORTED_EXTENSIONS:
|
||||
return self
|
||||
|
||||
try:
|
||||
with path.open(encoding="utf-8") as f:
|
||||
prompt_dict = yaml.safe_load(f) if path.suffix.lower() in (".yaml", ".yml") else json.load(f)
|
||||
except (json.JSONDecodeError, yaml.YAMLError, OSError):
|
||||
return self
|
||||
|
||||
return self.load_prompt_dict(prompt_dict, overwrite)
|
||||
|
||||
def load_prompt_by_class(self, cls: type, overwrite: bool = True) -> "PromptHandler":
|
||||
"""Load prompts from ``<class_module>.yaml`` (or ``.yml``) next to `cls`."""
|
||||
try:
|
||||
base_path = Path(inspect.getfile(cls)).with_suffix("")
|
||||
except (TypeError, OSError):
|
||||
return self
|
||||
|
||||
for ext in (".yaml", ".yml"):
|
||||
if (prompt_path := base_path.with_suffix(ext)).exists():
|
||||
return self.load_prompt_by_file(prompt_path, overwrite)
|
||||
|
||||
return self
|
||||
|
||||
def load_prompt_dict(self, prompt_dict: dict | None = None, overwrite: bool = True) -> "PromptHandler":
|
||||
"""Merge string entries from `prompt_dict` into the in-memory store."""
|
||||
if not isinstance(prompt_dict, dict):
|
||||
return self
|
||||
|
||||
for key, value in prompt_dict.items():
|
||||
if isinstance(value, str) and (overwrite or key not in self.data):
|
||||
self.data[key] = value
|
||||
|
||||
return self
|
||||
|
||||
def get_prompt(self, prompt_name: str) -> str:
|
||||
"""Return the template, preferring the language-suffixed variant when set."""
|
||||
for key in (f"{prompt_name}_{self.language}", prompt_name) if self.language else (prompt_name,):
|
||||
if key in self.data:
|
||||
return self.data[key].strip()
|
||||
|
||||
raise KeyError(f"Prompt '{prompt_name}' not found. Available: {list(self.data.keys())[:10]}")
|
||||
|
||||
def has_prompt(self, prompt_name: str) -> bool:
|
||||
"""True if either the localized or bare prompt is registered."""
|
||||
keys = (f"{prompt_name}_{self.language}", prompt_name) if self.language else (prompt_name,)
|
||||
return any(k in self.data for k in keys)
|
||||
|
||||
def list_prompts(self, language_filter: str | None = None) -> list[str]:
|
||||
"""List all keys, optionally filtered to those ending with ``_<language>``."""
|
||||
if language_filter is None:
|
||||
return list(self.data.keys())
|
||||
suffix = f"_{language_filter.strip()}"
|
||||
return [k for k in self.data if k.endswith(suffix)]
|
||||
|
||||
def prompt_format(self, prompt_name: str, validate: bool = True, **kwargs) -> str:
|
||||
"""Render a prompt: strip inactive flag-lines, then ``str.format`` it.
|
||||
|
||||
Boolean kwargs are treated as flags controlling ``[flag]`` line filtering.
|
||||
Remaining kwargs become positional substitutions for ``{var}`` placeholders.
|
||||
With `validate=True`, missing substitutions raise ``ValueError``.
|
||||
"""
|
||||
prompt = self.get_prompt(prompt_name)
|
||||
flags = {k: v for k, v in kwargs.items() if isinstance(v, bool)}
|
||||
formats = {k: v for k, v in kwargs.items() if not isinstance(v, bool)}
|
||||
|
||||
# Keep lines without flags; otherwise keep when at least one flag is enabled.
|
||||
if flags:
|
||||
lines = []
|
||||
for line in prompt.split("\n"):
|
||||
active_flags = _FLAG_PATTERN.findall(line)
|
||||
cleaned = _FLAG_PATTERN.sub("", line).lstrip()
|
||||
if not active_flags or any(flags.get(f, False) for f in active_flags):
|
||||
lines.append(cleaned)
|
||||
prompt = "\n".join(lines)
|
||||
|
||||
if validate:
|
||||
required = {f for _, f, _, _ in Formatter().parse(prompt) if f is not None}
|
||||
if missing := required - set(formats.keys()):
|
||||
raise ValueError(f"Missing format variables for '{prompt_name}': {sorted(missing)}")
|
||||
|
||||
return prompt.format(**formats).strip() if formats else prompt
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"PromptHandler(language='{self.language}', num_prompts={len(self.data)})"
|
||||
79
reme4/components/runtime_context.py
Normal file
79
reme4/components/runtime_context.py
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
"""Per-request runtime context shared across steps and jobs."""
|
||||
|
||||
import asyncio
|
||||
|
||||
from ..enumeration import ChunkEnum
|
||||
from ..schema import Response, StreamChunk
|
||||
|
||||
|
||||
class RuntimeContext:
|
||||
"""Scratch space for a single execution.
|
||||
|
||||
Holds the response object, an optional stream queue, and a free-form
|
||||
data dict accessed via mapping-style operators.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
response: Response | None = None,
|
||||
stream_queue: asyncio.Queue | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.response: Response = response or Response()
|
||||
self.stream_queue: asyncio.Queue | None = stream_queue
|
||||
self.data: dict = kwargs
|
||||
|
||||
def get(self, key: str, default=None):
|
||||
return self.data.get(key, default)
|
||||
|
||||
def update(self, data: dict) -> "RuntimeContext":
|
||||
self.data.update(data)
|
||||
return self
|
||||
|
||||
def __getitem__(self, key: str):
|
||||
return self.data[key]
|
||||
|
||||
def __setitem__(self, key: str, value):
|
||||
self.data[key] = value
|
||||
|
||||
def __delitem__(self, key: str):
|
||||
del self.data[key]
|
||||
|
||||
def __contains__(self, key: str) -> bool:
|
||||
return key in self.data
|
||||
|
||||
@property
|
||||
def stream(self) -> bool:
|
||||
return self.stream_queue is not None
|
||||
|
||||
@classmethod
|
||||
def from_context(cls, context: "RuntimeContext | None" = None, **kwargs) -> "RuntimeContext":
|
||||
# Reuse the existing context (merging kwargs) or create a new one.
|
||||
if context is None:
|
||||
return cls(**kwargs)
|
||||
context.update(kwargs)
|
||||
return context
|
||||
|
||||
async def _enqueue(self, chunk: StreamChunk) -> None:
|
||||
if self.stream_queue is None:
|
||||
raise RuntimeError("Stream queue not initialized")
|
||||
await self.stream_queue.put(chunk)
|
||||
|
||||
async def add_stream_string(self, chunk: str, chunk_type: ChunkEnum) -> "RuntimeContext":
|
||||
# Emit a text chunk to the stream queue.
|
||||
await self._enqueue(StreamChunk(chunk_type=chunk_type, chunk=chunk))
|
||||
return self
|
||||
|
||||
async def add_stream_done(self) -> "RuntimeContext":
|
||||
# Emit the terminal DONE marker to close the stream.
|
||||
await self._enqueue(StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True))
|
||||
return self
|
||||
|
||||
def apply_mapping(self, mapping: dict[str, str]) -> "RuntimeContext":
|
||||
# Copy data[source] into data[target] for each {source: target} pair.
|
||||
if not mapping:
|
||||
return self
|
||||
for source, target in mapping.items():
|
||||
if source in self.data:
|
||||
self.data[target] = self.data[source]
|
||||
return self
|
||||
11
reme4/components/service/__init__.py
Normal file
11
reme4/components/service/__init__.py
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
"""Service components for exposing jobs via different protocols."""
|
||||
|
||||
from .base_service import BaseService
|
||||
from .http_service import HttpService
|
||||
from .mcp_service import MCPService
|
||||
|
||||
__all__ = [
|
||||
"BaseService",
|
||||
"HttpService",
|
||||
"MCPService",
|
||||
]
|
||||
41
reme4/components/service/base_service.py
Normal file
41
reme4/components/service/base_service.py
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
from abc import abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ..job.base_job import BaseJob
|
||||
from ...enumeration import ComponentEnum
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...application import Application
|
||||
|
||||
|
||||
class BaseService(BaseComponent):
|
||||
"""Base class for services that expose jobs via HTTP, MCP, etc."""
|
||||
|
||||
component_type = ComponentEnum.SERVICE
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.service = None
|
||||
|
||||
@abstractmethod
|
||||
def build_service(self, app: "Application") -> None: ...
|
||||
|
||||
@abstractmethod
|
||||
def add_job(self, job: BaseJob) -> None: ...
|
||||
|
||||
@abstractmethod
|
||||
def start_service(self, app: "Application") -> None: ...
|
||||
|
||||
def add_jobs(self, app: "Application") -> None:
|
||||
for name, job in app.context.jobs.items():
|
||||
try:
|
||||
self.add_job(job)
|
||||
self.logger.info(f"Added job: {name}")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to add job {name}: {e}")
|
||||
|
||||
def run_app(self, app: "Application") -> None:
|
||||
self.build_service(app)
|
||||
self.add_jobs(app)
|
||||
self.start_service(app)
|
||||
85
reme4/components/service/http_service.py
Normal file
85
reme4/components/service/http_service.py
Normal file
|
|
@ -0,0 +1,85 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from .base_service import BaseService
|
||||
from ..component_registry import R
|
||||
from ..job import BaseJob, StreamJob
|
||||
from ...constants import REME_DEFAULT_HOST, REME_DEFAULT_PORT, REME_SERVICE_INFO
|
||||
from ...schema import Request, Response
|
||||
from ...utils import execute_stream_task
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...application import Application
|
||||
|
||||
|
||||
@R.register("http")
|
||||
class HttpService(BaseService):
|
||||
"""HTTP service: normal jobs -> JSON endpoints, stream jobs -> SSE endpoints."""
|
||||
|
||||
def __init__(self, host: str = REME_DEFAULT_HOST, port: int = REME_DEFAULT_PORT, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.host: str = host
|
||||
self.port: int = port
|
||||
|
||||
def _add_job(self, job: BaseJob) -> None:
|
||||
async def execute_endpoint(request: Request) -> Response:
|
||||
return await job(**request.model_dump(exclude_none=True))
|
||||
|
||||
self.service.post(path=f"/{job.name}", response_model=Response, description=job.description)(execute_endpoint)
|
||||
|
||||
def _add_stream_job(self, job: StreamJob) -> None:
|
||||
async def execute_stream_endpoint(request: Request) -> StreamingResponse:
|
||||
stream_queue = asyncio.Queue()
|
||||
task = asyncio.create_task(job(stream_queue=stream_queue, **request.model_dump(exclude_none=True)))
|
||||
|
||||
async def generate_stream() -> AsyncGenerator[bytes, None]:
|
||||
async for chunk in execute_stream_task(
|
||||
stream_queue=stream_queue,
|
||||
task=task,
|
||||
task_name=job.name,
|
||||
output_format="bytes",
|
||||
):
|
||||
assert isinstance(chunk, bytes)
|
||||
yield chunk
|
||||
|
||||
return StreamingResponse(generate_stream(), media_type="text/event-stream")
|
||||
|
||||
self.service.post(f"/{job.name}")(execute_stream_endpoint)
|
||||
|
||||
def add_job(self, job: BaseJob) -> None:
|
||||
if isinstance(job, StreamJob):
|
||||
self._add_stream_job(job)
|
||||
else:
|
||||
self._add_job(job)
|
||||
|
||||
def build_service(self, app: "Application") -> None:
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: FastAPI):
|
||||
await app.start()
|
||||
service_info = json.dumps({"host": self.host, "port": self.port})
|
||||
os.environ[REME_SERVICE_INFO] = service_info
|
||||
self.logger.info(f"ReMe Service started: {REME_SERVICE_INFO}={service_info}")
|
||||
yield
|
||||
await app.close()
|
||||
|
||||
self.service = FastAPI(title=app.config.app_name, lifespan=lifespan)
|
||||
self.service.add_middleware(
|
||||
CORSMiddleware, # type: ignore[arg-type]
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
self.service.post("/health")(lambda: {"status": "healthy"})
|
||||
|
||||
def start_service(self, app: "Application") -> None:
|
||||
uvicorn.run(self.service, host=self.host, port=self.port, **self.kwargs)
|
||||
69
reme4/components/service/mcp_service.py
Normal file
69
reme4/components/service/mcp_service.py
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
import json
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.server.server import Transport
|
||||
from fastmcp.tools import FunctionTool
|
||||
|
||||
from .base_service import BaseService
|
||||
from ..component_registry import R
|
||||
from ..job import StreamJob, BaseJob
|
||||
from ...constants import REME_DEFAULT_HOST, REME_DEFAULT_PORT, REME_SERVICE_INFO
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...application import Application
|
||||
|
||||
|
||||
@R.register("mcp")
|
||||
class MCPService(BaseService):
|
||||
"""Expose jobs as MCP (Model Context Protocol) tools."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
transport: Transport = "sse",
|
||||
host: str = REME_DEFAULT_HOST,
|
||||
port: int = REME_DEFAULT_PORT,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.transport: Transport = transport
|
||||
self.host: str = host
|
||||
self.port: int = port
|
||||
|
||||
def build_service(self, app: "Application") -> None:
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: FastMCP):
|
||||
await app.start()
|
||||
service_info = json.dumps({"host": self.host, "port": self.port})
|
||||
os.environ[REME_SERVICE_INFO] = service_info
|
||||
self.logger.info(f"ReMe MCP Service started: {REME_SERVICE_INFO}={service_info}")
|
||||
yield
|
||||
await app.close()
|
||||
|
||||
self.service = FastMCP(name=app.config.app_name, lifespan=lifespan)
|
||||
|
||||
def add_job(self, job: "BaseJob") -> None:
|
||||
if isinstance(job, StreamJob):
|
||||
return
|
||||
|
||||
async def execute_tool(**kwargs):
|
||||
response = await job(**kwargs)
|
||||
return response.answer
|
||||
|
||||
self.service.add_tool(
|
||||
FunctionTool(
|
||||
name=job.name,
|
||||
description=job.description,
|
||||
fn=execute_tool,
|
||||
parameters=job.parameters or None,
|
||||
),
|
||||
)
|
||||
|
||||
def start_service(self, app: "Application") -> None:
|
||||
transport_kwargs = {}
|
||||
if self.transport != "stdio":
|
||||
transport_kwargs["host"] = self.host
|
||||
transport_kwargs["port"] = self.port
|
||||
self.service.run(transport=self.transport, show_banner=False, **transport_kwargs)
|
||||
11
reme4/components/tokenizer/__init__.py
Normal file
11
reme4/components/tokenizer/__init__.py
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
"""Tokenizer component module."""
|
||||
|
||||
from .base_tokenizer import BaseTokenizer
|
||||
from .jieba_tokenizer import JiebaTokenizer
|
||||
from .regex_tokenizer import RegexTokenizer
|
||||
|
||||
__all__ = [
|
||||
"BaseTokenizer",
|
||||
"JiebaTokenizer",
|
||||
"RegexTokenizer",
|
||||
]
|
||||
43
reme4/components/tokenizer/base_tokenizer.py
Normal file
43
reme4/components/tokenizer/base_tokenizer.py
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
"""Abstract base class for tokenizers."""
|
||||
|
||||
import aiofiles
|
||||
from abc import abstractmethod
|
||||
from pathlib import Path
|
||||
|
||||
from ..base_component import BaseComponent
|
||||
from ...enumeration import ComponentEnum
|
||||
|
||||
|
||||
class BaseTokenizer(BaseComponent):
|
||||
"""Base tokenizer. Subclasses must implement `tokenize`. Loads stopwords on start."""
|
||||
|
||||
component_type = ComponentEnum.TOKENIZER
|
||||
DEFAULT_STOPWORDS_PATH = Path(__file__).parent / "stopwords"
|
||||
|
||||
def __init__(self, stopwords_path: str | Path | None = None, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.stopwords_path = Path(stopwords_path) if stopwords_path else self.DEFAULT_STOPWORDS_PATH
|
||||
self._stopwords: set[str] = set()
|
||||
|
||||
async def _start(self) -> None:
|
||||
"""Load stopwords from file."""
|
||||
if not self.stopwords_path.exists():
|
||||
self.logger.warning(f"Stopwords file not found: {self.stopwords_path}")
|
||||
return
|
||||
async with aiofiles.open(self.stopwords_path, encoding="utf-8") as f:
|
||||
content = await f.read()
|
||||
self._stopwords = {line.strip().lower() for line in content.splitlines() if line.strip()}
|
||||
self.logger.info(f"Loaded {len(self._stopwords)} stopwords from {self.stopwords_path}")
|
||||
|
||||
async def _close(self) -> None:
|
||||
"""Clear stopwords."""
|
||||
self._stopwords.clear()
|
||||
|
||||
@property
|
||||
def stopwords(self) -> set[str]:
|
||||
"""Get the loaded stopwords."""
|
||||
return self._stopwords
|
||||
|
||||
@abstractmethod
|
||||
def tokenize(self, texts: list[str], **kwargs) -> list[list[str]]:
|
||||
"""Tokenize a list of texts."""
|
||||
25
reme4/components/tokenizer/jieba_tokenizer.py
Normal file
25
reme4/components/tokenizer/jieba_tokenizer.py
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
"""Jieba tokenizer for Chinese text segmentation."""
|
||||
|
||||
from .base_tokenizer import BaseTokenizer
|
||||
from ..component_registry import R
|
||||
|
||||
|
||||
@R.register("jieba")
|
||||
class JiebaTokenizer(BaseTokenizer):
|
||||
"""Tokenizer using jieba for Chinese text segmentation."""
|
||||
|
||||
def __init__(self, filter_stopwords: bool = True, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.filter_stopwords = filter_stopwords
|
||||
|
||||
def tokenize(self, texts: list[str], lower: bool = True, **kwargs) -> list[list[str]]:
|
||||
"""Tokenize texts using jieba."""
|
||||
import jieba
|
||||
|
||||
result = []
|
||||
for text in texts:
|
||||
tokens = [x.lower() for x in jieba.cut(text)]
|
||||
if self.filter_stopwords and self._stopwords:
|
||||
tokens = [t for t in tokens if t not in self._stopwords]
|
||||
result.append(tokens)
|
||||
return result
|
||||
31
reme4/components/tokenizer/regex_tokenizer.py
Normal file
31
reme4/components/tokenizer/regex_tokenizer.py
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
"""Regex tokenizer with Chinese character splitting."""
|
||||
|
||||
import re
|
||||
from .base_tokenizer import BaseTokenizer
|
||||
from ..component_registry import R
|
||||
|
||||
|
||||
@R.register("regex")
|
||||
class RegexTokenizer(BaseTokenizer):
|
||||
"""Tokenizer using regex: splits Chinese chars individually, extracts non-Chinese words."""
|
||||
|
||||
WORD_PATTERN = re.compile(r"(?u)\b\w\w+\b") # 2+ char words
|
||||
CHINESE_PATTERN = re.compile(r"[一-鿿]") # single Chinese char
|
||||
|
||||
def __init__(self, filter_stopwords: bool = True, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.filter_stopwords = filter_stopwords
|
||||
|
||||
def tokenize(self, texts: list[str], lower: bool = True, **kwargs) -> list[list[str]]:
|
||||
"""Tokenize texts. Extracts Chinese chars, then non-Chinese words from remaining text."""
|
||||
result = []
|
||||
for text in texts:
|
||||
# Extract Chinese chars individually, then non-Chinese words
|
||||
tokens = self.CHINESE_PATTERN.findall(text)
|
||||
tokens.extend(self.WORD_PATTERN.findall(self.CHINESE_PATTERN.sub(" ", text)))
|
||||
if lower:
|
||||
tokens = [t.lower() for t in tokens]
|
||||
if self.filter_stopwords and self._stopwords:
|
||||
tokens = [t for t in tokens if t not in self._stopwords]
|
||||
result.append(tokens)
|
||||
return result
|
||||
1395
reme4/components/tokenizer/stopwords
Normal file
1395
reme4/components/tokenizer/stopwords
Normal file
File diff suppressed because it is too large
Load diff
5
reme4/config/__init__.py
Normal file
5
reme4/config/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
"""Config"""
|
||||
|
||||
from .config_parser import parse_args
|
||||
|
||||
__all__ = ["parse_args"]
|
||||
201
reme4/config/config_parser.py
Normal file
201
reme4/config/config_parser.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
"""Parser for YAML config with CLI argument overrides."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
# Config files are looked up relative to this module's directory
|
||||
_CONFIG_DIR = Path(__file__).parent
|
||||
# Extensions in priority order: yaml > yml > json when stems collide
|
||||
_SUPPORTED_EXTS = (".yaml", ".yml", ".json")
|
||||
_ENV_VAR_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-([^}]*))?}")
|
||||
# Strings like "007" / "00501" must stay as strings, not be coerced to numbers
|
||||
_LEADING_ZERO_RE = re.compile(r"^-?0\d")
|
||||
|
||||
|
||||
def _repl(m: re.Match) -> str:
|
||||
name: str = m.group(1)
|
||||
# group(2) is None when the placeholder has no `:-default` part
|
||||
default: str | None = m.group(2)
|
||||
v = os.environ.get(name)
|
||||
if v is None:
|
||||
if default is not None:
|
||||
return default
|
||||
raise ValueError(f"Config references undefined env var: {name}")
|
||||
return v
|
||||
|
||||
|
||||
def _expand_env_vars(value: Any) -> Any:
|
||||
"""Recursively expand `${VAR}` / `${VAR:-default}` placeholders in strings."""
|
||||
if isinstance(value, str):
|
||||
return _ENV_VAR_RE.sub(_repl, value)
|
||||
if isinstance(value, dict):
|
||||
return {k: _expand_env_vars(v) for k, v in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [_expand_env_vars(v) for v in value]
|
||||
return value
|
||||
|
||||
|
||||
def _discover_configs() -> dict[str, Path]:
|
||||
"""Pre-scan config directory: maps file stem (name without ext) -> Path."""
|
||||
discovered: dict[str, Path] = {}
|
||||
if _CONFIG_DIR.is_dir():
|
||||
# Sort by ext priority so registration order is deterministic across filesystems
|
||||
files = sorted(
|
||||
(p for p in _CONFIG_DIR.iterdir() if p.is_file() and p.suffix in _SUPPORTED_EXTS),
|
||||
key=lambda p: (_SUPPORTED_EXTS.index(p.suffix), p.name),
|
||||
)
|
||||
for p in files:
|
||||
discovered.setdefault(p.stem, p)
|
||||
return discovered
|
||||
|
||||
|
||||
_CONFIG_REGISTRY = _discover_configs()
|
||||
|
||||
|
||||
def parse_dot_notation(dot_list: list[str]) -> dict:
|
||||
"""Parse "key.subkey=value" strings into nested dict."""
|
||||
result: dict = {}
|
||||
for item in dot_list:
|
||||
if "=" not in item:
|
||||
raise ValueError(f"Invalid dot notation format (missing '='): {item}")
|
||||
key_path, value_str = item.split("=", 1)
|
||||
keys = key_path.split(".")
|
||||
current = result
|
||||
for key in keys[:-1]:
|
||||
if key in current and not isinstance(current[key], dict):
|
||||
raise ValueError(f"Cannot set nested key '{key_path}': '{key}' is already a value")
|
||||
current = current.setdefault(key, {})
|
||||
# Symmetric to the prefix check above: refuse scalar-over-dict overwrite
|
||||
last_key = keys[-1]
|
||||
if last_key in current and isinstance(current[last_key], dict):
|
||||
raise ValueError(f"Cannot overwrite nested dict at '{key_path}' with scalar value")
|
||||
current[last_key] = _convert_value(value_str)
|
||||
return result
|
||||
|
||||
|
||||
def _convert_value(value_str: str) -> Any:
|
||||
"""Convert string to appropriate Python type.
|
||||
|
||||
Only converts "true"/"false" (case-insensitive) to boolean.
|
||||
Use JSON format (e.g., '"yes"', '"no"') to preserve these as strings.
|
||||
Leading-zero strings (e.g., "007", "00501") are kept as strings.
|
||||
"""
|
||||
s = value_str.strip()
|
||||
lower = s.lower()
|
||||
|
||||
# Handle special values (null, bool)
|
||||
if lower in ("none", "null"):
|
||||
return None
|
||||
if lower == "true":
|
||||
return True
|
||||
if lower == "false":
|
||||
return False
|
||||
|
||||
# Skip int/float for leading-zero strings to keep zip codes / ids intact
|
||||
if not _LEADING_ZERO_RE.match(s):
|
||||
for converter in (int, float):
|
||||
try:
|
||||
return converter(s)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
# JSON handles lists, dicts, and explicitly-quoted strings
|
||||
try:
|
||||
return json.loads(s)
|
||||
except (ValueError, json.JSONDecodeError):
|
||||
pass
|
||||
|
||||
# Fallback to original string
|
||||
return s
|
||||
|
||||
|
||||
def _load_config(name_or_path: str, encoding: str = "utf-8") -> dict:
|
||||
"""Load a YAML or JSON config file.
|
||||
|
||||
First check if name_or_path matches a pre-discovered config (key in _CONFIG_REGISTRY).
|
||||
If not, treat as a file path and load directly.
|
||||
"""
|
||||
# 1. Try pre-discovered configs first
|
||||
if name_or_path in _CONFIG_REGISTRY:
|
||||
return _read_config_file(_CONFIG_REGISTRY[name_or_path], encoding)
|
||||
|
||||
# 2. Treat as file path
|
||||
p = Path(name_or_path)
|
||||
if p.suffix in _SUPPORTED_EXTS:
|
||||
if not p.exists():
|
||||
raise FileNotFoundError(f"Config file not found: {p}")
|
||||
return _read_config_file(p, encoding)
|
||||
|
||||
known = ", ".join(sorted(_CONFIG_REGISTRY)) if _CONFIG_REGISTRY else "none"
|
||||
raise FileNotFoundError(f"Config file not found: {name_or_path}. Available: {known}")
|
||||
|
||||
|
||||
def _read_config_file(path: Path, encoding: str = "utf-8") -> dict:
|
||||
"""Read YAML or JSON file based on extension. Expands ${ENV_VAR}."""
|
||||
with path.open(encoding=encoding) as f:
|
||||
if path.suffix == ".json":
|
||||
result = json.load(f)
|
||||
else:
|
||||
result = yaml.safe_load(f)
|
||||
if result is None:
|
||||
return {}
|
||||
return _expand_env_vars(result)
|
||||
|
||||
|
||||
def _deep_merge(base: dict, update: dict) -> dict:
|
||||
"""Recursively merge dicts."""
|
||||
result = base.copy()
|
||||
for k, v in update.items():
|
||||
if k in result and isinstance(result[k], dict) and isinstance(v, dict):
|
||||
result[k] = _deep_merge(result[k], v)
|
||||
else:
|
||||
result[k] = v
|
||||
return result
|
||||
|
||||
|
||||
def _strip_arg_dashes(arg: str) -> str:
|
||||
"""Strip a single leading `--` or `-` prefix (not all leading dashes)."""
|
||||
if arg.startswith("--"):
|
||||
return arg[2:]
|
||||
if arg.startswith("-"):
|
||||
return arg[1:]
|
||||
return arg
|
||||
|
||||
|
||||
def parse_args(*args, **kwargs) -> tuple[str, dict]:
|
||||
"""Parse CLI args: first arg is action, rest are config overrides.
|
||||
|
||||
Usage: reme app config=paw.yaml service.name=test
|
||||
Returns: (action, merged_config_dict)
|
||||
"""
|
||||
if not args:
|
||||
raise ValueError("No arguments provided")
|
||||
|
||||
first = _strip_arg_dashes(args[0])
|
||||
if "=" in first:
|
||||
raise ValueError(f"First argument must be action, got: {args[0]}")
|
||||
|
||||
action = first
|
||||
configs: list[dict] = []
|
||||
|
||||
for raw in args[1:]:
|
||||
arg = _strip_arg_dashes(raw)
|
||||
if arg.startswith("config="):
|
||||
path = arg.split("=", 1)[1].strip()
|
||||
if path:
|
||||
configs.append(_load_config(path))
|
||||
elif "=" in arg:
|
||||
configs.append(parse_dot_notation([arg]))
|
||||
|
||||
configs.append(kwargs)
|
||||
|
||||
merged: dict = {}
|
||||
for cfg in configs:
|
||||
merged = _deep_merge(merged, cfg)
|
||||
|
||||
return action, merged
|
||||
7
reme4/constants.py
Normal file
7
reme4/constants.py
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
"""Constants"""
|
||||
|
||||
REME_SERVICE_INFO = "REME_SERVICE_INFO"
|
||||
|
||||
REME_DEFAULT_HOST = "127.0.0.1"
|
||||
|
||||
REME_DEFAULT_PORT = 2333
|
||||
9
reme4/enumeration/__init__.py
Normal file
9
reme4/enumeration/__init__.py
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
"""Enumeration"""
|
||||
|
||||
from .chunk_enum import ChunkEnum
|
||||
from .component_enum import ComponentEnum
|
||||
|
||||
__all__ = [
|
||||
"ChunkEnum",
|
||||
"ComponentEnum",
|
||||
]
|
||||
21
reme4/enumeration/chunk_enum.py
Normal file
21
reme4/enumeration/chunk_enum.py
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
"""Chunk enumeration module."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class ChunkEnum(str, Enum):
|
||||
"""Enumeration of possible chunk categories for stream processing."""
|
||||
|
||||
THINK = "think"
|
||||
|
||||
CONTENT = "content"
|
||||
|
||||
TOOL_CALL = "tool_call"
|
||||
|
||||
TOOL_RESULT = "tool_result"
|
||||
|
||||
USAGE = "usage"
|
||||
|
||||
ERROR = "error"
|
||||
|
||||
DONE = "done"
|
||||
37
reme4/enumeration/component_enum.py
Normal file
37
reme4/enumeration/component_enum.py
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
"""Component enumeration module."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class ComponentEnum(str, Enum):
|
||||
"""Enumeration of component types for dependency injection and registration."""
|
||||
|
||||
BASE = "base"
|
||||
|
||||
AS_LLM = "as_llm"
|
||||
|
||||
AS_LLM_FORMATTER = "as_llm_formatter"
|
||||
|
||||
AS_TOKEN_COUNTER = "as_token_counter"
|
||||
|
||||
EMBEDDING_MODEL = "embedding_model"
|
||||
|
||||
FILE_PARSER = "file_parser"
|
||||
|
||||
FILE_STORE = "file_store"
|
||||
|
||||
FILE_GRAPH = "file_graph"
|
||||
|
||||
FILE_WATCHER = "file_watcher"
|
||||
|
||||
KEYWORD_INDEX = "keyword_index"
|
||||
|
||||
SERVICE = "service"
|
||||
|
||||
CLIENT = "client"
|
||||
|
||||
STEP = "step"
|
||||
|
||||
JOB = "job"
|
||||
|
||||
TOKENIZER = "tokenizer"
|
||||
27
reme4/reme.py
Normal file
27
reme4/reme.py
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
import asyncio
|
||||
import sys
|
||||
|
||||
from .application import Application
|
||||
from .components import R
|
||||
from .config import parse_args
|
||||
from .enumeration import ComponentEnum
|
||||
|
||||
|
||||
class ReMe(Application):
|
||||
"""ReMe memory management application."""
|
||||
|
||||
|
||||
def main():
|
||||
action, config = parse_args(sys.argv[1:])
|
||||
if action == "start":
|
||||
reme = ReMe(**config)
|
||||
reme.run_app()
|
||||
else:
|
||||
backend: str = config.pop("backend", "http")
|
||||
client_cls = R.get(ComponentEnum.CLIENT, backend)
|
||||
client = client_cls(action=action, **config)
|
||||
asyncio.run(client())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
25
reme4/schema/__init__.py
Normal file
25
reme4/schema/__init__.py
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
"""Schema"""
|
||||
|
||||
from .application_config import ApplicationConfig, ComponentConfig, JobConfig
|
||||
from .emb_node import EmbNode
|
||||
from .file_chunk import FileChunk
|
||||
from .file_front_matter import FileFrontMatter
|
||||
from .file_link import FileLink
|
||||
from .file_node import FileNode
|
||||
from .request import Request
|
||||
from .response import Response
|
||||
from .stream_chunk import StreamChunk
|
||||
|
||||
__all__ = [
|
||||
"ApplicationConfig",
|
||||
"ComponentConfig",
|
||||
"JobConfig",
|
||||
"EmbNode",
|
||||
"FileChunk",
|
||||
"FileFrontMatter",
|
||||
"FileLink",
|
||||
"FileNode",
|
||||
"Request",
|
||||
"Response",
|
||||
"StreamChunk",
|
||||
]
|
||||
42
reme4/schema/application_config.py
Normal file
42
reme4/schema/application_config.py
Normal file
|
|
@ -0,0 +1,42 @@
|
|||
"""Application configuration models."""
|
||||
|
||||
import os
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from ..enumeration import ComponentEnum
|
||||
|
||||
|
||||
class ComponentConfig(BaseModel):
|
||||
"""Base config for a component; extra fields allowed for backend-specific options."""
|
||||
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
backend: str = Field(default="", description="Backend implementation class name")
|
||||
|
||||
|
||||
class JobConfig(ComponentConfig):
|
||||
"""Config for a job — an ordered sequence of step components."""
|
||||
|
||||
name: str = Field(default="", description="Unique job identifier")
|
||||
description: str = Field(default="", description="Human-readable description")
|
||||
parameters: dict = Field(default_factory=dict, description="Job-level parameters")
|
||||
steps: list[ComponentConfig] = Field(default_factory=list, description="Ordered step configs")
|
||||
|
||||
|
||||
class ApplicationConfig(BaseModel):
|
||||
"""Root config for the ReMe application."""
|
||||
|
||||
app_name: str = Field(default=os.getenv("APP_NAME", "ReMe"), description="Application display name")
|
||||
working_dir: str = Field(default=".reme", description="Working directory for runtime files")
|
||||
enable_logo: bool = Field(default=False, description="Show ASCII logo on startup")
|
||||
language: str = Field(default="", description="Default language for LLM interactions")
|
||||
log_to_console: bool = Field(default=True, description="Log to console")
|
||||
log_to_file: bool = Field(default=True, description="Log to file")
|
||||
mcp_servers: dict[str, dict] = Field(default_factory=dict, description="MCP server configs by name")
|
||||
service: ComponentConfig = Field(default_factory=ComponentConfig, description="Service endpoint config")
|
||||
jobs: list[JobConfig] = Field(default_factory=list, description="Job definitions")
|
||||
components: dict[ComponentEnum, dict[str, ComponentConfig]] = Field(
|
||||
default_factory=dict,
|
||||
description="Component registry keyed by type then name",
|
||||
)
|
||||
32
reme4/schema/emb_node.py
Normal file
32
reme4/schema/emb_node.py
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
"""Embedding node — base record carrying text and its vector."""
|
||||
|
||||
from uuid import uuid4
|
||||
|
||||
import numpy as np
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
|
||||
|
||||
|
||||
class EmbNode(BaseModel):
|
||||
"""A text record with an optional embedding vector and metadata."""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
id: str = Field(default_factory=lambda: uuid4().hex, description="Unique node id")
|
||||
text: str = Field(default="", description="Text content")
|
||||
embedding: np.ndarray | None = Field(default=None, description="Embedding vector (float16)")
|
||||
metadata: dict = Field(default_factory=dict, description="Arbitrary metadata")
|
||||
|
||||
@field_validator("embedding", mode="before")
|
||||
@classmethod
|
||||
def validate_embedding(cls, v):
|
||||
# Coerce list/tuple inputs into a float16 ndarray for compact storage.
|
||||
if v is None:
|
||||
return v
|
||||
return np.array(v, dtype=np.float16)
|
||||
|
||||
@field_serializer("embedding")
|
||||
def serialize_embedding(self, v: np.ndarray | None, _info):
|
||||
# ndarray is not JSON-serializable; emit a plain list.
|
||||
if v is None:
|
||||
return None
|
||||
return v.tolist()
|
||||
26
reme4/schema/file_chunk.py
Normal file
26
reme4/schema/file_chunk.py
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
"""File chunk — an embedding node tied to a line range in a file."""
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from .emb_node import EmbNode
|
||||
|
||||
|
||||
class FileChunk(EmbNode):
|
||||
"""A chunk of a file with positional info and per-stage retrieval scores."""
|
||||
|
||||
path: str = Field(default="", description="Vault-relative file path")
|
||||
start_line: int = Field(default=0, description="Inclusive start line (0-based)")
|
||||
end_line: int = Field(default=0, description="Exclusive end line")
|
||||
scores: dict[str, float] = Field(default_factory=dict, description="Retrieval scores keyed by stage")
|
||||
|
||||
@property
|
||||
def score(self) -> float:
|
||||
"""Final aggregated score; 0.0 if not yet computed."""
|
||||
return self.scores.get("score", 0.0)
|
||||
|
||||
def set_hash_id(self):
|
||||
"""Replace ``id`` with a deterministic hash of (path, range, text)."""
|
||||
from ..utils import hash_text
|
||||
|
||||
self.id = hash_text(" ".join([self.path, str(self.start_line), str(self.end_line), self.text]))
|
||||
return self
|
||||
24
reme4/schema/file_front_matter.py
Normal file
24
reme4/schema/file_front_matter.py
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
"""FileFrontMatter — parsed Markdown front matter."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class FileFrontMatter(BaseModel):
|
||||
"""Markdown front matter; unknown keys are preserved as extras."""
|
||||
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
title: str = Field(default="", description="Document title")
|
||||
description: str = Field(default="", description="Document description")
|
||||
tags: list[str] | None = Field(default=None, description="Tags; None if absent")
|
||||
|
||||
@property
|
||||
def model_extra(self) -> dict[str, Any] | None:
|
||||
"""Get extra fields set during validation.
|
||||
|
||||
Returns:
|
||||
A dictionary of extra fields, or `None` if `config.extra` is not set to `"allow"`.
|
||||
"""
|
||||
return self.__pydantic_extra__
|
||||
20
reme4/schema/file_link.py
Normal file
20
reme4/schema/file_link.py
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
"""FileLink"""
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class FileLink(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
path: str = Field(
|
||||
default=...,
|
||||
description="Wikilink target — raw text pre-resolution, vault-relative path after.",
|
||||
)
|
||||
anchor: str | None = Field(
|
||||
default=None,
|
||||
description="Heading or block anchor (text after '#'); None if absent.",
|
||||
)
|
||||
predicate: str | None = Field(
|
||||
default=None,
|
||||
description="Dataview-style typed-link predicate; None for bare [[X]].",
|
||||
)
|
||||
16
reme4/schema/file_node.py
Normal file
16
reme4/schema/file_node.py
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
"""File node — a file's metadata, links, and chunk references in the graph."""
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .file_front_matter import FileFrontMatter
|
||||
from .file_link import FileLink
|
||||
|
||||
|
||||
class FileNode(BaseModel):
|
||||
"""A vault file as a graph node."""
|
||||
|
||||
path: str = Field(default=..., description="Vault-relative file path")
|
||||
st_mtime: float = Field(default=..., description="Filesystem mtime (seconds)")
|
||||
links: list[FileLink] = Field(default_factory=list, description="Outgoing wikilinks")
|
||||
chunk_ids: list[str] = Field(default_factory=list, description="Owned FileChunk ids")
|
||||
front_matter: FileFrontMatter = Field(default_factory=FileFrontMatter, description="Parsed front matter")
|
||||
11
reme4/schema/request.py
Normal file
11
reme4/schema/request.py
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
"""Request schema for service endpoints."""
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class Request(BaseModel):
|
||||
"""Incoming service request; extra fields are allowed for endpoint-specific payloads."""
|
||||
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
metadata: dict = Field(default_factory=dict, description="Request metadata for context")
|
||||
15
reme4/schema/response.py
Normal file
15
reme4/schema/response.py
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
"""Response schema for service endpoints and LLM calls."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class Response(BaseModel):
|
||||
"""Standard response envelope; extra fields allowed for endpoint-specific output."""
|
||||
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
answer: str | Any = Field(default="", description="Response content or result data")
|
||||
success: bool = Field(default=True, description="Whether the operation succeeded")
|
||||
metadata: dict = Field(default_factory=dict, description="Additional context and diagnostics")
|
||||
14
reme4/schema/stream_chunk.py
Normal file
14
reme4/schema/stream_chunk.py
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
"""Stream chunk schema for incremental responses (e.g. LLM streaming)."""
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ..enumeration import ChunkEnum
|
||||
|
||||
|
||||
class StreamChunk(BaseModel):
|
||||
"""A single chunk in a streaming response sequence."""
|
||||
|
||||
chunk_type: ChunkEnum = Field(default=ChunkEnum.CONTENT, description="Type of chunk content")
|
||||
chunk: str | dict | list = Field(default="", description="Chunk payload")
|
||||
done: bool = Field(default=False, description="Whether this is the final chunk")
|
||||
metadata: dict = Field(default_factory=dict, description="Chunk metadata")
|
||||
7
reme4/steps/__init__.py
Normal file
7
reme4/steps/__init__.py
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
"""steps"""
|
||||
|
||||
from .base_step import BaseStep
|
||||
|
||||
__all__ = [
|
||||
"BaseStep",
|
||||
]
|
||||
118
reme4/steps/base_step.py
Normal file
118
reme4/steps/base_step.py
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
"""Base step class for LLM workflow execution."""
|
||||
|
||||
import copy
|
||||
from abc import abstractmethod, ABC
|
||||
from typing import TypeVar
|
||||
|
||||
from agentscope.formatter import FormatterBase
|
||||
from agentscope.model import ChatModelBase
|
||||
from agentscope.token import TokenCounterBase
|
||||
|
||||
from ..components import ApplicationContext
|
||||
from ..components.embedding import BaseEmbeddingModel
|
||||
from ..components.file_parser import BaseFileParser
|
||||
from ..components.file_store import BaseFileStore
|
||||
from ..components.prompt_handler import PromptHandler
|
||||
from ..components.runtime_context import RuntimeContext
|
||||
from ..enumeration import ComponentEnum
|
||||
from ..utils import get_logger
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class BaseStep(ABC):
|
||||
"""Composable unit of an LLM workflow."""
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
# Snapshot init args so copy() can rebuild an equivalent instance later.
|
||||
instance = object.__new__(cls)
|
||||
instance._init_args = copy.copy(args)
|
||||
instance._init_kwargs = copy.copy(kwargs)
|
||||
return instance
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str | None = None,
|
||||
backend: str = "",
|
||||
app_context: "ApplicationContext | None" = None,
|
||||
language: str = "",
|
||||
prompt_dict: dict[str, str] | None = None,
|
||||
input_mapping: dict[str, str] | None = None,
|
||||
output_mapping: dict[str, str] | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.name: str = name or self.__class__.__name__
|
||||
self.backend: str = backend
|
||||
self.app_context: "ApplicationContext | None" = app_context
|
||||
self.language: str = language
|
||||
self.input_mapping = input_mapping
|
||||
self.output_mapping = output_mapping
|
||||
self.kwargs: dict = kwargs
|
||||
self.context: RuntimeContext | None = None
|
||||
|
||||
self.logger = get_logger()
|
||||
if hasattr(self.logger, "bind"):
|
||||
self.logger = self.logger.bind(component=self.name)
|
||||
|
||||
# Load class-level prompts first, then overlay caller-provided overrides.
|
||||
self.prompt = PromptHandler(language=self.language)
|
||||
self.prompt.load_prompt_by_class(self.__class__).load_prompt_dict(prompt_dict)
|
||||
|
||||
@abstractmethod
|
||||
async def execute(self):
|
||||
"""Run the step's logic against ``self.context``."""
|
||||
|
||||
async def __call__(self, context: RuntimeContext | None = None, **kwargs):
|
||||
# Build runtime context, then apply key remapping around execute().
|
||||
self.context = RuntimeContext.from_context(context, **kwargs)
|
||||
assert self.context is not None
|
||||
if self.input_mapping:
|
||||
self.context.apply_mapping(self.input_mapping)
|
||||
result = await self.execute()
|
||||
if self.output_mapping:
|
||||
self.context.apply_mapping(self.output_mapping)
|
||||
return result
|
||||
|
||||
def _resolve(self, key: str, base_cls: type[T], comp_enum: ComponentEnum, attr: str | None = None) -> T:
|
||||
"""Return a kwargs-supplied instance, or look one up by name in the app registry."""
|
||||
value = self.kwargs.get(key, "default")
|
||||
if isinstance(value, base_cls):
|
||||
return value
|
||||
assert self.app_context is not None
|
||||
comp = self.app_context.components[comp_enum][value]
|
||||
return getattr(comp, attr) if attr else comp
|
||||
|
||||
@property
|
||||
def as_llm(self) -> ChatModelBase:
|
||||
return self._resolve("as_llm", ChatModelBase, ComponentEnum.AS_LLM, "model")
|
||||
|
||||
@property
|
||||
def as_llm_formatter(self) -> FormatterBase:
|
||||
return self._resolve("as_llm_formatter", FormatterBase, ComponentEnum.AS_LLM_FORMATTER, "formatter")
|
||||
|
||||
@property
|
||||
def as_token_counter(self) -> TokenCounterBase:
|
||||
return self._resolve("as_token_counter", TokenCounterBase, ComponentEnum.AS_TOKEN_COUNTER, "token_counter")
|
||||
|
||||
@property
|
||||
def file_parser(self) -> BaseFileParser:
|
||||
return self._resolve("file_parser", BaseFileParser, ComponentEnum.FILE_PARSER)
|
||||
|
||||
@property
|
||||
def file_store(self) -> BaseFileStore:
|
||||
return self._resolve("file_store", BaseFileStore, ComponentEnum.FILE_STORE)
|
||||
|
||||
@property
|
||||
def embedding(self) -> BaseEmbeddingModel:
|
||||
return self._resolve("embedding", BaseEmbeddingModel, ComponentEnum.EMBEDDING_MODEL)
|
||||
|
||||
def prompt_format(self, prompt_name: str, **kwargs) -> str:
|
||||
return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs)
|
||||
|
||||
def get_prompt(self, prompt_name: str) -> str:
|
||||
return self.prompt.get_prompt(prompt_name=prompt_name)
|
||||
|
||||
def copy(self, **kwargs) -> "BaseStep":
|
||||
"""Construct a new instance from the original init args, applying overrides."""
|
||||
return self.__class__(*self._init_args, **{**self._init_kwargs, **kwargs})
|
||||
15
reme4/utils/__init__.py
Normal file
15
reme4/utils/__init__.py
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
"""Utility modules."""
|
||||
|
||||
from .common_utils import hash_text, execute_stream_task
|
||||
from .logger_utils import get_logger
|
||||
from .logo_utils import print_logo
|
||||
from .similarity_utils import cosine_similarity, batch_cosine_similarity
|
||||
|
||||
__all__ = [
|
||||
"hash_text",
|
||||
"execute_stream_task",
|
||||
"get_logger",
|
||||
"print_logo",
|
||||
"cosine_similarity",
|
||||
"batch_cosine_similarity",
|
||||
]
|
||||
105
reme4/utils/common_utils.py
Normal file
105
reme4/utils/common_utils.py
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
"""Common utilities: hashing and async stream task execution."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any, Literal
|
||||
|
||||
from .logger_utils import get_logger
|
||||
from ..enumeration import ChunkEnum
|
||||
from ..schema import StreamChunk
|
||||
|
||||
|
||||
def hash_text(text: str, encoding: str = "utf-8") -> str:
|
||||
"""Return SHA-256 hex digest of text."""
|
||||
return hashlib.sha256(text.encode(encoding)).hexdigest()
|
||||
|
||||
|
||||
def _format_chunk(
|
||||
chunk: StreamChunk,
|
||||
output_format: Literal["str", "bytes", "chunk"],
|
||||
) -> str | bytes | StreamChunk:
|
||||
"""Render a StreamChunk in the requested transport format."""
|
||||
if output_format == "chunk":
|
||||
return chunk
|
||||
data = "data:[DONE]\n\n" if chunk.done else f"data:{chunk.model_dump_json()}\n\n"
|
||||
return data.encode() if output_format == "bytes" else data
|
||||
|
||||
|
||||
async def execute_stream_task(
|
||||
stream_queue: asyncio.Queue[StreamChunk],
|
||||
task: asyncio.Task[Any],
|
||||
task_name: str | None = None,
|
||||
output_format: Literal["str", "bytes", "chunk"] = "str",
|
||||
) -> AsyncGenerator[str | bytes | StreamChunk, None]:
|
||||
"""Yield chunks from stream_queue while monitoring task; cancels task on exit.
|
||||
|
||||
output_format: "str"/"bytes" emit SSE frames, "chunk" emits raw StreamChunk.
|
||||
"""
|
||||
logger = get_logger()
|
||||
consumer: asyncio.Task[StreamChunk] | None = None
|
||||
try:
|
||||
while True:
|
||||
consumer = get_chunk = asyncio.create_task(stream_queue.get())
|
||||
done, _pending = await asyncio.wait({get_chunk, task}, return_when=asyncio.FIRST_COMPLETED)
|
||||
|
||||
# Producer still running — relay the next chunk and continue.
|
||||
if task not in done:
|
||||
chunk = get_chunk.result()
|
||||
yield _format_chunk(chunk, output_format)
|
||||
if chunk.done:
|
||||
return
|
||||
continue
|
||||
|
||||
# Producer finished. Capture any pending chunk, then stop the consumer wait
|
||||
# so we can inspect task state safely.
|
||||
pending_chunk: StreamChunk | None = None
|
||||
if get_chunk in done:
|
||||
pending_chunk = get_chunk.result()
|
||||
else:
|
||||
get_chunk.cancel()
|
||||
try:
|
||||
await get_chunk
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# Surface task failure first — an exception trumps trailing data.
|
||||
if task.cancelled():
|
||||
msg = f"Task cancelled: {task_name}" if task_name else "Task cancelled"
|
||||
raise asyncio.CancelledError(msg)
|
||||
exc = task.exception()
|
||||
if exc is not None:
|
||||
log_msg = f"Task error in {task_name}: {exc}" if task_name else f"Task error: {exc}"
|
||||
logger.error(log_msg, exc_info=exc)
|
||||
raise exc
|
||||
|
||||
# Producer ended cleanly — flush pending + drain queue so no chunk is lost,
|
||||
# then emit the terminal sentinel.
|
||||
if pending_chunk is not None:
|
||||
yield _format_chunk(pending_chunk, output_format)
|
||||
if pending_chunk.done:
|
||||
return
|
||||
while not stream_queue.empty():
|
||||
chunk = stream_queue.get_nowait()
|
||||
yield _format_chunk(chunk, output_format)
|
||||
if chunk.done:
|
||||
return
|
||||
|
||||
yield _format_chunk(StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True), output_format)
|
||||
return
|
||||
|
||||
finally:
|
||||
# Cancel consumer wait if still pending (e.g. on consumer aclose).
|
||||
if consumer is not None and not consumer.done():
|
||||
consumer.cancel()
|
||||
try:
|
||||
await consumer
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
# Cancel producer task if still running to avoid resource leaks.
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
108
reme4/utils/logger_utils.py
Normal file
108
reme4/utils/logger_utils.py
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
"""Logger utilities supporting both loguru and standard logging backends."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from logging.handlers import TimedRotatingFileHandler
|
||||
|
||||
_logger = None
|
||||
|
||||
_LOGURU_FORMAT = "{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}"
|
||||
_STDLIB_FORMAT = "%(asctime)s | %(levelname)s | %(filename)s:%(lineno)d | %(funcName)s | %(message)s"
|
||||
_STDLIB_DATEFMT = "%Y-%m-%d %H:%M:%S"
|
||||
|
||||
|
||||
def _enable_loguru() -> bool:
|
||||
return os.getenv("REME_DISABLE_LOGURU", "").lower() != "true"
|
||||
|
||||
|
||||
def _init_loguru(log_dir: str, level: str, log_to_console: bool, log_to_file: bool):
|
||||
from loguru import logger
|
||||
|
||||
logger.remove()
|
||||
|
||||
if log_to_console:
|
||||
logger.add(
|
||||
sink=sys.stdout,
|
||||
level=level,
|
||||
format=_LOGURU_FORMAT,
|
||||
colorize=True,
|
||||
)
|
||||
|
||||
if log_to_file:
|
||||
try:
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
log_filepath = os.path.join(log_dir, f"{current_ts}.log")
|
||||
|
||||
logger.add(
|
||||
log_filepath,
|
||||
level=level,
|
||||
rotation="00:00",
|
||||
retention="7 days",
|
||||
compression="zip",
|
||||
encoding="utf-8",
|
||||
format=_LOGURU_FORMAT,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error configuring file logging: {e}")
|
||||
|
||||
return logger
|
||||
|
||||
|
||||
def _init_stdlib(log_dir: str, level: str, log_to_console: bool, log_to_file: bool):
|
||||
logger = logging.getLogger("reme")
|
||||
logger.setLevel(level)
|
||||
logger.propagate = False
|
||||
|
||||
for handler in list(logger.handlers):
|
||||
logger.removeHandler(handler)
|
||||
|
||||
formatter = logging.Formatter(_STDLIB_FORMAT, datefmt=_STDLIB_DATEFMT)
|
||||
|
||||
if log_to_console:
|
||||
console_handler = logging.StreamHandler(sys.stdout)
|
||||
console_handler.setLevel(level)
|
||||
console_handler.setFormatter(formatter)
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
if log_to_file:
|
||||
try:
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
log_filepath = os.path.join(log_dir, f"{current_ts}.log")
|
||||
|
||||
file_handler = TimedRotatingFileHandler(
|
||||
log_filepath,
|
||||
when="midnight",
|
||||
backupCount=7,
|
||||
encoding="utf-8",
|
||||
)
|
||||
file_handler.setLevel(level)
|
||||
file_handler.setFormatter(formatter)
|
||||
logger.addHandler(file_handler)
|
||||
except Exception as e:
|
||||
logger.error(f"Error configuring file logging: {e}")
|
||||
|
||||
return logger
|
||||
|
||||
|
||||
def get_logger(
|
||||
log_dir: str = "logs",
|
||||
level: str = "INFO",
|
||||
log_to_console: bool = True,
|
||||
log_to_file: bool = True,
|
||||
force_init: bool = False,
|
||||
):
|
||||
"""Return the global logger, initializing sinks on first call (or when force_init)."""
|
||||
global _logger
|
||||
|
||||
if _logger is not None and not force_init:
|
||||
return _logger
|
||||
|
||||
if _enable_loguru():
|
||||
_logger = _init_loguru(log_dir, level, log_to_console, log_to_file)
|
||||
else:
|
||||
_logger = _init_stdlib(log_dir, level, log_to_console, log_to_file)
|
||||
return _logger
|
||||
87
reme4/utils/logo_utils.py
Normal file
87
reme4/utils/logo_utils.py
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
"""Startup banner with ASCII logo and service metadata."""
|
||||
|
||||
import importlib.metadata
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from rich.console import Console, Group
|
||||
from rich.panel import Panel
|
||||
from rich.table import Table
|
||||
from rich.text import Text
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..schema import ApplicationConfig
|
||||
|
||||
|
||||
def get_version(package_name: str) -> str:
|
||||
"""Return installed package version, or empty string if not installed."""
|
||||
try:
|
||||
return importlib.metadata.version(package_name)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
return ""
|
||||
|
||||
|
||||
def print_logo(app_config: "ApplicationConfig"):
|
||||
"""Print gradient ASCII logo and runtime config (backend, URL, versions)."""
|
||||
ascii_art = [
|
||||
r" ██████╗ ███████╗ ███╗ ███╗ ███████╗ ",
|
||||
r" ██╔══██╗ ██╔════╝ ████╗ ████║ ██╔════╝ ",
|
||||
r" ██████╔╝ █████╗ ██╔████╔██║ █████╗ ",
|
||||
r" ██╔══██╗ ██╔══╝ ██║╚██╔╝██║ ██╔══╝ ",
|
||||
r" ██║ ██║ ███████╗ ██║ ╚═╝ ██║ ███████╗ ",
|
||||
r" ╚═╝ ╚═╝ ╚══════╝ ╚═╝ ╚═╝ ╚══════╝ ",
|
||||
]
|
||||
|
||||
start_color = (85, 239, 196)
|
||||
end_color = (162, 155, 254)
|
||||
|
||||
logo_text = Text()
|
||||
for line in ascii_art:
|
||||
line_len = max(1, len(line) - 1)
|
||||
for i, char in enumerate(line):
|
||||
ratio = i / line_len
|
||||
rgb = tuple(int(s + (e - s) * ratio) for s, e in zip(start_color, end_color))
|
||||
logo_text.append(char, style=f"bold rgb({rgb[0]},{rgb[1]},{rgb[2]})")
|
||||
logo_text.append("\n")
|
||||
|
||||
info_table = Table.grid(padding=(0, 1))
|
||||
info_table.add_column(style="bold", justify="center")
|
||||
info_table.add_column(style="bold cyan", justify="left")
|
||||
info_table.add_column(style="white", justify="left")
|
||||
|
||||
# service is a ComponentConfig with extra="allow"; backend-specific fields live in model_extra.
|
||||
service = app_config.service
|
||||
backend = service.backend
|
||||
extra = service.model_extra or {}
|
||||
|
||||
info_table.add_row("📦", "Backend:", backend)
|
||||
|
||||
match backend:
|
||||
case "http":
|
||||
host = extra.get("host", "localhost")
|
||||
port = extra.get("port", 8000)
|
||||
info_table.add_row("🔗", "URL:", f"http://{host}:{port}")
|
||||
info_table.add_row("📚", "FastAPI:", Text(get_version("fastapi"), style="dim"))
|
||||
case "mcp":
|
||||
transport = extra.get("transport", "stdio")
|
||||
info_table.add_row("🚌", "Transport:", transport)
|
||||
if transport != "stdio":
|
||||
host = extra.get("host", "localhost")
|
||||
port = extra.get("port", 8000)
|
||||
url = f"http://{host}:{port}"
|
||||
if transport == "sse":
|
||||
url += "/sse"
|
||||
info_table.add_row("🔗", "URL:", url)
|
||||
info_table.add_row("📚", "FastMCP:", Text(get_version("fastmcp"), style="dim"))
|
||||
|
||||
info_table.add_row("🚀", "ReMe:", Text(get_version("reme-ai"), style="dim"))
|
||||
|
||||
panel = Panel(
|
||||
Group(logo_text, info_table),
|
||||
title=app_config.app_name,
|
||||
title_align="left",
|
||||
border_style="dim",
|
||||
padding=(1, 4),
|
||||
expand=False,
|
||||
)
|
||||
|
||||
Console().print(Group("\n", panel, "\n"))
|
||||
35
reme4/utils/similarity_utils.py
Normal file
35
reme4/utils/similarity_utils.py
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
"""Cosine similarity for single vectors and batched matrices."""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
|
||||
"""Cosine similarity of two equal-length vectors; returns 0.0 if either has zero norm."""
|
||||
if len(vec1) != len(vec2):
|
||||
raise ValueError(f"Vectors must have same length: {len(vec1)} != {len(vec2)}")
|
||||
|
||||
dot_product = sum(a * b for a, b in zip(vec1, vec2))
|
||||
magnitude1 = sum(a * a for a in vec1) ** 0.5
|
||||
magnitude2 = sum(b * b for b in vec2) ** 0.5
|
||||
|
||||
if magnitude1 == 0 or magnitude2 == 0:
|
||||
return 0.0
|
||||
|
||||
return dot_product / (magnitude1 * magnitude2)
|
||||
|
||||
|
||||
def batch_cosine_similarity(nd_array1: np.ndarray, nd_array2: np.ndarray) -> np.ndarray:
|
||||
"""Pairwise cosine similarity matrix between two batches; output shape (N1, N2)."""
|
||||
if nd_array1.shape[1] != nd_array2.shape[1]:
|
||||
raise ValueError(
|
||||
f"Embedding dimensions must match: {nd_array1.shape[1]} != {nd_array2.shape[1]}",
|
||||
)
|
||||
|
||||
dot_products = np.dot(nd_array1, nd_array2.T)
|
||||
norms1 = np.linalg.norm(nd_array1, axis=1)
|
||||
norms2 = np.linalg.norm(nd_array2, axis=1)
|
||||
norm_products = np.outer(norms1, norms2)
|
||||
# Guard against zero-norm rows to keep division finite.
|
||||
norm_products = np.where(norm_products == 0, 1e-10, norm_products)
|
||||
|
||||
return dot_products / norm_products
|
||||
Loading…
Add table
Reference in a new issue