Skip to content

Commit

Permalink
add Cloudflare AI Gateway engine
Browse files Browse the repository at this point in the history
add Cloudflare AI Gateway engine

add settings for Cloudflare AI Gateway engine

set utf8 encode for data, fix non english char cause 500 error

format json data

fixed indentation and config format error

fix line-length limitation in CI

reformatted code for CI

reformatted code for CI

limit system prompts to less 120 chars

cleanup unused variable & format code
  • Loading branch information
hezhijie0327 authored and return42 committed Sep 23, 2024
1 parent 14241e7 commit 6be56ae
Show file tree
Hide file tree
Showing 2 changed files with 85 additions and 0 deletions.
68 changes: 68 additions & 0 deletions searx/engines/cloudflareai.py
Original file line number Diff line number Diff line change
@@ -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
17 changes: 17 additions & 0 deletions searx/settings.yml
Original file line number Diff line number Diff line change
Expand Up @@ -482,6 +482,23 @@ engines:
# to show premium or plus results too:
# skip_premium: false

- name: cloudflareai
engine: cloudflareai
shortcut: cfai
# get api token and accont id from https://developers.cloudflare.com/workers-ai/get-started/rest-api/
cf_account_id: 'your_cf_accout_id'
cf_ai_api: 'your_cf_api'
# create your ai gateway by https://developers.cloudflare.com/ai-gateway/get-started/creating-gateway/
cf_ai_gateway: 'your_cf_ai_gateway_name'
# find the model name from https://developers.cloudflare.com/workers-ai/models/#text-generation
cf_ai_model: 'ai_model_name'
# custom your preferences
# cf_ai_model_display_name: 'Cloudflare AI'
# cf_ai_model_assistant: 'prompts_for_assistant_role'
# cf_ai_model_system: 'prompts_for_system_role'
timeout: 30
disabled: true

# - name: core.ac.uk
# engine: core
# categories: science
Expand Down

0 comments on commit 6be56ae

Please sign in to comment.