LangChain: Email Writer

# Install: pip install langchain langchain-openai langchain-community chromadb

from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
import os

# --- Config ---
CONTEXT_DIR = "./email-context"  # past emails, docs, templates
CHROMA_DIR = "./chroma_email"

# --- Load context documents ---
loader = DirectoryLoader(CONTEXT_DIR, glob="**/*.{md,txt}", loader_cls=TextLoader)
docs = loader.load()

splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
chunks = splitter.split_documents(docs)

embeddings = OpenAIEmbeddings()
vectorstore = Chroma.from_documents(chunks, embeddings, persist_directory=CHROMA_DIR)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# --- Email writing chain ---
prompt = ChatPromptTemplate.from_messages([
    ("system", """You are a professional email writer. Draft emails that are clear, concise, and appropriate for the context. Use the reference material below for tone and context.

Reference material:
{context}"""),
    ("human", """Write an email with these details:
To: {recipient}
Subject: {subject}
Key points: {key_points}
Tone: {tone}"""),
])

llm = ChatOpenAI(model="gpt-4o", temperature=0.3)
chain = prompt | llm | StrOutputParser()

def draft_email(recipient: str, subject: str, key_points: str, tone: str = "professional") -> str:
    context_docs = retriever.invoke(f"{recipient} {subject}")
    context = "\n".join(d.page_content for d in context_docs)
    return chain.invoke({
        "context": context,
        "recipient": recipient,
        "subject": subject,
        "key_points": key_points,
        "tone": tone,
    })

if __name__ == "__main__":
    print("Email Writer ready.")
    recipient = input("To: ")
    subject = input("Subject: ")
    key_points = input("Key points: ")
    tone = input("Tone (professional/casual/formal): ") or "professional"
    print("\n--- Draft ---\n")
    print(draft_email(recipient, subject, key_points, tone))