main.py

from fastapi import FastAPI, HTTPException, Query
from pydantic import BaseModel
from services import is_stammering
import database  

app = FastAPI(title="RAG Translation API")

class TranslationPair(BaseModel):
    source_language: str
    target_language: str
    sentence: str
    translation: str

# --- Endpoints ---

@app.post("/pairs")
async def add_pair(pair: TranslationPair):
    """Adds a new translation pair to the database[cite: 27, 28]."""
    try:
        database.add_translation_pair(
            pair.source_language, 
            pair.target_language, 
            pair.sentence, 
            pair.translation
        )
        return "ok"  
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/prompt")
async def get_prompt(
    source_language: str = Query(..., alias="source_language"),
    target_language: str = Query(..., alias="target_language"),
    input_sentence: str = Query(..., alias="input_sentence")
):
    
    # 1. Retrieve the top 4 similar pairs from ChromaDB 
    similar_pairs = database.query_similar_pairs(
        source_language, target_language, input_sentence
    )
    
    # 2. Construct the prompt for the LLM 
    # This format helps the LLM understand the context of previous translations
    suggestions_text = "\n".join([
        f"Source: {p['source']}\nTranslation: {p['target']}" 
        for p in similar_pairs
    ])
    
    full_prompt = (
        f"Translate the following sentence from {source_language} to {target_language}.\n\n"
        f"Here are some similar examples to help you:\n"
        f"{suggestions_text}\n\n"
        f"Now, translate this: {input_sentence}"
    )
    
    return {"prompt": full_prompt}

@app.get("/stammering")
async def detect_stammering(sentence: str, translation: str):
    """
    INPUT: an input sentence and its translation.
    OUTPUT: boolean indicating if stammering was detected.
    """
    detected = is_stammering(translation)
    return detected