|
@@ -0,0 +1,43 @@
|
|
|
+from typing import Any
|
|
|
+
|
|
|
+from core.helper import ssrf_proxy
|
|
|
+from core.tools.entities.tool_entities import ToolInvokeMessage
|
|
|
+from core.tools.tool.builtin_tool import BuiltinTool
|
|
|
+
|
|
|
+
|
|
|
+class JinaTokenizerTool(BuiltinTool):
|
|
|
+ _jina_tokenizer_endpoint = 'https://tokenize.jina.ai/'
|
|
|
+
|
|
|
+ def _invoke(
|
|
|
+ self,
|
|
|
+ user_id: str,
|
|
|
+ tool_parameters: dict[str, Any],
|
|
|
+ ) -> ToolInvokeMessage:
|
|
|
+ content = tool_parameters['content']
|
|
|
+ body = {
|
|
|
+ "content": content
|
|
|
+ }
|
|
|
+
|
|
|
+ headers = {
|
|
|
+ 'Content-Type': 'application/json'
|
|
|
+ }
|
|
|
+
|
|
|
+ if 'api_key' in self.runtime.credentials and self.runtime.credentials.get('api_key'):
|
|
|
+ headers['Authorization'] = "Bearer " + self.runtime.credentials.get('api_key')
|
|
|
+
|
|
|
+ if tool_parameters.get('return_chunks', False):
|
|
|
+ body['return_chunks'] = True
|
|
|
+
|
|
|
+ if tool_parameters.get('return_tokens', False):
|
|
|
+ body['return_tokens'] = True
|
|
|
+
|
|
|
+ if tokenizer := tool_parameters.get('tokenizer'):
|
|
|
+ body['tokenizer'] = tokenizer
|
|
|
+
|
|
|
+ response = ssrf_proxy.post(
|
|
|
+ self._jina_tokenizer_endpoint,
|
|
|
+ headers=headers,
|
|
|
+ json=body,
|
|
|
+ )
|
|
|
+
|
|
|
+ return self.create_json_message(response.json())
|