from fastapi import FastAPI, APIRouter
from pydantic import BaseModel
import ollama  # Make sure Ollama is installed and available

app = FastAPI()

# Define request model for the input data
class UserInput(BaseModel):
    question: str

# Function to generate responses using Ollama
def get_ollama_response(user_input: str) -> str:
    try:
        # Run Ollama model for generating a response
        response = ollama.chat(model="llama", messages=[{"role": "user", "content": user_input}])
        return response['text']  # Ensure you're extracting the response text from Ollama's response
    except Exception as e:
        return f"Error processing request: {str(e)}"

# Create an API router
router = APIRouter()

@router.post("/ollama-response")
async def ollama_response(user_input: UserInput):
    response = get_ollama_response(user_input.question)
    return {"response": response}

# Include router into the FastAPI app
app.include_router(router)