Spaces:
Running
Running
File size: 2,620 Bytes
388aa42 | 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 | """
RAG Retrieval Agent
Dedicated agent for vector database queries
Uses FAISS for local vector storage
"""
import json
from rag.scheme_vectorstore import load_scheme_vectorstore
from rag.exam_vectorstore import load_exam_vectorstore
def run_rag_agent(query: str, database: str = "schemes", k: int = 5) -> dict:
"""
Performs RAG retrieval from specified vectorstore
Args:
query: Search query
database: "schemes" or "exams"
k: Number of documents to retrieve
Returns:
Retrieved documents dictionary
"""
try:
if database == "schemes":
vectorstore = load_scheme_vectorstore()
elif database == "exams":
vectorstore = load_exam_vectorstore()
else:
return {
"error": f"Invalid database: {database}. Use 'schemes' or 'exams'",
"documents": []
}
# Similarity search
docs = vectorstore.similarity_search(query, k=k)
# Format results
formatted_docs = []
for i, doc in enumerate(docs):
formatted_docs.append({
"id": i + 1,
"content": doc.page_content,
"metadata": doc.metadata,
"source": doc.metadata.get('source', 'Unknown')
})
return {
"query": query,
"database": database,
"documents_found": len(formatted_docs),
"documents": formatted_docs
}
except FileNotFoundError as e:
return {
"error": f"Vectorstore not found for {database}. Please build it first.",
"documents": []
}
except Exception as e:
return {
"error": str(e),
"documents": []
}
def hybrid_rag_search(query: str, k: int = 3) -> dict:
"""
Searches both scheme and exam databases
Args:
query: Search query
k: Number of documents per database
Returns:
Combined results from both databases
"""
scheme_results = run_rag_agent(query, database="schemes", k=k)
exam_results = run_rag_agent(query, database="exams", k=k)
return {
"query": query,
"scheme_results": scheme_results,
"exam_results": exam_results
}
if __name__ == "__main__":
# Test the agent
result = run_rag_agent("agricultural schemes for farmers", database="schemes", k=3)
print(json.dumps(result, indent=2))
|