""" QA service layer for Grant Analyst. Provides streaming and non-streaming QA responses with prompt injection hardening. API and UI should use this service instead of calling LLM/search directly. """ from __future__ import annotations import json import logging from typing import Iterable, List, Optional from .models import QARequest, QAChunk, ChunkType, Grant, CitationInfo from .search.service import search_grants from .llm_client import LLMClient from .config import get_settings from .cache.memo import cache_get, cache_put logger = logging.getLogger(__name__) # Shared LLM client singleton _llm_client: LLMClient | None = None def get_llm_client() -> LLMClient: """Get or create shared LLM client instance.""" global _llm_client if _llm_client is None: _llm_client = LLMClient() return _llm_client # Context limits for LLM calls MAX_GRANTS = 5 MAX_DESC_CHARS = 1200 # Concise system prompt - reduced token usage SYSTEM_PROMPT = """You are a specialised grant assistant focused on Innovate UK and related UK/EU funding schemes. Use ONLY the retrieved grant records and their metadata as your primary evidence. If required information is not present in the retrieved context, say you are unsure rather than guessing. When you answer: - Be concise and structured. - Prioritise: grant name, funder, key eligibility, funding amount/range, deadline, and URL. - Highlight constraints or caveats clearly. - If multiple grants are relevant, list them clearly instead of writing long paragraphs. """ def _sanitize_grant_text(text: str) -> str: """ Sanitize grant text to remove potential prompt injection attempts. Strips lines that look like instructions or system prompts. Args: text: Raw text from grant data Returns: Sanitized text safe for inclusion in prompts """ if not text: return "" # Keywords that indicate potential injection injection_keywords = [ "ignore previous", "ignore all previous", "system prompt", "you are now", "forget everything", "new instructions", "disregard", "override", "act as", ] lines = text.split("\n") safe_lines = [] for line in lines: line_lower = line.lower().strip() # Skip lines that look like injection attempts if any(keyword in line_lower for keyword in injection_keywords): logger.warning(f"Filtered potential injection: {line[:50]}...") continue safe_lines.append(line) return "\n".join(safe_lines) def format_grant_context(grants: List[Grant]) -> str: """ Format top search results into a compact context string for the LLM. Limits both the number of grants and description length to control token usage. Args: grants: List of Grant objects Returns: Formatted, compact context string """ lines: list[str] = [] for i, grant in enumerate(grants[:MAX_GRANTS]): # Extract key attributes title = grant.title or "Untitled" funder = grant.source or grant.programme or "" ref_id = grant.id or "" url = grant.url or "" deadline = str(grant.close_date) if grant.close_date else "" status = grant.status or "" # Extract and truncate description/summary desc = grant.summary or grant.scope or "" if desc and len(desc) > MAX_DESC_CHARS: # Sanitize and truncate desc = _sanitize_grant_text(desc[:MAX_DESC_CHARS]).rstrip() + "..." elif desc: desc = _sanitize_grant_text(desc) # Format funding amount amount = "" if grant.funding: funding_parts = [] if grant.funding.min is not None: funding_parts.append(f"£{grant.funding.min:,.0f}") if grant.funding.max is not None: funding_parts.append(f"£{grant.funding.max:,.0f}") if funding_parts: amount = " - ".join(funding_parts) # Build compact block block = [ f"Grant {i+1}: {title}", f" Reference: {ref_id}" if ref_id else "", f" Funder: {funder}" if funder else "", f" Status: {status}" if status else "", f" Funding: {amount}" if amount else "", f" Deadline: {deadline}" if deadline else "", f" URL: {url}" if url else "", f" Summary: {desc}" if desc else "", ] lines.append("\n".join([ln for ln in block if ln])) return "\n\n".join(lines) def stream_answer(req: QARequest) -> Iterable[QAChunk]: """ Generate streaming QA response. Yields QAChunk objects that can be serialized to NDJSON or SSE. Args: req: Validated QA request Yields: QAChunk objects of various types (metadata, token, citations, done, error) """ settings = get_settings() try: # Initialize LLM client llm = get_llm_client() # Send metadata yield QAChunk(type=ChunkType.METADATA, session_id=req.session_id, query=req.query) # Search for relevant grants logger.debug(f"Query: {req.query[:100]}") hits = search_grants(req.query, req.filters, limit=10) logger.info(f"Search found {len(hits)} hits") if not hits: # No grants found yield QAChunk( type=ChunkType.TOKEN, content="I couldn't find any grants matching your query. Try different keywords or broader terms.", ) yield QAChunk(type=ChunkType.DONE, latency_ms=0) return # Extract grants and build compact context grants = [hit.grant for hit in hits] context = format_grant_context(grants) # Build prompt user_prompt = f"""User query: {req.query} Relevant grant opportunities: {context} Based on the grants above, answer the user's query concisely and accurately. Cite specific grants by ID and title. If none of the grants are truly relevant, say so.""" messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_prompt}, ] # Stream LLM response (log moved to debug to reduce overhead) logger.debug(f"Streaming LLM for query: {req.query[:50]}") # Use streaming chat for token in llm.chat( messages, stream=True, max_tokens=1200, model_type="analyzer", # Use analyzer model for QA ): yield QAChunk(type=ChunkType.TOKEN, content=token) # Send citations citations = [ {"grant_id": grant.id, "title": grant.title, "url": grant.url, "score": hit.score} for grant, hit in zip(grants[:5], hits[:5]) ] yield QAChunk(type=ChunkType.CITATIONS, citations=citations) # Send completion yield QAChunk(type=ChunkType.DONE) except Exception as e: logger.error(f"Error in stream_answer: {e}", exc_info=True) yield QAChunk(type=ChunkType.ERROR, error=str(e)) def answer_question(req: QARequest) -> dict: """ Generate non-streaming QA response with memoization. Collects all chunks from stream_answer and returns a complete response. Caches results to avoid redundant LLM calls for identical queries. Args: req: Validated QA request Returns: Dict with answer, citations, and metadata """ import time start_time = time.time() # Build cache key from normalized query and filters normalized_query = req.query.strip() filters_dict = req.filters.dict() if req.filters else {} cache_key = f"qa:{normalized_query}|filters:{json.dumps(filters_dict, sort_keys=True)}" # Try cache first cached = cache_get(cache_key) if cached: logger.debug(f"Cache hit for query: {normalized_query[:50]}") return cached answer_parts = [] citations = [] error = None success = True try: for chunk in stream_answer(req): if chunk.type == ChunkType.TOKEN: if chunk.content: answer_parts.append(chunk.content) elif chunk.type == ChunkType.CITATIONS: citations = chunk.citations or [] elif chunk.type == ChunkType.ERROR: error = chunk.error success = False break except Exception as e: error = str(e) success = False logger.error(f"Error in answer_question: {e}", exc_info=True) latency_ms = int((time.time() - start_time) * 1000) result = { "session_id": req.session_id or "unknown", "query": req.query, "answer": "".join(answer_parts), "citations": citations, "latency_ms": latency_ms, "success": success, "error": error, } # Cache successful responses if success and not error: try: cache_put(cache_key, result) logger.debug(f"Cached response for: {normalized_query[:50]}") except Exception as cache_err: logger.debug(f"Cache storage failed: {cache_err}") return result