diff --git a/devchat/openai/openai_chat.py b/devchat/openai/openai_chat.py index 515272a0..757574d7 100644 --- a/devchat/openai/openai_chat.py +++ b/devchat/openai/openai_chat.py @@ -1,4 +1,4 @@ -from typing import Any, Optional, Union, List, Dict, Iterator +from typing import Optional, Union, List, Dict, Iterator from pydantic import BaseModel, Field, Extra import openai from devchat.chat import Chat diff --git a/devchat/openai/openai_prompt.py b/devchat/openai/openai_prompt.py index 03741961..46c5b5f4 100644 --- a/devchat/openai/openai_prompt.py +++ b/devchat/openai/openai_prompt.py @@ -100,16 +100,16 @@ def append_new(self, message_type: str, content: str, self._new_messages[message_type].append(message) self._request_tokens += num_tokens return True - + def set_functions(self, functions, available_tokens: int = math.inf): num_tokens = message_tokens({"functions": json.dumps(functions)}, self.model) if num_tokens > available_tokens: return False - + self._new_messages[Message.FUNCTION] = functions self._request_tokens += num_tokens return True - + def get_functions(self): return self._new_messages.get(Message.FUNCTION, None)