mirror of
https://github.com/hansjone/oclaw.git
synced 2026-10-09 01:50:44 +08:00
101 lines
3 KiB
Python
101 lines
3 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from oclaw.platform.embeddings.embedding_client import build_default_embedding_client
|
|
from oclaw.platform.persistence.sqlite_store import SqliteStore
|
|
from oclaw.platform.config.paths import PROJECT_ROOT
|
|
|
|
|
|
def _default_workspace_root() -> Path:
|
|
import os
|
|
|
|
override = (os.getenv("AIA_WORKSPACE_ROOT") or "").strip()
|
|
if override:
|
|
return Path(override).expanduser().resolve()
|
|
return Path(PROJECT_ROOT).resolve()
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class IndexStats:
|
|
files_seen: int
|
|
chunks_upserted: int
|
|
embeddings_upserted: int
|
|
|
|
|
|
def _chunk_id(source: str, text: str) -> str:
|
|
raw = f"{source}\n{text}".encode("utf-8", errors="ignore")
|
|
return hashlib.sha1(raw).hexdigest()
|
|
|
|
|
|
def _iter_files(root: Path, patterns: list[str], *, max_files: int) -> list[Path]:
|
|
files: list[Path] = []
|
|
for pat in patterns:
|
|
for p in root.glob(pat):
|
|
if p.is_dir():
|
|
continue
|
|
files.append(p)
|
|
if len(files) >= max_files:
|
|
return files
|
|
# de-dup by path
|
|
uniq = []
|
|
seen = set()
|
|
for p in files:
|
|
rp = str(p)
|
|
if rp in seen:
|
|
continue
|
|
seen.add(rp)
|
|
uniq.append(p)
|
|
return uniq[:max_files]
|
|
|
|
|
|
def index_workspace(
|
|
store: SqliteStore,
|
|
*,
|
|
root: str | None = None,
|
|
patterns: list[str] | None = None,
|
|
max_files: int = 200,
|
|
max_chars_per_chunk: int = 1600,
|
|
) -> IndexStats:
|
|
r = Path(root).resolve() if root else _default_workspace_root()
|
|
pats = patterns or ["**/*.py", "README*.md", "**/*.md"]
|
|
files = _iter_files(r, pats, max_files=max(1, int(max_files)))
|
|
client = build_default_embedding_client()
|
|
chunks_upserted = 0
|
|
embeds_upserted = 0
|
|
|
|
for p in files:
|
|
try:
|
|
text = p.read_text(encoding="utf-8", errors="replace")
|
|
except Exception:
|
|
continue
|
|
rel = str(p.relative_to(r))
|
|
# naive chunking by chars (good enough for MVP)
|
|
for idx in range(0, len(text), max(200, int(max_chars_per_chunk))):
|
|
chunk = text[idx : idx + int(max_chars_per_chunk)].strip()
|
|
if not chunk:
|
|
continue
|
|
source = f"workspace:{rel}#c{idx//max_chars_per_chunk:04d}"
|
|
cid = _chunk_id(source, chunk)
|
|
store.upsert_knowledge_chunk(
|
|
chunk_id=cid,
|
|
source=source,
|
|
content=chunk,
|
|
metadata={"source": source, "path": rel, "offset": idx},
|
|
)
|
|
chunks_upserted += 1
|
|
try:
|
|
emb = client.embed(chunk[:8000])
|
|
store.upsert_knowledge_embedding(chunk_id=cid, model=emb.model, vector=emb.vector)
|
|
embeds_upserted += 1
|
|
except Exception:
|
|
continue
|
|
|
|
return IndexStats(files_seen=len(files), chunks_upserted=chunks_upserted, embeddings_upserted=embeds_upserted)
|
|
|
|
|
|
__all__ = ["IndexStats", "index_workspace"]
|
|
|