118 lines
3.8 KiB
Python
118 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
from pathlib import Path
|
|
|
|
from rag_agent.config import load_config
|
|
from rag_agent.ingest.chunker import chunk_text
|
|
from rag_agent.ingest.file_loader import iter_text_files
|
|
from rag_agent.ingest.git_watcher import filter_existing, filter_removed, get_changed_files
|
|
from rag_agent.index.embeddings import get_embedding_client
|
|
from rag_agent.index.postgres import (
|
|
connect,
|
|
delete_document,
|
|
ensure_schema,
|
|
get_or_create_story,
|
|
get_story_id,
|
|
replace_chunks,
|
|
upsert_document,
|
|
)
|
|
from rag_agent.agent.pipeline import StubLLMClient, answer_query
|
|
|
|
|
|
def _file_version(path: Path) -> str:
|
|
stat = path.stat()
|
|
payload = f"{path.as_posix()}:{stat.st_mtime_ns}:{stat.st_size}"
|
|
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def cmd_index(args: argparse.Namespace) -> None:
|
|
config = load_config()
|
|
conn = connect(config.db_dsn)
|
|
ensure_schema(conn, config.embeddings_dim)
|
|
story_id = get_or_create_story(conn, args.story)
|
|
embedding_client = get_embedding_client(config.embeddings_dim)
|
|
|
|
if args.changed:
|
|
changed_files = get_changed_files(config.repo_path, args.base_ref, args.head_ref)
|
|
removed = filter_removed(changed_files)
|
|
existing = filter_existing(changed_files)
|
|
else:
|
|
removed = []
|
|
existing = [path for path in Path(config.repo_path).rglob("*") if path.is_file()]
|
|
|
|
for path in removed:
|
|
delete_document(conn, story_id, str(path))
|
|
|
|
for path, text in iter_text_files(existing, config.allowed_extensions):
|
|
chunks = chunk_text(text, config.chunk_size, config.chunk_overlap)
|
|
if not chunks:
|
|
continue
|
|
embeddings = embedding_client.embed_texts([chunk.text for chunk in chunks])
|
|
document_id = upsert_document(
|
|
conn, story_id, str(path), _file_version(path)
|
|
)
|
|
replace_chunks(conn, document_id, chunks, embeddings)
|
|
|
|
|
|
def cmd_ask(args: argparse.Namespace) -> None:
|
|
config = load_config()
|
|
conn = connect(config.db_dsn)
|
|
ensure_schema(conn, config.embeddings_dim)
|
|
story_id = None
|
|
if args.story:
|
|
story_id = get_story_id(conn, args.story)
|
|
if story_id is None:
|
|
raise SystemExit(f"Story not found: {args.story}")
|
|
embedding_client = get_embedding_client(config.embeddings_dim)
|
|
llm_client = StubLLMClient()
|
|
answer = answer_query(
|
|
conn,
|
|
config,
|
|
embedding_client,
|
|
llm_client,
|
|
args.question,
|
|
top_k=args.top_k,
|
|
story_id=story_id,
|
|
)
|
|
print(answer)
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(prog="rag-agent")
|
|
sub = parser.add_subparsers(dest="command", required=True)
|
|
|
|
index_parser = sub.add_parser("index", help="Index files into Postgres")
|
|
index_parser.add_argument(
|
|
"--story",
|
|
required=True,
|
|
help="Story slug (e.g. branch name or story id); documents are tied to this story",
|
|
)
|
|
index_parser.add_argument("--changed", action="store_true", help="Index only changed files")
|
|
index_parser.add_argument("--base-ref", default="HEAD~1", help="Base git ref")
|
|
index_parser.add_argument("--head-ref", default="HEAD", help="Head git ref")
|
|
index_parser.set_defaults(func=cmd_index)
|
|
|
|
ask_parser = sub.add_parser("ask", help="Ask a question")
|
|
ask_parser.add_argument("question", help="Question text")
|
|
ask_parser.add_argument(
|
|
"--story",
|
|
default=None,
|
|
help="Limit retrieval to this story slug (optional)",
|
|
)
|
|
ask_parser.add_argument("--top-k", type=int, default=5, help="Top K chunks to retrieve")
|
|
ask_parser.set_defaults(func=cmd_ask)
|
|
|
|
return parser
|
|
|
|
|
|
def main() -> None:
|
|
parser = build_parser()
|
|
args = parser.parse_args()
|
|
args.func(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|