API Integration with Ollama¶

import requests
import json
from typing import Dict, List, Optional, Iterator

class OllamaClient:
    def __init__(self, base_url: str = "http://localhost:11434"):
        self.base_url = base_url.rstrip('/')
        self.session = requests.Session()

    def generate(self, 
                 model: str, 
                 prompt: str, 
                 stream: bool = False,
                 **kwargs) -> Dict:
        """Generate text completion"""
        url = f"{self.base_url}/api/generate"
        data = {
            "model": model,
            "prompt": prompt,
            "stream": stream,
            **kwargs
        }

        response = self.session.post(url, json=data)
        response.raise_for_status()

        if stream:
            return self._handle_stream(response)
        else:
            return response.json()

    def chat(self, 
             model: str, 
             messages: List[Dict[str, str]], 
             stream: bool = False) -> Dict:
        """Chat completion with conversation context"""
        url = f"{self.base_url}/api/chat"
        data = {
            "model": model,
            "messages": messages,
            "stream": stream
        }

        response = self.session.post(url, json=data)
        response.raise_for_status()

        if stream:
            return self._handle_stream(response)
        else:
            return response.json()

    def _handle_stream(self, response) -> Iterator[Dict]:
        """Handle streaming responses"""
        for line in response.iter_lines():
            if line:
                yield json.loads(line.decode('utf-8'))

    def list_models(self) -> List[str]:
        """List available models"""
        url = f"{self.base_url}/api/tags"
        response = self.session.get(url)
        response.raise_for_status()

        data = response.json()
        return [model['name'] for model in data.get('models', [])]

    def model_info(self, model: str) -> Dict:
        """Get model information"""
        url = f"{self.base_url}/api/show"
        data = {"name": model}

        response = self.session.post(url, json=data)
        response.raise_for_status()
        return response.json()

# Usage examples
client = OllamaClient()

# Simple generation
response = client.generate(
    model="codellama:7b",
    prompt="def fibonacci(n):",
    temperature=0.1
)
print(response['response'])

# Chat conversation
messages = [
    {"role": "user", "content": "How do I optimize this Python function?"},
    {"role": "assistant", "content": "I'd be happy to help! Please share the function you'd like to optimize."},
    {"role": "user", "content": "def slow_function():\n    time.sleep(1)\n    return 'done'"}
]

chat_response = client.chat("codellama:7b", messages)
print(chat_response['message']['content'])

# Streaming response
for chunk in client.generate("codellama:7b", "Explain async/await", stream=True):
    if 'response' in chunk:
        print(chunk['response'], end='', flush=True)

import asyncio
import aiohttp
import json
from typing import AsyncIterator

class AsyncOllamaClient:
    def __init__(self, base_url: str = "http://localhost:11434"):
        self.base_url = base_url.rstrip('/')

    async def generate(self, 
                      model: str, 
                      prompt: str, 
                      stream: bool = False) -> Dict:
        """Async generation with connection pooling"""
        url = f"{self.base_url}/api/generate"
        data = {
            "model": model,
            "prompt": prompt,
            "stream": stream
        }

        async with aiohttp.ClientSession() as session:
            async with session.post(url, json=data) as response:
                if stream:
                    return self._handle_stream(response)
                else:
                    return await response.json()

    async def _handle_stream(self, response) -> AsyncIterator[Dict]:
        """Handle async streaming responses"""
        async for line in response.content:
            if line:
                yield json.loads(line.decode('utf-8'))

    async def batch_generate(self, 
                           model: str, 
                           prompts: List[str]) -> List[Dict]:
        """Process multiple prompts concurrently"""
        tasks = [self.generate(model, prompt) for prompt in prompts]
        return await asyncio.gather(*tasks)

# Usage
async def main():
    client = AsyncOllamaClient()

    # Batch processing
    prompts = [
        "def quicksort(arr):",
        "def binary_search(arr, target):",
        "def fibonacci(n):"
    ]

    results = await client.batch_generate("codellama:7b", prompts)
    for prompt, result in zip(prompts, results):
        print(f"Prompt: {prompt}")
        print(f"Response: {result['response'][:100]}...")
        print("-" * 50)

# Run async client
asyncio.run(main())

import time
import logging
from typing import Optional, Dict, Any

class ResilientOllamaClient:
    def __init__(self, base_url="http://localhost:11434", max_retries=3):
        self.base_url = base_url
        self.max_retries = max_retries
        self.session = requests.Session()

        # Configure logging
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)

    def generate_with_retry(self, 
                           model: str, 
                           prompt: str, 
                           **kwargs) -> Optional[Dict[str, Any]]:
        """Generate with automatic retry and fallback"""

        # Try primary model
        for attempt in range(self.max_retries):
            try:
                response = self.generate(model, prompt, **kwargs)
                return response
            except requests.exceptions.ConnectionError:
                self.logger.warning(f"Connection failed, attempt {attempt + 1}/{self.max_retries}")
                time.sleep(2 ** attempt)  # Exponential backoff
            except requests.exceptions.Timeout:
                self.logger.warning(f"Request timeout, attempt {attempt + 1}/{self.max_retries}")
                time.sleep(1)
            except Exception as e:
                self.logger.error(f"Unexpected error: {e}")
                if attempt == self.max_retries - 1:
                    raise

        # Try fallback model if primary fails
        if ":" in model:
            base_model = model.split(":")[0]
            fallback_models = [f"{base_model}:7b", "codellama:7b", "mistral:7b"]

            for fallback in fallback_models:
                if fallback != model:
                    try:
                        self.logger.info(f"Trying fallback model: {fallback}")
                        return self.generate(fallback, prompt, **kwargs)
                    except Exception as e:
                        self.logger.warning(f"Fallback {fallback} failed: {e}")
                        continue

        return None

    def health_check(self) -> bool:
        """Check if Ollama service is healthy"""
        try:
            response = self.session.get(f"{self.base_url}/api/tags", timeout=5)
            return response.status_code == 200
        except Exception:
            return False

    def wait_for_service(self, timeout: int = 60) -> bool:
        """Wait for Ollama service to become available"""
        start_time = time.time()

        while time.time() - start_time < timeout:
            if self.health_check():
                return True

            self.logger.info("Waiting for Ollama service...")
            time.sleep(2)

        return False

# Usage with error handling
client = ResilientOllamaClient()

# Wait for service to be ready
if not client.wait_for_service():
    print("Ollama service not available")
    exit(1)

# Generate with automatic retry and fallback
response = client.generate_with_retry(
    model="codellama:13b",
    prompt="Write a function to handle errors gracefully"
)

if response:
    print(response['response'])
else:
    print("Failed to generate response after all retries")

import asyncio
import time
from asyncio import Queue, Semaphore
from dataclasses import dataclass
from typing import List, Callable, Any

@dataclass
class Task:
    id: str
    model: str
    prompt: str
    callback: Callable[[Any], None]
    priority: int = 0

class OllamaTaskQueue:
    def __init__(self, max_concurrent=2, rate_limit_per_minute=60):
        self.max_concurrent = max_concurrent
        self.rate_limit = rate_limit_per_minute
        self.semaphore = Semaphore(max_concurrent)
        self.queue = Queue()
        self.request_times = []

        self.client = AsyncOllamaClient()

    async def add_task(self, task: Task):
        """Add task to queue"""
        await self.queue.put(task)

    async def process_tasks(self):
        """Process tasks with rate limiting and concurrency control"""
        while True:
            try:
                # Wait for task
                task = await self.queue.get()

                # Check rate limit
                await self._check_rate_limit()

                # Process with concurrency control
                await self.semaphore.acquire()
                asyncio.create_task(self._process_single_task(task))

            except Exception as e:
                print(f"Error processing task: {e}")

    async def _check_rate_limit(self):
        """Enforce rate limiting"""
        now = time.time()
        self.request_times = [t for t in self.request_times if now - t < 60]

        if len(self.request_times) >= self.rate_limit:
            sleep_time = 60 - (now - self.request_times[0])
            if sleep_time > 0:
                await asyncio.sleep(sleep_time)

    async def _process_single_task(self, task: Task):
        """Process individual task"""
        try:
            self.request_times.append(time.time())

            response = await self.client.generate(task.model, task.prompt)
            task.callback(response)

        except Exception as e:
            task.callback({"error": str(e)})
        finally:
            self.semaphore.release()
            self.queue.task_done()

# Usage
async def main():
    queue = OllamaTaskQueue(max_concurrent=2, rate_limit_per_minute=30)

    # Start queue processor
    asyncio.create_task(queue.process_tasks())

    # Add tasks
    def handle_response(response):
        if "error" in response:
            print(f"Error: {response['error']}")
        else:
            print(f"Response: {response['response'][:100]}...")

    for i in range(10):
        task = Task(
            id=f"task_{i}",
            model="codellama:7b",
            prompt=f"Write function number {i}",
            callback=handle_response
        )
        await queue.add_task(task)

    # Wait for all tasks to complete
    await queue.queue.join()

asyncio.run(main())