diff --git a/app/shellgen/backends/ollama.py b/app/shellgen/backends/ollama.py new file mode 100644 index 0000000..4b6a70a --- /dev/null +++ b/app/shellgen/backends/ollama.py @@ -0,0 +1,94 @@ +"""Ollama backend implementation.""" + +import ollama as ollama_client +from typing import Optional + +from .base import LLMBackend + + +class OllamaBackend(LLMBackend): + """Ollama API backend for local LLM inference.""" + + def __init__( + self, + host: str = "localhost:11434", + model: str = "codellama", + temperature: float = 0.1, + max_tokens: int = 500, + ): + """Initialize the Ollama backend. + + Args: + host: Ollama server host. + model: Model name to use. + temperature: Generation temperature. + max_tokens: Maximum tokens to generate. + """ + self.host = host + self._model = model + self.temperature = temperature + self.max_tokens = max_tokens + self._client = None + + @property + def client(self): + """Get or create Ollama client.""" + if self._client is None: + self._client = ollama_client.Client(host=self.host) + return self._client + + def generate(self, prompt: str) -> str: + """Generate response using Ollama API. + + Args: + prompt: The prompt to send. + + Returns: + Generated response text. + """ + try: + response = self.client.generate( + model=self._model, + prompt=prompt, + options={ + "temperature": self.temperature, + "num_predict": self.max_tokens, + }, + ) + + return response.get("response", "") + + except Exception as e: + raise ConnectionError(f"Ollama API error: {e}") + + def is_available(self) -> bool: + """Check if Ollama is running and model is available. + + Returns: + True if backend is available. + """ + try: + self.client.ps() + return True + except Exception: + return False + + def get_model_name(self) -> str: + """Get the current model name. + + Returns: + Model name string. + """ + return self._model + + def set_model(self, model: str) -> None: + """Set the model to use. + + Args: + model: Name of the model. + """ + self._model = model + + def close(self) -> None: + """Clean up client resources.""" + self._client = None