File size: 2,376 Bytes
c4070e8 |
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 |
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import uvicorn
import os
from dotenv import load_dotenv
load_dotenv()
import json
import google.generativeai as genai
from tools import calculator, email_drafter, pdf_qa, code_inspector
# Configure Gemini API
try:
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
model = genai.GenerativeModel('gemini-2.5-flash')
except Exception as e:
print(f"Error configuring Gemini or loading model: {e}")
model = None # Set model to None if configuration fails
app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")
@app.get("/", response_class=HTMLResponse)
async def read_root():
with open("templates/index.html") as f:
return HTMLResponse(content=f.read(), status_code=200)
@app.post("/mcp")
async def mcp_endpoint(prompt: str = Form(...), file: UploadFile = File(None)):
if file:
file_path = f"temp_{file.filename}"
with open(file_path, "wb") as buffer:
buffer.write(await file.read())
if "pdf" in file.content_type:
response = pdf_qa.answer_pdf_question(file_path, prompt, model)
else:
response = "Unsupported file type."
os.remove(file_path)
return JSONResponse(content={"response": response})
tool_prompts = json.load(open("prompts/tool_prompts.json"))
if "sum" in prompt.lower():
response = calculator.calculate(prompt, "sum")
elif "multiply" in prompt.lower():
response = calculator.calculate(prompt, "multiply")
elif "subtract" in prompt.lower():
response = calculator.calculate(prompt, "subtract")
elif "divide" in prompt.lower():
response = calculator.calculate(prompt, "divide")
elif "draft an email" in prompt.lower():
response = email_drafter.draft_email(prompt, model)
elif "check code" in prompt.lower():
response = code_inspector.check_code(prompt, model)
else:
response = "I'm sorry, I don't understand that request."
return JSONResponse(content={"response": response})
if __name__ == "__main__":
uvicorn.run(app, host="127.0.0.1", port=7860)
|