Spaces:
Sleeping
Sleeping
File size: 6,020 Bytes
f97126e | 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | import contextvars
from typing import Literal, List
from langchain_core.tools import tool
from src.retriever import graph
from src.retriever.state import RetrievedNode
# ContextVar to collect retrieved nodes dynamically across tool executions for CLI/serve metadata
retrieved_nodes_var = contextvars.ContextVar("retrieved_nodes", default=None)
def _format_nodes(nodes: List[RetrievedNode]) -> str:
"""Helper to format retrieved nodes for LLM observation."""
if not nodes:
return "No matching legal sections or procedures found."
# Store in ContextVar if active
collected = retrieved_nodes_var.get()
if collected is not None:
for node in nodes:
if not any(x["node_id"] == node["node_id"] for x in collected):
collected.append(node)
formatted = []
for node in nodes:
fmt = (
f"ID: {node['node_id']}\n"
f"Act/Corpus: {node['act_code']}\n"
f"Title: {node['title']}\n"
f"Summary: {node.get('summary', '')}\n"
f"Content:\n{node.get('content', '')}\n"
f"Retrieval Method: {node.get('retrieval_method', 'unknown')}\n"
f"Score: {node.get('score', 1.0):.2f}"
)
formatted.append(fmt)
return "\n---\n" + "\n---\n".join(formatted)
@tool
async def search_statutes(
statute_code: Literal["BNS", "BNSS", "BSA"],
query: str,
method: Literal["tree", "bm25", "hybrid"] = "hybrid"
) -> str:
"""
Searches a specific Indian statutory act (BNS, BNSS, or BSA) for sections relevant to a query.
Use this when you need legal definitions, criminal offenses, punishments, or trial procedures.
Parameters:
- statute_code:
- 'BNS' (Bharatiya Nyaya Sanhita - Criminal offences, penalties, punishments, murder, robbery, etc.)
- 'BNSS' (Bharatiya Nagarik Suraksha Sanhita - Criminal trial procedures, arrests, bail, FIRs, investigations)
- 'BSA' (Bharatiya Sakshya Adhiniyam - Evidence act, witnesses, confessions, burden of proof)
- query: Specific keywords, section numbers, or legal scenario to search.
- method: Search methodology ('tree' for hierarchical navigation, 'bm25' for keyword search, 'hybrid' for both).
"""
if not graph._corpus_index:
return "Error: Corpus index not initialized. Ensure retriever.load() has been called."
nodes = []
seen_ids = set()
# 1. Tree Navigation (guided search)
if method in ["tree", "hybrid"] and graph._tree_navigator:
try:
tree_nodes = await graph._tree_navigator.navigate(query, statute_code)
for n in tree_nodes:
if n["node_id"] not in seen_ids:
nodes.append(n)
seen_ids.add(n["node_id"])
except Exception as e:
print(f"[Tool: search_statutes] Tree Nav failed: {e}")
# 2. BM25 (keyword search)
if method in ["bm25", "hybrid"] and graph._bm25_index:
try:
bm25_nodes = graph._bm25_index.search(
query,
graph._corpus_index,
top_k=5,
act_filter=[statute_code]
)
for n in bm25_nodes:
if n["node_id"] not in seen_ids:
nodes.append(n)
seen_ids.add(n["node_id"])
except Exception as e:
print(f"[Tool: search_statutes] BM25 failed: {e}")
return _format_nodes(nodes)
@tool
async def search_police_sop(query: str) -> str:
"""
Searches the Police Standard Operating Procedures (SOP) manual for operational guidelines,
patrol duties, checklists, timelines, and practical steps taken by police officers.
Use this when the query relates to how a police officer should register an FIR, conduct an arrest,
handle electronic evidence, or maintain a police station diary.
"""
if not graph._sop_retriever or not graph._corpus_index:
return "Error: SOP retriever not initialized."
try:
nodes = await graph._sop_retriever.retrieve(query, top_k=5)
return _format_nodes(nodes)
except Exception as e:
return f"Error searching Police SOP: {e}"
@tool
async def enrich_with_cross_references(section_id: str) -> str:
"""
Fetches other legal sections that are cross-referenced or linked to a specific section ID (e.g., 'BNSS_S35').
Use this when you have retrieved a section and want to follow its legal citations to other acts or sections
(for example, connecting a police procedure in the SOP to a section in the BNSS).
"""
if not graph._cross_ref_linker or not graph._corpus_index:
return "Error: Cross-reference linker not initialized."
node = graph._corpus_index.get_node(section_id)
if not node:
return f"Error: Legal section '{section_id}' not found in the index."
# Construct a skeleton RetrievedNode representation for the linker to consume
act_code = section_id.split("_")[0]
p_node: RetrievedNode = {
"node_id": section_id,
"act_code": act_code,
"title": node.get("title", ""),
"summary": node.get("summary", ""),
"content": node.get("content", ""),
"score": 1.0,
"node_type": node.get("node_type", "section"),
"page_range": node.get("metadata", {}).get("page_range", []),
"cross_act_refs": node.get("metadata", {}).get("cross_act_refs", []),
"internal_refs": node.get("metadata", {}).get("internal_refs", []),
"retrieval_method": "direct_lookup"
}
try:
enriched_nodes = graph._cross_ref_linker.enrich([p_node], max_links_per_node=5)
return _format_nodes(enriched_nodes)
except Exception as e:
return f"Error resolving cross references: {e}"
|