chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,166 @@
|
||||
import os
|
||||
import json
|
||||
from typing import Dict, Any, List, Optional, Type
|
||||
from pydantic import BaseModel, Field
|
||||
from crewai.tools import BaseTool
|
||||
|
||||
|
||||
class RAGInput(BaseModel):
|
||||
"""Input schema for RAG search tool"""
|
||||
query: str = Field(..., description="The search query for retrieval.")
|
||||
top_k: int = Field(default=3, description="Maximum number of retrieved results to fetch.")
|
||||
document_paths: List[str] = Field(default=None, description="Optional list of document paths to load. Only needed if no documents are already loaded in the vector database.")
|
||||
|
||||
|
||||
class RAGTool(BaseTool):
|
||||
name: str = "rag_search"
|
||||
description: str = "Search through research documents for relevant information."
|
||||
rag_pipeline: Any = Field(..., description="RAG pipeline instance")
|
||||
args_schema: Type[BaseModel] = RAGInput
|
||||
|
||||
def _run(self, query: str, top_k: int = 3, document_paths: List[str] = None):
|
||||
try:
|
||||
doc_count = self.rag_pipeline.vector_db.get_collection_count()
|
||||
if doc_count == 0:
|
||||
if not document_paths:
|
||||
return json.dumps({
|
||||
"status": "INSUFFICIENT_CONTEXT",
|
||||
"source_used": "RAG",
|
||||
"answer": "No documents have been loaded into the RAG system. Please provide document_paths to load documents first, or ensure documents have been previously loaded.",
|
||||
"citations": [],
|
||||
"confidence": 0.0,
|
||||
"retrieval_metadata": {
|
||||
"retrieved_chunks": 0,
|
||||
"top_scores": [],
|
||||
"document_count": 0
|
||||
}
|
||||
})
|
||||
|
||||
# load documents
|
||||
load_result = self._load_documents(document_paths)
|
||||
if load_result["status"] == "ERROR":
|
||||
return json.dumps(load_result, indent=2)
|
||||
|
||||
doc_count = self.rag_pipeline.vector_db.get_collection_count()
|
||||
if doc_count == 0:
|
||||
return json.dumps({
|
||||
"status": "INSUFFICIENT_CONTEXT",
|
||||
"source_used": "RAG",
|
||||
"answer": "Failed to load documents into the RAG system.",
|
||||
"citations": [],
|
||||
"confidence": 0.0,
|
||||
"retrieval_metadata": {
|
||||
"retrieved_chunks": 0,
|
||||
"top_scores": [],
|
||||
"document_count": 0
|
||||
}
|
||||
})
|
||||
|
||||
# Retrieve relevant context (no generation)
|
||||
context_results = self.rag_pipeline.retrieve_context(query, top_k=top_k)
|
||||
if not context_results:
|
||||
return json.dumps({
|
||||
"status": "INSUFFICIENT_CONTEXT",
|
||||
"source_used": "RAG",
|
||||
"answer": f"No relevant context found for query: '{query}'",
|
||||
"citations": [],
|
||||
"confidence": 0.0,
|
||||
"retrieval_metadata": {
|
||||
"retrieved_chunks": 0,
|
||||
"top_scores": [],
|
||||
"document_count": doc_count
|
||||
}
|
||||
})
|
||||
|
||||
context_blocks = []
|
||||
citations = []
|
||||
for i, result in enumerate(context_results):
|
||||
chunk_text = result.get("text", "")
|
||||
score = result.get("score", 0.0)
|
||||
page_number = result.get("page_number", 0)
|
||||
chunk_index = result.get("chunk_index", i)
|
||||
source_file = result.get("source_file", "unknown")
|
||||
|
||||
filename = source_file.split("/")[-1] if source_file != "unknown" else "unknown"
|
||||
context_blocks.append(f"**Context {i+1} (Score: {score:.3f}, Page {page_number}, Chunk {chunk_index})**\n{chunk_text[:500]}...")
|
||||
citations.append({
|
||||
"label": f"{filename} - Page {page_number}, Chunk {chunk_index}",
|
||||
"locator": f"page_{page_number}_chunk_{chunk_index}",
|
||||
"page_number": page_number,
|
||||
"chunk_index": chunk_index,
|
||||
"source_file": filename,
|
||||
"score": score,
|
||||
"content": chunk_text
|
||||
})
|
||||
|
||||
answer = "\n\n".join(context_blocks)
|
||||
|
||||
return json.dumps({
|
||||
"status": "OK",
|
||||
"source_used": "RAG",
|
||||
"answer": f"Retrieved {len(context_results)} relevant context chunks:\n\n{answer}",
|
||||
"citations": citations,
|
||||
"confidence": max([r.get("score", 0.0) for r in context_results]),
|
||||
"retrieval_metadata": {
|
||||
"retrieved_chunks": len(context_results),
|
||||
"top_scores": [r.get("score", 0.0) for r in context_results],
|
||||
"document_count": doc_count
|
||||
},
|
||||
"raw_context": context_results # Include raw context for further processing
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return json.dumps({
|
||||
"status": "ERROR",
|
||||
"source_used": "RAG",
|
||||
"answer": f"RAG search failed: {str(e)}",
|
||||
"citations": [],
|
||||
"confidence": 0.0,
|
||||
"retrieval_metadata": {
|
||||
"retrieved_chunks": 0,
|
||||
"top_scores": [],
|
||||
"document_count": 0
|
||||
},
|
||||
"error": str(e)
|
||||
})
|
||||
|
||||
def _load_documents(self, document_paths: List[str]) -> Dict[str, Any]:
|
||||
"""Helper method to load documents into the RAG pipeline"""
|
||||
try:
|
||||
if not document_paths:
|
||||
return {
|
||||
"status": "ERROR",
|
||||
"source_used": "RAG",
|
||||
"answer": "No document paths provided",
|
||||
"citations": [],
|
||||
"confidence": 0.0
|
||||
}
|
||||
|
||||
missing_files = [path for path in document_paths if not os.path.exists(path)]
|
||||
if missing_files:
|
||||
return {
|
||||
"status": "ERROR",
|
||||
"source_used": "RAG",
|
||||
"answer": f"Missing files: {missing_files}",
|
||||
"citations": [],
|
||||
"confidence": 0.0
|
||||
}
|
||||
|
||||
results = self.rag_pipeline.process_documents(document_paths)
|
||||
return {
|
||||
"status": "OK",
|
||||
"source_used": "RAG",
|
||||
"answer": f"Successfully processed {len(results['processed_docs'])} documents with {results['total_chunks']} total chunks.",
|
||||
"citations": [{"label": f"Processed: {doc['path']}", "locator": doc['path']} for doc in results['processed_docs']],
|
||||
"confidence": 0.96
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "ERROR",
|
||||
"source_used": "RAG",
|
||||
"answer": f"Document processing failed: {str(e)}",
|
||||
"citations": [],
|
||||
"confidence": 0.0,
|
||||
"error": str(e)
|
||||
}
|
||||
Reference in New Issue
Block a user