app.py

from fastapi import FastAPI, Depends, HTTPException
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
import time
import os
import base64
import uuid
from datetime import datetime, timedelta, timezone
from google import genai
from google.genai import types
from typing import List
import firebase_admin
from firebase_admin import credentials, auth

certificate = {
  "type": "service_account",
  "project_id": os.environ.get('project'),
  "private_key_id": os.environ.get('key_id'),
  "private_key": os.environ.get('private_key').replace('\\n', '\n'),
  "client_email": os.environ.get('client_email'),
  "client_id": os.environ.get('client_id'),
  "auth_uri": "https://accounts.google.com/o/oauth2/auth",
  "token_uri": "https://oauth2.googleapis.com/token",
  "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
  "client_x509_cert_url": os.environ.get('cert_url'),
  "universe_domain": "googleapis.com"
}

cred = credentials.Certificate(certificate)
firebase_admin.initialize_app(cred)


app = FastAPI()
security = HTTPBearer()


Tokens = []
user = []
# In-memory storage
sessions = {}  # Format: {uid: {conv_id: [history_list]}}


user_data = [] #[[email, name, gender, nationality, model name, model gender]]
user_inst = [] #[[system prompt, About user, demanded tone, custum instruction]]


Api_key = os.getenv('API_KEY')
System_instruction = '''**System Prompt for a Programmer-Oriented Coding Assistant:**\n\n> You are a highly focused, fast, and expert-level coding assistant built for professional programmers.\n> Your primary role is **to assist with code writing, debugging, refactoring, optimization, and architecture**.\n> Avoid unnecessary explanations unless asked. Do not teach—**support the user like a senior pair programmer** who assumes context and skill. Prioritize clean, correct, and efficient code.\n\n> Always:\n> * Get straight to the point.\n> * Suggest the most practical and scalable solution.\n> * Respond with complete code blocks when needed.\n> * Use strong defaults and modern conventions.\n> * Assume the user knows what they're doing.\n> * Think ahead: anticipate potential pitfalls or better approaches.\n> * Give fast, minimal answers when asked for quick help.\n\n> Only elaborate if specifically requested (e.g., “explain,” “why,” “teach,” “verbose”)'''

client = genai.Client(api_key=Api_key)


class ChatRequest(BaseModel):
    prompt: str
    conv_id: str

class NewConv(BaseModel):
    prompt: str

class VerifyRequest(BaseModel):
    uid: str
    idToken: str

class UpdateSystemRequest(BaseModel):
    key: str
    value: str

def update_details(email, name):
    user_data.append([email, name,'','','',''])
    user_inst.append(['','','',''])
    
    

def verify_access_token(auth: HTTPAuthorizationCredentials = Depends(security)):
    token = auth.credentials
    if token not in Tokens:
        raise HTTPException(
            status_code=401, 
            detail="Invalid or expired session token. Please login again."
        )
    return token
    
    
@app.post("/system/userInfo")
async def returnUserInfo(token: str = Depends(verify_access_token)):
    i = Tokens.index(token)
    return {'user_email': user_data[i][0], 'user_name': user_data[i][1], 'user_gender': user_data[i][2], 'assistant_nickname': user_data[i][4], 'assistant_gender': user_data[i][5], 'about_user': user_inst[i][1], 'assistant_tone': user_inst[i][2], 'custom_instruction': user_inst[i][3]}
    
@app.post("/update/system/userInfo")
async def update_Details(info: UpdateSystemRequest, token: str = Depends(verify_access_token)):
    i = Tokens.index(token)
    cat = info.key
    val = info.value
    if cat == 'user_name':
        user_data[i][1] = val
    if cat == 'user_gender':
        user_data[i][2] = val
    if cat == 'country':
        user_data[i][3] = val
    if cat == 'assistant_nickname':
        user_data[i][4] = val
    if cat == 'assistant_gender':
        user_data[i][5] = val
    if cat == 'system_prompt':
        user_inst[i][0] =val
    if cat == 'about_user':
        user_inst[i][1] =val
    if cat == 'assistant_tone':
        user_inst[i][2] =val
    if cat == 'custom_instruction':
        user_inst[i][3] =val
    else:
        raise HTTPException(status_code=401, detail="Invalid Data.")


@app.post("/api/verify")
def check_token(data: VerifyRequest):
    decoded_token = auth.verify_id_token(data.idToken)
    uid = decoded_token['uid']
    email = decoded_token['email']
    name = decoded_token['name']

    if decoded_token['uid'] == data.uid and decoded_token['email_verified']:
        new_token = base64.urlsafe_b64encode(uuid.uuid4().bytes).rstrip(b'=').decode()

        if uid not in user:
            user.append(uid)
            Tokens.append(new_token)
            sessions[uid] = {}
            update_details(email, name)
        else:
            idx = user.index(uid)
            user_data[idx][0] = email
            Tokens[idx] = new_token
        return {"customToken": new_token}
    else:
        raise HTTPException(status_code=401, detail="Invalid token. Verification failed.")        


def call_gemini(history: List[types.Content]):
    try:
        response = client.models.generate_content(
            model="gemma-3-4b-it",
            contents=history
        )
        return response.text
    except Exception as e:
        print(f"GenAI Error: {e}")
        raise HTTPException(status_code=500, detail="AI Generation Failed")
        

def gen_title(user, model):
    try:
        response = client.models.generate_content(
            model="gemma-3-1b-it",
            contents= title_prompt(user, model)
        )
        return response.text
    except Exception as e:
        return user

def title_prompt(user, model):
    return f'''---\n**System Role:**\n\n> You are a specialized utility model. Your sole task is to generate a concise, descriptive title (2–6 words) for a chat conversation based on the provided user query and model response.\n\n**Constraints:**\n\n* **Output only the title.** Do not include labels, quotes, full-stop or introductory text.\n* **Focus on the core intent** of the conversation.\n* **Avoid generic terms** like "Chat about..." or "Request for..."\n* **Style:** Professional, catchy, and informative.\n\n---\n\n**User Query:** {user}\n\n**Model Response:** {model}'''

def getSystemPrompt():
    text = ''
    return text #TODO: System Prompt builder later

def getConvId():
    u = uuid.uuid4().bytes[:12]
    return base64.urlsafe_b64encode(u).rstrip(b'=').decode()

def currentTime():
    # IST is UTC + 5:30
    ist_offset = timezone(timedelta(hours=5, minutes=30))
    current_time = datetime.now(ist_offset)
    return current_time.strftime("%I:%M %p")


@app.post("/new_conversation")
async def handleNewConv(new_conv: NewConv, token: str = Depends(verify_access_token)):
    a = time.time()
    i = Tokens.index(token)
    convs = sessions.get(user[i])
    history = [
        types.Content(role='model', parts=[types.Part(text=getSystemPrompt())]),
        types.Content(role='user', parts=[types.Part(text= f'time {currentTime()}\n{new_conv.prompt}')])
    ]
    text = call_gemini(history)
    if text:
        conv_id = getConvId()
        raw_history = [["model", getSystemPrompt()]]
        raw_history.append(["user", new_conv.prompt])
        raw_history.append(["model", text])
        convs[conv_id] = raw_history
        return {"title": gen_title(new_conv.prompt,text), "text": text, "conv_id": conv_id}


@app.post("/gen_resp")
async def handleChat(chat_request: ChatRequest, token: str = Depends(verify_access_token)):
    a= time.time()
    conv_id = chat_request.conv_id
    i = Tokens.index(token)
    convs = sessions.get(user[i])
    if conv_id not in convs:
        raise HTTPException(status_code=404, detail="Conversation not found")
    # Prepare history
    raw_history = convs[conv_id]
    history = [
        types.Content(role=k[0], parts=[types.Part(text=k[1])]) 
        for k in raw_history
    ]
    history.append(types.Content(role='user', parts=[types.Part(text= f'time {currentTime()}\n{chat_request.prompt}')]))
    # Generate response
    text = call_gemini(history)
    if text:
        raw_history.append(["user", chat_request.prompt])
        raw_history.append(["model", text])
        return {"text": text}