Spaces:
Running on Zero
Running on Zero
| import logging | |
| import os | |
| import re | |
| from functools import lru_cache | |
| from pathlib import Path | |
| from typing import TypedDict | |
| from dotenv import load_dotenv | |
| from groq import Groq | |
| from langchain_community.vectorstores import FAISS | |
| from langchain_huggingface import HuggingFaceEmbeddings | |
| from langchain_core.documents import Document | |
| from langchain_text_splitters import RecursiveCharacterTextSplitter | |
| from langgraph.graph import END, START, StateGraph | |
| load_dotenv() | |
| logger = logging.getLogger(__name__) | |
| BASE_DIR = Path(__file__).resolve().parent | |
| FAQ_PATH = BASE_DIR / "support_faq.txt" | |
| EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2" | |
| GROQ_MODEL_NAME = os.getenv("GROQ_MODEL", "openai/gpt-oss-20b") | |
| TOP_K = 3 | |
| MAX_NEW_TOKENS = 120 | |
| MAX_RETRIES = 1 | |
| # Similarity score threshold for FAISS distance strategy: | |
| # lower L2 distance means a closer match. This value is a starting point | |
| # and should be checked against the user's own FAQ test questions. | |
| MAX_FAISS_DISTANCE = 1.25 | |
| FALLBACK_ANSWER = ( | |
| "I couldn't find a reliable answer in the available FAQ. " | |
| "Please contact the support team through the Help section. " | |
| "Never share your password, OTP, or full payment-card details." | |
| ) | |
| class SupportState(TypedDict, total=False): | |
| question: str | |
| intent: str | |
| search_query: str | |
| context: str | |
| sources: list[str] | |
| draft: str | |
| retry_count: int | |
| final_answer: str | |
| relevant: bool | |
| class SupportRAG: | |
| """Initialize and manage the complete RAG pipeline.""" | |
| def __init__(self): | |
| if not FAQ_PATH.exists(): | |
| raise FileNotFoundError( | |
| f"Cannot find {FAQ_PATH}. Place support_faq.txt beside support_rag.py." | |
| ) | |
| faq_text = FAQ_PATH.read_text(encoding="utf-8").strip() | |
| if not faq_text: | |
| raise ValueError("support_faq.txt is empty.") | |
| logger.info("Loading FAQ and building FAISS index.") | |
| document = Document( | |
| page_content=faq_text, | |
| metadata={"source": FAQ_PATH.name}, | |
| ) | |
| splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=450, | |
| chunk_overlap=60, | |
| ) | |
| chunks = splitter.split_documents([document]) | |
| if not chunks: | |
| raise ValueError("No text chunks were created from support_faq.txt.") | |
| # Force embeddings onto CPU. In ZeroGPU Spaces, Sentence Transformers | |
| # may incorrectly auto-detect the emulated CUDA device and fail during | |
| # CUDA initialization. Groq handles text generation remotely, so CPU | |
| # embeddings are sufficient for this small FAQ index. | |
| self.embedding_model = HuggingFaceEmbeddings( | |
| model_name=EMBEDDING_MODEL_NAME, | |
| model_kwargs={"device": "cpu"}, | |
| ) | |
| self.vector_store = FAISS.from_documents(chunks, self.embedding_model) | |
| self.retriever = self.vector_store.as_retriever( | |
| search_type="similarity", | |
| search_kwargs={"k": TOP_K}, | |
| ) | |
| # The language model runs through Groq's hosted API. | |
| # Configure GROQ_API_KEY as a Hugging Face Space Secret. | |
| api_key = os.getenv("GROQ_API_KEY") | |
| if not api_key: | |
| raise RuntimeError( | |
| "GROQ_API_KEY is missing. Add it under Space Settings > Secrets." | |
| ) | |
| self.groq_client = Groq(api_key=api_key) | |
| logger.info("Using Groq model: %s", GROQ_MODEL_NAME) | |
| self.graph = self._build_graph() | |
| logger.info("Customer Support RAG pipeline initialized.") | |
| def _normalise(text: str) -> str: | |
| """Lowercase text and collapse punctuation/whitespace.""" | |
| return re.sub(r"\s+", " ", re.sub(r"[^a-z0-9\s]", " ", text.lower())).strip() | |
| def _classify_question(question: str) -> str: | |
| """Classify intent using simple keyword rules.""" | |
| text = question.lower() | |
| categories = { | |
| "order_tracking": [ | |
| "track", "tracking", "shipment", "delivery", | |
| "where is my order", "order status", "dispatch", | |
| ], | |
| "returns_refunds": [ | |
| "return", "refund", "money back", "exchange", | |
| ], | |
| "account_access": [ | |
| "password", "sign in", "sign-in", "login", | |
| "log in", "account", "otp", "verification code", | |
| ], | |
| "billing_payments": [ | |
| "payment", "charged", "charge", "billing", | |
| "invoice", "card", "transaction", "double charge", | |
| ], | |
| "technical_support": [ | |
| "website", "site", "error", "loading", | |
| "crash", "technical", "not working", "browser", "cache", | |
| ], | |
| } | |
| for category, keywords in categories.items(): | |
| if any(keyword in text for keyword in keywords): | |
| return category | |
| return "general_support" | |
| def classify_intent(self, state: SupportState) -> SupportState: | |
| question = state["question"].strip() | |
| return { | |
| **state, | |
| "question": question, | |
| "intent": self._classify_question(question), | |
| "search_query": question, | |
| "retry_count": 0, | |
| "sources": [], | |
| "relevant": False, | |
| } | |
| def retrieve_context(self, state: SupportState) -> SupportState: | |
| """Retrieve the best FAQ chunks and assess their distance.""" | |
| query = state.get("search_query") or state["question"] | |
| # FAISS returns L2 distance for this vector store; smaller is closer. | |
| results = self.vector_store.similarity_search_with_score(query, k=TOP_K) | |
| context_parts = [] | |
| sources = [] | |
| best_distance = float("inf") | |
| for doc, distance in results: | |
| text = doc.page_content.strip() | |
| best_distance = min(best_distance, float(distance)) | |
| if text: | |
| context_parts.append(text) | |
| source = doc.metadata.get("source", "unknown") | |
| if source not in sources: | |
| sources.append(source) | |
| context = "\n\n".join(context_parts) | |
| relevant = bool(context) and best_distance <= MAX_FAISS_DISTANCE | |
| logger.info( | |
| "Retrieved %d chunks for query %r; best FAISS distance=%.4f; relevant=%s", | |
| len(results), query, best_distance, relevant, | |
| ) | |
| return { | |
| **state, | |
| "context": context, | |
| "sources": sources if relevant else [], | |
| "relevant": relevant, | |
| } | |
| def _generate(self, prompt: str) -> str: | |
| """Generate a concise answer using Groq's hosted chat-completions API.""" | |
| messages = [ | |
| { | |
| "role": "system", | |
| "content": ( | |
| "You are a customer support assistant. Answer only from the FAQ " | |
| "context provided. Do not use outside knowledge, invent policies, " | |
| "or claim access to live systems. If the context does not answer " | |
| "the question, say you cannot confirm it. Never request passwords, " | |
| "OTPs, or full card details. Give a concise answer in 2 to 4 sentences." | |
| ), | |
| }, | |
| {"role": "user", "content": prompt}, | |
| ] | |
| response = self.groq_client.chat.completions.create( | |
| model=GROQ_MODEL_NAME, | |
| messages=messages, | |
| temperature=0, | |
| max_tokens=MAX_NEW_TOKENS, | |
| ) | |
| content = response.choices[0].message.content | |
| return (content or "").strip() | |
| def draft_answer(self, state: SupportState) -> SupportState: | |
| """Generate an answer only if retrieval passed the relevance check.""" | |
| context = state.get("context", "").strip() | |
| if not context or not state.get("relevant", False): | |
| return {**state, "draft": "", "final_answer": ""} | |
| prompt = ( | |
| f"Customer question:\n{state['question']}\n\n" | |
| f"FAQ context:\n{context}\n\n" | |
| "Instructions:\n" | |
| "- Answer the question directly using only the FAQ context.\n" | |
| "- For returns, explain the Return Request steps and relevant eligibility.\n" | |
| "- If provided, include the refund processing time.\n" | |
| "- Do not claim to know live order, refund, or payment status.\n" | |
| "- If the FAQ does not specify the answer, say it is not specified.\n" | |
| "- Keep the answer concise, in 2 to 4 short sentences.\n" | |
| "Answer:" | |
| ) | |
| try: | |
| answer = self._generate(prompt) | |
| except Exception: | |
| logger.exception("Answer generation failed.") | |
| answer = "" | |
| answer = answer.strip() | |
| logger.info("Generated draft length=%d", len(answer)) | |
| return {**state, "draft": answer} | |
| def review_answer(self, state: SupportState) -> SupportState: | |
| """Accept a draft only when retrieval passed the relevance check.""" | |
| answer = state.get("draft", "").strip() | |
| context = state.get("context", "").strip() | |
| if answer and context and state.get("relevant", False): | |
| return {**state, "final_answer": answer} | |
| return {**state, "final_answer": ""} | |
| def route_after_review(self, state: SupportState) -> str: | |
| # An irrelevant query should go directly to fallback rather than retrying | |
| # with a broader query that could retrieve unrelated chunks. | |
| if not state.get("relevant", False): | |
| return "fallback" | |
| if state.get("final_answer"): | |
| return "finish" | |
| if state.get("retry_count", 0) < MAX_RETRIES: | |
| return "retry" | |
| return "fallback" | |
| def rewrite_query(self, state: SupportState) -> SupportState: | |
| """Add intent-specific terms, then retrieve again.""" | |
| additions = { | |
| "order_tracking": "order tracking shipment delivery status dispatch", | |
| "returns_refunds": "refund request return policy eligibility refund processing time", | |
| "account_access": "account login password recovery sign in reset email", | |
| "billing_payments": "billing payment duplicate charge transaction", | |
| "technical_support": "website technical problem troubleshooting browser", | |
| "general_support": "customer support FAQ policy instructions", | |
| } | |
| question = state["question"] | |
| extra_terms = additions.get( | |
| state.get("intent", "general_support"), | |
| additions["general_support"], | |
| ) | |
| new_state = { | |
| **state, | |
| "search_query": f"{question} {extra_terms}", | |
| "retry_count": state.get("retry_count", 0) + 1, | |
| "final_answer": "", | |
| "draft": "", | |
| } | |
| return self.retrieve_context(new_state) | |
| def fallback_answer(self, state: SupportState) -> SupportState: | |
| return {**state, "final_answer": FALLBACK_ANSWER, "sources": []} | |
| def _build_graph(self): | |
| workflow = StateGraph(SupportState) | |
| workflow.add_node("classify", self.classify_intent) | |
| workflow.add_node("retrieve", self.retrieve_context) | |
| workflow.add_node("draft", self.draft_answer) | |
| workflow.add_node("review", self.review_answer) | |
| workflow.add_node("retry", self.rewrite_query) | |
| workflow.add_node("fallback", self.fallback_answer) | |
| workflow.add_edge(START, "classify") | |
| workflow.add_edge("classify", "retrieve") | |
| workflow.add_conditional_edges( | |
| "retrieve", | |
| lambda state: "draft" if state.get("relevant", False) else "fallback", | |
| {"draft": "draft", "fallback": "fallback"}, | |
| ) | |
| workflow.add_edge("draft", "review") | |
| workflow.add_conditional_edges( | |
| "review", | |
| self.route_after_review, | |
| {"finish": END, "retry": "retry", "fallback": "fallback"}, | |
| ) | |
| workflow.add_edge("retry", "draft") | |
| workflow.add_edge("fallback", END) | |
| return workflow.compile() | |
| def ask(self, question: str) -> dict: | |
| result = self.graph.invoke({"question": question}) | |
| return { | |
| "intent": result.get("intent", "general_support"), | |
| "answer": result.get("final_answer") or FALLBACK_ANSWER, | |
| "sources": result.get("sources", []), | |
| "retry_count": result.get("retry_count", 0), | |
| } | |
| def get_support_rag() -> SupportRAG: | |
| """Initialize the pipeline once and reuse it.""" | |
| return SupportRAG() | |
| def ask_support(question: str) -> dict: | |
| """Entry point used by app.py.""" | |
| return get_support_rag().ask(question) |