app.py

import streamlit as st
import os
import chromadb
from chromadb.utils.embedding_functions import OllamaEmbeddingFunction
import ollama

# -----------------------------------
# PATHS
# -----------------------------------
DB_PATH = "./story_db"
STORY_DIR = "./my_novel_pages"

os.makedirs(STORY_DIR, exist_ok=True)

# -----------------------------------
# PAGE CONFIG
# -----------------------------------
st.set_page_config(page_title="Private Novel Strategist", layout="wide")
st.title("🖋️ Private Novel Strategist (M1 Optimized – phi3)")

# -----------------------------------
# CHROMA SETUP (OLLAMA EMBEDDINGS)
# -----------------------------------
client = chromadb.PersistentClient(path=DB_PATH)

embedding_function = OllamaEmbeddingFunction(
    url="http://localhost:11434/api/embeddings",
    model_name="nomic-embed-text"
)

collection = client.get_or_create_collection(
    name="novel_memory",
    embedding_function=embedding_function
)

# -----------------------------------
# TEXT SPLITTER (SMALLER = FASTER)
# -----------------------------------
def split_text(text, chunk_size=600, overlap=100):
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - overlap
    return chunks

# -----------------------------------
# SYNC FUNCTION
# -----------------------------------
def sync_brain():
    global collection

    try:
        client.delete_collection("novel_memory")
    except:
        pass

    collection = client.get_or_create_collection(
        name="novel_memory",
        embedding_function=embedding_function
    )

    for filename in os.listdir(STORY_DIR):
        if filename.endswith(".txt"):
            with open(os.path.join(STORY_DIR, filename), "r", encoding="utf-8") as f:
                text = f.read()

            chunks = split_text(text)

            for i, chunk in enumerate(chunks):
                collection.add(
                    documents=[chunk],
                    ids=[f"{filename}_{i}"]
                )

# -----------------------------------
# SIDEBAR
# -----------------------------------
with st.sidebar:
    st.header("Project Management")

    if st.button("Sync New Chapters/Notes"):
        with st.spinner("Indexing manuscript..."):
            sync_brain()
        st.success("Memory Updated!")

    st.info(f"Directory: {os.path.abspath(STORY_DIR)}")

# -----------------------------------
# RETRIEVAL (Lower k = Faster)
# -----------------------------------
def retrieve_context(query, k=2):
    results = collection.query(
        query_texts=[query],
        n_results=k
    )
    return "\n\n".join(results["documents"][0])

# -----------------------------------
# CHAT MEMORY
# -----------------------------------
if "messages" not in st.session_state:
    st.session_state.messages = []

for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

# -----------------------------------
# CHAT INPUT
# -----------------------------------
if user_input := st.chat_input("Ask about your plot or characters..."):
    st.session_state.messages.append({"role": "user", "content": user_input})
    st.chat_message("user").write(user_input)

    with st.chat_message("assistant"):

        if collection.count() == 0:
            st.warning("⚠ Please sync your chapters first.")
        else:
            context = retrieve_context(user_input)

            prompt = f"""
You are a private novel strategist.

Use the provided context to answer.
Stay consistent with the story world.

Context:
{context}

Question:
{user_input}

Answer:
"""

            # STREAMING RESPONSE (Feels Instant)
            response = ollama.chat(
                model="phi3",
                messages=[{"role": "user", "content": prompt}],
                stream=True,
                options={
                    "temperature": 0.7,
                    "num_predict": 300  # limit output length for speed
                }
            )

            full_response = ""
            message_placeholder = st.empty()

            for chunk in response:
                content = chunk["message"]["content"]
                full_response += content
                message_placeholder.markdown(full_response)

            st.session_state.messages.append(
                {"role": "assistant", "content": full_response}
            )