Spaces:
Sleeping
Sleeping
File size: 4,778 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 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 | from .base import BaseRAGTechnique
from ..services.embed_service import get_embedding
from ..utils.json_utils import extract_json_block, repair_json
from typing import List, Dict, Any
import json
class AgenticRAG(BaseRAGTechnique):
async def retrieve(self, query: str, document_id: str, top_k: int, **kwargs) -> List[Dict[str, Any]]:
# This is the "agent loop"
await self.emit("AGENT_INIT", "#22C55E", "Qwen3 agent ready with 4 tools")
all_collected_chunks = []
conversation_history = []
system_prompt = f"""
You are an intelligent RAG agent. You have access to a document (ID: {document_id}).
Your goal is to answer the user query: "{query}"
Available tools:
1. search_docs(query: str, top_k: int) -> list of chunks
2. filter_search(filters: dict, query: str) -> list of chunks. filters can include "page" or "section".
3. get_page(page_num: int) -> text of that page
4. finish(answer: str) -> finish with final answer
Respond ONLY with a JSON object:
{{
"thought": "your reasoning",
"tool": "tool_name",
"args": {{ ... }}
}}
"""
for i in range(5): # Max 5 iterations
await self.emit("PLAN", "#8B5CF6", f"Agent iteration {i+1}: Thinking...")
agent_prompt = f"{system_prompt}\n\nHistory: {json.dumps(conversation_history)}\n\nAction:"
response_text = self.llm.generate(agent_prompt)
try:
action_data = repair_json(extract_json_block(response_text))
thought = action_data.get("thought", "")
tool = action_data.get("tool", "")
args = action_data.get("args", {})
await self.emit("PLAN", "#8B5CF6", f"Thought: {thought[:100]}...")
if tool == "finish":
self.final_agent_answer = args.get("answer", "")
break
# Execute Tool
await self.emit("TOOL", "#D97706", f"Tool call: {tool}({json.dumps(args)})")
observation = ""
if tool == "search_docs":
q = args.get("query", query)
tk = args.get("top_k", top_k)
q_vec = get_embedding(q)
results = await self.supabase.vector_search(q_vec, document_id, self.user_id, tk)
all_collected_chunks.extend(results)
observation = f"Found {len(results)} chunks."
elif tool == "filter_search":
f = args.get("filters", {})
q = args.get("query", query)
matching_ids = await self.supabase.filter_chunk_ids(document_id, self.user_id, f)
if matching_ids:
q_vec = get_embedding(q)
results = await self.supabase.vector_search(q_vec, document_id, self.user_id, top_k, filter_chunk_ids=matching_ids)
all_collected_chunks.extend(results)
observation = f"Filtered search found {len(results)} chunks."
else:
observation = "No chunks matched the filters."
elif tool == "get_page":
p = args.get("page_num")
results = await self.supabase.filter_chunk_ids(document_id, self.user_id, {"page": p})
if results:
chunks = await self.supabase.get_chunks_by_ids(results)
all_collected_chunks.extend(chunks)
observation = f"Retrieved page {p}."
else:
observation = f"Page {p} not found."
await self.emit("OBSERVE", "#22C55E", observation)
conversation_history.append({"action": action_data, "observation": observation})
except Exception as e:
logger.error(f"Agent error decoding JSON: {e}")
conversation_history.append({"error": f"Invalid JSON response from your side. Use the required JSON format. error: {str(e)}"})
await self.emit("FINAL", "#22C55E", "Answer generated after tool usage.")
return all_collected_chunks
async def generate(self, query: str, chunks: List[Dict[str, Any]]) -> str:
# If the agent finished with a final answer, use it.
if hasattr(self, "final_agent_answer") and self.final_agent_answer:
return self.final_agent_answer
await self.emit("GENERATE", "#7C3AED", "Qwen3 generating final summary...")
prompt = self.build_prompt(query, chunks)
return self.llm.generate(prompt)
|