Spaces:
Sleeping
Sleeping
File size: 4,278 Bytes
2f27a28 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """RAG pipeline: Chroma retrieval over Stack Overflow Python Q&A + Groq Llama generation."""
import logging
import chromadb
from groq import AsyncGroq
from app import config
logger = logging.getLogger(__name__)
SYSTEM_PROMPT = """\
You are a Python programming Q&A assistant for data science learners. You answer \
questions using the provided Stack Overflow excerpts as your primary source of truth.
Rules:
- Ground your answer in the provided context. When you use information from an \
excerpt, cite it inline as [1], [2], etc. matching the excerpt numbers.
- Include short, runnable code examples where they help.
- If the context does not contain enough information, say so explicitly and then \
give your best general answer, clearly marked as not sourced from the context.
- If the question is not about Python programming, politely say you only answer \
Python programming questions and do not attempt to answer it.
- Be concise and accurate. Prefer modern Python 3 idioms; if an excerpt shows \
Python 2 syntax, modernise it and mention that you did.
"""
USER_PROMPT_TEMPLATE = """\
Stack Overflow excerpts:
{context}
Learner's question: {question}
"""
class Retriever:
"""Thin wrapper around a persistent Chroma collection of SO Python Q&A pairs."""
def __init__(self, chroma_dir: str = config.CHROMA_DIR,
collection_name: str = config.COLLECTION_NAME):
client = chromadb.PersistentClient(path=chroma_dir)
self.collection = client.get_collection(collection_name)
def count(self) -> int:
return self.collection.count()
def query(self, question: str, k: int = config.TOP_K) -> list[dict]:
res = self.collection.query(
query_texts=[question],
n_results=k,
include=["documents", "metadatas", "distances"],
)
hits = []
for doc, meta, dist in zip(res["documents"][0], res["metadatas"][0], res["distances"][0]):
hits.append({
"text": doc,
"title": meta["title"],
"url": meta["url"],
"answer_score": meta["answer_score"],
"tags": meta.get("tags", ""),
# Cosine distance -> similarity in [0, 1].
"relevance": round(max(0.0, 1.0 - dist), 4),
"distance": dist,
})
return hits
def build_context(hits: list[dict]) -> str:
blocks = []
for i, h in enumerate(hits, start=1):
blocks.append(f"[{i}] {h['title']} (answer score: {h['answer_score']})\n{h['text']}")
return "\n\n---\n\n".join(blocks)
class RAGPipeline:
def __init__(self, retriever: Retriever):
self.retriever = retriever
self.llm = AsyncGroq(api_key=config.GROQ_API_KEY)
async def ask(self, question: str, top_k: int = config.TOP_K) -> dict:
hits = self.retriever.query(question, k=top_k)
# RELEVANCE_THRESHOLD is a minimum similarity (0-1); distance = 1 - similarity.
relevant = [h for h in hits if h["relevance"] >= config.RELEVANCE_THRESHOLD]
grounded = len(relevant) > 0
# When nothing passes the threshold, pass all hits anyway —
# the prompt instructs the model to flag unsourced answers and decline off-topic ones.
context_hits = relevant if grounded else hits
completion = await self.llm.chat.completions.create(
model=config.GROQ_MODEL,
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": USER_PROMPT_TEMPLATE.format(
context=build_context(context_hits), question=question)},
],
temperature=0.2,
max_tokens=1024,
timeout=config.LLM_TIMEOUT_SECONDS,
)
answer = completion.choices[0].message.content
return {
"answer": answer,
"sources": [
{
"title": h["title"],
"url": h["url"],
"relevance": h["relevance"],
"answer_score": h["answer_score"],
}
for h in context_hits
],
"grounded": grounded,
"model": config.GROQ_MODEL,
}
|