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)