Spaces:
Running
Running
| """Case query tool for smolagents. | |
| Allows juror agents to query any aspect of the case. | |
| """ | |
| from typing import TYPE_CHECKING | |
| from smolagents import Tool | |
| if TYPE_CHECKING: | |
| from case_db.index import CaseIndex | |
| from case_db.models import CriminalCase | |
| class CaseQueryTool(Tool): | |
| """Tool for agents to query any aspect of the case. | |
| General-purpose tool that allows agents to ask questions about | |
| the case - timeline, people, charges, context, connections between facts. | |
| Examples of queries an agent might make: | |
| - "What are the charges against the defendant?" | |
| - "Who discovered the body?" | |
| - "What is the defendant's background?" | |
| - "What time did the incident occur?" | |
| - "What is the connection between the witness and defendant?" | |
| """ | |
| name = "case_query" | |
| description = """Query any aspect of the case. Use this tool to: | |
| - Understand the timeline of events | |
| - Learn about people involved (defendant, witnesses, victims) | |
| - Clarify the charges and legal context | |
| - Find connections between facts | |
| - Get background information | |
| - Understand who said what | |
| This is a general-purpose case query tool. Ask any question about | |
| the case in natural language, such as "What happened on the night | |
| of the incident?" or "What is the relationship between the defendant | |
| and the victim?".""" | |
| inputs = { | |
| "query": { | |
| "type": "string", | |
| "description": "Question about the case (natural language)" | |
| } | |
| } | |
| output_type = "string" | |
| def __init__(self, case_index: "CaseIndex"): | |
| """Initialize with a case index. | |
| Args: | |
| case_index: LlamaIndex CaseIndex for semantic search | |
| """ | |
| super().__init__() | |
| self.case_index = case_index | |
| def forward(self, query: str) -> str: | |
| """Execute the case query. | |
| Args: | |
| query: Natural language question about the case | |
| Returns: | |
| Answer synthesized from case documents | |
| """ | |
| if self.case_index is None: | |
| return "Error: No case index available. Case query disabled." | |
| try: | |
| result = self.case_index.query(query) | |
| return result if result else "No information found for this query." | |
| except Exception as e: | |
| return f"Error querying case: {str(e)}" | |
| class CaseQueryToolMock(Tool): | |
| """Mock case query tool for testing without LlamaIndex. | |
| Returns case information based on simple pattern matching. | |
| """ | |
| name = "case_query" | |
| description = "Query case information (mock version for testing)" | |
| inputs = { | |
| "query": { | |
| "type": "string", | |
| "description": "Question about the case" | |
| } | |
| } | |
| output_type = "string" | |
| def __init__(self, case: "CriminalCase"): | |
| """Initialize with a case object. | |
| Args: | |
| case: CriminalCase object | |
| """ | |
| super().__init__() | |
| self.case = case | |
| def forward(self, query: str) -> str: | |
| """Simple pattern-based case query. | |
| Args: | |
| query: Question about the case | |
| Returns: | |
| Relevant case information | |
| """ | |
| if not self.case: | |
| return "No case information available." | |
| query_lower = query.lower() | |
| if "charge" in query_lower: | |
| return f"Charges: {', '.join(self.case.charges)}" | |
| if "defendant" in query_lower or "accused" in query_lower: | |
| if self.case.defendant: | |
| return f"Defendant: {self.case.defendant.name}. {self.case.defendant.background}" | |
| return "Defendant information not available." | |
| if "witness" in query_lower: | |
| if self.case.witnesses: | |
| summaries = [f"- {w.name} ({w.role}): {w.testimony_summary}" | |
| for w in self.case.witnesses] | |
| return "Witnesses:\n" + "\n".join(summaries) | |
| return "No witness information available." | |
| if "evidence" in query_lower: | |
| if self.case.evidence: | |
| summaries = [f"- [{e.type}] {e.description}" | |
| for e in self.case.evidence] | |
| return "Evidence:\n" + "\n".join(summaries) | |
| return "No evidence information available." | |
| if "summary" in query_lower or "what happened" in query_lower: | |
| return self.case.summary | |
| # Default: return case summary | |
| return f"Case: {self.case.title}\n\n{self.case.summary}" | |