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 = """\
Ask cross-article questions across 63 essays, 1992–2007. Built with cognee.