Update app.py

app.py

@@ -1,16 +1,54 @@
-from fastapi import FastAPI, Response, Cookie
-from pydantic import BaseModel, Field
-import time, os, base64, uuid, requests
+from fastapi import FastAPI, Depends, HTTPException
+from fastapi.security import HTTPAuthorizationCredentials
+from pydantic import BaseModel
+import time
+import os
+import base64
+import uuid
 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'),
+  "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)
+
+
+value = os.environ.get('YOUR_ENV_KEY')
 
 app = FastAPI()
 
 
 Tokens = []
-# In-memory storage (Better to use a Dict)
-sessions = {}  # Format: {token: {conv_id: [history_list]}}
+user = []
+# In-memory storage
+sessions = {}  # Format: {uid: {conv_id: [history_list]}}
+
+
+user_details = [] #[[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
@@ -19,84 +57,98 @@ class ChatRequest(BaseModel):
 class NewConv(BaseModel):
     prompt: str
 
-def gen(prompt):
+class VerifyRequest(BaseModel):
+    uid: str
+    idToken: str
+
+
+@app.post("/api/verify")
+def check_token(data: VerifyRequest):
+    decoded_token = auth.verify_id_token(data.idToken)
+    uid = decoded_token['uid']
+
+    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] = {}
+        else:
+            idx = user.index(uid)
+            Tokens[idx] = 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=prompt)
+        response = client.models.generate_content(
+            model="gemma-3-4b-it",
+            contents=history
+        )
         return response.text
     except Exception as e:
-        print(f"API Error: {e}")
-        return None
+        print(f"GenAI Error: {e}")
+        raise HTTPException(status_code=500, detail="AI Generation Failed")
+
 
 def getSystemPrompt():
-    return text #TODO: System Prompt builder
+    text = ''
+    return text #TODO: System Prompt builder later
 
 def getConvId():
     u = uuid.uuid4().bytes[:12]
     return base64.urlsafe_b64encode(u).rstrip(b'=').decode()
 
+
+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("/new_conversation")
-async def handleNewConv(new_conv: NewConv, token: str = Cookie(None)):
+async def handleNewConv(new_conv: NewConv, token: str = Depends(verify_access_token)):
     a = time.time()
-    user_convs = sessions.get(token)
-    if not user_convs:
-        return "Error" # TODO: Error implementation
+    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=new_conv.prompt)])
     ]
-    text = gen(history)
+    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])
-        user_convs[conv_id] = raw_history
-        return text
-
-    return None # TODO: Error implementation
-
-
-@app.post("/new_conv")
-async def handleNewConv(new_conv: NewConv):
-    a = time.time()
-
-    return {'text': str(a), 'conv_id': getConvId()} # TODO: Error implementation
-
-
-@app.post("/gen_resp")
-async def handleNewConv(new_conv: NewConv):
-    a = time.time()
-
-    return {'text': str(a)} # TODO: Error implementation
+        convs[conv_id] = raw_history
+        return {"title": new_conv.prompt, "text": text, "conv_id": conv_id}
 
 
 @app.post("/response")
-async def handleChat(chat_request: ChatRequest, token: str = Cookie(None)):
+async def handleChat(chat_request: ChatRequest, token: str = Depends(verify_access_token)):
     a= time.time()
     conv_id = chat_request.conv_id
-    # Check if token and conversation exist
-    user_convs = sessions.get(token)
-    if not user_convs or conv_id not in user_convs:
-        return "Please stop. Just refresh the page." # TODO: Error implementation
+    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 = user_convs[conv_id]
+    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=chat_request.prompt)]))
     # Generate response
-    text = gen(history)
+    text = call_gemini(history)
     if text:
         raw_history.append(["user", chat_request.prompt])
         raw_history.append(["model", text])
-        # Optional: Update history with new response here
-        return text
-
-    return None # TODO: Error implementation
-
-
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+        return {"text": text}