Spaces:
Sleeping
Sleeping
File size: 2,506 Bytes
8a8f3ed |
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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import datetime
import json
import os
from typing import Dict, List, Optional
from langchain.schema import HumanMessage, AIMessage
from langchain.memory import ConversationBufferMemory
class CaptionHistory:
"""
Manages caption generation history using Langchain
"""
def __init__(self):
self.memory = ConversationBufferMemory(
return_messages=True,
memory_key="chat_history"
)
self.history_file = "caption_history.json"
self.load_history() # Load existing history on initialization
def add_interaction(self, image_name: str, model: str,
caption: str, timestamp: str|None = None):
if not timestamp:
timestamp = datetime.datetime.now().isoformat()
interaction = {
"timestamp": timestamp,
"image_name": image_name,
"model": model,
"caption": caption
}
# Add to langchain memory
human_msg = HumanMessage(
content=f"Generate caption for {image_name} using {model}"
)
ai_msg = AIMessage(content=caption)
self.memory.chat_memory.add_user_message(human_msg.content)
self.memory.chat_memory.add_ai_message(ai_msg.content)
# Save the file
self.save_interaction(interaction)
def get_history(self) -> List[Optional[Dict[str, str]]]:
try:
with open(self.history_file, mode="r") as f:
return json.load(f)
except FileNotFoundError:
return []
def save_interaction(self, interaction: Dict[str, str]) -> None:
history = self.get_history()
history.append(interaction)
with open(self.history_file, mode="w") as f:
json.dump(history, f, indent=2)
def load_history(self):
"""Fixed: Proper string formatting in f-strings"""
history = self.get_history()
for item in history:
human_msg = HumanMessage(
content=f"Generate caption for {item['image_name']} using {item['model']}" # Fixed: proper quotes
)
ai_msg = AIMessage(content=item["caption"])
self.memory.chat_memory.add_user_message(human_msg.content)
self.memory.chat_memory.add_ai_message(ai_msg.content)
def clear_history(self):
self.memory.clear()
if os.path.exists(self.history_file):
os.remove(self.history_file) |