Spaces:

zaheer140
/

RagAgent

Sleeping

RagAgent / rag_engine.py

Update rag_engine.py

e4f8454 verified about 2 months ago

5.81 kB

from typing import List, Dict, Any
from rag.ocr import extract_text as ocr_extract_text
from rag.chunker import chunk_text as chunker_chunk_text
from rag.retriever import get_pinecone_store as retriever_get_pinecone_store, retrieve_hybrid
from rag.reranker import rerank
from rag.engine import agent_stream as engine_agent_stream

# Re-expose the system prompt for any direct imports (though now managed in rag/prompt.py)
from rag.prompt import SYSTEM_PROMPT

def extract_text(file_bytes: bytes) -> List[Dict[str, Any]]:
    """
    Extracts text from PDF pages (Proxy to rag/ocr.py).
    """
    return ocr_extract_text(file_bytes)

def chunk_text(pages: List[Dict[str, Any]], filename: str) -> List[Dict[str, Any]]:
    """
    Splits pages into recursive/semantic chunks (Proxy to rag/chunker.py).
    """
    return chunker_chunk_text(pages, filename)

def save_to_pinecone(machine_id: str, chunks: List[Dict[str, Any]]):
    """
    Processes raw chunks (classification, metadata generation, sequential linking)
    and saves them to both Pinecone and the local cache (Proxy to rag/ingestion.py logic).
    """
    if not chunks:
        return
        
    filename = chunks[0]["filename"]
    # Use first chunk text as sample for classification
    sample_text = chunks[0]["text"] if chunks else ""
    
    from rag.classifier import classify_document
    from rag.metadata import generate_metadata
    from langchain_core.documents import Document
    from config import STORE_DIR
    import json
    
    # 1. Automatically classify document type
    doc_type = classify_document(filename, sample_text)
    print(f"RAG Proxy: Classified '{filename}' as '{doc_type}'")
    
    # 2. Generate heading detection and sequential links
    enriched_chunks = generate_metadata(chunks, machine_id, doc_type)
    
    # 3. Upload to Pinecone
    try:
        vector_store = retriever_get_pinecone_store(machine_id)
        documents = [
            Document(
                page_content=c["text"],
                metadata={
                    "filename": c["filename"],
                    "page": c["page"],
                    "chunk_id": c["chunk_id"],
                    "document_type": c["document_type"],
                    "section": c["section"],
                    "heading": c["heading"],
                    "section_number": c.get("section_number", ""),
                    "previous_chunk": c["previous_chunk"] or "",
                    "next_chunk": c["next_chunk"] or "",
                    "machine_id": machine_id
                }
            )
            for c in enriched_chunks
        ]
        vector_store.add_documents(documents)
        print(f"RAG Proxy: Successfully uploaded {len(documents)} chunks to Pinecone.")
    except Exception as e:
        print(f"RAG Proxy: Failed to upload to Pinecone: {e}")
        
    # 4. Save to local chunks cache (for BM25 indexing & context expansion)
    from rag.retriever import load_local_chunks, clear_bm25_cache
    local_chunks = load_local_chunks(machine_id)
    # Remove existing chunks for this file if we are re-uploading
    local_chunks = [c for c in local_chunks if c["filename"] != filename]
    local_chunks.extend(enriched_chunks)
    
    cache_dir = STORE_DIR / machine_id
    cache_dir.mkdir(parents=True, exist_ok=True)
    cache_path = cache_dir / "chunks_cache.json"
    with open(cache_path, "w", encoding="utf-8") as f:
        json.dump(local_chunks, f, indent=2, ensure_ascii=False)

    try:
        from database import save_chunks_to_db
        save_chunks_to_db(machine_id, enriched_chunks)
    except Exception as e:
        print(f"RAG Proxy: Failed to save chunks to PostgreSQL: {e}", flush=True)

    # 5. Clear the BM25 cache for this machine to force rebuild on next query
    clear_bm25_cache(machine_id)

    # 6. Invalidate cached chat answers for this machine — a cached "not found" or a
    # cached answer built from the old chunk set must not outlive this new upload.
    from rag.engine import clear_query_cache
    clear_query_cache(machine_id)


def get_pinecone_store(machine_id: str):
    """
    Retrieves the Pinecone Vector Store for the machine namespace.
    """
    return retriever_get_pinecone_store(machine_id)

def retrieve(query: str, machine_id: str) -> List[Dict[str, Any]]:
    """
    Retrieves, reranks, and expands adjacent context in correct reading order.
    Used by the AI task and alert extraction endpoints in main.py.
    """
    from rag.classifier import classify_question
    from rag.query_expander import expand_query
    from rag.retriever import expand_context
    
    # 1. Classify query intent
    intent = classify_question(query)
    
    # 2. Expand queries
    queries = expand_query(query)
    
    # 3. Perform hybrid retrieval
    merged = retrieve_hybrid(queries, machine_id, intent)
    if not merged:
        return []
        
    # 4. Rerank to top 12 (increased recall to capture entire tables)
    top_reranked = rerank(query, merged, top_n=12)
    
    # 5. Run context expansion to pull adjacent table parts and sort them in correct page order
    expanded = expand_context(top_reranked, machine_id)
    
    # 6. Map back to the exact format main.py expects (in correct reading sequence)
    results = []
    for idx, c in enumerate(expanded):
        results.append({
            "id": idx,
            "text": c["text"],
            "filename": c["filename"],
            "score": c.get("score", 0.0),
            "page": c["page"]
        })
    return results

async def agent_stream(message: str, machine_id: str, history: List[Dict[str, Any]]):
    """
    Streams RAG v2 responses using SSE (Proxy to rag/engine.py).
    """
    async for event in engine_agent_stream(message, machine_id, history):
        yield event

	from typing import List, Dict, Any
	from rag.ocr import extract_text as ocr_extract_text
	from rag.chunker import chunk_text as chunker_chunk_text
	from rag.retriever import get_pinecone_store as retriever_get_pinecone_store, retrieve_hybrid
	from rag.reranker import rerank
	from rag.engine import agent_stream as engine_agent_stream

	# Re-expose the system prompt for any direct imports (though now managed in rag/prompt.py)
	from rag.prompt import SYSTEM_PROMPT

	def extract_text(file_bytes: bytes) -> List[Dict[str, Any]]:
	"""
	Extracts text from PDF pages (Proxy to rag/ocr.py).
	"""
	return ocr_extract_text(file_bytes)

	def chunk_text(pages: List[Dict[str, Any]], filename: str) -> List[Dict[str, Any]]:
	"""
	Splits pages into recursive/semantic chunks (Proxy to rag/chunker.py).
	"""
	return chunker_chunk_text(pages, filename)

	def save_to_pinecone(machine_id: str, chunks: List[Dict[str, Any]]):
	"""
	Processes raw chunks (classification, metadata generation, sequential linking)
	and saves them to both Pinecone and the local cache (Proxy to rag/ingestion.py logic).
	"""
	if not chunks:
	return

	filename = chunks[0]["filename"]
	# Use first chunk text as sample for classification
	sample_text = chunks[0]["text"] if chunks else ""

	from rag.classifier import classify_document
	from rag.metadata import generate_metadata
	from langchain_core.documents import Document
	from config import STORE_DIR
	import json

	# 1. Automatically classify document type
	doc_type = classify_document(filename, sample_text)
	print(f"RAG Proxy: Classified '{filename}' as '{doc_type}'")

	# 2. Generate heading detection and sequential links
	enriched_chunks = generate_metadata(chunks, machine_id, doc_type)

	# 3. Upload to Pinecone
	try:
	vector_store = retriever_get_pinecone_store(machine_id)
	documents = [
	Document(
	page_content=c["text"],
	metadata={
	"filename": c["filename"],
	"page": c["page"],
	"chunk_id": c["chunk_id"],
	"document_type": c["document_type"],
	"section": c["section"],
	"heading": c["heading"],
	"section_number": c.get("section_number", ""),
	"previous_chunk": c["previous_chunk"] or "",
	"next_chunk": c["next_chunk"] or "",
	"machine_id": machine_id
	}
	)
	for c in enriched_chunks
	]
	vector_store.add_documents(documents)
	print(f"RAG Proxy: Successfully uploaded {len(documents)} chunks to Pinecone.")
	except Exception as e:
	print(f"RAG Proxy: Failed to upload to Pinecone: {e}")

	# 4. Save to local chunks cache (for BM25 indexing & context expansion)
	from rag.retriever import load_local_chunks, clear_bm25_cache
	local_chunks = load_local_chunks(machine_id)
	# Remove existing chunks for this file if we are re-uploading
	local_chunks = [c for c in local_chunks if c["filename"] != filename]
	local_chunks.extend(enriched_chunks)

	cache_dir = STORE_DIR / machine_id
	cache_dir.mkdir(parents=True, exist_ok=True)
	cache_path = cache_dir / "chunks_cache.json"
	with open(cache_path, "w", encoding="utf-8") as f:
	json.dump(local_chunks, f, indent=2, ensure_ascii=False)

	try:
	from database import save_chunks_to_db
	save_chunks_to_db(machine_id, enriched_chunks)
	except Exception as e:
	print(f"RAG Proxy: Failed to save chunks to PostgreSQL: {e}", flush=True)

	# 5. Clear the BM25 cache for this machine to force rebuild on next query
	clear_bm25_cache(machine_id)

	# 6. Invalidate cached chat answers for this machine — a cached "not found" or a
	# cached answer built from the old chunk set must not outlive this new upload.
	from rag.engine import clear_query_cache
	clear_query_cache(machine_id)


	def get_pinecone_store(machine_id: str):
	"""
	Retrieves the Pinecone Vector Store for the machine namespace.
	"""
	return retriever_get_pinecone_store(machine_id)

	def retrieve(query: str, machine_id: str) -> List[Dict[str, Any]]:
	"""
	Retrieves, reranks, and expands adjacent context in correct reading order.
	Used by the AI task and alert extraction endpoints in main.py.
	"""
	from rag.classifier import classify_question
	from rag.query_expander import expand_query
	from rag.retriever import expand_context

	# 1. Classify query intent
	intent = classify_question(query)

	# 2. Expand queries
	queries = expand_query(query)

	# 3. Perform hybrid retrieval
	merged = retrieve_hybrid(queries, machine_id, intent)
	if not merged:
	return []

	# 4. Rerank to top 12 (increased recall to capture entire tables)
	top_reranked = rerank(query, merged, top_n=12)

	# 5. Run context expansion to pull adjacent table parts and sort them in correct page order
	expanded = expand_context(top_reranked, machine_id)

	# 6. Map back to the exact format main.py expects (in correct reading sequence)
	results = []
	for idx, c in enumerate(expanded):
	results.append({
	"id": idx,
	"text": c["text"],
	"filename": c["filename"],
	"score": c.get("score", 0.0),
	"page": c["page"]
	})
	return results

	async def agent_stream(message: str, machine_id: str, history: List[Dict[str, Any]]):
	"""
	Streams RAG v2 responses using SSE (Proxy to rag/engine.py).
	"""
	async for event in engine_agent_stream(message, machine_id, history):
	yield event