#!/usr/bin/env python3 """ AITBC Ollama Client Plugin - Submit LLM inference jobs to the network """ import httpx import json import asyncio from typing import Dict, Any, List, Optional class OllamaClient: """Client for submitting Ollama jobs to AITBC network""" def __init__(self, coordinator_url: str, api_key: str): self.coordinator_url = coordinator_url self.api_key = api_key self.client = httpx.Client() def list_available_models(self) -> List[str]: """Get available models from miners""" try: # For now, return common Ollama models # In production, this would query the network for available models return [ "deepseek-r1:14b", "qwen2.5-coder:14b", "deepseek-coder-v2:latest", "gemma3:12b", "deepcoder:latest", "deepseek-coder:6.7b-base", "llama3.2:3b-instruct-q8_0", "mistral:latest", "llama3.2:latest", "gemma3:4b", "qwen2.5:1.5b", "gemma3:1b", "lauchacarro/qwen2.5-translator:latest" ] except Exception as e: print(f"Failed to get models: {e}") return [] def submit_generation( self, model: str, prompt: str, system_prompt: Optional[str] = None, temperature: float = 0.7, max_tokens: Optional[int] = None, ttl_seconds: int = 300 ) -> Optional[str]: """Submit a text generation job""" job_payload = { "type": "generate", "model": model, "prompt": prompt, "temperature": temperature, "max_tokens": max_tokens } if system_prompt: job_payload["system_prompt"] = system_prompt return self._submit_job(job_payload, ttl_seconds) def submit_chat( self, model: str, messages: List[Dict[str, str]], temperature: float = 0.7, max_tokens: Optional[int] = None, ttl_seconds: int = 300 ) -> Optional[str]: """Submit a chat completion job""" job_payload = { "type": "chat", "model": model, "messages": messages, "temperature": temperature, "max_tokens": max_tokens } return self._submit_job(job_payload, ttl_seconds) def submit_code_generation( self, model: str, prompt: str, language: Optional[str] = None, temperature: float = 0.3, max_tokens: Optional[int] = None, ttl_seconds: int = 600 ) -> Optional[str]: """Submit a code generation job""" system_prompt = f"You are a helpful coding assistant. Generate {language or 'Python'} code." if language: system_prompt += f" Use {language} syntax." job_payload = { "type": "generate", "model": model, "prompt": prompt, "system_prompt": system_prompt, "temperature": temperature, "max_tokens": max_tokens } return self._submit_job(job_payload, ttl_seconds) def _submit_job(self, payload: Dict[str, Any], ttl_seconds: int) -> Optional[str]: """Submit job to coordinator""" job_data = { "payload": payload, "ttl_seconds": ttl_seconds } try: response = self.client.post( f"{self.coordinator_url}/v1/jobs", headers={ "Content-Type": "application/json", "X-Api-Key": self.api_key }, json=job_data ) if response.status_code == 201: job = response.json() return job['job_id'] else: print(f"❌ Failed to submit job: {response.status_code}") print(f" Response: {response.text}") return None except Exception as e: print(f"❌ Error submitting job: {e}") return None def get_job_status(self, job_id: str) -> Optional[Dict[str, Any]]: """Get job status and result""" try: response = self.client.get( f"{self.coordinator_url}/v1/jobs/{job_id}", headers={"X-Api-Key": self.api_key} ) if response.status_code == 200: return response.json() else: print(f"❌ Failed to get status: {response.status_code}") return None except Exception as e: print(f"❌ Error getting status: {e}") return None def wait_for_result(self, job_id: str, timeout: int = 60) -> Optional[Dict[str, Any]]: """Wait for job completion and return result""" import time start_time = time.time() while time.time() - start_time < timeout: status = self.get_job_status(job_id) if status: if status['state'] == 'completed': return status elif status['state'] == 'failed': print(f"❌ Job failed: {status.get('error', 'Unknown error')}") return status elif status['state'] == 'expired': print("⏰ Job expired") return status time.sleep(2) print(f"⏰ Timeout waiting for job {job_id}") return None # CLI interface def main(): import argparse parser = argparse.ArgumentParser(description="AITBC Ollama Client") parser.add_argument("--url", default="http://localhost:8001", help="Coordinator URL") parser.add_argument("--api-key", default="REDACTED_CLIENT_KEY", help="API key") subparsers = parser.add_subparsers(dest="command", help="Commands") # List models models_parser = subparsers.add_parser("models", help="List available models") # Generate text gen_parser = subparsers.add_parser("generate", help="Generate text") gen_parser.add_argument("model", help="Model name") gen_parser.add_argument("prompt", help="Text prompt") gen_parser.add_argument("--system", help="System prompt") gen_parser.add_argument("--temp", type=float, default=0.7, help="Temperature") gen_parser.add_argument("--max-tokens", type=int, help="Max tokens") # Chat chat_parser = subparsers.add_parser("chat", help="Chat completion") chat_parser.add_argument("model", help="Model name") chat_parser.add_argument("message", help="Message") chat_parser.add_argument("--temp", type=float, default=0.7, help="Temperature") # Code generation code_parser = subparsers.add_parser("code", help="Generate code") code_parser.add_argument("model", help="Model name") code_parser.add_argument("prompt", help="Code description") code_parser.add_argument("--lang", default="python", help="Programming language") # Check status status_parser = subparsers.add_parser("status", help="Check job status") status_parser.add_argument("job_id", help="Job ID") args = parser.parse_args() if not args.command: parser.print_help() return client = OllamaClient(args.url, args.api_key) if args.command == "models": models = client.list_available_models() print("🤖 Available Models:") for model in models: print(f" • {model}") elif args.command == "generate": print(f"📝 Generating with {args.model}...") job_id = client.submit_generation( args.model, args.prompt, args.system, args.temp, args.max_tokens ) if job_id: print(f"✅ Job submitted: {job_id}") result = client.wait_for_result(job_id) if result and result['state'] == 'completed': print(f"\n📄 Result:") print(result.get('result', {}).get('output', 'No output')) elif args.command == "chat": print(f"💬 Chatting with {args.model}...") messages = [{"role": "user", "content": args.message}] job_id = client.submit_chat(args.model, messages, args.temp) if job_id: print(f"✅ Job submitted: {job_id}") result = client.wait_for_result(job_id) if result and result['state'] == 'completed': print(f"\n🤖 Response:") print(result.get('result', {}).get('output', 'No response')) elif args.command == "code": print(f"💻 Generating {args.lang} code with {args.model}...") job_id = client.submit_code_generation(args.model, args.prompt, args.lang) if job_id: print(f"✅ Job submitted: {job_id}") result = client.wait_for_result(job_id) if result and result['state'] == 'completed': print(f"\n💾 Generated Code:") print(result.get('result', {}).get('output', 'No code')) elif args.command == "status": status = client.get_job_status(args.job_id) if status: print(f"📊 Job {args.job_id}:") print(f" State: {status['state']}") print(f" Miner: {status.get('assigned_miner_id', 'None')}") if status['state'] == 'completed': print(f" Cost: {status.get('result', {}).get('cost', 0)} AITBC") if __name__ == "__main__": main()