From 1ff4cfdc42a133887187560a30e85b1d4e3655c2 Mon Sep 17 00:00:00 2001
From: xuyuan23 <643854343@qq.com>
Date: Mon, 31 Jul 2023 17:49:10 +0800
Subject: [PATCH] feat: support bard proxy server

support bard proxy server

Close #384
---
 pilot/model/proxy/llms/bard.py | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/pilot/model/proxy/llms/bard.py b/pilot/model/proxy/llms/bard.py
index badb0e912..7d3670ad6 100644
--- a/pilot/model/proxy/llms/bard.py
+++ b/pilot/model/proxy/llms/bard.py
@@ -1,4 +1,5 @@
 import bardapi
+import requests
 from typing import List
 from pilot.configs.config import Config
 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):
-    token = CFG.bard_proxy_api_key
-
     history = []
     messages: List[ModelMessage] = params["messages"]
     for message in messages:
@@ -35,8 +34,21 @@ def bard_generate_stream(model, tokenizer, params, device, context_len=2048):
     for msg in history:
         if msg.get("content"):
             msgs.append(msg["content"])
-    response = bardapi.core.Bard(token).get_answer("\n".join(msgs))
-    if response is not None and response.get("content") is not None:
-        yield str(response["content"])
+
+    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:
-        yield f"bard response error: {str(response)}"
+        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:
+            yield str(response["content"])
+        else:
+            yield f"bard response error: {str(response)}"