app.py

import os

# Env defaults MUST be set before importing cognee — cognee reads config at import time.
os.environ.setdefault(
    "SYSTEM_ROOT_DIRECTORY", os.path.join(os.path.dirname(__file__), ".cognee_system")
)
os.environ.setdefault("LLM_MODEL", "openai/gpt-4o-mini")
os.environ.setdefault("TELEMETRY_DISABLED", "1")

import asyncio  # noqa: E402
import logging  # noqa: E402
import re  # noqa: E402
from datetime import datetime, timezone  # noqa: E402

import gradio as gr  # noqa: E402
import cognee  # noqa: E402
from cognee.modules.search.types.SearchType import SearchType  # noqa: E402

# ---- Config ----

DATASET = "manovich"
MAX_QUESTION_LEN = 2000

# Wallet safety net — global circuit breaker across all users per UTC day.
# Resets at UTC midnight or on container restart. Not auth, just a hard ceiling.
MAX_QUESTIONS_PER_DAY = 500

# Frozen DB stats (from cognee logs during last ingest). Update if DB is rebuilt.
STATS = {
    "articles": 63,
    "nodes": 3566,
    "edges": 8205,
    "span": "1992–2007",
}

ERA_FILTERS = {
    "All years (1992–2007)": "",
    "Early (1992–1998)": (
        "When answering, focus on Manovich's writing between 1992 and 1998. "
    ),
    "Middle (1999–2003)": (
        "When answering, focus on Manovich's writing between 1999 and 2003. "
    ),
    "Late (2004–2007)": (
        "When answering, focus on Manovich's writing between 2004 and 2007. "
    ),
}

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

# Serialize cognee access — keeps behavior predictable under concurrent requests.
_search_lock = asyncio.Lock()

# Global per-day counter.
_daily_counter = {"date": None, "count": 0}


def _today_utc() -> str:
    return datetime.now(timezone.utc).date().isoformat()


def _daily_cap_hit() -> bool:
    if _daily_counter["date"] != _today_utc():
        _daily_counter["date"] = _today_utc()
        _daily_counter["count"] = 0
    return _daily_counter["count"] >= MAX_QUESTIONS_PER_DAY


def _bump_daily() -> None:
    if _daily_counter["date"] != _today_utc():
        _daily_counter["date"] = _today_utc()
        _daily_counter["count"] = 0
    _daily_counter["count"] += 1


# ---- Search + citation extraction ----


def _unwrap_answer(result) -> str:
    if hasattr(result, "search_result"):
        items = result.search_result
    elif isinstance(result, dict):
        items = result.get("search_result", [])
    else:
        return str(result)
    return "\n".join(items) if isinstance(items, list) else str(items)


_TITLE_RE = re.compile(r"^# (.+)$", re.MULTILINE)
_YEAR_RE = re.compile(r"_year:\s*(\d{4})_")


def _extract_sources(chunks_result) -> list[tuple[str, str | None]]:
    """Extract (title, year) tuples from a CHUNKS SearchResult.

    The articles were ingested with a '# Title' heading and a '_year: YYYY_' tag
    in their text, so we parse those out of each chunk. Chunks from the middle of
    an article won't have this metadata and are skipped."""
    if chunks_result is None:
        return []
    if isinstance(chunks_result, dict):
        items = chunks_result.get("search_result", [])
    elif hasattr(chunks_result, "search_result"):
        items = chunks_result.search_result
    else:
        return []
    if not items:
        return []

    out: list[tuple[str, str | None]] = []
    for item in items:
        text = item.get("text", "") if isinstance(item, dict) else getattr(item, "text", "")
        if not text:
            continue
        t_match = _TITLE_RE.search(text)
        if not t_match:
            continue
        title = t_match.group(1).strip()
        y_match = _YEAR_RE.search(text)
        year = y_match.group(1) if y_match else None
        out.append((title, year))
    return out


async def _search(question: str, era_prefix: str) -> tuple[str, list[tuple[str, str | None]]]:
    prefixed = (era_prefix + question) if era_prefix else question

    async with _search_lock:
        # Run sequentially — concurrent cognee.search calls can contend on Kuzu locks.
        answer_raw = await cognee.search(
            query_text=prefixed,
            query_type=SearchType.GRAPH_COMPLETION,
            datasets=[DATASET],
        )
        try:
            chunks_raw = await cognee.search(
                query_text=question,
                query_type=SearchType.CHUNKS,
                datasets=[DATASET],
            )
        except Exception:
            logger.exception("CHUNKS search failed — continuing without sources.")
            chunks_raw = []

    if not answer_raw:
        return "No results found.", []

    answer = _unwrap_answer(answer_raw[0])
    raw_sources = _extract_sources(chunks_raw[0]) if chunks_raw else []
    logger.info(
        "search: q=%r chunks_len=%d raw_sources=%d",
        question[:60],
        len(chunks_raw[0].get("search_result", [])) if chunks_raw and isinstance(chunks_raw[0], dict) else 0,
        len(raw_sources),
    )

    # Dedupe by title, preserving order (highest-similarity chunks come first).
    seen: set[str] = set()
    uniq: list[tuple[str, str | None]] = []
    for title, year in raw_sources:
        if title not in seen:
            seen.add(title)
            uniq.append((title, year))
    return answer, uniq[:3]


def _format_answer(answer: str, sources: list[tuple[str, str | None]]) -> str:
    if not sources:
        return answer
    lines = []
    for title, year in sources:
        if year:
            lines.append(f"- *{title}* ({year})")
        else:
            lines.append(f"- *{title}*")
    return f"{answer}\n\n---\n**Drawing from:**\n" + "\n".join(lines)


# ---- Gradio handlers ----


async def respond(message: str, history: list, era: str):
    """Async generator yielding incremental history updates."""
    history = history or []
    message = (message or "").strip()

    def _with(content: str) -> list:
        return history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": content},
        ]

    if not message:
        yield history
        return

    if len(message) > MAX_QUESTION_LEN:
        yield _with(f"Question too long (>{MAX_QUESTION_LEN} characters). Please shorten it.")
        return

    if _daily_cap_hit():
        yield _with("The demo has reached its daily usage cap. Please try again tomorrow.")
        return

    # Progress placeholder.
    progress = _with(
        f"*Searching across {STATS['articles']} articles and ~{STATS['edges']:,} relationships…*"
    )
    yield progress

    try:
        answer, sources = await _search(message, ERA_FILTERS.get(era, ""))
        _bump_daily()
        final = progress.copy()
        final[-1] = {"role": "assistant", "content": _format_answer(answer, sources)}
        yield final
    except Exception:
        logger.exception("Search failed for message: %r", message)
        err = progress.copy()
        err[-1] = {
            "role": "assistant",
            "content": "Sorry — something went wrong answering that. Please try again.",
        }
        yield err


def _clear_chat():
    return [], ""


def _load_from_query_string(request: gr.Request | None):
    """Preload question text from ?q=… for shareable links."""
    if request is None:
        return ""
    try:
        params = dict(request.query_params) if request.query_params else {}
    except Exception:
        params = {}
    return params.get("q", "") or ""


# ---- UI ----

HEADER_HTML = """\
<div id="mkg-header">
  <h1>Lev Manovich — Knowledge Graph</h1>
  <p class="mkg-sub">
    Ask cross-article questions across 63 essays, 1992–2007.
    Built with <a href="https://github.com/topoteretes/cognee" target="_blank">cognee</a>.
  </p>
</div>
"""

CUSTOM_CSS = """
/* Kill Gradio's oversized top padding so the header is visible on load. */
.gradio-container { padding-top: 0 !important; }
footer { display: none !important; }

#mkg-header {
    background: linear-gradient(135deg, #1e3a8a 0%, #4338ca 100%);
    color: white;
    padding: 24px 28px;
    margin: 0 0 16px 0;
    border-radius: 8px;
}
#mkg-header h1 {
    margin: 0 0 6px 0;
    font-size: 1.75rem;
    line-height: 1.2;
    color: white;
    font-weight: 600;
}
#mkg-header .mkg-sub {
    margin: 0;
    color: rgba(255,255,255,0.85);
    font-size: 0.95rem;
}
#mkg-header a {
    color: #c7d2fe;
    text-decoration: underline;
}
"""

ABOUT = f"""\
This demo queries a knowledge graph built from **{STATS['articles']} articles**
Lev Manovich published between **{STATS['span']}**. The graph contains roughly
**{STATS['nodes']:,} entities** and **{STATS['edges']:,} relationships**. Each
question triggers a graph traversal to gather relevant context, which an LLM
(GPT-4o-mini) then synthesizes into an answer.

Answers are not generated from any single article — they assemble evidence
across the corpus. The *Drawing from* lines beneath each answer show the
articles cognee leaned on most.
"""

FOOTER = """\
---

*Corpus: Lev Manovich's collected essays, 1992–2007. Knowledge graph extracted
and queried via [cognee](https://github.com/topoteretes/cognee). Answers synthesized
by GPT-4o-mini. This is a demo — answers may contain inaccuracies.*
"""

EXAMPLES = [
    "What does Manovich mean by the term 'velvet revolution'?",
    "What does Manovich mean by 'deep remixability'?",
    "What is 'navigable space' and why does it matter?",
    "What is the relationship between cinema and software in his work?",
    "How does Manovich define new media?",
    "What thinkers does Manovich draw on most, and for what purposes?",
    "How does his thinking change between 1992 and 2007?",
]

with gr.Blocks(
    title="Manovich Knowledge Graph",
    theme=gr.themes.Soft(),
    css=CUSTOM_CSS,
) as demo:
    gr.HTML(HEADER_HTML)
    with gr.Accordion("About this graph", open=False):
        gr.Markdown(ABOUT)

    chatbot = gr.Chatbot(type="messages", height=480, show_label=False)

    with gr.Row():
        era = gr.Dropdown(
            choices=list(ERA_FILTERS.keys()),
            value="All years (1992–2007)",
            label="Period filter (advisory — nudges the answer, doesn't hard-filter retrieval)",
            scale=1,
        )

    msg = gr.Textbox(
        placeholder="Ask a question about Manovich's work…",
        show_label=False,
        autofocus=True,
        lines=2,
    )

    with gr.Row():
        submit_btn = gr.Button("Ask", variant="primary")
        clear_btn = gr.Button("Clear")

    gr.Markdown("**Try one of these:**")
    gr.Examples(examples=EXAMPLES, inputs=msg)

    gr.Markdown(FOOTER)

    # Wire up handlers. .then(lambda: "", outputs=msg) clears the textbox after submit.
    msg.submit(respond, [msg, chatbot, era], [chatbot]).then(lambda: "", outputs=msg)
    submit_btn.click(respond, [msg, chatbot, era], [chatbot]).then(lambda: "", outputs=msg)
    clear_btn.click(_clear_chat, outputs=[chatbot, msg])

    # Shareable links: populate textbox from ?q= on first load.
    demo.load(_load_from_query_string, outputs=msg)

demo.queue(max_size=10, default_concurrency_limit=2)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)