Coverage for src/asketmc_bot/rag_langchain.py: 0%
164 statements
« prev ^ index » next coverage.py v7.5.1, created at 2026-04-08 17:14 +0000
1#!/usr/bin/env python3.10
2"""
3rag_langchain.py — LangGraph integration for the Asketmc RAG Discord bot.
5Design goals:
6- No side effects at import time (no global graph/pipeline objects).
7- Dependency injection (no imports from `main` that create cycles).
8- PEP 8/257 compliant, typed, idempotent nodes.
9- Structured logging; deterministic behavior; clear guards and fallbacks.
10- Safe import behavior when optional deps (langgraph/langsmith) are not installed.
12Fixes:
13- Make lemma-based filtering non-destructive:
14 - If lemmas are empty -> skip filter and pass-through candidates.
15 - If filter returns empty -> fallback to pass-through candidates.
16 - If filter errors -> fallback to pass-through candidates.
17- Make context building resilient:
18 - If filtered_nodes empty -> fallback to reranked/retrieved.
19 - If build_context returns empty with lemmas -> retry with empty-lemmas.
20- Add node-level debug counters (retrieved/reranked/filtered/context length).
21"""
23from __future__ import annotations
25import logging
26import os
27from typing import Any, Awaitable, Callable, Dict, List, Set, TypedDict, cast
29from asketmc_bot import config as cfg
31log = logging.getLogger("asketmc.rag.langgraph")
32log.setLevel(logging.DEBUG if getattr(cfg, "DEBUG", False) else logging.INFO)
34# --- Optional dependency: langgraph -----------------------------------------
35# Keep module importable even when langgraph isn't installed (dev/CI environments).
36try:
37 # pylint: disable=import-error
38 from langgraph.graph import END, START, StateGraph # type: ignore[import-not-found]
39except Exception: # pragma: no cover
40 END = "__end__" # type: ignore[assignment]
41 START = "__start__" # type: ignore[assignment]
42 StateGraph = Any # type: ignore[misc,assignment]
45class RagState(TypedDict, total=False):
46 """Graph state passed between LangGraph nodes."""
47 question: str
48 lemmas: Set[str]
49 retrieved_nodes: List[Any]
50 reranked_nodes: List[Any]
51 filtered_nodes: List[Any]
52 context: str
53 answer: str
56class RagDeps(TypedDict):
57 """
58 Dependencies required by the RAG pipeline.
59 Provide concrete callables and objects to avoid circular imports.
60 """
61 extract_lemmas: Callable[[str], Set[str]]
62 retriever: Any # must expose `aretrieve(q: str) -> Awaitable[List[Any]]`
63 rerank: Callable[[str, List[Any]], Awaitable[List[Any]]]
64 get_filtered_nodes: Callable[[List[Any], Set[str]], Awaitable[List[Any]]]
65 build_context: Callable[[List[Any], Set[str], int], str]
67 # Must accept sys_prompt + ctx_txt + q (discord user question).
68 # `timeout_sec` may be supported by implementation.
69 query_model: Callable[..., Awaitable[tuple[str, bool]]]
71 sys_prompt: str
72 ctx_len_local: int
73 ctx_len_remote: int
76def _maybe_init_langsmith() -> None:
77 """
78 Initialize LangSmith tracing if enabled via env (LANGSMITH_TRACING=1/true/yes).
80 This function must not raise; tracing is always best-effort.
81 """
82 val = (os.getenv("LANGSMITH_TRACING") or "").strip().lower()
83 if val not in {"1", "true", "yes"}:
84 return
86 try:
87 # pylint: disable=import-error,import-outside-toplevel
88 import langsmith # type: ignore[import-not-found,unused-import] # noqa: F401
89 log.info("LangSmith tracing enabled (module present).")
90 except Exception as exc: # pragma: no cover
91 log.warning("LangSmith tracing requested but not available: %s", exc)
94def _require_non_empty_question(state: RagState, *, node_name: str) -> str:
95 question = (state.get("question") or "").strip()
96 if not question:
97 raise ValueError(f"Empty question for {node_name} node.")
98 return question
101def _lemmas_node(deps: RagDeps) -> Callable[[RagState], Dict[str, Any]]:
102 def _impl(state: RagState) -> Dict[str, Any]:
103 question = _require_non_empty_question(state, node_name="lemmas")
104 try:
105 lemmas = deps["extract_lemmas"](question)
106 except Exception as exc:
107 log.warning("[lemmas] extract_lemmas failed: %s", exc)
108 lemmas = set()
109 out = set(lemmas) if lemmas else set()
110 log.debug("[lemmas] size=%d", len(out))
111 return {"lemmas": out}
113 return _impl
116def _retrieve_node(deps: RagDeps) -> Callable[[RagState], Awaitable[Dict[str, Any]]]:
117 async def _impl(state: RagState) -> Dict[str, Any]:
118 question = _require_non_empty_question(state, node_name="retrieve")
119 retrieved = await deps["retriever"].aretrieve(question)
120 out = list(retrieved) if retrieved else []
121 log.debug("[retrieve] nodes=%d", len(out))
122 return {"retrieved_nodes": out}
124 return _impl
127def _rerank_node(deps: RagDeps) -> Callable[[RagState], Awaitable[Dict[str, Any]]]:
128 async def _impl(state: RagState) -> Dict[str, Any]:
129 question = _require_non_empty_question(state, node_name="rerank")
130 retrieved = state.get("retrieved_nodes") or []
131 if not retrieved:
132 log.debug("[rerank] nodes=0 (skip)")
133 return {"reranked_nodes": []}
135 try:
136 ranked = await deps["rerank"](question, retrieved)
137 except Exception as exc:
138 log.warning("[rerank] failed: %s (fallback to retrieved order)", exc)
139 ranked = list(retrieved)
141 out = list(ranked) if ranked else list(retrieved)
142 log.debug("[rerank] nodes=%d", len(out))
143 return {"reranked_nodes": out}
145 return _impl
148def _filter_node(deps: RagDeps) -> Callable[[RagState], Awaitable[Dict[str, Any]]]:
149 async def _impl(state: RagState) -> Dict[str, Any]:
150 lemmas = cast(Set[str], state.get("lemmas") or set())
151 candidates = state.get("reranked_nodes") or state.get("retrieved_nodes") or []
152 if not candidates:
153 log.debug("[filter] candidates=0 (skip)")
154 return {"filtered_nodes": []}
156 # Non-destructive behavior:
157 # - empty lemmas => skip filter
158 if not lemmas:
159 log.debug("[filter] lemmas=0 -> pass-through candidates=%d", len(candidates))
160 return {"filtered_nodes": list(candidates)}
162 try:
163 filtered = await deps["get_filtered_nodes"](list(candidates), cast(Set[str], lemmas))
164 except Exception as exc:
165 log.warning("[filter] get_filtered_nodes failed: %s (pass-through)", exc)
166 return {"filtered_nodes": list(candidates)}
168 if not filtered:
169 log.debug("[filter] filtered=0 -> pass-through candidates=%d", len(candidates))
170 return {"filtered_nodes": list(candidates)}
172 out = list(filtered)
173 log.debug("[filter] nodes=%d", len(out))
174 return {"filtered_nodes": out}
176 return _impl
179def _context_node(
180 deps: RagDeps,
181 *,
182 use_remote_ctx: bool,
183) -> Callable[[RagState], Dict[str, Any]]:
184 def _impl(state: RagState) -> Dict[str, Any]:
185 lemmas = cast(Set[str], state.get("lemmas") or set())
187 # Prefer filtered, but never allow empty to kill context completely.
188 nodes = (
189 state.get("filtered_nodes")
190 or state.get("reranked_nodes")
191 or state.get("retrieved_nodes")
192 or []
193 )
195 limit = deps["ctx_len_remote"] if use_remote_ctx else deps["ctx_len_local"]
197 ctx_txt = ""
198 if nodes:
199 try:
200 ctx_txt = deps["build_context"](list(nodes), cast(Set[str], lemmas), int(limit))
201 except Exception as exc:
202 log.warning("[context] build_context failed: %s", exc)
203 ctx_txt = ""
205 # If context came back empty, retry without lemma constraints (common failure mode).
206 if not ctx_txt and lemmas:
207 try:
208 ctx_txt = deps["build_context"](list(nodes), set(), int(limit))
209 if ctx_txt:
210 log.debug("[context] recovered by retry with lemmas=0")
211 except Exception as exc:
212 log.warning("[context] retry build_context(lemmas=0) failed: %s", exc)
214 ctx_len = len((ctx_txt or "").strip())
215 log.debug(
216 "[context] nodes=%d limit=%d ctx_len=%d remote=%s",
217 len(nodes),
218 int(limit),
219 ctx_len,
220 bool(use_remote_ctx),
221 )
222 return {"context": ctx_txt or ""}
224 return _impl
227def _llm_node(deps: RagDeps) -> Callable[[RagState], Awaitable[Dict[str, Any]]]:
228 async def _impl(state: RagState) -> Dict[str, Any]:
229 question = _require_non_empty_question(state, node_name="llm")
230 context = (state.get("context") or "").strip()
231 sys_prompt = (deps.get("sys_prompt") or "").strip()
233 timeout_total = int(getattr(cfg, "HTTP_TIMEOUT_TOTAL", 240))
235 log.debug("[llm] q_len=%d ctx_len=%d timeout_sec=%d", len(question), len(context), timeout_total)
237 try:
238 text, _used_fallback = await deps["query_model"](
239 sys_prompt=sys_prompt,
240 ctx_txt=context,
241 q=question,
242 timeout_sec=timeout_total,
243 )
244 except TypeError:
245 text, _used_fallback = await deps["query_model"](sys_prompt, context, question, timeout_total)
247 answer = (text or "").strip() or "No answer."
248 log.debug("[llm] answer_len=%d", len(answer))
249 return {"answer": answer}
251 return _impl
254def build_rag_graph(
255 deps: RagDeps,
256 *,
257 use_remote_ctx: bool = False,
258) -> Any:
259 """
260 Build a LangGraph StateGraph for the RAG pipeline (no side effects, not compiled).
262 Returns:
263 StateGraph-like object (langgraph is optional at import time).
264 """
265 if StateGraph is Any: # pragma: no cover
266 raise RuntimeError("langgraph is not installed; cannot build graph.")
268 graph = StateGraph(RagState)
270 graph.add_node("lemmas", _lemmas_node(deps))
271 graph.add_node("retrieve", _retrieve_node(deps))
272 graph.add_node("rerank", _rerank_node(deps))
273 graph.add_node("filter", _filter_node(deps))
274 graph.add_node("context", _context_node(deps, use_remote_ctx=use_remote_ctx))
275 graph.add_node("llm", _llm_node(deps))
277 graph.add_edge(START, "lemmas")
278 graph.add_edge("lemmas", "retrieve")
279 graph.add_edge("retrieve", "rerank")
280 graph.add_edge("rerank", "filter")
281 graph.add_edge("filter", "context")
282 graph.add_edge("context", "llm")
283 graph.add_edge("llm", END)
285 return graph
288async def run_rag_pipeline(
289 question: str,
290 deps: RagDeps,
291 *,
292 use_remote_ctx: bool = False,
293 enable_tracing: bool = False,
294) -> str:
295 """
296 Convenience runner to build, compile, and invoke the graph once.
297 """
298 if not isinstance(question, str) or not question.strip():
299 raise ValueError("Question must be a non-empty string.")
301 if enable_tracing:
302 _maybe_init_langsmith()
304 graph = build_rag_graph(deps, use_remote_ctx=use_remote_ctx)
305 app = graph.compile()
307 init_state: RagState = {"question": question.strip()}
308 output: RagState = await app.ainvoke(init_state)
310 answer = (output.get("answer") or "").strip()
311 return answer or "No answer."
314__all__ = [
315 "RagState",
316 "RagDeps",
317 "build_rag_graph",
318 "run_rag_pipeline",
319]