Spaces:
Sleeping
Sleeping
| """Grounded answer generation with numbered citations, streamed token by token.""" | |
| import json | |
| import re | |
| from collections.abc import Iterator | |
| from typing import Protocol | |
| import httpx | |
| from .config import REFUSAL, Settings | |
| from .retrieve import Hit | |
| SYSTEM = "You are a careful assistant that answers questions about a company's documents, using only the sources given." | |
| # Small models follow instructions best when they come right before the answer, with a worked example. | |
| INSTRUCTIONS = ( | |
| "Answer the question in one or two plain sentences, using only the sources above. " | |
| "State the actual fact (numbers, names, conditions), then put the number of the source that supports it " | |
| 'in square brackets, for example: "Annual plans can be refunded in full within 30 days [2]." ' | |
| 'Do not add phrases like "according to the sources" or "this is supported by source 1". ' | |
| f'If the sources do not answer the question, reply exactly: "{REFUSAL}"' | |
| ) | |
| def build_messages(question: str, hits: list[Hit]) -> list[dict]: | |
| context = "\n\n".join(f"[{i}] {h.chunk.text}" for i, h in enumerate(hits, start=1)) | |
| return [ | |
| {"role": "system", "content": SYSTEM}, | |
| {"role": "user", "content": f"Sources:\n\n{context}\n\nQuestion: {question}\n\n{INSTRUCTIONS}"}, | |
| ] | |
| def tidy(answer: str) -> str: | |
| """Collapse citation runs ("[1] [1] [2]" → "[1][2]") and cut a degenerate tail of bare citations.""" | |
| answer = re.sub(r"(\[\d+\])(?:\s*\1)+", r"\1", answer) | |
| answer = re.sub( | |
| r"((?:\s*\[\d+\]){3,})\s*$", lambda m: "".join(dict.fromkeys(re.findall(r"\[\d+\]", m.group(1)))), answer | |
| ) | |
| return answer.strip() | |
| def citations(answer: str, n_sources: int) -> list[int]: | |
| """Source numbers cited in the answer, in order of first use; out-of-range numbers are ignored.""" | |
| seen: list[int] = [] | |
| for m in re.finditer(r"\[(\d+(?:\s*,\s*\d+)*)\]", answer): | |
| for num in m.group(1).split(","): | |
| n = int(num) | |
| if 1 <= n <= n_sources and n not in seen: | |
| seen.append(n) | |
| return seen | |
| def strip_think(tokens: Iterator[str]) -> Iterator[str]: | |
| """Drop a leading <think>…</think> block (reasoning models) from a token stream.""" | |
| buf, inside = "", None | |
| for tok in tokens: | |
| if inside is False: | |
| yield tok | |
| continue | |
| buf += tok | |
| if inside is None: | |
| if not "<think>".startswith(buf.lstrip()[:7]): | |
| inside = False | |
| yield buf | |
| continue | |
| if buf.lstrip().startswith("<think>"): | |
| inside = True | |
| if inside and "</think>" in buf: | |
| inside, rest = False, buf.split("</think>", 1)[1].lstrip() | |
| if rest: | |
| yield rest | |
| class LLM(Protocol): | |
| def stream(self, messages: list[dict]) -> Iterator[str]: ... | |
| class LlamaCppLLM: | |
| """Local GGUF model on CPU via llama.cpp (downloaded from the Hugging Face Hub on first use).""" | |
| def __init__(self, s: Settings) -> None: | |
| from llama_cpp import Llama | |
| self.s = s | |
| # n_threads_batch (prompt reading) defaults to every CPU the OS reports; in a container limited to | |
| # 2 cores that oversubscribes badly, so both phases use the same thread count. | |
| self.llm = Llama.from_pretrained( | |
| repo_id=s.llm_repo, | |
| filename=s.llm_file, | |
| n_ctx=s.llm_ctx, | |
| n_threads=s.llm_threads, | |
| n_threads_batch=s.llm_threads, | |
| verbose=False, | |
| ) | |
| def stream(self, messages: list[dict]) -> Iterator[str]: | |
| if "qwen3" in self.s.llm_file.lower(): # answer directly, skip Qwen3's thinking mode | |
| messages = [*messages[:-1], {**messages[-1], "content": messages[-1]["content"] + " /no_think"}] | |
| yield from strip_think(self._raw(messages)) | |
| def _raw(self, messages: list[dict]) -> Iterator[str]: | |
| for part in self.llm.create_chat_completion( | |
| messages=messages, | |
| stream=True, | |
| max_tokens=self.s.max_new_tokens, | |
| temperature=self.s.temperature, | |
| repeat_penalty=1.15, | |
| stop=["\n\n\n", "Question:", "Sources:"], | |
| ): | |
| delta = part["choices"][0]["delta"].get("content") | |
| if delta: | |
| yield delta | |
| class OpenAICompatibleLLM: | |
| """Any OpenAI-compatible chat endpoint (OpenAI, vLLM, Ollama, LM Studio…).""" | |
| def __init__(self, s: Settings) -> None: | |
| self.s = s | |
| def stream(self, messages: list[dict]) -> Iterator[str]: | |
| body = { | |
| "model": self.s.openai_model, | |
| "messages": messages, | |
| "stream": True, | |
| "max_tokens": self.s.max_new_tokens, | |
| "temperature": self.s.temperature, | |
| } | |
| headers = {"Authorization": f"Bearer {self.s.openai_api_key or ''}"} | |
| with httpx.stream( | |
| "POST", f"{self.s.openai_base_url}/chat/completions", json=body, headers=headers, timeout=60 | |
| ) as r: | |
| r.raise_for_status() | |
| for line in r.iter_lines(): | |
| if line.startswith("data: ") and line != "data: [DONE]": | |
| delta = json.loads(line[6:])["choices"][0]["delta"].get("content") | |
| if delta: | |
| yield delta | |
| class AnthropicLLM: | |
| def __init__(self, s: Settings) -> None: | |
| self.s = s | |
| def stream(self, messages: list[dict]) -> Iterator[str]: | |
| system = messages[0]["content"] | |
| body = { | |
| "model": self.s.anthropic_model, | |
| "system": system, | |
| "messages": messages[1:], | |
| "stream": True, | |
| "max_tokens": self.s.max_new_tokens, | |
| "temperature": self.s.temperature, | |
| } | |
| headers = {"x-api-key": self.s.anthropic_api_key or "", "anthropic-version": "2023-06-01"} | |
| with httpx.stream("POST", "https://api.anthropic.com/v1/messages", json=body, headers=headers, timeout=60) as r: | |
| r.raise_for_status() | |
| for line in r.iter_lines(): | |
| if line.startswith("data: "): | |
| event = json.loads(line[6:]) | |
| if event.get("type") == "content_block_delta": | |
| yield event["delta"].get("text", "") | |
| class ExtractiveLLM: | |
| """No model at all: quotes the best chunk. Used for tests and when RAG_LLM_BACKEND=none.""" | |
| def stream(self, messages: list[dict]) -> Iterator[str]: | |
| m = re.search(r"\[1\] (.*?)(?:\n\n\[2\]|\n\nQuestion:)", messages[1]["content"], re.S) | |
| body = m.group(1).strip() if m else "" | |
| yield f"{body.splitlines()[-1][:400]} [1]" if body else REFUSAL | |
| def make_llm(s: Settings) -> LLM | None: | |
| if s.llm_backend == "llamacpp": | |
| return LlamaCppLLM(s) | |
| if s.llm_backend == "openai": | |
| return OpenAICompatibleLLM(s) | |
| if s.llm_backend == "anthropic": | |
| return AnthropicLLM(s) | |
| return ExtractiveLLM() | |