zulip-platform-plugins/agent-zero-zulip/src/agent_zero_zulip/a2a_server.py

"""A2A server for kagentz — LiteLLM with session memory + homelab project context."""
import os, json, uuid, asyncio
from datetime import datetime
from typing import Optional
import uvicorn
from fastapi import FastAPI, Request
import httpx

app = FastAPI()
LITELLM_URL = os.getenv("LITELLM_URL", "https://litellm.sysloggh.net/v1")
LITELLM_KEY = os.getenv("LITELLM_KEY", "sk-akDPhdO8qlhYNtVet6KZog")
LITELLM_MODEL = os.getenv("LITELLM_MODEL", "syslog-auto")
NOW = datetime.utcnow().strftime("%Y-%m-%d %H:%M UTC")

# Session store: session_id -> list of {"role": ..., "content": ...}
sessions: dict[str, list[dict]] = {}
# Task results: task_id -> (state, result)
task_results: dict[str, tuple[str, Optional[dict]]] = {}

HOMELAB_CONTEXT = f"""You are kagentz, a homelab automation agent for Syslog Solution LLC.

## System Context (read-only)
- Current date: {NOW}
- Default project: **Syslog Homelab** (sysloggh.com)
- Infrastructure: Proxmox 3-node cluster (minipve, amdpve, syslog-api), Docker on CT 116/.7/.14
- Zulip chat: chat.sysloggh.net (agent-hub stream)
- LLM backend: LiteLLM at litellm.sysloggh.net
- Available tools: Zulip messaging, SSH to infrastructure hosts, Docker commands

## Personality
- Professional but conversational
- Default all responses to the homelab project context unless the user specifies otherwise
- If asked about project context, state "Syslog Homelab" as your home project
- Keep responses concise and actionable
- When you don't know something, say so rather than guessing

## Your Capabilities
- You can converse about homelab infrastructure: Proxmox, Docker, networking, monitoring
- You can execute basic diagnostics via LiteLLM
- You cannot directly SSH or run commands (that's handled by other agents)
- You can discuss Zulip platform integrations, agent architectures (pi, Hermes, Agent Zero)
"""

MAX_SESSION_HISTORY = 20  # keep last 20 messages per session
SESSION_TTL = 3600  # seconds before session garbage collection


def _get_session(session_id: str) -> list[dict]:
    """Get or create a session's message history."""
    if session_id not in sessions:
        sessions[session_id] = [{"role": "system", "content": HOMELAB_CONTEXT}]
    return sessions[session_id]


def _cleanup_stale_sessions():
    """Remove old sessions periodically."""
    # Simple TTL: sessions that haven't been touched
    pass  # Sessions are cleaned on restart (ephemeral)


@app.get("/.well-known/agent.json")
async def agent_card():
    return {
        "name": "kagentz",
        "version": "1.0.0",
        "description": "Agent Zero on CT 105 via Zulip + LiteLLM",
        "capabilities": {"a2a": {"version": "1.0"}},
        "skills": [{"id": "chat", "name": "Chat"}],
    }


@app.post("/a2a")
async def a2a_endpoint(request: Request):
    body = await request.json()
    method = body.get("method", "")
    
    if method == "tasks/send":
        params = body.get("params", {})
        msg = params.get("message", {})
        parts = msg.get("parts", [])
        text = " ".join(p.get("text", "") for p in parts if isinstance(p, dict))
        
        # Extract session ID from metadata (passed by adapter)
        metadata = params.get("metadata", {}) or {}
        session_id = metadata.get("session_id", f"anon-{uuid.uuid4().hex[:8]}")
        task_id = f"az-{uuid.uuid4().hex[:12]}"
        
        asyncio.create_task(call_litellm(task_id, session_id, text))
        return {
            "jsonrpc": "2.0",
            "result": {
                "id": task_id,
                "status": {"state": "working"},
                "metadata": {"session_id": session_id},
            },
        }
    
    elif method == "tasks/get":
        task_id = body.get("params", {}).get("id", "")
        entry = task_results.get(task_id)
        if not entry:
            return {
                "jsonrpc": "2.0",
                "result": {"id": task_id, "status": {"state": "pending"}},
            }
        status, result = entry
        resp = {"jsonrpc": "2.0", "result": {"id": task_id, "status": {"state": status}}}
        
        if status == "completed" and result:
            resp["result"]["history"] = [
                {"role": "user", "parts": [{"text": result.get("input", "")}]},
                {"role": "assistant", "parts": [{"text": result.get("output", "")}]},
            ]
            # Pass back the session_id so the adapter can reuse it
            if result.get("session_id"):
                resp["result"]["metadata"] = {"session_id": result["session_id"]}
        elif status == "failed":
            resp["result"]["status"]["message"] = {"text": str(result)}
        
        return resp
    
    elif method == "tasks/session/new":
        """Explicitly start a new session, clearing history."""
        metadata = body.get("params", {}).get("metadata", {}) or {}
        session_id = metadata.get("session_id", f"new-{uuid.uuid4().hex[:8]}")
        sessions.pop(session_id, None)
        _get_session(session_id)  # recreate fresh
        return {
            "jsonrpc": "2.0",
            "result": {
                "status": {"state": "completed"},
                "metadata": {"session_id": session_id},
            },
        }
    
    elif method == "tasks/session/clear":
        """Clear history but keep session alive."""
        session_id = body.get("params", {}).get("metadata", {}).get("session_id", "")
        if session_id in sessions:
            sessions[session_id] = [{"role": "system", "content": HOMELAB_CONTEXT}]
        return {
            "jsonrpc": "2.0",
            "result": {"status": {"state": "completed"}},
        }
    
    return {"jsonrpc": "2.0", "error": {"code": -32601, "message": "Method not found"}}


async def call_litellm(task_id: str, session_id: str, message: str):
    """Call LiteLLM with conversation history from session."""
    try:
        session = _get_session(session_id)
        
        # Add user message to history
        session.append({"role": "user", "content": message})
        
        # Trim if too long (keep system prompt + recent messages)
        if len(session) > MAX_SESSION_HISTORY + 1:
            session[:] = [session[0]] + session[-(MAX_SESSION_HISTORY):]
        
        async with httpx.AsyncClient(timeout=120.0) as client:
            resp = await client.post(
                f"{LITELLM_URL}/chat/completions",
                headers={
                    "Authorization": f"Bearer {LITELLM_KEY}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": LITELLM_MODEL,
                    "messages": session,
                    "max_tokens": 2000,
                },
            )
            if resp.is_success:
                output = resp.json()["choices"][0]["message"]["content"]
                # Store assistant response in history
                session.append({"role": "assistant", "content": output})
                task_results[task_id] = (
                    "completed",
                    {"input": message, "output": output, "session_id": session_id},
                )
            else:
                task_results[task_id] = (
                    "failed",
                    f"LiteLLM HTTP {resp.status_code}: {resp.text[:200]}",
                )
    except Exception as e:
        task_results[task_id] = ("failed", str(e))


if __name__ == "__main__":
    print(f"kagentz A2A server on port 8001 (model={LITELLM_MODEL})")
    print(f"Project context: Syslog Homelab")
    print(f"Session memory: {MAX_SESSION_HISTORY} messages per session")
    uvicorn.run(app, host="0.0.0.0", port=8001, log_level="info")