Former-commit-id: 6b2b03bb82
pull/47/head
Kye 2 years ago
parent 84cb8ecd72
commit 78f8b37e25

@ -21,8 +21,6 @@ import google.generativeai as genai
class ChatGooglePalmError(Exception): class ChatGooglePalmError(Exception):
"""Error raised when there is an issue with the Google PaLM API.""" """Error raised when there is an issue with the Google PaLM API."""
def _truncate_at_stop_tokens( def _truncate_at_stop_tokens(
text: str, text: str,
stop: Optional[List[str]], stop: Optional[List[str]],
@ -37,12 +35,8 @@ def _truncate_at_stop_tokens(
text = text[:stop_token_idx] text = text[:stop_token_idx]
return text return text
def _response_to_result(response: genai.types.ChatResponse, stop: Optional[List[str]]) -> Dict[str, Any]: def _response_to_result(response: genai.types.ChatResponse, stop: Optional[List[str]]) -> Dict[str, Any]:
"""Convert a PaLM chat response to a result dictionary.""" """Convert a PaLM chat response to a result dictionary."""
# Code will need to be rewritten due to removal of schema imports.
# Instead of creating a ChatResult object, you would directly return a dictionary.
result = { result = {
"id": response.id, "id": response.id,
"created": response.created, "created": response.created,
@ -64,8 +58,6 @@ def _response_to_result(response: genai.types.ChatResponse, stop: Optional[List[
def _messages_to_prompt_dict(messages: List[Dict[str, Any]]) -> Dict[str, Any]: def _messages_to_prompt_dict(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
"""Convert a list of message dictionaries to a prompt dictionary.""" """Convert a list of message dictionaries to a prompt dictionary."""
# Code will need to be rewritten due to removal of schema imports.
# Instead of using BaseMessage.from_dict, you would directly create a dictionary from each message dict.
prompt = {"messages": []} prompt = {"messages": []}
for message in messages: for message in messages:
prompt["messages"].append({ prompt["messages"].append({

Loading…
Cancel
Save