from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline
from typing import List, Dict

# Initialize the FastAPI app
app = FastAPI()

# Initialize the text generation pipeline
pipe = pipeline("text-generation", model="CyberNative-AI/Colibri_8b_v0.1")

# Define the input schema for FastAPI
class Message(BaseModel):
    role: str
    content: str

class MessagesInput(BaseModel):
    messages: List[Message]

@app.post("/generate/")
async def generate_response(messages_input: MessagesInput):
    try:
        # Convert messages to the expected format
        messages = [{"role": msg.role, "content": msg.content} for msg in messages_input.messages]

        # Generate response using the pipeline
        response = pipe(messages)
        
        # Extract generated text
        generated_text = response[0]["generated_text"]

        return {
            "response": generated_text
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Run the app
# To start the server, use the command: uvicorn filename:app --host 0.0.0.0 --port 8000