Browse Source

Add PubMed to tools (#2652)

crazywoola 1 year ago
parent
commit
b5c1bb346c

+ 7 - 4
api/core/tools/provider/_position.yaml

@@ -1,16 +1,19 @@
 - google
 - bing
+- duckduckgo
+- yahoo
 - wikipedia
+- arxiv
+- pubmed
 - dalle
 - azuredalle
+- stablediffusion
 - webscraper
+- youtube
 - wolframalpha
+- maths
 - github
 - chart
 - time
-- yahoo
-- stablediffusion
 - vectorizer
-- youtube
 - gaode
-- maths

File diff suppressed because it is too large
+ 1 - 0
api/core/tools/provider/builtin/pubmed/_assets/icon.svg


+ 20 - 0
api/core/tools/provider/builtin/pubmed/pubmed.py

@@ -0,0 +1,20 @@
+from core.tools.errors import ToolProviderCredentialValidationError
+from core.tools.provider.builtin.pubmed.tools.pubmed_search import PubMedSearchTool
+from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
+
+
+class PubMedProvider(BuiltinToolProviderController):
+    def _validate_credentials(self, credentials: dict) -> None:
+        try:
+            PubMedSearchTool().fork_tool_runtime(
+                meta={
+                    "credentials": credentials,
+                }
+            ).invoke(
+                user_id='',
+                tool_parameters={
+                    "query": "John Doe",
+                },
+            )
+        except Exception as e:
+            raise ToolProviderCredentialValidationError(str(e))

+ 10 - 0
api/core/tools/provider/builtin/pubmed/pubmed.yaml

@@ -0,0 +1,10 @@
+identity:
+  author: Pink Banana
+  name: pubmed
+  label:
+    en_US: PubMed
+    zh_Hans: PubMed
+  description:
+    en_US: A search engine for biomedical literature.
+    zh_Hans: 一款生物医学文献搜索引擎。
+  icon: icon.svg

+ 40 - 0
api/core/tools/provider/builtin/pubmed/tools/pubmed_search.py

@@ -0,0 +1,40 @@
+from typing import Any
+
+from langchain.tools import PubmedQueryRun
+from pydantic import BaseModel, Field
+
+from core.tools.entities.tool_entities import ToolInvokeMessage
+from core.tools.tool.builtin_tool import BuiltinTool
+
+
+class PubMedInput(BaseModel):
+    query: str = Field(..., description="Search query.")
+
+
+class PubMedSearchTool(BuiltinTool):
+    """
+    Tool for performing a search using PubMed search engine.
+    """
+
+    def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]:
+        """
+        Invoke the PubMed search tool.
+
+        Args:
+            user_id (str): The ID of the user invoking the tool.
+            tool_parameters (dict[str, Any]): The parameters for the tool invocation.
+
+        Returns:
+            ToolInvokeMessage | list[ToolInvokeMessage]: The result of the tool invocation.
+        """
+        query = tool_parameters.get('query', '')
+
+        if not query:
+            return self.create_text_message('Please input query')
+
+        tool = PubmedQueryRun(args_schema=PubMedInput)
+
+        result = tool.run(query)
+
+        return self.create_text_message(self.summary(user_id=user_id, content=result))
+    

+ 23 - 0
api/core/tools/provider/builtin/pubmed/tools/pubmed_search.yaml

@@ -0,0 +1,23 @@
+identity:
+  name: pubmed_search
+  author: Pink Banana
+  label:
+    en_US: PubMed Search
+    zh_Hans: PubMed 搜索
+description:
+  human:
+    en_US: PubMed® comprises more than 35 million citations for biomedical literature from MEDLINE, life science journals, and online books. Citations may include links to full text content from PubMed Central and publisher web sites.
+    zh_Hans: PubMed® 包含来自 MEDLINE、生命科学期刊和在线书籍的超过 3500 万篇生物医学文献引用。引用可能包括来自 PubMed Central 和出版商网站的全文内容链接。
+  llm: Perform searches on PubMed and get results.
+parameters:
+  - name: query
+    type: string
+    required: true
+    label:
+      en_US: Query string
+      zh_Hans: 查询语句
+    human_description:
+      en_US: The search query.
+      zh_Hans: 搜索查询语句。
+    llm_description: Key words for searching
+    form: llm