File size: 5,632 Bytes
fe9fe77 adcbb74 fe9fe77 adcbb74 6c0182b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
import os
import json
from dotenv import load_dotenv
from pydantic import BaseModel
import google.generativeai as genai
from fastapi import FastAPI, HTTPException
import uvicorn
from prompts import PROMPTS
from llm_pipeline import example_odia_answer_json, example_odia_question_json
# Setup
load_dotenv()
# Check for required environment variables
google_api_key = os.getenv("GOOGLE_API_KEY")
if not google_api_key:
raise ValueError("GOOGLE_API_KEY not found in environment variables")
genai.configure(api_key=google_api_key)
model = genai.GenerativeModel(os.getenv("LLM_MODEL", "gemini-pro"))
LANGUAGE = "Odia"
# Models
class QuestionRequest(BaseModel):
question: str
class LLMResponseModel(BaseModel):
question_content: str
answer_language: str = LANGUAGE
reasoning_content: str
answer_content: str
def create_prompt(user_odia_question: str) -> str:
SIMPLE_PROMPT = PROMPTS["odia_reasoning_generation_prompt"]
prompt = SIMPLE_PROMPT.format(
user_odia_question=user_odia_question,
example_odia_question_json=example_odia_question_json,
example_answer_json=example_odia_answer_json
)
return prompt
# Functions
def chat_with_model(prompt: str) -> str:
try:
response = model.generate_content(prompt)
return response.text if response.text else "Error: Empty response"
except Exception as e:
return f"Error: {str(e)}"
def clean_json_text(text: str) -> str:
if text.startswith("Error:"):
return text
# Remove markdown code blocks
text = text.strip()
if text.startswith("```"):
lines = text.split('\n')
if len(lines) > 2:
text = '\n'.join(lines[1:-1])
else:
text = text.strip("`").replace("json", "", 1).strip()
# Extract JSON content
first = text.find("{")
last = text.rfind("}")
if first != -1 and last != -1:
return text[first:last+1]
return text
def validate_output(raw_output: str, original_question: str):
cleaned = clean_json_text(raw_output)
if cleaned.startswith("Error:"):
return {
"question_content": original_question,
"answer_language": LANGUAGE,
"reasoning_content": f"Error occurred: {cleaned}",
"answer_content": "Unable to generate answer due to error",
"error": cleaned
}
try:
# Try to parse and validate JSON
parsed_data = json.loads(cleaned)
validated = LLMResponseModel(**parsed_data)
return validated.model_dump()
except json.JSONDecodeError as je:
return {
"question_content": original_question,
"answer_language": LANGUAGE,
"reasoning_content": f"JSON parsing failed: {str(je)}",
"answer_content": "Unable to parse model response",
"error": f"JSON Error: {str(je)}"
}
except Exception as e:
return {
"question_content": original_question,
"answer_language": LANGUAGE,
"reasoning_content": f"Validation failed: {str(e)}",
"answer_content": "Unable to validate model response",
"error": f"Validation Error: {str(e)}"
}
def run_pipeline(question: str):
try:
# Use simple prompt if PROMPTS not available
prompt =create_prompt(user_odia_question=question)
raw_output = chat_with_model(prompt)
return validate_output(raw_output, question)
except Exception as e:
return {
"question_content": question,
"answer_language": LANGUAGE,
"reasoning_content": f"Pipeline error: {str(e)}",
"answer_content": "Unable to process question",
"error": f"Pipeline Error: {str(e)}"
}
# API
app = FastAPI(title="Odia Question Answering API", version="0.1.0")
@app.get("/")
async def root():
return {"message": "Odia Question Answering API is running", "status": "healthy"}
@app.get("/health")
async def health_check():
try:
# Test model connectivity
test_response = model.generate_content("Test")
return {
"status": "healthy",
"model": os.getenv("LLM_MODEL", "gemini-pro"),
"api_configured": bool(google_api_key)
}
except Exception as e:
return {
"status": "unhealthy",
"error": str(e),
"api_configured": bool(google_api_key)
}
@app.post("/generate")
async def generate_answer(request: QuestionRequest):
try:
if not request.question.strip():
raise HTTPException(status_code=400, detail="Question cannot be empty")
result = run_pipeline(request.question.strip())
# Check for critical errors that should return 500
if "error" in result and any(err_type in result["error"] for err_type in ["Error: ", "Pipeline Error:"]):
raise HTTPException(status_code=500, detail=f"Processing failed: {result['error']}")
return {"success": True, "data": result}
except HTTPException:
raise
except Exception as e:
raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")
if __name__ == "__main__":
print("Starting Odia Question Answering API...")
print(f"Google API Key configured: {'Yes' if google_api_key else 'No'}")
print(f"Model: {os.getenv('LLM_MODEL', 'gemini-pro')}")
host = os.getenv("ANSWER_SERVICE_HOST", "0.0.0.0")
port = int(os.getenv("ANSWER_SERVICE_PORT", "9000"))
uvicorn.run(app, host=0.0.0.0, port=9000) |