update

agent.py

@@ -34,9 +34,17 @@ def extract_components_from_citation(citation: str) -> dict:
 
     return {'volume': int(volume_num), 'reporter': reporter, 'first_page': int(first_page)}
 
-def create_assistant_tools(cfg, agent_config):
+class AgentTools:
+    def __init__(self, _cfg, agent_config):
+        self.tools_factory = ToolsFactory()
+        self.agent_config = agent_config
+        self.cfg = _cfg
+        self.vec_factory = VectaraToolFactory(vectara_api_key=_cfg.api_key,
+                                              vectara_corpus_key=_cfg.corpus_key)
+
 
     def get_opinion_text(
+            self,
             case_citation: str = Field(description = citation_description),
             summarize: bool = Field(default=True, description="if True returns case summary, otherwise the full text of the case")
             ) -> str:
@@ -52,7 +60,7 @@ def create_assistant_tools(cfg, agent_config):
         citation_dict = extract_components_from_citation(case_citation)
         if not citation_dict:
             return f"Citation is invalid: {case_citation}."
-        summarize_text = ToolsCatalog(agent_config).summarize_text
+        summarize_text = ToolsCatalog(self.agent_config).summarize_text
         reporter = citation_dict['reporter']
         volume_num = citation_dict['volume']
         first_page = citation_dict['first_page']
@@ -73,6 +81,7 @@ def create_assistant_tools(cfg, agent_config):
         return output
 
     def get_case_document_pdf(
+            self,
             case_citation = Field(description = citation_description)
             ) -> str:
         """
@@ -92,6 +101,7 @@ def create_assistant_tools(cfg, agent_config):
         return f"https://static.case.law/{reporter}/{volume_num}.pdf#page={page_number}"
 
     def get_case_document_page(
+            self,
             case_citation = Field(description = citation_description)
             ) -> str:
         """
@@ -110,6 +120,7 @@ def create_assistant_tools(cfg, agent_config):
         return url
         
     def get_case_name(
+            self,
             case_citation = Field(description = citation_description)
             ) -> Tuple[str, str]:
         """
@@ -128,6 +139,7 @@ def create_assistant_tools(cfg, agent_config):
         return res["name"], res["name_abbreviation"]
 
     def get_cited_cases(
+            self,
             case_citation = Field(description = citation_description)
             ) -> List[dict]:
         """
@@ -147,7 +159,7 @@ def create_assistant_tools(cfg, agent_config):
         citations = res["cites_to"]
         res = []
         for citation in citations[:10]:
-            name, name_abbreviation = get_case_name(citation["cite"])
+            name, name_abbreviation = self.get_case_name(citation["cite"])
             res.append({
                 "citation": citation["cite"],
                 "name": name,
@@ -156,6 +168,7 @@ def create_assistant_tools(cfg, agent_config):
         return res
 
     def validate_url(
+            self,
             url = Field(description = "A web url pointing to case-law document")
         ) -> str:
         """
@@ -166,50 +179,50 @@ def create_assistant_tools(cfg, agent_config):
         document_pattern = re.compile(r'^https://case.law/caselaw/?reporter=.*')
         return "URL is valid" if bool(pdf_pattern.match(url)) | bool(document_pattern.match(url)) else "URL is bad"
 
-    class QueryCaselawArgs(BaseModel):
-        query: str = Field(..., description="The user query.")
+    def get_tools(self):
+        class QueryCaselawArgs(BaseModel):
+            query: str = Field(..., description="The user query.")
 
-    vec_factory = VectaraToolFactory(vectara_api_key=cfg.api_key,
-                                     vectara_corpus_key=cfg.corpus_key)
-    summarizer = 'vectara-experimental-summary-ext-2023-12-11-med-omni'
+        vec_factory = VectaraToolFactory(vectara_api_key=cfg.api_key,
+                                        vectara_corpus_key=cfg.corpus_key)
+        summarizer = 'vectara-experimental-summary-ext-2023-12-11-med-omni'
 
-    ask_caselaw = vec_factory.create_rag_tool(
-        tool_name = "ask_caselaw",
-        tool_description = "A tool for asking questions about case law in Alaska. ",
-        tool_args_schema = QueryCaselawArgs,
-        reranker = "chain", rerank_k = 100,
-        rerank_chain = [
-            {
-                "type": "slingshot",
-                "cutoff": 0.2
-            },
-            {
-                "type": "mmr",
-                "diversity_bias": 0.1
-            },
-            {
-                "type": "userfn",
-                "user_function": "max(1000 * get('$.score') - hours(seconds(to_unix_timestamp(now()) - to_unix_timestamp(datetime_parse(get('$.document_metadata.decision_date'), 'yyyy-MM-dd')))) / 24 / 365, 0)"
-            }
-        ],
-        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
-        summary_num_results = 15,
-        vectara_summarizer = summarizer,
-        include_citations = True,
-    )
+        ask_caselaw = vec_factory.create_rag_tool(
+            tool_name = "ask_caselaw",
+            tool_description = "A tool for asking questions about case law in Alaska. ",
+            tool_args_schema = QueryCaselawArgs,
+            reranker = "chain", rerank_k = 100,
+            rerank_chain = [
+                {
+                    "type": "slingshot",
+                    "cutoff": 0.2
+                },
+                {
+                    "type": "mmr",
+                    "diversity_bias": 0.1
+                },
+                {
+                    "type": "userfn",
+                    "user_function": "max(1000 * get('$.score') - hours(seconds(to_unix_timestamp(now()) - to_unix_timestamp(datetime_parse(get('$.document_metadata.decision_date'), 'yyyy-MM-dd')))) / 24 / 365, 0)"
+                }
+            ],
+            n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
+            summary_num_results = 15,
+            vectara_summarizer = summarizer,
+            include_citations = True,
+        )
 
-    tools_factory = ToolsFactory()
-    return (
-        [ask_caselaw] +
-        [tools_factory.create_tool(tool) for tool in [
-            get_opinion_text,
-            get_case_document_pdf,
-            get_case_document_page,
-            get_cited_cases,
-            get_case_name,
-            validate_url
-        ]] 
-    )
+        return (
+            [ask_caselaw] +
+            [self.tools_factory.create_tool(tool) for tool in [
+                self.get_opinion_text,
+                self.get_case_document_pdf,
+                self.get_case_document_page,
+                self.get_cited_cases,
+                self.get_case_name,
+                self.validate_url
+            ]] 
+        )
 
 def get_agent_config() -> OmegaConf:
     cfg = OmegaConf.create({
@@ -254,7 +267,7 @@ def initialize_agent(_cfg, agent_progress_callback=None):
     """
     agent_config = AgentConfig()
     agent = Agent(
-        tools=create_assistant_tools(_cfg, agent_config=agent_config),
+        tools=AgentTools(_cfg, agent_config).get_tools(),
         topic="Case law in Alaska",
         custom_instructions=legal_assistant_instructions,
         agent_progress_callback=agent_progress_callback,

requirements.txt

@@ -6,4 +6,4 @@ streamlit-feedback==0.1.3
 uuid==1.30
 langdetect==1.0.9
 langcodes==3.4.0
-vectara-agentic==0.2.0
+vectara-agentic==0.2.1