neorag 0.2.0__py3-none-any.whl

neorag/__init__.py

@@ -0,0 +1,7 @@
+"""NeoRAG: a self-hosted, fine-tunable RAG agent."""
+
+from neorag.config import load_config
+
+__version__ = "0.1.0"
+
+__all__ = ["load_config", "__version__"]

neorag/agent/direct.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from neorag.data.qa_dataset import RAG_QA_SYSTEM_PROMPT
+from neorag.retrieval.passage import Passage, format_context
+
+
+def direct_answer(llm, retriever, question: str) -> tuple[str, list[Passage]]:
+    passages = retriever.retrieve(question)
+    if not passages:
+        return "I could not find the answer in the provided documents.", []
+    context, _ = format_context(passages)
+    messages = [
+        {"role": "system", "content": RAG_QA_SYSTEM_PROMPT},
+        {"role": "user", "content": f"CONTEXT:\n{context}\n\nQUESTION: {question}"},
+    ]
+    answer = llm.generate(messages)
+    return answer, passages

neorag/agent/fast.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor
+
+from neorag.data.qa_dataset import RAG_QA_SYSTEM_PROMPT
+from neorag.retrieval.passage import Passage
+
+
+def _build_context(passages: list[Passage], max_chars: int) -> tuple[str, int]:
+    lines = []
+    used = 0
+    count = 0
+    for i, p in enumerate(passages, start=1):
+        header = f"[{i}] (source: {p.metadata.get('source', 'unknown')})"
+        block = f"{header}\n{p.text}"
+        if used + len(block) > max_chars:
+            break
+        lines.append(block)
+        used += len(block)
+        count = i
+    return "\n\n".join(lines), count
+
+
+def fast_answer(
+    llm,
+    retriever,
+    web_searcher,
+    question: str,
+    max_context_chars: int = 12000,
+    on_phase=None,
+) -> tuple[str, list[Passage]]:
+    phases = on_phase or (lambda s: None)
+
+    def fetch_web() -> list[dict]:
+        if web_searcher is None:
+            return []
+        try:
+            return web_searcher.fetch_pages(question)
+        except Exception:
+            return []
+
+    phases("searching local docs and the web")
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        local_future = pool.submit(retriever.retrieve, question)
+        web_future = pool.submit(fetch_web)
+        local_passages = list(local_future.result() or [])
+        web_pages = web_future.result() or []
+
+    phases(f"reading {len(local_passages) + len(web_pages)} sources")
+    web_passages = (
+        web_searcher.rank_pages_to_passages(question, web_pages) if web_searcher else []
+    )
+    passages = local_passages + web_passages
+    if not passages:
+        return "I could not find the answer in the provided documents or on the web.", []
+
+    context, count = _build_context(passages, max_context_chars)
+    passages = passages[:count]
+
+    phases("generating answer")
+    messages = [
+        {"role": "system", "content": RAG_QA_SYSTEM_PROMPT},
+        {"role": "user", "content": f"CONTEXT:\n{context}\n\nQUESTION: {question}"},
+    ]
+    answer = llm.generate(messages)
+    return answer, passages

neorag/agent/prompts.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from neorag.agent.tools import Tool
+
+_INTRO = """You are NeoRAG, a retrieval-grounded answering agent. You answer only from information returned by tools, never from imagination. The tools may search a private document collection, the public web, or both."""
+
+_ACTION_FORMAT = """Each turn you must output exactly ONE JSON object and nothing else. No markdown, no extra text. Format:
+
+{"thought": "<your reasoning>", "action": "<tool name or answer/give_up>", "action_input": {<tool-specific input>}}"""
+
+_RULES = """Rules:
+1. After each tool call you will receive an OBSERVATION containing numbered passages like [1], [2].
+2. In the final "answer", cite the numbers of the passages you used, e.g. "citations": [1, 3].
+3. Never invent facts. If observations do not contain the answer, try another tool or rephrase (up to the step limit) or choose "give_up".
+4. Keep queries short and keyword-like; rewrite failed queries using synonyms."""
+
+_TOOL_DOC_TEMPLATE = 'Tool "{name}": {description}\nInput example: {input_example}'
+
+
+def _default_input_example(tool: Any) -> str:
+    if "query" in tool.description.lower():
+        return '{"query": "<search words>"}'
+    return '{"value": "<input>"}'
+
+
+def build_system_prompt(tools: list[Tool]) -> str:
+    tool_docs = [
+        _TOOL_DOC_TEMPLATE.format(
+            name=t.name, description=t.description, input_example=_default_input_example(t)
+        )
+        for t in tools
+    ]
+    tool_section = "\n\n".join(tool_docs) if tool_docs else "No tools are available; answer or give_up directly."
+    return f"{_INTRO}\n\n{_ACTION_FORMAT}\n\n{_RULES}\n\nAvailable tools:\n{tool_section}\n"
+
+
+SYSTEM_PROMPT = (
+    f"{_INTRO}\n\n{_ACTION_FORMAT}\n\n{_RULES}\n\n"
+    'Tools:\nTool "search": Search the private document collection. Input example: {"query": "<search words>"}\n'
+)
+
+OBSERVATION_HEADER = "OBSERVATION (numbered passages retrieved for your query)"
+
+LAST_CHANCE_PROMPT = """Based only on the OBSERVATIONS above, now respond with exactly ONE JSON object using the "answer" action:
+
+{"thought": "<reasoning>", "action": "answer", "action_input": {"answer": "<final answer grounded in the passages>", "citations": [<numbers>]}}
+
+If the passages truly do not contain the answer, use the "give_up" action instead."""
+
+NO_RESULTS_OBSERVATION = "No passages matched this query. Try different or simpler keywords."

neorag/agent/react.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass, field
+
+from neorag.agent.prompts import (
+    LAST_CHANCE_PROMPT,
+    NO_RESULTS_OBSERVATION,
+    OBSERVATION_HEADER,
+    build_system_prompt,
+)
+from neorag.agent.tools import Tool
+from neorag.logging_utils import logger
+from neorag.retrieval.passage import Passage
+
+
+@dataclass
+class AgentResult:
+    answer: str
+    citations: list[int] = field(default_factory=list)
+    passages: dict[int, Passage] = field(default_factory=dict)
+    steps_used: int = 0
+    gave_up: bool = False
+    trace: list[dict] = field(default_factory=list)
+
+
+def extract_json_object(text: str) -> dict | None:
+    text = re.sub(r"```(?:json)?", "", text).strip()
+    start = text.find("{")
+    while start != -1:
+        depth = 0
+        in_string = False
+        escape = False
+        for i in range(start, len(text)):
+            ch = text[i]
+            if in_string:
+                if escape:
+                    escape = False
+                elif ch == "\\":
+                    escape = True
+                elif ch == '"':
+                    in_string = False
+                continue
+            if ch == '"':
+                in_string = True
+            elif ch == "{":
+                depth += 1
+            elif ch == "}":
+                depth -= 1
+                if depth == 0:
+                    candidate = text[start : i + 1]
+                    try:
+                        obj = json.loads(candidate)
+                        return obj if isinstance(obj, dict) else None
+                    except json.JSONDecodeError:
+                        break
+        start = text.find("{", start + 1)
+    return None
+
+
+def parse_action(text: str) -> dict:
+    obj = extract_json_object(text)
+    if not obj:
+        return {"thought": "", "action": "invalid", "action_input": {}}
+    action = str(obj.get("action", "invalid")).strip().lower()
+    action_input = obj.get("action_input", obj.get("action_input:", {}))
+    if not isinstance(action_input, dict):
+        action_input = {"value": action_input}
+    thought = str(obj.get("thought", ""))
+    return {"thought": thought, "action": action, "action_input": action_input}
+
+
+def _format_observation(fresh: list[tuple[int, Passage]]) -> str:
+    if not fresh:
+        return NO_RESULTS_OBSERVATION
+    lines = [
+        f"[{n}] (source: {p.metadata.get('source', 'unknown')}, score: {p.score:.3f})\n{p.text}"
+        for n, p in fresh
+    ]
+    return "\n\n".join(lines)
+
+
+class RAGAgent:
+    def __init__(self, llm, tools: list[Tool], max_steps: int = 5, system_prompt: str | None = None):
+        self.llm = llm
+        self.tools = {t.name: t for t in tools}
+        self.max_steps = max_steps
+        self.system_prompt = system_prompt
+
+    def run(self, question: str, on_step=None) -> AgentResult:
+        system_content = self.system_prompt or build_system_prompt(list(self.tools.values()))
+        messages: list[dict[str, str]] = [
+            {"role": "system", "content": system_content},
+            {"role": "user", "content": question},
+        ]
+        seen_ids: set[str] = set()
+        passage_registry: dict[int, Passage] = {}
+        next_number = 1
+        trace: list[dict] = []
+        gave_up = False
+        answer: str | None = None
+        citations: list[int] = []
+        tool_names = sorted(self.tools.keys())
+        consecutive_failures = 0
+
+        for step in range(1, self.max_steps + 1):
+            raw = self.llm.generate(messages)
+            parsed = parse_action(raw)
+            action = parsed["action"]
+            action_input = parsed["action_input"]
+            trace.append({"step": step, "thought": parsed["thought"], "action": action, "raw": raw})
+            if on_step is not None:
+                try:
+                    on_step({"step": step, "action": action, "thought": parsed["thought"]})
+                except Exception:
+                    pass
+            logger.info(f"[agent step {step}] action={action} thought={parsed['thought'][:120]}")
+
+            if action == "answer":
+                candidate = str(action_input.get("answer", "")).strip()
+                if candidate:
+                    answer = candidate
+                    citations = self._validate_citations(action_input.get("citations"), passage_registry)
+                    break
+                consecutive_failures += 1
+                if consecutive_failures >= 2:
+                    break
+
+            if action == "give_up":
+                gave_up = True
+                answer = str(action_input.get("answer", "")).strip() or (
+                    "I could not find the answer in the provided documents."
+                )
+                break
+
+            tool = self.tools.get(action)
+            if tool is not None:
+                result = tool.fn(action_input)
+                consecutive_failures = 0
+                fresh = [p for p in result.passages if p.id not in seen_ids]
+                for p in fresh:
+                    seen_ids.add(p.id)
+                numbered = []
+                for p in fresh:
+                    passage_registry[next_number] = p
+                    numbered.append((next_number, p))
+                    next_number += 1
+                observation = _format_observation(numbered)
+                messages.append({"role": "assistant", "content": raw})
+                messages.append({"role": "user", "content": f"{OBSERVATION_HEADER}:\n{observation}"})
+                continue
+
+            messages.append({"role": "assistant", "content": raw})
+            consecutive_failures += 1
+            if consecutive_failures >= 2:
+                break
+            messages.append(
+                {
+                    "role": "user",
+                    "content": (
+                        "That was not a valid action. Respond with exactly one JSON object using "
+                        f'one of these actions: {", ".join(tool_names + ["answer", "give_up"])}. '
+                        'Example: {"thought": "...", "action": "...", "action_input": {...}}'
+                    ),
+                }
+            )
+
+        if answer is None:
+            logger.info("Step limit reached; forcing a final grounded answer.")
+            messages.append({"role": "user", "content": LAST_CHANCE_PROMPT})
+            raw = self.llm.generate(messages)
+            parsed = parse_action(raw)
+            answer = str(parsed["action_input"].get("answer", "")).strip()
+            citations = self._validate_citations(
+                parsed["action_input"].get("citations"), passage_registry
+            )
+            if parsed["action"] == "give_up":
+                gave_up = True
+            if not answer:
+                gave_up = True
+                answer = "I could not find the answer in the provided documents."
+
+        return AgentResult(
+            answer=answer,
+            citations=citations,
+            passages=passage_registry,
+            steps_used=len(trace),
+            gave_up=gave_up,
+            trace=trace,
+        )
+
+    def _validate_citations(
+        self, raw_citations: object, registry: dict[int, Passage]
+    ) -> list[int]:
+        if isinstance(raw_citations, (int, float)):
+            raw_citations = [raw_citations]
+        if not isinstance(raw_citations, list):
+            return []
+        valid = []
+        for c in raw_citations:
+            try:
+                n = int(c)
+            except (TypeError, ValueError):
+                continue
+            if n in registry:
+                valid.append(n)
+        return sorted(set(valid))

neorag/agent/tools.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable
+
+from neorag.agent.prompts import NO_RESULTS_OBSERVATION
+from neorag.retrieval.passage import Passage, format_context
+from neorag.retrieval.retriever import Retriever
+
+
+@dataclass
+class ToolResult:
+    observation: str
+    passages: list[Passage]
+
+
+@dataclass
+class Tool:
+    name: str
+    description: str
+    fn: Callable[[dict], ToolResult]
+
+
+class SearchDocumentsTool:
+    def __init__(self, retriever: Retriever, top_k: int | None = None):
+        self.retriever = retriever
+        self.top_k = top_k
+
+    def __call__(self, action_input: dict) -> ToolResult:
+        query = str(action_input.get("query", "")).strip()
+        if not query:
+            return ToolResult(observation=NO_RESULTS_OBSERVATION, passages=[])
+        retriever = self.retriever
+        if self.top_k:
+            original_final = retriever.final_k
+            retriever.final_k = self.top_k
+            try:
+                passages = retriever.retrieve(query)
+            finally:
+                retriever.final_k = original_final
+        else:
+            passages = retriever.retrieve(query)
+        if not passages:
+            return ToolResult(observation=NO_RESULTS_OBSERVATION, passages=[])
+        text, _ = format_context(passages)
+        return ToolResult(observation=text, passages=passages)
+
+
+def build_web_tool(web_searcher, top_k: int | None = None) -> Tool:
+    from neorag.web.websearch import WebSearchTool
+
+    searcher = WebSearchTool(web_searcher, top_k=top_k)
+    return Tool(
+        name="web_search",
+        description=(
+            "Search the public web for fresh information. "
+            'Input: {"query": "<search words>"}. Returns numbered passages with URLs as sources.'
+        ),
+        fn=searcher,
+    )
+
+
+def default_tools(
+    retriever: Retriever,
+    web_searcher=None,
+    search_top_k: int | None = None,
+    web_top_k: int | None = None,
+) -> list[Tool]:
+    search = Tool(
+        name="search",
+        description=(
+            "Search the private document collection. "
+            'Input: {"query": "<search words>"}. Returns numbered passages.'
+        ),
+        fn=SearchDocumentsTool(retriever, top_k=search_top_k),
+    )
+    tools = [search]
+    if web_searcher is not None:
+        tools.append(build_web_tool(web_searcher, top_k=web_top_k))
+    return tools