from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

# Store: session_id → history object
store: dict[str, InMemoryChatMessageHistory] = {}

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

# Wrap any chain/LLM with message history
chain_with_history = RunnableWithMessageHistory(
    runnable=chain,           # Any LangChain Runnable
    get_session_history=get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
)

# Invoke with session config
response = chain_with_history.invoke(
    {"input": "What is my name?"},
    config={"configurable": {"session_id": "user-123"}},
)

from langchain_postgres import PGVector
from langchain_openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
connection = "postgresql+psycopg://user:pass@localhost:5432/dbname"

vector_store = PGVector(
    embeddings=embeddings,
    collection_name="agent_memories",
    connection=connection,
    use_jsonb=True,
)

# Store memories
vector_store.add_documents(docs, ids=[doc.metadata["id"] for doc in docs])

# Retrieve memories
results = vector_store.similarity_search("user preferences", k=5)
results_with_scores = vector_store.similarity_search_with_score("topic", k=3)

# Filter by metadata
filtered = vector_store.similarity_search(
    "query",
    k=5,
    filter={"user_id": {"$eq": "user-123"}},
)

# As retriever for RAG
retriever = vector_store.as_retriever(
    search_type="mmr",
    search_kwargs={"k": 5, "fetch_k": 20},
)

from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
import chromadb

# Option A: LangChain wrapper (simplest)
vector_store = Chroma(
    collection_name="agent_memories",
    embedding_function=OpenAIEmbeddings(),
    persist_directory="./chroma_db",
)

# Option B: Pass pre-configured client
persistent_client = chromadb.PersistentClient(path="./chroma_db")
vector_store = Chroma(
    client=persistent_client,
    collection_name="agent_memories",
    embedding_function=OpenAIEmbeddings(),
)

# Store memories
vector_store.add_documents(documents=docs, ids=uuids)

# Retrieve
results = vector_store.similarity_search("user preferences", k=5)
results_with_scores = vector_store.similarity_search_with_score("topic", k=3)
mmr_results = vector_store.max_marginal_relevance_search("query", k=5)

# As retriever
retriever = vector_store.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 5},
)

from langgraph.store.memory import InMemoryStore

# Production: replace with PostgresStore or RedisStore
store = InMemoryStore(index={"embed": embed_func, "dims": 1536})

# Write memory (namespace = folder, key = document ID)
store.put(("memories", user_id), "preference_1", {"content": "Prefers concise answers"})

# Read by key
item = store.get(("memories", user_id), "preference_1")

# Semantic search
results = store.search(
    ("memories", user_id),
    query="communication style",
    limit=5,
)