local_model.py

import requests
import json
from typing import Dict, Any, List
import os

class LocalModelManager:
    def __init__(self):
        self.base_url = os.getenv("OLLAMA_API_URL", "https://ollama-api-y8pj.onrender.com")
        self.api_generate = f"{self.base_url}/api/generate"
        self.api_tags = f"{self.base_url}/api/tags"
    
    def check_ollama_status(self) -> bool:
        try:
            response = requests.get(self.api_tags, timeout=20)
            return response.status_code == 200
        except:
            return False
    
    def get_available_models(self) -> List[str]:
        try:
            response = requests.get(self.api_tags, timeout=20)
            if response.status_code == 200:
                models = response.json().get("models", [])
                return [model["name"] for model in models]
        except:
            pass
        return ["gemma3:1b"]
    
    def query_model(self, prompt: str, model: str = "gemma3:1b") -> Dict[str, Any]:
        try:
            print(f"🤖 Querying {model} at {self.api_generate}")
            
            data = {
                "model": model,
                "prompt": prompt,
                "stream": False,
                "options": {
                    "temperature": 0.7,
                    "num_predict": 400
                }
            }
            
            headers = {
                "Content-Type": "application/json",
                "Accept": "application/json"
            }
            
            response = requests.post(
                self.api_generate,
                json=data,
                headers=headers,
                timeout=120
            )
            
            print(f"Response status: {response.status_code}")
            
            if response.status_code == 200:
                result = response.json()
                return {
                    "success": True,
                    "response": result.get("response", "No response generated"),
                    "model": model,
                    "error": None
                }
            else:
                print(f"Error response: {response.text}")
                return {
                    "success": False,
                    "response": None,
                    "model": model,
                    "error": f"HTTP {response.status_code}: {response.text}"
                }
            
        except requests.exceptions.ConnectionError:
            return {
                "success": False,
                "response": None,
                "model": model,
                "error": "Connection failed. Check if Ollama service is running."
            }
        except requests.exceptions.Timeout:
            return {
                "success": False,
                "response": None,
                "model": model,
                "error": "Request timed out. Model might be busy."
            }
        except Exception as e:
            return {
                "success": False,
                "response": None,
                "model": model,
                "error": f"Error: {str(e)}"
            }

local_model_manager = LocalModelManager()