github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-08-24 02:53:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support bard proxy server

Browse Source 
support bard proxy server

Close #384
This commit is contained in:
xuyuan23 2023-07-31 17:49:10 +08:00
parent f57e6c8ba1
commit 1ff4cfdc42
1 changed files with 18 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
pilot/model/proxy/llms/bard.py
View File

@ -1,4 +1,5 @@
import bardapi import bardapi
import requests
from typing import List from typing import List
from pilot.configs.config import Config from pilot.configs.config import Config
from pilot.scene.base_message import ModelMessage, ModelMessageRoleType from pilot.scene.base_message import ModelMessage, ModelMessageRoleType
@ -7,8 +8,6 @@ CFG = Config()
def bard_generate_stream(model, tokenizer, params, device, context_len=2048): def bard_generate_stream(model, tokenizer, params, device, context_len=2048):
token = CFG.bard_proxy_api_key
history = [] history = []
messages: List[ModelMessage] = params["messages"] messages: List[ModelMessage] = params["messages"]
for message in messages: for message in messages:
@ -35,7 +34,20 @@ def bard_generate_stream(model, tokenizer, params, device, context_len=2048):
for msg in history: for msg in history:
if msg.get("content"): if msg.get("content"):
msgs.append(msg["content"]) msgs.append(msg["content"])
response = bardapi.core.Bard(token).get_answer("\n".join(msgs))
if CFG.proxy_server_url is not None:
headers = {"Content-Type": "application/json"}
payloads = {"input": "\n".join(msgs)}
response = requests.post(
CFG.proxy_server_url, headers=headers, json=payloads, stream=False
)
if response.ok is True:
yield response.text
else:
yield f"bard proxy url request failed!, response = {str(response)}"
else:
response = bardapi.core.Bard(CFG.bard_proxy_api_key).get_answer("\n".join(msgs))
if response is not None and response.get("content") is not None: if response is not None and response.get("content") is not None:
yield str(response["content"]) yield str(response["content"])
else: else: