custom_ollama.py

import requests

# Function to call the FastAPI server
def query_ollama(prompt):
    url = "http://localhost:8000/generate/"
    payload = {"prompt": prompt}
    response = requests.post(url, json=payload)
    
    if response.status_code == 200:
        return response.json()["response"]
    else:
        raise Exception(f"Failed to generate response: {response.text}")