HarshitX/Multi_LLM_Image_Captioning

1import datetime
2import json
3import os
4from typing import Dict, List, Optional
5from langchain.schema import HumanMessage, AIMessage
6from langchain.memory import ConversationBufferMemory
7
8class CaptionHistory:
9    """
10    Manages caption generation history using Langchain
11    """
12    def __init__(self):
13        self.memory = ConversationBufferMemory(
14            return_messages=True,
15            memory_key="chat_history"
16        )
17        self.history_file = "caption_history.json"
18        self.load_history()  # Load existing history on initialization
19
20    def add_interaction(self, image_name: str, model: str,
21    caption: str, timestamp: str|None = None):
22        if not timestamp:
23            timestamp = datetime.datetime.now().isoformat()
24
25        interaction = {
26            "timestamp": timestamp,
27            "image_name": image_name,
28            "model": model,
29            "caption": caption
30        }
31
32        # Add to langchain memory
33        human_msg = HumanMessage(
34            content=f"Generate caption for {image_name} using {model}"
35        )
36        ai_msg = AIMessage(content=caption)
37
38        self.memory.chat_memory.add_user_message(human_msg.content)
39        self.memory.chat_memory.add_ai_message(ai_msg.content)
40
41        # Save the file 
42        self.save_interaction(interaction)
43
44    def get_history(self) -> List[Optional[Dict[str, str]]]:
45        try:
46            with open(self.history_file, mode="r") as f:
47                return json.load(f)
48        except FileNotFoundError:
49            return []
50
51    def save_interaction(self, interaction: Dict[str, str]) -> None:
52        history = self.get_history()
53        history.append(interaction)
54        with open(self.history_file, mode="w") as f:
55            json.dump(history, f, indent=2)
56
57    def load_history(self):
58        """Fixed: Proper string formatting in f-strings"""
59        history = self.get_history()
60        for item in history:
61            human_msg = HumanMessage(
62                content=f"Generate caption for {item['image_name']} using {item['model']}"  # Fixed: proper quotes
63            )
64            ai_msg = AIMessage(content=item["caption"]) 
65            self.memory.chat_memory.add_user_message(human_msg.content)
66            self.memory.chat_memory.add_ai_message(ai_msg.content)
67
68    def clear_history(self):
69        self.memory.clear()
70        if os.path.exists(self.history_file):
71            os.remove(self.history_file)