core/llm.py

import openai
import requests
import time
from typing import List, Dict, Optional
from utils.config import config

class LLMProvider:
    def __init__(self, model_name: str, timeout: int = 30, retries: int = 3):
        self.model_name = model_name
        self.timeout = timeout
        self.retries = retries

class OllamaProvider(LLMProvider):
    def generate_response(self, prompt: str, conversation_history: List[Dict]) -> Optional[str]:
        url = f"{config.ollama_host}/api/chat"
        messages = conversation_history
        
        payload = {
            "model": self.model_name,
            "messages": messages,
            "stream": False
        }
        
        for attempt in range(self.retries):
            try:
                response = requests.post(url, json=payload, timeout=self.timeout)
                response.raise_for_status()
                return response.json()["message"]["content"]
            except Exception as e:
                if attempt == self.retries - 1:
                    print(f"Error after {self.retries} attempts: {e}")
                    return None
                time.sleep(2 ** attempt)  # Exponential backoff
        return None

class HuggingFaceProvider(LLMProvider):
    def __init__(self, model_name: str, timeout: int = 30, retries: int = 3):
        super().__init__(model_name, timeout, retries)
        if not config.hf_token:
            raise ValueError("HF_TOKEN not set - required for Hugging Face provider")
        
        self.client = openai.OpenAI(
            base_url=config.hf_api_url,
            api_key=config.hf_token
        )
    
    def generate_response(self, prompt: str, conversation_history: List[Dict]) -> Optional[str]:
        try:
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=conversation_history,
                max_tokens=500,
                temperature=0.7
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"Hugging Face API error: {e}")
            return None

def send_to_ollama(prompt: str, conversation_history: List[Dict], ollama_url: str, model: str) -> Optional[str]:
    config.ollama_host = ollama_url
    provider = OllamaProvider(model)
    return provider.generate_response(prompt, conversation_history)

def send_to_hf(prompt: str, conversation_history: List[Dict]) -> Optional[str]:
    # Using a common model that works well for coaching
    provider = HuggingFaceProvider("meta-llama/Llama-2-7b-chat-hf")
    return provider.generate_response(prompt, conversation_history)