query/generator.py

"""
Generate responses using the Ollama API.
"""
import requests
from typing import Dict
from config import OLLAMA_BASE_URL


def generate(prompt_payload: Dict) -> str:
    """
    Send prompt to Ollama and return the response text.
    
    Args:
        prompt_payload: Dictionary containing the model and messages
        
    Returns:
        Response text from the LLM
    """
    try:
        # Send request to Ollama chat endpoint
        resp = requests.post(
            f"{OLLAMA_BASE_URL}/api/chat",
            json=prompt_payload
        )
        resp.raise_for_status()
        
        # Extract response content
        response_data = resp.json()
        return response_data["message"]["content"]
    except Exception as e:
        print(f"Error generating response: {e}")
        return "Sorry, I encountered an error while generating a response."


def main():
    """Test function to verify generator works correctly."""
    print("Generator module ready for use.")


if __name__ == "__main__":
    main()