diff --git a/docs4/expert.yaml b/docs4/expert.yaml new file mode 100644 index 00000000..1cb2e0e0 --- /dev/null +++ b/docs4/expert.yaml @@ -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 + `//`. + 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. diff --git a/docs4/obsidian.md b/docs4/obsidian.md new file mode 100644 index 00000000..4d342f53 --- /dev/null +++ b/docs4/obsidian.md @@ -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 信息，主动决定推送给用户的消息 | + + + + + diff --git a/docs4/reme_todo.md b/docs4/reme_todo.md new file mode 100644 index 00000000..a3fecfa6 --- /dev/null +++ b/docs4/reme_todo.md @@ -0,0 +1,86 @@ +ReMe重构 +1. 根目录 + a. ✅ vault_root → working_dir @sen +2. file_parser + a. 抽象基类 parse: @jinli + ⅰ. 输入是path：相对路径 + ⅱ. 输出是FileMetadata & list[FileChunks] & list[FileEdge] + b. default parser 兼容老方案 @jinli + ⅰ. 带overlap的chunking策略 ，不输出FileEdge + c. markdown parser @sen + ⅰ. 根据markdown ast做chunk，不需要overlap + ⅱ. 增加一个索引的chunk chunk_type @锦鲤 file_chunk_type content/index + ⅲ. 增加link的正则解析：predicate:: [[path#anchor]] +3. file_store @sen + a. 抽象存储： + ⅰ. filenode = file + path + st_mtime + metadata + list[FileEdge] + ⅱ. graph=dict[str, filenode] 内存+json + ⅲ. list[FileChunk] 存db + b. 抽象基类 + ⅰ. graph：fellow dict的操作 update/get/set + ⅱ. chunks dict[str, list[chunk]] + 1. delete_chunks_by_path + 2. update_chunks_by_path + 3. list_chunks_by_path + 4. vector_search/keyword_search + ⅲ. 手写一个bm25检索 + ⅳ. 【核心】检索机制 vector bm25 graph 如何进行融合 +4. file_watcher @jinli + a. 抽象基类 + ⅰ. on_start: + 1. file_store 的start 在前，加载graph，file_watcher在后，递归扫描目录 + a. 通过ms_time对比graph，on_change 进行改动 + ⅱ. on_change: + 1. 更新/增加: + a. delete_chunks_by_path 更新数据库 + b. upate_chunks_by_path 更新数据库 + c. 更新graph + 2. 删除 + a. delete_chunks_by_path 更新数据库 + +MemorySchema +1. markdown文件结构 @sen + a. formatter： + ⅰ. title + ⅱ. desc + ⅲ. tags + ⅳ. +2. memory文件结构目录 + a. MEMORY.md + b. msg/files -> daily/YYYYMMDD/YYYYMMDD.md + xxxx.md + ⅰ. YYYYMMDD.md + 1. xxx -> xxxx.md + 2. xxx -> xxxd.md + ⅱ. + c. daily -> topic/topic_l1/topic_l1.md + xxx.md + topic_l2 + d. proactive + +steps: +1. 治理（算法+LLM）： + a. 节点关联P0：现有的链接做补充，挖掘新的LLM的link + ⅰ. /Users/yuli/workspace/ReMe/reme2/component/edge_extractor/llm_edge_extractor.py + ⅱ. 移动到steps + b. 节点整合/节点拆分/节点归档 + c. 健康度检查 +2. retrieve 调用store的检索 +3. 原子steps：reme edit +4. 组合steps：总结： + a. - freq (every_n_turn、compact) -> daily_summarizer + b. topic (/dream ) -> topic_summarizer(daily_xx -> topic_xx) + c. proactive -> proactive_summarizer(personal_xxx -> proactive_query - pre_query + + + + + + + + + + + + + + + +Qwenpaw \ No newline at end of file diff --git a/docs4/service.yaml b/docs4/service.yaml new file mode 100644 index 00000000..bf2e3fe6 --- /dev/null +++ b/docs4/service.yaml @@ -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 `/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. diff --git a/reme4/__init__.py b/reme4/__init__.py index e69de29b..76b34dcf 100644 --- a/reme4/__init__.py +++ b/reme4/__init__.py @@ -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", +] diff --git a/reme4/application.py b/reme4/application.py new file mode 100644 index 00000000..f98eaa02 --- /dev/null +++ b/reme4/application.py @@ -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) diff --git a/reme4/components/__init__.py b/reme4/components/__init__.py new file mode 100644 index 00000000..148eeeaa --- /dev/null +++ b/reme4/components/__init__.py @@ -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", +] diff --git a/reme4/components/application_context.py b/reme4/components/application_context.py new file mode 100644 index 00000000..90a7f194 --- /dev/null +++ b/reme4/components/application_context.py @@ -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] = {} diff --git a/reme4/components/as_llm/__init__.py b/reme4/components/as_llm/__init__.py new file mode 100644 index 00000000..82ae553f --- /dev/null +++ b/reme4/components/as_llm/__init__.py @@ -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", +] diff --git a/reme4/components/as_llm_formatter/__init__.py b/reme4/components/as_llm_formatter/__init__.py new file mode 100644 index 00000000..81e4fe37 --- /dev/null +++ b/reme4/components/as_llm_formatter/__init__.py @@ -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", +] diff --git a/reme4/components/as_llm_formatter/reme_openai_chat_formatter.py b/reme4/components/as_llm_formatter/reme_openai_chat_formatter.py new file mode 100644 index 00000000..a0b47082 --- /dev/null +++ b/reme4/components/as_llm_formatter/reme_openai_chat_formatter.py @@ -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="The following are the image contents from the tool " + f"result of '{block['name']}':", + ), + *promoted_blocks, + TextBlock(type="text", text=""), + ] + 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 diff --git a/reme4/components/as_token_counter/__init__.py b/reme4/components/as_token_counter/__init__.py new file mode 100644 index 00000000..36bef1bc --- /dev/null +++ b/reme4/components/as_token_counter/__init__.py @@ -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", +] diff --git a/reme4/components/as_token_counter/estimate_token_counter.py b/reme4/components/as_token_counter/estimate_token_counter.py new file mode 100644 index 00000000..7a7e4fbe --- /dev/null +++ b/reme4/components/as_token_counter/estimate_token_counter.py @@ -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) diff --git a/reme4/components/base_component.py b/reme4/components/base_component.py new file mode 100644 index 00000000..97157ae5 --- /dev/null +++ b/reme4/components/base_component.py @@ -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"" + + 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() diff --git a/reme4/components/client/__init__.py b/reme4/components/client/__init__.py new file mode 100644 index 00000000..e5a1df0a --- /dev/null +++ b/reme4/components/client/__init__.py @@ -0,0 +1,6 @@ +"""Client components.""" + +from .base_client import BaseClient +from .http_client import HttpClient + +__all__ = ["BaseClient", "HttpClient"] diff --git a/reme4/components/client/base_client.py b/reme4/components/client/base_client.py new file mode 100644 index 00000000..42bae599 --- /dev/null +++ b/reme4/components/client/base_client.py @@ -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.""" diff --git a/reme4/components/client/http_client.py b/reme4/components/client/http_client.py new file mode 100644 index 00000000..7a11413a --- /dev/null +++ b/reme4/components/client/http_client.py @@ -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 diff --git a/reme4/components/component_registry.py b/reme4/components/component_registry.py new file mode 100644 index 00000000..c78ceb9c --- /dev/null +++ b/reme4/components/component_registry.py @@ -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() diff --git a/reme4/components/embedding/__init__.py b/reme4/components/embedding/__init__.py new file mode 100644 index 00000000..bf2d4b48 --- /dev/null +++ b/reme4/components/embedding/__init__.py @@ -0,0 +1,6 @@ +"""Embedding model implementations.""" + +from .base_embedding_model import BaseEmbeddingModel +from .openai_embedding_model import OpenAIEmbeddingModel + +__all__ = ["BaseEmbeddingModel", "OpenAIEmbeddingModel"] diff --git a/reme4/components/embedding/base_embedding_model.py b/reme4/components/embedding/base_embedding_model.py new file mode 100644 index 00000000..00a1d795 --- /dev/null +++ b/reme4/components/embedding/base_embedding_model.py @@ -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") diff --git a/reme4/components/embedding/openai_embedding_model.py b/reme4/components/embedding/openai_embedding_model.py new file mode 100644 index 00000000..a03c4a01 --- /dev/null +++ b/reme4/components/embedding/openai_embedding_model.py @@ -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 diff --git a/reme4/components/file_graph/__init__.py b/reme4/components/file_graph/__init__.py new file mode 100644 index 00000000..12b01355 --- /dev/null +++ b/reme4/components/file_graph/__init__.py @@ -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"] diff --git a/reme4/components/file_graph/base_file_graph.py b/reme4/components/file_graph/base_file_graph.py new file mode 100644 index 00000000..7ae47ec7 --- /dev/null +++ b/reme4/components/file_graph/base_file_graph.py @@ -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*.""" diff --git a/reme4/components/file_graph/local_file_graph.py b/reme4/components/file_graph/local_file_graph.py new file mode 100644 index 00000000..b686224e --- /dev/null +++ b/reme4/components/file_graph/local_file_graph.py @@ -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] diff --git a/reme4/components/file_graph/nx_file_graph.py b/reme4/components/file_graph/nx_file_graph.py new file mode 100644 index 00000000..f4f4ea13 --- /dev/null +++ b/reme4/components/file_graph/nx_file_graph.py @@ -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] diff --git a/reme4/components/file_parser/__init__.py b/reme4/components/file_parser/__init__.py new file mode 100644 index 00000000..8898bec4 --- /dev/null +++ b/reme4/components/file_parser/__init__.py @@ -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"] diff --git a/reme4/components/file_parser/bare_file_parser.py b/reme4/components/file_parser/bare_file_parser.py new file mode 100644 index 00000000..daadccb9 --- /dev/null +++ b/reme4/components/file_parser/bare_file_parser.py @@ -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=[]), [] diff --git a/reme4/components/file_parser/base_file_parser.py b/reme4/components/file_parser/base_file_parser.py new file mode 100644 index 00000000..2433a831 --- /dev/null +++ b/reme4/components/file_parser/base_file_parser.py @@ -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).""" diff --git a/reme4/components/file_parser/default_file_parser.py b/reme4/components/file_parser/default_file_parser.py new file mode 100644 index 00000000..53bb10ec --- /dev/null +++ b/reme4/components/file_parser/default_file_parser.py @@ -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 diff --git a/reme4/components/file_store/__init__.py b/reme4/components/file_store/__init__.py new file mode 100644 index 00000000..7739fb99 --- /dev/null +++ b/reme4/components/file_store/__init__.py @@ -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", +] diff --git a/reme4/components/file_store/base_file_store.py b/reme4/components/file_store/base_file_store.py new file mode 100644 index 00000000..05a132fd --- /dev/null +++ b/reme4/components/file_store/base_file_store.py @@ -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) diff --git a/reme4/components/file_store/local_file_store.py b/reme4/components/file_store/local_file_store.py new file mode 100644 index 00000000..ebda9562 --- /dev/null +++ b/reme4/components/file_store/local_file_store.py @@ -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 diff --git a/reme4/components/file_watcher/__init__.py b/reme4/components/file_watcher/__init__.py new file mode 100644 index 00000000..0b11d995 --- /dev/null +++ b/reme4/components/file_watcher/__init__.py @@ -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", +] diff --git a/reme4/components/file_watcher/base_file_watcher.py b/reme4/components/file_watcher/base_file_watcher.py new file mode 100644 index 00000000..6424ef5e --- /dev/null +++ b/reme4/components/file_watcher/base_file_watcher.py @@ -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.""" diff --git a/reme4/components/file_watcher/lite_file_watcher.py b/reme4/components/file_watcher/lite_file_watcher.py new file mode 100644 index 00000000..1d50711e --- /dev/null +++ b/reme4/components/file_watcher/lite_file_watcher.py @@ -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]) diff --git a/reme4/components/job/__init__.py b/reme4/components/job/__init__.py new file mode 100644 index 00000000..0387572d --- /dev/null +++ b/reme4/components/job/__init__.py @@ -0,0 +1,6 @@ +"""Job components for executing workflows.""" + +from .base_job import BaseJob +from .stream_job import StreamJob + +__all__ = ["BaseJob", "StreamJob"] diff --git a/reme4/components/job/base_job.py b/reme4/components/job/base_job.py new file mode 100644 index 00000000..e4681bf1 --- /dev/null +++ b/reme4/components/job/base_job.py @@ -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 diff --git a/reme4/components/job/stream_job.py b/reme4/components/job/stream_job.py new file mode 100644 index 00000000..89f1fab1 --- /dev/null +++ b/reme4/components/job/stream_job.py @@ -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() diff --git a/reme4/components/keyword_index/__init__.py b/reme4/components/keyword_index/__init__.py new file mode 100644 index 00000000..fc08222c --- /dev/null +++ b/reme4/components/keyword_index/__init__.py @@ -0,0 +1,4 @@ +from .base_keyword_index import BaseKeywordIndex +from .bm25_index import BM25Index + +__all__ = ["BaseKeywordIndex", "BM25Index"] diff --git a/reme4/components/keyword_index/base_keyword_index.py b/reme4/components/keyword_index/base_keyword_index.py new file mode 100644 index 00000000..05b63231 --- /dev/null +++ b/reme4/components/keyword_index/base_keyword_index.py @@ -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.""" diff --git a/reme4/components/keyword_index/bm25_index.py b/reme4/components/keyword_index/bm25_index.py new file mode 100644 index 00000000..863704d5 --- /dev/null +++ b/reme4/components/keyword_index/bm25_index.py @@ -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 = {} diff --git a/reme4/components/prompt_handler.py b/reme4/components/prompt_handler.py new file mode 100644 index 00000000..a9ae203a --- /dev/null +++ b/reme4/components/prompt_handler.py @@ -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 ``.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 ``_``.""" + 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)})" diff --git a/reme4/components/runtime_context.py b/reme4/components/runtime_context.py new file mode 100644 index 00000000..92ee2b79 --- /dev/null +++ b/reme4/components/runtime_context.py @@ -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 diff --git a/reme4/components/service/__init__.py b/reme4/components/service/__init__.py new file mode 100644 index 00000000..41d8f7b0 --- /dev/null +++ b/reme4/components/service/__init__.py @@ -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", +] diff --git a/reme4/components/service/base_service.py b/reme4/components/service/base_service.py new file mode 100644 index 00000000..e47e25dd --- /dev/null +++ b/reme4/components/service/base_service.py @@ -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) diff --git a/reme4/components/service/http_service.py b/reme4/components/service/http_service.py new file mode 100644 index 00000000..83ed7105 --- /dev/null +++ b/reme4/components/service/http_service.py @@ -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) diff --git a/reme4/components/service/mcp_service.py b/reme4/components/service/mcp_service.py new file mode 100644 index 00000000..f14dfa96 --- /dev/null +++ b/reme4/components/service/mcp_service.py @@ -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) diff --git a/reme4/components/tokenizer/__init__.py b/reme4/components/tokenizer/__init__.py new file mode 100644 index 00000000..ff364470 --- /dev/null +++ b/reme4/components/tokenizer/__init__.py @@ -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", +] diff --git a/reme4/components/tokenizer/base_tokenizer.py b/reme4/components/tokenizer/base_tokenizer.py new file mode 100644 index 00000000..ee652cb3 --- /dev/null +++ b/reme4/components/tokenizer/base_tokenizer.py @@ -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.""" diff --git a/reme4/components/tokenizer/jieba_tokenizer.py b/reme4/components/tokenizer/jieba_tokenizer.py new file mode 100644 index 00000000..b4f9219f --- /dev/null +++ b/reme4/components/tokenizer/jieba_tokenizer.py @@ -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 diff --git a/reme4/components/tokenizer/regex_tokenizer.py b/reme4/components/tokenizer/regex_tokenizer.py new file mode 100644 index 00000000..4379a60e --- /dev/null +++ b/reme4/components/tokenizer/regex_tokenizer.py @@ -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 diff --git a/reme4/components/tokenizer/stopwords b/reme4/components/tokenizer/stopwords new file mode 100644 index 00000000..fc3c61d0 --- /dev/null +++ b/reme4/components/tokenizer/stopwords @@ -0,0 +1,1395 @@ +-- +? +“ +” +》 +－－ +able +about +above +according +accordingly +across +actually +after +afterwards +again +against +ain't +all +allow +allows +almost +alone +along +already +also +although +always +am +among +amongst +an +and +another +any +anybody +anyhow +anyone +anything +anyway +anyways +anywhere +apart +appear +appreciate +appropriate +are +aren't +around +as +a's +aside +ask +asking +associated +at +available +away +awfully +be +became +because +become +becomes +becoming +been +before +beforehand +behind +being +believe +below +beside +besides +best +better +between +beyond +both +brief +but +by +came +can +cannot +cant +can't +cause +causes +certain +certainly +changes +clearly +c'mon +co +com +come +comes +concerning +consequently +consider +considering +contain +containing +contains +corresponding +could +couldn't +course +c's +currently +definitely +described +despite +did +didn't +different +do +does +doesn't +doing +done +don't +down +downwards +during +each +edu +eg +eight +either +else +elsewhere +enough +entirely +especially +et +etc +even +ever +every +everybody +everyone +everything +everywhere +ex +exactly +example +except +far +few +fifth +first +five +followed +following +follows +for +former +formerly +forth +four +from +further +furthermore +get +gets +getting +given +gives +go +goes +going +gone +got +gotten +greetings +had +hadn't +happens +hardly +has +hasn't +have +haven't +having +he +hello +help +hence +her +here +hereafter +hereby +herein +here's +hereupon +hers +herself +he's +hi +him +himself +his +hither +hopefully +how +howbeit +however +i'd +ie +if +ignored +i'll +i'm +immediate +in +inasmuch +inc +indeed +indicate +indicated +indicates +inner +insofar +instead +into +inward +is +isn't +it +it'd +it'll +its +it's +itself +i've +just +keep +keeps +kept +know +known +knows +last +lately +later +latter +latterly +least +less +lest +let +let's +like +liked +likely +little +look +looking +looks +ltd +mainly +many +may +maybe +me +mean +meanwhile +merely +might +more +moreover +most +mostly +much +must +my +myself +name +namely +nd +near +nearly +necessary +need +needs +neither +never +nevertheless +new +next +nine +no +nobody +non +none +noone +nor +normally +not +nothing +novel +now +nowhere +obviously +of +off +often +oh +ok +okay +old +on +once +one +ones +only +onto +or +other +others +otherwise +ought +our +ours +ourselves +out +outside +over +overall +own +particular +particularly +per +perhaps +placed +please +plus +possible +presumably +probably +provides +que +quite +qv +rather +rd +re +really +reasonably +regarding +regardless +regards +relatively +respectively +right +said +same +saw +say +saying +says +second +secondly +see +seeing +seem +seemed +seeming +seems +seen +self +selves +sensible +sent +serious +seriously +seven +several +shall +she +should +shouldn't +since +six +so +some +somebody +somehow +someone +something +sometime +sometimes +somewhat +somewhere +soon +sorry +specified +specify +specifying +still +sub +such +sup +sure +take +taken +tell +tends +th +than +thank +thanks +thanx +that +thats +that's +the +their +theirs +them +themselves +then +thence +there +thereafter +thereby +therefore +therein +theres +there's +thereupon +these +they +they'd +they'll +they're +they've +think +third +this +thorough +thoroughly +those +though +three +through +throughout +thru +thus +to +together +too +took +toward +towards +tried +tries +truly +try +trying +t's +twice +two +un +under +unfortunately +unless +unlikely +until +unto +up +upon +us +use +used +useful +uses +using +usually +value +various +very +via +viz +vs +want +wants +was +wasn't +way +we +we'd +welcome +well +we'll +went +were +we're +weren't +we've +what +whatever +what's +when +whence +whenever +where +whereafter +whereas +whereby +wherein +where's +whereupon +wherever +whether +which +while +whither +who +whoever +whole +whom +who's +whose +why +will +willing +wish +with +within +without +wonder +won't +would +wouldn't +yes +yet +you +you'd +you'll +your +you're +yours +yourself +yourselves +you've +zero +zt +ZT +zz +ZZ +一 +一下 +一些 +一切 +一则 +一天 +一定 +一方面 +一旦 +一时 +一来 +一样 +一次 +一片 +一直 +一致 +一般 +一起 +一边 +一面 +万一 +上下 +上升 +上去 +上来 +上述 +上面 +下列 +下去 +下来 +下面 +不一 +不久 +不仅 +不会 +不但 +不光 +不单 +不变 +不只 +不可 +不同 +不够 +不如 +不得 +不怕 +不惟 +不成 +不拘 +不敢 +不断 +不是 +不比 +不然 +不特 +不独 +不管 +不能 +不要 +不论 +不足 +不过 +不问 +与 +与其 +与否 +与此同时 +专门 +且 +两者 +严格 +严重 +个 +个人 +个别 +中小 +中间 +丰富 +临 +为 +为主 +为了 +为什么 +为什麽 +为何 +为着 +主张 +主要 +举行 +乃 +乃至 +么 +之 +之一 +之前 +之后 +之後 +之所以 +之类 +乌乎 +乎 +乘 +也 +也好 +也是 +也罢 +了 +了解 +争取 +于 +于是 +于是乎 +云云 +互相 +产生 +人们 +人家 +什么 +什么样 +什麽 +今后 +今天 +今年 +今後 +仍然 +从 +从事 +从而 +他 +他人 +他们 +他的 +代替 +以 +以上 +以下 +以为 +以便 +以免 +以前 +以及 +以后 +以外 +以後 +以来 +以至 +以至于 +以致 +们 +任 +任何 +任凭 +任务 +企图 +伟大 +似乎 +似的 +但 +但是 +何 +何况 +何处 +何时 +作为 +你 +你们 +你的 +使得 +使用 +例如 +依 +依照 +依靠 +促进 +保持 +俺 +俺们 +倘 +倘使 +倘或 +倘然 +倘若 +假使 +假如 +假若 +做到 +像 +允许 +充分 +先后 +先後 +先生 +全部 +全面 +兮 +共同 +关于 +其 +其一 +其中 +其二 +其他 +其余 +其它 +其实 +其次 +具体 +具体地说 +具体说来 +具有 +再者 +再说 +冒 +冲 +决定 +况且 +准备 +几 +几乎 +几时 +凭 +凭借 +出去 +出来 +出现 +分别 +则 +别 +别的 +别说 +到 +前后 +前者 +前进 +前面 +加之 +加以 +加入 +加强 +十分 +即 +即令 +即使 +即便 +即或 +即若 +却不 +原来 +又 +及 +及其 +及时 +及至 +双方 +反之 +反应 +反映 +反过来 +反过来说 +取得 +受到 +变成 +另 +另一方面 +另外 +只是 +只有 +只要 +只限 +叫 +叫做 +召开 +叮咚 +可 +可以 +可是 +可能 +可见 +各 +各个 +各人 +各位 +各地 +各种 +各级 +各自 +合理 +同 +同一 +同时 +同样 +后来 +后面 +向 +向着 +吓 +吗 +否则 +吧 +吧哒 +吱 +呀 +呃 +呕 +呗 +呜 +呜呼 +呢 +周围 +呵 +呸 +呼哧 +咋 +和 +咚 +咦 +咱 +咱们 +咳 +哇 +哈 +哈哈 +哉 +哎 +哎呀 +哎哟 +哗 +哟 +哦 +哩 +哪 +哪个 +哪些 +哪儿 +哪天 +哪年 +哪怕 +哪样 +哪边 +哪里 +哼 +哼唷 +唉 +啊 +啐 +啥 +啦 +啪达 +喂 +喏 +喔唷 +嗡嗡 +嗬 +嗯 +嗳 +嘎 +嘎登 +嘘 +嘛 +嘻 +嘿 +因 +因为 +因此 +因而 +固然 +在 +在下 +地 +坚决 +坚持 +基本 +处理 +复杂 +多 +多少 +多数 +多次 +大力 +大多数 +大大 +大家 +大批 +大约 +大量 +失去 +她 +她们 +她的 +好的 +好象 +如 +如上所述 +如下 +如何 +如其 +如果 +如此 +如若 +存在 +宁 +宁可 +宁愿 +宁肯 +它 +它们 +它们的 +它的 +安全 +完全 +完成 +实现 +实际 +宣布 +容易 +密切 +对 +对于 +对应 +将 +少数 +尔后 +尚且 +尤其 +就 +就是 +就是说 +尽 +尽管 +属于 +岂但 +左右 +巨大 +巩固 +己 +已经 +帮助 +常常 +并 +并不 +并不是 +并且 +并没有 +广大 +广泛 +应当 +应用 +应该 +开外 +开始 +开展 +引起 +强烈 +强调 +归 +当 +当前 +当时 +当然 +当着 +形成 +彻底 +彼 +彼此 +往 +往往 +待 +後来 +後面 +得 +得出 +得到 +心里 +必然 +必要 +必须 +怎 +怎么 +怎么办 +怎么样 +怎样 +怎麽 +总之 +总是 +总的来看 +总的来说 +总的说来 +总结 +总而言之 +恰恰相反 +您 +意思 +愿意 +慢说 +成为 +我 +我们 +我的 +或 +或是 +或者 +战斗 +所 +所以 +所有 +所谓 +打 +扩大 +把 +抑或 +拿 +按 +按照 +换句话说 +换言之 +据 +掌握 +接着 +接著 +故 +故此 +整个 +方便 +方面 +旁人 +无宁 +无法 +无论 +既 +既是 +既然 +时候 +明显 +明确 +是 +是否 +是的 +显然 +显著 +普通 +普遍 +更加 +曾经 +替 +最后 +最大 +最好 +最後 +最近 +最高 +有 +有些 +有关 +有利 +有力 +有所 +有效 +有时 +有点 +有的 +有着 +有著 +望 +朝 +朝着 +本 +本着 +来 +来着 +极了 +构成 +果然 +果真 +某 +某个 +某些 +根据 +根本 +欢迎 +正在 +正如 +正常 +此 +此外 +此时 +此间 +毋宁 +每 +每个 +每天 +每年 +每当 +比 +比如 +比方 +比较 +毫不 +没有 +沿 +沿着 +注意 +深入 +清楚 +满足 +漫说 +焉 +然则 +然后 +然後 +然而 +照 +照着 +特别是 +特殊 +特点 +现代 +现在 +甚么 +甚而 +甚至 +用 +由 +由于 +由此可见 +的 +的话 +目前 +直到 +直接 +相似 +相信 +相反 +相同 +相对 +相对而言 +相应 +相当 +相等 +省得 +看出 +看到 +看来 +看看 +看见 +真是 +真正 +着 +着呢 +矣 +知道 +确定 +离 +积极 +移动 +突出 +突然 +立即 +第 +等 +等等 +管 +紧接着 +纵 +纵令 +纵使 +纵然 +练习 +组成 +经 +经常 +经过 +结合 +结果 +给 +绝对 +继续 +继而 +维持 +综上所述 +罢了 +考虑 +者 +而 +而且 +而况 +而外 +而已 +而是 +而言 +联系 +能 +能否 +能够 +腾 +自 +自个儿 +自从 +自各儿 +自家 +自己 +自身 +至 +至于 +良好 +若 +若是 +若非 +范围 +莫若 +获得 +虽 +虽则 +虽然 +虽说 +行为 +行动 +表明 +表示 +被 +要 +要不 +要不是 +要不然 +要么 +要是 +要求 +规定 +觉得 +认为 +认真 +认识 +让 +许多 +论 +设使 +设若 +该 +说明 +诸位 +谁 +谁知 +赶 +起 +起来 +起见 +趁 +趁着 +越是 +跟 +转动 +转变 +转贴 +较 +较之 +边 +达到 +迅速 +过 +过去 +过来 +运用 +还是 +还有 +这 +这个 +这么 +这么些 +这么样 +这么点儿 +这些 +这会儿 +这儿 +这就是说 +这时 +这样 +这点 +这种 +这边 +这里 +这麽 +进入 +进步 +进而 +进行 +连 +连同 +适应 +适当 +适用 +逐步 +逐渐 +通常 +通过 +造成 +遇到 +遭到 +避免 +那 +那个 +那么 +那么些 +那么样 +那些 +那会儿 +那儿 +那时 +那样 +那边 +那里 +那麽 +部分 +鄙人 +采取 +里面 +重大 +重新 +重要 +鉴于 +问题 +防止 +阿 +附近 +限制 +除 +除了 +除此之外 +除非 +随 +随着 +随著 +集中 +需要 +非但 +非常 +非徒 +靠 +顺 +顺着 +首先 +高兴 +是不是 +说说 diff --git a/reme4/config/__init__.py b/reme4/config/__init__.py new file mode 100644 index 00000000..b08107fe --- /dev/null +++ b/reme4/config/__init__.py @@ -0,0 +1,5 @@ +"""Config""" + +from .config_parser import parse_args + +__all__ = ["parse_args"] diff --git a/reme4/config/config_parser.py b/reme4/config/config_parser.py new file mode 100644 index 00000000..3410608c --- /dev/null +++ b/reme4/config/config_parser.py @@ -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 diff --git a/reme4/constants.py b/reme4/constants.py new file mode 100644 index 00000000..fce66e64 --- /dev/null +++ b/reme4/constants.py @@ -0,0 +1,7 @@ +"""Constants""" + +REME_SERVICE_INFO = "REME_SERVICE_INFO" + +REME_DEFAULT_HOST = "127.0.0.1" + +REME_DEFAULT_PORT = 2333 diff --git a/reme4/enumeration/__init__.py b/reme4/enumeration/__init__.py new file mode 100644 index 00000000..9887ddee --- /dev/null +++ b/reme4/enumeration/__init__.py @@ -0,0 +1,9 @@ +"""Enumeration""" + +from .chunk_enum import ChunkEnum +from .component_enum import ComponentEnum + +__all__ = [ + "ChunkEnum", + "ComponentEnum", +] diff --git a/reme4/enumeration/chunk_enum.py b/reme4/enumeration/chunk_enum.py new file mode 100644 index 00000000..d397be9e --- /dev/null +++ b/reme4/enumeration/chunk_enum.py @@ -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" diff --git a/reme4/enumeration/component_enum.py b/reme4/enumeration/component_enum.py new file mode 100644 index 00000000..c3e30a89 --- /dev/null +++ b/reme4/enumeration/component_enum.py @@ -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" diff --git a/reme4/reme.py b/reme4/reme.py new file mode 100644 index 00000000..437ef0cf --- /dev/null +++ b/reme4/reme.py @@ -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() diff --git a/reme4/schema/__init__.py b/reme4/schema/__init__.py new file mode 100644 index 00000000..0a54891b --- /dev/null +++ b/reme4/schema/__init__.py @@ -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", +] diff --git a/reme4/schema/application_config.py b/reme4/schema/application_config.py new file mode 100644 index 00000000..b71e6d21 --- /dev/null +++ b/reme4/schema/application_config.py @@ -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", + ) diff --git a/reme4/schema/emb_node.py b/reme4/schema/emb_node.py new file mode 100644 index 00000000..d4191a26 --- /dev/null +++ b/reme4/schema/emb_node.py @@ -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() diff --git a/reme4/schema/file_chunk.py b/reme4/schema/file_chunk.py new file mode 100644 index 00000000..f7de723b --- /dev/null +++ b/reme4/schema/file_chunk.py @@ -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 diff --git a/reme4/schema/file_front_matter.py b/reme4/schema/file_front_matter.py new file mode 100644 index 00000000..16e81537 --- /dev/null +++ b/reme4/schema/file_front_matter.py @@ -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__ diff --git a/reme4/schema/file_link.py b/reme4/schema/file_link.py new file mode 100644 index 00000000..3d0b436f --- /dev/null +++ b/reme4/schema/file_link.py @@ -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]].", + ) diff --git a/reme4/schema/file_node.py b/reme4/schema/file_node.py new file mode 100644 index 00000000..9be276f8 --- /dev/null +++ b/reme4/schema/file_node.py @@ -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") diff --git a/reme4/schema/request.py b/reme4/schema/request.py new file mode 100644 index 00000000..c1d9b2ab --- /dev/null +++ b/reme4/schema/request.py @@ -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") diff --git a/reme4/schema/response.py b/reme4/schema/response.py new file mode 100644 index 00000000..89cbce33 --- /dev/null +++ b/reme4/schema/response.py @@ -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") diff --git a/reme4/schema/stream_chunk.py b/reme4/schema/stream_chunk.py new file mode 100644 index 00000000..b2fb0aef --- /dev/null +++ b/reme4/schema/stream_chunk.py @@ -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") diff --git a/reme4/steps/__init__.py b/reme4/steps/__init__.py new file mode 100644 index 00000000..2fde5b11 --- /dev/null +++ b/reme4/steps/__init__.py @@ -0,0 +1,7 @@ +"""steps""" + +from .base_step import BaseStep + +__all__ = [ + "BaseStep", +] diff --git a/reme4/steps/base_step.py b/reme4/steps/base_step.py new file mode 100644 index 00000000..16df084f --- /dev/null +++ b/reme4/steps/base_step.py @@ -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}) diff --git a/reme4/utils/__init__.py b/reme4/utils/__init__.py new file mode 100644 index 00000000..cc43fe4e --- /dev/null +++ b/reme4/utils/__init__.py @@ -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", +] diff --git a/reme4/utils/common_utils.py b/reme4/utils/common_utils.py new file mode 100644 index 00000000..46d6f97e --- /dev/null +++ b/reme4/utils/common_utils.py @@ -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 diff --git a/reme4/utils/logger_utils.py b/reme4/utils/logger_utils.py new file mode 100644 index 00000000..53d63b67 --- /dev/null +++ b/reme4/utils/logger_utils.py @@ -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 diff --git a/reme4/utils/logo_utils.py b/reme4/utils/logo_utils.py new file mode 100644 index 00000000..2a70ab94 --- /dev/null +++ b/reme4/utils/logo_utils.py @@ -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")) diff --git a/reme4/utils/similarity_utils.py b/reme4/utils/similarity_utils.py new file mode 100644 index 00000000..ec0f7100 --- /dev/null +++ b/reme4/utils/similarity_utils.py @@ -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