mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 17:00:23 +00:00
93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from paperforge.memory.db import get_connection, get_memory_db_path
|
|
|
|
|
|
def _build_collection_tree(conn) -> list[dict]:
|
|
"""Build collection hierarchy from papers.collection_path.
|
|
|
|
Each collection_path is pipe-separated, e.g. "骨科 | 骨折".
|
|
Returns flat list of top-level collections with sub-collections.
|
|
"""
|
|
rows = conn.execute(
|
|
"SELECT collection_path, COUNT(*) as cnt FROM papers "
|
|
"WHERE collection_path != '' "
|
|
"GROUP BY collection_path ORDER BY cnt DESC"
|
|
).fetchall()
|
|
top: dict[str, dict] = {}
|
|
for row in rows:
|
|
parts = [p.strip() for p in row["collection_path"].split("|") if p.strip()]
|
|
if not parts:
|
|
continue
|
|
root = parts[0]
|
|
if root not in top:
|
|
top[root] = {"name": root, "count": 0, "sub": []}
|
|
top[root]["count"] += row["cnt"]
|
|
if len(parts) > 1:
|
|
sub_name = parts[-1]
|
|
if sub_name not in top[root]["sub"]:
|
|
top[root]["sub"].append(sub_name)
|
|
for c in top.values():
|
|
c["sub"] = sorted(c["sub"])
|
|
return sorted(top.values(), key=lambda x: -x["count"])
|
|
|
|
|
|
def get_agent_context(vault: Path) -> dict | None:
|
|
"""Build agent context from paperforge.db — library stats + collection tree.
|
|
|
|
Returns None if DB is missing or query fails.
|
|
"""
|
|
db_path = get_memory_db_path(vault)
|
|
if not db_path.exists():
|
|
return None
|
|
|
|
conn = get_connection(db_path, read_only=True)
|
|
try:
|
|
total = conn.execute("SELECT COUNT(*) FROM papers").fetchone()[0]
|
|
|
|
domains = {
|
|
r["domain"]: r["cnt"]
|
|
for r in conn.execute(
|
|
"SELECT domain, COUNT(*) as cnt FROM papers GROUP BY domain ORDER BY cnt DESC"
|
|
).fetchall()
|
|
}
|
|
|
|
lifecycle_counts = {
|
|
r["lifecycle"]: r["cnt"]
|
|
for r in conn.execute(
|
|
"SELECT lifecycle, COUNT(*) as cnt FROM papers GROUP BY lifecycle"
|
|
).fetchall()
|
|
}
|
|
|
|
ocr_counts = {
|
|
r["ocr_status"]: r["cnt"]
|
|
for r in conn.execute(
|
|
"SELECT ocr_status, COUNT(*) as cnt FROM papers GROUP BY ocr_status"
|
|
).fetchall()
|
|
}
|
|
|
|
deep_counts = {
|
|
r["deep_reading_status"]: r["cnt"]
|
|
for r in conn.execute(
|
|
"SELECT deep_reading_status, COUNT(*) as cnt FROM papers GROUP BY deep_reading_status"
|
|
).fetchall()
|
|
}
|
|
|
|
collections = _build_collection_tree(conn)
|
|
|
|
return {
|
|
"library": {
|
|
"paper_count": total,
|
|
"domain_counts": domains,
|
|
"lifecycle_counts": lifecycle_counts,
|
|
"ocr_counts": ocr_counts,
|
|
"deep_reading_counts": deep_counts,
|
|
},
|
|
"collections": collections,
|
|
}
|
|
except Exception:
|
|
return None
|
|
finally:
|
|
conn.close()
|