import json import re import polars as pl from src.config import Config from src.llm import LocalLLM from src.query_router import QueryRouter from src.retriever import Retriever class AnswerGenerator: def __init__( self, config: Config, llm: LocalLLM, retriever: Retriever, facts_df: pl.DataFrame, pages_df: pl.DataFrame, ): self.config = config self.llm = llm self.router = QueryRouter() self.retriever = retriever self.facts_df = facts_df self.pages_df = pages_df def answer(self, query: str) -> str: route = self.router.route(query) retrieved = self.retriever.retrieve(query) numeric_results = None if route in ("numeric", "calculation", "range", "comparison"): numeric_results = self._search_facts(query) if route == "numeric": return self._answer_numeric(query, retrieved, numeric_results) elif route == "calculation": return self._answer_calculation(query, retrieved, numeric_results) elif route == "comparison": return self._answer_comparison(query, retrieved, numeric_results) elif route == "range": return self._answer_range(query, retrieved, numeric_results) elif route == "policy": return self._answer_policy(query, retrieved) elif route == "summary": return self._answer_summary(query, retrieved) else: return self._answer_general(query, retrieved) def _search_facts(self, query: str) -> list[dict]: if self.facts_df.is_empty(): return [] query_lower = query.lower() numbers = re.findall(r"\d+", query) query_numbers = [float(n) for n in numbers] relevant_facts = [] for row in self.facts_df.to_dicts(): entity = row.get("entity", "").lower() attribute = row.get("attribute", "").lower() entity_words = entity.split() if any(w in query_lower for w in entity_words if len(w) > 2): relevant_facts.append(row) continue if attribute in query_lower: relevant_facts.append(row) continue if query_numbers: mn = row.get("min_quantity") mx = row.get("max_quantity") if mn is not None and mx is not None: for qn in query_numbers: if mn <= qn <= mx: relevant_facts.append(row) break elif mn is not None and mx is None: if any(qn >= mn for qn in query_numbers): relevant_facts.append(row) return relevant_facts def _answer_numeric(self, query: str, retrieved: list[dict], facts: list[dict]) -> str: if facts: answer = self._build_numeric_answer(query, facts, retrieved) if answer: return answer context = self._format_context(retrieved[:5]) return self._llm_answer(query, context) def _answer_calculation(self, query: str, retrieved: list[dict], facts: list[dict]) -> str: if facts: calc_result = self._try_calculate(query, facts) if calc_result: return calc_result context = self._format_context(retrieved[:5]) return self._llm_answer( query, context, instruction="If the question requires calculation, perform the calculation step by step and show your reasoning.", ) def _answer_comparison(self, query: str, retrieved: list[dict], facts: list[dict]) -> str: if facts: items = set(f["entity"] for f in facts if f.get("entity")) if len(items) >= 2: comparison = self._build_comparison(query, facts) if comparison: return comparison context = self._format_context(retrieved[:7]) return self._llm_answer(query, context) def _answer_range(self, query: str, retrieved: list[dict], facts: list[dict]) -> str: if facts: mn = min((f.get("min_quantity") or 0 for f in facts if f.get("min_quantity") is not None), default=None) mx = max((f.get("max_quantity") or 0 for f in facts if f.get("max_quantity") is not None), default=None) vals = sorted(set(f["value"] for f in facts if f.get("value")), key=float) if mn is not None and mx is not None: return ( f"The available range is {mn} to {mx}.\n" + f"Prices in this range: {', '.join(vals)}\n" + self._cite_sources(facts) ) context = self._format_context(retrieved[:5]) return self._llm_answer(query, context) def _answer_policy(self, query: str, retrieved: list[dict]) -> str: context = self._format_context(retrieved[:7]) return self._llm_answer( query, context, instruction="Answer based on policy/terms information. If the answer depends on missing information (e.g., user's specific situation), ask a follow-up question to clarify.", ) def _answer_summary(self, query: str, retrieved: list[dict]) -> str: context = self._format_context(retrieved[:7]) return self._llm_answer(query, context) def _answer_general(self, query: str, retrieved: list[dict]) -> str: if not retrieved: return ( "I don't have enough information from the crawled data to answer that question. " "Please try crawling a website first, or rephrase your question." ) context = self._format_context(retrieved[:5]) return self._llm_answer(query, context) def _try_calculate(self, query: str, facts: list[dict]) -> str | None: numbers = re.findall(r"\d+[.,]?\d*", query) if not numbers: return None query_lower = query.lower() has_price_per = any("price" in f.get("attribute", "").lower() or "per" in f.get("attribute", "").lower() for f in facts) has_unit_price = any(f.get("unit") == "per_item" for f in facts) if has_price_per or has_unit_price: for f in facts: attr = f.get("attribute", "").lower() if "price" in attr or "per" in attr: try: unit_price = float(f["value"]) except (ValueError, TypeError): continue qty = None for n_str in numbers: try: candidate = float(n_str.replace(",", ".")) mn = f.get("min_quantity") mx = f.get("max_quantity") if mn is not None and mx is not None: if mn <= candidate <= mx: qty = candidate break else: qty = candidate break except ValueError: continue if qty is not None: total = qty * unit_price currency = f.get("currency", "") currency_symbol = {"USD": "$", "EUR": "€", "GBP": "£"}.get(currency, "") return ( f"**Calculation:** {qty} × {unit_price} = {total}\n\n" + f"For {qty} items at {currency_symbol}{unit_price} each, " + f"the total is **{currency_symbol}{total:.2f}**.\n\n" + self._cite_sources([f]) ) return None def _build_numeric_answer(self, query: str, facts: list[dict], retrieved: list[dict]) -> str: for f in facts: attr = f.get("attribute", "").lower() if "price" in attr: value = f.get("value", "") currency = f.get("currency", "USD") unit = f.get("unit", "") mn = f.get("min_quantity") mx = f.get("max_quantity") entity = f.get("entity", "") parts = [f"The price for **{entity}** is **{value} {currency}**"] if unit: parts.append(f"({unit})") if mn is not None and mx is not None: if mn != mx: parts.append(f"for quantities between {mn} and {mx}") else: parts.append(f"for quantity {mn}") parts.append(f"\n\n{self._cite_sources([f])}") return " ".join(parts) return "" def _build_comparison(self, query: str, facts: list[dict]) -> str: by_entity: dict[str, list[dict]] = {} for f in facts: ent = f.get("entity", "unknown") by_entity.setdefault(ent, []).append(f) lines = [] for entity, efacts in by_entity.items(): prices = [f for f in efacts if "price" in f.get("attribute", "").lower()] if prices: vals = [f"{p['value']} {p.get('currency', '')}" for p in prices] lines.append(f"- **{entity}**: {', '.join(vals)}") if lines: lines.append("") lines.append(self._cite_sources(facts)) return "\n".join(lines) return "" def _llm_answer(self, query: str, context: str, instruction: str = "") -> str: if not context: return ( "I don't have enough information from the crawled data to answer that question. " "Please try crawling a website first." ) system_prompt = "You are a precise QA assistant that answers questions based ONLY on the provided crawled content." if instruction: system_prompt += f"\n\n{instruction}" try: text = self.llm.generate( messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"}, ], max_tokens=self.config.llm_max_tokens, temperature=self.config.llm_temperature, ) if text: return text except Exception: pass return "I'm sorry, I couldn't generate an answer at this time. Please try again later." def _format_context(self, retrieved: list[dict]) -> str: if not retrieved: return "" parts = [] seen_urls = set() for i, r in enumerate(retrieved, 1): text = r.get("text", "")[:500] url = r.get("url", "") title = r.get("title", "") source = f"[Source: {title}]({url})" if url else "" if url: seen_urls.add(url) parts.append(f"[{i}] {text}\n{source}") return "\n\n---\n\n".join(parts) def _cite_sources(self, facts: list[dict]) -> str: urls = set() for f in facts: url = f.get("source_url", "") if url: urls.add(url) if urls: return "**Sources:** " + ", ".join(sorted(urls)) return ""