forked from fangtianchen/algonotes_rag
64 lines
1.8 KiB
Python
64 lines
1.8 KiB
Python
# scripts/rag.py
|
|
# CLI: single-shot RAG query (retrieve → rerank → generate).
|
|
|
|
import json
|
|
from dataclasses import dataclass, asdict
|
|
from typing import Any
|
|
|
|
from src.performance import PerfTimer
|
|
|
|
|
|
@dataclass
|
|
class AskResult:
|
|
question: str
|
|
answer: str
|
|
success: bool
|
|
error: str | None = None
|
|
|
|
@property
|
|
def dict(self) -> dict[str, Any]:
|
|
return asdict(self)
|
|
|
|
|
|
def ask_question(question: str, stream: bool = False) -> AskResult:
|
|
"""Ask a single RAG question and return the answer.
|
|
|
|
Uses ``create_rag_agent()`` which invokes the full pipeline:
|
|
retrieve similar notes → rerank → generate answer with citations.
|
|
|
|
Args:
|
|
question: The question to ask.
|
|
stream: If True, print answer incrementally to stdout.
|
|
|
|
Returns:
|
|
An AskResult with the full answer text.
|
|
"""
|
|
from src.rag.agent import create_rag_agent
|
|
|
|
agent = create_rag_agent()
|
|
cfg = {"configurable": {"thread_id": "ask"}}
|
|
|
|
with PerfTimer("pipeline") as t:
|
|
try:
|
|
agent_steps = 0
|
|
full_response = ""
|
|
for event in agent.stream(
|
|
{"messages": [("user", question)]},
|
|
cfg,
|
|
):
|
|
agent_steps += 1
|
|
if "model" in event:
|
|
content = event["model"]["messages"][-1].content
|
|
if content:
|
|
if stream:
|
|
print(content, flush=True, end="")
|
|
full_response += content
|
|
|
|
if stream:
|
|
print()
|
|
|
|
t.set_extra("agent_steps", agent_steps)
|
|
return AskResult(question=question, answer=full_response, success=True)
|
|
except Exception as e:
|
|
return AskResult(question=question, answer="", success=False, error=str(e))
|