|
import json |
|
import os.path |
|
|
|
from pathlib import Path |
|
from typing import List, Optional, Union |
|
|
|
from langchain.schema import AIMessage, HumanMessage, SystemMessage |
|
|
|
from gpt_engineer.core.ai import AI |
|
from gpt_engineer.core.token_usage import TokenUsageLog |
|
|
|
|
|
Message = Union[AIMessage, HumanMessage, SystemMessage] |
|
|
|
|
|
class CachingAI(AI): |
|
def __init__(self, *args, **kwargs): |
|
self.temperature = 0.1 |
|
self.azure_endpoint = "" |
|
self.streaming = False |
|
try: |
|
self.model_name = "gpt-4-1106-preview" |
|
self.llm = self._create_chat_model() |
|
except: |
|
self.model_name = "cached_response_model" |
|
self.llm = None |
|
self.streaming = False |
|
self.token_usage_log = TokenUsageLog("gpt-4-1106-preview") |
|
self.cache_file = Path(__file__).parent / "ai_cache.json" |
|
|
|
def next( |
|
self, |
|
messages: List[Message], |
|
prompt: Optional[str] = None, |
|
*, |
|
step_name: str, |
|
) -> List[Message]: |
|
""" |
|
Advances the conversation by sending message history |
|
to LLM and updating with the response. |
|
|
|
Parameters |
|
---------- |
|
messages : List[Message] |
|
The list of messages in the conversation. |
|
prompt : Optional[str], optional |
|
The prompt to use, by default None. |
|
step_name : str |
|
The name of the step. |
|
|
|
Returns |
|
------- |
|
List[Message] |
|
The updated list of messages in the conversation. |
|
""" |
|
""" |
|
Advances the conversation by sending message history |
|
to LLM and updating with the response. |
|
""" |
|
if prompt: |
|
messages.append(HumanMessage(content=prompt)) |
|
|
|
|
|
if os.path.isfile(self.cache_file): |
|
with open(self.cache_file, "r") as cache_file: |
|
cache = json.load(cache_file) |
|
else: |
|
cache = dict() |
|
|
|
messages_key = self.serialize_messages(messages) |
|
if messages_key not in cache: |
|
callbacks = [] |
|
print("calling backoff inference") |
|
response = self.backoff_inference(messages, callbacks) |
|
self.token_usage_log.update_log( |
|
messages=messages, answer=response.content, step_name=step_name |
|
) |
|
print("called backoff inference") |
|
print("cost in usd:", self.token_usage_log.usage_cost()) |
|
|
|
messages.append(response) |
|
cache[messages_key] = self.serialize_messages(messages) |
|
with open(self.cache_file, "w") as cache_file: |
|
json.dump(cache, cache_file) |
|
cache_file.write("\n") |
|
|
|
return self.deserialize_messages(cache[messages_key]) |
|
|
