Spaces:
Running
Running
changes
Browse files- agent.py +235 -0
- test_agent.py +64 -0
agent.py
ADDED
@@ -0,0 +1,235 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import re
|
3 |
+
import requests
|
4 |
+
import json
|
5 |
+
from typing import Tuple, List
|
6 |
+
|
7 |
+
from omegaconf import OmegaConf
|
8 |
+
|
9 |
+
from typing import Optional
|
10 |
+
from pydantic import Field, BaseModel
|
11 |
+
|
12 |
+
from vectara_agent.agent import Agent
|
13 |
+
from vectara_agent.tools import ToolsFactory, VectaraToolFactory
|
14 |
+
from vectara_agent.tools_catalog import summarize_text
|
15 |
+
|
16 |
+
from dotenv import load_dotenv
|
17 |
+
load_dotenv(override=True)
|
18 |
+
|
19 |
+
citation_description = '''
|
20 |
+
The citation for a particular case.
|
21 |
+
Citation must include the volume number, reporter, and first page. For example: 253 P.2d 136.
|
22 |
+
'''
|
23 |
+
|
24 |
+
def extract_components_from_citation(citation: str) -> Tuple[int, str, int]:
|
25 |
+
citation_components = citation.split(' ')
|
26 |
+
volume_num = citation_components[0]
|
27 |
+
reporter = '-'.join(citation_components[1:-1]).replace('.', '').lower()
|
28 |
+
first_page = citation_components[-1]
|
29 |
+
|
30 |
+
if not volume_num.isdigit():
|
31 |
+
raise ValueError("volume number must be a number.")
|
32 |
+
if not first_page.isdigit():
|
33 |
+
raise ValueError("first page number must be a number.")
|
34 |
+
|
35 |
+
return int(volume_num), reporter, int(first_page)
|
36 |
+
|
37 |
+
def create_assistant_tools(cfg):
|
38 |
+
|
39 |
+
def get_opinion_text(
|
40 |
+
case_citation = Field(description = citation_description),
|
41 |
+
summarize: Optional[bool] = False
|
42 |
+
) -> str:
|
43 |
+
"""
|
44 |
+
Given case citation, returns the full opinion/ruling text of the case.
|
45 |
+
if summarize is True, the text is summarized.
|
46 |
+
If there is more than one opinion for the case, the type of each opinion is returned with the text,
|
47 |
+
and the opinions (or their summaries) are separated by semicolons (;)
|
48 |
+
"""
|
49 |
+
volume_num, reporter, first_page = extract_components_from_citation(case_citation)
|
50 |
+
response = requests.get(f"https://static.case.law/{reporter}/{volume_num}/cases/{first_page:04d}-01.json")
|
51 |
+
if response.status_code != 200:
|
52 |
+
return "Case not found; please check the citation."
|
53 |
+
res = json.loads(response.text)
|
54 |
+
|
55 |
+
if len(res["casebody"]["opinions"]) == 1:
|
56 |
+
text = res["casebody"]["opinions"][0]["text"]
|
57 |
+
output = text if not summarize else summarize_text(text, "law")
|
58 |
+
else:
|
59 |
+
output = ""
|
60 |
+
for opinion in res["casebody"]["opinions"]:
|
61 |
+
text = opinion["text"] if not summarize else summarize_text(opinion["text"], "law")
|
62 |
+
output += f"Opinion type: {opinion['type']}, text: {text};"
|
63 |
+
|
64 |
+
return output
|
65 |
+
|
66 |
+
def get_case_document_pdf(
|
67 |
+
case_citation = Field(description = citation_description)
|
68 |
+
) -> str:
|
69 |
+
"""
|
70 |
+
Given a case citation, returns a valid web url to a pdf of the case record
|
71 |
+
"""
|
72 |
+
volume_num, reporter, first_page = extract_components_from_citation(case_citation)
|
73 |
+
response = requests.get(f"https://static.case.law/{reporter}/{volume_num}/cases/{first_page:04d}-01.json")
|
74 |
+
if response.status_code != 200:
|
75 |
+
return "Case not found; please check the citation."
|
76 |
+
res = json.loads(response.text)
|
77 |
+
page_number = res["first_page_order"]
|
78 |
+
return f"https://static.case.law/{reporter}/{volume_num}.pdf#page={page_number}"
|
79 |
+
|
80 |
+
def get_case_document_page(
|
81 |
+
case_citation = Field(description = citation_description)
|
82 |
+
) -> str:
|
83 |
+
"""
|
84 |
+
Given a case citation, returns a valid web url to a page with information about the case.
|
85 |
+
"""
|
86 |
+
volume_num, reporter, first_page = extract_components_from_citation(case_citation)
|
87 |
+
url = f"https://case.law/caselaw/?reporter={reporter}&volume={volume_num}&case={first_page:04d}-01"
|
88 |
+
response = requests.get(url)
|
89 |
+
if response.status_code != 200:
|
90 |
+
return "Case not found; please check the citation."
|
91 |
+
return url
|
92 |
+
|
93 |
+
def get_case_name(
|
94 |
+
case_citation = Field(description = citation_description)
|
95 |
+
) -> Tuple[str, str]:
|
96 |
+
"""
|
97 |
+
Given a case citation, returns its name and name abbreviation.
|
98 |
+
"""
|
99 |
+
volume_num, reporter, first_page = extract_components_from_citation(case_citation)
|
100 |
+
response = requests.get(f"https://static.case.law/{reporter}/{volume_num}/cases/{first_page:04d}-01.json")
|
101 |
+
if response.status_code != 200:
|
102 |
+
return "Case not found", "Case not found"
|
103 |
+
res = json.loads(response.text)
|
104 |
+
return res["name"], res["name_abbreviation"]
|
105 |
+
|
106 |
+
def get_cited_cases(
|
107 |
+
case_citation = Field(description = citation_description)
|
108 |
+
) -> List[dict]:
|
109 |
+
"""
|
110 |
+
Given a case citation, returns a list of cases that are cited by the opinion of this case.
|
111 |
+
The output is a list of cases, each a dict with the citation, name and name_abbreviation of the case.
|
112 |
+
"""
|
113 |
+
volume_num, reporter, first_page = extract_components_from_citation(case_citation)
|
114 |
+
response = requests.get(f"https://static.case.law/{reporter}/{volume_num}/cases/{first_page:04d}-01.json")
|
115 |
+
if response.status_code != 200:
|
116 |
+
return "Case not found; please check the citation."
|
117 |
+
res = json.loads(response.text)
|
118 |
+
citations = res["cites_to"]
|
119 |
+
res = []
|
120 |
+
for citation in citations[:10]:
|
121 |
+
name, name_abbreviation = get_case_name(citation["cite"])
|
122 |
+
res.append({
|
123 |
+
"citation": citation["cite"],
|
124 |
+
"name": name,
|
125 |
+
"name_abbreviation": name_abbreviation
|
126 |
+
})
|
127 |
+
return res
|
128 |
+
|
129 |
+
def validate_url(
|
130 |
+
url = Field(description = "A web url pointing to case-law document")
|
131 |
+
) -> str:
|
132 |
+
"""
|
133 |
+
Given a link, returns whether or not the link is valid.
|
134 |
+
If it is not valid, it should not be used in any output.
|
135 |
+
"""
|
136 |
+
pdf_pattern = re.compile(r'^https://static.case.law/.*')
|
137 |
+
document_pattern = re.compile(r'^https://case.law/caselaw/?reporter=.*')
|
138 |
+
return "URL is valid" if bool(pdf_pattern.match(url)) | bool(document_pattern.match(url)) else "URL is bad"
|
139 |
+
|
140 |
+
class QueryCaselawArgs(BaseModel):
|
141 |
+
query: str = Field(..., description="The user query.")
|
142 |
+
citations: Optional[str] = Field(default = None,
|
143 |
+
description = "The citation of the case. Optional.",
|
144 |
+
examples = ['253 P.2d 136', '10 Alaska 11', '6 C.M.A. 3'])
|
145 |
+
|
146 |
+
vec_factory = VectaraToolFactory(vectara_api_key=cfg.api_key,
|
147 |
+
vectara_customer_id=cfg.customer_id,
|
148 |
+
vectara_corpus_id=cfg.corpus_id)
|
149 |
+
tools_factory = ToolsFactory()
|
150 |
+
|
151 |
+
ask_caselaw = vec_factory.create_rag_tool(
|
152 |
+
tool_name = "ask_caselaw",
|
153 |
+
tool_description = """
|
154 |
+
Returns a response (str) to the user query base on case law in the state of Alaska.
|
155 |
+
If 'citations' is provided, filters the response based on information from that case.
|
156 |
+
The response includes metadata about the case such as title/name the ruling, the court,
|
157 |
+
the decision date, the judges, and the case citation.
|
158 |
+
You can use case citations from the metadata as input to other tools.
|
159 |
+
Use this tool for general case law queries.
|
160 |
+
""",
|
161 |
+
tool_args_schema = QueryCaselawArgs,
|
162 |
+
reranker = "multilingual_reranker_v1", rerank_k = 100,
|
163 |
+
n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.0,
|
164 |
+
summary_num_results = 10,
|
165 |
+
vectara_summarizer = 'vectara-summary-ext-24-05-med-omni',
|
166 |
+
include_citations = False,
|
167 |
+
)
|
168 |
+
|
169 |
+
return (
|
170 |
+
[tools_factory.create_tool(tool) for tool in [
|
171 |
+
get_opinion_text,
|
172 |
+
get_case_document_pdf,
|
173 |
+
get_case_document_page,
|
174 |
+
get_cited_cases,
|
175 |
+
get_case_name,
|
176 |
+
validate_url
|
177 |
+
]] +
|
178 |
+
tools_factory.standard_tools() +
|
179 |
+
tools_factory.legal_tools() +
|
180 |
+
tools_factory.guardrail_tools() +
|
181 |
+
[ask_caselaw]
|
182 |
+
)
|
183 |
+
|
184 |
+
def get_agent_config() -> OmegaConf:
|
185 |
+
cfg = OmegaConf.create({
|
186 |
+
'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
|
187 |
+
'corpus_id': str(os.environ['VECTARA_CORPUS_ID']),
|
188 |
+
'api_key': str(os.environ['VECTARA_API_KEY']),
|
189 |
+
'examples': os.environ.get('QUERY_EXAMPLES', None),
|
190 |
+
'title': "Legal Assistant",
|
191 |
+
'demo_welcome': "Welcome to the Legal Assistant demo.",
|
192 |
+
'demo_description': "This demo can help you prepare for a court case by providing you information about past court cases in Alaska.",
|
193 |
+
})
|
194 |
+
return cfg
|
195 |
+
|
196 |
+
def initialize_agent(_cfg, update_func=None):
|
197 |
+
|
198 |
+
legal_assistant_instructions = """
|
199 |
+
- You are a helpful legal assistant, with expertise in case law for the state of Alaska.
|
200 |
+
- The ask_caselaw tool is your primary tools for finding information about cases.
|
201 |
+
Do not use your own knowledge to answer questions.
|
202 |
+
- For a query with multiple sub-questions, break down the query into the sub-questions,
|
203 |
+
and make separate calls to the ask_caselaw tool to answer each sub-question,
|
204 |
+
then combine the answers to provide a complete response.
|
205 |
+
- If the ask_caselaw tool responds that it does not have enough information to answer the query,
|
206 |
+
try to rephrase the query and call the tool again.
|
207 |
+
- When presenting the output from ask_caselaw tool,
|
208 |
+
Extract metadata from the tool's response, and respond in this format:
|
209 |
+
'On <decision date>, the <court> ruled in <case name> that <judges ruling>. This opinion was authored by <judges>'.
|
210 |
+
- Citations include 3 components: volume number, reporter, and first page.
|
211 |
+
Here are some examples: '253 P.2d 136', '10 Alaska 11', '6 C.M.A. 3'
|
212 |
+
Never use your internal knowledge to contruct or guess what the citation is.
|
213 |
+
- If two cases have conflicting rulings, assume that the case with the more current ruling date is correct.
|
214 |
+
- If the response is based on cases that are older than 5 years, make sure to inform the user that the information may be outdated,
|
215 |
+
since some case opinions may no longer apply in law.
|
216 |
+
- To summarize the case, use the get_opinion_text with summarize set to True.
|
217 |
+
- If a user wants to learn more about a case, you can call the get_case_document_pdf tool with the citation to get a valid URL.
|
218 |
+
If this is unsuccessful, call the get_case_document_page tool instead.
|
219 |
+
The text displayed with this URL should be the name_abbreviation of the case (DON'T just say the info can be found here).
|
220 |
+
Don't call the get_case_document_page tool until after you have tried the get_case_document_pdf tool.
|
221 |
+
Don't provide URLs from any other tools. Do not generate URLs yourself.
|
222 |
+
- When presenting a URL in your response, use the validate_url tool.
|
223 |
+
- If a user wants to test their argument, use the ask_caselaw tool to gather information about cases related to their argument
|
224 |
+
and the critique_as_judge tool to determine whether their argument is sound or has issues that must be corrected.
|
225 |
+
- Never discuss politics, and always respond politely.
|
226 |
+
"""
|
227 |
+
|
228 |
+
agent = Agent(
|
229 |
+
tools=create_assistant_tools(_cfg),
|
230 |
+
topic="Case law in Alaska",
|
231 |
+
custom_instructions=legal_assistant_instructions,
|
232 |
+
update_func=update_func
|
233 |
+
)
|
234 |
+
|
235 |
+
return agent
|
test_agent.py
ADDED
@@ -0,0 +1,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
|
2 |
+
import unittest
|
3 |
+
import re
|
4 |
+
import os
|
5 |
+
|
6 |
+
from omegaconf import OmegaConf
|
7 |
+
from vectara_agent.agent import Agent
|
8 |
+
|
9 |
+
from app import initialize_agent
|
10 |
+
|
11 |
+
from dotenv import load_dotenv
|
12 |
+
load_dotenv(override=True)
|
13 |
+
|
14 |
+
class TestAgentResponses(unittest.TestCase):
|
15 |
+
|
16 |
+
def test_responses(self):
|
17 |
+
|
18 |
+
cfg = OmegaConf.create({
|
19 |
+
'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
|
20 |
+
'corpus_id': str(os.environ['VECTARA_CORPUS_ID']),
|
21 |
+
'api_key': str(os.environ['VECTARA_API_KEY']),
|
22 |
+
'examples': os.environ.get('QUERY_EXAMPLES', None)
|
23 |
+
})
|
24 |
+
|
25 |
+
agent = initialize_agent(_cfg=cfg)
|
26 |
+
self.assertIsInstance(agent, Agent)
|
27 |
+
|
28 |
+
# Test whether cases are real or fake
|
29 |
+
self.assertEqual(''.join(re.findall(r'[a-zA-Z]', agent.chat('Is the case Brown v. Board of Education, 347 U.S. 483 (1954), a real case? Say "yes" or "no" only.'))).lower(), 'yes')
|
30 |
+
self.assertEqual(''.join(re.findall(r'[a-zA-Z]', agent.chat('Is the case Bowers v. Hardwick, 478 U.S. 186 (1986), a real case? Say "yes" or "no" only.'))).lower(), 'yes')
|
31 |
+
self.assertEqual(''.join(re.findall(r'[a-zA-Z]', agent.chat('Is the case Columbia University v. Rodham, 564 U.S. 911 (2010), a real case? Say "yes" or "no" only.'))).lower(), 'no')
|
32 |
+
|
33 |
+
# Test case citation extraction
|
34 |
+
self.assertEqual(agent.chat('What is the citation for the case Brown v. Board of Education? Provide ONLY the citation in "<volume>, <reporter>, <page>" format, nothing else.'), '347 U.S. 483')
|
35 |
+
self.assertEqual(agent.chat('What is the citation for the case Bowers v. Hardwick? Provide ONLY the citation in "<volume>, <reporter>, <page>" format, nothing else.'), '478 U.S. 186')
|
36 |
+
self.assertEqual(agent.chat('What is the citation for the case McCulloch v. Maryland? Provide ONLY the citation in "<volume>, <reporter>, <page>" format, nothing else.'), '17 U.S. 316')
|
37 |
+
|
38 |
+
# Test opinion author identification
|
39 |
+
self.assertEqual(agent.chat('Who wrote the majority opinion in Brown v. Board of Education, 347 U.S. 483 (1954)? Provide the first and the last name of the judge ONLY.'), 'Earl Warren')
|
40 |
+
self.assertEqual(agent.chat('Who wrote the majority opinion in Bowers v. Hardwick, 478 U.S. 186 (1986)? Provide the first and the last name of the judge ONLY.'), 'Byron White')
|
41 |
+
self.assertEqual(agent.chat('Who wrote the majority opinion in McCulloch v. Maryland, 17 U.S. 316 (1819)? Provide the first and the last name of the judge ONLY.'), 'John Marshall')
|
42 |
+
|
43 |
+
# Test opinion text understanding
|
44 |
+
self.assertEqual(agent.chat("Did the court in Plessy v. Ferguson, 163 U.S. 537 (1896) affirm or reverse the lower court's decision? Say 'affirm' or 'reverse' only.").lower(), 'affirm')
|
45 |
+
self.assertEqual(agent.chat("Did the court in Bowers v. Hardwick, 478 U.S. 186 (1986) affirm or reverse the lower court's decision? Say 'affirm' or 'reverse' only.").lower(), 'reverse')
|
46 |
+
self.assertEqual(agent.chat("Did the court in McCulloch v. Maryland, 17 U.S. 316 (1819) affirm or reverse the lower court's decision? Say 'affirm' or 'reverse' only.").lower(), 'reverse')
|
47 |
+
|
48 |
+
# Test court identification
|
49 |
+
self.assertIn('united states court of appeals for the second circuit', agent.chat("Which court decided the case Viacom International Inc. v. YouTube, Inc., 676 F.3d 19 (2012)? Provide the name of the court ONLY, nothing else.").lower())
|
50 |
+
self.assertIn('united states court of appeals for the district of columbia circuit', agent.chat("Which court decided the case Durham v. United States, 214 F.2d 862 (1954)? Provide the name of the court ONLY, nothing else.").lower())
|
51 |
+
self.assertIn('supreme court', agent.chat("Which court decided the case Bowers v. Hardwick (1986)? Provide the name of the court ONLY, nothing else.").lower())
|
52 |
+
|
53 |
+
# Test overruling of case
|
54 |
+
self.assertIn(agent.chat("What year was Whitney v. California, 274 U.S. 357, overruled? Provide the year only."), ['1969', 'I don\'t know.']) # Our agent seems to not find the answer to this question, which I don't see as a problem (At least it's not hallucinating)
|
55 |
+
self.assertEqual(agent.chat("What year was Austin v. Michigan Chamber of Commerce, 494 U.S. 652, overruled? Provide the year only."), '2010')
|
56 |
+
|
57 |
+
# Compare two rulings
|
58 |
+
self.assertEqual(agent.chat('Do the cases Brown v. Board of Education, 347 U.S. 483 (1954) and Plessy v. Ferguson, 163 U.S. 537 (1896) agree or disagree with each other? Say "agree" or "disagree" only.').lower(), 'disagree')
|
59 |
+
# self.assertEqual(agent.chat('Do the cases Youngstown Sheet & Tube Co. v. Sawyer, 343 U.S. 579 (1952) and Medellin v. Texas, 552 U.S. 491 (2008) agree or disagree with each other? Say "agree" or "disagree" only.').lower(), 'agree') # Our agent thinks that these rulings disagree, so I commented out this test.
|
60 |
+
self.assertEqual(agent.chat('Do the cases Whitney v. California, 274 U.S. 357 (1927) and Brandenburg v. Ohio, 395 U.S. 444 (1969) agree or disagree with each other? Say "agree" or "disagree" only.').lower(), 'disagree')
|
61 |
+
|
62 |
+
|
63 |
+
if __name__ == "__main__":
|
64 |
+
unittest.main()
|