blob: 3379fefe0b20180326be1a14e7b90583cf56be70 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
|
from typing import List
import requests
from pydantic import Field
from ...core.main import ChatMessage
from .base import LLM
class GooglePaLMAPI(LLM):
"""
The Google PaLM API is currently in public preview, so production applications are not supported yet. However, you can [create an API key in Google MakerSuite](https://makersuite.google.com/u/2/app/apikey) and begin trying out the `chat-bison-001` model. Change `~/.continue/config.py` to look like this:
```python title="~/.continue/config.py"
from continuedev.core.models import Models
from continuedev.libs.llm.hf_inference_api import GooglePaLMAPI
config = ContinueConfig(
...
models=Models(
default=GooglePaLMAPI(
model="chat-bison-001"
api_key="<MAKERSUITE_API_KEY>",
)
)
```
"""
api_key: str = Field(..., description="Google PaLM API key")
model: str = "chat-bison-001"
async def _stream_complete(self, prompt, options):
api_url = f"https://generativelanguage.googleapis.com/v1beta2/models/{self.model}:generateMessage?key={self.api_key}"
body = {"prompt": {"messages": [{"content": prompt}]}}
response = requests.post(api_url, json=body)
yield response.json()["candidates"][0]["content"]
async def _stream_chat(self, messages: List[ChatMessage], options):
msg_lst = []
for message in messages:
msg_lst.append({"content": message["content"]})
api_url = f"https://generativelanguage.googleapis.com/v1beta2/models/{self.model}:generateMessage?key={self.api_key}"
body = {"prompt": {"messages": msg_lst}}
response = requests.post(api_url, json=body)
yield {
"content": response.json()["candidates"][0]["content"],
"role": "assistant",
}
|
