diff options
Diffstat (limited to 'searx/engines')
| -rw-r--r-- | searx/engines/cloudflareai.py | 68 |
1 files changed, 68 insertions, 0 deletions
diff --git a/searx/engines/cloudflareai.py b/searx/engines/cloudflareai.py new file mode 100644 index 000000000..fc1051efb --- /dev/null +++ b/searx/engines/cloudflareai.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: AGPL-3.0-or-later +"""Cloudflare AI engine""" + +from json import loads, dumps +from searx.exceptions import SearxEngineAPIException + +about = { + "website": 'https://ai.cloudflare.com', + "wikidata_id": None, + "official_api_documentation": 'https://developers.cloudflare.com/workers-ai', + "use_official_api": True, + "require_api_key": True, + "results": 'JSON', +} + +cf_account_id = '' +cf_ai_api = '' +cf_ai_gateway = '' + +cf_ai_model = '' +cf_ai_model_display_name = 'Cloudflare AI' + +# Assistant messages hint to the AI about the desired output format. Not all models support this role. +cf_ai_model_assistant = 'Keep your answers as short and effective as possible.' +# System messages define the AI's personality. You can use them to set rules and how you expect the AI to behave. +cf_ai_model_system = 'You are a self-aware language model who is honest and direct about any question from the user.' + + +def request(query, params): + + params['query'] = query + + params['url'] = f'https://gateway.ai.cloudflare.com/v1/{cf_account_id}/{cf_ai_gateway}/workers-ai/{cf_ai_model}' + + params['method'] = 'POST' + + params['headers']['Authorization'] = f'Bearer {cf_ai_api}' + params['headers']['Content-Type'] = 'application/json' + + params['data'] = dumps( + { + 'messages': [ + {'role': 'assistant', 'content': cf_ai_model_assistant}, + {'role': 'system', 'content': cf_ai_model_system}, + {'role': 'user', 'content': params['query']}, + ] + } + ).encode('utf-8') + + return params + + +def response(resp): + results = [] + json = loads(resp.text) + + if 'error' in json: + raise SearxEngineAPIException('Cloudflare AI error: ' + json['error']) + + if 'result' in json: + results.append( + { + 'content': json['result']['response'], + 'infobox': cf_ai_model_display_name, + } + ) + + return results |