File size: 16,569 Bytes

#!/usr/bin/env python
"""PALIMPSESTE — FastAPI web server with cognitive metrics + streaming.

Endpoints:
  GET  /                — serves the React production build (or fallback HTML)
  GET  /health          — health check
  GET  /stats           — model statistics
  GET  /metrics         — rich metrics for dashboard visualization
  POST /chat            — send message, get response + confidence + source + explanation
  POST /chat/stream     — SSE streaming response with token-level data
  POST /teach           — teach a new Q/A pair at runtime (O(1))
  GET  /memory/samples  — recent memory traces for visualization
  GET  /conversation    — current conversation transcript

Usage:
  pip install fastapi uvicorn
  python examples/api_server.py --model thefinalboss/palimpseste-max --port 3332
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from palimseste.hf import HFPalimpsesteLM
from palimseste.chat import Conversation
from palimseste.cognitive import CognitiveAgent, CognitiveResponse

try:
    from fastapi import FastAPI
    from fastapi.middleware.cors import CORSMiddleware
    from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
    from pydantic import BaseModel
    import uvicorn
    HAS_FASTAPI = True
except ImportError:
    HAS_FASTAPI = False


# ----------------------------------------------------------------- models
class ChatRequest(BaseModel):
    message: str
    temperature: float | None = 0.0
    max_tokens: int = 200


class TeachRequest(BaseModel):
    question: str
    answer: str


class LearnTextRequest(BaseModel):
    text: str
    tag: str | None = None


class DreamRequest(BaseModel):
    cycles: int = 3
    replay_batch: int = 200


def _chain_to_dict(chain):
    if chain is None:
        return None
    return {
        "success": chain.success,
        "answer": chain.answer,
        "original_question": chain.original_question,
        "n_hops": chain.n_hops,
        "steps": [
            {
                "sub_question": s.sub_question,
                "sub_answer": s.sub_answer,
                "resolved_question": s.resolved_question,
            }
            for s in chain.steps
        ],
    }


def _cog_response_to_dict(resp: CognitiveResponse, elapsed_ms: float) -> dict:
    return {
        "response": resp.text,
        "confidence": round(resp.confidence, 4),
        "source": resp.source,
        "explanation": resp.explanation,
        "chain": _chain_to_dict(resp.chain),
        "corrected_answer": resp.corrected_answer,
        "elapsed_ms": round(elapsed_ms, 1),
    }


# ----------------------------------------------------------------- fallback UI
FALLBACK_HTML = """<!DOCTYPE html>
<html>
<head><title>PALIMPSESTE</title></head>
<body style="font-family:sans-serif;background:#0a0a12;color:#e0e0e0;padding:40px">
<h1>PALIMPSESTE API</h1>
<p>API is running. Endpoints: /health /stats /metrics /chat /teach /memory/samples</p>
</body></html>"""


def create_app(lm: HFPalimpsesteLM, corpus_pairs=None) -> "FastAPI":
    if not HAS_FASTAPI:
        raise ImportError("FastAPI not installed. Run: pip install fastapi uvicorn")

    app = FastAPI(title="PALIMPSESTE", description="Hypervectorial Cortex API")

    # CORS — allow React dev server
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    conv = Conversation(model=lm, fuzzy_threshold=0.72, learn_live=True)
    if corpus_pairs:
        conv.register_questions(corpus_pairs)
    agent = CognitiveAgent(conv=conv)

    # Cortex features: expertise, dreaming, compositional reasoning
    from palimseste.cortex import InstantExpert, Dreamer, Composer
    from palimseste.reasoning import Reasoner
    expert = InstantExpert(lm=lm)
    dreamer = Dreamer(mem=lm.mem, phi=lm.phi)
    reasoner = Reasoner(conv=conv)
    composer = Composer(reasoner=reasoner)

    # Cortex features 4-7: multimodal, meta-learning, dual memory, analogies
    from palimseste.cortex import (
        MultimodalFusion, MetaLearner, DualMemory, Analogizer,
    )
    fusion = MultimodalFusion(mem=lm.mem, encoder=lm.encoder)
    meta_learner = MetaLearner(mem=lm.mem, phi=lm.phi)
    dual_mem = DualMemory(mem=lm.mem, encoder=lm.encoder)

    # try to serve React build if it exists
    react_build = Path(__file__).resolve().parent.parent / "web" / "dist"
    if react_build.exists():
        from fastapi.staticfiles import StaticFiles

        @app.get("/", response_class=HTMLResponse)
        async def root():
            index = react_build / "index.html"
            return HTMLResponse(index.read_text(encoding="utf-8"))

        app.mount("/assets", StaticFiles(directory=str(react_build / "assets")), name="assets")
    else:
        @app.get("/", response_class=HTMLResponse)
        async def root():
            return FALLBACK_HTML

    @app.get("/health")
    async def health():
        return {"status": "ok", "D": lm.config.D, "|M|": len(lm.mem)}

    @app.get("/stats")
    async def stats():
        s = lm.stats()
        s["turns"] = conv.turn_count
        s["n_corrections"] = agent.n_corrections
        return JSONResponse(s)

    @app.get("/metrics")
    async def metrics():
        """Rich metrics for the dashboard."""
        mem_stats = lm.mem.stats()
        s = lm.stats()
        return JSONResponse({
            "model": {
                "D": s["D"],
                "n_traces": s["n_traces"],
                "n_meta": s["n_meta"],
                "vocab_size": s["vocab_size"],
                "context_window": s["context_window"],
                "theoretical_capacity_log2": s["theoretical_capacity_log2"],
            },
            "memory": {
                "mean_weight": round(mem_stats.mean_weight, 4),
                "min_weight": round(mem_stats.min_weight, 4),
                "lsh_size": mem_stats.lsh_size,
            },
            "conversation": {
                "turns": conv.turn_count,
                "n_corrections": agent.n_corrections,
                "known_questions": len(conv._known_questions),
            },
            "long_context": lm.long_context_stats(),
            "config": {
                "kernel_radius": lm.config.kernel_radius,
                "kernel_min_weight": lm.config.kernel_min_weight,
                "temperature": lm.config.temperature,
                "model_type": lm.config.model_type,
            },
        })

    @app.get("/memory/samples")
    async def memory_samples(limit: int = 50):
        """Recent memory traces for visualization."""
        mem = lm.mem
        total = len(mem)
        n = min(limit, total)
        if n == 0:
            return JSONResponse({"traces": [], "total": 0})
        # sample the last n traces
        traces = []
        for i in range(total - n, total):
            t = mem._traces[i]
            traces.append({
                "id": t.id,
                "weight": round(t.weight, 4),
                "meta": t.meta,
                "tag": t.tag or "",
                "age": total - t.id,
            })
        return JSONResponse({"traces": traces, "total": total})

    @app.get("/conversation")
    async def conversation():
        """Current conversation transcript."""
        turns = []
        for turn in conv.history:
            turns.append({"role": turn.role, "text": turn.text})
        return JSONResponse({"turns": turns, "count": conv.turn_count})

    @app.post("/chat")
    async def chat(req: ChatRequest):
        t0 = time.perf_counter()
        resp = agent.respond(req.message, max_new_tokens=req.max_tokens,
                             temperature=req.temperature, seed=None)
        dt = (time.perf_counter() - t0) * 1000
        return JSONResponse(_cog_response_to_dict(resp, dt))

    @app.post("/chat/stream")
    async def chat_stream(req: ChatRequest):
        """SSE streaming with cognitive metadata."""
        t0 = time.perf_counter()

        def generate():
            # first get the full cognitive response
            resp = agent.respond(req.message, max_new_tokens=req.max_tokens,
                                 temperature=req.temperature, seed=None)
            dt = (time.perf_counter() - t0) * 1000

            # stream the text token by token (simulate streaming)
            tokens = resp.text.split()
            for i, token in enumerate(tokens):
                chunk = {
                    "type": "token",
                    "text": token + (" " if i < len(tokens) - 1 else ""),
                    "index": i,
                }
                yield f"data: {json.dumps(chunk)}\n\n"
                time.sleep(0.02)

            # send metadata at the end
            meta = _cog_response_to_dict(resp, dt)
            meta["type"] = "metadata"
            yield f"data: {json.dumps(meta)}\n\n"

        return StreamingResponse(generate(), media_type="text/event-stream")

    @app.post("/teach")
    async def teach(req: TeachRequest):
        n = len(lm.mem)
        resp = agent.teach(req.question, req.answer)
        return JSONResponse({
            "success": True,
            "message": resp.text,
            "tokens_written": len(lm.mem) - n,
            "explanation": resp.explanation,
        })

    # ================================================================ CORTEX

    @app.post("/learn-text")
    async def learn_text(req: LearnTextRequest):
        """Feature 1: Instant Expertise — ingest a document, become expert."""
        n_before = len(lm.mem)
        result = expert.learn_from_text(req.text, document_tag=req.tag, verbose=False)
        return JSONResponse({
            "success": True,
            "n_tokens": result.n_tokens,
            "n_facts": result.n_facts,
            "n_seconds": round(result.n_seconds, 2),
            "document_tag": result.document_tag,
            "facts": [{"q": q, "a": a} for q, a in result.facts[:20]],
            "tokens_written": len(lm.mem) - n_before,
        })

    @app.post("/dream")
    async def dream(req: DreamRequest):
        """Feature 2: Dream Consolidation — get smarter while idle."""
        n_before = dreamer.n_concepts
        result = dreamer.dream(n_cycles=req.cycles, replay_batch=req.replay_batch)
        return JSONResponse({
            "success": True,
            "n_concepts_promoted": result.n_concepts_promoted,
            "n_concepts_extracted": result.n_concepts_extracted,
            "n_total_concepts": dreamer.n_concepts,
            "n_new_concepts": dreamer.n_concepts - n_before,
            "n_seconds": round(result.n_seconds, 2),
            "concept_labels": result.concept_labels[:20],
            "new_connections": result.new_connections[:10],
        })

    @app.post("/reason")
    async def reason(req: ChatRequest):
        """Feature 3: Compositional Reasoning — decompose and resolve."""
        result = composer.reason(req.message)
        return JSONResponse({
            "answer": result.answer,
            "success": result.success,
            "n_decompositions": result.n_decompositions,
            "n_hops": result.n_hops,
            "n_seconds": round(result.n_seconds, 3),
            "steps": [
                {
                    "type": s.step_type,
                    "question": s.sub_question,
                    "answer": s.sub_answer,
                    "confidence": s.confidence,
                }
                for s in result.steps
            ],
        })

    @app.get("/cortex/status")
    async def cortex_status():
        """Status of all 7 cortex features."""
        return JSONResponse({
            "expertise": {
                "n_documents": expert.n_documents,
                "documents": list(expert.documents().keys()),
            },
            "dreamer": {
                "n_concepts": dreamer.n_concepts,
                "n_consolidated": dreamer.consolidator.n_promoted,
                "n_abstracted": dreamer.abstraction.n_concepts,
            },
            "multimodal": {
                "n_bindings": fusion.n_bindings,
            },
            "meta_learning": {
                "n_adaptations": meta_learner.n_adaptations,
                "current_config": meta_learner.current_config,
            },
            "dual_memory": {
                "n_episodic": dual_mem.n_episodic,
                "n_semantic": dual_mem.n_semantic,
                "total": dual_mem.total,
            },
        })

    # ================================================================ FEATURE 4-7

    @app.post("/adapt")
    async def adapt(domain: str | None = None):
        """Feature 5: Meta-learning — self-tune kernel parameters."""
        result = meta_learner.adapt(domain=domain)
        return JSONResponse({
            "accepted": result.accepted,
            "param_changed": result.param_changed,
            "energy_before": round(result.energy_before, 4),
            "energy_after": round(result.energy_after, 4),
            "rationale": result.rationale,
            "current_config": meta_learner.current_config,
        })

    @app.post("/memory/episodic")
    async def store_episodic(content: str, tag: str = ""):
        """Feature 6: Store an episodic memory."""
        record = dual_mem.store_episodic(content, tag=tag)
        return JSONResponse({
            "success": True,
            "type": record.memory_type,
            "n_episodic": dual_mem.n_episodic,
            "n_semantic": dual_mem.n_semantic,
        })

    @app.post("/memory/semantic")
    async def store_semantic(content: str, tag: str = ""):
        """Feature 6: Store a semantic memory (persistent fact)."""
        record = dual_mem.store_semantic(content, tag=tag)
        return JSONResponse({
            "success": True,
            "type": record.memory_type,
            "n_episodic": dual_mem.n_episodic,
            "n_semantic": dual_mem.n_semantic,
        })

    @app.post("/memory/recall")
    async def recall_memory(query: str, top_k: int = 5):
        """Feature 6: Recall memories by query."""
        results = dual_mem.recall(query, top_k=top_k)
        return JSONResponse({
            "results": [
                {
                    "content": r.content,
                    "type": r.memory_type,
                    "tag": r.tag,
                    "score": round(s, 4),
                    "timestamp": r.timestamp,
                }
                for r, s in results
            ],
            "n_results": len(results),
        })

    return app


def main():
    if not HAS_FASTAPI:
        print("FastAPI not installed. Run: pip install fastapi uvicorn")
        sys.exit(1)

    parser = argparse.ArgumentParser(description="PALIMPSESTE FastAPI server")
    parser.add_argument("--model", "-m", type=str, required=True,
                        help="path to model dir or HF Hub repo id")
    parser.add_argument("--port", "-p", type=int, default=3332)
    parser.add_argument("--host", type=str, default="0.0.0.0")
    args = parser.parse_args()

    print(f"loading model from {args.model} ...", file=sys.stderr)
    try:
        lm = HFPalimpsesteLM.from_pretrained(args.model)
    except Exception as e:
        print(f"Failed to load '{args.model}': {e}", file=sys.stderr)
        # Try local fallback
        local = Path(__file__).resolve().parent.parent / "palimpseste-max"
        if local.exists():
            print(f"Falling back to local model at {local} ...", file=sys.stderr)
            lm = HFPalimpsesteLM.from_pretrained(str(local))
        else:
            raise
    print(f"loaded: D={lm.config.D:,}  |M|={len(lm.mem):,}", file=sys.stderr)

    corpus_pairs = None
    try:
        sys.path.insert(0, str(Path(__file__).resolve().parent))
        from corpus_chat import get_corpus
        corpus_pairs = get_corpus()
        print(f"registered {len(corpus_pairs)} Q/A pairs for fuzzy matching",
              file=sys.stderr)
    except ImportError:
        pass

    app = create_app(lm, corpus_pairs)
    print(f"\nPALIMPSESTE API server starting on http://{args.host}:{args.port}",
          file=sys.stderr)
    print(f"Open http://localhost:{args.port} in your browser\n", file=sys.stderr)
    uvicorn.run(app, host=args.host, port=args.port, log_level="warning")


if __name__ == "__main__":
    main()