ahserver_g/ahserver/llm_client.py

180 lines
5.0 KiB
Python
Raw Normal View History

2024-04-09 18:10:52 +08:00
import re
2024-04-09 11:57:18 +08:00
import json
2024-04-17 18:31:29 +08:00
from traceback import print_exc
2024-04-09 11:57:18 +08:00
from aiohttp import web
from appPublic.dictObject import DictObject
from appPublic.httpclient import HttpClient, RESPONSE_TEXT, RESPONSE_JSON, RESPONSE_BIN,RESPONSE_FILE, RESPONSE_STREAM
from appPublic.argsConvert import ArgsConvert
class LlmProxy:
def __init__(self, processor, desc):
assert desc.name
2024-04-09 18:10:52 +08:00
self.name = desc.name
2024-04-09 11:57:18 +08:00
self.processor = processor
self.auth_api = desc.auth
self.desc = desc
self.api_name = desc.name
2024-04-09 18:40:09 +08:00
self.data = DictObject()
2024-04-09 11:57:18 +08:00
self.ac = ArgsConvert('${', '}')
2024-04-09 18:10:52 +08:00
def line_chunk_match(self, l):
2024-04-17 18:31:29 +08:00
if self.api.chunk_match:
match = re.search(self.api.chunk_match, l)
2024-04-09 18:10:52 +08:00
if match:
return match.group(1)
return l
2024-04-17 18:31:29 +08:00
async def write_chunk(self, ll):
def eq(a, b):
return a == b
def ne(a, b):
return a != b
opfuncs = {
'==':eq,
'!=':ne
}
2024-04-11 09:57:54 +08:00
try:
2024-04-17 18:31:29 +08:00
print('write_chunk(),l=', ll)
l = self.line_chunk_match(ll)
2024-04-11 09:57:54 +08:00
d = DictObject(** json.loads(l))
2024-04-17 18:42:45 +08:00
j = {}
for r in self.api.resp or []:
j[r.name] = d.get_data_by_keys(r.value);
2024-04-17 18:31:29 +08:00
if self.api.chunk_filter:
v = d.get_data_by_keys(self.api.chunk_filter.name)
v1 = self.api.chunk_filter.value
op = self.api.chunk_filter.op
f = opfuncs.get(op)
if f and f(v,v1):
2024-04-17 18:42:45 +08:00
j[self.api.chunk_filter.field] = ''
print('filtered j=', j)
2024-04-09 11:57:18 +08:00
2024-04-11 09:57:54 +08:00
jstr = json.dumps(j) + '\n'
bin = jstr.encode('utf-8')
await self.resp.write(bin)
2024-04-14 11:35:48 +08:00
await self.resp.drain()
2024-04-11 09:57:54 +08:00
except Exception as e:
2024-04-17 18:31:29 +08:00
print(f'Error:Write_chunk(),{l=} error:{e=}')
print_exc()
2024-04-09 11:57:18 +08:00
async def stream_handle(self, chunk):
chunk = chunk.decode('utf-8')
chunk = self.remain_str + chunk
lines = chunk.split('\n')
self.remain_str = lines[-1]
ls = lines[:-1]
for l in ls:
2024-04-09 18:10:52 +08:00
if l == '':
continue
2024-04-09 11:57:18 +08:00
await self.write_chunk(l)
2024-04-11 09:57:54 +08:00
async def get_apikey(self, apiname):
2024-04-09 18:10:52 +08:00
f = self.processor.run_ns.get_llm_user_apikey
if f:
# return a DictObject instance
2024-04-11 09:57:54 +08:00
return await f(apiname, self.user)
2024-04-09 18:10:52 +08:00
raise Exception('get_llm_user_apikey() function not found in ServerEnv')
2024-04-09 11:57:18 +08:00
async def do_auth(self, request):
2024-04-09 18:10:52 +08:00
d = self.desc.auth
2024-04-11 09:57:54 +08:00
self.data = self.get_data(self.name)
2024-04-09 18:10:52 +08:00
if self.data.authed:
return
2024-04-11 09:57:54 +08:00
self.data = await self.get_apikey(self.name)
2024-04-09 18:10:52 +08:00
if self.data is None:
2024-04-11 09:57:54 +08:00
raise Exception(f'user({self.user}) do not has a apikey for {self.name}')
2024-04-09 18:10:52 +08:00
method = d.get('method', 'POST')
headers = {}
for h in d.get('headers',{}):
headers[h.get('name')] = h.get('value')
mydata = {}
for p in d.data or []:
mydata[p.get('name')] = p.get('value')
myparams = {}
for p in d.params or []:
myparams[p.get('name')] = p.get('value')
url = d.get('url')
params = {}
_params = self.datalize(myparams, params)
_headers = self.datalize(headers, params)
_data = self.datalize(mydata, params)
hc = HttpClient()
resp_data = await hc.request(url, method, response_type=RESPONSE_JSON,
params=_params,
data=None if _data == {} else json.dumps(_data),
headers=_headers)
resp_data = DictObject(**resp_data)
for sd in d.set_data:
self.data[sd.name] = resp_data.get_data_by_keys(sd.field)
self.data.authed = True
2024-04-11 09:57:54 +08:00
self.set_data(self.name, self.data)
2024-04-09 18:10:52 +08:00
2024-04-11 09:57:54 +08:00
def data_key(self, apiname):
if self.user is None:
self.user = 'anonymous'
return apiname + '_a_' + self.user
2024-04-09 11:57:18 +08:00
2024-04-11 09:57:54 +08:00
def set_data(self, apiname, data):
2024-04-09 18:10:52 +08:00
request = self.processor.run_ns.request
app = request.app
2024-04-11 09:57:54 +08:00
app.set_data(self.data_key(apiname), data)
2024-04-09 18:10:52 +08:00
2024-04-11 09:57:54 +08:00
def get_data(self, apiname):
2024-04-09 18:10:52 +08:00
request = self.processor.run_ns.request
app = request.app
2024-04-11 09:57:54 +08:00
return app.get_data(self.data_key(apiname))
2024-04-09 18:10:52 +08:00
2024-04-09 11:57:18 +08:00
async def __call__(self, request, params):
2024-04-11 09:57:54 +08:00
self.user = await self.processor.run_ns.get_user()
2024-04-09 11:57:18 +08:00
mapi = params.mapi
stream = params.stream
self.resp = web.StreamResponse()
await self.resp.prepare(request)
if stream is None:
stream = True
self.remain_str = ''
if not self.desc[mapi]:
raise Exception(f'{mapi} not defined')
d = self.desc[mapi]
self.api = d
2024-04-09 18:10:52 +08:00
self.chunk_match = d.chunk_match
2024-04-09 11:57:18 +08:00
if self.api.need_auth and self.auth_api:
2024-04-09 18:10:52 +08:00
await self.do_auth(request)
2024-04-11 09:57:54 +08:00
else:
self.data = await self.get_apikey(self.name)
2024-04-09 11:57:18 +08:00
assert d.get('url')
method = d.get('method', 'POST')
headers = {}
for h in d.get('headers',{}):
headers[h.get('name')] = h.get('value')
mydata = {}
for p in d.get('data', {}):
mydata[p.get('name')] = p.get('value')
myparams = {}
for p in d.get('params', {}):
myparams[p.get('name')] = p.get('value')
url = d.get('url')
_params = self.datalize(myparams, params)
_headers = self.datalize(headers, params)
_data = self.datalize(mydata, params)
response_type = RESPONSE_STREAM
hc = HttpClient()
2024-04-11 16:22:48 +08:00
print(f'{url=},{method=},{_params=},{_data=},{_headers=}')
2024-04-09 11:57:18 +08:00
resp_data = await hc.request(url, method, response_type=response_type,
params=_params,
data=None if _data == {} else json.dumps(_data),
stream_func=self.stream_handle,
headers=_headers)
if self.remain_str != '':
2024-04-09 18:10:52 +08:00
await self.write_chunk(self.remain_str)
2024-04-09 11:57:18 +08:00
return self.resp
def datalize(self, dic, data={}):
mydata = self.data.copy()
mydata.update(data)
s1 = self.ac.convert(dic, mydata)
return s1