293 lines
9.7 KiB
Python
Executable File
293 lines
9.7 KiB
Python
Executable File
#!/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="${CLIENT_API_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()
|