Spaces:
Running
Running
File size: 2,469 Bytes
2ecc4a7 27d1bb9 2ecc4a7 | 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 | from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional
from ..services.supabase_client import supabase_service
from ..services.embed_service import get_embedding, truncate_to_1k
from ..services.llm_service import llm_service
from ..utils.ws_manager import ws_manager
import logging
logger = logging.getLogger(__name__)
class BaseRAGTechnique(ABC):
def __init__(self, job_id: str, user_id: str):
self.job_id = job_id
self.user_id = user_id
self.supabase = supabase_service
self.llm = llm_service
def build_prompt(self, query: str, chunks: List[Dict[str, Any]]) -> str:
context = "\n\n".join([c["text"] for c in chunks])
return f"""Context:
{context}
Question: {query}
Instructions:
1. If the user is just saying a general greeting (e.g., "hi", "hello"), respond politely without using the context.
2. If the user asks a general question about the document itself (e.g., "what is this document about?", "summarize"), summarize the provided context to answer.
3. Otherwise, answer the question based ONLY on the provided context. If the answer is not in the context, state that clearly."""
@abstractmethod
async def retrieve(self, query: str, document_id: str, top_k: int, **kwargs) -> List[Dict[str, Any]]:
pass
@abstractmethod
async def generate(self, query: str, chunks: List[Dict[str, Any]]) -> str:
pass
async def run(self, query: str, document_id: str, top_k: int = 5, **kwargs) -> Dict[str, Any]:
"""Execute the full RAG pipeline."""
try:
# 1. Retrieval
chunks = await self.retrieve(query, document_id, top_k, **kwargs)
# 2. Generation
answer = await self.generate(query, chunks)
return {
"answer": answer,
"sources": chunks,
"job_id": self.job_id
}
except Exception as e:
logger.error(f"RAG execution failed: {e}")
await self.emit("ERROR", "red", f"Critical error: {str(e)}")
raise
async def emit(self, step: str, color: str, detail: str, metadata: dict = {}):
"""Broadcast progress to frontend."""
await ws_manager.emit(self.job_id, self.user_id, {
"step": step,
"status": "running",
"color": color,
"detail": detail,
"metadata": metadata
})
|