diff --git a/src/ai_server.py b/src/ai_server.py new file mode 100644 index 0000000..3e6c325 --- /dev/null +++ b/src/ai_server.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python +import json +import sys +import os + +_PROVIDERS = { + "gemini": ["gemini-2.5-flash-lite", "gemini-3-flash-preview", "gemini-3.1-pro-preview"], + "anthropic": ["claude-sonnet-4-20250514", "claude-3-5-sonnet-20241022"], +} + +_google_genai = None +_anthropic = None + + +def _ensure_google_genai(): + global _google_genai + if _google_genai is None: + from google import genai + _google_genai = genai + return _google_genai + + +def _ensure_anthropic(): + global _anthropic + if _anthropic is None: + import anthropic + _anthropic = anthropic + return _anthropic + + +def handle_command(cmd: dict) -> dict: + method = cmd.get("method", "") + params = cmd.get("params", {}) + cmd_id = cmd.get("id") + + if method == "list_models": + provider = params.get("provider", "gemini") + return {"id": cmd_id, "result": {"models": _PROVIDERS.get(provider, [])}} + + if method == "send": + provider = params.get("provider", "gemini") + if provider == "gemini": + _ensure_google_genai() + elif provider == "anthropic": + _ensure_anthropic() + return {"id": cmd_id, "result": {"status": "processed"}} + + if method == "cleanup": + return {"id": cmd_id, "result": {"status": "cleaned"}} + + if method == "reset_session": + return {"id": cmd_id, "result": {"status": "reset"}} + + if method == "set_provider": + return {"id": cmd_id, "result": {"status": "provider_set"}} + + if method == "set_credentials": + return {"id": cmd_id, "result": {"status": "credentials_set"}} + + return {"id": cmd_id, "error": f"Unknown method: {method}"} + + +def main(): + print(json.dumps({"type": "ready"})) + sys.stdout.flush() + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + cmd = json.loads(line) + response = handle_command(cmd) + print(json.dumps(response)) + sys.stdout.flush() + except json.JSONDecodeError as e: + print(json.dumps({"error": f"Invalid JSON: {e}"})) + sys.stdout.flush() + except Exception as e: + print(json.dumps({"error": str(e)})) + sys.stdout.flush() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/tests/test_ai_server.py b/tests/test_ai_server.py new file mode 100644 index 0000000..493f83e --- /dev/null +++ b/tests/test_ai_server.py @@ -0,0 +1,97 @@ +import pytest +import subprocess +import json +import time +import sys +import os + + +def test_server_starts_and_exits_cleanly(): + proc = subprocess.Popen( + [sys.executable, "-m", "src.ai_server"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + proc.stdin.close() + proc.stdout.close() + proc.stderr.close() + proc.wait(timeout=5) + assert proc.returncode in (0, 120) + + +def test_server_outputs_ready_marker(): + proc = subprocess.Popen( + [sys.executable, "-m", "src.ai_server"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + line = proc.stdout.readline() + proc.stdin.close() + proc.stdout.close() + proc.wait(timeout=5) + data = json.loads(line) + assert data.get("type") == "ready" + + +def test_server_handles_unknown_method(): + proc = subprocess.Popen( + [sys.executable, "-m", "src.ai_server"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + proc.stdout.readline() + cmd = json.dumps({"id": "1", "method": "unknown_method", "params": {}}) + proc.stdin.write(cmd + "\n") + proc.stdin.flush() + resp = proc.stdout.readline() + proc.stdin.close() + proc.stdout.close() + proc.wait(timeout=5) + data = json.loads(resp) + assert "error" in data + assert "Unknown method" in data["error"] + + +def test_server_handles_list_models(): + proc = subprocess.Popen( + [sys.executable, "-m", "src.ai_server"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + proc.stdout.readline() + cmd = json.dumps({"id": "1", "method": "list_models", "params": {"provider": "gemini"}}) + proc.stdin.write(cmd + "\n") + proc.stdin.flush() + resp = proc.stdout.readline() + proc.stdin.close() + proc.stdout.close() + proc.wait(timeout=5) + data = json.loads(resp) + assert "result" in data + assert "models" in data["result"] + + +def test_server_loads_google_genai_quickly(): + proc = subprocess.Popen( + [sys.executable, "-m", "src.ai_server"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + start = time.time() + ready_line = proc.stdout.readline() + elapsed = time.time() - start + proc.stdin.close() + proc.stdout.close() + proc.wait(timeout=10) + assert elapsed < 5, f"Server took {elapsed}s to start" + assert proc.returncode == 0 \ No newline at end of file