Add shellgen core and backends modules
This commit is contained in:
94
app/shellgen/backends/ollama.py
Normal file
94
app/shellgen/backends/ollama.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user