Memory &
Context Management

from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_google_genai import ChatGoogleGenerativeAI
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Store for managing multiple conversation sessions
store = {}

def get_session_history(session_id: str) -> InMemoryChatMessageHistory:
    """Get or create a chat history for a session"""
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]

# Create a prompt with memory placeholder
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant."),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}")
])

# Initialize the model
llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    google_api_key=os.getenv("GOOGLE_API_KEY")
)

# Create a chain
chain = prompt | llm

# Wrap with message history
chain_with_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

# Have a conversation (note the session_id in config)
config = {"configurable": {"session_id": "alice_session"}}

response1 = chain_with_history.invoke(
    {"input": "Hi, my name is Alice"}, 
    config=config
)
print("AI:", response1.content)

# Continue conversation - history is automatically managed!
response2 = chain_with_history.invoke(
    {"input": "What's my name?"}, 
    config=config
)
print("AI:", response2.content)

# Check the conversation history
history = get_session_history("alice_session")
print("\nConversation history:")
for message in history.messages:
    print(f"{message.__class__.__name__}: {message.content}")

from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_google_genai import ChatGoogleGenerativeAI
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Store for managing multiple conversation sessions
store = {}

def get_session_history(session_id: str) -> InMemoryChatMessageHistory:
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]

# Set up the chain
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant."),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}")
])

llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    google_api_key=os.getenv("GOOGLE_API_KEY")
)

chain = prompt | llm

chain_with_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

# Simulate multiple users
users = [
    ("alice_session", "Hi, I'm Alice and I love Python"),
    ("bob_session", "Hello, I'm Bob and I prefer JavaScript"),
    ("alice_session", "What's my favorite language?"),
    ("bob_session", "What language do I prefer?"),
]

for session_id, message in users:
    config = {"configurable": {"session_id": session_id}}
    response = chain_with_history.invoke({"input": message}, config=config)
    print(f"[{session_id}] User: {message}")
    print(f"[{session_id}] AI: {response.content}")
    print()

# Show that memories are separate
print("\n=== Session Histories ===")
for session_id in ["alice_session", "bob_session"]:
    history = get_session_history(session_id)
    print(f"\n{session_id}:")
    for msg in history.messages:
        print(f"  {msg.__class__.__name__}: {msg.content[:50]}...")

from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.schema import HumanMessage, AIMessage
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

class WindowedChatMessageHistory(InMemoryChatMessageHistory):
    """Chat message history that only keeps last k exchanges"""
    
    def __init__(self, k: int = 3):
        super().__init__()
        self._k = k  # Number of exchanges to keep (1 exchange = 2 messages)
        
    def add_messages(self, messages):
        """Add messages and maintain window size"""
        super().add_messages(messages)
        # Keep only last k exchanges (2k messages)
        if len(self.messages) > self._k * 2:
            self.messages = self.messages[-(self._k * 2):]

# Store for managing sessions with windowed memory
store = {}

def get_windowed_session_history(session_id: str) -> WindowedChatMessageHistory:
    if session_id not in store:
        store[session_id] = WindowedChatMessageHistory(k=3)  # Keep last 3 exchanges
    return store[session_id]

# Set up the chain
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant."),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}")
])

llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    google_api_key=os.getenv("GOOGLE_API_KEY")
)

chain = prompt | llm

# Wrap with windowed message history
chain_with_window = RunnableWithMessageHistory(
    chain,
    get_windowed_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

# Simulate a longer conversation
config = {"configurable": {"session_id": "test_window"}}
responses = [
    "Tell me about Python",
    "What are decorators?",
    "Give me an example",
    "What about generators?",
    "How do they save memory?",
    "What did we discuss first?"  # This won't remember Python intro
]

for user_input in responses:
    print(f"\nUser: {user_input}")
    response = chain_with_window.invoke({"input": user_input}, config=config)
    print(f"AI: {response.content[:100]}...")  # Show first 100 chars

# Show what's in the windowed memory
print("\nWindow Memory (last 3 exchanges):")
history = get_windowed_session_history("test_window")
for msg in history.messages:
    print(f"{msg.__class__.__name__}: {msg.content[:50]}...")

from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.schema import SystemMessage, HumanMessage, AIMessage
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

class SummarizedChatMessageHistory(InMemoryChatMessageHistory):
    """Chat message history that summarizes old messages"""
    
    def __init__(self, llm, max_messages: int = 10):
        super().__init__()
        self._llm = llm
        self._max_messages = max_messages
        self._summary = ""
        
    def add_messages(self, messages):
        """Add messages and summarize if needed"""
        super().add_messages(messages)
        
        # If we exceed the limit, summarize older messages
        if len(self.messages) > self._max_messages:
            # Get messages to summarize (all but the last 4)
            messages_to_summarize = self.messages[:-4]
            
            # Create a summary prompt
            summary_prompt = f"""Previous summary: {self._summary}
            
New messages to add to summary:
{self._format_messages(messages_to_summarize)}

Create a concise summary of the conversation so far, preserving key information."""
            
            # Get summary from LLM
            summary_response = self._llm.invoke(summary_prompt)
            self._summary = summary_response.content
            
            # Keep only the summary and recent messages
            self.messages = [
                SystemMessage(content=f"Conversation summary: {self._summary}")
            ] + self.messages[-4:]
    
    def _format_messages(self, messages):
        """Format messages for summarization"""
        formatted = []
        for msg in messages:
            role = "Human" if isinstance(msg, HumanMessage) else "AI"
            formatted.append(f"{role}: {msg.content}")
        return "\n".join(formatted)

# Initialize model
llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    google_api_key=os.getenv("GOOGLE_API_KEY")
)

# Store for managing sessions with summary memory
store = {}

def get_summary_session_history(session_id: str) -> SummarizedChatMessageHistory:
    if session_id not in store:
        store[session_id] = SummarizedChatMessageHistory(llm, max_messages=10)
    return store[session_id]

# Set up the chain
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant."),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}")
])

chain = prompt | llm

# Wrap with summary message history
chain_with_summary = RunnableWithMessageHistory(
    chain,
    get_summary_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

# Have a detailed conversation
config = {"configurable": {"session_id": "test_summary"}}
conversations = [
    "I want to build a web scraper for job listings. I need to scrape Indeed and LinkedIn.",
    "It should save data to PostgreSQL and send email alerts for new matches.",
    "The tech stack will be Python with BeautifulSoup and Selenium.",
    "Add scheduling with cron jobs to run daily.",
    "Include error handling and logging.",
    "What database did I mention earlier?"  # Should remember PostgreSQL
]

for msg in conversations:
    response = chain_with_summary.invoke({"input": msg}, config=config)
    print(f"User: {msg}")
    print(f"AI: {response.content[:150]}...")
    print()

# Show the summarized memory
print("\nMemory State:")
history = get_summary_session_history("test_summary")
for msg in history.messages:
    if isinstance(msg, SystemMessage):
        print(f"Summary: {msg.content}")
    else:
        print(f"{msg.__class__.__name__}: {msg.content[:50]}...")

# Note: ConversationSummaryBufferMemory is deprecated, but the concept
# can be implemented using the modern approach by combining our
# WindowedChatMessageHistory and SummarizedChatMessageHistory patterns

from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.schema import SystemMessage, HumanMessage, AIMessage
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

class HybridChatMessageHistory(InMemoryChatMessageHistory):
    """Hybrid memory: summarize old messages, keep recent ones in full"""
    
    def __init__(self, llm, buffer_size: int = 4, max_messages: int = 6):
        super().__init__()
        self._llm = llm
        self._buffer_size = buffer_size  # Messages to keep in full
        self._max_messages = max_messages  # Total messages before summarizing
        self._summary = ""
        
    def add_messages(self, messages):
        """Add messages and apply hybrid strategy"""
        super().add_messages(messages)
        
        # If we exceed the limit, summarize older messages
        if len(self.messages) > self._max_messages:
            # Messages to summarize (all except the buffer)
            messages_to_summarize = self.messages[:-self._buffer_size]
            
            # Create summary including previous summary
            summary_prompt = f"""Previous summary: {self._summary}
            
New messages to add to summary:
{self._format_messages(messages_to_summarize)}

Create a concise summary preserving key technical details and decisions."""
            
            # Get summary
            summary_response = self._llm.invoke(summary_prompt)
            self._summary = summary_response.content
            
            # Keep summary + recent buffer
            self.messages = [
                SystemMessage(content=f"Previous conversation summary: {self._summary}")
            ] + self.messages[-self._buffer_size:]
    
    def _format_messages(self, messages):
        """Format messages for summarization"""
        formatted = []
        for msg in messages:
            role = "Human" if isinstance(msg, HumanMessage) else "AI"
            formatted.append(f"{role}: {msg.content}")
        return "\n".join(formatted)

# Initialize model
llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    google_api_key=os.getenv("GOOGLE_API_KEY")
)

# Store for sessions
store = {}

def get_hybrid_session_history(session_id: str) -> HybridChatMessageHistory:
    if session_id not in store:
        store[session_id] = HybridChatMessageHistory(llm, buffer_size=4, max_messages=6)
    return store[session_id]

# Set up chain
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful technical assistant."),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}")
])

chain = prompt | llm

chain_with_hybrid = RunnableWithMessageHistory(
    chain,
    get_hybrid_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

# Example: Technical discussion that evolves
config = {"configurable": {"session_id": "tech_discussion"}}
topics = [
    "I'm building a REST API with FastAPI",
    "I need authentication with JWT tokens",
    "The database will be PostgreSQL with SQLAlchemy",
    "How should I structure the project folders?",
    "What about testing strategies?",
    "Should I use Docker for deployment?"
]

for topic in topics:
    print(f"\nUser: {topic}")
    response = chain_with_hybrid.invoke({"input": topic}, config=config)
    print(f"AI: {response.content[:150]}...")

# Show hybrid memory state
print("\nHybrid Memory State:")
history = get_hybrid_session_history("tech_discussion")
for i, msg in enumerate(history.messages):
    if isinstance(msg, SystemMessage):
        print(f"[Summary] {msg.content[:100]}...")
    else:
        role = "Human" if isinstance(msg, HumanMessage) else "AI"
        print(f"[{role}] {msg.content[:50]}...")