add the law tools

agent.json

@@ -1,6 +1,7 @@
 {
     "tools": [
-        "web_search",
+        "law_tool",
+        "law_rag_query",
         "final_answer"
     ],
     "model": {
@@ -39,7 +40,6 @@
     "name": null,
     "description": null,
     "requirements": [
-        "duckduckgo_search",
         "smolagents"
     ],
     "authorized_imports": [

app.py

@@ -5,7 +5,8 @@ from smolagents import GradioUI, CodeAgent, HfApiModel
 # Get current directory path
 CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
 
-from tools.web_search import DuckDuckGoSearchTool as WebSearch
+from tools.law_tool import LawTool
+from tools.law_rag_query import LawRAGQuery
 from tools.final_answer import FinalAnswerTool as FinalAnswer
 
 
@@ -15,7 +16,8 @@ model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
 provider=None,
 )
 
-web_search = WebSearch()
+law_tool = LawTool()
+law_rag_query = LawRAGQuery()
 final_answer = FinalAnswer()
 
 
@@ -24,7 +26,7 @@ with open(os.path.join(CURRENT_DIR, "prompts.yaml"), 'r') as stream:
 
 agent = CodeAgent(
     model=model,
-    tools=[web_search],
+    tools=[law_tool, law_rag_query],
     managed_agents=[],
     max_steps=20,
     verbosity_level=1,

requirements.txt

@@ -1,3 +1,8 @@
 duckduckgo_search
 smolagents
-gradio[oauth]==5.23.2
+gradio
+datasets
+langchain
+langchain-chroma
+langchain-text-splitters
+datasets

tools/law_rag_query.py

@@ -0,0 +1,54 @@
+from smolagents import Tool
+from langchain_chroma import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from datasets import load_dataset
+import os
+
+class LawRAGQuery(Tool):
+    name = "law_rag_query"
+    description = """
+    This is a tool that returns law content by input a question. It will find the related law and return."""
+    inputs = {
+        "question": {
+            "type": "string",
+            "description": "the question",
+        }
+    }
+    output_type = "array"
+    vectorstore = None
+
+    def __init__(self):
+        dataset = load_dataset("robin0307/law", split='train')
+        law = dataset.to_pandas()
+        self.vectorstore = self.get_vectorstore("thenlper/gte-large-zh", list(law['content']))
+        super().__init__()
+
+    def get_vectorstore(self, model_path, data_list, path="chroma_db"):
+        embeddings = HuggingFaceEmbeddings(model_name=model_path)
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=2048, chunk_overlap=50)
+        chunks = [text_splitter.split_text(text) for text in data_list]
+
+        # Flatten the list
+        if os.path.isdir(path):
+            vectorstore = Chroma(embedding_function=embeddings, persist_directory=path)
+        else:
+            splits = [chunk for sublist in chunks for chunk in sublist]
+            vectorstore = Chroma.from_texts(texts=splits, embedding=embeddings, persist_directory=path)
+        print("count:", vectorstore._collection.count())
+        return vectorstore
+
+    def get_docs(self, input, k=10):
+        retrieved_documents = self.vectorstore.similarity_search_with_score(input, k=50)
+
+        results = []
+        for i, (doc, score) in enumerate(retrieved_documents):
+            results.append((doc.page_content, score))
+            if i >= k:
+                break
+        return results
+
+    def forward(self, question: str):
+        docs = self.get_docs(question)
+        return docs
+    

tools/law_tool.py

@@ -0,0 +1,31 @@
+from smolagents import Tool
+from datasets import load_dataset
+
+class LawTool(Tool):
+    name = "law_tool"
+    description = """
+    This is a tool that returns law content by input the category and number."""
+    inputs = {
+        "category": {
+            "type": "string",
+            "description": "the law category (such as 民法, 中華民國刑法, 民事訴訟法, 刑事訴訟法, 律師法 etc)",
+        },
+        "number": {
+            "type": "integer",
+            "description": "the law number (such as 23)"
+        }
+    }
+    output_type = "string"
+    law = None
+
+    def __init__(self):
+        dataset = load_dataset("robin0307/law", split='train')
+        self.law = dataset.to_pandas()
+        super().__init__()
+
+    def forward(self, category: str, number: int):
+        if category == "刑法":
+            category = "中華民國刑法"
+        
+        data = self.law.loc[(self.law["category"]==category) & (self.law["number"]==number), "content"].values[0]
+        return data

tools/web_search.py

@@ -1,27 +0,0 @@
-from typing import Any, Optional
-from smolagents.tools import Tool
-import duckduckgo_search
-
-class DuckDuckGoSearchTool(Tool):
-    name = "web_search"
-    description = "Performs a duckduckgo web search based on your query (think a Google search) then returns the top search results."
-    inputs = {'query': {'type': 'string', 'description': 'The search query to perform.'}}
-    output_type = "string"
-
-    def __init__(self, max_results=10, **kwargs):
-        super().__init__()
-        self.max_results = max_results
-        try:
-            from duckduckgo_search import DDGS
-        except ImportError as e:
-            raise ImportError(
-                "You must install package `duckduckgo_search` to run this tool: for instance run `pip install duckduckgo-search`."
-            ) from e
-        self.ddgs = DDGS(**kwargs)
-
-    def forward(self, query: str) -> str:
-        results = self.ddgs.text(query, max_results=self.max_results)
-        if len(results) == 0:
-            raise Exception("No results found! Try a less restrictive/shorter query.")
-        postprocessed_results = [f"[{result['title']}]({result['href']})\n{result['body']}" for result in results]
-        return "## Search Results\n\n" + "\n\n".join(postprocessed_results)