This commit is contained in:
jinli.yl 2026-05-15 23:17:49 +08:00
parent e5960d0ad0
commit c8e96b5ae8
76 changed files with 6141 additions and 0 deletions

378
docs4/expert.yaml Normal file
View 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
View 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
View 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. 抽象基类
. graphfellow 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 在前加载graphfile_watcher在后递归扫描目录
a. 通过ms_time对比graphon_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. 原子stepsreme 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
View 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.

View file

@ -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
View 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)

View 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",
]

View 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] = {}

View 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",
]

View 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",
]

View 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

View 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",
]

View 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)

View 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()

View file

@ -0,0 +1,6 @@
"""Client components."""
from .base_client import BaseClient
from .http_client import HttpClient
__all__ = ["BaseClient", "HttpClient"]

View 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."""

View 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

View 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()

View file

@ -0,0 +1,6 @@
"""Embedding model implementations."""
from .base_embedding_model import BaseEmbeddingModel
from .openai_embedding_model import OpenAIEmbeddingModel
__all__ = ["BaseEmbeddingModel", "OpenAIEmbeddingModel"]

View 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")

View 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

View 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"]

View 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*."""

View 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]

View 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]

View 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"]

View 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=[]), []

View 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)."""

View 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

View 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",
]

View 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)

View 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

View 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",
]

View 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."""

View 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])

View file

@ -0,0 +1,6 @@
"""Job components for executing workflows."""
from .base_job import BaseJob
from .stream_job import StreamJob
__all__ = ["BaseJob", "StreamJob"]

View 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

View 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()

View file

@ -0,0 +1,4 @@
from .base_keyword_index import BaseKeywordIndex
from .bm25_index import BM25Index
__all__ = ["BaseKeywordIndex", "BM25Index"]

View 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."""

View 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 = {}

View 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)})"

View 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

View 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",
]

View 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)

View 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)

View 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)

View 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",
]

View 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."""

View 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

View 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

File diff suppressed because it is too large Load diff

5
reme4/config/__init__.py Normal file
View file

@ -0,0 +1,5 @@
"""Config"""
from .config_parser import parse_args
__all__ = ["parse_args"]

View 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
View file

@ -0,0 +1,7 @@
"""Constants"""
REME_SERVICE_INFO = "REME_SERVICE_INFO"
REME_DEFAULT_HOST = "127.0.0.1"
REME_DEFAULT_PORT = 2333

View file

@ -0,0 +1,9 @@
"""Enumeration"""
from .chunk_enum import ChunkEnum
from .component_enum import ComponentEnum
__all__ = [
"ChunkEnum",
"ComponentEnum",
]

View 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"

View 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
View 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
View 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",
]

View 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
View 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()

View 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

View 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
View 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
View 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
View 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
View 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")

View 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
View file

@ -0,0 +1,7 @@
"""steps"""
from .base_step import BaseStep
__all__ = [
"BaseStep",
]

118
reme4/steps/base_step.py Normal file
View 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
View 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
View 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
View 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
View 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"))

View 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