| """ |
| Ollama Client for SPARKNET |
| Handles communication with local Ollama LLM models |
| """ |
|
|
| import ollama |
| from typing import List, Dict, Optional, Generator, Any |
| from loguru import logger |
| import json |
|
|
|
|
| class OllamaClient: |
| """Client for interacting with Ollama LLM models.""" |
|
|
| def __init__( |
| self, |
| host: str = "localhost", |
| port: int = 11434, |
| default_model: str = "llama3.2:latest", |
| timeout: int = 300, |
| ): |
| """ |
| Initialize Ollama client. |
| |
| Args: |
| host: Ollama server host |
| port: Ollama server port |
| default_model: Default model to use |
| timeout: Request timeout in seconds |
| """ |
| self.host = host |
| self.port = port |
| self.base_url = f"http://{host}:{port}" |
| self.default_model = default_model |
| self.timeout = timeout |
| self.client = ollama.Client(host=self.base_url) |
|
|
| logger.info(f"Initialized Ollama client: {self.base_url}") |
|
|
| def list_models(self) -> List[Dict[str, Any]]: |
| """ |
| List available models. |
| |
| Returns: |
| List of model information dictionaries |
| """ |
| try: |
| response = self.client.list() |
| models = response.get("models", []) |
| logger.info(f"Found {len(models)} available models") |
| return models |
| except Exception as e: |
| logger.error(f"Error listing models: {e}") |
| return [] |
|
|
| def pull_model(self, model_name: str) -> bool: |
| """ |
| Pull/download a model. |
| |
| Args: |
| model_name: Name of the model to pull |
| |
| Returns: |
| True if successful, False otherwise |
| """ |
| try: |
| logger.info(f"Pulling model: {model_name}") |
| self.client.pull(model_name) |
| logger.info(f"Successfully pulled model: {model_name}") |
| return True |
| except Exception as e: |
| logger.error(f"Error pulling model {model_name}: {e}") |
| return False |
|
|
| def generate( |
| self, |
| prompt: str, |
| model: Optional[str] = None, |
| system: Optional[str] = None, |
| temperature: float = 0.7, |
| max_tokens: Optional[int] = None, |
| stream: bool = False, |
| **kwargs, |
| ) -> str | Generator[str, None, None]: |
| """ |
| Generate completion from a prompt. |
| |
| Args: |
| prompt: Input prompt |
| model: Model to use (default: self.default_model) |
| system: System prompt |
| temperature: Sampling temperature |
| max_tokens: Maximum tokens to generate |
| stream: Whether to stream the response |
| **kwargs: Additional generation parameters |
| |
| Returns: |
| Generated text or generator if streaming |
| """ |
| model = model or self.default_model |
|
|
| options = { |
| "temperature": temperature, |
| } |
| if max_tokens: |
| options["num_predict"] = max_tokens |
|
|
| options.update(kwargs) |
|
|
| try: |
| logger.debug(f"Generating with model {model}, prompt length: {len(prompt)}") |
|
|
| if stream: |
| return self._generate_stream(prompt, model, system, options) |
| else: |
| response = self.client.generate( |
| model=model, |
| prompt=prompt, |
| system=system, |
| options=options, |
| ) |
| generated_text = response.get("response", "") |
| logger.debug(f"Generated {len(generated_text)} characters") |
| return generated_text |
|
|
| except Exception as e: |
| logger.error(f"Error generating completion: {e}") |
| return "" |
|
|
| def _generate_stream( |
| self, |
| prompt: str, |
| model: str, |
| system: Optional[str], |
| options: Dict, |
| ) -> Generator[str, None, None]: |
| """ |
| Generate streaming completion. |
| |
| Args: |
| prompt: Input prompt |
| model: Model to use |
| system: System prompt |
| options: Generation options |
| |
| Yields: |
| Generated text chunks |
| """ |
| try: |
| stream = self.client.generate( |
| model=model, |
| prompt=prompt, |
| system=system, |
| options=options, |
| stream=True, |
| ) |
|
|
| for chunk in stream: |
| if "response" in chunk: |
| yield chunk["response"] |
|
|
| except Exception as e: |
| logger.error(f"Error in streaming generation: {e}") |
| yield "" |
|
|
| def chat( |
| self, |
| messages: List[Dict[str, str]], |
| model: Optional[str] = None, |
| temperature: float = 0.7, |
| stream: bool = False, |
| **kwargs, |
| ) -> str | Generator[str, None, None]: |
| """ |
| Chat completion with conversation history. |
| |
| Args: |
| messages: List of message dicts with 'role' and 'content' |
| model: Model to use (default: self.default_model) |
| temperature: Sampling temperature |
| stream: Whether to stream the response |
| **kwargs: Additional chat parameters |
| |
| Returns: |
| Response text or generator if streaming |
| """ |
| model = model or self.default_model |
|
|
| options = { |
| "temperature": temperature, |
| } |
| options.update(kwargs) |
|
|
| try: |
| logger.debug(f"Chat with model {model}, {len(messages)} messages") |
|
|
| if stream: |
| return self._chat_stream(messages, model, options) |
| else: |
| response = self.client.chat( |
| model=model, |
| messages=messages, |
| options=options, |
| ) |
| message = response.get("message", {}) |
| content = message.get("content", "") |
| logger.debug(f"Chat response: {len(content)} characters") |
| return content |
|
|
| except Exception as e: |
| logger.error(f"Error in chat completion: {e}") |
| return "" |
|
|
| def _chat_stream( |
| self, |
| messages: List[Dict[str, str]], |
| model: str, |
| options: Dict, |
| ) -> Generator[str, None, None]: |
| """ |
| Streaming chat completion. |
| |
| Args: |
| messages: List of message dicts |
| model: Model to use |
| options: Chat options |
| |
| Yields: |
| Response text chunks |
| """ |
| try: |
| stream = self.client.chat( |
| model=model, |
| messages=messages, |
| options=options, |
| stream=True, |
| ) |
|
|
| for chunk in stream: |
| if "message" in chunk: |
| message = chunk["message"] |
| if "content" in message: |
| yield message["content"] |
|
|
| except Exception as e: |
| logger.error(f"Error in streaming chat: {e}") |
| yield "" |
|
|
| def embed( |
| self, |
| text: str | List[str], |
| model: str = "nomic-embed-text:latest", |
| ) -> List[List[float]]: |
| """ |
| Generate embeddings for text. |
| |
| Args: |
| text: Text or list of texts to embed |
| model: Embedding model to use |
| |
| Returns: |
| List of embedding vectors |
| """ |
| try: |
| if isinstance(text, str): |
| text = [text] |
|
|
| logger.debug(f"Generating embeddings for {len(text)} texts") |
|
|
| embeddings = [] |
| for t in text: |
| response = self.client.embeddings(model=model, prompt=t) |
| embedding = response.get("embedding", []) |
| embeddings.append(embedding) |
|
|
| logger.debug(f"Generated {len(embeddings)} embeddings") |
| return embeddings |
|
|
| except Exception as e: |
| logger.error(f"Error generating embeddings: {e}") |
| return [] |
|
|
| def count_tokens(self, text: str) -> int: |
| """ |
| Estimate token count for text. |
| Simple estimation: ~4 characters per token for English text. |
| |
| Args: |
| text: Text to count tokens for |
| |
| Returns: |
| Estimated token count |
| """ |
| |
| return len(text) // 4 |
|
|
| def is_available(self) -> bool: |
| """ |
| Check if Ollama server is available. |
| |
| Returns: |
| True if server is responding, False otherwise |
| """ |
| try: |
| self.list_models() |
| return True |
| except Exception: |
| return False |
|
|
|
|
| |
| _ollama_client: Optional[OllamaClient] = None |
|
|
|
|
| def get_ollama_client( |
| host: str = "localhost", |
| port: int = 11434, |
| default_model: str = "llama3.2:latest", |
| ) -> OllamaClient: |
| """Get or create the global Ollama client instance.""" |
| global _ollama_client |
| if _ollama_client is None: |
| _ollama_client = OllamaClient(host=host, port=port, default_model=default_model) |
| return _ollama_client |
|
|
