llmitate/app.py at main · TiElixir/llmitate · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
#!/usr/bin/env python3
"""
STEP 6 — Local chat UI (Streamlit) via Ollama
"""

from __future__ import annotations

import json
from pathlib import Path

import requests
import streamlit as st

from utils import load_config, summarize_persona_for_prompt

OLLAMA_URL = "http://localhost:11434"


def load_system_prompt(cfg: dict) -> str:
    profile_path = Path(cfg["output_dir"]) / "persona_profile.json"
    base = (
        "You are chatting in an Instagram DM. Reply exactly as the other person would — "
        "same slang, punctuation, emoji habits, and message length. "
        "Never act like an AI assistant. Stay in character.\n"
        "IMPORTANT: Use slang and common words naturally and sparingly. DO NOT spam or overuse the same catchphrases in every message."
    )
    if profile_path.exists():
        with open(profile_path, encoding="utf-8") as f:
            profile = json.load(f)
        summary = profile.get("system_prompt_summary") or summarize_persona_for_prompt(profile)
        return f"{base}\n\nStyle notes: {summary}"
    return base


def format_prompt(history: list[dict], system: str, user_message: str) -> str:
    parts = [
        "<|begin_of_text|>",
        "<|start_header_id|>system<|end_header_id|>\n\n",
        f"{system}<|eot_id|>",
    ]
    for msg in history:
        role = "user" if msg["role"] == "me" else "assistant"
        parts.append(f"<|start_header_id|>{role}<|end_header_id|>\n\n")
        parts.append(f"{msg['content']}<|eot_id|>")
    parts.append("<|start_header_id|>user<|end_header_id|>\n\n")
    parts.append(f"{user_message}<|eot_id|>")
    parts.append("<|start_header_id|>assistant<|end_header_id|>\n\n")
    return "".join(parts)


def generate(prompt: str, model: str, temperature: float, n: int = 1) -> list[str]:
    results: list[str] = []
    for _ in range(n):
        resp = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={
                "model": model,
                "prompt": prompt,
                "stream": False,
                "raw": True,
                "options": {"temperature": temperature, "num_predict": 256},
            },
            timeout=120,
        )
        resp.raise_for_status()
        results.append(resp.json().get("response", "").strip())
    return results


def main() -> None:
    cfg = load_config()
    model_name = cfg.get("ollama_model_name", "persona-imitate")
    target_name = cfg.get("target_sender", "them")
    default_temp = cfg.get("default_temperature", 0.7)
    vibe_n = cfg.get("vibe_check_candidates", 3)
    system_prompt = load_system_prompt(cfg)

    st.set_page_config(page_title="Persona Chat", page_icon="💬")
    st.title(f"Chat as {target_name}")
    st.caption("Local persona-mimicking chatbot · powered by Ollama")

    if "history" not in st.session_state:
        st.session_state.history = []

    with st.sidebar:
        temperature = st.slider("Temperature", 0.1, 1.5, default_temp, 0.05)
        st.markdown("**Low** = more in-character & repetitive  \n**High** = more creative")
        vibe_mode = st.toggle("Vibe check mode", value=False)
        if vibe_mode:
            st.info(f"Generates {vibe_n} candidate replies per message.")
        if st.button("Clear chat"):
            st.session_state.history = []
            st.rerun()

    for msg in st.session_state.history:
        with st.chat_message("user" if msg["role"] == "me" else "assistant"):
            st.write(msg["content"])

    user_input = st.chat_input("Type a message...")
    if user_input:
        st.session_state.history.append({"role": "me", "content": user_input})
        with st.chat_message("user"):
            st.write(user_input)

        prompt = format_prompt(st.session_state.history[:-1], system_prompt, user_input)
        try:
            n = vibe_n if vibe_mode else 1
            replies = generate(prompt, model_name, temperature, n=n)
        except requests.RequestException as exc:
            st.error(f"Ollama error: {exc}. Is Ollama running? Model created?")
            return

        if vibe_mode:
            with st.chat_message("assistant"):
                for i, reply in enumerate(replies, 1):
                    st.markdown(f"**Candidate {i}:** {reply}")
            reply = replies[0]
        else:
            reply = replies[0]
            with st.chat_message("assistant"):
                st.write(reply)

        st.session_state.history.append({"role": "them", "content": reply})


if __name__ == "__main__":
    main()