import requests
import json

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

    def chat(self, model, messages, stream=False, format=None):
        url = f"{self.base_url}/api/chat"
        payload = {
            "model": model,
            "messages": messages,
            "stream": stream,
        }
        if format:
            payload["format"] = format

        response = requests.post(url, json=payload, stream=stream)
        response.raise_for_status()

        if stream:
            for line in response:
                if line:
                  try:
                    yield json.loads(line)
                  except json.JSONDecodeError:
                    pass
        else:
            return response.json()