AgapHost/openai/cognee-llm/server.js

const http = require('http');
const fs = require('fs');
const path = require('path');
const crypto = require('crypto');
const { spawn } = require('child_process');

const PORT = 8011;
const MODEL_ID = 'cognee-llm';
const TIMEOUT_MS = 5 * 60 * 1000; // one-shot structured calls; generous but bounded
// Bounded parallelism: SPIKE-FINDINGS.md gate 5 flagged the Kimi subscription as a
// single-seat, interactive-oriented plan — batch cognify must not hammer it with
// unbounded concurrent CLI spawns (rate-limit/throttle risk on a shared live account).
const MAX_CONCURRENCY = 3;

const WORKSPACE = '/workspace';
fs.mkdirSync(WORKSPACE, { recursive: true });

// The CLI has no raw sampling-temperature knob (it's an agent loop, not a
// completions API) — "low temperature" for structured extraction is enforced
// via an instruction preamble instead, prepended to whatever system prompt
// the caller (Cognee) supplies.
const STRUCTURED_SYSTEM_PREAMBLE = [
  'You are a stateless structured-extraction engine.',
  'This is a one-shot call with no memory of prior calls: do not reference earlier turns.',
  'Respond deterministically and concisely. When asked for JSON, output raw JSON only',
  '- no prose, no markdown code fences, no commentary before or after.',
].join(' ');

// --- message helpers ---------------------------------------------------------
// Text only, no media parts: this wrapper's policy is no-media/no-MCP, unlike
// adolf-llm which persists inbound images and lets the CLI's ReadMediaFile
// tool read them.
function textOf(msg) {
  const c = msg.content;
  if (Array.isArray(c)) return c.map(p => (typeof p.text === 'string' ? p.text : '')).join('\n');
  return c == null ? '' : String(c);
}

function buildPrompt(messages) {
  const systemParts = messages.filter(m => m.role === 'system').map(textOf);
  const rest = messages.filter(m => m.role !== 'system');
  const preamble = [STRUCTURED_SYSTEM_PREAMBLE, ...systemParts].join('\n\n');
  const transcript = rest
    .map(m => `${m.role === 'assistant' ? 'Assistant' : 'User'}: ${textOf(m)}`)
    .join('\n\n');
  return `${preamble}\n\n${transcript}`.trim();
}

// --- bounded concurrency queue -----------------------------------------------
let active = 0;
const queue = [];
function drain() {
  if (queue.length && active < MAX_CONCURRENCY) queue.shift()();
}
function withSlot(fn) {
  return new Promise((resolve, reject) => {
    const run = () => {
      active++;
      fn().then(
        v => { active--; drain(); resolve(v); },
        e => { active--; drain(); reject(e); },
      );
    };
    if (active < MAX_CONCURRENCY) run();
    else queue.push(run);
  });
}

// --- kimi invocation: stateless one-shot, no resume --------------------------
// Fresh temp dir per call, NO -r/-S session flag, discard the dir after.
// Returns the assembled text from --output-format stream-json:
//   {"role":"assistant","content":"..."}
// (reuses the same parse core as kimi-agent/server.js's runKimi, minus the
// resume/session-id bookkeeping that wrapper needs and this one deliberately
// does not).
function runKimi({ prompt, cwd }) {
  return new Promise((resolve, reject) => {
    const args = ['-p', prompt, '--output-format', 'stream-json'];
    const child = spawn('kimi', args, { cwd, timeout: TIMEOUT_MS });

    let stdout = '';
    let stderr = '';
    child.stdout.on('data', d => { stdout += d; });
    child.stderr.on('data', d => { stderr += d; });

    child.on('error', reject);
    child.on('close', code => {
      const parts = [];
      for (const line of stdout.split('\n')) {
        const t = line.trim();
        if (!t) continue;
        let obj;
        try { obj = JSON.parse(t); } catch { continue; }
        if (obj.role === 'assistant' && obj.content) parts.push(obj.content);
      }
      const text = parts.join('').trim();
      if (!text && code !== 0) {
        reject(new Error(`kimi exited ${code}: ${stderr.slice(0, 2000)}`));
      } else {
        resolve(text);
      }
    });
  });
}

async function handleTurn(messages) {
  const prompt = buildPrompt(messages || []);
  const reqId = crypto.randomUUID();
  const dir = path.join(WORKSPACE, reqId);
  fs.mkdirSync(dir, { recursive: true });
  try {
    return await withSlot(() => runKimi({ prompt, cwd: dir }));
  } finally {
    // Stateless one-shot: nothing about this call is meant to survive it, so
    // the temp dir is discarded unconditionally, success or failure.
    fs.rm(dir, { recursive: true, force: true }, () => {});
  }
}

// --- OpenAI-compatible HTTP surface (non-streaming only) ---------------------
function completionBody(text) {
  return {
    id: `chatcmpl-${Date.now()}`,
    object: 'chat.completion',
    created: Math.floor(Date.now() / 1000),
    model: MODEL_ID,
    choices: [{
      index: 0,
      message: { role: 'assistant', content: text },
      finish_reason: 'stop',
    }],
    usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 },
  };
}

const server = http.createServer((req, res) => {
  if (req.method === 'GET' && req.url === '/v1/models') {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({
      object: 'list',
      data: [{ id: MODEL_ID, object: 'model', owned_by: 'moonshot' }],
    }));
    return;
  }

  if (req.method === 'POST' && req.url === '/v1/chat/completions') {
    let body = '';
    req.on('data', d => { body += d; });
    req.on('end', async () => {
      let parsed;
      try {
        parsed = JSON.parse(body);
      } catch {
        res.writeHead(400, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ error: 'invalid JSON body' }));
        return;
      }

      try {
        const text = await handleTurn(parsed.messages || []);
        // Non-streaming policy: always return the full body even if the
        // caller sets stream:true. Cognee's batch cognify has no use for SSE,
        // and a one-shot call has nothing to incrementally stream anyway.
        res.writeHead(200, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify(completionBody(text)));
      } catch (err) {
        res.writeHead(500, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ error: String(err.message || err) }));
      }
    });
    return;
  }

  res.writeHead(404, { 'Content-Type': 'application/json' });
  res.end(JSON.stringify({ error: 'not found' }));
});

server.listen(PORT, () => console.log(`cognee-llm wrapper listening on :${PORT}`));